From: Bhargava Jandhyala (bjandhya) Date: Thu, 7 Oct 2021 14:25:22 +0000 (+0000) Subject: Merge pull request #3088 in SNORT/snort3 from ~BSACHDEV/snort3:stress_smb2 to master X-Git-Tag: 3.1.15.0~18 X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=9655f03e47e13b46719cfcc81f61c860296d9f6a;p=thirdparty%2Fsnort3.git Merge pull request #3088 in SNORT/snort3 from ~BSACHDEV/snort3:stress_smb2 to master Squashed commit of the following: commit 4de99c141ba599b04b6bb43fbb6af18b63ae836a Author: Bhargava Jandhyala (bjandhya) Date: Mon Oct 4 14:05:53 2021 +0000 dce_smb: Optimised handling pruning of flows in stress environment --- diff --git a/src/file_api/file_flows.cc b/src/file_api/file_flows.cc index df7ad70ec..ac49d75b9 100644 --- a/src/file_api/file_flows.cc +++ b/src/file_api/file_flows.cc @@ -150,7 +150,7 @@ FilePolicyBase* FileFlows::get_file_policy(Flow* flow) void FileFlows::set_current_file_context(FileContext* ctx) { // If we finished processing a file context object last time, delete it - if (current_context_delete_pending) + if (current_context_delete_pending and (current_context != ctx)) { delete current_context; current_context_delete_pending = false; diff --git a/src/service_inspectors/dce_rpc/dce_smb2.cc b/src/service_inspectors/dce_rpc/dce_smb2.cc index 5fff3e76e..05ae281c9 100644 --- a/src/service_inspectors/dce_rpc/dce_smb2.cc +++ b/src/service_inspectors/dce_rpc/dce_smb2.cc @@ -186,6 +186,9 @@ void Dce2Smb2SessionData::process_command(const Smb2Hdr* smb_hdr, { const uint8_t* smb_data = (const uint8_t*)smb_hdr + SMB2_HEADER_LENGTH; uint16_t structure_size = alignedNtohs((const uint16_t*)smb_data); + uint16_t command = alignedNtohs(&(smb_hdr->command)); + uint64_t session_id = Smb2Sid(smb_hdr); + Dce2Smb2SessionTrackerPtr session = find_session(session_id); // Macro and shorthand to save some repetitive code // Should only be used in this function @@ -206,6 +209,11 @@ void Dce2Smb2SessionData::process_command(const Smb2Hdr* smb_hdr, SMB_DEBUG(dce_smb_trace, DEFAULT_TRACE_OPTION_ID, TRACE_ERROR_LEVEL, GET_CURRENT_PACKET, \ "%s : smb data beyond end detected\n", \ smb2_command_string[command]); \ + if (session) \ + { \ + session->set_do_not_delete(false); \ + session->set_prev_comand(SMB2_COM_MAX); \ + } \ return; \ } \ } @@ -217,6 +225,11 @@ void Dce2Smb2SessionData::process_command(const Smb2Hdr* smb_hdr, dce2_smb_stats.v2_ ## counter ## _err_resp++; \ SMB_DEBUG(dce_smb_trace, DEFAULT_TRACE_OPTION_ID, TRACE_INFO_LEVEL, GET_CURRENT_PACKET, "%s_RESP: error\n", \ smb2_command_string[command]); \ + if (session) \ + { \ + session->set_do_not_delete(false); \ + session->set_prev_comand(SMB2_COM_MAX); \ + } \ return; \ } \ } @@ -227,19 +240,34 @@ void Dce2Smb2SessionData::process_command(const Smb2Hdr* smb_hdr, SMB_DEBUG(dce_smb_trace, DEFAULT_TRACE_OPTION_ID, TRACE_ERROR_LEVEL \ , GET_CURRENT_PACKET, "%s: invalid struct size\n", \ smb2_command_string[command]); \ + if (session) \ + { \ + session->set_do_not_delete(false); \ + session->set_prev_comand(SMB2_COM_MAX); \ + } \ return; \ } - uint16_t command = alignedNtohs(&(smb_hdr->command)); - uint64_t session_id = Smb2Sid(smb_hdr); SMB_DEBUG(dce_smb_trace, DEFAULT_TRACE_OPTION_ID, TRACE_INFO_LEVEL, GET_CURRENT_PACKET, "%s : flow %" PRIu32 " mid %" PRIu64 " sid %" PRIu64 " tid %" PRIu32 "\n", (command < SMB2_COM_MAX ? smb2_command_string[command] : "unknown"), flow_key, Smb2Mid(smb_hdr), session_id, Smb2Tid(smb_hdr)); + + // Handling case of two threads trying to do close same session at a time + if (command == SMB2_COM_CLOSE and (session and session->get_prev_command() != SMB2_COM_MAX)) + { + session->set_do_not_delete(false); + return; + } + + if (session) + { + session->set_do_not_delete(true); + session->set_prev_comand(command); + } + // Try to find the session. // The case when session is not available will be handled per command. - Dce2Smb2SessionTrackerPtr session = find_session(session_id); - switch (command) { //commands processed by flow @@ -459,6 +487,11 @@ void Dce2Smb2SessionData::process_command(const Smb2Hdr* smb_hdr, dce2_smb_stats.v2_msgs_uninspected++; break; } + if (session) + { + session->set_prev_comand(SMB2_COM_MAX); + session->set_do_not_delete(false); + } } // This is the main entry point for SMB2 processing. @@ -467,6 +500,7 @@ void Dce2Smb2SessionData::process() Packet* p = DetectionEngine::get_current_packet(); const uint8_t* data_ptr = p->data; uint16_t data_len = p->dsize; + Dce2Smb2SessionTrackerPtr session = nullptr; // Process the header if (p->is_pdu_start()) @@ -502,6 +536,8 @@ void Dce2Smb2SessionData::process() SMB_DEBUG(dce_smb_trace, DEFAULT_TRACE_OPTION_ID, TRACE_ERROR_LEVEL, p, "bad next command offset\n"); dce2_smb_stats.v2_bad_next_cmd_offset++; + if (session) + session->set_do_not_delete(false); return; } if (next_command_offset) @@ -516,6 +552,8 @@ void Dce2Smb2SessionData::process() SMB_DEBUG(dce_smb_trace, DEFAULT_TRACE_OPTION_ID, TRACE_INFO_LEVEL, p, "compound request limit" " reached %" PRIu8 "\n",compound_request_index); + if (session) + session->set_do_not_delete(false); return; } } @@ -523,6 +561,12 @@ void Dce2Smb2SessionData::process() } else { + if ( tcp_file_tracker ) + { + session = find_session(tcp_file_tracker->get_session_id()); + if (session) + session->set_do_not_delete(true); + } tcp_file_tracker_mutex.lock(); if ( tcp_file_tracker and tcp_file_tracker->accepting_raw_data_from(flow_key)) { @@ -537,6 +581,8 @@ void Dce2Smb2SessionData::process() } tcp_file_tracker_mutex.unlock(); } + if (session) + session->set_do_not_delete(false); } void Dce2Smb2SessionData::set_reassembled_data(uint8_t* nb_ptr, uint16_t co_len) diff --git a/src/service_inspectors/dce_rpc/dce_smb2_file.h b/src/service_inspectors/dce_rpc/dce_smb2_file.h index 62ae34936..8187dcffb 100644 --- a/src/service_inspectors/dce_rpc/dce_smb2_file.h +++ b/src/service_inspectors/dce_rpc/dce_smb2_file.h @@ -42,10 +42,11 @@ 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, Dce2Smb2TreeTracker* 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), - direction(FILE_DOWNLOAD), parent_tree(p_tree) + direction(FILE_DOWNLOAD), parent_tree(p_tree), session_id(sid) { SMB_DEBUG(dce_smb_trace, DEFAULT_TRACE_OPTION_ID, TRACE_DEBUG_LEVEL, GET_CURRENT_PACKET, "file tracker %" PRIu64 " created\n", file_id); @@ -67,6 +68,7 @@ public: void set_direction(FileDirection dir) { direction = dir; } Dce2Smb2TreeTracker* get_parent() { return parent_tree; } uint64_t get_file_id() { return file_id; } + uint64_t get_session_id() { return session_id; } private: void file_detect(); @@ -81,6 +83,7 @@ private: FileDirection direction; Dce2Smb2TreeTracker* parent_tree; std::unordered_map > flow_state; + uint64_t session_id; std::mutex process_file_mutex; std::mutex flow_state_mutex; }; diff --git a/src/service_inspectors/dce_rpc/dce_smb2_session.cc b/src/service_inspectors/dce_rpc/dce_smb2_session.cc index aa67f7432..05ed9d899 100644 --- a/src/service_inspectors/dce_rpc/dce_smb2_session.cc +++ b/src/service_inspectors/dce_rpc/dce_smb2_session.cc @@ -160,6 +160,7 @@ Dce2Smb2TreeTracker* Dce2Smb2SessionTracker::connect_tree(const uint32_t tree_id void Dce2Smb2SessionTracker::clean_file_context_from_flow(Dce2Smb2FileTracker* file_tracker, uint64_t file_id, uint64_t file_name_hash) { + attached_flows_mutex.lock(); for (auto it_flow : attached_flows) { snort::FileFlows* file_flows = snort::FileFlows::get_file_flows( @@ -168,6 +169,7 @@ void Dce2Smb2SessionTracker::clean_file_context_from_flow(Dce2Smb2FileTracker* f file_flows->remove_processed_file_context(file_name_hash, file_id); it_flow.second->reset_matching_tcp_file_tracker(file_tracker); } + attached_flows_mutex.unlock(); } void Dce2Smb2SessionTracker::increase_size(const size_t size) @@ -191,6 +193,16 @@ void Dce2Smb2SessionTracker::unlink() // Session Tracker is created and destroyed only from session cache Dce2Smb2SessionTracker::~Dce2Smb2SessionTracker() { + if (!(fcfs_mutex.try_lock())) + return; + + if (do_not_delete ) + { + // Dont prune the session in LRU Cache + smb2_session_cache.find_id(get_key()); + fcfs_mutex.unlock(); + return; + } if (smb_module_is_up and (snort::is_packet_thread())) { SMB_DEBUG(dce_smb_trace, DEFAULT_TRACE_OPTION_ID, @@ -212,6 +224,7 @@ Dce2Smb2SessionTracker::~Dce2Smb2SessionTracker() { 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 ce31f10d1..bf6a21d62 100644 --- a/src/service_inspectors/dce_rpc/dce_smb2_session.h +++ b/src/service_inspectors/dce_rpc/dce_smb2_session.h @@ -36,7 +36,9 @@ public: session_id = key.sid; session_key = key; reload_prune = false; - SMB_DEBUG(dce_smb_trace, DEFAULT_TRACE_OPTION_ID, TRACE_DEBUG_LEVEL, GET_CURRENT_PACKET, + do_not_delete = false; + command_prev = SMB2_COM_MAX; + SMB_DEBUG(dce_smb_trace, DEFAULT_TRACE_OPTION_ID, TRACE_DEBUG_LEVEL, GET_CURRENT_PACKET, "session tracker %" PRIu64 "created\n", session_id); } @@ -70,17 +72,30 @@ public: void process(const uint16_t, uint8_t, const Smb2Hdr*, const uint8_t*, const uint32_t); void increase_size(const size_t size); void decrease_size(const size_t size); - void set_reload_prune() { reload_prune = true; } + void set_reload_prune(bool flag) { reload_prune = flag; } + uint64_t get_session_id() { return session_id; } + void set_do_not_delete(bool flag) { do_not_delete = flag; } + bool get_do_not_delete() { return do_not_delete; } + void set_prev_comand(uint16_t cmd) { command_prev = cmd; } + uint16_t get_prev_command() { return command_prev; } 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; Dce2Smb2TreeTracker* 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; Smb2SessionKey session_key; Dce2Smb2SessionDataMap attached_flows; Dce2Smb2TreeTrackerMap connected_trees; std::atomic reload_prune; std::mutex connected_trees_mutex; std::mutex attached_flows_mutex; + // fcfs_mutex is to make sure the mutex is taken at first come first basis if code + // is being hit by two different paths + std::mutex fcfs_mutex; }; #endif diff --git a/src/service_inspectors/dce_rpc/dce_smb2_session_cache.h b/src/service_inspectors/dce_rpc/dce_smb2_session_cache.h index 982a095e9..e82d260bf 100644 --- a/src/service_inspectors/dce_rpc/dce_smb2_session_cache.h +++ b/src/service_inspectors/dce_rpc/dce_smb2_session_cache.h @@ -42,23 +42,25 @@ public: using Data = std::shared_ptr; + Data find_id(Key key) + { + Data session = this->find(key); + return session; + } + Data find_session(Key key, Dce2Smb2SessionData* ssd) { - flow_mutex.lock(); Data session = this->find(key); if (session) session->attach_flow(ssd->get_flow_key(), ssd); - flow_mutex.unlock(); return session; } Data find_else_create_session(Key& key, Dce2Smb2SessionData* ssd) { Data new_session = Data(new Value(key)); - flow_mutex.lock(); Data session = this->find_else_insert(key, new_session, nullptr); session->attach_flow(ssd->get_flow_key(), ssd); - flow_mutex.unlock(); return session; } @@ -93,7 +95,7 @@ public: data.emplace_back(list_iter->second); // increase reference count // This instructs the session_tracker to take a lock before detaching // from ssd, when it is getting destroyed. - list_iter->second->set_reload_prune(); + list_iter->second->set_reload_prune(true); decrease_size(list_iter->second.get()); map.erase(list_iter->first); list.erase(list_iter); @@ -110,7 +112,6 @@ private: using LruBase::max_size; using LruBase::stats; using LruListIter = typename LruBase::LruListIter; - std::mutex flow_mutex; void increase_size(Value* value_ptr=nullptr) override { if (value_ptr) current_size += sizeof(*value_ptr); diff --git a/src/service_inspectors/dce_rpc/dce_smb2_tree.cc b/src/service_inspectors/dce_rpc/dce_smb2_tree.cc index 5dc9e10aa..59b435d4b 100644 --- a/src/service_inspectors/dce_rpc/dce_smb2_tree.cc +++ b/src/service_inspectors/dce_rpc/dce_smb2_tree.cc @@ -39,7 +39,8 @@ uint64_t Smb2Mid(const Smb2Hdr* hdr) Dce2Smb2FileTracker* Dce2Smb2TreeTracker::open_file(const uint64_t file_id, const uint32_t current_flow_key) { - Dce2Smb2FileTracker* ftracker = new Dce2Smb2FileTracker(file_id, current_flow_key, this); + Dce2Smb2FileTracker* ftracker = new Dce2Smb2FileTracker(file_id, current_flow_key, this, + this->get_parent()->get_session_id()); tree_tracker_mutex.lock(); opened_files.insert(std::make_pair(file_id, ftracker)); tree_tracker_mutex.unlock(); diff --git a/src/service_inspectors/dce_rpc/dce_smb_inspector.cc b/src/service_inspectors/dce_rpc/dce_smb_inspector.cc index f2ae621bd..76f192106 100644 --- a/src/service_inspectors/dce_rpc/dce_smb_inspector.cc +++ b/src/service_inspectors/dce_rpc/dce_smb_inspector.cc @@ -27,6 +27,7 @@ #include "dce_smb_module.h" #include "dce_smb_utils.h" #include "dce_smb2_session_cache.h" +#include "main/thread_config.h" #define DCE_SMB_PROTOCOL_ID "netbios-ssn" @@ -139,8 +140,9 @@ static Inspector* dce2_smb_ctor(Module* m) dce2SmbProtoConf config; mod->get_data(config); size_t max_smb_mem = DCE2_ScSmbMemcap(&config); + uint16_t num_threads = ThreadConfig::get_instance_max(); smb_module_is_up = true; - smb2_session_cache.reload_prune(max_smb_mem); + smb2_session_cache.reload_prune(max_smb_mem*num_threads); return new Dce2Smb(config); }