From 043f13a9f4eac6119f6f981731207d412f72236f Mon Sep 17 00:00:00 2001
From: Hiroshi Ichikawa <gimite@gmail.com>
Date: Tue, 30 Jul 2013 15:25:57 +0900
Subject: [PATCH] Fixing a bug handling received PONG frame. #21

---
 lib/web_socket.rb | 84 ++++++++++++++++++++++++++---------------------
 1 file changed, 46 insertions(+), 38 deletions(-)

diff --git a/lib/web_socket.rb b/lib/web_socket.rb
index 4fd2b2b..59e8236 100644
--- a/lib/web_socket.rb
+++ b/lib/web_socket.rb
@@ -197,46 +197,51 @@ def receive()
           end
         
         else
-          begin
-            bytes = read(2).unpack("C*")
-            fin = (bytes[0] & 0x80) != 0
-            opcode = bytes[0] & 0x0f
-            mask = (bytes[1] & 0x80) != 0
-            plength = bytes[1] & 0x7f
-            if plength == 126
-              bytes = read(2)
-              plength = bytes.unpack("n")[0]
-            elsif plength == 127
-              bytes = read(8)
-              (high, low) = bytes.unpack("NN")
-              plength = high * (2 ** 32) + low
-            end
-            if @server && !mask
-              # Masking is required.
-              @socket.close()
-              raise(WebSocket::Error, "received unmasked data")
-            end
-            mask_key = mask ? read(4).unpack("C*") : nil
-            payload = read(plength)
-            payload = apply_mask(payload, mask_key) if mask
-            case opcode
-              when OPCODE_TEXT
-                return force_encoding(payload, "UTF-8")
-              when OPCODE_BINARY
-                raise(WebSocket::Error, "received binary data, which is not supported")
-              when OPCODE_CLOSE
-                close(1005, "", :peer)
-                return nil
-              when OPCODE_PING
-                raise(WebSocket::Error, "received ping, which is not supported")
-              when OPCODE_PONG
-              else
-                raise(WebSocket::Error, "received unknown opcode: %d" % opcode)
+          while true
+            begin
+              bytes = read(2).unpack("C*")
+              fin = (bytes[0] & 0x80) != 0
+              opcode = bytes[0] & 0x0f
+              mask = (bytes[1] & 0x80) != 0
+              plength = bytes[1] & 0x7f
+              if plength == 126
+                bytes = read(2)
+                plength = bytes.unpack("n")[0]
+              elsif plength == 127
+                bytes = read(8)
+                (high, low) = bytes.unpack("NN")
+                plength = high * (2 ** 32) + low
+              end
+              if @server && !mask
+                # Masking is required.
+                @socket.close()
+                raise(WebSocket::Error, "received unmasked data")
+              end
+              mask_key = mask ? read(4).unpack("C*") : nil
+              payload = read(plength)
+              payload = apply_mask(payload, mask_key) if mask
+              if WebSocket.debug
+                $stderr.printf("recv_frame> opcode:%d fin:%d payload:%p\n" % [opcode, fin ? 1 : 0, payload])
+              end
+              case opcode
+                when OPCODE_TEXT
+                  return force_encoding(payload, "UTF-8")
+                when OPCODE_BINARY
+                  raise(WebSocket::Error, "received binary data, which is not supported")
+                when OPCODE_CLOSE
+                  close(1005, "", :peer)
+                  return nil
+                when OPCODE_PING
+                  raise(WebSocket::Error, "received ping, which is not supported")
+                when OPCODE_PONG
+                  next
+                else
+                  raise(WebSocket::Error, "received unknown opcode: %d" % opcode)
+              end
+            rescue EOFError
+              return nil
             end
-          rescue EOFError
-            return nil
           end
-        
       end
     end
     
@@ -319,6 +324,9 @@ def read_header()
     end
 
     def send_frame(opcode, payload, mask)
+      if WebSocket.debug
+        $stderr.printf("send_frame> opcode:%d masked:%d payload:%p\n" % [opcode, mask ? 1 : 0, payload])
+      end
       payload = force_encoding(payload.dup(), "ASCII-8BIT")
       # Setting StringIO's encoding to ASCII-8BIT.
       buffer = StringIO.new(force_encoding("", "ASCII-8BIT"))