Fix incorrect response streaming leftover bytes calculation (#63)

Author: csvance

src/Curl.jl

@@ -17,31 +17,35 @@ function write_callback(
     try
         req = unsafe_pointer_to_objref(req_p)::gRPCRequest
 
+        !isnothing(req.ex) && return typemax(Csize_t)
+
         n = size * count
         buf = unsafe_wrap(Array, convert(Ptr{UInt8}, data), (n,))
 
         handled_n_bytes_total = 0
         try 
             while !isnothing(buf) && handled_n_bytes_total < n
                 handled_n_bytes, buf = handle_write(req, buf)
-                handled_n_bytes_total += handled_n_bytes 
+                handled_n_bytes_total += handled_n_bytes
                 handled_n_bytes == 0 && break
             end
-        catch ex 
+        catch ex
             # Eat InvalidStateException raised on put! to closed channel
             !isa(ex, InvalidStateException) && rethrow(ex)
         end
 
-        # If there was an exception handle it 
         !isnothing(req.ex) && return typemax(Csize_t)
 
         # Check that we handled the correct number of bytes
         # If there was no exception in handle_write this should always match 
         if handled_n_bytes_total != n
-            req.ex = gRPCServiceCallException(
-                GRPC_INTERNAL,
-                "Recieved $(n) bytes from curl but only handled $(handled_n_bytes_total)",
-            )
+            if isnothing(req.ex)
+                req.ex = gRPCServiceCallException(
+                    GRPC_INTERNAL,
+                    "Recieved $(n) bytes from curl but only handled $(handled_n_bytes_total)",
+                )
+            end
+
             # If we are response streaming unblock the task waiting on response_c
             !isnothing(req.response_c) && close(req.response_c)
             return typemax(Csize_t)
@@ -337,7 +341,6 @@ function handle_write(req::gRPCRequest, buf::Vector{UInt8})::Tuple{Int64, Union{
             # Not enough data yet to read the entire header
             return write(req.response, buf), nothing
         else
-
             buf_header = buf[1:header_bytes_left]
             n = write(req.response, buf_header)
 
@@ -405,11 +408,11 @@ function handle_write(req::gRPCRequest, buf::Vector{UInt8})::Tuple{Int64, Union{
         req.response_length = 0
 
         # Handle the remaining data
-        leftover_bytes = message_bytes_left - n
+        leftover_bytes = length(buf) - n
 
         buf_leftover = nothing 
         if leftover_bytes > 0
-            buf_leftover = unsafe_wrap(Array, pointer(buf) + n, (length(leftover_bytes),))
+            buf_leftover = unsafe_wrap(Array, pointer(buf) + n, (leftover_bytes,))
         end
 
         return n, buf_leftover

test/runtests.jl

@@ -119,16 +119,6 @@ include("gen/test/test_pb.jl")
         end
     # end
 
-    # @testset "Max Message Size" begin 
-        # Create a client with much more restictive max message lengths
-        client = TestService_TestRPC_Client(_TEST_HOST, _TEST_PORT; max_send_message_length=1024, max_recieve_message_length=1024)
-
-        # Send too much 
-        @test_throws gRPCServiceCallException grpc_sync_request(client, TestRequest(1, zeros(UInt64, 1024)))
-        # Receive too much
-        @test_throws gRPCServiceCallException grpc_sync_request(client, TestRequest(1024, zeros(UInt64, 1)))
-    # end
-
     # @testset "Async Channels" begin 
         client = TestService_TestRPC_Client(_TEST_HOST, _TEST_PORT)
 
@@ -144,7 +134,7 @@ include("gen/test/test_pb.jl")
         end
     # end
 
-    @static if VERSION >= v"1.12"
+@static if VERSION >= v"1.12"
     # @testset "Response Streaming" begin
         N = 1000
 
@@ -203,6 +193,7 @@ include("gen/test/test_pb.jl")
             @test last(response.data) == i
         end
 
+
         close(request_c)
         grpc_async_await(req)
     # end
@@ -346,7 +337,7 @@ include("gen/test/test_pb.jl")
             @test isa(ex, gRPCServiceCallException)
         end
     # end
-    end
+end
 
     # @testset "Timeout Header Value Formatting" begin
         # Test integer seconds
@@ -379,5 +370,15 @@ include("gen/test/test_pb.jl")
         end
     # end
 
+    # @testset "Max Message Size" begin
+        # Create a client with much more restictive max message lengths
+        client = TestService_TestRPC_Client(_TEST_HOST, _TEST_PORT; max_send_message_length=1024, max_recieve_message_length=1024)
+
+        # Send too much
+        @test_throws gRPCServiceCallException grpc_sync_request(client, TestRequest(1, zeros(UInt64, 1024)))
+        # Receive too much
+        @test_throws gRPCServiceCallException grpc_sync_request(client, TestRequest(1024, zeros(UInt64, 1)))
+    # end
+
     grpc_shutdown()
 end