objectstore_server/
multipart.rs

1//! Types and utilities to support Multipart streaming responses.
2//!
3//! Compared to `axum_extra::response::MultipartForm`, this implementation supports attaching arbitrary headers to
4//! each part, as well as the possibility to convert a `Stream` of those parts to a streaming `Response`.
5
6use axum::body::Body;
7use axum::response::IntoResponse as _;
8use axum::response::Response;
9use bytes::{BufMut, Bytes, BytesMut};
10use futures::Stream;
11use futures::StreamExt;
12use futures::stream::BoxStream;
13use http::HeaderMap;
14use http::HeaderValue;
15use http::header::{CONTENT_DISPOSITION, CONTENT_TYPE};
16
17/// A part in a Multipart response.
18#[derive(Debug)]
19pub struct Part {
20    headers: HeaderMap,
21    body: Bytes,
22}
23
24impl Part {
25    /// Creates a new Multipart part with the given content type, body, and headers.
26    /// The name is hardcoded to "part".
27    pub fn new(body: Bytes, mut headers: HeaderMap, content_type: Option<HeaderValue>) -> Self {
28        headers.insert(
29            CONTENT_DISPOSITION,
30            HeaderValue::from_static("form-data; name=part"),
31        );
32        if let Some(content_type) = content_type {
33            headers.insert(CONTENT_TYPE, content_type);
34        }
35        Self { headers, body }
36    }
37}
38
39/// Converts a stream of [`Part`]s into a streaming `multipart/form-data` HTTP response.
40///
41/// Implemented for any `Stream<Item = T>` where `T: Into<Part>`. The boundary is
42/// formatted as `os-boundary-<hex>` using the provided `u128` value.
43pub trait IntoMultipartResponse {
44    /// Consumes `self` and returns a streaming `multipart/form-data` response.
45    ///
46    /// The `boundary` parameter is formatted as `os-boundary-<hex>` and embedded in the
47    /// `Content-Type` header. Each item yielded by the stream becomes one part, separated
48    /// by that boundary. A closing boundary is appended after the last item.
49    fn into_multipart_response(self, boundary: u128) -> Response;
50}
51
52impl<S, T> IntoMultipartResponse for S
53where
54    S: Stream<Item = T> + Send + 'static,
55    T: Into<Part> + Send,
56{
57    fn into_multipart_response(self, boundary: u128) -> Response {
58        let boundary_str = format!("os-boundary-{boundary:032x}");
59        let boundary = {
60            let mut bytes = BytesMut::with_capacity(boundary_str.len() + 4);
61            bytes.put(&b"--"[..]);
62            bytes.put(boundary_str.as_bytes());
63            bytes.put(&b"\r\n"[..]);
64            bytes.freeze()
65        };
66
67        let mut headers = HeaderMap::new();
68        headers.insert(
69            CONTENT_TYPE,
70            format!("multipart/form-data; boundary=\"{}\"", &boundary_str)
71                .parse()
72                .expect("valid header value, as we just defined it as \"os-boundary-X\" where X are hex digits"),
73        );
74
75        let body: BoxStream<Result<Bytes, std::convert::Infallible>> = async_stream::try_stream! {
76            let items = self;
77            futures::pin_mut!(items);
78            while let Some(item) = items.next().await {
79                yield boundary.clone();
80                let part = item.into();
81                yield serialize_headers(part.headers);
82                yield serialize_body(part.body);
83            }
84
85            let mut closing = BytesMut::with_capacity(boundary.len());
86            closing.put(boundary.slice(..boundary.len() - 2)); // don't take trailing \r\n
87            closing.put(&b"--"[..]);
88            yield closing.freeze();
89        }
90        .boxed();
91
92        (headers, Body::from_stream(body)).into_response()
93    }
94}
95
96fn serialize_headers(headers: HeaderMap) -> Bytes {
97    // https://github.com/hyperium/hyper/blob/0f0b6ed3ac55ac1682afd2104cb8d0385149249a/src/proto/h1/role.rs#L399
98    let mut res = BytesMut::with_capacity(30 + 30 * headers.len());
99    for (name, value) in &headers {
100        res.put(name.as_str().as_bytes());
101        res.put(&b": "[..]);
102        res.put(value.as_bytes());
103        res.put(&b"\r\n"[..]);
104    }
105    res.put(&b"\r\n"[..]);
106    res.freeze()
107}
108
109fn serialize_body(body: Bytes) -> Bytes {
110    let mut res = BytesMut::with_capacity(body.len() + 2);
111    res.put(body);
112    res.put(&b"\r\n"[..]);
113    res.freeze()
114}
115
116#[cfg(test)]
117mod tests {
118    use super::*;
119    use axum::body::{Body, to_bytes};
120    use axum::extract::{FromRequest, Multipart};
121    use axum::http::Request;
122
123    /// Validates that our `Multipart` streaming response produces a valid `multipart/form-data` response
124    /// that's parsed as expected by `axum::extract::Multipart`.
125    #[tokio::test]
126    async fn test_multipart_response() {
127        let mut extra_headers = HeaderMap::new();
128        extra_headers.insert("X-Custom-Header", "custom-value".parse().unwrap());
129        extra_headers.insert("X-File-Id", "12345".parse().unwrap());
130        let parts = vec![
131            Part::new(
132                Bytes::from(r#"{"key":"value"}"#),
133                HeaderMap::new(),
134                Some(HeaderValue::from_static("application/json")),
135            ),
136            Part::new(
137                Bytes::from(vec![0x00, 0x01, 0x02, 0xff, 0xfe]),
138                extra_headers,
139                Some(HeaderValue::from_static("application/octet-stream")),
140            ),
141        ];
142        let boundary: u128 = 0xdeadbeef;
143        let response = futures::stream::iter(parts).into_multipart_response(boundary);
144
145        let boundary = format!("os-boundary-{boundary:032x}");
146        let content_type_str = format!("multipart/form-data; boundary=\"{boundary}\"");
147        assert_eq!(
148            response
149                .headers()
150                .get(CONTENT_TYPE)
151                .unwrap()
152                .to_str()
153                .unwrap(),
154            &content_type_str
155        );
156
157        let body = to_bytes(response.into_body(), usize::MAX).await.unwrap();
158        let request = Request::builder()
159            .header(CONTENT_TYPE, &content_type_str)
160            .body(Body::from(body))
161            .unwrap();
162        let mut multipart = Multipart::from_request(request, &()).await.unwrap();
163
164        let field = multipart.next_field().await.unwrap().unwrap();
165        assert_eq!(field.name(), Some("part"));
166        assert_eq!(field.file_name(), None);
167        assert_eq!(field.content_type(), Some("application/json"));
168        assert_eq!(field.headers().len(), 2);
169        assert_eq!(field.bytes().await.unwrap(), r#"{"key":"value"}"#);
170
171        let field = multipart.next_field().await.unwrap().unwrap();
172        assert_eq!(field.name(), Some("part"));
173        assert_eq!(field.file_name(), None);
174        assert_eq!(field.content_type(), Some("application/octet-stream"));
175        assert_eq!(field.headers().len(), 4);
176        assert_eq!(
177            field.headers().get("X-Custom-Header").unwrap(),
178            "custom-value"
179        );
180        assert_eq!(field.headers().get("X-File-Id").unwrap(), "12345");
181        assert!(field.headers().get("content-disposition").is_some());
182        assert!(field.headers().get("content-type").is_some());
183        assert_eq!(
184            field.bytes().await.unwrap(),
185            vec![0x00, 0x01, 0x02, 0xff, 0xfe]
186        );
187
188        assert!(multipart.next_field().await.unwrap().is_none());
189    }
190}
objectstore_server/multipart.rs

objectstore_server/
multipart.rs
