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
39pub trait IntoMultipartResponse {
40    fn into_multipart_response(self, boundary: u128) -> Response;
41}
42
43impl<S, T> IntoMultipartResponse for S
44where
45    S: Stream<Item = T> + Send + 'static,
46    T: Into<Part> + Send,
47{
48    fn into_multipart_response(self, boundary: u128) -> Response {
49        let boundary_str = format!("os-boundary-{:032x}", boundary);
50        let boundary = {
51            let mut bytes = BytesMut::with_capacity(boundary_str.len() + 4);
52            bytes.put(&b"--"[..]);
53            bytes.put(boundary_str.as_bytes());
54            bytes.put(&b"\r\n"[..]);
55            bytes.freeze()
56        };
57
58        let mut headers = HeaderMap::new();
59        headers.insert(
60            CONTENT_TYPE,
61            format!("multipart/form-data; boundary=\"{}\"", &boundary_str)
62                .parse()
63                .expect("valid header value, as we just defined it as \"os-boundary-X\" where X are hex digits"),
64        );
65
66        let body: BoxStream<Result<bytes::Bytes, std::convert::Infallible>> =
67            async_stream::try_stream! {
68                let items = self;
69                futures::pin_mut!(items);
70                while let Some(item) = items.next().await {
71                    yield boundary.clone();
72                    let part = item.into();
73                    yield serialize_headers(part.headers);
74                    yield serialize_body(part.body);
75                }
76
77                let mut closing = BytesMut::with_capacity(boundary.len());
78                closing.put(boundary.slice(..boundary.len() - 2)); // don't take trailing \r\n
79                closing.put(&b"--"[..]);
80                yield closing.freeze();
81            }
82            .boxed();
83
84        (headers, Body::from_stream(body)).into_response()
85    }
86}
87
88fn serialize_headers(headers: HeaderMap) -> Bytes {
89    // https://github.com/hyperium/hyper/blob/0f0b6ed3ac55ac1682afd2104cb8d0385149249a/src/proto/h1/role.rs#L399
90    let mut res = BytesMut::with_capacity(30 + 30 * headers.len());
91    for (name, value) in &headers {
92        res.put(name.as_str().as_bytes());
93        res.put(&b": "[..]);
94        res.put(value.as_bytes());
95        res.put(&b"\r\n"[..]);
96    }
97    res.put(&b"\r\n"[..]);
98    res.freeze()
99}
100
101fn serialize_body(body: Bytes) -> Bytes {
102    let mut res = BytesMut::with_capacity(body.len() + 2);
103    res.put(body);
104    res.put(&b"\r\n"[..]);
105    res.freeze()
106}
107
108#[cfg(test)]
109mod tests {
110    use super::*;
111    use axum::body::{Body, to_bytes};
112    use axum::extract::{FromRequest, Multipart};
113    use axum::http::Request;
114
115    /// Validates that our `Multipart` streaming response produces a valid `multipart/form-data` response
116    /// that's parsed as expected by `axum::extract::Multipart`.
117    #[tokio::test]
118    async fn test_multipart_response() {
119        let mut extra_headers = HeaderMap::new();
120        extra_headers.insert("X-Custom-Header", "custom-value".parse().unwrap());
121        extra_headers.insert("X-File-Id", "12345".parse().unwrap());
122        let parts = vec![
123            Part::new(
124                Bytes::from(r#"{"key":"value"}"#),
125                HeaderMap::new(),
126                Some(HeaderValue::from_static("application/json")),
127            ),
128            Part::new(
129                Bytes::from(vec![0x00, 0x01, 0x02, 0xff, 0xfe]),
130                extra_headers,
131                Some(HeaderValue::from_static("application/octet-stream")),
132            ),
133        ];
134        let boundary: u128 = 0xdeadbeef;
135        let response = futures::stream::iter(parts).into_multipart_response(boundary);
136
137        let boundary = format!("os-boundary-{:032x}", boundary);
138        let content_type_str = format!("multipart/form-data; boundary=\"{}\"", boundary);
139        assert_eq!(
140            response
141                .headers()
142                .get(CONTENT_TYPE)
143                .unwrap()
144                .to_str()
145                .unwrap(),
146            &content_type_str
147        );
148
149        let body = to_bytes(response.into_body(), usize::MAX).await.unwrap();
150        let request = Request::builder()
151            .header(CONTENT_TYPE, &content_type_str)
152            .body(Body::from(body))
153            .unwrap();
154        let mut multipart = Multipart::from_request(request, &()).await.unwrap();
155
156        let field = multipart.next_field().await.unwrap().unwrap();
157        assert_eq!(field.name(), Some("part"));
158        assert_eq!(field.file_name(), None);
159        assert_eq!(field.content_type(), Some("application/json"));
160        assert_eq!(field.headers().len(), 2);
161        assert_eq!(field.bytes().await.unwrap(), r#"{"key":"value"}"#);
162
163        let field = multipart.next_field().await.unwrap().unwrap();
164        assert_eq!(field.name(), Some("part"));
165        assert_eq!(field.file_name(), None);
166        assert_eq!(field.content_type(), Some("application/octet-stream"));
167        assert_eq!(field.headers().len(), 4);
168        assert_eq!(
169            field.headers().get("X-Custom-Header").unwrap(),
170            "custom-value"
171        );
172        assert_eq!(field.headers().get("X-File-Id").unwrap(), "12345");
173        assert!(field.headers().get("content-disposition").is_some());
174        assert!(field.headers().get("content-type").is_some());
175        assert_eq!(
176            field.bytes().await.unwrap(),
177            vec![0x00, 0x01, 0x02, 0xff, 0xfe]
178        );
179
180        assert!(multipart.next_field().await.unwrap().is_none());
181    }
182}
objectstore_server/multipart.rs

objectstore_server/
multipart.rs
