From c2d9ca53c0ecf6ca660d17205ad0dc53bf0157ca Mon Sep 17 00:00:00 2001 From: "Y. T. Chung" Date: Thu, 27 Oct 2016 01:47:06 +0800 Subject: [PATCH] [#28] Add customized write_all and copy fror EncryptedWriter --- src/relay/tcprelay/local.rs | 2 +- src/relay/tcprelay/mod.rs | 3 +- src/relay/tcprelay/server.rs | 2 +- src/relay/tcprelay/stream.rs | 172 ++++++++++++++++++++++++++++------- 4 files changed, 143 insertions(+), 36 deletions(-) diff --git a/src/relay/tcprelay/local.rs b/src/relay/tcprelay/local.rs index e89be321..a7b9136b 100644 --- a/src/relay/tcprelay/local.rs +++ b/src/relay/tcprelay/local.rs @@ -97,7 +97,7 @@ impl Socks5RelayLocal { .and_then(move |(svr_s, w)| { super::proxy_server_handshake(svr_s, cloned_svr_cfg, addr).and_then(move |(svr_r, svr_w)| { let rhalf = svr_r.and_then(move |svr_r| copy(svr_r, w)); - let whalf = svr_w.and_then(move |svr_w| copy(r, svr_w)); + let whalf = svr_w.and_then(move |svr_w| svr_w.copy_from(r)); tunnel(cloned_addr, whalf, rhalf) }) diff --git a/src/relay/tcprelay/mod.rs b/src/relay/tcprelay/mod.rs index f05dc88d..ff8035c4 100644 --- a/src/relay/tcprelay/mod.rs +++ b/src/relay/tcprelay/mod.rs @@ -95,7 +95,8 @@ pub fn proxy_server_handshake(remote_stream: TcpStream, relay_addr); // Send relay address to remote - relay_addr.write_to(enc_w).and_then(flush) + let local_buf = Vec::new(); + relay_addr.write_to(local_buf).and_then(|buf| enc_w.write_all(buf)).and_then(|(enc_w, _)| flush(enc_w)) }) .boxed(); diff --git a/src/relay/tcprelay/server.rs b/src/relay/tcprelay/server.rs index 33cf0da1..36bc2e12 100644 --- a/src/relay/tcprelay/server.rs +++ b/src/relay/tcprelay/server.rs @@ -125,7 +125,7 @@ impl TcpRelayServer { let (svr_r, svr_w) = svr_s.split(); tunnel(cloned_addr, copy(r, svr_w), - w_fut.and_then(|w| copy(svr_r, w))) + w_fut.and_then(|w| w.copy_from(svr_r))) }) }); diff --git a/src/relay/tcprelay/stream.rs b/src/relay/tcprelay/stream.rs index cf1f87ca..1b6f1c8f 100644 --- a/src/relay/tcprelay/stream.rs +++ b/src/relay/tcprelay/stream.rs @@ -23,12 +23,15 @@ use std::io::{self, Read, BufRead, Write}; use std::cmp; +use std::mem; use crypto::{Cipher, CipherVariant}; +use futures::{Future, Poll}; + /// Reader wrapper that will decrypt data automatically pub struct DecryptedReader - where R: Read + 'static + where R: Read { reader: R, buffer: Vec, @@ -40,7 +43,7 @@ pub struct DecryptedReader const BUFFER_SIZE: usize = 2048; impl DecryptedReader - where R: Read + 'static + where R: Read { pub fn new(r: R, cipher: CipherVariant) -> DecryptedReader { DecryptedReader { @@ -76,7 +79,7 @@ impl DecryptedReader } impl BufRead for DecryptedReader - where R: Read + 'static + where R: Read { fn fill_buf<'b>(&'b mut self) -> io::Result<&'b [u8]> { while self.pos >= self.buffer.len() { @@ -114,7 +117,7 @@ impl BufRead for DecryptedReader } impl Read for DecryptedReader - where R: Read + 'static + where R: Read { fn read(&mut self, buf: &mut [u8]) -> io::Result { let nread = { @@ -128,34 +131,20 @@ impl Read for DecryptedReader /// Writer wrapper that will encrypt data automatically pub struct EncryptedWriter - where W: Write + 'static + where W: Write { writer: W, cipher: CipherVariant, - buffer: Vec, } impl EncryptedWriter - where W: Write + 'static + where W: Write { /// Creates a new EncryptedWriter pub fn new(w: W, cipher: CipherVariant) -> EncryptedWriter { EncryptedWriter { writer: w, cipher: cipher, - buffer: Vec::new(), - } - } - - /// Finalize the cipher, which will writes the final block into buffer - pub fn finalize(&mut self) -> io::Result<()> { - match self.cipher.finalize(&mut self.buffer) { - Ok(..) => { - self.writer - .write_all(&self.buffer[..]) - .and_then(|_| self.writer.flush()) - } - Err(err) => Err(From::from(err)), } } @@ -173,20 +162,45 @@ impl EncryptedWriter pub fn get_mut(&mut self) -> &mut W { &mut self.writer } + + fn cipher_update(&mut self, buf: &[u8], out: &mut Vec) -> io::Result<()> { + self.cipher.update(buf, out).map_err(From::from) + } + + fn cipher_finalize(&mut self, out: &mut Vec) -> io::Result<()> { + self.cipher.finalize(out).map_err(From::from) + } + + /// write_all + pub fn write_all>(self, buf: B) -> EncryptedWriteAll { + EncryptedWriteAll::Writing { + writer: self, + buf: buf, + pos: 0, + enc_buf: Vec::new(), + encrypted: false, + } + } + + /// Copy all data from reader + pub fn copy_from(self, r: R) -> EncryptedCopy { + EncryptedCopy { + reader: r, + writer: self, + read_done: false, + amt: 0, + pos: 0, + cap: 0, + buf: Vec::new(), + } + } } impl Write for EncryptedWriter - where W: Write + 'static + where W: Write { fn write(&mut self, buf: &[u8]) -> io::Result { - match self.cipher.update(buf, &mut self.buffer) { - Ok(..) => { - let len = try!(self.writer.write(&self.buffer[..])); - self.buffer.drain(..len); - Ok(len) - } - Err(err) => Err(From::from(err)), - } + self.writer.write(buf) } fn flush(&mut self) -> io::Result<()> { @@ -194,10 +208,102 @@ impl Write for EncryptedWriter } } -impl Drop for EncryptedWriter - where W: Write + 'static +/// write_all and encrypt data +pub enum EncryptedWriteAll + where W: Write, + B: AsRef<[u8]> { - fn drop(&mut self) { - let _ = self.finalize(); + Writing { + writer: EncryptedWriter, + buf: B, + pos: usize, + enc_buf: Vec, + encrypted: bool, + }, + Empty, +} + +impl Future for EncryptedWriteAll + where W: Write, + B: AsRef<[u8]> +{ + type Item = (EncryptedWriter, B); + type Error = io::Error; + + fn poll(&mut self) -> Poll { + match *self { + EncryptedWriteAll::Empty => panic!("poll after EncryptedWriteAll finished"), + EncryptedWriteAll::Writing { ref mut writer, ref buf, ref mut pos, ref mut enc_buf, ref mut encrypted } => { + if !*encrypted { + *encrypted = true; + try!(writer.cipher_update(buf.as_ref(), enc_buf)); + } + + while *pos < enc_buf.len() { + let n = try_nb!(writer.write(&enc_buf[*pos..])); + *pos += n; + if n == 0 { + let err = io::Error::new(io::ErrorKind::Other, "zero-length write"); + return Err(err); + } + } + } + } + + match mem::replace(self, EncryptedWriteAll::Empty) { + EncryptedWriteAll::Writing { writer, buf, .. } => Ok((writer, buf).into()), + EncryptedWriteAll::Empty => unreachable!(), + } } } + +/// Encrypted copy +pub struct EncryptedCopy { + reader: R, + writer: EncryptedWriter, + read_done: bool, + amt: u64, + pos: usize, + cap: usize, + buf: Vec, +} + +impl Future for EncryptedCopy { + type Item = u64; + type Error = io::Error; + + fn poll(&mut self) -> Poll { + let mut local_buf = [0u8; 2048]; + loop { + // If our buffer is empty, then we need to read some data to + // continue. + if self.pos == self.cap && !self.read_done { + let n = try_nb!(self.reader.read(&mut local_buf[..])); + if n == 0 { + self.read_done = true; + } else { + self.pos = 0; + + self.buf.clear(); + try!(self.writer.cipher_update(&local_buf[..n], &mut self.buf)); + self.cap = self.buf.len(); + } + } + + // If our buffer has some data, let's write it out! + while self.pos < self.cap { + let i = try_nb!(self.writer.write(&self.buf[self.pos..self.cap])); + self.pos += i; + self.amt += i as u64; + } + + // If we've written al the data and we've seen EOF, flush out the + // data and finish the transfer. + // done with the entire transfer. + if self.pos == self.cap && self.read_done { + try_nb!(self.writer.flush()); + return Ok(self.amt.into()); + } + } + } +} \ No newline at end of file