Add TLS support to the benchmark script
diff --git a/lib/rb/benchmark/benchmark.rb b/lib/rb/benchmark/benchmark.rb
index 3dc67dd..4a520a5 100644
--- a/lib/rb/benchmark/benchmark.rb
+++ b/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 @@
@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 @@
@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 @@
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 @@
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
diff --git a/lib/rb/benchmark/client.rb b/lib/rb/benchmark/client.rb
index 703dc8f..693bf60 100644
--- a/lib/rb/benchmark/client.rb
+++ b/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 @@
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
diff --git a/lib/rb/benchmark/server.rb b/lib/rb/benchmark/server.rb
index 74e13f4..153eb0f 100644
--- a/lib/rb/benchmark/server.rb
+++ b/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 @@
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 @@
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
diff --git a/lib/rb/benchmark/thin_server.rb b/lib/rb/benchmark/thin_server.rb
index 4de2eef..b9e2db2 100644
--- a/lib/rb/benchmark/thin_server.rb
+++ b/lib/rb/benchmark/thin_server.rb
@@ -18,6 +18,7 @@
#
$:.unshift File.dirname(__FILE__) + '/../lib'
+$:.unshift File.dirname(__FILE__) + '/../ext'
require 'thrift'
$:.unshift File.dirname(__FILE__) + "/gen-rb"
require 'benchmark_service'