Adapted examples and added TotalSpanner

- completed `fn invert_ranges`

tlsn/examples/twitter_dm.rs

@@ -13,7 +13,7 @@ use std::{
     ops::Range,
     sync::Arc,
 };
-use tlsn_core::span::{http::HttpSpanner, SpanCommit};
+use tlsn_core::span::{http::HttpSpanner, invert_ranges, SpanCommit, SpanError};
 use tokio::{fs::File, io::AsyncWriteExt as _};
 use tokio_rustls::TlsConnector;
 use tokio_util::compat::{FuturesAsyncReadCompatExt, TokioAsyncReadCompatExt};
@@ -119,7 +119,7 @@ async fn main() {
     let request = Request::builder()
         .uri(format!("https://{NOTARY_DOMAIN}:{NOTARY_PORT}/session"))
         .method("POST")
-        .header("Host", NOTARY_DOMAIN.clone())
+        .header("Host", NOTARY_DOMAIN)
         // Need to specify application/json for axum to parse it as json
         .header("Content-Type", "application/json")
         .body(Body::from(payload))
@@ -250,29 +250,9 @@ async fn main() {
     client_socket.close().await.unwrap();
 
     // The Prover task should be done now, so we can grab it.
-    let mut prover = prover_task.await.unwrap().unwrap();
-
-    // Identify the ranges in the transcript that contain secrets
-    let (public_ranges, private_ranges) = find_ranges(
-        prover.sent_transcript().data(),
-        &[
-            access_token.as_bytes(),
-            auth_token.as_bytes(),
-            csrf_token.as_bytes(),
-        ],
-    );
-
-    // Commit to the outbound transcript, isolating the data that contain secrets
-    for range in public_ranges.iter().chain(private_ranges.iter()) {
-        prover.add_commitment_sent(range.clone()).unwrap();
-    }
+    let prover = prover_task.await.unwrap().unwrap();
 
-    // Commit to the full received transcript in one shot, as we don't need to redact anything
-    let recv_len = prover.recv_transcript().data().len();
-    prover.add_commitment_recv(0..recv_len as u32).unwrap();
-
-    // Finalize, returning the notarized session
-    let notarized_session = prover.finalize(None).await.unwrap();
+    let notarized_session = prover.finalize(Box::new(TwitterSpanner)).await.unwrap();
 
     debug!("Notarization complete!");
 
@@ -287,70 +267,38 @@ async fn main() {
     .unwrap();
 }
 
-/// Find the ranges of the public and private parts of a sequence.
-///
-/// Returns a tuple of `(public, private)` ranges.
-fn find_ranges(seq: &[u8], sub_seq: &[&[u8]]) -> (Vec<Range<u32>>, Vec<Range<u32>>) {
-    let mut private_ranges = Vec::new();
-    for s in sub_seq {
-        for (idx, w) in seq.windows(s.len()).enumerate() {
-            if w == *s {
-                private_ranges.push(idx as u32..(idx + w.len()) as u32);
-            }
-        }
-    }
-
-    let mut sorted_ranges = private_ranges.clone();
-    sorted_ranges.sort_by_key(|r| r.start);
-
-    let mut public_ranges = Vec::new();
-    let mut last_end = 0;
-    for r in sorted_ranges {
-        if r.start > last_end {
-            public_ranges.push(last_end..r.start);
-        }
-        last_end = r.end;
-    }
-
-    if last_end < seq.len() as u32 {
-        public_ranges.push(last_end..seq.len() as u32);
-    }
-
-    (public_ranges, private_ranges)
-}
-
 /// Read a PEM-formatted file and return its buffer reader
 async fn read_pem_file(file_path: &str) -> Result<BufReader<StdFile>> {
     let key_file = File::open(file_path).await?.into_std().await;
     Ok(BufReader::new(key_file))
 }
 
-struct TwitterSpanner<'a, 'b> {
-    http: HttpSpanner<'a, 'b>,
-}
+struct TwitterSpanner;
 
-impl<'a, 'b> SpanCommit for TwitterSpanner<'a, 'b> {
-    fn span_request(&mut self, request: &[u8]) -> Vec<Range<usize>> {
+impl SpanCommit for TwitterSpanner {
+    fn span_request(&mut self, request: &[u8]) -> Result<Vec<Range<usize>>, SpanError> {
         let mut headers = vec![EMPTY_HEADER; 12];
-        self.http.parse_request(&mut headers, request);
+        let mut http_spanner = HttpSpanner::new();
+
+        http_spanner.parse_request(&mut headers, request).unwrap();
 
-        let cookie = self
-            .http
+        let cookie = http_spanner
             .header_value_span_request("Cookie", request)
             .unwrap();
-        let authorization = self
-            .http
+        let authorization = http_spanner
             .header_value_span_request("Authorization", request)
             .unwrap();
-        let csrf = self
-            .http
+        let csrf = http_spanner
             .header_value_span_request("X-Csrf-Token", request)
             .unwrap();
 
-        vec![cookie, authorization, csrf]
+        invert_ranges(vec![cookie, authorization, csrf], request.len())
     }
 
-    fn span_response(&mut self, response: &[u8]) -> Vec<Range<usize>> {
-        vec![0..response.len()]
+    fn span_response(&mut self, response: &[u8]) -> Result<Vec<Range<usize>>, SpanError> {
+        Ok(vec![Range {
+            start: 0,
+            end: response.len(),
+        }])
     }
 }

tlsn/tests-integration/Cargo.toml

@@ -13,6 +13,7 @@ publish = false
 tlsn-tls-core.workspace = true
 tlsn-prover.workspace = true
 tlsn-notary.workspace = true
+tlsn-core.workspace = true
 tls-server-fixture.workspace = true
 
 p256 = { workspace = true, features = ["ecdsa"] }

tlsn/tests-integration/tests/test.rs

@@ -1,6 +1,7 @@
 use futures::AsyncWriteExt;
 use hyper::{body::to_bytes, Body, Request, StatusCode};
 use tls_server_fixture::{bind_test_server_hyper, CA_CERT_DER, SERVER_DOMAIN};
+use tlsn_core::span::TotalSpanner;
 use tlsn_notary::{bind_notary, NotaryConfig};
 use tlsn_prover::{bind_prover, ProverConfig};
 use tokio::io::{AsyncRead, AsyncWrite};
@@ -76,15 +77,9 @@ async fn prover<T: AsyncWrite + AsyncRead + Send + Unpin + 'static>(notary_socke
 
     client_socket.close().await.unwrap();
 
-    let mut prover = prover_task.await.unwrap().unwrap();
+    let prover = prover_task.await.unwrap().unwrap();
 
-    let sent_len = prover.sent_transcript().data().len();
-    let recv_len = prover.recv_transcript().data().len();
-
-    prover.add_commitment_sent(0..sent_len as u32).unwrap();
-    prover.add_commitment_recv(0..recv_len as u32).unwrap();
-
-    _ = prover.finalize(None).await.unwrap();
+    _ = prover.finalize(Box::new(TotalSpanner)).await.unwrap();
 }
 
 #[instrument(skip(socket))]

tlsn/tlsn-core/src/span/mod.rs

@@ -15,11 +15,31 @@ pub mod json;
 /// created
 pub trait SpanCommit {
     /// Identify byte ranges in the request to commit to
-    fn span_request(&mut self, request: &[u8]) -> Vec<Range<usize>>;
+    fn span_request(&mut self, request: &[u8]) -> Result<Vec<Range<usize>>, SpanError>;
     /// Identify byte ranges in the response to commit to
-    fn span_response(&mut self, response: &[u8]) -> Vec<Range<usize>>;
+    fn span_response(&mut self, response: &[u8]) -> Result<Vec<Range<usize>>, SpanError>;
 }
 
+/// A Spanner that commits to the entire request and response
+pub struct TotalSpanner;
+
+impl SpanCommit for TotalSpanner {
+    fn span_request(&mut self, request: &[u8]) -> Result<Vec<Range<usize>>, SpanError> {
+        Ok(vec![Range {
+            start: 0,
+            end: request.len(),
+        }])
+    }
+
+    fn span_response(&mut self, response: &[u8]) -> Result<Vec<Range<usize>>, SpanError> {
+        Ok(vec![Range {
+            start: 0,
+            end: response.len(),
+        }])
+    }
+}
+
+/// Inverts a set of ranges, i.e. returns the complement of the ranges
 pub fn invert_ranges(
     ranges: Vec<Range<usize>>,
     len: usize,
@@ -45,15 +65,34 @@ pub fn invert_ranges(
     }
 
     // Now invert ranges
-    let mut inverted = (0..len).collect::<Vec<usize>>();
+    let mut inverted = vec![Range { start: 0, end: len }];
+
+    for range in ranges.iter() {
+        let inv = inverted
+            .iter_mut()
+            .find(|inv| range.start >= inv.start)
+            .unwrap();
+
+        let original_len = inv.end;
+        inv.end = range.start;
+
+        inverted.push(Range {
+            start: range.end,
+            end: original_len,
+        });
+    }
+
+    Ok(inverted)
 }
 
 /// An error that can occur during span creation
+#[allow(missing_docs)]
 #[derive(Debug, thiserror::Error)]
 pub enum SpanError {
-    /// The request or response could not be parsed
     #[error("Error during parsing")]
     ParseError,
     #[error("Found invalid ranges")]
     InvalidRange,
+    #[error("Custom error: {0}")]
+    Custom(String),
 }

tlsn/tlsn-prover/src/error.rs

@@ -22,6 +22,8 @@ pub enum ProverError {
     ServerNoCloseNotify,
     #[error(transparent)]
     CommitmentError(#[from] CommitmentError),
+    #[error(transparent)]
+    SpanError(#[from] tlsn_core::span::SpanError),
 }
 
 impl From<MpcTlsError> for ProverError {

tlsn/tlsn-prover/src/lib.rs

@@ -309,11 +309,11 @@ where
         mut spanner: Box<dyn SpanCommit>,
     ) -> Result<NotarizedSession, ProverError> {
         // Add commitments identified by the spanner
-        for range in spanner.span_request(self.state.transcript_tx.data()) {
-            self.add_commitment(range, Direction::Sent)?;
+        for range in spanner.span_request(self.state.transcript_tx.data())? {
+            self.add_commitment(range.start as u32..range.end as u32, Direction::Sent)?;
         }
-        for range in spanner.span_response(self.state.transcript_rx.data()) {
-            self.add_commitment(range, Direction::Received)?;
+        for range in spanner.span_response(self.state.transcript_rx.data())? {
+            self.add_commitment(range.start as u32..range.end as u32, Direction::Received)?;
         }
 
         let Notarize {