kcode-server-object-envelopes 0.1.1

use chrono::{DateTime, Utc};
use kcode_kweb_db::ObjectId;

use crate::codec::{Reader, push_string};
use crate::{Error, Result};

const PROVENANCE_MAGIC: &[u8; 8] = b"KPROV\0\x01\0";

/// One object referenced by a stored provenance payload.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct StoredArtifact {
    pub object_id: ObjectId,
    pub original_filename: String,
    pub media_type: String,
    pub role: String,
    pub byte_length: u64,
    pub sha256: [u8; 32],
}

/// KennedyServer provenance stored as one opaque Kweb object payload.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct StoredProvenance {
    pub data: String,
    pub source: String,
    pub source_created_at: DateTime<Utc>,
    pub artifacts: Vec<StoredArtifact>,
}

/// Encodes one provenance payload without storing it.
pub fn encode_provenance(value: &StoredProvenance) -> Result<Vec<u8>> {
    let artifact_count = u64::try_from(value.artifacts.len())
        .map_err(|_| Error::new("too many provenance artifacts"))?;
    let mut encoded_length = PROVENANCE_MAGIC
        .len()
        .checked_add(8 + 4)
        .ok_or_else(|| Error::new("provenance encoded length overflow"))?;
    encoded_length = checked_string_length(encoded_length, &value.source, "provenance source")?;
    encoded_length = checked_string_length(encoded_length, &value.data, "provenance data")?;
    encoded_length = encoded_length
        .checked_add(8)
        .ok_or_else(|| Error::new("provenance encoded length overflow"))?;

    for artifact in &value.artifacts {
        encoded_length = encoded_length
            .checked_add(artifact.object_id.to_bytes().len())
            .ok_or_else(|| Error::new("provenance encoded length overflow"))?;
        encoded_length = checked_string_length(
            encoded_length,
            &artifact.original_filename,
            "provenance artifact filename",
        )?;
        encoded_length = checked_string_length(
            encoded_length,
            &artifact.media_type,
            "provenance artifact media type",
        )?;
        encoded_length =
            checked_string_length(encoded_length, &artifact.role, "provenance artifact role")?;
        encoded_length = encoded_length
            .checked_add(8 + 32)
            .ok_or_else(|| Error::new("provenance encoded length overflow"))?;
    }

    let mut output = Vec::new();
    output
        .try_reserve_exact(encoded_length)
        .map_err(|_| Error::new("unable to allocate encoded provenance object"))?;
    output.extend_from_slice(PROVENANCE_MAGIC);
    output.extend_from_slice(&value.source_created_at.timestamp().to_be_bytes());
    output.extend_from_slice(
        &value
            .source_created_at
            .timestamp_subsec_nanos()
            .to_be_bytes(),
    );
    push_string(&mut output, &value.source, "provenance source")?;
    push_string(&mut output, &value.data, "provenance data")?;
    output.extend_from_slice(&artifact_count.to_be_bytes());
    for artifact in &value.artifacts {
        output.extend_from_slice(&artifact.object_id.to_bytes());
        push_string(
            &mut output,
            &artifact.original_filename,
            "provenance artifact filename",
        )?;
        push_string(
            &mut output,
            &artifact.media_type,
            "provenance artifact media type",
        )?;
        push_string(&mut output, &artifact.role, "provenance artifact role")?;
        output.extend_from_slice(&artifact.byte_length.to_be_bytes());
        output.extend_from_slice(&artifact.sha256);
    }
    Ok(output)
}

/// Decodes one provenance payload without reading its Kweb object.
pub fn decode_provenance(bytes: &[u8]) -> Result<StoredProvenance> {
    let mut input = Reader::new(bytes);
    let magic = input.take(PROVENANCE_MAGIC.len(), "provenance marker")?;
    if magic != PROVENANCE_MAGIC {
        return Err(Error::new("provenance object has unknown format"));
    }
    let seconds = i64::from_be_bytes(input.array("provenance timestamp seconds")?);
    let nanos = u32::from_be_bytes(input.array("provenance timestamp nanoseconds")?);
    let source_created_at = DateTime::<Utc>::from_timestamp(seconds, nanos)
        .ok_or_else(|| Error::new("provenance object has invalid timestamp"))?;
    let source = input.string("provenance source")?;
    let data = input.string("provenance data")?;
    let artifact_count = usize::try_from(input.u64("provenance artifact count")?)
        .map_err(|_| Error::new("provenance artifact count exceeds usize"))?;
    if artifact_count > bytes.len() / 50 {
        return Err(Error::new("provenance artifact count is impossible"));
    }
    let mut artifacts = Vec::new();
    artifacts
        .try_reserve_exact(artifact_count)
        .map_err(|_| Error::new("unable to allocate provenance artifacts"))?;
    for _ in 0..artifact_count {
        let object_id = ObjectId::from_bytes(input.array("provenance artifact object ID")?)
            .map_err(|error| Error::new(error.to_string()))?;
        artifacts.push(StoredArtifact {
            object_id,
            original_filename: input.string("provenance artifact filename")?,
            media_type: input.string("provenance artifact media type")?,
            role: input.string("provenance artifact role")?,
            byte_length: input.u64("provenance artifact byte length")?,
            sha256: input.array("provenance artifact SHA-256")?,
        });
    }
    input.finish("provenance object")?;
    Ok(StoredProvenance {
        data,
        source,
        source_created_at,
        artifacts,
    })
}

fn checked_string_length(current: usize, value: &str, label: &str) -> Result<usize> {
    u32::try_from(value.len()).map_err(|_| Error::new(format!("{label} exceeds u32")))?;
    current
        .checked_add(4)
        .and_then(|length| length.checked_add(value.len()))
        .ok_or_else(|| Error::new(format!("{label} encoded length overflow")))
}