tower_mcp/
mrtr.rs

1//! Server-side helpers for Multi Round-Trip Requests (SEP-2322).
2//!
3//! The final protocol carries continuation state through an untrusted client.
4//! [`RequestStateCodec`] produces versioned, expiring, HMAC-SHA256-protected
5//! tokens so stateless server instances can share continuation state safely by
6//! sharing the same key.
7
8use std::time::{Duration, SystemTime, UNIX_EPOCH};
9
10use base64::Engine;
11use base64::engine::general_purpose::URL_SAFE_NO_PAD;
12use serde::Serialize;
13use serde::de::DeserializeOwned;
14use sha2::{Digest, Sha256};
15
16use crate::protocol::InputResponses;
17
18const TOKEN_VERSION: &str = "v1";
19const SHA256_BLOCK_SIZE: usize = 64;
20const DEFAULT_MAX_TOKEN_BYTES: usize = 64 * 1024;
21
22/// MRTR continuation values supplied on a retry of the original request.
23///
24/// The router inserts this value into [`crate::RequestContext`] for
25/// `tools/call`, `prompts/get`, and `resources/read`. Handlers can use
26/// [`RequestContext::mrtr`](crate::RequestContext::mrtr) or the convenience
27/// accessors on the context.
28#[derive(Debug, Clone, Default)]
29pub struct MrtrRequest {
30    input_responses: Option<InputResponses>,
31    request_state: Option<String>,
32}
33
34impl MrtrRequest {
35    pub(crate) fn new(
36        input_responses: Option<InputResponses>,
37        request_state: Option<String>,
38    ) -> Self {
39        Self {
40            input_responses,
41            request_state,
42        }
43    }
44
45    /// Client responses keyed by the identifiers from the prior
46    /// `inputRequests` map.
47    pub fn input_responses(&self) -> Option<&InputResponses> {
48        self.input_responses.as_ref()
49    }
50
51    /// Opaque continuation token echoed by the client.
52    pub fn request_state(&self) -> Option<&str> {
53        self.request_state.as_deref()
54    }
55
56    /// Consume the continuation values.
57    pub fn into_parts(self) -> (Option<InputResponses>, Option<String>) {
58        (self.input_responses, self.request_state)
59    }
60}
61
62/// Errors produced while encoding or validating opaque MRTR request state.
63#[derive(Debug, thiserror::Error)]
64#[non_exhaustive]
65pub enum RequestStateError {
66    /// HMAC keys shorter than 256 bits do not meet this codec's minimum.
67    #[error("request-state key must be at least 32 bytes")]
68    WeakKey,
69    /// The configured TTL must allow the state to live for some amount of time.
70    #[error("request-state TTL must be greater than zero")]
71    ZeroTtl,
72    /// The serialized state exceeded the configured token-size limit.
73    #[error("request-state token exceeds the configured maximum of {0} bytes")]
74    TooLarge(usize),
75    /// The token did not have the versioned three-part wire shape.
76    #[error("request-state token is malformed")]
77    Malformed,
78    /// The token uses a codec version this server does not understand.
79    #[error("unsupported request-state token version")]
80    UnsupportedVersion,
81    /// The HMAC did not match the payload.
82    #[error("request-state integrity verification failed")]
83    Integrity,
84    /// The token is no longer valid.
85    #[error("request-state token has expired")]
86    Expired,
87    /// A token bound to one authorization subject was used by another.
88    #[error("request-state token is not bound to the current subject")]
89    SubjectMismatch,
90    /// The state value could not be serialized.
91    #[error("failed to serialize request state: {0}")]
92    Encode(#[source] serde_json::Error),
93    /// The state value could not be decoded as the expected type.
94    #[error("failed to decode request state: {0}")]
95    Decode(#[source] serde_json::Error),
96    /// The system clock is earlier than the Unix epoch.
97    #[error("system clock is earlier than the Unix epoch")]
98    Clock,
99}
100
101#[derive(Debug, Serialize, serde::Deserialize)]
102struct StateEnvelope<T> {
103    issued_at: u64,
104    expires_at: u64,
105    #[serde(default, skip_serializing_if = "Option::is_none")]
106    subject: Option<String>,
107    state: T,
108}
109
110/// HMAC-SHA256 codec for opaque, expiring MRTR `requestState` values.
111///
112/// Construct the codec with the same key and TTL on every server instance
113/// that may receive a retry. Use [`encode_for`](Self::encode_for) and
114/// [`decode_for`](Self::decode_for) when an authenticated subject is
115/// available; subject binding prevents one user from replaying another user's
116/// continuation token. Binding is intentionally explicit: authentication
117/// middleware may place any application-defined principal type in request
118/// extensions, so the transport-neutral codec cannot safely infer one.
119#[derive(Clone)]
120pub struct RequestStateCodec {
121    key: std::sync::Arc<[u8]>,
122    ttl: Duration,
123    max_token_bytes: usize,
124}
125
126impl std::fmt::Debug for RequestStateCodec {
127    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
128        f.debug_struct("RequestStateCodec")
129            .field("key", &"<redacted>")
130            .field("ttl", &self.ttl)
131            .field("max_token_bytes", &self.max_token_bytes)
132            .finish()
133    }
134}
135
136impl RequestStateCodec {
137    /// Create a codec from a shared key and token TTL.
138    ///
139    /// The key must contain at least 32 bytes of entropy. Configuration
140    /// secrets should be decoded to raw bytes before calling this constructor.
141    pub fn new(key: impl AsRef<[u8]>, ttl: Duration) -> Result<Self, RequestStateError> {
142        let key = key.as_ref();
143        if key.len() < 32 {
144            return Err(RequestStateError::WeakKey);
145        }
146        if ttl.is_zero() {
147            return Err(RequestStateError::ZeroTtl);
148        }
149        Ok(Self {
150            key: std::sync::Arc::from(key),
151            ttl,
152            max_token_bytes: DEFAULT_MAX_TOKEN_BYTES,
153        })
154    }
155
156    /// Set the maximum accepted and emitted token size.
157    pub fn with_max_token_bytes(mut self, max_token_bytes: usize) -> Self {
158        self.max_token_bytes = max_token_bytes;
159        self
160    }
161
162    /// Encode state without authorization-subject binding.
163    pub fn encode<T: Serialize>(&self, state: &T) -> Result<String, RequestStateError> {
164        self.encode_at(None, state, unix_seconds()?)
165    }
166
167    /// Encode state bound to an authenticated subject identifier.
168    pub fn encode_for<T: Serialize>(
169        &self,
170        subject: impl Into<String>,
171        state: &T,
172    ) -> Result<String, RequestStateError> {
173        self.encode_at(Some(subject.into()), state, unix_seconds()?)
174    }
175
176    /// Verify and decode state that was not subject-bound.
177    pub fn decode<T: DeserializeOwned>(&self, token: &str) -> Result<T, RequestStateError> {
178        self.decode_at(token, None, unix_seconds()?)
179    }
180
181    /// Verify and decode state for the current authenticated subject.
182    pub fn decode_for<T: DeserializeOwned>(
183        &self,
184        token: &str,
185        subject: &str,
186    ) -> Result<T, RequestStateError> {
187        self.decode_at(token, Some(subject), unix_seconds()?)
188    }
189
190    fn encode_at<T: Serialize>(
191        &self,
192        subject: Option<String>,
193        state: &T,
194        now: u64,
195    ) -> Result<String, RequestStateError> {
196        let ttl = self.ttl.as_secs();
197        let envelope = StateEnvelope {
198            issued_at: now,
199            expires_at: now.saturating_add(ttl),
200            subject,
201            state,
202        };
203        let payload = serde_json::to_vec(&envelope).map_err(RequestStateError::Encode)?;
204        let payload = URL_SAFE_NO_PAD.encode(payload);
205        let signed = format!("{TOKEN_VERSION}.{payload}");
206        let signature = URL_SAFE_NO_PAD.encode(hmac_sha256(&self.key, signed.as_bytes()));
207        let token = format!("{signed}.{signature}");
208        if token.len() > self.max_token_bytes {
209            return Err(RequestStateError::TooLarge(self.max_token_bytes));
210        }
211        Ok(token)
212    }
213
214    fn decode_at<T: DeserializeOwned>(
215        &self,
216        token: &str,
217        subject: Option<&str>,
218        now: u64,
219    ) -> Result<T, RequestStateError> {
220        if token.len() > self.max_token_bytes {
221            return Err(RequestStateError::TooLarge(self.max_token_bytes));
222        }
223        let mut parts = token.split('.');
224        let version = parts.next().ok_or(RequestStateError::Malformed)?;
225        let payload = parts.next().ok_or(RequestStateError::Malformed)?;
226        let signature = parts.next().ok_or(RequestStateError::Malformed)?;
227        if parts.next().is_some() {
228            return Err(RequestStateError::Malformed);
229        }
230        if version != TOKEN_VERSION {
231            return Err(RequestStateError::UnsupportedVersion);
232        }
233
234        let supplied_signature = URL_SAFE_NO_PAD
235            .decode(signature)
236            .map_err(|_| RequestStateError::Malformed)?;
237        let signed = format!("{version}.{payload}");
238        let expected_signature = hmac_sha256(&self.key, signed.as_bytes());
239        if !constant_time_eq(&supplied_signature, &expected_signature) {
240            return Err(RequestStateError::Integrity);
241        }
242
243        let payload = URL_SAFE_NO_PAD
244            .decode(payload)
245            .map_err(|_| RequestStateError::Malformed)?;
246        let envelope: StateEnvelope<T> =
247            serde_json::from_slice(&payload).map_err(RequestStateError::Decode)?;
248        if now > envelope.expires_at {
249            return Err(RequestStateError::Expired);
250        }
251        match (envelope.subject.as_deref(), subject) {
252            (None, None) => {}
253            (Some(expected), Some(actual)) if expected == actual => {}
254            _ => return Err(RequestStateError::SubjectMismatch),
255        }
256        Ok(envelope.state)
257    }
258}
259
260fn unix_seconds() -> Result<u64, RequestStateError> {
261    SystemTime::now()
262        .duration_since(UNIX_EPOCH)
263        .map(|duration| duration.as_secs())
264        .map_err(|_| RequestStateError::Clock)
265}
266
267fn hmac_sha256(key: &[u8], message: &[u8]) -> [u8; 32] {
268    let mut normalized = [0u8; SHA256_BLOCK_SIZE];
269    if key.len() > SHA256_BLOCK_SIZE {
270        normalized[..32].copy_from_slice(&Sha256::digest(key));
271    } else {
272        normalized[..key.len()].copy_from_slice(key);
273    }
274
275    let mut inner_pad = [0x36u8; SHA256_BLOCK_SIZE];
276    let mut outer_pad = [0x5cu8; SHA256_BLOCK_SIZE];
277    for ((inner, outer), key_byte) in inner_pad
278        .iter_mut()
279        .zip(outer_pad.iter_mut())
280        .zip(normalized)
281    {
282        *inner ^= key_byte;
283        *outer ^= key_byte;
284    }
285
286    let mut inner = Sha256::new();
287    inner.update(inner_pad);
288    inner.update(message);
289    let inner = inner.finalize();
290
291    let mut outer = Sha256::new();
292    outer.update(outer_pad);
293    outer.update(inner);
294    outer.finalize().into()
295}
296
297fn constant_time_eq(left: &[u8], right: &[u8]) -> bool {
298    if left.len() != right.len() {
299        return false;
300    }
301    left.iter()
302        .zip(right)
303        .fold(0u8, |difference, (left, right)| difference | (left ^ right))
304        == 0
305}
306
307#[cfg(test)]
308mod tests {
309    use super::*;
310
311    const KEY: &[u8; 32] = b"0123456789abcdef0123456789abcdef";
312
313    #[derive(Debug, PartialEq, serde::Serialize, serde::Deserialize)]
314    struct State {
315        round: u8,
316        value: String,
317    }
318
319    #[test]
320    fn round_trips_shared_state() {
321        let first = RequestStateCodec::new(KEY, Duration::from_secs(60)).unwrap();
322        let second = RequestStateCodec::new(KEY, Duration::from_secs(60)).unwrap();
323        let state = State {
324            round: 2,
325            value: "kept".into(),
326        };
327        let token = first.encode_at(None, &state, 100).unwrap();
328        assert_eq!(second.decode_at::<State>(&token, None, 120).unwrap(), state);
329    }
330
331    #[test]
332    fn rejects_tampering_expiry_and_wrong_subject() {
333        let codec = RequestStateCodec::new(KEY, Duration::from_secs(10)).unwrap();
334        let token = codec
335            .encode_at(
336                Some("alice".into()),
337                &State {
338                    round: 1,
339                    value: "x".into(),
340                },
341                100,
342            )
343            .unwrap();
344
345        assert!(matches!(
346            codec.decode_at::<State>(&format!("{token}x"), Some("alice"), 101),
347            Err(RequestStateError::Integrity | RequestStateError::Malformed)
348        ));
349        assert!(matches!(
350            codec.decode_at::<State>(&token, Some("bob"), 101),
351            Err(RequestStateError::SubjectMismatch)
352        ));
353        assert!(matches!(
354            codec.decode_at::<State>(&token, Some("alice"), 111),
355            Err(RequestStateError::Expired)
356        ));
357    }
358
359    #[test]
360    fn enforces_key_ttl_and_size_limits() {
361        assert!(matches!(
362            RequestStateCodec::new(b"short", Duration::from_secs(1)),
363            Err(RequestStateError::WeakKey)
364        ));
365        assert!(matches!(
366            RequestStateCodec::new(KEY, Duration::ZERO),
367            Err(RequestStateError::ZeroTtl)
368        ));
369        let codec = RequestStateCodec::new(KEY, Duration::from_secs(1))
370            .unwrap()
371            .with_max_token_bytes(8);
372        assert!(matches!(
373            codec.encode(&"too large"),
374            Err(RequestStateError::TooLarge(8))
375        ));
376    }
377}
tower_mcp/mrtr.rs

tower_mcp/
mrtr.rs
