Merge pull request #496 from koic/request_state_security

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

README.md

@@ -60,8 +60,11 @@ It implements the Model Context Protocol specification, handling model context r
   which the modern lifecycle forbids. On the retried request the handler re-runs from the start and reads the answers via
   `server_context.input_responses` / `server_context.input_response(key)` and the echoed opaque `server_context.request_state`
   (deterministic replay; the server holds no memory between rounds). The SDK rejects issuance on legacy requests and returns `-32021`
-  when an embedded request needs a client capability the request did not declare. Note that the echoed `requestState` arrives as
-  client-controlled input; treat it accordingly
+  when an embedded request needs a client capability the request did not declare. The echoed `requestState` arrives as
+  client-controlled input: pass `MCP::Server::RequestStateSecurity.new(key:)` (a 32-byte key) via `Server.new(request_state_security:)` to
+  have it sealed with AES-256-GCM and bound to a TTL plus the originating method, target, and arguments, all transparently to handlers.
+  Multi-process deployments must share the key across workers; without `request_state_security:` the state crosses the wire exactly as
+  the handler wrote it and protecting it is the handler author's responsibility
 - `ping` - Simple health check
 - `logging/setLevel` - Configures the minimum log level for the server
 - `tools/list` - Lists all registered tools and their schemas

conformance/server.rb

@@ -575,6 +575,40 @@ module Conformance
       end
     end
 
+    class TestInputRequiredResultTamperedState < MCP::Tool
+      tool_name "test_input_required_result_tampered_state"
+      description "A tool whose sealed requestState rejects tampered echoes (SEP-2322)"
+
+      class << self
+        # A tampered `requestState` never reaches this handler: the server-level unsealing rejects it with -32602
+        # before dispatch, so the tool body is a plain MRTR flow.
+        def call(server_context:, **_args)
+          state = Mrtr.parse_state(server_context.request_state)
+          confirmed = Mrtr.accepted_content(server_context.input_response("confirm"))
+
+          if state && state["kind"] == "tamper-check" && confirmed
+            return MCP::Tool::Response.new([MCP::Content::Text.new("requestState integrity verified").to_h])
+          end
+
+          MCP::Server::InputRequiredResult.new(
+            input_requests: {
+              confirm: {
+                method: "elicitation/create",
+                params: {
+                  message: "Confirm to continue",
+                  requestedSchema: {
+                    type: "object",
+                    properties: { ok: { type: "boolean" } },
+                  },
+                },
+              },
+            },
+            request_state: JSON.generate({ kind: "tamper-check", nonce: SecureRandom.hex(8) }),
+          )
+        end
+      end
+    end
+
     class TestInputRequiredResultCapabilities < MCP::Tool
       tool_name "test_input_required_result_capabilities"
       description "A tool that only embeds input requests the client's declared capabilities can fulfill (SEP-2322)"
@@ -869,6 +903,7 @@ module Conformance
           Tools::TestInputRequiredResultRequestState,
           Tools::TestInputRequiredResultMultipleInputs,
           Tools::TestInputRequiredResultMultiRound,
+          Tools::TestInputRequiredResultTamperedState,
           Tools::TestInputRequiredResultCapabilities,
           Tools::TestStreamingElicitation,
           Tools::TestLoggingTool,
@@ -882,6 +917,10 @@ module Conformance
         ],
         resources: resources,
         resource_templates: resource_templates,
+        # A per-boot random key is enough for conformance: every MRTR round trip completes
+        # within one server process, and sealing keeps the fixture's `requestState` values
+        # tamper-evident without any tool-level verification code.
+        request_state_security: MCP::Server::RequestStateSecurity.new(key: SecureRandom.bytes(32)),
         capabilities: {
           tools: { listChanged: true },
           prompts: { listChanged: true },

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

test/mcp/server/request_state_security_test.rb

@@ -0,0 +1,90 @@
+# frozen_string_literal: true
+
+require "test_helper"
+
+module MCP
+  class Server
+    class RequestStateSecurityTest < ActiveSupport::TestCase
+      include ActiveSupport::Testing::TimeHelpers
+
+      KEY = ("k" * 32).freeze
+
+      test "validates constructor arguments" do
+        assert_raises(ArgumentError) { RequestStateSecurity.new(key: "short") }
+        assert_raises(ArgumentError) { RequestStateSecurity.new(key: KEY, ttl: 0) }
+      end
+
+      test "seals and unseals a state bound to the originating request" do
+        security = RequestStateSecurity.new(key: KEY, audience: "my_server")
+        sealed = security.seal("plain-state", method: "tools/call", target: "my_tool", arguments_digest: "digest")
+
+        assert sealed.start_with?("v1.")
+        refute_includes sealed, "plain-state"
+        assert_equal "plain-state",
+          security.unseal(sealed, method: "tools/call", target: "my_tool", arguments_digest: "digest")
+      end
+
+      test "rejects expired tokens" do
+        security = RequestStateSecurity.new(key: KEY, ttl: 60)
+        sealed = security.seal("state", method: "tools/call", target: "t", arguments_digest: "d")
+
+        travel 120 do
+          assert_raises(RequestStateSecurity::InvalidStateError) do
+            security.unseal(sealed, method: "tools/call", target: "t", arguments_digest: "d")
+          end
+        end
+      end
+
+      test "rejects claims mismatches fail-closed" do
+        security = RequestStateSecurity.new(key: KEY, audience: "my_server")
+        sealed = security.seal("state", method: "tools/call", target: "my_tool", arguments_digest: "digest")
+
+        assert_raises(RequestStateSecurity::InvalidStateError) do
+          security.unseal(sealed, method: "prompts/get", target: "my_tool", arguments_digest: "digest")
+        end
+        assert_raises(RequestStateSecurity::InvalidStateError) do
+          security.unseal(sealed, method: "tools/call", target: "other_tool", arguments_digest: "digest")
+        end
+        assert_raises(RequestStateSecurity::InvalidStateError) do
+          security.unseal(sealed, method: "tools/call", target: "my_tool", arguments_digest: "other")
+        end
+
+        other_audience = RequestStateSecurity.new(key: KEY, audience: "other_server")
+        assert_raises(RequestStateSecurity::InvalidStateError) do
+          other_audience.unseal(sealed, method: "tools/call", target: "my_tool", arguments_digest: "digest")
+        end
+      end
+
+      test "rejects tampered and malformed tokens" do
+        security = RequestStateSecurity.new(key: KEY)
+        sealed = security.seal("state", method: "tools/call", target: "t", arguments_digest: "d")
+
+        tampered = sealed.dup
+        tampered[-1] = tampered[-1] == "A" ? "B" : "A"
+        assert_raises(RequestStateSecurity::InvalidStateError) do
+          security.unseal(tampered, method: "tools/call", target: "t", arguments_digest: "d")
+        end
+
+        assert_raises(RequestStateSecurity::InvalidStateError) do
+          security.unseal("not-a-token", method: "tools/call", target: "t", arguments_digest: "d")
+        end
+        assert_raises(RequestStateSecurity::InvalidStateError) do
+          security.unseal("v1.!!!", method: "tools/call", target: "t", arguments_digest: "d")
+        end
+        assert_raises(RequestStateSecurity::InvalidStateError) do
+          security.unseal(nil, method: "tools/call", target: "t", arguments_digest: "d")
+        end
+      end
+
+      test "rejects tokens sealed under a different key" do
+        sealed = RequestStateSecurity.new(key: KEY)
+          .seal("state", method: "tools/call", target: "t", arguments_digest: "d")
+        other = RequestStateSecurity.new(key: "x" * 32)
+
+        assert_raises(RequestStateSecurity::InvalidStateError) do
+          other.unseal(sealed, method: "tools/call", target: "t", arguments_digest: "d")
+        end
+      end
+    end
+  end
+end

test/mcp/server_test.rb

@@ -510,6 +510,80 @@ module MCP
       assert_equal "input_required", response.dig(:result, :resultType)
     end
 
+    test "#handle seals and unseals requestState transparently when request_state_security is set" do
+      security = Server::RequestStateSecurity.new(key: "k" * 32)
+      server = Server.new(name: "mrtr_test", tools: [], request_state_security: security)
+      seen_state = nil
+      server.define_tool(name: "mrtr_tool") do |server_context:|
+        if server_context.request_state
+          seen_state = server_context.request_state
+          Tool::Response.new([{ type: "text", text: "done" }])
+        else
+          Server::InputRequiredResult.new(request_state: "plain-state")
+        end
+      end
+
+      first = server.handle(modern_request("tools/call", { name: "mrtr_tool", arguments: {} }))
+      sealed = first.dig(:result, :requestState)
+
+      assert sealed.start_with?("v1.")
+      refute_equal "plain-state", sealed
+
+      retry_response = server.handle(modern_request("tools/call", {
+        name: "mrtr_tool",
+        arguments: {},
+        requestState: sealed,
+      }))
+
+      refute_nil retry_response[:result]
+      # The handler reads the plaintext it originally wrote.
+      assert_equal "plain-state", seen_state
+    end
+
+    test "#handle rejects a tampered or cross-call requestState echo with -32602" do
+      security = Server::RequestStateSecurity.new(key: "k" * 32)
+      server = Server.new(name: "mrtr_test", tools: [], request_state_security: security)
+      server.define_tool(name: "mrtr_tool") do |server_context:|
+        if server_context.request_state
+          Tool::Response.new([{ type: "text", text: "done" }])
+        else
+          Server::InputRequiredResult.new(request_state: "plain-state")
+        end
+      end
+
+      first = server.handle(modern_request("tools/call", { name: "mrtr_tool", arguments: {} }))
+      sealed = first.dig(:result, :requestState)
+
+      tampered = server.handle(modern_request("tools/call", {
+        name: "mrtr_tool",
+        arguments: {},
+        requestState: sealed.sub("v1.", "v1.x"),
+      }))
+      # Different arguments than the sealing call: the digest claim no longer matches.
+      cross_call = server.handle(modern_request("tools/call", {
+        name: "mrtr_tool",
+        arguments: { other: true },
+        requestState: sealed,
+      }))
+
+      assert_equal JsonRpcHandler::ErrorCode::INVALID_PARAMS, tampered.dig(:error, :code)
+      assert_equal "Invalid or expired requestState", tampered.dig(:error, :message)
+      assert_equal JsonRpcHandler::ErrorCode::INVALID_PARAMS, cross_call.dig(:error, :code)
+    end
+
+    test "#handle passes requestState through untouched without request_state_security" do
+      server = Server.new(name: "mrtr_test", tools: [])
+      seen_state = nil
+      server.define_tool(name: "mrtr_tool") do |server_context:|
+        seen_state = server_context.request_state
+        Tool::Response.new([{ type: "text", text: "done" }])
+      end
+
+      server.handle(modern_request("tools/call", { name: "mrtr_tool", requestState: "raw-state" }))
+
+      assert_equal "raw-state", seen_state
+    end
+
     test "#handle prompts/get serializes an input_required result on a modern request" do
       server = Server.new(name: "mrtr_test")
       server.define_prompt(name: "mrtr_prompt", arguments: []) do |_args, server_context:|