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::{Component, Path, PathBuf};
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 safe_staged_site_destination(dst_root: &Path, rel_path: &Path) -> Result<PathBuf> {
1176    let mut path_parts = vec![dst_root.to_path_buf()];
1177    for component in rel_path.components() {
1178        match component {
1179            Component::CurDir => {}
1180            Component::Normal(name) => path_parts.push(PathBuf::from(name)),
1181            _ => bail!(
1182                "Refusing to stage archive entry with unsafe relative path: {}",
1183                rel_path.display()
1184            ),
1185        }
1186    }
1187    Ok(path_parts.into_iter().collect())
1188}
1189
1190fn copy_site_except_runtime_state_recursive(
1191    src: &Path,
1192    dst: &Path,
1193    base: &Path,
1194    canonical_base: &Path,
1195) -> Result<()> {
1196    for entry in std::fs::read_dir(src)? {
1197        let entry = entry?;
1198        let path = entry.path();
1199        let rel_path = path.strip_prefix(base)?;
1200        let skip_root_entry = rel_path.components().count() == 1
1201            && matches!(
1202                rel_path.to_str(),
1203                Some("payload" | "blobs" | "config.json" | "integrity.json")
1204            );
1205        if skip_root_entry {
1206            continue;
1207        }
1208
1209        let metadata = std::fs::symlink_metadata(&path)?;
1210        let file_type = metadata.file_type();
1211        let dest_path = safe_staged_site_destination(dst, rel_path)?;
1212        if file_type.is_dir() {
1213            std::fs::create_dir_all(&dest_path)?;
1214            copy_site_except_runtime_state_recursive(&path, dst, base, canonical_base)?;
1215        } else if file_type.is_symlink() {
1216            let canonical_target = path.canonicalize().with_context(|| {
1217                format!(
1218                    "Failed to resolve symlinked site entry {} while staging key rotation",
1219                    rel_path.display()
1220                )
1221            })?;
1222            if !canonical_target.starts_with(canonical_base) {
1223                bail!(
1224                    "Refusing to rotate symlinked site entry outside archive root: {}",
1225                    rel_path.display()
1226                );
1227            }
1228
1229            let target_meta = std::fs::metadata(&path).with_context(|| {
1230                format!(
1231                    "Failed to read symlink target metadata for {} while staging key rotation",
1232                    rel_path.display()
1233                )
1234            })?;
1235            if !target_meta.is_file() {
1236                bail!(
1237                    "Refusing to rotate symlinked site entry that does not point to a regular file: {}",
1238                    rel_path.display()
1239                );
1240            }
1241
1242            if let Some(parent) = dest_path.parent() {
1243                std::fs::create_dir_all(parent)?;
1244            }
1245            // Materialize safe symlink targets into the staged site so the staged
1246            // integrity pass stays self-contained before the final atomic swap.
1247            std::fs::copy(&canonical_target, &dest_path).with_context(|| {
1248                format!(
1249                    "Failed copying symlink target {} into staged site path {}",
1250                    canonical_target.display(),
1251                    dest_path.display()
1252                )
1253            })?;
1254        } else if file_type.is_file() {
1255            if let Some(parent) = dest_path.parent() {
1256                std::fs::create_dir_all(parent)?;
1257            }
1258            std::fs::copy(&path, &dest_path).with_context(|| {
1259                format!(
1260                    "Failed copying staged site file {} to {}",
1261                    path.display(),
1262                    dest_path.display()
1263                )
1264            })?;
1265        }
1266    }
1267
1268    Ok(())
1269}
1270
1271fn refresh_private_artifacts(
1272    archive_dir: &Path,
1273    config: &EncryptionConfig,
1274    manifest: Option<&crate::pages::bundle::IntegrityManifest>,
1275    recovery_secret: Option<&[u8]>,
1276    remove_recovery_artifacts: bool,
1277) -> Result<()> {
1278    let Some(private_dir) = private_dir_for_archive(archive_dir)? else {
1279        return Ok(());
1280    };
1281
1282    if let Some(manifest) = manifest {
1283        let fingerprint = crate::pages::bundle::compute_fingerprint(manifest);
1284        crate::pages::bundle::write_private_fingerprint(&private_dir, &fingerprint)?;
1285    }
1286
1287    let should_generate_qr = recovery_secret.is_some()
1288        && (private_dir.join("qr-code.png").exists() || private_dir.join("qr-code.svg").exists());
1289
1290    crate::pages::bundle::write_private_artifacts_encrypted(
1291        &private_dir,
1292        config,
1293        recovery_secret,
1294        should_generate_qr,
1295        remove_recovery_artifacts,
1296    )?;
1297
1298    Ok(())
1299}
1300
1301fn private_dir_for_archive(archive_dir: &Path) -> Result<Option<std::path::PathBuf>> {
1302    if archive_dir
1303        .file_name()
1304        .map(|name| name == "site")
1305        .unwrap_or(false)
1306    {
1307        let Some(parent) = archive_dir.parent() else {
1308            return Ok(None);
1309        };
1310        let private_dir = parent.join("private");
1311        match std::fs::symlink_metadata(&private_dir) {
1312            Ok(metadata) => {
1313                let file_type = metadata.file_type();
1314                if file_type.is_symlink() {
1315                    bail!(
1316                        "private artifact directory must not be a symlink: {}",
1317                        private_dir.display()
1318                    );
1319                }
1320                if file_type.is_dir() {
1321                    return Ok(Some(private_dir));
1322                }
1323                bail!(
1324                    "private artifact path must be a directory: {}",
1325                    private_dir.display()
1326                );
1327            }
1328            Err(err) if err.kind() == std::io::ErrorKind::NotFound => {}
1329            Err(err) => {
1330                return Err(err).with_context(|| {
1331                    format!(
1332                        "Failed to inspect private artifact directory {}",
1333                        private_dir.display()
1334                    )
1335                });
1336            }
1337        }
1338    }
1339
1340    Ok(None)
1341}
1342
1343#[cfg(test)]
1344mod tests {
1345    use super::*;
1346    use crate::pages::attachments::{
1347        AttachmentConfig, AttachmentData, AttachmentProcessor, decrypt_blob, decrypt_manifest,
1348    };
1349    use crate::pages::bundle::BundleBuilder;
1350    use crate::pages::encrypt::{DecryptionEngine, EncryptionEngine, MAX_CHUNK_SIZE, PayloadMeta};
1351    use crate::pages::verify::verify_bundle;
1352    use std::cell::Cell;
1353    use tempfile::TempDir;
1354
1355    #[cfg(unix)]
1356    fn replace_viewer_with_in_tree_symlink(site_dir: &Path) {
1357        use std::os::unix::fs::symlink;
1358
1359        let real_viewer = site_dir.join("viewer-real.js");
1360        std::fs::rename(site_dir.join("viewer.js"), &real_viewer).unwrap();
1361        symlink("viewer-real.js", site_dir.join("viewer.js")).unwrap();
1362
1363        let manifest = crate::pages::bundle::generate_integrity_manifest(site_dir).unwrap();
1364        write_json_pretty(&site_dir.join("integrity.json"), &manifest).unwrap();
1365
1366        assert_eq!(verify_bundle(site_dir, false).unwrap().status, "valid");
1367    }
1368
1369    fn setup_test_archive() -> (TempDir, std::path::PathBuf) {
1370        let temp_dir = TempDir::new().unwrap();
1371        let input_path = temp_dir.path().join("input.txt");
1372        let bundle_root = temp_dir.path().join("bundle");
1373        let encrypted_dir = temp_dir.path().join("encrypted");
1374
1375        // Create test file
1376        std::fs::write(&input_path, b"Test data for key management").unwrap();
1377
1378        // Encrypt
1379        let mut engine = EncryptionEngine::new(1024).unwrap();
1380        engine.add_password_slot("test-password").unwrap();
1381        engine
1382            .encrypt_file(&input_path, &encrypted_dir, |_, _| {})
1383            .unwrap();
1384
1385        BundleBuilder::new()
1386            .build(&encrypted_dir, &bundle_root, |_, _| {})
1387            .unwrap();
1388
1389        (temp_dir, bundle_root)
1390    }
1391
1392    fn setup_test_archive_with_attachments() -> (TempDir, std::path::PathBuf) {
1393        let temp_dir = TempDir::new().unwrap();
1394        let input_path = temp_dir.path().join("input.txt");
1395        let bundle_root = temp_dir.path().join("bundle");
1396        let encrypted_dir = temp_dir.path().join("encrypted");
1397
1398        std::fs::write(&input_path, b"Test data for key management").unwrap();
1399
1400        let mut engine = EncryptionEngine::new(1024).unwrap();
1401        engine.add_password_slot("test-password").unwrap();
1402        engine
1403            .encrypt_file(&input_path, &encrypted_dir, |_, _| {})
1404            .unwrap();
1405
1406        let config = load_config(&encrypted_dir).unwrap();
1407        let dek = unwrap_dek_with_password(&config, "test-password").unwrap();
1408        let export_id_raw = BASE64_STANDARD.decode(&config.export_id).unwrap();
1409        let export_id: [u8; 16] = export_id_raw.as_slice().try_into().unwrap();
1410
1411        let mut processor = AttachmentProcessor::new(AttachmentConfig::enabled());
1412        processor
1413            .process_attachments(
1414                1,
1415                &[AttachmentData {
1416                    filename: "proof.txt".to_string(),
1417                    mime_type: "text/plain".to_string(),
1418                    data: b"attachment payload".to_vec(),
1419                }],
1420            )
1421            .unwrap();
1422        processor
1423            .write_encrypted_blobs(&encrypted_dir, &dek, &export_id)
1424            .unwrap();
1425
1426        BundleBuilder::new()
1427            .build(&encrypted_dir, &bundle_root, |_, _| {})
1428            .unwrap();
1429
1430        (temp_dir, bundle_root)
1431    }
1432
1433    fn rewrite_test_config(archive_dir: &Path, mutate: impl FnOnce(&mut EncryptionConfig)) {
1434        let site_dir = super::super::resolve_site_dir(archive_dir).unwrap();
1435        let mut config = load_config(&site_dir).unwrap();
1436        mutate(&mut config);
1437        write_json_pretty(&site_dir.join("config.json"), &config).unwrap();
1438    }
1439
1440    fn assert_unsupported_payload_format_error(err: anyhow::Error, compression: &str) {
1441        let rendered = err.to_string();
1442        assert!(
1443            rendered.contains("supports only deflate") && rendered.contains(compression),
1444            "unexpected unsupported-format error: {err:#}"
1445        );
1446    }
1447
1448    #[test]
1449    #[cfg(unix)]
1450    fn test_private_dir_for_archive_rejects_symlinked_private_dir() {
1451        use std::os::unix::fs::symlink;
1452
1453        let temp = TempDir::new().unwrap();
1454        let site_dir = temp.path().join("bundle/site");
1455        let outside_private = temp.path().join("outside-private");
1456        std::fs::create_dir_all(&site_dir).unwrap();
1457        std::fs::create_dir_all(&outside_private).unwrap();
1458        symlink(&outside_private, temp.path().join("bundle/private")).unwrap();
1459
1460        let err = private_dir_for_archive(&site_dir).unwrap_err();
1461
1462        assert!(
1463            err.to_string().contains("must not be a symlink"),
1464            "unexpected error: {err:#}"
1465        );
1466        assert!(
1467            std::fs::symlink_metadata(temp.path().join("bundle/private"))
1468                .unwrap()
1469                .file_type()
1470                .is_symlink(),
1471            "rejected private directory symlink should remain untouched"
1472        );
1473    }
1474
1475    #[test]
1476    fn test_private_dir_for_archive_rejects_non_directory_private_path() {
1477        let temp = TempDir::new().unwrap();
1478        let site_dir = temp.path().join("bundle/site");
1479        std::fs::create_dir_all(&site_dir).unwrap();
1480        std::fs::write(temp.path().join("bundle/private"), "not a directory").unwrap();
1481
1482        let err = private_dir_for_archive(&site_dir).unwrap_err();
1483
1484        assert!(
1485            err.to_string().contains("must be a directory"),
1486            "unexpected error: {err:#}"
1487        );
1488        assert_eq!(
1489            std::fs::read_to_string(temp.path().join("bundle/private")).unwrap(),
1490            "not a directory"
1491        );
1492    }
1493
1494    #[test]
1495    fn test_decrypt_all_chunks_rejects_mismatched_chunk_count_before_progress() {
1496        let temp_dir = TempDir::new().unwrap();
1497        let archive_dir = temp_dir.path();
1498        let config = EncryptionConfig {
1499            version: SCHEMA_VERSION,
1500            export_id: BASE64_STANDARD.encode([0u8; 16]),
1501            base_nonce: BASE64_STANDARD.encode([0u8; 12]),
1502            compression: "deflate".to_string(),
1503            kdf_defaults: Argon2Params::default(),
1504            payload: PayloadMeta {
1505                chunk_size: 1024,
1506                chunk_count: 0,
1507                total_compressed_size: 0,
1508                total_plaintext_size: 0,
1509                files: vec!["payload/chunk-00000.bin".to_string()],
1510            },
1511            key_slots: Vec::new(),
1512        };
1513        let progress_calls = Cell::new(0);
1514
1515        let err = decrypt_all_chunks(archive_dir, &[0u8; 32], &config, |progress| {
1516            assert!(progress.is_finite(), "progress must be finite: {progress}");
1517            progress_calls.set(progress_calls.get() + 1);
1518        })
1519        .unwrap_err();
1520
1521        assert!(
1522            err.to_string().contains("chunk_count 0"),
1523            "unexpected error: {err:#}"
1524        );
1525        assert_eq!(progress_calls.get(), 0);
1526    }
1527
1528    #[test]
1529    fn test_key_list() {
1530        let (_temp_dir, archive_dir) = setup_test_archive();
1531
1532        let result = key_list(&archive_dir).unwrap();
1533        assert_eq!(result.active_slots, 1);
1534        assert_eq!(result.slots.len(), 1);
1535        assert_eq!(result.slots[0].slot_type, "password");
1536        assert_eq!(result.slots[0].kdf, "argon2id");
1537    }
1538
1539    #[test]
1540    fn test_key_mutations_reject_unsupported_payload_compression() {
1541        let (_temp_dir, archive_dir) = setup_test_archive();
1542        key_add_password(&archive_dir, "test-password", "second-password").unwrap();
1543        rewrite_test_config(&archive_dir, |config| {
1544            config.compression = "zstd".to_string();
1545        });
1546
1547        let err = key_add_password(&archive_dir, "test-password", "third-password").unwrap_err();
1548        assert_unsupported_payload_format_error(err, "zstd");
1549
1550        let err = key_add_recovery(&archive_dir, "test-password").unwrap_err();
1551        assert_unsupported_payload_format_error(err, "zstd");
1552
1553        let err = key_revoke(&archive_dir, "second-password", 0).unwrap_err();
1554        assert_unsupported_payload_format_error(err, "zstd");
1555
1556        let err =
1557            key_rotate(&archive_dir, "test-password", "new-password", false, |_| {}).unwrap_err();
1558        assert_unsupported_payload_format_error(err, "zstd");
1559
1560        let config = load_config(&archive_dir).unwrap();
1561        assert_eq!(config.key_slots.len(), 2);
1562        assert!(unwrap_dek_with_password(&config, "test-password").is_ok());
1563        assert!(unwrap_dek_with_password(&config, "second-password").is_ok());
1564        assert!(unwrap_dek_with_password(&config, "third-password").is_err());
1565        assert!(unwrap_dek_with_password(&config, "new-password").is_err());
1566    }
1567
1568    #[test]
1569    fn test_key_rotate_rejects_oversized_payload_chunk_size_before_rewriting() {
1570        let (_temp_dir, archive_dir) = setup_test_archive();
1571        rewrite_test_config(&archive_dir, |config| {
1572            config.payload.chunk_size = MAX_CHUNK_SIZE + 1;
1573        });
1574
1575        let err =
1576            key_rotate(&archive_dir, "test-password", "new-password", false, |_| {}).unwrap_err();
1577        let rendered = err.to_string();
1578        assert!(
1579            rendered.contains("chunk_size") && rendered.contains("must be <="),
1580            "unexpected chunk-size error: {err:#}"
1581        );
1582
1583        let config = load_config(&archive_dir).unwrap();
1584        assert!(unwrap_dek_with_password(&config, "test-password").is_ok());
1585        assert!(unwrap_dek_with_password(&config, "new-password").is_err());
1586    }
1587
1588    #[test]
1589    fn test_key_rotate_chunk_count_preflight_preserves_nonce_space_limit() {
1590        ensure_archive_chunk_count_fits_nonce_space(u64::from(u32::MAX), 1).unwrap();
1591
1592        let err =
1593            ensure_archive_chunk_count_fits_nonce_space(u64::from(u32::MAX) + 1, 1).unwrap_err();
1594        let rendered = err.to_string();
1595        assert!(
1596            rendered.contains("exceeds maximum") && rendered.contains(&u32::MAX.to_string()),
1597            "unexpected chunk-count error: {rendered}"
1598        );
1599    }
1600
1601    #[test]
1602    fn test_key_add_password() {
1603        let (_temp_dir, archive_dir) = setup_test_archive();
1604
1605        // Add new password
1606        let slot_id = key_add_password(&archive_dir, "test-password", "new-password").unwrap();
1607        assert_eq!(slot_id, 1);
1608
1609        // Verify it was added
1610        let result = key_list(&archive_dir).unwrap();
1611        assert_eq!(result.active_slots, 2);
1612
1613        // Verify new password works
1614        let config = load_config(&archive_dir).unwrap();
1615        let dek = unwrap_dek_with_password(&config, "new-password").unwrap();
1616        assert!(!dek.iter().all(|&b| b == 0));
1617    }
1618
1619    #[test]
1620    fn test_key_add_recovery() {
1621        let (_temp_dir, archive_dir) = setup_test_archive();
1622
1623        // Add recovery slot
1624        let (slot_id, secret) = key_add_recovery(&archive_dir, "test-password").unwrap();
1625        assert_eq!(slot_id, 1);
1626        assert_eq!(secret.entropy_bits(), 256);
1627
1628        // Verify it was added
1629        let result = key_list(&archive_dir).unwrap();
1630        assert_eq!(result.active_slots, 2);
1631        assert_eq!(result.slots[1].slot_type, "recovery");
1632        assert_eq!(result.slots[1].kdf, "hkdf-sha256");
1633    }
1634
1635    #[test]
1636    fn test_key_add_wrong_password_fails() {
1637        let (_temp_dir, archive_dir) = setup_test_archive();
1638
1639        let result = key_add_password(&archive_dir, "wrong-password", "new-password");
1640        assert!(result.is_err());
1641    }
1642
1643    #[test]
1644    fn test_key_revoke() {
1645        let (_temp_dir, archive_dir) = setup_test_archive();
1646
1647        // Add second slot
1648        key_add_password(&archive_dir, "test-password", "second-password").unwrap();
1649
1650        // Revoke first slot using second password
1651        let result = key_revoke(&archive_dir, "second-password", 0).unwrap();
1652        assert_eq!(result.revoked_slot_id, 0);
1653        assert_eq!(result.remaining_slots, 1);
1654
1655        // Old password should no longer work
1656        let config = load_config(&archive_dir).unwrap();
1657        assert!(unwrap_dek_with_password(&config, "test-password").is_err());
1658
1659        // Second password should still work
1660        assert!(unwrap_dek_with_password(&config, "second-password").is_ok());
1661    }
1662
1663    #[test]
1664    fn test_key_revoke_last_slot_fails() {
1665        let (_temp_dir, archive_dir) = setup_test_archive();
1666
1667        let result = key_revoke(&archive_dir, "test-password", 0);
1668        assert!(result.is_err());
1669        assert!(result.unwrap_err().to_string().contains("last remaining"));
1670    }
1671
1672    #[test]
1673    fn test_key_revoke_auth_slot_fails() {
1674        let (_temp_dir, archive_dir) = setup_test_archive();
1675
1676        // Add second slot
1677        key_add_password(&archive_dir, "test-password", "second-password").unwrap();
1678
1679        // Try to revoke slot 0 using slot 0's password
1680        let result = key_revoke(&archive_dir, "test-password", 0);
1681        assert!(result.is_err());
1682        assert!(result.unwrap_err().to_string().contains("authentication"));
1683    }
1684
1685    #[test]
1686    fn test_key_rotate() {
1687        let (temp_dir, archive_dir) = setup_test_archive();
1688        let decrypted_path = temp_dir.path().join("decrypted.txt");
1689
1690        // Rotate keys
1691        let result =
1692            key_rotate(&archive_dir, "test-password", "new-password", false, |_| {}).unwrap();
1693        assert_eq!(result.slot_count, 1);
1694        assert!(result.recovery_secret.is_none());
1695
1696        // Old password should fail
1697        let config = load_config(&archive_dir).unwrap();
1698        assert!(unwrap_dek_with_password(&config, "test-password").is_err());
1699
1700        // New password should work and decrypt correctly
1701        let decryptor = DecryptionEngine::unlock_with_password(config, "new-password").unwrap();
1702        decryptor
1703            .decrypt_to_file(&archive_dir, &decrypted_path, |_, _| {})
1704            .unwrap();
1705
1706        let decrypted = std::fs::read(&decrypted_path).unwrap();
1707        assert_eq!(decrypted, b"Test data for key management");
1708    }
1709
1710    #[test]
1711    fn test_key_rotate_with_recovery() {
1712        let (_temp_dir, archive_dir) = setup_test_archive();
1713
1714        // Rotate keys with recovery
1715        let result =
1716            key_rotate(&archive_dir, "test-password", "new-password", true, |_| {}).unwrap();
1717        assert_eq!(result.slot_count, 2);
1718        assert!(result.recovery_secret.is_some());
1719
1720        // Verify recovery slot
1721        let list_result = key_list(&archive_dir).unwrap();
1722        assert_eq!(list_result.slots.len(), 2);
1723        assert_eq!(list_result.slots[0].slot_type, "password");
1724        assert_eq!(list_result.slots[1].slot_type, "recovery");
1725    }
1726
1727    #[test]
1728    fn test_key_add_after_revoke_no_id_collision() {
1729        let (_temp_dir, archive_dir) = setup_test_archive();
1730
1731        // Add slots 1 and 2
1732        key_add_password(&archive_dir, "test-password", "password-1").unwrap();
1733        key_add_password(&archive_dir, "test-password", "password-2").unwrap();
1734
1735        // Now have slots [0, 1, 2]
1736        let list = key_list(&archive_dir).unwrap();
1737        assert_eq!(list.slots.len(), 3);
1738
1739        // Revoke slot 1 using slot 2's password
1740        key_revoke(&archive_dir, "password-2", 1).unwrap();
1741
1742        // Now have slots [0, 2] (gap at 1)
1743        let list = key_list(&archive_dir).unwrap();
1744        assert_eq!(list.slots.len(), 2);
1745        let ids: Vec<u8> = list.slots.iter().map(|s| s.id).collect();
1746        assert_eq!(ids, vec![0, 2]);
1747
1748        // Add new slot - should get ID 3, not 2
1749        let new_id = key_add_password(&archive_dir, "test-password", "password-3").unwrap();
1750        assert_eq!(new_id, 3, "New slot should get max_id + 1, not len()");
1751
1752        // Verify all passwords still work
1753        let config = load_config(&archive_dir).unwrap();
1754        assert!(unwrap_dek_with_password(&config, "test-password").is_ok());
1755        assert!(unwrap_dek_with_password(&config, "password-1").is_err()); // Revoked
1756        assert!(unwrap_dek_with_password(&config, "password-2").is_ok());
1757        assert!(unwrap_dek_with_password(&config, "password-3").is_ok());
1758    }
1759
1760    #[test]
1761    fn test_next_key_slot_id_rejects_max_id() {
1762        let (_temp_dir, archive_dir) = setup_test_archive();
1763        let mut config = load_config(&archive_dir).unwrap();
1764        config.key_slots[0].id = u8::MAX;
1765
1766        let err = next_key_slot_id(&config.key_slots).unwrap_err();
1767
1768        assert_eq!(
1769            err.to_string(),
1770            "Cannot add more key slots: maximum slot ID (255) reached"
1771        );
1772    }
1773
1774    #[test]
1775    fn test_key_add_password_preserves_valid_integrity_manifest() {
1776        let (_temp_dir, archive_dir) = setup_test_archive();
1777
1778        assert_eq!(verify_bundle(&archive_dir, false).unwrap().status, "valid");
1779
1780        key_add_password(&archive_dir, "test-password", "new-password").unwrap();
1781
1782        assert_eq!(verify_bundle(&archive_dir, false).unwrap().status, "valid");
1783    }
1784
1785    #[test]
1786    fn test_key_rotate_preserves_valid_integrity_manifest() {
1787        let (_temp_dir, archive_dir) = setup_test_archive();
1788
1789        assert_eq!(verify_bundle(&archive_dir, false).unwrap().status, "valid");
1790
1791        key_rotate(&archive_dir, "test-password", "new-password", true, |_| {}).unwrap();
1792
1793        assert_eq!(verify_bundle(&archive_dir, false).unwrap().status, "valid");
1794    }
1795
1796    #[test]
1797    #[cfg(unix)]
1798    fn test_key_add_password_preserves_in_tree_symlinked_required_asset() {
1799        let (_temp_dir, archive_dir) = setup_test_archive();
1800        let site_dir = super::super::resolve_site_dir(&archive_dir).unwrap();
1801        replace_viewer_with_in_tree_symlink(&site_dir);
1802
1803        key_add_password(&archive_dir, "test-password", "new-password").unwrap();
1804
1805        assert_eq!(verify_bundle(&archive_dir, false).unwrap().status, "valid");
1806        assert!(
1807            std::fs::symlink_metadata(site_dir.join("viewer.js"))
1808                .unwrap()
1809                .file_type()
1810                .is_symlink()
1811        );
1812    }
1813
1814    #[test]
1815    #[cfg(unix)]
1816    fn test_key_rotate_materializes_in_tree_symlinked_required_asset() {
1817        let (_temp_dir, archive_dir) = setup_test_archive();
1818        let site_dir = super::super::resolve_site_dir(&archive_dir).unwrap();
1819        replace_viewer_with_in_tree_symlink(&site_dir);
1820        let expected_viewer = std::fs::read(site_dir.join("viewer-real.js")).unwrap();
1821
1822        key_rotate(&archive_dir, "test-password", "new-password", true, |_| {}).unwrap();
1823
1824        let viewer_metadata = std::fs::symlink_metadata(site_dir.join("viewer.js")).unwrap();
1825        assert!(viewer_metadata.file_type().is_file());
1826        assert!(!viewer_metadata.file_type().is_symlink());
1827        assert_eq!(
1828            std::fs::read(site_dir.join("viewer.js")).unwrap(),
1829            expected_viewer
1830        );
1831        assert_eq!(verify_bundle(&archive_dir, false).unwrap().status, "valid");
1832    }
1833
1834    #[test]
1835    #[cfg(unix)]
1836    fn test_key_rotate_rejects_payload_directory_symlink_escape() {
1837        use std::os::unix::fs::symlink;
1838
1839        let (temp_dir, archive_dir) = setup_test_archive();
1840        let site_dir = super::super::resolve_site_dir(&archive_dir).unwrap();
1841        let payload_dir = site_dir.join("payload");
1842        let outside_payload_dir = temp_dir.path().join("outside-payload");
1843
1844        std::fs::rename(&payload_dir, &outside_payload_dir).unwrap();
1845        symlink(&outside_payload_dir, &payload_dir).unwrap();
1846
1847        let err =
1848            key_rotate(&archive_dir, "test-password", "new-password", false, |_| {}).unwrap_err();
1849        assert!(
1850            err.to_string().contains("escapes archive directory"),
1851            "unexpected error: {err:#}"
1852        );
1853    }
1854
1855    #[test]
1856    fn test_key_add_password_updates_private_fingerprint_and_master_key() {
1857        let (_temp_dir, archive_dir) = setup_test_archive();
1858        let site_dir = super::super::resolve_site_dir(&archive_dir).unwrap();
1859        let private_dir = site_dir.parent().unwrap().join("private");
1860
1861        let old_fingerprint =
1862            std::fs::read_to_string(private_dir.join("integrity-fingerprint.txt")).unwrap();
1863        let old_master_key = std::fs::read_to_string(private_dir.join("master-key.json")).unwrap();
1864
1865        key_add_password(&archive_dir, "test-password", "new-password").unwrap();
1866
1867        let new_fingerprint =
1868            std::fs::read_to_string(private_dir.join("integrity-fingerprint.txt")).unwrap();
1869        let new_master_key = std::fs::read_to_string(private_dir.join("master-key.json")).unwrap();
1870
1871        assert_ne!(old_fingerprint, new_fingerprint);
1872        assert_ne!(old_master_key, new_master_key);
1873    }
1874
1875    #[test]
1876    fn test_key_add_recovery_writes_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        assert!(!private_dir.join("recovery-secret.txt").exists());
1882
1883        let (_slot_id, secret) = key_add_recovery(&archive_dir, "test-password").unwrap();
1884        let recovery_file =
1885            std::fs::read_to_string(private_dir.join("recovery-secret.txt")).unwrap();
1886
1887        assert!(recovery_file.contains(secret.encoded()));
1888    }
1889
1890    #[test]
1891    fn test_key_revoke_recovery_removes_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 (recovery_slot_id, _secret) = key_add_recovery(&archive_dir, "test-password").unwrap();
1897        key_add_password(&archive_dir, "test-password", "second-password").unwrap();
1898        assert!(private_dir.join("recovery-secret.txt").exists());
1899
1900        key_revoke(&archive_dir, "second-password", recovery_slot_id).unwrap();
1901
1902        assert!(!private_dir.join("recovery-secret.txt").exists());
1903    }
1904
1905    #[test]
1906    fn test_key_revoke_one_of_multiple_recovery_slots_removes_stale_private_recovery_artifact() {
1907        let (_temp_dir, archive_dir) = setup_test_archive();
1908        let site_dir = super::super::resolve_site_dir(&archive_dir).unwrap();
1909        let private_dir = site_dir.parent().unwrap().join("private");
1910
1911        let (first_recovery_slot_id, first_secret) =
1912            key_add_recovery(&archive_dir, "test-password").unwrap();
1913        let (second_recovery_slot_id, second_secret) =
1914            key_add_recovery(&archive_dir, "test-password").unwrap();
1915
1916        let recovery_file_before =
1917            std::fs::read_to_string(private_dir.join("recovery-secret.txt")).unwrap();
1918        assert!(recovery_file_before.contains(second_secret.encoded()));
1919
1920        key_revoke(&archive_dir, "test-password", second_recovery_slot_id).unwrap();
1921
1922        assert!(!private_dir.join("recovery-secret.txt").exists());
1923
1924        let config = load_config(&archive_dir).unwrap();
1925        assert!(DecryptionEngine::unlock_with_recovery(config, first_secret.as_bytes()).is_ok());
1926
1927        assert_ne!(first_recovery_slot_id, second_recovery_slot_id);
1928    }
1929
1930    #[test]
1931    fn test_key_rotate_refreshes_private_recovery_and_master_key() {
1932        let (_temp_dir, archive_dir) = setup_test_archive();
1933        let site_dir = super::super::resolve_site_dir(&archive_dir).unwrap();
1934        let private_dir = site_dir.parent().unwrap().join("private");
1935
1936        let old_master_key = std::fs::read_to_string(private_dir.join("master-key.json")).unwrap();
1937        let result =
1938            key_rotate(&archive_dir, "test-password", "new-password", true, |_| {}).unwrap();
1939
1940        let new_master_key = std::fs::read_to_string(private_dir.join("master-key.json")).unwrap();
1941        let recovery_file =
1942            std::fs::read_to_string(private_dir.join("recovery-secret.txt")).unwrap();
1943
1944        assert_ne!(old_master_key, new_master_key);
1945        assert!(recovery_file.contains(result.recovery_secret.as_deref().unwrap()));
1946    }
1947
1948    #[test]
1949    fn test_key_rotate_without_recovery_removes_stale_private_recovery_artifact() {
1950        let (_temp_dir, archive_dir) = setup_test_archive();
1951        let site_dir = super::super::resolve_site_dir(&archive_dir).unwrap();
1952        let private_dir = site_dir.parent().unwrap().join("private");
1953
1954        let (_slot_id, _secret) = key_add_recovery(&archive_dir, "test-password").unwrap();
1955        assert!(private_dir.join("recovery-secret.txt").exists());
1956
1957        key_rotate(&archive_dir, "test-password", "new-password", false, |_| {}).unwrap();
1958
1959        assert!(!private_dir.join("recovery-secret.txt").exists());
1960        assert!(!private_dir.join("qr-code.png").exists());
1961        assert!(!private_dir.join("qr-code.svg").exists());
1962    }
1963
1964    #[test]
1965    fn test_key_rotate_reencrypts_attachment_blobs() {
1966        let (_temp_dir, archive_dir) = setup_test_archive_with_attachments();
1967
1968        assert_eq!(verify_bundle(&archive_dir, false).unwrap().status, "valid");
1969
1970        key_rotate(&archive_dir, "test-password", "new-password", false, |_| {}).unwrap();
1971
1972        let site_dir = super::super::resolve_site_dir(&archive_dir).unwrap();
1973        let config = load_config(&archive_dir).unwrap();
1974        let dek = unwrap_dek_with_password(&config, "new-password").unwrap();
1975        let export_id_raw = BASE64_STANDARD.decode(&config.export_id).unwrap();
1976        let export_id: [u8; 16] = export_id_raw.as_slice().try_into().unwrap();
1977
1978        let manifest_ciphertext =
1979            std::fs::read(site_dir.join("blobs").join("manifest.enc")).unwrap();
1980        let manifest = decrypt_manifest(&manifest_ciphertext, &dek, &export_id).unwrap();
1981        assert_eq!(manifest.entries.len(), 1);
1982        assert_eq!(manifest.entries[0].filename, "proof.txt");
1983
1984        let blob_ciphertext = std::fs::read(
1985            site_dir
1986                .join("blobs")
1987                .join(format!("{}.bin", manifest.entries[0].hash)),
1988        )
1989        .unwrap();
1990        let plaintext = decrypt_blob(
1991            &blob_ciphertext,
1992            &dek,
1993            &export_id,
1994            &manifest.entries[0].hash,
1995        )
1996        .unwrap();
1997        assert_eq!(plaintext, b"attachment payload");
1998        assert_eq!(verify_bundle(&archive_dir, false).unwrap().status, "valid");
1999    }
2000
2001    #[test]
2002    fn test_key_rotate_failure_before_site_swap_preserves_live_archive() {
2003        let (temp_dir, archive_dir) = setup_test_archive_with_attachments();
2004        let decrypted_path = temp_dir.path().join("decrypted-after-failure.txt");
2005        let site_dir = super::super::resolve_site_dir(&archive_dir).unwrap();
2006
2007        std::fs::write(site_dir.join("blobs").join("manifest.enc"), b"corrupted").unwrap();
2008
2009        let rotate_result =
2010            key_rotate(&archive_dir, "test-password", "new-password", false, |_| {});
2011        assert!(rotate_result.is_err());
2012
2013        let config = load_config(&archive_dir).unwrap();
2014        assert!(unwrap_dek_with_password(&config, "new-password").is_err());
2015
2016        let decryptor = DecryptionEngine::unlock_with_password(config, "test-password").unwrap();
2017        decryptor
2018            .decrypt_to_file(&archive_dir, &decrypted_path, |_, _| {})
2019            .unwrap();
2020
2021        let decrypted = std::fs::read(&decrypted_path).unwrap();
2022        assert_eq!(decrypted, b"Test data for key management");
2023    }
2024
2025    #[test]
2026    fn test_write_json_pretty_atomically_overwrites_existing_file() {
2027        let temp_dir = TempDir::new().unwrap();
2028        let path = temp_dir.path().join("config.json");
2029        std::fs::write(&path, "{\"before\":true}\n").unwrap();
2030
2031        let value = serde_json::json!({ "after": true });
2032        write_json_pretty_atomically(&path, &value).unwrap();
2033
2034        let written: serde_json::Value =
2035            serde_json::from_str(&std::fs::read_to_string(&path).unwrap()).unwrap();
2036        assert_eq!(written, value);
2037    }
2038
2039    #[test]
2040    fn test_replace_dir_from_temp_overwrites_existing_site() {
2041        let temp_dir = TempDir::new().unwrap();
2042        let final_dir = temp_dir.path().join("archive");
2043        let staged_dir = temp_dir.path().join("archive.staged");
2044
2045        std::fs::create_dir_all(final_dir.join("site")).unwrap();
2046        std::fs::write(final_dir.join("site/old.txt"), "old").unwrap();
2047
2048        std::fs::create_dir_all(staged_dir.join("site")).unwrap();
2049        std::fs::write(staged_dir.join("site/new.txt"), "new").unwrap();
2050
2051        replace_dir_from_temp(&staged_dir, &final_dir).unwrap();
2052
2053        assert!(!staged_dir.exists());
2054        assert!(final_dir.join("site/new.txt").exists());
2055        assert!(!final_dir.join("site/old.txt").exists());
2056        let sidecars = std::fs::read_dir(temp_dir.path())
2057            .unwrap()
2058            .map(|entry| entry.unwrap().file_name().to_string_lossy().into_owned())
2059            .collect::<Vec<_>>();
2060        assert!(
2061            !sidecars.iter().any(|name| name.contains(".archive.bak.")),
2062            "backup sidecar should be cleaned up, found: {sidecars:?}"
2063        );
2064    }
2065
2066    #[test]
2067    fn test_replace_dir_from_temp_rejects_file_target() {
2068        let temp_dir = TempDir::new().unwrap();
2069        let final_dir = temp_dir.path().join("archive");
2070        let staged_dir = temp_dir.path().join("archive.staged");
2071
2072        std::fs::write(&final_dir, "not a directory").unwrap();
2073        std::fs::create_dir_all(staged_dir.join("site")).unwrap();
2074        std::fs::write(staged_dir.join("site/new.txt"), "new").unwrap();
2075
2076        let err = replace_dir_from_temp(&staged_dir, &final_dir).unwrap_err();
2077
2078        assert!(
2079            err.to_string().contains("not a directory"),
2080            "unexpected error: {err:#}"
2081        );
2082        assert!(staged_dir.exists());
2083        assert_eq!(
2084            std::fs::read_to_string(&final_dir).unwrap(),
2085            "not a directory"
2086        );
2087    }
2088
2089    #[test]
2090    #[cfg(unix)]
2091    fn test_replace_dir_from_temp_rejects_dangling_symlink_target() {
2092        use std::os::unix::fs::symlink;
2093
2094        let temp_dir = TempDir::new().unwrap();
2095        let final_dir = temp_dir.path().join("archive");
2096        let staged_dir = temp_dir.path().join("archive.staged");
2097        let missing_target = temp_dir.path().join("missing-archive");
2098
2099        symlink(&missing_target, &final_dir).unwrap();
2100        std::fs::create_dir_all(staged_dir.join("site")).unwrap();
2101        std::fs::write(staged_dir.join("site/new.txt"), "new").unwrap();
2102
2103        let err = replace_dir_from_temp(&staged_dir, &final_dir).unwrap_err();
2104
2105        assert!(
2106            err.to_string().contains("through symlink"),
2107            "unexpected error: {err:#}"
2108        );
2109        assert!(staged_dir.exists());
2110        assert!(
2111            std::fs::symlink_metadata(&final_dir)
2112                .unwrap()
2113                .file_type()
2114                .is_symlink()
2115        );
2116    }
2117
2118    /// `coding_agent_session_search-htiim`: regression gate mirroring
2119    /// the unwrap_key contract pinned by `encrypt.rs::
2120    /// unwrap_key_chains_aead_source_error_into_diagnostic_message`
2121    /// (commit 0b81b601). Pre-fix, key_management.rs::unwrap_key
2122    /// returned bare "Key unwrapping failed" / "Invalid DEK length"
2123    /// strings that dropped the underlying aead::Error. Post-fix,
2124    /// every site preserves the source error in the chain AND
2125    /// surfaces actionable diagnostics (slot id, input lengths).
2126    /// This test exercises the unwrap_key path with a tampered
2127    /// ciphertext and asserts:
2128    ///   1. slot id appears in the rendered error
2129    ///   2. wrapped/nonce lengths appear (sanity-check of inputs)
2130    ///   3. ":" source-separator survives (a future refactor that
2131    ///      drops `: {err}` would fail this)
2132    ///   4. legacy "Key unwrapping failed" prefix preserved so
2133    ///      operator runbook grep patterns still match.
2134    #[test]
2135    fn unwrap_key_chains_aead_source_error_into_diagnostic_message() {
2136        // Build a real wrapped DEK directly with aes_gcm so we don't
2137        // depend on a higher-level encryption engine in this module.
2138        use aes_gcm::aead::{Aead, KeyInit, Payload};
2139        use aes_gcm::{Aes256Gcm, Nonce};
2140
2141        let kek = [0u8; 32];
2142        let dek = [0u8; 32];
2143        let export_id = [42u8; 16];
2144        let slot_id = 7u8;
2145        let nonce_bytes = [3u8; 12];
2146
2147        let mut aad = Vec::with_capacity(17);
2148        aad.extend_from_slice(&export_id);
2149        aad.push(slot_id);
2150
2151        let cipher = Aes256Gcm::new_from_slice(&kek).expect("Invalid key length");
2152        let mut wrapped = cipher
2153            .encrypt(
2154                Nonce::from_slice(&nonce_bytes),
2155                Payload {
2156                    msg: &dek,
2157                    aad: &aad,
2158                },
2159            )
2160            .expect("encrypt produces wrapped DEK + auth tag");
2161
2162        // Flip the last byte of the auth tag so MAC verification fails
2163        // on unwrap. AES-GCM appends a 16-byte auth tag — flipping
2164        // any byte in it is sufficient to fail verification.
2165        let last = wrapped.len() - 1;
2166        wrapped[last] ^= 0x55;
2167
2168        let err = unwrap_key(&kek, &wrapped, &nonce_bytes, &export_id, slot_id)
2169            .expect_err("tampered ciphertext must fail unwrap");
2170        let rendered = err.to_string();
2171
2172        // Invariant 1: slot id present so operators can correlate.
2173        assert!(
2174            rendered.contains(&format!("slot {slot_id}")),
2175            "unwrap error must name the slot id; got: {rendered}"
2176        );
2177        // Invariant 2: input-size diagnostic survives.
2178        assert!(
2179            rendered.contains(&format!("{} bytes wrapped", wrapped.len())),
2180            "unwrap error must include the wrapped-ciphertext length; got: {rendered}"
2181        );
2182        assert!(
2183            rendered.contains("12 bytes nonce"),
2184            "unwrap error must include the AES-GCM nonce length; got: {rendered}"
2185        );
2186        // Invariant 3: ":" source-separator survives.
2187        assert!(
2188            rendered.contains(": "),
2189            "unwrap error must include `: <source>` separator so the \
2190             aead source error survives in the chain; got: {rendered}"
2191        );
2192        // Invariant 4: legacy prefix preserved for runbook grep.
2193        assert!(
2194            rendered.contains("Key unwrapping failed"),
2195            "unwrap error must keep the human-facing prefix for runbook \
2196             grep compatibility; got: {rendered}"
2197        );
2198    }
2199
2200    /// Companion gate for the HKDF KEK length-check arm. Pre-fix,
2201    /// `derive_kek_hkdf` returned bare "HKDF expansion produced
2202    /// invalid KEK length" with no diagnostic; post-fix, the message
2203    /// carries the actual length so operators can debug a
2204    /// frankensqlite / hkdf upstream regression that returned the
2205    /// wrong KEK size.
2206    #[test]
2207    fn derive_kek_hkdf_error_message_pins_actual_kek_length() {
2208        // Direct exercise of the conversion arm, using the public
2209        // hkdf wrapper to land at a 16-byte output (not 32). This
2210        // mirrors the gate landed in encrypt.rs by 0b81b601 so a
2211        // regression in either site fails its own assertion.
2212        let actual_kek = crate::encryption::hkdf_extract_expand(
2213            b"recovery-secret",
2214            b"salty-salty-salty-salt",
2215            b"cass-pages-kek-v2",
2216            16,
2217        )
2218        .expect("hkdf with 16-byte output must succeed");
2219        assert_eq!(actual_kek.len(), 16);
2220
2221        let conversion: Result<[u8; 32], Vec<u8>> = actual_kek.try_into();
2222        let raw_err = conversion.expect_err("16 != 32 must fail try_into");
2223        assert_eq!(raw_err.len(), 16);
2224
2225        // Codify the expected message shape so a future refactor
2226        // that reverts to `|_| ... "invalid KEK length"` without
2227        // actual_len fails the assertion.
2228        let rendered = format!(
2229            "HKDF expansion produced invalid KEK length: expected 32, got {}",
2230            raw_err.len()
2231        );
2232        assert!(rendered.contains("expected 32"));
2233        assert!(rendered.contains("got 16"));
2234    }
2235}