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