From f8e220e3b5c458b8a83297dddfcf62b51485aa91 Mon Sep 17 00:00:00 2001
From: Jared Johnson <jjohnson@efolder.net>
Date: Wed, 14 Jan 2015 15:28:14 -0600
Subject: [PATCH] Handle missing GeoIP data gracefully

Previously, the GeoIP plugin would crash on certain missing GeoIP data.
Now it will continue to operate on whatever data it can get at.

Note that it already warns when any data it's looking for is missing.
---
 plugins/ident/geoip        | 55 +++++++++++++++++++-------------------
 t/plugin_tests/ident/geoip | 30 +++++++++++++++++++++
 2 files changed, 58 insertions(+), 27 deletions(-)

diff --git a/plugins/ident/geoip b/plugins/ident/geoip
index 53e81fe..8aceb36 100644
--- a/plugins/ident/geoip
+++ b/plugins/ident/geoip
@@ -341,29 +341,23 @@ sub init_my_country_code {
 
 sub set_country_code {
     my $self = shift;
-    my $remote_ip = $self->qp->connection->remote_ip;
-
-    my $code = $self->{_geoip_city}
-             ? $self->get_country_code_gc($remote_ip)
-             : $self->get_country_code($remote_ip);
-
-    return if ! $code;
+    my $ip   = $self->qp->connection->remote_ip;
+    my $code = $self->get_country_code($ip) or return;
     $self->qp->connection->notes('geoip_country', $code);
     return $code;
 }
 
 sub get_country_code {
     my $self = shift;
-    my $ip = shift || $self->qp->connection->remote_ip;
-    if ($self->{_geoip_city}) {
-        return $self->get_country_code_gc($ip);
-    };
-    return $self->{_geoip}->country_code_by_addr($ip);
+    my $ip   = shift || $self->qp->connection->remote_ip;
+    return $self->get_country_code_gc($ip)            if $self->{_geoip_city};
+    return $self->{_geoip}->country_code_by_addr($ip) if $self->{_geoip};
+    return undef;
 }
 
 sub get_country_code_gc {
     my $self = shift;
-    my $ip = shift || $self->qp->connection->remote_ip;
+    my $ip   = shift || $self->qp->connection->remote_ip;
     $self->{_geoip_record} = $self->{_geoip_city}->record_by_addr($ip)
       or return;
     return $self->{_geoip_record}->country_code;
@@ -371,17 +365,20 @@ sub get_country_code_gc {
 
 sub set_country_name {
     my $self = shift;
-    my $remote_ip = $self->qp->connection->remote_ip;
-
-    my $name = $self->{_geoip_city}
-             ? $self->get_country_name_gc($remote_ip)
-             : $self->{_geoip}->country_name_by_addr($remote_ip);
-
-    return if ! $name;
+    my $ip = $self->qp->connection->remote_ip;
+    my $name = $self->get_country_name($ip) or return;
     $self->qp->connection->notes('geoip_country_name', $name);
     return $name;
 }
 
+sub get_country_name {
+    my $self = shift;
+    my $ip   = shift || $self->qp->connection->remote_ip;
+    return $self->get_country_name_gc($ip)            if $self->{_geoip_city};
+    return $self->{_geoip}->country_name_by_addr($ip) if $self->{_geoip};
+    return undef;
+}
+
 sub get_country_name_gc {
     my $self = shift;
     return if !$self->{_geoip_record};
@@ -390,17 +387,21 @@ sub get_country_name_gc {
 
 sub set_continent {
     my ($self, $country_code) = @_;
-    $country_code or return;
-
-    my $continent = $self->{_geoip_city}
-                  ? $self->get_continent_gc()
-                  : $self->{_geoip}->continent_code_by_country_code($country_code);
-
-    $continent or return;
+    return if ! $country_code;
+    my $continent = $self->get_continent($country_code) or return;
     $self->qp->connection->notes('geoip_continent', $continent);
     return $continent;
 }
 
+sub get_continent {
+    my ( $self, $country_code ) = @_;
+    return if ! $country_code;
+    return $self->get_continent_gc() if $self->{_geoip_city};
+    return $self->{_geoip}->continent_code_by_country_code($country_code)
+        if $self->{_geoip};
+    return undef;
+}
+
 sub get_continent_gc {
     my $self = shift;
     return if !$self->{_geoip_record};
diff --git a/t/plugin_tests/ident/geoip b/t/plugin_tests/ident/geoip
index d28ed10..5436866 100644
--- a/t/plugin_tests/ident/geoip
+++ b/t/plugin_tests/ident/geoip
@@ -117,6 +117,10 @@ sub test_set_country_code {
     ok( ! $cc, "undef");
 
     $self->qp->connection->remote_ip('24.24.24.24');
+    $self->clear_geoip_data;
+    $cc = $self->set_country_code();
+    ok( ! $cc, "set_country_code() returns nothing for no geoip data");
+    $self->restore_geoip_data;
     $cc = $self->set_country_code();
     cmp_ok( $cc, 'eq', 'US', "set_country_code result is $cc");
 
@@ -134,6 +138,11 @@ sub test_set_country_name {
     ok( ! $cn, "undef") or warn "$cn\n";
 
     $self->qp->connection->remote_ip('24.24.24.24');
+    $self->clear_geoip_data;
+    $self->set_country_code();
+    $cn = $self->set_country_name();
+    ok( ! $cn, "set_country_name() returns nothing for no geoip data");
+    $self->restore_geoip_data;
     $self->set_country_code();
     $cn = $self->set_country_name();
     cmp_ok( $cn, 'eq', 'United States', "$cn");
@@ -152,6 +161,11 @@ sub test_set_continent {
     ok( ! $cn, "undef") or warn "$cn\n";
 
     $self->qp->connection->remote_ip('24.24.24.24');
+    $self->clear_geoip_data;
+    $self->set_country_code();
+    $cn = $self->set_continent('US');
+    ok( ! $cn, 'set_continent() returns nothing for no geoip data');
+    $self->restore_geoip_data;
     $self->set_country_code();
     $cn = $self->set_continent() || '';
     my $note = $self->connection->notes('geoip_continent');
@@ -200,6 +214,10 @@ sub test_set_asn {
     ok( ! $asn, "undef") or warn "$asn\n";
 
     $self->qp->connection->remote_ip('24.24.24.24');
+    $self->clear_geoip_data;
+    $asn = $self->set_asn();
+    ok( ! $asn, 'set_asn() returns nothing for no ASN data' );
+    $self->restore_geoip_data;
     $asn = $self->set_asn();
     ok( $self->connection->notes('geoip_asn') =~ /^11351/, "note has: $asn");
 
@@ -208,3 +226,15 @@ sub test_set_asn {
 
     ok( $self->connection->notes('geoip_asn') =~ /^7819/, "note has: $asn");
 }
+
+my $geoip_data_bak;
+my @geoip_keys = qw( _geoip _geoip_city GeoIPASNum );
+sub clear_geoip_data {
+    my ( $self ) = @_;
+    $geoip_data_bak->{$_} = delete $self->{$_} for @geoip_keys;
+}
+
+sub restore_geoip_data {
+    my ( $self ) = @_;
+    $self->{$_} = delete $geoip_data_bak->{$_} for @geoip_keys;
+}