feat: exec tty support

This commit is contained in:
Alex Zenla
2024-04-22 23:02:14 +00:00
parent 284ed8f17b
commit 2c9152d433
10 changed files with 286 additions and 106 deletions

13
Cargo.lock generated
View File

@ -1479,6 +1479,7 @@ dependencies = [
"nix 0.28.0", "nix 0.28.0",
"oci-spec", "oci-spec",
"path-absolutize", "path-absolutize",
"pty-process",
"rtnetlink", "rtnetlink",
"serde", "serde",
"serde_json", "serde_json",
@ -2221,6 +2222,17 @@ dependencies = [
"prost", "prost",
] ]
[[package]]
name = "pty-process"
version = "0.4.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "8749b545e244c90bf74a5767764cc2194f1888bb42f84015486a64c82bea5cc0"
dependencies = [
"libc",
"rustix",
"tokio",
]
[[package]] [[package]]
name = "quote" name = "quote"
version = "1.0.35" version = "1.0.35"
@ -2456,6 +2468,7 @@ checksum = "65e04861e65f21776e67888bfbea442b3642beaa0138fdb1dd7a84a52dffdb89"
dependencies = [ dependencies = [
"bitflags 2.5.0", "bitflags 2.5.0",
"errno", "errno",
"itoa",
"libc", "libc",
"linux-raw-sys", "linux-raw-sys",
"windows-sys 0.52.0", "windows-sys 0.52.0",

View File

@ -91,6 +91,10 @@ features = ["derive"]
version = "0.13.1" version = "0.13.1"
features = ["derive"] features = ["derive"]
[workspace.dependencies.pty-process]
version = "0.4.0"
features = ["async"]
[workspace.dependencies.reqwest] [workspace.dependencies.reqwest]
version = "0.12.4" version = "0.12.4"
default-features = false default-features = false

View File

@ -21,6 +21,8 @@ pub struct ExecCommand {
env: Option<Vec<String>>, env: Option<Vec<String>>,
#[arg(short = 'w', long, help = "Working directory")] #[arg(short = 'w', long, help = "Working directory")]
working_directory: Option<String>, working_directory: Option<String>,
#[arg(short = 't', long, help = "Allocate tty")]
tty: bool,
#[arg(help = "Guest to exec inside, either the name or the uuid")] #[arg(help = "Guest to exec inside, either the name or the uuid")]
guest: String, guest: String,
#[arg( #[arg(
@ -47,14 +49,16 @@ impl ExecCommand {
command: self.command, command: self.command,
working_directory: self.working_directory.unwrap_or_default(), working_directory: self.working_directory.unwrap_or_default(),
}), }),
data: vec![], tty: self.tty,
stdin: vec![],
stdin_closed: false,
}; };
let stream = StdioConsoleStream::stdin_stream_exec(initial).await; let stream = StdioConsoleStream::stdin_stream_exec(initial).await;
let response = client.exec_guest(Request::new(stream)).await?.into_inner(); let response = client.exec_guest(Request::new(stream)).await?.into_inner();
let result = StdioConsoleStream::exec_output(self.tty, response).await;
let code = StdioConsoleStream::exec_output(response).await?; StdioConsoleStream::restore_terminal_mode();
let code = result?;
std::process::exit(code); std::process::exit(code);
} }
} }

View File

@ -68,7 +68,13 @@ impl StdioConsoleStream {
if size == 1 && buffer[0] == 0x1d { if size == 1 && buffer[0] == 0x1d {
break; break;
} }
yield ExecGuestRequest { guest_id: String::default(), task: None, data };
let closed = size == 0;
yield ExecGuestRequest { guest_id: String::default(), task: None, tty: false, stdin: data, stdin_closed: closed };
if closed {
break;
}
} }
} }
} }
@ -90,7 +96,11 @@ impl StdioConsoleStream {
Ok(()) Ok(())
} }
pub async fn exec_output(mut stream: Streaming<ExecGuestReply>) -> Result<i32> { pub async fn exec_output(tty: bool, mut stream: Streaming<ExecGuestReply>) -> Result<i32> {
if tty && stdin().is_tty() {
enable_raw_mode()?;
StdioConsoleStream::register_terminal_restore_hook()?;
}
let mut stdout = stdout(); let mut stdout = stdout();
let mut stderr = stderr(); let mut stderr = stderr();
while let Some(reply) = stream.next().await { while let Some(reply) = stream.next().await {

View File

@ -210,15 +210,28 @@ impl ControlService for DaemonControlService {
.collect(), .collect(),
command: task.command, command: task.command,
working_directory: task.working_directory, working_directory: task.working_directory,
tty: request.tty,
})), })),
})), })),
}; };
let (request_stdin, request_stdin_closed) = (request.stdin.clone(), request.stdin_closed);
let output = try_stream! { let output = try_stream! {
let mut handle = idm.send_stream(idm_request).await.map_err(|x| ApiError { let mut handle = idm.send_stream(idm_request).await.map_err(|x| ApiError {
message: x.to_string(), message: x.to_string(),
})?; })?;
if !request_stdin.is_empty() {
let _ = handle.update(IdmRequest {
request: Some(IdmRequestType::ExecStream(ExecStreamRequestUpdate {
update: Some(Update::Stdin(ExecStreamRequestStdin {
data: request_stdin,
closed: request_stdin_closed,
})),
}))}).await;
}
loop { loop {
select! { select! {
x = input.next() => if let Some(update) = x { x = input.next() => if let Some(update) = x {
@ -227,11 +240,12 @@ impl ControlService for DaemonControlService {
}.into()); }.into());
if let Ok(update) = update { if let Ok(update) = update {
if !update.data.is_empty() { if !update.stdin.is_empty() {
let _ = handle.update(IdmRequest { let _ = handle.update(IdmRequest {
request: Some(IdmRequestType::ExecStream(ExecStreamRequestUpdate { request: Some(IdmRequestType::ExecStream(ExecStreamRequestUpdate {
update: Some(Update::Stdin(ExecStreamRequestStdin { update: Some(Update::Stdin(ExecStreamRequestStdin {
data: update.data, data: update.stdin,
closed: update.stdin_closed,
})), })),
}))}).await; }))}).await;
} }

View File

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

View File

@ -1,4 +1,4 @@
use std::{collections::HashMap, process::Stdio}; use std::{collections::HashMap, process::Stdio, time::Duration};
use anyhow::{anyhow, Result}; use anyhow::{anyhow, Result};
use krata::idm::{ use krata::idm::{
@ -9,10 +9,12 @@ use krata::idm::{
}, },
internal::{response::Response as ResponseType, Request, Response}, internal::{response::Response as ResponseType, Request, Response},
}; };
use pty_process::{Pty, Size};
use tokio::{ use tokio::{
io::{AsyncReadExt, AsyncWriteExt}, io::{AsyncReadExt, AsyncWriteExt},
join, join,
process::Command, process::{Child, Command},
time::sleep,
}; };
pub struct GuestExecTask { pub struct GuestExecTask {
@ -58,26 +60,125 @@ impl GuestExecTask {
start.working_directory.clone() start.working_directory.clone()
}; };
let mut child = Command::new(exe) 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 (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;
};
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 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 !update.data.is_empty() && write.write_all(&update.data).await.is_err() {
break;
}
if update.closed {
break;
}
}
});
let mut result = child.inner.wait().await;
if result.is_err() {
sleep(Duration::from_millis(10)).await;
if let Ok(Some(status)) = child.inner.try_wait() {
result = Ok(status);
}
}
let code = result.as_ref().ok().and_then(|x| x.code()).unwrap_or(-1);
let error = result
.as_ref()
.map_err(|x| x.to_string())
.err()
.unwrap_or_default();
let _ = pty_read_task.await;
stdin_task.abort();
let _ = stdin_task.await;
let response = Response {
response: Some(ResponseType::ExecStream(ExecStreamResponseUpdate {
exited: true,
exit_code: code,
error,
stdout: vec![],
stderr: vec![],
})),
};
self.handle.respond(response).await?;
child.kill = false;
} else {
let mut child = ChildDropGuard {
inner: Command::new(exe)
.args(cmd) .args(cmd)
.envs(env) .envs(env)
.current_dir(dir) .current_dir(dir)
.stdin(Stdio::piped()) .stdin(Stdio::piped())
.stdout(Stdio::piped()) .stdout(Stdio::piped())
.stderr(Stdio::piped()) .stderr(Stdio::piped())
.kill_on_drop(true)
.spawn() .spawn()
.map_err(|error| anyhow!("failed to spawn: {}", error))?; .map_err(|error| anyhow!("failed to spawn: {}", error))?,
kill: true,
};
let mut stdin = child let mut stdin = child
.inner
.stdin .stdin
.take() .take()
.ok_or_else(|| anyhow!("stdin was missing"))?; .ok_or_else(|| anyhow!("stdin was missing"))?;
let mut stdout = child let mut stdout = child
.inner
.stdout .stdout
.take() .take()
.ok_or_else(|| anyhow!("stdout was missing"))?; .ok_or_else(|| anyhow!("stdout was missing"))?;
let mut stderr = child let mut stderr = child
.inner
.stderr .stderr
.take() .take()
.ok_or_else(|| anyhow!("stderr was missing"))?; .ok_or_else(|| anyhow!("stderr was missing"))?;
@ -144,14 +245,29 @@ impl GuestExecTask {
continue; continue;
}; };
if stdin.write_all(&update.data).await.is_err() { if !update.data.is_empty() && stdin.write_all(&update.data).await.is_err() {
break;
}
if update.closed {
break; break;
} }
} }
}); });
let exit = child.wait().await?; let mut result = child.inner.wait().await;
let code = exit.code().unwrap_or(-1); if result.is_err() {
sleep(Duration::from_millis(10)).await;
if let Ok(Some(status)) = child.inner.try_wait() {
result = Ok(status);
}
}
let code = result.as_ref().ok().and_then(|x| x.code()).unwrap_or(-1);
let error = result
.as_ref()
.map_err(|x| x.to_string())
.err()
.unwrap_or_default();
let _ = join!(stdout_task, stderr_task); let _ = join!(stdout_task, stderr_task);
stdin_task.abort(); stdin_task.abort();
@ -160,13 +276,27 @@ impl GuestExecTask {
response: Some(ResponseType::ExecStream(ExecStreamResponseUpdate { response: Some(ResponseType::ExecStream(ExecStreamResponseUpdate {
exited: true, exited: true,
exit_code: code, exit_code: code,
error: String::new(), error,
stdout: vec![], stdout: vec![],
stderr: vec![], stderr: vec![],
})), })),
}; };
self.handle.respond(response).await?; self.handle.respond(response).await?;
child.kill = false;
}
Ok(()) 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());
}
}
}

View File

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

View File

@ -67,7 +67,9 @@ message ListGuestsReply {
message ExecGuestRequest { message ExecGuestRequest {
string guest_id = 1; string guest_id = 1;
krata.v1.common.GuestTaskSpec task = 2; krata.v1.common.GuestTaskSpec task = 2;
bytes data = 3; bytes stdin = 3;
bool stdin_closed = 4;
bool tty = 5;
} }
message ExecGuestReply { message ExecGuestReply {

View File

@ -28,5 +28,5 @@ build_and_run() {
fi fi
RUST_TARGET="$(./hack/build/target.sh)" RUST_TARGET="$(./hack/build/target.sh)"
./hack/build/cargo.sh build ${CARGO_BUILD_FLAGS} --bin "${EXE_TARGET}" ./hack/build/cargo.sh build ${CARGO_BUILD_FLAGS} --bin "${EXE_TARGET}"
exec sudo sh -c "RUST_LOG='${RUST_LOG}' 'target/${RUST_TARGET}/debug/${EXE_TARGET}' $*" exec sudo RUST_LOG="${RUST_LOG}" "target/${RUST_TARGET}/debug/${EXE_TARGET}" "${@}"
} }