Skip to main content

mobius_gateway/
server.rs

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