diff --git a/src/storage/backend/sqlite/SQLite.cc b/src/storage/backend/sqlite/SQLite.cc index 79b9328981..e4fc4afd9b 100644 --- a/src/storage/backend/sqlite/SQLite.cc +++ b/src/storage/backend/sqlite/SQLite.cc @@ -42,7 +42,7 @@ ErrorResult SQLite::DoOpen(RecordValPtr options) { } std::string create = "create table if not exists " + table_name + " ("; - create.append("key_str text primary key, value_str text not null);"); + create.append("key_str text primary key, value_str text not null, expire_time real);"); char* errorMsg = nullptr; if ( int res = sqlite3_exec(db, create.c_str(), NULL, NULL, &errorMsg); res != SQLITE_OK ) { @@ -54,12 +54,15 @@ ErrorResult SQLite::DoOpen(RecordValPtr options) { } static std::map statements = - {{"put", util::fmt("insert into %s (key_str, value_str) values(?, ?)", table_name.c_str())}, - {"put_update", util::fmt("insert into %s (key_str, value_str) values(?, ?) ON CONFLICT(key_str) " - "DO UPDATE SET value_str=?", - table_name.c_str())}, + {{"put", util::fmt("insert into %s (key_str, value_str, expire_time) values(?, ?, ?)", table_name.c_str())}, + {"put_update", + util::fmt("insert into %s (key_str, value_str, expire_time) values(?, ?, ?) ON CONFLICT(key_str) " + "DO UPDATE SET value_str=?", + table_name.c_str())}, {"get", util::fmt("select value_str from %s where key_str=?", table_name.c_str())}, - {"erase", util::fmt("delete from %s where key_str=?", table_name.c_str())}}; + {"erase", util::fmt("delete from %s where key_str=?", table_name.c_str())}, + {"expire", util::fmt("delete from %s where expire_time > 0 and expire_time != 0 and expire_time <= ?", + table_name.c_str())}}; for ( const auto& [key, stmt] : statements ) { sqlite3_stmt* ps; @@ -123,8 +126,13 @@ ErrorResult SQLite::DoPut(ValPtr key, ValPtr value, bool overwrite, double expir return res; } + if ( auto res = checkError(sqlite3_bind_double(stmt, 3, expiration_time)); res.has_value() ) { + sqlite3_reset(stmt); + return res; + } + if ( overwrite ) { - if ( auto res = checkError(sqlite3_bind_text(stmt, 3, value_str.data(), value_str.size(), SQLITE_STATIC)); + if ( auto res = checkError(sqlite3_bind_text(stmt, 4, value_str.data(), value_str.size(), SQLITE_STATIC)); res.has_value() ) { sqlite3_reset(stmt); return res; @@ -155,7 +163,7 @@ ValResult SQLite::DoGet(ValPtr key, ValResultCallback* cb) { if ( auto res = checkError(sqlite3_bind_text(stmt, 1, key_str.data(), key_str.size(), SQLITE_STATIC)); res.has_value() ) { sqlite3_reset(stmt); - return nonstd::unexpected(res.value()); + return zeek::unexpected(res.value()); } int errorcode = sqlite3_step(stmt); @@ -200,6 +208,23 @@ ErrorResult SQLite::DoErase(ValPtr key, ErrorResultCallback* cb) { return std::nullopt; } +/** + * Removes any entries in the backend that have expired. Can be overridden by + * derived classes. + */ +void SQLite::Expire() { + auto stmt = prepared_stmts["expire"]; + + if ( auto res = checkError(sqlite3_bind_double(stmt, 1, util::current_time())); res.has_value() ) { + sqlite3_reset(stmt); + // TODO: do something with the error here? + } + + if ( auto res = checkError(sqlite3_step(stmt)); res.has_value() ) { + // TODO: do something with the error here? + } +} + // returns true in case of error ErrorResult SQLite::checkError(int code) { if ( code != SQLITE_OK && code != SQLITE_DONE ) { diff --git a/src/storage/backend/sqlite/SQLite.h b/src/storage/backend/sqlite/SQLite.h index 7899d7108b..1a3a713909 100644 --- a/src/storage/backend/sqlite/SQLite.h +++ b/src/storage/backend/sqlite/SQLite.h @@ -48,7 +48,11 @@ public: */ ErrorResult DoErase(ValPtr key, ErrorResultCallback* cb = nullptr) override; - // TODO: add support for checking for expired data + /** + * Removes any entries in the backend that have expired. Can be overridden by + * derived classes. + */ + void Expire() override; private: ErrorResult checkError(int code); diff --git a/testing/btest/Baseline/scripts.base.frameworks.storage.expiration/.stderr b/testing/btest/Baseline/scripts.base.frameworks.storage.expiration/.stderr new file mode 100644 index 0000000000..0494e4d891 --- /dev/null +++ b/testing/btest/Baseline/scripts.base.frameworks.storage.expiration/.stderr @@ -0,0 +1,3 @@ +### BTest baseline data generated by btest-diff. Do not edit. Use "btest -U/-u" to update. Requires BTest >= 0.63. +error in /Users/tim/Desktop/projects/storage-framework/testing/btest/.tmp/scripts.base.frameworks.storage.expiration/expiration.zeek, line 20: Failed to retrieve data: Failed to find row for key: no more rows available (Storage::get(backend, to_any_coerce key, F)) +received termination signal diff --git a/testing/btest/Baseline/scripts.base.frameworks.storage.expiration/out b/testing/btest/Baseline/scripts.base.frameworks.storage.expiration/out new file mode 100644 index 0000000000..f75c9fb9cc --- /dev/null +++ b/testing/btest/Baseline/scripts.base.frameworks.storage.expiration/out @@ -0,0 +1,5 @@ +### BTest baseline data generated by btest-diff. Do not edit. Use "btest -U/-u" to update. Requires BTest >= 0.63. +put result, T +get result, value7890 +get result same as inserted, T +get result, F diff --git a/testing/btest/scripts/base/frameworks/storage/expiration.zeek b/testing/btest/scripts/base/frameworks/storage/expiration.zeek new file mode 100644 index 0000000000..30caade463 --- /dev/null +++ b/testing/btest/scripts/base/frameworks/storage/expiration.zeek @@ -0,0 +1,42 @@ +# @TEST-DOC: Automatic expiration of stored data +# @TEST-EXEC: zcat <$TRACES/echo-connections.pcap.gz | zeek -b %INPUT > out +# @TEST-EXEC: btest-diff out +# @TEST-EXEC: btest-diff .stderr + +@load base/frameworks/storage +@load policy/frameworks/storage/backend/sqlite + +redef Storage::expire_interval = 5 secs; +redef exit_only_after_terminate = T; + +# Create a typename here that can be passed down into get(). +type str: string; + +global backend: opaque of Storage::BackendHandle; +global key: string = "key1234"; +global value: string = "value7890"; + +event check_removed() { + local res2 = Storage::get(backend, key, F); + print "get result", res2; + + Storage::close_backend(backend); + terminate(); +} + +event zeek_init() { + local opts : Storage::Backend::SQLite::Options; + opts$database_path = "storage-test.sqlite"; + opts$table_name = "testing"; + + backend = Storage::open_backend(Storage::SQLITE, opts, str, str); + + local res = Storage::put(backend, [$key=key, $value=value, $overwrite=T, $expire_time=2 secs, $async_mode=F]); + print "put result", res; + + local res2 = Storage::get(backend, key, F); + print "get result", res2; + print "get result same as inserted", value == (res2 as string); + + schedule 5 secs { check_removed() }; +}