Socks4とDNSキャッシュを実装

WinLinux1028 · WinLinux1028 · commit 4375beba2b35 · 2023-11-03T03:13:45.000+09:00
diff --git a/Cargo.toml b/Cargo.toml
@@ -21,3 +21,4 @@ futures-util = "0.3"
 bytes = "1"
 idna = "0.4"
 percent-encoding = "2"
+lru_time_cache = "0.11"
diff --git a/src/main.rs b/src/main.rs
@@ -13,8 +13,17 @@ use hyper::{
     service::{make_service_fn, service_fn},
     Body, Method, Request, Response, Server, StatusCode,
 };
-use std::{io::Write, time::Duration};
-use tokio::io::{AsyncRead, AsyncWrite};
+use lru_time_cache::LruCache;
+use std::{
+    hash::Hash,
+    io::Write,
+    net::{Ipv4Addr, Ipv6Addr},
+    time::Duration,
+};
+use tokio::{
+    io::{AsyncRead, AsyncWrite},
+    sync::Mutex,
+};
 
 use once_cell::sync::OnceCell;
 
@@ -81,15 +90,24 @@ async fn main() {
             let proxy_protocol_main = proxy_protocol[proxy_protocol.len() - 1];
             if proxy_protocol_main == "http" {
                 proxy_stack.push(Box::new(outbound::HttpProxy::new(proxy).unwrap()));
+            } else if proxy_protocol_main == "socks4" {
+                proxy_stack.push(Box::new(outbound::Socks4Proxy::new(proxy).unwrap()));
             } else {
                 panic!("This protocol can not use: {}", proxy_protocol_main);
             }
         }
     }
 
+    let dns_cache = if config.doh_endpoint.is_some() {
+        LruCache::with_expiry_duration_and_capacity(Duration::from_secs(7200), 65535)
+    } else {
+        LruCache::with_capacity(0)
+    };
+
     if PROXY
         .set(ProxyState {
             config,
+            dns_cache: Mutex::new(dns_cache),
             proxy_stack,
         })
         .is_err()
@@ -121,10 +139,18 @@ async fn handle(request: Request<Body>) -> Result<Response<Body>, Error> {
     }
 }
 
+#[allow(clippy::type_complexity)]
 struct ProxyState {
     config: Config,
+    dns_cache: Mutex<LruCache<String, (DnsCacheState<Ipv4Addr>, DnsCacheState<Ipv6Addr>)>>,
     proxy_stack: Vec<Box<dyn ProxyOutBound>>,
 }
 
+enum DnsCacheState<T: Hash> {
+    Some(T),
+    Fail,
+    None,
+}
+
 pub trait Stream: AsyncRead + AsyncWrite {}
 impl<RW> Stream for RW where RW: AsyncRead + AsyncWrite {}
diff --git a/src/outbound/socks4.rs b/src/outbound/socks4.rs
@@ -1,10 +1,14 @@
 use super::ProxyOutBound;
-use crate::{config::ProxyConfig, utils::SocketAddr, Connection, Error};
+use crate::{
+    config::ProxyConfig,
+    utils::{HostName, SocketAddr},
+    Connection, Error,
+};
 
 use std::str::FromStr;
 
 use async_trait::async_trait;
-use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
+use tokio::io::{AsyncReadExt, AsyncWriteExt};
 
 pub struct Socks4Proxy {
     addr: SocketAddr,
@@ -40,14 +44,57 @@ impl ProxyOutBound for Socks4Proxy {
         mut proxies: Box<dyn Iterator<Item = &Box<dyn ProxyOutBound>> + Send>,
         addr: &SocketAddr,
     ) -> Result<Connection, Error> {
-        let server = proxies
+        let ip;
+        let mut hostname = None;
+        match &addr.hostname {
+            HostName::V4(v4) => ip = *v4,
+            HostName::Domain(domain) => {
+                ip = "0.0.0.1".parse()?;
+                if domain.contains('\0') {
+                    return Err("".into());
+                }
+                hostname = Some(domain);
+            }
+            HostName::V6(_) => return Err("".into()),
+        }
+
+        let ip_octet = ip.octets();
+        if hostname.is_none()
+            && ip_octet[0] == 0
+            && ip_octet[1] == 0
+            && ip_octet[2] == 0
+            && ip_octet[3] != 0
+        {
+            return Err("".into());
+        }
+
+        let mut server = proxies
             .next()
             .ok_or("")?
             .connect(proxies, &self.addr)
             .await?;
-        let mut server = BufReader::new(server);
 
         server.write_all(&[4, 1]).await?;
+        server.write_all(&addr.port.to_be_bytes()).await?;
+        server.write_all(&ip_octet).await?;
+        if let Some(auth) = &self.auth {
+            server.write_all(auth.as_bytes()).await?
+        }
+        server.write_all(b"\0").await?;
+        if let Some(hostname) = hostname {
+            server.write_all(hostname.as_bytes()).await?;
+            server.write_all(b"\0").await?;
+        }
+        server.flush().await?;
+
+        if server.read_u8().await? != 0 {
+            return Err("".into());
+        }
+        if server.read_u8().await? != 90 {
+            return Err("".into());
+        }
+        let mut buf = [0; 6];
+        server.read_exact(&mut buf).await?;
 
         Ok(Box::new(server))
     }
diff --git a/src/utils/addr.rs b/src/utils/addr.rs
@@ -1,5 +1,5 @@
 use super::ParsedUri;
-use crate::{http_proxy, Connection, Error, PROXY};
+use crate::{http_proxy, Connection, DnsCacheState, Error, PROXY};
 
 use std::{
     fmt::{Display, Write},
@@ -150,16 +150,37 @@ impl HostName {
             _ => return Err("".into()),
         };
 
+        let proxy = PROXY.get().ok_or("")?;
+        let mut uri: ParsedUri =
+            Uri::from_str(proxy.config.doh_endpoint.as_ref().ok_or("")?)?.try_into()?;
+
+        let mut dns_cache = proxy.dns_cache.lock().await;
+        if let Some(cache_content) = dns_cache.get(domain) {
+            if qtype == QueryType::A {
+                match cache_content.0 {
+                    DnsCacheState::Some(s) => return Ok(Self::V4(s)),
+                    DnsCacheState::Fail => return Err("".into()),
+                    DnsCacheState::None => (),
+                }
+            } else if qtype == QueryType::AAAA {
+                match cache_content.1 {
+                    DnsCacheState::Some(s) => return Ok(Self::V6(s)),
+                    DnsCacheState::Fail => return Err("".into()),
+                    DnsCacheState::None => (),
+                }
+            } else {
+                return Err("".into());
+            }
+        }
+        drop(dns_cache);
+
         let mut query = dns_parser::Builder::new_query(0xabcd, true);
         query.add_question(domain, false, qtype, QueryClass::IN);
         let query = query.build().map_err(|_| "")?;
 
         let base64 = base64::engine::general_purpose::URL_SAFE_NO_PAD;
         let query = base64.encode(query);
 
-        let proxy = PROXY.get().ok_or("")?;
-        let mut uri: ParsedUri =
-            Uri::from_str(proxy.config.doh_endpoint.as_ref().ok_or("")?)?.try_into()?;
         if let Some(s) = uri.query.as_mut() {
             s.push_str(&format!("&dns={}", query));
         } else {
@@ -184,17 +205,33 @@ impl HostName {
         }
         let response_body = dns_parser::Packet::parse(&response_body)?;
 
+        let mut dns_cache = proxy.dns_cache.lock().await;
+        let cache_content = dns_cache
+            .entry(domain.to_string())
+            .or_insert((DnsCacheState::None, DnsCacheState::None));
+
         for answer in response_body.answers {
             if answer.cls != dns_parser::Class::IN {
                 continue;
             }
             match answer.data {
-                RData::A(addr) => return Ok(Self::V4(addr.0)),
-                RData::AAAA(addr) => return Ok(Self::V6(addr.0)),
+                RData::A(addr) => {
+                    cache_content.0 = DnsCacheState::Some(addr.0);
+                    return Ok(Self::V4(addr.0));
+                }
+                RData::AAAA(addr) => {
+                    cache_content.1 = DnsCacheState::Some(addr.0);
+                    return Ok(Self::V6(addr.0));
+                }
                 _ => continue,
             }
         }
 
+        if qtype == QueryType::A {
+            cache_content.0 = DnsCacheState::Fail;
+        } else if qtype == QueryType::AAAA {
+            cache_content.1 = DnsCacheState::Fail;
+        }
         Err("".into())
     }
 }
