decompress_unzstd.c 10 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350
  1. // SPDX-License-Identifier: GPL-2.0
  2. /*
  3. * Important notes about in-place decompression
  4. *
  5. * At least on x86, the kernel is decompressed in place: the compressed data
  6. * is placed to the end of the output buffer, and the decompressor overwrites
  7. * most of the compressed data. There must be enough safety margin to
  8. * guarantee that the write position is always behind the read position.
  9. *
  10. * The safety margin for ZSTD with a 128 KB block size is calculated below.
  11. * Note that the margin with ZSTD is bigger than with GZIP or XZ!
  12. *
  13. * The worst case for in-place decompression is that the beginning of
  14. * the file is compressed extremely well, and the rest of the file is
  15. * uncompressible. Thus, we must look for worst-case expansion when the
  16. * compressor is encoding uncompressible data.
  17. *
  18. * The structure of the .zst file in case of a compresed kernel is as follows.
  19. * Maximum sizes (as bytes) of the fields are in parenthesis.
  20. *
  21. * Frame Header: (18)
  22. * Blocks: (N)
  23. * Checksum: (4)
  24. *
  25. * The frame header and checksum overhead is at most 22 bytes.
  26. *
  27. * ZSTD stores the data in blocks. Each block has a header whose size is
  28. * a 3 bytes. After the block header, there is up to 128 KB of payload.
  29. * The maximum uncompressed size of the payload is 128 KB. The minimum
  30. * uncompressed size of the payload is never less than the payload size
  31. * (excluding the block header).
  32. *
  33. * The assumption, that the uncompressed size of the payload is never
  34. * smaller than the payload itself, is valid only when talking about
  35. * the payload as a whole. It is possible that the payload has parts where
  36. * the decompressor consumes more input than it produces output. Calculating
  37. * the worst case for this would be tricky. Instead of trying to do that,
  38. * let's simply make sure that the decompressor never overwrites any bytes
  39. * of the payload which it is currently reading.
  40. *
  41. * Now we have enough information to calculate the safety margin. We need
  42. * - 22 bytes for the .zst file format headers;
  43. * - 3 bytes per every 128 KiB of uncompressed size (one block header per
  44. * block); and
  45. * - 128 KiB (biggest possible zstd block size) to make sure that the
  46. * decompressor never overwrites anything from the block it is currently
  47. * reading.
  48. *
  49. * We get the following formula:
  50. *
  51. * safety_margin = 22 + uncompressed_size * 3 / 131072 + 131072
  52. * <= 22 + (uncompressed_size >> 15) + 131072
  53. */
  54. /*
  55. * Preboot environments #include "path/to/decompress_unzstd.c".
  56. * All of the source files we depend on must be #included.
  57. * zstd's only source dependeny is xxhash, which has no source
  58. * dependencies.
  59. *
  60. * When UNZSTD_PREBOOT is defined we declare __decompress(), which is
  61. * used for kernel decompression, instead of unzstd().
  62. *
  63. * Define __DISABLE_EXPORTS in preboot environments to prevent symbols
  64. * from xxhash and zstd from being exported by the EXPORT_SYMBOL macro.
  65. */
  66. #ifdef STATIC
  67. # define UNZSTD_PREBOOT
  68. # include "xxhash.c"
  69. # include "zstd/entropy_common.c"
  70. # include "zstd/fse_decompress.c"
  71. # include "zstd/huf_decompress.c"
  72. # include "zstd/zstd_common.c"
  73. # include "zstd/decompress.c"
  74. #endif
  75. #include <linux/decompress/mm.h>
  76. #include <linux/kernel.h>
  77. #include <linux/zstd.h>
  78. /* 128MB is the maximum window size supported by zstd. */
  79. #define ZSTD_WINDOWSIZE_MAX (1 << ZSTD_WINDOWLOG_MAX)
  80. /*
  81. * Size of the input and output buffers in multi-call mode.
  82. * Pick a larger size because it isn't used during kernel decompression,
  83. * since that is single pass, and we have to allocate a large buffer for
  84. * zstd's window anyway. The larger size speeds up initramfs decompression.
  85. */
  86. #define ZSTD_IOBUF_SIZE (1 << 17)
  87. static int INIT handle_zstd_error(size_t ret, void (*error)(char *x))
  88. {
  89. const int err = ZSTD_getErrorCode(ret);
  90. if (!ZSTD_isError(ret))
  91. return 0;
  92. switch (err) {
  93. case ZSTD_error_memory_allocation:
  94. error("ZSTD decompressor ran out of memory");
  95. break;
  96. case ZSTD_error_prefix_unknown:
  97. error("Input is not in the ZSTD format (wrong magic bytes)");
  98. break;
  99. case ZSTD_error_dstSize_tooSmall:
  100. case ZSTD_error_corruption_detected:
  101. case ZSTD_error_checksum_wrong:
  102. error("ZSTD-compressed data is corrupt");
  103. break;
  104. default:
  105. error("ZSTD-compressed data is probably corrupt");
  106. break;
  107. }
  108. return -1;
  109. }
  110. /*
  111. * Handle the case where we have the entire input and output in one segment.
  112. * We can allocate less memory (no circular buffer for the sliding window),
  113. * and avoid some memcpy() calls.
  114. */
  115. static int INIT decompress_single(const u8 *in_buf, long in_len, u8 *out_buf,
  116. long out_len, long *in_pos,
  117. void (*error)(char *x))
  118. {
  119. const size_t wksp_size = ZSTD_DCtxWorkspaceBound();
  120. void *wksp = large_malloc(wksp_size);
  121. ZSTD_DCtx *dctx = ZSTD_initDCtx(wksp, wksp_size);
  122. int err;
  123. size_t ret;
  124. if (dctx == NULL) {
  125. error("Out of memory while allocating ZSTD_DCtx");
  126. err = -1;
  127. goto out;
  128. }
  129. /*
  130. * Find out how large the frame actually is, there may be junk at
  131. * the end of the frame that ZSTD_decompressDCtx() can't handle.
  132. */
  133. ret = ZSTD_findFrameCompressedSize(in_buf, in_len);
  134. err = handle_zstd_error(ret, error);
  135. if (err)
  136. goto out;
  137. in_len = (long)ret;
  138. ret = ZSTD_decompressDCtx(dctx, out_buf, out_len, in_buf, in_len);
  139. err = handle_zstd_error(ret, error);
  140. if (err)
  141. goto out;
  142. if (in_pos != NULL)
  143. *in_pos = in_len;
  144. err = 0;
  145. out:
  146. if (wksp != NULL)
  147. large_free(wksp);
  148. return err;
  149. }
  150. static int INIT __unzstd(unsigned char *in_buf, long in_len,
  151. long (*fill)(void*, unsigned long),
  152. long (*flush)(void*, unsigned long),
  153. unsigned char *out_buf, long out_len,
  154. long *in_pos,
  155. void (*error)(char *x))
  156. {
  157. ZSTD_inBuffer in;
  158. ZSTD_outBuffer out;
  159. ZSTD_frameParams params;
  160. void *in_allocated = NULL;
  161. void *out_allocated = NULL;
  162. void *wksp = NULL;
  163. size_t wksp_size;
  164. ZSTD_DStream *dstream;
  165. int err;
  166. size_t ret;
  167. /*
  168. * ZSTD decompression code won't be happy if the buffer size is so big
  169. * that its end address overflows. When the size is not provided, make
  170. * it as big as possible without having the end address overflow.
  171. */
  172. if (out_len == 0)
  173. out_len = UINTPTR_MAX - (uintptr_t)out_buf;
  174. if (fill == NULL && flush == NULL)
  175. /*
  176. * We can decompress faster and with less memory when we have a
  177. * single chunk.
  178. */
  179. return decompress_single(in_buf, in_len, out_buf, out_len,
  180. in_pos, error);
  181. /*
  182. * If in_buf is not provided, we must be using fill(), so allocate
  183. * a large enough buffer. If it is provided, it must be at least
  184. * ZSTD_IOBUF_SIZE large.
  185. */
  186. if (in_buf == NULL) {
  187. in_allocated = large_malloc(ZSTD_IOBUF_SIZE);
  188. if (in_allocated == NULL) {
  189. error("Out of memory while allocating input buffer");
  190. err = -1;
  191. goto out;
  192. }
  193. in_buf = in_allocated;
  194. in_len = 0;
  195. }
  196. /* Read the first chunk, since we need to decode the frame header. */
  197. if (fill != NULL)
  198. in_len = fill(in_buf, ZSTD_IOBUF_SIZE);
  199. if (in_len < 0) {
  200. error("ZSTD-compressed data is truncated");
  201. err = -1;
  202. goto out;
  203. }
  204. /* Set the first non-empty input buffer. */
  205. in.src = in_buf;
  206. in.pos = 0;
  207. in.size = in_len;
  208. /* Allocate the output buffer if we are using flush(). */
  209. if (flush != NULL) {
  210. out_allocated = large_malloc(ZSTD_IOBUF_SIZE);
  211. if (out_allocated == NULL) {
  212. error("Out of memory while allocating output buffer");
  213. err = -1;
  214. goto out;
  215. }
  216. out_buf = out_allocated;
  217. out_len = ZSTD_IOBUF_SIZE;
  218. }
  219. /* Set the output buffer. */
  220. out.dst = out_buf;
  221. out.pos = 0;
  222. out.size = out_len;
  223. /*
  224. * We need to know the window size to allocate the ZSTD_DStream.
  225. * Since we are streaming, we need to allocate a buffer for the sliding
  226. * window. The window size varies from 1 KB to ZSTD_WINDOWSIZE_MAX
  227. * (8 MB), so it is important to use the actual value so as not to
  228. * waste memory when it is smaller.
  229. */
  230. ret = ZSTD_getFrameParams(&params, in.src, in.size);
  231. err = handle_zstd_error(ret, error);
  232. if (err)
  233. goto out;
  234. if (ret != 0) {
  235. error("ZSTD-compressed data has an incomplete frame header");
  236. err = -1;
  237. goto out;
  238. }
  239. if (params.windowSize > ZSTD_WINDOWSIZE_MAX) {
  240. error("ZSTD-compressed data has too large a window size");
  241. err = -1;
  242. goto out;
  243. }
  244. /*
  245. * Allocate the ZSTD_DStream now that we know how much memory is
  246. * required.
  247. */
  248. wksp_size = ZSTD_DStreamWorkspaceBound(params.windowSize);
  249. wksp = large_malloc(wksp_size);
  250. dstream = ZSTD_initDStream(params.windowSize, wksp, wksp_size);
  251. if (dstream == NULL) {
  252. error("Out of memory while allocating ZSTD_DStream");
  253. err = -1;
  254. goto out;
  255. }
  256. /*
  257. * Decompression loop:
  258. * Read more data if necessary (error if no more data can be read).
  259. * Call the decompression function, which returns 0 when finished.
  260. * Flush any data produced if using flush().
  261. */
  262. if (in_pos != NULL)
  263. *in_pos = 0;
  264. do {
  265. /*
  266. * If we need to reload data, either we have fill() and can
  267. * try to get more data, or we don't and the input is truncated.
  268. */
  269. if (in.pos == in.size) {
  270. if (in_pos != NULL)
  271. *in_pos += in.pos;
  272. in_len = fill ? fill(in_buf, ZSTD_IOBUF_SIZE) : -1;
  273. if (in_len < 0) {
  274. error("ZSTD-compressed data is truncated");
  275. err = -1;
  276. goto out;
  277. }
  278. in.pos = 0;
  279. in.size = in_len;
  280. }
  281. /* Returns zero when the frame is complete. */
  282. ret = ZSTD_decompressStream(dstream, &out, &in);
  283. err = handle_zstd_error(ret, error);
  284. if (err)
  285. goto out;
  286. /* Flush all of the data produced if using flush(). */
  287. if (flush != NULL && out.pos > 0) {
  288. if (out.pos != flush(out.dst, out.pos)) {
  289. error("Failed to flush()");
  290. err = -1;
  291. goto out;
  292. }
  293. out.pos = 0;
  294. }
  295. } while (ret != 0);
  296. if (in_pos != NULL)
  297. *in_pos += in.pos;
  298. err = 0;
  299. out:
  300. if (in_allocated != NULL)
  301. large_free(in_allocated);
  302. if (out_allocated != NULL)
  303. large_free(out_allocated);
  304. if (wksp != NULL)
  305. large_free(wksp);
  306. return err;
  307. }
  308. #ifndef UNZSTD_PREBOOT
  309. STATIC int INIT unzstd(unsigned char *buf, long len,
  310. long (*fill)(void*, unsigned long),
  311. long (*flush)(void*, unsigned long),
  312. unsigned char *out_buf,
  313. long *pos,
  314. void (*error)(char *x))
  315. {
  316. return __unzstd(buf, len, fill, flush, out_buf, 0, pos, error);
  317. }
  318. #else
  319. STATIC int INIT __decompress(unsigned char *buf, long len,
  320. long (*fill)(void*, unsigned long),
  321. long (*flush)(void*, unsigned long),
  322. unsigned char *out_buf, long out_len,
  323. long *pos,
  324. void (*error)(char *x))
  325. {
  326. return __unzstd(buf, len, fill, flush, out_buf, out_len, pos, error);
  327. }
  328. #endif