Merge pull request #195 from jaredj/missing-geoip-data

Handle missing GeoIP data gracefully

plugins/ident/geoip

@@ -341,13 +341,8 @@ sub init_my_country_code {
 
 sub set_country_code {
     my $self = shift;
-    my $remote_ip = $self->qp->connection->remote_ip;
+    my $ip   = $self->qp->connection->remote_ip;
-
+    my $code = $self->get_country_code($ip) or return;
-    my $code = $self->{_geoip_city}
-             ? $self->get_country_code_gc($remote_ip)
-             : $self->get_country_code($remote_ip);
-
-    return if ! $code;
     $self->qp->connection->notes('geoip_country', $code);
     return $code;
 }
@@ -357,8 +352,11 @@ sub get_country_code {
     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);
+    if ($self->{_geoip}) {
+        return $self->{_geoip}->country_code_by_addr($ip);
+    }
+    return undef;
 }
 
 sub get_country_code_gc {
@@ -371,17 +369,24 @@ sub get_country_code_gc {
 
 sub set_country_name {
     my $self = shift;
-    my $remote_ip = $self->qp->connection->remote_ip;
+    my $ip   = $self->qp->connection->remote_ip;
-
+    my $name = $self->get_country_name($ip) or return;
-    my $name = $self->{_geoip_city}
-             ? $self->get_country_name_gc($remote_ip)
-             : $self->{_geoip}->country_name_by_addr($remote_ip);
-
-    return if ! $name;
     $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;
+    if ($self->{_geoip_city}) {
+        return $self->get_country_name_gc($ip);
+    }
+    if ($self->{_geoip}) {
+        return $self->{_geoip}->country_name_by_addr($ip);
+    }
+    return undef;
+}
+
 sub get_country_name_gc {
     my $self = shift;
     return if !$self->{_geoip_record};
@@ -390,17 +395,24 @@ sub get_country_name_gc {
 
 sub set_continent {
     my ($self, $country_code) = @_;
-    $country_code or return;
+    return if !$country_code;
-
+    my $continent = $self->get_continent($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;
     $self->qp->connection->notes('geoip_continent', $continent);
     return $continent;
 }
 
+sub get_continent {
+    my ($self, $country_code) = @_;
+    return if !$country_code;
+    if ($self->{_geoip_city}) {
+        return $self->get_continent_gc();
+    }
+    if ($self->{_geoip}) {
+        return $self->{_geoip}->continent_code_by_country_code($country_code);
+    }
+    return undef;
+}
+
 sub get_continent_gc {
     my $self = shift;
     return if !$self->{_geoip_record};

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;
+}