Merge pull request #1396 from cberkhoff/s-routes

ipa-core/src/net/server/handlers/mod.rs

@@ -3,7 +3,7 @@ mod query;
 
 use axum::Router;
 
-use crate::net::{http_serde, transport::MpcHttpTransport};
+use crate::net::{http_serde, transport::MpcHttpTransport, ShardHttpTransport};
 
 pub fn mpc_router(transport: MpcHttpTransport) -> Router {
     echo::router().nest(
@@ -13,3 +13,10 @@ pub fn mpc_router(transport: MpcHttpTransport) -> Router {
             .merge(query::h2h_router(transport)),
     )
 }
+
+pub fn shard_router(transport: ShardHttpTransport) -> Router {
+    echo::router().nest(
+        http_serde::query::BASE_AXUM_PATH,
+        Router::new().merge(query::s2s_router(transport)),
+    )
+}

ipa-core/src/net/server/handlers/query/mod.rs

@@ -6,6 +6,8 @@ mod results;
 mod status;
 mod step;
 
+use std::marker::PhantomData;
+
 use axum::{
     response::{IntoResponse, Response},
     Router,
@@ -17,9 +19,9 @@ use futures_util::{
 use hyper::{Request, StatusCode};
 use tower::{layer::layer_fn, Service};
 
-use crate::{
-    helpers::HelperIdentity,
-    net::{server::ClientIdentity, transport::MpcHttpTransport},
+use crate::net::{
+    server::ClientIdentity, transport::MpcHttpTransport, ConnectionFlavor, Helper, Shard,
+    ShardHttpTransport,
 };
 
 /// Construct router for IPA query web service
@@ -45,9 +47,16 @@ pub fn query_router(transport: MpcHttpTransport) -> Router {
 // It might make sense to split the query and h2h handlers into two modules.
 pub fn h2h_router(transport: MpcHttpTransport) -> Router {
     Router::new()
-        .merge(prepare::router(transport.clone()))
-        .merge(step::router(transport))
-        .layer(layer_fn(HelperAuthentication::new))
+        .merge(step::router(transport.clone()))
+        .merge(prepare::router(transport.inner_transport))
+        .layer(layer_fn(HelperAuthentication::<_, Helper>::new))
+}
+
+/// Construct router for shard-to-shard communications similar to [`h2h_router`].
+pub fn s2s_router(transport: ShardHttpTransport) -> Router {
+    Router::new()
+        .merge(prepare::router(transport.inner_transport))
+        .layer(layer_fn(HelperAuthentication::<_, Shard>::new))
 }
 
 /// Returns HTTP 401 Unauthorized if the request does not have valid authentication.
@@ -63,18 +72,24 @@ pub fn h2h_router(transport: MpcHttpTransport) -> Router {
 /// requests would not have this request extension, causing axum to fail the request with
 /// `ExtensionRejection::MissingExtension`, however, this would return a 500 error instead of 401.
 #[derive(Clone)]
-pub struct HelperAuthentication<S> {
+pub struct HelperAuthentication<S, F: ConnectionFlavor> {
     inner: S,
+    flavor: PhantomData<F>,
 }
 
-impl<S> HelperAuthentication<S> {
+impl<S, F: ConnectionFlavor> HelperAuthentication<S, F> {
     fn new(inner: S) -> Self {
-        Self { inner }
+        Self {
+            inner,
+            flavor: PhantomData,
+        }
     }
 }
 
-impl<B, S: Service<Request<B>, Response = Response>> Service<Request<B>>
-    for HelperAuthentication<S>
+impl<B, S, F> Service<Request<B>> for HelperAuthentication<S, F>
+where
+    S: Service<Request<B>, Response = Response>,
+    F: ConnectionFlavor,
 {
     type Response = Response;
     type Error = S::Error;
@@ -88,7 +103,7 @@ impl<B, S: Service<Request<B>, Response = Response>> Service<Request<B>>
     }
 
     fn call(&mut self, req: Request<B>) -> Self::Future {
-        match req.extensions().get::<ClientIdentity<HelperIdentity>>() {
+        match req.extensions().get::<ClientIdentity<F::Identity>>() {
             Some(ClientIdentity(_)) => self.inner.call(req).left_future(),
             None => ready(Ok((
                 StatusCode::UNAUTHORIZED,

ipa-core/src/net/server/handlers/query/prepare.rs

@@ -1,26 +1,28 @@
+use std::sync::Arc;
+
 use axum::{extract::Path, response::IntoResponse, routing::post, Extension, Json, Router};
 use hyper::StatusCode;
 
 use crate::{
-    helpers::{query::PrepareQuery, BodyStream, HelperIdentity},
+    helpers::{query::PrepareQuery, BodyStream},
     net::{
         http_serde::{
             self,
             query::{prepare::RequestBody, QueryConfigQueryParams},
         },
         server::ClientIdentity,
-        transport::MpcHttpTransport,
-        Error,
+        transport::HttpTransport,
+        ConnectionFlavor, Error,
     },
     protocol::QueryId,
     query::PrepareQueryError,
 };
 
 /// Called by whichever peer helper is the leader for an individual query, to initiatialize
 /// processing of that query.
-async fn handler(
-    transport: Extension<MpcHttpTransport>,
-    _: Extension<ClientIdentity<HelperIdentity>>, // require that client is an authenticated helper
+async fn handler<F: ConnectionFlavor>(
+    transport: Extension<Arc<HttpTransport<F>>>,
+    _: Extension<ClientIdentity<F::Identity>>, // require that client is an authenticated helper
     Path(query_id): Path<QueryId>,
     QueryConfigQueryParams(config): QueryConfigQueryParams,
     Json(RequestBody { roles }): Json<RequestBody>,
@@ -30,7 +32,7 @@ async fn handler(
         config,
         roles,
     };
-    let _ = transport
+    let _ = Arc::clone(&transport)
         .dispatch(data, BodyStream::empty())
         .await
         .map_err(|e| Error::application(StatusCode::INTERNAL_SERVER_ERROR, e))?;
@@ -44,9 +46,9 @@ impl IntoResponse for PrepareQueryError {
     }
 }
 
-pub fn router(transport: MpcHttpTransport) -> Router {
+pub fn router<F: ConnectionFlavor>(transport: Arc<HttpTransport<F>>) -> Router {
     Router::new()
-        .route(http_serde::query::prepare::AXUM_PATH, post(handler))
+        .route(http_serde::query::prepare::AXUM_PATH, post(handler::<F>))
         .layer(Extension(transport))
 }
 

ipa-core/src/net/server/mod.rs

@@ -111,14 +111,15 @@ impl IpaHttpServer<Helper> {
 impl IpaHttpServer<Shard> {
     #[must_use]
     pub fn new_shards(
-        _transport: &ShardHttpTransport,
+        transport: &ShardHttpTransport,
         config: ServerConfig,
         network_config: NetworkConfig<Shard>,
     ) -> Self {
+        let router = handlers::shard_router(transport.clone());
         IpaHttpServer {
             config,
             network_config,
-            router: Router::new(),
+            router,
         }
     }
 }

ipa-core/src/net/transport.rs

@@ -38,13 +38,13 @@ pub struct HttpTransport<F: ConnectionFlavor> {
 /// HTTP transport for helper to helper traffic.
 #[derive(Clone)]
 pub struct MpcHttpTransport {
-    inner_transport: Arc<HttpTransport<Helper>>,
+    pub(super) inner_transport: Arc<HttpTransport<Helper>>,
 }
 
 /// A stub for HTTP transport implementation, suitable for serving shard-to-shard traffic
 #[derive(Clone)]
 pub struct ShardHttpTransport {
-    inner_transport: Arc<HttpTransport<Shard>>,
+    pub(super) inner_transport: Arc<HttpTransport<Shard>>,
 }
 
 impl RouteParams<RouteId, NoQueryId, NoStep> for QueryConfig {