#![cfg_attr(
not(test),
deny(
clippy::unwrap_used,
clippy::expect_used,
clippy::panic,
clippy::unreachable,
clippy::todo,
clippy::unimplemented,
clippy::indexing_slicing,
)
)]
use axum::extract::{FromRequest, FromRequestParts};
use axum::response::{IntoResponse, Response};
macro_rules! impl_extractor_deref {
($extractor:ident) => {
impl<T> std::ops::Deref for $extractor<T> {
type Target = T;
fn deref(&self) -> &Self::Target {
&self.0
}
}
impl<T> std::ops::DerefMut for $extractor<T> {
fn deref_mut(&mut self) -> &mut Self::Target {
&mut self.0
}
}
};
}
#[derive(Debug, Clone, Copy, Default)]
pub struct Form<T>(pub T);
impl_extractor_deref!(Form);
impl<S, T> FromRequest<S> for Form<T>
where
S: Send + Sync,
axum::extract::Form<T>: FromRequest<S, Rejection = axum::extract::rejection::FormRejection>,
{
type Rejection = crate::AutumnError;
async fn from_request(req: axum::extract::Request, state: &S) -> Result<Self, Self::Rejection> {
axum::extract::Form::from_request(req, state)
.await
.map(|axum::extract::Form(value)| Self(value))
.map_err(|err| rejection_to_error(err.status(), err.body_text()))
}
}
#[derive(Debug, Clone, Copy, Default)]
pub struct Json<T>(pub T);
impl_extractor_deref!(Json);
impl<S, T> FromRequest<S> for Json<T>
where
S: Send + Sync,
axum::extract::Json<T>: FromRequest<S, Rejection = axum::extract::rejection::JsonRejection>,
{
type Rejection = crate::AutumnError;
async fn from_request(req: axum::extract::Request, state: &S) -> Result<Self, Self::Rejection> {
axum::extract::Json::from_request(req, state)
.await
.map(|axum::extract::Json(value)| Self(value))
.map_err(|err| rejection_to_error(err.status(), err.body_text()))
}
}
impl<T> IntoResponse for Json<T>
where
axum::Json<T>: IntoResponse,
{
fn into_response(self) -> Response {
axum::Json(self.0).into_response()
}
}
#[derive(Debug, Clone, Copy, Default)]
pub struct Path<T>(pub T);
impl_extractor_deref!(Path);
impl<S, T> FromRequestParts<S> for Path<T>
where
S: Send + Sync,
axum::extract::Path<T>:
FromRequestParts<S, Rejection = axum::extract::rejection::PathRejection>,
{
type Rejection = crate::AutumnError;
async fn from_request_parts(
parts: &mut axum::http::request::Parts,
state: &S,
) -> Result<Self, Self::Rejection> {
axum::extract::Path::from_request_parts(parts, state)
.await
.map(|axum::extract::Path(value)| Self(value))
.map_err(|err| rejection_to_error(err.status(), err.body_text()))
}
}
#[derive(Debug, Clone, Copy, Default)]
pub struct Query<T>(pub T);
impl_extractor_deref!(Query);
impl<S, T> FromRequestParts<S> for Query<T>
where
S: Send + Sync,
axum::extract::Query<T>:
FromRequestParts<S, Rejection = axum::extract::rejection::QueryRejection>,
{
type Rejection = crate::AutumnError;
async fn from_request_parts(
parts: &mut axum::http::request::Parts,
state: &S,
) -> Result<Self, Self::Rejection> {
axum::extract::Query::from_request_parts(parts, state)
.await
.map(|axum::extract::Query(value)| Self(value))
.map_err(|err| rejection_to_error(err.status(), err.body_text()))
}
}
fn rejection_to_error(status: http::StatusCode, body_text: String) -> crate::AutumnError {
crate::AutumnError::bad_request_msg(body_text).with_status(status)
}
#[cfg(feature = "multipart")]
pub struct Multipart {
inner: axum::extract::Multipart,
config: crate::security::config::UploadConfig,
}
#[cfg(feature = "multipart")]
impl Multipart {
pub async fn next_field(&mut self) -> crate::AutumnResult<Option<MultipartField<'_>>> {
let Some(mut field) = self
.inner
.next_field()
.await
.map_err(|err| multipart_error_to_error(&err))?
else {
return Ok(None);
};
let needs_sniff = field.file_name().is_some()
&& (!self.config.allowed_mime_types.is_empty()
|| self.config.reject_on_content_type_mismatch);
let filename_is_empty = field.file_name().is_some_and(str::is_empty);
if !needs_sniff {
return Ok(Some(MultipartField::new(
field,
self.config.max_file_size_bytes,
)));
}
let mut prefix: Vec<u8> = Vec::with_capacity(SNIFF_PREFIX_BYTES);
while prefix.len() < SNIFF_PREFIX_BYTES {
match field
.chunk()
.await
.map_err(|err| multipart_error_to_error(&err))?
{
Some(chunk) => {
prefix.extend_from_slice(&chunk);
if prefix.len() > self.config.max_file_size_bytes {
return Err(file_too_large_error(self.config.max_file_size_bytes));
}
}
None => break,
}
}
let declared_essence = field.content_type().map(content_type_essence);
let sniffed = sniff_content_type(&prefix);
let is_empty_file_input = filename_is_empty && prefix.is_empty();
if !is_empty_file_input {
if !self.config.allowed_mime_types.is_empty() {
enforce_upload_allow_list(
&self.config.allowed_mime_types,
&prefix,
declared_essence,
sniffed,
)?;
}
if self.config.reject_on_content_type_mismatch {
enforce_content_type_match(declared_essence, sniffed)?;
}
}
Ok(Some(MultipartField {
inner: field,
max_file_size_bytes: self.config.max_file_size_bytes,
prefix,
sniffed_content_type: sniffed,
is_empty_optional_input: is_empty_file_input,
}))
}
}
#[cfg(feature = "multipart")]
fn sniff_content_type(prefix: &[u8]) -> Option<&'static str> {
infer::get(prefix).map(|kind| kind.mime_type())
}
#[cfg(feature = "multipart")]
fn content_type_essence(raw: &str) -> &str {
raw.split(';').next().unwrap_or("").trim()
}
#[cfg(feature = "multipart")]
fn prefix_looks_like_markup(prefix: &[u8]) -> bool {
let bytes = prefix.strip_prefix(&[0xEF, 0xBB, 0xBF]).unwrap_or(prefix);
bytes.iter().find(|byte| !byte.is_ascii_whitespace()) == Some(&b'<')
}
#[cfg(feature = "multipart")]
fn truncate_for_error(value: &str) -> String {
const MAX_CHARS: usize = 128;
value.chars().take(MAX_CHARS).collect()
}
#[cfg(feature = "multipart")]
fn is_signatureless_text_type(essence: &str) -> bool {
let lower = essence.to_ascii_lowercase();
lower.starts_with("text/") || matches!(lower.as_str(), "application/json" | "application/csv")
}
#[cfg(feature = "multipart")]
fn enforce_upload_allow_list(
allowed: &[String],
prefix: &[u8],
declared_essence: Option<&str>,
sniffed: Option<&str>,
) -> crate::AutumnResult<()> {
let in_list = |value: &str| {
allowed
.iter()
.any(|entry| entry.eq_ignore_ascii_case(value))
};
if let Some(sniffed) = sniffed {
if !in_list(sniffed) {
return Err(crate::AutumnError::bad_request_msg(format!(
"upload content type not allowed: sniffed={sniffed}"
)));
}
} else if prefix_looks_like_markup(prefix) {
return Err(crate::AutumnError::bad_request_msg(
"upload rejected: unrecognized file content looks like markup",
));
} else if !matches!(
declared_essence,
Some(essence) if is_signatureless_text_type(essence) && in_list(essence)
) {
return Err(crate::AutumnError::bad_request_msg(format!(
"upload content type could not be verified from its content: declared={}",
declared_essence.map_or_else(String::new, truncate_for_error)
)));
}
Ok(())
}
#[cfg(feature = "multipart")]
fn enforce_content_type_match(
declared_essence: Option<&str>,
sniffed: Option<&str>,
) -> crate::AutumnResult<()> {
let Some(declared_essence) = declared_essence else {
return Err(crate::AutumnError::bad_request_msg(
"content-type mismatch check enabled but the upload declared no content type",
));
};
match sniffed {
Some(sniffed) if declared_essence.eq_ignore_ascii_case(sniffed) => Ok(()),
Some(sniffed) => Err(crate::AutumnError::bad_request_msg(format!(
"declared content type {} does not match sniffed content type {sniffed}",
truncate_for_error(declared_essence)
))),
None => Err(crate::AutumnError::bad_request_msg(format!(
"cannot verify declared content type {}: file content is unrecognized",
truncate_for_error(declared_essence)
))),
}
}
#[cfg(feature = "multipart")]
impl<S> axum::extract::FromRequest<S> for Multipart
where
S: Send + Sync,
axum::extract::Multipart:
axum::extract::FromRequest<S, Rejection = axum::extract::multipart::MultipartRejection>,
{
type Rejection = crate::AutumnError;
async fn from_request(
mut req: axum::extract::Request,
state: &S,
) -> Result<Self, Self::Rejection> {
let config = req
.extensions()
.get::<crate::security::config::UploadConfig>()
.cloned()
.unwrap_or_default();
axum::extract::DefaultBodyLimit::max(config.max_request_size_bytes).apply(&mut req);
let inner = axum::extract::Multipart::from_request(req, state)
.await
.map_err(|err| multipart_rejection_to_error(&err))?;
Ok(Self { inner, config })
}
}
#[cfg(feature = "multipart")]
const SNIFF_PREFIX_BYTES: usize = 512;
#[cfg(feature = "multipart")]
pub struct MultipartField<'a> {
inner: axum::extract::multipart::Field<'a>,
max_file_size_bytes: usize,
prefix: Vec<u8>,
sniffed_content_type: Option<&'static str>,
is_empty_optional_input: bool,
}
#[cfg(all(feature = "multipart", feature = "storage"))]
struct MultipartFieldStreamState<'a> {
inner: axum::extract::multipart::Field<'a>,
prefix: Option<bytes::Bytes>,
total: usize,
max: usize,
errored: bool,
}
#[cfg(feature = "multipart")]
#[allow(clippy::elidable_lifetime_names)]
impl<'a> MultipartField<'a> {
const fn new(inner: axum::extract::multipart::Field<'a>, max_file_size_bytes: usize) -> Self {
Self {
inner,
max_file_size_bytes,
prefix: Vec::new(),
sniffed_content_type: None,
is_empty_optional_input: false,
}
}
#[must_use]
pub fn name(&self) -> Option<&str> {
self.inner.name()
}
#[must_use]
pub const fn sniffed_content_type(&self) -> Option<&str> {
self.sniffed_content_type
}
#[must_use]
pub fn file_name(&self) -> Option<&str> {
if self.is_empty_optional_input {
None
} else {
self.inner.file_name()
}
}
#[must_use]
pub fn content_type(&self) -> Option<&str> {
self.inner.content_type()
}
#[must_use]
pub fn with_max_bytes(mut self, max: usize) -> Self {
self.max_file_size_bytes = self.max_file_size_bytes.min(max);
self
}
pub async fn bytes_limited(mut self) -> crate::AutumnResult<Vec<u8>> {
let mut out = self.prefix;
let mut read = out.len();
if read > self.max_file_size_bytes {
return Err(file_too_large_error(self.max_file_size_bytes));
}
while let Some(chunk) = self
.inner
.chunk()
.await
.map_err(|err| multipart_error_to_error(&err))?
{
read += chunk.len();
if read > self.max_file_size_bytes {
return Err(file_too_large_error(self.max_file_size_bytes));
}
out.extend_from_slice(&chunk);
}
Ok(out)
}
#[cfg(feature = "storage")]
pub async fn save_to_blob_store<'b>(
self,
store: &'b (dyn crate::storage::BlobStore + '_),
key: impl Into<String>,
) -> crate::AutumnResult<crate::storage::Blob>
where
'a: 'b,
{
let key = key.into();
let content_type = self
.sniffed_content_type
.map(str::to_owned)
.or_else(|| self.inner.content_type().map(str::to_owned))
.unwrap_or_else(|| "application/octet-stream".to_owned());
let prefix = if self.prefix.is_empty() {
None
} else {
Some(bytes::Bytes::from(self.prefix))
};
let state = MultipartFieldStreamState {
inner: self.inner,
prefix,
total: 0,
max: self.max_file_size_bytes,
errored: false,
};
let stream = futures::stream::unfold(state, |mut state| async move {
if state.errored {
return None;
}
if let Some(prefix) = state.prefix.take() {
state.total = state.total.saturating_add(prefix.len());
if state.total > state.max {
let err = crate::storage::BlobStoreError::PayloadTooLarge(format!(
"uploaded file exceeds limit of {} bytes",
state.max,
));
state.errored = true;
return Some((Err(err), state));
}
return Some((Ok(prefix), state));
}
match state.inner.chunk().await {
Ok(Some(chunk)) => {
state.total = state.total.saturating_add(chunk.len());
if state.total > state.max {
let err = crate::storage::BlobStoreError::PayloadTooLarge(format!(
"uploaded file exceeds limit of {} bytes",
state.max,
));
state.errored = true;
Some((Err(err), state))
} else {
Some((Ok(chunk), state))
}
}
Ok(None) => None,
Err(err) => {
state.errored = true;
let mapped = blob_error_from_multipart(&err);
Some((Err(mapped), state))
}
}
});
let stream: crate::storage::ByteStream<'b> = Box::pin(stream);
store
.put_stream(&key, &content_type, stream)
.await
.map_err(crate::storage::BlobStoreError::into_autumn_error)
}
pub async fn save_to<P: AsRef<std::path::Path>>(
mut self,
path: P,
) -> crate::AutumnResult<usize> {
use tokio::io::AsyncWriteExt as _;
let path = path.as_ref();
let mut file = tokio::fs::File::create(path)
.await
.map_err(crate::AutumnError::internal_server_error)?;
let mut written = 0usize;
if !self.prefix.is_empty() {
written += self.prefix.len();
if written > self.max_file_size_bytes {
drop(file);
let _ = tokio::fs::remove_file(path).await;
return Err(file_too_large_error(self.max_file_size_bytes));
}
file.write_all(&self.prefix)
.await
.map_err(crate::AutumnError::internal_server_error)?;
}
while let Some(chunk) = self
.inner
.chunk()
.await
.map_err(|err| multipart_error_to_error(&err))?
{
written += chunk.len();
if written > self.max_file_size_bytes {
drop(file);
let _ = tokio::fs::remove_file(path).await;
return Err(file_too_large_error(self.max_file_size_bytes));
}
file.write_all(&chunk)
.await
.map_err(crate::AutumnError::internal_server_error)?;
}
file.flush()
.await
.map_err(crate::AutumnError::internal_server_error)?;
Ok(written)
}
}
#[cfg(feature = "multipart")]
fn multipart_rejection_to_error(
err: &axum::extract::multipart::MultipartRejection,
) -> crate::AutumnError {
crate::AutumnError::bad_request_msg(err.body_text()).with_status(err.status())
}
#[cfg(feature = "multipart")]
#[cfg(all(feature = "multipart", feature = "storage"))]
fn blob_error_from_multipart(
err: &axum::extract::multipart::MultipartError,
) -> crate::storage::BlobStoreError {
let status = err.status();
let body = err.body_text();
if status == http::StatusCode::PAYLOAD_TOO_LARGE {
crate::storage::BlobStoreError::PayloadTooLarge(body)
} else if status.is_client_error() {
crate::storage::BlobStoreError::InvalidInput(body)
} else {
crate::storage::BlobStoreError::Io(body)
}
}
#[cfg(feature = "multipart")]
fn multipart_error_to_error(err: &axum::extract::multipart::MultipartError) -> crate::AutumnError {
crate::AutumnError::bad_request_msg(err.body_text()).with_status(err.status())
}
#[cfg(feature = "multipart")]
fn file_too_large_error(max_file_size_bytes: usize) -> crate::AutumnError {
crate::AutumnError::bad_request_msg(format!(
"uploaded file exceeds limit of {max_file_size_bytes} bytes",
))
.with_status(http::StatusCode::PAYLOAD_TOO_LARGE)
}
pub use axum::extract::State;
#[cfg(all(test, feature = "multipart"))]
mod tests {
use super::*;
use axum::extract::FromRequest;
use axum::http::Request;
#[tokio::test]
async fn test_multipart_field_bytes_limited_success() {
let body = "--boundary\r\nContent-Disposition: form-data; name=\"file\"; filename=\"test.txt\"\r\n\r\nhello\r\n--boundary--\r\n";
let req = Request::builder()
.header("content-type", "multipart/form-data; boundary=boundary")
.body(axum::body::Body::from(body))
.unwrap();
let mut multipart = axum::extract::Multipart::from_request(req, &())
.await
.unwrap();
let field = multipart.next_field().await.unwrap().unwrap();
let wrapper = MultipartField::new(field, 100);
let bytes = wrapper.bytes_limited().await.unwrap();
assert_eq!(bytes, b"hello");
}
#[tokio::test]
async fn test_multipart_field_bytes_limited_too_large() {
let body = "--boundary\r\nContent-Disposition: form-data; name=\"file\"; filename=\"test.txt\"\r\n\r\nhello world\r\n--boundary--\r\n";
let req = Request::builder()
.header("content-type", "multipart/form-data; boundary=boundary")
.body(axum::body::Body::from(body))
.unwrap();
let mut multipart = axum::extract::Multipart::from_request(req, &())
.await
.unwrap();
let field = multipart.next_field().await.unwrap().unwrap();
let wrapper = MultipartField::new(field, 5);
let err = wrapper.bytes_limited().await.unwrap_err();
assert_eq!(err.status(), http::StatusCode::PAYLOAD_TOO_LARGE);
}
#[tokio::test]
async fn test_multipart_field_save_to_success() {
let body = "--boundary\r\nContent-Disposition: form-data; name=\"file\"; filename=\"test.txt\"\r\n\r\nfile content\r\n--boundary--\r\n";
let req = Request::builder()
.header("content-type", "multipart/form-data; boundary=boundary")
.body(axum::body::Body::from(body))
.unwrap();
let mut multipart = axum::extract::Multipart::from_request(req, &())
.await
.unwrap();
let field = multipart.next_field().await.unwrap().unwrap();
let wrapper = MultipartField::new(field, 100);
let dir = tempfile::tempdir().unwrap();
let file_path = dir.path().join("out.txt");
let written = wrapper.save_to(&file_path).await.unwrap();
assert_eq!(written, 12);
let content = std::fs::read_to_string(&file_path).unwrap();
assert_eq!(content, "file content");
}
#[tokio::test]
async fn test_multipart_field_save_to_too_large() {
let body = "--boundary\r\nContent-Disposition: form-data; name=\"file\"; filename=\"test.txt\"\r\n\r\nfile content\r\n--boundary--\r\n";
let req = Request::builder()
.header("content-type", "multipart/form-data; boundary=boundary")
.body(axum::body::Body::from(body))
.unwrap();
let mut multipart = axum::extract::Multipart::from_request(req, &())
.await
.unwrap();
let field = multipart.next_field().await.unwrap().unwrap();
let wrapper = MultipartField::new(field, 4);
let dir = tempfile::tempdir().unwrap();
let file_path = dir.path().join("out_large.txt");
let err = wrapper.save_to(&file_path).await.unwrap_err();
assert_eq!(err.status(), http::StatusCode::PAYLOAD_TOO_LARGE);
assert!(!file_path.exists());
}
#[cfg(feature = "storage")]
#[tokio::test]
async fn test_multipart_field_save_to_blob_store_success() {
use crate::storage::{BlobStore, LocalBlobStore, local::SigningKey};
use std::time::Duration;
let body = "--boundary\r\nContent-Disposition: form-data; name=\"file\"; filename=\"test.txt\"\r\nContent-Type: text/plain\r\n\r\nblob content\r\n--boundary--\r\n";
let req = Request::builder()
.header("content-type", "multipart/form-data; boundary=boundary")
.body(axum::body::Body::from(body))
.unwrap();
let mut multipart = axum::extract::Multipart::from_request(req, &())
.await
.unwrap();
let field = multipart.next_field().await.unwrap().unwrap();
let wrapper = MultipartField::new(field, 100);
let root = tempfile::tempdir().unwrap();
let store = LocalBlobStore::new(
"local",
root.path(),
"/blobs",
Duration::from_secs(3600),
SigningKey::random(),
vec![],
)
.unwrap();
let blob = wrapper.save_to_blob_store(&store, "myblob").await.unwrap();
assert_eq!(blob.key, "myblob");
assert_eq!(blob.content_type, "text/plain");
let bytes = store.get("myblob").await.unwrap();
assert_eq!(&bytes[..], b"blob content");
}
#[cfg(feature = "storage")]
#[tokio::test]
async fn test_multipart_field_save_to_blob_store_too_large() {
use crate::storage::{BlobStore, LocalBlobStore, local::SigningKey};
use std::time::Duration;
let body = "--boundary\r\nContent-Disposition: form-data; name=\"file\"; filename=\"test.txt\"\r\nContent-Type: text/plain\r\n\r\nblob content\r\n--boundary--\r\n";
let req = Request::builder()
.header("content-type", "multipart/form-data; boundary=boundary")
.body(axum::body::Body::from(body))
.unwrap();
let mut multipart = axum::extract::Multipart::from_request(req, &())
.await
.unwrap();
let field = multipart.next_field().await.unwrap().unwrap();
let wrapper = MultipartField::new(field, 4);
let root = tempfile::tempdir().unwrap();
let store = LocalBlobStore::new(
"local",
root.path(),
"/blobs",
Duration::from_secs(3600),
SigningKey::random(),
vec![],
)
.unwrap();
let err = wrapper
.save_to_blob_store(&store, "myblob")
.await
.unwrap_err();
assert_eq!(err.status(), http::StatusCode::PAYLOAD_TOO_LARGE);
let get_err = store.get("myblob").await.unwrap_err();
assert_eq!(get_err.status(), http::StatusCode::NOT_FOUND);
}
#[tokio::test]
async fn test_multipart_field_metadata() {
let body = "--boundary\r\nContent-Disposition: form-data; name=\"custom_name\"; filename=\"custom_file.png\"\r\nContent-Type: image/png\r\n\r\npng\r\n--boundary--\r\n";
let req = Request::builder()
.header("content-type", "multipart/form-data; boundary=boundary")
.body(axum::body::Body::from(body))
.unwrap();
let mut multipart = axum::extract::Multipart::from_request(req, &())
.await
.unwrap();
let field = multipart.next_field().await.unwrap().unwrap();
let wrapper = MultipartField::new(field, 100);
assert_eq!(wrapper.name(), Some("custom_name"));
assert_eq!(wrapper.file_name(), Some("custom_file.png"));
assert_eq!(wrapper.content_type(), Some("image/png"));
let tighter = wrapper.with_max_bytes(50);
assert_eq!(tighter.max_file_size_bytes, 50);
let not_tighter = tighter.with_max_bytes(200);
assert_eq!(not_tighter.max_file_size_bytes, 50); }
#[test]
fn sniff_content_type_recognizes_png() {
let png = [0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A, 0x00, 0x00];
assert_eq!(sniff_content_type(&png), Some("image/png"));
}
#[test]
fn sniff_content_type_unknown_is_none() {
assert_eq!(sniff_content_type(&[0x01, 0x02]), None);
assert_eq!(sniff_content_type(&[]), None);
}
#[test]
fn content_type_essence_strips_parameters() {
assert_eq!(
content_type_essence("image/png; charset=binary"),
"image/png"
);
assert_eq!(content_type_essence(" text/csv "), "text/csv");
assert_eq!(content_type_essence("application/json"), "application/json");
assert_eq!(content_type_essence(""), "");
}
#[test]
fn is_signatureless_text_type_is_case_insensitive() {
assert!(is_signatureless_text_type("text/csv"));
assert!(is_signatureless_text_type("Text/Csv"));
assert!(is_signatureless_text_type("TEXT/PLAIN"));
assert!(is_signatureless_text_type("application/json"));
assert!(is_signatureless_text_type("Application/JSON"));
assert!(is_signatureless_text_type("application/csv"));
assert!(!is_signatureless_text_type("image/png"));
assert!(!is_signatureless_text_type("Image/PNG"));
assert!(!is_signatureless_text_type("application/octet-stream"));
assert!(!is_signatureless_text_type("application/pdf"));
}
#[test]
fn prefix_looks_like_markup_detects_leading_angle_bracket() {
assert!(prefix_looks_like_markup(b"<!DOCTYPE html>"));
assert!(prefix_looks_like_markup(b"<svg onload=alert(1)>"));
assert!(prefix_looks_like_markup(b"<?xml version=\"1.0\"?>"));
assert!(prefix_looks_like_markup(b" \n\t<html>"));
assert!(prefix_looks_like_markup(&[0xEF, 0xBB, 0xBF, b'<', b'a']));
assert!(!prefix_looks_like_markup(b"a,b\n1,2"));
assert!(!prefix_looks_like_markup(br#"{"a":1}"#));
assert!(!prefix_looks_like_markup(&[0x01, 0x02]));
assert!(!prefix_looks_like_markup(b""));
}
#[tokio::test]
async fn next_field_exposes_sniffed_content_type_for_genuine_png() {
let mut body: Vec<u8> = Vec::new();
body.extend_from_slice(
b"--boundary\r\nContent-Disposition: form-data; name=\"file\"; \
filename=\"real.png\"\r\nContent-Type: application/octet-stream\r\n\r\n",
);
body.extend_from_slice(&[
0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A, 0x00, 0x00, 0x00, 0x0D, 0x49, 0x48,
0x44, 0x52,
]);
body.extend_from_slice(b"\r\n--boundary--\r\n");
let req = Request::builder()
.header("content-type", "multipart/form-data; boundary=boundary")
.body(axum::body::Body::from(body))
.unwrap();
let inner = axum::extract::Multipart::from_request(req, &())
.await
.unwrap();
let mut multipart = Multipart {
inner,
config: crate::security::config::UploadConfig {
allowed_mime_types: vec!["image/png".to_owned()],
..crate::security::config::UploadConfig::default()
},
};
let field = multipart.next_field().await.unwrap().unwrap();
assert_eq!(field.sniffed_content_type(), Some("image/png"));
let bytes = field.bytes_limited().await.unwrap();
assert_eq!(bytes.len(), 16);
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CurrentPath(pub String);
impl CurrentPath {
#[must_use]
pub fn as_str(&self) -> &str {
&self.0
}
}
impl<S> FromRequestParts<S> for CurrentPath
where
S: Send + Sync,
{
type Rejection = std::convert::Infallible;
async fn from_request_parts(
parts: &mut axum::http::request::Parts,
state: &S,
) -> Result<Self, Self::Rejection> {
let axum::extract::OriginalUri(uri) =
axum::extract::OriginalUri::from_request_parts(parts, state)
.await
.unwrap();
Ok(Self(uri.path().to_owned()))
}
}
use crate::security::trusted_proxies::ResolvedClientIdentity;
pub struct ClientAddr(pub std::net::IpAddr);
impl ClientAddr {
#[must_use]
pub const fn ip(&self) -> std::net::IpAddr {
self.0
}
}
impl<S> FromRequestParts<S> for ClientAddr
where
S: Send + Sync,
{
type Rejection = (axum::http::StatusCode, &'static str);
async fn from_request_parts(
parts: &mut axum::http::request::Parts,
_state: &S,
) -> Result<Self, Self::Rejection> {
parts
.extensions
.get::<ResolvedClientIdentity>()
.and_then(|id| id.addr)
.map(ClientAddr)
.ok_or((
axum::http::StatusCode::INTERNAL_SERVER_ERROR,
"ClientAddr not resolved. Is the TrustedProxiesLayer installed?",
))
}
}
impl<S> axum::extract::OptionalFromRequestParts<S> for ClientAddr
where
S: Send + Sync,
{
type Rejection = std::convert::Infallible;
async fn from_request_parts(
parts: &mut axum::http::request::Parts,
_state: &S,
) -> Result<Option<Self>, Self::Rejection> {
Ok(parts
.extensions
.get::<ResolvedClientIdentity>()
.and_then(|id| id.addr)
.map(ClientAddr))
}
}
pub struct ClientHost(pub String);
impl ClientHost {
#[must_use]
pub fn as_str(&self) -> &str {
&self.0
}
}
impl<S> FromRequestParts<S> for ClientHost
where
S: Send + Sync,
{
type Rejection = (axum::http::StatusCode, &'static str);
async fn from_request_parts(
parts: &mut axum::http::request::Parts,
_state: &S,
) -> Result<Self, Self::Rejection> {
parts
.extensions
.get::<ResolvedClientIdentity>()
.and_then(|id| id.host.clone())
.map(ClientHost)
.ok_or((
axum::http::StatusCode::INTERNAL_SERVER_ERROR,
"ClientHost not resolved. Is the TrustedProxiesLayer installed?",
))
}
}
impl<S> axum::extract::OptionalFromRequestParts<S> for ClientHost
where
S: Send + Sync,
{
type Rejection = std::convert::Infallible;
async fn from_request_parts(
parts: &mut axum::http::request::Parts,
_state: &S,
) -> Result<Option<Self>, Self::Rejection> {
Ok(parts
.extensions
.get::<ResolvedClientIdentity>()
.and_then(|id| id.host.clone())
.map(ClientHost))
}
}
pub struct ClientScheme(pub String);
impl ClientScheme {
#[must_use]
pub fn as_str(&self) -> &str {
&self.0
}
#[must_use]
pub fn is_https(&self) -> bool {
self.0.eq_ignore_ascii_case("https")
}
}
impl<S> FromRequestParts<S> for ClientScheme
where
S: Send + Sync,
{
type Rejection = (axum::http::StatusCode, &'static str);
async fn from_request_parts(
parts: &mut axum::http::request::Parts,
_state: &S,
) -> Result<Self, Self::Rejection> {
parts
.extensions
.get::<ResolvedClientIdentity>()
.map(|id| Self(id.scheme.clone().unwrap_or_else(|| "http".to_owned())))
.ok_or((
axum::http::StatusCode::INTERNAL_SERVER_ERROR,
"ClientScheme not resolved. Is the TrustedProxiesLayer installed?",
))
}
}
impl<S> axum::extract::OptionalFromRequestParts<S> for ClientScheme
where
S: Send + Sync,
{
type Rejection = std::convert::Infallible;
async fn from_request_parts(
parts: &mut axum::http::request::Parts,
_state: &S,
) -> Result<Option<Self>, Self::Rejection> {
Ok(parts
.extensions
.get::<ResolvedClientIdentity>()
.map(|id| Self(id.scheme.clone().unwrap_or_else(|| "http".to_owned()))))
}
}
#[cfg(test)]
mod trusted_proxy_extractor_tests {
use super::*;
use axum::Router;
use axum::body::Body;
use axum::routing::get;
use tower::ServiceExt;
fn make_identity(addr: &str, host: &str, scheme: &str) -> ResolvedClientIdentity {
ResolvedClientIdentity {
addr: Some(addr.parse().unwrap()),
host: Some(host.to_owned()),
scheme: Some(scheme.to_owned()),
}
}
#[tokio::test]
async fn client_addr_extractor_reads_from_extension() {
async fn handler(ClientAddr(ip): ClientAddr) -> String {
ip.to_string()
}
let app = Router::new().route("/", get(handler));
let mut req = axum::http::Request::builder()
.uri("/")
.body(Body::empty())
.unwrap();
req.extensions_mut()
.insert(make_identity("192.0.2.1", "app.example", "https"));
let resp = app.oneshot(req).await.unwrap();
assert_eq!(resp.status(), axum::http::StatusCode::OK);
let body = axum::body::to_bytes(resp.into_body(), 64).await.unwrap();
assert_eq!(&body[..], b"192.0.2.1");
}
#[tokio::test]
async fn client_host_extractor_reads_from_extension() {
async fn handler(ClientHost(host): ClientHost) -> String {
host
}
let app = Router::new().route("/", get(handler));
let mut req = axum::http::Request::builder()
.uri("/")
.body(Body::empty())
.unwrap();
req.extensions_mut()
.insert(make_identity("192.0.2.1", "app.example", "https"));
let resp = app.oneshot(req).await.unwrap();
let body = axum::body::to_bytes(resp.into_body(), 64).await.unwrap();
assert_eq!(&body[..], b"app.example");
}
#[tokio::test]
async fn client_scheme_extractor_reads_from_extension() {
async fn handler(ClientScheme(scheme): ClientScheme) -> String {
scheme
}
let app = Router::new().route("/", get(handler));
let mut req = axum::http::Request::builder()
.uri("/")
.body(Body::empty())
.unwrap();
req.extensions_mut()
.insert(make_identity("192.0.2.1", "app.example", "https"));
let resp = app.oneshot(req).await.unwrap();
let body = axum::body::to_bytes(resp.into_body(), 64).await.unwrap();
assert_eq!(&body[..], b"https");
}
#[tokio::test]
async fn client_addr_missing_returns_500() {
async fn handler(_: ClientAddr) -> &'static str {
"ok"
}
let app = Router::new().route("/", get(handler));
let req = axum::http::Request::builder()
.uri("/")
.body(Body::empty())
.unwrap();
let resp = app.oneshot(req).await.unwrap();
assert_eq!(resp.status(), axum::http::StatusCode::INTERNAL_SERVER_ERROR);
}
#[tokio::test]
async fn optional_client_addr_returns_none_when_missing() {
async fn handler(addr: Option<ClientAddr>) -> String {
if addr.is_some() {
"some".to_owned()
} else {
"none".to_owned()
}
}
let app = Router::new().route("/", get(handler));
let req = axum::http::Request::builder()
.uri("/")
.body(Body::empty())
.unwrap();
let resp = app.oneshot(req).await.unwrap();
let body = axum::body::to_bytes(resp.into_body(), 64).await.unwrap();
assert_eq!(&body[..], b"none");
}
#[tokio::test]
async fn current_path_extracts_uri_path() {
async fn handler(CurrentPath(path): CurrentPath) -> String {
path
}
let app = Router::new().route("/admin/posts", get(handler));
let req = axum::http::Request::builder()
.uri("/admin/posts")
.body(Body::empty())
.unwrap();
let resp = app.oneshot(req).await.unwrap();
let body = axum::body::to_bytes(resp.into_body(), 64).await.unwrap();
assert_eq!(&body[..], b"/admin/posts");
}
#[tokio::test]
async fn current_path_strips_query_string() {
async fn handler(CurrentPath(path): CurrentPath) -> String {
path
}
let app = Router::new().route("/admin/posts", get(handler));
let req = axum::http::Request::builder()
.uri("/admin/posts?page=2")
.body(Body::empty())
.unwrap();
let resp = app.oneshot(req).await.unwrap();
let body = axum::body::to_bytes(resp.into_body(), 64).await.unwrap();
assert_eq!(&body[..], b"/admin/posts");
}
#[tokio::test]
async fn current_path_preserves_prefix_under_nested_router() {
async fn handler(CurrentPath(path): CurrentPath) -> String {
path
}
let inner = Router::new().route("/posts", get(handler));
let app = Router::new().nest("/admin", inner);
let req = axum::http::Request::builder()
.uri("/admin/posts")
.body(Body::empty())
.unwrap();
let resp = app.oneshot(req).await.unwrap();
let body = axum::body::to_bytes(resp.into_body(), 64).await.unwrap();
assert_eq!(&body[..], b"/admin/posts");
}
}