Get multiplexing working

webb-gadget/src/lib.rs

@@ -22,6 +22,8 @@ pub enum Error {
     GadgetManagerError { err: GadgetError },
     InitError { err: String },
     WorkManagerError { err: WorkManagerError },
+    ProtocolRemoteError { err: String },
+    ClientError { err: String },
 }
 
 impl Display for Error {

zk-gadget/Cargo.toml

@@ -18,7 +18,8 @@ async-trait = { workspace = true }
 parking_lot = { workspace = true }
 log = { workspace = true }
 hex = { workspace = true }
-bytes = { workspace = true }
+bytes = { workspace = true, features = ["serde"] }
+sp-runtime = { workspace = true }
 
 [dev-dependencies]
 tonic = "0.10.2"

zk-gadget/src/client_ext.rs

@@ -0,0 +1,37 @@
+use crate::client_ext::job_types::{JobCircuitProperties, JobProperties};
+use async_trait::async_trait;
+use gadget_core::gadget::substrate::Client;
+use sp_runtime::traits::Block;
+
+#[async_trait]
+pub trait ClientWithApi<B: Block>: Client<B> + Clone {
+    async fn get_job_circuit_properties(
+        &self,
+        job_id: u64,
+    ) -> Result<JobCircuitProperties, webb_gadget::Error>;
+    async fn get_job_properties(&self, job_id: u64) -> Result<JobProperties, webb_gadget::Error>;
+}
+
+pub mod job_types {
+    pub struct ZkJob {
+        pub circuit: JobCircuitProperties,
+        pub properties: JobProperties,
+    }
+
+    pub struct JobCircuitProperties {
+        pub job_id: u64,
+        pub pk: Vec<u8>,
+        pub wasm_uri: String,
+        pub r1cs_uri: String,
+    }
+
+    pub struct JobProperties {
+        pub job_id: u64,
+        pub circuit_id: u64,
+        pub public_inputs: Vec<u8>,
+        pub pss: u64,
+        pub a_shares: Vec<u8>,
+        pub ax_shares: Vec<u8>,
+        pub qap_shares: Vec<u8>,
+    }
+}

zk-gadget/src/lib.rs

@@ -1,6 +1,7 @@
-use crate::module::ZkModule;
+use crate::client_ext::ClientWithApi;
+use crate::module::proto_gen::AsyncProtocolGenerator;
+use crate::module::{AdditionalProtocolParams, ZkModule};
 use crate::network::{RegistantId, ZkNetworkService};
-use gadget_core::gadget::substrate::Client;
 use gadget_core::Block;
 use mpc_net::prod::RustlsCertificate;
 use std::net::SocketAddr;
@@ -10,6 +11,8 @@ use webb_gadget::Error;
 pub mod module;
 pub mod network;
 
+pub mod client_ext;
+
 pub struct ZkGadgetConfig {
     pub king_bind_addr: Option<SocketAddr>,
     pub client_only_king_addr: Option<SocketAddr>,
@@ -19,14 +22,28 @@ pub struct ZkGadgetConfig {
     pub client_only_king_public_identity_der: Option<Vec<u8>>,
 }
 
-pub async fn run<C: Client<B>, B: Block>(config: ZkGadgetConfig, client: C) -> Result<(), Error> {
+pub async fn run<
+    C: ClientWithApi<B>,
+    B: Block,
+    T: AdditionalProtocolParams,
+    Gen: AsyncProtocolGenerator<T, Error, ZkNetworkService>,
+>(
+    config: ZkGadgetConfig,
+    client: C,
+    additional_protocol_params: T,
+    async_protocol_generator: Gen,
+) -> Result<(), Error> {
     // Create the zk gadget module
     let network = create_zk_network(&config).await?;
     log::info!("Created zk network for party {}", config.id);
 
     let zk_module = ZkModule {
         party_id: config.id,
-    }; // TODO: proper implementation
+        additional_protocol_params,
+        network: network.clone(),
+        client: client.clone(),
+        async_protocol_generator: Box::new(async_protocol_generator),
+    };
 
     // Plug the module into the webb gadget
     webb_gadget::run(network, zk_module, client).await

zk-gadget/src/module/mod.rs

@@ -1,4 +1,5 @@
-use crate::network::RegistantId;
+use crate::client_ext::ClientWithApi;
+use crate::network::{RegistantId, ZkNetworkService};
 use async_trait::async_trait;
 use gadget_core::job_manager::ProtocolWorkManager;
 use gadget_core::Block;
@@ -8,12 +9,22 @@ use webb_gadget::{BlockImportNotification, Error, FinalityNotification};
 
 pub mod proto_gen;
 
-pub struct ZkModule {
+pub struct ZkModule<T, C> {
     pub party_id: RegistantId,
+    pub additional_protocol_params: T,
+    pub client: C,
+    pub network: ZkNetworkService,
+    pub async_protocol_generator:
+        Box<dyn proto_gen::AsyncProtocolGenerator<T, Error, ZkNetworkService>>,
 }
 
+pub trait AdditionalProtocolParams: Send + Sync + Clone + 'static {}
+impl<T: Send + Sync + Clone + 'static> AdditionalProtocolParams for T {}
+
 #[async_trait]
-impl<B: Block> WebbGadgetModule<B> for ZkModule {
+impl<B: Block, T: AdditionalProtocolParams, C: ClientWithApi<B>> WebbGadgetModule<B>
+    for ZkModule<T, C>
+{
     async fn process_finality_notification(
         &self,
         _notification: FinalityNotification<B>,
@@ -24,8 +35,29 @@ impl<B: Block> WebbGadgetModule<B> for ZkModule {
             "Party {} received a finality notification at {now}",
             self.party_id
         );
+
         // TODO: call proto_gen::create_zk_async_protocol to generate a protocol, then
-        // push the returned remote and protocol to the job manager
+        // push the returned remote and protocol to the job manager. E.g.,:
+        let task_id = [0u8; 32];
+        let n_parties = 5; // TODO: get this from somewhere
+        let (remote, protocol) = proto_gen::create_zk_async_protocol(
+            0,
+            now,
+            0,
+            task_id,
+            self.party_id,
+            n_parties,
+            self.additional_protocol_params.clone(),
+            self.network.clone(),
+            &*self.async_protocol_generator,
+        );
+
+        if let Err(err) =
+            _job_manager.push_task(task_id, false, std::sync::Arc::new(remote), protocol)
+        {
+            log::error!("Failed to push task to job manager: {err:?}");
+        }
+
         Ok(())
     }
 
@@ -39,9 +71,9 @@ impl<B: Block> WebbGadgetModule<B> for ZkModule {
 
     async fn process_error(
         &self,
-        _error: Error,
+        error: Error,
         _job_manager: &ProtocolWorkManager<WebbWorkManager>,
     ) {
-        todo!()
+        log::error!("Party {} received an error: {:?}", self.party_id, error);
     }
 }