Add lookup_private_addresses option to GeoIP processor for internal I… (#6770)

Add lookup_private_addresses option to GeoIP processor for internal IP enrichment

Signed-off-by: Srikanth Padakanti <srikanth_padakanti@apple.com>

Author: srikanthpadakanti

data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/geoip/processor/GeoIPProcessor.java

@@ -148,7 +148,7 @@ private void processRecords(Collection<Record<Event>> records, GeoIPDatabaseRead
         //Lookup from DB
         if (ipAddress != null && !ipAddress.isEmpty()) {
           try {
-            final Optional<InetAddress> optionalInetAddress = GeoInetAddress.usableInetFromString(ipAddress);
+            final Optional<InetAddress> optionalInetAddress = GeoInetAddress.usableInetFromString(ipAddress, geoIPProcessorConfig.getLookupPrivateAddresses());
             if (optionalInetAddress.isPresent()) {
                 geoData = geoIPDatabaseReader.getGeoData(optionalInetAddress.get(), fields, databases);
                 if (geoData.isEmpty()) {

data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/geoip/processor/GeoIPProcessorConfig.java

@@ -52,6 +52,11 @@ public class GeoIPProcessorConfig {
     })
     private String whenCondition;
 
+    @JsonProperty("lookup_private_addresses")
+    @JsonPropertyDescription("When set to <code>true</code>, the processor will also enrich private/internal IP addresses (e.g. 10.x, 172.16.x, 192.168.x, loopback). " +
+            "This is useful when using custom MMDB files enriched with internal IPAM data. Default is <code>false</code>.")
+    private boolean lookupPrivateAddresses = false;
+
     /**
      * Get List of entries
      * @return List of EntryConfig
@@ -91,4 +96,8 @@ public List<String> getTagsOnNoValidIp() {
     public String getWhenCondition() {
         return whenCondition;
     }
+
+    public boolean getLookupPrivateAddresses() {
+        return lookupPrivateAddresses;
+    }
 }

data-prepper-plugins/geoip-processor/src/main/java/org/opensearch/dataprepper/plugins/geoip/processor/GeoInetAddress.java

@@ -19,13 +19,17 @@
 class GeoInetAddress {
 
     static Optional<InetAddress> usableInetFromString(final String ipAddress) {
+        return usableInetFromString(ipAddress, false);
+    }
+
+    static Optional<InetAddress> usableInetFromString(final String ipAddress, final boolean lookupPrivateAddresses) {
         final InetAddress address;
         try {
             address = InetAddresses.forString(ipAddress);
         } catch (final IllegalArgumentException e) {
             return Optional.empty();
         }
-        if (isPublicIpAddress(address))
+        if (isPublicIpAddress(address) || lookupPrivateAddresses)
             return Optional.of(address);
         return Optional.empty();
     }

data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/geoip/processor/GeoIPProcessorConfigTest.java

@@ -32,6 +32,7 @@ void testDefaultConfig() {
         assertThat(geoIPProcessorConfig.getTagsOnEngineFailure(), equalTo(null));
         assertThat(geoIPProcessorConfig.getTagsOnIPNotFound(), equalTo(null));
         assertThat(geoIPProcessorConfig.getWhenCondition(), equalTo(null));
+        assertThat(geoIPProcessorConfig.getLookupPrivateAddresses(), equalTo(false));
     }
 
     @Test
@@ -51,4 +52,10 @@ void testGetEntries() throws NoSuchFieldException, IllegalAccessException {
         assertThat(geoIPProcessorConfig.getTagsOnIPNotFound(), equalTo(tagsOnIPNotFound));
         assertThat(geoIPProcessorConfig.getWhenCondition(), equalTo(whenCondition));
     }
+
+    @Test
+    void testLookupPrivateAddressesWhenSet() throws NoSuchFieldException, IllegalAccessException {
+        ReflectivelySetField.setField(GeoIPProcessorConfig.class, geoIPProcessorConfig, "lookupPrivateAddresses", true);
+        assertThat(geoIPProcessorConfig.getLookupPrivateAddresses(), equalTo(true));
+    }
 }

data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/geoip/processor/GeoIPProcessorTest.java

@@ -47,6 +47,7 @@
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.lenient;
 import static org.mockito.Mockito.mockStatic;
@@ -313,7 +314,7 @@ void doExecuteTest_should_not_add_geodata_if_ip_address_is_not_public() {
         when(geoIPProcessorConfig.getTagsOnNoValidIp()).thenReturn(testTags);
 
         try (final MockedStatic<GeoInetAddress> ipValidationCheckMockedStatic = mockStatic(GeoInetAddress.class)) {
-            ipValidationCheckMockedStatic.when(() -> GeoInetAddress.usableInetFromString(any())).thenReturn(Optional.empty());
+            ipValidationCheckMockedStatic.when(() -> GeoInetAddress.usableInetFromString(any(), anyBoolean())).thenReturn(Optional.empty());
 
             when(geoIPProcessorConfig.getEntries()).thenReturn(List.of(entry));
             when(entry.getSource()).thenReturn(SOURCE);

data-prepper-plugins/geoip-processor/src/test/java/org/opensearch/dataprepper/plugins/geoip/processor/GeoInetAddressTest.java

@@ -45,4 +45,63 @@ void ipValidationcheckTest_invalid(String invalidIpAddress) {
         assertThat(actual, notNullValue());
         assertThat(actual.isPresent(), equalTo(false));
     }
-}
+
+    @ParameterizedTest
+    @ValueSource(strings = {"10.0.0.1", "172.16.0.1", "192.168.1.1"})
+    void privateIpAcceptedWhenLookupEnabled(String privateIpAddress) {
+        final Optional<InetAddress> actual = GeoInetAddress.usableInetFromString(privateIpAddress, true);
+        assertThat(actual, notNullValue());
+        assertThat(actual.isPresent(), equalTo(true));
+        assertThat(actual.get(), equalTo(InetAddresses.forString(privateIpAddress)));
+    }
+
+    @ParameterizedTest
+    @ValueSource(strings = {"10.0.0.1", "172.16.0.1", "192.168.1.1"})
+    void privateIpRejectedWhenLookupDisabled(String privateIpAddress) {
+        final Optional<InetAddress> actual = GeoInetAddress.usableInetFromString(privateIpAddress, false);
+        assertThat(actual, notNullValue());
+        assertThat(actual.isPresent(), equalTo(false));
+    }
+
+    @ParameterizedTest
+    @ValueSource(strings = {"127.0.0.1", "::1"})
+    void loopbackAcceptedWhenLookupEnabled(String loopbackAddress) {
+        final Optional<InetAddress> actual = GeoInetAddress.usableInetFromString(loopbackAddress, true);
+        assertThat(actual, notNullValue());
+        assertThat(actual.isPresent(), equalTo(true));
+    }
+
+    @ParameterizedTest
+    @ValueSource(strings = {"93.184.216.34", "2607:f8b0:4005:805::200e"})
+    void publicIpWorksInBothModes(String publicIpAddress) {
+        final Optional<InetAddress> withFlag = GeoInetAddress.usableInetFromString(publicIpAddress, true);
+        final Optional<InetAddress> withoutFlag = GeoInetAddress.usableInetFromString(publicIpAddress, false);
+        assertThat(withFlag.isPresent(), equalTo(true));
+        assertThat(withoutFlag.isPresent(), equalTo(true));
+    }
+
+    @ParameterizedTest
+    @ValueSource(strings = {"123", "255.255.255.999", "true", "[1,2,3]"})
+    void invalidIpRejectedInBothModes(String invalidIpAddress) {
+        final Optional<InetAddress> withFlag = GeoInetAddress.usableInetFromString(invalidIpAddress, true);
+        final Optional<InetAddress> withoutFlag = GeoInetAddress.usableInetFromString(invalidIpAddress, false);
+        assertThat(withFlag.isPresent(), equalTo(false));
+        assertThat(withoutFlag.isPresent(), equalTo(false));
+    }
+
+    @ParameterizedTest
+    @ValueSource(strings = {"224.0.0.1", "239.255.255.250"})
+    void multicastAcceptedWhenLookupEnabled(String multicastAddress) {
+        final Optional<InetAddress> actual = GeoInetAddress.usableInetFromString(multicastAddress, true);
+        assertThat(actual, notNullValue());
+        assertThat(actual.isPresent(), equalTo(true));
+    }
+
+    @ParameterizedTest
+    @ValueSource(strings = {"169.254.1.1", "fe80::1"})
+    void linkLocalAcceptedWhenLookupEnabled(String linkLocalAddress) {
+        final Optional<InetAddress> actual = GeoInetAddress.usableInetFromString(linkLocalAddress, true);
+        assertThat(actual, notNullValue());
+        assertThat(actual.isPresent(), equalTo(true));
+    }
+}