use alloc::{format, string::String, vec::Vec};
use mls_rs_core::{
error::IntoAnyError,
group::{EpochRecord, GroupState, GroupStateStorage},
key_package::{KeyPackageData, KeyPackageStorage},
};
use mls_rs_core::{
identity::BasicCredential,
mls_rs_codec::{MlsDecode, MlsEncode},
};
use zeroize::Zeroizing;
use crate::shared::Arc;
use crate::store::{LeafStore, StoreError, KEY_TYPE_GROUP_EPOCH, KEY_TYPE_GROUP_STATE};
use crate::store::{KEY_TYPE_KEY_PACKAGE, KEY_TYPE_PEER};
impl IntoAnyError for StoreError {}
pub(crate) const PRIOR_EPOCH_RETENTION: u64 = 3;
pub(crate) fn hex(bytes: &[u8]) -> String {
let mut out = String::with_capacity(bytes.len() * 2);
for b in bytes {
const DIGITS: &[u8; 16] = b"0123456789abcdef";
out.push(DIGITS[(b >> 4) as usize] as char);
out.push(DIGITS[(b & 0x0f) as usize] as char);
}
out
}
fn state_key(group_id: &[u8]) -> String {
hex(group_id)
}
fn epoch_key(group_id: &[u8], epoch_id: u64) -> String {
format!("{}:{}", hex(group_id), epoch_id)
}
fn max_epoch_key(group_id: &[u8]) -> String {
format!("{}:max", hex(group_id))
}
const STATE_MARKER_LEN: usize = 9;
fn encode_state_entry(marker: Option<u64>, state: &[u8]) -> Vec<u8> {
let mut out = Vec::with_capacity(STATE_MARKER_LEN + state.len());
match marker {
Some(epoch) => {
out.push(1);
out.extend_from_slice(&epoch.to_be_bytes());
}
None => out.extend_from_slice(&[0u8; STATE_MARKER_LEN]),
}
out.extend_from_slice(state);
out
}
fn decode_state_entry(raw: &[u8]) -> Result<(Option<u64>, &[u8]), StoreError> {
if raw.len() < STATE_MARKER_LEN {
return Err(StoreError::Corrupt(format!(
"group state entry is {} bytes, too short to carry its epoch marker",
raw.len()
)));
}
let (header, state) = raw.split_at(STATE_MARKER_LEN);
let marker = match header[0] {
0 => None,
1 => {
let mut epoch = [0u8; 8];
epoch.copy_from_slice(&header[1..STATE_MARKER_LEN]);
Some(u64::from_be_bytes(epoch))
}
other => {
return Err(StoreError::Corrupt(format!(
"group state entry carries epoch marker flag {other}, expected 0 or 1"
)))
}
};
Ok((marker, state))
}
pub(crate) fn state_marker(store: &Arc<dyn LeafStore>, group_key: &str) -> Option<u64> {
let raw = store.load(KEY_TYPE_GROUP_STATE, group_key).ok()??;
decode_state_entry(&raw).ok().and_then(|(marker, _)| marker)
}
#[derive(Clone)]
pub(crate) struct GroupStateAdapter {
store: Arc<dyn LeafStore>,
}
impl GroupStateAdapter {
pub(crate) fn new(store: Arc<dyn LeafStore>) -> Self {
Self { store }
}
}
impl GroupStateStorage for GroupStateAdapter {
type Error = StoreError;
fn state(&self, group_id: &[u8]) -> Result<Option<Zeroizing<Vec<u8>>>, Self::Error> {
let Some(raw) = self
.store
.load(KEY_TYPE_GROUP_STATE, &state_key(group_id))?
else {
return Ok(None);
};
let raw = Zeroizing::new(raw);
let (_, state) = decode_state_entry(&raw)?;
Ok(Some(Zeroizing::new(state.to_vec())))
}
fn epoch(
&self,
group_id: &[u8],
epoch_id: u64,
) -> Result<Option<Zeroizing<Vec<u8>>>, Self::Error> {
Ok(self
.store
.load(KEY_TYPE_GROUP_EPOCH, &epoch_key(group_id, epoch_id))?
.map(Zeroizing::new))
}
fn write(
&mut self,
state: GroupState,
epoch_inserts: Vec<EpochRecord>,
epoch_updates: Vec<EpochRecord>,
) -> Result<(), Self::Error> {
let previous = self.max_epoch_id(&state.id)?;
let mut marker = previous;
for record in epoch_inserts.iter().chain(epoch_updates.iter()) {
self.store.store(
KEY_TYPE_GROUP_EPOCH,
&epoch_key(&state.id, record.id),
&record.data,
)?;
marker = Some(marker.map_or(record.id, |held: u64| held.max(record.id)));
}
self.store.store(
KEY_TYPE_GROUP_STATE,
&state_key(&state.id),
&encode_state_entry(marker, &state.data),
)?;
let advanced_to = if marker == previous { None } else { marker };
if let Some(high) = advanced_to {
let _ = self.store.store(
KEY_TYPE_GROUP_EPOCH,
&max_epoch_key(&state.id),
&high.to_be_bytes(),
);
}
for record in epoch_inserts.iter() {
if let Some(expired) = record.id.checked_sub(PRIOR_EPOCH_RETENTION) {
let _ = self
.store
.delete(KEY_TYPE_GROUP_EPOCH, &epoch_key(&state.id, expired));
}
}
Ok(())
}
fn max_epoch_id(&self, group_id: &[u8]) -> Result<Option<u64>, Self::Error> {
let Some(raw) = self
.store
.load(KEY_TYPE_GROUP_STATE, &state_key(group_id))?
else {
return Ok(None);
};
let raw = Zeroizing::new(raw);
let (marker, _) = decode_state_entry(&raw)?;
Ok(marker)
}
}
const MAX_UNSPENT_KEY_PACKAGES: usize = 4;
const KEY_PACKAGE_INDEX: &str = "__index__";
#[derive(Clone)]
pub(crate) struct KeyPackageAdapter {
store: Arc<dyn LeafStore>,
}
impl KeyPackageAdapter {
pub(crate) fn new(store: Arc<dyn LeafStore>) -> Self {
Self { store }
}
fn index(&self) -> Result<Vec<String>, StoreError> {
let Some(raw) = self.store.load(KEY_TYPE_KEY_PACKAGE, KEY_PACKAGE_INDEX)? else {
return Ok(Vec::new());
};
Ok(serde_json::from_slice(&raw).unwrap_or_default())
}
fn save_index(&self, index: &[String]) -> Result<(), StoreError> {
let encoded = serde_json::to_vec(index)
.map_err(|e| StoreError::Store(format!("cannot encode key package index: {e}")))?;
self.store
.store(KEY_TYPE_KEY_PACKAGE, KEY_PACKAGE_INDEX, &encoded)
}
pub(crate) fn erase(&self, reference: &str) -> Result<(), StoreError> {
self.store.delete(KEY_TYPE_KEY_PACKAGE, reference)?;
if let Ok(mut index) = self.index() {
if let Some(at) = index.iter().position(|held| held == reference) {
index.remove(at);
let _ = self.save_index(&index);
}
}
Ok(())
}
}
impl KeyPackageStorage for KeyPackageAdapter {
type Error = StoreError;
fn delete(&mut self, id: &[u8]) -> Result<(), Self::Error> {
self.erase(&hex(id))
}
fn insert(&mut self, id: Vec<u8>, pkg: KeyPackageData) -> Result<(), Self::Error> {
let encoded = pkg
.mls_encode_to_vec()
.map_err(|e| StoreError::Store(format!("cannot encode key package data: {e:?}")))?;
let key = hex(&id);
let mut index = self.index()?;
if !index.iter().any(|held| held == &key) {
index.push(key.clone());
}
let evicted: Vec<String> = if index.len() > MAX_UNSPENT_KEY_PACKAGES {
index
.drain(..index.len() - MAX_UNSPENT_KEY_PACKAGES)
.collect()
} else {
Vec::new()
};
for stale in evicted {
self.store.delete(KEY_TYPE_KEY_PACKAGE, &stale)?;
}
self.save_index(&index)?;
self.store.store(KEY_TYPE_KEY_PACKAGE, &key, &encoded)
}
fn get(&self, id: &[u8]) -> Result<Option<KeyPackageData>, Self::Error> {
let Some(raw) = self.store.load(KEY_TYPE_KEY_PACKAGE, &hex(id))? else {
return Ok(None);
};
let decoded = KeyPackageData::mls_decode(&mut &raw[..])
.map_err(|e| StoreError::Corrupt(format!("key package data does not decode: {e:?}")))?;
Ok(Some(decoded))
}
}
#[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize)]
pub(crate) struct PeerRecord {
#[serde(default)]
pub(crate) env_versions: Vec<u8>,
#[serde(default)]
pub(crate) wire_versions: Vec<u8>,
#[serde(default)]
pub(crate) key_package_sent: bool,
#[serde(default)]
pub(crate) key_package_ref: Option<String>,
#[serde(default)]
pub(crate) last_reset_ms: i64,
#[serde(default)]
pub(crate) acknowledged: Vec<String>,
}
pub(crate) const REMEMBERED_ACKS: usize = 4;
impl PeerRecord {
pub(crate) fn was_acknowledged(&self, id: &str) -> bool {
self.acknowledged.iter().any(|seen| seen == id)
}
pub(crate) fn remember_acknowledged(&mut self, id: &str) {
if self.was_acknowledged(id) {
return;
}
if self.acknowledged.len() >= REMEMBERED_ACKS {
self.acknowledged.remove(0);
}
self.acknowledged.push(String::from(id));
}
pub(crate) fn remember_reset(&mut self, signed_ms: i64) {
self.last_reset_ms = signed_ms;
}
pub(crate) fn reset_is_unspent(&self, signed_ms: i64) -> bool {
signed_ms > self.last_reset_ms
}
pub(crate) fn load(store: &Arc<dyn LeafStore>, peer: &str) -> Result<Option<Self>, StoreError> {
let Some(raw) = store.load(KEY_TYPE_PEER, peer)? else {
return Ok(None);
};
serde_json::from_slice(&raw)
.map(Some)
.map_err(|e| StoreError::Corrupt(format!("peer record does not decode: {e}")))
}
pub(crate) fn save(&self, store: &Arc<dyn LeafStore>, peer: &str) -> Result<(), StoreError> {
let encoded = serde_json::to_vec(self)
.map_err(|e| StoreError::Store(format!("cannot encode peer record: {e}")))?;
store.store(KEY_TYPE_PEER, peer, &encoded)
}
}
pub(crate) fn credential_address(credential: &BasicCredential) -> Result<&str, crate::LeafError> {
core::str::from_utf8(&credential.identifier).map_err(|_| {
crate::LeafError::IdentityBinding(String::from(
"credential identifier is not valid UTF-8, so it names no address",
))
})
}