THRIFT-6120: Bound Ruby TLS accept handshakes

lib/rb/lib/thrift/server/nonblocking_server.rb

@@ -53,7 +53,12 @@ def serve
             break
           end
           next if rd.nil?
-          socket = @server_transport.accept
+          begin
+            socket = @server_transport.accept
+          rescue => e
+            next if defined?(OpenSSL::SSL::SSLError) && e.is_a?(OpenSSL::SSL::SSLError)
+            raise
+          end
           @logger.debug "Accepted socket: #{socket.inspect}"
           @io_manager.add_connection socket
         end

lib/rb/lib/thrift/server/thread_pool_server.rb

@@ -49,7 +49,12 @@ def serve
           Thread.new do
             begin
               loop do
-                client = @server_transport.accept
+                begin
+                  client = @server_transport.accept
+                rescue => e
+                  next if defined?(OpenSSL::SSL::SSLError) && e.is_a?(OpenSSL::SSL::SSLError)
+                  raise
+                end
                 trans = @transport_factory.get_transport(client)
                 prot = @protocol_factory.get_protocol(trans)
                 begin

lib/rb/lib/thrift/transport/ssl_server_socket.rb

@@ -19,6 +19,7 @@
 # under the License.
 #
 
+require 'io/wait'
 require 'socket'
 
 module Thrift
@@ -31,12 +32,71 @@ def initialize(host_or_port, port = nil, ssl_context = nil, client_timeout: DEFA
     attr_accessor :ssl_context
 
     def listen
-      socket = TCPServer.new(@host, @port)
-      @handle = OpenSSL::SSL::SSLServer.new(socket, @ssl_context)
+      tcp_server = TCPServer.new(@host, @port)
+      @handle = OpenSSL::SSL::SSLServer.new(tcp_server, @ssl_context).tap do |server|
+        server.start_immediately = false
+      end
+    end
+
+    def accept
+      return if @handle.nil?
+
+      transport = nil
+      ssl_socket = @handle.accept
+      deadline = Process.clock_gettime(Process::CLOCK_MONOTONIC) + @client_timeout unless @client_timeout.nil? || @client_timeout == 0
+      ssl_socket.to_io.setsockopt(::Socket::IPPROTO_TCP, ::Socket::TCP_NODELAY, 1)
+      accept_ssl_socket(ssl_socket, deadline)
+
+      transport = Socket.new.tap do |accepted_transport|
+        accepted_transport.timeout = @client_timeout
+        accepted_transport.handle = ssl_socket
+      end
+    ensure
+      # Thread#kill leaves $! nil, but marks the thread as aborting.
+      close_accepted_socket(ssl_socket) if transport.nil? || $! || Thread.current.status == "aborting"
     end
 
     def to_s
       "ssl(#{super.to_s})"
     end
+
+    private
+
+    def accept_ssl_socket(ssl_socket, deadline)
+      return ssl_socket.accept unless deadline
+
+      loop do
+        case ssl_socket.accept_nonblock(exception: false)
+        when ssl_socket
+          return ssl_socket
+        when :wait_readable
+          wait_for_handshake(ssl_socket, :read, deadline)
+        when :wait_writable
+          wait_for_handshake(ssl_socket, :write, deadline)
+        else
+          raise TransportException.new(TransportException::NOT_OPEN, "SSL server socket: Unexpected TLS handshake result")
+        end
+      end
+    end
+
+    def wait_for_handshake(ssl_socket, direction, deadline)
+      remaining = deadline - Process.clock_gettime(Process::CLOCK_MONOTONIC)
+      ready = if remaining > 0
+        if direction == :read
+          ssl_socket.to_io.wait_readable(remaining)
+        else
+          ssl_socket.to_io.wait_writable(remaining)
+        end
+      end
+      return if ready
+
+      raise OpenSSL::SSL::SSLError, "SSL server socket: Timed out accepting TLS connection"
+    end
+
+    def close_accepted_socket(ssl_socket)
+      ssl_socket&.close
+    rescue StandardError
+      nil
+    end
   end
 end

lib/rb/spec/nonblocking_server_spec.rb

@@ -265,6 +265,33 @@ def setup_client_thread(result)
     end
   end
 
+  describe "NonblockingServer accept errors" do
+    it "preserves the original error when OpenSSL is not loaded" do
+      hide_const("OpenSSL")
+      error = RuntimeError.new("plain accept failed")
+      server_transport = double(
+        "server transport",
+        :listen => nil,
+        :closed? => false,
+        :close => nil
+      )
+      allow(server_transport).to receive(:accept).and_raise(error)
+      io_manager = double("IOManager", :ensure_closed => nil)
+      server = Thrift::NonblockingServer.new(
+        double("processor"),
+        server_transport,
+        nil,
+        nil,
+        1,
+        Logger.new(IO::NULL)
+      )
+      allow(server).to receive(:start_io_manager).and_return(io_manager)
+      allow(server).to receive(:select).and_return([[server_transport], nil, nil])
+
+      expect { server.serve }.to raise_error(error)
+    end
+  end
+
   describe Thrift::NonblockingServer::IOManager do
     def build_io_manager
       logger = Logger.new(IO::NULL)
@@ -330,11 +357,18 @@ def build_io_manager
   end
 
   describe "#{Thrift::NonblockingServer} with TLS transport" do
+    let(:client_timeout) { Thrift::BaseServerTransport::DEFAULT_CLIENT_TIMEOUT }
+
     before(:each) do
       @port = available_port
       handler = Handler.new
       processor = SpecNamespace::NonblockingService::Processor.new(handler)
-      @transport = Thrift::SSLServerSocket.new('localhost', @port, create_server_ssl_context)
+      @transport = Thrift::SSLServerSocket.new(
+        'localhost',
+        @port,
+        create_server_ssl_context,
+        client_timeout: client_timeout
+      )
       transport_factory = Thrift::FramedTransportFactory.new
       logger = Logger.new(STDERR)
       logger.level = Logger::WARN
@@ -355,7 +389,7 @@ def build_io_manager
 
     after(:each) do
       @clients.each(&:close)
-      @server.shutdown if @server
+      @server.shutdown if @server && @server_thread&.alive?
       @server_thread.join(2) if @server_thread
       @transport.close if @transport
     end
@@ -370,6 +404,22 @@ def build_io_manager
       expect(@server_thread.join(2)).to be_an_instance_of(Thread)
     end
 
+    context "when a TLS handshake times out" do
+      let(:client_timeout) { 0.1 }
+
+      it "continues accepting connections" do
+        stalled_client = TCPSocket.new('localhost', @port)
+
+        expect(Timeout.timeout(1) { stalled_client.read(1) }).to be_nil
+        expect(@server_thread).to be_alive
+
+        client = setup_tls_client
+        expect(client.greeting(true)).to eq(SpecNamespace::Hello.new)
+      ensure
+        stalled_client&.close
+      end
+    end
+
     def setup_tls_client
       transport = Thrift::FramedTransport.new(
         Thrift::SSLSocket.new('localhost', @port, nil, create_client_ssl_context)

lib/rb/spec/server_spec.rb

@@ -234,5 +234,22 @@
       expect(@server_trans).to receive(:close)
       expect { @server.serve }.to(throw_symbol(:stop))
     end
+
+    it "should not enqueue TLS accept errors" do
+      exception_q = @server.instance_variable_get(:@exception_q)
+      expect(@server_trans).to receive(:listen).ordered
+      expect(@server_trans).to receive(:accept).ordered.and_raise(OpenSSL::SSL::SSLError)
+      expect(@server_trans).to receive(:accept).ordered.and_return(@client)
+      expect(@trans).to receive(:get_transport).once.with(@client).and_return(@trans)
+      expect(@prot).to receive(:get_protocol).once.with(@trans).and_return(@prot)
+      allow(Thread).to receive(:new).and_yield
+      expect(@processor).to receive(:process).once.with(@prot, @prot) { throw :stop }
+      expect(@trans).to receive(:close).once
+      expect(@server_trans).to receive(:close)
+
+      catch(:stop) { @server.serve }
+
+      expect(exception_q).to be_empty
+    end
   end
 end