Add options to limit extracted file sizes w/ 100MB default.

src/file_analysis/AnalyzerSet.cc

@@ -35,6 +35,14 @@ AnalyzerSet::~AnalyzerSet()
 	delete analyzer_hash;
 	}
 
+Analyzer* AnalyzerSet::Find(file_analysis::Tag tag, RecordVal* args)
+	{
+	HashKey* key = GetKey(tag, args);
+	Analyzer* rval = analyzer_map.Lookup(key);
+	delete key;
+	return rval;
+	}
+
 bool AnalyzerSet::Add(file_analysis::Tag tag, RecordVal* args)
 	{
 	HashKey* key = GetKey(tag, args);

src/file_analysis/AnalyzerSet.h

@@ -37,6 +37,14 @@ public:
 	 */
 	~AnalyzerSet();
 
+	/**
+	 * Looks up an analyzer by its tag and arguments.
+	 * @param tag an analyzer tag.
+	 * @param args an \c AnalyzerArgs record.
+	 * @return pointer to an analyzer instance, or a null pointer if not found.
+	 */
+	Analyzer* Find(file_analysis::Tag tag, RecordVal* args);
+
 	/**
 	 * Attach an analyzer to #file immediately.
 	 * @param tag the analyzer tag of the file analyzer to add.

src/file_analysis/File.cc

@@ -14,6 +14,8 @@
 #include "analyzer/Analyzer.h"
 #include "analyzer/Manager.h"
 
+#include "analyzer/extract/Extract.h"
+
 using namespace file_analysis;
 
 static Val* empty_connection_table()
@@ -203,6 +205,22 @@ void File::SetTimeoutInterval(double interval)
 	val->Assign(timeout_interval_idx, new Val(interval, TYPE_INTERVAL));
 	}
 
+bool File::SetExtractionLimit(RecordVal* args, uint64 bytes)
+	{
+	Analyzer* a = analyzers.Find(file_mgr->GetComponentTag("EXTRACT"), args);
+
+	if ( ! a )
+		return false;
+
+	Extract* e = dynamic_cast<Extract*>(a);
+
+	if ( ! e )
+		return false;
+
+	e->SetLimit(bytes);
+	return true;
+	}
+
 void File::IncrementByteCount(uint64 size, int field_idx)
 	{
 	uint64 old = LookupFieldDefaultCount(field_idx);
@@ -458,7 +476,7 @@ void File::FileEvent(EventHandlerPtr h, val_list* vl)
 			}
 		}
 
-	if ( h == file_new || h == file_timeout )
+	if ( h == file_new || h == file_timeout || h == file_extraction_limit )
 		{
 		// immediate feedback is required for these events.
 		mgr.Drain();

src/file_analysis/File.h

@@ -56,6 +56,14 @@ public:
 	 */
 	void SetTimeoutInterval(double interval);
 
+	/**
+	 * Change the maximum size that an attached extraction analyzer is allowed.
+	 * @param args the file extraction analyzer whose limit needs changed.
+	 * @param bytes new limit.
+	 * @return false if no extraction analyzer is active, else true.
+	 */
+	bool SetExtractionLimit(RecordVal* args, uint64 bytes);
+
 	/**
 	 * @return value of the "id" field from #val record.
 	 */

src/file_analysis/Manager.cc

@@ -184,6 +184,17 @@ bool Manager::SetTimeoutInterval(const string& file_id, double interval) const
 	return true;
 	}
 
+bool Manager::SetExtractionLimit(const string& file_id, RecordVal* args,
+                                 uint64 n) const
+	{
+	File* file = LookupFile(file_id);
+
+	if ( ! file )
+		return false;
+
+	return file->SetExtractionLimit(args, n);
+	}
+
 bool Manager::AddAnalyzer(const string& file_id, file_analysis::Tag tag,
                           RecordVal* args) const
 	{

src/file_analysis/Manager.h

@@ -173,6 +173,19 @@ public:
 	 */
 	bool SetTimeoutInterval(const string& file_id, double interval) const;
 
+	/**
+	 * Sets a limit on the maximum size allowed for extracting the file
+	 * to local disk;
+	 * @param file_id the file identifier/hash.
+	 * @param args a \c AnalyzerArgs value which describes a file analyzer,
+	 *        which should be a file extraction analyzer.
+	 * @param n the new extraction limit, in bytes.
+	 * @return false if file identifier and analyzer did not map to anything,
+	 *         else true.
+	 */
+	bool SetExtractionLimit(const string& file_id, RecordVal* args,
+	                        uint64 n) const;
+
 	/**
 	 * Queue attachment of an analzer to the file identifier.  Multiple
 	 * analyzers of a given type can be attached per file identifier at a time

src/file_analysis/analyzer/extract/CMakeLists.txt

@@ -5,4 +5,6 @@ include_directories(BEFORE ${CMAKE_CURRENT_SOURCE_DIR}
 
 bro_plugin_begin(Bro FileExtract)
 bro_plugin_cc(Extract.cc Plugin.cc ../../Analyzer.cc)
+bro_plugin_bif(events.bif)
+bro_plugin_bif(functions.bif)
 bro_plugin_end()

src/file_analysis/analyzer/extract/Extract.cc

@@ -4,13 +4,15 @@
 
 #include "Extract.h"
 #include "util.h"
+#include "Event.h"
 #include "file_analysis/Manager.h"
 
 using namespace file_analysis;
 
-Extract::Extract(RecordVal* args, File* file, const string& arg_filename)
+Extract::Extract(RecordVal* args, File* file, const string& arg_filename,
+                 uint64 arg_limit)
     : file_analysis::Analyzer(file_mgr->GetComponentTag("EXTRACT"), args, file),
-	  filename(arg_filename)
+      filename(arg_filename), limit(arg_limit)
 	{
 	fd = open(filename.c_str(), O_WRONLY | O_CREAT | O_TRUNC, 0666);
 
@@ -29,15 +31,51 @@ Extract::~Extract()
 		safe_close(fd);
 	}
 
-file_analysis::Analyzer* Extract::Instantiate(RecordVal* args, File* file)
+static Val* get_extract_field_val(RecordVal* args, const char* name)
 	{
 	using BifType::Record::Files::AnalyzerArgs;
-	Val* v = args->Lookup(AnalyzerArgs->FieldOffset("extract_filename"));
+	Val* rval = args->Lookup(AnalyzerArgs->FieldOffset(name));
 
-	if ( ! v )
+	if ( ! rval )
+		reporter->Error("File extraction analyzer missing arg field: %s", name);
+
+	return rval;
+	}
+
+file_analysis::Analyzer* Extract::Instantiate(RecordVal* args, File* file)
+	{
+	Val* fname = get_extract_field_val(args, "extract_filename");
+	Val* limit = get_extract_field_val(args, "extract_limit");
+
+	if ( ! fname || ! limit )
 		return 0;
 
-	return new Extract(args, file, v->AsString()->CheckString());
+	return new Extract(args, file, fname->AsString()->CheckString(),
+	                   limit->AsCount());
+	}
+
+static bool check_limit_exceeded(uint64 lim, uint64 off, uint64 len, uint64* n)
+	{
+	if ( lim == 0 )
+		{
+		*n = len;
+		return false;
+		}
+
+	if ( off >= lim )
+		{
+		*n = 0;
+		return true;
+		}
+
+	*n = lim - off;
+
+	if ( len > *n )
+		return true;
+	else
+		*n = len;
+
+	return false;
 	}
 
 bool Extract::DeliverChunk(const u_char* data, uint64 len, uint64 offset)
@@ -45,6 +83,26 @@ bool Extract::DeliverChunk(const u_char* data, uint64 len, uint64 offset)
 	if ( ! fd )
 		return false;
 
-	safe_pwrite(fd, data, len, offset);
-	return true;
+	uint64 towrite = 0;
+	bool limit_exceeded = check_limit_exceeded(limit, offset, len, &towrite);
+
+	if ( limit_exceeded && file_extraction_limit )
+		{
+		File* f = GetFile();
+		val_list* vl = new val_list();
+		vl->append(f->GetVal()->Ref());
+		vl->append(Args()->Ref());
+		vl->append(new Val(limit, TYPE_COUNT));
+		vl->append(new Val(offset, TYPE_COUNT));
+		vl->append(new Val(len, TYPE_COUNT));
+		f->FileEvent(file_extraction_limit, vl);
+
+		// Limit may have been modified by BIF, re-check it.
+		limit_exceeded = check_limit_exceeded(limit, offset, len, &towrite);
+		}
+
+	if ( towrite > 0 )
+		safe_pwrite(fd, data, towrite, offset);
+
+	return ( ! limit_exceeded );
 	}

src/file_analysis/analyzer/extract/Extract.h

@@ -9,6 +9,8 @@
 #include "File.h"
 #include "Analyzer.h"
 
+#include "analyzer/extract/events.bif.h"
+
 namespace file_analysis {
 
 /**
@@ -41,6 +43,13 @@ public:
 	 */
 	static file_analysis::Analyzer* Instantiate(RecordVal* args, File* file);
 
+	/**
+	 * Sets the maximum allowed extracted file size.  A value of zero means
+	 * "no limit".
+	 * @param bytes number of bytes allowed to be extracted
+	 */
+	void SetLimit(uint64 bytes) { limit = bytes; }
+
 protected:
 
 	/**
@@ -49,12 +58,15 @@ protected:
 	 * @param file the file to which the analyzer will be attached.
 	 * @param arg_filename a file system path which specifies the local file
 	 *        to which the contents of the file will be extracted/written.
+	 * @param arg_limit the maximum allowed file size.
 	 */
-	Extract(RecordVal* args, File* file, const string& arg_filename);
+	Extract(RecordVal* args, File* file, const string& arg_filename,
+	        uint64 arg_limit);
 
 private:
 	string filename;
 	int fd;
+	uint64 limit;
 };
 
 } // namespace file_analysis

src/file_analysis/analyzer/extract/Plugin.cc

@@ -18,6 +18,11 @@ protected:
 
 		AddComponent(new ::file_analysis::Component("EXTRACT",
 		        ::file_analysis::Extract::Instantiate));
+
+		extern std::list<std::pair<const char*, int> > __bif_events_init();
+		AddBifInitFunction(&__bif_events_init);
+		extern std::list<std::pair<const char*, int> > __bif_functions_init();
+		AddBifInitFunction(&__bif_functions_init);
 		}
 };
 

src/file_analysis/analyzer/extract/events.bif

@@ -0,0 +1,19 @@
+## This event is generated when a file extraction analyzer is about
+## to exceed the maximum permitted file size allowed by
+## *extract_size_limit* field of :bro:see:`Files::AnalyzerArgs`.
+## The analyzer is automatically removed from file *f*.
+##
+## f: The file.
+##
+## args: Arguments that identify a particular file extraction analyzer.
+##       This is only provided to be able to pass along to
+##       :bro:see:`FileExtract::set_limit`.
+##
+## limit: The limit, in bytes, the extracted file is about to breach.
+##
+## offset: The offset at which a file chunk is about to be written.
+##
+## len:: The length of the file chunk about to be written.
+##
+## .. bro:see:: Files::add_analyzer Files::ANALYZER_EXTRACT
+event file_extraction_limit%(f: fa_file, args: any, limit: count, offset: count, len: count%);

src/file_analysis/analyzer/extract/functions.bif

@@ -0,0 +1,19 @@
+##! Internal functions used by the extraction file analyzer.
+
+module FileExtract;
+
+%%{
+#include "file_analysis/Manager.h"
+%%}
+
+## :bro:see:`FileExtract::set_limit`.
+function FileExtract::__set_limit%(file_id: string, args: any, n: count%): bool
+    %{
+	using BifType::Record::Files::AnalyzerArgs;
+	RecordVal* rv = args->AsRecordVal()->CoerceTo(AnalyzerArgs);
+    bool result = file_mgr->SetExtractionLimit(file_id->CheckString(), rv, n);
+	Unref(rv);
+    return new Val(result, TYPE_BOOL);
+    %}
+
+module GLOBAL;