refactor: add connection info (#547)

Author: gusinacio

Cargo.lock

@@ -40,8 +40,8 @@ tap_core.workspace = true
 uuid.workspace = true
 alloy.workspace = true
 typed-builder.workspace = true
-tower_governor = "0.4.3"
-governor = "0.6.0"
+tower_governor = { version = "0.5.0", features = ["axum"] }
+governor = "0.8.0"
 tower-http = { version = "0.6.2", features = [
   "auth",
   "cors",

crates/service/Cargo.toml

@@ -1,7 +1,7 @@
 // Copyright 2023-, Edge & Node, GraphOps, and Semiotic Labs.
 // SPDX-License-Identifier: Apache-2.0
 
-use std::time::Duration;
+use std::{net::SocketAddr, time::Duration};
 
 use anyhow::anyhow;
 use axum::{extract::Request, serve, ServiceExt};
@@ -119,7 +119,7 @@ pub async fn run() -> anyhow::Result<()> {
     let app = router.create_router().await?;
     let router = NormalizePath::trim_trailing_slash(app);
     //
-    let service = ServiceExt::<Request>::into_make_service(router);
+    let service = ServiceExt::<Request>::into_make_service_with_connect_info::<SocketAddr>(router);
     Ok(serve(listener, service)
         .with_graceful_shutdown(shutdown_handler())
         .await?)

crates/service/src/service.rs

@@ -26,7 +26,7 @@ use reqwest::Method;
 use tap_core::{manager::Manager, receipt::checks::CheckList};
 use tower::ServiceBuilder;
 use tower_governor::{
-    governor::GovernorConfigBuilder, key_extractor::PeerIpKeyExtractor, GovernorLayer,
+    governor::GovernorConfigBuilder, key_extractor::SmartIpKeyExtractor, GovernorLayer,
 };
 use tower_http::{
     auth::AsyncRequireAuthorizationLayer,
@@ -434,12 +434,13 @@ impl ServiceRouter {
 fn create_rate_limiter(
     burst_per_millisecond: u64,
     burst_size: u32,
-) -> GovernorLayer<PeerIpKeyExtractor, NoOpMiddleware<QuantaInstant>> {
+) -> GovernorLayer<SmartIpKeyExtractor, NoOpMiddleware<QuantaInstant>> {
     GovernorLayer {
         config: Arc::new(
             GovernorConfigBuilder::default()
                 .per_millisecond(burst_per_millisecond)
                 .burst_size(burst_size)
+                .key_extractor(SmartIpKeyExtractor)
                 .finish()
                 .expect("Failed to set up rate limiting"),
         ),

crates/service/src/service/router.rs

@@ -1,10 +1,10 @@
 // Copyright 2023-, Edge & Node, GraphOps, and Semiotic Labs.
 // SPDX-License-Identifier: Apache-2.0
 
-use std::time::Duration;
+use std::{net::SocketAddr, time::Duration};
 
 use alloy::primitives::Address;
-use axum::{body::to_bytes, http::Request};
+use axum::{body::to_bytes, extract::ConnectInfo, http::Request, Extension};
 use axum_extra::headers::Header;
 use indexer_config::{BlockchainConfig, GraphNodeConfig, IndexerConfig, NonZeroGRT};
 use indexer_monitor::EscrowAccounts;
@@ -95,7 +95,20 @@ async fn full_integration_test(database: PgPool) {
         .allocations(allocations)
         .build();
 
-    let mut app = router.create_router().await.unwrap();
+    let socket_info = Extension(ConnectInfo(SocketAddr::from(([0, 0, 0, 0], 1337))));
+    let mut app = router.create_router().await.unwrap().layer(socket_info);
+
+    let res = app
+        .call(Request::get("/").body(String::new()).unwrap())
+        .await
+        .unwrap();
+
+    assert_eq!(res.status(), StatusCode::OK);
+
+    let graphql_response = res.into_body();
+    let bytes = to_bytes(graphql_response, usize::MAX).await.unwrap();
+    let res = String::from_utf8(bytes.into()).unwrap();
+    insta::assert_snapshot!(res);
 
     let receipt = create_signed_receipt(
         SignedReceiptRequest::builder()

crates/service/tests/router_test.rs

@@ -3,4 +3,4 @@ source: crates/service/tests/router_test.rs
 expression: res
 snapshot_kind: text
 ---
-{"message":"No Tap receipt was found in the request"}
+{"graphQLResponse":"\n                {\n                    \"data\": {\n                        \"graphNetwork\": {\n                            \"currentEpoch\": 960\n                        }\n                    }\n                }\n                ","attestation":null}

crates/service/tests/snapshots/router_test__full_integration_test-2.snap

@@ -0,0 +1,6 @@
+---
+source: crates/service/tests/router_test.rs
+expression: res
+snapshot_kind: text
+---
+{"message":"No Tap receipt was found in the request"}

crates/service/tests/snapshots/router_test__full_integration_test-3.snap

@@ -3,4 +3,4 @@ source: crates/service/tests/router_test.rs
 expression: res
 snapshot_kind: text
 ---
-{"graphQLResponse":"\n                {\n                    \"data\": {\n                        \"graphNetwork\": {\n                            \"currentEpoch\": 960\n                        }\n                    }\n                }\n                ","attestation":null}
+Service is up and running