feat: add extra_data to core_client decryption (#470)

* chore: change public decryption extra_data to &[u8] instead of Vec<u8>

* chore: pass extra_data from request in core-client

* chore: properly propagate extra_data

* chore: fix integration test build

* chore: fix test

* chore: more tests on extra_data

* chore: bump aws-lc-rs

* chore: update cargo audit

* ci: upgrade zizmor action

* chore: bump zismor version

* chore: unify parsing of extra_data

* chore: fix typos in docstrings

* chore: remove explicit zizmor version in workflow

* doc: update cargo audit comment

---------

Co-authored-by: Kelong Cong <kc1212@users.noreply.github.com>
Co-authored-by: Frederic Egmorte <frederic.egmorte@zama.ai>

Author: 3 people

.cargo/audit.toml

@@ -2,8 +2,17 @@
 # permanently specified in this file.
 
 [advisories]
-# The ignored vulnerability RUSTSEC-2023-0071 is not applicable in our use-case
-ignore = ["RUSTSEC-2023-0071"]
+# RUSTSEC-2023-0071 is not applicable in our use-case.
+# RUSTSEC-2026-0049 is not an urgent fix for us. It can only be triggered
+# when talking to AWS APIs by allowing the attacker to trick us into accepting
+# expired certificates. If AWS's TLS certificates are stolen, the consequences
+# are of a different magnitude and not something we can do much about.
+# For the time being we also require the "connector-hyper-0-14-x" feature
+# because we need to specify fixed server names for TLS connections (because
+# when socat proxies are used, the connection goes to localhost and if we don't
+# overload the server names, the handshake would fail because localhost isn't
+# `sts.*.amazonaws.com`).
+ignore = ["RUSTSEC-2023-0071", "RUSTSEC-2026-0049"]
 informational_warnings = ["unmaintained"]
 severity_threshold = "medium"
 

.cargo/deny.toml

@@ -100,7 +100,7 @@ allow = [
     "ISC",
     "MIT",
     "MPL-2.0",
-    "OpenSSL",
+    # "OpenSSL", no longer used
     # "Unicode-DFS-2016", no longer used
     "Unicode-3.0",
     "Unlicense",

.github/workflows/ci_lint.yml

@@ -51,7 +51,6 @@ jobs:
           persist-credentials: false
 
       - name: Run zizmor
-        uses: zizmorcore/zizmor-action@e673c3917a1aef3c65c972347ed84ccd013ecda4 # v0.2.0
+        uses: zizmorcore/zizmor-action@71321a20a9ded102f6e9ce5718a2fcec2c4f70d8 # v0.5.2
         with:
           persona: pedantic
-          version: 1.15.2

Cargo.lock

@@ -91,11 +91,11 @@ async-trait = "=0.1.89"  # Async trait support -  MEDIUM RISK: Reputable individ
 async_cell = "0.2.2"  # Async cell implementation - HIGH RISK: Individual maintainer, very low popularity
 attestation-doc-validation = { version = "=0.10.0" }  # AWS Nitro attestation validation - LOW RISK: Evervault (reputable security company), security-critical but trusted
 aws-config = { version = "=1.8.12" }  # AWS SDK configuration - LOW RISK: Official AWS SDK, actively maintained
-aws-lc-rs = { version = "=1.16.1" } # AWS-LC dependency pinned because of an earlier version vulnerability - LOW RISK: Official AWS SDK
+aws-lc-rs = { version = "=1.16.2" } # AWS-LC dependency pinned because of an earlier version vulnerability - LOW RISK: Official AWS SDK
 aws-nitro-enclaves-nsm-api = { version = "=0.4.0" }  # AWS Nitro Enclaves NSM API - LOW RISK: Official AWS SDK
 aws-sdk-kms = { version = "=1.98.0" }  # AWS KMS client - LOW RISK: Official AWS SDK for key management
 aws-sdk-s3 = { version = "=1.120.0" }  # AWS S3 client - LOW RISK: Official AWS SDK for object storage
-aws-smithy-runtime = { version = "=1.10.3", features = ["client", "connector-hyper-0-14-x"] }  # AWS Smithy runtime - LOW RISK: Official AWS runtime library
+aws-smithy-runtime = { version = "=1.10.3", features = ["client", "connector-hyper-0-14-x"] }  # AWS Smithy runtime - LOW RISK: Official AWS runtime library. Remove "connector-hyper-0-14-x" when possible as it introduces a potential vulnerability.
 aws-smithy-runtime-api = { version = "=1.11.6" }  # AWS Smithy runtime API - LOW RISK: Official AWS runtime API
 aws-smithy-types = { version = "=1.4.6" }  # AWS Smithy types - LOW RISK: Official AWS type definitions
 axum = { version = "=0.8.8", features = ["tokio"] }  # Web framework - LOW RISK: tokio-rs team, 168M+ downloads, actively maintained

Cargo.toml

@@ -33,6 +33,7 @@ pub(crate) async fn do_crsgen(
     destination_prefix: &Path,
     context_id: Option<ContextId>,
     epoch_id: Option<EpochId>,
+    extra_data: Vec<u8>,
 ) -> anyhow::Result<RequestId> {
     let req_id = RequestId::new_random(rng);
 
@@ -115,7 +116,7 @@ pub(crate) async fn do_crsgen(
         destination_prefix,
         req_id,
         domain,
-        vec![],
+        extra_data,
         resp_response_vec,
         cmd_conf.download_all,
     )

core-client/src/crsgen.rs

@@ -26,13 +26,13 @@ fn check_ext_pt_signature(
     external_handles: Vec<Vec<u8>>,
     domain: Eip712Domain,
     kms_addrs: &[alloy_primitives::Address],
-    extra_data: Vec<u8>,
+    extra_data: &[u8],
 ) -> anyhow::Result<()> {
     tracing::debug!(
         "Checking signature for PTs: {:?}, ext. handles: {:?}, extra_data: {}, ext. sig {}",
         plaintexts,
         external_handles,
-        hex::encode(&extra_data),
+        hex::encode(extra_data),
         hex::encode(external_sig)
     );
     let message = compute_public_decryption_message(external_handles, plaintexts, extra_data)?;
@@ -54,6 +54,7 @@ fn check_external_decryption_signature(
     external_handles: &[Vec<u8>],
     domain: &Eip712Domain,
     kms_addrs: &[alloy_primitives::Address],
+    extra_data: &[u8],
 ) -> anyhow::Result<()> {
     let mut results = Vec::new();
     for response in responses {
@@ -67,7 +68,7 @@ fn check_external_decryption_signature(
             external_handles.to_owned(),
             domain.clone(),
             kms_addrs,
-            vec![],
+            extra_data,
         )?;
 
         for (idx, pt) in payload.plaintexts.iter().enumerate() {
@@ -114,6 +115,7 @@ pub(crate) async fn do_public_decrypt<R: Rng + CryptoRng>(
     num_expected_responses: usize,
     inter_request_delay: tokio::time::Duration,
     parallel_requests: usize,
+    extra_data: Vec<u8>,
 ) -> anyhow::Result<Vec<(Option<RequestId>, String)>> {
     let mut timings_start = HashMap::new();
     let mut durations = Vec::new();
@@ -136,6 +138,7 @@ pub(crate) async fn do_public_decrypt<R: Rng + CryptoRng>(
         let core_endpoints_resp = core_endpoints_resp.clone();
         let ptxt = ptxt.clone();
         let kms_addrs = kms_addrs.clone();
+        let extra_data = extra_data.clone();
 
         // start timing measurement for this request
         timings_start.insert(req_id, tokio::time::Instant::now()); // start timing for this request
@@ -149,6 +152,7 @@ pub(crate) async fn do_public_decrypt<R: Rng + CryptoRng>(
                 context_id.as_ref(),
                 &key_id.into(),
                 epoch_id.as_ref(),
+                &extra_data,
             )?;
 
             // make parallel requests by calling [decrypt] in a thread
@@ -246,6 +250,7 @@ pub(crate) async fn do_user_decrypt<R: Rng + CryptoRng>(
     num_expected_responses: usize,
     inter_request_delay: tokio::time::Duration,
     parallel_requests: usize,
+    extra_data: Vec<u8>,
 ) -> anyhow::Result<Vec<(Option<RequestId>, String)>> {
     let mut join_set: JoinSet<Result<_, anyhow::Error>> = JoinSet::new();
     let mut timings_start = HashMap::new();
@@ -267,6 +272,7 @@ pub(crate) async fn do_user_decrypt<R: Rng + CryptoRng>(
         let core_endpoints_req = core_endpoints_req.clone();
         let core_endpoints_resp = core_endpoints_resp.clone();
         let original_plaintext = ptxt.clone();
+        let extra_data = extra_data.clone();
 
         // start timing measurement for this request
         timings_start.insert(req_id, tokio::time::Instant::now()); // start timing for this request
@@ -281,6 +287,7 @@ pub(crate) async fn do_user_decrypt<R: Rng + CryptoRng>(
                 context_id.as_ref(),
                 epoch_id.as_ref(),
                 PkeSchemeType::MlKem512,
+                &extra_data,
             )?;
 
             let (user_decrypt_req, enc_pk, enc_sk) = user_decrypt_req_tuple;
@@ -587,7 +594,8 @@ pub(crate) async fn get_public_decrypt_responses(
             .clone(),
     };
 
-    let (domain, external_handles) = if let Some(decryption_request) = dec_req.as_ref() {
+    let (domain, external_handles, extra_data) = if let Some(decryption_request) = dec_req.as_ref()
+    {
         let domain_msg = decryption_request
             .domain
             .as_ref()
@@ -599,7 +607,8 @@ pub(crate) async fn get_public_decrypt_responses(
             .iter()
             .map(|ct| ct.external_handle.clone())
             .collect();
-        (domain, external_handles)
+        let extra_data = decryption_request.extra_data.clone();
+        (domain, external_handles, extra_data)
     } else {
         //If the decryption request isn't provided we assume it was dummy domains and handles
         let num_handles = resp_response_vec
@@ -610,7 +619,16 @@ pub(crate) async fn get_public_decrypt_responses(
             .ok_or_else(|| anyhow::anyhow!("missing payload in first decryption response"))?
             .plaintexts
             .len();
-        (dummy_domain(), vec![dummy_handle(); num_handles])
+        let extra_data = resp_response_vec
+            .first()
+            .ok_or_else(|| anyhow::anyhow!("no public decryption responses available"))?
+            .extra_data
+            .clone();
+        (
+            dummy_domain(),
+            vec![dummy_handle(); num_handles],
+            extra_data,
+        )
     };
 
     // check the internal signatures
@@ -627,6 +645,7 @@ pub(crate) async fn get_public_decrypt_responses(
         &external_handles,
         &domain,
         kms_addrs,
+        &extra_data,
     )?;
 
     tracing::info!(

core-client/src/decrypt.rs

@@ -65,6 +65,7 @@ pub(crate) async fn do_keygen(
     insecure: bool,
     shared_config: &SharedKeyGenParameters,
     destination_prefix: &Path,
+    extra_data: Vec<u8>,
 ) -> anyhow::Result<RequestId> {
     let req_id = RequestId::new_random(rng);
 
@@ -162,7 +163,7 @@ pub(crate) async fn do_keygen(
         destination_prefix,
         req_id,
         domain,
-        vec![],
+        extra_data,
         resp_response_vec,
         cmd_conf.download_all,
         shared_config.compressed,

core-client/src/keygen.rs

@@ -494,6 +494,20 @@ impl CipherArguments {
             }
         }
     }
+
+    pub fn get_extra_data(&self) -> Vec<u8> {
+        let hex_str = match self {
+            CipherArguments::FromFile(cipher_file) => &cipher_file.extra_data,
+            CipherArguments::FromArgs(cipher_parameters) => &cipher_parameters.extra_data,
+        };
+        parse_extra_data(hex_str)
+    }
+}
+
+// Helper function to parse the extra data from the CLI arguments, with the same logic for both CipherParameters and CipherFile.
+// Defaults to an empty byte vector if the extra data is not provided or if the hex parsing fails.
+fn parse_extra_data(hex_str: &Option<String>) -> Vec<u8> {
+    parse_hex(hex_str.as_deref().unwrap_or("")).unwrap_or_default()
 }
 
 #[derive(Debug, Args, Clone, Serialize, Deserialize)]
@@ -553,6 +567,11 @@ pub struct CipherParameters {
     #[serde(skip_serializing, skip_deserializing)]
     #[clap(long, default_value_t = false)]
     pub compressed_keys: bool,
+    /// Optional extra data (hex-encoded) to include in the request.
+    /// Can optionally have a "0x" prefix.
+    #[serde(skip_serializing, skip_deserializing)]
+    #[clap(long)]
+    pub extra_data: Option<String>,
 }
 
 #[derive(Debug, Args, Clone)]
@@ -573,6 +592,10 @@ pub struct CipherFile {
     /// Number of requests to be sent in parallel (at most num_requests) before waiting for inter_request_delay_ms.
     #[clap(long, short = 'p', default_value_t = 0)]
     pub parallel_requests: usize,
+    /// Optional extra data (hex-encoded) to include in the request.
+    /// Can optionally have a "0x" prefix.
+    #[clap(long)]
+    pub extra_data: Option<String>,
 }
 
 #[derive(Debug, Serialize, Deserialize)]
@@ -601,6 +624,10 @@ pub struct SharedKeyGenParameters {
     pub use_existing_key_tag: bool,
     pub context_id: Option<ContextId>,
     pub epoch_id: Option<EpochId>,
+    /// Optional extra data (hex-encoded) to include in the request.
+    /// Can optionally have a "0x" prefix.
+    #[clap(long)]
+    pub extra_data: Option<String>,
 }
 
 #[derive(Debug, Parser, Clone)]
@@ -626,6 +653,10 @@ pub struct CrsParameters {
     pub epoch_id: Option<EpochId>,
     #[clap(long)]
     pub context_id: Option<ContextId>,
+    /// Optional extra data (hex-encoded) to include in the request.
+    /// Can optionally have a "0x" prefix.
+    #[clap(long)]
+    pub extra_data: Option<String>,
 }
 
 impl Default for CrsParameters {
@@ -634,6 +665,7 @@ impl Default for CrsParameters {
             max_num_bits: 2048,
             epoch_id: None,
             context_id: None,
+            extra_data: None,
         }
     }
 }
@@ -1667,6 +1699,7 @@ pub async fn execute_cmd(
                 num_expected_responses,
                 cipher_args.get_inter_request_delay_ms(),
                 cipher_args.get_parallel_requests(),
+                cipher_args.get_extra_data(),
             )
             .await?
         }
@@ -1742,6 +1775,7 @@ pub async fn execute_cmd(
                 num_expected_responses,
                 cipher_args.get_inter_request_delay_ms(),
                 cipher_args.get_parallel_requests(),
+                cipher_args.get_extra_data(),
             )
             .await?
         }
@@ -1767,6 +1801,7 @@ pub async fn execute_cmd(
                 false,
                 shared_args,
                 destination_prefix,
+                parse_extra_data(&shared_args.extra_data),
             )
             .await?;
 
@@ -1792,6 +1827,7 @@ pub async fn execute_cmd(
                 true,
                 shared_args,
                 destination_prefix,
+                parse_extra_data(&shared_args.extra_data),
             )
             .await?;
 
@@ -1801,6 +1837,7 @@ pub async fn execute_cmd(
             max_num_bits,
             epoch_id,
             context_id,
+            extra_data,
         }) => {
             let mut internal_client = internal_client.unwrap();
             tracing::info!(
@@ -1822,6 +1859,7 @@ pub async fn execute_cmd(
                 destination_prefix,
                 *context_id,
                 *epoch_id,
+                parse_extra_data(extra_data),
             )
             .await?;
             vec![(Some(req_id), "crsgen done".to_string())]
@@ -1830,6 +1868,7 @@ pub async fn execute_cmd(
             max_num_bits,
             epoch_id,
             context_id,
+            extra_data,
         }) => {
             let mut internal_client = internal_client.unwrap();
             tracing::info!(
@@ -1851,6 +1890,7 @@ pub async fn execute_cmd(
                 destination_prefix,
                 *context_id,
                 *epoch_id,
+                parse_extra_data(extra_data),
             )
             .await?;
             vec![(Some(req_id), "insecure crsgen done".to_string())]