Merge from main

ipa-core/Cargo.toml

@@ -82,6 +82,7 @@ ipa-step = { version = "*", path = "../ipa-step" }
 ipa-step-derive = { version = "*", path = "../ipa-step-derive" }
 
 aes = "0.8.3"
+assertions = "0.1.0"
 async-trait = "0.1.79"
 async-scoped = { version = "0.9.0", features = ["use-tokio"], optional = true }
 axum = { version = "0.7.5", optional = true, features = ["http2", "macros"] }

ipa-core/src/bin/helper.rs

@@ -127,7 +127,7 @@ async fn server(args: ServerArgs) -> Result<(), BoxError> {
                 }),
             )
         }
-        (None, None) => (ClientIdentity::Helper(my_identity), None),
+        (None, None) => (ClientIdentity::Header(my_identity), None),
         _ => panic!("should have been rejected by clap"),
     };
 

ipa-core/src/bin/report_collector.rs

@@ -22,7 +22,7 @@ use ipa_core::{
     config::{KeyRegistries, NetworkConfig},
     ff::{boolean_array::BA32, FieldType},
     helpers::query::{DpMechanism, IpaQueryConfig, QueryConfig, QuerySize, QueryType},
-    net::MpcHelperClient,
+    net::{Helper, MpcHelperClient},
     report::{EncryptedOprfReportStreams, DEFAULT_KEY_ID},
     test_fixture::{
         ipa::{ipa_in_the_clear, CappingOrder, IpaSecurityModel, TestRawDataRecord},
@@ -380,7 +380,7 @@ async fn ipa(
 
 async fn ipa_test(
     args: &Args,
-    network: &NetworkConfig,
+    network: &NetworkConfig<Helper>,
     security_model: IpaSecurityModel,
     ipa_query_config: IpaQueryConfig,
     helper_clients: &[MpcHelperClient; 3],

ipa-core/src/cli/clientconf.rs

@@ -186,7 +186,7 @@ fn assert_network_config(config_toml: &Map<String, Value>, config_str: &str) {
     else {
         panic!("peers section in toml config is not a table");
     };
-    for (i, peer_config_actual) in nw_config.peers.iter().enumerate() {
+    for (i, peer_config_actual) in nw_config.peers().iter().enumerate() {
         assert_peer_config(&peer_config_expected[i], peer_config_actual);
     }
 }

ipa-core/src/cli/crypto/encrypt.rs

@@ -244,7 +244,7 @@ this is not toml!
     }
 
     #[test]
-    #[should_panic = "invalid length 2, expected an array of length 3"]
+    #[should_panic = "Expected a Vec of length 3 but it was 2"]
     fn encrypt_incomplete_network_file() {
         let input_file = sample_data::write_csv(sample_data::test_ipa_data().take(10)).unwrap();
 

ipa-core/src/cli/playbook/mod.rs

@@ -20,7 +20,7 @@ use crate::{
     executor::IpaRuntime,
     ff::boolean_array::{BA20, BA3, BA8},
     helpers::query::DpMechanism,
-    net::{ClientIdentity, MpcHelperClient},
+    net::{ClientIdentity, Helper, MpcHelperClient},
     protocol::{dp::NoiseParams, ipa_prf::oprf_padding::insecure::OPRFPaddingDp},
 };
 
@@ -194,19 +194,19 @@ pub async fn make_clients(
     network_path: Option<&Path>,
     scheme: Scheme,
     wait: usize,
-) -> ([MpcHelperClient; 3], NetworkConfig) {
+) -> ([MpcHelperClient; 3], NetworkConfig<Helper>) {
     let mut wait = wait;
     let network = if let Some(path) = network_path {
         NetworkConfig::from_toml_str(&fs::read_to_string(path).unwrap()).unwrap()
     } else {
-        NetworkConfig {
-            peers: [
+        NetworkConfig::<Helper>::new_mpc(
+            vec![
                 PeerConfig::new("localhost:3000".parse().unwrap(), None),
                 PeerConfig::new("localhost:3001".parse().unwrap(), None),
                 PeerConfig::new("localhost:3002".parse().unwrap(), None),
             ],
-            client: ClientConfig::default(),
-        }
+            ClientConfig::default(),
+        )
     };
     let network = network.override_scheme(&scheme);
 

ipa-core/src/config.rs

@@ -1,13 +1,12 @@
 use std::{
-    array,
     borrow::{Borrow, Cow},
     fmt::{Debug, Formatter},
-    iter::Zip,
+    iter::zip,
     path::PathBuf,
-    slice,
     time::Duration,
 };
 
+use base64::{engine::general_purpose::STANDARD as BASE64, Engine as _};
 use hyper::{http::uri::Scheme, Uri};
 use hyper_util::client::legacy::Builder;
 use rustls_pemfile::Item;
@@ -22,6 +21,8 @@ use crate::{
         Deserializable as _, IpaPrivateKey, IpaPublicKey, KeyRegistry, PrivateKeyOnly,
         PublicKeyOnly, Serializable as _,
     },
+    net::{ConnectionFlavor, Helper, Shard},
+    sharding::ShardIndex,
 };
 
 pub type OwnedCertificate = CertificateDer<'static>;
@@ -37,23 +38,115 @@ pub enum Error {
     IOError(#[from] std::io::Error),
 }
 
-/// Configuration information describing a helper network.
+/// Configuration describing either 3 peers in a Ring or N shard peers. In a non-sharded case a
+/// single [`NetworkConfig`] represents the entire network. In a sharded case, each host should
+/// have one Ring and one Sharded configuration to know how to reach its peers.
 ///
 /// The most important thing this contains is discovery information for each of the participating
-/// helpers.
+/// peers.
 #[derive(Clone, Debug, Deserialize)]
-pub struct NetworkConfig {
-    /// Information about each helper participating in the network. The order that helpers are
-    /// listed here determines their assigned helper identities in the network. Note that while the
-    /// helper identities are stable, roles are assigned per query.
-    pub peers: [PeerConfig; 3],
+pub struct NetworkConfig<F: ConnectionFlavor = Helper> {
+    peers: Vec<PeerConfig>,
 
     /// HTTP client configuration.
     #[serde(default)]
     pub client: ClientConfig,
+
+    /// The identities of the index-matching peers. Separating this from [`Self::peers`](field) so
+    /// that parsing is easy to implement.
+    #[serde(skip)]
+    identities: Vec<F::Identity>,
+}
+
+impl<F: ConnectionFlavor> NetworkConfig<F> {
+    /// # Panics
+    /// If `PathAndQuery::from_str("")` fails
+    #[must_use]
+    pub fn override_scheme(self, scheme: &Scheme) -> Self {
+        Self {
+            peers: self
+                .peers
+                .into_iter()
+                .map(|mut peer| {
+                    let mut parts = peer.url.into_parts();
+                    parts.scheme = Some(scheme.clone());
+                    // `http::uri::Uri::from_parts()` requires that a URI have a path if it has a
+                    // scheme. If the URI does not have a scheme, it is not required to have a path.
+                    if parts.path_and_query.is_none() {
+                        parts.path_and_query = Some("".parse().unwrap());
+                    }
+                    peer.url = Uri::try_from(parts).unwrap();
+                    peer
+                })
+                .collect(),
+            ..self
+        }
+    }
+
+    #[must_use]
+    pub fn vec_peers(&self) -> Vec<PeerConfig> {
+        self.peers.clone()
+    }
+
+    #[must_use]
+    pub fn get_peer(&self, i: usize) -> Option<&PeerConfig> {
+        self.peers.get(i)
+    }
+
+    pub fn peers_iter(&self) -> std::slice::Iter<'_, PeerConfig> {
+        self.peers.iter()
+    }
+
+    /// We currently require an exact match with the peer cert (i.e. we don't support verifying
+    /// the certificate against a truststore and identifying the peer by the certificate
+    /// subject). This could be changed if the need arises.
+    #[must_use]
+    pub fn identify_cert(&self, cert: Option<&CertificateDer>) -> Option<F::Identity> {
+        let cert = cert?;
+        for (id, p) in zip(self.identities.iter(), self.peers.iter()) {
+            if p.certificate.as_ref() == Some(cert) {
+                return Some(*id);
+            }
+        }
+        // It might be nice to log something here. We could log the certificate base64?
+        tracing::error!(
+            "A client certificate was presented that does not match a known helper. Certificate: {}",
+            BASE64.encode(cert),
+        );
+        None
+    }
+}
+
+impl NetworkConfig<Shard> {
+    /// # Panics
+    /// In the unlikely event a usize cannot be turned into a u32
+    #[must_use]
+    pub fn new_shards(peers: Vec<PeerConfig>, client: ClientConfig) -> Self {
+        let identities = (0u32..peers.len().try_into().unwrap())
+            .map(ShardIndex::from)
+            .collect();
+        Self {
+            peers,
+            client,
+            identities,
+        }
+    }
 }
 
-impl NetworkConfig {
+impl NetworkConfig<Helper> {
+    /// Creates a new configuration for 3 MPC clients (ring) configuration.
+    /// # Panics
+    /// If the vector doesn't contain exactly 3 items.
+    #[must_use]
+    pub fn new_mpc(ring: Vec<PeerConfig>, client: ClientConfig) -> Self {
+        assert_eq!(3, ring.len());
+        Self {
+            peers: ring,
+            client,
+            identities: HelperIdentity::make_three().to_vec(),
+        }
+    }
+
     /// Reads config from string. Expects config to be toml format.
     /// To read file, use `fs::read_to_string`
     ///
@@ -62,49 +155,25 @@ impl NetworkConfig {
     pub fn from_toml_str(input: &str) -> Result<Self, Error> {
         use config::{Config, File, FileFormat};
 
-        let conf: Self = Config::builder()
+        let mut conf: Self = Config::builder()
             .add_source(File::from_str(input, FileFormat::Toml))
             .build()?
             .try_deserialize()?;
 
-        Ok(conf)
-    }
-
-    pub fn new(peers: [PeerConfig; 3], client: ClientConfig) -> Self {
-        Self { peers, client }
-    }
+        conf.identities = HelperIdentity::make_three().to_vec();
 
-    pub fn peers(&self) -> &[PeerConfig; 3] {
-        &self.peers
-    }
-
-    // Can maybe be replaced with array::zip when stable?
-    pub fn enumerate_peers(
-        &self,
-    ) -> Zip<array::IntoIter<HelperIdentity, 3>, slice::Iter<PeerConfig>> {
-        HelperIdentity::make_three()
-            .into_iter()
-            .zip(self.peers.iter())
+        Ok(conf)
     }
 
+    /// Clones the internal configs and returns them as an array.
     /// # Panics
-    /// If `PathAndQuery::from_str("")` fails
+    /// If the internal vector isn't of size 3.
     #[must_use]
-    pub fn override_scheme(self, scheme: &Scheme) -> NetworkConfig {
-        NetworkConfig {
-            peers: self.peers.map(|mut peer| {
-                let mut parts = peer.url.into_parts();
-                parts.scheme = Some(scheme.clone());
-                // `http::uri::Uri::from_parts()` requires that a URI have a path if it has a
-                // scheme. If the URI does not have a scheme, it is not required to have a path.
-                if parts.path_and_query.is_none() {
-                    parts.path_and_query = Some("".parse().unwrap());
-                }
-                peer.url = Uri::try_from(parts).unwrap();
-                peer
-            }),
-            ..self
-        }
+    pub fn peers(&self) -> [PeerConfig; 3] {
+        self.peers
+            .clone()
+            .try_into()
+            .unwrap_or_else(|v: Vec<_>| panic!("Expected a Vec of length 3 but it was {}", v.len()))
     }
 }
 
@@ -422,10 +491,11 @@ impl KeyRegistries {
     /// If network file is improperly formatted
     pub fn init_from(
         &mut self,
-        network: &NetworkConfig,
+        network: &NetworkConfig<Helper>,
     ) -> Option<[&KeyRegistry<PublicKeyOnly>; 3]> {
         // Get the configs, if all three peers have one
-        let configs = network.peers().iter().try_fold(Vec::new(), |acc, peer| {
+        let peers = network.peers();
+        let configs = peers.iter().try_fold(Vec::new(), |acc, peer| {
             if let (mut vec, Some(hpke_config)) = (acc, peer.hpke_config.as_ref()) {
                 vec.push(hpke_config);
                 Some(vec)
@@ -453,10 +523,12 @@ mod tests {
     use rand::rngs::StdRng;
     use rand_core::SeedableRng;
 
+    use super::{NetworkConfig, PeerConfig};
     use crate::{
         config::{ClientConfig, HpkeClientConfig, Http2Configurator, HttpClientConfigurator},
         helpers::HelperIdentity,
         net::test::TestConfigBuilder,
+        sharding::ShardIndex,
     };
 
     const URI_1: &str = "http://localhost:3000";
@@ -531,4 +603,13 @@ mod tests {
             }),
         );
     }
+
+    #[test]
+    fn indexing_peer_happy_case() {
+        let uri1 = URI_1.parse::<Uri>().unwrap();
+        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);
+    }
 }

ipa-core/src/helpers/buffers/unordered_receiver.rs

@@ -295,7 +295,6 @@ where
     inner: Arc<Mutex<OperatingState<S, C>>>,
 }
 
-#[allow(dead_code)]
 impl<S, C> UnorderedReceiver<S, C>
 where
     S: Stream<Item = C> + Send,

ipa-core/src/helpers/mod.rs

@@ -181,6 +181,10 @@ impl HelperIdentity {
     pub const TWO: Self = Self { id: 2 };
     pub const THREE: Self = Self { id: 3 };
 
+    pub const ONE_STR: &'static str = "A";
+    pub const TWO_STR: &'static str = "B";
+    pub const THREE_STR: &'static str = "C";
+
     /// Given a helper identity, return an array of the identities of the other two helpers.
     // The order that helpers are returned here is not intended to be meaningful, however,
     // it is currently used directly to determine the assignment of roles in