1 // Protocol Buffers - Google's data interchange format 2 // Copyright 2023 Google Inc. All rights reserved. 3 // 4 // Use of this source code is governed by a BSD-style 5 // license that can be found in the LICENSE file or at 6 // https://developers.google.com/open-source/licenses/bsd 7 8 #ifndef GOOGLE_PROTOBUF_IO_TEST_ZERO_COPY_STREAM_H__ 9 #define GOOGLE_PROTOBUF_IO_TEST_ZERO_COPY_STREAM_H__ 10 11 #include <deque> 12 #include <initializer_list> 13 #include <memory> 14 #include <string> 15 #include <utility> 16 #include <vector> 17 18 #include "absl/log/absl_check.h" 19 #include "google/protobuf/io/zero_copy_stream.h" 20 21 // Must be included last. 22 #include "google/protobuf/port_def.inc" 23 24 namespace google { 25 namespace protobuf { 26 namespace io { 27 namespace internal { 28 29 // Input stream used for testing the proper handling of input fragmentation. 30 // It also will assert the preconditions of the methods. 31 class TestZeroCopyInputStream final : public ZeroCopyInputStream { 32 public: 33 // The input stream will provide the buffers exactly as passed here. TestZeroCopyInputStream(std::initializer_list<std::string> buffers)34 TestZeroCopyInputStream(std::initializer_list<std::string> buffers) 35 : buffers_(buffers.begin(), buffers.end()) {} TestZeroCopyInputStream(const std::vector<std::string> & buffers)36 explicit TestZeroCopyInputStream(const std::vector<std::string>& buffers) 37 : buffers_(buffers.begin(), buffers.end()) {} 38 39 // Allow copy to be able to inspect the stream without consuming the original. TestZeroCopyInputStream(const TestZeroCopyInputStream & other)40 TestZeroCopyInputStream(const TestZeroCopyInputStream& other) 41 : ZeroCopyInputStream(), 42 buffers_(other.buffers_), 43 last_returned_buffer_( 44 other.last_returned_buffer_ 45 ? std::make_unique<std::string>(*other.last_returned_buffer_) 46 : nullptr), 47 byte_count_(other.byte_count_) {} 48 Next(const void ** data,int * size)49 bool Next(const void** data, int* size) override { 50 ABSL_CHECK(data) << "data must not be null"; 51 ABSL_CHECK(size) << "size must not be null"; 52 last_returned_buffer_ = nullptr; 53 54 // We are done 55 if (buffers_.empty()) return false; 56 57 last_returned_buffer_ = 58 std::make_unique<std::string>(std::move(buffers_.front())); 59 buffers_.pop_front(); 60 *data = last_returned_buffer_->data(); 61 *size = static_cast<int>(last_returned_buffer_->size()); 62 byte_count_ += *size; 63 return true; 64 } 65 BackUp(int count)66 void BackUp(int count) override { 67 ABSL_CHECK_GE(count, 0) << "count must not be negative"; 68 ABSL_CHECK(last_returned_buffer_ != nullptr) 69 << "The last call was not a successful Next()"; 70 ABSL_CHECK_LE(count, last_returned_buffer_->size()) 71 << "count must be within bounds of last buffer"; 72 buffers_.push_front( 73 last_returned_buffer_->substr(last_returned_buffer_->size() - count)); 74 last_returned_buffer_ = nullptr; 75 byte_count_ -= count; 76 } 77 Skip(int count)78 bool Skip(int count) override { 79 ABSL_CHECK_GE(count, 0) << "count must not be negative"; 80 last_returned_buffer_ = nullptr; 81 while (true) { 82 if (count == 0) return true; 83 if (buffers_.empty()) return false; 84 auto& front = buffers_[0]; 85 int front_size = static_cast<int>(front.size()); 86 if (front_size <= count) { 87 count -= front_size; 88 byte_count_ += front_size; 89 buffers_.pop_front(); 90 } else { 91 // The front is enough, just chomp from it. 92 front = front.substr(count); 93 byte_count_ += count; 94 return true; 95 } 96 } 97 } 98 ByteCount()99 int64_t ByteCount() const override { return byte_count_; } 100 101 private: 102 // For simplicity of implementation, we pop the elements from the from and 103 // move them to `last_returned_buffer_`. It makes it simpler to keep track of 104 // the state of the object. The extra cost is not relevant for testing. 105 std::deque<std::string> buffers_; 106 // absl::optional could work here, but std::unique_ptr makes it more likely 107 // for sanitizers to detect if the string is used after it is destroyed. 108 std::unique_ptr<std::string> last_returned_buffer_; 109 int64_t byte_count_ = 0; 110 }; 111 112 } // namespace internal 113 } // namespace io 114 } // namespace protobuf 115 } // namespace google 116 117 #include "google/protobuf/port_undef.inc" 118 119 #endif // GOOGLE_PROTOBUF_IO_TEST_ZERO_COPY_STREAM_H__ 120