From 4647d112c01152b42f52169b877ddcfc7cc1d1ea Mon Sep 17 00:00:00 2001 From: Alex Koshelev Date: Mon, 25 Mar 2024 13:33:32 -0700 Subject: [PATCH] Minor touches --- ipa-core/src/app.rs | 15 ++++++---- ipa-core/src/bin/helper.rs | 2 +- ipa-core/src/helpers/mod.rs | 2 +- ipa-core/src/helpers/transport/query/mod.rs | 2 +- ipa-core/src/net/client/mod.rs | 30 +++++++++++-------- .../src/net/server/handlers/query/prepare.rs | 2 +- ipa-core/src/net/transport.rs | 2 +- ipa-core/src/query/processor.rs | 8 ++--- ipa-core/src/test_fixture/world.rs | 11 +++++-- 9 files changed, 43 insertions(+), 31 deletions(-) diff --git a/ipa-core/src/app.rs b/ipa-core/src/app.rs index 1fa44657d..958a56dee 100644 --- a/ipa-core/src/app.rs +++ b/ipa-core/src/app.rs @@ -58,19 +58,16 @@ impl RequestHandler for QueryRequestHandler { } RouteId::ReceiveQuery => { let req = req.into::()?; - let transport = Clone::clone(self.transport.lock().unwrap().as_ref().unwrap()); - HelperResponse::from(qp.new_query(transport, req).await?) + HelperResponse::from(qp.new_query(self.transport(), req).await?) } RouteId::PrepareQuery => { let req = req.into::()?; - let transport = Clone::clone(self.transport.lock().unwrap().as_ref().unwrap()); - HelperResponse::from(qp.prepare(&transport, req)?) + HelperResponse::from(qp.prepare(&self.transport(), req)?) } RouteId::QueryInput => { let query_id = ext_query_id(&req)?; - let transport = Clone::clone(self.transport.lock().unwrap().as_ref().unwrap()); HelperResponse::from(qp.receive_inputs( - transport, + self.transport(), QueryInput { query_id, input_stream: data, @@ -89,6 +86,12 @@ impl RequestHandler for QueryRequestHandler { } } +impl QueryRequestHandler { + fn transport(&self) -> TransportImpl { + Clone::clone(self.transport.lock().unwrap().as_ref().unwrap()) + } +} + #[derive(Clone)] pub struct RequestHandlerSetup { qp: Arc, diff --git a/ipa-core/src/bin/helper.rs b/ipa-core/src/bin/helper.rs index 6964b94f6..32f49ae32 100644 --- a/ipa-core/src/bin/helper.rs +++ b/ipa-core/src/bin/helper.rs @@ -14,7 +14,7 @@ use ipa_core::{ }, config::{hpke_registry, HpkeServerConfig, NetworkConfig, ServerConfig, TlsConfig}, error::BoxError, - helpers::{HelperIdentity, RequestHandler}, + helpers::HelperIdentity, net::{ClientIdentity, HttpTransport, MpcHelperClient}, AppSetup, }; diff --git a/ipa-core/src/helpers/mod.rs b/ipa-core/src/helpers/mod.rs index e00cadd33..a10427cfb 100644 --- a/ipa-core/src/helpers/mod.rs +++ b/ipa-core/src/helpers/mod.rs @@ -234,7 +234,7 @@ pub enum Role { } #[derive(Clone, Debug)] -#[cfg_attr(test, derive(Copy, PartialEq, Eq))] +#[cfg_attr(test, derive(PartialEq, Eq))] #[cfg_attr( feature = "enable-serde", derive(serde::Serialize, serde::Deserialize), diff --git a/ipa-core/src/helpers/transport/query/mod.rs b/ipa-core/src/helpers/transport/query/mod.rs index e7b29b858..e8e1f35ba 100644 --- a/ipa-core/src/helpers/transport/query/mod.rs +++ b/ipa-core/src/helpers/transport/query/mod.rs @@ -100,7 +100,7 @@ pub enum QueryConfigError { } #[derive(Clone, Debug)] -#[cfg_attr(test, derive(Copy, PartialEq, Eq))] +#[cfg_attr(test, derive(PartialEq, Eq))] #[cfg_attr(feature = "enable-serde", derive(Serialize, Deserialize))] pub struct PrepareQuery { pub query_id: QueryId, diff --git a/ipa-core/src/net/client/mod.rs b/ipa-core/src/net/client/mod.rs index 482f27da4..296298860 100644 --- a/ipa-core/src/net/client/mod.rs +++ b/ipa-core/src/net/client/mod.rs @@ -416,8 +416,7 @@ fn make_http_connector() -> HttpConnector { connector } -#[cfg(all(test, feature = "real-world-infra"))] // TODO fix back - // #[cfg(all(test, web_test))] // TODO fix back +#[cfg(all(test, web_test))] pub(crate) mod tests { use std::{ fmt::Debug, @@ -432,11 +431,11 @@ pub(crate) mod tests { use crate::{ ff::{FieldType, Fp31}, helpers::{ - make_boxed_handler, query::QueryType::TestMultiply, routing::Addr, ApiError, - BodyStream, BytesStream, HelperResponse, PanickingHandler, RequestHandler, - RoleAssignment, Transport, MESSAGE_PAYLOAD_SIZE_BYTES, + make_boxed_handler, query::QueryType::TestMultiply, BytesStream, HelperResponse, + PanickingHandler, RequestHandler, RoleAssignment, Transport, + MESSAGE_PAYLOAD_SIZE_BYTES, }, - net::{test::TestServer, HttpTransport}, + net::test::TestServer, protocol::step::StepNarrow, query::ProtocolResult, secret_sharing::replicated::semi_honest::AdditiveShare as Replicated, @@ -517,7 +516,7 @@ pub(crate) mod tests { let output = test_query_command( |client| async move { client.echo(expected_output).await.unwrap() }, - || Box::new(PanickingHandler::default()), + || Box::::default(), ) .await; assert_eq!(expected_output, &output); @@ -550,13 +549,14 @@ pub(crate) mod tests { #[tokio::test] async fn prepare() { - let input = PrepareQuery { - query_id: QueryId, - config: QueryConfig::new(TestMultiply, FieldType::Fp31, 1).unwrap(), - roles: RoleAssignment::new(HelperIdentity::make_three()), - }; + let config = QueryConfig::new(TestMultiply, FieldType::Fp31, 1).unwrap(); let handler = move || { make_boxed_handler(move |addr, _| async move { + let input = PrepareQuery { + query_id: QueryId, + config, + roles: RoleAssignment::new(HelperIdentity::make_three()), + }; let prepare_query = addr.into::().unwrap(); assert_eq!(prepare_query, input); @@ -566,7 +566,11 @@ pub(crate) mod tests { test_query_command( |client| { - let req = input.clone(); + let req = PrepareQuery { + query_id: QueryId, + config, + roles: RoleAssignment::new(HelperIdentity::make_three()), + }; async move { client.prepare_query(req).await.unwrap() } }, handler, diff --git a/ipa-core/src/net/server/handlers/query/prepare.rs b/ipa-core/src/net/server/handlers/query/prepare.rs index 447d22e6b..9f55793b0 100644 --- a/ipa-core/src/net/server/handlers/query/prepare.rs +++ b/ipa-core/src/net/server/handlers/query/prepare.rs @@ -71,7 +71,7 @@ mod tests { config: QueryConfig::new(TestMultiply, FieldType::Fp31, 1).unwrap(), roles: RoleAssignment::new(HelperIdentity::make_three()), }); - let expected_prepare_query = req.data; + let expected_prepare_query = req.data.clone(); let TestServer { transport, .. } = TestServer::builder() .with_request_handler(Box::new( move |addr: Addr, _data: BodyStream| { diff --git a/ipa-core/src/net/transport.rs b/ipa-core/src/net/transport.rs index 66bfb66a3..36ae1b231 100644 --- a/ipa-core/src/net/transport.rs +++ b/ipa-core/src/net/transport.rs @@ -237,7 +237,7 @@ mod tests { use crate::{ config::{NetworkConfig, ServerConfig}, ff::{FieldType, Fp31, Serializable}, - helpers::query::QueryType::TestMultiply, + helpers::query::{QueryInput, QueryType::TestMultiply}, net::{ client::ClientIdentity, test::{get_test_identity, TestConfig, TestConfigBuilder, TestServer}, diff --git a/ipa-core/src/query/processor.rs b/ipa-core/src/query/processor.rs index 9b8374aeb..81508ddcb 100644 --- a/ipa-core/src/query/processor.rs +++ b/ipa-core/src/query/processor.rs @@ -149,13 +149,13 @@ impl Processor { let prepare_request = PrepareQuery { query_id, config: req, - roles, + roles: roles.clone(), }; // Inform other parties about new query. If any of them rejects it, this join will fail try_join( - transport.send(left, prepare_request, stream::empty()), - transport.send(right, prepare_request, stream::empty()), + transport.send(left, prepare_request.clone(), stream::empty()), + transport.send(right, prepare_request.clone(), stream::empty()), ) .await .map_err(NewQueryError::Transport)?; @@ -511,7 +511,7 @@ mod tests { let req = prepare_query(identities); let transport = network.transport(identities[1]); let processor = Processor::default(); - processor.prepare(&transport, req).unwrap(); + processor.prepare(&transport, req.clone()).unwrap(); assert!(matches!( processor.prepare(&transport, req), Err(PrepareQueryError::AlreadyRunning) diff --git a/ipa-core/src/test_fixture/world.rs b/ipa-core/src/test_fixture/world.rs index c7ce12a57..3cd252fae 100644 --- a/ipa-core/src/test_fixture/world.rs +++ b/ipa-core/src/test_fixture/world.rs @@ -539,9 +539,14 @@ impl ShardWorld { let participants = make_participants(&mut StdRng::seed_from_u64(config.seed + shard_seed)); let network = InMemoryMpcNetwork::default(); - let mut gateways = network - .transports() - .map(|t| Gateway::new(QueryId, config.gateway_config, *config.role_assignment(), t)); + let mut gateways = network.transports().map(|t| { + Gateway::new( + QueryId, + config.gateway_config, + config.role_assignment().clone(), + t, + ) + }); // The name for `g` is too complicated and depends on features enabled #[allow(clippy::redundant_closure_for_method_calls)]