From 55f0ff2e447b71d61645c6865ba9d013b98f220a Mon Sep 17 00:00:00 2001 From: th4s Date: Wed, 8 Nov 2023 11:41:13 +0100 Subject: [PATCH] Fix prover closing connection too early --- tlsn/tlsn-prover/src/tls/notarize.rs | 4 +++- tlsn/tlsn-prover/src/tls/prove.rs | 9 ++++++--- 2 files changed, 9 insertions(+), 4 deletions(-) diff --git a/tlsn/tlsn-prover/src/tls/notarize.rs b/tlsn/tlsn-prover/src/tls/notarize.rs index 2824f271be..f0caed6b97 100644 --- a/tlsn/tlsn-prover/src/tls/notarize.rs +++ b/tlsn/tlsn-prover/src/tls/notarize.rs @@ -90,8 +90,10 @@ impl Prover { let (notary_encoder_seed, SignedSessionHeader { header, signature }) = futures::select_biased! { res = notarize_fut => res?, _ = ot_fut => return Err(OTShutdownError)?, - _ = mux_fut => return Err(std::io::Error::from(std::io::ErrorKind::UnexpectedEof))?, + _ = &mut mux_fut => return Err(std::io::Error::from(std::io::ErrorKind::UnexpectedEof))?, }; + // Wait for the notary to correctly close the connection + mux_fut.await?; // Check the header is consistent with the Prover's view header diff --git a/tlsn/tlsn-prover/src/tls/prove.rs b/tlsn/tlsn-prover/src/tls/prove.rs index 51a13b4999..ff3460172a 100644 --- a/tlsn/tlsn-prover/src/tls/prove.rs +++ b/tlsn/tlsn-prover/src/tls/prove.rs @@ -159,7 +159,7 @@ impl Prover { handshake_decommitment, }; - let mut verify_fut = Box::pin(async move { + let mut finalize_fut = Box::pin(async move { let mut channel = verify_mux.get_channel("finalize").await?; _ = vm @@ -183,11 +183,14 @@ impl Prover { .fuse(); futures::select_biased! { - res = verify_fut => res?, + res = finalize_fut => res?, _ = ot_fut => return Err(OTShutdownError)?, - _ = mux_fut => return Err(std::io::Error::from(std::io::ErrorKind::UnexpectedEof))?, + _ = &mut mux_fut => return Err(std::io::Error::from(std::io::ErrorKind::UnexpectedEof))?, }; + // We need to wait for the verifier to correctly close the connection. Otherwise the prover + // would rush ahead and close the connection before the verifier has finished. + mux_fut.await?; Ok(()) } }