Merge pull request #49 from worldcoin/0xkitsune/utils

feat(utils): utilities to seed MPC and iris code database

Cargo.toml

@@ -22,6 +22,8 @@ indicatif = "0.17.7"
 indoc = "2.0.4"
 itertools = "0.12.0"
 metrics = "0.21.1"
+# TODO: only include this for migrate-codes bin
+mongodb = "2.8.1"
 ordered-float = "4.2.0"
 rand = "0.8.5"
 rayon = "1.8.1"
@@ -55,7 +57,7 @@ path = "bin/mpc_node.rs"
 
 [[bin]]
 name = "utils"
-path = "bin/utils.rs"
+path = "bin/utils/utils.rs"
 
 [[bin]]
 name = "e2e"

bin/utils/README.md

@@ -0,0 +1,40 @@
+# MPC Utils
+
+
+## Seed MPC Databases
+Seed the coordinator and participant databases with random shares and masks. When seeding multiple participant databases, specify a `--participant-db-url` for each participant db.
+
+```bash
+cargo run --bin utils seed-mpc-db  --coordinator-db-url <COORDINATOR_DB_URL> --participant-db-url <PARTICIPANT_0_DB_URL> --participant-db-url <PARTICIPANT_1_DB_URL> --num-templates <NUM_TEMPLATES> 
+```
+
+## Seed Iris Database
+Seed the iris code database with random iris code entries. Each entry includes a `signup_id`, `serial_id`, `iris_code_left`, `mask_code_left`, `iris_code_right` and `mask_code_right`.
+
+```bash
+cargo run --bin utils seed-iris-db  --iris-code-db <IRIS_CODE_DB> --num-templates <NUM_TEMPLATES> 
+```
+
+
+## Send a Random Query via SQS
+Send a random iris code template to the MPC setup via SQS. Ensure that you are logged into the proper AWS account for the queue url you are using. If running with Localstack, ensure to also pass `--endpoint-url <ENDPOINT_URL>` and `--region <AWS_REGION>`.
+
+```bash
+cargo run --bin utils sqs-query  --queue-url <SQS_QUEUE_URL> 
+```
+
+
+## Receive SQS Results
+Receive messages from a queue. This is useful when inspecting the results queue after sending a query to the MPC setup. If the queue is empty, the program will continually check for new message every second. Once messages are in the queue, the program will receive all messages, print them to the terminal and then delete the message from the queue. Ensure that you are logged into the proper AWS account for the queue url you are using. If running with Localstack, ensure to also pass `--endpoint-url <ENDPOINT_URL>` and `--region <AWS_REGION>`.
+
+```bash
+cargo run --bin utils sqs-receive  --queue-url <SQS_QUEUE_URL> 
+```
+
+## Generate Mock Templates
+
+Generate random templates and write the resulting items to an output file.
+
+```bash
+cargo run --bin utils generate-mock-templates --num-templates <NUM_TEMPLATES> --output <OUTPUT_FILE>
+```

bin/utils/generate_mock_templates.rs

@@ -0,0 +1,28 @@
+use clap::Args;
+use mpc::template::Template;
+use rand::{thread_rng, Rng};
+
+#[derive(Debug, Clone, Args)]
+pub struct GenerateMockTemplates {
+    #[clap(short, long)]
+    pub output: String,
+
+    #[clap(short, long)]
+    pub num_templates: usize,
+}
+
+pub async fn generate_mock_templates(
+    args: &GenerateMockTemplates,
+) -> eyre::Result<()> {
+    let mut rng = thread_rng();
+
+    let templates: Vec<Template> =
+        (0..args.num_templates).map(|_| rng.gen()).collect();
+
+    let json = serde_json::to_string(&templates)?;
+
+    //write to file
+    std::fs::write(&args.output, json)?;
+
+    Ok(())
+}