Repository: cloudflare/daphne
Pull request #116 (2022-09-09 01:58:49)

Allow rwlock-guarded data to be passed by reference
Partially addresses issue #82. There are two commits in this PR (review them separately):

1. Allow `HpkeDecrypter::get_hpke_config_for` to handle mutex guard.

Change the return type of `HpkeDecrypter::get_hpke_config_for` to
`HpkeDecryptor::WrappedHpkeConfig: AsRef<HpkeConfig>`. This allows
Daphne-Worker to return a reference to an HPKE config protected by a
mutex.

For `DaphneWorkerConfig`, the return type is
`GuardedHpkeReceiverConfig<'a>`, which is required to live as long as
the mutex guard for the underlying data structure (a map from config IDs
to configs). The compiler does not know whether this object lives as
long as `DaphneWorkerConfig`, so it is necessary to add a lifetime to
the `HpkeDecrypter` trait.

Because `get_hpke_config_for()` is called at most once per HTTP request,
and because the HPKE config needs to be fetched from KV for each
request, this change does not save any overhead. (It doesn't add any
either.) However, since plan to move the DAP task configs to KV as part
of issue#82, this pattern will need to be replicated for
`DapAggregator::get_task_config_for()` as well.

2. Make `DapAggregator::get_task_config_for()` async.

Allows us to use KV to back task storage from behind a mutex, which will be
needed for issue#84. Copying each time we need to fetch a task config would
be very expensive.


tanyav2: Given that these are configs with what appears to be significantly more reads than writes, any reason why you're using mutex and not a rwlock? 
cjpatton: > Given that these are configs with what appears to be significantly more reads than writes, any reason why you're using mutex and not a rwlock?

Ohhh, actually we are use `RwLock`, I just have been conflating these 😬 
tanyav2: > > Given that these are configs with what appears to be significantly more reads than writes, any reason why you're using mutex and not a rwlock?
> 
> Ohhh, actually we are use `RwLock`, I just have been conflating these grimacing

oops I'd just started looking at this, based it on the PR message
cjpatton: @armfazh asks (internally): 

> what I didn't understand so well is why to store references for that specific use case.

The high-level goal here is to plumb a reference to an HPKE config (or DAP task config) through the trait API when the underlying data structure is protected by an `RwLock`. I believe an explicit lifetime is needed to tell the compiler that the thing that's returned (i.e., an implementation of `AsRef<HpkeConfig>`) lives as long as the underlying data.
cjpatton: [Here](https://play.rust-lang.org/?version=stable&mode=debug&edition=2021&gist=c7c9e8b3438d0156ec4d2e972d5beb0c) is a minimal working example if anyone wants to play around with this. This makes me realize that we probably also want to add the lifetime to the `Id` reference as well, I'll update this PR accordingly.
cjpatton: Since there are no objects, and no one has any better ideas here, I'm gonna merge as-is. Thanks all for the comments.

Patch:
diff --git a/daphne/src/hpke.rs b/daphne/src/hpke.rs
index 8885e27..6d3e7df 100644
--- a/daphne/src/hpke.rs
+++ b/daphne/src/hpke.rs
@@ -81,9 +81,15 @@ impl HpkeConfig {
 
 /// HPKE decrypter functionality.
 #[async_trait(?Send)]
-pub trait HpkeDecrypter {
+pub trait HpkeDecrypter<'a> {
+    /// Return type of `get_hpke_config_for()`, wraps a reference to an HPKE config.
+    type WrappedHpkeConfig: AsRef<HpkeConfig>;
+
     /// Look up the HPKE configuration to use for the given task ID.
-    async fn get_hpke_config_for(&self, task_id: &Id) -> Result<Option<HpkeConfig>, DapError>;
+    async fn get_hpke_config_for(
+        &'a self,
+        task_id: &Id,
+    ) -> Result<Option<Self::WrappedHpkeConfig>, DapError>;
 
     /// Returns `true` if a ciphertext with the HPKE config ID can be consumed in the current task.
     async fn can_hpke_decrypt(&self, task_id: &Id, config_id: u8) -> Result<bool, DapError>;
@@ -252,8 +258,13 @@ impl HpkeReceiverConfig {
 }
 
 #[async_trait(?Send)]
-impl HpkeDecrypter for HpkeReceiverConfig {
-    async fn get_hpke_config_for(&self, _task_id: &Id) -> Result<Option<HpkeConfig>, DapError> {
+impl<'a> HpkeDecrypter<'a> for HpkeReceiverConfig {
+    type WrappedHpkeConfig = HpkeConfig;
+
+    async fn get_hpke_config_for(
+        &'a self,
+        _task_id: &Id,
+    ) -> Result<Option<Self::WrappedHpkeConfig>, DapError> {
         unreachable!("not implemented");
     }
 
diff --git a/daphne/src/lib.rs b/daphne/src/lib.rs
index b376162..9a07a2b 100644
--- a/daphne/src/lib.rs
+++ b/daphne/src/lib.rs
@@ -325,6 +325,12 @@ impl DapTaskConfig {
     }
 }
 
+impl AsRef<DapTaskConfig> for DapTaskConfig {
+    fn as_ref(&self) -> &Self {
+        self
+    }
+}
+
 #[derive(Deserialize, Serialize)]
 struct ShadowDapTaskConfig {
     version: DapVersion,
diff --git a/daphne/src/messages.rs b/daphne/src/messages.rs
index 5e1a342..5f90529 100644
--- a/daphne/src/messages.rs
+++ b/daphne/src/messages.rs
@@ -616,6 +616,12 @@ pub struct HpkeConfig {
     pub(crate) public_key: Vec<u8>,
 }
 
+impl AsRef<HpkeConfig> for HpkeConfig {
+    fn as_ref(&self) -> &Self {
+        self
+    }
+}
+
 impl Encode for HpkeConfig {
     fn encode(&self, bytes: &mut Vec<u8>) {
         self.id.encode(bytes);
diff --git a/daphne/src/roles.rs b/daphne/src/roles.rs
index 18e9ed3..54b761d 100644
--- a/daphne/src/roles.rs
+++ b/daphne/src/roles.rs
@@ -79,7 +79,10 @@ pub trait DapAuthorizedSender<S> {
 
 /// DAP Aggregator functionality.
 #[async_trait(?Send)]
-pub trait DapAggregator<S>: HpkeDecrypter + Sized {
+pub trait DapAggregator<'a, S>: HpkeDecrypter<'a> + Sized {
+    /// Return type of `get_task_config_for()`, wraps a reference to task configuration.
+    type WrappedDapTaskConfig: AsRef<DapTaskConfig>;
+
     /// Decide whether the given DAP request is authorized.
     async fn authorized(&self, req: &DapRequest<S>) -> Result<bool, DapError>;
 
@@ -87,7 +90,10 @@ pub trait DapAggregator<S>: HpkeDecrypter + Sized {
     fn get_global_config(&self) -> &DapGlobalConfig;
 
     /// Look up the DAP task configuration for the given task ID.
-    fn get_task_config_for(&self, task_id: &Id) -> Option<&DapTaskConfig>;
+    async fn get_task_config_for(
+        &'a self,
+        task_id: &Id,
+    ) -> Result<Option<Self::WrappedDapTaskConfig>, DapError>;
 
     /// Get the current time (number of seconds since the beginning of UNIX time).
     fn get_current_time(&self) -> u64;
@@ -119,7 +125,7 @@ pub trait DapAggregator<S>: HpkeDecrypter + Sized {
         -> Result<(), DapError>;
 
     /// Handle HTTP GET to `/hpke_config?task_id=<task_id>`.
-    async fn http_get_hpke_config(&self, req: &DapRequest<S>) -> Result<DapResponse, DapAbort> {
+    async fn http_get_hpke_config(&'a self, req: &DapRequest<S>) -> Result<DapResponse, DapAbort> {
         // Check whether the DAP version indicated by the sender is supported.
         if req.version == DapVersion::Unknown {
             return Err(DapAbort::InvalidProtocolVersion);
@@ -145,10 +151,11 @@ pub trait DapAggregator<S>: HpkeDecrypter + Sized {
         if let Some(ref task_id) = id {
             let task_config = self
                 .get_task_config_for(task_id)
+                .await?
                 .ok_or(DapAbort::UnrecognizedTask)?;
 
             // Check whether the DAP version in the request matches the task config.
-            if task_config.version != req.version {
+            if task_config.as_ref().version != req.version {
                 return Err(DapAbort::InvalidProtocolVersion);
             }
 
@@ -158,7 +165,7 @@ pub trait DapAggregator<S>: HpkeDecrypter + Sized {
                 .ok_or(DapAbort::UnrecognizedTask)?;
             Ok(DapResponse {
                 media_type: Some(MEDIA_TYPE_HPKE_CONFIG),
-                payload: hpke_config.get_encoded(),
+                payload: hpke_config.as_ref().get_encoded(),
             })
         } else {
             Err(DapAbort::BadRequest("missing query parameter".into()))
@@ -192,7 +199,7 @@ macro_rules! leader_post {
 
 /// DAP Leader functionality.
 #[async_trait(?Send)]
-pub trait DapLeader<S>: DapAuthorizedSender<S> + DapAggregator<S> {
+pub trait DapLeader<'a, S>: DapAuthorizedSender<S> + DapAggregator<'a, S> {
     /// Data type used to guide selection of a set of reports for aggregation.
     type ReportSelector;
 
@@ -235,7 +242,7 @@ pub trait DapLeader<S>: DapAuthorizedSender<S> + DapAggregator<S> {
 
     /// Handle HTTP POST to `/upload`. The input is the encoded report sent in the body of the HTTP
     /// request.
-    async fn http_post_upload(&self, req: &DapRequest<S>) -> Result<(), DapAbort> {
+    async fn http_post_upload(&'a self, req: &DapRequest<S>) -> Result<(), DapAbort> {
         // Check whether the DAP version indicated by the sender is supported.
         if req.version == DapVersion::Unknown {
             return Err(DapAbort::InvalidProtocolVersion);
@@ -244,10 +251,11 @@ pub trait DapLeader<S>: DapAuthorizedSender<S> + DapAggregator<S> {
         let report = Report::get_decoded(req.payload.as_ref())?;
         let task_config = self
             .get_task_config_for(&report.task_id)
+            .await?
             .ok_or(DapAbort::UnrecognizedTask)?;
 
         // Check whether the DAP version in the request matches the task config.
-        if task_config.version != req.version {
+        if task_config.as_ref().version != req.version {
             return Err(DapAbort::InvalidProtocolVersion);
         }
 
@@ -273,7 +281,7 @@ pub trait DapLeader<S>: DapAuthorizedSender<S> + DapAggregator<S> {
     /// Handle HTTP POST to `/collect`. The input is a [`CollectReq`](crate::messages::CollectReq).
     /// The return value is a URI that the Collector can poll later on to get the corresponding
     /// [`CollectResp`](crate::messages::CollectResp).
-    async fn http_post_collect(&self, req: &DapRequest<S>) -> Result<Url, DapAbort> {
+    async fn http_post_collect(&'a self, req: &DapRequest<S>) -> Result<Url, DapAbort> {
         // Check whether the DAP version indicated by the sender is supported.
         if req.version == DapVersion::Unknown {
             return Err(DapAbort::InvalidProtocolVersion);
@@ -284,9 +292,11 @@ pub trait DapLeader<S>: DapAuthorizedSender<S> + DapAggregator<S> {
         }
 
         let collect_req = CollectReq::get_decoded(req.payload.as_ref())?;
-        let task_config = self
+        let wrapped_task_config = self
             .get_task_config_for(&collect_req.task_id)
+            .await?
             .ok_or(DapAbort::UnrecognizedTask)?;
+        let task_config = wrapped_task_config.as_ref();
         check_batch_param!(
             task_config,
             collect_req.batch_interval,
@@ -473,7 +483,7 @@ pub trait DapLeader<S>: DapAuthorizedSender<S> + DapAggregator<S> {
     /// create a bottleneck. Such deployments can improve throughput by running many aggregation
     /// jobs in parallel.
     async fn process(
-        &self,
+        &'a self,
         selector: &Self::ReportSelector,
     ) -> Result<DapLeaderProcessTelemetry, DapAbort> {
         let mut telem = DapLeaderProcessTelemetry::default();
@@ -484,12 +494,14 @@ pub trait DapLeader<S>: DapAuthorizedSender<S> + DapAggregator<S> {
         for (task_id, reports) in self.get_reports(selector).await?.into_iter() {
             let task_config = self
                 .get_task_config_for(&task_id)
+                .await?
                 .ok_or(DapAbort::UnrecognizedTask)?;
 
             telem.reports_processed += reports.len() as u64;
             if !reports.is_empty() {
-                telem.reports_aggregated +=
-                    self.run_agg_job(&task_id, task_config, reports).await?;
+                telem.reports_aggregated += self
+                    .run_agg_job(&task_id, task_config.as_ref(), reports)
+                    .await?;
             }
         }
 
@@ -500,10 +512,11 @@ pub trait DapLeader<S>: DapAuthorizedSender<S> + DapAggregator<S> {
         for (collect_id, collect_req) in self.get_pending_collect_jobs().await? {
             let task_config = self
                 .get_task_config_for(&collect_req.task_id)
+                .await?
                 .ok_or(DapAbort::UnrecognizedTask)?;
 
             telem.reports_collected += self
-                .run_collect_job(&collect_id, task_config, collect_req)
+                .run_collect_job(&collect_id, task_config.as_ref(), collect_req)
                 .await?;
         }
 
@@ -513,7 +526,7 @@ pub trait DapLeader<S>: DapAuthorizedSender<S> + DapAggregator<S> {
 
 /// DAP Helper functionality.
 #[async_trait(?Send)]
-pub trait DapHelper<S>: DapAggregator<S> {
+pub trait DapHelper<'a, S>: DapAggregator<'a, S> {
     /// Update the metadata for the given set of report shares, marking them as aggregated. The
     /// return value is the subset of the report shares that will not be aggregated due to a
     /// transition failure. This occurs if, for example, the report was previously aggregated, but
@@ -545,7 +558,7 @@ pub trait DapHelper<S>: DapAggregator<S> {
     /// AggregateContinueReq and the response is an AggregateResp.
     ///
     /// This is called during the Initialization and Continuation phases.
-    async fn http_post_aggregate(&self, req: &DapRequest<S>) -> Result<DapResponse, DapAbort> {
+    async fn http_post_aggregate(&'a self, req: &DapRequest<S>) -> Result<DapResponse, DapAbort> {
         // Check whether the DAP version indicated by the sender is supported.
         if req.version == DapVersion::Unknown {
             return Err(DapAbort::InvalidProtocolVersion);
@@ -558,9 +571,11 @@ pub trait DapHelper<S>: DapAggregator<S> {
         match req.media_type {
             Some(MEDIA_TYPE_AGG_INIT_REQ) => {
                 let agg_init_req = AggregateInitializeReq::get_decoded(&req.payload)?;
-                let task_config = self
+                let wrapped_task_config = self
                     .get_task_config_for(&agg_init_req.task_id)
+                    .await?
                     .ok_or(DapAbort::UnrecognizedTask)?;
+                let task_config = wrapped_task_config.as_ref();
                 let helper_state =
                     self.get_helper_state(&agg_init_req.task_id, &agg_init_req.agg_job_id);
 
@@ -630,9 +645,11 @@ pub trait DapHelper<S>: DapAggregator<S> {
             }
             Some(MEDIA_TYPE_AGG_CONT_REQ) => {
                 let agg_cont_req = AggregateContinueReq::get_decoded(&req.payload)?;
-                let task_config = self
+                let wrapped_task_config = self
                     .get_task_config_for(&agg_cont_req.task_id)
+                    .await?
                     .ok_or(DapAbort::UnrecognizedTask)?;
+                let task_config = wrapped_task_config.as_ref();
 
                 // Check whether the DAP version in the request matches the task config.
                 if task_config.version != req.version {
@@ -671,7 +688,7 @@ pub trait DapHelper<S>: DapAggregator<S> {
     ///
     /// This is called during the Collection phase.
     async fn http_post_aggregate_share(
-        &self,
+        &'a self,
         req: &DapRequest<S>,
     ) -> Result<DapResponse, DapAbort> {
         // Check whether the DAP version indicated by the sender is supported.
@@ -684,9 +701,11 @@ pub trait DapHelper<S>: DapAggregator<S> {
         }
 
         let agg_share_req = AggregateShareReq::get_decoded(&req.payload)?;
-        let task_config = self
+        let wrapped_task_config = self
             .get_task_config_for(&agg_share_req.task_id)
+            .await?
             .ok_or(DapAbort::UnrecognizedTask)?;
+        let task_config = wrapped_task_config.as_ref();
         check_batch_param!(
             task_config,
             agg_share_req.batch_interval,
diff --git a/daphne/src/vdaf/mod.rs b/daphne/src/vdaf/mod.rs
index 085e57c..705f6fa 100644
--- a/daphne/src/vdaf/mod.rs
+++ b/daphne/src/vdaf/mod.rs
@@ -213,7 +213,7 @@ impl VdafConfig {
     #[allow(clippy::too_many_arguments)]
     pub(crate) async fn consume_report_share(
         &self,
-        decrypter: &impl HpkeDecrypter,
+        decrypter: &impl HpkeDecrypter<'_>,
         is_leader: bool,
         verify_key: &VdafVerifyKey,
         task_id: &Id,
@@ -285,7 +285,7 @@ impl VdafConfig {
     /// * `reports` is the set of reports uploaded by Clients.
     pub(crate) async fn produce_agg_init_req(
         &self,
-        decrypter: &impl HpkeDecrypter,
+        decrypter: &impl HpkeDecrypter<'_>,
         verify_key: &VdafVerifyKey,
         task_id: &Id,
         agg_job_id: &Id,
@@ -384,7 +384,7 @@ impl VdafConfig {
     /// is the transition failure the Helper is to transmit.
     pub(crate) async fn handle_agg_init_req(
         &self,
-        decrypter: &impl HpkeDecrypter,
+        decrypter: &impl HpkeDecrypter<'_>,
         verify_key: &VdafVerifyKey,
         agg_init_req: &AggregateInitializeReq,
     ) -> Result<DapHelperTransition<AggregateResp>, DapAbort> {
@@ -734,7 +734,7 @@ impl VdafConfig {
     // TODO spec: Allow the collector to have multiple HPKE public keys (the way Aggregators do).
     pub async fn consume_encrypted_agg_shares(
         &self,
-        decrypter: &impl HpkeDecrypter,
+        decrypter: &impl HpkeDecrypter<'_>,
         task_id: &Id,
         batch_interval: &Interval,
         encrypted_agg_shares: Vec<HpkeCiphertext>,
diff --git a/daphne_worker/src/config.rs b/daphne_worker/src/config.rs
index e91399e..f21390d 100644
--- a/daphne_worker/src/config.rs
+++ b/daphne_worker/src/config.rs
@@ -17,7 +17,7 @@ use daphne::{
     auth::BearerToken,
     constants,
     hpke::HpkeReceiverConfig,
-    messages::{Id, Interval, Nonce},
+    messages::{HpkeConfig, Id, Interval, Nonce},
     DapError, DapGlobalConfig, DapRequest, DapTaskConfig, DapVersion,
 };
 use matchit::Router;
@@ -264,10 +264,10 @@ impl<D> DaphneWorkerConfig<D> {
 
     // Get a reference to the HPKE receiver configs, ensuring that the config indicated by
     // `hpke_config_id` is cached (if it exists).
-    pub(crate) async fn cached_hpke_receiver_configs(
+    pub(crate) async fn hpke_receiver_config(
         &self,
         hpke_config_id: u8,
-    ) -> Result<RwLockReadGuard<HashMap<u8, HpkeReceiverConfig>>> {
+    ) -> Result<Option<GuardedHpkeReceiverConfig>> {
         // If the config is cached, then return immediately.
         {
             let hpke_receiver_configs = self.hpke_receiver_configs.read().map_err(|e| {
@@ -278,7 +278,10 @@ impl<D> DaphneWorkerConfig<D> {
             })?;
 
             if hpke_receiver_configs.get(&hpke_config_id).is_some() {
-                return Ok(hpke_receiver_configs);
+                return Ok(Some(GuardedHpkeReceiverConfig {
+                    hpke_receiver_configs,
+                    hpke_config_id,
+                }));
             }
         }
 
@@ -305,7 +308,15 @@ impl<D> DaphneWorkerConfig<D> {
                 e
             ))
         })?;
-        Ok(hpke_receiver_configs)
+
+        if hpke_receiver_configs.get(&hpke_config_id).is_some() {
+            Ok(Some(GuardedHpkeReceiverConfig {
+                hpke_receiver_configs,
+                hpke_config_id,
+            }))
+        } else {
+            Ok(None)
+        }
     }
 
     /// Clear all persistant durable objects storage.
@@ -366,6 +377,30 @@ impl<D> DaphneWorkerConfig<D> {
     }
 }
 
+/// RwLockReadGuard'ed HPKE receiver config.
+pub(crate) struct GuardedHpkeReceiverConfig<'a> {
+    hpke_receiver_configs: RwLockReadGuard<'a, HashMap<u8, HpkeReceiverConfig>>,
+    hpke_config_id: u8,
+}
+
+impl AsRef<HpkeConfig> for GuardedHpkeReceiverConfig<'_> {
+    fn as_ref(&self) -> &HpkeConfig {
+        &self
+            .hpke_receiver_configs
+            .get(&self.hpke_config_id)
+            .unwrap()
+            .config
+    }
+}
+
+impl GuardedHpkeReceiverConfig<'_> {
+    pub(crate) fn decryptor(&self) -> &HpkeReceiverConfig {
+        self.hpke_receiver_configs
+            .get(&self.hpke_config_id)
+            .unwrap()
+    }
+}
+
 /// Deployment types for Daphne-Worker. This defines overrides used to control inter-Aggregator
 /// communication.
 #[derive(Debug, Default)]
diff --git a/daphne_worker/src/dap.rs b/daphne_worker/src/dap.rs
index 45f1d25..2e71179 100644
--- a/daphne_worker/src/dap.rs
+++ b/daphne_worker/src/dap.rs
@@ -7,7 +7,10 @@
 //! draft-ietf-ppm-dap-01.
 
 use crate::{
-    config::{DaphneWorkerConfig, DaphneWorkerDeployment, KV_KEY_PREFIX_HPKE_RECEIVER_CONFIG},
+    config::{
+        DaphneWorkerConfig, DaphneWorkerDeployment, GuardedHpkeReceiverConfig,
+        KV_KEY_PREFIX_HPKE_RECEIVER_CONFIG,
+    },
     dap_err,
     durable::{
         aggregate_store::{
@@ -40,8 +43,8 @@ use daphne::{
     constants,
     hpke::{HpkeDecrypter, HpkeReceiverConfig},
     messages::{
-        CollectReq, CollectResp, HpkeCiphertext, HpkeConfig, Id, Interval, Nonce, Report,
-        ReportShare, TransitionFailure,
+        CollectReq, CollectResp, HpkeCiphertext, Id, Interval, Nonce, Report, ReportShare,
+        TransitionFailure,
     },
     roles::{DapAggregator, DapAuthorizedSender, DapHelper, DapLeader},
     DapAggregateShare, DapCollectJob, DapError, DapGlobalConfig, DapHelperState, DapOutputShare,
@@ -66,11 +69,13 @@ pub(crate) fn dap_response_to_worker(resp: DapResponse) -> Result<Response> {
 }
 
 #[async_trait(?Send)]
-impl<D> HpkeDecrypter for DaphneWorkerConfig<D> {
+impl<'a, D> HpkeDecrypter<'a> for DaphneWorkerConfig<D> {
+    type WrappedHpkeConfig = GuardedHpkeReceiverConfig<'a>;
+
     async fn get_hpke_config_for(
-        &self,
+        &'a self,
         _task_id: &Id,
-    ) -> std::result::Result<Option<HpkeConfig>, DapError> {
+    ) -> std::result::Result<Option<GuardedHpkeReceiverConfig<'a>>, DapError> {
         let kv_store = self.kv("DAP_HPKE_RECEIVER_CONFIG_STORE").map_err(dap_err)?;
         let keys = kv_store
             .list()
@@ -80,7 +85,7 @@ impl<D> HpkeDecrypter for DaphneWorkerConfig<D> {
             .await
             .map_err(|e| DapError::Fatal(format!("kv_store: {}", e)))?;
 
-        if keys.keys.is_empty() {
+        let hpke_config_id = if keys.keys.is_empty() {
             // Generate a new random HPKE config ID and key for KV.
             let hpke_config_id: u8 = rand::thread_rng().gen();
             let new_key = format!("{}/{}", KV_KEY_PREFIX_HPKE_RECEIVER_CONFIG, hpke_config_id);
@@ -106,20 +111,16 @@ impl<D> HpkeDecrypter for DaphneWorkerConfig<D> {
                 .await
                 .map_err(|e| DapError::Fatal(format!("kv_store: {}", e)))?;
 
-            Ok(Some(hpke_receiver_config.config))
+            hpke_config_id
         } else {
-            // Return the first HPKE receiver config.
-            let builder = kv_store.get(&keys.keys[0].name);
-            if let Some(hpke_receiver_config) = builder
-                .json::<HpkeReceiverConfig>()
-                .await
-                .map_err(|e| DapError::Fatal(format!("kv_store: {}", e)))?
-            {
-                Ok(Some(hpke_receiver_config.config))
-            } else {
-                Ok(None)
-            }
-        }
+            // Return the first HPKE receiver config in the list.
+            parse_hpke_config_id_from_kv_key_name(keys.keys[0].name.as_str())?
+        };
+
+        Ok(self
+            .hpke_receiver_config(hpke_config_id)
+            .await
+            .map_err(dap_err)?)
     }
 
     async fn can_hpke_decrypt(
@@ -127,11 +128,11 @@ impl<D> HpkeDecrypter for DaphneWorkerConfig<D> {
         _task_id: &Id,
         config_id: u8,
     ) -> std::result::Result<bool, DapError> {
-        let hpke_receiver_configs = self
-            .cached_hpke_receiver_configs(config_id)
+        Ok(self
+            .hpke_receiver_config(config_id)
             .await
-            .map_err(dap_err)?;
-        Ok(hpke_receiver_configs.get(&config_id).is_some())
+            .map_err(dap_err)?
+            .is_some())
     }
 
     async fn hpke_decrypt(
@@ -141,18 +142,43 @@ impl<D> HpkeDecrypter for DaphneWorkerConfig<D> {
         aad: &[u8],
         ciphertext: &HpkeCiphertext,
     ) -> std::result::Result<Vec<u8>, DapError> {
-        let hpke_receiver_configs = self
-            .cached_hpke_receiver_configs(ciphertext.config_id)
+        if let Some(hpke_receiver_config) = self
+            .hpke_receiver_config(ciphertext.config_id)
             .await
-            .map_err(dap_err)?;
-        if let Some(hpke_receiver_config) = hpke_receiver_configs.get(&ciphertext.config_id) {
-            Ok(hpke_receiver_config.decrypt(info, aad, &ciphertext.enc, &ciphertext.payload)?)
+            .map_err(dap_err)?
+        {
+            Ok(hpke_receiver_config.decryptor().decrypt(
+                info,
+                aad,
+                &ciphertext.enc,
+                &ciphertext.payload,
+            )?)
         } else {
             Err(DapError::Transition(TransitionFailure::HpkeUnknownConfigId))
         }
     }
 }
 
+fn parse_hpke_config_id_from_kv_key_name(name: &str) -> std::result::Result<u8, DapError> {
+    let mut iter = name.split('/');
+    if let Some(prefix) = iter.next() {
+        if prefix == KV_KEY_PREFIX_HPKE_RECEIVER_CONFIG {
+            if let Some(config_id_str) = iter.next() {
+                if iter.next().is_none() {
+                    if let Ok(config_id) = config_id_str.parse::<u8>() {
+                        return Ok(config_id);
+                    }
+                }
+            }
+        }
+    }
+
+    Err(DapError::Fatal(format!(
+        "malformed kv_store HPKE receiver config key: '{}'",
+        name
+    )))
+}
+
 #[async_trait(?Send)]
 impl<D> BearerTokenProvider for DaphneWorkerConfig<D> {
     async fn get_leader_bearer_token_for(
@@ -186,7 +212,9 @@ impl<D> DapAuthorizedSender<BearerToken> for DaphneWorkerConfig<D> {
 }
 
 #[async_trait(?Send)]
-impl<D> DapAggregator<BearerToken> for DaphneWorkerConfig<D> {
+impl<'a, D> DapAggregator<'a, BearerToken> for DaphneWorkerConfig<D> {
+    type WrappedDapTaskConfig = &'a DapTaskConfig;
+
     async fn authorized(
         &self,
         req: &DapRequest<BearerToken>,
@@ -198,8 +226,11 @@ impl<D> DapAggregator<BearerToken> for DaphneWorkerConfig<D> {
         &self.global_config
     }
 
-    fn get_task_config_for(&self, task_id: &Id) -> Option<&DapTaskConfig> {
-        self.tasks.get(task_id)
+    async fn get_task_config_for(
+        &'a self,
+        task_id: &Id,
+    ) -> std::result::Result<Option<&'a DapTaskConfig>, DapError> {
+        Ok(self.tasks.get(task_id))
     }
 
     fn get_current_time(&self) -> u64 {
@@ -357,7 +388,7 @@ impl<D> DapAggregator<BearerToken> for DaphneWorkerConfig<D> {
 }
 
 #[async_trait(?Send)]
-impl<D> DapLeader<BearerToken> for DaphneWorkerConfig<D> {
+impl<'a, D> DapLeader<'a, BearerToken> for DaphneWorkerConfig<D> {
     type ReportSelector = InternalAggregateInfo;
 
     async fn put_report(&self, report: &Report) -> std::result::Result<(), DapError> {
@@ -610,7 +641,7 @@ impl<D> DapLeader<BearerToken> for DaphneWorkerConfig<D> {
 }
 
 #[async_trait(?Send)]
-impl<D> DapHelper<BearerToken> for DaphneWorkerConfig<D> {
+impl<'a, D> DapHelper<'a, BearerToken> for DaphneWorkerConfig<D> {
     async fn mark_aggregated(
         &self,
         task_id: &Id,


Test patch:
diff --git a/daphne/src/roles_test.rs b/daphne/src/roles_test.rs
index f57cac4..1efcb70 100644
--- a/daphne/src/roles_test.rs
+++ b/daphne/src/roles_test.rs
@@ -44,7 +44,7 @@ macro_rules! get_reports {
 impl MockAggregator {
     fn gen_test_upload_req(&self, report: Report) -> DapRequest<BearerToken> {
         let task_id = self.nominal_task_id();
-        let task_config = self.get_task_config_for(task_id).unwrap();
+        let task_config = self.tasks.get(task_id).unwrap();
         let version = task_config.version.clone();
 
         DapRequest {
@@ -59,7 +59,7 @@ impl MockAggregator {
     fn gen_test_agg_init_req(&self, report_shares: Vec<ReportShare>) -> DapRequest<BearerToken> {
         let mut rng = thread_rng();
         let task_id = self.nominal_task_id();
-        let task_config = self.get_task_config_for(task_id).unwrap();
+        let task_config = self.tasks.get(task_id).unwrap();
         let version = task_config.version.clone();
 
         DapRequest {
@@ -83,7 +83,7 @@ impl MockAggregator {
         transitions: Vec<Transition>,
     ) -> DapRequest<BearerToken> {
         let task_id = self.nominal_task_id();
-        let task_config = self.get_task_config_for(task_id).unwrap();
+        let task_config = self.tasks.get(task_id).unwrap();
         let version = task_config.version.clone();
 
         DapRequest {
@@ -106,7 +106,7 @@ impl MockAggregator {
         checksum: [u8; 32],
     ) -> DapRequest<BearerToken> {
         let task_id = self.nominal_task_id();
-        let task_config = self.get_task_config_for(task_id).unwrap();
+        let task_config = self.tasks.get(task_id).unwrap();
         let version = task_config.version.clone();
 
         DapRequest {
@@ -127,7 +127,7 @@ impl MockAggregator {
 
     fn gen_test_collect_req(&self) -> DapRequest<BearerToken> {
         let task_id = self.nominal_task_id();
-        let task_config = self.get_task_config_for(task_id).unwrap();
+        let task_config = self.tasks.get(task_id).unwrap();
         let version = task_config.version.clone();
 
         DapRequest {
@@ -481,7 +481,7 @@ async fn http_post_aggregate_failure_report_replayed() {
 async fn http_post_aggregate_failure_batch_collected() {
     let helper = MockAggregator::new();
     let task_id = helper.nominal_task_id();
-    let task_config = helper.get_task_config_for(task_id).unwrap();
+    let task_config = helper.tasks.get(task_id).unwrap();
 
     let report = helper.gen_test_report(task_id);
     let report_shares = vec![ReportShare {
@@ -655,7 +655,7 @@ async fn get_reports_empty_response() {
 async fn poll_collect_job_test_results() {
     let leader = MockAggregator::new();
     let task_id = leader.nominal_task_id();
-    let task_config = leader.get_task_config_for(task_id).unwrap();
+    let task_config = leader.tasks.get(task_id).unwrap();
 
     // Collector: Create a CollectReq.
     let collector_collect_req = CollectReq {
@@ -714,7 +714,7 @@ async fn poll_collect_job_test_results() {
 async fn http_post_collect_fail_invalid_batch_interval() {
     let leader = MockAggregator::new();
     let task_id = leader.nominal_task_id();
-    let task_config = leader.get_task_config_for(task_id).unwrap();
+    let task_config = leader.tasks.get(task_id).unwrap();
     let now = leader.now;
 
     // Collector: Create a CollectReq with a very large batch interval.
@@ -799,7 +799,7 @@ async fn http_post_collect_fail_invalid_batch_interval() {
 async fn http_post_collect_succeed_max_batch_interval() {
     let leader = MockAggregator::new();
     let task_id = leader.nominal_task_id();
-    let task_config = leader.get_task_config_for(task_id).unwrap();
+    let task_config = leader.tasks.get(task_id).unwrap();
     let now = leader.now;
 
     // Collector: Create a CollectReq with a very large batch interval.
@@ -831,7 +831,7 @@ async fn http_post_collect_succeed_max_batch_interval() {
 async fn http_post_collect_fail_overlapping_batch_interval() {
     let leader = MockAggregator::new();
     let task_id = leader.nominal_task_id();
-    let task_config = leader.get_task_config_for(task_id).unwrap();
+    let task_config = leader.tasks.get(task_id).unwrap();
     let helper = MockAggregator::new();
     let now = leader.now;
 
@@ -897,7 +897,7 @@ async fn http_post_collect_fail_overlapping_batch_interval() {
 async fn http_post_collect_success() {
     let leader = MockAggregator::new();
     let task_id = leader.nominal_task_id();
-    let task_config = leader.get_task_config_for(task_id).unwrap();
+    let task_config = leader.tasks.get(task_id).unwrap();
 
     // Collector: Create a CollectReq.
     let collector_collect_req = CollectReq {
@@ -942,7 +942,7 @@ async fn http_post_collect_success() {
 async fn http_post_fail_wrong_dap_version() {
     let leader = MockAggregator::new();
     let task_id = leader.nominal_task_id();
-    let task_config = leader.get_task_config_for(task_id).unwrap();
+    let task_config = leader.tasks.get(task_id).unwrap();
 
     // Send a request with the wrong DAP version.
     let report = leader.gen_test_report(task_id);
@@ -975,7 +975,7 @@ async fn successful_http_post_upload() {
 async fn e2e() {
     let leader = MockAggregator::new();
     let task_id = leader.nominal_task_id();
-    let task_config = leader.get_task_config_for(task_id).unwrap();
+    let task_config = leader.tasks.get(task_id).unwrap();
 
     let helper = MockAggregator::new();
 
diff --git a/daphne/src/testing.rs b/daphne/src/testing.rs
index b8d277c..660cb68 100644
--- a/daphne/src/testing.rs
+++ b/daphne/src/testing.rs
@@ -94,7 +94,7 @@ pub(crate) struct MockAggregateInfo {
 pub(crate) struct MockAggregator {
     pub(crate) now: u64,
     pub(crate) global_config: DapGlobalConfig,
-    tasks: HashMap<Id, DapTaskConfig>,
+    pub(crate) tasks: HashMap<Id, DapTaskConfig>,
     hpke_receiver_config_list: Vec<HpkeReceiverConfig>,
     pub(crate) report_store: Arc<Mutex<HashMap<Id, ReportStore>>>,
     leader_state_store: Arc<Mutex<HashMap<Id, LeaderState>>>,
@@ -196,14 +196,19 @@ impl BearerTokenProvider for MockAggregator {
 }
 
 #[async_trait(?Send)]
-impl HpkeDecrypter for MockAggregator {
-    async fn get_hpke_config_for(&self, _task_id: &Id) -> Result<Option<HpkeConfig>, DapError> {
+impl<'a> HpkeDecrypter<'a> for MockAggregator {
+    type WrappedHpkeConfig = &'a HpkeConfig;
+
+    async fn get_hpke_config_for(
+        &'a self,
+        _task_id: &Id,
+    ) -> Result<Option<&'a HpkeConfig>, DapError> {
         if self.hpke_receiver_config_list.is_empty() {
             return Ok(None);
         }
 
         // Always advertise the first HPKE config in the list.
-        Ok(Some(self.hpke_receiver_config_list[0].config.clone()))
+        Ok(Some(&self.hpke_receiver_config_list[0].config))
     }
 
     async fn can_hpke_decrypt(&self, _task_id: &Id, config_id: u8) -> Result<bool, DapError> {
@@ -239,7 +244,9 @@ impl DapAuthorizedSender<BearerToken> for MockAggregator {
 }
 
 #[async_trait(?Send)]
-impl DapAggregator<BearerToken> for MockAggregator {
+impl<'a> DapAggregator<'a, BearerToken> for MockAggregator {
+    type WrappedDapTaskConfig = &'a DapTaskConfig;
+
     async fn authorized(&self, req: &DapRequest<BearerToken>) -> Result<bool, DapError> {
         self.bearer_token_authorized(req).await
     }
@@ -248,8 +255,11 @@ impl DapAggregator<BearerToken> for MockAggregator {
         &self.global_config
     }
 
-    fn get_task_config_for(&self, task_id: &Id) -> Option<&DapTaskConfig> {
-        self.tasks.get(task_id)
+    async fn get_task_config_for(
+        &'a self,
+        task_id: &Id,
+    ) -> Result<Option<&'a DapTaskConfig>, DapError> {
+        Ok(self.tasks.get(task_id))
     }
 
     fn get_current_time(&self) -> u64 {
@@ -288,6 +298,7 @@ impl DapAggregator<BearerToken> for MockAggregator {
     ) -> Result<(), DapError> {
         let task_config = self
             .get_task_config_for(task_id)
+            .await?
             .ok_or_else(|| DapError::fatal("task not found"))?;
 
         let agg_shares =
@@ -376,7 +387,7 @@ impl DapAggregator<BearerToken> for MockAggregator {
 }
 
 #[async_trait(?Send)]
-impl DapHelper<BearerToken> for MockAggregator {
+impl<'a> DapHelper<'a, BearerToken> for MockAggregator {
     async fn mark_aggregated(
         &self,
         task_id: &Id,
@@ -384,6 +395,7 @@ impl DapHelper<BearerToken> for MockAggregator {
     ) -> Result<HashMap<Nonce, TransitionFailure>, DapError> {
         let task_config = self
             .get_task_config_for(task_id)
+            .await?
             .ok_or_else(|| DapError::fatal("task not found"))?;
         let mut early_fails: HashMap<Nonce, TransitionFailure> = HashMap::new();
 
@@ -482,13 +494,14 @@ impl DapHelper<BearerToken> for MockAggregator {
 }
 
 #[async_trait(?Send)]
-impl DapLeader<BearerToken> for MockAggregator {
+impl<'a> DapLeader<'a, BearerToken> for MockAggregator {
     type ReportSelector = MockAggregateInfo;
 
     async fn put_report(&self, report: &Report) -> Result<(), DapError> {
         let task_id = &report.task_id;
         let task_config = self
             .get_task_config_for(task_id)
+            .await?
             .ok_or_else(|| DapError::fatal("task not found"))?;
         let bucket_info = BucketInfo::new(task_config, task_id, &report.nonce);
 
@@ -566,6 +579,7 @@ impl DapLeader<BearerToken> for MockAggregator {
         let mut rng = thread_rng();
         let task_config = self
             .get_task_config_for(&collect_req.task_id)
+            .await?
             .ok_or_else(|| DapError::fatal("task not found"))?;
 
         let mut leader_state_store_mutex_guard = self