Add notary client crate

notary-server/Cargo.toml

@@ -45,8 +45,9 @@ ws_stream_tungstenite = { version = "0.13.0", features = ["tokio_io"] }
 # specify vendored feature to use statically linked copy of OpenSSL
 hyper-tls = { version = "0.6.0", features = ["vendored"] }
 rstest = "0.18"
+tls-server-fixture = { path = "../components/tls/tls-server-fixture" }
 tlsn-notary-client = { path = "../tlsn/tlsn-notary-client" }
 tlsn-prover = { path = "../tlsn/tlsn-prover", features = ["tracing"] }
-tls-server-fixture = { path = "../components/tls/tls-server-fixture" }
+tlsn-tls-client-async = { path = "../components/tls/tls-client-async" }
 tlsn-tls-core = { path = "../components/tls/tls-core" }
 tokio-native-tls = { version = "0.3.1", features = ["vendored"] }

notary-server/tests/integration_test.rs

@@ -10,23 +10,30 @@ use hyper_util::{
 };
 use rstest::rstest;
 use std::{string::String, time::Duration};
-use tls_core::anchors::RootCertStore as TlsClientRootCertStore;
+use tls_client_async::TlsConnection;
+use tls_core::{anchors::RootCertStore, key::Certificate};
 use tls_server_fixture::{bind_test_server_hyper, CA_CERT_DER, SERVER_DOMAIN};
-use tlsn_notary_client::client::NotaryClient;
-use tlsn_prover::tls::{state::Setup, Prover, ProverConfig};
-use tokio_util::compat::{FuturesAsyncReadCompatExt, TokioAsyncReadCompatExt};
+use tlsn_notary_client::{NotaryClient, NotaryConnection};
+use tlsn_prover::tls::{Prover, ProverConfig};
+use tokio::{
+    io::{AsyncRead, AsyncWrite},
+    net::TcpStream,
+};
+use tokio_util::compat::{Compat, FuturesAsyncReadCompatExt, TokioAsyncReadCompatExt};
 use tracing::debug;
 use ws_stream_tungstenite::WsStream;
 
 use notary_server::{
-    run_server, AuthorizationProperties, LoggingProperties, NotarizationProperties,
+    read_pem_file, run_server, AuthorizationProperties, LoggingProperties, NotarizationProperties,
     NotarizationSessionRequest, NotarizationSessionResponse, NotaryServerProperties,
     NotarySigningKeyProperties, ServerProperties, TLSProperties,
 };
 
+const MAX_SENT_DATA: usize = 1 << 13;
+const MAX_RECV_DATA: usize = 1 << 13;
+
+const NOTARY_CA_CERT_PATH: &str = "./fixture/tls/rootCA.crt";
 const NOTARY_CA_CERT_BYTES: &[u8] = include_bytes!("../fixture/tls/rootCA.crt");
-const MAX_SENT: usize = 1 << 13;
-const MAX_RECV: usize = 1 << 13;
 const API_KEY: &str = "test_api_key_0";
 
 fn get_server_config(port: u16, tls_enabled: bool, auth_enabled: bool) -> NotaryServerProperties {
@@ -83,80 +90,109 @@ async fn setup_config_and_server(
     notary_config
 }
 
-fn get_server_root_cert_store() -> TlsClientRootCertStore {
-    let mut root_store = tls_core::anchors::RootCertStore::empty();
-    root_store
-        .add(&tls_core::key::Certificate(CA_CERT_DER.to_vec()))
-        .unwrap();
-    root_store
-}
-
-async fn tcp_prover(
-    notary_config: NotaryServerProperties,
-    server_root_store: TlsClientRootCertStore,
-) -> Prover<Setup> {
+async fn tcp_prover(notary_config: NotaryServerProperties) -> (TcpStream, String) {
     let notary_client = NotaryClient::builder()
         .host(&notary_config.server.host)
         .port(notary_config.server.port)
-        .max_sent_data(MAX_SENT)
-        .max_recv_data(MAX_RECV)
-        // set this to None to turn off TLS with notary server
+        .max_sent_data(MAX_SENT_DATA)
+        .max_recv_data(MAX_RECV_DATA)
+        .root_cert_store(None)
         .notary_dns(None)
-        // set this to None to turn off TLS with notary server
-        .notary_root_cert_store(None)
-        .server_dns(SERVER_DOMAIN)
-        .server_root_cert_store(server_root_store)
         .build()
         .unwrap();
 
-    notary_client.setup_tcp_prover().await.unwrap()
+    if let (NotaryConnection::Tcp(notary_socket), session_id) =
+        notary_client.request_notarization().await.unwrap()
+    {
+        (notary_socket, session_id)
+    } else {
+        panic!("Invalid notary connection received: TLS");
+    }
 }
 
-async fn tls_prover(
-    notary_config: NotaryServerProperties,
-    server_root_store: TlsClientRootCertStore,
-) -> Prover<Setup> {
+async fn tls_prover(notary_config: NotaryServerProperties) -> (Compat<TlsConnection>, String) {
+    let mut certificate_file_reader = read_pem_file(NOTARY_CA_CERT_PATH).await.unwrap();
+    let mut certificates: Vec<Certificate> = rustls_pemfile::certs(&mut certificate_file_reader)
+        .unwrap()
+        .into_iter()
+        .map(Certificate)
+        .collect();
+    let certificate = certificates.remove(0);
+
+    let mut root_cert_store = RootCertStore::empty();
+    root_cert_store.add(&certificate).unwrap();
+
     let mut notary_client_builder = NotaryClient::builder();
 
     notary_client_builder
         .host(&notary_config.server.host)
         .port(notary_config.server.port)
-        .max_sent_data(MAX_SENT)
-        .max_recv_data(MAX_RECV)
-        .server_dns(SERVER_DOMAIN)
-        .server_root_cert_store(server_root_store);
+        .max_sent_data(MAX_SENT_DATA)
+        .max_recv_data(MAX_RECV_DATA)
+        .root_cert_store(Some(root_cert_store))
+        .notary_dns(Some(notary_config.server.name));
 
     if notary_config.authorization.enabled {
         notary_client_builder.api_key(API_KEY);
     }
 
     let notary_client = notary_client_builder.build().unwrap();
 
-    notary_client.setup_tls_prover().await.unwrap()
+    if let (NotaryConnection::Tls(notary_socket), session_id) =
+        notary_client.request_notarization().await.unwrap()
+    {
+        (notary_socket, session_id)
+    } else {
+        panic!("Invalid notary connection received: TCP");
+    }
 }
 
 #[rstest]
 #[case::with_tls_and_auth(
-    tls_prover(setup_config_and_server(100, 7047, true, true).await, get_server_root_cert_store())
+    tls_prover(setup_config_and_server(100, 7047, true, true).await)
 )]
 #[case::with_tls_and_no_auth(
-    tls_prover(setup_config_and_server(100, 7048, true, false).await, get_server_root_cert_store())
+    tls_prover(setup_config_and_server(100, 7048, true, false).await)
 )]
 #[case::without_tls(
-    tcp_prover(setup_config_and_server(100, 7049, false, false).await, get_server_root_cert_store())
+    tcp_prover(setup_config_and_server(100, 7049, false, false).await)
 )]
 #[awt]
 #[tokio::test]
-async fn test_tcp_prover(
+async fn test_tcp_prover<S: AsyncWrite + AsyncRead + Send + Unpin + 'static>(
     #[future]
     #[case]
-    prover: Prover<Setup>,
+    requested_notarization: (S, String),
 ) {
+    let (notary_socket, session_id) = requested_notarization;
+
+    let mut root_cert_store = RootCertStore::empty();
+    root_cert_store
+        .add(&tls_core::key::Certificate(CA_CERT_DER.to_vec()))
+        .unwrap();
+
+    // Prover config using the session_id returned from calling /session endpoint in notary client
+    let prover_config = ProverConfig::builder()
+        .id(session_id)
+        .server_dns(SERVER_DOMAIN)
+        .max_sent_data(MAX_SENT_DATA)
+        .max_recv_data(MAX_RECV_DATA)
+        .root_cert_store(root_cert_store)
+        .build()
+        .unwrap();
+
+    // Create a new prover
+    let prover = Prover::new(prover_config)
+        .setup(notary_socket.compat())
+        .await
+        .unwrap();
+
     // Connect to the Server
     let (client_socket, server_socket) = tokio::io::duplex(2 << 16);
     let server_task = tokio::spawn(bind_test_server_hyper(server_socket.compat()));
 
     let (tls_connection, prover_fut) = prover.connect(client_socket.compat()).await.unwrap();
+
     // Spawn the Prover task to be run concurrently
     let prover_task = tokio::spawn(prover_fut);
 
@@ -235,8 +271,8 @@ async fn test_websocket_prover() {
     // Build the HTTP request to configure notarization
     let payload = serde_json::to_string(&NotarizationSessionRequest {
         client_type: notary_server::ClientType::Websocket,
-        max_sent_data: Some(MAX_SENT),
-        max_recv_data: Some(MAX_RECV),
+        max_sent_data: Some(MAX_SENT_DATA),
+        max_recv_data: Some(MAX_RECV_DATA),
     })
     .unwrap();
 
@@ -315,8 +351,8 @@ async fn test_websocket_prover() {
         .id(notarization_response.session_id)
         .server_dns(SERVER_DOMAIN)
         .root_cert_store(root_store)
-        .max_sent_data(MAX_SENT)
-        .max_recv_data(MAX_RECV)
+        .max_sent_data(MAX_SENT_DATA)
+        .max_recv_data(MAX_RECV_DATA)
         .build()
         .unwrap();
 

tlsn/Cargo.toml

@@ -62,7 +62,6 @@ bincode = "1"
 hex = "0.4"
 bytes = "1.4"
 opaque-debug = "0.3"
-eyre = "0.6.8"
 
 tracing = "0.1"
 tracing-subscriber = "0.3"

tlsn/examples/Cargo.toml

@@ -8,9 +8,9 @@ version = "0.0.0"
 tlsn-core.workspace = true
 tlsn-notary-client.workspace = true
 tlsn-prover = { workspace = true, features = ["tracing"] }
+tlsn-tls-core.workspace = true
 tlsn-verifier.workspace = true
 
-eyre.workspace = true
 futures.workspace = true
 http-body-util.workspace = true
 hyper = { workspace = true, features = ["client", "http1"] }
@@ -33,6 +33,7 @@ chrono = "0.4"
 dotenv = "0.15.0"
 elliptic-curve = { version = "0.13.5", features = ["pkcs8"] }
 regex = "1.10.3"
+rustls-pemfile = { version = "1.0.2" }
 serde = { version = "1.0.147", features = ["derive"] }
 serde_json = "1.0"
 

tlsn/examples/discord/discord_dm.rs

@@ -6,8 +6,11 @@ use http_body_util::{BodyExt, Empty};
 use hyper::{body::Bytes, Request, StatusCode};
 use hyper_util::rt::TokioIo;
 use std::{env, ops::Range, str};
+use tls_core::anchors::RootCertStore;
 use tlsn_core::proof::TlsProof;
-use tlsn_examples::request_notarization;
+use tlsn_examples::parse_cert;
+use tlsn_notary_client::{NotaryClient, NotaryConnection};
+use tlsn_prover::tls::{Prover, ProverConfig};
 use tokio::io::AsyncWriteExt as _;
 use tokio_util::compat::{FuturesAsyncReadCompatExt, TokioAsyncReadCompatExt};
 use tracing::debug;
@@ -19,6 +22,10 @@ const SERVER_DOMAIN: &str = "discord.com";
 const NOTARY_HOST: &str = "127.0.0.1";
 const NOTARY_PORT: u16 = 7047;
 
+// Setting to enable connecting to local notary server via TLS
+const NOTARY_CA_CERT_PATH: &str = "../../../notary-server/fixture/tls/rootCA.crt";
+const NOTARY_DNS: &str = "tlsnotaryserver.io";
+
 // P/S: If the following limits are increased, please ensure max-transcript-size of
 // the notary server's config (../../../notary-server) is increased too, where
 // max-transcript-size = MAX_SENT_DATA + MAX_RECV_DATA
@@ -38,16 +45,44 @@ async fn main() {
     let auth_token = env::var("AUTHORIZATION").unwrap();
     let user_agent = env::var("USER_AGENT").unwrap();
 
-    // Create a new prover
-    let prover = request_notarization(
-        NOTARY_HOST,
-        NOTARY_PORT,
-        Some(MAX_SENT_DATA),
-        Some(MAX_RECV_DATA),
-        SERVER_DOMAIN,
-    )
-    .await;
+    // Setup a client connection to the notary server via TLS
+    let root_ca_cert = parse_cert(NOTARY_CA_CERT_PATH).await;
+    let mut root_cert_store = RootCertStore::empty();
+    root_cert_store.add(&root_ca_cert).unwrap();
+
+    let notary_client = NotaryClient::builder()
+        .host(NOTARY_HOST)
+        .port(NOTARY_PORT)
+        .max_sent_data(MAX_SENT_DATA)
+        .max_recv_data(MAX_RECV_DATA)
+        .root_cert_store(Some(root_cert_store))
+        .notary_dns(Some(NOTARY_DNS.to_string()))
+        .build()
+        .unwrap();
+
+    // Send requests for configuration and notarization to the notary server
+    let (NotaryConnection::Tls(notary_socket), session_id) =
+        notary_client.request_notarization().await.unwrap()
+    else {
+        panic!("Invalid notary connection received: TCP");
+    };
+
+    // Configure a new prover with the unique session id returned from notary client
+    let prover_config = ProverConfig::builder()
+        .id(session_id)
+        .server_dns(SERVER_DOMAIN)
+        .max_sent_data(MAX_SENT_DATA)
+        .max_recv_data(MAX_RECV_DATA)
+        .build()
+        .unwrap();
+
+    // Create a new prover and set up the MPC backend.
+    let prover = Prover::new(prover_config)
+        .setup(notary_socket.compat())
+        .await
+        .unwrap();
 
+    // Open a new socket to the application server
     let client_socket = tokio::net::TcpStream::connect((SERVER_DOMAIN, 443))
         .await
         .unwrap();

tlsn/examples/src/lib.rs

@@ -1,8 +1,9 @@
 use elliptic_curve::pkcs8::DecodePrivateKey;
 use futures::{AsyncRead, AsyncWrite};
-use tlsn_notary_client::client::NotaryClient;
-use tlsn_prover::tls::{state::Setup, Prover};
+use std::io::BufReader;
+use tls_core::key::Certificate;
 use tlsn_verifier::tls::{Verifier, VerifierConfig};
+use tokio::fs::File;
 
 /// Runs a simple Notary with the provided connection to the Prover.
 pub async fn run_notary<T: AsyncWrite + AsyncRead + Send + Unpin + 'static>(conn: T) {
@@ -23,31 +24,14 @@ pub async fn run_notary<T: AsyncWrite + AsyncRead + Send + Unpin + 'static>(conn
         .unwrap();
 }
 
-/// Requests notarization from the Notary server.
-pub async fn request_notarization(
-    host: &str,
-    port: u16,
-    max_sent_data: Option<usize>,
-    max_recv_data: Option<usize>,
-    server_dns: &str,
-) -> Prover<Setup> {
-    let mut notary_client_builder = NotaryClient::builder();
-
-    notary_client_builder
-        .host(host)
-        .port(port)
-        .server_dns(server_dns);
-
-    if let Some(max_sent_data) = max_sent_data {
-        notary_client_builder.max_sent_data(max_sent_data);
-    }
-
-    if let Some(max_recv_data) = max_recv_data {
-        notary_client_builder.max_recv_data(max_recv_data);
-    }
-
-    let notary_client = notary_client_builder.build().unwrap();
-
-    // Setup tls connection to the notary server
-    notary_client.setup_tls_prover().await.unwrap()
+/// Parse certificate as tls-core's Certificate struct, so that one can use tls-client's RootCertStore to add the cert
+pub async fn parse_cert(file_path: &str) -> Certificate {
+    let key_file = File::open(file_path).await.unwrap().into_std().await;
+    let mut certificate_file_reader = BufReader::new(key_file);
+    let mut certificates: Vec<Certificate> = rustls_pemfile::certs(&mut certificate_file_reader)
+        .unwrap()
+        .into_iter()
+        .map(Certificate)
+        .collect();
+    certificates.remove(0)
 }