Add type field to session::Key to help avoid collisions in map

This commit is contained in:
Tim Wojtulewicz 2021-05-20 11:00:11 -07:00
parent 30ab914cd8
commit 3a8047f535
4 changed files with 20 additions and 13 deletions

View file

@ -4,7 +4,8 @@
namespace zeek::session::detail {
Key::Key(const void* session, size_t size, bool copy) : size(size)
Key::Key(const void* session, size_t size, size_t type, bool copy) :
size(size), type(type)
{
data = reinterpret_cast<const uint8_t*>(session);
@ -63,6 +64,8 @@ bool Key::operator<(const Key& rhs) const
{
if ( size != rhs.size )
return size < rhs.size;
else if ( type != rhs.type )
return type < rhs.type;
return memcmp(data, rhs.data, size) < 0;
}

View file

@ -20,16 +20,22 @@ namespace zeek::session::detail {
class Key final {
public:
const static size_t CONNECTION_KEY_TYPE=0;
/**
* Create a new session key from a data pointer.
*
* @param session A pointer to the data for the key.
* @param size The size of the key data, in bytes.
* @param type An identifier for the type of this key. The value used should be
* unique across all types of session keys. CONNECTION_KEY_TYPE (0) is used by
* Connection sessions and is reserved. This value is used to avoid collisions
* when doing comparisons of the memory stored by keys.
* @param copy Flag for whether the data should be copied into the Key
* during construction. This defaults to false because normally the only time
* data is copied into the key is when it's inserted into the session map.
*/
Key(const void* key_data, size_t size, bool copy=false);
Key(const void* key_data, size_t size, size_t type, bool copy=false);
~Key();
@ -55,6 +61,7 @@ public:
private:
const uint8_t* data = nullptr;
size_t size = 0;
size_t type = CONNECTION_KEY_TYPE;
bool copied = false;
};

View file

@ -149,23 +149,17 @@ Connection* Manager::FindConnection(Val* v)
htons((unsigned short) resp_portv->Port()),
orig_portv->PortType(), false);
detail::Key key(&conn_key, sizeof(conn_key), false);
Connection* conn = nullptr;
auto it = session_map.find(key);
if ( it != session_map.end() )
conn = static_cast<Connection*>(it->second);
return conn;
return FindConnection(conn_key);
}
Connection* Manager::FindConnection(const zeek::detail::ConnKey& conn_key)
{
detail::Key key(&conn_key, sizeof(conn_key), false);
detail::Key key(&conn_key, sizeof(conn_key),
detail::Key::CONNECTION_KEY_TYPE, false);
auto it = session_map.find(key);
if ( it != session_map.end() )
return dynamic_cast<Connection*>(it->second);
return static_cast<Connection*>(it->second);
return nullptr;
}