Add support for websocket PING and PONG frames

runtime/bin/websocket_impl.dart

 // Copyright (c) 2012, the Dart project authors.  Please see the AUTHORS file
 // for details. All rights reserved. Use of this source code is governed by a
 // BSD-style license that can be found in the LICENSE file.
 
 final String _webSocketGUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
 
 class _WebSocketMessageType {
   static final int NONE = 0;
   static final int BINARY = 1;
   static final int TEXT = 2;
-  static final int CLOSE = 3;
 }
 
 
 class _WebSocketOpcode {
   static final int CONTINUATION = 0;
   static final int TEXT = 1;
   static final int BINARY = 2;
   static final int RESERVED_3 = 3;
   static final int RESERVED_4 = 4;
   static final int RESERVED_5 = 5;
@@ skipping 25 matching lines @@
 class _WebSocketProtocolProcessor {
   static final int START = 0;
   static final int LEN_FIRST = 1;
   static final int LEN_REST = 2;
   static final int MASK = 3;
   static final int PAYLOAD = 4;
   static final int CLOSED = 5;
   static final int FAILURE = 6;
 
   _WebSocketProtocolProcessor() {
-    _reset();
+    _prepareForNextFrame();
     _currentMessageType = _WebSocketMessageType.NONE;
   }
 
   /**
    * Process data received from the underlying communication channel.
    */
   void update(List<int> buffer, int offset, int count) {
     int index = offset;
     int lastIndex = offset + count;
     try {
@@ skipping 30 matching lines @@
               if (_currentMessageType != _WebSocketMessageType.NONE) {
                 throw new WebSocketException("Protocol error");
               }
               _currentMessageType = _WebSocketMessageType.BINARY;
               if (onMessageStart != null) {
                 onMessageStart(_WebSocketMessageType.BINARY);
               }
               break;
 
             case _WebSocketOpcode.CLOSE:
-              if (_currentMessageType != _WebSocketMessageType.NONE) {
-                throw new WebSocketException("Protocol error");
-              }
-              _currentMessageType = _WebSocketMessageType.CLOSE;
-              break;
-
             case _WebSocketOpcode.PING:
-              // TODO(sgjesse): Handle ping.
-              throw new UnsupportedOperationException("Web socket PING");
-              break;
-
             case _WebSocketOpcode.PONG:
-              // TODO(sgjesse): Handle pong.
-              throw new UnsupportedOperationException("Web socket PONG");
+              // Control frames cannot be fragmented.
+              if (!_fin) throw new WebSocketException("Protocol error");
               break;
 
             default:
               throw new WebSocketException("Protocol error");
               break;
             }
             _state = LEN_FIRST;
             break;
 
           case LEN_FIRST:
             _masked = (byte & 0x80) != 0;
             _len = byte & 0x7F;
+            if (_isControlFrame() && _len > 126) {
+              throw new WebSocketException("Protocol error");
+            }
             if (_len < 126) {
               _lengthDone();
             } else if (_len == 126) {
               _len = 0;
               _remainingLenBytes = 2;
               _state = LEN_REST;
             } else if (_len == 127) {
               _len = 0;
               _remainingLenBytes = 8;
               _state = LEN_REST;
@@ skipping 17 matching lines @@
             break;
 
           case PAYLOAD:
             // The payload is not handled one byte at a time but in blocks.
             int payload;
             if (lastIndex - index <= _remainingPayloadBytes) {
               payload = lastIndex - index;
             } else {
               payload = _remainingPayloadBytes;
             }
+            _remainingPayloadBytes -= payload;
+
             // Unmask payload if masked.
             if (_masked) {
               for (int i = 0; i < payload; i++) {
                 int maskingByte =
                     ((_maskingKey >> ((3 - _unmaskingIndex) * 8)) & 0xFF);
                 buffer[index + i] = buffer[index + i] ^ maskingByte;
                 _unmaskingIndex = (_unmaskingIndex + 1) % 4;
               }
             }
 
-            switch (_currentMessageType) {
-              case _WebSocketMessageType.NONE:
-                throw new WebSocketException("Protocol error");
-                break;
+            if (_isControlFrame()) {
+              if (payload > 0) {
+                // Allocate a buffer for collecting the control frame
+                // payload if any.
+                if (_controlPayload == null) {
+                  _controlPayload = new List<int>();
+                }
+                _controlPayload.addAll(buffer.getRange(index, payload));
+                index += payload;
+              }
 
-              case _WebSocketMessageType.TEXT:
-              case _WebSocketMessageType.BINARY:
-                if (onMessageData != null) {
-                  onMessageData(buffer, index, payload);
-                }
-                _remainingPayloadBytes -= payload;
-                index += payload;
-                if (_remainingPayloadBytes == 0) {
-                  _frameEnd();
-                }
-                break;
+              if (_remainingPayloadBytes == 0) {
+                _controlFrameEnd();
+              }
+            } else {
+              switch (_currentMessageType) {
+                case _WebSocketMessageType.NONE:
+                  throw new WebSocketException("Protocol error");
+                  break;
 
-              case _WebSocketMessageType.CLOSE:
-                // Allocate a buffer for holding the close payload if any.
-                if (_closePayload == null) {
-                  _closePayload = new List<int>();
-                }
-                _closePayload.addAll(buffer.getRange(index, payload));
-                _remainingPayloadBytes -= payload;
-                index += payload;
-                if (_fin) {
-                  if (_remainingPayloadBytes != 0) {
-                    throw new WebSocketException("Protocol error");
+                case _WebSocketMessageType.TEXT:
+                case _WebSocketMessageType.BINARY:
+                  if (onMessageData != null) {
+                    onMessageData(buffer, index, payload);
                   }
-                  int status;
-                  String reason;
-                  if (_closePayload.length > 0) {
-                    if (_closePayload.length == 1) {
-                      throw new WebSocketException("Protocol error");
-                    }
-                    status = _closePayload[0] << 8 | _closePayload[1];
-                    if (_closePayload.length > 2) {
-                      var decoder = _StringDecoders.decoder(Encoding.UTF_8);
-                      decoder.write(_closePayload.getRange(
-                          2, _closePayload.length - 2));
-                      reason = decoder.decoded;
-                    }
+                  index += payload;
+                  if (_remainingPayloadBytes == 0) {
+                    _messageFrameEnd();
                   }
-                  if (onClosed != null) onClosed(status, reason);
-                  _currentMessageType = _WebSocketMessageType.NONE;
-                  _state = CLOSED;
-                }
-                break;
+                  break;
 
-              default:
-                throw new WebSocketException("Protocol error");
-                break;
+                default:
+                  throw new WebSocketException("Protocol error");
+                  break;
+
+              }
             }
 
             // Hack - as we always do index++ below.
             index--;
-            break;
-
-          default:
-            throw new WebSocketException("Protocol error");
-            break;
         }
 
         // Move to the next byte.
         index++;
       }
     } catch (var e) {
       _reportError(e);
     }
   }
 
@@ skipping 15 matching lines @@
       _startPayload();
     }
   }
 
   void _maskDone() {
     _remainingPayloadBytes = _len;
     _startPayload();
   }
 
   void _startPayload() {
-    // Check whether there is any payload. If not indicate empty
-    // message or close without state and reason.
+    // Check whether there is any payload. If not handle callbacks

[Anders Johnsen, 2012/05/22 08:54:34]

Weird sentence. 

[Søren Gjesse, 2012/05/22 10:34:28]

Rephrased.

+    // without going through the PAYLOAD state..
     if (_remainingPayloadBytes == 0) {
-      if (_currentMessageType ==_WebSocketMessageType.CLOSE) {
-        if (onClosed != null) onClosed(null, null);
-        _state = CLOSED;
+      if (_isControlFrame()) {
+        switch (_opcode) {
+          case _WebSocketOpcode.CLOSE:
+            if (onClosed != null) onClosed(null, null);
+            _state = CLOSED;
+            break;
+          case _WebSocketOpcode.PING:
+            if (onPing != null) onPing(null);
+            break;
+          case _WebSocketOpcode.PONG:
+            if (onPong != null) onPong(null);
+            break;
+        }
+        _prepareForNextFrame();
       } else {
-        _frameEnd();
+        _messageFrameEnd();
       }
     } else {
       _state = PAYLOAD;
     }
   }
 
-  void _frameEnd() {
-    if (_remainingPayloadBytes != 0) {
-      throw new WebSocketException("Protocol error");
-    }
+  void _messageFrameEnd() {
     if (_fin) {
       if (onMessageEnd != null) onMessageEnd();
       _currentMessageType = _WebSocketMessageType.NONE;
     }
-    _reset();
+    _prepareForNextFrame();
   }
 
-  void _reset() {
-    _state = START;
+  void _controlFrameEnd() {
+    switch (_opcode) {
+      case _WebSocketOpcode.CLOSE:
+        int status;
+        String reason;
+        if (_controlPayload.length > 0) {
+          if (_controlPayload.length == 1) {
+            throw new WebSocketException("Protocol error");
+          }
+          status = _controlPayload[0] << 8 | _controlPayload[1];
+          if (_controlPayload.length > 2) {
+            var decoder = _StringDecoders.decoder(Encoding.UTF_8);
+            decoder.write(_controlPayload.getRange(
+                          2, _controlPayload.length - 2));

[Mads Ager (google), 2012/05/22 09:04:13]

Maybe just use four-space indentation here. The arguments are currently aligned
as arguments to write.

[Søren Gjesse, 2012/05/22 10:34:28]

Done.

+            reason = decoder.decoded;
+          }
+        }
+        if (onClosed != null) onClosed(status, reason);
+        _state = CLOSED;
+        break;
+
+      case _WebSocketOpcode.PING:
+        if (onPing != null) onPing(_controlPayload);
+        break;
+
+      case _WebSocketOpcode.PONG:
+        if (onPong != null) onPong(_controlPayload);
+        break;
+    }
+    _prepareForNextFrame();
+  }
+
+  bool _isControlFrame() {
+    return _opcode == _WebSocketOpcode.CLOSE ||
+           _opcode == _WebSocketOpcode.PING ||
+           _opcode == _WebSocketOpcode.PONG;
+  }
+
+  void _prepareForNextFrame() {
+    if (_state != CLOSED && _state != FAILURE) _state = START;
     _fin = null;
     _opcode = null;
     _len = null;
     _masked = null;
     _maskingKey = 0;
     _remainingLenBytes = null;
     _remainingMaskingKeyBytes = null;
     _remainingPayloadBytes = null;
     _unmaskingIndex = 0;
+    _controlPayload = null;
   }
 
   void _reportError(e) {
     // Report the error through the error callback if any. Otherwise
     // throw the error.
     if (onError != null) {
       onError(e);
       _state = FAILURE;
     } else {
       throw e;
     }
   }
 
   int _state;
   bool _fin;
   int _opcode;
   int _len;
   bool _masked;
   int _maskingKey;
   int _remainingLenBytes;
   int _remainingMaskingKeyBytes;
   int _remainingPayloadBytes;
   int _unmaskingIndex;
 
   int _currentMessageType;
-  List<int> _closePayload;
+  List<int> _controlPayload;
 
   Function onMessageStart;
   Function onMessageData;
   Function onMessageEnd;
+  Function onPing;
+  Function onPong;
   Function onClosed;
   Function onError;
 }
 
 
 class _WebSocketConnectionBase  {
   void _socketConnected(Socket socket) {
     _socket = socket;
     _socket.onError = (e) {
       _reportError(e);
       _socket.close();
     };
   }
 
   void _startProcessing(List<int> unparsedData) {
     _WebSocketProtocolProcessor processor = new _WebSocketProtocolProcessor();
     processor.onMessageStart = _onWebSocketMessageStart;
     processor.onMessageData = _onWebSocketMessageData;
     processor.onMessageEnd = _onWebSocketMessageEnd;
+    processor.onPing = _onWebSocketPing;
+    processor.onPong = _onWebSocketPong;
     processor.onClosed = _onWebSocketClosed;
     processor.onError = _onWebSocketError;
     if (unparsedData != null) {
       processor.update(unparsedData, 0, unparsedData.length);
     }
     _socket.onData = () {
       int available = _socket.available();
       List<int> data = new List<int>(available);
       int read = _socket.readList(data, 0, available);
       processor.update(data, 0, read);
@@ skipping 107 matching lines @@
       if (_currentMessageType == _WebSocketMessageType.TEXT) {
         _onMessage(_decoder.decoded);
       } else {
         _onMessage(_outputStream.contents());
       }
     }
     _decoder = null;
     _outputStream = null;
   }
 
+  _onWebSocketPing(List<int> payload) {
+    _sendFrame(_WebSocketOpcode.PONG, payload);
+  }
+
+  _onWebSocketPong(List<int> payload) {
+    // Currently pong messages are ignored.
+  }
+
   _onWebSocketClosed(int status, String reason) {
     _closeReceived = true;
     if (_onClosed != null) _onClosed(status, reason);
     if (_closeSent) {
       // Got close frame in response to close frame. Now close the socket.
       if (_closeTimer != null) _closeTimer.cancel();
       _socket.close();
     } else {
       close(status);
     }
   }
 
   _onWebSocketError(e) {
     _reportError(e);
     _socket.close();
   }
 
-  _sendFrame(int opcode, List<int> data) {
+  _sendFrame(int opcode, [List<int> data]) {
     bool mask = false;  // Masking not implemented for server.
     int dataLength = data == null ? 0 : data.length;
     // Determine the header size.
     int headerSize = (mask) ? 6 : 2;
     if (dataLength > 65535) {
       headerSize += 8;
     } else if (dataLength > 125) {
       headerSize += 2;
     }
     List<int> header = new List<int>(headerSize);
@@ skipping 332 matching lines @@
 
 class _WebSocketCloseEvent implements CloseEvent {
   _WebSocketCloseEvent(this._wasClean, this._code, this._reason);
   bool get wasClean() => _wasClean;
   int get code() => _code;
   String get reason() => _reason;
   bool _wasClean;
   int _code;
   String _reason;
 }