diff --git a/src/session/Key.cc b/src/session/Key.cc index 5c7317505c..f87acc76b4 100644 --- a/src/session/Key.cc +++ b/src/session/Key.cc @@ -70,4 +70,14 @@ bool Key::operator<(const Key& rhs) const return memcmp(data, rhs.data, size) < 0; } +bool Key::operator==(const Key& rhs) const + { + if ( size != rhs.size ) + return false; + else if ( type != rhs.type ) + return false; + + return memcmp(data, rhs.data, size) == 0; + } + } // namespace zeek::session::detail diff --git a/src/session/Key.h b/src/session/Key.h index 7201fbcaa5..ca37d31088 100644 --- a/src/session/Key.h +++ b/src/session/Key.h @@ -4,9 +4,12 @@ #include #include +#include "zeek/Hash.h" namespace zeek::session::detail { +struct KeyHash; + /** * This type is used as the key for the map in SessionManager. It represents a * raw block of memory that points to a key of some type for a session, such as @@ -57,12 +60,23 @@ public: void CopyData(); bool operator<(const Key& rhs) const; + bool operator==(const Key& rhs) const; + + std::size_t Hash() const { + return zeek::detail::HashKey::HashBytes(data, size); + } private: + friend struct KeyHash; + const uint8_t* data = nullptr; size_t size = 0; size_t type = CONNECTION_KEY_TYPE; bool copied = false; }; +struct KeyHash { + std::size_t operator()(const Key& k) const { return k.Hash(); } +}; + } // namespace zeek::session::detail diff --git a/src/session/Manager.cc b/src/session/Manager.cc index 527adfed74..1792a212ed 100644 --- a/src/session/Manager.cc +++ b/src/session/Manager.cc @@ -218,9 +218,17 @@ void Manager::Insert(Session* s, bool remove_existing) void Manager::Drain() { - for ( const auto& entry : session_map ) + std::vector keys; + keys.reserve(session_map.size()); + + for ( auto& entry : session_map ) + keys.push_back(&(entry.first)); + std::sort(keys.begin(), keys.end(), [](const detail::Key* a, const detail::Key* b) { + return *a < *b; }); + + for ( const auto* k : keys ) { - Session* tc = entry.second; + Session* tc = session_map.at(*k); tc->Done(); tc->RemovalEvent(); } diff --git a/src/session/Manager.h b/src/session/Manager.h index 976cc7f10d..e824e0b14b 100644 --- a/src/session/Manager.h +++ b/src/session/Manager.h @@ -3,7 +3,7 @@ #pragma once #include // for u_char -#include +#include #include #include "zeek/Frag.h" @@ -119,7 +119,7 @@ public: private: - using SessionMap = std::map; + using SessionMap = std::unordered_map; // Inserts a new connection into the sessions map. If a connection with // the same key already exists in the map, it will be overwritten by