// Example fuzzer for PNG using protos.
#include <string>
#include <sstream>
#include <fstream>
#include <zlib.h>  // for crc32

#include "libprotobuf-mutator/src/libfuzzer/libfuzzer_macro.h"
#include "png_fuzz_proto.pb.h"

static void WriteInt(std::stringstream &out, uint32_t x) {
  x = __builtin_bswap32(x);
  out.write((char *)&x, sizeof(x));
}

static void WriteByte(std::stringstream &out, uint8_t x) {
  out.write((char *)&x, sizeof(x));
}

static std::string Compress(const std::string &s) {
  std::string out(s.size() + 100, '\0');
  size_t out_len = out.size();
  compress((uint8_t *)&out[0], &out_len, (uint8_t *)s.data(), s.size());
  out.resize(out_len);
  return out;
}

// Chunk is written as:
//  * 4-byte length
//  * 4-byte type
//  * the data itself
//  * 4-byte crc (of type and data)
static void WriteChunk(std::stringstream &out, const char *type,
                       const std::string &chunk, bool compress = false) {
  std::string compressed;
  const std::string *s = &chunk;
  if (compress) {
    compressed = Compress(chunk);
    s = &compressed;
  }
  uint32_t len = s->size();
  uint32_t crc = crc32(crc32(0, (const unsigned char *)type, 4),
                       (const unsigned char *)s->data(), s->size());
  WriteInt(out, len);
  out.write(type, 4);
  out.write(s->data(), s->size());
  WriteInt(out, crc);
}

std::string ProtoToPng(const PngProto &png_proto) {
  std::stringstream all;
  const unsigned char header[] = {0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a};
  all.write((const char*)header, sizeof(header));
  std::stringstream ihdr_str;
  auto &ihdr = png_proto.ihdr();
  // Avoid large images.
  // They may have interesting bugs, but OOMs are going to kill fuzzing.
  uint32_t w = std::min(ihdr.width(), 4096U);
  uint32_t h = std::min(ihdr.height(), 4096U);
  WriteInt(ihdr_str, w);
  WriteInt(ihdr_str, h);
  WriteInt(ihdr_str, ihdr.other1());
  WriteByte(ihdr_str, ihdr.other2());
  WriteChunk(all, "IHDR", ihdr_str.str());

  for (size_t i = 0, n = png_proto.chunks_size(); i < n; i++) {
    auto &chunk = png_proto.chunks(i);
    if (chunk.has_plte()) {
      WriteChunk(all, "PLTE", chunk.plte().data());
    } else if (chunk.has_idat()) {
      WriteChunk(all, "IDAT", chunk.idat().data(), true);
    } else if (chunk.has_iccp()) {
      std::stringstream iccp_str;
      iccp_str << "xyz";  // don't fuzz iCCP name field.
      WriteByte(iccp_str, 0);
      WriteByte(iccp_str, 0);
      auto compressed_data = Compress(chunk.iccp().data());
      iccp_str.write(compressed_data.data(), compressed_data.size());
      WriteChunk(all, "iCCP", iccp_str.str());
    } else if (chunk.has_other_chunk()) {
      auto &other_chunk = chunk.other_chunk();
      char type[5] = {0};
      if (other_chunk.has_known_type()) {
        static const char * known_chunks[] = {
            "bKGD", "cHRM", "dSIG", "eXIf", "gAMA", "hIST", "iCCP",
            "iTXt", "pHYs", "sBIT", "sPLT", "sRGB", "sTER", "tEXt",
            "tIME", "tRNS", "zTXt", "sCAL", "pCAL", "oFFs",
        };
        size_t known_chunks_size =
            sizeof(known_chunks) / sizeof(known_chunks[0]);
        size_t chunk_idx = other_chunk.known_type() % known_chunks_size;
        memcpy(type, known_chunks[chunk_idx], 4);
      } else if (other_chunk.has_unknown_type()) {
        uint32_t unknown_type_int = other_chunk.unknown_type();
        memcpy(type, &unknown_type_int, 4);
      } else {
        continue;
      }
      type[4] = 0;
      WriteChunk(all, type, other_chunk.data());
    }
  }
  WriteChunk(all, "IEND", "");

  std::string res = all.str();
  if (const char *dump_path = getenv("PROTO_FUZZER_DUMP_PATH")) {
    // With libFuzzer binary run this to generate a PNG file x.png:
    // PROTO_FUZZER_DUMP_PATH=x.png ./a.out proto-input
    std::ofstream of(dump_path);
    of.write(res.data(), res.size());
  }
  return res;
}

// The actual fuzz target that consumes the PNG data.
extern "C" int FuzzPNG(const uint8_t* data, size_t size);

DEFINE_PROTO_FUZZER(const PngProto &png_proto) {
  auto s = ProtoToPng(png_proto);
  FuzzPNG((const uint8_t*)s.data(), s.size());
}