use axum::{ Router, extract::{ConnectInfo, Path, State}, http::{HeaderMap, HeaderValue, StatusCode, header}, response::{IntoResponse, Response}, routing::get, }; use chrono::{DateTime, Utc}; use libservice::ServiceConnectInfo; use lru::LruCache; use maud::Markup; use parking_lot::Mutex; use std::{collections::HashMap, num::NonZero, pin::Pin, sync::Arc, time::Instant}; use tower_http::compression::{CompressionLayer, DefaultPredicate}; use tracing::{trace, warn}; use crate::{ClientInfo, RequestContext, page::Page}; pub struct PageServer { /// If true, expired pages will be rerendered before being sent to the user. /// If false, requests never trigger rerenders. We rely on the rerender task. /// /// If true, we deliver fresher pages but delay responses. /// TODO: replace this with a smarter rendering strategy? never_rerender_on_request: bool, /// Map of `{ route: page }` pages: Arc>>>, /// Map of `{ route: (page data, expire time) }` /// /// We use an LruCache for bounded memory usage. html_cache: Mutex)>>, /// Called whenever we need to render a page. /// - this method should call `page.generate_html()`, /// - wrap the result in ``, /// - and add `` /// ``` render_page: Box< dyn Send + Sync + for<'a> Fn( &'a Page, &'a RequestContext, ) -> Pin + 'a + Send + Sync>>, >, } impl PageServer { pub fn new( render_page: Box< dyn Send + Sync + for<'a> Fn( &'a Page, &'a RequestContext, ) -> Pin + 'a + Send + Sync>>, >, ) -> Arc { #[expect(clippy::unwrap_used)] let cache_size = NonZero::new(128).unwrap(); Arc::new(Self { pages: Arc::new(Mutex::new(HashMap::new())), html_cache: Mutex::new(LruCache::new(cache_size)), render_page, never_rerender_on_request: true, }) } pub fn add_page(&self, route: impl Into, page: Page) -> &Self { #[expect(clippy::expect_used)] let route = route .into() .strip_prefix("/") .expect("page route must start with /") .to_owned(); self.pages.lock().insert(route, Arc::new(page)); self } /// Re-render the page at `route`, regardless of cache state. /// Does nothing if there is no page at `route`. /// /// Returns the rendered page's content. async fn render_page( &self, reason: &'static str, route: &str, req_ctx: RequestContext, ) -> Option<(String, Option>)> { let now = Utc::now(); let start = Instant::now(); let page = match self.pages.lock().get(route) { Some(x) => x.clone(), None => return None, }; trace!( message = "Rendering page", route, reason, lock_time_ms = start.elapsed().as_millis() ); let html = (self.render_page)(&page, &req_ctx).await.0; let mut expires = None; if let Some(ttl) = page.html_ttl { expires = Some(now + ttl); self.html_cache .lock() .put((route.to_owned(), req_ctx), (html.clone(), now + ttl)); } let elapsed = start.elapsed().as_millis(); trace!(message = "Rendered page", route, reason, time_ms = elapsed); return Some((html, expires)); } async fn handler( Path(route): Path, State(state): State>, ConnectInfo(addr): ConnectInfo, headers: HeaderMap, ) -> Response { let client_info = ClientInfo::from_headers(&headers); let ua = headers .get("user-agent") .and_then(|x| x.to_str().ok()) .unwrap_or(""); trace!( message = "Serving route", route, addr = ?addr.addr, user_agent = ua, device_type = ?client_info.device_type ); // Normalize url with redirect if route.ends_with('/') || route.contains("//") || route.starts_with('/') { let mut new_route = route.clone(); while new_route.contains("//") { new_route = new_route.replace("//", "/"); } let new_route = new_route.trim_matches('/'); trace!( message = "Redirecting", route, new_route, addr = ?addr.addr, user_agent = ua, device_type = ?client_info.device_type ); let mut headers = HeaderMap::with_capacity(2); let new_route = match HeaderValue::from_str(&format!("/{new_route}")) { Ok(x) => x, Err(_) => { // Be extra careful, this is user-provided data return StatusCode::BAD_REQUEST.into_response(); } }; headers.append(header::LOCATION, new_route); headers.append("Accept-CH", HeaderValue::from_static("Sec-CH-UA-Mobile")); return (StatusCode::PERMANENT_REDIRECT, headers).into_response(); } let req_ctx = RequestContext { client_info, route: format!("/{route}"), }; let cache_key = (route.clone(), req_ctx.clone()); let now = Utc::now(); let mut html_expires = None; // Get from cache, if available if let Some((html, expires)) = state.html_cache.lock().get(&cache_key) && (*expires > now || state.never_rerender_on_request) { html_expires = Some((html.clone(), Some(*expires))); }; if html_expires.is_none() { html_expires = match state.render_page("request", &route, req_ctx).await { Some(x) => Some(x.clone()), None => { trace!( message = "Not found", route, addr = ?addr.addr, user_agent = ua, device_type = ?client_info.device_type ); return (StatusCode::NOT_FOUND, "page doesn't exist").into_response(); } }; } #[expect(clippy::unwrap_used)] let (html, expires) = html_expires.unwrap(); let mut headers = HeaderMap::with_capacity(3); headers.append( header::CONTENT_TYPE, HeaderValue::from_static("text/html; charset=utf-8"), ); let max_age = match expires { Some(expires) => (expires - now).num_seconds().max(1), None => 1, }; #[expect(clippy::unwrap_used)] headers.append( header::CACHE_CONTROL, // immutable; public/private HeaderValue::from_str(&format!("immutable, public, max-age={}", max_age)).unwrap(), ); headers.append("Accept-CH", HeaderValue::from_static("Sec-CH-UA-Mobile")); return (headers, html).into_response(); } pub fn into_router(self: Arc) -> Router<()> { let compression: CompressionLayer = CompressionLayer::new() .br(true) .deflate(true) .gzip(true) .zstd(true) .compress_when(DefaultPredicate::new()); Router::new() .route( "/", get(|state, conn, headers| async { Self::handler(Path(String::new()), state, conn, headers).await }), ) .route("/{*path}", get(Self::handler)) .layer(compression) .with_state(self) } }