From a4b8aa1f30b04ed7abaf6659a2e82e9f9e7d0660 Mon Sep 17 00:00:00 2001 From: Tim Wojtulewicz Date: Thu, 1 Aug 2019 14:20:10 -0700 Subject: [PATCH] Change FragReassembler to use a tuple as a key and use std::map for fragments in Sessions --- src/Frag.cc | 3 +-- src/Frag.h | 10 ++++++--- src/Sessions.cc | 55 ++++++++++++++----------------------------------- src/Sessions.h | 4 ++-- 4 files changed, 26 insertions(+), 46 deletions(-) diff --git a/src/Frag.cc b/src/Frag.cc index 4201aa3faa..51f82e84f3 100644 --- a/src/Frag.cc +++ b/src/Frag.cc @@ -27,7 +27,7 @@ void FragTimer::Dispatch(double t, int /* is_expire */) FragReassembler::FragReassembler(NetSessions* arg_s, const IP_Hdr* ip, const u_char* pkt, - HashKey* k, double t) + const FragReassemblerKey& k, double t) : Reassembler(0, REASSEM_FRAG) { s = arg_s; @@ -68,7 +68,6 @@ FragReassembler::~FragReassembler() DeleteTimer(); delete [] proto_hdr; delete reassembled_pkt; - delete key; } void FragReassembler::AddFragment(double t, const IP_Hdr* ip, const u_char* pkt) diff --git a/src/Frag.h b/src/Frag.h index bf451cd5ed..8957237f03 100644 --- a/src/Frag.h +++ b/src/Frag.h @@ -3,6 +3,8 @@ #ifndef frag_h #define frag_h +#include + #include "util.h" #include "IP.h" #include "Net.h" @@ -17,10 +19,12 @@ class FragTimer; typedef void (FragReassembler::*frag_timer_func)(double t); +using FragReassemblerKey = std::tuple; + class FragReassembler : public Reassembler { public: FragReassembler(NetSessions* s, const IP_Hdr* ip, const u_char* pkt, - HashKey* k, double t); + const FragReassemblerKey& k, double t); ~FragReassembler() override; void AddFragment(double t, const IP_Hdr* ip, const u_char* pkt); @@ -30,7 +34,7 @@ public: void ClearTimer() { expire_timer = 0; } const IP_Hdr* ReassembledPkt() { return reassembled_pkt; } - HashKey* Key() const { return key; } + const FragReassemblerKey& Key() const { return key; } protected: void BlockInserted(DataBlock* start_block) override; @@ -43,7 +47,7 @@ protected: NetSessions* s; uint64_t frag_size; // size of fully reassembled fragment uint16_t next_proto; // first IPv6 fragment header's next proto field - HashKey* key; + FragReassemblerKey key; FragTimer* expire_timer; }; diff --git a/src/Sessions.cc b/src/Sessions.cc index 4ca4ec6ef6..34c159f772 100644 --- a/src/Sessions.cc +++ b/src/Sessions.cc @@ -87,17 +87,6 @@ void IPTunnelTimer::Dispatch(double t, int is_expire) NetSessions::NetSessions() { - TypeList* t = new TypeList(); - t->Append(base_type(TYPE_ADDR)); // source IP address - t->Append(base_type(TYPE_ADDR)); // dest IP address - t->Append(base_type(TYPE_COUNT)); // source and dest ports - - ch = new CompositeHash(t); - - Unref(t); - - fragments.SetDeleteFunc(bro_obj_delete_func); - if ( stp_correlate_pair ) stp_manager = new analyzer::stepping_stone::SteppingStoneManager(); else @@ -131,7 +120,6 @@ NetSessions::NetSessions() NetSessions::~NetSessions() { - delete ch; delete packet_filter; delete pkt_profiler; Unref(arp_analyzer); @@ -144,6 +132,8 @@ NetSessions::~NetSessions() Unref(entry.second); for ( const auto& entry : icmp_conns ) Unref(entry.second); + for ( const auto& entry : fragments ) + Unref(entry.second); } void NetSessions::Done() @@ -939,27 +929,22 @@ FragReassembler* NetSessions::NextFragment(double t, const IP_Hdr* ip, { uint32_t frag_id = ip->ID(); - ListVal* key = new ListVal(TYPE_ANY); - key->Append(new AddrVal(ip->SrcAddr())); - key->Append(new AddrVal(ip->DstAddr())); - key->Append(val_mgr->GetCount(frag_id)); + FragReassemblerKey key = std::make_tuple(ip->SrcAddr(), ip->DstAddr(), frag_id); - HashKey* h = ch->ComputeHash(key, 1); - if ( ! h ) - reporter->InternalError("hash computation failed"); + FragReassembler* f = nullptr; + auto it = fragments.find(key); + if ( it != fragments.end() ) + f = it->second; - FragReassembler* f = fragments.Lookup(h); if ( ! f ) { - f = new FragReassembler(this, ip, pkt, h, t); - fragments.Insert(h, f); - Unref(key); + f = new FragReassembler(this, ip, pkt, key, t); + fragments[key] = f; + if ( fragments.size() > stats.max_fragments ) + stats.max_fragments = fragments.size(); return f; } - delete h; - Unref(key); - f->AddFragment(t, ip, pkt); return f; } @@ -1097,15 +1082,8 @@ void NetSessions::Remove(FragReassembler* f) if ( ! f ) return; - HashKey* k = f->Key(); - - if ( k ) - { - if ( ! fragments.RemoveEntry(k) ) - reporter->InternalWarning("fragment reassembler not in dict"); - } - else - reporter->InternalWarning("missing fragment reassembler hash key"); + if ( fragments.erase(f->Key()) == 0 ) + reporter->InternalWarning("fragment reassembler not in dict"); Unref(f); } @@ -1188,13 +1166,13 @@ void NetSessions::GetStats(SessionStats& s) const s.cumulative_UDP_conns = stats.cumulative_UDP_conns; s.num_ICMP_conns = icmp_conns.size(); s.cumulative_ICMP_conns = stats.cumulative_ICMP_conns; - s.num_fragments = fragments.Length(); + s.num_fragments = fragments.size(); s.num_packets = num_packets_processed; s.max_TCP_conns = stats.max_TCP_conns; s.max_UDP_conns = stats.max_UDP_conns; s.max_ICMP_conns = stats.max_ICMP_conns; - s.max_fragments = fragments.MaxLength(); + s.max_fragments = stats.max_fragments; } Connection* NetSessions::NewConn(const ConnIDKey& k, double t, const ConnID* id, @@ -1481,11 +1459,10 @@ unsigned int NetSessions::MemoryAllocation() return ConnectionMemoryUsage() + padded_sizeof(*this) - + ch->MemoryAllocation() + padded_sizeof(tcp_conns.size() * (sizeof(ConnectionMap::key_type) + sizeof(ConnectionMap::value_type))) + padded_sizeof(udp_conns.size() * (sizeof(ConnectionMap::key_type) + sizeof(ConnectionMap::value_type))) + padded_sizeof(icmp_conns.size() * (sizeof(ConnectionMap::key_type) + sizeof(ConnectionMap::value_type))) - + fragments.MemoryAllocation() - padded_sizeof(fragments) + + padded_sizeof(fragments.size() * (sizeof(FragmentMap::key_type) + sizeof(FragmentMap::value_type))) // FIXME: MemoryAllocation() not implemented for rest. ; } diff --git a/src/Sessions.h b/src/Sessions.h index 20ee491ff3..372982ff18 100644 --- a/src/Sessions.h +++ b/src/Sessions.h @@ -173,6 +173,7 @@ protected: friend class IPTunnelTimer; using ConnectionMap = std::map; + using FragmentMap = std::map; Connection* NewConn(const ConnIDKey& k, double t, const ConnID* id, const u_char* data, int proto, uint32_t flow_label, @@ -221,11 +222,10 @@ protected: // the new one. void InsertConnection(ConnectionMap* m, const ConnIDKey& key, Connection* conn); - CompositeHash* ch; ConnectionMap tcp_conns; ConnectionMap udp_conns; ConnectionMap icmp_conns; - PDict fragments; + FragmentMap fragments; SessionStats stats;