diff --git a/src/broker/Manager.cc b/src/broker/Manager.cc index 18403b9f9c..b9c95125e1 100644 --- a/src/broker/Manager.cc +++ b/src/broker/Manager.cc @@ -1608,15 +1608,66 @@ StoreHandleVal* Manager::MakeMaster(const string& name, broker::backend type, iosource_mgr->RegisterFd(handle->proxy.mailbox().descriptor(), this); CheckForwarding(name); - if ( bstate->endpoint.use_real_time() ) - return handle; + if ( ! bstate->endpoint.use_real_time() ) + // Wait for master to become available/responsive. + // Possibly avoids timeouts in scripts during unit tests. + handle->store.exists(""); + + BrokerStoreToZeekTable(name, handle); - // Wait for master to become available/responsive. - // Possibly avoids timeouts in scripts during unit tests. - handle->store.exists(""); return handle; } +void Manager::BrokerStoreToZeekTable(const std::string& name, const StoreHandleVal* handle) + { + // consider if it might be wise to disable &on_change while filling the table + if ( ! handle->forward_to ) + return; + + auto keys = handle->store.keys(); + if ( ! keys ) + return; + + auto set = caf::get_if(&(keys->get_data())); + auto table = handle->forward_to; + const auto& its = table->GetType()->AsTableType()->IndexTypes(); + bool is_set = table->GetType()->IsSet(); + assert( its.size() == 1 ); + for ( const auto key : *set ) + { + auto zeek_key = data_to_val(key, its[0].get()); + if ( ! zeek_key ) + { + reporter->Error("Failed to convert key %s while importing broker store to table for store %s. Aborting import.", to_string(key).c_str(), name.c_str()); + // just abort - this probably means the types are incompatible + return; + } + + if ( is_set ) + { + table->Assign(zeek_key, nullptr, false); + continue; + } + + auto value = handle->store.get(key); + if ( ! value ) + { + reporter->Error("Failed to load value for key %s while importing broker store %s to table", to_string(key).c_str(), name.c_str()); + continue; + } + + auto zeek_value = data_to_val(*value, table->GetType()->Yield().get()); + if ( ! zeek_value ) + { + reporter->Error("Could not convert %s to table value while trying to import broker store %s. Aborting import.", to_string(value).c_str(), name.c_str()); + return; + } + + table->Assign(zeek_key, zeek_value, false); + } + return; + } + StoreHandleVal* Manager::MakeClone(const string& name, double resync_interval, double stale_interval, double mutation_buffer_interval) diff --git a/src/broker/Manager.h b/src/broker/Manager.h index 6cfcc05a9b..9ccefde4b0 100644 --- a/src/broker/Manager.h +++ b/src/broker/Manager.h @@ -355,7 +355,11 @@ private: void ProcessError(broker::error err); void ProcessStoreResponse(StoreHandleVal*, broker::store::response response); void FlushPendingQueries(); + // Check if a broker store is associated to a table on the Zeek side. void CheckForwarding(const std::string& name); + // Send the content of a broker store to the backing table. This is typically used + // when a master/clone is created. + void BrokerStoreToZeekTable(const std::string& name, const StoreHandleVal* handle); void Error(const char* format, ...) __attribute__((format (printf, 2, 3))); diff --git a/testing/btest/Baseline/broker.store.brokerstore-attr-persistence/output1 b/testing/btest/Baseline/broker.store.brokerstore-attr-persistence/output1 new file mode 100644 index 0000000000..1710727ce2 --- /dev/null +++ b/testing/btest/Baseline/broker.store.brokerstore-attr-persistence/output1 @@ -0,0 +1,20 @@ +{ +[b] = 3, +[whatever] = 5, +[a] = 5 +} +{ +I am really a set!, +Believe me - I am a set, +I am a set! +} +{ +[b] = [a=2, b=d, c={ +elem1, +elem2 +}], +[a] = [a=1, b=c, c={ +elem1, +elem2 +}] +} diff --git a/testing/btest/Baseline/broker.store.brokerstore-attr-persistence/output2 b/testing/btest/Baseline/broker.store.brokerstore-attr-persistence/output2 new file mode 100644 index 0000000000..1710727ce2 --- /dev/null +++ b/testing/btest/Baseline/broker.store.brokerstore-attr-persistence/output2 @@ -0,0 +1,20 @@ +{ +[b] = 3, +[whatever] = 5, +[a] = 5 +} +{ +I am really a set!, +Believe me - I am a set, +I am a set! +} +{ +[b] = [a=2, b=d, c={ +elem1, +elem2 +}], +[a] = [a=1, b=c, c={ +elem1, +elem2 +}] +} diff --git a/testing/btest/broker/store/brokerstore-attr-persistence.zeek b/testing/btest/broker/store/brokerstore-attr-persistence.zeek new file mode 100644 index 0000000000..5004b12b91 --- /dev/null +++ b/testing/btest/broker/store/brokerstore-attr-persistence.zeek @@ -0,0 +1,80 @@ +# @TEST-PORT: BROKER_PORT + +# @TEST-EXEC: zeek -B broker -b one.zeek > output1 +# @TEST-EXEC: zeek -B broker -b two.zeek > output2 +# @TEST-EXEC: btest-diff output1 +# @TEST-EXEC: btest-diff output2 +# @TEST-EXEC: diff output1 output2 + +# the first test writes out the sqlite files... + +@TEST-START-FILE one.zeek + +module TestModule; + +global tablestore: opaque of Broker::Store; +global setstore: opaque of Broker::Store; +global recordstore: opaque of Broker::Store; + +type testrec: record { + a: count; + b: string; + c: set[string]; +}; + +global t: table[string] of count &broker_store="table"; +global s: set[string] &broker_store="set"; +global r: table[string] of testrec &broker_store="rec"; + +event zeek_init() + { + tablestore = Broker::create_master("table", Broker::SQLITE); + setstore = Broker::create_master("set", Broker::SQLITE); + recordstore = Broker::create_master("rec", Broker::SQLITE); + t["a"] = 5; + t["b"] = 3; + t["c"] = 4; + t["whatever"] = 5; + delete t["c"]; + add s["I am a set!"]; + add s["I am really a set!"]; + add s["Believe me - I am a set"]; + r["a"] = testrec($a=1, $b="b", $c=set("elem1", "elem2")); + r["a"] = testrec($a=1, $b="c", $c=set("elem1", "elem2")); + r["b"] = testrec($a=2, $b="d", $c=set("elem1", "elem2")); + print t; + print s; + print r; + } + +@TEST-END-FILE +@TEST-START-FILE two.zeek + +# the second one reads them in again + +module TestModule; + +global tablestore: opaque of Broker::Store; +global setstore: opaque of Broker::Store; +global recordstore: opaque of Broker::Store; + +type testrec: record { + a: count; + b: string; + c: set[string]; +}; + +global t: table[string] of count &broker_store="table"; +global s: set[string] &broker_store="set"; +global r: table[string] of testrec &broker_store="rec"; + +event zeek_init() + { + tablestore = Broker::create_master("table", Broker::SQLITE); + setstore = Broker::create_master("set", Broker::SQLITE); + recordstore = Broker::create_master("rec", Broker::SQLITE); + print t; + print s; + print r; + } +@TEST-END-FILE