diff --git a/cli/healthz_test.go b/cli/healthz_test.go new file mode 100644 index 0000000000..0adf246652 --- /dev/null +++ b/cli/healthz_test.go @@ -0,0 +1,44 @@ +package cli + +import ( + "errors" + "net/http" + "net/http/httptest" + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestHealthz(t *testing.T) { + failing := healthz(map[string]interface{}{ + "mock": healthcheckFunc(func() error { + return errors.New("health check failed") + })}) + ok := healthz(map[string]interface{}{ + "mock": healthcheckFunc(func() error { + return nil + })}) + + var httpTests = []struct { + wantHeader int + handler http.Handler + }{ + {200, ok}, + {500, failing}, + } + for _, tt := range httpTests { + t.Run("", func(t *testing.T) { + rr := httptest.NewRecorder() + req := httptest.NewRequest("GET", "/healthz", nil) + tt.handler.ServeHTTP(rr, req) + assert.Equal(t, rr.Code, tt.wantHeader) + }) + } + +} + +type healthcheckFunc func() error + +func (fn healthcheckFunc) HealthCheck() error { + return fn() +} diff --git a/cli/serve.go b/cli/serve.go index 72b7b83319..da90a855bb 100644 --- a/cli/serve.go +++ b/cli/serve.go @@ -1,6 +1,7 @@ package cli import ( + "encoding/json" "flag" "fmt" "net/http" @@ -141,7 +142,13 @@ the way that the kolide server works. apiHandler = service.WithSetup(svc, logger, apiHandler) } } + + healthCheckers := map[string]interface{}{ + "datastore": ds, + "query_result_store": resultStore, + } http.Handle("/api/", apiHandler) + http.Handle("/healthz", healthz(healthCheckers)) http.Handle("/version", version.Handler()) http.Handle("/metrics", prometheus.Handler()) http.Handle("/assets/", service.ServeStaticAssets("/assets/")) @@ -176,3 +183,33 @@ the way that the kolide server works. return serveCmd } + +// healthz is an http handler which responds with either +// 200 OK if the server can successfuly communicate with it's backends or +// 500 if any of the backends are reporting an issue. +func healthz(deps map[string]interface{}) http.HandlerFunc { + type healthChecker interface { + HealthCheck() error + } + + return func(w http.ResponseWriter, r *http.Request) { + errs := make(map[string]string) + for name, dep := range deps { + if hc, ok := dep.(healthChecker); ok { + err := hc.HealthCheck() + if err != nil { + errs[name] = err.Error() + } + } + } + + if len(errs) > 0 { + w.WriteHeader(http.StatusInternalServerError) + enc := json.NewEncoder(w) + enc.SetIndent("", " ") + enc.Encode(map[string]interface{}{ + "errors": errs, + }) + } + } +} diff --git a/server/datastore/mysql/datastore.go b/server/datastore/mysql/datastore.go index cb2c3d563e..dacb155c74 100644 --- a/server/datastore/mysql/datastore.go +++ b/server/datastore/mysql/datastore.go @@ -113,6 +113,12 @@ func (d *Datastore) Drop() error { } +// HealthCheck returns an error if the MySQL backend is not healthy. +func (d *Datastore) HealthCheck() error { + _, err := d.db.Exec("select 1") + return err +} + // Close frees resources associated with underlying mysql connection func (d *Datastore) Close() error { return d.db.Close() diff --git a/server/pubsub/redis_query_results.go b/server/pubsub/redis_query_results.go index cc0e223f7c..f5b0cccfdb 100644 --- a/server/pubsub/redis_query_results.go +++ b/server/pubsub/redis_query_results.go @@ -154,3 +154,13 @@ func (r *redisQueryResults) ReadChannel(ctx context.Context, query kolide.Distri }() return outChannel, nil } + +// HealthCheck verifies that the redis backend can be pinged, returning an error +// otherwise. +func (r *redisQueryResults) HealthCheck() error { + conn := r.pool.Get() + defer conn.Close() + + _, err := conn.Do("PING") + return err +}