test: payload encoding / decoding (#96)

waku/v2/node/waku_payload.go

@@ -6,6 +6,7 @@ import (
 	"crypto/ecdsa"
 	crand "crypto/rand"
 	"encoding/binary"
+	"fmt"
 	mrand "math/rand"
 
 	"errors"
@@ -70,22 +71,38 @@ func (payload Payload) Encode(version uint32) ([]byte, error) {
 		case Symmetric:
 			encoded, err := encryptSymmetric(data, payload.Key.SymKey)
 			if err != nil {
-				return nil, errors.New("Couldn't encrypt using symmetric key")
+				return nil, fmt.Errorf("couldn't encrypt using symmetric key: %w", err)
 			} else {
 				return encoded, nil
 			}
 		case Asymmetric:
 			encoded, err := encryptAsymmetric(data, &payload.Key.PubKey)
 			if err != nil {
-				return nil, errors.New("Couldn't encrypt using asymmetric key")
+				return nil, fmt.Errorf("couldn't encrypt using asymmetric key: %w", err)
 			} else {
 				return encoded, nil
 			}
 		case None:
-			return nil, errors.New("Non supported KeyKind")
+			return nil, errors.New("non supported KeyKind")
 		}
 	}
-	return nil, errors.New("Unsupported WakuMessage version")
+	return nil, errors.New("unsupported wakumessage version")
+}
+
+func EncodeWakuMessage(message *pb.WakuMessage, keyInfo *KeyInfo) error {
+	msgPayload := message.Payload
+	payload := Payload{
+		Data: msgPayload,
+		Key:  keyInfo,
+	}
+
+	encodedBytes, err := payload.Encode(message.Version)
+	if err != nil {
+		return err
+	}
+
+	message.Payload = encodedBytes
+	return nil
 }
 
 // Decodes a WakuMessage depending on the version parameter.
@@ -98,12 +115,12 @@ func DecodePayload(message *pb.WakuMessage, keyInfo *KeyInfo) (*DecodedPayload,
 		switch keyInfo.Kind {
 		case Symmetric:
 			if keyInfo.SymKey == nil {
-				return nil, errors.New("Symmetric key is required")
+				return nil, errors.New("symmetric key is required")
 			}
 
 			decodedData, err := decryptSymmetric(message.Payload, keyInfo.SymKey)
 			if err != nil {
-				return nil, errors.New("Couldn't decrypt using symmetric key")
+				return nil, fmt.Errorf("couldn't decrypt using symmetric key: %w", err)
 			}
 
 			decodedPayload, err := validateAndParse(decodedData)
@@ -114,12 +131,12 @@ func DecodePayload(message *pb.WakuMessage, keyInfo *KeyInfo) (*DecodedPayload,
 			return decodedPayload, nil
 		case Asymmetric:
 			if keyInfo.PrivKey == nil {
-				return nil, errors.New("Private key is required")
+				return nil, errors.New("private key is required")
 			}
 
 			decodedData, err := decryptAsymmetric(message.Payload, keyInfo.PrivKey)
 			if err != nil {
-				return nil, errors.New("Couldn't decrypt using asymmetric key")
+				return nil, fmt.Errorf("couldn't decrypt using asymmetric key: %w", err)
 			}
 
 			decodedPayload, err := validateAndParse(decodedData)
@@ -129,10 +146,20 @@ func DecodePayload(message *pb.WakuMessage, keyInfo *KeyInfo) (*DecodedPayload,
 
 			return decodedPayload, nil
 		case None:
-			return nil, errors.New("Non supported KeyKind")
+			return nil, errors.New("non supported KeyKind")
 		}
 	}
-	return nil, errors.New("Unsupported WakuMessage version")
+	return nil, errors.New("unsupported wakumessage version")
+}
+
+func DecodeWakuMessage(message *pb.WakuMessage, keyInfo *KeyInfo) error {
+	decodedPayload, err := DecodePayload(message, keyInfo)
+	if err != nil {
+		return err
+	}
+
+	message.Payload = decodedPayload.Data
+	return nil
 }
 
 const aesNonceLength = 12

waku/v2/node/waku_payload_encoding_test.go

@@ -0,0 +1,163 @@
+package node
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/ethereum/go-ethereum/crypto"
+	"github.com/status-im/go-waku/waku/v2/protocol/pb"
+	"github.com/stretchr/testify/require"
+)
+
+func createTestMsg() *pb.WakuMessage {
+	message := new(pb.WakuMessage)
+	message.Payload = []byte{0, 1, 2}
+	message.Version = 0
+	message.Timestamp = float64(123456)
+	return message
+}
+
+func TestEncodeDecodePayload(t *testing.T) {
+	data := []byte{0, 1, 2}
+	version := uint32(0)
+
+	keyInfo := new(KeyInfo)
+	keyInfo.Kind = None
+
+	p := Payload{
+		Data: data,
+		Key:  keyInfo,
+	}
+
+	encodedPayload, err := p.Encode(version)
+	require.NoError(t, err)
+	require.Equal(t, data, encodedPayload)
+
+	message := new(pb.WakuMessage)
+	message.Payload = encodedPayload
+	message.Version = version
+	message.Timestamp = float64(123456)
+
+	decodedPayload, err := DecodePayload(message, keyInfo)
+	require.NoError(t, err)
+	require.Equal(t, data, decodedPayload.Data)
+}
+
+func TestEncodeDecodeVersion0(t *testing.T) {
+	message := createTestMsg()
+
+	keyInfo := new(KeyInfo)
+	keyInfo.Kind = None
+
+	err := EncodeWakuMessage(message, keyInfo)
+	require.NoError(t, err)
+
+	err = DecodeWakuMessage(message, keyInfo)
+	require.NoError(t, err)
+}
+
+func generateSymKey() ([]byte, error) {
+	key, err := generateSecureRandomData(aesKeyLength)
+	if err != nil {
+		return nil, err
+	} else if !validateDataIntegrity(key, aesKeyLength) {
+		return nil, fmt.Errorf("error in generating symkey: crypto/rand failed to generate random data")
+	}
+
+	return key, nil
+}
+
+func TestEncodeDecodeVersion1Symmetric(t *testing.T) {
+	message := createTestMsg()
+	message.Version = 1
+	data := message.Payload
+
+	keyInfo := new(KeyInfo)
+	keyInfo.Kind = Symmetric
+
+	var err error
+	keyInfo.SymKey, err = generateSymKey()
+	require.NoError(t, err)
+
+	err = EncodeWakuMessage(message, keyInfo)
+	require.NoError(t, err)
+	require.NotEqual(t, data, message.Payload)
+	require.NotNil(t, message.Payload)
+
+	err = DecodeWakuMessage(message, keyInfo)
+	require.NoError(t, err)
+	require.Equal(t, data, message.Payload)
+}
+
+func TestEncodeDecodeVersion1Asymmetric(t *testing.T) {
+	message := createTestMsg()
+	message.Version = 1
+	data := message.Payload
+
+	privKey, err := crypto.GenerateKey()
+	require.NoError(t, err)
+
+	keyInfo := new(KeyInfo)
+	keyInfo.Kind = Asymmetric
+	keyInfo.PubKey = privKey.PublicKey
+
+	err = EncodeWakuMessage(message, keyInfo)
+	require.NoError(t, err)
+	require.NotEqual(t, data, message.Payload)
+	require.NotNil(t, message.Payload)
+
+	keyInfo.PrivKey = privKey
+	err = DecodeWakuMessage(message, keyInfo)
+	require.NoError(t, err)
+	require.Equal(t, data, message.Payload)
+}
+
+func TestEncodeDecodeIncorrectKey(t *testing.T) {
+	message := createTestMsg()
+	message.Version = 1
+
+	privKey, err := crypto.GenerateKey()
+	require.NoError(t, err)
+
+	symKey, err := generateSymKey()
+	require.NoError(t, err)
+
+	keyInfo := new(KeyInfo)
+	keyInfo.Kind = Asymmetric
+	keyInfo.SymKey = symKey
+
+	err = EncodeWakuMessage(message, keyInfo)
+	require.Error(t, err)
+
+	keyInfo.SymKey = nil
+	keyInfo.PrivKey = privKey
+
+	err = EncodeWakuMessage(message, keyInfo)
+	require.Error(t, err)
+}
+
+func TestEncodeUnsupportedVersion(t *testing.T) {
+	message := createTestMsg()
+	message.Version = 99
+
+	keyInfo := new(KeyInfo)
+	keyInfo.Kind = None
+
+	err := EncodeWakuMessage(message, keyInfo)
+	require.Error(t, err)
+	require.Equal(t, "unsupported wakumessage version", err.Error())
+}
+
+func TestDecodeUnsupportedVersion(t *testing.T) {
+	message := createTestMsg()
+	message.Version = 99
+
+	keyInfo := new(KeyInfo)
+	keyInfo.Kind = None
+
+	decodedPayload, err := DecodePayload(message, keyInfo)
+
+	require.Nil(t, decodedPayload)
+	require.Error(t, err)
+	require.Equal(t, "unsupported wakumessage version", err.Error())
+}