/*

 * HEIF codec.

 * Copyright (c) 2017 struktur AG, Joachim Bauch <bauch@struktur.de>

 *

 * This file is part of libheif.

 *

 * libheif is free software: you can redistribute it and/or modify

 * it under the terms of the GNU Lesser General Public License as

 * published by the Free Software Foundation, either version 3 of

 * the License, or (at your option) any later version.

 *

 * libheif is distributed in the hope that it will be useful,

 * but WITHOUT ANY WARRANTY; without even the implied warranty of

 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the

 * GNU Lesser General Public License for more details.

 *

 * You should have received a copy of the GNU Lesser General Public License

 * along with libheif.  If not, see <http://www.gnu.org/licenses/>.

 */

#include <assert.h>

#include <limits.h>

#include <stdlib.h>

#include "libheif/heif.h"

static const enum heif_colorspace kFuzzColorSpace = heif_colorspace_YCbCr;

static const enum heif_chroma kFuzzChroma = heif_chroma_420;

static void TestDecodeImage(struct heif_context* ctx,

                            const struct heif_image_handle* handle, size_t filesize)

{

  struct heif_image* image = nullptr;

  struct heif_error err;

  bool primary = heif_image_handle_is_primary_image(handle);

  (void) primary;

  int width = heif_image_handle_get_width(handle);

  int height = heif_image_handle_get_height(handle);

  (void)width; (void)height;

  assert(width >= 0);

  assert(height >= 0);

  int metadata_count = heif_image_handle_get_number_of_metadata_blocks(handle, nullptr);

  assert(metadata_count >= 0);

  assert(static_cast<size_t>(metadata_count) < filesize / sizeof(heif_item_id));

  heif_item_id* metadata_ids = static_cast<heif_item_id*>(malloc(metadata_count * sizeof(heif_item_id)));

  assert(metadata_ids);

  int metadata_ids_count = heif_image_handle_get_list_of_metadata_block_IDs(handle, nullptr, metadata_ids,

                                                                            metadata_count);

  assert(metadata_count == metadata_ids_count);

  (void)metadata_ids_count;

  for (int i = 0; i < metadata_count; i++) {

    heif_image_handle_get_metadata_type(handle, metadata_ids[i]);

    heif_image_handle_get_metadata_content_type(handle, metadata_ids[i]);

    size_t metadata_size = heif_image_handle_get_metadata_size(handle, metadata_ids[i]);

    // This assertion is invalid. Metadata can in fact be larger than the file if there are several

    // overlapping iloc extents. Does not make much sense, but it is technically valid.

    //assert(metadata_size < filesize);

    uint8_t* metadata_data = static_cast<uint8_t*>(malloc(metadata_size));

    assert(metadata_data);

    heif_image_handle_get_metadata(handle, metadata_ids[i], metadata_data);

    free(metadata_data);

  }

  free(metadata_ids);

  err = heif_decode_image(handle, &image, kFuzzColorSpace, kFuzzChroma, nullptr);

  if (err.code != heif_error_Ok) {

    heif_image_release(image);

    return;

  }

  assert(heif_image_get_colorspace(image) == kFuzzColorSpace);

  assert(heif_image_get_chroma_format(image) == kFuzzChroma);

  // TODO(fancycode): Should we also check the planes?

  heif_image_release(image);

}

static int clip_int(size_t size)

{

  return size > INT_MAX ? INT_MAX : static_cast<int>(size);

}

extern "C" int LLVMFuzzerTestOneInput(const uint8_t* data, size_t size)

{

  struct heif_context* ctx;

  struct heif_error err;

  struct heif_image_handle* primary_handle = nullptr;

  int images_count;

  heif_item_id* image_IDs = NULL;

  bool explicit_init = size == 0 || data[size - 1] & 1;

  if (explicit_init) {

    heif_init(nullptr);

  }

  heif_check_filetype(data, clip_int(size));

  heif_main_brand(data, clip_int(size));

  heif_get_file_mime_type(data, clip_int(size));

  ctx = heif_context_alloc();

  assert(ctx);

  auto* limits = heif_context_get_security_limits(ctx);

  limits->max_total_memory = UINT64_C(2) * 1024 * 1024 * 1024;

  limits->max_memory_block_size = 128 * 1024 * 1024; // 128 MB

  // Also bound the image size itself. libheif's memory accounting only covers its own

  // allocations, but the codec libraries allocate their frame buffers before libheif

  // sees anything (dav1d needs ~12.5 bytes/pixel for a frame header alone). Under MSan,

  // where every allocation costs about three times its size in RSS, a 56 Mpixel AV1 frame

  // header was enough to exceed libFuzzer's 2560 MB limit.

  limits->max_image_size_pixels = 16 * 1024 * 1024; // 16 Mpixel

  err = heif_context_read_from_memory(ctx, data, size, nullptr);

  if (err.code != heif_error_Ok) {

    // Not a valid HEIF file passed (which is most likely while fuzzing).

    goto quit;

  }

  err = heif_context_get_primary_image_handle(ctx, &primary_handle);

  if (err.code == heif_error_Ok) {

    assert(heif_image_handle_is_primary_image(primary_handle));

    TestDecodeImage(ctx, primary_handle, size);

    heif_image_handle_release(primary_handle);

    primary_handle = nullptr;

  }

  images_count = heif_context_get_number_of_top_level_images(ctx);

  if (!images_count) {

    // File doesn't contain any images.

    goto quit;

  }

  image_IDs = (heif_item_id*) malloc(images_count * sizeof(heif_item_id));

  assert(image_IDs);

  images_count = heif_context_get_list_of_top_level_image_IDs(ctx, image_IDs, images_count);

  if (!images_count) {

    // Could not get list of image ids.

    goto quit;

  }

  for (int i = 0; i < images_count; ++i) {

    struct heif_image_handle* image_handle = nullptr;

    err = heif_context_get_image_handle(ctx, image_IDs[i], &image_handle);

    if (err.code != heif_error_Ok) {

      heif_image_handle_release(image_handle);

      // Ignore, we are only interested in crashes here.

      continue;

    }

    TestDecodeImage(ctx, image_handle, size);

    int num_thumbnails = heif_image_handle_get_number_of_thumbnails(image_handle);

    for (int t = 0; t < num_thumbnails; ++t) {

      struct heif_image_handle* thumbnail_handle = nullptr;

      heif_image_handle_get_thumbnail(image_handle, t, &thumbnail_handle);

      if (thumbnail_handle) {

        TestDecodeImage(ctx, thumbnail_handle, size);

        heif_image_handle_release(thumbnail_handle);

      }

    }

    heif_image_handle_release(image_handle);

  }

  quit:

  heif_image_handle_release(primary_handle);

  heif_context_free(ctx);

  free(image_IDs);

  if (explicit_init) {

    heif_deinit();

  }

  return 0;

}