use snafu::Snafu;
pub trait Reads {
fn read(&mut self, buf: &mut [u8]) -> Result<usize, StreamError>;
fn read_exact(&mut self, buf: &mut [u8]) -> Result<(), StreamError> {
let mut read = 0;
while read < buf.len() {
read += self.read(&mut buf[read..])?;
}
Ok(())
}
}
#[cfg(not(any(feature = "std", test)))]
impl Reads for &[u8] {
fn read(&mut self, buf: &mut [u8]) -> Result<usize, StreamError> {
let amt = core::cmp::min(buf.len(), self.len());
let (a, b) = self.split_at(amt);
if amt == 1 {
buf[0] = a[0];
} else {
buf[..amt].copy_from_slice(a);
}
*self = b;
Ok(amt)
}
fn read_exact(&mut self, buf: &mut [u8]) -> Result<(), StreamError> {
if buf.len() > self.len() {
return Err(StreamError::Closed);
}
let (a, b) = self.split_at(buf.len());
if buf.len() == 1 {
buf[0] = a[0];
} else {
buf.copy_from_slice(a);
}
*self = b;
Ok(())
}
}
#[cfg(not(any(feature = "std", test)))]
impl<R: Reads + ?Sized> Reads for &mut R {
fn read(&mut self, buf: &mut [u8]) -> Result<usize, StreamError> {
(**self).read(buf)
}
fn read_exact(&mut self, buf: &mut [u8]) -> Result<(), StreamError> {
(**self).read_exact(buf)
}
}
#[cfg(any(feature = "std", test))]
impl<T> Reads for T
where
T: std::io::Read,
{
fn read(&mut self, buf: &mut [u8]) -> Result<usize, StreamError> {
self.read(buf).map_err(|e| match e.kind() {
std::io::ErrorKind::ConnectionReset
| std::io::ErrorKind::ConnectionAborted
| std::io::ErrorKind::BrokenPipe => StreamError::Closed,
std::io::ErrorKind::UnexpectedEof => StreamError::Empty,
_ => StreamError::Other {
message: "Unexpected IO Error",
},
})
}
fn read_exact(&mut self, buf: &mut [u8]) -> Result<(), StreamError> {
self.read_exact(buf).map_err(|e| match e.kind() {
std::io::ErrorKind::ConnectionReset
| std::io::ErrorKind::ConnectionAborted
| std::io::ErrorKind::BrokenPipe => StreamError::Closed,
std::io::ErrorKind::UnexpectedEof => StreamError::Empty,
_ => StreamError::Other {
message: "Unexpected IO Error",
},
})
}
}
pub trait Writes {
fn write(&mut self, buf: &[u8]) -> Result<usize, StreamError>;
fn write_all(&mut self, buf: &[u8]) -> Result<(), StreamError> {
let mut written = 0;
while written < buf.len() {
written += self.write(&buf[written..])?;
}
Ok(())
}
}
#[cfg_attr(
not(any(
feature = "langs-python",
feature = "langs-sql",
feature = "langs-typescript",
feature = "langs-open-api",
test
)),
allow(dead_code)
)]
pub(crate) struct FmtWriter<'w, W: Writes> {
writes: &'w mut W,
}
impl<'w, W: Writes> From<&'w mut W> for FmtWriter<'w, W> {
fn from(value: &'w mut W) -> Self {
Self { writes: value }
}
}
impl<W: Writes> core::fmt::Write for FmtWriter<'_, W> {
fn write_str(&mut self, s: &str) -> core::fmt::Result {
match self.writes.write_all(s.as_bytes()) {
Ok(_) => Ok(()),
Err(_) => Err(core::fmt::Error),
}
}
}
#[cfg(not(any(feature = "std", test)))]
impl Writes for alloc::vec::Vec<u8> {
fn write(&mut self, buf: &[u8]) -> Result<usize, StreamError> {
self.extend_from_slice(buf);
Ok(buf.len())
}
fn write_all(&mut self, buf: &[u8]) -> Result<(), StreamError> {
self.extend_from_slice(buf);
Ok(())
}
}
#[cfg(any(feature = "std", test))]
impl<T> Writes for T
where
T: std::io::Write,
{
fn write(&mut self, buf: &[u8]) -> Result<usize, StreamError> {
let written = self.write(buf).map_err(|e| match e.kind() {
std::io::ErrorKind::ConnectionReset
| std::io::ErrorKind::ConnectionAborted
| std::io::ErrorKind::BrokenPipe
| std::io::ErrorKind::UnexpectedEof => StreamError::Closed,
_ => StreamError::Other {
message: "Unexpected IO Error",
},
})?;
if written == 0 {
Err(StreamError::Closed)
} else {
Ok(written)
}
}
fn write_all(&mut self, buf: &[u8]) -> Result<(), StreamError> {
self.write_all(buf).map_err(|e| match e.kind() {
std::io::ErrorKind::ConnectionReset
| std::io::ErrorKind::ConnectionAborted
| std::io::ErrorKind::BrokenPipe
| std::io::ErrorKind::UnexpectedEof => StreamError::Closed,
_ => StreamError::Other {
message: "Unexpected IO Error",
},
})
}
}
#[derive(Debug, Clone, PartialEq, Eq, Snafu)]
pub enum StreamError {
Empty,
Closed,
#[snafu(display("{message}"))]
Other { message: &'static str },
}