From 769db97b20e047d403f22e38f4b0d779227aaca8 Mon Sep 17 00:00:00 2001
From: Mark Kittisopikul <markkitt@gmail.com>
Date: Sun, 2 Jun 2024 01:19:59 -0400
Subject: [PATCH 1/2] Only reset session on error in process, override
 transcode with pledge

---
 src/compression.jl | 27 ++++++++++++++++++++++-----
 src/libzstd.jl     |  6 ++++++
 2 files changed, 28 insertions(+), 5 deletions(-)

diff --git a/src/compression.jl b/src/compression.jl
index cabc3f9..25062bf 100644
--- a/src/compression.jl
+++ b/src/compression.jl
@@ -77,6 +77,25 @@ end
 
 # Methods
 # -------
+function TranscodingStreams.transcode(
+    codec::ZstdCompressor,
+    input::TranscodingStreams.Buffer,
+    output::Union{TranscodingStreams.Buffer,Nothing} = nothing,
+)
+    code = LibZstd.ZSTD_CCtx_setPledgedSrcSize(codec.cstream, TranscodingStreams.buffersize(input))
+    if iserror(code)
+        zstderror(codec.cstream, code)
+    end
+    return invoke(
+        TranscodingStreams.transcode,
+        Tuple{
+              TranscodingStreams.Codec,
+              TranscodingStreams.Buffer,
+              typeof(output)
+        },
+        codec, input, output
+    )
+end
 
 function TranscodingStreams.initialize(codec::ZstdCompressor)
     code = initialize!(codec.cstream, codec.level)
@@ -102,11 +121,8 @@ function TranscodingStreams.finalize(codec::ZstdCompressor)
 end
 
 function TranscodingStreams.startproc(codec::ZstdCompressor, mode::Symbol, error::Error)
-    code = reset!(codec.cstream, 0 #=unknown source size=#)
-    if iserror(code)
-        error[] = ErrorException("zstd error")
-        return :error
-    end
+    reset!(codec.cstream.ibuffer)
+    reset!(codec.cstream.obuffer)
     return :ok
 end
 
@@ -146,6 +162,7 @@ function TranscodingStreams.process(codec::ZstdCompressor, input::Memory, output
     if iserror(code)
         ptr = LibZstd.ZSTD_getErrorName(code)
         error[] = ErrorException("zstd error: " * unsafe_string(ptr))
+        reset!(cstream)
         return Δin, Δout, :error
     else
         return Δin, Δout, input.size == 0 && code == 0 ? :end : :ok
diff --git a/src/libzstd.jl b/src/libzstd.jl
index 79d021d..2fb07ca 100644
--- a/src/libzstd.jl
+++ b/src/libzstd.jl
@@ -60,6 +60,12 @@ function initialize!(cstream::CStream, level::Integer)
     return LibZstd.ZSTD_initCStream(cstream, level)
 end
 
+function reset!(cstream::CStream)
+    # Resetting session never fails.
+    # Also this will reset the pledged size
+    return LibZstd.ZSTD_CCtx_reset(cstream, LibZstd.ZSTD_reset_session_only)
+end
+
 function reset!(cstream::CStream, srcsize::Integer)
     # ZSTD_resetCStream is deprecated
     # https://github.com/facebook/zstd/blob/9d2a45a705e22ad4817b41442949cd0f78597154/lib/zstd.h#L2253-L2272

From 874e49c5be866470c6878b8676b49cf8f69a026e Mon Sep 17 00:00:00 2001
From: Mark Kittisopikul <markkitt@gmail.com>
Date: Sun, 2 Jun 2024 02:06:15 -0400
Subject: [PATCH 2/2] Add test and find_decompressed_size(::Vector{UInt8})

---
 src/libzstd.jl   | 5 +++++
 test/runtests.jl | 1 +
 2 files changed, 6 insertions(+)

diff --git a/src/libzstd.jl b/src/libzstd.jl
index 2fb07ca..84f7fd4 100644
--- a/src/libzstd.jl
+++ b/src/libzstd.jl
@@ -169,6 +169,11 @@ end
 const ZSTD_CONTENTSIZE_UNKNOWN = Culonglong(0) - 1
 const ZSTD_CONTENTSIZE_ERROR   = Culonglong(0) - 2
 
+# ZSTD_findDecompressedSize gets the decompressed size of all available frames
+# Alternatively, consider ZSTD_getFrameContentSize for a single frame
 function find_decompressed_size(src::Ptr, size::Integer)
     return LibZstd.ZSTD_findDecompressedSize(src, size)
 end
+function find_decompressed_size(src::Vector{UInt8})
+    return LibZstd.ZSTD_findDecompressedSize(src, length(src))
+end
diff --git a/test/runtests.jl b/test/runtests.jl
index 7ca5875..0aef654 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -32,6 +32,7 @@ Random.seed!(1234)
             for n in 0:L-1
                 @test_throws ErrorException transcode(ZstdDecompressor, compressed_data[1:n])
             end
+            @test CodecZstd.find_decompressed_size(compressed_data) == length(uncompressed_data)
             @test transcode(ZstdDecompressor, compressed_data) == uncompressed_data
         end
     end