Add tests

Author: nhz2

src/compression.jl

@@ -98,7 +98,7 @@ function TranscodingStreams.startproc(codec::ZstdCompressor, mode::Symbol, error
         codec.cstream.ptr = ptr
         i_code = initialize!(codec.cstream, codec.level)
         if iserror(i_code)
-            error[] = ErrorException("zstd error")
+            error[] = ErrorException("zstd initialization error")
             return :error
         end
     end

src/decompression.jl

@@ -33,16 +33,6 @@ end
 # Methods
 # -------
 
-function TranscodingStreams.initialize(codec::ZstdDecompressor)
-    code = initialize!(codec.dstream)
-    if iserror(code)
-        zstderror(codec.dstream, code)
-    end
-    reset!(codec.dstream.ibuffer)
-    reset!(codec.dstream.obuffer)
-    return
-end
-
 function TranscodingStreams.finalize(codec::ZstdDecompressor)
     if codec.dstream.ptr != C_NULL
         code = free!(codec.dstream)
@@ -51,12 +41,22 @@ function TranscodingStreams.finalize(codec::ZstdDecompressor)
         end
         codec.dstream.ptr = C_NULL
     end
-    reset!(codec.dstream.ibuffer)
-    reset!(codec.dstream.obuffer)
-    return
+    nothing
 end
 
 function TranscodingStreams.startproc(codec::ZstdDecompressor, mode::Symbol, error::Error)
+    if codec.dstream.ptr == C_NULL
+        ptr = LibZstd.ZSTD_createDStream()
+        if ptr == C_NULL
+            throw(OutOfMemoryError())
+        end
+        codec.dstream.ptr = ptr
+        i_code = initialize!(codec.dstream)
+        if iserror(i_code)
+            error[] = ErrorException("zstd initialization error")
+            return :error
+        end
+    end
     code = reset!(codec.dstream)
     if iserror(code)
         error[] = ErrorException("zstd error")
@@ -66,6 +66,9 @@ function TranscodingStreams.startproc(codec::ZstdDecompressor, mode::Symbol, err
 end
 
 function TranscodingStreams.process(codec::ZstdDecompressor, input::Memory, output::Memory, error::Error)
+    if codec.dstream.ptr == C_NULL
+        error("startproc must be called before process")
+    end
     dstream = codec.dstream
     dstream.ibuffer.src = input.ptr
     dstream.ibuffer.size = input.size

src/libzstd.jl

@@ -123,11 +123,7 @@ mutable struct DStream
     obuffer::OutBuffer
 
     function DStream()
-        ptr = LibZstd.ZSTD_createDStream()
-        if ptr == C_NULL
-            throw(OutOfMemoryError())
-        end
-        return new(ptr, InBuffer(), OutBuffer())
+        return new(C_NULL, InBuffer(), OutBuffer())
     end
 end
 Base.unsafe_convert(::Type{Ptr{LibZstd.ZSTD_DStream}}, dstream::DStream) = dstream.ptr

test/runtests.jl

@@ -158,4 +158,72 @@ include("utils.jl")
 
     include("compress_endOp.jl")
     include("static_only_tests.jl")
+
+    @testset "reusing a compressor" begin
+        compressor = ZstdCompressor()
+        x = rand(UInt8, 1000)
+        TranscodingStreams.initialize(compressor)
+        ret1 = transcode(compressor, x)
+        TranscodingStreams.finalize(compressor)
+
+        # compress again using the same compressor
+        TranscodingStreams.initialize(compressor) # segfault happens here!
+        ret2 = transcode(compressor, x)
+        ret3 = transcode(compressor, x)
+        TranscodingStreams.finalize(compressor)
+
+        @test transcode(ZstdDecompressor, ret1) == x
+        @test transcode(ZstdDecompressor, ret2) == x
+        @test transcode(ZstdDecompressor, ret3) == x
+        @test ret1 == ret2
+        @test ret1 == ret3
+
+        decompressor = ZstdDecompressor()
+        TranscodingStreams.initialize(decompressor)
+        @test transcode(decompressor, ret1) == x
+        TranscodingStreams.finalize(decompressor)
+
+        TranscodingStreams.initialize(decompressor)
+        @test transcode(decompressor, ret1) == x
+        TranscodingStreams.finalize(decompressor)
+    end
+
+    @testset "use after free doesn't segfault" begin
+        @testset "$(Codec)" for Codec in (ZstdCompressor, ZstdDecompressor)
+            codec = Codec()
+            TranscodingStreams.initialize(codec)
+            TranscodingStreams.finalize(codec)
+            data = [0x00,0x01]
+            GC.@preserve data let m = TranscodingStreams.Memory(pointer(data), length(data))
+                try
+                    TranscodingStreams.expectedsize(codec, m)
+                catch
+                end
+                try
+                    TranscodingStreams.minoutsize(codec, m)
+                catch
+                end
+                try
+                    TranscodingStreams.initialize(codec)
+                catch
+                end
+                try
+                    TranscodingStreams.process(codec, m, m, TranscodingStreams.Error())
+                catch
+                end
+                try
+                    TranscodingStreams.startproc(codec, :read, TranscodingStreams.Error())
+                catch
+                end
+                try
+                    TranscodingStreams.process(codec, m, m, TranscodingStreams.Error())
+                catch
+                end
+                try
+                    TranscodingStreams.finalize(codec)
+                catch
+                end
+            end
+        end
+    end
 end