engram-rs 1.2.0

use crate::error::{EngramError, Result};
use ed25519_dalek::{Signature, Signer, SigningKey, Verifier, VerifyingKey};
use sha2::{Digest, Sha256};
use std::io::{Read, Seek, SeekFrom, Write};

/// Signature block magic: "SIGN"
pub const SIGNATURE_BLOCK_MAGIC: [u8; 4] = [0x53, 0x49, 0x47, 0x4E];

/// Fixed size of the signature block (104 bytes):
/// - Magic "SIGN": 4 bytes
/// - Algorithm: 1 byte (0x01 = Ed25519)
/// - Public key: 32 bytes
/// - Signature: 64 bytes
/// - Reserved: 3 bytes
pub const SIGNATURE_BLOCK_SIZE: usize = 104;

/// Algorithm identifier for Ed25519
pub const ALGORITHM_ED25519: u8 = 0x01;

/// Archive-level digital signature block
///
/// Written between the central directory and the ENDR record.
/// Signs a SHA-256 hash of all archive content preceding the signature block
/// (header + file entries + central directory).
#[derive(Debug, Clone)]
pub struct SignatureBlock {
    /// Ed25519 public key (32 bytes)
    pub public_key: [u8; 32],
    /// Ed25519 signature (64 bytes)
    pub signature: [u8; 64],
}

impl SignatureBlock {
    /// Sign the archive content
    ///
    /// Computes SHA-256 over all bytes from the start of the file to the current
    /// file position (which should be the end of the central directory), then
    /// signs that hash with the provided Ed25519 key.
    ///
    /// The file cursor must be positioned at the end of the central directory
    /// (where the signature block will be written).
    pub fn sign_archive<S: Read + Seek>(file: &mut S, signing_key: &SigningKey) -> Result<Self> {
        // Record where we are (end of content to sign)
        let content_end = file.stream_position()?;

        // Hash everything from byte 0 to content_end
        file.seek(SeekFrom::Start(0))?;
        let mut hasher = Sha256::new();
        let mut buf = [0u8; 8192];
        let mut remaining = content_end;

        while remaining > 0 {
            let to_read = std::cmp::min(remaining, buf.len() as u64) as usize;
            file.read_exact(&mut buf[..to_read])?;
            hasher.update(&buf[..to_read]);
            remaining -= to_read as u64;
        }

        let hash: [u8; 32] = hasher.finalize().into();

        // Sign the hash
        let signature = signing_key.sign(&hash);
        let public_key = signing_key.verifying_key();

        // Seek back to where we need to write
        file.seek(SeekFrom::Start(content_end))?;

        Ok(Self {
            public_key: public_key.to_bytes(),
            signature: signature.to_bytes(),
        })
    }

    /// Write the signature block
    pub fn write_to<W: Write>(&self, mut writer: W) -> Result<usize> {
        let mut bytes_written = 0;

        // Magic "SIGN"
        writer.write_all(&SIGNATURE_BLOCK_MAGIC)?;
        bytes_written += 4;

        // Algorithm (Ed25519 = 0x01)
        writer.write_all(&[ALGORITHM_ED25519])?;
        bytes_written += 1;

        // Public key (32 bytes)
        writer.write_all(&self.public_key)?;
        bytes_written += 32;

        // Signature (64 bytes)
        writer.write_all(&self.signature)?;
        bytes_written += 64;

        // Reserved (3 bytes)
        writer.write_all(&[0u8; 3])?;
        bytes_written += 3;

        debug_assert_eq!(bytes_written, SIGNATURE_BLOCK_SIZE);
        Ok(bytes_written)
    }

    /// Read a signature block from a reader
    pub fn read_from<R: Read>(mut reader: R) -> Result<Self> {
        // Read and verify magic
        let mut magic = [0u8; 4];
        reader.read_exact(&mut magic)?;
        if magic != SIGNATURE_BLOCK_MAGIC {
            return Err(EngramError::InvalidFormat(
                "Invalid signature block magic (expected SIGN)".to_string(),
            ));
        }

        // Read algorithm
        let mut algo = [0u8; 1];
        reader.read_exact(&mut algo)?;
        if algo[0] != ALGORITHM_ED25519 {
            return Err(EngramError::InvalidSignature);
        }

        // Read public key
        let mut public_key = [0u8; 32];
        reader.read_exact(&mut public_key)?;

        // Read signature
        let mut signature = [0u8; 64];
        reader.read_exact(&mut signature)?;

        // Skip reserved
        let mut reserved = [0u8; 3];
        reader.read_exact(&mut reserved)?;

        Ok(Self {
            public_key,
            signature,
        })
    }

    /// Verify the archive signature against the file content
    ///
    /// Computes SHA-256 over the archive content (from byte 0 up to the
    /// signature block position) and verifies it against the stored signature.
    ///
    /// `content_end` is the byte offset where the signature block starts
    /// (i.e., everything before the signature block is the signed content).
    pub fn verify<S: Read + Seek>(
        &self,
        file: &mut S,
        content_end: u64,
        verifying_key: &VerifyingKey,
    ) -> Result<bool> {
        // Hash everything from byte 0 to content_end
        file.seek(SeekFrom::Start(0))?;
        let mut hasher = Sha256::new();
        let mut buf = [0u8; 8192];
        let mut remaining = content_end;

        while remaining > 0 {
            let to_read = std::cmp::min(remaining, buf.len() as u64) as usize;
            file.read_exact(&mut buf[..to_read])?;
            hasher.update(&buf[..to_read]);
            remaining -= to_read as u64;
        }

        let hash: [u8; 32] = hasher.finalize().into();

        // Verify signature
        let signature = Signature::from_bytes(&self.signature);
        match verifying_key.verify(&hash, &signature) {
            Ok(()) => Ok(true),
            Err(_) => Ok(false),
        }
    }

    /// Get the public key as a VerifyingKey
    pub fn verifying_key(&self) -> Result<VerifyingKey> {
        VerifyingKey::from_bytes(&self.public_key).map_err(|_| EngramError::InvalidPublicKey)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use rand::rngs::OsRng;
    use std::io::Cursor;

    #[test]
    fn test_signature_block_roundtrip() {
        let block = SignatureBlock {
            public_key: [0xAA; 32],
            signature: [0xBB; 64],
        };

        let mut buf = Vec::new();
        let written = block.write_to(&mut buf).unwrap();
        assert_eq!(written, SIGNATURE_BLOCK_SIZE);
        assert_eq!(buf.len(), SIGNATURE_BLOCK_SIZE);

        let parsed = SignatureBlock::read_from(&buf[..]).unwrap();
        assert_eq!(parsed.public_key, block.public_key);
        assert_eq!(parsed.signature, block.signature);
    }

    #[test]
    fn test_sign_and_verify() {
        let mut csprng = OsRng;
        let signing_key = SigningKey::generate(&mut csprng);
        let verifying_key = signing_key.verifying_key();

        // Create fake archive content
        let content = b"fake archive header and file data and central directory";
        let mut cursor = Cursor::new(Vec::new());
        cursor.write_all(content).unwrap();

        // Sign
        let block = SignatureBlock::sign_archive(&mut cursor, &signing_key).unwrap();

        // Verify
        let content_end = content.len() as u64;
        let result = block
            .verify(&mut cursor, content_end, &verifying_key)
            .unwrap();
        assert!(result, "Signature should verify");
    }

    #[test]
    fn test_verify_fails_with_wrong_key() {
        let mut csprng = OsRng;
        let signing_key = SigningKey::generate(&mut csprng);
        let wrong_key = SigningKey::generate(&mut csprng);
        let wrong_verifying_key = wrong_key.verifying_key();

        let content = b"some archive content";
        let mut cursor = Cursor::new(Vec::new());
        cursor.write_all(content).unwrap();

        let block = SignatureBlock::sign_archive(&mut cursor, &signing_key).unwrap();

        let result = block
            .verify(&mut cursor, content.len() as u64, &wrong_verifying_key)
            .unwrap();
        assert!(!result, "Signature should NOT verify with wrong key");
    }

    #[test]
    fn test_verify_fails_with_tampered_content() {
        let mut csprng = OsRng;
        let signing_key = SigningKey::generate(&mut csprng);
        let verifying_key = signing_key.verifying_key();

        let content = b"original archive content";
        let mut cursor = Cursor::new(Vec::new());
        cursor.write_all(content).unwrap();

        let block = SignatureBlock::sign_archive(&mut cursor, &signing_key).unwrap();

        // Tamper with content
        cursor.seek(SeekFrom::Start(0)).unwrap();
        cursor.write_all(b"TAMPERED").unwrap();

        let result = block
            .verify(&mut cursor, content.len() as u64, &verifying_key)
            .unwrap();
        assert!(!result, "Signature should NOT verify after tampering");
    }

    #[test]
    fn test_invalid_magic_rejected() {
        let mut data = vec![0xFFu8; SIGNATURE_BLOCK_SIZE];
        data[4] = ALGORITHM_ED25519;

        let result = SignatureBlock::read_from(&data[..]);
        assert!(result.is_err());
    }
}