websocket_inflater.cc 8.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280
  1. // Copyright 2013 The Chromium Authors. All rights reserved.
  2. // Use of this source code is governed by a BSD-style license that can be
  3. // found in the LICENSE file.
  4. #include "net/websockets/websocket_inflater.h"
  5. #include <algorithm>
  6. #include <vector>
  7. #include "base/check_op.h"
  8. #include "net/base/io_buffer.h"
  9. #include "third_party/zlib/zlib.h"
  10. namespace net {
  11. namespace {
  12. class ShrinkableIOBufferWithSize : public IOBufferWithSize {
  13. public:
  14. explicit ShrinkableIOBufferWithSize(size_t size) : IOBufferWithSize(size) {}
  15. void Shrink(int new_size) {
  16. CHECK_GE(new_size, 0);
  17. CHECK_LE(new_size, size_);
  18. size_ = new_size;
  19. }
  20. private:
  21. ~ShrinkableIOBufferWithSize() override = default;
  22. };
  23. } // namespace
  24. WebSocketInflater::WebSocketInflater()
  25. : input_queue_(kDefaultInputIOBufferCapacity),
  26. output_buffer_(kDefaultBufferCapacity) {}
  27. WebSocketInflater::WebSocketInflater(size_t input_queue_capacity,
  28. size_t output_buffer_capacity)
  29. : input_queue_(input_queue_capacity),
  30. output_buffer_(output_buffer_capacity) {
  31. DCHECK_GT(input_queue_capacity, 0u);
  32. DCHECK_GT(output_buffer_capacity, 0u);
  33. }
  34. bool WebSocketInflater::Initialize(int window_bits) {
  35. DCHECK_LE(8, window_bits);
  36. DCHECK_GE(15, window_bits);
  37. stream_ = std::make_unique<z_stream>();
  38. memset(stream_.get(), 0, sizeof(*stream_));
  39. int result = inflateInit2(stream_.get(), -window_bits);
  40. if (result != Z_OK) {
  41. inflateEnd(stream_.get());
  42. stream_.reset();
  43. return false;
  44. }
  45. return true;
  46. }
  47. WebSocketInflater::~WebSocketInflater() {
  48. if (stream_) {
  49. inflateEnd(stream_.get());
  50. stream_.reset();
  51. }
  52. }
  53. bool WebSocketInflater::AddBytes(const char* data, size_t size) {
  54. if (!size)
  55. return true;
  56. if (!input_queue_.IsEmpty()) {
  57. // choked
  58. input_queue_.Push(data, size);
  59. return true;
  60. }
  61. int result = InflateWithFlush(data, size);
  62. if (stream_->avail_in > 0)
  63. input_queue_.Push(&data[size - stream_->avail_in], stream_->avail_in);
  64. return result == Z_OK || result == Z_BUF_ERROR;
  65. }
  66. bool WebSocketInflater::Finish() {
  67. return AddBytes("\x00\x00\xff\xff", 4);
  68. }
  69. scoped_refptr<IOBufferWithSize> WebSocketInflater::GetOutput(size_t size) {
  70. auto buffer = base::MakeRefCounted<ShrinkableIOBufferWithSize>(size);
  71. size_t num_bytes_copied = 0;
  72. while (num_bytes_copied < size && output_buffer_.Size() > 0) {
  73. size_t num_bytes_to_copy =
  74. std::min(output_buffer_.Size(), size - num_bytes_copied);
  75. output_buffer_.Read(&buffer->data()[num_bytes_copied], num_bytes_to_copy);
  76. num_bytes_copied += num_bytes_to_copy;
  77. int result = InflateChokedInput();
  78. if (result != Z_OK && result != Z_BUF_ERROR)
  79. return nullptr;
  80. }
  81. buffer->Shrink(num_bytes_copied);
  82. return buffer;
  83. }
  84. int WebSocketInflater::InflateWithFlush(const char* next_in, size_t avail_in) {
  85. int result = Inflate(next_in, avail_in, Z_NO_FLUSH);
  86. if (result != Z_OK && result != Z_BUF_ERROR)
  87. return result;
  88. if (CurrentOutputSize() > 0)
  89. return result;
  90. // CurrentOutputSize() == 0 means there is no data to be output,
  91. // so we should make sure it by using Z_SYNC_FLUSH.
  92. return Inflate(reinterpret_cast<const char*>(stream_->next_in),
  93. stream_->avail_in,
  94. Z_SYNC_FLUSH);
  95. }
  96. int WebSocketInflater::Inflate(const char* next_in,
  97. size_t avail_in,
  98. int flush) {
  99. stream_->next_in = reinterpret_cast<Bytef*>(const_cast<char*>(next_in));
  100. stream_->avail_in = avail_in;
  101. int result = Z_BUF_ERROR;
  102. do {
  103. std::pair<char*, size_t> tail = output_buffer_.GetTail();
  104. if (!tail.second)
  105. break;
  106. stream_->next_out = reinterpret_cast<Bytef*>(tail.first);
  107. stream_->avail_out = tail.second;
  108. result = inflate(stream_.get(), flush);
  109. output_buffer_.AdvanceTail(tail.second - stream_->avail_out);
  110. if (result == Z_STREAM_END) {
  111. // Received a block with BFINAL set to 1. Reset the decompression state.
  112. result = inflateReset(stream_.get());
  113. } else if (tail.second == stream_->avail_out) {
  114. break;
  115. }
  116. } while (result == Z_OK || result == Z_BUF_ERROR);
  117. return result;
  118. }
  119. int WebSocketInflater::InflateChokedInput() {
  120. if (input_queue_.IsEmpty())
  121. return InflateWithFlush(nullptr, 0);
  122. int result = Z_BUF_ERROR;
  123. while (!input_queue_.IsEmpty()) {
  124. std::pair<char*, size_t> top = input_queue_.Top();
  125. result = InflateWithFlush(top.first, top.second);
  126. input_queue_.Consume(top.second - stream_->avail_in);
  127. if (result != Z_OK && result != Z_BUF_ERROR)
  128. return result;
  129. if (stream_->avail_in > 0) {
  130. // There are some data which are not consumed.
  131. break;
  132. }
  133. }
  134. return result;
  135. }
  136. WebSocketInflater::OutputBuffer::OutputBuffer(size_t capacity)
  137. : capacity_(capacity),
  138. buffer_(capacity_ + 1) // 1 for sentinel
  139. {}
  140. WebSocketInflater::OutputBuffer::~OutputBuffer() = default;
  141. size_t WebSocketInflater::OutputBuffer::Size() const {
  142. return (tail_ + buffer_.size() - head_) % buffer_.size();
  143. }
  144. std::pair<char*, size_t> WebSocketInflater::OutputBuffer::GetTail() {
  145. DCHECK_LT(tail_, buffer_.size());
  146. return std::make_pair(&buffer_[tail_],
  147. std::min(capacity_ - Size(), buffer_.size() - tail_));
  148. }
  149. void WebSocketInflater::OutputBuffer::Read(char* dest, size_t size) {
  150. DCHECK_LE(size, Size());
  151. size_t num_bytes_copied = 0;
  152. if (tail_ < head_) {
  153. size_t num_bytes_to_copy = std::min(size, buffer_.size() - head_);
  154. DCHECK_LT(head_, buffer_.size());
  155. memcpy(&dest[num_bytes_copied], &buffer_[head_], num_bytes_to_copy);
  156. AdvanceHead(num_bytes_to_copy);
  157. num_bytes_copied += num_bytes_to_copy;
  158. }
  159. if (num_bytes_copied == size)
  160. return;
  161. DCHECK_LE(head_, tail_);
  162. size_t num_bytes_to_copy = size - num_bytes_copied;
  163. DCHECK_LE(num_bytes_to_copy, tail_ - head_);
  164. DCHECK_LT(head_, buffer_.size());
  165. memcpy(&dest[num_bytes_copied], &buffer_[head_], num_bytes_to_copy);
  166. AdvanceHead(num_bytes_to_copy);
  167. num_bytes_copied += num_bytes_to_copy;
  168. DCHECK_EQ(size, num_bytes_copied);
  169. return;
  170. }
  171. void WebSocketInflater::OutputBuffer::AdvanceHead(size_t advance) {
  172. DCHECK_LE(advance, Size());
  173. head_ = (head_ + advance) % buffer_.size();
  174. }
  175. void WebSocketInflater::OutputBuffer::AdvanceTail(size_t advance) {
  176. DCHECK_LE(advance + Size(), capacity_);
  177. tail_ = (tail_ + advance) % buffer_.size();
  178. }
  179. WebSocketInflater::InputQueue::InputQueue(size_t capacity)
  180. : capacity_(capacity) {}
  181. WebSocketInflater::InputQueue::~InputQueue() = default;
  182. std::pair<char*, size_t> WebSocketInflater::InputQueue::Top() {
  183. DCHECK(!IsEmpty());
  184. if (buffers_.size() == 1) {
  185. return std::make_pair(&buffers_.front()->data()[head_of_first_buffer_],
  186. tail_of_last_buffer_ - head_of_first_buffer_);
  187. }
  188. return std::make_pair(&buffers_.front()->data()[head_of_first_buffer_],
  189. capacity_ - head_of_first_buffer_);
  190. }
  191. void WebSocketInflater::InputQueue::Push(const char* data, size_t size) {
  192. if (!size)
  193. return;
  194. size_t num_copied_bytes = 0;
  195. if (!IsEmpty())
  196. num_copied_bytes += PushToLastBuffer(data, size);
  197. while (num_copied_bytes < size) {
  198. DCHECK(IsEmpty() || tail_of_last_buffer_ == capacity_);
  199. buffers_.push_back(base::MakeRefCounted<IOBufferWithSize>(capacity_));
  200. tail_of_last_buffer_ = 0;
  201. num_copied_bytes +=
  202. PushToLastBuffer(&data[num_copied_bytes], size - num_copied_bytes);
  203. }
  204. }
  205. void WebSocketInflater::InputQueue::Consume(size_t size) {
  206. DCHECK(!IsEmpty());
  207. DCHECK_LE(size + head_of_first_buffer_, capacity_);
  208. head_of_first_buffer_ += size;
  209. if (head_of_first_buffer_ == capacity_) {
  210. buffers_.pop_front();
  211. head_of_first_buffer_ = 0;
  212. }
  213. if (buffers_.size() == 1 && head_of_first_buffer_ == tail_of_last_buffer_) {
  214. buffers_.pop_front();
  215. head_of_first_buffer_ = 0;
  216. tail_of_last_buffer_ = 0;
  217. }
  218. }
  219. size_t WebSocketInflater::InputQueue::PushToLastBuffer(const char* data,
  220. size_t size) {
  221. DCHECK(!IsEmpty());
  222. size_t num_bytes_to_copy = std::min(size, capacity_ - tail_of_last_buffer_);
  223. if (!num_bytes_to_copy)
  224. return 0;
  225. IOBufferWithSize* buffer = buffers_.back().get();
  226. memcpy(&buffer->data()[tail_of_last_buffer_], data, num_bytes_to_copy);
  227. tail_of_last_buffer_ += num_bytes_to_copy;
  228. return num_bytes_to_copy;
  229. }
  230. } // namespace net