diff --git a/internal/cmd/server/http.go b/internal/cmd/server/http.go index 41389e9..fef05fd 100644 --- a/internal/cmd/server/http.go +++ b/internal/cmd/server/http.go @@ -95,15 +95,12 @@ func HTTPHandler( return HTTPTelemetry( http.TimeoutHandler( httpx.BodyReaderHandler( - httpx.MaxBytesHandler( - httpx.NotifyHandler( - notify, - httpdb.Handler( - &httpdb.Routes{DB: db, Origin: origin}, - http.HandlerFunc(HTTPDefault), - ), + httpx.NotifyHandler( + notify, + httpdb.Handler( + &httpdb.Routes{DB: db, Origin: origin}, + http.HandlerFunc(HTTPDefault), ), - httpMaxBodyBytes, ), httpMaxBodyBytes, ), diff --git a/pkg/httpx/handlers.go b/pkg/httpx/handlers.go index 8370f35..dda8efb 100644 --- a/pkg/httpx/handlers.go +++ b/pkg/httpx/handlers.go @@ -5,6 +5,7 @@ import ( "context" "crypto/tls" "encoding/base64" + "errors" "io" "net" "net/http" @@ -15,20 +16,6 @@ import ( "github.com/nt0xa/sonar/pkg/netx" ) -func MaxBytesHandler(h http.Handler, n int64) http.Handler { - return &maxBytesHandler{h, n} -} - -type maxBytesHandler struct { - h http.Handler - n int64 -} - -func (h *maxBytesHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) { - r.Body = http.MaxBytesReader(w, r.Body, h.n) - h.h.ServeHTTP(w, r) -} - type NotifyFunc func( ctx context.Context, remoteAddr net.Addr, @@ -103,7 +90,21 @@ func NotifyHandler(notify NotifyFunc, next http.Handler) http.Handler { // BodyReaderHandler reads body so it will appear in request log. func BodyReaderHandler(next http.Handler, maxMemory int64) http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - body, _ := io.ReadAll(r.Body) + if maxMemory > 0 { + r.Body = http.MaxBytesReader(w, r.Body, maxMemory) + } + + body, err := io.ReadAll(r.Body) + if err != nil { + if _, ok := errors.AsType[*http.MaxBytesError](err); ok { + http.Error(w, http.StatusText(http.StatusRequestEntityTooLarge), http.StatusRequestEntityTooLarge) + return + } + + http.Error(w, http.StatusText(http.StatusBadRequest), http.StatusBadRequest) + return + } + r.Body = io.NopCloser(bytes.NewBuffer(body)) next.ServeHTTP(w, r) diff --git a/pkg/httpx/httpx_test.go b/pkg/httpx/httpx_test.go index 9eeebde..a479984 100644 --- a/pkg/httpx/httpx_test.go +++ b/pkg/httpx/httpx_test.go @@ -10,6 +10,7 @@ import ( "mime/multipart" "net" "net/http" + "net/http/httptest" "net/http/httptrace" "net/url" "os" @@ -57,25 +58,22 @@ func TestMain(m *testing.M) { h := http.TimeoutHandler( httpx.BodyReaderHandler( - httpx.MaxBytesHandler( - httpx.NotifyHandler( - func( - ctx context.Context, - remoteAddr net.Addr, - receivedAt *time.Time, - secure bool, - read, written, combined []byte, - meta *httpx.Meta, - ) { - notifier.Notify(remoteAddr, combined, secure) - }, - http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - w.Header().Set("Content-Type", "text/html; charset=utf-8") - w.WriteHeader(200) - _, _ = w.Write([]byte("test")) - }), - ), - 1<<20, + httpx.NotifyHandler( + func( + ctx context.Context, + remoteAddr net.Addr, + receivedAt *time.Time, + secure bool, + read, written, combined []byte, + meta *httpx.Meta, + ) { + notifier.Notify(remoteAddr, combined, secure) + }, + http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/html; charset=utf-8") + w.WriteHeader(200) + _, _ = w.Write([]byte("test")) + }), ), 1<<20, ), @@ -402,6 +400,40 @@ func TestHTTPX(t *testing.T) { notifier.AssertExpectations(t) } +func TestBodyReaderHandlerLimitsBodyBeforeBuffering(t *testing.T) { + called := false + h := httpx.BodyReaderHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + called = true + w.WriteHeader(http.StatusOK) + }), 4) + + req, err := http.NewRequest(http.MethodPost, "/", strings.NewReader("12345")) + require.NoError(t, err) + + rr := httptest.NewRecorder() + h.ServeHTTP(rr, req) + + assert.False(t, called) + assert.Equal(t, http.StatusRequestEntityTooLarge, rr.Code) +} + +func TestBodyReaderHandlerReplaysLimitedBody(t *testing.T) { + h := httpx.BodyReaderHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, err := io.ReadAll(r.Body) + require.NoError(t, err) + assert.Equal(t, "1234", string(body)) + w.WriteHeader(http.StatusNoContent) + }), 4) + + req, err := http.NewRequest(http.MethodPost, "/", strings.NewReader("1234")) + require.NoError(t, err) + + rr := httptest.NewRecorder() + h.ServeHTTP(rr, req) + + assert.Equal(t, http.StatusNoContent, rr.Code) +} + func TestKeepAlive(t *testing.T) { clientTrace := &httptrace.ClientTrace{ GotConn: func(info httptrace.GotConnInfo) {