StephenDev0/idevice
1
0
Fork 0
mirror of https://github.com/jkcoxson/idevice.git synced 2026-03-02 14:36:16 +01:00
Code Issues Packages Projects Releases Wiki Activity

Read IP packets outside of tokio::select to avoid cancel

Browse Source
This commit is contained in:
Jackson Coxson
2025-08-20 12:41:09 -06:00
parent f388aaaf2d
commit 6ce6777479
3 changed files with 175 additions and 111 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
idevice/src/tcp/adapter.rs
View File

@@ -112,7 +112,7 @@ impl ConnectionState {
#[derive(Debug)]
pub struct Adapter {
/// The underlying transport connection
peer: Box<dyn ReadWrite>,
pub(crate) peer: Box<dyn ReadWrite>,
/// The local IP address
host_ip: IpAddr,
/// The remote peer's IP address
@@ -538,9 +538,15 @@ impl Adapter {
}
pub(crate) async fn process_tcp_packet(&mut self) -> Result<(), std::io::Error> {
loop {
let ip_packet = self.read_ip_packet().await?;
let res = TcpPacket::parse(&ip_packet)?;
self.process_tcp_packet_from_payload(&ip_packet).await
}
pub(crate) async fn process_tcp_packet_from_payload(
&mut self,
payload: &[u8],
) -> Result<(), std::io::Error> {
let res = TcpPacket::parse(payload)?;
let mut ack_me = None;
if let Some(state) = self.states.get(&res.destination_port) {
@@ -554,14 +560,14 @@ impl Adapter {
debug!("responding to keep-alive probe");
let port = res.destination_port;
self.ack(port).await?;
break;
return Ok(());
}
}
if let Some(state) = self.states.get_mut(&res.destination_port) {
if state.peer_seq > res.sequence_number {
// ignore retransmission
continue;
return Ok(());
}
state.peer_seq = res.sequence_number + res.payload.len() as u32;
@@ -594,8 +600,6 @@ impl Adapter {
if let Some(a) = ack_me {
self.ack(a).await?;
}
break;
}
Ok(())
}

58
idevice/src/tcp/handle.rs
View File

@@ -8,13 +8,16 @@ use std::{collections::HashMap, path::PathBuf, sync::Mutex, task::Poll};
use crossfire::{AsyncRx, MTx, Tx, mpsc, spsc, stream::AsyncStream};
use futures::{StreamExt, stream::FuturesUnordered};
use log::trace;
use log::{debug, trace};
use tokio::{
io::{AsyncRead, AsyncWrite},
io::{AsyncRead, AsyncReadExt, AsyncWrite},
sync::oneshot,
};
use crate::tcp::adapter::ConnectionStatus;
use crate::tcp::{
adapter::ConnectionStatus,
packets::{IpParseError, Ipv6Packet},
};
pub type ConnectToPortRes =
oneshot::Sender<Result<(u16, AsyncRx<Result<Vec<u8>, std::io::Error>>), std::io::Error>>;
@@ -50,6 +53,9 @@ impl AdapterHandle {
tokio::spawn(async move {
let mut handles: HashMap<u16, Tx<Result<Vec<u8>, std::io::Error>>> = HashMap::new();
let mut tick = tokio::time::interval(std::time::Duration::from_millis(1));
let mut read_buf = [0u8; 4096];
let mut bytes_in_buf = 0;
loop {
tokio::select! {
// check for messages for us
@@ -96,16 +102,25 @@ impl AdapterHandle {
}
}
r = adapter.process_tcp_packet() => {
if let Err(e) = r {
// propagate error to all streams; close them
for (hp, tx) in handles.drain() {
let _ = tx.send(Err(e.kind().into())); // or clone/convert
let _ = adapter.close(hp).await;
result = adapter.peer.read(&mut read_buf[bytes_in_buf..]) => {
match result {
Ok(0) => {
debug!("Underlying stream closed (EOF)");
break; // Exit the main actor loop
}
break;
Ok(s) => {
bytes_in_buf += s;
loop {
match Ipv6Packet::parse(&read_buf[..bytes_in_buf]) {
IpParseError::Ok { packet, bytes_consumed } => {
// We got a full packet! Process it.
if let Err(e) = adapter.process_tcp_packet_from_payload(&packet.payload).await {
debug!("CRITICAL: Failed to process IP packet: {e:?}");
}
// And remove it from the buffer by shifting the remaining bytes
read_buf.copy_within(bytes_consumed..bytes_in_buf, 0);
bytes_in_buf -= bytes_consumed;
// Push any newly available bytes to per-conn channels
let mut dead = Vec::new();
for (&hp, tx) in &handles {
@@ -144,6 +159,29 @@ impl AdapterHandle {
let _ = adapter.close(hp).await;
}
}
IpParseError::NotEnough => {
// Buffer doesn't have a full packet, wait for the next read
break;
}
IpParseError::Invalid => {
// Corrupted data, close the connection
// ... (error handling) ...
return;
}
}
}
}
Err(e) => {
debug!("Failed to read: {e:?}, closing stack");
for (hp, tx) in handles.drain() {
let _ = tx.send(Err(e.kind().into())); // or clone/convert
let _ = adapter.close(hp).await;
}
break;
}
}
}
_ = tick.tick() => {
let _ = adapter.write_buffer_flush().await;

34
idevice/src/tcp/packets.rs
View File

@@ -6,6 +6,7 @@ use std::{
sync::Arc,
};
use log::debug;
use tokio::{
io::{AsyncRead, AsyncReadExt},
sync::Mutex,
@@ -109,6 +110,7 @@ impl Ipv4Packet {
let ihl = (version_ihl & 0x0F) * 4;
if version != 4 || ihl < 20 {
debug!("Got an invalid IPv4 header");
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidData,
"Invalid IPv4 header",
@@ -220,21 +222,34 @@ pub struct Ipv6Packet {
pub payload: Vec<u8>,
}
#[derive(Debug, Clone)]
pub(crate) enum IpParseError<T> {
Ok { packet: T, bytes_consumed: usize },
NotEnough,
Invalid,
}
impl Ipv6Packet {
pub fn parse(packet: &[u8]) -> Option<Self> {
pub(crate) fn parse(packet: &[u8]) -> IpParseError<Ipv6Packet> {
if packet.len() < 40 {
return None;
return IpParseError::NotEnough;
}
let version = packet[0] >> 4;
if version != 6 {
return None;
return IpParseError::Invalid;
}
let traffic_class = ((packet[0] & 0x0F) << 4) | (packet[1] >> 4);
let flow_label =
((packet[1] as u32 & 0x0F) << 16) | ((packet[2] as u32) << 8) | packet[3] as u32;
let payload_length = u16::from_be_bytes([packet[4], packet[5]]);
let total_packet_len = 40 + payload_length as usize;
if packet.len() < total_packet_len {
return IpParseError::NotEnough;
}
let next_header = packet[6];
let hop_limit = packet[7];
let source = Ipv6Addr::new(
@@ -258,9 +273,10 @@ impl Ipv6Packet {
u16::from_be_bytes([packet[36], packet[37]]),
u16::from_be_bytes([packet[38], packet[39]]),
);
let payload = packet[40..].to_vec();
let payload = packet[40..total_packet_len].to_vec();
Some(Self {
IpParseError::Ok {
packet: Self {
version,
traffic_class,
flow_label,
@@ -270,7 +286,9 @@ impl Ipv6Packet {
source,
destination,
payload,
})
},
bytes_consumed: total_packet_len,
}
}
pub async fn from_reader<R: AsyncRead + Unpin>(
@@ -278,14 +296,17 @@ impl Ipv6Packet {
log: &Option<Arc<Mutex<tokio::fs::File>>>,
) -> Result<Self, std::io::Error> {
let mut log_packet = Vec::new();
let mut header = [0u8; 40]; // IPv6 header size is fixed at 40 bytes
reader.read_exact(&mut header).await?;
if log.is_some() {
log_packet.extend_from_slice(&header);
}
let version = header[0] >> 4;
if version != 6 {
debug!("Got an invalid IPv6 header");
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidData,
"Invalid IPv6 header",
@@ -457,6 +478,7 @@ pub struct TcpPacket {
impl TcpPacket {
pub fn parse(packet: &[u8]) -> Result<Self, std::io::Error> {
if packet.len() < 20 {
debug!("Got an invalid TCP header");
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidData,
"Not enough bytes for TCP header",