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        let response_bytes = response
206            .headers()
207            .get(http::header::CONTENT_LENGTH)
208            .and_then(|v| v.to_str().ok())
209            .and_then(|v| v.parse::<u64>().ok())
210            .unwrap_or(0);
211        info!(
212            client_ip = %client_ip,
213            method = %method,
214            uri = %uri,
215            status = status,
216            duration_ms = duration.as_millis() as u64,
217            response_bytes = response_bytes,
218            "http request",
219        );
220    }
221
222    state.active_connections.fetch_sub(1, Ordering::Relaxed);
223    response
224}
225
226async fn handle_inner(state: &AppState, req: Request<Body>) -> Response<Body> {
227    let max_body = state.config.max_request_size;
228    let payload = match encode_request(req, max_body).await {
229        Ok(p) => p,
230        Err(e) => {
231            error!(error = ?e, "encode request");
232            return Response::builder()
233                .status(500)
234                .body(Body::from("encode error"))
235                .unwrap();
236        }
237    };
238
239    let response_value = match state.executor.execute_value("http.handle", payload).await {
240        Ok(v) => v,
241        Err(e) => {
242            error!(error = ?e, "dispatch to worker");
243            return Response::builder()
244                .status(502)
245                .body(Body::from("worker error"))
246                .unwrap();
247        }
248    };
249
250    match decode_response(response_value) {
251        Ok(r) => r,
252        Err(e) => {
253            error!(error = ?e, "decode response");
254            Response::builder()
255                .status(500)
256                .body(Body::from("decode error"))
257                .unwrap()
258        }
259    }
260}
261
262/// Resolve the real client IP from X-Forwarded-For if the peer is a trusted proxy.
263///
264/// Walks the X-Forwarded-For chain from right to left, stopping at the first
265/// IP that is NOT in a trusted subnet. This is the standard secure algorithm
266/// (rightmost non-trusted).
267pub fn resolve_client_ip(peer_ip: IpAddr, xff: Option<&str>, trusted: &[IpNet]) -> IpAddr {
268    if trusted.is_empty() {
269        return peer_ip;
270    }
271
272    if !is_trusted(peer_ip, trusted) {
273        return peer_ip;
274    }
275
276    let Some(xff) = xff else {
277        return peer_ip;
278    };
279
280    let addrs: Vec<&str> = xff.split(',').map(|s| s.trim()).collect();
281
282    // Walk from right to left — the rightmost non-trusted IP is the client.
283    for addr_str in addrs.iter().rev() {
284        if let Ok(ip) = addr_str.parse::<IpAddr>() {
285            if !is_trusted(ip, trusted) {
286                return ip;
287            }
288        }
289    }
290
291    // All IPs in the chain are trusted — use peer IP.
292    peer_ip
293}
294
295fn is_trusted(ip: IpAddr, trusted: &[IpNet]) -> bool {
296    trusted.iter().any(|net| net.contains(&ip))
297}
298
299fn build_compression_layer(
300    config: &crate::config::CompressionConfig,
301) -> tower_http::compression::CompressionLayer {
302    use crate::config::CompressionAlgorithm;
303    use tower_http::compression::CompressionLayer;
304
305    let mut layer = CompressionLayer::new()
306        .no_gzip()
307        .no_br()
308        .no_zstd()
309        .no_deflate();
310
311    for algo in &config.algorithms {
312        layer = match algo {
313            CompressionAlgorithm::Gzip => layer.gzip(true),
314            CompressionAlgorithm::Br => layer.br(true),
315            CompressionAlgorithm::Zstd => layer.zstd(true),
316            CompressionAlgorithm::Deflate => layer.deflate(true),
317        };
318    }
319
320    layer
321}
folk_plugin_http/server.rs

folk_plugin_http/
server.rs
