From 9136f80d36f1949c713ee1a1f8dfbbe30001e7cb Mon Sep 17 00:00:00 2001
From: DimAn <diman@diman.io>
Date: Wed, 12 Apr 2023 15:32:22 +0200
Subject: [PATCH] validator: add `set-public-tpu-address` command (#30452)

---
 gossip/src/cluster_info.rs         |  9 ++++++++-
 validator/src/admin_rpc_service.rs | 31 ++++++++++++++++++++++++++++++
 validator/src/cli.rs               | 13 +++++++++++++
 validator/src/main.rs              | 26 +++++++++++++++++++++++++
 4 files changed, 78 insertions(+), 1 deletion(-)

diff --git a/gossip/src/cluster_info.rs b/gossip/src/cluster_info.rs
index 03f2ff588..749d6aafd 100644
--- a/gossip/src/cluster_info.rs
+++ b/gossip/src/cluster_info.rs
@@ -24,7 +24,7 @@ use {
         cluster_info_metrics::{
             submit_gossip_stats, Counter, GossipStats, ScopedTimer, TimedGuard,
         },
-        contact_info::{ContactInfo, LegacyContactInfo},
+        contact_info::{ContactInfo, Error as ContactInfoError, LegacyContactInfo},
         crds::{Crds, Cursor, GossipRoute},
         crds_gossip::CrdsGossip,
         crds_gossip_error::CrdsGossipError,
@@ -677,6 +677,13 @@ impl ClusterInfo {
         self.push_self();
     }
 
+    pub fn set_tpu(&self, tpu_addr: SocketAddr) -> Result<(), ContactInfoError> {
+        self.my_contact_info.write().unwrap().set_tpu(tpu_addr)?;
+        self.insert_self();
+        self.push_self();
+        Ok(())
+    }
+
     pub fn lookup_contact_info<F, Y>(&self, id: &Pubkey, map: F) -> Option<Y>
     where
         F: FnOnce(&LegacyContactInfo) -> Y,
diff --git a/validator/src/admin_rpc_service.rs b/validator/src/admin_rpc_service.rs
index 0916577bc..099f34be4 100644
--- a/validator/src/admin_rpc_service.rs
+++ b/validator/src/admin_rpc_service.rs
@@ -226,6 +226,13 @@ pub trait AdminRpc {
         secondary_index: RpcAccountIndex,
         max_entries: usize,
     ) -> Result<Vec<(String, usize)>>;
+
+    #[rpc(meta, name = "setPublicTpuAddress")]
+    fn set_public_tpu_address(
+        &self,
+        meta: Self::Metadata,
+        public_tpu_addr: SocketAddr,
+    ) -> Result<()>;
 }
 
 pub struct AdminRpcImpl;
@@ -588,6 +595,30 @@ impl AdminRpc for AdminRpcImpl {
             Ok(largest_keys)
         })
     }
+
+    fn set_public_tpu_address(
+        &self,
+        meta: Self::Metadata,
+        public_tpu_addr: SocketAddr,
+    ) -> Result<()> {
+        debug!("set_public_tpu_address rpc request received: {public_tpu_addr}");
+
+        meta.with_post_init(|post_init| {
+            post_init
+                .cluster_info
+                .set_tpu(public_tpu_addr)
+                .map_err(|err| {
+                    error!("Failed to set public TPU address to {public_tpu_addr}: {err}");
+                    jsonrpc_core::error::Error::internal_error()
+                })?;
+            warn!(
+                "Public TPU addresses set to {} (udp) and {} (quic)",
+                post_init.cluster_info.my_contact_info().tpu().unwrap(),
+                post_init.cluster_info.my_contact_info().tpu_quic().unwrap(),
+            );
+            Ok(())
+        })
+    }
 }
 
 impl AdminRpcImpl {
diff --git a/validator/src/cli.rs b/validator/src/cli.rs
index d07b81045..05f50fda7 100644
--- a/validator/src/cli.rs
+++ b/validator/src/cli.rs
@@ -1614,6 +1614,19 @@ pub fn app<'a>(version: &'a str, default_args: &'a DefaultArgs) -> App<'a, 'a> {
                 )
                 .after_help("Note: If this command exits with a non-zero status \
                          then this not a good time for a restart")
+        ).
+        subcommand(
+            SubCommand::with_name("set-public-tpu-address")
+                .about("Specify TPU address to advertise in gossip")
+                .arg(
+                    Arg::with_name("public_tpu_addr")
+                        .index(1)
+                        .value_name("HOST:PORT")
+                        .takes_value(true)
+                        .required(true)
+                        .validator(solana_net_utils::is_host_port)
+                        .help("TPU address to advertise in gossip")
+                ),
         );
 }
 
diff --git a/validator/src/main.rs b/validator/src/main.rs
index 0923b5185..d1296347c 100644
--- a/validator/src/main.rs
+++ b/validator/src/main.rs
@@ -835,6 +835,32 @@ pub fn main() {
                 _ => unreachable!(),
             }
         }
+        ("set-public-tpu-address", Some(subcommand_matches)) => {
+            let public_tpu_addr: SocketAddr = subcommand_matches
+                .value_of("public_tpu_addr")
+                .map(|public_tpu_addr| {
+                    solana_net_utils::parse_host_port(public_tpu_addr).unwrap_or_else(|err| {
+                        eprintln!("Failed to parse --set-public-tpu-address HOST:PORT {err}");
+                        exit(1);
+                    })
+                })
+                .unwrap();
+
+            let admin_client = admin_rpc_service::connect(&ledger_path);
+            admin_rpc_service::runtime()
+                .block_on(async move {
+                    admin_client
+                        .await?
+                        .set_public_tpu_address(public_tpu_addr)
+                        .await
+                })
+                .unwrap_or_else(|err| {
+                    println!("setPublicTpuAddress request failed: {err}");
+                    exit(1);
+                });
+
+            return;
+        }
         _ => unreachable!(),
     };