From e1f230fb66574c91b07e3a0cf441d357fb2f1d3d Mon Sep 17 00:00:00 2001
From: Nicolas Mattia <nicolas@dfinity.org>
Date: Fri, 26 Nov 2021 16:50:06 +0100
Subject: [PATCH] Return 'BadChallenge' on bad CAPTCHA

---
 .../generated/internet_identity_idl.js        |  1 +
 .../generated/internet_identity_types.d.ts    |  3 ++-
 src/frontend/src/flows/confirmRegister.ts     | 24 +++++++++++++++----
 src/frontend/src/flows/loginUnknown.ts        |  8 +++++++
 src/frontend/src/utils/iiConnection.ts        |  7 +++++-
 src/internet_identity/internet_identity.did   |  2 ++
 src/internet_identity/src/main.rs             | 18 +++++++-------
 7 files changed, 49 insertions(+), 14 deletions(-)

diff --git a/src/frontend/generated/internet_identity_idl.js b/src/frontend/generated/internet_identity_idl.js
index e8d0bf5ff3..bed3fb426a 100644
--- a/src/frontend/generated/internet_identity_idl.js
+++ b/src/frontend/generated/internet_identity_idl.js
@@ -82,6 +82,7 @@ export const idlFactory = ({ IDL }) => {
     'chars' : IDL.Text,
   });
   const RegisterResponse = IDL.Variant({
+    'bad_challenge' : IDL.Null,
     'canister_full' : IDL.Null,
     'registered' : IDL.Record({ 'user_number' : UserNumber }),
   });
diff --git a/src/frontend/generated/internet_identity_types.d.ts b/src/frontend/generated/internet_identity_types.d.ts
index c6cbdcf783..5f34b793a3 100644
--- a/src/frontend/generated/internet_identity_types.d.ts
+++ b/src/frontend/generated/internet_identity_types.d.ts
@@ -50,7 +50,8 @@ export interface ProofOfWork { 'nonce' : bigint, 'timestamp' : Timestamp }
 export type PublicKey = Array<number>;
 export type Purpose = { 'authentication' : null } |
   { 'recovery' : null };
-export type RegisterResponse = { 'canister_full' : null } |
+export type RegisterResponse = { 'bad_challenge' : null } |
+  { 'canister_full' : null } |
   { 'registered' : { 'user_number' : UserNumber } };
 export type SessionKey = PublicKey;
 export interface SignedDelegation {
diff --git a/src/frontend/src/flows/confirmRegister.ts b/src/frontend/src/flows/confirmRegister.ts
index 860582cd91..de018adbaa 100644
--- a/src/frontend/src/flows/confirmRegister.ts
+++ b/src/frontend/src/flows/confirmRegister.ts
@@ -21,7 +21,7 @@ const pageContent = html`
       <p class="captcha-status-text">…</p>
       <img id="captchaImg" />
       <input id="captchaInput" />
-      <p>Please confirm to add your device.</p>
+      <p class="confirm-paragraph">Please confirm to add your device.</p>
       <button type="submit" class="primary" id="confirmRegisterButton" disabled>
         Confirm
       </button>
@@ -58,6 +58,13 @@ const tryRegister = (
       displayUserNumber(result.userNumber).then(() => {
         func(apiResultToLoginResult(result));
       });
+    } else if (result.kind == "badChallenge") {
+      const confirmParagraph = document.querySelector(
+        ".confirm-paragraph"
+      ) as HTMLElement;
+      confirmParagraph.innerHTML =
+        "The value you entered is incorrect. A new challenge is generated.";
+      requestCaptcha();
     } else {
       // Currently if something goes wrong we only tell the user that
       // something went wrong and then reload the page.
@@ -84,6 +91,16 @@ const requestCaptcha = (): void => {
   ) as HTMLElement;
   captchaStatusText.innerHTML = "Creating CAPTCHA challenge…";
 
+  const captchaInput = document.querySelector(
+    "#captchaInput"
+  ) as HTMLFormElement;
+  captchaInput.disabled = true;
+
+  const confirmRegisterButton = form.querySelector(
+    "#confirmRegisterButton"
+  ) as HTMLFormElement;
+  confirmRegisterButton.disabled = true;
+
   // Wrap this in a promise to avoid slowing things down
   const makePow: Promise<ProofOfWork> = new Promise((resolve) => {
     const now_in_ns = BigInt(Date.now()) * BigInt(1000000);
@@ -100,15 +117,14 @@ const requestCaptcha = (): void => {
           "src",
           `data:image/png;base64, ${captchaResp.png_base64}`
         );
-        const confirmRegisterButton = form.querySelector(
-          "#confirmRegisterButton"
-        ) as HTMLFormElement;
         confirmRegisterButton.setAttribute(
           "data-captcha-key",
           `${captchaResp.challenge_key}`
         );
         captchaStatusText.innerHTML = "Please type in the characters you see.";
         confirmRegisterButton.disabled = false;
+        captchaInput.disabled = false;
+        captchaInput.value = "";
       }
     });
 };
diff --git a/src/frontend/src/flows/loginUnknown.ts b/src/frontend/src/flows/loginUnknown.ts
index 89731459b2..1f3e15fd1c 100644
--- a/src/frontend/src/flows/loginUnknown.ts
+++ b/src/frontend/src/flows/loginUnknown.ts
@@ -217,6 +217,14 @@ export const apiResultToLoginResult = (result: ApiResult): LoginResult => {
           "Failed to register with Internet Identity, because there is no space left at the moment. We're working on increasing the capacity.",
       };
     }
+    case "badChallenge": {
+      return {
+        tag: "err",
+        title: "Failed to register",
+        message:
+          "Failed to register with Internet Identity, because the challenge wasn't successful",
+      };
+    }
     case "seedPhraseFail": {
       return {
         tag: "err",
diff --git a/src/frontend/src/utils/iiConnection.ts b/src/frontend/src/utils/iiConnection.ts
index 8162db1c26..b9f0f36591 100644
--- a/src/frontend/src/utils/iiConnection.ts
+++ b/src/frontend/src/utils/iiConnection.ts
@@ -57,13 +57,16 @@ export type RegisterResult =
   | LoginSuccess
   | AuthFail
   | ApiError
-  | RegisterNoSpace;
+  | RegisterNoSpace
+  | BadChallenge;
 
 type LoginSuccess = {
   kind: "loginSuccess";
   connection: IIConnection;
   userNumber: bigint;
 };
+
+type BadChallenge = { kind: "badChallenge" };
 type UnknownUser = { kind: "unknownUser"; userNumber: bigint };
 type AuthFail = { kind: "authFail"; error: Error };
 type ApiError = { kind: "apiError"; error: Error };
@@ -121,6 +124,8 @@ export class IIConnection {
         connection: new IIConnection(identity, delegationIdentity, actor),
         userNumber,
       };
+    } else if (hasOwnProperty(registerResponse, "bad_challenge")) {
+      return { kind: "badChallenge" };
     } else {
       console.error("unexpected register response", registerResponse);
       throw Error("unexpected register response");
diff --git a/src/internet_identity/internet_identity.did b/src/internet_identity/internet_identity.did
index fba7609a68..eae2cc0c04 100644
--- a/src/internet_identity/internet_identity.did
+++ b/src/internet_identity/internet_identity.did
@@ -67,6 +67,8 @@ type RegisterResponse = variant {
   registered: record { user_number: UserNumber; };
   // No more registrations are possible in this instance of the II service canister.
   canister_full;
+  // The challenge was not successful.
+  bad_challenge;
 };
 
 type Delegation = record {
diff --git a/src/internet_identity/src/main.rs b/src/internet_identity/src/main.rs
index 1868c304b5..662d8dc5d5 100644
--- a/src/internet_identity/src/main.rs
+++ b/src/internet_identity/src/main.rs
@@ -155,7 +155,8 @@ enum RegisterResponse {
     Registered { user_number: UserNumber },
     #[serde(rename = "canister_full")]
     CanisterFull,
-    // TODO: add BadCaptcha here
+    #[serde(rename = "bad_challenge")]
+    BadChallenge,
 }
 
 mod hash;
@@ -289,7 +290,10 @@ async fn init_salt() {
 
 #[update]
 async fn register(device_data: DeviceData, challenge_result: ChallengeAttempt) -> RegisterResponse {
-    check_challenge(challenge_result);
+    if let Err(()) = check_challenge(challenge_result) {
+        return RegisterResponse::BadChallenge;
+    }
+
     check_entry_limits(&device_data);
 
     if caller() != Principal::self_authenticating(device_data.pubkey.clone()) {
@@ -533,20 +537,18 @@ fn create_captcha<T: RngCore>(rng: T) -> (Base64, String) {
 }
 
 // just traps if challenge isn't OK, because when in rome...
-fn check_challenge(res: ChallengeAttempt) {
+fn check_challenge(res: ChallengeAttempt) -> Result<(),()> {
 
     STATE.with(|s| {
         let mut inflight_challenges = s.inflight_challenges.borrow_mut();
         match inflight_challenges.remove(&res.key) {
             Some(challenge) => {
                 if res.chars !=  challenge.chars {
-                    // NOTE: we _could_ show the expected chars here (the key has been
-                    // removed already so the user won't be able to re-submit the answer using that
-                    // key).
-                    trap("CAPTCHA challenge failed");
+                    return Err(());
                 }
+                return Ok(());
             },
-            None =>  trap("Could not find a CAPTCHA challenge with that key") ,
+            None =>  Err(()),
         }
     })
 }