From f8c623074be8a736beaa6b5934ec1abfd08aaaf1 Mon Sep 17 00:00:00 2001
From: Jon Chambers <jon@signal.org>
Date: Mon, 17 May 2021 18:40:07 -0400
Subject: [PATCH] Introduce an ASN-to-IP manager.

---
 .../WhisperServerConfiguration.java           | 17 +++-
 .../textsecuregcm/WhisperServerService.java   |  5 +-
 ...va => MonitoredS3ObjectConfiguration.java} |  2 +-
 .../textsecuregcm/util/AsnManager.java        | 98 +++++++++++++++++++
 .../textsecuregcm/util/AsnTable.java          |  6 ++
 .../util/TorExitNodeManager.java              |  4 +-
 .../textsecuregcm/util/AsnManagerTest.java    | 45 +++++++++
 .../util/TorExitNodeManagerTest.java          |  4 +-
 8 files changed, 171 insertions(+), 10 deletions(-)
 rename service/src/main/java/org/whispersystems/textsecuregcm/configuration/{TorExitNodeConfiguration.java => MonitoredS3ObjectConfiguration.java} (96%)
 create mode 100644 service/src/main/java/org/whispersystems/textsecuregcm/util/AsnManager.java
 create mode 100644 service/src/test/java/org/whispersystems/textsecuregcm/util/AsnManagerTest.java

diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerConfiguration.java b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerConfiguration.java
index 26aa42b8d..428838940 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerConfiguration.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerConfiguration.java
@@ -40,7 +40,7 @@ import org.whispersystems.textsecuregcm.configuration.RemoteConfigConfiguration;
 import org.whispersystems.textsecuregcm.configuration.SecureBackupServiceConfiguration;
 import org.whispersystems.textsecuregcm.configuration.SecureStorageServiceConfiguration;
 import org.whispersystems.textsecuregcm.configuration.TestDeviceConfiguration;
-import org.whispersystems.textsecuregcm.configuration.TorExitNodeConfiguration;
+import org.whispersystems.textsecuregcm.configuration.MonitoredS3ObjectConfiguration;
 import org.whispersystems.textsecuregcm.configuration.TurnConfiguration;
 import org.whispersystems.textsecuregcm.configuration.TwilioConfiguration;
 import org.whispersystems.textsecuregcm.configuration.UnidentifiedDeliveryConfiguration;
@@ -259,7 +259,12 @@ public class WhisperServerConfiguration extends Configuration {
   @Valid
   @NotNull
   @JsonProperty
-  private TorExitNodeConfiguration tor;
+  private MonitoredS3ObjectConfiguration torExitNodeList;
+
+  @Valid
+  @NotNull
+  @JsonProperty
+  private MonitoredS3ObjectConfiguration asnTable;
 
   @Valid
   @NotNull
@@ -450,8 +455,12 @@ public class WhisperServerConfiguration extends Configuration {
     return reportMessageDynamoDb;
   }
 
-  public TorExitNodeConfiguration getTorExitNodeConfiguration() {
-    return tor;
+  public MonitoredS3ObjectConfiguration getTorExitNodeListConfiguration() {
+    return torExitNodeList;
+  }
+
+  public MonitoredS3ObjectConfiguration getAsnTableConfiguration() {
+    return asnTable;
   }
 
   public DonationConfiguration getDonationConfiguration() {
diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
index 6c7024cfc..73010709a 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
@@ -186,6 +186,7 @@ import org.whispersystems.textsecuregcm.storage.ReportMessageManager;
 import org.whispersystems.textsecuregcm.storage.ReservedUsernames;
 import org.whispersystems.textsecuregcm.storage.Usernames;
 import org.whispersystems.textsecuregcm.storage.UsernamesManager;
+import org.whispersystems.textsecuregcm.util.AsnManager;
 import org.whispersystems.textsecuregcm.util.Constants;
 import org.whispersystems.textsecuregcm.util.TorExitNodeManager;
 import org.whispersystems.textsecuregcm.websocket.AuthenticatedConnectListener;
@@ -436,7 +437,8 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
     GCMSender                  gcmSender                  = new GCMSender(gcmSenderExecutor, accountsManager, config.getGcmConfiguration().getApiKey());
     RateLimiters               rateLimiters               = new RateLimiters(config.getLimitsConfiguration(), dynamicConfigurationManager, rateLimitersCluster);
     ProvisioningManager        provisioningManager        = new ProvisioningManager(pubSubManager);
-    TorExitNodeManager         torExitNodeManager         = new TorExitNodeManager(recurringJobExecutor, config.getTorExitNodeConfiguration());
+    TorExitNodeManager         torExitNodeManager         = new TorExitNodeManager(recurringJobExecutor, config.getTorExitNodeListConfiguration());
+    AsnManager                 asnManager                 = new AsnManager(recurringJobExecutor, config.getAsnTableConfiguration());
 
     AccountAuthenticator                  accountAuthenticator                  = new AccountAuthenticator(accountsManager);
     DisabledPermittedAccountAuthenticator disabledPermittedAccountAuthenticator = new DisabledPermittedAccountAuthenticator(accountsManager);
@@ -489,6 +491,7 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
     environment.lifecycle().manage(clientPresenceManager);
     environment.lifecycle().manage(currencyManager);
     environment.lifecycle().manage(torExitNodeManager);
+    environment.lifecycle().manage(asnManager);
 
     AWSCredentials         credentials               = new BasicAWSCredentials(config.getCdnConfiguration().getAccessKey(), config.getCdnConfiguration().getAccessSecret());
     AWSCredentialsProvider credentialsProvider       = new AWSStaticCredentialsProvider(credentials);
diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/configuration/TorExitNodeConfiguration.java b/service/src/main/java/org/whispersystems/textsecuregcm/configuration/MonitoredS3ObjectConfiguration.java
similarity index 96%
rename from service/src/main/java/org/whispersystems/textsecuregcm/configuration/TorExitNodeConfiguration.java
rename to service/src/main/java/org/whispersystems/textsecuregcm/configuration/MonitoredS3ObjectConfiguration.java
index 1314e45da..dd82d5f25 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/configuration/TorExitNodeConfiguration.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/configuration/MonitoredS3ObjectConfiguration.java
@@ -11,7 +11,7 @@ import javax.validation.Valid;
 import javax.validation.constraints.NotBlank;
 import java.time.Duration;
 
-public class TorExitNodeConfiguration {
+public class MonitoredS3ObjectConfiguration {
 
   @JsonProperty
   @NotBlank
diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/util/AsnManager.java b/service/src/main/java/org/whispersystems/textsecuregcm/util/AsnManager.java
new file mode 100644
index 000000000..c373ba988
--- /dev/null
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/util/AsnManager.java
@@ -0,0 +1,98 @@
+/*
+ * Copyright 2013-2021 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.textsecuregcm.util;
+
+import static com.codahale.metrics.MetricRegistry.name;
+
+import com.amazonaws.services.s3.model.S3Object;
+import com.google.common.annotations.VisibleForTesting;
+import io.dropwizard.lifecycle.Managed;
+import io.micrometer.core.instrument.Counter;
+import io.micrometer.core.instrument.Metrics;
+import io.micrometer.core.instrument.Timer;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.InputStreamReader;
+import java.net.Inet4Address;
+import java.net.UnknownHostException;
+import java.util.Optional;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.zip.GZIPInputStream;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.whispersystems.textsecuregcm.configuration.MonitoredS3ObjectConfiguration;
+
+public class AsnManager implements Managed {
+
+  private final S3ObjectMonitor asnTableMonitor;
+
+  private final AtomicReference<AsnTable> asnTable = new AtomicReference<>(AsnTable.EMPTY);
+
+  private static final Timer REFRESH_TIMER = Metrics.timer(name(AsnManager.class, "refresh"));
+  private static final Counter REFRESH_ERRORS = Metrics.counter(name(AsnManager.class, "refreshErrors"));
+
+  private static final Logger log = LoggerFactory.getLogger(AsnManager.class);
+
+  public AsnManager(
+      final ScheduledExecutorService scheduledExecutorService,
+      final MonitoredS3ObjectConfiguration configuration) {
+
+    this.asnTableMonitor = new S3ObjectMonitor(
+        configuration.getS3Region(),
+        configuration.getS3Bucket(),
+        configuration.getObjectKey(),
+        configuration.getMaxSize(),
+        scheduledExecutorService,
+        configuration.getRefreshInterval(),
+        this::handleAsnTableChanged);
+  }
+
+  @Override
+  public void start() throws Exception {
+    try {
+      handleAsnTableChanged(asnTableMonitor.getObject());
+    } catch (final Exception e) {
+      log.warn("Failed to load initial IP-to-ASN map", e);
+    }
+
+    asnTableMonitor.start();
+  }
+
+  @Override
+  public void stop() throws Exception {
+    asnTableMonitor.stop();
+  }
+
+  public Optional<Long> getAsn(final String address) {
+    try {
+      return asnTable.get().getAsn((Inet4Address) Inet4Address.getByName(address));
+    } catch (final UnknownHostException e) {
+      log.warn("Could not parse \"{}\" as an Inet4Address", address);
+      return Optional.empty();
+    }
+  }
+
+  private void handleAsnTableChanged(final S3Object asnTableObject) {
+    REFRESH_TIMER.record(() -> {
+      try {
+        handleAsnTableChanged(new GZIPInputStream(asnTableObject.getObjectContent()));
+      } catch (final IOException e) {
+        log.error("Retrieved object was not a gzip archive", e);
+      }
+    });
+  }
+
+  @VisibleForTesting
+  void handleAsnTableChanged(final InputStream inputStream) {
+    try (final InputStreamReader reader = new InputStreamReader(inputStream)) {
+      asnTable.set(new AsnTable(reader));
+    } catch (final Exception e) {
+      REFRESH_ERRORS.increment();
+      log.warn("Failed to refresh IP-to-ASN table", e);
+    }
+  }
+}
diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/util/AsnTable.java b/service/src/main/java/org/whispersystems/textsecuregcm/util/AsnTable.java
index b3387565d..ab2e406fb 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/util/AsnTable.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/util/AsnTable.java
@@ -44,6 +44,8 @@ class AsnTable {
     }
   }
 
+  public static final AsnTable EMPTY = new AsnTable();
+
   public AsnTable(final Reader tsvReader) throws IOException {
     final TreeMap<Long, AsnRange> treeMap = new TreeMap<>();
 
@@ -60,6 +62,10 @@ class AsnTable {
     asnBlocksByFirstIp = treeMap;
   }
 
+  private AsnTable() {
+    asnBlocksByFirstIp = new TreeMap<>();
+  }
+
   public Optional<Long> getAsn(final Inet4Address address) {
     final long addressAsLong = ipToLong(address);
 
diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/util/TorExitNodeManager.java b/service/src/main/java/org/whispersystems/textsecuregcm/util/TorExitNodeManager.java
index 412752c37..00e758972 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/util/TorExitNodeManager.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/util/TorExitNodeManager.java
@@ -24,7 +24,7 @@ import java.util.concurrent.atomic.AtomicReference;
 import java.util.stream.Collectors;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
-import org.whispersystems.textsecuregcm.configuration.TorExitNodeConfiguration;
+import org.whispersystems.textsecuregcm.configuration.MonitoredS3ObjectConfiguration;
 
 /**
  * A utility for checking whether IP addresses belong to Tor exit nodes using the "bulk exit list."
@@ -44,7 +44,7 @@ public class TorExitNodeManager implements Managed {
 
   public TorExitNodeManager(
       final ScheduledExecutorService scheduledExecutorService,
-      final TorExitNodeConfiguration configuration) {
+      final MonitoredS3ObjectConfiguration configuration) {
 
     this.exitListMonitor = new S3ObjectMonitor(
         configuration.getS3Region(),
diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/util/AsnManagerTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/util/AsnManagerTest.java
new file mode 100644
index 000000000..ad756362e
--- /dev/null
+++ b/service/src/test/java/org/whispersystems/textsecuregcm/util/AsnManagerTest.java
@@ -0,0 +1,45 @@
+/*
+ * Copyright 2013-2021 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.textsecuregcm.util;
+
+import org.junit.jupiter.api.Test;
+import org.whispersystems.textsecuregcm.configuration.MonitoredS3ObjectConfiguration;
+
+import java.io.ByteArrayInputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.InputStreamReader;
+import java.net.Inet4Address;
+import java.nio.charset.StandardCharsets;
+import java.util.Optional;
+import java.util.concurrent.ScheduledExecutorService;
+
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+import static org.junit.jupiter.api.Assertions.*;
+import static org.mockito.Mockito.mock;
+
+class AsnManagerTest {
+
+  @Test
+  void getAsn() throws IOException {
+    final MonitoredS3ObjectConfiguration configuration = new MonitoredS3ObjectConfiguration();
+    configuration.setS3Region("ap-northeast-3");
+
+    final AsnManager asnManager = new AsnManager(mock(ScheduledExecutorService.class), configuration);
+
+    assertEquals(Optional.empty(), asnManager.getAsn("10.0.0.1"));
+
+    try (final InputStream tableInputStream = getClass().getResourceAsStream("ip2asn-test.tsv")) {
+      asnManager.handleAsnTableChanged(tableInputStream);
+    }
+
+    assertEquals(Optional.of(7922L), asnManager.getAsn("50.79.54.1"));
+    assertEquals(Optional.of(7552L), asnManager.getAsn("27.79.32.1"));
+    assertEquals(Optional.empty(), asnManager.getAsn("32.79.117.1"));
+    assertEquals(Optional.empty(), asnManager.getAsn("10.0.0.1"));
+  }
+}
diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/util/TorExitNodeManagerTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/util/TorExitNodeManagerTest.java
index 9d2922d1f..76f25d0e6 100644
--- a/service/src/test/java/org/whispersystems/textsecuregcm/util/TorExitNodeManagerTest.java
+++ b/service/src/test/java/org/whispersystems/textsecuregcm/util/TorExitNodeManagerTest.java
@@ -13,14 +13,14 @@ import java.io.ByteArrayInputStream;
 import java.nio.charset.StandardCharsets;
 import java.util.concurrent.ScheduledExecutorService;
 import org.junit.Test;
-import org.whispersystems.textsecuregcm.configuration.TorExitNodeConfiguration;
+import org.whispersystems.textsecuregcm.configuration.MonitoredS3ObjectConfiguration;
 import org.whispersystems.textsecuregcm.redis.AbstractRedisClusterTest;
 
 public class TorExitNodeManagerTest extends AbstractRedisClusterTest {
 
   @Test
   public void testIsTorExitNode() {
-    final TorExitNodeConfiguration configuration = new TorExitNodeConfiguration();
+    final MonitoredS3ObjectConfiguration configuration = new MonitoredS3ObjectConfiguration();
     configuration.setS3Region("ap-northeast-3");
 
     final TorExitNodeManager torExitNodeManager =