-
Notifications
You must be signed in to change notification settings - Fork 47
/
websocket_handler_test.exs
167 lines (127 loc) · 4.5 KB
/
websocket_handler_test.exs
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
defmodule WebSocketHandlerTest do
use ExUnit.Case, async: true
defmodule WebSocketHandler do
@behaviour :cowboy_websocket
# We never actually call this; it's just here to quell compiler warnings
@impl true
def init(req, state), do: {:cowboy_websocket, req, state}
@impl true
def websocket_init(_opts), do: {:ok, :init}
@impl true
def websocket_handle({:text, "state"}, state), do: {[{:text, inspect(state)}], state}
def websocket_handle({:text, "whoami"}, state),
do: {[{:text, :erlang.pid_to_list(self())}], state}
@impl true
def websocket_info(msg, state), do: {[{:text, inspect(msg)}], state}
end
@protocol_options [
idle_timeout: 1000,
request_timeout: 1000
]
setup_all do
{:ok, _} = Plug.Cowboy.http(__MODULE__, [], port: 8083, protocol_options: @protocol_options)
on_exit(fn -> :ok = Plug.Cowboy.shutdown(__MODULE__.HTTP) end)
{:ok, port: 8083}
end
@behaviour Plug
@impl Plug
def init(arg), do: arg
@impl Plug
def call(conn, _opts) do
conn = Plug.Conn.fetch_query_params(conn)
handler = conn.query_params["handler"] |> String.to_atom()
Plug.Conn.upgrade_adapter(conn, :websocket, {handler, [], [timeout: 1000]})
end
test "websocket_init and websocket_handle are called", context do
client = tcp_client(context)
http1_handshake(client, WebSocketHandler)
send_text_frame(client, "state")
{:ok, result} = recv_text_frame(client)
assert result == inspect(:init)
end
test "websocket_info is called", context do
client = tcp_client(context)
http1_handshake(client, WebSocketHandler)
send_text_frame(client, "whoami")
{:ok, pid} = recv_text_frame(client)
pid = pid |> String.to_charlist() |> :erlang.list_to_pid()
Process.send(pid, "hello info", [])
{:ok, response} = recv_text_frame(client)
assert response == inspect("hello info")
end
# Simple WebSocket client
def tcp_client(context) do
{:ok, socket} = :gen_tcp.connect('localhost', context[:port], active: false, mode: :binary)
socket
end
def http1_handshake(client, module, params \\ []) do
params = params |> Keyword.put(:handler, module)
:gen_tcp.send(client, """
GET /?#{URI.encode_query(params)} HTTP/1.1\r
Host: server.example.com\r
Upgrade: websocket\r
Connection: Upgrade\r
Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==\r
Sec-WebSocket-Version: 13\r
\r
""")
{:ok, response} = :gen_tcp.recv(client, 234)
[
"HTTP/1.1 101 Switching Protocols",
"cache-control: max-age=0, private, must-revalidate",
"connection: Upgrade",
"date: " <> _date,
"sec-websocket-accept: s3pPLMBiTxaQ9kYGzzhZRbK+xOo=",
"server: Cowboy",
"upgrade: websocket",
"",
""
] = String.split(response, "\r\n")
end
defp recv_text_frame(client) do
{:ok, 0x8, 0x1, body} = recv_frame(client)
{:ok, body}
end
defp recv_frame(client) do
{:ok, header} = :gen_tcp.recv(client, 2)
<<flags::4, opcode::4, 0::1, length::7>> = header
{:ok, data} =
case length do
0 ->
{:ok, <<>>}
126 ->
{:ok, <<length::16>>} = :gen_tcp.recv(client, 2)
:gen_tcp.recv(client, length)
127 ->
{:ok, <<length::64>>} = :gen_tcp.recv(client, 8)
:gen_tcp.recv(client, length)
length ->
:gen_tcp.recv(client, length)
end
{:ok, flags, opcode, data}
end
defp send_text_frame(client, data, flags \\ 0x8) do
send_frame(client, flags, 0x1, data)
end
defp send_frame(client, flags, opcode, data) do
mask = :rand.uniform(1_000_000)
masked_data = mask(data, mask)
mask_flag_and_size =
case byte_size(masked_data) do
size when size <= 125 -> <<1::1, size::7>>
size when size <= 65_535 -> <<1::1, 126::7, size::16>>
size -> <<1::1, 127::7, size::64>>
end
:gen_tcp.send(client, [<<flags::4, opcode::4>>, mask_flag_and_size, <<mask::32>>, masked_data])
end
# Note that masking is an involution, so we don't need a separate unmask function
defp mask(payload, mask, acc \\ <<>>)
defp mask(payload, mask, acc) when is_integer(mask), do: mask(payload, <<mask::32>>, acc)
defp mask(<<h::32, rest::binary>>, <<mask::32>>, acc) do
mask(rest, mask, acc <> <<Bitwise.bxor(h, mask)::32>>)
end
defp mask(<<h::8, rest::binary>>, <<current::8, mask::24>>, acc) do
mask(rest, <<mask::24, current::8>>, acc <> <<Bitwise.bxor(h, current)::8>>)
end
defp mask(<<>>, _mask, acc), do: acc
end