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_deflater.h"
6
7#include <string.h>
8#include <algorithm>
9#include <deque>
10#include <vector>
11
12#include "base/logging.h"
13#include "net/base/io_buffer.h"
14#include "third_party/zlib/zlib.h"
15
16namespace net {
17
18WebSocketDeflater::WebSocketDeflater(ContextTakeOverMode mode)
19    : mode_(mode), are_bytes_added_(false) {}
20
21WebSocketDeflater::~WebSocketDeflater() {
22  if (stream_) {
23    deflateEnd(stream_.get());
24    stream_.reset(NULL);
25  }
26}
27
28bool WebSocketDeflater::Initialize(int window_bits) {
29  DCHECK(!stream_);
30  stream_.reset(new z_stream);
31
32  DCHECK_LE(8, window_bits);
33  DCHECK_GE(15, window_bits);
34  memset(stream_.get(), 0, sizeof(*stream_));
35  int result = deflateInit2(stream_.get(),
36                            Z_DEFAULT_COMPRESSION,
37                            Z_DEFLATED,
38                            -window_bits,  // Negative value for raw deflate
39                            8,  // default mem level
40                            Z_DEFAULT_STRATEGY);
41  if (result != Z_OK) {
42    deflateEnd(stream_.get());
43    stream_.reset();
44    return false;
45  }
46  const size_t kFixedBufferSize = 4096;
47  fixed_buffer_.resize(kFixedBufferSize);
48  return true;
49}
50
51bool WebSocketDeflater::AddBytes(const char* data, size_t size) {
52  if (!size)
53    return true;
54
55  are_bytes_added_ = true;
56  stream_->next_in = reinterpret_cast<Bytef*>(const_cast<char*>(data));
57  stream_->avail_in = size;
58
59  int result = Deflate(Z_NO_FLUSH);
60  DCHECK(result != Z_BUF_ERROR || !stream_->avail_in);
61  return result == Z_BUF_ERROR;
62}
63
64bool WebSocketDeflater::Finish() {
65  if (!are_bytes_added_) {
66    // Since consecutive calls of deflate with Z_SYNC_FLUSH and no input
67    // lead to an error, we create and return the output for the empty input
68    // manually.
69    buffer_.push_back('\x02');
70    buffer_.push_back('\x00');
71    ResetContext();
72    return true;
73  }
74  stream_->next_in = NULL;
75  stream_->avail_in = 0;
76
77  int result = Deflate(Z_SYNC_FLUSH);
78  // Deflate returning Z_BUF_ERROR means that it's successfully flushed and
79  // blocked for input data.
80  if (result != Z_BUF_ERROR) {
81    ResetContext();
82    return false;
83  }
84  // Remove 4 octets from the tail as the specification requires.
85  if (CurrentOutputSize() < 4) {
86    ResetContext();
87    return false;
88  }
89  buffer_.resize(buffer_.size() - 4);
90  ResetContext();
91  return true;
92}
93
94void WebSocketDeflater::PushSyncMark() {
95  DCHECK(!are_bytes_added_);
96  const char data[] = {'\x00', '\x00', '\xff', '\xff'};
97  buffer_.insert(buffer_.end(), &data[0], &data[sizeof(data)]);
98}
99
100scoped_refptr<IOBufferWithSize> WebSocketDeflater::GetOutput(size_t size) {
101  std::deque<char>::iterator begin = buffer_.begin();
102  std::deque<char>::iterator end = begin + std::min(size, buffer_.size());
103
104  scoped_refptr<IOBufferWithSize> result = new IOBufferWithSize(end - begin);
105  std::copy(begin, end, result->data());
106  buffer_.erase(begin, end);
107  return result;
108}
109
110void WebSocketDeflater::ResetContext() {
111  if (mode_ == DO_NOT_TAKE_OVER_CONTEXT)
112    deflateReset(stream_.get());
113  are_bytes_added_ = false;
114}
115
116int WebSocketDeflater::Deflate(int flush) {
117  int result = Z_OK;
118  do {
119    stream_->next_out = reinterpret_cast<Bytef*>(&fixed_buffer_[0]);
120    stream_->avail_out = fixed_buffer_.size();
121    result = deflate(stream_.get(), flush);
122    size_t size = fixed_buffer_.size() - stream_->avail_out;
123    buffer_.insert(buffer_.end(), &fixed_buffer_[0], &fixed_buffer_[0] + size);
124  } while (result == Z_OK);
125  return result;
126}
127
128}  // namespace net
129