Only require captcha if threshold rate is exceeded (#2638)

* Only require captcha if threshold rate is exceeded

This PR enables the dynamic captcha feature: I.e. if the captcha is
configured to be dynamic _and_ the configured threshold has been exceeded,
then require a captcha to be solved. Otherwise skip the captcha entirely.

* Implement review feedback

src/canister_tests/src/framework.rs

@@ -209,6 +209,20 @@ pub fn arg_with_anchor_range(
     })
 }
 
+pub fn arg_with_dynamic_captcha() -> Option<InternetIdentityInit> {
+    Some(InternetIdentityInit {
+        captcha_config: Some(CaptchaConfig {
+            max_unsolved_captchas: 50,
+            captcha_trigger: CaptchaTrigger::Dynamic {
+                threshold_pct: 20,
+                current_rate_sampling_interval_s: 10,
+                reference_rate_sampling_interval_s: 100,
+            },
+        }),
+        ..InternetIdentityInit::default()
+    })
+}
+
 pub fn archive_wasm_hash(wasm: &Vec<u8>) -> [u8; 32] {
     let mut hasher = Sha256::new();
     hasher.update(wasm);

src/internet_identity/src/anchor_management/registration/registration_flow_v2.rs

@@ -1,7 +1,8 @@
 use crate::anchor_management::registration::captcha::{
-    check_captcha_solution, create_captcha, make_rng, Base64,
+    check_captcha_solution, create_captcha, make_rng,
 };
 use crate::anchor_management::registration::rate_limit::process_rate_limit;
+use crate::anchor_management::registration::Base64;
 use crate::anchor_management::{activity_bookkeeping, post_operation_bookkeeping};
 use crate::state;
 use crate::state::flow_states::RegistrationFlowState;
@@ -12,11 +13,24 @@ use ic_cdk::caller;
 use internet_identity_interface::archive::types::{DeviceDataWithoutAlias, Operation};
 use internet_identity_interface::internet_identity::types::IdRegFinishError::IdentityLimitReached;
 use internet_identity_interface::internet_identity::types::{
-    CheckCaptchaArg, CheckCaptchaError, DeviceData, DeviceWithUsage, IdRegFinishArg,
-    IdRegFinishError, IdRegFinishResult, IdRegNextStepResult, IdRegStartError, IdentityNumber,
-    RegistrationFlowNextStep,
+    CaptchaTrigger, CheckCaptchaArg, CheckCaptchaError, DeviceData, DeviceWithUsage,
+    IdRegFinishArg, IdRegFinishError, IdRegFinishResult, IdRegNextStepResult, IdRegStartError,
+    IdentityNumber, RegistrationFlowNextStep, StaticCaptchaTrigger,
 };
 
+impl RegistrationFlowState {
+    pub fn to_flow_next_step(&self) -> RegistrationFlowNextStep {
+        match self {
+            RegistrationFlowState::CheckCaptcha {
+                captcha_png_base64, ..
+            } => RegistrationFlowNextStep::CheckCaptcha {
+                captcha_png_base64: captcha_png_base64.0.clone(),
+            },
+            RegistrationFlowState::FinishRegistration { .. } => RegistrationFlowNextStep::Finish,
+        }
+    }
+}
+
 impl From<RegistrationFlowState> for RegistrationFlowNextStep {
     fn from(value: RegistrationFlowState) -> Self {
         match value {
@@ -38,17 +52,41 @@ pub async fn identity_registration_start() -> Result<IdRegNextStepResult, IdRegS
         return Err(IdRegStartError::InvalidCaller);
     }
 
-    let (captcha_png_base64, flow_state) = captcha_flow_state(time()).await;
+    let now = time();
+    let flow_state = match captcha_required() {
+        true => captcha_flow_state(now).await.1,
+        false => RegistrationFlowState::FinishRegistration {
+            flow_created_timestamp_ns: now,
+        },
+    };
+
+    let next_step_result = IdRegNextStepResult {
+        next_step: flow_state.to_flow_next_step(),
+    };
+
     state::with_flow_states_mut(|flow_states| {
         flow_states.new_registration_flow(caller, flow_state)
     })
     .map_err(|_| IdRegStartError::AlreadyInProgress)?;
 
-    Ok(IdRegNextStepResult {
-        next_step: RegistrationFlowNextStep::CheckCaptcha {
-            captcha_png_base64: captcha_png_base64.0,
+    Ok(next_step_result)
+}
+
+fn captcha_required() -> bool {
+    let captcha_config = state::persistent_state(|ps| ps.captcha_config.clone());
+    match captcha_config.captcha_trigger {
+        CaptchaTrigger::Static(static_trigger) => match static_trigger {
+            StaticCaptchaTrigger::CaptchaEnabled => true,
+            StaticCaptchaTrigger::CaptchaDisabled => false,
         },
-    })
+        CaptchaTrigger::Dynamic { .. } => {
+            state::storage_borrow(|storage| storage.registration_rates.registration_rates())
+                .map(|rates| rates.captcha_required())
+                // unreachable because it is only `None` on static captcha
+                // default to true nonetheless because requiring a captcha is better than panicking
+                .unwrap_or(true)
+        }
+    }
 }
 
 async fn captcha_flow_state(flow_created_timestamp_ns: u64) -> (Base64, RegistrationFlowState) {

src/internet_identity/src/storage/registration_rates.rs

@@ -23,6 +23,12 @@ pub struct NormalizedRegistrationRates {
     pub captcha_threshold_rate: f64,
 }
 
+impl NormalizedRegistrationRates {
+    pub fn captcha_required(&self) -> bool {
+        self.current_rate_per_second > self.captcha_threshold_rate
+    }
+}
+
 struct DynamicCaptchaConfig {
     reference_rate_retention_ns: u64,
     current_rate_retention_ns: u64,

src/internet_identity/tests/integration/v2_api/authn_method_test_helpers.rs

@@ -3,7 +3,8 @@ use canister_tests::framework::{test_principal, time};
 use ic_cdk::api::management_canister::main::CanisterId;
 use internet_identity_interface::internet_identity::types::{
     AuthnMethod, AuthnMethodData, AuthnMethodProtection, AuthnMethodPurpose,
-    AuthnMethodSecuritySettings, IdentityNumber, MetadataEntryV2, PublicKeyAuthn, WebAuthn,
+    AuthnMethodSecuritySettings, IdentityNumber, MetadataEntryV2, PublicKeyAuthn,
+    RegistrationFlowNextStep, WebAuthn,
 };
 use pocket_ic::PocketIc;
 use serde_bytes::ByteBuf;
@@ -59,13 +60,16 @@ pub fn create_identity_with_authn_method(
 ) -> IdentityNumber {
     // unique flow principal as the time changes every round
     let flow_principal = test_principal(time(env));
-    api_v2::identity_registration_start(env, canister_id, flow_principal)
+    let result = api_v2::identity_registration_start(env, canister_id, flow_principal)
         .expect("API call failed")
         .expect("registration start failed");
 
-    api_v2::check_captcha(env, canister_id, flow_principal, "a".to_string())
-        .expect("API call failed")
-        .expect("check_captcha failed");
+    // supply captcha only if required
+    if let RegistrationFlowNextStep::CheckCaptcha { .. } = result.next_step {
+        api_v2::check_captcha(env, canister_id, flow_principal, "a".to_string())
+            .expect("API call failed")
+            .expect("check_captcha failed");
+    }
 
     api_v2::identity_registration_finish(env, canister_id, flow_principal, authn_method)
         .expect("API call failed")

src/internet_identity/tests/integration/v2_api/identity_register.rs

@@ -1,3 +1,5 @@
+mod dynamic_captcha;
+
 use crate::v2_api::authn_method_test_helpers::{
     create_identity_with_authn_method, test_authn_method,
 };
@@ -9,8 +11,9 @@ use canister_tests::framework::{
 };
 use internet_identity_interface::internet_identity::types::IdentityInfoError::Unauthorized;
 use internet_identity_interface::internet_identity::types::{
-    CheckCaptchaError, IdRegFinishError, IdRegStartError, MetadataEntryV2, RateLimitConfig,
-    RegistrationFlowNextStep,
+    CaptchaConfig, CaptchaTrigger, CheckCaptchaError, IdRegFinishError, IdRegStartError,
+    InternetIdentityInit, MetadataEntryV2, RateLimitConfig, RegistrationFlowNextStep,
+    StaticCaptchaTrigger,
 };
 use pocket_ic::CallError;
 use serde_bytes::ByteBuf;
@@ -271,3 +274,25 @@ fn should_trigger_rate_limit_on_too_many_flows() -> Result<(), CallError> {
     );
     Ok(())
 }
+
+#[test]
+fn should_not_require_captcha_when_disabled() {
+    let env = env();
+    let canister_id = install_ii_canister_with_arg(
+        &env,
+        II_WASM.clone(),
+        Some(InternetIdentityInit {
+            captcha_config: Some(CaptchaConfig {
+                max_unsolved_captchas: 50,
+                captcha_trigger: CaptchaTrigger::Static(StaticCaptchaTrigger::CaptchaDisabled),
+            }),
+            ..InternetIdentityInit::default()
+        }),
+    );
+
+    let result = api_v2::identity_registration_start(&env, canister_id, test_principal(0))
+        .expect("API call failed")
+        .expect("registration start failed");
+
+    assert!(matches!(result.next_step, RegistrationFlowNextStep::Finish));
+}

src/internet_identity/tests/integration/v2_api/identity_register/dynamic_captcha.rs

@@ -0,0 +1,68 @@
+use crate::v2_api::authn_method_test_helpers::{
+    create_identity_with_authn_method, test_authn_method,
+};
+use canister_tests::api::internet_identity::api_v2;
+use canister_tests::framework::{
+    arg_with_dynamic_captcha, env, install_ii_canister_with_arg, test_principal, II_WASM,
+};
+use internet_identity_interface::internet_identity::types::RegistrationFlowNextStep;
+use std::time::Duration;
+
+#[test]
+fn should_not_require_captcha_below_threshold_rate() {
+    let env = env();
+    let canister_id =
+        install_ii_canister_with_arg(&env, II_WASM.clone(), arg_with_dynamic_captcha());
+    let authn_method = test_authn_method();
+
+    let flow_principal = test_principal(0);
+    let result = api_v2::identity_registration_start(&env, canister_id, flow_principal)
+        .expect("API call failed")
+        .expect("registration start failed");
+
+    assert!(matches!(result.next_step, RegistrationFlowNextStep::Finish));
+
+    api_v2::identity_registration_finish(&env, canister_id, flow_principal, &authn_method)
+        .expect("API call failed")
+        .expect("registration finish failed");
+}
+
+#[test]
+fn should_require_captcha_above_threshold_rate() {
+    let env = env();
+    let canister_id =
+        install_ii_canister_with_arg(&env, II_WASM.clone(), arg_with_dynamic_captcha());
+    let authn_method = test_authn_method();
+
+    // initialize a base rate of one registration every 2 seconds
+    for _ in 0..10 {
+        create_identity_with_authn_method(&env, canister_id, &authn_method);
+        env.advance_time(Duration::from_secs(2))
+    }
+
+    // Double the rate of registrations to one per second
+    // The 20% threshold rate should allow 5 registrations before the captcha kicks in
+    for i in 0..5 {
+        let flow_principal = test_principal(i);
+        let result = api_v2::identity_registration_start(&env, canister_id, flow_principal)
+            .expect("API call failed")
+            .expect("registration start failed");
+
+        assert!(matches!(result.next_step, RegistrationFlowNextStep::Finish));
+
+        api_v2::identity_registration_finish(&env, canister_id, flow_principal, &authn_method)
+            .expect("API call failed")
+            .expect("registration finish failed");
+        env.advance_time(Duration::from_secs(1));
+    }
+
+    let result = api_v2::identity_registration_start(&env, canister_id, test_principal(99))
+        .expect("API call failed")
+        .expect("registration start failed");
+
+    // captcha kicks in
+    assert!(matches!(
+        result.next_step,
+        RegistrationFlowNextStep::CheckCaptcha { .. }
+    ));
+}