Skip to content

Commit

Permalink
handle error instead of panic when creating and sending icmp packet
Browse files Browse the repository at this point in the history
  • Loading branch information
Arian8j2 committed Oct 15, 2024
1 parent 475225a commit a4a4954
Show file tree
Hide file tree
Showing 3 changed files with 17 additions and 8 deletions.
2 changes: 1 addition & 1 deletion forwarder/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -135,7 +135,7 @@ fn peers_thread(
// client <--server socket--- peer <----- remote
server_socket
.send_to(buffer, peer.get_client_addr())
.unwrap();
.ok();
});
poll.poll(peers, on_peer_recv)?;
Ok(())
Expand Down
4 changes: 3 additions & 1 deletion forwarder/src/poll/udp.rs
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,9 @@ impl Poll for UdpPoll {
let peers = peers.read();
for event in &events {
let port = event.token().0 as u16;
let peer = peers.find_peer_with_port(&port).unwrap();
let Some(peer) = peers.find_peer_with_port(&port) else {
continue;
};
peer.set_used();
// each epoll event may result in multiple readiness events
while let Ok(size) = peer.socket.recv(&mut buffer) {
Expand Down
19 changes: 13 additions & 6 deletions forwarder/src/socket/icmp.rs
Original file line number Diff line number Diff line change
Expand Up @@ -50,7 +50,7 @@ impl IcmpSocket {

impl SocketTrait for IcmpSocket {
fn send_to(&self, buffer: &[u8], to: &SocketAddr) -> io::Result<usize> {
let packet = craft_icmp_packet(buffer, &self.udp_socket_addr, to);
let packet = craft_icmp_packet(buffer, &self.udp_socket_addr, to)?;
let mut to_addr = *to;
// in linux `send_to` on icmpv6 socket requires destination port to be zero
to_addr.set_port(0);
Expand All @@ -74,6 +74,7 @@ impl SocketTrait for IcmpSocket {
let payload_len = packet.payload.len();
buffer[..payload_len].copy_from_slice(packet.payload);

// doesn't panic because from_addr is either ipv6 or ipv4
let mut from_addr = from_addr.as_socket().unwrap();
from_addr.set_port(packet.src_port);
return Ok((payload_len, from_addr));
Expand Down Expand Up @@ -111,8 +112,10 @@ impl NonBlockingSocketTrait for NonBlockingIcmpSocket {
}

fn send(&self, buffer: &[u8]) -> io::Result<usize> {
let dst_addr = self.connected_addr.unwrap();
let packet = craft_icmp_packet(buffer, &self.icmp_socket.udp_socket_addr, &dst_addr);
let dst_addr = self
.connected_addr
.ok_or_else(|| Into::<io::Error>::into(io::ErrorKind::NotConnected))?;
let packet = craft_icmp_packet(buffer, &self.icmp_socket.udp_socket_addr, &dst_addr)?;
self.icmp_socket.socket.send(&packet)
}

Expand All @@ -130,7 +133,11 @@ impl NonBlockingSocketTrait for NonBlockingIcmpSocket {
}
}

fn craft_icmp_packet(payload: &[u8], source_addr: &SocketAddr, dst_addr: &SocketAddr) -> Vec<u8> {
fn craft_icmp_packet(
payload: &[u8],
source_addr: &SocketAddr,
dst_addr: &SocketAddr,
) -> io::Result<Vec<u8>> {
let echo_header = IcmpEchoHeader {
id: dst_addr.port(),
seq: source_addr.port(),
Expand All @@ -146,15 +153,15 @@ fn craft_icmp_packet(payload: &[u8], source_addr: &SocketAddr, dst_addr: &Socket
let source_ip = as_socket_addr_v6(*source_addr).ip().octets();
let destination_ip = as_socket_addr_v6(*dst_addr).ip().octets();
Icmpv6Header::with_checksum(icmp_type, source_ip, destination_ip, payload)
.unwrap()
.map_err(|_| Into::<io::Error>::into(io::ErrorKind::InvalidInput))?
.to_bytes()
.to_vec()
};

let mut header_and_payload = Vec::with_capacity(icmp_header.len() + payload.len());
header_and_payload.extend_from_slice(&icmp_header);
header_and_payload.extend_from_slice(payload);
header_and_payload
Ok(header_and_payload)
}

pub struct IcmpPacket<'a> {
Expand Down

0 comments on commit a4a4954

Please sign in to comment.