diff --git a/src/file_analysis/Manager.cc b/src/file_analysis/Manager.cc index 5ff7bd7186..3f04ebfc2b 100644 --- a/src/file_analysis/Manager.cc +++ b/src/file_analysis/Manager.cc @@ -54,8 +54,11 @@ void Manager::Terminate() { vector keys; - for ( IDMap::iterator it = id_map.begin(); it != id_map.end(); ++it ) - keys.push_back(it->first); + IterCookie* it = id_map.InitForIteration(); + HashKey* key; + + while ( id_map.NextEntry(key, it) ) + keys.push_back(static_cast(key->Key())); for ( size_t i = 0; i < keys.size(); ++i ) Timeout(keys[i], true); @@ -249,11 +252,12 @@ File* Manager::GetFile(const string& file_id, Connection* conn, if ( IsIgnored(file_id) ) return 0; - File* rval = id_map[file_id]; + File* rval = id_map.Lookup(file_id.c_str()); if ( ! rval ) { - rval = id_map[file_id] = new File(file_id, conn, tag, is_orig); + rval = new File(file_id, conn, tag, is_orig); + id_map.Insert(file_id.c_str(), rval); rval->ScheduleInactivityTimer(); if ( IsIgnored(file_id) ) @@ -272,12 +276,7 @@ File* Manager::GetFile(const string& file_id, Connection* conn, File* Manager::LookupFile(const string& file_id) const { - IDMap::const_iterator it = id_map.find(file_id); - - if ( it == id_map.end() ) - return 0; - - return it->second; + return id_map.Lookup(file_id.c_str()); } void Manager::Timeout(const string& file_id, bool is_terminating) @@ -308,37 +307,38 @@ void Manager::Timeout(const string& file_id, bool is_terminating) bool Manager::IgnoreFile(const string& file_id) { - if ( id_map.find(file_id) == id_map.end() ) + if ( ! id_map.Lookup(file_id.c_str()) ) return false; DBG_LOG(DBG_FILE_ANALYSIS, "Ignore FileID %s", file_id.c_str()); - ignored.insert(file_id); - + delete ignored.Insert(file_id.c_str(), new bool); return true; } bool Manager::RemoveFile(const string& file_id) { - IDMap::iterator it = id_map.find(file_id); + HashKey key(file_id.c_str()); + // Can't remove from the dictionary/map right away as invoking EndOfFile + // may cause some events to be executed which actually depend on the file + // still being in the dictionary/map. + File* f = static_cast(id_map.Lookup(&key)); - if ( it == id_map.end() ) + if ( ! f ) return false; DBG_LOG(DBG_FILE_ANALYSIS, "Remove FileID %s", file_id.c_str()); - it->second->EndOfFile(); - - delete it->second; - id_map.erase(file_id); - ignored.erase(file_id); - + f->EndOfFile(); + delete f; + id_map.Remove(&key); + delete static_cast(ignored.Remove(&key)); return true; } bool Manager::IsIgnored(const string& file_id) { - return ignored.find(file_id) != ignored.end(); + return ignored.Lookup(file_id.c_str()) != 0; } string Manager::GetFileID(analyzer::Tag tag, Connection* c, bool is_orig) diff --git a/src/file_analysis/Manager.h b/src/file_analysis/Manager.h index bb6aaab971..2137e81389 100644 --- a/src/file_analysis/Manager.h +++ b/src/file_analysis/Manager.h @@ -4,10 +4,9 @@ #define FILE_ANALYSIS_MANAGER_H #include -#include -#include #include +#include "Dict.h" #include "Net.h" #include "Conn.h" #include "Val.h" @@ -27,6 +26,9 @@ namespace file_analysis { +declare(PDict,bool); +declare(PDict,File); + /** * Main entry point for interacting with file analysis. */ @@ -288,8 +290,8 @@ public: protected: friend class FileTimer; - typedef set IDSet; - typedef map IDMap; + typedef PDict(bool) IDSet; + typedef PDict(File) IDMap; /** * Create a new file to be analyzed or retrieve an existing one. @@ -361,8 +363,8 @@ protected: private: - IDMap id_map; /**< Map file ID to file_analysis::File records. */ - IDSet ignored; /**< Ignored files. Will be finally removed on EOF. */ + PDict(File) id_map; /**< Map file ID to file_analysis::File records. */ + PDict(bool) ignored; /**< Ignored files. Will be finally removed on EOF. */ string current_file_id; /**< Hash of what get_file_handle event sets. */ RuleFileMagicState* magic_state; /**< File magic signature match state. */