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 284
