ingest-geoip: establish boundaries (#109655)

Author: yaauie

modules/ingest-geoip/src/internalClusterTest/java/org/elasticsearch/ingest/geoip/DatabaseNodeServiceIT.java

@@ -84,7 +84,7 @@ private void assertValidDatabase(DatabaseNodeService databaseNodeService, String
         IpDatabase database = databaseNodeService.getDatabase(databaseFileName);
         assertNotNull(database);
         assertThat(database.getDatabaseType(), equalTo(databaseType));
-        CountryResponse countryResponse = database.getCountry("89.160.20.128");
+        CountryResponse countryResponse = database.getResponse("89.160.20.128", GeoIpTestUtils::getCountry);
         assertNotNull(countryResponse);
         Country country = countryResponse.getCountry();
         assertNotNull(country);

modules/ingest-geoip/src/internalClusterTest/java/org/elasticsearch/ingest/geoip/ReloadingDatabasesWhilePerformingGeoLookupsIT.java

@@ -205,10 +205,10 @@ private static DatabaseNodeService createRegistry(Path geoIpConfigDir, Path geoI
     private static void lazyLoadReaders(DatabaseNodeService databaseNodeService) throws IOException {
         if (databaseNodeService.get("GeoLite2-City.mmdb") != null) {
             databaseNodeService.get("GeoLite2-City.mmdb").getDatabaseType();
-            databaseNodeService.get("GeoLite2-City.mmdb").getCity("2.125.160.216");
+            databaseNodeService.get("GeoLite2-City.mmdb").getResponse("2.125.160.216", GeoIpTestUtils::getCity);
         }
         databaseNodeService.get("GeoLite2-City-Test.mmdb").getDatabaseType();
-        databaseNodeService.get("GeoLite2-City-Test.mmdb").getCity("2.125.160.216");
+        databaseNodeService.get("GeoLite2-City-Test.mmdb").getResponse("2.125.160.216", GeoIpTestUtils::getCity);
     }
 
 }

modules/ingest-geoip/src/main/java/org/elasticsearch/ingest/geoip/Database.java

@@ -9,7 +9,6 @@
 
 package org.elasticsearch.ingest.geoip;
 
-import org.elasticsearch.common.Strings;
 import org.elasticsearch.core.Nullable;
 
 import java.util.Arrays;
@@ -19,12 +18,10 @@
 import java.util.Set;
 
 /**
- * A high-level representation of a kind of geoip database that is supported by the {@link GeoIpProcessor}.
+ * A high-level representation of a kind of ip location database that is supported by the {@link GeoIpProcessor}.
  * <p>
  * A database has a set of properties that are valid to use with it (see {@link Database#properties()}),
  * as well as a list of default properties to use if no properties are specified (see {@link Database#defaultProperties()}).
- * <p>
- * See especially {@link Database#getDatabase(String, String)} which is used to obtain instances of this class.
  */
 enum Database {
 
@@ -142,61 +139,6 @@ enum Database {
         )
     );
 
-    private static final String CITY_DB_SUFFIX = "-City";
-    private static final String COUNTRY_DB_SUFFIX = "-Country";
-    private static final String ASN_DB_SUFFIX = "-ASN";
-    private static final String ANONYMOUS_IP_DB_SUFFIX = "-Anonymous-IP";
-    private static final String CONNECTION_TYPE_DB_SUFFIX = "-Connection-Type";
-    private static final String DOMAIN_DB_SUFFIX = "-Domain";
-    private static final String ENTERPRISE_DB_SUFFIX = "-Enterprise";
-    private static final String ISP_DB_SUFFIX = "-ISP";
-
-    @Nullable
-    private static Database getMaxmindDatabase(final String databaseType) {
-        if (databaseType.endsWith(Database.CITY_DB_SUFFIX)) {
-            return Database.City;
-        } else if (databaseType.endsWith(Database.COUNTRY_DB_SUFFIX)) {
-            return Database.Country;
-        } else if (databaseType.endsWith(Database.ASN_DB_SUFFIX)) {
-            return Database.Asn;
-        } else if (databaseType.endsWith(Database.ANONYMOUS_IP_DB_SUFFIX)) {
-            return Database.AnonymousIp;
-        } else if (databaseType.endsWith(Database.CONNECTION_TYPE_DB_SUFFIX)) {
-            return Database.ConnectionType;
-        } else if (databaseType.endsWith(Database.DOMAIN_DB_SUFFIX)) {
-            return Database.Domain;
-        } else if (databaseType.endsWith(Database.ENTERPRISE_DB_SUFFIX)) {
-            return Database.Enterprise;
-        } else if (databaseType.endsWith(Database.ISP_DB_SUFFIX)) {
-            return Database.Isp;
-        } else {
-            return null; // no match was found
-        }
-    }
-
-    /**
-     * Parses the passed-in databaseType (presumably from the passed-in databaseFile) and return the Database instance that is
-     * associated with that databaseType.
-     *
-     * @param databaseType the database type String from the metadata of the database file
-     * @param databaseFile the database file from which the database type was obtained
-     * @throws IllegalArgumentException if the databaseType is not associated with a Database instance
-     * @return the Database instance that is associated with the databaseType
-     */
-    public static Database getDatabase(final String databaseType, final String databaseFile) {
-        Database database = null;
-
-        if (Strings.hasText(databaseType)) {
-            database = getMaxmindDatabase(databaseType);
-        }
-
-        if (database == null) {
-            throw new IllegalArgumentException("Unsupported database type [" + databaseType + "] for file [" + databaseFile + "]");
-        }
-
-        return database;
-    }
-
     private final Set<Property> properties;
     private final Set<Property> defaultProperties;
 

modules/ingest-geoip/src/main/java/org/elasticsearch/ingest/geoip/DatabaseReaderLazyLoader.java

@@ -9,47 +9,32 @@
 
 package org.elasticsearch.ingest.geoip;
 
-import com.maxmind.db.DatabaseRecord;
-import com.maxmind.db.Network;
 import com.maxmind.db.NoCache;
 import com.maxmind.db.Reader;
-import com.maxmind.geoip2.model.AnonymousIpResponse;
-import com.maxmind.geoip2.model.AsnResponse;
-import com.maxmind.geoip2.model.CityResponse;
-import com.maxmind.geoip2.model.ConnectionTypeResponse;
-import com.maxmind.geoip2.model.CountryResponse;
-import com.maxmind.geoip2.model.DomainResponse;
-import com.maxmind.geoip2.model.EnterpriseResponse;
-import com.maxmind.geoip2.model.IspResponse;
 
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.apache.lucene.util.SetOnce;
 import org.elasticsearch.ExceptionsHelper;
 import org.elasticsearch.common.CheckedBiFunction;
 import org.elasticsearch.common.CheckedSupplier;
-import org.elasticsearch.common.network.InetAddresses;
-import org.elasticsearch.common.network.NetworkAddress;
 import org.elasticsearch.core.Booleans;
 import org.elasticsearch.core.IOUtils;
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.core.SuppressForbidden;
 
 import java.io.File;
 import java.io.IOException;
-import java.net.InetAddress;
 import java.nio.file.Files;
 import java.nio.file.Path;
-import java.util.List;
 import java.util.Objects;
-import java.util.Optional;
 import java.util.concurrent.atomic.AtomicInteger;
 
 /**
  * Facilitates lazy loading of the database reader, so that when the geoip plugin is installed, but not used,
  * no memory is being wasted on the database reader.
  */
-class DatabaseReaderLazyLoader implements IpDatabase {
+public class DatabaseReaderLazyLoader implements IpDatabase {
 
     private static final boolean LOAD_DATABASE_ON_HEAP = Booleans.parseBoolean(System.getProperty("es.geoip.load_db_on_heap", "false"));
 
@@ -96,94 +81,6 @@ public final String getDatabaseType() throws IOException {
         return databaseType.get();
     }
 
-    @Nullable
-    @Override
-    public CityResponse getCity(String ipAddress) {
-        return getResponse(ipAddress, (reader, ip) -> lookup(reader, ip, CityResponse.class, CityResponse::new));
-    }
-
-    @Nullable
-    @Override
-    public CountryResponse getCountry(String ipAddress) {
-        return getResponse(ipAddress, (reader, ip) -> lookup(reader, ip, CountryResponse.class, CountryResponse::new));
-    }
-
-    @Nullable
-    @Override
-    public AsnResponse getAsn(String ipAddress) {
-        return getResponse(
-            ipAddress,
-            (reader, ip) -> lookup(
-                reader,
-                ip,
-                AsnResponse.class,
-                (response, responseIp, network, locales) -> new AsnResponse(response, responseIp, network)
-            )
-        );
-    }
-
-    @Nullable
-    @Override
-    public AnonymousIpResponse getAnonymousIp(String ipAddress) {
-        return getResponse(
-            ipAddress,
-            (reader, ip) -> lookup(
-                reader,
-                ip,
-                AnonymousIpResponse.class,
-                (response, responseIp, network, locales) -> new AnonymousIpResponse(response, responseIp, network)
-            )
-        );
-    }
-
-    @Nullable
-    @Override
-    public ConnectionTypeResponse getConnectionType(String ipAddress) {
-        return getResponse(
-            ipAddress,
-            (reader, ip) -> lookup(
-                reader,
-                ip,
-                ConnectionTypeResponse.class,
-                (response, responseIp, network, locales) -> new ConnectionTypeResponse(response, responseIp, network)
-            )
-        );
-    }
-
-    @Nullable
-    @Override
-    public DomainResponse getDomain(String ipAddress) {
-        return getResponse(
-            ipAddress,
-            (reader, ip) -> lookup(
-                reader,
-                ip,
-                DomainResponse.class,
-                (response, responseIp, network, locales) -> new DomainResponse(response, responseIp, network)
-            )
-        );
-    }
-
-    @Nullable
-    @Override
-    public EnterpriseResponse getEnterprise(String ipAddress) {
-        return getResponse(ipAddress, (reader, ip) -> lookup(reader, ip, EnterpriseResponse.class, EnterpriseResponse::new));
-    }
-
-    @Nullable
-    @Override
-    public IspResponse getIsp(String ipAddress) {
-        return getResponse(
-            ipAddress,
-            (reader, ip) -> lookup(
-                reader,
-                ip,
-                IspResponse.class,
-                (response, responseIp, network, locales) -> new IspResponse(response, responseIp, network)
-            )
-        );
-    }
-
     boolean preLookup() {
         return currentUsages.updateAndGet(current -> current < 0 ? current : current + 1) > 0;
     }
@@ -199,14 +96,12 @@ int current() {
         return currentUsages.get();
     }
 
+    @Override
     @Nullable
-    private <RESPONSE> RESPONSE getResponse(
-        String ipAddress,
-        CheckedBiFunction<Reader, String, Optional<RESPONSE>, Exception> responseProvider
-    ) {
+    public <RESPONSE> RESPONSE getResponse(String ipAddress, CheckedBiFunction<Reader, String, RESPONSE, Exception> responseProvider) {
         return cache.putIfAbsent(ipAddress, databasePath.toString(), ip -> {
             try {
-                return responseProvider.apply(get(), ipAddress).orElse(null);
+                return responseProvider.apply(get(), ipAddress);
             } catch (Exception e) {
                 throw ExceptionsHelper.convertToRuntime(e);
             }
@@ -263,23 +158,6 @@ private static File pathToFile(Path databasePath) {
         return databasePath.toFile();
     }
 
-    @FunctionalInterface
-    private interface ResponseBuilder<RESPONSE> {
-        RESPONSE build(RESPONSE response, String responseIp, Network network, List<String> locales);
-    }
-
-    private <RESPONSE> Optional<RESPONSE> lookup(Reader reader, String ip, Class<RESPONSE> clazz, ResponseBuilder<RESPONSE> builder)
-        throws IOException {
-        InetAddress inetAddress = InetAddresses.forString(ip);
-        DatabaseRecord<RESPONSE> record = reader.getRecord(inetAddress, clazz);
-        RESPONSE result = record.getData();
-        if (result == null) {
-            return Optional.empty();
-        } else {
-            return Optional.of(builder.build(result, NetworkAddress.format(inetAddress), record.getNetwork(), List.of("en")));
-        }
-    }
-
     long getBuildDateMillis() throws IOException {
         if (buildDate.get() == null) {
             synchronized (buildDate) {

modules/ingest-geoip/src/main/java/org/elasticsearch/ingest/geoip/GeoIpCache.java

@@ -26,7 +26,7 @@
  * cost of deserialization for each lookup (cached or not). This comes at slight expense of higher memory usage, but significant
  * reduction of CPU usage.
  */
-final class GeoIpCache {
+public final class GeoIpCache {
 
     /**
      * Internal-only sentinel object for recording that a result from the geoip database was null (i.e. there was no result). By caching