Skip to main content

mobius_gateway/
server.rs

1//! Authenticated raw, WebSocket-loopback, and TLS gateway listeners.
2
3mod dispatch;
4mod responses;
5mod transport;
6mod voice;
7
8use std::collections::BTreeMap;
9use std::fs;
10use std::future::Future;
11use std::path::{Path, PathBuf};
12use std::sync::{Arc, Mutex};
13use std::time::Duration;
14
15use chrono::Utc;
16use futures_util::StreamExt as _;
17use mobius::agent::validate_submission;
18use mobius::backend::session_files::{PendingSessionFileWrite, SessionFileStore};
19use mobius::protocol::Op;
20use rustls::ServerConfig;
21use rustls::pki_types::{CertificateDer, PrivateKeyDer};
22use tokio::io::{AsyncRead, AsyncWrite};
23use tokio::net::{TcpListener, TcpStream};
24use tokio::sync::broadcast;
25use tokio::task::JoinSet;
26use tokio::time::Instant;
27use tokio_rustls::TlsAcceptor;
28use tokio_tungstenite::accept_hdr_async_with_config;
29use tokio_tungstenite::tungstenite::handshake::server::{
30    Callback, ErrorResponse, Request, Response,
31};
32use tokio_tungstenite::tungstenite::http::StatusCode;
33use tokio_tungstenite::tungstenite::http::header::{HOST, ORIGIN};
34use tokio_tungstenite::tungstenite::protocol::WebSocketConfig;
35
36use crate::auth::{AuthStore, ClientIdentity, PairingGrant};
37use crate::bots::BotStore;
38use crate::config::{ConfigStore, CredentialStore, GatewayConfig, TlsConfig};
39use crate::host::{GatewayHost, HostHandle, Rejection};
40use crate::wire::{
41    ClientFrame, ClientKind, ClientMessage, ClientStatus, DirectoryEntry, DirectoryListing,
42    FrameReader, MAX_FRAME_BYTES, ProfileSnapshot, ServerFrame, ServerMessage, framed_to_websocket,
43    read_frame, read_frame_with_limit, validate_version, websocket_error, websocket_to_framed,
44    write_frame,
45};
46use crate::{Error, Result};
47
48use self::dispatch::*;
49use self::responses::*;
50use self::transport::*;
51
52const PRE_AUTH_TIMEOUT: Duration = Duration::from_secs(5);
53const MAX_AUTHENTICATED_CONNECTIONS: usize = 32;
54const MAX_PRE_AUTH_CONNECTIONS: usize = 8;
55const MAX_CONNECTIONS: usize = MAX_AUTHENTICATED_CONNECTIONS + MAX_PRE_AUTH_CONNECTIONS;
56const INACTIVITY_TIMEOUT: Duration = Duration::from_secs(72 * 60 * 60);
57const ROUTINE_TICK: Duration = Duration::from_secs(15);
58const MAX_DIRECTORY_ENTRIES: usize = 512;
59const MAX_PENDING_UPLOADS: usize = 8;
60const WEBSOCKET_BRIDGE_BYTES: usize = 16 * 1024;
61
62const _: () = assert!(MAX_FRAME_BYTES <= u32::MAX as usize);
63
64/// Fully assembled machine gateway and its chat registry.
65pub struct GatewayServer {
66    config: GatewayConfig,
67    listener: TcpListener,
68    auth: Arc<AuthStore>,
69    host: GatewayHost,
70    bots: Arc<BotStore>,
71}
72
73impl GatewayServer {
74    /// Opens protected state and the machine-wide chat registry.
75    pub async fn open(state_dir: PathBuf) -> Result<Self> {
76        let (store, config) = ConfigStore::open(state_dir)?;
77        let listener = TcpListener::bind(config.listen).await?;
78        Self::assemble(store, config, listener).await
79    }
80
81    /// Binds and initializes a fresh local gateway before exposing its one-use pairing grant.
82    pub async fn bootstrap(
83        state_dir: PathBuf,
84        listen: std::net::SocketAddr,
85    ) -> Result<(Self, PairingGrant)> {
86        let listener = TcpListener::bind(listen).await?;
87        let listen = listener.local_addr()?;
88        let (store, config) = ConfigStore::initialize(state_dir, listen, None)?;
89        let initialized_state = store.state_dir().to_path_buf();
90        let result = match AuthStore::initialize(store.auth_path()) {
91            Ok((_, grant)) => Self::assemble(store, config, listener)
92                .await
93                .map(|server| (server, grant)),
94            Err(error) => Err(error),
95        };
96        match result {
97            Ok(result) => Ok(result),
98            Err(error) => {
99                fs::remove_dir_all(&initialized_state).map_err(|cleanup| {
100                    Error::Config(format!(
101                        "{error}; failed to remove incomplete gateway state at {}: {cleanup}",
102                        initialized_state.display()
103                    ))
104                })?;
105                Err(error)
106            }
107        }
108    }
109
110    async fn assemble(
111        store: ConfigStore,
112        config: GatewayConfig,
113        listener: TcpListener,
114    ) -> Result<Self> {
115        let auth = Arc::new(AuthStore::open(store.auth_path())?);
116        let credentials = Arc::new(CredentialStore::open(store.credentials_path())?);
117        let bots = Arc::new(BotStore::open(store.state_dir())?);
118        let host =
119            GatewayHost::start(store, config.clone(), credentials, Arc::clone(&bots)).await?;
120        Ok(Self {
121            config,
122            listener,
123            auth,
124            host,
125            bots,
126        })
127    }
128
129    /// Serves until a process shutdown signal or 72 hours of inactivity.
130    pub async fn serve(self) -> Result<()> {
131        let websocket_host = self.configured_websocket_host()?;
132        self.serve_with_host(websocket_host).await
133    }
134
135    /// Serves Cloudflare WebSockets using the resolved public hostname.
136    pub(crate) async fn serve_cloudflare(self, hostname: String) -> Result<()> {
137        let cloudflare = self.config.cloudflare.as_ref().ok_or_else(|| {
138            Error::Config("a Cloudflare hostname requires tunnel configuration".into())
139        })?;
140        if cloudflare
141            .hostname()
142            .is_some_and(|configured| configured != hostname)
143        {
144            return Err(Error::Config(
145                "runtime Cloudflare hostname does not match gateway configuration".into(),
146            ));
147        }
148        self.serve_with_host(Some(hostname)).await
149    }
150
151    async fn serve_with_host(self, websocket_host: Option<String>) -> Result<()> {
152        #[cfg(unix)]
153        {
154            use tokio::signal::unix::{SignalKind, signal};
155
156            let mut interrupts = signal(SignalKind::interrupt())?;
157            let mut terminations = signal(SignalKind::terminate())?;
158            self.serve_until_inactive_with_host(
159                async move {
160                    tokio::select! {
161                        _ = interrupts.recv() => {}
162                        _ = terminations.recv() => {}
163                    }
164                },
165                INACTIVITY_TIMEOUT,
166                websocket_host,
167            )
168            .await
169        }
170        #[cfg(not(unix))]
171        self.serve_until_inactive_with_host(
172            async {
173                let _ = tokio::signal::ctrl_c().await;
174            },
175            INACTIVITY_TIMEOUT,
176            websocket_host,
177        )
178        .await
179    }
180
181    /// Serves until shutdown or the same inactivity policy as [`Self::serve`].
182    ///
183    /// Signal shutdown and await this future to close connections, finish routine
184    /// dispatch, and stop resident sessions through their normal cleanup. Merely
185    /// dropping the future does not perform graceful shutdown.
186    pub async fn serve_until(self, shutdown: impl Future<Output = ()>) -> Result<()> {
187        let websocket_host = self.configured_websocket_host()?;
188        self.serve_until_inactive_with_host(shutdown, INACTIVITY_TIMEOUT, websocket_host)
189            .await
190    }
191
192    #[cfg(test)]
193    async fn serve_until_inactive(
194        self,
195        shutdown: impl Future<Output = ()>,
196        inactivity_timeout: Duration,
197    ) -> Result<()> {
198        let websocket_host = self.configured_websocket_host()?;
199        self.serve_until_inactive_with_host(shutdown, inactivity_timeout, websocket_host)
200            .await
201    }
202
203    async fn serve_until_inactive_with_host(
204        self,
205        shutdown: impl Future<Output = ()>,
206        inactivity_timeout: Duration,
207        websocket_host: Option<String>,
208    ) -> Result<()> {
209        self.config.validate()?;
210        let tls = self.config.tls.as_ref().map(tls_acceptor).transpose()?;
211        if tls.is_none() && !self.listener.local_addr()?.ip().is_loopback() {
212            return Err(Error::Config(
213                "plaintext listeners are restricted to loopback".into(),
214            ));
215        }
216        let mut connections = JoinSet::new();
217        let mut routine_dispatchers = JoinSet::new();
218        let connection_admission =
219            ConnectionAdmission::new(MAX_PRE_AUTH_CONNECTIONS, MAX_AUTHENTICATED_CONNECTIONS);
220        let client_connections = Arc::new(ClientConnections::default());
221        let (client_revocations, _) = broadcast::channel(MAX_CONNECTIONS);
222        let mut has_active_routines = self.bots.has_active_routines(Utc::now().timestamp())?;
223        let inactivity = tokio::time::sleep(inactivity_timeout);
224        tokio::pin!(inactivity);
225        let mut routine_timer = tokio::time::interval(ROUTINE_TICK);
226        routine_timer.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
227        tokio::pin!(shutdown);
228        let result = async {
229            loop {
230            tokio::select! {
231                biased;
232                () = &mut shutdown => break Ok(()),
233                _ = routine_timer.tick() => {
234                    let now = Utc::now().timestamp();
235                    let routines_active = self.bots.has_active_routines(now)?;
236                    if has_active_routines && !routines_active && connections.is_empty() {
237                        inactivity.as_mut().reset(tokio::time::Instant::now() + inactivity_timeout);
238                    }
239                    has_active_routines = routines_active;
240                    let due = self.bots.take_due(now)?;
241                    if !due.is_empty() {
242                        let host = self.host.clone();
243                        routine_dispatchers.spawn(async move {
244                            for (routine_id, run) in due {
245                                if let Err(error) = host.run_due_routine(routine_id.clone(), run).await {
246                                    eprintln!(
247                                        "routine run failed: routine_id={routine_id} code={} message={}",
248                                        error.code, error.message
249                                    );
250                                }
251                            }
252                        });
253                    }
254                }
255                Some(_) = connections.join_next(), if !connections.is_empty() => {
256                    if connections.is_empty() {
257                        has_active_routines =
258                            self.bots.has_active_routines(Utc::now().timestamp())?;
259                        if !has_active_routines {
260                            inactivity.as_mut().reset(tokio::time::Instant::now() + inactivity_timeout);
261                        }
262                    }
263                }
264                Some(_) = routine_dispatchers.join_next(), if !routine_dispatchers.is_empty() => {}
265                accepted = async {
266                    let admission = connection_admission.admit().await;
267                    self.listener.accept().await.map(|accepted| (accepted, admission))
268                }, if connections.len() < MAX_CONNECTIONS => {
269                    let ((stream, peer), admission) = accepted?;
270                    let auth = Arc::clone(&self.auth);
271                    let host = self.host.clone();
272                    let bots = Arc::clone(&self.bots);
273                    let client_connections = Arc::clone(&client_connections);
274                    let client_revocations = client_revocations.clone();
275                    let tls = tls.clone();
276                    let websocket_host = websocket_host.clone();
277                    connections.spawn(async move {
278                        let auth_deadline = Instant::now() + PRE_AUTH_TIMEOUT;
279                        let connection = ConnectionContext {
280                            local: peer.ip().is_loopback(),
281                            auth,
282                            host,
283                            bots,
284                            client_connections,
285                            client_revocations,
286                            admission,
287                        };
288                        let result = if let Some(tls) = tls {
289                            let stream = match tokio::time::timeout_at(
290                                auth_deadline,
291                                tls.accept(stream),
292                            )
293                            .await
294                            {
295                                Ok(Ok(stream)) => stream,
296                                Ok(Err(error)) => {
297                                    eprintln!("gateway TLS handshake failed: {:?}", error.kind());
298                                    return;
299                                }
300                                Err(_) => {
301                                    eprintln!("gateway TLS handshake timed out");
302                                    return;
303                                }
304                            };
305                            serve_connection(stream, connection, auth_deadline, None).await
306                        } else {
307                            serve_plaintext_connection(
308                                stream,
309                                connection,
310                                PlaintextHandshake {
311                                    expected_websocket_host: websocket_host,
312                                    auth_deadline,
313                                },
314                            )
315                            .await
316                        };
317                        if let Err(error) = result {
318                            eprintln!("gateway connection failed: {}", connection_diagnostic(&error));
319                        }
320                    });
321                }
322                () = &mut inactivity, if connections.is_empty() && !has_active_routines => {
323                    has_active_routines = self.bots.has_active_routines(Utc::now().timestamp())?;
324                    if !has_active_routines {
325                        break Ok(());
326                    }
327                }
328            }
329            }
330        }
331        .await;
332        connections.shutdown().await;
333        while routine_dispatchers.join_next().await.is_some() {}
334        self.host.shutdown().await;
335        result
336    }
337
338    fn configured_websocket_host(&self) -> Result<Option<String>> {
339        self.config
340            .cloudflare
341            .as_ref()
342            .map(|cloudflare| {
343                cloudflare.hostname().map(str::to_owned).ok_or_else(|| {
344                    Error::Config(
345                        "quick tunnel hostname is unavailable before cloudflared starts".into(),
346                    )
347                })
348            })
349            .transpose()
350    }
351
352    /// Returns the bound address from persisted configuration.
353    #[must_use]
354    pub const fn listen_addr(&self) -> std::net::SocketAddr {
355        self.config.listen
356    }
357}
358
359#[cfg(test)]
360mod tests;