Skip to main content

alibi_core/utils/
jwe.rs

1//! Authenticated direct-key JWE shared by account and session cookies.
2use crate::{
3    AuthError, AuthResult,
4    utils::json::{JsValue, parse_value},
5};
6use aes_gcm::aes::{
7    Aes256,
8    cipher::{BlockCipherDecrypt, BlockCipherEncrypt, KeyInit},
9};
10use base64::{
11    Engine, alphabet,
12    engine::{
13        DecodePaddingMode,
14        general_purpose::{GeneralPurpose, GeneralPurposeConfig, URL_SAFE_NO_PAD as BASE64},
15    },
16};
17use chrono::Utc;
18use hkdf::Hkdf;
19use hmac::{Hmac, Mac};
20use hybrid_array::Array;
21use serde_json::json;
22use sha2::{Digest, Sha256, Sha512};
23
24const INFO: &[u8] = b"BetterAuth.js Generated Encryption Key";
25
26fn invalid() -> AuthError {
27    AuthError::bad_request("Invalid encrypted cookie")
28}
29
30fn decode_segment(value: &str) -> AuthResult<Vec<u8>> {
31    // Pinned JOSE uses base64url followed by atob: only these five ASCII
32    // whitespace characters are ignored; padding is optional but exact.
33    let compact: Vec<_> = value
34        .bytes()
35        .filter(|byte| !matches!(byte, b'\t' | b'\n' | b'\x0c' | b'\r' | b' '))
36        .collect();
37    let unpadded = compact
38        .iter()
39        .rposition(|byte| *byte != b'=')
40        .map_or(0, |index| index + 1);
41    let padding = compact.len() - unpadded;
42    if padding != 0 && (compact.len() % 4 != 0 || padding != (4 - unpadded % 4) % 4) {
43        return Err(invalid());
44    }
45    GeneralPurpose::new(
46        &alphabet::URL_SAFE,
47        GeneralPurposeConfig::new()
48            .with_decode_padding_mode(DecodePaddingMode::RequireNone)
49            .with_decode_allow_trailing_bits(true),
50    )
51    .decode(compact.get(..unpadded).ok_or_else(invalid)?)
52    .map_err(|_| invalid())
53}
54
55fn key(secret: &str, salt: &str) -> AuthResult<[u8; 64]> {
56    let mut key = [0; 64];
57    Hkdf::<Sha256>::new(Some(salt.as_bytes()), secret.as_bytes())
58        .expand(INFO, &mut key)
59        .map_err(|_| invalid())?;
60    Ok(key)
61}
62
63fn thumbprint(key: &[u8]) -> String {
64    let canonical = format!("{{\"k\":\"{}\",\"kty\":\"oct\"}}", BASE64.encode(key));
65    BASE64.encode(Sha256::digest(canonical.as_bytes()))
66}
67
68fn authentication(
69    key: &[u8],
70    header: &str,
71    iv: &[u8],
72    ciphertext: &[u8],
73) -> AuthResult<Hmac<Sha512>> {
74    let mut mac = Hmac::<Sha512>::new_from_slice(key.get(..32).ok_or_else(invalid)?)
75        .map_err(|_| invalid())?;
76    mac.update(header.as_bytes());
77    mac.update(iv);
78    mac.update(ciphertext);
79    mac.update(
80        &(u64::try_from(header.len())
81            .map_err(|_| invalid())?
82            .checked_mul(8)
83            .ok_or_else(invalid)?)
84        .to_be_bytes(),
85    );
86    Ok(mac)
87}
88
89pub fn encode(
90    secret: &str,
91    salt: &str,
92    payload: &serde_json::Value,
93    max_age: f64,
94) -> AuthResult<String> {
95    let key = key(secret, salt)?;
96    let header = BASE64.encode(serde_json::to_vec(
97        &json!({"alg":"dir","enc":"A256CBC-HS512","kid":thumbprint(&key)}),
98    )?);
99    let now = Utc::now().timestamp();
100    let mut claims = serde_json::to_value(payload)?;
101    let claims = claims.as_object_mut().ok_or_else(invalid)?;
102    _ = claims.insert("iat".into(), json!(now));
103    let expiry = serde_json::Number::from(now).as_f64().ok_or_else(invalid)? + max_age;
104    if !expiry.is_finite() {
105        return Err(AuthError::internal(
106            "Invalid encrypted-cookie expiration time",
107        ));
108    }
109    _ = claims.insert("exp".into(), json!(expiry));
110    _ = claims.insert("jti".into(), json!(uuid::Uuid::new_v4().to_string()));
111    let mut ciphertext = crate::utils::json::to_vec(claims)?;
112    let padding = 16 - ciphertext.len() % 16;
113    ciphertext.resize(
114        ciphertext.len() + padding,
115        u8::try_from(padding).map_err(|error| AuthError::Internal(error.to_string()))?,
116    );
117    let mut iv = [0; 16];
118    rand::fill(&mut iv);
119    let cipher =
120        Aes256::new_from_slice(key.get(32..).ok_or_else(invalid)?).map_err(|_| invalid())?;
121    let mut previous = iv;
122    for block in ciphertext.as_chunks_mut::<16>().0 {
123        for (byte, chained) in block.iter_mut().zip(previous) {
124            *byte ^= chained;
125        }
126        let block_array: &mut Array<u8, _> = block.into();
127        cipher.encrypt_block(block_array);
128        previous.copy_from_slice(block);
129    }
130    let tag = authentication(&key, &header, &iv, &ciphertext)?
131        .finalize()
132        .into_bytes();
133    Ok(format!(
134        "{header}..{}.{}.{}",
135        BASE64.encode(iv),
136        BASE64.encode(ciphertext),
137        BASE64.encode(tag.get(..32).ok_or_else(invalid)?)
138    ))
139}
140
141/// Authenticate the protected header, ciphertext and registered time claims.
142///
143/// # Errors
144/// Returns an error for malformed, unauthenticated or expired data.
145pub fn decode(secret: &str, salt: &str, token: &str) -> AuthResult<serde_json::Value> {
146    // Cookie values arrive URI-encoded; the source cookie parser decodes them
147    // before JOSE sees the protected header and its authenticated spelling.
148    let token = percent_encoding::percent_decode_str(token)
149        .decode_utf8()
150        .map_err(|_| invalid())?;
151    decode_parsed(secret, salt, &token)
152}
153
154/// Authenticate a compact JWE after the owning cookie parser decoded its value.
155///
156/// # Errors
157/// Returns an error for malformed, unauthenticated or expired data.
158pub fn decode_parsed(secret: &str, salt: &str, token: &str) -> AuthResult<serde_json::Value> {
159    let parts: Vec<_> = token.split('.').collect();
160    let [header, encrypted_key, iv, ciphertext, tag] = parts.as_slice() else {
161        return Err(invalid());
162    };
163    if !encrypted_key.is_empty() {
164        return Err(invalid());
165    }
166    let key = key(secret, salt)?;
167    let header_bytes = decode_segment(header)?;
168    let header_data = parse_value(std::str::from_utf8(&header_bytes).map_err(|_| invalid())?)
169        .map_err(|_| invalid())?;
170    if header_data.get("alg").and_then(JsValue::as_str) != Some("dir")
171        || header_data.get("enc").and_then(JsValue::as_str) != Some("A256CBC-HS512")
172        || header_data.get("crit").is_some()
173        || header_data
174            .get("zip")
175            .is_some_and(|zip| zip.as_str() != Some("DEF"))
176    {
177        return Err(invalid());
178    }
179    if let Some(kid) = header_data.get("kid")
180        && kid.as_str() != Some(thumbprint(&key).as_str())
181    {
182        return Err(invalid());
183    }
184    let iv = decode_segment(iv)?;
185    let mut ciphertext = decode_segment(ciphertext)?;
186    let tag = decode_segment(tag)?;
187    if iv.len() != 16 || ciphertext.is_empty() || ciphertext.len() % 16 != 0 || tag.len() != 32 {
188        return Err(invalid());
189    }
190    authentication(&key, header, &iv, &ciphertext)?
191        .verify_truncated_left(&tag)
192        .map_err(|_| invalid())?;
193    let cipher =
194        Aes256::new_from_slice(key.get(32..).ok_or_else(invalid)?).map_err(|_| invalid())?;
195    let mut previous: [u8; 16] = iv.try_into().map_err(|_| invalid())?;
196    for block in ciphertext.as_chunks_mut::<16>().0 {
197        let mut encrypted = [0; 16];
198        encrypted.copy_from_slice(block);
199        let block_array: &mut Array<u8, _> = block.into();
200        cipher.decrypt_block(block_array);
201        for (byte, previous) in block.iter_mut().zip(previous) {
202            *byte ^= previous;
203        }
204        previous = encrypted;
205    }
206    let padding = usize::from(*ciphertext.last().ok_or_else(invalid)?);
207    if padding == 0
208        || padding > 16
209        || ciphertext
210            .get(ciphertext.len().saturating_sub(padding)..)
211            .ok_or_else(invalid)?
212            .iter()
213            .any(|byte| usize::from(*byte) != padding)
214    {
215        return Err(invalid());
216    }
217    ciphertext.truncate(ciphertext.len() - padding);
218    if header_data.get("zip").is_some() {
219        // JOSE 6.2.12 authenticates and decrypts before bounded raw DEFLATE.
220        let mut inflated = Vec::with_capacity(250_001);
221        let mut inflater = flate2::Decompress::new(false);
222        let status = inflater
223            .decompress_vec(&ciphertext, &mut inflated, flate2::FlushDecompress::Finish)
224            .map_err(|_| invalid())?;
225        if status != flate2::Status::StreamEnd || inflated.len() > 250_000 {
226            return Err(invalid());
227        }
228        ciphertext = inflated;
229    }
230    let claims = parse_value(std::str::from_utf8(&ciphertext).map_err(|_| invalid())?)
231        .map_err(|_| invalid())?;
232    let now = serde_json::Number::from(Utc::now().timestamp())
233        .as_f64()
234        .ok_or_else(invalid)?;
235    for name in ["iat", "exp", "nbf"] {
236        if let Some(value) = claims.get(name) {
237            let value = value.as_f64().ok_or_else(invalid)?;
238            if name == "exp" && value <= now - 15.0 || name == "nbf" && value > now + 15.0 {
239                return Err(invalid());
240            }
241        }
242    }
243    crate::utils::json::from_slice(&ciphertext).map_err(|_| invalid())
244}
245
246// LCOV_EXCL_START
247#[cfg(test)]
248mod tests {
249    #![allow(
250        clippy::indexing_slicing,
251        reason = "test fixtures index known-length tokens"
252    )]
253    use super::*;
254    use std::io::Write as _;
255
256    const SECRET: &str = "jwe-test-secret";
257    const SALT: &str = "jwe-test-salt";
258
259    fn seal(header: &serde_json::Value, plaintext: &[u8]) -> String {
260        let key = key(SECRET, SALT).unwrap();
261        let header = BASE64.encode(serde_json::to_vec(header).unwrap());
262        let mut ciphertext = plaintext.to_vec();
263        let padding = 16 - ciphertext.len() % 16;
264        ciphertext.resize(ciphertext.len() + padding, u8::try_from(padding).unwrap());
265        let iv = [7u8; 16];
266        let cipher = Aes256::new_from_slice(&key[32..]).unwrap();
267        let mut previous = iv;
268        for block in ciphertext.as_chunks_mut::<16>().0 {
269            for (byte, previous) in block.iter_mut().zip(previous) {
270                *byte ^= previous;
271            }
272            let block_array: &mut Array<u8, _> = block.into();
273            cipher.encrypt_block(block_array);
274            previous.copy_from_slice(block);
275        }
276        let tag = authentication(&key, &header, &iv, &ciphertext)
277            .unwrap()
278            .finalize()
279            .into_bytes();
280        format!(
281            "{header}..{}.{}.{}",
282            BASE64.encode(iv),
283            BASE64.encode(ciphertext),
284            BASE64.encode(&tag[..32])
285        )
286    }
287
288    fn deflate(data: &[u8]) -> Vec<u8> {
289        let mut encoder =
290            flate2::write::DeflateEncoder::new(Vec::new(), flate2::Compression::default());
291        encoder.write_all(data).unwrap();
292        encoder.finish().unwrap()
293    }
294
295    fn header_with_zip() -> serde_json::Value {
296        json!({"alg":"dir","enc":"A256CBC-HS512","zip":"DEF"})
297    }
298
299    #[test]
300    fn deflated_payload_roundtrips() {
301        let claims = br#"{"sub":"user-1"}"#;
302        let token = seal(&header_with_zip(), &deflate(claims));
303        assert_eq!(
304            decode_parsed(SECRET, SALT, &token).unwrap(),
305            json!({"sub":"user-1"})
306        );
307    }
308
309    #[test]
310    fn rejects_corrupt_and_oversized_deflate_streams() {
311        let garbage = seal(&header_with_zip(), &[0xff, 0xff, 0xff, 0xff]);
312        assert!(decode_parsed(SECRET, SALT, &garbage).is_err());
313
314        let truncated = seal(&header_with_zip(), &deflate(br#"{"sub":"user-1"}"#)[..4]);
315        assert!(decode_parsed(SECRET, SALT, &truncated).is_err());
316
317        let big = format!(r#"{{"v":"{}"}}"#, "a".repeat(250_001));
318        let oversized = seal(&header_with_zip(), &deflate(big.as_bytes()));
319        assert!(decode_parsed(SECRET, SALT, &oversized).is_err());
320    }
321
322    #[test]
323    fn rejects_malformed_compact_tokens() {
324        let valid = encode(SECRET, SALT, &json!({"a":1}), 60.0).unwrap();
325        assert!(decode_parsed(SECRET, SALT, &valid).is_ok());
326        assert!(decode_parsed(SECRET, SALT, "a.b.c").is_err());
327        assert!(decode_parsed(SECRET, SALT, "a.b.c.d.e").is_err());
328
329        let mut parts: Vec<_> = valid.split('.').collect();
330        let short_iv = BASE64.encode([0u8; 8]);
331        parts[2] = &short_iv;
332        assert!(decode_parsed(SECRET, SALT, &parts.join(".")).is_err());
333
334        let mut parts: Vec<_> = valid.split('.').collect();
335        let bad_padding = format!("{}==", parts[3]);
336        parts[3] = &bad_padding;
337        assert!(decode_parsed(SECRET, SALT, &parts.join(".")).is_err());
338    }
339}
340// LCOV_EXCL_STOP