From c10d9603ad332f76f482e8088e57b00b9a683177 Mon Sep 17 00:00:00 2001
From: Chris Eager <chris@signal.org>
Date: Mon, 29 Mar 2021 13:49:42 -0500
Subject: [PATCH] Add metric for international PreKey requests

---
 .../controllers/KeysController.java           | 50 ++++++++++++-------
 1 file changed, 33 insertions(+), 17 deletions(-)

diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/controllers/KeysController.java b/service/src/main/java/org/whispersystems/textsecuregcm/controllers/KeysController.java
index 76f145f97..06da186bb 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/controllers/KeysController.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/controllers/KeysController.java
@@ -4,8 +4,27 @@
  */
 package org.whispersystems.textsecuregcm.controllers;
 
+import static com.codahale.metrics.MetricRegistry.name;
+
 import com.codahale.metrics.annotation.Timed;
 import io.dropwizard.auth.Auth;
+import java.util.Collections;
+import java.util.LinkedList;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import javax.validation.Valid;
+import javax.ws.rs.Consumes;
+import javax.ws.rs.GET;
+import javax.ws.rs.HeaderParam;
+import javax.ws.rs.PUT;
+import javax.ws.rs.Path;
+import javax.ws.rs.PathParam;
+import javax.ws.rs.Produces;
+import javax.ws.rs.WebApplicationException;
+import javax.ws.rs.core.MediaType;
+import javax.ws.rs.core.Response;
+import io.micrometer.core.instrument.Metrics;
 import org.whispersystems.textsecuregcm.auth.AmbiguousIdentifier;
 import org.whispersystems.textsecuregcm.auth.Anonymous;
 import org.whispersystems.textsecuregcm.auth.DisabledPermittedAccount;
@@ -22,23 +41,7 @@ import org.whispersystems.textsecuregcm.storage.Account;
 import org.whispersystems.textsecuregcm.storage.AccountsManager;
 import org.whispersystems.textsecuregcm.storage.Device;
 import org.whispersystems.textsecuregcm.storage.KeysDynamoDb;
-
-import javax.validation.Valid;
-import javax.ws.rs.Consumes;
-import javax.ws.rs.GET;
-import javax.ws.rs.HeaderParam;
-import javax.ws.rs.PUT;
-import javax.ws.rs.Path;
-import javax.ws.rs.PathParam;
-import javax.ws.rs.Produces;
-import javax.ws.rs.WebApplicationException;
-import javax.ws.rs.core.MediaType;
-import javax.ws.rs.core.Response;
-import java.util.Collections;
-import java.util.LinkedList;
-import java.util.List;
-import java.util.Map;
-import java.util.Optional;
+import org.whispersystems.textsecuregcm.util.Util;
 
 @SuppressWarnings("OptionalUsedAsFieldOrParameterType")
 @Path("/v2/keys")
@@ -49,6 +52,11 @@ public class KeysController {
   private final AccountsManager             accounts;
   private final DirectoryQueue              directoryQueue;
 
+  private static final String INTERNATIONAL_PREKEY_REQUEST_COUNTER_NAME =
+      name(KeysController.class, "internationalPreKeyGet");
+
+  private static final String SOURCE_COUNTRY_TAG_NAME = "sourceCountry";
+
   public KeysController(RateLimiters rateLimiters, KeysDynamoDb keysDynamoDb, AccountsManager accounts, DirectoryQueue directoryQueue) {
     this.rateLimiters                = rateLimiters;
     this.keysDynamoDb                = keysDynamoDb;
@@ -119,6 +127,14 @@ public class KeysController {
 
     if (account.isPresent()) {
       rateLimiters.getPreKeysLimiter().validate(account.get().getNumber() + "." + account.get().getAuthenticatedDevice().get().getId() +  "__" + target.get().getNumber() + "." + deviceId);
+
+      final String accountCountryCode = Util.getCountryCode(account.get().getNumber());
+      final String targetCountryCode = Util.getCountryCode(target.get().getNumber());
+
+      if (!accountCountryCode.equals(targetCountryCode)) {
+        Metrics.counter(INTERNATIONAL_PREKEY_REQUEST_COUNTER_NAME, SOURCE_COUNTRY_TAG_NAME, accountCountryCode)
+            .increment();
+      }
     }
 
     Map<Long, PreKey>        preKeysByDeviceId = getLocalKeys(target.get(), deviceId);