anvil_server/
ws.rs

1use crate::{PubSubRpcHandler, error::RequestError, pubsub::PubSubConnection};
2use anvil_rpc::request::Request;
3use axum::{
4    extract::{
5        State, WebSocketUpgrade,
6        ws::{Message, WebSocket},
7    },
8    response::Response,
9};
10use futures::{Sink, Stream, ready};
11use std::{
12    pin::Pin,
13    task::{Context, Poll},
14};
15
16#[pin_project::pin_project]
17struct SocketConn(#[pin] WebSocket);
18
19impl Stream for SocketConn {
20    type Item = Result<Option<Request>, RequestError>;
21
22    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
23        match ready!(self.project().0.poll_next(cx)) {
24            Some(msg) => Poll::Ready(Some(on_message(msg))),
25            _ => Poll::Ready(None),
26        }
27    }
28}
29
30impl Sink<String> for SocketConn {
31    type Error = axum::Error;
32
33    fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
34        self.project().0.poll_ready(cx)
35    }
36
37    fn start_send(self: Pin<&mut Self>, item: String) -> Result<(), Self::Error> {
38        self.project().0.start_send(Message::Text(item.into()))
39    }
40
41    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
42        self.project().0.poll_flush(cx)
43    }
44
45    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
46        self.project().0.poll_close(cx)
47    }
48}
49
50fn on_message(msg: Result<Message, axum::Error>) -> Result<Option<Request>, RequestError> {
51    match msg? {
52        Message::Text(text) => Ok(Some(serde_json::from_str(&text)?)),
53        Message::Binary(data) => {
54            // the binary payload type is the request as-is but as bytes, if this is a valid
55            // `Request` then we can deserialize the Json from the data Vec
56            Ok(Some(serde_json::from_slice(&data)?))
57        }
58        Message::Close(_) => {
59            trace!(target: "rpc::ws", "ws client disconnected");
60            Err(RequestError::Disconnect)
61        }
62        _ => Ok(None),
63    }
64}
65
66/// Handles incoming Websocket upgrade
67///
68/// This is the entrypoint invoked by the axum server for a websocket request
69pub async fn handle_ws<Http, Ws: PubSubRpcHandler>(
70    ws: WebSocketUpgrade,
71    State((_, handler)): State<(Http, Ws)>,
72) -> Response {
73    ws.on_upgrade(|socket| PubSubConnection::new(SocketConn(socket), handler))
74}
anvil_server/ws.rs

anvil_server/
ws.rs
