1#[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
48pub const HEARTBEAT_TIMEOUT: Duration = Duration::from_secs(90);
50
51const ENROLL_TIMEOUT: Duration = Duration::from_secs(10);
53
54#[derive(Debug, Clone)]
56pub struct RelayConfig {
57 pub bind: SocketAddr,
59 pub enroll_token: String,
61 pub rate_limit: RateLimitConfig,
68 #[cfg(feature = "tls")]
70 pub tls: Option<crate::tls::TlsFiles>,
71 pub public_base: Option<String>,
77}
78
79impl RelayConfig {
80 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 #[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 pub fn without_rate_limit(mut self) -> Self {
101 self.rate_limit.enabled = false;
102 self
103 }
104
105 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 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 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#[derive(Debug, Clone)]
131pub struct RelayState {
132 config: RelayConfig,
133 devices: DeviceRegistry,
134}
135
136impl RelayState {
137 pub fn new(config: RelayConfig) -> Self {
139 Self {
140 config,
141 devices: DeviceRegistry::new(),
142 }
143 }
144
145 pub fn devices(&self) -> &DeviceRegistry {
147 &self.devices
148 }
149}
150
151pub 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
172pub 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 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 let service = router.into_make_service_with_connect_info::<SocketAddr>();
201
202 #[cfg(feature = "tls")]
203 if let Some(files) = tls {
204 let config = crate::tls::acceptor(files.load()?);
207 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
223fn 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
236fn 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 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
260fn 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
272async 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
282async fn control_session(socket: WebSocket, state: RelayState, observed: Option<String>) {
284 let (mut sink, mut stream) = socket.split();
285
286 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 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 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 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 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 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 _ => 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
419async 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
456async fn data_handler(ws: WebSocketUpgrade, State(state): State<RelayState>) -> Response {
463 ws.on_upgrade(move |socket| attach_data_connection(socket, state))
464}
465
466async 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 if let Some(mut extra) = device.offer(socket).await {
497 let _ = extra.close().await;
498 }
499}
500
501async 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 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 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 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
592fn 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
604async 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 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 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
671const MAX_BODY: usize = 8 * 1024 * 1024;
673
674async 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
721async 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
734async 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
744fn 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 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}