Pass key/value types for validation when opening backends

scripts/base/frameworks/storage/main.zeek

@@ -7,13 +7,22 @@ module Storage;
 export {
 	## Opens a new backend connection based on a configuration object.
 	##
-	## btype: A tag indicating what type of backend should be opened.
+	## btype: A tag indicating what type of backend should be opened. These are
+	##        defined by the backend plugins loaded.
 	##
 	## options: A record containing the configuration for the connection.
 	##
-	## Returns: A handle to the new backend connection, or null if the
-	##          connection failed.
-	global open_backend: function(btype: Storage::Backend, options: any): opaque of Storage::BackendHandle;
+	## key_type: The script-level type of keys stored in the backend. Used for
+	##           validation of keys passed to other framework methods.
+	##
+	## val_type: The script-level type of keys stored in the backend. Used for
+	##           validation of values passed to :zeek:see:`Storage::put` as well as
+	##           for type conversions for return values from :zeek:see:`Storage::get`.
+	##
+	## Returns: A handle to the new backend connection, or ``F`` if the connection
+	##          failed.
+	global open_backend: function(btype: Storage::Backend, options: any, key_type: any,
+	                              val_type: any): opaque of Storage::BackendHandle;
 
 	## Closes an existing backend connection.
 	##
@@ -30,10 +39,13 @@ export {
 	##
 	## value: A corresponding value.
 	##
-	## overwrite: A flag indicating whether this value should overwrite an
-	##            existing entry for the key.
+	## overwrite: A flag indicating whether this value should overwrite an existing
+	##            entry for the key.
 	##
-	## Returns: A boolean indicating success or failure of the operation.
+	## Returns: A boolean indicating success or failure of the operation. Type
+	##          comparison failures against the types passed to
+	##          :zeek:see:`Storage::open_backend` for the backend will cause false to
+	##          be returned.
 	global put: function(backend: opaque of Storage::BackendHandle, key: any, value: any, overwrite: bool): bool;
 
 	## Gets an entry from the backend.
@@ -42,12 +54,11 @@ export {
 	##
 	## key: The key to look up.
 	##
-	## val_type: The type of the value to return.
-	##
-	## Returns: A boolean indicating success or failure of the
-	##          operation. Type conversion failures for the value will
-	##          return false.
-	global get: function(backend: opaque of Storage::BackendHandle, key: any, val_type: any): any;
+	## Returns: A boolean indicating success or failure of the operation.  Type
+	##          comparison failures against the types passed to
+	##          :zeek:see:`Storage::open_backend` for the backend will cause false to
+	##          be returned.
+	global get: function(backend: opaque of Storage::BackendHandle, key: any): any;
 
 	## Erases an entry from the backend.
 	##
@@ -55,13 +66,16 @@ export {
 	##
 	## key: The key to erase.
 	##
-	## Returns: A boolean indicating success or failure of the operation.
+	## Returns: A boolean indicating success or failure of the operation.  Type
+	##          comparison failures against the types passed to
+	##          :zeek:see:`Storage::open_backend` for the backend will cause false to
+	##          be returned.
 	global erase: function(backend: opaque of Storage::BackendHandle, key: any): bool;
 }
 
-function open_backend(btype: Storage::Backend, options: any): opaque of Storage::BackendHandle
+function open_backend(btype: Storage::Backend, options: any, key_type: any, val_type: any): opaque of Storage::BackendHandle
 {
-	return Storage::__open_backend(btype, options);
+	return Storage::__open_backend(btype, options, key_type, val_type);
 }
 
 function close_backend(backend: opaque of Storage::BackendHandle): bool
@@ -74,9 +88,9 @@ function put(backend: opaque of Storage::BackendHandle, key: any, value: any, ov
 	return Storage::__put(backend, key, value, overwrite);
 }
 
-function get(backend: opaque of Storage::BackendHandle, key: any, val_type: any): any
+function get(backend: opaque of Storage::BackendHandle, key: any): any
 {
-	return Storage::__get(backend, key, val_type);
+	return Storage::__get(backend, key);
 }
 
 function erase(backend: opaque of Storage::BackendHandle, key: any): bool

src/storage/Backend.cc

@@ -2,27 +2,48 @@
 
 #include "zeek/storage/Backend.h"
 
+#include "zeek/Desc.h"
 #include "zeek/broker/Data.h"
 
 namespace zeek::storage {
 
-ErrorResult Backend::Open(RecordValPtr options) { return DoOpen(std::move(options)); }
+ErrorResult Backend::Open(RecordValPtr options, TypePtr kt, TypePtr vt) {
+    key_type = std::move(kt);
+    val_type = std::move(vt);
+
+    return DoOpen(std::move(options));
+}
 
 ErrorResult Backend::Put(ValPtr key, ValPtr value, bool overwrite) {
     // The intention for this method is to do some other heavy lifting in regard
     // to backends that need to pass data through the manager instead of directly
     // through the workers. For the first versions of the storage framework it
     // just calls the backend itself directly.
+    if ( ! same_type(key->GetType(), key_type) )
+        return util::fmt("type of key passed (%s) does not match backend's key type (%s)",
+                         obj_desc_short(key->GetType().get()).c_str(), key_type->GetName().c_str());
+    if ( ! same_type(value->GetType(), val_type) )
+        return util::fmt("type of value passed (%s) does not match backend's value type (%s)",
+                         obj_desc_short(value->GetType().get()).c_str(), val_type->GetName().c_str());
+
     return DoPut(std::move(key), std::move(value), overwrite);
 }
 
-ValResult Backend::Get(ValPtr key, TypePtr value_type) {
+ValResult Backend::Get(ValPtr key) {
     // See the note in Put().
-    return DoGet(std::move(key), std::move(value_type));
+    if ( ! same_type(key->GetType(), key_type) )
+        return zeek::unexpected<std::string>(util::fmt("type of key passed (%s) does not match backend's key type (%s)",
+                                                       key->GetType()->GetName().c_str(), key_type->GetName().c_str()));
+
+    return DoGet(std::move(key));
 }
 
 ErrorResult Backend::Erase(ValPtr key) {
     // See the note in Put().
+    if ( ! same_type(key->GetType(), key_type) )
+        return util::fmt("type of key passed (%s) does not match backend's key type (%s)",
+                         key->GetType()->GetName().c_str(), key_type->GetName().c_str());
+
     return DoErase(std::move(key));
 }
 

src/storage/Backend.h

@@ -41,13 +41,10 @@ public:
      * Retrieve a value from the backend for a provided key.
      *
      * @param key the key to lookup in the backend.
-     * @param value_type The script-land type to be used when retrieving values
-     * from the backend.
-     * @return A result pair containing a ValPtr with the resulting value or
-     * nullptr retrieval failed, and a string with the error message if the
-     * operation failed.
+     * @return A std::expected containing either a valid ValPtr with the result
+     * of the operation or a string containing an error message for failure.
      */
-    ValResult Get(ValPtr key, TypePtr value_type);
+    ValResult Get(ValPtr key);
 
     /**
      * Erases the value for a key from the backend.
@@ -79,10 +76,14 @@ protected:
      * Called by the manager system to open the backend.
      *
      * @param options A record storing configuration options for the backend.
-     * @return A result pair containing a bool with the success state, and a
-     * possible error string if the operation failed.
+     * @param kt The script-side type of the keys stored in the backend. Used for
+     * validation of types.
+     * @param vt The script-side type of the values stored in the backend. Used for
+     * validation of types and conversion during retrieval.
+     * @return An optional value potentially containing an error string if
+     * needed. Will be unset if the operation succeeded.
      */
-    ErrorResult Open(RecordValPtr options);
+    ErrorResult Open(RecordValPtr options, TypePtr kt, TypePtr vt);
 
     /**
      * Finalizes the backend when it's being closed. Can be overridden by
@@ -103,13 +104,16 @@ protected:
     /**
      * The workhorse method for Get().
      */
-    virtual ValResult DoGet(ValPtr key, TypePtr vt) = 0;
+    virtual ValResult DoGet(ValPtr key) = 0;
 
     /**
      * The workhorse method for Erase().
      */
     virtual ErrorResult DoErase(ValPtr key) = 0;
 
+    TypePtr key_type;
+    TypePtr val_type;
+
     std::string tag;
 };
 

src/storage/Manager.cc

@@ -10,7 +10,8 @@ Manager::Manager() : plugin::ComponentManager<storage::Component>("Storage", "Ba
 
 void Manager::InitPostScript() { detail::backend_opaque = make_intrusive<OpaqueType>("Storage::Backend"); }
 
-zeek::expected<BackendPtr, std::string> Manager::OpenBackend(const Tag& type, RecordValPtr options) {
+zeek::expected<BackendPtr, std::string> Manager::OpenBackend(const Tag& type, RecordValPtr options, TypePtr key_type,
+                                                             TypePtr val_type) {
     Component* c = Lookup(type);
     if ( ! c ) {
         return zeek::unexpected<std::string>(
@@ -32,7 +33,7 @@ zeek::expected<BackendPtr, std::string> Manager::OpenBackend(const Tag& type, Re
             util::fmt("Failed to instantiate backend %s", GetComponentName(type).c_str()));
     }
 
-    if ( auto res = bp->Open(std::move(options)); res.has_value() ) {
+    if ( auto res = bp->Open(std::move(options), std::move(key_type), std::move(val_type)); res.has_value() ) {
         return zeek::unexpected<std::string>(
             util::fmt("Failed to open backend %s: %s", GetComponentName(type).c_str(), res.value().c_str()));
     }

src/storage/Manager.h

@@ -23,12 +23,17 @@ public:
      * Opens a new storage backend.
      *
      * @param type The tag for the type of backend being opened.
-     * @param options A record val representing the configuration for this
-     * type of backend.
-     * @return A pair containing a pointer to a backend and a string for
-     * returning error messages if needed.
+     * @param options A record val representing the configuration for this type of
+     * backend.
+     * @param key_type The script-side type of the keys stored in the backend. Used for
+     * validation of types.
+     * @param val_type The script-side type of the values stored in the backend. Used for
+     * validation of types and conversion during retrieval.
+     * @return An optional value potentially containing an error string if needed. Will be
+     * unset if the operation succeeded.
      */
-    zeek::expected<BackendPtr, std::string> OpenBackend(const Tag& type, RecordValPtr options);
+    zeek::expected<BackendPtr, std::string> OpenBackend(const Tag& type, RecordValPtr options, TypePtr key_type,
+                                                        TypePtr val_type);
 
     /**
      * Closes a storage backend.

src/storage/storage.bif

@@ -14,14 +14,16 @@ event Storage::backend_opened%(%);
 # Generated when a backend connection is lost
 event Storage::backend_lost%(%);
 
-function Storage::__open_backend%(btype: Storage::Backend, options: any%): opaque of Storage::BackendHandle
+function Storage::__open_backend%(btype: Storage::Backend, options: any, key_type: any, val_type: any%): opaque of Storage::BackendHandle
 	%{
 	auto btype_val = IntrusivePtr<EnumVal>{NewRef{}, btype->AsEnumVal()};
 	Tag tag{btype_val};
 
-	auto options_val = IntrusivePtr<RecordVal>{NewRef{}, options->AsRecordVal()};
+	auto kt = key_type->AsTypeVal()->GetType()->AsTypeType()->GetType();
+	auto vt = val_type->AsTypeVal()->GetType()->AsTypeType()->GetType();
 
-	auto b = storage_mgr->OpenBackend(tag, options_val);
+	auto options_val = IntrusivePtr<RecordVal>{NewRef{}, options->AsRecordVal()};
+	auto b = storage_mgr->OpenBackend(tag, options_val, kt, vt);
 
 	if ( ! b.has_value() ) {
 		emit_builtin_error(b.error().c_str());
@@ -70,7 +72,7 @@ function Storage::__put%(backend: opaque of Storage::BackendHandle, key: any, va
 	return val_mgr->Bool(true);
 	%}
 
-function Storage::__get%(backend: opaque of Storage::BackendHandle, key: any, val_type: any%): any
+function Storage::__get%(backend: opaque of Storage::BackendHandle, key: any%): any
 	%{
 	auto b = dynamic_cast<storage::detail::BackendHandleVal*>(backend);
 	if ( ! b ) {
@@ -83,8 +85,7 @@ function Storage::__get%(backend: opaque of Storage::BackendHandle, key: any, va
 	// TODO: add support for when statements (see broker/store.bif)
 
 	auto key_v = IntrusivePtr<Val>{NewRef{}, key};
-	auto vt = val_type->AsTypeVal()->GetType()->AsTypeType()->GetType();
-	auto result = b->backend->Get(key_v, vt);
+	auto result = b->backend->Get(key_v);
 	if ( ! result.has_value() ) {
 		emit_builtin_error(util::fmt("Failed to retrieve data: %s", result.error().c_str()));
 		return val_mgr->Bool(false);

testing/btest/Baseline/plugins.storage/zeek-stderr

@@ -1,4 +1,4 @@
 ### BTest baseline data generated by btest-diff. Do not edit. Use "btest -U/-u" to update. Requires BTest >= 0.63.
-error in <...>/storage.zeek, line 37: Failed to retrieve data: Failed to find key (Storage::get(b, to_any_coerce key, to_any_coerce str))
-error in <...>/storage.zeek, line 50: Failed to open backend STORAGEDUMMY: open_fail was set to true, returning error (Storage::open_backend(Storage::STORAGEDUMMY, to_any_coerce opts))
+error in <...>/storage.zeek, line 37: Failed to retrieve data: Failed to find key (Storage::get(b, to_any_coerce key))
+error in <...>/storage.zeek, line 50: Failed to open backend STORAGEDUMMY: open_fail was set to true, returning error (Storage::open_backend(Storage::STORAGEDUMMY, to_any_coerce opts, to_any_coerce str, to_any_coerce str))
 error in <...>/storage.zeek, line 51: Invalid storage handle (Storage::close_backend(b2) and F)

testing/btest/plugins/storage-plugin/src/StorageDummy.cc

@@ -46,13 +46,13 @@ zeek::storage::ErrorResult StorageDummy::DoPut(zeek::ValPtr key, zeek::ValPtr va
 /**
  * The workhorse method for Get(). This must be implemented for plugins.
  */
-zeek::storage::ValResult StorageDummy::DoGet(zeek::ValPtr key, zeek::TypePtr vt) {
+zeek::storage::ValResult StorageDummy::DoGet(zeek::ValPtr key) {
     auto json_key = key->ToJSON();
     auto it = data.find(json_key->ToStdString());
     if ( it == data.end() )
         return zeek::unexpected<std::string>("Failed to find key");
 
-    auto val = zeek::detail::ValFromJSON(it->second.c_str(), vt, zeek::Func::nil);
+    auto val = zeek::detail::ValFromJSON(it->second.c_str(), val_type, zeek::Func::nil);
     if ( std::holds_alternative<zeek::ValPtr>(val) ) {
         zeek::ValPtr val_v = std::get<zeek::ValPtr>(val);
         return val_v;

testing/btest/plugins/storage-plugin/src/StorageDummy.h

@@ -41,7 +41,7 @@ public:
     /**
      * The workhorse method for Get().
      */
-    zeek::storage::ValResult DoGet(zeek::ValPtr key, zeek::TypePtr vt) override;
+    zeek::storage::ValResult DoGet(zeek::ValPtr key) override;
 
     /**
      * The workhorse method for Erase().

testing/btest/plugins/storage.zeek

@@ -26,20 +26,20 @@ event zeek_init() {
 
 	# Test basic operation. The second get() should return an error
 	# as the key should have been erased.
-	local b = Storage::open_backend(Storage::STORAGEDUMMY, opts);
+	local b = Storage::open_backend(Storage::STORAGEDUMMY, opts, str, str);
 	local put_res = Storage::put(b, key, value, F);
-	local get_res = Storage::get(b, key, str);
+	local get_res = Storage::get(b, key);
 	if ( get_res is bool ) {
 		print("Got an invalid value in response!");
 	}
 
 	local erase_res = Storage::erase(b, key);
-	get_res = Storage::get(b, key, str);
+	get_res = Storage::get(b, key);
 	Storage::close_backend(b);
 
 	# Test attempting to use the closed handle.
 	put_res = Storage::put(b, "a", "b", F);
-	get_res = Storage::get(b, "a", str);
+	get_res = Storage::get(b, "a");
 	erase_res = Storage::erase(b, "a");
 
 	print(fmt("results of trying to use closed handle: get: %d, put: %d, erase: %d",
@@ -47,6 +47,6 @@ event zeek_init() {
 
 	# Test failing to open the handle and test closing an invalid handle.
 	opts$open_fail = T;
-	local b2 = Storage::open_backend(Storage::STORAGEDUMMY, opts);
+	local b2 = Storage::open_backend(Storage::STORAGEDUMMY, opts, str, str);
 	Storage::close_backend(b2);
 }