luckyimnotanigger/tun2proxy
1
0
Fork 0
mirror of https://github.com/tun2proxy/tun2proxy.git synced 2025-05-03 20:04:44 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

read_data_from_tcp_stream

Browse source
This commit is contained in:
ssrlive 2023-11-10 13:40:38 +08:00
parent ca94d15597
commit 25257e9a27
1 changed files with 30 additions and 14 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

44
src/tun2proxy.rs
View file

@ -585,10 +585,15 @@ impl<'a> TunToProxy<'a> {
state.dns_over_tcp_expiry = Some(Self::common_udp_life_timeout()); state.dns_over_tcp_expiry = Some(Self::common_udp_life_timeout());
let mut vecbuf = vec![]; let mut vecbuf = vec![];
Self::read_data_from_tcp_stream(&mut state.mio_stream, &mut state.is_tcp_closed, |data| { Self::read_data_from_tcp_stream(
vecbuf.extend_from_slice(data); &mut state.mio_stream,
Ok(()) IP_PACKAGE_MAX_SIZE,
})?; &mut state.is_tcp_closed,
|data| {
vecbuf.extend_from_slice(data);
Ok(())
},
)?;
let data_event = IncomingDataEvent { let data_event = IncomingDataEvent {
direction: IncomingDirection::FromServer, direction: IncomingDirection::FromServer,
@ -1084,13 +1089,18 @@ impl<'a> TunToProxy<'a> {
let mut vecbuf = vec![]; let mut vecbuf = vec![];
use std::io::{Error, ErrorKind}; use std::io::{Error, ErrorKind};
let r = Self::read_data_from_tcp_stream(&mut state.mio_stream, &mut state.is_tcp_closed, |data| { let r = Self::read_data_from_tcp_stream(
vecbuf.extend_from_slice(data); &mut state.mio_stream,
if vecbuf.len() >= IP_PACKAGE_MAX_SIZE { IP_PACKAGE_MAX_SIZE,
return Err(Error::new(ErrorKind::OutOfMemory, "IP_PACKAGE_MAX_SIZE exceeded")); &mut state.is_tcp_closed,
} |data| {
Ok(()) vecbuf.extend_from_slice(data);
}); if vecbuf.len() >= IP_PACKAGE_MAX_SIZE {
return Err(Error::new(ErrorKind::OutOfMemory, "IP_PACKAGE_MAX_SIZE exceeded"));
}
Ok(())
},
);
let len = vecbuf.len(); let len = vecbuf.len();
if let Err(error) = r { if let Err(error) = r {
if error.kind() == ErrorKind::OutOfMemory { if error.kind() == ErrorKind::OutOfMemory {
@ -1207,11 +1217,17 @@ impl<'a> TunToProxy<'a> {
Ok(()) Ok(())
} }
fn read_data_from_tcp_stream<F>(stream: &mut TcpStream, is_closed: &mut bool, mut cb: F) -> std::io::Result<()> fn read_data_from_tcp_stream<F>(
stream: &mut dyn std::io::Read,
buffer_size: usize,
is_closed: &mut bool,
mut callback: F,
) -> std::io::Result<()>
where where
F: FnMut(&mut [u8]) -> std::io::Result<()>, F: FnMut(&mut [u8]) -> std::io::Result<()>,
{ {
let mut tmp = [0_u8; IP_PACKAGE_MAX_SIZE]; assert!(buffer_size > 0);
let mut tmp = vec![0_u8; buffer_size];
loop { loop {
match stream.read(&mut tmp) { match stream.read(&mut tmp) {
Ok(0) => { Ok(0) => {
@ -1220,7 +1236,7 @@ impl<'a> TunToProxy<'a> {
break; break;
} }
Ok(read_result) => { Ok(read_result) => {
cb(&mut tmp[0..read_result])?; callback(&mut tmp[0..read_result])?;
} }
Err(error) => { Err(error) => {
if error.kind() == std::io::ErrorKind::WouldBlock { if error.kind() == std::io::ErrorKind::WouldBlock {