libsoliton/soliton/tests/integration_storage.rs

#![allow(deprecated)] // Tests exercise from_bytes directly.

use soliton::primitives::random;
use soliton::primitives::xwing;
use soliton::ratchet::RatchetState;
use soliton::storage::{StorageKey, StorageKeyRing, decrypt_blob, encrypt_blob};

#[test]
fn key_rotation_lifecycle() {
    let key_v1 = StorageKey::new(1, random::random_array()).unwrap();
    let key_v2 = StorageKey::new(2, random::random_array()).unwrap();

    let mut ring = StorageKeyRing::new(key_v1).unwrap();

    // Encrypt blobs with v1.
    let blob_a = encrypt_blob(
        ring.active_key().unwrap(),
        b"message A",
        "channel-1",
        "seg-0",
        false,
    )
    .unwrap();
    let blob_b = encrypt_blob(
        ring.active_key().unwrap(),
        b"message B",
        "channel-1",
        "seg-1",
        true,
    )
    .unwrap();

    // Add v2 as the new active key.
    ring.add_key(key_v2, true).unwrap();
    assert_eq!(ring.active_key().unwrap().version(), 2);

    // Encrypt new blob with v2.
    let blob_c = encrypt_blob(
        ring.active_key().unwrap(),
        b"message C",
        "channel-2",
        "seg-0",
        false,
    )
    .unwrap();

    // Decrypt all blobs — v1 blobs should still work.
    let pt_a = decrypt_blob(&ring, &blob_a, "channel-1", "seg-0").unwrap();
    assert_eq!(&*pt_a, b"message A");
    let pt_b = decrypt_blob(&ring, &blob_b, "channel-1", "seg-1").unwrap();
    assert_eq!(&*pt_b, b"message B");
    let pt_c = decrypt_blob(&ring, &blob_c, "channel-2", "seg-0").unwrap();
    assert_eq!(&*pt_c, b"message C");

    // Remove v1.
    assert!(ring.remove_key(1).unwrap());

    // v1 blobs can no longer be decrypted.
    assert!(decrypt_blob(&ring, &blob_a, "channel-1", "seg-0").is_err());

    // v2 blobs still work.
    let pt_c2 = decrypt_blob(&ring, &blob_c, "channel-2", "seg-0").unwrap();
    assert_eq!(&*pt_c2, b"message C");
}

#[test]
fn ratchet_state_storage_round_trip() {
    // Create a ratchet state with some conversation history.
    let (ek_pk, ek_sk) = xwing::keygen().unwrap();
    let rk: [u8; 32] = random::random_array();
    let ck: [u8; 32] = random::random_array();
    let fp_a = [0xAAu8; 32];
    let fp_b = [0xBBu8; 32];
    let mut alice = RatchetState::init_alice(rk, ck, fp_a, fp_b, ek_pk.clone(), ek_sk).unwrap();
    let mut bob = RatchetState::init_bob(rk, ck, fp_b, fp_a, ek_pk).unwrap();

    // Exchange some messages.
    let enc = alice.encrypt(b"hello").unwrap();
    bob.decrypt(&enc.header, &enc.ciphertext).unwrap();
    let enc = bob.encrypt(b"world").unwrap();
    alice.decrypt(&enc.header, &enc.ciphertext).unwrap();

    // Serialize the ratchet state.
    let alice_bytes = alice.to_bytes().unwrap().0;

    // Encrypt as a storage blob.
    let storage_key = StorageKey::new(1, random::random_array()).unwrap();
    let blob = encrypt_blob(&storage_key, &alice_bytes, "session-store", "alice-1", true).unwrap();

    // Decrypt the blob.
    let ring = StorageKeyRing::new(storage_key).unwrap();
    let decrypted = decrypt_blob(&ring, &blob, "session-store", "alice-1").unwrap();

    // Verify serialization fidelity: the round-tripped bytes must be identical
    // (excluding the epoch field, which advances on every to_bytes() call).
    let alice2 = RatchetState::from_bytes(&decrypted).unwrap();
    let alice2_bytes = alice2.to_bytes().unwrap().0;
    assert_eq!(alice_bytes[0], alice2_bytes[0], "version must match");
    assert_eq!(
        alice_bytes[9..],
        alice2_bytes[9..],
        "ratchet state serialization round-trip must be identical (post-epoch fields)",
    );
    // Epoch must advance by exactly 1 per to_bytes() call.
    let epoch1 = u64::from_be_bytes(alice_bytes[1..9].try_into().unwrap());
    let epoch2 = u64::from_be_bytes(alice2_bytes[1..9].try_into().unwrap());
    assert_eq!(epoch2, epoch1 + 1);

    // Verify continued operation after deserialization.
    let mut alice2 = RatchetState::from_bytes(&decrypted).unwrap();
    let enc = alice2.encrypt(b"after storage").unwrap();
    let pt = bob.decrypt(&enc.header, &enc.ciphertext).unwrap();
    assert_eq!(&*pt, b"after storage");
}