futu_cache/trd_cache/

cipher_exchange.rs

use std::collections::{BTreeMap, BTreeSet};
use std::sync::atomic::Ordering;

use dashmap::mapref::entry::Entry;
use futu_core::trade_broker::broker_id_for_security_firm_like_cpp;

use super::{AccKey, TrdCache};

#[derive(Clone, Copy, PartialEq, Eq)]
pub(super) enum CipherFreshness {
    Unbound,
    Pending {
        broker_id: u32,
        connection_epoch: u64,
        request_id: u64,
    },
    Fresh {
        broker_id: u32,
        connection_epoch: u64,
    },
    Failed {
        broker_id: u32,
        connection_epoch: u64,
    },
}

/// Secret-bearing cache record. Deliberately does not implement `Debug`.
pub(super) struct CipherRecord {
    revision: u64,
    cipher: Option<Vec<u8>>,
    freshness: CipherFreshness,
}

/// One CMD2902 request row bound to the account revision captured at send time.
/// Deliberately does not implement `Debug` because it owns cipher bytes.
pub struct CipherExchangeAccount {
    account_id: AccKey,
    revision: u64,
    trade_cipher: Vec<u8>,
}

impl CipherExchangeAccount {
    #[must_use]
    pub const fn account_id(&self) -> AccKey {
        self.account_id
    }

    #[must_use]
    pub const fn revision(&self) -> u64 {
        self.revision
    }

    #[must_use]
    pub fn trade_cipher(&self) -> &[u8] {
        &self.trade_cipher
    }
}

/// In-flight CMD2902 ownership token. `request_id` distinguishes overlapping
/// requests on the same connection epoch; account revision rejects a response
/// racing with unlock/lock/reload.
pub struct CipherExchangeLease {
    broker_id: u32,
    connection_epoch: u64,
    request_id: u64,
    broker_generation: u64,
    accounts: Vec<CipherExchangeAccount>,
}

impl CipherExchangeLease {
    #[must_use]
    pub const fn broker_id(&self) -> u32 {
        self.broker_id
    }

    #[must_use]
    pub const fn connection_epoch(&self) -> u64 {
        self.connection_epoch
    }

    #[must_use]
    pub const fn request_id(&self) -> u64 {
        self.request_id
    }

    #[must_use]
    pub const fn broker_generation(&self) -> u64 {
        self.broker_generation
    }

    #[must_use]
    pub fn accounts(&self) -> &[CipherExchangeAccount] {
        &self.accounts
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CipherExchangeError {
    DuplicateAccount,
    EmptyCipher,
    AccountSetMismatch,
    LeaseStale,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct CipherExchangePublishReport {
    published: usize,
    stale: usize,
    broker_generation: u64,
}

impl CipherExchangePublishReport {
    #[must_use]
    pub const fn published(self) -> usize {
        self.published
    }

    #[must_use]
    pub const fn stale(self) -> usize {
        self.stale
    }

    #[must_use]
    pub const fn broker_generation(self) -> u64 {
        self.broker_generation
    }
}

impl TrdCache {
    #[must_use]
    pub fn get_cipher(&self, acc_id: AccKey) -> Option<Vec<u8>> {
        self.cipher_records
            .get(&acc_id)
            .and_then(|record| record.cipher.clone())
    }

    /// Publish a newly unlocked cipher and invalidate every in-flight CMD2902
    /// lease for this account by incrementing its revision.
    pub fn set_cipher(&self, acc_id: AccKey, cipher: Vec<u8>) {
        let broker_id = self.broker_id_for_cipher_account(acc_id);
        let _publication = self.cipher_publication_lock.lock();
        match self.cipher_records.entry(acc_id) {
            Entry::Occupied(mut entry) => {
                let record = entry.get_mut();
                record.revision = record.revision.saturating_add(1);
                record.cipher = Some(cipher);
                record.freshness = CipherFreshness::Unbound;
            }
            Entry::Vacant(entry) => {
                entry.insert(CipherRecord {
                    revision: 1,
                    cipher: Some(cipher),
                    freshness: CipherFreshness::Unbound,
                });
            }
        }
        if let Some(broker_id) = broker_id {
            self.bump_broker_cipher_generation_locked(broker_id);
        }
    }

    #[must_use]
    pub fn get_cipher_revision(&self, acc_id: AccKey) -> u64 {
        self.cipher_records
            .get(&acc_id)
            .map_or(0, |record| record.revision)
    }

    #[must_use]
    pub fn cipher_count(&self) -> usize {
        self.cipher_records
            .iter()
            .filter(|record| record.cipher.is_some())
            .count()
    }

    /// Clear one account cipher while retaining a revision tombstone.
    #[must_use]
    pub fn clear_cipher(&self, acc_id: AccKey) -> bool {
        let broker_id = self.broker_id_for_cipher_account(acc_id);
        let _publication = self.cipher_publication_lock.lock();
        let Some(mut record) = self.cipher_records.get_mut(&acc_id) else {
            return false;
        };
        if record.cipher.take().is_none() {
            return false;
        }
        record.revision = record.revision.saturating_add(1);
        record.freshness = CipherFreshness::Failed {
            broker_id: 0,
            connection_epoch: 0,
        };
        if let Some(broker_id) = broker_id {
            self.bump_broker_cipher_generation_locked(broker_id);
        }
        true
    }

    #[must_use]
    pub fn clear_all_ciphers_and_bump_versions(&self) -> (usize, Vec<(u64, u64)>) {
        let acc_ids: Vec<_> = self
            .cipher_records
            .iter()
            .filter_map(|record| record.cipher.as_ref().map(|_| *record.key()))
            .collect();
        let mut bumped = Vec::with_capacity(acc_ids.len());
        for acc_id in acc_ids {
            if self.clear_cipher(acc_id) {
                bumped.push((acc_id, self.bump_cipher_state_version(acc_id)));
            }
        }
        (bumped.len(), bumped)
    }

    /// Freeze the exact real-account cipher set for one broker connection.
    /// C++ sorts the broker account list before building CMD2902; Rust preserves
    /// deterministic account order for the same wire behavior.
    #[must_use]
    pub fn begin_cipher_exchange(
        &self,
        broker_id: u32,
        connection_epoch: u64,
    ) -> CipherExchangeLease {
        let _publication = self.cipher_publication_lock.lock();
        let request_id = self
            .next_cipher_exchange_request_id
            .fetch_add(1, Ordering::SeqCst)
            .wrapping_add(1);
        let mut account_ids: Vec<_> = self
            .accounts
            .iter()
            .filter(|account| account.trd_env == 1)
            .filter(|account| {
                account
                    .security_firm
                    .and_then(broker_id_for_security_firm_like_cpp)
                    == Some(broker_id)
            })
            .map(|account| *account.key())
            .collect();
        account_ids.sort_unstable();

        let mut accounts = Vec::new();
        for account_id in account_ids {
            let Some(mut record) = self.cipher_records.get_mut(&account_id) else {
                continue;
            };
            let Some(cipher) = record
                .cipher
                .as_ref()
                .filter(|cipher| !cipher.is_empty())
                .cloned()
            else {
                continue;
            };
            let revision = record.revision;
            record.freshness = CipherFreshness::Pending {
                broker_id,
                connection_epoch,
                request_id,
            };
            accounts.push(CipherExchangeAccount {
                account_id,
                revision,
                trade_cipher: cipher,
            });
        }

        CipherExchangeLease {
            broker_id,
            connection_epoch,
            request_id,
            broker_generation: self.broker_cipher_generation_locked(broker_id),
            accounts,
        }
    }

    /// Complete CMD2902 only after validating an exact, duplicate-free,
    /// non-empty response set. Each account publish is then guarded by the
    /// captured revision + broker + connection epoch + request id CAS tuple.
    pub fn complete_cipher_exchange(
        &self,
        lease: &CipherExchangeLease,
        response: Vec<(AccKey, Vec<u8>)>,
    ) -> Result<CipherExchangePublishReport, CipherExchangeError> {
        let expected: BTreeSet<_> = lease.accounts.iter().map(|row| row.account_id).collect();
        let mut response_by_account = BTreeMap::new();
        for (account_id, cipher) in response {
            if cipher.is_empty() {
                self.fail_cipher_exchange(lease);
                return Err(CipherExchangeError::EmptyCipher);
            }
            if response_by_account.insert(account_id, cipher).is_some() {
                self.fail_cipher_exchange(lease);
                return Err(CipherExchangeError::DuplicateAccount);
            }
        }
        let actual: BTreeSet<_> = response_by_account.keys().copied().collect();
        if actual != expected {
            self.fail_cipher_exchange(lease);
            return Err(CipherExchangeError::AccountSetMismatch);
        }

        let _publication = self.cipher_publication_lock.lock();
        if self.broker_cipher_generation_locked(lease.broker_id) != lease.broker_generation {
            self.fail_cipher_exchange_locked(lease);
            return Err(CipherExchangeError::LeaseStale);
        }
        let mut published = 0;
        let mut stale = 0;
        for requested in &lease.accounts {
            let Some(mut record) = self.cipher_records.get_mut(&requested.account_id) else {
                stale += 1;
                continue;
            };
            let owns_lease = record.revision == requested.revision
                && record.freshness
                    == (CipherFreshness::Pending {
                        broker_id: lease.broker_id,
                        connection_epoch: lease.connection_epoch,
                        request_id: lease.request_id,
                    });
            if !owns_lease {
                stale += 1;
                continue;
            }
            record.revision = record.revision.saturating_add(1);
            record.cipher = response_by_account.remove(&requested.account_id);
            record.freshness = CipherFreshness::Fresh {
                broker_id: lease.broker_id,
                connection_epoch: lease.connection_epoch,
            };
            published += 1;
            self.bump_broker_cipher_generation_locked(lease.broker_id);
        }
        Ok(CipherExchangePublishReport {
            published,
            stale,
            broker_generation: self.broker_cipher_generation_locked(lease.broker_id),
        })
    }

    /// Mark only the still-owned lease rows failed. A newer unlock/exchange is
    /// left untouched.
    pub fn fail_cipher_exchange(&self, lease: &CipherExchangeLease) {
        let _publication = self.cipher_publication_lock.lock();
        self.fail_cipher_exchange_locked(lease);
    }

    fn fail_cipher_exchange_locked(&self, lease: &CipherExchangeLease) {
        for requested in &lease.accounts {
            let Some(mut record) = self.cipher_records.get_mut(&requested.account_id) else {
                continue;
            };
            if record.revision == requested.revision
                && record.freshness
                    == (CipherFreshness::Pending {
                        broker_id: lease.broker_id,
                        connection_epoch: lease.connection_epoch,
                        request_id: lease.request_id,
                    })
            {
                record.freshness = CipherFreshness::Failed {
                    broker_id: lease.broker_id,
                    connection_epoch: lease.connection_epoch,
                };
            }
        }
    }

    /// Run route publication only while the broker's cipher generation still
    /// equals the completed/empty lease generation. Cipher set/clear/publish
    /// share this lock, closing the last-check-to-install race.
    pub fn with_current_broker_cipher_generation<T>(
        &self,
        broker_id: u32,
        expected_generation: u64,
        publish: impl FnOnce() -> T,
    ) -> Option<T> {
        let _publication = self.cipher_publication_lock.lock();
        (self.broker_cipher_generation_locked(broker_id) == expected_generation).then(publish)
    }

    /// Financial write gate: pending/failed/stale-epoch ciphers are not usable.
    /// A newly unlocked cipher is initially unbound and remains usable until a
    /// reconnect lease explicitly transitions it to Pending.
    #[must_use]
    pub fn get_cipher_for_broker_epoch(
        &self,
        acc_id: AccKey,
        broker_id: u32,
        connection_epoch: u64,
    ) -> Option<Vec<u8>> {
        let account = self.accounts.get(&acc_id)?;
        if account.trd_env != 1
            || account
                .security_firm
                .and_then(broker_id_for_security_firm_like_cpp)
                != Some(broker_id)
        {
            return None;
        }
        let record = self.cipher_records.get(&acc_id)?;
        let usable = match record.freshness {
            CipherFreshness::Unbound => true,
            CipherFreshness::Fresh {
                broker_id: fresh_broker,
                connection_epoch: fresh_epoch,
            } => fresh_broker == broker_id && fresh_epoch == connection_epoch,
            CipherFreshness::Pending { .. } | CipherFreshness::Failed { .. } => false,
        };
        usable.then(|| record.cipher.clone()).flatten()
    }

    fn broker_id_for_cipher_account(&self, acc_id: AccKey) -> Option<u32> {
        let account = self.accounts.get(&acc_id)?;
        (account.trd_env == 1)
            .then(|| {
                account
                    .security_firm
                    .and_then(broker_id_for_security_firm_like_cpp)
            })
            .flatten()
    }

    fn broker_cipher_generation_locked(&self, broker_id: u32) -> u64 {
        self.cipher_broker_generations
            .get(&broker_id)
            .map_or(0, |generation| *generation.value())
    }

    fn bump_broker_cipher_generation_locked(&self, broker_id: u32) -> u64 {
        let mut generation = self.cipher_broker_generations.entry(broker_id).or_insert(0);
        *generation = generation.saturating_add(1);
        *generation
    }
}