cookie_agent_server/
transport.rs

1use async_trait::async_trait;
2use cookie_agent_protocol::{MessageFrame, Transport, TransportError};
3use futures_util::{SinkExt as _, StreamExt as _};
4use tokio::sync::mpsc;
5use tokio_tungstenite::{
6    MaybeTlsStream, WebSocketStream,
7    tungstenite::{Message, client::IntoClientRequest as _},
8};
9use zeroize::Zeroizing;
10
11use crate::validate_websocket_url;
12
13pub struct InProcessStream {
14    sender: mpsc::Sender<MessageFrame>,
15    receiver: mpsc::Receiver<MessageFrame>,
16}
17
18#[must_use]
19pub fn in_process_pair(capacity: usize) -> (InProcessStream, InProcessStream) {
20    let (client_to_server_tx, client_to_server_rx) = mpsc::channel(capacity);
21    let (server_to_client_tx, server_to_client_rx) = mpsc::channel(capacity);
22    (
23        InProcessStream {
24            sender: client_to_server_tx,
25            receiver: server_to_client_rx,
26        },
27        InProcessStream {
28            sender: server_to_client_tx,
29            receiver: client_to_server_rx,
30        },
31    )
32}
33
34#[async_trait]
35impl Transport for InProcessStream {
36    async fn send(&mut self, frame: MessageFrame) -> Result<(), TransportError> {
37        self.sender
38            .send(frame)
39            .await
40            .map_err(|_| TransportError::Closed)
41    }
42
43    async fn recv(&mut self) -> Result<Option<MessageFrame>, TransportError> {
44        Ok(self.receiver.recv().await)
45    }
46}
47
48/// Tokio WebSocket transport for protocol clients.
49pub struct WebSocketTransport {
50    socket: WebSocketStream<MaybeTlsStream<tokio::net::TcpStream>>,
51}
52
53impl WebSocketTransport {
54    /// Connect to a validated daemon endpoint with an explicit per-run bearer token.
55    pub async fn connect_with_token(url: &str, token: &str) -> Result<Self, TransportError> {
56        validate_websocket_url(url).map_err(|error| TransportError::Other(error.to_string()))?;
57        let request = authenticated_request(url, token)?;
58        let (socket, _) = tokio_tungstenite::connect_async(request)
59            .await
60            .map_err(|error| websocket_error(&error))?;
61        Ok(Self { socket })
62    }
63}
64
65fn websocket_error(error: &tokio_tungstenite::tungstenite::Error) -> TransportError {
66    let mut message = error.to_string();
67    if let tokio_tungstenite::tungstenite::Error::Http(response) = error
68        && let Some(body) = response.body()
69    {
70        message.push_str(&format!(
71            "\nResponse body:\n{}",
72            cookie_agent_protocol::diagnostics::sanitize(&String::from_utf8_lossy(body), 4096)
73        ));
74    }
75    TransportError::Other(cookie_agent_protocol::diagnostics::sanitize(&message, 4096))
76}
77
78#[async_trait]
79impl Transport for WebSocketTransport {
80    async fn send(&mut self, frame: MessageFrame) -> Result<(), TransportError> {
81        let text = match frame {
82            MessageFrame::Text(text) => text,
83            MessageFrame::Value(value) => serde_json::to_string(&value)?,
84        };
85        self.socket
86            .send(Message::Text(text.into()))
87            .await
88            .map_err(|error| websocket_error(&error))
89    }
90
91    async fn recv(&mut self) -> Result<Option<MessageFrame>, TransportError> {
92        loop {
93            match self.socket.next().await {
94                Some(Ok(Message::Text(text))) => {
95                    return Ok(Some(MessageFrame::Text(text.to_string())));
96                }
97                Some(Ok(Message::Close(_))) | None => return Ok(None),
98                Some(Ok(_)) => {}
99                Some(Err(error)) => {
100                    return Err(websocket_error(&error));
101                }
102            }
103        }
104    }
105}
106
107pub(crate) fn authenticated_request(
108    url: &str,
109    token: &str,
110) -> Result<tokio_tungstenite::tungstenite::http::Request<()>, TransportError> {
111    if token.len() != crate::token::TOKEN_ENCODED_BYTES
112        || !token
113            .bytes()
114            .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'_'))
115    {
116        return Err(TransportError::Other(
117            "invalid daemon authentication token".into(),
118        ));
119    }
120    let mut request = url
121        .into_client_request()
122        .map_err(|error| TransportError::Other(error.to_string()))?;
123    let authorization = Zeroizing::new(format!("Bearer {token}"));
124    let value = authorization
125        .parse()
126        .map_err(|_| TransportError::Other("invalid authorization header".into()))?;
127    request.headers_mut().insert("authorization", value);
128    Ok(request)
129}
130
131#[cfg(test)]
132mod tests {
133    use super::authenticated_request;
134
135    #[test]
136    fn websocket_rejection_preserves_decoded_body_and_context() {
137        let response = tokio_tungstenite::tungstenite::http::Response::builder().status(403)
138            .body(Some(br#"{"message":"review\u002dsecret","password":"review\"secret-tail","reason":"Useful cause"}"#.to_vec())).unwrap();
139        let error = super::websocket_error(&tokio_tungstenite::tungstenite::Error::Http(Box::new(
140            response,
141        )))
142        .to_string();
143        assert!(error.contains("403"));
144        assert!(error.contains("Useful cause"));
145        assert!(error.contains("review-secret"));
146        assert!(error.contains("secret-tail"));
147    }
148
149    #[test]
150    fn websocket_rejection_shows_body_including_echoed_token() {
151        let response = tokio_tungstenite::tungstenite::http::Response::builder()
152            .status(403)
153            .header("set-cookie", "private-cookie")
154            .body(Some(
155                b"Gateway denied the connection; echoed opaque-known-value".to_vec(),
156            ))
157            .unwrap();
158        let error = super::websocket_error(&tokio_tungstenite::tungstenite::Error::Http(Box::new(
159            response,
160        )))
161        .to_string();
162        assert!(error.contains("403"));
163        assert!(error.contains("Gateway denied the connection"));
164        assert!(error.contains("opaque-known-value"));
165        assert!(!error.contains("private-cookie"));
166    }
167
168    #[test]
169    fn websocket_auth_uses_a_bearer_header_without_url_credentials() {
170        let token = "A".repeat(43);
171        let request =
172            authenticated_request("ws://127.0.0.1:7419/ws", &token).expect("authenticated request");
173        assert_eq!(request.uri().to_string(), "ws://127.0.0.1:7419/ws");
174        assert_eq!(
175            request
176                .headers()
177                .get("authorization")
178                .and_then(|value| value.to_str().ok()),
179            Some(format!("Bearer {token}").as_str())
180        );
181        assert!(!request.uri().to_string().contains(&token));
182        assert!(authenticated_request("ws://127.0.0.1:7419/ws", "sentinel-secret").is_err());
183    }
184}
cookie_agent_server/transport.rs

cookie_agent_server/
transport.rs
