KeepAlive: codec which produces valid CBOR encoding

IntersectMBO · Apr 15, 2021 · abcf8ed · abcf8ed
1 parent 5729847
commit abcf8ed
Show file tree

Hide file tree

Showing 2 changed files with 72 additions and 10 deletions.
diff --git a/ouroboros-network/protocol-tests/Ouroboros/Network/Protocol/KeepAlive/Test.hs b/ouroboros-network/protocol-tests/Ouroboros/Network/Protocol/KeepAlive/Test.hs
@@ -50,14 +50,17 @@ import Test.Tasty.QuickCheck (testProperty)
 
 tests :: TestTree
 tests = testGroup "Ouroboros.Network.Protocol.KeepAlive"
-  [ testProperty "direct"         prop_direct
-  , testProperty "connect"        prop_connect
-  , testProperty "channel ST"     prop_channel_ST
-  , testProperty "channel IO"     prop_channel_IO
-  , testProperty "codec"          prop_codec
-  , testProperty "codec 2-splits" prop_codec_splits2
-  , testProperty "codec 3-splits" (withMaxSuccess 33 prop_codec_splits3)
-  , testProperty "byteLimits"     prop_byteLimits
+  [ testProperty "direct"            prop_direct
+  , testProperty "connect"           prop_connect
+  , testProperty "channel ST"        prop_channel_ST
+  , testProperty "channel IO"        prop_channel_IO
+  , testProperty "codec"             prop_codec
+  , testProperty "codec 2-splits"    prop_codec_splits2
+  , testProperty "codec 3-splits"    (withMaxSuccess 33 prop_codec_splits3)
+  , testProperty "codec v2"          prop_codec_v2
+  , testProperty "codec v2 2-splits" prop_codec_v2_splits2
+  , testProperty "codec v2 3-splits" (withMaxSuccess 33 prop_codec_v2_splits3)
+  , testProperty "byteLimits"        prop_byteLimits
   ]
 
 --
@@ -155,6 +158,18 @@ prop_codec_splits3 :: AnyMessageAndAgency KeepAlive -> Bool
 prop_codec_splits3 msg =
     runST (prop_codec_splitsM splits3 codecKeepAlive msg)
 
+prop_codec_v2 :: AnyMessageAndAgency KeepAlive -> Bool
+prop_codec_v2 msg =
+    runST (prop_codecM codecKeepAlive_v2 msg)
+
+prop_codec_v2_splits2 :: AnyMessageAndAgency KeepAlive -> Bool
+prop_codec_v2_splits2 msg =
+    runST (prop_codec_splitsM splits2 codecKeepAlive_v2 msg)
+
+prop_codec_v2_splits3 :: AnyMessageAndAgency KeepAlive -> Bool
+prop_codec_v2_splits3 msg =
+    runST (prop_codec_splitsM splits3 codecKeepAlive_v2 msg)
+
 
 prop_byteLimits :: AnyMessageAndAgency KeepAlive
                          -> Bool

diff --git a/ouroboros-network/src/Ouroboros/Network/Protocol/KeepAlive/Codec.hs b/ouroboros-network/src/Ouroboros/Network/Protocol/KeepAlive/Codec.hs
@@ -8,6 +8,7 @@
 
 module Ouroboros.Network.Protocol.KeepAlive.Codec
   ( codecKeepAlive
+  , codecKeepAlive_v2
   , codecKeepAliveId
   , byteLimitsKeepAlive
   , timeLimitsKeepAlive
@@ -17,10 +18,13 @@ import           Control.Monad.Class.MonadST
 import           Control.Monad.Class.MonadTime (DiffTime)
 
 import           Data.ByteString.Lazy (ByteString)
+import           Text.Printf
 
-import qualified Codec.CBOR.Encoding as CBOR (Encoding, encodeWord, encodeWord16)
+import qualified Codec.CBOR.Encoding as CBOR ( Encoding, encodeListLen
+                                             , encodeWord, encodeWord16 )
 import qualified Codec.CBOR.Read     as CBOR
-import qualified Codec.CBOR.Decoding as CBOR (Decoder, decodeWord, decodeWord16)
+import qualified Codec.CBOR.Decoding as CBOR ( Decoder, decodeListLen
+                                             , decodeWord, decodeWord16 )
 
 import           Network.TypedProtocol.Core
 
@@ -64,6 +68,49 @@ codecKeepAlive = mkCodecCborLazyBS encodeMsg decodeMsg
            fail ("codecKeepAlive.StServer: unexpected key: " ++ show key)
 
 
+codecKeepAlive_v2
+  :: forall m.
+     MonadST m
+  => Codec KeepAlive CBOR.DeserialiseFailure m ByteString
+codecKeepAlive_v2 = mkCodecCborLazyBS encodeMsg decodeMsg
+   where
+     encodeMsg :: forall (pr :: PeerRole) st st'.
+                  PeerHasAgency pr st
+               -> Message KeepAlive st st'
+               -> CBOR.Encoding
+     encodeMsg (ClientAgency TokClient) (MsgKeepAlive (Cookie c)) =
+          CBOR.encodeListLen 2
+       <> CBOR.encodeWord 0
+       <> CBOR.encodeWord16 c
+     encodeMsg (ServerAgency TokServer) (MsgKeepAliveResponse (Cookie c)) =
+          CBOR.encodeListLen 2
+       <> CBOR.encodeWord 1
+       <> CBOR.encodeWord16 c
+     encodeMsg (ClientAgency TokClient) MsgDone =
+          CBOR.encodeListLen 1
+       <> CBOR.encodeWord 2
+
+     decodeMsg :: forall (pr :: PeerRole) s (st :: KeepAlive).
+                  PeerHasAgency pr st
+               -> CBOR.Decoder s (SomeMessage st)
+     decodeMsg stok = do
+       len <- CBOR.decodeListLen
+       key <- CBOR.decodeWord
+       case (stok, len, key) of
+         (ClientAgency TokClient, 2, 0) -> do
+             cookie <- CBOR.decodeWord16
+             return (SomeMessage $ MsgKeepAlive $ Cookie cookie)
+         (ServerAgency TokServer, 2, 1) -> do
+             cookie <- CBOR.decodeWord16
+             return (SomeMessage $ MsgKeepAliveResponse $ Cookie cookie)
+         (ClientAgency TokClient, 1, 2) -> pure (SomeMessage MsgDone)
+
+         (ClientAgency TokClient, _, _) ->
+           fail (printf "codecKeepAlive (%s) unexpected key (%d, %d)" (show stok) key len)
+         (ServerAgency TokServer, _, _ ) ->
+           fail (printf "codecKeepAlive (%s) unexpected key (%d, %d)" (show stok) key len)
+
+
 byteLimitsKeepAlive :: (bytes -> Word) -> ProtocolSizeLimits KeepAlive bytes
 byteLimitsKeepAlive = ProtocolSizeLimits sizeLimitForState
   where