use std::io::{self, Read, Seek, SeekFrom, Write};
use std::process::{Command, Stdio};
use tempfile::NamedTempFile;
use crate::reducers::{RecoveryMarker, StreamReducer};
use crate::rules::{PackRequest, Registry, SelectedRule, Selection};
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct ExpectedSelection {
pub pack_id: String,
pub rule_id: String,
pub source_digest: [u8; 32],
}
struct CapturedStream {
output: CapturedOutput,
raw: Option<RawSpool>,
archive_error: Option<String>,
raw_emitted: bool,
}
enum CapturedOutput {
Reduced(Box<StreamReducer>),
Passthrough,
}
struct RawSpool {
file: NamedTempFile,
bytes_written: usize,
#[cfg(test)]
fail_after: Option<usize>,
}
struct SpoolWriteError {
source: io::Error,
bytes_written: usize,
}
impl RawSpool {
fn new() -> io::Result<Self> {
Ok(Self {
file: NamedTempFile::new()?,
bytes_written: 0,
#[cfg(test)]
fail_after: None,
})
}
#[cfg(test)]
fn failing_after(limit: usize) -> io::Result<Self> {
let mut spool = Self::new()?;
spool.fail_after = Some(limit);
Ok(spool)
}
fn write_chunk(&mut self, body: &[u8]) -> Result<(), SpoolWriteError> {
#[cfg(test)]
let permitted = self.fail_after.map_or(body.len(), |limit| {
limit.saturating_sub(self.bytes_written).min(body.len())
});
#[cfg(not(test))]
let permitted = body.len();
let mut offset = 0;
while offset < permitted {
match self.file.write(&body[offset..permitted]) {
Ok(0) => {
return Err(SpoolWriteError {
source: io::Error::new(
io::ErrorKind::WriteZero,
"archive spool wrote zero bytes",
),
bytes_written: offset,
});
}
Ok(count) => {
offset += count;
self.bytes_written = self.bytes_written.saturating_add(count);
}
Err(source) => {
return Err(SpoolWriteError {
source,
bytes_written: offset,
});
}
}
}
if permitted != body.len() {
return Err(SpoolWriteError {
source: io::Error::other("simulated archive spool failure"),
bytes_written: offset,
});
}
Ok(())
}
fn as_file_mut(&mut self) -> &mut std::fs::File {
self.file.as_file_mut()
}
fn prepared_path(&mut self) -> Result<std::path::PathBuf, String> {
self.file
.flush()
.map_err(|error| format!("could not flush archive spool: {error}"))?;
Ok(self.file.path().to_path_buf())
}
}
pub fn run(
arguments: &[String],
archive_key: Option<&crate::archive::ArchiveKey>,
) -> Result<i32, String> {
run_with_rules(arguments, archive_key, &[], None)
}
pub fn run_with_rules(
arguments: &[String],
archive_key: Option<&crate::archive::ArchiveKey>,
packs: &[PackRequest],
expected: Option<&ExpectedSelection>,
) -> Result<i32, String> {
if arguments.is_empty() {
return Err("child command is missing".to_owned());
}
let selected = select_rule(arguments, packs, expected)?;
let passthrough = selected.is_none();
let (stdout_spool, stderr_spool) = if archive_key.is_some() {
(
Some(
RawSpool::new()
.map_err(|error| format!("could not create stdout archive spool: {error}"))?,
),
Some(
RawSpool::new()
.map_err(|error| format!("could not create stderr archive spool: {error}"))?,
),
)
} else {
(None, None)
};
let stdout_output = selected
.as_ref()
.map(|selected| StreamReducer::new(&selected.rule))
.transpose()?;
let stderr_output = selected
.as_ref()
.map(|selected| StreamReducer::new(&selected.rule))
.transpose()?;
let mut child = Command::new(&arguments[0])
.args(&arguments[1..])
.stdin(Stdio::inherit())
.stdout(Stdio::piped())
.stderr(Stdio::piped())
.spawn()
.map_err(|error| format!("could not start {}: {error}", arguments[0]))?;
let stdout = child
.stdout
.take()
.ok_or_else(|| "could not capture child stdout".to_owned())?;
let stderr = child
.stderr
.take()
.ok_or_else(|| "could not capture child stderr".to_owned())?;
let stdout_thread = std::thread::spawn(move || {
capture(
stdout,
stdout_spool,
io::stdout(),
stdout_output.map_or(CapturedOutput::Passthrough, |reducer| {
CapturedOutput::Reduced(Box::new(reducer))
}),
)
});
let stderr_thread = std::thread::spawn(move || {
capture(
stderr,
stderr_spool,
io::stderr(),
stderr_output.map_or(CapturedOutput::Passthrough, |reducer| {
CapturedOutput::Reduced(Box::new(reducer))
}),
)
});
let status = child
.wait()
.map_err(|error| format!("could not wait for child: {error}"))?;
finish_run(
status,
stdout_thread,
stderr_thread,
archive_key,
passthrough,
)
}
type CaptureThread = std::thread::JoinHandle<io::Result<CapturedStream>>;
fn finish_run(
status: std::process::ExitStatus,
stdout_thread: CaptureThread,
stderr_thread: CaptureThread,
archive_key: Option<&crate::archive::ArchiveKey>,
passthrough: bool,
) -> Result<i32, String> {
let mut stdout = join_capture(stdout_thread, "stdout")?;
let mut stderr = join_capture(stderr_thread, "stderr")?;
let succeeded = status.success();
let archive_ref = if passthrough {
None
} else if let Some(key) = archive_key {
match crate::archive::Archive::open_read_only().and_then(|archive| archive.archive_ref(key))
{
Ok(value) => Some(value),
Err(error) => {
restore_after_capture_error(&mut stdout, &mut stderr)?;
eprintln!("yarp: archive failed after command execution: {error}");
return Ok(exit_code(status));
}
}
} else {
None
};
let stdout_after = finish_output(
std::mem::replace(&mut stdout.output, CapturedOutput::Passthrough),
succeeded,
archive_ref.as_deref().map(|archive_ref| RecoveryMarker {
archive_ref,
source: "stdout",
completeness: "complete",
}),
);
let stderr_after = finish_output(
std::mem::replace(&mut stderr.output, CapturedOutput::Passthrough),
succeeded,
archive_ref.as_deref().map(|archive_ref| RecoveryMarker {
archive_ref,
source: "stderr",
completeness: "complete",
}),
);
let capture_errors = capture_errors(&stdout, &stderr);
if !capture_errors.is_empty() {
restore_after_capture_error(&mut stdout, &mut stderr)?;
eprintln!(
"yarp: archive failed after command execution: {}",
capture_errors.join("; ")
);
return Ok(exit_code(status));
}
if let Some(key) = archive_key
&& let Err(error) = archive_captures(
key,
&mut stdout,
&mut stderr,
stdout_after.as_deref(),
stderr_after.as_deref(),
passthrough,
)
{
restore_after_capture_error(&mut stdout, &mut stderr)?;
eprintln!("yarp: archive failed after command execution: {error}");
return Ok(exit_code(status));
}
if passthrough {
restore_after_capture_error(&mut stdout, &mut stderr)?;
}
if let Some(stdout_after) = stdout_after {
io::stdout()
.write_all(&stdout_after)
.map_err(|error| format!("could not write stdout: {error}"))?;
}
if let Some(stderr_after) = stderr_after {
io::stderr()
.write_all(&stderr_after)
.map_err(|error| format!("could not write stderr: {error}"))?;
}
Ok(exit_code(status))
}
fn capture_errors(stdout: &CapturedStream, stderr: &CapturedStream) -> Vec<String> {
[
stdout
.archive_error
.as_ref()
.map(|error| format!("stdout: {error}")),
stderr
.archive_error
.as_ref()
.map(|error| format!("stderr: {error}")),
]
.into_iter()
.flatten()
.collect()
}
fn restore_after_capture_error(
stdout: &mut CapturedStream,
stderr: &mut CapturedStream,
) -> Result<(), String> {
if !stdout.raw_emitted {
emit_raw(stdout, &mut io::stdout(), "stdout")?;
stdout.raw_emitted = true;
}
if !stderr.raw_emitted {
emit_raw(stderr, &mut io::stderr(), "stderr")?;
stderr.raw_emitted = true;
}
Ok(())
}
fn archive_captures(
key: &crate::archive::ArchiveKey,
stdout: &mut CapturedStream,
stderr: &mut CapturedStream,
stdout_after: Option<&[u8]>,
stderr_after: Option<&[u8]>,
passthrough: bool,
) -> Result<(), String> {
let stdout_path = stdout
.raw
.as_mut()
.ok_or_else(|| "stdout archive spool is missing".to_owned())?
.prepared_path()?;
let stderr_path = stderr
.raw
.as_mut()
.ok_or_else(|| "stderr archive spool is missing".to_owned())?
.prepared_path()?;
let captured_at_ms = unix_time_ms();
let operation = if passthrough {
crate::archive_protocol::ArchiveOperation::CapturePassthroughStreams {
request_id: 1,
schema_version: crate::archive_protocol::INGEST_SCHEMA_VERSION,
session: key.session.clone(),
source_call_id: key.source_call_id.clone(),
captured_at_ms,
stdout_path,
stderr_path,
}
} else {
let stdout_after = spool_bytes(stdout_after.unwrap_or_default(), "stdout")?;
let stderr_after = spool_bytes(stderr_after.unwrap_or_default(), "stderr")?;
let operation = crate::archive_protocol::ArchiveOperation::CaptureStreams {
request_id: 1,
schema_version: crate::archive_protocol::INGEST_SCHEMA_VERSION,
session: key.session.clone(),
source_call_id: key.source_call_id.clone(),
captured_at_ms,
stdout_before_path: stdout_path,
stderr_before_path: stderr_path,
stdout_after_path: stdout_after.path().to_path_buf(),
stderr_after_path: stderr_after.path().to_path_buf(),
};
return crate::archive_client::execute(operation).map(|_| ());
};
crate::archive_client::execute(operation).map(|_| ())
}
fn spool_bytes(body: &[u8], label: &str) -> Result<NamedTempFile, String> {
let mut file = NamedTempFile::new()
.map_err(|error| format!("could not create {label} archive spool: {error}"))?;
file.write_all(body)
.and_then(|()| file.flush())
.map_err(|error| format!("could not write {label} archive spool: {error}"))?;
Ok(file)
}
fn select_rule(
arguments: &[String],
packs: &[PackRequest],
expected: Option<&ExpectedSelection>,
) -> Result<Option<SelectedRule>, String> {
let mut registry = match Registry::load(packs) {
Ok(registry) => registry,
Err(_) if expected.is_some() => return Ok(None),
Err(error) => return Err(error),
};
let selection = match registry.select(arguments) {
Ok(selection) => selection,
Err(_) if expected.is_some() => return Ok(None),
Err(error) => return Err(error),
};
match (selection, expected) {
(Selection::Reduce(selected), Some(expected))
if selected.pack_id == expected.pack_id
&& selected.rule.id == expected.rule_id
&& selected.source_digest == expected.source_digest =>
{
Ok(Some(selected))
}
(Selection::Reduce(selected), None) => Ok(Some(selected)),
(_, Some(_)) => Ok(None),
(Selection::Unsupported, None) => Err("command is not on the YARP allowlist".to_owned()),
(Selection::Transform(_), None) => {
Err("pipeline transforms cannot run as standalone YARP rules".to_owned())
}
(Selection::Passthrough(_), None) => {
Err("command is protected by a YARP pass-through rule".to_owned())
}
(Selection::Ambiguous(ids), None) => Err(format!(
"command matches multiple YARP reduction rules: {}",
ids.join(", ")
)),
}
}
fn capture(
mut reader: impl Read,
mut raw: Option<RawSpool>,
mut fallback: impl Write,
mut output: CapturedOutput,
) -> io::Result<CapturedStream> {
let passthrough = matches!(output, CapturedOutput::Passthrough);
let mut archive_error = None;
let mut raw_emitted = passthrough && raw.is_none();
let mut buffer = [0_u8; 8 * 1024];
loop {
let count = reader.read(&mut buffer)?;
if count == 0 {
break;
}
let chunk = &buffer[..count];
if passthrough && raw_emitted {
fallback.write_all(chunk)?;
}
if !passthrough && raw_emitted {
fallback.write_all(chunk)?;
} else if let Some(spool) = &mut raw
&& let Err(error) = spool.write_chunk(chunk)
{
if !raw_emitted {
rewind_and_copy(spool.as_file_mut(), &mut fallback)?;
fallback.write_all(&chunk[error.bytes_written..])?;
}
archive_error = Some(format!("could not write archive spool: {}", error.source));
raw = None;
raw_emitted = true;
}
if let CapturedOutput::Reduced(reducer) = &mut output {
reducer.push(chunk);
}
}
if let Some(spool) = &mut raw
&& let Err(error) = spool.file.flush()
{
if !raw_emitted {
rewind_and_copy(spool.as_file_mut(), &mut fallback)?;
}
archive_error = Some(format!("could not flush archive spool: {error}"));
raw = None;
raw_emitted = true;
}
if raw_emitted {
fallback.flush()?;
}
Ok(CapturedStream {
output,
raw,
archive_error,
raw_emitted,
})
}
fn finish_output(
output: CapturedOutput,
success: bool,
recovery: Option<RecoveryMarker<'_>>,
) -> Option<Vec<u8>> {
match output {
CapturedOutput::Reduced(reducer) => Some((*reducer).finish(success, recovery)),
CapturedOutput::Passthrough => None,
}
}
fn join_capture(
thread: std::thread::JoinHandle<io::Result<CapturedStream>>,
stream: &str,
) -> Result<CapturedStream, String> {
thread
.join()
.map_err(|_| format!("{stream} capture thread panicked"))?
.map_err(|error| format!("could not read child {stream}: {error}"))
}
fn rewind_and_copy(file: &mut std::fs::File, output: &mut impl Write) -> io::Result<()> {
file.seek(SeekFrom::Start(0))?;
io::copy(file, output)?;
Ok(())
}
fn copy_raw(file: &mut std::fs::File, output: &mut impl Write, stream: &str) -> Result<(), String> {
rewind_and_copy(file, output)
.map_err(|error| format!("could not restore raw {stream}: {error}"))
}
fn emit_raw(
captured: &mut CapturedStream,
output: &mut impl Write,
stream: &str,
) -> Result<(), String> {
let spool = captured
.raw
.as_mut()
.ok_or_else(|| format!("raw {stream} archive spool is missing"))?;
copy_raw(spool.as_file_mut(), output, stream)?;
output
.flush()
.map_err(|error| format!("could not flush raw {stream}: {error}"))
}
fn unix_time_ms() -> i64 {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map_or(0, |duration| {
i64::try_from(duration.as_millis()).unwrap_or(i64::MAX)
})
}
fn exit_code(status: std::process::ExitStatus) -> i32 {
if let Some(code) = status.code() {
return code;
}
#[cfg(unix)]
{
use std::os::unix::process::ExitStatusExt;
status.signal().map_or(1, |signal| 128 + signal)
}
#[cfg(not(unix))]
{
1
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::io::Cursor;
fn selected_rule(arguments: &[&str]) -> SelectedRule {
let mut registry = Registry::builtins_only();
let arguments = arguments
.iter()
.map(ToString::to_string)
.collect::<Vec<_>>();
let Selection::Reduce(selected) = registry.select(&arguments).expect("selection") else {
panic!("expected reduction rule");
};
selected
}
#[test]
fn capture_reduces_with_the_selected_rule() {
let selected = selected_rule(&["cargo", "test"]);
let input = " Compiling crate\n".repeat(400);
let reducer = StreamReducer::new(&selected.rule).expect("reducer");
let captured = capture(
Cursor::new(input.as_bytes()),
None,
io::sink(),
CapturedOutput::Reduced(Box::new(reducer)),
)
.expect("capture");
let output = finish_output(captured.output, true, None).expect("reduced output");
assert!(output.len() < input.len());
}
#[test]
fn passthrough_capture_emits_exact_bytes() {
let input = b"exact\xffbytes\n";
let mut output = Vec::new();
let captured = capture(
Cursor::new(input),
None,
&mut output,
CapturedOutput::Passthrough,
)
.expect("capture");
assert_eq!(output, input);
assert!(finish_output(captured.output, true, None).is_none());
}
#[test]
fn archived_passthrough_waits_for_capture_acknowledgement() {
let input = b"deferred passthrough\n";
let spool = RawSpool::new().expect("spool");
let mut output = Vec::new();
let mut captured = capture(
Cursor::new(input),
Some(spool),
&mut output,
CapturedOutput::Passthrough,
)
.expect("capture");
assert!(output.is_empty());
assert!(!captured.raw_emitted);
emit_raw(&mut captured, &mut output, "stdout").expect("emit after acknowledgement");
assert_eq!(output, input);
}
#[test]
fn archive_spool_failure_emits_the_exact_raw_stream() {
let input = b"raw output that exceeds the simulated limit\n";
let selected = selected_rule(&["git", "status"]);
let spool = RawSpool::failing_after(7).expect("spool");
let mut output = Vec::new();
let captured = capture(
Cursor::new(input),
Some(spool),
&mut output,
CapturedOutput::Reduced(Box::new(
StreamReducer::new(&selected.rule).expect("reducer"),
)),
)
.expect("capture");
assert_eq!(output, input);
assert!(captured.raw_emitted);
assert!(captured.raw.is_none());
assert!(
captured
.archive_error
.is_some_and(|error| error.contains("archive spool"))
);
}
#[test]
fn rejects_unselected_direct_children() {
let result = run(&["cat".to_owned(), ".env".to_owned()], None);
assert_eq!(
result,
Err("pipeline transforms cannot run as standalone YARP rules".to_owned())
);
let result = run(&["echo".to_owned(), "secret".to_owned()], None);
assert_eq!(
result,
Err("command is not on the YARP allowlist".to_owned())
);
}
#[test]
fn rewrite_disagreement_fails_open() {
let expected = ExpectedSelection {
pack_id: "yarp-builtins".to_owned(),
rule_id: "git/status".to_owned(),
source_digest: [0_u8; 32],
};
let selected = select_rule(
&["git".to_owned(), "status".to_owned()],
&[],
Some(&expected),
)
.expect("selection");
assert!(selected.is_none());
}
}