gunnar-crypt 1.5.1

//! age's hybrid post-quantum recipient, `mlkem768x25519` (c2sp.org/age):
//! an `age1pq1…` public key and its `AGE-SECRET-KEY-PQ-1…` identity.
//!
//! **Why it is here and not in age.** No age release carries it yet: age
//! 0.12.1 (2026-07) has only the encryption-only `tag`/`tagpq` types, and it
//! pulls ml-kem 0.2, which cannot share a lock with russh's ml-kem 0.3. rage's
//! native `pq` type is an open pull request (str4d/rage#632). This module is
//! the same stanza, on age 0.11.5's public [`age::Recipient`] and
//! [`age::Identity`] traits, and it is meant to be replaced by age's own type
//! the day a release carries one. Nothing in it is new cryptography:
//!
//! * the KEM is X-Wing, a.k.a. MLKEM768-X25519 (HPKE KEM id 0x647a), from the
//!   released `hpke` 0.14.1 and `x-wing` 0.1.1 (RustCrypto, ml-kem 0.3);
//! * the wrap is HPKE (RFC 9180) base mode with HKDF-SHA256 and
//!   ChaCha20Poly1305, info `age-encryption.org/mlkem768x25519`, empty aad —
//!   exactly the spec's parameters;
//! * an identity is the 32-byte X-Wing decapsulation seed, bech32 with the
//!   `AGE-SECRET-KEY-PQ-` prefix; a recipient is the 1216-byte encapsulation
//!   key, bech32 with `age1pq` (no length limit, as age encodes it).
//!
//! The stanza is `-> mlkem768x25519 <base64 enc, 1120 bytes>` with a 32-byte
//! body (the 16-byte file key and its tag). Every C2SP CCTV `hybrid*` vector
//! is run against it (`tests/pq.rs`), and Go age v1.3 opens what it wraps and
//! the other way round (the interop row in the design document).
//!
//! **One rule the vectors pin**: an X25519 half that is a low-order point
//! makes the X25519 shared secret all-zero, and the spec refuses it as a
//! HEADER failure. hpke 0.14.1's X-Wing accepts such a share, so the
//! decapsulation is asked first of x-wing's `DecapsulationKeyRejectNonContrib`
//! (its documented strict variant) and a refusal there is a header failure.
//!
//! The recipient labels its stanza `postquantum`, as age's own will: age
//! itself then refuses to put it in one file with a classical recipient.

use std::collections::HashSet;
use std::fmt;

use age::secrecy::{ExposeSecret as _, SecretString};
use age_core::format::{FileKey, Stanza, FILE_KEY_BYTES};
use base64::engine::general_purpose::STANDARD_NO_PAD;
use base64::Engine as _;
use bech32::{FromBase32 as _, ToBase32 as _, Variant};
use hpke::{Deserializable as _, Kem as _, Serializable as _};
use x_wing::TryDecapsulate as _;
use zeroize::{Zeroize as _, Zeroizing};

pub use crate::members::PQ_STANZA as STANZA_TYPE;

/// The HPKE `info` of the stanza (c2sp.org/age, `mlkem768x25519`).
const INFO: &[u8] = b"age-encryption.org/mlkem768x25519";
/// The recipient's bech32 HRP.
const RECIPIENT_HRP: &str = "age1pq";
/// The identity's bech32 HRP (written upper-case).
const IDENTITY_HRP: &str = "age-secret-key-pq-";
/// The encapsulated key: an ML-KEM-768 ciphertext and an X25519 share.
const ENC_BYTES: usize = 1120;
/// The label age gives a post-quantum recipient's stanzas.
pub const LABEL: &str = "postquantum";

type Kem = hpke::kem::XWing;
type Aead = hpke::aead::ChaCha20Poly1305;
type Kdf = hpke::kdf::HkdfSha256;

/// The OS random source, for HPKE's encapsulation, through the `getrandom`
/// this crate already uses on both targets.
struct OsRng;

impl hpke::rand_core::TryRng for OsRng {
    type Error = core::convert::Infallible;
    fn try_next_u32(&mut self) -> Result<u32, Self::Error> {
        Ok(u32::from_le_bytes(crate::random_bytes::<4>()))
    }
    fn try_next_u64(&mut self) -> Result<u64, Self::Error> {
        Ok(u64::from_le_bytes(crate::random_bytes::<8>()))
    }
    fn try_fill_bytes(&mut self, dst: &mut [u8]) -> Result<(), Self::Error> {
        getrandom::getrandom(dst).expect("the OS random source is available");
        Ok(())
    }
}
impl hpke::rand_core::TryCryptoRng for OsRng {}

/// An `age1pq1…` public key.
#[derive(Clone, PartialEq, Eq)]
pub struct Recipient(<Kem as hpke::Kem>::PublicKey);

impl std::str::FromStr for Recipient {
    type Err = &'static str;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let bytes = crate::members::pq_public_key_bytes(s).ok_or("not an age1pq1… key")?;
        <Kem as hpke::Kem>::PublicKey::from_bytes(&bytes)
            .map(Recipient)
            .map_err(|_| "not a valid ML-KEM-768 + X25519 public key")
    }
}

impl fmt::Display for Recipient {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let s = bech32::encode(RECIPIENT_HRP, self.0.to_bytes().to_base32(), Variant::Bech32)
            .map_err(|_| fmt::Error)?;
        f.write_str(&s)
    }
}

impl fmt::Debug for Recipient {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{self}")
    }
}

impl age::Recipient for Recipient {
    fn wrap_file_key(
        &self,
        file_key: &FileKey,
    ) -> Result<(Vec<Stanza>, HashSet<String>), age::EncryptError> {
        let (enc, body) = hpke::single_shot_seal_with_rng::<Aead, Kdf, Kem>(
            &hpke::OpModeS::Base,
            &self.0,
            INFO,
            file_key.expose_secret(),
            &[],
            &mut OsRng,
        )
        .map_err(|e| age::EncryptError::Io(std::io::Error::other(e.to_string())))?;
        Ok((
            vec![Stanza {
                tag: STANZA_TYPE.to_string(),
                args: vec![STANDARD_NO_PAD.encode(enc.to_bytes())],
                body,
            }],
            HashSet::from([LABEL.to_string()]),
        ))
    }
}

/// An `AGE-SECRET-KEY-PQ-1…` identity: the 32-byte X-Wing seed.
#[derive(Clone)]
pub struct Identity(<Kem as hpke::Kem>::PrivateKey);

impl Identity {
    /// A new identity from the OS random source.
    pub fn generate() -> Identity {
        let seed = Zeroizing::new(crate::random_bytes::<32>());
        Identity::from_seed(&seed)
    }

    fn from_seed(seed: &[u8; 32]) -> Identity {
        Identity(<Kem as hpke::Kem>::PrivateKey::from_bytes(seed).expect("32 bytes is a seed"))
    }

    /// The `AGE-SECRET-KEY-PQ-1…` line. Secret: written only to the 0600
    /// identity file on the member's own machine.
    pub fn to_secret_string(&self) -> SecretString {
        let seed = Zeroizing::new(self.0.to_bytes());
        let mut lower = bech32::encode(IDENTITY_HRP, seed.as_slice().to_base32(), Variant::Bech32)
            .expect("the HRP is valid");
        let upper = lower.to_uppercase();
        lower.zeroize();
        SecretString::from(upper)
    }

    /// Its public key.
    pub fn to_public(&self) -> Recipient {
        Recipient(Kem::sk_to_pk(&self.0))
    }
}

impl std::str::FromStr for Identity {
    type Err = &'static str;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let s = s.trim();
        if !s.starts_with("AGE-SECRET-KEY-PQ-1") {
            return Err("not an AGE-SECRET-KEY-PQ-1… identity");
        }
        let (hrp, data, variant) = bech32::decode(s).map_err(|_| "invalid bech32")?;
        if hrp != IDENTITY_HRP || variant != Variant::Bech32 {
            return Err("not an AGE-SECRET-KEY-PQ-1… identity");
        }
        let bytes = Zeroizing::new(Vec::<u8>::from_base32(&data).map_err(|_| "invalid bech32")?);
        let seed: &[u8; 32] = bytes
            .as_slice()
            .try_into()
            .map_err(|_| "a post-quantum identity is 32 bytes")?;
        Ok(Identity::from_seed(seed))
    }
}

impl age::Identity for Identity {
    fn unwrap_stanza(&self, stanza: &Stanza) -> Option<Result<FileKey, age::DecryptError>> {
        if stanza.tag != STANZA_TYPE {
            return None;
        }
        let header = || Some(Err(age::DecryptError::InvalidHeader));
        // Exactly one canonical base64 argument of 1120 bytes, and a body of
        // the file key plus its tag (c2sp.org/age#mlkem768x25519-recipient-stanza).
        let [arg] = &stanza.args[..] else {
            return header();
        };
        // 1120 bytes are 1494 base64 characters: anything else is refused
        // before it is decoded, however long it is.
        if arg.len() != 1494 {
            return header();
        }
        let Some(enc) = STANDARD_NO_PAD
            .decode(arg)
            .ok()
            .filter(|b| b.len() == ENC_BYTES && STANDARD_NO_PAD.encode(b) == *arg)
        else {
            return header();
        };
        if stanza.body.len() != FILE_KEY_BYTES + 16 {
            return header();
        }
        // The X25519 half must contribute (the spec's rule; hpke's X-Wing
        // would accept an all-zero shared secret).
        let seed = Zeroizing::new(<[u8; 32]>::from(self.0.to_bytes()));
        let strict = x_wing::DecapsulationKeyRejectNonContrib::from(*seed);
        let Ok(ct) = x_wing::Ciphertext::try_from(enc.as_slice()) else {
            return header();
        };
        if strict.try_decapsulate(&ct).is_err() {
            return header();
        }
        let Ok(enc) = <Kem as hpke::Kem>::EncappedKey::from_bytes(&enc) else {
            return header();
        };
        match hpke::single_shot_open::<Aead, Kdf, Kem>(
            &hpke::OpModeR::Base,
            &self.0,
            &enc,
            INFO,
            &stanza.body,
            &[],
        ) {
            // A decapsulation failure would happen for every identity: the
            // header is bad. An AEAD failure only means "not this identity".
            Err(hpke::HpkeError::DecapError) => header(),
            Err(_) => None,
            Ok(mut pt) => {
                let key = FileKey::init_with_mut(|k| k.copy_from_slice(&pt));
                pt.zeroize();
                Some(Ok(key))
            }
        }
    }
}

/// Is `text` a post-quantum identity file (age-keygen's layout: `#` comment
/// lines and one `AGE-SECRET-KEY-PQ-1…` line)?
pub fn is_identity_file(text: &str) -> bool {
    identity_line(text).is_some()
}

fn identity_line(text: &str) -> Option<&str> {
    text.lines()
        .map(str::trim)
        .find(|l| l.starts_with("AGE-SECRET-KEY-PQ-1"))
}

/// The identity in an identity file.
pub fn parse_identity_file(text: &str) -> Result<Identity, &'static str> {
    identity_line(text)
        .ok_or("no AGE-SECRET-KEY-PQ-1… line")?
        .parse()
}

/// The identity file `gunnar crypt keygen` writes: age-keygen's layout, so
/// `age -d -i` reads it too.
pub fn identity_file(identity: &Identity, created: &str) -> Zeroizing<String> {
    Zeroizing::new(format!(
        "# created: {created}\n# public key: {}\n{}\n",
        identity.to_public(),
        identity.to_secret_string().expose_secret()
    ))
}