net/websockets/websocket_inflater.cc

   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:
  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
  35 WebSocketInflater::WebSocketInflater()
  36     : input_queue_(kDefaultInputIOBufferCapacity),
  37       output_buffer_(kDefaultBufferCapacity) {}
  38
  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
  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
  61 WebSocketInflater::~WebSocketInflater() {
  62   if (stream_) {
  63     inflateEnd(stream_.get());
  64     stream_.reset();
  65   }
  66 }
  67
  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
  85 bool WebSocketInflater::Finish() {
  86   return AddBytes("\x00\x00\xff\xff", 4);
  87 }
  88
  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
 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
 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
 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
 169 WebSocketInflater::OutputBuffer::OutputBuffer(size_t capacity)
 170     : capacity_(capacity),
 171       buffer_(capacity_ + 1),  // 1 for sentinel
 172       head_(0),
 173       tail_(0) {}
 174
 175 WebSocketInflater::OutputBuffer::~OutputBuffer() {}
 176
 177 size_t WebSocketInflater::OutputBuffer::Size() const {
 178   return (tail_ + buffer_.size() - head_) % buffer_.size();
 179 }
 180
 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
 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
 212 void WebSocketInflater::OutputBuffer::AdvanceHead(size_t advance) {
 213   DCHECK_LE(advance, Size());
 214   head_ = (head_ + advance) % buffer_.size();
 215 }
 216
 217 void WebSocketInflater::OutputBuffer::AdvanceTail(size_t advance) {
 218   DCHECK_LE(advance + Size(), capacity_);
 219   tail_ = (tail_ + advance) % buffer_.size();
 220 }
 221
 222 WebSocketInflater::InputQueue::InputQueue(size_t capacity)
 223     : capacity_(capacity), head_of_first_buffer_(0), tail_of_last_buffer_(0) {}
 224
 225 WebSocketInflater::InputQueue::~InputQueue() {}
 226
 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
 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
 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
 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