sigoden · sigoden · Mar 15, 2024 · Mar 15, 2024 · Mar 15, 2024
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -38,14 +38,18 @@ serde = { version = "1.0.197", features = ["derive"] }
 serde_json = { version = "1.0.114", features = ["preserve_order"] }
 time = "0.3.34"
 tokio = { version = "1", features = ["rt-multi-thread", "net", "macros", "fs", "io-util", "signal"]}
+tokio-graceful = "0.1.6"
 tokio-util = { version = "0.7",  features = ["io-util", "compat"] }
 tokio-rustls = "0.25.0"
 tokio-stream = { version = "0.1.14", default-features = false, features = ["sync"] }
 tokio-tungstenite = { version = "0.21.0", features = ["rustls-tls-webpki-roots"] }
 url = "2.5.0"
 
 [dev-dependencies]
+async-http-proxy = { version = "1.2.5", features = ["runtime-tokio"] }
+insta = "1.36.1"
 pretty_assertions = "1.4.0"
+reqwest = { version = "0.11", features = ["rustls-tls", "stream"], default-features = false }
 
 [profile.release]
 lto = true

diff --git a/assets/index.html b/assets/index.html
@@ -32,7 +32,7 @@
             flex-direction: column;
             flex: 1;
             border: 1px solid #ccc;
-            width: 100%;
+            width: 50%;
         }
 
         .left-panel {
@@ -223,7 +223,6 @@
     </style>
     <script src="https://unpkg.com/@highlightjs/[email protected]/highlight.min.js"></script>
     <script src="https://unpkg.com/clipboard@2/dist/clipboard.min.js"></script>
-    <script src="https://unpkg.com/[email protected]/dist/sanitize-html.min.js"></script>
 </head>
 
 <body>
@@ -757,7 +756,6 @@
             let { encode, value } = body;
             let contentType = getContentType(headers);
             if (encode == "utf8") {
-                value = sanitizeHtml(value);
                 let language = highlightLanguage(contentType);
                 const highlightedCode = hljs.getLanguage(language)
                     ? hljs.highlight(value, { language }).value
@@ -780,7 +778,7 @@
                         filename = disposition.match(/filename="(.*?)"/)[1] || filename;
                     }
                     let downloadUrl = URL.createObjectURL(base64ToBlob(value, contentType || "application/octet-stream"));
-                    return `<div class="media-view"><a href="${downloadUrl}" download="${filename}">Download the binary body</a></div>`;
+                    return `<div class="media-view"><a href="${downloadUrl}" download="${filename}">Download</a></div>`;
                 }
             }
             return `<div class="code-view">${value}</div>`;
@@ -805,8 +803,6 @@
             let value = data.body.value;
             if (data.body.encode == "base64") {
                 value = `data:;base64,${value}`;
-            } else {
-                value = sanitizeHtml(value);
             }
             return `<div class="ws-msg">
                 <small>${arrow}<span class="ws-date">${date}</span></small>

diff --git a/src/certificate_authority.rs → src/cert.rs b/src/certificate_authority.rs → src/cert.rs
@@ -8,55 +8,62 @@ use rcgen::{
 };
 use rsa::pkcs8::EncodePrivateKey;
 use rsa::RsaPrivateKey;
-use std::{fs, io::Cursor, sync::Arc};
+use std::{fs, io::Cursor, path::Path, sync::Arc};
 use time::{Duration, OffsetDateTime};
 use tokio_rustls::rustls::{
     pki_types::{CertificateDer, PrivateKeyDer},
     ServerConfig,
 };
 
-const CA_CERT_FILENAME: &str = "proxyfor-ca-cert.cer";
-const KEY_FILENAME: &str = "proxyfor-key.pem";
 const TTL_SECS: i64 = 365 * 24 * 60 * 60;
 const CACHE_TTL: u64 = TTL_SECS as u64 / 2;
 const NOT_BEFORE_OFFSET: i64 = 60;
 
-pub fn load_ca() -> Result<CertificateAuthority> {
-    let mut config_dir = dirs::home_dir().ok_or_else(|| anyhow!("No home dir"))?;
-    config_dir.push(".proxyfor");
-    if !config_dir.exists() {
-        fs::create_dir_all(&config_dir)
-            .with_context(|| format!("Failed to create config dir '{}'", config_dir.display()))?;
-    }
-
-    let ca_file = config_dir.join(CA_CERT_FILENAME);
-    let key_file = config_dir.join(KEY_FILENAME);
-    let (key, ca_cert, ca_data) = if !ca_file.exists() {
+pub fn init_ca(ca_cert_file: &Path, private_key_file: &Path) -> Result<CertificateAuthority> {
+    let (private_key, ca_cert, ca_data) = if !ca_cert_file.exists() {
         let key = gen_private_key().with_context(|| "Failed to generate private key")?;
         let ca_cert = gen_ca_cert(&key).with_context(|| "Failed to generate CA certificate")?;
         let ca_data = ca_cert
             .serialize_pem()
             .with_context(|| "Failed to generate CA certificate")?;
-        fs::write(&ca_file, &ca_data)
-            .with_context(|| format!("Failed to save CA certificate to '{}'", ca_file.display()))?;
-        fs::write(&key_file, key.serialize_pem())
-            .with_context(|| format!("Failed to save private key to '{}'", key_file.display()))?;
+        fs::write(ca_cert_file, &ca_data).with_context(|| {
+            format!(
+                "Failed to save CA certificate to '{}'",
+                ca_cert_file.display()
+            )
+        })?;
+        fs::write(private_key_file, key.serialize_pem()).with_context(|| {
+            format!(
+                "Failed to save private key to '{}'",
+                private_key_file.display()
+            )
+        })?;
         (key, ca_cert, ca_data)
     } else {
-        let key_err = || format!("Failed to read private key at '{}'", key_file.display());
-        let key_data = fs::read_to_string(&key_file).with_context(key_err)?;
+        let key_err = || {
+            format!(
+                "Failed to read private key at '{}'",
+                private_key_file.display()
+            )
+        };
+        let key_data = fs::read_to_string(private_key_file).with_context(key_err)?;
         let key = KeyPair::from_pem(&key_data).with_context(key_err)?;
         let key_clone = KeyPair::from_pem(&key_data).with_context(key_err)?;
 
-        let ca_err = || format!("Failed to read CA certificate at '{}'", ca_file.display());
-        let ca_data = fs::read_to_string(&ca_file).with_context(ca_err)?;
+        let ca_err = || {
+            format!(
+                "Failed to read CA certificate at '{}'",
+                ca_cert_file.display()
+            )
+        };
+        let ca_data = fs::read_to_string(ca_cert_file).with_context(ca_err)?;
         let ca_params =
             CertificateParams::from_ca_cert_pem(&ca_data, key_clone).with_context(ca_err)?;
         let ca_cert = Certificate::from_params(ca_params).with_context(ca_err)?;
         (key, ca_cert, ca_data)
     };
 
-    let mut key_reader = Cursor::new(key.serialize_pem());
+    let mut key_reader = Cursor::new(private_key.serialize_pem());
     let key_der = rustls_pemfile::read_one(&mut key_reader)
         .ok()
         .flatten()
@@ -68,7 +75,7 @@ pub fn load_ca() -> Result<CertificateAuthority> {
         })
         .ok_or_else(|| anyhow!("Invalid private key"))?;
 
-    let ca = CertificateAuthority::new(key, key_der, ca_cert, ca_data, 1_000);
+    let ca = CertificateAuthority::new(private_key, key_der, ca_cert, ca_data, 1_000);
     Ok(ca)
 }
 
@@ -127,7 +134,7 @@ impl CertificateAuthority {
         Ok(server_cfg)
     }
 
-    fn gen_cert(&self, authority: &Authority) -> Result<CertificateDer<'static>> {
+    pub fn gen_cert(&self, authority: &Authority) -> Result<CertificateDer<'static>> {
         let mut params = CertificateParams::default();
         params.serial_number = Some(thread_rng().gen::<u64>().into());
 

diff --git a/src/cli.rs b/src/cli.rs
diff --git a/src/filter.rs b/src/filter.rs
@@ -1,21 +1,21 @@
 use anyhow::{Context, Result};
 use fancy_regex::Regex;
 
-pub(crate) fn parse_filters(filters: &[String]) -> Result<Vec<Filter>> {
+pub fn parse_title_filters(filters: &[String]) -> Result<Vec<TitleFilter>> {
     let mut output = vec![];
     for filter in filters {
         if let Some(re) = filter.strip_prefix('/').and_then(|v| v.strip_suffix('/')) {
             let regex = Regex::new(&re.to_lowercase())
                 .with_context(|| format!("Invalid filter regex: {filter}"))?;
-            output.push(Filter::Regex(regex));
+            output.push(TitleFilter::Regex(regex));
         } else {
-            output.push(Filter::Plain(filter.to_lowercase()));
+            output.push(TitleFilter::Plain(filter.to_lowercase()));
         }
     }
     Ok(output)
 }
 
-pub(crate) fn is_match_title(filters: &[Filter], title: &str) -> bool {
+pub fn is_match_title(filters: &[TitleFilter], title: &str) -> bool {
     if filters.is_empty() {
         true
     } else {
@@ -24,7 +24,7 @@ pub(crate) fn is_match_title(filters: &[Filter], title: &str) -> bool {
     }
 }
 
-pub(crate) fn is_match_type(filters: &[String], content_type: &str) -> bool {
+pub fn is_match_type(filters: &[String], content_type: &str) -> bool {
     if filters.is_empty() {
         true
     } else {
@@ -34,23 +34,23 @@ pub(crate) fn is_match_type(filters: &[String], content_type: &str) -> bool {
 }
 
 #[derive(Debug)]
-pub(crate) enum Filter {
+pub enum TitleFilter {
     Regex(Regex),
     Plain(String),
 }
 
-impl Filter {
-    pub(crate) fn is_match(&self, value: &str) -> bool {
+impl TitleFilter {
+    pub fn is_match(&self, value: &str) -> bool {
         match self {
-            Filter::Regex(v) => v.is_match(value).unwrap_or_default(),
-            Filter::Plain(v) => value.contains(v),
+            TitleFilter::Regex(v) => v.is_match(value).unwrap_or_default(),
+            TitleFilter::Plain(v) => value.contains(v),
         }
     }
 }
 
 #[test]
 fn test_filters() {
-    let filters = parse_filters(&[
+    let filters = parse_title_filters(&[
         "postman-echo.com".to_string(),
         "/^(GET|POST) https:\\/\\/httpbin.org/".to_string(),
     ])

diff --git a/src/lib.rs b/src/lib.rs
@@ -0,0 +1,8 @@
+pub mod cert;
+pub mod filter;
+pub mod server;
+pub mod traffic;
+
+mod recorder;
+mod rewind;
+mod state;