// Copyright 2015 The Chromium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "net/filter/brotli_filter.h"

#include "base/bit_cast.h"
#include "base/macros.h"
#include "base/metrics/histogram_macros.h"
#include "base/numerics/safe_conversions.h"
#include "base/numerics/safe_math.h"
#include "third_party/brotli/dec/decode.h"

namespace net {

// BrotliFilter applies Brotli content decoding to a data stream.
// Brotli format specification: http://www.ietf.org/id/draft-alakuijala-brotli
//
// BrotliFilter is a subclass of Filter. See the latter's header file filter.h
// for sample usage.
class BrotliFilter : public Filter {
 public:
  BrotliFilter(FilterType type)
      : Filter(type),
        decoding_status_(DecodingStatus::DECODING_IN_PROGRESS),
        used_memory_(0),
        used_memory_maximum_(0),
        consumed_bytes_(0),
        produced_bytes_(0) {
    brotli_state_ = BrotliCreateState(BrotliFilter::AllocateMemory,
                                      BrotliFilter::FreeMemory, this);
    CHECK(brotli_state_);
  }

  ~BrotliFilter() override {
    BrotliDestroyState(brotli_state_);
    brotli_state_ = nullptr;
    DCHECK(used_memory_ == 0);

    UMA_HISTOGRAM_ENUMERATION(
        "BrotliFilter.Status", static_cast<int>(decoding_status_),
        static_cast<int>(DecodingStatus::DECODING_STATUS_COUNT));
    if (decoding_status_ == DecodingStatus::DECODING_DONE) {
      // CompressionPercent is undefined when there is no output produced.
      if (produced_bytes_ != 0) {
        UMA_HISTOGRAM_PERCENTAGE(
            "BrotliFilter.CompressionPercent",
            static_cast<int>((consumed_bytes_ * 100) / produced_bytes_));
      }
    }

    // All code here is for gathering stats, and can be removed when
    // BrotliFilter is considered stable.
    static const int kBuckets = 48;
    static const int64_t kMaxKb = 1 << (kBuckets / 3);  // 64MiB in KiB
    UMA_HISTOGRAM_CUSTOM_COUNTS("BrotliFilter.UsedMemoryKB",
                                used_memory_maximum_ / 1024, 1, kMaxKb,
                                kBuckets);
  }

  // Decodes the pre-filter data and writes the output into the |dest_buffer|
  // passed in.
  // The function returns FilterStatus. See filter.h for its description.
  //
  // Upon entry, |*dest_len| is the total size (in number of chars) of the
  // destination buffer. Upon exit, |*dest_len| is the actual number of chars
  // written into the destination buffer.
  //
  // This function will fail if there is no pre-filter data in the
  // |stream_buffer_|. On the other hand, |*dest_len| can be 0 upon successful
  // return. For example, decompressor may process some pre-filter data
  // but not produce output yet.
  FilterStatus ReadFilteredData(char* dest_buffer, int* dest_len) override {
    if (!dest_buffer || !dest_len)
      return Filter::FILTER_ERROR;

    if (decoding_status_ == DecodingStatus::DECODING_DONE) {
      *dest_len = 0;
      return Filter::FILTER_DONE;
    }

    if (decoding_status_ != DecodingStatus::DECODING_IN_PROGRESS)
      return Filter::FILTER_ERROR;

    size_t output_buffer_size = base::checked_cast<size_t>(*dest_len);
    size_t input_buffer_size = base::checked_cast<size_t>(stream_data_len_);

    size_t available_in = input_buffer_size;
    const uint8_t* next_in = bit_cast<uint8_t*>(next_stream_data_);
    size_t available_out = output_buffer_size;
    uint8_t* next_out = bit_cast<uint8_t*>(dest_buffer);
    size_t total_out = 0;
    BrotliResult result =
        BrotliDecompressStream(&available_in, &next_in, &available_out,
                               &next_out, &total_out, brotli_state_);

    CHECK(available_in <= input_buffer_size);
    CHECK(available_out <= output_buffer_size);
    consumed_bytes_ += input_buffer_size - available_in;
    produced_bytes_ += output_buffer_size - available_out;

    base::CheckedNumeric<size_t> safe_bytes_written(output_buffer_size);
    safe_bytes_written -= available_out;
    int bytes_written =
        base::checked_cast<int>(safe_bytes_written.ValueOrDie());

    switch (result) {
      case BROTLI_RESULT_NEEDS_MORE_OUTPUT:
      // Fall through.
      case BROTLI_RESULT_SUCCESS:
        *dest_len = bytes_written;
        stream_data_len_ = base::checked_cast<int>(available_in);
        next_stream_data_ = bit_cast<char*>(next_in);
        if (result == BROTLI_RESULT_SUCCESS) {
          decoding_status_ = DecodingStatus::DECODING_DONE;
          return Filter::FILTER_DONE;
        }
        return Filter::FILTER_OK;

      case BROTLI_RESULT_NEEDS_MORE_INPUT:
        *dest_len = bytes_written;
        stream_data_len_ = 0;
        next_stream_data_ = nullptr;
        return Filter::FILTER_NEED_MORE_DATA;

      default:
        decoding_status_ = DecodingStatus::DECODING_ERROR;
        return Filter::FILTER_ERROR;
    }
  }

 private:
  static void* AllocateMemory(void* opaque, size_t size) {
    BrotliFilter* filter = reinterpret_cast<BrotliFilter*>(opaque);
    return filter->AllocateMemoryInternal(size);
  }

  static void FreeMemory(void* opaque, void* address) {
    BrotliFilter* filter = reinterpret_cast<BrotliFilter*>(opaque);
    filter->FreeMemoryInternal(address);
  }

  void* AllocateMemoryInternal(size_t size) {
    size_t* array = reinterpret_cast<size_t*>(malloc(size + sizeof(size_t)));
    if (!array)
      return nullptr;
    used_memory_ += size;
    if (used_memory_maximum_ < used_memory_)
      used_memory_maximum_ = used_memory_;
    array[0] = size;
    return &array[1];
  }

  void FreeMemoryInternal(void* address) {
    if (!address)
      return;
    size_t* array = reinterpret_cast<size_t*>(address);
    used_memory_ -= array[-1];
    free(&array[-1]);
  }

  // Reported in UMA and must be kept in sync with the histograms.xml file.
  enum class DecodingStatus : int {
    DECODING_IN_PROGRESS = 0,
    DECODING_DONE,
    DECODING_ERROR,

    DECODING_STATUS_COUNT
    // DECODING_STATUS_COUNT must always be the last element in this enum.
  };

  // Tracks the status of decoding.
  // This variable is updated only by ReadFilteredData.
  DecodingStatus decoding_status_;

  BrotliState* brotli_state_;

  size_t used_memory_;
  size_t used_memory_maximum_;
  size_t consumed_bytes_;
  size_t produced_bytes_;

  DISALLOW_COPY_AND_ASSIGN(BrotliFilter);
};

Filter* CreateBrotliFilter(Filter::FilterType type_id) {
  return new BrotliFilter(type_id);
}

}  // namespace net