• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1 /* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
2 
3 Licensed under the Apache License, Version 2.0 (the "License");
4 you may not use this file except in compliance with the License.
5 You may obtain a copy of the License at
6 
7     http://www.apache.org/licenses/LICENSE-2.0
8 
9 Unless required by applicable law or agreed to in writing, software
10 distributed under the License is distributed on an "AS IS" BASIS,
11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 See the License for the specific language governing permissions and
13 limitations under the License.
14 ==============================================================================*/
15 #include "tensorflow/lite/experimental/acceleration/mini_benchmark/libjpeg_decoder.h"
16 
17 #include <setjmp.h>
18 
19 #include <algorithm>
20 #include <cctype>
21 #include <cstddef>
22 #include <cstdio>
23 #include <cstring>
24 #include <functional>
25 #include <limits>
26 #include <memory>
27 #include <string>
28 
29 #include "absl/strings/match.h"
30 #include "absl/strings/string_view.h"
31 #include "tensorflow/lite/c/c_api_types.h"
32 #include "tensorflow/lite/experimental/acceleration/mini_benchmark/decode_jpeg_status.h"
33 #include "tensorflow/lite/experimental/acceleration/mini_benchmark/jpeg_decompress_buffered_struct.h"
34 #include "tensorflow/lite/experimental/acceleration/mini_benchmark/jpeg_header_parser.h"
35 #include "tensorflow/lite/experimental/acceleration/mini_benchmark/libjpeg_handle.h"
36 #include "tensorflow/lite/minimal_logging.h"
37 #include "tensorflow/lite/string_type.h"
38 #include "tensorflow/lite/string_util.h"
39 
40 namespace tflite {
41 namespace acceleration {
42 namespace decode_jpeg_kernel {
43 
44 // Limiting max image size to 10,000x10,000x3
45 // This size would fit on 32 bit systems.
46 // static
47 const size_t LibjpegDecoder::kMaxImageHeight = 10000;
48 // static
49 const size_t LibjpegDecoder::kMaxImageWidth = 10000;
50 
51 constexpr char kSizeMismatchError[] =
52     "JPEG parameter struct mismatch: library thinks size is ";
53 
Impl(size_t decompress_struct_size,const LibjpegHandle * handle)54 LibjpegDecoder::Impl::Impl(size_t decompress_struct_size,
55                            const LibjpegHandle* handle)
56     : decompress_struct_size_(decompress_struct_size),
57       handle_(handle),
58       cinfo_(decompress_struct_size) {
59   cinfo_.get()->err = handle->jpeg_std_error_(&jerr_);
60   jerr_.error_exit = ErrorExit;
61   cinfo_.get()->client_data = this;
62 }
63 
ErrorExit(j_common_ptr cinfo)64 void LibjpegDecoder::Impl::ErrorExit(j_common_ptr cinfo) {
65   Impl* const impl = reinterpret_cast<Impl*>(cinfo->client_data);
66   char message[JMSG_LENGTH_MAX];
67   cinfo->err->format_message(cinfo, message);
68   impl->status_.code = kTfLiteError;
69   impl->status_.error_message = message;
70   // Libjpeg aborts the program in case of any errors by using longjmp and then
71   // calling exit(). The only way to avoid this, is to transfer the control flow
72   // to the caller by using setjmp/longjmp.
73   // Note: Ensure that function containing the corresponding setjmp() is
74   // guaranteed not to have completed execution.
75   // https://wiki.sei.cmu.edu/confluence/display/c/MSC22-C.+Use+the+setjmp%28%29%2C+longjmp%28%29+facility+securely
76   longjmp(impl->env_, 1);
77 }
78 
ExtractSizeFromErrorMessage(const std::string & error_message,size_t & expected_size)79 Status ExtractSizeFromErrorMessage(const std::string& error_message,
80                                    size_t& expected_size) {
81   Status status;
82   // Special error handling for struct mismatch issues.
83   // If there's a mismatch, set `expected_size` with the expected
84   // size. Error messages are like this: "JPEG parameter struct
85   // mismatch: library thinks size is 480, caller expects 464".
86   static const int kExpLengthStart = strlen(kSizeMismatchError);
87   int end = kExpLengthStart;
88   while (end < error_message.length() && std::isdigit(error_message[end])) {
89     end++;
90   }
91   if (end > kExpLengthStart) {
92     expected_size = std::stoi(error_message.substr(kExpLengthStart, end));
93   } else {
94     status.code = kTfLiteError;
95     status.error_message =
96         "Couldn't parse the size from message: \'" + error_message + "\'";
97   }
98   return status;
99 }
100 
Create(Status & status)101 std::unique_ptr<LibjpegDecoder> LibjpegDecoder::Create(Status& status) {
102   std::unique_ptr<LibjpegDecoder> decoder(
103       new LibjpegDecoder(LibCHandle::Create(status)));
104   if (status.code != kTfLiteOk) {
105     return nullptr;
106   }
107   decoder->libjpeg_handle_ = LibjpegHandle::Create(status);
108   if (decoder->libjpeg_handle_ == nullptr) {
109     return nullptr;
110   }
111 
112   // Tries to probe the libjpeg library to get the expected size of
113   // `jpeg_decompress_struct`.
114   Impl impl(sizeof(jpeg_decompress_struct), decoder->libjpeg_handle_.get());
115   impl.jpeg_CreateDecompress(LibjpegHandle::kLibjpegVersion,
116                              sizeof(jpeg_decompress_struct));
117   status = impl.status();
118   if (status.code == kTfLiteOk) {
119     decoder->expected_size_for_decompress_struct_ =
120         sizeof(jpeg_decompress_struct);
121     return decoder;
122   }
123   if (!absl::StrContains(status.error_message, kSizeMismatchError)) {
124     return nullptr;
125   }
126   status = ExtractSizeFromErrorMessage(
127       status.error_message, decoder->expected_size_for_decompress_struct_);
128   if (status.code != kTfLiteOk) {
129     return nullptr;
130   }
131   return decoder;
132 }
133 
134 namespace {
135 
JpegHeaderToString(const JpegHeader & header)136 std::string JpegHeaderToString(const JpegHeader& header) {
137   return "(" + std::to_string(header.height) + ", " +
138          std::to_string(header.width) + ", " + std::to_string(header.channels) +
139          ", " + std::to_string(header.bits_per_sample) + ")";
140 }
141 
142 }  // namespace
143 
DecodeImage(const tflite::StringRef & encoded,const JpegHeader & expected_image_dimensions,unsigned char * decoded,const size_t & decoded_size) const144 Status LibjpegDecoder::DecodeImage(const tflite::StringRef& encoded,
145                                    const JpegHeader& expected_image_dimensions,
146                                    unsigned char* decoded,
147                                    const size_t& decoded_size) const {
148   if (expected_image_dimensions.bits_per_sample != 8) {
149     return {kTfLiteError, "Supporting only images with 8 bits per sample"};
150   }
151   if (expected_image_dimensions.channels != 1 &&
152       expected_image_dimensions.channels != 3) {
153     return {kTfLiteError, "Supporting only images with 1 or 3 channels"};
154   }
155   if (expected_image_dimensions.width > kMaxImageWidth ||
156       expected_image_dimensions.height > kMaxImageHeight) {
157     return {kTfLiteError, "Image is too big, dimensions (" +
158                               std::to_string(expected_image_dimensions.width) +
159                               "," +
160                               std::to_string(expected_image_dimensions.width) +
161                               ") larger than the maximum allowed (" +
162                               std::to_string(kMaxImageWidth) + ", " +
163                               std::to_string(kMaxImageHeight) + ")"};
164   }
165   // We match the buffer size and the expected size of the decoded image from
166   // the header to prevent buffer overflows.
167   JpegHeader header;
168   Status read_header_status = ReadJpegHeader(encoded, &header);
169   if (read_header_status.code != kTfLiteOk) {
170     return read_header_status;
171   }
172 
173   if (expected_image_dimensions.channels != header.channels ||
174       expected_image_dimensions.width != header.width ||
175       expected_image_dimensions.height != header.height ||
176       expected_image_dimensions.bits_per_sample != header.bits_per_sample) {
177     return {kTfLiteError, "Decoded image size " + JpegHeaderToString(header) +
178                               " is different from provided image size " +
179                               JpegHeaderToString(expected_image_dimensions)};
180   }
181 
182   size_t header_image_size = static_cast<size_t>(header.width) *
183                              static_cast<size_t>(header.height) *
184                              static_cast<size_t>(header.channels);
185 
186   if (header_image_size != decoded_size) {
187     return {kTfLiteError, "Size of buffer(" + std::to_string(decoded_size) +
188                               ") for storing decoded image must be equal to "
189                               "the size of decoded image(" +
190                               std::to_string(header_image_size) + ")."};
191   }
192 
193   // Dropping constness as fmemopen requires non-const buffers.
194   char* image_buffer = const_cast<char*>(encoded.str);
195   size_t image_size = encoded.len;
196   std::unique_ptr<FILE, std::function<void(FILE*)>> file(
197       libc_handle_.fmemopen(image_buffer, image_size, "r"),
198       [](FILE* f) { fclose(f); });
199   if (file == nullptr) {
200     return {kTfLiteError, "Fmemopen failed."};
201   }
202 
203   Impl impl(expected_size_for_decompress_struct_, libjpeg_handle_.get());
204   if (impl.jpeg_CreateDecompress(LibjpegHandle::kLibjpegVersion,
205                                  expected_size_for_decompress_struct_)) {
206     return impl.status();
207   }
208   if (impl.jpeg_stdio_src(file.get())) {
209     return impl.status();
210   }
211   // jpeg_read_header() must be called before calling jpeg_start_decompress().
212   // It initialises decompression parameters to default values.
213   // jpeg_read_header() should not be relied upon for getting image information
214   // (width, height) etc. Fields populated by jpeg_read_header() such as
215   // `image_width` and `image_height` come after the `jpeg_common_fields` and
216   // these may have been shifted in the struct on some builds of libjpeg.
217   // See go/libjpeg-android.
218   int read_header_result = 0;
219   if (impl.jpeg_read_header(read_header_result, true) != kTfLiteOk) {
220     return impl.status();
221   }
222   if (read_header_result != JPEG_HEADER_OK) {
223     return {kTfLiteError, "Failed call jpeg_read_header"};
224   }
225   boolean start_decompress_result = false;
226   if (impl.jpeg_start_decompress(start_decompress_result) != kTfLiteOk) {
227     return impl.status();
228   }
229   if (!start_decompress_result) {
230     return {kTfLiteError, "Failed call jpeg_start_decompress_"};
231   }
232 
233   size_t height = header.height;
234   size_t row_stride = header.width * header.channels;
235 
236   // Decoding the data in a buffer as large as the largest allowed JPEG
237   // image row to avoid overflows in case  we are reading a wrong value for the
238   // image size in the header or we are receiving an image with an header
239   // deliberately incorrect to cause a buffer overflow.
240   // Using 4 channels because the decode color process can handle 3 or 4
241   // channels:
242   // https://github.com/libjpeg-turbo/libjpeg-turbo/blob/main/jdcolor.c#L383
243   const size_t kMaxImageSize = JPEG_MAX_DIMENSION * 4;
244   // Initializing the buffer in case we are trying to read more data than
245   // actually available to avoid having access to uninitialized memory.
246   // Libjpeg turbo actually fills the unread bytes with zeros
247   // https://github.com/libjpeg-turbo/libjpeg-turbo/blob/main/jdhuff.c#L360
248   // but we don't know what the library on the target system would do.
249   // Output tensor data stored as RGBRGBRGB..., row-wise for all images.
250   std::vector<unsigned char> decode_buffer(kMaxImageSize);
251   // Do not rely on fields such as `output_scanline` from
252   // `jpeg_decompress_struct` as these would have been shifted. See
253   // go/libjpeg-android.
254   unsigned char* buffer_array[1];
255   buffer_array[0] = decode_buffer.data();
256   size_t decoded_offset = 0;
257   while (height--) {
258     // According to the documentation, jpeg_read_scanlines returns the number
259     // of lines read
260     // https://android.googlesource.com/platform/external/jpeg/+/c6859b743e7248b9f401264aac939a5af0d63799/libjpeg.doc#655
261     // In case of premature ending of the image, the implementation of
262     // jpeg_read_scanlines in the version of JPEG Turbo we are using to test
263     // emits a warning ("Corrupt JPEG data: premature end of data segment")
264     // but doesn't fail and consider the line as successfully read.
265     // See test
266     // LibjpegDecoderTest::DoesNotFailDecodingAnImageWithLessDataThanDeclaredInJpegHeader
267     unsigned int num_of_scanlines_read = 0;
268     if (impl.jpeg_read_scanlines(num_of_scanlines_read, buffer_array, 1) !=
269         kTfLiteOk) {
270       return impl.status();
271     }
272 
273     if (num_of_scanlines_read != 1) {
274       return {kTfLiteError, "Expected " + std::to_string(header.height) +
275                                 " lines but found only " +
276                                 std::to_string(header.height - height) +
277                                 " read scanlines is " +
278                                 std::to_string(num_of_scanlines_read)};
279     }
280 
281     std::copy_n(buffer_array[0], row_stride, decoded + decoded_offset);
282 
283     decoded_offset += row_stride;
284   }
285   boolean finish_decompress_result = false;
286   if (impl.jpeg_finish_decompress(finish_decompress_result) != kTfLiteOk) {
287     return impl.status();
288   }
289   if (!finish_decompress_result) {
290     return {kTfLiteError, "Failed call jpeg_finish_decompress_"};
291   }
292   if (impl.jpeg_destroy_decompress() != kTfLiteOk) {
293     return impl.status();
294   }
295   return impl.status();
296 }
297 }  // namespace decode_jpeg_kernel
298 }  // namespace acceleration
299 }  // namespace tflite
300