diff --git a/src/relay/tcprelay/mod.rs b/src/relay/tcprelay/mod.rs index 05b85ea0..8f80959b 100644 --- a/src/relay/tcprelay/mod.rs +++ b/src/relay/tcprelay/mod.rs @@ -21,7 +21,7 @@ use bytes::BytesMut; use futures::future::FusedFuture; use futures::{ready, select, Future}; use log::trace; -use tokio::io::{ReadHalf, WriteHalf}; +use tokio::io::{BufReader, ReadHalf, WriteHalf}; use tokio::net::TcpStream; use tokio::prelude::*; use tokio::time::{self, Delay}; @@ -47,7 +47,7 @@ const BUFFER_SIZE: usize = 8 * 1024; // 8K buffer /// The only feature: Supports timeout pub struct Connection { // Actual connection socket - stream: S, + stream: BufReader, // Timer instance // Read and Write operations shares the same timer timer: Option, @@ -55,15 +55,20 @@ pub struct Connection { timeout: Option, } -impl Connection { +impl Connection +where + S: AsyncRead, +{ pub fn new(stream: S, timeout: Option) -> Connection { Connection { - stream, + stream: BufReader::new(stream), timer: None, timeout, } } +} +impl Connection { fn make_timeout_error() -> io::Error { use std::io::ErrorKind; ErrorKind::TimedOut.into() @@ -100,17 +105,23 @@ where } } -impl Deref for Connection { +impl Deref for Connection +where + S: AsyncRead, +{ type Target = S; fn deref(&self) -> &Self::Target { - &self.stream + self.stream.get_ref() } } -impl DerefMut for Connection { +impl DerefMut for Connection +where + S: AsyncRead, +{ fn deref_mut(&mut self) -> &mut Self::Target { - &mut self.stream + self.stream.get_mut() } } @@ -136,7 +147,7 @@ where impl AsyncWrite for Connection where - S: AsyncWrite + Unpin, + S: AsyncRead + AsyncWrite + Unpin, { fn poll_write(mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8]) -> Poll> { match Pin::new(&mut self.stream).poll_write(cx, buf) { @@ -198,7 +209,7 @@ async fn connect_proxy_server(context: SharedContext, svr_cfg: Arc let mut last_err: Option = None; for addr in &vec_ipaddr { match try_timeout(TcpStream::connect(addr), timeout).await { - Ok(s) => return Ok(STcpStream::new(s, timeout)), + Ok(s) => return Ok(STcpStream::new(BufReader::new(s), timeout)), Err(e) => { error!( "Failed to connect {}:{}, resolved address {}, try another (err: {})", diff --git a/src/relay/tcprelay/server.rs b/src/relay/tcprelay/server.rs index 4b8fe611..743c6af5 100644 --- a/src/relay/tcprelay/server.rs +++ b/src/relay/tcprelay/server.rs @@ -1,26 +1,21 @@ //! Relay for TCP server that running on the server side -use std::{io, net::SocketAddr, sync::Arc}; +use std::io; +use std::net::SocketAddr; +use std::sync::Arc; -use crate::relay::socks5::Address; +use futures::future::{self, Either}; +use futures::stream::{FuturesUnordered, StreamExt}; +use log::{debug, error, info, trace}; +use tokio; +use tokio::net::{TcpListener, TcpStream}; use crate::context::SharedContext; +use crate::relay::socks5::Address; -use futures::{ - future::{self, Either}, - stream::{FuturesUnordered, StreamExt}, -}; -use log::{debug, error, info, trace}; -use tokio::{ - self, - net::{TcpListener, TcpStream}, -}; - -use super::{ - context::{SharedTcpServerContext, TcpServerContext}, - monitor::TcpMonStream, - CryptoStream, STcpStream, -}; +use super::context::{SharedTcpServerContext, TcpServerContext}; +use super::monitor::TcpMonStream; +use super::{CryptoStream, STcpStream}; #[allow(clippy::cognitive_complexity)] async fn handle_client(