├── .DS_Store ├── Part1-Fuzzing ├── .DS_Store ├── README.md └── 模糊测试.pptx ├── Part2-AFL-pattern ├── .DS_Store ├── AFL详解.pptx └── README.md ├── Part3-ok-file-formats ├── .DS_Store ├── README.md ├── creash │ ├── id^%000000,sig^%11,src^%001329,op^%havoc,rep^%4 │ ├── id^%000001,sig^%11,src^%001377,op^%havoc,rep^%4 │ ├── id^%000002,sig^%11,src^%001377,op^%havoc,rep^%8 │ ├── id^%000003,sig^%06,src^%001377,op^%havoc,rep^%4 │ ├── id^%000004,sig^%11,src^%001377,op^%havoc,rep^%2 │ └── id^%000005,sig^%11,src^%001377,op^%havoc,rep^%2 ├── fuzz_png ├── img │ ├── .DS_Store │ ├── AFL分支.png │ ├── AFL流程.png │ ├── Fuzz流程.png │ ├── README.md │ ├── afl-plot.png │ ├── crashes.png │ ├── cycledone.png │ ├── example.png │ ├── fuzz过程.png │ ├── ok_file_formats.png │ ├── test-program.png │ ├── test-readme.png │ ├── 分类.png │ └── 多进程.png ├── main.c ├── ok-file-formats模糊测试.pptx ├── ok_png.c ├── ok_png.h └── process_of_fuzz.md ├── README.md ├── img ├── AFL分支.png └── AFL流程.png └── 夏令营培训大纲v2-张涛-模糊测试初探.docx /.DS_Store: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WayneDevMaze/Fuzzing-Introduction/528d82a492d88941f9b4d4731569f21b5fc90533/.DS_Store -------------------------------------------------------------------------------- /Part1-Fuzzing/.DS_Store: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WayneDevMaze/Fuzzing-Introduction/528d82a492d88941f9b4d4731569f21b5fc90533/Part1-Fuzzing/.DS_Store -------------------------------------------------------------------------------- /Part1-Fuzzing/README.md: -------------------------------------------------------------------------------- 1 | # 模糊测试介绍 2 | 什么是模糊测试技术呢，就是用大量的随机输入去测试软件的健壮性等，最早是软件测试的概念。后来出现了afl，便出现了以覆盖率为导向的模糊测试技术。 3 | 4 | ## AFL 5 | ![img](../img/AFL流程.png) 6 | 看图可以很好的理解，最开始给afl初始种子，然后经过变异得到输入，将输入给到待测试程序，并通过插桩技术对程序跟踪，或者覆盖率，个人认为现在这个覆盖率已经变成了一个大的概念，可以是块覆盖、可以是边覆盖、可以是很多其他有利于程序输入的信息，只要是有利于变异和种子筛选的都可以算，然后覆盖率指导之前的流程，周而复始，直到遇到崩溃情况，便将其保存。 7 | 这种覆盖率的想法使得输入不再是随机盲目的，可以说afl的出现使得模糊测试进入了新的阶段，虽然afl在2.52b版本之后就停止更新了，但是Google目前还一直维护着，而且后来出现了很多分支。今天用到的模糊测试工具就是AFL，其实其他工具使用思路也大同小异，都可以借鉴。 8 | 9 | ## 挑选合适的Fuzzer 10 | 这一步我单独写一个章节是因为这里很重要，虽然说我们一直讨论的是用AFL相关，但是实际上在此基础上有很多很多的分支，我才疏学浅，这里就列一小部分（有兴趣可以找这三四年顶会的综述论文看一下，一定会有收获的） 11 | ![AFL分支](../img/AFL分支.png) 12 | 关于AFL的分支，我这里分为四类（并不一定是afl转来了，但一定是受其思想影响的进步）： 13 | #### 移植 14 | 1.afl多是用在Linux的c语言类程序文件模糊测试，所以就有人做了平台移植：最经典的就是winafl，在Windows平台下的模糊测试工具； 15 | 2.还有针对语言的移植，比如针对rust语言的/afl.rs、针对python语言的python-afl、针对Java语言的java-afl、针对C#语言的sharpfuzz； 16 | #### 性能的改进 17 | 1.这部分是指速度性能的改进，既有针对有源码的libfuzzer，通过用llvm、clang框架大大提升了性能，这部分Google自家的fuzzer做的算是很好了，还有honggfuzz，都是对afl的很大改进； 18 | 2.也有针对无源码的，像Qemu、PIN、Dyninst这些工具的引入，尽管说afl自带的系统里有qemu等但是效率不理想，所以就有一些别的大牛单独做的新的分支，甚至还有跟Frida结合的应用； 19 | #### 提升效率 20 | 这部分每年都会看见论文的新花样，针对变异策略的、针对种子筛选的层出不穷，afl-smart、ijon等等，不管是加入符号执行还是对策略改进，实际上对afl给改动并不大。主要还是从理论上进行了设想，提出了新思路； 21 | #### 重构 22 | 1.基本上跟原来对afl差别很大了，一种是对整个变异策略全部重建，比如Angora就是用rust语言重新写对，并且用了自己新对变异策略； 23 | 2.另一种是加入新元素，比如Neuzz这个就是加入了机器学习，再比如今年出现了一些对蓝牙的、对use设备对模糊测试； 24 | 总结下来就是一二类属于实际应用类，比较注重运用，怎么用的舒服，用得快怎么来，三四类属于理论类，比较注意想法。这篇文章的fuzz目标是文件类型的，而且是个小库，就用afl就够了，算是抛砖引玉了。 25 | 26 | ## 使用 27 | AFL使用流程：https://www.cnblogs.com/wayne-tao/p/11739420.html 28 | 29 | -------------------------------------------------------------------------------- /Part1-Fuzzing/模糊测试.pptx: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WayneDevMaze/Fuzzing-Introduction/528d82a492d88941f9b4d4731569f21b5fc90533/Part1-Fuzzing/模糊测试.pptx -------------------------------------------------------------------------------- /Part2-AFL-pattern/.DS_Store: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WayneDevMaze/Fuzzing-Introduction/528d82a492d88941f9b4d4731569f21b5fc90533/Part2-AFL-pattern/.DS_Store -------------------------------------------------------------------------------- /Part2-AFL-pattern/AFL详解.pptx: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WayneDevMaze/Fuzzing-Introduction/528d82a492d88941f9b4d4731569f21b5fc90533/Part2-AFL-pattern/AFL详解.pptx -------------------------------------------------------------------------------- /Part2-AFL-pattern/README.md: -------------------------------------------------------------------------------- 1 | ## 2.1 AFL代码模块 2 | AFL头文件介绍： 3 | https://www.cnblogs.com/wayne-tao/p/12091747.html 4 | 文件变异策略： 5 | https://www.cnblogs.com/wayne-tao/p/12019499.html 6 | ## 2.1 AFL周边工具 7 | afl-cmin： 8 | https://www.cnblogs.com/wayne-tao/p/11889718.html 9 | afl-tmin： 10 | https://www.cnblogs.com/wayne-tao/p/11889718.html 11 | afl-whatsup： 12 | https://www.cnblogs.com/wayne-tao/p/13569939.html 13 | afl-plot： 14 | https://www.cnblogs.com/wayne-tao/p/13569953.html 15 | -------------------------------------------------------------------------------- /Part3-ok-file-formats/.DS_Store: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WayneDevMaze/Fuzzing-Introduction/528d82a492d88941f9b4d4731569f21b5fc90533/Part3-ok-file-formats/.DS_Store -------------------------------------------------------------------------------- /Part3-ok-file-formats/README.md: -------------------------------------------------------------------------------- 1 | ## 参照此大纲 2 | https://bbs.pediy.com/thread-261253.htm 3 | -------------------------------------------------------------------------------- /Part3-ok-file-formats/creash/id^%000000,sig^%11,src^%001329,op^%havoc,rep^%4: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WayneDevMaze/Fuzzing-Introduction/528d82a492d88941f9b4d4731569f21b5fc90533/Part3-ok-file-formats/creash/id^%000000,sig^%11,src^%001329,op^%havoc,rep^%4 -------------------------------------------------------------------------------- /Part3-ok-file-formats/creash/id^%000001,sig^%11,src^%001377,op^%havoc,rep^%4: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WayneDevMaze/Fuzzing-Introduction/528d82a492d88941f9b4d4731569f21b5fc90533/Part3-ok-file-formats/creash/id^%000001,sig^%11,src^%001377,op^%havoc,rep^%4 -------------------------------------------------------------------------------- /Part3-ok-file-formats/creash/id^%000002,sig^%11,src^%001377,op^%havoc,rep^%8: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WayneDevMaze/Fuzzing-Introduction/528d82a492d88941f9b4d4731569f21b5fc90533/Part3-ok-file-formats/creash/id^%000002,sig^%11,src^%001377,op^%havoc,rep^%8 -------------------------------------------------------------------------------- /Part3-ok-file-formats/creash/id^%000003,sig^%06,src^%001377,op^%havoc,rep^%4: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WayneDevMaze/Fuzzing-Introduction/528d82a492d88941f9b4d4731569f21b5fc90533/Part3-ok-file-formats/creash/id^%000003,sig^%06,src^%001377,op^%havoc,rep^%4 -------------------------------------------------------------------------------- /Part3-ok-file-formats/creash/id^%000004,sig^%11,src^%001377,op^%havoc,rep^%2: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WayneDevMaze/Fuzzing-Introduction/528d82a492d88941f9b4d4731569f21b5fc90533/Part3-ok-file-formats/creash/id^%000004,sig^%11,src^%001377,op^%havoc,rep^%2 -------------------------------------------------------------------------------- /Part3-ok-file-formats/creash/id^%000005,sig^%11,src^%001377,op^%havoc,rep^%2: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WayneDevMaze/Fuzzing-Introduction/528d82a492d88941f9b4d4731569f21b5fc90533/Part3-ok-file-formats/creash/id^%000005,sig^%11,src^%001377,op^%havoc,rep^%2 -------------------------------------------------------------------------------- /Part3-ok-file-formats/fuzz_png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WayneDevMaze/Fuzzing-Introduction/528d82a492d88941f9b4d4731569f21b5fc90533/Part3-ok-file-formats/fuzz_png -------------------------------------------------------------------------------- /Part3-ok-file-formats/img/.DS_Store: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WayneDevMaze/Fuzzing-Introduction/528d82a492d88941f9b4d4731569f21b5fc90533/Part3-ok-file-formats/img/.DS_Store -------------------------------------------------------------------------------- /Part3-ok-file-formats/img/AFL分支.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WayneDevMaze/Fuzzing-Introduction/528d82a492d88941f9b4d4731569f21b5fc90533/Part3-ok-file-formats/img/AFL分支.png -------------------------------------------------------------------------------- /Part3-ok-file-formats/img/AFL流程.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WayneDevMaze/Fuzzing-Introduction/528d82a492d88941f9b4d4731569f21b5fc90533/Part3-ok-file-formats/img/AFL流程.png -------------------------------------------------------------------------------- /Part3-ok-file-formats/img/Fuzz流程.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WayneDevMaze/Fuzzing-Introduction/528d82a492d88941f9b4d4731569f21b5fc90533/Part3-ok-file-formats/img/Fuzz流程.png -------------------------------------------------------------------------------- /Part3-ok-file-formats/img/README.md: -------------------------------------------------------------------------------- 1 | # IMG 2 | 给md文件用的图片库 -------------------------------------------------------------------------------- /Part3-ok-file-formats/img/afl-plot.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WayneDevMaze/Fuzzing-Introduction/528d82a492d88941f9b4d4731569f21b5fc90533/Part3-ok-file-formats/img/afl-plot.png -------------------------------------------------------------------------------- /Part3-ok-file-formats/img/crashes.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WayneDevMaze/Fuzzing-Introduction/528d82a492d88941f9b4d4731569f21b5fc90533/Part3-ok-file-formats/img/crashes.png -------------------------------------------------------------------------------- /Part3-ok-file-formats/img/cycledone.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WayneDevMaze/Fuzzing-Introduction/528d82a492d88941f9b4d4731569f21b5fc90533/Part3-ok-file-formats/img/cycledone.png -------------------------------------------------------------------------------- /Part3-ok-file-formats/img/example.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WayneDevMaze/Fuzzing-Introduction/528d82a492d88941f9b4d4731569f21b5fc90533/Part3-ok-file-formats/img/example.png -------------------------------------------------------------------------------- /Part3-ok-file-formats/img/fuzz过程.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WayneDevMaze/Fuzzing-Introduction/528d82a492d88941f9b4d4731569f21b5fc90533/Part3-ok-file-formats/img/fuzz过程.png -------------------------------------------------------------------------------- /Part3-ok-file-formats/img/ok_file_formats.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WayneDevMaze/Fuzzing-Introduction/528d82a492d88941f9b4d4731569f21b5fc90533/Part3-ok-file-formats/img/ok_file_formats.png -------------------------------------------------------------------------------- /Part3-ok-file-formats/img/test-program.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WayneDevMaze/Fuzzing-Introduction/528d82a492d88941f9b4d4731569f21b5fc90533/Part3-ok-file-formats/img/test-program.png -------------------------------------------------------------------------------- /Part3-ok-file-formats/img/test-readme.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WayneDevMaze/Fuzzing-Introduction/528d82a492d88941f9b4d4731569f21b5fc90533/Part3-ok-file-formats/img/test-readme.png -------------------------------------------------------------------------------- /Part3-ok-file-formats/img/分类.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WayneDevMaze/Fuzzing-Introduction/528d82a492d88941f9b4d4731569f21b5fc90533/Part3-ok-file-formats/img/分类.png -------------------------------------------------------------------------------- /Part3-ok-file-formats/img/多进程.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WayneDevMaze/Fuzzing-Introduction/528d82a492d88941f9b4d4731569f21b5fc90533/Part3-ok-file-formats/img/多进程.png -------------------------------------------------------------------------------- /Part3-ok-file-formats/main.c: -------------------------------------------------------------------------------- 1 | #include 2 | #include "ok_png.h" 3 | 4 | int main(int _argc, char **_argv) { 5 | FILE *file = fopen(_argv[1], "rb"); 6 | ok_png image = ok_png_read(file, OK_PNG_COLOR_FORMAT_RGBA | OK_PNG_PREMULTIPLIED_ALPHA | OK_PNG_FLIP_Y); 7 | fclose(file); 8 | if (image.data) { 9 | printf("Got image! Size: %li x %li\n", (long)image.width, (long)image.height); 10 | free(image.data); 11 | } 12 | return 0; 13 | } 14 | -------------------------------------------------------------------------------- /Part3-ok-file-formats/ok-file-formats模糊测试.pptx: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WayneDevMaze/Fuzzing-Introduction/528d82a492d88941f9b4d4731569f21b5fc90533/Part3-ok-file-formats/ok-file-formats模糊测试.pptx -------------------------------------------------------------------------------- /Part3-ok-file-formats/ok_png.c: -------------------------------------------------------------------------------- 1 | /* 2 | ok-file-formats 3 | https://github.com/brackeen/ok-file-formats 4 | Copyright (c) 2014-2020 David Brackeen 5 | 6 | Permission is hereby granted, free of charge, to any person obtaining a copy of this software and 7 | associated documentation files (the "Software"), to deal in the Software without restriction, 8 | including without limitation the rights to use, copy, modify, merge, publish, distribute, 9 | sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all copies or 13 | substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT 16 | NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND 17 | NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, 18 | DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT 19 | OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 20 | */ 21 | 22 | #include "ok_png.h" 23 | #include 24 | #include 25 | 26 | #if __STDC_VERSION__ >= 199901L 27 | #define RESTRICT restrict 28 | #else 29 | #define RESTRICT 30 | #endif 31 | 32 | #ifndef min 33 | #define min(a, b) ((a) < (b) ? (a) : (b)) 34 | #endif 35 | 36 | #define OK_SIZE_MAX (~(size_t)0) 37 | 38 | #define PNG_TYPE(a, b, c, d) ((a << 24) | (b << 16) | (c << 8) | d) 39 | 40 | static const uint32_t OK_PNG_CHUNK_IHDR = PNG_TYPE('I', 'H', 'D', 'R'); 41 | static const uint32_t OK_PNG_CHUNK_PLTE = PNG_TYPE('P', 'L', 'T', 'E'); 42 | static const uint32_t OK_PNG_CHUNK_TRNS = PNG_TYPE('t', 'R', 'N', 'S'); 43 | static const uint32_t OK_PNG_CHUNK_IDAT = PNG_TYPE('I', 'D', 'A', 'T'); 44 | static const uint32_t OK_PNG_CHUNK_IEND = PNG_TYPE('I', 'E', 'N', 'D'); 45 | static const uint32_t OK_PNG_CHUNK_CGBI = PNG_TYPE('C', 'g', 'B', 'I'); 46 | 47 | static const uint8_t OK_PNG_COLOR_TYPE_GRAYSCALE = 0; 48 | static const uint8_t OK_PNG_COLOR_TYPE_RGB = 2; 49 | static const uint8_t OK_PNG_COLOR_TYPE_PALETTE = 3; 50 | static const uint8_t OK_PNG_COLOR_TYPE_GRAYSCALE_WITH_ALPHA = 4; 51 | static const uint8_t OK_PNG_COLOR_TYPE_RGB_WITH_ALPHA = 6; 52 | static const uint8_t OK_PNG_SAMPLES_PER_PIXEL[] = {1, 0, 3, 1, 2, 0, 4}; 53 | 54 | #ifndef OK_NO_DEFAULT_ALLOCATOR 55 | 56 | static void *ok_stdlib_alloc(void *user_data, size_t size) { 57 | (void)user_data; 58 | return malloc(size); 59 | } 60 | 61 | static void ok_stdlib_free(void *user_data, void *memory) { 62 | (void)user_data; 63 | free(memory); 64 | } 65 | 66 | const ok_png_allocator OK_PNG_DEFAULT_ALLOCATOR = { 67 | .alloc = ok_stdlib_alloc, 68 | .free = ok_stdlib_free, 69 | .image_alloc = NULL 70 | }; 71 | 72 | #endif 73 | 74 | typedef enum { 75 | OK_PNG_FILTER_NONE = 0, 76 | OK_PNG_FILTER_SUB, 77 | OK_PNG_FILTER_UP, 78 | OK_PNG_FILTER_AVG, 79 | OK_PNG_FILTER_PAETH, 80 | OK_PNG_NUM_FILTERS 81 | } ok_png_filter_type; 82 | 83 | typedef struct { 84 | // Image 85 | ok_png *png; 86 | 87 | // Allocator 88 | ok_png_allocator allocator; 89 | void *allocator_user_data; 90 | 91 | // Input 92 | ok_png_input input; 93 | void *input_user_data; 94 | 95 | // Decode options 96 | ok_png_decode_flags decode_flags; 97 | 98 | // Decoding 99 | ok_inflater *inflater; 100 | size_t inflater_bytes_read; 101 | uint8_t *inflate_buffer; 102 | uint8_t *curr_scanline; 103 | uint8_t *prev_scanline; 104 | uint32_t scanline; 105 | uint8_t interlace_pass; // 0 for uninitialized, 1 for non-interlaced, 1..7 for interlaced 106 | bool ready_for_next_interlace_pass; 107 | uint8_t *temp_data_row; 108 | bool decoding_completed; 109 | 110 | // PNG data 111 | uint8_t bit_depth; 112 | uint8_t color_type; 113 | uint8_t interlace_method; 114 | uint8_t palette[256 * 4]; 115 | uint32_t palette_length; 116 | uint16_t single_transparent_color_key[3]; 117 | bool has_single_transparent_color; 118 | bool is_ios_format; 119 | 120 | } ok_png_decoder; 121 | 122 | #define ok_alloc(decoder, size) (decoder)->allocator.alloc((decoder)->allocator_user_data, (size)) 123 | #define ok_png_error(png, error_code, message) ok_png_set_error((png), (error_code)) 124 | 125 | static void ok_png_set_error(ok_png *png, ok_png_error error_code) { 126 | if (png) { 127 | png->width = 0; 128 | png->height = 0; 129 | png->error_code = error_code; 130 | } 131 | } 132 | 133 | static bool ok_read(ok_png_decoder *decoder, uint8_t *buffer, size_t length) { 134 | if (decoder->input.read(decoder->input_user_data, buffer, length) == length) { 135 | return true; 136 | } else { 137 | ok_png_error(decoder->png, OK_PNG_ERROR_IO, "Read error: error calling input function."); 138 | return false; 139 | } 140 | } 141 | 142 | static bool ok_seek(ok_png_decoder *decoder, long length) { 143 | if (decoder->input.seek(decoder->input_user_data, length)) { 144 | return true; 145 | } else { 146 | ok_png_error(decoder->png, OK_PNG_ERROR_IO, "Seek error: error calling input function."); 147 | return false; 148 | } 149 | } 150 | 151 | #ifndef OK_NO_STDIO 152 | 153 | static size_t ok_file_read(void *user_data, uint8_t *buffer, size_t length) { 154 | return fread(buffer, 1, length, (FILE *)user_data); 155 | } 156 | 157 | static bool ok_file_seek(void *user_data, long count) { 158 | return fseek((FILE *)user_data, count, SEEK_CUR) == 0; 159 | } 160 | 161 | static const ok_png_input OK_PNG_FILE_INPUT = { 162 | .read = ok_file_read, 163 | .seek = ok_file_seek, 164 | }; 165 | 166 | #endif 167 | 168 | static void ok_png_decode(ok_png *png, ok_png_decode_flags decode_flags, 169 | ok_png_input input, void *input_user_data, 170 | ok_png_allocator allocator, void *allocator_user_data); 171 | 172 | // Public API 173 | 174 | #if !defined(OK_NO_STDIO) && !defined(OK_NO_DEFAULT_ALLOCATOR) 175 | 176 | ok_png ok_png_read(FILE *file, ok_png_decode_flags decode_flags) { 177 | return ok_png_read_with_allocator(file, decode_flags, OK_PNG_DEFAULT_ALLOCATOR, NULL); 178 | } 179 | 180 | #endif 181 | 182 | #if !defined(OK_NO_STDIO) 183 | 184 | ok_png ok_png_read_with_allocator(FILE *file, ok_png_decode_flags decode_flags, 185 | ok_png_allocator allocator, void *allocator_user_data) { 186 | ok_png png = { 0 }; 187 | if (file) { 188 | ok_png_decode(&png, decode_flags, OK_PNG_FILE_INPUT, file, allocator, allocator_user_data); 189 | } else { 190 | ok_png_error(&png, OK_PNG_ERROR_API, "File not found"); 191 | } 192 | return png; 193 | } 194 | 195 | #endif 196 | 197 | ok_png ok_png_read_from_input(ok_png_decode_flags decode_flags, 198 | ok_png_input input_callbacks, void *input_callbacks_user_data, 199 | ok_png_allocator allocator, void *allocator_user_data) { 200 | ok_png png = { 0 }; 201 | ok_png_decode(&png, decode_flags, input_callbacks, input_callbacks_user_data, 202 | allocator, allocator_user_data); 203 | return png; 204 | } 205 | 206 | // Main read functions 207 | 208 | static inline uint16_t readBE16(const uint8_t *data) { 209 | return (uint16_t)((data[0] << 8) | data[1]); 210 | } 211 | 212 | static inline uint32_t readBE32(const uint8_t *data) { 213 | return (((uint32_t)data[0] << 24) | 214 | ((uint32_t)data[1] << 16) | 215 | ((uint32_t)data[2] << 8) | 216 | ((uint32_t)data[3] << 0)); 217 | } 218 | 219 | static inline void ok_png_premultiply(uint8_t *dst) { 220 | const uint8_t a = dst[3]; 221 | if (a == 0) { 222 | dst[0] = 0; 223 | dst[1] = 0; 224 | dst[2] = 0; 225 | } else if (a < 255) { 226 | dst[0] = (a * dst[0] + 127) / 255; 227 | dst[1] = (a * dst[1] + 127) / 255; 228 | dst[2] = (a * dst[2] + 127) / 255; 229 | } 230 | } 231 | 232 | static inline void ok_png_unpremultiply(uint8_t *dst) { 233 | const uint8_t a = dst[3]; 234 | if (a > 0 && a < 255) { 235 | dst[0] = 255 * dst[0] / a; 236 | dst[1] = 255 * dst[1] / a; 237 | dst[2] = 255 * dst[2] / a; 238 | } 239 | } 240 | 241 | static bool ok_png_read_header(ok_png_decoder *decoder, uint32_t chunk_length) { 242 | ok_png *png = decoder->png; 243 | if (chunk_length != 13) { 244 | ok_png_error(png, OK_PNG_ERROR_INVALID, "Invalid IHDR chunk length"); 245 | return false; 246 | } 247 | uint8_t chunk_data[13]; 248 | if (!ok_read(decoder, chunk_data, sizeof(chunk_data))) { 249 | return false; 250 | } 251 | png->width = readBE32(chunk_data); 252 | png->height = readBE32(chunk_data + 4); 253 | png->bpp = 4; // Always decoding to 32-bit color 254 | decoder->bit_depth = chunk_data[8]; 255 | decoder->color_type = chunk_data[9]; 256 | uint8_t compression_method = chunk_data[10]; 257 | uint8_t filter_method = chunk_data[11]; 258 | decoder->interlace_method = chunk_data[12]; 259 | uint64_t stride = (uint64_t)png->width * png->bpp; 260 | 261 | if (compression_method != 0) { 262 | ok_png_error(png, OK_PNG_ERROR_INVALID, "Invalid compression method"); 263 | return false; 264 | } else if (filter_method != 0) { 265 | ok_png_error(png, OK_PNG_ERROR_INVALID, "Invalid filter method"); 266 | return false; 267 | } else if (decoder->interlace_method != 0 && decoder->interlace_method != 1) { 268 | ok_png_error(png, OK_PNG_ERROR_INVALID, "Invalid interlace method"); 269 | return false; 270 | } else if (stride > UINT32_MAX) { 271 | ok_png_error(png, OK_PNG_ERROR_UNSUPPORTED, "Width too large"); 272 | return false; 273 | } 274 | 275 | const int c = decoder->color_type; 276 | const int b = decoder->bit_depth; 277 | const bool valid = 278 | (c == OK_PNG_COLOR_TYPE_GRAYSCALE && (b == 1 || b == 2 || b == 4 || b == 8 || b == 16)) || 279 | (c == OK_PNG_COLOR_TYPE_RGB && (b == 8 || b == 16)) || 280 | (c == OK_PNG_COLOR_TYPE_PALETTE && (b == 1 || b == 2 || b == 4 || b == 8)) || 281 | (c == OK_PNG_COLOR_TYPE_GRAYSCALE_WITH_ALPHA && (b == 8 || b == 16)) || 282 | (c == OK_PNG_COLOR_TYPE_RGB_WITH_ALPHA && (b == 8 || b == 16)); 283 | 284 | if (!valid) { 285 | ok_png_error(png, OK_PNG_ERROR_INVALID, "Invalid combination of color type and bit depth"); 286 | return false; 287 | } 288 | 289 | png->stride = (uint32_t)stride; 290 | png->has_alpha = (c == OK_PNG_COLOR_TYPE_GRAYSCALE_WITH_ALPHA || 291 | c == OK_PNG_COLOR_TYPE_RGB_WITH_ALPHA); 292 | decoder->interlace_pass = 0; 293 | decoder->ready_for_next_interlace_pass = true; 294 | return true; 295 | } 296 | 297 | static bool ok_png_read_palette(ok_png_decoder *decoder, uint32_t chunk_length) { 298 | ok_png *png = decoder->png; 299 | decoder->palette_length = chunk_length / 3; 300 | 301 | if (decoder->palette_length > 256 || decoder->palette_length * 3 != chunk_length) { 302 | ok_png_error(png, OK_PNG_ERROR_INVALID, "Invalid palette chunk length"); 303 | return false; 304 | } 305 | const bool src_is_bgr = decoder->is_ios_format; 306 | const bool dst_is_bgr = (decoder->decode_flags & OK_PNG_COLOR_FORMAT_BGRA) != 0; 307 | const bool should_byteswap = src_is_bgr != dst_is_bgr; 308 | uint8_t *dst = decoder->palette; 309 | uint8_t buffer[256 * 3]; 310 | if (!ok_read(decoder, buffer, 3 * decoder->palette_length)) { 311 | return false; 312 | } 313 | uint8_t *in = buffer; 314 | if (should_byteswap) { 315 | for (uint32_t i = 0; i < decoder->palette_length; i++, in += 3, dst += 4) { 316 | dst[0] = in[2]; 317 | dst[1] = in[1]; 318 | dst[2] = in[0]; 319 | dst[3] = 0xff; 320 | } 321 | } else { 322 | for (uint32_t i = 0; i < decoder->palette_length; i++, in += 3, dst += 4) { 323 | dst[0] = in[0]; 324 | dst[1] = in[1]; 325 | dst[2] = in[2]; 326 | dst[3] = 0xff; 327 | } 328 | } 329 | return true; 330 | } 331 | 332 | static bool ok_png_read_transparency(ok_png_decoder *decoder, uint32_t chunk_length) { 333 | ok_png *png = decoder->png; 334 | png->has_alpha = true; 335 | 336 | if (decoder->color_type == OK_PNG_COLOR_TYPE_PALETTE) { 337 | if (chunk_length > decoder->palette_length || chunk_length > 256) { 338 | ok_png_error(png, OK_PNG_ERROR_INVALID, 339 | "Invalid transparency length for palette color type"); 340 | return false; 341 | } 342 | 343 | const bool should_premultiply = (decoder->decode_flags & OK_PNG_PREMULTIPLIED_ALPHA) != 0; 344 | uint8_t *dst = decoder->palette; 345 | uint8_t buffer[256]; 346 | if (!ok_read(decoder, buffer, chunk_length)) { 347 | return false; 348 | } 349 | for (uint32_t i = 0; i < chunk_length; i++) { 350 | dst[3] = buffer[i]; 351 | if (should_premultiply) { 352 | ok_png_premultiply(dst); 353 | } 354 | dst += 4; 355 | } 356 | return true; 357 | } else if (decoder->color_type == OK_PNG_COLOR_TYPE_GRAYSCALE) { 358 | if (chunk_length != 2) { 359 | ok_png_error(png, OK_PNG_ERROR_INVALID, 360 | "Invalid transparency length for grayscale color type"); 361 | return false; 362 | } else { 363 | uint8_t buffer[2]; 364 | if (!ok_read(decoder, buffer, sizeof(buffer))) { 365 | return false; 366 | } 367 | const uint16_t v = readBE16(buffer); 368 | decoder->single_transparent_color_key[0] = v; 369 | decoder->single_transparent_color_key[1] = v; 370 | decoder->single_transparent_color_key[2] = v; 371 | decoder->has_single_transparent_color = true; 372 | return true; 373 | } 374 | } else if (decoder->color_type == OK_PNG_COLOR_TYPE_RGB) { 375 | if (chunk_length != 6) { 376 | ok_png_error(png, OK_PNG_ERROR_INVALID, 377 | "Invalid transparency length for truecolor color type"); 378 | return false; 379 | } else { 380 | uint8_t buffer[6]; 381 | if (!ok_read(decoder, buffer, sizeof(buffer))) { 382 | return false; 383 | } 384 | decoder->single_transparent_color_key[0] = readBE16(buffer + 0); 385 | decoder->single_transparent_color_key[1] = readBE16(buffer + 2); 386 | decoder->single_transparent_color_key[2] = readBE16(buffer + 4); 387 | decoder->has_single_transparent_color = true; 388 | return true; 389 | } 390 | } else { 391 | ok_png_error(png, OK_PNG_ERROR_INVALID, 392 | "Invalid transparency for color type"); 393 | return false; 394 | } 395 | } 396 | 397 | static inline int ok_png_paeth_predictor(int a, int b, int c) { 398 | int p = a + b - c; 399 | int pa = abs(p - a); 400 | int pb = abs(p - b); 401 | int pc = abs(p - c); 402 | if (pa <= pb && pa <= pc) { 403 | return a; 404 | } else if (pb <= pc) { 405 | return b; 406 | } else { 407 | return c; 408 | } 409 | } 410 | 411 | static void ok_png_decode_filter(uint8_t * RESTRICT curr, const uint8_t * RESTRICT prev, 412 | size_t length, int filter, uint8_t bpp) { 413 | switch (filter) { 414 | case OK_PNG_FILTER_NONE: 415 | // Do nothing 416 | break; 417 | case OK_PNG_FILTER_SUB: { 418 | // Input = Sub 419 | // Raw(x) = Sub(x) + Raw(x-bpp) 420 | // For all x < 0, assume Raw(x) = 0. 421 | for (size_t i = bpp; i < length; i++) { 422 | curr[i] = curr[i] + curr[i - bpp]; 423 | } 424 | break; 425 | } 426 | case OK_PNG_FILTER_UP: { 427 | // Input = Up 428 | // Raw(x) = Up(x) + Prior(x) 429 | for (size_t i = 0; i < length; i++) { 430 | curr[i] = curr[i] + prev[i]; 431 | } 432 | break; 433 | } 434 | case OK_PNG_FILTER_AVG: { 435 | // Input = Average 436 | // Raw(x) = Average(x) + floor((Raw(x-bpp)+Prior(x))/2) 437 | for (size_t i = 0; i < bpp; i++) { 438 | curr[i] = curr[i] + (prev[i] >> 1); 439 | } 440 | for (size_t j = bpp; j < length; j++) { 441 | curr[j] = curr[j] + ((curr[j - bpp] + prev[j]) >> 1); 442 | } 443 | break; 444 | } 445 | case OK_PNG_FILTER_PAETH: { 446 | // Input = Paeth 447 | // Raw(x) = Paeth(x) + PaethPredictor(Raw(x-bpp), Prior(x), Prior(x-bpp)) 448 | for (size_t i = 0; i < bpp; i++) { 449 | curr[i] += prev[i]; 450 | } 451 | for (size_t j = bpp; j < length; j++) { 452 | curr[j] += ok_png_paeth_predictor(curr[j - bpp], prev[j], prev[j - bpp]); 453 | } 454 | break; 455 | } 456 | } 457 | } 458 | 459 | static void ok_png_transform_scanline(ok_png_decoder *decoder, const uint8_t *src, uint32_t width) { 460 | ok_png *png = decoder->png; 461 | const bool dst_flip_y = (decoder->decode_flags & OK_PNG_FLIP_Y) != 0; 462 | uint8_t *dst_start; 463 | uint8_t *dst_end; 464 | if (decoder->interlace_method == 0) { 465 | const uint32_t dst_y = 466 | (dst_flip_y ? (png->height - decoder->scanline - 1) : decoder->scanline); 467 | dst_start = png->data + (dst_y * png->stride); 468 | } else if (decoder->interlace_pass == 7) { 469 | const uint32_t t_scanline = decoder->scanline * 2 + 1; 470 | const uint32_t dst_y = dst_flip_y ? (png->height - t_scanline - 1) : t_scanline; 471 | dst_start = png->data + (dst_y * png->stride); 472 | } else { 473 | dst_start = decoder->temp_data_row; 474 | } 475 | dst_end = dst_start + width * png->bpp; 476 | 477 | const int c = decoder->color_type; 478 | const int d = decoder->bit_depth; 479 | const bool t = decoder->has_single_transparent_color; 480 | const bool has_full_alpha = (c == OK_PNG_COLOR_TYPE_GRAYSCALE_WITH_ALPHA || 481 | c == OK_PNG_COLOR_TYPE_RGB_WITH_ALPHA); 482 | const bool src_is_premultiplied = decoder->is_ios_format; 483 | const bool dst_is_premultiplied = (decoder->decode_flags & OK_PNG_PREMULTIPLIED_ALPHA) != 0; 484 | const bool src_is_bgr = decoder->is_ios_format; 485 | const bool dst_is_bgr = (decoder->decode_flags & OK_PNG_COLOR_FORMAT_BGRA) != 0; 486 | bool should_byteswap = ((c == OK_PNG_COLOR_TYPE_RGB || c == OK_PNG_COLOR_TYPE_RGB_WITH_ALPHA) && 487 | src_is_bgr != dst_is_bgr); 488 | 489 | // Simple transforms 490 | if (c == OK_PNG_COLOR_TYPE_GRAYSCALE && d == 8 && !t) { 491 | uint8_t *dst = dst_start; 492 | while (dst < dst_end) { 493 | uint8_t v = *src++; 494 | *dst++ = v; 495 | *dst++ = v; 496 | *dst++ = v; 497 | *dst++ = 0xff; 498 | } 499 | } else if (c == OK_PNG_COLOR_TYPE_PALETTE && d == 8) { 500 | uint8_t *dst = dst_start; 501 | const uint8_t *palette = decoder->palette; 502 | while (dst < dst_end) { 503 | memcpy(dst, palette + *src * 4, 4); 504 | dst += 4; 505 | src++; 506 | } 507 | } else if (c == OK_PNG_COLOR_TYPE_RGB && d == 8 && !t) { 508 | if (should_byteswap) { 509 | uint8_t *dst = dst_start; 510 | while (dst < dst_end) { 511 | dst[0] = src[2]; 512 | dst[1] = src[1]; 513 | dst[2] = src[0]; 514 | dst[3] = 0xff; 515 | src += 3; 516 | dst += 4; 517 | } 518 | should_byteswap = false; 519 | } else { 520 | uint8_t *dst = dst_start; 521 | while (dst < dst_end) { 522 | memcpy(dst, src, 3); 523 | dst[3] = 0xff; 524 | src += 3; 525 | dst += 4; 526 | } 527 | } 528 | } else if (c == OK_PNG_COLOR_TYPE_GRAYSCALE_WITH_ALPHA && d == 8) { 529 | uint8_t *dst = dst_start; 530 | while (dst < dst_end) { 531 | uint8_t v = *src++; 532 | uint8_t a = *src++; 533 | *dst++ = v; 534 | *dst++ = v; 535 | *dst++ = v; 536 | *dst++ = a; 537 | } 538 | } else if (c == OK_PNG_COLOR_TYPE_RGB_WITH_ALPHA && d == 8) { 539 | memcpy(dst_start, src, width * 4); 540 | } else { 541 | // Complex transforms: 1-, 2-, 4- and 16-bit, and 8-bit with single-color transparency. 542 | const uint8_t *palette = decoder->palette; 543 | const int bitmask = (1 << d) - 1; 544 | int bit = 8 - d; 545 | uint16_t tr = decoder->single_transparent_color_key[0]; 546 | uint16_t tg = decoder->single_transparent_color_key[1]; 547 | uint16_t tb = decoder->single_transparent_color_key[2]; 548 | if (d <= 8) { 549 | tr = (uint16_t)((tr & bitmask) * (255 / bitmask)); 550 | tg = (uint16_t)((tg & bitmask) * (255 / bitmask)); 551 | tb = (uint16_t)((tb & bitmask) * (255 / bitmask)); 552 | } 553 | uint8_t *dst = dst_start; 554 | while (dst < dst_end) { 555 | uint16_t r = 0; 556 | uint16_t g = 0; 557 | uint16_t b = 0; 558 | uint16_t a = 0xffff; 559 | 560 | if (d < 8) { 561 | if (bit < 0) { 562 | bit = 8 - d; 563 | src++; 564 | } 565 | int v = (*src >> bit) & bitmask; 566 | if (c == OK_PNG_COLOR_TYPE_GRAYSCALE) { 567 | r = g = b = (uint16_t)(v * (255 / bitmask)); 568 | } else { 569 | const uint8_t *psrc = palette + (v * 4); 570 | r = *psrc++; 571 | g = *psrc++; 572 | b = *psrc++; 573 | a = *psrc++; 574 | } 575 | bit -= d; 576 | } else if (d == 8) { 577 | if (c == OK_PNG_COLOR_TYPE_GRAYSCALE) { 578 | r = g = b = *src++; 579 | } else if (c == OK_PNG_COLOR_TYPE_PALETTE) { 580 | const uint8_t *psrc = palette + (*src++ * 4); 581 | r = *psrc++; 582 | g = *psrc++; 583 | b = *psrc++; 584 | a = *psrc++; 585 | } else if (c == OK_PNG_COLOR_TYPE_GRAYSCALE_WITH_ALPHA) { 586 | r = g = b = *src++; 587 | a = *src++; 588 | } else if (c == OK_PNG_COLOR_TYPE_RGB) { 589 | r = *src++; 590 | g = *src++; 591 | b = *src++; 592 | } else if (c == OK_PNG_COLOR_TYPE_RGB_WITH_ALPHA) { 593 | r = *src++; 594 | g = *src++; 595 | b = *src++; 596 | a = *src++; 597 | } 598 | } else if (d == 16) { 599 | if (c == OK_PNG_COLOR_TYPE_GRAYSCALE) { 600 | r = g = b = readBE16(src); 601 | src += 2; 602 | } else if (c == OK_PNG_COLOR_TYPE_GRAYSCALE_WITH_ALPHA) { 603 | r = g = b = readBE16(src); 604 | a = readBE16(src + 2); 605 | src += 4; 606 | } else if (c == OK_PNG_COLOR_TYPE_RGB) { 607 | r = readBE16(src); 608 | g = readBE16(src + 2); 609 | b = readBE16(src + 4); 610 | src += 6; 611 | } else if (c == OK_PNG_COLOR_TYPE_RGB_WITH_ALPHA) { 612 | r = readBE16(src); 613 | g = readBE16(src + 2); 614 | b = readBE16(src + 4); 615 | a = readBE16(src + 6); 616 | src += 8; 617 | } 618 | } 619 | 620 | if (t && r == tr && g == tg && b == tb) { 621 | a = 0; 622 | if (dst_is_premultiplied) { 623 | r = b = g = 0; 624 | } 625 | } 626 | 627 | if (d == 16) { 628 | // This is libpng's formula for scaling 16-bit to 8-bit 629 | r = (r * 255 + 32895) >> 16; 630 | g = (g * 255 + 32895) >> 16; 631 | b = (b * 255 + 32895) >> 16; 632 | a = (a * 255 + 32895) >> 16; 633 | } 634 | 635 | if (should_byteswap) { 636 | *dst++ = (uint8_t)b; 637 | *dst++ = (uint8_t)g; 638 | *dst++ = (uint8_t)r; 639 | *dst++ = (uint8_t)a; 640 | } else { 641 | *dst++ = (uint8_t)r; 642 | *dst++ = (uint8_t)g; 643 | *dst++ = (uint8_t)b; 644 | *dst++ = (uint8_t)a; 645 | } 646 | } 647 | should_byteswap = false; 648 | } 649 | 650 | // Color format convert: Premultiply, un-premultiply, and swap if needed 651 | if (should_byteswap) { 652 | if (has_full_alpha && src_is_premultiplied && !dst_is_premultiplied) { 653 | // Convert from BGRA_PRE to RGBA 654 | for (uint8_t *dst = dst_start; dst < dst_end; dst += 4) { 655 | const uint8_t v = dst[0]; 656 | dst[0] = dst[2]; 657 | dst[2] = v; 658 | ok_png_unpremultiply(dst); 659 | } 660 | } else if (has_full_alpha && !src_is_premultiplied && dst_is_premultiplied) { 661 | // Convert from BGRA to RGBA_PRE 662 | for (uint8_t *dst = dst_start; dst < dst_end; dst += 4) { 663 | const uint8_t v = dst[0]; 664 | dst[0] = dst[2]; 665 | dst[2] = v; 666 | ok_png_premultiply(dst); 667 | } 668 | } else { 669 | // Convert from BGRA to RGBA (or BGRA_PRE to RGBA_PRE) 670 | for (uint8_t *dst = dst_start; dst < dst_end; dst += 4) { 671 | const uint8_t v = dst[0]; 672 | dst[0] = dst[2]; 673 | dst[2] = v; 674 | } 675 | } 676 | } else if (has_full_alpha) { 677 | if (src_is_premultiplied && !dst_is_premultiplied) { 678 | // Convert from RGBA_PRE to RGBA 679 | for (uint8_t *dst = dst_start; dst < dst_end; dst += 4) { 680 | ok_png_unpremultiply(dst); 681 | } 682 | } else if (!src_is_premultiplied && dst_is_premultiplied) { 683 | // Convert from RGBA to RGBA_PRE 684 | for (uint8_t *dst = dst_start; dst < dst_end; dst += 4) { 685 | ok_png_premultiply(dst); 686 | } 687 | } else { 688 | // Do nothing: Already in correct format, RGBA or RGBA_PRE 689 | } 690 | } 691 | 692 | // If interlaced, copy from the temp buffer 693 | if (decoder->interlace_method == 1 && decoder->interlace_pass < 7) { 694 | const int i = decoder->interlace_pass; 695 | const uint32_t s = decoder->scanline; 696 | // 1 2 3 4 5 6 7 697 | static const uint32_t dst_x[] = {0, 0, 4, 0, 2, 0, 1, 0 }; 698 | static const uint32_t dst_dx[] = {0, 8, 8, 4, 4, 2, 2, 1 }; 699 | const uint32_t dst_y[] = {0, s*8, s*8, 4+s*8, s*4, 2+s*4, s*2, 1+s*2 }; 700 | 701 | uint32_t x = dst_x[i]; 702 | uint32_t y = dst_y[i]; 703 | uint32_t dx = 4 * dst_dx[i]; 704 | if (dst_flip_y) { 705 | y = (png->height - y - 1); 706 | } 707 | 708 | src = dst_start; 709 | uint8_t *src_end = dst_end; 710 | uint8_t *dst = png->data + (y * png->stride) + (x * 4); 711 | while (src < src_end) { 712 | memcpy(dst, src, 4); 713 | dst += dx; 714 | src += 4; 715 | } 716 | } 717 | } 718 | 719 | static uint32_t ok_png_get_width_for_pass(const ok_png_decoder *decoder) { 720 | const uint32_t w = decoder->png->width; 721 | if (decoder->interlace_method == 0) { 722 | return w; 723 | } 724 | 725 | switch (decoder->interlace_pass) { 726 | case 1: return (w + 7) / 8; 727 | case 2: return (w + 3) / 8; 728 | case 3: return (w + 3) / 4; 729 | case 4: return (w + 1) / 4; 730 | case 5: return (w + 1) / 2; 731 | case 6: return w / 2; 732 | case 7: return w; 733 | default: return 0; 734 | } 735 | } 736 | 737 | static uint32_t ok_png_get_height_for_pass(const ok_png_decoder *decoder) { 738 | const uint32_t h = decoder->png->height; 739 | if (decoder->interlace_method == 0) { 740 | return h; 741 | } 742 | 743 | switch (decoder->interlace_pass) { 744 | case 1: return (h + 7) / 8; 745 | case 2: return (h + 7) / 8; 746 | case 3: return (h + 3) / 8; 747 | case 4: return (h + 3) / 4; 748 | case 5: return (h + 1) / 4; 749 | case 6: return (h + 1) / 2; 750 | case 7: return h / 2; 751 | default: return 0; 752 | } 753 | } 754 | 755 | static bool ok_png_read_data(ok_png_decoder *decoder, uint32_t bytes_remaining) { 756 | ok_png *png = decoder->png; 757 | size_t inflate_buffer_size = 64 * 1024; 758 | size_t num_passes = decoder->interlace_method == 0 ? 1 : 7; 759 | uint8_t bits_per_pixel = decoder->bit_depth * OK_PNG_SAMPLES_PER_PIXEL[decoder->color_type]; 760 | uint8_t bytes_per_pixel = (bits_per_pixel + 7) / 8; 761 | uint64_t max_bytes_per_scanline = 1 + ((uint64_t)png->width * bits_per_pixel + 7) / 8; 762 | size_t platform_max_bytes_per_scanline = (size_t)max_bytes_per_scanline; 763 | 764 | // Create buffers 765 | if (!png->data) { 766 | if (decoder->allocator.image_alloc) { 767 | decoder->allocator.image_alloc(decoder->allocator_user_data, 768 | png->width, png->height, png->bpp, 769 | &png->data, &png->stride); 770 | } else { 771 | uint64_t size = (uint64_t)png->stride * png->height; 772 | size_t platform_size = (size_t)size; 773 | if (platform_size == size) { 774 | png->data = ok_alloc(decoder, platform_size); 775 | } 776 | } 777 | if (!png->data) { 778 | ok_png_error(png, OK_PNG_ERROR_ALLOCATION, "Couldn't allocate memory for image"); 779 | return false; 780 | } 781 | if (png->stride < png->width * png->bpp) { 782 | ok_png_error(png, OK_PNG_ERROR_API, "Invalid stride"); 783 | return false; 784 | } 785 | } 786 | if (!decoder->prev_scanline) { 787 | if (max_bytes_per_scanline == platform_max_bytes_per_scanline) { 788 | decoder->prev_scanline = ok_alloc(decoder, platform_max_bytes_per_scanline); 789 | } 790 | } 791 | if (!decoder->curr_scanline) { 792 | if (max_bytes_per_scanline == platform_max_bytes_per_scanline) { 793 | decoder->curr_scanline = ok_alloc(decoder, platform_max_bytes_per_scanline); 794 | } 795 | } 796 | if (!decoder->inflate_buffer) { 797 | decoder->inflate_buffer = ok_alloc(decoder, inflate_buffer_size); 798 | } 799 | if (decoder->interlace_method == 1 && !decoder->temp_data_row) { 800 | decoder->temp_data_row = ok_alloc(decoder, png->width * png->bpp); 801 | } 802 | if (!decoder->curr_scanline || !decoder->prev_scanline || !decoder->inflate_buffer || 803 | (decoder->interlace_method == 1 && !decoder->temp_data_row)) { 804 | ok_png_error(png, OK_PNG_ERROR_ALLOCATION, "Couldn't allocate buffers"); 805 | return false; 806 | } 807 | 808 | // Setup inflater 809 | if (!decoder->inflater) { 810 | decoder->inflater = ok_inflater_init(decoder->is_ios_format, 811 | decoder->allocator, decoder->allocator_user_data); 812 | if (!decoder->inflater) { 813 | ok_png_error(png, OK_PNG_ERROR_ALLOCATION, "Couldn't init inflater"); 814 | return false; 815 | } 816 | } 817 | 818 | // Sanity check - this happened with one file in the PNG suite 819 | if (decoder->decoding_completed) { 820 | if (bytes_remaining > 0) { 821 | return ok_seek(decoder, (long)bytes_remaining); 822 | } else { 823 | return true; 824 | } 825 | } 826 | 827 | // Read data 828 | uint32_t curr_width = ok_png_get_width_for_pass(decoder); 829 | uint32_t curr_height = ok_png_get_height_for_pass(decoder); 830 | size_t curr_bytes_per_scanline = (size_t)(1 + ((uint64_t)curr_width * bits_per_pixel + 7) / 8); 831 | while (true) { 832 | // Setup pass 833 | while (decoder->ready_for_next_interlace_pass) { 834 | decoder->ready_for_next_interlace_pass = false; 835 | decoder->scanline = 0; 836 | decoder->interlace_pass++; 837 | if (decoder->interlace_pass == num_passes + 1) { 838 | // Done decoding - skip any remaining chunk data 839 | decoder->decoding_completed = true; 840 | if (bytes_remaining > 0) { 841 | return ok_seek(decoder, (long)bytes_remaining); 842 | } else { 843 | return true; 844 | } 845 | } 846 | curr_width = ok_png_get_width_for_pass(decoder); 847 | curr_height = ok_png_get_height_for_pass(decoder); 848 | curr_bytes_per_scanline = (size_t)(1 + ((uint64_t)curr_width * bits_per_pixel + 7) / 8); 849 | if (curr_width == 0 || curr_height == 0) { 850 | // No data for this pass - happens if width or height <= 4 851 | decoder->ready_for_next_interlace_pass = true; 852 | } else { 853 | memset(decoder->curr_scanline, 0, curr_bytes_per_scanline); 854 | memset(decoder->prev_scanline, 0, curr_bytes_per_scanline); 855 | decoder->inflater_bytes_read = 0; 856 | } 857 | } 858 | 859 | // Read compressed data 860 | if (ok_inflater_needs_input(decoder->inflater)) { 861 | if (bytes_remaining == 0) { 862 | // Need more data, but there is no remaining data in this chunk. 863 | // There may be another IDAT chunk. 864 | return true; 865 | } 866 | const size_t len = min(inflate_buffer_size, bytes_remaining); 867 | if (!ok_read(decoder, decoder->inflate_buffer, len)) { 868 | return false; 869 | } 870 | bytes_remaining -= len; 871 | ok_inflater_set_input(decoder->inflater, decoder->inflate_buffer, len); 872 | } 873 | 874 | // Decompress data 875 | size_t len = ok_inflater_inflate(decoder->inflater, 876 | decoder->curr_scanline + decoder->inflater_bytes_read, 877 | curr_bytes_per_scanline - decoder->inflater_bytes_read); 878 | if (len == OK_SIZE_MAX) { 879 | ok_png_error(png, OK_PNG_ERROR_INFLATER, "Inflater error"); 880 | return false; 881 | } 882 | decoder->inflater_bytes_read += len; 883 | if (decoder->inflater_bytes_read == curr_bytes_per_scanline) { 884 | // Apply filter 885 | const int filter = decoder->curr_scanline[0]; 886 | if (filter > 0 && filter < OK_PNG_NUM_FILTERS) { 887 | ok_png_decode_filter(decoder->curr_scanline + 1, decoder->prev_scanline + 1, 888 | curr_bytes_per_scanline - 1, filter, bytes_per_pixel); 889 | } else if (filter != 0) { 890 | ok_png_error(png, OK_PNG_ERROR_INVALID, "Invalid filter type"); 891 | return false; 892 | } 893 | 894 | // Transform 895 | ok_png_transform_scanline(decoder, decoder->curr_scanline + 1, curr_width); 896 | 897 | // Setup for next scanline or pass 898 | decoder->scanline++; 899 | if (decoder->scanline == curr_height) { 900 | decoder->ready_for_next_interlace_pass = true; 901 | } else { 902 | uint8_t *temp = decoder->curr_scanline; 903 | decoder->curr_scanline = decoder->prev_scanline; 904 | decoder->prev_scanline = temp; 905 | decoder->inflater_bytes_read = 0; 906 | } 907 | } 908 | } 909 | } 910 | 911 | static void ok_png_decode2(ok_png_decoder *decoder) { 912 | ok_png *png = decoder->png; 913 | 914 | uint8_t png_header[8]; 915 | if (!ok_read(decoder, png_header, sizeof(png_header))) { 916 | return; 917 | } 918 | uint8_t png_signature[8] = {137, 80, 78, 71, 13, 10, 26, 10}; 919 | if (memcmp(png_header, png_signature, 8) != 0) { 920 | ok_png_error(decoder->png, OK_PNG_ERROR_INVALID, "Invalid signature (not a PNG file)"); 921 | return; 922 | } 923 | 924 | // When info_only is true, we only care about the IHDR chunk and whether or not 925 | // the tRNS chunk exists. 926 | bool info_only = (decoder->decode_flags & OK_PNG_INFO_ONLY) != 0; 927 | bool hdr_found = false; 928 | bool end_found = false; 929 | while (!end_found) { 930 | uint8_t chunk_header[8]; 931 | uint8_t chunk_footer[4]; 932 | if (!ok_read(decoder, chunk_header, sizeof(chunk_header))) { 933 | return; 934 | } 935 | const uint32_t chunk_length = readBE32(chunk_header); 936 | const uint32_t chunk_type = readBE32(chunk_header + 4); 937 | bool success = false; 938 | 939 | if (!hdr_found && chunk_type != OK_PNG_CHUNK_CGBI && chunk_type != OK_PNG_CHUNK_IHDR) { 940 | ok_png_error(png, OK_PNG_ERROR_INVALID, "IHDR chunk must appear first"); 941 | return; 942 | } 943 | if (chunk_type == OK_PNG_CHUNK_IHDR) { 944 | hdr_found = true; 945 | success = ok_png_read_header(decoder, chunk_length); 946 | if (success && info_only) { 947 | // If the png has alpha, then we have all the info we need. 948 | // Otherwise, continue scanning to see if the tRNS chunk exists. 949 | if (png->has_alpha) { 950 | return; 951 | } 952 | } 953 | } else if (chunk_type == OK_PNG_CHUNK_CGBI) { 954 | success = ok_seek(decoder, (long)chunk_length); 955 | decoder->is_ios_format = true; 956 | } else if (chunk_type == OK_PNG_CHUNK_PLTE && !info_only) { 957 | success = ok_png_read_palette(decoder, chunk_length); 958 | } else if (chunk_type == OK_PNG_CHUNK_TRNS) { 959 | if (info_only) { 960 | // No need to parse this chunk, we have all the info we need. 961 | png->has_alpha = true; 962 | return; 963 | } else { 964 | success = ok_png_read_transparency(decoder, chunk_length); 965 | } 966 | } else if (chunk_type == OK_PNG_CHUNK_IDAT) { 967 | if (info_only) { 968 | // Both IHDR and tRNS must come before IDAT, so we have all the info we need. 969 | return; 970 | } 971 | success = ok_png_read_data(decoder, chunk_length); 972 | } else if (chunk_type == OK_PNG_CHUNK_IEND) { 973 | success = ok_seek(decoder, (long)chunk_length); 974 | end_found = true; 975 | } else { 976 | // Ignore this chunk 977 | success = ok_seek(decoder, (long)chunk_length); 978 | } 979 | 980 | if (!success) { 981 | return; 982 | } 983 | 984 | // Read the footer (CRC) and ignore it 985 | if (!ok_read(decoder, chunk_footer, sizeof(chunk_footer))) { 986 | return; 987 | } 988 | } 989 | 990 | // Sanity check 991 | if (!decoder->decoding_completed) { 992 | ok_png_error(png, OK_PNG_ERROR_INVALID, "Missing imaga data"); 993 | } 994 | } 995 | 996 | void ok_png_decode(ok_png *png, ok_png_decode_flags decode_flags, 997 | ok_png_input input, void *input_user_data, 998 | ok_png_allocator allocator, void *allocator_user_data) { 999 | if (!input.read || !input.seek) { 1000 | ok_png_error(png, OK_PNG_ERROR_API, 1001 | "Invalid argument: input read and seek functions must not be NULL"); 1002 | return; 1003 | } 1004 | 1005 | if (!allocator.alloc || !allocator.free) { 1006 | ok_png_error(png, OK_PNG_ERROR_API, 1007 | "Invalid argument: allocator alloc and free functions must not be NULL"); 1008 | return; 1009 | } 1010 | 1011 | ok_png_decoder *decoder = allocator.alloc(allocator_user_data, sizeof(ok_png_decoder)); 1012 | if (!decoder) { 1013 | ok_png_error(png, OK_PNG_ERROR_ALLOCATION, "Couldn't allocate decoder."); 1014 | return; 1015 | } 1016 | memset(decoder, 0, sizeof(ok_png_decoder)); 1017 | 1018 | decoder->png = png; 1019 | decoder->decode_flags = decode_flags; 1020 | decoder->input = input; 1021 | decoder->input_user_data = input_user_data; 1022 | decoder->allocator = allocator; 1023 | decoder->allocator_user_data = allocator_user_data; 1024 | 1025 | ok_png_decode2(decoder); 1026 | 1027 | // Cleanup decoder 1028 | ok_inflater_free(decoder->inflater); 1029 | allocator.free(allocator_user_data, decoder->curr_scanline); 1030 | allocator.free(allocator_user_data, decoder->prev_scanline); 1031 | allocator.free(allocator_user_data, decoder->inflate_buffer); 1032 | allocator.free(allocator_user_data, decoder->temp_data_row); 1033 | allocator.free(allocator_user_data, decoder); 1034 | } 1035 | 1036 | // 1037 | // Inflater 1038 | // Written from RFC 1950 and RFC 1951. 1039 | // Some of the comments are copy-and-paste from the RFCs. 1040 | // 1041 | 1042 | #define BUFFER_SIZE_BITS 16 1043 | #if BUFFER_SIZE_BITS != 16 1044 | #error The circular buffer pointers, buffer_start_pos, buffer_end_pos, and buffer_offset, \ 1045 | only work with 16-bit buffers. 1046 | #endif 1047 | 1048 | // 32k for back buffer, 32k for forward buffer 1049 | #define BUFFER_SIZE (1 << BUFFER_SIZE_BITS) 1050 | #define BUFFER_SIZE_MASK (BUFFER_SIZE - 1) 1051 | 1052 | #define BLOCK_TYPE_NO_COMPRESSION 0 1053 | #define BLOCK_TYPE_FIXED_HUFFMAN 1 1054 | #define BLOCK_TYPE_DYNAMIC_HUFFMAN 2 1055 | 1056 | typedef enum { 1057 | OK_INFLATER_STATE_READY_FOR_HEAD = 0, 1058 | OK_INFLATER_STATE_READY_FOR_NEXT_BLOCK, 1059 | OK_INFLATER_STATE_READING_STORED_BLOCK_HEADER, 1060 | OK_INFLATER_STATE_READING_STORED_BLOCK, 1061 | OK_INFLATER_STATE_READING_DYNAMIC_BLOCK_HEADER, 1062 | OK_INFLATER_STATE_READING_DYNAMIC_CODE_LENGTHS, 1063 | OK_INFLATER_STATE_READING_DYNAMIC_LITERAL_TREE, 1064 | OK_INFLATER_STATE_READING_DYNAMIC_DISTANCE_TREE, 1065 | OK_INFLATER_STATE_READING_DYNAMIC_COMPRESSED_BLOCK, 1066 | OK_INFLATER_STATE_READING_FIXED_COMPRESSED_BLOCK, 1067 | OK_INFLATER_STATE_READING_DYNAMIC_DISTANCE, 1068 | OK_INFLATER_STATE_READING_FIXED_DISTANCE, 1069 | OK_INFLATER_STATE_DONE, 1070 | OK_INFLATER_STATE_ERROR, 1071 | } ok_inflater_state; 1072 | 1073 | #define VALUE_BITS 9 1074 | #define VALUE_BIT_MASK ((1 << VALUE_BITS) - 1) 1075 | #define MAX_NUM_CODES 289 1076 | #define MAX_CODE_LENGTH 16 1077 | 1078 | static const int OK_INFLATER_DISTANCE_TABLE[] = { 1079 | 1, 2, 3, 4, 5, 7, 9, 13, 17, 25, 1080 | 33, 49, 65, 97, 129, 193, 257, 385, 513, 769, 1081 | 1025, 1537, 2049, 3073, 4097, 6145, 8193, 12289, 16385, 24577 1082 | }; 1083 | static const int OK_INFLATER_DISTANCE_TABLE_LENGTH = 30; 1084 | 1085 | static const int OK_INFLATER_LENGTH_TABLE[] = { 1086 | 3, 4, 5, 6, 7, 8, 9, 10, 11, 13, 1087 | 15, 17, 19, 23, 27, 31, 35, 43, 51, 59, 1088 | 67, 83, 99, 115, 131, 163, 195, 227, 258 1089 | }; 1090 | 1091 | static const int OK_INFLATER_BIT_LENGTH_TABLE[] = { 1092 | 16, 17, 18, 0, 8, 7, 9, 6, 10, 5, 11, 4, 12, 3, 13, 2, 14, 1, 15 1093 | }; 1094 | static const int OK_INFLATER_BIT_LENGTH_TABLE_LENGTH = 19; 1095 | 1096 | typedef struct { 1097 | uint16_t lookup_table[1 << (MAX_CODE_LENGTH - 1)]; 1098 | unsigned int bits; 1099 | } ok_inflater_huffman_tree; 1100 | 1101 | struct ok_inflater { 1102 | // Options 1103 | bool nowrap; 1104 | 1105 | // Allocator 1106 | ok_png_allocator allocator; 1107 | void *allocator_user_data; 1108 | 1109 | // Input 1110 | const uint8_t *input; 1111 | const uint8_t *input_end; 1112 | uint32_t input_buffer; 1113 | unsigned int input_buffer_bits; 1114 | 1115 | // Inflate data 1116 | uint8_t *buffer; 1117 | uint16_t buffer_start_pos; 1118 | uint16_t buffer_end_pos; 1119 | bool final_block; 1120 | ok_inflater_state state; 1121 | int state_count; 1122 | int state_distance; 1123 | 1124 | // For dynamic blocks 1125 | int num_literal_codes; 1126 | int num_distance_codes; 1127 | int num_code_length_codes; 1128 | uint8_t tree_codes[MAX_NUM_CODES]; 1129 | ok_inflater_huffman_tree *code_length_huffman; 1130 | int huffman_code; 1131 | 1132 | // Huffman 1133 | ok_inflater_huffman_tree *literal_huffman; 1134 | ok_inflater_huffman_tree *distance_huffman; 1135 | ok_inflater_huffman_tree *fixed_literal_huffman; 1136 | ok_inflater_huffman_tree *fixed_distance_huffman; 1137 | }; 1138 | 1139 | #define ok_inflater_error(inflater, message) ok_inflater_set_error(inflater) 1140 | 1141 | static void ok_inflater_set_error(ok_inflater *inflater) { 1142 | if (inflater) { 1143 | inflater->state = OK_INFLATER_STATE_ERROR; 1144 | } 1145 | } 1146 | 1147 | // 1148 | // Buffer writing / flushing 1149 | // 1150 | 1151 | // is_buffer_full is commented out because it is not used, but it helps to understand how the buffer works. 1152 | //static inline bool is_buffer_full(const ok_inflater *inflater) { 1153 | // return (uint16_t)(inflater->buffer_end_pos + 1) == inflater->buffer_start_pos; 1154 | //} 1155 | 1156 | // Number of bytes that can be written until full or buffer wrap 1157 | static inline uint16_t ok_inflater_can_write(const ok_inflater *inflater) { 1158 | if (inflater->buffer_start_pos == 0) { 1159 | return -inflater->buffer_end_pos - 1; 1160 | } else if (inflater->buffer_start_pos > inflater->buffer_end_pos) { 1161 | return inflater->buffer_start_pos - inflater->buffer_end_pos - 1; 1162 | } else { 1163 | return -inflater->buffer_end_pos; 1164 | } 1165 | } 1166 | 1167 | static inline uint16_t ok_inflater_can_write_total(const ok_inflater *inflater) { 1168 | return inflater->buffer_start_pos - inflater->buffer_end_pos - 1; 1169 | } 1170 | 1171 | static inline void ok_inflater_write_byte(ok_inflater *inflater, const uint8_t b) { 1172 | inflater->buffer[inflater->buffer_end_pos & BUFFER_SIZE_MASK] = b; 1173 | inflater->buffer_end_pos++; 1174 | } 1175 | 1176 | static inline int ok_inflater_write_bytes(ok_inflater *inflater, const uint8_t *src, int len) { 1177 | int bytes_remaining = len; 1178 | while (bytes_remaining > 0) { 1179 | int n = min(bytes_remaining, ok_inflater_can_write(inflater)); 1180 | if (n == 0) { 1181 | return len - bytes_remaining; 1182 | } 1183 | memcpy(inflater->buffer + inflater->buffer_end_pos, src, (size_t)n); 1184 | inflater->buffer_end_pos += n; 1185 | bytes_remaining -= n; 1186 | src += n; 1187 | } 1188 | return len; 1189 | } 1190 | 1191 | static inline int ok_inflater_write_byte_n(ok_inflater *inflater, const uint8_t b, int len) { 1192 | int bytes_remaining = len; 1193 | while (bytes_remaining > 0) { 1194 | int n = min(bytes_remaining, ok_inflater_can_write(inflater)); 1195 | if (n == 0) { 1196 | return len - bytes_remaining; 1197 | } 1198 | memset(inflater->buffer + inflater->buffer_end_pos, b, (size_t)n); 1199 | inflater->buffer_end_pos += n; 1200 | bytes_remaining -= n; 1201 | } 1202 | return len; 1203 | } 1204 | 1205 | static inline uint16_t ok_inflater_can_flush_total(const ok_inflater *inflater) { 1206 | return inflater->buffer_end_pos - inflater->buffer_start_pos; 1207 | } 1208 | 1209 | // Number of bytes that can be flushed until empty of buffer wrap 1210 | static inline uint16_t ok_inflater_can_flush(const ok_inflater *inflater) { 1211 | if (inflater->buffer_start_pos <= inflater->buffer_end_pos) { 1212 | return inflater->buffer_end_pos - inflater->buffer_start_pos; 1213 | } else { 1214 | return -inflater->buffer_start_pos; 1215 | } 1216 | } 1217 | 1218 | static inline size_t ok_inflater_flush(ok_inflater *inflater, uint8_t *dst, size_t len) { 1219 | size_t bytes_remaining = len; 1220 | while (bytes_remaining > 0) { 1221 | size_t n = min(bytes_remaining, ok_inflater_can_flush(inflater)); 1222 | if (n == 0) { 1223 | return len - bytes_remaining; 1224 | } 1225 | memcpy(dst, inflater->buffer + inflater->buffer_start_pos, n); 1226 | inflater->buffer_start_pos += n; 1227 | bytes_remaining -= n; 1228 | dst += n; 1229 | } 1230 | return len; 1231 | } 1232 | 1233 | // Read from input 1234 | 1235 | static inline void ok_inflater_skip_byte_align(ok_inflater *inflater) { 1236 | unsigned int skip_bits = inflater->input_buffer_bits & 7; 1237 | inflater->input_buffer >>= skip_bits; 1238 | inflater->input_buffer_bits -= skip_bits; 1239 | } 1240 | 1241 | static inline bool ok_inflater_load_bits(ok_inflater *inflater, unsigned int num_bits) { 1242 | while (inflater->input_buffer_bits < num_bits) { 1243 | if (inflater->input == inflater->input_end) { 1244 | return false; 1245 | } 1246 | uint32_t input = *inflater->input++; 1247 | inflater->input_buffer |= input << inflater->input_buffer_bits; 1248 | inflater->input_buffer_bits += 8; 1249 | } 1250 | return true; 1251 | } 1252 | 1253 | // Assumes at least num_bits bits are loaded into buffer (call load_bits first) 1254 | static inline uint32_t ok_inflater_read_bits(ok_inflater *inflater, unsigned int num_bits) { 1255 | uint32_t ans = inflater->input_buffer & ((1 << num_bits) - 1); 1256 | inflater->input_buffer >>= num_bits; 1257 | inflater->input_buffer_bits -= num_bits; 1258 | return ans; 1259 | } 1260 | 1261 | // Assumes at least num_bits bits are loaded into buffer (call load_bits first) 1262 | static inline uint32_t ok_inflater_peek_bits(ok_inflater *inflater, unsigned int num_bits) { 1263 | return inflater->input_buffer & ((1 << num_bits) - 1); 1264 | } 1265 | 1266 | // Huffman 1267 | 1268 | static inline uint32_t ok_inflater_reverse_bits(uint32_t value, unsigned int num_bits) { 1269 | uint32_t rev_value = value & 1; 1270 | for (unsigned int i = num_bits - 1; i > 0; i--) { 1271 | value >>= 1; 1272 | rev_value <<= 1; 1273 | rev_value |= value & 1; 1274 | } 1275 | return rev_value; 1276 | } 1277 | 1278 | static int ok_inflater_decode_literal(ok_inflater *inflater, const uint16_t *tree_lookup_table, 1279 | unsigned int tree_bits) { 1280 | if (!ok_inflater_load_bits(inflater, tree_bits)) { 1281 | return -1; 1282 | } 1283 | uint32_t p = ok_inflater_peek_bits(inflater, tree_bits); 1284 | uint16_t value = tree_lookup_table[p]; 1285 | ok_inflater_read_bits(inflater, value >> VALUE_BITS); 1286 | return value & VALUE_BIT_MASK; 1287 | } 1288 | 1289 | static void ok_inflater_make_huffman_tree_from_array(ok_inflater_huffman_tree *tree, 1290 | const uint8_t *code_length, int length) { 1291 | tree->bits = 1; 1292 | 1293 | // Count the number of codes for each code length. 1294 | // Let code_length_count[n] be the number of codes of length n, n >= 1. 1295 | unsigned int code_length_count[MAX_CODE_LENGTH]; 1296 | int i; 1297 | for (i = 0; i < MAX_CODE_LENGTH; i++) { 1298 | code_length_count[i] = 0; 1299 | } 1300 | for (i = 0; i < length; i++) { 1301 | code_length_count[code_length[i]]++; 1302 | } 1303 | 1304 | // Find the numerical value of the smallest code for each code length: 1305 | unsigned int next_code[MAX_CODE_LENGTH]; 1306 | unsigned int code = 0; 1307 | for (i = 1; i < MAX_CODE_LENGTH; i++) { 1308 | code = (code + code_length_count[i - 1]) << 1; 1309 | next_code[i] = code; 1310 | if (code_length_count[i] != 0) { 1311 | tree->bits = (unsigned int)i; 1312 | } 1313 | } 1314 | 1315 | // Init lookup table 1316 | const unsigned int max = 1 << tree->bits; 1317 | memset(tree->lookup_table, 0, sizeof(tree->lookup_table[0]) * max); 1318 | 1319 | // Assign numerical values to all codes, using consecutive values for all 1320 | // codes of the same length with the base values determined at step 2. 1321 | // Codes that are never used (which have a bit length of zero) must not be 1322 | // assigned a value. 1323 | for (i = 0; i < length; i++) { 1324 | unsigned int len = code_length[i]; 1325 | if (len != 0) { 1326 | code = next_code[len]; 1327 | next_code[len]++; 1328 | 1329 | unsigned int value = (unsigned int)i | (len << VALUE_BITS); 1330 | tree->lookup_table[ok_inflater_reverse_bits(code, len)] = (uint16_t)value; 1331 | } 1332 | } 1333 | 1334 | // Fill in the missing parts of the lookup table 1335 | int next_limit = 1; 1336 | int num_bits = 0; 1337 | int mask = 0; 1338 | for (i = 1; i < (int)max; i++) { 1339 | if (i == next_limit) { 1340 | mask = (1 << num_bits) - 1; 1341 | num_bits++; 1342 | next_limit <<= 1; 1343 | } 1344 | if (tree->lookup_table[i] == 0) { 1345 | tree->lookup_table[i] = tree->lookup_table[i & mask]; 1346 | } 1347 | } 1348 | } 1349 | 1350 | static bool ok_inflater_inflate_huffman_tree(ok_inflater *inflater, ok_inflater_huffman_tree *tree, 1351 | ok_inflater_huffman_tree *code_length_huffman, 1352 | int num_codes) { 1353 | if (num_codes < 0 || num_codes >= MAX_NUM_CODES) { 1354 | ok_inflater_error(inflater, "Invalid num_codes"); 1355 | return false; 1356 | } 1357 | const uint16_t *tree_lookup_table = code_length_huffman->lookup_table; 1358 | const unsigned int tree_bits = code_length_huffman->bits; 1359 | // 0 - 15: Represent code lengths of 0 - 15 1360 | // 16: Copy the previous code length 3 - 6 times. 1361 | // (2 bits of length) 1362 | // 17: Repeat a code length of 0 for 3 - 10 times. 1363 | // (3 bits of length) 1364 | // 18: Repeat a code length of 0 for 11 - 138 times 1365 | // (7 bits of length) 1366 | while (inflater->state_count < num_codes) { 1367 | if (inflater->huffman_code < 0) { 1368 | inflater->huffman_code = ok_inflater_decode_literal(inflater, tree_lookup_table, 1369 | tree_bits); 1370 | if (inflater->huffman_code < 0) { 1371 | return false; 1372 | } 1373 | } 1374 | if (inflater->huffman_code <= 15) { 1375 | inflater->tree_codes[inflater->state_count++] = (uint8_t)inflater->huffman_code; 1376 | } else { 1377 | int value = 0; 1378 | int len; 1379 | unsigned int len_bits; 1380 | switch (inflater->huffman_code) { 1381 | case 16: 1382 | len = 3; 1383 | len_bits = 2; 1384 | if (inflater->state_count == 0) { 1385 | ok_inflater_error(inflater, "Invalid previous code"); 1386 | return false; 1387 | } 1388 | value = inflater->tree_codes[inflater->state_count - 1]; 1389 | break; 1390 | case 17: 1391 | len = 3; 1392 | len_bits = 3; 1393 | break; 1394 | case 18: 1395 | len = 11; 1396 | len_bits = 7; 1397 | break; 1398 | default: 1399 | ok_inflater_error(inflater, "Invalid huffman code"); 1400 | return false; 1401 | } 1402 | if (!ok_inflater_load_bits(inflater, len_bits)) { 1403 | return false; 1404 | } 1405 | len += ok_inflater_read_bits(inflater, len_bits); 1406 | if (len > num_codes - inflater->state_count) { 1407 | ok_inflater_error(inflater, "Invalid length"); 1408 | return false; 1409 | } 1410 | memset(inflater->tree_codes + inflater->state_count, value, (size_t)len); 1411 | inflater->state_count += len; 1412 | } 1413 | inflater->huffman_code = -1; 1414 | } 1415 | ok_inflater_make_huffman_tree_from_array(tree, inflater->tree_codes, num_codes); 1416 | return true; 1417 | } 1418 | 1419 | // Inflate 1420 | 1421 | static bool ok_inflater_zlib_header(ok_inflater *inflater) { 1422 | if (!ok_inflater_load_bits(inflater, 16)) { 1423 | return false; 1424 | } else { 1425 | uint32_t compression_method = ok_inflater_read_bits(inflater, 4); 1426 | uint32_t compression_info = ok_inflater_read_bits(inflater, 4); 1427 | uint32_t flag_check = ok_inflater_read_bits(inflater, 5); 1428 | uint32_t flag_dict = ok_inflater_read_bits(inflater, 1); 1429 | uint32_t flag_compression_level = ok_inflater_read_bits(inflater, 2); 1430 | 1431 | uint32_t bits = ((compression_info << 12) | (compression_method << 8) | 1432 | (flag_compression_level << 6) | (flag_dict << 5) | flag_check); 1433 | if (bits % 31 != 0) { 1434 | ok_inflater_error(inflater, "Invalid zlib header"); 1435 | return false; 1436 | } 1437 | if (compression_method != 8) { 1438 | ok_inflater_error(inflater, "Invalid inflater compression method"); 1439 | return false; 1440 | } 1441 | if (compression_info > 7) { 1442 | ok_inflater_error(inflater, "Invalid window size"); 1443 | return false; 1444 | } 1445 | if (flag_dict) { 1446 | ok_inflater_error(inflater, "Needs external dictionary"); 1447 | return false; 1448 | } 1449 | 1450 | inflater->state = OK_INFLATER_STATE_READY_FOR_NEXT_BLOCK; 1451 | return true; 1452 | } 1453 | } 1454 | 1455 | static bool ok_inflater_init_fixed_huffman(ok_inflater *inflater) { 1456 | if (!inflater->fixed_literal_huffman) { 1457 | ok_inflater_huffman_tree *tree = ok_alloc(inflater, sizeof(ok_inflater_huffman_tree)); 1458 | if (tree) { 1459 | uint8_t code_length[288]; 1460 | int i; 1461 | for (i = 0; i < 144; i++) { 1462 | code_length[i] = 8; 1463 | } 1464 | for (i = 144; i < 256; i++) { 1465 | code_length[i] = 9; 1466 | } 1467 | for (i = 256; i < 280; i++) { 1468 | code_length[i] = 7; 1469 | } 1470 | for (i = 280; i < 288; i++) { 1471 | code_length[i] = 8; 1472 | } 1473 | ok_inflater_make_huffman_tree_from_array(tree, code_length, 1474 | sizeof(code_length) / sizeof(code_length[0])); 1475 | inflater->fixed_literal_huffman = tree; 1476 | } 1477 | } 1478 | if (!inflater->fixed_distance_huffman) { 1479 | ok_inflater_huffman_tree *tree = ok_alloc(inflater, sizeof(ok_inflater_huffman_tree)); 1480 | if (tree) { 1481 | uint8_t distance_code_length[32]; 1482 | for (int i = 0; i < 32; i++) { 1483 | distance_code_length[i] = 5; 1484 | } 1485 | ok_inflater_make_huffman_tree_from_array(tree, distance_code_length, 32); 1486 | inflater->fixed_distance_huffman = tree; 1487 | } 1488 | } 1489 | return inflater->fixed_literal_huffman && inflater->fixed_distance_huffman; 1490 | } 1491 | 1492 | static bool ok_inflater_next_block(ok_inflater *inflater) { 1493 | if (inflater->final_block) { 1494 | inflater->state = OK_INFLATER_STATE_DONE; 1495 | ok_inflater_skip_byte_align(inflater); 1496 | return true; 1497 | } else if (!ok_inflater_load_bits(inflater, 3)) { 1498 | return false; 1499 | } else { 1500 | inflater->final_block = ok_inflater_read_bits(inflater, 1); 1501 | uint32_t block_type = ok_inflater_read_bits(inflater, 2); 1502 | switch (block_type) { 1503 | case BLOCK_TYPE_NO_COMPRESSION: 1504 | inflater->state = OK_INFLATER_STATE_READING_STORED_BLOCK_HEADER; 1505 | break; 1506 | case BLOCK_TYPE_DYNAMIC_HUFFMAN: 1507 | inflater->state = OK_INFLATER_STATE_READING_DYNAMIC_BLOCK_HEADER; 1508 | break; 1509 | case BLOCK_TYPE_FIXED_HUFFMAN: { 1510 | if (!ok_inflater_init_fixed_huffman(inflater)) { 1511 | ok_inflater_error(inflater, "Couldn't initilize fixed huffman trees"); 1512 | return false; 1513 | } 1514 | inflater->state = OK_INFLATER_STATE_READING_FIXED_COMPRESSED_BLOCK; 1515 | inflater->huffman_code = -1; 1516 | break; 1517 | } 1518 | default: 1519 | ok_inflater_error(inflater, "Invalid block type"); 1520 | break; 1521 | } 1522 | return true; 1523 | } 1524 | } 1525 | 1526 | static bool ok_inflater_stored_block_header(ok_inflater *inflater) { 1527 | ok_inflater_skip_byte_align(inflater); 1528 | if (!ok_inflater_load_bits(inflater, 32)) { 1529 | return false; 1530 | } else { 1531 | uint32_t len = ok_inflater_read_bits(inflater, 16); 1532 | uint32_t clen = ok_inflater_read_bits(inflater, 16); 1533 | if ((len & 0xffff) != ((~clen) & 0xffff)) { 1534 | ok_inflater_error(inflater, "Invalid stored block"); 1535 | return false; 1536 | } else if (len == 0) { 1537 | inflater->state = OK_INFLATER_STATE_READY_FOR_NEXT_BLOCK; 1538 | return true; 1539 | } else { 1540 | inflater->state = OK_INFLATER_STATE_READING_STORED_BLOCK; 1541 | inflater->state_count = (int)len; 1542 | return true; 1543 | } 1544 | } 1545 | } 1546 | 1547 | static bool ok_inflater_stored_block(ok_inflater *inflater) { 1548 | const intptr_t can_read = inflater->input_end - inflater->input; 1549 | if (can_read == 0) { 1550 | return false; 1551 | } else { 1552 | int len = ok_inflater_write_bytes(inflater, inflater->input, 1553 | min((int)can_read, inflater->state_count)); 1554 | if (len == 0) { 1555 | // Buffer full 1556 | return false; 1557 | } 1558 | inflater->input += len; 1559 | inflater->state_count -= len; 1560 | if (inflater->state_count == 0) { 1561 | inflater->state = OK_INFLATER_STATE_READY_FOR_NEXT_BLOCK; 1562 | } 1563 | return true; 1564 | } 1565 | } 1566 | 1567 | static int ok_inflater_decode_distance(ok_inflater *inflater, 1568 | const ok_inflater_huffman_tree *tree) { 1569 | if (!ok_inflater_load_bits(inflater, tree->bits)) { 1570 | return -1; 1571 | } 1572 | uint32_t p = ok_inflater_peek_bits(inflater, tree->bits); 1573 | uint16_t tree_value = tree->lookup_table[p]; 1574 | int value = tree_value & VALUE_BIT_MASK; 1575 | unsigned int value_bits = tree_value >> VALUE_BITS; 1576 | if (value < 4) { 1577 | ok_inflater_read_bits(inflater, value_bits); 1578 | return value + 1; 1579 | } else if (value >= OK_INFLATER_DISTANCE_TABLE_LENGTH) { 1580 | ok_inflater_error(inflater, "Invalid distance"); 1581 | return -1; 1582 | } else { 1583 | unsigned int extra_bits = (unsigned int)((value >> 1) - 1); 1584 | if (!ok_inflater_load_bits(inflater, value_bits + extra_bits)) { 1585 | return -1; 1586 | } 1587 | int d = (int)(ok_inflater_read_bits(inflater, value_bits + extra_bits) >> value_bits); 1588 | return OK_INFLATER_DISTANCE_TABLE[value] + d; 1589 | } 1590 | } 1591 | 1592 | static int ok_inflater_decode_length(ok_inflater *inflater, int value) { 1593 | if (value < 8) { 1594 | return value + 3; 1595 | } else { 1596 | int len = OK_INFLATER_LENGTH_TABLE[value]; 1597 | unsigned int extra_bits = (unsigned int)((value >> 2) - 1); 1598 | if (extra_bits <= 5) { 1599 | if (!ok_inflater_load_bits(inflater, extra_bits)) { 1600 | return -1; 1601 | } 1602 | len += ok_inflater_read_bits(inflater, extra_bits); 1603 | } 1604 | return len; 1605 | } 1606 | } 1607 | 1608 | static bool ok_inflater_distance_with_tree(ok_inflater *inflater, 1609 | const ok_inflater_huffman_tree *tree) { 1610 | if (inflater->state_count < 0) { 1611 | inflater->state_count = ok_inflater_decode_length(inflater, inflater->huffman_code); 1612 | if (inflater->state_count < 0) { 1613 | // Needs input 1614 | return false; 1615 | } 1616 | inflater->huffman_code = -1; 1617 | } 1618 | if (inflater->state_distance < 0) { 1619 | inflater->state_distance = ok_inflater_decode_distance(inflater, tree); 1620 | if (inflater->state_distance < 0) { 1621 | // Needs input 1622 | return false; 1623 | } 1624 | } 1625 | 1626 | // Copy len bytes from offset to buffer_end_pos 1627 | int buffer_offset = (inflater->buffer_end_pos - inflater->state_distance) & BUFFER_SIZE_MASK; 1628 | if (inflater->state_distance == 1) { 1629 | // Optimization: can use memset 1630 | int n = inflater->state_count; 1631 | int n2 = ok_inflater_write_byte_n(inflater, inflater->buffer[buffer_offset], n); 1632 | inflater->state_count -= n2; 1633 | if (n2 != n) { 1634 | // Full buffer 1635 | return false; 1636 | } 1637 | } else if (buffer_offset + inflater->state_count < BUFFER_SIZE) { 1638 | // Optimization: the offset won't wrap 1639 | int bytes_copyable = inflater->state_distance; 1640 | while (inflater->state_count > 0) { 1641 | int n = min(inflater->state_count, bytes_copyable); 1642 | int n2 = ok_inflater_write_bytes(inflater, inflater->buffer + buffer_offset, n); 1643 | inflater->state_count -= n2; 1644 | bytes_copyable += n2; 1645 | if (n2 != n) { 1646 | // Full buffer 1647 | return false; 1648 | } 1649 | } 1650 | } else { 1651 | // This could be optimized, but it happens rarely, so it's probably not worth it 1652 | while (inflater->state_count > 0) { 1653 | int n = min(inflater->state_count, inflater->state_distance); 1654 | n = min(n, (BUFFER_SIZE - buffer_offset)); 1655 | int n2 = ok_inflater_write_bytes(inflater, inflater->buffer + buffer_offset, n); 1656 | inflater->state_count -= n2; 1657 | buffer_offset = (buffer_offset + n2) & BUFFER_SIZE_MASK; 1658 | if (n2 != n) { 1659 | // Full buffer 1660 | return false; 1661 | } 1662 | } 1663 | } 1664 | return true; 1665 | } 1666 | 1667 | static bool ok_inflater_distance(ok_inflater *inflater) { 1668 | bool is_fixed = inflater->state == OK_INFLATER_STATE_READING_FIXED_DISTANCE; 1669 | const ok_inflater_huffman_tree *distance_tree = 1670 | (is_fixed ? inflater->fixed_distance_huffman : inflater->distance_huffman); 1671 | if (ok_inflater_distance_with_tree(inflater, distance_tree)) { 1672 | if (is_fixed) { 1673 | inflater->state = OK_INFLATER_STATE_READING_FIXED_COMPRESSED_BLOCK; 1674 | } else { 1675 | inflater->state = OK_INFLATER_STATE_READING_DYNAMIC_COMPRESSED_BLOCK; 1676 | } 1677 | return true; 1678 | } else { 1679 | return false; 1680 | } 1681 | } 1682 | 1683 | static bool ok_inflater_dynamic_block_header(ok_inflater *inflater) { 1684 | if (!ok_inflater_load_bits(inflater, 14)) { 1685 | return false; 1686 | } else { 1687 | inflater->num_literal_codes = (int)ok_inflater_read_bits(inflater, 5) + 257; 1688 | inflater->num_distance_codes = (int)ok_inflater_read_bits(inflater, 5) + 1; 1689 | inflater->num_code_length_codes = (int)ok_inflater_read_bits(inflater, 4) + 4; 1690 | 1691 | for (int i = inflater->num_code_length_codes; i < OK_INFLATER_BIT_LENGTH_TABLE_LENGTH; i++) { 1692 | inflater->tree_codes[OK_INFLATER_BIT_LENGTH_TABLE[i]] = 0; 1693 | } 1694 | 1695 | inflater->state = OK_INFLATER_STATE_READING_DYNAMIC_CODE_LENGTHS; 1696 | inflater->state_count = inflater->num_code_length_codes; 1697 | return true; 1698 | } 1699 | } 1700 | 1701 | static bool ok_inflater_dynamic_block_code_lengths(ok_inflater *inflater) { 1702 | while (inflater->state_count > 0) { 1703 | if (!ok_inflater_load_bits(inflater, 3)) { 1704 | return false; 1705 | } 1706 | int index = inflater->num_code_length_codes - inflater->state_count; 1707 | inflater->tree_codes[OK_INFLATER_BIT_LENGTH_TABLE[index]] = 1708 | (uint8_t)ok_inflater_read_bits(inflater, 3); 1709 | inflater->state_count--; 1710 | } 1711 | ok_inflater_make_huffman_tree_from_array(inflater->code_length_huffman, 1712 | inflater->tree_codes, 1713 | OK_INFLATER_BIT_LENGTH_TABLE_LENGTH); 1714 | 1715 | inflater->state = OK_INFLATER_STATE_READING_DYNAMIC_LITERAL_TREE; 1716 | inflater->huffman_code = -1; 1717 | inflater->state_count = 0; 1718 | return true; 1719 | } 1720 | 1721 | static bool ok_inflater_compressed_block(ok_inflater *inflater) { 1722 | const bool is_fixed = inflater->state == OK_INFLATER_STATE_READING_FIXED_COMPRESSED_BLOCK; 1723 | const ok_inflater_huffman_tree *literal_tree = 1724 | (is_fixed ? inflater->fixed_literal_huffman : inflater->literal_huffman); 1725 | const ok_inflater_huffman_tree *distance_tree = 1726 | (is_fixed ? inflater->fixed_distance_huffman : inflater->distance_huffman); 1727 | 1728 | // decode literal/length value from input stream 1729 | 1730 | size_t max_write = ok_inflater_can_write_total(inflater); 1731 | const uint16_t *tree_lookup_table = literal_tree->lookup_table; 1732 | const unsigned int tree_bits = literal_tree->bits; 1733 | while (max_write > 0) { 1734 | int value = ok_inflater_decode_literal(inflater, tree_lookup_table, tree_bits); 1735 | if (value < 0) { 1736 | // Needs input 1737 | return false; 1738 | } else if (value < 256) { 1739 | ok_inflater_write_byte(inflater, (uint8_t)value); 1740 | max_write--; 1741 | } else if (value == 256) { 1742 | inflater->state = OK_INFLATER_STATE_READY_FOR_NEXT_BLOCK; 1743 | return true; 1744 | } else if (value < 286) { 1745 | inflater->huffman_code = value - 257; 1746 | inflater->state_count = -1; 1747 | inflater->state_distance = -1; 1748 | if (ok_inflater_distance_with_tree(inflater, distance_tree)) { 1749 | max_write = ok_inflater_can_write_total(inflater); 1750 | } else { 1751 | if (is_fixed) { 1752 | inflater->state = OK_INFLATER_STATE_READING_FIXED_DISTANCE; 1753 | } else { 1754 | inflater->state = OK_INFLATER_STATE_READING_DYNAMIC_DISTANCE; 1755 | } 1756 | return false; 1757 | } 1758 | } else { 1759 | ok_inflater_error(inflater, "Invalid inflater literal"); 1760 | return false; 1761 | } 1762 | } 1763 | // Output buffer full 1764 | return false; 1765 | } 1766 | 1767 | static bool ok_inflater_literal_tree(ok_inflater *inflater) { 1768 | bool done = ok_inflater_inflate_huffman_tree(inflater, inflater->literal_huffman, 1769 | inflater->code_length_huffman, 1770 | inflater->num_literal_codes); 1771 | if (done) { 1772 | inflater->state = OK_INFLATER_STATE_READING_DYNAMIC_DISTANCE_TREE; 1773 | inflater->huffman_code = -1; 1774 | inflater->state_count = 0; 1775 | return true; 1776 | } else { 1777 | return false; 1778 | } 1779 | } 1780 | 1781 | static bool ok_inflater_distance_tree(ok_inflater *inflater) { 1782 | bool done = ok_inflater_inflate_huffman_tree(inflater, inflater->distance_huffman, 1783 | inflater->code_length_huffman, 1784 | inflater->num_distance_codes); 1785 | if (done) { 1786 | inflater->state = OK_INFLATER_STATE_READING_DYNAMIC_COMPRESSED_BLOCK; 1787 | inflater->huffman_code = -1; 1788 | return true; 1789 | } else { 1790 | return false; 1791 | } 1792 | } 1793 | 1794 | static bool ok_inflater_noop(ok_inflater *inflater) { 1795 | (void)inflater; 1796 | return false; 1797 | } 1798 | 1799 | static bool (*OK_INFLATER_STATE_FUNCTIONS[])(ok_inflater *) = { 1800 | ok_inflater_zlib_header, 1801 | ok_inflater_next_block, 1802 | ok_inflater_stored_block_header, 1803 | ok_inflater_stored_block, 1804 | ok_inflater_dynamic_block_header, 1805 | ok_inflater_dynamic_block_code_lengths, 1806 | ok_inflater_literal_tree, 1807 | ok_inflater_distance_tree, 1808 | ok_inflater_compressed_block, 1809 | ok_inflater_compressed_block, 1810 | ok_inflater_distance, 1811 | ok_inflater_distance, 1812 | ok_inflater_noop, 1813 | ok_inflater_noop 1814 | }; 1815 | 1816 | // Public Inflater API 1817 | 1818 | ok_inflater *ok_inflater_init(bool nowrap, ok_png_allocator allocator, void *allocator_user_data) { 1819 | ok_inflater *inflater = allocator.alloc(allocator_user_data, sizeof(ok_inflater)); 1820 | if (inflater) { 1821 | memset(inflater, 0, sizeof(ok_inflater)); 1822 | inflater->nowrap = nowrap; 1823 | inflater->allocator = allocator; 1824 | inflater->allocator_user_data = allocator_user_data; 1825 | inflater->state = (nowrap ? OK_INFLATER_STATE_READY_FOR_NEXT_BLOCK : 1826 | OK_INFLATER_STATE_READY_FOR_HEAD); 1827 | inflater->buffer = ok_alloc(inflater, BUFFER_SIZE); 1828 | inflater->code_length_huffman = ok_alloc(inflater, sizeof(ok_inflater_huffman_tree)); 1829 | inflater->literal_huffman = ok_alloc(inflater, sizeof(ok_inflater_huffman_tree)); 1830 | inflater->distance_huffman = ok_alloc(inflater, sizeof(ok_inflater_huffman_tree)); 1831 | 1832 | if (!inflater->buffer || 1833 | !inflater->code_length_huffman || 1834 | !inflater->literal_huffman || 1835 | !inflater->distance_huffman) { 1836 | ok_inflater_free(inflater); 1837 | inflater = NULL; 1838 | } 1839 | } 1840 | return inflater; 1841 | } 1842 | 1843 | void ok_inflater_reset(ok_inflater *inflater) { 1844 | if (inflater) { 1845 | inflater->input = NULL; 1846 | inflater->input_end = NULL; 1847 | inflater->input_buffer = 0; 1848 | inflater->input_buffer_bits = 0; 1849 | 1850 | inflater->buffer_start_pos = 0; 1851 | inflater->buffer_end_pos = 0; 1852 | inflater->final_block = false; 1853 | inflater->state = (inflater->nowrap ? OK_INFLATER_STATE_READY_FOR_NEXT_BLOCK : 1854 | OK_INFLATER_STATE_READY_FOR_HEAD); 1855 | } 1856 | } 1857 | 1858 | void ok_inflater_free(ok_inflater *inflater) { 1859 | if (inflater) { 1860 | ok_png_allocator allocator = inflater->allocator; 1861 | void *allocator_user_data = inflater->allocator_user_data; 1862 | allocator.free(allocator_user_data, inflater->buffer); 1863 | allocator.free(allocator_user_data, inflater->code_length_huffman); 1864 | allocator.free(allocator_user_data, inflater->literal_huffman); 1865 | allocator.free(allocator_user_data, inflater->distance_huffman); 1866 | allocator.free(allocator_user_data, inflater->fixed_literal_huffman); 1867 | allocator.free(allocator_user_data, inflater->fixed_distance_huffman); 1868 | allocator.free(allocator_user_data, inflater); 1869 | } 1870 | } 1871 | 1872 | bool ok_inflater_needs_input(const ok_inflater *inflater) { 1873 | return inflater && 1874 | inflater->state != OK_INFLATER_STATE_ERROR && 1875 | ok_inflater_can_flush_total(inflater) == 0 && 1876 | inflater->input == inflater->input_end; 1877 | } 1878 | 1879 | void ok_inflater_set_input(ok_inflater *inflater, const uint8_t *buffer, size_t buffer_length) { 1880 | if (inflater) { 1881 | if (inflater->input == inflater->input_end) { 1882 | inflater->input = buffer; 1883 | inflater->input_end = inflater->input + buffer_length; 1884 | } else { 1885 | ok_inflater_error(inflater, "ok_inflater_set_input was called with unread input data."); 1886 | } 1887 | } 1888 | } 1889 | 1890 | size_t ok_inflater_inflate(ok_inflater *inflater, uint8_t *dst, size_t dst_len) { 1891 | if (!inflater || inflater->state == OK_INFLATER_STATE_ERROR || 1892 | inflater->state == OK_INFLATER_STATE_DONE) { 1893 | return OK_SIZE_MAX; 1894 | } 1895 | 1896 | // Each state function should return false if input is needed or the buffer is full. 1897 | // Run until one condition occurs: 1898 | // 1. Output buffer can be filled, 1899 | // 2. Internal buffer is full, 1900 | // 3. Needs more input, 1901 | // 4. Done inflating, or 1902 | // 5. An error occured. 1903 | while (ok_inflater_can_flush_total(inflater) < dst_len && 1904 | (*OK_INFLATER_STATE_FUNCTIONS[inflater->state])(inflater)) { 1905 | } 1906 | return ok_inflater_flush(inflater, dst, dst_len); 1907 | } 1908 | -------------------------------------------------------------------------------- /Part3-ok-file-formats/ok_png.h: -------------------------------------------------------------------------------- 1 | /* 2 | ok-file-formats 3 | https://github.com/brackeen/ok-file-formats 4 | Copyright (c) 2014-2020 David Brackeen 5 | 6 | Permission is hereby granted, free of charge, to any person obtaining a copy of this software and 7 | associated documentation files (the "Software"), to deal in the Software without restriction, 8 | including without limitation the rights to use, copy, modify, merge, publish, distribute, 9 | sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all copies or 13 | substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT 16 | NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND 17 | NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, 18 | DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT 19 | OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 20 | */ 21 | 22 | #ifndef OK_PNG_H 23 | #define OK_PNG_H 24 | 25 | /** 26 | * @file 27 | * Functions to read any PNG file. 28 | * 29 | * This PNG decoder: 30 | * - Supports all PNG color types and bit depths. 31 | * - Supports Apple's proprietary PNG extensions for iOS. 32 | * - Options to premultiply alpha and flip data vertically. 33 | * - Option to get image dimensions without decoding. 34 | * - Returns data in RGBA or BGRA format. 35 | * 36 | * Caveats: 37 | * - No gamma conversion. 38 | * - No CRC or ADLER32 checks. 39 | * - Ignores all chunks not related to the image. 40 | * 41 | * Example: 42 | * 43 | * #include 44 | * #include "ok_png.h" 45 | * 46 | * int main() { 47 | * FILE *file = fopen("my_image.png", "rb"); 48 | * ok_png image = ok_png_read(file, OK_PNG_COLOR_FORMAT_RGBA); 49 | * fclose(file); 50 | * if (image.data) { 51 | * printf("Got image! Size: %li x %li\n", (long)image.width, (long)image.height); 52 | * free(image.data); 53 | * } 54 | * return 0; 55 | * } 56 | */ 57 | 58 | #include 59 | #include 60 | #ifndef OK_NO_STDIO 61 | #include 62 | #endif 63 | 64 | #ifdef __cplusplus 65 | extern "C" { 66 | #endif 67 | 68 | typedef enum { 69 | OK_PNG_SUCCESS = 0, 70 | OK_PNG_ERROR_API, // Invalid argument sent to public API function 71 | OK_PNG_ERROR_INVALID, // Not a valid PNG file 72 | OK_PNG_ERROR_INFLATER, // Couldn't inflate data 73 | OK_PNG_ERROR_UNSUPPORTED, // Unsupported PNG file (width > 1073741824) 74 | OK_PNG_ERROR_ALLOCATION, // Couldn't allocate memory 75 | OK_PNG_ERROR_IO, // Couldn't read or seek the file 76 | } ok_png_error; 77 | 78 | /** 79 | * The data returned from #ok_png_read(). 80 | */ 81 | typedef struct { 82 | uint32_t width; 83 | uint32_t height; 84 | uint32_t stride; 85 | uint8_t bpp; // Always 4 86 | bool has_alpha; 87 | ok_png_error error_code:16; 88 | uint8_t *data; 89 | } ok_png; 90 | 91 | /** 92 | * PNG decode flags. Use `OK_PNG_COLOR_FORMAT_RGBA` for the most cases. 93 | */ 94 | typedef enum { 95 | /// Set for RGBA color format. This is the default format for standard PNG files. 96 | OK_PNG_COLOR_FORMAT_RGBA = (0 << 0), 97 | /// Set for BGRA color format. 98 | OK_PNG_COLOR_FORMAT_BGRA = (1 << 0), 99 | /// Set for premultiplied alpha. The default format on iOS devices is 100 | /// `(OK_PNG_COLOR_FORMAT_RGBA | OK_PNG_PREMULTIPLIED_ALPHA)` 101 | OK_PNG_PREMULTIPLIED_ALPHA = (1 << 1), 102 | /// Set to flip the image data along the horizontal axis, so that the first row of data is 103 | /// the last row in the image. 104 | OK_PNG_FLIP_Y = (1 << 2), 105 | /// Set to read an image's dimensions and color format without reading the image data. 106 | OK_PNG_INFO_ONLY = (1 << 3) 107 | } ok_png_decode_flags; 108 | 109 | // MARK: Reading from a FILE 110 | 111 | #if !defined(OK_NO_STDIO) && !defined(OK_NO_DEFAULT_ALLOCATOR) 112 | 113 | /** 114 | * Reads a PNG image using the default "stdlib" allocator. 115 | * On success, #ok_png.data contains the packed image data, with a size of 116 | * (`width * height * 4`). On failure, #ok_png.data is `NULL` and #ok_png.error_code is nonzero. 117 | * 118 | * The returned `data` must be freed by the caller (using stdlib's `free()`). 119 | * 120 | * @param file The file to read. 121 | * @param decode_flags The PNG decode flags. Use `OK_PNG_COLOR_FORMAT_RGBA` for the most cases. 122 | * @return a #ok_png object. 123 | */ 124 | ok_png ok_png_read(FILE *file, ok_png_decode_flags decode_flags); 125 | 126 | #endif 127 | 128 | // MARK: Reading from a FILE, using a custom allocator 129 | 130 | typedef struct { 131 | /** 132 | * Allocates uninitilized memory. 133 | * 134 | * @param user_data The pointer passed to #ok_png_read_with_allocator. 135 | * @param size The size of the memory to allocate. 136 | * @return the pointer to the newly allocated memory, or `NULL` if the memory could not be allocated. 137 | */ 138 | void *(*alloc)(void *user_data, size_t size); 139 | 140 | /** 141 | * Frees memory previously allocated with `alloc`. 142 | * 143 | * @param user_data The pointer passed to #ok_png_read_with_allocator. 144 | * @param memory The memory to free. 145 | */ 146 | void (*free)(void *user_data, void *memory); 147 | 148 | /** 149 | * Allocates memory for the decoded image. 150 | * This function may be `NULL`, in which case `alloc` is used instead. 151 | * 152 | * @param user_data The pointer passed to #ok_png_read_with_allocator. 153 | * @param width The image's width, in pixels. 154 | * @param height The image's height, in pixels. 155 | * @param bpp The image's number of bytes per pixel. 156 | * @param out_buffer The buffer to output data. The buffer must have a minimum size of 157 | * (`out_stride * height`) bytes. Set to `NULL` if the memory could not be allocated. 158 | * @param out_stride The stride of the buffer, in bytes. 159 | * By default, `out_stride` is `width * bpp`, but can be changed if needed. 160 | * The value must be greater than or equal to `width * bpp` or an error will occur. 161 | */ 162 | void (*image_alloc)(void *user_data, uint32_t width, uint32_t height, uint8_t bpp, 163 | uint8_t **out_buffer, uint32_t *out_stride); 164 | } ok_png_allocator; 165 | 166 | #if !defined(OK_NO_DEFAULT_ALLOCATOR) 167 | 168 | /// The default allocator using stdlib's `malloc` and `free`. 169 | extern const ok_png_allocator OK_PNG_DEFAULT_ALLOCATOR; 170 | 171 | #endif 172 | 173 | #if !defined(OK_NO_STDIO) 174 | 175 | /** 176 | * Reads a PNG image using a custom allocator. 177 | * On success, #ok_png.data contains the packed image data, with a size of 178 | * (`width * height * 4`). On failure, #ok_png.data is `NULL` and #ok_png.error_code is nonzero. 179 | * 180 | * The returned `data` must be freed by the caller. 181 | * 182 | * @param file The file to read. 183 | * @param decode_flags The PNG decode flags. Use `OK_PNG_COLOR_FORMAT_RGBA` for the most cases. 184 | * @param allocator The allocator to use. 185 | * @param allocator_user_data The pointer to pass to the allocator functions. 186 | * If using `OK_PNG_DEFAULT_ALLOCATOR`, this value should be `NULL`. 187 | * @return a #ok_png object. 188 | */ 189 | ok_png ok_png_read_with_allocator(FILE *file, ok_png_decode_flags decode_flags, 190 | ok_png_allocator allocator, void *allocator_user_data); 191 | 192 | #endif 193 | 194 | // MARK: Reading from custom input 195 | 196 | typedef struct { 197 | /** 198 | * Reads bytes from its source (typically `user_data`), copying the data to `buffer`. 199 | * 200 | * @param user_data The parameter that was passed to the #ok_png_read_from_input() 201 | * @param buffer The data buffer to copy bytes to. 202 | * @param count The number of bytes to read. 203 | * @return The number of bytes read. 204 | */ 205 | size_t (*read)(void *user_data, uint8_t *buffer, size_t count); 206 | 207 | /** 208 | * Skips bytes from its source (typically `user_data`). 209 | * 210 | * @param user_data The parameter that was passed to the #ok_png_read_from_input(). 211 | * @param count The number of bytes to skip. 212 | * @return `true` if success. 213 | */ 214 | bool (*seek)(void *user_data, long count); 215 | } ok_png_input; 216 | 217 | /** 218 | * Reads a PNG image. On success, #ok_png.data contains the packed image data, with a size of 219 | * (`width * height * 4`). On failure, #ok_png.data is `NULL` and #ok_png.error_code is nonzero. 220 | * 221 | * The returned `data` must be freed by the caller. 222 | * 223 | * @param decode_flags The PNG decode flags. Use `OK_PNG_COLOR_FORMAT_RGBA` for the most cases. 224 | * @param input_callbacks The custom input functions. 225 | * @param input_callbacks_user_data The parameter to be passed to the input's `read` and `seek` functions. 226 | * @param allocator The allocator to use. 227 | * @param allocator_user_data The pointer to pass to the allocator functions. 228 | * If using `OK_PNG_DEFAULT_ALLOCATOR`, this value should be `NULL`. 229 | * @return a #ok_png object. 230 | */ 231 | ok_png ok_png_read_from_input(ok_png_decode_flags decode_flags, 232 | ok_png_input input_callbacks, void *input_callbacks_user_data, 233 | ok_png_allocator allocator, void *allocator_user_data); 234 | 235 | // MARK: Inflater 236 | 237 | typedef struct ok_inflater ok_inflater; 238 | 239 | /** 240 | * Creates an inflater. Returns `NULL` on error. The inflater should be freed with 241 | * #ok_inflater_free(). 242 | * 243 | * @param nowrap If `true`, the inflater assumes there is no zlib wrapper around the data. 244 | */ 245 | ok_inflater *ok_inflater_init(bool nowrap, ok_png_allocator allocator, void *allocator_user_data); 246 | 247 | /** 248 | * Resets the inflater to work with a new stream. 249 | */ 250 | void ok_inflater_reset(ok_inflater *inflater); 251 | 252 | /** 253 | * Returns true if the inflater needs more input. 254 | */ 255 | bool ok_inflater_needs_input(const ok_inflater *inflater); 256 | 257 | /** 258 | * Sets the input for the inflater. Only call this function if #ok_inflater_needs_input() returns 259 | * `true`, otherwise, an error may occur. 260 | * 261 | * @param inflater The inflater. 262 | * @param buffer The input buffer. 263 | * @param buffer_length The length of the input buffer. 264 | */ 265 | void ok_inflater_set_input(ok_inflater *inflater, const uint8_t *buffer, size_t buffer_length); 266 | 267 | /** 268 | * Inflates at most `dst_length` bytes. Returns the number of bytes inflated, or `SIZE_MAX` 269 | * if an error occured. 270 | * 271 | * @param inflater The inflater. 272 | * @param dst The destination buffer to inflate bytes to. 273 | * @param dst_length The maximum number of bytes to inflate. 274 | */ 275 | size_t ok_inflater_inflate(ok_inflater *inflater, uint8_t *dst, size_t dst_length); 276 | 277 | /** 278 | * Frees the inflater. 279 | */ 280 | void ok_inflater_free(ok_inflater *inflater); 281 | 282 | #ifdef __cplusplus 283 | } 284 | #endif 285 | 286 | #endif 287 | -------------------------------------------------------------------------------- /Part3-ok-file-formats/process_of_fuzz.md: -------------------------------------------------------------------------------- 1 | # 一个完整Fuzz过程 —— ok-file-formats 2 | 3 | ## 前言 4 | 本篇是之前写的《AFL笔记》（[上](https://bbs.pediy.com/thread-257399.htm)、[下](https://bbs.pediy.com/thread-259982.htm)）的后续，去年开始断断续续学习模糊测试写了两篇侧重分析Fuzz代码的笔记，这篇主要写一下这半年（很水的半年）的Fuzz过程总结的一些经验，用Git上的一个例子做解释：[ok_file_formats](https://github.com/brackeen/ok-file-formats)。有说的不对的地方，欢迎斧正，因为都是个人总结的一家之言，大家尽管吐槽就好。 5 | 6 | P.S.吐槽一下看雪的编辑器，没有找到草稿箱功能，之前写了一半，今天终于结束了实验课想补完，打开发现 GG，什么都没有了，只能重新整理思路从头再写一遍了。 7 | 8 | ## 一、模糊测试 9 | 什么是模糊测试技术呢，就是用大量的随机输入去测试软件的健壮性等，最早是软件测试的概念。后来出现了afl，便出现了以覆盖率为导向的模糊测试技术。 10 | 11 | ![AFL-覆盖率为导向的模糊测试](img/AFL流程.png) 12 | 13 | 看图可以很好的理解，最开始给afl初始种子，然后经过变异得到输入，将输入给到待测试程序，并通过插桩技术对程序跟踪，或者覆盖率，个人认为现在这个覆盖率已经变成了一个大的概念，可以是块覆盖、可以是边覆盖、可以是很多其他有利于程序输入的信息，只要是有利于变异和种子筛选的都可以算，然后覆盖率指导之前的流程，周而复始，直到遇到崩溃情况，便将其保存。 14 | 这种覆盖率的想法使得输入不再是随机盲目的，可以说afl的出现使得模糊测试进入了新的阶段，虽然afl在2.52b版本之后就停止更新了，但是Google目前还一直维护着，而且后来出现了很多分支。今天用到的模糊测试工具就是AFL，其实其他工具使用思路也大同小异，都可以借鉴。 15 | 16 | ## 二、Fuzz过程图 17 | 要想走完一个完整的Fuzz过程，不是仅仅知道上面说的原理就可以的做到的，还是要从整体的流程上进行说明才行。 18 | 19 | ![完整的Fuzz过程图](img/Fuzz流程.png) 20 | 21 | 整个流程是从确定fuzz目标开始，到单个crash分析结束，主要包含三部分： 22 | 1. 浅绿色：主线流程，一般情况下确定好fuzz目标之后，就挑选合适到fuzzer，然后进行fuzz，这个fuzz的过程可长可短，视情况而定。最后可能一个crash都没有，也可能有好多crash，甚至会出现上百个，虽然fuzzer会说是unique crash，实际上很多情况下会出现很多触发相同漏洞的crash，这时候就需要进行一定的处理，得到真正的unique crash，然后再对单个结果进行分析。 23 | 2. 粉红色：辅助流程，当确定了fuzz目标之后需要分析源代码，然后编写两段程序，分别是给fuzzer用的程序和验证漏洞用的内存检测程序。 24 | 3. 黄色：小工具类的使用，比如afl-whatsup、afl-plot、afl-cmin、afl-tmin、afl多线程等等。 25 | 26 | ## 三、选择目标程序 27 | ### 3.1 类型 28 | 通常情况下要进行fuzz之前，应该先了解你的目标是什么样子的，代码量大还是小，开源还是闭源，属于文件读取类型还是协议类型，等等等等。下面列了一个表格是我所接触范围内的一个分类，当然随着研究的继续，这个表格会继续扩充。 29 | 30 | ![类型](img/分类.png) 31 | 32 | >1). 工业类型 33 |  1.1). 大型公开库：主要是指那些调试安装都很麻烦的库，这类库的 fuzz 难点有三： 34 |  第一、漏洞少，这些库虽然都是开源的，但是因为用到的人很多，基本上都维护好多年了，很完善了已经，剩下的漏洞不多了，简单的洞肯定都被人挖出来了，所以跑出漏洞的概率有点点看运气； 35 |  第二、调试麻烦，随着项目的维护，很多公开库已经不再是当初单纯的单一功能，而且安装、使用都需要一定的学习时间成本，如果本身对这一类又不是很熟悉的情况下，浪费的时间会很多； 36 |  第三、体量大，因为项目涉及函数超级多，所以即使得到了crash，在分析crash是否可用的时候也会非常麻烦，有可能这个crash是因为引用的别的库的旧版本导致的。 37 | 因此不管是从哪方面讲，大型公开库的Fuzz都是很有难度的，但是这也意味着一旦发现漏洞，那将是很宝贵的。下面举两个不同类型的例子的坑： 38 |   1.1.1). 开源底层库 - OpenCV，刚过年的时候我尝试过用libfuzzer对其进行模糊测试，可以说这个库完全满足上面三点（笑哭），刚开始在编译阶段就把我卡了好久，再加上那时候也不太懂，就瞎搞，换了好几个docker容器，最后好不容易搭好了环境，开开心心的跑读取函数，最后有几个crash，一跟踪发现是个其他库的，而且还不是漏洞，就很气。。。 39 |   1.1.2). 使用量大 - ffmpeg，有一类库虽然体积不大，但是用户基数超级大，ffmpeg就是属于这一类，毕竟微信传视频用的都是这个库，但是这一类的库就是上面说的第一条，漏洞极少，被这么多用户进行了检验，要想跑出洞，真的是超级难。 40 |  1.2). 小型公开库：与大型相对，就是指那类关注量小，使用量少的库。这一类库一般fuzz难度小，好编译，而且写测试程序也好写。我一般找这些库的方式有两种： 41 |   1.2.1). 公开托管平台比如github、dockerhub上面的公开库或者软件，比如ok_file_mormats就是我等下要举的例子，就是来自GitHub。比如我要fuzz图像类型的库，就搜索png、jpg、img等等这类关键词，然后按照star多少的顺序排序，然后假设是想要C语言的库就在左侧选中只要C； 42 |   1.2.1). 再有就是网站公开的一些库，这一类有自己的补丁版本，有小公司或者小团队维护，而且漏洞上报也不方便，我一般不选这一类，有这精力还不如搞大的； 43 |  1.3). 公开软件：这类我目前接触的不多，之后应该可能看看研究研究，我主要分了三类软件：一是针对操作系统，比如Linux、unix；二是针对常用系统软件，比如vim、gedit；三是通用应用软件，比如浏览器火狐、Chrome（还推荐一个方式，可以查Chrome用了哪些库，去测试那些库准没错）。 44 | 2). 学术类型 45 |  2.1). 性能测试，一般采用统一数据集，像lava-m（基本上近几年的Fuzzing论文都跑这个库做性能测试），为了证明在普遍认同的同一情况下自己的fuzzer比别人好； 46 |  2.2). 统一对比，一般是用公开库，这时候一般用真实的库来对比，一般就放自己效果好的库，不考虑版本问题，不一定会用最新版本的库； 47 |  2.3). 为了CVE，一般是为了跑出真实的漏洞，挑选的多是上面工业类型里的小型库，有的时候一些顶会文章会涉及到大型库，不过感觉今年fuzzing的申请CVE开始变得好难了，不好搞呀； 48 | 49 | 最后挑选的是一个文件转换库，其实严格意义上也不是个库了，算是个提供很多文件转换标准的函数汇总吧。第一是因为这个比较简单，内容少，好分析；第二是因为作者自己就做了一个test的文件夹，放了很多testcase，就免去了我们自己搜集了。 50 | 51 | ![ok file formats](img/ok_file_formats.png) 52 | 53 | ### 3.2 种子准备 54 | 当经过挑选之后找到了合适的目标程序，那么紧接着下一步一定是找初始种子了，这个种子的好坏其实说重要还是蛮重要的，如果fuzz的目标比较小，像这个ok-file-formats，可能用afl自带的图片testcase就够了，不过这个作者提供了一些之前能触发漏洞的testcase，可以说是非常有帮助了。 55 | 当然并不是所以的软件作者都这么有良心，最好的办法就是积累，在fuzz的时候有同类型文件触发漏洞的可以考虑留下，给其他同类型库用，遇到合适的留下，久而久之就有不小的testcase库啦。 56 | 57 | ## 四、挑选合适的Fuzzer 58 | 这一步我单独写一个章节是因为这里很重要，虽然说我们一直讨论的是用AFL相关，但是实际上在此基础上有很多很多的分支，我才疏学浅，这里就列一小部分（有兴趣可以找这三四年顶会的综述论文看一下，一定会有收获的） 59 | 60 | ![AFL分支](img/AFL分支.png) 61 | 62 | 关于AFL的分支，我这里分为四类（并不一定是afl转来了，但一定是受其思想影响的进步）： 63 | #### 移植 64 | 1. afl多是用在Linux的c语言类程序文件模糊测试，所以就有人做了平台移植：最经典的就是winafl，在Windows平台下的模糊测试工具； 65 | 2. 还有针对语言的移植，比如针对rust语言的/afl.rs、针对python语言的python-afl、针对Java语言的java-afl、针对C#语言的sharpfuzz； 66 | #### 性能的改进 67 | 1. 这部分是指速度性能的改进，既有针对有源码的libfuzzer，通过用llvm、clang框架大大提升了性能，这部分Google自家的fuzzer做的算是很好了，还有honggfuzz，都是对afl的很大改进； 68 | 2. 也有针对无源码的，像Qemu、PIN、Dyninst这些工具的引入，尽管说afl自带的系统里有qemu等但是效率不理想，所以就有一些别的大牛单独做的新的分支，甚至还有跟Frida结合的应用； 69 | #### 提升效率 70 | 这部分每年都会看见论文的新花样，针对变异策略的、针对种子筛选的层出不穷，afl-smart、ijon等等，不管是加入符号执行还是对策略改进，实际上对afl给改动并不大。主要还是从理论上进行了设想，提出了新思路； 71 | #### 重构 72 | 1. 基本上跟原来对afl差别很大了，一种是对整个变异策略全部重建，比如Angora就是用rust语言重新写对，并且用了自己新对变异策略； 73 | 2. 另一种是加入新元素，比如Neuzz这个就是加入了机器学习，再比如今年出现了一些对蓝牙的、对use设备对模糊测试； 74 | 75 | 总结下来就是一二类属于实际应用类，比较注重运用，怎么用的舒服，用得快怎么来，三四类属于理论类，比较注意想法。这篇文章的fuzz目标是文件类型的，而且是个小库，就用afl就够了，算是抛砖引玉了。 76 | 77 | ## 五、开始进行Fuzzing 78 | 前面都算是准备工作，接下来是真正的写代码的部分，先把代码clone下来： 79 | `git clone https://github.com/brackeen/ok-file-formats.git` 80 | 查看readme和GitHub上浏览一下结构，分析一下，可以发现两个比较有意思的点： 81 | 1. 项目总的readme不仅给出了项目介绍，还给出了一个简单的例子，这个例子就是我们的突破口，极大减少了学习时间，接下来的程序编写就可以借鉴这里的例子。 82 | ![exampel](img/example.png) 83 | 2. 另外一个惊喜就是，在test文件夹下的readme，给出了关于如何fuzz的介绍，以及这些testcase怎么用，哪里来的等等。 84 | ![test-readme](img/test-readme.png) 85 | 就在我写这篇文章的时候发现，作者把我之前提交issue描述中有用的部分给加到描述里面了，有点小开心。 86 | 87 | ### 5.1 编写Fuzz程序 88 |  接下来言归正传，现在知道了一个简单的例子，首先就来编写Fuzz用的程序，编写的这个程序的目的就是为了把要测试的函数加进去。我们照着葫芦画瓢，写一个读jpg文件的程序： 89 | ```C 90 | #include 91 | #include 92 | #include "ok_jpg.h" 93 | #include "ok_jpg.c" 94 | 95 | int main(int _argc, char **_argv) { 96 | //把第一个参数当作文件名，这个名就是个相对路径，不过在fuzz的时候可以用@@来代替 97 | FILE *file = fopen(_argv[1], "rb"); 98 | //读取jpg格式图片 99 | ok_jpg image = ok_jpg_read(file, OK_JPG_COLOR_FORMAT_RGBA); 100 | fclose(file); 101 | if (image.data) { 102 | printf("Got image! Size: %li x %li\n", (long)image.width, (long)image.height); 103 | free(image.data); 104 | } 105 | return 0; 106 | } 107 | ``` 108 | 把作者给的例子简单改一下，原来用的是绝对的文件名，这里改为传入的第一个参数，关于这个`_argv`,[0]、[1]的区别有点像shell里的$0、$1，传入的第一个参数是`_argv[1]`，而不是`_argv[0]`。然后下一行的读取，换成了在 ok_jpg.h 里的 ok_jpg_read 函数。 109 | 编写完之后开始编译代码，这里项目很小，就不整那些花里胡哨的CMake了（以后有机会再专门分享一篇关于fuzz的时候编译的总结），直接用 afl-gcc 进行编译： 110 | ``` 111 | afl-gcc -g -o fuzz_jpg main.c ok_jpg.c ok_jpg.h 112 | ``` 113 | 编译完成之后，就得到将要进行fuzz的目标程序 fuzz_jpg。 114 | 115 | ### 5.2 用 AFL 进行Fuzz 116 | 117 | 开始fuzz之前，可以先测试一下程序是否正常： 118 | ``` 119 | ./fuzz_jpg jpg/2001-stargate.jpg 120 | ``` 121 | ![测试程序](img/test-program.png) 122 | 使用命令开始fuzz： 123 | ``` 124 | afl-fuzz -m none -t 1000 -i jpg -o out ./fuzz_jpg @@ 125 | ``` 126 | ![模糊测试](img/fuzz过程.png) 127 | 前面也说到了，第一个参数可以用@@来代替，@@所代表传入的是变异后的输入的文件相对路径。 128 | 129 | 关于AFL的使用我就不详细说了，这方面的资料太多了，如果没有用过，可以先看一下这篇，试着走一遍流程比较好：https://www.cnblogs.com/wayne-tao/p/11739420.html ，我在这篇文章后面写了一些对于报错的处理，如果还遇到什么奇奇怪怪的问题可以在评论区说一下，大家探讨探讨。 130 | 131 | ### 5.3 性能优化 132 | 在Fuzz的过程中会遇到一些效率问题，有的则可以通过改进来解决。 133 | #### 种子筛选 134 | 像5.2的图里所示，我用到的测试用例是jpg文件夹下的图片，这些图片都是从这个项目的test文件夹下找的，下面我们可以用 afl-cmin 和 afl-tmin 工具分别对这些测试用例实现“测试集精简化”、“测试用例最小化”，命令如下（在跟afl-fuzz一样的目录下）： 135 | ```bash 136 | #cmin 137 | afl-cmin -i jpg -o newjpg ./fuzz_jpg @@ 138 | #处理完成 139 | mv jpg cmin_jpg 140 | mv newjpg jpg 141 | #tmin 142 | afl-tmin -d 1 -i jpg -o newjpg ./fuzz_jpg @@ 143 | #处理完成 144 | mv jpg tmin_jpg 145 | mv newjpg jpg 146 | ``` 147 | 关于cmin和tmin的使用如果不太懂可以看一下我的这一篇文章：https://www.cnblogs.com/wayne-tao/p/11889718.html ，他们的使用做了详细解释。 148 | 149 | >注：我这里的用了一个新的参数 -d，因为原版只有针对单个文件的，我稍微改了下，可以对文件夹处理。 150 | 151 | #### 多线程 152 | afl 自带了多线程的方式，通过设置主进程（-M）和从属进程（-S）来设置多个进程一起Fuzz，参数后面加名称（这个名称会在-o输出文件夹里单独用这个作为文件夹命令）： 153 | ```bash 154 | afl-fuzz -m none -t 500 -i jpg -o out -M fuzzM ./fuzz_jpg @@ 155 | ``` 156 | 因为服务器的核比较多，我一般是先写一个简单的脚本，然后再整，这样快一点点 157 | ![多进程](img/多进程.png) 158 | 159 | ### 5.4 查看Fuzz状态 160 | 进入Fuzz阶段之后，看到crash就停吗？万一没有看到crash呢？那么多进程要一个一个检查吗？ 161 | 肯定不是的，首先要确定一个fuzz过程是不是该停止了最直观的办法是看右上角的cycle done，不是看数字，而是看颜色，当变成绿色的时候基本上就该停了，没多少新路径可以给你用了，这时候不管有没有crash其实再继续意义也不大了，不能说不可能有crash，只是几率极低了。 162 | 163 | ![cycledone](img/cycledone.png) 164 | 165 | #### afl-whatsup 166 | 开了很多进程，比如30个，即使是用了screen，也不能一下一下的点这看吧，所以用afl-whatsup是最方便的方式，可以看到实时的fuzz进展。 167 | ```bash 168 | afl-whatsup out_dir 169 | #这里的out_dir是-o参数后面的，而不是-M或者-S后面的 170 | ``` 171 | 172 | #### afl-plot 173 | 还有一个工具afl-plot，论文上会经常看到一些图标，而且很多论文图表格式都一致，其实是afl-plot工具制作的，他会制作一个网页，下面带着三张图，就像这样 174 | 175 | ![afl-plot工具](img/afl-plot.png) 176 | 177 | 可以看到在26号凌晨五六点的时候其实我关了的话也可以，后面基本上没什么执行效率的波动了，所以这个形式比看颜色应该是稍微精准那么一点点。 178 | 179 | ## 六、结果分析 180 | 到这里开始Fuzz了，也说了一下在这过程中的一些小细节，那么如果跑出了crash该怎么处理呢，这一步就比较重要了，这一块算是自己摸索出来的一些小经验，希望可以帮到大家。 181 | ### 6.1 验证程序 182 | 首先，拿到crash文件之后（就是那个id开头的那一串），要想知道这个crash好不好，有没有触发漏洞需要验证一下，怎么验证呢，这时候就需要一个叫做Asan的工具（Address Sanitizer），目前好像是Google在管理，而且已经集成到gcc自带了。是一个内存检测的工具，可以把程序异常很好的展现出来。 183 | 同样是之前编写的fuzz程序的那段代码，我们这一次不用afl-gcc编译了，用gcc编译，因为目的不同，这次是为了验证漏洞，所以加上Asan的选项： 184 | ```bash 185 | gcc -g -fsanitize=address -fno-omit-frame-pointer -O1 -o Asanjpg main.c ok_jpg.c ok_jpg.h 186 | ``` 187 | 这时候就得到了一个内存检测的验证程序了：Asanjpg，验证程序地址：https://github.com/WayneDevMaze/afl_test_works/blob/master/OKfileformat/Asanjpg/Asanjpg 。 188 | 189 | 一般情况下就是这么个套路，验证程序跟Fuzz测试程序在核心函数上没有区别，区别就是编译过程的不同，有时候也会在文件读取方式那里稍微做一下改动。 190 | 191 | ### 6.2 crash处理 192 | 当得到验证程序之后，就可以开始处理crash了，crash其实就是 out_dir/master/crash 下的文件，以id开头，有编号，我们把这个文件copy出来之后，放到Asanjpg同文件夹下（我这里就用已经验证过的了，自己做测试的时候只需要改文件名即可，[crash/poc地址](https://github.com/WayneDevMaze/afl_test_works/blob/master/OKfileformat/Asanjpg/crash/jpg-heap-buffer-overflow-1)），并运行： 193 | ```bash 194 | ./Asanjpg crash/jpg-heap-buffer-overflow-1 195 | ``` 196 | 可以看到内存检测的输出： 197 | ```c 198 | ==98287==ERROR: AddressSanitizer: heap-buffer-overflow on address 0x631000039680 at pc 0x562394639b54 bp 0x7ffee24654e0 sp 0x7ffee24654d0 199 | READ of size 2 at 0x631000039680 thread T0 200 | #0 0x562394639b53 in ok_jpg_decode_block_subsequent_scan /root/study/ok-file-formats/afl-test/ok_jpg.c:1102 201 | #1 0x56239463b11f in ok_jpg_decode_scan /root/study/ok-file-formats/afl-test/ok_jpg.c:1238 202 | #2 0x56239463fc60 in ok_jpg_read_sos /root/study/ok-file-formats/afl-test/ok_jpg.c:1734 203 | #3 0x562394640d3c in ok_jpg_decode2 /root/study/ok-file-formats/afl-test/ok_jpg.c:1900 204 | #4 0x562394641605 in ok_jpg_decode /root/study/ok-file-formats/afl-test/ok_jpg.c:1990 205 | #5 0x5623946308a4 in ok_jpg_read_with_allocator /root/study/ok-file-formats/afl-test/ok_jpg.c:268 206 | #6 0x56239463071b in ok_jpg_read /root/study/ok-file-formats/afl-test/ok_jpg.c:257 207 | #7 0x56239462fd5e in main /root/study/ok-file-formats/afl-test/main.c:8 208 | #8 0x7fe63f9a4b96 in __libc_start_main (/lib/x86_64-linux-gnu/libc.so.6+0x21b96) 209 | #9 0x56239462fb29 in _start (/root/study/ok-file-formats/afl-test/Asanjpg/Asanjpg+0x2b29) 210 | 211 | 0x631000039680 is located 113 bytes to the right of 69135-byte region [0x631000028800,0x63100003960f) 212 | allocated by thread T0 here: 213 | #0 0x7fe63fe52b40 in __interceptor_malloc (/usr/lib/x86_64-linux-gnu/libasan.so.4+0xdeb40) 214 | #1 0x56239462ff00 in ok_stdlib_alloc /root/study/ok-file-formats/afl-test/ok_jpg.c:55 215 | #2 0x56239463eb20 in ok_jpg_read_sof /root/study/ok-file-formats/afl-test/ok_jpg.c:1595 216 | #3 0x562394640ac2 in ok_jpg_decode2 /root/study/ok-file-formats/afl-test/ok_jpg.c:1884 217 | #4 0x562394641605 in ok_jpg_decode /root/study/ok-file-formats/afl-test/ok_jpg.c:1990 218 | #5 0x5623946308a4 in ok_jpg_read_with_allocator /root/study/ok-file-formats/afl-test/ok_jpg.c:268 219 | #6 0x56239463071b in ok_jpg_read /root/study/ok-file-formats/afl-test/ok_jpg.c:257 220 | #7 0x56239462fd5e in main /root/study/ok-file-formats/afl-test/main.c:8 221 | #8 0x7fe63f9a4b96 in __libc_start_main (/lib/x86_64-linux-gnu/libc.so.6+0x21b96) 222 | 223 | SUMMARY: AddressSanitizer: heap-buffer-overflow /root/study/ok-file-formats/afl-test/ok_jpg.c:1102 in ok_jpg_decode_block_subsequent_scan 224 | Shadow bytes around the buggy address: 225 | 0x0c627ffff280: 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 226 | 0x0c627ffff290: 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 227 | 0x0c627ffff2a0: 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 228 | 0x0c627ffff2b0: 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 229 | 0x0c627ffff2c0: 00 07 fa fa fa fa fa fa fa fa fa fa fa fa fa fa 230 | =>0x0c627ffff2d0:[fa]fa fa fa fa fa fa fa fa fa fa fa fa fa fa fa 231 | 0x0c627ffff2e0: fa fa fa fa fa fa fa fa fa fa fa fa fa fa fa fa 232 | 0x0c627ffff2f0: fa fa fa fa fa fa fa fa fa fa fa fa fa fa fa fa 233 | 0x0c627ffff300: fa fa fa fa fa fa fa fa fa fa fa fa fa fa fa fa 234 | 0x0c627ffff310: fa fa fa fa fa fa fa fa fa fa fa fa fa fa fa fa 235 | 0x0c627ffff320: fa fa fa fa fa fa fa fa fa fa fa fa fa fa fa fa 236 | Shadow byte legend (one shadow byte represents 8 application bytes): 237 | Addressable: 00 238 | Partially addressable: 01 02 03 04 05 06 07 239 | Heap left redzone: fa 240 | Freed heap region: fd 241 | Stack left redzone: f1 242 | Stack mid redzone: f2 243 | Stack right redzone: f3 244 | Stack after return: f5 245 | Stack use after scope: f8 246 | Global redzone: f9 247 | Global init order: f6 248 | Poisoned by user: f7 249 | Container overflow: fc 250 | Array cookie: ac 251 | Intra object redzone: bb 252 | ASan internal: fe 253 | Left alloca redzone: ca 254 | Right alloca redzone: cb 255 | ==98287==ABORTING 256 | ``` 257 | 258 | 第一、可以从SUMMARY看出这个crash是不是有价值，有时候报一些莫名其妙的没有价值的错还需要用gdb进一步调试，当然我一般懒得调试（好吧其实是我技术不到位分析不出啥东西来）； 259 | 第二、在第一段中给出了触发位置的详细调用过程，如果想写exp，这一块比较重要，我现在不是很了解这部分，不敢说太多； 260 | 第三、并不是所有的crash都可以报这种错误，有时候就是仅仅是文件类型不匹配等等问题，这时候crash其实就没什么用； 261 | 第四、**重复crash**，这里着重说一下，我现在在做的一个小工具（正在验证中），请看截图： 262 | 263 | ![crashes](img/crashes.png) 264 | 265 | 可以看到20个线程，最后生成了2262个crash，经过前十几个的验证，发现重复性特别高，虽然crash文件内容不一样，但是触发的漏洞位置一模一样，这种crash其实并不unique，而且很多crash并不能触发真正的漏洞，是假crash，所以最近在做一个筛选出真正独一无二crash的工具。 266 | 267 | ### 6.3 漏洞提交issue 268 | 当得到6.2中的Asan报告之后，就可以认为触发了堆溢出漏洞，可以在GitHub上给作者提交issue了，最后贴一下我提交issue的大概内容（加注释） 269 | [issue地址](https://github.com/brackeen/ok-file-formats/issues/7) 270 | 271 | > 272 | ># Describe 【漏洞描述】 273 | >A heap-buffer-overflow was discovered in ok_file_formats. The issue is being triggered in function ok_jpg_decode_block_subsequent_scan() at ok_jpg.c:1102 274 | >**这部分放漏洞的总述，把漏洞类型、漏洞触发函数位置、漏洞所在文件列出来；** 275 | ># Reproduce 【复现过程所需内容】 276 | >**test program** 277 | >```C 278 | >#include 279 | >int main(int _argc, char **_argv) {} 280 | >``` 281 | >**我一般是把验证程序放在这里，好让作者可以自己reproduce-复现；** 282 | >Tested in Ubuntu 18.04, 64bit. 283 | >Compile test program with address sanitizer with this command: 284 | >```bash 285 | >gcc -g -fsanitize=address -fno-omit-frame-pointer -O1 -o Asanjpg main.c ok_jpg.c ok_jpg.h 286 | >``` 287 | >**告诉作者使用的，以及编译的环境和方法；** 288 | >You can get program [here](https://github.com/WayneDevMaze/afl_test_works/blob/master/OKfileformat/Asanjpg/Asanjpg). 289 | >**为了方便作者验证漏洞，我也会把编译好的程序放出来** 290 | ># ASan Reports 【验证程序报的内存报告】 291 | >```bash 292 | >./Asanjpg crash/jpg-heap-buffer-overflow-1 293 | >``` 294 | >Get ASan reports 295 | >```bash 296 | >==98287==ABORTING 297 | >``` 298 | >**这里就放Asanjpg验证之后报的内容，就是6.2里报的内存错误跟踪信息** 299 | ># Poc 【其实就是crash文件】 300 | >Poc file is [here](https://github.com/WayneDevMaze/afl_test_works/blob/master/OKfileformat/Asanjpg/crash/jpg-heap-buffer-overflow-1). 301 | >**最后一定要放出poc文件，总结下来一定要有的：漏洞描述、复现验证程序及源代码、报错信息（不一定非得是ASan的结果，也可以是gdb调试信息）、POC文件（触发漏洞的crash文件）** 302 | >   303 | 304 | 当然如果Fuzz的程序比较知名，可以尝试申请CVE，我这个尝试申请了，但是没有回信，应该是凉了。 305 | 306 | ## 结语 307 | 总体来说，整个Fuzz过程熟悉一遍之后就变得不那么复杂了，还记得刚过年的时候我还是一头雾水，现在对小型库的fuzz已经比较有信心了，希望大家看完本文能有所帮助。 308 | 309 | [GitHub-AFL源码阅读笔记](https://github.com/WayneDevMaze/Chinese_noted_AFL)文中说的想要做的工具，完成后应该会放在这，欢迎star； 310 | [Fuzz学习过程随笔](https://www.cnblogs.com/wayne-tao/tag/AFL学习/)记录自己在学习过程中遇到的坑和一些心得； -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # 模糊测试初探 2 | 按照三部分对模糊测试进行介绍和实践： 3 | 4 | ---- 5 | 6 | ## 一、模糊测试简介及AFL入门（1.5h） 7 | 1.1 模糊测试 8 | 1.2 目前主流的模糊测试 9 | 1.3 AFL介绍 10 | 1.4 小例子初探 11 | ## 二、AFL详解（1.5h） 12 | 2.1 AFL代码模块介绍 13 | 2.2 AFL周边工具 14 | 15 | ---- 16 | 17 | ## 三、库实战测试 —— ok-file-formats（3h） 18 | 3.1 实战过程分析（0.5h） 19 | 3.2 程序分析（0.5h） 20 | 3.3 模糊测试（0.5h） 21 | 3.4 结果分析（0.5h） 22 | 3.5 学员实操及问题解答（1h） 23 | 24 | ---- 25 | 26 | ## 准备工作： 27 | Linux环境下学习，最好是Ubuntu（版本无所谓） 28 | 1. afl下载安装，参考[第一次测试](https://www.cnblogs.com/wayne-tao/p/11739420.html) 29 | 2. ok-file-formats下载，https://github.com/brackeen/ok-file-formats 30 | 31 | ## 扩展阅读： 32 | 1. AFL源码解读，https://github.com/WayneDevMaze/Chinese_noted_AFL 33 | 2. 相关修改学习笔记，https://www.cnblogs.com/wayne-tao/tag/AFL学习/ 34 | 3. 完整的fuzz流程，https://bbs.pediy.com/thread-261253.htm 35 | -------------------------------------------------------------------------------- /img/AFL分支.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WayneDevMaze/Fuzzing-Introduction/528d82a492d88941f9b4d4731569f21b5fc90533/img/AFL分支.png -------------------------------------------------------------------------------- /img/AFL流程.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WayneDevMaze/Fuzzing-Introduction/528d82a492d88941f9b4d4731569f21b5fc90533/img/AFL流程.png -------------------------------------------------------------------------------- /夏令营培训大纲v2-张涛-模糊测试初探.docx: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/WayneDevMaze/Fuzzing-Introduction/528d82a492d88941f9b4d4731569f21b5fc90533/夏令营培训大纲v2-张涛-模糊测试初探.docx --------------------------------------------------------------------------------