Cleanup IP version selection

protocols/bgp/server/bmp_router.go

@@ -234,12 +234,8 @@ func (r *router) processPeerUpNotification(msg *bmppkt.PeerUpNotification) error
 	}
 
 	addrLen := net.IPv4len
-	for i := 0; i < net.IPv6len-net.IPv4len; i++ {
-		if msg.PerPeerHeader.PeerAddress[i] == 0 {
-			continue
-		}
+	if msg.PerPeerHeader.GetIPVersion() == 6 {
 		addrLen = net.IPv6len
-		break
 	}
 
 	// bnet.IPFromBytes can only fail if length of argument is not 4 or 16. However, length is ensured here.

protocols/bmp/packet/per_peer_header.go

@@ -57,3 +57,11 @@ func decodePerPeerHeader(buf *bytes.Buffer) (*PerPeerHeader, error) {
 
 	return p, nil
 }
+
+// GetIPVersion gets the IP version of the BGP session
+func (p *PerPeerHeader) GetIPVersion() uint8 {
+	if p.PeerFlags>>7 == 1 {
+		return 6
+	}
+	return 4
+}

protocols/bmp/packet/per_peer_header_test.go

@@ -113,4 +113,47 @@ func TestDecodePerPeerHeader(t *testing.T) {
 
 		assert.Equalf(t, test.expected, p, "Test %q", test.name)
 	}
+
+}
+
+func TestGetIPVersion(t *testing.T) {
+	tests := []struct {
+		name     string
+		p        *PerPeerHeader
+		expected uint8
+	}{
+		{
+			name: "IPv4",
+			p: &PerPeerHeader{
+				PeerFlags: 0,
+			},
+			expected: 4,
+		},
+		{
+			name: "IPv4 #2",
+			p: &PerPeerHeader{
+				PeerFlags: 127,
+			},
+			expected: 4,
+		},
+		{
+			name: "IPv6",
+			p: &PerPeerHeader{
+				PeerFlags: 128,
+			},
+			expected: 6,
+		},
+		{
+			name: "IPv6 #2",
+			p: &PerPeerHeader{
+				PeerFlags: 129,
+			},
+			expected: 6,
+		},
+	}
+
+	for _, test := range tests {
+		v := test.p.GetIPVersion()
+		assert.Equal(t, test.expected, v)
+	}
 }