axum: factor out non-generic parts of request conversion

This saves ~10% in binary size!
This commit is contained in:
Lambda 2024-05-31 10:23:49 +00:00
parent 624ff57414
commit bf1d54defc

View file

@ -15,12 +15,13 @@ use axum_extra::{
typed_header::TypedHeaderRejectionReason, typed_header::TypedHeaderRejectionReason,
TypedHeader, TypedHeader,
}; };
use bytes::{BufMut, BytesMut}; use bytes::{BufMut, Bytes, BytesMut};
use http::{Request, StatusCode}; use http::{Request, StatusCode};
use http_body_util::BodyExt; use http_body_util::BodyExt;
use ruma::{ use ruma::{
api::{ api::{
client::error::ErrorKind, AuthScheme, IncomingRequest, OutgoingResponse, client::error::ErrorKind, AuthScheme, IncomingRequest, Metadata,
OutgoingResponse,
}, },
CanonicalJsonValue, OwnedDeviceId, OwnedServerName, OwnedUserId, UserId, CanonicalJsonValue, OwnedDeviceId, OwnedServerName, OwnedUserId, UserId,
}; };
@ -37,18 +38,24 @@ enum Token {
None, None,
} }
#[async_trait] /// Return value of [`ar_from_request_inner()`], used to construct an [`Ar`].
impl<T, S> FromRequest<S> for Ar<T> struct ArPieces {
where sender_user: Option<OwnedUserId>,
T: IncomingRequest, sender_device: Option<OwnedDeviceId>,
{ sender_servername: Option<OwnedServerName>,
type Rejection = Error; json_body: Option<CanonicalJsonValue>,
appservice_info: Option<RegistrationInfo>,
path_params: Path<Vec<String>>,
http_request: Request<Bytes>,
}
/// Non-generic part of [`Ar::from_request()`]. Splitting this out reduces
/// binary size by ~10%.
#[allow(clippy::too_many_lines)] #[allow(clippy::too_many_lines)]
async fn from_request( async fn ar_from_request_inner(
req: axum::extract::Request, req: axum::extract::Request,
_state: &S, metadata: Metadata,
) -> Result<Self, Self::Rejection> { ) -> Result<ArPieces> {
#[derive(Deserialize)] #[derive(Deserialize)]
struct QueryParams { struct QueryParams {
access_token: Option<String>, access_token: Option<String>,
@ -68,7 +75,6 @@ where
(parts, body) (parts, body)
}; };
let metadata = T::METADATA;
let auth_header: Option<TypedHeader<Authorization<Bearer>>> = let auth_header: Option<TypedHeader<Authorization<Bearer>>> =
parts.extract().await?; parts.extract().await?;
let path_params: Path<Vec<String>> = parts.extract().await?; let path_params: Path<Vec<String>> = parts.extract().await?;
@ -177,10 +183,7 @@ where
.extract::<TypedHeader<Authorization<XMatrix>>>() .extract::<TypedHeader<Authorization<XMatrix>>>()
.await .await
.map_err(|e| { .map_err(|e| {
warn!( warn!("Missing or invalid Authorization header: {}", e);
"Missing or invalid Authorization header: {}",
e
);
let msg = match e.reason() { let msg = match e.reason() {
TypedHeaderRejectionReason::Missing => { TypedHeaderRejectionReason::Missing => {
@ -208,9 +211,7 @@ where
let mut request_map = BTreeMap::from_iter([ let mut request_map = BTreeMap::from_iter([
( (
"method".to_owned(), "method".to_owned(),
CanonicalJsonValue::String( CanonicalJsonValue::String(parts.method.to_string()),
parts.method.to_string(),
),
), ),
( (
"uri".to_owned(), "uri".to_owned(),
@ -239,8 +240,7 @@ where
]); ]);
if let Some(json_body) = &json_body { if let Some(json_body) = &json_body {
request_map request_map.insert("content".to_owned(), json_body.clone());
.insert("content".to_owned(), json_body.clone());
}; };
let keys_result = services() let keys_result = services()
@ -268,24 +268,20 @@ where
keys, keys,
)]); )]);
match ruma::signatures::verify_json( match ruma::signatures::verify_json(&pub_key_map, &request_map)
&pub_key_map, {
&request_map,
) {
Ok(()) => (None, None, Some(x_matrix.origin), None), Ok(()) => (None, None, Some(x_matrix.origin), None),
Err(e) => { Err(e) => {
warn!( warn!(
"Failed to verify json request from {}: \ "Failed to verify json request from {}: {}\n{:?}",
{}\n{:?}",
x_matrix.origin, e, request_map x_matrix.origin, e, request_map
); );
if parts.uri.to_string().contains('@') { if parts.uri.to_string().contains('@') {
warn!( warn!(
"Request uri contained '@' character. \ "Request uri contained '@' character. Make \
Make sure your reverse proxy gives \ sure your reverse proxy gives Grapevine the \
Grapevine the raw uri (apache: use \ raw uri (apache: use nocanon)"
nocanon)"
); );
} }
@ -308,8 +304,7 @@ where
) => { ) => {
return Err(Error::BadRequest( return Err(Error::BadRequest(
ErrorKind::Unauthorized, ErrorKind::Unauthorized,
"Only server signatures should be used on this \ "Only server signatures should be used on this endpoint.",
endpoint.",
)); ));
} }
(AuthScheme::AppserviceToken, Token::User(_)) => { (AuthScheme::AppserviceToken, Token::User(_)) => {
@ -327,10 +322,7 @@ where
if let Some(CanonicalJsonValue::Object(json_body)) = &mut json_body { if let Some(CanonicalJsonValue::Object(json_body)) = &mut json_body {
let user_id = sender_user.clone().unwrap_or_else(|| { let user_id = sender_user.clone().unwrap_or_else(|| {
UserId::parse_with_server_name( UserId::parse_with_server_name("", services().globals.server_name())
"",
services().globals.server_name(),
)
.expect("we know this is valid") .expect("we know this is valid")
}); });
@ -347,8 +339,7 @@ where
) )
}); });
if let Some(CanonicalJsonValue::Object(initial_request)) = if let Some(CanonicalJsonValue::Object(initial_request)) = uiaa_request
uiaa_request
{ {
for (key, value) in initial_request { for (key, value) in initial_request {
json_body.entry(key).or_insert(value); json_body.entry(key).or_insert(value);
@ -360,15 +351,39 @@ where
.expect("value serialization can't fail"); .expect("value serialization can't fail");
body = buf.into_inner().freeze(); body = buf.into_inner().freeze();
} }
let http_request = http_request.body(body).unwrap();
let http_request = http_request.body(&*body).unwrap();
debug!("{:?}", http_request); debug!("{:?}", http_request);
let body = T::try_from_http_request(http_request, &path_params) Ok(ArPieces {
sender_user,
sender_device,
sender_servername,
json_body,
appservice_info,
path_params,
http_request,
})
}
#[async_trait]
impl<T, S> FromRequest<S> for Ar<T>
where
T: IncomingRequest,
{
type Rejection = Error;
async fn from_request(
req: axum::extract::Request,
_state: &S,
) -> Result<Self, Self::Rejection> {
let pieces = ar_from_request_inner(req, T::METADATA).await?;
let body =
T::try_from_http_request(pieces.http_request, &pieces.path_params)
.map_err(|e| { .map_err(|e| {
warn!("try_from_http_request failed: {:?}", e); warn!("try_from_http_request failed: {:?}", e);
debug!("JSON body: {:?}", json_body); debug!("JSON body: {:?}", pieces.json_body);
Error::BadRequest( Error::BadRequest(
ErrorKind::BadJson, ErrorKind::BadJson,
"Failed to deserialize request.", "Failed to deserialize request.",
@ -377,11 +392,11 @@ where
Ok(Ar { Ok(Ar {
body, body,
sender_user, sender_user: pieces.sender_user,
sender_device, sender_device: pieces.sender_device,
sender_servername, sender_servername: pieces.sender_servername,
json_body, json_body: pieces.json_body,
appservice_info, appservice_info: pieces.appservice_info,
}) })
} }
} }