agent-pay 0.1.0

//! Compact JWS over Ed25519 with JCS-canonical payload.

use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;

use base64::engine::general_purpose::URL_SAFE_NO_PAD;
use base64::Engine;
use serde_json::{json, Value};

use crate::error::Error;
use crate::jcs::canonical_json;
use crate::keys::{ed25519_sign, ed25519_verify};

pub type ResolveKey = Arc<
    dyn (Fn(String) -> Pin<Box<dyn Future<Output = Result<[u8; 32], Error>> + Send>>) + Send + Sync,
>;

fn b64url_encode(data: &[u8]) -> String {
    URL_SAFE_NO_PAD.encode(data)
}

fn b64url_decode(s: &str) -> Result<Vec<u8>, Error> {
    URL_SAFE_NO_PAD
        .decode(s)
        .map_err(|e| Error::Jws(format!("base64url: {e}")))
}

pub async fn sign_compact(
    payload: &Value,
    private_key: &[u8; 32],
    kid: &str,
) -> Result<String, Error> {
    let header = json!({ "alg": "EdDSA", "typ": "JWS", "kid": kid });
    let header_bytes = canonical_json(&header)?;
    let payload_bytes = canonical_json(payload)?;
    let header_b64 = b64url_encode(&header_bytes);
    let payload_b64 = b64url_encode(&payload_bytes);
    let signing_input = format!("{header_b64}.{payload_b64}");
    let sig = ed25519_sign(private_key, signing_input.as_bytes());
    Ok(format!(
        "{header_b64}.{payload_b64}.{}",
        b64url_encode(&sig)
    ))
}

pub async fn verify_compact(
    token: &str,
    resolve_key: &ResolveKey,
) -> Result<(Value, String), Error> {
    let parts: Vec<&str> = token.split('.').collect();
    if parts.len() != 3 {
        return Err(Error::Jws("compact JWS must have 3 parts".into()));
    }
    let header_b64 = parts[0];
    let payload_b64 = parts[1];
    let sig_b64 = parts[2];
    let header_bytes = b64url_decode(header_b64)?;
    let header: Value = serde_json::from_slice(&header_bytes)?;
    let alg = header.get("alg").and_then(|v| v.as_str()).unwrap_or("");
    if alg != "EdDSA" {
        return Err(Error::Jws(format!("unsupported JWS alg: {alg}")));
    }
    let kid = header
        .get("kid")
        .and_then(|v| v.as_str())
        .ok_or_else(|| Error::Jws("JWS header missing kid".into()))?
        .to_string();
    let public_key = (resolve_key)(kid.clone()).await?;
    let signing_input = format!("{header_b64}.{payload_b64}");
    let sig = b64url_decode(sig_b64)?;
    if !ed25519_verify(&public_key, signing_input.as_bytes(), &sig) {
        return Err(Error::Jws("JWS signature verification failed".into()));
    }
    let payload_bytes = b64url_decode(payload_b64)?;
    let payload: Value = serde_json::from_slice(&payload_bytes)?;
    Ok((payload, kid))
}