1 /************************************************************************
2 * Author    : Tiago Dionizio <tiago.dionizio@gmail.com>                 *
3 * Library   : lzlib - Lua 5 interface to access zlib library functions  *
4 *                                                                       *
5 * Permission is hereby granted, free of charge, to any person obtaining *
6 * a copy of this software and associated documentation files (the       *
7 * "Software"), to deal in the Software without restriction, including   *
8 * without limitation the rights to use, copy, modify, merge, publish,   *
9 * distribute, sublicense, and/or sell copies of the Software, and to    *
10 * permit persons to whom the Software is furnished to do so, subject to *
11 * the following conditions:                                             *
12 *                                                                       *
13 * The above copyright notice and this permission notice shall be        *
14 * included in all copies or substantial portions of the Software.       *
15 *                                                                       *
16 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,       *
17 * EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF    *
18 * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.*
19 * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY  *
20 * CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,  *
21 * TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE     *
22 * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.                *
23 ************************************************************************/
24 
25 #include <stdlib.h>
26 #include <string.h>
27 
28 #include "lua.h"
29 #include "lauxlib.h"
30 
31 #include "zlib.h"
32 
33 /*
34 ** =========================================================================
35 ** compile time options wich determine available functionality
36 ** =========================================================================
37 */
38 
39 /* TODO
40 
41 - also call flush on table/userdata when flush function is detected
42 - remove io_cb check inflate_block if condition
43 - only set eos when ZSTREAM_END is reached
44 - check for stream errors to close stream when really needed
45 
46 */
47 
48 
49 /*
50 ** =========================================================================
51 ** zlib stream metamethods
52 ** =========================================================================
53 */
54 #define ZSTREAMMETA "zlib:zstream"
55 
56 #define LZ_ANY     -1
57 #define LZ_NONE    0
58 #define LZ_DEFLATE 1
59 #define LZ_INFLATE 2
60 
61 #if 0
62     #define LZ_BUFFER_SIZE LUAL_BUFFERSIZE
63 #else
64     #define LZ_BUFFER_SIZE 8192
65 #endif
66 
67 typedef struct {
68     /* zlib structures */
69     z_stream zstream;
70     /* stream state. LZ_DEFLATE | LZ_INFLATE */
71     int state;
72     int error;
73     int peek;
74     int eos;
75     /* user callback source for reading/writing */
76     int io_cb;
77     /* input buffer */
78     int i_buffer_ref;
79     size_t i_buffer_pos;
80     size_t i_buffer_len;
81     const char *i_buffer;
82     /* output buffer */
83     size_t o_buffer_len;
84     size_t o_buffer_max;
85     char o_buffer[LZ_BUFFER_SIZE];
86     /* dictionary */
87     const Bytef *dictionary;
88     size_t dictionary_len;
89 } lz_stream;
90 
91 
92 /* forward declarations */
93 static int lzstream_docompress(lua_State *L, lz_stream *s, int from, int to, int flush);
94 
95 
lzstream_new(lua_State * L,int src)96 static lz_stream *lzstream_new(lua_State *L, int src) {
97     lz_stream *s = (lz_stream*)lua_newuserdata(L, sizeof(lz_stream));
98 
99     luaL_getmetatable(L, ZSTREAMMETA);
100     lua_setmetatable(L, -2);        /* set metatable */
101 
102     s->state = LZ_NONE;
103     s->error = Z_OK;
104     s->eos = 0;
105     s->io_cb = LUA_REFNIL;
106 
107     s->i_buffer = NULL;
108     s->i_buffer_ref = LUA_REFNIL;
109     s->i_buffer_pos = 0;
110     s->i_buffer_len = 0;
111 
112     s->peek = 0;
113     s->o_buffer_len = 0;
114     s->o_buffer_max = sizeof(s->o_buffer) / sizeof(s->o_buffer[0]);
115 
116     s->zstream.zalloc = Z_NULL;
117     s->zstream.zfree = Z_NULL;
118 
119     /* prepare source */
120     if (lua_isstring(L, src)) {
121         lua_pushvalue(L, src);
122         s->i_buffer_ref = luaL_ref(L, LUA_REGISTRYINDEX);
123         s->i_buffer = lua_tolstring(L, src, &s->i_buffer_len);
124     } else {
125         /* table | function | userdata */
126         lua_pushvalue(L, src);
127         s->io_cb = luaL_ref(L, LUA_REGISTRYINDEX);
128     }
129     return s;
130 }
131 
lzstream_cleanup(lua_State * L,lz_stream * s)132 static void lzstream_cleanup(lua_State *L, lz_stream *s) {
133     if (s && s->state != LZ_NONE) {
134         if (s->state == LZ_INFLATE) {
135             inflateEnd(&s->zstream);
136         }
137         if (s->state == LZ_DEFLATE) {
138             deflateEnd(&s->zstream);
139         }
140 
141         luaL_unref(L, LUA_REGISTRYINDEX, s->io_cb);
142         luaL_unref(L, LUA_REGISTRYINDEX, s->i_buffer_ref);
143         s->state = LZ_NONE;
144     }
145 }
146 
147 /* ====================================================================== */
148 
lzstream_get(lua_State * L,int index)149 static lz_stream *lzstream_get(lua_State *L, int index) {
150     lz_stream *s = (lz_stream*)luaL_checkudata(L, index, ZSTREAMMETA);
151     if (s == NULL) luaL_argerror(L, index, "bad zlib stream");
152     return s;
153 }
154 
lzstream_check(lua_State * L,int index,int state)155 static lz_stream *lzstream_check(lua_State *L, int index, int state) {
156     lz_stream *s = lzstream_get(L, index);
157     if ((state != LZ_ANY && s->state != state) || s->state == LZ_NONE) {
158         luaL_argerror(L, index, "attempt to use invalid zlib stream");
159     }
160     return s;
161 }
162 
163 /* ====================================================================== */
164 
lzstream_tostring(lua_State * L)165 static int lzstream_tostring(lua_State *L) {
166     lz_stream *s = (lz_stream*)luaL_checkudata(L, 1, ZSTREAMMETA);
167     if (s == NULL) luaL_argerror(L, 1, "bad zlib stream");
168 
169     if (s->state == LZ_NONE) {
170         lua_pushstring(L, "zlib stream (closed)");
171     } else if (s->state == LZ_DEFLATE) {
172         lua_pushfstring(L, "zlib deflate stream (%p)", (void*)s);
173     } else if (s->state == LZ_INFLATE) {
174         lua_pushfstring(L, "zlib inflate stream (%p)", (void*)s);
175     } else {
176         lua_pushfstring(L, "%p", (void*)s);
177     }
178 
179     return 1;
180 }
181 
182 /* ====================================================================== */
183 
lzstream_gc(lua_State * L)184 static int lzstream_gc(lua_State *L) {
185     lz_stream *s = lzstream_get(L, 1);
186     lzstream_cleanup(L, s);
187     return 0;
188 }
189 
190 /* ====================================================================== */
191 
lzstream_close(lua_State * L)192 static int lzstream_close(lua_State *L) {
193     lz_stream *s = lzstream_get(L, 1);
194 
195     if (s->state == LZ_DEFLATE) {
196         lua_settop(L, 0);
197         lua_pushliteral(L, "");
198         return lzstream_docompress(L, s, 1, 1, Z_FINISH);
199     }
200 
201     lzstream_cleanup(L, s);
202     lua_pushboolean(L, 1);
203     return 1;
204 }
205 
206 /* ====================================================================== */
207 
lzstream_adler(lua_State * L)208 static int lzstream_adler(lua_State *L) {
209     lz_stream *s = lzstream_check(L, 1, LZ_ANY);
210     lua_pushnumber(L, s->zstream.adler);
211     return 1;
212 }
213 
214 /* ====================================================================== */
215 
216 /*
217     zlib.deflate(
218         sink: function | { write: function [, close: function, flush: function] },
219         compression level, [Z_DEFAILT_COMPRESSION]
220         method, [Z_DEFLATED]
221         windowBits, [15]
222         memLevel, [8]
223         strategy, [Z_DEFAULT_STRATEGY]
224         dictionary: [""]
225     )
226 */
lzlib_deflate(lua_State * L)227 static int lzlib_deflate(lua_State *L) {
228     int level, method, windowBits, memLevel, strategy;
229     lz_stream *s;
230     const char *dictionary;
231     size_t dictionary_len;
232 
233     if (lua_istable(L, 1) || lua_isuserdata(L, 1)) {
234         /* is there a :write function? */
235         lua_getfield(L, 1, "write");
236         if (!lua_isfunction(L, -1)) {
237             luaL_argerror(L, 1, "output parameter does not provide :write function");
238         }
239         lua_pop(L, 1);
240     }
241     else if (!lua_isfunction(L, 1)) {
242         luaL_argerror(L, 1, "output parameter must be a function, table or userdata value");
243     }
244 
245     level = (int) luaL_optinteger(L, 2, Z_DEFAULT_COMPRESSION);
246     method = (int) luaL_optinteger(L, 3, Z_DEFLATED);
247     windowBits = (int) luaL_optinteger(L, 4, 15);
248     memLevel = (int) luaL_optinteger(L, 5, 8);
249     strategy = (int) luaL_optinteger(L, 6, Z_DEFAULT_STRATEGY);
250     dictionary = luaL_optlstring(L, 7, NULL, &dictionary_len);
251 
252     s = lzstream_new(L, 1);
253 
254     if (deflateInit2(&s->zstream, level, method, windowBits, memLevel, strategy) != Z_OK) {
255         lua_pushliteral(L, "call to deflateInit2 failed");
256         lua_error(L);
257     }
258 
259     if (dictionary) {
260         if (deflateSetDictionary(&s->zstream, (const Bytef *) dictionary, dictionary_len) != Z_OK) {
261             lua_pushliteral(L, "call to deflateSetDictionnary failed");
262             lua_error(L);
263         }
264     }
265 
266     s->state = LZ_DEFLATE;
267     return 1;
268 }
269 
270 /*
271     zlib.inflate(
272         source: string | function | { read: function, close: function },
273         windowBits: number, [15]
274         dictionary: [""]
275     )
276 */
lzlib_inflate(lua_State * L)277 static int lzlib_inflate(lua_State *L)
278 {
279     int windowBits;
280     lz_stream *s;
281     int have_peek = 0;
282     const char *dictionary;
283     size_t dictionary_len;
284 
285     if (lua_istable(L, 1) || lua_isuserdata(L, 1)) {
286         /* is there a :read function? */
287         lua_getfield(L, 1, "read");
288         if (!lua_isfunction(L, -1)) {
289             luaL_argerror(L, 1, "input parameter does not provide :read function");
290         }
291         lua_pop(L, 1);
292         /* check for peek function */
293         lua_getfield(L, 1, "peek");
294         have_peek = lua_isfunction(L, -1);
295         lua_pop(L, 1);
296     }
297     else if (!lua_isstring(L, 1) && !lua_isfunction(L, 1)) {
298         luaL_argerror(L, 1, "input parameter must be a string, function, table or userdata value");
299     }
300 
301     windowBits = (int) luaL_optinteger(L, 2, 15);
302     dictionary = luaL_optlstring(L, 3, NULL, &dictionary_len);
303 
304     s = lzstream_new(L, 1);
305 
306     if (windowBits > 0 && windowBits < 16) {
307         windowBits |= 32;
308     }
309 
310     if (inflateInit2(&s->zstream, windowBits) != Z_OK) {
311         lua_pushliteral(L, "call to inflateInit2 failed");
312         lua_error(L);
313     }
314 
315     if (dictionary) {
316         s->dictionary = (const Bytef *) dictionary;
317         s->dictionary_len = dictionary_len;
318     }
319 
320     s->peek = have_peek;
321     s->state = LZ_INFLATE;
322     return 1;
323 }
324 
325 /* ====================================================================== */
326 
lz_pushresult(lua_State * L,lz_stream * s)327 static int lz_pushresult (lua_State *L, lz_stream *s) {
328     if (s->error == Z_OK) {
329         lua_pushboolean(L, 1);
330         return 1;
331     } else {
332         lua_pushnil(L);
333         lua_pushstring(L, zError(s->error));
334         lua_pushinteger(L, s->error);
335         return 3;
336     }
337 }
338 
339 /*
340     Get block to process:
341         - top of stack gets
342 */
lzstream_fetch_block(lua_State * L,lz_stream * s,int hint)343 static const char* lzstream_fetch_block(lua_State *L, lz_stream *s, int hint) {
344     if (s->i_buffer_pos >= s->i_buffer_len) {
345         luaL_unref(L, LUA_REGISTRYINDEX, s->i_buffer_ref);
346         s->i_buffer_ref = LUA_NOREF;
347         s->i_buffer = NULL;
348 
349         lua_rawgeti(L, LUA_REGISTRYINDEX, s->io_cb);
350         if (!lua_isnil(L, -1)) {
351             if (lua_isfunction(L, -1)) {
352                 lua_pushinteger(L, hint);
353                 lua_call(L, 1, 1);
354             } else {
355                 lua_getfield(L, -1, (s->peek ? "peek" : "read"));
356                 lua_insert(L, -2);
357                 lua_pushinteger(L, hint);
358                 lua_call(L, 2, 1);
359             }
360 
361             if (lua_isstring(L, -1)) {
362                 s->i_buffer_pos = 0;
363                 s->i_buffer = lua_tolstring(L, -1, &s->i_buffer_len);
364                 if (s->i_buffer_len > 0) {
365                     s->i_buffer_ref = luaL_ref(L, LUA_REGISTRYINDEX);
366                 } else {
367                     lua_pop(L, 1);
368                 }
369             } else if (lua_isnil(L, -1)) {
370                 lua_pop(L, 1);
371             } else {
372                 lua_pushliteral(L, "deflate callback must return string or nil");
373                 lua_error(L);
374             }
375         } else {
376             lua_pop(L, 1);
377         }
378     }
379 
380     return s->i_buffer;
381 }
382 
lzstream_inflate_block(lua_State * L,lz_stream * s)383 static int lzstream_inflate_block(lua_State *L, lz_stream *s) {
384     if (lzstream_fetch_block(L, s, LZ_BUFFER_SIZE) || !s->eos) {
385         int r;
386 
387         if (s->i_buffer_len == s->i_buffer_pos) {
388             s->zstream.next_in = NULL;
389             s->zstream.avail_in = 0;
390         } else {
391             s->zstream.next_in = (unsigned char*)(s->i_buffer + s->i_buffer_pos);
392             s->zstream.avail_in = s->i_buffer_len - s->i_buffer_pos;
393         }
394 
395         s->zstream.next_out = (unsigned char*)s->o_buffer + s->o_buffer_len;
396         s->zstream.avail_out = s->o_buffer_max - s->o_buffer_len;
397 
398         /* munch some more */
399         r = inflate(&s->zstream, Z_SYNC_FLUSH);
400 
401         if (r == Z_NEED_DICT) {
402             if (s->dictionary == NULL) {
403                 lua_pushliteral(L, "no inflate dictionary provided");
404                 lua_error(L);
405             }
406 
407             if (inflateSetDictionary(&s->zstream, s->dictionary, s->dictionary_len) != Z_OK) {
408                 lua_pushliteral(L, "call to inflateSetDictionnary failed");
409                 lua_error(L);
410             }
411 
412             r = inflate(&s->zstream, Z_SYNC_FLUSH);
413         }
414 
415         if (r != Z_OK && r != Z_STREAM_END && r != Z_BUF_ERROR) {
416             lzstream_cleanup(L, s);
417             s->error = r;
418             #if 1
419             lua_pushfstring(L, "failed to decompress [%d]", r);
420             lua_error(L);
421             #endif
422         }
423 
424         if (r == Z_STREAM_END) {
425             luaL_unref(L, LUA_REGISTRYINDEX, s->i_buffer_ref);
426             s->i_buffer_ref = LUA_NOREF;
427             s->i_buffer = NULL;
428 
429             s->eos = 1;
430         }
431 
432         /* number of processed bytes */
433         if (s->peek) {
434             size_t processed = s->i_buffer_len - s->i_buffer_pos - s->zstream.avail_in;
435 
436             lua_rawgeti(L, LUA_REGISTRYINDEX, s->io_cb);
437             lua_getfield(L, -1, "read");
438             lua_insert(L, -2);
439             lua_pushinteger(L, processed);
440             lua_call(L, 2, 0);
441         }
442 
443         s->i_buffer_pos = s->i_buffer_len - s->zstream.avail_in;
444         s->o_buffer_len = s->o_buffer_max - s->zstream.avail_out;
445     }
446 
447     return s->o_buffer_len;
448 }
449 
450 /*
451 ** Remove n bytes from the output buffer.
452 */
lzstream_remove(lz_stream * s,size_t n)453 static void lzstream_remove(lz_stream *s, size_t n) {
454     memmove(s->o_buffer, s->o_buffer + n, s->o_buffer_len - n);
455     s->o_buffer_len -= n;
456 }
457 
458 /*
459 ** Copy at most n bytes to buffer b and remove them from the
460 ** output stream buffer.
461 */
lzstream_flush_buffer(lua_State * L,lz_stream * s,size_t n,luaL_Buffer * b)462 static int lzstream_flush_buffer(lua_State *L, lz_stream *s, size_t n, luaL_Buffer *b) {
463     /* check output */
464     if (n > s->o_buffer_len) {
465         n = s->o_buffer_len;
466     }
467 
468     if (n > 0) {
469         lua_pushlstring(L, s->o_buffer, n);
470         luaL_addvalue(b);
471 
472         lzstream_remove(s, n);
473     }
474 
475     return n;
476 }
477 
478 /*
479     z:read(
480         {number | '*l' | '*a'}*
481     )
482 */
lz_test_eof(lua_State * L,lz_stream * s)483 static int lz_test_eof(lua_State *L, lz_stream *s) {
484     lua_pushlstring(L, NULL, 0);
485     if (s->o_buffer_len > 0) {
486         return 1;
487     } else if (s->eos) {
488         return 0;
489     } else {
490         return lzstream_inflate_block(L, s);
491     }
492 }
493 
lz_read_line(lua_State * L,lz_stream * s)494 static int lz_read_line(lua_State *L, lz_stream *s) {
495     luaL_Buffer b;
496     size_t l = 0, n;
497 
498     luaL_buffinit(L, &b);
499 
500     if (s->o_buffer_len > 0 || !s->eos) do {
501         char *p = s->o_buffer;
502         size_t len = s->o_buffer_len;
503 
504         /* find newline in output buffer */
505         for (n = 0; n < len; ++n, ++p) {
506             if (*p == '\n' || *p == '\r') {
507                 int eat_nl = *p == '\r';
508                 luaL_addlstring(&b, s->o_buffer, n);
509                 lzstream_remove(s, n+1);
510                 l += n;
511 
512                 if (eat_nl && lzstream_inflate_block(L, s)) {
513                     if (s->o_buffer_len > 0 && *s->o_buffer == '\n') {
514                         lzstream_remove(s, 1);
515                     }
516                 }
517 
518                 luaL_pushresult(&b);
519                 return 1;
520             }
521         }
522 
523         if (len > 0) {
524             luaL_addlstring(&b, s->o_buffer, len);
525             lzstream_remove(s, len);
526             l += len;
527         }
528     } while (lzstream_inflate_block(L, s));
529 
530     luaL_pushresult(&b);
531     return l > 0 || !s->eos || s->o_buffer_len > 0;
532 }
533 
534 
lz_read_chars(lua_State * L,lz_stream * s,size_t n)535 static int lz_read_chars(lua_State *L, lz_stream *s, size_t n) {
536     size_t len;
537     luaL_Buffer b;
538     luaL_buffinit(L, &b);
539 
540     if (s->o_buffer_len > 0 || !s->eos) do {
541         size_t rlen = lzstream_flush_buffer(L, s, n, &b);
542         n -= rlen;
543     } while (n > 0 && lzstream_inflate_block(L, s));
544 
545     luaL_pushresult(&b);
546     lua_tolstring(L, -1, &len);
547     return n == 0 || len > 0;
548 }
549 
lzstream_decompress(lua_State * L)550 static int lzstream_decompress(lua_State *L) {
551     lz_stream *s = lzstream_check(L, 1, LZ_INFLATE);
552     int nargs = lua_gettop(L) - 1;
553     int success;
554     int n;
555     if (nargs == 0) {  /* no arguments? */
556         success = lz_read_line(L, s);
557         n = 3;  /* to return 1 result */
558     }
559     else {  /* ensure stack space for all results and for auxlib's buffer */
560         luaL_checkstack(L, nargs+LUA_MINSTACK, "too many arguments");
561         success = 1;
562         for (n = 2; nargs-- && success; n++) {
563             if (lua_type(L, n) == LUA_TNUMBER) {
564                 size_t l = (size_t)lua_tointeger(L, n);
565                 success = (l == 0) ? lz_test_eof(L, s) : lz_read_chars(L, s, l);
566             }
567             else {
568                 const char *p = lua_tostring(L, n);
569                 luaL_argcheck(L, p && p[0] == '*', n, "invalid option");
570                 switch (p[1]) {
571                     case 'l':  /* line */
572                         success = lz_read_line(L, s);
573                         break;
574                     case 'a':  /* file */
575                         lz_read_chars(L, s, ~((size_t)0));  /* read MAX_SIZE_T chars */
576                         success = 1; /* always success */
577                         break;
578                     default:
579                         return luaL_argerror(L, n, "invalid format");
580                 }
581             }
582         }
583     }
584     if (s->error != Z_OK) {
585         return lz_pushresult(L, s);
586     }
587     if (!success) {
588         lua_pop(L, 1);  /* remove last result */
589         lua_pushnil(L);  /* push nil instead */
590     }
591     return n - 2;
592 }
593 
594 
lzstream_readline(lua_State * L)595 static int lzstream_readline(lua_State *L) {
596     lz_stream *s;
597     int sucess;
598 
599     s = lzstream_check(L, lua_upvalueindex(1), LZ_INFLATE);
600     sucess = lz_read_line(L, s);
601 
602     if (s->error != Z_OK) {
603         return lz_pushresult(L, s);
604     }
605 
606     if (sucess) {
607         return 1;
608     } else {
609         /* EOF */
610         return 0;
611     }
612 }
613 
lzstream_lines(lua_State * L)614 static int lzstream_lines(lua_State *L) {
615     lzstream_check(L, 1, LZ_INFLATE);
616     lua_settop(L, 1);
617     lua_pushcclosure(L, lzstream_readline, 1);
618     return 1;
619 }
620 
621 /* ====================================================================== */
622 
lzstream_docompress(lua_State * L,lz_stream * s,int from,int to,int flush)623 static int lzstream_docompress(lua_State *L, lz_stream *s, int from, int to, int flush) {
624     int r, arg;
625     int self = 0;
626     size_t b_size = s->o_buffer_max;
627     unsigned char *b = (unsigned char *)s->o_buffer;
628 
629     /* number of processed bytes */
630     lua_rawgeti(L, LUA_REGISTRYINDEX, s->io_cb);
631     if (!lua_isfunction(L, -1)) {
632         self = 1;
633         lua_getfield(L, -1, "write");
634     }
635 
636     for (arg = from; arg <= to; arg++) {
637         s->zstream.next_in = (unsigned char*)luaL_checklstring(L, arg, (size_t*)&s->zstream.avail_in);
638 
639         do {
640             s->zstream.next_out = b;
641             s->zstream.avail_out = b_size;
642 
643             /* bake some more */
644             r = deflate(&s->zstream, flush);
645             if (r != Z_OK && r != Z_STREAM_END && r != Z_BUF_ERROR) {
646                 lzstream_cleanup(L, s);
647                 lua_pushboolean(L, 0);
648                 lua_pushfstring(L, "failed to compress [%d]", r);
649                 return 2;
650             }
651 
652             if (s->zstream.avail_out != b_size) {
653                 /* write output */
654                 lua_pushvalue(L, -1); /* function */
655                 if (self) lua_pushvalue(L, -3); /* self */
656                 lua_pushlstring(L, (char*)b, b_size - s->zstream.avail_out); /* data */
657                 lua_call(L, (self ? 2 : 1), 0);
658             }
659 
660             if (r == Z_STREAM_END) {
661                 lzstream_cleanup(L, s);
662                 break;
663             }
664 
665             /* process all input */
666         } while (s->zstream.avail_in > 0 || s->zstream.avail_out == 0);
667     }
668 
669     lua_pushboolean(L, 1);
670     return 1;
671 }
672 
lzstream_compress(lua_State * L)673 static int lzstream_compress(lua_State *L) {
674     lz_stream *s = lzstream_check(L, 1, LZ_DEFLATE);
675     return lzstream_docompress(L, s, 2, lua_gettop(L), Z_NO_FLUSH);
676 }
677 
678 
679 /* ====================================================================== */
680 
lzstream_flush(lua_State * L)681 static int lzstream_flush(lua_State *L) {
682     static int flush_values[] = { Z_SYNC_FLUSH, Z_FULL_FLUSH, Z_FINISH };
683     static const char *const flush_opts[] = { "sync", "full", "finish" };
684 
685     lz_stream *s = lzstream_check(L, 1, LZ_DEFLATE);
686     int flush = luaL_checkoption(L, 2, flush_opts[0], flush_opts);
687 
688     lua_settop(L, 0);
689     lua_pushliteral(L, "");
690     return lzstream_docompress(L, s, 1, 1, flush_values[flush]);
691 }
692 
693 /*
694 ** =========================================================================
695 ** zlib functions
696 ** =========================================================================
697 */
698 
lzlib_version(lua_State * L)699 static int lzlib_version(lua_State *L)
700 {
701     lua_pushstring(L, zlibVersion());
702     return 1;
703 }
704 
705 /* ====================================================================== */
lzlib_adler32(lua_State * L)706 static int lzlib_adler32(lua_State *L)
707 {
708     if (lua_gettop(L) == 0)
709     {
710         /* adler32 initial value */
711         lua_pushnumber(L, adler32(0L, Z_NULL, 0));
712     }
713     else
714     {
715         /* update adler32 checksum */
716         size_t len;
717         int adler = (int) luaL_checkinteger(L, 1);
718         const unsigned char* buf = (unsigned char*)luaL_checklstring(L, 2, &len);
719 
720         lua_pushnumber(L, adler32(adler, buf, len));
721     }
722     return 1;
723 }
724 
725 /* ====================================================================== */
lzlib_crc32(lua_State * L)726 static int lzlib_crc32(lua_State *L)
727 {
728     if (lua_gettop(L) == 0)
729     {
730         /* crc32 initial value */
731         lua_pushnumber(L, crc32(0L, Z_NULL, 0));
732     }
733     else
734     {
735         /* update crc32 checksum */
736         size_t len;
737         int crc = (int) luaL_checkinteger(L, 1);
738         const unsigned char* buf = (unsigned char*)luaL_checklstring(L, 2, &len);
739 
740         lua_pushnumber(L, crc32(crc, buf, len));
741     }
742     return 1;
743 }
744 
745 /* ====================================================================== */
746 
747 
lzlib_compress(lua_State * L)748 static int lzlib_compress(lua_State *L) {
749     size_t avail_in;
750     const char *next_in = luaL_checklstring(L, 1, &avail_in);
751     int level = (int) luaL_optinteger(L, 2, Z_DEFAULT_COMPRESSION);
752     int method = (int) luaL_optinteger(L, 3, Z_DEFLATED);
753     int windowBits = (int) luaL_optinteger(L, 4, 15);
754     int memLevel = (int) luaL_optinteger(L, 5, 8);
755     int strategy = (int) luaL_optinteger(L, 6, Z_DEFAULT_STRATEGY);
756 
757     int ret;
758     luaL_Buffer b;
759     z_stream zs;
760 
761     luaL_buffinit(L, &b);
762 
763     zs.zalloc = Z_NULL;
764     zs.zfree = Z_NULL;
765 
766     zs.next_out = Z_NULL;
767     zs.avail_out = 0;
768     zs.next_in = Z_NULL;
769     zs.avail_in = 0;
770 
771     ret = deflateInit2(&zs, level, method, windowBits, memLevel, strategy);
772 
773     if (ret != Z_OK)
774     {
775         lua_pushnil(L);
776         lua_pushnumber(L, ret);
777         return 2;
778     }
779 
780     zs.next_in = (unsigned char*)next_in;
781     zs.avail_in = avail_in;
782 
783     for(;;)
784     {
785         zs.next_out = (unsigned char*)luaL_prepbuffer(&b);
786         zs.avail_out = LUAL_BUFFERSIZE;
787 
788         /* munch some more */
789         ret = deflate(&zs, Z_FINISH);
790 
791         /* push gathered data */
792         luaL_addsize(&b, LUAL_BUFFERSIZE - zs.avail_out);
793 
794         /* done processing? */
795         if (ret == Z_STREAM_END)
796             break;
797 
798         /* error condition? */
799         if (ret != Z_OK)
800             break;
801     }
802 
803     /* cleanup */
804     deflateEnd(&zs);
805 
806     luaL_pushresult(&b);
807     lua_pushnumber(L, ret);
808     return 2;
809 }
810 
811 /* ====================================================================== */
812 
lzlib_decompress(lua_State * L)813 static int lzlib_decompress(lua_State *L)
814 {
815     size_t avail_in;
816     const char *next_in = luaL_checklstring(L, 1, &avail_in);
817     int windowBits = (int) luaL_optinteger(L, 2, 15);
818 
819     int ret;
820     luaL_Buffer b;
821     z_stream zs;
822 
823     luaL_buffinit(L, &b);
824 
825     zs.zalloc = Z_NULL;
826     zs.zfree = Z_NULL;
827 
828     zs.next_out = Z_NULL;
829     zs.avail_out = 0;
830     zs.next_in = Z_NULL;
831     zs.avail_in = 0;
832 
833     ret = inflateInit2(&zs, windowBits);
834 
835     if (ret != Z_OK) {
836         lua_pushliteral(L, "failed to initialize zstream structures");
837         lua_error(L);
838     }
839 
840     zs.next_in = (unsigned char*)next_in;
841     zs.avail_in = avail_in;
842 
843     for (;;) {
844         zs.next_out = (unsigned char*)luaL_prepbuffer(&b);
845         zs.avail_out = LUAL_BUFFERSIZE;
846 
847         /* bake some more */
848         ret = inflate(&zs, Z_FINISH);
849 
850         /* push gathered data */
851         luaL_addsize(&b, LUAL_BUFFERSIZE - zs.avail_out);
852 
853         /* done processing? */
854         if (ret == Z_STREAM_END)
855             break;
856 
857         if (ret != Z_OK && ret != Z_BUF_ERROR) {
858             /* cleanup */
859             inflateEnd(&zs);
860 
861             lua_pushliteral(L, "failed to process zlib stream");
862             lua_error(L);
863         }
864     }
865 
866     /* cleanup */
867     inflateEnd(&zs);
868 
869     luaL_pushresult(&b);
870     return 1;
871 }
872 
873 
874 /*
875 ** =========================================================================
876 ** Register functions
877 ** =========================================================================
878 */
879 
880 #if (LUA_VERSION_NUM >= 502)
881 
882 #define luaL_register(L,n,f)	luaL_setfuncs(L,f,0)
883 
884 #endif
885 
luaopen_zlib(lua_State * L)886 LUALIB_API int luaopen_zlib(lua_State *L)
887 {
888     const luaL_Reg lzstream_meta[] =
889     {
890         {"write",           lzstream_compress   },
891         {"read",            lzstream_decompress },
892         {"lines",           lzstream_lines      },
893         {"flush",           lzstream_flush      },
894         {"close",           lzstream_close      },
895 
896         {"adler",           lzstream_adler      },
897 
898         {"__tostring",      lzstream_tostring   },
899         {"__gc",            lzstream_gc         },
900         {NULL, NULL}
901     };
902 
903     const luaL_Reg zlib[] =
904     {
905         {"version",         lzlib_version       },
906         {"adler32",         lzlib_adler32       },
907         {"crc32",           lzlib_crc32         },
908 
909         {"deflate",         lzlib_deflate       },
910         {"inflate",         lzlib_inflate       },
911 
912         {"compress",        lzlib_compress      },
913         {"decompress",      lzlib_decompress    },
914 
915         {NULL, NULL}
916     };
917 
918     /* ====================================================================== */
919 
920     /* create new metatable for zlib compression structures */
921     luaL_newmetatable(L, ZSTREAMMETA);
922     lua_pushliteral(L, "__index");
923     lua_pushvalue(L, -2);               /* push metatable */
924     lua_rawset(L, -3);                  /* metatable.__index = metatable */
925 
926     /*
927     ** Stack: metatable
928     */
929     luaL_register(L, NULL, lzstream_meta);
930 
931     lua_pop(L, 1);                      /* remove metatable from stack */
932 
933     /*
934     ** Stack:
935     */
936     lua_newtable(L);
937 
938     lua_pushliteral (L, "_COPYRIGHT");
939     lua_pushliteral (L, "Copyright (C) 2003-2010 Tiago Dionizio");
940     lua_settable (L, -3);
941     lua_pushliteral (L, "_DESCRIPTION");
942     lua_pushliteral (L, "Lua 5 interface to access zlib library functions");
943     lua_settable (L, -3);
944     lua_pushliteral (L, "_VERSION");
945     lua_pushliteral (L, "lzlib 0.4-work3");
946     lua_settable (L, -3);
947 
948 #define PUSH_LITERAL(name) \
949     lua_pushliteral (L, #name); \
950     lua_pushinteger (L, Z_##name); \
951     lua_settable (L, -3);
952 
953 #define PUSH_NUMBER(name, value) \
954     lua_pushliteral (L, #name); \
955     lua_pushinteger (L, value); \
956     lua_settable (L, -3);
957 
958     PUSH_LITERAL(NO_COMPRESSION)
959     PUSH_LITERAL(BEST_SPEED)
960     PUSH_LITERAL(BEST_COMPRESSION)
961     PUSH_LITERAL(DEFAULT_COMPRESSION)
962 
963     PUSH_LITERAL(FILTERED)
964     PUSH_LITERAL(HUFFMAN_ONLY)
965     PUSH_LITERAL(RLE)
966     PUSH_LITERAL(FIXED)
967     PUSH_LITERAL(DEFAULT_STRATEGY)
968 
969     PUSH_NUMBER(MINIMUM_MEMLEVEL, 1)
970     PUSH_NUMBER(MAXIMUM_MEMLEVEL, 9)
971     PUSH_NUMBER(DEFAULT_MEMLEVEL, 8)
972 
973     PUSH_NUMBER(DEFAULT_WINDOWBITS, 15)
974     PUSH_NUMBER(MINIMUM_WINDOWBITS, 8)
975     PUSH_NUMBER(MAXIMUM_WINDOWBITS, 15)
976 
977     PUSH_NUMBER(GZIP_WINDOWBITS, 16)
978     PUSH_NUMBER(RAW_WINDOWBITS, -1)
979 
980     luaL_register(L, NULL, zlib);
981 
982     /*
983     ** Stack: zlib table
984     */
985     return 1;
986 }
987