|
| 1 | +package httpserver |
| 2 | + |
| 3 | +import ( |
| 4 | + "bytes" |
| 5 | + "io" |
| 6 | + "net/http" |
| 7 | + "net/http/httptest" |
| 8 | + "testing" |
| 9 | +) |
| 10 | + |
| 11 | +type countingBody struct { |
| 12 | + r io.Reader |
| 13 | + read int64 |
| 14 | +} |
| 15 | + |
| 16 | +func (c *countingBody) Read(p []byte) (int, error) { |
| 17 | + n, err := c.r.Read(p) |
| 18 | + c.read += int64(n) |
| 19 | + return n, err |
| 20 | +} |
| 21 | + |
| 22 | +func (c *countingBody) Close() error { return nil } |
| 23 | + |
| 24 | +// TestLogLayerBoundsChunkedBody checks that a body without Content-Length |
| 25 | +// (chunked) is not read past the dump limit before the next handler runs, |
| 26 | +// and that the next handler still gets the whole body. |
| 27 | +func TestLogLayerBoundsChunkedBody(t *testing.T) { |
| 28 | + EnableVerbose = true |
| 29 | + defer func() { EnableVerbose = false }() |
| 30 | + |
| 31 | + const limit = 1024 |
| 32 | + payload := bytes.Repeat([]byte("A"), 1<<20) |
| 33 | + body := &countingBody{r: bytes.NewReader(payload)} |
| 34 | + |
| 35 | + server := &HTTPServer{options: &Options{MaxDumpBodySize: limit}} |
| 36 | + var readBeforeHandler int64 |
| 37 | + var got []byte |
| 38 | + handler := server.loglayer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { |
| 39 | + readBeforeHandler = body.read |
| 40 | + got, _ = io.ReadAll(r.Body) |
| 41 | + w.WriteHeader(http.StatusUnauthorized) |
| 42 | + })) |
| 43 | + |
| 44 | + req := httptest.NewRequest(http.MethodPost, "http://example.com/", nil) |
| 45 | + req.Body = body |
| 46 | + req.ContentLength = -1 |
| 47 | + req.TransferEncoding = []string{"chunked"} |
| 48 | + handler.ServeHTTP(httptest.NewRecorder(), req) |
| 49 | + |
| 50 | + if readBeforeHandler > limit { |
| 51 | + t.Fatalf("loglayer read %d bytes before the next handler, limit %d", readBeforeHandler, limit) |
| 52 | + } |
| 53 | + if !bytes.Equal(got, payload) { |
| 54 | + t.Fatalf("next handler got %d bytes, want %d", len(got), len(payload)) |
| 55 | + } |
| 56 | +} |
| 57 | + |
| 58 | +// TestLogLayerSkipsBodyWithoutLog checks that the body is not read when no |
| 59 | +// log needs it. |
| 60 | +func TestLogLayerSkipsBodyWithoutLog(t *testing.T) { |
| 61 | + body := &countingBody{r: bytes.NewReader(bytes.Repeat([]byte("A"), 4096))} |
| 62 | + server := &HTTPServer{options: &Options{}} |
| 63 | + var readBeforeHandler int64 = -1 |
| 64 | + handler := server.loglayer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { |
| 65 | + readBeforeHandler = body.read |
| 66 | + })) |
| 67 | + req := httptest.NewRequest(http.MethodPost, "http://example.com/", nil) |
| 68 | + req.Body = body |
| 69 | + handler.ServeHTTP(httptest.NewRecorder(), req) |
| 70 | + if readBeforeHandler != 0 { |
| 71 | + t.Fatalf("loglayer read %d bytes with no log enabled", readBeforeHandler) |
| 72 | + } |
| 73 | +} |
0 commit comments