use bytes::{Bytes, BytesMut};
use futures_util::stream::Stream;
use http_body_util::BodyExt;
use hyper::body::Incoming;
use std::pin::Pin;
use crate::libpod::error::PodmanError;
#[derive(Debug)]
pub enum LogOutput {
StdOut { message: Bytes },
StdErr { message: Bytes },
}
pub type BoxStream<T> = Pin<Box<dyn Stream<Item = Result<T, PodmanError>> + Send>>;
const MAX_STREAM_BUF: usize = 256 * 1024 * 1024;
fn stream_buf_overflow() -> PodmanError {
PodmanError::Api {
status: 0,
message: format!("stream chunk exceeds the {MAX_STREAM_BUF} byte limit"),
}
}
pub fn parse_frame(buf: &mut BytesMut) -> Option<(u8, Bytes)> {
if buf.len() < 8 {
return None;
}
let size = u32::from_be_bytes([buf[4], buf[5], buf[6], buf[7]]) as usize;
if buf.len() < 8 + size {
return None;
}
let stream_type = buf[0];
let mut frame = buf.split_to(8 + size);
let payload = frame.split_off(8).freeze();
Some((stream_type, payload))
}
pub fn take_json_line(buf: &mut BytesMut) -> Option<Bytes> {
let nl = buf.iter().position(|&b| b == b'\n')?;
let mut line = buf.split_to(nl + 1);
line.truncate(nl); Some(line.freeze())
}
pub fn parse_multiplexed(body: Incoming) -> BoxStream<LogOutput> {
Box::pin(futures_util::stream::try_unfold(
(body, BytesMut::new()),
|(mut body, mut buf)| async move {
loop {
if let Some((stream_type, payload)) = parse_frame(&mut buf) {
let output = match stream_type {
1 => LogOutput::StdOut { message: payload },
2 => LogOutput::StdErr { message: payload },
_ => continue, };
return Ok(Some((output, (body, buf))));
}
match body.frame().await {
Some(Ok(frame)) => {
if let Ok(data) = frame.into_data() {
buf.extend_from_slice(&data);
if buf.len() > MAX_STREAM_BUF {
return Err(stream_buf_overflow());
}
}
}
Some(Err(e)) => return Err(PodmanError::from(e)),
None => return Ok(None),
}
}
},
))
}
pub fn parse_raw(body: Incoming) -> BoxStream<LogOutput> {
Box::pin(futures_util::stream::try_unfold(
body,
|mut body| async move {
loop {
match body.frame().await {
Some(Ok(frame)) => {
if let Ok(data) = frame.into_data() {
if !data.is_empty() {
return Ok(Some((LogOutput::StdOut { message: data }, body)));
}
}
}
Some(Err(e)) => return Err(PodmanError::from(e)),
None => return Ok(None),
}
}
},
))
}
pub fn parse_json_lines<T: serde::de::DeserializeOwned + Send + 'static>(
body: Incoming,
) -> BoxStream<T> {
Box::pin(futures_util::stream::try_unfold(
(body, BytesMut::new()),
|(mut body, mut buf)| async move {
loop {
if let Some(line) = take_json_line(&mut buf) {
if line.is_empty() {
continue;
}
let item: T = serde_json::from_slice(&line).map_err(PodmanError::Json)?;
return Ok(Some((item, (body, buf))));
}
match body.frame().await {
Some(Ok(frame)) => {
if let Ok(data) = frame.into_data() {
buf.extend_from_slice(&data);
if buf.len() > MAX_STREAM_BUF {
return Err(stream_buf_overflow());
}
}
}
Some(Err(e)) => return Err(PodmanError::from(e)),
None => {
let line = std::mem::take(&mut buf);
if !line.is_empty() {
let item: T =
serde_json::from_slice(&line).map_err(PodmanError::Json)?;
return Ok(Some((item, (body, buf))));
}
return Ok(None);
}
}
}
},
))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn parse_frame_incomplete_header() {
let mut buf = BytesMut::from(&[0x01, 0x00, 0x00, 0x00][..]);
assert!(parse_frame(&mut buf).is_none());
assert_eq!(buf.as_ref(), &[0x01, 0x00, 0x00, 0x00]);
}
#[test]
fn parse_frame_header_present_payload_missing() {
let mut buf = BytesMut::from(
&[
0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x05, b'a', b'b', b'c',
][..],
);
assert!(parse_frame(&mut buf).is_none());
assert_eq!(buf.len(), 11);
}
#[test]
fn parse_frame_stdout_complete() {
let mut buf = BytesMut::from(&[0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x05][..]);
buf.extend_from_slice(b"hello");
let (stype, data) = parse_frame(&mut buf).unwrap();
assert_eq!(stype, 1);
assert_eq!(data.as_ref(), b"hello");
assert!(buf.is_empty());
}
#[test]
fn parse_frame_stderr_complete() {
let mut buf = BytesMut::from(&[0x02, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x03][..]);
buf.extend_from_slice(b"err");
let (stype, data) = parse_frame(&mut buf).unwrap();
assert_eq!(stype, 2);
assert_eq!(data.as_ref(), b"err");
assert!(buf.is_empty());
}
#[test]
fn parse_frame_zero_length_payload() {
let mut buf = BytesMut::from(&[0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00][..]);
let (stype, data) = parse_frame(&mut buf).unwrap();
assert_eq!(stype, 1);
assert!(data.is_empty());
assert!(buf.is_empty());
}
#[test]
fn parse_frame_leaves_remainder() {
let mut buf =
BytesMut::from(&[0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x02, b'h', b'i'][..]);
buf.extend_from_slice(b"leftover");
let (_, data) = parse_frame(&mut buf).unwrap();
assert_eq!(data.as_ref(), b"hi");
assert_eq!(buf.as_ref(), b"leftover");
}
#[test]
fn parse_frame_two_frames_in_one_buffer_demux() {
let mut buf =
BytesMut::from(&[0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x02, b'h', b'i'][..]);
buf.extend_from_slice(&[0x02, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x02, b'e', b'r']);
let (stype1, data1) = parse_frame(&mut buf).unwrap();
assert_eq!(stype1, 1);
assert_eq!(data1.as_ref(), b"hi");
let (stype2, data2) = parse_frame(&mut buf).unwrap();
assert_eq!(stype2, 2);
assert_eq!(data2.as_ref(), b"er");
assert!(buf.is_empty());
assert!(parse_frame(&mut buf).is_none());
}
#[test]
fn parse_frame_split_across_reads_reassembles() {
let mut buf = BytesMut::from(&[0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x05][..]);
buf.extend_from_slice(b"hel");
assert!(parse_frame(&mut buf).is_none());
buf.extend_from_slice(b"lo");
let (stype, data) = parse_frame(&mut buf).unwrap();
assert_eq!(stype, 1);
assert_eq!(data.as_ref(), b"hello");
assert!(buf.is_empty());
}
#[test]
fn take_json_line_no_newline() {
let mut buf = BytesMut::from(&b"partial line"[..]);
assert!(take_json_line(&mut buf).is_none());
assert_eq!(buf.as_ref(), b"partial line");
}
#[test]
fn take_json_line_with_newline() {
let mut buf = BytesMut::from(&b"line1\nline2"[..]);
let line = take_json_line(&mut buf).unwrap();
assert_eq!(line.as_ref(), b"line1");
assert_eq!(buf.as_ref(), b"line2");
}
#[test]
fn take_json_line_empty_line() {
let mut buf = BytesMut::from(&b"\nnext"[..]);
let line = take_json_line(&mut buf).unwrap();
assert!(line.is_empty());
assert_eq!(buf.as_ref(), b"next");
}
#[test]
fn take_json_line_multiple_lines() {
let mut buf = BytesMut::from(&b"a\nb\nc"[..]);
assert_eq!(take_json_line(&mut buf).unwrap().as_ref(), b"a");
assert_eq!(take_json_line(&mut buf).unwrap().as_ref(), b"b");
assert!(take_json_line(&mut buf).is_none());
}
#[test]
fn take_json_line_multiple_lines_in_one_buffer_in_order() {
let mut buf = BytesMut::from(
&br#"{"a":1}
{"b":2}
{"c":3}
"#[..],
);
assert_eq!(take_json_line(&mut buf).unwrap().as_ref(), br#"{"a":1}"#);
assert_eq!(take_json_line(&mut buf).unwrap().as_ref(), br#"{"b":2}"#);
assert_eq!(take_json_line(&mut buf).unwrap().as_ref(), br#"{"c":3}"#);
assert!(take_json_line(&mut buf).is_none());
assert!(buf.is_empty());
}
#[test]
fn take_json_line_split_across_reads_reassembles() {
let mut buf = BytesMut::from(&b"{\"a\":"[..]);
assert!(take_json_line(&mut buf).is_none());
buf.extend_from_slice(b"1}\n");
assert_eq!(take_json_line(&mut buf).unwrap().as_ref(), br#"{"a":1}"#);
assert!(buf.is_empty());
}
fn cap_check(buf_len: usize) -> Option<PodmanError> {
if buf_len > MAX_STREAM_BUF {
Some(stream_buf_overflow())
} else {
None
}
}
#[test]
fn over_cap_buffer_is_rejected() {
match cap_check(MAX_STREAM_BUF + 1) {
Some(PodmanError::Api { status, message }) => {
assert_eq!(status, 0);
assert!(message.contains(&MAX_STREAM_BUF.to_string()));
}
other => panic!("expected Api overflow error, got {other:?}"),
}
}
#[test]
fn at_cap_buffer_is_accepted() {
assert!(cap_check(MAX_STREAM_BUF).is_none());
}
}