mirrors/krata
1
0
Fork 0
mirror of https://github.com/edera-dev/krata.git synced 2025-08-03 13:11:31 +00:00
Code Issues Packages Projects Releases Wiki Activity

feature(exec): implement tty support (fixes #335) (#336)

Browse Source
This commit is contained in:
Alex Zenla
2024-08-14 12:45:59 -07:00
committed by GitHub
parent 87530edf70
commit bf3b73bf24
13 changed files with 350 additions and 166 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
crates/ctl/src/cli/zone/exec.rs
View File

@ -21,6 +21,8 @@ pub struct ZoneExecCommand {
env: Option<Vec<String>>,
#[arg(short = 'w', long, help = "Working directory")]
working_directory: Option<String>,
#[arg(short = 't', long, help = "Allocate tty")]
tty: bool,
#[arg(help = "Zone to exec inside, either the name or the uuid")]
zone: String,
#[arg(
@ -46,8 +48,10 @@ impl ZoneExecCommand {
.collect(),
command: self.command,
working_directory: self.working_directory.unwrap_or_default(),
tty: self.tty,
}),
data: vec![],
stdin: vec![],
stdin_closed: false,
};
let stream = StdioConsoleStream::stdin_stream_exec(initial).await;
@ -57,7 +61,7 @@ impl ZoneExecCommand {
.await?
.into_inner();
let code = StdioConsoleStream::exec_output(response).await?;
let code = StdioConsoleStream::exec_output(response, self.tty).await?;
std::process::exit(code);
}
}

3
crates/ctl/src/cli/zone/launch.rs
View File

@ -59,6 +59,8 @@ pub struct ZoneLaunchCommand {
device: Vec<String>,
#[arg[short, long, help = "Environment variables set in the zone"]]
env: Option<Vec<String>>,
#[arg(short = 't', long, help = "Allocate tty for task")]
tty: bool,
#[arg(
short,
long,
@ -143,6 +145,7 @@ impl ZoneLaunchCommand {
.collect(),
command: self.command,
working_directory: self.working_directory.unwrap_or_default(),
tty: self.tty,
}),
annotations: vec![],
devices: self

2
crates/ctl/src/cli/zone/top.rs
View File

@ -112,7 +112,7 @@ impl ZoneTopApp {
}
fn render_frame(&mut self, frame: &mut Frame) {
frame.render_widget(self, frame.size());
frame.render_widget(self, frame.area());
}
fn handle_event(&mut self, event: Event) -> io::Result<()> {

14
crates/ctl/src/console.rs
View File

@ -62,11 +62,15 @@ impl StdioConsoleStream {
break;
}
};
let data = buffer[0..size].to_vec();
let stdin = buffer[0..size].to_vec();
if size == 1 && buffer[0] == 0x1d {
break;
}
yield ExecInsideZoneRequest { zone_id: String::default(), task: None, data };
let stdin_closed = size == 0;
yield ExecInsideZoneRequest { zone_id: String::default(), task: None, stdin, stdin_closed, };
if stdin_closed {
break;
}
}
}
}
@ -88,7 +92,11 @@ impl StdioConsoleStream {
Ok(())
}
pub async fn exec_output(mut stream: Streaming<ExecInsideZoneReply>) -> Result<i32> {
pub async fn exec_output(mut stream: Streaming<ExecInsideZoneReply>, raw: bool) -> Result<i32> {
if raw && stdin().is_tty() {
enable_raw_mode()?;
StdioConsoleStream::register_terminal_restore_hook()?;
}
let mut stdout = stdout();
let mut stderr = stderr();
while let Some(reply) = stream.next().await {

8
crates/daemon/src/control.rs
View File

@ -226,6 +226,7 @@ impl ControlService for DaemonControlService {
.collect(),
command: task.command,
working_directory: task.working_directory,
tty: task.tty,
})),
})),
};
@ -243,11 +244,12 @@ impl ControlService for DaemonControlService {
}.into());
if let Ok(update) = update {
if !update.data.is_empty() {
if !update.stdin.is_empty() {
let _ = handle.update(IdmRequest {
request: Some(IdmRequestType::ExecStream(ExecStreamRequestUpdate {
update: Some(Update::Stdin(ExecStreamRequestStdin {
data: update.data,
data: update.stdin,
closed: update.stdin_closed,
})),
}))}).await;
}
@ -263,7 +265,7 @@ impl ControlService for DaemonControlService {
error: update.error,
exit_code: update.exit_code,
stdout: update.stdout,
stderr: update.stderr
stderr: update.stderr,
};
yield reply;
},

2
crates/krata/proto/krata/idm/internal.proto
View File

@ -45,10 +45,12 @@ message ExecStreamRequestStart {
repeated ExecEnvVar environment = 1;
repeated string command = 2;
string working_directory = 3;
bool tty = 4;
}
message ExecStreamRequestStdin {
bytes data = 1;
bool closed = 2;
}
message ExecStreamRequestUpdate {

1
crates/krata/proto/krata/v1/common.proto
View File

@ -56,6 +56,7 @@ message ZoneTaskSpec {
repeated ZoneTaskSpecEnvVar environment = 1;
repeated string command = 2;
string working_directory = 3;
bool tty = 4;
}
message ZoneTaskSpecEnvVar {

3
crates/krata/proto/krata/v1/control.proto
View File

@ -84,7 +84,8 @@ message ListZonesReply {
message ExecInsideZoneRequest {
string zone_id = 1;
krata.v1.common.ZoneTaskSpec task = 2;
bytes data = 3;
bytes stdin = 3;
bool stdin_closed = 4;
}
message ExecInsideZoneReply {

1
crates/zone/Cargo.toml
View File

@ -22,6 +22,7 @@ nix = { workspace = true, features = ["ioctl", "process", "fs"] }
oci-spec = { workspace = true }
path-absolutize = { workspace = true }
platform-info = { workspace = true }
pty-process = { workspace = true, features = ["async"] }
rtnetlink = { workspace = true }
serde = { workspace = true }
serde_json = { workspace = true }

301
crates/zone/src/exec.rs
View File

@ -1,12 +1,6 @@
use std::{collections::HashMap, process::Stdio};
use anyhow::{anyhow, Result};
use tokio::{
io::{AsyncReadExt, AsyncWriteExt},
join,
process::Command,
};
use krata::idm::{
client::IdmClientStreamResponseHandle,
internal::{
@ -15,6 +9,14 @@ use krata::idm::{
},
internal::{response::Response as ResponseType, Request, Response},
};
use libc::c_int;
use pty_process::{Pty, Size};
use tokio::process::Child;
use tokio::{
io::{AsyncReadExt, AsyncWriteExt},
join,
process::Command,
};
use crate::childwait::ChildWait;
@ -52,7 +54,7 @@ impl ZoneExecTask {
if !env.contains_key("PATH") {
env.insert(
"PATH".to_string(),
"/bin:/usr/bin:/usr/local/bin".to_string(),
"/bin:/usr/bin:/usr/local/bin:/sbin:/usr/sbin".to_string(),
);
}
@ -63,112 +65,196 @@ impl ZoneExecTask {
};
let mut wait_subscription = self.wait.subscribe().await?;
let mut child = Command::new(exe)
.args(cmd)
.envs(env)
.current_dir(dir)
.stdin(Stdio::piped())
.stdout(Stdio::piped())
.stderr(Stdio::piped())
.kill_on_drop(true)
.spawn()
.map_err(|error| anyhow!("failed to spawn: {}", error))?;
let pid = child.id().ok_or_else(|| anyhow!("pid is not provided"))?;
let mut stdin = child
.stdin
.take()
.ok_or_else(|| anyhow!("stdin was missing"))?;
let mut stdout = child
.stdout
.take()
.ok_or_else(|| anyhow!("stdout was missing"))?;
let mut stderr = child
.stderr
.take()
.ok_or_else(|| anyhow!("stderr was missing"))?;
let stdout_handle = self.handle.clone();
let stdout_task = tokio::task::spawn(async move {
let mut stdout_buffer = vec![0u8; 8 * 1024];
loop {
let Ok(size) = stdout.read(&mut stdout_buffer).await else {
break;
};
if size > 0 {
let response = Response {
response: Some(ResponseType::ExecStream(ExecStreamResponseUpdate {
exited: false,
exit_code: 0,
error: String::new(),
stdout: stdout_buffer[0..size].to_vec(),
stderr: vec![],
})),
let code: c_int;
if start.tty {
let pty = Pty::new().map_err(|error| anyhow!("unable to allocate pty: {}", error))?;
pty.resize(Size::new(24, 80))?;
let mut child = ChildDropGuard {
inner: pty_process::Command::new(exe)
.args(cmd)
.envs(env)
.current_dir(dir)
.spawn(
&pty.pts()
.map_err(|error| anyhow!("unable to allocate pts: {}", error))?,
)
.map_err(|error| anyhow!("failed to spawn: {}", error))?,
kill: true,
};
let pid = child
.inner
.id()
.ok_or_else(|| anyhow!("pid is not provided"))?;
let (mut read, mut write) = pty.into_split();
let pty_read_handle = self.handle.clone();
let pty_read_task = tokio::task::spawn(async move {
let mut stdout_buffer = vec![0u8; 8 * 1024];
loop {
let Ok(size) = read.read(&mut stdout_buffer).await else {
break;
};
let _ = stdout_handle.respond(response).await;
} else {
break;
if size > 0 {
let response = Response {
response: Some(ResponseType::ExecStream(ExecStreamResponseUpdate {
exited: false,
exit_code: 0,
error: String::new(),
stdout: stdout_buffer[0..size].to_vec(),
stderr: vec![],
})),
};
let _ = pty_read_handle.respond(response).await;
} else {
break;
}
}
}
});
});
let stderr_handle = self.handle.clone();
let stderr_task = tokio::task::spawn(async move {
let mut stderr_buffer = vec![0u8; 8 * 1024];
loop {
let Ok(size) = stderr.read(&mut stderr_buffer).await else {
break;
};
if size > 0 {
let response = Response {
response: Some(ResponseType::ExecStream(ExecStreamResponseUpdate {
exited: false,
exit_code: 0,
error: String::new(),
stdout: vec![],
stderr: stderr_buffer[0..size].to_vec(),
})),
let stdin_task = tokio::task::spawn(async move {
loop {
let Some(request) = receiver.recv().await else {
break;
};
let _ = stderr_handle.respond(response).await;
} else {
break;
let Some(RequestType::ExecStream(update)) = request.request else {
continue;
};
let Some(Update::Stdin(update)) = update.update else {
continue;
};
if !update.data.is_empty() && write.write_all(&update.data).await.is_err() {
break;
}
if update.closed {
break;
}
}
}
});
});
let stdin_task = tokio::task::spawn(async move {
loop {
let Some(request) = receiver.recv().await else {
break;
};
let Some(RequestType::ExecStream(update)) = request.request else {
continue;
};
let Some(Update::Stdin(update)) = update.update else {
continue;
};
if stdin.write_all(&update.data).await.is_err() {
break;
code = loop {
if let Ok(event) = wait_subscription.recv().await {
if event.pid.as_raw() as u32 == pid {
break event.status;
}
}
}
});
};
let data_task = tokio::task::spawn(async move {
let _ = join!(stdout_task, stderr_task);
child.kill = false;
let _ = join!(pty_read_task);
stdin_task.abort();
});
} else {
let mut child = Command::new(exe)
.args(cmd)
.envs(env)
.current_dir(dir)
.stdin(Stdio::piped())
.stdout(Stdio::piped())
.stderr(Stdio::piped())
.kill_on_drop(true)
.spawn()
.map_err(|error| anyhow!("failed to spawn: {}", error))?;
let code = loop {
if let Ok(event) = wait_subscription.recv().await {
if event.pid.as_raw() as u32 == pid {
break event.status;
let pid = child.id().ok_or_else(|| anyhow!("pid is not provided"))?;
let mut stdin = child
.stdin
.take()
.ok_or_else(|| anyhow!("stdin was missing"))?;
let mut stdout = child
.stdout
.take()
.ok_or_else(|| anyhow!("stdout was missing"))?;
let mut stderr = child
.stderr
.take()
.ok_or_else(|| anyhow!("stderr was missing"))?;
let stdout_handle = self.handle.clone();
let stdout_task = tokio::task::spawn(async move {
let mut stdout_buffer = vec![0u8; 8 * 1024];
loop {
let Ok(size) = stdout.read(&mut stdout_buffer).await else {
break;
};
if size > 0 {
let response = Response {
response: Some(ResponseType::ExecStream(ExecStreamResponseUpdate {
exited: false,
exit_code: 0,
error: String::new(),
stdout: stdout_buffer[0..size].to_vec(),
stderr: vec![],
})),
};
let _ = stdout_handle.respond(response).await;
} else {
break;
}
}
}
};
data_task.await?;
});
let stderr_handle = self.handle.clone();
let stderr_task = tokio::task::spawn(async move {
let mut stderr_buffer = vec![0u8; 8 * 1024];
loop {
let Ok(size) = stderr.read(&mut stderr_buffer).await else {
break;
};
if size > 0 {
let response = Response {
response: Some(ResponseType::ExecStream(ExecStreamResponseUpdate {
exited: false,
exit_code: 0,
error: String::new(),
stdout: vec![],
stderr: stderr_buffer[0..size].to_vec(),
})),
};
let _ = stderr_handle.respond(response).await;
} else {
break;
}
}
});
let stdin_task = tokio::task::spawn(async move {
loop {
let Some(request) = receiver.recv().await else {
break;
};
let Some(RequestType::ExecStream(update)) = request.request else {
continue;
};
let Some(Update::Stdin(update)) = update.update else {
continue;
};
if stdin.write_all(&update.data).await.is_err() {
break;
}
}
});
let data_task = tokio::task::spawn(async move {
let _ = join!(stdout_task, stderr_task);
stdin_task.abort();
});
code = loop {
if let Ok(event) = wait_subscription.recv().await {
if event.pid.as_raw() as u32 == pid {
break event.status;
}
}
};
data_task.await?;
}
let response = Response {
response: Some(ResponseType::ExecStream(ExecStreamResponseUpdate {
exited: true,
@ -183,3 +269,16 @@ impl ZoneExecTask {
Ok(())
}
}
struct ChildDropGuard {
pub inner: Child,
pub kill: bool,
}
impl Drop for ChildDropGuard {
fn drop(&mut self) {
if self.kill {
drop(self.inner.start_kill());
}
}
}

10
crates/zone/src/metrics.rs
View File

@ -2,7 +2,7 @@ use std::{ops::Add, path::Path};
use anyhow::Result;
use krata::idm::internal::{MetricFormat, MetricNode};
use sysinfo::Process;
use sysinfo::{Process, ProcessesToUpdate};
pub struct MetricsCollector {}
@ -38,7 +38,7 @@ impl MetricsCollector {
}
fn collect_processes(&self, sysinfo: &mut sysinfo::System) -> Result<MetricNode> {
sysinfo.refresh_processes();
sysinfo.refresh_processes(ProcessesToUpdate::All);
let mut processes = Vec::new();
let mut sysinfo_processes = sysinfo.processes().values().collect::<Vec<_>>();
sysinfo_processes.sort_by_key(|x| x.pid());
@ -70,7 +70,11 @@ impl MetricsCollector {
metrics.push(MetricNode::raw_value("cwd", working_directory));
}
let cmdline = process.cmd().to_vec();
let cmdline = process
.cmd()
.iter()
.map(|x| x.to_string_lossy().to_string())
.collect::<Vec<_>>();
metrics.push(MetricNode::raw_value("cmdline", cmdline));
metrics.push(MetricNode::structural(
"memory",