net/websockets/websocket_unittest.cc

   1 // Copyright (c) 2009 The Chromium Authors. All rights reserved.
   2 // Use of this source code is governed by a BSD-style license that can be
   3 // found in the LICENSE file.
   4
   5 #include <string>
   6 #include <vector>
   7
   8 #include "base/callback.h"
   9 #include "net/base/completion_callback.h"
  10 #include "net/base/io_buffer.h"
  11 #include "net/base/mock_host_resolver.h"
  12 #include "net/base/test_completion_callback.h"
  13 #include "net/socket/socket_test_util.h"
  14 #include "net/url_request/url_request_unittest.h"
  15 #include "net/websockets/websocket.h"
  16 #include "testing/gtest/include/gtest/gtest.h"
  17 #include "testing/gmock/include/gmock/gmock.h"
  18 #include "testing/platform_test.h"
  19
  20 struct WebSocketEvent {
  21   enum EventType {
  22     EVENT_OPEN, EVENT_MESSAGE, EVENT_ERROR, EVENT_CLOSE,
  23   };
  24
  25   WebSocketEvent(EventType type, net::WebSocket* websocket,
  26                  const std::string& websocket_msg, bool websocket_flag)
  27       : event_type(type), socket(websocket), msg(websocket_msg),
  28         flag(websocket_flag) {}
  29
  30   EventType event_type;
  31   net::WebSocket* socket;
  32   std::string msg;
  33   bool flag;
  34 };
  35
  36 class WebSocketEventRecorder : public net::WebSocketDelegate {
  37  public:
  38   explicit WebSocketEventRecorder(net::CompletionCallback* callback)
  39       : onopen_(NULL),
  40         onmessage_(NULL),
  41         onerror_(NULL),
  42         onclose_(NULL),
  43         callback_(callback) {}
  44   virtual ~WebSocketEventRecorder() {
  45     delete onopen_;
  46     delete onmessage_;
  47     delete onerror_;
  48     delete onclose_;
  49   }
  50
  51   void SetOnOpen(Callback1<WebSocketEvent*>::Type* callback) {
  52     onopen_ = callback;
  53   }
  54   void SetOnMessage(Callback1<WebSocketEvent*>::Type* callback) {
  55     onmessage_ = callback;
  56   }
  57   void SetOnClose(Callback1<WebSocketEvent*>::Type* callback) {
  58     onclose_ = callback;
  59   }
  60
  61   virtual void OnOpen(net::WebSocket* socket) {
  62     events_.push_back(
  63         WebSocketEvent(WebSocketEvent::EVENT_OPEN, socket,
  64                        std::string(), false));
  65     if (onopen_)
  66       onopen_->Run(&events_.back());
  67   }
  68
  69   virtual void OnMessage(net::WebSocket* socket, const std::string& msg) {
  70     events_.push_back(
  71         WebSocketEvent(WebSocketEvent::EVENT_MESSAGE, socket, msg, false));
  72     if (onmessage_)
  73       onmessage_->Run(&events_.back());
  74   }
  75   virtual void OnError(net::WebSocket* socket) {
  76     events_.push_back(
  77         WebSocketEvent(WebSocketEvent::EVENT_ERROR, socket,
  78                        std::string(), false));
  79     if (onerror_)
  80       onerror_->Run(&events_.back());
  81   }
  82   virtual void OnClose(net::WebSocket* socket, bool was_clean) {
  83     events_.push_back(
  84         WebSocketEvent(WebSocketEvent::EVENT_CLOSE, socket,
  85                        std::string(), was_clean));
  86     if (onclose_)
  87       onclose_->Run(&events_.back());
  88     if (callback_)
  89       callback_->Run(net::OK);
  90   }
  91
  92   void DoClose(WebSocketEvent* event) {
  93     event->socket->Close();
  94   }
  95
  96   const std::vector<WebSocketEvent>& GetSeenEvents() const {
  97     return events_;
  98   }
  99
 100  private:
 101   std::vector<WebSocketEvent> events_;
 102   Callback1<WebSocketEvent*>::Type* onopen_;
 103   Callback1<WebSocketEvent*>::Type* onmessage_;
 104   Callback1<WebSocketEvent*>::Type* onerror_;
 105   Callback1<WebSocketEvent*>::Type* onclose_;
 106   net::CompletionCallback* callback_;
 107
 108   DISALLOW_COPY_AND_ASSIGN(WebSocketEventRecorder);
 109 };
 110
 111 namespace net {
 112
 113 class WebSocketTest : public PlatformTest {
 114  protected:
 115   void InitReadBuf(WebSocket* websocket) {
 116     // Set up |current_read_buf_|.
 117     websocket->current_read_buf_ = new GrowableIOBuffer();
 118   }
 119   void SetReadConsumed(WebSocket* websocket, int consumed) {
 120     websocket->read_consumed_len_ = consumed;
 121   }
 122   void AddToReadBuf(WebSocket* websocket, const char* data, int len) {
 123     websocket->AddToReadBuffer(data, len);
 124   }
 125
 126   void TestProcessFrameData(WebSocket* websocket,
 127                             const char* expected_remaining_data,
 128                             int expected_remaining_len) {
 129     websocket->ProcessFrameData();
 130
 131     const char* actual_remaining_data =
 132         websocket->current_read_buf_->StartOfBuffer()
 133         + websocket->read_consumed_len_;
 134     int actual_remaining_len =
 135         websocket->current_read_buf_->offset() - websocket->read_consumed_len_;
 136
 137     EXPECT_EQ(expected_remaining_len, actual_remaining_len);
 138     EXPECT_TRUE(!memcmp(expected_remaining_data, actual_remaining_data,
 139                         expected_remaining_len));
 140   }
 141 };
 142
 143 TEST_F(WebSocketTest, Connect) {
 144   MockClientSocketFactory mock_socket_factory;
 145   MockRead data_reads[] = {
 146     MockRead("HTTP/1.1 101 Web Socket Protocol Handshake\r\n"
 147              "Upgrade: WebSocket\r\n"
 148              "Connection: Upgrade\r\n"
 149              "WebSocket-Origin: http://example.com\r\n"
 150              "WebSocket-Location: ws://example.com/demo\r\n"
 151              "WebSocket-Protocol: sample\r\n"
 152              "\r\n"),
 153     // Server doesn't close the connection after handshake.
 154     MockRead(true, ERR_IO_PENDING),
 155   };
 156   MockWrite data_writes[] = {
 157     MockWrite("GET /demo HTTP/1.1\r\n"
 158               "Upgrade: WebSocket\r\n"
 159               "Connection: Upgrade\r\n"
 160               "Host: example.com\r\n"
 161               "Origin: http://example.com\r\n"
 162               "WebSocket-Protocol: sample\r\n"
 163               "\r\n"),
 164   };
 165   StaticSocketDataProvider data(data_reads, arraysize(data_reads),
 166                                 data_writes, arraysize(data_writes));
 167   mock_socket_factory.AddSocketDataProvider(&data);
 168
 169   WebSocket::Request* request(
 170       new WebSocket::Request(GURL("ws://example.com/demo"),
 171                              "sample",
 172                              "http://example.com",
 173                              "ws://example.com/demo",
 174                              WebSocket::DRAFT75,
 175                              new TestURLRequestContext()));
 176   request->SetHostResolver(new MockHostResolver());
 177   request->SetClientSocketFactory(&mock_socket_factory);
 178
 179   TestCompletionCallback callback;
 180
 181   scoped_ptr<WebSocketEventRecorder> delegate(
 182       new WebSocketEventRecorder(&callback));
 183   delegate->SetOnOpen(NewCallback(delegate.get(),
 184                                   &WebSocketEventRecorder::DoClose));
 185
 186   scoped_refptr<WebSocket> websocket(
 187       new WebSocket(request, delegate.get()));
 188
 189   EXPECT_EQ(WebSocket::INITIALIZED, websocket->ready_state());
 190   websocket->Connect();
 191
 192   callback.WaitForResult();
 193
 194   const std::vector<WebSocketEvent>& events = delegate->GetSeenEvents();
 195   EXPECT_EQ(2U, events.size());
 196
 197   EXPECT_EQ(WebSocketEvent::EVENT_OPEN, events[0].event_type);
 198   EXPECT_EQ(WebSocketEvent::EVENT_CLOSE, events[1].event_type);
 199 }
 200
 201 TEST_F(WebSocketTest, ServerSentData) {
 202   MockClientSocketFactory mock_socket_factory;
 203   static const char kMessage[] = "Hello";
 204   static const char kFrame[] = "\x00Hello\xff";
 205   static const int kFrameLen = sizeof(kFrame) - 1;
 206   MockRead data_reads[] = {
 207     MockRead("HTTP/1.1 101 Web Socket Protocol Handshake\r\n"
 208              "Upgrade: WebSocket\r\n"
 209              "Connection: Upgrade\r\n"
 210              "WebSocket-Origin: http://example.com\r\n"
 211              "WebSocket-Location: ws://example.com/demo\r\n"
 212              "WebSocket-Protocol: sample\r\n"
 213              "\r\n"),
 214     MockRead(true, kFrame, kFrameLen),
 215     // Server doesn't close the connection after handshake.
 216     MockRead(true, ERR_IO_PENDING),
 217   };
 218   MockWrite data_writes[] = {
 219     MockWrite("GET /demo HTTP/1.1\r\n"
 220               "Upgrade: WebSocket\r\n"
 221               "Connection: Upgrade\r\n"
 222               "Host: example.com\r\n"
 223               "Origin: http://example.com\r\n"
 224               "WebSocket-Protocol: sample\r\n"
 225               "\r\n"),
 226   };
 227   StaticSocketDataProvider data(data_reads, arraysize(data_reads),
 228                                 data_writes, arraysize(data_writes));
 229   mock_socket_factory.AddSocketDataProvider(&data);
 230
 231   WebSocket::Request* request(
 232       new WebSocket::Request(GURL("ws://example.com/demo"),
 233                              "sample",
 234                              "http://example.com",
 235                              "ws://example.com/demo",
 236                              WebSocket::DRAFT75,
 237                              new TestURLRequestContext()));
 238   request->SetHostResolver(new MockHostResolver());
 239   request->SetClientSocketFactory(&mock_socket_factory);
 240
 241   TestCompletionCallback callback;
 242
 243   scoped_ptr<WebSocketEventRecorder> delegate(
 244       new WebSocketEventRecorder(&callback));
 245   delegate->SetOnMessage(NewCallback(delegate.get(),
 246                                      &WebSocketEventRecorder::DoClose));
 247
 248   scoped_refptr<WebSocket> websocket(
 249       new WebSocket(request, delegate.get()));
 250
 251   EXPECT_EQ(WebSocket::INITIALIZED, websocket->ready_state());
 252   websocket->Connect();
 253
 254   callback.WaitForResult();
 255
 256   const std::vector<WebSocketEvent>& events = delegate->GetSeenEvents();
 257   EXPECT_EQ(3U, events.size());
 258
 259   EXPECT_EQ(WebSocketEvent::EVENT_OPEN, events[0].event_type);
 260   EXPECT_EQ(WebSocketEvent::EVENT_MESSAGE, events[1].event_type);
 261   EXPECT_EQ(kMessage, events[1].msg);
 262   EXPECT_EQ(WebSocketEvent::EVENT_CLOSE, events[2].event_type);
 263 }
 264
 265 TEST_F(WebSocketTest, ProcessFrameDataForLengthCalculation) {
 266   WebSocket::Request* request(
 267       new WebSocket::Request(GURL("ws://example.com/demo"),
 268                              "sample",
 269                              "http://example.com",
 270                              "ws://example.com/demo",
 271                              WebSocket::DRAFT75,
 272                              new TestURLRequestContext()));
 273   TestCompletionCallback callback;
 274   scoped_ptr<WebSocketEventRecorder> delegate(
 275       new WebSocketEventRecorder(&callback));
 276
 277   scoped_refptr<WebSocket> websocket(
 278       new WebSocket(request, delegate.get()));
 279
 280   // Frame data: skip length 1 ('x'), and try to skip length 129
 281   // (1 * 128 + 1) bytes after \x81\x01, but buffer is too short to skip.
 282   static const char kTestLengthFrame[] =
 283       "\x80\x01x\x80\x81\x01\x01\x00unexpected data\xFF";
 284   const int kTestLengthFrameLength = sizeof(kTestLengthFrame) - 1;
 285   InitReadBuf(websocket.get());
 286   AddToReadBuf(websocket.get(), kTestLengthFrame, kTestLengthFrameLength);
 287   SetReadConsumed(websocket.get(), 0);
 288
 289   static const char kExpectedRemainingFrame[] =
 290       "\x80\x81\x01\x01\x00unexpected data\xFF";
 291   const int kExpectedRemainingLength = sizeof(kExpectedRemainingFrame) - 1;
 292   TestProcessFrameData(websocket.get(),
 293                        kExpectedRemainingFrame, kExpectedRemainingLength);
 294   // No onmessage event expected.
 295   const std::vector<WebSocketEvent>& events = delegate->GetSeenEvents();
 296   EXPECT_EQ(1U, events.size());
 297
 298   EXPECT_EQ(WebSocketEvent::EVENT_ERROR, events[0].event_type);
 299
 300   websocket->DetachDelegate();
 301 }
 302
 303 TEST_F(WebSocketTest, ProcessFrameDataForUnterminatedString) {
 304   WebSocket::Request* request(
 305       new WebSocket::Request(GURL("ws://example.com/demo"),
 306                              "sample",
 307                              "http://example.com",
 308                              "ws://example.com/demo",
 309                              WebSocket::DRAFT75,
 310                              new TestURLRequestContext()));
 311   TestCompletionCallback callback;
 312   scoped_ptr<WebSocketEventRecorder> delegate(
 313       new WebSocketEventRecorder(&callback));
 314
 315   scoped_refptr<WebSocket> websocket(
 316       new WebSocket(request, delegate.get()));
 317
 318   static const char kTestUnterminatedFrame[] =
 319       "\x00unterminated frame";
 320   const int kTestUnterminatedFrameLength = sizeof(kTestUnterminatedFrame) - 1;
 321   InitReadBuf(websocket.get());
 322   AddToReadBuf(websocket.get(), kTestUnterminatedFrame,
 323                kTestUnterminatedFrameLength);
 324   SetReadConsumed(websocket.get(), 0);
 325   TestProcessFrameData(websocket.get(),
 326                        kTestUnterminatedFrame, kTestUnterminatedFrameLength);
 327   {
 328     // No onmessage event expected.
 329     const std::vector<WebSocketEvent>& events = delegate->GetSeenEvents();
 330     EXPECT_EQ(0U, events.size());
 331   }
 332
 333   static const char kTestTerminateFrame[] = " is terminated in next read\xff";
 334   const int kTestTerminateFrameLength = sizeof(kTestTerminateFrame) - 1;
 335   AddToReadBuf(websocket.get(), kTestTerminateFrame,
 336                kTestTerminateFrameLength);
 337   TestProcessFrameData(websocket.get(), "", 0);
 338
 339   static const char kExpectedMsg[] =
 340       "unterminated frame is terminated in next read";
 341   {
 342     const std::vector<WebSocketEvent>& events = delegate->GetSeenEvents();
 343     EXPECT_EQ(1U, events.size());
 344
 345     EXPECT_EQ(WebSocketEvent::EVENT_MESSAGE, events[0].event_type);
 346     EXPECT_EQ(kExpectedMsg, events[0].msg);
 347   }
 348
 349   websocket->DetachDelegate();
 350 }
 351
 352 }  // namespace net