1use crate::response::ApiResponse;
2use axum::{
3 body::Bytes,
4 extract::{FromRequest, OptionalFromRequest, Request},
5 response::IntoResponse,
6};
7use serde::de::DeserializeOwned;
8use std::{str::FromStr, sync::LazyLock};
9
10pub struct PayloadRejection(anyhow::Error);
11
12impl IntoResponse for PayloadRejection {
13 fn into_response(self) -> axum::response::Response {
14 ApiResponse::error(format!("invalid payload: {}", self.0))
15 .with_status(axum::http::StatusCode::BAD_REQUEST)
16 .into_response()
17 }
18}
19
20impl From<anyhow::Error> for PayloadRejection {
21 fn from(err: anyhow::Error) -> Self {
22 Self(err)
23 }
24}
25
26static AVAILABLE_DESERIALIZERS: LazyLock<[mime::Mime; 4]> = LazyLock::new(|| {
27 [
28 mime::APPLICATION_JSON,
29 mime::APPLICATION_MSGPACK,
30 mime::TEXT_XML,
31 mime::Mime::from_str("application/yaml").unwrap(),
32 ]
33});
34
35const MAX_XML_NESTING_DEPTH: usize = 256;
36
37fn check_xml_nesting_depth(content: &[u8]) -> Result<(), anyhow::Error> {
38 let mut reader = quick_xml::Reader::from_reader(content);
39 let mut buffer = Vec::new();
40 let mut depth: usize = 0;
41
42 loop {
43 match reader.read_event_into(&mut buffer) {
44 Ok(quick_xml::events::Event::Start(_)) => {
45 depth += 1;
46
47 if depth > MAX_XML_NESTING_DEPTH {
48 return Err(anyhow::anyhow!(
49 "xml nesting exceeds the maximum depth of {MAX_XML_NESTING_DEPTH}"
50 ));
51 }
52 }
53 Ok(quick_xml::events::Event::End(_)) => depth = depth.saturating_sub(1),
54 Ok(quick_xml::events::Event::Eof) | Err(_) => break,
55 _ => {}
56 }
57
58 buffer.clear();
59 }
60
61 Ok(())
62}
63
64pub struct Payload<T: DeserializeOwned>(pub T);
66
67impl<T: DeserializeOwned> Payload<T> {
68 pub fn into_inner(self) -> T {
69 self.0
70 }
71
72 pub fn from_bytes(
73 content_type: mime::Mime,
74 mut bytes: Bytes,
75 ) -> Result<Self, PayloadRejection> {
76 match content_type.essence_str() {
77 m if m == mime::APPLICATION_JSON.essence_str() => {
78 if bytes.is_empty() {
79 bytes = Bytes::from_static(b"{}");
80 }
81
82 let value = serde_json::from_slice(&bytes).map_err(anyhow::Error::from)?;
83 Ok(Payload(value))
84 }
85 m if m == mime::APPLICATION_MSGPACK.essence_str() => {
86 if bytes.is_empty() {
87 bytes = Bytes::from_static(&[0x80]);
88 }
89
90 let mut de = rmp_serde::Deserializer::new(bytes.as_ref()).with_human_readable();
91 let value = T::deserialize(&mut de).map_err(anyhow::Error::from)?;
92 Ok(Payload(value))
93 }
94 m if m == mime::TEXT_XML.essence_str() => {
95 if bytes.is_empty() {
96 bytes = Bytes::from_static(b"<root></root>");
97 }
98
99 check_xml_nesting_depth(bytes.as_ref())?;
100
101 let value =
102 quick_xml::de::from_reader(bytes.as_ref()).map_err(anyhow::Error::from)?;
103 Ok(Payload(value))
104 }
105 "application/yaml" => {
106 if bytes.is_empty() {
107 bytes = Bytes::from_static(b"{}");
108 }
109
110 let value = serde_saphyr::from_slice_with_options(
111 &bytes,
112 crate::utils::yaml_deserialize_options(),
113 )
114 .map_err(anyhow::Error::from)?;
115 Ok(Payload(value))
116 }
117 _ => Err(PayloadRejection(anyhow::anyhow!(
118 "unsupported content type"
119 ))),
120 }
121 }
122}
123
124impl<T: DeserializeOwned, S: Send + Sync> FromRequest<S> for Payload<T> {
125 type Rejection = PayloadRejection;
126
127 async fn from_request(req: Request, state: &S) -> Result<Self, Self::Rejection> {
128 let content_type = req
129 .headers()
130 .get(axum::http::header::CONTENT_TYPE)
131 .and_then(|v| v.to_str().ok())
132 .and_then(|s| s.parse::<mime::Mime>().ok());
133
134 let Some(content_type) = content_type else {
135 return Err(PayloadRejection(anyhow::anyhow!(
136 "missing content type header"
137 )));
138 };
139
140 if !AVAILABLE_DESERIALIZERS.contains(&content_type) {
141 return Err(PayloadRejection(anyhow::anyhow!(
142 "unsupported content type"
143 )));
144 }
145
146 let bytes = match Bytes::from_request(req, state).await {
147 Ok(b) => b,
148 Err(_) => return Err(PayloadRejection(anyhow::anyhow!("failed to read body"))),
149 };
150 Self::from_bytes(content_type, bytes)
151 }
152}
153
154impl<T, S> OptionalFromRequest<S> for Payload<T>
155where
156 T: DeserializeOwned,
157 S: Send + Sync,
158{
159 type Rejection = PayloadRejection;
160
161 async fn from_request(req: Request, state: &S) -> Result<Option<Self>, Self::Rejection> {
162 let content_type = req
163 .headers()
164 .get(axum::http::header::CONTENT_TYPE)
165 .and_then(|v| v.to_str().ok())
166 .and_then(|s| s.parse::<mime::Mime>().ok());
167 let Some(content_type) = content_type else {
168 return Ok(None);
169 };
170
171 if !AVAILABLE_DESERIALIZERS.contains(&content_type) {
172 return Err(PayloadRejection(anyhow::anyhow!(
173 "unsupported content type"
174 )));
175 }
176
177 let bytes = match Bytes::from_request(req, state).await {
178 Ok(b) => b,
179 Err(_) => return Err(PayloadRejection(anyhow::anyhow!("failed to read body"))),
180 };
181 Self::from_bytes(content_type, bytes).map(Some)
182 }
183}
184
185#[cfg(test)]
186mod tests {
187 use super::*;
188
189 #[test]
190 fn xml_nesting_within_the_limit_is_accepted() {
191 let body = "<a>".repeat(MAX_XML_NESTING_DEPTH);
192
193 assert!(check_xml_nesting_depth(body.as_bytes()).is_ok());
194 }
195
196 #[test]
197 fn xml_nesting_beyond_the_limit_is_rejected() {
198 let body = "<a>".repeat(MAX_XML_NESTING_DEPTH + 1);
199
200 assert!(check_xml_nesting_depth(body.as_bytes()).is_err());
201 }
202
203 #[test]
204 fn xml_siblings_do_not_accumulate_depth() {
205 let body = format!(
206 "<root>{}</root>",
207 "<a></a>".repeat(MAX_XML_NESTING_DEPTH * 4)
208 );
209
210 assert!(check_xml_nesting_depth(body.as_bytes()).is_ok());
211 }
212
213 #[test]
214 fn xml_payloads_within_the_limit_still_deserialize() {
215 let payload = match Payload::<serde_json::Value>::from_bytes(
216 mime::TEXT_XML,
217 Bytes::from_static(b"<root><nested><key>value</key></nested></root>"),
218 ) {
219 Ok(payload) => payload,
220 Err(err) => panic!("payload rejected: {}", err.0),
221 };
222
223 assert_eq!(payload.0["nested"]["key"]["$text"], "value");
224 }
225
226 #[test]
227 fn deeply_nested_xml_payloads_are_rejected_before_deserialization() {
228 let body = format!("<root>{}", "<a>".repeat(MAX_XML_NESTING_DEPTH + 1));
229
230 assert!(
231 Payload::<serde_json::Value>::from_bytes(mime::TEXT_XML, Bytes::from(body)).is_err()
232 );
233 }
234
235 fn yaml_payload(
237 body: impl Into<Bytes>,
238 ) -> Result<Payload<serde_json::Value>, PayloadRejection> {
239 Payload::from_bytes(
240 mime::Mime::from_str("application/yaml").unwrap(),
241 body.into(),
242 )
243 }
244
245 #[test]
246 fn yaml_payloads_deserialize() {
247 let payload = match yaml_payload("name: test\nnested:\n list: [1, 2]\n glob: \"*leaf\"\n")
248 {
249 Ok(payload) => payload,
250 Err(err) => panic!("payload rejected: {}", err.0),
251 };
252
253 assert_eq!(payload.0["name"], "test");
254 assert_eq!(payload.0["nested"]["list"], serde_json::json!([1, 2]));
255 assert_eq!(payload.0["nested"]["glob"], "*leaf");
256 }
257
258 #[test]
259 fn empty_yaml_body_deserializes_as_an_empty_object() {
260 let payload = match yaml_payload("") {
261 Ok(payload) => payload,
262 Err(err) => panic!("payload rejected: {}", err.0),
263 };
264
265 assert_eq!(payload.0, serde_json::json!({}));
266 }
267
268 #[test]
269 fn yaml_alias_expansion_is_rejected() {
270 let leaf = vec!["1"; 1000].join(",");
271 let aliases = vec!["*leaf"; 500].join(",");
272 let body = format!("seed: &leaf [{leaf}]\nx: [{aliases}]\n");
273
274 assert!(yaml_payload(body).is_err());
275 }
276
277 #[test]
278 fn yaml_merge_key_aliases_are_rejected() {
279 assert!(yaml_payload("base: &base\n a: 1\nderived:\n <<: *base\n b: 2\n").is_err());
280 }
281
282 #[test]
283 fn unused_yaml_anchors_are_accepted() {
284 let payload = match yaml_payload("seed: &leaf [1, 2]\nother: 3\n") {
285 Ok(payload) => payload,
286 Err(err) => panic!("payload rejected: {}", err.0),
287 };
288
289 assert_eq!(payload.0["seed"], serde_json::json!([1, 2]));
290 }
291}