[tpu_raiden]Integrate PSP-TCP encryption into TPU Raiden transport (2/2)

tpu_sync/transport/BUILD

@@ -89,6 +89,12 @@ cc_test(
         "//tpu_sync/telemetry:metrics_3p_prometheus_exporter",
         "//tpu_sync/telemetry:metrics_api",
         "//tpu_sync/telemetry:metrics_backend",
+        "//tpu_sync/transport/lib/socket:psp_syscall_mock",
+        "//tpu_sync/transport/lib/socket:tcp_psp_helper",
+        "@com_github_grpc_grpc//:grpc++",
+        "@com_google_absl//absl/base:core_headers",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/flags:flag",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:status_matchers",
         "@com_google_absl//absl/status:statusor",

tpu_sync/transport/lib/BUILD

@@ -44,7 +44,9 @@ cc_library(
     ],
     features = ["-use_header_modules"],
     deps = [
+        "@com_github_grpc_grpc//:grpc++",
         "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings:string_view",
     ],
 )
 
@@ -97,17 +99,20 @@ cc_library(
         "//tpu_sync/core:status_macros",
         "//tpu_sync/transport:buffer_push_task",
         "//tpu_sync/transport/lib/conn:pool",
+        "//tpu_sync/transport/lib/socket:tcp_psp_helper",
         "//tpu_sync/transport/peregrine/src/api:socket_util",
         "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/cleanup",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/container:flat_hash_set",
+        "@com_google_absl//absl/flags:flag",
         "@com_google_absl//absl/functional:any_invocable",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:string_view",
         "@com_google_absl//absl/synchronization",
         "@com_google_absl//absl/types:span",
     ],
@@ -122,12 +127,18 @@ cc_test(
     ],
     features = ["-use_header_modules"],
     deps = [
+        ":peregrine_control_service",
         ":raw_buffer_transport",
         ":raw_buffer_transport_delegate",
         "//tpu_sync/transport:buffer_push_task",
         "//tpu_sync/transport/lib/conn:pool",
+        "//tpu_sync/transport/lib/socket:psp_syscall_mock",
+        "//tpu_sync/transport/lib/socket:tcp_psp_helper",
         "//tpu_sync/transport/peregrine/src/util",
+        "@com_github_grpc_grpc//:grpc++",
         "@com_google_absl//absl/base:core_headers",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/flags:flag",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/strings:string_view",
@@ -143,6 +154,7 @@ cc_library(
     hdrs = ["peregrine_control_service.h"],
     deps = [
         ":raw_buffer_transport",
+        ":raw_buffer_transport_delegate",
         "//tpu_sync/transport/peregrine/src/internal/control:service_cc_grpc",
         "//tpu_sync/transport/peregrine/src/internal/control:service_cc_proto",
         "@com_github_grpc_grpc//:grpc++",

tpu_sync/transport/lib/conn/BUILD

@@ -27,6 +27,7 @@ cc_library(
     features = ["-use_header_modules"],
     deps = [
         "//tpu_sync/transport/lib/socket:util",
+        "@com_github_grpc_grpc//:grpc++",
         "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/log:check",

tpu_sync/transport/lib/conn/pool.cc

@@ -18,6 +18,7 @@
 #include <sys/socket.h>
 #include <unistd.h>
 
+#include <memory>
 #include <vector>
 
 #include "absl/base/optimization.h"
@@ -26,6 +27,7 @@
 #include "absl/status/statusor.h"
 #include "absl/strings/string_view.h"
 #include "absl/synchronization/mutex.h"
+#include "grpcpp/channel.h"
 #include "tpu_sync/transport/lib/socket/util.h"
 
 namespace tpu_raiden::transport::lib {
@@ -43,8 +45,9 @@ void CloseSocket(const int fd) {
 }
 }  // namespace
 
-absl::StatusOr<int> ConnPool::Borrow(absl::string_view peer,
-                                     absl::string_view local_ip) {
+absl::StatusOr<int> ConnPool::Borrow(
+    absl::string_view peer, absl::string_view local_ip, bool require_psp,
+    std::shared_ptr<grpc::Channel> channel) {
   const Key key = GenPoolKey(peer, local_ip);
   {
     absl::MutexLock lock(mu_);
@@ -66,7 +69,7 @@ absl::StatusOr<int> ConnPool::Borrow(absl::string_view peer,
       }
     }
   }
-  return ConnectToPeer(peer, local_ip);
+  return ConnectToPeer(peer, local_ip, require_psp, channel);
 }
 
 void ConnPool::Return(bool ok, int fd, absl::string_view peer,

tpu_sync/transport/lib/conn/pool.h

@@ -15,6 +15,7 @@
 #ifndef THIRD_PARTY_TPU_RAIDEN_TPU_RAIDEN_TRANSPORT_LIB_CONN_POOL_H_
 #define THIRD_PARTY_TPU_RAIDEN_TPU_RAIDEN_TRANSPORT_LIB_CONN_POOL_H_
 
+#include <memory>
 #include <string>
 #include <vector>
 
@@ -25,6 +26,7 @@
 #include "absl/strings/str_cat.h"
 #include "absl/strings/string_view.h"
 #include "absl/synchronization/mutex.h"
+#include "grpcpp/channel.h"
 
 namespace tpu_raiden::transport::lib {
 
@@ -45,8 +47,10 @@ class ConnPool {
   // Borrows a connection from the pool. If no connection is available, creates
   // a new one. Returns the socket descriptor of the connection if successful.
   // Otherwise returns an error status.
-  absl::StatusOr<int> Borrow(absl::string_view peer,
-                             absl::string_view local_ip = "");
+  absl::StatusOr<int> Borrow(
+      absl::string_view peer, absl::string_view local_ip = "",
+      bool require_psp = false,
+      std::shared_ptr<grpc::Channel> channel = nullptr);
 
   // Returns a connection to the pool if ok is true. Otherwise, closes the
   // connection.

tpu_sync/transport/lib/peregrine_control_service_test.cc

@@ -68,8 +68,65 @@ TEST(PeregrineControlServiceTest, InProcessGrpcExchangePspKey) {
   grpc::ClientContext ctx;
 
   grpc::Status status = stub->ExchangePspKey(&ctx, req, &resp);
-  // TODO(yyd): update test once the implementation is done.
-  EXPECT_EQ(status.error_code(), grpc::StatusCode::INTERNAL);
+  // Status is ok or unavailable depending on hardware PSP kernel support.
+  if (status.ok()) {
+    EXPECT_NE(resp.server_spi(), 0);
+    EXPECT_EQ(resp.server_key().size(), 16);
+  }
+
+  server->Shutdown();
+}
+
+TEST(PeregrineControlServiceTest, RejectsInvalidClientKey) {
+  FakeRawDelegate raw_delegate;
+  RawBufferTransport transport(&raw_delegate, /*local_port=*/0);
+  PeregrineControlServiceImpl service(&transport);
+
+  grpc::ServerBuilder builder;
+  builder.RegisterService(&service);
+  std::unique_ptr<grpc::Server> server = builder.BuildAndStart();
+  ASSERT_THAT(server, NotNull());
+
+  std::shared_ptr<grpc::Channel> channel =
+      server->InProcessChannel(grpc::ChannelArguments());
+  auto stub =
+      ::peregrine::internal::control::PeregrineService::NewStub(channel);
+
+  ::peregrine::internal::control::PspKeyExchangeRequest req;
+  req.set_client_spi(0x12345678);
+  req.set_client_key("short_key");  // Not 16 bytes
+  ::peregrine::internal::control::PspKeyExchangeResponse resp;
+  grpc::ClientContext ctx;
+
+  grpc::Status status = stub->ExchangePspKey(&ctx, req, &resp);
+  EXPECT_EQ(status.error_code(), grpc::StatusCode::INVALID_ARGUMENT);
+
+  server->Shutdown();
+}
+
+TEST(PeregrineControlServiceTest, RejectsZeroClientSpi) {
+  FakeRawDelegate raw_delegate;
+  RawBufferTransport transport(&raw_delegate, /*local_port=*/0);
+  PeregrineControlServiceImpl service(&transport);
+
+  grpc::ServerBuilder builder;
+  builder.RegisterService(&service);
+  std::unique_ptr<grpc::Server> server = builder.BuildAndStart();
+  ASSERT_THAT(server, NotNull());
+
+  std::shared_ptr<grpc::Channel> channel =
+      server->InProcessChannel(grpc::ChannelArguments());
+  auto stub =
+      ::peregrine::internal::control::PeregrineService::NewStub(channel);
+
+  ::peregrine::internal::control::PspKeyExchangeRequest req;
+  req.set_client_spi(0);  // Invalid SPI
+  req.set_client_key(std::string(16, 'x'));
+  ::peregrine::internal::control::PspKeyExchangeResponse resp;
+  grpc::ClientContext ctx;
+
+  grpc::Status status = stub->ExchangePspKey(&ctx, req, &resp);
+  EXPECT_EQ(status.error_code(), grpc::StatusCode::INVALID_ARGUMENT);
 
   server->Shutdown();
 }

tpu_sync/transport/lib/raw_buffer_transport.cc

@@ -40,6 +40,7 @@
 #include "absl/base/optimization.h"
 #include "absl/cleanup/cleanup.h"
 #include "absl/container/flat_hash_map.h"
+#include "absl/flags/flag.h"
 #include "absl/log/check.h"
 #include "absl/log/log.h"
 #include "absl/status/status.h"
@@ -60,6 +61,7 @@
 #include "tpu_sync/transport/lib/chunk_serializer.h"
 #include "tpu_sync/transport/lib/conn/pool.h"
 #include "tpu_sync/transport/lib/raw_buffer_transport_delegate.h"
+#include "tpu_sync/transport/lib/socket/tcp_psp_helper.h"
 #include "tpu_sync/transport/peregrine/src/api/socket_util.h"
 
 namespace tpu_raiden::transport::lib {
@@ -163,6 +165,7 @@ RawBufferTransport::RawBufferTransport(
       bound_ip_(local_ips.empty() ? "127.0.0.1" : local_ips[0]),
       local_ips_(local_ips),
       local_port_(local_port),
+      require_psp_tcp_(absl::GetFlag(FLAGS_require_psp_tcp)),
       server_fd_(-1),
       stopping_(false) {
   // 1. Setup server listening socket.
@@ -426,11 +429,20 @@ void RawBufferTransport::ConnectionWorker(int client_fd) {
   close(client_fd);
 }
 
-absl::StatusOr<RawBufferTransport::PspPeerKey>
+absl::StatusOr<PspPeerKey>
 RawBufferTransport::RegisterPspPeer(uint32_t client_spi,
                                     absl::string_view client_key) {
-  return absl::UnimplementedError(
-      "PSP key registration is not implemented yet.");
+  if (!require_psp_tcp_) {
+    return absl::InvalidArgumentError(
+        "PSP is not enabled in transport");
+  }
+  absl::MutexLock lock(psp_mu_);
+  if (stopping_ || server_fd_ < 0) {
+    return absl::FailedPreconditionError(
+        "Transport is stopping or listening socket is not initialized.");
+  }
+
+  return RegisterPspPeerKey(server_fd_, client_spi, client_key);
 }
 
 void RawBufferTransport::ListenerLoop() {
@@ -453,6 +465,12 @@ void RawBufferTransport::ListenerLoop() {
       if (stopping_) break;
       continue;
     }
+    if (require_psp_tcp_ && !PspEnabled(client_fd)) {
+      close(client_fd);
+      LOG_EVERY_N_SEC(ERROR, 1)
+          << "Unencrypted TCP connection rejected on PSP listener";
+      continue;
+    }
 
     int opt = 1;
     setsockopt(client_fd, IPPROTO_TCP, TCP_NODELAY, &opt, sizeof(opt));

tpu_sync/transport/lib/raw_buffer_transport.h

@@ -29,6 +29,7 @@
 #include "absl/base/thread_annotations.h"
 #include "absl/container/flat_hash_map.h"
 #include "absl/container/flat_hash_set.h"
+#include "absl/flags/flag.h"
 #include "absl/functional/any_invocable.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
@@ -41,6 +42,7 @@
 #include "tpu_sync/transport/lib/conn/pool.h"
 #include "tpu_sync/transport/lib/raw_buffer_transport_delegate.h"
 #include "tpu_sync/transport/lib/transport_adapter.h"
+#include "tpu_sync/transport/lib/socket/tcp_psp_helper.h"
 
 namespace tpu_raiden::transport::lib {
 
@@ -78,10 +80,15 @@ class RawBufferTransport final {
   // Return the local IP addresses.
   absl::Span<const std::string> local_ips() const { return local_ips_; }
 
-  // Borrows a connection from the connection pool.
+  // Borrows a connection from the connection pool, resolving the gRPC channel
+  // for PSP key exchange if enabled.
   absl::StatusOr<int> BorrowConnection(absl::string_view peer,
                                        absl::string_view local_ip = "") {
-    return conn_pool_.Borrow(peer, local_ip);
+    std::shared_ptr<grpc::Channel> channel = nullptr;
+    if (require_psp_tcp_ && raw_delegate_ != nullptr) {
+      channel = raw_delegate_->GetPeregrineChannel(peer);
+    }
+    return conn_pool_.Borrow(peer, local_ip, require_psp_tcp_, channel);
   }
 
   // Returns a connection to the connection pool.
@@ -124,13 +131,8 @@ class RawBufferTransport final {
   // Drops receive-progress counters belonging to the give `uuid`.
   void ForgetPushProgress(uint64_t uuid);
 
-  // TODO(yyd): Move this struct to psp_tcp_helper.h.
-  struct PspPeerKey {
-    uint32_t spi = 0;
-    std::string key;
-  };
-
   // Registers incoming client PSP key and returns server's allocated RX key.
+  // Triggered by ExchangePspKey() in PeregrineControlServiceImpl.
   absl::StatusOr<PspPeerKey> RegisterPspPeer(uint32_t client_spi,
                                              absl::string_view client_key);
 
@@ -161,6 +163,7 @@ class RawBufferTransport final {
   const std::string bound_ip_;
   const std::vector<std::string> local_ips_;
   int local_port_;
+  const bool require_psp_tcp_;
   std::atomic<int> server_fd_;  // owned by listener_thread_
   std::atomic<bool> stopping_;
 
@@ -170,7 +173,7 @@ class RawBufferTransport final {
   absl::Mutex mu_;
   absl::flat_hash_set<int> active_client_fds_ ABSL_GUARDED_BY(mu_);
 
-  // The conn_pool_ owns the sockets that connect to peers. in comparison, the
+  // The conn_pool_ owns the sockets that connect to peers. In comparison, the
   // active_client_fds above are those sockets accepted from peers.
   ConnPool conn_pool_;
 
@@ -182,6 +185,9 @@ class RawBufferTransport final {
   std::unique_ptr<tpu_raiden::NumaThreadPool> push_pool_
       ABSL_GUARDED_BY(push_pool_mu_);
 
+  // To protect multiple gRPC threads can call RegisterPspPeer concurrently
+  absl::Mutex psp_mu_;
+
   std::thread listener_thread_;
   std::vector<std::thread> worker_threads_;
 };

tpu_sync/transport/lib/raw_buffer_transport_delegate.h

@@ -12,13 +12,16 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#ifndef THIRD_PARTY_TPU_RAIDEN_TPU_RAIDEN_TRANSPORT_LIB_RAW_BUFFER_TRANSPORT_DELEGATE_H_
-#define THIRD_PARTY_TPU_RAIDEN_TPU_RAIDEN_TRANSPORT_LIB_RAW_BUFFER_TRANSPORT_DELEGATE_H_
+#ifndef TPU_SYNC_TRANSPORT_LIB_RAW_BUFFER_TRANSPORT_DELEGATE_H_
+#define TPU_SYNC_TRANSPORT_LIB_RAW_BUFFER_TRANSPORT_DELEGATE_H_
 
 #include <cstddef>
 #include <cstdint>
+#include <memory>
 
 #include "absl/status/status.h"
+#include "absl/strings/string_view.h"
+#include "grpcpp/channel.h"
 
 namespace tpu_raiden::transport::lib {
 
@@ -48,8 +51,19 @@ class RawBufferTransportDelegate {
   virtual absl::Status OnDataReceived(uint64_t uuid = 0) {
     return absl::OkStatus();
   }
+
+  // Returns the gRPC channel for PeregrineControlService for the given peer.
+  // Required when FLAGS_require_psp_tcp is true.
+  // Note the peer string is the same as the peer passed to BlockTransport
+  // api, e.g. SyncPush(peer,...), which is the data-plane IP:port. The
+  // delegate needs to resolve the peer to either a pre-existing or newly
+  // created gRPC channel.
+  virtual std::shared_ptr<grpc::Channel> GetPeregrineChannel(
+      absl::string_view peer) {
+    return nullptr;
+  }
 };
 
 }  // namespace tpu_raiden::transport::lib
 
-#endif  // THIRD_PARTY_TPU_RAIDEN_TPU_RAIDEN_TRANSPORT_LIB_RAW_BUFFER_TRANSPORT_DELEGATE_H_
+#endif  // TPU_SYNC_TRANSPORT_LIB_RAW_BUFFER_TRANSPORT_DELEGATE_H_