1
use axum::{
2
    extract::{
3
        rejection::JsonRejection, FromRequest, FromRequestParts, Json as AxumJson,
4
        Path as AxumPath, Query as AxumQuery, Request,
5
    },
6
    http::{header, request::Parts},
7
    response::{IntoResponse, Response},
8
};
9
use bytes::{BufMut, BytesMut};
10
use serde::{de::DeserializeOwned, Serialize};
11

            
12
use crate::{constants::ContentType, err::ErrResp};
13

            
14
/// JSON Extractor / Response.
15
///
16
/// This is the customized [`axum::extract::Json`] version to respose error with [`ErrResp`].
17
pub struct Json<T>(pub T);
18

            
19
/// Path Extractor / Response.
20
///
21
/// This is the customized [`axum::extract::Path`] version to respose error with [`ErrResp`].
22
pub struct Path<T>(pub T);
23

            
24
/// Query Extractor / Response.
25
///
26
/// This is the customized [`axum::extract::Query`] version to respose error with [`ErrResp`].
27
pub struct Query<T>(pub T);
28

            
29
impl<S, T> FromRequest<S> for Json<T>
30
where
31
    AxumJson<T>: FromRequest<S, Rejection = JsonRejection>,
32
    S: Send + Sync,
33
{
34
    type Rejection = ErrResp;
35

            
36
4
    async fn from_request(req: Request, state: &S) -> Result<Self, Self::Rejection> {
37
4
        match AxumJson::<T>::from_request(req, state).await {
38
2
            Err(e) => Err(ErrResp::ErrParam(Some(e.to_string()))),
39
2
            Ok(value) => Ok(Self(value.0)),
40
        }
41
4
    }
42
}
43

            
44
impl<T> IntoResponse for Json<T>
45
where
46
    T: Serialize,
47
{
48
2
    fn into_response(self) -> Response {
49
2
        // Use a small initial capacity of 128 bytes like serde_json::to_vec
50
2
        // https://docs.rs/serde_json/1.0.82/src/serde_json/ser.rs.html#2189
51
2
        let mut buf = BytesMut::with_capacity(128).writer();
52
2
        match serde_json::to_writer(&mut buf, &self.0) {
53
            Err(e) => ErrResp::ErrUnknown(Some(e.to_string())).into_response(),
54
2
            Ok(()) => (
55
2
                [(header::CONTENT_TYPE, ContentType::JSON)],
56
2
                buf.into_inner().freeze(),
57
2
            )
58
2
                .into_response(),
59
        }
60
2
    }
61
}
62

            
63
impl<T, S> FromRequestParts<S> for Path<T>
64
where
65
    T: DeserializeOwned + Send,
66
    S: Send + Sync,
67
{
68
    type Rejection = ErrResp;
69

            
70
4
    async fn from_request_parts(parts: &mut Parts, state: &S) -> Result<Self, Self::Rejection> {
71
4
        match AxumPath::from_request_parts(parts, state).await {
72
2
            Err(e) => Err(ErrResp::ErrParam(Some(e.to_string()))),
73
2
            Ok(value) => Ok(Self(value.0)),
74
        }
75
4
    }
76
}
77

            
78
impl<T, S> FromRequestParts<S> for Query<T>
79
where
80
    T: DeserializeOwned,
81
    S: Send + Sync,
82
{
83
    type Rejection = ErrResp;
84

            
85
4
    async fn from_request_parts(parts: &mut Parts, state: &S) -> Result<Self, Self::Rejection> {
86
4
        match AxumQuery::from_request_parts(parts, state).await {
87
2
            Err(e) => Err(ErrResp::ErrParam(Some(e.to_string()))),
88
2
            Ok(value) => Ok(Self(value.0)),
89
        }
90
4
    }
91
}
92

            
93
/// Parse Authorization header content. Returns `None` means no Authorization header.
94
6
pub fn parse_header_auth(req: &Request) -> Result<Option<String>, ErrResp> {
95
6
    let mut auth_all = req.headers().get_all(header::AUTHORIZATION).iter();
96
6
    let auth = match auth_all.next() {
97
2
        None => return Ok(None),
98
4
        Some(auth) => match auth.to_str() {
99
            Err(e) => return Err(ErrResp::ErrParam(Some(e.to_string()))),
100
4
            Ok(auth) => auth,
101
4
        },
102
4
    };
103
4
    if auth_all.next() != None {
104
2
        return Err(ErrResp::ErrParam(Some(
105
2
            "invalid multiple Authorization header".to_string(),
106
2
        )));
107
2
    }
108
2
    Ok(Some(auth.to_string()))
109
6
}