1use 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#[derive(Debug, Deserialize)]
43pub struct ServiceAccount {
44 pub client_email: String,
45 pub private_key: String,
46 pub token_uri: String,
47 }
49
50#[derive(Debug, Serialize)]
52struct JwtClaims {
53 iss: String,
54 scope: String,
55 aud: String,
56 exp: u64,
57 iat: u64,
58}
59
60pub 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
77fn 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
119pub 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
125pub 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
136pub 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
163pub 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}