luckyimnotanigger/tun2proxy
1
0
Fork 0
mirror of https://github.com/tun2proxy/tun2proxy.git synced 2025-04-21 14:29:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

read_data_from_tcp_stream for callback

Browse source
This commit is contained in:
ssrlive 2023-09-28 11:48:54 +08:00
parent 03aa70f3c2
commit 05cb35fabb
1 changed files with 17 additions and 7 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

24
src/tun2proxy.rs
View file

@ -567,7 +567,11 @@ impl<'a> TunToProxy<'a> {
assert!(state.dns_over_tcp_expiry.is_some()); assert!(state.dns_over_tcp_expiry.is_some());
state.dns_over_tcp_expiry = Some(Self::common_udp_life_timeout()); state.dns_over_tcp_expiry = Some(Self::common_udp_life_timeout());
let vecbuf = Self::read_data_from_tcp_stream(&mut state.mio_stream)?; let mut vecbuf = vec![];
Self::read_data_from_tcp_stream(&mut state.mio_stream, |data| {
vecbuf.extend_from_slice(data);
Ok(())
})?;
let data_event = IncomingDataEvent { let data_event = IncomingDataEvent {
direction: IncomingDirection::FromServer, direction: IncomingDirection::FromServer,
@ -1029,7 +1033,11 @@ impl<'a> TunToProxy<'a> {
let state = self.connection_map.get_mut(&conn_info).ok_or(e)?; let state = self.connection_map.get_mut(&conn_info).ok_or(e)?;
// TODO: Move this reading process to its own function. // TODO: Move this reading process to its own function.
let vecbuf = Self::read_data_from_tcp_stream(&mut state.mio_stream)?; let mut vecbuf = vec![];
Self::read_data_from_tcp_stream(&mut state.mio_stream, |data| {
vecbuf.extend_from_slice(data);
Ok(())
})?;
let data_event = IncomingDataEvent { let data_event = IncomingDataEvent {
direction: IncomingDirection::FromServer, direction: IncomingDirection::FromServer,
@ -1095,17 +1103,19 @@ impl<'a> TunToProxy<'a> {
Ok(()) Ok(())
} }
fn read_data_from_tcp_stream(stream: &mut TcpStream) -> Result<Vec<u8>> { fn read_data_from_tcp_stream<F>(stream: &mut TcpStream, mut callback: F) -> Result<()>
let mut vecbuf = Vec::<u8>::new(); where
F: FnMut(&mut [u8]) -> Result<()>,
{
let mut tmp: [u8; 4096] = [0_u8; 4096];
loop { loop {
let mut tmp: [u8; 4096] = [0_u8; 4096];
match stream.read(&mut tmp) { match stream.read(&mut tmp) {
Ok(0) => { Ok(0) => {
// The tcp connection closed // The tcp connection closed
break; break;
} }
Ok(read_result) => { Ok(read_result) => {
vecbuf.extend_from_slice(&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 {
@ -1120,7 +1130,7 @@ impl<'a> TunToProxy<'a> {
} }
}; };
} }
Ok(vecbuf) Ok(())
} }
#[cfg(any(target_os = "linux", target_os = "macos"))] #[cfg(any(target_os = "linux", target_os = "macos"))]