anubis-age 1.4.0

//! The age file format.

use std::io::{self, BufRead, Read, Write};

use anubis_core::format::{grease_the_joint, Stanza};

use crate::{
    error::DecryptError,
    primitives::{HmacKey, HmacWriter},
    EncryptError,
};

#[cfg(feature = "async")]
use futures::io::{
    AsyncBufRead, AsyncBufReadExt, AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt,
};

const AGE_MAGIC: &[u8] = b"anubis-encryption.org/";
const V1_MAGIC: &[u8] = b"v1";
const MAC_TAG: &[u8] = b"---";
const ENCODED_MAC_LENGTH: usize = 86; // Base64 encoding of 64 bytes (SHA-512)

#[derive(Debug, PartialEq)]
pub(crate) struct HeaderV1 {
    pub(crate) recipients: Vec<Stanza>,
    pub(crate) mac: [u8; 64], // HMAC-SHA512 produces 64 bytes
    /// The serialized bytes of this header. Set if we parsed this header from a reader,
    /// to handle the possibility that the header is not round-trip canonical, such as it
    /// containing a legacy stanza with a body of length 0 mod 64.
    ///
    /// We do not write this back out in `Header::write`, because we never write headers
    /// that we have parsed, and headers we generate for writing will never have this set.
    encoded_bytes: Option<Vec<u8>>,
}

impl HeaderV1 {
    pub(crate) fn new(recipients: Vec<Stanza>, mac_key: HmacKey) -> Result<Self, EncryptError> {
        let mut header = HeaderV1 {
            recipients,
            mac: [0; 64], // HMAC-SHA512 produces 64 bytes
            encoded_bytes: None,
        };

        // Keep the joint well oiled!
        header.recipients.push(grease_the_joint());

        let mut mac = HmacWriter::new(mac_key);
        cookie_factory::gen(write::header_v1_minus_mac(&header), &mut mac)
            .expect("can serialize Header into HmacWriter");
        header
            .mac
            .copy_from_slice(mac.finalize().into_bytes().as_slice());

        Ok(header)
    }

    pub(crate) fn verify_mac(&self, mac_key: HmacKey) -> Result<(), hmac::digest::MacError> {
        let mut mac = HmacWriter::new(mac_key);
        if let Some(bytes) = &self.encoded_bytes {
            // The MAC covers all bytes up to the end of the --- prefix.
            mac.write_all(&bytes[..bytes.len() - ENCODED_MAC_LENGTH - 2])
                .expect("can serialize Header into HmacWriter");
        } else {
            cookie_factory::gen(write::header_v1_minus_mac(self), &mut mac)
                .expect("can serialize Header into HmacWriter");
        }
        mac.verify(&self.mac)
    }

    /// Enforces structural requirements on the v1 header.
    ///
    /// Always returns `true` in ML-KEM-only mode.
    pub(crate) fn is_valid(&self) -> bool {
        true
    }
}

impl Header {
    pub(crate) fn read<R: Read>(mut input: R) -> Result<Self, DecryptError> {
        let mut data = vec![];
        loop {
            match read::header(&data) {
                Ok((_, mut header)) => {
                    if let Header::V1(h) = &mut header {
                        h.encoded_bytes = Some(data);
                    }
                    break Ok(header);
                }
                Err(nom::Err::Incomplete(nom::Needed::Size(n))) => {
                    // Read the needed additional bytes. We need to be careful how the
                    // parser is constructed, because if we read more than we need, the
                    // remainder of the input will be truncated.
                    let m = data.len();
                    let new_len = m + n.get();
                    data.resize(new_len, 0);
                    input.read_exact(&mut data[m..new_len])?;
                }
                Err(_) => {
                    break Err(DecryptError::InvalidHeader);
                }
            }
        }
    }

    pub(crate) fn read_buffered<R: BufRead>(mut input: R) -> Result<Self, DecryptError> {
        let mut data = vec![];
        loop {
            match read::header(&data) {
                Ok((_, mut header)) => {
                    if let Header::V1(h) = &mut header {
                        h.encoded_bytes = Some(data);
                    }
                    break Ok(header);
                }
                Err(nom::Err::Incomplete(nom::Needed::Size(_))) => {
                    // As we have a buffered reader, we can leverage the fact that the
                    // currently-defined header formats are newline-separated, to more
                    // efficiently read data for the parser to consume.
                    if input.read_until(b'\n', &mut data)? == 0 {
                        break Err(DecryptError::Io(io::Error::new(
                            io::ErrorKind::UnexpectedEof,
                            "Incomplete header",
                        )));
                    }
                }
                Err(_) => {
                    break Err(DecryptError::InvalidHeader);
                }
            }
        }
    }

    #[cfg(feature = "async")]
    #[cfg_attr(docsrs, doc(cfg(feature = "async")))]
    pub(crate) async fn read_async<R: AsyncRead + Unpin>(
        mut input: R,
    ) -> Result<Self, DecryptError> {
        let mut data = vec![];
        loop {
            match read::header(&data) {
                Ok((_, mut header)) => {
                    if let Header::V1(h) = &mut header {
                        h.encoded_bytes = Some(data);
                    }
                    break Ok(header);
                }
                Err(nom::Err::Incomplete(nom::Needed::Size(n))) => {
                    // Read the needed additional bytes. We need to be careful how the
                    // parser is constructed, because if we read more than we need, the
                    // remainder of the input will be truncated.
                    let m = data.len();
                    let new_len = m + n.get();
                    data.resize(new_len, 0);
                    input.read_exact(&mut data[m..new_len]).await?;
                }
                Err(_) => {
                    break Err(DecryptError::InvalidHeader);
                }
            }
        }
    }

    #[cfg(feature = "async")]
    #[cfg_attr(docsrs, doc(cfg(feature = "async")))]
    pub(crate) async fn read_async_buffered<R: AsyncBufRead + Unpin>(
        mut input: R,
    ) -> Result<Self, DecryptError> {
        let mut data = vec![];
        loop {
            match read::header(&data) {
                Ok((_, mut header)) => {
                    if let Header::V1(h) = &mut header {
                        h.encoded_bytes = Some(data);
                    }
                    break Ok(header);
                }
                Err(nom::Err::Incomplete(nom::Needed::Size(_))) => {
                    // As we have a buffered reader, we can leverage the fact that the
                    // currently-defined header formats are newline-separated, to more
                    // efficiently read data for the parser to consume.
                    if input.read_until(b'\n', &mut data).await? == 0 {
                        break Err(DecryptError::Io(io::Error::new(
                            io::ErrorKind::UnexpectedEof,
                            "Incomplete header",
                        )));
                    }
                }
                Err(_) => {
                    break Err(DecryptError::InvalidHeader);
                }
            }
        }
    }

    pub(crate) fn write<W: Write>(&self, mut output: W) -> io::Result<()> {
        cookie_factory::gen(write::header(self), &mut output)
            .map(|_| ())
            .map_err(|e| {
                io::Error::new(
                    io::ErrorKind::Other,
                    format!("failed to write header: {}", e),
                )
            })
    }

    #[cfg(feature = "async")]
    #[cfg_attr(docsrs, doc(cfg(feature = "async")))]
    pub(crate) async fn write_async<W: AsyncWrite + Unpin>(&self, mut output: W) -> io::Result<()> {
        let mut buf = vec![];
        cookie_factory::gen(write::header(self), &mut buf)
            .map(|_| ())
            .map_err(|e| {
                io::Error::new(
                    io::ErrorKind::Other,
                    format!("failed to write header: {}", e),
                )
            })?;

        output.write_all(&buf).await
    }
}

#[derive(Debug, PartialEq)]
pub(crate) enum Header {
    V1(HeaderV1),
    Unknown(String),
}

mod read {
    use anubis_core::format::read::{arbitrary_string, legacy_age_stanza};
    use nom::{
        branch::alt,
        bytes::streaming::{tag, take},
        character::streaming::newline,
        combinator::{map, map_opt},
        multi::many1,
        sequence::{pair, preceded, terminated},
        IResult,
    };

    use super::*;
    use crate::util::read::base64_arg;

    fn header_v1(input: &[u8]) -> IResult<&[u8], HeaderV1> {
        preceded(
            pair(tag(V1_MAGIC), newline),
            map(
                pair(
                    many1(legacy_age_stanza),
                    preceded(
                        pair(tag(MAC_TAG), tag(b" ")),
                        terminated(
                            map_opt(take(ENCODED_MAC_LENGTH), |tag| {
                                base64_arg::<_, 64, 86>(&tag) // SHA-512: 64 bytes, buffer size 86
                            }),
                            newline,
                        ),
                    ),
                ),
                |(recipients, mac)| HeaderV1 {
                    recipients: recipients.into_iter().map(Stanza::from).collect(),
                    mac,
                    encoded_bytes: None,
                },
            ),
        )(input)
    }

    /// From the anubis specification:
    /// ```text
    /// The first line of the header is anubis-encryption.org/ followed by an arbitrary
    /// version string. ... We describe version v1, other versions can change anything
    /// after the first line.
    /// ```
    pub(super) fn header(input: &[u8]) -> IResult<&[u8], Header> {
        preceded(
            tag(AGE_MAGIC),
            alt((
                map(header_v1, Header::V1),
                map(terminated(arbitrary_string, newline), |s| {
                    Header::Unknown(s.to_string())
                }),
            )),
        )(input)
    }
}

mod write {
    use anubis_core::format::write::age_stanza;
    use cookie_factory::{
        combinator::{slice, string},
        multi::all,
        sequence::tuple,
        SerializeFn, WriteContext,
    };
    use std::io::Write;

    use super::*;
    use crate::util::write::encoded_data;

    fn recipient_stanza<'a, W: 'a + Write>(r: &'a Stanza) -> impl SerializeFn<W> + 'a {
        move |w: WriteContext<W>| {
            let args: Vec<_> = r.args.iter().map(|s| s.as_str()).collect();
            let writer = age_stanza(&r.tag, &args, &r.body);
            writer(w)
        }
    }

    pub(super) fn header_v1_minus_mac<'a, W: 'a + Write>(
        h: &'a HeaderV1,
    ) -> impl SerializeFn<W> + 'a {
        tuple((
            slice(AGE_MAGIC),
            slice(V1_MAGIC),
            string("\n"),
            all(h.recipients.iter().map(move |r| recipient_stanza(r))),
            slice(MAC_TAG),
        ))
    }

    fn header_v1<'a, W: 'a + Write>(h: &'a HeaderV1) -> impl SerializeFn<W> + 'a {
        tuple((
            header_v1_minus_mac(h),
            string(" "),
            encoded_data(&h.mac),
            string("\n"),
        ))
    }

    pub(super) fn header<'a, W: 'a + Write>(h: &'a Header) -> impl SerializeFn<W> + 'a {
        move |w: WriteContext<W>| match h {
            Header::V1(v1) => header_v1(v1)(w),
            Header::Unknown(version) => tuple((slice(AGE_MAGIC), slice(version), string("\n")))(w),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::Header;

    // Header parsing is tested extensively in the protocol round-trip tests

    #[test]
    #[ignore] // Legacy test disabled - we only support ML-KEM now, not X25519
    fn parse_legacy_header() {
        // This test is kept for reference but ignored since we no longer support
        // X25519 or other legacy recipient types
    }
}