66package transparent
77
88import (
9+ "bytes"
10+ "encoding/json"
11+ "fmt"
12+ "io"
13+ "math"
14+ "mime"
915 "net/http"
16+ "strings"
1017
1118 "github.com/stacklok/toolhive/pkg/transport/types"
1219)
1320
21+ // maxJSONRPCResponseBytes caps how much of an upstream JSON-RPC response the proxy
22+ // will buffer for structural validation. Matches existing streamable-HTTP body
23+ // limits elsewhere in the codebase (pkg/vmcp/client, pkg/vmcp/session/internal/backend).
24+ const maxJSONRPCResponseBytes = 100 << 20 // 100 MiB
25+
26+ // JSON-RPC error code returned to clients when the proxy rejects a malformed
27+ // upstream response. -32000..-32099 is the implementation-defined server-error
28+ // range in the JSON-RPC 2.0 spec; -32603 is reserved for internal JSON-RPC
29+ // implementation errors and is not appropriate for a policy-level rejection.
30+ const jsonRPCInvalidUpstreamCode = - 32000
31+
1432// ResponseProcessor defines the interface for processing and modifying HTTP responses
1533// based on transport-specific requirements.
1634type ResponseProcessor interface {
@@ -22,12 +40,38 @@ type ResponseProcessor interface {
2240 ShouldProcess (resp * http.Response ) bool
2341}
2442
25- // NoOpResponseProcessor is a processor that does nothing .
26- // Used for transports that don't require response processing (e.g., streamable-http) .
43+ // NoOpResponseProcessor is the default processor for non-SSE transports .
44+ // It validates JSON-RPC responses for streamable HTTP and otherwise leaves responses unchanged .
2745type NoOpResponseProcessor struct {}
2846
29- // ProcessResponse is a no-op implementation.
30- func (* NoOpResponseProcessor ) ProcessResponse (_ * http.Response ) error {
47+ // ProcessResponse validates JSON-RPC responses when applicable.
48+ func (* NoOpResponseProcessor ) ProcessResponse (resp * http.Response ) error {
49+ if ! shouldValidateJSONRPCResponse (resp ) {
50+ return nil
51+ }
52+
53+ // Read one byte past the cap so we can detect oversize without allocating beyond it.
54+ body , err := io .ReadAll (io .LimitReader (resp .Body , maxJSONRPCResponseBytes + 1 ))
55+ if err != nil {
56+ return fmt .Errorf ("failed to read upstream response body: %w" , err )
57+ }
58+ _ = resp .Body .Close ()
59+
60+ if len (body ) > maxJSONRPCResponseBytes {
61+ writeInvalidUpstreamJSONRPCResponse (resp , fmt .Errorf (
62+ "upstream JSON-RPC response exceeds maximum allowed size of %d bytes" , maxJSONRPCResponseBytes ))
63+ return nil
64+ }
65+
66+ if err := validateJSONRPCResponse (body ); err != nil {
67+ writeInvalidUpstreamJSONRPCResponse (resp , err )
68+ return nil
69+ }
70+
71+ // The reverse proxy still needs a readable body after validation.
72+ resp .Body = io .NopCloser (bytes .NewReader (body ))
73+ resp .ContentLength = int64 (len (body ))
74+ resp .Header .Set ("Content-Length" , fmt .Sprintf ("%d" , len (body )))
3175 return nil
3276}
3377
@@ -36,6 +80,154 @@ func (*NoOpResponseProcessor) ShouldProcess(_ *http.Response) bool {
3680 return false
3781}
3882
83+ func shouldValidateJSONRPCResponse (resp * http.Response ) bool {
84+ if resp == nil || resp .Body == nil || resp .Request == nil {
85+ return false
86+ }
87+ if resp .Request .Method != http .MethodPost || resp .StatusCode != http .StatusOK {
88+ return false
89+ }
90+ if ! hasIdentityContentEncoding (resp .Header .Get ("Content-Encoding" )) {
91+ // Content-Encoding semantics (RFC 9110): media-type rules apply after decoding.
92+ // Validating a still-encoded body would mis-classify legitimate gzip JSON-RPC
93+ // frames as invalid. Skip rather than introduce decompression here.
94+ return false
95+ }
96+ if ! requestLooksLikeMCP (resp .Request ) {
97+ // Narrow validation to traffic that carries an MCP streamable-HTTP signal,
98+ // so non-MCP application/json POSTs flowing through the catch-all are not
99+ // rewritten. Backward-compat clients omitting MCP-Protocol-Version on the
100+ // initial initialize will pass through unchanged.
101+ return false
102+ }
103+ contentType := strings .ToLower (resp .Header .Get ("Content-Type" ))
104+ mediaType , _ , err := mime .ParseMediaType (contentType )
105+ if err != nil {
106+ return false
107+ }
108+ return mediaType == "application/json" || mediaType == "application/json-rpc"
109+ }
110+
111+ func hasIdentityContentEncoding (value string ) bool {
112+ v := strings .TrimSpace (strings .ToLower (value ))
113+ return v == "" || v == "identity"
114+ }
115+
116+ func requestLooksLikeMCP (req * http.Request ) bool {
117+ if req == nil {
118+ return false
119+ }
120+ return req .Header .Get ("MCP-Protocol-Version" ) != "" || req .Header .Get ("Mcp-Session-Id" ) != ""
121+ }
122+
123+ func validateJSONRPCResponse (body []byte ) error {
124+ var payload any
125+ dec := json .NewDecoder (bytes .NewReader (body ))
126+ if err := dec .Decode (& payload ); err != nil {
127+ return fmt .Errorf ("invalid JSON body: %w" , err )
128+ }
129+ if dec .More () {
130+ return fmt .Errorf ("JSON-RPC response must contain a single JSON value" )
131+ }
132+ if err := dec .Decode (& struct {}{}); err != io .EOF {
133+ return fmt .Errorf ("JSON-RPC response must contain a single JSON value" )
134+ }
135+
136+ switch value := payload .(type ) {
137+ case map [string ]any :
138+ return validateJSONRPCResponseObject (value )
139+ case []any :
140+ if len (value ) == 0 {
141+ return fmt .Errorf ("JSON-RPC batch response must not be empty" )
142+ }
143+ for i , item := range value {
144+ obj , ok := item .(map [string ]any )
145+ if ! ok {
146+ return fmt .Errorf ("JSON-RPC batch item %d must be an object" , i )
147+ }
148+ if err := validateJSONRPCResponseObject (obj ); err != nil {
149+ return fmt .Errorf ("JSON-RPC batch item %d is invalid: %w" , i , err )
150+ }
151+ }
152+ return nil
153+ default :
154+ return fmt .Errorf ("JSON-RPC response must be an object or array" )
155+ }
156+ }
157+
158+ func validateJSONRPCResponseObject (obj map [string ]any ) error {
159+ if obj ["jsonrpc" ] != "2.0" {
160+ return fmt .Errorf (`JSON-RPC response must include "jsonrpc":"2.0"` )
161+ }
162+
163+ if _ , ok := obj ["id" ]; ! ok {
164+ return fmt .Errorf ("JSON-RPC response must include id" )
165+ }
166+ if ! isValidJSONRPCID (obj ["id" ]) {
167+ return fmt .Errorf ("JSON-RPC response id must be string, number, or null" )
168+ }
169+
170+ _ , hasResult := obj ["result" ]
171+ _ , hasError := obj ["error" ]
172+ if hasResult == hasError {
173+ return fmt .Errorf ("JSON-RPC response must include exactly one of result or error" )
174+ }
175+ if hasError {
176+ if errObj , ok := obj ["error" ].(map [string ]any ); ! ok || ! isValidJSONRPCError (errObj ) {
177+ return fmt .Errorf ("JSON-RPC error response must include error.code and error.message" )
178+ }
179+ }
180+
181+ return nil
182+ }
183+
184+ func isValidJSONRPCID (id any ) bool {
185+ switch id .(type ) {
186+ case nil , string , float64 :
187+ return true
188+ default :
189+ return false
190+ }
191+ }
192+
193+ func isValidJSONRPCError (errObj map [string ]any ) bool {
194+ code , codeOK := errObj ["code" ].(float64 )
195+ if ! codeOK || math .Trunc (code ) != code {
196+ // JSON-RPC 2.0 requires error.code to be an integer.
197+ return false
198+ }
199+ _ , messageOK := errObj ["message" ].(string )
200+ return messageOK
201+ }
202+
203+ func writeInvalidUpstreamJSONRPCResponse (resp * http.Response , validationErr error ) {
204+ body , err := json .Marshal (map [string ]any {
205+ "jsonrpc" : "2.0" ,
206+ "error" : map [string ]any {
207+ "code" : jsonRPCInvalidUpstreamCode ,
208+ "message" : "Invalid upstream JSON-RPC response" ,
209+ "data" : validationErr .Error (),
210+ },
211+ "id" : nil ,
212+ })
213+ if err != nil {
214+ body = []byte (`{"jsonrpc":"2.0","error":{"code":-32000,"message":"Invalid upstream JSON-RPC response"},"id":null}` )
215+ }
216+
217+ resp .StatusCode = http .StatusBadGateway
218+ resp .Status = fmt .Sprintf ("%d %s" , http .StatusBadGateway , http .StatusText (http .StatusBadGateway ))
219+ resp .Body = io .NopCloser (bytes .NewReader (body ))
220+ resp .ContentLength = int64 (len (body ))
221+
222+ // Replace headers wholesale so upstream session/cookie/cache metadata is not
223+ // smuggled into the proxy-generated error. Only carry the fields needed to
224+ // describe this synthetic body.
225+ resp .Header = http.Header {}
226+ resp .Header .Set ("Content-Type" , "application/json" )
227+ resp .Header .Set ("Content-Length" , fmt .Sprintf ("%d" , len (body )))
228+ resp .Trailer = nil
229+ }
230+
39231// createResponseProcessor is a factory function that creates the appropriate
40232// response processor based on transport type.
41233func createResponseProcessor (
0 commit comments