Add support for configuring tasks with DP strategy (#1173)

* Add support for configuring tasks with DP strategy * Check entire validation error object in tests * Add support to divviup-client * Add re-exports * Add support to the CLI * Mark new enums as non-exhaustive * Update OpenAPI schema
divviup · Jul 23, 2024 · 93f503b · 93f503b
1 parent db5b5da
commit 93f503b
Show file tree

Hide file tree

Showing 16 changed files with 800 additions and 103 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -2,15 +2,15 @@
 members = [".", "migration", "client", "test-support", "cli"]
 
 [workspace.package]
-version = "0.3.22"
+version = "0.4.0"
 edition = "2021"
 license = "MPL-2.0"
 homepage = "https://divviup.org"
 repository = "https://github.com/divviup/divviup-api"
 
 [workspace.dependencies]
-divviup-client = { path = "./client", version = "0.3.22" }
-divviup-cli = { path = "./cli", version = "0.3.22" }
+divviup-client = { path = "./client", version = "0.4.0" }
+divviup-cli = { path = "./cli", version = "0.4.0" }
 divviup-api.path = "."
 test-support.path = "./test-support"
 

diff --git a/cli/src/tasks.rs b/cli/src/tasks.rs
@@ -1,6 +1,9 @@
 use crate::{CliResult, DetermineAccountId, Error, Output};
 use clap::Subcommand;
-use divviup_client::{DivviupClient, Histogram, NewTask, SumVec, Uuid, Vdaf};
+use divviup_client::{
+    dp_strategy::{self, PureDpBudget, PureDpDiscreteLaplace},
+    BigUint, DivviupClient, Histogram, NewTask, Ratio, SumVec, Uuid, Vdaf,
+};
 use humantime::{Duration, Timestamp};
 use std::time::SystemTime;
 use time::OffsetDateTime;
@@ -14,6 +17,11 @@ pub enum VdafName {
     SumVec,
 }
 
+#[derive(clap::ValueEnum, Clone, Debug)]
+pub enum DpStrategy {
+    PureDpDiscreteLaplace,
+}
+
 #[derive(Subcommand, Debug)]
 pub enum TaskAction {
     /// list all tasks for the target account
@@ -52,6 +60,10 @@ pub enum TaskAction {
         bits: Option<u8>,
         #[arg(long)]
         chunk_length: Option<u64>,
+        #[arg(long, requires = "differential_privacy_epsilon")]
+        differential_privacy_strategy: Option<DpStrategy>,
+        #[arg(long, requires = "differential_privacy_strategy")]
+        differential_privacy_epsilon: Option<f64>,
     },
 
     /// rename a task
@@ -108,24 +120,65 @@ impl TaskAction {
                 length,
                 bits,
                 chunk_length,
+                differential_privacy_strategy,
+                differential_privacy_epsilon,
             } => {
                 let vdaf = match vdaf {
-                    VdafName::Count => Vdaf::Count,
+                    VdafName::Count => {
+                        if differential_privacy_strategy.is_some()
+                            || differential_privacy_epsilon.is_some()
+                        {
+                            return Err(Error::Other(
+                                "differential privacy noise is not yet supported with Prio3Count"
+                                    .into(),
+                            ));
+                        }
+                        Vdaf::Count
+                    }
                     VdafName::Histogram => {
+                        let dp_strategy =
+                            match (differential_privacy_strategy, differential_privacy_epsilon) {
+                                (None, None) => dp_strategy::Prio3Histogram::NoDifferentialPrivacy,
+                                (None, Some(_)) => {
+                                    return Err(Error::Other(
+                                        "missing differential-privacy-strategy".into(),
+                                    ))
+                                }
+                                (Some(_), None) => {
+                                    return Err(Error::Other(
+                                        "missing differential-privacy-epsilon".into(),
+                                    ))
+                                }
+                                (Some(DpStrategy::PureDpDiscreteLaplace), Some(epsilon)) => {
+                                    dp_strategy::Prio3Histogram::PureDpDiscreteLaplace(
+                                        PureDpDiscreteLaplace {
+                                            budget: PureDpBudget {
+                                                epsilon: float_to_biguint_ratio(epsilon)
+                                                    .ok_or_else(|| {
+                                                        Error::Other("invalid epsilon".into())
+                                                    })?,
+                                            },
+                                        },
+                                    )
+                                }
+                            };
                         match (length, categorical_buckets, continuous_buckets) {
                             (Some(length), None, None) => Vdaf::Histogram(Histogram::Length {
                                 length,
                                 chunk_length,
+                                dp_strategy,
                             }),
                             (None, Some(buckets), None) => {
                                 Vdaf::Histogram(Histogram::Categorical {
                                     buckets,
                                     chunk_length,
+                                    dp_strategy,
                                 })
                             }
                             (None, None, Some(buckets)) => Vdaf::Histogram(Histogram::Continuous {
                                 buckets,
                                 chunk_length,
+                                dp_strategy,
                             }),
                             (None, None, None) => {
                                 return Err(Error::Other("continuous-buckets, categorical-buckets, or length are required for histogram vdaf".into()));
@@ -135,15 +188,66 @@ impl TaskAction {
                             }
                         }
                     }
-                    VdafName::Sum => Vdaf::Sum {
-                        bits: bits.unwrap(),
-                    },
-                    VdafName::CountVec => Vdaf::CountVec {
-                        length: length.unwrap(),
-                        chunk_length,
-                    },
+                    VdafName::Sum => {
+                        if differential_privacy_strategy.is_some()
+                            || differential_privacy_epsilon.is_some()
+                        {
+                            return Err(Error::Other(
+                                "differential privacy noise is not yet supported with Prio3Sum"
+                                    .into(),
+                            ));
+                        }
+                        Vdaf::Sum {
+                            bits: bits.unwrap(),
+                        }
+                    }
+                    VdafName::CountVec => {
+                        if differential_privacy_strategy.is_some()
+                            || differential_privacy_epsilon.is_some()
+                        {
+                            return Err(Error::Other(
+                                "differential privacy noise is not supported with Prio3CountVec"
+                                    .into(),
+                            ));
+                        }
+                        Vdaf::CountVec {
+                            length: length.unwrap(),
+                            chunk_length,
+                        }
+                    }
                     VdafName::SumVec => {
-                        Vdaf::SumVec(SumVec::new(bits.unwrap(), length.unwrap(), chunk_length))
+                        let dp_strategy =
+                            match (differential_privacy_strategy, differential_privacy_epsilon) {
+                                (None, None) => dp_strategy::Prio3SumVec::NoDifferentialPrivacy,
+                                (None, Some(_)) => {
+                                    return Err(Error::Other(
+                                        "missing differential-privacy-strategy".into(),
+                                    ))
+                                }
+                                (Some(_), None) => {
+                                    return Err(Error::Other(
+                                        "missing differential-privacy-epsilon".into(),
+                                    ))
+                                }
+                                (Some(DpStrategy::PureDpDiscreteLaplace), Some(epsilon)) => {
+                                    dp_strategy::Prio3SumVec::PureDpDiscreteLaplace(
+                                        PureDpDiscreteLaplace {
+                                            budget: PureDpBudget {
+                                                epsilon: float_to_biguint_ratio(epsilon)
+                                                    .ok_or_else(|| {
+                                                        Error::Other("invalid epsilon".into())
+                                                    })?,
+                                            },
+                                        },
+                                    )
+                                }
+                            };
+                        Vdaf::SumVec(SumVec::new(
+                            bits.unwrap(),
+                            length.unwrap(),
+                            chunk_length,
+                            dp_strategy,
+                        ))
                     }
                 };
 
@@ -201,3 +305,12 @@ impl TaskAction {
         Ok(())
     }
 }
+
+fn float_to_biguint_ratio(value: f64) -> Option<Ratio<BigUint>> {
+    let signed_ratio = Ratio::from_float(value)?;
+    let unsigned_ratio = Ratio::new(
+        signed_ratio.numer().clone().try_into().ok()?,
+        signed_ratio.denom().clone().try_into().ok()?,
+    );
+    Some(unsigned_ratio)
+}
diff --git a/client/Cargo.toml b/client/Cargo.toml
@@ -27,6 +27,8 @@ time = { version = "0.3.36", features = ["serde", "serde-well-known"] }
 log = "0.4.22"
 pad-adapter = "0.1.1"
 janus_messages = "0.7.25"
+num-bigint = "0.4.6"
+num-rational = "0.4.2"
 
 [dev-dependencies]
 divviup-api.workspace = true

diff --git a/client/src/dp_strategy.rs b/client/src/dp_strategy.rs
@@ -0,0 +1,31 @@
+use num_bigint::BigUint;
+use num_rational::Ratio;
+use serde::{Deserialize, Serialize};
+
+#[derive(Serialize, Deserialize, Debug, Clone, Eq, PartialEq, Default)]
+#[serde(tag = "dp_strategy")]
+#[non_exhaustive]
+pub enum Prio3Histogram {
+    #[default]
+    NoDifferentialPrivacy,
+    PureDpDiscreteLaplace(PureDpDiscreteLaplace),
+}
+
+#[derive(Serialize, Deserialize, Debug, Clone, Eq, PartialEq, Default)]
+#[serde(tag = "dp_strategy")]
+#[non_exhaustive]
+pub enum Prio3SumVec {
+    #[default]
+    NoDifferentialPrivacy,
+    PureDpDiscreteLaplace(PureDpDiscreteLaplace),
+}
+
+#[derive(Serialize, Deserialize, Debug, Clone, Eq, PartialEq)]
+pub struct PureDpDiscreteLaplace {
+    pub budget: PureDpBudget,
+}
+
+#[derive(Serialize, Deserialize, Debug, Clone, Eq, PartialEq)]
+pub struct PureDpBudget {
+    pub epsilon: Ratio<BigUint>,
+}
diff --git a/client/src/lib.rs b/client/src/lib.rs
@@ -13,6 +13,7 @@ mod account;
 mod aggregator;
 mod api_token;
 mod collector_credentials;
+pub mod dp_strategy;
 mod membership;
 mod protocol;
 mod task;
@@ -38,6 +39,8 @@ pub use janus_messages::{
     HpkeConfig, HpkePublicKey,
 };
 pub use membership::Membership;
+pub use num_bigint::BigUint;
+pub use num_rational::Ratio;
 pub use protocol::Protocol;
 pub use task::{Histogram, NewTask, SumVec, Task, Vdaf};
 pub use time::OffsetDateTime;

diff --git a/client/src/task.rs b/client/src/task.rs
@@ -3,6 +3,8 @@ use serde::{Deserialize, Serialize};
 use time::OffsetDateTime;
 use uuid::Uuid;
 
+use crate::dp_strategy;
+
 #[derive(Serialize, Deserialize, Clone, Debug, Eq, PartialEq)]
 pub struct Task {
     pub id: String,
@@ -79,14 +81,20 @@ pub enum Histogram {
     Categorical {
         buckets: Vec<String>,
         chunk_length: Option<u64>,
+        #[serde(default)]
+        dp_strategy: dp_strategy::Prio3Histogram,
     },
     Continuous {
         buckets: Vec<u64>,
         chunk_length: Option<u64>,
+        #[serde(default)]
+        dp_strategy: dp_strategy::Prio3Histogram,
     },
     Length {
         length: u64,
         chunk_length: Option<u64>,
+        #[serde(default)]
+        dp_strategy: dp_strategy::Prio3Histogram,
     },
 }
 
@@ -110,22 +118,39 @@ impl Histogram {
         .map(|c| c as usize)
         .unwrap_or_else(|| optimal_chunk_length(self.length()))
     }
+
+    /// The differential privacy strategy.
+    pub fn dp_strategy(&self) -> &dp_strategy::Prio3Histogram {
+        match self {
+            Histogram::Categorical { dp_strategy, .. }
+            | Histogram::Continuous { dp_strategy, .. }
+            | Histogram::Length { dp_strategy, .. } => dp_strategy,
+        }
+    }
 }
 
-#[derive(Serialize, Deserialize, Debug, Clone, Copy, Eq, PartialEq)]
+#[derive(Serialize, Deserialize, Debug, Clone, Eq, PartialEq)]
 pub struct SumVec {
     pub bits: u8,
     pub length: u64,
     chunk_length: Option<u64>,
+    #[serde(default)]
+    pub dp_strategy: dp_strategy::Prio3SumVec,
 }
 
 impl SumVec {
     /// Create a new SumVec
-    pub fn new(bits: u8, length: u64, chunk_length: Option<u64>) -> Self {
+    pub fn new(
+        bits: u8,
+        length: u64,
+        chunk_length: Option<u64>,
+        dp_strategy: dp_strategy::Prio3SumVec,
+    ) -> Self {
         Self {
             bits,
             length,
             chunk_length,
+            dp_strategy,
         }
     }