Skip to main content

coding_agent_search/pages/
key_management.rs

1//! Key management operations for encrypted pages archives.
2//!
3//! Provides CLI operations to manage key slots in an encrypted archive:
4//! - `list`: Show all key slots
5//! - `add`: Add a new password or recovery key slot
6//! - `revoke`: Remove a key slot
7//! - `rotate`: Full key rotation (regenerate DEK, re-encrypt payload)
8//!
9//! # Security Model
10//!
11//! The archive uses envelope encryption with multiple key slots (like LUKS):
12//! - A random Data Encryption Key (DEK) encrypts the payload
13//! - Each key slot wraps the DEK with a Key Encryption Key (KEK)
14//! - KEK is derived from password (Argon2id) or recovery secret (HKDF-SHA256)
15//! - Add/revoke only modifies config.json; payload unchanged
16//! - Rotate re-encrypts entire payload with new DEK
17
18use crate::pages::attachments::reencrypt_blobs_into_dir;
19use crate::pages::encrypt::{
20    Argon2Params, EncryptionConfig, KdfAlgorithm, KeySlot, SlotType, load_config,
21    validate_supported_payload_format,
22};
23use crate::pages::qr::RecoverySecret;
24use aes_gcm::{
25    Aes256Gcm, Nonce,
26    aead::{Aead, KeyInit, Payload},
27};
28use anyhow::{Context, Result, bail};
29use argon2::{Algorithm, Argon2, Params, Version};
30use base64::prelude::*;
31use chrono::{DateTime, Utc};
32use flate2::{Compression, read::DeflateDecoder, write::DeflateEncoder};
33use rand::Rng;
34use serde::Serialize;
35use std::fs::File;
36use std::io::{BufWriter, Read, Write};
37use std::path::Path;
38use tracing::info;
39
40/// Argon2id default parameters
41#[cfg(not(test))]
42const ARGON2_MEMORY_KB: u32 = 65536; // 64 MB
43#[cfg(test)]
44const ARGON2_MEMORY_KB: u32 = 64;
45#[cfg(not(test))]
46const ARGON2_ITERATIONS: u32 = 3;
47#[cfg(test)]
48const ARGON2_ITERATIONS: u32 = 1;
49#[cfg(not(test))]
50const ARGON2_PARALLELISM: u32 = 4;
51#[cfg(test)]
52const ARGON2_PARALLELISM: u32 = 1;
53
54/// Schema version for encryption
55const SCHEMA_VERSION: u8 = 2;
56const MAX_ARCHIVE_CHUNKS: u64 = u32::MAX as u64;
57
58fn max_encryptable_plaintext_bytes(chunk_size: usize) -> u64 {
59    MAX_ARCHIVE_CHUNKS.saturating_mul(chunk_size as u64)
60}
61
62fn ensure_archive_chunk_count_fits_nonce_space(chunk_count: u64, chunk_size: usize) -> Result<()> {
63    if chunk_count > MAX_ARCHIVE_CHUNKS {
64        bail!(
65            "File too large: exceeds maximum of {} chunks ({} bytes with current chunk size)",
66            u32::MAX,
67            max_encryptable_plaintext_bytes(chunk_size)
68        );
69    }
70    Ok(())
71}
72
73fn ensure_can_write_archive_chunk(chunk_index: u32, chunk_size: usize) -> Result<()> {
74    if chunk_index == u32::MAX {
75        bail!(
76            "File too large: exceeds maximum of {} chunks ({} bytes with current chunk size)",
77            u32::MAX,
78            max_encryptable_plaintext_bytes(chunk_size)
79        );
80    }
81    Ok(())
82}
83
84/// Result of listing key slots
85#[derive(Debug, Clone, Serialize)]
86pub struct KeyListResult {
87    pub slots: Vec<KeySlotInfo>,
88    pub active_slots: usize,
89    pub dek_created_at: Option<String>,
90    pub export_id: String,
91}
92
93/// Information about a single key slot
94#[derive(Debug, Clone, Serialize)]
95pub struct KeySlotInfo {
96    pub id: u8,
97    pub slot_type: String,
98    pub kdf: String,
99    #[serde(skip_serializing_if = "Option::is_none")]
100    pub kdf_params: Option<Argon2Params>,
101    #[serde(skip_serializing_if = "Option::is_none")]
102    pub label: Option<String>,
103}
104
105/// Result of adding a key slot
106#[derive(Debug)]
107pub enum AddKeyResult {
108    Password { slot_id: u8 },
109    Recovery { slot_id: u8, secret: RecoverySecret },
110}
111
112/// Result of revoking a key slot
113#[derive(Debug, Serialize)]
114pub struct RevokeResult {
115    pub revoked_slot_id: u8,
116    pub remaining_slots: usize,
117}
118
119/// Result of key rotation
120#[derive(Debug, Serialize)]
121pub struct RotateResult {
122    pub new_dek_created_at: DateTime<Utc>,
123    pub slot_count: usize,
124    #[serde(skip_serializing_if = "Option::is_none")]
125    pub recovery_secret: Option<String>,
126}
127
128/// List all key slots in an archive
129pub fn key_list(archive_dir: &Path) -> Result<KeyListResult> {
130    let archive_dir = super::resolve_site_dir(archive_dir)?;
131    let config = load_config(&archive_dir)?;
132
133    let slots: Vec<KeySlotInfo> = config
134        .key_slots
135        .iter()
136        .map(|slot| KeySlotInfo {
137            id: slot.id,
138            slot_type: match slot.slot_type {
139                SlotType::Password => "password".to_string(),
140                SlotType::Recovery => "recovery".to_string(),
141            },
142            kdf: match slot.kdf {
143                KdfAlgorithm::Argon2id => "argon2id".to_string(),
144                KdfAlgorithm::HkdfSha256 => "hkdf-sha256".to_string(),
145            },
146            kdf_params: slot.argon2_params.clone(),
147            label: None, // Labels stored in encrypted metadata (future)
148        })
149        .collect();
150
151    Ok(KeyListResult {
152        active_slots: slots.len(),
153        slots,
154        dek_created_at: None, // Would need to store in config
155        export_id: config.export_id,
156    })
157}
158
159/// Add a new password-based key slot
160pub fn key_add_password(
161    archive_dir: &Path,
162    current_password: &str,
163    new_password: &str,
164) -> Result<u8> {
165    let archive_dir = super::resolve_site_dir(archive_dir)?;
166    let config_path = archive_dir.join("config.json");
167    let mut config = load_config(&archive_dir)?;
168    validate_supported_payload_format(&config)?;
169
170    // Unlock with current password to get DEK
171    let dek = zeroize::Zeroizing::new(unwrap_dek_with_password(&config, current_password)?);
172
173    // Create new slot (use max ID + 1 since IDs are stable after revocation)
174    // If no slots exist, start at 0; otherwise use max + 1
175    let slot_id = next_key_slot_id(&config.key_slots)?;
176    let new_slot = create_password_slot(new_password, &dek, &config.export_id, slot_id)?;
177
178    config.key_slots.push(new_slot);
179
180    // Write updated config
181    write_json_pretty_atomically(&config_path, &config)?;
182
183    // Update integrity.json if present
184    let manifest = regenerate_integrity_manifest(&archive_dir)?;
185    refresh_private_artifacts(&archive_dir, &config, manifest.as_ref(), None, false)?;
186
187    info!(slot_id, "Added password key slot");
188    Ok(slot_id)
189}
190
191/// Add a new recovery secret key slot
192pub fn key_add_recovery(
193    archive_dir: &Path,
194    current_password: &str,
195) -> Result<(u8, RecoverySecret)> {
196    let archive_dir = super::resolve_site_dir(archive_dir)?;
197    let config_path = archive_dir.join("config.json");
198    let mut config = load_config(&archive_dir)?;
199    validate_supported_payload_format(&config)?;
200
201    // Unlock with current password to get DEK
202    let dek = zeroize::Zeroizing::new(unwrap_dek_with_password(&config, current_password)?);
203
204    // Generate recovery secret
205    let secret = RecoverySecret::generate();
206
207    // Create new slot (use max ID + 1 since IDs are stable after revocation)
208    // If no slots exist, start at 0; otherwise use max + 1
209    let slot_id = next_key_slot_id(&config.key_slots)?;
210    let new_slot = create_recovery_slot(secret.as_bytes(), &dek, &config.export_id, slot_id)?;
211
212    config.key_slots.push(new_slot);
213
214    // Write updated config
215    write_json_pretty_atomically(&config_path, &config)?;
216
217    // Update integrity.json if present
218    let manifest = regenerate_integrity_manifest(&archive_dir)?;
219    refresh_private_artifacts(
220        &archive_dir,
221        &config,
222        manifest.as_ref(),
223        Some(secret.as_bytes()),
224        false,
225    )?;
226
227    info!(slot_id, "Added recovery key slot");
228    Ok((slot_id, secret))
229}
230
231fn next_key_slot_id(key_slots: &[KeySlot]) -> Result<u8> {
232    match key_slots.iter().map(|s| s.id).max() {
233        Some(max_id) => max_id.checked_add(1).ok_or_else(|| {
234            anyhow::anyhow!("Cannot add more key slots: maximum slot ID (255) reached")
235        }),
236        None => Ok(0),
237    }
238}
239
240/// Revoke a key slot
241pub fn key_revoke(
242    archive_dir: &Path,
243    current_password: &str,
244    slot_id_to_revoke: u8,
245) -> Result<RevokeResult> {
246    let archive_dir = super::resolve_site_dir(archive_dir)?;
247    let config_path = archive_dir.join("config.json");
248    let mut config = load_config(&archive_dir)?;
249    validate_supported_payload_format(&config)?;
250
251    // Safety: Cannot revoke last slot
252    if config.key_slots.len() <= 1 {
253        anyhow::bail!("Cannot revoke the last remaining key slot. Add another key first.");
254    }
255
256    // Find which slot authenticates with this password
257    let (auth_slot_id, dek) = unwrap_dek_with_slot_id(&config, current_password)?;
258    let mut _dek = zeroize::Zeroizing::new(dek); // Zeroize on drop
259
260    // Verify they aren't trying to revoke the slot they are currently using
261    if auth_slot_id == slot_id_to_revoke {
262        bail!(
263            "Cannot revoke slot {} used for authentication. Use a different password.",
264            slot_id_to_revoke
265        );
266    }
267
268    // Verify slot exists
269    if !config.key_slots.iter().any(|s| s.id == slot_id_to_revoke) {
270        bail!("Slot {} not found", slot_id_to_revoke);
271    }
272
273    let revoked_slot_is_recovery = config
274        .key_slots
275        .iter()
276        .find(|s| s.id == slot_id_to_revoke)
277        .map(|s| s.slot_type == SlotType::Recovery)
278        .unwrap_or(false);
279
280    // Remove the slot (keeping IDs stable - they're part of the AAD binding)
281    config.key_slots.retain(|s| s.id != slot_id_to_revoke);
282
283    // Write updated config
284    write_json_pretty_atomically(&config_path, &config)?;
285
286    // Update integrity.json if present
287    let manifest = regenerate_integrity_manifest(&archive_dir)?;
288    let has_recovery_slot = config
289        .key_slots
290        .iter()
291        .any(|slot| slot.slot_type == SlotType::Recovery);
292    refresh_private_artifacts(
293        &archive_dir,
294        &config,
295        manifest.as_ref(),
296        None,
297        revoked_slot_is_recovery || !has_recovery_slot,
298    )?;
299
300    info!(slot_id = slot_id_to_revoke, "Revoked key slot");
301    Ok(RevokeResult {
302        revoked_slot_id: slot_id_to_revoke,
303        remaining_slots: config.key_slots.len(),
304    })
305}
306
307/// Full key rotation - regenerate DEK and re-encrypt payload
308pub fn key_rotate(
309    archive_dir: &Path,
310    old_password: &str,
311    new_password: &str,
312    keep_recovery: bool,
313    progress: impl Fn(f32),
314) -> Result<RotateResult> {
315    let archive_dir = super::resolve_site_dir(archive_dir)?;
316    let config = load_config(&archive_dir)?;
317    validate_supported_payload_format(&config)?;
318    let old_export_id_raw = BASE64_STANDARD.decode(&config.export_id)?;
319    let old_export_id: [u8; 16] = old_export_id_raw.as_slice().try_into().map_err(|err| {
320        // [coding_agent_session_search-htiim] Chain the underlying
321        // TryFromSliceError so a debug-mode error chain shows the
322        // exact conversion that failed in addition to the
323        // human-readable length mismatch.
324        anyhow::anyhow!(
325            "invalid export_id length: expected 16, got {}: {err}",
326            old_export_id_raw.len()
327        )
328    })?;
329
330    // 1. Decrypt payload with old password
331    let old_dek = zeroize::Zeroizing::new(unwrap_dek_with_password(&config, old_password)?);
332    let plaintext =
333        zeroize::Zeroizing::new(decrypt_all_chunks(&archive_dir, &old_dek, &config, |p| {
334            progress(p * 0.5)
335        })?);
336
337    // 2. Generate new DEK and export_id
338    let mut new_dek = zeroize::Zeroizing::new([0u8; 32]);
339    let mut new_export_id = [0u8; 16];
340    let mut new_base_nonce = [0u8; 12];
341    let mut rng = rand::rng();
342    rng.fill_bytes(new_dek.as_mut());
343    rng.fill_bytes(&mut new_export_id);
344    rng.fill_bytes(&mut new_base_nonce);
345
346    let staged_site_dir = unique_atomic_sidecar_path(&archive_dir, "rotate", "site");
347    copy_site_except_runtime_state(&archive_dir, &staged_site_dir)?;
348
349    // 3. Re-encrypt payload with new DEK into the staged site
350    let chunk_count = encrypt_all_chunks(
351        &plaintext,
352        &new_dek,
353        &new_export_id,
354        &new_base_nonce,
355        config.payload.chunk_size,
356        &staged_site_dir.join("payload"),
357        |p| progress(0.5 + p * 0.5),
358    )?;
359
360    reencrypt_blobs_into_dir(
361        &archive_dir,
362        &staged_site_dir,
363        &old_dek,
364        &old_export_id,
365        &new_dek,
366        &new_export_id,
367    )?;
368
369    // 4. Create new key slots
370    let mut new_slots = vec![create_password_slot(
371        new_password,
372        &new_dek,
373        &BASE64_STANDARD.encode(new_export_id),
374        0,
375    )?];
376
377    let mut recovery_secret_encoded: Option<String> = None;
378    let mut recovery_secret_bytes: Option<Vec<u8>> = None;
379    if keep_recovery {
380        let secret = RecoverySecret::generate();
381        new_slots.push(create_recovery_slot(
382            secret.as_bytes(),
383            &new_dek,
384            &BASE64_STANDARD.encode(new_export_id),
385            1,
386        )?);
387        recovery_secret_bytes = Some(secret.as_bytes().to_vec());
388        recovery_secret_encoded = Some(secret.encoded().to_string());
389    }
390
391    // 5. Write new config
392    let new_config = EncryptionConfig {
393        version: config.version,
394        export_id: BASE64_STANDARD.encode(new_export_id),
395        base_nonce: BASE64_STANDARD.encode(new_base_nonce),
396        compression: config.compression,
397        kdf_defaults: Argon2Params::default(),
398        payload: crate::pages::encrypt::PayloadMeta {
399            chunk_size: config.payload.chunk_size,
400            chunk_count,
401            total_compressed_size: 0, // Recalculated
402            total_plaintext_size: plaintext.len() as u64,
403            files: (0..chunk_count)
404                .map(|i| format!("payload/chunk-{:05}.bin", i))
405                .collect(),
406        },
407        key_slots: new_slots.clone(),
408    };
409
410    write_json_pretty(&staged_site_dir.join("config.json"), &new_config)?;
411
412    // 6. Regenerate integrity.json for the staged site, then swap atomically
413    let manifest = crate::pages::bundle::generate_integrity_manifest(&staged_site_dir)?;
414    write_json_pretty(&staged_site_dir.join("integrity.json"), &manifest)?;
415    sync_tree(&staged_site_dir)?;
416    replace_dir_from_temp(&staged_site_dir, &archive_dir)?;
417    refresh_private_artifacts(
418        &archive_dir,
419        &new_config,
420        Some(&manifest),
421        recovery_secret_bytes.as_deref(),
422        !keep_recovery,
423    )?;
424
425    Ok(RotateResult {
426        new_dek_created_at: chrono::Utc::now(),
427        slot_count: new_slots.len(),
428        recovery_secret: recovery_secret_encoded,
429    })
430}
431
432// ============================================================================
433// Helper functions
434// ============================================================================
435
436/// Unwrap DEK using password (tries all password slots)
437fn unwrap_dek_with_password(config: &EncryptionConfig, password: &str) -> Result<[u8; 32]> {
438    let export_id = BASE64_STANDARD.decode(&config.export_id)?;
439
440    for slot in &config.key_slots {
441        if slot.slot_type != SlotType::Password {
442            continue;
443        }
444
445        let salt = BASE64_STANDARD.decode(&slot.salt)?;
446        let wrapped_dek = BASE64_STANDARD.decode(&slot.wrapped_dek)?;
447        let nonce = BASE64_STANDARD.decode(&slot.nonce)?;
448
449        if let Ok(kek) = derive_kek_argon2id(password, &salt) {
450            let result = unwrap_key(&kek, &wrapped_dek, &nonce, &export_id, slot.id);
451            if let Ok(dek) = result {
452                return Ok(dek);
453            }
454        }
455    }
456
457    bail!("Invalid password or no matching key slot")
458}
459
460/// Unwrap DEK and return which slot was used
461fn unwrap_dek_with_slot_id(config: &EncryptionConfig, password: &str) -> Result<(u8, [u8; 32])> {
462    let export_id = BASE64_STANDARD.decode(&config.export_id)?;
463
464    for slot in &config.key_slots {
465        if slot.slot_type != SlotType::Password {
466            continue;
467        }
468
469        let salt = BASE64_STANDARD.decode(&slot.salt)?;
470        let wrapped_dek = BASE64_STANDARD.decode(&slot.wrapped_dek)?;
471        let nonce = BASE64_STANDARD.decode(&slot.nonce)?;
472
473        if let Ok(kek) = derive_kek_argon2id(password, &salt) {
474            let result = unwrap_key(&kek, &wrapped_dek, &nonce, &export_id, slot.id);
475            if let Ok(dek) = result {
476                return Ok((slot.id, dek));
477            }
478        }
479    }
480
481    bail!("Invalid password or no matching key slot")
482}
483
484/// Derive KEK from password using Argon2id
485fn derive_kek_argon2id(password: &str, salt: &[u8]) -> Result<zeroize::Zeroizing<[u8; 32]>> {
486    let params = Params::new(
487        ARGON2_MEMORY_KB,
488        ARGON2_ITERATIONS,
489        ARGON2_PARALLELISM,
490        Some(32),
491    )
492    .map_err(|e| anyhow::anyhow!("Invalid Argon2 parameters: {:?}", e))?;
493
494    let argon2 = Argon2::new(Algorithm::Argon2id, Version::V0x13, params);
495
496    let mut kek = zeroize::Zeroizing::new([0u8; 32]);
497    argon2
498        .hash_password_into(password.as_bytes(), salt, kek.as_mut())
499        .map_err(|e| anyhow::anyhow!("Argon2 derivation failed: {:?}", e))?;
500
501    Ok(kek)
502}
503
504/// Derive KEK from recovery secret using HKDF-SHA256
505fn derive_kek_hkdf(secret: &[u8], salt: &[u8]) -> Result<zeroize::Zeroizing<[u8; 32]>> {
506    let kek = crate::encryption::hkdf_extract_expand(secret, salt, b"cass-pages-kek-v2", 32)
507        .map_err(|e| anyhow::anyhow!("HKDF extract+expand failed for recovery secret KEK: {e}"))?;
508    let actual_len = kek.len();
509    let kek: [u8; 32] = kek.try_into().map_err(|_| {
510        // [coding_agent_session_search-htiim] Capture actual_len BEFORE
511        // try_into consumes the Vec so the message can show the actual
512        // KEK length operators / future contributors need to debug a
513        // frankensqlite / hkdf upstream regression.
514        anyhow::anyhow!(
515            "HKDF expansion produced invalid KEK length: expected 32, got {}",
516            actual_len
517        )
518    })?;
519    Ok(zeroize::Zeroizing::new(kek))
520}
521
522/// Unwrap DEK with KEK
523fn unwrap_key(
524    kek: &[u8; 32],
525    wrapped: &[u8],
526    nonce: &[u8],
527    export_id: &[u8],
528    slot_id: u8,
529) -> Result<[u8; 32]> {
530    let cipher = Aes256Gcm::new_from_slice(kek).expect("Invalid key length");
531
532    let actual_nonce_len = nonce.len();
533    let nonce: &[u8; 12] = nonce.try_into().map_err(|err| {
534        // [coding_agent_session_search-htiim] Chain TryFromSliceError so
535        // debug-mode chains preserve the source while the operator
536        // sees the exact-length diagnostic.
537        anyhow::anyhow!(
538            "invalid nonce length: expected 12, got {}: {err}",
539            actual_nonce_len
540        )
541    })?;
542
543    // AAD: export_id || slot_id
544    let mut aad = Vec::with_capacity(export_id.len() + 1);
545    aad.extend_from_slice(export_id);
546    aad.push(slot_id);
547
548    let dek = cipher
549        .decrypt(
550            Nonce::from_slice(nonce),
551            Payload {
552                msg: wrapped,
553                aad: &aad,
554            },
555        )
556        .map_err(|err| {
557            // [coding_agent_session_search-htiim] Chain the underlying
558            // aead error so operators can distinguish "wrong password
559            // (KEK derivation succeeded but DEK MAC failed)" from
560            // "corrupt key slot ciphertext" from "wrong AAD (slot id /
561            // export id mismatch)". The aead crate's Display impl
562            // remains opaque about the specific sub-failure (timing-
563            // attack hardening), but the source error type IS preserved
564            // so debug-mode error chains can show whether the failure
565            // came from the cipher layer vs a subsequent layer. Slot
566            // id is included so operators can correlate with the
567            // recovery / password slot they were attempting. Mirrors
568            // the encrypt.rs::unwrap_key fix landed in 0b81b601.
569            anyhow::anyhow!(
570                "Key unwrapping failed for slot {} ({} bytes wrapped, {} bytes nonce, \
571                 {} bytes aad): {}",
572                slot_id,
573                wrapped.len(),
574                actual_nonce_len,
575                aad.len(),
576                err
577            )
578        })?;
579
580    let dek_len = dek.len();
581    dek.try_into().map_err(|_| {
582        anyhow::anyhow!(
583            "Invalid DEK length after unwrap: expected 32, got {}",
584            dek_len
585        )
586    })
587}
588
589/// Create a password-based key slot
590fn create_password_slot(
591    password: &str,
592    dek: &[u8; 32],
593    export_id_b64: &str,
594    slot_id: u8,
595) -> Result<KeySlot> {
596    let export_id = BASE64_STANDARD.decode(export_id_b64)?;
597
598    // Generate salt
599    let mut salt = [0u8; 32];
600    let mut rng = rand::rng();
601    rng.fill_bytes(&mut salt);
602
603    // Derive KEK from password
604    let kek = derive_kek_argon2id(password, &salt)?;
605
606    // Wrap DEK
607    let result = wrap_key(&kek, dek, &export_id, slot_id);
608
609    let (wrapped_dek, nonce) = result?;
610
611    Ok(KeySlot {
612        id: slot_id,
613        slot_type: SlotType::Password,
614        kdf: KdfAlgorithm::Argon2id,
615        salt: BASE64_STANDARD.encode(salt),
616        wrapped_dek: BASE64_STANDARD.encode(&wrapped_dek),
617        nonce: BASE64_STANDARD.encode(nonce),
618        argon2_params: Some(Argon2Params::default()),
619    })
620}
621
622/// Create a recovery secret key slot
623fn create_recovery_slot(
624    secret: &[u8],
625    dek: &[u8; 32],
626    export_id_b64: &str,
627    slot_id: u8,
628) -> Result<KeySlot> {
629    let export_id = BASE64_STANDARD.decode(export_id_b64)?;
630
631    // Generate salt
632    let mut salt = [0u8; 16];
633    let mut rng = rand::rng();
634    rng.fill_bytes(&mut salt);
635
636    // Derive KEK from recovery secret
637    let kek = derive_kek_hkdf(secret, &salt)?;
638
639    // Wrap DEK
640    let result = wrap_key(&kek, dek, &export_id, slot_id);
641
642    let (wrapped_dek, nonce) = result?;
643
644    Ok(KeySlot {
645        id: slot_id,
646        slot_type: SlotType::Recovery,
647        kdf: KdfAlgorithm::HkdfSha256,
648        salt: BASE64_STANDARD.encode(salt),
649        wrapped_dek: BASE64_STANDARD.encode(&wrapped_dek),
650        nonce: BASE64_STANDARD.encode(nonce),
651        argon2_params: None,
652    })
653}
654
655/// Wrap DEK with KEK using AES-256-GCM
656fn wrap_key(
657    kek: &[u8; 32],
658    dek: &[u8; 32],
659    export_id: &[u8],
660    slot_id: u8,
661) -> Result<(Vec<u8>, [u8; 12])> {
662    let cipher = Aes256Gcm::new_from_slice(kek).expect("Invalid key length");
663
664    let mut nonce = [0u8; 12];
665    let mut rng = rand::rng();
666    rng.fill_bytes(&mut nonce);
667
668    // AAD: export_id || slot_id
669    let mut aad = Vec::with_capacity(export_id.len() + 1);
670    aad.extend_from_slice(export_id);
671    aad.push(slot_id);
672
673    let wrapped = cipher
674        .encrypt(
675            Nonce::from_slice(&nonce),
676            Payload {
677                msg: dek,
678                aad: &aad,
679            },
680        )
681        .map_err(|e| anyhow::anyhow!("Key wrapping failed: {}", e))?;
682
683    Ok((wrapped, nonce))
684}
685
686/// Decrypt all chunks and return plaintext
687fn decrypt_all_chunks(
688    archive_dir: &Path,
689    dek: &[u8; 32],
690    config: &EncryptionConfig,
691    progress: impl Fn(f32),
692) -> Result<Vec<u8>> {
693    let cipher = Aes256Gcm::new_from_slice(dek).expect("Invalid key length");
694    let base_nonce_raw = BASE64_STANDARD.decode(&config.base_nonce)?;
695    let base_nonce: [u8; 12] = base_nonce_raw.as_slice().try_into().map_err(|err| {
696        // [coding_agent_session_search-htiim] Chain TryFromSliceError so
697        // debug chains preserve the source. Length diagnostic is the
698        // operator-facing signal.
699        anyhow::anyhow!(
700            "invalid base_nonce length: expected 12, got {}: {err}",
701            base_nonce_raw.len()
702        )
703    })?;
704    let export_id_raw = BASE64_STANDARD.decode(&config.export_id)?;
705    let export_id: [u8; 16] = export_id_raw.as_slice().try_into().map_err(|err| {
706        // [coding_agent_session_search-htiim] Chain TryFromSliceError.
707        anyhow::anyhow!(
708            "invalid export_id length: expected 16, got {}: {err}",
709            export_id_raw.len()
710        )
711    })?;
712    let canonical_archive_dir = archive_dir.canonicalize().with_context(|| {
713        format!(
714            "Failed to resolve archive root {} before decrypting chunks",
715            archive_dir.display()
716        )
717    })?;
718
719    let mut plaintext = Vec::new();
720
721    if config.payload.chunk_count != config.payload.files.len() {
722        bail!(
723            "Invalid config: payload chunk_count {} does not match file list length {}",
724            config.payload.chunk_count,
725            config.payload.files.len()
726        );
727    }
728
729    for (chunk_index, chunk_file) in config.payload.files.iter().enumerate() {
730        progress(chunk_index as f32 / config.payload.chunk_count as f32);
731
732        let expected_chunk_file = format!("payload/chunk-{chunk_index:05}.bin");
733        if chunk_file != &expected_chunk_file {
734            bail!(
735                "Invalid chunk path in config.json: expected {}, got {}",
736                expected_chunk_file,
737                chunk_file
738            );
739        }
740        let chunk_path = archive_dir.join(chunk_file);
741        let chunk_meta = std::fs::symlink_metadata(&chunk_path).with_context(|| {
742            format!(
743                "Failed to inspect encrypted chunk {} at {}",
744                chunk_index,
745                chunk_path.display()
746            )
747        })?;
748        if chunk_meta.file_type().is_symlink() {
749            bail!("Encrypted chunk must not be a symlink: {}", chunk_file);
750        }
751        if !chunk_meta.file_type().is_file() {
752            bail!("Encrypted chunk must be a regular file: {}", chunk_file);
753        }
754
755        let canonical_chunk_path = chunk_path.canonicalize().with_context(|| {
756            format!(
757                "Failed to resolve encrypted chunk {} at {}",
758                chunk_index,
759                chunk_path.display()
760            )
761        })?;
762        if !canonical_chunk_path.starts_with(&canonical_archive_dir) {
763            bail!(
764                "Encrypted chunk path escapes archive directory: {}",
765                chunk_file
766            );
767        }
768
769        let ciphertext = std::fs::read(&canonical_chunk_path)?;
770
771        // Derive nonce
772        let nonce = derive_chunk_nonce(&base_nonce, chunk_index as u32);
773
774        // Build AAD
775        let aad = build_chunk_aad(&export_id, chunk_index as u32);
776
777        // Decrypt
778        let compressed = cipher
779            .decrypt(
780                Nonce::from_slice(&nonce),
781                Payload {
782                    msg: &ciphertext,
783                    aad: &aad,
784                },
785            )
786            .map_err(|err| {
787                // [coding_agent_session_search-htiim] Chain the aead error
788                // so operators can correlate: which chunk failed, how
789                // big the ciphertext was, and what the cipher layer
790                // reported. The aead crate keeps the sub-failure type
791                // opaque (timing-attack hardening) but the source is
792                // preserved in the error chain. Mirrors encrypt.rs::
793                // decrypt_all_chunks fix landed in 0b81b601.
794                anyhow::anyhow!(
795                    "Decryption failed for chunk {} ({} bytes ciphertext): {}",
796                    chunk_index,
797                    ciphertext.len(),
798                    err
799                )
800            })?;
801
802        // Decompress
803        let mut decoder = DeflateDecoder::new(&compressed[..]);
804        let mut chunk_plaintext = Vec::new();
805        decoder.read_to_end(&mut chunk_plaintext)?;
806
807        plaintext.extend(chunk_plaintext);
808    }
809
810    progress(1.0);
811    Ok(plaintext)
812}
813
814/// Encrypt plaintext into chunks
815fn encrypt_all_chunks(
816    plaintext: &[u8],
817    dek: &[u8; 32],
818    export_id: &[u8; 16],
819    base_nonce: &[u8; 12],
820    chunk_size: usize,
821    payload_dir: &Path,
822    progress: impl Fn(f32),
823) -> Result<usize> {
824    std::fs::create_dir_all(payload_dir)?;
825
826    let cipher = Aes256Gcm::new_from_slice(dek).expect("Invalid key length");
827    if chunk_size == 0 {
828        anyhow::bail!("chunk_size must be > 0");
829    }
830    let total_chunks = plaintext.len().div_ceil(chunk_size);
831    ensure_archive_chunk_count_fits_nonce_space(total_chunks as u64, chunk_size)?;
832    let mut chunk_index = 0u32;
833
834    for (i, chunk) in plaintext.chunks(chunk_size).enumerate() {
835        progress(i as f32 / total_chunks as f32);
836        ensure_can_write_archive_chunk(chunk_index, chunk_size)?;
837
838        // Compress
839        let mut compressed = Vec::new();
840        {
841            let mut encoder = DeflateEncoder::new(&mut compressed, Compression::default());
842            encoder.write_all(chunk)?;
843            encoder.finish()?;
844        }
845
846        // Derive nonce
847        let nonce = derive_chunk_nonce(base_nonce, chunk_index);
848
849        // Build AAD
850        let aad = build_chunk_aad(export_id, chunk_index);
851
852        // Encrypt
853        let ciphertext = cipher
854            .encrypt(
855                Nonce::from_slice(&nonce),
856                Payload {
857                    msg: &compressed,
858                    aad: &aad,
859                },
860            )
861            .map_err(|e| anyhow::anyhow!("Encryption failed: {}", e))?;
862
863        // Write chunk file
864        let chunk_filename = format!("chunk-{:05}.bin", chunk_index);
865        let chunk_path = payload_dir.join(&chunk_filename);
866        let mut chunk_file = File::create(&chunk_path)?;
867        chunk_file.write_all(&ciphertext)?;
868
869        chunk_index = chunk_index.checked_add(1).ok_or_else(|| {
870            anyhow::anyhow!(
871                "File too large: exceeds maximum of {} chunks ({} bytes with current chunk size)",
872                u32::MAX,
873                (u32::MAX as u64) * (chunk_size as u64)
874            )
875        })?;
876    }
877
878    progress(1.0);
879    Ok(chunk_index as usize)
880}
881
882/// Derive chunk nonce from base nonce and chunk index
883fn derive_chunk_nonce(base_nonce: &[u8; 12], chunk_index: u32) -> [u8; 12] {
884    let mut nonce = *base_nonce;
885    // Set the last 4 bytes to the chunk index (big-endian)
886    nonce[8..12].copy_from_slice(&chunk_index.to_be_bytes());
887    nonce
888}
889
890/// Build AAD for chunk encryption
891fn build_chunk_aad(export_id: &[u8; 16], chunk_index: u32) -> Vec<u8> {
892    let mut aad = Vec::with_capacity(21);
893    aad.extend_from_slice(export_id);
894    aad.extend_from_slice(&chunk_index.to_be_bytes());
895    aad.push(SCHEMA_VERSION);
896    aad
897}
898
899/// Regenerate entire integrity.json
900fn regenerate_integrity_manifest(
901    archive_dir: &Path,
902) -> Result<Option<crate::pages::bundle::IntegrityManifest>> {
903    let integrity_path = archive_dir.join("integrity.json");
904    if !integrity_path.exists() {
905        return Ok(None);
906    }
907
908    let integrity = crate::pages::bundle::generate_integrity_manifest(archive_dir)?;
909    write_json_pretty(&integrity_path, &integrity)?;
910
911    Ok(Some(integrity))
912}
913
914fn write_json_pretty_atomically<T: Serialize>(path: &Path, value: &T) -> Result<()> {
915    let temp_path = unique_atomic_temp_path(path);
916    {
917        let file = File::create(&temp_path)?;
918        let mut writer = BufWriter::new(file);
919        serde_json::to_writer_pretty(&mut writer, value)?;
920        writer.flush()?;
921        writer.get_ref().sync_all()?;
922    }
923    replace_file_from_temp(&temp_path, path)
924}
925
926fn write_json_pretty<T: Serialize>(path: &Path, value: &T) -> Result<()> {
927    let file = File::create(path)?;
928    let mut writer = BufWriter::new(file);
929    serde_json::to_writer_pretty(&mut writer, value)?;
930    writer.flush()?;
931    writer.get_ref().sync_all()?;
932    Ok(())
933}
934
935fn replace_file_from_temp(temp_path: &Path, final_path: &Path) -> Result<()> {
936    if cfg!(windows) {
937        match std::fs::rename(temp_path, final_path) {
938            Ok(()) => {
939                sync_parent_directory(final_path)?;
940                Ok(())
941            }
942            Err(first_err) if final_path.exists() => {
943                let backup_path = unique_atomic_backup_path(final_path);
944                std::fs::rename(final_path, &backup_path).map_err(|backup_err| {
945                    let _ = std::fs::remove_file(temp_path);
946                    anyhow::anyhow!(
947                        "failed replacing {} with {}: {}; failed moving existing file to backup {}: {}",
948                        final_path.display(),
949                        temp_path.display(),
950                        first_err,
951                        backup_path.display(),
952                        backup_err
953                    )
954                })?;
955
956                match std::fs::rename(temp_path, final_path) {
957                    Ok(()) => {
958                        let _ = std::fs::remove_file(&backup_path);
959                        sync_parent_directory(final_path)?;
960                        Ok(())
961                    }
962                    Err(second_err) => match std::fs::rename(&backup_path, final_path) {
963                        Ok(()) => {
964                            let _ = std::fs::remove_file(temp_path);
965                            sync_parent_directory(final_path)?;
966                            anyhow::bail!(
967                                "failed replacing {} with {}: {}; restored original file",
968                                final_path.display(),
969                                temp_path.display(),
970                                second_err
971                            );
972                        }
973                        Err(restore_err) => {
974                            anyhow::bail!(
975                                "failed replacing {} with {}: {}; restore error: {}; temp file retained at {}",
976                                final_path.display(),
977                                temp_path.display(),
978                                second_err,
979                                restore_err,
980                                temp_path.display()
981                            );
982                        }
983                    },
984                }
985            }
986            Err(err) => Err(err.into()),
987        }
988    } else {
989        std::fs::rename(temp_path, final_path)?;
990        sync_parent_directory(final_path)?;
991        Ok(())
992    }
993}
994
995#[cfg(not(windows))]
996fn sync_parent_directory(path: &Path) -> Result<()> {
997    let Some(parent) = path.parent() else {
998        return Ok(());
999    };
1000    std::fs::File::open(parent)?.sync_all()?;
1001    Ok(())
1002}
1003
1004#[cfg(windows)]
1005fn sync_parent_directory(_path: &Path) -> Result<()> {
1006    Ok(())
1007}
1008
1009fn unique_atomic_temp_path(path: &Path) -> std::path::PathBuf {
1010    unique_atomic_sidecar_path(path, "tmp", "config.json")
1011}
1012
1013fn unique_atomic_backup_path(path: &Path) -> std::path::PathBuf {
1014    unique_atomic_sidecar_path(path, "bak", "config.json")
1015}
1016
1017fn unique_atomic_sidecar_path(
1018    path: &Path,
1019    suffix: &str,
1020    fallback_name: &str,
1021) -> std::path::PathBuf {
1022    static NEXT_NONCE: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1023
1024    let timestamp = std::time::SystemTime::now()
1025        .duration_since(std::time::UNIX_EPOCH)
1026        .unwrap_or_default()
1027        .as_nanos();
1028    let nonce = NEXT_NONCE.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1029    let file_name = path
1030        .file_name()
1031        .and_then(|name| name.to_str())
1032        .unwrap_or(fallback_name);
1033
1034    path.with_file_name(format!(
1035        ".{file_name}.{suffix}.{}.{}.{}",
1036        std::process::id(),
1037        timestamp,
1038        nonce
1039    ))
1040}
1041
1042fn replace_dir_from_temp(temp_dir: &Path, final_dir: &Path) -> Result<()> {
1043    if !ensure_replaceable_site_dir(final_dir)? {
1044        std::fs::rename(temp_dir, final_dir).with_context(|| {
1045            format!(
1046                "failed renaming staged site {} into place at {}",
1047                temp_dir.display(),
1048                final_dir.display()
1049            )
1050        })?;
1051        sync_parent_directory(final_dir)?;
1052        return Ok(());
1053    }
1054
1055    let backup_dir = unique_atomic_sidecar_path(final_dir, "bak", "site");
1056    std::fs::rename(final_dir, &backup_dir).with_context(|| {
1057        format!(
1058            "failed preparing backup {} before replacing {}",
1059            backup_dir.display(),
1060            final_dir.display()
1061        )
1062    })?;
1063
1064    match std::fs::rename(temp_dir, final_dir) {
1065        Ok(()) => {
1066            sync_parent_directory(final_dir)?;
1067            let _ = std::fs::remove_dir_all(&backup_dir);
1068            sync_parent_directory(final_dir)?;
1069            Ok(())
1070        }
1071        Err(second_err) => match std::fs::rename(&backup_dir, final_dir) {
1072            Ok(()) => {
1073                let _ = std::fs::remove_dir_all(temp_dir);
1074                sync_parent_directory(final_dir)?;
1075                anyhow::bail!(
1076                    "failed replacing {} with {}: {}; restored original site",
1077                    final_dir.display(),
1078                    temp_dir.display(),
1079                    second_err
1080                )
1081            }
1082            Err(restore_err) => anyhow::bail!(
1083                "failed replacing {} with {}: {}; restore error: {}; staged site retained at {}",
1084                final_dir.display(),
1085                temp_dir.display(),
1086                second_err,
1087                restore_err,
1088                temp_dir.display()
1089            ),
1090        },
1091    }
1092}
1093
1094fn ensure_replaceable_site_dir(path: &Path) -> Result<bool> {
1095    match std::fs::symlink_metadata(path) {
1096        Ok(metadata) => {
1097            let file_type = metadata.file_type();
1098            if file_type.is_symlink() {
1099                bail!(
1100                    "Refusing to replace site directory through symlink: {}",
1101                    path.display()
1102                );
1103            }
1104            if !file_type.is_dir() {
1105                bail!(
1106                    "Refusing to replace site directory because it is not a directory: {}",
1107                    path.display()
1108                );
1109            }
1110            Ok(true)
1111        }
1112        Err(err) if err.kind() == std::io::ErrorKind::NotFound => Ok(false),
1113        Err(err) => Err(err).with_context(|| {
1114            format!(
1115                "Failed inspecting site directory before replacement: {}",
1116                path.display()
1117            )
1118        }),
1119    }
1120}
1121
1122#[cfg(not(windows))]
1123fn sync_tree(path: &Path) -> Result<()> {
1124    sync_tree_inner(path)?;
1125    sync_parent_directory(path)
1126}
1127
1128#[cfg(windows)]
1129fn sync_tree(_path: &Path) -> Result<()> {
1130    Ok(())
1131}
1132
1133#[cfg(not(windows))]
1134fn sync_tree_inner(path: &Path) -> Result<()> {
1135    let metadata = std::fs::symlink_metadata(path)
1136        .with_context(|| format!("Failed reading metadata for {}", path.display()))?;
1137    let file_type = metadata.file_type();
1138    if file_type.is_symlink() {
1139        return Ok(());
1140    }
1141    if file_type.is_file() {
1142        std::fs::File::open(path)
1143            .with_context(|| format!("Failed opening {} for sync", path.display()))?
1144            .sync_all()
1145            .with_context(|| format!("Failed syncing {}", path.display()))?;
1146        return Ok(());
1147    }
1148    if file_type.is_dir() {
1149        for entry in std::fs::read_dir(path)
1150            .with_context(|| format!("Failed reading directory {}", path.display()))?
1151        {
1152            let entry = entry.with_context(|| format!("Failed walking {}", path.display()))?;
1153            sync_tree_inner(&entry.path())?;
1154        }
1155        std::fs::File::open(path)
1156            .with_context(|| format!("Failed opening directory {} for sync", path.display()))?
1157            .sync_all()
1158            .with_context(|| format!("Failed syncing directory {}", path.display()))?;
1159    }
1160    Ok(())
1161}
1162
1163fn copy_site_except_runtime_state(src: &Path, dst: &Path) -> Result<()> {
1164    std::fs::create_dir_all(dst)
1165        .with_context(|| format!("Failed to create staged site directory {}", dst.display()))?;
1166    let canonical_base = src.canonicalize().with_context(|| {
1167        format!(
1168            "Failed to resolve archive root {} before staging key rotation",
1169            src.display()
1170        )
1171    })?;
1172    copy_site_except_runtime_state_recursive(src, dst, src, &canonical_base)
1173}
1174
1175fn copy_site_except_runtime_state_recursive(
1176    src: &Path,
1177    dst: &Path,
1178    base: &Path,
1179    canonical_base: &Path,
1180) -> Result<()> {
1181    for entry in std::fs::read_dir(src)? {
1182        let entry = entry?;
1183        let path = entry.path();
1184        let rel_path = path.strip_prefix(base)?;
1185        let skip_root_entry = rel_path.components().count() == 1
1186            && matches!(
1187                rel_path.to_str(),
1188                Some("payload" | "blobs" | "config.json" | "integrity.json")
1189            );
1190        if skip_root_entry {
1191            continue;
1192        }
1193
1194        let metadata = std::fs::symlink_metadata(&path)?;
1195        let file_type = metadata.file_type();
1196        let dest_path = dst.join(rel_path);
1197        if file_type.is_dir() {
1198            std::fs::create_dir_all(&dest_path)?;
1199            copy_site_except_runtime_state_recursive(&path, dst, base, canonical_base)?;
1200        } else if file_type.is_symlink() {
1201            let canonical_target = path.canonicalize().with_context(|| {
1202                format!(
1203                    "Failed to resolve symlinked site entry {} while staging key rotation",
1204                    rel_path.display()
1205                )
1206            })?;
1207            if !canonical_target.starts_with(canonical_base) {
1208                bail!(
1209                    "Refusing to rotate symlinked site entry outside archive root: {}",
1210                    rel_path.display()
1211                );
1212            }
1213
1214            let target_meta = std::fs::metadata(&path).with_context(|| {
1215                format!(
1216                    "Failed to read symlink target metadata for {} while staging key rotation",
1217                    rel_path.display()
1218                )
1219            })?;
1220            if !target_meta.is_file() {
1221                bail!(
1222                    "Refusing to rotate symlinked site entry that does not point to a regular file: {}",
1223                    rel_path.display()
1224                );
1225            }
1226
1227            if let Some(parent) = dest_path.parent() {
1228                std::fs::create_dir_all(parent)?;
1229            }
1230            // Materialize safe symlink targets into the staged site so the staged
1231            // integrity pass stays self-contained before the final atomic swap.
1232            std::fs::copy(&canonical_target, &dest_path).with_context(|| {
1233                format!(
1234                    "Failed copying symlink target {} into staged site path {}",
1235                    canonical_target.display(),
1236                    dest_path.display()
1237                )
1238            })?;
1239        } else if file_type.is_file() {
1240            if let Some(parent) = dest_path.parent() {
1241                std::fs::create_dir_all(parent)?;
1242            }
1243            std::fs::copy(&path, &dest_path).with_context(|| {
1244                format!(
1245                    "Failed copying staged site file {} to {}",
1246                    path.display(),
1247                    dest_path.display()
1248                )
1249            })?;
1250        }
1251    }
1252
1253    Ok(())
1254}
1255
1256fn refresh_private_artifacts(
1257    archive_dir: &Path,
1258    config: &EncryptionConfig,
1259    manifest: Option<&crate::pages::bundle::IntegrityManifest>,
1260    recovery_secret: Option<&[u8]>,
1261    remove_recovery_artifacts: bool,
1262) -> Result<()> {
1263    let Some(private_dir) = private_dir_for_archive(archive_dir)? else {
1264        return Ok(());
1265    };
1266
1267    if let Some(manifest) = manifest {
1268        let fingerprint = crate::pages::bundle::compute_fingerprint(manifest);
1269        crate::pages::bundle::write_private_fingerprint(&private_dir, &fingerprint)?;
1270    }
1271
1272    let should_generate_qr = recovery_secret.is_some()
1273        && (private_dir.join("qr-code.png").exists() || private_dir.join("qr-code.svg").exists());
1274
1275    crate::pages::bundle::write_private_artifacts_encrypted(
1276        &private_dir,
1277        config,
1278        recovery_secret,
1279        should_generate_qr,
1280        remove_recovery_artifacts,
1281    )?;
1282
1283    Ok(())
1284}
1285
1286fn private_dir_for_archive(archive_dir: &Path) -> Result<Option<std::path::PathBuf>> {
1287    if archive_dir
1288        .file_name()
1289        .map(|name| name == "site")
1290        .unwrap_or(false)
1291    {
1292        let Some(parent) = archive_dir.parent() else {
1293            return Ok(None);
1294        };
1295        let private_dir = parent.join("private");
1296        match std::fs::symlink_metadata(&private_dir) {
1297            Ok(metadata) => {
1298                let file_type = metadata.file_type();
1299                if file_type.is_symlink() {
1300                    bail!(
1301                        "private artifact directory must not be a symlink: {}",
1302                        private_dir.display()
1303                    );
1304                }
1305                if file_type.is_dir() {
1306                    return Ok(Some(private_dir));
1307                }
1308                bail!(
1309                    "private artifact path must be a directory: {}",
1310                    private_dir.display()
1311                );
1312            }
1313            Err(err) if err.kind() == std::io::ErrorKind::NotFound => {}
1314            Err(err) => {
1315                return Err(err).with_context(|| {
1316                    format!(
1317                        "Failed to inspect private artifact directory {}",
1318                        private_dir.display()
1319                    )
1320                });
1321            }
1322        }
1323    }
1324
1325    Ok(None)
1326}
1327
1328#[cfg(test)]
1329mod tests {
1330    use super::*;
1331    use crate::pages::attachments::{
1332        AttachmentConfig, AttachmentData, AttachmentProcessor, decrypt_blob, decrypt_manifest,
1333    };
1334    use crate::pages::bundle::BundleBuilder;
1335    use crate::pages::encrypt::{DecryptionEngine, EncryptionEngine, MAX_CHUNK_SIZE, PayloadMeta};
1336    use crate::pages::verify::verify_bundle;
1337    use std::cell::Cell;
1338    use tempfile::TempDir;
1339
1340    #[cfg(unix)]
1341    fn replace_viewer_with_in_tree_symlink(site_dir: &Path) {
1342        use std::os::unix::fs::symlink;
1343
1344        let real_viewer = site_dir.join("viewer-real.js");
1345        std::fs::rename(site_dir.join("viewer.js"), &real_viewer).unwrap();
1346        symlink("viewer-real.js", site_dir.join("viewer.js")).unwrap();
1347
1348        let manifest = crate::pages::bundle::generate_integrity_manifest(site_dir).unwrap();
1349        write_json_pretty(&site_dir.join("integrity.json"), &manifest).unwrap();
1350
1351        assert_eq!(verify_bundle(site_dir, false).unwrap().status, "valid");
1352    }
1353
1354    fn setup_test_archive() -> (TempDir, std::path::PathBuf) {
1355        let temp_dir = TempDir::new().unwrap();
1356        let input_path = temp_dir.path().join("input.txt");
1357        let bundle_root = temp_dir.path().join("bundle");
1358        let encrypted_dir = temp_dir.path().join("encrypted");
1359
1360        // Create test file
1361        std::fs::write(&input_path, b"Test data for key management").unwrap();
1362
1363        // Encrypt
1364        let mut engine = EncryptionEngine::new(1024).unwrap();
1365        engine.add_password_slot("test-password").unwrap();
1366        engine
1367            .encrypt_file(&input_path, &encrypted_dir, |_, _| {})
1368            .unwrap();
1369
1370        BundleBuilder::new()
1371            .build(&encrypted_dir, &bundle_root, |_, _| {})
1372            .unwrap();
1373
1374        (temp_dir, bundle_root)
1375    }
1376
1377    fn setup_test_archive_with_attachments() -> (TempDir, std::path::PathBuf) {
1378        let temp_dir = TempDir::new().unwrap();
1379        let input_path = temp_dir.path().join("input.txt");
1380        let bundle_root = temp_dir.path().join("bundle");
1381        let encrypted_dir = temp_dir.path().join("encrypted");
1382
1383        std::fs::write(&input_path, b"Test data for key management").unwrap();
1384
1385        let mut engine = EncryptionEngine::new(1024).unwrap();
1386        engine.add_password_slot("test-password").unwrap();
1387        engine
1388            .encrypt_file(&input_path, &encrypted_dir, |_, _| {})
1389            .unwrap();
1390
1391        let config = load_config(&encrypted_dir).unwrap();
1392        let dek = unwrap_dek_with_password(&config, "test-password").unwrap();
1393        let export_id_raw = BASE64_STANDARD.decode(&config.export_id).unwrap();
1394        let export_id: [u8; 16] = export_id_raw.as_slice().try_into().unwrap();
1395
1396        let mut processor = AttachmentProcessor::new(AttachmentConfig::enabled());
1397        processor
1398            .process_attachments(
1399                1,
1400                &[AttachmentData {
1401                    filename: "proof.txt".to_string(),
1402                    mime_type: "text/plain".to_string(),
1403                    data: b"attachment payload".to_vec(),
1404                }],
1405            )
1406            .unwrap();
1407        processor
1408            .write_encrypted_blobs(&encrypted_dir, &dek, &export_id)
1409            .unwrap();
1410
1411        BundleBuilder::new()
1412            .build(&encrypted_dir, &bundle_root, |_, _| {})
1413            .unwrap();
1414
1415        (temp_dir, bundle_root)
1416    }
1417
1418    fn rewrite_test_config(archive_dir: &Path, mutate: impl FnOnce(&mut EncryptionConfig)) {
1419        let site_dir = super::super::resolve_site_dir(archive_dir).unwrap();
1420        let mut config = load_config(&site_dir).unwrap();
1421        mutate(&mut config);
1422        write_json_pretty(&site_dir.join("config.json"), &config).unwrap();
1423    }
1424
1425    fn assert_unsupported_payload_format_error(err: anyhow::Error, compression: &str) {
1426        let rendered = err.to_string();
1427        assert!(
1428            rendered.contains("supports only deflate") && rendered.contains(compression),
1429            "unexpected unsupported-format error: {err:#}"
1430        );
1431    }
1432
1433    #[test]
1434    #[cfg(unix)]
1435    fn test_private_dir_for_archive_rejects_symlinked_private_dir() {
1436        use std::os::unix::fs::symlink;
1437
1438        let temp = TempDir::new().unwrap();
1439        let site_dir = temp.path().join("bundle/site");
1440        let outside_private = temp.path().join("outside-private");
1441        std::fs::create_dir_all(&site_dir).unwrap();
1442        std::fs::create_dir_all(&outside_private).unwrap();
1443        symlink(&outside_private, temp.path().join("bundle/private")).unwrap();
1444
1445        let err = private_dir_for_archive(&site_dir).unwrap_err();
1446
1447        assert!(
1448            err.to_string().contains("must not be a symlink"),
1449            "unexpected error: {err:#}"
1450        );
1451        assert!(
1452            std::fs::symlink_metadata(temp.path().join("bundle/private"))
1453                .unwrap()
1454                .file_type()
1455                .is_symlink(),
1456            "rejected private directory symlink should remain untouched"
1457        );
1458    }
1459
1460    #[test]
1461    fn test_private_dir_for_archive_rejects_non_directory_private_path() {
1462        let temp = TempDir::new().unwrap();
1463        let site_dir = temp.path().join("bundle/site");
1464        std::fs::create_dir_all(&site_dir).unwrap();
1465        std::fs::write(temp.path().join("bundle/private"), "not a directory").unwrap();
1466
1467        let err = private_dir_for_archive(&site_dir).unwrap_err();
1468
1469        assert!(
1470            err.to_string().contains("must be a directory"),
1471            "unexpected error: {err:#}"
1472        );
1473        assert_eq!(
1474            std::fs::read_to_string(temp.path().join("bundle/private")).unwrap(),
1475            "not a directory"
1476        );
1477    }
1478
1479    #[test]
1480    fn test_decrypt_all_chunks_rejects_mismatched_chunk_count_before_progress() {
1481        let temp_dir = TempDir::new().unwrap();
1482        let archive_dir = temp_dir.path();
1483        let config = EncryptionConfig {
1484            version: SCHEMA_VERSION,
1485            export_id: BASE64_STANDARD.encode([0u8; 16]),
1486            base_nonce: BASE64_STANDARD.encode([0u8; 12]),
1487            compression: "deflate".to_string(),
1488            kdf_defaults: Argon2Params::default(),
1489            payload: PayloadMeta {
1490                chunk_size: 1024,
1491                chunk_count: 0,
1492                total_compressed_size: 0,
1493                total_plaintext_size: 0,
1494                files: vec!["payload/chunk-00000.bin".to_string()],
1495            },
1496            key_slots: Vec::new(),
1497        };
1498        let progress_calls = Cell::new(0);
1499
1500        let err = decrypt_all_chunks(archive_dir, &[0u8; 32], &config, |progress| {
1501            assert!(progress.is_finite(), "progress must be finite: {progress}");
1502            progress_calls.set(progress_calls.get() + 1);
1503        })
1504        .unwrap_err();
1505
1506        assert!(
1507            err.to_string().contains("chunk_count 0"),
1508            "unexpected error: {err:#}"
1509        );
1510        assert_eq!(progress_calls.get(), 0);
1511    }
1512
1513    #[test]
1514    fn test_key_list() {
1515        let (_temp_dir, archive_dir) = setup_test_archive();
1516
1517        let result = key_list(&archive_dir).unwrap();
1518        assert_eq!(result.active_slots, 1);
1519        assert_eq!(result.slots.len(), 1);
1520        assert_eq!(result.slots[0].slot_type, "password");
1521        assert_eq!(result.slots[0].kdf, "argon2id");
1522    }
1523
1524    #[test]
1525    fn test_key_mutations_reject_unsupported_payload_compression() {
1526        let (_temp_dir, archive_dir) = setup_test_archive();
1527        key_add_password(&archive_dir, "test-password", "second-password").unwrap();
1528        rewrite_test_config(&archive_dir, |config| {
1529            config.compression = "zstd".to_string();
1530        });
1531
1532        let err = key_add_password(&archive_dir, "test-password", "third-password").unwrap_err();
1533        assert_unsupported_payload_format_error(err, "zstd");
1534
1535        let err = key_add_recovery(&archive_dir, "test-password").unwrap_err();
1536        assert_unsupported_payload_format_error(err, "zstd");
1537
1538        let err = key_revoke(&archive_dir, "second-password", 0).unwrap_err();
1539        assert_unsupported_payload_format_error(err, "zstd");
1540
1541        let err =
1542            key_rotate(&archive_dir, "test-password", "new-password", false, |_| {}).unwrap_err();
1543        assert_unsupported_payload_format_error(err, "zstd");
1544
1545        let config = load_config(&archive_dir).unwrap();
1546        assert_eq!(config.key_slots.len(), 2);
1547        assert!(unwrap_dek_with_password(&config, "test-password").is_ok());
1548        assert!(unwrap_dek_with_password(&config, "second-password").is_ok());
1549        assert!(unwrap_dek_with_password(&config, "third-password").is_err());
1550        assert!(unwrap_dek_with_password(&config, "new-password").is_err());
1551    }
1552
1553    #[test]
1554    fn test_key_rotate_rejects_oversized_payload_chunk_size_before_rewriting() {
1555        let (_temp_dir, archive_dir) = setup_test_archive();
1556        rewrite_test_config(&archive_dir, |config| {
1557            config.payload.chunk_size = MAX_CHUNK_SIZE + 1;
1558        });
1559
1560        let err =
1561            key_rotate(&archive_dir, "test-password", "new-password", false, |_| {}).unwrap_err();
1562        let rendered = err.to_string();
1563        assert!(
1564            rendered.contains("chunk_size") && rendered.contains("must be <="),
1565            "unexpected chunk-size error: {err:#}"
1566        );
1567
1568        let config = load_config(&archive_dir).unwrap();
1569        assert!(unwrap_dek_with_password(&config, "test-password").is_ok());
1570        assert!(unwrap_dek_with_password(&config, "new-password").is_err());
1571    }
1572
1573    #[test]
1574    fn test_key_rotate_chunk_count_preflight_preserves_nonce_space_limit() {
1575        ensure_archive_chunk_count_fits_nonce_space(u64::from(u32::MAX), 1).unwrap();
1576
1577        let err =
1578            ensure_archive_chunk_count_fits_nonce_space(u64::from(u32::MAX) + 1, 1).unwrap_err();
1579        let rendered = err.to_string();
1580        assert!(
1581            rendered.contains("exceeds maximum") && rendered.contains(&u32::MAX.to_string()),
1582            "unexpected chunk-count error: {rendered}"
1583        );
1584    }
1585
1586    #[test]
1587    fn test_key_add_password() {
1588        let (_temp_dir, archive_dir) = setup_test_archive();
1589
1590        // Add new password
1591        let slot_id = key_add_password(&archive_dir, "test-password", "new-password").unwrap();
1592        assert_eq!(slot_id, 1);
1593
1594        // Verify it was added
1595        let result = key_list(&archive_dir).unwrap();
1596        assert_eq!(result.active_slots, 2);
1597
1598        // Verify new password works
1599        let config = load_config(&archive_dir).unwrap();
1600        let dek = unwrap_dek_with_password(&config, "new-password").unwrap();
1601        assert!(!dek.iter().all(|&b| b == 0));
1602    }
1603
1604    #[test]
1605    fn test_key_add_recovery() {
1606        let (_temp_dir, archive_dir) = setup_test_archive();
1607
1608        // Add recovery slot
1609        let (slot_id, secret) = key_add_recovery(&archive_dir, "test-password").unwrap();
1610        assert_eq!(slot_id, 1);
1611        assert_eq!(secret.entropy_bits(), 256);
1612
1613        // Verify it was added
1614        let result = key_list(&archive_dir).unwrap();
1615        assert_eq!(result.active_slots, 2);
1616        assert_eq!(result.slots[1].slot_type, "recovery");
1617        assert_eq!(result.slots[1].kdf, "hkdf-sha256");
1618    }
1619
1620    #[test]
1621    fn test_key_add_wrong_password_fails() {
1622        let (_temp_dir, archive_dir) = setup_test_archive();
1623
1624        let result = key_add_password(&archive_dir, "wrong-password", "new-password");
1625        assert!(result.is_err());
1626    }
1627
1628    #[test]
1629    fn test_key_revoke() {
1630        let (_temp_dir, archive_dir) = setup_test_archive();
1631
1632        // Add second slot
1633        key_add_password(&archive_dir, "test-password", "second-password").unwrap();
1634
1635        // Revoke first slot using second password
1636        let result = key_revoke(&archive_dir, "second-password", 0).unwrap();
1637        assert_eq!(result.revoked_slot_id, 0);
1638        assert_eq!(result.remaining_slots, 1);
1639
1640        // Old password should no longer work
1641        let config = load_config(&archive_dir).unwrap();
1642        assert!(unwrap_dek_with_password(&config, "test-password").is_err());
1643
1644        // Second password should still work
1645        assert!(unwrap_dek_with_password(&config, "second-password").is_ok());
1646    }
1647
1648    #[test]
1649    fn test_key_revoke_last_slot_fails() {
1650        let (_temp_dir, archive_dir) = setup_test_archive();
1651
1652        let result = key_revoke(&archive_dir, "test-password", 0);
1653        assert!(result.is_err());
1654        assert!(result.unwrap_err().to_string().contains("last remaining"));
1655    }
1656
1657    #[test]
1658    fn test_key_revoke_auth_slot_fails() {
1659        let (_temp_dir, archive_dir) = setup_test_archive();
1660
1661        // Add second slot
1662        key_add_password(&archive_dir, "test-password", "second-password").unwrap();
1663
1664        // Try to revoke slot 0 using slot 0's password
1665        let result = key_revoke(&archive_dir, "test-password", 0);
1666        assert!(result.is_err());
1667        assert!(result.unwrap_err().to_string().contains("authentication"));
1668    }
1669
1670    #[test]
1671    fn test_key_rotate() {
1672        let (temp_dir, archive_dir) = setup_test_archive();
1673        let decrypted_path = temp_dir.path().join("decrypted.txt");
1674
1675        // Rotate keys
1676        let result =
1677            key_rotate(&archive_dir, "test-password", "new-password", false, |_| {}).unwrap();
1678        assert_eq!(result.slot_count, 1);
1679        assert!(result.recovery_secret.is_none());
1680
1681        // Old password should fail
1682        let config = load_config(&archive_dir).unwrap();
1683        assert!(unwrap_dek_with_password(&config, "test-password").is_err());
1684
1685        // New password should work and decrypt correctly
1686        let decryptor = DecryptionEngine::unlock_with_password(config, "new-password").unwrap();
1687        decryptor
1688            .decrypt_to_file(&archive_dir, &decrypted_path, |_, _| {})
1689            .unwrap();
1690
1691        let decrypted = std::fs::read(&decrypted_path).unwrap();
1692        assert_eq!(decrypted, b"Test data for key management");
1693    }
1694
1695    #[test]
1696    fn test_key_rotate_with_recovery() {
1697        let (_temp_dir, archive_dir) = setup_test_archive();
1698
1699        // Rotate keys with recovery
1700        let result =
1701            key_rotate(&archive_dir, "test-password", "new-password", true, |_| {}).unwrap();
1702        assert_eq!(result.slot_count, 2);
1703        assert!(result.recovery_secret.is_some());
1704
1705        // Verify recovery slot
1706        let list_result = key_list(&archive_dir).unwrap();
1707        assert_eq!(list_result.slots.len(), 2);
1708        assert_eq!(list_result.slots[0].slot_type, "password");
1709        assert_eq!(list_result.slots[1].slot_type, "recovery");
1710    }
1711
1712    #[test]
1713    fn test_key_add_after_revoke_no_id_collision() {
1714        let (_temp_dir, archive_dir) = setup_test_archive();
1715
1716        // Add slots 1 and 2
1717        key_add_password(&archive_dir, "test-password", "password-1").unwrap();
1718        key_add_password(&archive_dir, "test-password", "password-2").unwrap();
1719
1720        // Now have slots [0, 1, 2]
1721        let list = key_list(&archive_dir).unwrap();
1722        assert_eq!(list.slots.len(), 3);
1723
1724        // Revoke slot 1 using slot 2's password
1725        key_revoke(&archive_dir, "password-2", 1).unwrap();
1726
1727        // Now have slots [0, 2] (gap at 1)
1728        let list = key_list(&archive_dir).unwrap();
1729        assert_eq!(list.slots.len(), 2);
1730        let ids: Vec<u8> = list.slots.iter().map(|s| s.id).collect();
1731        assert_eq!(ids, vec![0, 2]);
1732
1733        // Add new slot - should get ID 3, not 2
1734        let new_id = key_add_password(&archive_dir, "test-password", "password-3").unwrap();
1735        assert_eq!(new_id, 3, "New slot should get max_id + 1, not len()");
1736
1737        // Verify all passwords still work
1738        let config = load_config(&archive_dir).unwrap();
1739        assert!(unwrap_dek_with_password(&config, "test-password").is_ok());
1740        assert!(unwrap_dek_with_password(&config, "password-1").is_err()); // Revoked
1741        assert!(unwrap_dek_with_password(&config, "password-2").is_ok());
1742        assert!(unwrap_dek_with_password(&config, "password-3").is_ok());
1743    }
1744
1745    #[test]
1746    fn test_next_key_slot_id_rejects_max_id() {
1747        let (_temp_dir, archive_dir) = setup_test_archive();
1748        let mut config = load_config(&archive_dir).unwrap();
1749        config.key_slots[0].id = u8::MAX;
1750
1751        let err = next_key_slot_id(&config.key_slots).unwrap_err();
1752
1753        assert_eq!(
1754            err.to_string(),
1755            "Cannot add more key slots: maximum slot ID (255) reached"
1756        );
1757    }
1758
1759    #[test]
1760    fn test_key_add_password_preserves_valid_integrity_manifest() {
1761        let (_temp_dir, archive_dir) = setup_test_archive();
1762
1763        assert_eq!(verify_bundle(&archive_dir, false).unwrap().status, "valid");
1764
1765        key_add_password(&archive_dir, "test-password", "new-password").unwrap();
1766
1767        assert_eq!(verify_bundle(&archive_dir, false).unwrap().status, "valid");
1768    }
1769
1770    #[test]
1771    fn test_key_rotate_preserves_valid_integrity_manifest() {
1772        let (_temp_dir, archive_dir) = setup_test_archive();
1773
1774        assert_eq!(verify_bundle(&archive_dir, false).unwrap().status, "valid");
1775
1776        key_rotate(&archive_dir, "test-password", "new-password", true, |_| {}).unwrap();
1777
1778        assert_eq!(verify_bundle(&archive_dir, false).unwrap().status, "valid");
1779    }
1780
1781    #[test]
1782    #[cfg(unix)]
1783    fn test_key_add_password_preserves_in_tree_symlinked_required_asset() {
1784        let (_temp_dir, archive_dir) = setup_test_archive();
1785        let site_dir = super::super::resolve_site_dir(&archive_dir).unwrap();
1786        replace_viewer_with_in_tree_symlink(&site_dir);
1787
1788        key_add_password(&archive_dir, "test-password", "new-password").unwrap();
1789
1790        assert_eq!(verify_bundle(&archive_dir, false).unwrap().status, "valid");
1791        assert!(
1792            std::fs::symlink_metadata(site_dir.join("viewer.js"))
1793                .unwrap()
1794                .file_type()
1795                .is_symlink()
1796        );
1797    }
1798
1799    #[test]
1800    #[cfg(unix)]
1801    fn test_key_rotate_materializes_in_tree_symlinked_required_asset() {
1802        let (_temp_dir, archive_dir) = setup_test_archive();
1803        let site_dir = super::super::resolve_site_dir(&archive_dir).unwrap();
1804        replace_viewer_with_in_tree_symlink(&site_dir);
1805        let expected_viewer = std::fs::read(site_dir.join("viewer-real.js")).unwrap();
1806
1807        key_rotate(&archive_dir, "test-password", "new-password", true, |_| {}).unwrap();
1808
1809        let viewer_metadata = std::fs::symlink_metadata(site_dir.join("viewer.js")).unwrap();
1810        assert!(viewer_metadata.file_type().is_file());
1811        assert!(!viewer_metadata.file_type().is_symlink());
1812        assert_eq!(
1813            std::fs::read(site_dir.join("viewer.js")).unwrap(),
1814            expected_viewer
1815        );
1816        assert_eq!(verify_bundle(&archive_dir, false).unwrap().status, "valid");
1817    }
1818
1819    #[test]
1820    #[cfg(unix)]
1821    fn test_key_rotate_rejects_payload_directory_symlink_escape() {
1822        use std::os::unix::fs::symlink;
1823
1824        let (temp_dir, archive_dir) = setup_test_archive();
1825        let site_dir = super::super::resolve_site_dir(&archive_dir).unwrap();
1826        let payload_dir = site_dir.join("payload");
1827        let outside_payload_dir = temp_dir.path().join("outside-payload");
1828
1829        std::fs::rename(&payload_dir, &outside_payload_dir).unwrap();
1830        symlink(&outside_payload_dir, &payload_dir).unwrap();
1831
1832        let err =
1833            key_rotate(&archive_dir, "test-password", "new-password", false, |_| {}).unwrap_err();
1834        assert!(
1835            err.to_string().contains("escapes archive directory"),
1836            "unexpected error: {err:#}"
1837        );
1838    }
1839
1840    #[test]
1841    fn test_key_add_password_updates_private_fingerprint_and_master_key() {
1842        let (_temp_dir, archive_dir) = setup_test_archive();
1843        let site_dir = super::super::resolve_site_dir(&archive_dir).unwrap();
1844        let private_dir = site_dir.parent().unwrap().join("private");
1845
1846        let old_fingerprint =
1847            std::fs::read_to_string(private_dir.join("integrity-fingerprint.txt")).unwrap();
1848        let old_master_key = std::fs::read_to_string(private_dir.join("master-key.json")).unwrap();
1849
1850        key_add_password(&archive_dir, "test-password", "new-password").unwrap();
1851
1852        let new_fingerprint =
1853            std::fs::read_to_string(private_dir.join("integrity-fingerprint.txt")).unwrap();
1854        let new_master_key = std::fs::read_to_string(private_dir.join("master-key.json")).unwrap();
1855
1856        assert_ne!(old_fingerprint, new_fingerprint);
1857        assert_ne!(old_master_key, new_master_key);
1858    }
1859
1860    #[test]
1861    fn test_key_add_recovery_writes_private_recovery_artifact() {
1862        let (_temp_dir, archive_dir) = setup_test_archive();
1863        let site_dir = super::super::resolve_site_dir(&archive_dir).unwrap();
1864        let private_dir = site_dir.parent().unwrap().join("private");
1865
1866        assert!(!private_dir.join("recovery-secret.txt").exists());
1867
1868        let (_slot_id, secret) = key_add_recovery(&archive_dir, "test-password").unwrap();
1869        let recovery_file =
1870            std::fs::read_to_string(private_dir.join("recovery-secret.txt")).unwrap();
1871
1872        assert!(recovery_file.contains(secret.encoded()));
1873    }
1874
1875    #[test]
1876    fn test_key_revoke_recovery_removes_private_recovery_artifact() {
1877        let (_temp_dir, archive_dir) = setup_test_archive();
1878        let site_dir = super::super::resolve_site_dir(&archive_dir).unwrap();
1879        let private_dir = site_dir.parent().unwrap().join("private");
1880
1881        let (recovery_slot_id, _secret) = key_add_recovery(&archive_dir, "test-password").unwrap();
1882        key_add_password(&archive_dir, "test-password", "second-password").unwrap();
1883        assert!(private_dir.join("recovery-secret.txt").exists());
1884
1885        key_revoke(&archive_dir, "second-password", recovery_slot_id).unwrap();
1886
1887        assert!(!private_dir.join("recovery-secret.txt").exists());
1888    }
1889
1890    #[test]
1891    fn test_key_revoke_one_of_multiple_recovery_slots_removes_stale_private_recovery_artifact() {
1892        let (_temp_dir, archive_dir) = setup_test_archive();
1893        let site_dir = super::super::resolve_site_dir(&archive_dir).unwrap();
1894        let private_dir = site_dir.parent().unwrap().join("private");
1895
1896        let (first_recovery_slot_id, first_secret) =
1897            key_add_recovery(&archive_dir, "test-password").unwrap();
1898        let (second_recovery_slot_id, second_secret) =
1899            key_add_recovery(&archive_dir, "test-password").unwrap();
1900
1901        let recovery_file_before =
1902            std::fs::read_to_string(private_dir.join("recovery-secret.txt")).unwrap();
1903        assert!(recovery_file_before.contains(second_secret.encoded()));
1904
1905        key_revoke(&archive_dir, "test-password", second_recovery_slot_id).unwrap();
1906
1907        assert!(!private_dir.join("recovery-secret.txt").exists());
1908
1909        let config = load_config(&archive_dir).unwrap();
1910        assert!(DecryptionEngine::unlock_with_recovery(config, first_secret.as_bytes()).is_ok());
1911
1912        assert_ne!(first_recovery_slot_id, second_recovery_slot_id);
1913    }
1914
1915    #[test]
1916    fn test_key_rotate_refreshes_private_recovery_and_master_key() {
1917        let (_temp_dir, archive_dir) = setup_test_archive();
1918        let site_dir = super::super::resolve_site_dir(&archive_dir).unwrap();
1919        let private_dir = site_dir.parent().unwrap().join("private");
1920
1921        let old_master_key = std::fs::read_to_string(private_dir.join("master-key.json")).unwrap();
1922        let result =
1923            key_rotate(&archive_dir, "test-password", "new-password", true, |_| {}).unwrap();
1924
1925        let new_master_key = std::fs::read_to_string(private_dir.join("master-key.json")).unwrap();
1926        let recovery_file =
1927            std::fs::read_to_string(private_dir.join("recovery-secret.txt")).unwrap();
1928
1929        assert_ne!(old_master_key, new_master_key);
1930        assert!(recovery_file.contains(result.recovery_secret.as_deref().unwrap()));
1931    }
1932
1933    #[test]
1934    fn test_key_rotate_without_recovery_removes_stale_private_recovery_artifact() {
1935        let (_temp_dir, archive_dir) = setup_test_archive();
1936        let site_dir = super::super::resolve_site_dir(&archive_dir).unwrap();
1937        let private_dir = site_dir.parent().unwrap().join("private");
1938
1939        let (_slot_id, _secret) = key_add_recovery(&archive_dir, "test-password").unwrap();
1940        assert!(private_dir.join("recovery-secret.txt").exists());
1941
1942        key_rotate(&archive_dir, "test-password", "new-password", false, |_| {}).unwrap();
1943
1944        assert!(!private_dir.join("recovery-secret.txt").exists());
1945        assert!(!private_dir.join("qr-code.png").exists());
1946        assert!(!private_dir.join("qr-code.svg").exists());
1947    }
1948
1949    #[test]
1950    fn test_key_rotate_reencrypts_attachment_blobs() {
1951        let (_temp_dir, archive_dir) = setup_test_archive_with_attachments();
1952
1953        assert_eq!(verify_bundle(&archive_dir, false).unwrap().status, "valid");
1954
1955        key_rotate(&archive_dir, "test-password", "new-password", false, |_| {}).unwrap();
1956
1957        let site_dir = super::super::resolve_site_dir(&archive_dir).unwrap();
1958        let config = load_config(&archive_dir).unwrap();
1959        let dek = unwrap_dek_with_password(&config, "new-password").unwrap();
1960        let export_id_raw = BASE64_STANDARD.decode(&config.export_id).unwrap();
1961        let export_id: [u8; 16] = export_id_raw.as_slice().try_into().unwrap();
1962
1963        let manifest_ciphertext =
1964            std::fs::read(site_dir.join("blobs").join("manifest.enc")).unwrap();
1965        let manifest = decrypt_manifest(&manifest_ciphertext, &dek, &export_id).unwrap();
1966        assert_eq!(manifest.entries.len(), 1);
1967        assert_eq!(manifest.entries[0].filename, "proof.txt");
1968
1969        let blob_ciphertext = std::fs::read(
1970            site_dir
1971                .join("blobs")
1972                .join(format!("{}.bin", manifest.entries[0].hash)),
1973        )
1974        .unwrap();
1975        let plaintext = decrypt_blob(
1976            &blob_ciphertext,
1977            &dek,
1978            &export_id,
1979            &manifest.entries[0].hash,
1980        )
1981        .unwrap();
1982        assert_eq!(plaintext, b"attachment payload");
1983        assert_eq!(verify_bundle(&archive_dir, false).unwrap().status, "valid");
1984    }
1985
1986    #[test]
1987    fn test_key_rotate_failure_before_site_swap_preserves_live_archive() {
1988        let (temp_dir, archive_dir) = setup_test_archive_with_attachments();
1989        let decrypted_path = temp_dir.path().join("decrypted-after-failure.txt");
1990        let site_dir = super::super::resolve_site_dir(&archive_dir).unwrap();
1991
1992        std::fs::write(site_dir.join("blobs").join("manifest.enc"), b"corrupted").unwrap();
1993
1994        let rotate_result =
1995            key_rotate(&archive_dir, "test-password", "new-password", false, |_| {});
1996        assert!(rotate_result.is_err());
1997
1998        let config = load_config(&archive_dir).unwrap();
1999        assert!(unwrap_dek_with_password(&config, "new-password").is_err());
2000
2001        let decryptor = DecryptionEngine::unlock_with_password(config, "test-password").unwrap();
2002        decryptor
2003            .decrypt_to_file(&archive_dir, &decrypted_path, |_, _| {})
2004            .unwrap();
2005
2006        let decrypted = std::fs::read(&decrypted_path).unwrap();
2007        assert_eq!(decrypted, b"Test data for key management");
2008    }
2009
2010    #[test]
2011    fn test_write_json_pretty_atomically_overwrites_existing_file() {
2012        let temp_dir = TempDir::new().unwrap();
2013        let path = temp_dir.path().join("config.json");
2014        std::fs::write(&path, "{\"before\":true}\n").unwrap();
2015
2016        let value = serde_json::json!({ "after": true });
2017        write_json_pretty_atomically(&path, &value).unwrap();
2018
2019        let written: serde_json::Value =
2020            serde_json::from_str(&std::fs::read_to_string(&path).unwrap()).unwrap();
2021        assert_eq!(written, value);
2022    }
2023
2024    #[test]
2025    fn test_replace_dir_from_temp_overwrites_existing_site() {
2026        let temp_dir = TempDir::new().unwrap();
2027        let final_dir = temp_dir.path().join("archive");
2028        let staged_dir = temp_dir.path().join("archive.staged");
2029
2030        std::fs::create_dir_all(final_dir.join("site")).unwrap();
2031        std::fs::write(final_dir.join("site/old.txt"), "old").unwrap();
2032
2033        std::fs::create_dir_all(staged_dir.join("site")).unwrap();
2034        std::fs::write(staged_dir.join("site/new.txt"), "new").unwrap();
2035
2036        replace_dir_from_temp(&staged_dir, &final_dir).unwrap();
2037
2038        assert!(!staged_dir.exists());
2039        assert!(final_dir.join("site/new.txt").exists());
2040        assert!(!final_dir.join("site/old.txt").exists());
2041        let sidecars = std::fs::read_dir(temp_dir.path())
2042            .unwrap()
2043            .map(|entry| entry.unwrap().file_name().to_string_lossy().into_owned())
2044            .collect::<Vec<_>>();
2045        assert!(
2046            !sidecars.iter().any(|name| name.contains(".archive.bak.")),
2047            "backup sidecar should be cleaned up, found: {sidecars:?}"
2048        );
2049    }
2050
2051    #[test]
2052    fn test_replace_dir_from_temp_rejects_file_target() {
2053        let temp_dir = TempDir::new().unwrap();
2054        let final_dir = temp_dir.path().join("archive");
2055        let staged_dir = temp_dir.path().join("archive.staged");
2056
2057        std::fs::write(&final_dir, "not a directory").unwrap();
2058        std::fs::create_dir_all(staged_dir.join("site")).unwrap();
2059        std::fs::write(staged_dir.join("site/new.txt"), "new").unwrap();
2060
2061        let err = replace_dir_from_temp(&staged_dir, &final_dir).unwrap_err();
2062
2063        assert!(
2064            err.to_string().contains("not a directory"),
2065            "unexpected error: {err:#}"
2066        );
2067        assert!(staged_dir.exists());
2068        assert_eq!(
2069            std::fs::read_to_string(&final_dir).unwrap(),
2070            "not a directory"
2071        );
2072    }
2073
2074    #[test]
2075    #[cfg(unix)]
2076    fn test_replace_dir_from_temp_rejects_dangling_symlink_target() {
2077        use std::os::unix::fs::symlink;
2078
2079        let temp_dir = TempDir::new().unwrap();
2080        let final_dir = temp_dir.path().join("archive");
2081        let staged_dir = temp_dir.path().join("archive.staged");
2082        let missing_target = temp_dir.path().join("missing-archive");
2083
2084        symlink(&missing_target, &final_dir).unwrap();
2085        std::fs::create_dir_all(staged_dir.join("site")).unwrap();
2086        std::fs::write(staged_dir.join("site/new.txt"), "new").unwrap();
2087
2088        let err = replace_dir_from_temp(&staged_dir, &final_dir).unwrap_err();
2089
2090        assert!(
2091            err.to_string().contains("through symlink"),
2092            "unexpected error: {err:#}"
2093        );
2094        assert!(staged_dir.exists());
2095        assert!(
2096            std::fs::symlink_metadata(&final_dir)
2097                .unwrap()
2098                .file_type()
2099                .is_symlink()
2100        );
2101    }
2102
2103    /// `coding_agent_session_search-htiim`: regression gate mirroring
2104    /// the unwrap_key contract pinned by `encrypt.rs::
2105    /// unwrap_key_chains_aead_source_error_into_diagnostic_message`
2106    /// (commit 0b81b601). Pre-fix, key_management.rs::unwrap_key
2107    /// returned bare "Key unwrapping failed" / "Invalid DEK length"
2108    /// strings that dropped the underlying aead::Error. Post-fix,
2109    /// every site preserves the source error in the chain AND
2110    /// surfaces actionable diagnostics (slot id, input lengths).
2111    /// This test exercises the unwrap_key path with a tampered
2112    /// ciphertext and asserts:
2113    ///   1. slot id appears in the rendered error
2114    ///   2. wrapped/nonce lengths appear (sanity-check of inputs)
2115    ///   3. ":" source-separator survives (a future refactor that
2116    ///      drops `: {err}` would fail this)
2117    ///   4. legacy "Key unwrapping failed" prefix preserved so
2118    ///      operator runbook grep patterns still match.
2119    #[test]
2120    fn unwrap_key_chains_aead_source_error_into_diagnostic_message() {
2121        // Build a real wrapped DEK directly with aes_gcm so we don't
2122        // depend on a higher-level encryption engine in this module.
2123        use aes_gcm::aead::{Aead, KeyInit, Payload};
2124        use aes_gcm::{Aes256Gcm, Nonce};
2125
2126        let kek = [0u8; 32];
2127        let dek = [0u8; 32];
2128        let export_id = [42u8; 16];
2129        let slot_id = 7u8;
2130        let nonce_bytes = [3u8; 12];
2131
2132        let mut aad = Vec::with_capacity(17);
2133        aad.extend_from_slice(&export_id);
2134        aad.push(slot_id);
2135
2136        let cipher = Aes256Gcm::new_from_slice(&kek).expect("Invalid key length");
2137        let mut wrapped = cipher
2138            .encrypt(
2139                Nonce::from_slice(&nonce_bytes),
2140                Payload {
2141                    msg: &dek,
2142                    aad: &aad,
2143                },
2144            )
2145            .expect("encrypt produces wrapped DEK + auth tag");
2146
2147        // Flip the last byte of the auth tag so MAC verification fails
2148        // on unwrap. AES-GCM appends a 16-byte auth tag — flipping
2149        // any byte in it is sufficient to fail verification.
2150        let last = wrapped.len() - 1;
2151        wrapped[last] ^= 0x55;
2152
2153        let err = unwrap_key(&kek, &wrapped, &nonce_bytes, &export_id, slot_id)
2154            .expect_err("tampered ciphertext must fail unwrap");
2155        let rendered = err.to_string();
2156
2157        // Invariant 1: slot id present so operators can correlate.
2158        assert!(
2159            rendered.contains(&format!("slot {slot_id}")),
2160            "unwrap error must name the slot id; got: {rendered}"
2161        );
2162        // Invariant 2: input-size diagnostic survives.
2163        assert!(
2164            rendered.contains(&format!("{} bytes wrapped", wrapped.len())),
2165            "unwrap error must include the wrapped-ciphertext length; got: {rendered}"
2166        );
2167        assert!(
2168            rendered.contains("12 bytes nonce"),
2169            "unwrap error must include the AES-GCM nonce length; got: {rendered}"
2170        );
2171        // Invariant 3: ":" source-separator survives.
2172        assert!(
2173            rendered.contains(": "),
2174            "unwrap error must include `: <source>` separator so the \
2175             aead source error survives in the chain; got: {rendered}"
2176        );
2177        // Invariant 4: legacy prefix preserved for runbook grep.
2178        assert!(
2179            rendered.contains("Key unwrapping failed"),
2180            "unwrap error must keep the human-facing prefix for runbook \
2181             grep compatibility; got: {rendered}"
2182        );
2183    }
2184
2185    /// Companion gate for the HKDF KEK length-check arm. Pre-fix,
2186    /// `derive_kek_hkdf` returned bare "HKDF expansion produced
2187    /// invalid KEK length" with no diagnostic; post-fix, the message
2188    /// carries the actual length so operators can debug a
2189    /// frankensqlite / hkdf upstream regression that returned the
2190    /// wrong KEK size.
2191    #[test]
2192    fn derive_kek_hkdf_error_message_pins_actual_kek_length() {
2193        // Direct exercise of the conversion arm, using the public
2194        // hkdf wrapper to land at a 16-byte output (not 32). This
2195        // mirrors the gate landed in encrypt.rs by 0b81b601 so a
2196        // regression in either site fails its own assertion.
2197        let actual_kek = crate::encryption::hkdf_extract_expand(
2198            b"recovery-secret",
2199            b"salty-salty-salty-salt",
2200            b"cass-pages-kek-v2",
2201            16,
2202        )
2203        .expect("hkdf with 16-byte output must succeed");
2204        assert_eq!(actual_kek.len(), 16);
2205
2206        let conversion: Result<[u8; 32], Vec<u8>> = actual_kek.try_into();
2207        let raw_err = conversion.expect_err("16 != 32 must fail try_into");
2208        assert_eq!(raw_err.len(), 16);
2209
2210        // Codify the expected message shape so a future refactor
2211        // that reverts to `|_| ... "invalid KEK length"` without
2212        // actual_len fails the assertion.
2213        let rendered = format!(
2214            "HKDF expansion produced invalid KEK length: expected 32, got {}",
2215            raw_err.len()
2216        );
2217        assert!(rendered.contains("expected 32"));
2218        assert!(rendered.contains("got 16"));
2219    }
2220}