use std::sync::Arc;
use base64::engine::general_purpose::URL_SAFE_NO_PAD as B64URL;
use base64::Engine as _;
use jsonwebtoken::{Algorithm, DecodingKey, EncodingKey, Header, Validation};
use rainier_support::{Error, Result};
use serde::de::DeserializeOwned;
use serde::Serialize;
use serde_json::{json, Value};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum JwtAlgorithm {
Rs256,
Es256,
}
impl JwtAlgorithm {
fn as_jsonwebtoken(self) -> Algorithm {
match self {
Self::Rs256 => Algorithm::RS256,
Self::Es256 => Algorithm::ES256,
}
}
pub fn name(self) -> &'static str {
match self {
Self::Rs256 => "RS256",
Self::Es256 => "ES256",
}
}
}
pub struct JwtKey {
kid: String,
algorithm: JwtAlgorithm,
encoding: EncodingKey,
decoding: DecodingKey,
jwk: Value,
}
impl JwtKey {
pub fn rs256_from_pem(kid: impl Into<String>, pem: &str) -> Result<Self> {
use rsa::pkcs1::DecodeRsaPrivateKey;
use rsa::pkcs8::DecodePrivateKey;
use rsa::traits::PublicKeyParts;
use rsa::RsaPrivateKey;
let private = RsaPrivateKey::from_pkcs8_pem(pem)
.or_else(|_| RsaPrivateKey::from_pkcs1_pem(pem))
.map_err(|e| Error::internal(format!("could not read the RSA private key: {e}")))?;
let public = private.to_public_key();
let jwk = json!({
"kty": "RSA",
"n": B64URL.encode(public.n().to_bytes_be()),
"e": B64URL.encode(public.e().to_bytes_be()),
});
Ok(Self {
kid: kid.into(),
algorithm: JwtAlgorithm::Rs256,
encoding: EncodingKey::from_rsa_pem(pem.as_bytes())
.map_err(|e| Error::internal(format!("could not load the RSA key: {e}")))?,
decoding: DecodingKey::from_rsa_pem(public_pem(&public)?.as_bytes())
.map_err(|e| Error::internal(format!("could not load the RSA key: {e}")))?,
jwk,
})
}
pub fn es256_from_pem(kid: impl Into<String>, pem: &str) -> Result<Self> {
use p256::elliptic_curve::sec1::ToEncodedPoint;
use p256::pkcs8::DecodePrivateKey;
use p256::SecretKey;
let secret = SecretKey::from_pkcs8_pem(pem)
.map_err(|e| Error::internal(format!("could not read the P-256 private key: {e}")))?;
let point = secret.public_key().to_encoded_point(false);
let jwk = json!({
"kty": "EC",
"crv": "P-256",
"x": B64URL.encode(point.x().ok_or_else(|| Error::internal("no x coordinate"))?),
"y": B64URL.encode(point.y().ok_or_else(|| Error::internal("no y coordinate"))?),
});
let public_pem = {
use p256::pkcs8::EncodePublicKey;
secret
.public_key()
.to_public_key_pem(Default::default())
.map_err(|e| Error::internal(format!("could not encode the public key: {e}")))?
};
Ok(Self {
kid: kid.into(),
algorithm: JwtAlgorithm::Es256,
encoding: EncodingKey::from_ec_pem(pem.as_bytes())
.map_err(|e| Error::internal(format!("could not load the P-256 key: {e}")))?,
decoding: DecodingKey::from_ec_pem(public_pem.as_bytes())
.map_err(|e| Error::internal(format!("could not load the P-256 key: {e}")))?,
jwk,
})
}
pub fn generate_rs256(kid: impl Into<String>, bits: usize) -> Result<Self> {
use rsa::pkcs8::{EncodePrivateKey, LineEnding};
use rsa::RsaPrivateKey;
let private = RsaPrivateKey::new(&mut rand::thread_rng(), bits)
.map_err(|e| Error::internal(format!("could not generate an RSA key: {e}")))?;
let pem = private
.to_pkcs8_pem(LineEnding::LF)
.map_err(|e| Error::internal(format!("could not encode the key: {e}")))?;
Self::rs256_from_pem(kid, &pem)
}
pub fn generate_es256(kid: impl Into<String>) -> Result<Self> {
use p256::pkcs8::{EncodePrivateKey, LineEnding};
let secret = p256::SecretKey::random(&mut rand::thread_rng());
let pem = secret
.to_pkcs8_pem(LineEnding::LF)
.map_err(|e| Error::internal(format!("could not encode the key: {e}")))?;
Self::es256_from_pem(kid, &pem)
}
pub fn kid(&self) -> &str {
&self.kid
}
pub fn algorithm(&self) -> JwtAlgorithm {
self.algorithm
}
pub fn to_jwk(&self) -> Value {
let mut jwk = self.jwk.clone();
jwk["kid"] = json!(self.kid);
jwk["alg"] = json!(self.algorithm.name());
jwk["use"] = json!("sig");
jwk
}
}
impl std::fmt::Debug for JwtKey {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("JwtKey")
.field("kid", &self.kid)
.field("algorithm", &self.algorithm)
.finish()
}
}
#[derive(Debug, Default)]
pub struct JwtKeyRing {
keys: Vec<Arc<JwtKey>>,
}
impl JwtKeyRing {
pub fn new(key: JwtKey) -> Self {
Self { keys: vec![Arc::new(key)] }
}
#[must_use = "this returns a ring with the key added"]
pub fn with_previous(mut self, key: JwtKey) -> Self {
self.keys.push(Arc::new(key));
self
}
pub fn current(&self) -> Option<&Arc<JwtKey>> {
self.keys.first()
}
pub fn find(&self, kid: &str) -> Option<&Arc<JwtKey>> {
self.keys.iter().find(|key| key.kid == kid)
}
pub fn ids(&self) -> Vec<&str> {
self.keys.iter().map(|key| key.kid.as_str()).collect()
}
pub fn jwks(&self) -> Value {
json!({ "keys": self.keys.iter().map(|key| key.to_jwk()).collect::<Vec<_>>() })
}
}
pub struct Jwt {
ring: JwtKeyRing,
issuer: Option<String>,
audience: Vec<String>,
leeway: u64,
}
impl Jwt {
pub fn new(ring: JwtKeyRing) -> Self {
Self { ring, issuer: None, audience: Vec::new(), leeway: 60 }
}
#[must_use = "this returns a configured signer rather than configuring in place"]
pub fn issued_by(mut self, issuer: impl Into<String>) -> Self {
self.issuer = Some(issuer.into());
self
}
#[must_use = "this returns a configured signer rather than configuring in place"]
pub fn for_audience(mut self, audience: impl Into<String>) -> Self {
self.audience.push(audience.into());
self
}
#[must_use = "this returns a configured signer rather than configuring in place"]
pub fn leeway(mut self, seconds: u64) -> Self {
self.leeway = seconds;
self
}
pub fn ring(&self) -> &JwtKeyRing {
&self.ring
}
pub fn jwks(&self) -> Value {
self.ring.jwks()
}
pub fn sign<C: Serialize>(&self, claims: &C) -> Result<String> {
let key = self.ring.current().ok_or_else(|| {
Error::internal("the JWT key ring is empty, so nothing can be signed")
})?;
let mut header = Header::new(key.algorithm.as_jsonwebtoken());
header.kid = Some(key.kid.clone());
jsonwebtoken::encode(&header, claims, &key.encoding)
.map_err(|e| Error::internal(format!("could not sign the token: {e}")))
}
pub fn verify<C: DeserializeOwned>(&self, token: &str) -> Result<C> {
let header = jsonwebtoken::decode_header(token)
.map_err(|_| Error::unauthenticated("This token is not readable."))?;
let kid = header.kid.ok_or_else(|| Error::unauthenticated("This token names no key."))?;
let key = self.ring.find(&kid).ok_or_else(|| {
Error::unauthenticated("This token names a key this service does not hold.")
})?;
let mut validation = Validation::new(key.algorithm.as_jsonwebtoken());
validation.leeway = self.leeway;
if let Some(issuer) = &self.issuer {
validation.set_issuer(&[issuer]);
}
if !self.audience.is_empty() {
validation.set_audience(&self.audience);
} else {
validation.validate_aud = false;
}
jsonwebtoken::decode::<C>(token, &key.decoding, &validation)
.map(|data| data.claims)
.map_err(|e| Error::unauthenticated(format!("This token is not valid: {e}")))
}
pub fn kid_of(token: &str) -> Option<String> {
jsonwebtoken::decode_header(token).ok()?.kid
}
}
impl std::fmt::Debug for Jwt {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Jwt")
.field("keys", &self.ring.ids())
.field("issuer", &self.issuer)
.field("audience", &self.audience)
.finish()
}
}
fn public_pem(public: &rsa::RsaPublicKey) -> Result<String> {
use rsa::pkcs8::{EncodePublicKey, LineEnding};
public
.to_public_key_pem(LineEnding::LF)
.map_err(|e| Error::internal(format!("could not encode the public key: {e}")))
}
#[cfg(test)]
mod tests {
use super::*;
use serde::Deserialize;
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
struct Claims {
sub: String,
exp: i64,
#[serde(skip_serializing_if = "Option::is_none")]
iss: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
aud: Option<String>,
}
fn in_an_hour() -> i64 {
std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).unwrap().as_secs() as i64
+ 3600
}
fn claims() -> Claims {
Claims { sub: "user-42".into(), exp: in_an_hour(), iss: None, aud: None }
}
fn rsa_key() -> &'static JwtKey {
use std::sync::OnceLock;
static KEY: OnceLock<JwtKey> = OnceLock::new();
KEY.get_or_init(|| JwtKey::generate_rs256("test-rsa", 2048).expect("generate"))
}
fn rsa_jwt() -> Jwt {
Jwt::new(JwtKeyRing::new(JwtKey::rs256_from_pem("test-rsa", &rsa_pem()).expect("from pem")))
}
fn rsa_pem() -> String {
use rsa::pkcs8::{EncodePrivateKey, LineEnding};
use std::sync::OnceLock;
static PEM: OnceLock<String> = OnceLock::new();
PEM.get_or_init(|| {
let private = rsa::RsaPrivateKey::new(&mut rand::thread_rng(), 2048).expect("generate");
private.to_pkcs8_pem(LineEnding::LF).expect("encode").to_string()
})
.clone()
}
#[test]
fn an_rs256_token_round_trips() {
let jwt = rsa_jwt();
let token = jwt.sign(&claims()).unwrap();
assert_eq!(jwt.verify::<Claims>(&token).unwrap().sub, "user-42");
}
#[test]
fn an_es256_token_round_trips() {
let jwt = Jwt::new(JwtKeyRing::new(JwtKey::generate_es256("test-ec").unwrap()));
let token = jwt.sign(&claims()).unwrap();
assert_eq!(jwt.verify::<Claims>(&token).unwrap().sub, "user-42");
}
#[test]
fn the_header_names_the_key() {
let token = rsa_jwt().sign(&claims()).unwrap();
assert_eq!(Jwt::kid_of(&token).as_deref(), Some("test-rsa"));
}
#[test]
fn a_tampered_token_does_not_verify() {
let jwt = rsa_jwt();
let token = jwt.sign(&claims()).unwrap();
let mut parts: Vec<&str> = token.split('.').collect();
let payload = parts[1].to_string();
let tampered_payload = format!("{}A", &payload[..payload.len() - 1]);
parts[1] = &tampered_payload;
assert!(jwt.verify::<Claims>(&parts.join(".")).is_err());
}
#[test]
fn an_expired_token_does_not_verify() {
let jwt = rsa_jwt();
let expired = Claims { exp: 1, ..claims() };
let token = jwt.sign(&expired).unwrap();
assert!(jwt.verify::<Claims>(&token).is_err());
}
#[test]
fn a_token_from_another_key_does_not_verify() {
let token = rsa_jwt().sign(&claims()).unwrap();
let stranger = Jwt::new(JwtKeyRing::new(
JwtKey::rs256_from_pem("test-rsa", &other_pem()).expect("from pem"),
));
assert!(stranger.verify::<Claims>(&token).is_err());
}
fn other_pem() -> String {
use rsa::pkcs8::{EncodePrivateKey, LineEnding};
use std::sync::OnceLock;
static PEM: OnceLock<String> = OnceLock::new();
PEM.get_or_init(|| {
let private = rsa::RsaPrivateKey::new(&mut rand::thread_rng(), 2048).expect("generate");
private.to_pkcs8_pem(LineEnding::LF).expect("encode").to_string()
})
.clone()
}
#[test]
fn an_unknown_kid_is_refused_without_trying_every_key() {
let jwt = rsa_jwt();
let token = jwt.sign(&claims()).unwrap();
let other = Jwt::new(JwtKeyRing::new(JwtKey::generate_es256("somebody-else").unwrap()));
let error = other.verify::<Claims>(&token).unwrap_err();
assert!(error.message().contains("does not hold"), "{}", error.message());
}
#[test]
fn a_rotation_keeps_verifying_the_previous_keys_tokens() {
let old = JwtKey::rs256_from_pem("old", &rsa_pem()).unwrap();
let signed_before = Jwt::new(JwtKeyRing::new(old)).sign(&claims()).unwrap();
let rotated = Jwt::new(
JwtKeyRing::new(JwtKey::generate_es256("new").unwrap())
.with_previous(JwtKey::rs256_from_pem("old", &rsa_pem()).unwrap()),
);
assert!(rotated.verify::<Claims>(&signed_before).is_ok());
let signed_after = rotated.sign(&claims()).unwrap();
assert_eq!(Jwt::kid_of(&signed_after).as_deref(), Some("new"));
}
#[test]
fn the_jwks_lists_every_key_that_can_still_verify() {
let ring = JwtKeyRing::new(JwtKey::generate_es256("new").unwrap())
.with_previous(JwtKey::rs256_from_pem("old", &rsa_pem()).unwrap());
let jwks = ring.jwks();
let keys = jwks["keys"].as_array().unwrap();
assert_eq!(keys.len(), 2);
assert_eq!(keys[0]["kid"], "new");
assert_eq!(keys[1]["kid"], "old");
}
#[test]
fn a_jwks_entry_is_the_shape_a_relying_party_reads() {
let jwks = JwtKeyRing::new(JwtKey::rs256_from_pem("k1", &rsa_pem()).unwrap()).jwks();
let key = &jwks["keys"][0];
assert_eq!(key["kty"], "RSA");
assert_eq!(key["alg"], "RS256");
assert_eq!(key["use"], "sig");
assert_eq!(key["kid"], "k1");
assert!(key["n"].as_str().unwrap().len() > 300);
assert_eq!(key["e"].as_str().unwrap(), "AQAB", "65537, as every RSA key uses");
assert!(!key["n"].as_str().unwrap().contains('='), "base64url is unpadded");
}
#[test]
fn an_ec_jwks_entry_carries_its_curve_and_point() {
let jwks = JwtKeyRing::new(JwtKey::generate_es256("k1").unwrap()).jwks();
let key = &jwks["keys"][0];
assert_eq!(key["kty"], "EC");
assert_eq!(key["crv"], "P-256");
assert_eq!(key["alg"], "ES256");
assert_eq!(key["x"].as_str().unwrap().len(), 43);
assert_eq!(key["y"].as_str().unwrap().len(), 43);
}
#[test]
fn the_issuer_is_checked_when_it_is_configured() {
let signer = rsa_jwt().issued_by("https://id.example.com");
let token = signer
.sign(&Claims { iss: Some("https://id.example.com".into()), ..claims() })
.unwrap();
assert!(signer.verify::<Claims>(&token).is_ok());
let impostor = signer
.sign(&Claims { iss: Some("https://evil.example.com".into()), ..claims() })
.unwrap();
assert!(signer.verify::<Claims>(&impostor).is_err());
}
#[test]
fn the_audience_is_checked_when_it_is_configured() {
let signer = rsa_jwt().for_audience("api");
let right = signer.sign(&Claims { aud: Some("api".into()), ..claims() }).unwrap();
assert!(signer.verify::<Claims>(&right).is_ok());
let wrong = signer.sign(&Claims { aud: Some("analytics".into()), ..claims() }).unwrap();
assert!(signer.verify::<Claims>(&wrong).is_err());
}
#[test]
fn a_token_with_an_audience_verifies_when_none_is_configured() {
let jwt = rsa_jwt();
let token = jwt.sign(&Claims { aud: Some("api".into()), ..claims() }).unwrap();
assert!(jwt.verify::<Claims>(&token).is_ok());
}
#[test]
fn nonsense_is_refused_without_panicking() {
let jwt = rsa_jwt();
for token in ["", "not.a.token", "a.b", &"x".repeat(5000)] {
assert!(jwt.verify::<Claims>(token).is_err(), "{token:?}");
}
}
#[test]
fn an_empty_ring_cannot_sign() {
let jwt = Jwt::new(JwtKeyRing::default());
let error = jwt.sign(&claims()).unwrap_err();
assert!(error.message().contains("empty"), "{}", error.message());
}
#[test]
fn a_generated_key_can_be_read_back_from_its_own_pem() {
let key = rsa_key();
assert_eq!(key.kid(), "test-rsa");
assert_eq!(key.algorithm(), JwtAlgorithm::Rs256);
assert_eq!(key.to_jwk()["kty"], "RSA");
}
}