SparsePostQuantumRatchet/src/chain.rs
gram-signal 63d9b93634
Some checks failed
hax / fstar-type-checking (push) Has been cancelled
proverif / proofs (push) Has been cancelled
CI / cargo test (push) Has been cancelled
CI / MSRV (push) Has been cancelled
CI / cross test polynomial i586 (push) Has been cancelled
CI / cross test polynomial i686 (push) Has been cancelled
CI / cross test polynomial aarch64 (push) Has been cancelled
Improve Hax proofs.
This PR improves the hax proofs by removing some admitted parts. It focuses on the v1 module and chain.rs. The main improvements come from the use of refinement types, and from some hax improvements around loops and more (not released yet which is why this is a draft PR for now).

Co-authored-by: Maxime Buyse <maxime@cryspen.com>
Co-authored-by: maximebuyse <45398004+maximebuyse@users.noreply.github.com>
2026-02-20 13:27:45 -08:00

707 lines
25 KiB
Rust

// Copyright 2025 Signal Messenger, LLC
// SPDX-License-Identifier: AGPL-3.0-only
use super::{Direction, Epoch, EpochSecret, Error};
use crate::kdf;
use crate::proto::pq_ratchet as pqrpb;
use crate::proto::pq_ratchet::ChainParams as ChainParamsPB;
use std::cmp::Ordering;
use std::collections::VecDeque;
/// Parameters for controlling the behavior of PQR key chains.
/// It's recommended to use the Default API for overriding values,
/// as future values may be added to this struct, and Default allows
/// them to be added in a backwards-compatible fashion.
/// IE: let params = ChainParams{max_jump: 10, ..Default::default()};
#[derive(Clone, Copy)]
pub struct ChainParams {
/// Disallow requesting a key that is more than MAX_JUMP ahead of `ctr`.
/// If zero, defaults to the current library-compiled default value.
pub max_jump: u32,
/// Keep around keys back to at least `ctr - MAX_OOO_KEYS`, in case an out-of-order
/// message comes in. Messages older than this that arrive out-of-order
/// will not be able to be decrypted and will return Error::KeyTrimmed.
/// If zero, defaults to the current library-compiled default value.
pub max_ooo_keys: u32,
}
impl Default for ChainParams {
fn default() -> Self {
DEFAULT_CHAIN_PARAMS
}
}
const DEFAULT_CHAIN_PARAMS: ChainParams = ChainParams {
max_jump: 25_000,
max_ooo_keys: 2_000,
};
impl ChainParams {
pub(crate) fn into_pb(self) -> ChainParamsPB {
ChainParamsPB {
max_jump: if self.max_jump == DEFAULT_CHAIN_PARAMS.max_jump {
0
} else {
self.max_jump
},
max_ooo_keys: if self.max_ooo_keys == DEFAULT_CHAIN_PARAMS.max_ooo_keys {
0
} else {
self.max_ooo_keys
},
}
}
/// Public wrapper for test utilities and benchmarks.
/// For internal use, call `into_pb` directly.
#[cfg(feature = "test-utils")]
pub fn into_pb_test(self) -> ChainParamsPB {
self.into_pb()
}
}
impl ChainParamsPB {
// The Default for protobufs is to have everything be zeros. Therefore,
// we use some getter functions locally to apply sane defaults to values that
// are not explicitly set.
fn max_jump_or_default(&self) -> u32 {
if self.max_jump > 0 {
self.max_jump
} else {
DEFAULT_CHAIN_PARAMS.max_jump
}
}
fn max_ooo_keys_or_default(&self) -> u32 {
if self.max_ooo_keys > 0 {
self.max_ooo_keys
} else {
DEFAULT_CHAIN_PARAMS.max_ooo_keys
}
}
/// When the size of our key history exceeds this amount, we run a
/// garbage collection on it.
fn trim_size(&self) -> usize {
let max_ooo = self.max_ooo_keys_or_default() as usize;
hax_lib::assume!(max_ooo < 390451572);
max_ooo * 11 / 10 + 1
}
}
struct KeyHistory {
// Keys are stored as [u8; 4][u8; 32], where the first is the index as a BE32
// and the second is the key.
// data.len() <= KEY_SIZE*TRIM_SIZE
data: Vec<u8>,
}
/// ChainEpochDirection keeps track of keys related to either half of send/recv.
struct ChainEpochDirection {
ctr: u32,
// next.len() == 32
next: Vec<u8>,
prev: KeyHistory,
}
/// ChainEpoch keeps state on a single epoch's keys.
struct ChainEpoch {
send: ChainEpochDirection,
recv: ChainEpochDirection,
}
/// Chain keeps track of keys for all epochs.
pub struct Chain {
dir: Direction,
current_epoch: Epoch,
send_epoch: Epoch,
links: VecDeque<ChainEpoch>, // stores [link[current_epoch-N] .. link[current_epoch]]
// next_root.len() == 32
next_root: Vec<u8>,
params: pqrpb::ChainParams,
}
/// We keep around this many epochs to keep prior to the current send epoch.
/// We'll always keep the send epoch and any subsequent epochs.
const EPOCHS_TO_KEEP_PRIOR_TO_SEND_EPOCH: usize = 1;
#[hax_lib::attributes]
impl KeyHistory {
/// Size in bytes of a single key stored within a KeyHistory.
const KEY_SIZE: usize = 4 + 32;
fn new() -> Self {
Self {
data: Vec::with_capacity(Self::KEY_SIZE * 2),
}
}
#[hax_lib::requires(_params.trim_size() < 119304647 && self.data.len() <= KeyHistory::KEY_SIZE * _params.trim_size())]
fn add(&mut self, k: (u32, [u8; 32]), _params: &pqrpb::ChainParams) {
self.data.extend_from_slice(&k.0.to_be_bytes()[..]);
self.data.extend_from_slice(&k.1[..]);
}
#[hax_lib::opaque] // ordering of slices needed
fn gc(&mut self, current_key: u32, params: &pqrpb::ChainParams) {
if self.data.len() >= params.trim_size() * Self::KEY_SIZE {
// We assume that k.0 is the highest key index we've ever seen, and base
// our trimming on that.
assert!(current_key >= params.max_ooo_keys_or_default());
let trim_horizon = &(current_key - params.max_ooo_keys_or_default()).to_be_bytes()[..];
// This does a single O(n) pass over our list, dropping all keys less than
// our computed trim horizon.
let mut i: usize = 0;
while i < self.data.len() {
if matches!(
trim_horizon.cmp(&self.data[i..i + 4]),
std::cmp::Ordering::Greater
) {
self.remove(i, params);
// Don't advance i here; we could have replaced the value there-in
// with another old key.
} else {
i += Self::KEY_SIZE;
}
}
}
}
fn clear(&mut self) {
self.data.clear();
}
#[hax_lib::requires(my_array_index <= self.data.len() && _params.trim_size() < 119304647 && self.data.len() <= KeyHistory::KEY_SIZE * _params.trim_size())]
fn remove(&mut self, mut my_array_index: usize, _params: &pqrpb::ChainParams) {
if my_array_index + Self::KEY_SIZE < self.data.len() {
let new_end = self.data.len() - Self::KEY_SIZE;
self.data.copy_within(new_end.., my_array_index);
my_array_index = new_end;
}
self.data.truncate(my_array_index);
}
#[hax_lib::opaque] // needs a model of step_by loop with return
fn get(
&mut self,
at: u32,
current_ctr: u32,
params: &pqrpb::ChainParams,
) -> Result<Vec<u8>, Error> {
assert_eq!(self.data.len() % Self::KEY_SIZE, 0);
if at + (params.max_ooo_keys_or_default()) < current_ctr {
// We've already discarded this because it's too old.
return Err(Error::KeyTrimmed(at));
}
let want = at.to_be_bytes();
for i in (0..self.data.len()).step_by(Self::KEY_SIZE) {
if self.data[i..i + 4] == want {
let out = self.data[i + 4..i + Self::KEY_SIZE].to_vec();
self.remove(i, params);
return Ok(out);
}
}
// This is a key we should have and we don't, so it must have already
// been requested.
Err(Error::KeyAlreadyRequested(at))
}
}
#[hax_lib::attributes]
impl ChainEpochDirection {
fn new(k: &[u8]) -> Self {
Self {
ctr: 0,
prev: KeyHistory::new(),
next: k.to_vec(),
}
}
#[hax_lib::requires(self.next.len() > 0 && self.ctr < u32::MAX)]
fn next_key(&mut self) -> (u32, Vec<u8>) {
let (idx, key) = Self::next_key_internal(&mut self.next, &mut self.ctr);
(idx, key.to_vec())
}
#[hax_lib::requires(next.len() > 0 && *ctr < u32::MAX)]
#[hax_lib::ensures(|_| *future(ctr) == ctr + 1)]
fn next_key_internal(next: &mut [u8], ctr: &mut u32) -> (u32, [u8; 32]) {
assert!(!next.is_empty());
*ctr += 1;
let mut genr8r = [0u8; 64];
kdf::hkdf_to_slice(
&[0u8; 32], // 32 is the hash output length
&*next,
[
ctr.to_be_bytes().as_slice(),
b"Signal PQ Ratchet V1 Chain Next",
]
.concat()
.as_slice(),
&mut genr8r,
);
next.copy_from_slice(&genr8r[..32]);
(*ctr, genr8r[32..].try_into().expect("correct size"))
}
fn key(&mut self, at: u32, params: &pqrpb::ChainParams) -> Result<Vec<u8>, Error> {
match at.cmp(&self.ctr) {
Ordering::Greater => {
if at - self.ctr > params.max_jump_or_default() {
return Err(Error::KeyJump(self.ctr, at));
}
}
Ordering::Less => {
return self.prev.get(at, self.ctr, params);
}
Ordering::Equal => {
// We've already returned this key once, we won't do it again.
return Err(Error::KeyAlreadyRequested(at));
}
}
hax_lib::assume!(
params.max_ooo_keys_or_default() < 390451572 && self.ctr <= u32::MAX - 390451572
);
if at > self.ctr + params.max_ooo_keys_or_default() {
// We're about to make all currently-held keys obsolete - just remove
// them all.
self.prev.clear();
}
while at > self.ctr + 1 {
hax_lib::loop_invariant!(self.ctr < u32::MAX);
hax_lib::loop_decreases!(u32::MAX - self.ctr);
hax_lib::assume!(self.next.len() > 0);
let k = Self::next_key_internal(&mut self.next, &mut self.ctr);
hax_lib::assume!(
params.max_ooo_keys_or_default() < 390451572 && self.ctr <= u32::MAX - 390451572
);
// Only add keys into our history if we're not going to immediately GC them.
if self.ctr + params.max_ooo_keys_or_default() >= at {
hax_lib::assume!(
params.trim_size() < 119304647
&& self.prev.data.len() <= KeyHistory::KEY_SIZE * params.trim_size()
);
self.prev.add(k, params);
}
}
// After we've potentially added some new keys, see if there's any we
// want to throw away.
self.prev.gc(self.ctr, params);
hax_lib::assume!(self.next.len() > 0);
Ok(Self::next_key_internal(&mut self.next, &mut self.ctr)
.1
.to_vec())
}
fn into_pb(self) -> pqrpb::chain::epoch::EpochDirection {
pqrpb::chain::epoch::EpochDirection {
ctr: self.ctr,
next: self.next,
prev: self.prev.data,
}
}
fn from_pb(pb: pqrpb::chain::epoch::EpochDirection) -> Result<Self, Error> {
Ok(Self {
ctr: pb.ctr,
next: pb.next,
prev: KeyHistory { data: pb.prev },
})
}
fn clear_next(&mut self) {
self.next.clear();
}
}
#[hax_lib::attributes]
impl Chain {
#[hax_lib::requires(genr8r.len() == 96)]
fn ced_for_direction(genr8r: &[u8], dir: &Direction) -> ChainEpochDirection {
ChainEpochDirection::new(match dir {
Direction::A2B => &genr8r[32..64],
Direction::B2A => &genr8r[64..96],
})
}
pub fn new(initial_key: &[u8], dir: Direction, params: ChainParamsPB) -> Result<Self, Error> {
let mut genr8r = [0u8; 96];
kdf::hkdf_to_slice(
&[0u8; 32],
initial_key,
b"Signal PQ Ratchet V1 Chain Start",
&mut genr8r,
);
Ok(Self {
dir,
current_epoch: 0,
send_epoch: 0,
links: VecDeque::from([ChainEpoch {
send: Self::ced_for_direction(&genr8r, &dir),
recv: Self::ced_for_direction(&genr8r, &dir.switch()),
}]),
next_root: genr8r[0..32].to_vec(),
params,
})
}
pub fn add_epoch(&mut self, epoch_secret: EpochSecret) {
// This assume could be turned into a precondition but it uses private fields
hax_lib::assume!(
self.current_epoch < u64::MAX && epoch_secret.epoch == self.current_epoch + 1
);
assert!(epoch_secret.epoch == self.current_epoch + 1);
let mut genr8r = [0u8; 96];
kdf::hkdf_to_slice(
&self.next_root,
&epoch_secret.secret,
b"Signal PQ Ratchet V1 Chain Add Epoch",
&mut genr8r,
);
self.current_epoch = epoch_secret.epoch;
self.next_root = genr8r[0..32].to_vec();
self.links.push_back(ChainEpoch {
send: Self::ced_for_direction(&genr8r, &self.dir),
recv: Self::ced_for_direction(&genr8r, &self.dir.switch()),
});
}
#[hax_lib::ensures(|res| if let Ok(v) = res {v < self.links.len()} else {true})]
fn epoch_idx(&mut self, epoch: Epoch) -> Result<usize, Error> {
if epoch > self.current_epoch {
return Err(Error::EpochOutOfRange(epoch));
}
let back = (self.current_epoch - epoch) as usize;
let links = self.links.len();
if back >= links {
return Err(Error::EpochOutOfRange(epoch));
}
Ok(links - 1 - back)
}
pub fn send_key(&mut self, epoch: Epoch) -> Result<(u32, Vec<u8>), Error> {
if epoch < self.send_epoch {
return Err(Error::SendKeyEpochDecreased(self.send_epoch, epoch));
}
let mut epoch_index = self.epoch_idx(epoch)?;
if self.send_epoch != epoch {
self.send_epoch = epoch;
while epoch_index > EPOCHS_TO_KEEP_PRIOR_TO_SEND_EPOCH {
hax_lib::loop_decreases!(epoch_index);
self.links.pop_front();
epoch_index -= 1;
}
for i in 0..epoch_index {
hax_lib::assume!(i < self.links.len());
self.links[i].send.clear_next();
}
}
hax_lib::assume!(
epoch_index < self.links.len()
&& self.links[epoch_index].send.next.len() > 0
&& self.links[epoch_index].send.ctr < u32::MAX
);
Ok(self.links[epoch_index].send.next_key())
}
pub fn recv_key(&mut self, epoch: Epoch, index: u32) -> Result<Vec<u8>, Error> {
let epoch_index = self.epoch_idx(epoch)?;
self.links[epoch_index].recv.key(index, &self.params)
}
#[hax_lib::opaque] // into_iter and map
pub(crate) fn into_pb(self) -> pqrpb::Chain {
pqrpb::Chain {
direction: self.dir.into(),
current_epoch: self.current_epoch,
send_epoch: self.send_epoch,
links: self
.links
.into_iter()
.map(|link| pqrpb::chain::Epoch {
send: Some(link.send.into_pb()),
recv: Some(link.recv.into_pb()),
})
.collect::<Vec<_>>(),
next_root: self.next_root,
params: Some(self.params),
}
}
#[hax_lib::opaque] // into_iter and map
pub(crate) fn from_pb(pb: pqrpb::Chain) -> Result<Self, Error> {
Ok(Self {
dir: pb.direction.try_into().map_err(|_| Error::StateDecode)?,
current_epoch: pb.current_epoch,
send_epoch: pb.send_epoch,
next_root: pb.next_root,
links: pb
.links
.into_iter()
.map(|link| {
Ok::<ChainEpoch, Error>(ChainEpoch {
send: ChainEpochDirection::from_pb(link.send.ok_or(Error::StateDecode)?)?,
recv: ChainEpochDirection::from_pb(link.recv.ok_or(Error::StateDecode)?)?,
})
})
.collect::<Result<VecDeque<_>, _>>()?,
params: pb.params.ok_or(Error::StateDecode)?,
})
}
}
#[cfg(test)]
mod test {
use super::*;
use crate::{Direction, EpochSecret, Error};
use proptest::prelude::*;
use rand::TryRngCore;
use rand::seq::SliceRandom;
#[test]
fn directions_match() {
let mut a2b = Chain::new(b"1", Direction::A2B, ChainParams::default().into_pb()).unwrap();
let mut b2a = Chain::new(b"1", Direction::B2A, ChainParams::default().into_pb()).unwrap();
let sk1 = a2b.send_key(0).unwrap();
assert_eq!(sk1.0, 1);
assert_eq!(sk1.1, b2a.recv_key(0, 1).unwrap());
a2b.add_epoch(EpochSecret {
epoch: 1,
secret: vec![2],
});
b2a.add_epoch(EpochSecret {
epoch: 1,
secret: vec![2],
});
let sk2 = a2b.send_key(1).unwrap();
assert_eq!(sk2.0, 1);
assert_eq!(sk2.1, b2a.recv_key(1, 1).unwrap());
for _i in 2..10 {
a2b.send_key(1).unwrap();
}
let sk3 = a2b.send_key(1).unwrap();
assert_eq!(sk3.0, 10);
assert_eq!(sk3.1, b2a.recv_key(1, 10).unwrap());
}
#[test]
fn previously_returned_key() {
let mut a2b = Chain::new(b"1", Direction::A2B, ChainParams::default().into_pb()).unwrap();
a2b.recv_key(0, 2).expect("should get key first time");
assert!(matches!(
a2b.recv_key(0, 2),
Err(Error::KeyAlreadyRequested(2))
));
}
#[test]
fn very_old_keys_are_trimmed() {
let params = ChainParams {
max_jump: 10,
max_ooo_keys: 10,
}
.into_pb();
let mut a2b = Chain::new(b"1", Direction::A2B, params).unwrap();
a2b.recv_key(0, 10).expect("should allow this jump");
a2b.recv_key(0, 12).expect("should allow progression");
assert!(matches!(a2b.recv_key(0, 1), Err(Error::KeyTrimmed(1))));
}
#[test]
fn out_of_order_keys() {
let max_ooo = DEFAULT_CHAIN_PARAMS.max_ooo_keys;
let mut a2b = Chain::new(b"1", Direction::A2B, ChainParams::default().into_pb()).unwrap();
let mut b2a = Chain::new(b"1", Direction::B2A, ChainParams::default().into_pb()).unwrap();
let mut keys = Vec::with_capacity(max_ooo as usize);
for _i in 0..(max_ooo as usize) {
keys.push(a2b.send_key(0).unwrap());
}
let mut rng = rand::rngs::OsRng.unwrap_err();
keys.shuffle(&mut rng);
for (idx, key) in keys {
assert_eq!(b2a.recv_key(0, idx).unwrap(), key);
}
}
#[test]
fn clear_old_send_keys() {
let mut a2b = Chain::new(b"1", Direction::A2B, ChainParams::default().into_pb()).unwrap();
a2b.send_key(0).unwrap();
a2b.send_key(0).unwrap();
a2b.add_epoch(EpochSecret {
epoch: 1,
secret: vec![2],
});
a2b.send_key(1).unwrap();
assert!(matches!(
a2b.send_key(0).unwrap_err(),
Error::SendKeyEpochDecreased(1, 0)
));
}
#[derive(Clone, Debug)]
enum KeyHistoryAction {
AddNewNext,
AddNewSkip,
RequestStored(usize),
RequestNotStored(usize),
GarbageCollect,
}
impl KeyHistoryAction {
fn strategy() -> impl Strategy<Value = Self> {
proptest::prop_oneof![
Just(Self::AddNewNext),
Just(Self::AddNewSkip),
any::<usize>().prop_map(Self::RequestStored),
any::<usize>().prop_map(Self::RequestNotStored),
Just(Self::GarbageCollect),
]
}
}
#[test]
fn key_history_prop_test() {
let _ = env_logger::builder().is_test(true).try_init();
proptest!(|(actions in proptest::collection::vec(KeyHistoryAction::strategy(), ..25))| {
let mut kh = KeyHistory::new();
let mut stored = vec![];
let mut not_stored = vec![];
let mut ctr = 0u32;
let params = pqrpb::ChainParams {
max_ooo_keys: 3u32,
max_jump: 5u32,
};
log::debug!("========= STARTING =========");
for action in actions {
match action {
KeyHistoryAction::AddNewNext => {
log::debug!("adding {}", ctr);
kh.add((ctr, [1u8; 32]), &params);
stored.push(ctr);
ctr += 1;
}
KeyHistoryAction::AddNewSkip => {
log::debug!("skipping {}", ctr);
not_stored.push(ctr);
ctr += 1;
}
KeyHistoryAction::RequestStored(i) => {
if !stored.is_empty() {
let k = stored.swap_remove(i % stored.len());
log::debug!("requesting stored {}", k);
kh.get(k, ctr, &params).unwrap();
not_stored.push(k);
}
}
KeyHistoryAction::RequestNotStored(i) => {
if !not_stored.is_empty() {
let k = not_stored.swap_remove(i % not_stored.len());
log::debug!("requesting not stored {}", k);
kh.get(k, ctr, &params).unwrap_err();
}
}
KeyHistoryAction::GarbageCollect => {
log::debug!("gc at {}", ctr);
kh.gc(ctr, &params);
}
}
let mut fell_off = vec![];
(fell_off, stored) = stored.into_iter().partition(
|n| n + params.max_ooo_keys < ctr);
if !fell_off.is_empty() {
log::debug!("fell off: {:?}", fell_off);
}
not_stored.extend(fell_off.into_iter());
}
});
}
#[derive(Clone, Debug)]
enum CEDAction {
NextKey,
NextKeyAt(u32),
RequestStored(usize),
RequestNotStored(usize),
TooHigh,
}
impl CEDAction {
fn strategy() -> impl Strategy<Value = Self> {
proptest::prop_oneof![
Just(Self::NextKey),
any::<u32>().prop_map(Self::NextKeyAt),
any::<usize>().prop_map(Self::RequestStored),
any::<usize>().prop_map(Self::RequestNotStored),
Just(Self::TooHigh),
]
}
}
#[test]
fn ced_prop_test() {
let _ = env_logger::builder().is_test(true).try_init();
proptest!(|(actions in proptest::collection::vec(CEDAction::strategy(), ..25))| {
let mut ced = ChainEpochDirection::new(&[1u8; 32]);
let mut stored = vec![];
let mut not_stored = vec![0];
let mut ctr = 0u32;
let params = pqrpb::ChainParams {
max_ooo_keys: 3u32,
max_jump: 5u32,
};
log::debug!("========= STARTING =========");
for action in actions {
match action {
CEDAction::NextKey => {
ctr += 1;
log::debug!("next_key {}", ctr);
let k = ced.next_key().0;
not_stored.push(k);
}
CEDAction::NextKeyAt(i) => {
let jump = i % params.max_jump;
ctr += 1;
for at in ctr..ctr+jump {
stored.push(at);
}
ctr += jump;
log::debug!("next_key_at {}", ctr);
ced.key(ctr, &params).unwrap();
not_stored.push(ctr);
}
CEDAction::RequestStored(i) => {
if !stored.is_empty() {
let k = stored.swap_remove(i % stored.len());
log::debug!("requesting stored {}", k);
ced.key(k, &params).unwrap();
not_stored.push(k);
}
}
CEDAction::RequestNotStored(i) => {
if !not_stored.is_empty() {
let k = not_stored.swap_remove(i % not_stored.len());
log::debug!("requesting not stored {}", k);
ced.key(k, &params).unwrap_err();
}
}
CEDAction::TooHigh => {
let high = ctr + params.max_jump + 1;
log::debug!("too high {}", high);
ced.key(high, &params).unwrap_err();
}
}
let mut fell_off = vec![];
(fell_off, stored) = stored.into_iter().partition(
|n| n + params.max_ooo_keys < ctr);
if !fell_off.is_empty() {
log::debug!("fell off: {:?}", fell_off);
}
not_stored.extend(fell_off.into_iter());
}
});
}
}