std.compress.zstandard: fix crashes

2 files changed, 10 insertions, 6 deletions

diff --git a/lib/std/compress/zstandard/decode/block.zig b/lib/std/compress/zstandard/decode/block.zig
index ab1d476dd6..8f97bea399 100644
--- a/lib/std/compress/zstandard/decode/block.zig
+++ b/lib/std/compress/zstandard/decode/block.zig
@@ -981,6 +981,8 @@ fn decodeStreams(size_format: u2, stream_data: []const u8) !LiteralsSection.Stre
     const stream_3_start = stream_2_start + stream_2_length;
     const stream_4_start = stream_3_start + stream_3_length;
 
+    if (stream_data.len < stream_4_start) return error.MalformedLiteralsSection;
+
     return .{ .four = .{
         stream_data[stream_1_start .. stream_1_start + stream_1_length],
         stream_data[stream_2_start .. stream_2_start + stream_2_length],
diff --git a/lib/std/compress/zstandard/decode/huffman.zig b/lib/std/compress/zstandard/decode/huffman.zig
index 01913c7044..f5639e7721 100644
--- a/lib/std/compress/zstandard/decode/huffman.zig
+++ b/lib/std/compress/zstandard/decode/huffman.zig
@@ -59,7 +59,7 @@ fn assignWeights(huff_bits: *readers.ReverseBitReader, accuracy_log: usize, entr
     var even_state: u32 = huff_bits.readBitsNoEof(u32, accuracy_log) catch return error.MalformedHuffmanTree;
     var odd_state: u32 = huff_bits.readBitsNoEof(u32, accuracy_log) catch return error.MalformedHuffmanTree;
 
-    while (i < 255) {
+    while (i < 254) {
         const even_data = entries[even_state];
         var read_bits: usize = 0;
         const even_bits = huff_bits.readBits(u32, even_data.bits, &read_bits) catch unreachable;
@@ -78,7 +78,7 @@ fn assignWeights(huff_bits: *readers.ReverseBitReader, accuracy_log: usize, entr
         weights[i] = std.math.cast(u4, odd_data.symbol) orelse return error.MalformedHuffmanTree;
         i += 1;
         if (read_bits < odd_data.bits) {
-            if (i == 256) return error.MalformedHuffmanTree;
+            if (i == 255) return error.MalformedHuffmanTree;
             weights[i] = std.math.cast(u4, entries[even_state].symbol) orelse return error.MalformedHuffmanTree;
             i += 1;
             break;
@@ -147,16 +147,18 @@ fn assignSymbols(weight_sorted_prefixed_symbols: []LiteralsSection.HuffmanTree.P
 }
 
 fn buildHuffmanTree(weights: *[256]u4, symbol_count: usize) error{MalformedHuffmanTree}!LiteralsSection.HuffmanTree {
-    var weight_power_sum: u16 = 0;
+    var weight_power_sum_big: u32 = 0;
     for (weights[0 .. symbol_count - 1]) |value| {
         if (value > 0) {
-            weight_power_sum += @as(u16, 1) << (value - 1);
+            weight_power_sum_big += @as(u16, 1) << (value - 1);
         }
     }
-    if (weight_power_sum >= 1 << 11) return error.MalformedHuffmanTree;
+    if (weight_power_sum_big >= 1 << 11) return error.MalformedHuffmanTree;
+    const weight_power_sum = @intCast(u16, weight_power_sum_big);
 
     // advance to next power of two (even if weight_power_sum is a power of 2)
-    const max_number_of_bits = std.math.log2_int(u16, weight_power_sum) + 1;
+    // TODO: is it valid to have weight_power_sum == 0?
+    const max_number_of_bits = if (weight_power_sum == 0) 1 else std.math.log2_int(u16, weight_power_sum) + 1;
     const next_power_of_two = @as(u16, 1) << max_number_of_bits;
     weights[symbol_count - 1] = std.math.log2_int(u16, next_power_of_two - weight_power_sum) + 1;