feat: accept sessionId in url

Required by the browser to pass session id with Websocket API

src/domain/notary.rs

@@ -23,6 +23,14 @@ pub struct NotarizationSessionRequest {
     pub max_transcript_size: Option<usize>,
 }
 
+/// Request object of the /notarize API
+#[derive(Debug, Clone, Serialize, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct NotarizationNotarizeRequest {
+    /// Session id that is returned from /session API
+    pub session_id: Option<String>,
+}
+
 /// Types of client that the prover is using
 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
 pub enum ClientType {

src/service.rs

@@ -4,7 +4,7 @@ pub mod websocket;
 
 use async_trait::async_trait;
 use axum::{
-    extract::{rejection::JsonRejection, FromRequestParts, State},
+    extract::{rejection::JsonRejection, FromRequestParts, Query, State},
     http::{header, request::Parts, HeaderMap, StatusCode},
     response::{IntoResponse, Json, Response},
 };
@@ -17,7 +17,10 @@ use tracing::{debug, error, info, trace};
 use uuid::Uuid;
 
 use crate::{
-    domain::notary::{NotarizationSessionRequest, NotarizationSessionResponse, NotaryGlobals},
+    domain::notary::{
+        NotarizationNotarizeRequest, NotarizationSessionRequest, NotarizationSessionResponse,
+        NotaryGlobals,
+    },
     error::NotaryServerError,
     service::{
         axum_websocket::{header_eq, WebSocketUpgrade},
@@ -68,6 +71,7 @@ pub async fn upgrade_protocol(
     protocol_upgrade: ProtocolUpgrade,
     mut headers: HeaderMap,
     State(notary_globals): State<NotaryGlobals>,
+    Query(params): Query<NotarizationNotarizeRequest>,
 ) -> Response {
     info!("Received upgrade protocol request");
     // Extract the session_id from the headers
@@ -81,9 +85,13 @@ pub async fn upgrade_protocol(
             }
         },
         None => {
-            let err_msg = "Missing X-Session-Id in upgrade protocol request".to_string();
-            error!(err_msg);
-            return NotaryServerError::BadProverRequest(err_msg).into_response();
+            if let Some(param_session_id) = params.session_id {
+                param_session_id
+            } else {
+                let err_msg = "Either X-Session-Id or session_id parameter are required in upgrade protocol request".to_string();
+                error!(err_msg);
+                return NotaryServerError::BadProverRequest(err_msg).into_response();
+            }
         }
     };
     // Fetch the configuration data from the store using the session_id

tests/integration_test.rs

@@ -421,3 +421,100 @@ async fn test_websocket_prover() {
 
     debug!("Done notarization!");
 }
+
+#[tokio::test]
+async fn test_websocket_prover_session_id_in_uri() {
+    // This test is copied from `test_websocket_prover` except that the session id
+    // is passed in the URI instead of the header
+
+    // Notary server configuration setup
+    let notary_config = setup_config_and_server(100, 7049).await;
+    let notary_host = notary_config.server.host.clone();
+    let notary_port = notary_config.server.port;
+
+    // Connect to the notary server via TLS-WebSocket
+    // Try to avoid dealing with transport layer directly to mimic the limitation of a browser extension that uses websocket
+    //
+    // Establish TLS setup for connections later
+    let certificate =
+        tokio_native_tls::native_tls::Certificate::from_pem(NOTARY_CA_CERT_BYTES).unwrap();
+    let notary_tls_connector = tokio_native_tls::native_tls::TlsConnector::builder()
+        .add_root_certificate(certificate)
+        .use_sni(false)
+        .danger_accept_invalid_certs(true)
+        .build()
+        .unwrap();
+
+    // Call the /session HTTP API to configure notarization and obtain session id
+    let mut hyper_http_connector = HttpConnector::new();
+    hyper_http_connector.enforce_http(false);
+    let mut hyper_tls_connector =
+        HttpsConnector::from((hyper_http_connector, notary_tls_connector.clone().into()));
+    hyper_tls_connector.https_only(true);
+    let https_client = Client::builder().build::<_, hyper::Body>(hyper_tls_connector);
+
+    // Build the HTTP request to configure notarization
+    let payload = serde_json::to_string(&NotarizationSessionRequest {
+        client_type: notary_server::ClientType::Websocket,
+        max_transcript_size: Some(notary_config.notarization.max_transcript_size),
+    })
+    .unwrap();
+
+    let request = Request::builder()
+        .uri(format!("https://{notary_host}:{notary_port}/session"))
+        .method("POST")
+        .header("Host", notary_host.clone())
+        // Need to specify application/json for axum to parse it as json
+        .header("Content-Type", "application/json")
+        .body(Body::from(payload))
+        .unwrap();
+
+    debug!("Sending request");
+
+    let response = https_client.request(request).await.unwrap();
+
+    debug!("Sent request");
+
+    assert!(response.status() == StatusCode::OK);
+
+    debug!("Response OK");
+
+    // Pretty printing :)
+    let payload = to_bytes(response.into_body()).await.unwrap().to_vec();
+    let notarization_response =
+        serde_json::from_str::<NotarizationSessionResponse>(&String::from_utf8_lossy(&payload))
+            .unwrap();
+
+    debug!("Notarization response: {:?}", notarization_response,);
+
+    // Connect to the Notary via TLS-Websocket
+    //
+    // Note: This will establish a new TLS-TCP connection instead of reusing the previous TCP connection
+    // used in the previous HTTP POST request because we cannot claim back the tcp connection used in hyper
+    // client while using its high level request function — there does not seem to have a crate that can let you
+    // make a request without establishing TCP connection where you can claim the TCP connection later after making the request
+    let request = http::Request::builder()
+        // Note that the session id is passed in the URI instead of the header
+        .uri(format!(
+            "wss://{}:{}/notarize?sessionId={}",
+            notary_host,
+            notary_port,
+            notarization_response.session_id.clone()
+        ))
+        .header("Host", notary_host.clone())
+        .header("Sec-WebSocket-Key", uuid::Uuid::new_v4().to_string())
+        .header("Sec-WebSocket-Version", "13")
+        .header("Connection", "Upgrade")
+        .header("Upgrade", "Websocket")
+        .body(())
+        .unwrap();
+
+    // If Websocket connection fails it will panic here
+    connect_async_with_tls_connector_and_config(
+        request,
+        Some(notary_tls_connector.into()),
+        Some(WebSocketConfig::default()),
+    )
+    .await
+    .unwrap();
+}