Send first packet with IV/Nonce to reduce latency

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "shadowsocks-rust"
-version = "1.8.0"
+version = "1.8.1"
 authors = ["Y. T. CHUNG <zonyitoo@gmail.com>"]
 description = "shadowsocks is a fast tunnel proxy that helps you bypass firewalls."
 repository = "https://github.com/zonyitoo/shadowsocks-rust"

src/relay/tcprelay/aead.rs

@@ -43,7 +43,7 @@ use std::{
 };
 
 use byteorder::{BigEndian, ByteOrder};
-use bytes::{BufMut, BytesMut};
+use bytes::{BufMut, Bytes, BytesMut};
 use futures::ready;
 use tokio::prelude::*;
 
@@ -217,15 +217,17 @@ pub struct EncryptedWriter {
     cipher: BoxAeadEncryptor,
     tag_size: usize,
     steps: EncryptWriteStep,
+    nonce: Option<Bytes>,
 }
 
 impl EncryptedWriter {
     /// Creates a new EncryptedWriter
-    pub fn new(t: CipherType, key: &[u8], nonce: &[u8]) -> EncryptedWriter {
+    pub fn new(t: CipherType, key: &[u8], nonce: Bytes) -> EncryptedWriter {
         EncryptedWriter {
-            cipher: crypto::new_aead_encryptor(t, key, nonce),
+            cipher: crypto::new_aead_encryptor(t, key, &nonce),
             tag_size: t.tag_size(),
             steps: EncryptWriteStep::Nothing,
+            nonce: Some(nonce),
         }
     }
 
@@ -253,7 +255,18 @@ impl EncryptedWriter {
                     let output_length = self.buffer_size(data);
                     let data_length = data.len() as u16;
 
-                    let mut buf = BytesMut::with_capacity(output_length);
+                    // Send the first packet with nonce
+                    let nonce_length = match self.nonce {
+                        Some(ref n) => n.len(),
+                        None => 0,
+                    };
+
+                    let mut buf = BytesMut::with_capacity(nonce_length + output_length);
+
+                    // Put nonce first
+                    if let Some(n) = self.nonce.take() {
+                        buf.extend(n);
+                    }
 
                     let mut data_len_buf = [0u8; 2];
                     BigEndian::write_u16(&mut data_len_buf, data_length);

src/relay/tcprelay/crypto_io.rs

@@ -9,7 +9,6 @@ use std::{
 };
 
 use byte_string::ByteStr;
-use bytes::Bytes;
 use futures::ready;
 use log::trace;
 use tokio::prelude::*;
@@ -36,18 +35,12 @@ enum ReadStatus {
     Established,
 }
 
-enum WriteStatus {
-    SendIv(Bytes, usize),
-    Established,
-}
-
 pub struct CryptoStream<S> {
     stream: S,
     dec: Option<DecryptedReader>,
-    enc: Option<EncryptedWriter>,
+    enc: EncryptedWriter,
     svr_cfg: Arc<ServerConfig>,
     read_status: ReadStatus,
-    write_status: WriteStatus,
 }
 
 impl<S: Unpin> Unpin for CryptoStream<S> {}
@@ -60,7 +53,7 @@ impl<S> CryptoStream<S> {
             CipherCategory::Aead => method.salt_size(),
         };
 
-        let local_iv = match method.category() {
+        let iv = match method.category() {
             CipherCategory::Stream => {
                 let local_iv = method.gen_init_vec();
                 trace!("Generated Stream cipher IV {:?}", local_iv);
@@ -73,13 +66,18 @@ impl<S> CryptoStream<S> {
             }
         };
 
+        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)),
+        };
+
         CryptoStream {
             stream,
             dec: None,
-            enc: None,
+            enc,
             svr_cfg,
             read_status: ReadStatus::WaitIv(vec![0u8; prev_len], 0usize),
-            write_status: WriteStatus::SendIv(local_iv, 0usize),
         }
     }
 }
@@ -133,41 +131,9 @@ impl<S> CryptoStream<S>
 where
     S: AsyncWrite + Unpin,
 {
-    fn poll_write_handshake(&mut self, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
-        if let WriteStatus::SendIv(ref iv, ref mut pos) = self.write_status {
-            while *pos < iv.len() {
-                let n = ready!(Pin::new(&mut self.stream).poll_write(cx, &iv[*pos..]))?;
-                if n == 0 {
-                    use std::io::ErrorKind;
-                    return Poll::Ready(Err(ErrorKind::UnexpectedEof.into()));
-                }
-                *pos += n;
-            }
-
-            let method = self.svr_cfg.method();
-            let enc = match method.category() {
-                CipherCategory::Stream => {
-                    trace!("Sent Stream cipher IV {:?}", ByteStr::new(&iv));
-                    EncryptedWriter::Stream(StreamEncryptedWriter::new(method, self.svr_cfg.key(), &iv))
-                }
-                CipherCategory::Aead => {
-                    trace!("Sent AEAD cipher salt {:?}", ByteStr::new(&iv));
-                    EncryptedWriter::Aead(AeadEncryptedWriter::new(method, self.svr_cfg.key(), &iv))
-                }
-            };
-
-            self.enc = Some(enc);
-            self.write_status = WriteStatus::Established;
-        }
-
-        Poll::Ready(Ok(()))
-    }
-
     fn priv_poll_write(mut self: Pin<&mut Self>, ctx: &mut Context<'_>, buf: &[u8]) -> Poll<io::Result<usize>> {
-        ready!(self.poll_write_handshake(ctx))?;
-
         let stream = unsafe { &mut *(&mut self.stream as *mut _) };
-        match *self.enc.as_mut().unwrap() {
+        match self.enc {
             EncryptedWriter::Aead(ref mut w) => w.poll_write_encrypted(ctx, stream, buf),
             EncryptedWriter::Stream(ref mut w) => w.poll_write_encrypted(ctx, stream, buf),
         }

src/relay/tcprelay/stream.rs

@@ -9,7 +9,7 @@ use std::{
 };
 
 use crate::crypto::{new_stream, BoxStreamCipher, CipherType, CryptoMode};
-use bytes::{BufMut, BytesMut};
+use bytes::{BufMut, Bytes, BytesMut};
 use futures::ready;
 use tokio::prelude::*;
 
@@ -94,14 +94,16 @@ enum EncryptWriteStep {
 pub struct EncryptedWriter {
     cipher: BoxStreamCipher,
     steps: EncryptWriteStep,
+    iv: Option<Bytes>,
 }
 
 impl EncryptedWriter {
     /// Creates a new EncryptedWriter
-    pub fn new(t: CipherType, key: &[u8], iv: &[u8]) -> EncryptedWriter {
+    pub fn new(t: CipherType, key: &[u8], iv: Bytes) -> EncryptedWriter {
         EncryptedWriter {
-            cipher: new_stream(t, key, iv, CryptoMode::Encrypt),
+            cipher: new_stream(t, key, &iv, CryptoMode::Encrypt),
             steps: EncryptWriteStep::Nothing,
+            iv: Some(iv),
         }
     }
 
@@ -122,7 +124,19 @@ impl EncryptedWriter {
         loop {
             match self.steps {
                 EncryptWriteStep::Nothing => {
-                    let mut buf = BytesMut::with_capacity(self.buffer_size(data));
+                    // Send the first packet with iv
+                    let iv_length = match self.iv {
+                        Some(ref i) => i.len(),
+                        None => 0,
+                    };
+
+                    let mut buf = BytesMut::with_capacity(iv_length + self.buffer_size(data));
+
+                    // Put iv first
+                    if let Some(i) = self.iv.take() {
+                        buf.extend(i);
+                    }
+
                     self.cipher_update(data, &mut buf)?;
 
                     self.steps = EncryptWriteStep::Writing(buf, 0);