loco_openapi/
auth.rs

1use std::sync::OnceLock;
2
3use utoipa::{
4    openapi::security::{ApiKey, ApiKeyValue, HttpAuthScheme, HttpBuilder, SecurityScheme},
5    Modify,
6};
7
8// Import Loco types for conversion
9use loco_rs::{app::AppContext, config::JWTLocation as LocoJWTLocation};
10
11// Our own JWTLocation enum that doesn't depend on Loco
12#[derive(Default, Debug, Clone, PartialEq, Eq)]
13pub enum JWTLocation {
14    #[default]
15    Bearer,
16    Query(String),
17    Cookie(String),
18}
19
20// Implement From trait for conversion from Loco type to our type
21impl From<&LocoJWTLocation> for JWTLocation {
22    fn from(loco_location: &LocoJWTLocation) -> Self {
23        match loco_location {
24            LocoJWTLocation::Bearer => Self::Bearer,
25            LocoJWTLocation::Query { name } => Self::Query(name.clone()),
26            LocoJWTLocation::Cookie { name } => Self::Cookie(name.clone()),
27        }
28    }
29}
30
31// Support conversion from JWTLocationConfig (Single/Multiple) to JWTLocation
32impl From<&loco_rs::config::JWTLocationConfig> for JWTLocation {
33    fn from(cfg: &loco_rs::config::JWTLocationConfig) -> Self {
34        match cfg {
35            loco_rs::config::JWTLocationConfig::Single(loc) => Self::from(loc),
36            loco_rs::config::JWTLocationConfig::Multiple(locs) => {
37                locs.first().map_or(Self::Bearer, Self::from)
38            }
39        }
40    }
41}
42
43// Direct conversion from AppContext to JWTLocation for ease of use
44impl From<&AppContext> for JWTLocation {
45    fn from(ctx: &AppContext) -> Self {
46        ctx.config
47            .auth
48            .as_ref()
49            .and_then(|auth| auth.jwt.as_ref())
50            .and_then(|jwt| jwt.location.as_ref())
51            .map_or(Self::Bearer, Self::from)
52    }
53}
54
55static JWT_LOCATION: OnceLock<Option<JWTLocation>> = OnceLock::new();
56
57// Main API for working with JWT location - independent from Loco
58pub fn set_jwt_location(jwt_location: JWTLocation) -> &'static Option<JWTLocation> {
59    JWT_LOCATION.get_or_init(|| Some(jwt_location))
60}
61
62pub fn get_jwt_location() -> Option<&'static JWTLocation> {
63    JWT_LOCATION.get().unwrap_or(&None).as_ref()
64}
65
66// Security implementation using our JWTLocation
67pub struct SecurityAddon;
68
69impl Modify for SecurityAddon {
70    fn modify(&self, openapi: &mut utoipa::openapi::OpenApi) {
71        if let Some(jwt_location) = get_jwt_location() {
72            if let Some(components) = openapi.components.as_mut() {
73                components.add_security_schemes_from_iter([
74                    (
75                        "jwt_token",
76                        match jwt_location {
77                            JWTLocation::Bearer => SecurityScheme::Http(
78                                HttpBuilder::new()
79                                    .scheme(HttpAuthScheme::Bearer)
80                                    .bearer_format("JWT")
81                                    .build(),
82                            ),
83                            JWTLocation::Query(name) => {
84                                SecurityScheme::ApiKey(ApiKey::Query(ApiKeyValue::new(name)))
85                            }
86                            JWTLocation::Cookie(name) => {
87                                SecurityScheme::ApiKey(ApiKey::Cookie(ApiKeyValue::new(name)))
88                            }
89                        },
90                    ),
91                    (
92                        "api_key",
93                        SecurityScheme::ApiKey(ApiKey::Header(ApiKeyValue::new("apikey"))),
94                    ),
95                ]);
96            }
97        }
98    }
99}
100
101#[cfg(test)]
102mod tests {
103    use super::*;
104
105    #[test]
106    fn test_default_jwt_location() {
107        assert_eq!(JWTLocation::default(), JWTLocation::Bearer);
108    }
109
110    #[test]
111    fn test_set_get_jwt_location() {
112        set_jwt_location(JWTLocation::Bearer);
113        assert_eq!(get_jwt_location(), Some(&JWTLocation::Bearer));
114    }
115
116    #[test]
117    fn test_from_loco_jwt_location() {
118        let loco_bearer = LocoJWTLocation::Bearer;
119        assert_eq!(JWTLocation::from(&loco_bearer), JWTLocation::Bearer);
120
121        let loco_query = LocoJWTLocation::Query {
122            name: "token".to_string(),
123        };
124        assert_eq!(
125            JWTLocation::from(&loco_query),
126            JWTLocation::Query("token".to_string())
127        );
128
129        let loco_cookie = LocoJWTLocation::Cookie {
130            name: "auth".to_string(),
131        };
132        assert_eq!(
133            JWTLocation::from(&loco_cookie),
134            JWTLocation::Cookie("auth".to_string())
135        );
136    }
137}
loco_openapi/auth.rs

loco_openapi/
auth.rs
