// 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/extras/dec/jpg.h"

#include <cstdint>

#include "lib/extras/dec/color_hints.h"

#include "lib/extras/packed_image.h"

#include "lib/extras/size_constraints.h"

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

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

#if !JPEGXL_ENABLE_JPEG

namespace jxl {

namespace extras {

bool CanDecodeJPG() { return false; }

Status DecodeImageJPG(const Span<const uint8_t> bytes,

                      const ColorHints& color_hints, PackedPixelFile* ppf,

                      const SizeConstraints* constraints,

                      const JPGDecompressParams* dparams) {

  return false;

}

}  // namespace extras

}  // namespace jxl

#else  // JPEGXL_ENABLE_JPEG

#include <jxl/codestream_header.h>

#include <jxl/color_encoding.h>

#include <jxl/types.h>

#include <algorithm>

#include <cstring>

#include <memory>

#include <numeric>

#include <utility>

#include <vector>

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

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

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

namespace jxl {

namespace extras {

namespace {

constexpr unsigned char kICCSignature[12] = {

    0x49, 0x43, 0x43, 0x5F, 0x50, 0x52, 0x4F, 0x46, 0x49, 0x4C, 0x45, 0x00};

constexpr int kICCMarker = JPEG_APP0 + 2;

constexpr unsigned char kExifSignature[6] = {0x45, 0x78, 0x69,

                                             0x66, 0x00, 0x00};

constexpr int kExifMarker = JPEG_APP0 + 1;

inline bool IsJPG(const Span<const uint8_t> bytes) {

  if (bytes.size() < 2) return false;

  if (bytes[0] != 0xFF || bytes[1] != 0xD8) return false;

  return true;

}

bool MarkerIsICC(const jpeg_saved_marker_ptr marker) {

  return marker->marker == kICCMarker &&

         marker->data_length >= sizeof kICCSignature + 2 &&

         std::equal(std::begin(kICCSignature), std::end(kICCSignature),

                    marker->data);

}

bool MarkerIsExif(const jpeg_saved_marker_ptr marker) {

  return marker->marker == kExifMarker &&

         marker->data_length >= sizeof kExifSignature + 2 &&

         std::equal(std::begin(kExifSignature), std::end(kExifSignature),

                    marker->data);

}

Status ReadICCProfile(jpeg_decompress_struct* const cinfo,

                      std::vector<uint8_t>* const icc) {

  constexpr size_t kICCSignatureSize = sizeof kICCSignature;

  // ICC signature + uint8_t index + uint8_t max_index.

  constexpr size_t kICCHeadSize = kICCSignatureSize + 2;

  // Markers are 1-indexed, and we keep them that way in this vector to get a

  // convenient 0 at the front for when we compute the offsets later.

  std::vector<size_t> marker_lengths;

  int num_markers = 0;

  int seen_markers_count = 0;

  bool has_num_markers = false;

  for (jpeg_saved_marker_ptr marker = cinfo->marker_list; marker != nullptr;

       marker = marker->next) {

    // marker is initialized by libjpeg, which we are not instrumenting with

    // msan.

    msan::UnpoisonMemory(marker, sizeof(*marker));

    msan::UnpoisonMemory(marker->data, marker->data_length);

    if (!MarkerIsICC(marker)) continue;

    const int current_marker = marker->data[kICCSignatureSize];

    if (current_marker == 0) {

      return JXL_FAILURE("inconsistent JPEG ICC marker numbering");

    }

    const int current_num_markers = marker->data[kICCSignatureSize + 1];

    if (current_marker > current_num_markers) {

      return JXL_FAILURE("inconsistent JPEG ICC marker numbering");

    }

    if (has_num_markers) {

      if (current_num_markers != num_markers) {

        return JXL_FAILURE("inconsistent numbers of JPEG ICC markers");

      }

    } else {

      num_markers = current_num_markers;

      has_num_markers = true;

      marker_lengths.resize(num_markers + 1);

    }

    size_t marker_length = marker->data_length - kICCHeadSize;

    if (marker_length == 0) {

      // NB: if we allow empty chunks, then the next check is incorrect.

      return JXL_FAILURE("Empty ICC chunk");

    }

    if (marker_lengths[current_marker] != 0) {

      return JXL_FAILURE("duplicate JPEG ICC marker number");

    }

    marker_lengths[current_marker] = marker_length;

    seen_markers_count++;

  }

  if (marker_lengths.empty()) {

    // Not an error.

    return false;

  }

  if (seen_markers_count != num_markers) {

    JXL_ENSURE(has_num_markers);

    return JXL_FAILURE("Incomplete set of ICC chunks");

  }

  std::vector<size_t> offsets = std::move(marker_lengths);

  std::partial_sum(offsets.begin(), offsets.end(), offsets.begin());

  icc->resize(offsets.back());

  for (jpeg_saved_marker_ptr marker = cinfo->marker_list; marker != nullptr;

       marker = marker->next) {

    if (!MarkerIsICC(marker)) continue;

    const uint8_t* first = marker->data + kICCHeadSize;

    uint8_t current_marker = marker->data[kICCSignatureSize];

    size_t offset = offsets[current_marker - 1];

    size_t marker_length = offsets[current_marker] - offset;

    std::copy_n(first, marker_length, icc->data() + offset);

  }

  return true;

}

void ReadExif(jpeg_decompress_struct* const cinfo,

              std::vector<uint8_t>* const exif) {

  constexpr size_t kExifSignatureSize = sizeof kExifSignature;

  for (jpeg_saved_marker_ptr marker = cinfo->marker_list; marker != nullptr;

       marker = marker->next) {

    // marker is initialized by libjpeg, which we are not instrumenting with

    // msan.

    msan::UnpoisonMemory(marker, sizeof(*marker));

    msan::UnpoisonMemory(marker->data, marker->data_length);

    if (!MarkerIsExif(marker)) continue;

    size_t marker_length = marker->data_length - kExifSignatureSize;

    exif->resize(marker_length);

    std::copy_n(marker->data + kExifSignatureSize, marker_length, exif->data());

    return;

  }

}

void MyErrorExit(j_common_ptr cinfo) {

  jmp_buf* env = static_cast<jmp_buf*>(cinfo->client_data);

  (*cinfo->err->output_message)(cinfo);

  jpeg_destroy_decompress(reinterpret_cast<j_decompress_ptr>(cinfo));

  longjmp(*env, 1);

}

void MyOutputMessage(j_common_ptr cinfo) {

  if (JXL_IS_DEBUG_BUILD) {

    char buf[JMSG_LENGTH_MAX + 1];

    (*cinfo->err->format_message)(cinfo, buf);

    buf[JMSG_LENGTH_MAX] = 0;

    JXL_WARNING("%s", buf);

  }

}

Status UnmapColors(uint8_t* row, size_t xsize, int components,

                   JSAMPARRAY colormap, size_t num_colors) {

  JXL_ENSURE(colormap != nullptr);

  std::vector<uint8_t> tmp(xsize * components);

  for (size_t x = 0; x < xsize; ++x) {

    JXL_ENSURE(row[x] < num_colors);

    for (int c = 0; c < components; ++c) {

      tmp[x * components + c] = colormap[c][row[x]];

    }

  }

  memcpy(row, tmp.data(), tmp.size());

  return true;

}

}  // namespace

bool CanDecodeJPG() { return true; }

Status DecodeImageJPG(const Span<const uint8_t> bytes,

                      const ColorHints& color_hints, PackedPixelFile* ppf,

                      const SizeConstraints* constraints,

                      const JPGDecompressParams* dparams) {

  // Don't do anything for non-JPEG files (no need to report an error)

  if (!IsJPG(bytes)) return false;

  // TODO(veluca): use JPEGData also for pixels?

  // We need to declare all the non-trivial destructor local variables before

  // the call to setjmp().

  std::unique_ptr<JSAMPLE[]> row;

  const auto try_catch_block = [&]() -> bool {

    jpeg_decompress_struct cinfo = {};

    // Setup error handling in jpeg library so we can deal with broken jpegs in

    // the fuzzer.

    jpeg_error_mgr jerr;

    jmp_buf env;

    cinfo.err = jpeg_std_error(&jerr);

    jerr.error_exit = &MyErrorExit;

    jerr.output_message = &MyOutputMessage;

    if (setjmp(env)) {

      return false;

    }

    cinfo.client_data = static_cast<void*>(&env);

    jpeg_create_decompress(&cinfo);

    jpeg_mem_src(&cinfo, reinterpret_cast<const unsigned char*>(bytes.data()),

                 bytes.size());

    jpeg_save_markers(&cinfo, kICCMarker, 0xFFFF);

    jpeg_save_markers(&cinfo, kExifMarker, 0xFFFF);

    const auto failure = [&cinfo](const char* str) -> Status {

      jpeg_abort_decompress(&cinfo);

      jpeg_destroy_decompress(&cinfo);

      return JXL_FAILURE("%s", str);

    };

    int read_header_result = jpeg_read_header(&cinfo, TRUE);

    // TODO(eustas): what about JPEG_HEADER_TABLES_ONLY?

    if (read_header_result == JPEG_SUSPENDED) {

      return failure("truncated JPEG input");

    }

    if (!VerifyDimensions(constraints, cinfo.image_width, cinfo.image_height)) {

      return failure("image too big");

    }

    // Might cause CPU-zip bomb.

    if (cinfo.arith_code) {

      return failure("arithmetic code JPEGs are not supported");

    }

    int nbcomp = cinfo.num_components;

    if (nbcomp != 1 && nbcomp != 3) {

      return failure("unsupported number of components in JPEG");

    }

    if (ReadICCProfile(&cinfo, &ppf->icc)) {

      ppf->primary_color_representation = PackedPixelFile::kIccIsPrimary;

    } else {

      ppf->primary_color_representation =

          PackedPixelFile::kColorEncodingIsPrimary;

      ppf->icc.clear();

      // Default to SRGB

      // Actually, (cinfo.output_components == nbcomp) will be checked after

      // `jpeg_start_decompress`.

      ppf->color_encoding.color_space =

          (nbcomp == 1) ? JXL_COLOR_SPACE_GRAY : JXL_COLOR_SPACE_RGB;

      ppf->color_encoding.white_point = JXL_WHITE_POINT_D65;

      ppf->color_encoding.primaries = JXL_PRIMARIES_SRGB;

      ppf->color_encoding.transfer_function = JXL_TRANSFER_FUNCTION_SRGB;

      ppf->color_encoding.rendering_intent = JXL_RENDERING_INTENT_PERCEPTUAL;

    }

    ReadExif(&cinfo, &ppf->metadata.exif);

    if (!ApplyColorHints(color_hints, /*color_already_set=*/true,

                         /*is_gray=*/false, ppf)) {

      return failure("ApplyColorHints failed");

    }

    ppf->info.xsize = cinfo.image_width;

    ppf->info.ysize = cinfo.image_height;

    // Original data is uint, so exponent_bits_per_sample = 0.

    ppf->info.bits_per_sample = BITS_IN_JSAMPLE;

    static_assert(BITS_IN_JSAMPLE == 8 || BITS_IN_JSAMPLE == 16,

                  "Only 8/16 bit samples are supported");

    ppf->info.exponent_bits_per_sample = 0;

    ppf->info.uses_original_profile = JXL_TRUE;

    // No alpha in JPG

    ppf->info.alpha_bits = 0;

    ppf->info.alpha_exponent_bits = 0;

    ppf->info.num_color_channels = nbcomp;

    ppf->info.orientation = JXL_ORIENT_IDENTITY;

    if (dparams && dparams->num_colors > 0) {

      cinfo.quantize_colors = TRUE;

      cinfo.desired_number_of_colors = dparams->num_colors;

      cinfo.two_pass_quantize = static_cast<boolean>(dparams->two_pass_quant);

      cinfo.dither_mode = static_cast<J_DITHER_MODE>(dparams->dither_mode);

    }

    jpeg_start_decompress(&cinfo);

    JXL_ENSURE(cinfo.out_color_components == nbcomp);

    JxlDataType data_type =

        ppf->info.bits_per_sample <= 8 ? JXL_TYPE_UINT8 : JXL_TYPE_UINT16;

    const JxlPixelFormat format{

        /*num_channels=*/static_cast<uint32_t>(nbcomp),

        data_type,

        /*endianness=*/JXL_NATIVE_ENDIAN,

        /*align=*/0,

    };

    ppf->frames.clear();

    // Allocates the frame buffer.

    {

      JXL_ASSIGN_OR_RETURN(

          PackedFrame frame,

          PackedFrame::Create(cinfo.image_width, cinfo.image_height, format));

      ppf->frames.emplace_back(std::move(frame));

    }

    const auto& frame = ppf->frames.back();

    JXL_ENSURE(sizeof(JSAMPLE) * cinfo.out_color_components *

                   cinfo.image_width <=

               frame.color.stride);

    if (cinfo.quantize_colors) {

      JSAMPLE** colormap = cinfo.colormap;

      jxl::msan::UnpoisonMemory(reinterpret_cast<void*>(colormap),

                                cinfo.out_color_components * sizeof(JSAMPLE*));

      for (int c = 0; c < cinfo.out_color_components; ++c) {

        jxl::msan::UnpoisonMemory(

            reinterpret_cast<void*>(colormap[c]),

            cinfo.actual_number_of_colors * sizeof(JSAMPLE));

      }

    }

    if (dparams && dparams->num_colors > 0) {

      JXL_ENSURE(cinfo.colormap != nullptr);

    }

    for (size_t y = 0; y < cinfo.image_height; ++y) {

      JSAMPROW rows[] = {reinterpret_cast<JSAMPLE*>(

          static_cast<uint8_t*>(frame.color.pixels()) +

          frame.color.stride * y)};

      jpeg_read_scanlines(&cinfo, rows, 1);

      msan::UnpoisonMemory(rows[0], sizeof(JSAMPLE) * cinfo.output_components *

                                        cinfo.image_width);

      if (dparams && dparams->num_colors > 0) {

        JXL_RETURN_IF_ERROR(

            UnmapColors(rows[0], cinfo.output_width, cinfo.out_color_components,

                        cinfo.colormap, cinfo.actual_number_of_colors));

      }

    }

    jpeg_finish_decompress(&cinfo);

    jpeg_destroy_decompress(&cinfo);

    return true;

  };

  return try_catch_block();

}

}  // namespace extras

}  // namespace jxl

#endif  // JPEGXL_ENABLE_JPEG