From 2c045e5de7ebe832abf87b775b4eddbcab5885b8 Mon Sep 17 00:00:00 2001 From: yuroitaki <25913766+yuroitaki@users.noreply.github.com> Date: Tue, 22 Oct 2024 15:08:19 +0800 Subject: [PATCH] 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 | 1 + crates/notary/server/src/config.rs | 3 ++ crates/notary/server/src/service.rs | 30 +++++++++++-------- crates/notary/server/src/service/tcp.rs | 10 +------ crates/notary/server/src/service/websocket.rs | 10 +------ .../notary/tests-integration/tests/notary.rs | 6 ++-- 6 files changed, 27 insertions(+), 33 deletions(-) diff --git a/crates/notary/server/config/config.yaml b/crates/notary/server/config/config.yaml index 5916b9bdf1..ac680e6021 100644 --- a/crates/notary/server/config/config.yaml +++ b/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 diff --git a/crates/notary/server/src/config.rs b/crates/notary/server/src/config.rs index c2f3963182..b5181fb699 100644 --- a/crates/notary/server/src/config.rs +++ b/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)] diff --git a/crates/notary/server/src/service.rs b/crates/notary/server/src/service.rs index dd93d36157..e58deed74f 100644 --- a/crates/notary/server/src/service.rs +++ b/crates/notary/server/src/service.rs @@ -2,8 +2,6 @@ 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}, @@ -11,10 +9,15 @@ use axum::{ 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( socket: T, - crypto_provider: Arc, + 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( 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(()) } diff --git a/crates/notary/server/src/service/tcp.rs b/crates/notary/server/src/service/tcp.rs index e68f79370d..e2b433de6b 100644 --- a/crates/notary/server/src/service/tcp.rs +++ b/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!"); } diff --git a/crates/notary/server/src/service/websocket.rs b/crates/notary/server/src/service/websocket.rs index f0f1ea5163..9c53b45f5d 100644 --- a/crates/notary/server/src/service/websocket.rs +++ b/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!"); } diff --git a/crates/notary/tests-integration/tests/notary.rs b/crates/notary/tests-integration/tests/notary.rs index da067cd58b..4476e3b8ab 100644 --- a/crates/notary/tests-integration/tests/notary.rs +++ b/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( .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)