cache: support caching multiple records for the same question

src/cache.rs

@@ -43,40 +43,68 @@ impl DnsCache {
     pub async fn get_cache(
         &self,
         question: &Question<Dname<Vec<u8>>>,
-    ) -> Option<Record<Dname<Vec<u8>>, UnknownRecordData<Vec<u8>>>> {
-        let (value, metadata): (Option<Vec<u8>>, Option<DnsCacheMetadata>) = self
+    ) -> Option<Vec<Record<Dname<Vec<u8>>, UnknownRecordData<Vec<u8>>>>> {
+        // One question can have multiple cached records; so we list by prefix
+        // Note that list_prefix returns 1000 records at maximum by default
+        // We don't expect one question to have that many answers, so it
+        // should be fine
+        let keys = self
             .store
-            .get_buf_metadata(&Self::question_to_key(question))
-            .await;
-        if value.is_none() || metadata.is_none() {
+            .list_prefix(&Self::question_to_key_prefix(question))
+            .await
+            .ok()?
+            .keys;
+        if keys.len() == 0 {
             return None;
         }
 
-        let (value, metadata) = (value.unwrap(), metadata.unwrap());
-        let elapsed_since_creation = (Date::now() / 1000f64) as u64 - metadata.created_ts;
-        // Calculate the remaining TTL correctly
-        // don't just return the original TTL blindly
-        let remaining_ttl = if elapsed_since_creation > metadata.ttl as u64 {
-            0
-        } else {
-            metadata.ttl as u64 - elapsed_since_creation
-        };
+        // If there are keys available, then return all of the cached records
+        let mut ret = Vec::new();
 
-        Some(Record::new(
-            question.qname().to_owned(),
-            question.qclass(),
-            remaining_ttl as u32,
-            UnknownRecordData::from_octets(question.qtype(), value),
-        ))
+        for k in keys {
+            let (value, metadata): (Option<Vec<u8>>, Option<DnsCacheMetadata>) =
+                self.store.get_buf_metadata(&k.name).await;
+            if value.is_none() || metadata.is_none() {
+                continue;
+            }
+
+            let (value, metadata) = (value.unwrap(), metadata.unwrap());
+            let elapsed_since_creation = (Date::now() / 1000f64) as u64 - metadata.created_ts;
+            // Calculate the remaining TTL correctly
+            // don't just return the original TTL blindly
+            let remaining_ttl = if elapsed_since_creation > metadata.ttl as u64 {
+                0
+            } else {
+                metadata.ttl as u64 - elapsed_since_creation
+            };
+
+            ret.push(Record::new(
+                question.qname().to_owned(),
+                question.qclass(),
+                remaining_ttl as u32,
+                UnknownRecordData::from_octets(question.qtype(), value),
+            ));
+        }
+
+        Some(ret)
     }
 
     fn record_to_key(record: &Record<Dname<Vec<u8>>, UnknownRecordData<Vec<u8>>>) -> String {
-        format!("{};{};{}", record.owner(), record.rtype(), record.class())
+        format!(
+            "{};{};{};{}",
+            record.owner(),
+            record.rtype(),
+            record.class(),
+            // We need to append the hash of the record data to the key
+            // because one question might have multiple answers
+            // When reading, we need to list the keys first
+            crate::util::hash_buf(record.data().data())
+        )
     }
 
-    fn question_to_key(question: &Question<Dname<Vec<u8>>>) -> String {
+    fn question_to_key_prefix(question: &Question<Dname<Vec<u8>>>) -> String {
         format!(
-            "{};{};{}",
+            "{};{};{};",
             question.qname(),
             question.qtype(),
             question.qclass()

src/client.rs

@@ -195,7 +195,7 @@ impl Client {
         let mut remaining = Vec::new();
         for q in questions {
             match self.cache.get_cache(&q).await {
-                Some(ans) => answers.push(ans),
+                Some(mut ans) => answers.append(&mut ans),
                 None => remaining.push(q),
             }
         }

src/kv.rs

@@ -27,6 +27,8 @@ extern "C" {
     ) -> Promise;
     #[wasm_bindgen(method, js_name = "getWithMetadata")]
     pub fn get_with_metadata_opts(this: &JsKvNamespace, key: &str, opts: JsValue) -> Promise;
+    #[wasm_bindgen(method)]
+    pub fn list(this: &JsKvNamespace, opts: JsValue) -> Promise;
 }
 
 // wasm-bindgen types are not Send + Sync, thus not usable in async_static
@@ -50,6 +52,25 @@ pub struct KvGetOptions {
     data_type: String,
 }
 
+#[derive(Serialize)]
+pub struct KvListOptions {
+    prefix: Option<String>,
+    limit: Option<u64>, // 1000 is default
+    cursor: Option<String>,
+}
+
+#[derive(Deserialize, Debug)]
+pub struct KvListKey {
+    pub name: String,
+}
+
+#[derive(Deserialize, Debug)]
+pub struct KvListResult {
+    pub keys: Vec<KvListKey>,
+    pub list_complete: bool,
+    pub cursor: Option<String>,
+}
+
 pub struct KvNamespace {
     inner: JsKvNamespace,
 }
@@ -123,6 +144,23 @@ impl KvNamespace {
             },
         )
     }
+
+    // List KV keys by prefix only
+    pub async fn list_prefix(&self, prefix: &str) -> Result<KvListResult, String> {
+        let promise = self.inner.list(
+            JsValue::from_serde(&KvListOptions {
+                prefix: Some(prefix.to_string()),
+                limit: None,
+                cursor: None,
+            })
+            .unwrap(),
+        );
+        let res = JsFuture::from(promise)
+            .await
+            .map_err(|_| "Could not list KV by prefix".to_string())?;
+        res.into_serde()
+            .map_err(|_| "Could not parse return value from KV listing".to_string())
+    }
 }
 
 #[wasm_bindgen]

src/util.rs

@@ -1,6 +1,7 @@
 use domain::base::message::Message;
 use js_sys::{Math, Promise};
 use std::ops::Add;
+use std::{collections::hash_map::DefaultHasher, hash::Hasher};
 use wasm_bindgen::prelude::*;
 use wasm_bindgen::JsValue;
 use wasm_bindgen_futures::JsFuture;
@@ -49,3 +50,11 @@ impl FromFloat<f64> for u16 {
         f as u16
     }
 }
+
+// Calculate a hash value from a u8 slice
+// used for generating answer cache keys
+pub fn hash_buf(buf: &[u8]) -> u64 {
+    let mut hasher = DefaultHasher::new();
+    hasher.write(buf);
+    hasher.finish()
+}