Fix domain sniffing/override

Author: eycorsican

leaf/src/app/dispatcher.rs

@@ -209,15 +209,23 @@ impl Dispatcher {
                         match kind {
                             sniff::SniffKind::Tls => {
                                 if do_tls {
-                                    sess.tls_sniffed_domain = Some(domain);
+                                    sess.tls_sniffed_domain = Some(domain.clone());
                                 }
                             }
                             sniff::SniffKind::Http => {
                                 if do_http {
-                                    sess.http_sniffed_domain = Some(domain);
+                                    sess.http_sniffed_domain = Some(domain.clone());
                                 }
                             }
                         }
+
+                        if option::DOMAIN_OVERRIDE.load(std::sync::atomic::Ordering::Relaxed) {
+                            if let Ok(dest) = SocksAddr::try_from((domain, sess.destination.port()))
+                            {
+                                debug!("override destination with sniffed domain={}", dest);
+                                sess.destination = dest;
+                            }
+                        }
                     }
                 }
                 Err(e) => {

leaf/src/app/router.rs

@@ -564,9 +564,7 @@ impl Router {
     }
 
     pub async fn pick_route<'a>(&'a self, sess: &'a Session) -> Result<Option<&'a String>> {
-        let effective_dest = sess
-            .effective_destination()
-            .unwrap_or_else(|_| std::borrow::Cow::Borrowed(&sess.destination));
+        let effective_dest = &sess.destination;
         for rule in &self.rules {
             if rule.apply(sess) {
                 return Ok(Some(&rule.target));

leaf/src/option/mod.rs

@@ -124,6 +124,12 @@ lazy_static! {
         AtomicBool::new(v)
     };
 
+    /// Override the original destination with the sniffed domain.
+    pub static ref DOMAIN_OVERRIDE: AtomicBool = {
+        let v: bool = get_env_var_or("DOMAIN_OVERRIDE", false);
+        AtomicBool::new(v)
+    };
+
     /// Turn on DNS sniffing, if the destination is an IP, we try to find the
     /// domain from the DNS cache.
     pub static ref DNS_DOMAIN_SNIFFING: AtomicBool = {

leaf/src/proxy/failover/stream.rs

@@ -148,7 +148,7 @@ impl OutboundStreamHandler for Handler {
 
         if let Some(cache) = &self.cache {
             // Try the cached actor first if exists.
-            let cache_key = sess.effective_destination()?.to_string();
+            let cache_key = sess.destination.to_string();
             if let Some(idx) = cache.lock().await.get(&cache_key) {
                 let a = &self.actors[*idx];
                 debug!(

leaf/src/proxy/mod.rs

@@ -403,7 +403,7 @@ pub async fn connect_stream_outbound(
             Ok(Some(new_tcp_stream(dns_client, &addr, &port).await?))
         }
         OutboundConnect::Direct => {
-            let dest = sess.effective_destination()?;
+            let dest = &sess.destination;
             trace!("connect stream direct dst={}", &dest);
             Ok(Some(
                 new_tcp_stream(dns_client, &dest.host(), &dest.port()).await?,
@@ -437,7 +437,7 @@ pub async fn connect_datagram_outbound(
                 Ok(Some(OutboundTransport::Stream(stream)))
             }
         },
-        OutboundConnect::Direct => match sess.effective_destination()?.as_ref() {
+        OutboundConnect::Direct => match &sess.destination {
             SocksAddr::Domain(domain, port) => {
                 let socket = new_udp_socket(&crate::option::UNSPECIFIED_BIND_ADDR).await?;
                 Ok(Some(OutboundTransport::Datagram(Box::new(

leaf/src/proxy/shadowsocks/outbound/datagram.rs

@@ -42,7 +42,7 @@ impl OutboundDatagramHandler for Handler {
 
         let dgram = ShadowedDatagram::new(&self.cipher, &self.password)?;
 
-        let destination = match sess.effective_destination()?.as_ref() {
+        let destination = match &sess.destination {
             SocksAddr::Domain(domain, port) => Some(SocksAddr::Domain(domain.to_owned(), *port)),
             _ => None,
         };

leaf/src/proxy/shadowsocks/outbound/stream.rs

@@ -61,7 +61,7 @@ impl OutboundStreamHandler for Handler {
             self.prefix.as_ref().cloned(),
         )?;
         let mut buf = BytesMut::new();
-        sess.effective_destination()?
+        sess.destination
             .write_buf(&mut buf, SocksAddrWireType::PortLast);
 
         let payload = peek_tcp_one_off(lhs).await;

leaf/src/proxy/socks/outbound/stream.rs

@@ -34,7 +34,7 @@ impl OutboundStreamHandler for Handler {
                 password: auth_password.to_owned(),
             }),
         };
-        match sess.effective_destination()?.as_ref() {
+        match &sess.destination {
             SocksAddr::Ip(a) => {
                 let _ = async_socks5::connect(&mut stream, a.to_owned(), auth)
                     .map_err(io::Error::other)

leaf/src/proxy/tls/outbound/stream.rs

@@ -206,7 +206,7 @@ impl OutboundStreamHandler for Handler {
         let name = if !&self.server_name.is_empty() {
             self.server_name.clone()
         } else {
-            sess.effective_destination()?.host()
+            sess.destination.host()
         };
         if let Some(stream) = stream {
             #[cfg(feature = "rustls-tls")]

leaf/src/proxy/trojan/outbound/datagram.rs

@@ -42,11 +42,11 @@ impl OutboundDatagramHandler for Handler {
         buf.put_slice(password.as_bytes());
         buf.put_slice(b"\r\n");
         buf.put_u8(0x03); // udp
-        sess.effective_destination()?
+        sess.destination
             .write_buf(&mut buf, SocksAddrWireType::PortLast);
         buf.put_slice(b"\r\n");
 
-        let destination = match sess.effective_destination()?.as_ref() {
+        let destination = match &sess.destination {
             SocksAddr::Domain(domain, port) => Some(SocksAddr::Domain(domain.to_owned(), *port)),
             _ => None,
         };