From f4d9421836b9fa796be7d2c441a8389819524bc7 Mon Sep 17 00:00:00 2001
From: KonradStaniec <konrad.staniec@gmail.com>
Date: Wed, 18 Aug 2021 09:23:57 +0200
Subject: [PATCH] Fix format of content message (#803)

---
 fluffy/network/state/content.nim         | 21 +++++++++++++++++++--
 fluffy/network/state/messages.nim        |  5 +----
 fluffy/network/state/portal_protocol.nim |  5 ++++-
 fluffy/tests/test_portal_encoding.nim    | 11 +++++++----
 4 files changed, 31 insertions(+), 11 deletions(-)

diff --git a/fluffy/network/state/content.nim b/fluffy/network/state/content.nim
index 0738b21d0..79d3aaf01 100644
--- a/fluffy/network/state/content.nim
+++ b/fluffy/network/state/content.nim
@@ -10,6 +10,7 @@
 {.push raises: [Defect].}
 
 import
+  std/[options, sugar],
   nimcrypto/[sha2, hash], stew/objects,
   eth/ssz/ssz_serialization, eth/trie/[hexary, db]
 
@@ -61,11 +62,24 @@ func fromSszBytes*(T: type ContentType, data: openArray[byte]):
 
   contentType
 
-func toContentId*(contentKey: ContentKey): ContentId =
+func encodeKey*(contentKey: ContentKey): seq[byte] =
+  SSZ.encode(contentKey)
+
+# TODO consider if more powerfull error handling is necessary here.
+func decodeKey*(contentKey: ByteList): Option[ContentKey] = 
+  try:
+    some(SSZ.decode(contentKey.asSeq(), ContentKey))
+  except SszError:
+    return none[ContentKey]()
+
+func encodeKeyAsList*(contentKey: ContentKey): ByteList =
+  List.init(encodeKey(contentKey), 2048)
+
+func toContentId*(contentKey: ByteList): ContentId =
   # TODO: Hash function to be defined, sha256 used now, might be confusing
   # with keccak256 that is used for the actual nodes:
   # https://github.com/ethereum/stateless-ethereum-specs/blob/master/state-network.md#content
-  sha2.sha_256.digest(SSZ.encode(contentKey))
+  sha2.sha_256.digest(contentKey.asSeq())
 
 type
   ContentStorage* = object
@@ -84,3 +98,6 @@ proc getContent*(storage: ContentStorage, key: ContentKey): Option[seq[byte]] =
     some(val)
   else:
     none(seq[byte])
+
+proc getContent*(storage: ContentStorage, contentKey: ByteList): Option[seq[byte]] =
+  decodeKey(contentKey).flatMap((key: ContentKey) => getContent(storage, key))
diff --git a/fluffy/network/state/messages.nim b/fluffy/network/state/messages.nim
index 7bdd3aec8..0031775dd 100644
--- a/fluffy/network/state/messages.nim
+++ b/fluffy/network/state/messages.nim
@@ -48,10 +48,7 @@ type
     # also be limited to 300 bytes instead of 2048
 
   FindContentMessage* = object
-    contentKey*: ContentKey
-  # TODO: According to the specification, this is actually a ByteList from the
-  # serialized ContentKey container, which will result in a different
-  # serialization
+    contentKey*: ByteList
 
   FoundContentMessage* = object
     enrs*: List[ByteList, 32]
diff --git a/fluffy/network/state/portal_protocol.nim b/fluffy/network/state/portal_protocol.nim
index baa0f418f..470577a04 100644
--- a/fluffy/network/state/portal_protocol.nim
+++ b/fluffy/network/state/portal_protocol.nim
@@ -176,9 +176,12 @@ proc findNode*(p: PortalProtocol, dst: Node, distances: List[uint16, 256]):
   # TODO Add nodes validation
   return await reqResponse[FindNodeMessage, NodesMessage](p, dst, PortalProtocolId, fn)
 
+# TODO It maybe better to accept bytelist, to not tie network layer to particular
+# content
 proc findContent*(p: PortalProtocol, dst: Node, contentKey: ContentKey):
     Future[PortalResult[FoundContentMessage]] {.async.} =
-  let fc = FindContentMessage(contentKey: contentKey)
+  let encKey = encodeKeyAsList(contentKey)
+  let fc = FindContentMessage(contentKey: encKey)
 
   trace "Send message request", dstId = dst.id, kind = MessageKind.findcontent
   return await reqResponse[FindContentMessage, FoundContentMessage](p, dst, PortalProtocolId, fc)
diff --git a/fluffy/tests/test_portal_encoding.nim b/fluffy/tests/test_portal_encoding.nim
index 520886a38..bed2bd194 100644
--- a/fluffy/tests/test_portal_encoding.nim
+++ b/fluffy/tests/test_portal_encoding.nim
@@ -113,18 +113,21 @@ suite "Portal Protocol Message Encodings":
         networkId: 0'u16,
         contentType: ContentType.Account,
         nodeHash: nodeHash)
-      fn = FindContentMessage(contentKey: contentKey)
+
+      contentEncoded: ByteList = encodeKeyAsList(contentKey)
+      
+      fn = FindContentMessage(contentKey: contentEncoded)
 
     let encoded = encodeMessage(fn)
-    check encoded.toHex == "050000010000000000000000000000000000000000000000000000000000000000000000"
+    check encoded.toHex == "05040000000000010000000000000000000000000000000000000000000000000000000000000000"
 
     let decoded = decodeMessage(encoded)
     check decoded.isOk()
-
+    
     let message = decoded.get()
     check:
       message.kind == findcontent
-      message.findcontent.contentKey == contentKey
+      message.findcontent.contentKey == contentEncoded
 
   test "FoundContent Response - payload":
     let