1mod 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
62pub struct GatewayServer {
64 config: GatewayConfig,
65 listener: TcpListener,
66 auth: Arc<AuthStore>,
67 host: GatewayHost,
68 cron: Arc<CronStore>,
69}
70
71impl GatewayServer {
72 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 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 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 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 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 #[must_use]
321 pub const fn listen_addr(&self) -> std::net::SocketAddr {
322 self.config.listen
323 }
324}
325
326#[cfg(test)]
327mod tests;