test: add basic auth server tests (#2278)

crates/rpc/rpc-builder/Cargo.toml

@@ -39,6 +39,8 @@ reth-transaction-pool = { path = "../../transaction-pool", features = ["test-uti
 reth-provider = { path = "../../storage/provider", features = ["test-utils"] }
 reth-network-api = { path = "../../net/network-api", features = ["test-utils"] }
 reth-interfaces = { path = "../../interfaces", features = ["test-utils"] }
+reth-beacon-consensus = { path = "../../consensus/beacon" }
+reth-payload-builder = { path = "../../payload/builder", features = ["test-utils"] }
 
 tokio = { version = "1", features = ["rt", "rt-multi-thread"] }
 serde_json = "1.0.94"

crates/rpc/rpc-builder/src/auth.rs

@@ -2,17 +2,24 @@ use crate::{
     constants,
     error::{RpcError, ServerKind},
 };
+use hyper::header::AUTHORIZATION;
 pub use jsonrpsee::server::ServerBuilder;
-use jsonrpsee::server::{RpcModule, ServerHandle};
+use jsonrpsee::{
+    http_client::HeaderMap,
+    server::{RpcModule, ServerHandle},
+};
 use reth_network_api::{NetworkInfo, Peers};
 use reth_provider::{BlockProvider, EvmEnvProvider, HeaderProvider, StateProviderFactory};
 use reth_rpc::{
-    eth::cache::EthStateCache, AuthLayer, EthApi, EthFilter, JwtAuthValidator, JwtSecret,
+    eth::cache::EthStateCache, AuthLayer, Claims, EthApi, EthFilter, JwtAuthValidator, JwtSecret,
 };
 use reth_rpc_api::{servers::*, EngineApiServer};
 use reth_tasks::TaskSpawner;
 use reth_transaction_pool::TransactionPool;
-use std::net::{IpAddr, Ipv4Addr, SocketAddr};
+use std::{
+    net::{IpAddr, Ipv4Addr, SocketAddr},
+    time::{Duration, SystemTime, UNIX_EPOCH},
+};
 
 /// Configure and launch a _standalone_ auth server with `engine` and a _new_ `eth` namespace.
 #[allow(clippy::too_many_arguments)]
@@ -24,7 +31,7 @@ pub async fn launch<Client, Pool, Network, Tasks, EngineApi>(
     engine_api: EngineApi,
     socket_addr: SocketAddr,
     secret: JwtSecret,
-) -> Result<ServerHandle, RpcError>
+) -> Result<AuthServerHandle, RpcError>
 where
     Client: BlockProvider
         + HeaderProvider
@@ -52,7 +59,7 @@ pub async fn launch_with_eth_api<Client, Pool, Network, EngineApi>(
     engine_api: EngineApi,
     socket_addr: SocketAddr,
     secret: JwtSecret,
-) -> Result<ServerHandle, RpcError>
+) -> Result<AuthServerHandle, RpcError>
 where
     Client: BlockProvider
         + HeaderProvider
@@ -73,7 +80,7 @@ where
 
     // Create auth middleware.
     let middleware =
-        tower::ServiceBuilder::new().layer(AuthLayer::new(JwtAuthValidator::new(secret)));
+        tower::ServiceBuilder::new().layer(AuthLayer::new(JwtAuthValidator::new(secret.clone())));
 
     // By default, both http and ws are enabled.
     let server = ServerBuilder::new()
@@ -82,7 +89,10 @@ where
         .await
         .map_err(|err| RpcError::from_jsonrpsee_error(err, ServerKind::Auth(socket_addr)))?;
 
-    Ok(server.start(module)?)
+    let local_addr = server.local_addr()?;
+
+    let handle = server.start(module)?;
+    Ok(AuthServerHandle { handle, local_addr, secret })
 }
 
 /// Server configuration for the auth server.
@@ -103,12 +113,12 @@ impl AuthServerConfig {
     }
 
     /// Convenience function to start a server in one step.
-    pub async fn start(self, module: AuthRpcModule) -> Result<ServerHandle, RpcError> {
+    pub async fn start(self, module: AuthRpcModule) -> Result<AuthServerHandle, RpcError> {
         let Self { socket_addr, secret } = self;
 
         // Create auth middleware.
-        let middleware =
-            tower::ServiceBuilder::new().layer(AuthLayer::new(JwtAuthValidator::new(secret)));
+        let middleware = tower::ServiceBuilder::new()
+            .layer(AuthLayer::new(JwtAuthValidator::new(secret.clone())));
 
         // By default, both http and ws are enabled.
         let server =
@@ -116,7 +126,10 @@ impl AuthServerConfig {
                 |err| RpcError::from_jsonrpsee_error(err, ServerKind::Auth(socket_addr)),
             )?;
 
-        Ok(server.start(module.inner)?)
+        let local_addr = server.local_addr()?;
+
+        let handle = server.start(module.inner)?;
+        Ok(AuthServerHandle { handle, local_addr, secret })
     }
 }
 
@@ -172,8 +185,93 @@ pub struct AuthRpcModule {
 // === impl TransportRpcModules ===
 
 impl AuthRpcModule {
+    /// Create a new `AuthRpcModule` with the given `engine_api`.
+    pub fn new<EngineApi>(engine: EngineApi) -> Self
+    where
+        EngineApi: EngineApiServer,
+    {
+        let mut module = RpcModule::new(());
+        module.merge(engine.into_rpc()).expect("No conflicting methods");
+        Self { inner: module }
+    }
+
+    /// Get a reference to the inner `RpcModule`.
+    pub fn module_mut(&mut self) -> &mut RpcModule<()> {
+        &mut self.inner
+    }
+
     /// Convenience function for starting a server
-    pub async fn start_server(self, config: AuthServerConfig) -> Result<ServerHandle, RpcError> {
+    pub async fn start_server(
+        self,
+        config: AuthServerConfig,
+    ) -> Result<AuthServerHandle, RpcError> {
         config.start(self).await
     }
 }
+
+/// A handle to the spawned auth server.
+///
+/// When this type is dropped or [AuthServerHandle::stop] has been called the server will be
+/// stopped.
+#[derive(Clone, Debug)]
+#[must_use = "Server stops if dropped"]
+pub struct AuthServerHandle {
+    local_addr: SocketAddr,
+    handle: ServerHandle,
+    secret: JwtSecret,
+}
+
+// === impl AuthServerHandle ===
+
+impl AuthServerHandle {
+    /// Returns the [`SocketAddr`] of the http server if started.
+    pub fn local_addr(&self) -> SocketAddr {
+        self.local_addr
+    }
+
+    /// Tell the server to stop without waiting for the server to stop.
+    pub fn stop(self) -> Result<(), RpcError> {
+        Ok(self.handle.stop()?)
+    }
+
+    /// Returns the url to the http server
+    pub fn http_url(&self) -> String {
+        format!("http://{}", self.local_addr)
+    }
+
+    /// Returns the url to the ws server
+    pub fn ws_url(&self) -> String {
+        format!("ws://{}", self.local_addr)
+    }
+
+    fn bearer(&self) -> String {
+        format!(
+            "Bearer {}",
+            self.secret
+                .encode(&Claims {
+                    iat: (SystemTime::now().duration_since(UNIX_EPOCH).unwrap() +
+                        Duration::from_secs(60))
+                    .as_secs(),
+                    exp: None,
+                })
+                .unwrap()
+        )
+    }
+
+    /// Returns a http client connected to the server.
+    pub fn http_client(&self) -> jsonrpsee::http_client::HttpClient {
+        jsonrpsee::http_client::HttpClientBuilder::default()
+            .set_headers(HeaderMap::from_iter([(AUTHORIZATION, self.bearer().parse().unwrap())]))
+            .build(self.http_url())
+            .expect("Failed to create http client")
+    }
+
+    /// Returns a ws client connected to the server.
+    pub async fn ws_client(&self) -> jsonrpsee::ws_client::WsClient {
+        jsonrpsee::ws_client::WsClientBuilder::default()
+            .set_headers(HeaderMap::from_iter([(AUTHORIZATION, self.bearer().parse().unwrap())]))
+            .build(self.ws_url())
+            .await
+            .expect("Failed to create ws client")
+    }
+}

crates/rpc/rpc-builder/src/lib.rs

@@ -1255,7 +1255,7 @@ impl fmt::Debug for RpcServer {
 ///
 /// When this type is dropped or [RpcServerHandle::stop] has been called the server will be stopped.
 #[derive(Clone)]
-#[must_use = "Server stop if dropped"]
+#[must_use = "Server stops if dropped"]
 pub struct RpcServerHandle {
     /// The address of the http/ws server
     http_local_addr: Option<SocketAddr>,

crates/rpc/rpc-builder/tests/it/auth.rs

@@ -0,0 +1,44 @@
+//! Auth server tests
+
+use crate::utils::launch_auth;
+use jsonrpsee::core::client::{ClientT, SubscriptionClientT};
+use reth_primitives::Block;
+use reth_rpc::JwtSecret;
+use reth_rpc_api::clients::EngineApiClient;
+use reth_rpc_types::engine::{ForkchoiceState, PayloadId, TransitionConfiguration};
+
+#[allow(unused_must_use)]
+async fn test_basic_engine_calls<C>(client: &C)
+where
+    C: ClientT + SubscriptionClientT + Sync,
+{
+    let block = Block::default().seal_slow();
+    EngineApiClient::new_payload_v1(client, block.clone().into()).await;
+    EngineApiClient::new_payload_v2(client, block.into()).await;
+    EngineApiClient::fork_choice_updated_v1(client, ForkchoiceState::default(), None).await;
+    EngineApiClient::get_payload_v1(client, PayloadId::new([0, 0, 0, 0, 0, 0, 0, 0])).await;
+    EngineApiClient::get_payload_v2(client, PayloadId::new([0, 0, 0, 0, 0, 0, 0, 0])).await;
+    EngineApiClient::get_payload_bodies_by_hash_v1(client, vec![]).await;
+    EngineApiClient::get_payload_bodies_by_range_v1(client, 0u64.into(), 1u64.into()).await;
+    EngineApiClient::exchange_transition_configuration(client, TransitionConfiguration::default())
+        .await;
+    EngineApiClient::exchange_capabilities(client, vec![]).await;
+}
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_auth_endpoints_http() {
+    reth_tracing::init_test_tracing();
+    let secret = JwtSecret::random();
+    let handle = launch_auth(secret).await;
+    let client = handle.http_client();
+    test_basic_engine_calls(&client).await
+}
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_auth_endpoints_ws() {
+    reth_tracing::init_test_tracing();
+    let secret = JwtSecret::random();
+    let handle = launch_auth(secret).await;
+    let client = handle.ws_client().await;
+    test_basic_engine_calls(&client).await
+}

crates/rpc/rpc-builder/tests/it/main.rs

@@ -1,3 +1,4 @@
+mod auth;
 mod http;
 mod serde;
 mod startup;

crates/rpc/rpc-builder/tests/it/utils.rs

@@ -1,18 +1,43 @@
+use reth_beacon_consensus::BeaconConsensusEngineHandle;
 use reth_network_api::test_utils::NoopNetwork;
+use reth_payload_builder::test_utils::spawn_test_payload_service;
+use reth_primitives::MAINNET;
 use reth_provider::test_utils::{NoopProvider, TestCanonStateSubscriptions};
+use reth_rpc::JwtSecret;
 use reth_rpc_builder::{
+    auth::{AuthRpcModule, AuthServerConfig, AuthServerHandle},
     RpcModuleBuilder, RpcModuleSelection, RpcServerConfig, RpcServerHandle,
     TransportRpcModuleConfig,
 };
+use reth_rpc_engine_api::EngineApi;
 use reth_tasks::TokioTaskExecutor;
 use reth_transaction_pool::test_utils::{testing_pool, TestPool};
-use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4};
+use std::{
+    net::{Ipv4Addr, SocketAddr, SocketAddrV4},
+    sync::Arc,
+};
+use tokio::sync::mpsc::unbounded_channel;
 
 /// Localhost with port 0 so a free port is used.
 pub fn test_address() -> SocketAddr {
     SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::UNSPECIFIED, 0))
 }
 
+/// Launches a new server for the auth module
+pub async fn launch_auth(secret: JwtSecret) -> AuthServerHandle {
+    let config = AuthServerConfig::builder(secret).socket_addr(test_address()).build();
+    let (tx, _rx) = unbounded_channel();
+    let beacon_engine_handle = BeaconConsensusEngineHandle::new(tx);
+    let engine_api = EngineApi::new(
+        NoopProvider::default(),
+        Arc::new(MAINNET.clone()),
+        beacon_engine_handle,
+        spawn_test_payload_service().into(),
+    );
+    let module = AuthRpcModule::new(engine_api);
+    module.start_server(config).await.unwrap()
+}
+
 /// Launches a new server with http only with the given modules
 pub async fn launch_http(modules: impl Into<RpcModuleSelection>) -> RpcServerHandle {
     let builder = test_rpc_builder();