Commit 078c9915 authored by Martin Redmond's avatar Martin Redmond Committed by Matt Holt

proxy: custom upstream health check by body string, closes #324 (#1691)

parent bf7b2548
package proxy
import (
"bytes"
"fmt"
"io"
"io/ioutil"
......@@ -44,6 +45,7 @@ type staticUpstream struct {
Timeout time.Duration
Host string
Port string
ContentString string
}
WithoutPathPrefix string
IgnoredSubPaths []string
......@@ -337,6 +339,11 @@ func parseBlock(c *caddyfile.Dispenser, u *staticUpstream) error {
return c.Errf("invalid health_check_port '%s'", port)
}
u.HealthCheck.Port = port
case "health_check_contains":
if !c.NextArg() {
return c.ArgErr()
}
u.HealthCheck.ContentString = c.Val()
case "header_upstream":
var header, value string
if !c.Args(&header, &value) {
......@@ -402,27 +409,42 @@ func (u *staticUpstream) healthCheck() {
}
hostURL += u.HealthCheck.Path
var unhealthy bool
unhealthy := func() bool {
// set up request, needed to be able to modify headers
// possible errors are bad HTTP methods or un-parsable urls
req, err := http.NewRequest("GET", hostURL, nil)
if err != nil {
unhealthy = true
} else {
return true
}
// set host for request going upstream
if u.HealthCheck.Host != "" {
req.Host = u.HealthCheck.Host
}
if r, err := u.HealthCheck.Client.Do(req); err == nil {
r, err := u.HealthCheck.Client.Do(req)
if err != nil {
return true
}
defer func() {
io.Copy(ioutil.Discard, r.Body)
r.Body.Close()
unhealthy = r.StatusCode < 200 || r.StatusCode >= 400
} else {
unhealthy = true
}()
if r.StatusCode < 200 || r.StatusCode >= 400 {
return true
}
if u.HealthCheck.ContentString == "" { // don't check for content string
return false
}
// TODO ReadAll will be replaced if deemed necessary
// See https://github.com/mholt/caddy/pull/1691
buf, err := ioutil.ReadAll(r.Body)
if err != nil {
return true
}
if bytes.Contains(buf, []byte(u.HealthCheck.ContentString)) {
return false
}
return true
}()
if unhealthy {
atomic.StoreInt32(&host.Unhealthy, 1)
} else {
......
......@@ -448,3 +448,56 @@ func TestHealthCheckPort(t *testing.T) {
})
}
func TestHealthCheckContentString(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
fmt.Fprintf(w, "blablabla good blablabla")
r.Body.Close()
}))
_, port, err := net.SplitHostPort(server.Listener.Addr().String())
if err != nil {
t.Fatal(err)
}
defer server.Close()
tests := []struct {
config string
shouldContain bool
}{
{"proxy / localhost:" + port +
" { health_check /testhealth " +
" health_check_contains good\n}",
true,
},
{"proxy / localhost:" + port + " {\n health_check /testhealth health_check_port " + port +
" \n health_check_contains bad\n}",
false,
},
}
for i, test := range tests {
u, err := NewStaticUpstreams(caddyfile.NewDispenser("Testfile", strings.NewReader(test.config)), "")
if err != nil {
t.Error("Expected no error. Test %d Got:", i, err.Error())
}
for _, upstream := range u {
staticUpstream, ok := upstream.(*staticUpstream)
if !ok {
t.Errorf("Type mismatch: %#v", upstream)
continue
}
staticUpstream.healthCheck()
for _, host := range staticUpstream.Hosts {
if test.shouldContain && atomic.LoadInt32(&host.Unhealthy) == 0 {
// healthcheck url was hit and the required test string was found
continue
}
if !test.shouldContain && atomic.LoadInt32(&host.Unhealthy) != 0 {
// healthcheck url was hit and the required string was not found
continue
}
t.Errorf("Health check bad response")
}
upstream.Stop()
}
}
}
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment