From 5a0395963fa64dd20044dabbd2a7187350500a28 Mon Sep 17 00:00:00 2001 From: Sean Matt Date: Mon, 15 Jul 2024 08:44:51 -0400 Subject: [PATCH] feat: implement generalizable rpc middleware (#9429) --- crates/rpc/rpc-builder/src/lib.rs | 97 +++++++++++++++++++++++-------- 1 file changed, 72 insertions(+), 25 deletions(-) diff --git a/crates/rpc/rpc-builder/src/lib.rs b/crates/rpc/rpc-builder/src/lib.rs index 8a6dce5ae6..f02afd9836 100644 --- a/crates/rpc/rpc-builder/src/lib.rs +++ b/crates/rpc/rpc-builder/src/lib.rs @@ -143,7 +143,10 @@ use error::{ConflictingModules, RpcError, ServerKind}; use http::{header::AUTHORIZATION, HeaderMap}; use jsonrpsee::{ core::RegisterMethodError, - server::{AlreadyStoppedError, IdProvider, RpcServiceBuilder, ServerHandle}, + server::{ + middleware::rpc::{RpcService, RpcServiceT}, + AlreadyStoppedError, IdProvider, RpcServiceBuilder, ServerHandle, + }, Methods, RpcModule, }; use reth_engine_primitives::EngineTypes; @@ -169,6 +172,7 @@ use reth_rpc_layer::{AuthLayer, Claims, JwtAuthValidator, JwtSecret}; use reth_tasks::{pool::BlockingTaskGuard, TaskSpawner, TokioTaskExecutor}; use reth_transaction_pool::{noop::NoopTransactionPool, TransactionPool}; use serde::{Deserialize, Serialize}; +use tower::Layer; use tower_http::cors::CorsLayer; use crate::{ @@ -1116,8 +1120,8 @@ where /// /// Once the [`RpcModule`] is built via [`RpcModuleBuilder`] the servers can be started, See also /// [`ServerBuilder::build`] and [`Server::start`](jsonrpsee::server::Server::start). -#[derive(Default, Debug)] -pub struct RpcServerConfig { +#[derive(Debug)] +pub struct RpcServerConfig { /// Configs for JSON-RPC Http. http_server_config: Option>, /// Allowed CORS Domains for http @@ -1136,10 +1140,31 @@ pub struct RpcServerConfig { ipc_endpoint: Option, /// JWT secret for authentication jwt_secret: Option, + /// Configurable RPC middleware + #[allow(dead_code)] + rpc_middleware: RpcServiceBuilder, } // === impl RpcServerConfig === +impl Default for RpcServerConfig { + /// Create a new config instance + fn default() -> Self { + Self { + http_server_config: None, + http_cors_domains: None, + http_addr: None, + ws_server_config: None, + ws_cors_domains: None, + ws_addr: None, + ipc_server_config: None, + ipc_endpoint: None, + jwt_secret: None, + rpc_middleware: RpcServiceBuilder::new(), + } + } +} + impl RpcServerConfig { /// Creates a new config with only http set pub fn http(config: ServerBuilder) -> Self { @@ -1166,29 +1191,56 @@ impl RpcServerConfig { self } + /// Configures the ws server + /// + /// Note: this always configures an [`EthSubscriptionIdProvider`] [`IdProvider`] for + /// convenience. To set a custom [`IdProvider`], please use [`Self::with_id_provider`]. + pub fn with_ws(mut self, config: ServerBuilder) -> Self { + self.ws_server_config = Some(config.set_id_provider(EthSubscriptionIdProvider::default())); + self + } + + /// Configures the ipc server + /// + /// Note: this always configures an [`EthSubscriptionIdProvider`] [`IdProvider`] for + /// convenience. To set a custom [`IdProvider`], please use [`Self::with_id_provider`]. + pub fn with_ipc(mut self, config: IpcServerBuilder) -> Self { + self.ipc_server_config = Some(config.set_id_provider(EthSubscriptionIdProvider::default())); + self + } +} + +impl RpcServerConfig { + /// Configure rpc middleware + pub fn set_rpc_middleware(self, rpc_middleware: RpcServiceBuilder) -> RpcServerConfig { + RpcServerConfig { + http_server_config: self.http_server_config, + http_cors_domains: self.http_cors_domains, + http_addr: self.http_addr, + ws_server_config: self.ws_server_config, + ws_cors_domains: self.ws_cors_domains, + ws_addr: self.ws_addr, + ipc_server_config: self.ipc_server_config, + ipc_endpoint: self.ipc_endpoint, + jwt_secret: self.jwt_secret, + rpc_middleware, + } + } + /// Configure the cors domains for http _and_ ws pub fn with_cors(self, cors_domain: Option) -> Self { self.with_http_cors(cors_domain.clone()).with_ws_cors(cors_domain) } - /// Configure the cors domains for HTTP - pub fn with_http_cors(mut self, cors_domain: Option) -> Self { - self.http_cors_domains = cors_domain; - self - } - /// Configure the cors domains for WS pub fn with_ws_cors(mut self, cors_domain: Option) -> Self { self.ws_cors_domains = cors_domain; self } - /// Configures the ws server - /// - /// Note: this always configures an [`EthSubscriptionIdProvider`] [`IdProvider`] for - /// convenience. To set a custom [`IdProvider`], please use [`Self::with_id_provider`]. - pub fn with_ws(mut self, config: ServerBuilder) -> Self { - self.ws_server_config = Some(config.set_id_provider(EthSubscriptionIdProvider::default())); + /// Configure the cors domains for HTTP + pub fn with_http_cors(mut self, cors_domain: Option) -> Self { + self.http_cors_domains = cors_domain; self } @@ -1210,15 +1262,6 @@ impl RpcServerConfig { self } - /// Configures the ipc server - /// - /// Note: this always configures an [`EthSubscriptionIdProvider`] [`IdProvider`] for - /// convenience. To set a custom [`IdProvider`], please use [`Self::with_id_provider`]. - pub fn with_ipc(mut self, config: IpcServerBuilder) -> Self { - self.ipc_server_config = Some(config.set_id_provider(EthSubscriptionIdProvider::default())); - self - } - /// Sets a custom [`IdProvider`] for all configured transports. /// /// By default all transports use [`EthSubscriptionIdProvider`] @@ -1292,7 +1335,11 @@ impl RpcServerConfig { /// If both http and ws are on the same port, they are combined into one server. /// /// Returns the [`RpcServerHandle`] with the handle to the started servers. - pub async fn start(self, modules: &TransportRpcModules) -> Result { + pub async fn start(self, modules: &TransportRpcModules) -> Result + where + RpcMiddleware: for<'a> Layer> + Clone + Send + 'static, + >::Service: Send + std::marker::Sync, + { let mut http_handle = None; let mut ws_handle = None; let mut ipc_handle = None;