use core::fmt;
pub use crate::error::Full;
use crate::{class, error::WriteError, text};
pub trait Sink {
fn put(&mut self, bytes: &[u8]) -> Result<(), Full>;
fn remaining(&self) -> Option<usize> {
None
}
}
impl<S: Sink + ?Sized> Sink for &mut S {
#[inline]
fn put(&mut self, bytes: &[u8]) -> Result<(), Full> {
(**self).put(bytes)
}
#[inline]
fn remaining(&self) -> Option<usize> {
(**self).remaining()
}
}
#[cfg(feature = "alloc")]
#[cfg_attr(docsrs, doc(cfg(feature = "alloc")))]
impl Sink for alloc::vec::Vec<u8> {
#[inline]
fn put(&mut self, bytes: &[u8]) -> Result<(), Full> {
self.extend_from_slice(bytes);
Ok(())
}
}
pub struct SliceSink<'a> {
buf: &'a mut [u8],
len: usize,
}
impl<'a> SliceSink<'a> {
#[must_use]
pub fn new(buf: &'a mut [u8]) -> Self {
Self { buf, len: 0 }
}
#[must_use]
pub fn written(&self) -> &[u8] {
self.buf.get(..self.len).unwrap_or(&[])
}
#[must_use]
pub fn len(&self) -> usize {
self.len
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.len == 0
}
pub fn clear(&mut self) {
self.len = 0;
}
}
impl Sink for SliceSink<'_> {
#[inline]
fn put(&mut self, bytes: &[u8]) -> Result<(), Full> {
let end = self.len.checked_add(bytes.len()).ok_or(Full)?;
let dst = self.buf.get_mut(self.len..end).ok_or(Full)?;
dst.copy_from_slice(bytes);
self.len = end;
Ok(())
}
#[inline]
fn remaining(&self) -> Option<usize> {
Some(self.buf.len().saturating_sub(self.len))
}
}
impl fmt::Debug for SliceSink<'_> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("SliceSink")
.field("len", &self.len)
.field("capacity", &self.buf.len())
.finish()
}
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash)]
pub struct CountSink {
n: usize,
}
impl CountSink {
#[must_use]
pub const fn new() -> Self {
Self { n: 0 }
}
#[must_use]
pub const fn count(&self) -> usize {
self.n
}
}
impl Sink for CountSink {
#[inline]
fn put(&mut self, bytes: &[u8]) -> Result<(), Full> {
self.n = self.n.saturating_add(bytes.len());
Ok(())
}
}
#[inline]
pub(crate) fn reserve(out: &impl Sink, len: usize) -> Result<(), WriteError> {
match out.remaining() {
Some(room) if room < len => Err(WriteError::Full),
_ => Ok(()),
}
}
pub fn write_decimal_u64(n: u64, out: &mut impl Sink) -> Result<(), WriteError> {
let mut digits = [0u8; 20];
let mut i = digits.len();
let mut v = n;
loop {
i -= 1;
if let Some(d) = digits.get_mut(i) {
*d = b'0' + (v % 10) as u8;
}
v /= 10;
if v == 0 || i == 0 {
break;
}
}
out.put(digits.get(i..).unwrap_or(&[]))?;
Ok(())
}
pub fn write_hex_u64(n: u64, out: &mut impl Sink) -> Result<(), WriteError> {
const HEX: &[u8; 16] = b"0123456789abcdef";
let mut digits = [0u8; 16];
let mut i = digits.len();
let mut v = n;
loop {
i -= 1;
if let (Some(d), Some(&h)) = (digits.get_mut(i), HEX.get((v & 0xf) as usize)) {
*d = h;
}
v >>= 4;
if v == 0 || i == 0 {
break;
}
}
out.put(digits.get(i..).unwrap_or(&[]))?;
Ok(())
}
pub fn write_token(s: &[u8], out: &mut impl Sink) -> Result<(), WriteError> {
if !text::is_token(s) {
return Err(WriteError::NotToken);
}
out.put(s)?;
Ok(())
}
pub(crate) fn quoted_len(s: &[u8]) -> Result<usize, WriteError> {
let mut len = s.len().saturating_add(2);
for (i, &b) in s.iter().enumerate() {
match b {
b'"' | b'\\' => len = len.saturating_add(1),
b'\t' => {}
_ if class::is_quoted_pair_char(b) => {}
_ => return Err(WriteError::forbidden(i)),
}
}
Ok(len)
}
pub fn write_quoted(s: &[u8], out: &mut impl Sink) -> Result<(), WriteError> {
let len = quoted_len(s)?;
reserve(out, len)?;
put_quoted(s, out)
}
fn put_quoted(s: &[u8], out: &mut impl Sink) -> Result<(), WriteError> {
out.put(b"\"")?;
let mut rest = s;
while let Some(i) = rest.iter().position(|&b| b == b'"' || b == b'\\') {
let (plain, tail) = rest.split_at(i);
out.put(plain)?;
if let Some((escaped, after)) = tail.split_first() {
out.put(&[b'\\', *escaped])?;
rest = after;
} else {
rest = &[];
}
}
out.put(rest)?;
out.put(b"\"")?;
Ok(())
}
pub(crate) fn word_len(s: &[u8]) -> Result<usize, WriteError> {
if text::is_token(s) {
Ok(s.len())
} else {
quoted_len(s)
}
}
pub fn write_word(s: &[u8], out: &mut impl Sink) -> Result<(), WriteError> {
if text::is_token(s) {
out.put(s)?;
Ok(())
} else {
write_quoted(s, out)
}
}
pub fn write_field_value(s: &[u8], out: &mut impl Sink) -> Result<(), WriteError> {
if let Err(e) = crate::scan::check_field_value(s) {
return Err(WriteError::forbidden(e.offset()));
}
out.put(s)?;
Ok(())
}
pub fn write_param(name: &[u8], value: &[u8], out: &mut impl Sink) -> Result<(), WriteError> {
if !text::is_token(name) {
return Err(WriteError::NotToken);
}
let value_len = word_len(value)?;
reserve(out, name.len().saturating_add(3).saturating_add(value_len))?;
out.put(b"; ")?;
out.put(name)?;
out.put(b"=")?;
if text::is_token(value) {
out.put(value)?;
Ok(())
} else {
put_quoted(value, out)
}
}
#[cfg(test)]
mod tests {
use std::vec::Vec;
use super::*;
fn written(f: impl Fn(&mut SliceSink<'_>) -> Result<(), WriteError>) -> Vec<u8> {
let mut buf = [0u8; 256];
let mut out = SliceSink::new(&mut buf);
f(&mut out).unwrap();
out.written().to_vec()
}
#[test]
fn test_number_writers_known_values() {
assert_eq!(written(|o| write_decimal_u64(0, o)), b"0");
assert_eq!(written(|o| write_decimal_u64(10, o)), b"10");
assert_eq!(
written(|o| write_decimal_u64(u64::MAX, o)),
b"18446744073709551615"
);
assert_eq!(written(|o| write_hex_u64(0, o)), b"0");
assert_eq!(written(|o| write_hex_u64(u64::MAX, o)), b"ffffffffffffffff");
assert_eq!(written(|o| write_hex_u64(0x10, o)), b"10");
}
#[test]
fn test_writers_refuse_crlf_in_every_position() {
let bad: Vec<u8> = (0u8..0x20).filter(|&b| b != b'\t').chain([0x7f]).collect();
for &b in &bad {
for pos in 0..3 {
let mut input = *b"abc";
input[pos] = b;
let mut buf = [0u8; 64];
let mut out = SliceSink::new(&mut buf);
let at = pos as u32;
assert_eq!(
write_quoted(&input, &mut out),
Err(WriteError::Forbidden { at })
);
assert_eq!(
write_word(&input, &mut out),
Err(WriteError::Forbidden { at })
);
assert_eq!(
write_field_value(&input, &mut out),
Err(WriteError::Forbidden { at })
);
assert_eq!(
write_param(b"n", &input, &mut out),
Err(WriteError::Forbidden { at })
);
assert_eq!(write_token(&input, &mut out), Err(WriteError::NotToken));
assert!(out.is_empty());
}
}
}
#[test]
fn test_quoted_escapes_only_quote_and_backslash() {
assert_eq!(written(|o| write_quoted(b"", o)), b"\"\"");
assert_eq!(
written(|o| write_quoted(b"a\\b\"c\td\x80", o)),
b"\"a\\\\b\\\"c\td\x80\""
);
}
#[test]
fn test_slice_sink_one_byte_short_writes_nothing() {
for (input, needed) in [(&b"a\"b"[..], 6usize), (b"x y", 5)] {
let mut buf = [0u8; 16];
let mut out = SliceSink::new(&mut buf[..needed - 1]);
assert_eq!(write_quoted(input, &mut out), Err(WriteError::Full));
assert!(out.is_empty());
let mut out = SliceSink::new(&mut buf[..needed]);
assert_eq!(write_quoted(input, &mut out), Ok(()));
}
let mut buf = [0u8; 9];
let mut out = SliceSink::new(&mut buf);
assert_eq!(
write_param(b"name", b"a b", &mut out),
Err(WriteError::Full)
);
assert!(out.is_empty());
let mut buf = [0u8; 3];
let mut out = SliceSink::new(&mut buf);
assert_eq!(write_decimal_u64(1234, &mut out), Err(WriteError::Full));
assert!(out.is_empty());
}
#[test]
fn test_count_sink_matches_slice_output() {
let mut count = CountSink::new();
write_param(b"a", b"x\"y", &mut count).unwrap();
assert_eq!(
count.count(),
written(|o| write_param(b"a", b"x\"y", o)).len()
);
let mut c = CountSink { n: usize::MAX };
c.put(b"x").unwrap();
assert_eq!(c.count(), usize::MAX);
}
#[test]
fn test_param_name_must_be_token() {
let mut buf = [0u8; 16];
let mut out = SliceSink::new(&mut buf);
assert_eq!(
write_param(b"a b", b"x", &mut out),
Err(WriteError::NotToken)
);
assert_eq!(write_param(b"", b"x", &mut out), Err(WriteError::NotToken));
}
#[test]
fn test_sink_for_mut_reference_and_debug() {
let mut buf = [0u8; 4];
let mut out = SliceSink::new(&mut buf);
{
let r = &mut out;
write_token(b"ab", &mut &mut *r).unwrap();
assert_eq!(Sink::remaining(&r), Some(2));
}
assert_eq!(std::format!("{out:?}"), "SliceSink { len: 2, capacity: 4 }");
}
#[cfg(feature = "alloc")]
#[test]
fn test_vec_sink_grows() {
let mut v = Vec::new();
write_quoted(b"a b", &mut v).unwrap();
assert_eq!(v, b"\"a b\"");
assert_eq!(Sink::remaining(&v), None);
}
}