Merge pull request #520 from koic/bound_oauth_response_bodies

Bound OAuth response bodies in the client

README.md

@@ -3002,6 +3002,11 @@ the HTTP client connects to a moment later. The same-origin rule is what protect
 If you replace the OAuth HTTP client through `MCP::Client::OAuth::Flow.new(http_client_factory:)`, do not add redirect-following middleware. Every check above runs against
 the URL as written, so a connection that follows a `3xx` on its own would reach hosts these rules just refused.
 
+The SDK also bounds what those endpoints may return. A discovery, dynamic client registration, token, or token exchange response is refused once it passes 4 MiB,
+measured as the body arrives rather than after it has been buffered, so a compressed body that expands past the limit is refused partway through the expansion.
+Unlike the transport's `max_message_bytes:`, this limit is not configurable: these documents run to kilobytes in normal operation, and a connection supplied through
+`http_client_factory:` is bounded as well, so there is no way to opt out of it.
+
 #### Customizing the Faraday Connection
 
 You can pass a block to `MCP::Client::HTTP.new` to customize the underlying Faraday connection.

lib/mcp/client/oauth.rb

@@ -1,5 +1,6 @@
 # frozen_string_literal: true
 
+require_relative "oauth/bounded_body"
 require_relative "oauth/discovery"
 require_relative "oauth/flow"
 require_relative "oauth/in_memory_storage"

lib/mcp/client/oauth/bounded_body.rb

@@ -0,0 +1,67 @@
+# frozen_string_literal: true
+
+module MCP
+  class Client
+    module OAuth
+      # Bounds an OAuth response body while it arrives, rather than after it has been buffered.
+      # Discovery documents, registration responses, and token responses are all small by definition,
+      # so a body that keeps growing is never something worth holding in memory. Matches the 4 MiB cap of
+      # `MCP::Client::HTTP::MAX_MESSAGE_BYTES` and `MCP::Client::Stdio::MAX_LINE_BYTES`.
+      class BoundedBody
+        MAX_RESPONSE_BYTES = 4 * 1024 * 1024
+
+        # Raised while the body is read. Each caller translates it into its own error type,
+        # so this never reaches an embedder.
+        class TooLargeError < StandardError; end
+
+        # What the OAuth code reads from a response. The Faraday response itself is not passed on,
+        # so a later caller cannot reach the unbounded `response.body` by accident.
+        Response = Struct.new(:status, :body)
+
+        def initialize(max_bytes: MAX_RESPONSE_BYTES)
+          @max_bytes = max_bytes
+          @buffer = +""
+        end
+
+        # Faraday `on_data` streaming callback. The chunks arrive decompressed: the default `Net::HTTP` adapter negotiates
+        # `Accept-Encoding` itself and reads the body through `Net::HTTPResponse#inflater`, so a small compressed body
+        # that expands past the cap is refused partway through the expansion rather than after it. That holds only while
+        # the connection leaves `Accept-Encoding` to the adapter; see `Flow#default_http_client`.
+        def on_data
+          proc do |chunk, _received_bytes, _env|
+            @buffer << chunk
+
+            raise TooLargeError, too_large_message if @buffer.bytesize > @max_bytes
+          end
+        end
+
+        # The status paired with the bounded body. Adapters that ignore `on_data` leave the buffer empty and deliver
+        # the whole body in `response.body`, so that path is measured here instead; the bytes are already allocated by then,
+        # but refusing them still keeps an over-cap document out of `JSON.parse`.
+        def response_for(response)
+          Response.new(response.status, bounded_body(response))
+        end
+
+        private
+
+        def bounded_body(response)
+          return @buffer unless @buffer.empty?
+
+          body = response.body
+          body = body.is_a?(String) ? body : body.to_s
+          raise TooLargeError, too_large_message if body.bytesize > @max_bytes
+
+          body
+        end
+
+        def too_large_message
+          # Not "the authorization server": protected resource metadata comes from the MCP server's own origin,
+          # so this message covers endpoints on both sides of the flow.
+          "Response body from the OAuth endpoint exceeds #{@max_bytes} bytes"
+        end
+      end
+
+      private_constant :BoundedBody
+    end
+  end
+end

lib/mcp/client/oauth/flow.rb

@@ -1021,26 +1021,54 @@ module MCP
         end
 
         def http_get(url)
-          http_client.get(url)
+          bounded_request do |on_data|
+            http_client.get(url) do |req|
+              req.options.on_data = on_data
+            end
+          end
         end
 
         def http_post_json(url, body)
-          http_client.post(url) do |req|
-            req.headers["Content-Type"] = "application/json"
-            req.headers["Accept"] = "application/json"
-            req.body = JSON.generate(body)
+          bounded_request do |on_data|
+            http_client.post(url) do |req|
+              req.headers["Content-Type"] = "application/json"
+              req.headers["Accept"] = "application/json"
+              req.options.on_data = on_data
+              req.body = JSON.generate(body)
+            end
           end
         end
 
         def http_post_form(url, form, headers: {})
-          http_client.post(url) do |req|
-            req.headers["Content-Type"] = "application/x-www-form-urlencoded"
-            req.headers["Accept"] = "application/json"
-            headers.each { |key, value| req.headers[key] = value }
-            req.body = URI.encode_www_form(form)
+          bounded_request do |on_data|
+            http_client.post(url) do |req|
+              req.headers["Content-Type"] = "application/x-www-form-urlencoded"
+              req.headers["Accept"] = "application/json"
+
+              headers.each do |key, value|
+                req.headers[key] = value
+              end
+
+              req.options.on_data = on_data
+              req.body = URI.encode_www_form(form)
+            end
           end
         end
 
+        # Issues a request with the response body bounded as it arrives, and returns the status paired with
+        # that body. An over-cap response is refused rather than truncated: a partial discovery or token document
+        # cannot be validated, and `fetch_metadata_json` must not fall through to the next candidate URL either,
+        # since the same server would serve the same body.
+        def bounded_request
+          bounded = BoundedBody.new
+
+          response = yield(bounded.on_data)
+
+          bounded.response_for(response)
+        rescue BoundedBody::TooLargeError => e
+          raise AuthorizationError, "#{e.message}."
+        end
+
         def http_client
           @http_client ||= @http_client_factory.call
         end
@@ -1050,8 +1078,14 @@ module MCP
         # that transparently followed a `3xx` would let a server reach a host the checks just refused.
         # A caller passing `http_client_factory:` takes on that responsibility: add redirect following here
         # and the guards above only cover the first hop.
+        #
+        # `Accept-Encoding` is deliberately left unset. `Net::HTTP::GenericRequest` negotiates it and decodes
+        # the response only while the caller has not claimed that header; assigning it turns `decode_content` off,
+        # which would silently move `BoundedBody`'s cap onto compressed bytes and let a small body expand past it
+        # after the check.
         def default_http_client
           require "faraday"
+
           Faraday.new do |faraday|
             faraday.headers["Accept"] = "application/json"
           end

lib/mcp/client/oauth/id_jag_token_exchange.rb

@@ -34,10 +34,13 @@ module MCP
           def request(token_endpoint:, id_token:, client_id:, audience:, resource:, http_client: nil)
             http_client ||= default_http_client
 
+            bounded = BoundedBody.new
+
             response = begin
-              http_client.post(token_endpoint) do |req|
+              raw_response = http_client.post(token_endpoint) do |req|
                 req.headers["Content-Type"] = "application/x-www-form-urlencoded"
                 req.headers["Accept"] = "application/json"
+                req.options.on_data = bounded.on_data
                 req.body = URI.encode_www_form(
                   "grant_type" => GRANT_TYPE,
                   "subject_token" => id_token,
@@ -48,6 +51,10 @@ module MCP
                   "client_id" => client_id,
                 )
               end
+
+              bounded.response_for(raw_response)
+            rescue BoundedBody::TooLargeError => e
+              raise ExchangeError, "#{e.message}."
             rescue Faraday::Error => e
               raise ExchangeError, "Token exchange request to #{token_endpoint} failed: #{e.class}: #{e.message}."
             end
@@ -88,8 +95,12 @@ module MCP
             assertion
           end
 
+          # `Accept-Encoding` is deliberately left unset, for the same reason as `Flow#default_http_client`:
+          # claiming that header turns Net::HTTP's `decode_content` off and would move `BoundedBody`'s cap
+          # onto compressed bytes.
           def default_http_client
             require "faraday"
+
             Faraday.new do |faraday|
               faraday.headers["Accept"] = "application/json"
             end

test/mcp/client/oauth/bounded_body_test.rb

@@ -0,0 +1,150 @@
+# frozen_string_literal: true
+
+require "test_helper"
+require "faraday"
+require "socket"
+require "stringio"
+require "zlib"
+require "mcp/client/oauth"
+
+module MCP
+  class Client
+    module OAuth
+      class BoundedBodyTest < Minitest::Test
+        def test_on_data_rejects_chunks_that_together_exceed_the_limit
+          bounded = bounded_body(max_bytes: 64)
+
+          error = assert_raises(too_large_error) do
+            4.times do
+              bounded.on_data.call("a" * 32, 32, nil)
+            end
+          end
+
+          assert_includes(error.message, "exceeds 64 bytes")
+        end
+
+        def test_on_data_accepts_a_body_that_stops_at_the_limit
+          bounded = bounded_body(max_bytes: 64)
+
+          2.times do
+            bounded.on_data.call("a" * 32, 32, nil)
+          end
+
+          assert_equal("a" * 64, bounded.response_for(stub_response(body: "")).body)
+        end
+
+        def test_response_for_rejects_an_over_limit_body_when_the_adapter_did_not_stream
+          bounded = bounded_body(max_bytes: 64)
+
+          error = assert_raises(too_large_error) do
+            bounded.response_for(stub_response(body: "a" * 65))
+          end
+
+          assert_includes(error.message, "exceeds 64 bytes")
+        end
+
+        def test_response_for_falls_back_to_the_response_body_when_the_adapter_did_not_stream
+          bounded = bounded_body(max_bytes: 64)
+
+          assert_equal("a" * 64, bounded.response_for(stub_response(body: "a" * 64)).body)
+        end
+
+        def test_response_for_coerces_a_non_string_body
+          bounded = bounded_body(max_bytes: 64)
+
+          assert_equal("", bounded.response_for(stub_response(body: nil)).body)
+        end
+
+        def test_response_for_carries_the_status_through
+          bounded = bounded_body(max_bytes: 64)
+
+          assert_equal(404, bounded.response_for(stub_response(status: 404, body: "")).status)
+        end
+
+        # Served over a real socket rather than WebMock: WebMock hands back the stubbed body verbatim
+        # whatever its `Content-Encoding`, so the adapter's inflater never runs and a stubbed version
+        # of this test would assert the opposite of what it looks like.
+        def test_the_cap_counts_decompressed_bytes
+          bounded = bounded_body(max_bytes: 1024 * 1024)
+          compressed = gzip("a" * (5 * 1024 * 1024))
+
+          assert_operator(compressed.bytesize, :<, 1024 * 1024, "the compressed body must fit under the cap")
+
+          serving_gzip(compressed) do |url|
+            assert_raises(too_large_error) do
+              Faraday.new.get(url) do |req|
+                req.options.on_data = bounded.on_data
+                req.options.timeout = 5
+              end
+            end
+          end
+        end
+
+        private
+
+        def gzip(content)
+          io = StringIO.new
+
+          writer = Zlib::GzipWriter.new(io)
+          writer.write(content)
+          writer.close
+
+          io.string
+        end
+
+        def serving_gzip(body)
+          # Other test files load `webmock/minitest`, which blocks real connections process-wide.
+          # This is the one place that needs the adapter's own request path.
+          WebMock.disable! if defined?(WebMock)
+
+          server = TCPServer.new("127.0.0.1", 0)
+          thread = Thread.new do
+            socket = server.accept
+
+            loop do
+              line = socket.gets
+
+              break if line.nil? || line == "\r\n"
+            end
+
+            socket.write(
+              "HTTP/1.1 200 OK\r\n" \
+                "Content-Type: application/json\r\n" \
+                "Content-Encoding: gzip\r\n" \
+                "Content-Length: #{body.bytesize}\r\n" \
+                "Connection: close\r\n\r\n",
+            )
+            socket.write(body)
+
+            socket.close
+          rescue IOError, SystemCallError
+            nil
+          end
+
+          # Scoped to start after both are assigned, so the cleanup below never runs against `nil`
+          # and a failure to open the socket surfaces as itself rather than as a `NoMethodError`.
+          begin
+            yield("http://127.0.0.1:#{server.addr[1]}/")
+          ensure
+            thread.kill
+            server.close
+          end
+        ensure
+          WebMock.enable! if defined?(WebMock)
+        end
+
+        def bounded_body(max_bytes:)
+          OAuth.const_get(:BoundedBody).new(max_bytes: max_bytes)
+        end
+
+        def too_large_error
+          OAuth.const_get(:BoundedBody)::TooLargeError
+        end
+
+        def stub_response(body:, status: 200)
+          Struct.new(:status, :body).new(status, body)
+        end
+      end
+    end
+  end
+end

test/mcp/client/oauth/flow_test.rb

@@ -948,6 +948,69 @@ module MCP
           assert_match(/authorization_servers/i, error.message)
         end
 
+        def test_run_refuses_an_authorization_server_metadata_body_over_the_cap
+          stub_request(:get, @as_metadata_url).to_return(
+            status: 200,
+            headers: { "Content-Type" => "application/json" },
+            body: over_cap_body,
+          )
+
+          error = assert_raises(Flow::AuthorizationError) { run_authorization_flow }
+
+          assert_match(/exceeds \d+ bytes/, error.message)
+        end
+
+        def test_run_refuses_a_protected_resource_metadata_body_over_the_cap
+          # PRM discovery failures select the legacy path by design, so the refusal shows up as
+          # the fallback rather than as a raise. The padded document is valid JSON naming
+          # an authorization server, so contacting that server is exactly what would happen
+          # if the body had been read: the assertion below fails if the cap stops working.
+          stub_request(:any, %r{\Ahttps://srv\.example\.com/}).to_return(status: 404)
+          stub_request(:get, @prm_url).to_return(
+            status: 200,
+            headers: { "Content-Type" => "application/json" },
+            body: JSON.generate(
+              resource: "https://srv.example.com/mcp",
+              authorization_servers: [@auth_base],
+              padding: "a" * (4 * 1024 * 1024),
+            ),
+          )
+
+          assert_raises(Flow::AuthorizationError) { run_authorization_flow }
+
+          assert_not_requested(:get, @as_metadata_url)
+        end
+
+        def test_run_refuses_a_dynamic_client_registration_body_over_the_cap
+          stub_request(:post, "#{@auth_base}/register").to_return(
+            status: 201,
+            headers: { "Content-Type" => "application/json" },
+            body: over_cap_body,
+          )
+
+          error = assert_raises(Flow::AuthorizationError) { run_authorization_flow }
+
+          assert_match(/exceeds \d+ bytes/, error.message)
+        end
+
+        def test_run_refuses_a_token_endpoint_body_over_the_cap
+          stub_request(:post, "#{@auth_base}/token").to_return(
+            status: 200,
+            headers: { "Content-Type" => "application/json" },
+            body: over_cap_body,
+          )
+
+          error = assert_raises(Flow::AuthorizationError) { run_authorization_flow }
+
+          assert_match(/exceeds \d+ bytes/, error.message)
+        end
+
+        # One byte past the cap, so the three request wrappers are each proven to install the streaming callback.
+        # A wrapper that forgot it would buffer this whole body and parse it.
+        def over_cap_body
+          "a" * (4 * 1024 * 1024 + 1)
+        end
+
         def test_run_raises_when_token_response_is_not_a_json_object
           # The token endpoint MUST return a JSON object per RFC 6749 §5.1.
           # A non-object body would otherwise be persisted into the provider

test/mcp/client/oauth/id_jag_token_exchange_test.rb

@@ -101,6 +101,17 @@ module MCP
           error = assert_raises(IDJAGTokenExchange::ExchangeError) { request_exchange }
           assert_match(/not a JSON object/, error.message)
         end
+
+        def test_request_refuses_a_response_body_over_the_cap
+          stub_request(:post, @token_endpoint).to_return(
+            status: 200,
+            headers: { "Content-Type" => "application/json" },
+            body: "a" * (4 * 1024 * 1024 + 1),
+          )
+
+          error = assert_raises(IDJAGTokenExchange::ExchangeError) { request_exchange }
+          assert_match(/exceeds \d+ bytes/, error.message)
+        end
       end
     end
   end