From 7aeb5f1e81e39fe1dbf76ffbbb03f6bcf5572c5a Mon Sep 17 00:00:00 2001 From: zonyitoo Date: Sat, 4 Jan 2020 22:55:22 +0800 Subject: [PATCH] Send first packet with IV/Nonce to reduce latency --- Cargo.toml | 2 +- src/relay/tcprelay/aead.rs | 21 ++++++++++--- src/relay/tcprelay/crypto_io.rs | 54 ++++++--------------------------- src/relay/tcprelay/stream.rs | 22 +++++++++++--- 4 files changed, 46 insertions(+), 53 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 21d696ed..6ce929bb 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "shadowsocks-rust" -version = "1.8.0" +version = "1.8.1" authors = ["Y. T. CHUNG "] description = "shadowsocks is a fast tunnel proxy that helps you bypass firewalls." repository = "https://github.com/zonyitoo/shadowsocks-rust" diff --git a/src/relay/tcprelay/aead.rs b/src/relay/tcprelay/aead.rs index 3b9df435..91d01906 100644 --- a/src/relay/tcprelay/aead.rs +++ b/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, } 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); diff --git a/src/relay/tcprelay/crypto_io.rs b/src/relay/tcprelay/crypto_io.rs index 999a3966..1da8afd7 100644 --- a/src/relay/tcprelay/crypto_io.rs +++ b/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 { stream: S, dec: Option, - enc: Option, + enc: EncryptedWriter, svr_cfg: Arc, read_status: ReadStatus, - write_status: WriteStatus, } impl Unpin for CryptoStream {} @@ -60,7 +53,7 @@ impl CryptoStream { 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 CryptoStream { } }; + 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 CryptoStream where S: AsyncWrite + Unpin, { - fn poll_write_handshake(&mut self, cx: &mut Context<'_>) -> Poll> { - 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> { - 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), } diff --git a/src/relay/tcprelay/stream.rs b/src/relay/tcprelay/stream.rs index debef19e..4fafc2c9 100644 --- a/src/relay/tcprelay/stream.rs +++ b/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, } 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);