Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 5 additions & 8 deletions internal/cmd/server/http.go
Original file line number Diff line number Diff line change
Expand Up @@ -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,
),
Expand Down
31 changes: 16 additions & 15 deletions pkg/httpx/handlers.go
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ import (
"context"
"crypto/tls"
"encoding/base64"
"errors"
"io"
"net"
"net/http"
Expand All @@ -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,
Expand Down Expand Up @@ -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)
Expand Down
70 changes: 51 additions & 19 deletions pkg/httpx/httpx_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ import (
"mime/multipart"
"net"
"net/http"
"net/http/httptest"
"net/http/httptrace"
"net/url"
"os"
Expand Down Expand Up @@ -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("<html><body>test</body></html>"))
}),
),
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("<html><body>test</body></html>"))
}),
),
1<<20,
),
Expand Down Expand Up @@ -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) {
Expand Down
Loading