Wraps request body with http.MaxBytesReader to limit incoming payload size. Without this, any endpoint accepting a body is vulnerable to large uploads consuming all available memory. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
62 lines
1.5 KiB
Go
62 lines
1.5 KiB
Go
package server_test
|
|
|
|
import (
|
|
"io"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"testing"
|
|
|
|
"git.codelab.vc/pkg/httpx/server"
|
|
)
|
|
|
|
func TestMaxBodySize(t *testing.T) {
|
|
t.Run("allows body within limit", func(t *testing.T) {
|
|
handler := server.MaxBodySize(1024)(
|
|
http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
b, err := io.ReadAll(r.Body)
|
|
if err != nil {
|
|
http.Error(w, err.Error(), http.StatusBadRequest)
|
|
return
|
|
}
|
|
w.WriteHeader(http.StatusOK)
|
|
w.Write(b)
|
|
}),
|
|
)
|
|
|
|
body := strings.NewReader("hello")
|
|
w := httptest.NewRecorder()
|
|
req := httptest.NewRequest(http.MethodPost, "/", body)
|
|
handler.ServeHTTP(w, req)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("got status %d, want %d", w.Code, http.StatusOK)
|
|
}
|
|
if w.Body.String() != "hello" {
|
|
t.Fatalf("got body %q, want %q", w.Body.String(), "hello")
|
|
}
|
|
})
|
|
|
|
t.Run("rejects body exceeding limit", func(t *testing.T) {
|
|
handler := server.MaxBodySize(5)(
|
|
http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
_, err := io.ReadAll(r.Body)
|
|
if err != nil {
|
|
http.Error(w, "body too large", http.StatusRequestEntityTooLarge)
|
|
return
|
|
}
|
|
w.WriteHeader(http.StatusOK)
|
|
}),
|
|
)
|
|
|
|
body := strings.NewReader("this is longer than 5 bytes")
|
|
w := httptest.NewRecorder()
|
|
req := httptest.NewRequest(http.MethodPost, "/", body)
|
|
handler.ServeHTTP(w, req)
|
|
|
|
if w.Code != http.StatusRequestEntityTooLarge {
|
|
t.Fatalf("got status %d, want %d", w.Code, http.StatusRequestEntityTooLarge)
|
|
}
|
|
})
|
|
}
|