diff --git a/src/gateway.rs b/src/gateway.rs index 7a14981..d9bb9a3 100644 --- a/src/gateway.rs +++ b/src/gateway.rs @@ -5,6 +5,9 @@ use pnet::packet::Packet; use pnet::packet::MutablePacket; use crate::interface; +const TIMEOUT: u64 = 3000; +const RETRY: u32 = 3; + /// Struct of default Gateway information pub struct Gateway { pub ip: Option, @@ -17,56 +20,69 @@ pub fn get_default_gateway() -> Gateway { ip: None, mac: None, }; - let default_gateway_ip: Option = match get_default_gateway_ip(){ - Ok(gateway_ip) => Some(gateway_ip), + let ip: Option = match get_default_gateway_ip() { + Ok(ip) => Some(ip), Err(_) => None, }; - gateway.ip = default_gateway_ip.clone(); - if let Some(gateway_ip) = default_gateway_ip.clone(){ - let default_gateway_mac: Option = match get_default_gateway_mac(gateway_ip.to_string()) { - Ok(gateway_mac) => Some(gateway_mac), + gateway.ip = ip.clone(); + if let Some(gateway_ip) = ip { + let mac: Option = match get_default_gateway_mac(gateway_ip) { + Ok(mac) => Some(mac), Err(_) => None, }; - gateway.mac = default_gateway_mac; + gateway.mac = mac; } return gateway; } /// Get default Gateway IP address -pub fn get_default_gateway_ip() -> Result{ - send_udp_packet(); - let timeout = Duration::from_millis(3000); +pub fn get_default_gateway_ip() -> Result { + match send_udp_packet() { + Ok(_) => (), + Err(e) => return Err(format!("Failed to send UDP packet {}", e)), + } + let timeout = Duration::from_millis(TIMEOUT); let r = receive_icmp_packets(pnet::packet::icmp::IcmpTypes::TimeExceeded, &timeout); return r; } /// Get default Gateway MAC address -pub fn get_default_gateway_mac(gateway_ip: String) -> Result{ - match gateway_ip.parse::(){ +pub fn get_default_gateway_mac(gateway_ip: String) -> Result { + match gateway_ip.parse::() { Ok(ipv4_addr) => { - if let Some(gateway_mac) = get_mac_through_arp(ipv4_addr){ + if let Some(gateway_mac) = get_mac_through_arp(ipv4_addr) { return Ok(gateway_mac); }else{ return Err(String::from("Failed to get gateway mac address")); } }, - Err(_) => return Err(String::from("Failed to get gateway mac address")), + Err(_) => return Err(String::from("Invalid IPv4 address")), } } -fn send_udp_packet(){ +fn send_udp_packet() -> Result<(), String> { let buf = [0u8; 0]; let socket = match UdpSocket::bind("0.0.0.0:0") { Ok(s) => s, - Err(e) => panic!("Error happened {}", e), + Err(e) => return Err(format!("Failed to create UDP socket {}", e)), }; let dest: &str = "1.1.1.1:80"; - socket.set_ttl(1).unwrap(); - socket.send_to(&buf, dest).unwrap(); + match socket.set_ttl(1) { + Ok(_) => (), + Err(e) => return Err(format!("Failed to set TTL {}", e)), + } + match socket.send_to(&buf, dest) { + Ok(_) => (), + Err(e) => return Err(format!("Failed to send data {}", e)), + } + Ok(()) } -fn receive_icmp_packets(icmp_type: pnet::packet::icmp::IcmpType, timeout: &Duration) -> Result{ - let default_idx = interface::get_default_interface_index().unwrap(); +fn receive_icmp_packets(icmp_type: pnet::packet::icmp::IcmpType, timeout: &Duration) -> Result { + let default_idx = match interface::get_default_interface_index() { + Some(idx) => idx, + None => return Err(String::from("Failed to get default interface")), + }; let interfaces = pnet::datalink::interfaces(); let interface = interfaces.into_iter().filter(|interface: &pnet::datalink::NetworkInterface| interface.index == default_idx).next().expect("Failed to get Interface"); let (mut _tx, mut rx) = match pnet::datalink::channel(&interface, Default::default()) { @@ -82,15 +98,18 @@ fn receive_packets(rx: &mut Box, icmp_type loop { match rx.next() { Ok(frame) => { - let frame = pnet::packet::ethernet::EthernetPacket::new(frame).unwrap(); + let frame = match pnet::packet::ethernet::EthernetPacket::new(frame) { + Some(f) => f, + None => return Err(String::from("Failed to read packet")), + }; match frame.get_ethertype() { pnet::packet::ethernet::EtherTypes::Ipv4 => { - if let Some(ip_addr) = ipv4_handler(&frame, icmp_type){ + if let Some(ip_addr) = ipv4_handler(&frame, icmp_type) { return Ok(ip_addr); } }, pnet::packet::ethernet::EtherTypes::Ipv6 => { - if let Some(ip_addr) = ipv6_handler(&frame, icmp_type){ + if let Some(ip_addr) = ipv6_handler(&frame, icmp_type) { return Ok(ip_addr); } }, @@ -104,13 +123,16 @@ fn receive_packets(rx: &mut Box, icmp_type if Instant::now().duration_since(start_time) > *timeout { return Err(String::from("timeout")); }else{ - send_udp_packet(); + match send_udp_packet() { + Ok(_) => (), + Err(e) => return Err(format!("Failed to send UDP packet {}", e)), + } } } } fn ipv4_handler(ethernet: &pnet::packet::ethernet::EthernetPacket, icmp_type: pnet::packet::icmp::IcmpType) -> Option { - if let Some(packet) = pnet::packet::ipv4::Ipv4Packet::new(ethernet.payload()){ + if let Some(packet) = pnet::packet::ipv4::Ipv4Packet::new(ethernet.payload()) { match packet.get_next_level_protocol() { pnet::packet::ip::IpNextHeaderProtocols::Icmp => { return icmp_handler(&packet, icmp_type); @@ -125,7 +147,7 @@ fn ipv4_handler(ethernet: &pnet::packet::ethernet::EthernetPacket, icmp_type: pn } fn ipv6_handler(ethernet: &pnet::packet::ethernet::EthernetPacket, icmp_type: pnet::packet::icmp::IcmpType) -> Option { - if let Some(packet) = pnet::packet::ipv6::Ipv6Packet::new(ethernet.payload()){ + if let Some(packet) = pnet::packet::ipv6::Ipv6Packet::new(ethernet.payload()) { match packet.get_next_header() { pnet::packet::ip::IpNextHeaderProtocols::Icmpv6 => { return icmpv6_handler(&packet, icmp_type); @@ -140,7 +162,7 @@ fn ipv6_handler(ethernet: &pnet::packet::ethernet::EthernetPacket, icmp_type: pn } fn icmp_handler(ip_packet: &pnet::packet::ipv4::Ipv4Packet, icmp_type: pnet::packet::icmp::IcmpType) -> Option { - if let Some(packet) = pnet::packet::icmp::IcmpPacket::new(ip_packet.payload()){ + if let Some(packet) = pnet::packet::icmp::IcmpPacket::new(ip_packet.payload()) { if packet.get_icmp_type() == icmp_type { let ipv4_addr = ip_packet.get_source(); return Some(ipv4_addr.to_string()) @@ -153,7 +175,7 @@ fn icmp_handler(ip_packet: &pnet::packet::ipv4::Ipv4Packet, icmp_type: pnet::pac } fn icmpv6_handler(ip_packet: &pnet::packet::ipv6::Ipv6Packet, icmp_type: pnet::packet::icmp::IcmpType) -> Option { - if let Some(packet) = pnet::packet::icmp::IcmpPacket::new(ip_packet.payload()){ + if let Some(packet) = pnet::packet::icmp::IcmpPacket::new(ip_packet.payload()) { if packet.get_icmp_type() == icmp_type { let ipv6_addr = ip_packet.get_source(); return Some(ipv6_addr.to_string()) @@ -166,7 +188,10 @@ fn icmpv6_handler(ip_packet: &pnet::packet::ipv6::Ipv6Packet, icmp_type: pnet::p } fn get_mac_through_arp(dst_ip: Ipv4Addr) -> Option { - let default_idx = interface::get_default_interface_index().unwrap(); + let default_idx = match interface::get_default_interface_index() { + Some(idx) => idx, + None => return None, + }; let interfaces = pnet::datalink::interfaces(); let interface = interfaces.into_iter().filter(|interface: &pnet::datalink::NetworkInterface| interface.index == default_idx).next().expect("Failed to get Interface"); let src_ip = interface.ips.iter().find(|ip| ip.is_ipv4()) @@ -180,37 +205,56 @@ fn get_mac_through_arp(dst_ip: Ipv4Addr) -> Option { Ok(_) => panic!("Unknown channel type"), Err(e) => panic!("Error happened {}", e), }; - + let iface_mac: pnet::datalink::MacAddr = match interface.mac { + Some(mac) => mac, + None => return None, + }; let mut ethernet_buffer = [0u8; 42]; - let mut ethernet_packet = pnet::packet::ethernet::MutableEthernetPacket::new(&mut ethernet_buffer).unwrap(); - + let mut ethernet_packet = match pnet::packet::ethernet::MutableEthernetPacket::new(&mut ethernet_buffer) { + Some(ethernet_packet) => ethernet_packet, + None => return None, + }; ethernet_packet.set_destination(pnet::datalink::MacAddr::broadcast()); - ethernet_packet.set_source(interface.mac.unwrap()); + ethernet_packet.set_source(iface_mac); ethernet_packet.set_ethertype(pnet::packet::ethernet::EtherTypes::Arp); let mut arp_buffer = [0u8; 28]; - let mut arp_packet = pnet::packet::arp::MutableArpPacket::new(&mut arp_buffer).unwrap(); + let mut arp_packet = match pnet::packet::arp::MutableArpPacket::new(&mut arp_buffer) { + Some(arp_packet) => arp_packet, + None => return None, + }; arp_packet.set_hardware_type(pnet::packet::arp::ArpHardwareTypes::Ethernet); arp_packet.set_protocol_type(pnet::packet::ethernet::EtherTypes::Ipv4); arp_packet.set_hw_addr_len(6); arp_packet.set_proto_addr_len(4); arp_packet.set_operation(pnet::packet::arp::ArpOperations::Request); - arp_packet.set_sender_hw_addr(interface.mac.unwrap()); + arp_packet.set_sender_hw_addr(iface_mac); arp_packet.set_sender_proto_addr(src_ip); arp_packet.set_target_hw_addr(pnet::datalink::MacAddr::zero()); arp_packet.set_target_proto_addr(dst_ip); ethernet_packet.set_payload(arp_packet.packet_mut()); - sender.send_to(ethernet_packet.packet(), None).unwrap().unwrap(); + match sender.send_to(ethernet_packet.packet(), None) { + Some(s) => { + match s { + Ok(_) => (), + Err(_) => return None, + } + }, + None => return None, + } let mut target_mac_addr: pnet::datalink::MacAddr = pnet::datalink::MacAddr::zero(); - for _x in 0..2 { + for _x in 0..(RETRY - 1) { let buf = receiver.next().unwrap(); - let arp = pnet::packet::arp::ArpPacket::new(&buf[pnet::packet::ethernet::MutableEthernetPacket::minimum_packet_size()..]).unwrap(); - if arp.get_sender_hw_addr() != interface.mac.unwrap() { + let arp = match pnet::packet::arp::ArpPacket::new(&buf[pnet::packet::ethernet::MutableEthernetPacket::minimum_packet_size()..]) { + Some(arp) => arp, + None => return None, + }; + if arp.get_sender_hw_addr() != iface_mac { target_mac_addr = arp.get_sender_hw_addr(); break; } diff --git a/src/interface.rs b/src/interface.rs index cd48ed6..fe621ce 100644 --- a/src/interface.rs +++ b/src/interface.rs @@ -14,7 +14,7 @@ pub struct Interface { } /// Get default Interface -pub fn get_default_interface()-> Option { +pub fn get_default_interface() -> Option { let local_ip = get_local_ipaddr(); let all_interfaces = datalink::interfaces(); if let Some(local_ip) = local_ip { @@ -61,8 +61,8 @@ pub fn get_default_interface_index() -> Option { let local_ip = get_local_ipaddr(); let all_interfaces = datalink::interfaces(); if let Some(local_ip) = local_ip { - for iface in all_interfaces{ - for ip in iface.ips{ + for iface in all_interfaces { + for ip in iface.ips { if local_ip == ip.ip().to_string() { return Some(iface.index) } @@ -79,8 +79,8 @@ pub fn get_default_interface_name() -> Option { let local_ip = get_local_ipaddr(); let all_interfaces = datalink::interfaces(); if let Some(local_ip) = local_ip { - for iface in all_interfaces{ - for ip in iface.ips{ + for iface in all_interfaces { + for ip in iface.ips { if local_ip == ip.ip().to_string() { return Some(iface.name) }