From b4b395019d6d2440ac7dc7a29f9e985344dda714 Mon Sep 17 00:00:00 2001
From: srivignessh <sriknowledge@gmail.com>
Date: Mon, 18 Nov 2019 01:45:10 -0800
Subject: [PATCH] Add TLS encryption support to yb-ts-cli

Add --certs_dir plumbing to yb-ts-cli from yb-admin.

Fixes #2877

TODO: Unit test
---
 src/yb/tools/ts-cli.cc | 45 +++++++++++++++++++++++++++++++-----------
 1 file changed, 33 insertions(+), 12 deletions(-)

diff --git a/src/yb/tools/ts-cli.cc b/src/yb/tools/ts-cli.cc
index 8d6bebd05252..a5ed8dc956d2 100644
--- a/src/yb/tools/ts-cli.cc
+++ b/src/yb/tools/ts-cli.cc
@@ -57,6 +57,10 @@
 #include "yb/util/net/sockaddr.h"
 #include "yb/rpc/messenger.h"
 #include "yb/rpc/rpc_controller.h"
+#include "yb/rpc/secure_stream.h"
+#include "yb/server/secure.h"
+#include "yb/util/env_util.h"
+#include "yb/util/path_util.h"
 
 using yb::HostPort;
 using yb::rpc::Messenger;
@@ -94,6 +98,9 @@ DEFINE_bool(force, false, "If true, allows the set_flag command to set a flag "
             "which is not explicitly marked as runtime-settable. Such flag changes may be "
             "simply ignored on the server, or may cause the server to crash.");
 
+DEFINE_string(certs_dir_name, "",
+              "Directory with certificates to use for secure server connection.");
+
 // Check that the value of argc matches what's expected, otherwise return a
 // non-zero exit code. Should be used in main().
 #define CHECK_ARGC_OR_RETURN_WITH_USAGE(op, expected) \
@@ -129,7 +136,9 @@ class TsAdminClient {
  public:
   // Creates an admin client for host/port combination e.g.,
   // "localhost" or "127.0.0.1:7050".
-  TsAdminClient(std::string addr, int64_t timeout_millis);
+  TsAdminClient(std::string addr,
+                int64_t timeout_millis,
+                std::string certs_dir);
 
   ~TsAdminClient();
 
@@ -171,7 +180,9 @@ class TsAdminClient {
  private:
   std::string addr_;
   MonoDelta timeout_;
+  std::string certs_dir_;
   bool initted_;
+  std::unique_ptr<rpc::SecureContext> secure_context_;
   std::unique_ptr<rpc::Messenger> messenger_;
   shared_ptr<server::GenericServiceProxy> generic_proxy_;
   gscoped_ptr<tserver::TabletServerServiceProxy> ts_proxy_;
@@ -180,9 +191,12 @@ class TsAdminClient {
   DISALLOW_COPY_AND_ASSIGN(TsAdminClient);
 };
 
-TsAdminClient::TsAdminClient(string addr, int64_t timeout_millis)
+TsAdminClient::TsAdminClient(string addr,
+                             int64_t timeout_millis,
+                             string certs_dir)
     : addr_(std::move(addr)),
       timeout_(MonoDelta::FromMilliseconds(timeout_millis)),
+      certs_dir_(std::move(certs_dir)),
       initted_(false) {}
 
 TsAdminClient::~TsAdminClient() {
@@ -194,15 +208,22 @@ TsAdminClient::~TsAdminClient() {
 Status TsAdminClient::Init() {
   CHECK(!initted_);
 
-  HostPort host_port;
-  RETURN_NOT_OK(host_port.ParseString(addr_, tserver::TabletServer::kDefaultPort));
-  messenger_ = VERIFY_RESULT(MessengerBuilder("ts-cli").Build());
-
-  rpc::ProxyCache proxy_cache(messenger_.get());
-
-  generic_proxy_.reset(new server::GenericServiceProxy(&proxy_cache, host_port));
-  ts_proxy_.reset(new TabletServerServiceProxy(&proxy_cache, host_port));
-  ts_admin_proxy_.reset(new TabletServerAdminServiceProxy(&proxy_cache, host_port));
+  if(!certs_dir_.empty()) {
+    HostPort host_port;
+    RETURN_NOT_OK(host_port.ParseString(addr_, tserver::TabletServer::kDefaultPort));
+    MessengerBuilder messenger_builder("ts-cli");
+    FLAGS_use_client_to_server_encryption = true;
+    FLAGS_certs_dir = certs_dir_;
+    secure_context_ = VERIFY_RESULT(server::SetupSecureContext(
+        "", "", server::SecureContextType::kClientToServer, &messenger_builder));
+    messenger_ = VERIFY_RESULT(messenger_builder.Build());
+    
+    rpc::ProxyCache proxy_cache(messenger_.get());
+
+    generic_proxy_.reset(new server::GenericServiceProxy(&proxy_cache, host_port));
+    ts_proxy_.reset(new TabletServerServiceProxy(&proxy_cache, host_port));
+    ts_admin_proxy_.reset(new TabletServerAdminServiceProxy(&proxy_cache, host_port));
+  }
 
   initted_ = true;
 
@@ -390,7 +411,7 @@ static int TsCliMain(int argc, char** argv) {
 
   string op = GetOp(argc, argv);
 
-  TsAdminClient client(addr, FLAGS_timeout_ms);
+  TsAdminClient client(addr, FLAGS_timeout_ms, FLAGS_certs_dir_name);
 
   RETURN_NOT_OK_PREPEND_FROM_MAIN(client.Init(),
                                   "Unable to establish connection to " + addr);