diff --git a/enc/encode.cc b/enc/encode.cc index 51b11efab..df6a59948 100644 --- a/enc/encode.cc +++ b/enc/encode.cc @@ -687,8 +687,16 @@ bool BrotliCompressor::WriteBrotliData(const bool is_last, assert(input_pos_ - last_flush_pos_ <= 1u << 24); const uint32_t metablock_size = static_cast(input_pos_ - last_flush_pos_); - const size_t max_out_size = 2 * metablock_size + 500; - uint8_t* storage = GetBrotliStorage(max_out_size); + size_t metadata_size = params_.comment ? strlen(params_.comment) : 0; + size_t encoded_metadata_size = metadata_size ? 6 + metadata_size : 0; + const size_t max_out_size = encoded_metadata_size + 2 * metablock_size + 500; + uint8_t* output_storage = GetBrotliStorage(max_out_size); + if (encoded_metadata_size) { + if (!WriteMetadata(metadata_size, (const uint8_t*)params_.comment, false, &encoded_metadata_size, output_storage)) + return false; + params_.comment = 0; + } + uint8_t* storage = output_storage + encoded_metadata_size; storage[0] = last_byte_; size_t storage_ix = last_byte_bits_; bool font_mode = params_.mode == BrotliParams::MODE_FONT; @@ -711,8 +719,8 @@ bool BrotliCompressor::WriteBrotliData(const bool is_last, // Save the state of the distance cache in case we need to restore it for // emitting an uncompressed block. memcpy(saved_dist_cache_, dist_cache_, sizeof(dist_cache_)); - *output = &storage[0]; - *out_size = storage_ix >> 3; + *output = &output_storage[0]; + *out_size = encoded_metadata_size + (storage_ix >> 3); return true; } diff --git a/enc/encode.h b/enc/encode.h index 1b0885e53..a12f205a2 100644 --- a/enc/encode.h +++ b/enc/encode.h @@ -34,7 +34,8 @@ struct BrotliParams { enable_dictionary(true), enable_transforms(false), greedy_block_split(false), - enable_context_modeling(true) {} + enable_context_modeling(true), + comment(0) {} enum Mode { // Default compression mode. The compressor does not know anything in @@ -62,6 +63,7 @@ struct BrotliParams { bool enable_transforms; bool greedy_block_split; bool enable_context_modeling; + const char* comment; }; // An instance can not be reused for multiple brotli streams. diff --git a/python/bro.py b/python/bro.py old mode 100755 new mode 100644 index c6f74ce94..a3739eb21 --- a/python/bro.py +++ b/python/bro.py @@ -81,6 +81,8 @@ def main(args=None): 'on the quality. Defaults to 0.') params.add_argument('--custom-dictionary', metavar="FILE", type=str, dest='dictfile', help='Custom dictionary file.', default = None) + params.add_argument('--comment', type=str, dest='comment', + help='Archive comment.', default = None) # set default values using global DEFAULT_PARAMS dictionary parser.set_defaults(**DEFAULT_PARAMS) @@ -113,6 +115,10 @@ def main(args=None): else: custom_dictionary = '' + if options.comment: + archive_comment = options.comment + else: + archive_comment = '' try: if options.decompress: @@ -120,7 +126,7 @@ def main(args=None): else: data = brotli.compress( data, mode=options.mode, quality=options.quality, - lgwin=options.lgwin, lgblock=options.lgblock, dictionary=custom_dictionary) + lgwin=options.lgwin, lgblock=options.lgblock, dictionary=custom_dictionary, comment=archive_comment) except brotli.error as e: parser.exit(1,'bro: error: %s: %s' % (e, options.infile or 'sys.stdin')) diff --git a/python/brotlimodule.cc b/python/brotlimodule.cc index fbc73ca50..e10ef6dca 100644 --- a/python/brotlimodule.cc +++ b/python/brotlimodule.cc @@ -107,6 +107,8 @@ PyDoc_STRVAR(compress__doc__, " quality. Defaults to 0.\n" " dictionary (bytes, optional): Custom dictionary. Only last sliding window\n" " size bytes will be used.\n" +" comment (string, optional): A comment that will be added as the leading\n" +" metadata metablock.\n" "\n" "Returns:\n" " The compressed byte string.\n" @@ -118,6 +120,7 @@ static PyObject* brotli_compress(PyObject *self, PyObject *args, PyObject *keywd PyObject *ret = NULL; uint8_t *input, *output, *custom_dictionary; size_t length, output_length, custom_dictionary_length; + const char *comment = 0; BrotliParams::Mode mode = (BrotliParams::Mode) -1; int quality = -1; int lgwin = -1; @@ -125,19 +128,20 @@ static PyObject* brotli_compress(PyObject *self, PyObject *args, PyObject *keywd int ok; static const char *kwlist[] = { - "string", "mode", "quality", "lgwin", "lgblock", "dictionary", NULL}; + "string", "mode", "quality", "lgwin", "lgblock", "dictionary", "comment", NULL}; custom_dictionary = NULL; custom_dictionary_length = 0; - ok = PyArg_ParseTupleAndKeywords(args, keywds, "s#|O&O&O&O&s#:compress", + ok = PyArg_ParseTupleAndKeywords(args, keywds, "s#|O&O&O&O&s#s:compress", const_cast(kwlist), &input, &length, &mode_convertor, &mode, &quality_convertor, &quality, &lgwin_convertor, &lgwin, &lgblock_convertor, &lgblock, - &custom_dictionary, &custom_dictionary_length); + &custom_dictionary, &custom_dictionary_length, + &comment); if (!ok) return NULL; @@ -153,6 +157,7 @@ static PyObject* brotli_compress(PyObject *self, PyObject *args, PyObject *keywd params.lgwin = lgwin; if (lgblock != -1) params.lgblock = lgblock; + params.comment = comment; if (custom_dictionary_length == 0) { ok = BrotliCompressBuffer(params, length, input,