libjxl

FORK: libjxl patches used on blog
git clone https://git.neptards.moe/blog/libjxl.git
Log | Files | Refs | Submodules | README | LICENSE

jpegli_dec_fuzzer_corpus.cc (11596B)


      1 // Copyright (c) the JPEG XL Project Authors. All rights reserved.
      2 //
      3 // Use of this source code is governed by a BSD-style
      4 // license that can be found in the LICENSE file.
      5 
      6 #include <setjmp.h>
      7 #include <stdint.h>
      8 #include <stdio.h>
      9 #include <stdlib.h>
     10 #include <sys/stat.h>
     11 #include <sys/types.h>
     12 #if defined(_WIN32) || defined(_WIN64)
     13 #include "third_party/dirent.h"
     14 #else
     15 #include <dirent.h>
     16 #include <unistd.h>
     17 #endif
     18 
     19 #include <algorithm>
     20 #include <iostream>
     21 #include <mutex>
     22 #include <random>
     23 #include <vector>
     24 
     25 #include "lib/jpegli/encode.h"
     26 #include "lib/jxl/base/data_parallel.h"
     27 #include "lib/jxl/base/random.h"
     28 #include "tools/file_io.h"
     29 #include "tools/thread_pool_internal.h"
     30 
     31 namespace {
     32 
     33 const size_t kMaxWidth = 50000;
     34 const size_t kMaxHeight = 50000;
     35 const size_t kMaxPixels = 20 * (1 << 20);  // 20 MP
     36 
     37 std::mutex stderr_mutex;
     38 
     39 std::vector<uint8_t> GetSomeTestImage(size_t xsize, size_t ysize,
     40                                       size_t num_channels, uint16_t seed) {
     41   // Cause more significant image difference for successive seeds.
     42   jxl::Rng generator(seed);
     43 
     44   // Returns random integer in interval [0, max_value)
     45   auto rng = [&generator](size_t max_value) -> size_t {
     46     return generator.UniformU(0, max_value);
     47   };
     48 
     49   // Dark background gradient color
     50   uint16_t r0 = rng(32768);
     51   uint16_t g0 = rng(32768);
     52   uint16_t b0 = rng(32768);
     53   uint16_t r1 = rng(32768);
     54   uint16_t g1 = rng(32768);
     55   uint16_t b1 = rng(32768);
     56 
     57   // Circle with different color
     58   size_t circle_x = rng(xsize);
     59   size_t circle_y = rng(ysize);
     60   size_t circle_r = rng(std::min(xsize, ysize));
     61 
     62   // Rectangle with random noise
     63   size_t rect_x0 = rng(xsize);
     64   size_t rect_y0 = rng(ysize);
     65   size_t rect_x1 = rng(xsize);
     66   size_t rect_y1 = rng(ysize);
     67   if (rect_x1 < rect_x0) std::swap(rect_x0, rect_y1);
     68   if (rect_y1 < rect_y0) std::swap(rect_y0, rect_y1);
     69 
     70   size_t num_pixels = xsize * ysize;
     71   std::vector<uint8_t> pixels(num_pixels * num_channels);
     72   // Create pixel content to test.
     73   for (size_t y = 0; y < ysize; y++) {
     74     for (size_t x = 0; x < xsize; x++) {
     75       uint16_t r = r0 * (ysize - y - 1) / ysize + r1 * y / ysize;
     76       uint16_t g = g0 * (ysize - y - 1) / ysize + g1 * y / ysize;
     77       uint16_t b = b0 * (ysize - y - 1) / ysize + b1 * y / ysize;
     78       // put some shape in there for visual debugging
     79       if ((x - circle_x) * (x - circle_x) + (y - circle_y) * (y - circle_y) <
     80           circle_r * circle_r) {
     81         r = (65535 - x * y) ^ seed;
     82         g = (x << 8) + y + seed;
     83         b = (y << 8) + x * seed;
     84       } else if (x > rect_x0 && x < rect_x1 && y > rect_y0 && y < rect_y1) {
     85         r = rng(65536);
     86         g = rng(65536);
     87         b = rng(65536);
     88       }
     89       size_t i = (y * xsize + x) * num_channels;
     90       pixels[i + 0] = (r >> 8);
     91       if (num_channels == 3) {
     92         pixels[i + 1] = (g >> 8);
     93         pixels[i + 2] = (b >> 8);
     94       }
     95     }
     96   }
     97   return pixels;
     98 }
     99 
    100 // ImageSpec needs to be a packed struct to allow us to use the raw memory of
    101 // the struct for hashing to create a consistent id.
    102 #pragma pack(push, 1)
    103 struct ImageSpec {
    104   bool Validate() const {
    105     if (width > kMaxWidth || height > kMaxHeight ||
    106         width * height > kMaxPixels) {
    107       return false;
    108     }
    109     return true;
    110   }
    111 
    112   friend std::ostream& operator<<(std::ostream& o, const ImageSpec& spec) {
    113     o << "ImageSpec<"
    114       << "size=" << spec.width << "x" << spec.height
    115       << " * chan=" << spec.num_channels << " q=" << spec.quality
    116       << " p=" << spec.progressive_level << " r=" << spec.restart_interval
    117       << ">";
    118     return o;
    119   }
    120 
    121   void SpecHash(uint8_t hash[16]) const {
    122     const uint8_t* from = reinterpret_cast<const uint8_t*>(this);
    123     std::seed_seq hasher(from, from + sizeof(*this));
    124     uint32_t* to = reinterpret_cast<uint32_t*>(hash);
    125     hasher.generate(to, to + 4);
    126   }
    127 
    128   uint32_t width = 256;
    129   uint32_t height = 256;
    130   uint32_t num_channels = 3;
    131   uint32_t quality = 90;
    132   uint32_t sampling = 0x11111111;
    133   uint32_t progressive_level = 2;
    134   uint32_t restart_interval = 0;
    135   uint32_t fraction = 100;
    136   // The seed for the PRNG.
    137   uint32_t seed = 7777;
    138 };
    139 #pragma pack(pop)
    140 static_assert(sizeof(ImageSpec) % 4 == 0, "Add padding to ImageSpec.");
    141 
    142 bool EncodeWithJpegli(const ImageSpec& spec, const std::vector<uint8_t>& pixels,
    143                       std::vector<uint8_t>* compressed) {
    144   uint8_t* buffer = nullptr;
    145   unsigned long buffer_size = 0;
    146   jpeg_compress_struct cinfo;
    147   const auto try_catch_block = [&]() -> bool {
    148     jpeg_error_mgr jerr;
    149     jmp_buf env;
    150     cinfo.err = jpegli_std_error(&jerr);
    151     if (setjmp(env)) {
    152       return false;
    153     }
    154     cinfo.client_data = reinterpret_cast<void*>(&env);
    155     cinfo.err->error_exit = [](j_common_ptr cinfo) {
    156       (*cinfo->err->output_message)(cinfo);
    157       jmp_buf* env = reinterpret_cast<jmp_buf*>(cinfo->client_data);
    158       jpegli_destroy(cinfo);
    159       longjmp(*env, 1);
    160     };
    161     jpegli_create_compress(&cinfo);
    162     jpegli_mem_dest(&cinfo, &buffer, &buffer_size);
    163     cinfo.image_width = spec.width;
    164     cinfo.image_height = spec.height;
    165     cinfo.input_components = spec.num_channels;
    166     cinfo.in_color_space = spec.num_channels == 1 ? JCS_GRAYSCALE : JCS_RGB;
    167     jpegli_set_defaults(&cinfo);
    168     jpegli_set_quality(&cinfo, spec.quality, TRUE);
    169     uint32_t sampling = spec.sampling;
    170     for (int c = 0; c < cinfo.num_components; ++c) {
    171       cinfo.comp_info[c].h_samp_factor = sampling & 0xf;
    172       cinfo.comp_info[c].v_samp_factor = (sampling >> 4) & 0xf;
    173       sampling >>= 8;
    174     }
    175     jpegli_set_progressive_level(&cinfo, spec.progressive_level);
    176     cinfo.restart_interval = spec.restart_interval;
    177     jpegli_start_compress(&cinfo, TRUE);
    178     size_t stride = cinfo.image_width * cinfo.input_components;
    179     std::vector<uint8_t> row_bytes(stride);
    180     for (size_t y = 0; y < cinfo.image_height; ++y) {
    181       memcpy(row_bytes.data(), &pixels[y * stride], stride);
    182       JSAMPROW row[] = {row_bytes.data()};
    183       jpegli_write_scanlines(&cinfo, row, 1);
    184     }
    185     jpegli_finish_compress(&cinfo);
    186     return true;
    187   };
    188   bool success = try_catch_block();
    189   jpegli_destroy_compress(&cinfo);
    190   if (success) {
    191     buffer_size = buffer_size * spec.fraction / 100;
    192     compressed->assign(buffer, buffer + buffer_size);
    193   }
    194   if (buffer) std::free(buffer);
    195   return success;
    196 }
    197 
    198 bool GenerateFile(const char* output_dir, const ImageSpec& spec,
    199                   bool regenerate, bool quiet) {
    200   // Compute a checksum of the ImageSpec to name the file. This is just to keep
    201   // the output of this program repeatable.
    202   uint8_t checksum[16];
    203   spec.SpecHash(checksum);
    204   std::string hash_str(sizeof(checksum) * 2, ' ');
    205   static const char* hex_chars = "0123456789abcdef";
    206   for (size_t i = 0; i < sizeof(checksum); i++) {
    207     hash_str[2 * i] = hex_chars[checksum[i] >> 4];
    208     hash_str[2 * i + 1] = hex_chars[checksum[i] % 0x0f];
    209   }
    210   std::string output_fn = std::string(output_dir) + "/" + hash_str + ".jpg";
    211 
    212   // Don't regenerate files if they already exist on disk to speed-up
    213   // consecutive calls when --regenerate is not used.
    214   struct stat st;
    215   if (!regenerate && stat(output_fn.c_str(), &st) == 0 && S_ISREG(st.st_mode)) {
    216     return true;
    217   }
    218 
    219   if (!quiet) {
    220     std::unique_lock<std::mutex> lock(stderr_mutex);
    221     std::cerr << "Generating " << spec << " as " << hash_str << std::endl;
    222   }
    223 
    224   uint8_t hash[16];
    225   spec.SpecHash(hash);
    226   std::mt19937 mt(spec.seed);
    227 
    228   std::vector<uint8_t> pixels =
    229       GetSomeTestImage(spec.width, spec.height, spec.num_channels, spec.seed);
    230   std::vector<uint8_t> compressed;
    231   JXL_CHECK(EncodeWithJpegli(spec, pixels, &compressed));
    232 
    233   // Append 4 bytes with the flags used by jpegli_dec_fuzzer to select the
    234   // decoding output.
    235   std::uniform_int_distribution<> dis256(0, 255);
    236   for (size_t i = 0; i < 4; ++i) {
    237     compressed.push_back(dis256(mt));
    238   }
    239 
    240   if (!jpegxl::tools::WriteFile(output_fn, compressed)) {
    241     return false;
    242   }
    243   if (!quiet) {
    244     std::unique_lock<std::mutex> lock(stderr_mutex);
    245     std::cerr << "Stored " << output_fn << " size: " << compressed.size()
    246               << std::endl;
    247   }
    248 
    249   return true;
    250 }
    251 
    252 void Usage() {
    253   fprintf(stderr,
    254           "Use: fuzzer_corpus [-r] [-q] [-j THREADS] [output_dir]\n"
    255           "\n"
    256           "  -r Regenerate files if already exist.\n"
    257           "  -q Be quiet.\n"
    258           "  -j THREADS Number of parallel jobs to run.\n");
    259 }
    260 
    261 }  // namespace
    262 
    263 int main(int argc, const char** argv) {
    264   const char* dest_dir = nullptr;
    265   bool regenerate = false;
    266   bool quiet = false;
    267   size_t num_threads = std::thread::hardware_concurrency();
    268   for (int optind = 1; optind < argc;) {
    269     if (!strcmp(argv[optind], "-r")) {
    270       regenerate = true;
    271       optind++;
    272     } else if (!strcmp(argv[optind], "-q")) {
    273       quiet = true;
    274       optind++;
    275     } else if (!strcmp(argv[optind], "-j")) {
    276       optind++;
    277       if (optind < argc) {
    278         num_threads = atoi(argv[optind++]);
    279       } else {
    280         fprintf(stderr, "-j needs an argument value.\n");
    281         Usage();
    282         return 1;
    283       }
    284     } else if (dest_dir == nullptr) {
    285       dest_dir = argv[optind++];
    286     } else {
    287       fprintf(stderr, "Unknown parameter: \"%s\".\n", argv[optind]);
    288       Usage();
    289       return 1;
    290     }
    291   }
    292   if (!dest_dir) {
    293     dest_dir = "corpus";
    294   }
    295 
    296   struct stat st;
    297   memset(&st, 0, sizeof(st));
    298   if (stat(dest_dir, &st) != 0 || !S_ISDIR(st.st_mode)) {
    299     fprintf(stderr, "Output path \"%s\" is not a directory.\n", dest_dir);
    300     Usage();
    301     return 1;
    302   }
    303 
    304   std::mt19937 mt(77777);
    305 
    306   std::vector<std::pair<uint32_t, uint32_t>> image_sizes = {
    307       {8, 8},     {32, 32},   {128, 128}, {10000, 1}, {10000, 2}, {1, 10000},
    308       {2, 10000}, {555, 256}, {257, 513}, {512, 265}, {264, 520},
    309   };
    310   std::vector<uint32_t> sampling_ratios = {
    311       0x11111111,  // 444
    312       0x11111112,  // 422
    313       0x11111121,  // 440
    314       0x11111122,  // 420
    315       0x11222211,  // luma subsampling
    316   };
    317 
    318   ImageSpec spec;
    319   std::vector<ImageSpec> specs;
    320   for (auto img_size : image_sizes) {
    321     spec.width = img_size.first;
    322     spec.height = img_size.second;
    323     for (uint32_t num_channels : {1, 3}) {
    324       spec.num_channels = num_channels;
    325       for (uint32_t sampling : sampling_ratios) {
    326         spec.sampling = sampling;
    327         if (num_channels == 1 && sampling != 0x11111111) continue;
    328         for (uint32_t restart : {0, 1, 1024}) {
    329           spec.restart_interval = restart;
    330           for (uint32_t prog_level : {0, 1, 2}) {
    331             spec.progressive_level = prog_level;
    332             for (uint32_t quality : {10, 90, 100}) {
    333               spec.quality = quality;
    334               for (uint32_t fraction : {10, 70, 100}) {
    335                 spec.fraction = fraction;
    336                 spec.seed = mt() % 777777;
    337                 if (!spec.Validate()) {
    338                   if (!quiet) {
    339                     std::cerr << "Skipping " << spec << std::endl;
    340                   }
    341                 } else {
    342                   specs.push_back(spec);
    343                 }
    344               }
    345             }
    346           }
    347         }
    348       }
    349     }
    350   }
    351 
    352   jpegxl::tools::ThreadPoolInternal pool{num_threads};
    353   const auto generate = [&specs, dest_dir, regenerate, quiet](
    354                             const uint32_t task, size_t /* thread */) {
    355     const ImageSpec& spec = specs[task];
    356     GenerateFile(dest_dir, spec, regenerate, quiet);
    357   };
    358   if (!RunOnPool(&pool, 0, specs.size(), jxl::ThreadPool::NoInit, generate,
    359                  "FuzzerCorpus")) {
    360     std::cerr << "Error generating fuzzer corpus" << std::endl;
    361     return 1;
    362   }
    363   std::cerr << "Finished generating fuzzer corpus" << std::endl;
    364   return 0;
    365 }