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
7#include "base/message_loop.h"
8#include "googleurl/src/gurl.h"
9#include "net/base/address_list.h"
10#include "net/base/sys_addrinfo.h"
11#include "net/base/test_completion_callback.h"
12#include "net/socket_stream/socket_stream.h"
13#include "net/websockets/websocket_throttle.h"
14#include "testing/gtest/include/gtest/gtest.h"
15#include "testing/platform_test.h"
16
17class DummySocketStreamDelegate : public net::SocketStream::Delegate {
18 public:
19  DummySocketStreamDelegate() {}
20  virtual ~DummySocketStreamDelegate() {}
21  virtual void OnConnected(
22      net::SocketStream* socket, int max_pending_send_allowed) {}
23  virtual void OnSentData(net::SocketStream* socket, int amount_sent) {}
24  virtual void OnReceivedData(net::SocketStream* socket,
25                              const char* data, int len) {}
26  virtual void OnClose(net::SocketStream* socket) {}
27};
28
29namespace net {
30
31class WebSocketThrottleTest : public PlatformTest {
32 protected:
33  struct addrinfo *AddAddr(int a1, int a2, int a3, int a4,
34                           struct addrinfo* next) {
35    struct addrinfo* addrinfo = new struct addrinfo;
36    memset(addrinfo, 0, sizeof(struct addrinfo));
37    addrinfo->ai_family = AF_INET;
38    int addrlen = sizeof(struct sockaddr_in);
39    addrinfo->ai_addrlen = addrlen;
40    addrinfo->ai_addr = reinterpret_cast<sockaddr*>(new char[addrlen]);
41    memset(addrinfo->ai_addr, 0, sizeof(addrlen));
42    struct sockaddr_in* addr =
43        reinterpret_cast<sockaddr_in*>(addrinfo->ai_addr);
44    int addrint = ((a1 & 0xff) << 24) |
45        ((a2 & 0xff) << 16) |
46        ((a3 & 0xff) <<  8) |
47        ((a4 & 0xff));
48    memcpy(&addr->sin_addr, &addrint, sizeof(int));
49    addrinfo->ai_next = next;
50    return addrinfo;
51  }
52  void DeleteAddrInfo(struct addrinfo* head) {
53    if (!head)
54      return;
55    struct addrinfo* next;
56    for (struct addrinfo* a = head; a != NULL; a = next) {
57      next = a->ai_next;
58      delete [] a->ai_addr;
59      delete a;
60    }
61  }
62
63  static void SetAddressList(SocketStream* socket, struct addrinfo* head) {
64    socket->CopyAddrInfo(head);
65  }
66};
67
68TEST_F(WebSocketThrottleTest, Throttle) {
69  WebSocketThrottle::Init();
70  DummySocketStreamDelegate delegate;
71
72  WebSocketThrottle* throttle = Singleton<WebSocketThrottle>::get();
73
74  EXPECT_EQ(throttle,
75            SocketStreamThrottle::GetSocketStreamThrottleForScheme("ws"));
76  EXPECT_EQ(throttle,
77            SocketStreamThrottle::GetSocketStreamThrottleForScheme("wss"));
78
79  // For host1: 1.2.3.4, 1.2.3.5, 1.2.3.6
80  struct addrinfo* addr = AddAddr(1, 2, 3, 4, NULL);
81  addr = AddAddr(1, 2, 3, 5, addr);
82  addr = AddAddr(1, 2, 3, 6, addr);
83  scoped_refptr<SocketStream> s1 =
84      new SocketStream(GURL("ws://host1/"), &delegate);
85  WebSocketThrottleTest::SetAddressList(s1, addr);
86  DeleteAddrInfo(addr);
87
88  TestCompletionCallback callback_s1;
89  EXPECT_EQ(OK, throttle->OnStartOpenConnection(s1, &callback_s1));
90
91  // For host2: 1.2.3.4
92  addr = AddAddr(1, 2, 3, 4, NULL);
93  scoped_refptr<SocketStream> s2 =
94      new SocketStream(GURL("ws://host2/"), &delegate);
95  WebSocketThrottleTest::SetAddressList(s2, addr);
96  DeleteAddrInfo(addr);
97
98  TestCompletionCallback callback_s2;
99  EXPECT_EQ(ERR_IO_PENDING, throttle->OnStartOpenConnection(s2, &callback_s2));
100
101  // For host3: 1.2.3.5
102  addr = AddAddr(1, 2, 3, 5, NULL);
103  scoped_refptr<SocketStream> s3 =
104      new SocketStream(GURL("ws://host3/"), &delegate);
105  WebSocketThrottleTest::SetAddressList(s3, addr);
106  DeleteAddrInfo(addr);
107
108  TestCompletionCallback callback_s3;
109  EXPECT_EQ(ERR_IO_PENDING, throttle->OnStartOpenConnection(s3, &callback_s3));
110
111  // For host4: 1.2.3.4, 1.2.3.6
112  addr = AddAddr(1, 2, 3, 4, NULL);
113  addr = AddAddr(1, 2, 3, 6, addr);
114  scoped_refptr<SocketStream> s4 =
115      new SocketStream(GURL("ws://host4/"), &delegate);
116  WebSocketThrottleTest::SetAddressList(s4, addr);
117  DeleteAddrInfo(addr);
118
119  TestCompletionCallback callback_s4;
120  EXPECT_EQ(ERR_IO_PENDING, throttle->OnStartOpenConnection(s4, &callback_s4));
121
122  static const char kHeader[] = "HTTP/1.1 101 Web Socket Protocol\r\n";
123  EXPECT_EQ(OK,
124            throttle->OnRead(s1.get(), kHeader, sizeof(kHeader) - 1, NULL));
125  EXPECT_FALSE(callback_s2.have_result());
126  EXPECT_FALSE(callback_s3.have_result());
127  EXPECT_FALSE(callback_s4.have_result());
128
129  static const char kHeader2[] =
130      "Upgrade: WebSocket\r\n"
131      "Connection: Upgrade\r\n"
132      "WebSocket-Origin: http://www.google.com\r\n"
133      "WebSocket-Location: ws://websocket.chromium.org\r\n"
134      "\r\n";
135  EXPECT_EQ(OK,
136            throttle->OnRead(s1.get(), kHeader2, sizeof(kHeader2) - 1, NULL));
137  MessageLoopForIO::current()->RunAllPending();
138  EXPECT_TRUE(callback_s2.have_result());
139  EXPECT_TRUE(callback_s3.have_result());
140  EXPECT_FALSE(callback_s4.have_result());
141
142  throttle->OnClose(s1.get());
143  MessageLoopForIO::current()->RunAllPending();
144  EXPECT_FALSE(callback_s4.have_result());
145  s1->DetachDelegate();
146
147  throttle->OnClose(s2.get());
148  MessageLoopForIO::current()->RunAllPending();
149  EXPECT_TRUE(callback_s4.have_result());
150  s2->DetachDelegate();
151
152  throttle->OnClose(s3.get());
153  MessageLoopForIO::current()->RunAllPending();
154  s3->DetachDelegate();
155  throttle->OnClose(s4.get());
156  s4->DetachDelegate();
157}
158
159}
160