From b2c4e97cba099bb421575a0ef4b3c0175b00b99f Mon Sep 17 00:00:00 2001
From: nhz2 <nhz2@cornell.edu>
Date: Sat, 5 Oct 2024 12:49:50 -0400
Subject: [PATCH] test errors

---
 src/compression.jl |  5 ++++-
 test/runtests.jl   | 54 ++++++++++++++++++++++++++++++++++++++++------
 2 files changed, 51 insertions(+), 8 deletions(-)

diff --git a/src/compression.jl b/src/compression.jl
index 11e014b..6a2b7b8 100644
--- a/src/compression.jl
+++ b/src/compression.jl
@@ -112,6 +112,9 @@ end
 
 if isdefined(TranscodingStreams, :pledgeinsize)
     function TranscodingStreams.pledgeinsize(codec::ZstdCompressor, insize::Int64, error::Error)::Symbol
+        if codec.cstream.ptr == C_NULL
+            Base.error("`startproc` must be called before `pledgeinsize`")
+        end
         srcsize = if signbit(insize)
             ZSTD_CONTENTSIZE_UNKNOWN
         else
@@ -129,7 +132,7 @@ end
 
 function TranscodingStreams.process(codec::ZstdCompressor, input::Memory, output::Memory, error::Error)
     if codec.cstream.ptr == C_NULL
-        error("startproc must be called before process")
+        Base.error("`startproc` must be called before `process`")
     end
     cstream = codec.cstream
     ibuffer_starting_pos = UInt(0)
diff --git a/test/runtests.jl b/test/runtests.jl
index c2c195d..b705864 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -168,13 +168,47 @@ include("utils.jl")
         if isdefined(TranscodingStreams, :pledgeinsize)
             # when pledgeinsize is available transcode should save the 
             # decompressed size in a header
-            for n in [0:30; 1000; 1000000; 10000000;]
-                a = transcode(ZstdCompressor, rand(UInt8, n))
-                @test CodecZstd.find_decompressed_size(a) == n
-                a = transcode(ZstdCompressor, rand(0x00:0x01, n))
-                @test CodecZstd.find_decompressed_size(a) == n
-                a = transcode(ZstdCompressor, zeros(UInt8, n))
-                @test CodecZstd.find_decompressed_size(a) == n
+            for n in [0:30; 1000; 1000000;]
+                v = transcode(ZstdCompressor, rand(UInt8, n))
+                @test CodecZstd.find_decompressed_size(v) == n
+            end
+
+            # Test what happens if pledgeinsize promise is broken
+            d1 = zeros(UInt8, 10000)
+            d2 = zeros(UInt8, 10000)
+            GC.@preserve d1 d2 begin
+                @testset "too many bytes" begin
+                    m1 = TranscodingStreams.Memory(pointer(d1), 1000)
+                    m2 = TranscodingStreams.Memory(pointer(d2), 1000)
+                    codec = ZstdCompressor()
+                    e = TranscodingStreams.Error()
+                    @test TranscodingStreams.startproc(codec, :read, e) === :ok
+                    @test TranscodingStreams.pledgeinsize(codec, 10, e) === :ok
+                    @test TranscodingStreams.process(codec, m1, m2, e) === (0, 0, :error)
+                    @test e[] == ErrorException("zstd error: Src size is incorrect")
+                end
+                @testset "too few bytes" begin
+                    m1 = TranscodingStreams.Memory(pointer(d1), 10)
+                    m2 = TranscodingStreams.Memory(pointer(d2), 1000)
+                    codec = ZstdCompressor()
+                    e = TranscodingStreams.Error()
+                    @test TranscodingStreams.startproc(codec, :read, e) === :ok
+                    @test TranscodingStreams.pledgeinsize(codec, 10000, e) === :ok
+                    @test TranscodingStreams.process(codec, m1, m2, e)[3] === :ok
+                    m1 = TranscodingStreams.Memory(pointer(d1), 0)
+                    @test TranscodingStreams.process(codec, m1, m2, e)[3] === :error
+                    @test e[] == ErrorException("zstd error: Src size is incorrect")
+                end
+                @testset "set pledgeinsize after process" begin
+                    m1 = TranscodingStreams.Memory(pointer(d1), 1000)
+                    m2 = TranscodingStreams.Memory(pointer(d2), 1000)
+                    codec = ZstdCompressor()
+                    e = TranscodingStreams.Error()
+                    @test TranscodingStreams.startproc(codec, :read, e) === :ok
+                    @test TranscodingStreams.process(codec, m1, m2, e)[3] === :ok
+                    @test TranscodingStreams.pledgeinsize(codec, 10000, e) === :error
+                    @test e[] == ErrorException("zstd error setting pledged source size")
+                end
             end
         end
     end
@@ -218,6 +252,12 @@ include("utils.jl")
             TranscodingStreams.finalize(codec)
             data = [0x00,0x01]
             GC.@preserve data let m = TranscodingStreams.Memory(pointer(data), length(data))
+                if isdefined(TranscodingStreams, :pledgeinsize)
+                    try
+                        TranscodingStreams.pledgeinsize(codec, 10, TranscodingStreams.Error())
+                    catch
+                    end
+                end
                 try
                     TranscodingStreams.expectedsize(codec, m)
                 catch