From a8ffd5af791de43b782d85a8a4104e5b366fe476 Mon Sep 17 00:00:00 2001
From: Brandon Williams <brandon@mystenlabs.com>
Date: Thu, 2 Jan 2025 15:29:54 -0600
Subject: [PATCH] sui-node: use sui-http for rpc http service

---
 Cargo.lock                 |  1 +
 crates/sui-node/Cargo.toml |  1 +
 crates/sui-node/src/lib.rs | 35 ++++++++++++++++-------------------
 3 files changed, 18 insertions(+), 19 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 6fca532f7f4f9..a73baef1b0246 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -14678,6 +14678,7 @@ dependencies = [
  "sui-archival",
  "sui-config",
  "sui-core",
+ "sui-http",
  "sui-json-rpc",
  "sui-json-rpc-api",
  "sui-macros",
diff --git a/crates/sui-node/Cargo.toml b/crates/sui-node/Cargo.toml
index 430d73f963abe..db60f9770d470 100644
--- a/crates/sui-node/Cargo.toml
+++ b/crates/sui-node/Cargo.toml
@@ -55,6 +55,7 @@ telemetry-subscribers.workspace = true
 fastcrypto.workspace = true
 fastcrypto-zkp.workspace = true
 move-vm-profiler.workspace = true
+sui-http.workspace = true
 
 [target.'cfg(msim)'.dependencies]
 sui-simulator.workspace = true
diff --git a/crates/sui-node/src/lib.rs b/crates/sui-node/src/lib.rs
index 8e84d01b2e29b..1de90befbcb14 100644
--- a/crates/sui-node/src/lib.rs
+++ b/crates/sui-node/src/lib.rs
@@ -17,7 +17,6 @@ use mysten_network::server::SUI_TLS_SERVER_NAME;
 use prometheus::Registry;
 use std::collections::{BTreeSet, HashMap, HashSet};
 use std::fmt;
-use std::net::SocketAddr;
 use std::path::PathBuf;
 use std::str::FromStr;
 #[cfg(msim)]
@@ -229,7 +228,7 @@ pub struct SuiNode {
     config: NodeConfig,
     validator_components: Mutex<Option<ValidatorComponents>>,
     /// The http server responsible for serving JSON-RPC as well as the experimental rest service
-    _http_server: Option<tokio::task::JoinHandle<()>>,
+    _http_server: Option<sui_http::ServerHandle>,
     state: Arc<AuthorityState>,
     transaction_orchestrator: Option<Arc<TransactiondOrchestrator<NetworkAuthorityClient>>>,
     registry_service: RegistryService,
@@ -2020,7 +2019,7 @@ pub async fn build_http_server(
     prometheus_registry: &Registry,
     _custom_runtime: Option<Handle>,
     software_version: &'static str,
-) -> Result<Option<tokio::task::JoinHandle<()>>> {
+) -> Result<Option<sui_http::ServerHandle>> {
     // Validators do not expose these APIs
     if config.consensus_config().is_some() {
         return Ok(None);
@@ -2125,25 +2124,23 @@ pub async fn build_http_server(
         rpc_service.into_router().await
     };
 
-    router = router.merge(rpc_router);
-
-    let listener = tokio::net::TcpListener::bind(&config.json_rpc_address)
-        .await
-        .unwrap();
-    let addr = listener.local_addr().unwrap();
+    let layers = ServiceBuilder::new()
+        .map_request(|mut request: axum::http::Request<_>| {
+            if let Some(connect_info) = request.extensions().get::<sui_http::ConnectInfo>() {
+                let axum_connect_info = axum::extract::ConnectInfo(connect_info.remote_addr);
+                request.extensions_mut().insert(axum_connect_info);
+            }
+            request
+        })
+        .layer(axum::middleware::from_fn(server_timing_middleware));
 
-    router = router.layer(axum::middleware::from_fn(server_timing_middleware));
+    router = router.merge(rpc_router).layer(layers);
 
-    let handle = tokio::spawn(async move {
-        axum::serve(
-            listener,
-            router.into_make_service_with_connect_info::<SocketAddr>(),
-        )
-        .await
-        .unwrap()
-    });
+    let handle = sui_http::Builder::new()
+        .serve(&config.json_rpc_address, router)
+        .map_err(|e| anyhow::anyhow!("{e}"))?;
 
-    info!(local_addr =? addr, "Sui JSON-RPC server listening on {addr}");
+    info!(local_addr =? handle.local_addr(), "Sui JSON-RPC server listening on {}", handle.local_addr());
 
     Ok(Some(handle))
 }