packet: improve graceful restart cap usability

Signed-off-by: ISHIDA Wataru <ishida.wataru@lab.ntt.co.jp>

2 files changed, 49 insertions, 23 deletions

diff --git a/packet/bgp.go b/packet/bgp.go
index d47e3b65..c1dd4fcb 100644
--- a/packet/bgp.go
+++ b/packet/bgp.go
@@ -314,13 +314,13 @@ type CapCarryingLabelInfo struct {
 	DefaultParameterCapability
 }
 
-type CapGracefulRestartTuples struct {
+type CapGracefulRestartTuple struct {
 	AFI   uint16
 	SAFI  uint8
 	Flags uint8
 }
 
-func (c CapGracefulRestartTuples) MarshalJSON() ([]byte, error) {
+func (c *CapGracefulRestartTuple) MarshalJSON() ([]byte, error) {
 	return json.Marshal(struct {
 		RouteFamily RouteFamily `json:"route_family"`
 		Flags       uint8       `json:"flags"`
@@ -330,28 +330,38 @@ func (c CapGracefulRestartTuples) MarshalJSON() ([]byte, error) {
 	})
 }
 
-type CapGracefulRestartValue struct {
-	Flags  uint8                      `json:"flags"`
-	Time   uint16                     `json:"time"`
-	Tuples []CapGracefulRestartTuples `json:"tuples"`
+func NewCapGracefulRestartTuple(rf RouteFamily, forward bool) *CapGracefulRestartTuple {
+	afi, safi := RouteFamilyToAfiSafi(rf)
+	flags := 0
+	if forward {
+		flags = 0x80
+	}
+	return &CapGracefulRestartTuple{
+		AFI:   afi,
+		SAFI:  safi,
+		Flags: uint8(flags),
+	}
 }
 
 type CapGracefulRestart struct {
 	DefaultParameterCapability
-	CapValue CapGracefulRestartValue
+	Flags  uint8
+	Time   uint16
+	Tuples []*CapGracefulRestartTuple
 }
 
 func (c *CapGracefulRestart) DecodeFromBytes(data []byte) error {
 	c.DefaultParameterCapability.DecodeFromBytes(data)
 	data = data[2:]
 	restart := binary.BigEndian.Uint16(data[0:2])
-	c.CapValue.Flags = uint8(restart >> 12)
-	c.CapValue.Time = restart & 0xfff
+	c.Flags = uint8(restart >> 12)
+	c.Time = restart & 0xfff
 	data = data[2:]
+	c.Tuples = make([]*CapGracefulRestartTuple, 0, len(data)/4)
 	for len(data) >= 4 {
-		t := CapGracefulRestartTuples{binary.BigEndian.Uint16(data[0:2]),
+		t := &CapGracefulRestartTuple{binary.BigEndian.Uint16(data[0:2]),
 			data[2], data[3]}
-		c.CapValue.Tuples = append(c.CapValue.Tuples, t)
+		c.Tuples = append(c.Tuples, t)
 		data = data[4:]
 	}
 	return nil
@@ -359,8 +369,8 @@ func (c *CapGracefulRestart) DecodeFromBytes(data []byte) error {
 
 func (c *CapGracefulRestart) Serialize() ([]byte, error) {
 	buf := make([]byte, 2)
-	binary.BigEndian.PutUint16(buf[0:], uint16(c.CapValue.Flags)<<12|c.CapValue.Time)
-	for _, t := range c.CapValue.Tuples {
+	binary.BigEndian.PutUint16(buf[0:], uint16(c.Flags)<<12|c.Time)
+	for _, t := range c.Tuples {
 		tbuf := make([]byte, 4)
 		binary.BigEndian.PutUint16(tbuf[0:2], t.AFI)
 		tbuf[2] = t.SAFI
@@ -371,16 +381,32 @@ func (c *CapGracefulRestart) Serialize() ([]byte, error) {
 	return c.DefaultParameterCapability.Serialize()
 }
 
-func NewCapGracefulRestart(flags uint8, time uint16, tuples []CapGracefulRestartTuples) *CapGracefulRestart {
+func (c *CapGracefulRestart) MarshalJSON() ([]byte, error) {
+	return json.Marshal(struct {
+		Code   BGPCapabilityCode          `json:"code"`
+		Flags  uint8                      `json:"flags"`
+		Time   uint16                     `json:"time"`
+		Tuples []*CapGracefulRestartTuple `json:"tuples"`
+	}{
+		Code:   c.Code(),
+		Flags:  c.Flags,
+		Time:   c.Time,
+		Tuples: c.Tuples,
+	})
+}
+
+func NewCapGracefulRestart(restarting bool, time uint16, tuples []*CapGracefulRestartTuple) *CapGracefulRestart {
+	flags := 0
+	if restarting {
+		flags = 0x08
+	}
 	return &CapGracefulRestart{
-		DefaultParameterCapability{
+		DefaultParameterCapability: DefaultParameterCapability{
 			CapCode: BGP_CAP_GRACEFUL_RESTART,
 		},
-		CapGracefulRestartValue{
-			flags,
-			time,
-			tuples,
-		},
+		Flags:  uint8(flags),
+		Time:   time,
+		Tuples: tuples,
 	}
 }
 
diff --git a/packet/bgp_test.go b/packet/bgp_test.go
index d8e41e68..f243a785 100644
--- a/packet/bgp_test.go
+++ b/packet/bgp_test.go
@@ -26,10 +26,10 @@ func open() *BGPMessage {
 		[]ParameterCapabilityInterface{NewCapRouteRefresh()})
 	p2 := NewOptionParameterCapability(
 		[]ParameterCapabilityInterface{NewCapMultiProtocol(RF_IPv4_UC)})
-	g := CapGracefulRestartTuples{4, 2, 3}
+	g := &CapGracefulRestartTuple{4, 2, 3}
 	p3 := NewOptionParameterCapability(
-		[]ParameterCapabilityInterface{NewCapGracefulRestart(2, 100,
-			[]CapGracefulRestartTuples{g})})
+		[]ParameterCapabilityInterface{NewCapGracefulRestart(false, 100,
+			[]*CapGracefulRestartTuple{g})})
 	p4 := NewOptionParameterCapability(
 		[]ParameterCapabilityInterface{NewCapFourOctetASNumber(100000)})
 	p5 := NewOptionParameterCapability(