use std::time::Duration;
use crate::cbor::Value;
use crate::connection::{FrameStream, RecvFrameError, SendFrameError, Session};
use crate::frame::{self, StreamEncoding, StreamMode, StreamRole};
use crate::identity::KeyPair;
pub struct StreamHandle {
stream: FrameStream,
pub stream_id: [u8; 16],
pub mode: StreamMode,
seq_out: u64,
}
#[derive(Debug)]
pub enum OpenError {
OpenStream(quinn::ConnectionError),
Send(SendFrameError),
}
impl std::fmt::Display for OpenError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
OpenError::OpenStream(e) => write!(f, "opening a dedicated stream: {e}"),
OpenError::Send(e) => write!(f, "sending stream_open: {e}"),
}
}
}
impl std::error::Error for OpenError {}
#[derive(Debug)]
pub enum AcceptError {
AcceptStream(quinn::ConnectionError),
Timeout,
Recv(RecvFrameError),
Parse(frame::ParseStreamOpenError),
}
impl std::fmt::Display for AcceptError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
AcceptError::AcceptStream(e) => write!(f, "accepting a dedicated stream: {e}"),
AcceptError::Timeout => write!(f, "no inbound stream within the given timeout"),
AcceptError::Recv(e) => write!(f, "reading the stream's first frame: {e}"),
AcceptError::Parse(e) => write!(f, "expected a stream_open frame: {e}"),
}
}
}
impl std::error::Error for AcceptError {}
#[derive(Debug, Clone)]
pub enum StreamItem {
Data {
seq: u64,
encoding: StreamEncoding,
body: Value,
},
Eof,
}
#[derive(Debug)]
pub enum RecvStreamError {
Recv(RecvFrameError),
Parse(frame::ParseStreamEventError),
PeerAborted {
code: String,
message: String,
},
StreamIdMismatch,
UnexpectedFrame,
}
impl std::fmt::Display for RecvStreamError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
RecvStreamError::Recv(e) => write!(f, "{e}"),
RecvStreamError::Parse(e) => write!(f, "{e}"),
RecvStreamError::PeerAborted { code, message } => {
write!(f, "peer aborted the stream: {code} ({message})")
}
RecvStreamError::StreamIdMismatch => {
write!(f, "received a frame for a different stream_id")
}
RecvStreamError::UnexpectedFrame => {
write!(f, "received a frame not valid in this context")
}
}
}
}
impl std::error::Error for RecvStreamError {}
impl StreamHandle {
pub async fn open(
session: &mut Session,
procedure: &str,
realm: [u8; 32],
mode: StreamMode,
args: Value,
deadline_ms: i128,
identity: &KeyPair,
) -> Result<Self, OpenError> {
let mut stream = session
.open_dedicated_stream()
.await
.map_err(OpenError::OpenStream)?;
let stream_id: [u8; 16] = rand::random();
let spec = frame::StreamOpenSpec::new(
stream_id,
procedure,
realm,
mode,
args,
deadline_ms,
identity.node_id(),
);
let signed = frame::sign(frame::stream_open(&spec), identity);
stream.send_frame(signed).await.map_err(OpenError::Send)?;
Ok(Self {
stream,
stream_id,
mode,
seq_out: 0,
})
}
pub async fn accept(
session: &mut Session,
timeout: Duration,
) -> Result<(Self, frame::StreamOpenInfo), AcceptError> {
let mut stream = tokio::time::timeout(timeout, session.accept_dedicated_stream())
.await
.map_err(|_| AcceptError::Timeout)?
.map_err(AcceptError::AcceptStream)?;
let first = stream.recv_frame().await.map_err(AcceptError::Recv)?;
let open = frame::parse_stream_open(&first).map_err(AcceptError::Parse)?;
let handle = Self {
stream,
stream_id: open.stream_id,
mode: open.mode,
seq_out: 0,
};
Ok((handle, open))
}
pub async fn send_reply(
&mut self,
payload: Value,
identity: &KeyPair,
) -> Result<(), SendFrameError> {
let spec = frame::StreamReplySpec::new(self.stream_id, payload, identity.node_id());
let signed = frame::sign(frame::stream_reply(&spec), identity);
self.stream.send_frame(signed).await
}
pub async fn send_data(
&mut self,
encoding: StreamEncoding,
body: Value,
identity: &KeyPair,
) -> Result<(), SendFrameError> {
let spec = frame::StreamDataSpec::new(
self.stream_id,
self.seq_out,
encoding,
body,
Some(identity.public_bytes()),
);
self.seq_out += 1;
let signed = frame::sign(frame::stream_data(&spec), identity);
self.stream.send_frame(signed).await
}
pub async fn close_send(&mut self, identity: &KeyPair) -> Result<(), SendFrameError> {
let spec = frame::StreamEndSpec::new(
self.stream_id,
StreamRole::Send,
Some(identity.public_bytes()),
);
let signed = frame::sign(frame::stream_end(&spec), identity);
self.stream.send_frame(signed).await
}
pub async fn recv(&mut self, timeout: Duration) -> Result<StreamItem, RecvStreamError> {
let value = self
.stream
.recv_frame_timeout(timeout)
.await
.map_err(RecvStreamError::Recv)?;
match frame::parse_stream_event(&value).map_err(RecvStreamError::Parse)? {
frame::StreamEvent::Data {
stream_id,
seq,
encoding,
body,
} => {
self.check_stream_id(stream_id)?;
Ok(StreamItem::Data {
seq,
encoding,
body,
})
}
frame::StreamEvent::End { stream_id, role: _ } => {
self.check_stream_id(stream_id)?;
Ok(StreamItem::Eof)
}
frame::StreamEvent::Error {
stream_id,
code,
message,
} => {
self.check_stream_id(stream_id)?;
Err(RecvStreamError::PeerAborted { code, message })
}
frame::StreamEvent::Reply { .. } => Err(RecvStreamError::UnexpectedFrame),
}
}
pub async fn await_reply(
&mut self,
timeout: Duration,
) -> Result<(Value, [u8; 32]), RecvStreamError> {
let value = self
.stream
.recv_frame_timeout(timeout)
.await
.map_err(RecvStreamError::Recv)?;
match frame::parse_stream_event(&value).map_err(RecvStreamError::Parse)? {
frame::StreamEvent::Reply {
stream_id,
payload,
responded_by,
} => {
self.check_stream_id(stream_id)?;
Ok((payload, responded_by))
}
frame::StreamEvent::Error {
stream_id,
code,
message,
} => {
self.check_stream_id(stream_id)?;
Err(RecvStreamError::PeerAborted { code, message })
}
frame::StreamEvent::Data { .. } | frame::StreamEvent::End { .. } => {
Err(RecvStreamError::UnexpectedFrame)
}
}
}
fn check_stream_id(&self, stream_id: [u8; 16]) -> Result<(), RecvStreamError> {
if stream_id == self.stream_id {
Ok(())
} else {
Err(RecvStreamError::StreamIdMismatch)
}
}
pub async fn abort(
mut self,
code: impl Into<String>,
message: impl Into<String>,
identity: &KeyPair,
) {
let spec = frame::StreamErrorSpec::new(
self.stream_id,
code,
message,
Some(identity.public_bytes()),
);
let signed = frame::sign(frame::stream_error(&spec), identity);
let _ = self.stream.send_frame(signed).await;
}
}