add recv/send bytes support (net i/o)

README.md

@@ -466,8 +466,10 @@ The first `[[columns]]` is shown at left side, and the last is shown at right si
 | Priority             | pri                   | Priority                                      | o     | o     | o       | o       |
 | Processor            | psr                   | Currently assigned processor                  | o     |       |         | o       |
 | ReadBytes            | -not supported-       | Read bytes from storage                       | o     | o     | o       | o       |
+| RecvBytes            | -not supported-       | Received bytes per second (Windows 11)        |       |       | o       |         |
 | RtPriority           | rtprio                | Kernel base priority (Windows) / real-time priority (Linux) | o     |       | o       |         |
 | SecContext           | label                 | Security context                              | o     |       |         |         |
+| SendBytes            | -not supported-       | Sent bytes per second (Windows 11)            |       |       | o       |         |
 | Separator            | -not supported-       | Show `\|` for column separation               | o     | o     | o       | o       |
 | Session              | sid                   | Session ID                                    | o     | o     | o       | o       |
 | ShdPnd               | pending               | Pending signal mask for process               | o     |       |         | o       |
@@ -511,7 +513,7 @@ The first `[[columns]]` is shown at left side, and the last is shown at right si
 | VmTotal              | -not supported-       | Total virtual memory size                     | *[^*] | o     | *[^*]   | *[^*]   |
 | VoluntaryContextSw   | -not supported-       | Voluntary context switch count                | o     |       |         | o       |
 | Wchan                | wchan                 | Process sleeping kernel function              | o     |       |         | o       |
-| WorkDir              | -not supported-       | Current working directory                     | o     |       |         |         |
+| WorkDir              | -not supported-       | Current working directory                     | o     |       | o       |         |
 | WriteByte            | -not supported-       | Write bytes to storage                        | o     | o     | o       | o       |
 
 [^*]: Alias for VmRss on these platforms

src/columns/os_windows.rs

@@ -14,7 +14,9 @@ pub mod ppid;
 pub mod priority;
 pub mod rt_priority;
 pub mod read_bytes;
+pub mod recv_bytes;
 pub mod separator;
+pub mod send_bytes;
 pub mod session;
 pub mod slot;
 pub mod start_time;
@@ -52,7 +54,9 @@ pub use self::ppid::Ppid;
 pub use self::priority::Priority;
 pub use self::rt_priority::RtPriority;
 pub use self::read_bytes::ReadBytes;
+pub use self::recv_bytes::RecvBytes;
 pub use self::separator::Separator;
+pub use self::send_bytes::SendBytes;
 pub use self::session::Session;
 pub use self::slot::Slot;
 pub use self::start_time::StartTime;
@@ -101,7 +105,9 @@ pub enum ConfigColumnKind {
     Ppid,
     Priority,
     ReadBytes,
+    RecvBytes,
     RtPriority,
+    SendBytes,
     Separator,
     Session,
     Slot,
@@ -154,7 +160,9 @@ pub fn gen_column(
         ConfigColumnKind::Ppid => Box::new(Ppid::new(header)),
         ConfigColumnKind::Priority => Box::new(Priority::new(header)),
         ConfigColumnKind::ReadBytes => Box::new(ReadBytes::new(header)),
+        ConfigColumnKind::RecvBytes => Box::new(RecvBytes::new(header)),
         ConfigColumnKind::RtPriority => Box::new(RtPriority::new(header)),
+        ConfigColumnKind::SendBytes => Box::new(SendBytes::new(header)),
         ConfigColumnKind::Separator => Box::new(Separator::new(separator)),
         ConfigColumnKind::Session => Box::new(Session::new(header)),
         ConfigColumnKind::Slot => Box::new(Slot::new()),
@@ -228,10 +236,18 @@ pub static KIND_LIST: Lazy<BTreeMap<ConfigColumnKind, (&'static str, &'static st
                 ConfigColumnKind::ReadBytes,
                 ("ReadBytes", "Read bytes from storage"),
             ),
+            (
+                ConfigColumnKind::RecvBytes,
+                ("RecvBytes", "Received bytes per second (Windows 11)"),
+            ),
             (
                 ConfigColumnKind::RtPriority,
                 ("RtPriority", "Kernel base priority (0-31)"),
             ),
+            (
+                ConfigColumnKind::SendBytes,
+                ("SendBytes", "Sent bytes per second (Windows 11)"),
+            ),
             (
                 ConfigColumnKind::Separator,
                 ("Separator", "Show | for column separation"),
@@ -407,6 +423,18 @@ numeric_search = false
 nonnumeric_search = false
 align = "Right"
 [[columns]]
+kind = "RecvBytes"
+style = "ByUnit"
+numeric_search = false
+nonnumeric_search = false
+align = "Right"
+[[columns]]
+kind = "SendBytes"
+style = "ByUnit"
+numeric_search = false
+nonnumeric_search = false
+align = "Right"
+[[columns]]
 kind = "Slot"
 style = "ByUnit"
 numeric_search = false

src/columns/recv_bytes.rs

@@ -0,0 +1,55 @@
+use crate::process::ProcessInfo;
+use crate::util::bytify;
+use crate::{Column, column_default};
+use std::cmp;
+use std::collections::HashMap;
+
+/// Bytes received per second.
+///
+/// Windows only: the counters come from `ProcessNetworkIoInformation`
+/// (`PROCESS_NETWORK_COUNTERS::BytesIn`), so the column is declared by
+/// `os_windows.rs` alone. It reads a rate, which needs the two samples
+/// `collect_proc` takes; a process that could not be opened, or a Windows
+/// older than 11, reports zero.
+pub struct RecvBytes {
+    header: String,
+    unit: String,
+    fmt_contents: HashMap<i32, String>,
+    raw_contents: HashMap<i32, u64>,
+    width: usize,
+}
+
+impl RecvBytes {
+    pub fn new(header: Option<String>) -> Self {
+        let header = header.unwrap_or_else(|| String::from("Recv"));
+        let unit = String::from("[B/s]");
+        Self {
+            fmt_contents: HashMap::new(),
+            raw_contents: HashMap::new(),
+            width: 0,
+            header,
+            unit,
+        }
+    }
+}
+
+impl Column for RecvBytes {
+    fn add(&mut self, proc: &ProcessInfo) {
+        let interval_ms =
+            (proc.interval.as_secs() * 1000 + u64::from(proc.interval.subsec_millis())).max(1);
+        let io = proc
+            .net_info
+            .curr_recv
+            .saturating_sub(proc.net_info.prev_recv)
+            * 1000
+            / interval_ms;
+
+        let raw_content = io;
+        let fmt_content = bytify(raw_content);
+
+        self.fmt_contents.insert(proc.pid, fmt_content);
+        self.raw_contents.insert(proc.pid, raw_content);
+    }
+
+    column_default!(u64, true);
+}

src/columns/send_bytes.rs

@@ -0,0 +1,55 @@
+use crate::process::ProcessInfo;
+use crate::util::bytify;
+use crate::{Column, column_default};
+use std::cmp;
+use std::collections::HashMap;
+
+/// Bytes sent per second.
+///
+/// Windows only: the counters come from `ProcessNetworkIoInformation`
+/// (`PROCESS_NETWORK_COUNTERS::BytesOut`), so the column is declared by
+/// `os_windows.rs` alone. It reads a rate, which needs the two samples
+/// `collect_proc` takes; a process that could not be opened, or a Windows
+/// older than 11, reports zero.
+pub struct SendBytes {
+    header: String,
+    unit: String,
+    fmt_contents: HashMap<i32, String>,
+    raw_contents: HashMap<i32, u64>,
+    width: usize,
+}
+
+impl SendBytes {
+    pub fn new(header: Option<String>) -> Self {
+        let header = header.unwrap_or_else(|| String::from("Send"));
+        let unit = String::from("[B/s]");
+        Self {
+            fmt_contents: HashMap::new(),
+            raw_contents: HashMap::new(),
+            width: 0,
+            header,
+            unit,
+        }
+    }
+}
+
+impl Column for SendBytes {
+    fn add(&mut self, proc: &ProcessInfo) {
+        let interval_ms =
+            (proc.interval.as_secs() * 1000 + u64::from(proc.interval.subsec_millis())).max(1);
+        let io = proc
+            .net_info
+            .curr_send
+            .saturating_sub(proc.net_info.prev_send)
+            * 1000
+            / interval_ms;
+
+        let raw_content = io;
+        let fmt_content = bytify(raw_content);
+
+        self.fmt_contents.insert(proc.pid, fmt_content);
+        self.raw_contents.insert(proc.pid, raw_content);
+    }
+
+    column_default!(u64, true);
+}

src/process/ntapi.rs

@@ -123,6 +123,24 @@ pub struct PROCESS_DISK_COUNTERS {
     pub FlushOperationCount: u64,
 }
 
+/// `PROCESS_NETWORK_COUNTERS` - the output of `ProcessNetworkIoInformation`
+/// (114).
+///
+/// Both counters are cumulative for the lifetime of the process, so a rate
+/// needs two samples taken an interval apart.
+#[repr(C)]
+#[allow(non_snake_case)]
+pub struct PROCESS_NETWORK_COUNTERS {
+    /// Total bytes received by the process.
+    pub BytesIn: u64,
+    /// Total bytes sent by the process.
+    pub BytesOut: u64,
+}
+
+const _: () = assert!(size_of::<PROCESS_NETWORK_COUNTERS>() == 16);
+const _: () = assert!(offset_of!(PROCESS_NETWORK_COUNTERS, BytesIn) == 0);
+const _: () = assert!(offset_of!(PROCESS_NETWORK_COUNTERS, BytesOut) == 8);
+
 /// The leading, version-stable part of `SYSTEM_PROCESS_INFORMATION_EXTENSION`,
 /// as returned by `SystemFullProcessInformation` (148).
 ///
@@ -463,6 +481,9 @@ impl std::hash::Hash for SID_MAX {
 /// the process extension, which includes the user SID.
 pub const SYSTEM_FULL_PROCESS_INFORMATION_CLASS: i32 = 148;
 
+// ProcessInformationClass. Windows 11
+pub const PROCESS_NETWORK_IO_COUNTERS_CLASS: i32 = 114;
+
 /// Guard against a nonsensical `UNICODE_STRING::Length` turning into a huge
 /// allocation. Real command lines are capped well below this.
 const MAX_COMMAND_LINE_BYTES: usize = 64 * 1024;
@@ -1092,6 +1113,61 @@ pub fn process_image_machine(handle: HANDLE) -> Option<u16> {
     Some(info.Machine)
 }
 
+// ---------------------------------------------------------------------------
+// ProcessNetworkIoInformation
+// ---------------------------------------------------------------------------
+
+/// What one `ProcessNetworkIoInformation` query produced: the status, how much
+/// the kernel wrote, and the counters themselves.
+///
+/// The status is handed back rather than swallowed so a caller can tell the
+/// two failures apart: "this Windows does not implement the class" leaves the
+/// columns at zero, while `STATUS_INFO_LENGTH_MISMATCH` would mean this
+/// structure is the wrong size.
+struct NetworkCountersQuery {
+    status: i32,
+    written: u32,
+    counters: PROCESS_NETWORK_COUNTERS,
+}
+
+fn query_network_counters(handle: HANDLE) -> NetworkCountersQuery {
+    // SAFETY: a zeroed `PROCESS_NETWORK_COUNTERS` is a valid output buffer -
+    // both fields are integers.
+    let mut info: PROCESS_NETWORK_COUNTERS = unsafe { std::mem::zeroed() };
+    let mut ret_len: u32 = 0;
+
+    let status = unsafe {
+        NtQueryInformationProcess(
+            handle,
+            PROCESS_NETWORK_IO_COUNTERS_CLASS,
+            ptr::addr_of_mut!(info).cast::<c_void>(),
+            size_of::<PROCESS_NETWORK_COUNTERS>() as u32,
+            ptr::addr_of_mut!(ret_len),
+        )
+    };
+
+    NetworkCountersQuery {
+        status,
+        written: ret_len,
+        counters: info,
+    }
+}
+
+/// Reads the cumulative network traffic of `handle`.
+///
+/// `None` when the class is not implemented - it is a Windows 11 addition, so
+/// an older build answers `STATUS_INVALID_INFO_CLASS` - when the handle lacks
+/// the rights, or when the kernel wrote less than the whole structure.
+pub fn process_network_counters(handle: HANDLE) -> Option<PROCESS_NETWORK_COUNTERS> {
+    let query = query_network_counters(handle);
+    if query.status != STATUS_SUCCESS
+        || (query.written as usize) < size_of::<PROCESS_NETWORK_COUNTERS>()
+    {
+        return None;
+    }
+    Some(query.counters)
+}
+
 /// Reads the PEB base address of `handle`.
 ///
 /// `ProcessBasicInformation` only needs `PROCESS_QUERY_LIMITED_INFORMATION`.
@@ -1604,4 +1680,28 @@ mod tests {
         let handle = unsafe { GetCurrentProcess() };
         assert_eq!(process_image_machine(handle), Some(OWN_MACHINE));
     }
+
+    /// `ProcessNetworkIoInformation` is a Windows 11 addition, so whether the
+    /// counters come back at all depends on the build. What has to hold on
+    /// every build is that the 16-byte structure is accepted: a wrong
+    /// `PROCESS_NETWORK_COUNTERS` layout is answered with
+    /// `STATUS_INFO_LENGTH_MISMATCH`, which would leave both network columns
+    /// empty without a word of explanation.
+    #[test]
+    fn network_counters_layout_is_accepted() {
+        use windows_sys::Win32::System::Threading::GetCurrentProcess;
+
+        // SAFETY: the pseudo-handle needs no closing.
+        let handle = unsafe { GetCurrentProcess() };
+        let query = query_network_counters(handle);
+
+        assert_ne!(query.status, STATUS_INFO_LENGTH_MISMATCH);
+        if query.status == STATUS_SUCCESS {
+            assert_eq!(
+                query.written as usize,
+                size_of::<PROCESS_NETWORK_COUNTERS>()
+            );
+            assert!(process_network_counters(handle).is_some());
+        }
+    }
 }

src/process/windows.rs

@@ -6,6 +6,7 @@ use std::ffi::c_void;
 use std::mem::{MaybeUninit, zeroed};
 use std::path::PathBuf;
 use std::ptr;
+use std::sync::atomic::{AtomicU8, Ordering};
 use std::thread;
 use std::time::{Duration, Instant};
 use windows_sys::Win32::Foundation::{CloseHandle, FALSE, HANDLE};
@@ -64,6 +65,7 @@ pub struct ProcessInfo {
     pub cpu_info: CpuInfo,
     pub memory_info: MemoryInfo,
     pub disk_info: DiskInfo,
+    pub net_info: NetInfo,
     pub user: Option<SID_MAX>,
     pub groups: Vec<SID_MAX>,
     pub priority: i32,
@@ -119,6 +121,18 @@ pub struct DiskInfo {
     pub curr_write: u64,
 }
 
+/// Cumulative network traffic sampled at both ends of the interval.
+///
+/// `ProcessNetworkIoInformation` is a Windows 11 addition and needs a process
+/// handle, so both the class and the process can be missing: the counters stay
+/// at zero then and the rate comes out as zero.
+pub struct NetInfo {
+    pub prev_recv: u64,
+    pub prev_send: u64,
+    pub curr_recv: u64,
+    pub curr_send: u64,
+}
+
 pub struct CpuInfo {
     pub prev_sys: u64,
     pub prev_user: u64,
@@ -168,11 +182,9 @@ pub fn collect_proc(
             // A recycled pid would otherwise pair up with an unrelated process.
             .filter(|p| p.create_time == proc.create_time || proc.create_time == 0);
 
-        let (prev_sys, prev_user, prev_read, prev_write) = match prev {
-            Some(p) => (p.kernel_time, p.user_time, p.read, p.write),
-            // Started between the two samples: report no delta.
-            None => (proc.kernel_time, proc.user_time, proc.read, proc.write),
-        };
+        // A process that started between the two samples pairs with itself,
+        // which reports no delta.
+        let prev = prev.unwrap_or(&proc);
 
         let handles = ProcHandles::open(proc.pid);
 
@@ -199,18 +211,25 @@ pub fn collect_proc(
             ppid: proc.ppid,
             start_time: filetime_to_local(proc.create_time),
             cpu_info: CpuInfo {
-                prev_sys,
-                prev_user,
+                prev_sys: prev.kernel_time,
+                prev_user: prev.user_time,
                 curr_sys: proc.kernel_time,
                 curr_user: proc.user_time,
             },
-            memory_info: proc.memory_info,
             disk_info: DiskInfo {
-                prev_read,
-                prev_write,
+                prev_read: prev.read,
+                prev_write: prev.write,
                 curr_read: proc.read,
                 curr_write: proc.write,
             },
+            net_info: NetInfo {
+                prev_recv: prev.recv,
+                prev_send: prev.send,
+                curr_recv: proc.recv,
+                curr_send: proc.send,
+            },
+            // Last: this moves out of `proc`, which `prev` may still borrow.
+            memory_info: proc.memory_info,
             user,
             groups: groups.unwrap_or_default(),
             priority,
@@ -276,6 +295,14 @@ pub fn collect_proc(
                     curr_read: 0,
                     curr_write: 0,
                 },
+                // Network traffic is only accounted per process, so a thread
+                // row reports none rather than repeating its process'.
+                net_info: NetInfo {
+                    prev_recv: 0,
+                    prev_send: 0,
+                    curr_recv: 0,
+                    curr_send: 0,
+                },
                 user,
                 groups,
                 priority: thread.priority,
@@ -305,6 +332,10 @@ struct ProcSnapshot {
     user_time: u64,
     read: u64,
     write: u64,
+    /// Cumulative network bytes received / sent, when the class that reports
+    /// them is implemented and the process could be opened.
+    recv: u64,
+    send: u64,
     memory_info: MemoryInfo,
     base_priority: i32,
     /// Scheduler state of the most active thread.
@@ -337,11 +368,24 @@ fn take_snapshot(with_thread: bool) -> SystemSnapshot {
         return SystemSnapshot { procs, threads };
     };
 
+    // `ProcessNetworkIoInformation` is a Windows 11 addition: on an older
+    // build no handle is opened for it at all.
+    let want_net = network_counters_supported();
+
     for entry in buffer.iter() {
         let info = entry.info();
+        let pid = info.UniqueProcessId as usize as i32;
         let ppid = info.InheritedFromUniqueProcessId as usize as i32;
         let image_name = entry.image_name();
 
+        // Unlike the rest of the snapshot, the network counters are read from
+        // the process itself, which costs an `OpenProcess` each.
+        let (recv, send) = if want_net {
+            network_counters(pid)
+        } else {
+            (0, 0)
+        };
+
         // Hide the kernel's own processes. The full snapshot's classification is
         // authoritative - a non-zero `SYSTEM_PROCESS_CLASSIFICATION` marks System,
         // Secure System, Registry and Memory Compression, which carry real parent
@@ -353,7 +397,7 @@ fn take_snapshot(with_thread: bool) -> SystemSnapshot {
             || (ppid == 4 && !image_name.contains('.'));
 
         procs.push(ProcSnapshot {
-            pid: info.UniqueProcessId as usize as i32,
+            pid,
             ppid,
             thread_count: info.NumberOfThreads as i32,
             image_name,
@@ -363,6 +407,8 @@ fn take_snapshot(with_thread: bool) -> SystemSnapshot {
             user_time: info.UserTime as u64,
             read: info.IoCounters.ReadTransferCount,
             write: info.IoCounters.WriteTransferCount,
+            recv,
+            send,
             memory_info: MemoryInfo {
                 page_fault_count: u64::from(info.VirtualMemoryCounters.PageFaultCount),
                 peak_working_set_size: info.VirtualMemoryCounters.PeakWorkingSetSize as u64,
@@ -394,6 +440,43 @@ fn take_snapshot(with_thread: bool) -> SystemSnapshot {
     SystemSnapshot { procs, threads }
 }
 
+/// Whether `ProcessNetworkIoInformation` answers on this machine.
+///
+/// The class is implemented from Windows 11 on, so the answer depends on the
+/// OS build alone: it is probed once against this process - whose handle
+/// carries every access right, so a failure means the class is missing rather
+/// than the handle being too weak - and remembered. An older build then pays
+/// nothing instead of an `OpenProcess` and a rejected query per process and
+/// per sample.
+fn network_counters_supported() -> bool {
+    const UNKNOWN: u8 = 0;
+    const YES: u8 = 1;
+    const NO: u8 = 2;
+    static SUPPORTED: AtomicU8 = AtomicU8::new(UNKNOWN);
+
+    match SUPPORTED.load(Ordering::Acquire) {
+        YES => true,
+        NO => false,
+        _ => {
+            // SAFETY: the pseudo-handle needs no closing.
+            let handle = unsafe { GetCurrentProcess() };
+            let yes = ntapi::process_network_counters(handle).is_some();
+            SUPPORTED.store(if yes { YES } else { NO }, Ordering::Release);
+            yes
+        }
+    }
+}
+
+/// The cumulative network counters of `pid`, or `(0, 0)` when the handle
+/// cannot be opened or the class is not implemented.
+fn network_counters(pid: i32) -> (u64, u64) {
+    let handles = ProcHandles::open(pid);
+    match handles.any().and_then(ntapi::process_network_counters) {
+        Some(counters) => (counters.BytesIn, counters.BytesOut),
+        None => (0, 0),
+    }
+}
+
 /// Idle and System have no command line and, depending on privileges, no
 /// readable image name either.
 fn image_fallback(proc: &ProcSnapshot) -> String {