Enable IPv6 and IPv4 labeled prefixes to match against prefix set in policy

4 files changed, 52 insertions, 4 deletions

diff --git a/internal/pkg/table/path.go b/internal/pkg/table/path.go
index 805fa63c..15fcfa05 100644
--- a/internal/pkg/table/path.go
+++ b/internal/pkg/table/path.go
@@ -1233,12 +1233,12 @@ func nlriToIPNet(nlri bgp.AddrPrefixInterface) *net.IPNet {
 	case *bgp.LabeledIPAddrPrefix:
 		return &net.IPNet{
 			IP:   net.IP(T.Prefix.To4()),
-			Mask: net.CIDRMask(int(T.Length), 32),
+			Mask: net.CIDRMask(int(T.Length)-T.Labels.Len()*8, 32),
 		}
 	case *bgp.LabeledIPv6AddrPrefix:
 		return &net.IPNet{
-			IP:   net.IP(T.Prefix.To4()),
-			Mask: net.CIDRMask(int(T.Length), 128),
+			IP:   net.IP(T.Prefix.To16()),
+			Mask: net.CIDRMask(int(T.Length)-T.Labels.Len()*8, 128),
 		}
 	}
 	return nil
diff --git a/internal/pkg/table/path_test.go b/internal/pkg/table/path_test.go
index 421502fa..50681a0b 100644
--- a/internal/pkg/table/path_test.go
+++ b/internal/pkg/table/path_test.go
@@ -2,6 +2,7 @@
 package table
 
 import (
+	"net"
 	"testing"
 	"time"
 
@@ -363,3 +364,22 @@ func TestReplaceAS(t *testing.T) {
 	assert.Equal(t, list[2], uint32(10))
 	assert.Equal(t, list[3], uint32(2))
 }
+
+func TestNLRIToIPNet(t *testing.T) {
+	_, n1, _ := net.ParseCIDR("30.30.30.0/24")
+	ipNet := nlriToIPNet(bgp.NewIPAddrPrefix(24, "30.30.30.0"))
+	assert.Equal(t, n1, ipNet)
+
+	_, n2, _ := net.ParseCIDR("2806:106e:19::/48")
+	ipNet = nlriToIPNet(bgp.NewIPv6AddrPrefix(48, "2806:106e:19::"))
+	assert.Equal(t, n2, ipNet)
+
+	labels := bgp.NewMPLSLabelStack(100, 200)
+	_, n3, _ := net.ParseCIDR("30.30.30.0/24")
+	ipNet = nlriToIPNet(bgp.NewLabeledIPAddrPrefix(24, "30.30.30.0", *labels))
+	assert.Equal(t, n3, ipNet)
+
+	_, n4, _ := net.ParseCIDR("2806:106e:19::/48")
+	ipNet = nlriToIPNet(bgp.NewLabeledIPv6AddrPrefix(48, "2806:106e:19::", *labels))
+	assert.Equal(t, n4, ipNet)
+}
diff --git a/internal/pkg/table/policy.go b/internal/pkg/table/policy.go
index d2e4f5ef..f4bcf0a4 100644
--- a/internal/pkg/table/policy.go
+++ b/internal/pkg/table/policy.go
@@ -1439,7 +1439,10 @@ func (c *PrefixCondition) Option() MatchOption {
 // subsequent comparison is skipped if that matches the conditions.
 // If PrefixList's length is zero, return true.
 func (c *PrefixCondition) Evaluate(path *Path, _ *PolicyOptions) bool {
-	if path.GetRouteFamily() != c.set.family {
+	pathAfi, _ := bgp.RouteFamilyToAfiSafi(path.GetRouteFamily())
+	cAfi, _ := bgp.RouteFamilyToAfiSafi(c.set.family)
+
+	if cAfi != pathAfi {
 		return false
 	}
 
diff --git a/internal/pkg/table/policy_test.go b/internal/pkg/table/policy_test.go
index 9f74f32c..fbe9d06b 100644
--- a/internal/pkg/table/policy_test.go
+++ b/internal/pkg/table/policy_test.go
@@ -3024,6 +3024,31 @@ func TestPrefixSetMatchV4withV6Prefix(t *testing.T) {
 	assert.False(t, m.Evaluate(path, nil))
 }
 
+func TestPrefixSetMatchV6LabeledwithV6Prefix(t *testing.T) {
+	p1 := config.Prefix{
+		IpPrefix:        "2806:106e:19::/48",
+		MasklengthRange: "48..48",
+	}
+	ps, err := NewPrefixSet(config.PrefixSet{
+		PrefixSetName: "ps1",
+		PrefixList:    []config.Prefix{p1},
+	})
+	assert.Nil(t, err)
+	m := &PrefixCondition{
+		set: ps,
+	}
+
+	labels := bgp.NewMPLSLabelStack(100, 200)
+	n1 := bgp.NewLabeledIPv6AddrPrefix(48, "2806:106e:19::", *labels)
+	path := NewPath(nil, n1, false, []bgp.PathAttributeInterface{}, time.Now(), false)
+	assert.True(t, m.Evaluate(path, nil))
+
+	labels = bgp.NewMPLSLabelStack(100, 200)
+	n2 := bgp.NewLabeledIPv6AddrPrefix(48, "1806:106e:19::", *labels)
+	path = NewPath(nil, n2, false, []bgp.PathAttributeInterface{}, time.Now(), false)
+	assert.False(t, m.Evaluate(path, nil))
+}
+
 func TestLargeCommunityMatchAction(t *testing.T) {
 	coms := []*bgp.LargeCommunity{
 		&bgp.LargeCommunity{ASN: 100, LocalData1: 100, LocalData2: 100},