cookie_agent_server/service/
websocket.rs

1use std::{
2    net::{IpAddr, Ipv4Addr, SocketAddr},
3    sync::Arc,
4};
5
6use async_trait::async_trait;
7use axum::{
8    Router,
9    extract::{ConnectInfo, State, WebSocketUpgrade},
10    http::{HeaderMap, StatusCode, header},
11    response::{IntoResponse, Response},
12    routing::get,
13};
14use futures_util::StreamExt;
15use subtle::ConstantTimeEq;
16use tokio::net::TcpListener;
17use zeroize::Zeroizing;
18
19use super::{RunningServer, Server, ServerError};
20use crate::{MessageFrame, MessageStream, TransportError};
21
22struct WebSocketStream {
23    socket: axum::extract::ws::WebSocket,
24}
25
26#[async_trait]
27impl MessageStream for WebSocketStream {
28    async fn send(&mut self, frame: MessageFrame) -> Result<(), TransportError> {
29        let text = match frame {
30            MessageFrame::Text(text) => text,
31            MessageFrame::Value(value) => {
32                serde_json::to_string(&value).map_err(|_| TransportError::Closed)?
33            }
34        };
35        self.socket
36            .send(axum::extract::ws::Message::Text(text.into()))
37            .await
38            .map_err(|error| TransportError::Other(error.to_string()))
39    }
40
41    async fn recv(&mut self) -> Result<Option<MessageFrame>, TransportError> {
42        loop {
43            match self.socket.next().await {
44                Some(Ok(axum::extract::ws::Message::Text(text))) => {
45                    return Ok(Some(MessageFrame::Text(text.to_string())));
46                }
47                Some(Ok(axum::extract::ws::Message::Close(_))) | None => return Ok(None),
48                Some(Ok(_)) => {}
49                Some(Err(error)) => return Err(TransportError::Other(error.to_string())),
50            }
51        }
52    }
53}
54
55impl Server {
56    pub fn router(self: Arc<Self>, token: Arc<Zeroizing<String>>) -> Router {
57        Router::new()
58            .route("/ws", get(websocket_upgrade))
59            .with_state(WebSocketState {
60                server: self,
61                token,
62            })
63    }
64
65    pub async fn serve(
66        self: Arc<Self>,
67        port: u16,
68        token: Zeroizing<String>,
69    ) -> Result<RunningServer, ServerError> {
70        let router = self.clone().router(Arc::new(token));
71        let listener = TcpListener::bind(SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), port))
72            .await
73            .map_err(ServerError::Listen)?;
74        let address = listener.local_addr().map_err(ServerError::Listen)?;
75        let shutdown = self.shutdown.clone();
76        let task = tokio::spawn(async move {
77            let _ = axum::serve(
78                listener,
79                router.into_make_service_with_connect_info::<SocketAddr>(),
80            )
81            .with_graceful_shutdown(shutdown.cancelled_owned())
82            .await;
83        });
84        Ok(RunningServer { address, task })
85    }
86}
87
88#[derive(Clone)]
89struct WebSocketState {
90    server: Arc<Server>,
91    token: Arc<Zeroizing<String>>,
92}
93
94async fn websocket_upgrade(
95    State(state): State<WebSocketState>,
96    ConnectInfo(peer): ConnectInfo<SocketAddr>,
97    headers: HeaderMap,
98    upgrade: WebSocketUpgrade,
99) -> Response {
100    if !peer.ip().is_loopback() || headers.contains_key(header::ORIGIN) {
101        return StatusCode::FORBIDDEN.into_response();
102    }
103    if !authorized(&headers, state.token.as_str()) {
104        return StatusCode::UNAUTHORIZED.into_response();
105    }
106    let server = state.server;
107    upgrade.on_upgrade(move |socket| async move {
108        let _ = server.serve_stream(WebSocketStream { socket }).await;
109    })
110}
111
112pub(crate) fn authorized(headers: &HeaderMap, expected: &str) -> bool {
113    let mut values = headers.get_all(header::AUTHORIZATION).iter();
114    let Some(value) = values.next() else {
115        return false;
116    };
117    if values.next().is_some() {
118        return false;
119    }
120    let Ok(value) = value.to_str() else {
121        return false;
122    };
123    let Some(token) = value.strip_prefix("Bearer ") else {
124        return false;
125    };
126    token.len() == expected.len() && bool::from(token.as_bytes().ct_eq(expected.as_bytes()))
127}
128
129#[cfg(test)]
130mod tests {
131    use axum::http::{HeaderMap, HeaderValue, header};
132
133    use super::authorized;
134
135    fn bearer(value: &str) -> HeaderMap {
136        let mut headers = HeaderMap::new();
137        headers.insert(
138            header::AUTHORIZATION,
139            HeaderValue::from_str(value).expect("authorization header"),
140        );
141        headers
142    }
143
144    #[test]
145    fn in_memory_token_is_validated_without_any_file_read() {
146        let token = "A".repeat(43);
147        assert!(authorized(&bearer(&format!("Bearer {token}")), &token));
148        assert!(!authorized(&bearer("Bearer wrong"), &token));
149        assert!(!authorized(
150            &bearer(&format!("Bearer {}", "A".repeat(42))),
151            &token
152        ));
153        assert!(!authorized(&bearer("Basic abc"), &token));
154        assert!(!authorized(&HeaderMap::new(), &token));
155    }
156}
cookie_agent_server/service/websocket.rs

cookie_agent_server/service/
websocket.rs
