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
15namespace net {
16
17namespace {
18
19class ShrinkableIOBufferWithSize : public IOBufferWithSize {
20 public:
21  explicit ShrinkableIOBufferWithSize(int size)
22      : IOBufferWithSize(size) {}
23
24  void Shrink(int new_size) {
25    DCHECK_LE(new_size, size_);
26    size_ = new_size;
27  }
28
29 private:
30  virtual ~ShrinkableIOBufferWithSize() {}
31};
32
33}  // namespace
34
35WebSocketInflater::WebSocketInflater()
36    : input_queue_(kDefaultInputIOBufferCapacity),
37      output_buffer_(kDefaultBufferCapacity) {}
38
39WebSocketInflater::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
47bool 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
61WebSocketInflater::~WebSocketInflater() {
62  if (stream_) {
63    inflateEnd(stream_.get());
64    stream_.reset();
65  }
66}
67
68bool 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
85bool WebSocketInflater::Finish() {
86  return AddBytes("\x00\x00\xff\xff", 4);
87}
88
89scoped_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
107int 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
121int 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
147int 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
169WebSocketInflater::OutputBuffer::OutputBuffer(size_t capacity)
170    : capacity_(capacity),
171      buffer_(capacity_ + 1),  // 1 for sentinel
172      head_(0),
173      tail_(0) {}
174
175WebSocketInflater::OutputBuffer::~OutputBuffer() {}
176
177size_t WebSocketInflater::OutputBuffer::Size() const {
178  return (tail_ + buffer_.size() - head_) % buffer_.size();
179}
180
181std::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
187void 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
212void WebSocketInflater::OutputBuffer::AdvanceHead(size_t advance) {
213  DCHECK_LE(advance, Size());
214  head_ = (head_ + advance) % buffer_.size();
215}
216
217void WebSocketInflater::OutputBuffer::AdvanceTail(size_t advance) {
218  DCHECK_LE(advance + Size(), capacity_);
219  tail_ = (tail_ + advance) % buffer_.size();
220}
221
222WebSocketInflater::InputQueue::InputQueue(size_t capacity)
223    : capacity_(capacity), head_of_first_buffer_(0), tail_of_last_buffer_(0) {}
224
225WebSocketInflater::InputQueue::~InputQueue() {}
226
227std::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
237void 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
255void 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
271size_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