From 854ab353b39f137d6a2602abe812a6944e71f32d Mon Sep 17 00:00:00 2001
From: adel-signal <adel@signal.org>
Date: Thu, 2 May 2024 12:34:34 -0700
Subject: [PATCH] calling: update TurnCallRouter to shuffle instance IPs to
 prevent allocation skew

Co-authored-by: Jonathan Klabunde Tomer <125505367+jkt-signal@users.noreply.github.com>
---
 .../textsecuregcm/WhisperServerService.java        |  3 ++-
 .../calls/routing/TurnCallRouter.java              | 14 +++++++++++---
 .../whispersystems/textsecuregcm/util/Util.java    |  9 +++------
 .../calls/routing/TurnCallRouterTest.java          |  4 +++-
 4 files changed, 19 insertions(+), 11 deletions(-)

diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
index 1ffd0058d..427c3dbd0 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
@@ -747,7 +747,8 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
         callRoutingTableManager,
         manualCallRoutingTableManager,
         configTurnRouter,
-        geoIpCityDatabaseManager
+        geoIpCityDatabaseManager,
+        false
     );
 
     final ClientConnectionManager clientConnectionManager = new ClientConnectionManager();
diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/calls/routing/TurnCallRouter.java b/service/src/main/java/org/whispersystems/textsecuregcm/calls/routing/TurnCallRouter.java
index 7d6a4c1ad..23af12dd4 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/calls/routing/TurnCallRouter.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/calls/routing/TurnCallRouter.java
@@ -32,19 +32,23 @@ public class TurnCallRouter {
   private final Supplier<CallRoutingTable> manualRouting;
   private final DynamicConfigTurnRouter configTurnRouter;
   private final Supplier<DatabaseReader> geoIp;
+  // controls whether instance IPs are shuffled. using if & boolean is ~5x faster than a function pointer
+  private final boolean stableSelect;
 
   public TurnCallRouter(
       @Nonnull Supplier<CallDnsRecords> callDnsRecords,
       @Nonnull Supplier<CallRoutingTable> performanceRouting,
       @Nonnull Supplier<CallRoutingTable> manualRouting,
       @Nonnull DynamicConfigTurnRouter configTurnRouter,
-      @Nonnull Supplier<DatabaseReader> geoIp
+      @Nonnull Supplier<DatabaseReader> geoIp,
+      boolean stableSelect
   ) {
     this.performanceRouting = performanceRouting;
     this.callDnsRecords = callDnsRecords;
     this.manualRouting = manualRouting;
     this.configTurnRouter = configTurnRouter;
     this.geoIp = geoIp;
+    this.stableSelect = stableSelect;
   }
 
   /**
@@ -139,10 +143,10 @@ public class TurnCallRouter {
 
     CallDnsRecords dnsRecords = this.callDnsRecords.get();
     List<InetAddress> ipv4Selection = datacenters.stream()
-        .flatMap(dc -> Util.randomNOfStable(dnsRecords.aByRegion().get(dc), limit).stream())
+        .flatMap(dc -> randomNOf(dnsRecords.aByRegion().get(dc), limit, stableSelect).stream())
         .toList();
     List<InetAddress> ipv6Selection = datacenters.stream()
-        .flatMap(dc -> Util.randomNOfStable(dnsRecords.aaaaByRegion().get(dc), limit).stream())
+        .flatMap(dc -> randomNOf(dnsRecords.aaaaByRegion().get(dc), limit, stableSelect).stream())
         .toList();
 
     // increase numV4 if not enough v6 options. vice-versa is also true
@@ -157,6 +161,10 @@ public class TurnCallRouter {
     ).toList();
   }
 
+  private static <E> List<E> randomNOf(List<E> values, int n, boolean stableSelect) {
+    return stableSelect ? Util.randomNOfStable(values, n) : Util.randomNOfShuffled(values, n);
+  }
+
   private static List<String> getUrlsForInstances(List<String> instanceIps) {
     return instanceIps.stream().flatMap(ip -> Stream.of(
             String.format("turn:%s", ip),
diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/util/Util.java b/service/src/main/java/org/whispersystems/textsecuregcm/util/Util.java
index 72d01d117..60301a1e9 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/util/Util.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/util/Util.java
@@ -169,21 +169,18 @@ public class Util {
   }
 
   /**
-   * Chooses min(values.size(), n) random values.
+   * Chooses min(values.size(), n) random values in shuffled order.
    * <br>
    * Copies the input Array - use for small lists only or for when n/values.size() is near 1.
    */
-  public static <E> List<E> randomNOf(List<E> values, int n) {
+  public static <E> List<E> randomNOfShuffled(List<E> values, int n) {
     if(values == null || values.isEmpty()) {
       return Collections.emptyList();
     }
 
     List<E> result = new ArrayList<>(values);
-    if(n >= values.size()) {
-      return result;
-    }
-
     Collections.shuffle(result);
+
     return result.stream().limit(n).toList();
   }
 
diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/calls/routing/TurnCallRouterTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/calls/routing/TurnCallRouterTest.java
index a338d24a1..458d0e11b 100644
--- a/service/src/test/java/org/whispersystems/textsecuregcm/calls/routing/TurnCallRouterTest.java
+++ b/service/src/test/java/org/whispersystems/textsecuregcm/calls/routing/TurnCallRouterTest.java
@@ -108,7 +108,9 @@ public class TurnCallRouterTest {
         () -> performanceTable,
         () -> manualTable,
          configTurnRouter,
-        () -> geoIp
+        () -> geoIp,
+        // set to true so the return values are predictable
+        true
     );
   }