1 | /* |
2 | * Wslay - The WebSocket Library |
3 | * |
4 | * Copyright (c) 2011, 2012 Tatsuhiro Tsujikawa |
5 | * |
6 | * Permission is hereby granted, free of charge, to any person obtaining |
7 | * a copy of this software and associated documentation files (the |
8 | * "Software"), to deal in the Software without restriction, including |
9 | * without limitation the rights to use, copy, modify, merge, publish, |
10 | * distribute, sublicense, and/or sell copies of the Software, and to |
11 | * permit persons to whom the Software is furnished to do so, subject to |
12 | * the following conditions: |
13 | * |
14 | * The above copyright notice and this permission notice shall be |
15 | * included in all copies or substantial portions of the Software. |
16 | * |
17 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, |
18 | * EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF |
19 | * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND |
20 | * NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE |
21 | * LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION |
22 | * OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION |
23 | * WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. |
24 | */ |
25 | #include "wslay_frame.h" |
26 | |
27 | #include <stddef.h> |
28 | #include <string.h> |
29 | #include <assert.h> |
30 | |
31 | #include "wslay_net.h" |
32 | |
33 | #define wslay_min(A, B) (((A) < (B)) ? (A) : (B)) |
34 | |
35 | int wslay_frame_context_init(wslay_frame_context_ptr *ctx, |
36 | const struct wslay_frame_callbacks *callbacks, |
37 | void *user_data) { |
38 | *ctx = malloc(sizeof(struct wslay_frame_context)); |
39 | if (*ctx == NULL) { |
40 | return -1; |
41 | } |
42 | memset(*ctx, 0, sizeof(struct wslay_frame_context)); |
43 | (*ctx)->istate = RECV_HEADER1; |
44 | (*ctx)->ireqread = 2; |
45 | (*ctx)->ostate = PREP_HEADER; |
46 | (*ctx)->user_data = user_data; |
47 | (*ctx)->ibufmark = (*ctx)->ibuflimit = (*ctx)->ibuf; |
48 | (*ctx)->callbacks = *callbacks; |
49 | return 0; |
50 | } |
51 | |
52 | void wslay_frame_context_free(wslay_frame_context_ptr ctx) { free(ctx); } |
53 | |
54 | ssize_t wslay_frame_send(wslay_frame_context_ptr ctx, |
55 | struct wslay_frame_iocb *iocb) { |
56 | if (iocb->data_length > iocb->payload_length) { |
57 | return WSLAY_ERR_INVALID_ARGUMENT; |
58 | } |
59 | if (ctx->ostate == PREP_HEADER) { |
60 | uint8_t *hdptr = ctx->oheader; |
61 | memset(ctx->oheader, 0, sizeof(ctx->oheader)); |
62 | *hdptr |= (uint8_t)((uint8_t)(iocb->fin << 7) & 0x80u); |
63 | *hdptr |= (uint8_t)((uint8_t)(iocb->rsv << 4) & 0x70u); |
64 | /* Suppress stubborn gcc-10 warning */ |
65 | *hdptr |= (uint8_t)((uint8_t)(iocb->opcode << 0) & 0xfu); |
66 | ++hdptr; |
67 | *hdptr |= (uint8_t)((uint8_t)(iocb->mask << 7) & 0x80u); |
68 | if (wslay_is_ctrl_frame(iocb->opcode) && iocb->payload_length > 125) { |
69 | return WSLAY_ERR_INVALID_ARGUMENT; |
70 | } |
71 | if (iocb->payload_length < 126) { |
72 | *hdptr |= (uint8_t)iocb->payload_length; |
73 | ++hdptr; |
74 | } else if (iocb->payload_length < (1 << 16)) { |
75 | uint16_t len = htons((uint16_t)iocb->payload_length); |
76 | *hdptr |= 126; |
77 | ++hdptr; |
78 | memcpy(hdptr, &len, 2); |
79 | hdptr += 2; |
80 | } else if (iocb->payload_length < (1ull << 63)) { |
81 | uint64_t len = hton64(iocb->payload_length); |
82 | *hdptr |= 127; |
83 | ++hdptr; |
84 | memcpy(hdptr, &len, 8); |
85 | hdptr += 8; |
86 | } else { |
87 | /* Too large payload length */ |
88 | return WSLAY_ERR_INVALID_ARGUMENT; |
89 | } |
90 | if (iocb->mask) { |
91 | if (ctx->callbacks.genmask_callback(ctx->omaskkey, 4, ctx->user_data) != |
92 | 0) { |
93 | return WSLAY_ERR_INVALID_CALLBACK; |
94 | } else { |
95 | ctx->omask = 1; |
96 | memcpy(hdptr, ctx->omaskkey, 4); |
97 | hdptr += 4; |
98 | } |
99 | } |
100 | ctx->ostate = SEND_HEADER; |
101 | ctx->oheadermark = ctx->oheader; |
102 | ctx->oheaderlimit = hdptr; |
103 | ctx->opayloadlen = iocb->payload_length; |
104 | ctx->opayloadoff = 0; |
105 | } |
106 | if (ctx->ostate == SEND_HEADER) { |
107 | ptrdiff_t len = ctx->oheaderlimit - ctx->oheadermark; |
108 | ssize_t r; |
109 | int flags = 0; |
110 | if (iocb->data_length > 0) { |
111 | flags |= WSLAY_MSG_MORE; |
112 | } |
113 | r = ctx->callbacks.send_callback(ctx->oheadermark, (size_t)len, flags, |
114 | ctx->user_data); |
115 | if (r > 0) { |
116 | if (r > len) { |
117 | return WSLAY_ERR_INVALID_CALLBACK; |
118 | } else { |
119 | ctx->oheadermark += r; |
120 | if (ctx->oheadermark == ctx->oheaderlimit) { |
121 | ctx->ostate = SEND_PAYLOAD; |
122 | } else { |
123 | return WSLAY_ERR_WANT_WRITE; |
124 | } |
125 | } |
126 | } else { |
127 | return WSLAY_ERR_WANT_WRITE; |
128 | } |
129 | } |
130 | if (ctx->ostate == SEND_PAYLOAD) { |
131 | size_t totallen = 0; |
132 | if (iocb->data_length > 0) { |
133 | if (ctx->omask) { |
134 | uint8_t temp[4096]; |
135 | const uint8_t *datamark = iocb->data, |
136 | *datalimit = iocb->data + iocb->data_length; |
137 | while (datamark < datalimit) { |
138 | size_t datalen = (size_t)(datalimit - datamark); |
139 | const uint8_t *writelimit = |
140 | datamark + wslay_min(sizeof(temp), datalen); |
141 | size_t writelen = (size_t)(writelimit - datamark); |
142 | ssize_t r; |
143 | size_t i; |
144 | for (i = 0; i < writelen; ++i) { |
145 | temp[i] = datamark[i] ^ ctx->omaskkey[(ctx->opayloadoff + i) % 4]; |
146 | } |
147 | r = ctx->callbacks.send_callback(temp, writelen, 0, ctx->user_data); |
148 | if (r > 0) { |
149 | if ((size_t)r > writelen) { |
150 | return WSLAY_ERR_INVALID_CALLBACK; |
151 | } else { |
152 | datamark += r; |
153 | ctx->opayloadoff += (uint64_t)r; |
154 | totallen += (size_t)r; |
155 | } |
156 | } else { |
157 | if (totallen > 0) { |
158 | break; |
159 | } else { |
160 | return WSLAY_ERR_WANT_WRITE; |
161 | } |
162 | } |
163 | } |
164 | } else { |
165 | ssize_t r; |
166 | r = ctx->callbacks.send_callback(iocb->data, iocb->data_length, 0, |
167 | ctx->user_data); |
168 | if (r > 0) { |
169 | if ((size_t)r > iocb->data_length) { |
170 | return WSLAY_ERR_INVALID_CALLBACK; |
171 | } else { |
172 | ctx->opayloadoff += (uint64_t)r; |
173 | totallen = (size_t)r; |
174 | } |
175 | } else { |
176 | return WSLAY_ERR_WANT_WRITE; |
177 | } |
178 | } |
179 | } |
180 | if (ctx->opayloadoff == ctx->opayloadlen) { |
181 | ctx->ostate = PREP_HEADER; |
182 | } |
183 | return (ssize_t)totallen; |
184 | } |
185 | return WSLAY_ERR_INVALID_ARGUMENT; |
186 | } |
187 | |
188 | ssize_t wslay_frame_write(wslay_frame_context_ptr ctx, |
189 | struct wslay_frame_iocb *iocb, uint8_t *buf, |
190 | size_t buflen, size_t *pwpayloadlen) { |
191 | uint8_t *buf_last = buf; |
192 | size_t i; |
193 | size_t hdlen; |
194 | |
195 | *pwpayloadlen = 0; |
196 | |
197 | if (iocb->data_length > iocb->payload_length) { |
198 | return WSLAY_ERR_INVALID_ARGUMENT; |
199 | } |
200 | |
201 | switch (ctx->ostate) { |
202 | case PREP_HEADER: |
203 | case PREP_HEADER_NOBUF: |
204 | hdlen = 2; |
205 | if (iocb->payload_length < 126) { |
206 | /* nothing to do */ |
207 | } else if (iocb->payload_length < (1 << 16)) { |
208 | hdlen += 2; |
209 | } else if (iocb->payload_length < (1ull << 63)) { |
210 | hdlen += 8; |
211 | } |
212 | if (iocb->mask) { |
213 | hdlen += 4; |
214 | } |
215 | |
216 | if (buflen < hdlen) { |
217 | ctx->ostate = PREP_HEADER_NOBUF; |
218 | return 0; |
219 | } |
220 | |
221 | memset(buf_last, 0, hdlen); |
222 | *buf_last |= (uint8_t)((uint8_t)(iocb->fin << 7) & 0x80u); |
223 | *buf_last |= (uint8_t)((uint8_t)(iocb->rsv << 4) & 0x70u); |
224 | /* Suppress stubborn gcc-10 warning */ |
225 | *buf_last |= (uint8_t)((uint8_t)(iocb->opcode << 0) & 0xfu); |
226 | ++buf_last; |
227 | *buf_last |= (uint8_t)((uint8_t)(iocb->mask << 7) & 0x80u); |
228 | if (wslay_is_ctrl_frame(iocb->opcode) && iocb->payload_length > 125) { |
229 | return WSLAY_ERR_INVALID_ARGUMENT; |
230 | } |
231 | if (iocb->payload_length < 126) { |
232 | *buf_last |= (uint8_t)iocb->payload_length; |
233 | ++buf_last; |
234 | } else if (iocb->payload_length < (1 << 16)) { |
235 | uint16_t len = htons((uint16_t)iocb->payload_length); |
236 | *buf_last |= 126; |
237 | ++buf_last; |
238 | memcpy(buf_last, &len, 2); |
239 | buf_last += 2; |
240 | } else if (iocb->payload_length < (1ull << 63)) { |
241 | uint64_t len = hton64(iocb->payload_length); |
242 | *buf_last |= 127; |
243 | ++buf_last; |
244 | memcpy(buf_last, &len, 8); |
245 | buf_last += 8; |
246 | } else { |
247 | /* Too large payload length */ |
248 | return WSLAY_ERR_INVALID_ARGUMENT; |
249 | } |
250 | if (iocb->mask) { |
251 | if (ctx->callbacks.genmask_callback(ctx->omaskkey, 4, ctx->user_data) != |
252 | 0) { |
253 | return WSLAY_ERR_INVALID_CALLBACK; |
254 | } else { |
255 | ctx->omask = 1; |
256 | memcpy(buf_last, ctx->omaskkey, 4); |
257 | buf_last += 4; |
258 | } |
259 | } |
260 | ctx->ostate = SEND_PAYLOAD; |
261 | ctx->opayloadlen = iocb->payload_length; |
262 | ctx->opayloadoff = 0; |
263 | |
264 | buflen -= (size_t)(buf_last - buf); |
265 | /* fall through */ |
266 | case SEND_PAYLOAD: |
267 | if (iocb->data_length > 0) { |
268 | size_t writelen = wslay_min(buflen, iocb->data_length); |
269 | |
270 | if (ctx->omask) { |
271 | for (i = 0; i < writelen; ++i) { |
272 | *buf_last++ = |
273 | iocb->data[i] ^ ctx->omaskkey[(ctx->opayloadoff + i) % 4]; |
274 | } |
275 | } else { |
276 | memcpy(buf_last, iocb->data, writelen); |
277 | buf_last += writelen; |
278 | } |
279 | |
280 | ctx->opayloadoff += writelen; |
281 | *pwpayloadlen = writelen; |
282 | } |
283 | |
284 | if (ctx->opayloadoff == ctx->opayloadlen) { |
285 | ctx->ostate = PREP_HEADER; |
286 | } |
287 | |
288 | return buf_last - buf; |
289 | default: |
290 | return WSLAY_ERR_INVALID_ARGUMENT; |
291 | } |
292 | } |
293 | |
294 | static void wslay_shift_ibuf(wslay_frame_context_ptr ctx) { |
295 | ptrdiff_t len = ctx->ibuflimit - ctx->ibufmark; |
296 | memmove(ctx->ibuf, ctx->ibufmark, (size_t)len); |
297 | ctx->ibuflimit = ctx->ibuf + len; |
298 | ctx->ibufmark = ctx->ibuf; |
299 | } |
300 | |
301 | static ssize_t wslay_recv(wslay_frame_context_ptr ctx) { |
302 | ssize_t r; |
303 | if (ctx->ibufmark != ctx->ibuf) { |
304 | wslay_shift_ibuf(ctx); |
305 | } |
306 | r = ctx->callbacks.recv_callback( |
307 | ctx->ibuflimit, (size_t)(ctx->ibuf + sizeof(ctx->ibuf) - ctx->ibuflimit), |
308 | 0, ctx->user_data); |
309 | if (r > 0) { |
310 | ctx->ibuflimit += r; |
311 | } else { |
312 | r = WSLAY_ERR_WANT_READ; |
313 | } |
314 | return r; |
315 | } |
316 | |
317 | #define WSLAY_AVAIL_IBUF(ctx) ((size_t)(ctx->ibuflimit - ctx->ibufmark)) |
318 | |
319 | ssize_t wslay_frame_recv(wslay_frame_context_ptr ctx, |
320 | struct wslay_frame_iocb *iocb) { |
321 | ssize_t r; |
322 | if (ctx->istate == RECV_HEADER1) { |
323 | uint8_t fin, opcode, rsv, payloadlen; |
324 | if (WSLAY_AVAIL_IBUF(ctx) < ctx->ireqread) { |
325 | if ((r = wslay_recv(ctx)) <= 0) { |
326 | return r; |
327 | } |
328 | } |
329 | if (WSLAY_AVAIL_IBUF(ctx) < ctx->ireqread) { |
330 | return WSLAY_ERR_WANT_READ; |
331 | } |
332 | fin = (ctx->ibufmark[0] >> 7) & 1; |
333 | rsv = (ctx->ibufmark[0] >> 4) & 7; |
334 | opcode = ctx->ibufmark[0] & 0xfu; |
335 | ctx->iom.opcode = opcode; |
336 | ctx->iom.fin = fin; |
337 | ctx->iom.rsv = rsv; |
338 | ++ctx->ibufmark; |
339 | ctx->imask = (ctx->ibufmark[0] >> 7) & 1; |
340 | payloadlen = ctx->ibufmark[0] & 0x7fu; |
341 | ++ctx->ibufmark; |
342 | if (wslay_is_ctrl_frame(opcode) && (payloadlen > 125 || !fin)) { |
343 | return WSLAY_ERR_PROTO; |
344 | } |
345 | if (payloadlen == 126) { |
346 | ctx->istate = RECV_EXT_PAYLOADLEN; |
347 | ctx->ireqread = 2; |
348 | } else if (payloadlen == 127) { |
349 | ctx->istate = RECV_EXT_PAYLOADLEN; |
350 | ctx->ireqread = 8; |
351 | } else { |
352 | ctx->ipayloadlen = payloadlen; |
353 | ctx->ipayloadoff = 0; |
354 | if (ctx->imask) { |
355 | ctx->istate = RECV_MASKKEY; |
356 | ctx->ireqread = 4; |
357 | } else { |
358 | ctx->istate = RECV_PAYLOAD; |
359 | } |
360 | } |
361 | } |
362 | if (ctx->istate == RECV_EXT_PAYLOADLEN) { |
363 | if (WSLAY_AVAIL_IBUF(ctx) < ctx->ireqread) { |
364 | if ((r = wslay_recv(ctx)) <= 0) { |
365 | return r; |
366 | } |
367 | if (WSLAY_AVAIL_IBUF(ctx) < ctx->ireqread) { |
368 | return WSLAY_ERR_WANT_READ; |
369 | } |
370 | } |
371 | ctx->ipayloadlen = 0; |
372 | ctx->ipayloadoff = 0; |
373 | memcpy((uint8_t *)&ctx->ipayloadlen + (8 - ctx->ireqread), ctx->ibufmark, |
374 | ctx->ireqread); |
375 | ctx->ipayloadlen = ntoh64(ctx->ipayloadlen); |
376 | ctx->ibufmark += ctx->ireqread; |
377 | if (ctx->ireqread == 8) { |
378 | if (ctx->ipayloadlen < (1 << 16) || ctx->ipayloadlen & (1ull << 63)) { |
379 | return WSLAY_ERR_PROTO; |
380 | } |
381 | } else if (ctx->ipayloadlen < 126) { |
382 | return WSLAY_ERR_PROTO; |
383 | } |
384 | if (ctx->imask) { |
385 | ctx->istate = RECV_MASKKEY; |
386 | ctx->ireqread = 4; |
387 | } else { |
388 | ctx->istate = RECV_PAYLOAD; |
389 | } |
390 | } |
391 | if (ctx->istate == RECV_MASKKEY) { |
392 | if (WSLAY_AVAIL_IBUF(ctx) < ctx->ireqread) { |
393 | if ((r = wslay_recv(ctx)) <= 0) { |
394 | return r; |
395 | } |
396 | if (WSLAY_AVAIL_IBUF(ctx) < ctx->ireqread) { |
397 | return WSLAY_ERR_WANT_READ; |
398 | } |
399 | } |
400 | memcpy(ctx->imaskkey, ctx->ibufmark, 4); |
401 | ctx->ibufmark += 4; |
402 | ctx->istate = RECV_PAYLOAD; |
403 | } |
404 | if (ctx->istate == RECV_PAYLOAD) { |
405 | uint8_t *readlimit, *readmark; |
406 | uint64_t rempayloadlen = ctx->ipayloadlen - ctx->ipayloadoff; |
407 | if (WSLAY_AVAIL_IBUF(ctx) == 0 && rempayloadlen > 0) { |
408 | if ((r = wslay_recv(ctx)) <= 0) { |
409 | return r; |
410 | } |
411 | } |
412 | readmark = ctx->ibufmark; |
413 | readlimit = WSLAY_AVAIL_IBUF(ctx) < rempayloadlen |
414 | ? ctx->ibuflimit |
415 | : ctx->ibufmark + rempayloadlen; |
416 | if (ctx->imask) { |
417 | for (; ctx->ibufmark != readlimit; ++ctx->ibufmark, ++ctx->ipayloadoff) { |
418 | ctx->ibufmark[0] ^= ctx->imaskkey[ctx->ipayloadoff % 4]; |
419 | } |
420 | } else { |
421 | ctx->ibufmark = readlimit; |
422 | ctx->ipayloadoff += (uint64_t)(readlimit - readmark); |
423 | } |
424 | iocb->fin = ctx->iom.fin; |
425 | iocb->rsv = ctx->iom.rsv; |
426 | iocb->opcode = ctx->iom.opcode; |
427 | iocb->payload_length = ctx->ipayloadlen; |
428 | iocb->mask = ctx->imask; |
429 | iocb->data = readmark; |
430 | iocb->data_length = (size_t)(ctx->ibufmark - readmark); |
431 | if (ctx->ipayloadlen == ctx->ipayloadoff) { |
432 | ctx->istate = RECV_HEADER1; |
433 | ctx->ireqread = 2; |
434 | } |
435 | return (ssize_t)iocb->data_length; |
436 | } |
437 | return WSLAY_ERR_INVALID_ARGUMENT; |
438 | } |
439 | |