use std::io::{BufRead, Write};
use gunnar_pktline::{
decode, Decoded, Packet, BAND_ERROR, BAND_PACK, BAND_PROGRESS, SIDEBAND_CHUNK,
};
use crate::{Error, Result};
pub const TRANSPORT_BUFFER_BYTES: usize = 256 * 1024;
pub(crate) fn read_packet(input: &mut dyn BufRead) -> Result<Option<Packet>> {
let mut head = [0u8; 4];
match read_up_to(input, &mut head)? {
0 => return Ok(None),
4 => {}
_ => return Err(Error::UnexpectedEof("a pkt-line header")),
}
match decode(&head)? {
Decoded::Packet(pkt, _) => return Ok(Some(pkt)),
Decoded::Incomplete => {}
}
let len = line_length(&head);
let mut line = vec![0u8; len];
line[..4].copy_from_slice(&head);
input.read_exact(&mut line[4..]).map_err(|err| {
if err.kind() == std::io::ErrorKind::UnexpectedEof {
Error::UnexpectedEof("a pkt-line payload")
} else {
Error::Io(err)
}
})?;
match decode(&line)? {
Decoded::Packet(pkt, used) if used == len => Ok(Some(pkt)),
_ => Err(Error::protocol(format!(
"a {len}-byte pkt-line did not decode as one packet"
))),
}
}
pub(crate) fn expect_packet(input: &mut dyn BufRead, what: &'static str) -> Result<Packet> {
read_packet(input)?.ok_or(Error::UnexpectedEof(what))
}
fn line_length(head: &[u8; 4]) -> usize {
head.iter().fold(0usize, |acc, &b| {
let nibble = match b {
b'0'..=b'9' => b - b'0',
b'a'..=b'f' => b - b'a' + 10,
b'A'..=b'F' => b - b'A' + 10,
_ => 0,
};
(acc << 4) | nibble as usize
})
}
fn read_up_to(input: &mut dyn BufRead, buf: &mut [u8]) -> Result<usize> {
let mut filled = 0;
while filled < buf.len() {
match input.read(&mut buf[filled..])? {
0 => break,
n => filled += n,
}
}
Ok(filled)
}
pub(crate) fn read_section(
input: &mut dyn BufRead,
what: &'static str,
budget: &mut crate::limits::Budget,
) -> Result<Vec<Vec<u8>>> {
let mut lines = Vec::new();
loop {
match expect_packet(input, what)? {
Packet::Flush => return Ok(lines),
Packet::Data(data) => {
budget.spend(data.len(), what)?;
let text = data.strip_suffix(b"\n").unwrap_or(&data);
lines.push(text.to_vec());
}
other => {
return Err(Error::protocol(format!(
"unexpected {other:?} while reading {what}"
)))
}
}
}
}
pub(crate) struct SidebandWriter<'a> {
out: &'a mut dyn Write,
band: u8,
buf: Vec<u8>,
}
impl<'a> SidebandWriter<'a> {
pub(crate) fn new(out: &'a mut dyn Write, band: u8) -> Self {
SidebandWriter {
out,
band,
buf: Vec::with_capacity(SIDEBAND_CHUNK),
}
}
fn emit(&mut self) -> std::io::Result<()> {
if self.buf.is_empty() {
return Ok(());
}
gunnar_pktline::write_sideband(&mut self.out, self.band, &self.buf).map_err(as_io)?;
self.buf.clear();
Ok(())
}
}
impl Write for SidebandWriter<'_> {
fn write(&mut self, data: &[u8]) -> std::io::Result<usize> {
if data.is_empty() {
return Ok(0);
}
let room = SIDEBAND_CHUNK - self.buf.len();
let take = room.min(data.len());
self.buf.extend_from_slice(&data[..take]);
if self.buf.len() == SIDEBAND_CHUNK {
self.emit()?;
}
Ok(take)
}
fn flush(&mut self) -> std::io::Result<()> {
self.emit()?;
self.out.flush()
}
}
fn as_io(err: gunnar_pktline::Error) -> std::io::Error {
match err {
gunnar_pktline::Error::Io(err) => err,
other => std::io::Error::other(other.to_string()),
}
}
pub(crate) enum PackOutput<'a> {
Raw(&'a mut dyn Write),
Sideband(SidebandWriter<'a>),
}
impl<'a> PackOutput<'a> {
pub(crate) fn new(out: &'a mut dyn Write, sideband: bool) -> Self {
if sideband {
PackOutput::Sideband(SidebandWriter::new(out, BAND_PACK))
} else {
PackOutput::Raw(out)
}
}
}
impl Write for PackOutput<'_> {
fn write(&mut self, data: &[u8]) -> std::io::Result<usize> {
match self {
PackOutput::Raw(out) => out.write(data),
PackOutput::Sideband(out) => out.write(data),
}
}
fn flush(&mut self) -> std::io::Result<()> {
match self {
PackOutput::Raw(out) => out.flush(),
PackOutput::Sideband(out) => out.flush(),
}
}
}
pub(crate) fn text(out: &mut dyn Write, line: &[u8]) -> Result<()> {
let mut sink: &mut dyn Write = out;
gunnar_pktline::write_text(&mut sink, line)?;
Ok(())
}
pub(crate) fn data(out: &mut dyn Write, payload: &[u8]) -> Result<()> {
let mut sink: &mut dyn Write = out;
gunnar_pktline::write_data(&mut sink, payload)?;
Ok(())
}
pub(crate) fn flush(out: &mut dyn Write) -> Result<()> {
let mut sink: &mut dyn Write = out;
gunnar_pktline::write_flush(&mut sink)?;
Ok(())
}
pub(crate) fn delim(out: &mut dyn Write) -> Result<()> {
let mut sink: &mut dyn Write = out;
gunnar_pktline::write_delim(&mut sink)?;
Ok(())
}
pub(crate) fn sideband(out: &mut dyn Write, band: u8, payload: &[u8]) -> Result<()> {
let mut sink: &mut dyn Write = out;
gunnar_pktline::write_sideband(&mut sink, band, payload)?;
Ok(())
}
pub(crate) fn progress(out: &mut dyn Write, enabled: bool, message: &str) -> Result<()> {
if !enabled || message.is_empty() {
return Ok(());
}
sideband(out, BAND_PROGRESS, message.as_bytes())?;
Ok(())
}
pub(crate) fn sideband_error(out: &mut dyn Write, message: &str) -> Result<()> {
let mut text = message.to_owned();
if !text.ends_with('\n') {
text.push('\n');
}
sideband(out, BAND_ERROR, text.as_bytes())?;
Ok(())
}
pub fn err_packet(out: &mut dyn Write, message: &str) -> Result<()> {
let mut line = String::from("ERR ");
line.push_str(message);
text(out, line.as_bytes())?;
Ok(())
}
#[cfg(test)]
mod transport_buffer_tests {
use std::io::{BufWriter, Write};
use super::{PackOutput, TRANSPORT_BUFFER_BYTES};
#[derive(Default)]
struct Calls {
sizes: Vec<usize>,
}
impl Write for Calls {
fn write(&mut self, data: &[u8]) -> std::io::Result<usize> {
self.sizes.push(data.len());
Ok(data.len())
}
fn flush(&mut self) -> std::io::Result<()> {
Ok(())
}
}
fn writes_under(capacity: usize, bytes: usize) -> Vec<usize> {
let payload = vec![0x5Au8; bytes];
let mut calls = Calls::default();
{
let mut buffered = BufWriter::with_capacity(capacity, &mut calls);
{
let mut pack = PackOutput::new(&mut buffered, true);
pack.write_all(&payload)
.expect("a counting sink cannot fail");
pack.flush().expect("a counting sink cannot fail");
}
buffered.flush().expect("a counting sink cannot fail");
}
calls.sizes
}
#[test]
fn transport_buffer_holds_whole_sideband_packets() {
const MIB: usize = 1024 * 1024;
let red = writes_under(8 * 1024, MIB);
assert!(
red.iter().any(|n| *n <= 8),
"the 8 KiB arm is supposed to expose a 5-byte write per sideband packet and did not, \
so this test is no longer measuring what it says: {red:?}"
);
assert!(
red.len() >= 32,
"1 MiB is 16 sideband packets and the 8 KiB arm should cost two writes each; it cost \
{} — the red arm has stopped being red",
red.len()
);
let green = writes_under(TRANSPORT_BUFFER_BYTES, MIB);
assert!(
green.iter().all(|n| *n > 8),
"a buffer of {TRANSPORT_BUFFER_BYTES} bytes still handed the stream a write of {:?} \
bytes, which is a pkt-line header on its own — the payload no longer fits the \
buffer",
green.iter().min()
);
assert!(
green.len() * 4 <= red.len(),
"{} write(s) at {TRANSPORT_BUFFER_BYTES} bytes against {} at 8 KiB is under the 4x \
this buffer exists to buy",
green.len(),
red.len()
);
}
}