Skip to content

Commit 9c788ee

Browse files
chrisbeardchrisbeard
authored
Fix[bmqp]: cap PUT decompression for v1 msg properties (bloomberg#1628)
Signed-off-by: Christopher Beard <cbeard9@bloomberg.net>
1 parent ec10f62 commit 9c788ee

7 files changed

Lines changed: 211 additions & 22 deletions

File tree

‎src/groups/bmq/bmqp/bmqp_compression.cpp‎

Lines changed: 46 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -114,16 +114,20 @@ struct ZLib {
114114

115115
/// Apply the operation given by the specified `zlibMethod` and
116116
/// `zlibEndMethod` on the specified `input` using the specified
117-
/// `stream`, and write the result to the specified `output`. Return 0
118-
/// on success and non-zero otherwise, in which case a message is
119-
/// written to the specified `errorStream` if it is non-zero.
117+
/// `stream`, and write the result to the specified `output`. If the
118+
/// specified `maxOutputSize` is non-zero, fail as soon as the
119+
/// accumulated output would exceed `maxOutputSize` bytes; a value
120+
/// of 0 means no limit is enforced. Return 0 on success and
121+
/// non-zero otherwise, in which case a message is written to the
122+
/// specified `errorStream` if it is non-zero.
120123
static int writeOutput(bdlbb::Blob* output,
121124
bdlbb::BlobBufferFactory* factory,
122125
z_stream* stream,
123126
bsl::ostream* errorStream,
124127
const bdlbb::Blob& input,
125128
ZlibStreamMethod zlibMethod,
126-
ZlibEndStreamMethod zlibEndMethod);
129+
ZlibEndStreamMethod zlibEndMethod,
130+
bsls::Types::Uint64 maxOutputSize);
127131
};
128132

129133
// ===========
@@ -226,13 +230,15 @@ int ZLib::writeOutput(bdlbb::Blob* output,
226230
bsl::ostream* errorStream,
227231
const bdlbb::Blob& input,
228232
ZlibStreamMethod zlibMethod,
229-
ZlibEndStreamMethod zlibEndMethod)
233+
ZlibEndStreamMethod zlibEndMethod,
234+
bsls::Types::Uint64 maxOutputSize)
230235
{
231236
enum RcEnum {
232237
rc_SUCCESS = 0,
233238
rc_STREAM_INIT_FAILURE = -1,
234239
rc_STREAM_PROCESS_FAILURE = -2,
235-
rc_STREAM_END_FAILURE = -3
240+
rc_STREAM_END_FAILURE = -3,
241+
rc_MAX_SIZE_EXCEEDED = -4
236242
};
237243

238244
bdlbb::BlobBuffer inBuffer;
@@ -258,6 +264,20 @@ int ZLib::writeOutput(bdlbb::Blob* output,
258264
zlibEndMethod(stream);
259265
return rc_STREAM_PROCESS_FAILURE; // RETURN
260266
}
267+
268+
// Fail if the accumulated output would exceed the cap. The current
269+
// 'outBuffer' has not been appended to 'output' yet, so account
270+
// for the bytes already written into it.
271+
if (maxOutputSize != 0 &&
272+
static_cast<bsls::Types::Uint64>(output->length()) +
273+
(outBuffer.size() - stream->avail_out) >
274+
maxOutputSize) {
275+
setError(errorStream,
276+
"Decompressed output exceeds maximum size",
277+
0);
278+
zlibEndMethod(stream);
279+
return rc_MAX_SIZE_EXCEEDED; // RETURN
280+
}
261281
}
262282

263283
// Continue to write output data until the stream reaches its end, or the
@@ -269,6 +289,19 @@ int ZLib::writeOutput(bdlbb::Blob* output,
269289
advanceOutput(output, &outBuffer, factory, stream);
270290
lastSize = stream->avail_out;
271291
result = zlibMethod(stream, Z_FINISH);
292+
293+
// Fail closed if the accumulated output would exceed the cap (see the
294+
// explanation above).
295+
if (maxOutputSize != 0 &&
296+
static_cast<bsls::Types::Uint64>(output->length()) +
297+
(outBuffer.size() - stream->avail_out) >
298+
maxOutputSize) {
299+
setError(errorStream,
300+
"Decompressed output exceeds maximum size",
301+
0);
302+
zlibEndMethod(stream);
303+
return rc_MAX_SIZE_EXCEEDED; // RETURN
304+
}
272305
} while ((Z_BUF_ERROR == result || Z_OK == result) &&
273306
lastSize != stream->avail_out);
274307

@@ -386,6 +419,7 @@ int Compression::decompress(bdlbb::Blob* output,
386419
bdlbb::BlobBufferFactory* factory,
387420
bmqt::CompressionAlgorithmType::Enum algorithm,
388421
const bdlbb::Blob& input,
422+
bsls::Types::Uint64 maxOutputSize,
389423
bsl::ostream* errorStream,
390424
bslma::Allocator* allocator)
391425
{
@@ -396,6 +430,7 @@ int Compression::decompress(bdlbb::Blob* output,
396430
return Compression_Impl::decompressZlib(output,
397431
factory,
398432
input,
433+
maxOutputSize,
399434
errorStream,
400435
allocator); // RETURN
401436
case bmqt::CompressionAlgorithmType::e_NONE:
@@ -449,12 +484,14 @@ int Compression_Impl::compressZlib(bdlbb::Blob* output,
449484
errorStream,
450485
input,
451486
&::deflate,
452-
&::deflateEnd);
487+
&::deflateEnd,
488+
0); // no output cap for compression
453489
}
454490

455491
int Compression_Impl::decompressZlib(bdlbb::Blob* output,
456492
bdlbb::BlobBufferFactory* factory,
457493
const bdlbb::Blob& input,
494+
bsls::Types::Uint64 maxOutputSize,
458495
bsl::ostream* errorStream,
459496
bslma::Allocator* allocator)
460497
{
@@ -481,7 +518,8 @@ int Compression_Impl::decompressZlib(bdlbb::Blob* output,
481518
errorStream,
482519
input,
483520
&::inflate,
484-
&::inflateEnd);
521+
&::inflateEnd,
522+
maxOutputSize);
485523
}
486524

487525
} // close package namespace

‎src/groups/bmq/bmqp/bmqp_compression.h‎

Lines changed: 20 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@
3737
#include <bdlbb_blob.h>
3838
#include <bsl_ostream.h>
3939
#include <bslma_allocator.h>
40+
#include <bsls_types.h>
4041

4142
namespace BloombergLP {
4243

@@ -88,16 +89,21 @@ struct Compression {
8889
/// specified `algorithm`, and load the uncompressed data into specified
8990
/// `output`, using the specified `factory` to supply the needed data
9091
/// buffers. Return 0 on success, and non-zero otherwise. Optionally
91-
/// specify an `errorStream` to record details on any errors that may
92-
/// occur during this operation. Also, optionally specify `allocator`
93-
/// which will be used to supply memory. Also note, that any existing
94-
/// data in the specified `output` will be preserved.
92+
/// specify a `maxOutputSize`, in bytes, beyond which decompression fails
93+
/// closed with a non-zero return code rather than continuing to allocate
94+
/// output buffers; a value of 0 (the default) means no limit is enforced.
95+
/// This bounds the memory consumed when decompressing highly-compressible
96+
/// input. Optionally specify an `errorStream` to record details
97+
/// on any errors that may occur during this operation. Also, optionally
98+
/// specify `allocator` which will be used to supply memory. Also note,
99+
/// that any existing data in the specified `output` will be preserved.
95100
static int decompress(bdlbb::Blob* output,
96101
bdlbb::BlobBufferFactory* factory,
97102
bmqt::CompressionAlgorithmType::Enum algorithm,
98103
const bdlbb::Blob& input,
99-
bsl::ostream* errorStream = 0,
100-
bslma::Allocator* allocator = 0);
104+
bsls::Types::Uint64 maxOutputSize = 0,
105+
bsl::ostream* errorStream = 0,
106+
bslma::Allocator* allocator = 0);
101107
};
102108

103109
// ======================
@@ -131,13 +137,17 @@ struct Compression_Impl {
131137
/// Decompress the data within the specified `input` as according to the
132138
/// Zlib algorithm, and load the uncompressed data into the specified
133139
/// `output` blob, using the specified `factory` to supply needed data
134-
/// buffers. Return 0 on success, and non-zero otherwise. Specify an
135-
/// `errorStream` to record details on any errors that may occur during
136-
/// this operation. Also, specify `allocator` which will be used to
137-
/// supply memory. Return 0 on success, and non-zero otherwise.
140+
/// buffers. Return 0 on success, and non-zero otherwise. Optionally
141+
/// specify a `maxOutputSize`, in bytes, beyond which decompression fails
142+
/// closed with a non-zero return code rather than continuing to allocate
143+
/// output buffers; a value of 0 (the default) means no limit is enforced.
144+
/// Specify an `errorStream` to record details on any errors that may
145+
/// occur during this operation. Also, specify `allocator` which will be
146+
/// used to supply memory. Return 0 on success, and non-zero otherwise.
138147
static int decompressZlib(bdlbb::Blob* output,
139148
bdlbb::BlobBufferFactory* factory,
140149
const bdlbb::Blob& input,
150+
bsls::Types::Uint64 maxOutputSize,
141151
bsl::ostream* errorStream,
142152
bslma::Allocator* allocator);
143153
};

‎src/groups/bmq/bmqp/bmqp_compression.t.cpp‎

Lines changed: 117 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -293,6 +293,7 @@ static void eZlibCompressDecompressHelper(
293293
&bufferFactory,
294294
algorithm,
295295
compressed,
296+
0, // no output cap
296297
&error,
297298
bmqtst::TestHelperUtil::allocator());
298299
*decompressionTime = bsls::TimeUtil::getTimer() - startTime;
@@ -336,6 +337,7 @@ eZlibCompressDecompressHelper(bsls::Types::Int64* compressionTime,
336337
&decompressed,
337338
&bufferFactory,
338339
compressed,
340+
0, // no output cap
339341
&error,
340342
bmqtst::TestHelperUtil::allocator());
341343
*decompressionTime = bsls::TimeUtil::getTimer() - startTime;
@@ -377,6 +379,7 @@ static void eZlibCompressionRatioHelper(bsls::Types::Int64* inputSize,
377379
&decompressed,
378380
&bufferFactory,
379381
compressed,
382+
0, // no output cap
380383
&error,
381384
bmqtst::TestHelperUtil::allocator());
382385
BMQTST_ASSERT_EQ(rc, 0);
@@ -544,6 +547,7 @@ static void test1_breathingTest()
544547
&decompressed,
545548
&bufferFactory,
546549
compressed,
550+
0, // no output cap
547551
&error,
548552
bmqtst::TestHelperUtil::allocator());
549553
BMQTST_ASSERT_EQ(rc, 0);
@@ -613,6 +617,7 @@ static void test1_breathingTest()
613617
&decompressed,
614618
&bufferFactory,
615619
compressed,
620+
0, // no output cap
616621
&error,
617622
bmqtst::TestHelperUtil::allocator());
618623
BMQTST_ASSERT_EQ(rc, 0);
@@ -692,6 +697,7 @@ static void test2_compression_cluster_message()
692697
&decompressed,
693698
&bufferFactory,
694699
compressed,
700+
0, // no output cap
695701
&error,
696702
bmqtst::TestHelperUtil::allocator());
697703
BMQTST_ASSERT_EQ(rc, 0);
@@ -752,6 +758,116 @@ static void test3_compression_decompression_none()
752758
}
753759
}
754760

761+
static void test4_decompressionSizeLimit()
762+
// ------------------------------------------------------------------------
763+
// DECOMPRESSION OUTPUT SIZE LIMIT
764+
//
765+
// Concerns:
766+
// A highly-compressible input must not be allowed to expand without
767+
// bound. When a maximum output size is supplied, decompression must fail
768+
// closed once the accumulated output would exceed the cap, rather than
769+
// continuing to allocate memory.
770+
//
771+
// Plan:
772+
// - Compress a large run of zero bytes (which zlib shrinks dramatically).
773+
// - Decompress with no cap and confirm the full round-trip (control).
774+
// - Decompress with a cap below the true output size and confirm it fails
775+
// with a non-zero code while producing output bounded near the cap (not
776+
// the full expansion).
777+
// - Decompress with a cap above the true output size and confirm success.
778+
//
779+
// Testing:
780+
// bmqp::Compression::decompress with a non-zero maxOutputSize
781+
// ------------------------------------------------------------------------
782+
{
783+
bmqtst::TestHelper::printTestName("DECOMPRESSION SIZE LIMIT TEST");
784+
785+
// Use small blob buffers so the enforced cap is tight (the overrun beyond
786+
// the cap is bounded by a single buffer).
787+
bdlbb::PooledBlobBufferFactory bufferFactory(
788+
1024,
789+
bmqtst::TestHelperUtil::allocator());
790+
791+
bmqu::MemOutStream error(bmqtst::TestHelperUtil::allocator());
792+
793+
// Build a large, highly-compressible input (all zero bytes) that expands
794+
// far beyond the cap we will impose.
795+
const int k_INPUT_SIZE = 8 * 1024 * 1024; // 8 MB
796+
bsl::string zeros(k_INPUT_SIZE, '\0', bmqtst::TestHelperUtil::allocator());
797+
798+
bdlbb::Blob input(&bufferFactory, bmqtst::TestHelperUtil::allocator());
799+
bdlbb::BlobUtil::append(&input, zeros.data(), k_INPUT_SIZE);
800+
801+
bdlbb::Blob compressed(&bufferFactory,
802+
bmqtst::TestHelperUtil::allocator());
803+
int rc = bmqp::Compression::compress(
804+
&compressed,
805+
&bufferFactory,
806+
bmqt::CompressionAlgorithmType::e_ZLIB,
807+
input,
808+
&error,
809+
bmqtst::TestHelperUtil::allocator());
810+
BMQTST_ASSERT_EQ(rc, 0);
811+
// Highly-compressible: compressed form is a tiny fraction of the input.
812+
BMQTST_ASSERT_LT(compressed.length(), k_INPUT_SIZE);
813+
814+
{
815+
PVV("Control: no cap decompresses fully");
816+
bdlbb::Blob decompressed(&bufferFactory,
817+
bmqtst::TestHelperUtil::allocator());
818+
rc = bmqp::Compression::decompress(
819+
&decompressed,
820+
&bufferFactory,
821+
bmqt::CompressionAlgorithmType::e_ZLIB,
822+
compressed,
823+
0, // no cap
824+
&error,
825+
bmqtst::TestHelperUtil::allocator());
826+
BMQTST_ASSERT_EQ(rc, 0);
827+
BMQTST_ASSERT_EQ(decompressed.length(), k_INPUT_SIZE);
828+
}
829+
830+
{
831+
PVV("Cap below decompressed size fails closed");
832+
const bsls::Types::Uint64 k_CAP = 1024 * 1024; // 1 MB < 8 MB
833+
bdlbb::Blob decompressed(&bufferFactory,
834+
bmqtst::TestHelperUtil::allocator());
835+
rc = bmqp::Compression::decompress(
836+
&decompressed,
837+
&bufferFactory,
838+
bmqt::CompressionAlgorithmType::e_ZLIB,
839+
compressed,
840+
k_CAP,
841+
&error,
842+
bmqtst::TestHelperUtil::allocator());
843+
// Fails with a non-zero code ...
844+
BMQTST_ASSERT_NE(rc, 0);
845+
// ... and stops well short of fully expanding the input, bounded
846+
// near the cap rather than the full 8 MB.
847+
BMQTST_ASSERT_LT(decompressed.length(), k_INPUT_SIZE);
848+
BMQTST_ASSERT_LE(
849+
static_cast<bsls::Types::Uint64>(decompressed.length()),
850+
k_CAP + 1024);
851+
}
852+
853+
{
854+
PVV("Cap above decompressed size succeeds");
855+
const bsls::Types::Uint64 k_CAP = 64 * 1024 * 1024; // 64 MB > 8 MB
856+
bdlbb::Blob decompressed(&bufferFactory,
857+
bmqtst::TestHelperUtil::allocator());
858+
rc = bmqp::Compression::decompress(
859+
&decompressed,
860+
&bufferFactory,
861+
bmqt::CompressionAlgorithmType::e_ZLIB,
862+
compressed,
863+
k_CAP,
864+
&error,
865+
bmqtst::TestHelperUtil::allocator());
866+
BMQTST_ASSERT_EQ(rc, 0);
867+
BMQTST_ASSERT_EQ(decompressed.length(), k_INPUT_SIZE);
868+
}
869+
}
870+
755871
// ============================================================================
756872
// PERFORMANCE TESTS
757873
// ----------------------------------------------------------------------------
@@ -1113,6 +1229,7 @@ int main(int argc, char* argv[])
11131229
case 1: test1_breathingTest(); break;
11141230
case 2: test2_compression_cluster_message(); break;
11151231
case 3: test3_compression_decompression_none(); break;
1232+
case 4: test4_decompressionSizeLimit(); break;
11161233
case -1:
11171234
BMQTST_BENCHMARK_WITH_ARGS(
11181235
testN1_performanceCompressionDecompressionDefault,

‎src/groups/bmq/bmqp/bmqp_protocolutil.cpp‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -457,6 +457,7 @@ int ProtocolUtil::convertToOld(bdlbb::Blob* dst,
457457
factory,
458458
cat,
459459
compressed,
460+
0, // no output cap
460461
&error,
461462
allocator);
462463
}
@@ -518,7 +519,8 @@ int ProtocolUtil::parse(bdlbb::Blob* messagePropertiesOutput,
518519
bool haveNewMessageProperties,
519520
bmqt::CompressionAlgorithmType::Enum cat,
520521
bdlbb::BlobBufferFactory* blobBufferFactory,
521-
bslma::Allocator* allocator)
522+
bslma::Allocator* allocator,
523+
bsls::Types::Uint64 maxDecompressedSize)
522524
{
523525
// This is to capture parsing and de-compressing MessageProperties in a
524526
// single place.
@@ -619,6 +621,7 @@ int ProtocolUtil::parse(bdlbb::Blob* messagePropertiesOutput,
619621
blobBufferFactory,
620622
cat,
621623
bufferCompressed,
624+
maxDecompressedSize,
622625
&error,
623626
allocator);
624627

0 commit comments

Comments
 (0)