feat: add client-side proxy auth and request parsing timeout

CONFIGURATION.md

@@ -43,6 +43,10 @@ remote = "https://your-server-domain/YOUR_SECRET_PATH"
 # address = "your-server-ip"
 # public_key = "your-public-key"
 
+# [client.auth]
+# username = "proxyuser"
+# password = "proxypass"
+
 [auth]
 token = "your-token"
 
@@ -221,4 +225,4 @@ upstream httproxy_backend {
     server unix:/dev/shm/httproxy.sock;
     keepalive 32;
 }
-```
+```

src/bypass/mod.rs

@@ -197,6 +197,12 @@ pub struct BypassRulesBuilder {
     ip_table: IpCidrTable,
 }
 
+impl Default for BypassRulesBuilder {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
 impl BypassRulesBuilder {
     pub fn new() -> Self {
         Self {

src/client/connection.rs

@@ -7,7 +7,10 @@ use tokio::net::TcpStream;
 use tracing::{Instrument, info, warn};
 
 use crate::client::{
-    constants::{CONNECT_RESPONSE, DOWNLOAD_CONNECT_TIMEOUT, EARLY_READ_WINDOW},
+    constants::{
+        CONNECT_RESPONSE, DOWNLOAD_CONNECT_TIMEOUT, EARLY_READ_WINDOW,
+        PROXY_AUTH_REQUIRED_RESPONSE, PROXY_REQUEST_PARSE_TIMEOUT,
+    },
     handshake::{self, try_pq_connect},
     proxy,
     state::SharedState,
@@ -23,7 +26,27 @@ pub async fn handle_connection(
     let (mut read_half, mut write_half) = socket.into_split();
 
     let mut buffer = BytesMut::with_capacity(16 * 1024);
-    let (method, header_len, url) = proxy::parse_proxy_request(&mut read_half, &mut buffer).await?;
+
+    let (method, header_len, url) = loop {
+        let (method, header_len, url, proxy_auth_header) = tokio::time::timeout(
+            PROXY_REQUEST_PARSE_TIMEOUT,
+            proxy::parse_proxy_request(&mut read_half, &mut buffer),
+        )
+        .await
+        .map_err(|_| anyhow!("proxy request parse timeout"))??;
+
+        if let Some((ref expected_auth, _)) = state.proxy_auth
+            && proxy_auth_header
+                .as_ref()
+                .is_none_or(|h| h.trim() != expected_auth.as_str())
+        {
+            write_half.write_all(PROXY_AUTH_REQUIRED_RESPONSE).await?;
+            write_half.flush().await?;
+            buffer.advance(header_len);
+            continue;
+        }
+        break (method, header_len, url);
+    };
 
     if method == "CONNECT" {
         buffer.advance(header_len);

src/client/constants.rs

@@ -1,9 +1,12 @@
 use std::time::Duration;
 
 pub const CONNECT_RESPONSE: &[u8] = b"HTTP/1.1 200 Connection Established\r\n\r\n";
+pub const PROXY_AUTH_REQUIRED_RESPONSE: &[u8] =
+    b"HTTP/1.1 407 Proxy Authentication Required\r\nProxy-Authenticate: Basic realm=\"httproxy\"\r\nContent-Length: 0\r\n\r\n";
 pub const EARLY_READ_WINDOW: Duration = Duration::from_millis(2);
 
 pub const DOWNLOAD_CONNECT_TIMEOUT: Duration = Duration::from_secs(15);
+pub const PROXY_REQUEST_PARSE_TIMEOUT: Duration = Duration::from_secs(10);
 pub const UPLOAD_REQUEST_TIMEOUT: Duration = Duration::from_secs(30);
 
 pub const MAX_BATCH_BYTES: usize = 1024 * 1024;

src/client/handshake.rs

@@ -20,6 +20,7 @@ use crate::client::constants::{
     DECODE_BUF_CAPACITY, DOWNLOAD_CONNECT_TIMEOUT, MIN_PADDING, PADDING_POOL,
 };
 
+#[allow(clippy::too_many_arguments)]
 pub async fn try_pq_connect(
     http_client: &Arc<wreq::Client>,
     state: &Arc<SharedState>,
@@ -139,6 +140,7 @@ pub async fn try_pq_connect(
     result
 }
 
+#[allow(clippy::explicit_auto_deref)]
 pub async fn full_handshake(
     http_client: &Arc<wreq::Client>,
     state: &Arc<SharedState>,

src/client/mod.rs

@@ -10,6 +10,7 @@ use crate::config::ClientTopConfig;
 use crate::crypto;
 
 use anyhow::{Context, Result};
+use base64::Engine;
 use std::sync::Arc;
 use tokio::sync::{Mutex, OnceCell};
 
@@ -36,12 +37,22 @@ pub fn build_state(cfg: &ClientTopConfig) -> Result<Arc<state::SharedState>> {
     let remote: url::Url = cfg.client.remote.parse().context("invalid server URL")?;
     let remote_str = remote.as_str().to_owned();
 
+    let proxy_auth = cfg.client.auth.as_ref().map(|a| {
+        let expected = format!(
+            "Basic {}",
+            base64::engine::general_purpose::STANDARD
+                .encode(format!("{}:{}", a.username, a.password))
+        );
+        (expected, a.username.clone())
+    });
+
     Ok(Arc::new(state::SharedState {
         remote_str,
         auth_header: format!("Bearer {}", cfg.auth.token),
         traffic_config: cfg.traffic_shaping.clone(),
         bypass,
         server_public_key,
+        proxy_auth,
         initial_master: Mutex::new(None),
         handshake_lock: OnceCell::new(),
     }))

src/client/proxy.rs

@@ -7,7 +7,7 @@ use url::Url;
 pub async fn parse_proxy_request(
     reader: &mut (impl AsyncReadExt + Unpin),
     buffer: &mut BytesMut,
-) -> Result<(String, usize, String)> {
+) -> Result<(String, usize, String, Option<String>)> {
     const MAX_HEADER_LEN: usize = 16 * 1024;
 
     loop {
@@ -17,10 +17,12 @@ pub async fn parse_proxy_request(
         let mut headers = [httparse::EMPTY_HEADER; 64];
         let mut req = httparse::Request::new(&mut headers);
         if let httparse::Status::Complete(amt) = req.parse(buffer)? {
+            let proxy_auth = extract_header(req.headers, "proxy-authorization");
             return Ok((
                 req.method.context("no method")?.to_owned(),
                 amt,
                 req.path.context("no path")?.to_owned(),
+                proxy_auth,
             ));
         }
         if buffer.len() > MAX_HEADER_LEN {
@@ -29,6 +31,13 @@ pub async fn parse_proxy_request(
     }
 }
 
+fn extract_header(headers: &[httparse::Header<'_>], name: &str) -> Option<String> {
+    headers
+        .iter()
+        .find(|h| h.name.eq_ignore_ascii_case(name))
+        .map(|h| String::from_utf8_lossy(h.value).into_owned())
+}
+
 #[inline]
 pub fn resolve_target_host(method: &str, url_str: &str) -> Result<String> {
     if method == "CONNECT" {

src/client/state.rs

@@ -7,6 +7,8 @@ use zeroize::Zeroizing;
 use crate::bypass::BypassRules;
 use crate::shaper::TrafficConfig;
 
+pub type InitialMasterEntry = (String, Zeroizing<[u8; 32]>, Instant);
+
 pub struct ManualResolver {
     pub target_addr: String,
 }
@@ -37,6 +39,7 @@ pub struct SharedState {
     pub traffic_config: TrafficConfig,
     pub bypass: Option<Arc<BypassRules>>,
     pub server_public_key: Option<x25519_dalek::PublicKey>,
-    pub initial_master: Mutex<Option<(String, Zeroizing<[u8; 32]>, Instant)>>,
+    pub proxy_auth: Option<(String, String)>,
+    pub initial_master: Mutex<Option<InitialMasterEntry>>,
     pub handshake_lock: OnceCell<tokio::sync::Mutex<()>>,
 }

src/config/mod.rs

@@ -55,6 +55,15 @@ pub struct ClientSection {
     pub address: Option<String>,
     #[serde(default)]
     pub public_key: Option<String>,
+    #[serde(default)]
+    pub auth: Option<ClientProxyAuth>,
+}
+
+#[derive(Deserialize, Debug, Clone)]
+#[serde(deny_unknown_fields)]
+pub struct ClientProxyAuth {
+    pub username: String,
+    pub password: String,
 }
 
 #[derive(Deserialize, Debug)]

src/server/handlers.rs

@@ -190,6 +190,7 @@ async fn handle_plaintext_download(
     Ok(resp)
 }
 
+#[allow(clippy::explicit_auto_deref)]
 async fn handle_fresh_handshake(
     state: Arc<AppState>,
     headers: HeaderMap,
@@ -315,6 +316,7 @@ async fn handle_fresh_handshake(
     Ok(resp)
 }
 
+#[allow(clippy::explicit_auto_deref)]
 async fn handle_pq_download(
     state: Arc<AppState>,
     cookie_val: &str,

src/server/janitor.rs

@@ -26,7 +26,7 @@ pub async fn stream_janitor(streams: Arc<DashMap<String, Arc<UploadStream>>>) {
 }
 
 pub async fn master_and_nonce_janitor(
-    master_store: Arc<DashMap<String, (String, zeroize::Zeroizing<[u8; 32]>, std::time::Instant)>>,
+    master_store: Arc<DashMap<String, super::MasterStoreEntry>>,
     used_nonces: Arc<DashMap<String, DashSet<[u8; 16]>>>,
 ) {
     let mut interval = tokio::time::interval(NONCE_CLEANUP_INTERVAL);