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 <gtest/gtest.h>
20
21 #include <grpc/byte_buffer.h>
22 #include <grpc/slice.h>
23 #include <grpcpp/impl/grpc_library.h>
24 #include <grpcpp/impl/proto_utils.h>
25
26 #include "test/core/util/test_config.h"
27
28 namespace grpc {
29
30 namespace internal {
31
32 // Provide access to ProtoBufferWriter internals.
33 class ProtoBufferWriterPeer {
34 public:
ProtoBufferWriterPeer(ProtoBufferWriter * writer)35 explicit ProtoBufferWriterPeer(ProtoBufferWriter* writer) : writer_(writer) {}
have_backup() const36 bool have_backup() const { return writer_->have_backup_; }
backup_slice() const37 const grpc_slice& backup_slice() const { return writer_->backup_slice_; }
slice() const38 const grpc_slice& slice() const { return writer_->slice_; }
39
40 private:
41 ProtoBufferWriter* writer_;
42 };
43
44 // Provide access to ByteBuffer internals.
45 class GrpcByteBufferPeer {
46 public:
GrpcByteBufferPeer(ByteBuffer * bb)47 explicit GrpcByteBufferPeer(ByteBuffer* bb) : bb_(bb) {}
c_buffer()48 grpc_byte_buffer* c_buffer() { return bb_->c_buffer(); }
49
50 private:
51 ByteBuffer* bb_;
52 };
53
54 class ProtoUtilsTest : public ::testing::Test {
55 protected:
SetUpTestSuite()56 static void SetUpTestSuite() {
57 // Ensure the ProtoBufferWriter internals are initialized.
58 grpc::internal::GrpcLibrary lib;
59 grpc_init();
60 }
61
TearDownTestSuite()62 static void TearDownTestSuite() { grpc_shutdown(); }
63 };
64
65 // Regression test for a memory corruption bug where a series of
66 // ProtoBufferWriter Next()/Backup() invocations could result in a dangling
67 // pointer returned by Next() due to the interaction between grpc_slice inlining
68 // and GRPC_SLICE_START_PTR.
TEST_F(ProtoUtilsTest,TinyBackupThenNext)69 TEST_F(ProtoUtilsTest, TinyBackupThenNext) {
70 ByteBuffer bp;
71 const int block_size = 1024;
72 ProtoBufferWriter writer(&bp, block_size, 8192);
73 ProtoBufferWriterPeer peer(&writer);
74
75 void* data;
76 int size;
77 // Allocate a slice.
78 ASSERT_TRUE(writer.Next(&data, &size));
79 EXPECT_EQ(block_size, size);
80 // Return a single byte.
81 writer.BackUp(1);
82 EXPECT_FALSE(peer.have_backup());
83 // On the next allocation, the returned slice is non-inlined.
84 ASSERT_TRUE(writer.Next(&data, &size));
85 EXPECT_TRUE(peer.slice().refcount != nullptr);
86 EXPECT_EQ(block_size, size);
87 }
88
89 namespace {
90
91 // Set backup_size to 0 to indicate no backup is needed.
BufferWriterTest(int block_size,int total_size,int backup_size)92 void BufferWriterTest(int block_size, int total_size, int backup_size) {
93 ByteBuffer bb;
94 ProtoBufferWriter writer(&bb, block_size, total_size);
95
96 int written_size = 0;
97 void* data;
98 int size = 0;
99 bool backed_up_entire_slice = false;
100
101 while (written_size < total_size) {
102 EXPECT_TRUE(writer.Next(&data, &size));
103 EXPECT_GT(size, 0);
104 EXPECT_TRUE(data);
105 int write_size = size;
106 bool should_backup = false;
107 if (backup_size > 0 && size > backup_size) {
108 write_size = size - backup_size;
109 should_backup = true;
110 } else if (size == backup_size && !backed_up_entire_slice) {
111 // only backup entire slice once.
112 backed_up_entire_slice = true;
113 should_backup = true;
114 write_size = 0;
115 }
116 // May need a last backup.
117 if (write_size + written_size > total_size) {
118 write_size = total_size - written_size;
119 should_backup = true;
120 backup_size = size - write_size;
121 ASSERT_GT(backup_size, 0);
122 }
123 for (int i = 0; i < write_size; i++) {
124 (static_cast<uint8_t*>(data))[i] = written_size % 128;
125 written_size++;
126 }
127 if (should_backup) {
128 writer.BackUp(backup_size);
129 }
130 }
131 EXPECT_EQ(bb.Length(), (size_t)total_size);
132
133 grpc_byte_buffer_reader reader;
134 GrpcByteBufferPeer peer(&bb);
135 grpc_byte_buffer_reader_init(&reader, peer.c_buffer());
136 int read_bytes = 0;
137 while (read_bytes < total_size) {
138 grpc_slice s;
139 EXPECT_TRUE(grpc_byte_buffer_reader_next(&reader, &s));
140 for (size_t i = 0; i < GRPC_SLICE_LENGTH(s); i++) {
141 EXPECT_EQ(GRPC_SLICE_START_PTR(s)[i], read_bytes % 128);
142 read_bytes++;
143 }
144 grpc_slice_unref(s);
145 }
146 EXPECT_EQ(read_bytes, total_size);
147 grpc_byte_buffer_reader_destroy(&reader);
148 }
149
150 class WriterTest : public ::testing::Test {
151 protected:
SetUpTestSuite()152 static void SetUpTestSuite() {
153 grpc::internal::GrpcLibrary lib;
154 // Ensure the ProtoBufferWriter internals are initialized.
155 grpc_init();
156 }
157
TearDownTestSuite()158 static void TearDownTestSuite() { grpc_shutdown(); }
159 };
160
TEST_F(WriterTest,TinyBlockTinyBackup)161 TEST_F(WriterTest, TinyBlockTinyBackup) {
162 for (int i = 2; i < static_cast<int> GRPC_SLICE_INLINED_SIZE; i++) {
163 BufferWriterTest(i, 256, 1);
164 }
165 }
166
TEST_F(WriterTest,SmallBlockTinyBackup)167 TEST_F(WriterTest, SmallBlockTinyBackup) { BufferWriterTest(64, 256, 1); }
168
TEST_F(WriterTest,SmallBlockNoBackup)169 TEST_F(WriterTest, SmallBlockNoBackup) { BufferWriterTest(64, 256, 0); }
170
TEST_F(WriterTest,SmallBlockFullBackup)171 TEST_F(WriterTest, SmallBlockFullBackup) { BufferWriterTest(64, 256, 64); }
172
TEST_F(WriterTest,LargeBlockTinyBackup)173 TEST_F(WriterTest, LargeBlockTinyBackup) { BufferWriterTest(4096, 8192, 1); }
174
TEST_F(WriterTest,LargeBlockNoBackup)175 TEST_F(WriterTest, LargeBlockNoBackup) { BufferWriterTest(4096, 8192, 0); }
176
TEST_F(WriterTest,LargeBlockFullBackup)177 TEST_F(WriterTest, LargeBlockFullBackup) { BufferWriterTest(4096, 8192, 4096); }
178
TEST_F(WriterTest,LargeBlockLargeBackup)179 TEST_F(WriterTest, LargeBlockLargeBackup) {
180 BufferWriterTest(4096, 8192, 4095);
181 }
182
183 } // namespace
184 } // namespace internal
185 } // namespace grpc
186
main(int argc,char ** argv)187 int main(int argc, char** argv) {
188 grpc::testing::TestEnvironment env(&argc, argv);
189 ::testing::InitGoogleTest(&argc, argv);
190 return RUN_ALL_TESTS();
191 }
192