huge message and payload refactor.

pkg/network/message.go

@@ -14,6 +14,7 @@ import (
 const (
 	// The minimum size of a valid message.
 	minMessageSize = 24
+	cmdSize        = 12
 )
 
 // NetMode type that is compatible with netModes below.
@@ -53,14 +54,14 @@ type Message struct {
 	Magic NetMode
 	// Command is utf8 code, of which the length is 12 bytes,
 	// the extra part is filled with 0.
-	Command []byte
+	Command [cmdSize]byte
 	// Length of the payload
 	Length uint32
 	// Checksum is the first 4 bytes of the value that two times SHA256
 	// hash of the payload
 	Checksum uint32
 	// Payload send with the message.
-	Payload payload.Payloader
+	Payload payload.Payload
 }
 
 type commandType string
@@ -80,7 +81,7 @@ const (
 	cmdTX                     = "tx"
 )
 
-func newMessage(magic NetMode, cmd commandType, p payload.Payloader) *Message {
+func newMessage(magic NetMode, cmd commandType, p payload.Payload) *Message {
 	var (
 		size     uint32
 		checksum []byte
@@ -88,15 +89,18 @@ func newMessage(magic NetMode, cmd commandType, p payload.Payloader) *Message {
 
 	if p != nil {
 		size = p.Size()
-		b, _ := p.MarshalBinary()
-		checksum = sumSHA256(sumSHA256(b))
+		buf := new(bytes.Buffer)
+		if err := p.EncodeBinary(buf); err != nil {
+			panic(err)
+		}
+		checksum = sumSHA256(sumSHA256(buf.Bytes()))
 	} else {
 		checksum = sumSHA256(sumSHA256([]byte{}))
 	}
 
 	return &Message{
 		Magic:    magic,
-		Command:  cmdToByteSlice(cmd),
+		Command:  cmdToByteArray(cmd),
 		Length:   size,
 		Payload:  p,
 		Checksum: binary.LittleEndian.Uint32(checksum[:4]),
@@ -105,7 +109,7 @@ func newMessage(magic NetMode, cmd commandType, p payload.Payloader) *Message {
 
 // Converts the 12 byte command slice to a commandType.
 func (m *Message) commandType() commandType {
-	cmd := string(bytes.TrimRight(m.Command, "\x00"))
+	cmd := cmdByteArrayToString(m.Command)
 	switch cmd {
 	case "version":
 		return cmdVersion
@@ -136,31 +140,20 @@ func (m *Message) commandType() commandType {
 
 // decode a Message from the given reader.
 func (m *Message) decode(r io.Reader) error {
-	// 24 bytes for the fixed sized fields.
-	buf := make([]byte, minMessageSize)
-	n, err := r.Read(buf)
-	if err != nil {
-		return err
-	}
-
-	if n != minMessageSize {
-		return fmt.Errorf("Expected to read exactly %d bytes got %d", minMessageSize, n)
-	}
-
-	m.Magic = NetMode(binary.LittleEndian.Uint32(buf[0:4]))
-	m.Command = buf[4:16]
-	m.Length = binary.LittleEndian.Uint32(buf[16:20])
-	m.Checksum = binary.LittleEndian.Uint32(buf[20:24])
+	binary.Read(r, binary.LittleEndian, &m.Magic)
+	binary.Read(r, binary.LittleEndian, &m.Command)
+	binary.Read(r, binary.LittleEndian, &m.Length)
+	binary.Read(r, binary.LittleEndian, &m.Checksum)
 
 	// return if their is no payload.
 	if m.Length == 0 {
 		return nil
 	}
 
-	return m.unmarshalPayload(r)
+	return m.decodePayload(r)
 }
 
-func (m *Message) unmarshalPayload(r io.Reader) error {
+func (m *Message) decodePayload(r io.Reader) error {
 	pbuf := make([]byte, m.Length)
 	n, err := r.Read(pbuf)
 	if err != nil {
@@ -176,23 +169,24 @@ func (m *Message) unmarshalPayload(r io.Reader) error {
 		return errors.New("checksum mismatch error")
 	}
 
-	var p payload.Payloader
+	rr := bytes.NewReader(pbuf)
+	var p payload.Payload
 	switch m.commandType() {
 	case cmdVersion:
 		p = &payload.Version{}
-		if err := p.UnmarshalBinary(pbuf); err != nil {
-			return err
-		}
-	case cmdInv:
-		p = &payload.Inventory{}
-		if err := p.UnmarshalBinary(pbuf); err != nil {
-			return err
-		}
-	case cmdAddr:
-		p = &payload.AddressList{}
-		if err := p.UnmarshalBinary(pbuf); err != nil {
+		if err := p.DecodeBinary(rr); err != nil {
 			return err
 		}
+		// case cmdInv:
+		// 	p = &payload.Inventory{}
+		// 	if err := p.UnmarshalBinary(pbuf); err != nil {
+		// 		return err
+		// 	}
+		// case cmdAddr:
+		// 	p = &payload.AddressList{}
+		// 	if err := p.UnmarshalBinary(pbuf); err != nil {
+		// 		return err
+		// 	}
 	}
 
 	m.Payload = p
@@ -202,37 +196,13 @@ func (m *Message) unmarshalPayload(r io.Reader) error {
 
 // encode a Message to any given io.Writer.
 func (m *Message) encode(w io.Writer) error {
-	buf := make([]byte, minMessageSize+m.Length)
-
-	binary.LittleEndian.PutUint32(buf[0:4], uint32(m.Magic))
-	copy(buf[4:16], m.Command)
-	binary.LittleEndian.PutUint32(buf[16:20], m.Length)
-	binary.LittleEndian.PutUint32(buf[20:24], m.Checksum)
+	binary.Write(w, binary.LittleEndian, m.Magic)
+	binary.Write(w, binary.LittleEndian, m.Command)
+	binary.Write(w, binary.LittleEndian, m.Length)
+	binary.Write(w, binary.LittleEndian, m.Checksum)
 
 	if m.Payload != nil {
-		payload, err := m.Payload.MarshalBinary()
-		if err != nil {
-			return err
-		}
-		copy(buf[minMessageSize:minMessageSize+m.Length], payload)
-	}
-
-	n, err := w.Write(buf)
-	if err != nil {
-		return err
-	}
-
-	// safety check to if we have written enough bytes.
-	if m.Length > 0 {
-		expectWritten := minMessageSize + m.Length
-		if uint32(n) != expectWritten {
-			return fmt.Errorf("expected to written exactly %d did %d", expectWritten, n)
-		}
-	} else {
-		expectWritten := minMessageSize
-		if n != expectWritten {
-			return fmt.Errorf("expected to written exactly %d did %d", expectWritten, n)
-		}
+		return m.Payload.EncodeBinary(w)
 	}
 
 	return nil
@@ -240,21 +210,31 @@ func (m *Message) encode(w io.Writer) error {
 
 // convert a command (string) to a byte slice filled with 0 bytes till
 // size 12.
-func cmdToByteSlice(cmd commandType) []byte {
+func cmdToByteArray(cmd commandType) [cmdSize]byte {
 	cmdLen := len(cmd)
-	if cmdLen > 12 {
+	if cmdLen > cmdSize {
 		panic("exceeded command max length of size 12")
 	}
 
 	// The command can have max 12 bytes, rest is filled with 0.
-	b := []byte(cmd)
-	for i := 0; i < 12-cmdLen; i++ {
-		b = append(b, byte('\x00'))
+	b := [cmdSize]byte{}
+	for i := 0; i < cmdLen; i++ {
+		b[i] = cmd[i]
 	}
 
 	return b
 }
 
+func cmdByteArrayToString(cmd [cmdSize]byte) string {
+	buf := []byte{}
+	for i := 0; i < cmdSize; i++ {
+		if cmd[i] != 0 {
+			buf = append(buf, cmd[i])
+		}
+	}
+	return string(buf)
+}
+
 func sumSHA256(b []byte) []byte {
 	h := sha256.New()
 	h.Write(b)

pkg/network/payload/payload.go

@@ -0,0 +1,10 @@
+package payload
+
+import "io"
+
+// Payload is anything that can be binary encoded/decoded.
+type Payload interface {
+	EncodeBinary(io.Writer) error
+	DecodeBinary(io.Reader) error
+	Size() uint32
+}

pkg/network/payload/payloader.go

@@ -1,13 +0,0 @@
-package payload
-
-import (
-	"encoding"
-)
-
-// Payloader is anything that can be binary marshaled and unmarshaled.
-// Every payload embbedded in messages need to satisfy the Payloader interface.
-type Payloader interface {
-	encoding.BinaryMarshaler
-	encoding.BinaryUnmarshaler
-	Size() uint32
-}

pkg/network/payload/version.go

@@ -2,9 +2,11 @@ package payload
 
 import (
 	"encoding/binary"
+	"io"
 )
 
 const (
+	lenUA          = 12
 	minVersionSize = 27
 )
 
@@ -21,7 +23,7 @@ type Version struct {
 	// it's used to distinguish the node from public IP
 	Nonce uint32
 	// client id
-	UserAgent []byte
+	UserAgent [lenUA]byte
 	// Height of the block chain
 	StartHeight uint32
 	// Whether to receive and forward
@@ -36,63 +38,53 @@ func NewVersion(id uint32, p uint16, ua string, h uint32, r bool) *Version {
 		Timestamp:   12345,
 		Port:        p,
 		Nonce:       id,
-		UserAgent:   []byte(ua),
+		UserAgent:   uaToByteArray(ua),
 		StartHeight: 0,
 		Relay:       r,
 	}
 }
 
-// UnmarshalBinary implements the Payloader interface.
-func (p *Version) UnmarshalBinary(b []byte) error {
-	// Length of the user agent should be calculated dynamicaly.
+// DecodeBinary implements the Payload interface.
+func (p *Version) DecodeBinary(r io.Reader) error {
+	// TODO: Length of the user agent should be calculated dynamicaly.
 	// There is no information about the size or format of this.
 	// the only thing we know is by looking at the #c source code.
 	// /NEO:{0}/ => /NEO:2.6.0/
-	lenUA := len(b) - minVersionSize
+	err := binary.Read(r, binary.LittleEndian, &p.Version)
+	err = binary.Read(r, binary.LittleEndian, &p.Services)
+	err = binary.Read(r, binary.LittleEndian, &p.Timestamp)
+	err = binary.Read(r, binary.LittleEndian, &p.Port)
+	err = binary.Read(r, binary.LittleEndian, &p.Nonce)
+	err = binary.Read(r, binary.LittleEndian, &p.UserAgent)
+	err = binary.Read(r, binary.LittleEndian, &p.StartHeight)
+	err = binary.Read(r, binary.LittleEndian, &p.Relay)
 
-	p.Version = binary.LittleEndian.Uint32(b[0:4])
-	p.Services = binary.LittleEndian.Uint64(b[4:12])
-	p.Timestamp = binary.LittleEndian.Uint32(b[12:16])
-	// FIXME: port's byteorder should be big endian according to the docs.
-	// but when connecting to the privnet docker image it's little endian.
-	p.Port = binary.LittleEndian.Uint16(b[16:18])
-	p.Nonce = binary.LittleEndian.Uint32(b[18:22])
-	p.UserAgent = b[22 : 22+lenUA]
-	curlen := 22 + lenUA
-	p.StartHeight = binary.LittleEndian.Uint32(b[curlen : curlen+4])
-	p.Relay = b[len(b)-1 : len(b)][0] == 1
-
-	return nil
+	return err
 }
 
-// MarshalBinary implements the Payloader interface.
-func (p *Version) MarshalBinary() ([]byte, error) {
-	b := make([]byte, p.Size())
+// EncodeBinary implements the Payload interface.
+func (p *Version) EncodeBinary(w io.Writer) error {
+	err := binary.Write(w, binary.LittleEndian, p.Version)
+	err = binary.Write(w, binary.LittleEndian, p.Services)
+	err = binary.Write(w, binary.LittleEndian, p.Timestamp)
+	err = binary.Write(w, binary.LittleEndian, p.Port)
+	err = binary.Write(w, binary.LittleEndian, p.Nonce)
+	err = binary.Write(w, binary.LittleEndian, p.UserAgent)
+	err = binary.Write(w, binary.LittleEndian, p.StartHeight)
+	err = binary.Write(w, binary.LittleEndian, p.Relay)
 
-	binary.LittleEndian.PutUint32(b[0:4], p.Version)
-	binary.LittleEndian.PutUint64(b[4:12], p.Services)
-	binary.LittleEndian.PutUint32(b[12:16], p.Timestamp)
-	// FIXME: byte order (little / big)?
-	binary.LittleEndian.PutUint16(b[16:18], p.Port)
-	binary.LittleEndian.PutUint32(b[18:22], p.Nonce)
-	copy(b[22:22+len(p.UserAgent)], p.UserAgent) //
-	curLen := 22 + len(p.UserAgent)
-	binary.LittleEndian.PutUint32(b[curLen:curLen+4], p.StartHeight)
-
-	// yikes
-	var bln []byte
-	if p.Relay {
-		bln = []byte{1}
-	} else {
-		bln = []byte{0}
-	}
-
-	copy(b[curLen+4:len(b)], bln)
-
-	return b, nil
+	return err
 }
 
 // Size implements the payloader interface.
 func (p *Version) Size() uint32 {
 	return uint32(minVersionSize + len(p.UserAgent))
 }
+
+func uaToByteArray(ua string) [lenUA]byte {
+	buf := [lenUA]byte{}
+	for i := 0; i < lenUA; i++ {
+		buf[i] = ua[i]
+	}
+	return buf
+}

pkg/network/payload/version_test.go

@@ -7,15 +7,23 @@ import (
 )
 
 func TestVersionEncodeDecode(t *testing.T) {
-	p := NewVersion(3000, "/NEO/", 0, true)
+	version := NewVersion(13337, 3000, "./NEO:0.0.1/", 0, true)
 
 	buf := new(bytes.Buffer)
-	p.Encode(buf)
+	if err := version.EncodeBinary(buf); err != nil {
+		t.Fatal(err)
+	}
 
-	pd := &Version{}
-	pd.Decode(buf)
+	versionDecoded := &Version{}
+	if err := versionDecoded.DecodeBinary(buf); err != nil {
+		t.Fatal(err)
+	}
 
-	if !reflect.DeepEqual(p, pd) {
-		t.Fatalf("expect %v to be equal to %v", p, pd)
+	if !reflect.DeepEqual(version, versionDecoded) {
+		t.Fatalf("expected both version payload to be equal: %+v and %+v", version, versionDecoded)
+	}
+
+	if version.Size() != uint32(minVersionSize+len(version.UserAgent)) {
+		t.Fatalf("Expected version size of %d", minVersionSize+lenUA)
 	}
 }

pkg/network/server.go

@@ -181,9 +181,9 @@ func (s *Server) processMessage(msg *Message, peer *Peer) error {
 		return s.handleVersionCmd(msg.Payload.(*payload.Version), peer)
 	case cmdVerack:
 	case cmdGetAddr:
-		return s.handleGetAddrCmd(msg, peer)
+		// return s.handleGetAddrCmd(msg, peer)
 	case cmdAddr:
-		return s.handleAddrCmd(msg.Payload.(*payload.AddressList), peer)
+		// return s.handleAddrCmd(msg.Payload.(*payload.AddressList), peer)
 	case cmdGetHeaders:
 	case cmdHeaders:
 	case cmdGetBlocks: