read code

This commit is contained in:
ssrlive 2024-10-21 17:26:35 +08:00
parent 442e752cba
commit be45bc8a22
2 changed files with 40 additions and 51 deletions

View file

@ -240,24 +240,21 @@ where
let mut ip_stack = ipstack::IpStack::new(ipstack_config, device); let mut ip_stack = ipstack::IpStack::new(ipstack_config, device);
#[cfg(feature = "udpgw")] #[cfg(feature = "udpgw")]
let udpgw_client = match args.udpgw_server { let udpgw_client = args.udpgw_server.as_ref().map(|addr| {
None => None, log::info!("UDPGW enabled");
Some(addr) => { let client = Arc::new(UdpGwClient::new(
log::info!("UDPGW enabled"); mtu,
let client = Arc::new(UdpGwClient::new( args.udpgw_max_connections.unwrap_or(100),
mtu, UDPGW_KEEPALIVE_TIME,
args.udpgw_max_connections.unwrap_or(100), args.udp_timeout,
UDPGW_KEEPALIVE_TIME, *addr,
args.udp_timeout, ));
addr, let client_keepalive = client.clone();
)); tokio::spawn(async move {
let client_keepalive = client.clone(); client_keepalive.heartbeat_task().await;
tokio::spawn(async move { });
client_keepalive.heartbeat_task().await; client
}); });
Some(client)
}
};
loop { loop {
let virtual_dns = virtual_dns.clone(); let virtual_dns = virtual_dns.clone();

View file

@ -446,44 +446,36 @@ impl UdpGwClient {
/// # Returns /// # Returns
/// - `Result<UdpGwResponse>`: Returns a result type containing the parsed UDP gateway response, or an error if one occurs. /// - `Result<UdpGwResponse>`: Returns a result type containing the parsed UDP gateway response, or an error if one occurs.
pub(crate) async fn recv_udpgw_packet(udp_mtu: u16, udp_timeout: u64, stream: &mut UdpGwClientStreamReader) -> Result<UdpGwResponse> { pub(crate) async fn recv_udpgw_packet(udp_mtu: u16, udp_timeout: u64, stream: &mut UdpGwClientStreamReader) -> Result<UdpGwResponse> {
let result = match tokio::time::timeout( let result = tokio::time::timeout(
tokio::time::Duration::from_secs(udp_timeout + 2), tokio::time::Duration::from_secs(udp_timeout + 2),
stream.inner.read(&mut stream.recv_buf[..2]), stream.inner.read(&mut stream.recv_buf[..2]),
) )
.await .await
{ .map_err(std::io::Error::from)?;
Ok(ret) => ret, let n = result?;
Err(_e) => { if n == 0 {
return Err(("wait tcp data timeout").into()); return Ok(UdpGwResponse::TcpClose);
}
};
match result {
Ok(0) => Ok(UdpGwResponse::TcpClose),
Ok(n) => {
if n < std::mem::size_of::<PackLenHeader>() {
return Err("received PackLenHeader error".into());
}
let packet_len = u16::from_le_bytes([stream.recv_buf[0], stream.recv_buf[1]]);
if packet_len > udp_mtu {
return Err("packet too long".into());
}
let mut left_len: usize = packet_len as usize;
let mut recv_len = 0;
while left_len > 0 {
if let Ok(len) = stream.inner.read(&mut stream.recv_buf[recv_len..left_len]).await {
if len == 0 {
return Ok(UdpGwResponse::TcpClose);
}
recv_len += len;
left_len -= len;
} else {
return Ok(UdpGwResponse::TcpClose);
}
}
return UdpGwClient::parse_udp_response(udp_mtu, packet_len as usize, stream);
}
Err(_) => Err("tcp read error".into()),
} }
if n < std::mem::size_of::<PackLenHeader>() {
return Err("received PackLenHeader error".into());
}
let packet_len = u16::from_le_bytes([stream.recv_buf[0], stream.recv_buf[1]]);
if packet_len > udp_mtu {
return Err("packet too long".into());
}
let mut left_len: usize = packet_len as usize;
let mut recv_len = 0;
while left_len > 0 {
let Ok(len) = stream.inner.read(&mut stream.recv_buf[recv_len..left_len]).await else {
return Ok(UdpGwResponse::TcpClose);
};
if len == 0 {
return Ok(UdpGwResponse::TcpClose);
}
recv_len += len;
left_len -= len;
}
UdpGwClient::parse_udp_response(udp_mtu, packet_len as usize, stream)
} }
/// Sends a UDP gateway packet. /// Sends a UDP gateway packet.