atagen

malice

public
Star0
Watch0
e153508681c5c91d170d860a395f773ef22afa18
2
ssh://brewery@pls.lobotomise.me:33/atagen/malice.git
crates/malice-static/src/lib.rsRust163 lines5.3 KB
1
pub mod idle;
2
pub mod push;
3
4
use std::{future::Future, net::SocketAddr, time::Duration};
5
6
use axum::{
7
    Router,
8
    http::{HeaderName, HeaderValue, StatusCode, header},
9
    middleware,
10
    middleware::Next,
11
    response::{IntoResponse, Response},
12
    routing::any,
13
};
14
use listenfd::ListenFd;
15
use tower_http::{
16
    services::{ServeDir, ServeFile},
17
    trace::TraceLayer,
18
};
19
20
use crate::idle::IdleTracker;
21
22
fn bind_addr() -> SocketAddr {
23
    std::env::var("MALICE_BIND")
24
        .unwrap_or_else(|_| "127.0.0.1:3000".to_owned())
25
        .parse()
26
        .expect("MALICE_BIND must be host:port")
27
}
28
29
fn idle_exit_secs() -> Option<u64> {
30
    std::env::var("MALICE_IDLE_EXIT_SECS")
31
        .ok()
32
        .and_then(|value| value.trim().parse().ok())
33
}
34
35
/// Takes over a systemd-activated listener from fd 3 when present, so the
36
/// socket unit can hold the port before the process ever starts. Falls back
37
/// to a plain bind otherwise.
38
pub async fn bind_listener() -> tokio::net::TcpListener {
39
    let mut listenfd = ListenFd::from_env();
40
    if let Some(listener) = listenfd
41
        .take_tcp_listener(0)
42
        .expect("inspecting inherited sockets")
43
    {
44
        listener
45
            .set_nonblocking(true)
46
            .expect("setting inherited socket nonblocking");
47
        return tokio::net::TcpListener::from_std(listener)
48
            .expect("adopting inherited socket into tokio");
49
    }
50
    let addr = bind_addr();
51
    tracing::info!("listening on {addr}");
52
    tokio::net::TcpListener::bind(addr)
53
        .await
54
        .expect("binding the listener must succeed")
55
}
56
57
/// Serves `app` on `listener`, exiting on ctrl-c or, when
58
/// `MALICE_IDLE_EXIT_SECS` is set, after that long with no request in flight
59
/// or started. `on_shutdown` runs after the exit signal either way.
60
pub async fn serve(
61
    listener: tokio::net::TcpListener,
62
    app: Router,
63
    on_shutdown: impl Future<Output = ()> + Send + 'static,
64
) {
65
    let idle = idle_exit_secs().map(|secs| (IdleTracker::new(), secs));
66
    let app = match &idle {
67
        Some((tracker, _)) => {
68
            app.layer(middleware::from_fn_with_state(tracker.clone(), idle::track))
69
        }
70
        None => app,
71
    };
72
    axum::serve(listener, app)
73
        .with_graceful_shutdown(async move {
74
            match idle {
75
                Some((tracker, secs)) => {
76
                    tokio::select! {
77
                        _ = tokio::signal::ctrl_c() => {}
78
                        _ = tracker.wait_idle(Duration::from_secs(secs)) => {
79
                            tracing::info!("idle for {secs}s with no requests, exiting");
80
                        }
81
                    }
82
                }
83
                None => {
84
                    let _ = tokio::signal::ctrl_c().await;
85
                }
86
            }
87
            on_shutdown.await;
88
        })
89
        .await
90
        .expect("serving must succeed");
91
}
92
93
async fn api_not_found() -> StatusCode {
94
    StatusCode::NOT_FOUND
95
}
96
97
async fn missing_frontend() -> Response {
98
    (
99
        StatusCode::SERVICE_UNAVAILABLE,
100
        "frontend assets missing; set MALICE_WEB_DIR or use the Nix package",
101
    )
102
        .into_response()
103
}
104
105
fn is_no_cache_path(path: &str) -> bool {
106
    matches!(path, "/" | "/index.html" | "/sw.js" | "/manifest.webmanifest")
107
}
108
109
fn is_immutable_path(path: &str) -> bool {
110
    path.starts_with("/assets/") || path.starts_with("/twemoji/")
111
}
112
113
/// index.html and sw.js must never be cached, or a stale service worker or
114
/// app shell survives a rebuild. Content-hashed assets are cached forever.
115
pub async fn static_headers(request: axum::extract::Request, next: Next) -> Response {
116
    let no_cache = is_no_cache_path(request.uri().path());
117
    let immutable = is_immutable_path(request.uri().path());
118
    let mut request = request;
119
    if no_cache {
120
        let headers = request.headers_mut();
121
        headers.remove(header::IF_MODIFIED_SINCE);
122
        headers.remove(header::IF_NONE_MATCH);
123
    }
124
    let mut response = next.run(request).await;
125
    response.headers_mut().remove(header::LAST_MODIFIED);
126
    if immutable && response.status().is_success() {
127
        response.headers_mut().insert(
128
            header::CACHE_CONTROL,
129
            HeaderValue::from_static("public, max-age=31536000, immutable"),
130
        );
131
    }
132
    if no_cache {
133
        let headers = response.headers_mut();
134
        headers.remove(header::ETAG);
135
        headers.insert(header::CACHE_CONTROL, HeaderValue::from_static("no-store"));
136
        headers.insert(
137
            HeaderName::from_static("service-worker-allowed"),
138
            HeaderValue::from_static("/"),
139
        );
140
    }
141
    response
142
}
143
144
/// Serves only the built frontend: no `/api` routes, no session state, no
145
/// data dir. Used for the socket-activated on-device build.
146
pub fn static_router(web_dir: Option<String>) -> Router {
147
    let router = Router::new()
148
        .route("/api", any(api_not_found))
149
        .route("/api/{*rest}", any(api_not_found))
150
        .layer(TraceLayer::new_for_http());
151
    match web_dir {
152
        Some(dir) => {
153
            let files = ServeDir::new(&dir)
154
                .precompressed_br()
155
                .precompressed_gzip()
156
                .not_found_service(ServeFile::new(format!("{dir}/index.html")));
157
            router
158
                .fallback_service(files)
159
                .layer(middleware::from_fn(static_headers))
160
        }
161
        None => router.fallback(missing_frontend),
162
    }
163
}