Skip to main content

axum_image/
extract.rs

1use axum::{
2    body::Bytes,
3    extract::{FromRequest, Request},
4    http::{StatusCode, header::CONTENT_TYPE},
5};
6use axum_extra::extract::Multipart;
7use serde::de::DeserializeOwned;
8
9/// An image extractor accepting:
10/// * `multipart/form-data`
11/// * `image/png`
12/// * `image/jpeg`
13/// * `image/avif`
14/// * `image/webp`
15pub struct Image(pub Bytes, pub String);
16
17impl<S> FromRequest<S> for Image
18where
19    Bytes: FromRequest<S>,
20    S: Send + Sync,
21{
22    type Rejection = StatusCode;
23
24    async fn from_request(req: Request, state: &S) -> Result<Self, Self::Rejection> {
25        let Some(content_type) = req.headers().get(CONTENT_TYPE) else {
26            return Err(StatusCode::BAD_REQUEST);
27        };
28
29        let content_type = content_type.to_str().unwrap();
30        let content_type_string = content_type.to_string();
31
32        let body = if content_type.starts_with("multipart/form-data") {
33            let mut multipart = Multipart::from_request(req, state)
34                .await
35                .map_err(|_| StatusCode::BAD_REQUEST)?;
36
37            let Ok(Some(field)) = multipart.next_field().await else {
38                return Err(StatusCode::BAD_REQUEST);
39            };
40
41            field.bytes().await.map_err(|_| StatusCode::BAD_REQUEST)?
42        } else if (content_type == "image/avif")
43            | (content_type == "image/jpeg")
44            | (content_type == "image/png")
45            | (content_type == "image/webp")
46            | (content_type == "image/gif")
47        {
48            Bytes::from_request(req, state)
49                .await
50                .map_err(|_| StatusCode::BAD_REQUEST)?
51        } else {
52            return Err(StatusCode::BAD_REQUEST);
53        };
54
55        Ok(Self(body, content_type_string))
56    }
57}
58
59/// A file extractor accepting:
60/// * `multipart/form-data`
61///
62/// Will also attempt to parse out the **last** field in the multipart upload
63/// as the given struct from JSON. Every other field is put into a vector of bytes,
64/// as they are seen as raw binary data.
65pub struct JsonMultipart<T: DeserializeOwned>(pub Vec<Bytes>, pub T);
66
67impl<S, T> FromRequest<S> for JsonMultipart<T>
68where
69    Bytes: FromRequest<S>,
70    S: Send + Sync,
71    T: DeserializeOwned,
72{
73    type Rejection = (StatusCode, String);
74
75    async fn from_request(req: Request, state: &S) -> Result<Self, Self::Rejection> {
76        let Some(content_type) = req.headers().get(CONTENT_TYPE) else {
77            return Err((
78                StatusCode::BAD_REQUEST,
79                "no content type header".to_string(),
80            ));
81        };
82
83        let content_type = content_type.to_str().unwrap();
84
85        if !content_type.starts_with("multipart/form-data") {
86            return Err((
87                StatusCode::BAD_REQUEST,
88                "expected multipart/form-data".to_string(),
89            ));
90        }
91
92        let mut multipart = Multipart::from_request(req, state).await.map_err(|_| {
93            (
94                StatusCode::BAD_REQUEST,
95                "could not read multipart".to_string(),
96            )
97        })?;
98
99        let mut body: Vec<Bytes> = {
100            let mut out = Vec::new();
101
102            while let Ok(Some(field)) = multipart.next_field().await {
103                out.push(
104                    field
105                        .bytes()
106                        .await
107                        .map_err(|e| (StatusCode::BAD_REQUEST, e.to_string()))?,
108                );
109            }
110
111            out
112        };
113
114        let last = match body.pop() {
115            Some(b) => b,
116            None => {
117                return Err((
118                    StatusCode::BAD_REQUEST,
119                    "could not read json data".to_string(),
120                ));
121            }
122        };
123
124        let json: T = match serde_json::from_str(&match String::from_utf8(last.to_vec()) {
125            Ok(s) => s,
126            Err(_) => return Err((StatusCode::BAD_REQUEST, "json data isn't utf8".to_string())),
127        }) {
128            Ok(s) => s,
129            Err(e) => {
130                return Err((StatusCode::BAD_REQUEST, e.to_string()));
131            }
132        };
133
134        Ok(Self(body, json))
135    }
136}