Improve web-socket implementation by using making _WebSocketProtocolProcessor a StreamTransformer.

sdk/lib/io/websocket_impl.dart

 // Copyright (c) 2013, 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.
 
 part of dart.io;
 
 const String _webSocketGUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
 
 class _WebSocketMessageType {
   static const int NONE = 0;
@@ skipping 15 matching lines @@
   static const int PING = 9;
   static const int PONG = 10;
   static const int RESERVED_B = 11;
   static const int RESERVED_C = 12;
   static const int RESERVED_D = 13;
   static const int RESERVED_E = 14;
   static const int RESERVED_F = 15;
 }
 
 /**
- * The web socket protocol processor handles the protocol byte stream
- * which is supplied through the [:update:] and [:closed:]
- * methods. As the protocol is processed the following callbacks are
- * called:
+ * The web socket protocol transformer handles the protocol byte stream
+ * which is supplied through the [:handleData:]. As the protocol is processed,
+ * it'll output frame data as either a List<int> or String.
  *
- *   [:onMessageStart:]
- *   [:onMessageData:]
- *   [:onMessageEnd:]
- *   [:onClosed:]
- *
+ * Important infomation about usage: Be sure you use unsubscribeOnError, so the
+ * socket will be closed when the processer encounter an error. Not using it
+ * will lead to undefined behaviour.
  */
-class _WebSocketProtocolProcessor {
+class _WebSocketProtocolTransformer extends StreamEventTransformer {
   static const int START = 0;
   static const int LEN_FIRST = 1;
   static const int LEN_REST = 2;
   static const int MASK = 3;
   static const int PAYLOAD = 4;
   static const int CLOSED = 5;
   static const int FAILURE = 6;
 
-  _WebSocketProtocolProcessor() {
+  _WebSocketProtocolTransformer() {
     _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;
+  void handleData(List<int> buffer, StreamSink sink) {
+    int count = buffer.length;
+    int index = 0;
+    int lastIndex = count;
     try {
       if (_state == CLOSED) {
         throw new WebSocketException("Data on closed connection");
       }
       if (_state == FAILURE) {
         throw new WebSocketException("Data on failed connection");
       }
       while ((index < lastIndex) && _state != CLOSED && _state != FAILURE) {
         int byte = buffer[index];
         switch (_state) {
           case START:
             _fin = (byte & 0x80) != 0;
             _opcode = (byte & 0xF);
             switch (_opcode) {
             case _WebSocketOpcode.CONTINUATION:
               if (_currentMessageType == _WebSocketMessageType.NONE) {
                 throw new WebSocketException("Protocol error");
               }
               break;
 
             case _WebSocketOpcode.TEXT:
               if (_currentMessageType != _WebSocketMessageType.NONE) {
                 throw new WebSocketException("Protocol error");
               }
               _currentMessageType = _WebSocketMessageType.TEXT;
-              if (onMessageStart != null) {
-                onMessageStart(_WebSocketMessageType.TEXT);
-              }
+              _buffer = new StringBuffer();
               break;
 
             case _WebSocketOpcode.BINARY:
               if (_currentMessageType != _WebSocketMessageType.NONE) {
                 throw new WebSocketException("Protocol error");
               }
               _currentMessageType = _WebSocketMessageType.BINARY;
-              if (onMessageStart != null) {
-                onMessageStart(_WebSocketMessageType.BINARY);
-              }
+              _buffer = new _BufferList();
               break;
 
             case _WebSocketOpcode.CLOSE:
             case _WebSocketOpcode.PING:
             case _WebSocketOpcode.PONG:
               // Control frames cannot be fragmented.
               if (!_fin) throw new WebSocketException("Protocol error");
               break;
 
             default:
               throw new WebSocketException("Protocol error");
             }
             _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();
+              _lengthDone(sink);
             } else if (_len == 126) {
               _len = 0;
               _remainingLenBytes = 2;
               _state = LEN_REST;
             } else if (_len == 127) {
               _len = 0;
               _remainingLenBytes = 8;
               _state = LEN_REST;
             }
             break;
 
           case LEN_REST:
             _len = _len << 8 | byte;
             _remainingLenBytes--;
             if (_remainingLenBytes == 0) {
-              _lengthDone();
+              _lengthDone(sink);
             }
             break;
 
           case MASK:
             _maskingKey = _maskingKey << 8 | byte;
             _remainingMaskingKeyBytes--;
             if (_remainingMaskingKeyBytes == 0) {
-              _maskDone();
+              _maskDone(sink);
             }
             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;
@@ skipping 15 matching lines @@
                 // 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;
               }
 
               if (_remainingPayloadBytes == 0) {
-                _controlFrameEnd();
+                _controlFrameEnd(sink);
               }
             } else {
               switch (_currentMessageType) {
                 case _WebSocketMessageType.NONE:
                   throw new WebSocketException("Protocol error");
 
                 case _WebSocketMessageType.TEXT:
-                case _WebSocketMessageType.BINARY:
-                  if (onMessageData != null) {
-                    onMessageData(buffer, index, payload);
-                  }
+                  _buffer.add(_decodeString(buffer.getRange(index, payload)));
                   index += payload;
                   if (_remainingPayloadBytes == 0) {
-                    _messageFrameEnd();
+                    _messageFrameEnd(sink);
                   }
                   break;
 
+                case _WebSocketMessageType.BINARY:
+                  _buffer.add(buffer.getRange(index, payload));
+                  index += payload;
+                  if (_remainingPayloadBytes == 0) {
+                    _messageFrameEnd(sink);
+                  }
+                  break;
+
                 default:
                   throw new WebSocketException("Protocol error");
               }
             }
 
             // Hack - as we always do index++ below.
             index--;
             break;
         }
 
         // Move to the next byte.
         index++;
       }
     } catch (e) {
-      if (onClosed != null) onClosed(WebSocketStatus.PROTOCOL_ERROR,
-                                      "Protocol error");
       _state = FAILURE;
+      sink.signalError(e);
     }
   }
 
-  /**
-   * Indicate that the underlying communication channel has been closed.
-   */
-  void closed() {
-    if (_state == START || _state == CLOSED || _state == FAILURE) return;
-    if (onClosed != null) onClosed(WebSocketStatus.ABNORMAL_CLOSURE,
-                                    "Connection closed unexpectedly");
-    _state = CLOSED;
-  }
-
-  void _lengthDone() {
+  void _lengthDone(StreamSink sink) {
     if (_masked) {
       _state = MASK;
       _remainingMaskingKeyBytes = 4;
     } else {
       _remainingPayloadBytes = _len;
-      _startPayload();
+      _startPayload(sink);
     }
   }
 
-  void _maskDone() {
+  void _maskDone(StreamSink sink) {
     _remainingPayloadBytes = _len;
-    _startPayload();
+    _startPayload(sink);
   }
 
-  void _startPayload() {
+  void _startPayload(StreamSink sink) {
     // If there is no actual payload perform perform callbacks without
     // going through the PAYLOAD state.
     if (_remainingPayloadBytes == 0) {
       if (_isControlFrame()) {
         switch (_opcode) {
           case _WebSocketOpcode.CLOSE:
-            if (onClosed != null) onClosed(1005, "");
             _state = CLOSED;
+            sink.close();
             break;
           case _WebSocketOpcode.PING:
-            if (onPing != null) onPing(null);
+            // TODO(ajohnsen): Handle ping.
             break;
           case _WebSocketOpcode.PONG:
-            if (onPong != null) onPong(null);
+            // TODO(ajohnsen): Handle pong.
             break;
         }
         _prepareForNextFrame();
       } else {
-        _messageFrameEnd();
+        _messageFrameEnd(sink);
       }
     } else {
       _state = PAYLOAD;
     }
   }
 
-  void _messageFrameEnd() {
+  void _messageFrameEnd(StreamSink sink) {
     if (_fin) {
-      if (onMessageEnd != null) onMessageEnd();
+      switch (_currentMessageType) {
+        case _WebSocketMessageType.TEXT:
+          sink.add(_buffer.toString());
+          break;
+        case _WebSocketMessageType.BINARY:
+          if (_buffer.length == 0) {
+            sink.add(const []);
+          } else {
+            sink.add(_buffer.readBytes(_buffer.length));
+          }
+          break;
+      }
+      _buffer = null;
       _currentMessageType = _WebSocketMessageType.NONE;
     }
     _prepareForNextFrame();
   }
 
-  void _controlFrameEnd() {
+  void _controlFrameEnd(StreamSink sink) {
     switch (_opcode) {
       case _WebSocketOpcode.CLOSE:
-        int status = WebSocketStatus.NO_STATUS_RECEIVED;
-        String reason = "";
+        closeCode = WebSocketStatus.NO_STATUS_RECEIVED;
         if (_controlPayload.length > 0) {
           if (_controlPayload.length == 1) {
             throw new WebSocketException("Protocol error");
           }
-          status = _controlPayload[0] << 8 | _controlPayload[1];
-          if (status == WebSocketStatus.NO_STATUS_RECEIVED) {
+          closeCode = _controlPayload[0] << 8 | _controlPayload[1];
+          if (closeCode == WebSocketStatus.NO_STATUS_RECEIVED) {
             throw new WebSocketException("Protocol error");
           }
           if (_controlPayload.length > 2) {
-            reason = _decodeString(
+            closeReason = _decodeString(
                 _controlPayload.getRange(2, _controlPayload.length - 2));
           }
         }
-        if (onClosed != null) onClosed(status, reason);
         _state = CLOSED;
+        sink.close();
         break;
 
       case _WebSocketOpcode.PING:
-        if (onPing != null) onPing(_controlPayload);
+        // TODO(ajohnsen): Handle ping.
         break;
 
       case _WebSocketOpcode.PONG:
-        if (onPong != null) onPong(_controlPayload);
+        // TODO(ajohnsen): Handle pong.
         break;
     }
     _prepareForNextFrame();
   }
 
   bool _isControlFrame() {
     return _opcode == _WebSocketOpcode.CLOSE ||
            _opcode == _WebSocketOpcode.PING ||
            _opcode == _WebSocketOpcode.PONG;
   }
@@ skipping 18 matching lines @@
   int _len;
   bool _masked;
   int _maskingKey;
   int _remainingLenBytes;
   int _remainingMaskingKeyBytes;
   int _remainingPayloadBytes;
   int _unmaskingIndex;
 
   int _currentMessageType;
   List<int> _controlPayload;
+  var _buffer;  // Either StringBuffer or _BufferList.
 
-  Function onMessageStart;
-  Function onMessageData;
-  Function onMessageEnd;
-  Function onPing;
-  Function onPong;
-  Function onClosed;
+  int closeCode = WebSocketStatus.NO_STATUS_RECEIVED;
+  String closeReason = "";
 }
 
 
 class _WebSocketTransformerImpl implements WebSocketTransformer {
   final StreamController<WebSocket> _controller =
       new StreamController<WebSocket>();
 
   Stream<WebSocket> bind(Stream<HttpRequest> stream) {
     stream.listen((request) {
         _upgrade(request)
@@ skipping 56 matching lines @@
       return false;
     }
     return true;
   }
 }
 
 
 class _WebSocketImpl extends Stream implements WebSocket {
   final StreamController _controller = new StreamController();
 
-  final _WebSocketProtocolProcessor _processor =
-      new _WebSocketProtocolProcessor();
-
   final Socket _socket;
   int _readyState = WebSocket.CONNECTING;
   bool _writeClosed = false;
   int _closeCode;
   String _closeReason;
 
   static final HttpClient _httpClient = new HttpClient();
 
   static Future<WebSocket> connect(String url, [protocols]) {
     Uri uri = Uri.parse(url);
@@ skipping 61 matching lines @@
           }
         }
         return response.detachSocket()
             .then((socket) => new _WebSocketImpl._fromSocket(socket));
       });
   }
 
   _WebSocketImpl._fromSocket(Socket this._socket) {
     _readyState = WebSocket.OPEN;
 
-    int type;
-    var data;
-    _processor.onMessageStart = (int t) {
-      type = t;
-      if (type == _WebSocketMessageType.TEXT) {
-        data = new StringBuffer();
-      } else {
-        data = [];
-      }
-    };
-    _processor.onMessageData = (buffer, offset, count) {
-      if (type == _WebSocketMessageType.TEXT) {
-        data.add(_decodeString(buffer.getRange(offset, count)));
-      } else {
-        data.addAll(buffer.getRange(offset, count));
-      }
-    };
-    _processor.onMessageEnd = () {
-      if (type == _WebSocketMessageType.TEXT) {
-        _controller.add(data.toString());
-      } else {
-        _controller.add(data);
-      }
-    };
-    _processor.onClosed = (code, reason) {
-      bool clean = true;
-      if (_readyState == WebSocket.OPEN) {
-        _readyState = WebSocket.CLOSING;
-        if (code != WebSocketStatus.NO_STATUS_RECEIVED) {
-          _close(code);
-        } else {
-          _close();
-          clean = false;
-        }
-        _readyState = WebSocket.CLOSED;
-      }
-      if (_readyState == WebSocket.CLOSED) return;
-      _closeCode = code;
-      _closeReason = reason;
-      _controller.close();
-    };
-
-    _socket.listen(
-        (data) => _processor.update(data, 0, data.length),
-        onDone: () => _processor.closed(),
-        onError: (error) => _controller.signalError(error));
+    bool closed = false;
+    var transformer = new _WebSocketProtocolTransformer();
+    _socket.transform(transformer).listen(
+        (data) {
+          _controller.add(data);
+        },
+        onError: (error) {
+          if (closed) return;
+          closed = true;
+          _controller.signalError(error);
+          _controller.close();
+        },
+        onDone: () {
+          if (closed) return;
+          closed = true;
+          if (_readyState == WebSocket.OPEN) {
+            _readyState = WebSocket.CLOSING;
+            if (transformer.closeCode != WebSocketStatus.NO_STATUS_RECEIVED) {
+              _close(transformer.closeCode);
+            } else {
+              _close();
+            }
+            _readyState = WebSocket.CLOSED;
+          }
+          _closeCode = transformer.closeCode;
+          _closeReason = transformer.closeReason;
+          _controller.close();
+          if (_writeClosed) _socket.destroy();
+        },
+        unsubscribeOnError: true);
 
     _socket.done
         .catchError((error) {
-          if (_readyState == WebSocket.CLOSED) return;
+          if (closed) return;
+          closed = true;
           _readyState = WebSocket.CLOSED;
-          _closeCode = ABNORMAL_CLOSURE;
+          _closeCode = WebSocketStatus.ABNORMAL_CLOSURE;
           _controller.signalError(error);
           _controller.close();
-          _socket.destroy();
         })
         .whenComplete(() {
           _writeClosed = true;
         });
   }
 
   StreamSubscription listen(void onData(message),
                             {void onError(AsyncError error),
                              void onDone(),
                              bool unsubscribeOnError}) {
@@ skipping 59 matching lines @@
       } else {
         if (message is !List<int>) {
           throw new ArgumentError(message);
         }
         opcode = _WebSocketOpcode.BINARY;
         data = message;
       }
     } else {
       opcode = _WebSocketOpcode.TEXT;
     }
-    try {
-      _sendFrame(opcode, data);
-    } catch (_) {
-      // The socket can be closed before _socket.done have a chance
-      // to complete.
-    }
+    _sendFrame(opcode, data);
   }
 
   void _sendFrame(int opcode, [List<int> data]) {
     if (_writeClosed) return;
     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;
@@ skipping 12 matching lines @@
       lengthBytes = 8;
     } else if (dataLength > 125) {
       header[index++] = 126;
       lengthBytes = 2;
     }
     // Write the length in network byte order into the header.
     for (int i = 0; i < lengthBytes; i++) {
       header[index++] = dataLength >> (((lengthBytes - 1) - i) * 8) & 0xFF;
     }
     assert(index == headerSize);
-    _socket.add(header);
-    if (data != null) {
-      _socket.add(data);
+    try {
+      _socket.add(header);
+      if (data != null) {
+        _socket.add(data);
+      }
+    } catch (_) {
+      // The socket can be closed before _socket.done have a chance
+      // to complete.
+      _writeClosed = true;
     }
   }
 }