#![allow(deprecated)]
use crate::kv_reduce::KvOp;
use aes_gcm::aead::{Aead, AeadCore, KeyInit, OsRng, Payload};
use aes_gcm::{Aes256Gcm, Nonce};
use k256::ecdsa::SigningKey;
use sha3::{Digest, Keccak256};
const OP_VERSION: u8 = 1;
const NONCE_LEN: usize = 12;
const EPH_PUB_LEN: usize = 33;
const GCM_TAG_LEN: usize = 16;
pub fn derive_room_key(identity_secret: &[u8; 32], room_id: u64) -> [u8; 32] {
let mut h = Keccak256::new();
h.update(b"localharness/sessionroom/v1/key");
h.update(identity_secret);
h.update(room_id.to_be_bytes());
h.finalize().into()
}
pub fn room_recipient(room_id: u64) -> [u8; 20] {
let mut h = Keccak256::new();
h.update(b"localharness/sessionroom/v1/recipient");
h.update(room_id.to_be_bytes());
let d = h.finalize();
let mut addr = [0u8; 20];
addr.copy_from_slice(&d[12..]);
addr
}
pub fn encode_op(op: &KvOp) -> Option<Vec<u8>> {
let key = op.key.as_bytes();
if key.len() > u16::MAX as usize {
return None;
}
if let Some(v) = &op.value {
if v.len() > u32::MAX as usize {
return None;
}
}
let mut out = Vec::with_capacity(1 + 2 + key.len() + 1 + 4 + 8 + 20 + 8);
out.push(OP_VERSION);
out.extend_from_slice(&(key.len() as u16).to_be_bytes());
out.extend_from_slice(key);
match &op.value {
Some(v) => {
out.push(1);
out.extend_from_slice(&(v.len() as u32).to_be_bytes());
out.extend_from_slice(v);
}
None => out.push(0),
}
out.extend_from_slice(&op.lamport.to_be_bytes());
out.extend_from_slice(&op.writer);
out.extend_from_slice(&op.ts.to_be_bytes());
Some(out)
}
pub fn decode_op(bytes: &[u8]) -> Option<KvOp> {
let mut i = 0usize;
let take = |i: &mut usize, n: usize| -> Option<&[u8]> {
let end = i.checked_add(n)?;
let s = bytes.get(*i..end)?;
*i = end;
Some(s)
};
if *take(&mut i, 1)?.first()? != OP_VERSION {
return None;
}
let key_len = u16::from_be_bytes(take(&mut i, 2)?.try_into().ok()?) as usize;
let key = String::from_utf8(take(&mut i, key_len)?.to_vec()).ok()?;
let value = match take(&mut i, 1)?[0] {
0 => None,
1 => {
let val_len = u32::from_be_bytes(take(&mut i, 4)?.try_into().ok()?) as usize;
Some(take(&mut i, val_len)?.to_vec())
}
_ => return None,
};
let lamport = u64::from_be_bytes(take(&mut i, 8)?.try_into().ok()?);
let writer: [u8; 20] = take(&mut i, 20)?.try_into().ok()?;
let ts = u64::from_be_bytes(take(&mut i, 8)?.try_into().ok()?);
if i != bytes.len() {
return None; }
Some(KvOp {
key,
value,
lamport,
writer,
ts,
})
}
pub fn seal_op(
op: &KvOp,
k_room: &[u8; 32],
writer_key: &SigningKey,
room_id: u64,
) -> Option<Vec<u8>> {
let cipher = Aes256Gcm::new_from_slice(k_room).ok()?;
let nonce = Aes256Gcm::generate_nonce(&mut OsRng);
let ct = cipher.encrypt(&nonce, encode_op(op)?.as_slice()).ok()?;
let mut sealed = Vec::with_capacity(NONCE_LEN + ct.len());
sealed.extend_from_slice(nonce.as_slice());
sealed.extend_from_slice(&ct);
Some(crate::signaling_seal::seal_envelope(
writer_key,
&room_recipient(room_id),
&sealed,
))
}
pub fn open_op(
blob: &[u8],
k_room: &[u8; 32],
writer_addr: &[u8; 20],
room_id: u64,
) -> Option<KvOp> {
let sealed = crate::signaling_seal::open_envelope(blob, writer_addr, &room_recipient(room_id))?;
if sealed.len() < NONCE_LEN {
return None;
}
let (nonce_bytes, ct) = sealed.split_at(NONCE_LEN);
let cipher = Aes256Gcm::new_from_slice(k_room).ok()?;
let plaintext = cipher.decrypt(Nonce::from_slice(nonce_bytes), ct).ok()?;
let op = decode_op(&plaintext)?;
(op.writer == *writer_addr).then_some(op)
}
const GRANT_AAD_TAG: &[u8] = b"localharness/sessionroom/v1/keygrant";
pub fn key_grant_seal(k_room: &[u8; 32], recipient_pubkey_sec1: &[u8]) -> Option<Vec<u8>> {
let (eph_pub, eph_signer) = crate::wallet::ephemeral_keypair();
let key = crate::wallet::ecdh_shared_key(&eph_signer, recipient_pubkey_sec1).ok()?;
let cipher = Aes256Gcm::new_from_slice(&key).ok()?;
let nonce = Aes256Gcm::generate_nonce(&mut OsRng);
let aad = grant_aad(&eph_pub);
let ct = cipher
.encrypt(
&nonce,
Payload {
msg: k_room.as_slice(),
aad: &aad,
},
)
.ok()?;
let mut out = Vec::with_capacity(eph_pub.len() + NONCE_LEN + ct.len());
out.extend_from_slice(&eph_pub);
out.extend_from_slice(nonce.as_slice());
out.extend_from_slice(&ct);
Some(out)
}
pub fn key_grant_open(sealed: &[u8], recipient_key: &SigningKey) -> Option<[u8; 32]> {
if sealed.len() < EPH_PUB_LEN + NONCE_LEN + GCM_TAG_LEN {
return None;
}
let (eph_pub, rest) = sealed.split_at(EPH_PUB_LEN);
let (nonce_bytes, ct) = rest.split_at(NONCE_LEN);
let key = crate::wallet::ecdh_shared_key(recipient_key, eph_pub).ok()?;
let cipher = Aes256Gcm::new_from_slice(&key).ok()?;
let aad = grant_aad(eph_pub);
let plaintext = cipher
.decrypt(
Nonce::from_slice(nonce_bytes),
Payload { msg: ct, aad: &aad },
)
.ok()?;
plaintext.try_into().ok()
}
fn grant_aad(eph_pub: &[u8]) -> Vec<u8> {
let mut aad = Vec::with_capacity(GRANT_AAD_TAG.len() + eph_pub.len());
aad.extend_from_slice(GRANT_AAD_TAG);
aad.extend_from_slice(eph_pub);
aad
}
#[cfg(test)]
mod tests {
use super::*;
fn key(b: u8) -> SigningKey {
SigningKey::from_slice(&[b; 32]).unwrap()
}
fn sample(writer: [u8; 20]) -> KvOp {
KvOp {
key: "score".into(),
value: Some(b"42".to_vec()),
lamport: 7,
writer,
ts: 1234,
}
}
#[test]
fn encode_decode_round_trip_value_and_tombstone() {
let v = sample([9u8; 20]);
assert_eq!(decode_op(&encode_op(&v).unwrap()).unwrap(), v);
let t = KvOp {
value: None,
..sample([3u8; 20])
};
assert_eq!(decode_op(&encode_op(&t).unwrap()).unwrap(), t);
let mut buf = encode_op(&v).unwrap();
buf.push(0xff);
assert!(decode_op(&buf).is_none());
assert!(decode_op(&encode_op(&v).unwrap()[..3]).is_none());
}
#[test]
fn encode_op_rejects_oversize_key_instead_of_truncating() {
let big = KvOp {
key: "k".repeat(u16::MAX as usize + 1),
..sample([1u8; 20])
};
assert!(encode_op(&big).is_none(), "oversize key must not encode");
assert!(
seal_op(&big, &[7u8; 32], &key(1), 1).is_none(),
"seal must propagate the rejection"
);
let edge = KvOp {
key: "k".repeat(u16::MAX as usize),
..sample([1u8; 20])
};
assert_eq!(decode_op(&encode_op(&edge).unwrap()).unwrap(), edge);
}
#[test]
fn derive_room_key_deterministic_and_room_unique() {
let secret = [5u8; 32];
assert_eq!(derive_room_key(&secret, 1), derive_room_key(&secret, 1));
assert_ne!(derive_room_key(&secret, 1), derive_room_key(&secret, 2));
assert_ne!(derive_room_key(&secret, 1), derive_room_key(&[6u8; 32], 1));
assert_ne!(room_recipient(1), room_recipient(2));
}
#[test]
fn seal_open_round_trip() {
let wk = key(1);
let waddr = crate::wallet::address(&wk);
let k = derive_room_key(&[7u8; 32], 42);
let op = sample(waddr);
let blob = seal_op(&op, &k, &wk, 42).unwrap();
assert_eq!(open_op(&blob, &k, &waddr, 42).unwrap(), op);
}
#[test]
fn wrong_key_rejected() {
let wk = key(1);
let waddr = crate::wallet::address(&wk);
let op = sample(waddr);
let blob = seal_op(&op, &derive_room_key(&[7u8; 32], 42), &wk, 42).unwrap();
assert!(open_op(&blob, &derive_room_key(&[8u8; 32], 42), &waddr, 42).is_none());
}
#[test]
fn cross_room_replay_rejected() {
let wk = key(1);
let waddr = crate::wallet::address(&wk);
let k = derive_room_key(&[7u8; 32], 42);
let blob = seal_op(&sample(waddr), &k, &wk, 42).unwrap();
assert!(open_op(&blob, &k, &waddr, 43).is_none());
}
#[test]
fn wrong_writer_rejected() {
let wk = key(1);
let waddr = crate::wallet::address(&wk);
let k = derive_room_key(&[7u8; 32], 42);
let blob = seal_op(&sample(waddr), &k, &wk, 42).unwrap();
let other = crate::wallet::address(&key(2));
assert!(open_op(&blob, &k, &other, 42).is_none());
}
#[test]
fn tampered_blob_rejected() {
let wk = key(1);
let waddr = crate::wallet::address(&wk);
let k = derive_room_key(&[7u8; 32], 42);
let mut blob = seal_op(&sample(waddr), &k, &wk, 42).unwrap();
let last = blob.len() - 1;
blob[last] ^= 0xff;
assert!(open_op(&blob, &k, &waddr, 42).is_none());
}
#[test]
fn key_grant_round_trip() {
let member = key(11);
let member_pub = crate::wallet::pubkey_compressed(&member);
let k_room = [0x5Au8; 32];
let grant = key_grant_seal(&k_room, &member_pub).unwrap();
assert_eq!(grant.len(), EPH_PUB_LEN + NONCE_LEN + 32 + GCM_TAG_LEN);
assert_ne!(&grant[..EPH_PUB_LEN], member_pub.as_slice());
assert_eq!(key_grant_open(&grant, &member).unwrap(), k_room);
}
#[test]
fn key_grant_works_with_uncompressed_recipient_pubkey() {
use k256::ecdsa::VerifyingKey;
let member = key(12);
let uncompressed = VerifyingKey::from(&member)
.to_encoded_point(false)
.as_bytes()
.to_vec();
let k_room = [0x7Cu8; 32];
let grant = key_grant_seal(&k_room, &uncompressed).unwrap();
assert_eq!(key_grant_open(&grant, &member).unwrap(), k_room);
}
#[test]
fn key_grant_each_seal_uses_fresh_ephemeral() {
let member = key(13);
let member_pub = crate::wallet::pubkey_compressed(&member);
let k_room = [0x33u8; 32];
let g1 = key_grant_seal(&k_room, &member_pub).unwrap();
let g2 = key_grant_seal(&k_room, &member_pub).unwrap();
assert_ne!(g1, g2);
assert_ne!(&g1[..EPH_PUB_LEN], &g2[..EPH_PUB_LEN]); assert_eq!(key_grant_open(&g1, &member).unwrap(), k_room);
assert_eq!(key_grant_open(&g2, &member).unwrap(), k_room);
}
#[test]
fn key_grant_wrong_key_rejected() {
let member_a = key(14);
let member_b = key(15);
let a_pub = crate::wallet::pubkey_compressed(&member_a);
let grant = key_grant_seal(&[0x01u8; 32], &a_pub).unwrap();
assert!(key_grant_open(&grant, &member_b).is_none());
assert!(key_grant_open(&grant, &member_a).is_some());
}
#[test]
fn key_grant_tampered_any_byte_rejected() {
let member = key(16);
let member_pub = crate::wallet::pubkey_compressed(&member);
let grant = key_grant_seal(&[0x9Eu8; 32], &member_pub).unwrap();
for i in 0..grant.len() {
let mut t = grant.clone();
t[i] ^= 0x01;
assert!(
key_grant_open(&t, &member).is_none(),
"tampered grant byte {i} must be rejected"
);
}
}
#[test]
fn key_grant_malformed_inputs_rejected() {
let member = key(17);
assert!(key_grant_open(b"", &member).is_none());
assert!(key_grant_open(&[0u8; EPH_PUB_LEN], &member).is_none());
assert!(key_grant_open(&[0u8; EPH_PUB_LEN + NONCE_LEN], &member).is_none());
assert!(key_grant_open(&[0u8; EPH_PUB_LEN + NONCE_LEN + GCM_TAG_LEN], &member).is_none());
}
#[test]
fn key_grant_bad_recipient_pubkey_returns_none() {
assert!(key_grant_seal(&[0x11u8; 32], &[0u8; 33]).is_none());
assert!(key_grant_seal(&[0x11u8; 32], b"").is_none());
}
#[test]
fn granted_key_decrypts_room_ops_end_to_end() {
let creator = crate::wallet::generate().signer.clone();
let member = key(18);
let member_pub = crate::wallet::pubkey_compressed(&member);
let k_room: [u8; 32] = {
let mut k = [0u8; 32];
for (i, b) in k.iter_mut().enumerate() {
*b = (i as u8).wrapping_mul(7).wrapping_add(3);
}
k
};
let grant = key_grant_seal(&k_room, &member_pub).unwrap();
let recovered = key_grant_open(&grant, &member).unwrap();
assert_eq!(recovered, k_room);
let waddr = crate::wallet::address(&creator);
let op = sample(waddr);
let blob = seal_op(&op, &k_room, &creator, 99).unwrap();
assert_eq!(open_op(&blob, &recovered, &waddr, 99).unwrap(), op);
assert!(open_op(&blob, &derive_room_key(&[0u8; 32], 99), &waddr, 99).is_none());
}
#[test]
fn key_grant_does_not_leak_k_room_plaintext() {
let member = key(19);
let member_pub = crate::wallet::pubkey_compressed(&member);
let k_room = [0xC4u8; 32];
let grant = key_grant_seal(&k_room, &member_pub).unwrap();
assert!(
!grant.windows(k_room.len()).any(|w| w == k_room),
"raw K_room leaked into the grant"
);
}
}