From afda5ca98f6af1c1523382ee4a8271910bcdbc3a Mon Sep 17 00:00:00 2001
From: Jon Chambers <jon@signal.org>
Date: Wed, 2 Nov 2022 17:08:19 -0400
Subject: [PATCH] Add a test for checking push challenge tokens

---
 .../controllers/AccountController.java        | 11 ++++++---
 .../controllers/AccountControllerTest.java    | 23 +++++++++++++++++++
 2 files changed, 31 insertions(+), 3 deletions(-)

diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/controllers/AccountController.java b/service/src/main/java/org/whispersystems/textsecuregcm/controllers/AccountController.java
index 1bb9eace8..fca43db92 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/controllers/AccountController.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/controllers/AccountController.java
@@ -787,14 +787,18 @@ public class AccountController {
     }
   }
 
-  private boolean pushChallengeMatches(
+  @VisibleForTesting
+  static boolean pushChallengeMatches(
       final String number,
       final Optional<String> pushChallenge,
       final Optional<StoredVerificationCode> storedVerificationCode) {
+
     final String countryCode = Util.getCountryCode(number);
     final String region = Util.getRegion(number);
-    Optional<String> storedPushChallenge = storedVerificationCode.map(StoredVerificationCode::pushCode);
-    boolean match = Optionals.zipWith(pushChallenge, storedPushChallenge, String::equals).orElse(false);
+    final Optional<String> storedPushChallenge = storedVerificationCode.map(StoredVerificationCode::pushCode);
+
+    final boolean match = Optionals.zipWith(pushChallenge, storedPushChallenge, String::equals).orElse(false);
+
     Metrics.counter(PUSH_CHALLENGE_COUNTER_NAME,
             COUNTRY_CODE_TAG_NAME, countryCode,
             REGION_TAG_NAME, region,
@@ -802,6 +806,7 @@ public class AccountController {
             CHALLENGE_PRESENT_TAG_NAME, Boolean.toString(pushChallenge.isPresent()),
             CHALLENGE_MATCH_TAG_NAME, Boolean.toString(match))
         .increment();
+
     return match;
   }
 
diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/controllers/AccountControllerTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/controllers/AccountControllerTest.java
index adaa31172..84aa4fb34 100644
--- a/service/src/test/java/org/whispersystems/textsecuregcm/controllers/AccountControllerTest.java
+++ b/service/src/test/java/org/whispersystems/textsecuregcm/controllers/AccountControllerTest.java
@@ -6,6 +6,7 @@
 package org.whispersystems.textsecuregcm.controllers;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyList;
 import static org.mockito.ArgumentMatchers.anyString;
@@ -48,6 +49,7 @@ import java.util.UUID;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.TimeUnit;
 import java.util.stream.Stream;
+import javax.annotation.Nullable;
 import javax.ws.rs.client.Entity;
 import javax.ws.rs.core.MediaType;
 import javax.ws.rs.core.Response;
@@ -1940,4 +1942,25 @@ class AccountControllerTest {
     assertThat(response.getStatus()).isEqualTo(413);
     assertThat(response.getHeaderString("Retry-After")).isEqualTo(String.valueOf(Duration.ofSeconds(13).toSeconds()));
   }
+
+  @ParameterizedTest
+  @MethodSource
+  void pushTokensMatch(@Nullable final String pushChallenge, @Nullable final StoredVerificationCode storedVerificationCode, final boolean expectMatch) {
+    final String number = "+18005550123";
+    final Optional<String> maybePushChallenge = Optional.ofNullable(pushChallenge);
+    final Optional<StoredVerificationCode> maybeStoredVerificationCode = Optional.ofNullable(storedVerificationCode);
+
+    assertEquals(expectMatch, AccountController.pushChallengeMatches(number, maybePushChallenge, maybeStoredVerificationCode));
+  }
+
+  private static Stream<Arguments> pushTokensMatch() {
+    return Stream.of(
+        Arguments.of(null, null, false),
+        Arguments.of("123456", null, false),
+        Arguments.of(null, new StoredVerificationCode(null, 0, null, null, null), false),
+        Arguments.of(null, new StoredVerificationCode(null, 0, "123456", null, null), false),
+        Arguments.of("654321", new StoredVerificationCode(null, 0, "123456", null, null), false),
+        Arguments.of("123456", new StoredVerificationCode(null, 0, "123456", null, null), true)
+    );
+  }
 }