Repository: cloudflare/daphne
Pull request #643 (2024-07-30 16:36:51)

Add api for costumizing the audit log
Also removes `fn host()` from the Aggregator trait as well as the
unused env variable from the service config.



Patch:
diff --git a/crates/daphne-server/src/lib.rs b/crates/daphne-server/src/lib.rs
index f1c08f3..d992c96 100644
--- a/crates/daphne-server/src/lib.rs
+++ b/crates/daphne-server/src/lib.rs
@@ -3,7 +3,12 @@
 
 use std::sync::Arc;
 
-use daphne::{auth::BearerToken, roles::leader::in_memory_leader::InMemoryLeaderState, DapError};
+use daphne::{
+    audit_log::{AuditLog, NoopAuditLog},
+    auth::BearerToken,
+    roles::leader::in_memory_leader::InMemoryLeaderState,
+    DapError,
+};
 use daphne_service_utils::{config::DaphneServiceConfig, metrics::DaphneServiceMetrics};
 use futures::lock::Mutex;
 use serde::{Deserialize, Serialize};
@@ -50,7 +55,6 @@ mod storage_proxy_connection;
 ///     default_num_agg_span_shards: NonZeroUsize::new(2).unwrap(),
 /// };
 /// let service_config = DaphneServiceConfig {
-///     env: "some-machine-identifier".into(),
 ///     role: DapRole::Helper,
 ///     global,
 ///     base_url: None,
@@ -74,6 +78,7 @@ pub struct App {
     cache: RwLock<kv::Cache>,
     metrics: Box<dyn DaphneServiceMetrics>,
     service_config: DaphneServiceConfig,
+    audit_log: Box<dyn AuditLog + Send + Sync>,
 
     /// Volatile memory for the Leader, including the work queue, pending reports, and pending
     /// colleciton requests. Note that in a production Leader, it is necessary to store this state
@@ -112,11 +117,19 @@ impl App {
             http: reqwest::Client::new(),
             cache: Default::default(),
             metrics: Box::new(daphne_service_metrics),
+            audit_log: Box::new(NoopAuditLog),
             service_config,
             test_leader_state: Default::default(),
         })
     }
 
+    pub fn set_audit_log<A>(&mut self, audit_log: A)
+    where
+        A: AuditLog + Send + Sync + 'static,
+    {
+        self.audit_log = Box::new(audit_log);
+    }
+
     pub(crate) fn durable(&self) -> Do<'_> {
         Do::new(&self.storage_proxy_config, &self.http)
     }
diff --git a/crates/daphne-server/src/roles/aggregator.rs b/crates/daphne-server/src/roles/aggregator.rs
index 17ee7fd..aa00463 100644
--- a/crates/daphne-server/src/roles/aggregator.rs
+++ b/crates/daphne-server/src/roles/aggregator.rs
@@ -5,7 +5,7 @@ use std::{borrow::Cow, future::ready, num::NonZeroUsize, ops::Range, time::Syste
 
 use axum::async_trait;
 use daphne::{
-    audit_log::{AuditLog, NoopAuditLog},
+    audit_log::AuditLog,
     auth::{BearerToken, BearerTokenProvider},
     error::DapAbort,
     fatal_error,
@@ -384,11 +384,7 @@ impl DapAggregator<DaphneAuth> for crate::App {
     }
 
     fn audit_log(&self) -> &dyn AuditLog {
-        &NoopAuditLog
-    }
-
-    fn host(&self) -> &str {
-        &self.service_config.env
+        &*self.audit_log
     }
 }
 
diff --git a/crates/daphne-service-utils/src/config.rs b/crates/daphne-service-utils/src/config.rs
index d025426..4b729c3 100644
--- a/crates/daphne-service-utils/src/config.rs
+++ b/crates/daphne-service-utils/src/config.rs
@@ -36,8 +36,6 @@ pub type HpkeRecieverConfigList = Vec<HpkeReceiverConfig>;
 /// Daphne service configuration, including long-lived parameters used across DAP tasks.
 #[derive(Serialize, Deserialize, Debug, Clone)]
 pub struct DaphneServiceConfig {
-    pub env: String,
-
     /// Indicates the role the service should play.
     pub role: DapRole,
 
diff --git a/crates/daphne/src/audit_log.rs b/crates/daphne/src/audit_log.rs
index a4fbf68..88ed4b4 100644
--- a/crates/daphne/src/audit_log.rs
+++ b/crates/daphne/src/audit_log.rs
@@ -3,19 +3,13 @@
 
 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,
+        vdaf_step: u8,
     );
 }
 
@@ -25,11 +19,10 @@ 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,
+        _vdaf_step: u8,
     ) {
     }
 }
diff --git a/crates/daphne/src/roles/aggregator.rs b/crates/daphne/src/roles/aggregator.rs
index d9b9828..0c1cb2b 100644
--- a/crates/daphne/src/roles/aggregator.rs
+++ b/crates/daphne/src/roles/aggregator.rs
@@ -161,10 +161,6 @@ pub trait DapAggregator<S: Sync>: HpkeProvider + DapReportInitializer + Sized {
 
     /// Access the audit log.
     fn audit_log(&self) -> &dyn AuditLog;
-
-    /// Return the hostname of the request URL. The value is "unspecified-host" if the URL does not
-    /// indicate a hostname.
-    fn host(&self) -> &str;
 }
 
 /// Handle request for the Aggregator's HPKE configuration.
diff --git a/crates/daphne/src/roles/helper.rs b/crates/daphne/src/roles/helper.rs
index 64d5bc9..6b081f7 100644
--- a/crates/daphne/src/roles/helper.rs
+++ b/crates/daphne/src/roles/helper.rs
@@ -9,7 +9,6 @@ use tracing::error;
 
 use super::{check_batch, check_request_content_type, resolve_taskprov, DapAggregator};
 use crate::{
-    audit_log::AggregationJobAuditAction,
     constants::DapMediaType,
     error::DapAbort,
     messages::{
@@ -78,7 +77,6 @@ pub async fn handle_agg_job_init_req<'req, S: Sync, A: DapHelper<S>>(
         &agg_job_init_req.agg_param,
     )?;
 
-    let prep_init_count = agg_job_init_req.prep_inits.len();
     let part_batch_sel = agg_job_init_req.part_batch_sel.clone();
     let initialized_reports = task_config
         .consume_agg_job_req(
@@ -107,11 +105,10 @@ pub async fn handle_agg_job_init_req<'req, S: Sync, A: DapHelper<S>>(
     };
 
     aggregator.audit_log().on_aggregation_job(
-        aggregator.host(),
         task_id,
         task_config,
-        prep_init_count as u64,
-        AggregationJobAuditAction::Init,
+        agg_job_resp.transitions.len() as u64,
+        0, /* vdaf step */
     );
 
     metrics.inbound_req_inc(DaphneRequestType::Aggregate);


Test patch:
diff --git a/crates/daphne/src/testing/mod.rs b/crates/daphne/src/testing/mod.rs
index b685ed6..798a9f8 100644
--- a/crates/daphne/src/testing/mod.rs
+++ b/crates/daphne/src/testing/mod.rs
@@ -7,7 +7,7 @@
 pub mod report_generator;
 
 use crate::{
-    audit_log::{AggregationJobAuditAction, AuditLog},
+    audit_log::AuditLog,
     auth::{BearerToken, BearerTokenProvider},
     constants::DapMediaType,
     fatal_error,
@@ -488,11 +488,10 @@ impl MockAuditLog {
 impl AuditLog for MockAuditLog {
     fn on_aggregation_job(
         &self,
-        _host: &str,
         _task_id: &TaskId,
         _task_config: &DapTaskConfig,
         _report_count: u64,
-        _action: AggregationJobAuditAction,
+        _vdaf_step: u8,
     ) {
         self.0.fetch_add(1, Ordering::Relaxed);
     }
@@ -1003,10 +1002,6 @@ impl DapAggregator<BearerToken> for InMemoryAggregator {
     fn audit_log(&self) -> &dyn AuditLog {
         &self.audit_log
     }
-
-    fn host(&self) -> &str {
-        "unspecified-host"
-    }
 }
 
 #[async_trait]