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" | source = "registry+https://github.com/rust-lang/crates.io-index" | ||||||
| checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a" | checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a" | ||||||
| 
 | 
 | ||||||
|  | [[package]] | ||||||
|  | name = "bitflags" | ||||||
|  | version = "2.4.2" | ||||||
|  | source = "registry+https://github.com/rust-lang/crates.io-index" | ||||||
|  | checksum = "ed570934406eb16438a4e976b1b4500774099c13b8cb96eec99f620f05090ddf" | ||||||
|  | 
 | ||||||
| [[package]] | [[package]] | ||||||
| name = "block-buffer" | name = "block-buffer" | ||||||
| version = "0.10.4" | version = "0.10.4" | ||||||
|  | @ -281,6 +287,16 @@ dependencies = [ | ||||||
|  "subtle", |  "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]] | [[package]] | ||||||
| name = "fiat-crypto" | name = "fiat-crypto" | ||||||
| version = "0.2.5" | version = "0.2.5" | ||||||
|  | @ -435,9 +451,15 @@ checksum = "e2abad23fbc42b3700f2f279844dc832adb2b2eb069b2df918f455c4e18cc646" | ||||||
| 
 | 
 | ||||||
| [[package]] | [[package]] | ||||||
| name = "libc" | name = "libc" | ||||||
| version = "0.2.151" | version = "0.2.152" | ||||||
| source = "registry+https://github.com/rust-lang/crates.io-index" | 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]] | [[package]] | ||||||
| name = "lock_api" | name = "lock_api" | ||||||
|  | @ -687,6 +709,7 @@ dependencies = [ | ||||||
|  "bytes", |  "bytes", | ||||||
|  "chrono", |  "chrono", | ||||||
|  "futures", |  "futures", | ||||||
|  |  "rustix", | ||||||
|  "serde", |  "serde", | ||||||
|  "serde_json", |  "serde_json", | ||||||
|  "sysinfo", |  "sysinfo", | ||||||
|  | @ -743,7 +766,7 @@ version = "0.4.1" | ||||||
| source = "registry+https://github.com/rust-lang/crates.io-index" | source = "registry+https://github.com/rust-lang/crates.io-index" | ||||||
| checksum = "4722d768eff46b75989dd134e5c353f0d6296e5aaa3132e776cbdb56be7731aa" | checksum = "4722d768eff46b75989dd134e5c353f0d6296e5aaa3132e776cbdb56be7731aa" | ||||||
| dependencies = [ | dependencies = [ | ||||||
|  "bitflags", |  "bitflags 1.3.2", | ||||||
| ] | ] | ||||||
| 
 | 
 | ||||||
| [[package]] | [[package]] | ||||||
|  | @ -804,6 +827,20 @@ dependencies = [ | ||||||
|  "semver", |  "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]] | [[package]] | ||||||
| name = "rustls" | name = "rustls" | ||||||
| version = "0.21.10" | version = "0.21.10" | ||||||
|  | @ -884,7 +921,7 @@ version = "2.9.2" | ||||||
| source = "registry+https://github.com/rust-lang/crates.io-index" | source = "registry+https://github.com/rust-lang/crates.io-index" | ||||||
| checksum = "05b64fb303737d99b81884b2c63433e9ae28abebe5eb5045dcdd175dc2ecf4de" | checksum = "05b64fb303737d99b81884b2c63433e9ae28abebe5eb5045dcdd175dc2ecf4de" | ||||||
| dependencies = [ | dependencies = [ | ||||||
|  "bitflags", |  "bitflags 1.3.2", | ||||||
|  "core-foundation", |  "core-foundation", | ||||||
|  "core-foundation-sys", |  "core-foundation-sys", | ||||||
|  "libc", |  "libc", | ||||||
|  |  | ||||||
|  | @ -9,6 +9,7 @@ async-nats = "0.33.0" | ||||||
| bytes = "1.5.0" | bytes = "1.5.0" | ||||||
| chrono = { version = "0.4.33", default-features = false, features = ["now", "serde"] } | chrono = { version = "0.4.33", default-features = false, features = ["now", "serde"] } | ||||||
| futures = { version = "0.3.30", default-features = false, features = ["std"] } | 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 = { version = "1.0.195", features = ["derive"] } | ||||||
| serde_json = "1.0.111" | serde_json = "1.0.111" | ||||||
| sysinfo = { version = "0.30.5", default-features = false } | sysinfo = { version = "0.30.5", default-features = false } | ||||||
|  |  | ||||||
|  | @ -1,6 +1,6 @@ | ||||||
| //! System health information and checking
 | //! System health information and checking
 | ||||||
| 
 | 
 | ||||||
| use std::{sync::Arc, time::Duration}; | use std::{collections::HashMap, sync::Arc, time::Duration}; | ||||||
| 
 | 
 | ||||||
| use serde::{Deserialize, Serialize}; | use serde::{Deserialize, Serialize}; | ||||||
| use tokio::sync::watch; | use tokio::sync::watch; | ||||||
|  | @ -9,7 +9,7 @@ use crate::messaging::{Client, Message}; | ||||||
| 
 | 
 | ||||||
| const MEMORY_USAGE_CRITICAL_THRESHOLD: f64 = 90.0; | const MEMORY_USAGE_CRITICAL_THRESHOLD: f64 = 90.0; | ||||||
| const CPU_USAGE_CRITICAL_THRESHOLD: f32 = 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 { | pub struct System { | ||||||
|     sys: sysinfo::System, |     sys: sysinfo::System, | ||||||
|  | @ -33,7 +33,7 @@ impl System { | ||||||
|                 .with_cpu(CpuRefreshKind::everything()), |                 .with_cpu(CpuRefreshKind::everything()), | ||||||
|         ); |         ); | ||||||
| 
 | 
 | ||||||
|         // self.disks.refresh_list();
 |         self.disks.refresh_list(); | ||||||
|     } |     } | ||||||
| 
 | 
 | ||||||
|     pub fn system(&self) -> &sysinfo::System { |     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 { | pub enum Status { | ||||||
|     #[default] |     #[default] | ||||||
|     Normal, |     Normal, | ||||||
|     Critical, |     Critical, | ||||||
| } | } | ||||||
| 
 | 
 | ||||||
|  | #[derive(Clone, Debug, Default, Serialize, Deserialize)] | ||||||
|  | pub enum DiskStatus { | ||||||
|  |     #[default] | ||||||
|  |     Normal, | ||||||
|  |     // HighUsage,
 | ||||||
|  |     VeryHighUsage, | ||||||
|  | } | ||||||
|  | 
 | ||||||
| #[derive(Clone, Debug, Default, Serialize, Deserialize)] | #[derive(Clone, Debug, Default, Serialize, Deserialize)] | ||||||
| pub struct Health { | pub struct Health { | ||||||
|     cpu_status: Status, |     cpu_status: Status, | ||||||
|     memory_status: Status, |     memory_status: Status, | ||||||
|     disk_status: Status, |     disk_status: HashMap<String, DiskStatus>, | ||||||
| } | } | ||||||
| 
 | 
 | ||||||
| #[derive(Clone)] | #[derive(Clone)] | ||||||
|  | @ -80,13 +88,12 @@ impl HealthMonitor { | ||||||
| 
 | 
 | ||||||
|         let cpu_usage = sys.global_cpu_info().cpu_usage(); |         let cpu_usage = sys.global_cpu_info().cpu_usage(); | ||||||
| 
 | 
 | ||||||
|         // for d in system.disks().list() {
 |         let disks_usage = system.disks().list().iter().map(|dk| { | ||||||
|         //     let _avail = if d.total_space() > 0 {
 |             ( | ||||||
|         //         (d.available_space() * 100 / d.total_space()) as u8
 |                 dk.name().to_str().unwrap_or("<INVALID DISK NAME DETECTED>"), | ||||||
|         //     } else {
 |                 (dk.total_space() - dk.available_space()) as f64 / dk.total_space() as f64 * 100.0, | ||||||
|         //         0 as u8
 |             ) | ||||||
|         //     };
 |         }); | ||||||
|         // }
 |  | ||||||
| 
 | 
 | ||||||
|         self.0.send_if_modified(|health| { |         self.0.send_if_modified(|health| { | ||||||
|             let cpu_changed = match health.cpu_status { |             let cpu_changed = match health.cpu_status { | ||||||
|  | @ -113,7 +120,42 @@ impl HealthMonitor { | ||||||
|                 _ => false, |                 _ => 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
 |     // Forever refresh system resources and monitor changes
 | ||||||
|     std::thread::spawn(move || loop { |     std::thread::spawn(move || loop { | ||||||
|         const REFRESH_INTERVAL: Duration = Duration::from_secs(1); |         const REFRESH_INTERVAL: Duration = Duration::from_secs(5); | ||||||
|         system.refresh_resources(); |         system.refresh_resources(); | ||||||
|         health_monitor.check_system(&system); |         health_monitor.check_system(&system); | ||||||
|         std::thread::sleep(REFRESH_INTERVAL); |         std::thread::sleep(REFRESH_INTERVAL); | ||||||
|  |  | ||||||
|  | @ -5,6 +5,7 @@ use tracing::Level; | ||||||
| 
 | 
 | ||||||
| mod health; | mod health; | ||||||
| mod messaging; | mod messaging; | ||||||
|  | mod pty; | ||||||
| mod services; | mod services; | ||||||
| 
 | 
 | ||||||
| #[tokio::main] | #[tokio::main] | ||||||
|  | @ -28,7 +29,7 @@ async fn run() -> anyhow::Result<()> { | ||||||
|     let _health_monitor = init_health_subsystem(client.clone()).await; |     let _health_monitor = init_health_subsystem(client.clone()).await; | ||||||
|     tracing::info!("initialized health system"); |     tracing::info!("initialized health system"); | ||||||
| 
 | 
 | ||||||
|     init_services(client).await; |     init_services(client).await?; | ||||||
|     tracing::info!("initialized services"); |     tracing::info!("initialized services"); | ||||||
| 
 | 
 | ||||||
|     tracing::info!("agent is ready"); |     tracing::info!("agent is ready"); | ||||||
|  |  | ||||||
|  | @ -1,9 +1,8 @@ | ||||||
| use std::fmt::Display; | use std::fmt::{Debug, Display}; | ||||||
| 
 | 
 | ||||||
| use anyhow::{anyhow, Result}; | use anyhow::{anyhow, Result}; | ||||||
| use bytes::Bytes; | use bytes::Bytes; | ||||||
| use futures::Stream; | use futures::Stream; | ||||||
| use serde::Deserialize; |  | ||||||
| use tokio_stream::StreamExt; | use tokio_stream::StreamExt; | ||||||
| 
 | 
 | ||||||
| use crate::health::Health; | use crate::health::Health; | ||||||
|  | @ -18,6 +17,7 @@ pub struct Client { | ||||||
| pub enum Subject { | pub enum Subject { | ||||||
|     Health, |     Health, | ||||||
|     Exec, |     Exec, | ||||||
|  |     OpenTerminal, | ||||||
| } | } | ||||||
| 
 | 
 | ||||||
| impl Display for Subject { | impl Display for Subject { | ||||||
|  | @ -25,10 +25,24 @@ impl Display for Subject { | ||||||
|         match self { |         match self { | ||||||
|             Subject::Health => write!(f, "health"), |             Subject::Health => write!(f, "health"), | ||||||
|             Subject::Exec => write!(f, "exec"), |             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)] | #[derive(Debug)] | ||||||
| pub struct Message { | pub struct Message { | ||||||
|     subject: Subject, |     subject: Subject, | ||||||
|  | @ -40,27 +54,19 @@ impl Message { | ||||||
|     fn from_transport(msg: async_nats::Message) -> Result<Self> { |     fn from_transport(msg: async_nats::Message) -> Result<Self> { | ||||||
|         let suffix = msg.subject.split_terminator('.').last().unwrap_or_default(); |         let suffix = msg.subject.split_terminator('.').last().unwrap_or_default(); | ||||||
| 
 | 
 | ||||||
|         match suffix { |         Ok(Message { | ||||||
|             "exec" => Ok(Message { |             subject: suffix.try_into()?, | ||||||
|                 subject: Subject::Exec, |  | ||||||
|             payload: msg.payload, |             payload: msg.payload, | ||||||
|             reply: msg.reply, |             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 { |     pub fn subject(&self) -> &Subject { | ||||||
|         &self.subject |         &self.subject | ||||||
|     } |     } | ||||||
| 
 | 
 | ||||||
|     pub fn parse_payload<'a, T: Deserialize<'a>>(&'a self) -> Result<T> { |     pub fn body(&self) -> Bytes { | ||||||
|         Ok(serde_json::from_slice(&self.payload[..])?) |         self.payload.clone() | ||||||
|     } |     } | ||||||
| 
 | 
 | ||||||
|     pub fn health(health: Health) -> Result<Message> { |     pub fn health(health: Health) -> Result<Message> { | ||||||
|  | @ -125,3 +131,9 @@ impl Client { | ||||||
|         Ok(subscriber) |         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