use std::io::{self, Read, Write};
use weir_core::{
Durability, Envelope, HEADER_LEN, Header, MAX_PAYLOAD_HARD_CAP, MessageType, NackReason,
};
const MAX_RESPONSE_PAYLOAD_LEN: usize = 2;
#[cfg(unix)]
pub type DefaultTransport = std::os::unix::net::UnixStream;
#[cfg(not(unix))]
pub type DefaultTransport = std::net::TcpStream;
#[derive(Debug)]
#[non_exhaustive]
pub enum ClientError {
Io(io::Error),
Nack(NackReason),
VersionMismatch {
daemon_version: u8,
},
UnknownNack(u8),
PayloadTooLarge {
len: usize,
limit: usize,
},
Protocol(String),
NoDefaultDurability,
EmptyPayload,
}
impl ClientError {
#[must_use]
pub fn is_recoverable(&self) -> bool {
match self {
Self::PayloadTooLarge { .. } | Self::NoDefaultDurability | Self::EmptyPayload => true,
Self::Nack(NackReason::InternalError) => true,
Self::Nack(_) => false,
Self::VersionMismatch { .. } | Self::UnknownNack(_) => false,
Self::Io(_) | Self::Protocol(_) => false,
#[allow(unreachable_patterns)]
_ => false,
}
}
}
impl std::fmt::Display for ClientError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Io(e) => write!(f, "socket I/O error: {e}"),
Self::Nack(r) => write!(f, "server nack: {r:?}"),
Self::VersionMismatch { daemon_version } => write!(
f,
"wire version mismatch: daemon speaks v{daemon_version}, this client speaks v{}",
weir_core::WIRE_VERSION
),
Self::UnknownNack(b) => write!(f, "server nack with unknown reason {b:#04x}"),
Self::PayloadTooLarge { len, limit } => write!(
f,
"payload too large: {len} bytes exceeds the {limit}-byte protocol hard cap"
),
Self::Protocol(msg) => write!(f, "protocol error: {msg}"),
Self::NoDefaultDurability => {
write!(f, "push_default called but no default durability was set")
}
Self::EmptyPayload => write!(
f,
"payload is empty; rejected locally (an empty payload is the WAB end-of-records sentinel)"
),
}
}
}
impl std::error::Error for ClientError {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
match self {
Self::Io(e) => Some(e),
_ => None,
}
}
}
impl From<io::Error> for ClientError {
fn from(e: io::Error) -> Self {
Self::Io(e)
}
}
pub struct WeirClient<S = DefaultTransport> {
pub(crate) stream: S,
pub(crate) default_durability: Option<Durability>,
pub(crate) poisoned: bool,
pub(crate) closed_after_nack: bool,
}
impl<S> std::fmt::Debug for WeirClient<S> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("WeirClient")
.field("default_durability", &self.default_durability)
.field("poisoned", &self.poisoned)
.field("closed_after_nack", &self.closed_after_nack)
.finish_non_exhaustive()
}
}
impl<S: Read + Write> WeirClient<S> {
pub fn set_default_durability(&mut self, durability: Durability) {
self.default_durability = Some(durability);
}
#[must_use = "a dropped push result hides whether the record was acked; an \
unhandled Nack also closes the connection"]
pub fn push(
&mut self,
payload: impl AsRef<[u8]>,
durability: Durability,
) -> Result<(), ClientError> {
self.ensure_usable()?;
let bytes = payload.as_ref();
if bytes.is_empty() {
return Err(ClientError::EmptyPayload);
}
if bytes.len() > MAX_PAYLOAD_HARD_CAP {
return Err(ClientError::PayloadTooLarge {
len: bytes.len(),
limit: MAX_PAYLOAD_HARD_CAP,
});
}
let payload = weir_core::Payload::copy_from_slice(bytes);
let header = Header::new(MessageType::Push, durability, 0);
let frame = Envelope::new(header, payload).encode();
if let Err(write_err) = self.stream.write_all(&frame) {
if let Ok(resp) = self.read_response()
&& resp.header().message_type() == MessageType::Nack
{
self.closed_after_nack = true;
return Err(nack_error(resp.payload()));
}
self.poisoned = true;
return Err(write_err.into());
}
let resp = self.read_response()?;
match resp.header().message_type() {
MessageType::Ack => Ok(()),
MessageType::Nack => Err(self.surface_nack(resp.payload())),
other => {
self.poisoned = true;
Err(ClientError::Protocol(format!(
"expected Ack or Nack, got {other:?}"
)))
}
}
}
fn note_nack(&mut self, err: ClientError) -> ClientError {
if !matches!(err, ClientError::Nack(NackReason::InternalError)) {
self.closed_after_nack = true;
}
err
}
fn surface_nack(&mut self, payload: &[u8]) -> ClientError {
match nack_error(payload) {
ClientError::Protocol(msg) => {
self.poisoned = true;
ClientError::Protocol(msg)
}
other => self.note_nack(other),
}
}
pub fn push_default(&mut self, payload: impl AsRef<[u8]>) -> Result<(), ClientError> {
let d = self
.default_durability
.ok_or(ClientError::NoDefaultDurability)?;
self.push(payload, d)
}
pub fn health_check(&mut self) -> Result<(), ClientError> {
self.ensure_usable()?;
let header = Header::new(MessageType::HealthCheck, Durability::Durable, 0);
let frame = Envelope::new(header, vec![]).encode();
if let Err(e) = self.stream.write_all(&frame) {
self.poisoned = true;
return Err(e.into());
}
let resp = self.read_response()?;
match resp.header().message_type() {
MessageType::HealthCheckResponse => Ok(()),
MessageType::Nack => Err(self.surface_nack(resp.payload())),
other => {
self.poisoned = true;
Err(ClientError::Protocol(format!(
"expected HealthCheckResponse, got {other:?}"
)))
}
}
}
#[must_use]
pub fn is_poisoned(&self) -> bool {
self.poisoned || self.closed_after_nack
}
pub(crate) fn ensure_usable(&self) -> Result<(), ClientError> {
if self.closed_after_nack {
return Err(ClientError::Protocol(
"connection closed by the daemon after a Nack; \
drop this client and reconnect before sending again"
.into(),
));
}
if self.poisoned {
return Err(ClientError::Protocol(
"client connection poisoned by a prior response-read error/timeout; \
drop this client and reconnect before sending again"
.into(),
));
}
Ok(())
}
fn read_response(&mut self) -> Result<Envelope, ClientError> {
let result = self.read_response_inner();
if result.is_err() {
self.poisoned = true;
}
result
}
fn read_response_inner(&mut self) -> Result<Envelope, ClientError> {
let mut header_buf = [0u8; HEADER_LEN];
self.stream.read_exact(&mut header_buf)?;
let header =
Header::decode(&header_buf).map_err(|e| ClientError::Protocol(e.to_string()))?;
let payload_len = header.payload_len() as usize;
if payload_len > MAX_RESPONSE_PAYLOAD_LEN {
return Err(ClientError::Protocol(format!(
"response payload_len {payload_len} exceeds the {MAX_RESPONSE_PAYLOAD_LEN}-byte \
response cap; refusing to allocate"
)));
}
let mut payload_buf = vec![0u8; payload_len];
if payload_len > 0 {
self.stream.read_exact(&mut payload_buf)?;
}
let mut crc_buf = [0u8; 4];
self.stream.read_exact(&mut crc_buf)?;
let expected = u32::from_le_bytes(crc_buf);
let computed = crc32fast::hash(&payload_buf);
if expected != computed {
return Err(ClientError::Protocol(format!(
"response payload CRC mismatch: expected {expected:#010x}, computed {computed:#010x}"
)));
}
Ok(Envelope::new(header, payload_buf))
}
#[cfg_attr(not(feature = "tls"), allow(dead_code))]
pub(crate) fn from_parts(stream: S, default_durability: Option<Durability>) -> Self {
Self {
stream,
default_durability,
poisoned: false,
closed_after_nack: false,
}
}
}
pub(crate) fn nack_error(payload: &[u8]) -> ClientError {
match payload.first().copied() {
Some(b) => match NackReason::try_from(b) {
Ok(NackReason::VersionMismatch) => match payload.get(1).copied() {
Some(daemon_version) => ClientError::VersionMismatch { daemon_version },
None => ClientError::Nack(NackReason::VersionMismatch),
},
Ok(r) => ClientError::Nack(r),
Err(_) => ClientError::UnknownNack(b),
},
None => ClientError::Protocol("Nack frame had empty payload".into()),
}
}