Add missing non_exhaustive to enums (#1957)

* Add non_exhaustive to ExtensionType

* Add non_exhaustive to query type

* Add non_exhaustive to HPKE types

* Errant unwrap in aggregator API, should propagate error

* Missing unwrap

aggregator_api/src/lib.rs

@@ -5,9 +5,11 @@ mod routes;
 mod tests;
 
 use async_trait::async_trait;
-use janus_aggregator_core::datastore;
-use janus_aggregator_core::{datastore::Datastore, instrumented};
-use janus_core::{http::extract_bearer_token, task::AuthenticationToken, time::Clock};
+use janus_aggregator_core::{
+    datastore::{self, Datastore},
+    instrumented,
+};
+use janus_core::{hpke, http::extract_bearer_token, task::AuthenticationToken, time::Clock};
 use janus_messages::{HpkeConfigId, RoleParseError, TaskId};
 use ring::constant_time;
 use routes::*;
@@ -162,6 +164,8 @@ enum Error {
     Url(#[from] url::ParseError),
     #[error(transparent)]
     Role(#[from] RoleParseError),
+    #[error(transparent)]
+    Hpke(#[from] hpke::Error),
 }
 
 #[async_trait]
@@ -200,6 +204,9 @@ impl Handler for Error {
             Self::Role(err) => conn
                 .with_status(Status::BadRequest)
                 .with_body(err.to_string()),
+            Self::Hpke(err) => conn
+                .with_status(Status::BadRequest)
+                .with_body(err.to_string()),
         }
         .halt()
     }

aggregator_api/src/routes.rs

@@ -146,12 +146,14 @@ pub(super) async fn post_task<C: Clock>(
         _ => unreachable!(),
     };
 
+    // Unwrap safety: we always use a supported KEM.
     let hpke_keys = Vec::from([generate_hpke_config_and_private_key(
         random(),
         HpkeKemId::X25519HkdfSha256,
         HpkeKdfId::HkdfSha256,
         HpkeAeadId::Aes128Gcm,
-    )]);
+    )
+    .unwrap()]);
 
     let task = Arc::new(
         Task::new(
@@ -321,7 +323,7 @@ pub(super) async fn put_global_hpke_config<C: Clock>(
         req.kem_id.unwrap_or(HpkeKemId::X25519HkdfSha256),
         req.kdf_id.unwrap_or(HpkeKdfId::HkdfSha256),
         req.aead_id.unwrap_or(HpkeAeadId::Aes128Gcm),
-    );
+    )?;
 
     let inserted_keypair = ds
         .run_tx_with_name("put_global_hpke_config", |tx| {

aggregator_api/src/tests.rs

@@ -206,6 +206,7 @@ async fn post_task_bad_role() {
             HpkeKdfId::HkdfSha256,
             HpkeAeadId::Aes128Gcm,
         )
+        .unwrap()
         .config()
         .clone(),
         aggregator_auth_token: Some(aggregator_auth_token),
@@ -246,6 +247,7 @@ async fn post_task_unauthorized() {
             HpkeKdfId::HkdfSha256,
             HpkeAeadId::Aes128Gcm,
         )
+        .unwrap()
         .config()
         .clone(),
         aggregator_auth_token: Some(aggregator_auth_token),
@@ -287,6 +289,7 @@ async fn post_task_helper_no_optional_fields() {
             HpkeKdfId::HkdfSha256,
             HpkeAeadId::Aes128Gcm,
         )
+        .unwrap()
         .config()
         .clone(),
         aggregator_auth_token: None,
@@ -366,6 +369,7 @@ async fn post_task_helper_with_aggregator_auth_token() {
             HpkeKdfId::HkdfSha256,
             HpkeAeadId::Aes128Gcm,
         )
+        .unwrap()
         .config()
         .clone(),
         aggregator_auth_token: Some(aggregator_auth_token),
@@ -408,6 +412,7 @@ async fn post_task_idempotence() {
             HpkeKdfId::HkdfSha256,
             HpkeAeadId::Aes128Gcm,
         )
+        .unwrap()
         .config()
         .clone(),
         aggregator_auth_token: Some(aggregator_auth_token.clone()),
@@ -488,6 +493,7 @@ async fn post_task_leader_all_optional_fields() {
             HpkeKdfId::HkdfSha256,
             HpkeAeadId::Aes128Gcm,
         )
+        .unwrap()
         .config()
         .clone(),
         aggregator_auth_token: Some(aggregator_auth_token.clone()),
@@ -577,6 +583,7 @@ async fn post_task_leader_no_aggregator_auth_token() {
             HpkeKdfId::HkdfSha256,
             HpkeAeadId::Aes128Gcm,
         )
+        .unwrap()
         .config()
         .clone(),
         aggregator_auth_token: None,
@@ -861,7 +868,8 @@ async fn get_global_hpke_configs() {
         HpkeKemId::P256HkdfSha256,
         HpkeKdfId::HkdfSha384,
         HpkeAeadId::Aes128Gcm,
-    );
+    )
+    .unwrap();
     ds.run_tx(|tx| {
         let keypair1 = keypair1.clone();
         let keypair2 = keypair2.clone();
@@ -962,7 +970,8 @@ async fn get_global_hpke_config() {
         HpkeKemId::P256HkdfSha256,
         HpkeKdfId::HkdfSha384,
         HpkeAeadId::Aes128Gcm,
-    );
+    )
+    .unwrap();
     ds.run_tx(|tx| {
         let keypair1 = keypair1.clone();
         let keypair2 = keypair2.clone();

aggregator_core/src/task.rs

@@ -550,12 +550,14 @@ impl SerializedTask {
         }
 
         if self.hpke_keys.is_empty() {
+            // Unwrap safety: we always use a supported KEM.
             let hpke_keypair = generate_hpke_config_and_private_key(
                 random(),
                 HpkeKemId::X25519HkdfSha256,
                 HpkeKdfId::HkdfSha256,
                 HpkeAeadId::Aes128Gcm,
-            );
+            )
+            .unwrap();
 
             self.hpke_keys = Vec::from([hpke_keypair]);
         }

collector/src/lib.rs

@@ -25,7 +25,7 @@
 //!     HpkeKemId::X25519HkdfSha256,
 //!     HpkeKdfId::HkdfSha256,
 //!     HpkeAeadId::Aes128Gcm,
-//! );
+//! ).unwrap();
 //! let parameters = CollectorParameters::new(
 //!     task_id,
 //!     "https://example.com/dap/".parse().unwrap(),

core/src/hpke.rs

@@ -22,6 +22,8 @@ pub enum Error {
     Hpke(#[from] HpkeError),
     #[error("invalid HPKE configuration: {0}")]
     InvalidConfiguration(&'static str),
+    #[error("unsupported KEM")]
+    UnsupportedKem,
 }
 
 fn hpke_dispatch_config_from_hpke_config(
@@ -200,21 +202,23 @@ pub fn open(
 }
 
 /// Generate a new HPKE keypair and return it as an HpkeConfig (public portion) and
-/// HpkePrivateKey (private portion).
+/// HpkePrivateKey (private portion). This function errors if the supplied key
+/// encapsulated mechanism is not supported by the underlying HPKE library.
 pub fn generate_hpke_config_and_private_key(
     hpke_config_id: HpkeConfigId,
     kem_id: HpkeKemId,
     kdf_id: HpkeKdfId,
     aead_id: HpkeAeadId,
-) -> HpkeKeypair {
+) -> Result<HpkeKeypair, Error> {
     let Keypair {
         private_key,
         public_key,
     } = match kem_id {
         HpkeKemId::X25519HkdfSha256 => Kem::X25519HkdfSha256.gen_keypair(),
         HpkeKemId::P256HkdfSha256 => Kem::DhP256HkdfSha256.gen_keypair(),
+        _ => return Err(Error::UnsupportedKem),
     };
-    HpkeKeypair::new(
+    Ok(HpkeKeypair::new(
         HpkeConfig::new(
             hpke_config_id,
             kem_id,
@@ -223,7 +227,7 @@ pub fn generate_hpke_config_and_private_key(
             HpkePublicKey::from(public_key),
         ),
         HpkePrivateKey::new(private_key),
-    )
+    ))
 }
 
 /// An HPKE configuration and its corresponding private key.
@@ -313,6 +317,7 @@ pub mod test_util {
             HpkeKdfId::HkdfSha256,
             HpkeAeadId::Aes128Gcm,
         )
+        .unwrap()
     }
 
     pub fn generate_test_hpke_config_and_private_key_with_id(id: u8) -> HpkeKeypair {
@@ -322,6 +327,7 @@ pub mod test_util {
             HpkeKdfId::HkdfSha256,
             HpkeAeadId::Aes128Gcm,
         )
+        .unwrap()
     }
 }
 

interop_binaries/src/lib.rs

@@ -352,6 +352,7 @@ impl HpkeConfigRegistry {
         self.keypairs
             .entry(id)
             .or_insert_with(|| {
+                // Unwrap safety: we always use a supported KEM.
                 generate_hpke_config_and_private_key(
                     id,
                     // These algorithms should be broadly compatible with other DAP implementations, since they
@@ -360,6 +361,7 @@ impl HpkeConfigRegistry {
                     HpkeKdfId::HkdfSha256,
                     HpkeAeadId::Aes128Gcm,
                 )
+                .unwrap()
             })
             .clone()
     }

messages/src/lib.rs

@@ -744,6 +744,7 @@ impl<'de> Deserialize<'de> for TaskId {
 /// DAP protocol message representing an HPKE key encapsulation mechanism.
 #[derive(Clone, Copy, Debug, PartialEq, Eq, TryFromPrimitive, Serialize, Deserialize)]
 #[repr(u16)]
+#[non_exhaustive]
 pub enum HpkeKemId {
     /// NIST P-256 keys and HKDF-SHA256.
     P256HkdfSha256 = 0x0010,
@@ -773,6 +774,7 @@ impl Decode for HpkeKemId {
 /// DAP protocol message representing an HPKE key derivation function.
 #[derive(Clone, Copy, Debug, PartialEq, Eq, TryFromPrimitive, Serialize, Deserialize)]
 #[repr(u16)]
+#[non_exhaustive]
 pub enum HpkeKdfId {
     /// HMAC Key Derivation Function SHA256.
     HkdfSha256 = 0x0001,
@@ -804,6 +806,7 @@ impl Decode for HpkeKdfId {
 /// DAP protocol message representing an HPKE AEAD.
 #[derive(Clone, Copy, Debug, PartialEq, Eq, TryFromPrimitive, Serialize, Deserialize)]
 #[repr(u16)]
+#[non_exhaustive]
 pub enum HpkeAeadId {
     /// AES-128-GCM.
     Aes128Gcm = 0x0001,
@@ -886,6 +889,7 @@ impl Decode for Extension {
 /// DAP protocol message representing the type of an extension included in a client report.
 #[derive(Clone, Copy, Debug, Hash, Eq, PartialEq, TryFromPrimitive)]
 #[repr(u16)]
+#[non_exhaustive]
 pub enum ExtensionType {
     Tbd = 0,
 }
@@ -2021,6 +2025,7 @@ pub mod query_type {
     /// DAP protocol message representing the type of a query.
     #[derive(Copy, Clone, Debug, PartialEq, Eq, TryFromPrimitive, Serialize, Deserialize)]
     #[repr(u8)]
+    #[non_exhaustive]
     pub enum Code {
         Reserved = 0,
         TimeInterval = 1,

tools/src/bin/hpke_keygen.rs

@@ -19,7 +19,7 @@ fn main() -> Result<()> {
         options.kem.into(),
         options.kdf.into(),
         options.aead.into(),
-    );
+    )?;
 
     let mut writer = stdout().lock();