test: 33 regression tests for all security fixes

- handlers_kms_test.go: input validation, orgID from context not header - auth_test.go: JWT iss/aud enforcement, refresh/access token separation - ratelimit_test.go: clientIP uses RemoteAddr, ignores XFF - server_test.go: maxBodySize middleware enforcement
2026-09-24 00:26:24 +00:00 · 2026-04-06 23:35:05 -07:00
parent bb9fc6a76c
commit ffa8a543e2
4 changed files with 723 additions and 0 deletions
@@ -167,3 +167,150 @@ func TestInvalidJWT_EmptyString(t *testing.T) {
 		t.Fatal("validateJWT should reject empty string")
 	}
 }
+
+// === JWT Issuer / Audience regression tests ===
+
+func TestAccessToken_IssuerAndAudience(t *testing.T) {
+	s := newTestServer("iss-aud-secret")
+
+	tokenStr, err := s.generateJWT("u1", "o1", "admin")
+	if err != nil {
+		t.Fatalf("generateJWT: %v", err)
+	}
+
+	// Parse without validation to inspect raw claims.
+	parser := jwt.NewParser(jwt.WithoutClaimsValidation())
+	token, _, err := parser.ParseUnverified(tokenStr, &JWTClaims{})
+	if err != nil {
+		t.Fatalf("ParseUnverified: %v", err)
+	}
+	claims := token.Claims.(*JWTClaims)
+
+	if claims.Issuer != "mpc.lux.network" {
+		t.Errorf("access token Issuer = %q, want %q", claims.Issuer, "mpc.lux.network")
+	}
+	found := false
+	for _, aud := range claims.Audience {
+		if aud == "mpc-api" {
+			found = true
+		}
+	}
+	if !found {
+		t.Errorf("access token Audience = %v, want to contain %q", claims.Audience, "mpc-api")
+	}
+}
+
+func TestRefreshToken_IssuerAndAudience(t *testing.T) {
+	s := newTestServer("iss-aud-secret")
+
+	tokenStr, err := s.generateRefreshToken("u1", "o1", "admin")
+	if err != nil {
+		t.Fatalf("generateRefreshToken: %v", err)
+	}
+
+	parser := jwt.NewParser(jwt.WithoutClaimsValidation())
+	token, _, err := parser.ParseUnverified(tokenStr, &JWTClaims{})
+	if err != nil {
+		t.Fatalf("ParseUnverified: %v", err)
+	}
+	claims := token.Claims.(*JWTClaims)
+
+	if claims.Issuer != "mpc.lux.network" {
+		t.Errorf("refresh token Issuer = %q, want %q", claims.Issuer, "mpc.lux.network")
+	}
+	found := false
+	for _, aud := range claims.Audience {
+		if aud == "mpc-refresh" {
+			found = true
+		}
+	}
+	if !found {
+		t.Errorf("refresh token Audience = %v, want to contain %q", claims.Audience, "mpc-refresh")
+	}
+}
+
+func TestValidateJWT_RejectsWrongIssuer(t *testing.T) {
+	s := newTestServer("iss-test-secret")
+
+	// Create a token with wrong issuer.
+	claims := JWTClaims{
+		UserID: "u1",
+		OrgID:  "o1",
+		Role:   "admin",
+		RegisteredClaims: jwt.RegisteredClaims{
+			Issuer:    "evil.example.com",
+			Audience:  jwt.ClaimStrings{"mpc-api"},
+			ExpiresAt: jwt.NewNumericDate(time.Now().Add(24 * time.Hour)),
+			IssuedAt:  jwt.NewNumericDate(time.Now()),
+		},
+	}
+	token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
+	tokenStr, err := token.SignedString(s.jwtSecret)
+	if err != nil {
+		t.Fatalf("sign token: %v", err)
+	}
+
+	_, err = s.validateJWT(tokenStr)
+	if err == nil {
+		t.Fatal("validateJWT should reject token with wrong issuer")
+	}
+}
+
+func TestValidateJWT_RejectsWrongAudience(t *testing.T) {
+	s := newTestServer("aud-test-secret")
+
+	// Create a token with wrong audience.
+	claims := JWTClaims{
+		UserID: "u1",
+		OrgID:  "o1",
+		Role:   "admin",
+		RegisteredClaims: jwt.RegisteredClaims{
+			Issuer:    "mpc.lux.network",
+			Audience:  jwt.ClaimStrings{"wrong-audience"},
+			ExpiresAt: jwt.NewNumericDate(time.Now().Add(24 * time.Hour)),
+			IssuedAt:  jwt.NewNumericDate(time.Now()),
+		},
+	}
+	token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
+	tokenStr, err := token.SignedString(s.jwtSecret)
+	if err != nil {
+		t.Fatalf("sign token: %v", err)
+	}
+
+	_, err = s.validateJWT(tokenStr)
+	if err == nil {
+		t.Fatal("validateJWT should reject token with wrong audience")
+	}
+}
+
+func TestValidateJWT_RejectsRefreshToken(t *testing.T) {
+	s := newTestServer("cross-token-secret")
+
+	// A refresh token (audience "mpc-refresh") must NOT pass validateJWT
+	// (which expects audience "mpc-api").
+	refreshStr, err := s.generateRefreshToken("u1", "o1", "admin")
+	if err != nil {
+		t.Fatalf("generateRefreshToken: %v", err)
+	}
+
+	_, err = s.validateJWT(refreshStr)
+	if err == nil {
+		t.Fatal("validateJWT should reject a refresh token (audience mpc-refresh)")
+	}
+}
+
+func TestValidateRefreshToken_RejectsAccessToken(t *testing.T) {
+	s := newTestServer("cross-token-secret")
+
+	// An access token (audience "mpc-api") must NOT pass validateRefreshToken
+	// (which expects audience "mpc-refresh").
+	accessStr, err := s.generateJWT("u1", "o1", "admin")
+	if err != nil {
+		t.Fatalf("generateJWT: %v", err)
+	}
+
+	_, err = s.validateRefreshToken(accessStr)
+	if err == nil {
+		t.Fatal("validateRefreshToken should reject an access token (audience mpc-api)")
+	}
+}
@@ -0,0 +1,307 @@
+package api
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+)
+
+// mockMPCBackend implements MPCBackend for KMS handler tests.
+// Only used where validation passes and the handler calls MPC.
+type mockMPCBackend struct {
+	keygenResult *KeygenResult
+	signResult   *SignResult
+	keygenErr    error
+	signErr      error
+	reshareErr   error
+}
+
+func (m *mockMPCBackend) TriggerKeygen(orgID, walletID string) (*KeygenResult, error) {
+	if m.keygenErr != nil {
+		return nil, m.keygenErr
+	}
+	return m.keygenResult, nil
+}
+func (m *mockMPCBackend) TriggerSign(orgID, walletID string, payload []byte) (*SignResult, error) {
+	if m.signErr != nil {
+		return nil, m.signErr
+	}
+	return m.signResult, nil
+}
+func (m *mockMPCBackend) TriggerReshare(orgID, walletID string, newThreshold int, newParticipants []string) error {
+	return m.reshareErr
+}
+func (m *mockMPCBackend) ExportKeyShare(orgID, walletID string) ([]byte, error) {
+	return nil, nil
+}
+func (m *mockMPCBackend) GetClusterStatus() *ClusterStatus {
+	return &ClusterStatus{Ready: true}
+}
+
+func kmsTestServer() *Server {
+	return &Server{
+		jwtSecret: []byte("kms-test-secret"),
+		mpc:       &mockMPCBackend{},
+	}
+}
+
+func kmsRequest(t *testing.T, method, path string, body interface{}) *http.Request {
+	t.Helper()
+	var buf bytes.Buffer
+	if body != nil {
+		if err := json.NewEncoder(&buf).Encode(body); err != nil {
+			t.Fatalf("encode request body: %v", err)
+		}
+	}
+	req := httptest.NewRequest(method, path, &buf)
+	req.Header.Set("Content-Type", "application/json")
+	// Inject orgID into context (simulating authMiddleware).
+	ctx := context.WithValue(req.Context(), ctxOrgID, "test-org")
+	ctx = context.WithValue(ctx, ctxUserID, "test-user")
+	ctx = context.WithValue(ctx, ctxRole, "admin")
+	return req.WithContext(ctx)
+}
+
+func decodeKMSError(t *testing.T, rec *httptest.ResponseRecorder) string {
+	t.Helper()
+	var body map[string]string
+	if err := json.NewDecoder(rec.Body).Decode(&body); err != nil {
+		t.Fatalf("decode error response: %v", err)
+	}
+	return body["error"]
+}
+
+// --- handleKMSGenerate validation ---
+
+func TestKMSGenerate_EmptyValidatorID(t *testing.T) {
+	s := kmsTestServer()
+	req := kmsRequest(t, http.MethodPost, "/api/v1/keys/generate", kmsGenerateRequest{
+		ValidatorID: "",
+		Threshold:   3,
+		Parties:     5,
+	})
+	rec := httptest.NewRecorder()
+
+	s.handleKMSGenerate(rec, req)
+
+	if rec.Code != http.StatusBadRequest {
+		t.Fatalf("status = %d, want %d", rec.Code, http.StatusBadRequest)
+	}
+	if msg := decodeKMSError(t, rec); msg != "validator_id is required" {
+		t.Fatalf("error = %q, want %q", msg, "validator_id is required")
+	}
+}
+
+func TestKMSGenerate_ThresholdLessThanTwo(t *testing.T) {
+	s := kmsTestServer()
+	req := kmsRequest(t, http.MethodPost, "/api/v1/keys/generate", kmsGenerateRequest{
+		ValidatorID: "val-1",
+		Threshold:   1,
+		Parties:     5,
+	})
+	rec := httptest.NewRecorder()
+
+	s.handleKMSGenerate(rec, req)
+
+	if rec.Code != http.StatusBadRequest {
+		t.Fatalf("status = %d, want %d", rec.Code, http.StatusBadRequest)
+	}
+	if msg := decodeKMSError(t, rec); msg != "threshold must be >= 2" {
+		t.Fatalf("error = %q, want %q", msg, "threshold must be >= 2")
+	}
+}
+
+func TestKMSGenerate_PartiesLessThanThreshold(t *testing.T) {
+	s := kmsTestServer()
+	req := kmsRequest(t, http.MethodPost, "/api/v1/keys/generate", kmsGenerateRequest{
+		ValidatorID: "val-1",
+		Threshold:   4,
+		Parties:     3,
+	})
+	rec := httptest.NewRecorder()
+
+	s.handleKMSGenerate(rec, req)
+
+	if rec.Code != http.StatusBadRequest {
+		t.Fatalf("status = %d, want %d", rec.Code, http.StatusBadRequest)
+	}
+	if msg := decodeKMSError(t, rec); msg != "parties must be >= threshold" {
+		t.Fatalf("error = %q, want %q", msg, "parties must be >= threshold")
+	}
+}
+
+func TestKMSGenerate_ThresholdZero(t *testing.T) {
+	s := kmsTestServer()
+	req := kmsRequest(t, http.MethodPost, "/api/v1/keys/generate", kmsGenerateRequest{
+		ValidatorID: "val-1",
+		Threshold:   0,
+		Parties:     5,
+	})
+	rec := httptest.NewRecorder()
+
+	s.handleKMSGenerate(rec, req)
+
+	if rec.Code != http.StatusBadRequest {
+		t.Fatalf("status = %d, want %d", rec.Code, http.StatusBadRequest)
+	}
+	if msg := decodeKMSError(t, rec); msg != "threshold must be >= 2" {
+		t.Fatalf("error = %q, want %q", msg, "threshold must be >= 2")
+	}
+}
+
+// --- handleKMSSign validation ---
+
+func TestKMSSign_EmptyMessage(t *testing.T) {
+	s := kmsTestServer()
+	req := kmsRequest(t, http.MethodPost, "/api/v1/keys/val-1/sign", kmsSignRequest{
+		KeyType: "bls",
+		Message: nil,
+	})
+	rec := httptest.NewRecorder()
+
+	s.handleKMSSign(rec, req)
+
+	if rec.Code != http.StatusBadRequest {
+		t.Fatalf("status = %d, want %d", rec.Code, http.StatusBadRequest)
+	}
+	if msg := decodeKMSError(t, rec); msg != "message is required" {
+		t.Fatalf("error = %q, want %q", msg, "message is required")
+	}
+}
+
+func TestKMSSign_EmptyMessageBytes(t *testing.T) {
+	s := kmsTestServer()
+	req := kmsRequest(t, http.MethodPost, "/api/v1/keys/val-1/sign", kmsSignRequest{
+		KeyType: "bls",
+		Message: []byte{},
+	})
+	rec := httptest.NewRecorder()
+
+	s.handleKMSSign(rec, req)
+
+	if rec.Code != http.StatusBadRequest {
+		t.Fatalf("status = %d, want %d", rec.Code, http.StatusBadRequest)
+	}
+	if msg := decodeKMSError(t, rec); msg != "message is required" {
+		t.Fatalf("error = %q, want %q", msg, "message is required")
+	}
+}
+
+// --- handleKMSRotate validation ---
+
+func TestKMSRotate_EmptyRequest(t *testing.T) {
+	s := kmsTestServer()
+	req := kmsRequest(t, http.MethodPost, "/api/v1/keys/val-1/rotate", kmsRotateRequest{})
+	rec := httptest.NewRecorder()
+
+	s.handleKMSRotate(rec, req)
+
+	if rec.Code != http.StatusBadRequest {
+		t.Fatalf("status = %d, want %d", rec.Code, http.StatusBadRequest)
+	}
+	if msg := decodeKMSError(t, rec); msg != "new_threshold or new_participants required" {
+		t.Fatalf("error = %q, want %q", msg, "new_threshold or new_participants required")
+	}
+}
+
+func TestKMSRotate_ZeroThresholdNoParticipants(t *testing.T) {
+	s := kmsTestServer()
+	req := kmsRequest(t, http.MethodPost, "/api/v1/keys/val-1/rotate", kmsRotateRequest{
+		NewThreshold:    0,
+		NewParticipants: nil,
+	})
+	rec := httptest.NewRecorder()
+
+	s.handleKMSRotate(rec, req)
+
+	if rec.Code != http.StatusBadRequest {
+		t.Fatalf("status = %d, want %d", rec.Code, http.StatusBadRequest)
+	}
+}
+
+// --- Verify handlers use getOrgID (context-based) not X-Org-ID header ---
+
+func TestKMSGenerate_UsesContextOrgID(t *testing.T) {
+	// Regression: orgID must come from JWT context (getOrgID), not from
+	// a client-supplied header. If someone sets X-Org-ID header to a
+	// different org, the handler must ignore it and use the context value.
+	s := kmsTestServer()
+	s.mpc = &mockMPCBackend{
+		keygenResult: &KeygenResult{
+			WalletID:    "w-1",
+			ECDSAPubKey: "pub-ecdsa",
+			EDDSAPubKey: "pub-eddsa",
+		},
+	}
+
+	req := kmsRequest(t, http.MethodPost, "/api/v1/keys/generate", kmsGenerateRequest{
+		ValidatorID: "val-ctx-test",
+		Threshold:   2,
+		Parties:     3,
+	})
+	// Attacker tries to override org via header.
+	req.Header.Set("X-Org-ID", "attacker-org")
+	rec := httptest.NewRecorder()
+
+	s.handleKMSGenerate(rec, req)
+
+	// Handler should use context org ("test-org" from kmsRequest), not "attacker-org".
+	// If s.db is nil the handler skips the duplicate check and proceeds to keygen.
+	// The mock MPC backend succeeds. We verify the handler did not error on org mismatch.
+	if rec.Code != http.StatusCreated {
+		t.Fatalf("status = %d, want %d (body: %s)", rec.Code, http.StatusCreated, rec.Body.String())
+	}
+}
+
+// --- handleKMSGenerate invalid JSON body ---
+
+func TestKMSGenerate_InvalidJSON(t *testing.T) {
+	s := kmsTestServer()
+	req := httptest.NewRequest(http.MethodPost, "/api/v1/keys/generate", bytes.NewReader([]byte("not json")))
+	req.Header.Set("Content-Type", "application/json")
+	ctx := context.WithValue(req.Context(), ctxOrgID, "test-org")
+	req = req.WithContext(ctx)
+	rec := httptest.NewRecorder()
+
+	s.handleKMSGenerate(rec, req)
+
+	if rec.Code != http.StatusBadRequest {
+		t.Fatalf("status = %d, want %d", rec.Code, http.StatusBadRequest)
+	}
+	if msg := decodeKMSError(t, rec); msg != "invalid request body" {
+		t.Fatalf("error = %q, want %q", msg, "invalid request body")
+	}
+}
+
+func TestKMSSign_InvalidJSON(t *testing.T) {
+	s := kmsTestServer()
+	req := httptest.NewRequest(http.MethodPost, "/api/v1/keys/val-1/sign", bytes.NewReader([]byte("{bad")))
+	req.Header.Set("Content-Type", "application/json")
+	ctx := context.WithValue(req.Context(), ctxOrgID, "test-org")
+	req = req.WithContext(ctx)
+	rec := httptest.NewRecorder()
+
+	s.handleKMSSign(rec, req)
+
+	if rec.Code != http.StatusBadRequest {
+		t.Fatalf("status = %d, want %d", rec.Code, http.StatusBadRequest)
+	}
+}
+
+func TestKMSRotate_InvalidJSON(t *testing.T) {
+	s := kmsTestServer()
+	req := httptest.NewRequest(http.MethodPost, "/api/v1/keys/val-1/rotate", bytes.NewReader([]byte("[")))
+	req.Header.Set("Content-Type", "application/json")
+	ctx := context.WithValue(req.Context(), ctxOrgID, "test-org")
+	req = req.WithContext(ctx)
+	rec := httptest.NewRecorder()
+
+	s.handleKMSRotate(rec, req)
+
+	if rec.Code != http.StatusBadRequest {
+		t.Fatalf("status = %d, want %d", rec.Code, http.StatusBadRequest)
+	}
+}
@@ -0,0 +1,121 @@
+package api
+
+import (
+	"net/http/httptest"
+	"testing"
+)
+
+func TestClientIP_IPv4WithPort(t *testing.T) {
+	req := httptest.NewRequest("GET", "/test", nil)
+	req.RemoteAddr = "192.168.1.100:12345"
+
+	got := clientIP(req)
+	if got != "192.168.1.100" {
+		t.Errorf("clientIP = %q, want %q", got, "192.168.1.100")
+	}
+}
+
+func TestClientIP_IPv6WithPort(t *testing.T) {
+	req := httptest.NewRequest("GET", "/test", nil)
+	req.RemoteAddr = "[::1]:8080"
+
+	got := clientIP(req)
+	if got != "::1" {
+		t.Errorf("clientIP = %q, want %q", got, "::1")
+	}
+}
+
+func TestClientIP_IPv6Full(t *testing.T) {
+	req := httptest.NewRequest("GET", "/test", nil)
+	req.RemoteAddr = "[2001:db8::1]:443"
+
+	got := clientIP(req)
+	if got != "2001:db8::1" {
+		t.Errorf("clientIP = %q, want %q", got, "2001:db8::1")
+	}
+}
+
+func TestClientIP_NoPort(t *testing.T) {
+	// net.SplitHostPort fails when there's no port; clientIP falls back to
+	// returning RemoteAddr as-is.
+	req := httptest.NewRequest("GET", "/test", nil)
+	req.RemoteAddr = "10.0.0.1"
+
+	got := clientIP(req)
+	if got != "10.0.0.1" {
+		t.Errorf("clientIP = %q, want %q", got, "10.0.0.1")
+	}
+}
+
+func TestClientIP_IgnoresXForwardedFor(t *testing.T) {
+	// clientIP must use RemoteAddr, NOT X-Forwarded-For. The chi RealIP
+	// middleware already copies X-Forwarded-For into RemoteAddr before the
+	// rate limiter runs. clientIP itself must not re-parse headers — doing
+	// so would let attackers spoof their IP and bypass rate limiting.
+	req := httptest.NewRequest("GET", "/test", nil)
+	req.RemoteAddr = "10.0.0.1:9999"
+	req.Header.Set("X-Forwarded-For", "8.8.8.8, 1.1.1.1")
+
+	got := clientIP(req)
+	if got != "10.0.0.1" {
+		t.Errorf("clientIP = %q, want %q (should ignore X-Forwarded-For)", got, "10.0.0.1")
+	}
+}
+
+func TestClientIP_IgnoresXRealIP(t *testing.T) {
+	req := httptest.NewRequest("GET", "/test", nil)
+	req.RemoteAddr = "172.16.0.5:1234"
+	req.Header.Set("X-Real-IP", "8.8.8.8")
+
+	got := clientIP(req)
+	if got != "172.16.0.5" {
+		t.Errorf("clientIP = %q, want %q (should ignore X-Real-IP)", got, "172.16.0.5")
+	}
+}
+
+func TestRateLimiter_AllowsUnderLimit(t *testing.T) {
+	rl := &RateLimiter{
+		limit:      5,
+		windowSecs: 60,
+	}
+
+	for i := 0; i < 5; i++ {
+		if !rl.allow("192.168.1.1") {
+			t.Fatalf("request %d should be allowed", i+1)
+		}
+	}
+}
+
+func TestRateLimiter_BlocksOverLimit(t *testing.T) {
+	rl := &RateLimiter{
+		limit:      3,
+		windowSecs: 60,
+	}
+
+	for i := 0; i < 3; i++ {
+		rl.allow("192.168.1.1")
+	}
+
+	if rl.allow("192.168.1.1") {
+		t.Fatal("4th request should be rate limited")
+	}
+}
+
+func TestRateLimiter_DifferentIPsIndependent(t *testing.T) {
+	rl := &RateLimiter{
+		limit:      2,
+		windowSecs: 60,
+	}
+
+	// Exhaust limit for IP A.
+	rl.allow("10.0.0.1")
+	rl.allow("10.0.0.1")
+	if rl.allow("10.0.0.1") {
+		t.Fatal("IP A 3rd request should be limited")
+	}
+
+	// IP B should still be allowed.
+	if !rl.allow("10.0.0.2") {
+		t.Fatal("IP B should not be affected by IP A's rate limit")
+	}
+}
@@ -0,0 +1,148 @@
+package api
+
+import (
+	"bytes"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+)
+
+func TestMaxBodySize_SmallBody(t *testing.T) {
+	// A body under the limit should pass through without error.
+	limit := int64(1 << 20) // 1 MB
+	handler := maxBodySize(limit)(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		body, err := io.ReadAll(r.Body)
+		if err != nil {
+			writeError(w, http.StatusRequestEntityTooLarge, "body too large")
+			return
+		}
+		w.WriteHeader(http.StatusOK)
+		w.Write(body)
+	}))
+
+	payload := bytes.Repeat([]byte("a"), 1024) // 1 KB
+	req := httptest.NewRequest(http.MethodPost, "/test", bytes.NewReader(payload))
+	rec := httptest.NewRecorder()
+
+	handler.ServeHTTP(rec, req)
+
+	if rec.Code != http.StatusOK {
+		t.Fatalf("status = %d, want %d", rec.Code, http.StatusOK)
+	}
+	if rec.Body.Len() != 1024 {
+		t.Fatalf("body length = %d, want %d", rec.Body.Len(), 1024)
+	}
+}
+
+func TestMaxBodySize_ExactlyAtLimit(t *testing.T) {
+	limit := int64(1024)
+	handler := maxBodySize(limit)(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		_, err := io.ReadAll(r.Body)
+		if err != nil {
+			writeError(w, http.StatusRequestEntityTooLarge, "body too large")
+			return
+		}
+		w.WriteHeader(http.StatusOK)
+	}))
+
+	payload := bytes.Repeat([]byte("x"), 1024)
+	req := httptest.NewRequest(http.MethodPost, "/test", bytes.NewReader(payload))
+	rec := httptest.NewRecorder()
+
+	handler.ServeHTTP(rec, req)
+
+	if rec.Code != http.StatusOK {
+		t.Fatalf("status = %d, want %d", rec.Code, http.StatusOK)
+	}
+}
+
+func TestMaxBodySize_OverLimit(t *testing.T) {
+	limit := int64(1024)
+	handler := maxBodySize(limit)(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		_, err := io.ReadAll(r.Body)
+		if err != nil {
+			// http.MaxBytesReader returns an error when the body exceeds the limit.
+			writeError(w, http.StatusRequestEntityTooLarge, "body too large")
+			return
+		}
+		// If we reach here, the middleware failed to enforce the limit.
+		w.WriteHeader(http.StatusOK)
+	}))
+
+	// 2 KB body against 1 KB limit.
+	payload := bytes.Repeat([]byte("x"), 2048)
+	req := httptest.NewRequest(http.MethodPost, "/test", bytes.NewReader(payload))
+	rec := httptest.NewRecorder()
+
+	handler.ServeHTTP(rec, req)
+
+	if rec.Code != http.StatusRequestEntityTooLarge {
+		t.Fatalf("status = %d, want %d", rec.Code, http.StatusRequestEntityTooLarge)
+	}
+}
+
+func TestMaxBodySize_1MB_Production(t *testing.T) {
+	// Verify the production limit (1 MB = 1 << 20) rejects 2 MB.
+	limit := int64(1 << 20)
+	handler := maxBodySize(limit)(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		_, err := io.ReadAll(r.Body)
+		if err != nil {
+			writeError(w, http.StatusRequestEntityTooLarge, "body too large")
+			return
+		}
+		w.WriteHeader(http.StatusOK)
+	}))
+
+	// 2 MB payload.
+	payload := bytes.Repeat([]byte("A"), 2<<20)
+	req := httptest.NewRequest(http.MethodPost, "/test", bytes.NewReader(payload))
+	rec := httptest.NewRecorder()
+
+	handler.ServeHTTP(rec, req)
+
+	if rec.Code != http.StatusRequestEntityTooLarge {
+		t.Fatalf("status = %d, want %d", rec.Code, http.StatusRequestEntityTooLarge)
+	}
+}
+
+func TestMaxBodySize_EmptyBody(t *testing.T) {
+	limit := int64(1024)
+	handler := maxBodySize(limit)(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		body, err := io.ReadAll(r.Body)
+		if err != nil {
+			writeError(w, http.StatusRequestEntityTooLarge, "body too large")
+			return
+		}
+		if len(body) != 0 {
+			t.Errorf("body length = %d, want 0", len(body))
+		}
+		w.WriteHeader(http.StatusOK)
+	}))
+
+	req := httptest.NewRequest(http.MethodPost, "/test", nil)
+	rec := httptest.NewRecorder()
+
+	handler.ServeHTTP(rec, req)
+
+	if rec.Code != http.StatusOK {
+		t.Fatalf("status = %d, want %d", rec.Code, http.StatusOK)
+	}
+}
+
+func TestMaxBodySize_GetRequest(t *testing.T) {
+	// GET requests have no body; middleware should not interfere.
+	limit := int64(1024)
+	handler := maxBodySize(limit)(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusOK)
+	}))
+
+	req := httptest.NewRequest(http.MethodGet, "/test", nil)
+	rec := httptest.NewRecorder()
+
+	handler.ServeHTTP(rec, req)
+
+	if rec.Code != http.StatusOK {
+		t.Fatalf("status = %d, want %d", rec.Code, http.StatusOK)
+	}
+}