rialo-feature-management-program-interface 0.8.0-alpha.0

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

//! Feature state management
//!
//! Production schema for the on-chain feature registry, with two-phase
//! sticky-upsert validation. See `runtime/execution/NORTHSTAR-protocol-upgrades.md`
//! for the surrounding design.

extern crate alloc;

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

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

/// Sanity cap on `start_time_ms` / `end_time_ms`: roughly 100 years in
/// milliseconds.
///
/// Catches accidental seconds-vs-microseconds (or larger) unit confusion: a
/// well-formed ms timestamp falls comfortably below this cap, but a
/// microsecond or nanosecond value will overshoot it. Internal to the crate;
/// not part of the public API.
const MAX_TIMESTAMP_MS: u64 = 100 * 365 * 24 * 60 * 60 * 1000;

/// Per-feature entry stored in `FeaturesState`.
///
/// * Times are Unix epoch milliseconds, matching the unit Rialo writes into
///   `Clock::unix_timestamp` from the subdag-fed timestamp.
/// * `end_time_ms == None` marks a sticky entry: once it activates it stays
///   active forever and the management program rejects any Upsert that
///   would deactivate it.
#[derive(Clone, Debug, PartialEq, Eq, BorshSerialize, BorshDeserialize)]
pub struct FeatureEntry {
    /// Activation time (Unix epoch ms).
    pub start_time_ms: u64,
    /// Deactivation time (Unix epoch ms). `None` marks a sticky entry that
    /// never expires.
    pub end_time_ms: Option<u64>,
}

impl FeatureEntry {
    /// `true` iff this entry is sticky (`end_time_ms == None`).
    pub fn is_sticky(&self) -> bool {
        self.end_time_ms.is_none()
    }
}

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 and the per-feature entry map.
#[derive(Clone, Debug, PartialEq, BorshSerialize, BorshDeserialize)]
pub struct FeaturesState {
    /// Authority pubkey.
    authority: Pubkey,
    /// Per-feature entries, keyed by feature name.
    ///
    /// `pub(crate)` so external callers cannot bypass the two-phase sticky
    /// validation in `upsert` by writing directly. Read access through
    /// `Self::entries()`; tests construct fixtures via
    /// `Self::insert_for_test()`.
    pub(crate) entries: BTreeMap<String, FeatureEntry>,
}

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

impl FeaturesState {
    /// Create a new state with the given authority 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,
            entries: 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 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 at `current_time_ms`.
    ///
    /// Windowed entries are active iff `start_time_ms <= current_time_ms <
    /// end_time_ms`. Sticky entries (`end_time_ms == None`) are active iff
    /// `start_time_ms <= current_time_ms`.
    pub fn is_active(&self, feature_name: &str, current_time_ms: u64) -> bool {
        let Some(entry) = self.entries.get(feature_name) else {
            return false;
        };
        if current_time_ms < entry.start_time_ms {
            return false;
        }
        match entry.end_time_ms {
            Some(end) => current_time_ms < end,
            None => true,
        }
    }

    pub fn get(&self, feature_name: &str) -> Option<&FeatureEntry> {
        self.entries.get(feature_name)
    }

    /// Read-only view of every entry. Use `upsert` / `remove` to mutate so
    /// the two-phase sticky validation is not bypassed.
    pub fn entries(&self) -> &BTreeMap<String, FeatureEntry> {
        &self.entries
    }

    /// Test-only direct insert that skips `upsert`'s sticky validation.
    /// Use to seed fixtures that need a specific shape (e.g. an
    /// already-active entry before the test calls `upsert`); never call
    /// from production code 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. The function is documented as
    /// fixtures-only and never invoked by production paths.
    pub fn insert_for_test(&mut self, name: String, entry: FeatureEntry) {
        self.entries.insert(name, entry);
    }

    /// Add or update a feature entry under two-phase sticky validation.
    ///
    /// Phase is determined by `current_time_ms` relative to the existing
    /// entry's `start_time_ms`:
    ///
    /// * **Scheduled phase** (`current_time_ms < existing.start_time_ms`):
    ///   - A sticky entry cannot become windowed (caller may not pass a
    ///     concrete `end_time_ms` over an existing `None`).
    ///   - A **windowed→sticky upgrade is permitted**: an existing
    ///     windowed entry may be upserted with `end_time_ms = None` to
    ///     graduate the feature to a permanent gate.
    ///   - Whenever the new entry is sticky, `start_time_ms` may only
    ///     move *forward* relative to the existing entry. This applies
    ///     to both sticky→sticky and windowed→sticky upgrades, so the
    ///     operator-notice window the sticky rules protect cannot be
    ///     collapsed by re-classifying a windowed entry.
    ///   - Windowed→windowed modifications are otherwise unconstrained.
    /// * **Active phase** (`current_time_ms >= existing.start_time_ms`):
    ///   - Sticky entries are fully frozen; any modification is rejected
    ///     (idempotent no-op rewrites are accepted).
    ///   - Windowed entries stay mutable — `end_time_ms` may be reduced
    ///     (including to the past, which effectively deactivates) and
    ///     `start_time_ms` may be shifted forward.
    ///
    /// On a fresh insert, only invariants on the entry itself are checked
    /// (range, count cap).
    pub fn upsert(
        &mut self,
        name: String,
        entry: FeatureEntry,
        current_time_ms: u64,
    ) -> Result<(), FeatureManagementError> {
        Self::validate_entry_shape(&entry)?;

        match self.entries.get(&name) {
            None => {
                if self.entries.len() >= crate::MAX_FEATURE_COUNT {
                    return Err(FeatureManagementError::MaxFeatureCountExceeded);
                }
            }
            Some(existing) => {
                if current_time_ms >= existing.start_time_ms {
                    // Active-phase freeze applies only to sticky entries.
                    // Windowed entries remain mutable post-activation —
                    // operators can shorten an active window or roll it back
                    // by setting `end_time_ms` to the past.
                    if existing.is_sticky() {
                        if existing == &entry {
                            // No-op writes are permitted (idempotent re-broadcast).
                            return Ok(());
                        }
                        return Err(FeatureManagementError::EntryFrozen);
                    }
                } else {
                    // Scheduled phase: enforce sticky invariants.
                    if existing.is_sticky() && !entry.is_sticky() {
                        return Err(FeatureManagementError::StickyCannotBecomeWindowed);
                    }
                    // Forward-only `start_time_ms` applies whenever the *new*
                    // entry is sticky — sticky→sticky preservation and
                    // windowed→sticky upgrade both protect the operator-notice
                    // window from being pulled earlier.
                    if entry.is_sticky() && entry.start_time_ms < existing.start_time_ms {
                        return Err(FeatureManagementError::StickyStartTimeBackward);
                    }
                }
            }
        }

        self.entries.insert(name, entry);
        Ok(())
    }

    /// Remove a feature entry. Sticky entries are not removable in either
    /// phase — `StickyEntryNotRemovable` is returned in that case.
    pub fn remove(&mut self, name: &str) -> Result<(), FeatureManagementError> {
        if let Some(existing) = self.entries.get(name) {
            if existing.is_sticky() {
                return Err(FeatureManagementError::StickyEntryNotRemovable);
            }
            self.entries.remove(name);
        }
        Ok(())
    }

    fn validate_entry_shape(entry: &FeatureEntry) -> Result<(), FeatureManagementError> {
        if entry.start_time_ms > MAX_TIMESTAMP_MS {
            return Err(FeatureManagementError::TimeOutOfRange);
        }
        if let Some(end) = entry.end_time_ms {
            if end > MAX_TIMESTAMP_MS {
                return Err(FeatureManagementError::TimeOutOfRange);
            }
            if entry.start_time_ms >= end {
                return Err(FeatureManagementError::InvalidTimeRange);
            }
        }
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn windowed(start: u64, end: u64) -> FeatureEntry {
        FeatureEntry {
            start_time_ms: start,
            end_time_ms: Some(end),
        }
    }

    fn stick(start: u64) -> FeatureEntry {
        FeatureEntry {
            start_time_ms: start,
            end_time_ms: None,
        }
    }

    #[test]
    fn new_is_empty() {
        let authority = Pubkey::new_unique();
        let state = FeaturesState::new(authority);
        assert_eq!(state.get_authority(), &authority);
        assert!(state.entries.is_empty());
    }

    #[test]
    fn insert_windowed_ok() {
        let mut state = FeaturesState::new_for_test();
        state
            .upsert("foo".to_string(), windowed(100, 200), 50)
            .unwrap();
        assert_eq!(state.get("foo"), Some(&windowed(100, 200)));
    }

    #[test]
    fn insert_sticky_ok() {
        let mut state = FeaturesState::new_for_test();
        state.upsert("foo".to_string(), stick(100), 50).unwrap();
        assert_eq!(state.get("foo"), Some(&stick(100)));
    }

    #[test]
    fn is_active_window_semantics() {
        let mut state = FeaturesState::new_for_test();
        state
            .upsert("foo".to_string(), windowed(100, 200), 50)
            .unwrap();
        assert!(!state.is_active("foo", 99));
        assert!(state.is_active("foo", 100));
        assert!(state.is_active("foo", 199));
        assert!(!state.is_active("foo", 200));
        assert!(!state.is_active("missing", 150));
    }

    #[test]
    fn insert_rejects_invalid_range() {
        let mut state = FeaturesState::new_for_test();
        let err = state
            .upsert("foo".to_string(), windowed(200, 100), 0)
            .unwrap_err();
        assert_eq!(err, FeatureManagementError::InvalidTimeRange);
        let err = state
            .upsert("foo".to_string(), windowed(100, 100), 0)
            .unwrap_err();
        assert_eq!(err, FeatureManagementError::InvalidTimeRange);
    }

    #[test]
    fn insert_rejects_time_out_of_range() {
        let mut state = FeaturesState::new_for_test();
        let too_far = MAX_TIMESTAMP_MS + 1;
        let err = state
            .upsert("foo".to_string(), windowed(too_far, too_far + 100), 0)
            .unwrap_err();
        assert_eq!(err, FeatureManagementError::TimeOutOfRange);
        let err = state
            .upsert("bar".to_string(), windowed(0, too_far), 0)
            .unwrap_err();
        assert_eq!(err, FeatureManagementError::TimeOutOfRange);
    }

    #[test]
    fn scheduled_sticky_cannot_become_windowed() {
        let mut state = FeaturesState::new_for_test();
        state.upsert("foo".to_string(), stick(100), 50).unwrap();
        let err = state
            .upsert("foo".to_string(), windowed(100, 200), 50)
            .unwrap_err();
        assert_eq!(err, FeatureManagementError::StickyCannotBecomeWindowed);
    }

    #[test]
    fn scheduled_sticky_start_time_only_forward() {
        let mut state = FeaturesState::new_for_test();
        state.upsert("foo".to_string(), stick(100), 50).unwrap();
        let err = state.upsert("foo".to_string(), stick(99), 50).unwrap_err();
        assert_eq!(err, FeatureManagementError::StickyStartTimeBackward);
        // Delay (forward move) is permitted.
        state.upsert("foo".to_string(), stick(150), 50).unwrap();
        assert_eq!(state.get("foo").unwrap().start_time_ms, 150);
    }

    #[test]
    fn sticky_is_active_after_start() {
        let mut state = FeaturesState::new_for_test();
        state.upsert("foo".to_string(), stick(100), 50).unwrap();
        assert!(!state.is_active("foo", 99));
        assert!(state.is_active("foo", 100));
        assert!(state.is_active("foo", u64::MAX));
    }

    #[test]
    fn scheduled_windowed_can_upgrade_to_stick() {
        let mut state = FeaturesState::new_for_test();
        state
            .upsert("foo".to_string(), windowed(100, 200), 50)
            .unwrap();
        // Same `start_time_ms` (operator-notice window preserved) — allowed.
        state.upsert("foo".to_string(), stick(100), 50).unwrap();
        let entry = state.get("foo").unwrap();
        assert!(entry.is_sticky());
        assert_eq!(entry.start_time_ms, 100);
    }

    #[test]
    fn scheduled_windowed_upgrade_to_sticky_can_delay() {
        let mut state = FeaturesState::new_for_test();
        state
            .upsert("foo".to_string(), windowed(100, 200), 50)
            .unwrap();
        // Forward move is permitted on upgrade — operator can announce the
        // permanent activation later than the original windowed start.
        state.upsert("foo".to_string(), stick(150), 50).unwrap();
        let entry = state.get("foo").unwrap();
        assert!(entry.is_sticky());
        assert_eq!(entry.start_time_ms, 150);
    }

    #[test]
    fn scheduled_windowed_upgrade_to_sticky_rejects_backward_start() {
        let mut state = FeaturesState::new_for_test();
        state
            .upsert("foo".to_string(), windowed(100, 200), 50)
            .unwrap();
        // Pulling activation earlier on an upgrade would collapse the
        // operator-notice window the sticky rules exist to protect.
        let err = state.upsert("foo".to_string(), stick(99), 50).unwrap_err();
        assert_eq!(err, FeatureManagementError::StickyStartTimeBackward);
        // The existing windowed entry stays intact.
        let entry = state.get("foo").unwrap();
        assert!(!entry.is_sticky());
        assert_eq!(entry.start_time_ms, 100);
    }

    #[test]
    fn active_phase_freezes_sticky_entry() {
        // Sticky entries become immutable once `current_time_ms` crosses
        // their `start_time_ms`; subsequent upserts that change anything
        // are rejected.
        let mut state = FeaturesState::new_for_test();
        state.upsert("foo".to_string(), stick(100), 50).unwrap();
        let err = state
            .upsert("foo".to_string(), stick(110), 150)
            .unwrap_err();
        assert_eq!(err, FeatureManagementError::EntryFrozen);
    }

    #[test]
    fn active_phase_allows_windowed_modification() {
        // Windowed entries stay mutable post-activation: operators can
        // shorten the window or roll it back by setting `end_time_ms` to
        // a past timestamp. Only sticky entries freeze.
        let mut state = FeaturesState::new_for_test();
        state
            .upsert("foo".to_string(), windowed(100, 200), 50)
            .unwrap();
        // Extend the window.
        state
            .upsert("foo".to_string(), windowed(100, 300), 150)
            .unwrap();
        assert_eq!(state.get("foo").unwrap().end_time_ms, Some(300));
        // Roll back end_time_ms to the past — effectively deactivates.
        state
            .upsert("foo".to_string(), windowed(100, 140), 150)
            .unwrap();
        assert_eq!(state.get("foo").unwrap().end_time_ms, Some(140));
    }

    #[test]
    fn active_phase_sticky_no_op_write_ok() {
        let mut state = FeaturesState::new_for_test();
        state.upsert("foo".to_string(), stick(100), 50).unwrap();
        // Identical entry during active phase succeeds (idempotent rebroadcast).
        state.upsert("foo".to_string(), stick(100), 150).unwrap();
    }

    #[test]
    fn windowed_scheduled_phase_freely_modifiable() {
        let mut state = FeaturesState::new_for_test();
        state
            .upsert("foo".to_string(), windowed(100, 200), 50)
            .unwrap();
        state
            .upsert("foo".to_string(), windowed(120, 220), 50)
            .unwrap();
        assert_eq!(state.get("foo"), Some(&windowed(120, 220)));
    }

    #[test]
    fn remove_sticky_rejected() {
        let mut state = FeaturesState::new_for_test();
        state.upsert("foo".to_string(), stick(100), 50).unwrap();
        assert_eq!(
            state.remove("foo"),
            Err(FeatureManagementError::StickyEntryNotRemovable)
        );
        assert!(state.get("foo").is_some());
    }

    #[test]
    fn remove_windowed_ok() {
        let mut state = FeaturesState::new_for_test();
        state
            .upsert("foo".to_string(), windowed(100, 200), 50)
            .unwrap();
        state.remove("foo").unwrap();
        assert!(state.get("foo").is_none());
    }

    #[test]
    fn remove_missing_is_noop() {
        let mut state = FeaturesState::new_for_test();
        state.remove("never-existed").unwrap();
    }

    #[test]
    fn max_feature_count_enforced() {
        let mut state = FeaturesState::new_for_test();
        for i in 0..crate::MAX_FEATURE_COUNT {
            let name = alloc::format!("f{i}");
            state.upsert(name, windowed(100, 200), 50).unwrap();
        }
        let err = state
            .upsert("overflow".to_string(), windowed(100, 200), 50)
            .unwrap_err();
        assert_eq!(err, FeatureManagementError::MaxFeatureCountExceeded);
    }

    #[test]
    fn serialize_round_trip() {
        let mut state = FeaturesState::new_for_test();
        state
            .upsert("foo".to_string(), windowed(100, 200), 50)
            .unwrap();
        state.upsert("bar".to_string(), stick(300), 50).unwrap();
        let bytes = state.serialize().unwrap();
        let back = FeaturesState::deserialize(&bytes).unwrap();
        assert_eq!(state, back);
    }

    #[test]
    fn wire_format_golden() {
        // Pins the borsh layout of the production schema so an accidental
        // schema change (variant reorder, field rename, type swap) trips this
        // test rather than silently breaking on-disk state.
        let mut state = FeaturesState {
            authority: Pubkey::new_from_array([7u8; 32]),
            entries: BTreeMap::new(),
        };
        state.entries.insert("alpha".to_string(), windowed(1, 2));
        state.entries.insert("beta".to_string(), stick(1_000));

        // Layout:
        //   authority:        32 bytes
        //   entries (BTreeMap<String, FeatureEntry>) borsh layout:
        //     u32 LE len
        //     repeat: { u32 LE name_len, name bytes, FeatureEntry }
        //   FeatureEntry: u64 LE start_time_ms, Option<u64> end_time_ms
        //                 (Option borsh: u8 tag 0=None, 1=Some + u64 LE)
        let mut expected: Vec<u8> = Vec::new();
        expected.extend_from_slice(&[7u8; 32]);
        expected.extend_from_slice(&2u32.to_le_bytes()); // 2 entries
                                                         // "alpha" -> windowed(1, 2)
        expected.extend_from_slice(&5u32.to_le_bytes());
        expected.extend_from_slice(b"alpha");
        expected.extend_from_slice(&1u64.to_le_bytes());
        expected.push(1); // Some
        expected.extend_from_slice(&2u64.to_le_bytes());
        // "beta" -> stick(1_000)
        expected.extend_from_slice(&4u32.to_le_bytes());
        expected.extend_from_slice(b"beta");
        expected.extend_from_slice(&1_000u64.to_le_bytes());
        expected.push(0); // None

        let actual = state.serialize().unwrap();
        assert_eq!(actual, expected);

        // And it must round-trip.
        let back = FeaturesState::deserialize(&actual).unwrap();
        assert_eq!(back, state);
    }

    #[test]
    fn set_authority_works() {
        let mut state = FeaturesState::new_for_test();
        let new_auth = Pubkey::new_unique();
        state.set_authority(new_auth);
        assert_eq!(state.get_authority(), &new_auth);
    }
}