lattice_ai/mcp/
transport.rs

1//! WebSocket accept + authorization handshake.
2//!
3//! Upgrades an accepted TCP stream to a WebSocket, gating on the
4//! `x-claude-code-ide-authorization` header. A mismatched token is
5//! rejected during the handshake with HTTP 401 — the connection never
6//! reaches the MCP loop. Loopback bind + this token are the security
7//! boundary (design §4).
8
9use tokio::net::TcpStream;
10use tokio_tungstenite::WebSocketStream;
11use tokio_tungstenite::tungstenite::handshake::server::{ErrorResponse, Request, Response};
12use tokio_tungstenite::tungstenite::http::StatusCode;
13
14use crate::mcp::auth;
15use crate::mcp::error::Result;
16
17/// Accept a WebSocket connection on `stream`, requiring the
18/// `x-claude-code-ide-authorization` header to match `expected_token`
19/// (constant-time). Rejects with HTTP 401 otherwise — the returned future
20/// resolves to an error and the connection is dropped.
21// The handshake callback's `Result<Response, ErrorResponse>` return type is
22// dictated by tokio-tungstenite's `Callback` trait; `ErrorResponse`
23// (`http::Response<Option<String>>`) is inherently large, so we can't shrink
24// the Err variant here (clippy `result_large_err`).
25#[allow(clippy::result_large_err)]
26pub async fn accept(stream: TcpStream, expected_token: &str) -> Result<WebSocketStream<TcpStream>> {
27    let expected = expected_token.to_string();
28    let callback = move |request: &Request,
29                         response: Response|
30          -> std::result::Result<Response, ErrorResponse> {
31        let provided = request
32            .headers()
33            .get(auth::AUTH_HEADER)
34            .and_then(|v| v.to_str().ok())
35            .unwrap_or("");
36        if auth::header_matches(&expected, provided) {
37            Ok(response)
38        } else {
39            let mut err = ErrorResponse::new(Some("authorization rejected".to_string()));
40            *err.status_mut() = StatusCode::UNAUTHORIZED;
41            Err(err)
42        }
43    };
44    let ws = tokio_tungstenite::accept_hdr_async(stream, callback).await?;
45    Ok(ws)
46}
lattice_ai/mcp/transport.rs

lattice_ai/mcp/
transport.rs
