use std::collections::{HashMap, HashSet, VecDeque};
use std::io::{self, Write};
use std::sync::{Arc, Mutex};
use anyhow::Result;
#[cfg(not(miri))]
use anyhow::bail;
#[cfg(not(miri))]
use oxdock_pipe::OsPipeWriter;
use oxdock_pipe::{
KeeperGuard, Materialized, PipeHandle, PipeInfo, PipeInner, SharedInput, SharedOutput,
inspect as inspect_handle, materialize, peek as peek_handle, script_backend,
};
use oxdock_process::{CommandStderr, CommandStdin, CommandStdout};
#[derive(Clone, Default)]
pub(super) struct PipeRegistry;
#[cfg(not(miri))]
fn spent_handle(idx: usize) -> anyhow::Error {
anyhow::anyhow!(
"step {}: OS pipe handle has already been consumed by another binding; declare a fresh LET $x: PIPE for a new session",
idx + 1,
)
}
impl PipeRegistry {
pub(super) fn ensure_handle(handle: &PipeHandle, promote: bool) -> Result<()> {
materialize(handle, promote)?;
Ok(())
}
pub fn peek_pipe_content(handle: &PipeHandle) -> Result<Vec<u8>> {
peek_handle(handle)
}
pub(super) fn inspect_pipe(handle: &PipeHandle) -> PipeInfo {
inspect_handle(handle)
}
pub(super) fn pin_keeper(handle: &PipeHandle) -> Result<Option<KeeperGuard>> {
Ok(script_backend(handle).map(KeeperGuard::new))
}
pub(super) fn resolve_stdin(
idx: usize,
handle: &PipeHandle,
direct: bool,
promote: bool,
) -> Result<(CommandStdin, Option<Arc<PipeInner>>)> {
let _ = idx;
let _ = direct;
match materialize(handle, promote)? {
Materialized::Script(backend) => {
Ok((CommandStdin::Stream(backend.reader_handle()), Some(backend)))
}
#[cfg(not(miri))]
Materialized::Os(entry) => {
if direct {
return Ok((CommandStdin::OsPipe(entry.reader.clone()), None));
}
let owned = entry.reader.take().map_err(|_| spent_handle(idx))?;
Ok((
CommandStdin::Stream(Arc::new(std::sync::Mutex::new(owned))),
None,
))
}
}
}
pub(super) fn resolve_stdout(
idx: usize,
handle: &PipeHandle,
direct: bool,
promote: bool,
) -> Result<(StreamHandle, Option<Arc<PipeInner>>)> {
let _ = idx;
let _ = direct;
match materialize(handle, promote)? {
Materialized::Script(backend) => {
Ok((StreamHandle::Stream(backend.writer_handle()), Some(backend)))
}
#[cfg(not(miri))]
Materialized::Os(entry) => {
if direct {
return Ok((StreamHandle::Os(entry.writer.clone()), None));
}
let owned = entry.writer.take().map_err(|_| spent_handle(idx))?;
Ok((
StreamHandle::Stream(Arc::new(std::sync::Mutex::new(owned))),
None,
))
}
}
}
pub(super) fn resolve_stderr(
idx: usize,
handle: &PipeHandle,
direct: bool,
promote: bool,
) -> Result<StreamHandle> {
let _ = idx;
let _ = direct;
match materialize(handle, promote)? {
Materialized::Script(backend) => Ok(StreamHandle::Stream(backend.writer_handle())),
#[cfg(not(miri))]
Materialized::Os(entry) => {
if direct {
return Ok(StreamHandle::Os(entry.writer.clone()));
}
let owned = entry.writer.take().map_err(|_| spent_handle(idx))?;
Ok(StreamHandle::Stream(Arc::new(std::sync::Mutex::new(owned))))
}
}
}
}
#[derive(Clone, Default)]
pub struct ExecIo {
stdin: Option<SharedInput>,
stdout: Option<SharedOutput>,
stderr: Option<SharedOutput>,
inherit_env_overrides: HashMap<String, String>,
inherit_env_removed: HashSet<String>,
}
pub const CHUNK_SIZE: usize = 8192;
const MIN_RING_CAPACITY: usize = 1024;
pub(crate) struct SlidingWindow {
pub(crate) needle: Vec<u8>,
ring: VecDeque<u8>,
pub matched: bool,
}
impl SlidingWindow {
pub fn new(needle: Vec<u8>) -> Self {
Self {
ring: VecDeque::with_capacity(needle.len().max(MIN_RING_CAPACITY)),
needle,
matched: false,
}
}
pub fn push_chunk(&mut self, chunk: &[u8]) {
if self.matched {
return;
}
let limit = self.needle.len().max(MIN_RING_CAPACITY);
for &byte in chunk {
self.ring.push_back(byte);
if self.ring.len() > limit {
self.ring.pop_front();
}
self.check_match();
}
}
pub fn update_needle(&mut self, new_needle: Vec<u8>) {
if self.matched {
return;
}
self.needle = new_needle;
self.check_match();
}
fn check_match(&mut self) {
if self.matched || self.ring.len() < self.needle.len() {
return;
}
let start = self.ring.len() - self.needle.len();
if self
.ring
.iter()
.skip(start)
.zip(self.needle.iter())
.all(|(a, b)| a == b)
{
self.matched = true;
}
}
pub fn ring_buffer(&self) -> Vec<u8> {
self.ring.iter().copied().collect()
}
}
#[derive(Clone)]
pub(super) enum StreamHandle {
Stream(SharedOutput),
#[cfg(not(miri))]
Os(OsPipeWriter),
}
impl StreamHandle {
pub(super) fn to_stdout(&self) -> CommandStdout {
match self {
StreamHandle::Stream(writer) => CommandStdout::Stream(writer.clone()),
#[cfg(not(miri))]
StreamHandle::Os(writer) => CommandStdout::OsPipe(writer.clone()),
}
}
pub(super) fn to_stderr(&self) -> CommandStderr {
match self {
StreamHandle::Stream(writer) => CommandStderr::Stream(writer.clone()),
#[cfg(not(miri))]
StreamHandle::Os(writer) => CommandStderr::OsPipe(writer.clone()),
}
}
}
pub(super) fn write_stdout<F>(handle: Option<StreamHandle>, op: F) -> Result<()>
where
F: FnOnce(&mut dyn Write) -> Result<()>,
{
match handle {
Some(StreamHandle::Stream(writer)) => {
if let Ok(mut guard) = writer.lock() {
op(&mut *guard)?;
}
Ok(())
}
#[cfg(not(miri))]
Some(StreamHandle::Os(_)) => {
bail!("cannot write DSL output to a live OS pipe")
}
None => {
let mut stdout = io::stdout();
op(&mut stdout)
}
}
}
pub struct PipeStream {
reader: Option<SharedInput>,
writer: Option<SharedOutput>,
}
impl PipeStream {
pub fn reader(reader: SharedInput) -> Self {
Self {
reader: Some(reader),
writer: None,
}
}
pub fn writer(writer: SharedOutput) -> Self {
Self {
reader: None,
writer: Some(writer),
}
}
pub fn pair(reader: SharedInput, writer: SharedOutput) -> Self {
Self {
reader: Some(reader),
writer: Some(writer),
}
}
}
impl std::io::Read for PipeStream {
fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
let Some(reader) = &self.reader else {
return Err(std::io::Error::new(
std::io::ErrorKind::NotConnected,
"pipe stream has no reader half",
));
};
let mut guard = reader
.lock()
.map_err(|_| std::io::Error::other("pipe reader lock poisoned"))?;
guard.read(buf)
}
}
impl std::io::Write for PipeStream {
fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
let Some(writer) = &self.writer else {
return Err(std::io::Error::new(
std::io::ErrorKind::NotConnected,
"pipe stream has no writer half",
));
};
let mut guard = writer
.lock()
.map_err(|_| std::io::Error::other("pipe writer lock poisoned"))?;
guard.write(buf)
}
fn flush(&mut self) -> std::io::Result<()> {
let Some(writer) = &self.writer else {
return Err(std::io::Error::new(
std::io::ErrorKind::NotConnected,
"pipe stream has no writer half",
));
};
let mut guard = writer
.lock()
.map_err(|_| std::io::Error::other("pipe writer lock poisoned"))?;
guard.flush()
}
}
pub(crate) const EXACT_STDOUT_CAP: usize = 8 * 1024 * 1024;
pub(crate) struct ExactCapture {
pub(crate) bytes: Vec<u8>,
pub(crate) overflowed: bool,
}
impl ExactCapture {
pub fn new() -> Self {
Self {
bytes: Vec::new(),
overflowed: false,
}
}
pub fn push_chunk(&mut self, chunk: &[u8]) {
if self.overflowed {
return;
}
if self.bytes.len() + chunk.len() > EXACT_STDOUT_CAP {
self.overflowed = true;
return;
}
self.bytes.extend_from_slice(chunk);
}
}
#[derive(Clone, Copy)]
enum TeeStream {
Stdout,
Stderr,
}
struct TeeWriter {
inner: Option<SharedOutput>,
stream: TeeStream,
windows: Arc<Mutex<HashMap<(usize, usize), SlidingWindow>>>,
exact: Arc<Mutex<HashMap<usize, ExactCapture>>>,
}
impl Write for TeeWriter {
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
match &self.inner {
Some(inner) => {
let mut guard = inner
.lock()
.map_err(|_| io::Error::other("stdout sink poisoned"))?;
guard.write_all(buf)?;
}
None => match self.stream {
TeeStream::Stdout => io::stdout().write_all(buf)?,
TeeStream::Stderr => io::stderr().write_all(buf)?,
},
}
if let Ok(mut windows) = self.windows.lock() {
for window in windows.values_mut() {
window.push_chunk(buf);
}
}
if let Ok(mut exact) = self.exact.lock() {
for capture in exact.values_mut() {
capture.push_chunk(buf);
}
}
Ok(buf.len())
}
fn flush(&mut self) -> io::Result<()> {
match &self.inner {
Some(inner) => {
let mut guard = inner
.lock()
.map_err(|_| io::Error::other("stdout sink poisoned"))?;
guard.flush()?;
}
None => match self.stream {
TeeStream::Stdout => io::stdout().flush()?,
TeeStream::Stderr => io::stderr().flush()?,
},
}
Ok(())
}
}
pub(crate) fn teed_stdout(
sink: Option<SharedOutput>,
windows: Arc<Mutex<HashMap<(usize, usize), SlidingWindow>>>,
exact: Arc<Mutex<HashMap<usize, ExactCapture>>>,
) -> SharedOutput {
Arc::new(Mutex::new(TeeWriter {
inner: sink,
stream: TeeStream::Stdout,
windows,
exact,
}))
}
pub(crate) fn teed_stderr(
sink: Option<SharedOutput>,
windows: Arc<Mutex<HashMap<(usize, usize), SlidingWindow>>>,
) -> SharedOutput {
Arc::new(Mutex::new(TeeWriter {
inner: sink,
stream: TeeStream::Stderr,
windows,
exact: Arc::new(Mutex::new(HashMap::new())),
}))
}
impl ExecIo {
pub fn new() -> Self {
Self::default()
}
pub fn set_stdin(&mut self, stdin: Option<SharedInput>) {
self.stdin = stdin;
}
pub fn set_stdout(&mut self, stdout: Option<SharedOutput>) {
self.stdout = stdout.clone();
if self.stderr.is_none() {
self.stderr = stdout;
}
}
pub fn set_stderr(&mut self, stderr: Option<SharedOutput>) {
self.stderr = stderr;
}
pub fn insert_inherit_env<S: Into<String>, V: Into<String>>(&mut self, key: S, value: V) {
let key = key.into();
self.inherit_env_removed.remove(&key);
self.inherit_env_overrides.insert(key, value.into());
}
pub fn remove_inherit_env<S: Into<String>>(&mut self, key: S) {
let key = key.into();
self.inherit_env_overrides.remove(&key);
self.inherit_env_removed.insert(key);
}
pub fn inherit_env_value(&self, key: &str) -> Option<&String> {
self.inherit_env_overrides.get(key)
}
pub fn inherit_env_is_removed(&self, key: &str) -> bool {
self.inherit_env_removed.contains(key)
}
pub fn inherit_env_overrides(&self) -> &std::collections::HashMap<String, String> {
&self.inherit_env_overrides
}
pub(super) fn ensure_handle(&self, handle: &PipeHandle, promote: bool) -> Result<()> {
PipeRegistry::ensure_handle(handle, promote)
}
pub(super) fn inspect_pipe(&self, handle: &PipeHandle) -> PipeInfo {
PipeRegistry::inspect_pipe(handle)
}
pub fn peek_pipe_content(&self, handle: &PipeHandle) -> Result<Vec<u8>> {
PipeRegistry::peek_pipe_content(handle)
}
pub(super) fn pin_keeper(&self, handle: &PipeHandle) -> Result<Option<KeeperGuard>> {
PipeRegistry::pin_keeper(handle)
}
pub(super) fn resolve_stdin(
&self,
idx: usize,
handle: &PipeHandle,
direct: bool,
promote: bool,
) -> Result<(CommandStdin, Option<Arc<PipeInner>>)> {
PipeRegistry::resolve_stdin(idx, handle, direct, promote)
}
pub(super) fn resolve_stdout(
&self,
idx: usize,
handle: &PipeHandle,
direct: bool,
promote: bool,
) -> Result<(StreamHandle, Option<Arc<PipeInner>>)> {
PipeRegistry::resolve_stdout(idx, handle, direct, promote)
}
pub(super) fn resolve_stderr(
&self,
idx: usize,
handle: &PipeHandle,
direct: bool,
promote: bool,
) -> Result<StreamHandle> {
PipeRegistry::resolve_stderr(idx, handle, direct, promote)
}
pub fn stdin(&self) -> Option<SharedInput> {
self.stdin.clone()
}
pub fn stdout(&self) -> Option<SharedOutput> {
self.stdout.clone()
}
pub fn stderr(&self) -> Option<SharedOutput> {
self.stderr.clone().or_else(|| self.stdout.clone())
}
}
pub(super) fn assemble_default_io(
stdin: Option<SharedInput>,
stdout: Option<SharedOutput>,
) -> ExecIo {
let mut io = ExecIo::new();
io.set_stdin(stdin);
io.set_stdout(stdout.clone());
io.set_stderr(stdout);
io
}