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 ab2e406fb..d4cabe989 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/util/AsnTable.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/util/AsnTable.java @@ -10,6 +10,8 @@ import java.io.IOException; import java.io.Reader; import java.net.Inet4Address; import java.nio.ByteBuffer; +import java.util.HashMap; +import java.util.Map; import java.util.NavigableMap; import java.util.Optional; import java.util.TreeMap; @@ -21,7 +23,8 @@ import org.apache.commons.csv.CSVRecord; * Allows IP->ASN lookup operations using data from https://iptoasn.com/. */ class AsnTable { - private final NavigableMap asnBlocksByFirstIp; + private final NavigableMap asnBlocksByFirstIp = new TreeMap<>(); + private final Map countryCodesByAsn = new HashMap<>(); private static class AsnRange { private final long rangeStart; @@ -47,23 +50,20 @@ class AsnTable { public static final AsnTable EMPTY = new AsnTable(); public AsnTable(final Reader tsvReader) throws IOException { - final TreeMap treeMap = new TreeMap<>(); - try (final CSVParser csvParser = CSVFormat.TDF.parse(tsvReader)) { for (final CSVRecord record : csvParser) { final long start = Long.parseLong(record.get(0), 10); final long end = Long.parseLong(record.get(1), 10); final long asn = Long.parseLong(record.get(2), 10); + final String countryCode = record.get(3); - treeMap.put(start, new AsnRange(start, end, asn)); + asnBlocksByFirstIp.put(start, new AsnRange(start, end, asn)); + countryCodesByAsn.put(asn, countryCode); } } - - asnBlocksByFirstIp = treeMap; } private AsnTable() { - asnBlocksByFirstIp = new TreeMap<>(); } public Optional getAsn(final Inet4Address address) { @@ -74,6 +74,10 @@ class AsnTable { .map(entry -> entry.getValue().getAsn()); } + public Optional getCountryCode(final long asn) { + return Optional.ofNullable(countryCodesByAsn.get(asn)); + } + @VisibleForTesting static long ipToLong(final Inet4Address address) { final ByteBuffer buffer = ByteBuffer.allocate(8); diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/util/AsnTableTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/util/AsnTableTest.java index edf19385b..fa3f198b5 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/util/AsnTableTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/util/AsnTableTest.java @@ -30,6 +30,17 @@ class AsnTableTest { } } + @Test + void getCountryCode() throws IOException { + try (final InputStreamReader reader = new InputStreamReader(getClass().getResourceAsStream("ip2asn-test.tsv"))) { + final AsnTable asnTable = new AsnTable(reader); + + assertEquals(Optional.of("US"), asnTable.getCountryCode(7922)); + assertEquals(Optional.of("VN"), asnTable.getCountryCode(7552)); + assertEquals(Optional.empty(), asnTable.getCountryCode(1234)); + } + } + @Test void ipToLong() throws UnknownHostException { assertEquals(0x00000000ffffffffL, AsnTable.ipToLong((Inet4Address) Inet4Address.getByName("255.255.255.255")));