spiral-auth 0.1.12

//! A client that authenticates with a private key whose public half it registered, such as a
//! workload: `client_credentials` with an RFC 7523 `private_key_jwt` client assertion.

use std::fmt::{Debug, Formatter, Result as FmtResult};
use std::sync::Arc;
use std::time::{Duration, SystemTime, UNIX_EPOCH};

use async_trait::async_trait;
use base64::Engine as _;
use base64::engine::general_purpose::URL_SAFE_NO_PAD;
use reqwest_middleware::ClientWithMiddleware as ReqwestClient;
use rustls::SignatureScheme;
use rustls::pki_types::{PrivateKeyDer, PrivatePkcs1KeyDer, PrivatePkcs8KeyDer, PrivateSec1KeyDer};
use rustls::sign::SigningKey as RustlsSigningKey;
use serde::Deserialize;
use serde_json::json;
use url::Url;
use uuid::Uuid;

use super::{ClientGrant, request_token, token_endpoint};
use crate::tls::crypto_provider;
use crate::{
    CLIENT_ASSERTION_TYPE_JWT_BEARER, CredentialError, CredentialResult,
    OAuth2ClientAuthentication, OAuth2Grant, OAuth2TokenRequest, OAuth2TokenResponse, Token,
};

/// Each assertion is used once, right after it is signed.
const ASSERTION_LIFETIME: Duration = Duration::from_secs(60);

/// A client ID and the private key, with its `kid`, that authenticates it.
pub struct PrivateKeyCredentials {
    client_id: String,
    kid: String,
    key: SigningKey,
}

impl PrivateKeyCredentials {
    /// Read a private key in JWK form (RFC 7517): RSA, P-256, P-384 or Ed25519. It must carry
    /// the `kid` it was registered with.
    pub fn from_jwk(client_id: String, jwk: &str) -> CredentialResult<Self> {
        let jwk: PrivateJwk = serde_json::from_str(jwk).map_err(|error| {
            CredentialError::invalid(format!("the private key is not a JWK: {error}"))
        })?;
        let Some(kid) = jwk.kid.clone() else {
            return Err(CredentialError::invalid(
                "the private key JWK must name the kid it was registered with",
            ));
        };
        Ok(Self {
            client_id,
            kid,
            key: SigningKey::from_jwk(&jwk)?,
        })
    }

    /// A JWT that asserts this client's identity to `token_endpoint`.
    fn assertion(&self, token_endpoint: &Url) -> CredentialResult<String> {
        let now = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .expect("the clock is after the Unix epoch");
        let header = json!({ "alg": self.key.algorithm(), "typ": "JWT", "kid": self.kid });
        // RFC 7523 §3: the client is both issuer and subject; the audience is the token
        // endpoint the assertion is presented to.
        let claims = json!({
            "iss": self.client_id.as_str(),
            "sub": self.client_id.as_str(),
            "aud": token_endpoint.as_str(),
            "jti": Uuid::new_v4().to_string(),
            "iat": now.as_secs(),
            "exp": (now + ASSERTION_LIFETIME).as_secs(),
        });
        let signing_input = format!(
            "{}.{}",
            URL_SAFE_NO_PAD.encode(header.to_string()),
            URL_SAFE_NO_PAD.encode(claims.to_string())
        );
        let signature = self.key.sign(signing_input.as_bytes())?;
        Ok(format!(
            "{signing_input}.{}",
            URL_SAFE_NO_PAD.encode(signature)
        ))
    }
}

impl Debug for PrivateKeyCredentials {
    fn fmt(&self, f: &mut Formatter<'_>) -> FmtResult {
        f.debug_struct("PrivateKeyCredentials")
            .field("client_id", &self.client_id)
            .field("kid", &self.kid)
            .finish_non_exhaustive()
    }
}

/// The members of a private JWK this client can sign with, base64url-encoded.
#[derive(Deserialize)]
struct PrivateJwk {
    kty: String,
    kid: Option<String>,
    crv: Option<String>,
    n: Option<String>,
    e: Option<String>,
    d: Option<String>,
    p: Option<String>,
    q: Option<String>,
    dp: Option<String>,
    dq: Option<String>,
    qi: Option<String>,
    x: Option<String>,
    y: Option<String>,
}

/// A JWS algorithm this client signs assertions with (RFC 7518 §3, RFC 8037 §3.1).
#[derive(Clone, Copy)]
enum Algorithm {
    Rs256,
    Es256,
    Es384,
    EdDsa,
}

impl Algorithm {
    /// The JWS `alg` header value.
    fn name(self) -> &'static str {
        match self {
            Self::Rs256 => "RS256",
            Self::Es256 => "ES256",
            Self::Es384 => "ES384",
            Self::EdDsa => "EdDSA",
        }
    }

    fn scheme(self) -> SignatureScheme {
        match self {
            Self::Rs256 => SignatureScheme::RSA_PKCS1_SHA256,
            Self::Es256 => SignatureScheme::ECDSA_NISTP256_SHA256,
            Self::Es384 => SignatureScheme::ECDSA_NISTP384_SHA384,
            Self::EdDsa => SignatureScheme::ED25519,
        }
    }
}

/// A private key loaded into the TLS crypto provider, which signs with it.
struct SigningKey {
    algorithm: Algorithm,
    key: Arc<dyn RustlsSigningKey>,
}

impl SigningKey {
    fn from_jwk(jwk: &PrivateJwk) -> CredentialResult<Self> {
        let (algorithm, der) = match (jwk.kty.as_str(), jwk.crv.as_deref()) {
            ("RSA", _) => (Algorithm::Rs256, rsa_private_key(jwk)?),
            ("EC", Some("P-256")) => (Algorithm::Es256, ec_private_key(jwk, P256)?),
            ("EC", Some("P-384")) => (Algorithm::Es384, ec_private_key(jwk, P384)?),
            ("OKP", Some("Ed25519")) => (Algorithm::EdDsa, ed25519_private_key(jwk)?),
            (kty, crv) => {
                return Err(CredentialError::invalid(format!(
                    "unsupported private key type `{kty}` with curve {crv:?}; use RSA, P-256, P-384 or Ed25519"
                )));
            }
        };
        let provider = crypto_provider()?;
        let key = provider
            .key_provider
            .load_private_key(der)
            .map_err(|error| {
                CredentialError::invalid(format!(
                    "the private key JWK is not a usable key: {error}"
                ))
            })?;
        Ok(Self { algorithm, key })
    }

    /// The JWS `alg` this key signs with.
    fn algorithm(&self) -> &'static str {
        self.algorithm.name()
    }

    fn sign(&self, message: &[u8]) -> CredentialResult<Vec<u8>> {
        let failed = || CredentialError::invalid("signing the client assertion failed");
        let signer = self
            .key
            .choose_scheme(&[self.algorithm.scheme()])
            .ok_or_else(failed)?;
        let signature = signer.sign(message).map_err(|_| failed())?;
        match self.algorithm {
            // RFC 7518 §3.4: a JWS carries the fixed-width R || S, not the DER TLS signature.
            Algorithm::Es256 => ecdsa_fixed(&signature, 32).ok_or_else(failed),
            Algorithm::Es384 => ecdsa_fixed(&signature, 48).ok_or_else(failed),
            Algorithm::Rs256 | Algorithm::EdDsa => Ok(signature),
        }
    }
}

/// DER object identifier of the P-256 curve, 1.2.840.10045.3.1.7.
const P256: &[u8] = &[0x06, 0x08, 0x2a, 0x86, 0x48, 0xce, 0x3d, 0x03, 0x01, 0x07];
/// DER object identifier of the P-384 curve, 1.3.132.0.34.
const P384: &[u8] = &[0x06, 0x05, 0x2b, 0x81, 0x04, 0x00, 0x22];
/// DER object identifier of Ed25519, 1.3.101.112.
const ED25519: &[u8] = &[0x06, 0x03, 0x2b, 0x65, 0x70];

/// RFC 8017 §A.1.2 `RSAPrivateKey`.
fn rsa_private_key(jwk: &PrivateJwk) -> CredentialResult<PrivateKeyDer<'static>> {
    let mut fields = der_unsigned(&[0]);
    for (name, value) in [
        ("n", &jwk.n),
        ("e", &jwk.e),
        ("d", &jwk.d),
        ("p", &jwk.p),
        ("q", &jwk.q),
        ("dp", &jwk.dp),
        ("dq", &jwk.dq),
        ("qi", &jwk.qi),
    ] {
        fields.extend(der_unsigned(&member(name, value)?));
    }
    Ok(PrivateKeyDer::Pkcs1(PrivatePkcs1KeyDer::from(der(
        0x30, &fields,
    ))))
}

/// RFC 5915 §3 `ECPrivateKey` on the curve `curve`, with its public point.
fn ec_private_key(jwk: &PrivateJwk, curve: &[u8]) -> CredentialResult<PrivateKeyDer<'static>> {
    // SEC 1 §2.3.3: an uncompressed point is 0x04 followed by both coordinates.
    let mut point = vec![0x00, 0x04];
    point.extend(member("x", &jwk.x)?);
    point.extend(member("y", &jwk.y)?);
    let mut fields = der_unsigned(&[1]);
    fields.extend(der(0x04, &member("d", &jwk.d)?));
    fields.extend(der(0xa0, curve));
    fields.extend(der(0xa1, &der(0x03, &point)));
    Ok(PrivateKeyDer::Sec1(PrivateSec1KeyDer::from(der(
        0x30, &fields,
    ))))
}

/// RFC 8410 §7 PKCS #8 Ed25519 private key.
fn ed25519_private_key(jwk: &PrivateJwk) -> CredentialResult<PrivateKeyDer<'static>> {
    let mut fields = der_unsigned(&[0]);
    fields.extend(der(0x30, ED25519));
    fields.extend(der(0x04, &der(0x04, &member("d", &jwk.d)?)));
    Ok(PrivateKeyDer::Pkcs8(PrivatePkcs8KeyDer::from(der(
        0x30, &fields,
    ))))
}

/// A DER element with tag `tag` and contents `content`.
fn der(tag: u8, content: &[u8]) -> Vec<u8> {
    let mut element = vec![tag];
    let length = content.len().to_be_bytes();
    match length.iter().position(|byte| *byte != 0) {
        Some(first) if content.len() >= 0x80 => {
            element.push(0x80 | (length.len() - first) as u8);
            element.extend(&length[first..]);
        }
        _ => element.push(content.len() as u8),
    }
    element.extend(content);
    element
}

/// A DER INTEGER holding the unsigned big-endian `value`.
fn der_unsigned(value: &[u8]) -> Vec<u8> {
    let significant = value
        .iter()
        .position(|byte| *byte != 0)
        .map_or(&[0][..], |first| &value[first..]);
    let mut content = Vec::with_capacity(significant.len() + 1);
    if significant[0] & 0x80 != 0 {
        content.push(0);
    }
    content.extend(significant);
    der(0x02, &content)
}

/// The fixed-width `R || S` of a DER `ECDSA-Sig-Value`, each `width` bytes.
fn ecdsa_fixed(signature: &[u8], width: usize) -> Option<Vec<u8>> {
    let (sequence, rest) = der_element(signature, 0x30)?;
    if !rest.is_empty() {
        return None;
    }
    let (r, rest) = der_element(sequence, 0x02)?;
    let (s, rest) = der_element(rest, 0x02)?;
    if !rest.is_empty() {
        return None;
    }
    let mut fixed = Vec::with_capacity(2 * width);
    for integer in [r, s] {
        let magnitude = &integer[integer.iter().position(|byte| *byte != 0)?..];
        if magnitude.len() > width {
            return None;
        }
        fixed.resize(fixed.len() + width - magnitude.len(), 0);
        fixed.extend(magnitude);
    }
    Some(fixed)
}

/// The contents of the DER element with tag `tag` at the start of `input`, and what follows it.
fn der_element(input: &[u8], tag: u8) -> Option<(&[u8], &[u8])> {
    let (&found, input) = input.split_first()?;
    if found != tag {
        return None;
    }
    let (&first, input) = input.split_first()?;
    let (length, input) = match first {
        length @ 0..=0x7f => (usize::from(length), input),
        0x81 => {
            let (&length, input) = input.split_first()?;
            (usize::from(length), input)
        }
        _ => return None,
    };
    (input.len() >= length).then(|| input.split_at(length))
}

/// The base64url-decoded value of the JWK member `name`.
fn member(name: &str, value: &Option<String>) -> CredentialResult<Vec<u8>> {
    let value = value
        .as_deref()
        .ok_or_else(|| CredentialError::invalid(format!("the private key JWK has no `{name}`")))?;
    URL_SAFE_NO_PAD.decode(value).map_err(|error| {
        CredentialError::invalid(format!("the private key JWK's `{name}`: {error}"))
    })
}

/// Obtains a client's tokens from the authorization server with its private key, read from
/// `SPIRAL_CLIENT_ID` and `SPIRAL_PRIVATE_KEY` unless given.
pub struct PrivateKeyJwtProvider {
    client: ReqwestClient,
    token_endpoint: Url,
    credentials: Option<PrivateKeyCredentials>,
    resource: Option<String>,
}

impl PrivateKeyJwtProvider {
    /// Authenticate to the authorization server at `auth_url`.
    pub fn new(client: impl Into<ReqwestClient>, auth_url: &Url) -> CredentialResult<Self> {
        Ok(Self {
            client: client.into(),
            token_endpoint: token_endpoint(auth_url)?,
            credentials: None,
            resource: None,
        })
    }

    /// Request tokens for the resource server `resource` (RFC 8707), one of the client's
    /// audiences.
    pub fn with_resource(mut self, resource: impl Into<String>) -> Self {
        self.resource = Some(resource.into());
        self
    }

    /// Authenticate as `credentials` rather than from the environment.
    pub fn with_credentials(mut self, credentials: PrivateKeyCredentials) -> Self {
        self.credentials = Some(credentials);
        self
    }

    fn assertion(&self) -> CredentialResult<Option<(String, String)>> {
        if let Some(credentials) = &self.credentials {
            let assertion = credentials.assertion(&self.token_endpoint)?;
            return Ok(Some((credentials.client_id.clone(), assertion)));
        }
        let Ok(jwk) = std::env::var("SPIRAL_PRIVATE_KEY") else {
            return Ok(None);
        };
        let Ok(client_id) = std::env::var("SPIRAL_CLIENT_ID") else {
            return Err(CredentialError::invalid(
                "SPIRAL_PRIVATE_KEY is set but SPIRAL_CLIENT_ID is not",
            ));
        };
        let credentials = PrivateKeyCredentials::from_jwk(client_id, &jwk)?;
        let assertion = credentials.assertion(&self.token_endpoint)?;
        Ok(Some((credentials.client_id, assertion)))
    }
}

impl Debug for PrivateKeyJwtProvider {
    fn fmt(&self, f: &mut Formatter<'_>) -> FmtResult {
        f.debug_struct("PrivateKeyJwtProvider")
            .field("credentials", &self.credentials)
            .finish_non_exhaustive()
    }
}

impl PrivateKeyJwtProvider {
    /// A usable token, or `None` when this provider has no credentials to offer.
    pub async fn token(&self) -> CredentialResult<Option<Token>> {
        let Some((client_id, assertion)) = self.assertion()? else {
            return Ok(None);
        };
        let request = OAuth2TokenRequest {
            client: OAuth2ClientAuthentication {
                client_id: Some(client_id),
                client_assertion_type: Some(CLIENT_ASSERTION_TYPE_JWT_BEARER.to_owned()),
                client_assertion: Some(assertion),
                ..Default::default()
            },
            grant: OAuth2Grant::ClientCredentials {
                resource: self.resource.clone(),
                authorization_details: None,
            },
        };
        let response = request_token(&self.client, &self.token_endpoint, &request).await?;
        Ok(Some(response.into_access_token()))
    }
}

impl_token_source!(PrivateKeyJwtProvider);

#[async_trait]
impl ClientGrant for PrivateKeyJwtProvider {
    async fn grant(
        &self,
        resource: &str,
        authorization_details: &serde_json::Value,
    ) -> CredentialResult<OAuth2TokenResponse> {
        let Some((client_id, assertion)) = self.assertion()? else {
            return Err(CredentialError::invalid(
                "a client grant needs the client's private key",
            ));
        };
        let request = OAuth2TokenRequest {
            client: OAuth2ClientAuthentication {
                client_id: Some(client_id),
                client_assertion_type: Some(CLIENT_ASSERTION_TYPE_JWT_BEARER.to_owned()),
                client_assertion: Some(assertion),
                ..Default::default()
            },
            grant: OAuth2Grant::ClientCredentials {
                resource: Some(resource.to_owned()),
                authorization_details: Some(authorization_details.to_string()),
            },
        };
        request_token(&self.client, &self.token_endpoint, &request).await
    }
}

#[cfg(test)]
mod tests {
    use aws_lc_rs::encoding::{AsBigEndian as _, Curve25519SeedBin, EcPrivateKeyBin};
    use aws_lc_rs::signature::{
        ECDSA_P256_SHA256_FIXED, ECDSA_P256_SHA256_FIXED_SIGNING, ED25519, EcdsaKeyPair,
        Ed25519KeyPair, KeyPair as _, RSA_PKCS1_2048_8192_SHA256, RsaPublicKeyComponents,
        UnparsedPublicKey,
    };
    use serde_json::{Map, Value};
    use wiremock::matchers::{method, path};
    use wiremock::{Mock, MockServer, ResponseTemplate};

    use super::*;
    use crate::http::test_client;

    /// A new P-256 private JWK named `key-1`, and its public key as an uncompressed point.
    fn p256_jwk() -> (String, Vec<u8>) {
        let key = EcdsaKeyPair::generate(&ECDSA_P256_SHA256_FIXED_SIGNING).unwrap();
        let public = key.public_key().as_ref().to_vec();
        let private: EcPrivateKeyBin<'_> = key.private_key().as_be_bytes().unwrap();
        let jwk = json!({
            "kty": "EC",
            "crv": "P-256",
            "kid": "key-1",
            "x": URL_SAFE_NO_PAD.encode(&public[1..33]),
            "y": URL_SAFE_NO_PAD.encode(&public[33..]),
            "d": URL_SAFE_NO_PAD.encode(private.as_ref()),
        });
        (jwk.to_string(), public)
    }

    /// Load `jwk` through the TLS crypto provider and sign `message` as a JWS.
    fn sign(jwk: &Value, message: &[u8]) -> (&'static str, Vec<u8>) {
        let key = SigningKey::from_jwk(&serde_json::from_value(jwk.clone()).unwrap()).unwrap();
        (key.algorithm(), key.sign(message).unwrap())
    }

    #[test]
    fn signs_rsa_and_ed25519_keys_verifiably() {
        let message = b"header.claims";

        let rsa: Value =
            serde_json::from_str(include_str!("testdata/rsa-private.jwk.json")).unwrap();
        let (algorithm, signature) = sign(&rsa, message);
        assert_eq!(algorithm, "RS256");
        let component = |name: &str| URL_SAFE_NO_PAD.decode(rsa[name].as_str().unwrap()).unwrap();
        RsaPublicKeyComponents {
            n: component("n"),
            e: component("e"),
        }
        .verify(&RSA_PKCS1_2048_8192_SHA256, message, &signature)
        .unwrap();

        let ed25519 = Ed25519KeyPair::generate().unwrap();
        let seed: Curve25519SeedBin<'_> = ed25519.seed().unwrap().as_be_bytes().unwrap();
        let public = ed25519.public_key().as_ref().to_vec();
        let jwk = json!({
            "kty": "OKP",
            "crv": "Ed25519",
            "x": URL_SAFE_NO_PAD.encode(&public),
            "d": URL_SAFE_NO_PAD.encode(seed.as_ref()),
        });
        let (algorithm, signature) = sign(&jwk, message);
        assert_eq!(algorithm, "EdDSA");
        UnparsedPublicKey::new(&ED25519, &public)
            .verify(message, &signature)
            .unwrap();
    }

    fn decode_segment(segment: &str) -> Map<String, Value> {
        serde_json::from_slice(&URL_SAFE_NO_PAD.decode(segment).unwrap()).unwrap()
    }

    #[tokio::test]
    async fn signs_a_single_use_assertion_for_the_token_endpoint() {
        let server = MockServer::start().await;
        Mock::given(method("POST"))
            .and(path("/oauth/token"))
            .respond_with(ResponseTemplate::new(200).set_body_json(json!({
                "access_token": "issued",
                "token_type": "Bearer",
                "expires_in": 900,
            })))
            .mount(&server)
            .await;
        let (jwk, public) = p256_jwk();
        let credentials = PrivateKeyCredentials::from_jwk("work_ingest".to_owned(), &jwk).unwrap();
        let provider =
            PrivateKeyJwtProvider::new(test_client(), &Url::parse(&server.uri()).unwrap())
                .unwrap()
                .with_credentials(credentials);

        let token = provider.token().await.unwrap().unwrap();

        assert_eq!(token.expose_secret(), "issued");
        let requests = server.received_requests().await.unwrap();
        let form: Map<String, Value> = serde_urlencoded::from_bytes(&requests[0].body).unwrap();
        assert_eq!(form["grant_type"], "client_credentials");
        assert_eq!(
            form["client_assertion_type"],
            CLIENT_ASSERTION_TYPE_JWT_BEARER
        );
        let assertion = form["client_assertion"].as_str().unwrap();
        let (signing_input, signature) = assertion.rsplit_once('.').unwrap();
        UnparsedPublicKey::new(&ECDSA_P256_SHA256_FIXED, &public)
            .verify(
                signing_input.as_bytes(),
                &URL_SAFE_NO_PAD.decode(signature).unwrap(),
            )
            .unwrap();
        let (header, claims) = signing_input.split_once('.').unwrap();
        let header = decode_segment(header);
        assert_eq!(header["alg"], "ES256");
        assert_eq!(header["kid"], "key-1");
        let claims = decode_segment(claims);
        assert_eq!(claims["iss"], "work_ingest");
        assert_eq!(claims["sub"], "work_ingest");
        assert_eq!(claims["aud"], format!("{}/oauth/token", server.uri()));
        assert!(claims["jti"].is_string());
    }

    #[test]
    fn rejects_a_key_without_a_kid() {
        let (jwk, _) = p256_jwk();
        let mut jwk: Map<String, Value> = serde_json::from_str(&jwk).unwrap();
        jwk.remove("kid");

        let error = PrivateKeyCredentials::from_jwk(
            "work_ingest".to_owned(),
            &Value::Object(jwk).to_string(),
        )
        .unwrap_err();

        assert!(error.to_string().contains("kid"), "{error}");
    }
}