folk_plugin_http/
server.rs

1use std::net::{IpAddr, SocketAddr};
2use std::sync::Arc;
3use std::sync::atomic::{AtomicU64, Ordering};
4use std::time::Instant;
5
6use anyhow::Result;
7use axum::Router;
8use axum::body::Body;
9use axum::extract::{ConnectInfo, State};
10use axum::http::{self, Request, Response};
11use axum::routing::any;
12use ipnet::IpNet;
13use tokio::net::TcpListener;
14use tokio::sync::watch;
15use tower::ServiceBuilder;
16use tower_http::limit::RequestBodyLimitLayer;
17use tower_http::timeout::TimeoutLayer;
18use tracing::{error, info};
19
20use crate::config::HttpConfig;
21use crate::payload::{decode_response, encode_request};
22
23#[derive(Clone)]
24struct AppState {
25    executor: Arc<dyn folk_api::Executor>,
26    config: Arc<HttpConfig>,
27    active_connections: Arc<AtomicU64>,
28}
29
30pub struct HttpServer {
31    config: HttpConfig,
32    executor: Arc<dyn folk_api::Executor>,
33    active_connections: Arc<AtomicU64>,
34}
35
36impl HttpServer {
37    pub fn new(
38        config: HttpConfig,
39        executor: Arc<dyn folk_api::Executor>,
40        active_connections: Arc<AtomicU64>,
41    ) -> Self {
42        Self {
43            config,
44            executor,
45            active_connections,
46        }
47    }
48
49    pub async fn run(self, shutdown: watch::Receiver<bool>) -> Result<()> {
50        let state = AppState {
51            executor: self.executor.clone(),
52            config: Arc::new(self.config.clone()),
53            active_connections: self.active_connections.clone(),
54        };
55
56        let mut app = Router::new()
57            .route("/{*path}", any(handle))
58            .route("/", any(handle))
59            .with_state(state)
60            .layer(
61                ServiceBuilder::new()
62                    .layer(RequestBodyLimitLayer::new(self.config.max_request_size))
63                    .layer(TimeoutLayer::with_status_code(
64                        http::StatusCode::GATEWAY_TIMEOUT,
65                        self.config.write_timeout,
66                    )),
67            );
68
69        if self.config.compression.enabled {
70            app = app.layer(build_compression_layer(&self.config.compression));
71        }
72
73        #[cfg(feature = "tls")]
74        if let Some(ref tls) = self.config.tls {
75            return self.run_tls(app, tls, shutdown).await;
76        }
77
78        #[cfg(feature = "h2c")]
79        if self.config.h2c {
80            return self.run_h2c(app, shutdown).await;
81        }
82
83        self.run_plain(app, shutdown).await
84    }
85
86    async fn run_plain(&self, app: Router, shutdown: watch::Receiver<bool>) -> Result<()> {
87        let listener = TcpListener::bind(self.config.listen).await?;
88
89        axum::serve(
90            listener,
91            app.into_make_service_with_connect_info::<SocketAddr>(),
92        )
93        .with_graceful_shutdown(shutdown_signal(shutdown))
94        .await?;
95
96        Ok(())
97    }
98
99    #[cfg(feature = "tls")]
100    async fn run_tls(
101        &self,
102        app: Router,
103        tls: &crate::config::TlsConfig,
104        shutdown: watch::Receiver<bool>,
105    ) -> Result<()> {
106        use axum_server::Handle;
107        use axum_server::tls_rustls::RustlsConfig;
108
109        let rustls_config = RustlsConfig::from_pem_file(&tls.cert, &tls.key).await?;
110
111        info!(cert = %tls.cert.display(), "TLS enabled");
112
113        let handle = Handle::new();
114        let shutdown_handle = handle.clone();
115        tokio::spawn(async move {
116            shutdown_signal(shutdown).await;
117            shutdown_handle.graceful_shutdown(None);
118        });
119
120        axum_server::bind_rustls(self.config.listen, rustls_config)
121            .handle(handle)
122            .serve(app.into_make_service_with_connect_info::<SocketAddr>())
123            .await?;
124
125        Ok(())
126    }
127
128    #[cfg(feature = "h2c")]
129    async fn run_h2c(&self, app: Router, mut shutdown: watch::Receiver<bool>) -> Result<()> {
130        use hyper_util::rt::{TokioExecutor, TokioIo};
131        use hyper_util::server::conn::auto::Builder as AutoBuilder;
132
133        info!("h2c (HTTP/2 cleartext) enabled");
134
135        let listener = TcpListener::bind(self.config.listen).await?;
136        let builder = Arc::new(AutoBuilder::new(TokioExecutor::new()));
137
138        loop {
139            tokio::select! {
140                result = listener.accept() => {
141                    let (stream, remote_addr) = result?;
142                    let app = app.clone();
143                    let builder = builder.clone();
144                    tokio::spawn(async move {
145                        let svc = hyper::service::service_fn(move |mut req: Request<hyper::body::Incoming>| {
146                            // Inject ConnectInfo manually since we bypass axum::serve.
147                            req.extensions_mut().insert(ConnectInfo(remote_addr));
148                            let app = app.clone();
149                            async move {
150                                let resp = tower::Service::call(&mut app.clone(), req).await;
151                                resp.map_err(|e| match e {})
152                            }
153                        });
154                        let _ = builder.serve_connection_with_upgrades(TokioIo::new(stream), svc).await;
155                    });
156                }
157                _ = async {
158                    loop {
159                        if shutdown.changed().await.is_err() || *shutdown.borrow() {
160                            break;
161                        }
162                    }
163                } => {
164                    break;
165                }
166            }
167        }
168
169        Ok(())
170    }
171}
172
173async fn shutdown_signal(mut shutdown: watch::Receiver<bool>) {
174    loop {
175        if shutdown.changed().await.is_err() || *shutdown.borrow() {
176            break;
177        }
178    }
179}
180
181async fn handle(
182    State(state): State<AppState>,
183    connect_info: ConnectInfo<SocketAddr>,
184    req: Request<Body>,
185) -> Response<Body> {
186    state.active_connections.fetch_add(1, Ordering::Relaxed);
187    let start = Instant::now();
188    let method = req.method().clone();
189    let uri = req.uri().clone();
190    let peer_addr = connect_info.0;
191
192    let client_ip = resolve_client_ip(
193        peer_addr.ip(),
194        req.headers()
195            .get("x-forwarded-for")
196            .and_then(|v| v.to_str().ok()),
197        &state.config.trusted_proxies,
198    );
199
200    let response = handle_inner(&state, req).await;
201
202    if state.config.access_log {
203        let duration = start.elapsed();
204        let status = response.status().as_u16();
205        info!(
206            client_ip = %client_ip,
207            method = %method,
208            uri = %uri,
209            status = status,
210            duration_ms = duration.as_millis() as u64,
211            "http request",
212        );
213    }
214
215    state.active_connections.fetch_sub(1, Ordering::Relaxed);
216    response
217}
218
219async fn handle_inner(state: &AppState, req: Request<Body>) -> Response<Body> {
220    let max_body = state.config.max_request_size;
221    let payload = match encode_request(req, max_body).await {
222        Ok(p) => p,
223        Err(e) => {
224            error!(error = ?e, "encode request");
225            return Response::builder()
226                .status(500)
227                .body(Body::from("encode error"))
228                .unwrap();
229        }
230    };
231
232    let response_value = match state.executor.execute_value("http.handle", payload).await {
233        Ok(v) => v,
234        Err(e) => {
235            error!(error = ?e, "dispatch to worker");
236            return Response::builder()
237                .status(502)
238                .body(Body::from("worker error"))
239                .unwrap();
240        }
241    };
242
243    match decode_response(response_value) {
244        Ok(r) => r,
245        Err(e) => {
246            error!(error = ?e, "decode response");
247            Response::builder()
248                .status(500)
249                .body(Body::from("decode error"))
250                .unwrap()
251        }
252    }
253}
254
255/// Resolve the real client IP from X-Forwarded-For if the peer is a trusted proxy.
256///
257/// Walks the X-Forwarded-For chain from right to left, stopping at the first
258/// IP that is NOT in a trusted subnet. This is the standard secure algorithm
259/// (rightmost non-trusted).
260pub fn resolve_client_ip(peer_ip: IpAddr, xff: Option<&str>, trusted: &[IpNet]) -> IpAddr {
261    if trusted.is_empty() {
262        return peer_ip;
263    }
264
265    if !is_trusted(peer_ip, trusted) {
266        return peer_ip;
267    }
268
269    let Some(xff) = xff else {
270        return peer_ip;
271    };
272
273    let addrs: Vec<&str> = xff.split(',').map(|s| s.trim()).collect();
274
275    // Walk from right to left — the rightmost non-trusted IP is the client.
276    for addr_str in addrs.iter().rev() {
277        if let Ok(ip) = addr_str.parse::<IpAddr>() {
278            if !is_trusted(ip, trusted) {
279                return ip;
280            }
281        }
282    }
283
284    // All IPs in the chain are trusted — use peer IP.
285    peer_ip
286}
287
288fn is_trusted(ip: IpAddr, trusted: &[IpNet]) -> bool {
289    trusted.iter().any(|net| net.contains(&ip))
290}
291
292fn build_compression_layer(
293    config: &crate::config::CompressionConfig,
294) -> tower_http::compression::CompressionLayer {
295    use crate::config::CompressionAlgorithm;
296    use tower_http::compression::CompressionLayer;
297
298    let mut layer = CompressionLayer::new()
299        .no_gzip()
300        .no_br()
301        .no_zstd()
302        .no_deflate();
303
304    for algo in &config.algorithms {
305        layer = match algo {
306            CompressionAlgorithm::Gzip => layer.gzip(true),
307            CompressionAlgorithm::Br => layer.br(true),
308            CompressionAlgorithm::Zstd => layer.zstd(true),
309            CompressionAlgorithm::Deflate => layer.deflate(true),
310        };
311    }
312
313    layer
314}
folk_plugin_http/server.rs

folk_plugin_http/
server.rs
