web3-proxy/web3_proxy/src/frontend/rpc_proxy_ws.rs

560 lines
19 KiB
Rust
Raw Normal View History

2022-10-18 00:47:58 +03:00
//! Take a user's WebSocket JSON-RPC requests and either respond from local data or proxy the request to a backend rpc server.
//!
//! WebSockets are the preferred method of receiving requests, but not all clients have good support.
2022-11-20 01:05:51 +03:00
use super::authorization::{ip_is_authorized, key_is_authorized, Authorization, RequestMetadata};
use super::errors::{Web3ProxyError, Web3ProxyResponse};
use crate::jsonrpc::JsonRpcId;
use crate::stats::RpcQueryStats;
2022-11-20 01:05:51 +03:00
use crate::{
app::Web3ProxyApp,
2023-03-20 05:14:46 +03:00
frontend::errors::Web3ProxyResult,
2022-11-20 01:05:51 +03:00
jsonrpc::{JsonRpcForwardedResponse, JsonRpcForwardedResponseEnum, JsonRpcRequest},
};
2022-10-20 23:26:14 +03:00
use axum::headers::{Origin, Referer, UserAgent};
2022-05-20 08:30:54 +03:00
use axum::{
2022-05-29 20:28:41 +03:00
extract::ws::{Message, WebSocket, WebSocketUpgrade},
2022-08-06 04:17:25 +03:00
extract::Path,
response::{IntoResponse, Redirect},
Extension, TypedHeader,
2022-05-20 08:30:54 +03:00
};
use axum_client_ip::InsecureClientIp;
use axum_macros::debug_handler;
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},
};
2022-08-12 22:07:14 +03:00
use handlebars::Handlebars;
use hashbrown::HashMap;
use http::StatusCode;
2023-03-20 05:14:46 +03:00
use log::{info, trace, warn};
2022-12-24 04:32:58 +03:00
use serde_json::json;
2022-08-27 05:13:36 +03:00
use std::sync::Arc;
2022-07-09 01:14:45 +03:00
use std::{str::from_utf8_mut, sync::atomic::AtomicUsize};
2023-01-19 03:17:43 +03:00
use tokio::sync::{broadcast, OwnedSemaphorePermit, RwLock};
2022-06-16 05:53:37 +03:00
2023-03-03 04:39:50 +03:00
#[derive(Copy, Clone, Debug)]
2023-01-17 09:54:40 +03:00
pub enum ProxyMode {
/// send to the "best" synced server
Best,
/// send to all synced servers and return the fastest non-error response (reverts do not count as errors here)
Fastest(usize),
/// send to all servers for benchmarking. return the fastest non-error response
Versus,
2023-03-03 04:39:50 +03:00
/// send all requests and responses to kafka
Debug,
}
impl Default for ProxyMode {
fn default() -> Self {
Self::Best
}
2023-01-17 09:54:40 +03:00
}
2022-10-18 00:47:58 +03:00
/// Public entrypoint for WebSocket JSON-RPC requests.
2023-01-17 09:54:40 +03:00
/// Queries a single server at a time
#[debug_handler]
2022-09-24 08:53:45 +03:00
pub async fn websocket_handler(
2022-07-07 06:22:09 +03:00
Extension(app): Extension<Arc<Web3ProxyApp>>,
ip: InsecureClientIp,
2023-01-17 09:54:40 +03:00
origin: Option<TypedHeader<Origin>>,
ws_upgrade: Option<WebSocketUpgrade>,
) -> Web3ProxyResponse {
2023-01-19 03:17:43 +03:00
_websocket_handler(ProxyMode::Best, app, ip, origin, ws_upgrade).await
2023-01-17 09:54:40 +03:00
}
/// Public entrypoint for WebSocket JSON-RPC requests that uses all synced servers.
/// Queries all synced backends with every request! This might get expensive!
#[debug_handler]
pub async fn fastest_websocket_handler(
Extension(app): Extension<Arc<Web3ProxyApp>>,
ip: InsecureClientIp,
2023-01-17 09:54:40 +03:00
origin: Option<TypedHeader<Origin>>,
ws_upgrade: Option<WebSocketUpgrade>,
) -> Web3ProxyResponse {
2023-01-17 09:54:40 +03:00
// TODO: get the fastest number from the url params (default to 0/all)
// TODO: config to disable this
_websocket_handler(ProxyMode::Fastest(0), app, ip, origin, ws_upgrade).await
}
/// Public entrypoint for WebSocket JSON-RPC requests that uses all synced servers.
/// Queries **all** backends with every request! This might get expensive!
#[debug_handler]
pub async fn versus_websocket_handler(
Extension(app): Extension<Arc<Web3ProxyApp>>,
ip: InsecureClientIp,
2023-01-17 09:54:40 +03:00
origin: Option<TypedHeader<Origin>>,
ws_upgrade: Option<WebSocketUpgrade>,
) -> Web3ProxyResponse {
2023-01-17 09:54:40 +03:00
// TODO: config to disable this
_websocket_handler(ProxyMode::Versus, app, ip, origin, ws_upgrade).await
}
async fn _websocket_handler(
proxy_mode: ProxyMode,
app: Arc<Web3ProxyApp>,
InsecureClientIp(ip): InsecureClientIp,
2022-10-21 23:59:05 +03:00
origin: Option<TypedHeader<Origin>>,
2022-08-10 05:37:34 +03:00
ws_upgrade: Option<WebSocketUpgrade>,
) -> Web3ProxyResponse {
let origin = origin.map(|x| x.0);
2023-03-03 04:39:50 +03:00
let (authorization, _semaphore) = ip_is_authorized(&app, ip, origin, proxy_mode).await?;
let authorization = Arc::new(authorization);
2022-08-11 03:16:13 +03:00
match ws_upgrade {
2022-08-21 12:39:38 +03:00
Some(ws) => Ok(ws
2023-03-03 04:39:50 +03:00
.on_upgrade(move |socket| proxy_web3_socket(app, authorization, socket))
2022-08-21 12:39:38 +03:00
.into_response()),
None => {
2022-10-18 00:47:58 +03:00
if let Some(redirect) = &app.config.redirect_public_url {
// this is not a websocket. redirect to a friendly page
2022-11-28 22:59:42 +03:00
Ok(Redirect::permanent(redirect).into_response())
2022-10-18 00:47:58 +03:00
} else {
2023-03-20 05:14:46 +03:00
Err(Web3ProxyError::WebsocketOnly)
2022-10-18 00:47:58 +03:00
}
}
}
2022-08-04 04:10:27 +03:00
}
2022-10-18 00:47:58 +03:00
/// Authenticated entrypoint for WebSocket JSON-RPC requests. Web3 wallets use this.
/// Rate limit and billing based on the api key in the url.
/// Can optionally authorized based on origin, referer, or user agent.
#[debug_handler]
2022-09-24 08:53:45 +03:00
pub async fn websocket_handler_with_key(
2022-08-04 04:10:27 +03:00
Extension(app): Extension<Arc<Web3ProxyApp>>,
ip: InsecureClientIp,
2023-01-17 09:54:40 +03:00
Path(rpc_key): Path<String>,
origin: Option<TypedHeader<Origin>>,
referer: Option<TypedHeader<Referer>>,
user_agent: Option<TypedHeader<UserAgent>>,
ws_upgrade: Option<WebSocketUpgrade>,
) -> Web3ProxyResponse {
2023-01-17 09:54:40 +03:00
_websocket_handler_with_key(
2023-01-20 05:32:31 +03:00
ProxyMode::Best,
2023-01-17 09:54:40 +03:00
app,
ip,
rpc_key,
origin,
referer,
user_agent,
ws_upgrade,
)
.await
}
2023-03-03 04:39:50 +03:00
#[debug_handler]
pub async fn debug_websocket_handler_with_key(
Extension(app): Extension<Arc<Web3ProxyApp>>,
ip: InsecureClientIp,
Path(rpc_key): Path<String>,
origin: Option<TypedHeader<Origin>>,
referer: Option<TypedHeader<Referer>>,
user_agent: Option<TypedHeader<UserAgent>>,
ws_upgrade: Option<WebSocketUpgrade>,
) -> Web3ProxyResponse {
2023-03-03 04:39:50 +03:00
_websocket_handler_with_key(
ProxyMode::Debug,
app,
ip,
rpc_key,
origin,
referer,
user_agent,
ws_upgrade,
)
.await
}
2023-01-17 09:54:40 +03:00
#[debug_handler]
pub async fn fastest_websocket_handler_with_key(
Extension(app): Extension<Arc<Web3ProxyApp>>,
ip: InsecureClientIp,
2022-10-27 03:12:42 +03:00
Path(rpc_key): Path<String>,
origin: Option<TypedHeader<Origin>>,
referer: Option<TypedHeader<Referer>>,
user_agent: Option<TypedHeader<UserAgent>>,
2022-08-11 03:16:13 +03:00
ws_upgrade: Option<WebSocketUpgrade>,
) -> Web3ProxyResponse {
2023-01-17 09:54:40 +03:00
// TODO: get the fastest number from the url params (default to 0/all)
_websocket_handler_with_key(
ProxyMode::Fastest(0),
app,
ip,
rpc_key,
origin,
referer,
user_agent,
ws_upgrade,
)
.await
}
#[debug_handler]
pub async fn versus_websocket_handler_with_key(
Extension(app): Extension<Arc<Web3ProxyApp>>,
ip: InsecureClientIp,
2023-01-17 09:54:40 +03:00
Path(rpc_key): Path<String>,
origin: Option<TypedHeader<Origin>>,
referer: Option<TypedHeader<Referer>>,
user_agent: Option<TypedHeader<UserAgent>>,
ws_upgrade: Option<WebSocketUpgrade>,
) -> Web3ProxyResponse {
2023-01-17 09:54:40 +03:00
_websocket_handler_with_key(
ProxyMode::Versus,
app,
ip,
rpc_key,
origin,
referer,
user_agent,
ws_upgrade,
)
.await
}
#[allow(clippy::too_many_arguments)]
async fn _websocket_handler_with_key(
proxy_mode: ProxyMode,
app: Arc<Web3ProxyApp>,
InsecureClientIp(ip): InsecureClientIp,
2023-01-17 09:54:40 +03:00
rpc_key: String,
origin: Option<TypedHeader<Origin>>,
referer: Option<TypedHeader<Referer>>,
user_agent: Option<TypedHeader<UserAgent>>,
ws_upgrade: Option<WebSocketUpgrade>,
) -> Web3ProxyResponse {
2022-10-27 03:12:42 +03:00
let rpc_key = rpc_key.parse()?;
2022-09-24 08:53:45 +03:00
let (authorization, _semaphore) = key_is_authorized(
&app,
2022-10-27 03:12:42 +03:00
rpc_key,
ip,
origin.map(|x| x.0),
2023-03-03 04:39:50 +03:00
proxy_mode,
referer.map(|x| x.0),
user_agent.map(|x| x.0),
)
.await?;
2022-08-04 04:10:27 +03:00
2022-11-24 14:04:10 +03:00
trace!("websocket_handler_with_key {:?}", authorization);
2022-11-21 01:37:12 +03:00
let authorization = Arc::new(authorization);
2022-08-11 03:16:13 +03:00
match ws_upgrade {
2023-03-03 04:39:50 +03:00
Some(ws_upgrade) => {
Ok(ws_upgrade.on_upgrade(move |socket| proxy_web3_socket(app, authorization, socket)))
}
2022-08-11 03:16:13 +03:00
None => {
// if no websocket upgrade, this is probably a user loading the url with their browser
2022-11-21 01:37:12 +03:00
// TODO: rate limit here? key_is_authorized might be enough
match (
&app.config.redirect_public_url,
&app.config.redirect_rpc_key_url,
2023-01-19 03:17:43 +03:00
authorization.checks.rpc_secret_key_id,
) {
(None, None, _) => Err(Web3ProxyError::StatusCode(
2022-11-21 01:37:12 +03:00
StatusCode::BAD_REQUEST,
"this page is for rpcs".to_string(),
None,
)),
(Some(redirect_public_url), _, None) => {
2022-11-28 22:59:42 +03:00
Ok(Redirect::permanent(redirect_public_url).into_response())
2022-10-21 23:59:05 +03:00
}
(_, Some(redirect_rpc_key_url), rpc_key_id) => {
let reg = Handlebars::new();
2023-01-19 03:17:43 +03:00
if authorization.checks.rpc_secret_key_id.is_none() {
2022-11-21 01:37:12 +03:00
// i don't think this is possible
Err(Web3ProxyError::StatusCode(
2022-11-21 01:37:12 +03:00
StatusCode::UNAUTHORIZED,
"AUTHORIZATION header required".to_string(),
None,
))
} else {
let redirect_rpc_key_url = reg
.render_template(
redirect_rpc_key_url,
&json!({ "rpc_key_id": rpc_key_id }),
)
.expect("templating should always work");
// this is not a websocket. redirect to a page for this user
2022-11-28 22:59:42 +03:00
Ok(Redirect::permanent(&redirect_rpc_key_url).into_response())
}
}
// any other combinations get a simple error
_ => Err(Web3ProxyError::StatusCode(
StatusCode::BAD_REQUEST,
"this page is for rpcs".to_string(),
None,
)),
2022-10-18 00:47:58 +03:00
}
2022-08-11 03:16:13 +03:00
}
}
2022-05-29 20:28:41 +03:00
}
async fn proxy_web3_socket(
app: Arc<Web3ProxyApp>,
authorization: Arc<Authorization>,
socket: WebSocket,
) {
2022-05-29 20:28:41 +03:00
// 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_sender, response_receiver) = flume::unbounded::<Message>();
2022-05-29 20:28:41 +03:00
tokio::spawn(write_web3_socket(response_receiver, ws_tx));
2023-03-03 04:39:50 +03:00
tokio::spawn(read_web3_socket(app, authorization, ws_rx, response_sender));
2022-05-29 20:28:41 +03:00
}
2022-07-25 03:27:00 +03:00
/// websockets support a few more methods than http clients
2022-05-31 04:55:04 +03:00
async fn handle_socket_payload(
2022-06-14 10:13:42 +03:00
app: Arc<Web3ProxyApp>,
authorization: &Arc<Authorization>,
2022-05-31 04:55:04 +03:00
payload: &str,
2022-07-09 01:14:45 +03:00
response_sender: &flume::Sender<Message>,
subscription_count: &AtomicUsize,
2023-01-19 03:17:43 +03:00
subscriptions: Arc<RwLock<HashMap<String, AbortHandle>>>,
) -> (Message, Option<OwnedSemaphorePermit>) {
let (authorization, semaphore) = match authorization.check_again(&app).await {
Ok((a, s)) => (a, s),
Err(err) => {
let (_, err) = err.into_response_parts();
let err = serde_json::to_string(&err).expect("to_string should always work here");
return (Message::Text(err), None);
}
};
// TODO: do any clients send batches over websockets?
2022-05-31 04:55:04 +03:00
let (id, response) = match serde_json::from_str::<JsonRpcRequest>(payload) {
Ok(json_request) => {
let id = json_request.id.clone();
2022-05-31 04:55:04 +03:00
2023-03-20 05:14:46 +03:00
let response: Web3ProxyResult<JsonRpcForwardedResponseEnum> = match &json_request.method
[..]
{
"eth_subscribe" => {
2023-01-17 09:54:40 +03:00
// TODO: how can we subscribe with proxy_mode?
2022-11-20 01:05:51 +03:00
match app
.eth_subscribe(
authorization.clone(),
json_request,
subscription_count,
response_sender.clone(),
)
2022-11-20 01:05:51 +03:00
.await
{
Ok((handle, response)) => {
// TODO: better key
2023-01-19 03:17:43 +03:00
let mut x = subscriptions.write().await;
x.insert(
2022-09-30 07:18:18 +03:00
response
.result
.as_ref()
// TODO: what if there is an error?
.expect("response should always have a result, not an error")
.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" => {
2022-11-20 01:05:51 +03:00
// TODO: move this logic into the app?
let request_bytes = json_request.num_bytes();
let request_metadata = Arc::new(RequestMetadata::new(request_bytes));
2022-11-20 01:05:51 +03:00
let subscription_id = json_request.params.unwrap().to_string();
2022-05-31 04:55:04 +03:00
2023-01-19 03:17:43 +03:00
let mut x = subscriptions.write().await;
2022-11-20 01:05:51 +03:00
// TODO: is this the right response?
2023-01-19 03:17:43 +03:00
let partial_response = match x.remove(&subscription_id) {
None => false,
Some(handle) => {
handle.abort();
true
}
};
2022-05-31 04:55:04 +03:00
2023-01-19 03:17:43 +03:00
drop(x);
2022-07-22 22:30:39 +03:00
let response =
JsonRpcForwardedResponse::from_value(json!(partial_response), id.clone());
2022-05-31 04:55:04 +03:00
2022-11-20 01:05:51 +03:00
if let Some(stat_sender) = app.stat_sender.as_ref() {
let response_stat = RpcQueryStats::new(
Some(json_request.method.clone()),
2022-11-20 01:05:51 +03:00
authorization.clone(),
request_metadata,
response.num_bytes(),
);
if let Err(err) = stat_sender.send_async(response_stat.into()).await {
// TODO: what should we do?
warn!("stat_sender failed during eth_unsubscribe: {:?}", err);
}
}
Ok(response.into())
}
_ => app
2023-03-03 04:39:50 +03:00
.proxy_web3_rpc(authorization.clone(), json_request.into())
.await
2023-03-20 05:14:46 +03:00
.map(|(response, _)| response),
2022-05-31 04:55:04 +03:00
};
(id, response)
}
Err(err) => {
let id = JsonRpcId::None.to_raw_value();
2022-05-31 04:55:04 +03:00
(id, Err(err.into()))
}
};
let response_str = match response {
Ok(x) => serde_json::to_string(&x).expect("to_string should always work here"),
2022-05-31 04:55:04 +03:00
Err(err) => {
2023-03-20 05:14:46 +03:00
let (_, mut response) = err.into_response_parts();
response.id = id;
serde_json::to_string(&response).expect("to_string should always work here")
2022-05-31 04:55:04 +03:00
}
};
2023-01-19 03:17:43 +03:00
(Message::Text(response_str), semaphore)
2022-05-31 04:55:04 +03:00
}
2022-05-29 20:28:41 +03:00
async fn read_web3_socket(
2022-07-07 06:22:09 +03:00
app: Arc<Web3ProxyApp>,
authorization: Arc<Authorization>,
2022-05-29 20:28:41 +03:00
mut ws_rx: SplitStream<WebSocket>,
2022-07-09 01:14:45 +03:00
response_sender: flume::Sender<Message>,
2022-05-29 20:28:41 +03:00
) {
2023-01-19 03:17:43 +03:00
// TODO: need a concurrent hashmap
let subscriptions = Arc::new(RwLock::new(HashMap::new()));
let subscription_count = Arc::new(AtomicUsize::new(1));
let (close_sender, mut close_receiver) = broadcast::channel(1);
loop {
tokio::select! {
msg = ws_rx.next() => {
if let Some(Ok(msg)) = msg {
// spawn so that we can serve responses from this loop even faster
// TODO: only do these clones if the msg is text/binary?
let close_sender = close_sender.clone();
let app = app.clone();
let authorization = authorization.clone();
let response_sender = response_sender.clone();
let subscriptions = subscriptions.clone();
let subscription_count = subscription_count.clone();
let f = async move {
let mut _semaphore = None;
// new message from our client. forward to a backend and then send it through response_tx
let response_msg = match msg {
Message::Text(payload) => {
let (msg, s) = handle_socket_payload(
app.clone(),
&authorization,
&payload,
&response_sender,
&subscription_count,
subscriptions,
)
.await;
_semaphore = s;
msg
}
Message::Ping(x) => {
trace!("ping: {:?}", x);
Message::Pong(x)
}
Message::Pong(x) => {
trace!("pong: {:?}", x);
return;
}
Message::Close(_) => {
info!("closing websocket connection");
// TODO: do something to close subscriptions?
let _ = close_sender.send(true);
return;
}
Message::Binary(mut payload) => {
let payload = from_utf8_mut(&mut payload).unwrap();
let (msg, s) = handle_socket_payload(
app.clone(),
&authorization,
payload,
&response_sender,
&subscription_count,
subscriptions,
)
.await;
_semaphore = s;
msg
}
};
2022-05-29 20:28:41 +03:00
2023-01-19 03:17:43 +03:00
if response_sender.send_async(response_msg).await.is_err() {
let _ = close_sender.send(true);
return;
};
_semaphore = None;
};
tokio::spawn(f);
} else {
break;
}
}
_ = close_receiver.recv() => {
2022-05-30 04:28:22 +03:00
break;
}
2023-01-19 03:17:43 +03:00
}
2022-05-29 20:28:41 +03:00
}
}
async fn write_web3_socket(
response_rx: flume::Receiver<Message>,
mut ws_tx: SplitSink<WebSocket, Message>,
) {
2022-07-09 01:14:45 +03:00
// TODO: increment counter for open websockets
2022-11-16 11:34:17 +03:00
// TODO: is there any way to make this stream receive.
2022-05-29 20:28:41 +03:00
while let Ok(msg) = response_rx.recv_async().await {
2022-08-11 03:16:13 +03:00
// a response is ready
// TODO: poke rate limits for this user?
// forward the response to through the websocket
2022-06-16 05:53:37 +03:00
if let Err(err) = ws_tx.send(msg).await {
2022-11-16 10:19:56 +03:00
// this is common. it happens whenever a client disconnects
trace!("unable to write to websocket: {:?}", err);
2022-05-29 20:28:41 +03:00
break;
};
}
2022-07-09 01:14:45 +03:00
// TODO: decrement counter for open websockets
2022-05-29 20:28:41 +03:00
}