axum_security/jwt/
builder.rs

1use std::{borrow::Cow, error::Error, fmt::Display, marker::PhantomData, sync::Arc};
2
3use axum::http::{HeaderName, header::AUTHORIZATION};
4#[cfg(feature = "cookie")]
5use cookie_monster::CookieBuilder;
6use jsonwebtoken::{DecodingKey, EncodingKey, Header, Validation};
7
8#[cfg(feature = "cookie")]
9use crate::cookie_options::CookieOptionsBuilder;
10use crate::{
11    jwt::{ExtractFrom, JwtContext, JwtContextInner},
12    utils::get_env,
13};
14
15static PREFIX_BEARER: &str = "Bearer ";
16static PREFIX_NONE: &str = "";
17
18/// Builder for [`JwtContext`](super::JwtContext).
19///
20/// At minimum, provide signing keys via [`jwt_secret`](JwtContextBuilder::jwt_secret)
21/// or separate [`encoding_key`](JwtContextBuilder::encoding_key) /
22/// [`decoding_key`](JwtContextBuilder::decoding_key) calls.
23///
24/// By default, tokens are extracted from the `Authorization: Bearer <token>` header.
25/// Call [`extract_cookie`](JwtContextBuilder::extract_cookie) to extract from a cookie instead.
26pub struct JwtContextBuilder {
27    encoding_key: Option<EncodingKey>,
28    decoding_key: Option<DecodingKey>,
29    jwt_header: Header,
30    validation: Validation,
31    extract: ExtractFromBuilder,
32}
33
34impl Default for JwtContextBuilder {
35    fn default() -> Self {
36        Self::new()
37    }
38}
39
40impl JwtContextBuilder {
41    pub fn new() -> Self {
42        JwtContextBuilder {
43            encoding_key: None,
44            decoding_key: None,
45            jwt_header: Header::default(),
46            validation: Validation::default(),
47            extract: ExtractFromBuilder::header_with_prefix(AUTHORIZATION, PREFIX_BEARER),
48        }
49    }
50
51    pub fn encoding_key(mut self, encoding_key: EncodingKey) -> Self {
52        self.encoding_key = Some(encoding_key);
53        self
54    }
55
56    pub fn decoding_key(mut self, decoding_key: DecodingKey) -> Self {
57        self.decoding_key = Some(decoding_key);
58        self
59    }
60
61    pub fn jwt_secret(self, jwt_secret: impl AsRef<[u8]>) -> Self {
62        let jwt_secret = jwt_secret.as_ref();
63        self.encoding_key(EncodingKey::from_secret(jwt_secret))
64            .decoding_key(DecodingKey::from_secret(jwt_secret))
65    }
66
67    pub fn jwt_secret_env(self, name: &str) -> Self {
68        self.jwt_secret(get_env(name))
69    }
70
71    pub fn validation(mut self, validation: Validation) -> Self {
72        self.validation = validation;
73        self
74    }
75
76    pub fn jwt_header(mut self, header: Header) -> Self {
77        self.jwt_header = header;
78        self
79    }
80
81    pub fn extract_header_with_prefix(
82        mut self,
83        header: impl AsRef<[u8]>,
84        prefix: impl Into<Cow<'static, str>>,
85    ) -> Self {
86        self.extract = ExtractFromBuilder::header_with_prefix(
87            HeaderName::from_bytes(header.as_ref())
88                .expect("header value contains invalid characters"),
89            prefix.into(),
90        );
91        self
92    }
93
94    pub fn extract_header(mut self, header: impl AsRef<str>) -> Self {
95        self.extract = ExtractFromBuilder::header_with_prefix(
96            HeaderName::from_bytes(header.as_ref().as_bytes())
97                .expect("header value contains invalid characters"),
98            PREFIX_NONE,
99        );
100        self
101    }
102
103    #[cfg(feature = "cookie")]
104    pub fn extract_cookie(mut self, cookie_name: impl Into<Cow<'static, str>>) -> Self {
105        self.extract = ExtractFromBuilder::cookie(cookie_name.into());
106        self
107    }
108
109    #[cfg(feature = "cookie")]
110    pub fn use_dev_cookie(mut self, dev_mode: bool) -> Self {
111        self.extract = self.extract.with_cookie(|mut c| {
112            c.dev = dev_mode;
113            c
114        });
115        self
116    }
117
118    #[cfg(feature = "cookie")]
119    pub fn cookie(mut self, f: impl FnOnce(CookieBuilder) -> CookieBuilder) -> Self {
120        self.extract = self.extract.with_cookie(|mut c| {
121            c.cookie = (f)(c.cookie);
122            c
123        });
124        self
125    }
126
127    #[cfg(feature = "cookie")]
128    pub fn dev_cookie(mut self, f: impl FnOnce(CookieBuilder) -> CookieBuilder) -> Self {
129        self.extract = self.extract.with_cookie(|mut c| {
130            c.dev_cookie = (f)(c.dev_cookie);
131            c
132        });
133        self
134    }
135
136    pub fn try_build<T>(self) -> Result<JwtContext<T>, JwtBuilderError> {
137        let encoding_key = self
138            .encoding_key
139            .ok_or(JwtBuilderError::EncodingKeyMissing)?;
140
141        let decoding_key = self
142            .decoding_key
143            .ok_or(JwtBuilderError::DecodingKeyMissing)?;
144
145        let extract = self.extract.into_extract();
146
147        Ok(JwtContext(Arc::new(JwtContextInner {
148            encoding_key,
149            decoding_key,
150            jwt_header: self.jwt_header,
151            validation: self.validation,
152            extract,
153            data: PhantomData,
154        })))
155    }
156
157    pub fn build<T>(self) -> JwtContext<T> {
158        self.try_build().unwrap()
159    }
160}
161
162/// Error returned by [`JwtContextBuilder::try_build`].
163#[derive(Debug)]
164pub enum JwtBuilderError {
165    /// No encoding key was provided.
166    EncodingKeyMissing,
167    /// No decoding key was provided.
168    DecodingKeyMissing,
169}
170
171impl Display for JwtBuilderError {
172    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
173        match self {
174            JwtBuilderError::EncodingKeyMissing => f.write_str("Encoding key is missing"),
175            JwtBuilderError::DecodingKeyMissing => f.write_str("Decoding key is missing"),
176        }
177    }
178}
179
180impl Error for JwtBuilderError {}
181
182pub(crate) enum ExtractFromBuilder {
183    #[cfg(feature = "cookie")]
184    Cookie(Box<CookieOptionsBuilder>),
185    Header {
186        header: HeaderName,
187        prefix: Cow<'static, str>,
188    },
189}
190
191impl ExtractFromBuilder {
192    #[cfg(feature = "cookie")]
193    fn cookie(name: Cow<'static, str>) -> Self {
194        let mut builder = CookieOptionsBuilder::new();
195        builder.set_name(name);
196        ExtractFromBuilder::Cookie(builder.into())
197    }
198
199    fn into_extract(self) -> ExtractFrom {
200        match self {
201            #[cfg(feature = "cookie")]
202            ExtractFromBuilder::Cookie(builder) => ExtractFrom::Cookie(builder.build().into()),
203            ExtractFromBuilder::Header { header, prefix } => ExtractFrom::Header { header, prefix },
204        }
205    }
206
207    #[cfg(feature = "cookie")]
208    fn with_cookie(self, f: impl FnOnce(CookieOptionsBuilder) -> CookieOptionsBuilder) -> Self {
209        match self {
210            ExtractFromBuilder::Cookie(cookie) => ExtractFromBuilder::Cookie(f(*cookie).into()),
211            this => this,
212        }
213    }
214
215    fn header_with_prefix(header: HeaderName, prefix: impl Into<Cow<'static, str>>) -> Self {
216        ExtractFromBuilder::Header {
217            header,
218            prefix: prefix.into(),
219        }
220    }
221}
222
223#[cfg(test)]
224mod jwt_builder {
225    use crate::jwt::{DecodingKey, EncodingKey, JwtBuilderError, JwtContext};
226
227    #[test]
228    fn encoding_key_missing() {
229        let result = JwtContext::builder()
230            .decoding_key(DecodingKey::from_secret(b"test"))
231            .try_build::<()>();
232
233        assert!(matches!(result, Err(JwtBuilderError::EncodingKeyMissing)));
234    }
235
236    #[test]
237    fn decoding_key_missing() {
238        let result = JwtContext::builder()
239            .encoding_key(EncodingKey::from_secret(b"test"))
240            .try_build::<()>();
241
242        assert!(matches!(result, Err(JwtBuilderError::DecodingKeyMissing)));
243    }
244}
axum_security/jwt/builder.rs

axum_security/jwt/
builder.rs
