From f33446ee253f3da42c1ba1f9ea4ae0228e5cc110 Mon Sep 17 00:00:00 2001
From: games647 <games647@users.noreply.github.com>
Date: Sat, 17 Mar 2018 03:00:24 +0100
Subject: [PATCH] Migrate to GEO IP 2 (Related #1471) (#1529)

* Migrate to GEO IP 2, because support will drop in April

* Change all links of maxmind to https
* Update maxmind database dependency and add javatar to extract
the database from the tar archive
(now only a small difference in jar file size -> ~80KB smaller)
* Verify downloaded archive using MD5 (There are no other checksums available)
* Migrate to Java NIO instead of old java file I/O (Feedback?)
* Internal Optional usage for nullable values (Feedback?)

Minor:
* Schedule a Bukkit async task instead of creating a thread manually
* Validate ip input string
* Extract validation into single method
* Close all resources safely using try-resources

* More https links

* Add documentation

* Set the same last modification as in the tar archive

* Fix tests

* Comment how the legacy API responded to unknown entries

* Document missing function param

* Document our maxmind dependency modifications

* Include time unit into constant

* More logging for downloading the database

* Add missing return if the database cannot be found

* Delete temporarily file after working with it
---
 README.md                                     |   4 +-
 pom.xml                                       |  28 +-
 src/main/java/fr/xephi/authme/AuthMe.java     |   7 +-
 .../fr/xephi/authme/service/GeoIpService.java | 267 +++++++++++++-----
 .../properties/ProtectionSettings.java        |   2 +-
 .../authme/service/GeoIpServiceTest.java      |  61 ++--
 6 files changed, 262 insertions(+), 107 deletions(-)

diff --git a/README.md b/README.md
index 255f36cf9..24514408c 100644
--- a/README.md
+++ b/README.md
@@ -27,7 +27,7 @@ You can also create your own translation file and, if you want, you can share it
 <ul>
   <li><strong>E-Mail Recovery System !!!</strong></li>
   <li>Username spoofing protection.</li>
-  <li>Countries Whitelist/Blacklist! <a href="http://dev.maxmind.com/geoip/legacy/codes/iso3166/">(country codes)</a></li>
+  <li>Countries Whitelist/Blacklist! <a href="https://dev.maxmind.com/geoip/legacy/codes/iso3166/">(country codes)</a></li>
   <li><strong>Built-in AntiBot System!</strong></li>
   <li><strong>ForceLogin Feature: Admins can login with all account via console command!</strong></li>
   <li><strong>Avoid the "Logged in from another location" message!</strong></li>
@@ -126,4 +126,4 @@ Credits for the old version of the plugin: d4rkwarriors, fabe1337, Whoami2 and p
 Thanks also to: AS1LV3RN1NJA, Hoeze and eprimex
 
 ##### GeoIP License:
-This product uses data from the GeoLite API created by MaxMind, available at http://www.maxmind.com
+This product uses data from the GeoLite API created by MaxMind, available at https://www.maxmind.com
diff --git a/pom.xml b/pom.xml
index b518d97af..6f5b90609 100644
--- a/pom.xml
+++ b/pom.xml
@@ -267,8 +267,12 @@
                             <shadedPattern>fr.xephi.authme.libs.slf4j.slf4j</shadedPattern>
                         </relocation>
                         <relocation>
-                            <pattern>com.maxmind.geoip</pattern>
-                            <shadedPattern>fr.xephi.authme.libs.maxmind.geoip</shadedPattern>
+                            <pattern>com.maxmind.db</pattern>
+                            <shadedPattern>fr.xephi.authme.libs.maxmind</shadedPattern>
+                        </relocation>
+                        <relocation>
+                            <pattern>com.ice.tar</pattern>
+                            <shadedPattern>fr.xephi.authme.libs.tar</shadedPattern>
                         </relocation>
                         <relocation>
                             <pattern>net.ricecode.similarity</pattern>
@@ -399,13 +403,19 @@
             <optional>true</optional>
         </dependency>
 
-        <!-- Maxmind GeoIp API -->
+        <!-- MaxMind GEO IP with our modifications to use GSON in replacement of the big Jackson dependency -->
+        <!-- GSON is already included and therefore it reduces the file size in comparison to the original version -->
         <dependency>
-            <groupId>com.maxmind.geoip</groupId>
-            <artifactId>geoip-api</artifactId>
-            <version>1.3.1</version>
-            <scope>compile</scope>
-            <optional>true</optional>
+            <groupId>com.maxmind.db</groupId>
+            <artifactId>maxmind-db-gson</artifactId>
+            <version>2.0.2-SNAPSHOT</version>
+        </dependency>
+
+        <!-- Library for tar archives -->
+        <dependency>
+            <groupId>javatar</groupId>
+            <artifactId>javatar</artifactId>
+            <version>2.5</version>
         </dependency>
 
         <!-- Java Email Library -->
@@ -526,7 +536,7 @@
         <dependency>
             <groupId>com.comphenix.protocol</groupId>
             <artifactId>ProtocolLib-API</artifactId>
-            <version>4.4.0-SNAPSHOT</version>
+            <version>4.3.0</version>
             <scope>provided</scope>
             <exclusions>
                 <exclusion>
diff --git a/src/main/java/fr/xephi/authme/AuthMe.java b/src/main/java/fr/xephi/authme/AuthMe.java
index d547f3ae0..305aadc38 100644
--- a/src/main/java/fr/xephi/authme/AuthMe.java
+++ b/src/main/java/fr/xephi/authme/AuthMe.java
@@ -2,7 +2,9 @@ package fr.xephi.authme;
 
 import ch.jalu.injector.Injector;
 import ch.jalu.injector.InjectorBuilder;
+
 import com.google.common.annotations.VisibleForTesting;
+
 import fr.xephi.authme.api.NewAPI;
 import fr.xephi.authme.command.CommandHandler;
 import fr.xephi.authme.datasource.DataSource;
@@ -33,6 +35,9 @@ import fr.xephi.authme.settings.properties.SecuritySettings;
 import fr.xephi.authme.task.CleanupTask;
 import fr.xephi.authme.task.purge.PurgeService;
 import fr.xephi.authme.util.ExceptionUtils;
+
+import java.io.File;
+
 import org.apache.commons.lang.SystemUtils;
 import org.bukkit.Server;
 import org.bukkit.command.Command;
@@ -43,8 +48,6 @@ import org.bukkit.plugin.java.JavaPlugin;
 import org.bukkit.plugin.java.JavaPluginLoader;
 import org.bukkit.scheduler.BukkitScheduler;
 
-import java.io.File;
-
 import static fr.xephi.authme.service.BukkitService.TICKS_PER_MINUTE;
 import static fr.xephi.authme.util.Utils.isClassLoaded;
 
diff --git a/src/main/java/fr/xephi/authme/service/GeoIpService.java b/src/main/java/fr/xephi/authme/service/GeoIpService.java
index 9c43ee177..b973b6337 100644
--- a/src/main/java/fr/xephi/authme/service/GeoIpService.java
+++ b/src/main/java/fr/xephi/authme/service/GeoIpService.java
@@ -1,46 +1,81 @@
 package fr.xephi.authme.service;
 
 import com.google.common.annotations.VisibleForTesting;
-import com.maxmind.geoip.LookupService;
+import com.google.common.hash.HashCode;
+import com.google.common.hash.HashFunction;
+import com.google.common.hash.Hashing;
+import com.google.common.io.Resources;
+import com.ice.tar.TarEntry;
+import com.ice.tar.TarInputStream;
+import com.maxmind.db.GeoIp2Provider;
+import com.maxmind.db.Reader;
+import com.maxmind.db.Reader.FileMode;
+import com.maxmind.db.cache.CHMCache;
+import com.maxmind.db.model.Country;
+import com.maxmind.db.model.CountryResponse;
+
 import fr.xephi.authme.ConsoleLogger;
 import fr.xephi.authme.initialization.DataFolder;
 import fr.xephi.authme.util.FileUtils;
 import fr.xephi.authme.util.InternetProtocolUtils;
 
-import javax.inject.Inject;
+import java.io.BufferedInputStream;
 import java.io.File;
-import java.io.FileOutputStream;
 import java.io.IOException;
-import java.io.InputStream;
 import java.io.OutputStream;
+import java.net.InetAddress;
 import java.net.URL;
-import java.net.URLConnection;
-import java.util.concurrent.TimeUnit;
+import java.net.UnknownHostException;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.StandardCopyOption;
+import java.nio.file.attribute.FileTime;
+import java.time.Duration;
+import java.time.Instant;
+import java.util.Objects;
+import java.util.Optional;
 import java.util.zip.GZIPInputStream;
 
-import static com.maxmind.geoip.LookupService.GEOIP_MEMORY_CACHE;
+import javax.inject.Inject;
 
 public class GeoIpService {
-    private static final String LICENSE =
-        "[LICENSE] This product uses data from the GeoLite API created by MaxMind, available at http://www.maxmind.com";
-    private static final String GEOIP_URL =
-        "http://geolite.maxmind.com/download/geoip/database/GeoLiteCountry/GeoIP.dat.gz";
-    private LookupService lookupService;
-    private Thread downloadTask;
 
-    private final File dataFile;
+    private static final String LICENSE =
+            "[LICENSE] This product includes GeoLite2 data created by MaxMind, available at https://www.maxmind.com";
+
+    private static final String DATABASE_NAME = "GeoLite2-Country";
+    private static final String DATABASE_EXT = ".mmdb";
+    private static final String DATABASE_FILE = DATABASE_NAME + DATABASE_EXT;
+
+    private static final String ARCHIVE_FILE = DATABASE_NAME + ".tar.gz";
+
+    private static final String ARCHIVE_URL = "https://geolite.maxmind.com/download/geoip/database/" + ARCHIVE_FILE;
+    private static final String CHECKSUM_URL = ARCHIVE_URL + ".md5";
+
+    private static final int UPDATE_INTERVAL_DAYS = 30;
+
+    private final Path dataFile;
+    private final BukkitService bukkitService;
+
+    private GeoIp2Provider databaseReader;
+    private volatile boolean downloading;
 
     @Inject
-    GeoIpService(@DataFolder File dataFolder) {
-        this.dataFile = new File(dataFolder, "GeoIP.dat");
+    GeoIpService(@DataFolder File dataFolder, BukkitService bukkitService) {
+        this.bukkitService = bukkitService;
+        this.dataFile = dataFolder.toPath().resolve(DATABASE_FILE);
+
         // Fires download of recent data or the initialization of the look up service
         isDataAvailable();
     }
 
     @VisibleForTesting
-    GeoIpService(@DataFolder File dataFolder, LookupService lookupService) {
-        this.dataFile = dataFolder;
-        this.lookupService = lookupService;
+    GeoIpService(@DataFolder File dataFolder, BukkitService bukkitService, GeoIp2Provider reader) {
+        this.bukkitService = bukkitService;
+        this.dataFile = dataFolder.toPath().resolve(DATABASE_FILE);
+
+        this.databaseReader = reader;
     }
 
     /**
@@ -49,94 +84,186 @@ public class GeoIpService {
      * @return True if the data is available, false otherwise.
      */
     private synchronized boolean isDataAvailable() {
-        if (downloadTask != null && downloadTask.isAlive()) {
+        if (downloading) {
+            // we are currently downloading the database
             return false;
         }
-        if (lookupService != null) {
+
+        if (databaseReader != null) {
+            // everything is initialized
             return true;
         }
 
-        if (dataFile.exists()) {
-            boolean dataIsOld = (System.currentTimeMillis() - dataFile.lastModified()) > TimeUnit.DAYS.toMillis(30);
-            if (!dataIsOld) {
-                try {
-                    lookupService = new LookupService(dataFile, GEOIP_MEMORY_CACHE);
+        if (Files.exists(dataFile)) {
+            try {
+                FileTime lastModifiedTime = Files.getLastModifiedTime(dataFile);
+                if (Duration.between(lastModifiedTime.toInstant(), Instant.now()).toDays() <= UPDATE_INTERVAL_DAYS) {
+                    databaseReader = new Reader(dataFile.toFile(), FileMode.MEMORY, new CHMCache());
                     ConsoleLogger.info(LICENSE);
+
+                    // don't fire the update task - we are up to date
                     return true;
-                } catch (IOException e) {
-                    ConsoleLogger.logException("Failed to load GeoLiteAPI database", e);
-                    return false;
+                } else {
+                    ConsoleLogger.debug("GEO Ip database is older than " + UPDATE_INTERVAL_DAYS + " Days");
                 }
-            } else {
-                FileUtils.delete(dataFile);
+            } catch (IOException ioEx) {
+                ConsoleLogger.logException("Failed to load GeoLiteAPI database", ioEx);
+                return false;
             }
         }
-        // Ok, let's try to download the data file!
-        downloadTask = createDownloadTask();
-        downloadTask.start();
+
+        // File is outdated or doesn't exist - let's try to download the data file!
+        startDownloadTask();
         return false;
     }
 
     /**
      * Create a thread which will attempt to download new data from the GeoLite website.
-     *
-     * @return the generated download thread
      */
-    private Thread createDownloadTask() {
-        return new Thread(new Runnable() {
-            @Override
-            public void run() {
-                try {
-                    URL downloadUrl = new URL(GEOIP_URL);
-                    URLConnection conn = downloadUrl.openConnection();
-                    conn.setConnectTimeout(10000);
-                    conn.connect();
-                    InputStream input = conn.getInputStream();
-                    if (conn.getURL().toString().endsWith(".gz")) {
-                        input = new GZIPInputStream(input);
-                    }
-                    OutputStream output = new FileOutputStream(dataFile);
-                    byte[] buffer = new byte[2048];
-                    int length = input.read(buffer);
-                    while (length >= 0) {
-                        output.write(buffer, 0, length);
-                        length = input.read(buffer);
-                    }
-                    output.close();
-                    input.close();
-                } catch (IOException e) {
-                    ConsoleLogger.logException("Could not download GeoLiteAPI database", e);
+    private void startDownloadTask() {
+        downloading = true;
+
+        // use bukkit's cached threads
+        bukkitService.runTaskAsynchronously(() -> {
+            ConsoleLogger.info("Downloading GEO IP database, because the old database is outdated or doesn't exist");
+
+            Path tempFile = null;
+            try {
+                // download database to temporarily location
+                tempFile = Files.createTempFile(ARCHIVE_FILE, null);
+                try (OutputStream out = Files.newOutputStream(tempFile)) {
+                    Resources.copy(new URL(ARCHIVE_URL), out);
+                }
+
+                // MD5 checksum verification
+                String targetChecksum = Resources.toString(new URL(CHECKSUM_URL), StandardCharsets.UTF_8);
+                if (!verifyChecksum(Hashing.md5(), tempFile, targetChecksum)) {
+                    return;
+                }
+
+                // tar extract database and copy to target destination
+                if (!extractDatabase(tempFile, dataFile)) {
+                    ConsoleLogger.warning("Cannot find database inside downloaded GEO IP file at " + tempFile);
+                    return;
+                }
+
+                ConsoleLogger.info("Successfully downloaded new GEO IP database to " + dataFile);
+
+                //only set this value to false on success otherwise errors could lead to endless download triggers
+                downloading = false;
+            } catch (IOException ioEx) {
+                ConsoleLogger.logException("Could not download GeoLiteAPI database", ioEx);
+            } finally {
+                // clean up
+                if (tempFile != null) {
+                    FileUtils.delete(tempFile.toFile());
                 }
             }
         });
     }
 
+    /**
+     * Verify if the expected checksum is equal to the checksum of the given file.
+     *
+     * @param function the checksum function like MD5, SHA256 used to generate the checksum from the file
+     * @param file the file we want to calculate the checksum from
+     * @param expectedChecksum the expected checksum
+     * @return true if equal, false otherwise
+     * @throws IOException on I/O error reading the file
+     */
+    private boolean verifyChecksum(HashFunction function, Path file, String expectedChecksum) throws IOException {
+        HashCode actualHash = function.hashBytes(Files.readAllBytes(file));
+        HashCode expectedHash = HashCode.fromString(expectedChecksum);
+        if (Objects.equals(actualHash, expectedHash)) {
+            return true;
+        }
+
+        ConsoleLogger.warning("GEO IP checksum verification failed");
+        ConsoleLogger.warning("Expected: " + expectedHash + " Actual: " + actualHash);
+        return false;
+    }
+
+    /**
+     * Extract the database from the tar archive. Existing outputFile will be replaced if it already exists.
+     *
+     * @param tarInputFile gzipped tar input file where the database is
+     * @param outputFile destination file for the database
+     * @return true if the database was found, false otherwise
+     * @throws IOException on I/O error reading the tar archive or writing the output
+     */
+    private boolean extractDatabase(Path tarInputFile, Path outputFile) throws IOException {
+        // .gz -> gzipped file
+        try (BufferedInputStream in = new BufferedInputStream(Files.newInputStream(tarInputFile));
+             TarInputStream tarIn = new TarInputStream(new GZIPInputStream(in))) {
+            TarEntry entry;
+            while ((entry = tarIn.getNextEntry()) != null) {
+                if (!entry.isDirectory()) {
+                    // filename including folders (absolute path inside the archive)
+                    String filename = entry.getName();
+                    if (filename.endsWith(DATABASE_EXT)) {
+                        // found the database file
+                        Files.copy(tarIn, outputFile, StandardCopyOption.REPLACE_EXISTING);
+
+                        // update the last modification date to be same as in the archive
+                        Files.setLastModifiedTime(outputFile, FileTime.from(entry.getModTime().toInstant()));
+                        return true;
+                    }
+                }
+            }
+        }
+
+        return false;
+    }
+
     /**
      * Get the country code of the given IP address.
      *
      * @param ip textual IP address to lookup.
-     *
      * @return two-character ISO 3166-1 alpha code for the country.
      */
     public String getCountryCode(String ip) {
-        if (!InternetProtocolUtils.isLocalAddress(ip) && isDataAvailable()) {
-            return lookupService.getCountry(ip).getCode();
-        }
-        return "--";
+        return getCountry(ip).map(Country::getIsoCode).orElse("--");
     }
 
     /**
      * Get the country name of the given IP address.
      *
      * @param ip textual IP address to lookup.
-     *
      * @return The name of the country.
      */
     public String getCountryName(String ip) {
-        if (!InternetProtocolUtils.isLocalAddress(ip) && isDataAvailable()) {
-            return lookupService.getCountry(ip).getName();
-        }
-        return "N/A";
+        return getCountry(ip).map(Country::getName).orElse("N/A");
     }
 
+    /**
+     * Get the country of the given IP address
+     *
+     * @param ip textual IP address to lookup
+     * @return the wrapped Country model or {@link Optional#empty()} if
+     * <ul>
+     *     <li>Database reader isn't initialized</li>
+     *     <li>MaxMind has no record about this IP address</li>
+     *     <li>IP address is local</li>
+     *     <li>Textual representation is not a valid IP address</li>
+     * </ul>
+     */
+    private Optional<Country> getCountry(String ip) {
+        if (ip == null || ip.isEmpty() || InternetProtocolUtils.isLocalAddress(ip) || !isDataAvailable()) {
+            return Optional.empty();
+        }
+
+        try {
+            InetAddress address = InetAddress.getByName(ip);
+
+            //Reader.getCountry() can be null for unknown addresses
+            return Optional.ofNullable(databaseReader.getCountry(address)).map(CountryResponse::getCountry);
+        } catch (UnknownHostException e) {
+            // Ignore invalid ip addresses
+            // Legacy GEO IP Database returned a unknown country object with Country-Code: '--' and Country-Name: 'N/A'
+        } catch (IOException ioEx) {
+            ConsoleLogger.logException("Cannot lookup country for " + ip + " at GEO IP database", ioEx);
+        }
+
+        return Optional.empty();
+    }
 }
diff --git a/src/main/java/fr/xephi/authme/settings/properties/ProtectionSettings.java b/src/main/java/fr/xephi/authme/settings/properties/ProtectionSettings.java
index f0a6c2f74..469d983af 100644
--- a/src/main/java/fr/xephi/authme/settings/properties/ProtectionSettings.java
+++ b/src/main/java/fr/xephi/authme/settings/properties/ProtectionSettings.java
@@ -22,7 +22,7 @@ public final class ProtectionSettings implements SettingsHolder {
 
     @Comment({
         "Countries allowed to join the server and register. For country codes, see",
-        "http://dev.maxmind.com/geoip/legacy/codes/iso3166/",
+        "https://dev.maxmind.com/geoip/legacy/codes/iso3166/",
         "PLEASE USE QUOTES!"})
     public static final Property<List<String>> COUNTRIES_WHITELIST =
         newListProperty("Protection.countries", "US", "GB");
diff --git a/src/test/java/fr/xephi/authme/service/GeoIpServiceTest.java b/src/test/java/fr/xephi/authme/service/GeoIpServiceTest.java
index 60e3aa680..47dfb54e4 100644
--- a/src/test/java/fr/xephi/authme/service/GeoIpServiceTest.java
+++ b/src/test/java/fr/xephi/authme/service/GeoIpServiceTest.java
@@ -1,7 +1,13 @@
 package fr.xephi.authme.service;
 
-import com.maxmind.geoip.Country;
-import com.maxmind.geoip.LookupService;
+import com.maxmind.db.GeoIp2Provider;
+import com.maxmind.db.model.Country;
+import com.maxmind.db.model.CountryResponse;
+
+import java.io.File;
+import java.io.IOException;
+import java.net.InetAddress;
+
 import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
@@ -10,13 +16,11 @@ import org.junit.runner.RunWith;
 import org.mockito.Mock;
 import org.mockito.junit.MockitoJUnitRunner;
 
-import java.io.File;
-import java.io.IOException;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.BDDMockito.given;
 
 import static org.hamcrest.Matchers.equalTo;
 import static org.junit.Assert.assertThat;
-import static org.mockito.ArgumentMatchers.anyString;
-import static org.mockito.BDDMockito.given;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.verify;
@@ -29,8 +33,12 @@ public class GeoIpServiceTest {
 
     private GeoIpService geoIpService;
     private File dataFolder;
+
     @Mock
-    private LookupService lookupService;
+    private GeoIp2Provider lookupService;
+
+    @Mock
+    private BukkitService bukkitService;
 
     @Rule
     public TemporaryFolder temporaryFolder = new TemporaryFolder();
@@ -38,20 +46,24 @@ public class GeoIpServiceTest {
     @Before
     public void initializeGeoLiteApi() throws IOException {
         dataFolder = temporaryFolder.newFolder();
-        geoIpService = new GeoIpService(dataFolder, lookupService);
+        geoIpService = new GeoIpService(dataFolder, bukkitService, lookupService);
     }
 
     @Test
-    public void shouldGetCountry() {
+    public void shouldGetCountry() throws Exception {
         // given
-        String ip = "123.45.67.89";
+        InetAddress ip = InetAddress.getByName("123.45.67.89");
         String countryCode = "XX";
+
         Country country = mock(Country.class);
-        given(country.getCode()).willReturn(countryCode);
-        given(lookupService.getCountry(ip)).willReturn(country);
+        given(country.getIsoCode()).willReturn(countryCode);
+
+        CountryResponse response = mock(CountryResponse.class);
+        given(response.getCountry()).willReturn(country);
+        given(lookupService.getCountry(ip)).willReturn(response);
 
         // when
-        String result = geoIpService.getCountryCode(ip);
+        String result = geoIpService.getCountryCode(ip.getHostAddress());
 
         // then
         assertThat(result, equalTo(countryCode));
@@ -59,7 +71,7 @@ public class GeoIpServiceTest {
     }
 
     @Test
-    public void shouldNotLookUpCountryForLocalhostIp() {
+    public void shouldNotLookUpCountryForLocalhostIp() throws Exception  {
         // given
         String ip = "127.0.0.1";
 
@@ -68,20 +80,24 @@ public class GeoIpServiceTest {
 
         // then
         assertThat(result, equalTo("--"));
-        verify(lookupService, never()).getCountry(anyString());
+        verify(lookupService, never()).getCountry(any());
     }
 
     @Test
-    public void shouldLookUpCountryName() {
+    public void shouldLookUpCountryName() throws Exception {
         // given
-        String ip = "24.45.167.89";
+        InetAddress ip = InetAddress.getByName("24.45.167.89");
         String countryName = "Ecuador";
+
         Country country = mock(Country.class);
         given(country.getName()).willReturn(countryName);
-        given(lookupService.getCountry(ip)).willReturn(country);
+
+        CountryResponse response = mock(CountryResponse.class);
+        given(response.getCountry()).willReturn(country);
+        given(lookupService.getCountry(ip)).willReturn(response);
 
         // when
-        String result = geoIpService.getCountryName(ip);
+        String result = geoIpService.getCountryName(ip.getHostAddress());
 
         // then
         assertThat(result, equalTo(countryName));
@@ -89,16 +105,15 @@ public class GeoIpServiceTest {
     }
 
     @Test
-    public void shouldNotLookUpCountryNameForLocalhostIp() {
+    public void shouldNotLookUpCountryNameForLocalhostIp() throws Exception {
         // given
-        String ip = "127.0.0.1";
+        InetAddress ip = InetAddress.getByName("127.0.0.1");
 
         // when
-        String result = geoIpService.getCountryName(ip);
+        String result = geoIpService.getCountryName(ip.getHostAddress());
 
         // then
         assertThat(result, equalTo("N/A"));
         verify(lookupService, never()).getCountry(ip);
     }
-
 }