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

a3s_boot/adapters/
axum.rs
