From 63072ebe75cd2ebb1c58d0f23a393bc3cb787b7b Mon Sep 17 00:00:00 2001
From: anthdm <anthony@academiclabs.co>
Date: Thu, 1 Feb 2018 10:56:33 +0100
Subject: [PATCH] Added getheaders payload + abstracted the fields.

---
 pkg/network/payload/getblocks.go       | 32 +++++++++++++---------
 pkg/network/payload/getheaders.go      | 17 ++++++++++++
 pkg/network/payload/getheaders_test.go | 37 ++++++++++++++++++++++++++
 3 files changed, 73 insertions(+), 13 deletions(-)
 create mode 100644 pkg/network/payload/getheaders.go
 create mode 100644 pkg/network/payload/getheaders_test.go

diff --git a/pkg/network/payload/getblocks.go b/pkg/network/payload/getblocks.go
index 909335b9a..81371cb25 100644
--- a/pkg/network/payload/getblocks.go
+++ b/pkg/network/payload/getblocks.go
@@ -7,24 +7,17 @@ import (
 	. "github.com/anthdm/neo-go/pkg/util"
 )
 
-// GetBlocks payload
-type GetBlocks struct {
+// HashStartStop contains fields and methods to be shared with the
+// "GetBlocks" and "GetHeaders" payload.
+type HashStartStop struct {
 	// hash of latest block that node requests
 	HashStart []Uint256
 	// hash of last block that node requests
 	HashStop Uint256
 }
 
-// NewGetBlocks return a pointer to a GetBlocks object.
-func NewGetBlocks(start []Uint256, stop Uint256) *GetBlocks {
-	return &GetBlocks{
-		HashStart: start,
-		HashStop:  stop,
-	}
-}
-
 // DecodeBinary implements the payload interface.
-func (p *GetBlocks) DecodeBinary(r io.Reader) error {
+func (p *HashStartStop) DecodeBinary(r io.Reader) error {
 	var lenStart uint8
 
 	err := binary.Read(r, binary.LittleEndian, &lenStart)
@@ -36,7 +29,7 @@ func (p *GetBlocks) DecodeBinary(r io.Reader) error {
 }
 
 // EncodeBinary implements the payload interface.
-func (p *GetBlocks) EncodeBinary(w io.Writer) error {
+func (p *HashStartStop) EncodeBinary(w io.Writer) error {
 	err := binary.Write(w, binary.LittleEndian, uint8(len(p.HashStart)))
 	err = binary.Write(w, binary.LittleEndian, p.HashStart)
 	err = binary.Write(w, binary.LittleEndian, p.HashStop)
@@ -45,4 +38,17 @@ func (p *GetBlocks) EncodeBinary(w io.Writer) error {
 }
 
 // Size implements the payload interface.
-func (p *GetBlocks) Size() uint32 { return 0 }
+func (p *HashStartStop) Size() uint32 { return 0 }
+
+// GetBlocks payload
+type GetBlocks struct {
+	HashStartStop
+}
+
+// NewGetBlocks return a pointer to a GetBlocks object.
+func NewGetBlocks(start []Uint256, stop Uint256) *GetBlocks {
+	p := &GetBlocks{}
+	p.HashStart = start
+	p.HashStop = stop
+	return p
+}
diff --git a/pkg/network/payload/getheaders.go b/pkg/network/payload/getheaders.go
new file mode 100644
index 000000000..9ade274f7
--- /dev/null
+++ b/pkg/network/payload/getheaders.go
@@ -0,0 +1,17 @@
+package payload
+
+import "github.com/anthdm/neo-go/pkg/util"
+
+// GetHeaders payload is the same as the "GetBlocks" payload.
+type GetHeaders struct {
+	HashStartStop
+}
+
+// NewGetHeaders return a pointer to a GetHeaders object.
+func NewGetHeaders(start []util.Uint256, stop util.Uint256) *GetHeaders {
+	p := &GetHeaders{}
+	p.HashStart = start
+	p.HashStop = stop
+
+	return p
+}
diff --git a/pkg/network/payload/getheaders_test.go b/pkg/network/payload/getheaders_test.go
new file mode 100644
index 000000000..32cbb1b86
--- /dev/null
+++ b/pkg/network/payload/getheaders_test.go
@@ -0,0 +1,37 @@
+package payload
+
+import (
+	"bytes"
+	"crypto/sha256"
+	"reflect"
+	"testing"
+
+	"github.com/anthdm/neo-go/pkg/util"
+)
+
+func TestGetHeadersEncodeDecode(t *testing.T) {
+	start := []util.Uint256{
+		sha256.Sum256([]byte("a")),
+		sha256.Sum256([]byte("b")),
+	}
+	stop := sha256.Sum256([]byte("c"))
+
+	p := NewGetHeaders(start, stop)
+	buf := new(bytes.Buffer)
+	if err := p.EncodeBinary(buf); err != nil {
+		t.Fatal(err)
+	}
+
+	if have, want := buf.Len(), 1+64+32; have != want {
+		t.Fatalf("expecting a length of %d got %d", want, have)
+	}
+
+	pDecode := &GetHeaders{}
+	if err := pDecode.DecodeBinary(buf); err != nil {
+		t.Fatal(err)
+	}
+
+	if !reflect.DeepEqual(p, pDecode) {
+		t.Fatalf("expecting both getheaders payloads to be equal %v and %v", p, pDecode)
+	}
+}