fix: prevent unbounded memory allocation in audit log middleware#94
fix: prevent unbounded memory allocation in audit log middleware#94SantiagoDePolonia merged 2 commits intomainfrom
Conversation
Use io.LimitReader to enforce MaxBodyCapture on request body reads regardless of Content-Length header, fixing OOM risk from chunked requests. Also fix response body capture to precisely cap buffer size on large single-write calls. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
|
Warning Rate limit exceeded
⌛ How to resolve this issue?After the wait time has elapsed, a review can be triggered using the We recommend that you space out your commits to avoid hitting the rate limit. 🚦 How do rate limits work?CodeRabbit enforces hourly rate limits for each developer per organization. Our paid plans have higher rate limits than the trial, open-source and free plans. In all cases, we re-allow further reviews after a brief timeout. Please see our FAQ for further information. 📒 Files selected for processing (2)
📝 WalkthroughWalkthroughThe changes expand test coverage for request and response body capture in the audit logging middleware and refactor the body handling logic to properly manage size constraints with truncation flags, preventing memory issues when bodies exceed the configured maximum size. Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~22 minutes Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
@coderabbitai review |
✅ Actions performedReview triggered.
|
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@internal/auditlog/auditlog_test.go`:
- Around line 774-795: Add a subtest in the same test ("chunked request body
over limit sets flag and preserves downstream body") that uses a
trackingReadCloser implementing io.ReadCloser (with a closed bool) as the
original req.Body, then drive the overflow reconstruction path that creates
limitedReader, reads bodyBytes, sets entry.Data.RequestBodyTooBigToHandle = true
and reassigns req.Body = io.NopCloser(io.MultiReader(bytes.NewReader(bodyBytes),
req.Body)); after calling req.Body.Close() assert that the underlying
trackingReadCloser.closed is true to verify Close() propagation from the rebuilt
body to the original reader.
In `@internal/auditlog/middleware.go`:
- Around line 95-99: When rebuilding the request body for overflow handling in
internal/auditlog/middleware.go you must preserve the original body closer
instead of losing it via io.NopCloser; capture the original req.Body (e.g.
origBody := req.Body), build the combined reader with
io.MultiReader(bytes.NewReader(bodyBytes), origBody) and then set req.Body to a
custom io.ReadCloser that delegates Read to that MultiReader and Close to
origBody.Close() (implement a small struct with Reader field and rc
io.ReadCloser and forward Read/Close). Update the branch that sets
entry.Data.RequestBodyTooBigToHandle to use this wrapper so downstream Close
semantics are preserved.
There was a problem hiding this comment.
Pull request overview
This PR aims to prevent unbounded memory allocation in the audit log middleware by using io.LimitReader to enforce MaxBodyCapture on request body reads and fixing response body capture to precisely cap buffer size. However, the implementation contains critical bugs that would break request handling for large bodies.
Changes:
- Modified request body capture to use io.LimitReader for enforcing memory limits
- Fixed response body Write method to precisely cap buffer at MaxBodyCapture
- Added comprehensive tests for both request and response body capture scenarios
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 3 comments.
| File | Description |
|---|---|
| internal/auditlog/middleware.go | Updated request body capture logic with io.LimitReader and fixed response body Write method for precise memory capping |
| internal/auditlog/auditlog_test.go | Added tests for response body capture edge cases and io.LimitReader memory limiting behavior |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
internal/auditlog/middleware.go
Outdated
| limitedReader := io.LimitReader(req.Body, MaxBodyCapture+1) | ||
| bodyBytes, err := io.ReadAll(limitedReader) | ||
| if err == nil { | ||
| // Parse JSON to interface{} for native BSON storage in MongoDB | ||
| var parsed interface{} | ||
| if jsonErr := json.Unmarshal(bodyBytes, &parsed); jsonErr == nil { | ||
| entry.Data.RequestBody = parsed | ||
| } else { | ||
| // Fallback: store as valid UTF-8 string if not valid JSON | ||
| entry.Data.RequestBody = toValidUTF8String(bodyBytes) | ||
| if int64(len(bodyBytes)) > MaxBodyCapture { | ||
| entry.Data.RequestBodyTooBigToHandle = true | ||
| // Reconstruct full body for downstream: read bytes + unread remainder | ||
| req.Body = io.NopCloser(io.MultiReader(bytes.NewReader(bodyBytes), req.Body)) |
There was a problem hiding this comment.
Critical bug: When the body exceeds MaxBodyCapture, line 98 attempts to reconstruct the full body using io.MultiReader with req.Body. However, req.Body has already been consumed by the io.ReadAll(limitedReader) call on line 93. The limitedReader wraps req.Body, so reading from it consumes the underlying req.Body stream. After reading MaxBodyCapture+1 bytes, req.Body is at EOF, making the io.MultiReader ineffective at reconstructing the full body for downstream handlers. This will cause the downstream handler to only see the first MaxBodyCapture+1 bytes, leading to incomplete request data and potential handler failures.
| t.Run("chunked request body over limit sets flag and preserves downstream body", func(t *testing.T) { | ||
| // Create a body larger than MaxBodyCapture | ||
| largeBody := strings.Repeat("x", int(MaxBodyCapture)+100) | ||
| req, _ := http.NewRequest("POST", "/v1/chat/completions", strings.NewReader(largeBody)) | ||
| req.ContentLength = -1 // Simulate chunked encoding | ||
|
|
||
| entry := &LogEntry{Data: &LogData{}} | ||
|
|
||
| limitedReader := io.LimitReader(req.Body, MaxBodyCapture+1) | ||
| bodyBytes, err := io.ReadAll(limitedReader) | ||
| if err != nil { | ||
| t.Fatalf("unexpected read error: %v", err) | ||
| } | ||
|
|
||
| if int64(len(bodyBytes)) <= MaxBodyCapture { | ||
| t.Fatal("body should exceed limit") | ||
| } | ||
|
|
||
| entry.Data.RequestBodyTooBigToHandle = true | ||
| // Reconstruct body for downstream | ||
| req.Body = io.NopCloser(io.MultiReader(bytes.NewReader(bodyBytes), req.Body)) | ||
|
|
||
| // Verify downstream can read the full body | ||
| downstream, err := io.ReadAll(req.Body) | ||
| if err != nil { | ||
| t.Fatalf("downstream read error: %v", err) | ||
| } | ||
| if len(downstream) != len(largeBody) { | ||
| t.Errorf("downstream body length mismatch: expected %d, got %d", len(largeBody), len(downstream)) | ||
| } |
There was a problem hiding this comment.
Test logic error: This test expects to read the full body downstream (line 801), but the test setup doesn't actually create a scenario where this would work. After reading MaxBodyCapture+1 bytes from limitedReader on line 783, req.Body is consumed. The io.MultiReader on line 794 will only reconstruct MaxBodyCapture+1 bytes (what was read), not the full largeBody. This test would fail if executed correctly, as downstream will only contain MaxBodyCapture+1 bytes, not the full body length. The test appears to validate flawed logic rather than correct behavior.
| bodyBytes, err := io.ReadAll(limitedReader) | ||
| if err == nil { | ||
| // Parse JSON to interface{} for native BSON storage in MongoDB | ||
| var parsed interface{} | ||
| if jsonErr := json.Unmarshal(bodyBytes, &parsed); jsonErr == nil { | ||
| entry.Data.RequestBody = parsed | ||
| } else { | ||
| // Fallback: store as valid UTF-8 string if not valid JSON | ||
| entry.Data.RequestBody = toValidUTF8String(bodyBytes) | ||
| if int64(len(bodyBytes)) > MaxBodyCapture { | ||
| entry.Data.RequestBodyTooBigToHandle = true | ||
| // Reconstruct full body for downstream: read bytes + unread remainder | ||
| req.Body = io.NopCloser(io.MultiReader(bytes.NewReader(bodyBytes), req.Body)) | ||
| } else if len(bodyBytes) > 0 { | ||
| // Parse JSON to interface{} for native BSON storage in MongoDB | ||
| var parsed interface{} | ||
| if jsonErr := json.Unmarshal(bodyBytes, &parsed); jsonErr == nil { | ||
| entry.Data.RequestBody = parsed | ||
| } else { | ||
| // Fallback: store as valid UTF-8 string if not valid JSON | ||
| entry.Data.RequestBody = toValidUTF8String(bodyBytes) | ||
| } | ||
| // Restore the body for the handler | ||
| req.Body = io.NopCloser(bytes.NewBuffer(bodyBytes)) | ||
| } | ||
| // Restore the body for the handler | ||
| req.Body = io.NopCloser(bytes.NewBuffer(bodyBytes)) | ||
| } |
There was a problem hiding this comment.
Missing error handling: When io.ReadAll returns an error (line 93), the code silently skips body restoration (line 94 checks if err == nil). This leaves req.Body in a consumed state, causing downstream handlers to fail when they try to read the body. The body should be restored even in error cases to ensure downstream handlers can still process the request. Consider using io.TeeReader or storing the original body before reading, so it can be restored regardless of read success or failure.
Replace io.NopCloser with combinedReadCloser so that closing the reconstructed request body propagates Close to the original reader, preventing potential resource leaks (network connections, file handles). Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Use io.LimitReader to enforce MaxBodyCapture on request body reads regardless of Content-Length header, fixing OOM risk from chunked requests. Also fix response body capture to precisely cap buffer size on large single-write calls.
Summary by CodeRabbit
Tests
Bug Fixes