1/*
2 * libwebsockets - small server side websockets and web server implementation
3 *
4 * Copyright (C) 2010 - 2019 Andy Green <andy@warmcat.com>
5 *
6 * Permission is hereby granted, free of charge, to any person obtaining a copy
7 * of this software and associated documentation files (the "Software"), to
8 * deal in the Software without restriction, including without limitation the
9 * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
10 * sell copies of the Software, and to permit persons to whom the Software is
11 * furnished to do so, subject to the following conditions:
12 *
13 * The above copyright notice and this permission notice shall be included in
14 * all copies or substantial portions of the Software.
15 *
16 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
17 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
18 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
19 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
20 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
21 * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
22 * IN THE SOFTWARE.
23 */
24
25#include "private-lib-core.h"
26#include "extension-permessage-deflate.h"
27#include <stdio.h>
28#include <string.h>
29#include <assert.h>
30
31#define LWS_ZLIB_MEMLEVEL 8
32
33const struct lws_ext_options lws_ext_pm_deflate_options[] = {
34	/* public RFC7692 settings */
35	{ "server_no_context_takeover", EXTARG_NONE },
36	{ "client_no_context_takeover", EXTARG_NONE },
37	{ "server_max_window_bits",	EXTARG_OPT_DEC },
38	{ "client_max_window_bits",	EXTARG_OPT_DEC },
39	/* ones only user code can set */
40	{ "rx_buf_size",		EXTARG_DEC },
41	{ "tx_buf_size",		EXTARG_DEC },
42	{ "compression_level",		EXTARG_DEC },
43	{ "mem_level",			EXTARG_DEC },
44	{ NULL, 0 }, /* sentinel */
45};
46
47static void
48lws_extension_pmdeflate_restrict_args(struct lws *wsi,
49				      struct lws_ext_pm_deflate_priv *priv)
50{
51	int n, extra;
52
53	/* cap the RX buf at the nearest power of 2 to protocol rx buf */
54
55	n = (int)wsi->a.context->pt_serv_buf_size;
56	if (wsi->a.protocol->rx_buffer_size)
57		n = (int)wsi->a.protocol->rx_buffer_size;
58
59	extra = 7;
60	while (n >= 1 << (extra + 1))
61		extra++;
62
63	if (extra < priv->args[PMD_RX_BUF_PWR2]) {
64		priv->args[PMD_RX_BUF_PWR2] = (unsigned char)extra;
65		lwsl_wsi_info(wsi, " Capping pmd rx to %d", 1 << extra);
66	}
67}
68
69static unsigned char trail[] = { 0, 0, 0xff, 0xff };
70
71LWS_VISIBLE int
72lws_extension_callback_pm_deflate(struct lws_context *context,
73				  const struct lws_extension *ext,
74				  struct lws *wsi,
75				  enum lws_extension_callback_reasons reason,
76				  void *user, void *in, size_t len)
77{
78	struct lws_ext_pm_deflate_priv *priv =
79				     (struct lws_ext_pm_deflate_priv *)user;
80	struct lws_ext_pm_deflate_rx_ebufs *pmdrx =
81				(struct lws_ext_pm_deflate_rx_ebufs *)in;
82	struct lws_ext_option_arg *oa;
83	int n, ret = 0, was_fin = 0, m;
84	unsigned int pen = 0;
85	int penbits = 0;
86
87	switch (reason) {
88	case LWS_EXT_CB_NAMED_OPTION_SET:
89		oa = in;
90		if (!oa->option_name)
91			break;
92		lwsl_wsi_ext(wsi, "named option set: %s", oa->option_name);
93		for (n = 0; n < (int)LWS_ARRAY_SIZE(lws_ext_pm_deflate_options);
94		     n++)
95			if (!strcmp(lws_ext_pm_deflate_options[n].name,
96				    oa->option_name))
97				break;
98
99		if (n == (int)LWS_ARRAY_SIZE(lws_ext_pm_deflate_options))
100			break;
101		oa->option_index = n;
102
103		/* fallthru */
104
105	case LWS_EXT_CB_OPTION_SET:
106		oa = in;
107		lwsl_wsi_ext(wsi, "option set: idx %d, %s, len %d",
108			 oa->option_index, oa->start, oa->len);
109		if (oa->start)
110			priv->args[oa->option_index] = (unsigned char)atoi(oa->start);
111		else
112			priv->args[oa->option_index] = 1;
113
114		if (priv->args[PMD_CLIENT_MAX_WINDOW_BITS] == 8)
115			priv->args[PMD_CLIENT_MAX_WINDOW_BITS] = 9;
116
117		lws_extension_pmdeflate_restrict_args(wsi, priv);
118		break;
119
120	case LWS_EXT_CB_OPTION_CONFIRM:
121		if (priv->args[PMD_SERVER_MAX_WINDOW_BITS] < 8 ||
122		    priv->args[PMD_SERVER_MAX_WINDOW_BITS] > 15 ||
123		    priv->args[PMD_CLIENT_MAX_WINDOW_BITS] < 8 ||
124		    priv->args[PMD_CLIENT_MAX_WINDOW_BITS] > 15)
125			return -1;
126		break;
127
128	case LWS_EXT_CB_CLIENT_CONSTRUCT:
129	case LWS_EXT_CB_CONSTRUCT:
130
131		n = (int)context->pt_serv_buf_size;
132		if (wsi->a.protocol->rx_buffer_size)
133			n = (int)wsi->a.protocol->rx_buffer_size;
134
135		if (n < 128) {
136			lwsl_wsi_info(wsi, " permessage-deflate requires the protocol "
137				  "(%s) to have an RX buffer >= 128",
138				  wsi->a.protocol->name);
139			return -1;
140		}
141
142		/* fill in **user */
143		priv = lws_zalloc(sizeof(*priv), "pmd priv");
144		*((void **)user) = priv;
145		lwsl_wsi_ext(wsi, "LWS_EXT_CB_*CONSTRUCT");
146		memset(priv, 0, sizeof(*priv));
147
148		/* fill in pointer to options list */
149		if (in)
150			*((const struct lws_ext_options **)in) =
151					lws_ext_pm_deflate_options;
152
153		/* fallthru */
154
155	case LWS_EXT_CB_OPTION_DEFAULT:
156
157		/* set the public, RFC7692 defaults... */
158
159		priv->args[PMD_SERVER_NO_CONTEXT_TAKEOVER] = 0,
160		priv->args[PMD_CLIENT_NO_CONTEXT_TAKEOVER] = 0;
161		priv->args[PMD_SERVER_MAX_WINDOW_BITS] = 15;
162		priv->args[PMD_CLIENT_MAX_WINDOW_BITS] = 15;
163
164		/* ...and the ones the user code can override */
165
166		priv->args[PMD_RX_BUF_PWR2] = 10; /* ie, 1024 */
167		priv->args[PMD_TX_BUF_PWR2] = 10; /* ie, 1024 */
168		priv->args[PMD_COMP_LEVEL] = 1;
169		priv->args[PMD_MEM_LEVEL] = 8;
170
171		lws_extension_pmdeflate_restrict_args(wsi, priv);
172		break;
173
174	case LWS_EXT_CB_DESTROY:
175		lwsl_wsi_ext(wsi, "LWS_EXT_CB_DESTROY");
176		lws_free(priv->buf_rx_inflated);
177		lws_free(priv->buf_tx_deflated);
178		if (priv->rx_init)
179			(void)inflateEnd(&priv->rx);
180		if (priv->tx_init)
181			(void)deflateEnd(&priv->tx);
182		lws_free(priv);
183
184		return ret;
185
186
187	case LWS_EXT_CB_PAYLOAD_RX:
188		/*
189		 * ie, we are INFLATING
190		 */
191		lwsl_wsi_ext(wsi, " LWS_EXT_CB_PAYLOAD_RX: in %d, existing in %d",
192			 pmdrx->eb_in.len, priv->rx.avail_in);
193
194		/*
195		 * If this frame is not marked as compressed,
196		 * there is nothing we should do with it
197		 */
198
199		if (!(wsi->ws->rsv_first_msg & 0x40) || (wsi->ws->opcode & 8))
200			/*
201			 * This is a bit different than DID_NOTHING... we have
202			 * identified using ext-private bits in the packet, or
203			 * by it being a control fragment that we SHOULD not do
204			 * anything to it, parent should continue as if we
205			 * processed it
206			 */
207			return PMDR_NOTHING_WE_SHOULD_DO;
208
209		/*
210		 * we shouldn't come back in here if we already applied the
211		 * trailer for this compressed packet
212		 */
213		if (!wsi->ws->pmd_trailer_application)
214			return PMDR_DID_NOTHING;
215
216		pmdrx->eb_out.len = 0;
217
218		lwsl_wsi_ext(wsi, "LWS_EXT_CB_PAYLOAD_RX: in %d, "
219			 "existing avail in %d, pkt fin: %d",
220			 pmdrx->eb_in.len, priv->rx.avail_in, wsi->ws->final);
221
222		/* if needed, initialize the inflator */
223
224		if (!priv->rx_init) {
225			if (inflateInit2(&priv->rx,
226			     -priv->args[PMD_SERVER_MAX_WINDOW_BITS]) != Z_OK) {
227				lwsl_wsi_err(wsi, "iniflateInit failed");
228				return PMDR_FAILED;
229			}
230			priv->rx_init = 1;
231			if (!priv->buf_rx_inflated)
232				priv->buf_rx_inflated = lws_malloc(
233					(unsigned int)(LWS_PRE + 7 + 5 +
234					    (1 << priv->args[PMD_RX_BUF_PWR2])),
235					    "pmd rx inflate buf");
236			if (!priv->buf_rx_inflated) {
237				lwsl_wsi_err(wsi, "OOM");
238				return PMDR_FAILED;
239			}
240		}
241
242#if 0
243		/*
244		 * don't give us new input while we still work through
245		 * the last input
246		 */
247
248		if (priv->rx.avail_in && pmdrx->eb_in.token &&
249					 pmdrx->eb_in.len) {
250			lwsl_wsi_warn(wsi, "priv->rx.avail_in %d while getting new in",
251					priv->rx.avail_in);
252	//		assert(0);
253		}
254#endif
255		if (!priv->rx.avail_in && pmdrx->eb_in.token && pmdrx->eb_in.len) {
256			priv->rx.next_in = (unsigned char *)pmdrx->eb_in.token;
257			priv->rx.avail_in = (uInt)pmdrx->eb_in.len;
258		}
259
260		priv->rx.next_out = priv->buf_rx_inflated + LWS_PRE;
261		pmdrx->eb_out.token = priv->rx.next_out;
262		priv->rx.avail_out = (uInt)(1 << priv->args[PMD_RX_BUF_PWR2]);
263
264		/* so... if...
265		 *
266		 *  - he has no remaining input content for this message, and
267		 *
268		 *  - and this is the final fragment, and
269		 *
270		 *  - we used everything that could be drained on the input side
271		 *
272		 * ...then put back the 00 00 FF FF the sender stripped as our
273		 * input to zlib
274		 */
275		if (!priv->rx.avail_in &&
276		    wsi->ws->final &&
277		    !wsi->ws->rx_packet_length &&
278		    wsi->ws->pmd_trailer_application) {
279			lwsl_wsi_ext(wsi, "trailer apply 1");
280			was_fin = 1;
281			wsi->ws->pmd_trailer_application = 0;
282			priv->rx.next_in = trail;
283			priv->rx.avail_in = sizeof(trail);
284		}
285
286		/*
287		 * if after all that there's nothing pending and nothing to give
288		 * him right now, bail without having done anything
289		 */
290
291		if (!priv->rx.avail_in)
292			return PMDR_DID_NOTHING;
293
294		n = inflate(&priv->rx, was_fin ? Z_SYNC_FLUSH : Z_NO_FLUSH);
295		lwsl_wsi_ext(wsi, "inflate ret %d, avi %d, avo %d, wsifinal %d", n,
296			 priv->rx.avail_in, priv->rx.avail_out, wsi->ws->final);
297		switch (n) {
298		case Z_NEED_DICT:
299		case Z_STREAM_ERROR:
300		case Z_DATA_ERROR:
301		case Z_MEM_ERROR:
302			lwsl_wsi_err(wsi, "zlib error inflate %d: \"%s\"",
303				  n, priv->rx.msg);
304			return PMDR_FAILED;
305		}
306
307		/*
308		 * track how much input was used, and advance it
309		 */
310
311		pmdrx->eb_in.token = pmdrx->eb_in.token +
312				         ((unsigned int)pmdrx->eb_in.len - (unsigned int)priv->rx.avail_in);
313		pmdrx->eb_in.len = (int)priv->rx.avail_in;
314
315		lwsl_wsi_debug(wsi, "%d %d %d %d %d",
316				priv->rx.avail_in,
317				wsi->ws->final,
318				(int)wsi->ws->rx_packet_length,
319				was_fin,
320				wsi->ws->pmd_trailer_application);
321
322		if (!priv->rx.avail_in &&
323		    wsi->ws->final &&
324		    !wsi->ws->rx_packet_length &&
325		    !was_fin &&
326		    wsi->ws->pmd_trailer_application) {
327			lwsl_wsi_ext(wsi, "RX trailer apply 2");
328
329			/* we overallocated just for this situation where
330			 * we might issue something */
331			priv->rx.avail_out += 5;
332
333			was_fin = 1;
334			wsi->ws->pmd_trailer_application = 0;
335			priv->rx.next_in = trail;
336			priv->rx.avail_in = sizeof(trail);
337			n = inflate(&priv->rx, Z_SYNC_FLUSH);
338			lwsl_wsi_ext(wsi, "RX trailer infl ret %d, avi %d, avo %d",
339				 n, priv->rx.avail_in, priv->rx.avail_out);
340			switch (n) {
341			case Z_NEED_DICT:
342			case Z_STREAM_ERROR:
343			case Z_DATA_ERROR:
344			case Z_MEM_ERROR:
345				lwsl_wsi_info(wsi, "zlib error inflate %d: %s",
346					  n, priv->rx.msg);
347				return -1;
348			}
349
350			assert(priv->rx.avail_out);
351		}
352
353		pmdrx->eb_out.len = lws_ptr_diff(priv->rx.next_out,
354						 pmdrx->eb_out.token);
355		priv->count_rx_between_fin = priv->count_rx_between_fin + (size_t)pmdrx->eb_out.len;
356
357		lwsl_wsi_ext(wsi, "  RX leaving with new effbuff len %d, "
358			 "rx.avail_in=%d, TOTAL RX since FIN %lu",
359			 pmdrx->eb_out.len, priv->rx.avail_in,
360			 (unsigned long)priv->count_rx_between_fin);
361
362		if (was_fin) {
363			lwsl_wsi_ext(wsi, "was_fin");
364			priv->count_rx_between_fin = 0;
365			if (priv->args[PMD_SERVER_NO_CONTEXT_TAKEOVER]) {
366				lwsl_wsi_ext(wsi, "PMD_SERVER_NO_CONTEXT_TAKEOVER");
367				(void)inflateEnd(&priv->rx);
368				priv->rx_init = 0;
369			}
370
371			return PMDR_EMPTY_FINAL;
372		}
373
374		if (priv->rx.avail_in)
375			return PMDR_HAS_PENDING;
376
377		return PMDR_EMPTY_NONFINAL;
378
379	case LWS_EXT_CB_PAYLOAD_TX:
380
381		/*
382		 * ie, we are DEFLATING
383		 *
384		 * initialize us if needed
385		 */
386
387		if (!priv->tx_init) {
388			n = deflateInit2(&priv->tx, priv->args[PMD_COMP_LEVEL],
389					 Z_DEFLATED,
390					 -priv->args[PMD_SERVER_MAX_WINDOW_BITS +
391						(wsi->a.vhost->listen_port <= 0)],
392					 priv->args[PMD_MEM_LEVEL],
393					 Z_DEFAULT_STRATEGY);
394			if (n != Z_OK) {
395				lwsl_wsi_ext(wsi, "inflateInit2 failed %d", n);
396				return PMDR_FAILED;
397			}
398			priv->tx_init = 1;
399		}
400
401		if (!priv->buf_tx_deflated)
402			priv->buf_tx_deflated = lws_malloc((unsigned int)(LWS_PRE + 7 + 5 +
403					    (1 << priv->args[PMD_TX_BUF_PWR2])),
404					    "pmd tx deflate buf");
405		if (!priv->buf_tx_deflated) {
406			lwsl_wsi_err(wsi, "OOM");
407			return PMDR_FAILED;
408		}
409
410		/* hook us up with any deflated input that the caller has */
411
412		if (pmdrx->eb_in.token) {
413
414			assert(!priv->tx.avail_in);
415
416			priv->count_tx_between_fin = priv->count_tx_between_fin + (size_t)pmdrx->eb_in.len;
417			lwsl_wsi_ext(wsi, "TX: eb_in length %d, "
418				    "TOTAL TX since FIN: %d",
419				    pmdrx->eb_in.len,
420				    (int)priv->count_tx_between_fin);
421			priv->tx.next_in = (unsigned char *)pmdrx->eb_in.token;
422			priv->tx.avail_in = (uInt)pmdrx->eb_in.len;
423		}
424
425		priv->tx.next_out = priv->buf_tx_deflated + LWS_PRE + 5;
426		pmdrx->eb_out.token = priv->tx.next_out;
427		priv->tx.avail_out = (uInt)(1 << priv->args[PMD_TX_BUF_PWR2]);
428
429		pen = 0;
430		penbits = 0;
431		deflatePending(&priv->tx, &pen, &penbits);
432		pen = pen | (unsigned int)penbits;
433
434		if (!priv->tx.avail_in && (len & LWS_WRITE_NO_FIN)) {
435			lwsl_wsi_ext(wsi, "no available in, pen: %u", pen);
436
437			if (!pen)
438				return PMDR_DID_NOTHING;
439		}
440
441		m = Z_NO_FLUSH;
442		if (!(len & LWS_WRITE_NO_FIN)) {
443			lwsl_wsi_ext(wsi, "deflate with SYNC_FLUSH, pkt len %d",
444					(int)wsi->ws->rx_packet_length);
445			m = Z_SYNC_FLUSH;
446		}
447
448		n = deflate(&priv->tx, m);
449		if (n == Z_STREAM_ERROR) {
450			lwsl_wsi_notice(wsi, "Z_STREAM_ERROR");
451			return PMDR_FAILED;
452		}
453
454		pen = (!priv->tx.avail_out) && n != Z_STREAM_END;
455
456		lwsl_wsi_ext(wsi, "deflate ret %d, len 0x%x", n,
457				(unsigned int)len);
458
459		if ((len & 0xf) == LWS_WRITE_TEXT)
460			priv->tx_first_frame_type = LWSWSOPC_TEXT_FRAME;
461		if ((len & 0xf) == LWS_WRITE_BINARY)
462			priv->tx_first_frame_type = LWSWSOPC_BINARY_FRAME;
463
464		pmdrx->eb_out.len = lws_ptr_diff(priv->tx.next_out,
465						 pmdrx->eb_out.token);
466
467		if (m == Z_SYNC_FLUSH && !(len & LWS_WRITE_NO_FIN) && !pen &&
468		    pmdrx->eb_out.len < 4) {
469			lwsl_wsi_err(wsi, "FAIL want to trim out length %d",
470					(int)pmdrx->eb_out.len);
471			assert(0);
472		}
473
474		if (!(len & LWS_WRITE_NO_FIN) &&
475		    m == Z_SYNC_FLUSH &&
476		    !pen &&
477		    pmdrx->eb_out.len >= 4) {
478			// lwsl_wsi_err(wsi, "Trimming 4 from end of write");
479			priv->tx.next_out -= 4;
480			priv->tx.avail_out += 4;
481			priv->count_tx_between_fin = 0;
482
483			assert(priv->tx.next_out[0] == 0x00 &&
484			       priv->tx.next_out[1] == 0x00 &&
485			       priv->tx.next_out[2] == 0xff &&
486			       priv->tx.next_out[3] == 0xff);
487		}
488
489
490		/*
491		 * track how much input was used and advance it
492		 */
493
494		pmdrx->eb_in.token = pmdrx->eb_in.token +
495					((unsigned int)pmdrx->eb_in.len - (unsigned int)priv->tx.avail_in);
496		pmdrx->eb_in.len = (int)priv->tx.avail_in;
497
498		priv->compressed_out = 1;
499		pmdrx->eb_out.len = lws_ptr_diff(priv->tx.next_out,
500						 pmdrx->eb_out.token);
501
502		lwsl_wsi_ext(wsi, "  TX rewritten with new eb_in len %d, "
503				"eb_out len %d, deflatePending %d",
504				pmdrx->eb_in.len, pmdrx->eb_out.len, pen);
505
506		if (pmdrx->eb_in.len || pen)
507			return PMDR_HAS_PENDING;
508
509		if (!(len & LWS_WRITE_NO_FIN))
510			return PMDR_EMPTY_FINAL;
511
512		return PMDR_EMPTY_NONFINAL;
513
514	case LWS_EXT_CB_PACKET_TX_PRESEND:
515		if (!priv->compressed_out)
516			break;
517		priv->compressed_out = 0;
518
519		/*
520		 * we may have not produced any output for the actual "first"
521		 * write... in that case, we need to fix up the inappropriate
522		 * use of CONTINUATION when the first real write does come.
523		 */
524		if (priv->tx_first_frame_type & 0xf) {
525			*pmdrx->eb_in.token = (unsigned char)((((unsigned char)*pmdrx->eb_in.token) & (unsigned char)~0xf) |
526				((unsigned char)priv->tx_first_frame_type & (unsigned char)0xf));
527			/*
528			 * We have now written the "first" fragment, only
529			 * do that once
530			 */
531			priv->tx_first_frame_type = 0;
532		}
533
534		n = *(pmdrx->eb_in.token) & 15;
535
536		/* set RSV1, but not on CONTINUATION */
537		if (n == LWSWSOPC_TEXT_FRAME || n == LWSWSOPC_BINARY_FRAME)
538			*pmdrx->eb_in.token |= 0x40;
539
540		lwsl_wsi_ext(wsi, "PRESEND compressed: ws frame 0x%02X, len %d",
541			    ((*pmdrx->eb_in.token) & 0xff),
542			    pmdrx->eb_in.len);
543
544		if (((*pmdrx->eb_in.token) & 0x80) &&	/* fin */
545		    priv->args[PMD_CLIENT_NO_CONTEXT_TAKEOVER]) {
546			lwsl_wsi_debug(wsi, "PMD_CLIENT_NO_CONTEXT_TAKEOVER");
547			(void)deflateEnd(&priv->tx);
548			priv->tx_init = 0;
549		}
550
551		break;
552
553	default:
554		break;
555	}
556
557	return 0;
558}
559
560