leyline/core/websocket/
split.rs

1use futures_util::stream::{SplitSink, SplitStream};
2use futures_util::{SinkExt, StreamExt};
3use tokio_tungstenite::WebSocketStream;
4use tokio_tungstenite::tungstenite::Message;
5
6use super::WsMessage;
7use super::handshake::wire_error;
8use crate::core::error::Result;
9use crate::h2::client::H2ConnectStream;
10use crate::tls::TlsIo;
11
12pub(super) enum WsSinkInner {
13    H1(SplitSink<WebSocketStream<TlsIo>, Message>),
14    H2(SplitSink<WebSocketStream<H2ConnectStream>, Message>),
15}
16
17pub struct WsSink {
18    pub(super) inner: WsSinkInner,
19}
20
21impl WsSink {
22    pub async fn send(&mut self, msg: WsMessage) -> Result<()> {
23        let msg = msg.into_wire();
24        match &mut self.inner {
25            WsSinkInner::H1(s) => s.send(msg).await,
26            WsSinkInner::H2(s) => s.send(msg).await,
27        }
28        .map_err(|e| wire_error("ws send", e))
29    }
30
31    pub async fn close(&mut self) -> Result<()> {
32        match &mut self.inner {
33            WsSinkInner::H1(s) => s.close().await,
34            WsSinkInner::H2(s) => s.close().await,
35        }
36        .map_err(|e| wire_error("ws close", e))
37    }
38}
39
40pub(super) enum WsStreamInner {
41    H1(SplitStream<WebSocketStream<TlsIo>>),
42    H2(SplitStream<WebSocketStream<H2ConnectStream>>),
43}
44
45pub struct WsStream {
46    pub(super) inner: WsStreamInner,
47}
48
49impl WsStream {
50    pub async fn recv(&mut self) -> Result<Option<WsMessage>> {
51        let next = match &mut self.inner {
52            WsStreamInner::H1(s) => s.next().await,
53            WsStreamInner::H2(s) => s.next().await,
54        };
55        match next {
56            Some(Ok(msg)) => Ok(Some(WsMessage::wire(msg))),
57            Some(Err(e)) => Err(wire_error("ws recv", e)),
58            None => Ok(None),
59        }
60    }
61}
leyline/core/websocket/split.rs

leyline/core/websocket/
split.rs
