tor-browser

enc_icc_codec.cc (17955B)

---

// Copyright (c) the JPEG XL Project Authors. All rights reserved.
//
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

#include "lib/jxl/enc_icc_codec.h"

#include <jxl/memory_manager.h>

#include <cstdint>
#include <limits>
#include <map>
#include <vector>

#include "lib/jxl/base/status.h"
#include "lib/jxl/enc_ans.h"
#include "lib/jxl/enc_aux_out.h"
#include "lib/jxl/fields.h"
#include "lib/jxl/icc_codec_common.h"
#include "lib/jxl/padded_bytes.h"

namespace jxl {
namespace {

// Unshuffles or de-interleaves bytes, for example with width 2, turns
// "AaBbCcDc" into "ABCDabcd", this for example de-interleaves UTF-16 bytes into
// first all the high order bytes, then all the low order bytes.
// Transposes a matrix of width columns and ceil(size / width) rows. There are
// size elements, size may be < width * height, if so the
// last elements of the bottom row are missing, the missing spots are
// transposed along with the filled spots, and the result has the missing
// elements at the bottom of the rightmost column. The input is the input matrix
// in scanline order, the output is the result matrix in scanline order, with
// missing elements skipped over (this may occur at multiple positions).
Status Unshuffle(JxlMemoryManager* memory_manager, uint8_t* data, size_t size,
                size_t width) {
 size_t height = (size + width - 1) / width;  // amount of rows of input
 PaddedBytes result(memory_manager);
 JXL_ASSIGN_OR_RETURN(result,
                      PaddedBytes::WithInitialSpace(memory_manager, size));

 // i = input index, j output index
 size_t s = 0;
 size_t j = 0;
 for (size_t i = 0; i < size; i++) {
   result[j] = data[i];
   j += height;
   if (j >= size) j = ++s;
 }

 for (size_t i = 0; i < size; i++) {
   data[i] = result[i];
 }
 return true;
}

// This is performed by the encoder, the encoder must be able to encode any
// random byte stream (not just byte streams that are a valid ICC profile), so
// an error returned by this function is an implementation error.
Status PredictAndShuffle(size_t stride, size_t width, int order, size_t num,
                        const uint8_t* data, size_t size, size_t* pos,
                        PaddedBytes* result) {
 JXL_RETURN_IF_ERROR(CheckOutOfBounds(*pos, num, size));
 JxlMemoryManager* memory_manager = result->memory_manager();
 // Required by the specification, see decoder. stride * 4 must be < *pos.
 if (!*pos || ((*pos - 1u) >> 2u) < stride) {
   return JXL_FAILURE("Invalid stride");
 }
 if (*pos < stride * 4) return JXL_FAILURE("Too large stride");
 size_t start = result->size();
 for (size_t i = 0; i < num; i++) {
   uint8_t predicted =
       LinearPredictICCValue(data, *pos, i, stride, width, order);
   JXL_RETURN_IF_ERROR(result->push_back(data[*pos + i] - predicted));
 }
 *pos += num;
 if (width > 1) {
   JXL_RETURN_IF_ERROR(
       Unshuffle(memory_manager, result->data() + start, num, width));
 }
 return true;
}

inline Status EncodeVarInt(uint64_t value, PaddedBytes* data) {
 size_t pos = data->size();
 JXL_RETURN_IF_ERROR(data->resize(data->size() + 9));
 size_t output_size = data->size();
 uint8_t* output = data->data();

 // While more than 7 bits of data are left,
 // store 7 bits and set the next byte flag
 while (value > 127) {
   // TODO(eustas): should it be `<` ?
   JXL_ENSURE(pos <= output_size);
   // |128: Set the next byte flag
   output[pos++] = (static_cast<uint8_t>(value & 127)) | 128;
   // Remove the seven bits we just wrote
   value >>= 7;
 }
 // TODO(eustas): should it be `<` ?
 JXL_ENSURE(pos <= output_size);
 output[pos++] = static_cast<uint8_t>(value & 127);

 return data->resize(pos);
}

constexpr size_t kSizeLimit = std::numeric_limits<uint32_t>::max() >> 2;

}  // namespace

// Outputs a transformed form of the given icc profile. The result itself is
// not particularly smaller than the input data in bytes, but it will be in a
// form that is easier to compress (more zeroes, ...) and will compress better
// with brotli.
Status PredictICC(const uint8_t* icc, size_t size, PaddedBytes* result) {
 JxlMemoryManager* memory_manager = result->memory_manager();
 PaddedBytes commands{memory_manager};
 PaddedBytes data{memory_manager};

 static_assert(sizeof(size_t) >= 4, "size_t is too short");
 // Fuzzer expects that PredictICC can accept any input,
 // but 1GB should be enough for any purpose.
 if (size > kSizeLimit) {
   return JXL_FAILURE("ICC profile is too large");
 }

 JXL_RETURN_IF_ERROR(EncodeVarInt(size, result));

 // Header
 PaddedBytes header{memory_manager};
 JXL_RETURN_IF_ERROR(header.append(ICCInitialHeaderPrediction(size)));
 for (size_t i = 0; i < kICCHeaderSize && i < size; i++) {
   ICCPredictHeader(icc, size, header.data(), i);
   JXL_RETURN_IF_ERROR(data.push_back(icc[i] - header[i]));
 }
 if (size <= kICCHeaderSize) {
   JXL_RETURN_IF_ERROR(EncodeVarInt(0, result));  // 0 commands
   for (uint8_t b : data) {
     JXL_RETURN_IF_ERROR(result->push_back(b));
   }
   return true;
 }

 std::vector<Tag> tags;
 std::vector<size_t> tagstarts;
 std::vector<size_t> tagsizes;
 std::map<size_t, size_t> tagmap;

 // Tag list
 size_t pos = kICCHeaderSize;
 if (pos + 4 <= size) {
   uint64_t numtags = DecodeUint32(icc, size, pos);
   pos += 4;
   JXL_RETURN_IF_ERROR(EncodeVarInt(numtags + 1, &commands));
   uint64_t prevtagstart = kICCHeaderSize + numtags * 12;
   uint32_t prevtagsize = 0;
   for (size_t i = 0; i < numtags; i++) {
     if (pos + 12 > size) break;

     Tag tag = DecodeKeyword(icc, size, pos + 0);
     uint32_t tagstart = DecodeUint32(icc, size, pos + 4);
     uint32_t tagsize = DecodeUint32(icc, size, pos + 8);
     pos += 12;

     tags.push_back(tag);
     tagstarts.push_back(tagstart);
     tagsizes.push_back(tagsize);
     tagmap[tagstart] = tags.size() - 1;

     uint8_t tagcode = kCommandTagUnknown;
     for (size_t j = 0; j < kNumTagStrings; j++) {
       if (tag == *kTagStrings[j]) {
         tagcode = j + kCommandTagStringFirst;
         break;
       }
     }

     if (tag == kRtrcTag && pos + 24 < size) {
       bool ok = true;
       ok &= DecodeKeyword(icc, size, pos + 0) == kGtrcTag;
       ok &= DecodeKeyword(icc, size, pos + 12) == kBtrcTag;
       if (ok) {
         for (size_t kk = 0; kk < 8; kk++) {
           if (icc[pos - 8 + kk] != icc[pos + 4 + kk]) ok = false;
           if (icc[pos - 8 + kk] != icc[pos + 16 + kk]) ok = false;
         }
       }
       if (ok) {
         tagcode = kCommandTagTRC;
         pos += 24;
         i += 2;
       }
     }

     if (tag == kRxyzTag && pos + 24 < size) {
       bool ok = true;
       ok &= DecodeKeyword(icc, size, pos + 0) == kGxyzTag;
       ok &= DecodeKeyword(icc, size, pos + 12) == kBxyzTag;
       uint32_t offsetr = tagstart;
       uint32_t offsetg = DecodeUint32(icc, size, pos + 4);
       uint32_t offsetb = DecodeUint32(icc, size, pos + 16);
       uint32_t sizer = tagsize;
       uint32_t sizeg = DecodeUint32(icc, size, pos + 8);
       uint32_t sizeb = DecodeUint32(icc, size, pos + 20);
       ok &= sizer == 20;
       ok &= sizeg == 20;
       ok &= sizeb == 20;
       ok &= (offsetg == offsetr + 20);
       ok &= (offsetb == offsetr + 40);
       if (ok) {
         tagcode = kCommandTagXYZ;
         pos += 24;
         i += 2;
       }
     }

     uint8_t command = tagcode;
     uint64_t predicted_tagstart = prevtagstart + prevtagsize;
     if (predicted_tagstart != tagstart) command |= kFlagBitOffset;
     size_t predicted_tagsize = prevtagsize;
     if (tag == kRxyzTag || tag == kGxyzTag || tag == kBxyzTag ||
         tag == kKxyzTag || tag == kWtptTag || tag == kBkptTag ||
         tag == kLumiTag) {
       predicted_tagsize = 20;
     }
     if (predicted_tagsize != tagsize) command |= kFlagBitSize;
     JXL_RETURN_IF_ERROR(commands.push_back(command));
     if (tagcode == 1) {
       JXL_RETURN_IF_ERROR(AppendKeyword(tag, &data));
     }
     if (command & kFlagBitOffset)
       JXL_RETURN_IF_ERROR(EncodeVarInt(tagstart, &commands));
     if (command & kFlagBitSize)
       JXL_RETURN_IF_ERROR(EncodeVarInt(tagsize, &commands));

     prevtagstart = tagstart;
     prevtagsize = tagsize;
   }
 }
 // Indicate end of tag list or varint indicating there's none
 JXL_RETURN_IF_ERROR(commands.push_back(0));

 // Main content
 // The main content in a valid ICC profile contains tagged elements, with the
 // tag types (4 letter names) given by the tag list above, and the tag list
 // pointing to the start and indicating the size of each tagged element. It is
 // allowed for tagged elements to overlap, e.g. the curve for R, G and B could
 // all point to the same one.
 Tag tag;
 size_t tagstart = 0;
 size_t tagsize = 0;
 size_t clutstart = 0;

 // Should always check tag_sane before doing math with tagsize.
 const auto tag_sane = [&tagsize]() {
   return (tagsize > 8) && (tagsize < kSizeLimit);
 };

 size_t last0 = pos;
 // This loop appends commands to the output, processing some sub-section of a
 // current tagged element each time. We need to keep track of the tagtype of
 // the current element, and update it when we encounter the boundary of a
 // next one.
 // It is not required that the input data is a valid ICC profile, if the
 // encoder does not recognize the data it will still be able to output bytes
 // but will not predict as well.
 while (pos <= size) {
   size_t last1 = pos;
   PaddedBytes commands_add{memory_manager};
   PaddedBytes data_add{memory_manager};

   // This means the loop brought the position beyond the tag end.
   // If tagsize is nonsensical, any pos looks "ok-ish".
   if ((pos > tagstart + tagsize) && (tagsize < kSizeLimit)) {
     tag = {{0, 0, 0, 0}};  // nonsensical value
   }

   if (commands_add.empty() && data_add.empty() && tagmap.count(pos) &&
       pos + 4 <= size) {
     size_t index = tagmap[pos];
     tag = DecodeKeyword(icc, size, pos);
     tagstart = tagstarts[index];
     tagsize = tagsizes[index];

     if (tag == kMlucTag && tag_sane() && pos + tagsize <= size &&
         icc[pos + 4] == 0 && icc[pos + 5] == 0 && icc[pos + 6] == 0 &&
         icc[pos + 7] == 0) {
       size_t num = tagsize - 8;
       JXL_RETURN_IF_ERROR(commands_add.push_back(kCommandTypeStartFirst + 3));
       pos += 8;
       JXL_RETURN_IF_ERROR(commands_add.push_back(kCommandShuffle2));
       JXL_RETURN_IF_ERROR(EncodeVarInt(num, &commands_add));
       size_t start = data_add.size();
       for (size_t i = 0; i < num; i++) {
         JXL_RETURN_IF_ERROR(data_add.push_back(icc[pos]));
         pos++;
       }
       JXL_RETURN_IF_ERROR(
           Unshuffle(memory_manager, data_add.data() + start, num, 2));
     }

     if (tag == kCurvTag && tag_sane() && pos + tagsize <= size &&
         icc[pos + 4] == 0 && icc[pos + 5] == 0 && icc[pos + 6] == 0 &&
         icc[pos + 7] == 0) {
       size_t num = tagsize - 8;
       if (num > 16 && num < (1 << 28) && pos + num <= size && pos > 0) {
         JXL_RETURN_IF_ERROR(
             commands_add.push_back(kCommandTypeStartFirst + 5));
         pos += 8;
         JXL_RETURN_IF_ERROR(commands_add.push_back(kCommandPredict));
         int order = 1;
         int width = 2;
         int stride = width;
         JXL_RETURN_IF_ERROR(
             commands_add.push_back((order << 2) | (width - 1)));
         JXL_RETURN_IF_ERROR(EncodeVarInt(num, &commands_add));
         JXL_RETURN_IF_ERROR(PredictAndShuffle(stride, width, order, num, icc,
                                               size, &pos, &data_add));
       }
     }
   }

   if (tag == kMab_Tag || tag == kMba_Tag) {
     Tag subTag = DecodeKeyword(icc, size, pos);
     if (pos + 12 < size && (subTag == kCurvTag || subTag == kVcgtTag) &&
         DecodeUint32(icc, size, pos + 4) == 0) {
       uint32_t num = DecodeUint32(icc, size, pos + 8) * 2;
       if (num > 16 && num < (1 << 28) && pos + 12 + num <= size) {
         pos += 12;
         last1 = pos;
         JXL_RETURN_IF_ERROR(commands_add.push_back(kCommandPredict));
         int order = 1;
         int width = 2;
         int stride = width;
         JXL_RETURN_IF_ERROR(
             commands_add.push_back((order << 2) | (width - 1)));
         JXL_RETURN_IF_ERROR(EncodeVarInt(num, &commands_add));
         JXL_RETURN_IF_ERROR(PredictAndShuffle(stride, width, order, num, icc,
                                               size, &pos, &data_add));
       }
     }

     if (pos == tagstart + 24 && pos + 4 < size) {
       // Note that this value can be remembered for next iterations of the
       // loop, so the "pos == clutstart" if below can trigger during a later
       // iteration.
       clutstart = tagstart + DecodeUint32(icc, size, pos);
     }

     if (pos == clutstart && clutstart + 16 < size) {
       size_t numi = icc[tagstart + 8];
       size_t numo = icc[tagstart + 9];
       size_t width = icc[clutstart + 16];
       size_t stride = width * numo;
       size_t num = width * numo;
       for (size_t i = 0; i < numi && clutstart + i < size; i++) {
         num *= icc[clutstart + i];
       }
       if ((width == 1 || width == 2) && num > 64 && num < (1 << 28) &&
           pos + num <= size && pos > stride * 4) {
         JXL_RETURN_IF_ERROR(commands_add.push_back(kCommandPredict));
         int order = 1;
         uint8_t flags =
             (order << 2) | (width - 1) | (stride == width ? 0 : 16);
         JXL_RETURN_IF_ERROR(commands_add.push_back(flags));
         if (flags & 16) {
           JXL_RETURN_IF_ERROR(EncodeVarInt(stride, &commands_add));
         }
         JXL_RETURN_IF_ERROR(EncodeVarInt(num, &commands_add));
         JXL_RETURN_IF_ERROR(PredictAndShuffle(stride, width, order, num, icc,
                                               size, &pos, &data_add));
       }
     }
   }

   if (commands_add.empty() && data_add.empty() && tag == kGbd_Tag &&
       tag_sane() && pos == tagstart + 8 && pos + tagsize - 8 <= size &&
       pos > 16) {
     size_t width = 4;
     size_t order = 0;
     size_t stride = width;
     size_t num = tagsize - 8;
     uint8_t flags = (order << 2) | (width - 1) | (stride == width ? 0 : 16);
     JXL_RETURN_IF_ERROR(commands_add.push_back(kCommandPredict));
     JXL_RETURN_IF_ERROR(commands_add.push_back(flags));
     if (flags & 16) {
       JXL_RETURN_IF_ERROR(EncodeVarInt(stride, &commands_add));
     }
     JXL_RETURN_IF_ERROR(EncodeVarInt(num, &commands_add));
     JXL_RETURN_IF_ERROR(PredictAndShuffle(stride, width, order, num, icc,
                                           size, &pos, &data_add));
   }

   if (commands_add.empty() && data_add.empty() && pos + 20 <= size) {
     Tag subTag = DecodeKeyword(icc, size, pos);
     if (subTag == kXyz_Tag && DecodeUint32(icc, size, pos + 4) == 0) {
       JXL_RETURN_IF_ERROR(commands_add.push_back(kCommandXYZ));
       pos += 8;
       for (size_t j = 0; j < 12; j++) {
         JXL_RETURN_IF_ERROR(data_add.push_back(icc[pos++]));
       }
     }
   }

   if (commands_add.empty() && data_add.empty() && pos + 8 <= size) {
     if (DecodeUint32(icc, size, pos + 4) == 0) {
       Tag subTag = DecodeKeyword(icc, size, pos);
       for (size_t i = 0; i < kNumTypeStrings; i++) {
         if (subTag == *kTypeStrings[i]) {
           JXL_RETURN_IF_ERROR(
               commands_add.push_back(kCommandTypeStartFirst + i));
           pos += 8;
           break;
         }
       }
     }
   }

   if (!(commands_add.empty() && data_add.empty()) || pos == size) {
     if (last0 < last1) {
       JXL_RETURN_IF_ERROR(commands.push_back(kCommandInsert));
       JXL_RETURN_IF_ERROR(EncodeVarInt(last1 - last0, &commands));
       while (last0 < last1) {
         JXL_RETURN_IF_ERROR(data.push_back(icc[last0++]));
       }
     }
     for (uint8_t b : commands_add) {
       JXL_RETURN_IF_ERROR(commands.push_back(b));
     }
     for (uint8_t b : data_add) {
       JXL_RETURN_IF_ERROR(data.push_back(b));
     }
     last0 = pos;
   }
   if (commands_add.empty() && data_add.empty()) {
     pos++;
   }
 }

 JXL_RETURN_IF_ERROR(EncodeVarInt(commands.size(), result));
 for (uint8_t b : commands) {
   JXL_RETURN_IF_ERROR(result->push_back(b));
 }
 for (uint8_t b : data) {
   JXL_RETURN_IF_ERROR(result->push_back(b));
 }

 return true;
}

Status WriteICC(const Span<const uint8_t> icc, BitWriter* JXL_RESTRICT writer,
               LayerType layer, AuxOut* JXL_RESTRICT aux_out) {
 if (icc.empty()) return JXL_FAILURE("ICC must be non-empty");
 JxlMemoryManager* memory_manager = writer->memory_manager();
 PaddedBytes enc{memory_manager};
 JXL_RETURN_IF_ERROR(PredictICC(icc.data(), icc.size(), &enc));
 std::vector<std::vector<Token>> tokens(1);
 JXL_RETURN_IF_ERROR(writer->WithMaxBits(128, layer, aux_out, [&] {
   return U64Coder::Write(enc.size(), writer);
 }));

 for (size_t i = 0; i < enc.size(); i++) {
   tokens[0].emplace_back(
       ICCANSContext(i, i > 0 ? enc[i - 1] : 0, i > 1 ? enc[i - 2] : 0),
       enc[i]);
 }
 HistogramParams params;
 params.lz77_method = enc.size() < 4096 ? HistogramParams::LZ77Method::kOptimal
                                        : HistogramParams::LZ77Method::kLZ77;
 EntropyEncodingData code;
 std::vector<uint8_t> context_map;
 params.force_huffman = true;
 JXL_ASSIGN_OR_RETURN(
     size_t cost,
     BuildAndEncodeHistograms(memory_manager, params, kNumICCContexts, tokens,
                              &code, &context_map, writer, layer, aux_out));
 (void)cost;
 JXL_RETURN_IF_ERROR(
     WriteTokens(tokens[0], code, context_map, 0, writer, layer, aux_out));
 return true;
}

}  // namespace jxl