// Copyright 2019 Google Inc. All Rights Reserved.

// Licensed under the Apache License, Version 2.0 (the "License");

#include <algorithm>

#include <iostream>

#include <sstream>

#include <string>

#include <vector>

#include <cstring>

#include <cassert>

#include <random>

#include <zlib.h>

// A simple class for parsing, serializing, and mutating an PNG file.

// https://en.wikipedia.org/wiki/Portable_Network_Graphics

// It is an example of a custom mutator for libFuzzer

// (https://llvm.org/docs/LibFuzzer.html) used for

// "structure-aware coverage-guided fuzzing".

//

// If you have a non structure-aware fuzz target for any API that handles

// PNG inputs, you can turn that fuzz target into a structure-aware one

// by defining PNG_MUTATOR_DEFINE_LIBFUZZER_CUSTOM_MUTATOR and then

// including this file.

class PngMutator {

  using V = std::vector<uint8_t>;

 public:

  // Parse the input stream as a PNG file,

  // put every chunk into its own vector,

  // uncompress chunk data when needed,

  // merge the IDAT chunks into one vector.

  PngMutator(std::istream &in) {

    ihdr_.resize(13);

    Read4(in);

    Read4(in); // Skip the 8-byte magic value.

    // read IHDR.

    if (ReadInteger(in) != 13) return;

    if (Read4(in) != Type("IHDR")) return;

    // Read 13 values.

    in.read((char*)ihdr_.data(), ihdr_.size());

    Read4(in);  // ignore CRC

    ssize_t idat_idx = -1;

    while (in) {

      uint32_t len = ReadInteger(in);

      uint32_t type = Read4(in);

      if (type == Type("IEND")) break;  // do nothing

      char chunk_name[5];

      memcpy(chunk_name, &type, 4);

      chunk_name[4] = 0;

      if (len > (1 << 20)) return;

      V v(len);

      in.read((char *)v.data(), len);

      Read4(in);  // ignore CRC

      if (type == Type("IDAT")) {

        if (idat_idx != -1)

          Append(&chunks_[idat_idx].v, v);

        else {

          idat_idx = chunks_.size();

          chunks_.push_back({type, v});

        }

      } else if (type == Type("iCCP")) {

        auto it = v.begin();

        while (it < v.end() && isprint(*it)) it++;

        if (it < v.end() && !*it) it++;

        if (it < v.end() && !*it) it++;

        v = V(it, v.end());

        auto uncompressed = Uncompress(v);

        chunks_.push_back({type, uncompressed});

        auto compressed = Compress(uncompressed);

      } else {

        chunks_.push_back({type, v});

      }

      //  std::cerr << "CHUNK: " << chunk_name << std::endl;

    }

    if (idat_idx != -1)

      chunks_[idat_idx].v = Uncompress(chunks_[idat_idx].v);

  }

  // Write back the PNG file.

  void Serialize(std::ostream &out) {

    const unsigned char header[] = {0x89, 0x50, 0x4e, 0x47,

                                    0x0d, 0x0a, 0x1a, 0x0a};

    out.write((const char*)header, sizeof(header));

    WriteChunk(out, "IHDR", ihdr_);

    for (auto &ch : chunks_) {

      if (ch.type == Type("iCCP")) {

        V v;

        v.push_back('x');  // assuming the iCCP name doesn't matter.

        v.push_back(0);

        v.push_back(0);

        auto compressed = Compress(ch.v);

        Append(&v, compressed);

        WriteChunk(out, ch.type, v);

      } else {

        WriteChunk(out, ch.type, ch.v);

      }

    }

    WriteChunk(out, "IEND", {});

  }

  // Raw byte array mutator, like that provided by libFuzzer.

  using Mutator = size_t (*)(uint8_t *Data, size_t Size, size_t MaxSize);

  // Mutate the in-memory representation of a PNG file.

  // Given the same Seed, the same mutation is performed.

  void Mutate(Mutator m, unsigned int Seed) {

    std::minstd_rand rnd(Seed);

    auto M = [&](V *v) {

      if (v->empty())

        v->resize(v->size() + 1 + rnd() % 256);

      v->resize(m(v->data(), v->size(), v->size()));

    };

    switch (rnd() % 6) {

      // Mutate IHDR.

      case 0:

        m(ihdr_.data(), ihdr_.size(), ihdr_.size());

        break;

      // Mutate some other chunk.

      case 1:

        if (!chunks_.empty()) M(&chunks_[rnd() % chunks_.size()].v);

        break;

      // Shuffle the chunks.

      case 2:

        std::shuffle(chunks_.begin(), chunks_.end(), rnd);

        break;

      // Delete a random chunk.

      case 3:

        if (!chunks_.empty())

         chunks_.erase(chunks_.begin() + rnd() % chunks_.size());

        break;

      // Insert a random chunk with one of the known types, or a random type.

      case 4: {

        static const char *types[] = {

            "IATx", "sTER", "hIST", "sPLT", "mkBF", "mkBS", "mkTS", "prVW",

            "oFFs", "iDOT", "zTXt", "mkBT", "acTL", "iTXt", "sBIT", "tIME",

            "iCCP", "vpAg", "tRNS", "cHRM", "PLTE", "bKGD", "gAMA", "sRGB",

            "pHYs", "fdAT", "fcTL", "tEXt", "IDAT",

            "pCAL", "sCAL", "eXIf",

            "fUZz", // special chunk for extra fuzzing hints.

        };

        static const size_t n_types = sizeof(types) / sizeof(types[0]);

        uint32_t type =

            (rnd() % 10 <= 8) ? Type(types[rnd() % n_types]) : (uint32_t)rnd();

        size_t len = rnd() % 256;

        if (type == Type("fUZz"))

          len = 16;

        V v(len);

        for (auto &b : v) b = rnd();

        size_t pos = rnd() % (chunks_.size() + 1);

        chunks_.insert(chunks_.begin() + pos, {type, v});

      } break;

      // Any more interesting mutations with a PNG file?

      case 5: {

        auto it = std::find_if(

            chunks_.begin(), chunks_.end(),

            [](const Chunk &ch) { return ch.type == Type("fUZz"); });

        if (it != chunks_.end())

          m(it->v.data(), it->v.size(), it->v.size());

      }

    }

  }

  // Takes a random chunk from p and inserts into *this.

  void CrossOver(const PngMutator &p, unsigned int Seed) {

    if (p.chunks_.empty()) return;

    std::minstd_rand rnd(Seed);

    size_t idx = rnd() % p.chunks_.size();

    auto &ch = p.chunks_[idx];

    size_t pos = rnd() % (chunks_.size() + 1);

    chunks_.insert(chunks_.begin() + pos, ch);

  }

 private:

  void Append(V *to, const V &from) {

    to->insert(to->end(), from.begin(), from.end());

  }

  uint32_t Read4(std::istream &in) {

    uint32_t res = 0;

    in.read((char *)&res, sizeof(res));

    return res;

  }

  uint32_t ReadInteger(std::istream &in) {

    return __builtin_bswap32(Read4(in));

  }

  static uint32_t Type(const char *tagname) {

    uint32_t res;

    assert(strlen(tagname) == 4);

    memcpy(&res, tagname, 4);

    return res;

  }

  void WriteInt(std::ostream &out, uint32_t x) {

    x = __builtin_bswap32(x);

    out.write((char *)&x, sizeof(x));

  }

  // Chunk is written as:

  //  * 4-byte length

  //  * 4-byte type

  //  * the data itself

  //  * 4-byte crc (of type and data)

  void WriteChunk(std::ostream &out, const char *type, const V &chunk,

                  bool compress = false) {

    V compressed;

    const V *v = &chunk;

    if (compress) {

      compressed = Compress(chunk);

      v = &compressed;

    }

    uint32_t len = v->size();

    uint32_t crc = crc32(0, (const unsigned char *)type, 4);

    if (v->size())

      crc = crc32(crc, (const unsigned char *)v->data(), v->size());

    WriteInt(out, len);

    out.write(type, 4);

    out.write((const char*)v->data(), v->size());

    WriteInt(out, crc);

  }

  void WriteChunk(std::ostream &out, uint32_t type, const V &chunk) {

    char type_s[5];

    memcpy(type_s, &type, 4);

    type_s[4] = 0;

    WriteChunk(out, type_s, chunk);

  }

  V Uncompress(const V &compressed) {

    V v;

    static const size_t kMaxBuffer = 1 << 28;

    for (size_t sz = compressed.size() * 4; sz < kMaxBuffer; sz *= 2) {

      v.resize(sz);

      unsigned long len = sz;

      auto res =

          uncompress(v.data(), &len, compressed.data(), compressed.size());

      if (res == Z_BUF_ERROR) continue;

      if (res != Z_OK) return {};

      v.resize(len);

      break;

    }

    return v;

  }

  V Compress(const V &uncompressed) {

    V v;

    static const size_t kMaxBuffer = 1 << 28;

    for (size_t sz = uncompressed.size(); sz < kMaxBuffer; sz *= 2) {

      v.resize(sz);

      unsigned long len = sz;

      auto res =

          compress(v.data(), &len, uncompressed.data(), uncompressed.size());

      if (res == Z_BUF_ERROR) continue;

      if (res != Z_OK) return {};

      v.resize(len);

      break;

    }

    return v;

  }

  void PrintHex(const V &v, size_t max_n) {

    for (size_t i = 0; i < max_n && i < v.size(); i++) {

      std::cerr << "0x" << std::hex << (unsigned)v[i] << " " << std::dec;

    }

    std::cerr << std::endl;

  }

  V ihdr_;

  struct Chunk {

    uint32_t type;

    V v;

  };

  std::vector<Chunk> chunks_;

};

#ifdef PNG_MUTATOR_DEFINE_LIBFUZZER_CUSTOM_MUTATOR

extern "C" size_t LLVMFuzzerMutate(uint8_t *Data, size_t Size, size_t MaxSize);

#if STANDALONE_TARGET

size_t LLVMFuzzerMutate(uint8_t *Data, size_t Size, size_t MaxSize) {

  assert(false && "LLVMFuzzerMutate should not be called from StandaloneFuzzTargetMain");

  return 0;

}

#endif

extern "C" size_t LLVMFuzzerCustomMutator(uint8_t *Data, size_t Size,

                                          size_t MaxSize, unsigned int Seed) {

  std::string s(reinterpret_cast<const char*>(Data), Size);

  std::stringstream in(s);

  std::stringstream out;

  PngMutator p(in);

  p.Mutate(LLVMFuzzerMutate, Seed);

  p.Serialize(out);

  const auto &str = out.str();

  if (str.size() > MaxSize) return Size;

  memcpy(Data, str.data(), str.size());

  return str.size();

}

extern "C" size_t LLVMFuzzerCustomCrossOver(const uint8_t *Data1, size_t Size1,

                                            const uint8_t *Data2, size_t Size2,

                                            uint8_t *Out, size_t MaxOutSize,

                                            unsigned int Seed) {

  std::stringstream in1(

      std::string(reinterpret_cast<const char *>(Data1), Size1));

  std::stringstream in2(

      std::string(reinterpret_cast<const char *>(Data2), Size2));

  PngMutator p1(in1);

  PngMutator p2(in2);

  p1.CrossOver(p2, Seed);

  std::stringstream out;

  p1.Serialize(out);

  const auto &str = out.str();

  if (str.size() > MaxOutSize) return 0;

  memcpy(Out, str.data(), str.size());

  return str.size();

}

#endif  // PNG_MUTATOR_DEFINE_LIBFUZZER_CUSTOM_MUTATOR