Merge pull request #560 from koic/let_a_provider_customize_the_oauth_http_client

Let a provider customize the HTTP client the OAuth flow uses

lib/mcp/client/oauth/bounded_body.rb

@@ -26,7 +26,7 @@ module MCP
         # 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`.
+        # the connection leaves `Accept-Encoding` to the adapter; see `Flow.build_http_client`.
         def on_data
           proc do |chunk, _received_bytes, _env|
             @buffer << chunk

lib/mcp/client/oauth/client_credentials_provider.rb

@@ -46,6 +46,8 @@ module MCP
       #   server requires beyond the grant itself (Auth0's `audience`, for example).
       #   A key in `Flow::RESERVED_TOKEN_REQUEST_PARAMS` raises `Flow::InvalidTokenRequestParamsError`.
       #   The Hash is copied and frozen. See `StorageBackedProvider#token_request_params`.
+      # - `http_client_customizer` - Callable invoked with the `Faraday::Connection` the flow builds for
+      #   its own requests; see `StorageBackedProvider#http_client_customizer`.
       class ClientCredentialsProvider
         include StorageBackedProvider
 
@@ -66,7 +68,8 @@ module MCP
           scope: nil,
           storage: nil,
           authorization_request_validator: nil,
-          token_request_params: nil
+          token_request_params: nil,
+          http_client_customizer: nil
         )
           if blank?(client_id)
             raise InvalidCredentialsError, "client_id is required for the client_credentials grant."
@@ -104,6 +107,8 @@ module MCP
             client_information["client_secret"] = client_secret
           end
 
+          http_client_customizer = validated_http_client_customizer(http_client_customizer)
+
           @client_id = client_id
           @private_key = private_key
           @signing_algorithm = signing_algorithm
@@ -111,6 +116,7 @@ module MCP
           @storage = storage || InMemoryStorage.new
           @authorization_request_validator = authorization_request_validator
           @token_request_params = frozen_token_request_params(token_request_params)
+          @http_client_customizer = http_client_customizer
           @storage.save_client_information(client_information)
         end
 

lib/mcp/client/oauth/cross_app_access_provider.rb

@@ -29,6 +29,8 @@ module MCP
       #   the authorization server requires beyond the grant itself. A key in `Flow::RESERVED_TOKEN_REQUEST_PARAMS` raises
       #   `Flow::InvalidTokenRequestParamsError`. The Hash is copied and frozen.
       #   See `StorageBackedProvider#token_request_params`.
+      # - `http_client_customizer` - Callable invoked with the `Faraday::Connection` the flow builds for its own requests;
+      #   see `StorageBackedProvider#http_client_customizer`.
       #
       # https://github.com/modelcontextprotocol/modelcontextprotocol/issues/990
       class CrossAppAccessProvider
@@ -46,7 +48,8 @@ module MCP
           scope: nil,
           storage: nil,
           authorization_request_validator: nil,
-          token_request_params: nil
+          token_request_params: nil,
+          http_client_customizer: nil
         )
           if blank?(client_id)
             raise InvalidConfigurationError, "client_id is required for the jwt-bearer grant."
@@ -60,11 +63,14 @@ module MCP
             raise InvalidConfigurationError, "assertion_provider must be callable as `call(audience:, resource:)` and return the ID-JAG assertion."
           end
 
+          http_client_customizer = validated_http_client_customizer(http_client_customizer)
+
           @assertion_provider = assertion_provider
           @scope = scope
           @storage = storage || InMemoryStorage.new
           @authorization_request_validator = authorization_request_validator
           @token_request_params = frozen_token_request_params(token_request_params)
+          @http_client_customizer = http_client_customizer
           @storage.save_client_information(
             "client_id" => client_id,
             "client_secret" => client_secret,

lib/mcp/client/oauth/flow.rb

@@ -63,6 +63,71 @@ module MCP
         # a failed refresh as a reason to run the interactive flow.
         class InvalidTokenRequestParamsError < ArgumentError; end
 
+        # Raised by `RequestedOriginGuard` when middleware added through the provider's `http_client_customizer`
+        # would send a request to an origin other than the one the flow validated, or has dropped the record of
+        # the URL the flow asked for. An `ArgumentError` because the middleware is a configuration mistake,
+        # and deliberately outside `AuthorizationError`, which discovery treats as "nothing published"
+        # and `MCP::Client::HTTP` treats on a failed refresh as a reason to run the interactive flow.
+        class DestinationMismatchError < ArgumentError; end
+
+        # Faraday middleware registered on the connection `build_http_client` assembles before the customizer
+        # is invoked, so with the usual `use` it sits ahead of the customizer's middleware and sees the URL exactly
+        # as the flow requested it, which it records on the request environment for `RequestedOriginGuard`.
+        # The guard covers what happens to a request after that record; middleware inserted ahead of it with
+        # `builder.insert(0, ...)` that rewrites the URL before it or rebuilds the environment is outside the guard.
+        # The record lives on the environment, not in `env.request.context`: that slot belongs to the application,
+        # which may fill it on the connection or replace it from a middleware of its own. Only the first URL seen
+        # on an environment is recorded: a middleware inserted ahead of this one that re-enters the stack after
+        # a `3xx` with the same environment, or with its `dup`, which shares the record, cannot replace it with
+        # the redirected URL.
+        class RequestedURLStamp
+          KEY = :mcp_oauth_requested_url
+
+          def initialize(app)
+            @app = app
+          end
+
+          def call(env)
+            env[KEY] ||= env.url.to_s
+            @app.call(env)
+          end
+        end
+
+        # Faraday middleware registered last on that connection, so it sees `env.url` after any customizer-added
+        # middleware has rewritten it or followed a redirect. A request that would leave the origin the flow asked
+        # for is refused before it reaches the adapter, since every destination check ran against the URL as
+        # written; a same-origin change stays with the server those checks admitted. The record survives
+        # the `env.dup` that redirect-following middleware performs, and a request that arrives without it is
+        # refused as well, so a middleware that rebuilds the environment fails closed rather than open.
+        # The origin boundary resembles the one the Python SDK keeps for its own auth requests, which follows
+        # a redirect itself only within the origin; this flow follows none.
+        class RequestedOriginGuard
+          def initialize(app)
+            @app = app
+          end
+
+          def call(env)
+            requested = env[RequestedURLStamp::KEY]
+            unless requested
+              raise DestinationMismatchError, <<~MESSAGE
+                Request to #{Discovery.canonicalize_origin_and_path(env.url.to_s).inspect} carries no record of \
+                the URL the flow asked for; middleware that rebuilds the request environment is refused.
+              MESSAGE
+            end
+
+            unless Discovery.same_origin?(env.url.to_s, requested)
+              raise DestinationMismatchError, <<~MESSAGE
+                Request to #{Discovery.canonicalize_origin_and_path(requested).inspect} would be sent to \
+                #{Discovery.canonicalize_origin_and_path(env.url.to_s).inspect}, on a different origin; \
+                middleware that follows redirects or rewrites URLs is refused.
+              MESSAGE
+            end
+
+            @app.call(env)
+          end
+        end
+        private_constant :RequestedURLStamp, :RequestedOriginGuard
+
         class << self
           # Returns why `params` cannot ride a token request as `token_request_params`, or `nil` when it can.
           # Shared by the provider constructors and the flow, which both refuse the value with `InvalidTokenRequestParamsError`,
@@ -82,6 +147,33 @@ module MCP
 
             nil
           end
+
+          # Builds the connection the flow uses for its own requests: the SDK's defaults, `RequestedURLStamp`,
+          # then `customizer` (a provider's `http_client_customizer`, called with the `Faraday::Connection`),
+          # then `RequestedOriginGuard` last so it sees what the customizer's middleware does to each request
+          # after the stamp recorded it. Every request on the connection passes through both, so a caller using
+          # it directly is held to the same origin rule.
+          #
+          # Deliberately built without redirect-following middleware. Every destination check in this class runs
+          # against the URL as written, before the request goes out, so a connection that transparently followed
+          # a `3xx` would let a server reach a host the checks just refused. The guard turns following at
+          # the middleware level into a refusal; following inside an adapter stays invisible, so a customizer must
+          # not enable it.
+          #
+          # `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 build_http_client(customizer = nil)
+            require "faraday"
+
+            Faraday.new do |faraday|
+              faraday.headers["Accept"] = "application/json"
+              faraday.use(RequestedURLStamp)
+              customizer&.call(faraday)
+              faraday.use(RequestedOriginGuard)
+            end
+          end
         end
 
         def initialize(provider:, http_client_factory: nil)
@@ -1298,22 +1390,18 @@ module MCP
           @http_client ||= @http_client_factory.call
         end
 
-        # Deliberately built without redirect-following middleware. Every destination check in
-        # this class runs against the URL as written, before the request goes out, so a connection
-        # 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.
+        # A connection supplied through `http_client_factory:` replaces this one, the provider's customizer and
+        # `RequestedOriginGuard` included, so that caller takes on the redirect responsibility described on
+        # `build_http_client`; `bounded_request` caps its responses all the same.
         def default_http_client
-          require "faraday"
+          self.class.build_http_client(provider_http_client_customizer)
+        end
 
-          Faraday.new do |faraday|
-            faraday.headers["Accept"] = "application/json"
-          end
+        # `nil` for a provider that predates the hook or leaves it unset.
+        def provider_http_client_customizer
+          return unless @provider.respond_to?(:http_client_customizer)
+
+          @provider.http_client_customizer
         end
 
         def response_body_string(response)

lib/mcp/client/oauth/id_jag_token_exchange.rb

@@ -95,7 +95,7 @@ module MCP
             assertion
           end
 
-          # `Accept-Encoding` is deliberately left unset, for the same reason as `Flow#default_http_client`:
+          # `Accept-Encoding` is deliberately left unset, for the same reason as `Flow.build_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

lib/mcp/client/oauth/provider.rb

@@ -53,6 +53,8 @@ module MCP
       #   the grant itself. The authorization request is not affected. A key in `Flow::RESERVED_TOKEN_REQUEST_PARAMS`
       #   raises `Flow::InvalidTokenRequestParamsError`.
       #   The Hash is copied and frozen. See `StorageBackedProvider#token_request_params`.
+      # - `http_client_customizer` - Callable invoked with the `Faraday::Connection` the flow builds for
+      #   its own requests; see `StorageBackedProvider#http_client_customizer`.
       class Provider
         include StorageBackedProvider
 
@@ -93,7 +95,8 @@ module MCP
           storage: nil,
           client_id_metadata_document_url: nil,
           authorization_request_validator: nil,
-          token_request_params: nil
+          token_request_params: nil,
+          http_client_customizer: nil
         )
           unless Discovery.secure_url?(redirect_uri)
             raise InsecureRedirectURIError,
@@ -115,6 +118,8 @@ module MCP
                 "per the MCP authorization specification and `draft-ietf-oauth-client-id-metadata-document`."
           end
 
+          http_client_customizer = validated_http_client_customizer(http_client_customizer)
+
           @client_metadata = client_metadata
           @redirect_uri = redirect_uri
           @redirect_handler = redirect_handler
@@ -124,6 +129,7 @@ module MCP
           @client_id_metadata_document_url = client_id_metadata_document_url
           @authorization_request_validator = authorization_request_validator
           @token_request_params = frozen_token_request_params(token_request_params)
+          @http_client_customizer = http_client_customizer
         end
 
         # Identifies the OAuth flow this provider drives.

lib/mcp/client/oauth/storage_backed_provider.rb

@@ -28,6 +28,13 @@ module MCP
         # with `Flow::InvalidTokenRequestParamsError`.
         attr_reader :token_request_params
 
+        # Optional callable invoked with the `Faraday::Connection` the flow builds for its own requests,
+        # after the SDK's defaults and before its origin guard (see `Flow.build_http_client`), so an application can
+        # add middleware to discovery, registration, and token requests or swap their adapter. `nil` (the default)
+        # keeps the default connection. The transport's own connection and customizer block never serve these requests:
+        # they are bound to the MCP server URL and carry headers meant for that server.
+        attr_reader :http_client_customizer
+
         def access_token
           tokens&.dig("access_token") || tokens&.dig(:access_token)
         end
@@ -66,6 +73,14 @@ module MCP
 
           params.each_with_object({}) { |(key, value), copy| copy[key.dup.freeze] = value.dup.freeze }.freeze
         end
+
+        # Returns `customizer` when it is `nil` or callable and raises otherwise, so a connection object passed
+        # in place of a callable fails at construction rather than at the first `401`.
+        def validated_http_client_customizer(customizer)
+          return customizer if customizer.nil? || customizer.respond_to?(:call)
+
+          raise ArgumentError, "http_client_customizer must respond to call (got #{customizer.class})."
+        end
       end
     end
   end