Skip to main content

shell_tunnel/relay/
mod.rs

1//! Self-hosted relay: reaching a device that dialled out to you.
2//!
3//! The relay is the alternative to a third-party tunnel. A device opens one
4//! outbound WebSocket to it — no inbound port, no NAT configuration — and the
5//! relay routes public traffic back down that connection.
6//!
7//! It runs from the same binary (`shell-tunnel relay`), so an operator never
8//! has to match versions between two programs.
9//!
10//! What the relay deliberately does *not* do: interpret capability tokens.
11//! Enrollment decides which devices may attach; the capability token in each
12//! proxied request stays end-to-end between client and device. The relay is a
13//! router, not a second security boundary.
14
15#[cfg(feature = "relay-client")]
16pub mod client;
17pub mod protocol;
18pub mod proxy;
19pub mod registry;
20
21use std::net::SocketAddr;
22use std::sync::Arc;
23use std::time::Duration;
24
25use axum::{
26    body::Bytes,
27    extract::{
28        ws::{Message, WebSocket, WebSocketUpgrade},
29        FromRequestParts, Request, State,
30    },
31    http::{HeaderMap, StatusCode},
32    response::{IntoResponse, Response},
33    routing::{any, get},
34    Router,
35};
36use futures_util::{SinkExt, StreamExt};
37
38use crate::error::ShellTunnelError;
39use crate::security::{generate_api_key, rate_limit_middleware, RateLimitConfig, RateLimiter};
40use protocol::{reject, DeviceMessage, RelayMessage, PROTOCOL_VERSION};
41use proxy::{
42    is_forwardable, split_device_path, ProxyRequest, ProxyResponse, POOL_WAIT, REQUEST_TIMEOUT,
43};
44use registry::{Device, DeviceRegistry};
45
46pub use registry::{DeviceRegistry as Registry, POOL_TARGET};
47
48/// How long a device may go without a heartbeat before it is considered gone.
49pub const HEARTBEAT_TIMEOUT: Duration = Duration::from_secs(90);
50
51/// How long to wait for the enrollment frame before dropping a connection.
52const ENROLL_TIMEOUT: Duration = Duration::from_secs(10);
53
54/// Relay server settings.
55#[derive(Debug, Clone)]
56pub struct RelayConfig {
57    /// Address to listen on.
58    pub bind: SocketAddr,
59    /// Secret a device must present to attach.
60    pub enroll_token: String,
61    /// Per-IP request limiting for this relay.
62    ///
63    /// The relay is the only place this can work for proxied traffic: a device
64    /// replays requests to its own loopback listener, so *its* limiter sees
65    /// 127.0.0.1 for every caller and cannot tell them apart. Here the real
66    /// client address is still visible.
67    pub rate_limit: RateLimitConfig,
68    /// Serve HTTPS directly instead of relying on a reverse proxy.
69    #[cfg(feature = "tls")]
70    pub tls: Option<crate::tls::TlsFiles>,
71    /// Public base URL of this relay, when the operator states it explicitly.
72    ///
73    /// Left unset, the relay derives it from each connection's `Host` (and
74    /// `X-Forwarded-*`) headers, so a relay behind TLS termination still tells
75    /// devices an address that actually works.
76    pub public_base: Option<String>,
77}
78
79impl RelayConfig {
80    /// Create a configuration with the given bind address and token.
81    pub fn new(bind: SocketAddr, enroll_token: impl Into<String>) -> Self {
82        Self {
83            bind,
84            enroll_token: enroll_token.into(),
85            rate_limit: RateLimitConfig::default(),
86            #[cfg(feature = "tls")]
87            tls: None,
88            public_base: None,
89        }
90    }
91
92    /// Terminate TLS in-process using these files.
93    #[cfg(feature = "tls")]
94    pub fn with_tls(mut self, files: crate::tls::TlsFiles) -> Self {
95        self.tls = Some(files);
96        self
97    }
98
99    /// Turn per-IP request limiting off.
100    pub fn without_rate_limit(mut self) -> Self {
101        self.rate_limit.enabled = false;
102        self
103    }
104
105    /// Set the public base URL advertised to devices.
106    pub fn with_public_base(mut self, base: impl Into<String>) -> Self {
107        self.public_base = Some(base.into().trim_end_matches('/').to_string());
108        self
109    }
110
111    /// The base URL to advertise, preferring what the operator configured.
112    ///
113    /// `observed` is what the connection itself says this relay is reachable at.
114    /// Falling back to the bind address is a last resort — it is right only when
115    /// nothing is in front of the relay.
116    pub fn public_base_or(&self, observed: Option<String>) -> String {
117        self.public_base
118            .clone()
119            .or(observed)
120            .unwrap_or_else(|| format!("http://{}", self.bind))
121    }
122
123    /// Public URL that routes to `device_id`.
124    pub fn public_url_for(&self, device_id: &str, observed: Option<String>) -> String {
125        format!("{}/d/{}", self.public_base_or(observed), device_id)
126    }
127}
128
129/// Shared relay state.
130#[derive(Debug, Clone)]
131pub struct RelayState {
132    config: RelayConfig,
133    devices: DeviceRegistry,
134}
135
136impl RelayState {
137    /// Create state for `config`.
138    pub fn new(config: RelayConfig) -> Self {
139        Self {
140            config,
141            devices: DeviceRegistry::new(),
142        }
143    }
144
145    /// The device registry.
146    pub fn devices(&self) -> &DeviceRegistry {
147        &self.devices
148    }
149}
150
151/// Build the relay router.
152///
153/// Every route but `/health` is rate limited per client IP. Enrolment is the
154/// reason: without a limit, a weak enrolment token can be guessed at line speed,
155/// and the relay is the only place that sees who is asking.
156pub fn relay_router(state: RelayState) -> Router {
157    let limiter = Arc::new(RateLimiter::new(state.config.rate_limit.clone()));
158
159    Router::new()
160        .route("/health", get(|| async { "OK" }))
161        .route("/relay/v1/control", get(control_handler))
162        .route("/relay/v1/data", get(data_handler))
163        .route("/relay/v1/devices", get(devices_handler))
164        .route("/d/{*rest}", any(proxy_handler))
165        .layer(axum::middleware::from_fn_with_state(
166            limiter,
167            rate_limit_middleware,
168        ))
169        .with_state(state)
170}
171
172/// Run the relay server until shutdown.
173pub async fn serve_relay(config: RelayConfig) -> crate::Result<()> {
174    let bind = config.bind;
175    #[cfg(feature = "tls")]
176    let tls = config.tls.clone();
177    let state = RelayState::new(config);
178    let router = relay_router(state.clone());
179
180    // A device that vanished without closing its socket looks identical to an
181    // idle one, so entries are reaped on heartbeat staleness instead.
182    let sweeper = state.devices().clone();
183    tokio::spawn(async move {
184        let mut ticker = tokio::time::interval(HEARTBEAT_TIMEOUT / 3);
185        loop {
186            ticker.tick().await;
187            for id in sweeper.evict_stale(HEARTBEAT_TIMEOUT) {
188                tracing::info!(target: "relay", device_id = %id, "device evicted (no heartbeat)");
189            }
190        }
191    });
192
193    tracing::info!("relay listening on {}", bind);
194
195    let listener = tokio::net::TcpListener::bind(bind)
196        .await
197        .map_err(ShellTunnelError::Io)?;
198    // Connection info is what the rate limiter keys on; without it every caller
199    // would look identical.
200    let service = router.into_make_service_with_connect_info::<SocketAddr>();
201
202    #[cfg(feature = "tls")]
203    if let Some(files) = tls {
204        // Loaded before serving so a bad certificate stops startup rather than
205        // failing every connection at handshake time.
206        let config = crate::tls::acceptor(files.load()?);
207        // Renewal should not require a restart.
208        crate::tls::watch(files, config.clone());
209        let std_listener = listener.into_std().map_err(ShellTunnelError::Io)?;
210        return axum_server::from_tcp_rustls(std_listener, config)
211            .map_err(ShellTunnelError::Io)?
212            .serve(service)
213            .await
214            .map_err(|e| ShellTunnelError::Io(std::io::Error::other(e.to_string())));
215    }
216
217    axum::serve(listener, service)
218        .await
219        .map_err(|e| ShellTunnelError::Io(std::io::Error::other(e.to_string())))?;
220    Ok(())
221}
222
223/// Whether `name` is usable as a routing key in `/d/<name>/…`.
224///
225/// Deliberately narrow: the name lands in a URL path, so anything that could
226/// need escaping, traverse a path, or collide with the relay's own routes is
227/// rejected rather than sanitized.
228fn is_valid_device_name(name: &str) -> bool {
229    !name.is_empty()
230        && name.len() <= 64
231        && name
232            .chars()
233            .all(|c| c.is_ascii_alphanumeric() || c == '-' || c == '_')
234}
235
236/// Work out how this relay was addressed, from the connection's own headers.
237///
238/// A relay behind TLS termination sees plain HTTP on a loopback port, so the
239/// scheme and host it should advertise are only knowable from what the proxy
240/// forwards.
241fn observed_base(headers: &HeaderMap, tls: bool) -> Option<String> {
242    let host = headers
243        .get("x-forwarded-host")
244        .or_else(|| headers.get(axum::http::header::HOST))
245        .and_then(|value| value.to_str().ok())?;
246    if host.is_empty() {
247        return None;
248    }
249    // A proxy's own statement wins; failing that, the relay knows whether it
250    // terminated TLS itself. Guessing `http` while serving HTTPS would advertise
251    // a URL that the relay itself refuses.
252    let scheme = headers
253        .get("x-forwarded-proto")
254        .and_then(|value| value.to_str().ok())
255        .map(|proto| proto.split(',').next().unwrap_or(proto).trim().to_string())
256        .unwrap_or_else(|| if tls { "https" } else { "http" }.to_string());
257    Some(format!("{scheme}://{host}"))
258}
259
260/// Whether this relay terminates TLS itself.
261fn serves_tls(_state: &RelayState) -> bool {
262    #[cfg(feature = "tls")]
263    {
264        _state.config.tls.is_some()
265    }
266    #[cfg(not(feature = "tls"))]
267    {
268        false
269    }
270}
271
272/// Upgrade a device's outbound connection into the control channel.
273async fn control_handler(
274    ws: WebSocketUpgrade,
275    State(state): State<RelayState>,
276    headers: HeaderMap,
277) -> impl IntoResponse {
278    let observed = observed_base(&headers, serves_tls(&state));
279    ws.on_upgrade(move |socket| control_session(socket, state, observed))
280}
281
282/// Enroll a device, then serve its heartbeats until the connection ends.
283async fn control_session(socket: WebSocket, state: RelayState, observed: Option<String>) {
284    let (mut sink, mut stream) = socket.split();
285
286    // An unauthenticated peer must not be able to hold a connection open
287    // indefinitely, so enrollment is bounded in time.
288    let first = match tokio::time::timeout(ENROLL_TIMEOUT, stream.next()).await {
289        Ok(Some(Ok(Message::Text(text)))) => text,
290        _ => return,
291    };
292
293    let enroll = match serde_json::from_str::<DeviceMessage>(&first) {
294        Ok(DeviceMessage::Enroll {
295            enroll_token,
296            version,
297            label,
298            device_name,
299        }) => (enroll_token, version, label, device_name),
300        _ => {
301            reject_and_close(
302                &mut sink,
303                reject::BAD_HANDSHAKE,
304                "expected an enroll message",
305            )
306            .await;
307            return;
308        }
309    };
310    let (enroll_token, version, label, device_name) = enroll;
311
312    if version != PROTOCOL_VERSION {
313        reject_and_close(
314            &mut sink,
315            reject::UNSUPPORTED_VERSION,
316            &format!("relay speaks protocol version {PROTOCOL_VERSION}"),
317        )
318        .await;
319        return;
320    }
321
322    if !constant_time_eq(&enroll_token, &state.config.enroll_token) {
323        // No detail about *why*: a device that guessed wrong learns nothing.
324        tracing::debug!(target: "relay", "enrollment rejected: bad token");
325        reject_and_close(&mut sink, reject::BAD_TOKEN, "enrollment refused").await;
326        return;
327    }
328
329    // A named device keeps one URL across reconnects, which is what makes the
330    // relay usable when whoever calls the device cannot read its console. An
331    // unnamed one gets a random id, which nobody can guess but which changes
332    // every time it attaches.
333    let device_id = match device_name {
334        Some(name) if !is_valid_device_name(&name) => {
335            reject_and_close(
336                &mut sink,
337                reject::BAD_DEVICE_NAME,
338                "device names may use letters, digits, '-' and '_' (1-64 characters)",
339            )
340            .await;
341            return;
342        }
343        // Re-attaching under an existing name replaces the old entry rather
344        // than being refused: after a network drop the relay still holds a
345        // connection it cannot know is dead, and refusing would lock the device
346        // out until the heartbeat timeout expired. Only holders of the enrol
347        // token can do this, which is the same trust level as attaching at all.
348        Some(name) => name,
349        None => generate_api_key(),
350    };
351    let public_url = state.config.public_url_for(&device_id, observed);
352    let registry::DeviceHandles {
353        device,
354        mut refill_rx,
355    } = state.devices.attach(&device_id, label.clone());
356    tracing::info!(
357        target: "relay",
358        device_id = %device_id,
359        label = label.as_deref().unwrap_or("-"),
360        "device attached"
361    );
362
363    let enrolled = RelayMessage::Enrolled {
364        device_id: device_id.clone(),
365        public_url,
366    };
367    if send_json(&mut sink, &enrolled).await.is_err() {
368        state.devices.detach(&device_id);
369        return;
370    }
371
372    // Fill the pool up front so the first request does not pay for a handshake.
373    let fill = RelayMessage::OpenData {
374        count: registry::POOL_TARGET,
375    };
376    if send_json(&mut sink, &fill).await.is_err() {
377        state.devices.detach(&device_id);
378        return;
379    }
380
381    // The control channel multiplexes nothing but coordination: device
382    // heartbeats one way, pool-refill requests the other.
383    loop {
384        tokio::select! {
385            incoming = stream.next() => {
386                let Some(Ok(message)) = incoming else { break };
387                match message {
388                    Message::Text(text) => match serde_json::from_str::<DeviceMessage>(&text) {
389                        Ok(DeviceMessage::Heartbeat) => {
390                            device.touch();
391                            if send_json(&mut sink, &RelayMessage::HeartbeatAck).await.is_err() {
392                                break;
393                            }
394                        }
395                        // A second enrollment on an attached connection is a
396                        // protocol error, not a re-key: ignore it rather than
397                        // reassigning an id.
398                        _ => continue,
399                    },
400                    Message::Close(_) => break,
401                    _ => continue,
402                }
403            }
404            refill = refill_rx.recv() => {
405                if refill.is_none() {
406                    break;
407                }
408                if send_json(&mut sink, &RelayMessage::OpenData { count: 1 }).await.is_err() {
409                    break;
410                }
411            }
412        }
413    }
414
415    state.devices.detach(&device_id);
416    tracing::info!(target: "relay", device_id = %device_id, "device detached");
417}
418
419/// List the devices currently attached.
420///
421/// Authenticated with the enrolment token, because the answer is only useful to
422/// whoever operates this relay — and anyone holding that token could attach a
423/// device anyway, so listing them reveals nothing new.
424async fn devices_handler(State(state): State<RelayState>, headers: HeaderMap) -> Response {
425    let presented = headers
426        .get(axum::http::header::AUTHORIZATION)
427        .and_then(|value| value.to_str().ok())
428        .and_then(|value| value.strip_prefix("Bearer "))
429        .unwrap_or("");
430    if !constant_time_eq(presented, &state.config.enroll_token) {
431        return StatusCode::UNAUTHORIZED.into_response();
432    }
433
434    let base = state
435        .config
436        .public_base_or(observed_base(&headers, serves_tls(&state)));
437    let devices: Vec<_> = state
438        .devices
439        .list()
440        .into_iter()
441        .map(|device| {
442            let url = format!("{}/d/{}", base, device.id);
443            serde_json::json!({
444                "id": device.id,
445                "label": device.label,
446                "attached_secs": device.attached_secs,
447                "last_seen_secs": device.last_seen_secs,
448                "public_url": url,
449            })
450        })
451        .collect();
452
453    axum::Json(serde_json::json!({ "devices": devices })).into_response()
454}
455
456/// Accept a data connection and park it in its device's pool.
457///
458/// The connection authenticates itself in its first frame rather than in the
459/// URL: query strings land in the access logs of the reverse proxies this relay
460/// is meant to sit behind, so a token there would be written to disk in
461/// plaintext on exactly the deployments that follow our own TLS advice.
462async fn data_handler(ws: WebSocketUpgrade, State(state): State<RelayState>) -> Response {
463    ws.on_upgrade(move |socket| attach_data_connection(socket, state))
464}
465
466/// Read the attach frame, verify it, and hand the socket to the device's pool.
467async fn attach_data_connection(mut socket: WebSocket, state: RelayState) {
468    let first = tokio::time::timeout(ENROLL_TIMEOUT, socket.recv()).await;
469    let Ok(Some(Ok(Message::Text(text)))) = first else {
470        let _ = socket.close().await;
471        return;
472    };
473
474    let Ok(DeviceMessage::Attach {
475        device_id,
476        enroll_token,
477    }) = serde_json::from_str::<DeviceMessage>(&text)
478    else {
479        let _ = socket.close().await;
480        return;
481    };
482
483    if !constant_time_eq(&enroll_token, &state.config.enroll_token) {
484        tracing::debug!(target: "relay", "data connection rejected: bad token");
485        let _ = socket.close().await;
486        return;
487    }
488
489    let Some(device) = state.devices.get(&device_id) else {
490        let _ = socket.close().await;
491        return;
492    };
493
494    // A pool that is already full means the device over-supplied; closing the
495    // extra socket is better than holding it open forever.
496    if let Some(mut extra) = device.offer(socket).await {
497        let _ = extra.close().await;
498    }
499}
500
501/// Forward a public request to the addressed device and return its response.
502async fn proxy_handler(State(state): State<RelayState>, request: Request) -> Response {
503    let path_and_query = request
504        .uri()
505        .path_and_query()
506        .map(|p| p.as_str().to_string())
507        .unwrap_or_else(|| request.uri().path().to_string());
508
509    let Some((device_id, tail)) = split_device_path(&path_and_query) else {
510        return StatusCode::NOT_FOUND.into_response();
511    };
512
513    let Some(device) = state.devices.get(device_id) else {
514        // The device is not attached: this is the relay reporting a missing
515        // upstream, which is exactly what 502 means.
516        return (StatusCode::BAD_GATEWAY, "device is not connected").into_response();
517    };
518
519    let method = request.method().to_string();
520    let headers: Vec<(String, String)> = request
521        .headers()
522        .iter()
523        .filter(|(name, _)| is_forwardable(name.as_str()))
524        .filter_map(|(name, value)| {
525            value
526                .to_str()
527                .ok()
528                .map(|v| (name.as_str().to_string(), v.to_string()))
529        })
530        .collect();
531
532    // A WebSocket upgrade cannot be answered by buffering: the exchange has no
533    // end until one side closes. Because one request already owns one data
534    // connection for its lifetime, the same socket simply becomes the pipe —
535    // the connection-per-request model pays off here rather than needing a
536    // second mechanism.
537    if is_websocket_upgrade(request.headers()) {
538        let (mut parts, _) = request.into_parts();
539        let upgrade = match WebSocketUpgrade::from_request_parts(&mut parts, &state).await {
540            Ok(upgrade) => upgrade,
541            Err(rejection) => return rejection.into_response(),
542        };
543        let proxied = ProxyRequest {
544            method,
545            path: tail,
546            headers,
547            websocket: true,
548        };
549        return upgrade.on_upgrade(move |client| pipe_websocket(client, device, proxied));
550    }
551
552    let body = match axum::body::to_bytes(request.into_body(), MAX_BODY).await {
553        Ok(body) => body,
554        Err(_) => return StatusCode::PAYLOAD_TOO_LARGE.into_response(),
555    };
556
557    let Some(conn) = device.take(POOL_WAIT).await else {
558        // The device is attached but has no spare connection. 503 with a
559        // Retry-After is the honest answer: try again shortly.
560        return (
561            StatusCode::SERVICE_UNAVAILABLE,
562            [("retry-after", "1")],
563            "no data connection available",
564        )
565            .into_response();
566    };
567
568    match tokio::time::timeout(
569        REQUEST_TIMEOUT,
570        forward(
571            conn,
572            ProxyRequest {
573                method,
574                path: tail,
575                headers,
576                websocket: false,
577            },
578            body,
579        ),
580    )
581    .await
582    {
583        Ok(Ok(response)) => response,
584        Ok(Err(reason)) => {
585            tracing::debug!(target: "relay", device_id = %device.id, reason, "proxy failed");
586            (StatusCode::BAD_GATEWAY, "device did not answer").into_response()
587        }
588        Err(_) => (StatusCode::GATEWAY_TIMEOUT, "device timed out").into_response(),
589    }
590}
591
592/// Whether these headers ask to switch protocols to WebSocket.
593fn is_websocket_upgrade(headers: &HeaderMap) -> bool {
594    let header_contains = |name: axum::http::HeaderName, needle: &str| {
595        headers
596            .get(name)
597            .and_then(|value| value.to_str().ok())
598            .is_some_and(|value| value.to_ascii_lowercase().contains(needle))
599    };
600    header_contains(axum::http::header::UPGRADE, "websocket")
601        && header_contains(axum::http::header::CONNECTION, "upgrade")
602}
603
604/// Join a client's WebSocket to the device over one data connection.
605///
606/// The relay has already answered 101 by the time this runs — axum completes the
607/// handshake before invoking the callback — so a device that then refuses simply
608/// results in the client's socket closing.
609async fn pipe_websocket(mut client: WebSocket, device: Arc<Device>, request: ProxyRequest) {
610    let Some(mut conn) = device.take(POOL_WAIT).await else {
611        tracing::debug!(target: "relay", device_id = %device.id, "no data connection for websocket");
612        let _ = client.close().await;
613        return;
614    };
615
616    let Ok(header) = serde_json::to_string(&request) else {
617        let _ = client.close().await;
618        return;
619    };
620    if conn.send(Message::Text(header.into())).await.is_err() {
621        let _ = client.close().await;
622        return;
623    }
624
625    // The device answers with the status its own server returned; anything but
626    // a switch means the upgrade did not happen there.
627    let switched = matches!(
628        conn.recv().await,
629        Some(Ok(Message::Text(ref text)))
630            if serde_json::from_str::<ProxyResponse>(text)
631                .map(|response| response.status == 101)
632                .unwrap_or(false)
633    );
634    if !switched {
635        let _ = client.close().await;
636        let _ = conn.close().await;
637        return;
638    }
639
640    // From here the two sockets are the same conversation: copy frames until
641    // either end hangs up.
642    loop {
643        tokio::select! {
644            from_client = client.recv() => {
645                match from_client {
646                    Some(Ok(Message::Close(_))) | None | Some(Err(_)) => break,
647                    Some(Ok(message)) => {
648                        if conn.send(message).await.is_err() {
649                            break;
650                        }
651                    }
652                }
653            }
654            from_device = conn.recv() => {
655                match from_device {
656                    Some(Ok(Message::Close(_))) | None | Some(Err(_)) => break,
657                    Some(Ok(message)) => {
658                        if client.send(message).await.is_err() {
659                            break;
660                        }
661                    }
662                }
663            }
664        }
665    }
666
667    let _ = client.close().await;
668    let _ = conn.close().await;
669}
670
671/// Largest request body the relay will buffer before forwarding.
672const MAX_BODY: usize = 8 * 1024 * 1024;
673
674/// Drive one request/response exchange over a dedicated data connection.
675///
676/// Wire shape: request header (text) → request body (binary) → response header
677/// (text) → response body (binary frames) → close.
678async fn forward(
679    mut conn: WebSocket,
680    request: ProxyRequest,
681    body: Bytes,
682) -> Result<Response, &'static str> {
683    let header = serde_json::to_string(&request).map_err(|_| "request-encode")?;
684    conn.send(Message::Text(header.into()))
685        .await
686        .map_err(|_| "request-header-send")?;
687    conn.send(Message::Binary(body))
688        .await
689        .map_err(|_| "request-body-send")?;
690
691    let head: ProxyResponse = loop {
692        match conn.recv().await {
693            Some(Ok(Message::Text(text))) => {
694                break serde_json::from_str(&text).map_err(|_| "response-decode")?
695            }
696            Some(Ok(_)) => continue,
697            _ => return Err("response-header-missing"),
698        }
699    };
700
701    let mut body = Vec::new();
702    while let Some(Ok(message)) = conn.recv().await {
703        match message {
704            Message::Binary(chunk) => body.extend_from_slice(&chunk),
705            Message::Close(_) => break,
706            _ => continue,
707        }
708    }
709
710    let mut response = Response::builder().status(head.status);
711    for (name, value) in head.headers {
712        if is_forwardable(&name) {
713            response = response.header(name, value);
714        }
715    }
716    response
717        .body(axum::body::Body::from(body))
718        .map_err(|_| "response-build")
719}
720
721/// Send a rejection and close, best-effort.
722async fn reject_and_close<S>(sink: &mut S, code: &str, message: &str)
723where
724    S: SinkExt<Message> + Unpin,
725{
726    let rejected = RelayMessage::Rejected {
727        code: code.to_string(),
728        message: message.to_string(),
729    };
730    let _ = send_json(sink, &rejected).await;
731    let _ = sink.close().await;
732}
733
734/// Serialize and send one protocol message.
735async fn send_json<S, T>(sink: &mut S, message: &T) -> Result<(), ()>
736where
737    S: SinkExt<Message> + Unpin,
738    T: serde::Serialize,
739{
740    let json = serde_json::to_string(message).map_err(|_| ())?;
741    sink.send(Message::Text(json.into())).await.map_err(|_| ())
742}
743
744/// Compare secrets without leaking their contents through timing.
745///
746/// The token is short and comparisons are rare, but an early-exit `==` on a
747/// shared secret is the kind of detail that is cheap to get right and awkward
748/// to retrofit.
749fn constant_time_eq(a: &str, b: &str) -> bool {
750    let (a, b) = (a.as_bytes(), b.as_bytes());
751    if a.len() != b.len() {
752        return false;
753    }
754    a.iter().zip(b).fold(0u8, |acc, (x, y)| acc | (x ^ y)) == 0
755}
756
757#[cfg(test)]
758mod tests {
759    use super::*;
760
761    fn config() -> RelayConfig {
762        RelayConfig::new("127.0.0.1:0".parse().unwrap(), "secret")
763    }
764
765    #[test]
766    fn public_url_uses_the_device_path_prefix() {
767        let config = config().with_public_base("https://relay.example.com/");
768        assert_eq!(
769            config.public_url_for("dev-1", None),
770            "https://relay.example.com/d/dev-1"
771        );
772    }
773
774    #[test]
775    fn public_base_defaults_to_the_bind_address() {
776        let config = RelayConfig::new("127.0.0.1:8443".parse().unwrap(), "secret");
777        assert_eq!(
778            config.public_url_for("d", None),
779            "http://127.0.0.1:8443/d/d"
780        );
781    }
782
783    #[test]
784    fn an_observed_address_is_used_when_the_operator_configured_none() {
785        let config = config();
786        assert_eq!(
787            config.public_url_for("dev-1", Some("https://relay.example.com".into())),
788            "https://relay.example.com/d/dev-1"
789        );
790    }
791
792    #[test]
793    fn a_configured_base_wins_over_what_the_connection_observed() {
794        let config = config().with_public_base("https://canonical.example");
795        assert_eq!(
796            config.public_url_for("dev-1", Some("https://whatever.invalid".into())),
797            "https://canonical.example/d/dev-1"
798        );
799    }
800
801    #[test]
802    fn the_forwarded_scheme_and_host_are_preferred_over_the_direct_host() {
803        let mut headers = HeaderMap::new();
804        headers.insert(axum::http::header::HOST, "127.0.0.1:8443".parse().unwrap());
805        assert_eq!(
806            observed_base(&headers, false).as_deref(),
807            Some("http://127.0.0.1:8443")
808        );
809
810        headers.insert("x-forwarded-proto", "https".parse().unwrap());
811        headers.insert("x-forwarded-host", "relay.example.com".parse().unwrap());
812        assert_eq!(
813            observed_base(&headers, false).as_deref(),
814            Some("https://relay.example.com")
815        );
816    }
817
818    #[test]
819    fn a_proxy_chain_scheme_takes_the_first_entry() {
820        let mut headers = HeaderMap::new();
821        headers.insert(
822            axum::http::header::HOST,
823            "relay.example.com".parse().unwrap(),
824        );
825        headers.insert("x-forwarded-proto", "https, http".parse().unwrap());
826        assert_eq!(
827            observed_base(&headers, false).as_deref(),
828            Some("https://relay.example.com")
829        );
830    }
831
832    #[test]
833    fn no_host_header_means_nothing_observed() {
834        assert!(observed_base(&HeaderMap::new(), false).is_none());
835    }
836
837    #[test]
838    fn terminating_tls_makes_the_advertised_url_https() {
839        // Advertising http:// while refusing plaintext would hand out a URL the
840        // relay itself rejects.
841        let mut headers = HeaderMap::new();
842        headers.insert(
843            axum::http::header::HOST,
844            "relay.example.com".parse().unwrap(),
845        );
846        assert_eq!(
847            observed_base(&headers, true).as_deref(),
848            Some("https://relay.example.com")
849        );
850    }
851
852    #[test]
853    fn device_names_must_be_url_path_safe() {
854        assert!(is_valid_device_name("build-box"));
855        assert!(is_valid_device_name("laptop_2"));
856        assert!(is_valid_device_name("a"));
857
858        assert!(!is_valid_device_name(""));
859        assert!(!is_valid_device_name("has space"));
860        assert!(!is_valid_device_name("../escape"));
861        assert!(!is_valid_device_name("slash/inside"));
862        assert!(!is_valid_device_name("querylike?x=1"));
863        assert!(!is_valid_device_name(&"x".repeat(65)));
864    }
865
866    #[test]
867    fn constant_time_eq_matches_equality() {
868        assert!(constant_time_eq("abc", "abc"));
869        assert!(!constant_time_eq("abc", "abd"));
870        assert!(!constant_time_eq("abc", "ab"));
871        assert!(constant_time_eq("", ""));
872    }
873}