mobius_gateway/
server.rs

1//! Authenticated raw, WebSocket-loopback, and TLS gateway listeners.
2
3use crate::host::session_file_rejection;
4
5mod desktop;
6mod policy;
7pub use policy::ConnectionPolicy;
8mod dispatch;
9mod responses;
10mod transport;
11mod view;
12mod voice;
13
14use crate::telemetry::{StopCause, Trigger};
15use std::collections::{BTreeMap, BTreeSet};
16use std::fs;
17use std::future::Future;
18use std::path::{Path, PathBuf};
19use std::sync::{Arc, Mutex};
20use std::time::{Duration, SystemTime, UNIX_EPOCH};
21
22use chrono::Utc;
23use mobius::agent::validate_submission;
24use mobius::backend::session_files::{PendingSessionFileWrite, SessionFileStore};
25use mobius::protocol::Op;
26use rustls::ServerConfig;
27use rustls::pki_types::{CertificateDer, PrivateKeyDer};
28use tokio::io::{AsyncRead, AsyncWrite};
29use tokio::net::{TcpListener, TcpStream};
30use tokio::sync::broadcast;
31use tokio::task::JoinSet;
32use tokio::time::Instant;
33use tokio_rustls::TlsAcceptor;
34use tokio_tungstenite::accept_hdr_async_with_config;
35use tokio_tungstenite::tungstenite::handshake::server::{
36    Callback, ErrorResponse, Request, Response,
37};
38use tokio_tungstenite::tungstenite::http::StatusCode;
39use tokio_tungstenite::tungstenite::http::header::{HOST, ORIGIN, SEC_WEBSOCKET_PROTOCOL};
40use tokio_tungstenite::tungstenite::protocol::WebSocketConfig;
41use uuid::Uuid;
42
43use crate::auth::{AuthStore, ClientIdentity, PairingGrant};
44use crate::bots::BotStore;
45use crate::config::{ConfigStore, CredentialStore, GatewayConfig, TlsConfig};
46use crate::host::{GatewayHost, HostHandle, Rejection};
47use crate::wire::{
48    CatalogHint, ClientFrame, ClientKind, ClientMessage, ClientStatus, DirectoryEntry,
49    DirectoryListing, FrameReader, GatewayNotification, MAX_FRAME_BYTES, ProfileSnapshot,
50    ServerFrame, ServerMessage, SharedFrame, read_frame, read_frame_with_limit, validate_version,
51    websocket_error, write_frame,
52};
53use crate::{Error, Result};
54
55use self::dispatch::*;
56use self::responses::*;
57use self::transport::*;
58use self::view::ClientView;
59
60const ROUTINE_TICK: Duration = Duration::from_secs(15);
61const MAX_DIRECTORY_ENTRIES: usize = 512;
62use crate::wire::WEBSOCKET_BRIDGE_BYTES;
63const ACCESS_EXPIRY_ENV: &str = "MOBIUS_GATEWAY_ACCESS_EXPIRES_AT";
64
65const _: () = assert!(MAX_FRAME_BYTES <= u32::MAX as usize);
66
67/// Fully assembled machine gateway and its chat registry.
68pub struct GatewayServer {
69    config: GatewayConfig,
70    listener: TcpListener,
71    access_lease: Option<AccessLease>,
72    auth: Arc<AuthStore>,
73    host: GatewayHost,
74    bots: Arc<BotStore>,
75    ready: Option<tokio::sync::oneshot::Sender<()>>,
76}
77
78impl GatewayServer {
79    /// Opens protected state and the machine-wide chat registry.
80    /// # Errors
81    ///
82    /// Returns an error if the resource cannot be read, decoded, or validated.
83    pub async fn open(state_dir: PathBuf) -> Result<Self> {
84        let (store, config) = ConfigStore::open(state_dir)?;
85        let listener = TcpListener::bind(config.listen).await?;
86        Self::assemble(store, config, listener).await
87    }
88
89    /// Binds and initializes a fresh local gateway before exposing its one-use pairing grant.
90    /// # Errors
91    ///
92    /// Returns an error if configuration is invalid or a required resource cannot be initialized.
93    pub async fn bootstrap(
94        state_dir: PathBuf,
95        listen: std::net::SocketAddr,
96    ) -> Result<(Self, PairingGrant)> {
97        let listener = TcpListener::bind(listen).await?;
98        let listen = listener.local_addr()?;
99        let (store, config) = ConfigStore::initialize(state_dir, listen, None)?;
100        let initialized_state = store.state_dir().to_path_buf();
101        let result = match AuthStore::initialize(store.auth_path(), config.auth) {
102            Ok((_, grant)) => Self::assemble(store, config, listener)
103                .await
104                .map(|server| (server, grant)),
105            Err(error) => Err(error),
106        };
107        match result {
108            Ok(result) => Ok(result),
109            Err(error) => {
110                fs::remove_dir_all(&initialized_state).map_err(|cleanup| {
111                    Error::Config(format!(
112                        "{error}; failed to remove incomplete gateway state at {}: {cleanup}",
113                        initialized_state.display()
114                    ))
115                })?;
116                Err(error)
117            }
118        }
119    }
120
121    async fn assemble(
122        store: ConfigStore,
123        config: GatewayConfig,
124        listener: TcpListener,
125    ) -> Result<Self> {
126        let access_lease = configured_access_lease(&config.runtime)?;
127        let auth = Arc::new(AuthStore::open(store.auth_path(), config.auth)?);
128        let credentials = Arc::new(CredentialStore::open(store.credentials_path())?);
129        let bots = Arc::new(BotStore::open(store.state_dir())?);
130        bots.sync_telemetry_cursors(&config.telemetry.sinks)?;
131        let host =
132            GatewayHost::start(store, config.clone(), credentials, Arc::clone(&bots)).await?;
133        Ok(Self {
134            config,
135            listener,
136            access_lease,
137            auth,
138            host,
139            bots,
140            ready: None,
141        })
142    }
143
144    pub(crate) fn notify_ready(&mut self) -> tokio::sync::oneshot::Receiver<()> {
145        let (sender, receiver) = tokio::sync::oneshot::channel();
146        self.ready = Some(sender);
147        receiver
148    }
149
150    /// Serves until shutdown or the configured idle interval (72 hours by default).
151    /// # Errors
152    ///
153    /// Returns an error if validation or an operation required to complete the request fails.
154    pub async fn serve(self) -> Result<()> {
155        let websocket_host = self.configured_websocket_host()?;
156        self.serve_with_host(websocket_host).await
157    }
158
159    /// Serves Cloudflare WebSockets using the resolved public hostname.
160    pub(crate) async fn serve_cloudflare(self, hostname: String) -> Result<()> {
161        let cloudflare = self.config.cloudflare.as_ref().ok_or_else(|| {
162            Error::Config("a Cloudflare hostname requires tunnel configuration".into())
163        })?;
164        if cloudflare
165            .hostname()
166            .is_some_and(|configured| configured != hostname)
167        {
168            return Err(Error::Config(
169                "runtime Cloudflare hostname does not match gateway configuration".into(),
170            ));
171        }
172        self.serve_with_host(Some(hostname)).await
173    }
174
175    async fn serve_with_host(self, websocket_host: Option<String>) -> Result<()> {
176        let inactivity_timeout = Duration::from_secs(self.config.runtime.idle_exit_seconds);
177        #[cfg(unix)]
178        {
179            use tokio::signal::unix::{SignalKind, signal};
180
181            let mut interrupts = signal(SignalKind::interrupt())?;
182            let mut terminations = signal(SignalKind::terminate())?;
183            self.serve_until_inactive_with_host(
184                async move {
185                    tokio::select! {
186                        _ = interrupts.recv() => {}
187                        _ = terminations.recv() => {}
188                    }
189                },
190                inactivity_timeout,
191                websocket_host,
192            )
193            .await
194        }
195        #[cfg(not(unix))]
196        self.serve_until_inactive_with_host(
197            async {
198                let _ = tokio::signal::ctrl_c().await;
199            },
200            inactivity_timeout,
201            websocket_host,
202        )
203        .await
204    }
205
206    /// Serves until shutdown or the same inactivity policy as [`Self::serve`].
207    ///
208    /// Signal shutdown and await this future to close connections, finish routine
209    /// dispatch, and stop resident sessions through their normal cleanup. Merely
210    /// dropping the future does not perform graceful shutdown.
211    /// # Errors
212    ///
213    /// Returns an error if validation or an operation required to complete the request fails.
214    pub async fn serve_until(self, shutdown: impl Future<Output = ()>) -> Result<()> {
215        let websocket_host = self.configured_websocket_host()?;
216        let inactivity_timeout = Duration::from_secs(self.config.runtime.idle_exit_seconds);
217        self.serve_until_inactive_with_host(shutdown, inactivity_timeout, websocket_host)
218            .await
219    }
220
221    #[cfg(test)]
222    async fn serve_until_inactive(
223        self,
224        shutdown: impl Future<Output = ()>,
225        inactivity_timeout: Duration,
226    ) -> Result<()> {
227        let websocket_host = self.configured_websocket_host()?;
228        self.serve_until_inactive_with_host(shutdown, inactivity_timeout, websocket_host)
229            .await
230    }
231
232    async fn serve_until_inactive_with_host(
233        mut self,
234        shutdown: impl Future<Output = ()>,
235        inactivity_timeout: Duration,
236        websocket_host: Option<String>,
237    ) -> Result<()> {
238        self.config.validate()?;
239        let tls = self.config.tls.as_ref().map(tls_acceptor).transpose()?;
240        if tls.is_none() && !self.listener.local_addr()?.ip().is_loopback() {
241            return Err(Error::Config(
242                "plaintext listeners are restricted to loopback".into(),
243            ));
244        }
245        let ingress = match self.config.runtime.ingress {
246            Some(address) => Some(TcpListener::bind(address).await?),
247            None => None,
248        };
249        let mut telemetry_tasks = JoinSet::new();
250        let mut stop_cause = StopCause::Signal;
251        let mut connections = JoinSet::new();
252        let mut routine_dispatchers = JoinSet::new();
253        let mut next_nudge = Instant::now();
254        let connection_admission = ConnectionAdmission::new(
255            self.config.connections.pre_authentication,
256            self.config.connections.authenticated,
257        );
258        let client_connections = Arc::new(ClientConnections::default());
259        let mut activity_hook =
260            crate::telemetry::ActivityHook::new(self.config.telemetry.activity_hook.as_ref())?;
261        let (client_revocations, _) = broadcast::channel(self.config.connections.total());
262        let mut has_active_routines = self.bots.has_active_routines(Utc::now().timestamp())?
263            || self.bots.has_pending_deliveries()?
264            || self.bots.has_monitored_sessions()?;
265        let inactivity = tokio::time::sleep(inactivity_timeout);
266        tokio::pin!(inactivity);
267        let mut routine_timer = tokio::time::interval(ROUTINE_TICK);
268        routine_timer.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
269        tokio::pin!(shutdown);
270        if let Some(ready) = self.ready.take() {
271            let _ = ready.send(());
272        }
273        crate::telemetry::Telemetry::tick(&self.host, 0, Trigger::Start, &mut telemetry_tasks)
274            .await;
275        let mut access_expired = false;
276        let serving = async {
277            loop {
278                tokio::select! {
279                    biased;
280                    _ = async {
281                        tokio::time::sleep_until(next_nudge).await;
282                        self.host.telemetry.notify.notified().await;
283                    }, if routine_dispatchers.is_empty() => {
284                        next_nudge = Instant::now() + Duration::from_secs(1);
285                        let host = self.host.clone();
286                        routine_dispatchers.spawn(async move {
287                            if let Err(error) = host.dispatch_bot_events().await { eprintln!("Bot delivery failed: code={}", error.code); }
288                        });
289                    }
290                    _ = routine_timer.tick() => {
291                        if self.access_lease.is_some_and(AccessLease::expired) {
292                            access_expired = true;
293                            break Ok(());
294                        }
295                        self.tick_telemetry(&client_connections, Trigger::Interval, &mut telemetry_tasks).await;
296                        // Keep reservation and the shutdown decision under the same
297                        // admission gate; future schedules alone do not keep it open.
298                        let Ok(_admission) = self.host.begin_mutation().await else { continue; };
299                        let now = Utc::now().timestamp();
300                        let poll = self.bots.poll_due(now)?;
301                        let routines_active = poll.active || self.bots.has_pending_deliveries()? || self.bots.has_monitored_sessions()?;
302                        if has_active_routines && !routines_active && connections.is_empty() {
303                            inactivity.as_mut().reset(tokio::time::Instant::now() + inactivity_timeout);
304                        }
305                        has_active_routines = routines_active;
306                        let host = self.host.clone();
307                        // One event delivery worker at a time; all retries retain their original message ID.
308                        if routine_dispatchers.is_empty() {
309                            routine_dispatchers.spawn(async move {
310                                if let Err(error) = host.dispatch_bot_events().await { eprintln!("Bot delivery failed: code={}",error.code); }
311                            });
312                        }
313                        // The clock only committed schedule.due events. The same
314                        // hook action worker starts scheduled and manually requested runs.
315                        let _ = poll.events;
316
317                    }
318                    Some(result) = telemetry_tasks.join_next(), if !telemetry_tasks.is_empty() => {
319                        if let Err(error) = result { eprintln!("telemetry worker failed: {error}"); }
320                    }
321                    Some(_) = connections.join_next(), if !connections.is_empty() => {
322                        if connections.is_empty() {
323                            has_active_routines =
324                                self.bots.has_active_routines(Utc::now().timestamp())? || self.bots.has_pending_deliveries()? || self.bots.has_monitored_sessions()?;
325                            if !has_active_routines {
326                                inactivity.as_mut().reset(tokio::time::Instant::now() + inactivity_timeout);
327                            }
328                        }
329                    }
330                    Some(_) = routine_dispatchers.join_next(), if !routine_dispatchers.is_empty() => {
331                        self.tick_telemetry(&client_connections, Trigger::Interval, &mut telemetry_tasks).await;
332                    }
333                    accepted = async {
334                        let admission = connection_admission.admit().await;
335                        tokio::select! {
336                            accepted = self.listener.accept() => accepted.map(|accepted| (accepted, admission, false)),
337                            accepted = async { match &ingress { Some(listener) => listener.accept().await, None => std::future::pending().await } } => accepted.map(|accepted| (accepted, admission, true)),
338                        }
339                    }, if connections.len() < self.config.connections.total() => {
340                        let ((stream, peer), admission, ingress_connection) = accepted?;
341                        if self.access_lease.is_some_and(AccessLease::expired) {
342                            access_expired = true;
343                            break Ok(());
344                        }
345                        let auth = Arc::clone(&self.auth);
346                        let host = self.host.clone();
347                        let bots = Arc::clone(&self.bots);
348                        let client_connections = Arc::clone(&client_connections);
349                        let client_revocations = client_revocations.clone();
350                        let tls = tls.clone();
351                        let websocket_host = websocket_host.clone();
352                        connections.spawn(async move {
353                            let auth_deadline = Instant::now() + Duration::from_secs(self.config.connections.authentication_timeout_seconds);
354                            let connection = ConnectionContext {
355                                local: peer.ip().is_loopback(),
356                                desktop_transport: tls.is_some(),
357                                auth,
358                                host,
359                                bots,
360                                client_connections,
361                                client_revocations,
362                                admission,
363                                access_lease: self.access_lease,
364                            };
365                            let result = if ingress_connection {
366                                serve_websocket(stream, connection, PlaintextHandshake { expected_websocket_host: None, auth_deadline }).await
367                            } else if let Some(tls) = tls {
368                                let stream = match tokio::time::timeout_at(
369                                    auth_deadline,
370                                    tls.accept(stream),
371                                )
372                                .await
373                                {
374                                    Ok(Ok(stream)) => stream,
375                                    Ok(Err(error)) => {
376                                        eprintln!("gateway TLS handshake failed: {:?}", error.kind());
377                                        return;
378                                    }
379                                    Err(_) => {
380                                        eprintln!("gateway TLS handshake timed out");
381                                        return;
382                                    }
383                                };
384                                serve_connection(stream, connection, auth_deadline, None).await
385                            } else {
386                                serve_plaintext_connection(
387                                    stream,
388                                    connection,
389                                    PlaintextHandshake {
390                                        expected_websocket_host: websocket_host,
391                                        auth_deadline,
392                                    },
393                                )
394                                .await
395                            };
396                            if let Err(error) = result {
397                                eprintln!("gateway connection failed: {}", connection_diagnostic(&error));
398                            }
399                        });
400                    }
401                    () = &mut inactivity, if connections.is_empty() && !has_active_routines && self.config.runtime.idle_exit_seconds != 0 => {
402                        has_active_routines = self.bots.has_active_routines(Utc::now().timestamp())? || self.bots.has_pending_deliveries()? || self.bots.has_monitored_sessions()?;
403                        if !has_active_routines {
404                            if crate::telemetry::Telemetry::pending(&self.host).await {
405                                inactivity.as_mut().reset(Instant::now() + ROUTINE_TICK);
406                                continue;
407                            }
408                            stop_cause = StopCause::Idle;
409                            break Ok(());
410                        }
411                    }
412                }
413            }
414        };
415        let (result, lease_expired) = {
416            // Poll activity alongside the whole loop, including awaits inside
417            // dispatch branches. This scope cancels both futures before actor cleanup.
418            let activity = activity_hook.run(&self.host, || client_connections.native_count());
419            tokio::pin!(activity, serving);
420            tokio::select! {
421                biased;
422                () = &mut shutdown => (Ok(()), false),
423                () = async {
424                    match self.access_lease {
425                        Some(lease) => tokio::time::sleep_until(lease.deadline).await,
426                        None => std::future::pending().await,
427                    }
428                } => (Ok(()), true),
429                () = &mut activity => (Err(Error::Config("telemetry activity worker stopped".into())), false),
430                result = &mut serving => (result, false),
431            }
432        };
433        access_expired |= lease_expired;
434        // Snapshot workers can also hold host state while waiting for actors.
435        telemetry_tasks.shutdown().await;
436        connections.shutdown().await;
437        if access_expired {
438            routine_dispatchers.shutdown().await;
439        } else {
440            while routine_dispatchers.join_next().await.is_some() {}
441        }
442        self.host.shutdown().await;
443        let cause = if access_expired {
444            StopCause::LeaseExpired
445        } else if result.is_err() {
446            StopCause::Error
447        } else {
448            stop_cause
449        };
450        crate::telemetry::Telemetry::stop(&self.host, cause, &mut telemetry_tasks).await;
451        activity_hook.stop().await;
452        result
453    }
454
455    async fn tick_telemetry(
456        &self,
457        connections: &ClientConnections,
458        trigger: Trigger,
459        tasks: &mut JoinSet<()>,
460    ) {
461        match connections.native_count() {
462            Ok(clients) => {
463                crate::telemetry::Telemetry::tick(&self.host, clients, trigger, tasks).await
464            }
465            Err(error) => eprintln!("telemetry client count unavailable: {error}"),
466        }
467    }
468
469    fn configured_websocket_host(&self) -> Result<Option<String>> {
470        self.config
471            .cloudflare
472            .as_ref()
473            .map(|cloudflare| {
474                cloudflare.hostname().map(str::to_owned).ok_or_else(|| {
475                    Error::Config(
476                        "quick tunnel hostname is unavailable before cloudflared starts".into(),
477                    )
478                })
479            })
480            .transpose()
481    }
482
483    /// Returns the bound address from persisted configuration.
484    #[must_use]
485    pub const fn listen_addr(&self) -> std::net::SocketAddr {
486        self.config.listen
487    }
488}
489
490#[derive(Clone, Copy)]
491struct AccessLease {
492    expires_at: SystemTime,
493    deadline: Instant,
494}
495
496impl AccessLease {
497    fn expired(self) -> bool {
498        SystemTime::now() >= self.expires_at || Instant::now() >= self.deadline
499    }
500}
501
502fn configured_access_lease(runtime: &crate::config::RuntimeConfig) -> Result<Option<AccessLease>> {
503    let expires_at = match std::env::var(ACCESS_EXPIRY_ENV) {
504        Ok(value) => Some(value),
505        Err(std::env::VarError::NotPresent) => None,
506        Err(_) => return Err(Error::Config("gateway access expiry is invalid".into())),
507    };
508    access_lease(
509        expires_at.as_deref(),
510        runtime.require_access_lease,
511        Duration::from_secs(runtime.access_grace_seconds),
512        SystemTime::now(),
513    )
514}
515
516fn access_lease(
517    expires_at: Option<&str>,
518    required: bool,
519    grace: Duration,
520    now: SystemTime,
521) -> Result<Option<AccessLease>> {
522    let Some(expires_at) = expires_at else {
523        return if required {
524            Err(Error::Config("gateway access expiry is required".into()))
525        } else {
526            Ok(None)
527        };
528    };
529    if expires_at.is_empty() || !expires_at.bytes().all(|byte| byte.is_ascii_digit()) {
530        return Err(Error::Config("gateway access expiry is invalid".into()));
531    }
532    let seconds = expires_at
533        .parse::<u64>()
534        .map_err(|_| Error::Config("gateway access expiry is invalid".into()))?;
535    let expires_at = UNIX_EPOCH
536        .checked_add(Duration::from_secs(seconds))
537        .and_then(|expiry| expiry.checked_add(grace))
538        .ok_or_else(|| Error::Config("gateway access expiry is invalid".into()))?;
539    let remaining = expires_at
540        .duration_since(now)
541        .map_err(|_| Error::Config("gateway access has expired".into()))?;
542    if remaining.is_zero() {
543        return Err(Error::Config("gateway access has expired".into()));
544    }
545    let deadline = Instant::now()
546        .checked_add(remaining)
547        .ok_or_else(|| Error::Config("gateway access expiry is invalid".into()))?;
548    Ok(Some(AccessLease {
549        expires_at,
550        deadline,
551    }))
552}
553
554#[cfg(test)]
555mod tests;
mobius_gateway/server.rs

mobius_gateway/
server.rs
