fix: use default epoch_ID if not set (#483)

* fix: use default epoch_ID if not set

* chore: tests around missing / malformed epoch_id for key and crs_gen

Author: dd23

core/service/src/engine/centralized/service/crs_gen.rs

@@ -440,7 +440,38 @@ mod tests {
             assert_eq!(err.code(), tonic::Code::InvalidArgument);
         }
 
-        // missing epoch ID
+        // TODO: re-enable the following tests after the default epoch fallback is removed in validation
+        // https://github.com/zama-ai/kms-internal/issues/2758
+        //
+        // // missing epoch ID
+        // {
+        //     let request = CrsGenRequest {
+        //         request_id: Some(req_id.into()),
+        //         epoch_id: None, // missing
+        //         context_id: None,
+        //         params: FheParameter::Test.into(),
+        //         domain: Some(domain),
+        //         extra_data: vec![],
+        //         max_num_bits: None,
+        //     };
+        //     let err = crs_gen_impl(&kms, Request::new(request), false)
+        //         .await
+        //         .unwrap_err();
+        //     assert_eq!(err.code(), tonic::Code::InvalidArgument);
+        // }
+    }
+
+    // test the missing epoch ID case that is currently allowed with a warning, to make sure the default epoch fallback works, and to prepare for the future removal of the default epoch fallback
+    // TODO: remove following tests after the default epoch fallback is no longer used in validation
+    // https://github.com/zama-ai/kms-internal/issues/2758
+    #[tokio::test]
+    async fn default_epoch_id() {
+        let mut rng = AesRng::seed_from_u64(54321);
+        let (kms, _) = setup_central_test_kms(&mut rng).await;
+        let req_id = derive_request_id("test_crs_gen_default_epoch").unwrap();
+        let domain = alloy_to_protobuf_domain(&dummy_domain()).unwrap();
+
+        // missing epoch ID should use default epoch with a warning, but not fail
         {
             let request = CrsGenRequest {
                 request_id: Some(req_id.into()),
@@ -451,10 +482,21 @@ mod tests {
                 extra_data: vec![],
                 max_num_bits: None,
             };
-            let err = crs_gen_impl(&kms, Request::new(request), false)
+            let _req = crs_gen_impl(&kms, Request::new(request), false)
                 .await
-                .unwrap_err();
-            assert_eq!(err.code(), tonic::Code::InvalidArgument);
+                .unwrap();
+            let res = get_crs_gen_result_impl(&kms, Request::new(req_id.into()), false)
+                .await
+                .unwrap();
+
+            assert_eq!(
+                res.into_inner()
+                    .request_id
+                    .unwrap()
+                    .request_id
+                    .to_ascii_lowercase(),
+                req_id.as_str()
+            );
         }
     }
 

core/service/src/engine/centralized/service/key_gen.rs

@@ -690,6 +690,32 @@ pub(crate) mod tests {
             .unwrap_err();
             assert_eq!(err.code(), tonic::Code::InvalidArgument);
         }
+
+        // invalid epoch ID should fail
+        {
+            let request = KeyGenRequest {
+                params: Some(FheParameter::Test.into()),
+                keyset_config: None,
+                keyset_added_info: None,
+                request_id: Some(request_id.into()),
+                context_id: None,
+                preproc_id: Some(preproc_id.into()),
+                domain: Some(domain.clone()),
+                epoch_id: Some(kms_grpc::kms::v1::RequestId {
+                    request_id: "invalid-epoch-id".to_string(),
+                }),
+                extra_data: vec![],
+            };
+            let err = key_gen_impl(
+                &kms,
+                tonic::Request::new(request),
+                #[cfg(feature = "insecure")]
+                true,
+            )
+            .await
+            .unwrap_err();
+            assert_eq!(err.code(), tonic::Code::InvalidArgument);
+        }
     }
 
     #[tokio::test]

core/service/src/engine/validation_non_wasm.rs

@@ -761,9 +761,6 @@ fn unpack_crs_gen_request(req: CrsGenRequest) -> anyhow::Result<VerifiedCrsGenRe
         "Received new crs generation request"
     );
 
-    let epoch_id =
-        parse_optional_grpc_request_id(&req.epoch_id, RequestIdParsingErr::CrsGenRequest)?;
-
     // This verification is more strict than the checks in [compute_witness_dim] below
     // because it only allows powers of 2. But there are no strong reasons
     // to use max_num_bits that are not powers of 2 so we enforce it here.
@@ -779,12 +776,17 @@ fn unpack_crs_gen_request(req: CrsGenRequest) -> anyhow::Result<VerifiedCrsGenRe
     let witness_dim = compute_witness_dim(&crs_params, req.max_num_bits.map(|x| x as usize))?;
 
     // TODO(zama-ai/kms-internal/issues/2758)
-    // remove the default context when all of context is ready
+    // remove the default context and epoch when all of context is ready
     let context_id = match &req.context_id {
         Some(ctx) => parse_grpc_request_id(ctx, RequestIdParsingErr::Context)?,
         None => *DEFAULT_MPC_CONTEXT,
     };
 
+    let epoch_id = match &req.epoch_id {
+        Some(epoch) => parse_grpc_request_id(epoch, RequestIdParsingErr::Epoch)?,
+        None => *DEFAULT_EPOCH_ID,
+    };
+
     let eip712_domain = optional_protobuf_to_alloy_domain(req.domain.as_ref())?;
 
     Ok(VerifiedCrsGenRequest {