Ruby TLS Sample Echo Server w/ Proxy Protocol v1 Support

tls_proxy_enabled_server.rb

#!/usr/bin/env ruby

require 'openssl'
require 'socket'
require 'timeout'

Thread.abort_on_exception = true

SERVER_KEY = OpenSSL::PKey::RSA.new(2048)

SERVER_CERT = OpenSSL::X509::Certificate.new
SERVER_CERT.version = 2
SERVER_CERT.not_after = (Time.now + (86400 * 365 * 3))
SERVER_CERT.not_before = (Time.now - 86400)

ef = OpenSSL::X509::ExtensionFactory.new
ef.subject_certificate = SERVER_CERT
ef.issuer_certificate = SERVER_CERT

SERVER_CERT.serial = 1
SERVER_CERT.subject = OpenSSL::X509::Name.parse('CN=*')
SERVER_CERT.issuer = SERVER_CERT.subject
SERVER_CERT.public_key = SERVER_KEY.public_key

SERVER_CERT.add_extension(ef.create_extension('basicConstraints', 'CA:true,pathlen:0', true))
SERVER_CERT.add_extension(ef.create_extension('extendedKeyUsage', 'serverAuth,clientAuth'))
SERVER_CERT.add_extension(ef.create_extension('keyUsage', 'cRLSign,keyCertSign,digitalSignature,nonRepudiation', true))
SERVER_CERT.add_extension(ef.create_extension('subjectKeyIdentifier', 'hash', false))

SERVER_CERT.sign(SERVER_KEY, OpenSSL::Digest::SHA384.new)

sock = TCPServer.new('::', 4443)

sock.setsockopt(Socket::SOL_SOCKET, Socket::SO_REUSEADDR, 1)
sock.setsockopt(Socket::SOL_SOCKET, Socket::SO_REUSEPORT, 1)
sock.setsockopt(SOL_SOCKET, SO_KEEPALIVE, 1) if defined?(SO_KEEPALIVE)

sock.listen(10)

cert_store = OpenSSL::X509::Store.new
cert_store.set_default_paths

ssl_context = OpenSSL::SSL::SSLContext.new(:TLSv1_2_server)
ssl_context.cert_store = cert_store
ssl_context.key = SERVER_KEY
ssl_context.cert = SERVER_CERT
ssl_context.ciphers = ssl_context.ciphers.select { |c| c[3] >= 256 }

PROXY_PREAMBLE = 'PROXY '

client_threads = []

begin
  loop do
    if IO.select([sock], nil, nil, 5)
      unencrypted_socket = sock.accept

      client_thread = Thread.new(unencrypted_socket, ssl_context) do |sock, ssl_context|
        conn_data = {}
        begin
          # We need to wait for the socket to become ready before continuing
          # on. This also catches early proxy keepalive disconnects.
          unless IO.select([sock], nil, nil, 1)
            # Disconnect / never became ready
            next
          end

          _, conn_data[:server_port], _, conn_data[:server_ip] = sock.addr(:numeric)

          if sock.recv(PROXY_PREAMBLE.length, Socket::MSG_PEEK) == PROXY_PREAMBLE
            raw_preamble = ''
            loop do
              raw_preamble << sock.recv(1)
              break if raw_preamble[-1] == "\n"
            end

            _, _, conn_data[:client_ip], conn_data[:proxy_srv_ip], conn_data[:client_port], conn_data[:proxy_srv_port] = raw_preamble.split(' ')
            _, conn_data[:proxy_client_port], _, conn_data[:proxy_client_ip] = sock.peeraddr(:numeric)
          else
            _, conn_data[:client_port], _, conn_data[:client_ip] = sock.peeraddr(:numeric)
          end
          puts conn_data.inspect

          ssl_socket = OpenSSL::SSL::SSLSocket.new(sock, ssl_context)
          ssl_socket.sync_close = true

          Timeout.timeout(3) { ssl_socket.accept }

          loop do
            if IO.select([ssl_socket], nil, nil, 5)
              # Simple echo server
              if (data = ssl_socket.read_nonblock(1024))
                ssl_socket.write(data)
              end
            end
          end
        rescue EOFError
          # Connection closed
        rescue IOError, OpenSSL::SSL::SSLError
          # Do nothing, let it close
          puts 'Encountered IO error, client likely gone'
        rescue Timeout::Error
          # Also do nothing... let it close
          puts 'Timed out attempting to establish TLS session'
        rescue Errno::ENOTCONN
          # Client disconnected before we could get information about it
          # (likely a keepalive message)
        ensure
          ssl_socket.close if ssl_socket && !ssl_socket.closed?
          sock.close unless if sock && !sock.closed?
        end
      end

      client_threads << client_thread
      client_threads.reject(&:alive?).map(&:join)
      client_threads.select!(&:alive?)
    end
  end
  end
rescue Interrupt
  client_threads.map(&:kill).map(&:join)
  sock.close
end