From a89a576b122ac020ba2e5fcfcbafe311fb495d29 Mon Sep 17 00:00:00 2001
From: Matthias Seitz <matthias.seitz@outlook.de>
Date: Fri, 19 Jul 2024 17:05:37 +0200
Subject: [PATCH] test: add test for rpc middleware (#9620)

---
 crates/rpc/rpc-builder/tests/it/main.rs       |  1 +
 crates/rpc/rpc-builder/tests/it/middleware.rs | 80 +++++++++++++++++++
 2 files changed, 81 insertions(+)
 create mode 100644 crates/rpc/rpc-builder/tests/it/middleware.rs

diff --git a/crates/rpc/rpc-builder/tests/it/main.rs b/crates/rpc/rpc-builder/tests/it/main.rs
index 65ddebb3fd..a64ad1da2f 100644
--- a/crates/rpc/rpc-builder/tests/it/main.rs
+++ b/crates/rpc/rpc-builder/tests/it/main.rs
@@ -1,5 +1,6 @@
 mod auth;
 mod http;
+mod middleware;
 mod serde;
 mod startup;
 pub mod utils;
diff --git a/crates/rpc/rpc-builder/tests/it/middleware.rs b/crates/rpc/rpc-builder/tests/it/middleware.rs
new file mode 100644
index 0000000000..59cc86d4dc
--- /dev/null
+++ b/crates/rpc/rpc-builder/tests/it/middleware.rs
@@ -0,0 +1,80 @@
+use crate::utils::{test_address, test_rpc_builder};
+use jsonrpsee::{
+    server::{middleware::rpc::RpcServiceT, RpcServiceBuilder},
+    types::Request,
+    MethodResponse,
+};
+use reth_rpc::EthApi;
+use reth_rpc_builder::{RpcServerConfig, TransportRpcModuleConfig};
+use reth_rpc_eth_api::EthApiClient;
+use reth_rpc_server_types::RpcModuleSelection;
+use std::{
+    future::Future,
+    pin::Pin,
+    sync::{
+        atomic::{AtomicUsize, Ordering},
+        Arc,
+    },
+};
+use tower::Layer;
+
+#[derive(Clone, Default)]
+struct MyMiddlewareLayer {
+    count: Arc<AtomicUsize>,
+}
+
+impl<S> Layer<S> for MyMiddlewareLayer {
+    type Service = MyMiddlewareService<S>;
+
+    fn layer(&self, inner: S) -> Self::Service {
+        MyMiddlewareService { service: inner, count: self.count.clone() }
+    }
+}
+
+#[derive(Clone)]
+struct MyMiddlewareService<S> {
+    service: S,
+    count: Arc<AtomicUsize>,
+}
+
+impl<'a, S> RpcServiceT<'a> for MyMiddlewareService<S>
+where
+    S: RpcServiceT<'a> + Send + Sync + Clone + 'static,
+{
+    type Future = Pin<Box<dyn Future<Output = MethodResponse> + Send + 'a>>;
+
+    fn call(&self, req: Request<'a>) -> Self::Future {
+        tracing::info!("MyMiddleware processed call {}", req.method);
+        let count = self.count.clone();
+        let service = self.service.clone();
+        Box::pin(async move {
+            let rp = service.call(req).await;
+            // Modify the state.
+            count.fetch_add(1, Ordering::Relaxed);
+            rp
+        })
+    }
+}
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_rpc_middleware() {
+    let builder = test_rpc_builder();
+    let modules = builder.build(
+        TransportRpcModuleConfig::set_http(RpcModuleSelection::All),
+        Box::new(EthApi::with_spawner),
+    );
+
+    let mylayer = MyMiddlewareLayer::default();
+
+    let handle = RpcServerConfig::http(Default::default())
+        .with_http_address(test_address())
+        .set_rpc_middleware(RpcServiceBuilder::new().layer(mylayer.clone()))
+        .start(&modules)
+        .await
+        .unwrap();
+
+    let client = handle.http_client().unwrap();
+    EthApiClient::protocol_version(&client).await.unwrap();
+    let count = mylayer.count.load(Ordering::Relaxed);
+    assert_eq!(count, 1);
+}