Merge pull request #6 from JuliaComputing/tan/grpc

add message size limit checks

Author: tanmaykm

README.md

@@ -140,6 +140,9 @@ gRPCController(;
     [ revocation::Bool = true, ]
     [ request_timeout::Real = Inf, ]
     [ connect_timeout::Real = 0, ]
+    [ max_message_length = DEFAULT_MAX_MESSAGE_LENGTH, ]
+    [ max_recv_message_length = 0, ]
+    [ max_send_message_length = 0, ]
     [ verbose::Bool = false, ]
 )
 ```
@@ -155,6 +158,11 @@ gRPCController(;
 - `request_timeout`: request timeout (seconds)
 - `connect_timeout`: connect timeout (seconds) (default is 300 seconds, same
    as setting this to 0)
+- `max_message_length`: maximum message length (default is 4MB)
+- `max_recv_message_length`: maximum message length to receive (default is
+   `max_message_length`, same as setting this to 0)
+- `max_send_message_length`: maximum message length to send (default is
+   `max_message_length`, same as setting this to 0)
 - `verbose`: whether to print out verbose communication logs (default false)
 
 ### `gRPCChannel`
@@ -178,15 +186,6 @@ the server.
 - `success`: whether the request was completed successfully.
 - `message`: any error message if request was not successful
 
-### `gRPCException`
-
-Every gRPC request returns the result and a future representing the status
-of the gRPC request. Use the `gRPCCheck` method on the status future to check
-the request status and throw a `gRPCException` if it is not successful.
-
-A `gRPCException` has a member named `message` that may contain an error
-message if request was not successful.
-
 ### `gRPCCheck`
 
 ```julia
@@ -196,3 +195,36 @@ gRPCCheck(status; throw_error::Bool=true)
 Method to check the response of a gRPC request and raise a `gRPCException`
 if it has failed. If `throw_error` is set to false, returns `true` or `false`
 indicating success instead.
+
+### `gRPCException`
+
+Every gRPC request returns the result and a future representing the status
+of the gRPC request. Use the `gRPCCheck` method on the status future to check
+the request status and throw a `gRPCException` if it is not successful.
+
+The abstract `gRPCException` type has the following concrete implementations:
+
+- `gRPCMessageTooLargeException`
+- `gRPCServiceCallException`
+
+### `gRPCMessageTooLargeException`
+
+A `gRPMessageTooLargeException` exception is thrown when a message is
+encountered that has a size greater than the limit configured.
+Specifically, `max_recv_message_length` while receiving  and
+`max_send_message_length` while sending.
+
+A `gRPMessageTooLargeException` has the following members:
+
+- `limit`: the limit value that was exceeded
+- `encountered`: the amount of data that was actually received
+    or sent before this error was triggered. Note that this may
+    not correspond to the full size of the data, as error may be
+    thrown before actually materializing the complete data.
+
+### `gRPCServiceCallException`
+
+A `gRPCServiceCallException` is thrown if a gRPC request is not successful.
+It has the following members:
+
+- `message`: any error message if request was not successful

src/curl.jl

@@ -18,10 +18,10 @@ function buffer_send_data(input::Channel{T}) where T <: ProtoType
 end
 =#
 
-function send_data(easy::Curl.Easy, input::Channel{T}) where T <: ProtoType
+function send_data(easy::Curl.Easy, input::Channel{T}, max_send_message_length::Int) where T <: ProtoType
     while true
         yield()
-        data = isready(input) ? to_delimited_message_bytes(take!(input)) : isopen(input) ? UInt8[] : nothing
+        data = isready(input) ? to_delimited_message_bytes(take!(input), max_send_message_length) : isopen(input) ? UInt8[] : nothing
         easy.input === nothing && break
         easy.input = data
         Curl.curl_easy_pause(easy.handle, Curl.CURLPAUSE_CONT)
@@ -63,7 +63,7 @@ function easy_handle(maxage::Clong, keepalive::Clong, negotiation::Symbol, revoc
     easy
 end
 
-function recv_data(easy::Curl.Easy, output::Channel{T}) where T <: ProtoType
+function recv_data(easy::Curl.Easy, output::Channel{T}, max_recv_message_length::Int) where T <: ProtoType
     iob = PipeBuffer()
     waiting_for_header = true
     msgsize = 0
@@ -78,6 +78,11 @@ function recv_data(easy::Curl.Easy, output::Channel{T}) where T <: ProtoType
                 if bytesavailable(iob) >= 5
                     compressed = read(iob, UInt8)       # compression
                     datalen = ntoh(read(iob, UInt32))   # message length
+
+                    if datalen > max_recv_message_length
+                        throw(gRPCMessageTooLargeException(max_recv_message_length, datalen))
+                    end
+
                     waiting_for_header = false
                 else
                     need_more = true
@@ -114,10 +119,12 @@ end
 function grpc_request(downloader::Downloader, url::String, input::Channel{T1}, output::Channel{T2};
         maxage::Clong = typemax(Clong),
         keepalive::Clong = 60,
-	negotiation::Symbol = :http2_prior_knowledge,
-	revocation::Bool = true,
+        negotiation::Symbol = :http2_prior_knowledge,
+        revocation::Bool = true,
         request_timeout::Real = Inf,
         connect_timeout::Real = 0,
+        max_recv_message_length::Int = DEFAULT_MAX_RECV_MESSAGE_LENGTH,
+        max_send_message_length::Int = DEFAULT_MAX_SEND_MESSAGE_LENGTH,
         verbose::Bool = false)::gRPCStatus where {T1 <: ProtoType, T2 <: ProtoType}
     Curl.with_handle(easy_handle(maxage, keepalive, negotiation, revocation)) do easy
         # setup the request
@@ -131,14 +138,54 @@ function grpc_request(downloader::Downloader, url::String, input::Channel{T1}, o
         # do the request
         Curl.add_handle(downloader.multi, easy)
 
-        try
-            # do send recv data
-            Base.Experimental.@sync begin
-                @async recv_data(easy, output)
-                @async send_data(easy, input)
-            end
-        finally # ensure handle is removed
+        function cleanup()
             Curl.remove_handle(downloader.multi, easy)
+            # though remove_handle sets easy.handle to C_NULL, it does not close output and progress channels
+            # we need to close them here to unblock anything waiting on them
+            close(easy.output)
+            close(easy.progress)
+            close(output)
+            close(input)
+            nothing
+        end
+
+        # do send recv data
+        if VERSION < v"1.5"
+            cleaned_up = false
+            exception = nothing
+            cleanup_once = (ex)->begin
+                if !cleaned_up
+                    cleaned_up = true
+                    exception = ex
+                    cleanup()
+                end
+            end
+
+            @sync begin
+                @async try
+                    recv_data(easy, output, max_recv_message_length)
+                catch ex
+                    cleanup_once(ex)
+                end
+                @async try
+                    send_data(easy, input, max_send_message_length)
+                catch ex
+                    cleanup_once(ex)
+                end
+            end
+
+            if exception !== nothing
+                throw(exception)
+            end
+        else
+            try
+                Base.Experimental.@sync begin
+                    @async recv_data(easy, output, max_recv_message_length)
+                    @async send_data(easy, input, max_send_message_length)
+                end
+            finally # ensure handle is removed
+                cleanup()
+            end
         end
 
         (easy.code == CURLE_OK) ? gRPCStatus(true, "") : gRPCStatus(false, Curl.get_curl_errstr(easy))

src/gRPCClient.jl

@@ -7,8 +7,11 @@ using ProtoBuf
 import Downloads: Curl
 import ProtoBuf: call_method
 
-export gRPCController, gRPCChannel, gRPCException, gRPCStatus, gRPCCheck
+export gRPCController, gRPCChannel, gRPCException, gRPCServiceCallException, gRPCMessageTooLargeException, gRPCStatus, gRPCCheck
 
+abstract type gRPCException <: Exception end
+
+include("limitio.jl")
 include("curl.jl")
 include("grpc.jl")
 include("generate.jl")

src/grpc.jl

@@ -12,39 +12,56 @@
 struct gRPCStatus
     success::Bool
     message::String
+    exception::Union{Nothing,Exception}
+end
+
+gRPCStatus(success::Bool, message::AbstractString) = gRPCStatus(success, string(message), nothing)
+function gRPCStatus(status_future)
+    try
+        fetch(status_future)
+    catch ex
+        task_exception = isa(ex, TaskFailedException) ? ex.task.exception : ex
+        while isa(task_exception, TaskFailedException)
+            task_exception = task_exception.task.exception
+        end
+        gRPCStatus(false, string(task_exception), task_exception)
+    end
 end
 
 """
-    struct gRPCException
+    struct gRPCServiceCallException
         message::String
     end
 
-Every gRPC request returns the result and a future representing the status
-of the gRPC request. Use the `gRPCCheck` method on the status future to check
-the request status and throw a `gRPCException` if it is not successful.
-
-A `gRPCException` has the following members:
+A `gRPCServiceCallException` is thrown if a gRPC request is not successful.
+It has the following members:
 
 - `message`: any error message if request was not successful
 """
-struct gRPCException <: Exception
+struct gRPCServiceCallException <: gRPCException
     message::String
 end
 
+Base.show(io::IO, m::gRPCServiceCallException) = print(io, "gRPCServiceCallException - $(m.message)")
+
 """
     gRPCCheck(status; throw_error::Bool=true)
 
-Check the response of a gRPC request and raise a `gRPCException` if it has
-failed. If `throw_error` is set to false, returns `true` or `false` indicating
-success instead.
+Every gRPC request returns the result and a future representing the status
+of the gRPC request. Check the response of a gRPC request and raise a
+`gRPCException` if it has failed. If `throw_error` is set to false, this
+returns `true` or `false` indicating success instead.
 """
-gRPCCheck(status; throw_error::Bool=true) = gRPCCheck(fetch(status); throw_error=throw_error)
+gRPCCheck(status_future; throw_error::Bool=true) = gRPCCheck(gRPCStatus(status_future); throw_error=throw_error)
 function gRPCCheck(status::gRPCStatus; throw_error::Bool=true)
-    if throw_error
-        status.success || throw(gRPCException(status.message))
-    else
-        status.success
+    if throw_error && !status.success
+        if status.exception === nothing
+            throw(gRPCServiceCallException(status.message))
+        else
+            throw(status.exception)
+        end
     end
+    status.success
 end
 
 """
@@ -55,6 +72,9 @@ end
         [ revocation::Bool = true, ]
         [ request_timeout::Real = Inf, ]
         [ connect_timeout::Real = 0, ]
+        [ max_message_length = DEFAULT_MAX_MESSAGE_LENGTH, ]
+        [ max_recv_message_length = 0, ]
+        [ max_send_message_length = 0, ]
         [ verbose::Bool = false, ]
     )
 
@@ -70,6 +90,11 @@ Contains settings to control the behavior of gRPC requests.
 - `request_timeout`: request timeout (seconds)
 - `connect_timeout`: connect timeout (seconds) (default is 300 seconds, same
    as setting this to 0)
+- `max_message_length`: maximum message length (default is 4MB)
+- `max_recv_message_length`: maximum message length to receive (default is
+   `max_message_length`, same as setting this to 0)
+- `max_send_message_length`: maximum message length to send (default is
+   `max_message_length`, same as setting this to 0)
 - `verbose`: whether to print out verbose communication logs (default false)
 """
 struct gRPCController <: ProtoRpcController
@@ -79,6 +104,8 @@ struct gRPCController <: ProtoRpcController
     revocation::Bool
     request_timeout::Real
     connect_timeout::Real
+    max_recv_message_length::Int
+    max_send_message_length::Int
     verbose::Bool
 
     function gRPCController(;
@@ -88,9 +115,18 @@ struct gRPCController <: ProtoRpcController
             revocation::Bool = true,
             request_timeout::Real = Inf,
             connect_timeout::Real = 0,
+            max_message_length::Integer = DEFAULT_MAX_MESSAGE_LENGTH,
+            max_recv_message_length::Integer = 0,
+            max_send_message_length::Integer = 0,
             verbose::Bool = false
         )
-        new(maxage, keepalive, negotiation, revocation, request_timeout, connect_timeout, verbose)
+        if maxage < 0 || keepalive < 0 || request_timeout < 0 || connect_timeout < 0 || 
+            max_message_length < 0 || max_recv_message_length < 0 || max_send_message_length < 0
+            throw(ArgumentError("Invalid gRPCController parameter"))
+        end
+        (max_recv_message_length == 0) && (max_recv_message_length = max_message_length)
+        (max_send_message_length == 0) && (max_send_message_length = max_message_length)
+        new(maxage, keepalive, negotiation, revocation, request_timeout, connect_timeout, max_recv_message_length, max_send_message_length, verbose)
     end
 end
 
@@ -118,13 +154,15 @@ struct gRPCChannel <: ProtoRpcChannel
     end
 end
 
-function to_delimited_message_bytes(msg)
+function to_delimited_message_bytes(msg, max_message_length::Int)
     iob = IOBuffer()
-    write(iob, UInt8(0))                # compression
-    write(iob, hton(UInt32(0)))         # message length (placeholder)
-    data_len = writeproto(iob, msg)     # message bytes
-    seek(iob, 1)                        # seek out the message length placeholder
-    write(iob, hton(UInt32(data_len)))  # fill the message length
+    limitiob = LimitIO(iob, max_message_length)
+    write(limitiob, UInt8(0))                   # compression
+    write(limitiob, hton(UInt32(0)))            # message length (placeholder)
+    data_len = writeproto(limitiob, msg)        # message bytes
+
+    seek(iob, 1)                                # seek out the message length placeholder
+    write(iob, hton(UInt32(data_len)))          # fill the message length
     take!(iob)
 end
 
@@ -155,6 +193,8 @@ function call_method(channel::gRPCChannel, service::ServiceDescriptor, method::M
         revocation = controller.revocation,
         request_timeout = controller.request_timeout,
         connect_timeout = controller.connect_timeout,
+        max_recv_message_length = controller.max_recv_message_length,
+        max_send_message_length = controller.max_send_message_length,
         verbose = controller.verbose,
     )
     outchannel, status_future