fix(notary): implement timeout for notarization (#639)

* Add timeout.

* Fmt.

* Fix grammar.

* Move limit to config.

* Remove extra space.

---------

Co-authored-by: yuroitaki <>

crates/notary/server/config/config.yaml

@@ -14,6 +14,7 @@ server:
 notarization:
   max_sent_data: 4096
   max_recv_data: 16384
+  timeout: 1800
 
 tls:
   enabled: true

crates/notary/server/src/config.rs

@@ -30,6 +30,9 @@ pub struct NotarizationProperties {
     pub max_sent_data: usize,
     /// Global limit for maximum number of bytes that can be received
     pub max_recv_data: usize,
+    /// Number of seconds before notarization timeouts to prevent unreleased
+    /// memory
+    pub timeout: u64,
 }
 
 #[derive(Clone, Debug, Deserialize, Default)]

crates/notary/server/src/service.rs

@@ -2,19 +2,22 @@ pub mod axum_websocket;
 pub mod tcp;
 pub mod websocket;
 
-use std::sync::Arc;
-
 use async_trait::async_trait;
 use axum::{
     extract::{rejection::JsonRejection, FromRequestParts, Query, State},
     http::{header, request::Parts, StatusCode},
     response::{IntoResponse, Json, Response},
 };
 use axum_macros::debug_handler;
+use eyre::eyre;
+use std::time::Duration;
 use tlsn_common::config::ProtocolConfigValidator;
-use tlsn_core::{attestation::AttestationConfig, CryptoProvider};
+use tlsn_core::attestation::AttestationConfig;
 use tlsn_verifier::{Verifier, VerifierConfig};
-use tokio::io::{AsyncRead, AsyncWrite};
+use tokio::{
+    io::{AsyncRead, AsyncWrite},
+    time::timeout,
+};
 use tokio_util::compat::TokioAsyncReadCompatExt;
 use tracing::{debug, error, info, trace};
 use uuid::Uuid;
@@ -180,13 +183,13 @@ pub async fn initialize(
 /// Run the notarization
 pub async fn notary_service<T: AsyncWrite + AsyncRead + Send + Unpin + 'static>(
     socket: T,
-    crypto_provider: Arc<CryptoProvider>,
+    notary_globals: NotaryGlobals,
     session_id: &str,
-    max_sent_data: usize,
-    max_recv_data: usize,
 ) -> Result<(), NotaryServerError> {
     debug!(?session_id, "Starting notarization...");
 
+    let crypto_provider = notary_globals.crypto_provider.clone();
+
     let att_config = AttestationConfig::builder()
         .supported_signature_algs(Vec::from_iter(crypto_provider.signer.supported_algs()))
         .build()
@@ -195,16 +198,19 @@ pub async fn notary_service<T: AsyncWrite + AsyncRead + Send + Unpin + 'static>(
     let config = VerifierConfig::builder()
         .protocol_config_validator(
             ProtocolConfigValidator::builder()
-                .max_sent_data(max_sent_data)
-                .max_recv_data(max_recv_data)
+                .max_sent_data(notary_globals.notarization_config.max_sent_data)
+                .max_recv_data(notary_globals.notarization_config.max_recv_data)
                 .build()?,
         )
         .crypto_provider(crypto_provider)
         .build()?;
 
-    Verifier::new(config)
-        .notarize(socket.compat(), &att_config)
-        .await?;
+    timeout(
+        Duration::from_secs(notary_globals.notarization_config.timeout),
+        Verifier::new(config).notarize(socket.compat(), &att_config),
+    )
+    .await
+    .map_err(|_| eyre!("Timeout reached before notarization completes"))??;
 
     Ok(())
 }

crates/notary/server/src/service/tcp.rs

@@ -87,15 +87,7 @@ pub async fn tcp_notarize(
     session_id: String,
 ) {
     debug!(?session_id, "Upgraded to tcp connection");
-    match notary_service(
-        stream,
-        notary_globals.crypto_provider.clone(),
-        &session_id,
-        notary_globals.notarization_config.max_sent_data,
-        notary_globals.notarization_config.max_recv_data,
-    )
-    .await
-    {
+    match notary_service(stream, notary_globals, &session_id).await {
         Ok(_) => {
             info!(?session_id, "Successful notarization using tcp!");
         }

crates/notary/server/src/service/websocket.rs

@@ -16,15 +16,7 @@ pub async fn websocket_notarize(
     // Wrap the websocket in WsStream so that we have AsyncRead and AsyncWrite
     // implemented
     let stream = WsStream::new(socket.into_inner());
-    match notary_service(
-        stream,
-        notary_globals.crypto_provider.clone(),
-        &session_id,
-        notary_globals.notarization_config.max_sent_data,
-        notary_globals.notarization_config.max_recv_data,
-    )
-    .await
-    {
+    match notary_service(stream, notary_globals, &session_id).await {
         Ok(_) => {
             info!(?session_id, "Successful notarization using websocket!");
         }

crates/notary/tests-integration/tests/notary.rs

@@ -48,6 +48,7 @@ fn get_server_config(port: u16, tls_enabled: bool, auth_enabled: bool) -> Notary
         notarization: NotarizationProperties {
             max_sent_data: 1 << 13,
             max_recv_data: 1 << 14,
+            timeout: 1800,
         },
         tls: TLSProperties {
             enabled: tls_enabled,
@@ -191,8 +192,7 @@ async fn test_tcp_prover<S: AsyncWrite + AsyncRead + Send + Unpin + 'static>(
         .build()
         .unwrap();
 
-    // Prover config using the session_id returned from calling /session endpoint in
-    // notary client.
+    // Set up prover config.
     let prover_config = ProverConfig::builder()
         .server_name(SERVER_DOMAIN)
         .protocol_config(protocol_config)
@@ -386,7 +386,7 @@ async fn test_websocket_prover() {
         .build()
         .unwrap();
 
-    // Basic default prover config — use the responded session id from notary server
+    // Set up prover config.
     let prover_config = ProverConfig::builder()
         .server_name(SERVER_DOMAIN)
         .protocol_config(protocol_config)