dosh/src/exec_service.rs

use anyhow::{Context, Result, bail};
use serde::{Deserialize, Serialize};

pub const EXEC_STREAM_SENTINEL: &str = "@dosh-exec";
pub const FRAME_MAX_LEN: usize = 1024 * 1024;

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct ExecRequest {
    pub command: String,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub enum ExecResponse {
    Stdout(Vec<u8>),
    Stderr(Vec<u8>),
    Exit { code: Option<i32> },
    Error { message: String },
}

#[derive(Default)]
pub struct FrameDecoder {
    buf: Vec<u8>,
}

impl FrameDecoder {
    pub fn push(&mut self, bytes: &[u8]) -> Result<Vec<Vec<u8>>> {
        self.buf.extend_from_slice(bytes);
        let mut frames = Vec::new();
        loop {
            if self.buf.len() < 4 {
                break;
            }
            let len = u32::from_be_bytes(self.buf[..4].try_into().unwrap()) as usize;
            if len > FRAME_MAX_LEN {
                self.buf.clear();
                bail!("exec protocol frame too large: {len} bytes");
            }
            if self.buf.len() < 4 + len {
                break;
            }
            frames.push(self.buf[4..4 + len].to_vec());
            self.buf.drain(..4 + len);
        }
        Ok(frames)
    }
}

pub fn encode_request(request: &ExecRequest) -> Result<Vec<u8>> {
    encode_frame(&bincode::serialize(request).context("encode exec request")?)
}

pub fn encode_response(response: &ExecResponse) -> Result<Vec<u8>> {
    encode_frame(&bincode::serialize(response).context("encode exec response")?)
}

pub fn decode_request(frame: &[u8]) -> Result<ExecRequest> {
    bincode::deserialize(frame).context("decode exec request")
}

pub fn decode_response(frame: &[u8]) -> Result<ExecResponse> {
    bincode::deserialize(frame).context("decode exec response")
}

fn encode_frame(payload: &[u8]) -> Result<Vec<u8>> {
    if payload.len() > FRAME_MAX_LEN {
        bail!("exec protocol frame too large: {} bytes", payload.len());
    }
    let mut out = Vec::with_capacity(4 + payload.len());
    out.extend_from_slice(&(payload.len() as u32).to_be_bytes());
    out.extend_from_slice(payload);
    Ok(out)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn oversized_frame_clears_decoder_state() {
        let mut decoder = FrameDecoder::default();
        let too_large = ((FRAME_MAX_LEN + 1) as u32).to_be_bytes();
        assert!(decoder.push(&too_large).is_err());

        let valid = encode_response(&ExecResponse::Exit { code: Some(0) }).unwrap();
        let frames = decoder.push(&valid).unwrap();
        assert_eq!(frames.len(), 1);
        assert_eq!(
            decode_response(&frames[0]).unwrap(),
            ExecResponse::Exit { code: Some(0) }
        );
    }
}