IPinfo geolocation support (#114311) (#114373)

joegallo · web-flow · commit 80fb2ca371d0 · 2024-10-09T10:48:17.000+11:00
diff --git a/modules/ingest-geoip/src/main/java/org/elasticsearch/ingest/geoip/Database.java b/modules/ingest-geoip/src/main/java/org/elasticsearch/ingest/geoip/Database.java
@@ -169,7 +169,19 @@ enum Database {
             Property.TYPE
         ),
         Set.of(Property.IP, Property.ASN, Property.ORGANIZATION_NAME, Property.NETWORK)
-    );
+    ),
+    CityV2(
+        Set.of(
+            Property.IP,
+            Property.COUNTRY_ISO_CODE,
+            Property.REGION_NAME,
+            Property.CITY_NAME,
+            Property.TIMEZONE,
+            Property.LOCATION,
+            Property.POSTAL_CODE
+        ),
+        Set.of(Property.COUNTRY_ISO_CODE, Property.REGION_NAME, Property.CITY_NAME, Property.LOCATION)
+    ),;
 
     private final Set<Property> properties;
     private final Set<Property> defaultProperties;
diff --git a/modules/ingest-geoip/src/main/java/org/elasticsearch/ingest/geoip/IpinfoIpDataLookups.java b/modules/ingest-geoip/src/main/java/org/elasticsearch/ingest/geoip/IpinfoIpDataLookups.java
@@ -58,6 +58,25 @@ static Long parseAsn(final String asn) {
         }
     }
 
+    /**
+     * Lax-ly parses a string that contains a double into a Double (or null, if such parsing isn't possible).
+     * @param latlon a potentially empty (or null) string that is expected to contain a parsable double
+     * @return the parsed double
+     */
+    static Double parseLocationDouble(final String latlon) {
+        if (latlon == null || Strings.hasText(latlon) == false) {
+            return null;
+        } else {
+            String stripped = latlon.trim();
+            try {
+                return Double.parseDouble(stripped);
+            } catch (NumberFormatException e) {
+                logger.trace("Unable to parse non-compliant location string [{}]", latlon);
+                return null;
+            }
+        }
+    }
+
     public record AsnResult(
         Long asn,
         @Nullable String country, // not present in the free asn database
@@ -88,6 +107,31 @@ public record CountryResult(
         public CountryResult {}
     }
 
+    public record GeolocationResult(
+        String city,
+        String country,
+        Double latitude,
+        Double longitude,
+        String postalCode,
+        String region,
+        String timezone
+    ) {
+        @SuppressWarnings("checkstyle:RedundantModifier")
+        @MaxMindDbConstructor
+        public GeolocationResult(
+            @MaxMindDbParameter(name = "city") String city,
+            @MaxMindDbParameter(name = "country") String country,
+            @MaxMindDbParameter(name = "latitude") String latitude,
+            @MaxMindDbParameter(name = "longitude") String longitude,
+            // @MaxMindDbParameter(name = "network") String network, // for now we're not exposing this
+            @MaxMindDbParameter(name = "postal_code") String postalCode,
+            @MaxMindDbParameter(name = "region") String region,
+            @MaxMindDbParameter(name = "timezone") String timezone
+        ) {
+            this(city, country, parseLocationDouble(latitude), parseLocationDouble(longitude), postalCode, region, timezone);
+        }
+    }
+
     static class Asn extends AbstractBase<AsnResult> {
         Asn(Set<Database.Property> properties) {
             super(properties, AsnResult.class);
@@ -183,6 +227,65 @@ protected Map<String, Object> transform(final Result<CountryResult> result) {
         }
     }
 
+    static class Geolocation extends AbstractBase<GeolocationResult> {
+        Geolocation(final Set<Database.Property> properties) {
+            super(properties, GeolocationResult.class);
+        }
+
+        @Override
+        protected Map<String, Object> transform(final Result<GeolocationResult> result) {
+            GeolocationResult response = result.result;
+
+            Map<String, Object> data = new HashMap<>();
+            for (Database.Property property : this.properties) {
+                switch (property) {
+                    case IP -> data.put("ip", result.ip);
+                    case COUNTRY_ISO_CODE -> {
+                        String countryIsoCode = response.country;
+                        if (countryIsoCode != null) {
+                            data.put("country_iso_code", countryIsoCode);
+                        }
+                    }
+                    case REGION_NAME -> {
+                        String subdivisionName = response.region;
+                        if (subdivisionName != null) {
+                            data.put("region_name", subdivisionName);
+                        }
+                    }
+                    case CITY_NAME -> {
+                        String cityName = response.city;
+                        if (cityName != null) {
+                            data.put("city_name", cityName);
+                        }
+                    }
+                    case TIMEZONE -> {
+                        String locationTimeZone = response.timezone;
+                        if (locationTimeZone != null) {
+                            data.put("timezone", locationTimeZone);
+                        }
+                    }
+                    case POSTAL_CODE -> {
+                        String postalCode = response.postalCode;
+                        if (postalCode != null) {
+                            data.put("postal_code", postalCode);
+                        }
+                    }
+                    case LOCATION -> {
+                        Double latitude = response.latitude;
+                        Double longitude = response.longitude;
+                        if (latitude != null && longitude != null) {
+                            Map<String, Object> locationObject = new HashMap<>();
+                            locationObject.put("lat", latitude);
+                            locationObject.put("lon", longitude);
+                            data.put("location", locationObject);
+                        }
+                    }
+                }
+            }
+            return data;
+        }
+    }
+
     /**
      * Just a little record holder -- there's the data that we receive via the binding to our record objects from the Reader via the
      * getRecord call, but then we also need to capture the passed-in ip address that came from the caller as well as the network for
diff --git a/modules/ingest-geoip/src/test/java/org/elasticsearch/ingest/geoip/IpinfoIpDataLookupsTests.java b/modules/ingest-geoip/src/test/java/org/elasticsearch/ingest/geoip/IpinfoIpDataLookupsTests.java
@@ -38,6 +38,7 @@
 import static java.util.Map.entry;
 import static org.elasticsearch.ingest.geoip.GeoIpTestUtils.copyDatabase;
 import static org.elasticsearch.ingest.geoip.IpinfoIpDataLookups.parseAsn;
+import static org.elasticsearch.ingest.geoip.IpinfoIpDataLookups.parseLocationDouble;
 import static org.hamcrest.Matchers.empty;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.is;
@@ -72,6 +73,10 @@ public void testDatabasePropertyInvariants() {
         // the second ASN variant database is like a specialization of the ASN database
         assertThat(Sets.difference(Database.Asn.properties(), Database.AsnV2.properties()), is(empty()));
         assertThat(Database.Asn.defaultProperties(), equalTo(Database.AsnV2.defaultProperties()));
+
+        // the second City variant database is like a version of the ordinary City database but lacking many fields
+        assertThat(Sets.difference(Database.CityV2.properties(), Database.City.properties()), is(empty()));
+        assertThat(Sets.difference(Database.CityV2.defaultProperties(), Database.City.defaultProperties()), is(empty()));
     }
 
     public void testParseAsn() {
@@ -88,6 +93,18 @@ public void testParseAsn() {
         assertThat(parseAsn("anythingelse"), nullValue());
     }
 
+    public void testParseLocationDouble() {
+        // expected case: "123.45" is 123.45
+        assertThat(parseLocationDouble("123.45"), equalTo(123.45));
+        // defensive cases: null and empty becomes null, this is not expected fwiw
+        assertThat(parseLocationDouble(null), nullValue());
+        assertThat(parseLocationDouble(""), nullValue());
+        // defensive cases: we strip whitespace
+        assertThat(parseLocationDouble("  -123.45  "), equalTo(-123.45));
+        // bottom case: a non-parsable string is null
+        assertThat(parseLocationDouble("anythingelse"), nullValue());
+    }
+
     public void testAsn() throws IOException {
         assumeFalse("https://github.com/elastic/elasticsearch/issues/114266", Constants.WINDOWS);
         Path configDir = tmpDir;
@@ -100,7 +117,7 @@ public void testAsn() throws IOException {
 
         // this is the 'free' ASN database (sample)
         try (DatabaseReaderLazyLoader loader = configDatabases.getDatabase("ip_asn_sample.mmdb")) {
-            IpDataLookup lookup = new IpinfoIpDataLookups.Asn(Set.of(Database.Property.values()));
+            IpDataLookup lookup = new IpinfoIpDataLookups.Asn(Database.AsnV2.properties());
             Map<String, Object> data = lookup.getData(loader, "5.182.109.0");
             assertThat(
                 data,
@@ -118,7 +135,7 @@ public void testAsn() throws IOException {
 
         // this is the non-free or 'standard' ASN database (sample)
         try (DatabaseReaderLazyLoader loader = configDatabases.getDatabase("asn_sample.mmdb")) {
-            IpDataLookup lookup = new IpinfoIpDataLookups.Asn(Set.of(Database.Property.values()));
+            IpDataLookup lookup = new IpinfoIpDataLookups.Asn(Database.AsnV2.properties());
             Map<String, Object> data = lookup.getData(loader, "23.53.116.0");
             assertThat(
                 data,
@@ -185,7 +202,7 @@ public void testCountry() throws IOException {
 
         // this is the 'free' Country database (sample)
         try (DatabaseReaderLazyLoader loader = configDatabases.getDatabase("ip_country_sample.mmdb")) {
-            IpDataLookup lookup = new IpinfoIpDataLookups.Country(Set.of(Database.Property.values()));
+            IpDataLookup lookup = new IpinfoIpDataLookups.Country(Database.Country.properties());
             Map<String, Object> data = lookup.getData(loader, "4.221.143.168");
             assertThat(
                 data,
@@ -202,6 +219,74 @@ public void testCountry() throws IOException {
         }
     }
 
+    public void testGeolocation() throws IOException {
+        assumeFalse("https://github.com/elastic/elasticsearch/issues/114266", Constants.WINDOWS);
+        Path configDir = tmpDir;
+        copyDatabase("ipinfo/ip_geolocation_sample.mmdb", configDir.resolve("ip_geolocation_sample.mmdb"));
+
+        GeoIpCache cache = new GeoIpCache(1000); // real cache to test purging of entries upon a reload
+        ConfigDatabases configDatabases = new ConfigDatabases(configDir, cache);
+        configDatabases.initialize(resourceWatcherService);
+
+        // this is the non-free or 'standard' Geolocation database (sample)
+        try (DatabaseReaderLazyLoader loader = configDatabases.getDatabase("ip_geolocation_sample.mmdb")) {
+            IpDataLookup lookup = new IpinfoIpDataLookups.Geolocation(Database.CityV2.properties());
+            Map<String, Object> data = lookup.getData(loader, "2.124.90.182");
+            assertThat(
+                data,
+                equalTo(
+                    Map.ofEntries(
+                        entry("ip", "2.124.90.182"),
+                        entry("country_iso_code", "GB"),
+                        entry("region_name", "England"),
+                        entry("city_name", "London"),
+                        entry("timezone", "Europe/London"),
+                        entry("postal_code", "E1W"),
+                        entry("location", Map.of("lat", 51.50853, "lon", -0.12574))
+                    )
+                )
+            );
+        }
+    }
+
+    public void testGeolocationInvariants() {
+        assumeFalse("https://github.com/elastic/elasticsearch/issues/114266", Constants.WINDOWS);
+        Path configDir = tmpDir;
+        copyDatabase("ipinfo/ip_geolocation_sample.mmdb", configDir.resolve("ip_geolocation_sample.mmdb"));
+
+        {
+            final Set<String> expectedColumns = Set.of(
+                "network",
+                "city",
+                "region",
+                "country",
+                "postal_code",
+                "timezone",
+                "latitude",
+                "longitude"
+            );
+
+            Path databasePath = configDir.resolve("ip_geolocation_sample.mmdb");
+            assertDatabaseInvariants(databasePath, (ip, row) -> {
+                assertThat(row.keySet(), equalTo(expectedColumns));
+                {
+                    String latitude = (String) row.get("latitude");
+                    assertThat(latitude, equalTo(latitude.trim()));
+                    Double parsed = parseLocationDouble(latitude);
+                    assertThat(parsed, notNullValue());
+                    assertThat(latitude, equalTo(Double.toString(parsed))); // reverse it
+                }
+                {
+                    String longitude = (String) row.get("longitude");
+                    assertThat(longitude, equalTo(longitude.trim()));
+                    Double parsed = parseLocationDouble(longitude);
+                    assertThat(parsed, notNullValue());
+                    assertThat(longitude, equalTo(Double.toString(parsed))); // reverse it
+                }
+            });
+        }
+    }
+
     private static void assertDatabaseInvariants(final Path databasePath, final BiConsumer<InetAddress, Map<String, Object>> rowConsumer) {
         try (Reader reader = new Reader(pathToFile(databasePath))) {
             Networks<?> networks = reader.networks(Map.class);
diff --git a/modules/ingest-geoip/src/test/java/org/elasticsearch/ingest/geoip/MaxMindSupportTests.java b/modules/ingest-geoip/src/test/java/org/elasticsearch/ingest/geoip/MaxMindSupportTests.java
@@ -361,7 +361,7 @@ public class MaxMindSupportTests extends ESTestCase {
 
     private static final Set<Class<? extends AbstractResponse>> KNOWN_UNSUPPORTED_RESPONSE_CLASSES = Set.of(IpRiskResponse.class);
 
-    private static final Set<Database> KNOWN_UNSUPPORTED_DATABASE_VARIANTS = Set.of(Database.AsnV2);
+    private static final Set<Database> KNOWN_UNSUPPORTED_DATABASE_VARIANTS = Set.of(Database.AsnV2, Database.CityV2);
 
     public void testMaxMindSupport() {
         for (Database databaseType : Database.values()) {
diff --git a/modules/ingest-geoip/src/test/resources/ipinfo/ip_geolocation_sample.mmdb b/modules/ingest-geoip/src/test/resources/ipinfo/ip_geolocation_sample.mmdb
