Skip to main content

objectstore_server/extractors/
body.rs

1//! Axum extractor for bandwidth-metered request bodies.
2
3use axum::extract::rejection::PathRejection;
4use axum::extract::{FromRequest, FromRequestParts, Path, Request};
5use futures_util::{StreamExt, TryStreamExt};
6use objectstore_service::id::ObjectContext;
7use objectstore_service::stream::{ClientError, ClientStream};
8
9use super::id::ContextParams;
10use crate::state::ServiceState;
11
12/// An extractor that converts the request body into a metered [`ClientStream`].
13///
14/// Extracts the [`ObjectContext`] from the request path to attribute bandwidth to the correct
15/// per-usecase and per-scope accumulators in addition to the global accumulator.
16pub struct MeteredBody(pub ClientStream);
17
18impl std::fmt::Debug for MeteredBody {
19    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
20        f.debug_struct("MeteredBody").finish()
21    }
22}
23
24impl FromRequest<ServiceState> for MeteredBody {
25    type Rejection = PathRejection;
26
27    async fn from_request(request: Request, state: &ServiceState) -> Result<Self, Self::Rejection> {
28        let (mut parts, body) = request.into_parts();
29        let Path(params) =
30            <Path<ContextParams> as FromRequestParts<ServiceState>>::from_request_parts(
31                &mut parts, state,
32            )
33            .await?;
34        let context = ObjectContext {
35            usecase: params.usecase,
36            scopes: params.scopes,
37        };
38        let stream = body.into_data_stream().map_err(ClientError::new).boxed();
39        let stream = state.meter_stream(stream, &context).boxed();
40        Ok(Self(stream))
41    }
42}