From 3a8047f5356bbd73ec2036df254837eabdafc821 Mon Sep 17 00:00:00 2001
From: Tim Wojtulewicz <tim@corelight.com>
Date: Thu, 20 May 2021 11:00:11 -0700
Subject: [PATCH] Add type field to session::Key to help avoid collisions in
 map

---
 src/Conn.h             |  5 ++++-
 src/session/Key.cc     |  5 ++++-
 src/session/Key.h      |  9 ++++++++-
 src/session/Manager.cc | 14 ++++----------
 4 files changed, 20 insertions(+), 13 deletions(-)

diff --git a/src/Conn.h b/src/Conn.h
index 20a5b138c4..19d5d943d3 100644
--- a/src/Conn.h
+++ b/src/Conn.h
@@ -113,7 +113,10 @@ public:
 	// should be marked invalid.
 	const detail::ConnKey& Key() const	{ return key; }
 	session::detail::Key SessionKey(bool copy) const override
-		{ return session::detail::Key{&key, sizeof(key), copy}; }
+		{
+		return session::detail::Key{
+			&key, sizeof(key), session::detail::Key::CONNECTION_KEY_TYPE, copy};
+		}
 
 	const IPAddr& OrigAddr() const		{ return orig_addr; }
 	const IPAddr& RespAddr() const		{ return resp_addr; }
diff --git a/src/session/Key.cc b/src/session/Key.cc
index 6044f493f8..5c7317505c 100644
--- a/src/session/Key.cc
+++ b/src/session/Key.cc
@@ -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;
 	}
diff --git a/src/session/Key.h b/src/session/Key.h
index f07ec8996b..7201fbcaa5 100644
--- a/src/session/Key.h
+++ b/src/session/Key.h
@@ -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;
 };
 
diff --git a/src/session/Manager.cc b/src/session/Manager.cc
index df027c82cf..82ee11b217 100644
--- a/src/session/Manager.cc
+++ b/src/session/Manager.cc
@@ -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;
 	}