antimatter 2.0.13

use crate::capsule::capsule::SealedCapsule;
use crate::capsule::common::{
    deserialize_domain_id, serialize_domain_id, CapsuleError, FileHeader, BUNDLE_MAGIC_BYTES,
};
use crate::capsule::stream::{download_from_s3, upload_to_s3};
use chrono::Utc;
use ciborium::{de::from_reader, ser::into_writer};
use serde_tuple::{Deserialize_tuple, Serialize_tuple};
use std::fs::File;
use std::io::{Cursor, Read, Write};
use std::path::Path;

#[derive(Serialize_tuple, Deserialize_tuple)]
pub struct BundleHeader {
    // domain_id is the authorized domain that created the bundle.
    #[serde(
        serialize_with = "serialize_domain_id",
        deserialize_with = "deserialize_domain_id"
    )]
    pub domain_id: String,
    pub created: i64,
}

#[derive(Serialize_tuple, Deserialize_tuple)]
pub struct CapsuleBundle {
    pub header: BundleHeader,
    pub body: Vec<SealedCapsule>,
}

impl CapsuleBundle {
    /// Creates a new `CapsuleBundle` used for storing multiple `Capsule`s with
    /// their metadata.
    ///
    /// # Arguments
    ///
    /// * `domain_id` - The domain ID associated with the `Session`.
    ///
    /// # Returns
    ///
    /// A new `CapsuleBundle`.
    pub fn new(domain_id: String) -> Self {
        CapsuleBundle {
            header: BundleHeader {
                created: Utc::now().timestamp_millis(),
                domain_id,
            },
            body: vec![],
        }
    }

    /// Reads and creates a `CapsuleBundle` object from the provided reader.
    ///
    /// # Arguments
    ///
    /// * `R` - A reader to the capsule's byte buffer.
    ///
    /// # Returns
    ///
    /// A new `CapsuleBundle`.
    pub fn import_from_reader<R: Read>(mut r: R) -> Result<Self, CapsuleError> {
        let fh: FileHeader = from_reader(&mut r)
            .map_err(|e| CapsuleError::CBORDecodeFailed(format!("decoding header: {}", e)))?;
        if fh.magic != BUNDLE_MAGIC_BYTES {
            return Err(CapsuleError::BadMagic("invalid magic".to_string()));
        } else if fh.version != 1 {
            return Err(CapsuleError::UnsupportedVersion(format!(
                "unrecognized file format version {}",
                fh.version
            )));
        }
        Self::import_from_reader_without_header(&mut r)
    }

    pub fn import_from_reader_without_header<R: Read>(mut r: R) -> Result<Self, CapsuleError> {
        from_reader(&mut r)
            .map_err(|e| CapsuleError::CBORDecodeFailed(format!("decoding v1 bundle: {}", e)))
    }

    /// Reads and creates a `CapsuleBundle` object from a byte vector.
    ///
    /// # Arguments
    ///
    /// * `bytes` - A byte vector containing the `CapsuleBundle`.
    ///
    /// # Returns
    ///
    /// A new `CapsuleBundle`.
    pub fn import_from_bytes(bytes: Vec<u8>) -> Result<Self, CapsuleError> {
        Self::import_from_reader(Cursor::new(bytes))
    }

    /// Reads and creates a `CapsuleBundle` object from file.
    ///
    /// # Arguments
    ///
    /// * `path` - The path to the `CapsuleBundle` file.
    ///
    /// # Returns
    ///
    /// A new `CapsuleBundle`.
    pub fn import_from_file<P: AsRef<Path>>(path: P) -> Result<Self, CapsuleError> {
        let file = File::open(path).map_err(|e| CapsuleError::FileIOError(format!("{}", e)))?;
        Self::import_from_reader(file)
    }

    /// Reads and creates a `CapsuleBundle` object from S3.
    ///
    /// # Arguments
    ///
    /// * `bucket` - The S3 bucket containing the `CapsuleBundle`'s resource.
    /// * `key` - The S3 bucket key for the `CapsuleBundle`'s resource.
    ///
    /// # Returns
    ///
    /// A new `CapsuleBundle`.
    pub async fn import_from_s3(bucket: String, key: String) -> Result<Self, CapsuleError> {
        let raw = download_from_s3(bucket, key)
            .await
            .map_err(|e| CapsuleError::StreamReadFailure(format!("{}", e)))?;
        Self::import_from_bytes(raw)
    }

    /// Writes a `CapsuleBundle`'s bytes to the given writer.
    ///
    /// # Arguments
    ///
    /// * `W` - A writer to the capsule's destination.
    pub fn export_to_writer<W: Write>(&self, mut w: W) -> Result<(), CapsuleError> {
        into_writer(&FileHeader::new(1), &mut w)
            .map_err(|e| CapsuleError::CBOREncodeFailed(format!("encoding header: {}", e)))?;
        into_writer(self, &mut w)
            .map_err(|e| CapsuleError::CBOREncodeFailed(format!("encoding v1 bundle: {}", e)))?;
        Ok(())
    }

    /// Writes a `CapsuleBundle`'s bytes to the a byte vector.
    ///
    /// # Returns
    ///
    /// A byte vector containing the `CapsuleBundle`.
    pub fn export_to_bytes(&self) -> Result<Vec<u8>, CapsuleError> {
        let mut buffer = Cursor::new(Vec::new());
        self.export_to_writer(&mut buffer)?;
        Ok(buffer.into_inner())
    }

    /// Writes a `CapsuleBundle`'s bytes to the destination file path.
    ///
    /// # Arguments
    ///
    /// * `path` - The destination file path.
    ///
    /// # Returns
    ///
    /// Number of bytes written to path.
    pub fn export_to_file<P: AsRef<Path>>(&self, path: P) -> Result<u64, CapsuleError> {
        let mut file =
            File::create(path).map_err(|e| CapsuleError::FileIOError(format!("{}", e)))?;
        self.export_to_writer(&mut file)?;
        Ok(file
            .metadata()
            .map_err(|e| CapsuleError::FileIOError(format!("getting file size: {}", e)))?
            .len())
    }

    /// Writes a `CapsuleBundle`'s bytes to S3.
    ///
    /// # Arguments
    ///
    /// * `bucket` - The S3 bucket containing the `CapsuleBundle`'s resource.
    /// * `key` - The S3 bucket key for the `CapsuleBundle`'s resource.
    pub async fn export_to_s3(&mut self, bucket: String, key: String) -> Result<(), CapsuleError> {
        let bytes = self.export_to_bytes()?;
        upload_to_s3(bucket, key, bytes)
            .await
            .map_err(|e| CapsuleError::StreamWriteFailure(format!("{}", e)))
    }

    /// Adds a `Capsule` to a `CapsuleBundle`.
    ///
    /// # Arguments
    ///
    /// `capsule` - The `Capsule` to add to the `CapsuleBundle`.
    pub fn add_capsule(&mut self, capsule: SealedCapsule) {
        self.body.push(capsule)
    }

    /// Collects all capsule IDs for `Capsule`s in this `CapsuleBundle`.
    ///
    /// # Returns
    ///
    /// A vector of all capsule IDs in this `CapsuleBundle`.
    pub fn capsule_ids(self) -> Vec<String> {
        self.body
            .iter()
            .map(|cap| cap.header.capsule_id.clone())
            .collect()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::capsule::capsule::tests::*;
    use tempfile::NamedTempFile;

    #[test]
    fn test_create_and_add_capsule() {
        let domain_id = format!("dm-{}", generate_random_base58_string(11));
        let mut bundle = CapsuleBundle::new(domain_id.clone());

        let capsule = create_sealed_capsule();
        bundle.add_capsule(capsule);

        assert_eq!(bundle.header.domain_id, domain_id);
        assert!(!bundle.body.is_empty());
    }

    #[test]
    fn test_export_import_to_from_bytes() {
        let domain_id = format!("dm-{}", generate_random_base58_string(11));
        let mut bundle = CapsuleBundle::new(domain_id);
        bundle.add_capsule(create_sealed_capsule());

        let exported_bytes = bundle.export_to_bytes().expect("Export to bytes failed");
        assert!(!exported_bytes.is_empty());

        let imported_bundle =
            CapsuleBundle::import_from_bytes(exported_bytes).expect("Import from bytes failed");
        assert_eq!(imported_bundle.header.domain_id, bundle.header.domain_id);
        assert_eq!(imported_bundle.body.len(), bundle.body.len());
        assert!(!imported_bundle.capsule_ids().is_empty());
    }

    #[test]
    fn test_export_import_to_from_file() {
        let domain_id = format!("dm-{}", generate_random_base58_string(11));
        let mut bundle = CapsuleBundle::new(domain_id);
        bundle.add_capsule(create_sealed_capsule());

        let temp_file = NamedTempFile::new().expect("Failed to create temp file");
        let size = CapsuleBundle::export_to_file(&bundle, temp_file.path())
            .expect("Export to file failed");

        // Make sure the file is not empty
        let metadata = std::fs::metadata(temp_file.path()).expect("Failed to read metadata");
        assert_eq!(metadata.len(), size);

        let imported_bundle =
            CapsuleBundle::import_from_file(temp_file.path()).expect("Import from file failed");
        assert_eq!(imported_bundle.header.domain_id, bundle.header.domain_id);
        assert_eq!(imported_bundle.body.len(), bundle.body.len());
    }

    #[test]
    fn test_is_capsule_true() {
        let domain_id = format!("dm-{}", generate_random_base58_string(11));
        let mut bundle = CapsuleBundle::new(domain_id.clone());

        let capsule = create_sealed_capsule();
        bundle.add_capsule(capsule);

        let bytes = bundle.export_to_bytes().expect("export bytes failed");
        assert!(FileHeader::is_capsule(Cursor::new(bytes)).unwrap().1);
    }

    #[test]
    fn test_is_capsule_false() {
        let bytes = generate_random_vec(16);
        assert!(!FileHeader::is_capsule(Cursor::new(bytes)).unwrap().1);

        let bytes_short = generate_random_vec(1);
        assert!(!FileHeader::is_capsule(Cursor::new(bytes_short)).unwrap().1);
    }

    #[test]
    fn test_is_capsule_valid() {
        let bytes = vec![
            130, 136, 24, 249, 24, 216, 24, 132, 24, 83, 24, 144, 24, 201, 2, 24, 104, 3, 131,
        ];
        assert!(FileHeader::is_capsule(Cursor::new(bytes)).unwrap().1);
    }

    #[test]
    fn test_is_capsule_invalid() {
        let bytes = vec![
            130, 136, 24, 55, 24, 216, 24, 132, 24, 83, 24, 144, 24, 201, 2, 24, 104, 3, 131,
        ];
        assert!(!FileHeader::is_capsule(Cursor::new(bytes)).unwrap().1);
    }
}