poem_openapi/payload/
base64_payload.rs

1use std::ops::{Deref, DerefMut};
2
3use base64::engine::{general_purpose::STANDARD, Engine};
4use bytes::Bytes;
5use futures_util::TryFutureExt;
6use poem::{IntoResponse, Request, RequestBody, Response, Result};
7
8use crate::{
9    error::ParseRequestPayloadError,
10    payload::{ParsePayload, Payload},
11    registry::{MetaMediaType, MetaResponse, MetaResponses, MetaSchema, MetaSchemaRef, Registry},
12    ApiResponse,
13};
14
15/// A binary payload encoded with `base64`.
16///
17/// # Examples
18///
19/// ```rust
20/// use poem::{
21///     error::BadRequest,
22///     http::{Method, StatusCode, Uri},
23///     test::TestClient,
24///     Body, IntoEndpoint, Request, Result,
25/// };
26/// use poem_openapi::{
27///     payload::{Base64, Json},
28///     OpenApi, OpenApiService,
29/// };
30/// use tokio::{io::AsyncReadExt, sync::Mutex};
31///
32/// #[derive(Default)]
33/// struct MyApi {
34///     data: Mutex<Vec<u8>>,
35/// }
36///
37/// #[OpenApi]
38/// impl MyApi {
39///     #[oai(path = "/upload", method = "post")]
40///     async fn upload_binary(&self, data: Base64<Vec<u8>>) -> Json<usize> {
41///         let len = data.len();
42///         assert_eq!(data.0, b"abcdef");
43///         *self.data.lock().await = data.0;
44///         Json(len)
45///     }
46///
47///     #[oai(path = "/download", method = "get")]
48///     async fn download_binary(&self) -> Base64<Vec<u8>> {
49///         Base64(self.data.lock().await.clone())
50///     }
51/// }
52///
53/// let api = OpenApiService::new(MyApi::default(), "Demo", "0.1.0");
54/// let cli = TestClient::new(api);
55///
56/// # tokio::runtime::Runtime::new().unwrap().block_on(async {
57/// let resp = cli
58///     .post("/upload")
59///     .content_type("text/plain")
60///     .body("YWJjZGVm")
61///     .send()
62///     .await;
63/// resp.assert_status_is_ok();
64/// resp.assert_text("6").await;
65///
66/// let resp = cli.get("/download").send().await;
67/// resp.assert_status_is_ok();
68/// resp.assert_text("YWJjZGVm").await;
69/// # });
70/// ```
71#[derive(Debug, Clone, Eq, PartialEq)]
72pub struct Base64<T>(pub T);
73
74impl<T> Deref for Base64<T> {
75    type Target = T;
76
77    fn deref(&self) -> &Self::Target {
78        &self.0
79    }
80}
81
82impl<T> DerefMut for Base64<T> {
83    fn deref_mut(&mut self) -> &mut Self::Target {
84        &mut self.0
85    }
86}
87
88impl<T: Send> Payload for Base64<T> {
89    const CONTENT_TYPE: &'static str = "text/plain; charset=utf-8";
90
91    fn check_content_type(content_type: &str) -> bool {
92        matches!(content_type.parse::<mime::Mime>(), Ok(content_type) if content_type.type_() == "text"
93                && (content_type.subtype() == "plain"
94                || content_type
95                    .suffix()
96                    .is_some_and(|v| v == "plain")))
97    }
98
99    fn schema_ref() -> MetaSchemaRef {
100        MetaSchemaRef::Inline(Box::new(MetaSchema {
101            format: Some("string"),
102            ..MetaSchema::new("bytes")
103        }))
104    }
105}
106
107async fn read_base64(body: &mut RequestBody) -> Result<Vec<u8>> {
108    let body = async move { body.take() }
109        .and_then(|body| body.into_vec())
110        .await
111        .map_err(|err| ParseRequestPayloadError {
112            reason: err.to_string(),
113        })?;
114    let data = STANDARD
115        .decode(body)
116        .map_err(|err| ParseRequestPayloadError {
117            reason: err.to_string(),
118        })?;
119    Ok(data)
120}
121
122impl ParsePayload for Base64<Vec<u8>> {
123    const IS_REQUIRED: bool = true;
124
125    async fn from_request(_request: &Request, body: &mut RequestBody) -> Result<Self> {
126        read_base64(body).await.map(Self)
127    }
128}
129
130impl ParsePayload for Base64<Bytes> {
131    const IS_REQUIRED: bool = true;
132
133    async fn from_request(_request: &Request, body: &mut RequestBody) -> Result<Self> {
134        read_base64(body).await.map(|data| Self(data.into()))
135    }
136}
137
138impl<T: AsRef<[u8]> + Send> IntoResponse for Base64<T> {
139    fn into_response(self) -> Response {
140        Response::builder()
141            .content_type(Self::CONTENT_TYPE)
142            .body(STANDARD.encode(self.0.as_ref()))
143    }
144}
145
146impl<T: AsRef<[u8]> + Send> ApiResponse for Base64<T> {
147    fn meta() -> MetaResponses {
148        MetaResponses {
149            responses: vec![MetaResponse {
150                description: "",
151                status: Some(200),
152                status_range: None,
153                content: vec![MetaMediaType {
154                    content_type: Self::CONTENT_TYPE,
155                    schema: Self::schema_ref(),
156                }],
157                headers: vec![],
158            }],
159        }
160    }
161
162    fn register(_registry: &mut Registry) {}
163}
164
165impl_apirequest_for_payload!(Base64<Vec<u8>>);
166impl_apirequest_for_payload!(Base64<Bytes>);
poem_openapi/payload/base64_payload.rs

poem_openapi/payload/
base64_payload.rs
