From 3d87ebc2672b5456314354722d9d6354c09b285f Mon Sep 17 00:00:00 2001
From: bruce-riley <96066700+bruce-riley@users.noreply.github.com>
Date: Wed, 28 Jun 2023 14:51:30 -0500
Subject: [PATCH] CCQ: Message tweaks (#3147)

* CCQ: Message tweaks

Change-Id: I5d97fa0a32b42f1a763edf572b3d08365a7f5b20

* Add comments

Change-Id: I35a9fac0fbaae7d175e2529b842c5cef4ac386f6

* Fix comment paste error

Change-Id: I4dd23ea375e481d46e24a11d70b590221548c92e

* Address review comments

Change-Id: Ia516085fdcc182a5566b31342a5aa2480d37eedc
---
 node/hack/query/send_req.go |  4 +++-
 node/pkg/query/request.go   | 23 +++++++++++++++++++++++
 node/pkg/query/response.go  | 36 ++++++++++++++++++++++++++++++++++++
 3 files changed, 62 insertions(+), 1 deletion(-)

diff --git a/node/hack/query/send_req.go b/node/hack/query/send_req.go
index 6dd95de0d2..08518005fb 100644
--- a/node/hack/query/send_req.go
+++ b/node/hack/query/send_req.go
@@ -348,7 +348,9 @@ func sendQueryAndGetRsp(queryRequest *query.QueryRequest, sk *ecdsa.PrivateKey,
 		var isMatchingResponse bool
 		switch m := msg.Message.(type) {
 		case *gossipv1.GossipMessage_SignedQueryResponse:
-			logger.Info("query response received", zap.Any("response", m.SignedQueryResponse))
+			logger.Info("query response received", zap.Any("response", m.SignedQueryResponse),
+				zap.String("responseBytes", hexutil.Encode(m.SignedQueryResponse.QueryResponse)),
+				zap.String("sigBytes", hexutil.Encode(m.SignedQueryResponse.Signature)))
 			var response query.QueryResponsePublication
 			err := response.Unmarshal(m.SignedQueryResponse.QueryResponse)
 			if err != nil {
diff --git a/node/pkg/query/request.go b/node/pkg/query/request.go
index 79b25361a3..b2c6f33279 100644
--- a/node/pkg/query/request.go
+++ b/node/pkg/query/request.go
@@ -112,6 +112,7 @@ func (queryRequest *QueryRequest) Marshal() ([]byte, error) {
 
 	buf := new(bytes.Buffer)
 
+	vaa.MustWrite(buf, binary.BigEndian, uint8(1))           // version
 	vaa.MustWrite(buf, binary.BigEndian, queryRequest.Nonce) // uint32
 
 	vaa.MustWrite(buf, binary.BigEndian, uint8(len(queryRequest.PerChainQueries)))
@@ -134,6 +135,15 @@ func (queryRequest *QueryRequest) Unmarshal(data []byte) error {
 
 // UnmarshalFromReader deserializes the binary representation of a query request from an existing reader
 func (queryRequest *QueryRequest) UnmarshalFromReader(reader *bytes.Reader) error {
+	var version uint8
+	if err := binary.Read(reader, binary.BigEndian, &version); err != nil {
+		return fmt.Errorf("failed to read message version: %w", err)
+	}
+
+	if version != 1 {
+		return fmt.Errorf("unsupported message version: %d", version)
+	}
+
 	if err := binary.Read(reader, binary.BigEndian, &queryRequest.Nonce); err != nil {
 		return fmt.Errorf("failed to read request nonce: %w", err)
 	}
@@ -207,6 +217,13 @@ func (perChainQuery *PerChainQueryRequest) Marshal() ([]byte, error) {
 	if err != nil {
 		return nil, err
 	}
+
+	// Write the length of the query to facilitate on-chain parsing.
+	if len(queryBuf) > math.MaxUint32 {
+		return nil, fmt.Errorf("query too long")
+	}
+	vaa.MustWrite(buf, binary.BigEndian, uint32(len(queryBuf)))
+
 	buf.Write(queryBuf)
 	return buf.Bytes(), nil
 }
@@ -233,6 +250,12 @@ func (perChainQuery *PerChainQueryRequest) UnmarshalFromReader(reader *bytes.Rea
 		return err
 	}
 
+	// Skip the query length.
+	var queryLength uint32
+	if err := binary.Read(reader, binary.BigEndian, &queryLength); err != nil {
+		return fmt.Errorf("failed to read query length: %w", err)
+	}
+
 	switch queryType {
 	case EthCallQueryRequestType:
 		q := EthCallQueryRequest{}
diff --git a/node/pkg/query/response.go b/node/pkg/query/response.go
index b644ea9633..326195b981 100644
--- a/node/pkg/query/response.go
+++ b/node/pkg/query/response.go
@@ -97,12 +97,21 @@ func (msg *QueryResponsePublication) Marshal() ([]byte, error) {
 
 	buf := new(bytes.Buffer)
 
+	vaa.MustWrite(buf, binary.BigEndian, uint8(1)) // version
+
 	// Source
 	// TODO: support writing off-chain and on-chain requests
 	// Here, unset represents an off-chain request
 	vaa.MustWrite(buf, binary.BigEndian, vaa.ChainIDUnset)
 
 	buf.Write(msg.Request.Signature[:])
+
+	// Write the length of the request to facilitate on-chain parsing.
+	if len(msg.Request.QueryRequest) > math.MaxUint32 {
+		return nil, fmt.Errorf("request too long")
+	}
+	vaa.MustWrite(buf, binary.BigEndian, uint32(len(msg.Request.QueryRequest)))
+
 	buf.Write(msg.Request.QueryRequest)
 
 	// Per chain responses
@@ -122,6 +131,15 @@ func (msg *QueryResponsePublication) Marshal() ([]byte, error) {
 func (msg *QueryResponsePublication) Unmarshal(data []byte) error {
 	reader := bytes.NewReader(data[:])
 
+	var version uint8
+	if err := binary.Read(reader, binary.BigEndian, &version); err != nil {
+		return fmt.Errorf("failed to read message version: %w", err)
+	}
+
+	if version != 1 {
+		return fmt.Errorf("unsupported message version: %d", version)
+	}
+
 	// Request
 	requestChain := vaa.ChainID(0)
 	if err := binary.Read(reader, binary.BigEndian, &requestChain); err != nil {
@@ -139,6 +157,12 @@ func (msg *QueryResponsePublication) Unmarshal(data []byte) error {
 	}
 	signedQueryRequest.Signature = signature[:]
 
+	// Skip the query length.
+	queryRequestLen := uint32(0)
+	if err := binary.Read(reader, binary.BigEndian, &queryRequestLen); err != nil {
+		return fmt.Errorf("failed to read length of query request: %w", err)
+	}
+
 	queryRequest := QueryRequest{}
 	err := queryRequest.UnmarshalFromReader(reader)
 	if err != nil {
@@ -260,6 +284,12 @@ func (perChainResponse *PerChainQueryResponse) Marshal() ([]byte, error) {
 	if err != nil {
 		return nil, err
 	}
+
+	// Write the length of the response to facilitate on-chain parsing.
+	if len(respBuf) > math.MaxUint32 {
+		return nil, fmt.Errorf("response is too long")
+	}
+	vaa.MustWrite(buf, binary.BigEndian, uint32(len(respBuf)))
 	buf.Write(respBuf)
 	return buf.Bytes(), nil
 }
@@ -286,6 +316,12 @@ func (perChainResponse *PerChainQueryResponse) UnmarshalFromReader(reader *bytes
 		return err
 	}
 
+	// Skip the response length.
+	var respLength uint32
+	if err := binary.Read(reader, binary.BigEndian, &respLength); err != nil {
+		return fmt.Errorf("failed to read response length: %w", err)
+	}
+
 	switch queryType {
 	case EthCallQueryRequestType:
 		r := EthCallQueryResponse{}