Branch data Line data Source code
1 : : /*
2 : : * This file is part of the MicroPython project, http://micropython.org/
3 : : *
4 : : * The MIT License (MIT)
5 : : *
6 : : * Copyright (c) 2023 Jim Mussared
7 : : *
8 : : * Based on extmod/modzlib.c
9 : : * Copyright (c) 2014-2016 Paul Sokolovsky
10 : : * Copyright (c) 2021-2023 Damien P. George
11 : : *
12 : : * Permission is hereby granted, free of charge, to any person obtaining a copy
13 : : * of this software and associated documentation files (the "Software"), to deal
14 : : * in the Software without restriction, including without limitation the rights
15 : : * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
16 : : * copies of the Software, and to permit persons to whom the Software is
17 : : * furnished to do so, subject to the following conditions:
18 : : *
19 : : * The above copyright notice and this permission notice shall be included in
20 : : * all copies or substantial portions of the Software.
21 : : *
22 : : * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
23 : : * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
24 : : * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
25 : : * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
26 : : * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
27 : : * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
28 : : * THE SOFTWARE.
29 : : */
30 : :
31 : : #include <stdio.h>
32 : : #include <string.h>
33 : :
34 : : #include "py/runtime.h"
35 : : #include "py/stream.h"
36 : : #include "py/mperrno.h"
37 : :
38 : : #if MICROPY_PY_DEFLATE
39 : :
40 : : #include "lib/uzlib/uzlib.h"
41 : :
42 : : #if 0 // print debugging info
43 : : #define DEBUG_printf DEBUG_printf
44 : : #else // don't print debugging info
45 : : #define DEBUG_printf(...) (void)0
46 : : #endif
47 : :
48 : : typedef enum {
49 : : DEFLATEIO_FORMAT_MIN = 0,
50 : : DEFLATEIO_FORMAT_AUTO = DEFLATEIO_FORMAT_MIN, // Read mode this means auto-detect zlib/gzip, write mode this means RAW.
51 : : DEFLATEIO_FORMAT_RAW = 1,
52 : : DEFLATEIO_FORMAT_ZLIB = 2,
53 : : DEFLATEIO_FORMAT_GZIP = 3,
54 : : DEFLATEIO_FORMAT_MAX = DEFLATEIO_FORMAT_GZIP,
55 : : } deflateio_format_t;
56 : :
57 : : // This is used when the wbits is unset in the DeflateIO constructor. Default
58 : : // to the smallest window size (faster compression, less RAM usage, etc).
59 : : const int DEFLATEIO_DEFAULT_WBITS = 8;
60 : :
61 : : typedef struct {
62 : : void *window;
63 : : uzlib_uncomp_t decomp;
64 : : bool eof;
65 : : } mp_obj_deflateio_read_t;
66 : :
67 : : #if MICROPY_PY_DEFLATE_COMPRESS
68 : : typedef struct {
69 : : void *window;
70 : : size_t input_len;
71 : : uint32_t input_checksum;
72 : : uzlib_lz77_state_t lz77;
73 : : } mp_obj_deflateio_write_t;
74 : : #endif
75 : :
76 : : typedef struct {
77 : : mp_obj_base_t base;
78 : : mp_obj_t stream;
79 : : uint8_t format : 2;
80 : : uint8_t window_bits : 4;
81 : : bool close : 1;
82 : : mp_obj_deflateio_read_t *read;
83 : : #if MICROPY_PY_DEFLATE_COMPRESS
84 : : mp_obj_deflateio_write_t *write;
85 : : #endif
86 : : } mp_obj_deflateio_t;
87 : :
88 : 10500 : static int deflateio_read_stream(void *data) {
89 : 10500 : mp_obj_deflateio_t *self = data;
90 : 10500 : const mp_stream_p_t *stream = mp_get_stream(self->stream);
91 : 10500 : int err;
92 : 10500 : byte c;
93 : 10500 : mp_uint_t out_sz = stream->read(self->stream, &c, 1, &err);
94 [ + + ]: 10498 : if (out_sz == MP_STREAM_ERROR) {
95 : 2 : mp_raise_OSError(err);
96 : : }
97 [ + + ]: 10496 : if (out_sz == 0) {
98 : 2 : mp_raise_type(&mp_type_EOFError);
99 : : }
100 : 10494 : return c;
101 : : }
102 : :
103 : 376 : static bool deflateio_init_read(mp_obj_deflateio_t *self) {
104 [ + + ]: 376 : if (self->read) {
105 : : return true;
106 : : }
107 : :
108 : 142 : mp_get_stream_raise(self->stream, MP_STREAM_OP_READ);
109 : :
110 : 142 : self->read = m_new_obj(mp_obj_deflateio_read_t);
111 : 142 : memset(&self->read->decomp, 0, sizeof(self->read->decomp));
112 : 142 : self->read->decomp.source_read_data = self;
113 : 142 : self->read->decomp.source_read_cb = deflateio_read_stream;
114 : 142 : self->read->eof = false;
115 : :
116 : : // Don't modify self->window_bits as it may also be used for write.
117 : 142 : int wbits = self->window_bits;
118 : :
119 [ + + ]: 142 : if (self->format == DEFLATEIO_FORMAT_RAW) {
120 [ + + ]: 72 : if (wbits == 0) {
121 : : // The docs recommends always setting wbits explicitly when using
122 : : // RAW, but we still allow a default.
123 : 26 : wbits = DEFLATEIO_DEFAULT_WBITS;
124 : : }
125 : : } else {
126 : : // Parse the header if we're in NONE/ZLIB/GZIP modes.
127 : 70 : int header_wbits;
128 : 70 : int header_type = uzlib_parse_zlib_gzip_header(&self->read->decomp, &header_wbits);
129 [ + + ]: 68 : if (header_type < 0) {
130 : : // Stream header was invalid.
131 : 10 : return false;
132 : : }
133 [ + + + + : 62 : if ((self->format == DEFLATEIO_FORMAT_ZLIB && header_type != UZLIB_HEADER_ZLIB) || (self->format == DEFLATEIO_FORMAT_GZIP && header_type != UZLIB_HEADER_GZIP)) {
+ + + + ]
134 : : // Not what we expected.
135 : : return false;
136 : : }
137 : : // header_wbits will either be 15 (gzip) or 8-15 (zlib).
138 [ + + - + ]: 58 : if (wbits == 0 || header_wbits < wbits) {
139 : : // If the header specified something lower, then use that instead.
140 : : // No point doing a bigger allocation than we need to.
141 : 54 : wbits = header_wbits;
142 : : }
143 : : }
144 : :
145 : 130 : size_t window_len = (size_t)1 << wbits;
146 : 130 : self->read->window = m_new(uint8_t, window_len);
147 : :
148 : 130 : uzlib_uncompress_init(&self->read->decomp, self->read->window, window_len);
149 : :
150 : 130 : return true;
151 : : }
152 : :
153 : : #if MICROPY_PY_DEFLATE_COMPRESS
154 : 10202 : static void deflateio_out_byte(void *data, uint8_t b) {
155 : 10202 : mp_obj_deflateio_t *self = data;
156 : 10202 : const mp_stream_p_t *stream = mp_get_stream(self->stream);
157 : 10202 : int err;
158 : 10202 : mp_uint_t ret = stream->write(self->stream, &b, 1, &err);
159 [ + + ]: 10200 : if (ret == MP_STREAM_ERROR) {
160 : 2 : mp_raise_OSError(err);
161 : : }
162 : 10198 : }
163 : :
164 : 120 : static bool deflateio_init_write(mp_obj_deflateio_t *self) {
165 [ + + ]: 120 : if (self->write) {
166 : : return true;
167 : : }
168 : :
169 : 82 : const mp_stream_p_t *stream = mp_get_stream_raise(self->stream, MP_STREAM_OP_WRITE);
170 : :
171 : 82 : int wbits = self->window_bits;
172 [ + + ]: 82 : if (wbits == 0) {
173 : : // Same default wbits for all formats.
174 : 42 : wbits = DEFLATEIO_DEFAULT_WBITS;
175 : : }
176 : :
177 : : // Allocate the large window before allocating the mp_obj_deflateio_write_t, in case the
178 : : // window allocation fails the mp_obj_deflateio_t object will remain in a consistent state.
179 : 82 : size_t window_len = 1 << wbits;
180 : 82 : uint8_t *window = m_new(uint8_t, window_len);
181 : :
182 : 80 : self->write = m_new_obj(mp_obj_deflateio_write_t);
183 : 80 : self->write->window = window;
184 : 80 : self->write->input_len = 0;
185 : :
186 : 80 : uzlib_lz77_init(&self->write->lz77, self->write->window, window_len);
187 : 80 : self->write->lz77.dest_write_data = self;
188 : 80 : self->write->lz77.dest_write_cb = deflateio_out_byte;
189 : :
190 : : // Write header if needed.
191 : 80 : mp_uint_t ret = 0;
192 : 80 : int err;
193 [ + + ]: 80 : if (self->format == DEFLATEIO_FORMAT_ZLIB) {
194 : : // -----CMF------ ----------FLG---------------
195 : : // CINFO(5) CM(3) FLEVEL(2) FDICT(1) FCHECK(5)
196 : 8 : uint8_t buf[] = { 0x08, 0x80 }; // CM=2 (deflate), FLEVEL=2 (default), FDICT=0 (no dictionary)
197 : 8 : buf[0] |= MAX(wbits - 8, 1) << 4; // base-2 logarithm of the LZ77 window size, minus eight.
198 : 8 : buf[1] |= 31 - ((buf[0] * 256 + buf[1]) % 31); // (CMF*256 + FLG) % 31 == 0.
199 : 8 : ret = stream->write(self->stream, buf, sizeof(buf), &err);
200 : :
201 : 8 : self->write->input_checksum = 1; // ADLER32
202 [ + + ]: 72 : } else if (self->format == DEFLATEIO_FORMAT_GZIP) {
203 : : // ID1(8) ID2(8) CM(8) ---FLG--- MTIME(32) XFL(8) OS(8)
204 : : // FLG: x x x FCOMMENT FNAME FEXTRA FHCRC FTEXT
205 : 6 : uint8_t buf[] = { 0x1f, 0x8b, 0x08, 0x00, 0x00, 0x00, 0x00, 0x00, 0x04, 0x03 }; // MTIME=0, XFL=4 (fastest), OS=3 (unix)
206 : 6 : ret = stream->write(self->stream, buf, sizeof(buf), &err);
207 : :
208 : 6 : self->write->input_checksum = ~0; // CRC32
209 : : }
210 [ + + ]: 14 : if (ret == MP_STREAM_ERROR) {
211 : : return false;
212 : : }
213 : :
214 : : // Write starting block.
215 : 76 : uzlib_start_block(&self->write->lz77);
216 : :
217 : 76 : return true;
218 : : }
219 : : #endif
220 : :
221 : 250 : static mp_obj_t deflateio_make_new(const mp_obj_type_t *type, size_t n_args, size_t n_kw, const mp_obj_t *args_in) {
222 : : // args: stream, format=NONE, wbits=0, close=False
223 : 250 : mp_arg_check_num(n_args, n_kw, 1, 4, false);
224 : :
225 [ + + ]: 250 : mp_int_t format = n_args > 1 ? mp_obj_get_int(args_in[1]) : DEFLATEIO_FORMAT_AUTO;
226 [ + + ]: 208 : mp_int_t wbits = n_args > 2 ? mp_obj_get_int(args_in[2]) : 0;
227 : :
228 [ + + ]: 250 : if (format < DEFLATEIO_FORMAT_MIN || format > DEFLATEIO_FORMAT_MAX) {
229 : 8 : mp_raise_ValueError(MP_ERROR_TEXT("format"));
230 : : }
231 [ + + ]: 242 : if (wbits != 0 && (wbits < 5 || wbits > 15)) {
232 : 16 : mp_raise_ValueError(MP_ERROR_TEXT("wbits"));
233 : : }
234 : :
235 : 226 : mp_obj_deflateio_t *self = mp_obj_malloc(mp_obj_deflateio_t, type);
236 : 226 : self->stream = args_in[0];
237 : 226 : self->format = format;
238 : 226 : self->window_bits = wbits;
239 : 226 : self->read = NULL;
240 : : #if MICROPY_PY_DEFLATE_COMPRESS
241 : 226 : self->write = NULL;
242 : : #endif
243 [ + + - + ]: 226 : self->close = n_args > 3 ? mp_obj_is_true(args_in[3]) : false;
244 : :
245 : 226 : return MP_OBJ_FROM_PTR(self);
246 : : }
247 : :
248 : 378 : static mp_uint_t deflateio_read(mp_obj_t o_in, void *buf, mp_uint_t size, int *errcode) {
249 : 378 : mp_obj_deflateio_t *self = MP_OBJ_TO_PTR(o_in);
250 : :
251 [ + + + + ]: 378 : if (self->stream == MP_OBJ_NULL || !deflateio_init_read(self)) {
252 : 12 : *errcode = MP_EINVAL;
253 : 12 : return MP_STREAM_ERROR;
254 : : }
255 : :
256 [ + + ]: 364 : if (self->read->eof) {
257 : : return 0;
258 : : }
259 : :
260 : 272 : self->read->decomp.dest = buf;
261 : 272 : self->read->decomp.dest_limit = (uint8_t *)buf + size;
262 : 272 : int st = uzlib_uncompress_chksum(&self->read->decomp);
263 [ + + ]: 268 : if (st == UZLIB_DONE) {
264 : 98 : self->read->eof = true;
265 : : }
266 [ + + ]: 268 : if (st < 0) {
267 : 26 : DEBUG_printf("uncompress error=" INT_FMT "\n", st);
268 : 26 : *errcode = MP_EINVAL;
269 : 26 : return MP_STREAM_ERROR;
270 : : }
271 : 242 : return self->read->decomp.dest - (uint8_t *)buf;
272 : : }
273 : :
274 : : #if MICROPY_PY_DEFLATE_COMPRESS
275 : 124 : static mp_uint_t deflateio_write(mp_obj_t self_in, const void *buf, mp_uint_t size, int *errcode) {
276 : 124 : mp_obj_deflateio_t *self = MP_OBJ_TO_PTR(self_in);
277 : :
278 [ + + + + ]: 124 : if (self->stream == MP_OBJ_NULL || !deflateio_init_write(self)) {
279 : 8 : *errcode = MP_EINVAL;
280 : 8 : return MP_STREAM_ERROR;
281 : : }
282 : :
283 : 114 : self->write->input_len += size;
284 [ + + ]: 114 : if (self->format == DEFLATEIO_FORMAT_ZLIB) {
285 : 6 : self->write->input_checksum = uzlib_adler32(buf, size, self->write->input_checksum);
286 [ + + ]: 108 : } else if (self->format == DEFLATEIO_FORMAT_GZIP) {
287 : 4 : self->write->input_checksum = uzlib_crc32(buf, size, self->write->input_checksum);
288 : : }
289 : :
290 : 114 : uzlib_lz77_compress(&self->write->lz77, buf, size);
291 : 114 : return size;
292 : : }
293 : :
294 : 4 : static inline void put_le32(char *buf, uint32_t value) {
295 : 4 : buf[0] = value & 0xff;
296 : 4 : buf[1] = value >> 8 & 0xff;
297 : 4 : buf[2] = value >> 16 & 0xff;
298 : 4 : buf[3] = value >> 24 & 0xff;
299 : : }
300 : :
301 : 6 : static inline void put_be32(char *buf, uint32_t value) {
302 : 6 : buf[3] = value & 0xff;
303 : 6 : buf[2] = value >> 8 & 0xff;
304 : 6 : buf[1] = value >> 16 & 0xff;
305 : 6 : buf[0] = value >> 24 & 0xff;
306 : : }
307 : : #endif
308 : :
309 : 210 : static mp_uint_t deflateio_ioctl(mp_obj_t self_in, mp_uint_t request, uintptr_t arg, int *errcode) {
310 [ + - ]: 210 : if (request == MP_STREAM_CLOSE) {
311 : 210 : mp_obj_deflateio_t *self = MP_OBJ_TO_PTR(self_in);
312 : :
313 : 210 : mp_uint_t ret = 0;
314 : :
315 [ + - ]: 210 : if (self->stream != MP_OBJ_NULL) {
316 : : #if MICROPY_PY_DEFLATE_COMPRESS
317 [ + + ]: 210 : if (self->write) {
318 : 72 : uzlib_finish_block(&self->write->lz77);
319 : :
320 [ + + ]: 72 : const mp_stream_p_t *stream = mp_get_stream(self->stream);
321 : :
322 : : // Write footer if needed.
323 [ + + ]: 72 : if (self->format == DEFLATEIO_FORMAT_ZLIB || self->format == DEFLATEIO_FORMAT_GZIP) {
324 : 10 : char footer[8];
325 : 10 : size_t footer_len;
326 [ + + ]: 10 : if (self->format == DEFLATEIO_FORMAT_ZLIB) {
327 : 6 : put_be32(&footer[0], self->write->input_checksum);
328 : 6 : footer_len = 4;
329 : : } else { // DEFLATEIO_FORMAT_GZIP
330 : 4 : put_le32(&footer[0], ~self->write->input_checksum);
331 : 4 : put_le32(&footer[4], self->write->input_len);
332 : 4 : footer_len = 8;
333 : : }
334 [ + + ]: 10 : if (stream->write(self->stream, footer, footer_len, errcode) == MP_STREAM_ERROR) {
335 : 4 : ret = MP_STREAM_ERROR;
336 : : }
337 : : }
338 : : }
339 : : #endif
340 : :
341 : : // Only close the stream if required. e.g. when using io.BytesIO
342 : : // it needs to stay open so that getvalue() can be called.
343 [ + + ]: 210 : if (self->close) {
344 : 2 : mp_stream_close(self->stream);
345 : : }
346 : :
347 : : // Either way, free the reference to the stream.
348 : 210 : self->stream = MP_OBJ_NULL;
349 : : }
350 : :
351 : 210 : return ret;
352 : : } else {
353 : 0 : *errcode = MP_EINVAL;
354 : 0 : return MP_STREAM_ERROR;
355 : : }
356 : : }
357 : :
358 : : static const mp_stream_p_t deflateio_stream_p = {
359 : : .read = deflateio_read,
360 : : #if MICROPY_PY_DEFLATE_COMPRESS
361 : : .write = deflateio_write,
362 : : #endif
363 : : .ioctl = deflateio_ioctl,
364 : : };
365 : :
366 : : #if !MICROPY_ENABLE_DYNRUNTIME
367 : : static const mp_rom_map_elem_t deflateio_locals_dict_table[] = {
368 : : { MP_ROM_QSTR(MP_QSTR_read), MP_ROM_PTR(&mp_stream_read_obj) },
369 : : { MP_ROM_QSTR(MP_QSTR_readinto), MP_ROM_PTR(&mp_stream_readinto_obj) },
370 : : { MP_ROM_QSTR(MP_QSTR_readline), MP_ROM_PTR(&mp_stream_unbuffered_readline_obj) },
371 : : #if MICROPY_PY_DEFLATE_COMPRESS
372 : : { MP_ROM_QSTR(MP_QSTR_write), MP_ROM_PTR(&mp_stream_write_obj) },
373 : : #endif
374 : : { MP_ROM_QSTR(MP_QSTR_close), MP_ROM_PTR(&mp_stream_close_obj) },
375 : : { MP_ROM_QSTR(MP_QSTR___enter__), MP_ROM_PTR(&mp_identity_obj) },
376 : : { MP_ROM_QSTR(MP_QSTR___exit__), MP_ROM_PTR(&mp_stream___exit___obj) },
377 : : };
378 : : static MP_DEFINE_CONST_DICT(deflateio_locals_dict, deflateio_locals_dict_table);
379 : :
380 : : static MP_DEFINE_CONST_OBJ_TYPE(
381 : : deflateio_type,
382 : : MP_QSTR_DeflateIO,
383 : : MP_TYPE_FLAG_NONE,
384 : : make_new, deflateio_make_new,
385 : : protocol, &deflateio_stream_p,
386 : : locals_dict, &deflateio_locals_dict
387 : : );
388 : :
389 : : static const mp_rom_map_elem_t mp_module_deflate_globals_table[] = {
390 : : { MP_ROM_QSTR(MP_QSTR___name__), MP_ROM_QSTR(MP_QSTR_deflate) },
391 : : { MP_ROM_QSTR(MP_QSTR_DeflateIO), MP_ROM_PTR(&deflateio_type) },
392 : : { MP_ROM_QSTR(MP_QSTR_AUTO), MP_ROM_INT(DEFLATEIO_FORMAT_AUTO) },
393 : : { MP_ROM_QSTR(MP_QSTR_RAW), MP_ROM_INT(DEFLATEIO_FORMAT_RAW) },
394 : : { MP_ROM_QSTR(MP_QSTR_ZLIB), MP_ROM_INT(DEFLATEIO_FORMAT_ZLIB) },
395 : : { MP_ROM_QSTR(MP_QSTR_GZIP), MP_ROM_INT(DEFLATEIO_FORMAT_GZIP) },
396 : : };
397 : : static MP_DEFINE_CONST_DICT(mp_module_deflate_globals, mp_module_deflate_globals_table);
398 : :
399 : : const mp_obj_module_t mp_module_deflate = {
400 : : .base = { &mp_type_module },
401 : : .globals = (mp_obj_dict_t *)&mp_module_deflate_globals,
402 : : };
403 : :
404 : : MP_REGISTER_MODULE(MP_QSTR_deflate, mp_module_deflate);
405 : : #endif // !MICROPY_ENABLE_DYNRUNTIME
406 : :
407 : : // Source files #include'd here to make sure they're compiled in
408 : : // only if the module is enabled.
409 : :
410 : : #include "lib/uzlib/tinflate.c"
411 : : #include "lib/uzlib/header.c"
412 : : #include "lib/uzlib/adler32.c"
413 : : #include "lib/uzlib/crc32.c"
414 : :
415 : : #if MICROPY_PY_DEFLATE_COMPRESS
416 : : #include "lib/uzlib/lz77.c"
417 : : #endif
418 : :
419 : : #endif // MICROPY_PY_DEFLATE
|
