feat(rpc-builder): add tower layer for updating bearer token in auth client (#8010)

crates/e2e-test-utils/src/engine_api.rs

@@ -1,5 +1,5 @@
 use crate::traits::PayloadEnvelopeExt;
-use jsonrpsee::http_client::HttpClient;
+use jsonrpsee::http_client::{transport::HttpBackend, HttpClient};
 use reth::{
     api::{EngineTypes, PayloadBuilderAttributes},
     providers::CanonStateNotificationStream,
@@ -10,12 +10,13 @@ use reth::{
 };
 use reth_payload_builder::PayloadId;
 use reth_primitives::B256;
+use reth_rpc::AuthClientService;
 use std::marker::PhantomData;
 
 /// Helper for engine api operations
 pub struct EngineApiTestContext<E> {
     pub canonical_stream: CanonStateNotificationStream,
-    pub engine_api_client: HttpClient,
+    pub engine_api_client: HttpClient<AuthClientService<HttpBackend>>,
     pub _marker: PhantomData<E>,
 }
 

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

@@ -15,6 +15,7 @@ use jsonrpsee::{
 };
 pub use reth_ipc::server::Builder as IpcServerBuilder;
 
+use jsonrpsee::http_client::transport::HttpBackend;
 use reth_engine_primitives::EngineTypes;
 use reth_evm::ConfigureEvm;
 use reth_network_api::{NetworkInfo, Peers};
@@ -27,16 +28,13 @@ use reth_rpc::{
         cache::EthStateCache, gas_oracle::GasPriceOracle, EthFilterConfig, FeeHistoryCache,
         FeeHistoryCacheConfig,
     },
-    AuthLayer, Claims, EngineEthApi, EthApi, EthFilter, EthSubscriptionIdProvider,
-    JwtAuthValidator, JwtSecret,
+    secret_to_bearer_header, AuthClientLayer, AuthClientService, AuthLayer, EngineEthApi, EthApi,
+    EthFilter, EthSubscriptionIdProvider, JwtAuthValidator, JwtSecret,
 };
 use reth_rpc_api::servers::*;
 use reth_tasks::{pool::BlockingTaskPool, TaskSpawner};
 use reth_transaction_pool::TransactionPool;
-use std::{
-    net::{IpAddr, Ipv4Addr, SocketAddr},
-    time::{Duration, SystemTime, UNIX_EPOCH},
-};
+use std::net::{IpAddr, Ipv4Addr, SocketAddr};
 use tower::layer::util::Identity;
 
 /// Configure and launch a _standalone_ auth server with `engine` and a _new_ `eth` namespace.
@@ -397,32 +395,27 @@ impl AuthServerHandle {
         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 {
+    pub fn http_client(
+        &self,
+    ) -> jsonrpsee::http_client::HttpClient<AuthClientService<HttpBackend>> {
+        // Create a middleware that adds a new JWT token to every request.
+        let secret_layer = AuthClientLayer::new(self.secret.clone());
+        let middleware = tower::ServiceBuilder::default().layer(secret_layer);
         jsonrpsee::http_client::HttpClientBuilder::default()
-            .set_headers(HeaderMap::from_iter([(AUTHORIZATION, self.bearer().parse().unwrap())]))
+            .set_http_middleware(middleware)
             .build(self.http_url())
             .expect("Failed to create http client")
     }
 
-    /// Returns a ws client connected to the server.
+    /// Returns a ws client connected to the server. Note that the connection can only be
+    /// be established within 1 minute due to the JWT token expiration.
     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())]))
+            .set_headers(HeaderMap::from_iter([(
+                AUTHORIZATION,
+                secret_to_bearer_header(&self.secret),
+            )]))
             .build(self.ws_url())
             .await
             .expect("Failed to create ws client")

crates/rpc/rpc/src/layers/auth_client_layer.rs

@@ -0,0 +1,79 @@
+use crate::{Claims, JwtSecret};
+use http::HeaderValue;
+use hyper::{header::AUTHORIZATION, service::Service};
+use std::{
+    task::{Context, Poll},
+    time::{Duration, SystemTime, UNIX_EPOCH},
+};
+use tower::Layer;
+
+/// A layer that adds a new JWT token to every request using AuthClientService.
+#[derive(Debug)]
+pub struct AuthClientLayer {
+    secret: JwtSecret,
+}
+
+impl AuthClientLayer {
+    /// Create a new AuthClientLayer with the given `secret`.
+    pub fn new(secret: JwtSecret) -> Self {
+        Self { secret }
+    }
+}
+
+impl<S> Layer<S> for AuthClientLayer {
+    type Service = AuthClientService<S>;
+
+    fn layer(&self, inner: S) -> Self::Service {
+        AuthClientService::new(self.secret.clone(), inner)
+    }
+}
+
+/// Automatically authenticates every client request with the given `secret`.
+#[derive(Debug, Clone)]
+pub struct AuthClientService<S> {
+    secret: JwtSecret,
+    inner: S,
+}
+
+impl<S> AuthClientService<S> {
+    fn new(secret: JwtSecret, inner: S) -> Self {
+        Self { secret, inner }
+    }
+}
+
+impl<S, B> Service<hyper::Request<B>> for AuthClientService<S>
+where
+    S: Service<hyper::Request<B>>,
+{
+    type Response = S::Response;
+    type Error = S::Error;
+    type Future = S::Future;
+
+    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+        self.inner.poll_ready(cx)
+    }
+
+    fn call(&mut self, mut request: hyper::Request<B>) -> Self::Future {
+        request.headers_mut().insert(AUTHORIZATION, secret_to_bearer_header(&self.secret));
+        self.inner.call(request)
+    }
+}
+
+/// Helper function to convert a secret into a Bearer auth header value with claims according to
+/// <https://github.com/ethereum/execution-apis/blob/main/src/engine/authentication.md#jwt-claims>.
+/// The token is valid for 60 seconds.
+pub fn secret_to_bearer_header(secret: &JwtSecret) -> HeaderValue {
+    format!(
+        "Bearer {}",
+        secret
+            .encode(&Claims {
+                iat: (SystemTime::now().duration_since(UNIX_EPOCH).unwrap() +
+                    Duration::from_secs(60))
+                .as_secs(),
+                exp: None,
+            })
+            .unwrap()
+    )
+    .parse()
+    .unwrap()
+}

crates/rpc/rpc/src/layers/mod.rs

@@ -1,8 +1,11 @@
 use http::{HeaderMap, Response};
 
+mod auth_client_layer;
 mod auth_layer;
 mod jwt_secret;
 mod jwt_validator;
+
+pub use auth_client_layer::{secret_to_bearer_header, AuthClientLayer, AuthClientService};
 pub use auth_layer::AuthLayer;
 pub use jwt_secret::{Claims, JwtError, JwtSecret};
 pub use jwt_validator::JwtAuthValidator;

crates/rpc/rpc/src/lib.rs

@@ -41,7 +41,10 @@ pub use admin::AdminApi;
 pub use debug::DebugApi;
 pub use engine::{EngineApi, EngineEthApi};
 pub use eth::{EthApi, EthApiSpec, EthFilter, EthPubSub, EthSubscriptionIdProvider};
-pub use layers::{AuthLayer, AuthValidator, Claims, JwtAuthValidator, JwtError, JwtSecret};
+pub use layers::{
+    secret_to_bearer_header, AuthClientLayer, AuthClientService, AuthLayer, AuthValidator, Claims,
+    JwtAuthValidator, JwtError, JwtSecret,
+};
 pub use net::NetApi;
 pub use otterscan::OtterscanApi;
 pub use reth::RethApi;