1 /*
2  * %CopyrightBegin%
3  *
4  * Copyright Ericsson 2017. All Rights Reserved.
5  *
6  * Licensed under the Apache License, Version 2.0 (the "License");
7  * you may not use this file except in compliance with the License.
8  * You may obtain a copy of the License at
9  *
10  *     http://www.apache.org/licenses/LICENSE-2.0
11  *
12  * Unless required by applicable law or agreed to in writing, software
13  * distributed under the License is distributed on an "AS IS" BASIS,
14  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15  * See the License for the specific language governing permissions and
16  * limitations under the License.
17  *
18  * %CopyrightEnd%
19  */
20 
21 #define STATIC_ERLANG_NIF 1
22 
23 #include <stdio.h>
24 #include <zlib.h>
25 
26 #include "erl_nif.h"
27 #include "config.h"
28 #include "sys.h"
29 
30 #ifdef VALGRIND
31 #  include <valgrind/memcheck.h>
32 #endif
33 
34 #define INFL_DICT_SZ    (32768)
35 
36 /* NIF interface declarations */
37 static int load(ErlNifEnv *env, void** priv_data, ERL_NIF_TERM load_info);
38 static int upgrade(ErlNifEnv *env, void** priv_data, void** old_priv_data, ERL_NIF_TERM load_info);
39 static void unload(ErlNifEnv *env, void* priv_data);
40 
41 static ErlNifResourceType *rtype_zlib;
42 
43 static ERL_NIF_TERM am_not_on_controlling_process;
44 
45 static ERL_NIF_TERM am_not_initialized;
46 static ERL_NIF_TERM am_already_initialized;
47 
48 static ERL_NIF_TERM am_ok;
49 static ERL_NIF_TERM am_error;
50 
51 static ERL_NIF_TERM am_continue;
52 static ERL_NIF_TERM am_finished;
53 
54 static ERL_NIF_TERM am_not_supported;
55 static ERL_NIF_TERM am_need_dictionary;
56 
57 static ERL_NIF_TERM am_empty;
58 
59 static ERL_NIF_TERM am_stream_end;
60 static ERL_NIF_TERM am_stream_error;
61 static ERL_NIF_TERM am_data_error;
62 static ERL_NIF_TERM am_mem_error;
63 static ERL_NIF_TERM am_buf_error;
64 static ERL_NIF_TERM am_version_error;
65 static ERL_NIF_TERM am_unknown_error;
66 
67 typedef enum {
68     ST_NONE    = 0,
69     ST_DEFLATE = 1,
70     ST_INFLATE = 2,
71     ST_CLOSED = 3
72 } zlib_state_t;
73 
74 /* Controls what to do when the user attempts to decompress more data after
75  * Z_STREAM_END has been returned:
76  *
77  * - 'cut' wipes all further input and returns empty results until reset by
78  * the user. This is the default behavior, matching that of the old driver.
79  * - 'reset' resets the state without discarding any input, making it possible
80  * to decompress blindly concatenated streams.
81  * - 'error' crashes with a data error. */
82 typedef enum {
83     EOS_BEHAVIOR_ERROR = 0,
84     EOS_BEHAVIOR_RESET = 1,
85     EOS_BEHAVIOR_CUT = 2
86 } zlib_eos_behavior_t;
87 
88 typedef struct {
89     z_stream s;
90     zlib_state_t state;
91 
92     zlib_eos_behavior_t eos_behavior;
93 
94     /* These refer to the plaintext CRC, and are only needed for zlib:crc32/1
95      * which is deprecated. */
96     uLong input_crc;
97     uLong output_crc;
98     int want_input_crc;
99     int want_output_crc;
100 
101     int is_raw_stream;
102 
103     int eos_seen;
104 
105     /* DEPRECATED */
106     int inflateChunk_buffer_size;
107 
108     ErlNifPid controlling_process;
109     ErlNifMutex *controller_lock;
110 
111     ErlNifIOQueue *input_queue;
112 
113     ErlNifEnv *stash_env;
114     ERL_NIF_TERM stash_term;
115 } zlib_data_t;
116 
117 /* The NIFs: */
118 
119 static ERL_NIF_TERM zlib_open(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]);
120 static ERL_NIF_TERM zlib_close(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]);
121 static ERL_NIF_TERM zlib_set_controller(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]);
122 
123 static ERL_NIF_TERM zlib_deflateInit(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]);
124 static ERL_NIF_TERM zlib_deflateSetDictionary(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]);
125 static ERL_NIF_TERM zlib_deflateReset(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]);
126 static ERL_NIF_TERM zlib_deflateEnd(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]);
127 static ERL_NIF_TERM zlib_deflateParams(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]);
128 static ERL_NIF_TERM zlib_deflate(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]);
129 
130 static ERL_NIF_TERM zlib_inflateInit(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]);
131 static ERL_NIF_TERM zlib_inflateSetDictionary(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]);
132 static ERL_NIF_TERM zlib_inflateGetDictionary(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]);
133 static ERL_NIF_TERM zlib_inflateReset(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]);
134 static ERL_NIF_TERM zlib_inflateEnd(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]);
135 static ERL_NIF_TERM zlib_inflate(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]);
136 
137 static ERL_NIF_TERM zlib_crc32(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]);
138 
139 static ERL_NIF_TERM zlib_clearStash(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]);
140 static ERL_NIF_TERM zlib_setStash(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]);
141 static ERL_NIF_TERM zlib_getStash(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]);
142 
143 static ERL_NIF_TERM zlib_getBufSize(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]);
144 static ERL_NIF_TERM zlib_setBufSize(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]);
145 
146 static ERL_NIF_TERM zlib_enqueue_input(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]);
147 
148 static ErlNifFunc nif_funcs[] = {
149     {"close_nif", 1, zlib_close},
150     {"open_nif", 0, zlib_open},
151 
152     {"set_controller_nif", 2, zlib_set_controller},
153 
154     /* deflate */
155     {"deflateInit_nif", 6, zlib_deflateInit},
156     {"deflateSetDictionary_nif", 2, zlib_deflateSetDictionary},
157     {"deflateReset_nif", 1, zlib_deflateReset},
158     {"deflateEnd_nif", 1, zlib_deflateEnd},
159     {"deflateParams_nif", 3, zlib_deflateParams},
160     {"deflate_nif", 4, zlib_deflate},
161 
162     /* inflate */
163     {"inflateInit_nif", 3, zlib_inflateInit},
164     {"inflateSetDictionary_nif", 2, zlib_inflateSetDictionary},
165     {"inflateGetDictionary_nif", 1, zlib_inflateGetDictionary},
166     {"inflateReset_nif", 1, zlib_inflateReset},
167     {"inflateEnd_nif", 1, zlib_inflateEnd},
168     {"inflate_nif", 4, zlib_inflate},
169 
170     /* running checksum */
171     {"crc32_nif", 1, zlib_crc32},
172 
173     /* The stash keeps a single term alive across calls, and is used in
174      * exception_on_need_dict/1 to retain the old error behavior, and for
175      * saving data flushed through deflateParams/3. */
176     {"getStash_nif", 1, zlib_getStash},
177     {"clearStash_nif", 1, zlib_clearStash},
178     {"setStash_nif", 2, zlib_setStash},
179 
180     /* DEPRECATED: buffer size for inflateChunk */
181     {"getBufSize_nif", 1, zlib_getBufSize},
182     {"setBufSize_nif", 2, zlib_setBufSize},
183 
184     {"enqueue_nif", 2, zlib_enqueue_input},
185 };
186 
187 ERL_NIF_INIT(zlib, nif_funcs, load, NULL, upgrade, unload)
188 
189 static void gc_zlib(ErlNifEnv *env, void* data);
190 
load(ErlNifEnv * env,void ** priv_data,ERL_NIF_TERM load_info)191 static int load(ErlNifEnv *env, void** priv_data, ERL_NIF_TERM load_info)
192 {
193     am_not_on_controlling_process =
194         enif_make_atom(env, "not_on_controlling_process");
195 
196     am_not_initialized = enif_make_atom(env, "not_initialized");
197     am_already_initialized = enif_make_atom(env, "already_initialized");
198 
199     am_ok = enif_make_atom(env, "ok");
200     am_error = enif_make_atom(env, "error");
201 
202     am_continue = enif_make_atom(env, "continue");
203     am_finished = enif_make_atom(env, "finished");
204 
205     am_not_supported = enif_make_atom(env, "not_supported");
206     am_need_dictionary = enif_make_atom(env, "need_dictionary");
207 
208     am_empty = enif_make_atom(env, "empty");
209 
210     am_stream_end = enif_make_atom(env, "stream_end");
211     am_stream_error = enif_make_atom(env, "stream_error");
212     am_data_error = enif_make_atom(env, "data_error");
213     am_mem_error = enif_make_atom(env, "mem_error");
214     am_buf_error = enif_make_atom(env, "buf_error");
215     am_version_error = enif_make_atom(env, "version_error");
216     am_unknown_error = enif_make_atom(env, "unknown_error");
217 
218     rtype_zlib = enif_open_resource_type(env, NULL,
219             "gc_zlib", gc_zlib, ERL_NIF_RT_CREATE, NULL);
220     *priv_data = NULL;
221 
222     return 0;
223 }
224 
unload(ErlNifEnv * env,void * priv_data)225 static void unload(ErlNifEnv *env, void* priv_data)
226 {
227 
228 }
229 
upgrade(ErlNifEnv * env,void ** priv_data,void ** old_priv_data,ERL_NIF_TERM load_info)230 static int upgrade(ErlNifEnv *env, void** priv_data, void** old_priv_data, ERL_NIF_TERM load_info)
231 {
232     if(*old_priv_data != NULL) {
233         return -1; /* Don't know how to do that */
234     }
235     if(*priv_data != NULL) {
236         return -1; /* Don't know how to do that */
237     }
238     if(load(env, priv_data, load_info)) {
239         return -1;
240     }
241     return 0;
242 }
243 
zlib_alloc(void * data,unsigned int items,unsigned int size)244 static void* zlib_alloc(void* data, unsigned int items, unsigned int size)
245 {
246     return (void*) enif_alloc(items * size);
247 }
248 
zlib_free(void * data,void * addr)249 static void zlib_free(void* data, void* addr)
250 {
251     enif_free(addr);
252 }
253 
zlib_return(ErlNifEnv * env,int code)254 static ERL_NIF_TERM zlib_return(ErlNifEnv *env, int code) {
255     ERL_NIF_TERM reason;
256     switch(code) {
257     case Z_OK:
258         reason = am_ok;
259         break;
260     case Z_STREAM_END:
261         reason = am_stream_end;
262         break;
263     case Z_ERRNO:
264         reason = enif_make_int(env, errno);
265         break;
266     case Z_STREAM_ERROR:
267         reason = enif_raise_exception(env, am_stream_error);
268         break;
269     case Z_DATA_ERROR:
270         reason = enif_raise_exception(env, am_data_error);
271         break;
272     case Z_MEM_ERROR:
273         reason = am_mem_error;
274         break;
275     case Z_BUF_ERROR:
276         reason = am_buf_error;
277         break;
278     case Z_VERSION_ERROR:
279         reason = am_version_error;
280         break;
281     default:
282         reason = am_unknown_error;
283         break;
284     }
285     return reason;
286 }
287 
zlib_internal_close(zlib_data_t * d)288 static void zlib_internal_close(zlib_data_t *d) {
289     if(d->state == ST_DEFLATE) {
290         deflateEnd(&d->s);
291     } else if(d->state == ST_INFLATE) {
292         inflateEnd(&d->s);
293     }
294 
295     if(d->state != ST_CLOSED) {
296         if(d->stash_env != NULL) {
297             enif_free_env(d->stash_env);
298         }
299 
300         d->state = ST_CLOSED;
301     }
302 }
303 
gc_zlib(ErlNifEnv * env,void * data)304 static void gc_zlib(ErlNifEnv *env, void* data) {
305     zlib_data_t *d = (zlib_data_t*)data;
306 
307     enif_mutex_destroy(d->controller_lock);
308     enif_ioq_destroy(d->input_queue);
309 
310     zlib_internal_close(d);
311 
312     (void)env;
313 }
314 
get_zlib_data(ErlNifEnv * env,ERL_NIF_TERM opaque,zlib_data_t ** d)315 static int get_zlib_data(ErlNifEnv *env, ERL_NIF_TERM opaque, zlib_data_t **d) {
316     return enif_get_resource(env, opaque, rtype_zlib, (void **)d);
317 }
318 
zlib_process_check(ErlNifEnv * env,zlib_data_t * d)319 static int zlib_process_check(ErlNifEnv *env, zlib_data_t *d) {
320     int is_controlling_process;
321     ErlNifPid current_process;
322 
323     enif_self(env, &current_process);
324 
325     enif_mutex_lock(d->controller_lock);
326 
327     is_controlling_process = enif_is_identical(
328         enif_make_pid(env, &current_process),
329         enif_make_pid(env, &d->controlling_process));
330 
331     enif_mutex_unlock(d->controller_lock);
332 
333     return is_controlling_process;
334 }
335 
zlib_reset_input(zlib_data_t * d)336 static void zlib_reset_input(zlib_data_t *d) {
337     enif_ioq_destroy(d->input_queue);
338     d->input_queue = enif_ioq_create(ERL_NIF_IOQ_NORMAL);
339 
340     if(d->stash_env != NULL) {
341         enif_free_env(d->stash_env);
342         d->stash_env = NULL;
343         d->stash_term = NIL;
344     }
345 }
346 
zlib_flush_queue(int (* codec)(z_stream *,int),ErlNifEnv * env,zlib_data_t * d,size_t input_limit,ErlNifBinary * output_buffer,int flush,size_t * bytes_produced,size_t * bytes_consumed,size_t * bytes_remaining)347 static int zlib_flush_queue(int (*codec)(z_stream*, int), ErlNifEnv *env,
348         zlib_data_t *d, size_t input_limit, ErlNifBinary *output_buffer, int flush,
349         size_t *bytes_produced, size_t *bytes_consumed, size_t *bytes_remaining) {
350 
351     int vec_len, vec_idx;
352     SysIOVec *input_vec;
353     int res;
354 
355     input_vec = enif_ioq_peek(d->input_queue, &vec_len);
356     vec_idx = 0;
357     res = Z_OK;
358 
359     *bytes_produced = 0;
360     *bytes_consumed = 0;
361 
362     d->s.avail_out = output_buffer->size;
363     d->s.next_out = output_buffer->data;
364 
365     while(res == Z_OK && vec_idx < vec_len && *bytes_consumed < input_limit) {
366         size_t timeslice_percent, block_consumed, block_size;
367 
368         block_size = MIN(input_vec[vec_idx].iov_len, input_limit);
369 
370         d->s.next_in = input_vec[vec_idx].iov_base;
371         d->s.avail_in = block_size;
372 
373         res = codec(&d->s, Z_NO_FLUSH);
374 
375         ASSERT(d->s.avail_in == 0 || d->s.avail_out == 0 || res != Z_OK);
376 
377         block_consumed = block_size - d->s.avail_in;
378         *bytes_consumed += block_consumed;
379 
380         if(d->want_input_crc) {
381             d->input_crc =
382                 crc32(d->input_crc, input_vec[vec_idx].iov_base, block_consumed);
383         }
384 
385         timeslice_percent = (100 * block_consumed) / input_limit;
386         if(enif_consume_timeslice(env, MAX(1, timeslice_percent))) {
387             break;
388         }
389 
390         vec_idx++;
391     }
392 
393     if(!enif_ioq_deq(d->input_queue, *bytes_consumed, bytes_remaining)) {
394         *bytes_remaining = 0;
395         res = Z_BUF_ERROR;
396     }
397 
398     if(res == Z_OK && flush != Z_NO_FLUSH && (*bytes_remaining == 0)) {
399         d->s.next_in = NULL;
400         d->s.avail_in = 0;
401 
402         res = codec(&d->s, flush);
403     }
404 
405     *bytes_produced = output_buffer->size - d->s.avail_out;
406 
407     return res;
408 }
409 
zlib_codec(int (* codec)(z_stream *,int),ErlNifEnv * env,zlib_data_t * d,int input_chunk_size,int output_chunk_size,int flush)410 static ERL_NIF_TERM zlib_codec(int (*codec)(z_stream*, int),
411                                ErlNifEnv *env, zlib_data_t *d,
412                                int input_chunk_size,
413                                int output_chunk_size,
414                                int flush) {
415 
416     size_t bytes_produced, bytes_consumed, bytes_remaining;
417     ErlNifBinary output_buffer;
418     int res;
419 
420     if(!enif_alloc_binary(output_chunk_size, &output_buffer)) {
421         return zlib_return(env, Z_MEM_ERROR);
422     }
423 
424     res = zlib_flush_queue(codec, env, d, input_chunk_size, &output_buffer,
425         flush, &bytes_produced, &bytes_consumed, &bytes_remaining);
426 
427     if(res < 0 && res != Z_BUF_ERROR) {
428         enif_release_binary(&output_buffer);
429         return zlib_return(env, res);
430     }
431 
432     if(res == Z_STREAM_END) {
433         d->eos_seen = 1;
434     }
435 
436     if(d->want_output_crc) {
437         d->output_crc =
438             crc32(d->output_crc, output_buffer.data, bytes_produced);
439     }
440 
441     if(bytes_consumed == 0 && bytes_produced == 0 && bytes_remaining != 0) {
442         /* Die if we've made zero progress; this should not happen on
443          * well-formed input. */
444 
445         enif_release_binary(&output_buffer);
446         return zlib_return(env, Z_DATA_ERROR);
447     } else {
448         ERL_NIF_TERM flushed_output;
449 
450         if(bytes_produced > 0) {
451             if(bytes_produced < output_buffer.size) {
452                 enif_realloc_binary(&output_buffer, bytes_produced);
453             }
454 
455             flushed_output =
456                 enif_make_list1(env, enif_make_binary(env, &output_buffer));
457         } else {
458             enif_release_binary(&output_buffer);
459             flushed_output = enif_make_list(env, 0);
460         }
461 
462         if(bytes_remaining == 0 && bytes_produced < output_chunk_size) {
463             return enif_make_tuple2(env, am_finished, flushed_output);
464         } else if(res != Z_NEED_DICT) {
465             return enif_make_tuple2(env, am_continue, flushed_output);
466         }
467 
468         return enif_make_tuple3(env, am_need_dictionary,
469             enif_make_int(env, d->s.adler), flushed_output);
470     }
471 }
472 
473 /* zlib nifs */
474 
zlib_getStash(ErlNifEnv * env,int argc,const ERL_NIF_TERM argv[])475 static ERL_NIF_TERM zlib_getStash(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) {
476     zlib_data_t *d;
477 
478     if(argc != 1 || !get_zlib_data(env, argv[0], &d)) {
479         return enif_make_badarg(env);
480     } else if(!zlib_process_check(env, d)) {
481         return enif_raise_exception(env, am_not_on_controlling_process);
482     }
483 
484     if(d->stash_env == NULL) {
485         return am_empty;
486     }
487 
488     return enif_make_tuple2(env, am_ok, enif_make_copy(env, d->stash_term));
489 }
490 
zlib_clearStash(ErlNifEnv * env,int argc,const ERL_NIF_TERM argv[])491 static ERL_NIF_TERM zlib_clearStash(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) {
492     zlib_data_t *d;
493 
494     if(argc != 1 || !get_zlib_data(env, argv[0], &d)) {
495         return enif_make_badarg(env);
496     } else if(!zlib_process_check(env, d)) {
497         return enif_raise_exception(env, am_not_on_controlling_process);
498     } else if(d->stash_env == NULL) {
499         return enif_raise_exception(env, am_error);
500     }
501 
502     enif_free_env(d->stash_env);
503     d->stash_env = NULL;
504     d->stash_term = NIL;
505 
506     return am_ok;
507 }
508 
zlib_setStash(ErlNifEnv * env,int argc,const ERL_NIF_TERM argv[])509 static ERL_NIF_TERM zlib_setStash(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) {
510     zlib_data_t *d;
511 
512     if(argc != 2 || !get_zlib_data(env, argv[0], &d)) {
513         return enif_make_badarg(env);
514     } else if(!zlib_process_check(env, d)) {
515         return enif_raise_exception(env, am_not_on_controlling_process);
516     } else if(d->stash_env != NULL) {
517         return enif_raise_exception(env, am_error);
518     }
519 
520     d->stash_env = enif_alloc_env();
521     d->stash_term = enif_make_copy(d->stash_env, argv[1]);
522 
523     return am_ok;
524 }
525 
zlib_open(ErlNifEnv * env,int argc,const ERL_NIF_TERM argv[])526 static ERL_NIF_TERM zlib_open(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) {
527     zlib_data_t *d;
528     ERL_NIF_TERM result;
529 
530     d = (zlib_data_t *) enif_alloc_resource(rtype_zlib, sizeof(zlib_data_t));
531 
532     memset(&d->s, 0, sizeof(z_stream));
533 
534     enif_self(env, &d->controlling_process);
535 
536     d->input_queue = enif_ioq_create(ERL_NIF_IOQ_NORMAL);
537 
538     d->controller_lock = enif_mutex_create("zlib_controller_lock");
539 
540     d->s.zalloc = zlib_alloc;
541     d->s.zfree  = zlib_free;
542     d->s.opaque = d;
543     d->s.data_type = Z_BINARY;
544 
545     d->eos_behavior = EOS_BEHAVIOR_CUT;
546     d->eos_seen = 0;
547 
548     d->state = ST_NONE;
549 
550     d->want_output_crc = 0;
551     d->want_input_crc = 0;
552     d->is_raw_stream = 0;
553 
554     d->output_crc = crc32(0L, Z_NULL, 0);
555     d->input_crc = crc32(0L, Z_NULL, 0);
556 
557     d->stash_env = NULL;
558     d->stash_term = NIL;
559 
560     d->inflateChunk_buffer_size = 4000;
561 
562     result = enif_make_resource(env, d);
563     enif_release_resource(d);
564 
565     return result;
566 }
567 
zlib_close(ErlNifEnv * env,int argc,const ERL_NIF_TERM argv[])568 static ERL_NIF_TERM zlib_close(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) {
569     zlib_data_t *d;
570 
571     /* strictly speaking not needed since the gc will handle this */
572     if(argc != 1 || !get_zlib_data(env, argv[0], &d)) {
573         return enif_make_badarg(env);
574     } else if(!zlib_process_check(env, d)) {
575         return enif_raise_exception(env, am_not_on_controlling_process);
576     } else if(d->state == ST_CLOSED) {
577         return enif_raise_exception(env, am_not_initialized);
578     }
579 
580     zlib_internal_close(d);
581 
582     return am_ok;
583 }
584 
zlib_set_controller(ErlNifEnv * env,int argc,const ERL_NIF_TERM argv[])585 static ERL_NIF_TERM zlib_set_controller(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) {
586     zlib_data_t *d;
587 
588     ErlNifPid new_owner;
589 
590     if(argc != 2 || !get_zlib_data(env, argv[0], &d)
591                  || !enif_get_local_pid(env, argv[1], &new_owner)) {
592         return enif_make_badarg(env);
593     } else if(!zlib_process_check(env, d)) {
594         return enif_raise_exception(env, am_not_on_controlling_process);
595     }
596 
597     enif_mutex_lock(d->controller_lock);
598 
599     d->controlling_process = new_owner;
600 
601     enif_mutex_unlock(d->controller_lock);
602 
603     return am_ok;
604 }
605 
606 /* deflate */
607 
zlib_deflateInit(ErlNifEnv * env,int argc,const ERL_NIF_TERM argv[])608 static ERL_NIF_TERM zlib_deflateInit(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) {
609     zlib_data_t *d;
610     int level, method, windowBits, memLevel, strategy, res;
611 
612     if(argc != 6 || !get_zlib_data(env, argv[0], &d)
613                  || !enif_get_int(env, argv[1], &level)
614                  || !enif_get_int(env, argv[2], &method)
615                  || !enif_get_int(env, argv[3], &windowBits)
616                  || !enif_get_int(env, argv[4], &memLevel)
617                  || !enif_get_int(env, argv[5], &strategy)) {
618         return enif_make_badarg(env);
619     } else if(!zlib_process_check(env, d)) {
620         return enif_raise_exception(env, am_not_on_controlling_process);
621     } else if(d->state != ST_NONE) {
622         return enif_raise_exception(env, am_already_initialized);
623     }
624 
625     res = deflateInit2(&d->s, level, method, windowBits, memLevel, strategy);
626 
627     if(res == Z_OK) {
628         d->state = ST_DEFLATE;
629         d->eos_seen = 0;
630 
631         d->is_raw_stream = (windowBits < 0);
632 
633         d->want_output_crc = 0;
634         d->want_input_crc = d->is_raw_stream;
635 
636         d->output_crc = crc32(0L, Z_NULL, 0);
637         d->input_crc = crc32(0L, Z_NULL, 0);
638     }
639 
640     return zlib_return(env, res);
641 }
642 
zlib_deflateSetDictionary(ErlNifEnv * env,int argc,const ERL_NIF_TERM argv[])643 static ERL_NIF_TERM zlib_deflateSetDictionary(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) {
644     zlib_data_t *d;
645     ErlNifBinary bin;
646     int res;
647 
648     if(argc != 2 || !get_zlib_data(env, argv[0], &d)
649                  || !enif_inspect_iolist_as_binary(env, argv[1], &bin)) {
650         return enif_make_badarg(env);
651     } else if(!zlib_process_check(env, d)) {
652         return enif_raise_exception(env, am_not_on_controlling_process);
653     } else if(d->state != ST_DEFLATE) {
654         return enif_raise_exception(env, am_not_initialized);
655     }
656 
657     if((res = deflateSetDictionary(&d->s, bin.data, bin.size)) == Z_OK) {
658         uLong checksum = d->s.adler;
659 
660         /* d->s.adler is not updated in raw deflate mode, so we'll calculate it
661          * ourselves in case the user wants to rely on that behavior. */
662         if(d->is_raw_stream) {
663             checksum = adler32(0, bin.data, bin.size);
664         }
665 
666         return enif_make_int(env, checksum);
667     }
668 
669     return zlib_return(env, res);
670 }
671 
zlib_deflateReset(ErlNifEnv * env,int argc,const ERL_NIF_TERM argv[])672 static ERL_NIF_TERM zlib_deflateReset(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) {
673     zlib_data_t *d;
674     int res;
675 
676     if(argc != 1 || !get_zlib_data(env, argv[0], &d)) {
677         return enif_make_badarg(env);
678     } else if(!zlib_process_check(env, d)) {
679         return enif_raise_exception(env, am_not_on_controlling_process);
680     } else if(d->state != ST_DEFLATE) {
681         return enif_raise_exception(env, am_not_initialized);
682     }
683 
684     res = deflateReset(&d->s);
685 
686     d->input_crc = crc32(0L, Z_NULL, 0);
687     d->eos_seen = 0;
688 
689     zlib_reset_input(d);
690 
691     return zlib_return(env, res);
692 }
693 
zlib_deflateEnd(ErlNifEnv * env,int argc,const ERL_NIF_TERM argv[])694 static ERL_NIF_TERM zlib_deflateEnd(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) {
695     zlib_data_t *d;
696     int res;
697 
698     if(argc != 1 || !get_zlib_data(env, argv[0], &d)) {
699         return enif_make_badarg(env);
700     } else if(!zlib_process_check(env, d)) {
701         return enif_raise_exception(env, am_not_on_controlling_process);
702     } else if(d->state != ST_DEFLATE) {
703         return enif_raise_exception(env, am_not_initialized);
704     }
705 
706     res = deflateEnd(&d->s);
707 
708     if(res == Z_OK && enif_ioq_size(d->input_queue) > 0) {
709         res = Z_DATA_ERROR;
710     }
711 
712     zlib_reset_input(d);
713     d->state = ST_NONE;
714 
715     return zlib_return(env, res);
716 }
717 
zlib_deflateParams(ErlNifEnv * env,int argc,const ERL_NIF_TERM argv[])718 static ERL_NIF_TERM zlib_deflateParams(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) {
719     zlib_data_t *d;
720 
721     int res, level, strategy;
722     Bytef dummy_buffer;
723 
724     if(argc != 3 || !get_zlib_data(env, argv[0], &d)
725                  || !enif_get_int(env, argv[1], &level)
726                  || !enif_get_int(env, argv[2], &strategy)) {
727         return enif_make_badarg(env);
728     } else if(!zlib_process_check(env, d)) {
729         return enif_raise_exception(env, am_not_on_controlling_process);
730     } else if(d->state != ST_DEFLATE) {
731         return enif_raise_exception(env, am_not_initialized);
732     }
733 
734     /* This is a bit of a hack; deflateParams flushes with Z_BLOCK which won't
735      * stop at a byte boundary, so we can't split this operation up, and we
736      * can't allocate a buffer large enough to fit it in one go since we have
737      * to support zlib versions that lack deflatePending.
738      *
739      * We therefore flush everything prior to this call to ensure that we are
740      * stopped on a byte boundary and have no pending data. We then hand it a
741      * dummy buffer to detect when this assumption doesn't hold (Hopefully
742      * never), and to smooth over an issue with zlib 1.2.11 which always
743      * returns Z_BUF_ERROR when d->s.avail_out is 0, regardless of whether
744      * there's any pending data or not. */
745 
746     d->s.next_out = &dummy_buffer;
747     d->s.avail_out = 1;
748 
749     res = deflateParams(&d->s, level, strategy);
750 
751     if(d->s.avail_out == 0) {
752         return zlib_return(env, Z_STREAM_ERROR);
753     }
754 
755     return zlib_return(env, res);
756 }
757 
zlib_deflate(ErlNifEnv * env,int argc,const ERL_NIF_TERM argv[])758 static ERL_NIF_TERM zlib_deflate(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) {
759     zlib_data_t *d;
760 
761     int input_chunk_size, output_chunk_size, flush;
762 
763     if(argc != 4 || !get_zlib_data(env, argv[0], &d)
764                  || !enif_get_int(env, argv[1], &input_chunk_size)
765                  || !enif_get_int(env, argv[2], &output_chunk_size)
766                  || !enif_get_int(env, argv[3], &flush)) {
767         return enif_make_badarg(env);
768     } else if(!zlib_process_check(env, d)) {
769         return enif_raise_exception(env, am_not_on_controlling_process);
770     } else if(d->state != ST_DEFLATE) {
771         return enif_raise_exception(env, am_not_initialized);
772     }
773 
774     return zlib_codec(&deflate, env, d, input_chunk_size, output_chunk_size, flush);
775 }
776 
777 /* inflate */
778 
zlib_inflateInit(ErlNifEnv * env,int argc,const ERL_NIF_TERM argv[])779 static ERL_NIF_TERM zlib_inflateInit(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) {
780     zlib_data_t *d;
781 
782     int windowBits, eosBehavior, res;
783 
784     if(argc != 3 || !get_zlib_data(env, argv[0], &d)
785                  || !enif_get_int(env, argv[1], &windowBits)
786                  || !enif_get_int(env, argv[2], &eosBehavior)) {
787         return enif_make_badarg(env);
788     } else if(!zlib_process_check(env, d)) {
789         return enif_raise_exception(env, am_not_on_controlling_process);
790     } else if(d->state != ST_NONE) {
791         return enif_raise_exception(env, am_already_initialized);
792     }
793 
794     res = inflateInit2(&d->s, windowBits);
795 
796     if(res == Z_OK) {
797         d->state = ST_INFLATE;
798 
799         d->eos_behavior = eosBehavior;
800         d->eos_seen = 0;
801 
802         d->is_raw_stream = (windowBits < 0);
803 
804         d->want_output_crc = d->is_raw_stream;
805         d->want_input_crc = 0;
806 
807         d->output_crc = crc32(0L, Z_NULL, 0);
808         d->input_crc = crc32(0L, Z_NULL, 0);
809     }
810 
811     return zlib_return(env, res);
812 }
813 
zlib_inflateSetDictionary(ErlNifEnv * env,int argc,const ERL_NIF_TERM argv[])814 static ERL_NIF_TERM zlib_inflateSetDictionary(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) {
815     zlib_data_t *d;
816     ErlNifBinary bin;
817     int res;
818 
819     if(argc != 2 || !get_zlib_data(env, argv[0], &d)
820                  || !enif_inspect_iolist_as_binary(env, argv[1], &bin)) {
821         return enif_make_badarg(env);
822     } else if(!zlib_process_check(env, d)) {
823         return enif_raise_exception(env, am_not_on_controlling_process);
824     } else if(d->state != ST_INFLATE) {
825         return enif_raise_exception(env, am_not_initialized);
826     }
827 
828     res = inflateSetDictionary(&d->s, bin.data, bin.size);
829 
830     return zlib_return(env, res);
831 }
832 
833 #ifdef HAVE_ZLIB_INFLATEGETDICTIONARY
834 /* Work around broken build system with runtime version test */
zlib_supports_inflateGetDictionary(void)835 static int zlib_supports_inflateGetDictionary(void) {
836     static int supportsGetDictionary = -1;
837 
838 #if defined(__APPLE__) && defined(__MACH__)
839     if(supportsGetDictionary < 0) {
840         unsigned int v[4] = {0, 0, 0, 0};
841         unsigned hexver;
842 
843         sscanf(zlibVersion(), "%u.%u.%u.%u", &v[0], &v[1], &v[2], &v[3]);
844 
845         hexver = (v[0] << (8*3)) | (v[1] << (8*2)) | (v[2] << (8)) | v[3];
846         supportsGetDictionary = (hexver >= 0x1020701); /* 1.2.7.1 */
847     }
848 #endif
849 
850     return supportsGetDictionary;
851 }
852 #endif
853 
zlib_inflateGetDictionary(ErlNifEnv * env,int argc,const ERL_NIF_TERM argv[])854 static ERL_NIF_TERM zlib_inflateGetDictionary(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) {
855     zlib_data_t *d;
856 
857     if(argc != 1 || !get_zlib_data(env, argv[0], &d)) {
858         return enif_make_badarg(env);
859     } else if(!zlib_process_check(env, d)) {
860         return enif_raise_exception(env, am_not_on_controlling_process);
861     } else if(d->state != ST_INFLATE) {
862         return enif_raise_exception(env, am_not_initialized);
863     }
864 
865 #ifdef HAVE_ZLIB_INFLATEGETDICTIONARY
866     if(zlib_supports_inflateGetDictionary()) {
867         ErlNifBinary obin;
868         uInt len;
869         int res;
870 
871         enif_alloc_binary(INFL_DICT_SZ, &obin);
872         len = 0;
873 
874         if((res = inflateGetDictionary(&d->s, obin.data, &len)) < 0) {
875             enif_release_binary(&obin);
876             return zlib_return(env, res);
877         }
878 
879         enif_realloc_binary(&obin, (size_t)len);
880         return enif_make_binary(env, &obin);
881     }
882 #endif
883 
884     return enif_raise_exception(env, am_not_supported);
885 }
886 
zlib_inflateReset(ErlNifEnv * env,int argc,const ERL_NIF_TERM argv[])887 static ERL_NIF_TERM zlib_inflateReset(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) {
888     zlib_data_t *d;
889     int res;
890 
891     if(argc != 1 || !get_zlib_data(env, argv[0], &d)) {
892         return enif_make_badarg(env);
893     } else if(!zlib_process_check(env, d)) {
894         return enif_raise_exception(env, am_not_on_controlling_process);
895     } else if(d->state != ST_INFLATE) {
896         return enif_raise_exception(env, am_not_initialized);
897     }
898 
899     res = inflateReset(&d->s);
900 
901     d->output_crc = crc32(0L, Z_NULL, 0);
902     d->eos_seen = 0;
903 
904     zlib_reset_input(d);
905 
906     return zlib_return(env, res);
907 }
908 
zlib_inflateEnd(ErlNifEnv * env,int argc,const ERL_NIF_TERM argv[])909 static ERL_NIF_TERM zlib_inflateEnd(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) {
910     zlib_data_t *d;
911     int res;
912 
913     if(argc != 1 || !get_zlib_data(env, argv[0], &d)) {
914         return enif_make_badarg(env);
915     } else if(!zlib_process_check(env, d)) {
916         return enif_raise_exception(env, am_not_on_controlling_process);
917     } else if(d->state != ST_INFLATE) {
918         return enif_raise_exception(env, am_not_initialized);
919     }
920 
921     res = inflateEnd(&d->s);
922 
923     if(res == Z_OK && (!d->eos_seen || enif_ioq_size(d->input_queue) > 0)) {
924         res = Z_DATA_ERROR;
925     }
926 
927     zlib_reset_input(d);
928     d->state = ST_NONE;
929 
930     return zlib_return(env, res);
931 }
932 
zlib_inflate(ErlNifEnv * env,int argc,const ERL_NIF_TERM argv[])933 static ERL_NIF_TERM zlib_inflate(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) {
934     zlib_data_t *d;
935 
936     int input_chunk_size, output_chunk_size, flush;
937 
938     if(argc != 4 || !get_zlib_data(env, argv[0], &d)
939                  || !enif_get_int(env, argv[1], &input_chunk_size)
940                  || !enif_get_int(env, argv[2], &output_chunk_size)
941                  || !enif_get_int(env, argv[3], &flush)) {
942         return enif_make_badarg(env);
943     } else if(!zlib_process_check(env, d)) {
944         return enif_raise_exception(env, am_not_on_controlling_process);
945     } else if(d->state != ST_INFLATE) {
946         return enif_raise_exception(env, am_not_initialized);
947     }
948 
949     if(d->eos_seen && enif_ioq_size(d->input_queue) > 0) {
950         int res;
951 
952         switch(d->eos_behavior) {
953         case EOS_BEHAVIOR_ERROR:
954             return zlib_return(env, Z_DATA_ERROR);
955         case EOS_BEHAVIOR_RESET:
956             res = inflateReset(&d->s);
957 
958             if(res != Z_OK) {
959                 return zlib_return(env, res);
960             }
961 
962             d->eos_seen = 0;
963 
964             break;
965         case EOS_BEHAVIOR_CUT:
966             zlib_reset_input(d);
967         }
968     }
969 
970     return zlib_codec(&inflate, env, d, input_chunk_size, output_chunk_size, flush);
971 }
972 
zlib_crc32(ErlNifEnv * env,int argc,const ERL_NIF_TERM argv[])973 static ERL_NIF_TERM zlib_crc32(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) {
974     zlib_data_t *d;
975 
976     if(argc != 1 || !get_zlib_data(env, argv[0], &d)) {
977         return enif_make_badarg(env);
978     } else if(!zlib_process_check(env, d)) {
979         return enif_raise_exception(env, am_not_on_controlling_process);
980     }
981 
982     if(d->state == ST_DEFLATE) {
983         return enif_make_ulong(env, d->input_crc);
984     } else if(d->state == ST_INFLATE) {
985         return enif_make_ulong(env, d->output_crc);
986     }
987 
988     return enif_raise_exception(env, am_not_initialized);
989 }
990 
zlib_getBufSize(ErlNifEnv * env,int argc,const ERL_NIF_TERM argv[])991 static ERL_NIF_TERM zlib_getBufSize(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) {
992     zlib_data_t *d;
993 
994     if(argc != 1 || !get_zlib_data(env, argv[0], &d)) {
995         return enif_make_badarg(env);
996     } else if(!zlib_process_check(env, d)) {
997         return enif_raise_exception(env, am_not_on_controlling_process);
998     }
999 
1000     return enif_make_int(env, d->inflateChunk_buffer_size);
1001 }
1002 
zlib_setBufSize(ErlNifEnv * env,int argc,const ERL_NIF_TERM argv[])1003 static ERL_NIF_TERM zlib_setBufSize(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) {
1004     zlib_data_t *d;
1005 
1006     if(argc != 2 || !get_zlib_data(env, argv[0], &d)) {
1007         return enif_make_badarg(env);
1008     } else if(!zlib_process_check(env, d)) {
1009         return enif_raise_exception(env, am_not_on_controlling_process);
1010     }
1011 
1012     if(!enif_get_int(env, argv[1], &d->inflateChunk_buffer_size)) {
1013         return enif_make_badarg(env);
1014     }
1015 
1016     return am_ok;
1017 }
1018 
zlib_enqueue_input(ErlNifEnv * env,int argc,const ERL_NIF_TERM argv[])1019 static ERL_NIF_TERM zlib_enqueue_input(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) {
1020     zlib_data_t *d;
1021 
1022     ErlNifIOVec prealloc, *iovec = &prealloc;
1023     ERL_NIF_TERM tail;
1024 
1025     if(argc != 2 || !get_zlib_data(env, argv[0], &d)) {
1026         return enif_make_badarg(env);
1027     } else if(!zlib_process_check(env, d)) {
1028         return enif_raise_exception(env, am_not_on_controlling_process);
1029     } else if(d->state != ST_DEFLATE && d->state != ST_INFLATE) {
1030         return enif_raise_exception(env, am_not_initialized);
1031     }
1032 
1033     if(!enif_inspect_iovec(env, 256, argv[1], &tail, &iovec)) {
1034         return enif_make_badarg(env);
1035     } else if(!enif_ioq_enqv(d->input_queue, iovec, 0)) {
1036         return enif_make_badarg(env);
1037     }
1038 
1039     if(!enif_is_empty_list(env, tail)) {
1040         return enif_make_tuple2(env, am_continue, tail);
1041     }
1042 
1043     return am_ok;
1044 }
1045