diff --git a/README.md b/README.md index f83cdaa5..d563a541 100644 --- a/README.md +++ b/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 diff --git a/conformance/server.rb b/conformance/server.rb index dca807db..ef138d0e 100644 --- a/conformance/server.rb +++ b/conformance/server.rb @@ -575,6 +575,40 @@ def second_round(name) 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 @@ def build_server Tools::TestInputRequiredResultRequestState, Tools::TestInputRequiredResultMultipleInputs, Tools::TestInputRequiredResultMultiRound, + Tools::TestInputRequiredResultTamperedState, Tools::TestInputRequiredResultCapabilities, Tools::TestStreamingElicitation, Tools::TestLoggingTool, @@ -882,6 +917,10 @@ def build_server ], 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 }, diff --git a/lib/mcp/server.rb b/lib/mcp/server.rb index 262352f6..eaca0def 100644 --- a/lib/mcp/server.rb +++ b/lib/mcp/server.rb @@ -12,6 +12,7 @@ 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 @@ class ValidationError < StandardError; end 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 @@ def initialize( page_size: nil, ttl_ms: nil, cache_scope: nil, + request_state_security: nil, transport: nil ) @description = description @@ -184,6 +186,7 @@ def initialize( 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 @@ -594,6 +597,8 @@ def handle_request(request, method, session: nil, related_request_id: nil) 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) @@ -628,7 +633,7 @@ def handle_request(request, method, session: nil, related_request_id: nil) # 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; @@ -707,7 +712,7 @@ def lift_request_envelope(params, method:, session:) # 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)", @@ -720,7 +725,76 @@ def serialize_input_required_result(result, envelope:, request:) 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: diff --git a/lib/mcp/server/request_state_security.rb b/lib/mcp/server/request_state_security.rb new file mode 100644 index 00000000..bc96a858 --- /dev/null +++ b/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.`, 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 diff --git a/test/mcp/server/request_state_security_test.rb b/test/mcp/server/request_state_security_test.rb new file mode 100644 index 00000000..29b7d794 --- /dev/null +++ b/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 diff --git a/test/mcp/server_test.rb b/test/mcp/server_test.rb index 69d9b8cb..162bdfb3 100644 --- a/test/mcp/server_test.rb +++ b/test/mcp/server_test.rb @@ -505,6 +505,80 @@ class ServerTest < ActiveSupport::TestCase 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:|