support option proxy url

src/main.rs

@@ -90,7 +90,25 @@ fn request_to_http_text(req: &Request) -> anyhow::Result<String> {
     Ok(result)
 }
 
-async fn rustls_req(url_str: &str, proxy: &str) -> anyhow::Result<std::time::Duration> {
+async fn connect_to(url: &str, proxy: Option<&str>) -> anyhow::Result<tokio::net::TcpStream> {
+    let url: Url = url.parse()?;
+    if let Some(proxy) = proxy {
+        let proxy_url: Url = proxy.parse()?;
+        let mut ac = tokio::net::TcpStream::connect((
+            proxy_url.host().unwrap().to_string(),
+            proxy_url.port().unwrap(),
+        ))
+        .await?;
+        client::connect(&mut ac, (url.host().unwrap().to_string(), 443), None).await?;
+        Ok(ac)
+    } else {
+        let port = url.port_or_known_default().unwrap();
+        let ac = tokio::net::TcpStream::connect((url.host().unwrap().to_string(), port))
+            .await?;
+        Ok(ac)
+    }
+}
+async fn rustls_req(url_str: &str, proxy: Option<&str>) -> anyhow::Result<std::time::Duration> {
     use rustls_pki_types::ServerName;
     use std::sync::Arc;
     use tokio_rustls::rustls::{ClientConfig, RootCertStore};
@@ -109,16 +127,17 @@ async fn rustls_req(url_str: &str, proxy: &str) -> anyhow::Result<std::time::Dur
 
     // let stream = TcpStream::connect(&addr).await?;
     // let mut stream = connector.connect(dnsname, stream).await?;
-    let proxy_url: Url = proxy.parse()?;
+    // let proxy_url: Url = proxy.parse()?;
     let now = std::time::Instant::now();
-    let mut ac = tokio::net::TcpStream::connect((
+    // let mut ac = tokio::net::TcpStream::connect((
-        proxy_url.host().unwrap().to_string(),
+    //     proxy_url.host().unwrap().to_string(),
-        proxy_url.port().unwrap(),
+    //     proxy_url.port().unwrap(),
-    ))
+    // ))
-    .await?;
+    // .await?;
-    log::debug!("tcp conn established time used: {:?}", now.elapsed());
+    // log::debug!("tcp conn established time used: {:?}", now.elapsed());
-    let r = client::connect(&mut ac, (url.host().unwrap().to_string(), 443), None).await?;
+    // let r = client::connect(&mut ac, (url.host().unwrap().to_string(), 443), None).await?;
-    log::debug!("socks5 conn established time used: {:?}", now.elapsed());
+    // log::debug!("socks5 conn established time used: {:?}", now.elapsed());
+    let ac = connect_to(url_str, proxy).await?;
     // let tls_conn = TlsConnector::from(native_tls::TlsConnector::new().unwrap());
     let mut tls_s = connector.connect(dnsname, ac).await?;
     // let url: Url = "https://cp.cloudflare.com/generate_204".parse().unwrap();
@@ -135,18 +154,19 @@ async fn rustls_req(url_str: &str, proxy: &str) -> anyhow::Result<std::time::Dur
     Ok(now.elapsed())
 }
 
-async fn http_req(url_str: &str, proxy: &str) -> anyhow::Result<std::time::Duration> {
+async fn http_req(url_str: &str, proxy: Option<&str>) -> anyhow::Result<std::time::Duration> {
     let url: Url = url_str.parse()?;
 
     // let stream = TcpStream::connect(&addr).await?;
     // let mut stream = connector.connect(dnsname, stream).await?;
-    let proxy_url: Url = proxy.parse()?;
+    // let proxy_url: Url = proxy.parse()?;
     let now = std::time::Instant::now();
-    let mut ac = tokio::net::TcpStream::connect((
+    // let mut ac = tokio::net::TcpStream::connect((
-        proxy_url.host().unwrap().to_string(),
+    //     proxy_url.host().unwrap().to_string(),
-        proxy_url.port().unwrap(),
+    //     proxy_url.port().unwrap(),
-    ))
+    // ))
-    .await?;
+    // .await?;
+    let mut ac = connect_to(url_str, proxy).await?;
     log::debug!("tcp conn established time used: {:?}", now.elapsed());
     let r = client::connect(
         &mut ac,
@@ -180,7 +200,7 @@ async fn main() {
         .unwrap_or_else(|| "20".to_owned())
         .parse()
         .expect("second parameter should interger");
-    let proxy = std::env::args().nth(3).unwrap();
+    let proxy = std::env::args().nth(3);
     let fmt = std::env::args().nth(4);
 
     let mut stats = IncrementalStats::new();
@@ -188,9 +208,9 @@ async fn main() {
     let mut err = 0;
     for i in 0..iter_num {
         let fut = if url.starts_with("https") {
-            Either::Left(rustls_req(&url, &proxy))
+            Either::Left(rustls_req(&url, proxy.as_deref()))
         } else {
-            Either::Right(http_req(&url, &proxy))
+            Either::Right(http_req(&url, proxy.as_deref()))
         };
         match tokio::time::timeout(std::time::Duration::from_secs(5), fut).await {
             Ok(Ok(duration)) => {