Add snapshot/range storage primitives, gate GET preconditions on served snapshot, support partial-decrypt Range GET for SSE-encrypted objects

crates/myfsio-crypto/src/aes_gcm.rs

@@ -145,6 +145,113 @@ pub fn decrypt_stream_chunked(
     Ok(chunk_count)
 }
 
+const GCM_TAG_LEN: usize = 16;
+
+pub fn decrypt_stream_chunked_range(
+    input_path: &Path,
+    output_path: &Path,
+    key: &[u8],
+    base_nonce: &[u8],
+    chunk_plain_size: usize,
+    plaintext_size: u64,
+    plain_start: u64,
+    plain_end_inclusive: u64,
+) -> Result<u64, CryptoError> {
+    if key.len() != 32 {
+        return Err(CryptoError::InvalidKeySize(key.len()));
+    }
+    if base_nonce.len() != 12 {
+        return Err(CryptoError::InvalidNonceSize(base_nonce.len()));
+    }
+    if chunk_plain_size == 0 {
+        return Err(CryptoError::EncryptionFailed(
+            "chunk_plain_size must be > 0".into(),
+        ));
+    }
+    if plaintext_size == 0 {
+        let _ = File::create(output_path)?;
+        return Ok(0);
+    }
+    if plain_start > plain_end_inclusive || plain_end_inclusive >= plaintext_size {
+        return Err(CryptoError::EncryptionFailed(format!(
+            "range [{}, {}] invalid for plaintext size {}",
+            plain_start, plain_end_inclusive, plaintext_size
+        )));
+    }
+
+    let key_arr: [u8; 32] = key.try_into().unwrap();
+    let nonce_arr: [u8; 12] = base_nonce.try_into().unwrap();
+    let cipher = Aes256Gcm::new(&key_arr.into());
+
+    let n = chunk_plain_size as u64;
+    let first_chunk = (plain_start / n) as u32;
+    let last_chunk = (plain_end_inclusive / n) as u32;
+    let total_chunks = plaintext_size.div_ceil(n) as u32;
+    let final_chunk_plain = plaintext_size - (total_chunks as u64 - 1) * n;
+
+    let mut infile = File::open(input_path)?;
+
+    let mut header = [0u8; HEADER_SIZE];
+    infile.read_exact(&mut header)?;
+    let stored_chunk_count = u32::from_be_bytes(header);
+    if stored_chunk_count != total_chunks {
+        return Err(CryptoError::EncryptionFailed(format!(
+            "chunk count mismatch: header says {}, plaintext_size implies {}",
+            stored_chunk_count, total_chunks
+        )));
+    }
+
+    let mut outfile = File::create(output_path)?;
+
+    let stride = n + GCM_TAG_LEN as u64 + HEADER_SIZE as u64;
+    let first_offset = HEADER_SIZE as u64 + first_chunk as u64 * stride;
+    infile.seek(SeekFrom::Start(first_offset))?;
+
+    let mut size_buf = [0u8; HEADER_SIZE];
+    let mut bytes_written: u64 = 0;
+
+    for chunk_index in first_chunk..=last_chunk {
+        infile.read_exact(&mut size_buf)?;
+        let ct_len = u32::from_be_bytes(size_buf) as usize;
+
+        let expected_plain = if chunk_index + 1 == total_chunks {
+            final_chunk_plain as usize
+        } else {
+            chunk_plain_size
+        };
+        let expected_ct = expected_plain + GCM_TAG_LEN;
+        if ct_len != expected_ct {
+            return Err(CryptoError::EncryptionFailed(format!(
+                "chunk {} stored length {} != expected {} (corrupt file or chunk_size mismatch)",
+                chunk_index, ct_len, expected_ct
+            )));
+        }
+
+        let mut encrypted = vec![0u8; ct_len];
+        infile.read_exact(&mut encrypted)?;
+
+        let nonce_bytes = derive_chunk_nonce(&nonce_arr, chunk_index)?;
+        let nonce = Nonce::from_slice(&nonce_bytes);
+        let decrypted = cipher
+            .decrypt(nonce, encrypted.as_ref())
+            .map_err(|_| CryptoError::DecryptionFailed(chunk_index))?;
+
+        let chunk_plain_start = chunk_index as u64 * n;
+        let chunk_plain_end_exclusive = chunk_plain_start + decrypted.len() as u64;
+
+        let slice_start = plain_start.saturating_sub(chunk_plain_start) as usize;
+        let slice_end = (plain_end_inclusive + 1).min(chunk_plain_end_exclusive);
+        let slice_end_local = (slice_end - chunk_plain_start) as usize;
+
+        if slice_end_local > slice_start {
+            outfile.write_all(&decrypted[slice_start..slice_end_local])?;
+            bytes_written += (slice_end_local - slice_start) as u64;
+        }
+    }
+
+    Ok(bytes_written)
+}
+
 pub async fn encrypt_stream_chunked_async(
     input_path: &Path,
     output_path: &Path,
@@ -230,6 +337,191 @@ mod tests {
         assert!(matches!(result, Err(CryptoError::InvalidKeySize(16))));
     }
 
+    fn write_file(path: &Path, data: &[u8]) {
+        std::fs::File::create(path).unwrap().write_all(data).unwrap();
+    }
+
+    fn make_encrypted_file(
+        dir: &Path,
+        data: &[u8],
+        key: &[u8; 32],
+        nonce: &[u8; 12],
+        chunk: usize,
+    ) -> std::path::PathBuf {
+        let input = dir.join("input.bin");
+        let encrypted = dir.join("encrypted.bin");
+        write_file(&input, data);
+        encrypt_stream_chunked(&input, &encrypted, key, nonce, Some(chunk)).unwrap();
+        encrypted
+    }
+
+    #[test]
+    fn test_range_within_single_chunk() {
+        let dir = tempfile::tempdir().unwrap();
+        let data: Vec<u8> = (0u8..=255).cycle().take(4096).collect();
+        let key = [0x33u8; 32];
+        let nonce = [0x07u8; 12];
+        let encrypted = make_encrypted_file(dir.path(), &data, &key, &nonce, 1024);
+        let out = dir.path().join("range.bin");
+
+        let n = decrypt_stream_chunked_range(
+            &encrypted,
+            &out,
+            &key,
+            &nonce,
+            1024,
+            data.len() as u64,
+            200,
+            399,
+        )
+        .unwrap();
+        assert_eq!(n, 200);
+        let got = std::fs::read(&out).unwrap();
+        assert_eq!(got, &data[200..400]);
+    }
+
+    #[test]
+    fn test_range_spanning_multiple_chunks() {
+        let dir = tempfile::tempdir().unwrap();
+        let data: Vec<u8> = (0..5000u32).map(|i| (i % 251) as u8).collect();
+        let key = [0x44u8; 32];
+        let nonce = [0x02u8; 12];
+        let encrypted = make_encrypted_file(dir.path(), &data, &key, &nonce, 512);
+        let out = dir.path().join("range.bin");
+
+        let n = decrypt_stream_chunked_range(
+            &encrypted,
+            &out,
+            &key,
+            &nonce,
+            512,
+            data.len() as u64,
+            100,
+            2999,
+        )
+        .unwrap();
+        assert_eq!(n, 2900);
+        let got = std::fs::read(&out).unwrap();
+        assert_eq!(got, &data[100..3000]);
+    }
+
+    #[test]
+    fn test_range_covers_final_partial_chunk() {
+        let dir = tempfile::tempdir().unwrap();
+        let data: Vec<u8> = (0..1300u32).map(|i| (i % 71) as u8).collect();
+        let key = [0x55u8; 32];
+        let nonce = [0x0au8; 12];
+        let encrypted = make_encrypted_file(dir.path(), &data, &key, &nonce, 512);
+        let out = dir.path().join("range.bin");
+
+        let n = decrypt_stream_chunked_range(
+            &encrypted,
+            &out,
+            &key,
+            &nonce,
+            512,
+            data.len() as u64,
+            900,
+            1299,
+        )
+        .unwrap();
+        assert_eq!(n, 400);
+        let got = std::fs::read(&out).unwrap();
+        assert_eq!(got, &data[900..1300]);
+    }
+
+    #[test]
+    fn test_range_full_object() {
+        let dir = tempfile::tempdir().unwrap();
+        let data: Vec<u8> = (0..2048u32).map(|i| (i % 13) as u8).collect();
+        let key = [0x11u8; 32];
+        let nonce = [0x33u8; 12];
+        let encrypted = make_encrypted_file(dir.path(), &data, &key, &nonce, 512);
+        let out = dir.path().join("range.bin");
+
+        let n = decrypt_stream_chunked_range(
+            &encrypted,
+            &out,
+            &key,
+            &nonce,
+            512,
+            data.len() as u64,
+            0,
+            data.len() as u64 - 1,
+        )
+        .unwrap();
+        assert_eq!(n, data.len() as u64);
+        let got = std::fs::read(&out).unwrap();
+        assert_eq!(got, data);
+    }
+
+    #[test]
+    fn test_range_wrong_key_fails() {
+        let dir = tempfile::tempdir().unwrap();
+        let data = b"range-auth-check".repeat(100);
+        let key = [0x66u8; 32];
+        let nonce = [0x09u8; 12];
+        let encrypted = make_encrypted_file(dir.path(), &data, &key, &nonce, 256);
+        let out = dir.path().join("range.bin");
+
+        let wrong = [0x67u8; 32];
+        let r = decrypt_stream_chunked_range(
+            &encrypted,
+            &out,
+            &wrong,
+            &nonce,
+            256,
+            data.len() as u64,
+            0,
+            data.len() as u64 - 1,
+        );
+        assert!(matches!(r, Err(CryptoError::DecryptionFailed(_))));
+    }
+
+    #[test]
+    fn test_range_out_of_bounds_rejected() {
+        let dir = tempfile::tempdir().unwrap();
+        let data = vec![0u8; 100];
+        let key = [0x22u8; 32];
+        let nonce = [0x44u8; 12];
+        let encrypted = make_encrypted_file(dir.path(), &data, &key, &nonce, 64);
+        let out = dir.path().join("range.bin");
+
+        let r = decrypt_stream_chunked_range(
+            &encrypted,
+            &out,
+            &key,
+            &nonce,
+            64,
+            data.len() as u64,
+            50,
+            200,
+        );
+        assert!(r.is_err());
+    }
+
+    #[test]
+    fn test_range_mismatched_chunk_size_detected() {
+        let dir = tempfile::tempdir().unwrap();
+        let data: Vec<u8> = (0..2048u32).map(|i| i as u8).collect();
+        let key = [0x77u8; 32];
+        let nonce = [0x88u8; 12];
+        let encrypted = make_encrypted_file(dir.path(), &data, &key, &nonce, 512);
+        let out = dir.path().join("range.bin");
+
+        let r = decrypt_stream_chunked_range(
+            &encrypted,
+            &out,
+            &key,
+            &nonce,
+            1024,
+            data.len() as u64,
+            0,
+            1023,
+        );
+        assert!(r.is_err());
+    }
+
     #[test]
     fn test_wrong_key_fails_decrypt() {
         let dir = tempfile::tempdir().unwrap();

crates/myfsio-crypto/src/encryption.rs

@@ -4,7 +4,9 @@ use rand::RngCore;
 use std::collections::HashMap;
 use std::path::Path;
 
-use crate::aes_gcm::{decrypt_stream_chunked, encrypt_stream_chunked, CryptoError};
+use crate::aes_gcm::{
+    decrypt_stream_chunked, decrypt_stream_chunked_range, encrypt_stream_chunked, CryptoError,
+};
 use crate::kms::KmsService;
 
 #[derive(Debug, Clone, PartialEq)]
@@ -37,6 +39,8 @@ pub struct EncryptionMetadata {
     pub nonce: String,
     pub encrypted_data_key: Option<String>,
     pub kms_key_id: Option<String>,
+    pub chunk_size: Option<usize>,
+    pub plaintext_size: Option<u64>,
 }
 
 impl EncryptionMetadata {
@@ -53,6 +57,15 @@ impl EncryptionMetadata {
         if let Some(ref kid) = self.kms_key_id {
             map.insert("x-amz-encryption-key-id".to_string(), kid.clone());
         }
+        if let Some(cs) = self.chunk_size {
+            map.insert("x-amz-encryption-chunk-size".to_string(), cs.to_string());
+        }
+        if let Some(ps) = self.plaintext_size {
+            map.insert(
+                "x-amz-encryption-plaintext-size".to_string(),
+                ps.to_string(),
+            );
+        }
         map
     }
 
@@ -64,6 +77,12 @@ impl EncryptionMetadata {
             nonce: nonce.clone(),
             encrypted_data_key: meta.get("x-amz-encrypted-data-key").cloned(),
             kms_key_id: meta.get("x-amz-encryption-key-id").cloned(),
+            chunk_size: meta
+                .get("x-amz-encryption-chunk-size")
+                .and_then(|s| s.parse().ok()),
+            plaintext_size: meta
+                .get("x-amz-encryption-plaintext-size")
+                .and_then(|s| s.parse().ok()),
         })
     }
 
@@ -76,6 +95,8 @@ impl EncryptionMetadata {
         meta.remove("x-amz-encryption-nonce");
         meta.remove("x-amz-encrypted-data-key");
         meta.remove("x-amz-encryption-key-id");
+        meta.remove("x-amz-encryption-chunk-size");
+        meta.remove("x-amz-encryption-plaintext-size");
     }
 }
 
@@ -212,6 +233,11 @@ impl EncryptionService {
             data_key
         };
 
+        let plaintext_size = tokio::fs::metadata(input_path)
+            .await
+            .map_err(CryptoError::Io)?
+            .len();
+
         let ip = input_path.to_owned();
         let op = output_path.to_owned();
         let ak = actual_key;
@@ -228,22 +254,23 @@ impl EncryptionService {
             nonce: B64.encode(nonce),
             encrypted_data_key,
             kms_key_id,
+            chunk_size: Some(chunk_size),
+            plaintext_size: Some(plaintext_size),
         })
     }
 
-    pub async fn decrypt_object(
+    async fn resolve_data_key(
         &self,
-        input_path: &Path,
-        output_path: &Path,
         enc_meta: &EncryptionMetadata,
         customer_key: Option<&[u8]>,
-    ) -> Result<(), CryptoError> {
+    ) -> Result<([u8; 32], [u8; 12]), CryptoError> {
         let nonce_bytes = B64
             .decode(&enc_meta.nonce)
             .map_err(|e| CryptoError::EncryptionFailed(format!("Bad nonce encoding: {}", e)))?;
         if nonce_bytes.len() != 12 {
             return Err(CryptoError::InvalidNonceSize(nonce_bytes.len()));
         }
+        let nonce: [u8; 12] = nonce_bytes.try_into().unwrap();
 
         let data_key: [u8; 32] = if let Some(ck) = customer_key {
             if ck.len() != 32 {
@@ -281,15 +308,62 @@ impl EncryptionService {
             self.unwrap_data_key(wrapped)?
         };
 
+        Ok((data_key, nonce))
+    }
+
+    pub async fn decrypt_object(
+        &self,
+        input_path: &Path,
+        output_path: &Path,
+        enc_meta: &EncryptionMetadata,
+        customer_key: Option<&[u8]>,
+    ) -> Result<(), CryptoError> {
+        let (data_key, nonce) = self.resolve_data_key(enc_meta, customer_key).await?;
+
         let ip = input_path.to_owned();
         let op = output_path.to_owned();
-        let nb: [u8; 12] = nonce_bytes.try_into().unwrap();
-        tokio::task::spawn_blocking(move || decrypt_stream_chunked(&ip, &op, &data_key, &nb))
+        tokio::task::spawn_blocking(move || decrypt_stream_chunked(&ip, &op, &data_key, &nonce))
             .await
             .map_err(|e| CryptoError::Io(std::io::Error::new(std::io::ErrorKind::Other, e)))??;
 
         Ok(())
     }
+
+    pub async fn decrypt_object_range(
+        &self,
+        input_path: &Path,
+        output_path: &Path,
+        enc_meta: &EncryptionMetadata,
+        customer_key: Option<&[u8]>,
+        plain_start: u64,
+        plain_end_inclusive: u64,
+    ) -> Result<u64, CryptoError> {
+        let chunk_size = enc_meta.chunk_size.ok_or_else(|| {
+            CryptoError::EncryptionFailed("chunk_size missing from encryption metadata".into())
+        })?;
+        let plaintext_size = enc_meta.plaintext_size.ok_or_else(|| {
+            CryptoError::EncryptionFailed("plaintext_size missing from encryption metadata".into())
+        })?;
+
+        let (data_key, nonce) = self.resolve_data_key(enc_meta, customer_key).await?;
+
+        let ip = input_path.to_owned();
+        let op = output_path.to_owned();
+        tokio::task::spawn_blocking(move || {
+            decrypt_stream_chunked_range(
+                &ip,
+                &op,
+                &data_key,
+                &nonce,
+                chunk_size,
+                plaintext_size,
+                plain_start,
+                plain_end_inclusive,
+            )
+        })
+        .await
+        .map_err(|e| CryptoError::Io(std::io::Error::new(std::io::ErrorKind::Other, e)))?
+    }
 }
 
 #[cfg(test)]
@@ -383,12 +457,26 @@ mod tests {
             nonce: "dGVzdG5vbmNlMTI=".to_string(),
             encrypted_data_key: Some("c29tZWtleQ==".to_string()),
             kms_key_id: None,
+            chunk_size: Some(65_536),
+            plaintext_size: Some(1_234_567),
         };
         let map = meta.to_metadata_map();
         let restored = EncryptionMetadata::from_metadata(&map).unwrap();
         assert_eq!(restored.algorithm, "AES256");
         assert_eq!(restored.nonce, meta.nonce);
         assert_eq!(restored.encrypted_data_key, meta.encrypted_data_key);
+        assert_eq!(restored.chunk_size, Some(65_536));
+        assert_eq!(restored.plaintext_size, Some(1_234_567));
+    }
+
+    #[test]
+    fn test_encryption_metadata_legacy_missing_sizes() {
+        let mut map = HashMap::new();
+        map.insert("x-amz-server-side-encryption".to_string(), "AES256".into());
+        map.insert("x-amz-encryption-nonce".to_string(), "aGVsbG8=".into());
+        let restored = EncryptionMetadata::from_metadata(&map).unwrap();
+        assert_eq!(restored.chunk_size, None);
+        assert_eq!(restored.plaintext_size, None);
     }
 
     #[test]