middleware 0.1.0

my custom middleware
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120

pub mod types;

use axum::{
    body::Body,
    extract::{FromRequestParts, Request, State},
    http::{self, StatusCode},
    middleware::Next,
    response::Response,
};
use axum_extra::extract::CookieJar;
use jsonwebtoken::{Algorithm, DecodingKey, Validation, decode, decode_header};
use serde::Deserialize;
use service_error::ServiceError;

use crate::types::{Jwks, JwksTrait};

#[derive(Debug, Deserialize, Clone)]
pub struct RealmAccess {
    pub roles: Vec<String>,
}

#[derive(Debug, Deserialize, Clone)]
pub struct ResourceRoles {
    pub roles: Vec<String>,
}

#[derive(Debug, Deserialize, Clone)]
pub struct Claims {
    pub sub: Option<String>,
    pub exp: Option<u64>,
    pub iat: Option<u64>,
    pub iss: Option<String>,
    pub aud: Option<String>,
    pub typ: Option<String>,
    pub azp: Option<String>,
    pub name: Option<String>,
    pub preferred_username: Option<String>,
    pub given_name: Option<String>,
    pub family_name: Option<String>,
    pub email: Option<String>,
    pub email_verified: Option<bool>,
    pub realm_access: Option<RealmAccess>,
    pub resource_access: Option<std::collections::HashMap<String, ResourceRoles>>,
}

#[derive(Debug)]
pub struct ClaimsExtension(pub Claims);

impl<S> FromRequestParts<S> for ClaimsExtension
where
    S: Send + Sync,
{
    type Rejection = ServiceError;

    async fn from_request_parts(
        parts: &mut http::request::Parts,
        _state: &S,
    ) -> Result<Self, Self::Rejection> {
        // Access to headers from req
        let ext = &parts.extensions;

        let claims = ext.get::<Claims>();

        match claims {
            Some(claims) => Ok(ClaimsExtension(claims.clone())),
            None => Err(ServiceError::Internal(
                "Claims not found in request extensions".to_string(),
            )),
        }
    }
}

pub async fn validate_auth_cookie<T>(
    State(state): State<std::sync::Arc<T>>,
    cookies: CookieJar,
    mut req: Request<Body>,
    next: Next,
) -> Result<Response, StatusCode>
where
    T: JwksTrait,
{
    let token = get_auth_token(&cookies)?;

    match validate_token(token, &state.jwks()) {
        Ok(claims) => {
            req.extensions_mut().insert(claims);
            Ok(next.run(req).await)
        }
        Err(_) => Err(StatusCode::UNAUTHORIZED),
    }
}

fn get_auth_token(cookies: &CookieJar) -> Result<&str, StatusCode> {
    match cookies.get("access_token") {
        Some(token) => Ok(token.value()),
        None => Err(StatusCode::UNAUTHORIZED),
    }
}

fn validate_token(token: &str, jwks: &Jwks) -> Result<Claims, StatusCode> {
    let header = decode_header(token).map_err(|_| StatusCode::UNAUTHORIZED)?;
    let kid = header.kid.ok_or(StatusCode::UNAUTHORIZED)?;

    // Find matching key
    let jwk = jwks
        .keys
        .iter()
        .find(|k| k.kid == kid)
        .ok_or(StatusCode::UNAUTHORIZED)?;

    let decoding_key =
        DecodingKey::from_rsa_components(&jwk.n, &jwk.e).map_err(|_| StatusCode::UNAUTHORIZED)?;

    let validation = Validation::new(Algorithm::RS256);

    let token_data = decode::<Claims>(token, &decoding_key, &validation)
        .map_err(|_| StatusCode::UNAUTHORIZED)?;

    Ok(token_data.claims)
}