yeyinlaoying/vnt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[mio] 简化代理

Browse Source
This commit is contained in:
lubeilin
2024-03-06 22:46:17 +08:00
parent f8ee8e242f
commit c0e930a3ef
2 changed files with 54 additions and 63 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vnt/src/ip_proxy/icmp_proxy.rs
+1 -1
View File
@@ -161,7 +161,7 @@ fn recv_handle(
Ok(mut ipv4_packet) => match icmp::IcmpPacket::new(ipv4_packet.payload()) { Ok(mut ipv4_packet) => match icmp::IcmpPacket::new(ipv4_packet.payload()) {
Ok(icmp_packet) => match icmp_packet.header_other() { Ok(icmp_packet) => match icmp_packet.header_other() {
HeaderOther::Identifier(id, seq) => { HeaderOther::Identifier(id, seq) => {
if let Some(dest_ip) = nat_map.lock().remove(&(peer_ip, id, seq)) { if let Some(dest_ip) = nat_map.lock().get(&(peer_ip, id, seq)).cloned() {
ipv4_packet.set_destination_ip(dest_ip); ipv4_packet.set_destination_ip(dest_ip);
ipv4_packet.update_checksum(); ipv4_packet.update_checksum();
vnt/src/ip_proxy/tcp_proxy.rs
+53 -62
View File
@@ -144,54 +144,44 @@ fn tcp_proxy(
let (stream1, stream2, buf1, buf2, state1, state2) = val.as_mut(index); let (stream1, stream2, buf1, buf2, state1, state2) = val.as_mut(index);
if event.is_readable() { if event.is_readable() {
if let Err(_) = readable_handle(stream1, stream2, buf1, state2) { if let Err(_) = readable_handle(stream1, stream2, buf1, state2) {
if buf1.is_empty() { *state1 |= READ_CLOSED;
let _ = stream2.shutdown(Shutdown::Write);
}
and_shutdown_state(state1, Shutdown::Read)
} }
} }
if event.is_writable() { if event.is_writable() {
let read = buf2.len() >= BUF_LEN; let read = buf2.len() >= BUF_LEN;
if let Err(_) = writable_handle(stream1, buf2) { if let Err(_) = writable_handle(stream1, buf2) {
buf2.clear(); *state1 |= WRITE_CLOSED;
let _ = stream2.shutdown(Shutdown::Read);
and_shutdown_state(state1, Shutdown::Write)
} else if read { } else if read {
if readable_handle(stream2, stream1, buf2, state1).is_err() { if readable_handle(stream2, stream1, buf2, state1).is_err() {
if buf2.is_empty() { *state2 |= READ_CLOSED;
let _ = stream1.shutdown(Shutdown::Write);
}
and_shutdown_state(state2, Shutdown::Read)
} }
} }
} }
if event.is_read_closed() { if event.is_read_closed() || event.is_error() {
*state1 |= READ_CLOSED;
}
if event.is_write_closed() || event.is_error() {
*state1 |= WRITE_CLOSED;
}
if is_write_closed(*state1) {
let _ = stream1.shutdown(Shutdown::Write);
let _ = stream2.shutdown(Shutdown::Read);
}
if is_read_closed(*state1) {
let _ = stream1.shutdown(Shutdown::Read);
if buf1.is_empty() { if buf1.is_empty() {
let _ = stream2.shutdown(Shutdown::Write); let _ = stream2.shutdown(Shutdown::Write);
} }
and_shutdown_state(state1, Shutdown::Read)
} }
if event.is_write_closed() { if (is_both_closed(*state1) && buf1.is_empty())
let _ = stream2.shutdown(Shutdown::Read); || (is_both_closed(*state2) && buf2.is_empty())
and_shutdown_state(state1, Shutdown::Write) || (is_write_closed(*state1) && is_write_closed(*state2)
} || (is_read_closed(*state1)
if let Some(state1) = state1 { && is_read_closed(*state2)
if let Some(state2) = state2 { && buf1.is_empty()
if (state1 == &Shutdown::Both && buf2.is_empty()))
&& (state2 == &Shutdown::Write || buf1.is_empty())) {
|| (state2 == &Shutdown::Both && state1 == &Shutdown::Write close(src_index, &mut tcp_map, &mut mapping);
|| buf2.is_empty())
{
close(src_index, &mut tcp_map, &mut mapping);
} else if state2 == state1 {
if state1 == &Shutdown::Both
|| state1 == &Shutdown::Write
|| (buf1.is_empty() && buf2.is_empty())
{
close(src_index, &mut tcp_map, &mut mapping);
}
}
}
} }
} }
} }
@@ -199,20 +189,6 @@ fn tcp_proxy(
} }
} }
fn and_shutdown_state(s1: &mut Option<Shutdown>, s2: Shutdown) {
if let Some(s1) = s1 {
if s1 == &Shutdown::Read && s2 == Shutdown::Read {
*s1 = Shutdown::Read
} else if s1 == &Shutdown::Write && s2 == Shutdown::Write {
*s1 = Shutdown::Write
} else {
*s1 = Shutdown::Both
}
} else {
s1.replace(s2);
}
}
fn accept_handle( fn accept_handle(
registry: &Registry, registry: &Registry,
tcp_listener: &TcpListener, tcp_listener: &TcpListener,
@@ -321,8 +297,8 @@ struct ProxyValue {
dest_fd: usize, dest_fd: usize,
src_buf: BytesMut, src_buf: BytesMut,
dest_buf: BytesMut, dest_buf: BytesMut,
src_state: Option<Shutdown>, src_state: u8,
dest_state: Option<Shutdown>, dest_state: u8,
} }
const BUF_LEN: usize = 65536; const BUF_LEN: usize = 65536;
@@ -336,8 +312,8 @@ impl ProxyValue {
dest_fd, dest_fd,
src_buf: BytesMut::with_capacity(BUF_LEN), src_buf: BytesMut::with_capacity(BUF_LEN),
dest_buf: BytesMut::with_capacity(BUF_LEN), dest_buf: BytesMut::with_capacity(BUF_LEN),
src_state: None, src_state: NORMAL,
dest_state: None, dest_state: NORMAL,
} }
} }
fn as_mut( fn as_mut(
@@ -348,8 +324,8 @@ impl ProxyValue {
&mut TcpStream, &mut TcpStream,
&mut BytesMut, &mut BytesMut,
&mut BytesMut, &mut BytesMut,
&mut Option<Shutdown>, &mut u8,
&mut Option<Shutdown>, &mut u8,
) { ) {
if index == self.src_fd { if index == self.src_fd {
( (
@@ -377,7 +353,7 @@ fn readable_handle(
stream1: &mut TcpStream, stream1: &mut TcpStream,
stream2: &mut TcpStream, stream2: &mut TcpStream,
mid_buf: &mut BytesMut, mid_buf: &mut BytesMut,
state2: &mut Option<Shutdown>, state2: &mut u8,
) -> io::Result<()> { ) -> io::Result<()> {
let mut buf = [0; BUF_LEN]; let mut buf = [0; BUF_LEN];
@@ -398,16 +374,14 @@ fn readable_handle(
match stream2.write(buf) { match stream2.write(buf) {
Ok(end) => { Ok(end) => {
if end == 0 { if end == 0 {
mid_buf.clear(); *state2 |= WRITE_CLOSED;
and_shutdown_state(state2, Shutdown::Write);
return Err(io::Error::from(io::ErrorKind::WriteZero)); return Err(io::Error::from(io::ErrorKind::WriteZero));
} }
buf = &buf[end..]; buf = &buf[end..];
} }
Err(e) => { Err(e) => {
if e.kind() != io::ErrorKind::WouldBlock { if e.kind() != io::ErrorKind::WouldBlock {
mid_buf.clear(); *state2 |= WRITE_CLOSED;
and_shutdown_state(state2, Shutdown::Write);
return Err(e); return Err(e);
} }
break; break;
@@ -454,10 +428,27 @@ fn close(
tcp_map: &mut HashMap<usize, ProxyValue>, tcp_map: &mut HashMap<usize, ProxyValue>,
mapping: &mut HashMap<usize, usize>, mapping: &mut HashMap<usize, usize>,
) { ) {
if let Some(mut val) = tcp_map.remove(&index) { if let Some(val) = tcp_map.remove(&index) {
let _ = val.src_stream.flush(); let _ = val.src_stream.shutdown(Shutdown::Both);
let _ = val.dest_stream.flush(); let _ = val.dest_stream.shutdown(Shutdown::Both);
mapping.remove(&val.src_fd); mapping.remove(&val.src_fd);
mapping.remove(&val.dest_fd); mapping.remove(&val.dest_fd);
} }
} }
const NORMAL: u8 = 0b00;
const READ_CLOSED: u8 = 0b01;
const WRITE_CLOSED: u8 = 0b10;
const BOTH_CLOSED: u8 = 0b11;
fn is_read_closed(state: u8) -> bool {
(state & READ_CLOSED == READ_CLOSED) || is_both_closed(state)
}
fn is_write_closed(state: u8) -> bool {
(state & WRITE_CLOSED == WRITE_CLOSED) || is_both_closed(state)
}
fn is_both_closed(state: u8) -> bool {
state & BOTH_CLOSED == BOTH_CLOSED
}