Improve error handling

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<String>,
@@ -17,25 +20,28 @@ pub fn get_default_gateway() -> Gateway {
         ip: None,
         mac: None,
     };
-    let default_gateway_ip: Option<String> = match get_default_gateway_ip(){
+    let ip: Option<String> = match get_default_gateway_ip() {
-        Ok(gateway_ip) => Some(gateway_ip),
+        Ok(ip) => Some(ip),
         Err(_) => None,
     };
-    gateway.ip = default_gateway_ip.clone();
+    gateway.ip = ip.clone();
-    if let Some(gateway_ip) = default_gateway_ip.clone(){
+    if let Some(gateway_ip) = ip {
-        let default_gateway_mac: Option<String> = match get_default_gateway_mac(gateway_ip.to_string()) {
+        let mac: Option<String> = match get_default_gateway_mac(gateway_ip) {
-            Ok(gateway_mac) => Some(gateway_mac),
+            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<String, String> {
-    send_udp_packet();
+    match send_udp_packet() {
-    let timeout = Duration::from_millis(3000);
+        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;
 }
@@ -50,23 +56,33 @@ pub fn get_default_gateway_mac(gateway_ip: String) -> Result<String,String>{
                 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();
+    match socket.set_ttl(1) {
-    socket.send_to(&buf, dest).unwrap();
+        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<String, String> {
-    let default_idx = interface::get_default_interface_index().unwrap();
+    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,7 +98,10 @@ fn receive_packets(rx: &mut Box<dyn pnet::datalink::DataLinkReceiver>, 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) {
@@ -104,7 +123,10 @@ fn receive_packets(rx: &mut Box<dyn pnet::datalink::DataLinkReceiver>, 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)),
+            }
         }
     }
 }
@@ -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<String> {
-    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<String> {
         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();
+        let arp = match pnet::packet::arp::ArpPacket::new(&buf[pnet::packet::ethernet::MutableEthernetPacket::minimum_packet_size()..]) {
-        if arp.get_sender_hw_addr() != interface.mac.unwrap() {
+            Some(arp) => arp,
+            None => return None,
+        };
+        if arp.get_sender_hw_addr() != iface_mac {
             target_mac_addr = arp.get_sender_hw_addr();
             break;
         }