From a337dcbc8eae6dd616c0f38d54b6e17ea6647f2f Mon Sep 17 00:00:00 2001
From: Carlo Mazzaferro <carlo.mazzaferro@gmail.com>
Date: Thu, 19 Sep 2024 10:17:24 +0200
Subject: [PATCH] Fix/seed v1 dbs (#416)

* minor changes to seeding

* bump releases

* fix dockerfile

* remove server shutdown
---
 Dockerfile                                     |  2 +-
 iris-mpc-upgrade/src/bin/seed_v1_dbs.rs        | 12 ++++++++++--
 iris-mpc-upgrade/src/bin/tcp_upgrade_client.rs | 13 +------------
 3 files changed, 12 insertions(+), 15 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index ebe21398e..f728c9268 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -38,7 +38,7 @@ COPY --from=build-app /src/gpu-iris-mpc/target/x86_64-unknown-linux-gnu/release/
 COPY --from=build-app /src/gpu-iris-mpc/target/x86_64-unknown-linux-gnu/release/key-manager /bin/key-manager
 COPY --from=build-app /src/gpu-iris-mpc/target/x86_64-unknown-linux-gnu/release/upgrade-server /bin/upgrade-server
 COPY --from=build-app /src/gpu-iris-mpc/target/x86_64-unknown-linux-gnu/release/upgrade-client /bin/upgrade-client
-COPY --from=build-app /src/gpu-iris-mpc/target/x86_64-unknown-linux-gnu/release/upgrade-client /bin/upgrade-checker
+COPY --from=build-app /src/gpu-iris-mpc/target/x86_64-unknown-linux-gnu/release/upgrade-checker /bin/upgrade-checker
 
 USER 65534
 ENTRYPOINT ["/bin/server"]
diff --git a/iris-mpc-upgrade/src/bin/seed_v1_dbs.rs b/iris-mpc-upgrade/src/bin/seed_v1_dbs.rs
index 6f12dbe2d..74a06a68e 100644
--- a/iris-mpc-upgrade/src/bin/seed_v1_dbs.rs
+++ b/iris-mpc-upgrade/src/bin/seed_v1_dbs.rs
@@ -64,11 +64,11 @@ async fn main() -> eyre::Result<()> {
     };
 
     let shares_db0 = Db::new(&shares_db_config0).await?;
-    let latest_shares_id_0 = shares_db0.fetch_latest_share_id().await?;
+    let mut latest_shares_id_0 = shares_db0.fetch_latest_share_id().await?;
     let shares_db1 = Db::new(&shares_db_config1).await?;
     let latest_shares_id_1 = shares_db1.fetch_latest_share_id().await?;
     let masks_db = Db::new(&masks_db_config).await?;
-    let latest_masks_id = masks_db.fetch_latest_mask_id().await?;
+    let mut latest_masks_id = masks_db.fetch_latest_mask_id().await?;
 
     if latest_shares_id_0 != latest_shares_id_1 {
         return Err(eyre::eyre!(
@@ -84,6 +84,14 @@ async fn main() -> eyre::Result<()> {
     let mut shares0 = Vec::with_capacity(args.fill_to as usize);
     let mut shares1 = Vec::with_capacity(args.fill_to as usize);
 
+    if latest_shares_id_0 == 0 {
+        latest_shares_id_0 += 1;
+    }
+
+    if latest_masks_id == 0 {
+        latest_masks_id += 1;
+    }
+
     for i in latest_shares_id_0..args.fill_to {
         let mut iris_code = rng.gen::<Template>();
         // fix the iris code mask to be valid: all chunks of 2 bits are equal, since
diff --git a/iris-mpc-upgrade/src/bin/tcp_upgrade_client.rs b/iris-mpc-upgrade/src/bin/tcp_upgrade_client.rs
index 24aa1adb3..3cade55f2 100644
--- a/iris-mpc-upgrade/src/bin/tcp_upgrade_client.rs
+++ b/iris-mpc-upgrade/src/bin/tcp_upgrade_client.rs
@@ -18,10 +18,7 @@ use rand::{Rng, SeedableRng};
 use rand_chacha::ChaCha20Rng;
 use rustls::{pki_types::ServerName, ClientConfig};
 use std::{array, convert::TryFrom, pin::Pin, sync::Arc};
-use tokio::{
-    io::{AsyncReadExt, AsyncWriteExt},
-    net::TcpStream,
-};
+use tokio::{io::AsyncWriteExt, net::TcpStream};
 use tokio_rustls::{client::TlsStream, TlsConnector};
 
 fn install_tracing() {
@@ -269,15 +266,7 @@ 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(())
 }