Axum post handler that accepts multipart and image/jpeg

[rustic1]

Using axum, I have a post handler that should accept a file that is uploaded.
The problem is that one of the clients uploads a single file as content type "image/jpeg", whereas my browser uploads a single file as "multipart" structure.

The two handlers below work when I use them individually. However, what I would like is a handler that can accept both upload "types" under the same route.

How can this be achieved?

pub async fn upload_jpeg(mut multipart: Multipart) -> impl IntoResponse {
    while let Some(mut field) = multipart.next_field().await.unwrap() {
        let name = field.name().unwrap().to_string();
        let data = field.bytes().await.unwrap();
        eprintln!("Length of `{}` is {} bytes", name, data.len());
        tokio::fs::write("upload.jpg", data).await.unwrap();
    }
    "Ok"
}

pub async fn upload_jpeg2(mut data: Bytes) -> impl IntoResponse {
    eprintln!("Length is {} bytes", data.len());
    tokio::fs::write("upload.jpg", data).await.unwrap();
    "Ok"
}

[jofas]

You can define your own extractors in Axum.

[rustic1]

Ok, I have tried and came up with this, which seems to work:

pub async fn upload_jpeg3(headers: HeaderMap,
                          data: Bytes)
                          -> impl IntoResponse {
    let content_type = headers[header::CONTENT_TYPE].to_str().unwrap();
    if content_type.starts_with("multipart") {
        let stream = once(async move { Result::<Bytes, Infallible>::Ok(Bytes::from(data)) });
        let boundary = multer::parse_boundary(content_type).unwrap();
        let mut multipart = multer::Multipart::new(stream, boundary);
        if let Ok(Some(field)) = multipart.next_field().await {
            tokio::fs::write("upload.jpg", field.bytes().await.unwrap()).await.unwrap();
        }
    } else if content_type == "image/jpeg" {
        tokio::fs::write("upload.jpg", data).await.unwrap();
    }

    (StatusCode::CREATED, "image uploaded".to_string())
}

Still a hack and error handling is missing, but I appreciate any comment on whether the general approach makes sense.

[jofas]

Here how I'd use an extractor instead (with rudimentary error handling):

use axum::{
    async_trait,
    body::Bytes,
    extract::{FromRequest, Multipart, Request},
    http::{header::CONTENT_TYPE, StatusCode},
};

struct Jpeg(Bytes);

#[async_trait]
impl<S> FromRequest<S> for Jpeg
where
    Bytes: FromRequest<S>,
    S: Send + Sync,
{
    type Rejection = StatusCode;

    async fn from_request(req: Request, state: &S) -> Result<Self, Self::Rejection> {
        let Some(content_type) = req.headers().get(CONTENT_TYPE) else {
            return Err(StatusCode::BAD_REQUEST);
        };

        let body = if content_type == "multipart/form-data" {
            let mut multipart = Multipart::from_request(req, state)
                .await
                .map_err(|_| StatusCode::BAD_REQUEST)?;

            let Ok(Some(field)) = multipart.next_field().await else {
                return Err(StatusCode::BAD_REQUEST);
            };

            field.bytes().await.map_err(|_| StatusCode::BAD_REQUEST)?
        } else if content_type == "image/jpeg" {
            Bytes::from_request(req, state)
                .await
                .map_err(|_| StatusCode::BAD_REQUEST)?
        } else {
            return Err(StatusCode::BAD_REQUEST);
        };

        Ok(Self(body))
    }
}

I think this is nicer in the case of mulipart/form-data compared to using the Bytes extractor in your endpoint to convert the bytes back into a stream that you can then treat as multipart data. In your endpoint, you can then do this:

pub async fn upload_jpeg3(jpeg: Jpeg) -> impl IntoResponse {
    tokio::fs::write("upload.jpg", jpeg.0).await.unwrap();
    (StatusCode::CREATED, "image uploaded".to_string())
}

in order to save the file to disk.

[rustic1]

Thanks, this helped tremendously.

The only part of your code that did not work (for me) was the comparison with "multipart/form-data" since the boundary code was on the header same line. I had to use:

        let Ok(content_type) = content_type.to_str() else {
            return Err(StatusCode::BAD_REQUEST);
        };

        let body = if content_type.starts_with("multipart/form-data") {