lib/invoker/power/balancer.rb in invoker-1.1.0 vs lib/invoker/power/balancer.rb in invoker-1.2.0.pre

@@ -1,19 +1,26 @@
 require 'em-proxy'
 require 'http-parser'
 
 module Invoker
   module Power
-    class BalancerConnection < EventMachine::ProxyServer::Connection
+    class InvokerHttpProxy < EventMachine::ProxyServer::Connection
       attr_accessor :host, :ip, :port
       def set_host(host, selected_backend)
         self.host = host
         self.ip = selected_backend[:host]
         self.port = selected_backend[:port]
       end
     end
 
+    class InvokerHttpsProxy < InvokerHttpProxy
+      def post_init
+        super
+        start_tls
+      end
+    end
+
     class BalancerParser
       attr_accessor :host, :parser
       def initialize
         @parser = HTTP::Parser.new()
         @header = {}
@@ -42,41 +49,58 @@
         @parser << data
       end
     end
 
     class Balancer
-      attr_accessor :connection, :http_parser, :session
+      attr_accessor :connection, :http_parser, :session, :protocol
       DEV_MATCH_REGEX = /([\w-]+)\.dev(\:\d+)?$/
 
       def self.run(options = {})
-        EventMachine.start_server('0.0.0.0', Invoker.config.http_port,
-                                  BalancerConnection, options) do |connection|
-          balancer = Balancer.new(connection)
+        start_http_proxy(InvokerHttpProxy, 'http', options)
+        start_http_proxy(InvokerHttpsProxy, 'https', options)
+      end
+
+      def self.start_http_proxy(proxy_class, protocol, options)
+        port = protocol == 'http' ? Invoker.config.http_port : Invoker.config.https_port
+        EventMachine.start_server('0.0.0.0', port,
+                                  proxy_class, options) do |connection|
+          balancer = Balancer.new(connection, protocol)
           balancer.install_callbacks
         end
       end
 
-      def initialize(connection)
+      def initialize(connection, protocol)
         @connection = connection
+        @protocol = protocol
         @http_parser = BalancerParser.new()
         @session = nil
         @buffer = []
       end
 
+      # insert X_FORWARDED_PROTO_ so as rails can identify the request as coming from
+      # https
+      def insert_forwarded_proto_header(data)
+        if data =~ /\r\n\r\n/ && data !~ /X_FORWARDED_PROTO/i && protocol == 'https'
+          data.gsub(/\r\n\r\n/, "\r\nX_FORWARDED_PROTO: #{protocol}\r\n\r\n")
+        else
+          data
+        end
+      end
+
       def install_callbacks
         http_parser.on_headers_complete { |header| headers_received(header) }
-        connection.on_data {|data| upstream_data(data) }
+        connection.on_data { |data| upstream_data(data) }
         connection.on_response { |backend, data| backend_data(backend, data) }
         connection.on_finish { |backend, name| frontend_disconnect(backend, name) }
       end
 
       def headers_received(header)
         @session = UUID.generate()
         dns_check_response = select_backend_config(header['Host'])
         if dns_check_response && dns_check_response.port
           connection.server(session, host: '0.0.0.0', port: dns_check_response.port)
-          connection.relay_to_servers(@buffer.join)
+          connection.relay_to_servers(insert_forwarded_proto_header(@buffer.join))
           @buffer = []
         else
           return_error_page(404)
           connection.close_connection_after_writing
         end
@@ -86,11 +110,11 @@
         unless session
           @buffer << data
           append_for_http_parsing(data)
           nil
         else
-          data
+          insert_forwarded_proto_header(data)
         end
       end
 
       def append_for_http_parsing(data)
         http_parser << data
@@ -105,10 +129,10 @@
       end
 
       def frontend_disconnect(backend, name)
         http_parser.reset
         unless @backend_data
-          Invoker::Logger.puts("\nApplication not running. Returning error page.".color(:red))
+          Invoker::Logger.puts("\nApplication #{name} not running. Returning error page.".color(:red))
           return_error_page(503)
         end
         @backend_data = false
         connection.close_connection_after_writing if backend == session
       end