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

src/internet_identity/src/sessions.rs

@@ -324,7 +324,10 @@ pub fn app_prepare_delegation(
     request: AppPrepareDelegationRequest,
 ) -> Result<AppPrepareDelegationResponse, AppSessionError> {
     let now = time();
-    let (account, session) = authorize_session(now)?;
+    let (key, account, session) = authorize_session(now)?;
+
+    storage_borrow_mut(|storage| storage.record_session_use(&key, now))
+        .map_err(|err| AppSessionError::InternalCanisterError(err.to_string()))?;
 
     let expiration = u64::min(
         now.saturating_add(APP_DELEGATION_TTL_NS),
@@ -354,7 +357,7 @@ pub fn app_get_delegation(
     request: AppGetDelegationRequest,
 ) -> Result<SignedDelegation, AppSessionError> {
     let now = time();
-    let (account, session) = authorize_session(now)?;
+    let (_, account, session) = authorize_session(now)?;
 
     if request.expiration > now.saturating_add(APP_DELEGATION_TTL_NS)
         || request.expiration > session.valid_till_ns
@@ -396,7 +399,9 @@ pub fn app_get_delegation(
 /// The session index is keyed by the principal a session's chain is rooted at, so a hit is
 /// itself the proof that the caller is that session: nothing is named in the request and
 /// nothing is attached to it.
-fn authorize_session(now: Timestamp) -> Result<(Account, SessionRecord), AppSessionError> {
+fn authorize_session(
+    now: Timestamp,
+) -> Result<(SessionRecordKey, Account, SessionRecord), AppSessionError> {
     let key = storage_borrow(|storage| storage.lookup_session_with_principal(caller()))
         .ok_or(AppSessionError::NoMatchingSession)?;
 
@@ -411,7 +416,7 @@ fn authorize_session(now: Timestamp) -> Result<(Account, SessionRecord), AppSess
     if session.is_expired_or_idle(now) {
         return Err(AppSessionError::NoMatchingSession);
     }
-    Ok((account, session))
+    Ok((key, account, session))
 }
 
 fn account_seed(account: &Account) -> Result<Hash, AppSessionError> {

src/internet_identity/src/storage.rs

@@ -3339,6 +3339,72 @@ 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.
+    ///
+    /// `false` where the identity holds no such session, which is not a failure — a
+    /// session can be revoked between one call and the next.
+    pub fn record_session_use(
+        &mut self,
+        key: &SessionRecordKey,
+        now: Timestamp,
+    ) -> Result<bool, StorageError> {
+        let SessionRecordKey {
+            anchor_number,
+            origin,
+            account_number,
+            session_id,
+        } = key;
+        let (anchor_number, account_number, session_id) =
+            (*anchor_number, *account_number, *session_id);
+
+        if self.lookup_application_number_with_origin(origin).is_none() {
+            return Ok(false);
+        }
+        let mut anchor = self.read(anchor_number)?;
+        let (mut account_references, config) = self.account_state_for_origin(anchor_number, origin);
+
+        let Some(write) = account_references
+            .iter_mut()
+            .find(|write| write.account_reference.account_number == account_number)
+        else {
+            return Ok(false);
+        };
+        let Some(session) = write
+            .account_reference
+            .sessions
+            .iter_mut()
+            .find(|session| session.session_id == session_id)
+        else {
+            return Ok(false);
+        };
+
+        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(true)
+    }
+
     /// Retires an application no anchor references any more. The number is never
     /// reissued.
     fn remove_unreferenced_application(

src/internet_identity/src/storage/anchor.rs

@@ -766,6 +766,20 @@ impl Anchor {
         }
     }
 
+    /// Advances a device's `last_used`. Reports whether anything changed, so an unknown
+    /// device or a repeat inside one message costs no anchor write.
+    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;
+            }
+        }
+    }
+
     /// Resolves the browser a sign-in came from by the public key it proved possession of.
     ///
     /// An entry is reached only by the successor it announced. Presenting it promotes that

src/internet_identity/src/storage/tests.rs

@@ -6338,6 +6338,209 @@ mod session_consent_change_tests {
     }
 }
 
+mod session_refresh_stamp_tests {
+    use super::held_references;
+    use super::params;
+    use crate::storage::account::{AccountReference, SessionRecord, SessionRecordKey};
+    use crate::storage::CreateSessionParams;
+    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, SessionRecordKey) {
+        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) -> SessionRecord {
+        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();
+
+        let stamped = storage.record_session_use(&key, 2_000).unwrap();
+
+        assert!(stamped);
+        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] {
+            assert!(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);
+
+        assert!(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);
+    }
+
+    #[test]
+    fn a_stamp_for_a_session_that_is_gone_writes_nothing() {
+        let (mut storage, anchor_number, _key) = storage_with_session();
+
+        let wrote = storage
+            .record_session_use(
+                &SessionRecordKey {
+                    anchor_number,
+                    origin: ORIGIN.to_string(),
+                    account_number: None,
+                    session_id: 9_999,
+                },
+                5_000,
+            )
+            .unwrap();
+
+        assert!(!wrote);
+    }
+
+    #[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, SessionRecordKey, 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();
+
+        let stamped = storage.record_session_use(&key, 9_000).unwrap();
+
+        assert!(stamped);
+        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};