Skip to main content

max / makenotwork

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