Add TLS support to the benchmark script

lib/rb/benchmark/.gitignore

@@ -0,0 +1 @@
+*.pem

lib/rb/benchmark/benchmark.rb

@@ -40,12 +40,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 +76,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,7 +95,7 @@ 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
 
@@ -249,18 +251,53 @@ def resolve_const(const)
   const and const.split('::').inject(Object) { |k,c| k.const_get(c) }
 end
 
+def generate_certificate
+  key = OpenSSL::PKey::EC.generate("prime256v1")
+
+  cert = OpenSSL::X509::Certificate.new
+  cert.version = 2
+  cert.serial = 1
+  cert.subject = OpenSSL::X509::Name.parse("/C=US/O=Benchmark/CN=localhost")
+  cert.issuer = cert.subject
+  cert.public_key = key
+  cert.not_before = Time.now
+  cert.not_after = Time.now + 3600
+
+  # Add extensions
+  ef = OpenSSL::X509::ExtensionFactory.new
+  ef.subject_certificate = cert
+  ef.issuer_certificate = cert
+  cert.add_extension(ef.create_extension("basicConstraints", "CA:TRUE", true))
+  cert.add_extension(ef.create_extension("subjectAltName", "DNS:localhost,IP:127.0.0.1", false))
+
+  cert.sign(key, OpenSSL::Digest.new("SHA256"))
+
+  [cert, key]
+end
+
+if ENV['THRIFT_TLS']
+  puts "Generating TLS certificate and key..."
+  require 'openssl'
+
+  cert, key = generate_certificate
+  File.write(File.expand_path("cert.pem", __dir__), cert.to_pem)
+  File.write(File.expand_path("key.pem", __dir__), key.to_pem)
+end
+
 puts "Starting server..."
 args = {}
 args[:interpreter] = ENV['THRIFT_SERVER_INTERPRETER'] || ENV['THRIFT_INTERPRETER'] || "ruby"
 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

@@ -23,17 +23,33 @@
 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_NONE
+          ctx.min_version = OpenSSL::SSL::TLS1_2_VERSION
+
+          # Load certificate chain and private key
+          certs = OpenSSL::X509::Certificate.load_file(File.expand_path("cert.pem", __dir__))
+          pkey = OpenSSL::PKey.read(File.binread(File.expand_path("key.pem", __dir__)))
+          ctx.add_certificate(certs.first, pkey, *certs[1..])
+
+          ctx
+        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 +84,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

@@ -36,10 +36,24 @@ 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_NONE
+        ctx.min_version = OpenSSL::SSL::TLS1_2_VERSION
+
+        certs = OpenSSL::X509::Certificate.load_file(File.expand_path("cert.pem", __dir__))
+        pkey = OpenSSL::PKey.read(File.binread(File.expand_path("key.pem", __dir__)))
+        ctx.add_certificate(certs.first, pkey, *certs[1..])
+
+        ctx
+      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 +82,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