zstd decompress error

core/runtime/common/uncompress_code_if_needed.cpp

@@ -7,12 +7,15 @@
 #include "runtime/common/uncompress_code_if_needed.hpp"
 
 #include <zstd.h>
+#include <zstd_errors.h>
 
 OUTCOME_CPP_DEFINE_CATEGORY(kagome::runtime, UncompressError, e) {
   using E = kagome::runtime::UncompressError;
   switch (e) {
     case E::ZSTD_ERROR:
       return "WASM code not compressed by zstd!";
+    case E::BOMB_SIZE_REACHED:
+      return "Code decompression failed. Maximum size reached - possible bomb";
   }
   return "Unknown error";
 }
@@ -30,25 +33,26 @@ namespace kagome::runtime {
 
   outcome::result<void> uncompressCodeIfNeeded(common::BufferView buf,
                                                common::Buffer &res) {
-    if (buf.size() > kZstdPrefixSize
-        && std::equal(buf.begin(),
-                      buf.begin() + kZstdPrefixSize,
-                      std::begin(kZstdPrefix),
-                      std::end(kZstdPrefix))) {
+    if (startsWith(buf, kZstdPrefix)) {
+      auto zstd = buf.subspan(std::size(kZstdPrefix));
       // here we can check that blob is really ZSTD compressed
       // but we don't use the result size, it's unknown for the WASM blob
       // @see ZSTD_CONTENTSIZE_UNKNOWN
-      auto check_size = ZSTD_getFrameContentSize(buf.data() + kZstdPrefixSize,
-                                                 buf.size() - kZstdPrefixSize);
+      auto check_size = ZSTD_getFrameContentSize(zstd.data(), zstd.size());
       if (check_size == ZSTD_CONTENTSIZE_ERROR) {
         return UncompressError::ZSTD_ERROR;
       }
       res.resize(kCodeBlobBombLimit);
-      auto size = ZSTD_decompress(res.data(),
-                                  kCodeBlobBombLimit,
-                                  buf.data() + kZstdPrefixSize,
-                                  buf.size() - kZstdPrefixSize);
+      auto size = ZSTD_decompress(
+          res.data(), kCodeBlobBombLimit, zstd.data(), zstd.size());
+      if (ZSTD_isError(size)) {
+        if (ZSTD_getErrorCode(size) == ZSTD_error_dstSize_tooSmall) {
+          return UncompressError::BOMB_SIZE_REACHED;
+        }
+        return UncompressError::ZSTD_ERROR;
+      }
       res.resize(size);
+      res.shrink_to_fit();
     } else {
       res = common::Buffer{buf};
     }

core/runtime/common/uncompress_code_if_needed.hpp

@@ -9,7 +9,10 @@
 #include "common/buffer.hpp"
 
 namespace kagome::runtime {
-  enum class UncompressError : uint8_t { ZSTD_ERROR };
+  enum class UncompressError : uint8_t {
+    ZSTD_ERROR,
+    BOMB_SIZE_REACHED,
+  };
 
   outcome::result<void> uncompressCodeIfNeeded(common::BufferView buf,
                                                common::Buffer &res);