Skip to content

Commit

Permalink
Improve null-handling in MaxmindDataAdapter (#68)
Browse files Browse the repository at this point in the history 
* Migrate from TestNg to JUnit

* Add GeoLite2 Country database for tests

Version: 20180102
Download URL: https://dev.maxmind.com/geoip/geoip2/geolite2/#Downloads

* Improve null-handling in MaxmindDataAdapter

Fix #67

* Use null-safe map (HashMap) in MaxmindDataAdapter#doGet()

* Properly handle absence of CityResponse.location
  • Loading branch information
@joschi @bernd
joschi authored and bernd committed Jan 24, 2018
1 parent 55ebba4 commit 80e50fa
Show file tree
Hide file tree
Showing 5 changed files with 276 additions and 63 deletions.
18 changes: 15 additions & 3 deletions pom.xml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -102,9 +102,21 @@
<version>2.8.1</version>
</dependency>
<dependency>
<groupId>org.testng</groupId>
<artifactId>testng</artifactId>
<version>6.9.10</version>
<groupId>junit</groupId>
<artifactId>junit</artifactId>
<version>${junit.version}</version>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.assertj</groupId>
<artifactId>assertj-core</artifactId>
<version>${assertj-core.version}</version>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.mockito</groupId>
<artifactId>mockito-core</artifactId>
<version>${mockito.version}</version>
<scope>test</scope>
</dependency>
<dependency>
Expand Down
69 changes: 52 additions & 17 deletions src/main/java/org/graylog/plugins/map/geoip/MaxmindDataAdapter.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,15 +6,17 @@
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;
import com.maxmind.geoip2.DatabaseReader;
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;
Expand All @@ -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;
Expand Down Expand Up @@ -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 {
Expand All @@ -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<Object, Object> 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<Object, Object> 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<Object, Object> 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<Object, Object> 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();
}
}
Expand All @@ -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<MaxmindDataAdapter> {
@Override
MaxmindDataAdapter create(@Assisted("id") String id,
Expand Down
77 changes: 34 additions & 43 deletions src/test/java/org/graylog/plugins/map/geoip/GeoIpResolverEngineTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,79 +23,71 @@
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";

assertNotNull(resolver.getIpFromFieldValue(ip));
}

@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<String, Object> messageFields = Maps.newHashMap();
Expand All @@ -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<String, Object> messageFields = Maps.newHashMap();
Expand All @@ -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");
Expand Down
Loading

0 comments on commit 80e50fa

Please sign in to comment.