11906: keepstore ping

Arvados-DCO-1.1-Signed-off-by: Radhika Chippada <radhika@curoverse.com>

diff --git a/services/keepstore/config.go b/services/keepstore/config.go
index 19a614a1fa6fc42fd3cbe0166a9ed6c6e465043c..0a3ece41ad9d88589513c30a878f172276afd93c 100644 (file)
--- a/services/keepstore/config.go
+++ b/services/keepstore/config.go
@@ -40,6 +40,8 @@ type Config struct {
        blobSigningKey  []byte
        systemAuthToken string
        debugLogf       func(string, ...interface{})
+
+       ManagementToken string
 }
 
 var theConfig = DefaultConfig()

diff --git a/services/keepstore/handler_test.go b/services/keepstore/handler_test.go
index 751a4a77e3c8ba66431ca30d1d58f3cde2208f5d..7429d7ad41660799aaa90b3ba188bafb225d3978 100644 (file)
--- a/services/keepstore/handler_test.go
+++ b/services/keepstore/handler_test.go
@@ -27,6 +27,7 @@ import (
        "time"
 
        "git.curoverse.com/arvados.git/sdk/go/arvados"
+       "git.curoverse.com/arvados.git/sdk/go/arvadostest"
 )
 
 // A RequestTester represents the parameters for an HTTP request to
@@ -827,6 +828,18 @@ func IssueRequest(rt *RequestTester) *httptest.ResponseRecorder {
        return response
 }
 
+func IssueHealthCheckRequest(rt *RequestTester) *httptest.ResponseRecorder {
+       response := httptest.NewRecorder()
+       body := bytes.NewReader(rt.requestBody)
+       req, _ := http.NewRequest(rt.method, rt.uri, body)
+       if rt.apiToken != "" {
+               req.Header.Set("Authorization", "Bearer "+rt.apiToken)
+       }
+       loggingRouter := MakeRESTRouter()
+       loggingRouter.ServeHTTP(response, req)
+       return response
+}
+
 // ExpectStatusCode checks whether a response has the specified status code,
 // and reports a test failure if not.
 func ExpectStatusCode(
@@ -1140,3 +1153,61 @@ func TestUntrashHandlerWithNoWritableVolumes(t *testing.T) {
                http.StatusNotFound,
                response)
 }
+
+func TestHealthCheckPing(t *testing.T) {
+       defer teardown()
+
+       KeepVM = MakeTestVolumeManager(2)
+       defer KeepVM.Close()
+
+       // ping when disabled
+       theConfig.ManagementToken = ""
+       pingReq := &RequestTester{
+               method: "GET",
+               uri:    "/_health/ping",
+       }
+       response := IssueHealthCheckRequest(pingReq)
+       ExpectStatusCode(t,
+               "disabled",
+               http.StatusNotFound,
+               response)
+
+       // ping with no token
+       theConfig.ManagementToken = arvadostest.ManagementToken
+       pingReq = &RequestTester{
+               method: "GET",
+               uri:    "/_health/ping",
+       }
+       response = IssueHealthCheckRequest(pingReq)
+       ExpectStatusCode(t,
+               "authorization required",
+               http.StatusUnauthorized,
+               response)
+
+       // ping with wrong token
+       pingReq = &RequestTester{
+               method:   "GET",
+               uri:      "/_health/ping",
+               apiToken: "youarenotwelcomehere",
+       }
+       response = IssueHealthCheckRequest(pingReq)
+       ExpectStatusCode(t,
+               "authorization error",
+               http.StatusForbidden,
+               response)
+
+       // ping with management token
+       pingReq = &RequestTester{
+               method:   "GET",
+               uri:      "/_health/ping",
+               apiToken: arvadostest.ManagementToken,
+       }
+       response = IssueHealthCheckRequest(pingReq)
+       ExpectStatusCode(t,
+               "",
+               http.StatusOK,
+               response)
+       if !strings.Contains(response.Body.String(), `{"health":"OK"}`) {
+               t.Errorf("expected response to include %s: got %s", `{"health":"OK"}`, response.Body.String())
+       }
+}

diff --git a/services/keepstore/handlers.go b/services/keepstore/handlers.go
index f197196c416ae3af23c7b68bca1d0af65376c6d1..5e85ed6216d69ba43deed5d744e6d6bb7c79c99f 100644 (file)
--- a/services/keepstore/handlers.go
+++ b/services/keepstore/handlers.go
@@ -78,6 +78,9 @@ func MakeRESTRouter() *router {
        // Untrash moves blocks from trash back into store
        rest.HandleFunc(`/untrash/{hash:[0-9a-f]{32}}`, UntrashHandler).Methods("PUT")
 
+       // Health check ping
+       rest.HandleFunc(`/_health/ping`, HealthCheckPingHandler).Methods("GET")
+
        // Any request which does not match any of these routes gets
        // 400 Bad Request.
        rest.NotFoundHandler = http.HandlerFunc(BadRequestHandler)
@@ -617,6 +620,39 @@ func UntrashHandler(resp http.ResponseWriter, req *http.Request) {
        }
 }
 
+var pingResponseOK = map[string]string{"health": "OK"}
+
+// HealthCheckPingHandler processes "GET /_health/ping" requests
+func HealthCheckPingHandler(resp http.ResponseWriter, req *http.Request) {
+       healthCheckDo(resp, req, pingResponseOK)
+}
+
+func healthCheckDo(resp http.ResponseWriter, req *http.Request, v interface{}) {
+       msg, code := healthCheckAuth(resp, req)
+       if msg != "" {
+               http.Error(resp, msg, code)
+               return
+       }
+
+       ok, err := json.Marshal(v)
+       if err != nil {
+               http.Error(resp, err.Error(), 500)
+       }
+
+       resp.Write(ok)
+}
+
+func healthCheckAuth(resp http.ResponseWriter, req *http.Request) (string, int) {
+       if theConfig.ManagementToken == "" {
+               return "disabled", http.StatusNotFound
+       } else if h := req.Header.Get("Authorization"); h == "" {
+               return "authorization required", http.StatusUnauthorized
+       } else if h != "Bearer "+theConfig.ManagementToken {
+               return "authorization error", http.StatusForbidden
+       }
+       return "", 0
+}
+
 // GetBlock and PutBlock implement lower-level code for handling
 // blocks by rooting through volumes connected to the local machine.
 // Once the handler has determined that system policy permits the