Repository: cloudflare/daphne
Pull request #32 (2022-06-30 18:25:27)

Implement HpkeDecryptor functionality for MockAggregator
This is the start of the process of addressing issue #20.

As for the first step, these couple of commits implements the HpkeDecryptor functionality for MockAggregator.
The test code for the HpkeDecryptor functionality is also added.

nakatsuka-y: I've added a new test case that ultimately uses [`cosume_report_share`](https://github.com/cloudflare/daphne/blob/82bbcc1a39258a7a3b11476ebc1a72e5d993200f/daphne/src/vdaf/mod.rs#L186), which must decrypt the input share data to consume it.

The check is based on whether the decryption fails.
I am currently working on a test case to check whether the decryption succeeds.
nakatsuka-y: I've implemented a test case where the decryption succeeds.
Hopefully the two test cases I've added in the past few commits allows others to understand how the mock_backend can be used.

Patch:
diff --git a/daphne/src/hpke.rs b/daphne/src/hpke.rs
index 6144f72..fe2fcc7 100644
--- a/daphne/src/hpke.rs
+++ b/daphne/src/hpke.rs
@@ -148,3 +148,10 @@ pub trait HpkeDecrypter {
         ciphertext: &HpkeCiphertext,
     ) -> Result<Vec<u8>, DapError>;
 }
+
+/// Struct that combines HpkeConfig and HpkeSecretKey
+#[derive(Deserialize, Serialize)]
+pub struct HpkeReceiverConfig {
+    pub config: HpkeConfig,
+    pub secret_key: HpkeSecretKey,
+}
diff --git a/daphne/src/messages.rs b/daphne/src/messages.rs
index 4218423..934ee20 100644
--- a/daphne/src/messages.rs
+++ b/daphne/src/messages.rs
@@ -474,7 +474,7 @@ impl Decode for AggregateShareResp {
 }
 
 /// Codepoint for KEM schemes compatible with HPKE.
-#[derive(Clone, Copy, Debug, PartialEq)]
+#[derive(Clone, Copy, Deserialize, Serialize, Debug, PartialEq)]
 pub enum HpkeKemId {
     X25519HkdfSha256,
     NotImplemented(u16),
@@ -505,7 +505,7 @@ impl Decode for HpkeKemId {
 }
 
 /// Codepoint for KDF schemes compatible with HPKE.
-#[derive(Clone, Copy, Debug, PartialEq)]
+#[derive(Clone, Copy, Deserialize, Serialize, Debug, PartialEq)]
 pub enum HpkeKdfId {
     HkdfSha256,
     NotImplemented(u16),
@@ -536,7 +536,7 @@ impl Decode for HpkeKdfId {
 }
 
 /// Codepoint for AEAD schemes compatible with HPKE.
-#[derive(Clone, Copy, Debug, PartialEq)]
+#[derive(Clone, Copy, Deserialize, Serialize, Debug, PartialEq)]
 pub enum HpkeAeadId {
     Aes128Gcm,
     NotImplemented(u16),
@@ -567,7 +567,7 @@ impl Decode for HpkeAeadId {
 }
 
 /// The HPKE public key configuration of a Server.
-#[derive(Clone, Debug, PartialEq)]
+#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
 pub struct HpkeConfig {
     pub id: u8,
     pub kem_id: HpkeKemId,
@@ -577,6 +577,7 @@ pub struct HpkeConfig {
     // serialized key. We can't do this with rust-hpke because <X25519HkdfSha256 as Kem>::PublicKey
     // doesn't implement Debug. Eventually we'll replace rust-hpke with a more ergonomic
     // implementation that does. For now we'll eat the copy.
+    #[serde(with = "hex")]
     pub(crate) public_key: Vec<u8>,
 }
 
diff --git a/daphne/src/roles.rs b/daphne/src/roles.rs
index 355269c..ec99adc 100644
--- a/daphne/src/roles.rs
+++ b/daphne/src/roles.rs
@@ -415,8 +415,10 @@ pub trait DapHelper<S>: DapAggregator<S> {
         agg_job_id: &Id,
     ) -> Result<DapHelperState, DapError>;
 
-    /// Handle an HTTP POST to `/aggregate`. The input is an AggregateInitializeReq and the
-    /// response is an AggregateResp.
+    /// Handle an HTTP POST to `/aggregate`. The input is either an AggregateInitializeReq or
+    /// 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> {
         if !self.authorized(req).await? {
             return Err(DapAbort::UnauthorizedRequest);
@@ -433,6 +435,7 @@ pub trait DapHelper<S>: DapAggregator<S> {
                     .mark_aggregated(&agg_init_req.task_id, &agg_init_req.report_shares)
                     .await?;
 
+                // This is where the encrypted report share is decrypted.
                 let transition = task_config.vdaf.handle_agg_init_req(
                     self,
                     &task_config.vdaf_verify_key,
@@ -494,6 +497,8 @@ pub trait DapHelper<S>: DapAggregator<S> {
 
     /// Handle an HTTP POST to `/aggregate_share`. The input is an AggregateShareReq and the
     /// response is an AggregateShareResp.
+    ///
+    /// This is called during the Collection phase.
     async fn http_post_aggregate_share(
         &self,
         req: &DapRequest<S>,


Test patch:
diff --git a/daphne/src/roles_test.rs b/daphne/src/roles_test.rs
index b1ac239..2570b0e 100644
--- a/daphne/src/roles_test.rs
+++ b/daphne/src/roles_test.rs
@@ -4,20 +4,21 @@
 use crate::{
     auth::{BearerToken, BearerTokenProvider},
     constants::{MEDIA_TYPE_AGG_INIT_REQ, MEDIA_TYPE_AGG_SHARE_REQ, MEDIA_TYPE_COLLECT_REQ},
-    hpke::HpkeDecrypter,
+    hpke::{HpkeDecrypter, HpkeReceiverConfig, HpkeSecretKey},
     messages::{
-        AggregateInitializeReq, AggregateShareReq, CollectReq, CollectResp, HpkeCiphertext,
-        HpkeConfig, Id, Interval, Nonce, Report, ReportShare, TransitionFailure,
+        AggregateInitializeReq, AggregateResp, AggregateShareReq, CollectReq, CollectResp,
+        HpkeCiphertext, HpkeConfig, Id, Interval, Nonce, Report, ReportShare, TransitionFailure,
+        TransitionVar,
     },
     roles::{DapAggregator, DapAuthorizedSender, DapHelper, DapLeader},
-    DapAbort, DapAggregateShare, DapCollectJob, DapError, DapHelperState, DapOutputShare,
-    DapRequest, DapResponse, DapTaskConfig,
+    DapAbort, DapAggregateShare, DapCollectJob, DapError, DapHelperState, DapMeasurement,
+    DapOutputShare, DapRequest, DapResponse, DapTaskConfig, Prio3Config, VdafConfig,
 };
 use assert_matches::assert_matches;
 use async_trait::async_trait;
-use prio::codec::Encode;
+use prio::codec::{Decode, Encode};
 use rand::prelude::*;
-use std::collections::HashMap;
+use std::{collections::HashMap, vec};
 use url::Url;
 
 const TASK_LIST: &str = r#"{
@@ -28,27 +29,85 @@ const TASK_LIST: &str = r#"{
         "min_batch_duration": 3600,
         "min_batch_size": 100,
         "vdaf": {
-            "prio3": {
-                "histogram": {
-                    "buckets": [77, 999]
-                }
-            }
+            "prio3": "count"
         },
         "vdaf_verify_key": "1fd8d30dc0e0b7ac81f0050fcab0782d"
     }
 }"#;
 
+// TODO(nakatsuka-y) Merge secret key into `HpkeReceiverConfig` and remove `HpkeSecretKey`.
+// This removes the redundant "id" field in both of these structs. See issue #12.
+const HPKE_RECEIVER_CONFIG_LIST: &str = r#"[
+    {
+        "config": {
+            "id": 23,
+            "kem_id": "X25519HkdfSha256",
+            "kdf_id": "HkdfSha256",
+            "aead_id": "Aes128Gcm",
+            "public_key": "5dc71373c6aa7b0af67944a370ab96d8b8216832579c19159ca35d10f25a2765"
+        },
+        "secret_key": {
+            "id": 23,
+            "sk": "888e94344585f44530d03e250268be6c6a5caca5314513dcec488cc431486c69"
+        }
+    },
+    {
+        "config": {
+            "id": 14,
+            "kem_id": "X25519HkdfSha256",
+            "kdf_id": "HkdfSha256",
+            "aead_id": "Aes128Gcm",
+            "public_key": "b07126295bcfcdeaec61b310fd7ffbf8c6ca7f6c17e3e0a80a5405a242e5084b"
+        },
+        "secret_key": {
+            "id": 14,
+            "sk": "b809a4df399548f56c3a15ebaa4925dd292637f0b7e2f6bc3ba60376b69aa05e"
+        }
+    }
+]"#;
+
 const LEADER_BEARER_TOKEN: &str = "ivA1e7LpnySDNn1AulaZggFLQ1n7jZ8GWOUO7GY4hgs=";
 const COLLECTOR_BEARER_TOKEN: &str = "syfRfvcvNFF5MJk4Y-B7xjRIqD_iNzhaaEB9mYqO9hk=";
 
 struct MockAggregator {
     tasks: HashMap<Id, DapTaskConfig>,
+    hpke_config_list: Vec<HpkeConfig>,
+    hpke_secret_key_list: Vec<HpkeSecretKey>,
 }
 
 impl MockAggregator {
     fn new() -> Self {
+        // Construct task list
         let tasks = serde_json::from_str(TASK_LIST).expect("failed to parse task list");
-        Self { tasks }
+
+        // Construct HPKE receiver config List
+        let hpke_receiver_config_list: Vec<HpkeReceiverConfig> =
+            serde_json::from_str(HPKE_RECEIVER_CONFIG_LIST)
+                .expect("failed to parse hpke_receiver_config_list");
+
+        let mut hpke_config_list: Vec<HpkeConfig> =
+            Vec::with_capacity(hpke_receiver_config_list.len());
+        let mut hpke_secret_key_list: Vec<HpkeSecretKey> =
+            Vec::with_capacity(hpke_receiver_config_list.len());
+        for receiver_config in hpke_receiver_config_list {
+            hpke_config_list.push(receiver_config.config);
+            hpke_secret_key_list.push(receiver_config.secret_key);
+        }
+
+        Self {
+            tasks,
+            hpke_config_list,
+            hpke_secret_key_list,
+        }
+    }
+
+    fn get_hpke_secret_key_for(&self, hpke_config_id: u8) -> Option<&HpkeSecretKey> {
+        for hpke_secret_key in self.hpke_secret_key_list.iter() {
+            if hpke_config_id == hpke_secret_key.id {
+                return Some(hpke_secret_key);
+            }
+        }
+        None
     }
 
     /// Task to use for nominal tests.
@@ -77,21 +136,30 @@ impl BearerTokenProvider for MockAggregator {
 
 impl HpkeDecrypter for MockAggregator {
     fn get_hpke_config_for(&self, _task_id: &Id) -> Option<&HpkeConfig> {
-        unreachable!("not implemented");
+        if self.hpke_config_list.is_empty() {
+            return None;
+        }
+
+        // Advertise the first HPKE config in the list.
+        Some(&self.hpke_config_list[0])
     }
 
-    fn can_hpke_decrypt(&self, _task_id: &Id, _config_id: u8) -> bool {
-        unreachable!("not implemented");
+    fn can_hpke_decrypt(&self, _task_id: &Id, config_id: u8) -> bool {
+        self.get_hpke_secret_key_for(config_id).is_some()
     }
 
     fn hpke_decrypt(
         &self,
         _task_id: &Id,
-        _info: &[u8],
-        _aad: &[u8],
-        _ciphertext: &HpkeCiphertext,
+        info: &[u8],
+        aad: &[u8],
+        ciphertext: &HpkeCiphertext,
     ) -> Result<Vec<u8>, DapError> {
-        unreachable!("not implemented");
+        if let Some(hpke_secret_key) = self.get_hpke_secret_key_for(ciphertext.config_id) {
+            Ok(hpke_secret_key.decrypt(info, aad, &ciphertext.enc, &ciphertext.payload)?)
+        } else {
+            Err(DapError::Transition(TransitionFailure::HpkeUnknownConfigId))
+        }
     }
 }
 
@@ -149,7 +217,10 @@ impl DapHelper<BearerToken> for MockAggregator {
         _task_id: &Id,
         _report_shares: &[ReportShare],
     ) -> Result<HashMap<Nonce, TransitionFailure>, DapError> {
-        unreachable!("not implemented");
+        // Return empty HashMap (for now).
+        // TODO(nakatsuka-y) Implement correct functionality.
+        let early_fails: HashMap<Nonce, TransitionFailure> = HashMap::new();
+        return Ok(early_fails);
     }
 
     async fn put_helper_state(
@@ -158,7 +229,9 @@ impl DapHelper<BearerToken> for MockAggregator {
         _agg_job_id: &Id,
         _helper_state: &DapHelperState,
     ) -> Result<(), DapError> {
-        unreachable!("not implemented");
+        // Return empty Ok (for now).
+        // TODO(nakatsuka-y) Implement correct functionality.
+        Ok(())
     }
 
     async fn get_helper_state(
@@ -318,3 +391,172 @@ async fn http_post_collect_unauthorized_request() {
         Err(DapAbort::UnauthorizedRequest)
     );
 }
+
+#[tokio::test]
+async fn http_post_aggregate_invalid_ciphertext() {
+    let helper = MockAggregator::new();
+    let task_id = helper.nominal_task_id();
+    let task_config = helper.get_task_config_for(task_id).unwrap();
+
+    let req = DapRequest {
+        media_type: Some(MEDIA_TYPE_AGG_INIT_REQ),
+        payload: AggregateInitializeReq {
+            task_id: task_id.clone(),
+            agg_job_id: Id([1; 32]),
+            agg_param: b"this is an aggregation parameter".to_vec(),
+            report_shares: vec![ReportShare {
+                nonce: Nonce {
+                    time: 1637361337,
+                    rand: 10496152761178246059,
+                },
+                ignored_extensions: b"these are extensions".to_vec(),
+                encrypted_input_share: HpkeCiphertext {
+                    config_id: 23,
+                    enc: b"invalid encapsulated key".to_vec(),
+                    payload: b"invalid ciphertext".to_vec(),
+                },
+            }],
+        }
+        .get_encoded(),
+        url: task_config.helper_url.join("/aggregate").unwrap(),
+        sender_auth: Some(BearerToken::from(LEADER_BEARER_TOKEN.to_string())),
+    };
+
+    // Get AggregateResp and then extract the transition data from inside.
+    let agg_resp =
+        AggregateResp::get_decoded(&helper.http_post_aggregate(&req).await.unwrap().payload)
+            .unwrap();
+    let transition = &agg_resp.transitions[0];
+
+    // Expect failure due to invalid ciphertext.
+    assert_matches!(
+        transition.var,
+        TransitionVar::Failed(TransitionFailure::HpkeDecryptError)
+    );
+}
+
+#[tokio::test]
+async fn http_post_aggregate_valid_ciphertext() {
+    let helper = MockAggregator::new();
+    let task_id = helper.nominal_task_id();
+    let task_config = helper.get_task_config_for(task_id).unwrap();
+
+    // Construct HPKE receiver config List.
+    let hpke_receiver_config_list: Vec<HpkeReceiverConfig> =
+        serde_json::from_str(HPKE_RECEIVER_CONFIG_LIST)
+            .expect("failed to parse hpke_receiver_config_list");
+
+    // Construct HPKE config list.
+    let mut hpke_config_list = Vec::with_capacity(hpke_receiver_config_list.len());
+    for receiver_config in hpke_receiver_config_list {
+        hpke_config_list.push(receiver_config.config);
+    }
+
+    // Construct report.
+    let vdaf_config: &VdafConfig = &VdafConfig::Prio3(Prio3Config::Count);
+    let report = vdaf_config
+        .produce_report(
+            &hpke_config_list,
+            1637361337,
+            task_id,
+            DapMeasurement::U64(1),
+        )
+        .unwrap();
+
+    // Construct DapRequest.
+    let req = DapRequest {
+        media_type: Some(MEDIA_TYPE_AGG_INIT_REQ),
+        payload: AggregateInitializeReq {
+            task_id: task_id.clone(),
+            agg_job_id: Id([1; 32]),
+            agg_param: b"this is an aggregation parameter".to_vec(),
+            report_shares: vec![ReportShare {
+                nonce: report.nonce,
+                ignored_extensions: report.ignored_extensions,
+                // 1st share is for Leader and the rest is for Helpers (note that there is only 1 helper).
+                encrypted_input_share: report.encrypted_input_shares[1].clone(),
+            }],
+        }
+        .get_encoded(),
+        url: task_config.helper_url.join("/aggregate").unwrap(),
+        sender_auth: Some(BearerToken::from(LEADER_BEARER_TOKEN.to_string())),
+    };
+
+    // Get AggregateResp and then extract the transition data from inside.
+    let agg_resp =
+        AggregateResp::get_decoded(&helper.http_post_aggregate(&req).await.unwrap().payload)
+            .unwrap();
+    let transition = &agg_resp.transitions[0];
+
+    // Expect success due to valid ciphertext.
+    assert_matches!(transition.var, TransitionVar::Continued(_));
+}
+
+#[test]
+fn hpke_decrypter() {
+    // Construct mock aggregator.
+    let aggregator = MockAggregator::new();
+    let task_id = aggregator.nominal_task_id();
+
+    // Initialize variables for mock report.
+    let info = b"info string";
+    let aad = b"associated data";
+    let plaintext = b"plaintext";
+    let hpke_receiver_config_list: Vec<HpkeReceiverConfig> =
+        serde_json::from_str(HPKE_RECEIVER_CONFIG_LIST)
+            .expect("failed to parse hpke_receiver_config_list");
+    let config = &hpke_receiver_config_list[0].config;
+    let (enc, ciphertext) = config.encrypt(info, aad, plaintext).unwrap();
+
+    // Construct mock report.
+    let report = Report {
+        task_id: Id([23; 32]),
+        nonce: Nonce {
+            time: 1637364244,
+            rand: 10496152761178246059,
+        },
+        ignored_extensions: b"some extension".to_vec(),
+        encrypted_input_shares: vec![HpkeCiphertext {
+            config_id: 23,
+            enc: enc,
+            payload: ciphertext,
+        }],
+    };
+
+    // Expect false due to non-existing config ID.
+    assert_eq!(aggregator.can_hpke_decrypt(&task_id, 0), false);
+
+    // Expect true due to existing config ID.
+    assert_eq!(
+        aggregator.can_hpke_decrypt(&task_id, report.encrypted_input_shares[0].config_id),
+        true
+    );
+
+    // Expect decryption to fail.
+    assert_matches!(
+        aggregator.hpke_decrypt(
+            &report.task_id,
+            info,
+            aad,
+            &HpkeCiphertext {
+                config_id: 0,
+                enc: vec![],
+                payload: b"ciphertext".to_vec(),
+            }
+        ),
+        Err(DapError::Transition(TransitionFailure::HpkeUnknownConfigId))
+    );
+
+    // Expect decryption to succeed.
+    assert_eq!(
+        aggregator
+            .hpke_decrypt(
+                &report.task_id,
+                info,
+                aad,
+                &report.encrypted_input_shares[0]
+            )
+            .unwrap(),
+        plaintext
+    );
+}