From b43596f7376100dbe00a7c960f65328ffcf2c0bf Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sosth=C3=A8ne=20Gu=C3=A9don?= <sosthene@nitrokey.com>
Date: Fri, 26 Jul 2024 11:01:46 +0200
Subject: [PATCH] Fix credential count and add back hard limit

---
 src/ctap2.rs                       | 18 ++++++-------
 src/ctap2/credential_management.rs | 41 +++++++++++++++++++++++++++++-
 2 files changed, 49 insertions(+), 10 deletions(-)

diff --git a/src/ctap2.rs b/src/ctap2.rs
index 89bec75..ee85736 100644
--- a/src/ctap2.rs
+++ b/src/ctap2.rs
@@ -1,5 +1,6 @@
 //! The `ctap_types::ctap2::Authenticator` implementation.
 
+use credential_management::CredentialManagement;
 use ctap_types::{
     ctap2::{
         self, client_pin::Permissions, AttestationFormatsPreference, AttestationStatement,
@@ -21,15 +22,9 @@ use trussed::{
 };
 
 use crate::{
-    constants,
+    constants::{self, MAX_RESIDENT_CREDENTIALS_GUESSTIMATE},
     credential::{self, Credential, FullCredential, Key, StrippedCredential},
-    format_hex,
-    state::{
-        self,
-        // // (2022-02-27): 9288 bytes
-        // MinCredentialHeap,
-    },
-    Result, SigningAlgorithm, TrussedRequirements, UserPresence,
+    format_hex, state, Result, SigningAlgorithm, TrussedRequirements, UserPresence,
 };
 
 #[allow(unused_imports)]
@@ -393,7 +388,12 @@ impl<UP: UserPresence, T: TrussedRequirements> Authenticator for crate::Authenti
             self.delete_resident_key_by_user_id(&rp_id_hash, &credential.user.id)
                 .ok();
 
-            let mut key_store_full = !self.can_fit(serialized_credential.len());
+            let mut key_store_full = !self.can_fit(serialized_credential.len())
+                || CredentialManagement::new(self).count_credentials()
+                    >= self
+                        .config
+                        .max_resident_credential_count
+                        .unwrap_or(MAX_RESIDENT_CREDENTIALS_GUESSTIMATE);
 
             if !key_store_full {
                 // then store key, making it resident
diff --git a/src/ctap2/credential_management.rs b/src/ctap2/credential_management.rs
index 08dccd9..da67813 100644
--- a/src/ctap2/credential_management.rs
+++ b/src/ctap2/credential_management.rs
@@ -65,13 +65,52 @@ where
         let mut response: Response = Default::default();
 
         let max_resident_credentials = self.estimate_remaining();
-        response.existing_resident_credentials_count = Some(0);
+        response.existing_resident_credentials_count = Some(self.count_credentials());
         response.max_possible_remaining_residential_credentials_count =
             Some(max_resident_credentials.try_into().unwrap_or(u32::MAX));
 
         response
     }
 
+    pub fn count_credentials(&mut self) -> u32 {
+        let dir = PathBuf::from(b"rk");
+        let maybe_first_rp =
+            syscall!(self
+                .trussed
+                .read_dir_first(Location::Internal, dir.clone(), None))
+            .entry;
+
+        let first_rp = match maybe_first_rp {
+            None => return 0,
+            Some(rp) => rp,
+        };
+
+        let (mut num_rks, _) = self.count_rp_rks(PathBuf::from(first_rp.path()));
+        let mut last_rp = PathBuf::from(first_rp.file_name());
+
+        loop {
+            syscall!(self
+                .trussed
+                .read_dir_first(Location::Internal, dir.clone(), Some(last_rp),))
+            .entry
+            .unwrap();
+            let maybe_next_rp = syscall!(self.trussed.read_dir_next()).entry;
+
+            match maybe_next_rp {
+                None => {
+                    return num_rks;
+                }
+                Some(rp) => {
+                    last_rp = PathBuf::from(rp.file_name());
+                    info!("counting..");
+                    let (this_rp_rk_count, _) = self.count_rp_rks(PathBuf::from(rp.path()));
+                    info!("{:?}", this_rp_rk_count);
+                    num_rks += this_rp_rk_count;
+                }
+            }
+        }
+    }
+
     pub fn first_relying_party(&mut self) -> Result<Response> {
         info!("first rp");