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