Fix ip check and add more tests (#4278)

Signed-off-by: Hai Yan <[email protected]>
opensearch-project · Mar 13, 2024 · 6611631 · 6611631
1 parent 794a6f6
commit 6611631
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 18 deletions.
diff --git a/...essor/src/main/java/org/opensearch/dataprepper/plugins/geoip/utils/IPValidationCheck.java b/...essor/src/main/java/org/opensearch/dataprepper/plugins/geoip/utils/IPValidationCheck.java
@@ -5,12 +5,12 @@
 
 package org.opensearch.dataprepper.plugins.geoip.utils;
 
+import com.google.common.net.InetAddresses;
 import org.opensearch.dataprepper.plugins.geoip.exception.InvalidIPAddressException;
 
 import java.net.Inet4Address;
 import java.net.Inet6Address;
 import java.net.InetAddress;
-import java.net.UnknownHostException;
 
 /**
  * Implementation of class for checking IP validation
@@ -27,8 +27,8 @@ public class IPValidationCheck {
     public static boolean isPublicIpAddress(final String ipAddress) {
         InetAddress address;
         try {
-            address = InetAddress.getByName(ipAddress);
-        } catch (final UnknownHostException e) {
+            address = InetAddresses.forString(ipAddress);
+        } catch (final IllegalArgumentException e) {
             return false;
         }
         if (address instanceof Inet6Address || address instanceof Inet4Address) {

diff --git a/...r/src/test/java/org/opensearch/dataprepper/plugins/geoip/utils/IPValidationCheckTest.java b/...r/src/test/java/org/opensearch/dataprepper/plugins/geoip/utils/IPValidationCheckTest.java
@@ -5,30 +5,32 @@
 
 package org.opensearch.dataprepper.plugins.geoip.utils;
 
-import org.junit.jupiter.api.Assertions;
-import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
 import org.mockito.junit.jupiter.MockitoExtension;
 
+import static org.hamcrest.CoreMatchers.is;
+import static org.hamcrest.MatcherAssert.assertThat;
+
 @ExtendWith(MockitoExtension.class)
 class IPValidationCheckTest {
 
-    private static final String PRIVATE_IP_ADDRESS = "192.168.29.233";
-    private static final String PUBLIC_IP_ADDRESS = "2001:4860:4860::8888";
-    private static final String INVALID_IP_ADDRESS = "255.255.255.999";
-
-    @Test
-    void ipValidationcheckTest_public() {
-        Assertions.assertTrue(IPValidationCheck.isPublicIpAddress(PUBLIC_IP_ADDRESS));
+    @ParameterizedTest
+    @ValueSource(strings = {"93.184.216.34", "172.217.0.0", "142.250.64.0", "2607:f8b0:4005:805::200e"})
+    void ipValidationcheckTest_public(String publicIpAddress) {
+        assertThat(IPValidationCheck.isPublicIpAddress(publicIpAddress), is(true));
     }
 
-    @Test
-    void ipValidationcheckTest_negative() {
-        Assertions.assertFalse(IPValidationCheck.isPublicIpAddress(PRIVATE_IP_ADDRESS));
+    @ParameterizedTest
+    @ValueSource(strings = {"10.0.0.0", "172.16.0.0", "192.168.0.0"})
+    void ipValidationcheckTest_negative(String privateIpAddress) {
+        assertThat(IPValidationCheck.isPublicIpAddress(privateIpAddress), is(false));
     }
 
-    @Test
-    void ipValidationcheckTest_invalid() {
-        Assertions.assertFalse(IPValidationCheck.isPublicIpAddress(INVALID_IP_ADDRESS));
+    @ParameterizedTest
+    @ValueSource(strings = {"123", "255.255.255.999", "true", "[1,2,3]"})
+    void ipValidationcheckTest_invalid(String invalidIpAddress) {
+        assertThat(IPValidationCheck.isPublicIpAddress(invalidIpAddress), is(false));
     }
 }