65 lines
2.0 KiB
Rust
65 lines
2.0 KiB
Rust
use async_tungstenite::tungstenite::{Error as WebSocketError, Message as WebSocketMessage};
|
|
use std::{
|
|
io,
|
|
pin::Pin,
|
|
task::{Context, Poll},
|
|
};
|
|
|
|
pub struct Channel {
|
|
tx: futures::channel::mpsc::UnboundedSender<WebSocketMessage>,
|
|
rx: futures::channel::mpsc::UnboundedReceiver<WebSocketMessage>,
|
|
}
|
|
|
|
impl Channel {
|
|
pub fn new() -> Self {
|
|
let (tx, rx) = futures::channel::mpsc::unbounded();
|
|
Self { tx, rx }
|
|
}
|
|
|
|
pub fn bidirectional() -> (Self, Self) {
|
|
let (a_tx, a_rx) = futures::channel::mpsc::unbounded();
|
|
let (b_tx, b_rx) = futures::channel::mpsc::unbounded();
|
|
let a = Self { tx: a_tx, rx: b_rx };
|
|
let b = Self { tx: b_tx, rx: a_rx };
|
|
(a, b)
|
|
}
|
|
}
|
|
|
|
impl futures::Sink<WebSocketMessage> for Channel {
|
|
type Error = WebSocketError;
|
|
|
|
fn poll_ready(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
|
|
Pin::new(&mut self.tx)
|
|
.poll_ready(cx)
|
|
.map_err(|err| io::Error::new(io::ErrorKind::Other, err).into())
|
|
}
|
|
|
|
fn start_send(mut self: Pin<&mut Self>, item: WebSocketMessage) -> Result<(), Self::Error> {
|
|
Pin::new(&mut self.tx)
|
|
.start_send(item)
|
|
.map_err(|err| io::Error::new(io::ErrorKind::Other, err).into())
|
|
}
|
|
|
|
fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
|
|
Pin::new(&mut self.tx)
|
|
.poll_flush(cx)
|
|
.map_err(|err| io::Error::new(io::ErrorKind::Other, err).into())
|
|
}
|
|
|
|
fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
|
|
Pin::new(&mut self.tx)
|
|
.poll_close(cx)
|
|
.map_err(|err| io::Error::new(io::ErrorKind::Other, err).into())
|
|
}
|
|
}
|
|
|
|
impl futures::Stream for Channel {
|
|
type Item = Result<WebSocketMessage, WebSocketError>;
|
|
|
|
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
|
|
Pin::new(&mut self.rx)
|
|
.poll_next(cx)
|
|
.map(|i| i.map(|i| Ok(i)))
|
|
}
|
|
}
|