feat: Switch the source of truth for the lookup_application_with_origin map (#3353)

* first

* second

* Refer to the new map for generating new app numbers

Author: aterga

src/internet_identity/src/storage.rs

@@ -463,9 +463,8 @@ impl<M: Memory + Clone> Storage<M> {
                 lookup_anchor_with_device_credential_memory,
             ),
 
-            // TODO[ID-352]: Change this to use `lookup_application_with_origin_memory`.
             lookup_application_with_origin_memory_wrapper: MemoryWrapper::new(
-                lookup_application_with_origin_memory_old.clone(),
+                lookup_application_with_origin_memory.clone(),
             ),
             // TODO[ID-354]: Remove this after the production data is migrated to the new map.
             lookup_application_with_origin_memory_old: StableBTreeMap::init(
@@ -737,22 +736,20 @@ impl<M: Memory + Clone> Storage<M> {
         let origin_hash = StorableOriginHash::from_origin(origin);
         let origin_sha256 = StorableOriginSha256::from_origin(origin);
 
-        // TODO[ID_352]: Switch to `lookup_application_with_origin_memory`.
         if let Some(existing_number) = self
-            .lookup_application_with_origin_memory_old
-            .get(&origin_hash)
+            .lookup_application_with_origin_memory
+            .get(&origin_sha256)
         {
             existing_number
         } else {
-            let new_number: ApplicationNumber =
-                self.lookup_application_with_origin_memory_old.len();
+            let new_number: ApplicationNumber = self.lookup_application_with_origin_memory.len();
 
-            // Update the source of truth.
+            // Update the old map (will eventually be removed).
             // TODO[ID_353]: Remove this line.
             self.lookup_application_with_origin_memory_old
                 .insert(origin_hash, new_number);
 
-            // Update the new map (will eventually become the source of truth).
+            // Update the source of truth.
             self.lookup_application_with_origin_memory
                 .insert(origin_sha256, new_number);
 
@@ -772,9 +769,8 @@ impl<M: Memory + Clone> Storage<M> {
         &self,
         origin: &FrontendHostname,
     ) -> Option<ApplicationNumber> {
-        // TODO[ID-352]: Start reading from the new map (`lookup_application_with_origin_memory`).
-        self.lookup_application_with_origin_memory_old
-            .get(&StorableOriginHash::from_origin(origin))
+        self.lookup_application_with_origin_memory
+            .get(&StorableOriginSha256::from_origin(origin))
     }
 
     /// Only used in tests.

src/internet_identity/src/storage/tests.rs

@@ -397,7 +397,7 @@ mod application_lookup_tests {
 
         assert_eq!(
             storage.lookup_application_number_with_origin(&origin),
-            expected_old,
+            expected_new,
             "Unexpected result from lookup_application_number_with_origin for origin {}",
             origin
         );
@@ -484,41 +484,47 @@ mod application_lookup_tests {
 
         // These origins have hash prefix collisions with 8-byte hash but not with full SHA-256.
         // Example taken from https://github.com/yugt/sha256-prefix-collision
-        let origin1 = "08RTz8".to_string();
-        let origin2 = "4iRDWF".to_string();
+        let origin0 = "08RTz8".to_string();
+        let origin1 = "4iRDWF".to_string();
 
         // Smoke test: these are indeed distinct origins.
-        assert_ne!(origin1, origin2);
+        assert_ne!(origin0, origin1);
         // However, the SHA256 prefix is the same.
         assert_eq!(
-            StorableOriginHash::from_origin(&origin1),
-            StorableOriginHash::from_origin(&origin2)
+            StorableOriginHash::from_origin(&origin0),
+            StorableOriginHash::from_origin(&origin1)
         );
         // But the entire SHA256 sum is different.
         assert_ne!(
-            StorableOriginSha256::from_origin(&origin1),
-            StorableOriginSha256::from_origin(&origin2)
+            StorableOriginSha256::from_origin(&origin0),
+            StorableOriginSha256::from_origin(&origin1)
         );
 
+        assert_application_lookup(&storage, &origin1, None, None);
+        assert_application_lookup(&storage, &origin0, None, None);
+
         // Test what happens if we insert both origins
-        let app_num1 = storage.lookup_or_insert_application_number_with_origin(&origin1);
-        let app_num2 = storage.lookup_or_insert_application_number_with_origin(&origin2);
+        let app_num0 = storage.lookup_or_insert_application_number_with_origin(&origin0);
 
-        // Should get the same application number, as we do not yet detect collisions (since the
-        // source of truth is still the old lookup map).
-        // TODO[ID-352]: This assertion should be changed to `assert_ne`.
-        assert_eq!(app_num1, app_num2);
-        assert_eq!(app_num1, 0);
+        assert_application_lookup(&storage, &origin1, None, Some(0));
+        assert_application_lookup(&storage, &origin0, Some(0), Some(0));
 
-        // Both should be stored correctly in new SHA-256 map
+        let app_num1 = storage.lookup_or_insert_application_number_with_origin(&origin1);
 
-        assert_application_lookup(&storage, &origin1, Some(0), Some(0));
-        assert_application_lookup(&storage, &origin2, None, Some(0));
+        // Should get the different application numbers, as we prevented the collision by using
+        // the full SHA-256, not just a 8-byte prefix thereof.
+        assert_eq!(app_num0, 0);
+        assert_eq!(app_num1, 1);
+
+        // Both should be stored correctly in new SHA-256 map, but the collision in the old map
+        // would cause the application to be overwritten (hence we expect app ID 1 for origin0).
+        assert_application_lookup(&storage, &origin1, Some(1), Some(1));
+        assert_application_lookup(&storage, &origin0, Some(0), Some(1));
 
         // Applications should be stored with correct origins
+        let stored_app0 = storage.stable_application_memory.get(&app_num0).unwrap();
         let stored_app1 = storage.stable_application_memory.get(&app_num1).unwrap();
-
-        // The first origin took place, not the second one
+        assert_eq!(stored_app0.origin, origin0);
         assert_eq!(stored_app1.origin, origin1);
     }