This commit is contained in:
Hendrik Eeckhaut
2025-09-23 11:04:53 +02:00
parent 0914b907f8
commit be3ee16bb7
7 changed files with 110 additions and 115 deletions

View File

@@ -0,0 +1,37 @@
use http::Uri;
/// Configuration constants for the TLSNotary demo server
/// Maximum number of bytes that can be sent from prover to server
pub const MAX_SENT_DATA: usize = 2048;
/// Maximum number of bytes that can be received by prover from server
pub const MAX_RECV_DATA: usize = 4096;
/// Secret key used in demo requests (should be redacted in proofs)
pub const SECRET: &str = "TLSNotary's private key 🤡";
/// Default server configuration
pub struct Config {
pub host: String,
pub port: u16,
pub server_url: Uri,
}
impl Default for Config {
fn default() -> Self {
Self {
host: "0.0.0.0".into(),
port: 9816,
server_url:
"https://raw.githubusercontent.com/tlsnotary/tlsn/refs/tags/v0.1.0-alpha.12/crates/server-fixture/server/src/data/1kb.json".parse::<Uri>().unwrap(),
}
}
}
impl Config {
pub fn server_domain(&self) -> String {
self.server_url
.host()
.expect("Server URL must have a valid domain")
.to_string()
}
}

View File

@@ -6,6 +6,7 @@ use axum::{
};
use axum_websocket::{WebSocket, WebSocketUpgrade};
use eyre::eyre;
use http::Uri;
use hyper::{body::Incoming, server::conn::http1};
use hyper_util::rt::TokioIo;
use std::{
@@ -18,36 +19,25 @@ use tracing::{debug, error, info};
use ws_stream_tungstenite::WsStream;
mod axum_websocket;
pub mod config;
pub mod prover;
pub mod verifier;
pub mod websocket_utils;
use prover::prover;
use verifier::verifier;
// Maximum number of bytes that can be sent from prover to server
const MAX_SENT_DATA: usize = 2048;
// Maximum number of bytes that can be received by prover from server
const MAX_RECV_DATA: usize = 4096;
const SECRET: &str = "TLSNotary's private key 🤡";
/// Global data that needs to be shared with the axum handlers
#[derive(Clone, Debug)]
struct ServerGlobals {
pub server_url: String,
pub server_domain: String,
pub server_uri: Uri,
}
pub async fn run_server(
prover_host: &str,
prover_port: u16,
server_url: &str,
server_domain: &str,
) -> Result<(), eyre::ErrReport> {
pub async fn run_server(config: &config::Config) -> Result<(), eyre::ErrReport> {
let prover_address = SocketAddr::new(
IpAddr::V4(prover_host.parse().map_err(|err| {
IpAddr::V4(config.host.parse().map_err(|err| {
eyre!("Failed to parse prover host address from server config: {err}")
})?),
prover_port,
config.port,
);
let listener = TcpListener::bind(prover_address)
.await
@@ -60,8 +50,7 @@ pub async fn run_server(
.route("/prove", get(ws_handler_prover))
.route("/verify", get(ws_handler_verifier))
.with_state(ServerGlobals {
server_url: server_url.to_string(),
server_domain: server_domain.to_string(),
server_uri: config.server_url.clone(),
});
loop {
@@ -116,7 +105,7 @@ async fn handle_socket_prover(socket: WebSocket, prover_globals: ServerGlobals)
let socket = socket.into_inner();
let socket = WsStream::new(socket);
let result = prover(socket, &prover_globals.server_url).await;
let result = prover(socket, &prover_globals.server_uri).await;
match result {
Ok(()) => {
info!("============================================");
@@ -132,9 +121,10 @@ async fn handle_socket_prover(socket: WebSocket, prover_globals: ServerGlobals)
async fn handle_socket_verifier(socket: WebSocket, prover_globals: ServerGlobals) {
let stream = WsStream::new(socket.into_inner());
match verifier(stream, &prover_globals.server_domain).await {
let domain = prover_globals.server_uri.authority().unwrap().host();
match verifier(stream, &domain).await {
Ok((sent, received)) => {
info!("Successfully verified {}", &prover_globals.server_domain);
info!("Successfully verified {}", &domain);
info!("Verified sent data:\n{}", sent,);
println!("Verified received data:\n{received}",);
}

View File

@@ -1,16 +1,8 @@
use tlsn_demo_server::run_server;
use tlsn_demo_server::{config::Config, run_server};
use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt, EnvFilter};
const TRACING_FILTER: &str = "INFO";
const PROVER_HOST: &str = "0.0.0.0";
const PROVER_PORT: u16 = 9816;
/// Make sure the following domain is the same in SERVER_URL that will be proven
const SERVER_URL: &str = "https://raw.githubusercontent.com/tlsnotary/tlsn/refs/tags/v0.1.0-alpha.12/crates/server-fixture/server/src/data/1kb.json";
const SERVER_DOMAIN: &str = "raw.githubusercontent.com";
#[tokio::main]
async fn main() -> Result<(), eyre::ErrReport> {
tracing_subscriber::registry()
@@ -18,7 +10,8 @@ async fn main() -> Result<(), eyre::ErrReport> {
.with(tracing_subscriber::fmt::layer())
.init();
run_server(PROVER_HOST, PROVER_PORT, SERVER_URL, SERVER_DOMAIN).await?;
let config: Config = Config::default();
run_server(&config).await?;
Ok(())
}

View File

@@ -1,5 +1,3 @@
use async_tungstenite::{tokio::connect_async_with_config, tungstenite::protocol::WebSocketConfig};
use eyre::eyre;
use http_body_util::Empty;
use hyper::{body::Bytes, Request, StatusCode, Uri};
use hyper_util::rt::TokioIo;
@@ -10,6 +8,9 @@ use spansy::{
Spanned,
};
use crate::config::{MAX_RECV_DATA, MAX_SENT_DATA, SECRET};
use crate::websocket_utils::create_websocket_request;
use async_tungstenite::{tokio::connect_async_with_config, tungstenite::protocol::WebSocketConfig};
use tlsn::{
config::ProtocolConfig,
connection::ServerName,
@@ -19,57 +20,30 @@ use tokio::io::{AsyncRead, AsyncWrite};
use tokio_util::compat::{FuturesAsyncReadCompatExt, TokioAsyncReadCompatExt};
use tracing::{debug, info};
use ws_stream_tungstenite::WsStream;
const TRACING_FILTER: &str = "INFO";
const VERIFIER_HOST: &str = "localhost";
const VERIFIER_PORT: u16 = 9816;
// Maximum number of bytes that can be sent from prover to server
const MAX_SENT_DATA: usize = 2048;
// Maximum number of bytes that can be received by prover from server
const MAX_RECV_DATA: usize = 4096;
const SECRET: &str = "TLSNotary's private key 🤡";
/// Make sure the following url's domain is the same as SERVER_DOMAIN on the verifier side
pub async fn run_prover_test(
server_host: &str,
server_port: u16,
server_url: &str,
) -> Result<(), eyre::ErrReport> {
pub async fn run_prover_test(config: &crate::config::Config) -> Result<(), eyre::ErrReport> {
info!("Sending websocket verify request to server...");
let request = http::Request::builder()
.uri(format!("ws://{server_host}:{server_port}/verify"))
.header("Host", server_host)
.header("Sec-WebSocket-Key", uuid::Uuid::new_v4().to_string())
.header("Sec-WebSocket-Version", "13")
.header("Connection", "Upgrade")
.header("Upgrade", "Websocket")
.body(())
.unwrap();
let (prover_ws_stream, _) =
connect_async_with_config(request, Some(WebSocketConfig::default()))
.await
.map_err(|e| eyre!("Failed to connect to server: {}", e))?;
let request = create_websocket_request(&config.host, config.port, "/verify");
let (ws_stream, _) = connect_async_with_config(request, Some(WebSocketConfig::default()))
.await
.map_err(|e| eyre::eyre!("Failed to connect to server: {}", e))?;
let server_ws_socket = WsStream::new(ws_stream);
info!("Websocket connection established with prover!");
let server_ws_socket = WsStream::new(prover_ws_stream);
prover(server_ws_socket, server_url).await?;
prover(server_ws_socket, &config.server_url).await?;
info!("Proving is successful!");
Ok(())
}
pub async fn prover<T: AsyncWrite + AsyncRead + Send + Unpin + 'static>(
verifier_socket: T,
server_uri: &str,
server_uri: &Uri,
) -> Result<(), eyre::ErrReport> {
debug!("Starting proving...");
let uri = server_uri.parse::<Uri>().unwrap();
assert_eq!(uri.scheme().unwrap().as_str(), "https");
let server_domain = uri.authority().unwrap().host();
let server_port = uri.port_u16().unwrap_or(443);
assert_eq!(server_uri.scheme().unwrap().as_str(), "https");
let server_domain = server_uri.authority().unwrap().host();
let server_port = server_uri.port_u16().unwrap_or(443);
// Create prover and connect to verifier.
let prover_config = ProverConfig::builder()
@@ -114,7 +88,7 @@ pub async fn prover<T: AsyncWrite + AsyncRead + Send + Unpin + 'static>(
// MPC-TLS: Send Request and wait for Response.
info!("Send Request and wait for Response");
let request = Request::builder()
.uri(uri.clone())
.uri(server_uri.clone())
.header("Host", server_domain)
.header("Connection", "close")
.header("Secret", SECRET)

View File

@@ -1,3 +1,5 @@
use crate::config::{Config, MAX_RECV_DATA, MAX_SENT_DATA};
use crate::websocket_utils::create_websocket_request;
use async_tungstenite::{tokio::connect_async_with_config, tungstenite::protocol::WebSocketConfig};
use eyre::eyre;
use tlsn::{
@@ -8,39 +10,18 @@ use tlsn::{
use tokio::io::{AsyncRead, AsyncWrite};
use tokio_util::compat::TokioAsyncReadCompatExt;
use tracing::{debug, info};
use uuid;
use ws_stream_tungstenite::WsStream;
// Maximum number of bytes that can be sent from prover to server
const MAX_SENT_DATA: usize = 2048;
// Maximum number of bytes that can be received by prover from server
const MAX_RECV_DATA: usize = 4096;
/// Connect to prover via websocket and run verification
pub async fn run_verifier_test(
prover_host: &str,
prover_port: u16,
server_domain: &str,
) -> Result<(), eyre::ErrReport> {
pub async fn run_verifier_test(config: &Config) -> Result<(), eyre::ErrReport> {
info!("Sending websocket request to prover...");
let request = http::Request::builder()
.uri(format!("ws://{prover_host}:{prover_port}/prove"))
.header("Host", prover_host)
.header("Sec-WebSocket-Key", uuid::Uuid::new_v4().to_string())
.header("Sec-WebSocket-Version", "13")
.header("Connection", "Upgrade")
.header("Upgrade", "Websocket")
.body(())
.unwrap();
let (prover_ws_stream, _) =
connect_async_with_config(request, Some(WebSocketConfig::default()))
.await
.map_err(|e| eyre!("Failed to connect to prover: {}", e))?;
let request = create_websocket_request(&config.host, config.port, "/prove");
let (ws_stream, _) = connect_async_with_config(request, Some(WebSocketConfig::default()))
.await
.map_err(|e| eyre!("Failed to connect to prover: {}", e))?;
let prover_ws_socket = WsStream::new(ws_stream);
info!("Websocket connection established with prover!");
let prover_ws_socket = WsStream::new(prover_ws_stream);
verifier(prover_ws_socket, server_domain).await?;
verifier(prover_ws_socket, &config.server_domain()).await?;
info!("Verification is successful!");
Ok(())
}

View File

@@ -0,0 +1,17 @@
/// Shared utilities for WebSocket connection handling
use uuid;
/// Create a WebSocket connection request with standard headers
/// This eliminates duplication of WebSocket request creation across modules
pub fn create_websocket_request(host: &str, port: u16, path: &str) -> http::Request<()> {
http::Request::builder()
.uri(format!("ws://{host}:{port}{path}"))
.header("Host", host)
.header("Sec-WebSocket-Key", uuid::Uuid::new_v4().to_string())
.header("Sec-WebSocket-Version", "13")
.header("Connection", "Upgrade")
.header("Upgrade", "Websocket")
.body(())
.unwrap()
}

View File

@@ -1,13 +1,11 @@
use std::time::Duration;
use tlsn_demo_server::{prover::run_prover_test, run_server, verifier::run_verifier_test};
use tlsn_demo_server::{
config::Config, prover::run_prover_test, run_server, verifier::run_verifier_test,
};
use tokio::time::timeout;
use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt, EnvFilter};
const TRACING_FILTER: &str = "INFO";
const PROVER_HOST: &str = "127.0.0.1";
const PROVER_PORT: u16 = 9817; // Use different port to avoid conflicts
const SERVER_DOMAIN: &str = "raw.githubusercontent.com";
const SERVER_URL: &str = "https://raw.githubusercontent.com/tlsnotary/tlsn/refs/tags/v0.1.0-alpha.12/crates/server-fixture/server/src/data/1kb.json";
#[tokio::test]
async fn test_prover_verifier_integration() {
@@ -19,7 +17,8 @@ async fn test_prover_verifier_integration() {
// Start prover server in background task
let server_task = tokio::spawn(async move {
run_server(PROVER_HOST, PROVER_PORT, SERVER_URL, SERVER_DOMAIN)
let config: Config = Config::default();
run_server(&config)
.await
.expect("Server should start successfully")
});
@@ -28,9 +27,10 @@ async fn test_prover_verifier_integration() {
tokio::time::sleep(Duration::from_millis(500)).await;
// Run verifier test with timeout
let config: Config = Config::default();
let verification_result = timeout(
Duration::from_secs(60), // Generous timeout for network operations
run_verifier_test(PROVER_HOST, PROVER_PORT, SERVER_DOMAIN),
run_verifier_test(&config),
)
.await;
@@ -61,7 +61,8 @@ async fn test_verifier_prover_integration() {
// Start prover server in background task
let server_task = tokio::spawn(async move {
run_server(PROVER_HOST, PROVER_PORT, SERVER_URL, SERVER_DOMAIN)
let config: Config = Config::default();
run_server(&config)
.await
.expect("Server should start successfully")
});
@@ -70,9 +71,10 @@ async fn test_verifier_prover_integration() {
tokio::time::sleep(Duration::from_millis(500)).await;
// Run verifier test with timeout
let config: Config = Config::default();
let prover_result = timeout(
Duration::from_secs(60), // Generous timeout for network operations
run_prover_test(PROVER_HOST, PROVER_PORT, SERVER_URL),
run_prover_test(&config),
)
.await;
@@ -96,11 +98,12 @@ async fn test_verifier_prover_integration() {
#[tokio::test]
async fn test_verifier_connection_failure() {
// Test that verifier handles connection failure gracefully
let result = timeout(
Duration::from_secs(5),
run_verifier_test("127.0.0.1", 9999, SERVER_DOMAIN), // Non-existent port
)
.await;
let config = Config {
port: 54321, // wrong port
..Config::default()
};
let result = timeout(Duration::from_secs(5), run_verifier_test(&config)).await;
// Should either timeout or return an error
match result {