1 /* 2 * 3 * Copyright 2017 gRPC authors. 4 * 5 * Licensed under the Apache License, Version 2.0 (the "License"); 6 * you may not use this file except in compliance with the License. 7 * You may obtain a copy of the License at 8 * 9 * http://www.apache.org/licenses/LICENSE-2.0 10 * 11 * Unless required by applicable law or agreed to in writing, software 12 * distributed under the License is distributed on an "AS IS" BASIS, 13 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 * See the License for the specific language governing permissions and 15 * limitations under the License. 16 * 17 */ 18 19 #include <grpc/impl/codegen/byte_buffer.h> 20 #include <grpc/slice.h> 21 #include <grpcpp/impl/codegen/grpc_library.h> 22 #include <grpcpp/impl/codegen/proto_utils.h> 23 #include <grpcpp/impl/grpc_library.h> 24 #include <gtest/gtest.h> 25 26 namespace grpc { 27 28 namespace internal { 29 30 // Provide access to ProtoBufferWriter internals. 31 class ProtoBufferWriterPeer { 32 public: 33 explicit ProtoBufferWriterPeer(ProtoBufferWriter* writer) : writer_(writer) {} 34 bool have_backup() const { return writer_->have_backup_; } 35 const grpc_slice& backup_slice() const { return writer_->backup_slice_; } 36 const grpc_slice& slice() const { return writer_->slice_; } 37 38 private: 39 ProtoBufferWriter* writer_; 40 }; 41 42 // Provide access to ByteBuffer internals. 43 class GrpcByteBufferPeer { 44 public: 45 explicit GrpcByteBufferPeer(ByteBuffer* bb) : bb_(bb) {} 46 grpc_byte_buffer* c_buffer() { return bb_->c_buffer(); } 47 48 private: 49 ByteBuffer* bb_; 50 }; 51 52 class ProtoUtilsTest : public ::testing::Test {}; 53 54 // Regression test for a memory corruption bug where a series of 55 // ProtoBufferWriter Next()/Backup() invocations could result in a dangling 56 // pointer returned by Next() due to the interaction between grpc_slice inlining 57 // and GRPC_SLICE_START_PTR. 58 TEST_F(ProtoUtilsTest, TinyBackupThenNext) { 59 ByteBuffer bp; 60 const int block_size = 1024; 61 ProtoBufferWriter writer(&bp, block_size, 8192); 62 ProtoBufferWriterPeer peer(&writer); 63 64 void* data; 65 int size; 66 // Allocate a slice. 67 ASSERT_TRUE(writer.Next(&data, &size)); 68 EXPECT_EQ(block_size, size); 69 // Return a single byte. 70 writer.BackUp(1); 71 EXPECT_FALSE(peer.have_backup()); 72 // On the next allocation, the returned slice is non-inlined. 73 ASSERT_TRUE(writer.Next(&data, &size)); 74 EXPECT_TRUE(peer.slice().refcount != nullptr); 75 EXPECT_EQ(block_size, size); 76 } 77 78 namespace { 79 80 // Set backup_size to 0 to indicate no backup is needed. 81 void BufferWriterTest(int block_size, int total_size, int backup_size) { 82 ByteBuffer bb; 83 ProtoBufferWriter writer(&bb, block_size, total_size); 84 85 int written_size = 0; 86 void* data; 87 int size = 0; 88 bool backed_up_entire_slice = false; 89 90 while (written_size < total_size) { 91 EXPECT_TRUE(writer.Next(&data, &size)); 92 EXPECT_GT(size, 0); 93 EXPECT_TRUE(data); 94 int write_size = size; 95 bool should_backup = false; 96 if (backup_size > 0 && size > backup_size) { 97 write_size = size - backup_size; 98 should_backup = true; 99 } else if (size == backup_size && !backed_up_entire_slice) { 100 // only backup entire slice once. 101 backed_up_entire_slice = true; 102 should_backup = true; 103 write_size = 0; 104 } 105 // May need a last backup. 106 if (write_size + written_size > total_size) { 107 write_size = total_size - written_size; 108 should_backup = true; 109 backup_size = size - write_size; 110 ASSERT_GT(backup_size, 0); 111 } 112 for (int i = 0; i < write_size; i++) { 113 (static_cast<uint8_t*>(data))[i] = written_size % 128; 114 written_size++; 115 } 116 if (should_backup) { 117 writer.BackUp(backup_size); 118 } 119 } 120 EXPECT_EQ(bb.Length(), (size_t)total_size); 121 122 grpc_byte_buffer_reader reader; 123 GrpcByteBufferPeer peer(&bb); 124 grpc_byte_buffer_reader_init(&reader, peer.c_buffer()); 125 int read_bytes = 0; 126 while (read_bytes < total_size) { 127 grpc_slice s; 128 EXPECT_TRUE(grpc_byte_buffer_reader_next(&reader, &s)); 129 for (size_t i = 0; i < GRPC_SLICE_LENGTH(s); i++) { 130 EXPECT_EQ(GRPC_SLICE_START_PTR(s)[i], read_bytes % 128); 131 read_bytes++; 132 } 133 grpc_slice_unref(s); 134 } 135 EXPECT_EQ(read_bytes, total_size); 136 grpc_byte_buffer_reader_destroy(&reader); 137 } 138 139 TEST(WriterTest, TinyBlockTinyBackup) { 140 for (int i = 2; i < static_cast<int> GRPC_SLICE_INLINED_SIZE; i++) { 141 BufferWriterTest(i, 256, 1); 142 } 143 } 144 145 TEST(WriterTest, SmallBlockTinyBackup) { BufferWriterTest(64, 256, 1); } 146 147 TEST(WriterTest, SmallBlockNoBackup) { BufferWriterTest(64, 256, 0); } 148 149 TEST(WriterTest, SmallBlockFullBackup) { BufferWriterTest(64, 256, 64); } 150 151 TEST(WriterTest, LargeBlockTinyBackup) { BufferWriterTest(4096, 8192, 1); } 152 153 TEST(WriterTest, LargeBlockNoBackup) { BufferWriterTest(4096, 8192, 0); } 154 155 TEST(WriterTest, LargeBlockFullBackup) { BufferWriterTest(4096, 8192, 4096); } 156 157 TEST(WriterTest, LargeBlockLargeBackup) { BufferWriterTest(4096, 8192, 4095); } 158 159 } // namespace 160 } // namespace internal 161 } // namespace grpc 162 163 int main(int argc, char** argv) { 164 // Ensure the ProtoBufferWriter internals are initialized. 165 grpc::internal::GrpcLibraryInitializer init; 166 init.summon(); 167 grpc::GrpcLibraryCodegen lib; 168 169 ::testing::InitGoogleTest(&argc, argv); 170 return RUN_ALL_TESTS(); 171 } 172