From 62b35e0fb484c05d95d244fbec385da66ada5e98 Mon Sep 17 00:00:00 2001
From: Frederik Rothenberger <frederik.rothenberger@dfinity.org>
Date: Mon, 22 May 2023 13:13:04 +0200
Subject: [PATCH] Refactor authentication check to no longer trap (#1623)

This PR changes the authentication check to return a `Result`
instead of trapping. This refactoring is a prerequisite to
introduce authentication failure as a proper error in the v2 API.
---
 src/internet_identity/src/main.rs | 24 ++++++++++++++++--------
 1 file changed, 16 insertions(+), 8 deletions(-)

diff --git a/src/internet_identity/src/main.rs b/src/internet_identity/src/main.rs
index fc4ddd9065..8f938ff016 100644
--- a/src/internet_identity/src/main.rs
+++ b/src/internet_identity/src/main.rs
@@ -202,7 +202,9 @@ fn get_anchor_info(anchor_number: AnchorNumber) -> IdentityAnchorInfo {
 #[query]
 #[candid_method(query)]
 fn get_principal(anchor_number: AnchorNumber, frontend: FrontendHostname) -> Principal {
-    trap_if_not_authenticated(anchor_number);
+    let Ok(_) = check_authentication(anchor_number) else {
+        trap(&format!("{} could not be authenticated.", caller()));
+    };
     delegation::get_principal(anchor_number, frontend)
 }
 
@@ -233,7 +235,9 @@ fn get_delegation(
     session_key: SessionKey,
     expiration: Timestamp,
 ) -> GetDelegationResponse {
-    trap_if_not_authenticated(anchor_number);
+    let Ok(_) = check_authentication(anchor_number) else {
+        trap(&format!("{} could not be authenticated.", caller()));
+    };
     delegation::get_delegation(anchor_number, frontend, session_key, expiration)
 }
 
@@ -426,7 +430,9 @@ fn update_root_hash() {
 /// Note: this function reads / writes the anchor from / to stable memory. It is intended to be used by functions that
 /// do not further modify the anchor.
 fn authenticate_and_record_activity(anchor_number: AnchorNumber) -> Option<IIDomain> {
-    let (mut anchor, device_key) = trap_if_not_authenticated(anchor_number);
+    let Ok((mut anchor, device_key)) = check_authentication(anchor_number) else {
+        trap(&format!("{} could not be authenticated.", caller()));
+    };
     let domain = anchor.device(&device_key).unwrap().ii_domain();
     anchor_management::activity_bookkeeping(&mut anchor, &device_key);
     state::storage_borrow_mut(|storage| storage.write(anchor_number, anchor)).unwrap_or_else(
@@ -449,7 +455,9 @@ fn authenticated_anchor_operation<R>(
     anchor_number: AnchorNumber,
     op: impl FnOnce(&mut Anchor) -> Result<(R, Operation), R>,
 ) -> R {
-    let (mut anchor, device_key) = trap_if_not_authenticated(anchor_number);
+    let Ok((mut anchor, device_key)) = check_authentication(anchor_number) else {
+        trap(&format!("{} could not be authenticated.", caller()));
+    };
     anchor_management::activity_bookkeeping(&mut anchor, &device_key);
 
     let result = op(&mut anchor);
@@ -469,8 +477,8 @@ fn authenticated_anchor_operation<R>(
 }
 
 /// Checks if the caller is authenticated against the anchor provided and returns a reference to the device used.
-/// Traps if the caller is not authenticated.
-fn trap_if_not_authenticated(anchor_number: AnchorNumber) -> (Anchor, DeviceKey) {
+/// Returns an error if the caller cannot be authenticated.
+fn check_authentication(anchor_number: AnchorNumber) -> Result<(Anchor, DeviceKey), ()> {
     let anchor = state::anchor(anchor_number);
     let caller = caller();
 
@@ -482,10 +490,10 @@ fn trap_if_not_authenticated(anchor_number: AnchorNumber) -> (Anchor, DeviceKey)
                     .is_ok()
             })
         {
-            return (anchor.clone(), device.pubkey.clone());
+            return Ok((anchor.clone(), device.pubkey.clone()));
         }
     }
-    trap(&format!("{} could not be authenticated.", caller))
+    Err(())
 }
 
 /// New v2 API that aims to eventually replace the current API.