remora/crates/remora-object-source/src/identity.rs

use std::fmt;

use remora_domain::MemberPath;
use remora_outbound::CanonicalUrl;
use remora_stream::{IntegrityDigest, Sha256Accountant};
use remora_target::{CompletedPart, ContentSha256, MultipartIdentity};

use crate::ObjectSourceError;

#[derive(Clone, Eq, Hash, PartialEq)]
pub struct ObjectValidator(String);

impl ObjectValidator {
    pub fn parse(value: &str) -> Result<Self, ObjectSourceError> {
        if value.is_empty() || value.len() > 1024 || value.chars().any(char::is_control) {
            return Err(ObjectSourceError::InvalidRequest);
        }
        Ok(Self(value.to_owned()))
    }

    #[must_use]
    pub fn as_str(&self) -> &str {
        &self.0
    }
}

impl fmt::Debug for ObjectValidator {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter.write_str("ObjectValidator([redacted])")
    }
}

#[derive(Clone, Eq, PartialEq)]
pub struct ObjectSourceIdentity {
    endpoint: CanonicalUrl,
    bucket: String,
    key: String,
    size: u64,
    version_id: Option<String>,
    validator: Option<ObjectValidator>,
}

impl ObjectSourceIdentity {
    pub fn from_parts(
        endpoint: CanonicalUrl,
        bucket: impl Into<String>,
        key: impl Into<String>,
        size: u64,
        version_id: Option<String>,
        validator: Option<ObjectValidator>,
    ) -> Result<Self, ObjectSourceError> {
        let bucket = bucket.into();
        let key = key.into();
        let version_id = version_id.filter(|value| !value.is_empty() && value != "null");
        if bucket.is_empty()
            || bucket.chars().any(char::is_control)
            || key.is_empty()
            || key.starts_with('/')
            || key.chars().any(char::is_control)
            || version_id
                .as_deref()
                .is_some_and(|value| value.chars().any(char::is_control))
        {
            return Err(ObjectSourceError::InvalidRequest);
        }
        Ok(Self {
            endpoint,
            bucket,
            key,
            size,
            version_id,
            validator,
        })
    }

    #[must_use]
    pub fn endpoint(&self) -> &CanonicalUrl {
        &self.endpoint
    }

    #[must_use]
    pub fn bucket(&self) -> &str {
        &self.bucket
    }

    #[must_use]
    pub fn key(&self) -> &str {
        &self.key
    }

    #[must_use]
    pub fn size(&self) -> u64 {
        self.size
    }

    #[must_use]
    pub fn version_id(&self) -> Option<&str> {
        self.version_id.as_deref()
    }

    #[must_use]
    pub fn validator(&self) -> Option<&ObjectValidator> {
        self.validator.as_ref()
    }

    #[must_use]
    pub fn is_resumable(&self) -> bool {
        self.version_id.is_some() || self.validator.is_some()
    }
}

impl fmt::Debug for ObjectSourceIdentity {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter
            .debug_struct("ObjectSourceIdentity")
            .field("size", &self.size)
            .field("versioned", &self.version_id.is_some())
            .field("has_validator", &self.validator.is_some())
            .finish_non_exhaustive()
    }
}

#[derive(Clone, Eq, PartialEq)]
pub struct DestinationResumeProof {
    endpoint: CanonicalUrl,
    identity: MultipartIdentity,
    completed_part_count: u32,
    expected_content_sha256: Option<ContentSha256>,
    completed_bytes: u64,
    completed_parts_integrity: IntegrityDigest,
}

impl DestinationResumeProof {
    pub fn new(
        endpoint: CanonicalUrl,
        identity: MultipartIdentity,
        completed_parts: Vec<CompletedPart>,
        expected_content_sha256: Option<ContentSha256>,
    ) -> Result<Self, ObjectSourceError> {
        if identity.target_store_id.is_empty()
            || identity.bucket.is_empty()
            || identity.key.is_empty()
            || identity.upload_id.is_empty()
        {
            return Err(ObjectSourceError::InvalidRequest);
        }
        let (completed_part_count, completed_bytes, completed_parts_integrity) =
            summarize_parts(&completed_parts)?;
        Ok(Self {
            endpoint,
            identity,
            completed_part_count,
            expected_content_sha256,
            completed_bytes,
            completed_parts_integrity,
        })
    }

    #[must_use]
    pub fn identity(&self) -> &MultipartIdentity {
        &self.identity
    }

    #[must_use]
    pub fn endpoint(&self) -> &CanonicalUrl {
        &self.endpoint
    }

    #[must_use]
    pub fn completed_part_count(&self) -> u32 {
        self.completed_part_count
    }

    #[must_use]
    pub fn expected_content_sha256(&self) -> Option<&ContentSha256> {
        self.expected_content_sha256.as_ref()
    }

    #[must_use]
    pub fn completed_bytes(&self) -> u64 {
        self.completed_bytes
    }

    pub(crate) fn matches_parts(&self, parts: &[CompletedPart]) -> bool {
        summarize_parts(parts).is_ok_and(|(count, bytes, integrity)| {
            count == self.completed_part_count
                && bytes == self.completed_bytes
                && integrity == self.completed_parts_integrity
        })
    }
}

impl fmt::Debug for DestinationResumeProof {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter
            .debug_struct("DestinationResumeProof")
            .field("completed_part_count", &self.completed_part_count)
            .field("completed_bytes", &self.completed_bytes)
            .finish_non_exhaustive()
    }
}

fn summarize_parts(
    parts: &[CompletedPart],
) -> Result<(u32, u64, IntegrityDigest), ObjectSourceError> {
    if parts.is_empty() || parts.len() > u32::MAX as usize {
        return Err(ObjectSourceError::InvalidRequest);
    }
    let mut previous = None;
    let mut completed_bytes = 0_u64;
    let mut accountant = Sha256Accountant::new();
    for part in parts {
        let size = part
            .size
            .filter(|size| *size > 0)
            .ok_or(ObjectSourceError::InvalidRequest)?;
        if part.part_number == 0
            || previous.is_some_and(|number| part.part_number <= number)
            || part.etag.is_empty()
        {
            return Err(ObjectSourceError::InvalidRequest);
        }
        previous = Some(part.part_number);
        completed_bytes = completed_bytes
            .checked_add(size)
            .ok_or(ObjectSourceError::InvalidRequest)?;
        let etag_length =
            u64::try_from(part.etag.len()).map_err(|_| ObjectSourceError::InvalidRequest)?;
        accountant.update(&part.part_number.to_be_bytes());
        accountant.update(&size.to_be_bytes());
        accountant.update(&etag_length.to_be_bytes());
        accountant.update(part.etag.as_bytes());
    }
    let count = u32::try_from(parts.len()).map_err(|_| ObjectSourceError::InvalidRequest)?;
    let integrity = accountant
        .finalize()
        .digest()
        .ok_or(ObjectSourceError::Internal)?
        .clone();
    Ok((count, completed_bytes, integrity))
}

#[derive(Clone, Eq, PartialEq)]
pub struct ObjectResumeCheckpoint {
    source: ObjectSourceIdentity,
    result_path: MemberPath,
    offset: u64,
    prefix_integrity: IntegrityDigest,
    destination: DestinationResumeProof,
}

impl ObjectResumeCheckpoint {
    pub fn new(
        source: ObjectSourceIdentity,
        result_path: MemberPath,
        offset: u64,
        prefix_integrity: IntegrityDigest,
        destination: DestinationResumeProof,
    ) -> Result<Self, ObjectSourceError> {
        if !source.is_resumable()
            || offset == 0
            || offset >= source.size
            || destination.completed_bytes != offset
        {
            return Err(ObjectSourceError::InvalidRequest);
        }
        Ok(Self {
            source,
            result_path,
            offset,
            prefix_integrity,
            destination,
        })
    }

    #[must_use]
    pub fn source(&self) -> &ObjectSourceIdentity {
        &self.source
    }

    #[must_use]
    pub fn result_path(&self) -> &MemberPath {
        &self.result_path
    }

    #[must_use]
    pub fn offset(&self) -> u64 {
        self.offset
    }

    #[must_use]
    pub fn prefix_integrity(&self) -> &IntegrityDigest {
        &self.prefix_integrity
    }

    #[must_use]
    pub fn destination(&self) -> &DestinationResumeProof {
        &self.destination
    }
}

impl fmt::Debug for ObjectResumeCheckpoint {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter
            .debug_struct("ObjectResumeCheckpoint")
            .field("source", &self.source)
            .field("result_path", &self.result_path.canonical())
            .field("offset", &self.offset)
            .field("destination", &self.destination)
            .finish_non_exhaustive()
    }
}