extern crate alloc;
#[cfg(test)]
use alloc::string::ToString;
use alloc::{
collections::{BTreeMap, BTreeSet},
string::String,
vec::Vec,
};
use borsh::{BorshDeserialize, BorshSerialize};
use rialo_s_pubkey::Pubkey;
use crate::error::FeatureManagementError;
#[derive(Clone, Debug, PartialEq, Eq, BorshSerialize, BorshDeserialize)]
pub struct ScheduledRequest {
pub names: Vec<String>,
pub fire_at_ms: u64,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct FeaturesState {
authority: Pubkey,
pending_authority: Option<Pubkey>,
pub(crate) entries: BTreeSet<String>,
pub(crate) pending: BTreeMap<u64, ScheduledRequest>,
}
#[cfg(test)]
pub const DETERMINISTIC_TEST_KEYPAIR: &str =
"57Vqb7tHij5NhQnTgrgYXA19pC8ZVHQoCHpapSiQ8LJaeUvTcSBzKoB6CazhR6VtxmyVAbWnoeDSzD1Vm672NaKp";
impl FeaturesState {
pub fn new(authority: Pubkey) -> Self {
Self {
authority,
pending_authority: None,
entries: BTreeSet::new(),
pending: BTreeMap::new(),
}
}
#[cfg(test)]
pub fn new_for_test() -> Self {
use rialo_s_keypair::Keypair;
use rialo_s_signer::Signer;
let pubkey: Pubkey = Keypair::from_base58_string(DETERMINISTIC_TEST_KEYPAIR)
.try_pubkey()
.expect("Failed to get pubkey from deterministic test keypair");
Self::new(pubkey)
}
pub fn get_authority(&self) -> &Pubkey {
&self.authority
}
pub fn set_authority(&mut self, new_authority: Pubkey) {
self.authority = new_authority;
}
pub fn pending_authority(&self) -> Option<&Pubkey> {
self.pending_authority.as_ref()
}
pub fn set_pending_authority(&mut self, pending: Option<Pubkey>) {
self.pending_authority = pending;
}
pub fn propose_transfer(
&mut self,
new_authority: Pubkey,
) -> Result<(), FeatureManagementError> {
if new_authority == self.authority {
return Err(FeatureManagementError::InvalidTransferTarget);
}
if self.pending_authority.is_some() {
return Err(FeatureManagementError::PendingTransferExists);
}
self.pending_authority = Some(new_authority);
Ok(())
}
pub fn accept_transfer(&mut self) -> Result<(), FeatureManagementError> {
let Some(pending) = self.pending_authority else {
return Err(FeatureManagementError::NoPendingTransfer);
};
self.authority = pending;
self.pending_authority = None;
Ok(())
}
pub fn cancel_transfer(&mut self) -> Result<(), FeatureManagementError> {
if self.pending_authority.is_none() {
return Err(FeatureManagementError::NoPendingTransfer);
}
self.pending_authority = None;
Ok(())
}
pub fn serialize(&self) -> Result<Vec<u8>, borsh::io::Error> {
FeaturesStateVersioned::from(self).serialize()
}
pub fn deserialize(data: &[u8]) -> Result<Self, borsh::io::Error> {
Ok(FeaturesStateVersioned::deserialize(data)?.into())
}
pub fn is_active(&self, feature_name: &str) -> bool {
self.entries.contains(feature_name)
}
pub fn entries(&self) -> &BTreeSet<String> {
&self.entries
}
pub fn insert_for_test(&mut self, name: String) {
self.entries.insert(name);
}
pub fn enable(&mut self, names: Vec<String>) -> Result<(), FeatureManagementError> {
if names.len() > crate::MAX_NAMES_PER_BATCH {
return Err(FeatureManagementError::TooManyNames);
}
for name in &names {
if !crate::validate_feature_name(name) {
return Err(FeatureManagementError::InvalidFeatureName);
}
}
let new_distinct = names
.iter()
.filter(|n| !self.entries.contains(*n))
.collect::<BTreeSet<_>>()
.len();
if self.entries.len().saturating_add(new_distinct) > crate::MAX_FEATURE_COUNT {
return Err(FeatureManagementError::MaxFeatureCountExceeded);
}
for name in names {
self.entries.insert(name);
}
Ok(())
}
pub fn schedule(
&mut self,
request_id: u64,
names: Vec<String>,
fire_at_ms: u64,
) -> Result<(), FeatureManagementError> {
if names.len() > crate::MAX_NAMES_PER_BATCH {
return Err(FeatureManagementError::TooManyNames);
}
for name in &names {
if !crate::validate_feature_name(name) {
return Err(FeatureManagementError::InvalidFeatureName);
}
}
if self.pending.contains_key(&request_id) {
return Err(FeatureManagementError::RequestAlreadyExists);
}
if self.pending.len() >= crate::MAX_PENDING_REQUESTS {
return Err(FeatureManagementError::TooManyPendingRequests);
}
self.pending
.insert(request_id, ScheduledRequest { names, fire_at_ms });
let too_large = self
.serialize()
.map(|bytes| bytes.len() > crate::MAX_FEATURES_STATE_SIZE)
.unwrap_or(true);
if too_large {
self.pending.remove(&request_id);
return Err(FeatureManagementError::PendingStateTooLarge);
}
Ok(())
}
pub fn fire_scheduled(&mut self, request_id: u64) -> Result<(), FeatureManagementError> {
let req = self
.pending
.remove(&request_id)
.ok_or(FeatureManagementError::RequestNotFound)?;
self.enable(req.names)
}
pub fn cancel(&mut self, request_id: u64) -> Result<ScheduledRequest, FeatureManagementError> {
self.pending
.remove(&request_id)
.ok_or(FeatureManagementError::RequestNotFound)
}
pub fn pending(&self) -> &BTreeMap<u64, ScheduledRequest> {
&self.pending
}
}
#[derive(Clone, Debug, PartialEq, Eq, BorshSerialize, BorshDeserialize)]
pub struct FeaturesStateV1 {
pub(crate) authority: Pubkey,
pub(crate) pending_authority: Option<Pubkey>,
pub(crate) entries: BTreeSet<String>,
pub(crate) pending: BTreeMap<u64, ScheduledRequest>,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum FeaturesStateVersioned {
V1(FeaturesStateV1),
}
impl FeaturesStateVersioned {
pub fn serialize(&self) -> Result<Vec<u8>, borsh::io::Error> {
match self {
FeaturesStateVersioned::V1(v1) => borsh::to_vec(v1),
}
}
pub fn deserialize(data: &[u8]) -> Result<Self, borsh::io::Error> {
Ok(FeaturesStateVersioned::V1(borsh::from_slice(data)?))
}
}
impl From<&FeaturesState> for FeaturesStateV1 {
fn from(s: &FeaturesState) -> Self {
FeaturesStateV1 {
authority: s.authority,
pending_authority: s.pending_authority,
entries: s.entries.clone(),
pending: s.pending.clone(),
}
}
}
impl From<FeaturesState> for FeaturesStateV1 {
fn from(s: FeaturesState) -> Self {
FeaturesStateV1 {
authority: s.authority,
pending_authority: s.pending_authority,
entries: s.entries,
pending: s.pending,
}
}
}
impl From<FeaturesStateV1> for FeaturesState {
fn from(v: FeaturesStateV1) -> Self {
FeaturesState {
authority: v.authority,
pending_authority: v.pending_authority,
entries: v.entries,
pending: v.pending,
}
}
}
impl From<&FeaturesState> for FeaturesStateVersioned {
fn from(s: &FeaturesState) -> Self {
FeaturesStateVersioned::V1(FeaturesStateV1::from(s))
}
}
impl From<FeaturesState> for FeaturesStateVersioned {
fn from(s: FeaturesState) -> Self {
FeaturesStateVersioned::V1(FeaturesStateV1::from(s))
}
}
impl From<FeaturesStateVersioned> for FeaturesState {
fn from(v: FeaturesStateVersioned) -> Self {
match v {
FeaturesStateVersioned::V1(v1) => FeaturesState::from(v1),
}
}
}
#[cfg(test)]
mod tests {
use alloc::{format, vec};
use super::*;
#[test]
fn test_new_has_empty_set() {
let state = FeaturesState::new_for_test();
assert!(state.entries().is_empty());
}
#[test]
fn test_enable_inserts_name() {
let mut state = FeaturesState::new_for_test();
state.enable(vec!["feat_a".to_string()]).unwrap();
assert!(state.is_active("feat_a"));
assert!(!state.is_active("feat_b"));
}
#[test]
fn test_enable_is_idempotent() {
let mut state = FeaturesState::new_for_test();
state.enable(vec!["feat_a".to_string()]).unwrap();
state.enable(vec!["feat_a".to_string()]).unwrap();
assert_eq!(state.entries().len(), 1);
}
#[test]
fn test_enable_multi_names() {
let mut state = FeaturesState::new_for_test();
state
.enable(vec!["a".to_string(), "b".to_string(), "c".to_string()])
.unwrap();
assert_eq!(state.entries().len(), 3);
assert!(state.is_active("a"));
assert!(state.is_active("b"));
assert!(state.is_active("c"));
}
#[test]
fn test_enable_dedup_within_batch() {
let mut state = FeaturesState::new_for_test();
state
.enable(vec!["a".to_string(), "a".to_string(), "b".to_string()])
.unwrap();
assert_eq!(state.entries().len(), 2);
}
fn fill_to(state: &mut FeaturesState, target: usize) {
let batch = crate::MAX_NAMES_PER_BATCH;
let mut filled = 0;
while filled < target {
let take = (target - filled).min(batch);
let names: Vec<String> = (filled..filled + take).map(|i| format!("f{i}")).collect();
state.enable(names).unwrap();
filled += take;
}
assert_eq!(state.entries().len(), target);
}
#[test]
fn test_enable_max_count_enforced() {
let mut state = FeaturesState::new_for_test();
fill_to(&mut state, crate::MAX_FEATURE_COUNT);
assert_eq!(
state.enable(vec!["overflow".to_string()]),
Err(FeatureManagementError::MaxFeatureCountExceeded)
);
}
#[test]
fn test_enable_intra_batch_dup_at_boundary() {
let mut state = FeaturesState::new_for_test();
fill_to(&mut state, crate::MAX_FEATURE_COUNT - 1);
state
.enable(vec!["x".to_string(), "x".to_string()])
.expect("intra-batch duplicate must not double-count against the cap");
assert_eq!(state.entries().len(), crate::MAX_FEATURE_COUNT);
}
#[test]
fn test_set_authority_works() {
let mut state = FeaturesState::new_for_test();
let new = Pubkey::new_from_array([7u8; 32]);
state.set_authority(new);
assert_eq!(state.get_authority(), &new);
}
#[test]
fn test_borsh_round_trip_empty() {
let state = FeaturesState::new_for_test();
let bytes = state.serialize().unwrap();
let back = FeaturesState::deserialize(&bytes).unwrap();
assert_eq!(state, back);
}
#[test]
fn test_borsh_round_trip_with_entries() {
let mut state = FeaturesState::new_for_test();
state
.enable(vec!["alpha".to_string(), "beta".to_string()])
.unwrap();
let bytes = state.serialize().unwrap();
let back = FeaturesState::deserialize(&bytes).unwrap();
assert_eq!(state, back);
}
#[test]
fn test_wire_format_golden_none_pending() {
let mut state = FeaturesState::new(Pubkey::new_from_array([1u8; 32]));
state.insert_for_test("a".to_string());
state.insert_for_test("bb".to_string());
let bytes = state.serialize().unwrap();
let mut expected = vec![0u8; 0];
expected.extend_from_slice(&[1u8; 32]); expected.push(0); expected.extend_from_slice(&2u32.to_le_bytes()); expected.extend_from_slice(&1u32.to_le_bytes()); expected.extend_from_slice(b"a");
expected.extend_from_slice(&2u32.to_le_bytes()); expected.extend_from_slice(b"bb");
expected.extend_from_slice(&0u32.to_le_bytes());
assert_eq!(bytes, expected, "wire format drift detected");
let back = FeaturesState::deserialize(&bytes).unwrap();
assert_eq!(state, back, "golden bytes must round-trip");
}
#[test]
fn test_wire_format_golden_some_pending() {
let mut state = FeaturesState::new(Pubkey::new_from_array([1u8; 32]));
state.set_pending_authority(Some(Pubkey::new_from_array([2u8; 32])));
state.insert_for_test("a".to_string());
let bytes = state.serialize().unwrap();
let mut expected = vec![0u8; 0];
expected.extend_from_slice(&[1u8; 32]); expected.push(1); expected.extend_from_slice(&[2u8; 32]); expected.extend_from_slice(&1u32.to_le_bytes()); expected.extend_from_slice(&1u32.to_le_bytes()); expected.extend_from_slice(b"a");
expected.extend_from_slice(&0u32.to_le_bytes());
assert_eq!(bytes, expected, "wire format drift detected");
let back = FeaturesState::deserialize(&bytes).unwrap();
assert_eq!(state, back, "golden bytes must round-trip");
}
#[test]
fn test_wire_format_golden_with_pending() {
let mut state = FeaturesState::new(Pubkey::new_from_array([1u8; 32]));
state
.schedule(9, vec!["feat_x".to_string()], 1_700_000_000_000)
.expect("schedule must succeed");
let bytes = state.serialize().unwrap();
let mut expected = vec![0u8; 0];
expected.extend_from_slice(&[1u8; 32]); expected.push(0); expected.extend_from_slice(&0u32.to_le_bytes()); expected.extend_from_slice(&1u32.to_le_bytes()); expected.extend_from_slice(&9u64.to_le_bytes()); expected.extend_from_slice(&1u32.to_le_bytes()); expected.extend_from_slice(&6u32.to_le_bytes()); expected.extend_from_slice(b"feat_x");
expected.extend_from_slice(&1_700_000_000_000u64.to_le_bytes());
assert_eq!(bytes, expected, "wire format drift detected");
let back = FeaturesState::deserialize(&bytes).unwrap();
assert_eq!(state, back, "golden bytes must round-trip");
}
#[test]
fn test_wire_format_golden_full_shape() {
let mut state = FeaturesState::new(Pubkey::new_from_array([7u8; 32]));
state.set_pending_authority(Some(Pubkey::new_from_array([8u8; 32])));
state.insert_for_test("a".to_string());
state.insert_for_test("bb".to_string());
state
.schedule(1, vec!["x".to_string()], 1_000)
.expect("schedule id 1");
state
.schedule(2, vec!["y".to_string(), "z".to_string()], 2_000)
.expect("schedule id 2");
let bytes = state.serialize().unwrap();
let mut expected = vec![0u8; 0];
expected.extend_from_slice(&[7u8; 32]); expected.push(1); expected.extend_from_slice(&[8u8; 32]); expected.extend_from_slice(&2u32.to_le_bytes()); expected.extend_from_slice(&1u32.to_le_bytes()); expected.extend_from_slice(b"a");
expected.extend_from_slice(&2u32.to_le_bytes()); expected.extend_from_slice(b"bb");
expected.extend_from_slice(&2u32.to_le_bytes()); expected.extend_from_slice(&1u64.to_le_bytes()); expected.extend_from_slice(&1u32.to_le_bytes()); expected.extend_from_slice(&1u32.to_le_bytes()); expected.extend_from_slice(b"x");
expected.extend_from_slice(&1_000u64.to_le_bytes()); expected.extend_from_slice(&2u64.to_le_bytes()); expected.extend_from_slice(&2u32.to_le_bytes()); expected.extend_from_slice(&1u32.to_le_bytes()); expected.extend_from_slice(b"y");
expected.extend_from_slice(&1u32.to_le_bytes()); expected.extend_from_slice(b"z");
expected.extend_from_slice(&2_000u64.to_le_bytes());
assert_eq!(bytes, expected, "wire format drift detected");
let back = FeaturesState::deserialize(&bytes).unwrap();
assert_eq!(state, back, "golden bytes must round-trip");
}
#[test]
fn test_pending_authority_round_trip() {
let mut state = FeaturesState::new_for_test();
assert!(state.pending_authority().is_none());
let p = Pubkey::new_from_array([3u8; 32]);
state.set_pending_authority(Some(p));
assert_eq!(state.pending_authority(), Some(&p));
state.set_pending_authority(None);
assert!(state.pending_authority().is_none());
}
#[test]
fn test_propose_then_accept_transfers_authority() {
let mut state = FeaturesState::new_for_test();
let original = *state.get_authority();
let new = Pubkey::new_from_array([4u8; 32]);
state.propose_transfer(new).unwrap();
assert_eq!(state.pending_authority(), Some(&new));
assert_eq!(state.get_authority(), &original);
state.accept_transfer().unwrap();
assert_eq!(state.get_authority(), &new);
assert!(state.pending_authority().is_none());
}
#[test]
fn test_propose_to_self_rejected() {
let mut state = FeaturesState::new_for_test();
let current = *state.get_authority();
assert_eq!(
state.propose_transfer(current),
Err(FeatureManagementError::InvalidTransferTarget)
);
assert!(state.pending_authority().is_none());
}
#[test]
fn test_propose_with_pending_rejected() {
let mut state = FeaturesState::new_for_test();
state.propose_transfer(Pubkey::new_unique()).unwrap();
assert_eq!(
state.propose_transfer(Pubkey::new_unique()),
Err(FeatureManagementError::PendingTransferExists)
);
}
#[test]
fn test_accept_with_no_pending_rejected() {
let mut state = FeaturesState::new_for_test();
assert_eq!(
state.accept_transfer(),
Err(FeatureManagementError::NoPendingTransfer)
);
}
#[test]
fn test_cancel_clears_pending() {
let mut state = FeaturesState::new_for_test();
let original = *state.get_authority();
state.propose_transfer(Pubkey::new_unique()).unwrap();
state.cancel_transfer().unwrap();
assert!(state.pending_authority().is_none());
assert_eq!(state.get_authority(), &original);
}
#[test]
fn test_cancel_with_no_pending_rejected() {
let mut state = FeaturesState::new_for_test();
assert_eq!(
state.cancel_transfer(),
Err(FeatureManagementError::NoPendingTransfer)
);
}
#[test]
fn test_full_cycle_propose_cancel_propose_accept() {
let mut state = FeaturesState::new_for_test();
let original = *state.get_authority();
let first = Pubkey::new_from_array([10u8; 32]);
let second = Pubkey::new_from_array([20u8; 32]);
state.propose_transfer(first).unwrap();
state.cancel_transfer().unwrap();
assert_eq!(state.get_authority(), &original);
assert!(state.pending_authority().is_none());
state.propose_transfer(second).unwrap();
state.accept_transfer().unwrap();
assert_eq!(state.get_authority(), &second);
assert!(state.pending_authority().is_none());
}
#[test]
fn test_accept_then_propose_works_with_new_authority() {
let mut state = FeaturesState::new_for_test();
let mid = Pubkey::new_from_array([11u8; 32]);
let final_target = Pubkey::new_from_array([22u8; 32]);
state.propose_transfer(mid).unwrap();
state.accept_transfer().unwrap();
assert_eq!(state.get_authority(), &mid);
state.propose_transfer(final_target).unwrap();
assert_eq!(state.pending_authority(), Some(&final_target));
}
#[test]
fn test_propose_does_not_touch_authority() {
let mut state = FeaturesState::new_for_test();
let original = *state.get_authority();
state
.propose_transfer(Pubkey::new_from_array([33u8; 32]))
.unwrap();
assert_eq!(state.get_authority(), &original);
}
#[test]
fn test_single_step_clear_pending_blocks_stale_accept() {
let mut state = FeaturesState::new_for_test();
let pending = Pubkey::new_from_array([44u8; 32]);
let single_step_target = Pubkey::new_from_array([55u8; 32]);
state.propose_transfer(pending).unwrap();
assert_eq!(state.pending_authority(), Some(&pending));
state.set_authority(single_step_target);
state.set_pending_authority(None);
assert_eq!(
state.accept_transfer(),
Err(FeatureManagementError::NoPendingTransfer)
);
assert_eq!(state.get_authority(), &single_step_target);
}
#[test]
fn test_enable_rejects_oversized_batch() {
let mut state = FeaturesState::new_for_test();
let names: Vec<String> = (0..=crate::MAX_NAMES_PER_BATCH)
.map(|i| format!("f{i}"))
.collect();
assert_eq!(
state.enable(names),
Err(FeatureManagementError::TooManyNames)
);
}
#[test]
fn test_enable_rejects_name_above_max_length() {
let mut state = FeaturesState::new_for_test();
let long_name = "a".repeat(crate::MAX_FEATURE_NAME_LENGTH + 1);
assert_eq!(
state.enable(vec![long_name]),
Err(FeatureManagementError::InvalidFeatureName)
);
}
#[test]
fn test_enable_accepts_name_at_exact_max_length() {
let mut state = FeaturesState::new_for_test();
let at_max = "a".repeat(crate::MAX_FEATURE_NAME_LENGTH);
state
.enable(vec![at_max])
.expect("name at exact MAX_FEATURE_NAME_LENGTH must be accepted");
assert_eq!(state.entries().len(), 1);
}
#[test]
fn test_enable_rejects_one_long_name_in_otherwise_valid_batch() {
let mut state = FeaturesState::new_for_test();
let names = vec![
"shortish".to_string(),
"a".repeat(crate::MAX_FEATURE_NAME_LENGTH + 1),
"also_shortish".to_string(),
];
assert_eq!(
state.enable(names),
Err(FeatureManagementError::InvalidFeatureName)
);
assert!(
state.entries().is_empty(),
"rejected enable must leave state untouched"
);
}
#[test]
fn test_enable_rejects_empty_name() {
let mut state = FeaturesState::new_for_test();
assert_eq!(
state.enable(vec![String::new()]),
Err(FeatureManagementError::InvalidFeatureName)
);
}
#[test]
fn test_enable_rejects_invalid_charset() {
let mut state = FeaturesState::new_for_test();
assert_eq!(
state.enable(vec!["has space".to_string()]),
Err(FeatureManagementError::InvalidFeatureName)
);
assert_eq!(
state.enable(vec!["dot.name".to_string()]),
Err(FeatureManagementError::InvalidFeatureName)
);
}
#[test]
fn schedule_rejects_duplicate_request_id() {
let mut state = FeaturesState::new_for_test();
state
.schedule(7, vec!["feat_a".to_string()], 1_000)
.expect("first schedule with id 7 must succeed");
assert_eq!(
state.schedule(7, vec!["feat_b".to_string()], 2_000),
Err(FeatureManagementError::RequestAlreadyExists)
);
assert_eq!(state.pending().len(), 1);
let entry = state.pending().get(&7).expect("entry must exist");
assert_eq!(entry.names, vec!["feat_a".to_string()]);
assert_eq!(entry.fire_at_ms, 1_000);
}
#[test]
fn schedule_duplicate_id_at_capacity_reports_duplicate_not_full() {
let mut state = FeaturesState::new_for_test();
for id in 0..crate::MAX_PENDING_REQUESTS as u64 {
state
.schedule(id, vec!["feat".to_string()], 1_000)
.expect("fill under cap");
}
assert_eq!(state.pending().len(), crate::MAX_PENDING_REQUESTS);
assert_eq!(
state.schedule(0, vec!["feat".to_string()], 2_000),
Err(FeatureManagementError::RequestAlreadyExists),
"a duplicate id at capacity is RequestAlreadyExists, not TooManyPendingRequests"
);
assert_eq!(
state.schedule(
crate::MAX_PENDING_REQUESTS as u64,
vec!["feat".to_string()],
2_000
),
Err(FeatureManagementError::TooManyPendingRequests)
);
}
#[test]
fn schedule_rejects_oversized_pending_state() {
let max_batch =
|| vec!["a".repeat(crate::MAX_FEATURE_NAME_LENGTH); crate::MAX_NAMES_PER_BATCH];
let mut state = FeaturesState::new_for_test();
state
.schedule(1, max_batch(), 1_000)
.expect("one max-size request must fit");
assert!(
state.serialize().expect("serialize").len() <= crate::MAX_FEATURES_STATE_SIZE,
"a single max-size request must stay within MAX_FEATURES_STATE_SIZE"
);
assert_eq!(
state.schedule(2, max_batch(), 2_000),
Err(FeatureManagementError::PendingStateTooLarge)
);
assert_eq!(state.pending().len(), 1);
assert!(state.pending().get(&2).is_none());
assert!(state.pending().len() < crate::MAX_PENDING_REQUESTS);
}
#[test]
fn fire_scheduled_drains_pending_and_activates() {
let mut state = FeaturesState::new_for_test();
state
.schedule(5, vec!["feat_a".to_string(), "feat_b".to_string()], 1_000)
.expect("schedule must succeed");
assert_eq!(state.pending().len(), 1);
state.fire_scheduled(5).expect("fire must succeed");
assert!(state.is_active("feat_a"));
assert!(state.is_active("feat_b"));
assert!(state.pending().get(&5).is_none());
assert!(state.pending().is_empty());
}
#[test]
fn fire_scheduled_unknown_request_id() {
let mut state = FeaturesState::new_for_test();
assert_eq!(
state.fire_scheduled(99),
Err(FeatureManagementError::RequestNotFound)
);
}
#[test]
fn test_enable_is_permissive_on_unknown_names() {
let mut state = FeaturesState::new_for_test();
state
.enable(vec!["totally_fabricated_name_not_in_any_binary".to_string()])
.expect("enable must be permissive on KNOWN_FEATURES membership");
}
#[test]
fn v1_serialize_is_byte_identical_to_bare_borsh() {
let mut state = FeaturesState::new(Pubkey::new_from_array([1u8; 32]));
state.set_pending_authority(Some(Pubkey::new_from_array([2u8; 32])));
state.insert_for_test("a".to_string());
state
.schedule(9, vec!["feat_x".to_string()], 1_700_000_000_000)
.expect("schedule must succeed");
let via_envelope = state.serialize().unwrap();
let bare_v1 = borsh::to_vec(&FeaturesStateV1::from(&state)).unwrap();
assert_eq!(
via_envelope, bare_v1,
"envelope V1 output must equal a bare borsh of FeaturesStateV1"
);
}
#[test]
fn versioned_envelope_round_trips_v1() {
let mut state = FeaturesState::new(Pubkey::new_from_array([5u8; 32]));
state.insert_for_test("alpha".to_string());
let versioned = FeaturesStateVersioned::from(&state);
let bytes = versioned.serialize().unwrap();
let back = FeaturesStateVersioned::deserialize(&bytes).unwrap();
assert_eq!(versioned, back);
match &back {
FeaturesStateVersioned::V1(_) => {}
}
assert_eq!(FeaturesState::from(back), state);
}
#[test]
fn features_state_seam_round_trips_through_envelope() {
let mut state = FeaturesState::new(Pubkey::new_from_array([7u8; 32]));
state.set_pending_authority(Some(Pubkey::new_from_array([8u8; 32])));
state
.enable(vec!["one".to_string(), "two".to_string()])
.unwrap();
state
.schedule(3, vec!["later".to_string()], 1_800_000_000_000)
.unwrap();
let bytes = state.serialize().unwrap();
let back = FeaturesState::deserialize(&bytes).unwrap();
assert_eq!(state, back);
}
mod proptests {
use alloc::collections::{BTreeMap, BTreeSet};
use proptest::{prelude::*, test_runner::TestCaseError};
use super::*;
prop_compose! {
fn valid_name()(s in "[a-zA-Z0-9_-]{1,16}") -> String {
s
}
}
prop_compose! {
fn valid_batch()(batch in prop::collection::vec(valid_name(), 0..=8)) -> Vec<String> {
batch
}
}
prop_compose! {
fn arbitrary_state()(
authority in any::<[u8; 32]>(),
pending_authority in proptest::option::of(any::<[u8; 32]>()),
entries in prop::collection::btree_set(valid_name(), 0..=8),
pending in prop::collection::btree_map(
any::<u64>(),
(prop::collection::vec(valid_name(), 0..=4), any::<u64>()),
0..=8,
),
) -> FeaturesState {
let mut state = FeaturesState::new(Pubkey::new_from_array(authority));
state.set_pending_authority(pending_authority.map(Pubkey::new_from_array));
state.entries = entries;
state.pending = pending
.into_iter()
.map(|(id, (names, fire_at_ms))| (id, ScheduledRequest { names, fire_at_ms }))
.collect::<BTreeMap<u64, ScheduledRequest>>();
state
}
}
proptest! {
#[test]
fn borsh_round_trips_arbitrary_state(state in arbitrary_state()) {
let bytes = state
.serialize()
.map_err(|e| TestCaseError::fail(e.to_string()))?;
prop_assert!(
bytes.len() <= crate::MAX_FEATURES_STATE_SIZE,
"serialized {} bytes exceeds MAX_FEATURES_STATE_SIZE",
bytes.len()
);
let back = FeaturesState::deserialize(&bytes)
.map_err(|e| TestCaseError::fail(e.to_string()))?;
prop_assert_eq!(state, back);
}
#[test]
fn enable_is_monotone_and_unions(batches in prop::collection::vec(valid_batch(), 0..=10)) {
let mut state = FeaturesState::new_for_test();
let mut expected: BTreeSet<String> = BTreeSet::new();
let mut prev_len = 0usize;
for batch in &batches {
let r = state.enable(batch.clone());
prop_assert!(r.is_ok(), "valid batch must enable: {r:?}");
expected.extend(batch.iter().cloned());
prop_assert!(state.entries().len() >= prev_len);
prev_len = state.entries().len();
}
prop_assert_eq!(state.entries(), &expected);
for name in &expected {
prop_assert!(state.is_active(name));
}
let all: Vec<String> = expected.iter().cloned().collect();
for chunk in all.chunks(crate::MAX_NAMES_PER_BATCH) {
let r = state.enable(chunk.to_vec());
prop_assert!(r.is_ok(), "re-enable is a no-op: {r:?}");
}
prop_assert_eq!(state.entries().len(), expected.len());
}
#[test]
fn is_active_equals_membership(
batch in valid_batch(),
probe in valid_name(),
) {
let oracle: BTreeSet<String> = batch.iter().cloned().collect();
let mut state = FeaturesState::new_for_test();
let r = state.enable(batch);
prop_assert!(r.is_ok(), "valid batch must enable: {r:?}");
prop_assert_eq!(state.is_active(&probe), oracle.contains(&probe));
}
#[test]
fn enable_rejects_oversized_batch(
extra in 1usize..=20,
) {
let mut state = FeaturesState::new_for_test();
let names: Vec<String> = (0..crate::MAX_NAMES_PER_BATCH + extra)
.map(|i| format!("f{i}"))
.collect();
prop_assert_eq!(state.enable(names), Err(FeatureManagementError::TooManyNames));
prop_assert!(state.entries().is_empty());
}
#[test]
fn enable_rejects_invalid_name(
bad in prop_oneof![
"[a-z]{0,4}[ .!@/][a-z]{0,4}",
Just("a".repeat(crate::MAX_FEATURE_NAME_LENGTH + 1)),
],
) {
prop_assume!(!crate::validate_feature_name(&bad));
let mut state = FeaturesState::new_for_test();
prop_assert_eq!(
state.enable(vec![bad]),
Err(FeatureManagementError::InvalidFeatureName)
);
prop_assert!(state.entries().is_empty());
}
#[test]
fn schedule_cancel_round_trip(
reqs in prop::collection::vec(
(any::<u64>(), valid_batch(), any::<u64>()),
0..=50,
),
) {
let mut state = FeaturesState::new_for_test();
let mut recorded: Vec<(u64, Vec<String>, u64)> = Vec::new();
for (id, names, fire_at_ms) in reqs {
if recorded.iter().any(|(rid, ..)| *rid == id) {
prop_assert_eq!(
state.schedule(id, names.clone(), fire_at_ms),
Err(FeatureManagementError::RequestAlreadyExists)
);
} else {
let r = state.schedule(id, names.clone(), fire_at_ms);
prop_assert!(r.is_ok(), "distinct id must schedule: {r:?}");
let entry = state.pending().get(&id);
prop_assert!(entry.is_some(), "entry recorded");
let entry = entry.unwrap();
prop_assert_eq!(&entry.names, &names);
prop_assert_eq!(entry.fire_at_ms, fire_at_ms);
recorded.push((id, names, fire_at_ms));
}
}
prop_assert_eq!(state.pending().len(), recorded.len());
let mut absent = 0u64;
while recorded.iter().any(|(rid, ..)| *rid == absent) {
absent = absent.wrapping_add(1);
}
prop_assert_eq!(
state.cancel(absent),
Err(FeatureManagementError::RequestNotFound)
);
for (id, names, fire_at_ms) in &recorded {
let got = state.cancel(*id);
prop_assert!(got.is_ok(), "recorded id must cancel: {got:?}");
let got = got.unwrap();
prop_assert_eq!(&got.names, names);
prop_assert_eq!(got.fire_at_ms, *fire_at_ms);
prop_assert_eq!(
state.cancel(*id),
Err(FeatureManagementError::RequestNotFound)
);
}
prop_assert!(state.pending().is_empty());
}
#[test]
fn schedule_rejects_over_capacity(
fire_at_ms in any::<u64>(),
) {
let mut state = FeaturesState::new_for_test();
for id in 0..crate::MAX_PENDING_REQUESTS as u64 {
let r = state.schedule(id, vec!["feat".to_string()], fire_at_ms);
prop_assert!(r.is_ok(), "id {id} under cap must schedule: {r:?}");
}
prop_assert_eq!(state.pending().len(), crate::MAX_PENDING_REQUESTS);
prop_assert_eq!(
state.schedule(
crate::MAX_PENDING_REQUESTS as u64,
vec!["feat".to_string()],
fire_at_ms,
),
Err(FeatureManagementError::TooManyPendingRequests)
);
prop_assert_eq!(state.pending().len(), crate::MAX_PENDING_REQUESTS);
}
#[test]
fn schedule_rejects_invalid_input(
extra in 1usize..=20,
bad in prop_oneof![
"[a-z]{0,4}[ .!@/][a-z]{0,4}",
Just("a".repeat(crate::MAX_FEATURE_NAME_LENGTH + 1)),
],
) {
prop_assume!(!crate::validate_feature_name(&bad));
let mut state = FeaturesState::new_for_test();
let oversized: Vec<String> = (0..crate::MAX_NAMES_PER_BATCH + extra)
.map(|i| format!("f{i}"))
.collect();
prop_assert_eq!(
state.schedule(1, oversized, 1_000),
Err(FeatureManagementError::TooManyNames)
);
prop_assert_eq!(
state.schedule(2, vec![bad], 1_000),
Err(FeatureManagementError::InvalidFeatureName)
);
prop_assert!(state.pending().is_empty());
}
}
}
}