// Copyright (c) the JPEG XL Project 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 "lib/jxl/box_content_decoder.h"

#include <brotli/decode.h>

#include <jxl/decode.h>

#include <algorithm>

#include <cstddef>

#include <cstdint>

#include <cstring>

#include "lib/jxl/base/sanitizers.h"

namespace jxl {

JxlBoxContentDecoder::JxlBoxContentDecoder() = default;

JxlBoxContentDecoder::~JxlBoxContentDecoder() { Reset(); }

void JxlBoxContentDecoder::Reset() {

  if (brotli_dec) {

    BrotliDecoderDestroyInstance(brotli_dec);

    brotli_dec = nullptr;

  }

  header_done_ = false;

  brob_decode_ = false;

  box_until_eof_ = false;

  remaining_ = 0;

  pos_ = 0;

}

void JxlBoxContentDecoder::StartBox(bool brob_decode, bool box_until_eof,

                                    size_t contents_size) {

  Reset();

  brob_decode_ = brob_decode;

  box_until_eof_ = box_until_eof;

  remaining_ = box_until_eof ? 0 : contents_size;

}

JxlDecoderStatus JxlBoxContentDecoder::Process(const uint8_t* next_in,

                                               size_t avail_in, size_t box_pos,

                                               uint8_t** next_out,

                                               size_t* avail_out) {

  // The caller provides `box_pos` as the current position (in bytes) within

  // the box contents corresponding to `next_in`. Our internal `pos_` tracks

  // how many bytes of box contents have been consumed/processed so far.

  // If `box_pos` ever exceeds `pos_`, adjusting by `pos_ - box_pos` would

  // underflow and cause out-of-bounds reads.

  if (box_pos > pos_) return JXL_DEC_ERROR;

  const size_t offset = pos_ - box_pos;

  if (offset > avail_in) return JXL_DEC_ERROR;

  next_in += offset;

  avail_in -= offset;

  if (brob_decode_) {

    if (!header_done_) {

      if (avail_in < 4) return JXL_DEC_NEED_MORE_INPUT;

      if (!box_until_eof_) {

        if (remaining_ < 4) return JXL_DEC_ERROR;

        remaining_ -= 4;

      }

      next_in += 4;

      avail_in -= 4;

      pos_ += 4;

      header_done_ = true;

    }

    if (!brotli_dec) {

      brotli_dec = BrotliDecoderCreateInstance(nullptr, nullptr, nullptr);

    }

    size_t avail_in_brotli = avail_in;

    if (!box_until_eof_) avail_in_brotli = std::min<size_t>(avail_in_brotli, remaining_);

    const uint8_t* next_in_before = next_in;

    uint8_t* next_out_before = *next_out;

    msan::MemoryIsInitialized(next_in, avail_in_brotli);

    BrotliDecoderResult res = BrotliDecoderDecompressStream(

        brotli_dec, &avail_in_brotli, &next_in, avail_out, next_out, nullptr);

    size_t consumed = next_in - next_in_before;

    size_t produced = *next_out - next_out_before;

    if (res == BROTLI_DECODER_RESULT_ERROR) {

      return JXL_DEC_ERROR;

    }

    msan::UnpoisonMemory(next_out_before, produced);

    pos_ += consumed;

    if (!box_until_eof_) {

      remaining_ -= consumed;

      if (res == BROTLI_DECODER_RESULT_NEEDS_MORE_INPUT && remaining_ == 0) {

        return JXL_DEC_ERROR;

      }

    }

    if (res == BROTLI_DECODER_RESULT_NEEDS_MORE_INPUT) {

      return JXL_DEC_NEED_MORE_INPUT;

    }

    if (res == BROTLI_DECODER_RESULT_NEEDS_MORE_OUTPUT) {

      return JXL_DEC_BOX_NEED_MORE_OUTPUT;

    }

    if (res == BROTLI_DECODER_RESULT_SUCCESS) {

      if (!box_until_eof_ && remaining_ != 0) return JXL_DEC_ERROR;

      return JXL_DEC_BOX_COMPLETE;

    }

    // unknown Brotli result

    return JXL_DEC_ERROR;

  } else {

    // remaining box bytes as seen from dec->file_pos

    size_t can_read = avail_in;

    if (!box_until_eof_) can_read = std::min<size_t>(can_read, remaining_);

    size_t to_write = std::min<size_t>(can_read, *avail_out);

    memcpy(*next_out, next_in, to_write);

    *next_out += to_write;

    *avail_out -= to_write;

    if (!box_until_eof_) remaining_ -= to_write;

    pos_ += to_write;

    if (to_write < can_read) return JXL_DEC_BOX_NEED_MORE_OUTPUT;

    if (!box_until_eof_ && remaining_ > 0) return JXL_DEC_NEED_MORE_INPUT;

    return JXL_DEC_BOX_COMPLETE;

  }

}

}  // namespace jxl