Skip to main content

dns_update/
jwt.rs

1/*
2 * Copyright Stalwart Labs LLC See the COPYING
3 * file at the top-level directory of this distribution.
4 *
5 * Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
6 * https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
7 * <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
8 * option. This file may not be copied, modified, or distributed
9 * except according to those terms.
10 */
11
12//! Generic JWT utility for providers that need JWT authentication.
13
14use crate::http::try_build_client;
15use encodify::{
16    Error as CodecError,
17    base64::{LENIENT, URL_SAFE_NO_PAD},
18    pem::STANDARD as PEM,
19};
20use serde::{Deserialize, Serialize};
21use std::time::{SystemTime, UNIX_EPOCH};
22
23#[cfg(feature = "ring")]
24use ring::{
25    rand::SystemRandom,
26    signature::{RSA_PKCS1_SHA256, RSA_PKCS1_SHA512, RSA_PSS_SHA256, RsaKeyPair},
27};
28
29#[cfg(all(feature = "aws-lc-rs", not(feature = "ring")))]
30use aws_lc_rs::{
31    rand::SystemRandom,
32    signature::{RSA_PKCS1_SHA256, RSA_PKCS1_SHA512, RSA_PSS_SHA256, RsaKeyPair},
33};
34
35#[derive(Debug, Clone, Copy)]
36pub enum JwtSignAlgorithm {
37    Rs256,
38    Ps256,
39}
40
41/// Service account JSON fields needed for JWT creation.
42#[derive(Debug, Deserialize)]
43pub struct ServiceAccount {
44    pub client_email: String,
45    pub private_key: String,
46    pub token_uri: String,
47    // other fields are ignored
48}
49
50/// Claims for Google OAuth2 JWT.
51#[derive(Debug, Serialize)]
52struct JwtClaims {
53    iss: String,
54    scope: String,
55    aud: String,
56    exp: u64,
57    iat: u64,
58}
59
60/// Decode a PEM-encoded RSA private key (PKCS#8 `BEGIN PRIVATE KEY` or PKCS#1
61/// `BEGIN RSA PRIVATE KEY`), or an unarmored base64 PKCS#8 body, and return
62/// PKCS#8 DER bytes.
63pub fn rsa_private_key_pem_to_pkcs8_der(pem: &str) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
64    match PEM.decode(pem) {
65        Ok(block) => match block.label {
66            "ENCRYPTED PRIVATE KEY" => Err("encrypted PEM private keys are not supported".into()),
67            "RSA PRIVATE KEY" => Ok(wrap_pkcs1_in_pkcs8(&block.contents)),
68            _ => Ok(block.contents),
69        },
70        Err(CodecError::NotFound) => LENIENT
71            .decode(pem)
72            .map_err(|e| format!("Invalid base64 in private key: {}", e).into()),
73        Err(e) => Err(format!("Invalid PEM private key: {}", e).into()),
74    }
75}
76
77/// Wrap a PKCS#1 `RSAPrivateKey` DER blob in a PKCS#8 `PrivateKeyInfo` envelope.
78fn wrap_pkcs1_in_pkcs8(pkcs1_der: &[u8]) -> Vec<u8> {
79    const RSA_ENCRYPTION_ALG_ID: &[u8] = &[
80        0x30, 0x0D, 0x06, 0x09, 0x2A, 0x86, 0x48, 0x86, 0xF7, 0x0D, 0x01, 0x01, 0x01, 0x05, 0x00,
81    ];
82    let mut octet_string = Vec::with_capacity(pkcs1_der.len() + 4);
83    octet_string.push(0x04);
84    write_der_length(&mut octet_string, pkcs1_der.len());
85    octet_string.extend_from_slice(pkcs1_der);
86
87    let mut inner = Vec::with_capacity(3 + RSA_ENCRYPTION_ALG_ID.len() + octet_string.len());
88    inner.extend_from_slice(&[0x02, 0x01, 0x00]);
89    inner.extend_from_slice(RSA_ENCRYPTION_ALG_ID);
90    inner.extend_from_slice(&octet_string);
91
92    let mut out = Vec::with_capacity(inner.len() + 4);
93    out.push(0x30);
94    write_der_length(&mut out, inner.len());
95    out.extend_from_slice(&inner);
96    out
97}
98
99fn write_der_length(out: &mut Vec<u8>, len: usize) {
100    if len < 0x80 {
101        out.push(len as u8);
102    } else if len < 0x100 {
103        out.extend_from_slice(&[0x81, len as u8]);
104    } else if len < 0x10000 {
105        out.extend_from_slice(&[0x82, (len >> 8) as u8, len as u8]);
106    } else if len < 0x1000000 {
107        out.extend_from_slice(&[0x83, (len >> 16) as u8, (len >> 8) as u8, len as u8]);
108    } else {
109        out.extend_from_slice(&[
110            0x84,
111            (len >> 24) as u8,
112            (len >> 16) as u8,
113            (len >> 8) as u8,
114            len as u8,
115        ]);
116    }
117}
118
119/// Parse an RSA private key PEM (PKCS#1 or PKCS#8) into an `RsaKeyPair`.
120pub fn parse_rsa_pkcs8_pem(pem: &str) -> Result<RsaKeyPair, Box<dyn std::error::Error>> {
121    let der_bytes = rsa_private_key_pem_to_pkcs8_der(pem)?;
122    RsaKeyPair::from_pkcs8(&der_bytes).map_err(|e| format!("Invalid PKCS#8 RSA key: {}", e).into())
123}
124
125/// Sign `data` with RSA-SHA256 using a key pair, returning the raw signature bytes.
126pub fn rsa_sha256_sign(
127    key_pair: &RsaKeyPair,
128    data: &[u8],
129) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
130    let mut signature = vec![0u8; signature_len(key_pair)];
131    let rng = SystemRandom::new();
132    key_pair.sign(&RSA_PKCS1_SHA256, &rng, data, &mut signature)?;
133    Ok(signature)
134}
135
136/// Create a signed JWT using the service account private key.
137/// Returns the JWT as a compact string.
138pub fn create_jwt(sa: &ServiceAccount, scopes: &str) -> Result<String, Box<dyn std::error::Error>> {
139    let header = serde_json::json!({"alg": "RS256", "typ": "JWT"});
140    let mut jwt = URL_SAFE_NO_PAD.encode(serde_json::to_string(&header)?);
141
142    let now = SystemTime::now().duration_since(UNIX_EPOCH)?.as_secs();
143    let exp = now + 3600;
144
145    let claims = JwtClaims {
146        iss: sa.client_email.clone(),
147        scope: scopes.to_string(),
148        aud: sa.token_uri.clone(),
149        exp,
150        iat: now,
151    };
152    jwt.push('.');
153    URL_SAFE_NO_PAD.encode_append(serde_json::to_string(&claims)?, &mut jwt);
154
155    let key_pair = parse_rsa_pkcs8_pem(&sa.private_key)?;
156    let signature = rsa_sha256_sign(&key_pair, jwt.as_bytes())?;
157    jwt.push('.');
158    URL_SAFE_NO_PAD.encode_append(signature, &mut jwt);
159
160    Ok(jwt)
161}
162
163/// Exchange a JWT for an OAuth2 access token.
164pub async fn exchange_jwt_for_token(
165    token_uri: &str,
166    jwt: &str,
167) -> Result<String, Box<dyn std::error::Error>> {
168    let client = try_build_client(None)?;
169    let params = [
170        ("grant_type", "urn:ietf:params:oauth:grant-type:jwt-bearer"),
171        ("assertion", jwt),
172    ];
173    let body = serde_urlencoded::to_string(params).map_err(|e| e.to_string())?;
174    let resp: serde_json::Value = client
175        .post(token_uri)
176        .header("Content-Type", "application/x-www-form-urlencoded")
177        .body(body)
178        .send()
179        .await?
180        .json()
181        .await?;
182    if let Some(token) = resp.get("access_token") {
183        Ok(token.as_str().unwrap_or_default().to_string())
184    } else {
185        Err("Failed to obtain access token".into())
186    }
187}
188
189#[cfg(feature = "ring")]
190fn signature_len(key_pair: &RsaKeyPair) -> usize {
191    key_pair.public().modulus_len()
192}
193
194#[cfg(all(feature = "aws-lc-rs", not(feature = "ring")))]
195fn signature_len(key_pair: &RsaKeyPair) -> usize {
196    key_pair.public_modulus_len()
197}
198
199pub fn rsa_sha512_sign(
200    private_key_pem: &str,
201    data: &[u8],
202) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
203    let key_pair = parse_rsa_pkcs8_pem(private_key_pem)?;
204    let mut signature = vec![0u8; signature_len(&key_pair)];
205    let rng = SystemRandom::new();
206    key_pair.sign(&RSA_PKCS1_SHA512, &rng, data, &mut signature)?;
207    Ok(signature)
208}
209
210pub fn sign_jwt(
211    header: &serde_json::Value,
212    claims: &serde_json::Value,
213    private_key_pem: &str,
214    algorithm: JwtSignAlgorithm,
215) -> Result<String, Box<dyn std::error::Error>> {
216    let mut jwt = URL_SAFE_NO_PAD.encode(serde_json::to_string(header)?);
217    jwt.push('.');
218    URL_SAFE_NO_PAD.encode_append(serde_json::to_string(claims)?, &mut jwt);
219
220    let key_pair = parse_rsa_pkcs8_pem(private_key_pem)?;
221    let mut signature = vec![0u8; signature_len(&key_pair)];
222    let rng = SystemRandom::new();
223    match algorithm {
224        JwtSignAlgorithm::Rs256 => {
225            key_pair.sign(&RSA_PKCS1_SHA256, &rng, jwt.as_bytes(), &mut signature)?;
226        }
227        JwtSignAlgorithm::Ps256 => {
228            key_pair.sign(&RSA_PSS_SHA256, &rng, jwt.as_bytes(), &mut signature)?;
229        }
230    }
231    jwt.push('.');
232    URL_SAFE_NO_PAD.encode_append(signature, &mut jwt);
233    Ok(jwt)
234}