From 3aa835a6f2ed9a47b82f3d7ae80a11e8beacb5b5 Mon Sep 17 00:00:00 2001
From: intractable <intractable@gmail.com>
Date: Fri, 20 Oct 2017 16:39:26 -0500
Subject: [PATCH] Expose `GRPC_ARG_MAX_RECEIVE_MESSAGE_LENGTH` channel arg
 (#35)

* Expose the max receive message length channel argument + unittest

* Remove unused `MultiWayIf` ext
---
 cbits/grpc_haskell.c                    |  2 ++
 include/grpc_haskell.h                  |  3 +-
 src/Network/GRPC/Unsafe/ChannelArgs.chs | 11 ++++++--
 tests/LowLevelTests.hs                  | 37 +++++++++++++++++++++++++
 4 files changed, 49 insertions(+), 4 deletions(-)

diff --git a/cbits/grpc_haskell.c b/cbits/grpc_haskell.c
index c6f0ad6..73dc24e 100644
--- a/cbits/grpc_haskell.c
+++ b/cbits/grpc_haskell.c
@@ -438,6 +438,8 @@ char* translate_arg_key(enum supported_arg_key key){
       return GRPC_ARG_PRIMARY_USER_AGENT_STRING;
     case user_agent_suffix_key:
       return GRPC_ARG_SECONDARY_USER_AGENT_STRING;
+    case max_receive_message_length_key:
+      return GRPC_ARG_MAX_RECEIVE_MESSAGE_LENGTH;
     default:
       return "unknown_arg_key";
   }
diff --git a/include/grpc_haskell.h b/include/grpc_haskell.h
index fa82397..09ab5a3 100644
--- a/include/grpc_haskell.h
+++ b/include/grpc_haskell.h
@@ -154,7 +154,8 @@ enum supported_arg_key {
   compression_algorithm_key = 0,
   compression_level_key,
   user_agent_prefix_key,
-  user_agent_suffix_key
+  user_agent_suffix_key,
+  max_receive_message_length_key,
 };
 
 grpc_arg* create_arg_array(size_t n);
diff --git a/src/Network/GRPC/Unsafe/ChannelArgs.chs b/src/Network/GRPC/Unsafe/ChannelArgs.chs
index f7cd1c2..7290f18 100644
--- a/src/Network/GRPC/Unsafe/ChannelArgs.chs
+++ b/src/Network/GRPC/Unsafe/ChannelArgs.chs
@@ -6,6 +6,7 @@ import           Control.Exception
 import           Control.Monad
 import           Foreign.Marshal.Alloc (malloc, free)
 import           Foreign.Storable
+import           Numeric.Natural
 
 #include <grpc/grpc.h>
 #include <grpc/status.h>
@@ -47,9 +48,10 @@ data ArgValue = StringArg String | IntArg Int
 -- | Supported arguments for a channel. More cases will be added as we figure
 -- out what they are.
 data Arg = CompressionAlgArg CompressionAlgorithm
-           | CompressionLevelArg CompressionLevel
-           | UserAgentPrefix String
-           | UserAgentSuffix String
+         | CompressionLevelArg CompressionLevel
+         | UserAgentPrefix String
+         | UserAgentSuffix String
+         | MaxReceiveMessageLength Natural
   deriving (Show, Eq)
 
 {#fun create_string_arg as ^ {`GrpcArg', `Int', `ArgKey', `String'} -> `()'#}
@@ -67,6 +69,9 @@ createArg array (UserAgentPrefix prefix) i =
   createStringArg array i UserAgentPrefixKey prefix
 createArg array (UserAgentSuffix suffix) i =
   createStringArg array i UserAgentSuffixKey suffix
+createArg array (MaxReceiveMessageLength n) i =
+  createIntArg array i MaxReceiveMessageLengthKey $
+    fromIntegral (min n (fromIntegral (maxBound :: Int)))
 
 createChannelArgs :: [Arg] -> IO GrpcChannelArgs
 createChannelArgs args = do
diff --git a/tests/LowLevelTests.hs b/tests/LowLevelTests.hs
index fc33846..58945a1 100644
--- a/tests/LowLevelTests.hs
+++ b/tests/LowLevelTests.hs
@@ -57,6 +57,7 @@ lowLevelTests = testGroup "Unit tests of low-level Haskell library"
   , testCustomUserAgent
   , testClientCompression
   , testClientServerCompression
+  , testClientMaxReceiveMessageLengthChannelArg
   , testClientStreaming
   , testClientStreamingUnregistered
   , testServerStreaming
@@ -766,6 +767,42 @@ testClientServerCompressionLvl =
         return ("hello", dummyMeta, StatusOk, StatusDetails "")
       return ()
 
+testClientMaxReceiveMessageLengthChannelArg :: TestTree
+testClientMaxReceiveMessageLengthChannelArg = do
+  testGroup "max receive message length channel arg (client channel)"
+    [ csTest' "payload size < small bound succeeds" shouldSucceed server
+    , csTest' "payload size > small bound fails"    shouldFail    server
+    ]
+  where
+    -- The server always sends a 4-byte payload
+    pay    = "four"
+    server = TestServer (ServerConfig "localhost" 50051 ["/foo"] [] [] [] [] Nothing) $ \s -> do
+      let rm = head (normalMethods s)
+      void $ serverHandleNormalCall s rm mempty $ \sc -> do
+        payload sc @?= pay
+        pure (pay, mempty, StatusOk, StatusDetails "")
+
+    clientMax n k = TestClient conf $ \c -> do
+      rm <- clientRegisterMethodNormal c "/foo"
+      clientRequest c rm 1 pay mempty >>= k
+      where
+        conf = ClientConfig "localhost" 50051 [MaxReceiveMessageLength n] Nothing
+
+    -- Expect success when the max recv payload size is set to 4 bytes, and we
+    -- are sent 4.
+    shouldSucceed = clientMax 4 $ checkReqRslt $ \NormalRequestResult{..} -> do
+      rspCode @?= StatusOk
+      rspBody @?= pay
+      details @?= ""
+
+    -- Expect failure when the max recv payload size is set to 3 bytes, and we
+    -- are sent 4.
+    shouldFail = clientMax 3 $ \case
+        Left (GRPCIOBadStatusCode StatusInvalidArgument _)
+          -> pure ()
+        rsp
+          -> clientFail ("Expected failure response, but got: " ++ show rsp)
+
 --------------------------------------------------------------------------------
 -- Utilities and helpers