diff --git a/crates/harness/executor/src/bench/verifier.rs b/crates/harness/executor/src/bench/verifier.rs index 985048d33..3a8c0b890 100644 --- a/crates/harness/executor/src/bench/verifier.rs +++ b/crates/harness/executor/src/bench/verifier.rs @@ -1,9 +1,10 @@ use anyhow::Result; +use futures::TryFutureExt; use harness_core::bench::Bench; use tlsn::{ config::verifier::VerifierConfig, - verifier::Verifier, + verifier::{Verifier, VerifierError}, webpki::{CertificateDer, RootCertStore}, }; use tlsn_server_fixture_certs::CA_CERT_DER; @@ -19,15 +20,16 @@ pub async fn bench_verifier(provider: &IoProvider, _config: &Bench) -> Result<() .build()?, ); - let verifier = verifier - .commit(provider.provide_proto_io().await?) - .await? - .accept() - .await? - .run() - .await?; - let (_, verifier) = verifier.verify().await?.accept().await?; - verifier.close().await?; + let mut prover_io = provider.provide_proto_io().await?; + let (mpc_conn, verifier) = verifier.commit_with(&mut prover_io).await?; + let mpc_fut = mpc_conn.into_future(prover_io).map_err(VerifierError::from); + let verifier = async { + let verifier = verifier.accept().await?.run().await?; + let (_, verifier) = verifier.verify().await?.accept().await?; + verifier.close().await + }; + + futures::try_join!(mpc_fut, verifier); Ok(()) } diff --git a/crates/tlsn/src/conn/mpc.rs b/crates/tlsn/src/conn/mpc.rs index 97510e7e5..45bf6111e 100644 --- a/crates/tlsn/src/conn/mpc.rs +++ b/crates/tlsn/src/conn/mpc.rs @@ -170,7 +170,7 @@ where if read > 0 { this.write_buf.extend(&tmp_write_buf[..read]); } else { - return this.duplex.as_mut().poll_close(cx); + return this.socket.as_mut().poll_close(cx); } } diff --git a/crates/tlsn/src/verifier.rs b/crates/tlsn/src/verifier.rs index a57289aae..5b18b7ab1 100644 --- a/crates/tlsn/src/verifier.rs +++ b/crates/tlsn/src/verifier.rs @@ -68,7 +68,7 @@ impl Verifier { /// /// * `socket` - The socket to the prover. #[instrument(parent = &self.span, level = "info", skip_all, err)] - pub async fn commit( + pub fn commit( self, ) -> Result, VerifierError>, VerifierError> { let (duplex_a, duplex_b) = futures_plex::duplex(BUF_CAP); diff --git a/crates/tlsn/tests/test.rs b/crates/tlsn/tests/test.rs index 365d573e5..183f4e56b 100644 --- a/crates/tlsn/tests/test.rs +++ b/crates/tlsn/tests/test.rs @@ -1,4 +1,4 @@ -use futures::{AsyncReadExt, AsyncWriteExt}; +use futures::{AsyncReadExt, AsyncWriteExt, TryFutureExt}; use rangeset::RangeSet; use tlsn::{ config::{ @@ -15,7 +15,7 @@ use tlsn::{ Direction, Transcript, TranscriptCommitConfig, TranscriptCommitment, TranscriptCommitmentKind, TranscriptSecret, }, - verifier::{Verifier, VerifierOutput}, + verifier::{Verifier, VerifierError, VerifierOutput}, webpki::{CertificateDer, RootCertStore}, }; use tlsn_core::ProverOutput; @@ -215,7 +215,7 @@ async fn prover( #[instrument(skip(socket))] async fn verifier( - socket: T, + mut socket: T, ) -> VerifierOutput { let verifier = Verifier::new( VerifierConfig::builder() @@ -226,19 +226,19 @@ async fn verifier( .unwrap(), ); - let verifier = verifier - .commit(socket.compat()) - .await - .unwrap() - .accept() - .await - .unwrap() - .run() - .await - .unwrap(); + let (mpc_conn, verifier) = verifier.commit_with((&mut socket).compat()).await.unwrap(); - let (output, verifier) = verifier.verify().await.unwrap().accept().await.unwrap(); - verifier.close().await.unwrap(); + let mpc_fut = mpc_conn + .into_future(socket.compat()) + .map_err(VerifierError::from); + let verifier = async { + let verifier = verifier.accept().await?.run().await?; + let (output, verifier) = verifier.verify().await?.accept().await?; + verifier.close().await?; + Ok(output) + }; + + let (_, output) = futures::try_join!(mpc_fut, verifier).unwrap(); output }