Direct send for none ciphers, prevent unnecessary data copy

Author: zonyitoo

src/crypto/cipher.rs

@@ -263,6 +263,8 @@ pub enum CipherType {
 /// Category of ciphers
 #[derive(Debug, Copy, Clone, Eq, PartialEq)]
 pub enum CipherCategory {
+    /// No encryption
+    None,
     /// Stream ciphers is used for OLD ShadowSocks protocol, which uses stream ciphers to encrypt data payloads
     Stream,
     /// AEAD ciphers is used in modern ShadowSocks protocol, which sends data in separate packets
@@ -542,6 +544,8 @@ impl CipherType {
     /// Get category of cipher
     pub fn category(self) -> CipherCategory {
         match self {
+            CipherType::None | CipherType::Plain => CipherCategory::None,
+
             CipherType::Aes128Gcm | CipherType::Aes256Gcm | CipherType::ChaCha20IetfPoly1305 => CipherCategory::Aead,
 
             #[cfg(feature = "sodium")]

src/relay/dnsrelay/mod.rs

@@ -31,7 +31,7 @@ fn should_forward_by_ptr_name(acl: &AccessControl, name: &Name) -> bool {
     let mut iter = name.iter().rev();
     let mut next = || match iter.next() {
         Some(label) => std::str::from_utf8(label).unwrap_or("*"),
-        None => "0",    // zero fill the missing labels
+        None => "0", // zero fill the missing labels
     };
     if !"arpa".eq_ignore_ascii_case(next()) {
         return acl.is_default_in_proxy_list();
@@ -56,7 +56,14 @@ fn should_forward_by_ptr_name(acl: &AccessControl, name: &Name) -> bool {
                 }
             }
             acl.check_ip_in_proxy_list(&IpAddr::V6(Ipv6Addr::new(
-                segments[0], segments[1], segments[2], segments[3], segments[4], segments[5], segments[6], segments[7]
+                segments[0],
+                segments[1],
+                segments[2],
+                segments[3],
+                segments[4],
+                segments[5],
+                segments[6],
+                segments[7],
             )))
         }
         _ => acl.is_default_in_proxy_list(),
@@ -118,7 +125,7 @@ fn should_forward_by_response(
                     } else {
                         acl.is_default_in_proxy_list()
                     }
-                }}
+                }};
             }
             macro_rules! examine_record {
                 ($rec:ident, $is_answer:expr) => {
@@ -132,8 +139,11 @@ fn should_forward_by_response(
                         continue;
                     }
                     if $is_answer && !query.query_type().is_any() && $rec.record_type() != query.query_type() {
-                        warn!("local DNS response has inconsistent answer type {} for query {}",
-                              $rec.record_type(), query);
+                        warn!(
+                            "local DNS response has inconsistent answer type {} for query {}",
+                            $rec.record_type(),
+                            query
+                        );
                         return true;
                     }
                     let forward = match $rec.rdata() {
@@ -153,7 +163,11 @@ fn should_forward_by_response(
             }
             for rec in local_response.answers() {
                 if !names.contains(rec.name()) {
-                    warn!("local DNS response contains unexpected name {} for query {}", rec.name(), query);
+                    warn!(
+                        "local DNS response contains unexpected name {} for query {}",
+                        rec.name(),
+                        query
+                    );
                     return true;
                 }
                 examine_record!(rec, true);
@@ -183,7 +197,10 @@ impl<Remote: upstream::Upstream> DnsRelay<Remote> {
         // Start querying name servers
         debug!(
             "attempting lookup of {:?} {} with ns {:?} and {:?}",
-            query.query_type(), query.name(), local, remote
+            query.query_type(),
+            query.name(),
+            local,
+            remote
         );
 
         let remote_response_fut = try_timeout(remote.lookup(query), Some(Duration::new(3, 0)));
@@ -274,7 +291,7 @@ impl<Remote: upstream::Upstream> DnsRelay<Remote> {
 
 async fn run_tcp<Remote: upstream::Upstream + Send + Sync + 'static>(
     relay: Arc<DnsRelay<Remote>>,
-    bind_addr: SocketAddr
+    bind_addr: SocketAddr,
 ) -> io::Result<()> {
     let mut listener = TcpListener::bind(&bind_addr).await?;
 
@@ -302,7 +319,7 @@ async fn run_tcp<Remote: upstream::Upstream + Send + Sync + 'static>(
 
 async fn run_udp<Remote: upstream::Upstream + Send + Sync + 'static>(
     relay: Arc<DnsRelay<Remote>>,
-    bind_addr: SocketAddr
+    bind_addr: SocketAddr,
 ) -> io::Result<()> {
     let socket = create_udp_socket(&bind_addr).await?;
 
@@ -379,12 +396,13 @@ pub async fn run(context: SharedContext) -> io::Result<()> {
             context: context.clone(),
             svr_cfg: move || balancer.pick_server().server_config().clone(),
             ns: config.remote_dns_addr.clone().expect("remote query DNS address"),
-        }
+        },
     });
 
     future::select(
         tokio::spawn(run_tcp(relay.clone(), bind_addr)),
         tokio::spawn(run_udp(relay, bind_addr)),
-    ).await;
+    )
+    .await;
     Ok(())
 }

src/relay/dnsrelay/upstream.rs

@@ -10,11 +10,11 @@ use byteorder::{BigEndian, ByteOrder};
 use rand::Rng;
 use tokio::{
     io::{AsyncReadExt, AsyncWriteExt},
-    net::UdpSocket
+    net::UdpSocket,
 };
 use trust_dns_proto::{
     op::{Message, Query},
-    rr::{DNSClass, RecordType, Name, RData},
+    rr::{DNSClass, Name, RData, RecordType},
 };
 
 #[cfg(unix)]
@@ -25,10 +25,7 @@ use tokio::net::UnixStream;
 use crate::{
     config::{Config, ServerConfig},
     context::SharedContext,
-    relay::{
-        socks5::Address,
-        tcprelay::ProxyStream,
-    },
+    relay::{socks5::Address, tcprelay::ProxyStream},
 };
 
 #[derive(Debug)]
@@ -118,8 +115,8 @@ pub async fn write_message<T: AsyncWriteExt + Unpin>(stream: &mut T, message: &M
 }
 
 async fn stream_lookup<T>(query: &Query, stream: &mut T) -> io::Result<Message>
-    where
-        T: AsyncReadExt + AsyncWriteExt + Unpin,
+where
+    T: AsyncReadExt + AsyncWriteExt + Unpin,
 {
     write_message(stream, &generate_query_message(query)).await?;
     read_message(stream).await
@@ -133,10 +130,14 @@ pub struct UdpUpstream {
 #[async_trait]
 impl Upstream for UdpUpstream {
     async fn lookup(&self, query: &Query) -> io::Result<Message> {
-        let mut socket = UdpSocket::bind(SocketAddr::new(match self.server {
-            SocketAddr::V4(..) => IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)),
-            SocketAddr::V6(..) => IpAddr::V6(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 0)),
-        }, 0)).await?;
+        let mut socket = UdpSocket::bind(SocketAddr::new(
+            match self.server {
+                SocketAddr::V4(..) => IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)),
+                SocketAddr::V6(..) => IpAddr::V6(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 0)),
+            },
+            0,
+        ))
+        .await?;
         socket.connect(self.server).await?;
         socket.send(&generate_query_message(query).to_vec()?).await?;
         let mut response = vec![0; 512];
@@ -165,14 +166,15 @@ pub struct ProxyTcpUpstream<F> {
 
 impl<F> Debug for ProxyTcpUpstream<F> {
     fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
-        f.debug_struct("ProxyTcpUpstream")
-            .field("ns", &self.ns)
-            .finish()
+        f.debug_struct("ProxyTcpUpstream").field("ns", &self.ns).finish()
     }
 }
 
 #[async_trait]
-impl<F> Upstream for ProxyTcpUpstream<F> where F: Fn() -> ServerConfig + Send + Sync {
+impl<F> Upstream for ProxyTcpUpstream<F>
+where
+    F: Fn() -> ServerConfig + Send + Sync,
+{
     async fn lookup(&self, query: &Query) -> io::Result<Message> {
         let mut stream = ProxyStream::connect_proxied(self.context.clone(), &(self.svr_cfg)(), &self.ns).await?;
         stream_lookup(query, &mut stream).await

src/relay/tcprelay/crypto_io.rs

@@ -28,11 +28,13 @@ use super::{
 };
 
 enum DecryptedReader {
+    None,
     Aead(AeadDecryptedReader),
     Stream(StreamDecryptedReader),
 }
 
 enum EncryptedWriter {
+    None,
     Aead(AeadEncryptedWriter),
     Stream(StreamEncryptedWriter),
 }
@@ -62,9 +64,14 @@ impl<S> CryptoStream<S> {
     /// Create a new CryptoStream with the underlying stream connection
     pub fn new(context: SharedContext, stream: S, svr_cfg: &ServerConfig) -> CryptoStream<S> {
         let method = svr_cfg.method();
+        if method.category() == CipherCategory::None {
+            return CryptoStream::<S>::new_none(stream);
+        }
+
         let prev_len = match method.category() {
             CipherCategory::Stream => method.iv_size(),
             CipherCategory::Aead => method.salt_size(),
+            CipherCategory::None => 0,
         };
 
         let iv = match method.category() {
@@ -92,12 +99,14 @@ impl<S> CryptoStream<S> {
                 trace!("generated AEAD cipher salt {:?}", local_salt);
                 local_salt
             }
+            CipherCategory::None => Bytes::new(),
         };
 
         let method = svr_cfg.method();
         let enc = match method.category() {
             CipherCategory::Stream => EncryptedWriter::Stream(StreamEncryptedWriter::new(method, svr_cfg.key(), iv)),
             CipherCategory::Aead => EncryptedWriter::Aead(AeadEncryptedWriter::new(method, svr_cfg.key(), iv)),
+            CipherCategory::None => EncryptedWriter::None,
         };
 
         CryptoStream {
@@ -108,6 +117,15 @@ impl<S> CryptoStream<S> {
         }
     }
 
+    fn new_none(stream: S) -> CryptoStream<S> {
+        CryptoStream {
+            stream,
+            dec: Some(DecryptedReader::None),
+            enc: EncryptedWriter::None,
+            read_status: ReadStatus::Established,
+        }
+    }
+
     /// Return a reference to the underlying stream
     pub fn get_ref(&self) -> &S {
         &self.stream
@@ -148,6 +166,7 @@ where
                     trace!("got AEAD cipher salt {:?}", ByteStr::new(&buf));
                     DecryptedReader::Aead(AeadDecryptedReader::new(method, key, &buf))
                 }
+                CipherCategory::None => DecryptedReader::None,
             };
 
             self.dec = Some(dec);
@@ -162,6 +181,7 @@ where
         ready!(this.poll_read_handshake(ctx))?;
 
         match *this.dec.as_mut().unwrap() {
+            DecryptedReader::None => Pin::new(&mut this.stream).poll_read(ctx, buf),
             DecryptedReader::Aead(ref mut r) => r.poll_read_decrypted(ctx, &mut this.stream, buf),
             DecryptedReader::Stream(ref mut r) => r.poll_read_decrypted(ctx, &mut this.stream, buf),
         }
@@ -175,6 +195,7 @@ where
     fn priv_poll_write(self: Pin<&mut Self>, ctx: &mut Context<'_>, buf: &[u8]) -> Poll<io::Result<usize>> {
         let this = self.get_mut();
         match this.enc {
+            EncryptedWriter::None => Pin::new(&mut this.stream).poll_write(ctx, buf),
             EncryptedWriter::Aead(ref mut w) => w.poll_write_encrypted(ctx, &mut this.stream, buf),
             EncryptedWriter::Stream(ref mut w) => w.poll_write_encrypted(ctx, &mut this.stream, buf),
         }