Repository: tangle-network/gadget (language: rust)

Issue:
Add pre-job and post-job hooks as an abstraction to Jobs submitted to the job manager
Many job types we use have a format: perform some checks to see if we need to run a job, and if so, run the job. The final pattern is when the job finishes, we want some finalization logic (e.g., for the DKG, we would submit the public key on-chain).

Currently, we require that jobs satisfy `T: SendFuture<'static, ()>`. Under the new model, we would have a trait:

```rust
pub enum ProceedWithExecution {
    True,
    False
}

#[async_trait]
pub trait ExecutableJob: SendFuture<'static, ()> {
    async fn pre_job_hook(&mut self) -> Result<ProceedWithExecution, Error>;
    async fn post_job_hook(&mut self) -> Result<(), Error>;
}
```

Thus, jobs that are submitted must satisfy:

```rust
fn push_job<J: ExecutableJob>(&self, job: J)
```


Discussion / hints:

Fix patch:
diff --git a/.dockerignore b/.dockerignore
new file mode 100644
index 000000000..2aaa48c80
--- /dev/null
+++ b/.dockerignore
@@ -0,0 +1,2 @@
+target/release
+target/debug
diff --git a/.gitattributes b/.gitattributes
new file mode 100644
index 000000000..2e2d98364
--- /dev/null
+++ b/.gitattributes
@@ -0,0 +1,2 @@
+target/zk/job_1.json filter=lfs diff=lfs merge=lfs -text
+target/zk/proving_key.bin filter=lfs diff=lfs merge=lfs -text
diff --git a/.github/workflows/validate_pr.yml b/.github/workflows/validate_pr.yml
index ed3fb4e63..41b4319de 100644
--- a/.github/workflows/validate_pr.yml
+++ b/.github/workflows/validate_pr.yml
@@ -9,6 +9,21 @@ concurrency:
   cancel-in-progress: true
 
 jobs:
+  integration-tests:
+    name: Integration Tests
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v2
+
+      - name: Install protobuf
+        run: sudo apt-get install protobuf-compiler
+
+      - name: Generate certs
+        run: bash ./playground/generate_certs.sh
+
+      - name: Run ZK MVP
+        run: docker-compose up --abort-on-container-exit
   formatting:
     name: Rustfmt
     runs-on: ubuntu-latest
diff --git a/README.md b/README.md
index 208e57ff8..f83e0110f 100644
--- a/README.md
+++ b/README.md
@@ -10,3 +10,10 @@ The core library is `gadget-core`. The core library allows gadgets to hold stand
 Once such blockchain is a substrate blockchain. This is where `webb-gadget` comes into play. The `webb-gadget` is the `core-gadget` endowed with a connection to a substrate blockchain, a networking layer to communicate with other gadgets, and a `WebbGadgetModule` that has application-specific logic. 
 
 Since `webb-gadget` allows varying connections to a substrate blockchain and differing network layers, we can thus design above it the `zk-gadget` and `tss-gadget`. These gadgets are endowed with the same functionalities as the `webb-gadget` but with a (potentially) different blockchain connection, networking layer, and application-specific logic.
+
+## Testing
+
+### ZK-Gadget
+To test the ZK-gadget, you can run `./playground/run.sh --create-certs`. Additionally, you may use docker-compose to set up a testnet automatically, running a ZK test protocol in the process:
+
+`docker-compose up --abort-on-container-exit`
\ No newline at end of file
diff --git a/docker-compose.yml b/docker-compose.yml
new file mode 100644
index 000000000..69b1c01c0
--- /dev/null
+++ b/docker-compose.yml
@@ -0,0 +1,30 @@
+version: '3.8'
+services:
+  king:
+    build:
+      context: .
+      dockerfile: ./docker/server/Dockerfile
+    environment:
+      - BIND_ADDRESS=0.0.0.0:5555
+    ports:
+      - 5555:5555
+    networks:
+      zk_network:
+        ipv4_address: 172.40.0.10
+
+  clients:
+    build:
+      context: .
+      dockerfile: ./docker/client/Dockerfile
+    environment:
+      - SERVER_IP=172.40.0.10
+      - SERVER_PORT=5555
+    networks:
+      - zk_network
+
+networks:
+  zk_network:
+    driver: bridge
+    ipam:
+      config:
+        - subnet: 172.40.0.0/16
\ No newline at end of file
diff --git a/docker/client/Dockerfile b/docker/client/Dockerfile
new file mode 100644
index 000000000..7992970f9
--- /dev/null
+++ b/docker/client/Dockerfile
@@ -0,0 +1,8 @@
+FROM rust:latest
+
+WORKDIR /usr/src/client
+COPY . .
+RUN apt-get update && apt-get install -y protobuf-compiler netcat-traditional
+RUN cargo build --release --example zknode
+RUN cargo run --release --example zk-test-client -- --circuit-id 1 --job-id 1 --wasm ./fixtures/sha256/circom.wasm --r1cs ./fixtures/sha256/circom.r1cs --input ./fixtures/sha256/input.json --public-inputs ./fixtures/sha256/public_inputs.json --output-dir ./target/zk --generate-proving-key
+ENTRYPOINT ["./docker/client/client.sh"]
\ No newline at end of file
diff --git a/docker/client/client.sh b/docker/client/client.sh
new file mode 100755
index 000000000..29b22a115
--- /dev/null
+++ b/docker/client/client.sh
@@ -0,0 +1,26 @@
+#!/usr/bin/env bash
+set -e
+trap "exit" INT TERM
+trap "kill 0" EXIT
+
+echo "Server addr: $SERVER_IP:$SERVER_PORT"
+
+# Retry loop
+while ! nc -z -w 5 $SERVER_IP $SERVER_PORT; do
+  echo "Waiting for server $SERVER_IP:$SERVER_PORT ... "
+  sleep 0.5
+done
+
+PROCS=()
+n=8
+for i in {1..7}
+  do
+    RUST_LOG=debug cargo run --release --example zknode -- --n $n --king-ip $SERVER_IP:$SERVER_PORT --public-identity-der certs/$i/cert.der --client-only-king-public-identity-der certs/0/cert.der --watch-dir target/zk --name node$i --i $i --private-identity-der certs/$i/key.der &
+    pid=$!
+    PROCS[$i]=$pid
+  done
+
+# Step 5: wait for all processes to finish
+for pid in ${PROCS[@]}; do
+    wait $pid
+done
\ No newline at end of file
diff --git a/docker/server/Dockerfile b/docker/server/Dockerfile
new file mode 100644
index 000000000..a83a3814a
--- /dev/null
+++ b/docker/server/Dockerfile
@@ -0,0 +1,8 @@
+FROM rust:latest
+
+WORKDIR /usr/src/server
+COPY . .
+RUN apt-get update && apt-get install -y protobuf-compiler
+RUN cargo build --release --example zknode
+RUN cargo run --release --example zk-test-client -- --circuit-id 1 --job-id 1 --wasm ./fixtures/sha256/circom.wasm --r1cs ./fixtures/sha256/circom.r1cs --input ./fixtures/sha256/input.json --public-inputs ./fixtures/sha256/public_inputs.json --output-dir ./target/zk --generate-proving-key
+CMD RUST_LOG=debug cargo run --release --example zknode -- --n 8 --king-ip $BIND_ADDRESS --public-identity-der certs/0/cert.der --watch-dir target/zk --name king --i 0 --private-identity-der certs/0/key.der
\ No newline at end of file
diff --git a/gadget-core/Cargo.toml b/gadget-core/Cargo.toml
index e8d8ab59e..f4479ffb4 100644
--- a/gadget-core/Cargo.toml
+++ b/gadget-core/Cargo.toml
@@ -13,7 +13,6 @@ substrate = [
 ]
 
 [dependencies]
-sync_wrapper = "0.1.2"
 parking_lot = { workspace = true }
 tokio = { workspace = true, features = ["sync", "time", "macros", "rt"] }
 hex = "0.4.3"
@@ -23,4 +22,5 @@ sp-runtime = { optional = true, workspace = true, default-features = false }
 sc-client-api = { optional = true, workspace = true, default-features = false }
 futures = { optional = true, workspace = true }
 
-[dev-dependencies]
\ No newline at end of file
+[dev-dependencies]
+tokio = { workspace = true, features = ["macros"]}
\ No newline at end of file
diff --git a/gadget-core/src/job.rs b/gadget-core/src/job.rs
new file mode 100644
index 000000000..29a11b23c
--- /dev/null
+++ b/gadget-core/src/job.rs
@@ -0,0 +1,324 @@
+use crate::job_manager::SendFuture;
+use async_trait::async_trait;
+use std::error::Error;
+use std::fmt::Display;
+use std::future::Future;
+use std::pin::Pin;
+use std::task::{Context, Poll};
+
+pub enum ProceedWithExecution {
+    True,
+    False,
+}
+
+#[derive(Debug)]
+pub struct JobError {
+    pub reason: String,
+}
+
+impl<T: Into<String>> From<T> for JobError {
+    fn from(value: T) -> Self {
+        Self {
+            reason: value.into(),
+        }
+    }
+}
+
+impl Display for JobError {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{reason}", reason = self.reason)
+    }
+}
+
+impl Error for JobError {}
+
+#[async_trait]
+pub trait ExecutableJob: SendFuture<'static, Result<(), JobError>> + Unpin {
+    async fn pre_job_hook(&mut self) -> Result<ProceedWithExecution, JobError>;
+    async fn post_job_hook(&mut self) -> Result<(), JobError>;
+    async fn execute(&mut self) -> Result<(), JobError> {
+        match self.pre_job_hook().await? {
+            ProceedWithExecution::True => {
+                let result = (&mut self).await;
+                let post_result = self.post_job_hook().await;
+                result.and(post_result)
+            }
+            ProceedWithExecution::False => Ok(()),
+        }
+    }
+}
+
+pub struct ExecutableJobWrapper<Pre: ?Sized, Protocol, Post: ?Sized> {
+    pre: Pin<Box<Pre>>,
+    protocol: Pin<Box<Protocol>>,
+    post: Pin<Box<Post>>,
+}
+
+#[async_trait]
+impl<Pre: ?Sized, Protocol, Post: ?Sized> ExecutableJob
+    for ExecutableJobWrapper<Pre, Protocol, Post>
+where
+    Pre: Future<Output = Result<ProceedWithExecution, JobError>> + Send + 'static,
+    Protocol: SendFuture<'static, Result<(), JobError>>,
+    Post: Future<Output = Result<(), JobError>> + Send + 'static,
+{
+    async fn pre_job_hook(&mut self) -> Result<ProceedWithExecution, JobError> {
+        self.pre.as_mut().await
+    }
+
+    async fn post_job_hook(&mut self) -> Result<(), JobError> {
+        self.post.as_mut().await
+    }
+}
+
+impl<Pre, Protocol, Post> ExecutableJobWrapper<Pre, Protocol, Post>
+where
+    Pre: Future<Output = Result<ProceedWithExecution, JobError>>,
+    Protocol: SendFuture<'static, Result<(), JobError>>,
+    Post: Future<Output = Result<(), JobError>>,
+{
+    pub fn new(pre: Pre, protocol: Protocol, post: Post) -> Self {
+        Self {
+            pre: Box::pin(pre),
+            protocol: Box::pin(protocol),
+            post: Box::pin(post),
+        }
+    }
+}
+
+impl<Pre: ?Sized, Protocol, Post: ?Sized> Future for ExecutableJobWrapper<Pre, Protocol, Post>
+where
+    Protocol: SendFuture<'static, Result<(), JobError>>,
+{
+    type Output = <Protocol as Future>::Output;
+
+    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        self.protocol.as_mut().poll(cx)
+    }
+}
+
+#[derive(Default)]
+pub struct JobBuilder {
+    pre: Option<Pin<Box<PreJobHook>>>,
+    post: Option<Pin<Box<PostJobHook>>>,
+}
+
+pub type PreJobHook = dyn SendFuture<'static, Result<ProceedWithExecution, JobError>>;
+pub type PostJobHook = dyn SendFuture<'static, Result<(), JobError>>;
+
+pub struct DefaultPreJobHook;
+impl Future for DefaultPreJobHook {
+    type Output = Result<ProceedWithExecution, JobError>;
+
+    fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Self::Output> {
+        Poll::Ready(Ok(ProceedWithExecution::True))
+    }
+}
+
+pub struct DefaultPostJobHook;
+impl Future for DefaultPostJobHook {
+    type Output = Result<(), JobError>;
+
+    fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Self::Output> {
+        Poll::Ready(Ok(()))
+    }
+}
+
+pub type BuiltExecutableJobWrapper<Protocol> = ExecutableJobWrapper<
+    dyn SendFuture<'static, Result<ProceedWithExecution, JobError>>,
+    Protocol,
+    dyn SendFuture<'static, Result<(), JobError>>,
+>;
+
+impl JobBuilder {
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    pub fn pre<Pre>(mut self, pre: Pre) -> Self
+    where
+        Pre: SendFuture<'static, Result<ProceedWithExecution, JobError>>,
+    {
+        self.pre = Some(Box::pin(pre));
+        self
+    }
+
+    pub fn post<Post>(mut self, post: Post) -> Self
+    where
+        Post: SendFuture<'static, Result<(), JobError>>,
+    {
+        self.post = Some(Box::pin(post));
+        self
+    }
+
+    pub fn build<Protocol>(self, protocol: Protocol) -> BuiltExecutableJobWrapper<Protocol>
+    where
+        Protocol: SendFuture<'static, Result<(), JobError>>,
+    {
+        let pre = if let Some(pre) = self.pre {
+            pre
+        } else {
+            Box::pin(DefaultPreJobHook)
+        };
+
+        let post = if let Some(post) = self.post {
+            post
+        } else {
+            Box::pin(DefaultPostJobHook)
+        };
+
+        ExecutableJobWrapper {
+            pre,
+            protocol: Box::pin(protocol),
+            post,
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::job::ExecutableJob;
+
+    #[tokio::test]
+    async fn test_executable_job_wrapper_proceed() {
+        let counter = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0));
+        let counter_clone = counter.clone();
+        let counter_clone2 = counter.clone();
+        let counter_final = counter.clone();
+
+        let pre = async move {
+            counter.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
+            Ok(super::ProceedWithExecution::True)
+        };
+
+        let protocol = async move {
+            counter_clone.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
+            Ok(())
+        };
+
+        let post = async move {
+            counter_clone2.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
+            Ok(())
+        };
+
+        let mut job = super::ExecutableJobWrapper::new(pre, protocol, post);
+        job.execute().await.unwrap();
+        assert_eq!(counter_final.load(std::sync::atomic::Ordering::SeqCst), 3);
+    }
+
+    #[tokio::test]
+    async fn test_executable_job_wrapper_no_proceed() {
+        let counter = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0));
+        let counter_clone = counter.clone();
+        let counter_clone2 = counter.clone();
+        let counter_final = counter.clone();
+
+        let pre = async move {
+            counter.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
+            Ok(super::ProceedWithExecution::False)
+        };
+
+        let protocol = async move {
+            counter_clone.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
+            Ok(())
+        };
+
+        let post = async move {
+            counter_clone2.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
+            Ok(())
+        };
+
+        let mut job = super::ExecutableJobWrapper::new(pre, protocol, post);
+        job.execute().await.unwrap();
+        assert_eq!(counter_final.load(std::sync::atomic::Ordering::SeqCst), 1);
+    }
+
+    #[tokio::test]
+    async fn test_job_builder() {
+        let counter = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0));
+        let counter_clone = counter.clone();
+        let counter_clone2 = counter.clone();
+        let counter_final = counter.clone();
+
+        let protocol = async move {
+            counter_clone.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
+            Ok(())
+        };
+
+        let mut job = super::JobBuilder::new()
+            .pre(async move {
+                counter.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
+                Ok(super::ProceedWithExecution::True)
+            })
+            .post(async move {
+                counter_clone2.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
+                Ok(())
+            })
+            .build(protocol);
+
+        job.execute().await.unwrap();
+        assert_eq!(counter_final.load(std::sync::atomic::Ordering::SeqCst), 3);
+    }
+
+    #[tokio::test]
+    async fn test_job_builder_no_pre() {
+        let counter = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0));
+        let counter_clone = counter.clone();
+        let counter_clone2 = counter.clone();
+        let counter_final = counter.clone();
+
+        let protocol = async move {
+            counter_clone.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
+            Ok(())
+        };
+
+        let mut job = super::JobBuilder::default()
+            .post(async move {
+                counter_clone2.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
+                Ok(())
+            })
+            .build(protocol);
+
+        job.execute().await.unwrap();
+        assert_eq!(counter_final.load(std::sync::atomic::Ordering::SeqCst), 2);
+    }
+
+    #[tokio::test]
+    async fn test_job_builder_no_post() {
+        let counter = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0));
+        let counter_clone = counter.clone();
+        let counter_final = counter.clone();
+
+        let protocol = async move {
+            counter_clone.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
+            Ok(())
+        };
+
+        let mut job = super::JobBuilder::default()
+            .pre(async move {
+                counter.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
+                Ok(super::ProceedWithExecution::True)
+            })
+            .build(protocol);
+
+        job.execute().await.unwrap();
+        assert_eq!(counter_final.load(std::sync::atomic::Ordering::SeqCst), 2);
+    }
+
+    #[tokio::test]
+    async fn test_job_builder_no_pre_no_post() {
+        let counter = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0));
+        let counter_clone = counter.clone();
+        let counter_final = counter.clone();
+
+        let protocol = async move {
+            counter_clone.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
+            Ok(())
+        };
+
+        let mut job = super::JobBuilder::default().build(protocol);
+
+        job.execute().await.unwrap();
+        assert_eq!(counter_final.load(std::sync::atomic::Ordering::SeqCst), 1);
+    }
+}
diff --git a/gadget-core/src/job_manager.rs b/gadget-core/src/job_manager.rs
index 436338a96..89dca5ec7 100644
--- a/gadget-core/src/job_manager.rs
+++ b/gadget-core/src/job_manager.rs
@@ -1,3 +1,4 @@
+use crate::job::ExecutableJob;
 use parking_lot::RwLock;
 use std::fmt::{Debug, Display};
 use std::future::Future;
@@ -5,10 +6,8 @@ use std::ops::{Add, Sub};
 use std::{
     collections::{HashMap, HashSet, VecDeque},
     hash::{Hash, Hasher},
-    pin::Pin,
     sync::Arc,
 };
-use sync_wrapper::SyncWrapper;
 
 #[derive(Copy, Clone, PartialEq, Eq)]
 pub enum PollMethod {
@@ -157,10 +156,9 @@ impl<WM: WorkManagerInterface> ProtocolWorkManager<WM> {
 
         if let PollMethod::Interval { millis } = poll_method {
             let this_worker = this.clone();
-            let handler = async move {
-                let job_receiver_worker = this_worker.clone();
-                let logger = job_receiver_worker.utility.clone();
+            let logger = this_worker.utility.clone();
 
+            let handler = async move {
                 let periodic_poller = async move {
                     let mut interval =
                         tokio::time::interval(std::time::Duration::from_millis(millis));
@@ -186,18 +184,18 @@ impl<WM: WorkManagerInterface> ProtocolWorkManager<WM> {
     }
 
     /// Pushes the task, but does not necessarily start it
-    pub fn push_task(
+    pub fn push_task<T: ExecutableJob>(
         &self,
         task_hash: WM::TaskID,
         force_start: bool,
         handle: Arc<dyn ProtocolRemote<WM>>,
-        task: Pin<Box<dyn SendFuture<'static, ()>>>,
+        task: T,
     ) -> Result<(), WorkManagerError> {
         let mut lock = self.inner.write();
         // set as primary, that way on drop, the async protocol ends
         handle.set_as_primary();
         let job = Job {
-            task: Arc::new(RwLock::new(Some(task.into()))),
+            task: Arc::new(parking_lot::Mutex::new(Some(Box::new(task)))),
             handle,
             task_hash,
             utility: self.utility.clone(),
@@ -394,11 +392,8 @@ impl<WM: WorkManagerInterface> ProtocolWorkManager<WM> {
         // Put the job inside here, that way the drop code does not get called right away,
         // killing the process
         lock.active_tasks.insert(job);
-        // run the task
-        let task = async move {
-            let task = task.write().take().expect("Should not happen");
-            task.into_inner().await
-        };
+        let mut task = task.lock().take().expect("Should not happen");
+        let task = async move { task.execute().await };
 
         // Spawn the task. When it finishes, it will clean itself up
         tokio::task::spawn(task);
@@ -485,7 +480,7 @@ pub struct Job<WM: WorkManagerInterface> {
     task_hash: WM::TaskID,
     utility: Arc<WM>,
     handle: Arc<dyn ProtocolRemote<WM>>,
-    task: Arc<RwLock<Option<SyncFuture<()>>>>,
+    task: Arc<parking_lot::Mutex<Option<Box<dyn ExecutableJob>>>>,
 }
 
 impl<WM: WorkManagerInterface> Job<WM> {
@@ -515,8 +510,6 @@ pub enum ShutdownReason {
 pub trait SendFuture<'a, T>: Send + Future<Output = T> + 'a {}
 impl<'a, F: Send + Future<Output = T> + 'a, T> SendFuture<'a, T> for F {}
 
-pub type SyncFuture<T> = SyncWrapper<Pin<Box<dyn SendFuture<'static, T>>>>;
-
 impl<WM: WorkManagerInterface> PartialEq for Job<WM> {
     fn eq(&self, other: &Self) -> bool {
         self.task_hash == other.task_hash
@@ -557,8 +550,8 @@ fn should_deliver<WM: WorkManagerInterface>(
 
 #[cfg(test)]
 mod tests {
-
     use super::*;
+    use crate::job::{BuiltExecutableJobWrapper, JobBuilder, JobError};
     use parking_lot::Mutex;
     use std::sync::atomic::{AtomicBool, Ordering};
     use std::sync::Arc;
@@ -654,7 +647,7 @@ mod tests {
         started_at: u64,
     ) -> (
         Arc<TestProtocolRemote>,
-        Pin<Box<dyn SendFuture<'static, ()>>>,
+        BuiltExecutableJobWrapper<impl SendFuture<'static, Result<(), JobError>>>,
         UnboundedReceiver<TestMessage>,
     ) {
         let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
@@ -662,11 +655,12 @@ mod tests {
         let is_done = Arc::new(AtomicBool::new(false));
         let remote =
             TestProtocolRemote::new(session_id, ssid, started_at, tx, start_tx, is_done.clone());
-        let task = async move {
+        let protocol = async move {
             start_rx.await.unwrap();
             is_done.store(true, Ordering::SeqCst);
+            Ok(())
         };
-        let task = Box::pin(task);
+        let task = JobBuilder::new().build(protocol);
         (remote, task, rx)
     }
 
diff --git a/gadget-core/src/lib.rs b/gadget-core/src/lib.rs
index f0906143e..9484e782b 100644
--- a/gadget-core/src/lib.rs
+++ b/gadget-core/src/lib.rs
@@ -1,5 +1,4 @@
-pub use sc_client_api::Backend;
-pub use sp_runtime::traits::Block;
-
 pub mod gadget;
 pub mod job_manager;
+
+pub mod job;
diff --git a/playground/Cargo.toml b/playground/Cargo.toml
index ff82dbd59..a5998a282 100644
--- a/playground/Cargo.toml
+++ b/playground/Cargo.toml
@@ -5,17 +5,17 @@ authors = ["Shady Khalifa <dev+github@shadykhalifa.me>"]
 license = "GPL-3.0-or-later WITH Classpath-exception-2.0"
 edition = "2021"
 
-[[bin]]
-name = "client"
-path = "./client.rs"
+[[example]]
+name = "zk-test-client"
+path = "examples/client.rs"
 
-[[bin]]
+[[example]]
 name = "zknode"
-path = "./node.rs"
+path = "examples/node.rs"
 
-[[bin]]
+[[example]]
 name = "certgen"
-path = "./certgen.rs"
+path = "examples/certgen.rs"
 
 [features]
 default = ["parallel"]
diff --git a/playground/README.md b/playground/README.md
index a8d9eea51..bd136c55f 100644
--- a/playground/README.md
+++ b/playground/README.md
@@ -6,7 +6,7 @@
 ```
 2. Generate the shares by running the client:
 ```sh
-cargo run --release --bin client -- --circuit-id 1 --job-id 1 --wasm ./fixtures/sha256/circom.wasm --r1cs ./fixtures/sha256/circom.r1cs --input ./fixtures/sha256/input.json --public-inputs ./fixtures/sha256/public_inputs.json --output-dir ./target/zk --generate-proving-key
+cargo run --release --example zk-test-client -- --circuit-id 1 --job-id 1 --wasm ./fixtures/sha256/circom.wasm --r1cs ./fixtures/sha256/circom.r1cs --input ./fixtures/sha256/input.json --public-inputs ./fixtures/sha256/public_inputs.json --output-dir ./target/zk --generate-proving-key
 ```
 You can change the `--circuit-id` and `--job-id` to any number you want. The `--output-dir` is where the shares will be generated. You can change the `--wasm`, `--r1cs`, `--input` and `--public-inputs` to any circuit you want to run.
 Note: `--generate-proving-key` will generate the proving key. If you already have the proving key, you can omit this flag.
@@ -17,15 +17,8 @@ You will need 8 terminals to run the zknode(s). In each terminal, run:
 **King**
 
 ```sh
-RUST_LOG=debug cargo run --release --bin zknode -- --king-ip 127.0.0.1:5555 \
-    --certs certs/0/cert.der \
-    --certs certs/1/cert.der \
-    --certs certs/2/cert.der \
-    --certs certs/3/cert.der \
-    --certs certs/4/cert.der \
-    --certs certs/5/cert.der \
-    --certs certs/6/cert.der \
-    --certs certs/7/cert.der \
+RUST_LOG=debug cargo run --release --example zknode -- --king-ip 127.0.0.1:5555 \
+    --public_identity_der certs/0/cert.der \
     --watch-dir target/zk \
     --name king --i 0 \
     --private-identity-der certs/0/key.der
@@ -34,15 +27,9 @@ RUST_LOG=debug cargo run --release --bin zknode -- --king-ip 127.0.0.1:5555 \
 **Node 1**
 
 ```sh
-RUST_LOG=debug cargo run --release --bin zknode -- --king-ip 127.0.0.1:5555 \
-    --certs certs/0/cert.der \
-    --certs certs/1/cert.der \
-    --certs certs/2/cert.der \
-    --certs certs/3/cert.der \
-    --certs certs/4/cert.der \
-    --certs certs/5/cert.der \
-    --certs certs/6/cert.der \
-    --certs certs/7/cert.der \
+RUST_LOG=debug cargo run --release --example zknode -- --king-ip 127.0.0.1:5555 \
+    --public_identity_der certs/1/cert.der \
+    --client_only_king_public_identity_der certs/0/cert.der \
     --watch-dir target/zk \
     --name node1 --i 1 \
     --private-identity-der certs/1/key.der
@@ -50,15 +37,9 @@ RUST_LOG=debug cargo run --release --bin zknode -- --king-ip 127.0.0.1:5555 \
 
 **Node 2**
 ```sh
-RUST_LOG=debug cargo run --release --bin zknode -- --king-ip 127.0.0.1:5555 \
-    --certs certs/0/cert.der \
-    --certs certs/1/cert.der \
-    --certs certs/2/cert.der \
-    --certs certs/3/cert.der \
-    --certs certs/4/cert.der \
-    --certs certs/5/cert.der \
-    --certs certs/6/cert.der \
-    --certs certs/7/cert.der \
+RUST_LOG=debug cargo run --release --example zknode -- --king-ip 127.0.0.1:5555 \
+    --public_identity_der certs/2/cert.der \
+    --client_only_king_public_identity_der certs/0/cert.der \
     --watch-dir target/zk \
     --name node2 --i 2 \
     --private-identity-der certs/2/key.der
@@ -66,15 +47,9 @@ RUST_LOG=debug cargo run --release --bin zknode -- --king-ip 127.0.0.1:5555 \
 
 **Node 3**
 ```sh
-RUST_LOG=debug cargo run --release --bin zknode -- --king-ip 127.0.0.1:5555 \
-    --certs certs/0/cert.der \
-    --certs certs/1/cert.der \
-    --certs certs/2/cert.der \
-    --certs certs/3/cert.der \
-    --certs certs/4/cert.der \
-    --certs certs/5/cert.der \
-    --certs certs/6/cert.der \
-    --certs certs/7/cert.der \
+RUST_LOG=debug cargo run --release --example zknode -- --king-ip 127.0.0.1:5555 \
+    --public_identity_der certs/3/cert.der \
+    --client_only_king_public_identity_der certs/0/cert.der \
     --watch-dir target/zk \
     --name node3 --i 3 \
     --private-identity-der certs/3/key.der
@@ -83,15 +58,9 @@ RUST_LOG=debug cargo run --release --bin zknode -- --king-ip 127.0.0.1:5555 \
 
 **Node 4**
 ```sh
-RUST_LOG=debug cargo run --release --bin zknode -- --king-ip 127.0.0.1:5555 \
-    --certs certs/0/cert.der \
-    --certs certs/1/cert.der \
-    --certs certs/2/cert.der \
-    --certs certs/3/cert.der \
-    --certs certs/4/cert.der \
-    --certs certs/5/cert.der \
-    --certs certs/6/cert.der \
-    --certs certs/7/cert.der \
+RUST_LOG=debug cargo run --release --example zknode -- --king-ip 127.0.0.1:5555 \
+    --public_identity_der certs/4/cert.der \
+    --client_only_king_public_identity_der certs/0/cert.der \
     --watch-dir target/zk \
     --name node4 --i 4 \
     --private-identity-der certs/4/key.der
@@ -99,15 +68,9 @@ RUST_LOG=debug cargo run --release --bin zknode -- --king-ip 127.0.0.1:5555 \
 
 **Node 5**
 ```sh
-RUST_LOG=debug cargo run --release --bin zknode -- --king-ip 127.0.0.1:5555 \
-    --certs certs/0/cert.der \
-    --certs certs/1/cert.der \
-    --certs certs/2/cert.der \
-    --certs certs/3/cert.der \
-    --certs certs/4/cert.der \
-    --certs certs/5/cert.der \
-    --certs certs/6/cert.der \
-    --certs certs/7/cert.der \
+RUST_LOG=debug cargo run --release --example zknode -- --king-ip 127.0.0.1:5555 \
+    --public_identity_der certs/5/cert.der \
+    --client_only_king_public_identity_der certs/0/cert.der \
     --watch-dir target/zk \
     --name node5 --i 5 \
     --private-identity-der certs/5/key.der
@@ -115,15 +78,9 @@ RUST_LOG=debug cargo run --release --bin zknode -- --king-ip 127.0.0.1:5555 \
 
 **Node 6**
 ```sh
-RUST_LOG=debug cargo run --release --bin zknode -- --king-ip 127.0.0.1:5555 \
-    --certs certs/0/cert.der \
-    --certs certs/1/cert.der \
-    --certs certs/2/cert.der \
-    --certs certs/3/cert.der \
-    --certs certs/4/cert.der \
-    --certs certs/5/cert.der \
-    --certs certs/6/cert.der \
-    --certs certs/7/cert.der \
+RUST_LOG=debug cargo run --release --example zknode -- --king-ip 127.0.0.1:5555 \
+    --public_identity_der certs/6/cert.der \
+    --client_only_king_public_identity_der certs/0/cert.der \
     --watch-dir target/zk \
     --name node6 --i 6 \
     --private-identity-der certs/6/key.der
@@ -131,15 +88,9 @@ RUST_LOG=debug cargo run --release --bin zknode -- --king-ip 127.0.0.1:5555 \
 
 **Node 7**
 ```sh
-RUST_LOG=debug cargo run --release --bin zknode -- --king-ip 127.0.0.1:5555 \
-    --certs certs/0/cert.der \
-    --certs certs/1/cert.der \
-    --certs certs/2/cert.der \
-    --certs certs/3/cert.der \
-    --certs certs/4/cert.der \
-    --certs certs/5/cert.der \
-    --certs certs/6/cert.der \
-    --certs certs/7/cert.der \
+RUST_LOG=debug cargo run --release --example zknode -- --king-ip 127.0.0.1:5555 \
+    --public_identity_der certs/7/cert.der \
+    --client_only_king_public_identity_der certs/0/cert.der \
     --watch-dir target/zk \
     --name node7 --i 7 \
     --private-identity-der certs/7/key.der
diff --git a/playground/certgen.rs b/playground/examples/certgen.rs
similarity index 91%
rename from playground/certgen.rs
rename to playground/examples/certgen.rs
index 70ab2ad1a..0014860d5 100644
--- a/playground/certgen.rs
+++ b/playground/examples/certgen.rs
@@ -27,6 +27,8 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
         .push(DnType::CommonName, "Master Cert");
     params.subject_alt_names = vec![
         SanType::IpAddress("127.0.0.1".parse()?),
+        SanType::IpAddress("0.0.0.0".parse()?),
+        SanType::IpAddress("172.40.0.10".parse()?), // For docker container
         SanType::DnsName("localhost".to_string()),
     ];
 
diff --git a/playground/client.rs b/playground/examples/client.rs
similarity index 100%
rename from playground/client.rs
rename to playground/examples/client.rs
diff --git a/playground/node.rs b/playground/examples/node.rs
similarity index 84%
rename from playground/node.rs
rename to playground/examples/node.rs
index 6bc5dd541..701e05660 100644
--- a/playground/node.rs
+++ b/playground/examples/node.rs
@@ -2,7 +2,6 @@ use std::{
     collections::HashMap,
     net::SocketAddr,
     path::PathBuf,
-    pin::Pin,
     str::FromStr,
     sync::{
         atomic::{AtomicU64, Ordering},
@@ -23,11 +22,11 @@ use ark_std::{cfg_iter, start_timer};
 use ark_std::{end_timer, Zero};
 use clap::Parser;
 use futures::{future, TryFutureExt};
+use gadget_core::job::{BuiltExecutableJobWrapper, JobBuilder, JobError};
 use gadget_core::{gadget::substrate::Client, job_manager::SendFuture};
 use groth16::proving_key::PackedProvingKeyShare;
 use mpc_net::prod::CertToDer;
 use mpc_net::MultiplexedStreamID;
-use rustls::RootCertStore;
 use sc_client_api::FinalizeSummary;
 use sc_utils::mpsc::{TracingUnboundedReceiver, TracingUnboundedSender};
 use secret_sharing::pss::PackedSharingParams;
@@ -65,6 +64,8 @@ struct Args {
     /// The Id of the node, also acts as the party index
     #[arg(long)]
     i: usize,
+    #[arg(long)]
+    n: usize,
     /// King's IP address
     #[arg(long)]
     king_ip: SocketAddr,
@@ -74,7 +75,9 @@ struct Args {
     /// All the nodes' public identity DER (including the current node)
     /// it must be in the same order as the nodes' Ids
     #[arg(long)]
-    certs: Vec<PathBuf>,
+    public_identity_der: PathBuf,
+    #[arg(long)]
+    client_only_king_public_identity_der: Option<PathBuf>,
     /// Watch directory for new jobs
     #[arg(long)]
     watch_dir: PathBuf,
@@ -84,10 +87,7 @@ struct Args {
 async fn main() -> anyhow::Result<()> {
     setup_log();
     let args = Args::parse();
-    assert!(
-        args.i < args.certs.len(),
-        "i must be less than n (certs.len())"
-    );
+    assert!(args.i < args.n, "i must be less than n");
 
     tracing::info!(node = args.name, "Starting zknode");
 
@@ -97,61 +97,39 @@ async fn main() -> anyhow::Result<()> {
         processing_semophore: Arc::new(Semaphore::new(1)),
         watch_dir: args.watch_dir,
     };
+
     let client = BlockchainClient::new(state.clone());
     let (done_tx, mut done_rx) = tokio::sync::mpsc::unbounded_channel::<()>();
 
-    let mut cert_store = RootCertStore::empty();
-    // add all certs except our own
-    let iter = args.certs.iter().enumerate().filter(|(i, _)| *i != args.i);
-    for (_, cert_path) in iter {
-        let cert = rustls::Certificate(std::fs::read(cert_path)?);
-        cert_store.add(&cert)?;
-    }
+    let (public_identity_der, private_identity_der) = {
+        let public_identity_der_path = args.public_identity_der;
+        let private_identity_der_path = args.private_identity_der;
+        (
+            std::fs::read(public_identity_der_path)?,
+            std::fs::read(private_identity_der_path)?,
+        )
+    };
 
     let gadget_config = if args.i == 0 {
-        let (king_public_identity_der, king_private_identity_der) = {
-            let king_public_identity_der_path = args
-                .certs
-                .get(args.i)
-                .expect("King's public identity DER path is required");
-            let king_private_identity_der_path = args.private_identity_der;
-            (
-                std::fs::read(king_public_identity_der_path)?,
-                std::fs::read(king_private_identity_der_path)?,
-            )
-        };
         zk_gadget::ZkGadgetConfig {
             king_bind_addr: Some(args.king_ip),
             client_only_king_addr: None,
             id: args.i as _,
-            n_parties: args.certs.len(),
-            public_identity_der: king_public_identity_der.clone(),
-            private_identity_der: king_private_identity_der.clone(),
+            n_parties: args.n,
+            public_identity_der,
+            private_identity_der,
             client_only_king_public_identity_der: None,
         }
     } else {
         let king_public_identity_der_path = args
-            .certs
-            .get(0)
+            .client_only_king_public_identity_der
             .expect("King's public identity DER path is required");
         let king_public_identity_der = std::fs::read(king_public_identity_der_path)?;
-        // read our identity
-        let (public_identity_der, private_identity_der) = {
-            let public_identity_der_path = args
-                .certs
-                .get(args.i)
-                .expect("Public identity DER path is required");
-            let private_identity_der_path = args.private_identity_der;
-            (
-                std::fs::read(public_identity_der_path)?,
-                std::fs::read(private_identity_der_path)?,
-            )
-        };
         zk_gadget::ZkGadgetConfig {
             king_bind_addr: None,
             client_only_king_addr: Some(args.king_ip),
             id: args.i as _,
-            n_parties: args.certs.len(),
+            n_parties: args.n,
             public_identity_der,
             private_identity_der,
             client_only_king_public_identity_der: Some(king_public_identity_der.clone()),
@@ -208,8 +186,20 @@ fn async_protocol_generator(
         BlockchainClient,
         TestBlock,
     >,
-) -> Pin<Box<dyn SendFuture<'static, Result<(), webb_gadget::Error>>>> {
-    Box::pin(async move {
+) -> BuiltExecutableJobWrapper<impl SendFuture<'static, Result<(), JobError>>> {
+    let stop_tx = params.extra_parameters.stop_tx.clone();
+    let party_id = params.party_id;
+    JobBuilder::default()
+        .post(async move {
+            if party_id == 0 {
+                // If we are the king, wait a little bit for the other nodes to safely shutdown
+                tokio::time::sleep(Duration::from_millis(1000)).await;
+            }
+
+            stop_tx.send(()).expect("Failed to send stop signal");
+            Ok(())
+        })
+        .build(async move {
         let _permit = match params.client.state.processing_semophore.try_acquire() {
             Ok(permit) => permit,
             Err(tokio::sync::TryAcquireError::NoPermits) => {
@@ -221,31 +211,23 @@ fn async_protocol_generator(
         let jobs = params.client.state.jobs.lock().await;
         let circuits = params.client.state.circuits.lock().await;
         let Some(job) = jobs.get(&params.associated_task_id).cloned() else {
-            return Err(webb_gadget::Error::ClientError {
-                err: format!(
-                    "Job with id {} not found",
-                    hex::encode(params.associated_task_id)
-                ),
-            });
+            return Err(JobError::from(format!(
+                "Job with id {} not found",
+                hex::encode(params.associated_task_id)
+            )));
         };
         let Some(circuit) = circuits.get(&job.circuit_id).cloned() else {
-            return Err(webb_gadget::Error::ClientError {
-                err: format!("Circuit with id {} not found", job.circuit_id),
-            });
+            return Err(JobError::from(format!("Circuit with id {} not found", job.circuit_id)));
         };
 
         drop(jobs);
         drop(circuits);
 
         let proving_key = tokio::fs::read(&circuit.pk_uri)
-            .map_err(|e| webb_gadget::Error::ClientError {
-                err: format!("Failed to read proving key: {e}"),
-            })
+            .map_err(|e| JobError::from(format!("Failed to read proving key: {e}")))
             .await?;
         let pk = ProvingKey::deserialize_compressed(proving_key.as_slice()).map_err(|e| {
-            webb_gadget::Error::ClientError {
-                err: format!("Failed to deserialize proving key: {e}"),
-            }
+            JobError::from(format!("Failed to deserialize proving key: {e}"))
         })?;
 
         let pp = PackedSharingParams::new(job.pss_l);
@@ -258,19 +240,9 @@ fn async_protocol_generator(
         let our_qap_share = job
             .qap_shares
             .get(params.party_id as usize)
-            .ok_or_else(|| webb_gadget::Error::ClientError {
-                err: format!("QAP Shares for party {} not found", params.party_id),
-            })?;
-        let our_a_share = job.a_shares.get(params.party_id as usize).ok_or_else(|| {
-            webb_gadget::Error::ClientError {
-                err: format!("a Shares for party {} not found", params.party_id),
-            }
-        })?;
-        let our_ax_share = job.ax_shares.get(params.party_id as usize).ok_or_else(|| {
-            webb_gadget::Error::ClientError {
-                err: format!("ax Shares for party {} not found", params.party_id),
-            }
-        })?;
+            .ok_or_else(|| JobError::from(format!("QAP Shares for party {} not found", params.party_id)))?;
+        let our_a_share = job.a_shares.get(params.party_id as usize).ok_or_else(|| JobError::from(format!("a Shares for party {} not found", params.party_id)))?;
+        let our_ax_share = job.ax_shares.get(params.party_id as usize).ok_or_else(|| JobError::from(format!("ax Shares for party {} not found", params.party_id)))?;
         let m = circuit.num_inputs + circuit.num_constraints;
         let domain = Radix2EvaluationDomain::new(m)
             .ok_or(SynthesisError::PolynomialDegreeTooLarge)
@@ -293,11 +265,7 @@ fn async_protocol_generator(
             domain,
         };
 
-        let crs_share = crs_shares.get(params.party_id as usize).ok_or_else(|| {
-            webb_gadget::Error::ClientError {
-                err: format!("CRS Shares for party {} not found", params.party_id),
-            }
-        })?;
+        let crs_share = crs_shares.get(params.party_id as usize).ok_or_else(|| JobError::from(format!("CRS Shares for party {} not found", params.party_id)))?;
         let a_share = cfg_iter!(our_a_share)
             .map(|f| F::deserialize_compressed(f.as_slice()))
             .collect::<Result<Vec<_>, _>>()
diff --git a/playground/generate_certs.sh b/playground/generate_certs.sh
index cd075ba3c..a47f9d4a4 100755
--- a/playground/generate_certs.sh
+++ b/playground/generate_certs.sh
@@ -1,12 +1,9 @@
+#!/bin/sh
 # A simple script that invokes a simple binary to generate a self-signed
 # certificate for use in the playground.
-
-# 1. Build the binary
-cargo build --bin certgen
-# 2. Run the binary
-# Generate 8 certs from 0 to 7
-for i in {0..7}
+n=8
+for i in $(seq 0 $(($n - 1)))
   do
-      ./target/debug/certgen --i $i
+      cargo run --release --example certgen -- --i $i
   done
 echo "Done generating certs"
diff --git a/playground/run.sh b/playground/run.sh
new file mode 100755
index 000000000..da6466ead
--- /dev/null
+++ b/playground/run.sh
@@ -0,0 +1,52 @@
+#!/bin/sh
+set -e
+trap "exit" INT TERM
+trap "kill 0" EXIT
+
+# Step 0: cd to git working dir
+cd "$(git rev-parse --show-toplevel)"
+
+# Step 1: check to see if arg "--create-certs" was passed
+if [ "$1" = "--create-certs" ]; then
+    # Step 2: if so, invoke the generate_certs.sh script
+    echo "Creating certs..."
+    ./playground/generate_certs.sh
+fi
+
+# Step 2: Run the client unless --no-client was passed as the second or first parameter
+if [ "$1" != "--no-client" ] && [ "$2" != "--no-client" ]; then
+    cargo run --release --example zk-test-client -- --circuit-id 1 --job-id 1 --wasm ./fixtures/sha256/circom.wasm --r1cs ./fixtures/sha256/circom.r1cs --input ./fixtures/sha256/input.json --public-inputs ./fixtures/sha256/public_inputs.json --output-dir ./target/zk --generate-proving-key
+fi
+
+# Step 3: for x in 0..=7, run a command
+PROCS=()
+SERVER_IP=127.0.0.1
+SERVER_PORT=5555
+n=8
+for i in $(seq 0 $(($n - 1)))
+  do
+    # if i is zero, run the king
+    if [ "$i" -eq "0" ]; then
+        # Step 4: Run the king
+        RUST_LOG=debug cargo run --release --example zknode -- --n $n --king-ip $SERVER_IP:$SERVER_PORT --public-identity-der certs/$i/cert.der --watch-dir target/zk --name king --i $i --private-identity-der certs/0/key.der &
+        pid=$!
+        PROCS[$i]=$pid
+
+        # But first, wait for the king to start
+        while ! nc -z -w 5 $SERVER_IP $SERVER_PORT; do
+          echo "Waiting for server $SERVER_IP:$SERVER_PORT ... "
+          sleep 0.5
+        done
+    else
+        # Step 4: Run the client
+        RUST_LOG=debug cargo run --release --example zknode -- --n $n --king-ip $SERVER_IP:$SERVER_PORT --public-identity-der certs/$i/cert.der --client-only-king-public-identity-der certs/0/cert.der --watch-dir target/zk --name node$i --i $i --private-identity-der certs/$i/key.der &
+        pid=$!
+        PROCS[$i]=$pid
+        sleep 0.2
+    fi
+  done
+
+# Step 5: wait for all processes to finish
+for pid in ${PROCS[@]}; do
+    wait $pid
+done
\ No newline at end of file
diff --git a/zk-gadget/build.rs b/zk-gadget/build.rs
new file mode 100644
index 000000000..0664fb281
--- /dev/null
+++ b/zk-gadget/build.rs
@@ -0,0 +1,4 @@
+fn main() -> Result<(), Box<dyn std::error::Error>> {
+    tonic_build::compile_protos("../protos/zk-gadget-test/test.proto")?;
+    Ok(())
+}
diff --git a/zk-gadget/src/lib.rs b/zk-gadget/src/lib.rs
index 6e684c42a..1f06c9242 100644
--- a/zk-gadget/src/lib.rs
+++ b/zk-gadget/src/lib.rs
@@ -2,8 +2,10 @@ use crate::client_ext::ClientWithApi;
 use crate::module::proto_gen::AsyncProtocolGenerator;
 use crate::module::{AdditionalProtocolParams, ZkModule};
 use crate::network::{RegistantId, ZkNetworkService};
-use gadget_core::Block;
+use gadget_core::job::JobError;
+use gadget_core::job_manager::SendFuture;
 use mpc_net::prod::RustlsCertificate;
+use sp_runtime::traits::Block;
 use std::net::SocketAddr;
 use tokio_rustls::rustls::{Certificate, PrivateKey, RootCertStore};
 use webb_gadget::Error;
@@ -27,7 +29,8 @@ pub async fn run<
     C: ClientWithApi<B>,
     B: Block,
     T: AdditionalProtocolParams,
-    Gen: AsyncProtocolGenerator<T, Error, ZkNetworkService, C, B>,
+    F: SendFuture<'static, Result<(), JobError>>,
+    Gen: AsyncProtocolGenerator<T, Error, ZkNetworkService, C, B, F>,
 >(
     config: ZkGadgetConfig,
     client: C,
@@ -40,6 +43,7 @@ pub async fn run<
 
     let zk_module = ZkModule {
         party_id: config.id,
+        // TODO: The below should be inside job metadata, but NOT determined at runtime
         n_parties: config.n_parties,
         additional_protocol_params,
         network: network.clone(),
diff --git a/zk-gadget/src/module/mod.rs b/zk-gadget/src/module/mod.rs
index 1b2e5a833..aac5cce69 100644
--- a/zk-gadget/src/module/mod.rs
+++ b/zk-gadget/src/module/mod.rs
@@ -1,30 +1,35 @@
 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;
+use gadget_core::job::JobError;
+use gadget_core::job_manager::{ProtocolWorkManager, SendFuture};
+use sp_runtime::traits::Block;
 use webb_gadget::gadget::work_manager::WebbWorkManager;
 use webb_gadget::gadget::WebbGadgetModule;
 use webb_gadget::{BlockImportNotification, Error, FinalityNotification};
 
 pub mod proto_gen;
 
-pub struct ZkModule<T, C, B> {
+pub struct ZkModule<T, C, B, F: SendFuture<'static, Result<(), JobError>>> {
     pub party_id: RegistantId,
     pub n_parties: usize,
     pub additional_protocol_params: T,
     pub client: C,
     pub network: ZkNetworkService,
     pub async_protocol_generator:
-        Box<dyn proto_gen::AsyncProtocolGenerator<T, Error, ZkNetworkService, C, B>>,
+        Box<dyn proto_gen::AsyncProtocolGenerator<T, Error, ZkNetworkService, C, B, F>>,
 }
 
 pub trait AdditionalProtocolParams: Send + Sync + Clone + 'static {}
 impl<T: Send + Sync + Clone + 'static> AdditionalProtocolParams for T {}
 
 #[async_trait]
-impl<B: Block, T: AdditionalProtocolParams, C: ClientWithApi<B>> WebbGadgetModule<B>
-    for ZkModule<T, C, B>
+impl<
+        B: Block,
+        T: AdditionalProtocolParams,
+        C: ClientWithApi<B>,
+        F: SendFuture<'static, Result<(), JobError>>,
+    > WebbGadgetModule<B> for ZkModule<T, C, B, F>
 {
     async fn process_finality_notification(
         &self,
diff --git a/zk-gadget/src/module/proto_gen.rs b/zk-gadget/src/module/proto_gen.rs
index e9d700c25..3f47eafad 100644
--- a/zk-gadget/src/module/proto_gen.rs
+++ b/zk-gadget/src/module/proto_gen.rs
@@ -2,6 +2,9 @@ use crate::client_ext::ClientWithApi;
 use crate::network::RegistantId;
 use async_trait::async_trait;
 use bytes::Bytes;
+use gadget_core::job::{
+    BuiltExecutableJobWrapper, ExecutableJob, JobBuilder, JobError, ProceedWithExecution,
+};
 use gadget_core::job_manager::{ProtocolRemote, SendFuture, ShutdownReason, WorkManagerInterface};
 use mpc_net::{MpcNet, MpcNetError, MultiplexedStreamID};
 use parking_lot::Mutex;
@@ -10,7 +13,6 @@ use sp_runtime::traits::Block;
 use std::collections::HashMap;
 use std::fmt::Debug;
 use std::marker::PhantomData;
-use std::pin::Pin;
 use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::Arc;
 use webb_gadget::gadget::message::GadgetProtocolMessage;
@@ -50,11 +52,10 @@ pub struct ZkProtocolRemote {
     pub is_done: Arc<AtomicBool>,
 }
 
-pub trait AsyncProtocolGenerator<B, E, N, C, Bl>:
-    Send
-    + Sync
-    + 'static
-    + Fn(ZkAsyncProtocolParameters<B, N, C, Bl>) -> Pin<Box<dyn SendFuture<'static, Result<(), E>>>>
+pub trait AsyncProtocolGenerator<B, E, N, C, Bl, F>:
+    Send + Sync + 'static + Fn(ZkAsyncProtocolParameters<B, N, C, Bl>) -> BuiltExecutableJobWrapper<F>
+where
+    F: SendFuture<'static, Result<(), JobError>>,
 {
 }
 impl<
@@ -63,13 +64,12 @@ impl<
         N: Network,
         Bl: Block,
         C: ClientWithApi<Bl>,
+        F: SendFuture<'static, Result<(), JobError>>,
         T: Send
             + Sync
             + 'static
-            + Fn(
-                ZkAsyncProtocolParameters<B, N, C, Bl>,
-            ) -> Pin<Box<dyn SendFuture<'static, Result<(), E>>>>,
-    > AsyncProtocolGenerator<B, E, N, C, Bl> for T
+            + Fn(ZkAsyncProtocolParameters<B, N, C, Bl>) -> BuiltExecutableJobWrapper<F>,
+    > AsyncProtocolGenerator<B, E, N, C, Bl, F> for T
 {
 }
 
@@ -80,6 +80,7 @@ pub fn create_zk_async_protocol<
     N: Network,
     C: ClientWithApi<Bl>,
     Bl: Block,
+    F: SendFuture<'static, Result<(), JobError>>,
 >(
     session_id: <WebbWorkManager as WorkManagerInterface>::SessionID,
     now: <WebbWorkManager as WorkManagerInterface>::Clock,
@@ -90,14 +91,18 @@ pub fn create_zk_async_protocol<
     extra_parameters: B,
     network: N,
     client: C,
-    proto_gen: &dyn AsyncProtocolGenerator<B, E, N, C, Bl>,
-) -> (ZkProtocolRemote, Pin<Box<dyn SendFuture<'static, ()>>>) {
+    proto_gen: &dyn AsyncProtocolGenerator<B, E, N, C, Bl, F>,
+) -> (
+    ZkProtocolRemote,
+    BuiltExecutableJobWrapper<impl SendFuture<'static, Result<(), JobError>>>,
+) {
     let is_done = Arc::new(AtomicBool::new(false));
     let (to_async_protocol, mut protocol_message_rx) = tokio::sync::mpsc::unbounded_channel();
     let (start_tx, start_rx) = tokio::sync::oneshot::channel();
     let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel();
 
     let proto_hash_hex = hex::encode(task_id);
+    let proto_hash_hex_clone = proto_hash_hex.clone();
 
     let mut txs = HashMap::new();
     let mut rxs = HashMap::new();
@@ -150,39 +155,59 @@ pub fn create_zk_async_protocol<
         is_done: is_done.clone(),
     };
 
-    let async_protocol = proto_gen(params);
+    let mut async_protocol = proto_gen(params);
+
+    let pre_hook = async move {
+        match start_rx.await {
+            Ok(_) => Ok(ProceedWithExecution::True),
+            Err(err) => {
+                log::error!(
+                    "Protocol {proto_hash_hex_clone} failed to receive start signal: {err:?}"
+                );
+                Ok(ProceedWithExecution::False)
+            }
+        }
+    };
+
+    let post_hook = async move {
+        // Mark the task as done
+        is_done.store(true, Ordering::SeqCst);
+        Ok(())
+    };
 
     // This wrapped future enables proper functionality between the async protocol and the
     // job manager
-    let wrapped_future = Box::pin(async move {
-        if let Err(err) = start_rx.await {
-            log::error!("Protocol {proto_hash_hex} failed to receive start signal: {err:?}");
-        } else {
-            tokio::select! {
-                res0 = async_protocol => {
-                    if let Err(err) = res0 {
-                        log::error!("Protocol {proto_hash_hex} failed: {err:?}");
-                    } else {
-                        log::info!("Protocol {proto_hash_hex} finished");
-                    }
-                },
-
-                res1 = shutdown_rx => {
-                    match res1 {
-                        Ok(reason) => {
-                            log::info!("Protocol {proto_hash_hex} shutdown: {reason:?}");
-                        },
-                        Err(err) => {
-                            log::error!("Protocol {proto_hash_hex} shutdown failed: {err:?}");
-                        },
-                    }
+    let wrapped_future = async move {
+        tokio::select! {
+            res0 = async_protocol.execute() => {
+                if let Err(err) = res0 {
+                    log::error!("Protocol {proto_hash_hex} failed: {err:?}");
+                    Err(JobError::from(err.to_string()))
+                } else {
+                    log::info!("Protocol {proto_hash_hex} finished");
+                    Ok(())
+                }
+            },
+
+            res1 = shutdown_rx => {
+                match res1 {
+                    Ok(reason) => {
+                        log::info!("Protocol {proto_hash_hex} shutdown: {reason:?}");
+                        Ok(())
+                    },
+                    Err(err) => {
+                        log::error!("Protocol {proto_hash_hex} shutdown failed: {err:?}");
+                        Err(JobError::from(err.to_string()))
+                    },
                 }
             }
         }
+    };
 
-        // Mark the task as done
-        is_done.store(true, Ordering::SeqCst);
-    });
+    let wrapped_future = JobBuilder::default()
+        .pre(pre_hook)
+        .post(post_hook)
+        .build(wrapped_future);
 
     (remote, wrapped_future)
 }
diff --git a/zk-gadget/src/network/mod.rs b/zk-gadget/src/network/mod.rs
index 954e68125..8b1d53d4e 100644
--- a/zk-gadget/src/network/mod.rs
+++ b/zk-gadget/src/network/mod.rs
@@ -251,13 +251,23 @@ pub enum RegistryPacket {
 }
 
 fn handle_stream_as_king(
-    stream: TlsStream<TcpStream>,
+    tls_acceptor: TlsAcceptor,
+    stream: TcpStream,
     peer_addr: SocketAddr,
     registrants: Arc<Mutex<HashMap<RegistantId, Registrant>>>,
     to_outbound_txs: Arc<RwLock<HashMap<RegistantId, UnboundedSender<RegistryPacket>>>>,
     to_gadget: UnboundedSender<RegistryPacket>,
 ) {
     tokio::task::spawn(async move {
+        let stream = match tls_acceptor.accept(stream).await {
+            Ok(stream) => stream,
+            Err(err) => {
+                log::error!("[Registry] Failed to upgrade connection from {peer_addr}: {err:?}");
+                return;
+            }
+        };
+
+        let stream = TlsStream::Server(stream);
         let wrapped_stream = mpc_net::multi::wrap_stream(stream);
         let (mut sink, mut stream) = wrapped_stream.split();
         let (to_outbound_tx, mut to_outbound_rx) = tokio::sync::mpsc::unbounded_channel();
@@ -272,7 +282,7 @@ fn handle_stream_as_king(
                 }
             }
 
-            log::error!("to_outbound_rx closed");
+            log::warn!("to_outbound_rx closed");
         });
 
         while let Some(Ok(message)) = stream.next().await {
@@ -313,7 +323,7 @@ fn handle_stream_as_king(
             registrants.remove(&id);
         }
 
-        log::error!("[Registry] Connection closed to peer {peer_addr}")
+        log::warn!("[Registry] Connection closed to peer {peer_addr}")
     });
 }
 
@@ -462,14 +472,9 @@ impl Network for ZkNetworkService {
 
                 while let Ok((stream, peer_addr)) = listener.accept().await {
                     log::info!("[Registry] Accepted connection from {peer_addr}, upgrading to TLS");
-                    let stream = tls_acceptor.accept(stream).await.map_err(|err| {
-                        Error::RegistryCreateError {
-                            err: format!("{err:?}"),
-                        }
-                    })?;
-
                     handle_stream_as_king(
-                        TlsStream::Server(stream),
+                        tls_acceptor.clone(),
+                        stream,
                         peer_addr,
                         registrants.clone(),
                         to_outbound_txs.clone(),


Test patch:
diff --git a/protos/zk-gadget-test/test.proto b/protos/zk-gadget-test/test.proto
new file mode 100644
index 000000000..8ad65f68b
--- /dev/null
+++ b/protos/zk-gadget-test/test.proto
@@ -0,0 +1,12 @@
+syntax = "proto3";
+package test_blockchain_zk;
+
+message GetLatestHeaderRequest {}
+message GetLatestHeaderResponse {
+  uint64 latest_block_number = 1;
+  uint64 session_id = 2;
+}
+
+service Auth {
+  rpc GetLatestHeader(GetLatestHeaderRequest) returns (GetLatestHeaderResponse) {}
+}
\ No newline at end of file
diff --git a/test-gadget/src/gadget.rs b/test-gadget/src/gadget.rs
index 4653ecabc..f443a5d2d 100644
--- a/test-gadget/src/gadget.rs
+++ b/test-gadget/src/gadget.rs
@@ -5,6 +5,7 @@ use crate::work_manager::{
 };
 use async_trait::async_trait;
 use gadget_core::gadget::manager::AbstractGadget;
+use gadget_core::job::{BuiltExecutableJobWrapper, JobBuilder, JobError};
 use gadget_core::job_manager::{PollMethod, ProtocolWorkManager, SendFuture, WorkManagerInterface};
 use parking_lot::RwLock;
 use std::pin::Pin;
@@ -133,6 +134,7 @@ impl<B: Send + Sync + Clone + 'static> AbstractGadget for TestGadget<B> {
     }
 }
 
+#[allow(clippy::type_complexity)]
 fn create_test_async_protocol<B: Send + Sync + 'static>(
     session_id: <TestWorkManager as WorkManagerInterface>::SessionID,
     now: <TestWorkManager as WorkManagerInterface>::Clock,
@@ -140,7 +142,10 @@ fn create_test_async_protocol<B: Send + Sync + 'static>(
     task_id: <TestWorkManager as WorkManagerInterface>::TaskID,
     test_bundle: B,
     proto_gen: &dyn AsyncProtocolGenerator<B>,
-) -> (TestProtocolRemote, Pin<Box<dyn SendFuture<'static, ()>>>) {
+) -> (
+    TestProtocolRemote,
+    BuiltExecutableJobWrapper<Pin<Box<dyn SendFuture<'static, Result<(), JobError>>>>>,
+) {
     let is_done = Arc::new(AtomicBool::new(false));
     let (to_async_protocol, protocol_message_rx) = tokio::sync::mpsc::unbounded_channel();
     let (start_tx, start_rx) = tokio::sync::oneshot::channel();
@@ -169,6 +174,7 @@ fn create_test_async_protocol<B: Send + Sync + 'static>(
     };
 
     let future = proto_gen(params);
+    let future = JobBuilder::default().build(future);
 
     (remote, future)
 }
diff --git a/test-gadget/src/work_manager.rs b/test-gadget/src/work_manager.rs
index 851238b72..e48ac526a 100644
--- a/test-gadget/src/work_manager.rs
+++ b/test-gadget/src/work_manager.rs
@@ -1,5 +1,6 @@
 use crate::error::TestError;
 use crate::message::TestProtocolMessage;
+use gadget_core::job::JobError;
 use gadget_core::job_manager::{ProtocolRemote, SendFuture, ShutdownReason, WorkManagerInterface};
 use parking_lot::{Mutex, RwLock};
 use std::pin::Pin;
@@ -131,12 +132,19 @@ pub struct TestAsyncProtocolParameters<B> {
 }
 
 pub trait AsyncProtocolGenerator<B>:
-    Send + Sync + Fn(TestAsyncProtocolParameters<B>) -> Pin<Box<dyn SendFuture<'static, ()>>>
+    Send
+    + Sync
+    + Fn(TestAsyncProtocolParameters<B>) -> Pin<Box<dyn SendFuture<'static, Result<(), JobError>>>>
 {
 }
+
 impl<
         B: Send + Sync,
-        T: Send + Sync + Fn(TestAsyncProtocolParameters<B>) -> Pin<Box<dyn SendFuture<'static, ()>>>,
+        T: Send
+            + Sync
+            + Fn(
+                TestAsyncProtocolParameters<B>,
+            ) -> Pin<Box<dyn SendFuture<'static, Result<(), JobError>>>>,
     > AsyncProtocolGenerator<B> for T
 {
 }
diff --git a/test-gadget/tests/basic.rs b/test-gadget/tests/basic.rs
index 40bcbc2e8..76789a1cb 100644
--- a/test-gadget/tests/basic.rs
+++ b/test-gadget/tests/basic.rs
@@ -1,5 +1,6 @@
 #[cfg(test)]
 mod tests {
+    use gadget_core::job::JobError;
     use gadget_core::job_manager::SendFuture;
     use std::error::Error;
     use std::pin::Pin;
@@ -35,7 +36,7 @@ mod tests {
 
     fn async_proto_generator(
         mut params: TestAsyncProtocolParameters<TestBundle>,
-    ) -> Pin<Box<dyn SendFuture<'static, ()>>> {
+    ) -> Pin<Box<dyn SendFuture<'static, Result<(), JobError>>>> {
         let n_peers = params.test_bundle.n_peers;
         let party_id = params.test_bundle.party_id;
         Box::pin(async move {
@@ -79,6 +80,7 @@ mod tests {
                 .count_finished_tx
                 .send(())
                 .expect("Didn't send on_finish signal");
+            Ok::<_, JobError>(())
         })
     }
 }
diff --git a/test-gadget/tests/dfft.rs b/test-gadget/tests/dfft.rs
index 8afbc0df3..e45d477c3 100644
--- a/test-gadget/tests/dfft.rs
+++ b/test-gadget/tests/dfft.rs
@@ -20,6 +20,7 @@ mod tests {
         dfft::{d_fft, fft_in_place_rearrange},
         utils::pack::transpose,
     };
+    use gadget_core::job::JobError;
     use mpc_net::{MpcNet, MpcNetError, MultiplexedStreamID};
     use secret_sharing::pss::PackedSharingParams;
     use serde::{Deserialize, Serialize};
@@ -120,7 +121,7 @@ mod tests {
 
     fn async_proto_generator(
         mut params: TestAsyncProtocolParameters<TestBundle>,
-    ) -> Pin<Box<dyn SendFuture<'static, ()>>> {
+    ) -> Pin<Box<dyn SendFuture<'static, Result<(), JobError>>>> {
         Box::pin(async move {
             params
                 .start_rx
@@ -186,6 +187,7 @@ mod tests {
             let dom = Radix2EvaluationDomain::<Fr>::new(M).unwrap();
             d_fft_test::<Fr, _>(&pp, &dom, &network).await;
             on_end_tx.send(()).expect("Failed to send on_end signal");
+            Ok(())
         })
     }
 
diff --git a/test-gadget/tests/dmsm.rs b/test-gadget/tests/dmsm.rs
index bbc361f8b..b13a9ce3d 100644
--- a/test-gadget/tests/dmsm.rs
+++ b/test-gadget/tests/dmsm.rs
@@ -18,6 +18,7 @@ mod tests {
     use ark_std::UniformRand;
     use async_trait::async_trait;
     use bytes::Bytes;
+    use gadget_core::job::JobError;
     use mpc_net::{MpcNet, MpcNetError, MultiplexedStreamID};
     use secret_sharing::pss::PackedSharingParams;
     use serde::{Deserialize, Serialize};
@@ -115,7 +116,7 @@ mod tests {
 
     fn async_proto_generator(
         mut params: TestAsyncProtocolParameters<TestBundle>,
-    ) -> Pin<Box<dyn SendFuture<'static, ()>>> {
+    ) -> Pin<Box<dyn SendFuture<'static, Result<(), JobError>>>> {
         Box::pin(async move {
             params
                 .start_rx
@@ -171,6 +172,7 @@ mod tests {
             let dom = Radix2EvaluationDomain::<Fr>::new(M).unwrap();
             d_msm_test::<ark_bls12_377::G1Projective, _>(pp, dom, network).await;
             on_end_tx.send(()).expect("Failed to send on_end signal");
+            Ok(())
         })
     }
 
diff --git a/test-gadget/tests/zk_network.rs b/test-gadget/tests/zk_network.rs
index 273f88d77..0a1fc7564 100644
--- a/test-gadget/tests/zk_network.rs
+++ b/test-gadget/tests/zk_network.rs
@@ -12,6 +12,7 @@ mod tests {
 
     use async_trait::async_trait;
     use bytes::Bytes;
+    use gadget_core::job::JobError;
     use mpc_net::{MpcNet, MpcNetError, MultiplexedStreamID};
     use serde::{Deserialize, Serialize};
     use test_gadget::message::{TestProtocolMessage, UserID};
@@ -40,7 +41,7 @@ mod tests {
 
     fn async_proto_generator(
         mut params: TestAsyncProtocolParameters<TestBundle>,
-    ) -> Pin<Box<dyn SendFuture<'static, ()>>> {
+    ) -> Pin<Box<dyn SendFuture<'static, Result<(), JobError>>>> {
         Box::pin(async move {
             params
                 .start_rx
@@ -88,7 +89,7 @@ mod tests {
                         bincode2::deserialize(&message.payload).expect("Failed to deser message");
                     let txs = &txs[&deserialized.source];
                     let tx = &txs[deserialized.sid as usize];
-                    tx.send(deserialized).expect("Failed to send message");
+                    tx.send(deserialized).expect("Failed to send");
                 }
             });
 
@@ -145,6 +146,7 @@ mod tests {
             }
 
             on_end_tx.send(()).expect("Failed to send on_end signal");
+            Ok(())
         })
     }
 
diff --git a/zk-gadget/tests/jobs.rs b/zk-gadget/tests/jobs.rs
new file mode 100644
index 000000000..3379dd4a7
--- /dev/null
+++ b/zk-gadget/tests/jobs.rs
@@ -0,0 +1,380 @@
+#[cfg(test)]
+mod tests {
+    use crate::tests::client::{BlockchainClient, TestBlock};
+    use futures_util::stream::FuturesUnordered;
+    use futures_util::TryStreamExt;
+    use gadget_core::job::{BuiltExecutableJobWrapper, JobBuilder, JobError};
+    use gadget_core::job_manager::SendFuture;
+    use mpc_net::{MpcNet, MultiplexedStreamID};
+    use std::error::Error;
+    use std::net::SocketAddr;
+    use std::time::Duration;
+    use tokio::sync::mpsc::UnboundedSender;
+    use tracing_subscriber::fmt::SubscriberBuilder;
+    use tracing_subscriber::util::SubscriberInitExt;
+    use tracing_subscriber::EnvFilter;
+    use zk_gadget::module::proto_gen::ZkAsyncProtocolParameters;
+    use zk_gadget::network::ZkNetworkService;
+
+    pub fn setup_log() {
+        let _ = SubscriberBuilder::default()
+            .with_env_filter(EnvFilter::from_default_env())
+            .finish()
+            .try_init();
+    }
+
+    pub mod server {
+        use parking_lot::Mutex;
+        use std::collections::HashMap;
+        use std::sync::Arc;
+        use tokio::sync::RwLock;
+        use tonic::{Request, Response, Status};
+        use zk_gadget::client_ext::job_types::ZkJob;
+
+        tonic::include_proto!("test_blockchain_zk");
+
+        #[derive(Default, Clone)]
+        pub struct TestBlockchainServer {
+            pub latest_header: Arc<Mutex<u64>>,
+            pub jobs: Arc<RwLock<HashMap<u64, ZkJob>>>,
+        }
+
+        #[tonic::async_trait]
+        impl auth_server::Auth for TestBlockchainServer {
+            async fn get_latest_header(
+                &self,
+                _request: Request<GetLatestHeaderRequest>,
+            ) -> Result<Response<GetLatestHeaderResponse>, Status> {
+                let latest_block_number = *self.latest_header.lock();
+                let session_id = 0; // For now
+                Ok(Response::new(GetLatestHeaderResponse {
+                    latest_block_number,
+                    session_id,
+                }))
+            }
+        }
+    }
+
+    pub mod client {
+        use async_trait::async_trait;
+        use gadget_core::gadget::substrate::Client;
+        use sc_client_api::FinalizeSummary;
+        use serde::Serialize;
+        use sp_runtime::app_crypto::sp_core::Encode;
+        use sp_runtime::codec::Decode;
+        use std::error::Error;
+        use std::net::SocketAddr;
+        use std::sync::Arc;
+        use std::time::Duration;
+        use tokio::sync::Mutex;
+        use uuid::Uuid;
+        use webb_gadget::{BlockImportNotification, FinalityNotification};
+        use zk_gadget::client_ext::job_types::{CircuitProperties, JobProperties};
+        use zk_gadget::client_ext::ClientWithApi;
+
+        #[derive(Clone)]
+        pub struct BlockchainClient {
+            client: Arc<Mutex<super::server::auth_client::AuthClient<tonic::transport::Channel>>>,
+            latest_received_header: Arc<Mutex<u64>>,
+        }
+
+        impl BlockchainClient {
+            pub async fn new(server_addr: SocketAddr) -> Result<Self, Box<dyn Error>> {
+                let addr = format!("http://{server_addr}");
+                let client = super::server::auth_client::AuthClient::connect(addr).await?;
+                Ok(Self {
+                    client: Arc::new(Mutex::new(client)),
+                    latest_received_header: Arc::new(Mutex::new(0)),
+                })
+            }
+        }
+
+        #[async_trait]
+        impl Client<TestBlock> for BlockchainClient {
+            async fn get_next_finality_notification(
+                &self,
+            ) -> Option<FinalityNotification<TestBlock>> {
+                loop {
+                    let response = self.get_latest_finality_notification().await?;
+                    let mut lock = self.latest_received_header.lock().await;
+
+                    if response.header.number > *lock {
+                        *lock = response.header.number;
+                        return Some(response);
+                    }
+
+                    // Wait some time before trying again
+                    tokio::time::sleep(Duration::from_millis(200)).await;
+                }
+            }
+
+            async fn get_latest_finality_notification(
+                &self,
+            ) -> Option<FinalityNotification<TestBlock>> {
+                let request = tonic::Request::new(super::server::GetLatestHeaderRequest {});
+                let response = self
+                    .client
+                    .lock()
+                    .await
+                    .get_latest_header(request)
+                    .await
+                    .ok()?
+                    .into_inner();
+
+                Some(block_number_to_finality_notification(
+                    response.latest_block_number,
+                ))
+            }
+
+            async fn get_next_block_import_notification(
+                &self,
+            ) -> Option<BlockImportNotification<TestBlock>> {
+                futures_util::future::pending().await
+            }
+        }
+
+        #[async_trait]
+        impl ClientWithApi<TestBlock> for BlockchainClient {
+            async fn get_job_circuit_properties(
+                &self,
+                _circuit_id: u64,
+            ) -> Result<Option<CircuitProperties>, webb_gadget::Error> {
+                unimplemented!()
+            }
+
+            async fn get_job_properties(
+                &self,
+                _job_id: u64,
+            ) -> Result<Option<JobProperties>, webb_gadget::Error> {
+                unimplemented!()
+            }
+
+            async fn get_next_job(&self) -> Result<Option<JobProperties>, webb_gadget::Error> {
+                Ok(Some(JobProperties {
+                    job_id: 0,
+                    circuit_id: 0,
+                    public_inputs: vec![],
+                    pss_l: 0,
+                    a_shares: vec![],
+                    ax_shares: vec![],
+                    qap_shares: vec![],
+                }))
+            }
+        }
+
+        pub type TestBlock = sp_runtime::testing::Block<XtDummy>;
+        #[derive(Encode, Decode, Serialize, Clone, Eq, PartialEq, Debug)]
+        pub struct XtDummy;
+
+        impl sp_runtime::traits::Extrinsic for XtDummy {
+            type Call = ();
+            type SignaturePayload = ();
+        }
+
+        fn block_number_to_finality_notification(
+            block_number: u64,
+        ) -> FinalityNotification<TestBlock> {
+            let header = sp_runtime::generic::Header::<u64, _>::new_from_number(block_number);
+            let mut slice = [0u8; 32];
+            slice[..8].copy_from_slice(&block_number.to_be_bytes());
+            // add random uuid to ensure uniqueness
+            slice[8..24].copy_from_slice(&Uuid::new_v4().to_u128_le().to_be_bytes());
+
+            let hash = sp_runtime::testing::H256::from(slice);
+            let summary = FinalizeSummary {
+                header,
+                finalized: vec![hash],
+                stale_heads: vec![],
+            };
+
+            let (tx, _rx) = sc_utils::mpsc::tracing_unbounded("mpsc_finality_notification", 999999);
+            FinalityNotification::<TestBlock>::from_summary(summary, tx)
+        }
+    }
+
+    pub async fn start_blockchain(
+        addr: SocketAddr,
+        block_duration: Duration,
+    ) -> Result<(), String> {
+        let zkp_service = server::TestBlockchainServer::default();
+        // TODO: insert some jobs into the hashmap inside the zkp service
+
+        let zkp_server = server::auth_server::AuthServer::new(zkp_service.clone())
+            .max_decoding_message_size(1024 * 1024 * 1024) // To allow large messages
+            .max_encoding_message_size(1024 * 1024 * 1024);
+
+        let ticker = async move {
+            loop {
+                tokio::time::sleep(block_duration).await;
+                let mut lock = zkp_service.latest_header.lock();
+                *lock += 1;
+            }
+        };
+
+        let server = tonic::transport::Server::builder()
+            .add_service(zkp_server)
+            .serve(addr);
+
+        tokio::select! {
+            res0 = server => {
+                res0.map_err(|err| err.to_string())?;
+            },
+            _ = ticker => {
+                Err("Ticker stopped")?;
+            }
+        }
+
+        Ok(())
+    }
+
+    fn generate_pub_key_and_priv_key_der<T: Into<String>>(subject_name: T) -> (Vec<u8>, Vec<u8>) {
+        let cert =
+            rcgen::generate_simple_self_signed([subject_name.into()]).expect("Should compile");
+        (
+            cert.serialize_der().expect("Should serialize"),
+            cert.serialize_private_key_der(),
+        )
+    }
+
+    const N: usize = 5;
+
+    #[tokio::test(flavor = "multi_thread")]
+    async fn test_zk_gadget() -> Result<(), Box<dyn Error>> {
+        const BLOCK_DURATION: Duration = Duration::from_millis(6000);
+
+        setup_log();
+
+        let server_addr: SocketAddr = "127.0.0.1:50051".parse()?;
+        let king_bind_addr_orig: SocketAddr = "127.0.0.1:50052".parse()?;
+        let (king_public_identity_der, king_private_identity_der) =
+            generate_pub_key_and_priv_key_der(king_bind_addr_orig.ip().to_string());
+        let blockchain_future = tokio::spawn(start_blockchain(server_addr, BLOCK_DURATION));
+        tokio::time::sleep(Duration::from_millis(100)).await; // wait for the server to come up
+
+        let (done_tx, mut done_rx) = tokio::sync::mpsc::unbounded_channel();
+        let zk_gadgets_futures = FuturesUnordered::new();
+
+        for party_id in 0..N {
+            let client = BlockchainClient::new(server_addr).await?;
+            let test_config = if party_id == 0 {
+                zk_gadget::ZkGadgetConfig {
+                    king_bind_addr: Some(king_bind_addr_orig),
+                    client_only_king_addr: None,
+                    id: party_id as _,
+                    n_parties: N,
+                    public_identity_der: king_public_identity_der.clone(),
+                    private_identity_der: king_private_identity_der.clone(),
+                    client_only_king_public_identity_der: None,
+                }
+            } else {
+                let (public_identity_der, private_identity_der) =
+                    generate_pub_key_and_priv_key_der("localhost");
+                zk_gadget::ZkGadgetConfig {
+                    king_bind_addr: None,
+                    client_only_king_addr: Some(king_bind_addr_orig),
+                    id: party_id as _,
+                    n_parties: N,
+                    public_identity_der,
+                    private_identity_der,
+                    client_only_king_public_identity_der: Some(king_public_identity_der.clone()),
+                }
+            };
+
+            let additional_parameters = AdditionalParams {
+                stop_tx: done_tx.clone(),
+            };
+
+            let zk_gadget_future = zk_gadget::run(
+                test_config,
+                client,
+                additional_parameters,
+                async_protocol_generator,
+            );
+            zk_gadgets_futures.push(Box::pin(zk_gadget_future));
+        }
+
+        let done_rx_future = async move {
+            let mut done_signals_received = 0;
+            while done_signals_received < N {
+                done_rx.recv().await.ok_or("Did not receive done signal")?;
+                log::info!("Received {}/{} done signals", done_signals_received, N);
+                done_signals_received += 1;
+            }
+
+            log::info!("Received ALL done signals");
+
+            Ok::<_, String>(())
+        };
+
+        let zk_gadget_future = zk_gadgets_futures.try_collect::<Vec<_>>();
+
+        tokio::select! {
+            res0 = blockchain_future => {
+                res0??;
+            },
+            res1 = zk_gadget_future => {
+                res1?;
+            },
+            res2 = done_rx_future => {
+                res2?;
+            }
+        }
+
+        Ok(())
+    }
+
+    #[derive(Clone)]
+    struct AdditionalParams {
+        pub stop_tx: UnboundedSender<()>,
+    }
+
+    fn async_protocol_generator(
+        params: ZkAsyncProtocolParameters<
+            AdditionalParams,
+            ZkNetworkService,
+            BlockchainClient,
+            TestBlock,
+        >,
+    ) -> BuiltExecutableJobWrapper<impl SendFuture<'static, Result<(), JobError>>> {
+        JobBuilder::default().build(async move {
+            if params.party_id == 0 {
+                // Receive N-1 messages from the other parties
+                for party_id in 0..N {
+                    let party_id = party_id as u32;
+                    if party_id != params.party_id {
+                        let _message = params
+                            .recv_from(party_id, MultiplexedStreamID::Zero)
+                            .await
+                            .expect("Should receive protocol message");
+                    }
+                }
+
+                for party_id in 0..N {
+                    let party_id = party_id as u32;
+                    if party_id != params.party_id {
+                        params
+                            .send_to(party_id, Default::default(), MultiplexedStreamID::Zero)
+                            .await
+                            .expect("Should send");
+                    }
+                }
+            } else {
+                params
+                    .send_to(0, Default::default(), MultiplexedStreamID::Zero)
+                    .await
+                    .expect("Should send");
+                let _message_from_king = params
+                    .recv_from(0, MultiplexedStreamID::Zero)
+                    .await
+                    .expect("Should receive protocol message");
+            }
+
+            params
+                .extra_parameters
+                .stop_tx
+                .send(())
+                .expect("Should send");
+            Ok(())
+        })
+    }
+}