lib/invoker/power/balancer.rb in invoker-1.4.1 vs lib/invoker/power/balancer.rb in invoker-1.5.1

@@ -20,11 +20,11 @@
         start_tls
       end
     end
 
     class Balancer
-      attr_accessor :connection, :http_parser, :session, :protocol
+      attr_accessor :connection, :http_parser, :session, :protocol, :upgraded_to
 
       def self.run(options = {})
         start_http_proxy(InvokerHttpProxy, 'http', options)
         start_http_proxy(InvokerHttpsProxy, 'https', options)
       end
@@ -41,10 +41,11 @@
       def initialize(connection, protocol)
         @connection = connection
         @protocol = protocol
         @http_parser = HttpParser.new(protocol)
         @session = nil
+        @upgraded_to = nil
         @buffer = []
       end
 
       def install_callbacks
         http_parser.on_headers_complete { |headers| headers_received(headers) }
@@ -69,21 +70,25 @@
           return
         end
 
         dns_check_response = UrlRewriter.new.select_backend_config(headers['Host'])
         if dns_check_response && dns_check_response.port
-          connection.server(session, host: '0.0.0.0', port: dns_check_response.port)
+          connection.server(session, host: dns_check_response.ip, port: dns_check_response.port)
         else
           return_error_page(404)
           http_parser.reset
           connection.close_connection_after_writing
         end
       end
 
       def upstream_data(data)
-        append_for_http_parsing(data)
-        nil
+        if upgraded_to == "websocket"
+          data
+        else
+          append_for_http_parsing(data)
+          nil
+        end
       end
 
       def append_for_http_parsing(data)
         http_parser << data
       rescue HTTP::Parser::Error
@@ -91,9 +96,16 @@
         connection.close_connection_after_writing
       end
 
       def backend_data(backend, data)
         @backend_data = true
+
+        # check backend data for websockets connection. check for upgrade headers
+        # - Upgarde: websocket\r\n
+        if data =~ /Upgrade: websocket/
+          @upgraded_to = "websocket"
+        end
+
         data
       end
 
       def frontend_disconnect(backend, name)
         http_parser.reset