connection.rb in tamashii-agent-0.1.11

- old
+ new

@@ -123,12 +123,11 @@
         @driver.binary(Packet.new(Type::AUTH_TOKEN, 0, [Type::CLIENT[:agent], @master.serial_number,Config.token].join(",")).dump)
       end
 
       def start_web_driver
         # TODO: Improve below code
-        socket = Config.use_ssl ? OpenSSL::SSL::SSLSocket.new(self) : self
-        @driver = WebSocket::Driver.client(socket)
+        @driver = WebSocket::Driver.client(self)
         @driver.on :open, proc { |e|
           logger.info "Server opened"
           self.auth_request
           send_auth_request
         }
@@ -139,20 +138,24 @@
         }
         @driver.on :message, proc { |e|
           pkt = Packet.load(e.data)
           process_packet(pkt) if pkt
         }
+        @driver.on :error, proc { |e|
+          logger.error("#{e.message}")
+        }
         @driver.start
         self.connect
       end
 
       def register_socket_io
         _monitor = @selector.register(@io, :r)
         _monitor.value = proc do
           begin
-            msg = @io.recv_nonblock(65535)
-            if msg.empty?
+            msg = @io.read_nonblock(4096, exception: false)
+            next if msg == :wait_readable
+            if msg.nil?
               # socket closed
               logger.info "No message received from server. Connection reset"
               close_socket_io
               self.reset
               sleep 1
@@ -167,10 +170,14 @@
         end
       end
 
       def try_create_socket
         logger.info "try to open socket..."
-        TCPSocket.new(@host, @port)
+        if Config.use_ssl
+          OpenSSL::SSL::SSLSocket.new(TCPSocket.new(@host, @port)).connect
+        else
+          TCPSocket.new(@host, @port)
+        end
       rescue
         nil
       end
 
       def close_socket_io