feat(notary): add concurrency limit (#770)

* feat(notary): add concurrency limit

* switch to 503 status code

* remove test-api feature

* improve naming and comments

* set default concurrency to 32

Author: themighty1

crates/notary/client/src/client.rs

@@ -4,7 +4,11 @@
 //! subsequent requests for notarization.
 
 use http_body_util::{BodyExt as _, Either, Empty, Full};
-use hyper::{body::Bytes, client::conn::http1::Parts, Request, StatusCode};
+use hyper::{
+    body::{Bytes, Incoming},
+    client::conn::http1::Parts,
+    Request, Response, StatusCode,
+};
 use hyper_util::rt::TokioIo;
 use notary_server::{ClientType, NotarizationSessionRequest, NotarizationSessionResponse};
 use std::{
@@ -16,6 +20,7 @@ use std::{
 use tokio::{
     io::{AsyncRead, AsyncWrite, ReadBuf},
     net::TcpStream,
+    time::{sleep, timeout, Duration},
 };
 use tokio_rustls::{
     client::TlsStream,
@@ -129,6 +134,14 @@ pub struct NotaryClient {
     /// in notary server.
     #[builder(setter(into, strip_option), default)]
     api_key: Option<String>,
+    /// The duration of notarization request timeout in seconds.
+    #[builder(default = "60")]
+    request_timeout: usize,
+    /// The number of seconds to wait between notarization request retries.
+    ///
+    /// By default uses the value suggested by the server.
+    #[builder(default = "None")]
+    request_retry_override: Option<u64>,
 }
 
 impl NotaryClientBuilder {
@@ -355,15 +368,56 @@ impl NotaryClient {
 
             debug!("Sending notarization request: {:?}", notarization_request);
 
-            let notarization_response = notary_request_sender
-                .send_request(notarization_request)
-                .await
-                .map_err(|err| {
-                    error!("Failed to send http request for notarization");
-                    ClientError::new(ErrorKind::Http, Some(Box::new(err)))
-                })?;
+            let notarize_with_retry_fut = async {
+                loop {
+                    let notarization_response = notary_request_sender
+                        .send_request(notarization_request.clone())
+                        .await
+                        .map_err(|err| {
+                            error!("Failed to send http request for notarization");
+                            ClientError::new(ErrorKind::Http, Some(Box::new(err)))
+                        })?;
+
+                    if notarization_response.status() == StatusCode::SWITCHING_PROTOCOLS {
+                        return Ok::<Response<Incoming>, ClientError>(notarization_response);
+                    } else if notarization_response.status() == StatusCode::SERVICE_UNAVAILABLE {
+                        let retry_after = self
+                            .request_retry_override
+                            .unwrap_or(parse_retry_after(&notarization_response)?);
+
+                        debug!("Retrying notarization request in {:?}", retry_after);
+
+                        sleep(Duration::from_secs(retry_after)).await;
+                    } else {
+                        return Err(ClientError::new(
+                            ErrorKind::Internal,
+                            Some(
+                                format!(
+                                    "Server sent unexpected status code {:?}",
+                                    notarization_response.status()
+                                )
+                                .into(),
+                            ),
+                        ));
+                    }
+                }
+            };
+
+            let notarization_response = timeout(
+                Duration::from_secs(self.request_timeout as u64),
+                notarize_with_retry_fut,
+            )
+            .await
+            .map_err(|_| {
+                ClientError::new(
+                    ErrorKind::Internal,
+                    Some(
+                        "Timed out while waiting for server to accept notarization request".into(),
+                    ),
+                )
+            })??;
 
-            debug!("Sent notarization request");
+            debug!("Notarization request was accepted by the server");
 
             if notarization_response.status() != StatusCode::SWITCHING_PROTOCOLS {
                 return Err(ClientError::new(
@@ -388,6 +442,17 @@ impl NotaryClient {
 
         Ok((notary_socket.into_inner(), session_id))
     }
+
+    /// Sets notarization request timeout duration in seconds.
+    pub fn request_timeout(&mut self, timeout: usize) {
+        self.request_timeout = timeout;
+    }
+
+    /// Sets the number of seconds to wait between notarization request
+    /// retries.
+    pub fn request_retry_override(&mut self, seconds: u64) {
+        self.request_retry_override = Some(seconds);
+    }
 }
 
 /// Default root store using mozilla certs.
@@ -403,3 +468,34 @@ fn default_root_store() -> RootCertStore {
 
     root_store
 }
+
+// Attempts to parse the value of the "Retry-After" header from the given
+// `response`.
+fn parse_retry_after(response: &Response<Incoming>) -> Result<u64, ClientError> {
+    let seconds = match response.headers().get("Retry-After") {
+        Some(value) => {
+            let value_str = value.to_str().map_err(|err| {
+                ClientError::new(
+                    ErrorKind::Internal,
+                    Some(format!("Invalid Retry-After header: {}", err).into()),
+                )
+            })?;
+
+            let seconds: u64 = value_str.parse().map_err(|err| {
+                ClientError::new(
+                    ErrorKind::Internal,
+                    Some(format!("Could not parse Retry-After header as number: {}", err).into()),
+                )
+            })?;
+            seconds
+        }
+        None => {
+            return Err(ClientError::new(
+                ErrorKind::Internal,
+                Some("The expected Retry-After header was not found in server response".into()),
+            ));
+        }
+    };
+
+    Ok(seconds)
+}

crates/notary/server/config/config.yaml

@@ -44,3 +44,5 @@ logging:
 authorization:
   enabled: false
   whitelist_csv_path: "./fixture/auth/whitelist.csv"
+
+concurrency: 32

crates/notary/server/src/config.rs

@@ -1,6 +1,6 @@
 use serde::Deserialize;
 
-#[derive(Clone, Debug, Deserialize, Default)]
+#[derive(Clone, Debug, Deserialize)]
 pub struct NotaryServerProperties {
     /// Name and address of the notary server
     pub server: ServerProperties,
@@ -14,6 +14,22 @@ pub struct NotaryServerProperties {
     pub logging: LoggingProperties,
     /// Setting for authorization
     pub authorization: AuthorizationProperties,
+    /// The maximum number of concurrent notarization sessions
+    pub concurrency: usize,
+}
+
+impl Default for NotaryServerProperties {
+    fn default() -> Self {
+        Self {
+            server: ServerProperties::default(),
+            notarization: NotarizationProperties::default(),
+            tls: TLSProperties::default(),
+            notary_key: NotarySigningKeyProperties::default(),
+            logging: LoggingProperties::default(),
+            authorization: AuthorizationProperties::default(),
+            concurrency: 32,
+        }
+    }
 }
 
 #[derive(Clone, Debug, Deserialize, Default)]

crates/notary/server/src/domain/notary.rs

@@ -4,6 +4,7 @@ use std::{
     sync::{Arc, Mutex},
 };
 use tlsn_core::CryptoProvider;
+use tokio::sync::Semaphore;
 
 use crate::{config::NotarizationProperties, domain::auth::AuthorizationWhitelistRecord};
 
@@ -53,19 +54,23 @@ pub struct NotaryGlobals {
     pub store: Arc<Mutex<HashMap<String, ()>>>,
     /// Whitelist of API keys for authorization purpose
     pub authorization_whitelist: Option<Arc<Mutex<HashMap<String, AuthorizationWhitelistRecord>>>>,
+    /// A semaphore to acquire a permit for notarization
+    pub semaphore: Arc<Semaphore>,
 }
 
 impl NotaryGlobals {
     pub fn new(
         crypto_provider: Arc<CryptoProvider>,
         notarization_config: NotarizationProperties,
         authorization_whitelist: Option<Arc<Mutex<HashMap<String, AuthorizationWhitelistRecord>>>>,
+        semaphore: Arc<Semaphore>,
     ) -> Self {
         Self {
             crypto_provider,
             notarization_config,
             store: Default::default(),
             authorization_whitelist,
+            semaphore,
         }
     }
 }

crates/notary/server/src/server.rs

@@ -49,6 +49,8 @@ use crate::{
 #[cfg(feature = "tee_quote")]
 use crate::tee::{generate_ephemeral_keypair, quote};
 
+use tokio::sync::Semaphore;
+
 /// Start a TCP server (with or without TLS) to accept notarization request for
 /// both TCP and WebSocket clients
 #[tracing::instrument(skip(config))]
@@ -114,6 +116,7 @@ pub async fn run_server(config: &NotaryServerProperties) -> Result<(), NotarySer
         Arc::new(crypto_provider),
         config.notarization.clone(),
         authorization_whitelist,
+        Arc::new(Semaphore::new(config.concurrency)),
     );
 
     // Parameters needed for the info endpoint

crates/notary/server/src/service.rs

@@ -3,6 +3,7 @@ pub mod tcp;
 pub mod websocket;
 
 use axum::{
+    body::Body,
     extract::{rejection::JsonRejection, FromRequestParts, Query, State},
     http::{header, request::Parts, StatusCode},
     response::{IntoResponse, Json, Response},
@@ -78,6 +79,17 @@ pub async fn upgrade_protocol(
     State(notary_globals): State<NotaryGlobals>,
     Query(params): Query<NotarizationRequestQuery>,
 ) -> Response {
+    let permit = if let Ok(permit) = notary_globals.semaphore.clone().try_acquire_owned() {
+        permit
+    } else {
+        // TODO: estimate the time more precisely to avoid unnecessary retries.
+        return Response::builder()
+            .status(StatusCode::SERVICE_UNAVAILABLE)
+            .header("Retry-After", 5)
+            .body(Body::default())
+            .expect("Builder should not fail");
+    };
+
     info!("Received upgrade protocol request");
     let session_id = params.session_id;
     // Check if session_id exists in the store, this also removes session_id from
@@ -96,12 +108,14 @@ pub async fn upgrade_protocol(
     // This completes the HTTP Upgrade request and returns a successful response to
     // the client, meanwhile initiating the websocket or tcp connection
     match protocol_upgrade {
-        ProtocolUpgrade::Ws(ws) => {
-            ws.on_upgrade(move |socket| websocket_notarize(socket, notary_globals, session_id))
-        }
-        ProtocolUpgrade::Tcp(tcp) => {
-            tcp.on_upgrade(move |stream| tcp_notarize(stream, notary_globals, session_id))
-        }
+        ProtocolUpgrade::Ws(ws) => ws.on_upgrade(move |socket| async move {
+            websocket_notarize(socket, notary_globals, session_id).await;
+            drop(permit);
+        }),
+        ProtocolUpgrade::Tcp(tcp) => tcp.on_upgrade(move |stream| async move {
+            tcp_notarize(stream, notary_globals, session_id).await;
+            drop(permit);
+        }),
     }
 }
 

crates/notary/tests-integration/Cargo.toml

@@ -14,6 +14,7 @@ tlsn-tls-core = { workspace = true }
 tlsn-core = { workspace = true }
 
 async-tungstenite = { workspace = true, features = ["tokio-native-tls"] }
+futures = { workspace = true }
 http = { workspace = true }
 http-body-util = { workspace = true }
 hyper = { workspace = true, features = ["client", "http1", "server"] }