Merge pull request #496 from koic/request_state_security
Add opt-in `requestState` sealing via `RequestStateSecurity`
This commit is contained in:
Koichi ITO
GitHub
commit
6dfef81a75
2 changed files with 417 additions and 6 deletions
90
test/mcp/server/request_state_security_test.rb
Normal file
90
test/mcp/server/request_state_security_test.rb
Normal file
|
|
@ -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
|
||||
|
|
@ -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:|
|
||||
|
|
|
|||
Loading…
Reference in a new issue