middleware/proxy: Allow non-HTTP upstreams to be health checked (#589)

Allow HTTP health check to be performed against a regular DNS upstream server.

TODO: Add tests.

2 files changed, 123 insertions, 5 deletions

diff --git a/middleware/proxy/upstream.go b/middleware/proxy/upstream.go
index 59e1a534f..cabcd9ba8 100644
--- a/middleware/proxy/upstream.go
+++ b/middleware/proxy/upstream.go
@@ -4,8 +4,10 @@ import (
 	"fmt"
 	"io"
 	"io/ioutil"
+	"log"
 	"net"
 	"net/http"
+	"net/url"
 	"strconv"
 	"strings"
 	"sync/atomic"
@@ -14,7 +16,6 @@ import (
 	"github.com/coredns/coredns/middleware"
 	"github.com/coredns/coredns/middleware/pkg/dnsutil"
 	"github.com/coredns/coredns/middleware/pkg/tls"
-
 	"github.com/mholt/caddy/caddyfile"
 	"github.com/miekg/dns"
 )
@@ -229,16 +230,38 @@ func parseBlock(c *caddyfile.Dispenser, u *staticUpstream) error {
 
 func (u *staticUpstream) healthCheck() {
 	for _, host := range u.Hosts {
-		port := ""
+		var hostName, checkPort string
+
+		// The DNS server might be an HTTP server.  If so, extract its name.
+		if url, err := url.Parse(host.Name); err == nil {
+			hostName = url.Host
+		} else {
+			hostName = host.Name
+		}
+
+		// Extract the port number from the parsed server name.
+		checkHostName, checkPort, err := net.SplitHostPort(hostName)
+		if err != nil {
+			checkHostName = hostName
+		}
+
 		if u.HealthCheck.Port != "" {
-			port = ":" + u.HealthCheck.Port
+			checkPort = u.HealthCheck.Port
 		}
-		hostURL := host.Name + port + u.HealthCheck.Path
+
+		hostURL := "http://" + net.JoinHostPort(checkHostName, checkPort) + u.HealthCheck.Path
+		host.Unhealthy = false
+
 		if r, err := http.Get(hostURL); err == nil {
 			io.Copy(ioutil.Discard, r.Body)
 			r.Body.Close()
-			host.Unhealthy = r.StatusCode < 200 || r.StatusCode >= 400
+			if r.StatusCode < 200 || r.StatusCode >= 400 {
+				log.Printf("[WARNING] Health check URL %s returned HTTP code %d\n",
+					hostURL, r.StatusCode)
+				host.Unhealthy = true
+			}
 		} else {
+			log.Printf("[WARNING] Health check probe failed: %v\n", err)
 			host.Unhealthy = true
 		}
 	}
diff --git a/test/proxy_http_health_test.go b/test/proxy_http_health_test.go
new file mode 100644
index 000000000..9bfc3b209
--- /dev/null
+++ b/test/proxy_http_health_test.go
@@ -0,0 +1,95 @@
+package test
+
+import (
+	"io"
+	"io/ioutil"
+	"log"
+	"net/http"
+	"net/http/httptest"
+	"net/url"
+	"testing"
+
+	"github.com/coredns/coredns/middleware/proxy"
+	"github.com/coredns/coredns/middleware/test"
+	"github.com/coredns/coredns/request"
+	"github.com/miekg/dns"
+)
+
+func TestProxyWithHTTPCheckOK(t *testing.T) {
+	log.SetOutput(ioutil.Discard)
+
+	healthCheckServer := httptest.NewServer(http.HandlerFunc(
+		func(w http.ResponseWriter, r *http.Request) {
+			w.WriteHeader(http.StatusOK)
+			io.WriteString(w, "OK\n")
+		}))
+	defer healthCheckServer.Close()
+
+	healthCheckURL, err := url.Parse(healthCheckServer.URL)
+	if err != nil {
+		t.Fatal(err)
+	}
+	healthCheckPort := healthCheckURL.Port()
+
+	name, rm, err := test.TempFile(".", exampleOrg)
+	if err != nil {
+		t.Fatalf("failed to create zone: %s", err)
+	}
+	defer rm()
+
+	// We have to bind to 127.0.0.1 because the server started by
+	// httptest.NewServer does, and the IP addresses of the backend
+	// DNS and HTTP servers must match.
+	authoritativeCorefile := `example.org:0 {
+	   bind 127.0.0.1
+       file ` + name + `
+}
+`
+
+	authoritativeInstance, err := CoreDNSServer(authoritativeCorefile)
+	if err != nil {
+		t.Fatalf("Could not get CoreDNS authoritative instance: %s", err)
+	}
+
+	authoritativeAddr, _ := CoreDNSServerPorts(authoritativeInstance, 0)
+	if authoritativeAddr == "" {
+		t.Fatalf("Could not get CoreDNS authoritative instance UDP listening port")
+	}
+	defer authoritativeInstance.Stop()
+
+	proxyCorefile := `example.org:0 {
+    proxy . ` + authoritativeAddr + ` {
+		health_check /health:` + healthCheckPort + ` 1s
+
+	}
+}
+`
+
+	proxyInstance, err := CoreDNSServer(proxyCorefile)
+	if err != nil {
+		t.Fatalf("Could not get CoreDNS proxy instance: %s", err)
+	}
+
+	proxyAddr, _ := CoreDNSServerPorts(proxyInstance, 0)
+	if proxyAddr == "" {
+		t.Fatalf("Could not get CoreDNS proxy instance UDP listening port")
+	}
+	defer proxyInstance.Stop()
+
+	p := proxy.NewLookup([]string{proxyAddr})
+	state := request.Request{W: &test.ResponseWriter{}, Req: new(dns.Msg)}
+	resp, err := p.Lookup(state, "example.org.", dns.TypeA)
+	if err != nil {
+		t.Fatal("Expected to receive reply, but didn't")
+	}
+	// expect answer section with A record in it
+	if len(resp.Answer) == 0 {
+		t.Fatalf("Expected to at least one RR in the answer section, got none: %s", resp)
+	}
+	if resp.Answer[0].Header().Rrtype != dns.TypeA {
+		t.Errorf("Expected RR to A, got: %d", resp.Answer[0].Header().Rrtype)
+	}
+	if resp.Answer[0].(*dns.A).A.String() != "127.0.0.1" {
+		t.Errorf("Expected 127.0.0.1, got: %s", resp.Answer[0].(*dns.A).A.String())
+	}
+}