From 08c643bd4b01ab7661bf492ef50eea5237c02bb6 Mon Sep 17 00:00:00 2001
From: Brandon Williams <brandon@mystenlabs.com>
Date: Mon, 11 Mar 2024 20:07:57 -0500
Subject: [PATCH] json-rpc: add option to only allow websocket traffic

---
 crates/sui-config/src/node.rs                             | 4 ++++
 crates/sui-node/src/lib.rs                                | 8 +++++++-
 crates/sui-swarm-config/src/node_config_builder.rs        | 2 ++
 .../snapshot_tests__network_config_snapshot_matches.snap  | 8 ++++++++
 4 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/crates/sui-config/src/node.rs b/crates/sui-config/src/node.rs
index 849d7d4d45b5a..a006154853287 100644
--- a/crates/sui-config/src/node.rs
+++ b/crates/sui-config/src/node.rs
@@ -79,6 +79,10 @@ pub struct NodeConfig {
     #[serde(default = "default_enable_index_processing")]
     pub enable_index_processing: bool,
 
+    // only alow websocket connections for jsonrpc traffic
+    #[serde(default)]
+    pub websocket_only: bool,
+
     #[serde(default)]
     pub grpc_load_shed: Option<bool>,
 
diff --git a/crates/sui-node/src/lib.rs b/crates/sui-node/src/lib.rs
index ffd88152b368c..ef171c13d06c1 100644
--- a/crates/sui-node/src/lib.rs
+++ b/crates/sui-node/src/lib.rs
@@ -28,6 +28,7 @@ use sui_core::consensus_adapter::SubmitToConsensus;
 use sui_core::epoch::randomness::RandomnessManager;
 use sui_core::execution_cache::ExecutionCacheMetrics;
 use sui_core::execution_cache::NotifyReadWrapper;
+use sui_json_rpc::ServerType;
 use sui_json_rpc_api::JsonRpcMetrics;
 use sui_network::randomness;
 use sui_protocol_config::ProtocolVersion;
@@ -1928,7 +1929,12 @@ pub fn build_http_server(
         ))?;
         server.register_module(MoveUtils::new(state))?;
 
-        server.to_router(None)?
+        let server_type = if config.websocket_only {
+            Some(ServerType::WebSocket)
+        } else {
+            None
+        };
+        server.to_router(server_type)?
     };
 
     router = router.merge(json_rpc_router);
diff --git a/crates/sui-swarm-config/src/node_config_builder.rs b/crates/sui-swarm-config/src/node_config_builder.rs
index 9a270267a8bb8..64e6b36a0dcd7 100644
--- a/crates/sui-swarm-config/src/node_config_builder.rs
+++ b/crates/sui-swarm-config/src/node_config_builder.rs
@@ -195,6 +195,7 @@ impl ValidatorConfigBuilder {
             zklogin_oauth_providers: default_zklogin_oauth_providers(),
             authority_overload_config: self.authority_overload_config.unwrap_or_default(),
             run_with_range: None,
+            websocket_only: false,
         }
     }
 
@@ -447,6 +448,7 @@ impl FullnodeConfigBuilder {
             zklogin_oauth_providers: default_zklogin_oauth_providers(),
             authority_overload_config: Default::default(),
             run_with_range: self.run_with_range,
+            websocket_only: false,
         }
     }
 }
diff --git a/crates/sui-swarm-config/tests/snapshots/snapshot_tests__network_config_snapshot_matches.snap b/crates/sui-swarm-config/tests/snapshots/snapshot_tests__network_config_snapshot_matches.snap
index 684593e4e60c3..8a1cd5854ae92 100644
--- a/crates/sui-swarm-config/tests/snapshots/snapshot_tests__network_config_snapshot_matches.snap
+++ b/crates/sui-swarm-config/tests/snapshots/snapshot_tests__network_config_snapshot_matches.snap
@@ -46,6 +46,7 @@ validator_configs:
           request_batches_rate_limit: ~
     enable-event-processing: false
     enable-index-processing: true
+    websocket-only: false
     grpc-load-shed: ~
     grpc-concurrency-limit: 20000000000
     p2p-config:
@@ -170,6 +171,7 @@ validator_configs:
           request_batches_rate_limit: ~
     enable-event-processing: false
     enable-index-processing: true
+    websocket-only: false
     grpc-load-shed: ~
     grpc-concurrency-limit: 20000000000
     p2p-config:
@@ -294,6 +296,7 @@ validator_configs:
           request_batches_rate_limit: ~
     enable-event-processing: false
     enable-index-processing: true
+    websocket-only: false
     grpc-load-shed: ~
     grpc-concurrency-limit: 20000000000
     p2p-config:
@@ -418,6 +421,7 @@ validator_configs:
           request_batches_rate_limit: ~
     enable-event-processing: false
     enable-index-processing: true
+    websocket-only: false
     grpc-load-shed: ~
     grpc-concurrency-limit: 20000000000
     p2p-config:
@@ -542,6 +546,7 @@ validator_configs:
           request_batches_rate_limit: ~
     enable-event-processing: false
     enable-index-processing: true
+    websocket-only: false
     grpc-load-shed: ~
     grpc-concurrency-limit: 20000000000
     p2p-config:
@@ -666,6 +671,7 @@ validator_configs:
           request_batches_rate_limit: ~
     enable-event-processing: false
     enable-index-processing: true
+    websocket-only: false
     grpc-load-shed: ~
     grpc-concurrency-limit: 20000000000
     p2p-config:
@@ -790,6 +796,7 @@ validator_configs:
           request_batches_rate_limit: ~
     enable-event-processing: false
     enable-index-processing: true
+    websocket-only: false
     grpc-load-shed: ~
     grpc-concurrency-limit: 20000000000
     p2p-config:
@@ -878,3 +885,4 @@ account_keys:
   - mfPjCoE6SX0Sl84MnmNS/LS+tfPpkn7I8tziuk2g0WM=
   - 5RWlYF22jS9i76zLl8jP2D3D8GC5ht+IP1dWUBGZxi8=
 genesis: "[fake genesis]"
+