use alloc::vec::Vec;
use liminal_protocol::wire::{FRAME_MAX, GENERIC_HEADER_LEN};
const FRAME_TYPE_DELIVER: u8 = 0x19;
const PAYLOAD_LENGTH_OFFSET: usize = 6;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum SocketFailure {
Transport,
UnsupportedTextMessage,
MessageBeyondBound,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum SocketEvent {
Opened,
Binary(Vec<u8>),
Closed,
Failed(SocketFailure),
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum SocketCommand {
Open,
SendBinary(Vec<u8>),
Close,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum FrameViolation {
EmptyMessage,
TruncatedHeader {
length: usize,
},
DeclaredBeyondBound {
declared_total: u64,
bound: u64,
},
TruncatedBody {
declared_total: u64,
actual: u64,
},
TrailingBytes {
declared_total: u64,
actual: u64,
},
UndecodableBody,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum TransportTerminal {
PeerClosed,
CloseCompleted,
SocketFailed(SocketFailure),
ProtocolViolation(FrameViolation),
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum DriverPhase {
Idle,
Opening,
Established,
Closing,
Terminated,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum FrameCorrelation {
CorrelatedResponse,
UnsolicitedDelivery,
UnsolicitedFrame,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum PostTerminalEvent {
Opened,
Binary,
Closed,
Failed,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum EventRefusal {
OpenedWithoutOpenCommand,
BinaryBeforeEstablished,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ResponseExpectation {
None,
Correlated,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum CommandRefusal {
InvalidPhase {
phase: DriverPhase,
},
ExchangeOutstanding,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum DriverOutput {
Opened,
Frame {
bytes: Vec<u8>,
correlation: FrameCorrelation,
},
Terminal(TransportTerminal),
PostTerminalIgnored(PostTerminalEvent),
Refused(EventRefusal),
}
#[derive(Clone, Debug, PartialEq, Eq)]
#[must_use]
pub struct DriverStep {
pub output: DriverOutput,
pub command: Option<SocketCommand>,
}
impl DriverStep {
const fn output(output: DriverOutput) -> Self {
Self {
output,
command: None,
}
}
const fn with_command(output: DriverOutput, command: SocketCommand) -> Self {
Self {
output,
command: Some(command),
}
}
}
#[derive(Debug)]
pub struct WebSocketFrameDriver {
phase: DriverPhase,
exchange_outstanding: bool,
}
impl WebSocketFrameDriver {
#[must_use]
pub const fn new() -> Self {
Self {
phase: DriverPhase::Idle,
exchange_outstanding: false,
}
}
#[must_use]
pub const fn phase(&self) -> DriverPhase {
self.phase
}
#[must_use]
pub const fn has_outstanding_exchange(&self) -> bool {
self.exchange_outstanding
}
pub const fn command_open(&mut self) -> Result<SocketCommand, CommandRefusal> {
match self.phase {
DriverPhase::Idle => {
self.phase = DriverPhase::Opening;
Ok(SocketCommand::Open)
}
DriverPhase::Opening
| DriverPhase::Established
| DriverPhase::Closing
| DriverPhase::Terminated => Err(CommandRefusal::InvalidPhase { phase: self.phase }),
}
}
pub fn command_send(
&mut self,
bytes: Vec<u8>,
expectation: ResponseExpectation,
) -> Result<SocketCommand, CommandRefusal> {
if self.phase != DriverPhase::Established {
return Err(CommandRefusal::InvalidPhase { phase: self.phase });
}
match expectation {
ResponseExpectation::Correlated => {
if self.exchange_outstanding {
return Err(CommandRefusal::ExchangeOutstanding);
}
self.exchange_outstanding = true;
}
ResponseExpectation::None => {}
}
Ok(SocketCommand::SendBinary(bytes))
}
pub const fn command_close(&mut self) -> Result<SocketCommand, CommandRefusal> {
match self.phase {
DriverPhase::Opening | DriverPhase::Established => {
self.phase = DriverPhase::Closing;
Ok(SocketCommand::Close)
}
DriverPhase::Idle | DriverPhase::Closing | DriverPhase::Terminated => {
Err(CommandRefusal::InvalidPhase { phase: self.phase })
}
}
}
pub fn handle_event(&mut self, event: SocketEvent) -> DriverStep {
if self.phase == DriverPhase::Terminated {
let kind = match event {
SocketEvent::Opened => PostTerminalEvent::Opened,
SocketEvent::Binary(_) => PostTerminalEvent::Binary,
SocketEvent::Closed => PostTerminalEvent::Closed,
SocketEvent::Failed(_) => PostTerminalEvent::Failed,
};
return DriverStep::output(DriverOutput::PostTerminalIgnored(kind));
}
match event {
SocketEvent::Opened => self.handle_opened(),
SocketEvent::Binary(bytes) => self.handle_binary(bytes),
SocketEvent::Closed => self.handle_closed(),
SocketEvent::Failed(failure) => self.handle_failed(failure),
}
}
const fn handle_opened(&mut self) -> DriverStep {
match self.phase {
DriverPhase::Opening => {
self.phase = DriverPhase::Established;
DriverStep::output(DriverOutput::Opened)
}
DriverPhase::Idle | DriverPhase::Established | DriverPhase::Closing => {
DriverStep::output(DriverOutput::Refused(
EventRefusal::OpenedWithoutOpenCommand,
))
}
DriverPhase::Terminated => unreachable_terminated(),
}
}
fn handle_binary(&mut self, bytes: Vec<u8>) -> DriverStep {
match self.phase {
DriverPhase::Established | DriverPhase::Closing => match validate_frame(&bytes) {
Ok(()) => {
let correlation = self.classify_frame(&bytes);
DriverStep::output(DriverOutput::Frame { bytes, correlation })
}
Err(violation) => self.mint_terminal(
TransportTerminal::ProtocolViolation(violation),
self.phase == DriverPhase::Established,
),
},
DriverPhase::Idle | DriverPhase::Opening => {
DriverStep::output(DriverOutput::Refused(EventRefusal::BinaryBeforeEstablished))
}
DriverPhase::Terminated => unreachable_terminated(),
}
}
const fn handle_closed(&mut self) -> DriverStep {
let terminal = match self.phase {
DriverPhase::Closing => TransportTerminal::CloseCompleted,
DriverPhase::Idle | DriverPhase::Opening | DriverPhase::Established => {
TransportTerminal::PeerClosed
}
DriverPhase::Terminated => return unreachable_terminated(),
};
self.mint_terminal(terminal, false)
}
const fn handle_failed(&mut self, failure: SocketFailure) -> DriverStep {
let emit_close = !matches!(self.phase, DriverPhase::Closing);
self.mint_terminal(TransportTerminal::SocketFailed(failure), emit_close)
}
const fn mint_terminal(&mut self, terminal: TransportTerminal, emit_close: bool) -> DriverStep {
self.phase = DriverPhase::Terminated;
self.exchange_outstanding = false;
if emit_close {
DriverStep::with_command(DriverOutput::Terminal(terminal), SocketCommand::Close)
} else {
DriverStep::output(DriverOutput::Terminal(terminal))
}
}
fn classify_frame(&mut self, bytes: &[u8]) -> FrameCorrelation {
if bytes.first().copied() == Some(FRAME_TYPE_DELIVER) {
return FrameCorrelation::UnsolicitedDelivery;
}
if self.exchange_outstanding {
self.exchange_outstanding = false;
return FrameCorrelation::CorrelatedResponse;
}
FrameCorrelation::UnsolicitedFrame
}
}
impl Default for WebSocketFrameDriver {
fn default() -> Self {
Self::new()
}
}
const fn unreachable_terminated() -> DriverStep {
DriverStep::output(DriverOutput::PostTerminalIgnored(PostTerminalEvent::Failed))
}
fn validate_frame(bytes: &[u8]) -> Result<(), FrameViolation> {
if bytes.is_empty() {
return Err(FrameViolation::EmptyMessage);
}
if bytes.len() < GENERIC_HEADER_LEN {
return Err(FrameViolation::TruncatedHeader {
length: bytes.len(),
});
}
let Some(length_bytes) = bytes.get(PAYLOAD_LENGTH_OFFSET..GENERIC_HEADER_LEN) else {
return Err(FrameViolation::TruncatedHeader {
length: bytes.len(),
});
};
let Ok(length_array) = <[u8; 4]>::try_from(length_bytes) else {
return Err(FrameViolation::TruncatedHeader {
length: bytes.len(),
});
};
let declared_payload = u64::from(u32::from_be_bytes(length_array));
let Ok(header_len) = u64::try_from(GENERIC_HEADER_LEN) else {
return Err(FrameViolation::TruncatedHeader {
length: bytes.len(),
});
};
let declared_total = declared_payload.saturating_add(header_len);
if declared_total > FRAME_MAX {
return Err(FrameViolation::DeclaredBeyondBound {
declared_total,
bound: FRAME_MAX,
});
}
let actual = u64::try_from(bytes.len()).unwrap_or(u64::MAX);
if actual < declared_total {
return Err(FrameViolation::TruncatedBody {
declared_total,
actual,
});
}
if actual > declared_total {
return Err(FrameViolation::TrailingBytes {
declared_total,
actual,
});
}
Ok(())
}