Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
121 changes: 62 additions & 59 deletions src/GeoDB.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -8,24 +8,23 @@

namespace visor::geo {

std::ostream &operator<<(std::ostream &os, const City &c)
{
return os << c.location;
}

MaxmindDB &GeoIP()
{
static MaxmindDB ip_db;
static MaxmindDB ip_db(MaxmindDB::Type::Geo);
return ip_db;
}

MaxmindDB &GeoASN()
{
static MaxmindDB asn_db;
static MaxmindDB asn_db(MaxmindDB::Type::Asn);
return asn_db;
}

std::ostream &operator<<(std::ostream &os, const visor::geo::City &c)
{
os << c.location;
return os;
}

void MaxmindDB::enable(const std::string &database_filename, int cache_size)
{
auto status = MMDB_open(database_filename.c_str(), MMDB_MODE_MMAP, &_mmdb);
Expand All @@ -34,7 +33,11 @@ void MaxmindDB::enable(const std::string &database_filename, int cache_size)
throw std::runtime_error(msg);
}
if (cache_size != 0) {
_lru_cache = std::make_unique<LRUList<std::string, std::variant<std::string, City>>>(cache_size);
if (_type == Type::Geo) {
_lru_geo_cache = std::make_unique<LRUList<std::string, City>>(cache_size);
} else if (_type == Type::Asn) {
_lru_asn_cache = std::make_unique<LRUList<std::string, std::string>>(cache_size);
}
}
_enabled = true;
}
Expand All @@ -50,7 +53,7 @@ City MaxmindDB::getGeoLoc(const struct sockaddr *sa) const
{

if (!_enabled) {
return City();
return {};
}

int mmdb_error;
Expand All @@ -67,33 +70,33 @@ City MaxmindDB::getGeoLoc(const struct sockaddr_in *sa4) const
{

if (!_enabled) {
return City();
return {};
}

std::string ip_address;
if (_lru_cache) {
if (_lru_geo_cache) {
ip_address = fmt::format_int(sa4->sin_addr.s_addr).str();
std::shared_lock lock(_cache_mutex);
if (auto geoloc = _lru_cache->getValue(ip_address); geoloc.has_value()) {
return std::get<City>(geoloc.value());
if (auto geoloc = _lru_geo_cache->getValue(ip_address); geoloc.has_value()) {
return geoloc.value();
}
}

int mmdb_error;

MMDB_lookup_result_s lookup = MMDB_lookup_sockaddr(&_mmdb, reinterpret_cast<const struct sockaddr *>(sa4), &mmdb_error);
if (mmdb_error != MMDB_SUCCESS || !lookup.found_entry) {
if (_lru_cache) {
if (_lru_geo_cache) {
std::unique_lock lock(_cache_mutex);
_lru_cache->put(ip_address, City{"Unknown", std::string(), std::string()});
_lru_geo_cache->put(ip_address, City{"Unknown", std::string(), std::string()});
}
return City{"Unknown", std::string(), std::string()};
}

if (_lru_cache) {
if (_lru_geo_cache) {
auto geoloc = _getGeoLoc(&lookup);
std::unique_lock lock(_cache_mutex);
_lru_cache->put(ip_address, geoloc);
_lru_geo_cache->put(ip_address, geoloc);
return geoloc;
}

Expand All @@ -104,33 +107,33 @@ City MaxmindDB::getGeoLoc(const struct sockaddr_in6 *sa6) const
{

if (!_enabled) {
return City();
return {};
}

std::string ip_address;
if (_lru_cache) {
if (_lru_geo_cache) {
ip_address = fmt::format("{}", sa6->sin6_addr.s6_addr);
std::shared_lock lock(_cache_mutex);
if (auto geoloc = _lru_cache->getValue(ip_address); geoloc.has_value()) {
return std::get<City>(geoloc.value());
if (auto geoloc = _lru_geo_cache->getValue(ip_address); geoloc.has_value()) {
return geoloc.value();
}
}

int mmdb_error;

MMDB_lookup_result_s lookup = MMDB_lookup_sockaddr(&_mmdb, reinterpret_cast<const struct sockaddr *>(sa6), &mmdb_error);
if (mmdb_error != MMDB_SUCCESS || !lookup.found_entry) {
if (_lru_cache) {
if (_lru_geo_cache) {
std::unique_lock lock(_cache_mutex);
_lru_cache->put(ip_address, City{"Unknown", std::string(), std::string()});
_lru_geo_cache->put(ip_address, City{"Unknown", std::string(), std::string()});
}
return City{"Unknown", std::string(), std::string()};
}

if (_lru_cache) {
if (_lru_geo_cache) {
auto geoloc = _getGeoLoc(&lookup);
std::unique_lock lock(_cache_mutex);
_lru_cache->put(ip_address, geoloc);
_lru_geo_cache->put(ip_address, geoloc);
return geoloc;
}

Expand All @@ -141,31 +144,31 @@ City MaxmindDB::getGeoLoc(const char *ip_address) const
{

if (!_enabled) {
return City();
return {};
}

if (_lru_cache) {
if (_lru_geo_cache) {
std::shared_lock lock(_cache_mutex);
if (auto geoloc = _lru_cache->getValue(ip_address); geoloc.has_value()) {
return std::get<City>(geoloc.value());
if (auto geoloc = _lru_geo_cache->getValue(ip_address); geoloc.has_value()) {
return geoloc.value();
}
}

int gai_error, mmdb_error;

MMDB_lookup_result_s lookup = MMDB_lookup_string(&_mmdb, ip_address, &gai_error, &mmdb_error);
if (0 != gai_error || MMDB_SUCCESS != mmdb_error || !lookup.found_entry) {
if (_lru_cache) {
if (_lru_geo_cache) {
std::unique_lock lock(_cache_mutex);
_lru_cache->put(ip_address, City{"Unknown", std::string(), std::string()});
_lru_geo_cache->put(ip_address, City{"Unknown", std::string(), std::string()});
}
return City{"Unknown", std::string(), std::string()};
}

if (_lru_cache) {
if (_lru_geo_cache) {
auto geoloc = _getGeoLoc(&lookup);
std::unique_lock lock(_cache_mutex);
_lru_cache->put(ip_address, geoloc);
_lru_geo_cache->put(ip_address, geoloc);
return geoloc;
}

Expand Down Expand Up @@ -243,7 +246,7 @@ std::string MaxmindDB::getASNString(const struct sockaddr *sa) const
{

if (!_enabled) {
return "";
return {};
}

int mmdb_error;
Expand All @@ -260,32 +263,32 @@ std::string MaxmindDB::getASNString(const struct sockaddr_in *sa4) const
{

if (!_enabled) {
return "";
return {};
}
std::string ip_address;
if (_lru_cache) {
if (_lru_asn_cache) {
ip_address = fmt::format_int(sa4->sin_addr.s_addr).str();
std::shared_lock lock(_cache_mutex);
if (auto asn = _lru_cache->getValue(ip_address); asn.has_value()) {
return std::get<std::string>(asn.value());
if (auto asn = _lru_asn_cache->getValue(ip_address); asn.has_value()) {
return asn.value();
}
}

int mmdb_error;

MMDB_lookup_result_s lookup = MMDB_lookup_sockaddr(&_mmdb, reinterpret_cast<const struct sockaddr *>(sa4), &mmdb_error);
if (mmdb_error != MMDB_SUCCESS || !lookup.found_entry) {
if (_lru_cache) {
if (_lru_asn_cache) {
std::unique_lock lock(_cache_mutex);
_lru_cache->put(ip_address, "Unknown");
_lru_asn_cache->put(ip_address, "Unknown");
}
return "Unknown";
}

if (_lru_cache) {
if (_lru_asn_cache) {
auto asn = _getASNString(&lookup);
std::unique_lock lock(_cache_mutex);
_lru_cache->put(ip_address, asn);
_lru_asn_cache->put(ip_address, asn);
return asn;
}

Expand All @@ -296,33 +299,33 @@ std::string MaxmindDB::getASNString(const struct sockaddr_in6 *sa6) const
{

if (!_enabled) {
return "";
return {};
}

std::string ip_address;
if (_lru_cache) {
if (_lru_asn_cache) {
ip_address = fmt::format("{}", sa6->sin6_addr.s6_addr);
std::shared_lock lock(_cache_mutex);
if (auto asn = _lru_cache->getValue(ip_address); asn.has_value()) {
return std::get<std::string>(asn.value());
if (auto asn = _lru_asn_cache->getValue(ip_address); asn.has_value()) {
return asn.value();
}
}

int mmdb_error;

MMDB_lookup_result_s lookup = MMDB_lookup_sockaddr(&_mmdb, reinterpret_cast<const struct sockaddr *>(sa6), &mmdb_error);
if (mmdb_error != MMDB_SUCCESS || !lookup.found_entry) {
if (_lru_cache) {
if (_lru_asn_cache) {
std::unique_lock lock(_cache_mutex);
_lru_cache->put(ip_address, "Unknown");
_lru_asn_cache->put(ip_address, "Unknown");
}
return "Unknown";
}

if (_lru_cache) {
if (_lru_asn_cache) {
auto asn = _getASNString(&lookup);
std::unique_lock lock(_cache_mutex);
_lru_cache->put(ip_address, asn);
_lru_asn_cache->put(ip_address, asn);
return asn;
}

Expand All @@ -333,31 +336,31 @@ std::string MaxmindDB::getASNString(const char *ip_address) const
{

if (!_enabled) {
return "";
return {};
}

if (_lru_cache) {
if (_lru_asn_cache) {
std::shared_lock lock(_cache_mutex);
if (auto asn = _lru_cache->getValue(ip_address); asn.has_value()) {
return std::get<std::string>(asn.value());
if (auto asn = _lru_asn_cache->getValue(ip_address); asn.has_value()) {
return asn.value();
}
}

int gai_error, mmdb_error;

MMDB_lookup_result_s lookup = MMDB_lookup_string(&_mmdb, ip_address, &gai_error, &mmdb_error);
if (0 != gai_error || MMDB_SUCCESS != mmdb_error || !lookup.found_entry) {
if (_lru_cache) {
if (_lru_asn_cache) {
std::unique_lock lock(_cache_mutex);
_lru_cache->put(ip_address, "Unknown");
_lru_asn_cache->put(ip_address, "Unknown");
}
return "Unknown";
}

if (_lru_cache) {
if (_lru_asn_cache) {
auto asn = _getASNString(&lookup);
std::unique_lock lock(_cache_mutex);
_lru_cache->put(ip_address, asn);
_lru_asn_cache->put(ip_address, asn);
return asn;
}

Expand Down
27 changes: 19 additions & 8 deletions src/GeoDB.h
Original file line number Diff line number Diff line change
Expand Up @@ -8,8 +8,8 @@
#pragma GCC diagnostic ignored "-Wpedantic"
#pragma GCC diagnostic ignored "-Wdeprecated-declarations"
#endif
#include <maxminddb.h>
#include <fmt/format.h>
#include <maxminddb.h>
#ifdef __GNUC__
#pragma GCC diagnostic pop
#endif
Expand All @@ -19,7 +19,7 @@
#include <ostream>
#include <shared_mutex>
#include <string>
#include <variant>
#include <tuple>

#include "VisorLRUList.h"

Expand All @@ -33,19 +33,28 @@ struct City {

bool operator==(const City &other) const
{
return (location == other.location
&& latitude == other.latitude
&& longitude == other.longitude);
return std::tie(location, latitude, longitude) == std::tie(other.location, other.latitude, other.longitude);
}
bool operator!=(const City &other) const
{
return !(*this == other);
}

friend std::ostream &operator<<(std::ostream &os, const City &c);
};

std::ostream &operator<<(std::ostream &os, const visor::geo::City &c);

class MaxmindDB
{
static constexpr size_t DEFAULT_CACHE_SIZE = 10000;

public:
enum class Type {
Asn,
Geo
};

MaxmindDB(Type type)
: _type(type){};
~MaxmindDB();

void enable(const std::string &database_filename, int cache_size = DEFAULT_CACHE_SIZE);
Expand All @@ -68,9 +77,11 @@ class MaxmindDB
std::string getASNString(const struct sockaddr_in6 *sa6) const;

private:
Type _type;
mutable MMDB_s _mmdb;
bool _enabled = false;
std::unique_ptr<LRUList<std::string, std::variant<std::string, City>>> _lru_cache;
std::unique_ptr<LRUList<std::string, City>> _lru_geo_cache;
std::unique_ptr<LRUList<std::string, std::string>> _lru_asn_cache;
mutable std::shared_mutex _cache_mutex;

City _getGeoLoc(MMDB_lookup_result_s *lookup) const;
Expand Down