diff --git a/src/cache.rs b/src/cache.rs index 7baf4f4..aae5f5c 100644 --- a/src/cache.rs +++ b/src/cache.rs @@ -43,40 +43,68 @@ impl DnsCache { pub async fn get_cache( &self, question: &Question>>, - ) -> Option>, UnknownRecordData>>> { - let (value, metadata): (Option>, Option) = self + ) -> Option>, UnknownRecordData>>>> { + // 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>, Option) = + 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>, UnknownRecordData>>) -> 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>>) -> String { + fn question_to_key_prefix(question: &Question>>) -> String { format!( - "{};{};{}", + "{};{};{};", question.qname(), question.qtype(), question.qclass() diff --git a/src/client.rs b/src/client.rs index 5ba4fd6..92ec845 100644 --- a/src/client.rs +++ b/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), } } diff --git a/src/kv.rs b/src/kv.rs index f719876..6aebad0 100644 --- a/src/kv.rs +++ b/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, + limit: Option, // 1000 is default + cursor: Option, +} + +#[derive(Deserialize, Debug)] +pub struct KvListKey { + pub name: String, +} + +#[derive(Deserialize, Debug)] +pub struct KvListResult { + pub keys: Vec, + pub list_complete: bool, + pub cursor: Option, +} + 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 { + 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] diff --git a/src/util.rs b/src/util.rs index 1c8bbf3..0550246 100644 --- a/src/util.rs +++ b/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 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() +}