Add get_single_kb_item to Context

rust/src/nasl/builtin/error.rs

@@ -8,8 +8,8 @@ use super::cryptographic::CryptographicError;
 use super::host::HostError;
 use super::http::HttpError;
 use super::isotime::IsotimeError;
-use super::knowledge_base::KBError;
 use super::regex::RegexError;
+use super::KBError;
 use super::sys::SysError;
 use super::{misc::MiscError, network::socket::SocketError, ssh::SshError, string::StringError};
 

rust/src/nasl/builtin/knowledge_base/mod.rs

@@ -6,18 +6,22 @@
 mod tests;
 
 use std::time::{SystemTime, UNIX_EPOCH};
+use thiserror::Error;
 
 use crate::function_set;
 use crate::nasl::syntax::NaslValue;
 use crate::nasl::utils::error::FnError;
 use crate::nasl::utils::Context;
 use crate::storage::{Field, Kb, Retrieve};
 use nasl_function_proc_macro::nasl_function;
-use thiserror::Error;
 
 #[derive(Debug, Error)]
-#[error("{0}")]
-pub struct KBError(pub String);
+pub enum KBError {
+    #[error("Knowledge base item does not exist: {0}")]
+    ItemNotFound(String),
+    #[error("Multiple entries found for knowledge base item {0} where a single one was expected.")]
+    MultipleItemsFound(String),
+}
 
 /// NASL function to set a value under name in a knowledge base
 /// Only pushes unique values for the given name.
@@ -26,17 +30,9 @@ fn set_kb_item(
     c: &Context,
     name: NaslValue,
     value: NaslValue,
-    expires: Option<NaslValue>,
+    expires: Option<u64>,
 ) -> Result<NaslValue, FnError> {
-    let expires = match expires {
-        Some(NaslValue::Number(x)) => Some(x),
-        Some(NaslValue::Exit(0)) => None,
-        None => None,
-        Some(x) => {
-            return Err(KBError(format!("expected expires to be a number but is {x}.")).into())
-        }
-    }
-    .map(|seconds| {
+    let expires = expires.map(|seconds| {
         let start = SystemTime::now();
         match start.duration_since(UNIX_EPOCH) {
             Ok(x) => x.as_secs() + seconds as u64,

rust/src/nasl/builtin/mod.rs

@@ -27,6 +27,7 @@ mod sys;
 mod tests;
 
 pub use error::BuiltinError;
+pub use knowledge_base::KBError;
 pub use host::HostError;
 
 use crate::nasl::syntax::{Loader, NoOpLoader};

rust/src/nasl/builtin/network/mod.rs

@@ -5,9 +5,6 @@
 use std::{fmt::Display, net::IpAddr};
 
 use crate::nasl::prelude::*;
-use crate::storage::{Field, Retrieve};
-
-use super::knowledge_base::KBError;
 
 #[allow(clippy::module_inception)]
 pub mod network;
@@ -80,7 +77,7 @@ impl Display for OpenvasEncaps {
 }
 
 pub fn get_retry(context: &Context) -> u8 {
-    if let Ok(Some(val)) = get_kb_item(context, "timeout_retry") {
+    if let Ok(val) = context.get_single_kb_item("timeout_retry") {
         match val {
             NaslValue::String(val) => val.parse::<u8>().unwrap_or(2),
             NaslValue::Number(val) => {
@@ -97,26 +94,6 @@ pub fn get_retry(context: &Context) -> u8 {
     }
 }
 
-pub fn get_kb_item(context: &Context, name: &str) -> Result<Option<NaslValue>, FnError> {
-    context
-        .retriever()
-        .retrieve(context.key(), Retrieve::KB(name.to_string()))
-        .map(|r| {
-            r.into_iter().find_map(|x| match x {
-                Field::KB(kb) => kb.value.into(),
-                _ => None,
-            })
-        })
-        .map(|x| x.map(|x| x.into()))
-        .map_err(|e| e.into())
-}
-
-pub fn get_kb_item_str(context: &Context, name: &str) -> Result<String, FnError> {
-    get_kb_item(context, name)?
-        .map(|x| x.to_string())
-        .ok_or_else(|| KBError(format!("KB key {} is not set", name)).into())
-}
-
 struct Port(u16);
 
 impl FromNaslValue<'_> for Port {

rust/src/nasl/builtin/network/socket.rs

@@ -16,7 +16,7 @@ use rustls::ClientConnection;
 use thiserror::Error;
 
 use super::{
-    get_kb_item, get_kb_item_str, get_retry,
+    get_retry,
     network_utils::{convert_timeout, ipstr2ipaddr},
     tcp::TcpConnection,
     tls::create_tls_client,
@@ -305,37 +305,17 @@ impl NaslSockets {
     /// - Secret/kdc_use_tcp
     #[nasl_function]
     fn open_sock_kdc(&mut self, context: &Context) -> Result<NaslValue, FnError> {
-        let hostname = get_kb_item_str(context, "Secret/kdc_hostname")?;
+        let hostname: String = context.get_single_kb_item("Secret/kdc_hostname")?;
 
         let ip = lookup_host(&hostname)
             .map_err(|_| SocketError::HostnameLookupFailed(hostname.clone()))?
             .into_iter()
             .next()
             .ok_or(SocketError::HostnameNoIpFound(hostname))?;
 
-        let port = get_kb_item(context, "Secret/kdc_port")?;
+        let port = context.get_single_kb_item::<Port>("Secret/kdc_port")?.0;
 
-        let port = match port {
-            Some(NaslValue::Number(x)) => {
-                if x <= 0 || x > 65535 {
-                    Err(SocketError::Diagnostic(
-                        "KB key 'Secret/kdc_port' out of range".to_string(),
-                    ))
-                } else {
-                    Ok(x as u16)
-                }
-            }
-            Some(_) => Err(SocketError::Diagnostic(
-                "KB key 'Secret/kdc_port' has wrong type".to_string(),
-            )),
-            None => Err(SocketError::Diagnostic(
-                "KB key 'Secret/kdc_port' is not set".to_string(),
-            )),
-        }?;
-
-        let use_tcp: bool = get_kb_item(context, "Secret/kdc_use_tcp")?
-            .map(|x| x.into())
-            .unwrap_or(false);
+        let use_tcp: bool = context.get_single_kb_item("Secret/kdc_use_tcp")?;
 
         let socket = if use_tcp {
             let tcp = TcpConnection::connect(
@@ -475,10 +455,10 @@ impl NaslSockets {
     /// Reads the information necessary for a TLS connection from the KB and
     /// return a TlsConfig on success.
     fn get_tls_conf(context: &Context) -> Result<TlsConfig, FnError> {
-        let cert_path = get_kb_item_str(context, "SSL/cert")?;
-        let key_path = get_kb_item_str(context, "SSL/key")?;
-        let password = get_kb_item_str(context, "SSL/password")?;
-        let cafile_path = get_kb_item_str(context, "SSL/CA")?;
+        let cert_path = context.get_single_kb_item("SSL/cert")?;
+        let key_path = context.get_single_kb_item("SSL/key")?;
+        let password = context.get_single_kb_item("SSL/password")?;
+        let cafile_path = context.get_single_kb_item("SSL/CA")?;
 
         Ok(TlsConfig {
             cert_path,

rust/src/nasl/utils/context.rs

@@ -4,10 +4,16 @@
 
 //! Defines the context used within the interpreter and utilized by the builtin functions
 
+use itertools::Itertools;
+
+use crate::nasl::builtin::KBError;
 use crate::nasl::syntax::{Loader, NaslValue, Statement};
-use crate::storage::{ContextKey, Dispatcher, Retriever};
+use crate::nasl::{FromNaslValue, WithErrorInfo};
+use crate::storage::{ContextKey, Dispatcher, Field, Retrieve, Retriever};
 
 use super::hosts::resolve;
+use super::error::ReturnBehavior;
+use super::FnError;
 use super::{executor::Executor, lookup_keys::FC_ANON_ARGS};
 
 /// Contexts are responsible to locate, add and delete everything that is declared within a NASL plugin
@@ -485,6 +491,40 @@ impl<'a> Context<'a> {
     pub fn loader(&self) -> &dyn Loader {
         self.loader
     }
+
+    /// Return a single item from the knowledge base.
+    /// If multiple entries are found (which would result
+    /// in forking the interpreter), return an error.
+    /// This function automatically converts the item
+    /// to a specific type via its `FromNaslValue` impl
+    /// and returns the appropriate error if necessary.
+    pub fn get_single_kb_item<T: for<'b> FromNaslValue<'b>>(
+        &self,
+        name: &str,
+    ) -> Result<T, FnError> {
+        // If we find multiple or no items at all, return an error that
+        // exits the script instead of continuing execution with a return
+        // value, since this is most likely an error in the feed.
+        let val = self
+            .get_single_kb_item_inner(name)
+            .map_err(|e| e.with(ReturnBehavior::ExitScript))?;
+        T::from_nasl_value(&val)
+    }
+
+    fn get_single_kb_item_inner<'kb>(&self, name: &str) -> Result<NaslValue, FnError> {
+        let result = self
+            .retriever()
+            .retrieve(&self.key, Retrieve::KB(name.to_string()))?;
+        let single_item = result
+            .filter_map(|field| match field {
+                Field::KB(kb) => Some(kb.value.into()),
+                _ => None,
+            })
+            .at_most_one()
+            .map_err(|_| KBError::MultipleItemsFound(name.to_string()))?
+            .ok_or_else(|| KBError::ItemNotFound(name.to_string()))?;
+        Ok(single_item)
+    }
 }
 
 impl From<&ContextType> for NaslValue {

rust/src/nasl/utils/error.rs

@@ -186,6 +186,16 @@ impl<T: Into<NaslValue>, E: Into<FnError>> WithErrorInfo<ReturnValue<T>> for E {
     }
 }
 
+impl<E: Into<FnError>> WithErrorInfo<ReturnBehavior> for E {
+    type Error = FnError;
+
+    fn with(self, val: ReturnBehavior) -> Self::Error {
+        let mut e = self.into();
+        e.return_behavior = val;
+        e
+    }
+}
+
 impl ArgumentError {
     /// Helper function to quickly construct a `WrongArgument` variant
     /// containing the name of the argument, the expected value and