Skip to content

Commit

Permalink
Fix MMDB code to re-open explicitly opened DBs correctly
Browse files Browse the repository at this point in the history 
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.
  • Loading branch information
@ckreibich
ckreibich committed Jan 11, 2024
1 parent d79e2c2 commit 2882cf3
Show file tree
Hide file tree
Showing 2 changed files with 161 additions and 151 deletions.
249 changes: 104 additions & 145 deletions src/MMDB.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
static std::unique_ptr<MMDB> mmdb_asn;
LocDB mmdb_loc;
AsnDB mmdb_asn;

static void report_msg(const char* format, ...) {
if ( zeek::run_state::network_time > msg_suppression_time + msg_suppression_duration ) {
Expand All @@ -45,16 +43,88 @@ 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);
static zeek::ValPtr mmdb_getvalue(MMDB_entry_data_s* entry_data, int status, int data_type) {
switch ( status ) {
case MMDB_SUCCESS:
if ( entry_data->has_data ) {
switch ( data_type ) {
case MMDB_DATA_TYPE_UTF8_STRING:
return zeek::make_intrusive<zeek::StringVal>(entry_data->data_size, entry_data->utf8_string);
break;

case MMDB_DATA_TYPE_DOUBLE:
return zeek::make_intrusive<zeek::DoubleVal>(entry_data->double_value);
break;

case MMDB_DATA_TYPE_UINT32: return zeek::val_mgr->Count(entry_data->uint32);

default: break;
}
}
break;

case MMDB_LOOKUP_PATH_DOES_NOT_MATCH_DATA_ERROR:
// key doesn't exist, nothing to do
break;

default: report_msg("MaxMind DB error [%s]", MMDB_strerror(status)); break;
}

return nullptr;
}

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 ) {
throw std::runtime_error(MMDB_strerror(status));
memset(&mmdb, 0, sizeof(mmdb));
report_msg("Failed to open MaxMind DB: %s [%s]", filename.data(), MMDB_strerror(status));
return false;
}

return true;
}

MMDB::~MMDB() { MMDB_close(&mmdb); }
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};
Expand All @@ -70,36 +140,24 @@ bool MMDB::Lookup(const zeek::IPAddr& addr, MMDB_lookup_result_s& result) {
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);
result = MMDB_lookup_sockaddr(&mmdb, (struct sockaddr*)&ss, &mmdb_error);

if ( MMDB_SUCCESS != mmdb_error ) {
lookup_error = true;
throw std::runtime_error(MMDB_strerror(mmdb_error));
report_msg("MaxMind DB lookup location error [%s]", MMDB_strerror(mmdb_error));
Close();
return false;
}

return result;
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::StaleDB() {
struct stat buf;

if ( lookup_error )
return true;
bool MMDB::IsStaleDB() {
if ( ! IsOpen() )
return false;

static double mmdb_stale_check_interval = zeek::id::find_val("mmdb_stale_check_interval")->AsInterval();

Expand All @@ -110,6 +168,7 @@ bool MMDB::StaleDB() {
return false;

last_check = zeek::run_state::network_time;
struct stat buf;

if ( 0 != stat(mmdb.filename, &buf) )
return true;
Expand All @@ -123,88 +182,7 @@ bool MMDB::StaleDB() {
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:
if ( entry_data->has_data ) {
switch ( data_type ) {
case MMDB_DATA_TYPE_UTF8_STRING:
return zeek::make_intrusive<zeek::StringVal>(entry_data->data_size, entry_data->utf8_string);
break;

case MMDB_DATA_TYPE_DOUBLE:
return zeek::make_intrusive<zeek::DoubleVal>(entry_data->double_value);
break;

case MMDB_DATA_TYPE_UINT32: return zeek::val_mgr->Count(entry_data->uint32);

default: break;
}
}
break;

case MMDB_LOOKUP_PATH_DOES_NOT_MATCH_DATA_ERROR:
// key doesn't exist, nothing to do
break;

default: report_msg("MaxMind DB error [%s]", MMDB_strerror(status)); break;
}

return nullptr;
}

static bool mmdb_try_open_loc() {
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();
Expand All @@ -218,12 +196,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;
}

Expand All @@ -232,20 +210,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;
}

static bool mmdb_try_open_asn() {
bool AsnDB::OpenFromScriptConfig() {
const auto& mmdb_dir_val = zeek::detail::global_scope()->Find("mmdb_dir")->GetVal();
std::string mmdb_dir = mmdb_dir_val->AsString()->CheckString();

Expand All @@ -255,7 +232,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;
}

Expand All @@ -264,7 +241,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;
}

Expand All @@ -274,15 +251,15 @@ 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
}

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
Expand All @@ -293,21 +270,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 ) {
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;
}
}
if ( ! mmdb_loc.EnsureLoaded() )
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;

Expand Down Expand Up @@ -354,21 +322,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 ) {
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;
}
}
if ( ! mmdb_asn.EnsureLoaded() )
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;

Expand Down
Loading

0 comments on commit 2882cf3

Please sign in to comment.