Add TLS support to the benchmark script

Author: kpumuk

lib/rb/benchmark/benchmark.rb

@@ -19,6 +19,7 @@
 
 require 'rubygems'
 $:.unshift File.dirname(__FILE__) + '/../lib'
+$:.unshift File.dirname(__FILE__) + '/../ext'
 require 'thrift'
 require 'stringio'
 
@@ -40,12 +41,13 @@ def initialize(opts)
     @interpreter = opts.fetch(:interpreter, "ruby")
     @host = opts.fetch(:host, ::HOST)
     @port = opts.fetch(:port, ::PORT)
+    @tls = opts.fetch(:tls, false)
   end
 
   def start
     return if @serverclass == Object
     args = (File.basename(@interpreter) == "jruby" ? "-J-server" : "")
-    @pipe = IO.popen("#{@interpreter} #{args} #{File.dirname(__FILE__)}/server.rb #{@host} #{@port} #{@serverclass.name}", "r+")
+    @pipe = IO.popen("#{@interpreter} #{args} #{File.dirname(__FILE__)}/server.rb #{"-tls" if @tls} #{@host} #{@port} #{@serverclass.name}", "r+")
     Marshal.load(@pipe) # wait until the server has started
     sleep 0.4 # give the server time to actually start spawning sockets
   end
@@ -75,6 +77,7 @@ def initialize(opts, server)
     @interpreter = opts.fetch(:interpreter, "ruby")
     @server = server
     @log_exceptions = opts.fetch(:log_exceptions, false)
+    @tls = opts.fetch(:tls, false)
   end
 
   def run
@@ -93,13 +96,15 @@ def run
   end
 
   def spawn
-    pipe = IO.popen("#{@interpreter} #{File.dirname(__FILE__)}/client.rb #{"-log-exceptions" if @log_exceptions} #{@host} #{@port} #{@clients_per_process} #{@calls_per_client}")
+    pipe = IO.popen("#{@interpreter} #{File.dirname(__FILE__)}/client.rb #{"-log-exceptions" if @log_exceptions} #{"-tls" if @tls} #{@host} #{@port} #{@clients_per_process} #{@calls_per_client}")
     @pool << pipe
   end
 
   def socket_class
     if @socket
       Thrift::UNIXSocket
+    elsif @tls
+      Thrift::SSLSocket
     else
       Thrift::Socket
     end
@@ -255,12 +260,14 @@ def resolve_const(const)
 args[:class] = resolve_const(ENV['THRIFT_SERVER']) || Thrift::NonblockingServer
 args[:host] = ENV['THRIFT_HOST'] || HOST
 args[:port] = (ENV['THRIFT_PORT'] || PORT).to_i
+args[:tls] = ENV['THRIFT_TLS'] == 'true'
 server = Server.new(args)
 server.start
 
 args = {}
 args[:host] = ENV['THRIFT_HOST'] || HOST
 args[:port] = (ENV['THRIFT_PORT'] || PORT).to_i
+args[:tls] = ENV['THRIFT_TLS'] == 'true'
 args[:num_processes] = (ENV['THRIFT_NUM_PROCESSES'] || 40).to_i
 args[:clients_per_process] = (ENV['THRIFT_NUM_CLIENTS'] || 5).to_i
 args[:calls_per_client] = (ENV['THRIFT_NUM_CALLS'] || 50).to_i

lib/rb/benchmark/client.rb

@@ -18,22 +18,41 @@
 #
 
 $:.unshift File.dirname(__FILE__) + '/../lib'
+$:.unshift File.dirname(__FILE__) + '/../ext'
 require 'thrift'
+require 'openssl'
 $:.unshift File.dirname(__FILE__) + "/gen-rb"
 require 'benchmark_service'
 
 class Client
-  def initialize(host, port, clients_per_process, calls_per_client, log_exceptions)
+  def initialize(host, port, clients_per_process, calls_per_client, log_exceptions, tls)
     @host = host
     @port = port
     @clients_per_process = clients_per_process
     @calls_per_client = calls_per_client
     @log_exceptions = log_exceptions
+    @tls = tls
   end
 
   def run
     @clients_per_process.times do
-      socket = Thrift::Socket.new(@host, @port)
+      socket = if @tls
+        ssl_context = OpenSSL::SSL::SSLContext.new.tap do |ctx|
+          ctx.verify_mode = OpenSSL::SSL::VERIFY_PEER
+          ctx.min_version = OpenSSL::SSL::TLS1_2_VERSION
+
+          keys_dir = File.expand_path("../../../test/keys", __dir__)
+          ctx.ca_file = File.join(keys_dir, "CA.pem")
+          ctx.cert = OpenSSL::X509::Certificate.new(File.open(File.join(keys_dir, "client.crt")))
+          ctx.cert_store = OpenSSL::X509::Store.new
+          ctx.cert_store.add_file(File.join(keys_dir, 'server.pem'))
+          ctx.key = OpenSSL::PKey::RSA.new(File.open(File.join(keys_dir, "client.key")))
+        end
+
+        Thrift::SSLSocket.new(@host, @port, nil, ssl_context)
+      else
+        Thrift::Socket.new(@host, @port)
+      end
       transport = Thrift::FramedTransport.new(socket)
       protocol = Thrift::BinaryProtocol.new(transport)
       client = ThriftBenchmark::BenchmarkService::Client.new(protocol)
@@ -68,7 +87,8 @@ def print_exception(e)
 end
 
 log_exceptions = true if ARGV[0] == '-log-exceptions' and ARGV.shift
+tls = true if ARGV[0] == '-tls' and ARGV.shift
 
 host, port, clients_per_process, calls_per_client = ARGV
 
-Client.new(host, port.to_i, clients_per_process.to_i, calls_per_client.to_i, log_exceptions).run
+Client.new(host, port.to_i, clients_per_process.to_i, calls_per_client.to_i, log_exceptions, tls).run

lib/rb/benchmark/server.rb

@@ -18,7 +18,9 @@
 #
 
 $:.unshift File.dirname(__FILE__) + '/../lib'
+$:.unshift File.dirname(__FILE__) + '/../ext'
 require 'thrift'
+require 'openssl'
 $:.unshift File.dirname(__FILE__) + "/gen-rb"
 require 'benchmark_service'
 
@@ -36,10 +38,26 @@ def fibonacci(n)
     end
   end
 
-  def self.start_server(host, port, serverClass)
+  def self.start_server(host, port, serverClass, tls)
     handler = BenchmarkHandler.new
     processor = ThriftBenchmark::BenchmarkService::Processor.new(handler)
-    transport = ServerSocket.new(host, port)
+    transport = if tls
+      ssl_context = OpenSSL::SSL::SSLContext.new.tap do |ctx|
+        ctx.verify_mode = OpenSSL::SSL::VERIFY_PEER
+        ctx.min_version = OpenSSL::SSL::TLS1_2_VERSION
+
+        keys_dir = File.expand_path("../../../test/keys", __dir__)
+        ctx.ca_file = File.join(keys_dir, "CA.pem")
+        ctx.cert = OpenSSL::X509::Certificate.new(File.open(File.join(keys_dir, "server.crt")))
+        ctx.cert_store = OpenSSL::X509::Store.new
+        ctx.cert_store.add_file(File.join(keys_dir, 'client.pem'))
+        ctx.key = OpenSSL::PKey::RSA.new(File.open(File.join(keys_dir, "server.key")))
+      end
+
+      Thrift::SSLServerSocket.new(host, port, ssl_context)
+    else
+      ServerSocket.new(host, port)
+    end
     transport_factory = FramedTransportFactory.new
     args = [processor, transport, transport_factory, nil, 20]
     if serverClass == NonblockingServer
@@ -68,9 +86,11 @@ def resolve_const(const)
   const and const.split('::').inject(Object) { |k,c| k.const_get(c) }
 end
 
+tls = true if ARGV[0] == '-tls' and ARGV.shift
+
 host, port, serverklass = ARGV
 
-Server.start_server(host, port.to_i, resolve_const(serverklass))
+Server.start_server(host, port.to_i, resolve_const(serverklass), tls)
 
 # let our host know that the interpreter has started
 # ideally we'd wait until the server was serving, but we don't have a hook for that