Bug fixes

4 files changed, 69 insertions, 38 deletions

diff --git a/where-rs/src/servers.rs b/where-rs/src/servers.rs
index aab4322..9f112f8 100644
--- a/where-rs/src/servers.rs
+++ b/where-rs/src/servers.rs
@@ -1,5 +1,5 @@
 use std::io::ErrorKind;
-use std::net::{SocketAddr, UdpSocket};
+use std::net::{SocketAddr, ToSocketAddrs, UdpSocket};
 use std::time::Duration;
 use where_shared::error::{WhereError, WhereResult};
 use where_shared::{MAX_PAYLOAD_LENGTH, SessionCollection, WHERED_MAGIC};
@@ -7,15 +7,17 @@ use crate::config::{GlobalConfig, Server};
 
 impl Server {
     fn get_address(&self, config: &GlobalConfig) -> WhereResult<SocketAddr> {
-        let res: SocketAddr = match self.endpoint.parse() {
-            Ok(addr) => addr,
+        let res: SocketAddr = match self.endpoint.to_socket_addrs() {
+            Ok(mut addr) => {
+                addr.find(|i| i.is_ipv4()).unwrap()
+            },
             Err(_) => {
                 let mut endpoint = self.endpoint.clone();
                 let port = config.port.to_string();
 
                 endpoint.push(':');
                 endpoint.push_str(&port);
-                endpoint.parse()?
+                endpoint.to_socket_addrs()?.find(|i| i.is_ipv4()).unwrap()
             }
         };
 
diff --git a/where-rs/src/ui.rs b/where-rs/src/ui.rs
index d2c1826..56d1b5d 100644
--- a/where-rs/src/ui.rs
+++ b/where-rs/src/ui.rs
@@ -26,19 +26,33 @@ pub fn print_summary(mut sessions: Vec<Session>, config: GlobalConfig) {
     let tty_padding = max_key_with_min(&sessions, |s| s.tty.len(), 4);
     let pid_padding = max_key_with_min(&sessions, |s| s.pid.abs().checked_ilog10().unwrap_or_default() + 1 + (s.pid < 0) as u32, 4);
 
-    println!("{:pad_0$} {:<pad_1$} {:<pad_2$} {:<pad_3$} {:<pad_4$} {:<pad_5$} Since",
-             "Act",
-             "Host",
-             "Source",
-             "User",
-             "TTY",
-             "PID",
-             pad_0 = ACTIVE_PADDING,
-             pad_1 = host_padding,
-             pad_2 = remote_padding,
-             pad_3 = username_padding,
-             pad_4 = tty_padding,
-             pad_5 = pid_padding as usize);
+    if config.include_inactive {
+        println!("{:pad_0$} {:<pad_1$} {:<pad_2$} {:<pad_3$} {:<pad_4$} {:<pad_5$} Since",
+                 "Act",
+                 "Host",
+                 "Source",
+                 "User",
+                 "TTY",
+                 "PID",
+                 pad_0 = ACTIVE_PADDING,
+                 pad_1 = host_padding,
+                 pad_2 = remote_padding,
+                 pad_3 = username_padding,
+                 pad_4 = tty_padding,
+                 pad_5 = pid_padding as usize);
+    } else {
+        println!("{:<pad_1$} {:<pad_2$} {:<pad_3$} {:<pad_4$} {:<pad_5$} Since",
+                 "Host",
+                 "Source",
+                 "User",
+                 "TTY",
+                 "PID",
+                 pad_1 = host_padding,
+                 pad_2 = remote_padding,
+                 pad_3 = username_padding,
+                 pad_4 = tty_padding,
+                 pad_5 = pid_padding as usize);
+    }
 
     for session in sessions {
         if !config.include_inactive && !session.active {
@@ -57,19 +71,34 @@ pub fn print_summary(mut sessions: Vec<Session>, config: GlobalConfig) {
         let datetime = DateTime::from_timestamp(session.login_time, 0).unwrap();
         let time = datetime.format("%Y-%m-%d %H:%M:%S");
 
-        println!(" {:<pad_0$} {:<pad_1$} {:<pad_2$} {:<pad_3$} {:<pad_4$} {:<pad_5$} {}",
-                 active,
-                 host,
-                 remote,
-                 session.tty,
-                 session.user,
-                 session.pid,
-                 time,
-                 pad_0 = ACTIVE_PADDING,
-                 pad_1 = host_padding,
-                 pad_2 = remote_padding,
-                 pad_3 = username_padding,
-                 pad_4 = tty_padding,
-                 pad_5 = pid_padding as usize);
+        if config.include_inactive {
+            println!(" {:<pad_0$} {:<pad_1$} {:<pad_2$} {:<pad_3$} {:<pad_4$} {:<pad_5$} {}",
+                     active,
+                     host,
+                     remote,
+                     session.tty,
+                     session.user,
+                     session.pid,
+                     time,
+                     pad_0 = ACTIVE_PADDING,
+                     pad_1 = host_padding,
+                     pad_2 = remote_padding,
+                     pad_3 = username_padding,
+                     pad_4 = tty_padding,
+                     pad_5 = pid_padding as usize);
+        } else {
+            println!("{:<pad_1$} {:<pad_2$} {:<pad_3$} {:<pad_4$} {:<pad_5$} {}",
+                     host,
+                     remote,
+                     session.tty,
+                     session.user,
+                     session.pid,
+                     time,
+                     pad_1 = host_padding,
+                     pad_2 = remote_padding,
+                     pad_3 = username_padding,
+                     pad_4 = tty_padding,
+                     pad_5 = pid_padding as usize);
+        }
     }
 }
diff --git a/where-shared/src/lib.rs b/where-shared/src/lib.rs
index dc28f9f..8ff483d 100644
--- a/where-shared/src/lib.rs
+++ b/where-shared/src/lib.rs
@@ -110,13 +110,13 @@ impl Session {
     pub fn from_udp_payload(cursor: &mut PayloadCursor, host: &str) -> WhereResult<Self> {
         let pid = parse::read_field(cursor, |buf| Ok(i32::from_be_bytes(buf)))?;
         let login_time = parse::read_field(cursor, |buf| Ok(i64::from_be_bytes(buf)))?;
-        let user = parse::read_string_field(cursor)?;
-        let tty = parse::read_string_field(cursor)?;
+        let user = parse::read_string_field(cursor, MAX_USER_TTY_LENGTH as u32)?;
+        let tty = parse::read_string_field(cursor, MAX_USER_TTY_LENGTH as u32)?;
 
         let remote = {
             let has_remote_tag = parse::read_bool_field(cursor)?;
             if has_remote_tag {
-                Some(parse::read_string_field(cursor)?)
+                Some(parse::read_string_field(cursor, MAX_REMOTE_LENGTH as u32)?)
             } else {
                 None
             }
@@ -161,6 +161,7 @@ impl Session {
                 let host_bytes = host.into_bytes();
                 let host_length = (host_bytes.len() as u32).to_be_bytes();
 
+                bytes.push(1u8);
                 bytes.extend(&host_length);
                 bytes.extend(&host_bytes);
             }
diff --git a/where-shared/src/parse.rs b/where-shared/src/parse.rs
index e624ce6..83542d8 100644
--- a/where-shared/src/parse.rs
+++ b/where-shared/src/parse.rs
@@ -1,7 +1,6 @@
 use std::io::Read;
 
 use crate::error::{EncodeDecodeError, WhereError, WhereResult};
-use crate::MAX_USER_TTY_LENGTH;
 use crate::PayloadCursor;
 
 pub fn read_field<const N: usize, F, T>(cursor: &mut PayloadCursor, convert_func: F) -> WhereResult<T>
@@ -31,11 +30,11 @@ pub fn read_bool_field(cursor: &mut PayloadCursor) -> WhereResult<bool> {
     Ok(value)
 }
 
-pub fn read_string_field(cursor: &mut PayloadCursor) -> WhereResult<String> {
+pub fn read_string_field(cursor: &mut PayloadCursor, max_length: u32) -> WhereResult<String> {
     let string_length = read_field(cursor, |buf| Ok(u32::from_be_bytes(buf)))?;
 
-    if string_length > MAX_USER_TTY_LENGTH as u32 {
-        return Err(WhereError::EncodeDecodeError(EncodeDecodeError::StringSizeLimitExceeded(string_length, MAX_USER_TTY_LENGTH)));
+    if string_length > max_length {
+        return Err(WhereError::EncodeDecodeError(EncodeDecodeError::StringSizeLimitExceeded(string_length, max_length as usize)));
     }
 
     let string = read_field_dynamic(cursor, string_length as usize, |buf| Ok(String::from_utf8(buf)?))?;