sdk/go/health/handler.go

   1 // Copyright (C) The Arvados Authors. All rights reserved.
   2 //
   3 // SPDX-License-Identifier: Apache-2.0
   4
   5 package health
   6
   7 import (
   8         "encoding/json"
   9         "errors"
  10         "net/http"
  11         "strings"
  12         "sync"
  13 )
  14
  15 // Func is a health-check function: it returns nil when healthy, an
  16 // error when not.
  17 type Func func() error
  18
  19 // Routes is a map of URI path to health-check function.
  20 type Routes map[string]Func
  21
  22 // Handler is an http.Handler that responds to authenticated
  23 // health-check requests with JSON responses like {"health":"OK"} or
  24 // {"health":"ERROR","error":"error text"}.
  25 //
  26 // Fields of a Handler should not be changed after the Handler is
  27 // first used.
  28 type Handler struct {
  29         setupOnce sync.Once
  30         mux       *http.ServeMux
  31
  32         // Authentication token. If empty, all requests will return 404.
  33         Token string
  34
  35         // Route prefix, typically "/_health/".
  36         Prefix string
  37
  38         // Map of URI paths to health-check Func. The prefix is
  39         // omitted: Routes["foo"] is the health check invoked by a
  40         // request to "{Prefix}/foo".
  41         //
  42         // If "ping" is not listed here, it will be added
  43         // automatically and will always return a "healthy" response.
  44         Routes Routes
  45
  46         // If non-nil, Log is called after handling each request. The
  47         // error argument is nil if the request was successfully
  48         // authenticated and served, even if the health check itself
  49         // failed.
  50         Log func(*http.Request, error)
  51 }
  52
  53 // ServeHTTP implements http.Handler.
  54 func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
  55         h.setupOnce.Do(h.setup)
  56         h.mux.ServeHTTP(w, r)
  57 }
  58
  59 func (h *Handler) setup() {
  60         h.mux = http.NewServeMux()
  61         prefix := h.Prefix
  62         if !strings.HasSuffix(prefix, "/") {
  63                 prefix = prefix + "/"
  64         }
  65         for name, fn := range h.Routes {
  66                 h.mux.Handle(prefix+name, h.healthJSON(fn))
  67         }
  68         if _, ok := h.Routes["ping"]; !ok {
  69                 h.mux.Handle(prefix+"ping", h.healthJSON(func() error { return nil }))
  70         }
  71 }
  72
  73 var (
  74         healthyBody     = []byte(`{"health":"OK"}` + "\n")
  75         errNotFound     = errors.New(http.StatusText(http.StatusNotFound))
  76         errUnauthorized = errors.New(http.StatusText(http.StatusUnauthorized))
  77         errForbidden    = errors.New(http.StatusText(http.StatusForbidden))
  78 )
  79
  80 func (h *Handler) healthJSON(fn Func) http.Handler {
  81         return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  82                 var err error
  83                 defer func() {
  84                         if h.Log != nil {
  85                                 h.Log(r, err)
  86                         }
  87                 }()
  88                 if h.Token == "" {
  89                         http.Error(w, "disabled", http.StatusNotFound)
  90                         err = errNotFound
  91                 } else if ah := r.Header.Get("Authorization"); ah == "" {
  92                         http.Error(w, "authorization required", http.StatusUnauthorized)
  93                         err = errUnauthorized
  94                 } else if ah != "Bearer "+h.Token {
  95                         http.Error(w, "authorization error", http.StatusForbidden)
  96                         err = errForbidden
  97                 } else if err = fn(); err == nil {
  98                         w.Header().Set("Content-Type", "application/json")
  99                         w.Write(healthyBody)
 100                 } else {
 101                         w.Header().Set("Content-Type", "application/json")
 102                         enc := json.NewEncoder(w)
 103                         err = enc.Encode(map[string]string{
 104                                 "health": "ERROR",
 105                                 "error":  err.Error(),
 106                         })
 107                 }
 108         })
 109 }