feat(sdk): reclaim session IO after completion

Retain the session driver result so native callers can recover the underlying stream and WASM callers can await exclusive reuse of their injected IoChannel. Exercise the handoff by exchanging application bytes after the SDK protocol completes.

Assisted-by: GPT-5
Signed-off-by: Wondertan <hlibwondertan@gmail.com>

Author: Wondertan

crates/harness/executor/src/provider.rs

@@ -83,10 +83,29 @@ mod wasm {
     use crate::io::Io;
     use anyhow::{Result, anyhow};
     use std::time::Duration;
+    use wasm_bindgen::prelude::*;
+    use wasm_bindgen_futures::JsFuture;
 
     const CHECK_WS_OPEN_DELAY_MS: usize = 50;
     const MAX_RETRIES: usize = 50;
 
+    #[wasm_bindgen]
+    extern "C" {
+        type JsIoChannel;
+
+        #[wasm_bindgen(js_namespace = globalThis, js_name = connectIoChannel)]
+        fn connect_io_channel(url: String) -> js_sys::Promise;
+
+        #[wasm_bindgen(method, js_name = isOpen)]
+        fn is_open(this: &JsIoChannel) -> bool;
+    }
+
+    async fn connect_js_io(url: String) -> Result<JsValue> {
+        JsFuture::from(connect_io_channel(url))
+            .await
+            .map_err(|error| anyhow!("failed to connect JS IO: {error:?}"))
+    }
+
     impl IoProvider {
         /// Provides a connection to the server.
         pub async fn provide_server_io(&self) -> Result<impl Io> {
@@ -102,6 +121,18 @@ mod wasm {
             Ok(io.into_io())
         }
 
+        /// Provides a JavaScript `IoChannel` backed by a real WebSocket.
+        pub async fn provide_server_js_io(&self) -> Result<JsValue> {
+            connect_js_io(format!(
+                "ws://{}:{}/tcp?addr={}%3A{}",
+                &self.config.app_proxy.0,
+                self.config.app_proxy.1,
+                &self.config.app.0,
+                self.config.app.1,
+            ))
+            .await
+        }
+
         /// Provides a connection to the verifier.
         pub async fn provide_proto_io(&self) -> Result<impl Io> {
             let url = format!(
@@ -135,5 +166,30 @@ mod wasm {
 
             Ok(io.into_io())
         }
+
+        /// Provides a JavaScript `IoChannel` backed by the protocol WebSocket.
+        pub async fn provide_proto_js_io(&self) -> Result<JsValue> {
+            let url = format!(
+                "ws://{}:{}/tcp?addr={}%3A{}",
+                &self.config.proto_proxy.0,
+                self.config.proto_proxy.1,
+                &self.config.proto_1.0,
+                self.config.proto_1.1,
+            );
+            let mut retries = 0;
+
+            loop {
+                let io = connect_js_io(url.clone()).await?;
+                std::thread::sleep(Duration::from_millis(CHECK_WS_OPEN_DELAY_MS as u64));
+                if io.unchecked_ref::<JsIoChannel>().is_open() {
+                    return Ok(io);
+                }
+
+                retries += 1;
+                if retries > MAX_RETRIES {
+                    return Err(anyhow!("verifier did not accept connection"));
+                }
+            }
+        }
     }
 }

crates/harness/executor/test_plugins/sdk_core.rs

@@ -1,6 +1,7 @@
-use tlsn_sdk_core::{
-    HttpRequest, NetworkSetting, ProverConfig, Reveal, SdkProver, SdkVerifier, VerifierConfig,
-};
+use futures::io::{AsyncReadExt, AsyncWriteExt};
+#[cfg(not(target_arch = "wasm32"))]
+use tlsn_sdk_core::{HttpRequest, NetworkSetting, ProverConfig, Reveal, SdkProver};
+use tlsn_sdk_core::{SdkVerifier, VerifierConfig};
 use tlsn_server_fixture_certs::{CA_CERT_DER, SERVER_DOMAIN};
 
 use crate::IoProvider;
@@ -13,6 +14,15 @@ const MAX_RECV_DATA: usize = 1 << 11;
 crate::test!("sdk_core", prover, verifier);
 
 async fn prover(provider: &IoProvider) {
+    #[cfg(target_arch = "wasm32")]
+    return prover_wasm(provider).await;
+
+    #[cfg(not(target_arch = "wasm32"))]
+    prover_core(provider).await;
+}
+
+#[cfg(not(target_arch = "wasm32"))]
+async fn prover_core(provider: &IoProvider) {
     let config = ProverConfig::builder(SERVER_DOMAIN)
         .max_sent_data(MAX_SENT_DATA)
         .max_recv_data(MAX_RECV_DATA)
@@ -60,6 +70,100 @@ async fn prover(provider: &IoProvider) {
         .unwrap();
 
     assert!(prover.is_complete());
+
+    let mut io = prover.finish().await.unwrap();
+    io.write_all(b"prover-finished").await.unwrap();
+    let mut response = [0; 17];
+    io.read_exact(&mut response).await.unwrap();
+    assert_eq!(&response, b"verifier-finished");
+}
+
+#[cfg(target_arch = "wasm32")]
+async fn prover_wasm(provider: &IoProvider) {
+    use js_sys::{Promise, Uint8Array};
+    use std::collections::HashMap;
+    use tlsn_wasm::{
+        prover::{JsProver, ProverConfig},
+        types::{HttpRequest, Method, Reveal},
+    };
+    use wasm_bindgen::{JsCast, prelude::*};
+    use wasm_bindgen_futures::JsFuture;
+
+    #[wasm_bindgen]
+    extern "C" {
+        type TestIoChannel;
+
+        #[wasm_bindgen(method)]
+        fn read(this: &TestIoChannel) -> Promise;
+
+        #[wasm_bindgen(method)]
+        fn write(this: &TestIoChannel, data: &Uint8Array) -> Promise;
+    }
+
+    let config: ProverConfig = serde_json::from_value(serde_json::json!({
+        "server_name": SERVER_DOMAIN,
+        "mode": "Mpc",
+        "max_sent_data": MAX_SENT_DATA,
+        "max_sent_records": null,
+        "max_recv_data_online": null,
+        "max_recv_data": MAX_RECV_DATA,
+        "max_recv_records_online": null,
+        "defer_decryption_from_start": true,
+        "network": "Latency",
+        "client_auth": null,
+        "root_certs": [CA_CERT_DER],
+    }))
+    .unwrap();
+    let mut prover = JsProver::new(config).unwrap();
+
+    let proto_io = provider.provide_proto_js_io().await.unwrap();
+    let retained_io = proto_io.unchecked_ref::<TestIoChannel>();
+    prover
+        .setup(proto_io.clone().unchecked_into())
+        .await
+        .unwrap();
+
+    let server_io = provider.provide_server_js_io().await.unwrap();
+    let response = prover
+        .send_request(
+            Some(server_io.unchecked_into()),
+            HttpRequest {
+                uri: format!(
+                    "https://{}/bytes?size={}",
+                    SERVER_DOMAIN,
+                    MAX_RECV_DATA - 256
+                ),
+                method: Method::GET,
+                headers: HashMap::from([
+                    ("Host".to_string(), SERVER_DOMAIN.as_bytes().to_vec()),
+                    ("Connection".to_string(), b"close".to_vec()),
+                ]),
+                body: None,
+            },
+        )
+        .await
+        .unwrap();
+    assert_eq!(response.status, 200);
+
+    let transcript = prover.transcript().unwrap();
+    prover
+        .reveal(
+            Reveal {
+                sent: vec![0..transcript.sent.len() - 1],
+                recv: vec![2..transcript.recv.len()],
+                server_identity: true,
+            },
+            None,
+        )
+        .await
+        .unwrap();
+    prover.finish().await.unwrap();
+
+    JsFuture::from(retained_io.write(&Uint8Array::from(b"prover-finished".as_slice())))
+        .await
+        .unwrap();
+    let response = JsFuture::from(retained_io.read()).await.unwrap();
+    assert_eq!(Uint8Array::new(&response).to_vec(), b"verifier-finished");
 }
 
 async fn verifier(provider: &IoProvider) {
@@ -86,4 +190,10 @@ async fn verifier(provider: &IoProvider) {
     assert_eq!(output.server_name.as_deref(), Some(SERVER_DOMAIN));
     assert!(output.transcript.is_some());
     assert!(verifier.is_complete());
+
+    let mut io = verifier.finish().await.unwrap();
+    let mut request = [0; 15];
+    io.read_exact(&mut request).await.unwrap();
+    assert_eq!(&request, b"prover-finished");
+    io.write_all(b"verifier-finished").await.unwrap();
 }

crates/harness/static/executor.js

@@ -1,6 +1,67 @@
 import * as Comlink from "./comlink.mjs";
 import initWasm, * as wasm from "./generated/harness_executor.js";
 
+class WebSocketIoChannel {
+    constructor(socket) {
+        this.socket = socket;
+        this.queue = [];
+        this.reader = null;
+
+        socket.binaryType = "arraybuffer";
+        socket.onmessage = ({ data }) => {
+            const bytes = new Uint8Array(data);
+            if (this.reader) {
+                const reader = this.reader;
+                this.reader = null;
+                reader.resolve(bytes);
+            } else {
+                this.queue.push(bytes);
+            }
+        };
+        socket.onclose = () => {
+            if (this.reader) {
+                const reader = this.reader;
+                this.reader = null;
+                reader.resolve(null);
+            }
+        };
+        socket.onerror = () => {
+            if (this.reader) {
+                const reader = this.reader;
+                this.reader = null;
+                reader.reject(new Error("WebSocket error"));
+            }
+        };
+    }
+
+    read() {
+        if (this.queue.length) return Promise.resolve(this.queue.shift());
+        if (this.socket.readyState === WebSocket.CLOSED) return Promise.resolve(null);
+        if (this.reader) return Promise.reject(new Error("concurrent read"));
+        return new Promise((resolve, reject) => { this.reader = { resolve, reject }; });
+    }
+
+    write(data) {
+        this.socket.send(data);
+        return Promise.resolve();
+    }
+
+    close() {
+        this.socket.close();
+        return Promise.resolve();
+    }
+
+    isOpen() {
+        return this.socket.readyState === WebSocket.OPEN;
+    }
+}
+
+globalThis.connectIoChannel = (url) => new Promise((resolve, reject) => {
+    const socket = new WebSocket(url);
+    socket.onopen = () => resolve(new WebSocketIoChannel(socket));
+    socket.onerror = () => reject(new Error(`failed to connect to ${url}`));
+});
+
 class Executor {
     executor;
 

crates/sdk-core/src/prover.rs

@@ -33,6 +33,7 @@ use crate::{
 pub struct SdkProver {
     config: ProverConfig,
     state: State,
+    driver_task: Option<crate::spawn::DriverTask>,
 }
 
 #[allow(clippy::large_enum_variant)]
@@ -91,6 +92,7 @@ impl SdkProver {
         Ok(SdkProver {
             config,
             state: State::Initialized,
+            driver_task: None,
         })
     }
 
@@ -111,15 +113,9 @@ impl SdkProver {
 
         info!("connecting to verifier");
 
-        let session = Session::new(verifier_io);
+        let session = Session::new(Box::new(verifier_io) as crate::spawn::BoxIo);
         let (driver, mut handle) = session.split();
-
-        crate::spawn::spawn(async move {
-            match driver.await {
-                Ok(_io) => tracing::warn!("session driver completed (mux closed)"),
-                Err(e) => tracing::error!("session driver error: {e}"),
-            }
-        });
+        self.driver_task = Some(crate::spawn::DriverTask::spawn(driver));
 
         let prover_config = tlsn::config::prover::ProverConfig::builder().build()?;
         let prover = handle.new_prover(prover_config)?;
@@ -444,6 +440,23 @@ impl SdkProver {
     pub fn is_complete(&self) -> bool {
         matches!(self.state, State::Complete)
     }
+
+    /// Waits for the session driver to stop and returns the underlying IO.
+    ///
+    /// This must be called after [`reveal`](Self::reveal). The returned stream
+    /// is no longer read from or written to by TLSNotary and can be reused by
+    /// the application.
+    pub async fn finish(&mut self) -> Result<Box<dyn Io>> {
+        if !self.is_complete() {
+            return Err(SdkError::invalid_state("prover is not complete"));
+        }
+
+        self.driver_task
+            .take()
+            .ok_or_else(|| SdkError::invalid_state("prover session already finished"))?
+            .finish()
+            .await
+    }
 }
 
 fn build_reveal_output(output: ProverOutput) -> Result<RevealOutput> {

crates/sdk-core/src/spawn.rs

@@ -2,6 +2,40 @@
 
 use std::future::Future;
 
+use futures::channel::oneshot;
+use tlsn::SessionDriver;
+
+use crate::{
+    error::{Result, SdkError},
+    io::Io,
+};
+
+pub(crate) type BoxIo = Box<dyn Io>;
+
+pub(crate) struct DriverTask(oneshot::Receiver<tlsn::Result<BoxIo>>);
+
+impl DriverTask {
+    pub(crate) fn spawn(driver: SessionDriver<BoxIo>) -> Self {
+        let (sender, receiver) = oneshot::channel();
+        spawn(async move {
+            let result = driver.await;
+            match &result {
+                Ok(_) => tracing::warn!("session driver completed (mux closed)"),
+                Err(error) => tracing::error!("session driver error: {error}"),
+            }
+            let _ = sender.send(result);
+        });
+        Self(receiver)
+    }
+
+    pub(crate) async fn finish(self) -> Result<BoxIo> {
+        self.0
+            .await
+            .map_err(|_| SdkError::internal("session driver task dropped"))?
+            .map_err(Into::into)
+    }
+}
+
 /// Spawns a future on the appropriate runtime.
 #[cfg(feature = "wasm")]
 pub(crate) fn spawn(future: impl Future<Output = ()> + 'static) {