Increased message_compress.c coverage

2 files changed, 58 insertions, 1 deletions

diff --git a/src/core/compression/message_compress.c b/src/core/compression/message_compress.c
index 209c1f0ff1..a723c47819 100644
--- a/src/core/compression/message_compress.c
+++ b/src/core/compression/message_compress.c
@@ -91,6 +91,14 @@ error:
   return 0;
 }
 
+static void *zalloc_gpr(void* opaque, unsigned int items, unsigned int size) {
+  return gpr_malloc(items * size);
+}
+
+static void zfree_gpr(void* opaque, void *address) {
+  gpr_free(address);
+}
+
 static int zlib_compress(gpr_slice_buffer* input, gpr_slice_buffer* output,
                          int gzip) {
   z_stream zs;
@@ -99,6 +107,8 @@ static int zlib_compress(gpr_slice_buffer* input, gpr_slice_buffer* output,
   size_t count_before = output->count;
   size_t length_before = output->length;
   memset(&zs, 0, sizeof(zs));
+  zs.zalloc = zalloc_gpr;
+  zs.zfree = zfree_gpr;
   r = deflateInit2(&zs, Z_DEFAULT_COMPRESSION, Z_DEFLATED, 15 | (gzip ? 16 : 0),
                    8, Z_DEFAULT_STRATEGY);
   if (r != Z_OK) {
@@ -125,6 +135,8 @@ static int zlib_decompress(gpr_slice_buffer* input, gpr_slice_buffer* output,
   size_t count_before = output->count;
   size_t length_before = output->length;
   memset(&zs, 0, sizeof(zs));
+  zs.zalloc = zalloc_gpr;
+  zs.zfree = zfree_gpr;
   r = inflateInit2(&zs, 15 | (gzip ? 16 : 0));
   if (r != Z_OK) {
     gpr_log(GPR_ERROR, "inflateInit2 returns %d", r);
@@ -150,7 +162,7 @@ static int copy(gpr_slice_buffer* input, gpr_slice_buffer* output) {
   return 1;
 }
 
-int compress_inner(grpc_compression_algorithm algorithm,
+static int compress_inner(grpc_compression_algorithm algorithm,
                    gpr_slice_buffer* input, gpr_slice_buffer* output) {
   switch (algorithm) {
     case GRPC_COMPRESS_NONE:
diff --git a/test/core/compression/message_compress_test.c b/test/core/compression/message_compress_test.c
index 4b45d22885..70359a3f69 100644
--- a/test/core/compression/message_compress_test.c
+++ b/test/core/compression/message_compress_test.c
@@ -168,6 +168,49 @@ static void test_bad_data(void) {
   gpr_slice_buffer_destroy(&output);
 }
 
+static void test_bad_compression_algorithm(void) {
+  gpr_slice_buffer input;
+  gpr_slice_buffer output;
+  int was_compressed;
+
+  gpr_slice_buffer_init(&input);
+  gpr_slice_buffer_init(&output);
+  gpr_slice_buffer_add(&input, gpr_slice_from_copied_string(
+                                   "Never gonna give you up"));
+  was_compressed =
+      grpc_msg_compress(GRPC_COMPRESS_ALGORITHMS_COUNT, &input, &output);
+  GPR_ASSERT(0 == was_compressed);
+
+  was_compressed =
+      grpc_msg_compress(GRPC_COMPRESS_ALGORITHMS_COUNT + 123, &input, &output);
+  GPR_ASSERT(0 == was_compressed);
+
+  gpr_slice_buffer_destroy(&input);
+  gpr_slice_buffer_destroy(&output);
+}
+
+static void test_bad_decompression_algorithm(void) {
+  gpr_slice_buffer input;
+  gpr_slice_buffer output;
+  int was_decompressed;
+
+  gpr_slice_buffer_init(&input);
+  gpr_slice_buffer_init(&output);
+  gpr_slice_buffer_add(&input,
+                       gpr_slice_from_copied_string(
+                           "I'm not really compressed but it doesn't matter"));
+  was_decompressed =
+      grpc_msg_decompress(GRPC_COMPRESS_ALGORITHMS_COUNT, &input, &output);
+  GPR_ASSERT(0 == was_decompressed);
+
+  was_decompressed =
+      grpc_msg_decompress(GRPC_COMPRESS_ALGORITHMS_COUNT + 123, &input, &output);
+  GPR_ASSERT(0 == was_decompressed);
+
+  gpr_slice_buffer_destroy(&input);
+  gpr_slice_buffer_destroy(&output);
+}
+
 int main(int argc, char **argv) {
   unsigned i, j, k, m;
   grpc_slice_split_mode uncompressed_split_modes[] = {
@@ -192,6 +235,8 @@ int main(int argc, char **argv) {
   }
 
   test_bad_data();
+  test_bad_compression_algorithm();
+  test_bad_decompression_algorithm();
   grpc_shutdown();
 
   return 0;