From 8c4cd1c154b598316a1d706740552ba80777c2ea Mon Sep 17 00:00:00 2001
From: Daniel Kales <11509575+dkales@users.noreply.github.com>
Date: Tue, 10 Sep 2024 12:10:59 +0200
Subject: [PATCH] add basic tls setup with self-signed certs

---
 Cargo.lock                                    | 184 +++++++++++++++++-
 iris-mpc-upgrade/Cargo.toml                   |   7 +-
 iris-mpc-upgrade/src/bin/Readme.md            |  30 +--
 iris-mpc-upgrade/src/bin/checker.rs           |  34 ++--
 iris-mpc-upgrade/src/bin/gen_certs.rs         |  29 +++
 iris-mpc-upgrade/src/bin/seed_v1_dbs.rs       |   6 +-
 .../src/bin/tcp_upgrade_client.rs             |  80 ++++----
 .../src/bin/tcp_upgrade_server.rs             |  36 +++-
 iris-mpc-upgrade/src/config.rs                |  38 ++--
 iris-mpc-upgrade/src/db.rs                    |   5 +-
 10 files changed, 355 insertions(+), 94 deletions(-)
 create mode 100644 iris-mpc-upgrade/src/bin/gen_certs.rs

diff --git a/Cargo.lock b/Cargo.lock
index 8002837b1..f43070fd5 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -221,6 +221,33 @@ dependencies = [
  "zeroize",
 ]
 
+[[package]]
+name = "aws-lc-rs"
+version = "1.9.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2f95446d919226d587817a7d21379e6eb099b97b45110a7f272a444ca5c54070"
+dependencies = [
+ "aws-lc-sys",
+ "mirai-annotations",
+ "paste",
+ "zeroize",
+]
+
+[[package]]
+name = "aws-lc-sys"
+version = "0.21.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "234314bd569802ec87011d653d6815c6d7b9ffb969e9fee5b8b20ef860e8dce9"
+dependencies = [
+ "bindgen",
+ "cc",
+ "cmake",
+ "dunce",
+ "fs_extra",
+ "libc",
+ "paste",
+]
+
 [[package]]
 name = "aws-runtime"
 version = "1.4.2"
@@ -763,6 +790,29 @@ dependencies = [
  "serde",
 ]
 
+[[package]]
+name = "bindgen"
+version = "0.69.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a00dc851838a2120612785d195287475a3ac45514741da670b735818822129a0"
+dependencies = [
+ "bitflags 2.6.0",
+ "cexpr",
+ "clang-sys",
+ "itertools 0.12.1",
+ "lazy_static",
+ "lazycell",
+ "log",
+ "prettyplease",
+ "proc-macro2",
+ "quote",
+ "regex",
+ "rustc-hash",
+ "shlex",
+ "syn 2.0.77",
+ "which",
+]
+
 [[package]]
 name = "bitflags"
 version = "1.3.2"
@@ -889,9 +939,20 @@ version = "1.1.15"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "57b6a275aa2903740dc87da01c62040406b8812552e97129a63ea8850a17c6e6"
 dependencies = [
+ "jobserver",
+ "libc",
  "shlex",
 ]
 
+[[package]]
+name = "cexpr"
+version = "0.6.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6fac387a98bb7c37292057cffc56d62ecb629900026402633ae9160df93a8766"
+dependencies = [
+ "nom",
+]
+
 [[package]]
 name = "cfg-if"
 version = "1.0.0"
@@ -949,6 +1010,17 @@ dependencies = [
  "inout",
 ]
 
+[[package]]
+name = "clang-sys"
+version = "1.8.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0b023947811758c97c59bf9d1c188fd619ad4718dcaa767947df1cadb14f39f4"
+dependencies = [
+ "glob",
+ "libc",
+ "libloading",
+]
+
 [[package]]
 name = "clap"
 version = "4.5.16"
@@ -989,6 +1061,15 @@ version = "0.7.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "1462739cb27611015575c0c11df5df7601141071f07518d56fcc1be504cbec97"
 
+[[package]]
+name = "cmake"
+version = "0.1.50"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a31c789563b815f77f4250caee12365734369f942439b7defd71e18a48197130"
+dependencies = [
+ "cc",
+]
+
 [[package]]
 name = "colorchoice"
 version = "1.0.2"
@@ -1420,6 +1501,12 @@ version = "0.15.7"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "1aaf95b3e5c8f23aa320147307562d361db0ae0d51242340f558153b4eb2439b"
 
+[[package]]
+name = "dunce"
+version = "1.0.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "92773504d58c093f6de2459af4af33faa518c13451eb8f2b5698ed3d36e7c813"
+
 [[package]]
 name = "ecdsa"
 version = "0.14.8"
@@ -1612,6 +1699,12 @@ dependencies = [
  "percent-encoding",
 ]
 
+[[package]]
+name = "fs_extra"
+version = "1.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c"
+
 [[package]]
 name = "funty"
 version = "2.0.0"
@@ -2447,15 +2540,16 @@ dependencies = [
  "mpc",
  "rand",
  "rand_chacha",
- "rustls 0.21.12",
+ "rcgen",
+ "rustls 0.23.12",
+ "rustls-pemfile 2.1.3",
  "serde",
  "serde-big-array",
  "sqlx",
  "tokio",
- "tokio-rustls 0.24.1",
+ "tokio-rustls 0.26.0",
  "tracing",
  "tracing-subscriber",
- "webpki-roots",
 ]
 
 [[package]]
@@ -2517,6 +2611,15 @@ version = "1.0.11"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "49f1f14873335454500d59611f1cf4a4b0f786f9ac11f4312a78e4cf2566695b"
 
+[[package]]
+name = "jobserver"
+version = "0.1.32"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "48d1dbcbbeb6a7fec7e059840aa538bd62aaccf972c7346c4d9d2059312853d0"
+dependencies = [
+ "libc",
+]
+
 [[package]]
 name = "js-sys"
 version = "0.3.70"
@@ -2546,6 +2649,12 @@ dependencies = [
  "spin",
 ]
 
+[[package]]
+name = "lazycell"
+version = "1.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "830d08ce1d1d941e6b30645f1a0eb5643013d835ce3779a5fc208261dbe10f55"
+
 [[package]]
 name = "libc"
 version = "0.2.158"
@@ -2810,6 +2919,12 @@ dependencies = [
  "windows-sys 0.52.0",
 ]
 
+[[package]]
+name = "mirai-annotations"
+version = "1.12.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c9be0862c1b3f26a88803c4a49de6889c10e608b3ee9344e6ef5b45fb37ad3d1"
+
 [[package]]
 name = "mongodb"
 version = "2.8.2"
@@ -3273,6 +3388,16 @@ dependencies = [
  "digest",
 ]
 
+[[package]]
+name = "pem"
+version = "3.0.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8e459365e590736a54c3fa561947c84837534b8e9af6fc5bf781307e82658fae"
+dependencies = [
+ "base64 0.22.1",
+ "serde",
+]
+
 [[package]]
 name = "pem-rfc7468"
 version = "0.7.0"
@@ -3460,6 +3585,16 @@ dependencies = [
  "zerocopy",
 ]
 
+[[package]]
+name = "prettyplease"
+version = "0.2.20"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5f12335488a2f3b0a83b14edad48dca9879ce89b2edd10e80237e4e852dd645e"
+dependencies = [
+ "proc-macro2",
+ "syn 2.0.77",
+]
+
 [[package]]
 name = "proc-macro2"
 version = "1.0.86"
@@ -3570,6 +3705,19 @@ dependencies = [
  "crossbeam-utils",
 ]
 
+[[package]]
+name = "rcgen"
+version = "0.13.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "54077e1872c46788540de1ea3d7f4ccb1983d12f9aa909b234468676c1a36779"
+dependencies = [
+ "pem",
+ "ring",
+ "rustls-pki-types",
+ "time",
+ "yasna",
+]
+
 [[package]]
 name = "redox_syscall"
 version = "0.4.1"
@@ -3828,6 +3976,12 @@ version = "0.1.24"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "719b953e2095829ee67db738b3bfa9fa368c94900df327b3f07fe6e794d2fe1f"
 
+[[package]]
+name = "rustc-hash"
+version = "1.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2"
+
 [[package]]
 name = "rustc_version"
 version = "0.2.3"
@@ -3887,6 +4041,8 @@ version = "0.23.12"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "c58f8c84392efc0a126acce10fa59ff7b3d2ac06ab451a33f2741989b806b044"
 dependencies = [
+ "aws-lc-rs",
+ "log",
  "once_cell",
  "rustls-pki-types",
  "rustls-webpki 0.102.7",
@@ -3947,6 +4103,7 @@ version = "0.102.7"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "84678086bd54edf2b415183ed7a94d0efb049f1b646a33e22a36f3794be6ae56"
 dependencies = [
+ "aws-lc-rs",
  "ring",
  "rustls-pki-types",
  "untrusted",
@@ -5365,6 +5522,18 @@ version = "0.25.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "5f20c57d8d7db6d3b86154206ae5d8fba62dd39573114de97c2cb0578251f8e1"
 
+[[package]]
+name = "which"
+version = "4.4.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "87ba24419a2078cd2b0f2ede2691b6c66d8e47836da3b6db8265ebad47afbfc7"
+dependencies = [
+ "either",
+ "home",
+ "once_cell",
+ "rustix",
+]
+
 [[package]]
 name = "whoami"
 version = "1.5.1"
@@ -5666,6 +5835,15 @@ dependencies = [
  "linked-hash-map",
 ]
 
+[[package]]
+name = "yasna"
+version = "0.5.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e17bb3549cc1321ae1296b9cdc2698e2b6cb1992adfa19a8c72e5b7a738f44cd"
+dependencies = [
+ "time",
+]
+
 [[package]]
 name = "zerocopy"
 version = "0.7.35"
diff --git a/iris-mpc-upgrade/Cargo.toml b/iris-mpc-upgrade/Cargo.toml
index 7413564a0..fb8df2e32 100644
--- a/iris-mpc-upgrade/Cargo.toml
+++ b/iris-mpc-upgrade/Cargo.toml
@@ -25,9 +25,10 @@ tracing-subscriber.workspace = true
 
 mpc-uniqueness-check = { package = "mpc", git = "https://github.com/worldcoin/mpc-uniqueness-check", rev = "4d38402" }
 indicatif = "0.17.8"
-tokio-rustls = "0.24.1"
-rustls = "0.21.12"
-webpki-roots = "0.25.4"
+tokio-rustls = "0.26.0"
+rustls = "0.23.12"
+rustls-pemfile = "2.1.3"
+rcgen = "0.13.1"
 
 [dev-dependencies]
 float_eq = "1"
diff --git a/iris-mpc-upgrade/src/bin/Readme.md b/iris-mpc-upgrade/src/bin/Readme.md
index e6e9e1103..b3aef3577 100644
--- a/iris-mpc-upgrade/src/bin/Readme.md
+++ b/iris-mpc-upgrade/src/bin/Readme.md
@@ -16,6 +16,14 @@ will bring up two old dbs on ports 6200 and 6201, and 3 new dbs on ports 6100,61
 cargo run --release --bin seed-v1-dbs -- --shares-db-urls postgres://postgres:postgres@localhost:6100/shares --shares-db-urls postgres://postgres:postgres@localhost:6101/shares --masks-db-url postgres://postgres:postgres@localhost:6100/masks --num-elements 10000
 ```
 
+## Generate Self-Signed certificates for Client-Server communication
+
+```bash
+cargo run --release --bin gen_certs -- --sans localhost --sans party0 --key-path key0.pem --cert-path cert0.pem
+cargo run --release --bin gen_certs -- --sans localhost --sans party1 --key-path key1.pem --cert-path cert1.pem
+cargo run --release --bin gen_certs -- --sans localhost --sans party2 --key-path key2.pem --cert-path cert2.pem
+```
+
 ## Upgrade for left eye
 
 ### Run the 3 upgrade servers
@@ -23,15 +31,15 @@ cargo run --release --bin seed-v1-dbs -- --shares-db-urls postgres://postgres:po
 Concurrently run:
 
 ```bash
-cargo run --release --bin upgrade-server -- --bind-addr 127.0.0.1:8000 --db-url postgres://postgres:postgres@localhost:6200/postgres --party-id 0 --eye left
+cargo run --release --bin upgrade-server -- --bind-addr 127.0.0.1:8000 --key key0.pem --cert-chain cert0.pem --db-url postgres://postgres:postgres@localhost:6200/postgres --party-id 0 --eye left
 ```
 
 ```bash
-cargo run --release --bin upgrade-server -- --bind-addr 127.0.0.1:8001 --db-url postgres://postgres:postgres@localhost:6201/postgres --party-id 1 --eye left
+cargo run --release --bin upgrade-server -- --bind-addr 127.0.0.1:8001 --key key1.pem --cert-chain cert1.pem --db-url postgres://postgres:postgres@localhost:6201/postgres --party-id 1 --eye left
 ```
 
 ```bash
-cargo run --release --bin upgrade-server -- --bind-addr 127.0.0.1:8002 --db-url postgres://postgres:postgres@localhost:6202/postgres --party-id 2 --eye left
+cargo run --release --bin upgrade-server -- --bind-addr 127.0.0.1:8002 --key key2.pem --cert-chain cert2.pem --db-url postgres://postgres:postgres@localhost:6202/postgres --party-id 2 --eye left
 ```
 
 ### Run the 2 upgrade clients
@@ -39,11 +47,11 @@ cargo run --release --bin upgrade-server -- --bind-addr 127.0.0.1:8002 --db-url
 Concurrently run:
 
 ```bash
-cargo run --release --bin upgrade-client -- --server1 127.0.0.1:8000 --server2 127.0.0.1:8001 --server3 127.0.0.1:8002 --db-start 0 --db-end 10000 --party-id 0 --eye left --shares-db-url postgres://postgres:postgres@localhost:6100/shares --masks-db-url postgres://postgres:postgres@localhost:6100/masks
+cargo run --release --bin upgrade-client -- --server1 localhost:8000 --server2 localhost:8001 --server3 localhost:8002 --trusted-cert cert0.pem --trusted-cert cert1.pem --trusted-cert cert2.pem --db-start 1 --db-end 10001 --party-id 0 --eye left --shares-db-url postgres://postgres:postgres@localhost:6100/shares --masks-db-url postgres://postgres:postgres@localhost:6100/masks
 ```
 
 ```bash
-cargo run --release --bin upgrade-client -- --server1 127.0.0.1:8000 --server2 127.0.0.1:8001 --server3 127.0.0.1:8002 --db-start 0 --db-end 10000 --party-id 1 --eye left --shares-db-url postgres://postgres:postgres@localhost:6101/shares --masks-db-url postgres://postgres:postgres@localhost:6100/masks
+cargo run --release --bin upgrade-client -- --server1 localhost:8000 --server2 localhost:8001 --server3 localhost:8002 --trusted-cert cert0.pem --trusted-cert cert1.pem --trusted-cert cert2.pem --db-start 1 --db-end 10001 --party-id 1 --eye left --shares-db-url postgres://postgres:postgres@localhost:6101/shares --masks-db-url postgres://postgres:postgres@localhost:6100/masks
 ```
 
 ## Upgrade for right eye
@@ -53,15 +61,15 @@ cargo run --release --bin upgrade-client -- --server1 127.0.0.1:8000 --server2 1
 Concurrently run:
 
 ```bash
-cargo run --release --bin upgrade-server -- --bind-addr 127.0.0.1:8000 --db-url postgres://postgres:postgres@localhost:6200/postgres --party-id 0 --eye right
+cargo run --release --bin upgrade-server -- --bind-addr 127.0.0.1:8000 --key key0.pem --cert-chain cert0.pem --db-url postgres://postgres:postgres@localhost:6200/postgres --party-id 0 --eye right
 ```
 
 ```bash
-cargo run --release --bin upgrade-server -- --bind-addr 127.0.0.1:8001 --db-url postgres://postgres:postgres@localhost:6201/postgres --party-id 1 --eye right
+cargo run --release --bin upgrade-server -- --bind-addr 127.0.0.1:8001 --key key1.pem --cert-chain cert1.pem --db-url postgres://postgres:postgres@localhost:6201/postgres --party-id 1 --eye right
 ```
 
 ```bash
-cargo run --release --bin upgrade-server -- --bind-addr 127.0.0.1:8002 --db-url postgres://postgres:postgres@localhost:6202/postgres --party-id 2 --eye right
+cargo run --release --bin upgrade-server -- --bind-addr 127.0.0.1:8002 --key key2.pem --cert-chain cert2.pem --db-url postgres://postgres:postgres@localhost:6202/postgres --party-id 2 --eye right
 ```
 
 ### Run the 2 upgrade clients
@@ -71,15 +79,15 @@ cargo run --release --bin upgrade-server -- --bind-addr 127.0.0.1:8002 --db-url
 Concurrently run:
 
 ```bash
-cargo run --release --bin upgrade-client -- --server1 127.0.0.1:8000 --server2 127.0.0.1:8001 --server3 127.0.0.1:8002 --db-start 0 --db-end 10000 --party-id 0 --eye right --shares-db-url postgres://postgres:postgres@localhost:6100/shares --masks-db-url postgres://postgres:postgres@localhost:6100/masks
+cargo run --release --bin upgrade-client -- --server1 localhost:8000 --server2 localhost:8001 --server3 localhost:8002 --trusted-cert cert0.pem --trusted-cert cert1.pem --trusted-cert cert2.pem --db-start 1 --db-end 10001 --party-id 0 --eye right --shares-db-url postgres://postgres:postgres@localhost:6100/shares --masks-db-url postgres://postgres:postgres@localhost:6100/masks
 ```
 
 ```bash
-cargo run --release --bin upgrade-client -- --server1 127.0.0.1:8000 --server2 127.0.0.1:8001 --server3 127.0.0.1:8002 --db-start 0 --db-end 10000 --party-id 1 --eye right --shares-db-url postgres://postgres:postgres@localhost:6101/shares --masks-db-url postgres://postgres:postgres@localhost:6100/masks
+cargo run --release --bin upgrade-client -- --server1 localhost:8000 --server2 localhost:8001 --server3 localhost:8002 --trusted-cert cert0.pem --trusted-cert cert1.pem --trusted-cert cert2.pem --db-start 1 --db-end 10001 --party-id 1 --eye right --shares-db-url postgres://postgres:postgres@localhost:6101/shares --masks-db-url postgres://postgres:postgres@localhost:6100/masks
 ```
 
 ## Check the upgrade was successful
 
 ```bash
-cargo run --release --bin upgrade-checker -- --num-elements 10000 --db-urls postgres://postgres:postgres@localhost:6100/postgres --db-urls postgres://postgres:postgres@localhost:6101/postgres --db-urls postgres://postgres:postgres@localhost:6100/postgres --db-urls postgres://postgres:postgres@localhost:6101/postgres --db-urls postgres://postgres:postgres@localhost:6200/postgres --db-urls postgres://postgres:postgres@localhost:6201/postgres --db-urls postgres://postgres:postgres@localhost:6202/postgres
+cargo run --release --bin upgrade-checker -- --num-elements 10000 --db-urls postgres://postgres:postgres@localhost:6100/shares --db-urls postgres://postgres:postgres@localhost:6101/shares --db-urls postgres://postgres:postgres@localhost:6100/masks --db-urls postgres://postgres:postgres@localhost:6100/shares --db-urls postgres://postgres:postgres@localhost:6101/shares --db-urls postgres://postgres:postgres@localhost:6100/masks --db-urls postgres://postgres:postgres@localhost:6200/postgres --db-urls postgres://postgres:postgres@localhost:6201/postgres --db-urls postgres://postgres:postgres@localhost:6202/postgres
 ```
diff --git a/iris-mpc-upgrade/src/bin/checker.rs b/iris-mpc-upgrade/src/bin/checker.rs
index 87353ed3f..51eefcd01 100644
--- a/iris-mpc-upgrade/src/bin/checker.rs
+++ b/iris-mpc-upgrade/src/bin/checker.rs
@@ -27,33 +27,35 @@ struct Args {
 async fn main() -> eyre::Result<()> {
     let args = Args::parse();
 
-    if args.db_urls.len() != 7 {
+    if args.db_urls.len() != 9 {
         return Err(eyre::eyre!(
-            "Expect 5 db urls to be provided: old_left_db0, old_left_db1, old_right_db0, \
-             old_right_db1, new_db0, new_db1, new_db2"
+            "Expect 9 db urls to be provided: old_left_db0, old_left_db1, old_left_mask_db, \
+             old_right_db0, old_right_db1, old_right_mask_db, new_db0, new_db1, new_db2"
         ));
     }
 
     let old_left_db0 = V1Db::new(&args.db_urls[0]).await?;
     let old_left_db1 = V1Db::new(&args.db_urls[1]).await?;
-    let old_right_db0 = V1Db::new(&args.db_urls[2]).await?;
-    let old_right_db1 = V1Db::new(&args.db_urls[3]).await?;
+    let old_left_mask_db = V1Db::new(&args.db_urls[2]).await?;
+    let old_right_db0 = V1Db::new(&args.db_urls[3]).await?;
+    let old_right_db1 = V1Db::new(&args.db_urls[4]).await?;
+    let old_right_mask_db = V1Db::new(&args.db_urls[5]).await?;
 
-    let new_db0 = Store::new(&args.db_urls[4], "upgrade").await?;
-    let new_db1 = Store::new(&args.db_urls[5], "upgrade").await?;
-    let new_db2 = Store::new(&args.db_urls[6], "upgrade").await?;
+    let new_db0 = Store::new(&args.db_urls[6], "upgrade").await?;
+    let new_db1 = Store::new(&args.db_urls[7], "upgrade").await?;
+    let new_db2 = Store::new(&args.db_urls[8], "upgrade").await?;
 
     // grab the old shares from the db and reconstruct them
     let old_left_shares0 = old_left_db0
-        .stream_shares(0..args.num_elements)
+        .stream_shares(1..args.num_elements + 1)
         .collect::<Vec<_>>()
         .await;
     let old_left_shares1 = old_left_db1
-        .stream_shares(0..args.num_elements)
+        .stream_shares(1..args.num_elements + 1)
         .collect::<Vec<_>>()
         .await;
-    let old_left_masks = old_left_db0
-        .stream_masks(0..args.num_elements)
+    let old_left_masks = old_left_mask_db
+        .stream_masks(1..args.num_elements + 1)
         .collect::<Vec<_>>()
         .await;
 
@@ -76,15 +78,15 @@ async fn main() -> eyre::Result<()> {
     .collect();
 
     let old_right_shares0 = old_right_db0
-        .stream_shares(0..args.num_elements)
+        .stream_shares(1..args.num_elements + 1)
         .collect::<Vec<_>>()
         .await;
     let old_right_shares1 = old_right_db1
-        .stream_shares(0..args.num_elements)
+        .stream_shares(1..args.num_elements + 1)
         .collect::<Vec<_>>()
         .await;
-    let old_right_masks = old_right_db0
-        .stream_masks(0..args.num_elements)
+    let old_right_masks = old_right_mask_db
+        .stream_masks(1..args.num_elements + 1)
         .collect::<Vec<_>>()
         .await;
 
diff --git a/iris-mpc-upgrade/src/bin/gen_certs.rs b/iris-mpc-upgrade/src/bin/gen_certs.rs
new file mode 100644
index 000000000..84bef840d
--- /dev/null
+++ b/iris-mpc-upgrade/src/bin/gen_certs.rs
@@ -0,0 +1,29 @@
+use clap::Parser;
+use eyre::{Context, Result};
+use std::path::PathBuf;
+
+/// Certificate Generator for MPC-NET
+#[derive(Debug, PartialEq, Parser)]
+struct CliArgs {
+    /// The path to the .der certificate file
+    #[clap(short, long)]
+    cert_path: PathBuf,
+    /// The path to the .der key file
+    #[clap(short, long)]
+    key_path:  PathBuf,
+    /// The subject alternative names for the certificate
+    #[clap(short, long)]
+    sans:      Vec<String>,
+}
+
+fn main() -> Result<()> {
+    let args = CliArgs::parse();
+
+    let cert =
+        rcgen::generate_simple_self_signed(args.sans).context("generating self-signed cert")?;
+    let key = cert.key_pair.serialize_pem();
+    std::fs::write(args.key_path, key).context("writing key file")?;
+    let cert = cert.cert.pem();
+    std::fs::write(args.cert_path, cert).context("writing certificate file")?;
+    Ok(())
+}
diff --git a/iris-mpc-upgrade/src/bin/seed_v1_dbs.rs b/iris-mpc-upgrade/src/bin/seed_v1_dbs.rs
index be665cd19..6e015cba0 100644
--- a/iris-mpc-upgrade/src/bin/seed_v1_dbs.rs
+++ b/iris-mpc-upgrade/src/bin/seed_v1_dbs.rs
@@ -54,12 +54,12 @@ async fn main() -> eyre::Result<()> {
     let mut shares0 = Vec::with_capacity(args.num_elements as usize);
     let mut shares1 = Vec::with_capacity(args.num_elements as usize);
 
-    for i in 0..args.num_elements {
+    for i in 1..args.num_elements + 1 {
         let mut iris_code = rng.gen::<Template>();
         // fix the iris code mask to be valid: all chunks of 2 bits are equal, since
         // they mask the real/imaginary party of the same bit
-        for i in (0..BITS).step_by(2) {
-            iris_code.mask.set(i + 1, iris_code.mask.get(i))
+        for j in (0..BITS).step_by(2) {
+            iris_code.mask.set(j + 1, iris_code.mask.get(j))
         }
         let encoded = mpc_uniqueness_check::distance::encode(&iris_code).share(2, &mut rng);
         masks.push((i, iris_code.mask));
diff --git a/iris-mpc-upgrade/src/bin/tcp_upgrade_client.rs b/iris-mpc-upgrade/src/bin/tcp_upgrade_client.rs
index 0ee3532de..f3764dc42 100644
--- a/iris-mpc-upgrade/src/bin/tcp_upgrade_client.rs
+++ b/iris-mpc-upgrade/src/bin/tcp_upgrade_client.rs
@@ -1,5 +1,5 @@
 use clap::Parser;
-use eyre::ContextCompat;
+use eyre::{Context, ContextCompat};
 use futures::{Stream, StreamExt};
 use futures_concurrency::future::Join;
 use indicatif::{ProgressBar, ProgressStyle};
@@ -16,18 +16,13 @@ use iris_mpc_upgrade::{
 use mpc_uniqueness_check::{bits::Bits, distance::EncodedBits};
 use rand::{Rng, SeedableRng};
 use rand_chacha::ChaCha20Rng;
-use std::{array, pin::Pin};
+use rustls::{pki_types::ServerName, ClientConfig};
+use std::{array, convert::TryFrom, pin::Pin, sync::Arc};
 use tokio::{
-    io::{BufWriter},
+    io::{AsyncReadExt, AsyncWriteExt},
+    net::TcpStream,
 };
-
-use tokio::net::TcpStream;
-use tokio_rustls::{rustls::{ServerName, OwnedTrustAnchor}, TlsConnector, rustls::ClientConfig, client::TlsStream};
-use std::sync::Arc;
-use webpki_roots::TLS_SERVER_ROOTS;
-use tokio::io::{AsyncWrite, AsyncRead, AsyncWriteExt, AsyncReadExt};
-use std::convert::TryFrom;
-use std::error::Error;
+use tokio_rustls::{client::TlsStream, TlsConnector};
 
 fn install_tracing() {
     use tracing_subscriber::{fmt, prelude::*, EnvFilter};
@@ -43,29 +38,21 @@ fn install_tracing() {
         .init();
 }
 
-async fn prepare_tls_stream_for_writing(address: &str) -> eyre::Result<TlsStream<TcpStream>> {
+async fn prepare_tls_stream_for_writing(
+    address: &str,
+    client_config: Arc<ClientConfig>,
+) -> eyre::Result<TlsStream<TcpStream>> {
     // Create a TCP connection
     let stream = TcpStream::connect(address).await?;
 
-    // Create a basic client config with the default root certificate store
-    let mut root_cert_store = rustls::RootCertStore::empty();
-    root_cert_store.add_trust_anchors(TLS_SERVER_ROOTS.iter().map(|ta| {
-        OwnedTrustAnchor::from_subject_spki_name_constraints(
-            ta.subject.as_ref(),
-            ta.spki.as_ref(),
-            ta.name_constraints.clone().map(|nc| nc.as_ref()),
-        )
-    }));
-
-    let config = ClientConfig::builder()
-        .with_safe_defaults()
-        .with_root_certificates(root_cert_store)
-        .with_no_client_auth();
-
-    let tls_connector = TlsConnector::from(Arc::new(config));
+    let tls_connector = TlsConnector::from(client_config);
 
     // Hostname for SNI (Server Name Indication)
-    let dns_name = ServerName::try_from("localhost").unwrap();
+    // throw away the port number if there is one (e.g. "localhost:8080" ->
+    // "localhost")
+    let address = address.split(":").next().context("splitting address")?;
+    let dns_name =
+        ServerName::try_from(address.to_owned()).context("trying to convert address to SNI")?;
 
     // Perform the TLS handshake to establish a secure connection
     let tls_stream: TlsStream<TcpStream> = tls_connector.connect(dns_name, stream).await?;
@@ -82,14 +69,30 @@ async fn main() -> eyre::Result<()> {
         panic!("Party id must be 0, 1");
     }
 
-    let tls_stream_1 = prepare_tls_stream_for_writing(&args.server1).await?;
-    let tls_stream_2 = prepare_tls_stream_for_writing(&args.server2).await?;
-    let tls_stream_3 = prepare_tls_stream_for_writing(&args.server3).await?;
+    // read the trusted cert
+    let mut root_cert_store = rustls::RootCertStore::empty();
+    for trusted_cert in &args.trusted_cert {
+        let cert_pem = tokio::fs::read(trusted_cert).await?;
+        if !trusted_cert.ends_with(".pem") {
+            tracing::warn!("Expected trusted_chain file to end in \".pem\"");
+        }
+        let cert_chain =
+            rustls_pemfile::certs(&mut &cert_pem[..]).collect::<Result<Vec<_>, _>>()?;
+        for cert in cert_chain {
+            root_cert_store.add(cert)?;
+        }
+    }
+    let client_config = Arc::new(
+        ClientConfig::builder()
+            .with_root_certificates(root_cert_store)
+            .with_no_client_auth(),
+    );
+
+    let mut server1 = prepare_tls_stream_for_writing(&args.server1, client_config.clone()).await?;
+    let mut server2 = prepare_tls_stream_for_writing(&args.server2, client_config.clone()).await?;
+    let mut server3 = prepare_tls_stream_for_writing(&args.server3, client_config).await?;
 
     tracing::info!("Connecting to servers and syncing migration task parameters...");
-    let mut server1 = BufWriter::new(tls_stream_1);
-    let mut server2 = BufWriter::new(tls_stream_2);
-    let mut server3 = BufWriter::new(tls_stream_3);
     server1.write_u8(args.party_id).await?;
     server2.write_u8(args.party_id).await?;
     server3.write_u8(args.party_id).await?;
@@ -274,6 +277,13 @@ async fn main() -> eyre::Result<()> {
         pb.inc(diff);
     }
     tracing::info!("Processing done!");
+    let mut buf = [0u8; 1];
+    server1.read_exact(&mut buf[..]).await?;
+    server2.read_exact(&mut buf[..]).await?;
+    server3.read_exact(&mut buf[..]).await?;
+    server1.shutdown().await?;
+    server2.shutdown().await?;
+    server3.shutdown().await?;
     pb.finish();
 
     Ok(())
diff --git a/iris-mpc-upgrade/src/bin/tcp_upgrade_server.rs b/iris-mpc-upgrade/src/bin/tcp_upgrade_server.rs
index 73b451695..49db6791d 100644
--- a/iris-mpc-upgrade/src/bin/tcp_upgrade_server.rs
+++ b/iris-mpc-upgrade/src/bin/tcp_upgrade_server.rs
@@ -1,6 +1,6 @@
 use axum::{routing::get, Router};
 use clap::Parser;
-use eyre::{bail, Context};
+use eyre::{bail, eyre, Context};
 use futures_concurrency::future::Join;
 use iris_mpc_common::{helpers::task_monitor::TaskMonitor, id::PartyID};
 use iris_mpc_store::Store;
@@ -9,15 +9,17 @@ use iris_mpc_upgrade::{
     packets::{MaskShareMessage, TwoToThreeIrisCodeMessage},
     IrisCodeUpgrader, NewIrisShareSink,
 };
+use rustls::ServerConfig;
 use std::{
     sync::{atomic::AtomicUsize, Arc},
     time::{Duration, Instant},
 };
 use tokio::{
-    io::{AsyncReadExt, BufReader},
+    io::{AsyncReadExt, AsyncWriteExt},
     sync::mpsc,
     task::JoinSet,
 };
+use tokio_rustls::TlsAcceptor;
 
 fn install_tracing() {
     use tracing_subscriber::{fmt, prelude::*, EnvFilter};
@@ -89,10 +91,33 @@ async fn main() -> eyre::Result<()> {
     tracing::info!("Spawned processing tasks");
 
     // listen for incoming connections from clients
+    if !args.key.ends_with(".pem") {
+        tracing::warn!("Expected key file to end in \".pem\"");
+    }
+    let key_pem = tokio::fs::read(&args.key).await?;
+    let key = rustls_pemfile::private_key(&mut &key_pem[..])?.ok_or(eyre!(
+        "could not parse {} as a valid private key",
+        args.key.display()
+    ))?;
+
+    if !args.cert_chain.ends_with(".pem") {
+        tracing::warn!("Expected cert_chain file to end in \".pem\"");
+    }
+    let cert_pem = tokio::fs::read(&args.cert_chain).await?;
+    let cert_chain = rustls_pemfile::certs(&mut &cert_pem[..]).collect::<Result<Vec<_>, _>>()?;
+    let server_config = Arc::new(
+        ServerConfig::builder()
+            .with_no_client_auth()
+            .with_single_cert(cert_chain, key)?,
+    );
+    let tls_acceptor = TlsAcceptor::from(server_config);
     let client_listener = tokio::net::TcpListener::bind(args.bind_addr).await?;
 
-    let mut client_stream1 = BufReader::new(client_listener.accept().await?.0);
-    let mut client_stream2 = BufReader::new(client_listener.accept().await?.0);
+    let client_stream1_raw = client_listener.accept().await?.0;
+    let mut client_stream1 = tls_acceptor.accept(client_stream1_raw).await?;
+    let client_stream2_raw = client_listener.accept().await?.0;
+    let mut client_stream2 = tls_acceptor.accept(client_stream2_raw).await?;
+
     tracing::info!("Both Clients connected");
     let id1 = client_stream1.read_u8().await?;
     let id2 = client_stream2.read_u8().await?;
@@ -188,6 +213,9 @@ async fn main() -> eyre::Result<()> {
         sending += start.elapsed();
     }
 
+    client_stream1.write_all(&[0]).await?;
+    client_stream2.write_all(&[0]).await?;
+
     tracing::debug!("Receiving took: {}s", receiving.as_secs_f64());
     tracing::debug!("Sending took: {}s", sending.as_secs_f64());
     // close all senders
diff --git a/iris-mpc-upgrade/src/config.rs b/iris-mpc-upgrade/src/config.rs
index 5e8ab12f1..1e5e4da50 100644
--- a/iris-mpc-upgrade/src/config.rs
+++ b/iris-mpc-upgrade/src/config.rs
@@ -2,8 +2,8 @@ use clap::Parser;
 use iris_mpc_common::id::PartyID;
 use std::{
     fmt::{self, Formatter},
-    io,
-    net::{SocketAddr, ToSocketAddrs},
+    net::SocketAddr,
+    path::PathBuf,
     str::FromStr,
 };
 
@@ -31,6 +31,12 @@ pub struct UpgradeServerConfig {
     #[clap(long)]
     pub bind_addr: SocketAddr,
 
+    #[clap(long)]
+    pub key: PathBuf,
+
+    #[clap(long)]
+    pub cert_chain: PathBuf,
+
     #[clap(long)]
     pub db_url: String,
 
@@ -48,6 +54,8 @@ impl std::fmt::Debug for UpgradeServerConfig {
     fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
         f.debug_struct("UpgradeServerConfig")
             .field("bind_addr", &self.bind_addr)
+            .field("key", &self.key)
+            .field("cert", &self.cert_chain)
             .field("db_url", &"<redacted>")
             .field("party_id", &self.party_id)
             .field("threads", &self.threads)
@@ -58,14 +66,17 @@ impl std::fmt::Debug for UpgradeServerConfig {
 
 #[derive(Parser)]
 pub struct UpgradeClientConfig {
-    #[clap(long,  default_value = "127.0.0.1:8000", value_parser = resolve_host)]
-    pub server1: SocketAddr,
+    #[clap(long, default_value = "localhost:8000")]
+    pub server1: String,
 
-    #[clap(long,  default_value = "127.0.0.1:8001", value_parser = resolve_host)]
-    pub server2: SocketAddr,
+    #[clap(long, default_value = "localhost:8001")]
+    pub server2: String,
 
-    #[clap(long,  default_value = "127.0.0.1:8002", value_parser = resolve_host)]
-    pub server3: SocketAddr,
+    #[clap(long, default_value = "localhost:8002")]
+    pub server3: String,
+
+    #[clap(long)]
+    pub trusted_cert: Vec<PathBuf>,
 
     #[clap(long)]
     pub db_start: u64,
@@ -89,22 +100,13 @@ pub struct UpgradeClientConfig {
     pub masks_db_url: String,
 }
 
-fn resolve_host(hostname_port: &str) -> io::Result<SocketAddr> {
-    let socketaddr = hostname_port.to_socket_addrs()?.next().ok_or_else(|| {
-        io::Error::new(
-            io::ErrorKind::AddrNotAvailable,
-            format!("Could not find destination {hostname_port}"),
-        )
-    })?;
-    Ok(socketaddr)
-}
-
 impl std::fmt::Debug for UpgradeClientConfig {
     fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
         f.debug_struct("UpgradeClientConfig")
             .field("server1", &self.server1)
             .field("server2", &self.server2)
             .field("server3", &self.server3)
+            .field("trusted_cert", &self.trusted_cert)
             .field("shares_db_url", &"<redacted>")
             .field("masks_db_url", &"<redacted>")
             .field("db_start", &self.db_start)
diff --git a/iris-mpc-upgrade/src/db.rs b/iris-mpc-upgrade/src/db.rs
index 44b9bf9de..732cea113 100644
--- a/iris-mpc-upgrade/src/db.rs
+++ b/iris-mpc-upgrade/src/db.rs
@@ -1,3 +1,4 @@
+use eyre::Context;
 use futures::Stream;
 use mpc_uniqueness_check::{bits::Bits, encoded_bits::EncodedBits};
 use sqlx::Postgres;
@@ -10,7 +11,9 @@ impl V1Db {
     pub async fn new(url: &str) -> eyre::Result<Self> {
         tracing::info!("Connecting to database");
 
-        let pool = sqlx::Pool::connect(url).await?;
+        let pool = sqlx::Pool::connect(url)
+            .await
+            .with_context(|| format!("when connecting to {}", &url))?;
 
         tracing::info!("Connected to database");