From 3ce9466ee2f812d8362b4d90b56836957fb3e909 Mon Sep 17 00:00:00 2001 From: shellrow <81893184+shellrow@users.noreply.github.com> Date: Sat, 13 Nov 2021 16:32:43 +0900 Subject: [PATCH] Performance improvement, Refactoring --- Cargo.toml | 3 +- examples/default_gateway.rs | 14 +- examples/default_interface.rs | 33 +++-- src/gateway.rs | 236 +++++++++------------------------- src/interface.rs | 61 ++++++--- src/lib.rs | 10 +- 6 files changed, 142 insertions(+), 215 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 038979a..3f744ea 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -11,7 +11,8 @@ categories = ["network-programming"] license = "MIT" [dependencies] -pnet = "0.28" +pnet_packet = "0.28" +pnet_datalink = "0.28" [target.'cfg(windows)'.dependencies] winapi = { version = "0.3.9", features = [ "std" ] } diff --git a/examples/default_gateway.rs b/examples/default_gateway.rs index e07d41b..3883198 100644 --- a/examples/default_gateway.rs +++ b/examples/default_gateway.rs @@ -1,8 +1,14 @@ use default_net; fn main(){ - let default_gateway = default_net::get_default_gateway(); - println!("Default Gateway"); - println!("IP {:?}", default_gateway.ip); - println!("MAC {:?}", default_gateway.mac); + match default_net::get_default_gateway() { + Ok(gateway) => { + println!("Default Gateway"); + println!("\tMAC: {}", gateway.mac_addr); + println!("\tIP: {}", gateway.ip_addr); + }, + Err(e) => { + println!("{}", e); + }, + } } diff --git a/examples/default_interface.rs b/examples/default_interface.rs index 3d88b0e..9b6f92e 100644 --- a/examples/default_interface.rs +++ b/examples/default_interface.rs @@ -1,15 +1,28 @@ use default_net; fn main(){ - if let Some(default_interface) = default_net::get_default_interface(){ - println!("Index {}", default_interface.index); - println!("Name {}", default_interface.name); - println!("MAC {:?}", default_interface.mac); - println!("IPv4 {:?}", default_interface.ipv4); - println!("IPv6 {:?}", default_interface.ipv6); - println!("Gateway IP {:?}", default_interface.gateway.ip); - println!("Gateway MAC {:?}", default_interface.gateway.mac); - }else{ - println!("Failed to get default interface info"); + match default_net::get_default_interface() { + Ok(default_interface) => { + println!("Default Interface"); + println!("\tIndex: {}", default_interface.index); + println!("\tName: {}", default_interface.name); + if let Some(mac_addr) = default_interface.mac_addr { + println!("\tMAC: {}", mac_addr); + }else{ + println!("\tMAC: (Failed to get mac address)"); + } + println!("\tIPv4: {:?}", default_interface.ipv4); + println!("\tIPv6: {:?}", default_interface.ipv6); + if let Some(gateway) = default_interface.gateway { + println!("Default Gateway"); + println!("\tMAC: {}", gateway.mac_addr); + println!("\tIP: {}", gateway.ip_addr); + }else { + println!("Default Gateway: (Not found)"); + } + }, + Err(e) => { + println!("{}", e); + }, } } diff --git a/src/gateway.rs b/src/gateway.rs index 9bb5b49..187c1db 100644 --- a/src/gateway.rs +++ b/src/gateway.rs @@ -1,63 +1,46 @@ use std::net::UdpSocket; use std::time::{Duration, Instant}; -use std::net::{IpAddr, Ipv4Addr}; -use pnet::packet::Packet; -use pnet::packet::MutablePacket; -use crate::interface; +use std::net::{IpAddr}; +use pnet_packet::Packet; +use crate::interface::{self, MacAddr}; const TIMEOUT: u64 = 3000; -const RETRY: u32 = 3; /// Struct of default Gateway information +#[derive(Clone, Debug)] pub struct Gateway { - pub ip: Option, - pub mac: Option, + pub mac_addr: MacAddr, + pub ip_addr: IpAddr, } /// Get default Gateway -pub fn get_default_gateway() -> Gateway { - let mut gateway: Gateway = Gateway { - ip: None, - mac: None, +pub fn get_default_gateway() -> 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 config = pnet_datalink::Config { + write_buffer_size: 4096, + read_buffer_size: 4096, + read_timeout: None, + write_timeout: None, + channel_type: pnet_datalink::ChannelType::Layer2, + bpf_fd_attempts: 1000, + linux_fanout: None, + promiscuous: false, }; - let ip: Option = match get_default_gateway_ip() { - Ok(ip) => Some(ip), - Err(_) => None, + let (mut _tx, mut rx) = match pnet_datalink::channel(&interface, config) { + Ok(pnet_datalink::Channel::Ethernet(tx, rx)) => (tx, rx), + Ok(_) => panic!("Unknown channel type"), + Err(e) => panic!("Error happened {}", e), }; - 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 = mac; - } - return gateway; -} - -/// Get default Gateway IP address -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::() { - Ok(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("Invalid IPv4 address")), - } + receive_packets(&mut rx) } fn send_udp_packet() -> Result<(), String> { @@ -78,49 +61,33 @@ fn send_udp_packet() -> Result<(), String> { Ok(()) } -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 config = pnet::datalink::Config { - write_buffer_size: 4096, - read_buffer_size: 4096, - read_timeout: None, - write_timeout: None, - channel_type: pnet::datalink::ChannelType::Layer2, - bpf_fd_attempts: 1000, - linux_fanout: None, - promiscuous: false, - }; - let (mut _tx, mut rx) = match pnet::datalink::channel(&interface, config) { - Ok(pnet::datalink::Channel::Ethernet(tx, rx)) => (tx, rx), - Ok(_) => panic!("Unknown channel type"), - Err(e) => panic!("Error happened {}", e), - }; - receive_packets(&mut rx, icmp_type, timeout) -} - -fn receive_packets(rx: &mut Box, icmp_type: pnet::packet::icmp::IcmpType, timeout: &Duration) -> Result{ +fn receive_packets(rx: &mut Box) -> Result{ + let timeout = Duration::from_millis(TIMEOUT); let start_time = Instant::now(); loop { match rx.next() { Ok(frame) => { - let frame = match pnet::packet::ethernet::EthernetPacket::new(frame) { + 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) { - return Ok(ip_addr); + pnet_packet::ethernet::EtherTypes::Ipv4 => { + if let Some(ip_addr) = ipv4_handler(&frame) { + let gateway = Gateway { + mac_addr: MacAddr::new(frame.get_source().octets()), + ip_addr: ip_addr, + }; + return Ok(gateway); } }, - pnet::packet::ethernet::EtherTypes::Ipv6 => { - if let Some(ip_addr) = ipv6_handler(&frame, icmp_type) { - return Ok(ip_addr); + pnet_packet::ethernet::EtherTypes::Ipv6 => { + if let Some(ip_addr) = ipv6_handler(&frame) { + let gateway = Gateway { + mac_addr: MacAddr::new(frame.get_source().octets()), + ip_addr: ip_addr, + }; + return Ok(gateway); } }, _ => {} @@ -130,8 +97,8 @@ fn receive_packets(rx: &mut Box, icmp_type return Err(format!("An error occurred while reading: {}", e)); } } - if Instant::now().duration_since(start_time) > *timeout { - return Err(String::from("timeout")); + if Instant::now().duration_since(start_time) > timeout { + return Err(String::from("Recieve timeout")); }else{ match send_udp_packet() { Ok(_) => (), @@ -141,11 +108,11 @@ fn receive_packets(rx: &mut Box, icmp_type } } -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()) { +fn ipv4_handler(ethernet: &pnet_packet::ethernet::EthernetPacket) -> Option { + 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); + pnet_packet::ip::IpNextHeaderProtocols::Icmp => { + return icmp_handler(&packet); }, _ => { None @@ -156,11 +123,11 @@ 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()) { +fn ipv6_handler(ethernet: &pnet_packet::ethernet::EthernetPacket) -> Option { + 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); + pnet_packet::ip::IpNextHeaderProtocols::Icmpv6 => { + return icmpv6_handler(&packet); }, _ => { None @@ -171,11 +138,11 @@ 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 packet.get_icmp_type() == icmp_type { +fn icmp_handler(ip_packet: &pnet_packet::ipv4::Ipv4Packet) -> Option { + if let Some(packet) = pnet_packet::icmp::IcmpPacket::new(ip_packet.payload()) { + if packet.get_icmp_type() == pnet_packet::icmp::IcmpTypes::TimeExceeded { let ipv4_addr = ip_packet.get_source(); - return Some(ipv4_addr.to_string()) + return Some(IpAddr::V4(ipv4_addr)) }else{ None } @@ -184,11 +151,11 @@ 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 packet.get_icmp_type() == icmp_type { +fn icmpv6_handler(ip_packet: &pnet_packet::ipv6::Ipv6Packet) -> Option { + if let Some(packet) = pnet_packet::icmpv6::Icmpv6Packet::new(ip_packet.payload()) { + if packet.get_icmpv6_type() == pnet_packet::icmpv6::Icmpv6Types::TimeExceeded { let ipv6_addr = ip_packet.get_source(); - return Some(ipv6_addr.to_string()) + return Some(IpAddr::V6(ipv6_addr)) }else{ None } @@ -196,82 +163,3 @@ fn icmpv6_handler(ip_packet: &pnet::packet::ipv6::Ipv6Packet, icmp_type: pnet::p None } } - -fn get_mac_through_arp(dst_ip: Ipv4Addr) -> Option { - 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()) - .map(|ip| match ip.ip() { - IpAddr::V4(ip) => ip, - _ => unreachable!(), - }) - .unwrap(); - let (mut sender, mut receiver) = match pnet::datalink::channel(&interface, Default::default()) { - Ok(pnet::datalink::Channel::Ethernet(tx, rx)) => (tx, rx), - 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 = 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(iface_mac); - ethernet_packet.set_ethertype(pnet::packet::ethernet::EtherTypes::Arp); - - let mut arp_buffer = [0u8; 28]; - 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(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()); - - 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..(RETRY - 1) { - let buf = receiver.next().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; - } - } - if target_mac_addr == pnet::datalink::MacAddr::zero() { - return None; - }else{ - return Some(target_mac_addr.to_string()); - } -} diff --git a/src/interface.rs b/src/interface.rs index fe621ce..c9f4e5d 100644 --- a/src/interface.rs +++ b/src/interface.rs @@ -1,28 +1,50 @@ use std::net::UdpSocket; use std::net::{IpAddr, Ipv4Addr, Ipv6Addr}; -use pnet::datalink; -use crate::gateway; +use crate::gateway::{self, Gateway}; + +#[derive(Clone, Debug)] +pub struct MacAddr(u8, u8, u8, u8, u8, u8); + +impl MacAddr { + pub fn new(octets: [u8; 6]) -> MacAddr { + MacAddr(octets[0], octets[1], octets[2], octets[3], octets[4], octets[5]) + } + pub fn octets(&self) -> [u8; 6] { + [self.0,self.1,self.2,self.3,self.4,self.5] + } + pub fn address(&self) -> String { + format!("{:02x}:{:02x}:{:02x}:{:02x}:{:02x}:{:02x}", self.0,self.1,self.2,self.3,self.4,self.5) + } +} + +impl std::fmt::Display for MacAddr { + fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + let _ = write!(f,"{:<02x}:{:<02x}:{:<02x}:{:<02x}:{:<02x}:{:<02x}",self.0,self.1,self.2,self.3,self.4,self.5); + Ok(()) + } +} /// Struct of default Network Interface information +#[derive(Clone, Debug)] pub struct Interface { pub index: u32, pub name: String, - pub mac: Option, + pub mac_addr: Option, pub ipv4: Vec, pub ipv6: Vec, - pub gateway: gateway::Gateway, + pub gateway: Option, } /// Get default Interface -pub fn get_default_interface() -> Option { +pub fn get_default_interface() -> Result { let local_ip = get_local_ipaddr(); - let all_interfaces = datalink::interfaces(); + let interfaces = pnet_datalink::interfaces(); if let Some(local_ip) = local_ip { - for iface in all_interfaces{ + for iface in interfaces{ for ip in &iface.ips{ if local_ip == ip.ip().to_string() { - let mac_addr: Option = match iface.mac { - Some(mac_addr) => Some(mac_addr.to_string()), + let mac_addr: Option = match iface.mac { + Some(mac_addr) => Some(MacAddr::new(mac_addr.octets())), None => None, }; let mut ipv4_vec: Vec = vec![]; @@ -37,31 +59,34 @@ pub fn get_default_interface() -> Option { }, } } - let default_gateway = gateway::get_default_gateway(); + let default_gateway: Option = match gateway::get_default_gateway() { + Ok(gateway) => Some(gateway), + Err(_) => None, + }; let interface: Interface = Interface{ index: iface.index, name: iface.name, - mac: mac_addr, + mac_addr: mac_addr, ipv4: ipv4_vec, ipv6: ipv6_vec, gateway: default_gateway, }; - return Some(interface); + return Ok(interface); } } } - return None; + return Err(String::from("")); }else{ - return None; + return Err(String::from("")); } } /// Get default Interface index pub fn get_default_interface_index() -> Option { let local_ip = get_local_ipaddr(); - let all_interfaces = datalink::interfaces(); + let interfaces = pnet_datalink::interfaces(); if let Some(local_ip) = local_ip { - for iface in all_interfaces { + for iface in interfaces { for ip in iface.ips { if local_ip == ip.ip().to_string() { return Some(iface.index) @@ -77,9 +102,9 @@ pub fn get_default_interface_index() -> Option { /// Get default Interface name pub fn get_default_interface_name() -> Option { let local_ip = get_local_ipaddr(); - let all_interfaces = datalink::interfaces(); + let interfaces = pnet_datalink::interfaces(); if let Some(local_ip) = local_ip { - for iface in all_interfaces { + for iface in interfaces { for ip in iface.ips { if local_ip == ip.ip().to_string() { return Some(iface.name) diff --git a/src/lib.rs b/src/lib.rs index bbc4be2..836c4ba 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,13 +1,7 @@ -mod interface; -mod gateway; +pub mod interface; +pub mod gateway; pub use interface::Interface; pub use interface::get_default_interface; -pub use interface::get_default_interface_index; -pub use interface::get_default_interface_name; - pub use gateway::Gateway; pub use gateway::get_default_gateway; -pub use gateway::get_default_gateway_ip; -pub use gateway::get_default_gateway_mac; -