web3-proxy/web3-proxy/src/frontend/ws_proxy.rs

160 lines
5.0 KiB
Rust
Raw Normal View History

2022-06-05 22:58:47 +03:00
use crate::{
app::Web3ProxyApp,
jsonrpc::{JsonRpcForwardedResponse, JsonRpcForwardedResponseEnum, JsonRpcRequest},
};
2022-05-20 08:30:54 +03:00
use axum::{
2022-05-29 20:28:41 +03:00
extract::ws::{Message, WebSocket, WebSocketUpgrade},
2022-05-20 08:30:54 +03:00
response::IntoResponse,
2022-06-05 22:58:47 +03:00
Extension,
2022-05-20 08:30:54 +03:00
};
2022-05-29 20:28:41 +03:00
use futures::SinkExt;
2022-05-31 04:55:04 +03:00
use futures::{
future::AbortHandle,
stream::{SplitSink, SplitStream, StreamExt},
};
use hashbrown::HashMap;
2022-05-20 08:30:54 +03:00
use serde_json::value::RawValue;
2022-05-31 04:55:04 +03:00
use std::str::from_utf8_mut;
2022-05-20 08:30:54 +03:00
use std::sync::Arc;
2022-06-05 22:58:47 +03:00
use tracing::{debug, error, info};
2022-05-20 08:30:54 +03:00
2022-06-05 22:58:47 +03:00
pub async fn websocket_handler(
2022-05-29 20:28:41 +03:00
app: Extension<Arc<Web3ProxyApp>>,
ws: WebSocketUpgrade,
) -> impl IntoResponse {
ws.on_upgrade(|socket| proxy_web3_socket(app, socket))
}
async fn proxy_web3_socket(app: Extension<Arc<Web3ProxyApp>>, socket: WebSocket) {
// split the websocket so we can read and write concurrently
let (ws_tx, ws_rx) = socket.split();
// create a channel for our reader and writer can communicate. todo: benchmark different channels
let (response_tx, response_rx) = flume::unbounded::<Message>();
tokio::spawn(write_web3_socket(response_rx, ws_tx));
tokio::spawn(read_web3_socket(app, ws_rx, response_tx));
}
2022-05-31 04:55:04 +03:00
async fn handle_socket_payload(
app: &Web3ProxyApp,
payload: &str,
response_tx: &flume::Sender<Message>,
subscriptions: &mut HashMap<String, AbortHandle>,
) -> Message {
let (id, response) = match serde_json::from_str::<JsonRpcRequest>(payload) {
Ok(payload) => {
let id = payload.id.clone();
let response: anyhow::Result<JsonRpcForwardedResponseEnum> = match &payload.method[..] {
"eth_subscribe" => {
let response = app.eth_subscribe(payload, response_tx.clone()).await;
2022-05-31 04:55:04 +03:00
match response {
Ok((handle, response)) => {
// TODO: better key
subscriptions
.insert(response.result.as_ref().unwrap().to_string(), handle);
2022-05-31 04:55:04 +03:00
Ok(response.into())
}
Err(err) => Err(err),
2022-05-31 04:55:04 +03:00
}
}
"eth_unsubscribe" => {
let subscription_id = payload.params.unwrap().to_string();
2022-05-31 04:55:04 +03:00
let partial_response = match subscriptions.remove(&subscription_id) {
None => "false",
Some(handle) => {
handle.abort();
"true"
}
};
2022-05-31 04:55:04 +03:00
let response = JsonRpcForwardedResponse::from_string(
partial_response.to_string(),
id.clone(),
);
2022-05-31 04:55:04 +03:00
Ok(response.into())
}
_ => app.proxy_web3_rpc(payload.into()).await,
2022-05-31 04:55:04 +03:00
};
(id, response)
}
Err(err) => {
// TODO: what should this id be?
2022-06-06 01:39:44 +03:00
let id = RawValue::from_string("null".to_string()).unwrap();
2022-05-31 04:55:04 +03:00
(id, Err(err.into()))
}
};
let response_str = match response {
Ok(x) => serde_json::to_string(&x),
Err(err) => {
// we have an anyhow error. turn it into
let response = JsonRpcForwardedResponse::from_anyhow_error(err, id);
serde_json::to_string(&response)
}
}
.unwrap();
Message::Text(response_str)
}
2022-05-29 20:28:41 +03:00
async fn read_web3_socket(
app: Extension<Arc<Web3ProxyApp>>,
mut ws_rx: SplitStream<WebSocket>,
response_tx: flume::Sender<Message>,
) {
let mut subscriptions = HashMap::new();
2022-05-29 20:28:41 +03:00
while let Some(Ok(msg)) = ws_rx.next().await {
// new message from our client. forward to a backend and then send it through response_tx
let response_msg = match msg {
Message::Text(payload) => {
2022-05-31 04:55:04 +03:00
handle_socket_payload(app.0.as_ref(), &payload, &response_tx, &mut subscriptions)
.await
2022-05-29 20:28:41 +03:00
}
Message::Ping(x) => Message::Pong(x),
2022-05-31 04:55:04 +03:00
Message::Pong(x) => {
debug!("pong: {:?}", x);
continue;
}
Message::Close(_) => {
info!("closing websocket connection");
break;
}
Message::Binary(mut payload) => {
let payload = from_utf8_mut(&mut payload).unwrap();
handle_socket_payload(app.0.as_ref(), payload, &response_tx, &mut subscriptions)
.await
}
2022-05-29 20:28:41 +03:00
};
2022-05-30 04:28:22 +03:00
match response_tx.send_async(response_msg).await {
Ok(_) => {}
Err(err) => {
error!("{}", err);
break;
}
2022-05-29 20:28:41 +03:00
};
}
}
async fn write_web3_socket(
response_rx: flume::Receiver<Message>,
mut ws_tx: SplitSink<WebSocket, Message>,
) {
while let Ok(msg) = response_rx.recv_async().await {
// a response is ready. write it to ws_tx
if ws_tx.send(msg).await.is_err() {
// TODO: log the error
break;
};
}
}