rialo-feature-management-program-interface 0.10.2

// Copyright (c) Subzero Labs, Inc.
// SPDX-License-Identifier: Apache-2.0

//! Feature state management
//!
//! On-chain state for the feature management program. Activation is
//! **presence-based**: a feature is active iff its name is in `entries`. No
//! per-entry payload, no clock. See `runtime/execution/NORTHSTAR-protocol-upgrades.md`.

extern crate alloc;

#[cfg(test)]
use alloc::string::ToString;
use alloc::{collections::BTreeSet, string::String, vec::Vec};

use borsh::{BorshDeserialize, BorshSerialize};
use rialo_s_pubkey::Pubkey;

use crate::error::FeatureManagementError;

/// The program's global state.
///
/// Stored under the `STORAGE_ACCOUNT_SEED` PDA owned by the feature
/// management program. Carries the authority pubkey, an optional pending
/// authority for the two-step transfer handshake, and the set of feature
/// names that are currently active.
#[derive(Clone, Debug, PartialEq, Eq, BorshSerialize, BorshDeserialize)]
pub struct FeaturesState {
    /// Authority pubkey.
    authority: Pubkey,
    /// Pubkey of the proposed next authority while a two-step transfer is
    /// in flight. `None` while no transfer is pending.
    ///
    /// Set by `ProposeAuthorityTransfer`, cleared by `CancelAuthorityTransfer`,
    /// and consumed by `AcceptAuthorityTransfer` (which promotes it to
    /// `authority`). The single-step `UpdateAuthority` path sets `authority`
    /// directly and leaves this slot untouched.
    pending_authority: Option<Pubkey>,
    /// Active feature names. Presence is the activation; there is no
    /// per-entry payload.
    ///
    /// `pub(crate)` so external callers cannot bypass `enable`'s validation
    /// by writing directly. Read access through `Self::entries()`; tests
    /// construct fixtures via `Self::insert_for_test()`.
    pub(crate) entries: BTreeSet<String>,
}

#[cfg(test)]
pub const DETERMINISTIC_TEST_KEYPAIR: &str =
    "57Vqb7tHij5NhQnTgrgYXA19pC8ZVHQoCHpapSiQ8LJaeUvTcSBzKoB6CazhR6VtxmyVAbWnoeDSzD1Vm672NaKp";

impl FeaturesState {
    /// Create a new state with the given authority, no pending transfer, and
    /// an empty entry set.
    ///
    /// Genesis intentionally pre-seeds no features — features active from
    /// genesis are part of the implementation, not feature flags.
    pub fn new(authority: Pubkey) -> Self {
        Self {
            authority,
            pending_authority: None,
            entries: BTreeSet::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;
    }

    /// Propose a two-step authority transfer.
    ///
    /// Sets `pending_authority` to `Some(new_authority)`. Returns
    /// `PendingTransferExists` if a previous proposal is still
    /// outstanding — callers must `cancel_transfer` first. Returns
    /// `InvalidTransferTarget` if `new_authority` equals the current
    /// authority (degenerate target; distinct from a signer mismatch,
    /// which is the processor's `Unauthorized`).
    ///
    /// Caller is responsible for verifying the current-authority signature
    /// before invoking this.
    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(())
    }

    /// Commit a previously-proposed authority transfer.
    ///
    /// On success the authority moves to the pending value and
    /// `pending_authority` clears. Returns `NoPendingTransfer` if nothing
    /// is pending.
    ///
    /// **Contract:** the caller MUST have verified that the transaction
    /// signer equals `pending_authority()` before invoking this. This
    /// method does not re-check; the processor proves signer == pending
    /// via `verify_authority(&pending, ...)` and then commits here. Code
    /// outside the processor that calls this without the signer check
    /// would let any keypair commit the pending authority.
    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(())
    }

    /// Cancel a previously-proposed authority transfer.
    ///
    /// Clears `pending_authority`. Returns `NoPendingTransfer` if nothing
    /// is pending. Caller is responsible for verifying the
    /// current-authority signature before invoking this.
    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> {
        borsh::to_vec(self)
    }

    pub fn deserialize(data: &[u8]) -> Result<Self, borsh::io::Error> {
        borsh::from_slice(data)
    }

    /// Whether `feature_name` is active.
    ///
    /// Activation is membership: a feature is active iff its name is in
    /// `entries`. No clock consult.
    pub fn is_active(&self, feature_name: &str) -> bool {
        self.entries.contains(feature_name)
    }

    /// Read-only view of the active entry set. Use `enable` to mutate so the
    /// `MAX_FEATURE_COUNT` cap stays enforced.
    pub fn entries(&self) -> &BTreeSet<String> {
        &self.entries
    }

    /// Test-only direct insert that skips `enable`'s validation. Use to
    /// seed fixtures that need specific names. Never call from production
    /// paths.
    ///
    /// Not `cfg(test)`-gated because cross-crate test crates (e.g.
    /// `svm-execution`'s bank tests) need to construct `FeaturesState`
    /// fixtures, and `cfg(test)` from the declaring crate is invisible to
    /// dependent test crates. Documented as fixtures-only and never invoked
    /// by production paths.
    pub fn insert_for_test(&mut self, name: String) {
        self.entries.insert(name);
    }

    /// Add one or more feature names to the active set.
    ///
    /// Idempotent: re-submitting an existing name is a no-op (presence is
    /// the activation, so a name that is already present is already
    /// active). Enforces:
    ///
    /// * `MAX_NAMES_PER_BATCH` per-instruction cap → `TooManyNames`.
    /// * `validate_feature_name` on every name (length, allowed charset,
    ///   no leading/trailing whitespace) → `InvalidFeatureName`. Owning
    ///   the check here keeps it the single source of truth — callers
    ///   reaching `enable` outside the processor path cannot smuggle in
    ///   malformed names.
    /// * `MAX_FEATURE_COUNT` cap on the resulting set size — checked
    ///   against the post-insert count of distinct new names so a single
    ///   batch cannot push the registry past the cap →
    ///   `MaxFeatureCountExceeded`.
    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);
            }
        }
        // Note: `enable` is permissive on `KNOWN_FEATURES` membership. The
        // on-chain program does not (and cannot, deterministically) reject
        // a name absent from a given binary's `KNOWN_FEATURES`; that's a
        // binary-level concept. An active feature unknown to the running
        // binary is caught by `Bank::can_process_block`, which halts the
        // node at the activation block. Operators staging a name that no
        // released binary yet recognises is therefore on the hook to ship
        // the binary first.
        // Count *distinct* names absent from the existing set. Without the
        // collect-into-set, intra-batch duplicates (`["x", "x"]`) would
        // double-count and reject a request that BTreeSet semantics would
        // otherwise accept.
        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(())
    }
}

#[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);
    }

    /// Helper: fill state by submitting MAX_NAMES_PER_BATCH-sized chunks
    /// until it carries `target` names.
    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() {
        // At `MAX_FEATURE_COUNT - 1`, enabling `["x", "x"]` must succeed:
        // BTreeSet inserts a single new name, so the post-insert count is
        // exactly `MAX_FEATURE_COUNT`. A naive occurrence count would
        // wrongly reject this.
        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);
    }

    /// Golden wire-format test. Pins the exact borsh byte layout so any
    /// accidental field reorder, type swap, or new field added without a
    /// schema bump trips here rather than silently mutating the on-chain
    /// shape.
    ///
    /// Layout (borsh):
    /// * `authority`: 32 bytes (Pubkey, raw)
    /// * `pending_authority`: 1-byte tag (0 = None, 1 = Some) + 32 bytes when Some
    /// * `entries`: `BTreeSet<String>` — 4-byte LE length, then each item as
    ///   4-byte LE length + utf-8 bytes, in sorted order
    #[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]); // authority
        expected.push(0); // pending_authority = None tag
        expected.extend_from_slice(&2u32.to_le_bytes()); // entries len
        expected.extend_from_slice(&1u32.to_le_bytes()); // "a" len
        expected.extend_from_slice(b"a");
        expected.extend_from_slice(&2u32.to_le_bytes()); // "bb" len
        expected.extend_from_slice(b"bb");

        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]); // authority
        expected.push(1); // pending_authority = Some tag
        expected.extend_from_slice(&[2u8; 32]); // pending_authority payload
        expected.extend_from_slice(&1u32.to_le_bytes()); // entries len
        expected.extend_from_slice(&1u32.to_le_bytes()); // "a" len
        expected.extend_from_slice(b"a");

        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() {
        // Operator cancels a proposal and replaces it with a different
        // target; the new target accepts. Guards against a cancelled
        // proposal leaking into the post-accept authority.
        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() {
        // After a transfer commits, the new authority can propose
        // another transfer — the pending slot is back to None
        // post-accept so no `PendingTransferExists` leakage.
        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() {
        // A proposal in flight must not modify the active authority field
        // — only Accept commits.
        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() {
        // Models the processor's `UpdateAuthority` sequence:
        // `set_authority(new)` + `set_pending_authority(None)`. After
        // that sequence the previously-pending party cannot displace
        // the just-set authority via `accept_transfer`. Guards the
        // displacement vector flagged in PR3790 review.
        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));

        // Operator picks the single-step path instead of accepting.
        state.set_authority(single_step_target);
        state.set_pending_authority(None);

        // Stale pending party cannot displace the new authority.
        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() {
        // The check is per-name, not per-batch: a single oversized name
        // among well-formed ones must still trip the rejection.
        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() {
        // `validate_feature_name` allows alphanumeric + `_` + `-` only.
        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 test_enable_is_permissive_on_unknown_names() {
        // Documents the design: `state.enable` does not check
        // `KNOWN_FEATURES`. A name unknown to a binary may be stored on
        // chain; the runtime's `can_process_block` is what halts the node
        // at the activation block if the name is also active. The on-chain
        // program lacks a deterministic view of every binary's
        // `KNOWN_FEATURES`, so the check cannot live here.
        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");
    }
}