chore(kms-connector): limit decryption attempts (#1329)

* chore(kms-connector): limit decryption attempts

* chore(kms-connector): code cleaning

Author: eudelins-zama

…c1940b5f4f6306113ff9648f690f1379240.json …430f7bcdfba000a3166c7d68fd54908f6c9.jsonkms-connector/.sqlx/query-12d768492b153b53cf4cb6c9f3e66c1940b5f4f6306113ff9648f690f1379240.json renamed to kms-connector/.sqlx/query-60f58397e8588260d3c79b89bfdef430f7bcdfba000a3166c7d68fd54908f6c9.json kms-connector/.sqlx/query-12d768492b153b53cf4cb6c9f3e66c1940b5f4f6306113ff9648f690f1379240.json renamed to kms-connector/.sqlx/query-60f58397e8588260d3c79b89bfdef430f7bcdfba000a3166c7d68fd54908f6c9.json

@@ -87,6 +87,10 @@ address = "0x5ffdaAB0373E62E2ea2944776209aEf29E631A64"
 # ENV: KMS_CONNECTOR_GRPC_REQUEST_RETRIES
 # grpc_request_retries = 3
 
+# The maximum number of decryption attempts before a decryption request is permanently removed (optional, defaults to 200)
+# ENV: KMS_CONNECTOR_MAX_DECRYPTION_ATTEMPTS
+# max_decryption_attempts = 200
+
 # Timeout to get public decryption responses from KMS Core in seconds (optional, defaults to 300s / 5min)
 # ENV: KMS_CONNECTOR_PUBLIC_DECRYPTION_TIMEOUT_SECS
 # public_decryption_timeout_secs = 300

…ff833085240442de258975cde391de302a2.json …20421f84523b5a88dcd3f44305effaa5c6e.jsonkms-connector/.sqlx/query-f9e4703b71b26e3c08ecb4960d8c9ff833085240442de258975cde391de302a2.json renamed to kms-connector/.sqlx/query-e0d8f33e5502609d10d0a141bb2bd20421f84523b5a88dcd3f44305effaa5c6e.json kms-connector/.sqlx/query-f9e4703b71b26e3c08ecb4960d8c9ff833085240442de258975cde391de302a2.json renamed to kms-connector/.sqlx/query-e0d8f33e5502609d10d0a141bb2bd20421f84523b5a88dcd3f44305effaa5c6e.json

@@ -0,0 +1,2 @@
+ALTER TABLE public_decryption_requests ADD COLUMN error_counter SMALLINT DEFAULT 0 NOT NULL;
+ALTER TABLE user_decryption_requests ADD COLUMN error_counter SMALLINT DEFAULT 0 NOT NULL;

kms-connector/config/kms-worker.toml

@@ -90,9 +90,9 @@ async fn publish_user_decryption(
 
     sqlx::query!(
         "INSERT INTO user_decryption_requests(\
-                decryption_id, sns_ct_materials, user_address, public_key, extra_data, otlp_context\
-            ) \
-            VALUES ($1, $2, $3, $4, $5, $6) ON CONFLICT DO NOTHING",
+            decryption_id, sns_ct_materials, user_address, public_key, extra_data, otlp_context\
+        ) \
+        VALUES ($1, $2, $3, $4, $5, $6) ON CONFLICT DO NOTHING",
         request.decryptionId.as_le_slice(),
         sns_ciphertexts_db as Vec<SnsCiphertextMaterialDbItem>,
         request.userAddress.as_slice(),

kms-connector/connector-db/migrations/20251113084850_add_decryption_error_counter.sql

@@ -40,6 +40,8 @@ pub struct Config {
     pub events_batch_size: u8,
     /// Number of retries for GRPC requests sent to the KMS Core.
     pub grpc_request_retries: u8,
+    /// The maximum number of decryption attempts.
+    pub max_decryption_attempts: u16,
 
     /// Number of retries for S3 ciphertext retrieval.
     pub s3_ciphertext_retrieval_retries: u8,
@@ -121,6 +123,7 @@ impl Config {
             service_name: raw_config.service_name,
             events_batch_size: raw_config.events_batch_size,
             grpc_request_retries: raw_config.grpc_request_retries,
+            max_decryption_attempts: raw_config.max_decryption_attempts,
             s3_ciphertext_retrieval_retries: raw_config.s3_ciphertext_retrieval_retries,
             s3_connect_timeout: s3_ciphertext_retrieval_timeout,
             task_limit: raw_config.task_limit,
@@ -158,6 +161,7 @@ mod tests {
             env::remove_var("KMS_CONNECTOR_SERVICE_NAME");
             env::remove_var("KMS_CONNECTOR_EVENTS_BATCH_SIZE");
             env::remove_var("KMS_CONNECTOR_GRPC_REQUEST_RETRIES");
+            env::remove_var("KMS_CONNECTOR_MAX_DECRYPTION_ATTEMPTS");
             env::remove_var("KMS_CONNECTOR_S3_CIPHERTEXT_RETRIEVAL_RETRIES");
             env::remove_var("KMS_CONNECTOR_S3_CONNECT_TIMEOUT");
         }
@@ -237,6 +241,7 @@ mod tests {
             env::set_var("KMS_CONNECTOR_SERVICE_NAME", "kms-connector-test");
             env::set_var("KMS_CONNECTOR_EVENTS_BATCH_SIZE", "15");
             env::set_var("KMS_CONNECTOR_GRPC_REQUEST_RETRIES", "5");
+            env::set_var("KMS_CONNECTOR_MAX_DECRYPTION_ATTEMPTS", "300");
             env::set_var("KMS_CONNECTOR_S3_CIPHERTEXT_RETRIEVAL_RETRIES", "5");
             env::set_var("KMS_CONNECTOR_S3_CONNECT_TIMEOUT", "4");
         }
@@ -266,6 +271,7 @@ mod tests {
         assert_eq!(config.service_name, "kms-connector-test");
         assert_eq!(config.events_batch_size, 15);
         assert_eq!(config.grpc_request_retries, 5);
+        assert_eq!(config.max_decryption_attempts, 300);
         assert_eq!(config.s3_ciphertext_retrieval_retries, 5);
         assert_eq!(config.s3_connect_timeout.as_secs(), 4);
 

kms-connector/crates/gw-listener/src/core/publish.rs

@@ -36,6 +36,8 @@ pub struct RawConfig {
     pub events_batch_size: u8,
     #[serde(default = "default_grpc_request_retries")]
     pub grpc_request_retries: u8,
+    #[serde(default = "default_max_decryption_attempts")]
+    pub max_decryption_attempts: u16,
     #[serde(default = "default_s3_ciphertext_retrieval_retries")]
     pub s3_ciphertext_retrieval_retries: u8,
     #[serde(default = "default_s3_connect_timeout")]
@@ -68,6 +70,10 @@ fn default_grpc_request_retries() -> u8 {
     3
 }
 
+fn default_max_decryption_attempts() -> u16 {
+    200
+}
+
 fn default_s3_ciphertext_retrieval_retries() -> u8 {
     3
 }
@@ -137,6 +143,7 @@ impl Default for RawConfig {
             service_name: "kms-connector".to_string(),
             events_batch_size: 10,
             grpc_request_retries: 3,
+            max_decryption_attempts: default_max_decryption_attempts(),
             s3_ciphertext_retrieval_retries: 3,
             s3_connect_timeout: 2,
             task_limit: default_task_limit(),

kms-connector/crates/kms-worker/src/core/config/parsed.rs

@@ -121,7 +121,7 @@ impl DbEventPicker {
                     LIMIT $1 FOR UPDATE SKIP LOCKED
                 ) AS req
                 WHERE public_decryption_requests.decryption_id = req.decryption_id
-                RETURNING req.decryption_id, sns_ct_materials, extra_data, otlp_context, already_sent
+                RETURNING req.decryption_id, sns_ct_materials, extra_data, otlp_context, already_sent, error_counter
             ",
         )
         .bind(self.events_batch_size as i16)
@@ -144,7 +144,7 @@ impl DbEventPicker {
                     LIMIT $1 FOR UPDATE SKIP LOCKED
                 ) AS req
                 WHERE user_decryption_requests.decryption_id = req.decryption_id
-                RETURNING req.decryption_id, sns_ct_materials, user_address, public_key, extra_data, otlp_context, already_sent
+                RETURNING req.decryption_id, sns_ct_materials, user_address, public_key, extra_data, otlp_context, already_sent, error_counter
             ",
         )
         .bind(self.events_batch_size as i16)