feat(be): cap sessions per identity

src/internet_identity/src/storage.rs

@@ -322,6 +322,21 @@ const MAX_EVICTABLE_DEFAULT_ACCOUNTS: u64 = 500;
 /// Eviction target, below the cap.
 const EVICTABLE_DEFAULT_ACCOUNTS_WATERMARK: u64 = MAX_EVICTABLE_DEFAULT_ACCOUNTS * 9 / 10;
 
+/// Session records one identity may hold, counted as stored rather than as live.
+///
+/// Counting what is stored is what makes the cap cheap to trigger on: a session expires with
+/// no write anywhere, so no counter can follow the live set — something would have to
+/// decrement at the moment of expiry, and nothing runs then. An expired record holds its slot
+/// until something reclaims it, and because it is the first thing reclaimed, a held slot is
+/// never taken from a session in use.
+///
+/// A bound on concurrent activity, not on history: every session expires within 30 days, so
+/// the set is the apps used in the last month times the browsers they were used from.
+pub const MAX_SESSIONS_PER_ANCHOR: u32 = 500;
+/// Reclaiming goes down to here rather than to the cap, so the pass that walks an identity's
+/// rows runs once and then not again for the next fifty sign-ins.
+pub const SESSIONS_WATERMARK_PER_ANCHOR: u32 = 450;
+
 /// Bounds one message's eviction work.
 const MAX_EVICTIONS_PER_CALL: u64 =
     MAX_EVICTABLE_DEFAULT_ACCOUNTS - EVICTABLE_DEFAULT_ACCOUNTS_WATERMARK;
@@ -1855,6 +1870,142 @@ impl<M: Memory + Clone> Storage<M> {
         Some(canister_sig_principal(canister_id(), seed.to_vec()))
     }
 
+    /// Frees a slot for one more session, and reports whether the anchor has one.
+    ///
+    /// The stored count is a trigger, never the thing the cap is enforced against: a
+    /// session can expire with no write anywhere, so the count drifts upwards. Once it
+    /// reaches the cap this recounts what the rows hold and reclaims against that, so an
+    /// admission is only ever granted against a number that was just counted.
+    fn ensure_session_slot(
+        &mut self,
+        anchor_number: AnchorNumber,
+        now: Timestamp,
+    ) -> Result<bool, StorageError> {
+        if self.read(anchor_number)?.session_count < MAX_SESSIONS_PER_ANCHOR {
+            return Ok(true);
+        }
+        Ok(self.reclaim_sessions(anchor_number, now)? < MAX_SESSIONS_PER_ANCHOR)
+    }
+
+    /// Replaces the count with a number that was counted rather than accumulated.
+    fn set_session_count(
+        &mut self,
+        anchor_number: AnchorNumber,
+        count: u32,
+    ) -> Result<(), StorageError> {
+        let mut anchor = self.read(anchor_number)?;
+        if anchor.session_count == count {
+            return Ok(());
+        }
+        anchor.session_count = count;
+        self.write(anchor)
+    }
+
+    /// Walks the anchor's rows once and reclaims down to the watermark, taking sessions in
+    /// [`SessionRecord::reclaim_order`]: dead ones first, then the least recently used.
+    ///
+    /// Returns what the rows actually hold once it is done, which is the number the cap is
+    /// enforced against. The stored counter is only ever a trigger for running this pass —
+    /// it can drift, this cannot, because it counts the sessions themselves.
+    ///
+    /// One pass per fifty sign-ins, because it reclaims to the watermark rather than to the
+    /// cap, and bounded by the same row limit account eviction uses.
+    fn reclaim_sessions(
+        &mut self,
+        anchor_number: AnchorNumber,
+        now: Timestamp,
+    ) -> Result<u32, StorageError> {
+        struct Candidate {
+            order: (bool, Timestamp, SessionDeviceId),
+            row: usize,
+            account_number: Option<AccountNumber>,
+            device_id: SessionDeviceId,
+        }
+
+        // Every row, not a bounded prefix of them: the number this returns is what the cap is
+        // enforced against, and a truncated scan would undercount, lower the counter to the
+        // undercount, and let the stored set climb past the cap from there. An identity's rows
+        // are already bounded — the row cap holds the evictable ones and the account cap holds
+        // the rest — and a sequential scan of them costs a fraction of the writes it saves.
+        let mut rows: Vec<(ApplicationNumber, Vec<AccountReference>)> = self
+            .stable_account_reference_list_memory
+            .range(
+                (anchor_number, ApplicationNumber::MIN)..=(anchor_number, ApplicationNumber::MAX),
+            )
+            .map(|((_, application_number), list)| (application_number, list.into()))
+            .collect();
+
+        let mut candidates: Vec<Candidate> = vec![];
+        for (row, (_, references)) in rows.iter().enumerate() {
+            for reference in references {
+                for session in &reference.sessions {
+                    candidates.push(Candidate {
+                        order: session.reclaim_order(now),
+                        row,
+                        account_number: reference.account_number,
+                        device_id: session.device_id,
+                    });
+                }
+            }
+        }
+        let stored = candidates.len() as u32;
+        candidates.sort_by_key(|candidate| candidate.order);
+
+        let surplus = stored.saturating_sub(SESSIONS_WATERMARK_PER_ANCHOR) as usize;
+        let victims = &candidates[..surplus.min(candidates.len())];
+        if victims.is_empty() {
+            self.set_session_count(anchor_number, stored)?;
+            return Ok(stored);
+        }
+
+        // One write per row rather than one per victim: the row is a single blob, so
+        // dropping several of its sessions one at a time would rewrite it several times.
+        let mut touched: Vec<usize> = victims.iter().map(|victim| victim.row).collect();
+        touched.sort_unstable();
+        touched.dedup();
+
+        let mut dropped_total = 0usize;
+        for row in touched {
+            let (application_number, references) = &mut rows[row];
+            let application_number = *application_number;
+            let mut removed: Vec<(Option<AccountNumber>, SessionRecord)> = vec![];
+            for reference in references.iter_mut() {
+                let account_number = reference.account_number;
+                reference.sessions.retain(|session| {
+                    // The row has to be part of the match: one browser holds one session per
+                    // account, but the same browser and the same account number appear in
+                    // every row, so matching on that pair alone reaches across applications.
+                    let doomed = victims.iter().any(|victim| {
+                        victim.row == row
+                            && victim.account_number == account_number
+                            && victim.device_id == session.device_id
+                    });
+                    if doomed {
+                        removed.push((account_number, session.clone()));
+                    }
+                    !doomed
+                });
+            }
+            if removed.is_empty() {
+                continue;
+            }
+            self.write_reference_list(anchor_number, application_number, references.clone())?;
+            for (account_number, session) in &removed {
+                self.unindex_sessions(
+                    anchor_number,
+                    application_number,
+                    *account_number,
+                    std::slice::from_ref(session),
+                );
+            }
+            dropped_total += removed.len();
+        }
+
+        let remaining = stored.saturating_sub(dropped_total as u32);
+        self.set_session_count(anchor_number, remaining)?;
+        Ok(remaining)
+    }
+
     /// Moves the count without considering the cap, for the paths that only remove.
     fn change_session_count(
         &mut self,
@@ -2001,6 +2152,12 @@ impl<M: Memory + Clone> Storage<M> {
             }
         };
 
+        // Reclaiming before the session is admitted rather than after it: the stored set
+        // never sits above the cap, not even for the rest of this message.
+        if !self.ensure_session_slot(anchor_number, now_ns)? {
+            return Err(StorageError::SessionCapNotReclaimed { anchor_number });
+        }
+
         let mut references: Vec<AccountReference> = self
             .lookup_account_references(anchor_number, application_number)
             .ok_or(StorageError::MissingAccount {