diff --git a/pom.xml b/pom.xml index 231ddc555038..16a187f0d029 100644 --- a/pom.xml +++ b/pom.xml @@ -102,9 +102,21 @@ 2.8.1 - org.testng - testng - 6.9.10 + junit + junit + ${junit.version} + test + + + org.assertj + assertj-core + ${assertj-core.version} + test + + + org.mockito + mockito-core + ${mockito.version} test diff --git a/src/main/java/org/graylog/plugins/map/geoip/MaxmindDataAdapter.java b/src/main/java/org/graylog/plugins/map/geoip/MaxmindDataAdapter.java index 2f29b38e159d..97c533646d95 100644 --- a/src/main/java/org/graylog/plugins/map/geoip/MaxmindDataAdapter.java +++ b/src/main/java/org/graylog/plugins/map/geoip/MaxmindDataAdapter.java @@ -6,8 +6,8 @@ import com.fasterxml.jackson.annotation.JsonTypeName; import com.fasterxml.jackson.databind.annotation.JsonDeserialize; import com.google.auto.value.AutoValue; +import com.google.common.annotations.VisibleForTesting; import com.google.common.collect.ArrayListMultimap; -import com.google.common.collect.ImmutableMap; import com.google.common.collect.Multimap; import com.google.common.net.InetAddresses; import com.google.inject.assistedinject.Assisted; @@ -15,6 +15,8 @@ import com.maxmind.geoip2.exception.AddressNotFoundException; import com.maxmind.geoip2.model.CityResponse; import com.maxmind.geoip2.model.CountryResponse; +import com.maxmind.geoip2.record.Country; +import com.maxmind.geoip2.record.Location; import org.graylog.autovalue.WithBeanGetter; import org.graylog.plugins.map.config.DatabaseType; import org.graylog2.plugin.lookup.LookupCachePurge; @@ -36,6 +38,8 @@ import java.nio.file.Files; import java.nio.file.Path; import java.nio.file.Paths; +import java.util.HashMap; +import java.util.Map; import java.util.Optional; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicReference; @@ -130,7 +134,7 @@ private DatabaseReader loadReader(File file) throws IOException { @Override protected LookupResult doGet(Object key) { - InetAddress addr; + final InetAddress addr; if (key instanceof InetAddress) { addr = (InetAddress) key; } else { @@ -147,39 +151,60 @@ protected LookupResult doGet(Object key) { case MAXMIND_CITY: try { final CityResponse city = reader.city(addr); - final String singleValue = String.join(",", city.getLocation().getLatitude().toString(), city.getLocation().getLongitude().toString()); - final ImmutableMap.Builder map = ImmutableMap.builder(); + if (city == null) { + LOG.debug("No city data for IP address {}, returning empty result.", addr); + return LookupResult.empty(); + } + + final Location location = city.getLocation(); + final Map map = new HashMap<>(); map.put("city", city.getCity()); map.put("continent", city.getContinent()); map.put("country", city.getCountry()); - map.put("location", city.getLocation()); + map.put("location", location); map.put("postal", city.getPostal()); map.put("registered_country", city.getRegisteredCountry()); map.put("represented_country", city.getRepresentedCountry()); map.put("subdivisions", city.getSubdivisions()); map.put("traits", city.getTraits()); - return LookupResult.multi(singleValue, map.build()); + + final String singleValue; + if (location == null || location.getLatitude() == null || location.getLongitude() == null) { + singleValue = null; + } else { + singleValue = location.getLatitude() + "," + location.getLongitude(); + } + return LookupResult.multi(singleValue, map); } catch (AddressNotFoundException nfe) { + LOG.debug("Unable to look up city data for IP address {}, returning empty result.", addr, nfe); return LookupResult.empty(); } catch (Exception e) { - LOG.warn("Unable to look up IP address, returning empty result.", e); + LOG.warn("Unable to look up city data for IP address {}, returning empty result.", addr, e); return LookupResult.empty(); } case MAXMIND_COUNTRY: try { - final CountryResponse country = reader.country(addr); - final String singleValue = country.getCountry().getIsoCode(); - final ImmutableMap.Builder map = ImmutableMap.builder(); - map.put("continent", country.getContinent()); - map.put("country", country.getCountry()); - map.put("registered_country", country.getRegisteredCountry()); - map.put("represented_country", country.getRepresentedCountry()); - map.put("traits", country.getTraits()); - return LookupResult.multi(singleValue, map.build()); + final CountryResponse countryResponse = reader.country(addr); + if (countryResponse == null) { + LOG.debug("No country data for IP address {}, returning empty result.", addr); + return LookupResult.empty(); + } + + final Country country = countryResponse.getCountry(); + final Map map = new HashMap<>(); + map.put("continent", countryResponse.getContinent()); + map.put("country", country); + map.put("registered_country", countryResponse.getRegisteredCountry()); + map.put("represented_country", countryResponse.getRepresentedCountry()); + map.put("traits", countryResponse.getTraits()); + + final String singleValue = country == null ? null : country.getIsoCode(); + return LookupResult.multi(singleValue, map); } catch (AddressNotFoundException nfe) { + LOG.debug("Unable to look up country data for IP address {}, returning empty result.", addr, nfe); return LookupResult.empty(); } catch (Exception e) { - LOG.warn("Unable to look up IP address, returning empty result.", e); + LOG.warn("Unable to look up country data for IP address {}, returning empty result.", addr, e); return LookupResult.empty(); } } @@ -192,6 +217,16 @@ public void set(Object key, Object value) { throw new UnsupportedOperationException(); } + @VisibleForTesting + void setDatabaseReader(DatabaseReader databaseReader) { + this.databaseReader.set(databaseReader); + } + + @VisibleForTesting + DatabaseReader getDatabaseReader() { + return databaseReader.get(); + } + public interface Factory extends LookupDataAdapter.Factory { @Override MaxmindDataAdapter create(@Assisted("id") String id, diff --git a/src/test/java/org/graylog/plugins/map/geoip/GeoIpResolverEngineTest.java b/src/test/java/org/graylog/plugins/map/geoip/GeoIpResolverEngineTest.java index 75ed924cde43..139a42645c79 100644 --- a/src/test/java/org/graylog/plugins/map/geoip/GeoIpResolverEngineTest.java +++ b/src/test/java/org/graylog/plugins/map/geoip/GeoIpResolverEngineTest.java @@ -23,61 +23,53 @@ import com.google.common.net.InetAddresses; import org.graylog.plugins.map.config.GeoIpResolverConfig; import org.graylog2.plugin.Message; -import org.testng.annotations.AfterMethod; -import org.testng.annotations.BeforeMethod; -import org.testng.annotations.Test; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; import java.net.URISyntaxException; import java.util.Map; import static com.codahale.metrics.MetricRegistry.name; -import static org.testng.Assert.assertEquals; -import static org.testng.Assert.assertFalse; -import static org.testng.Assert.assertNotNull; -import static org.testng.Assert.assertNull; -import static org.testng.Assert.assertTrue; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertTrue; public class GeoIpResolverEngineTest { private MetricRegistry metricRegistry; private GeoIpResolverConfig config; - @BeforeMethod - public void setUp() { + @Before + public void setUp() throws URISyntaxException { config = GeoIpResolverConfig.defaultConfig().toBuilder().enabled(true).dbPath(this.getTestDatabasePath()).build(); metricRegistry = new MetricRegistry(); } - @AfterMethod + @After public void tearDown() { metricRegistry.removeMatching(MetricFilter.ALL); metricRegistry = null; } - private String getTestDatabasePath() { - String path = ""; - - try { - path = this.getClass().getResource("/GeoLite2-City.mmdb").toURI().getPath(); - } catch (URISyntaxException e) { - System.err.println("Could not get test geo location database: " + e); - } - - return path; + private String getTestDatabasePath() throws URISyntaxException { + return this.getClass().getResource("/GeoLite2-City.mmdb").toURI().getPath(); } @Test - public void getIpFromFieldValue() throws Exception { + public void getIpFromFieldValue() { final GeoIpResolverEngine resolver = new GeoIpResolverEngine(config, metricRegistry); final String ip = "127.0.0.1"; - assertEquals(resolver.getIpFromFieldValue(ip), InetAddresses.forString(ip)); + assertEquals(InetAddresses.forString(ip), resolver.getIpFromFieldValue(ip)); assertNull(resolver.getIpFromFieldValue("Message from \"127.0.0.1\"")); assertNull(resolver.getIpFromFieldValue("Test message with no IP")); } @Test - public void trimFieldValueBeforeLookup() throws Exception { + public void trimFieldValueBeforeLookup() { final GeoIpResolverEngine resolver = new GeoIpResolverEngine(config, metricRegistry); final String ip = " 2001:4860:4860::8888\t\n"; @@ -85,17 +77,17 @@ public void trimFieldValueBeforeLookup() throws Exception { } @Test - public void extractGeoLocationInformation() throws Exception { + public void extractGeoLocationInformation() { final GeoIpResolverEngine resolver = new GeoIpResolverEngine(config, metricRegistry); - assertTrue(resolver.extractGeoLocationInformation("1.2.3.4").isPresent(), "Should extract geo location information from public addresses"); - assertFalse(resolver.extractGeoLocationInformation("192.168.0.1").isPresent(), "Should not extract geo location information from private addresses"); - assertFalse(resolver.extractGeoLocationInformation(42).isPresent(), "Should not extract geo location information numeric fields"); - assertTrue(resolver.extractGeoLocationInformation(InetAddresses.forString("1.2.3.4")).isPresent(), "Should extract geo location information IP address fields"); + assertTrue("Should extract geo location information from public addresses", resolver.extractGeoLocationInformation("1.2.3.4").isPresent()); + assertFalse("Should not extract geo location information from private addresses", resolver.extractGeoLocationInformation("192.168.0.1").isPresent()); + assertFalse("Should not extract geo location information numeric fields", resolver.extractGeoLocationInformation(42).isPresent()); + assertTrue("Should extract geo location information IP address fields", resolver.extractGeoLocationInformation(InetAddresses.forString("1.2.3.4")).isPresent()); } @Test - public void disabledFilterTest() throws Exception { + public void disabledFilterTest() { final GeoIpResolverEngine resolver = new GeoIpResolverEngine(config.toBuilder().enabled(false).build(), metricRegistry); final Map messageFields = Maps.newHashMap(); @@ -108,26 +100,25 @@ public void disabledFilterTest() throws Exception { final Message message = new Message(messageFields); final boolean filtered = resolver.filter(message); - assertFalse(filtered, "Message should not be filtered out"); - assertEquals(message.getFields().size(), messageFields.size(), "Filter should not add new message fields"); + assertFalse("Message should not be filtered out", filtered); + assertEquals("Filter should not add new message fields", message.getFields().size(), messageFields.size()); } private void assertFieldNotResolved(Message message, String fieldName, String errorMessage) { - assertNull(message.getField(fieldName + "_geolocation"), errorMessage + " coordinates in " + fieldName); - assertNull(message.getField(fieldName + "_country_code"), errorMessage + " country in " + fieldName); - assertNull(message.getField(fieldName + "_city_name"), errorMessage + " city in " + fieldName); + assertNull(errorMessage + " coordinates in " + fieldName, message.getField(fieldName + "_geolocation")); + assertNull(errorMessage + " country in " + fieldName, message.getField(fieldName + "_country_code")); + assertNull(errorMessage + " city in " + fieldName, message.getField(fieldName + "_city_name")); } private void assertFieldResolved(Message message, String fieldName, String errorMessage) { - assertNotNull(message.getField(fieldName + "_geolocation"), errorMessage + " coordinates in " + fieldName); - assertNotNull(message.getField(fieldName + "_country_code"), errorMessage + " country in " + fieldName); - assertNotNull(message.getField(fieldName + "_city_name"), errorMessage + " city in " + fieldName); - assertTrue(((String) message.getField(fieldName + "_geolocation")).contains(","), - "Location coordinates for " + fieldName + " should include a comma"); + assertNotNull(errorMessage + " coordinates in " + fieldName, message.getField(fieldName + "_geolocation")); + assertNotNull(errorMessage + " country in " + fieldName, message.getField(fieldName + "_country_code")); + assertNotNull(errorMessage + " city in " + fieldName, message.getField(fieldName + "_city_name")); + assertTrue("Location coordinates for " + fieldName + " should include a comma", ((String) message.getField(fieldName + "_geolocation")).contains(",")); } @Test - public void filterResolvesIpGeoLocation() throws Exception { + public void filterResolvesIpGeoLocation() { final GeoIpResolverEngine resolver = new GeoIpResolverEngine(config, metricRegistry); final Map messageFields = Maps.newHashMap(); @@ -141,8 +132,8 @@ public void filterResolvesIpGeoLocation() throws Exception { final Message message = new Message(messageFields); final boolean filtered = resolver.filter(message); - assertFalse(filtered, "Message should not be filtered out"); - assertEquals(metricRegistry.timer(name(GeoIpResolverEngine.class, "resolveTime")).getCount(), 3, "Should have looked up three IPs"); + assertFalse("Message should not be filtered out", filtered); + assertEquals("Should have looked up three IPs", metricRegistry.timer(name(GeoIpResolverEngine.class, "resolveTime")).getCount(), 3); assertFieldNotResolved(message, "source", "Should not have resolved private IP"); assertFieldNotResolved(message, "message", "Should have resolved public IP"); assertFieldNotResolved(message, "gl2_remote_ip", "Should not have resolved text with an IP"); diff --git a/src/test/java/org/graylog/plugins/map/geoip/MaxmindDataAdapterTest.java b/src/test/java/org/graylog/plugins/map/geoip/MaxmindDataAdapterTest.java new file mode 100644 index 000000000000..3c15339f2aeb --- /dev/null +++ b/src/test/java/org/graylog/plugins/map/geoip/MaxmindDataAdapterTest.java @@ -0,0 +1,175 @@ +package org.graylog.plugins.map.geoip; + +import com.codahale.metrics.MetricRegistry; +import com.google.common.collect.ImmutableMap; +import com.maxmind.geoip2.DatabaseReader; +import com.maxmind.geoip2.model.CityResponse; +import com.maxmind.geoip2.model.CountryResponse; +import com.maxmind.geoip2.record.City; +import org.graylog.plugins.map.config.DatabaseType; +import org.graylog2.plugin.lookup.LookupResult; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.Suite; + +import java.net.URISyntaxException; +import java.util.Map; +import java.util.concurrent.TimeUnit; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +@RunWith(Suite.class) +@Suite.SuiteClasses({ + MaxmindDataAdapterTest.CityDatabaseTest.class, + MaxmindDataAdapterTest.CountryDatabaseTest.class +}) +public class MaxmindDataAdapterTest { + private static final Map DB_PATH = ImmutableMap.of( + DatabaseType.MAXMIND_CITY, "/GeoLite2-City.mmdb", + DatabaseType.MAXMIND_COUNTRY, "/GeoLite2-Country.mmdb" + ); + + abstract static class Base { + MaxmindDataAdapter adapter; + private final DatabaseType databaseType; + + Base(DatabaseType databaseType) { + this.databaseType = databaseType; + } + + @Before + public void setUp() throws Exception { + final MaxmindDataAdapter.Config config = MaxmindDataAdapter.Config.builder() + .checkInterval(1L) + .checkIntervalUnit(TimeUnit.HOURS) + .dbType(databaseType) + .path(getTestDatabasePath(DB_PATH.get(databaseType))) + .type("test") + .build(); + adapter = new MaxmindDataAdapter("test", "test", config, new MetricRegistry()); + + adapter.doStart(); + } + + @After + public void tearDown() throws Exception { + adapter.doStop(); + } + + private String getTestDatabasePath(String name) throws URISyntaxException { + return this.getClass().getResource(name).toURI().getPath(); + } + + @Test + public void doGetSuccessfullyResolvesLoopBackToEmptyResult() { + final LookupResult lookupResult = adapter.doGet("127.0.0.1"); + assertThat(lookupResult.isEmpty()).isTrue(); + } + + @Test + public void doGetSuccessfullyResolvesRFC1918AddressToEmptyResult() { + final LookupResult lookupResult = adapter.doGet("192.168.23.42"); + assertThat(lookupResult.isEmpty()).isTrue(); + } + + @Test + public void doGetReturnsEmptyResultIfDatabaseReaderReturnsNull() { + final DatabaseReader mockDatabaseReader = mock(DatabaseReader.class); + final DatabaseReader oldDatabaseReader = adapter.getDatabaseReader(); + + try { + adapter.setDatabaseReader(mockDatabaseReader); + + final LookupResult lookupResult = adapter.doGet("127.0.0.1"); + assertThat(lookupResult.isEmpty()).isTrue(); + } finally { + adapter.setDatabaseReader(oldDatabaseReader); + } + } + + @Test + public void doGetReturnsEmptyResultForInvalidIPAddress() { + final LookupResult lookupResult = adapter.doGet("Foobar"); + assertThat(lookupResult.isEmpty()).isTrue(); + } + } + + public static class CityDatabaseTest extends Base { + public CityDatabaseTest() { + super(DatabaseType.MAXMIND_CITY); + } + + @Test + public void doGetSuccessfullyResolvesGooglePublicDNSAddress() { + // This test will possibly get flaky when the entry for 8.8.8.8 changes! + final LookupResult lookupResult = adapter.doGet("8.8.8.8"); + assertThat(lookupResult.isEmpty()).isFalse(); + assertThat(lookupResult.multiValue()) + .extracting("city") + .extracting("geoNameId") + .containsExactly(5375480); + assertThat(lookupResult.multiValue()) + .extracting("location") + .extracting("metroCode") + .containsExactly(807); + } + + @Test + public void doGetReturnsResultIfCityResponseFieldsAreNull() throws Exception { + final CityResponse cityResponse = new CityResponse(null, null, null, null, null, null, null, null, null, null); + final DatabaseReader mockDatabaseReader = mock(DatabaseReader.class); + when(mockDatabaseReader.city(any())).thenReturn(cityResponse); + final DatabaseReader oldDatabaseReader = adapter.getDatabaseReader(); + + try { + adapter.setDatabaseReader(mockDatabaseReader); + + final LookupResult lookupResult = adapter.doGet("127.0.0.1"); + assertThat(lookupResult.isEmpty()).isFalse(); + assertThat(lookupResult.singleValue()).isNull(); + } finally { + adapter.setDatabaseReader(oldDatabaseReader); + } + } + } + + public static class CountryDatabaseTest extends Base { + public CountryDatabaseTest() { + super(DatabaseType.MAXMIND_COUNTRY); + } + + @Test + public void doGetSuccessfullyResolvesGooglePublicDNSAddress() { + // This test will possibly get flaky when the entry for 8.8.8.8 changes! + final LookupResult lookupResult = adapter.doGet("8.8.8.8"); + assertThat(lookupResult.isEmpty()).isFalse(); + assertThat(lookupResult.multiValue()) + .extracting("country") + .extracting("geoNameId") + .containsExactly(6252001); + } + + @Test + public void doGetReturnsResultIfCountryResponseFieldsAreNull() throws Exception { + final CountryResponse countryResponse = new CountryResponse(null, null, null, null, null, null); + final DatabaseReader mockDatabaseReader = mock(DatabaseReader.class); + when(mockDatabaseReader.country(any())).thenReturn(countryResponse); + final DatabaseReader oldDatabaseReader = adapter.getDatabaseReader(); + + try { + adapter.setDatabaseReader(mockDatabaseReader); + + final LookupResult lookupResult = adapter.doGet("127.0.0.1"); + assertThat(lookupResult.isEmpty()).isFalse(); + assertThat(lookupResult.singleValue()).isNull(); + } finally { + adapter.setDatabaseReader(oldDatabaseReader); + } + } + } +} \ No newline at end of file diff --git a/src/test/resources/GeoLite2-Country.mmdb b/src/test/resources/GeoLite2-Country.mmdb new file mode 100644 index 000000000000..56e7fdc043b0 Binary files /dev/null and b/src/test/resources/GeoLite2-Country.mmdb differ