Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit 26ad0c0

Browse files
zhangxp1998zhangxp1998
committed
Address some comments on PR
1 parent 82e0292 commit 26ad0c0

7 files changed

Lines changed: 46 additions & 37 deletions

File tree

‎src/main.rs‎

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -79,10 +79,15 @@ fn try_main() -> Result<(), failure::Error> {
7979
Ok(())
8080
}
8181

82+
pub struct OpenSockets {
83+
sockets_to_procs: HashMap<LocalSocket, String>,
84+
connections: Vec<Connection>,
85+
}
86+
8287
pub struct OsInputOutput {
8388
pub network_interfaces: Vec<NetworkInterface>,
8489
pub network_frames: Vec<Box<dyn DataLinkReceiver>>,
85-
pub get_open_sockets: fn() -> (HashMap<LocalSocket, String>, std::vec::Vec<Connection>),
90+
pub get_open_sockets: fn() -> OpenSockets,
8691
pub keyboard_events: Box<dyn Iterator<Item = Event> + Send>,
8792
pub dns_client: Option<dns::Client>,
8893
pub on_winch: Box<OnSigWinch>,
@@ -138,7 +143,10 @@ where
138143
while running.load(Ordering::Acquire) {
139144
let render_start_time = Instant::now();
140145
let utilization = { network_utilization.lock().unwrap().clone_and_reset() };
141-
let (connections_to_procs, connections) = get_open_sockets();
146+
let OpenSockets {
147+
sockets_to_procs,
148+
connections,
149+
} = get_open_sockets();
142150
let mut ip_to_host = IpTable::new();
143151
if let Some(dns_client) = dns_client.as_mut() {
144152
ip_to_host = dns_client.cache();
@@ -152,7 +160,7 @@ where
152160
}
153161
{
154162
let mut ui = ui.lock().unwrap();
155-
ui.update_state(connections_to_procs, utilization, ip_to_host);
163+
ui.update_state(sockets_to_procs, utilization, ip_to_host);
156164
if raw_mode {
157165
ui.output_text(&mut write_to_stdout);
158166
} else {

‎src/network/connection.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ impl Protocol {
1717
"TCP" => Some(Protocol::Tcp),
1818
"UDP" => Some(Protocol::Udp),
1919
"ICMP" => Some(Protocol::Icmp),
20-
_ => panic!("{} is not a valid protocol!", string),
20+
_ => None,
2121
}
2222
}
2323
}

‎src/network/utilization.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,14 +2,14 @@ use crate::network::{Connection, Direction, Segment};
22

33
use ::std::collections::HashMap;
44

5-
#[derive(Clone, Debug)]
5+
#[derive(Clone)]
66
pub struct ConnectionInfo {
77
pub interface_name: String,
88
pub total_bytes_downloaded: u128,
99
pub total_bytes_uploaded: u128,
1010
}
1111

12-
#[derive(Clone, Debug)]
12+
#[derive(Clone)]
1313
pub struct Utilization {
1414
pub connections: HashMap<Connection, ConnectionInfo>,
1515
}

‎src/os/linux.rs‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,10 @@ use ::std::collections::HashMap;
22

33
use ::procfs::process::FDTarget;
44

5-
use crate::network::{Connection, LocalSocket, Protocol};
5+
use crate::network::{Connection, Protocol};
6+
use crate::OpenSockets;
67

7-
pub(crate) fn get_open_sockets() -> (HashMap<LocalSocket, String>, std::vec::Vec<Connection>) {
8+
pub(crate) fn get_open_sockets() -> OpenSockets {
89
let mut open_sockets = HashMap::new();
910
let mut connections = std::vec::Vec::new();
1011
let all_procs = procfs::process::all_processes().unwrap();
@@ -46,5 +47,8 @@ pub(crate) fn get_open_sockets() -> (HashMap<LocalSocket, String>, std::vec::Vec
4647
connections.push(connection);
4748
};
4849
}
49-
(open_sockets, connections)
50+
OpenSockets {
51+
sockets_to_procs: open_sockets,
52+
connections,
53+
}
5054
}

‎src/os/lsof_utils.rs‎

Lines changed: 11 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -21,16 +21,6 @@ lazy_static! {
2121
static ref LISTEN_REGEX: Regex = Regex::new(r"(.*):(.*)").unwrap();
2222
}
2323

24-
fn parse_ip_addr(ip: &str) -> IpAddr {
25-
if let Ok(v4addr) = ip.parse() {
26-
IpAddr::V4(v4addr)
27-
} else if let Ok(v6addr) = ip.parse() {
28-
IpAddr::V6(v6addr)
29-
} else {
30-
panic!("{} is not a valid IP address", ip)
31-
}
32-
}
33-
3424
fn get_null_addr(ip_type: &str) -> &str {
3525
if ip_type.contains('4') {
3626
"0.0.0.0"
@@ -39,16 +29,12 @@ fn get_null_addr(ip_type: &str) -> &str {
3929
}
4030
}
4131

42-
#[allow(clippy::needless_return)]
4332
impl RawConnection {
4433
pub fn new(raw_line: &str) -> Option<RawConnection> {
34+
// Example row
35+
// com.apple 664 user 198u IPv4 0xeb179a6650592b8d 0t0 TCP 192.168.1.187:58535->1.2.3.4:443 (ESTABLISHED)
4536
let columns: Vec<&str> = raw_line.split_ascii_whitespace().collect();
4637
if columns.len() < 9 {
47-
println!(
48-
"lsof's output string has {} columns, different than expected: {:#?}",
49-
columns.len(),
50-
columns
51-
);
5238
return None;
5339
}
5440
let process_name = columns[0].replace("\\x20", " ");
@@ -68,6 +54,10 @@ impl RawConnection {
6854
// let connection_state = columns[9];
6955
// If this socket is in a "connected" state
7056
if let Some(caps) = CONNECTION_REGEX.captures(connection_str) {
57+
// Example
58+
// 192.168.1.187:64230->0.1.2.3:5228
59+
// *:*
60+
// *:4567
7161
let local_ip = String::from(caps.get(1).unwrap().as_str());
7262
let local_port = String::from(caps.get(2).unwrap().as_str());
7363
let remote_ip = String::from(caps.get(3).unwrap().as_str());
@@ -110,23 +100,23 @@ impl RawConnection {
110100
}
111101

112102
pub fn get_protocol(&self) -> Protocol {
113-
return Protocol::from_str(&self.protocol).unwrap();
103+
Protocol::from_str(&self.protocol).unwrap()
114104
}
115105

116106
pub fn get_remote_ip(&self) -> IpAddr {
117-
return parse_ip_addr(&self.remote_ip);
107+
self.remote_ip.parse().unwrap()
118108
}
119109

120110
pub fn get_remote_port(&self) -> u16 {
121-
return self.remote_port.parse::<u16>().unwrap();
111+
self.remote_port.parse::<u16>().unwrap()
122112
}
123113

124114
pub fn get_local_ip(&self) -> IpAddr {
125-
return parse_ip_addr(&self.local_ip);
115+
self.local_ip.parse().unwrap()
126116
}
127117

128118
pub fn get_local_port(&self) -> u16 {
129-
return self.local_port.parse::<u16>().unwrap();
119+
self.local_port.parse::<u16>().unwrap()
130120
}
131121
}
132122

‎src/os/macos.rs‎

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
use ::std::collections::HashMap;
22

3-
use crate::network::{Connection, LocalSocket};
3+
use crate::network::Connection;
4+
use crate::OpenSockets;
45

56
use super::lsof_utils;
67
use std::net::SocketAddr;
@@ -14,8 +15,7 @@ struct RawConnection {
1415
process_name: String,
1516
}
1617

17-
#[allow(clippy::needless_return)]
18-
pub(crate) fn get_open_sockets() -> (HashMap<LocalSocket, String>, std::vec::Vec<Connection>) {
18+
pub(crate) fn get_open_sockets() -> OpenSockets {
1919
let mut open_sockets = HashMap::new();
2020
let mut connections_vec = std::vec::Vec::new();
2121

@@ -35,5 +35,8 @@ pub(crate) fn get_open_sockets() -> (HashMap<LocalSocket, String>, std::vec::Vec
3535
connections_vec.push(connection);
3636
}
3737

38-
return (open_sockets, connections_vec);
38+
OpenSockets {
39+
sockets_to_procs: open_sockets,
40+
connections: connections_vec,
41+
}
3942
}

‎src/tests/fakes/fake_input.rs‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -13,9 +13,10 @@ use ::termion::event::Event;
1313
use crate::{
1414
network::{
1515
dns::{self, Lookup},
16-
Connection, LocalSocket, Protocol,
16+
Connection, Protocol,
1717
},
1818
os::OnSigWinch,
19+
OpenSockets,
1920
};
2021

2122
pub struct KeyboardEvents {
@@ -85,7 +86,7 @@ impl DataLinkReceiver for NetworkFrames {
8586
}
8687
}
8788

88-
pub fn get_open_sockets() -> (HashMap<LocalSocket, String>, std::vec::Vec<Connection>) {
89+
pub fn get_open_sockets() -> OpenSockets {
8990
let mut open_sockets = HashMap::new();
9091
let local_ip = IpAddr::V4(Ipv4Addr::new(10, 0, 0, 2));
9192
open_sockets.insert(
@@ -145,7 +146,10 @@ pub fn get_open_sockets() -> (HashMap<LocalSocket, String>, std::vec::Vec<Connec
145146
connections.push(connection);
146147
}
147148

148-
(local_socket_to_procs, connections)
149+
OpenSockets {
150+
sockets_to_procs: local_socket_to_procs,
151+
connections,
152+
}
149153
}
150154

151155
pub fn get_interfaces() -> Vec<NetworkInterface> {

0 commit comments

Comments
 (0)