zstd_wrapper.c 3.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144
  1. // SPDX-License-Identifier: GPL-2.0-or-later
  2. /*
  3. * Squashfs - a compressed read only filesystem for Linux
  4. *
  5. * Copyright (c) 2016-present, Facebook, Inc.
  6. * All rights reserved.
  7. *
  8. * zstd_wrapper.c
  9. */
  10. #include <linux/mutex.h>
  11. #include <linux/bio.h>
  12. #include <linux/slab.h>
  13. #include <linux/zstd.h>
  14. #include <linux/vmalloc.h>
  15. #include "squashfs_fs.h"
  16. #include "squashfs_fs_sb.h"
  17. #include "squashfs.h"
  18. #include "decompressor.h"
  19. #include "page_actor.h"
  20. struct workspace {
  21. void *mem;
  22. size_t mem_size;
  23. size_t window_size;
  24. };
  25. static void *zstd_init(struct squashfs_sb_info *msblk, void *buff)
  26. {
  27. struct workspace *wksp = kmalloc(sizeof(*wksp), GFP_KERNEL);
  28. if (wksp == NULL)
  29. goto failed;
  30. wksp->window_size = max_t(size_t,
  31. msblk->block_size, SQUASHFS_METADATA_SIZE);
  32. wksp->mem_size = ZSTD_DStreamWorkspaceBound(wksp->window_size);
  33. wksp->mem = vmalloc(wksp->mem_size);
  34. if (wksp->mem == NULL)
  35. goto failed;
  36. return wksp;
  37. failed:
  38. ERROR("Failed to allocate zstd workspace\n");
  39. kfree(wksp);
  40. return ERR_PTR(-ENOMEM);
  41. }
  42. static void zstd_free(void *strm)
  43. {
  44. struct workspace *wksp = strm;
  45. if (wksp)
  46. vfree(wksp->mem);
  47. kfree(wksp);
  48. }
  49. static int zstd_uncompress(struct squashfs_sb_info *msblk, void *strm,
  50. struct bio *bio, int offset, int length,
  51. struct squashfs_page_actor *output)
  52. {
  53. struct workspace *wksp = strm;
  54. ZSTD_DStream *stream;
  55. size_t total_out = 0;
  56. int error = 0;
  57. ZSTD_inBuffer in_buf = { NULL, 0, 0 };
  58. ZSTD_outBuffer out_buf = { NULL, 0, 0 };
  59. struct bvec_iter_all iter_all = {};
  60. struct bio_vec *bvec = bvec_init_iter_all(&iter_all);
  61. stream = ZSTD_initDStream(wksp->window_size, wksp->mem, wksp->mem_size);
  62. if (!stream) {
  63. ERROR("Failed to initialize zstd decompressor\n");
  64. return -EIO;
  65. }
  66. out_buf.size = PAGE_SIZE;
  67. out_buf.dst = squashfs_first_page(output);
  68. for (;;) {
  69. size_t zstd_err;
  70. if (in_buf.pos == in_buf.size) {
  71. const void *data;
  72. int avail;
  73. if (!bio_next_segment(bio, &iter_all)) {
  74. error = -EIO;
  75. break;
  76. }
  77. avail = min(length, ((int)bvec->bv_len) - offset);
  78. data = page_address(bvec->bv_page) + bvec->bv_offset;
  79. length -= avail;
  80. in_buf.src = data + offset;
  81. in_buf.size = avail;
  82. in_buf.pos = 0;
  83. offset = 0;
  84. }
  85. if (out_buf.pos == out_buf.size) {
  86. out_buf.dst = squashfs_next_page(output);
  87. if (out_buf.dst == NULL) {
  88. /* Shouldn't run out of pages
  89. * before stream is done.
  90. */
  91. error = -EIO;
  92. break;
  93. }
  94. out_buf.pos = 0;
  95. out_buf.size = PAGE_SIZE;
  96. }
  97. total_out -= out_buf.pos;
  98. zstd_err = ZSTD_decompressStream(stream, &out_buf, &in_buf);
  99. total_out += out_buf.pos; /* add the additional data produced */
  100. if (zstd_err == 0)
  101. break;
  102. if (ZSTD_isError(zstd_err)) {
  103. ERROR("zstd decompression error: %d\n",
  104. (int)ZSTD_getErrorCode(zstd_err));
  105. error = -EIO;
  106. break;
  107. }
  108. }
  109. squashfs_finish_page(output);
  110. return error ? error : total_out;
  111. }
  112. const struct squashfs_decompressor squashfs_zstd_comp_ops = {
  113. .init = zstd_init,
  114. .free = zstd_free,
  115. .decompress = zstd_uncompress,
  116. .id = ZSTD_COMPRESSION,
  117. .name = "zstd",
  118. .supported = 1
  119. };