max / makenotwork

12.2 KB · 300 lines History Blame Raw

1	use multithreaded::{AppState, config::Config, csrf};
2	use sqlx::postgres::PgPoolOptions;
3	use tokio::net::TcpListener;
4	use tower_sessions::ExpiredDeletion;
5	use tower_sessions::SessionManagerLayer;
6	use tower_sessions::cookie::SameSite;
7	use tower_sessions_sqlx_store::PostgresStore;
8	use tracing_subscriber::EnvFilter;
9
10	#[tokio::main]
11	async fn main() {
12	tracing_subscriber::fmt()
13	.with_env_filter(EnvFilter::from_default_env())
14	.init();
15
16	dotenvy::dotenv().ok();
17
18	let database_url = std::env::var("DATABASE_URL").expect("DATABASE_URL must be set");
19
20	// Explicit pool sizing for the handler pool. Some handlers fan out ~3
21	// queries concurrently via `try_join!`, so 32 connections sustain ~10
22	// concurrent such requests before acquisition queues. Session I/O and its
23	// expiry sweep run on a separate pool (below), so they no longer contend
24	// here (M-Pf2). Bound acquisition so a burst sheds fast rather than hanging.
25	let pool = PgPoolOptions::new()
26	.max_connections(32)
27	.acquire_timeout(std::time::Duration::from_secs(10))
28	.connect(&database_url)
29	.await
30	.expect("failed to connect to database");
31
32	sqlx::migrate!()
33	.run(&pool)
34	.await
35	.expect("failed to run migrations");
36
37	tracing::info!("migrations applied");
38
39	// Seed initial data if --seed flag is passed, then exit
40	if std::env::args().any(\|a\| a == "--seed") {
41	multithreaded::seed::run(&pool).await;
42	tracing::info!("seed data inserted");
43	return;
44	}
45
46	let config = Config::from_env();
47	multithreaded::error_page::init(config.mnw_base_url.clone());
48
49	// Probe the host trust store now rather than on the first outbound request,
50	// so a host with no usable CA anchors says so in the journal at boot and on
51	// /api/health, instead of presenting as a login outage on a box that reports
52	// healthy. The answer is cached, so the health handler does not repeat it.
53	multithreaded::trust_store::anchors_ok();
54
55	// Optional S3 storage for image uploads
56	let s3 = if let Some(ref s3_config) = config.s3 {
57	match multithreaded::storage::S3Storage::new(s3_config).await {
58	Ok(client) => {
59	tracing::info!("S3 storage configured (bucket: {})", s3_config.bucket);
60	Some(std::sync::Arc::new(client))
61	}
62	Err(e) => {
63	tracing::warn!("S3 storage unavailable: {e}");
64	None
65	}
66	}
67	} else {
68	tracing::info!("S3 storage not configured (image uploads disabled)");
69	None
70	};
71
72	let state = AppState {
73	db: pool.clone(),
74	config,
75	http: multithreaded::tls::builder()
76	.timeout(std::time::Duration::from_secs(15))
77	.connect_timeout(std::time::Duration::from_secs(5))
78	.build()
79	.expect("failed to build HTTP client"),
80	link_preview: multithreaded::link_preview::LinkPreviewFetcher::Http(
81	multithreaded::link_preview::build_preview_client(),
82	),
83	s3,
84	chat: std::sync::Arc::new(multithreaded::chat::new_chat()),
85	chat_identities: multithreaded::chat::IdentityCache::new(),
86	};
87
88	// Session store backed by PostgreSQL, on its own small pool. Every authed
89	// request does session I/O and the hourly expiry sweep scans the session
90	// table; keeping them off the handler pool means a handler-query burst can't
91	// starve session reads (or vice versa) (M-Pf2).
92	let session_pool = PgPoolOptions::new()
93	.max_connections(6)
94	.acquire_timeout(std::time::Duration::from_secs(10))
95	.connect(&database_url)
96	.await
97	.expect("failed to connect session store pool");
98	let session_store = PostgresStore::new(session_pool);
99	session_store
100	.migrate()
101	.await
102	.expect("failed to migrate session store");
103
104	// Both background loops run under panic supervision: a panic inside a sweep
105	// would otherwise silently kill the worker for the life of the process
106	// (S3 orphans pile up / the session table grows unbounded) while the server
107	// keeps serving. `supervise` restarts the worker after a panic.
108	let deletion_store = session_store.clone();
109	let deletion_task = supervise("session-expiry-sweep", move \|\| {
110	deletion_store
111	.clone()
112	.continuously_delete_expired(tokio::time::Duration::from_hours(1))
113	});
114
115	// Reconcile sweep: purge S3 objects for removed images (backlog + retries).
116	// Only meaningful when S3 is configured.
117	let purge_task = state.s3.as_ref().map(\|s3\| {
118	let db = state.db.clone();
119	let s3 = s3.clone();
120	supervise("image-purge-sweep", move \|\| {
121	multithreaded::maintenance::continuously_purge_removed_images(
122	db.clone(),
123	s3.clone(),
124	tokio::time::Duration::from_hours(6),
125	)
126	})
127	});
128
129	// Chat retention: expire by age and trim each room to its message cap.
130	// Unconditional. A deployment where no community has chat enabled runs two
131	// indexed statements that match nothing, which is cheaper than the check.
132	let chat_sweep_db = state.db.clone();
133	let chat_sweep_task = supervise("chat-retention-sweep", move \|\| {
134	multithreaded::maintenance::continuously_sweep_chat(
135	chat_sweep_db.clone(),
136	multithreaded::maintenance::CHAT_SWEEP_INTERVAL,
137	)
138	});
139
140	// Chat send-rate buckets. Keyed by (user, room), so the key space is
141	// unbounded and attacker-influenced; this is what keeps it off the 512M cap.
142	let chat_limiter = state.chat.clone();
143	let chat_rate_task = supervise("chat-rate-bucket-sweep", move \|\| {
144	multithreaded::maintenance::continuously_sweep_chat_rate_limits(chat_limiter.clone())
145	});
146
147	let session_layer = SessionManagerLayer::new(session_store)
148	.with_name("mt_session")
149	.with_same_site(SameSite::Lax)
150	.with_expiry(tower_sessions::Expiry::OnInactivity(time::Duration::days(
151	7,
152	)))
153	.with_secure(state.config.cookie_secure);
154
155	// CSRF + session are scoped to the forum routes only; the internal API uses
156	// HMAC auth and must not run them.
157	let forum = multithreaded::routes::forum_routes(state.clone())
158	.layer(axum::middleware::from_fn(csrf::csrf_middleware))
159	.layer(session_layer);
160
161	// Security headers (CSP, nosniff, X-Frame, cache-control) wrap the WHOLE app,
162	// applied outermost so `/static` assets and the internal API get them too,
163	// not just forum routes (the layers used to sit inside `forum_routes`, before
164	// the merge/nest, so static responses shipped without nosniff/X-Frame/CSP).
165	let app = forum
166	// Internal API routes, HMAC auth only, no CSRF/session middleware
167	.merge(multithreaded::routes::internal::internal_routes(state))
168	// Assets are compiled into the binary, not read from a `static/` dir
169	// beside it, so mt ships as the single file Sando's companion mechanism
170	// installs and its markup can never disagree with its stylesheet. See
171	// multithreaded::static_assets.
172	.route(
173	"/static/{*path}",
174	axum::routing::get(multithreaded::static_assets::serve),
175	)
176	.layer(tower_http::set_header::SetResponseHeaderLayer::overriding(
177	axum::http::header::CONTENT_SECURITY_POLICY,
178	axum::http::HeaderValue::from_static(
179	"default-src 'self'; img-src 'self'; style-src 'self'; \
180	frame-ancestors 'none'; object-src 'none'; base-uri 'none'; \
181	form-action 'self'",
182	),
183	))
184	.layer(tower_http::set_header::SetResponseHeaderLayer::overriding(
185	axum::http::header::X_CONTENT_TYPE_OPTIONS,
186	axum::http::HeaderValue::from_static("nosniff"),
187	))
188	.layer(tower_http::set_header::SetResponseHeaderLayer::overriding(
189	axum::http::header::X_FRAME_OPTIONS,
190	axum::http::HeaderValue::from_static("DENY"),
191	))
192	.layer(
193	tower_http::set_header::SetResponseHeaderLayer::if_not_present(
194	axum::http::header::CACHE_CONTROL,
195	axum::http::HeaderValue::from_static("private, no-cache"),
196	),
197	)
198	// Outermost: a hard per-request deadline so a hung handler (e.g. a slow
199	// query that outlives the pool acquire, or a stuck await) sheds instead of
200	// pinning a connection forever. 30s clears any legitimate request, every
201	// outbound call is bounded to 5-15s and the DB acquire to 10s; the only
202	// request that transfers meaningful bytes is an image upload, capped at
203	// ~5 MB (needs ~1.4 Mbit/s to finish inside the window). Returns 408.
204	.layer(tower_http::timeout::TimeoutLayer::with_status_code(
205	axum::http::StatusCode::REQUEST_TIMEOUT,
206	std::time::Duration::from_secs(30),
207	));
208
209	// Default to loopback. Rate limiting uses TrustedProxyKeyExtractor, which
210	// honors CF-Connecting-IP / X-Forwarded-For only from a configured trusted
211	// proxy (TRUSTED_PROXIES, default loopback) and otherwise keys on the direct
212	// peer, so spoofed forwarding headers from a non-proxy client are ignored
213	// regardless of bind address. Binding to
214	// 127.0.0.1 is still the right default (only Caddy reaches the port); a
215	// tailnet-direct staging box sets HOST and lists its proxy in TRUSTED_PROXIES.
216	let host = std::env::var("HOST").unwrap_or_else(\|_\| "127.0.0.1".to_string());
217	let port = std::env::var("PORT").unwrap_or_else(\|_\| "3400".to_string());
218	let addr = format!("{host}:{port}");
219
220	let listener = TcpListener::bind(&addr).await.expect("failed to bind");
221
222	tracing::info!("listening on {}", listener.local_addr().unwrap());
223
224	axum::serve(
225	listener,
226	app.into_make_service_with_connect_info::<std::net::SocketAddr>(),
227	)
228	.with_graceful_shutdown(shutdown_signal())
229	.await
230	.expect("server error");
231
232	deletion_task.abort();
233	let _ = deletion_task.await;
234	if let Some(task) = purge_task {
235	task.abort();
236	let _ = task.await;
237	}
238	for task in [chat_sweep_task, chat_rate_task] {
239	task.abort();
240	let _ = task.await;
241	}
242	}
243
244	/// Spawn a never-returning background loop under panic supervision.
245	///
246	/// The inner future is run in its own task so a panic surfaces here as a
247	/// `JoinError` instead of silently killing the worker; on a panic, or an
248	/// unexpected normal return, it is logged and restarted after a short backoff.
249	/// `factory` is `Fn` so it can rebuild the future (cloning any captured handles)
250	/// on each restart. Aborting the returned handle (at shutdown) stops the loop.
251	fn supervise<F, Fut>(name: &'static str, factory: F) -> tokio::task::JoinHandle<()>
252	where
253	F: Fn() -> Fut + Send + 'static,
254	Fut: std::future::Future + Send + 'static,
255	Fut::Output: Send,
256	{
257	const RESTART_BACKOFF: std::time::Duration = std::time::Duration::from_secs(5);
258	tokio::task::spawn(async move {
259	loop {
260	match tokio::task::spawn(factory()).await {
261	Ok(_) => {
262	tracing::error!(
263	task = name,
264	"background task returned unexpectedly; restarting"
265	);
266	}
267	Err(e) if e.is_panic() => {
268	tracing::error!(task = name, "background task panicked; restarting");
269	}
270	// Cancelled: the supervisor itself is being aborted at shutdown.
271	Err(_) => return,
272	}
273	tokio::time::sleep(RESTART_BACKOFF).await;
274	}
275	})
276	}
277
278	async fn shutdown_signal() {
279	use tokio::signal;
280	let ctrl_c = async {
281	signal::ctrl_c()
282	.await
283	.expect("failed to install Ctrl+C handler");
284	};
285	#[cfg(unix)]
286	let terminate = async {
287	signal::unix::signal(signal::unix::SignalKind::terminate())
288	.expect("failed to install signal handler")
289	.recv()
290	.await;
291	};
292	#[cfg(not(unix))]
293	let terminate = std::future::pending::<()>();
294	tokio::select! {
295	() = ctrl_c => {},
296	() = terminate => {},
297	}
298	tracing::info!("Shutdown signal received");
299	}
300