polykit_cache/
verification.rs

1//! Integrity verification for uploaded artifacts.
2
3use polykit_core::error::{Error, Result};
4use polykit_core::remote_cache::Artifact;
5use sha2::{Digest, Sha256};
6
7/// Verifies an uploaded artifact before storage.
8pub struct Verifier {
9    max_artifact_size: u64,
10}
11
12impl Verifier {
13    /// Creates a new verifier.
14    pub fn new(max_artifact_size: u64) -> Self {
15        Self {
16            max_artifact_size,
17        }
18    }
19
20    /// Verifies an uploaded artifact.
21    ///
22    /// # Arguments
23    ///
24    /// * `data` - Compressed artifact data
25    /// * `expected_cache_key` - The cache key from the URL
26    ///
27    /// # Returns
28    ///
29    /// Returns the parsed artifact and computed hash if verification passes.
30    ///
31    /// # Errors
32    ///
33    /// Returns an error if verification fails.
34    pub fn verify_upload(
35        &self,
36        data: &[u8],
37        expected_cache_key: &str,
38    ) -> Result<(Artifact, String)> {
39        // Check size limit
40        if data.len() as u64 > self.max_artifact_size {
41            return Err(Error::Adapter {
42                package: "verification".to_string(),
43                message: format!(
44                    "Artifact size {} exceeds maximum {}",
45                    data.len(),
46                    self.max_artifact_size
47                ),
48            });
49        }
50
51        // Compute SHA-256 hash
52        let mut hasher = Sha256::new();
53        hasher.update(data);
54        let computed_hash = format!("{:x}", hasher.finalize());
55
56        // Parse artifact
57        let artifact = Artifact::from_compressed(data.to_vec())?;
58
59        // Verify artifact integrity
60        polykit_core::remote_cache::ArtifactVerifier::verify(&artifact, Some(&computed_hash))?;
61
62        // Verify cache key matches
63        let metadata = artifact.metadata();
64        if metadata.cache_key_hash != expected_cache_key {
65            return Err(Error::Adapter {
66                package: "verification".to_string(),
67                message: format!(
68                    "Cache key mismatch: expected {}, got {}",
69                    expected_cache_key, metadata.cache_key_hash
70                ),
71            });
72        }
73
74        // Verify manifest integrity (already done by ArtifactVerifier, but double-check)
75        let manifest = artifact.manifest();
76        if manifest.total_size == 0 && !manifest.files.is_empty() {
77            return Err(Error::Adapter {
78                package: "verification".to_string(),
79                message: "Manifest has files but total_size is 0".to_string(),
80            });
81        }
82
83        Ok((artifact, computed_hash))
84    }
85
86}
87
88#[cfg(test)]
89mod tests {
90    use super::*;
91    use polykit_core::remote_cache::Artifact;
92    use std::collections::BTreeMap;
93    use std::path::PathBuf;
94
95    #[test]
96    fn test_verify_valid_artifact() {
97        let verifier = Verifier::new(1024 * 1024);
98
99        let mut output_files = BTreeMap::new();
100        output_files.insert(PathBuf::from("file.txt"), b"content".to_vec());
101
102        let cache_key = "aabbccdd11223344556677889900aabbccddeeff";
103        let artifact = Artifact::new(
104            "test".to_string(),
105            "build".to_string(),
106            "echo".to_string(),
107            cache_key.to_string(),
108            output_files,
109        )
110        .unwrap();
111
112        let compressed = artifact.compressed_data().to_vec();
113        let result = verifier.verify_upload(&compressed, cache_key);
114
115        assert!(result.is_ok());
116    }
117
118    #[test]
119    fn test_verify_cache_key_mismatch() {
120        let verifier = Verifier::new(1024 * 1024);
121
122        let mut output_files = BTreeMap::new();
123        output_files.insert(PathBuf::from("file.txt"), b"content".to_vec());
124
125        let cache_key = "aabbccdd11223344556677889900aabbccddeeff";
126        let artifact = Artifact::new(
127            "test".to_string(),
128            "build".to_string(),
129            "echo".to_string(),
130            cache_key.to_string(),
131            output_files,
132        )
133        .unwrap();
134
135        let compressed = artifact.compressed_data().to_vec();
136        let result = verifier.verify_upload(&compressed, "different_key");
137
138        assert!(result.is_err());
139    }
140
141    #[test]
142    fn test_verify_size_limit() {
143        let verifier = Verifier::new(100); // Very small limit
144
145        let mut output_files = BTreeMap::new();
146        output_files.insert(PathBuf::from("file.txt"), vec![0u8; 1000]); // Large file
147
148        let cache_key = "aabbccdd11223344556677889900aabbccddeeff";
149        let artifact = Artifact::new(
150            "test".to_string(),
151            "build".to_string(),
152            "echo".to_string(),
153            cache_key.to_string(),
154            output_files,
155        )
156        .unwrap();
157
158        let compressed = artifact.compressed_data().to_vec();
159        let _result = verifier.verify_upload(&compressed, cache_key);
160
161        // Should fail due to size limit (even compressed, it might exceed)
162        // This test depends on compression ratio
163    }
164}
polykit_cache/verification.rs

polykit_cache/
verification.rs
