diff options
| l--------- | .direnv/flake-profile | 1 | ||||
| l--------- | .direnv/flake-profile-1-link | 1 | ||||
| -rw-r--r-- | .envrc | 2 | ||||
| -rw-r--r-- | .idea/discord.xml | 12 | ||||
| -rw-r--r-- | flake.lock | 65 | ||||
| -rw-r--r-- | flake.nix | 24 | ||||
| -rw-r--r-- | shell.nix | 34 | ||||
| -rw-r--r-- | where-rs/Cargo.toml | 2 | ||||
| -rw-r--r-- | where-rs/src/main.rs | 46 | ||||
| -rw-r--r-- | where-shared/src/error.rs | 67 | ||||
| -rw-r--r-- | where-shared/src/lib.rs | 301 | ||||
| -rw-r--r-- | whered/src/main.rs | 43 |
12 files changed, 490 insertions, 108 deletions
diff --git a/.direnv/flake-profile b/.direnv/flake-profile new file mode 120000 index 0000000..0c05709 --- /dev/null +++ b/.direnv/flake-profile @@ -0,0 +1 @@ +flake-profile-1-link
\ No newline at end of file diff --git a/.direnv/flake-profile-1-link b/.direnv/flake-profile-1-link new file mode 120000 index 0000000..7d76e2e --- /dev/null +++ b/.direnv/flake-profile-1-link @@ -0,0 +1 @@ +/nix/store/wdy5lg9svwcwzfnfmmpj34as2dhyf1w3-nix-shell-env
\ No newline at end of file @@ -0,0 +1,2 @@ +#!/usr/bin/env bash +use flake diff --git a/.idea/discord.xml b/.idea/discord.xml new file mode 100644 index 0000000..3aef922 --- /dev/null +++ b/.idea/discord.xml @@ -0,0 +1,12 @@ +<?xml version="1.0" encoding="UTF-8"?> +<project version="4"> + <component name="DiscordProjectSettings"> + <option name="show" value="PROJECT_FILES" /> + <option name="description" value="" /> + <option name="theme" value="material" /> + <option name="button1Title" value="" /> + <option name="button1Url" value="" /> + <option name="button2Title" value="" /> + <option name="button2Url" value="" /> + </component> +</project>
\ No newline at end of file diff --git a/flake.lock b/flake.lock new file mode 100644 index 0000000..5eac8ff --- /dev/null +++ b/flake.lock @@ -0,0 +1,65 @@ +{ + "nodes": { + "fenix": { + "inputs": { + "nixpkgs": [ + "nixpkgs" + ], + "rust-analyzer-src": "rust-analyzer-src" + }, + "locked": { + "lastModified": 1709792596, + "narHash": "sha256-DQL1KJ9AaoQjwMkZkSBrW9/az2dwbbBnhNIbIzgeDIE=", + "owner": "nix-community", + "repo": "fenix", + "rev": "221fedb628b1e86e88d8fbfbd20b699448e58fa9", + "type": "github" + }, + "original": { + "owner": "nix-community", + "repo": "fenix", + "type": "github" + } + }, + "nixpkgs": { + "locked": { + "lastModified": 1709703039, + "narHash": "sha256-6hqgQ8OK6gsMu1VtcGKBxKQInRLHtzulDo9Z5jxHEFY=", + "owner": "NixOS", + "repo": "nixpkgs", + "rev": "9df3e30ce24fd28c7b3e2de0d986769db5d6225d", + "type": "github" + }, + "original": { + "id": "nixpkgs", + "ref": "nixos-unstable", + "type": "indirect" + } + }, + "root": { + "inputs": { + "fenix": "fenix", + "nixpkgs": "nixpkgs" + } + }, + "rust-analyzer-src": { + "flake": false, + "locked": { + "lastModified": 1709749600, + "narHash": "sha256-jLGAyPPi4PVo/9y6ayWgdAtzPotFWvNTCxrDzxO+Uj0=", + "owner": "rust-lang", + "repo": "rust-analyzer", + "rev": "ce15e73a8ece3cf9e7958bbad3aedb03daa10135", + "type": "github" + }, + "original": { + "owner": "rust-lang", + "ref": "nightly", + "repo": "rust-analyzer", + "type": "github" + } + } + }, + "root": "root", + "version": 7 +} diff --git a/flake.nix b/flake.nix new file mode 100644 index 0000000..3780814 --- /dev/null +++ b/flake.nix @@ -0,0 +1,24 @@ +{ + inputs = { + nixpkgs.url = "nixpkgs/nixos-unstable"; + + fenix = { + url = "github:nix-community/fenix"; + inputs.nixpkgs.follows = "nixpkgs"; + }; + }; + + outputs = { self, nixpkgs, fenix }: + let + # TODO: Support more architectures + system = "x86_64-linux"; + overlays = [ fenix.overlays.default ]; + + pkgs = import nixpkgs { + inherit system overlays; + }; + in + { + devShells.${system}.default = pkgs.callPackage ./shell.nix {}; + }; +} diff --git a/shell.nix b/shell.nix new file mode 100644 index 0000000..3755099 --- /dev/null +++ b/shell.nix @@ -0,0 +1,34 @@ +{ mkShell, pkgsCross, fenix }: + +# Reference: https://github.com/nix-community/naersk/blob/aeb58d5e8faead8980a807c840232697982d47b9/examples/cross-windows/flake.nix + +let + channel = "stable"; + + targets = with fenix.targets; [ + x86_64-unknown-linux-gnu + ]; + + crossCompileTargets = with fenix.targets; [ + x86_64-pc-windows-gnu + ]; + + toolchain = fenix.combine (with fenix.${channel}; [ + cargo + rustc + ] ++ map (target: target.${channel}.toolchain) targets + ++ map (target: target.${channel}.rust-std) crossCompileTargets); + +in + +mkShell { + buildInputs = with pkgsCross.mingwW64; [ + stdenv.cc + ]; + + packages = [ toolchain ]; + + # Link to libpthreads manually otherwise build fails + # See: https://github.com/NixOS/nixpkgs/issues/139966#issuecomment-1385222547 + env.CARGO_TARGET_X86_64_PC_WINDOWS_GNU_RUSTFLAGS = "-L native=${pkgsCross.mingwW64.windows.pthreads}/lib"; +} diff --git a/where-rs/Cargo.toml b/where-rs/Cargo.toml index a2722d8..06cb9e1 100644 --- a/where-rs/Cargo.toml +++ b/where-rs/Cargo.toml @@ -11,4 +11,4 @@ path = "src/main.rs" [dependencies] where-shared = { path = "../where-shared" } -coreutils_core = "0.1.1"
\ No newline at end of file +coreutils_core = "0.1.1" diff --git a/where-rs/src/main.rs b/where-rs/src/main.rs index 0e23482..ba846a5 100644 --- a/where-rs/src/main.rs +++ b/where-rs/src/main.rs @@ -1,15 +1,41 @@ use std::net::UdpSocket; -use coreutils_core::ByteSlice; -use where_shared::*; +use std::io::ErrorKind; +use std::time::Duration; +use where_shared::error::{WhereError, WhereResult}; +use where_shared::{MAX_PAYLOAD_LENGTH, SessionCollection, WHERED_MAGIC}; + +pub const TIMEOUT: Duration = Duration::from_millis(2000); +pub const MAX_SEND_RETRIES: usize = 3; fn main() { - let socket = UdpSocket::bind("0.0.0.0:0").expect("Could not start a UDP socket."); - socket.send_to(&WHERED_MAGIC, "127.0.0.1:15").expect("Could not send data to the server."); + if let Err(e) = start_client() { + eprintln!("where: {}", e); + std::process::exit(1); + } +} + +fn start_client() -> WhereResult<()> { + println!("{:?}", process_server("127.0.0.1:15")?); + Ok(()) +} + +fn process_server(server: &str) -> WhereResult<SessionCollection> { + let socket = UdpSocket::bind("0.0.0.0:0")?; + socket.set_read_timeout(Some(TIMEOUT))?; + + let mut buf = [0; MAX_PAYLOAD_LENGTH]; + + for _ in 0..MAX_SEND_RETRIES { + socket.send_to(&WHERED_MAGIC, server)?; - let mut buf = [0; 1024]; - socket.recv_from(&mut buf).expect("No data to receive from the server."); + match socket.recv_from(&mut buf) { + Ok(_) => { + return Ok(SessionCollection::from_udp_payload(buf)?); + }, + Err(e) if e.kind() == ErrorKind::TimedOut || e.kind() == ErrorKind::WouldBlock => continue, + Err(e) => return Err(WhereError::from(e)), + } + } - /*let list = SessionCollection::from_bytes(buf.to_vec()); - println!("{:?}", list);*/ - println!("{}", String::from_utf8_lossy(&buf)); -}
\ No newline at end of file + Err(WhereError::TimedOut(server.to_string(), MAX_SEND_RETRIES, TIMEOUT)) +} diff --git a/where-shared/src/error.rs b/where-shared/src/error.rs new file mode 100644 index 0000000..e94b042 --- /dev/null +++ b/where-shared/src/error.rs @@ -0,0 +1,67 @@ +use std::fmt::Display; +use std::{fmt, io}; +use std::time::Duration; +use crate::{MAX_ENTRY_LENGTH, MAX_PAYLOAD_LENGTH}; + +pub enum WhereError { + EncodeDecodeError(EncodeDecodeError), + IOError(io::Error), + TimedOut(String, usize, Duration) +} + +pub enum EncodeDecodeError { + InvalidEntryLength(usize), + InvalidPayloadLength(usize), + BadMagic([u8; 4]), + IncorrectEntryCount, + StringSizeLimitExceeded(u32, usize), + NonbinaryBoolean, + EmptyRemote, + IOErrorWhileTranscoding(io::Error) +} + +pub type WhereResult<T> = Result<T, WhereError>; +pub type EncodeDecodeResult<T> = Result<T, EncodeDecodeError>; + +impl From<io::Error> for WhereError { + fn from(value: io::Error) -> Self { + Self::IOError(value) + } +} + +impl From<EncodeDecodeError> for WhereError { + fn from(value: EncodeDecodeError) -> Self { + Self::EncodeDecodeError(value) + } +} + +impl From<io::Error> for EncodeDecodeError { + fn from(value: io::Error) -> Self { + Self::IOErrorWhileTranscoding(value) + } +} + +impl Display for EncodeDecodeError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::InvalidEntryLength(s) => write!(f, "Invalid entry length: {s} but maximum is {MAX_ENTRY_LENGTH}"), + Self::InvalidPayloadLength(s) => write!(f, "Invalid full payload length: {s} but maximum is {MAX_PAYLOAD_LENGTH}"), + Self::BadMagic(m) => write!(f, "Invalid packet magic ({}), possible corruption or invalid server", String::from_utf8_lossy(m)), + Self::IncorrectEntryCount => write!(f, "Invalid amount of entries decoded"), + Self::StringSizeLimitExceeded(curr, max) => write!(f, "Exceeded length limit for payload string ({curr} > {max})"), + Self::NonbinaryBoolean => write!(f, "Boolean value is not 0 or 1"), + Self::EmptyRemote => write!(f, "Remote tag set but no remote host is present"), + Self::IOErrorWhileTranscoding(e) => write!(f, "Input/output error while encoding/decoding: {e}"), + } + } +} + +impl Display for WhereError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::EncodeDecodeError(e) => write!(f, "Encode/decode error: {e}"), + Self::IOError(e) => write!(f, "Input/output error: {e}"), + Self::TimedOut(server, max_retry, timeout) => write!(f, "Timed out waiting for data from {server} after {max_retry} attempts every {} ms", timeout.as_millis()) + } + } +}
\ No newline at end of file diff --git a/where-shared/src/lib.rs b/where-shared/src/lib.rs index 1ea21a7..d64b284 100644 --- a/where-shared/src/lib.rs +++ b/where-shared/src/lib.rs @@ -1,7 +1,16 @@ -use coreutils_core::ByteSlice; -use coreutils_core::os::utmpx::UtmpxKind; +use std::io::Cursor; +use coreutils_core::os::utmpx::*; +use std::io::Read; +use crate::error::{EncodeDecodeError, EncodeDecodeResult}; + +pub mod error; pub const WHERED_MAGIC: [u8; 4] = *b"WHRD"; +pub const MAX_USER_TTY_LENGTH: usize = 32; +pub const MAX_REMOTE_LENGTH: usize = 64; +pub const MAX_ENTRY_LENGTH: usize = MAX_REMOTE_LENGTH + MAX_USER_TTY_LENGTH * 2 + 25; +pub const MAX_PAYLOAD_LENGTH: usize = 65501; +pub const MAX_PAYLOAD_ENTRIES: usize = MAX_PAYLOAD_LENGTH / MAX_ENTRY_LENGTH; #[derive(Debug)] pub struct Session { @@ -10,112 +19,240 @@ pub struct Session { tty: String, remote: Option<String>, active: bool, - login: i64 + login_time: i64 } #[derive(Debug)] -pub struct SessionCollection<Session> { +pub struct SessionCollection { inner: Vec<Session> } -impl SessionCollection<Session> { - pub fn fetch() -> SessionCollection<Session> { - let mut output: SessionCollection<Session> = SessionCollection { - inner: vec![] - }; - let utmp = coreutils_core::os::utmpx::UtmpxSet::system(); +impl SessionCollection { + pub fn fetch() -> Self { + let inner: Vec<Session> = UtmpxSet::system() + .into_iter() + .filter(|utmpx| utmpx.entry_type() == UtmpxKind::UserProcess || utmpx.entry_type() == UtmpxKind::DeadProcess) + .map(Session::from) + .collect(); - for item in utmp { - if item.entry_type() != UtmpxKind::UserProcess && item.entry_type() != UtmpxKind::DeadProcess { - continue; - } - - let host = item.host().to_string(); - - output.inner.push(Session { - user: item.user().to_string(), - pid: item.process_id(), - tty: item.device_name().to_string(), - remote: if &host == "" { - None - } else { - Some(host) - }, - active: item.entry_type() == UtmpxKind::UserProcess, - login: item.timeval().tv_sec - }); + Self { + inner } - - output } - pub fn into_bytes(self) -> Vec<u8> { + pub fn to_udp_payload(self) -> EncodeDecodeResult<Vec<u8>> { + println!("Encoding payload with {} entries", self.inner.len()); + let mut bytes: Vec<u8> = vec![]; + bytes.extend(&WHERED_MAGIC); + + let entry_count = (self.inner.len() as u16).to_be_bytes(); + bytes.extend(&entry_count); for item in self.inner { - bytes.append(&mut item.into_bytes().to_vec()); - bytes.append(&mut vec![0, 0, 0, 0]); + let entry = item.to_udp_payload(); + + if entry.len() > MAX_ENTRY_LENGTH { + return Err(EncodeDecodeError::InvalidEntryLength(entry.len())); + } + + bytes.extend(entry); } - bytes + if bytes.len() > MAX_PAYLOAD_LENGTH { + Err(EncodeDecodeError::InvalidPayloadLength(bytes.len())) + } else { + Ok(bytes) + } } - /*pub fn from_bytes(bytes: Vec<u8>) -> SessionCollection<Session> { - let entries = bytes.split([0, 0]); - let mut final_entries: Vec<Session> = vec![]; + pub fn from_udp_payload(buffer: [u8; MAX_PAYLOAD_LENGTH]) -> EncodeDecodeResult<Self> { + let mut buf = Cursor::new(buffer); + let mut inner = vec![]; + let mut magic = [0u8; 4]; + let mut length = [0u8; 2]; - for item in entries { - let parts: Vec<&[u8]> = item.split(0).collect(); - if item.split(0).count() == 0 { - continue; - } + Session::read_field(&mut buf, &mut magic)?; + Session::read_field(&mut buf, &mut length)?; + let entry_count = u16::from_be_bytes(length); - /*final_entries.push(Session { - //pid: i32::from_be_bytes(parts[0].as_bytes()), - //login: i64::from_be_bytes(parts[1].as_bytes()), - pid: 0, - login: 0, - user: parts[2].to_string(), - tty: parts[3].to_string(), - remote: if parts[4] == b"\xff".to_str().unwrap() { - None - } else { - Some(parts[4].to_string()) - }, - active: parts[5] == b"\xff".to_str().unwrap() - });*/ + if magic != WHERED_MAGIC { + return Err(EncodeDecodeError::BadMagic(magic)); } - SessionCollection { - inner: final_entries + for _ in 0..entry_count { + inner.push(Session::from_udp_payload(&mut buf)?); } - }*/ + + if inner.len() != entry_count as usize { + return Err(EncodeDecodeError::IncorrectEntryCount); + } + + Ok(Self { + inner + }) + } } impl Session { - pub fn into_bytes(self) -> Vec<u8> { - let mut bytes: Vec<u8> = vec![]; - let mut host_bytes = self.remote.clone().unwrap_or(String::from("!")).as_bytes().to_vec(); - let mut full = vec![255]; - - bytes.append(&mut self.pid.to_be_bytes().to_vec()); - bytes.append(&mut self.login.to_be_bytes().to_vec()); - bytes.append(&mut (self.user.len() as u32).to_be_bytes().to_vec()); - bytes.append(&mut self.user.as_bytes().to_vec()); - bytes.append(&mut (self.tty.len() as u32).to_be_bytes().to_vec()); - bytes.append(&mut self.tty.as_bytes().to_vec()); - bytes.append(&mut (host_bytes.len() as u32).to_be_bytes().to_vec()); - bytes.append(if let Some(_) = self.remote { - &mut host_bytes - } else { - &mut full - }); - bytes.push(if self.active { - 1 + pub fn from_udp_payload(cursor: &mut Cursor<[u8; MAX_PAYLOAD_LENGTH]>) -> EncodeDecodeResult<Self> { + let mut username_length = [0u8; 4]; + let mut pid = [0u8; 4]; + let mut tty_length = [0u8; 4]; + let mut remote_tag = [0u8; 1]; + let mut remote_length = [0u8; 4]; + let mut active = [0u8; 1]; + let mut login_time = [0u8; 8]; + + Session::read_field(cursor, &mut pid)?; + Session::read_field(cursor, &mut login_time)?; + + Session::read_field(cursor, &mut username_length)?; + let username_length = u32::from_be_bytes(username_length); + if username_length as usize > MAX_USER_TTY_LENGTH { + return Err(EncodeDecodeError::StringSizeLimitExceeded(username_length, MAX_USER_TTY_LENGTH)); + } + + let mut user = vec![0u8; username_length as usize]; + Session::read_field(cursor, &mut user)?; + + Session::read_field(cursor, &mut tty_length)?; + let tty_length = u32::from_be_bytes(tty_length); + if tty_length as usize > MAX_USER_TTY_LENGTH { + return Err(EncodeDecodeError::StringSizeLimitExceeded(tty_length, MAX_USER_TTY_LENGTH)); + } + + let mut tty = vec![0u8; tty_length as usize]; + Session::read_field(cursor, &mut tty)?; + + Session::read_field(cursor, &mut remote_tag)?; + if remote_tag[0] > 1 { + return Err(EncodeDecodeError::NonbinaryBoolean); + } + + let has_remote_tag = remote_tag[0] == 1; + + let remote = if has_remote_tag { + Session::read_field(cursor, &mut remote_length)?; + let remote_length = u32::from_be_bytes(remote_length); + if remote_length as usize > MAX_USER_TTY_LENGTH { + return Err(EncodeDecodeError::StringSizeLimitExceeded(username_length, MAX_USER_TTY_LENGTH)); + } + + if remote_length == 0 { + return Err(EncodeDecodeError::EmptyRemote); + } + + let mut remote = vec![0u8; remote_length as usize]; + Session::read_field(cursor, &mut remote)?; + + Some(String::from_utf8_lossy(&remote).to_string()) } else { - 0 - }); + None + }; + + Session::read_field(cursor, &mut active)?; + if active[0] > 1 { + return Err(EncodeDecodeError::NonbinaryBoolean); + } + + let user = String::from_utf8_lossy(&user).to_string(); + let pid = i32::from_be_bytes(pid); + let tty = String::from_utf8_lossy(&tty).to_string(); + let active = active[0] == 1; + let login_time = i64::from_be_bytes(login_time); + + Ok(Self { + user, + pid, + tty, + remote, + active, + login_time + }) + } + + fn read_field(cursor: &mut Cursor<[u8; MAX_PAYLOAD_LENGTH]>, buffer: &mut [u8]) -> EncodeDecodeResult<()> { + cursor.read_exact(buffer)?; + Ok(()) + } + + /*fn read_field<T, F>(cursor: &mut Cursor<&[u8]>, convert_func: F) -> WhereResult<T> + where + N: const usize, + F: FnOnce([u8; N]) -> EncodeDecodeError, + { + let mut buf = [0u8; N]; + cursor.read_exact(&mut buf)?; + + let value = convert_func(buf)?; + Ok(value) + }*/ + + pub fn to_udp_payload(self) -> Vec<u8> { + let mut bytes: Vec<u8> = vec![]; + + let pid = self.pid.to_be_bytes(); + let login_time = self.login_time.to_be_bytes(); + let user_length = (self.user.len() as u32).to_be_bytes(); + let user = self.user.as_bytes(); + let tty_length = (self.tty.len() as u32).to_be_bytes(); + let tty = self.tty.as_bytes(); + let active = self.active as u8; + + bytes.extend(&pid); + bytes.extend(&login_time); + bytes.extend(&user_length); + bytes.extend(user); + bytes.extend(&tty_length); + bytes.extend(tty); + + match self.remote { + None => bytes.push(0u8), + Some(host) => { + let host_bytes = host.into_bytes(); + let host_length = (host_bytes.len() as u32).to_be_bytes(); + + bytes.extend(&host_length); + bytes.extend(&host_bytes); + } + } + + bytes.push(active); bytes } -}
\ No newline at end of file +} + +impl From<Utmpx> for Session { + fn from(utmpx: Utmpx) -> Self { + // BStr doesn't have a known size at compile time, so we can't use it instead of String + let mut host = utmpx.host().to_string(); + host.truncate(MAX_REMOTE_LENGTH); + + let mut user = utmpx.user().to_string(); + user.truncate(MAX_USER_TTY_LENGTH); + + let pid = utmpx.process_id(); + // In the case of a user session, this will always be a TTY + let mut tty = utmpx.device_name().to_string(); + tty.truncate(MAX_USER_TTY_LENGTH); + + let remote = if host.is_empty() { + None + } else { + Some(host) + }; + let active = utmpx.entry_type() == UtmpxKind::UserProcess; + let login_time = utmpx.timeval().tv_sec; + + Self { + user, + pid, + tty, + remote, + active, + login_time + } + } +} diff --git a/whered/src/main.rs b/whered/src/main.rs index efc3fcc..6990dd6 100644 --- a/whered/src/main.rs +++ b/whered/src/main.rs @@ -1,23 +1,36 @@ use std::net::UdpSocket; -use where_shared::*; +use where_shared::error::WhereResult; +use where_shared::{SessionCollection, WHERED_MAGIC}; fn main() { - let socket = UdpSocket::bind("0.0.0.0:15").expect("Could not bind to port 15, is another instance of whered running?"); + if let Err(e) = run_server() { + eprintln!("whered: {}", e); + std::process::exit(1); + } +} + +fn run_server() -> WhereResult<()> { + let socket = UdpSocket::bind("0.0.0.0:15")?; + println!("Now listening on 0.0.0.0:15"); loop { - let mut buf = [0; WHERED_MAGIC.len()]; - let Ok((_, src)) = socket.recv_from(&mut buf) else { - eprintln!("Failed to receive data from the client, ignoring"); - continue - }; - - println!("{src}: New client!"); - - let sessions = SessionCollection::fetch(); - if let Err(_) = socket.send_to(&*sessions.into_bytes(), src) { - eprintln!("{src}: Failed to send data back to the client, ignoring"); - } else { - println!("{src}: Completed request"); + if let Err(e) = handle_request(&socket) { + eprintln!("whered: {}", e); } } +} + +fn handle_request(socket: &UdpSocket) -> WhereResult<()> { + let mut buf = [0; WHERED_MAGIC.len()]; + + let (_, src) = socket.recv_from(&mut buf)?; + println!("{src}: New client!"); + + let sessions = SessionCollection::fetch(); + let buf = sessions.to_udp_payload()?; + + socket.send_to(&buf, src)?; + println!("{src}: Completed request within {} bytes", buf.len()); + + Ok(()) }
\ No newline at end of file |