From 3074bc15fa59eb79955be3a4d5d5f5516dd6f48f Mon Sep 17 00:00:00 2001
From: th4s <th4s@metavoid.xyz>
Date: Mon, 21 Aug 2023 22:58:39 +0200
Subject: [PATCH] WIP: Modifying twitter example and improving API...

---
 tlsn/examples/Cargo.toml       |  1 +
 tlsn/examples/twitter_dm.rs    | 32 ++++++++++++++++++++++++++++++++
 tlsn/tlsn-core/src/span/mod.rs | 34 ++++++++++++++++++++++++++++++++--
 3 files changed, 65 insertions(+), 2 deletions(-)

diff --git a/tlsn/examples/Cargo.toml b/tlsn/examples/Cargo.toml
index ff8f3a8e9c..2c7e572b97 100644
--- a/tlsn/examples/Cargo.toml
+++ b/tlsn/examples/Cargo.toml
@@ -39,6 +39,7 @@ rustls = { version = "0.21" }
 rustls-pemfile = { version = "1.0.2" }
 tokio-rustls = { version = "0.24.1" }
 dotenv = "0.15.0"
+httparse = "1"
 
 [[example]]
 name = "twitter_dm"
diff --git a/tlsn/examples/twitter_dm.rs b/tlsn/examples/twitter_dm.rs
index 2cf17208eb..4b1ff35025 100644
--- a/tlsn/examples/twitter_dm.rs
+++ b/tlsn/examples/twitter_dm.rs
@@ -1,6 +1,7 @@
 /// This prover implementation talks to the notary server implemented in https://github.com/tlsnotary/notary-server, instead of the simple_notary.rs in this example directory
 use eyre::Result;
 use futures::AsyncWriteExt;
+use httparse::EMPTY_HEADER;
 use hyper::{body::to_bytes, client::conn::Parts, Body, Request, StatusCode};
 use rustls::{Certificate, ClientConfig, RootCertStore};
 use serde::{Deserialize, Serialize};
@@ -12,6 +13,7 @@ use std::{
     ops::Range,
     sync::Arc,
 };
+use tlsn_core::span::{http::HttpSpanner, SpanCommit};
 use tokio::{fs::File, io::AsyncWriteExt as _};
 use tokio_rustls::TlsConnector;
 use tokio_util::compat::{FuturesAsyncReadCompatExt, TokioAsyncReadCompatExt};
@@ -322,3 +324,33 @@ 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>,
+}
+
+impl<'a, 'b> SpanCommit for TwitterSpanner<'a, 'b> {
+    fn span_request(&mut self, request: &[u8]) -> Vec<Range<usize>> {
+        let mut headers = vec![EMPTY_HEADER; 12];
+        self.http.parse_request(&mut headers, request);
+
+        let cookie = self
+            .http
+            .header_value_span_request("Cookie", request)
+            .unwrap();
+        let authorization = self
+            .http
+            .header_value_span_request("Authorization", request)
+            .unwrap();
+        let csrf = self
+            .http
+            .header_value_span_request("X-Csrf-Token", request)
+            .unwrap();
+
+        vec![cookie, authorization, csrf]
+    }
+
+    fn span_response(&mut self, response: &[u8]) -> Vec<Range<usize>> {
+        vec![0..response.len()]
+    }
+}
diff --git a/tlsn/tlsn-core/src/span/mod.rs b/tlsn/tlsn-core/src/span/mod.rs
index d504bf41d3..30d1c81a96 100644
--- a/tlsn/tlsn-core/src/span/mod.rs
+++ b/tlsn/tlsn-core/src/span/mod.rs
@@ -15,9 +15,37 @@ 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<u32>>;
+    fn span_request(&mut self, request: &[u8]) -> Vec<Range<usize>>;
     /// Identify byte ranges in the response to commit to
-    fn span_response(&mut self, response: &[u8]) -> Vec<Range<u32>>;
+    fn span_response(&mut self, response: &[u8]) -> Vec<Range<usize>>;
+}
+
+pub fn invert_ranges(
+    ranges: Vec<Range<usize>>,
+    len: usize,
+) -> Result<Vec<Range<usize>>, SpanError> {
+    for range in ranges.iter() {
+        // Check that there is no invalid or empty range
+        if range.start >= range.end {
+            return Err(SpanError::InvalidRange);
+        }
+
+        // Check that ranges are not out of bounds
+        if range.start >= len || range.end > len {
+            return Err(SpanError::InvalidRange);
+        }
+
+        // Check that ranges are not overlapping
+        if ranges
+            .iter()
+            .any(|r| r.start < range.end && r.end > range.start)
+        {
+            return Err(SpanError::InvalidRange);
+        }
+    }
+
+    // Now invert ranges
+    let mut inverted = (0..len).collect::<Vec<usize>>();
 }
 
 /// An error that can occur during span creation
@@ -26,4 +54,6 @@ pub enum SpanError {
     /// The request or response could not be parsed
     #[error("Error during parsing")]
     ParseError,
+    #[error("Found invalid ranges")]
+    InvalidRange,
 }