table: support AfiSafiIn Policy Condition match

2 files changed, 98 insertions, 0 deletions

diff --git a/table/policy.go b/table/policy.go
index 770eb81b..da564abd 100644
--- a/table/policy.go
+++ b/table/policy.go
@@ -157,6 +157,7 @@ const (
 	CONDITION_ROUTE_TYPE
 	CONDITION_LARGE_COMMUNITY
 	CONDITION_NEXT_HOP
+	CONDITION_AFI_SAFI_IN
 )
 
 type ActionType int
@@ -1955,6 +1956,58 @@ func NewRouteTypeCondition(c config.RouteType) (*RouteTypeCondition, error) {
 	}, nil
 }
 
+type AfiSafiInCondition struct {
+	routeFamilies []bgp.RouteFamily
+}
+
+func (c *AfiSafiInCondition) Type() ConditionType {
+	return CONDITION_AFI_SAFI_IN
+}
+
+func (c *AfiSafiInCondition) Evaluate(path *Path, _ *PolicyOptions) bool {
+	for _, rf := range c.routeFamilies {
+		if path.GetRouteFamily() == rf {
+			return true
+		}
+	}
+	return false
+}
+
+func (c *AfiSafiInCondition) Set() DefinedSet {
+	return nil
+}
+
+func (c *AfiSafiInCondition) Name() string { return "" }
+
+func (c *AfiSafiInCondition) String() string {
+	tmp := make([]string, 0, len(c.routeFamilies))
+	for _, afiSafi := range c.routeFamilies {
+		tmp = append(tmp, afiSafi.String())
+	}
+	return strings.Join(tmp, " ")
+}
+
+func NewAfiSafiInCondition(afiSafInConfig []config.AfiSafiType) (*AfiSafiInCondition, error) {
+	if afiSafInConfig == nil {
+		return nil, nil
+	}
+
+	routeFamilies := make([]bgp.RouteFamily, 0, len(afiSafInConfig))
+	for _, afiSafiValue := range afiSafInConfig {
+		if err := afiSafiValue.Validate(); err != nil {
+			return nil, err
+		}
+		rf, err := bgp.GetRouteFamily(string(afiSafiValue))
+		if err != nil {
+			return nil, err
+		}
+		routeFamilies = append(routeFamilies, rf)
+	}
+	return &AfiSafiInCondition{
+		routeFamilies: routeFamilies,
+	}, nil
+}
+
 type Action interface {
 	Type() ActionType
 	Apply(*Path, *PolicyOptions) *Path
@@ -2652,6 +2705,13 @@ func (s *Statement) ToConfig() *config.Statement {
 				case *RouteTypeCondition:
 					v := c.(*RouteTypeCondition)
 					cond.BgpConditions.RouteType = v.typ
+				case *AfiSafiInCondition:
+					v := c.(*AfiSafiInCondition)
+					res := make([]config.AfiSafiType, 0, len(v.routeFamilies))
+					for _, rf := range v.routeFamilies {
+						res = append(res, config.AfiSafiType(rf.String()))
+					}
+					cond.BgpConditions.AfiSafiInList = res
 				}
 			}
 			return cond
@@ -2853,6 +2913,9 @@ func NewStatement(c config.Statement) (*Statement, error) {
 		func() (Condition, error) {
 			return NewNextHopCondition(c.Conditions.BgpConditions.NextHopInList)
 		},
+		func() (Condition, error) {
+			return NewAfiSafiInCondition(c.Conditions.BgpConditions.AfiSafiInList)
+		},
 	}
 	cs = make([]Condition, 0, len(cfs))
 	for _, f := range cfs {
@@ -3227,6 +3290,7 @@ func (r *RoutingPolicy) validateCondition(v Condition) (err error) {
 			c.set = i.(*LargeCommunitySet)
 		}
 	case CONDITION_NEXT_HOP:
+	case CONDITION_AFI_SAFI_IN:
 	case CONDITION_AS_PATH_LENGTH:
 	case CONDITION_RPKI:
 	}
diff --git a/table/policy_test.go b/table/policy_test.go
index 9e990dc9..007936c1 100644
--- a/table/policy_test.go
+++ b/table/policy_test.go
@@ -3058,6 +3058,40 @@ func TestLargeCommunityMatchAction(t *testing.T) {
 	assert.Equal(t, m.Evaluate(p, nil), true)
 }
 
+func TestAfiSafiInMatchPath(t *testing.T) {
+	condition, err := NewAfiSafiInCondition([]config.AfiSafiType{config.AFI_SAFI_TYPE_L3VPN_IPV4_UNICAST, config.AFI_SAFI_TYPE_L3VPN_IPV6_UNICAST})
+
+	rtExtCom, err := bgp.ParseExtendedCommunity(bgp.EC_SUBTYPE_ROUTE_TARGET, "100:100")
+	assert.NoError(t, err)
+
+	prefixVPNv4 := bgp.NewLabeledVPNIPAddrPrefix(0, "1.1.1.0/24", *bgp.NewMPLSLabelStack(), bgp.NewRouteDistinguisherTwoOctetAS(100, 100))
+	prefixVPNv6 := bgp.NewLabeledVPNIPv6AddrPrefix(0, "2001:0db8:85a3:0000:0000:8a2e:0370:7334", *bgp.NewMPLSLabelStack(), bgp.NewRouteDistinguisherTwoOctetAS(200, 200))
+	prefixRTC := bgp.NewRouteTargetMembershipNLRI(100, nil)
+	prefixv4 := bgp.NewIPAddrPrefix(0, "1.1.1.0/24")
+	prefixv6 := bgp.NewIPv6AddrPrefix(0, "2001:0db8:85a3:0000:0000:8a2e:0370:7334")
+
+	pathVPNv4 := NewPath(nil, prefixVPNv4, false, []bgp.PathAttributeInterface{bgp.NewPathAttributeExtendedCommunities([]bgp.ExtendedCommunityInterface{rtExtCom})}, time.Time{}, false)
+	pathVPNv6 := NewPath(nil, prefixVPNv6, false, []bgp.PathAttributeInterface{bgp.NewPathAttributeExtendedCommunities([]bgp.ExtendedCommunityInterface{rtExtCom})}, time.Time{}, false)
+	pathv4 := NewPath(nil, prefixv4, false, []bgp.PathAttributeInterface{}, time.Time{}, false)
+	pathv6 := NewPath(nil, prefixv6, false, []bgp.PathAttributeInterface{}, time.Time{}, false)
+	pathRTC := NewPath(nil, prefixRTC, false, []bgp.PathAttributeInterface{}, time.Time{}, false)
+
+	type Entry struct {
+		path        *Path
+		shouldMatch bool
+	}
+
+	for _, entry := range []Entry{
+		{pathVPNv4, true},
+		{pathVPNv6, true},
+		{pathv4, false},
+		{pathv6, false},
+		{pathRTC, false},
+	} {
+		assert.Equal(t, condition.Evaluate(entry.path, nil), entry.shouldMatch)
+	}
+}
+
 func TestMultipleStatementPolicy(t *testing.T) {
 	r := NewRoutingPolicy()
 	rp := config.RoutingPolicy{