Merge pull request #313 from cira0100/ldn/room-handshake

2 files changed, 28 insertions, 5 deletions

diff --git a/src/common/zstd_compression.cpp b/src/common/zstd_compression.cpp
index 7665482ee6..9bb8d7b919 100644
--- a/src/common/zstd_compression.cpp
+++ b/src/common/zstd_compression.cpp
@@ -35,16 +35,26 @@ std::vector<u8> CompressDataZSTDDefault(const u8* source, std::size_t source_siz
 }
 
 std::vector<u8> DecompressDataZSTD(std::span<const u8> compressed) {
-    const std::size_t decompressed_size =
+    return TryDecompressDataZSTD(compressed, std::vector<u8>{}.max_size())
+        .value_or(std::vector<u8>{});
+}
+
+std::optional<std::vector<u8>> TryDecompressDataZSTD(std::span<const u8> compressed,
+                                                     std::size_t max_size) {
+    // The size comes from the frame header, so check it before allocating
+    const unsigned long long decompressed_size =
         ZSTD_getFrameContentSize(compressed.data(), compressed.size());
-    std::vector<u8> decompressed(decompressed_size);
+    if (decompressed_size == ZSTD_CONTENTSIZE_ERROR ||
+        decompressed_size == ZSTD_CONTENTSIZE_UNKNOWN || decompressed_size > max_size) {
+        return std::nullopt;
+    }
+    std::vector<u8> decompressed(static_cast<std::size_t>(decompressed_size));
 
     const std::size_t uncompressed_result_size = ZSTD_decompress(
         decompressed.data(), decompressed.size(), compressed.data(), compressed.size());
 
-    if (decompressed_size != uncompressed_result_size || ZSTD_isError(uncompressed_result_size)) {
-        // Decompression failed
-        return {};
+    if (ZSTD_isError(uncompressed_result_size) || uncompressed_result_size != decompressed_size) {
+        return std::nullopt;
     }
     return decompressed;
 }
diff --git a/src/common/zstd_compression.h b/src/common/zstd_compression.h
index 819d7d201a..c20d96adb0 100644
--- a/src/common/zstd_compression.h
+++ b/src/common/zstd_compression.h
@@ -6,6 +6,7 @@
 
 #pragma once
 
+#include <optional>
 #include <span>
 #include <vector>
 
@@ -46,6 +47,18 @@ namespace Common::Compression {
 [[nodiscard]] std::vector<u8> DecompressDataZSTD(std::span<const u8> compressed);
 
 /**
+ * Decompresses a single Zstandard frame that declares its content size.
+ *
+ * @param compressed the compressed source memory region.
+ * @param max_size   the largest content size to accept.
+ *
+ * @return the decompressed data, or nullopt if the frame is invalid, does not declare its size,
+ *         or declares more than max_size bytes.
+ */
+[[nodiscard]] std::optional<std::vector<u8>> TryDecompressDataZSTD(std::span<const u8> compressed,
+                                                                   std::size_t max_size);
+
+/**
  * Decompresses a single Zstandard frame into a caller-provided buffer of known size.
  *
  * @param compressed the compressed source memory region.