SQLite: Handle automated expiration

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<std::string, std::string> statements =
-        {{"put", util::fmt("insert into %s (key_str, value_str) values(?, ?)", 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) values(?, ?) ON CONFLICT(key_str) "
+         {"put_update",
-                                  "DO UPDATE SET value_str=?",
+          util::fmt("insert into %s (key_str, value_str, expire_time) values(?, ?, ?) ON CONFLICT(key_str) "
-                                  table_name.c_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<std::string>(res.value());
+        return zeek::unexpected<std::string>(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 ) {

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);

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

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

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() };
+}