From 5b721fb000811c446a38ffef7a24e62a780a34b7 Mon Sep 17 00:00:00 2001
From: Anna Shaleva <shaleva.ann@nspcc.ru>
Date: Fri, 22 May 2020 12:59:18 +0300
Subject: [PATCH] protocol: add capabilities to address payload

Part of #871
---
 pkg/network/payload/address.go      | 31 +++++++++++++++++++----------
 pkg/network/payload/address_test.go |  7 ++++++-
 2 files changed, 26 insertions(+), 12 deletions(-)

diff --git a/pkg/network/payload/address.go b/pkg/network/payload/address.go
index e3d34ae926..eb97071517 100644
--- a/pkg/network/payload/address.go
+++ b/pkg/network/payload/address.go
@@ -6,22 +6,26 @@ import (
 	"time"
 
 	"github.com/nspcc-dev/neo-go/pkg/io"
+	"github.com/nspcc-dev/neo-go/pkg/network/capability"
 )
 
 // AddressAndTime payload.
 type AddressAndTime struct {
-	Timestamp uint32
-	Services  uint64
-	IP        [16]byte
-	Port      uint16
+	Timestamp    uint32
+	IP           [16]byte
+	Capabilities []capability.Capability
 }
 
 // NewAddressAndTime creates a new AddressAndTime object.
 func NewAddressAndTime(e *net.TCPAddr, t time.Time) *AddressAndTime {
 	aat := AddressAndTime{
 		Timestamp: uint32(t.UTC().Unix()),
-		Services:  1,
-		Port:      uint16(e.Port),
+		Capabilities: []capability.Capability{
+			{
+				Type: capability.TCPServer,
+				Data: &capability.Server{Port: uint16(e.Port)},
+			},
+		},
 	}
 	copy(aat.IP[:], e.IP)
 	return &aat
@@ -30,17 +34,15 @@ func NewAddressAndTime(e *net.TCPAddr, t time.Time) *AddressAndTime {
 // DecodeBinary implements Serializable interface.
 func (p *AddressAndTime) DecodeBinary(br *io.BinReader) {
 	p.Timestamp = br.ReadU32LE()
-	p.Services = br.ReadU64LE()
 	br.ReadBytes(p.IP[:])
-	p.Port = br.ReadU16BE()
+	br.ReadArray(&p.Capabilities, capability.MaxCapabilities)
 }
 
 // EncodeBinary implements Serializable interface.
 func (p *AddressAndTime) EncodeBinary(bw *io.BinWriter) {
 	bw.WriteU32LE(p.Timestamp)
-	bw.WriteU64LE(p.Services)
 	bw.WriteBytes(p.IP[:])
-	bw.WriteU16BE(p.Port)
+	bw.WriteArray(p.Capabilities)
 }
 
 // IPPortString makes a string from IP and port specified.
@@ -48,7 +50,14 @@ func (p *AddressAndTime) IPPortString() string {
 	var netip = make(net.IP, 16)
 
 	copy(netip, p.IP[:])
-	port := strconv.Itoa(int(p.Port))
+	var tcpPort uint16
+	for _, cap := range p.Capabilities {
+		if cap.Type == capability.TCPServer {
+			tcpPort = cap.Data.(*capability.Server).Port
+			break
+		}
+	}
+	port := strconv.Itoa(int(tcpPort))
 	return netip.String() + ":" + port
 }
 
diff --git a/pkg/network/payload/address_test.go b/pkg/network/payload/address_test.go
index 30a79d02f4..fb8777aee1 100644
--- a/pkg/network/payload/address_test.go
+++ b/pkg/network/payload/address_test.go
@@ -7,6 +7,7 @@ import (
 	"time"
 
 	"github.com/nspcc-dev/neo-go/pkg/internal/testserdes"
+	"github.com/nspcc-dev/neo-go/pkg/network/capability"
 	"github.com/stretchr/testify/assert"
 )
 
@@ -21,7 +22,11 @@ func TestEncodeDecodeAddress(t *testing.T) {
 	aatip := make(net.IP, 16)
 	copy(aatip, addr.IP[:])
 	assert.Equal(t, e.IP, aatip)
-	assert.Equal(t, e.Port, int(addr.Port))
+	assert.Equal(t, 1, len(addr.Capabilities))
+	assert.Equal(t, capability.Capability{
+		Type: capability.TCPServer,
+		Data: &capability.Server{Port: uint16(e.Port)},
+	}, addr.Capabilities[0])
 
 	testserdes.EncodeDecodeBinary(t, addr, new(AddressAndTime))
 }