#include <algorithm>

#include <cassert>

#include <cstring>

#include <string>

#include <vector>

#include "png.h"

namespace {

struct PngReader {

  png_structp png_ptr = nullptr;

  png_infop info_ptr = nullptr;

  png_infop end_info = nullptr;

};

struct PngArrayStream {

  const uint8_t *data;

  size_t size;

  size_t pos;

};

void PngArrayStreamCallback(png_structp png_ptr, png_bytep data,

                            png_size_t size) {

  PngArrayStream *stream =

      static_cast<PngArrayStream *>(png_get_io_ptr(png_ptr));

  if (stream->pos + size > stream->size) {

    memset(data, 0, size);

    stream->pos = size;

  } else {

    memcpy(data, &stream->data[stream->pos], size);

    stream->pos += size;

  }

}

static bool PngVerboseWarnings = getenv("PNG_VERBOSE_WARNINGS") != nullptr;

void PngErrorHandler(png_structp png_ptr, png_const_charp error_message) {

  if (PngVerboseWarnings) fprintf(stderr, "%s\n", error_message);

  longjmp(png_jmpbuf(png_ptr), 1);

}

void PngWarningHandler(png_structp png_ptr, png_const_charp warning_message) {

  if (PngVerboseWarnings) fprintf(stderr, "%s\n", warning_message);

  longjmp(png_jmpbuf(png_ptr), 1);

}

}  // namespace

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

  const size_t kPngSignatureSize = 8;

  const size_t kIHDRSize = 4 + 4 + 13 + 4;

  const size_t kMaxImageSize = 1 << 20;

  const size_t kMaxHeight = 1 << 10;

  auto Read32 = [&](const uint8_t *p) {

    uint32_t res;

    assert(p >= data);

    assert(p + sizeof(res) < data + size);

    memcpy(&res, p, sizeof(res));

    return res;

  };

  if (size < kPngSignatureSize + kIHDRSize) return 0;

  if (png_sig_cmp(data, 0, kPngSignatureSize)) return 0;

  uint32_t width = __builtin_bswap32(Read32(data + kPngSignatureSize + 8));

  uint32_t height = __builtin_bswap32(Read32(data + kPngSignatureSize + 12));

  // Reject too large images because they will OOM.

  // Also reject images with a too large height, because large height

  // will cause too many mallocs.

  // These two heuristics are far from optimal and may cause us to lose some

  // coverage. But w/o them fuzzing is way too slow.

  if ((uint64_t)width * height > kMaxImageSize) return 0;

  if (height > kMaxHeight) return 0;

  // Find the fUZz chunk and it's contents.

  const size_t fUZz_chunk_size = 16;

  const uint8_t fUZz_signature[8] = {0,   0,   0,   fUZz_chunk_size,

                                     'f', 'U', 'Z', 'z'};

  const uint8_t *fUZz_beg =

      std::search(data, data + size, fUZz_signature,

                  fUZz_signature + sizeof(fUZz_signature));

  if (fUZz_beg + sizeof(fUZz_signature) + fUZz_chunk_size < data + size)

    fUZz_beg += sizeof(fUZz_signature);

  else

    fUZz_beg = nullptr;

  PngReader reader;

  reader.png_ptr =

      png_create_read_struct(PNG_LIBPNG_VER_STRING, NULL, NULL, NULL);

  assert(reader.png_ptr);

  reader.info_ptr = png_create_info_struct(reader.png_ptr);

  assert(reader.info_ptr);

  reader.end_info = png_create_info_struct(reader.png_ptr);

  assert(reader.end_info);

  png_set_error_fn(reader.png_ptr, png_get_error_ptr(reader.png_ptr),

                   PngErrorHandler, PngWarningHandler);

  PngArrayStream stream{data, size, 0};

  if (setjmp(png_jmpbuf(reader.png_ptr)) == 0) {

    png_set_read_fn(reader.png_ptr, &stream, PngArrayStreamCallback);

    // Take transforms from the fUZz chunk. By default, enable all.

    int transforms = fUZz_beg ? Read32(fUZz_beg) : ~0;

    png_read_png(reader.png_ptr, reader.info_ptr, transforms, nullptr);

  }

  png_destroy_read_struct(&reader.png_ptr, &reader.info_ptr, &reader.end_info);

  // Run the same image through another libpng API.

  // There is probably some redundancy here (I don't know what I am doing!)

  png_image image;

  memset(&image, 0, sizeof(image));

  image.version = PNG_IMAGE_VERSION;

  if (png_image_begin_read_from_memory(&image, data, size)) {

    const size_t kMaxBufferSize = 64 << 20;

    image.format = fUZz_beg ? Read32(fUZz_beg + 4) : PNG_FORMAT_RGBA;

    size_t image_size = PNG_IMAGE_SIZE(image);

    if (image_size <= kMaxBufferSize) {

      png_bytep buffer = new png_byte[image_size];

      const size_t kColorMapSize = 256 * 4;

      // Do we need to take color & colormap from the fuzzed input?

      png_color color = {1, 2, 3};

      png_uint_16 colormap[256*4] = {0};

      for (size_t i = 0; i < kColorMapSize; i++)

        colormap[i] = i;

      png_image_finish_read(&image, &color, buffer, 0, colormap);

      delete[] buffer;

    }

  }

  png_image_free(&image);

  return 0;

}

extern "C" const char *__asan_default_options() { 

  // TODO: remove this once 

  // https://bugs.chromium.org/p/oss-fuzz/issues/detail?id=12716  

  // is fixed.  

  return "detect_leaks=0";  

}