wyrd-client 0.3.0

//! Client-owned artifact provenance and aggregate manifest hashing.

use std::collections::BTreeMap;
use std::path::{Path, PathBuf};

use base64::Engine;
use wyrd_spec::error::WyrdError;
use wyrd_spec::registry::{CardSubmission, RelativeArtifactPath};

use crate::cards::error::RegistryEngineError;

/// Validate source bytes before any declaration network call and stamp the
/// manifest hash that the server validates at the declaration boundary.
///
/// Entries are sorted in place. An entry authored without `size_bytes` or
/// `sha256` takes the value computed from its source; a declared value must
/// match it. The manifest hash is stamped only after every present source is
/// filled in or matched.
///
/// # Errors
///
/// Returns an IO error when a source cannot be opened or read,
/// `RegistrySpecTooLarge` when a source size overflows `u64`,
/// `RegistryManifestHashMismatch` when a source's size or SHA-256 differs from
/// a value its entry declares, and a serialization error when the manifest
/// cannot be canonicalized.
///
/// # Cancellation
///
/// Cancelling drops the open file handle and may leave the entries sorted
/// without a stamped hash; no network call has been made.
pub(crate) async fn validate_and_stamp(
    submission: &mut CardSubmission,
    sources: &BTreeMap<RelativeArtifactPath, PathBuf>,
) -> Result<(), RegistryEngineError> {
    submission
        .artifacts
        .sort_by(|left, right| left.relative_path.cmp(&right.relative_path));
    for artifact in &mut submission.artifacts {
        let Some(source) = sources.get(&artifact.relative_path) else {
            continue;
        };
        let (actual_size, actual_sha256) = hash_source(source).await?;
        let declared_size = *artifact.size_bytes.get_or_insert(actual_size);
        if actual_size != declared_size {
            return Err(WyrdError::RegistryManifestHashMismatch {
                message: "local artifact size does not match its manifest".to_owned(),
                details: serde_json::json!({
                    "relative_path": artifact.relative_path,
                    "expected_size_bytes": declared_size,
                    "actual_size_bytes": actual_size,
                }),
            }
            .into());
        }
        if *artifact.sha256.get_or_insert_with(|| actual_sha256.clone()) != actual_sha256 {
            return Err(WyrdError::RegistryManifestHashMismatch {
                message: "local artifact digest does not match its manifest".to_owned(),
                details: serde_json::json!({ "relative_path": artifact.relative_path }),
            }
            .into());
        }
    }
    if submission.artifacts.is_empty() {
        submission.metadata.artifact_hash = None;
        return Ok(());
    }
    let bytes =
        serde_jcs::to_vec(&submission.artifacts).map_err(WyrdError::from_spec_serialization)?;
    submission.metadata.artifact_hash = Some(blake3::hash(&bytes).to_hex().to_string());
    Ok(())
}

/// Stream one local artifact file and return its byte length and base64
/// SHA-256 digest, the two values compared against the manifest entry by
/// [`validate_and_stamp`].
///
/// # Errors
///
/// Returns an IO error when the file cannot be opened or read.
async fn hash_source(source: &Path) -> Result<(u64, String), RegistryEngineError> {
    let (size, digest) = crate::storage::upload::sha256_file(source).await?;
    Ok((
        size,
        base64::engine::general_purpose::STANDARD.encode(digest),
    ))
}

/// Source validation and manifest-hash stamping.
#[cfg(test)]
mod tests {
    use std::collections::BTreeMap;

    use base64::Engine as _;
    use sha2::Digest;
    use tempfile::tempdir;
    use wyrd_spec::api_version::ApiVersion;
    use wyrd_spec::envelope::{CardKind, Metadata};
    use wyrd_spec::registry::{ArtifactManifestEntry, CardSubmission, RelativeArtifactPath};

    use super::validate_and_stamp;
    use crate::cards::error::RegistryEngineError;

    /// A minimal Prompt submission carrying one artifact manifest entry.
    fn submission(entry: ArtifactManifestEntry) -> CardSubmission {
        CardSubmission {
            api_version: ApiVersion::v1(),
            kind: CardKind::Prompt,
            metadata: Metadata {
                name: "hash-test".parse().expect("test name is valid"),
                version: Some("1.0.0".parse().expect("test version is valid")),
                bump: None,
                space: Some("default".parse().expect("test space is valid")),
                uid: None,
                labels: BTreeMap::default(),
                annotations: BTreeMap::default(),
                spec_hash: None,
                artifact_hash: None,
                origin: None,
            },
            spec: serde_json::json!({}),
            artifacts: vec![entry],
        }
    }

    /// Matching sources pass and the manifest hash is stamped over sorted entries.
    #[tokio::test]
    async fn validates_sources_and_stamps_sorted_manifest_hash() {
        let directory = tempdir().expect("temporary directory creates");
        let source = directory.path().join("weights.bin");
        let bytes = b"weights";
        tokio::fs::write(&source, bytes)
            .await
            .expect("artifact writes");
        let relative_path = RelativeArtifactPath::new("weights.bin").expect("path is valid");
        let entry = ArtifactManifestEntry {
            relative_path: relative_path.clone(),
            sha256: Some(
                base64::engine::general_purpose::STANDARD.encode(sha2::Sha256::digest(bytes)),
            ),
            size_bytes: Some(u64::try_from(bytes.len()).expect("test bytes length fits u64")),
            content_type: None,
        };
        let mut card = submission(entry);
        let sources = BTreeMap::from([(relative_path, source)]);

        validate_and_stamp(&mut card, &sources)
            .await
            .expect("manifest validates");

        assert!(card.metadata.artifact_hash.is_some());
    }

    /// An entry authored without digest or size gets both from its source and
    /// a stamped manifest hash; a declared digest that differs from the
    /// source is still refused.
    #[tokio::test]
    async fn missing_metadata_is_computed_and_declared_mismatch_is_refused() {
        let directory = tempdir().expect("temporary directory creates");
        let source = directory.path().join("weights.bin");
        tokio::fs::write(&source, b"weights")
            .await
            .expect("artifact writes");
        let relative_path = RelativeArtifactPath::new("weights.bin").expect("path is valid");
        let entry = ArtifactManifestEntry {
            relative_path: relative_path.clone(),
            sha256: None,
            size_bytes: None,
            content_type: None,
        };
        let sources = BTreeMap::from([(relative_path, source)]);
        let mut card = submission(entry.clone());

        validate_and_stamp(&mut card, &sources)
            .await
            .expect("missing metadata is computed");

        let computed = &card.artifacts[0];
        assert_eq!(computed.size_bytes, Some(7));
        assert_eq!(
            computed.sha256.as_deref(),
            Some(
                base64::engine::general_purpose::STANDARD
                    .encode(sha2::Sha256::digest(b"weights"))
                    .as_str()
            )
        );
        assert!(card.metadata.artifact_hash.is_some());

        let mut wrong = submission(ArtifactManifestEntry {
            sha256: Some("AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA=".to_owned()),
            ..entry
        });
        let error = validate_and_stamp(&mut wrong, &sources)
            .await
            .expect_err("declared digest mismatch must fail");
        assert!(matches!(
            error,
            RegistryEngineError::Wyrd(
                wyrd_spec::error::WyrdError::RegistryManifestHashMismatch { .. }
            )
        ));
    }

    /// A digest mismatch between the declared entry and the file is refused before upload.
    #[tokio::test]
    async fn rejects_source_digest_mismatch_before_upload() {
        let directory = tempdir().expect("temporary directory creates");
        let source = directory.path().join("weights.bin");
        tokio::fs::write(&source, b"actual")
            .await
            .expect("artifact writes");
        let relative_path = RelativeArtifactPath::new("weights.bin").expect("path is valid");
        let entry = ArtifactManifestEntry {
            relative_path: relative_path.clone(),
            sha256: Some("AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA=".to_owned()),
            size_bytes: Some(6),
            content_type: None,
        };
        let mut card = submission(entry);
        let sources = BTreeMap::from([(relative_path, source)]);

        let error = validate_and_stamp(&mut card, &sources)
            .await
            .expect_err("digest mismatch must fail");
        assert!(matches!(
            error,
            RegistryEngineError::Wyrd(
                wyrd_spec::error::WyrdError::RegistryManifestHashMismatch { .. }
            )
        ));
    }

    /// A size mismatch in either direction is refused before upload.
    #[tokio::test]
    async fn rejects_short_and_long_sources_before_upload() {
        for (bytes, expected_size) in [(b"short".as_slice(), 6_u64), (b"longer".as_slice(), 5)] {
            let directory = tempdir().expect("temporary directory creates");
            let source = directory.path().join("weights.bin");
            tokio::fs::write(&source, bytes)
                .await
                .expect("artifact writes");
            let relative_path = RelativeArtifactPath::new("weights.bin").expect("path is valid");
            let entry = ArtifactManifestEntry {
                relative_path: relative_path.clone(),
                sha256: Some(
                    base64::engine::general_purpose::STANDARD.encode(sha2::Sha256::digest(bytes)),
                ),
                size_bytes: Some(expected_size),
                content_type: None,
            };
            let mut card = submission(entry);
            let sources = BTreeMap::from([(relative_path, source)]);

            let error = validate_and_stamp(&mut card, &sources)
                .await
                .expect_err("size mismatch must fail");
            assert!(matches!(
                error,
                RegistryEngineError::Wyrd(
                    wyrd_spec::error::WyrdError::RegistryManifestHashMismatch { .. }
                )
            ));
        }
    }
}