dynamic-waas-sdk 0.0.3

//! Decrypt delegated-wallet webhook payloads.
//!
//! Mirrors `@dynamic-labs-wallet/node`'s `decryptDelegatedWebhookData`.
//! The Dynamic backend emits two encrypted blobs in delegated-wallet
//! webhooks — the encrypted server key share and the encrypted per-wallet
//! API key. Both use a hybrid scheme:
//!
//!   * a fresh AES-256-GCM key is generated server-side,
//!   * the payload is encrypted with that key under AES-GCM,
//!   * the AES key itself is encrypted with the customer's RSA public key
//!     using RSA-OAEP-SHA256.
//!
//! The customer's RSA private key is the secret that gates both
//! decryptions; this function does the unwrap end-to-end.

use aes_gcm::aead::{Aead, KeyInit, Payload};
use aes_gcm::{Aes256Gcm, Key, Nonce};
use base64::engine::general_purpose;
use base64::Engine;
use rsa::pkcs8::DecodePrivateKey;
use rsa::{Oaep, RsaPrivateKey};
use serde::Deserialize;
use sha2::Sha256;

use dynamic_waas_sdk_core::{Error, Result, ServerKeyShare};

/// Algorithm identifiers Dynamic's delegation webhook currently emits.
/// Both name the same hybrid scheme (RSA-OAEP-SHA256 unwrapping an
/// AES-256-GCM payload key) — the production server has used both
/// labels over time. We hard-fail on any other value to defend against
/// algorithm-substitution / downgrade attacks. The current decrypt path
/// is fixed to RSA-OAEP-SHA256 + AES-256-GCM and ignores the field for
/// selection; validating it at the perimeter keeps future changes honest.
const SUPPORTED_ALGS: &[&str] = &["HYBRID-RSA-AES-256", "RSA-OAEP"];

/// Wire format of one encrypted webhook field.
///
/// Matches `EncryptedDelegatedPayload` in the Node SDK exactly:
///   * `alg`   — algorithm identifier, validated against [`SUPPORTED_ALGS`]
///   * `iv`    — AES-GCM 96-bit nonce, base64url
///   * `ct`    — AES-GCM ciphertext (without the auth tag), base64url
///   * `tag`   — AES-GCM 128-bit auth tag, base64url
///   * `ek`    — RSA-OAEP-SHA256 wrapped AES key, base64url
///   * `kid`   — optional key identifier for key rotation
#[derive(Debug, Clone, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct EncryptedDelegatedPayload {
    pub alg: String,
    pub iv: String,
    pub ct: String,
    pub tag: String,
    pub ek: String,
    #[serde(default)]
    pub kid: Option<String>,
}

/// Result of decrypting a delegated-wallet webhook.
///
/// Both fields hold long-lived secret material. The custom `Debug` impl
/// redacts them so accidental `tracing::debug!({:?})` or panic
/// `unwrap`s can't spill the wallet API key or share secret into logs.
#[derive(Clone)]
pub struct DecryptedWebhookData {
    /// The server's share of the wallet key. Pass this in as
    /// `external_server_key_shares` to `DelegatedWalletClient` operations.
    pub server_key_share: ServerKeyShare,
    /// The per-wallet delegated API key. Use as the `wallet_api_key`
    /// constructor argument to `DelegatedWalletClient`.
    pub wallet_api_key: String,
}

impl std::fmt::Debug for DecryptedWebhookData {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("DecryptedWebhookData")
            .field("server_key_share", &"<redacted>")
            .field("wallet_api_key", &"<redacted>")
            .finish()
    }
}

/// Decrypt the two encrypted blobs delivered in a delegated-wallet
/// webhook (the server key share + the per-wallet API key) using the
/// customer's RSA private key.
///
/// `rsa_private_key_pem` is a PKCS#8 PEM-encoded RSA private key — the
/// counterpart to the public key configured on the Dynamic environment
/// settings.
///
/// # Errors
/// Returns [`Error::InvalidArgument`] on:
///   * malformed PEM
///   * base64-url decode failure on any field
///   * RSA-OAEP unwrap failure (wrong key, tampered `ek`)
///   * AES-GCM auth-tag mismatch (tampered ciphertext)
///   * non-UTF8 plaintext for the wallet API key
///   * non-JSON / mismatched schema for the share payload
pub fn decrypt_delegated_webhook_data(
    rsa_private_key_pem: &str,
    encrypted_server_key_share: &EncryptedDelegatedPayload,
    encrypted_wallet_api_key: &EncryptedDelegatedPayload,
) -> Result<DecryptedWebhookData> {
    let private_key = RsaPrivateKey::from_pkcs8_pem(rsa_private_key_pem)
        .map_err(|e| Error::InvalidArgument(format!("invalid RSA private key PEM: {e}")))?;

    let share_plaintext = decrypt_one(&private_key, encrypted_server_key_share)?;
    let api_key_plaintext = decrypt_one(&private_key, encrypted_wallet_api_key)?;

    let share_json = std::str::from_utf8(&share_plaintext)
        .map_err(|e| Error::InvalidArgument(format!("decrypted share is not valid UTF-8: {e}")))?;
    let server_key_share: ServerKeyShare = parse_share_json(share_json)?;

    let wallet_api_key = String::from_utf8(api_key_plaintext).map_err(|e| {
        Error::InvalidArgument(format!("decrypted wallet API key is not valid UTF-8: {e}"))
    })?;

    Ok(DecryptedWebhookData {
        server_key_share,
        wallet_api_key,
    })
}

fn decrypt_one(
    private_key: &RsaPrivateKey,
    payload: &EncryptedDelegatedPayload,
) -> Result<Vec<u8>> {
    if !SUPPORTED_ALGS.contains(&payload.alg.as_str()) {
        return Err(Error::InvalidArgument(format!(
            "unsupported delegation payload algorithm `{}`; expected one of {:?}",
            payload.alg, SUPPORTED_ALGS
        )));
    }
    let ek = base64_url_decode(&payload.ek, "ek")?;
    let iv = base64_url_decode(&payload.iv, "iv")?;
    let ct = base64_url_decode(&payload.ct, "ct")?;
    let tag = base64_url_decode(&payload.tag, "tag")?;

    // RSA-OAEP-SHA256 unwrap the AES key.
    let aes_key_bytes = private_key
        .decrypt(Oaep::new::<Sha256>(), &ek)
        .map_err(|e| Error::InvalidArgument(format!("RSA-OAEP unwrap failed: {e}")))?;
    if aes_key_bytes.len() != 32 {
        return Err(Error::InvalidArgument(format!(
            "unwrapped AES key has wrong length: {} (expected 32)",
            aes_key_bytes.len()
        )));
    }
    if iv.len() != 12 {
        return Err(Error::InvalidArgument(format!(
            "AES-GCM nonce has wrong length: {} (expected 12)",
            iv.len()
        )));
    }

    // AES-256-GCM decrypt. The Node code stores `ct` and `tag` in separate
    // base64url fields; the `aes-gcm` crate's `decrypt` expects them
    // concatenated as `ct || tag`, so we glue them back together.
    let mut ct_with_tag = ct;
    ct_with_tag.extend_from_slice(&tag);

    let cipher = Aes256Gcm::new(Key::<Aes256Gcm>::from_slice(&aes_key_bytes));
    let nonce = Nonce::from_slice(&iv);
    cipher
        .decrypt(
            nonce,
            Payload {
                msg: &ct_with_tag,
                aad: &[],
            },
        )
        .map_err(|e| Error::InvalidArgument(format!("AES-GCM decrypt failed: {e}")))
}

fn base64_url_decode(input: &str, field: &str) -> Result<Vec<u8>> {
    general_purpose::URL_SAFE_NO_PAD
        .decode(input.trim_end_matches('='))
        .map_err(|e| Error::InvalidArgument(format!("invalid base64url in `{field}`: {e}")))
}

/// Parse the decrypted share JSON, accepting either `snake_case` or
/// `camelCase` field names (the wire format from Dynamic's webhook may
/// differ from the Rust SDK's canonical `snake_case`).
fn parse_share_json(s: &str) -> Result<ServerKeyShare> {
    #[derive(Deserialize)]
    struct Wire {
        #[serde(alias = "keyShareId")]
        key_share_id: String,
        #[serde(alias = "secretShare")]
        secret_share: String,
        #[serde(default, alias = "pubKey")]
        pub_key: Option<String>,
    }

    let wire: Wire = serde_json::from_str(s).map_err(|e| {
        Error::InvalidArgument(format!(
            "decrypted share is not a valid ServerKeyShare JSON: {e}"
        ))
    })?;
    let mut sks = ServerKeyShare::new(wire.key_share_id, wire.secret_share);
    if let Some(pk) = wire.pub_key {
        sks = sks.with_pub_key(pk);
    }
    Ok(sks)
}

#[cfg(test)]
mod tests {
    use super::*;
    use rand::TryRngCore;
    use rsa::pkcs8::{EncodePrivateKey, LineEnding};
    use rsa::RsaPublicKey;

    /// Encrypt a single plaintext using the same wire format the Dynamic
    /// webhook uses, given an RSA public key.
    fn encrypt_for_test(rsa_pub: &RsaPublicKey, plaintext: &[u8]) -> EncryptedDelegatedPayload {
        let mut rng = rand::rngs::OsRng;
        let mut aes_key = [0u8; 32];
        rng.try_fill_bytes(&mut aes_key).unwrap();
        let mut iv = [0u8; 12];
        rng.try_fill_bytes(&mut iv).unwrap();

        // Wrap the AES key with RSA-OAEP-SHA256.
        let ek = rsa_pub
            .encrypt(&mut rsa::rand_core::OsRng, Oaep::new::<Sha256>(), &aes_key)
            .unwrap();

        // AES-256-GCM encrypt.
        let cipher = Aes256Gcm::new(Key::<Aes256Gcm>::from_slice(&aes_key));
        let ciphertext_and_tag = cipher
            .encrypt(
                Nonce::from_slice(&iv),
                Payload {
                    msg: plaintext,
                    aad: &[],
                },
            )
            .unwrap();
        // The aes-gcm crate concatenates `ct||tag` (tag is 16 bytes at the end).
        let split_at = ciphertext_and_tag.len() - 16;
        let ct = &ciphertext_and_tag[..split_at];
        let tag = &ciphertext_and_tag[split_at..];

        EncryptedDelegatedPayload {
            alg: SUPPORTED_ALGS[0].into(),
            iv: general_purpose::URL_SAFE_NO_PAD.encode(iv),
            ct: general_purpose::URL_SAFE_NO_PAD.encode(ct),
            tag: general_purpose::URL_SAFE_NO_PAD.encode(tag),
            ek: general_purpose::URL_SAFE_NO_PAD.encode(&ek),
            kid: None,
        }
    }

    fn fresh_keypair() -> (String, RsaPublicKey) {
        let priv_key = RsaPrivateKey::new(&mut rsa::rand_core::OsRng, 2048).unwrap();
        let pub_key = RsaPublicKey::from(&priv_key);
        let pem = priv_key.to_pkcs8_pem(LineEnding::LF).unwrap().to_string();
        (pem, pub_key)
    }

    #[test]
    fn round_trip_snake_case() {
        let (pem, pub_key) = fresh_keypair();
        let share_json = r#"{"key_share_id":"ks-1","secret_share":"deadbeef","pub_key":"0x04abc"}"#;
        let wallet_api_key = "test-fixture-wallet-key-snake";

        let enc_share = encrypt_for_test(&pub_key, share_json.as_bytes());
        let enc_wapi = encrypt_for_test(&pub_key, wallet_api_key.as_bytes());

        let out = decrypt_delegated_webhook_data(&pem, &enc_share, &enc_wapi).unwrap();
        assert_eq!(out.server_key_share.key_share_id, "ks-1");
        assert_eq!(out.server_key_share.secret_share, "deadbeef");
        assert_eq!(out.server_key_share.pub_key.as_deref(), Some("0x04abc"));
        assert_eq!(out.wallet_api_key, wallet_api_key);
    }

    #[test]
    fn round_trip_camel_case() {
        // Webhook server may emit camelCase; we accept both via serde alias.
        let (pem, pub_key) = fresh_keypair();
        let share_json = r#"{"keyShareId":"ks-2","secretShare":"feedface"}"#;
        let wallet_api_key = "test-fixture-wallet-key-camel";

        let enc_share = encrypt_for_test(&pub_key, share_json.as_bytes());
        let enc_wapi = encrypt_for_test(&pub_key, wallet_api_key.as_bytes());

        let out = decrypt_delegated_webhook_data(&pem, &enc_share, &enc_wapi).unwrap();
        assert_eq!(out.server_key_share.key_share_id, "ks-2");
        assert_eq!(out.server_key_share.secret_share, "feedface");
        assert!(out.server_key_share.pub_key.is_none());
        assert_eq!(out.wallet_api_key, wallet_api_key);
    }

    #[test]
    fn wrong_private_key_fails_cleanly() {
        let (_pem_a, pub_key_a) = fresh_keypair();
        let (pem_b, _pub_key_b) = fresh_keypair();
        let share_json = r#"{"key_share_id":"ks","secret_share":"x"}"#;
        let enc_share = encrypt_for_test(&pub_key_a, share_json.as_bytes());
        let enc_wapi = encrypt_for_test(&pub_key_a, b"key");

        let err = decrypt_delegated_webhook_data(&pem_b, &enc_share, &enc_wapi).unwrap_err();
        match err {
            Error::InvalidArgument(msg) => assert!(msg.contains("RSA-OAEP")),
            other => panic!("expected InvalidArgument, got {other:?}"),
        }
    }

    #[test]
    fn tampered_ciphertext_fails_cleanly() {
        let (pem, pub_key) = fresh_keypair();
        let share_json = r#"{"key_share_id":"ks","secret_share":"x"}"#;
        let mut enc_share = encrypt_for_test(&pub_key, share_json.as_bytes());
        let enc_wapi = encrypt_for_test(&pub_key, b"key");

        // Flip a bit in the ciphertext.
        let mut ct_bytes = general_purpose::URL_SAFE_NO_PAD
            .decode(&enc_share.ct)
            .unwrap();
        ct_bytes[0] ^= 0x01;
        enc_share.ct = general_purpose::URL_SAFE_NO_PAD.encode(&ct_bytes);

        let err = decrypt_delegated_webhook_data(&pem, &enc_share, &enc_wapi).unwrap_err();
        match err {
            Error::InvalidArgument(msg) => assert!(msg.contains("AES-GCM")),
            other => panic!("expected InvalidArgument, got {other:?}"),
        }
    }

    #[test]
    fn unexpected_alg_fails_before_any_crypto() {
        // Defends against algorithm-substitution / downgrade attacks: a
        // tampered webhook that swaps `alg` to something we don't speak
        // must be rejected at the perimeter, before any RSA or AES work.
        let (pem, pub_key) = fresh_keypair();
        let share_json = r#"{"key_share_id":"ks","secret_share":"x"}"#;
        let mut enc_share = encrypt_for_test(&pub_key, share_json.as_bytes());
        let enc_wapi = encrypt_for_test(&pub_key, b"key");
        enc_share.alg = "RSA-OAEP-256+A128GCM".into(); // weaker / unsupported

        let err = decrypt_delegated_webhook_data(&pem, &enc_share, &enc_wapi).unwrap_err();
        match err {
            Error::InvalidArgument(msg) => assert!(
                msg.contains("unsupported delegation payload algorithm"),
                "message should mention the rejected algorithm: {msg}"
            ),
            other => panic!("expected InvalidArgument, got {other:?}"),
        }
    }

    #[test]
    fn decrypted_data_debug_redacts_both_fields() {
        // Defense in depth: even though we never log this struct
        // ourselves, callers using `tracing::debug!({:?}, decrypted)`
        // or an `unwrap()` panic message must not leak the share or the
        // wallet API key.
        let data = DecryptedWebhookData {
            server_key_share: ServerKeyShare::new("ks", "super-secret-share-bytes"),
            wallet_api_key: "super-secret-wallet-api-key".to_string(),
        };
        let dbg = format!("{data:?}");
        assert!(
            !dbg.contains("super-secret-share-bytes"),
            "share secret leaked via Debug: {dbg}"
        );
        assert!(
            !dbg.contains("super-secret-wallet-api-key"),
            "wallet api key leaked via Debug: {dbg}"
        );
        assert!(dbg.contains("redacted"));
    }

    #[test]
    fn malformed_pem_fails_cleanly() {
        let dummy = EncryptedDelegatedPayload {
            alg: "x".into(),
            iv: String::new(),
            ct: String::new(),
            tag: String::new(),
            ek: String::new(),
            kid: None,
        };
        let err = decrypt_delegated_webhook_data("not a pem", &dummy, &dummy).unwrap_err();
        match err {
            Error::InvalidArgument(msg) => assert!(msg.contains("RSA private key PEM")),
            other => panic!("expected InvalidArgument, got {other:?}"),
        }
    }
}