Wrote tests for message encode and decode

message.go

@@ -5,27 +5,12 @@ import (
 	"fmt"
 )
 
+// MessageHE contains ciphertext and an encrypted header.
 type MessageHE struct {
 	Header     []byte `json:"header"`
 	Ciphertext []byte `json:"ciphertext"`
 }
 
-type MessageEncHeader []byte
-
-func (mh MessageEncHeader) Decode() (MessageHeader, error) {
-	// n (4 bytes) + pn (4 bytes) + dh (32 bytes)
-	if len(mh) != 40 {
-		return MessageHeader{}, fmt.Errorf("encoded message header must be 40 bytes, %d given", len(mh))
-	}
-	var dh Key
-	copy(dh[:], mh[8:40])
-	return MessageHeader{
-		DH: dh,
-		N:  binary.LittleEndian.Uint32(mh[0:4]),
-		PN: binary.LittleEndian.Uint32(mh[4:8]),
-	}, nil
-}
-
 // Message is a single message exchanged by the parties.
 type Message struct {
 	Header     MessageHeader `json:"header"`
@@ -45,9 +30,27 @@ type MessageHeader struct {
 }
 
 // Encode the header in the binary format.
-func (mh MessageHeader) Encode() []byte {
+func (mh MessageHeader) Encode() MessageEncHeader {
 	buf := make([]byte, 8)
 	binary.LittleEndian.PutUint32(buf[0:4], mh.N)
 	binary.LittleEndian.PutUint32(buf[4:8], mh.PN)
 	return append(buf, mh.DH[:]...)
 }
+
+// MessageEncHeader is a binary-encoded representation of a message header.
+type MessageEncHeader []byte
+
+// Decode message header out of the binary-encoded representation.
+func (mh MessageEncHeader) Decode() (MessageHeader, error) {
+	// n (4 bytes) + pn (4 bytes) + dh (32 bytes)
+	if len(mh) != 40 {
+		return MessageHeader{}, fmt.Errorf("encoded message header must be 40 bytes, %d given", len(mh))
+	}
+	var dh Key
+	copy(dh[:], mh[8:40])
+	return MessageHeader{
+		DH: dh,
+		N:  binary.LittleEndian.Uint32(mh[0:4]),
+		PN: binary.LittleEndian.Uint32(mh[4:8]),
+	}, nil
+}

message_test.go

@@ -0,0 +1,33 @@
+package doubleratchet
+
+import (
+	"github.com/stretchr/testify/require"
+	"testing"
+)
+
+func TestMessageHeader_EncodeAndDecode(t *testing.T) {
+	// Arrange.
+	mh := MessageHeader{
+		DH: Key{0xe3, 0xbe, 0xb9, 0x4e, 0x70, 0x17, 0x37, 0xc, 0x1, 0x8f, 0xa9, 0x7e, 0xef, 0x4, 0xfb, 0x23, 0xac, 0xea, 0x28, 0xf7, 0xa9, 0x56, 0xcc, 0x1d, 0x46, 0xf3, 0xb5, 0x1d, 0x7d, 0x7d, 0x5e, 0x2c},
+		N:  uint32(10),
+		PN: uint32(15),
+	}
+
+	// Act.
+	decoded, err := mh.Encode().Decode()
+
+	// Assert.
+	require.Nil(t, err)
+	require.Equal(t, mh, decoded)
+}
+
+func TestMessageEncHeader_Decode(t *testing.T) {
+	// Arrange.
+	mhe := MessageEncHeader{0xe3, 0xbe, 0xb9, 0x4e, 0x70, 0x17, 0x37, 0xc, 0x1, 0x8f, 0xa9, 0x7e, 0xef, 0x4, 0xfb, 0x23, 0xac, 0xea, 0x28, 0xf7, 0xa9, 0x56, 0xcc, 0x1d, 0x46, 0xf3, 0xb5, 0x1d, 0x7d, 0x7d, 0x5e, 0x2c}
+
+	// Act.
+	_, err := mhe.Decode()
+
+	// Assert.
+	require.NotNil(t, err)
+}