use base64::Engine as _;
use protocol::exec::{
EVENT_EXIT, EVENT_GAP, EVENT_OUTPUT, ExitEvent, GapEvent, OutputEvent, StreamKind,
};
use super::sse::{ExecEvent, Frame, MAX_PENDING_BYTES, SseParser, decode};
use crate::error::ErrorKind;
const LINE_ENDINGS: [&str; 3] = ["\n", "\r\n", "\r"];
enum Piece {
Event {
wire: Vec<u8>,
event: serde_json::Value,
},
BadBase64 { wire: Vec<u8>, offset: u64 },
Unknown { wire: Vec<u8> },
WrongShape { wire: Vec<u8> },
Fields { wire: Vec<u8>, frame: Option<Frame> },
Keepalive { wire: Vec<u8> },
Junk { wire: Vec<u8> },
}
impl Piece {
fn wire(&self) -> &[u8] {
match self {
Piece::Event { wire, .. }
| Piece::BadBase64 { wire, .. }
| Piece::Unknown { wire }
| Piece::WrongShape { wire }
| Piece::Fields { wire, .. }
| Piece::Keepalive { wire }
| Piece::Junk { wire } => wire,
}
}
}
fn one_line(bytes: &[u8]) -> String {
String::from_utf8_lossy(bytes).replace(['\r', '\n'], " ")
}
fn frame_wire(event: &str, data: &str, ending: &str) -> Vec<u8> {
format!("event: {event}{ending}data: {data}{ending}{ending}").into_bytes()
}
#[derive(Clone, Copy)]
enum Field {
U64,
OptionI32,
Bool,
Text,
Stream,
}
fn wrong(field: Field, pick: u64) -> &'static str {
let choices: &[&str] = match field {
Field::U64 => &["-1", "1.5", "18446744073709551616", "\"7\"", "null", "[1]"],
Field::OptionI32 => &["2147483648", "-2147483649", "1.5", "\"7\"", "[1]"],
Field::Bool => &["1", "\"true\"", "null", "[]"],
Field::Text => &["7", "null", "[\"a\"]", "{}"],
Field::Stream => &["\"stdin\"", "7", "null", "\"STDOUT\""],
};
choices[usize::try_from(pick % choices.len() as u64).expect("a small index")]
}
fn right(field: Field) -> &'static str {
match field {
Field::U64 => "7",
Field::OptionI32 => "null",
Field::Bool => "true",
Field::Text => "\"\"",
Field::Stream => "\"stdout\"",
}
}
fn wrong_shape(number: u64, bits: u64) -> (&'static str, String) {
let (event, fields): (&str, &[(&str, Field)]) = match number % 3 {
0 => (
EVENT_OUTPUT,
&[
("offset", Field::U64),
("stream", Field::Stream),
("output", Field::Text),
],
),
1 => (EVENT_GAP, &[("from", Field::U64), ("to", Field::U64)]),
_ => (
EVENT_EXIT,
&[
("exit_code", Field::OptionI32),
("signal", Field::OptionI32),
("timed_out", Field::Bool),
("truncated", Field::Bool),
("writers_may_be_alive", Field::Bool),
("offset", Field::U64),
],
),
};
let forced = usize::try_from(number / 3 % fields.len() as u64).expect("a small index");
let mut members = Vec::new();
for (at, (name, field)) in fields.iter().enumerate() {
let shift = 3 * u32::try_from(at).expect("a small index");
let value = if at == forced || bits >> shift & 1 == 1 {
wrong(*field, bits >> (shift + 1) ^ number)
} else if bits >> (shift + 2) & 1 == 1 {
continue;
} else {
right(*field)
};
members.push(format!("\"{name}\":{value}"));
}
(event, format!("{{{}}}", members.join(",")))
}
const FIELD_NAMES: [&str; 4] = ["event", "data", "id", "retry"];
fn chosen_lines(bytes: &[u8]) -> Vec<String> {
bytes
.split(|byte| *byte == 0)
.map(|segment| {
let (pick, rest) = segment.split_first().map_or((0, &[][..]), |(p, r)| (*p, r));
let value = one_line(rest);
let name = FIELD_NAMES
.get(usize::from(pick % 5))
.map_or_else(|| value.replace(':', "_"), |name| (*name).to_string());
let line = match pick / 5 % 3 {
0 => format!("{name}: {value}"),
1 => format!("{name}:{value}"),
_ => name,
};
if line.is_empty() {
":".to_string()
} else {
line
}
})
.collect()
}
fn frame_of(lines: &[String]) -> Option<Frame> {
let mut event = String::new();
let mut data = Vec::new();
for line in lines {
if line.starts_with(':') {
continue;
}
let (name, value) = line.split_once(':').unwrap_or((line, ""));
let value = value.strip_prefix(' ').unwrap_or(value);
match name {
"event" => event = value.to_string(),
"data" => data.push(value),
_ => {}
}
}
(!data.is_empty()).then(|| Frame {
event,
data: data.join("\n"),
})
}
fn piece(kind: u8, number: u64, other: u64, bytes: &[u8]) -> Piece {
let ending = LINE_ENDINGS[usize::from(kind / 9) % LINE_ENDINGS.len()];
match kind % 9 {
0 => {
let stream = StreamKind::ALL[usize::from(number % 2 == 1)];
let output = OutputEvent {
offset: number,
stream,
output: base64::engine::general_purpose::STANDARD.encode(bytes),
};
Piece::Event {
wire: frame_wire(EVENT_OUTPUT, &json(&output), ending),
event: output_value(stream, number, bytes),
}
}
1 => {
let gap = GapEvent {
from: number,
to: other,
};
Piece::Event {
wire: frame_wire(EVENT_GAP, &json(&gap), ending),
event: gap_value(gap.from, gap.to),
}
}
2 => {
let bit = |at: u32| number >> at & 1 == 1;
let [a, b, c, d, e, f, g, h] = other.to_le_bytes();
let exit = ExitEvent {
exit_code: bit(0).then_some(i32::from_le_bytes([a, b, c, d])),
signal: bit(1).then_some(i32::from_le_bytes([e, f, g, h])),
timed_out: bit(2),
truncated: bit(3),
writers_may_be_alive: bit(4),
offset: number,
};
Piece::Event {
wire: frame_wire(EVENT_EXIT, &json(&exit), ending),
event: exit_value(&exit),
}
}
3 => {
let output = OutputEvent {
offset: number,
stream: StreamKind::Stdout,
output: format!("{}!", one_line(bytes)),
};
Piece::BadBase64 {
wire: frame_wire(EVENT_OUTPUT, &json(&output), ending),
offset: number,
}
}
4 => Piece::Keepalive {
wire: format!(":{}{ending}{ending}", one_line(bytes)).into_bytes(),
},
5 => Piece::Unknown {
wire: frame_wire(&format!("x-{}", one_line(bytes)), &one_line(bytes), ending),
},
6 => {
let (event, data) = wrong_shape(number, other);
Piece::WrongShape {
wire: frame_wire(event, &data, ending),
}
}
7 => {
let lines = chosen_lines(bytes);
let mut wire = lines.join(ending);
wire.push_str(ending);
wire.push_str(ending);
Piece::Fields {
wire: wire.into_bytes(),
frame: frame_of(&lines),
}
}
_ => {
let mut wire = bytes.to_vec();
wire.extend_from_slice(b"\n\n");
Piece::Junk { wire }
}
}
}
fn json(event: &impl serde::Serialize) -> String {
serde_json::to_string(event).expect("a protocol event serializes")
}
fn output_value(stream: StreamKind, offset: u64, data: &[u8]) -> serde_json::Value {
serde_json::json!({ "output": { "stream": stream, "offset": offset, "data": data } })
}
fn gap_value(from: u64, to: u64) -> serde_json::Value {
serde_json::json!({ "gap": { "from": from, "to": to } })
}
fn exit_value(exit: &ExitEvent) -> serde_json::Value {
serde_json::json!({ "exit": exit })
}
fn value(event: &ExecEvent) -> serde_json::Value {
match event {
ExecEvent::Output {
stream,
offset,
data,
} => output_value(*stream, *offset, data),
ExecEvent::Gap { from, to } => gap_value(*from, *to),
ExecEvent::Exit(exit) => exit_value(exit),
}
}
fn encode(event: &ExecEvent) -> Frame {
let (name, data) = match event {
ExecEvent::Output {
stream,
offset,
data,
} => (
EVENT_OUTPUT,
json(&OutputEvent {
offset: *offset,
stream: *stream,
output: base64::engine::general_purpose::STANDARD.encode(data),
}),
),
ExecEvent::Gap { from, to } => (
EVENT_GAP,
json(&GapEvent {
from: *from,
to: *to,
}),
),
ExecEvent::Exit(exit) => (EVENT_EXIT, json(exit)),
};
Frame {
event: name.to_string(),
data,
}
}
fn render(frame: &Frame) -> Vec<u8> {
let mut wire = format!("event: {}\n", frame.event);
for line in frame.data.split('\n') {
wire.push_str("data: ");
wire.push_str(line);
wire.push('\n');
}
wire.push('\n');
wire.into_bytes()
}
fn check_piece(piece: &Piece) {
let frames = SseParser::new()
.feed(piece.wire())
.expect("one piece is far under the ceiling");
match piece {
Piece::Event { event, .. } => {
assert_eq!(frames.len(), 1, "BIND-22: a daemon frame is one frame");
let decoded = decode(&frames[0])
.expect("BIND-22: a daemon frame decodes")
.expect("BIND-22: a daemon frame is an event");
assert_eq!(
&value(&decoded),
event,
"BIND-22: the event it was built from"
);
}
Piece::BadBase64 { offset, .. } => {
assert_eq!(
frames.len(),
1,
"BIND-22: a corrupt output frame is still a frame"
);
let error = decode(&frames[0]).expect_err("BIND-22: corrupt output is an error");
assert!(
error.to_string().contains(&offset.to_string()),
"BIND-22: the error names the offset: {error}"
);
}
Piece::Unknown { .. } => {
assert_eq!(
frames.len(),
1,
"BIND-22: an unknown event is still a frame"
);
assert!(
decode(&frames[0])
.expect("an unknown event isn't an error")
.is_none(),
"BIND-22: an unknown event is dropped"
);
}
Piece::WrongShape { .. } => {
assert_eq!(
frames.len(),
1,
"BIND-22: a mistyped event is still a frame"
);
assert!(
decode(&frames[0])
.expect("a mistyped event isn't an error")
.is_none(),
"BIND-22: a mistyped event is dropped: {:?}",
frames[0]
);
}
Piece::Fields { frame, .. } => {
assert_eq!(
frames.as_slice(),
frame.as_slice(),
"BIND-22: chosen lines parse by the SSE field rules"
);
}
Piece::Keepalive { .. } => {
assert!(frames.is_empty(), "BIND-22: a comment is not a frame");
}
Piece::Junk { .. } => {}
}
}
fn check_round_trip(frame: &Frame) {
let mut parser = SseParser::new();
let again = parser
.feed(&render(frame))
.expect("one frame is far under the ceiling");
assert_eq!(
again,
std::slice::from_ref(frame),
"BIND-22: the frame round-trips"
);
assert_eq!(
parser.pending(),
0,
"BIND-22: a whole frame leaves nothing held"
);
match decode(frame) {
Ok(Some(event)) => {
let reencoded = decode(&encode(&event))
.expect("BIND-22: a re-encoded event decodes")
.expect("BIND-22: a re-encoded event is an event");
assert_eq!(
value(&reencoded),
value(&event),
"BIND-22: the event round-trips"
);
}
Ok(None) => {}
Err(error) => assert_eq!(
frame.event, EVENT_OUTPUT,
"BIND-22: only output fails to decode, not {frame:?}: {error}"
),
}
}
#[test]
fn hostile_stream_bytes_stay_bounded_and_every_event_round_trips() {
const ITERATIONS: usize = 128;
bolero::check!()
.with_iterations(ITERATIONS)
.with_shrink_time(std::time::Duration::ZERO)
.with_type::<(Vec<(u8, u64, u64, Vec<u8>)>, Vec<u16>, Option<(u8, u8, u8)>)>()
.for_each(|(pieces, splits, flood)| {
let pieces: Vec<Piece> = pieces
.iter()
.map(|(kind, number, other, bytes)| piece(*kind, *number, *other, bytes))
.collect();
let mut expected = Vec::new();
let mut whole = Vec::new();
for piece in &pieces {
check_piece(piece);
expected.extend(
SseParser::new()
.feed(piece.wire())
.expect("one piece is far under the ceiling"),
);
whole.extend_from_slice(piece.wire());
}
let mut one_read = SseParser::new();
let frames = one_read
.feed(&whole)
.expect("the pieces are far under the ceiling");
assert_eq!(
frames, expected,
"BIND-22: one read yields each piece's frames"
);
let residue = one_read.pending();
for frame in &expected {
check_round_trip(frame);
}
let mut sizes = splits.iter().cycle();
let mut parser = SseParser::new();
let mut emitted = Vec::new();
let mut at = 0;
while at < whole.len() {
let size = sizes
.next()
.map_or(whole.len(), |size| usize::from(*size) % 64 + 1);
let end = at + size.min(whole.len() - at);
emitted.extend(
parser
.feed(&whole[at..end])
.expect("the pieces are far under the ceiling"),
);
at = end;
}
assert_eq!(
emitted, expected,
"BIND-22: the chunks yield each piece's frames"
);
assert_eq!(
parser.pending(),
residue,
"BIND-22: the chunks yield each piece's frames and hold what one read holds"
);
let flood = flood.filter(|(gate, ..)| gate % 4 == 0).map_or_else(
Vec::new,
|(_, byte, over)| {
let byte = if byte == b'\n' || byte == b'\r' {
b'x'
} else {
byte
};
vec![byte; MAX_PENDING_BYTES - 4 + usize::from(over)]
},
);
let refusal_due = residue + flood.len() > MAX_PENDING_BYTES;
let framed = whole.len();
whole.extend_from_slice(&flood);
let floor = flood.len() / 1024;
let mut parser = SseParser::new();
let mut emitted = Vec::new();
let mut refused = false;
let mut at = 0;
let mut sizes = splits.iter().cycle();
while at < whole.len() {
let size = match sizes.next() {
None => whole.len(),
Some(size) if at < framed => usize::from(*size) % 64 + 1,
Some(size) => (usize::from(*size) + 1).max(floor),
};
let end = at + size.min(whole.len() - at);
match parser.feed(&whole[at..end]) {
Ok(frames) => {
assert!(
parser.pending() <= MAX_PENDING_BYTES,
"BIND-22: {} bytes held, over the {MAX_PENDING_BYTES}-byte ceiling",
parser.pending()
);
emitted.extend(frames);
}
Err(error) => {
assert_eq!(error.kind(), ErrorKind::Protocol, "{error}");
assert!(
!error.retryable(),
"BIND-22: a reconnect refills it: {error}"
);
assert_eq!(parser.pending(), 0, "BIND-22: refused bytes are dropped");
refused = true;
break;
}
}
at = end;
}
assert_eq!(
refused,
refusal_due,
"BIND-22: {residue} + {} undelimited bytes against the ceiling",
flood.len()
);
if refused {
assert!(
expected.starts_with(&emitted),
"BIND-22: a refused stream yields a prefix of its frames"
);
} else {
assert_eq!(
emitted, expected,
"BIND-22: a flood under the ceiling leaves the frames whole"
);
assert_eq!(parser.pending(), residue + flood.len());
}
});
}