From: Bhargava Jandhyala (bjandhya) Date: Tue, 31 May 2022 11:12:31 +0000 (+0000) Subject: Pull request #3444: dce_rpc: converting tree tracker to shared ptr X-Git-Tag: 3.1.31.0~6 X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=6dbefae00b116463e0193c3eb0f42f1cccef27af;p=thirdparty%2Fsnort3.git Pull request #3444: dce_rpc: converting tree tracker to shared ptr Merge in SNORT/snort3 from ~UMUNNIKR/snort3:tree_tracker_shared_ptr to master Squashed commit of the following: commit 7e04875cd7ad8cb7122469b985fe8f02575dba4d Author: Unnikrishnan M Date: Tue May 17 12:11:41 2022 +0530 dce_rpc: converting tree tracker to shared ptr --- diff --git a/src/service_inspectors/dce_rpc/dce_co.cc b/src/service_inspectors/dce_rpc/dce_co.cc index 69b65225e..5717f7b51 100644 --- a/src/service_inspectors/dce_rpc/dce_co.cc +++ b/src/service_inspectors/dce_rpc/dce_co.cc @@ -1370,7 +1370,7 @@ static Packet* dce_co_reassemble(DCE2_SsnData* sd, DCE2_CoTracker* cot, DCE2_RpktType rpkt_type; Packet* rpkt = DCE2_CoGetRpkt(sd, cot, co_rtype, &rpkt_type); - if (rpkt == nullptr) + if (!rpkt || !rpkt->data) { return nullptr; } @@ -2246,7 +2246,7 @@ static Packet* DCE2_CoGetSegRpkt(DCE2_SsnData* sd, case DCE2_TRANS_TYPE__SMB: rpkt = DCE2_GetRpkt(p, DCE2_RPKT_TYPE__SMB_CO_SEG, data_ptr, data_len); - if ( !rpkt ) + if ( !rpkt || !rpkt->data ) return nullptr; set_smb_reassembled_data(const_cast(rpkt->data), diff --git a/src/service_inspectors/dce_rpc/dce_smb2_file.cc b/src/service_inspectors/dce_rpc/dce_smb2_file.cc index 1ab85f677..1d4a2859d 100644 --- a/src/service_inspectors/dce_rpc/dce_smb2_file.cc +++ b/src/service_inspectors/dce_rpc/dce_smb2_file.cc @@ -303,6 +303,6 @@ Dce2Smb2FileTracker::~Dce2Smb2FileTracker(void) snort_free((void*)file_name); file_name = nullptr; - parent_tree = nullptr; + parent_tree.reset(); } diff --git a/src/service_inspectors/dce_rpc/dce_smb2_file.h b/src/service_inspectors/dce_rpc/dce_smb2_file.h index 4aed699bf..efb26aff8 100644 --- a/src/service_inspectors/dce_rpc/dce_smb2_file.h +++ b/src/service_inspectors/dce_rpc/dce_smb2_file.h @@ -27,6 +27,7 @@ #include class Dce2Smb2TreeTracker; +using Dce2Smb2TreeTrackerPtr = std::shared_ptr; typedef struct _tcp_flow_state { @@ -43,7 +44,7 @@ public: Dce2Smb2FileTracker(const Dce2Smb2FileTracker& arg) = delete; Dce2Smb2FileTracker& operator=(const Dce2Smb2FileTracker& arg) = delete; - Dce2Smb2FileTracker(uint64_t file_idv, const uint32_t flow_key, Dce2Smb2TreeTracker* p_tree, + Dce2Smb2FileTracker(uint64_t file_idv, const uint32_t flow_key, Dce2Smb2TreeTrackerPtr p_tree, uint64_t sid) : ignore(true), file_name_len(0), file_flow_key(flow_key), file_id(file_idv), file_size(0), file_name_hash(0), file_name(nullptr), @@ -67,8 +68,8 @@ public: void stop_accepting_raw_data_from(uint32_t); void set_direction(FileDirection dir) { direction = dir; } - Dce2Smb2TreeTracker* get_parent() { return parent_tree; } - void set_parent(Dce2Smb2TreeTracker* pt) { parent_tree = pt; } + Dce2Smb2TreeTrackerPtr get_parent() { return parent_tree; } + void set_parent(Dce2Smb2TreeTrackerPtr pt) { parent_tree = pt; } uint64_t get_file_id() { return file_id; } uint64_t get_file_name_hash() { return file_name_hash; } uint64_t get_session_id() { return session_id; } @@ -89,7 +90,7 @@ private: uint64_t file_name_hash; char* file_name; FileDirection direction; - Dce2Smb2TreeTracker* parent_tree; + Dce2Smb2TreeTrackerPtr parent_tree; std::unordered_map > flow_state; uint64_t session_id; std::mutex process_file_mutex; diff --git a/src/service_inspectors/dce_rpc/dce_smb2_session.cc b/src/service_inspectors/dce_rpc/dce_smb2_session.cc index 55b25f0a3..d2d2c6786 100644 --- a/src/service_inspectors/dce_rpc/dce_smb2_session.cc +++ b/src/service_inspectors/dce_rpc/dce_smb2_session.cc @@ -40,7 +40,7 @@ Dce2Smb2SessionData* Dce2Smb2SessionTracker::get_flow(uint32_t flow_key) return (it_flow != attached_flows.end()) ? it_flow->second : nullptr; } -Dce2Smb2TreeTracker* Dce2Smb2SessionTracker::find_tree_for_message( +Dce2Smb2TreeTrackerPtr Dce2Smb2SessionTracker::find_tree_for_message( const uint64_t message_id, const uint32_t flow_key) { std::lock_guard guard(connected_trees_mutex); @@ -53,10 +53,20 @@ Dce2Smb2TreeTracker* Dce2Smb2SessionTracker::find_tree_for_message( return nullptr; } +Dce2Smb2TreeTrackerPtr Dce2Smb2SessionTracker::find_tree_for_tree_id( + const uint32_t tree_id) +{ + std::lock_guard guard(connected_trees_mutex); + auto it_tree = connected_trees.find(tree_id); + if (it_tree != connected_trees.end()) + return it_tree->second; + return nullptr; +} + void Dce2Smb2SessionTracker::process(const uint16_t command, uint8_t command_type, const Smb2Hdr* smb_header, const uint8_t* end, const uint32_t current_flow_key) { - Dce2Smb2TreeTracker* tree = nullptr; + Dce2Smb2TreeTrackerPtr tree; uint32_t tree_id = Smb2Tid(smb_header); if (tree_id) @@ -121,7 +131,7 @@ void Dce2Smb2SessionTracker::process(const uint16_t command, uint8_t command_typ } } -Dce2Smb2TreeTracker* Dce2Smb2SessionTracker::connect_tree(const uint32_t tree_id, +Dce2Smb2TreeTrackerPtr Dce2Smb2SessionTracker::connect_tree(const uint32_t tree_id, const uint32_t current_flow_key, const uint8_t share_type) { Dce2Smb2SessionData* current_flow = get_flow(current_flow_key); @@ -135,7 +145,7 @@ Dce2Smb2TreeTracker* Dce2Smb2SessionTracker::connect_tree(const uint32_t tree_id dce2_smb_stats.v2_tree_cnct_ignored++; return nullptr; } - Dce2Smb2TreeTracker* tree = nullptr; + Dce2Smb2TreeTrackerPtr tree; connected_trees_mutex.lock(); auto it_tree = connected_trees.find(tree_id); if (it_tree != connected_trees.end()) @@ -143,7 +153,7 @@ Dce2Smb2TreeTracker* Dce2Smb2SessionTracker::connect_tree(const uint32_t tree_id connected_trees_mutex.unlock(); if (!tree) { - tree = new Dce2Smb2TreeTracker(tree_id, this, share_type); + tree = std::make_shared(tree_id, this, share_type); connected_trees_mutex.lock(); connected_trees.insert(std::make_pair(tree_id, tree)); connected_trees_mutex.unlock(); @@ -206,20 +216,17 @@ Dce2Smb2SessionTracker::~Dce2Smb2SessionTracker(void) return; } - std::vector all_trees; connected_trees_mutex.lock(); auto it_tree = connected_trees.begin(); while (it_tree != connected_trees.end()) { - all_trees.push_back(it_tree->second); - it_tree = connected_trees.erase(it_tree); + auto next_it_tree = std::next(it_tree); + it_tree->second->close_all_files(); + disconnect_tree(it_tree->second->get_tree_id()); + it_tree = next_it_tree; } connected_trees_mutex.unlock(); - for (Dce2Smb2TreeTracker* tree : all_trees) - { - delete tree; - } do_not_delete = false; fcfs_mutex.unlock(); } diff --git a/src/service_inspectors/dce_rpc/dce_smb2_session.h b/src/service_inspectors/dce_rpc/dce_smb2_session.h index 8a813e105..aa7c4ae44 100644 --- a/src/service_inspectors/dce_rpc/dce_smb2_session.h +++ b/src/service_inspectors/dce_rpc/dce_smb2_session.h @@ -45,11 +45,10 @@ public: } ~Dce2Smb2SessionTracker(); - Dce2Smb2TreeTracker* connect_tree(const uint32_t, const uint32_t, + Dce2Smb2TreeTrackerPtr connect_tree(const uint32_t, const uint32_t, uint8_t=SMB2_SHARE_TYPE_DISK); void disconnect_tree(uint32_t tree_id) { - std::lock_guard guard(connected_trees_mutex); connected_trees.erase(tree_id); decrease_size(sizeof(Dce2Smb2TreeTracker)); } @@ -90,12 +89,13 @@ public: dce2_smb_stats.total_encrypted_sessions++; } bool get_encryption_flag() { return encryption_flag; } + Dce2Smb2TreeTrackerPtr find_tree_for_tree_id(const uint32_t); private: // do_not_delete is to make sure when we are in processing we should not delete the context // which is being processed bool do_not_delete; bool file_context_cleaned; - Dce2Smb2TreeTracker* find_tree_for_message(const uint64_t, const uint32_t); + Dce2Smb2TreeTrackerPtr find_tree_for_message(const uint64_t, const uint32_t); uint64_t session_id; //to keep the tab of previous command uint16_t command_prev; diff --git a/src/service_inspectors/dce_rpc/dce_smb2_tree.cc b/src/service_inspectors/dce_rpc/dce_smb2_tree.cc index 2a7bb73e4..4d5c45160 100644 --- a/src/service_inspectors/dce_rpc/dce_smb2_tree.cc +++ b/src/service_inspectors/dce_rpc/dce_smb2_tree.cc @@ -40,8 +40,11 @@ uint64_t Smb2Mid(const Smb2Hdr* hdr) Dce2Smb2FileTrackerPtr Dce2Smb2TreeTracker::open_file(const uint64_t file_id, const uint32_t current_flow_key) { + Dce2Smb2TreeTrackerPtr tree_ptr = parent_session->find_tree_for_tree_id(tree_id); + if (!tree_ptr) + return nullptr; std::shared_ptr ftracker = std::make_shared ( - file_id, current_flow_key, this, this->get_parent()->get_session_id()); + file_id, current_flow_key, tree_ptr, this->get_parent()->get_session_id()); tree_tracker_mutex.lock(); opened_files.insert(std::make_pair(file_id, ftracker)); tree_tracker_mutex.unlock(); @@ -65,7 +68,7 @@ void Dce2Smb2TreeTracker::close_file(uint64_t file_id, bool destroy) if (it_file != opened_files.end()) { Dce2Smb2FileTrackerPtr file = it_file->second; - it_file->second->set_parent(nullptr); + it_file->second->get_parent().reset(); if (opened_files.erase(file_id) and destroy) { parent_session->decrease_size(sizeof(Dce2Smb2FileTracker)); @@ -76,6 +79,19 @@ void Dce2Smb2TreeTracker::close_file(uint64_t file_id, bool destroy) tree_tracker_mutex.unlock(); } +void Dce2Smb2TreeTracker::close_all_files() +{ + tree_tracker_mutex.lock(); + auto it_file = opened_files.begin(); + while (it_file != opened_files.end()) + { + get_parent()->clean_file_context_from_flow(it_file->second->get_file_id(), + it_file->second->get_file_name_hash()); + it_file = opened_files.erase(it_file); + } + tree_tracker_mutex.unlock(); +} + Dce2Smb2RequestTracker* Dce2Smb2TreeTracker::find_request(const uint64_t message_id, const uint32_t current_flow_key) { @@ -603,12 +619,11 @@ Dce2Smb2TreeTracker::~Dce2Smb2TreeTracker(void) { get_parent()->clean_file_context_from_flow(it_file.second->get_file_id(), it_file.second->get_file_name_hash()); - it_file.second->set_parent(nullptr); + it_file.second->get_parent().reset(); parent_session->decrease_size(sizeof(Dce2Smb2FileTracker)); } tree_tracker_mutex.unlock(); - parent_session->disconnect_tree(tree_id); } diff --git a/src/service_inspectors/dce_rpc/dce_smb2_tree.h b/src/service_inspectors/dce_rpc/dce_smb2_tree.h index d679d6d0a..0df0b3f4b 100644 --- a/src/service_inspectors/dce_rpc/dce_smb2_tree.h +++ b/src/service_inspectors/dce_rpc/dce_smb2_tree.h @@ -60,6 +60,7 @@ public: Dce2Smb2FileTrackerPtr open_file(const uint64_t, const uint32_t); void close_file(uint64_t, bool); + void close_all_files(); Dce2Smb2FileTrackerPtr find_file(uint64_t); Dce2Smb2RequestTracker* find_request(const uint64_t, const uint32_t); void process(uint16_t, uint8_t, const Smb2Hdr*, const uint8_t*, const uint32_t); @@ -98,8 +99,9 @@ private: std::mutex tree_tracker_mutex; }; +using Dce2Smb2TreeTrackerPtr = std::shared_ptr; using Dce2Smb2TreeTrackerMap = - std::unordered_map >; + std::unordered_map >; #endif