Merge pull request #496 from koic/request_state_security

Add opt-in `requestState` sealing via `RequestStateSecurity`

lib/mcp/server.rb

@@ -12,6 +12,7 @@ require_relative "server_context"
 require_relative "server/capabilities"
 require_relative "server/input_required_result"
 require_relative "server/pagination"
+require_relative "server/request_state_security"
 require_relative "server/transports"
 
 module MCP
@@ -145,7 +146,7 @@ module MCP
     CACHE_SCOPES = ["public", "private"].freeze
 
     attr_accessor :description, :icons, :name, :title, :version, :website_url, :instructions, :tools, :prompts, :resource_templates, :server_context, :configuration, :capabilities, :transport, :logging_message_notification
-    attr_reader :resources, :page_size, :client_capabilities, :ttl_ms, :cache_scope
+    attr_reader :resources, :page_size, :client_capabilities, :ttl_ms, :cache_scope, :request_state_security
 
     def initialize(
       description: nil,
@@ -165,6 +166,7 @@ module MCP
       page_size: nil,
       ttl_ms: nil,
       cache_scope: nil,
+      request_state_security: nil,
       transport: nil
     )
       @description = description
@@ -184,6 +186,7 @@ module MCP
       self.page_size = page_size
       self.ttl_ms = ttl_ms
       self.cache_scope = cache_scope
+      @request_state_security = request_state_security
       @configuration = MCP.configuration.merge(configuration)
       @client = nil
       @client_protocol_version = nil
@@ -606,6 +609,8 @@ module MCP
             session.configure_logging(request_logging) if request_logging.valid_level?
           end
 
+          params = unseal_request_state(params, method: method) if @request_state_security
+
           result = case method
           when Methods::INITIALIZE
             init(params, session: session)
@@ -640,7 +645,7 @@ module MCP
           # Runs after the cancellation check so a cancelled request stays suppressed
           # instead of turning into a gate error response.
           if result.is_a?(InputRequiredResult)
-            result = serialize_input_required_result(result, envelope: envelope, request: params)
+            result = serialize_input_required_result(result, envelope: envelope, request: params, method: method)
           end
 
           # SEP-2322 makes `resultType` REQUIRED on every result a 2026-07-28 server returns;
@@ -718,7 +723,7 @@ module MCP
     # a final result. The capability gate enforces the SEP-2575 rule that servers MUST NOT rely on
     # (or embed requests for) capabilities the client did not declare, and reports every missing capability at
     # once so the client sees the full set.
-    def serialize_input_required_result(result, envelope:, request:)
+    def serialize_input_required_result(result, envelope:, request:, method:)
       if envelope.nil?
         raise RequestHandlerError.new(
           "input_required results require the 2026-07-28 stateless lifecycle (SEP-2322)",
@@ -731,7 +736,76 @@ module MCP
       raise MissingRequiredClientCapabilityError.new(missing, request) unless missing.empty?
 
       add_instrumentation_data(input_required: true)
-      result.to_h
+      serialized = result.to_h
+
+      if @request_state_security && serialized[:requestState]
+        serialized = serialized.merge(requestState: @request_state_security.seal(
+          serialized[:requestState],
+          method: method,
+          target: mrtr_target(request),
+          arguments_digest: mrtr_arguments_digest(request),
+        ))
+      end
+
+      serialized
+    end
+
+    # Methods whose results may be `input_required` and whose retried requests carry
+    # `inputResponses`/`requestState` (SEP-2322).
+    MRTR_METHODS = [Methods::TOOLS_CALL, Methods::PROMPTS_GET, Methods::RESOURCES_READ].freeze
+
+    # Replaces a sealed client-echoed `requestState` with its verified plaintext before dispatch,
+    # so handlers always read the state they wrote. A tampered, expired, or cross-request token is
+    # rejected as invalid params, matching the Python SDK's "Invalid or expired requestState" behavior.
+    def unseal_request_state(params, method:)
+      return params unless MRTR_METHODS.include?(method)
+      return params unless params.is_a?(Hash)
+
+      sealed = params[:requestState] || params["requestState"]
+      return params unless sealed
+
+      plaintext = @request_state_security.unseal(
+        sealed,
+        method: method,
+        target: mrtr_target(params),
+        arguments_digest: mrtr_arguments_digest(params),
+      )
+      key = params.key?("requestState") ? "requestState" : :requestState
+      params.merge(key => plaintext)
+    rescue RequestStateSecurity::InvalidStateError => e
+      raise RequestHandlerError.new(
+        "Invalid or expired requestState",
+        params,
+        error_type: :invalid_params,
+        error_code: JsonRpcHandler::ErrorCode::INVALID_PARAMS,
+        original_error: e,
+      )
+    end
+
+    def mrtr_target(params)
+      return "" unless params.is_a?(Hash)
+
+      params[:name] || params["name"] || params[:uri] || params["uri"] || ""
+    end
+
+    # Digest of the originating arguments, binding a sealed state to retries of
+    # the same call with the same inputs. Keys are stringified and sorted recursively
+    # so symbol/string parses of identical JSON digest identically.
+    def mrtr_arguments_digest(params)
+      arguments = params.is_a?(Hash) ? params[:arguments] || params["arguments"] : nil
+      OpenSSL::Digest::SHA256.hexdigest(canonical_json(arguments || {}))
+    end
+
+    def canonical_json(value)
+      case value
+      when Hash
+        pairs = value.map { |key, nested| [key.to_s, nested] }.sort_by(&:first)
+        "{#{pairs.map { |key, nested| "#{key.to_json}:#{canonical_json(nested)}" }.join(",")}}"
+      when Array
+        "[#{value.map { |element| canonical_json(element) }.join(",")}]"
+      else
+        value.to_json
+      end
     end
 
     # Extracts the SEP-2322 retry fields a client sends when re-issuing a request:

lib/mcp/server/request_state_security.rb

@@ -0,0 +1,131 @@
+# frozen_string_literal: true
+
+require "json"
+require "openssl"
+
+module MCP
+  class Server
+    # Opt-in protection for the SEP-2322 `requestState` echo. The opaque continuation string leaves the server,
+    # sits in the client's hands, and comes back as client-controlled input, so it must be treated like
+    # any other untrusted data. Sealing encrypts the state with AES-256-GCM (clients cannot read it) and binds
+    # a claims envelope that unsealing verifies fail-closed:
+    #
+    # - `exp`: a TTL window (re-sealed each round)
+    # - `m` / `t`: the originating method and target (tool/prompt name or resource URI)
+    # - `a`: a digest of the originating arguments, so the state only resumes the same call with the same inputs
+    # - `aud`: an optional audience, so tokens cannot cross servers sharing a key
+    #
+    # Pass an instance via `Server.new(request_state_security:)` and the seal/unseal happens transparently;
+    # handlers keep reading plaintext through `server_context.request_state`. Without it, the state crosses
+    # the wire exactly as the handler wrote it (the author's responsibility, matching the Python SDK's low-level Server).
+    # The key must be shared across workers in multi-process
+    #
+    # deployments; a per-process random key makes retries that land on another worker fail with an invalid-state error,
+    # forcing clients to restart the flow.
+    #
+    # The token format is `v1.<base64url(iv || ciphertext || tag)>`, with the version prefix bound as GCM associated data,
+    # following the Python SDK's `AESGCMRequestStateCodec`.
+    class RequestStateSecurity
+      class InvalidStateError < StandardError; end
+
+      VERSION_PREFIX = "v1."
+      KEY_BYTES = 32
+      IV_BYTES = 12
+      TAG_BYTES = 16
+      DEFAULT_TTL = 300
+
+      def initialize(key:, ttl: DEFAULT_TTL, audience: nil)
+        unless key.is_a?(String) && key.bytesize == KEY_BYTES
+          raise ArgumentError, "key must be a #{KEY_BYTES}-byte String"
+        end
+        unless ttl.is_a?(Numeric) && ttl.positive?
+          raise ArgumentError, "ttl must be a positive number of seconds"
+        end
+
+        @key = key.dup.force_encoding(Encoding::BINARY).freeze
+        @ttl = ttl
+        @audience = audience
+      end
+
+      # Seals a plaintext state into an opaque token bound to the originating request.
+      def seal(state, method:, target:, arguments_digest:)
+        claims = {
+          v: 1,
+          exp: Time.now.to_i + @ttl,
+          m: method,
+          t: target,
+          a: arguments_digest,
+          aud: @audience,
+          s: state,
+        }.compact
+
+        cipher = OpenSSL::Cipher.new("aes-256-gcm").encrypt
+        cipher.key = @key
+        iv = cipher.random_iv
+        cipher.auth_data = VERSION_PREFIX
+        ciphertext = cipher.update(JSON.generate(claims)) + cipher.final
+
+        VERSION_PREFIX + base64url_encode(iv + ciphertext + cipher.auth_tag)
+      end
+
+      # Unseals a client-echoed token and verifies every claim, failing closed with
+      # `InvalidStateError` on tampering, expiry, or a claims mismatch.
+      def unseal(sealed, method:, target:, arguments_digest:)
+        unless sealed.is_a?(String) && sealed.start_with?(VERSION_PREFIX)
+          raise InvalidStateError, "malformed token"
+        end
+
+        blob = base64url_decode(sealed.delete_prefix(VERSION_PREFIX))
+        raise InvalidStateError, "malformed token" if blob.bytesize < IV_BYTES + TAG_BYTES
+
+        iv = blob.byteslice(0, IV_BYTES)
+        tag = blob.byteslice(-TAG_BYTES, TAG_BYTES)
+        ciphertext = blob.byteslice(IV_BYTES, blob.bytesize - IV_BYTES - TAG_BYTES)
+
+        cipher = OpenSSL::Cipher.new("aes-256-gcm").decrypt
+        cipher.key = @key
+        cipher.iv = iv
+        cipher.auth_tag = tag
+        cipher.auth_data = VERSION_PREFIX
+        plaintext = begin
+          cipher.update(ciphertext) + cipher.final
+        rescue OpenSSL::Cipher::CipherError
+          raise InvalidStateError, "authentication failed"
+        end
+
+        claims = begin
+          JSON.parse(plaintext, symbolize_names: true)
+        rescue JSON::ParserError
+          raise InvalidStateError, "malformed claims"
+        end
+
+        verify!(claims, method: method, target: target, arguments_digest: arguments_digest)
+        claims[:s]
+      end
+
+      private
+
+      def verify!(claims, method:, target:, arguments_digest:)
+        raise InvalidStateError, "unsupported version" unless claims[:v] == 1
+        raise InvalidStateError, "expired" unless claims[:exp].is_a?(Integer) && Time.now.to_i <= claims[:exp]
+        raise InvalidStateError, "method mismatch" unless claims[:m] == method
+        raise InvalidStateError, "target mismatch" unless claims[:t] == target
+        raise InvalidStateError, "arguments mismatch" unless claims[:a] == arguments_digest
+        raise InvalidStateError, "audience mismatch" unless claims[:aud] == @audience
+        raise InvalidStateError, "missing state" unless claims[:s].is_a?(String)
+      end
+
+      def base64url_encode(data)
+        [data].pack("m0").tr("+/", "-_").delete("=")
+      end
+
+      def base64url_decode(encoded)
+        padded = encoded.tr("-_", "+/")
+        padded += "=" * ((4 - padded.length % 4) % 4)
+        padded.unpack1("m0")
+      rescue ArgumentError
+        raise InvalidStateError, "malformed token"
+      end
+    end
+  end
+end