Repository: cloudflare/daphne
Pull request #334 (2023-06-02 09:47:38)

Adds new AuditLog trait to daphne
Enables implementations backed by services like: https://developers.cloudflare.com/analytics/analytics-engine/

* For now, there is just a basic handler for successful aggregation job requests.
* Adds a `Display` to `VdafConfig` for a reasonably succinct string repr, useful for event logging.
* Adds some assertions on the handler being invoked

mendess: Isn't it possible to leverage tracing somehow to achieve this? Adding this effectively means we have to visibility mechanisms

Patch:
diff --git a/daphne/src/audit_log.rs b/daphne/src/audit_log.rs
new file mode 100644
index 0000000..a4fbf68
--- /dev/null
+++ b/daphne/src/audit_log.rs
@@ -0,0 +1,35 @@
+// Copyright (c) 2023 Cloudflare, Inc. All rights reserved.
+// SPDX-License-Identifier: BSD-3-Clause
+
+use crate::{messages::TaskId, DapTaskConfig};
+
+pub enum AggregationJobAuditAction {
+    Init,
+    Continue,
+}
+
+pub trait AuditLog {
+    fn on_aggregation_job(
+        &self,
+        host: &str,
+        task_id: &TaskId,
+        task_config: &DapTaskConfig,
+        report_count: u64,
+        action: AggregationJobAuditAction,
+    );
+}
+
+/// Default implementation of the trait, which is a no-op.
+pub struct NoopAuditLog;
+
+impl AuditLog for NoopAuditLog {
+    fn on_aggregation_job(
+        &self,
+        _host: &str,
+        _task_id: &TaskId,
+        _task_config: &DapTaskConfig,
+        _report_count: u64,
+        _action: AggregationJobAuditAction,
+    ) {
+    }
+}
diff --git a/daphne/src/lib.rs b/daphne/src/lib.rs
index d83722e..c3bf3a1 100644
--- a/daphne/src/lib.rs
+++ b/daphne/src/lib.rs
@@ -62,7 +62,7 @@ use std::{
     borrow::Cow,
     cmp::{max, min},
     collections::{HashMap, HashSet},
-    fmt::Debug,
+    fmt::{Debug, Display},
 };
 use url::Url;
 
@@ -702,7 +702,7 @@ pub enum DapHelperTransition<M: Debug> {
     Finish(Vec<DapOutputShare>, M),
 }
 
-/// Specificaiton of a concrete VDAF.
+/// Specification of a concrete VDAF.
 #[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
 #[serde(rename_all = "snake_case")]
 pub enum VdafConfig {
@@ -718,6 +718,15 @@ impl std::str::FromStr for VdafConfig {
     }
 }
 
+impl Display for VdafConfig {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            VdafConfig::Prio3(prio3_config) => write!(f, "Prio3({prio3_config})"),
+            VdafConfig::Prio2 { dimension } => write!(f, "Prio2({dimension})"),
+        }
+    }
+}
+
 /// Supported data types for prio3.
 #[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
 #[serde(rename_all = "snake_case")]
@@ -739,6 +748,19 @@ pub enum Prio3Config {
     SumVec { bits: usize, len: usize },
 }
 
+impl Display for Prio3Config {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Prio3Config::Count => write!(f, "Count"),
+            Prio3Config::Histogram { buckets } => {
+                write!(f, "Histogram({})", buckets.len())
+            }
+            Prio3Config::Sum { bits } => write!(f, "Sum({bits})"),
+            Prio3Config::SumVec { bits, len } => write!(f, "SumVec({bits},{len})"),
+        }
+    }
+}
+
 /// DAP sender role.
 #[derive(Clone, Copy, Debug)]
 pub enum DapSender {
@@ -951,6 +973,7 @@ impl MetaAggregationJobId<'_> {
 }
 
 pub mod aborts;
+pub mod audit_log;
 pub mod auth;
 pub mod constants;
 #[cfg(test)]
diff --git a/daphne/src/roles.rs b/daphne/src/roles.rs
index a295e61..563bb56 100644
--- a/daphne/src/roles.rs
+++ b/daphne/src/roles.rs
@@ -4,6 +4,7 @@
 //! Trait definitions for Daphne backends.
 
 use crate::{
+    audit_log::{AggregationJobAuditAction, AuditLog},
     constants::DapMediaType,
     hpke::HpkeDecrypter,
     messages::{
@@ -200,6 +201,9 @@ pub trait DapAggregator<S>: HpkeDecrypter + Sized {
 
     /// Access the Prometheus metrics.
     fn metrics(&self) -> &DaphneMetrics;
+
+    /// Access the audit log.
+    fn audit_log(&self) -> &dyn AuditLog;
 }
 
 macro_rules! leader_post {
@@ -945,6 +949,14 @@ pub trait DapHelper<S>: DapAggregator<S> {
                     }
                 };
 
+                self.audit_log().on_aggregation_job(
+                    req.host(),
+                    task_id,
+                    task_config,
+                    agg_job_init_req.report_shares.len() as u64,
+                    AggregationJobAuditAction::Init,
+                );
+
                 metrics.agg_job_started_inc();
                 metrics.inbound_req_inc(DaphneRequestType::Aggregate);
                 Ok(DapResponse {
@@ -1016,6 +1028,14 @@ pub trait DapHelper<S>: DapAggregator<S> {
                     }
                 };
 
+                self.audit_log().on_aggregation_job(
+                    req.host(),
+                    task_id,
+                    task_config,
+                    out_shares_count,
+                    AggregationJobAuditAction::Continue,
+                );
+
                 metrics.report_inc_by("aggregated", out_shares_count);
                 metrics.agg_job_completed_inc();
                 metrics.inbound_req_inc(DaphneRequestType::Aggregate);
diff --git a/daphne_worker/src/config.rs b/daphne_worker/src/config.rs
index cfea532..3dba107 100644
--- a/daphne_worker/src/config.rs
+++ b/daphne_worker/src/config.rs
@@ -21,6 +21,7 @@ use crate::{
 };
 use daphne::{
     aborts::DapAbort,
+    audit_log::AuditLog,
     auth::BearerToken,
     constants::DapMediaType,
     hpke::HpkeReceiverConfig,
@@ -415,6 +416,9 @@ pub(crate) struct DaphneWorkerRequestState<'srv> {
 
     /// Error reporting for Daphne internal errors.
     pub(crate) error_reporter: &'srv dyn ErrorReporter,
+
+    /// Audit logging
+    pub(crate) audit_log: &'srv dyn AuditLog,
 }
 
 impl<'srv> DaphneWorkerRequestState<'srv> {
@@ -422,6 +426,7 @@ impl<'srv> DaphneWorkerRequestState<'srv> {
         isolate_state: &'srv DaphneWorkerIsolateState,
         req: &Request,
         error_reporter: &'srv dyn ErrorReporter,
+        audit_log: &'srv dyn AuditLog,
     ) -> Result<Self> {
         let prometheus_registry = Registry::new();
         let metrics = DaphneWorkerMetrics::register(&prometheus_registry, None)
@@ -438,6 +443,7 @@ impl<'srv> DaphneWorkerRequestState<'srv> {
             metrics,
             host,
             error_reporter,
+            audit_log,
         })
     }
 
diff --git a/daphne_worker/src/dap.rs b/daphne_worker/src/dap.rs
index ddd31d0..7082cc2 100644
--- a/daphne_worker/src/dap.rs
+++ b/daphne_worker/src/dap.rs
@@ -46,6 +46,7 @@ use crate::{
 };
 use async_trait::async_trait;
 use daphne::{
+    audit_log::AuditLog,
     auth::{BearerToken, BearerTokenProvider},
     constants::DapMediaType,
     hpke::HpkeDecrypter,
@@ -697,6 +698,10 @@ impl<'srv> DapAggregator<DaphneWorkerAuth> for DaphneWorker<'srv> {
     fn metrics(&self) -> &DaphneMetrics {
         &self.state.metrics.daphne
     }
+
+    fn audit_log(&self) -> &dyn AuditLog {
+        self.state.audit_log
+    }
 }
 
 #[async_trait(?Send)]
diff --git a/daphne_worker/src/lib.rs b/daphne_worker/src/lib.rs
index 7ebe177..8d80359 100644
--- a/daphne_worker/src/lib.rs
+++ b/daphne_worker/src/lib.rs
@@ -160,6 +160,7 @@ use crate::{
 };
 use daphne::{
     aborts::DapAbort,
+    audit_log::{AuditLog, NoopAuditLog},
     auth::BearerToken,
     constants::DapMediaType,
     messages::{CollectionJobId, Duration, TaskId, Time},
@@ -195,12 +196,16 @@ pub struct DaphneWorkerRouter<'srv> {
 
     /// Error reporting for Daphne. By default is a no-op.
     pub error_reporter: &'srv dyn error_reporting::ErrorReporter,
+
+    /// Audit log, used to record statistics of tasks processed.
+    pub audit_log: &'srv dyn AuditLog,
 }
 
 impl<'srv> Default for DaphneWorkerRouter<'srv> {
     fn default() -> Self {
         Self {
             error_reporter: &error_reporting::NoopErrorReporter {},
+            audit_log: &NoopAuditLog,
             enable_internal_test: false,
             enable_default_response: false,
         }
@@ -263,7 +268,8 @@ impl DaphneWorkerRouter<'_> {
         } else {
             ISOLATE_STATE.get_or_try_init(|| DaphneWorkerIsolateState::from_worker_env(&env))?
         };
-        let state = DaphneWorkerRequestState::new(shared_state, &req, self.error_reporter)?;
+        let state =
+            DaphneWorkerRequestState::new(shared_state, &req, self.error_reporter, self.audit_log)?;
 
         let router = Router::with_data(&state)
             .get_async("/:version/hpke_config", |req, ctx| async move {


Test patch:
diff --git a/daphne/src/roles_test.rs b/daphne/src/roles_test.rs
index 7890142..5831be1 100644
--- a/daphne/src/roles_test.rs
+++ b/daphne/src/roles_test.rs
@@ -17,7 +17,9 @@ use crate::{
     roles::{early_metadata_check, DapAggregator, DapAuthorizedSender, DapHelper, DapLeader},
     taskprov::TaskprovVersion,
     test_version, test_versions,
-    testing::{AggStore, DapBatchBucketOwned, MockAggregator, MockAggregatorReportSelector},
+    testing::{
+        AggStore, DapBatchBucketOwned, MockAggregator, MockAggregatorReportSelector, MockAuditLog,
+    },
     vdaf::VdafVerifyKey,
     DapAbort, DapAggregateShare, DapCollectJob, DapGlobalConfig, DapMeasurement, DapQueryConfig,
     DapRequest, DapResource, DapTaskConfig, DapVersion, MetaAggregationJobId, Prio3Config,
@@ -169,6 +171,7 @@ impl Test {
             collector_hpke_config: collector_hpke_receiver_config.config.clone(),
             taskprov_vdaf_verify_key_init,
             metrics: DaphneMetrics::register(&prometheus_registry, Some("test_helper")).unwrap(),
+            audit_log: MockAuditLog::default(),
             peer: None,
         });
 
@@ -189,6 +192,7 @@ impl Test {
             collector_hpke_config: collector_hpke_receiver_config.config,
             taskprov_vdaf_verify_key_init,
             metrics: DaphneMetrics::register(&prometheus_registry, Some("test_leader")).unwrap(),
+            audit_log: MockAuditLog::default(),
             peer: Some(Arc::clone(&helper)),
         });
 
@@ -541,6 +545,8 @@ async fn http_post_aggregate_invalid_batch_sel(version: DapVersion) {
         t.helper.http_post_aggregate(&req).await.unwrap_err(),
         DapAbort::QueryMismatch { .. }
     );
+
+    assert_eq!(t.helper.audit_log.invocations(), 0);
 }
 
 async_test_versions! { http_post_aggregate_invalid_batch_sel }
@@ -564,6 +570,8 @@ async fn http_post_aggregate_init_unauthorized_request(version: DapVersion) {
         t.helper.http_post_aggregate(&req).await,
         Err(DapAbort::UnauthorizedRequest { .. })
     );
+
+    assert_eq!(t.helper.audit_log.invocations(), 0);
 }
 
 async_test_versions! { http_post_aggregate_init_unauthorized_request }
@@ -589,6 +597,8 @@ async fn http_post_aggregate_init_expired_task(version: DapVersion) {
         agg_job_resp.transitions[0].var,
         TransitionVar::Failed(TransitionFailure::TaskExpired)
     );
+
+    assert_eq!(t.helper.audit_log.invocations(), 1);
 }
 
 async_test_versions! { http_post_aggregate_init_expired_task }
@@ -615,6 +625,10 @@ async fn http_post_aggregate_bad_round(version: DapVersion) {
         _ => panic!("agg_job_id resource missing!"),
     };
     let resp = t.helper.http_post_aggregate(&req).await.unwrap();
+
+    // AggregationJobInitReq succeeds
+    assert_eq!(t.helper.audit_log.invocations(), 1);
+
     let agg_job_resp = AggregationJobResp::get_decoded(&resp.payload).unwrap();
     assert_eq!(agg_job_resp.transitions.len(), 1);
     assert_matches!(agg_job_resp.transitions[0].var, TransitionVar::Continued(_));
@@ -630,6 +644,9 @@ async fn http_post_aggregate_bad_round(version: DapVersion) {
         t.helper.http_post_aggregate(&req).await,
         Err(DapAbort::RoundMismatch { .. })
     );
+
+    // AggregationJobContinueReq fails
+    assert_eq!(t.helper.audit_log.invocations(), 1);
 }
 
 async_test_versions! { http_post_aggregate_bad_round }
@@ -656,6 +673,10 @@ async fn http_post_aggregate_zero_round(version: DapVersion) {
         _ => panic!("agg_job_id resource missing!"),
     };
     let resp = t.helper.http_post_aggregate(&req).await.unwrap();
+
+    // AggregationJobInitReq succeeds
+    assert_eq!(t.helper.audit_log.invocations(), 1);
+
     let agg_job_resp = AggregationJobResp::get_decoded(&resp.payload).unwrap();
     assert_eq!(agg_job_resp.transitions.len(), 1);
     assert_matches!(agg_job_resp.transitions[0].var, TransitionVar::Continued(_));
@@ -671,6 +692,9 @@ async fn http_post_aggregate_zero_round(version: DapVersion) {
         t.helper.http_post_aggregate(&req).await,
         Err(DapAbort::UnrecognizedMessage)
     );
+
+    // AggregationJobContinueReq fails
+    assert_eq!(t.helper.audit_log.invocations(), 1);
 }
 
 async_test_versions! { http_post_aggregate_zero_round }
@@ -744,6 +768,8 @@ async fn http_post_aggregate_cont_unauthorized_request(version: DapVersion) {
         t.helper.http_post_aggregate(&req).await,
         Err(DapAbort::UnauthorizedRequest { .. })
     );
+
+    assert_eq!(t.helper.audit_log.invocations(), 0);
 }
 
 async_test_versions! { http_post_aggregate_cont_unauthorized_request }
@@ -1038,6 +1064,8 @@ async fn http_post_aggregate_failure_batch_collected(version: DapVersion) {
             .unwrap();
     let transition = &agg_job_resp.transitions[0];
 
+    assert_eq!(t.helper.audit_log.invocations(), 1);
+
     // Expect failure due to report store marked as collected.
     assert_matches!(
         transition.var,
@@ -1071,9 +1099,13 @@ async fn http_post_aggregate_abort_helper_state_overwritten(version: DapVersion)
     // Send aggregate request.
     let _ = t.helper.http_post_aggregate(&req).await;
 
+    assert_eq!(t.helper.audit_log.invocations(), 1);
+
     // Send another aggregate request.
     let err = t.helper.http_post_aggregate(&req).await.unwrap_err();
 
+    assert_eq!(t.helper.audit_log.invocations(), 1);
+
     // Expect failure due to overwriting existing helper state.
     assert_matches!(err, DapAbort::BadRequest(e) =>
         assert_eq!(e, "unexpected message for aggregation job (already exists)")
@@ -1089,6 +1121,8 @@ async fn http_post_aggregate_fail_send_cont_req(version: DapVersion) {
         .gen_test_agg_job_cont_req(&agg_job_id, Vec::default(), version)
         .await;
 
+    assert_eq!(t.helper.audit_log.invocations(), 0);
+
     // Send aggregate continue request to helper.
     let err = t.helper.http_post_aggregate(&req).await.unwrap_err();
 
diff --git a/daphne/src/testing.rs b/daphne/src/testing.rs
index f8abb35..862ba25 100644
--- a/daphne/src/testing.rs
+++ b/daphne/src/testing.rs
@@ -4,6 +4,7 @@
 //! Mock backend functionality to test DAP protocol.
 
 use crate::{
+    audit_log::{AggregationJobAuditAction, AuditLog},
     auth::{BearerToken, BearerTokenProvider},
     constants::DapMediaType,
     hpke::{HpkeDecrypter, HpkeReceiverConfig},
@@ -27,7 +28,10 @@ use std::{
     collections::{HashMap, HashSet, VecDeque},
     hash::Hash,
     ops::DerefMut,
-    sync::{Arc, Mutex},
+    sync::{
+        atomic::{AtomicU32, Ordering},
+        Arc, Mutex,
+    },
     time::SystemTime,
 };
 use url::Url;
@@ -84,6 +88,29 @@ impl<'a> DapBatchBucket<'a> {
 
 pub(crate) struct MockAggregatorReportSelector(pub(crate) TaskId);
 
+#[derive(Default)]
+pub(crate) struct MockAuditLog(AtomicU32);
+
+impl MockAuditLog {
+    #[cfg(test)]
+    pub(crate) fn invocations(&self) -> u32 {
+        self.0.load(Ordering::Relaxed)
+    }
+}
+
+impl AuditLog for MockAuditLog {
+    fn on_aggregation_job(
+        &self,
+        _host: &str,
+        _task_id: &TaskId,
+        _task_config: &DapTaskConfig,
+        _report_count: u64,
+        _action: AggregationJobAuditAction,
+    ) {
+        self.0.fetch_add(1, Ordering::Relaxed);
+    }
+}
+
 pub(crate) struct MockAggregator {
     pub(crate) global_config: DapGlobalConfig,
     pub(crate) tasks: Arc<Mutex<HashMap<TaskId, DapTaskConfig>>>,
@@ -97,6 +124,7 @@ pub(crate) struct MockAggregator {
     pub(crate) collector_hpke_config: HpkeConfig,
     pub(crate) taskprov_vdaf_verify_key_init: [u8; 32],
     pub(crate) metrics: DaphneMetrics,
+    pub(crate) audit_log: MockAuditLog,
 
     // Leader: Reference to peer. Used to simulate HTTP requests from Leader to Helper, i.e.,
     // implement `DapLeader::send_http_post()` for `MockAggregator`. Not set by the Helper.
@@ -544,6 +572,10 @@ impl DapAggregator<BearerToken> for MockAggregator {
     fn metrics(&self) -> &DaphneMetrics {
         &self.metrics
     }
+
+    fn audit_log(&self) -> &dyn AuditLog {
+        &self.audit_log
+    }
 }
 
 #[async_trait(?Send)]