Fix MMDB code to re-open explicitly opened DBs correctly

The filename from which a DB first gets opened (either via an explicitly
specified filename, or via the path sequence now configurable at the script
layer) is now "sticky", meaning re-opening won't switch to a different file.

This was easiest by moving most state into the MMDB class itself. The previous
approach of tracking the two DB instances via a smart pointer and blowing the
pointed-to objects away as needed is now instead one of two objects fixed over
the lifetime of Zeek, able to open/close/reopen their underlying Maxmind DBs.

The MMDB class now only has one Lookup() method since there was no need to break
them apart -- it saves the return of a MMDB_lookup_result_s over the stack and
there's no need for throwing an exception.

src/MMDB.cc

@@ -18,13 +18,11 @@ namespace zeek {
 
 static int msg_count = 0;
 static double msg_suppression_time = 0;
-static bool did_loc_db_error = false;
-static bool did_asn_db_error = false;
 static constexpr int msg_limit = 20;
 static constexpr double msg_suppression_duration = 300;
 
-static std::unique_ptr<MMDB> mmdb_loc;
+LocDB mmdb_loc;
-static std::unique_ptr<MMDB> mmdb_asn;
+AsnDB mmdb_asn;
 
 static void report_msg(const char* format, ...) {
     if ( zeek::run_state::network_time > msg_suppression_time + msg_suppression_duration ) {
@@ -45,135 +43,6 @@ static void report_msg(const char* format, ...) {
     zeek::reporter->Info("%s", msg.data());
 }
 
-MMDB::MMDB(const char* filename, struct stat info)
-    : file_info(info), lookup_error{false}, last_check{zeek::run_state::network_time} {
-    int status = MMDB_open(filename, MMDB_MODE_MMAP, &mmdb);
-
-    if ( MMDB_SUCCESS != status ) {
-        throw std::runtime_error(MMDB_strerror(status));
-    }
-}
-
-MMDB::~MMDB() { MMDB_close(&mmdb); }
-
-bool MMDB::Lookup(const zeek::IPAddr& addr, MMDB_lookup_result_s& result) {
-    struct sockaddr_storage ss = {0};
-
-    if ( IPv4 == addr.GetFamily() ) {
-        struct sockaddr_in* sa = (struct sockaddr_in*)&ss;
-        sa->sin_family = AF_INET;
-        addr.CopyIPv4(&sa->sin_addr);
-    }
-    else {
-        struct sockaddr_in6* sa = (struct sockaddr_in6*)&ss;
-        sa->sin6_family = AF_INET6;
-        addr.CopyIPv6(&sa->sin6_addr);
-    }
-
-    try {
-        result = Lookup((struct sockaddr*)&ss);
-    } catch ( const std::exception& e ) {
-        report_msg("MaxMind DB lookup location error [%s]", e.what());
-        return false;
-    }
-
-    return result.found_entry;
-}
-
-MMDB_lookup_result_s MMDB::Lookup(const struct sockaddr* const sa) {
-    int mmdb_error;
-    MMDB_lookup_result_s result = MMDB_lookup_sockaddr(&mmdb, sa, &mmdb_error);
-
-    if ( MMDB_SUCCESS != mmdb_error ) {
-        lookup_error = true;
-        throw std::runtime_error(MMDB_strerror(mmdb_error));
-    }
-
-    return result;
-}
-
-// Check to see if the Maxmind DB should be closed and reopened.  This will
-// happen if there was a lookup error or if the mmap'd file has been replaced
-// by an external process.
-bool MMDB::StaleDB() {
-    struct stat buf;
-
-    if ( lookup_error )
-        return true;
-
-    static double mmdb_stale_check_interval = zeek::id::find_val("mmdb_stale_check_interval")->AsInterval();
-
-    if ( mmdb_stale_check_interval < 0.0 )
-        return false;
-
-    if ( zeek::run_state::network_time - last_check < mmdb_stale_check_interval )
-        return false;
-
-    last_check = zeek::run_state::network_time;
-
-    if ( 0 != stat(mmdb.filename, &buf) )
-        return true;
-
-    if ( buf.st_ino != file_info.st_ino || buf.st_mtime != file_info.st_mtime ) {
-        report_msg("%s change detected for MaxMind DB [%s]",
-                   buf.st_ino != file_info.st_ino ? "Inode" : "Modification time", mmdb.filename);
-        return true;
-    }
-
-    return false;
-}
-
-const char* MMDB::Filename() { return mmdb.filename; }
-
-static bool mmdb_open(const char* filename, bool asn) {
-    struct stat buf;
-
-    if ( 0 != stat(filename, &buf) ) {
-        return false;
-    }
-
-    try {
-        if ( asn ) {
-            mmdb_asn.reset(new MMDB(filename, buf));
-        }
-        else {
-            mmdb_loc.reset(new MMDB(filename, buf));
-        }
-    }
-
-    catch ( const std::exception& e ) {
-        if ( asn )
-            did_asn_db_error = false;
-        else
-            did_loc_db_error = false;
-
-        report_msg("Failed to open MaxMind DB: %s [%s]", filename, e.what());
-        return false;
-    }
-
-    return true;
-}
-
-static bool mmdb_open_loc(const char* filename) { return mmdb_open(filename, false); }
-
-static bool mmdb_open_asn(const char* filename) { return mmdb_open(filename, true); }
-
-static void mmdb_check_loc() {
-    if ( mmdb_loc && mmdb_loc->StaleDB() ) {
-        report_msg("Closing stale MaxMind DB [%s]", mmdb_loc->Filename());
-        did_loc_db_error = false;
-        mmdb_loc.reset();
-    }
-}
-
-static void mmdb_check_asn() {
-    if ( mmdb_asn && mmdb_asn->StaleDB() ) {
-        report_msg("Closing stale MaxMind DB [%s]", mmdb_asn->Filename());
-        did_asn_db_error = false;
-        mmdb_asn.reset();
-    }
-}
-
 static zeek::ValPtr mmdb_getvalue(MMDB_entry_data_s* entry_data, int status, int data_type) {
     switch ( status ) {
         case MMDB_SUCCESS:
@@ -204,7 +73,115 @@ static zeek::ValPtr mmdb_getvalue(MMDB_entry_data_s* entry_data, int status, int
     return nullptr;
 }
 
-static bool mmdb_try_open_loc() {
+MMDB::MMDB() : mmdb{}, file_info{}, reported_error{false}, last_check{zeek::run_state::network_time} {}
+
+MMDB::~MMDB() { Close(); }
+
+bool MMDB::OpenFile(const char* a_filename) {
+    filename = a_filename;
+    Close();
+
+    if ( 0 != stat(a_filename, &file_info) ) {
+        return false;
+    }
+
+    int status = MMDB_open(a_filename, MMDB_MODE_MMAP, &mmdb);
+
+    if ( MMDB_SUCCESS != status ) {
+        memset(&mmdb, 0, sizeof(mmdb));
+        report_msg("Failed to open MaxMind DB: %s [%s]", filename.data(), MMDB_strerror(status));
+        return false;
+    }
+
+    return true;
+}
+
+void MMDB::Close() {
+    if ( IsOpen() ) {
+        MMDB_close(&mmdb);
+        memset(&mmdb, 0, sizeof(mmdb));
+        reported_error = false;
+    }
+}
+
+bool MMDB::EnsureLoaded() {
+    bool res = true;
+
+    if ( filename.empty() )
+        res = OpenFromScriptConfig();
+    else if ( ! IsOpen() )
+        res = OpenFile(filename.data());
+    else if ( IsStaleDB() ) {
+        report_msg("Closing stale MaxMind DB [%s]", filename.data());
+        if ( ! OpenFile(filename.data()) )
+            res = false;
+    }
+
+    if ( ! res && ! reported_error ) {
+        reported_error = true;
+        zeek::emit_builtin_error(zeek::util::fmt("Failed to open %s", Description()));
+    }
+
+    return res;
+}
+
+bool MMDB::Lookup(const zeek::IPAddr& addr, MMDB_lookup_result_s& result) {
+    struct sockaddr_storage ss = {0};
+
+    if ( IPv4 == addr.GetFamily() ) {
+        struct sockaddr_in* sa = (struct sockaddr_in*)&ss;
+        sa->sin_family = AF_INET;
+        addr.CopyIPv4(&sa->sin_addr);
+    }
+    else {
+        struct sockaddr_in6* sa = (struct sockaddr_in6*)&ss;
+        sa->sin6_family = AF_INET6;
+        addr.CopyIPv6(&sa->sin6_addr);
+    }
+
+    int mmdb_error;
+    result = MMDB_lookup_sockaddr(&mmdb, (struct sockaddr*)&ss, &mmdb_error);
+
+    if ( MMDB_SUCCESS != mmdb_error ) {
+        report_msg("MaxMind DB lookup location error [%s]", MMDB_strerror(mmdb_error));
+        Close();
+        return false;
+    }
+
+    return true;
+}
+
+// Check to see if the Maxmind DB should be closed and reopened.  This will
+// happen if there was a lookup error or if the mmap'd file has been replaced
+// by an external process.
+bool MMDB::IsStaleDB() {
+    if ( ! IsOpen() )
+        return false;
+
+    static double mmdb_stale_check_interval = zeek::id::find_val("mmdb_stale_check_interval")->AsInterval();
+
+    if ( mmdb_stale_check_interval < 0.0 )
+        return false;
+
+    if ( zeek::run_state::network_time - last_check < mmdb_stale_check_interval )
+        return false;
+
+    last_check = zeek::run_state::network_time;
+    struct stat buf;
+
+    if ( 0 != stat(mmdb.filename, &buf) )
+        return true;
+
+    if ( buf.st_ino != file_info.st_ino || buf.st_mtime != file_info.st_mtime ) {
+        report_msg("%s change detected for MaxMind DB [%s]",
+                   buf.st_ino != file_info.st_ino ? "Inode" : "Modification time", mmdb.filename);
+        return true;
+    }
+
+    return false;
+}
+
+bool LocDB::OpenFromScriptConfig() {
     // City database is always preferred over Country database.
     const auto& mmdb_dir_val = zeek::detail::global_scope()->Find("mmdb_dir")->GetVal();
     std::string mmdb_dir = mmdb_dir_val->AsString()->CheckString();
@@ -218,12 +195,12 @@ static bool mmdb_try_open_loc() {
     if ( ! mmdb_dir.empty() ) {
         auto d = mmdb_dir + "/" + mmdb_city_db;
 
-        if ( mmdb_open_loc(d.data()) )
+        if ( OpenFile(d.data()) )
             return true;
 
         d = mmdb_dir + "/" + mmdb_country_db;
 
-        if ( mmdb_open_loc(d.data()) )
+        if ( OpenFile(d.data()) )
             return true;
     }
 
@@ -232,20 +209,19 @@ static bool mmdb_try_open_loc() {
 
     for ( unsigned int i = 0; i < vv->Size(); ++i ) {
         auto d = std::string(vv->StringAt(i)->CheckString()) + "/" + mmdb_city_db;
-        if ( mmdb_open_loc(d.data()) )
+        if ( OpenFile(d.data()) )
             return true;
     }
 
     for ( unsigned int i = 0; i < vv->Size(); ++i ) {
         auto d = std::string(vv->StringAt(i)->CheckString()) + "/" + mmdb_country_db;
-        if ( mmdb_open_loc(d.data()) )
+        if ( OpenFile(d.data()) )
             return true;
     }
 
     return false;
 }
-
+bool AsnDB::OpenFromScriptConfig() {
-static bool mmdb_try_open_asn() {
     const auto& mmdb_dir_val = zeek::detail::global_scope()->Find("mmdb_dir")->GetVal();
     std::string mmdb_dir = mmdb_dir_val->AsString()->CheckString();
 
@@ -255,7 +231,7 @@ static bool mmdb_try_open_asn() {
     if ( ! mmdb_dir.empty() ) {
         auto d = mmdb_dir + "/" + mmdb_asn_db;
 
-        if ( mmdb_open_asn(d.data()) )
+        if ( OpenFile(d.data()) )
             return true;
     }
 
@@ -264,7 +240,7 @@ static bool mmdb_try_open_asn() {
 
     for ( unsigned int i = 0; i < vv->Size(); ++i ) {
         auto d = std::string(vv->StringAt(i)->CheckString()) + "/" + mmdb_asn_db;
-        if ( mmdb_open_loc(d.data()) )
+        if ( OpenFile(d.data()) )
             return true;
     }
 
@@ -274,7 +250,7 @@ static bool mmdb_try_open_asn() {
 
 ValPtr mmdb_open_location_db(StringVal* filename) {
 #ifdef USE_GEOIP
-    return zeek::val_mgr->Bool(mmdb_open_loc(filename->CheckString()));
+    return zeek::val_mgr->Bool(mmdb_loc.OpenFile(filename->CheckString()));
 #else
     return zeek::val_mgr->False();
 #endif
@@ -282,7 +258,7 @@ ValPtr mmdb_open_location_db(StringVal* filename) {
 
 ValPtr mmdb_open_asn_db(StringVal* filename) {
 #ifdef USE_GEOIP
-    return zeek::val_mgr->Bool(mmdb_open_asn(filename->CheckString()));
+    return zeek::val_mgr->Bool(mmdb_asn.OpenFile(filename->CheckString()));
 #else
     return zeek::val_mgr->False();
 #endif
@@ -293,21 +269,12 @@ RecordValPtr mmdb_lookup_location(AddrVal* addr) {
     auto location = zeek::make_intrusive<zeek::RecordVal>(geo_location);
 
 #ifdef USE_GEOIP
-    mmdb_check_loc();
+    if ( ! mmdb_loc.EnsureLoaded() )
-    if ( ! mmdb_loc ) {
+        return location;
-        if ( ! mmdb_try_open_loc() ) {
-            if ( ! did_loc_db_error ) {
-                did_loc_db_error = true;
-                zeek::emit_builtin_error("Failed to open GeoIP location database");
-            }
-
-            return location;
-        }
-    }
 
     MMDB_lookup_result_s result;
 
-    if ( mmdb_loc->Lookup(addr->AsAddr(), result) ) {
+    if ( mmdb_loc.Lookup(addr->AsAddr(), result) ) {
         MMDB_entry_data_s entry_data;
         int status;
 
@@ -354,21 +321,12 @@ RecordValPtr mmdb_lookup_autonomous_system(AddrVal* addr) {
     auto autonomous_system = zeek::make_intrusive<zeek::RecordVal>(geo_autonomous_system);
 
 #ifdef USE_GEOIP
-    mmdb_check_asn();
+    if ( ! mmdb_asn.EnsureLoaded() )
-    if ( ! mmdb_asn ) {
+        return autonomous_system;
-        if ( ! mmdb_try_open_asn() ) {
-            if ( ! did_asn_db_error ) {
-                did_asn_db_error = true;
-                zeek::emit_builtin_error("Failed to open GeoIP ASN database");
-            }
-
-            return autonomous_system;
-        }
-    }
 
     MMDB_lookup_result_s result;
 
-    if ( mmdb_asn->Lookup(addr->AsAddr(), result) ) {
+    if ( mmdb_asn.Lookup(addr->AsAddr(), result) ) {
         MMDB_entry_data_s entry_data;
         int status;
 

src/MMDB.h

@@ -10,25 +10,76 @@ namespace zeek {
 
 #include <maxminddb.h>
 
+// The MMDB class encapsulates a low-level libmaxmind MMDB_s structure. It
+// tracks whether that DB is currently loaded, and can open it from a file in
+// two ways: (1) via explicit specification of a filename, (2) by determining
+// the configuration from configuration settings in the script layer (mmdb_dir
+// etc). This configuration depends on whether this is a geolocation DB or an
+// ASN one, so details are left to derived classes below that specialize.
+//
+// The class tracks the inode and modification time of a DB file to detect
+// "stale" DBs, which get reloaded (from the same location in the file system)
+// upon the first lookup that detects staleness.
 class MMDB {
 public:
-    MMDB(const char* filename, struct stat info);
+    MMDB();
+    virtual ~MMDB();
 
-    ~MMDB();
+    // Implements the logic to determine a file system path for the DB from
+    // script-layer configuration settings, and opens the DB from there. Returns
+    // true if successful, false otherwise.
+    virtual bool OpenFromScriptConfig() = 0;
 
+    // Helper string to identify the type of DB, useful in error messages.
+    virtual const char* Description() = 0;
+
+    // Opens the DB at the given location, closing and cleaning up any currently
+    // opened DB if there is one. Returns true if successful, false otherwise.
+    bool OpenFile(const char* filename);
+
+    // Closes a currently opened DB, releasing its state. Safe to call on a
+    // closed DB.
+    void Close();
+
+    // Predicate; returns true if the DB is currently opened.
+    bool IsOpen() const { return mmdb.filename != nullptr; }
+
+    // Ensures that upon return the underlying DB file is loaded. When no
+    // filename is configured for the DB (i.e. OpenFile() has never been called
+    // on it), this triggers the script-level configuration lookup via
+    // OpenFromScriptConfig(). When a filename is available but it's not
+    // currently loaded, it does so. Finally, if there's a loaded DB but it's
+    // found to be stale, it gets reloaded. When the load operation succeeds, or
+    // the DB was already loaded and not stale, this returns true, and false if
+    // anything went wrong.
+    bool EnsureLoaded();
+
+    // Looks up a given IP address in the DB, storing the result in the provided
+    // result structure.
     bool Lookup(const zeek::IPAddr& addr, MMDB_lookup_result_s& result);
-    bool StaleDB();
-    const char* Filename();
 
 private:
-    MMDB_lookup_result_s Lookup(const struct sockaddr* const sa);
+    bool IsStaleDB();
 
+    std::string filename;
     MMDB_s mmdb;
     struct stat file_info;
-    bool lookup_error;
+    bool reported_error; // to ensure we emit builtin errors during opening only once.
     double last_check;
 };
 
+class LocDB : public MMDB {
+public:
+    bool OpenFromScriptConfig();
+    const char* Description() { return "GeoIP location database"; }
+};
+
+class AsnDB : public MMDB {
+public:
+    bool OpenFromScriptConfig();
+    const char* Description() { return "GeoIP ASN database"; }
+};
+
 #endif // USE_GEOIP
 
 ValPtr mmdb_open_location_db(zeek::StringVal* filename);