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