packet: fix PathAttributeAggregator

PathAttributeAggregator can handle both two-bytes and four-bytes AS
number. So its serializer needs to know which way should be used.

Signed-off-by: FUJITA Tomonori <fujita.tomonori@lab.ntt.co.jp>

2 files changed, 22 insertions, 6 deletions

diff --git a/packet/bgp.go b/packet/bgp.go
index fcbb3bd1..681c7a0b 100644
--- a/packet/bgp.go
+++ b/packet/bgp.go
@@ -20,6 +20,7 @@ import (
 	"fmt"
 	"math"
 	"net"
+	"reflect"
 )
 
 // move somewhere else
@@ -1395,6 +1396,7 @@ func NewPathAttributeAtomicAggregate() *PathAttributeAtomicAggregate {
 
 type PathAttributeAggregatorParam struct {
 	AS      uint32
+	askind  reflect.Kind
 	Address net.IP
 }
 
@@ -1408,29 +1410,42 @@ func (p *PathAttributeAggregator) DecodeFromBytes(data []byte) error {
 	if len(p.PathAttribute.Value) == 6 {
 		p.Value.AS = uint32(binary.BigEndian.Uint16(p.PathAttribute.Value[0:2]))
 		p.Value.Address = p.PathAttribute.Value[2:]
+		p.Value.askind = reflect.Uint16
 	} else {
 		p.Value.AS = binary.BigEndian.Uint32(p.PathAttribute.Value[0:4])
 		p.Value.Address = p.PathAttribute.Value[4:]
+		p.Value.askind = reflect.Uint32
 	}
 	return nil
 }
 
 func (p *PathAttributeAggregator) Serialize() ([]byte, error) {
-	buf := make([]byte, 6)
-	binary.BigEndian.PutUint16(buf, uint16(p.Value.AS))
-	copy(buf[2:], p.Value.Address)
+	var buf []byte
+	switch p.Value.askind {
+	case reflect.Uint16:
+		buf = make([]byte, 6)
+		binary.BigEndian.PutUint16(buf, uint16(p.Value.AS))
+		copy(buf[2:], p.Value.Address)
+	case reflect.Uint32:
+		buf = make([]byte, 8)
+		binary.BigEndian.PutUint32(buf, p.Value.AS)
+		copy(buf[4:], p.Value.Address)
+	}
+
 	p.PathAttribute.Value = buf
 	return p.PathAttribute.Serialize()
 }
 
-func NewPathAttributeAggregator(as uint16, address string) *PathAttributeAggregator {
+func NewPathAttributeAggregator(as interface{}, address string) *PathAttributeAggregator {
+	v := reflect.ValueOf(as)
 	return &PathAttributeAggregator{
 		PathAttribute: PathAttribute{
 			Flags: BGP_ATTR_FLAG_OPTIONAL | BGP_ATTR_FLAG_TRANSITIVE,
 			Type:  BGP_ATTR_TYPE_AGGREGATOR,
 		},
 		Value: PathAttributeAggregatorParam{
-			AS:      uint32(as),
+			AS:      uint32(v.Uint()),
+			askind:  v.Kind(),
 			Address: net.ParseIP(address).To4(),
 		},
 	}
diff --git a/packet/bgp_test.go b/packet/bgp_test.go
index 59734ba4..e09ffc44 100644
--- a/packet/bgp_test.go
+++ b/packet/bgp_test.go
@@ -89,7 +89,8 @@ func update() *BGPMessage {
 		NewPathAttributeMultiExitDisc(1 << 20),
 		NewPathAttributeLocalPref(1 << 22),
 		NewPathAttributeAtomicAggregate(),
-		NewPathAttributeAggregator(30002, "129.0.2.99"),
+		NewPathAttributeAggregator(uint32(30002), "129.0.2.99"),
+		NewPathAttributeAggregator(uint32(300020), "129.0.2.99"),
 		NewPathAttributeCommunities([]uint32{1, 3}),
 		NewPathAttributeOriginatorId("10.10.0.1"),
 		NewPathAttributeClusterList([]string{"10.10.0.2", "10.10.0.3"}),