Add transport abstraction for MCP proxy, fix no-params DLP bypass#103
Conversation
|
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. 📝 WalkthroughWalkthroughRefactors MCP stdio I/O to framed MessageReader/MessageWriter abstractions with StdioReader/StdioWriter, updates forwarding and proxy logic to use those interfaces, expands DLP/URL-decoding and scanner checks, and adapts tests (IPv4-only test servers, new test wrappers and additional unit tests). Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant StdioReader as "StdioReader\n(MessageReader)"
participant Proxy as "Proxy\n(ForwardScanned / RunProxy)"
participant StdioWriter as "StdioWriter\n(MessageWriter)"
participant Backend
Client->>StdioReader: send newline-delimited message
StdioReader->>Proxy: ReadMessage() -> []byte
Proxy->>Proxy: scan/process (DLP, injection, approver)
alt clean → forward
Proxy->>StdioWriter: WriteMessage(msg)
StdioWriter->>Backend: single Write(msg + "\n")
else block/strip/warn
Proxy->>Proxy: block/strip/log, optionally WriteMessage(strip/block)
end
Backend-->>StdioWriter: response message
StdioWriter->>Proxy: WriteMessage(response)
Proxy->>StdioReader: (route back) WriteMessage -> Client/stdout
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Suggested labels
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 3 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (3 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 |
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Fix all issues with AI agents
In `@internal/mcp/proxy.go`:
- Around line 37-45: The allocation math (len(msg)+1) in syncWriter.WriteMessage
can trigger a CodeQL overflow warning; replace the manual make+copy with a safe
append-based copy that avoids explicit size arithmetic: create a fresh slice
copy via buf := append(append([]byte(nil), msg...), '\n') (so you don't mutate
the caller's backing array), then write buf as before in
syncWriter.WriteMessage; keep the mutex locking/unlocking and return the write
error as currently implemented.
In `@internal/mcp/transport.go`:
- Around line 51-74: Wrap the scanner and write errors and add a size guard:
when checking sr.scanner.Err() return a wrapped error like fmt.Errorf("scanner
error: %w", err) instead of returning err directly; in StdioWriter.WriteMessage
enforce the same maxLineSize used by StdioReader (compare len(msg) to
maxLineSize and return a wrapped error such as fmt.Errorf("message too large: %d
> %d", len(msg), maxLineSize)) before writing; and wrap the underlying write
error from sw.w.Write with fmt.Errorf("writing message: %w", err). Use the
existing symbols sr.scanner.Err, StdioWriter.WriteMessage, and maxLineSize to
locate and apply these changes.
🧹 Nitpick comments (2)
internal/cli/cli_test.go (1)
146-151: Add a timeout to the DNS probe to avoid test hangs.A blocked resolver can stall indefinitely; a short timeout keeps the skip guard safe.
♻️ Suggested change
- if _, err := net.DefaultResolver.LookupHost(context.Background(), "example.com"); err != nil { + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + if _, err := net.DefaultResolver.LookupHost(ctx, "example.com"); err != nil { t.Skip("DNS unavailable (restricted environment)") }internal/mcp/transport_test.go (1)
11-107: Consider consolidating the StdioReader cases into a table-driven test.Several tests differ only by input/expected output, which fits the project’s test style guideline.
As per coding guidelines: Table-driven tests should be used where applicable in Go test files.
Introduce MessageReader/MessageWriter interfaces to decouple MCP scanning from stdio framing, enabling future HTTP transport support. Refactor ForwardScanned and ForwardScannedInput to use the new interfaces with StdioReader/StdioWriter concrete implementations. Fix a bypass where ScanRequest returned clean when JSON-RPC params was empty/null, allowing secrets in result/error/unknown fields to evade DLP and injection scanning. Now scans the full raw message when params is absent. Harden test infrastructure: force IPv4 loopback listeners in httptest servers to avoid failures in sandboxed environments without IPv6 support.
Address CodeQL overflow warning on len(msg)+1 allocation by adding maxLineSize bounds check to both syncWriter and StdioWriter. Wrap scanner errors with context. Add DNS timeout to test skip guard.
Replace string literal "audit" with ModeAudit constant to satisfy goconst linter (v2.9.0 detects existing constant).
Test the maxLineSize overflow guard on both syncWriter and StdioWriter to cover the error branch added in the previous commit.
luckyPipewrench
force-pushed
the
feat/mcp-http-proxy
branch
from
329ef4f to
8c145d3
Compare
There was a problem hiding this comment.
🧹 Nitpick comments (1)
internal/mcp/transport_test.go (1)
181-196: Consider testing the wrapped error message fromWriteMessage.
TestStdioWriter_TooLargecheckserr.Error()contains"message too large"— good. ButTestStdioWriter_PropagatesError(line 141) only checkserr != nil. Since the past review recommended wrapping write errors intransport.go, you may want to verify the wrapped error message here too, or at minimum useerrors.Is/errors.Asto confirm the underlying error is propagated.
- TestRunProxy_WithToolConfig: exercises tool scanning config init path - TestRunProxy_InputScanningBlocksNotification: verifies notifications with DLP matches are silently dropped (no error response sent) Improves RunProxy coverage from 91.1% to 95.6%.
- Use errors.Is for http.ErrServerClosed comparison (wrapped error safety) - Extract stripOrBlock helper to deduplicate strip-failure fallback logic - Increase maxDecodeRounds from 3 to 10 for multi-layer encoding - Add direct unit tests for stripOrBlock (valid strip + invalid JSON paths)
Red team finding: base64/hex-encoded secrets in individual MCP tool arguments bypassed DLP because ScanTextForDLP was only called on the joined string (not valid base64). Now scans each extracted string individually for encoded patterns. Display fix (review feedback): fetch proxy responses and logs now show the fully-decoded URL via normalizeDisplayURL, so operators see the actual resolved URL instead of partially-decoded intermediates.
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@internal/mcp/proxy.go`:
- Around line 34-48: In syncWriter.WriteMessage, after writing the buffer with
sw.w.Write(buf) wrap and return the underlying error using fmt.Errorf("writing
message: %w", err) instead of returning err directly; keep the existing
maxLineSize check and mutex behavior, and ensure fmt is imported if not already
for error wrapping.
🧹 Nitpick comments (1)
internal/proxy/proxy.go (1)
502-514:normalizeDisplayURLduplicatesiterativeDecodefromscanner.go— use a shared constant instead of hardcoded500.Both functions have identical logic: iteratively
QueryUnescapeuntil the string stabilizes.iterativeDecodeininternal/scanner/scanner.gousesmaxDecodeRounds = 500, butnormalizeDisplayURLhardcodes500. Sinceproxy.goalready imports the scanner package, the cleanest fix is to exportmaxDecodeRounds(or create an exportedIterativeDecodehelper) and reuse it to prevent drift.At minimum, define a named constant in
proxy.go:func normalizeDisplayURL(raw string) string { - for range 500 { + for range maxDisplayDecodeRounds {
Addresses CodeRabbit review: use fmt.Errorf wrapping on write errors to match codebase error-handling convention.
Export scanner.IterativeDecode and remove duplicate normalizeDisplayURL from proxy. Revert unrelated ScanStream comment addition.
3 participants
Summary
Test plan
Summary by CodeRabbit
Bug Fixes
Tests