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(|_error| 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(|_error| 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(|_error| 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(|_error| invalid())?
82            .checked_mul(8)
83            .ok_or_else(invalid)?)
84        .to_be_bytes(),
85    );
86    Ok(mac)
87}
88
89///
90/// # Errors
91/// Returns an error when validation, storage, or an application callback fails.
92pub fn encode(
93    secret: &str,
94    salt: &str,
95    payload: &serde_json::Value,
96    max_age: f64,
97) -> AuthResult<String> {
98    let key = key(secret, salt)?;
99    let header = BASE64.encode(serde_json::to_vec(
100        &json!({"alg":"dir","enc":"A256CBC-HS512","kid":thumbprint(&key)}),
101    )?);
102    let now = Utc::now().timestamp();
103    let mut claims = serde_json::to_value(payload)?;
104    let claims = claims.as_object_mut().ok_or_else(invalid)?;
105    drop(claims.insert("iat".into(), json!(now)));
106    let expiry = serde_json::Number::from(now).as_f64().ok_or_else(invalid)? + max_age;
107    if !expiry.is_finite() {
108        return Err(AuthError::internal(
109            "Invalid encrypted-cookie expiration time",
110        ));
111    }
112    drop(claims.insert("exp".into(), json!(expiry)));
113    drop(claims.insert("jti".into(), json!(uuid::Uuid::new_v4().to_string())));
114    let mut ciphertext = crate::utils::json::to_vec(claims)?;
115    let padding = 16 - ciphertext.len() % 16;
116    ciphertext.resize(
117        ciphertext.len() + padding,
118        u8::try_from(padding).map_err(|error| AuthError::Internal(error.to_string()))?,
119    );
120    let mut iv = [0; 16];
121    rand::fill(&mut iv);
122    let cipher =
123        Aes256::new_from_slice(key.get(32..).ok_or_else(invalid)?).map_err(|_error| invalid())?;
124    let mut previous = iv;
125    for block in ciphertext.as_chunks_mut::<16>().0 {
126        for (byte, previous_2_3) in block.iter_mut().zip(previous) {
127            *byte ^= previous_2_3;
128        }
129        let block_array: &mut Array<u8, _> = block.into();
130        cipher.encrypt_block(block_array);
131        previous.copy_from_slice(block);
132    }
133    let tag = authentication(&key, &header, &iv, &ciphertext)?
134        .finalize()
135        .into_bytes();
136    Ok(format!(
137        "{header}..{}.{}.{}",
138        BASE64.encode(iv),
139        BASE64.encode(ciphertext),
140        BASE64.encode(tag.get(..32).ok_or_else(invalid)?)
141    ))
142}
143
144/// Authenticate the protected header, ciphertext and registered time claims.
145///
146/// # Errors
147/// Returns an error for malformed, unauthenticated or expired data.
148pub fn decode(secret: &str, salt: &str, token: &str) -> AuthResult<serde_json::Value> {
149    // Cookie values arrive URI-encoded; the source cookie parser decodes them
150    // before JOSE sees the protected header and its authenticated spelling.
151    let token = percent_encoding::percent_decode_str(token)
152        .decode_utf8()
153        .map_err(|_error| invalid())?;
154    decode_parsed(secret, salt, &token)
155}
156
157/// Authenticate a compact JWE after the owning cookie parser decoded its value.
158///
159/// # Errors
160/// Returns an error for malformed, unauthenticated or expired data.
161pub fn decode_parsed(secret: &str, salt: &str, token: &str) -> AuthResult<serde_json::Value> {
162    let parts: Vec<_> = token.split('.').collect();
163    let [header, encrypted_key, iv, ciphertext, tag] = parts.as_slice() else {
164        return Err(invalid());
165    };
166    if !encrypted_key.is_empty() {
167        return Err(invalid());
168    }
169    let key = key(secret, salt)?;
170    let header_bytes = decode_segment(header)?;
171    let header_data = parse_value(std::str::from_utf8(&header_bytes).map_err(|_error| invalid())?)
172        .map_err(|_error| invalid())?;
173    if header_data.get("alg").and_then(JsValue::as_str) != Some("dir")
174        || header_data.get("enc").and_then(JsValue::as_str) != Some("A256CBC-HS512")
175        || header_data.get("crit").is_some()
176        || header_data
177            .get("zip")
178            .is_some_and(|zip| zip.as_str() != Some("DEF"))
179    {
180        return Err(invalid());
181    }
182    if let Some(kid) = header_data.get("kid")
183        && kid.as_str() != Some(thumbprint(&key).as_str())
184    {
185        return Err(invalid());
186    }
187    let iv = decode_segment(iv)?;
188    let mut ciphertext = decode_segment(ciphertext)?;
189    let tag = decode_segment(tag)?;
190    if iv.len() != 16 || ciphertext.is_empty() || ciphertext.len() % 16 != 0 || tag.len() != 32 {
191        return Err(invalid());
192    }
193    authentication(&key, header, &iv, &ciphertext)?
194        .verify_truncated_left(&tag)
195        .map_err(|_error| invalid())?;
196    let cipher =
197        Aes256::new_from_slice(key.get(32..).ok_or_else(invalid)?).map_err(|_error| invalid())?;
198    let mut previous: [u8; 16] = iv.try_into().map_err(|_error| invalid())?;
199    for block in ciphertext.as_chunks_mut::<16>().0 {
200        let mut encrypted = [0; 16];
201        encrypted.copy_from_slice(block);
202        let block_array: &mut Array<u8, _> = block.into();
203        cipher.decrypt_block(block_array);
204        for (byte, previous_2) in block.iter_mut().zip(previous) {
205            *byte ^= previous_2;
206        }
207        previous = encrypted;
208    }
209    let padding = usize::from(*ciphertext.last().ok_or_else(invalid)?);
210    if padding == 0
211        || padding > 16
212        || ciphertext
213            .get(ciphertext.len().saturating_sub(padding)..)
214            .ok_or_else(invalid)?
215            .iter()
216            .any(|byte| usize::from(*byte) != padding)
217    {
218        return Err(invalid());
219    }
220    ciphertext.truncate(ciphertext.len() - padding);
221    if header_data.get("zip").is_some() {
222        // JOSE 6.2.12 authenticates and decrypts before bounded raw DEFLATE.
223        let mut inflated = Vec::with_capacity(250_001);
224        let mut inflater = flate2::Decompress::new(false);
225        let status = inflater
226            .decompress_vec(&ciphertext, &mut inflated, flate2::FlushDecompress::Finish)
227            .map_err(|_error| invalid())?;
228        if status != flate2::Status::StreamEnd || inflated.len() > 250_000 {
229            return Err(invalid());
230        }
231        ciphertext = inflated;
232    }
233    let claims = parse_value(std::str::from_utf8(&ciphertext).map_err(|_error| invalid())?)
234        .map_err(|_error| invalid())?;
235    let now = serde_json::Number::from(Utc::now().timestamp())
236        .as_f64()
237        .ok_or_else(invalid)?;
238    for name in ["iat", "exp", "nbf"] {
239        if let Some(value) = claims.get(name) {
240            let value = value.as_f64().ok_or_else(invalid)?;
241            if name == "exp" && value <= now - 15.0 || name == "nbf" && value > now + 15.0 {
242                return Err(invalid());
243            }
244        }
245    }
246    crate::utils::json::from_slice(&ciphertext).map_err(|_error| invalid())
247}