From cfd75cf9eacabd611683fd9bc94d266291f1dedf Mon Sep 17 00:00:00 2001
From: Michael Beardsworth <beardsworth@intrinsic.ai>
Date: Mon, 6 May 2024 11:51:11 -0700
Subject: [PATCH] Allow programmatic configuration of unicast relays.

This change allows users to configure relays from code without having to
`setenv(GZ_RELAY)`.

Signed-off-by: Michael Beardsworth <beardsworth@intrinsic.ai>
---
 include/gz/transport/Discovery.hh  | 51 ++++++++++++++++++------------
 include/gz/transport/Node.hh       |  3 ++
 include/gz/transport/NodeShared.hh |  3 ++
 src/Node.cc                        |  8 +++++
 src/NodeShared.cc                  | 19 +++++++++++
 5 files changed, 64 insertions(+), 20 deletions(-)

diff --git a/include/gz/transport/Discovery.hh b/include/gz/transport/Discovery.hh
index 161ebc4e3..186af6020 100644
--- a/include/gz/transport/Discovery.hh
+++ b/include/gz/transport/Discovery.hh
@@ -740,6 +740,37 @@ namespace gz
         }
       }
 
+      /// \brief Register a new relay address.
+      /// \param[in] _ip New IP address.
+      public: void AddRelayAddress(const std::string &_ip)
+      {
+        // Sanity check: Make sure that this IP address is not already saved.
+        for (auto const &addr : this->relayAddrs)
+        {
+          if (addr.sin_addr.s_addr == inet_addr(_ip.c_str()))
+            return;
+        }
+
+        sockaddr_in addr;
+        memset(&addr, 0, sizeof(addr));
+        addr.sin_family = AF_INET;
+        addr.sin_addr.s_addr = inet_addr(_ip.c_str());
+        addr.sin_port = htons(static_cast<u_short>(this->port));
+
+        this->relayAddrs.push_back(addr);
+      }
+
+      public: std::vector<std::string> RelayAddresses() const
+      {
+        std::vector<std::string> result;
+
+        for (auto const &addr : this->relayAddrs) {
+          result.push_back(inet_ntoa(addr.sin_addr));
+        }
+
+        return result;
+      }
+
       /// \brief Broadcast periodic heartbeats.
       private: void UpdateHeartbeat()
       {
@@ -1420,26 +1451,6 @@ namespace gz
         return true;
       }
 
-      /// \brief Register a new relay address.
-      /// \param[in] _ip New IP address.
-      private: void AddRelayAddress(const std::string &_ip)
-      {
-        // Sanity check: Make sure that this IP address is not already saved.
-        for (auto const &addr : this->relayAddrs)
-        {
-          if (addr.sin_addr.s_addr == inet_addr(_ip.c_str()))
-            return;
-        }
-
-        sockaddr_in addr;
-        memset(&addr, 0, sizeof(addr));
-        addr.sin_family = AF_INET;
-        addr.sin_addr.s_addr = inet_addr(_ip.c_str());
-        addr.sin_port = htons(static_cast<u_short>(this->port));
-
-        this->relayAddrs.push_back(addr);
-      }
-
       /// \brief Default activity interval value (ms.).
       /// \sa ActivityInterval.
       /// \sa SetActivityInterval.
diff --git a/include/gz/transport/Node.hh b/include/gz/transport/Node.hh
index f4eb284b4..c7f805b13 100644
--- a/include/gz/transport/Node.hh
+++ b/include/gz/transport/Node.hh
@@ -762,6 +762,9 @@ namespace gz
       public: std::optional<TopicStatistics> TopicStats(
                   const std::string &_topic) const;
 
+      public: void AddGlobalRelay(const std::string& relay_address);
+      public: std::vector<std::string> GlobalRelays() const;
+
       /// \brief Get a pointer to the shared node (singleton shared by all the
       /// nodes).
       /// \return The pointer to the shared node.
diff --git a/include/gz/transport/NodeShared.hh b/include/gz/transport/NodeShared.hh
index 56df051dc..461266944 100644
--- a/include/gz/transport/NodeShared.hh
+++ b/include/gz/transport/NodeShared.hh
@@ -284,6 +284,9 @@ namespace gz
       public: std::optional<TopicStatistics> TopicStats(
                   const std::string &_topic) const;
 
+      public: void AddGlobalRelay(const std::string& relay_address);
+      public: std::vector<std::string> GlobalRelays() const;
+
       /// \brief Constructor.
       protected: NodeShared();
 
diff --git a/src/Node.cc b/src/Node.cc
index 4dbc29598..a286be9f2 100644
--- a/src/Node.cc
+++ b/src/Node.cc
@@ -1139,3 +1139,11 @@ bool Node::RequestRaw(const std::string &_topic,
   bool executed = this->Request(_topic, *req, _timeout, *res, _result);
   return executed && res->SerializeToString(&_response);
 }
+
+void Node::AddGlobalRelay(const std::string& relay_address) {
+  Shared()->AddGlobalRelay(relay_address);
+}
+
+std::vector<std::string> Node::GlobalRelays() const {
+  return Shared()->GlobalRelays();
+}
diff --git a/src/NodeShared.cc b/src/NodeShared.cc
index ce9e8a0c8..f3ce36791 100644
--- a/src/NodeShared.cc
+++ b/src/NodeShared.cc
@@ -1941,3 +1941,22 @@ int NodeSharedPrivate::NonNegativeEnvVar(const std::string &_envVar,
   }
   return numVal;
 }
+
+void NodeShared::AddGlobalRelay(const std::string& relay_address) {
+  dataPtr->msgDiscovery->AddRelayAddress(relay_address);
+  dataPtr->srvDiscovery->AddRelayAddress(relay_address);
+}
+
+std::vector<std::string> NodeShared::GlobalRelays() const {
+  // Merge relays from message and service discovery. They should be identical
+  // since they're typically build from the same sources.
+  //
+  // This is confusing - do we want to add different handling here?
+  auto msgRelays = dataPtr->msgDiscovery->RelayAddresses();
+  std::set<std::string> msgRelaySet(msgRelays.cbegin(), msgRelays.cend());
+  auto srvRelays = dataPtr->srvDiscovery->RelayAddresses();
+  std::set<std::string> srvRelaySet(srvRelays.cbegin(), srvRelays.cend());
+  srvRelaySet.merge(msgRelaySet);
+
+  return std::vector<std::string>(srvRelaySet.cbegin(), srvRelaySet.cend());
+}