feat(be): record that a session is still in use

src/internet_identity/src/sessions.rs

@@ -355,16 +355,33 @@ pub fn app_prepare_delegation(
 ) -> Result<AppPrepareDelegationResponse, AppSessionError> {
     let now = time();
     let AuthorizedSession {
-        account, session, ..
+        locator,
+        account,
+        session,
     } = authorize_session(now)?;
 
+    // Everything that can refuse, before the stamp. Returning `Err` on the IC commits
+    // whatever was written before it — only a trap rolls back — so a stamp above this
+    // would leave the session recorded as used while the caller is told the call failed.
+    // All three depend on values already in hand, so there is nothing to gain by
+    // computing them later.
     let expiration = u64::min(
         now.saturating_add(APP_DELEGATION_TTL_NS),
         session.valid_till_ns,
     );
     let seed = account_seed(&account)?;
     let access = DelegationAccess::from_read_only(session.read_only);
 
+    // A session revoked between `authorize_session` above and this stamp is a race, not
+    // an internal fault: the answer is the one the caller would have got a moment
+    // earlier.
+    storage_borrow_mut(|storage| storage.record_session_use(&locator, now)).map_err(
+        |err| match err {
+            StorageError::SessionNotFound { .. } => AppSessionError::NoSuchSession,
+            other => AppSessionError::InternalCanisterError(other.to_string()),
+        },
+    )?;
+
     state::signature_map_mut(|sigs| {
         add_delegation_signature(
             sigs,
@@ -440,8 +457,6 @@ pub fn app_get_delegation(
 /// three values a caller gathered: the principal lookup and the liveness check have both
 /// happened, and no field can be here without them.
 struct AuthorizedSession {
-    // Read by the refresh stamp, which lands one PR up.
-    #[allow(dead_code)]
     locator: SessionLocator,
     account: Account,
     session: Session,

src/internet_identity/src/storage.rs

@@ -3331,6 +3331,73 @@ impl<M: Memory + Clone> Storage<M> {
         })
     }
 
+    /// Records that a session was used: its own stamp, its account reference's, and the
+    /// browser's in the device registry.
+    ///
+    /// A locator naming nothing is [`StorageError::SessionNotFound`] rather than a
+    /// quiet non-event. Nobody makes a decision from it, and the one caller that could
+    /// have ignored it went on to mint a delegation for a session no list holds.
+    pub fn record_session_use(
+        &mut self,
+        key: &SessionLocator,
+        now: Timestamp,
+    ) -> Result<(), StorageError> {
+        let SessionLocator {
+            anchor_number,
+            origin,
+            account_number,
+            session_id,
+        } = key;
+        let (anchor_number, account_number, session_id) =
+            (*anchor_number, *account_number, *session_id);
+        let not_found = || StorageError::SessionNotFound {
+            anchor_number,
+            session_id,
+        };
+
+        if self.lookup_application_number_with_origin(origin).is_none() {
+            return Err(not_found());
+        }
+        let mut anchor = self.read(anchor_number)?;
+        let (mut account_references, config) = self.account_state_for_origin(anchor_number, origin);
+
+        let write = account_references
+            .iter_mut()
+            .find(|write| write.account_reference.account_number == account_number)
+            .ok_or_else(not_found)?;
+        let session = write
+            .account_reference
+            .sessions
+            .iter_mut()
+            .find(|session| session.session_id == session_id)
+            .ok_or_else(not_found)?;
+
+        session.last_refreshed_ns = Some(now);
+        let browser_id = session.browser_id;
+        write.account_reference.last_used = Some(now);
+
+        // This list is being rewritten anyway, so its dead sessions go now. It costs one
+        // pass over a list already in memory and no write of its own, and it means every
+        // list anyone still uses stays clean without anything having to sweep for it.
+        for write in account_references.iter_mut() {
+            write
+                .account_reference
+                .sessions
+                .retain(|session| !session.is_expired_or_idle(now));
+        }
+
+        // Stamped before the write rather than after, because the write is what stores the
+        // record. There is no second store: handing it over is handing over the storing of
+        // it, whatever was changed on it.
+        anchor.stamp_browser_use(browser_id, now);
+        self.write_account_state(
+            anchor,
+            now,
+            BTreeMap::from([(origin.clone(), Some((account_references, config)))]),
+        )?;
+        Ok(())
+    }
+
     /// Retires an application no anchor references any more. The number is never
     /// reissued.
     fn remove_unreferenced_application(
@@ -4276,6 +4343,13 @@ pub enum StorageError {
     SessionAlreadyOver {
         anchor_number: AnchorNumber,
     },
+    /// The session a caller named is not among the identity's — whether it never was, or
+    /// has since been revoked, replaced or pruned. Those are one observation rather than
+    /// three: a session that is not there cannot be told apart from one that never was.
+    SessionNotFound {
+        anchor_number: AnchorNumber,
+        session_id: SessionId,
+    },
     AnchorNumberOutOfRange {
         anchor_number: AnchorNumber,
         range: (AnchorNumber, AnchorNumber),
@@ -4368,6 +4442,13 @@ impl fmt::Display for StorageError {
                 f,
                 "a session for Identity Anchor {anchor_number} would be over before it started"
             ),
+            Self::SessionNotFound {
+                anchor_number,
+                session_id,
+            } => write!(
+                f,
+                "Identity Anchor {anchor_number} holds no session {session_id}"
+            ),
             Self::DeserializationError(err) => {
                 write!(f, "failed to deserialize a Candid value: {err}")
             }

src/internet_identity/src/storage/anchor.rs

@@ -766,6 +766,21 @@ impl Anchor {
         }
     }
 
+    /// Advances a browser's `last_used`, where the anchor holds that browser and the
+    /// stamp moves it forward. A browser no entry names, or a repeat inside one message,
+    /// leaves the registry as it is.
+    pub fn stamp_browser_use(&mut self, browser_id: BrowserId, now: Timestamp) {
+        if let Some(browser) = self
+            .browsers
+            .iter_mut()
+            .find(|browser| browser.id == browser_id)
+        {
+            if browser.last_used < now {
+                browser.last_used = now;
+            }
+        }
+    }
+
     /// What a caller outside storage may know about this anchor's browsers: an
     /// identifier, what the browser said it was, and when. The keys stay here — they
     /// are how a sign-in proves which entry it is, so handing them out would let

src/internet_identity/src/storage/tests.rs

@@ -6380,6 +6380,217 @@ mod session_consent_change_tests {
     }
 }
 
+mod session_refresh_stamp_tests {
+    use super::held_references;
+    use super::params;
+    use crate::storage::account::{AccountReference, Session, SessionLocator};
+    use crate::storage::{CreateSessionParams, StorageError};
+    use crate::Storage;
+    use ic_stable_structures::VectorMemory;
+    use internet_identity_interface::internet_identity::types::AnchorNumber;
+    use pretty_assertions::assert_eq;
+
+    const ORIGIN: &str = "https://example.com";
+
+    fn storage_with_session() -> (Storage<VectorMemory>, AnchorNumber, SessionLocator) {
+        let mut storage = Storage::new((10_000, 3_784_873), VectorMemory::default());
+        storage.update_salt([17u8; 32]);
+        let anchor = storage.allocate_anchor(0).unwrap();
+        let anchor_number = anchor.anchor_number();
+        storage.write(anchor).unwrap();
+        let (key, _) = storage
+            .create_session(CreateSessionParams {
+                valid_till_ns: u64::MAX,
+                ..params(anchor_number, 1, 1_000)
+            })
+            .unwrap();
+        (storage, anchor_number, key)
+    }
+
+    fn reference(storage: &Storage<VectorMemory>, anchor_number: AnchorNumber) -> AccountReference {
+        let application_number = storage
+            .lookup_application_number_with_origin(&ORIGIN.to_string())
+            .unwrap();
+        held_references(storage, anchor_number, application_number)
+            .into_iter()
+            .find(|reference| reference.account_number.is_none())
+            .unwrap()
+    }
+
+    fn session_of(storage: &Storage<VectorMemory>, anchor_number: AnchorNumber) -> Session {
+        reference(storage, anchor_number).sessions.remove(0)
+    }
+
+    #[test]
+    fn a_refresh_stamps_the_session_and_the_reference() {
+        let (mut storage, anchor_number, key) = storage_with_session();
+
+        storage.record_session_use(&key, 2_000).unwrap();
+
+        assert_eq!(
+            session_of(&storage, anchor_number).last_refreshed_ns,
+            Some(2_000)
+        );
+        assert_eq!(reference(&storage, anchor_number).last_used, Some(2_000));
+    }
+
+    #[test]
+    fn every_refresh_advances_the_stamp() {
+        let (mut storage, anchor_number, key) = storage_with_session();
+
+        for now in [1_001, 1_002, 1_003] {
+            storage.record_session_use(&key, now).unwrap();
+            assert_eq!(
+                session_of(&storage, anchor_number).last_refreshed_ns,
+                Some(now)
+            );
+        }
+    }
+
+    /// The list is rewritten anyway, so the refresh is where a dead sibling is collected —
+    /// index entry and session count included, since nothing else will come for them.
+    #[test]
+    fn a_refresh_collects_the_dead_sessions_beside_it() {
+        let (mut storage, anchor_number, key) = storage_with_session();
+
+        let (_, dead) = storage
+            .create_session(CreateSessionParams {
+                valid_till_ns: 1_500,
+                ..params(anchor_number, 9, 1_000)
+            })
+            .unwrap();
+        let dead_principal = storage
+            .lookup_session_with_principal_memory
+            .iter()
+            .find(|(_, handle)| handle.session_id == dead.session_id)
+            .map(|(principal, _)| principal)
+            .expect("the session should be indexed");
+        assert!(storage
+            .lookup_session_with_principal(dead_principal)
+            .is_some());
+        assert_eq!(storage.read(anchor_number).unwrap().session_count, 2);
+
+        storage.record_session_use(&key, 2_000).unwrap();
+
+        let sessions = reference(&storage, anchor_number).sessions;
+        assert_eq!(sessions.len(), 1, "the expired sibling was left behind");
+        // The first browser to sign in, so the first id the registry minted.
+        assert_eq!(sessions[0].browser_id, 0);
+        assert!(
+            storage
+                .lookup_session_with_principal(dead_principal)
+                .is_none(),
+            "the expired sibling's index entry outlived it"
+        );
+        assert_eq!(storage.read(anchor_number).unwrap().session_count, 1);
+    }
+
+    /// A locator naming nothing is refused rather than reported as a quiet non-event.
+    /// The caller that would have ignored a `false` here goes on to mint a delegation,
+    /// which is the one thing a session no list holds must not get.
+    #[test]
+    fn a_stamp_for_a_session_that_is_gone_is_refused() {
+        let (mut storage, anchor_number, _key) = storage_with_session();
+
+        let refused = storage.record_session_use(
+            &SessionLocator {
+                anchor_number,
+                origin: ORIGIN.to_string(),
+                account_number: None,
+                session_id: 9_999,
+            },
+            5_000,
+        );
+
+        assert!(
+            matches!(
+                refused,
+                Err(StorageError::SessionNotFound {
+                    anchor_number: refused_anchor,
+                    session_id: 9_999
+                }) if refused_anchor == anchor_number
+            ),
+            "a locator naming no session should be refused, got {refused:?}"
+        );
+    }
+
+    #[test]
+    fn stamping_leaves_a_second_device_alone() {
+        let (mut storage, anchor_number, key) = storage_with_session();
+        storage
+            .create_session(CreateSessionParams {
+                valid_till_ns: u64::MAX,
+                ..params(anchor_number, 2, 1_000)
+            })
+            .unwrap();
+        let now = 2_000;
+
+        storage.record_session_use(&key, now).unwrap();
+
+        let sessions = reference(&storage, anchor_number).sessions;
+        assert_eq!(sessions.len(), 2);
+        // The registry mints the ids, in the order the browsers first signed in.
+        let stamped = sessions.iter().find(|s| s.browser_id == 0).unwrap();
+        let untouched = sessions.iter().find(|s| s.browser_id == 1).unwrap();
+        assert_eq!(stamped.last_refreshed_ns, Some(now));
+        assert_eq!(untouched.last_refreshed_ns, None);
+    }
+
+    fn storage_with_registered_browser(
+    ) -> (Storage<VectorMemory>, AnchorNumber, SessionLocator, u32) {
+        let mut storage = Storage::new((10_000, 3_784_873), VectorMemory::default());
+        storage.update_salt([17u8; 32]);
+        let anchor = storage.allocate_anchor(0).unwrap();
+        let anchor_number = anchor.anchor_number();
+        storage.write(anchor).unwrap();
+        // The sign-in registers the browser, so nothing here puts one on the record by
+        // hand and then claims its id.
+        let (key, session) = storage
+            .create_session(CreateSessionParams {
+                valid_till_ns: u64::MAX,
+                ..params(anchor_number, 1, 1_000)
+            })
+            .unwrap();
+        (storage, anchor_number, key, session.browser_id)
+    }
+
+    fn browser_last_used(storage: &Storage<VectorMemory>, anchor_number: AnchorNumber) -> u64 {
+        storage.read(anchor_number).unwrap().browsers()[0].last_used
+    }
+
+    #[test]
+    fn a_refresh_advances_the_device_registry() {
+        let (mut storage, anchor_number, key, _browser_id) = storage_with_registered_browser();
+
+        storage.record_session_use(&key, 9_000).unwrap();
+
+        assert_eq!(browser_last_used(&storage, anchor_number), 9_000);
+    }
+
+    #[test]
+    fn a_refresh_leaves_the_device_enrolment_timestamp_alone() {
+        let (mut storage, anchor_number, key, _browser_id) = storage_with_registered_browser();
+
+        storage.record_session_use(&key, 9_000).unwrap();
+
+        let device = storage.read(anchor_number).unwrap().browsers()[0].clone();
+        assert_eq!(device.created_at, 1_000);
+        assert_eq!(device.last_used, 9_000);
+    }
+
+    #[test]
+    fn a_refresh_for_a_device_the_anchor_never_registered_still_stamps_the_session() {
+        let (mut storage, anchor_number, key) = storage_with_session();
+
+        storage.record_session_use(&key, 9_000).unwrap();
+
+        assert_eq!(
+            session_of(&storage, anchor_number).last_refreshed_ns,
+            Some(9_000)
+        );
+    }
+}
+
 mod browser_session_count_tests {
     use super::TEST_NOW;
     use super::{params, params_at};