Accept identity deletion requests (#320)

* rename SMPCRequest to UniquenessRequest * parse identity deletion requests and add them to batch * fmt * fmt * move request types into a switch case statement * delete sqs message on default case * log and continue upon invalid request instead of crashing
worldcoin · Sep 6, 2024 · b5b553d · b5b553d
1 parent 6d42f88
commit b5b553d
Show file tree

Hide file tree

Showing 7 changed files with 260 additions and 111 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/iris-mpc-common/Cargo.toml b/iris-mpc-common/Cargo.toml
@@ -8,6 +8,7 @@ edition = "2021"
 aws-config.workspace = true
 aws-sdk-kms.workspace = true
 aws-sdk-sns.workspace = true
+aws-sdk-sqs.workspace = true
 aws-sdk-s3.workspace = true
 aws-sdk-secretsmanager.workspace = true
 clap.workspace = true

diff --git a/iris-mpc-common/src/helpers/smpc_request.rs b/iris-mpc-common/src/helpers/smpc_request.rs
@@ -1,9 +1,15 @@
 use super::{key_pair::SharesDecodingError, sha256::calculate_sha256};
 use crate::helpers::key_pair::SharesEncryptionKeyPairs;
+use aws_sdk_sqs::{
+    error::SdkError,
+    operation::{delete_message::DeleteMessageError, receive_message::ReceiveMessageError},
+};
 use base64::{engine::general_purpose::STANDARD, Engine};
+use eyre::Report;
 use reqwest::Client;
 use serde::{Deserialize, Serialize};
 use std::sync::LazyLock;
+use thiserror::Error;
 use tokio_retry::{
     strategy::{jitter, FixedInterval},
     Retry,
@@ -27,14 +33,62 @@ pub struct SQSMessage {
     pub unsubscribe_url:   String,
 }
 
+pub const SMPC_REQUEST_TYPE_ATTRIBUTE: &str = "message_type";
+pub const IDENTITY_DELETION_REQUEST_TYPE: &str = "identity_deletion";
+pub const UNIQUENESS_REQUEST_TYPE: &str = "uniqueness";
+
 #[derive(Serialize, Deserialize, Debug, Clone)]
-pub struct SMPCRequest {
+pub struct UniquenessRequest {
     pub batch_size:              Option<usize>,
     pub signup_id:               String,
     pub s3_presigned_url:        String,
     pub iris_shares_file_hashes: [String; 3],
 }
 
+#[derive(Serialize, Deserialize, Debug, Clone)]
+pub struct IdentityDeletionRequest {
+    pub serial_id: u32,
+}
+
+#[derive(Error, Debug)]
+pub enum ReceiveRequestError {
+    #[error("Failed to read from request SQS: {0}")]
+    FailedToReadFromSQS(#[from] SdkError<ReceiveMessageError>),
+
+    #[error("Failed to delete request from SQS: {0}")]
+    FailedToDeleteFromSQS(#[from] SdkError<DeleteMessageError>),
+
+    #[error("Failed to mark request as deleted in the database: {0}")]
+    FailedToMarkRequestAsDeleted(#[from] Report),
+
+    #[error("Failed to parse {json_name} JSON: {err}")]
+    JsonParseError {
+        json_name: String,
+        err:       serde_json::Error,
+    },
+
+    #[error("Request does not contain a message type attribute")]
+    NoMessageTypeAttribute,
+
+    #[error("Request does not contain a string message type attribute")]
+    NoStringMessageTypeAttribute,
+
+    #[error("Message type attribute is not valid")]
+    InvalidMessageType,
+
+    #[error("Failed to join receive handle: {0}")]
+    FailedToJoinHandle(#[from] tokio::task::JoinError),
+}
+
+impl ReceiveRequestError {
+    pub fn json_parse_error(json_name: &str, err: serde_json::error::Error) -> Self {
+        ReceiveRequestError::JsonParseError {
+            json_name: json_name.to_string(),
+            err,
+        }
+    }
+}
+
 #[derive(Serialize, Deserialize, Debug, Clone)]
 pub struct SharesS3Object {
     pub iris_share_0: String,
@@ -65,7 +119,7 @@ impl SharesS3Object {
 
 static S3_HTTP_CLIENT: LazyLock<Client> = LazyLock::new(Client::new);
 
-impl SMPCRequest {
+impl UniquenessRequest {
     pub async fn get_iris_data_by_party_id(
         &self,
         party_id: usize,
@@ -152,6 +206,7 @@ impl SMPCRequest {
 
         Ok(iris_share)
     }
+
     pub fn validate_iris_share(
         &self,
         party_id: usize,

diff --git a/iris-mpc-common/tests/smpc_request.rs b/iris-mpc-common/tests/smpc_request.rs
@@ -4,7 +4,7 @@ mod tests {
     use iris_mpc_common::helpers::{
         key_pair::{SharesDecodingError, SharesEncryptionKeyPairs},
         sha256::calculate_sha256,
-        smpc_request::{IrisCodesJSON, SMPCRequest},
+        smpc_request::{IrisCodesJSON, UniquenessRequest},
     };
     use serde_json::json;
     use sodiumoxide::crypto::{box_::PublicKey, sealedbox};
@@ -40,17 +40,17 @@ mod tests {
         }
     }
 
-    fn get_mock_smpc_request_with_hashes(hashes: [String; 3]) -> SMPCRequest {
-        SMPCRequest {
+    fn get_mock_smpc_request_with_hashes(hashes: [String; 3]) -> UniquenessRequest {
+        UniquenessRequest {
             batch_size:              Some(1),
             signup_id:               "signup_mock".to_string(),
             s3_presigned_url:        "https://example.com/mock".to_string(),
             iris_shares_file_hashes: hashes,
         }
     }
 
-    fn get_mock_request() -> SMPCRequest {
-        SMPCRequest {
+    fn get_mock_request() -> UniquenessRequest {
+        UniquenessRequest {
             batch_size:              None,
             signup_id:               "test_signup_id".to_string(),
             s3_presigned_url:        "https://example.com/package".to_string(),
@@ -81,7 +81,7 @@ mod tests {
             .mount(&mock_server)
             .await;
 
-        let smpc_request = SMPCRequest {
+        let smpc_request = UniquenessRequest {
             batch_size:              None,
             signup_id:               "test_signup_id".to_string(),
             s3_presigned_url:        mock_server.uri().clone() + "/test_presign_url",

diff --git a/iris-mpc-gpu/src/server/mod.rs b/iris-mpc-gpu/src/server/mod.rs
@@ -27,15 +27,17 @@ pub struct BatchMetadata {
 
 #[derive(Default, Debug, Clone, PartialEq, Eq, Hash)]
 pub struct BatchQuery {
-    pub request_ids:   Vec<String>,
-    pub metadata:      Vec<BatchMetadata>,
-    pub query_left:    BatchQueryEntries,
-    pub db_left:       BatchQueryEntries,
-    pub store_left:    BatchQueryEntries,
-    pub query_right:   BatchQueryEntries,
-    pub db_right:      BatchQueryEntries,
-    pub store_right:   BatchQueryEntries,
-    pub valid_entries: Vec<bool>,
+    pub request_ids:                Vec<String>,
+    pub metadata:                   Vec<BatchMetadata>,
+    pub query_left:                 BatchQueryEntries,
+    pub db_left:                    BatchQueryEntries,
+    pub store_left:                 BatchQueryEntries,
+    pub query_right:                BatchQueryEntries,
+    pub db_right:                   BatchQueryEntries,
+    pub store_right:                BatchQueryEntries,
+    pub deletion_requests:          Vec<u32>,
+    pub deletion_requests_metadata: Vec<BatchMetadata>,
+    pub valid_entries:              Vec<bool>,
 }
 
 macro_rules! filter_by_indices {

diff --git a/iris-mpc/src/bin/client.rs b/iris-mpc/src/bin/client.rs
@@ -9,7 +9,7 @@ use iris_mpc_common::{
     helpers::{
         key_pair::download_public_key,
         sha256::calculate_sha256,
-        smpc_request::{IrisCodesJSON, ResultEvent, SMPCRequest},
+        smpc_request::{IrisCodesJSON, ResultEvent, UniquenessRequest},
         sqs_s3_helper::upload_file_and_generate_presigned_url,
     },
     iris_db::{db::IrisDB, iris::IrisCode},
@@ -343,7 +343,7 @@ async fn main() -> eyre::Result<()> {
                     }
                 };
 
-                let request_message = SMPCRequest {
+                let request_message = UniquenessRequest {
                     batch_size: None,
                     signup_id: request_id.to_string(),
                     s3_presigned_url: presigned_url,