Skip to main content

a3s_boot/adapters/
axum.rs

1use crate::{
2    BootApplication, BootError, BootRequest, BootResponse, HttpAdapter, HttpMethod, Result,
3    RouteDefinition, WebSocketGatewayDefinition, WebSocketMessage,
4};
5use axum::body::{to_bytes, Body, HttpBody};
6use axum::extract::ws::{Message as AxumWebSocketMessage, WebSocket, WebSocketUpgrade};
7use axum::extract::Request;
8use axum::http::{
9    header::{ALLOW, CONTENT_TYPE},
10    response::Builder as ResponseBuilder,
11    HeaderMap, HeaderName, HeaderValue, Method, StatusCode, Uri,
12};
13use axum::response::Response;
14use axum::routing::{on, MethodFilter, MethodRouter};
15use axum::Router;
16use futures_util::StreamExt;
17use std::error::Error;
18use std::net::SocketAddr;
19
20/// Axum-backed HTTP adapter.
21///
22/// Axum is the default adapter, not the Boot kernel. Applications can replace
23/// it by implementing [`HttpAdapter`] for another backend.
24#[derive(Debug, Clone)]
25pub struct AxumAdapter {
26    body_limit: usize,
27}
28
29impl AxumAdapter {
30    pub fn new() -> Self {
31        Self {
32            body_limit: 1024 * 1024,
33        }
34    }
35
36    pub fn with_body_limit(mut self, body_limit: usize) -> Self {
37        self.body_limit = body_limit;
38        self
39    }
40}
41
42impl Default for AxumAdapter {
43    fn default() -> Self {
44        Self::new()
45    }
46}
47
48impl HttpAdapter for AxumAdapter {
49    type Output = Router;
50
51    fn build(&self, app: BootApplication) -> Result<Self::Output> {
52        let fallback_routing = app.api_versioning().is_some()
53            || app.routes().iter().any(|route| {
54                route.host().is_some() || route.method().is_wildcard() || route.has_catch_all_path()
55            });
56        let mut router = if fallback_routing {
57            let body_limit = self.body_limit;
58            let app = app.clone();
59            Router::new().fallback(move |request: Request| {
60                let app = app.clone();
61                async move { dispatch_application(app, request, body_limit).await }
62            })
63        } else {
64            Router::new().fallback(not_found_fallback)
65        };
66
67        if !fallback_routing {
68            for route in app.routes().iter().cloned() {
69                let path = axum_route_path(route.path());
70                router = router.route(
71                    &path,
72                    route_to_method_router(route, app.clone(), self.body_limit),
73                );
74            }
75        }
76
77        for gateway in app.gateways().iter().cloned() {
78            let path = axum_route_path(gateway.path());
79            router = router.route(&path, gateway_to_method_router(gateway));
80        }
81        let body_limit = self.body_limit;
82        let app = app.clone();
83        Ok(router.method_not_allowed_fallback(move |request: Request| {
84            let app = app.clone();
85            async move { dispatch_application(app, request, body_limit).await }
86        }))
87    }
88
89    fn serve(
90        &self,
91        app: BootApplication,
92        addr: SocketAddr,
93    ) -> crate::BoxFuture<'static, Result<()>> {
94        let adapter = self.clone();
95        Box::pin(async move {
96            let router = adapter.build(app)?;
97            let listener = tokio::net::TcpListener::bind(addr).await?;
98            axum::serve(listener, router).await?;
99            Ok(())
100        })
101    }
102}
103
104fn gateway_to_method_router(gateway: WebSocketGatewayDefinition) -> MethodRouter {
105    axum::routing::get(move |ws: WebSocketUpgrade, uri: Uri, headers: HeaderMap| {
106        let gateway = gateway.clone();
107        async move {
108            ws.on_upgrade(move |socket| async move {
109                handle_websocket(socket, gateway, uri, headers).await;
110            })
111        }
112    })
113}
114
115fn route_to_method_router(
116    route: RouteDefinition,
117    app: BootApplication,
118    body_limit: usize,
119) -> MethodRouter {
120    let method = route.method();
121    let dispatch = move |request: Request| {
122        let route = route.clone();
123        let app = app.clone();
124        async move { dispatch_route(route, app, request, body_limit).await }
125    };
126
127    on(method_filter(method), dispatch)
128}
129
130fn method_filter(method: HttpMethod) -> MethodFilter {
131    match method {
132        HttpMethod::All => MethodFilter::GET
133            .or(MethodFilter::POST)
134            .or(MethodFilter::PUT)
135            .or(MethodFilter::PATCH)
136            .or(MethodFilter::DELETE)
137            .or(MethodFilter::OPTIONS)
138            .or(MethodFilter::HEAD),
139        HttpMethod::Get => MethodFilter::GET,
140        HttpMethod::Post => MethodFilter::POST,
141        HttpMethod::Put => MethodFilter::PUT,
142        HttpMethod::Patch => MethodFilter::PATCH,
143        HttpMethod::Delete => MethodFilter::DELETE,
144        HttpMethod::Options => MethodFilter::OPTIONS,
145        HttpMethod::Head => MethodFilter::HEAD,
146    }
147}
148
149fn axum_route_path(path: &str) -> String {
150    let path = path.strip_prefix('/').unwrap_or(path);
151    if path.is_empty() {
152        return "/".to_string();
153    }
154
155    let segments = path
156        .split('/')
157        .enumerate()
158        .map(|(index, segment)| {
159            if segment.starts_with("{*") && segment.ends_with('}') {
160                format!("{{*p{index}}}")
161            } else if segment.starts_with('{') && segment.ends_with('}') {
162                format!("{{p{index}}}")
163            } else {
164                segment.to_string()
165            }
166        })
167        .collect::<Vec<_>>();
168
169    format!("/{}", segments.join("/"))
170}
171
172async fn dispatch_route(
173    route: RouteDefinition,
174    app: BootApplication,
175    request: Request,
176    body_limit: usize,
177) -> Response {
178    let path = request.uri().path().to_string();
179    let is_head = request.method() == Method::HEAD;
180    let boot_request = match to_boot_request(request, body_limit).await {
181        Ok(request) => request,
182        Err(err) => return finalize_response(&app, &path, boot_error_response(err), is_head),
183    };
184
185    let response = match route.call(boot_request).await {
186        Ok(response) => to_axum_response(response),
187        Err(err) => boot_error_response(err),
188    };
189    finalize_response(&app, &path, response, is_head)
190}
191
192async fn dispatch_application(
193    app: BootApplication,
194    request: Request,
195    body_limit: usize,
196) -> Response {
197    let path = request.uri().path().to_string();
198    let is_head = request.method() == Method::HEAD;
199    let boot_request = match to_boot_request(request, body_limit).await {
200        Ok(request) => request,
201        Err(err) => return finalize_response(&app, &path, boot_error_response(err), is_head),
202    };
203
204    let response = to_axum_response(app.handle(boot_request).await);
205    finalize_response(&app, &path, response, is_head)
206}
207
208async fn not_found_fallback(request: Request) -> Response {
209    let is_head = request.method() == Method::HEAD;
210    let response = boot_error_response(BootError::NotFound(format!(
211        "{} {}",
212        request.method(),
213        request.uri().path()
214    )));
215    strip_head_body(is_head, response)
216}
217
218async fn to_boot_request(axum_request: Request, body_limit: usize) -> Result<BootRequest> {
219    let path = axum_request.uri().path().to_string();
220    let method =
221        HttpMethod::try_from(axum_request.method().clone()).map_err(|error| match error {
222            BootError::MethodNotAllowed(method) => {
223                BootError::MethodNotAllowed(format!("{method} {path}"))
224            }
225            error => error,
226        })?;
227    let query_string = axum_request.uri().query().map(str::to_string);
228    let mut headers = Vec::new();
229    for (name, value) in axum_request.headers() {
230        let value = value.to_str().map_err(|err| {
231            BootError::BadRequest(format!("invalid request header value for {name}: {err}"))
232        })?;
233        headers.push((name.as_str().to_string(), value.to_string()));
234    }
235    let mut boot_request = BootRequest::new(method, path);
236    if let Some(query_string) = query_string {
237        boot_request = boot_request.with_query_string(query_string);
238    }
239    for (name, value) in headers {
240        boot_request = if boot_request.header(&name).is_some() {
241            boot_request.append_header(name, value)
242        } else {
243            boot_request.with_header(name, value)
244        };
245    }
246
247    boot_request.validate_headers()?;
248    boot_request.validate_body_limit(body_limit)?;
249
250    let body = axum_request.into_body();
251    if body.size_hint().lower() > body_limit as u64 {
252        return Err(BootError::PayloadTooLarge(format!(
253            "request body exceeds {body_limit} bytes"
254        )));
255    }
256    let body = to_bytes(body, body_limit)
257        .await
258        .map_err(|err| map_body_error(err, body_limit))?
259        .to_vec();
260
261    let boot_request = boot_request.with_body(body);
262    boot_request.validate_with_body_limit(body_limit)?;
263    Ok(boot_request)
264}
265
266fn websocket_boot_request(uri: Uri, headers: HeaderMap) -> Result<BootRequest> {
267    let mut request = BootRequest::new(HttpMethod::Get, uri.path().to_string());
268    if let Some(query) = uri.query() {
269        request = request.with_query_string(query.to_string());
270    }
271    for (name, value) in &headers {
272        let value = value.to_str().map_err(|err| {
273            BootError::BadRequest(format!("invalid request header value for {name}: {err}"))
274        })?;
275        request = if request.header(name.as_str()).is_some() {
276            request.append_header(name.as_str(), value)
277        } else {
278            request.with_header(name.as_str(), value)
279        };
280    }
281    request.validate_headers()?;
282    Ok(request)
283}
284
285async fn handle_websocket(
286    mut socket: WebSocket,
287    gateway: WebSocketGatewayDefinition,
288    uri: Uri,
289    headers: HeaderMap,
290) {
291    let request = match websocket_boot_request(uri, headers) {
292        Ok(request) => request,
293        Err(error) => {
294            let _ = send_websocket_error(&mut socket, error).await;
295            return;
296        }
297    };
298    let connection = match gateway.connect(request) {
299        Ok(connection) => connection,
300        Err(error) => {
301            let _ = send_websocket_error(&mut socket, error).await;
302            return;
303        }
304    };
305
306    while let Some(message) = socket.next().await {
307        let message = match message {
308            Ok(message) => message,
309            Err(error) => {
310                let _ =
311                    send_websocket_error(&mut socket, BootError::Adapter(error.to_string())).await;
312                return;
313            }
314        };
315
316        let Some(message) = decode_websocket_message(message) else {
317            continue;
318        };
319        let message = match message {
320            Ok(message) => message,
321            Err(error) => {
322                if send_websocket_error(&mut socket, error).await.is_err() {
323                    return;
324                }
325                continue;
326            }
327        };
328
329        match connection.dispatch(message).await {
330            Ok(Some(reply)) => {
331                if send_websocket_message(&mut socket, reply).await.is_err() {
332                    return;
333                }
334            }
335            Ok(None) => {}
336            Err(error) => {
337                if send_websocket_error(&mut socket, error).await.is_err() {
338                    return;
339                }
340            }
341        }
342    }
343}
344
345fn decode_websocket_message(message: AxumWebSocketMessage) -> Option<Result<WebSocketMessage>> {
346    match message {
347        AxumWebSocketMessage::Text(text) => Some(
348            serde_json::from_str(&text).map_err(|error| BootError::BadRequest(error.to_string())),
349        ),
350        AxumWebSocketMessage::Binary(bytes) => Some(
351            serde_json::from_slice(&bytes)
352                .map_err(|error| BootError::BadRequest(error.to_string())),
353        ),
354        AxumWebSocketMessage::Close(_) => None,
355        AxumWebSocketMessage::Ping(_) | AxumWebSocketMessage::Pong(_) => None,
356    }
357}
358
359async fn send_websocket_message(
360    socket: &mut WebSocket,
361    message: WebSocketMessage,
362) -> std::result::Result<(), axum::Error> {
363    let text = serde_json::to_string(&message)
364        .unwrap_or_else(|error| format!(r#"{{"event":"error","data":"{error}"}}"#));
365    socket.send(AxumWebSocketMessage::Text(text.into())).await
366}
367
368async fn send_websocket_error(
369    socket: &mut WebSocket,
370    error: BootError,
371) -> std::result::Result<(), axum::Error> {
372    let message = WebSocketMessage::text("error", error.http_response_message());
373    send_websocket_message(socket, message).await
374}
375
376fn map_body_error(error: axum::Error, body_limit: usize) -> BootError {
377    if error
378        .source()
379        .is_some_and(|source| source.is::<http_body_util::LengthLimitError>())
380    {
381        BootError::PayloadTooLarge(format!("request body exceeds {body_limit} bytes"))
382    } else {
383        BootError::Adapter(error.to_string())
384    }
385}
386
387fn to_axum_response(response: BootResponse) -> Response {
388    if let Err(error) = response.validate() {
389        return error_response(StatusCode::INTERNAL_SERVER_ERROR, internal_message(error));
390    }
391    let is_streaming = response.is_streaming();
392    let status = match StatusCode::from_u16(response.status()) {
393        Ok(status) => status,
394        Err(err) => {
395            return error_response(
396                StatusCode::INTERNAL_SERVER_ERROR,
397                format!("invalid response status {}: {err}", response.status()),
398            );
399        }
400    };
401
402    let mut builder = Response::builder().status(status);
403    builder = match with_response_headers(builder, response.header_entries()) {
404        Ok(builder) => builder,
405        Err(message) => return error_response(StatusCode::INTERNAL_SERVER_ERROR, message),
406    };
407
408    let body = if is_streaming {
409        let Some(stream) = response.into_sse_stream() else {
410            return error_response(
411                StatusCode::INTERNAL_SERVER_ERROR,
412                "streaming response body has already been consumed".to_string(),
413            );
414        };
415        Body::from_stream(stream.map(|event| event.map(|event| event.encode())))
416    } else {
417        Body::from(response.into_body())
418    };
419
420    builder
421        .body(body)
422        .unwrap_or_else(|err| error_response(StatusCode::INTERNAL_SERVER_ERROR, err.to_string()))
423}
424
425fn with_response_headers<I, N, V>(
426    mut builder: ResponseBuilder,
427    headers: I,
428) -> std::result::Result<ResponseBuilder, String>
429where
430    I: IntoIterator<Item = (N, V)>,
431    N: AsRef<str>,
432    V: AsRef<str>,
433{
434    for (name, value) in headers {
435        let name = name.as_ref();
436        let value = value.as_ref();
437        let header_name = HeaderName::try_from(name)
438            .map_err(|err| format!("invalid response header name {name:?}: {err}"))?;
439        let header_value = HeaderValue::try_from(value)
440            .map_err(|err| format!("invalid response header value for {name:?}: {err}"))?;
441        builder = builder.header(header_name, header_value);
442    }
443
444    Ok(builder)
445}
446
447fn internal_message(error: BootError) -> String {
448    match error {
449        BootError::Internal(message) => message,
450        error => error.to_string(),
451    }
452}
453
454fn with_allow_header(app: &BootApplication, path: &str, mut response: Response) -> Response {
455    if response.status() != StatusCode::METHOD_NOT_ALLOWED {
456        return response;
457    }
458
459    let Some(allow) = app.allowed_methods_header(path) else {
460        return response;
461    };
462
463    match HeaderValue::try_from(allow) {
464        Ok(value) => {
465            response.headers_mut().insert(ALLOW, value);
466            response
467        }
468        Err(err) => error_response(
469            StatusCode::INTERNAL_SERVER_ERROR,
470            format!("invalid allow header value: {err}"),
471        ),
472    }
473}
474
475fn finalize_response(
476    app: &BootApplication,
477    path: &str,
478    response: Response,
479    is_head: bool,
480) -> Response {
481    strip_head_body(is_head, with_allow_header(app, path, response))
482}
483
484fn strip_head_body(is_head: bool, response: Response) -> Response {
485    if !is_head {
486        return response;
487    }
488
489    let (parts, _) = response.into_parts();
490    Response::from_parts(parts, Body::empty())
491}
492
493fn boot_error_response(error: BootError) -> Response {
494    to_axum_response(BootResponse::from_error(&error))
495}
496
497fn error_response(status: StatusCode, message: String) -> Response {
498    let mut response = Response::new(Body::from(message));
499    *response.status_mut() = status;
500    response.headers_mut().insert(
501        CONTENT_TYPE,
502        HeaderValue::from_static("text/plain; charset=utf-8"),
503    );
504    response
505}
506
507impl TryFrom<Method> for HttpMethod {
508    type Error = BootError;
509
510    fn try_from(method: Method) -> std::result::Result<Self, Self::Error> {
511        method.as_str().parse()
512    }
513}