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