-
Notifications
You must be signed in to change notification settings - Fork 105
/
client.lua
371 lines (288 loc) · 8.84 KB
/
client.lua
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
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
-- Copyright (C) Yichun Zhang (agentzh)
-- FIXME: this library is very rough and is currently just for testing
-- the websocket server.
local wbproto = require "resty.websocket.protocol"
local bit = require "bit"
local _recv_frame = wbproto.recv_frame
local _send_frame = wbproto.send_frame
local new_tab = wbproto.new_tab
local tcp = ngx.socket.tcp
local re_match = ngx.re.match
local encode_base64 = ngx.encode_base64
local concat = table.concat
local char = string.char
local str_find = string.find
local rand = math.random
local rshift = bit.rshift
local band = bit.band
local setmetatable = setmetatable
local type = type
local debug = ngx.config.debug
local ngx_log = ngx.log
local ngx_DEBUG = ngx.DEBUG
local assert = assert
local ssl_support = true
if not ngx.config
or not ngx.config.ngx_lua_version
or ngx.config.ngx_lua_version < 9011
then
ssl_support = false
end
local _M = new_tab(0, 13)
_M._VERSION = '0.10'
local mt = { __index = _M }
function _M.new(self, opts)
local sock, err = tcp()
if not sock then
return nil, err
end
local max_payload_len, send_unmasked, timeout
if opts then
max_payload_len = opts.max_payload_len
send_unmasked = opts.send_unmasked
timeout = opts.timeout
if timeout then
sock:settimeout(timeout)
end
end
return setmetatable({
sock = sock,
max_payload_len = max_payload_len or 65535,
send_unmasked = send_unmasked,
}, mt)
end
function _M.connect(self, uri, opts)
local sock = self.sock
if not sock then
return nil, "not initialized"
end
local m, err = re_match(uri, [[^(wss?)://([^:/]+)(?::(\d+))?(.*)]], "jo")
if not m then
if err then
return nil, "failed to match the uri: " .. err
end
return nil, "bad websocket uri"
end
local scheme = m[1]
local host = m[2]
local port = m[3]
local path = m[4]
-- ngx.say("host: ", host)
-- ngx.say("port: ", port)
if not port then
port = scheme == 'wss' and 443 or 80
end
if path == "" then
path = "/"
end
local ssl_verify, server_name, headers, proto_header, origin_header
local sock_opts = false
local client_cert, client_priv_key
if opts then
local protos = opts.protocols
if protos then
if type(protos) == "table" then
proto_header = "\r\nSec-WebSocket-Protocol: "
.. concat(protos, ",")
else
proto_header = "\r\nSec-WebSocket-Protocol: " .. protos
end
end
local origin = opts.origin
if origin then
origin_header = "\r\nOrigin: " .. origin
end
local pool = opts.pool
if pool then
sock_opts = { pool = pool }
end
client_cert = opts.client_cert
client_priv_key = opts.client_priv_key
if client_cert then
assert(client_priv_key,
"client_priv_key must be provided with client_cert")
end
if opts.ssl_verify or opts.server_name then
if not ssl_support then
return nil, "ngx_lua 0.9.11+ required for SSL sockets"
end
ssl_verify = opts.ssl_verify
server_name = opts.server_name or host
end
if opts.headers then
headers = opts.headers
if type(headers) ~= "table" then
return nil, "custom headers must be a table"
end
end
end
local ok, err
if sock_opts then
ok, err = sock:connect(host, port, sock_opts)
else
ok, err = sock:connect(host, port)
end
if not ok then
return nil, "failed to connect: " .. err
end
if scheme == "wss" then
if not ssl_support then
return nil, "ngx_lua 0.9.11+ required for SSL sockets"
end
if client_cert then
ok, err = sock:setclientcert(client_cert, client_priv_key)
if not ok then
return nil, "failed to set TLS client certificate: " .. err
end
end
ok, err = sock:sslhandshake(false, server_name, ssl_verify)
if not ok then
return nil, "ssl handshake failed: " .. err
end
end
-- check for connections from pool:
local count, err = sock:getreusedtimes()
if not count then
return nil, "failed to get reused times: " .. err
end
if count > 0 then
-- being a reused connection (must have done handshake)
return 1
end
local custom_headers
if headers then
custom_headers = concat(headers, "\r\n")
custom_headers = "\r\n" .. custom_headers
end
-- do the websocket handshake:
local bytes = char(rand(256) - 1, rand(256) - 1, rand(256) - 1,
rand(256) - 1, rand(256) - 1, rand(256) - 1,
rand(256) - 1, rand(256) - 1, rand(256) - 1,
rand(256) - 1, rand(256) - 1, rand(256) - 1,
rand(256) - 1, rand(256) - 1, rand(256) - 1,
rand(256) - 1)
local key = encode_base64(bytes)
local req = "GET " .. path .. " HTTP/1.1\r\nUpgrade: websocket\r\nHost: "
.. host .. ":" .. port
.. "\r\nSec-WebSocket-Key: " .. key
.. (proto_header or "")
.. "\r\nSec-WebSocket-Version: 13"
.. (origin_header or "")
.. "\r\nConnection: Upgrade"
.. (custom_headers or "")
.. "\r\n\r\n"
local bytes, err = sock:send(req)
if not bytes then
return nil, "failed to send the handshake request: " .. err
end
local header_reader = sock:receiveuntil("\r\n\r\n")
-- FIXME: check for too big response headers
local header, err, partial = header_reader()
if not header then
return nil, "failed to receive response header: " .. err
end
-- error("header: " .. header)
-- FIXME: verify the response headers
m, err = re_match(header, [[^\s*HTTP/1\.1\s+]], "jo")
if not m then
return nil, "bad HTTP response status line: " .. header
end
return 1
end
function _M.set_timeout(self, time)
local sock = self.sock
if not sock then
return nil, nil, "not initialized yet"
end
return sock:settimeout(time)
end
function _M.recv_frame(self)
if self.fatal then
return nil, nil, "fatal error already happened"
end
local sock = self.sock
if not sock then
return nil, nil, "not initialized yet"
end
local data, typ, err = _recv_frame(sock, self.max_payload_len, false)
if not data and not str_find(err, ": timeout", 1, true) then
self.fatal = true
end
return data, typ, err
end
local function send_frame(self, fin, opcode, payload)
if self.fatal then
return nil, "fatal error already happened"
end
if self.closed then
return nil, "already closed"
end
local sock = self.sock
if not sock then
return nil, "not initialized yet"
end
local bytes, err = _send_frame(sock, fin, opcode, payload,
self.max_payload_len,
not self.send_unmasked)
if not bytes then
self.fatal = true
end
return bytes, err
end
_M.send_frame = send_frame
function _M.send_text(self, data)
return send_frame(self, true, 0x1, data)
end
function _M.send_binary(self, data)
return send_frame(self, true, 0x2, data)
end
local function send_close(self, code, msg)
local payload
if code then
if type(code) ~= "number" or code > 0x7fff then
return nil, "bad status code"
end
payload = char(band(rshift(code, 8), 0xff), band(code, 0xff))
.. (msg or "")
end
if debug then
ngx_log(ngx_DEBUG, "sending the close frame")
end
local bytes, err = send_frame(self, true, 0x8, payload)
if not bytes then
self.fatal = true
end
self.closed = true
return bytes, err
end
_M.send_close = send_close
function _M.send_ping(self, data)
return send_frame(self, true, 0x9, data)
end
function _M.send_pong(self, data)
return send_frame(self, true, 0xa, data)
end
function _M.close(self)
if self.fatal then
return nil, "fatal error already happened"
end
local sock = self.sock
if not sock then
return nil, "not initialized"
end
if not self.closed then
local bytes, err = send_close(self)
if not bytes then
return nil, "failed to send close frame: " .. err
end
end
return sock:close()
end
function _M.set_keepalive(self, ...)
local sock = self.sock
if not sock then
return nil, "not initialized"
end
return sock:setkeepalive(...)
end
return _M