sdk/go/httpserver/logger_test.go

   1 // Copyright (C) The Arvados Authors. All rights reserved.
   2 //
   3 // SPDX-License-Identifier: Apache-2.0
   4
   5 package httpserver
   6
   7 import (
   8         "bytes"
   9         "context"
  10         "encoding/json"
  11         "fmt"
  12         "io/ioutil"
  13         "net"
  14         "net/http"
  15         "net/http/httptest"
  16         "testing"
  17         "time"
  18
  19         "git.arvados.org/arvados.git/sdk/go/ctxlog"
  20         "github.com/sirupsen/logrus"
  21         check "gopkg.in/check.v1"
  22 )
  23
  24 func Test(t *testing.T) {
  25         check.TestingT(t)
  26 }
  27
  28 var _ = check.Suite(&Suite{})
  29
  30 type Suite struct {
  31         ctx     context.Context
  32         log     *logrus.Logger
  33         logdata *bytes.Buffer
  34 }
  35
  36 func (s *Suite) SetUpTest(c *check.C) {
  37         s.logdata = bytes.NewBuffer(nil)
  38         s.log = logrus.New()
  39         s.log.Out = s.logdata
  40         s.log.Formatter = &logrus.JSONFormatter{
  41                 TimestampFormat: time.RFC3339Nano,
  42         }
  43         s.ctx = ctxlog.Context(context.Background(), s.log)
  44 }
  45
  46 func (s *Suite) TestWithDeadline(c *check.C) {
  47         req, err := http.NewRequest("GET", "https://foo.example/bar", nil)
  48         c.Assert(err, check.IsNil)
  49
  50         // Short timeout cancels context in <1s
  51         resp := httptest.NewRecorder()
  52         HandlerWithDeadline(time.Millisecond, http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
  53                 select {
  54                 case <-req.Context().Done():
  55                         w.Write([]byte("ok"))
  56                 case <-time.After(time.Second):
  57                         c.Error("timed out")
  58                 }
  59         })).ServeHTTP(resp, req.WithContext(s.ctx))
  60         c.Check(resp.Body.String(), check.Equals, "ok")
  61
  62         // Long timeout does not cancel context in <1ms
  63         resp = httptest.NewRecorder()
  64         HandlerWithDeadline(time.Second, http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
  65                 select {
  66                 case <-req.Context().Done():
  67                         c.Error("request context done too soon")
  68                 case <-time.After(time.Millisecond):
  69                         w.Write([]byte("ok"))
  70                 }
  71         })).ServeHTTP(resp, req.WithContext(s.ctx))
  72         c.Check(resp.Body.String(), check.Equals, "ok")
  73 }
  74
  75 func (s *Suite) TestExemptFromDeadline(c *check.C) {
  76         srv := Server{
  77                 Addr: ":",
  78                 Server: http.Server{
  79                         Handler: HandlerWithDeadline(time.Millisecond, http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
  80                                 if req.URL.Path == "/exempt" {
  81                                         ExemptFromDeadline(req)
  82                                 }
  83                                 time.Sleep(time.Second / 10)
  84                                 if req.Context().Err() != nil {
  85                                         w.WriteHeader(499)
  86                                         return
  87                                 }
  88                                 fmt.Fprint(w, "ok")
  89                         })),
  90                         BaseContext: func(net.Listener) context.Context { return s.ctx },
  91                 },
  92         }
  93         srv.Start()
  94         defer srv.Close()
  95
  96         resp, err := http.Get("http://" + srv.Addr + "/normal")
  97         c.Assert(err, check.IsNil)
  98         c.Check(resp.StatusCode, check.Equals, 499)
  99         body, err := ioutil.ReadAll(resp.Body)
 100         c.Check(err, check.IsNil)
 101         c.Check(string(body), check.Equals, "")
 102
 103         resp, err = http.Get("http://" + srv.Addr + "/exempt")
 104         c.Assert(err, check.IsNil)
 105         c.Check(resp.StatusCode, check.Equals, 200)
 106         body, err = ioutil.ReadAll(resp.Body)
 107         c.Check(err, check.IsNil)
 108         c.Check(string(body), check.Equals, "ok")
 109 }
 110
 111 func (s *Suite) TestLogRequests(c *check.C) {
 112         h := AddRequestIDs(LogRequests(
 113                 http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
 114                         w.Write([]byte("hello world"))
 115                 })))
 116
 117         req, err := http.NewRequest("GET", "https://foo.example/bar", nil)
 118         req.Header.Set("X-Forwarded-For", "1.2.3.4:12345")
 119         c.Assert(err, check.IsNil)
 120         resp := httptest.NewRecorder()
 121
 122         h.ServeHTTP(resp, req.WithContext(s.ctx))
 123
 124         dec := json.NewDecoder(s.logdata)
 125
 126         gotReq := make(map[string]interface{})
 127         err = dec.Decode(&gotReq)
 128         c.Check(err, check.IsNil)
 129         c.Logf("%#v", gotReq)
 130         c.Check(gotReq["RequestID"], check.Matches, "req-[a-z0-9]{20}")
 131         c.Check(gotReq["reqForwardedFor"], check.Equals, "1.2.3.4:12345")
 132         c.Check(gotReq["msg"], check.Equals, "request")
 133
 134         gotResp := make(map[string]interface{})
 135         err = dec.Decode(&gotResp)
 136         c.Check(err, check.IsNil)
 137         c.Logf("%#v", gotResp)
 138         c.Check(gotResp["RequestID"], check.Equals, gotReq["RequestID"])
 139         c.Check(gotResp["reqForwardedFor"], check.Equals, "1.2.3.4:12345")
 140         c.Check(gotResp["msg"], check.Equals, "response")
 141
 142         c.Assert(gotResp["time"], check.FitsTypeOf, "")
 143         _, err = time.Parse(time.RFC3339Nano, gotResp["time"].(string))
 144         c.Check(err, check.IsNil)
 145
 146         for _, key := range []string{"timeToStatus", "timeWriteBody", "timeTotal"} {
 147                 c.Assert(gotResp[key], check.FitsTypeOf, float64(0))
 148                 c.Check(gotResp[key].(float64), check.Not(check.Equals), float64(0))
 149         }
 150 }
 151
 152 func (s *Suite) TestLogErrorBody(c *check.C) {
 153         dec := json.NewDecoder(s.logdata)
 154
 155         for _, trial := range []struct {
 156                 label      string
 157                 statusCode int
 158                 sentBody   string
 159                 expectLog  bool
 160                 expectBody string
 161         }{
 162                 {"ok", 200, "hello world", false, ""},
 163                 {"redir", 302, "<a href='http://foo.example/baz'>redir</a>", false, ""},
 164                 {"4xx short body", 400, "oops", true, "oops"},
 165                 {"4xx long body", 400, fmt.Sprintf("%0*d", sniffBytes*2, 1), true, fmt.Sprintf("%0*d", sniffBytes, 0)},
 166                 {"5xx empty body", 500, "", true, ""},
 167         } {
 168                 comment := check.Commentf("in trial: %q", trial.label)
 169
 170                 req, err := http.NewRequest("GET", "https://foo.example/bar", nil)
 171                 c.Assert(err, check.IsNil)
 172                 resp := httptest.NewRecorder()
 173
 174                 LogRequests(
 175                         http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
 176                                 w.WriteHeader(trial.statusCode)
 177                                 w.Write([]byte(trial.sentBody))
 178                         }),
 179                 ).ServeHTTP(resp, req.WithContext(s.ctx))
 180
 181                 gotReq := make(map[string]interface{})
 182                 err = dec.Decode(&gotReq)
 183                 c.Check(err, check.IsNil)
 184                 c.Logf("%#v", gotReq)
 185                 gotResp := make(map[string]interface{})
 186                 err = dec.Decode(&gotResp)
 187                 c.Check(err, check.IsNil)
 188                 c.Logf("%#v", gotResp)
 189                 if trial.expectLog {
 190                         c.Check(gotResp["respBody"], check.Equals, trial.expectBody, comment)
 191                 } else {
 192                         c.Check(gotResp["respBody"], check.IsNil, comment)
 193                 }
 194         }
 195 }