addpath serialization clean up

protocols/bgp/packet/mp_reach_nlri.go

@@ -28,11 +28,7 @@ func (n *MultiProtocolReachNLRI) serialize(buf *bytes.Buffer, opt *EncodeOptions
 	tempBuf.WriteByte(0) // RESERVED
 
 	for cur := n.NLRI; cur != nil; cur = cur.Next {
-		if opt.UseAddPath {
-			n.NLRI.serializeAddPath(tempBuf)
-		} else {
-			n.NLRI.serialize(tempBuf)
-		}
+		cur.serialize(tempBuf, opt.UseAddPath)
 	}
 
 	buf.Write(tempBuf.Bytes())

protocols/bgp/packet/mp_unreach_nlri.go

@@ -21,10 +21,7 @@ func (n *MultiProtocolUnreachNLRI) serialize(buf *bytes.Buffer, opt *EncodeOptio
 	tempBuf.WriteByte(n.SAFI)
 
 	for cur := n.NLRI; cur != nil; cur = cur.Next {
-		if opt.UseAddPath {
-			tempBuf.Write(convert.Uint32Byte(cur.PathIdentifier))
-		}
-		tempBuf.Write(serializePrefix(cur.Prefix))
+		cur.serialize(tempBuf, opt.UseAddPath)
 	}
 
 	buf.Write(tempBuf.Bytes())

protocols/bgp/packet/nlri.go

@@ -89,20 +89,22 @@ func decodeNLRI(buf *bytes.Buffer, afi uint16, addPath bool) (*NLRI, uint8, erro
 	return nlri, consumed, nil
 }
 
-func (n *NLRI) serialize(buf *bytes.Buffer) uint8 {
-	buf.WriteByte(n.Prefix.Pfxlen())
-	b := n.Prefix.Addr().Bytes()
+func (n *NLRI) serialize(buf *bytes.Buffer, addPath bool) uint8 {
+	numBytes := uint8(0)
 
-	nBytes := BytesInAddr(n.Prefix.Pfxlen())
-	buf.Write(b[:nBytes])
+	if addPath {
+		buf.Write(convert.Uint32Byte(n.PathIdentifier))
+		numBytes += 4
+	}
 
-	return nBytes + 1
-}
+	buf.WriteByte(n.Prefix.Pfxlen())
+	numBytes++
 
-func (n *NLRI) serializeAddPath(buf *bytes.Buffer) uint8 {
-	buf.Write(convert.Uint32Byte(n.PathIdentifier))
+	pfxNumBytes := BytesInAddr(n.Prefix.Pfxlen())
+	buf.Write(n.Prefix.Addr().Bytes()[:pfxNumBytes])
+	numBytes += pfxNumBytes
 
-	return uint8(n.serialize(buf) + 4)
+	return numBytes
 }
 
 // BytesInAddr gets the amount of bytes needed to encode an NLRI of prefix length pfxlen

protocols/bgp/packet/nlri_test.go

@@ -220,6 +220,7 @@ func TestNLRISerialize(t *testing.T) {
 	tests := []struct {
 		name     string
 		nlri     *NLRI
+		addPath  bool
 		expected []byte
 	}{
 		{
@@ -243,51 +244,38 @@ func TestNLRISerialize(t *testing.T) {
 			},
 			expected: []byte{17, 100, 200, 128},
 		},
-	}
-
-	for _, test := range tests {
-		buf := bytes.NewBuffer(nil)
-		test.nlri.serialize(buf)
-		res := buf.Bytes()
-		assert.Equal(t, test.expected, res)
-	}
-}
-
-func TestNLRIAddPathSerialize(t *testing.T) {
-	tests := []struct {
-		name     string
-		nlri     *NLRI
-		expected []byte
-	}{
 		{
-			name: "Test #1",
+			name: "with add-path #1",
 			nlri: &NLRI{
 				PathIdentifier: 100,
 				Prefix:         bnet.NewPfx(bnet.IPv4FromOctets(1, 2, 3, 0), 25),
 			},
+			addPath:  true,
 			expected: []byte{0, 0, 0, 100, 25, 1, 2, 3, 0},
 		},
 		{
-			name: "Test #2",
+			name: "with add-path #2",
 			nlri: &NLRI{
 				PathIdentifier: 100,
 				Prefix:         bnet.NewPfx(bnet.IPv4FromOctets(1, 2, 3, 0), 24),
 			},
+			addPath:  true,
 			expected: []byte{0, 0, 0, 100, 24, 1, 2, 3},
 		},
 		{
-			name: "Test #3",
+			name: "with add-path #3",
 			nlri: &NLRI{
 				PathIdentifier: 100,
 				Prefix:         bnet.NewPfx(bnet.IPv4FromOctets(100, 200, 128, 0), 17),
 			},
+			addPath:  true,
 			expected: []byte{0, 0, 0, 100, 17, 100, 200, 128},
 		},
 	}
 
 	for _, test := range tests {
 		buf := bytes.NewBuffer(nil)
-		test.nlri.serializeAddPath(buf)
+		test.nlri.serialize(buf, test.addPath)
 		res := buf.Bytes()
 		assert.Equal(t, test.expected, res)
 	}

protocols/bgp/packet/update.go

@@ -18,18 +18,11 @@ type BGPUpdate struct {
 // SerializeUpdate serializes an BGPUpdate to wire format
 func (b *BGPUpdate) SerializeUpdate(opt *EncodeOptions) ([]byte, error) {
 	budget := MaxLen - MinLen
-	nlriLen := 0
 	buf := bytes.NewBuffer(nil)
 
 	withdrawBuf := bytes.NewBuffer(nil)
 	for withdraw := b.WithdrawnRoutes; withdraw != nil; withdraw = withdraw.Next {
-		if opt.UseAddPath {
-			nlriLen = int(withdraw.serializeAddPath(withdrawBuf))
-		} else {
-			nlriLen = int(withdraw.serialize(withdrawBuf))
-		}
-
-		budget -= nlriLen
+		budget -= int(withdraw.serialize(withdrawBuf, opt.UseAddPath))
 		if budget < 0 {
 			return nil, fmt.Errorf("update too long")
 		}
@@ -46,13 +39,7 @@ func (b *BGPUpdate) SerializeUpdate(opt *EncodeOptions) ([]byte, error) {
 
 	nlriBuf := bytes.NewBuffer(nil)
 	for nlri := b.NLRI; nlri != nil; nlri = nlri.Next {
-		if opt.UseAddPath {
-			nlriLen = int(nlri.serializeAddPath(nlriBuf))
-		} else {
-			nlriLen = int(nlri.serialize(nlriBuf))
-		}
-
-		budget -= nlriLen
+		budget -= int(nlri.serialize(nlriBuf, opt.UseAddPath))
 		if budget < 0 {
 			return nil, fmt.Errorf("update too long")
 		}
@@ -92,8 +79,7 @@ func (b *BGPUpdate) SerializeUpdateAddPath(opt *EncodeOptions) ([]byte, error) {
 
 	withdrawBuf := bytes.NewBuffer(nil)
 	for withdraw := b.WithdrawnRoutes; withdraw != nil; withdraw = withdraw.Next {
-		nlriLen := int(withdraw.serialize(withdrawBuf))
-		budget -= nlriLen
+		budget -= int(withdraw.serialize(withdrawBuf, opt.UseAddPath))
 		if budget < 0 {
 			return nil, fmt.Errorf("update too long")
 		}
@@ -110,8 +96,7 @@ func (b *BGPUpdate) SerializeUpdateAddPath(opt *EncodeOptions) ([]byte, error) {
 
 	nlriBuf := bytes.NewBuffer(nil)
 	for nlri := b.NLRI; nlri != nil; nlri = nlri.Next {
-		nlriLen := int(nlri.serialize(nlriBuf))
-		budget -= nlriLen
+		budget -= int(nlri.serialize(nlriBuf, opt.UseAddPath))
 		if budget < 0 {
 			return nil, fmt.Errorf("update too long")
 		}