more churn
This commit is contained in:
		
							parent
							
								
									67d147c7e3
								
							
						
					
					
						commit
						46b87d2559
					
				
					 10 changed files with 486 additions and 125 deletions
				
			
		
							
								
								
									
										45
									
								
								Cargo.lock
									
									
									
										generated
									
									
									
								
							
							
						
						
									
										45
									
								
								Cargo.lock
									
									
									
										generated
									
									
									
								
							|  | @ -105,6 +105,12 @@ version = "1.3.2" | |||
| source = "registry+https://github.com/rust-lang/crates.io-index" | ||||
| checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a" | ||||
| 
 | ||||
| [[package]] | ||||
| name = "bitflags" | ||||
| version = "2.4.2" | ||||
| source = "registry+https://github.com/rust-lang/crates.io-index" | ||||
| checksum = "ed570934406eb16438a4e976b1b4500774099c13b8cb96eec99f620f05090ddf" | ||||
| 
 | ||||
| [[package]] | ||||
| name = "block-buffer" | ||||
| version = "0.10.4" | ||||
|  | @ -281,6 +287,16 @@ dependencies = [ | |||
|  "subtle", | ||||
| ] | ||||
| 
 | ||||
| [[package]] | ||||
| name = "errno" | ||||
| version = "0.3.8" | ||||
| source = "registry+https://github.com/rust-lang/crates.io-index" | ||||
| checksum = "a258e46cdc063eb8519c00b9fc845fc47bcfca4130e2f08e88665ceda8474245" | ||||
| dependencies = [ | ||||
|  "libc", | ||||
|  "windows-sys 0.52.0", | ||||
| ] | ||||
| 
 | ||||
| [[package]] | ||||
| name = "fiat-crypto" | ||||
| version = "0.2.5" | ||||
|  | @ -435,9 +451,15 @@ checksum = "e2abad23fbc42b3700f2f279844dc832adb2b2eb069b2df918f455c4e18cc646" | |||
| 
 | ||||
| [[package]] | ||||
| name = "libc" | ||||
| version = "0.2.151" | ||||
| version = "0.2.152" | ||||
| source = "registry+https://github.com/rust-lang/crates.io-index" | ||||
| checksum = "302d7ab3130588088d277783b1e2d2e10c9e9e4a16dd9050e6ec93fb3e7048f4" | ||||
| checksum = "13e3bf6590cbc649f4d1a3eefc9d5d6eb746f5200ffb04e5e142700b8faa56e7" | ||||
| 
 | ||||
| [[package]] | ||||
| name = "linux-raw-sys" | ||||
| version = "0.4.13" | ||||
| source = "registry+https://github.com/rust-lang/crates.io-index" | ||||
| checksum = "01cda141df6706de531b6c46c3a33ecca755538219bd484262fa09410c13539c" | ||||
| 
 | ||||
| [[package]] | ||||
| name = "lock_api" | ||||
|  | @ -687,6 +709,7 @@ dependencies = [ | |||
|  "bytes", | ||||
|  "chrono", | ||||
|  "futures", | ||||
|  "rustix", | ||||
|  "serde", | ||||
|  "serde_json", | ||||
|  "sysinfo", | ||||
|  | @ -743,7 +766,7 @@ version = "0.4.1" | |||
| source = "registry+https://github.com/rust-lang/crates.io-index" | ||||
| checksum = "4722d768eff46b75989dd134e5c353f0d6296e5aaa3132e776cbdb56be7731aa" | ||||
| dependencies = [ | ||||
|  "bitflags", | ||||
|  "bitflags 1.3.2", | ||||
| ] | ||||
| 
 | ||||
| [[package]] | ||||
|  | @ -804,6 +827,20 @@ dependencies = [ | |||
|  "semver", | ||||
| ] | ||||
| 
 | ||||
| [[package]] | ||||
| name = "rustix" | ||||
| version = "0.38.30" | ||||
| source = "registry+https://github.com/rust-lang/crates.io-index" | ||||
| checksum = "322394588aaf33c24007e8bb3238ee3e4c5c09c084ab32bc73890b99ff326bca" | ||||
| dependencies = [ | ||||
|  "bitflags 2.4.2", | ||||
|  "errno", | ||||
|  "itoa", | ||||
|  "libc", | ||||
|  "linux-raw-sys", | ||||
|  "windows-sys 0.52.0", | ||||
| ] | ||||
| 
 | ||||
| [[package]] | ||||
| name = "rustls" | ||||
| version = "0.21.10" | ||||
|  | @ -884,7 +921,7 @@ version = "2.9.2" | |||
| source = "registry+https://github.com/rust-lang/crates.io-index" | ||||
| checksum = "05b64fb303737d99b81884b2c63433e9ae28abebe5eb5045dcdd175dc2ecf4de" | ||||
| dependencies = [ | ||||
|  "bitflags", | ||||
|  "bitflags 1.3.2", | ||||
|  "core-foundation", | ||||
|  "core-foundation-sys", | ||||
|  "libc", | ||||
|  |  | |||
|  | @ -9,6 +9,7 @@ async-nats = "0.33.0" | |||
| bytes = "1.5.0" | ||||
| chrono = { version = "0.4.33", default-features = false, features = ["now", "serde"] } | ||||
| futures = { version = "0.3.30", default-features = false, features = ["std"] } | ||||
| rustix = { version = "0.38.30", features = ["termios", "stdio", "pty", "process"] } | ||||
| serde = { version = "1.0.195", features = ["derive"] } | ||||
| serde_json = "1.0.111" | ||||
| sysinfo = { version = "0.30.5", default-features = false } | ||||
|  |  | |||
|  | @ -1,6 +1,6 @@ | |||
| //! System health information and checking
 | ||||
| 
 | ||||
| use std::{sync::Arc, time::Duration}; | ||||
| use std::{collections::HashMap, sync::Arc, time::Duration}; | ||||
| 
 | ||||
| use serde::{Deserialize, Serialize}; | ||||
| use tokio::sync::watch; | ||||
|  | @ -9,7 +9,7 @@ use crate::messaging::{Client, Message}; | |||
| 
 | ||||
| const MEMORY_USAGE_CRITICAL_THRESHOLD: f64 = 90.0; | ||||
| const CPU_USAGE_CRITICAL_THRESHOLD: f32 = 90.0; | ||||
| const DISK_USAGE_CRITICAL_THRESHOLD: f32 = 90.0; | ||||
| const DISK_USAGE_CRITICAL_THRESHOLD: f64 = 90.0; | ||||
| 
 | ||||
| pub struct System { | ||||
|     sys: sysinfo::System, | ||||
|  | @ -33,7 +33,7 @@ impl System { | |||
|                 .with_cpu(CpuRefreshKind::everything()), | ||||
|         ); | ||||
| 
 | ||||
|         // self.disks.refresh_list();
 | ||||
|         self.disks.refresh_list(); | ||||
|     } | ||||
| 
 | ||||
|     pub fn system(&self) -> &sysinfo::System { | ||||
|  | @ -45,18 +45,26 @@ impl System { | |||
|     } | ||||
| } | ||||
| 
 | ||||
| #[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] | ||||
| #[derive(Clone, Debug, Default, Serialize, Deserialize)] | ||||
| pub enum Status { | ||||
|     #[default] | ||||
|     Normal, | ||||
|     Critical, | ||||
| } | ||||
| 
 | ||||
| #[derive(Clone, Debug, Default, Serialize, Deserialize)] | ||||
| pub enum DiskStatus { | ||||
|     #[default] | ||||
|     Normal, | ||||
|     // HighUsage,
 | ||||
|     VeryHighUsage, | ||||
| } | ||||
| 
 | ||||
| #[derive(Clone, Debug, Default, Serialize, Deserialize)] | ||||
| pub struct Health { | ||||
|     cpu_status: Status, | ||||
|     memory_status: Status, | ||||
|     disk_status: Status, | ||||
|     disk_status: HashMap<String, DiskStatus>, | ||||
| } | ||||
| 
 | ||||
| #[derive(Clone)] | ||||
|  | @ -80,13 +88,12 @@ impl HealthMonitor { | |||
| 
 | ||||
|         let cpu_usage = sys.global_cpu_info().cpu_usage(); | ||||
| 
 | ||||
|         // for d in system.disks().list() {
 | ||||
|         //     let _avail = if d.total_space() > 0 {
 | ||||
|         //         (d.available_space() * 100 / d.total_space()) as u8
 | ||||
|         //     } else {
 | ||||
|         //         0 as u8
 | ||||
|         //     };
 | ||||
|         // }
 | ||||
|         let disks_usage = system.disks().list().iter().map(|dk| { | ||||
|             ( | ||||
|                 dk.name().to_str().unwrap_or("<INVALID DISK NAME DETECTED>"), | ||||
|                 (dk.total_space() - dk.available_space()) as f64 / dk.total_space() as f64 * 100.0, | ||||
|             ) | ||||
|         }); | ||||
| 
 | ||||
|         self.0.send_if_modified(|health| { | ||||
|             let cpu_changed = match health.cpu_status { | ||||
|  | @ -113,7 +120,42 @@ impl HealthMonitor { | |||
|                 _ => false, | ||||
|             }; | ||||
| 
 | ||||
|             cpu_changed || memory_changed | ||||
|             let mut disk_changed = false; | ||||
| 
 | ||||
|             disks_usage.for_each(|(name, usage)| match health.disk_status.get_mut(name) { | ||||
|                 Some(DiskStatus::Normal) if usage > DISK_USAGE_CRITICAL_THRESHOLD => { | ||||
|                     println!("{usage}"); | ||||
|                     health | ||||
|                         .disk_status | ||||
|                         .insert(name.to_owned(), DiskStatus::VeryHighUsage); | ||||
| 
 | ||||
|                     disk_changed = true; | ||||
|                 } | ||||
|                 Some(DiskStatus::VeryHighUsage) if usage <= DISK_USAGE_CRITICAL_THRESHOLD => { | ||||
|                     health | ||||
|                         .disk_status | ||||
|                         .insert(name.to_owned(), DiskStatus::Normal); | ||||
| 
 | ||||
|                     disk_changed = true; | ||||
|                 } | ||||
|                 None if usage > DISK_USAGE_CRITICAL_THRESHOLD => { | ||||
|                     health | ||||
|                         .disk_status | ||||
|                         .insert(name.to_owned(), DiskStatus::VeryHighUsage); | ||||
| 
 | ||||
|                     disk_changed = true; | ||||
|                 } | ||||
|                 None => { | ||||
|                     health | ||||
|                         .disk_status | ||||
|                         .insert(name.to_owned(), DiskStatus::Normal); | ||||
| 
 | ||||
|                     disk_changed = true; | ||||
|                 } | ||||
|                 _ => {} | ||||
|             }); | ||||
| 
 | ||||
|             cpu_changed || memory_changed || disk_changed | ||||
|         }); | ||||
|     } | ||||
| 
 | ||||
|  | @ -137,7 +179,7 @@ pub async fn init_health_subsystem(client: Client) -> HealthMonitor { | |||
| 
 | ||||
|     // Forever refresh system resources and monitor changes
 | ||||
|     std::thread::spawn(move || loop { | ||||
|         const REFRESH_INTERVAL: Duration = Duration::from_secs(1); | ||||
|         const REFRESH_INTERVAL: Duration = Duration::from_secs(5); | ||||
|         system.refresh_resources(); | ||||
|         health_monitor.check_system(&system); | ||||
|         std::thread::sleep(REFRESH_INTERVAL); | ||||
|  |  | |||
|  | @ -5,6 +5,7 @@ use tracing::Level; | |||
| 
 | ||||
| mod health; | ||||
| mod messaging; | ||||
| mod pty; | ||||
| mod services; | ||||
| 
 | ||||
| #[tokio::main] | ||||
|  | @ -28,7 +29,7 @@ async fn run() -> anyhow::Result<()> { | |||
|     let _health_monitor = init_health_subsystem(client.clone()).await; | ||||
|     tracing::info!("initialized health system"); | ||||
| 
 | ||||
|     init_services(client).await; | ||||
|     init_services(client).await?; | ||||
|     tracing::info!("initialized services"); | ||||
| 
 | ||||
|     tracing::info!("agent is ready"); | ||||
|  |  | |||
|  | @ -1,9 +1,8 @@ | |||
| use std::fmt::Display; | ||||
| use std::fmt::{Debug, Display}; | ||||
| 
 | ||||
| use anyhow::{anyhow, Result}; | ||||
| use bytes::Bytes; | ||||
| use futures::Stream; | ||||
| use serde::Deserialize; | ||||
| use tokio_stream::StreamExt; | ||||
| 
 | ||||
| use crate::health::Health; | ||||
|  | @ -18,6 +17,7 @@ pub struct Client { | |||
| pub enum Subject { | ||||
|     Health, | ||||
|     Exec, | ||||
|     OpenTerminal, | ||||
| } | ||||
| 
 | ||||
| impl Display for Subject { | ||||
|  | @ -25,10 +25,24 @@ impl Display for Subject { | |||
|         match self { | ||||
|             Subject::Health => write!(f, "health"), | ||||
|             Subject::Exec => write!(f, "exec"), | ||||
|             Subject::OpenTerminal => write!(f, "open_terminal"), | ||||
|         } | ||||
|     } | ||||
| } | ||||
| 
 | ||||
| impl TryFrom<&str> for Subject { | ||||
|     type Error = anyhow::Error; | ||||
| 
 | ||||
|     fn try_from(value: &str) -> Result<Self> { | ||||
|         Ok(match value { | ||||
|             "health" => Subject::Health, | ||||
|             "exec" => Subject::Exec, | ||||
|             "open_terminal" => Subject::OpenTerminal, | ||||
|             _ => return Err(anyhow!("unknown subject '{}'", value)), | ||||
|         }) | ||||
|     } | ||||
| } | ||||
| 
 | ||||
| #[derive(Debug)] | ||||
| pub struct Message { | ||||
|     subject: Subject, | ||||
|  | @ -40,27 +54,19 @@ impl Message { | |||
|     fn from_transport(msg: async_nats::Message) -> Result<Self> { | ||||
|         let suffix = msg.subject.split_terminator('.').last().unwrap_or_default(); | ||||
| 
 | ||||
|         match suffix { | ||||
|             "exec" => Ok(Message { | ||||
|                 subject: Subject::Exec, | ||||
|         Ok(Message { | ||||
|             subject: suffix.try_into()?, | ||||
|             payload: msg.payload, | ||||
|             reply: msg.reply, | ||||
|             }), | ||||
|             "health" => Ok(Message { | ||||
|                 subject: Subject::Health, | ||||
|                 payload: msg.payload, | ||||
|                 reply: msg.reply, | ||||
|             }), | ||||
|             _ => Err(anyhow!("unknown subject: {}", msg.subject)), | ||||
|         } | ||||
|         }) | ||||
|     } | ||||
| 
 | ||||
|     pub fn subject(&self) -> &Subject { | ||||
|         &self.subject | ||||
|     } | ||||
| 
 | ||||
|     pub fn parse_payload<'a, T: Deserialize<'a>>(&'a self) -> Result<T> { | ||||
|         Ok(serde_json::from_slice(&self.payload[..])?) | ||||
|     pub fn body(&self) -> Bytes { | ||||
|         self.payload.clone() | ||||
|     } | ||||
| 
 | ||||
|     pub fn health(health: Health) -> Result<Message> { | ||||
|  | @ -125,3 +131,9 @@ impl Client { | |||
|         Ok(subscriber) | ||||
|     } | ||||
| } | ||||
| 
 | ||||
| impl Debug for Client { | ||||
|     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { | ||||
|         write!(f, "Client {{ id: {} }}", self.id) | ||||
|     } | ||||
| } | ||||
|  |  | |||
							
								
								
									
										166
									
								
								agent/src/pty.rs
									
									
									
									
									
										Normal file
									
								
							
							
						
						
									
										166
									
								
								agent/src/pty.rs
									
									
									
									
									
										Normal file
									
								
							|  | @ -0,0 +1,166 @@ | |||
| use std::{io, task::ready}; | ||||
| 
 | ||||
| use rustix::{ | ||||
|     fd::OwnedFd, | ||||
|     fs::{fcntl_getfl, fcntl_setfl, OFlags}, | ||||
|     process::{ioctl_tiocsctty, setsid}, | ||||
|     pty::{grantpt, ioctl_tiocgptpeer, openpt, unlockpt, OpenptFlags}, | ||||
|     stdio::{dup2_stderr, dup2_stdin, dup2_stdout}, | ||||
|     termios::{tcsetwinsize, Winsize}, | ||||
| }; | ||||
| use tokio::{ | ||||
|     io::{unix::AsyncFd, AsyncRead, AsyncWrite}, | ||||
|     process::Child, | ||||
| }; | ||||
| 
 | ||||
| #[derive(Debug)] | ||||
| pub struct Pty { | ||||
|     fd: AsyncFd<OwnedFd>, | ||||
| } | ||||
| 
 | ||||
| impl Pty { | ||||
|     pub fn open() -> io::Result<Self> { | ||||
|         let master = openpt(OpenptFlags::RDWR | OpenptFlags::NOCTTY | OpenptFlags::CLOEXEC)?; | ||||
|         grantpt(&master)?; | ||||
|         unlockpt(&master)?; | ||||
| 
 | ||||
|         // Set nonblocking
 | ||||
|         let flags = fcntl_getfl(&master)?; | ||||
|         fcntl_setfl(&master, flags | OFlags::NONBLOCK)?; | ||||
| 
 | ||||
|         let fd = AsyncFd::new(master)?; | ||||
| 
 | ||||
|         Ok(Self { fd }) | ||||
|     } | ||||
| 
 | ||||
|     pub fn child(&self) -> io::Result<PtyChild> { | ||||
|         // NOTE: Linux v4.13 and above
 | ||||
|         let fd = ioctl_tiocgptpeer(&self.fd, OpenptFlags::RDWR | OpenptFlags::NOCTTY)?; | ||||
|         let child = PtyChild { fd }; | ||||
| 
 | ||||
|         Ok(child) | ||||
|     } | ||||
| 
 | ||||
|     pub fn resize_window(&self, rows: u16, cols: u16) -> io::Result<()> { | ||||
|         let winsize = Winsize { | ||||
|             ws_row: rows, | ||||
|             ws_col: cols, | ||||
|             ws_xpixel: 0, | ||||
|             ws_ypixel: 0, | ||||
|         }; | ||||
| 
 | ||||
|         tcsetwinsize(&self.fd, winsize)?; | ||||
| 
 | ||||
|         Ok(()) | ||||
|     } | ||||
| 
 | ||||
|     pub fn try_clone(&self) -> io::Result<Pty> { | ||||
|         let fd = self.fd.get_ref().try_clone()?; | ||||
| 
 | ||||
|         Ok(Pty { | ||||
|             fd: AsyncFd::new(fd)?, | ||||
|         }) | ||||
|     } | ||||
| } | ||||
| 
 | ||||
| impl AsyncRead for Pty { | ||||
|     fn poll_read( | ||||
|         self: std::pin::Pin<&mut Self>, | ||||
|         cx: &mut std::task::Context<'_>, | ||||
|         buf: &mut tokio::io::ReadBuf<'_>, | ||||
|     ) -> std::task::Poll<io::Result<()>> { | ||||
|         loop { | ||||
|             let mut guard = ready!(self.fd.poll_read_ready(cx)?); | ||||
| 
 | ||||
|             match guard.try_io(|inner| { | ||||
|                 let fd = inner.get_ref(); | ||||
|                 let n = rustix::io::read(fd, buf.initialize_unfilled())?; | ||||
|                 buf.advance(n); | ||||
| 
 | ||||
|                 Ok(()) | ||||
|             }) { | ||||
|                 Ok(result) => return std::task::Poll::Ready(result), | ||||
|                 Err(_would_block) => continue, | ||||
|             } | ||||
|         } | ||||
|     } | ||||
| } | ||||
| 
 | ||||
| impl AsyncWrite for Pty { | ||||
|     fn poll_write( | ||||
|         self: std::pin::Pin<&mut Self>, | ||||
|         cx: &mut std::task::Context<'_>, | ||||
|         buf: &[u8], | ||||
|     ) -> std::task::Poll<Result<usize, io::Error>> { | ||||
|         loop { | ||||
|             let mut guard = ready!(self.fd.poll_write_ready(cx))?; | ||||
| 
 | ||||
|             match guard.try_io(|inner| Ok(rustix::io::write(inner.get_ref(), buf)?)) { | ||||
|                 Ok(result) => return std::task::Poll::Ready(result), | ||||
|                 Err(_would_block) => continue, | ||||
|             } | ||||
|         } | ||||
|     } | ||||
| 
 | ||||
|     fn poll_flush( | ||||
|         self: std::pin::Pin<&mut Self>, | ||||
|         _cx: &mut std::task::Context<'_>, | ||||
|     ) -> std::task::Poll<Result<(), io::Error>> { | ||||
|         std::task::Poll::Ready(Ok(())) | ||||
|     } | ||||
| 
 | ||||
|     fn poll_shutdown( | ||||
|         self: std::pin::Pin<&mut Self>, | ||||
|         _cx: &mut std::task::Context<'_>, | ||||
|     ) -> std::task::Poll<Result<(), io::Error>> { | ||||
|         std::task::Poll::Ready(Ok(())) | ||||
|     } | ||||
| } | ||||
| 
 | ||||
| #[derive(Debug)] | ||||
| pub struct PtyChild { | ||||
|     fd: OwnedFd, | ||||
| } | ||||
| 
 | ||||
| impl PtyChild { | ||||
|     pub fn login_tty(&self) -> io::Result<()> { | ||||
|         setsid()?; | ||||
|         ioctl_tiocsctty(&self.fd)?; | ||||
|         dup2_stdin(&self.fd)?; | ||||
|         dup2_stdout(&self.fd)?; | ||||
|         dup2_stderr(&self.fd)?; | ||||
| 
 | ||||
|         Ok(()) | ||||
|     } | ||||
| } | ||||
| 
 | ||||
| pub fn open_shell(pty_child: PtyChild, shell: &str) -> io::Result<Child> { | ||||
|     let mut cmd = tokio::process::Command::new(shell); | ||||
| 
 | ||||
|     unsafe { | ||||
|         cmd.pre_exec(move || { | ||||
|             pty_child.login_tty()?; | ||||
|             Ok(()) | ||||
|         }); | ||||
|     } | ||||
| 
 | ||||
|     cmd.spawn() | ||||
| } | ||||
| 
 | ||||
| #[cfg(test)] | ||||
| mod test { | ||||
|     use rustix::fd::AsRawFd; | ||||
| 
 | ||||
|     use super::*; | ||||
| 
 | ||||
|     #[tokio::test] | ||||
|     async fn can_open_pty() { | ||||
|         let pty = Pty::open().unwrap(); | ||||
|         let child = pty.child().unwrap(); | ||||
| 
 | ||||
|         let master_fd = pty.fd.get_ref().as_raw_fd(); | ||||
|         let child_fd = child.fd.as_raw_fd(); | ||||
| 
 | ||||
|         assert!(master_fd != child_fd); | ||||
|     } | ||||
| } | ||||
|  | @ -1,89 +0,0 @@ | |||
| use std::process::Stdio; | ||||
| 
 | ||||
| use bytes::Bytes; | ||||
| use futures::{FutureExt, Stream}; | ||||
| use serde::Deserialize; | ||||
| use tokio::process::Command; | ||||
| use tokio_stream::StreamExt; | ||||
| use tokio_util::codec::{BytesCodec, FramedRead}; | ||||
| 
 | ||||
| use crate::{ | ||||
|     health::Health, | ||||
|     messaging::{Client, Message, Subject}, | ||||
| }; | ||||
| 
 | ||||
| pub async fn init_services(client: Client) { | ||||
|     let mut message_stream = client.subscribe().await.unwrap(); | ||||
| 
 | ||||
|     tokio::spawn(async move { | ||||
|         while let Some(message) = message_stream.next().await { | ||||
|             let client = client.clone(); | ||||
|             tokio::spawn(async move { | ||||
|                 if let Err(err) = handle_message(client, message).await { | ||||
|                     tracing::warn!("{err}"); | ||||
|                 } | ||||
|             }); | ||||
|         } | ||||
|     }); | ||||
| } | ||||
| 
 | ||||
| async fn handle_message(client: Client, message: Message) -> anyhow::Result<()> { | ||||
|     match message.subject() { | ||||
|         Subject::Exec => { | ||||
|             let stream = exec_handler(message.parse_payload()?).await?; | ||||
|             client.reply(message, stream).await; | ||||
|         } | ||||
|         Subject::Health => { | ||||
|             let health: Health = message.parse_payload()?; | ||||
|             tracing::info!(?health, "received a health"); | ||||
|         } | ||||
|     } | ||||
| 
 | ||||
|     Ok(()) | ||||
| } | ||||
| 
 | ||||
| /// An operating system program execution.
 | ||||
| #[derive(Debug, Deserialize)] | ||||
| struct ExecMessage { | ||||
|     user: String, | ||||
|     program: String, | ||||
|     #[serde(default)] | ||||
|     args: Vec<String>, | ||||
| } | ||||
| 
 | ||||
| async fn exec_handler(req: ExecMessage) -> anyhow::Result<impl Stream<Item = Bytes> + Unpin> { | ||||
|     // TODO: Tasks should be idempontent
 | ||||
|     // TODO: Root user should be able to run only specific programs
 | ||||
|     let mut cmd = if req.user != "root" { | ||||
|         let mut cmd = Command::new("sudo"); | ||||
|         cmd.arg("-iu").arg(&req.user).arg("--").arg(&req.program); | ||||
|         cmd | ||||
|     } else { | ||||
|         Command::new(&req.program) | ||||
|     }; | ||||
| 
 | ||||
|     let mut io = cmd | ||||
|         .args(&req.args) | ||||
|         .stdout(Stdio::piped()) | ||||
|         .stderr(Stdio::piped()) | ||||
|         .spawn()?; | ||||
| 
 | ||||
|     let stdout = FramedRead::new(io.stdout.take().unwrap(), BytesCodec::new()) | ||||
|         .map(|stdout| Bytes::from(stdout.unwrap())); | ||||
| 
 | ||||
|     let stderr = FramedRead::new(io.stderr.take().unwrap(), BytesCodec::new()) | ||||
|         .map(|stderr| Bytes::from(stderr.unwrap())); | ||||
| 
 | ||||
|     let exit = async move { | ||||
|         io.wait() | ||||
|             .await | ||||
|             .map(|exit| { | ||||
|                 let exit = exit.to_string(); | ||||
|                 Bytes::from(exit) | ||||
|             }) | ||||
|             .unwrap() | ||||
|     } | ||||
|     .into_stream(); | ||||
| 
 | ||||
|     Ok(Box::pin(stdout.merge(stderr).chain(exit))) | ||||
| } | ||||
							
								
								
									
										79
									
								
								agent/src/services/exec.rs
									
									
									
									
									
										Normal file
									
								
							
							
						
						
									
										79
									
								
								agent/src/services/exec.rs
									
									
									
									
									
										Normal file
									
								
							|  | @ -0,0 +1,79 @@ | |||
| use std::process::Stdio; | ||||
| 
 | ||||
| use bytes::Bytes; | ||||
| use futures::FutureExt; | ||||
| use serde::Deserialize; | ||||
| use tokio::process::Command; | ||||
| use tokio_stream::StreamExt; | ||||
| use tokio_util::codec::{BytesCodec, FramedRead}; | ||||
| 
 | ||||
| use super::Ctx; | ||||
| 
 | ||||
| #[derive(Debug, Deserialize)] | ||||
| pub struct ExecReq { | ||||
|     user: String, | ||||
|     program: String, | ||||
|     #[serde(default)] | ||||
|     args: Vec<String>, | ||||
| } | ||||
| 
 | ||||
| impl TryFrom<Bytes> for ExecReq { | ||||
|     type Error = serde_json::Error; | ||||
| 
 | ||||
|     fn try_from(value: Bytes) -> Result<Self, Self::Error> { | ||||
|         serde_json::from_slice(&value[..]) | ||||
|     } | ||||
| } | ||||
| 
 | ||||
| /// An operating system program execution.
 | ||||
| pub async fn exec(ctx: Ctx<ExecReq>) -> anyhow::Result<()> { | ||||
|     // TODO: Tasks should be idempontent
 | ||||
|     // TODO: Root user should be able to run only specific programs
 | ||||
| 
 | ||||
|     let mut cmd = if &ctx.body.user != "root" { | ||||
|         let mut cmd = Command::new("sudo"); | ||||
|         cmd.arg("-iu") | ||||
|             .arg(&ctx.body.user) | ||||
|             .arg("--") | ||||
|             .arg(&ctx.body.program); | ||||
|         cmd | ||||
|     } else { | ||||
|         Command::new(&ctx.body.program) | ||||
|     }; | ||||
| 
 | ||||
|     let mut io = cmd | ||||
|         .args(&ctx.body.args) | ||||
|         .stdout(Stdio::piped()) | ||||
|         .stderr(Stdio::piped()) | ||||
|         .spawn()?; | ||||
| 
 | ||||
|     let stdout = | ||||
|         FramedRead::new(io.stdout.take().unwrap(), BytesCodec::new()).filter_map(|stdout| { | ||||
|             stdout | ||||
|                 .map(|bytes| bytes.freeze()) | ||||
|                 .map_err(|err| tracing::error!(%err, "read error on stdout")) | ||||
|                 .ok() | ||||
|         }); | ||||
| 
 | ||||
|     let stderr = | ||||
|         FramedRead::new(io.stderr.take().unwrap(), BytesCodec::new()).filter_map(|stderr| { | ||||
|             stderr | ||||
|                 .map(|bytes| bytes.freeze()) | ||||
|                 .map_err(|err| tracing::error!(%err, "read error on stderr")) | ||||
|                 .ok() | ||||
|         }); | ||||
| 
 | ||||
|     let exit = async move { | ||||
|         io.wait() | ||||
|             .await | ||||
|             .map(|exit| { | ||||
|                 let exit = exit.to_string(); | ||||
|                 Bytes::from(exit) | ||||
|             }) | ||||
|             .unwrap() | ||||
|     } | ||||
|     .into_stream(); | ||||
| 
 | ||||
|     //     Ok(Box::pin(stdout.merge(stderr).chain(exit)))
 | ||||
|     Ok(()) | ||||
| } | ||||
							
								
								
									
										67
									
								
								agent/src/services/mod.rs
									
									
									
									
									
										Normal file
									
								
							
							
						
						
									
										67
									
								
								agent/src/services/mod.rs
									
									
									
									
									
										Normal file
									
								
							|  | @ -0,0 +1,67 @@ | |||
| use anyhow::Context; | ||||
| use bytes::Bytes; | ||||
| use thiserror::Error; | ||||
| use tokio_stream::StreamExt; | ||||
| 
 | ||||
| mod exec; | ||||
| mod terminal; | ||||
| 
 | ||||
| #[derive(Debug, Error)] | ||||
| enum ServiceError { | ||||
|     #[error("received an invalid body format for a valid message")] | ||||
|     BodyFormatError, | ||||
| } | ||||
| 
 | ||||
| struct Ctx<T> { | ||||
|     body: T, | ||||
| } | ||||
| 
 | ||||
| impl<T> Ctx<T> | ||||
| where | ||||
|     T: TryFrom<Bytes>, | ||||
| { | ||||
|     fn with_body(client: crate::messaging::Client, body: Bytes) -> Result<Self, ServiceError> { | ||||
|         Ok(Self { | ||||
|             body: body | ||||
|                 .try_into() | ||||
|                 .map_err(|_err| ServiceError::BodyFormatError)?, | ||||
|         }) | ||||
|     } | ||||
| } | ||||
| 
 | ||||
| async fn route_message( | ||||
|     client: crate::messaging::Client, | ||||
|     message: crate::messaging::Message, | ||||
| ) -> Result<(), ServiceError> { | ||||
|     match message.subject() { | ||||
|         crate::messaging::Subject::Health => {} | ||||
|         crate::messaging::Subject::Exec => { | ||||
|             let ctx = Ctx::with_body(client, message.body())?; | ||||
|             let _ = self::exec::exec(ctx).await; | ||||
|         } | ||||
|         crate::messaging::Subject::OpenTerminal => { | ||||
|             let ctx = Ctx::with_body(client, message.body())?; | ||||
|             let _ = self::terminal::open_terminal(ctx).await; | ||||
|         } | ||||
|     } | ||||
| 
 | ||||
|     Ok(()) | ||||
| } | ||||
| 
 | ||||
| pub async fn init_services(client: crate::messaging::Client) -> anyhow::Result<()> { | ||||
|     let mut message_stream = client | ||||
|         .subscribe() | ||||
|         .await | ||||
|         .context("could not initialize services system")?; | ||||
| 
 | ||||
|     tokio::spawn(async move { | ||||
|         while let Some(message) = message_stream.next().await { | ||||
|             // TODO: How do i handle this error?
 | ||||
|             if let Err(err) = route_message(client.clone(), message).await { | ||||
|                 tracing::warn!("{}", err); | ||||
|             }; | ||||
|         } | ||||
|     }); | ||||
| 
 | ||||
|     Ok(()) | ||||
| } | ||||
							
								
								
									
										45
									
								
								agent/src/services/terminal.rs
									
									
									
									
									
										Normal file
									
								
							
							
						
						
									
										45
									
								
								agent/src/services/terminal.rs
									
									
									
									
									
										Normal file
									
								
							|  | @ -0,0 +1,45 @@ | |||
| use bytes::Bytes; | ||||
| use futures::Stream; | ||||
| use serde::Deserialize; | ||||
| use tokio::io::AsyncWriteExt; | ||||
| use tokio_stream::StreamExt; | ||||
| use tokio_util::codec::{BytesCodec, FramedRead}; | ||||
| 
 | ||||
| use super::Ctx; | ||||
| 
 | ||||
| #[derive(Debug, Deserialize)] | ||||
| pub struct OpenTerminalMessage { | ||||
|     id: String, | ||||
| } | ||||
| 
 | ||||
| impl TryFrom<Bytes> for OpenTerminalMessage { | ||||
|     type Error = serde_json::Error; | ||||
| 
 | ||||
|     fn try_from(value: Bytes) -> Result<Self, Self::Error> { | ||||
|         serde_json::from_slice(&value[..]) | ||||
|     } | ||||
| } | ||||
| 
 | ||||
| pub async fn open_terminal(ctx: Ctx<OpenTerminalMessage>) -> anyhow::Result<()> { | ||||
|     let pty = crate::pty::Pty::open()?; | ||||
|     let mut pty_clone = pty.try_clone()?; | ||||
| 
 | ||||
|     tokio::spawn(async move { | ||||
|         while let Some(data) = tokio_stream::once(b"foo").next().await { | ||||
|             if let Err(err) = pty_clone.write_all(&data[..]).await { | ||||
|                 tracing::warn!(%err, "pseudoterminal write error"); | ||||
|             } | ||||
|         } | ||||
|     }); | ||||
| 
 | ||||
|     let _out_stream = FramedRead::new(pty, BytesCodec::new()).filter_map(|inner| { | ||||
|         inner | ||||
|             .map(|bytes| bytes.freeze()) | ||||
|             .map_err(|err| { | ||||
|                 tracing::warn!(%err, "pseudoterminal read error"); | ||||
|             }) | ||||
|             .ok() | ||||
|     }); | ||||
| 
 | ||||
|     Ok(()) | ||||
| } | ||||
		Loading…
	
		Reference in a new issue