Modernize various C++/Zeek-isms in the MMDB code.

This commit is contained in:
Christian Kreibich 2024-01-11 17:13:41 -08:00
parent e8f0f727cd
commit dbad072f06
3 changed files with 48 additions and 50 deletions

View file

@ -77,15 +77,15 @@ MMDB::MMDB() : mmdb{}, file_info{}, reported_error{false}, last_check{zeek::run_
MMDB::~MMDB() { Close(); } MMDB::~MMDB() { Close(); }
bool MMDB::OpenFile(const char* a_filename) { bool MMDB::OpenFile(const std::string& a_filename) {
filename = a_filename; filename = a_filename;
Close(); Close();
if ( 0 != stat(a_filename, &file_info) ) { if ( 0 != stat(filename.data(), &file_info) ) {
return false; return false;
} }
int status = MMDB_open(a_filename, MMDB_MODE_MMAP, &mmdb); int status = MMDB_open(a_filename.data(), MMDB_MODE_MMAP, &mmdb);
if ( MMDB_SUCCESS != status ) { if ( MMDB_SUCCESS != status ) {
memset(&mmdb, 0, sizeof(mmdb)); memset(&mmdb, 0, sizeof(mmdb));
@ -110,16 +110,16 @@ bool MMDB::EnsureLoaded() {
if ( filename.empty() ) if ( filename.empty() )
res = OpenFromScriptConfig(); res = OpenFromScriptConfig();
else if ( ! IsOpen() ) else if ( ! IsOpen() )
res = OpenFile(filename.data()); res = OpenFile(filename);
else if ( IsStaleDB() ) { else if ( IsStaleDB() ) {
report_msg("Closing stale MaxMind DB [%s]", filename.data()); report_msg("Closing stale MaxMind DB [%s]", filename.data());
if ( ! OpenFile(filename.data()) ) if ( ! OpenFile(filename) )
res = false; res = false;
} }
if ( ! res && ! reported_error ) { if ( ! res && ! reported_error ) {
reported_error = true; reported_error = true;
zeek::emit_builtin_error(zeek::util::fmt("Failed to open %s", Description())); zeek::emit_builtin_error(zeek::util::fmt("Failed to open %s", Description().data()));
} }
return res; return res;
@ -183,64 +183,62 @@ bool MMDB::IsStaleDB() {
bool LocDB::OpenFromScriptConfig() { bool LocDB::OpenFromScriptConfig() {
// City database is always preferred over Country database. // City database is always preferred over Country database.
const auto& mmdb_dir_val = zeek::detail::global_scope()->Find("mmdb_dir")->GetVal(); const auto& mmdb_dir_val = zeek::id::find_val<StringVal>("mmdb_dir");
std::string mmdb_dir = mmdb_dir_val->AsString()->CheckString(); std::string mmdb_dir{mmdb_dir_val->ToStdStringView()};
const auto& mmdb_city_db_val = zeek::detail::global_scope()->Find("mmdb_city_db")->GetVal(); const auto& mmdb_city_db_val = zeek::id::find_val<StringVal>("mmdb_city_db");
std::string mmdb_city_db = mmdb_city_db_val->AsString()->CheckString(); std::string mmdb_city_db{mmdb_city_db_val->ToStdStringView()};
const auto& mmdb_country_db_val = zeek::detail::global_scope()->Find("mmdb_country_db")->GetVal(); const auto& mmdb_country_db_val = zeek::id::find_val<StringVal>("mmdb_country_db");
std::string mmdb_country_db = mmdb_country_db_val->AsString()->CheckString(); std::string mmdb_country_db{mmdb_country_db_val->ToStdStringView()};
if ( ! mmdb_dir.empty() ) { if ( ! mmdb_dir.empty() ) {
auto d = mmdb_dir + "/" + mmdb_city_db; auto d = mmdb_dir + "/" + mmdb_city_db;
if ( OpenFile(d.data()) ) if ( OpenFile(d) )
return true; return true;
d = mmdb_dir + "/" + mmdb_country_db; d = mmdb_dir + "/" + mmdb_country_db;
if ( OpenFile(d.data()) ) if ( OpenFile(d) )
return true; return true;
} }
const auto& mmdb_dir_fallbacks_val = zeek::detail::global_scope()->Find("mmdb_dir_fallbacks")->GetVal(); const auto& mmdb_dir_fallbacks_val = zeek::id::find_val<VectorVal>("mmdb_dir_fallbacks");
auto* vv = mmdb_dir_fallbacks_val->AsVectorVal();
for ( unsigned int i = 0; i < vv->Size(); ++i ) { for ( unsigned int i = 0; i < mmdb_dir_fallbacks_val->Size(); ++i ) {
auto d = std::string(vv->StringAt(i)->CheckString()) + "/" + mmdb_city_db; auto d = mmdb_dir_fallbacks_val->StringValAt(i)->ToStdString() + "/" + mmdb_city_db;
if ( OpenFile(d.data()) ) if ( OpenFile(d) )
return true; return true;
} }
for ( unsigned int i = 0; i < vv->Size(); ++i ) { for ( unsigned int i = 0; i < mmdb_dir_fallbacks_val->Size(); ++i ) {
auto d = std::string(vv->StringAt(i)->CheckString()) + "/" + mmdb_country_db; auto d = mmdb_dir_fallbacks_val->StringValAt(i)->ToStdString() + "/" + mmdb_country_db;
if ( OpenFile(d.data()) ) if ( OpenFile(d) )
return true; return true;
} }
return false; return false;
} }
bool AsnDB::OpenFromScriptConfig() { bool AsnDB::OpenFromScriptConfig() {
const auto& mmdb_dir_val = zeek::detail::global_scope()->Find("mmdb_dir")->GetVal(); const auto& mmdb_dir_val = zeek::id::find_val<StringVal>("mmdb_dir");
std::string mmdb_dir = mmdb_dir_val->AsString()->CheckString(); std::string mmdb_dir{mmdb_dir_val->ToStdStringView()};
const auto& mmdb_asn_db_val = zeek::detail::global_scope()->Find("mmdb_asn_db")->GetVal(); const auto& mmdb_asn_db_val = zeek::id::find_val<StringVal>("mmdb_asn_db");
std::string mmdb_asn_db = mmdb_asn_db_val->AsString()->CheckString(); std::string mmdb_asn_db{mmdb_asn_db_val->ToStdStringView()};
if ( ! mmdb_dir.empty() ) { if ( ! mmdb_dir.empty() ) {
auto d = mmdb_dir + "/" + mmdb_asn_db; auto d = mmdb_dir + "/" + mmdb_asn_db;
if ( OpenFile(d.data()) ) if ( OpenFile(d) )
return true; return true;
} }
const auto& mmdb_dir_fallbacks_val = zeek::detail::global_scope()->Find("mmdb_dir_fallbacks")->GetVal(); const auto& mmdb_dir_fallbacks_val = zeek::id::find_val<VectorVal>("mmdb_dir_fallbacks");
auto* vv = mmdb_dir_fallbacks_val->AsVectorVal();
for ( unsigned int i = 0; i < vv->Size(); ++i ) { for ( unsigned int i = 0; i < mmdb_dir_fallbacks_val->Size(); ++i ) {
auto d = std::string(vv->StringAt(i)->CheckString()) + "/" + mmdb_asn_db; auto d = mmdb_dir_fallbacks_val->StringValAt(i)->ToStdString() + "/" + mmdb_asn_db;
if ( OpenFile(d.data()) ) if ( OpenFile(d) )
return true; return true;
} }
@ -248,23 +246,23 @@ bool AsnDB::OpenFromScriptConfig() {
} }
#endif // USE_GEOIP #endif // USE_GEOIP
ValPtr mmdb_open_location_db(StringVal* filename) { ValPtr mmdb_open_location_db(const StringValPtr& filename) {
#ifdef USE_GEOIP #ifdef USE_GEOIP
return zeek::val_mgr->Bool(mmdb_loc.OpenFile(filename->CheckString())); return zeek::val_mgr->Bool(mmdb_loc.OpenFile(filename->ToStdString()));
#else #else
return zeek::val_mgr->False(); return zeek::val_mgr->False();
#endif #endif
} }
ValPtr mmdb_open_asn_db(StringVal* filename) { ValPtr mmdb_open_asn_db(const StringValPtr& filename) {
#ifdef USE_GEOIP #ifdef USE_GEOIP
return zeek::val_mgr->Bool(mmdb_asn.OpenFile(filename->CheckString())); return zeek::val_mgr->Bool(mmdb_asn.OpenFile(filename->ToStdString()));
#else #else
return zeek::val_mgr->False(); return zeek::val_mgr->False();
#endif #endif
} }
RecordValPtr mmdb_lookup_location(AddrVal* addr) { RecordValPtr mmdb_lookup_location(const AddrValPtr& addr) {
static auto geo_location = zeek::id::find_type<zeek::RecordType>("geo_location"); static auto geo_location = zeek::id::find_type<zeek::RecordType>("geo_location");
auto location = zeek::make_intrusive<zeek::RecordVal>(geo_location); auto location = zeek::make_intrusive<zeek::RecordVal>(geo_location);
@ -316,7 +314,7 @@ RecordValPtr mmdb_lookup_location(AddrVal* addr) {
return location; return location;
} }
RecordValPtr mmdb_lookup_autonomous_system(AddrVal* addr) { RecordValPtr mmdb_lookup_autonomous_system(const AddrValPtr& addr) {
static auto geo_autonomous_system = zeek::id::find_type<zeek::RecordType>("geo_autonomous_system"); static auto geo_autonomous_system = zeek::id::find_type<zeek::RecordType>("geo_autonomous_system");
auto autonomous_system = zeek::make_intrusive<zeek::RecordVal>(geo_autonomous_system); auto autonomous_system = zeek::make_intrusive<zeek::RecordVal>(geo_autonomous_system);

View file

@ -31,11 +31,11 @@ public:
virtual bool OpenFromScriptConfig() = 0; virtual bool OpenFromScriptConfig() = 0;
// Helper string to identify the type of DB, useful in error messages. // Helper string to identify the type of DB, useful in error messages.
virtual const char* Description() = 0; virtual std::string_view Description() = 0;
// Opens the DB at the given location, closing and cleaning up any currently // 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. // opened DB if there is one. Returns true if successful, false otherwise.
bool OpenFile(const char* filename); bool OpenFile(const std::string& filename);
// Closes a currently opened DB, releasing its state. Safe to call on a // Closes a currently opened DB, releasing its state. Safe to call on a
// closed DB. // closed DB.
@ -71,21 +71,21 @@ private:
class LocDB : public MMDB { class LocDB : public MMDB {
public: public:
bool OpenFromScriptConfig(); bool OpenFromScriptConfig();
const char* Description() { return "GeoIP location database"; } std::string_view Description() { return "GeoIP location database"; }
}; };
class AsnDB : public MMDB { class AsnDB : public MMDB {
public: public:
bool OpenFromScriptConfig(); bool OpenFromScriptConfig();
const char* Description() { return "GeoIP ASN database"; } std::string_view Description() { return "GeoIP ASN database"; }
}; };
#endif // USE_GEOIP #endif // USE_GEOIP
ValPtr mmdb_open_location_db(zeek::StringVal* filename); ValPtr mmdb_open_location_db(const StringValPtr& filename);
ValPtr mmdb_open_asn_db(zeek::StringVal* filename); ValPtr mmdb_open_asn_db(const StringValPtr& filename);
RecordValPtr mmdb_lookup_location(zeek::AddrVal* addr); RecordValPtr mmdb_lookup_location(const AddrValPtr& addr);
RecordValPtr mmdb_lookup_autonomous_system(zeek::AddrVal* addr); RecordValPtr mmdb_lookup_autonomous_system(const AddrValPtr& addr);
} // namespace zeek } // namespace zeek

View file

@ -12,7 +12,7 @@
## .. zeek:see:: lookup_autonomous_system ## .. zeek:see:: lookup_autonomous_system
function mmdb_open_location_db%(f: string%) : bool function mmdb_open_location_db%(f: string%) : bool
%{ %{
return zeek::mmdb_open_location_db(f); return zeek::mmdb_open_location_db(StringValPtr(NewRef(), f));
%} %}
## Initializes MMDB for later use of lookup_autonomous_system. ## Initializes MMDB for later use of lookup_autonomous_system.
@ -25,7 +25,7 @@ function mmdb_open_location_db%(f: string%) : bool
## .. zeek:see:: lookup_autonomous_system ## .. zeek:see:: lookup_autonomous_system
function mmdb_open_asn_db%(f: string%) : bool function mmdb_open_asn_db%(f: string%) : bool
%{ %{
return zeek::mmdb_open_asn_db(f); return zeek::mmdb_open_asn_db(StringValPtr(NewRef(), f));
%} %}
## Performs a geo-lookup of an IP address. ## Performs a geo-lookup of an IP address.
@ -38,7 +38,7 @@ function mmdb_open_asn_db%(f: string%) : bool
## .. zeek:see:: lookup_autonomous_system ## .. zeek:see:: lookup_autonomous_system
function lookup_location%(a: addr%) : geo_location function lookup_location%(a: addr%) : geo_location
%{ %{
return zeek::mmdb_lookup_location(a); return zeek::mmdb_lookup_location(AddrValPtr(NewRef(), a));
%} %}
## Performs an lookup of AS number & organization of an IP address. ## Performs an lookup of AS number & organization of an IP address.
@ -51,5 +51,5 @@ function lookup_location%(a: addr%) : geo_location
## .. zeek:see:: lookup_location ## .. zeek:see:: lookup_location
function lookup_autonomous_system%(a: addr%) : geo_autonomous_system function lookup_autonomous_system%(a: addr%) : geo_autonomous_system
%{ %{
return zeek::mmdb_lookup_autonomous_system(a); return zeek::mmdb_lookup_autonomous_system(AddrValPtr(NewRef(), a));
%} %}