Merge pull request #1473 from akoshelev/https-sharded-shuffle-test-2

HTTPS test for sharded shuffle

ipa-core/src/app.rs

@@ -220,7 +220,6 @@ impl RequestHandler<HelperIdentity> for Inner {
         data: BodyStream,
     ) -> Result<HelperResponse, ApiError> {
         let qp = &self.query_processor;
-
         Ok(match req.route {
             r @ RouteId::Records => {
                 return Err(ApiError::BadRequest(

ipa-core/src/bin/helper.rs

@@ -371,7 +371,7 @@ pub async fn main() {
     let res = match args.command {
         None => server(args.server, handle).await,
         Some(HelperCommand::Keygen(args)) => keygen(&args),
-        Some(HelperCommand::TestSetup(args)) => test_setup(args),
+        Some(HelperCommand::TestSetup(args)) => test_setup(&args),
         Some(HelperCommand::Confgen(args)) => client_config_setup(args),
         Some(HelperCommand::ShardedConfgen(args)) => sharded_client_config_setup(args),
     };

ipa-core/src/cli/clientconf.rs

@@ -9,6 +9,7 @@ use crate::{
     },
     error::BoxError,
     helpers::HelperIdentity,
+    sharding::ShardIndex,
 };
 
 #[derive(Debug, Args)]
@@ -157,7 +158,7 @@ fn create_sharded_conf_from_files(
             let mut shard_dir = base_dir.clone();
             let id_nr: u8 = id.into();
             shard_dir.push(format!("helper{id_nr}"));
-            shard_dir.push(format!("shard{ix}"));
+            shard_dir.push(shard_conf_folder(ix));
 
             let host_name = find_file_with_extension(&shard_dir, "pem").unwrap();
             let tls_cert_file = shard_dir.join(format!("{host_name}.pem"));
@@ -223,3 +224,7 @@ fn gen_conf_from_args(
     );
     Ok(())
 }
+
+pub fn shard_conf_folder<I: TryInto<ShardIndex>>(shard_id: I) -> PathBuf {
+    format!("shard{}", shard_id.try_into().ok().unwrap()).into()
+}

ipa-core/src/cli/config_parse.rs

@@ -123,8 +123,8 @@ fn parse_sharded_network_toml(input: &str) -> Result<ShardedNetworkToml, Error>
 
 /// Generates client configuration file at the requested destination. The destination must exist
 /// before this function is called
-pub fn gen_client_config(
-    clients_conf: impl Iterator<Item = HelperClientConf>,
+pub fn gen_client_config<I: IntoIterator<Item = HelperClientConf>>(
+    clients_conf: I,
     use_http1: bool,
     conf_file: &mut File,
 ) -> Result<(), BoxError> {
@@ -352,7 +352,7 @@ pub fn sharded_server_from_toml_str(
             identities: shard_count.iter().collect(),
         };
         Ok((mpc_network, shard_network))
-    } else if missing_urls == [0, 1, 2] && shard_count == ShardIndex(1) {
+    } else if missing_urls == [0, 1, 2] && shard_count == ShardIndex::from(1) {
         // This is the special case we're dealing with a non-sharded, single ring MPC.
         // Since the shard network will be of size 1, it can't really communicate with anyone else.
         // Hence we just create a config where I'm the only shard. We take the MPC configuration

ipa-core/src/cli/test_setup.rs

@@ -1,11 +1,12 @@
 use std::{
     fs::{DirBuilder, File},
     iter::zip,
-    path::PathBuf,
+    path::{Path, PathBuf},
 };
 
 use clap::Args;
 
+use super::clientconf::shard_conf_folder;
 use crate::{
     cli::{
         config_parse::{gen_client_config, HelperClientConf},
@@ -66,7 +67,7 @@ impl TestSetupArgs {
 ///
 /// # Panics
 /// If something that shouldn't happen goes wrong.
-pub fn test_setup(args: TestSetupArgs) -> Result<(), BoxError> {
+pub fn test_setup(args: &TestSetupArgs) -> Result<(), BoxError> {
     assert_eq!(
         args.ports.len(),
         args.shard_ports.len(),
@@ -97,96 +98,72 @@ pub fn test_setup(args: TestSetupArgs) -> Result<(), BoxError> {
     }
 }
 
-fn sharded_keygen(args: TestSetupArgs) -> Result<(), BoxError> {
-    let localhost = String::from("localhost");
-    let keygen_args: Vec<_> = [1, 2, 3]
-        .into_iter()
-        .cycle()
-        .take(args.ports.len())
-        .enumerate()
-        .map(|(i, id)| {
-            let shard_dir = args.output_dir.join(format!("shard{i}"));
-            DirBuilder::new().create(&shard_dir)?;
-            Ok::<_, BoxError>(if i < 3 {
-                // Only leader shards generate MK keys
-                KeygenArgs {
-                    name: localhost.clone(),
-                    tls_cert: shard_dir.helper_tls_cert(id),
-                    tls_key: shard_dir.helper_tls_key(id),
-                    tls_expire_after: 365,
-                    mk_public_key: Some(shard_dir.helper_mk_public_key(id)),
-                    mk_private_key: Some(shard_dir.helper_mk_private_key(id)),
-                }
-            } else {
-                KeygenArgs {
-                    name: localhost.clone(),
-                    tls_cert: shard_dir.helper_tls_cert(id),
-                    tls_key: shard_dir.helper_tls_key(id),
-                    tls_expire_after: 365,
-                    mk_public_key: None,
-                    mk_private_key: None,
-                }
-            })
-        })
-        .collect::<Result<_, _>>()?;
-
-    for ka in &keygen_args {
-        keygen(ka)?;
-    }
+fn sharded_keygen(args: &TestSetupArgs) -> Result<(), BoxError> {
+    const RING_SIZE: usize = 3;
 
+    // we split all ports into chunks of 3 (for each MPC ring) and go over
+    // all of them, creating configuration
     let clients_config: Vec<_> = zip(
-        keygen_args.iter(),
-        zip(
-            keygen_args.clone().into_iter().take(3).cycle(),
-            zip(args.ports, args.shard_ports),
-        ),
-    )
-    .map(
-        |(keygen, (leader_keygen, (port, shard_port)))| HelperClientConf {
-            host: localhost.to_string(),
-            port,
-            shard_port,
-            tls_cert_file: keygen.tls_cert.clone(),
-            mk_public_key_file: leader_keygen.mk_public_key.clone().unwrap(),
-        },
+        args.ports.chunks(RING_SIZE),
+        args.shard_ports.chunks(RING_SIZE),
     )
-    .collect();
+    .enumerate()
+    .flat_map(|(shard_id, (mpc_ports, shard_ports))| {
+        let shard_dir = args.output_dir.join(shard_conf_folder(shard_id));
+        DirBuilder::new().create(&shard_dir)?;
+        make_client_configs(mpc_ports, shard_ports, &shard_dir)
+    })
+    .flatten()
+    .collect::<Vec<_>>();
 
     let mut conf_file = File::create(args.output_dir.join("network.toml"))?;
-    gen_client_config(clients_config.into_iter(), args.use_http1, &mut conf_file)
+    gen_client_config(clients_config, args.use_http1, &mut conf_file)
 }
 
 /// This generates directories and files needed to run a non-sharded MPC.
 /// The directory structure is flattened and does not include per-shard configuration.
-fn non_sharded_keygen(args: TestSetupArgs) -> Result<(), BoxError> {
+fn non_sharded_keygen(args: &TestSetupArgs) -> Result<(), BoxError> {
+    let client_configs = make_client_configs(&args.ports, &args.shard_ports, &args.output_dir)?;
+
+    let mut conf_file = File::create(args.output_dir.join("network.toml"))?;
+    gen_client_config(client_configs, args.use_http1, &mut conf_file)
+}
+
+fn make_client_configs(
+    mpc_ports: &[u16],
+    shard_ports: &[u16],
+    config_dir: &Path,
+) -> Result<Vec<HelperClientConf>, BoxError> {
+    assert_eq!(shard_ports.len(), mpc_ports.len());
+    assert_eq!(3, shard_ports.len());
+
     let localhost = String::from("localhost");
-    let clients_config: [_; 3] = zip([1, 2, 3], zip(args.ports, args.shard_ports))
-        .map(|(id, (port, shard_port))| {
+    zip(mpc_ports.iter(), shard_ports.iter())
+        .enumerate()
+        .map(|(i, (&mpc_port, &shard_port))| {
+            let id = u8::try_from(i + 1).unwrap();
+
+            // TODO: only leader shards should generate MK encryptions.
             let keygen_args = KeygenArgs {
                 name: localhost.clone(),
-                tls_cert: args.output_dir.helper_tls_cert(id),
-                tls_key: args.output_dir.helper_tls_key(id),
+                tls_cert: config_dir.helper_tls_cert(id),
+                tls_key: config_dir.helper_tls_key(id),
                 tls_expire_after: 365,
-                mk_public_key: Some(args.output_dir.helper_mk_public_key(id)),
-                mk_private_key: Some(args.output_dir.helper_mk_private_key(id)),
+                mk_public_key: Some(config_dir.helper_mk_public_key(id)),
+                mk_private_key: Some(config_dir.helper_mk_private_key(id)),
             };
 
             keygen(&keygen_args)?;
 
             Ok(HelperClientConf {
                 host: localhost.to_string(),
-                port,
+                port: mpc_port,
                 shard_port,
                 tls_cert_file: keygen_args.tls_cert,
                 mk_public_key_file: keygen_args.mk_public_key.unwrap(),
             })
         })
-        .collect::<Result<Vec<_>, BoxError>>()?
-        .try_into()
-        .unwrap();
-
-    let mut conf_file = File::create(args.output_dir.join("network.toml"))?;
-    gen_client_config(clients_config.into_iter(), args.use_http1, &mut conf_file)
+        .collect::<Result<_, _>>()
 }
 
 #[cfg(test)]
@@ -212,7 +189,7 @@ mod tests {
             ports: vec![3000, 3001, 3002, 3003, 3004, 3005],
             shard_ports: vec![6000, 6001, 6002, 6003, 6004, 6005],
         };
-        test_setup(args).unwrap();
+        test_setup(&args).unwrap();
         let network_config_path = outdir.join("network.toml");
         let network_config_string = &fs::read_to_string(network_config_path).unwrap();
         let _r = sharded_server_from_toml_str(
@@ -237,7 +214,7 @@ mod tests {
             ports: vec![],
             shard_ports: vec![],
         };
-        test_setup(args).unwrap();
+        test_setup(&args).unwrap();
     }
 
     #[test]
@@ -252,6 +229,6 @@ mod tests {
             ports: vec![3000, 3001],
             shard_ports: vec![6000, 6001, 6002],
         };
-        test_setup(args).unwrap();
+        test_setup(&args).unwrap();
     }
 }

ipa-core/src/config.rs

@@ -602,6 +602,6 @@ mod tests {
         let pc1 = PeerConfig::new(uri1, None);
         let client = ClientConfig::default();
         let conf = NetworkConfig::new_shards(vec![pc1.clone()], client);
-        assert_eq!(conf.peers[ShardIndex(0)].url, pc1.url);
+        assert_eq!(conf.peers[ShardIndex::FIRST].url, pc1.url);
     }
 }

ipa-core/src/net/test.rs

@@ -138,7 +138,7 @@ impl TestNetwork<Shard> {
     /// Creates all the shards for a helper and creates a network.
     fn new_shards(id: HelperIdentity, ports: Vec<Option<u16>>, conf: &TestConfigBuilder) -> Self {
         let servers: Vec<_> = (0..conf.shard_count)
-            .map(ShardIndex)
+            .map(ShardIndex::from)
             .zip(ports)
             .map(|(ix, p)| {
                 let sid = ShardedHelperIdentity::new(id, ix);
@@ -296,7 +296,7 @@ impl TestConfig {
     /// Creates a new [`TestConfig`] using the provided configuration.
     fn new(conf: &TestConfigBuilder) -> Self {
         let rings = (0..conf.shard_count)
-            .map(ShardIndex)
+            .map(ShardIndex::from)
             .map(|s| {
                 let ports = conf.get_ports_for_shard_index(s);
                 TestNetwork::<Helper>::new_mpc(s, ports, conf)
@@ -1049,7 +1049,7 @@ mod tests {
         let builder = TestConfigBuilder::with_http_and_default_test_ports();
         assert_eq!(
             vec![Some(3000), Some(3001), Some(3002)],
-            builder.get_ports_for_shard_index(ShardIndex(0))
+            builder.get_ports_for_shard_index(ShardIndex::FIRST)
         );
         assert_eq!(
             vec![Some(6001)],
@@ -1060,6 +1060,9 @@ mod tests {
     #[test]
     fn get_os_ports() {
         let builder = TestConfigBuilder::default();
-        assert_eq!(3, builder.get_ports_for_shard_index(ShardIndex(0)).len());
+        assert_eq!(
+            3,
+            builder.get_ports_for_shard_index(ShardIndex::FIRST).len()
+        );
     }
 }

ipa-core/src/protocol/hybrid/agg.rs

@@ -187,6 +187,7 @@ pub mod test {
     // the inputs are laid out to work with exactly 2 shards
     // as if it we're resharded by match_key/prf
     const SHARDS: usize = 2;
+    const SECOND_SHARD: ShardIndex = ShardIndex::from_u32(1);
 
     // we re-use these as the "prf" of the match_key
     // to avoid needing to actually do the prf here
@@ -374,8 +375,8 @@ pub mod test {
             let results: Vec<[Vec<[AggregateableHybridReport<BA8, BA3>; 2]>; 3]> = world
                 .malicious(records.clone().into_iter(), |ctx, input| {
                     let match_keys = match ctx.shard_id() {
-                        ShardIndex(0) => SHARD1_MKS,
-                        ShardIndex(1) => SHARD2_MKS,
+                        ShardIndex::FIRST => SHARD1_MKS,
+                        SECOND_SHARD => SHARD2_MKS,
                         _ => panic!("invalid shard_id"),
                     };
                     async move {
@@ -446,8 +447,8 @@ pub mod test {
             let results: Vec<[Vec<AggregateableHybridReport<BA8, BA3>>; 3]> = world
                 .malicious(records.clone().into_iter(), |ctx, input| {
                     let match_keys = match ctx.shard_id() {
-                        ShardIndex(0) => SHARD1_MKS,
-                        ShardIndex(1) => SHARD2_MKS,
+                        ShardIndex::FIRST => SHARD1_MKS,
+                        SECOND_SHARD => SHARD2_MKS,
                         _ => panic!("invalid shard_id"),
                     };
                     async move {
@@ -572,8 +573,8 @@ pub mod test {
             let _results: Vec<[Vec<AggregateableHybridReport<BA8, BA3>>; 3]> = world
                 .malicious(records.clone().into_iter(), |ctx, input| {
                     let match_keys = match ctx.shard_id() {
-                        ShardIndex(0) => SHARD1_MKS,
-                        ShardIndex(1) => SHARD2_MKS,
+                        ShardIndex::FIRST => SHARD1_MKS,
+                        SECOND_SHARD => SHARD2_MKS,
                         _ => panic!("invalid shard_id"),
                     };
                     async move {