policy: fix bug of handling multiple prefix-match with same IP prefix.

fix handling of prefix-match configuration like below

[[defined-sets.prefix-sets]]
prefix-set-name = "ps1"
[[defined-sets.prefix-sets.prefix-list]]
ip-prefix = "0.0.0.0/0"
masklength-range = "0..7"
[[defined-sets.prefix-sets.prefix-list]]
ip-prefix = "0.0.0.0/0"
masklength-range = "25..32"

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

2 files changed, 120 insertions, 11 deletions

diff --git a/table/policy.go b/table/policy.go
index 74ced9e1..2122e508 100644
--- a/table/policy.go
+++ b/table/policy.go
@@ -286,8 +286,15 @@ func (lhs *PrefixSet) Append(arg DefinedSet) error {
 	if !ok {
 		return fmt.Errorf("type cast failed")
 	}
-	rhs.tree.Walk(func(s string, v interface{}) bool {
-		lhs.tree.Insert(s, v)
+	rhs.tree.Walk(func(key string, v interface{}) bool {
+		w, ok := lhs.tree.Get(key)
+		if ok {
+			r := v.([]*Prefix)
+			l := w.([]*Prefix)
+			lhs.tree.Insert(key, append(l, r...))
+		} else {
+			lhs.tree.Insert(key, v)
+		}
 		return false
 	})
 	return nil
@@ -298,8 +305,31 @@ func (lhs *PrefixSet) Remove(arg DefinedSet) error {
 	if !ok {
 		return fmt.Errorf("type cast failed")
 	}
-	rhs.tree.Walk(func(s string, v interface{}) bool {
-		lhs.tree.Delete(s)
+	rhs.tree.Walk(func(key string, v interface{}) bool {
+		w, ok := lhs.tree.Get(key)
+		if !ok {
+			return false
+		}
+		r := v.([]*Prefix)
+		l := w.([]*Prefix)
+		new := make([]*Prefix, 0, len(l))
+		for _, lp := range l {
+			delete := false
+			for _, rp := range r {
+				if lp.Equal(rp) {
+					delete = true
+					break
+				}
+			}
+			if !delete {
+				new = append(new, lp)
+			}
+		}
+		if len(new) == 0 {
+			lhs.tree.Delete(key)
+		} else {
+			lhs.tree.Insert(key, new)
+		}
 		return false
 	})
 	return nil
@@ -317,8 +347,10 @@ func (lhs *PrefixSet) Replace(arg DefinedSet) error {
 func (s *PrefixSet) ToConfig() *config.PrefixSet {
 	list := make([]config.Prefix, 0, s.tree.Len())
 	s.tree.Walk(func(s string, v interface{}) bool {
-		p := v.(*Prefix)
-		list = append(list, config.Prefix{IpPrefix: p.Prefix.String(), MasklengthRange: fmt.Sprintf("%d..%d", p.MasklengthRangeMin, p.MasklengthRangeMax)})
+		ps := v.([]*Prefix)
+		for _, p := range ps {
+			list = append(list, config.Prefix{IpPrefix: p.Prefix.String(), MasklengthRange: fmt.Sprintf("%d..%d", p.MasklengthRangeMin, p.MasklengthRangeMax)})
+		}
 		return false
 	})
 	return &config.PrefixSet{
@@ -333,7 +365,14 @@ func NewPrefixSetFromApiStruct(name string, prefixes []*Prefix) (*PrefixSet, err
 	}
 	tree := radix.New()
 	for _, x := range prefixes {
-		tree.Insert(CidrToRadixkey(x.Prefix.String()), x)
+		key := CidrToRadixkey(x.Prefix.String())
+		d, ok := tree.Get(key)
+		if ok {
+			ps := d.([]*Prefix)
+			tree.Insert(key, append(ps, x))
+		} else {
+			tree.Insert(key, []*Prefix{x})
+		}
 	}
 	return &PrefixSet{
 		name: name,
@@ -355,7 +394,14 @@ func NewPrefixSet(c config.PrefixSet) (*PrefixSet, error) {
 		if err != nil {
 			return nil, err
 		}
-		tree.Insert(CidrToRadixkey(y.Prefix.String()), y)
+		key := CidrToRadixkey(y.Prefix.String())
+		d, ok := tree.Get(key)
+		if ok {
+			ps := d.([]*Prefix)
+			tree.Insert(key, append(ps, y))
+		} else {
+			tree.Insert(key, []*Prefix{y})
+		}
 	}
 	return &PrefixSet{
 		name: name,
@@ -970,9 +1016,14 @@ func (c *PrefixCondition) Evaluate(path *Path, _ *PolicyOptions) bool {
 	}
 
 	result := false
-	_, p, ok := c.set.tree.LongestPrefix(key)
-	if ok && p.(*Prefix).MasklengthRangeMin <= masklen && masklen <= p.(*Prefix).MasklengthRangeMax {
-		result = true
+	_, ps, ok := c.set.tree.LongestPrefix(key)
+	if ok {
+		for _, p := range ps.([]*Prefix) {
+			if p.MasklengthRangeMin <= masklen && masklen <= p.MasklengthRangeMax {
+				result = true
+				break
+			}
+		}
 	}
 
 	if c.option == MATCH_OPTION_INVERT {
diff --git a/table/policy_test.go b/table/policy_test.go
index d9767ad4..cfafb1ed 100644
--- a/table/policy_test.go
+++ b/table/policy_test.go
@@ -2815,3 +2815,61 @@ func createAs4Value(s string) uint32 {
 	lower, _ := strconv.Atoi(v[1])
 	return uint32(upper)<<16 + uint32(lower)
 }
+
+func TestPrefixSetMatch(t *testing.T) {
+	p1 := config.Prefix{
+		IpPrefix:        "0.0.0.0/0",
+		MasklengthRange: "0..7",
+	}
+	p2 := config.Prefix{
+		IpPrefix:        "0.0.0.0/0",
+		MasklengthRange: "25..32",
+	}
+	ps, err := NewPrefixSet(config.PrefixSet{
+		PrefixSetName: "ps1",
+		PrefixList:    []config.Prefix{p1, p2},
+	})
+	assert.Nil(t, err)
+	m := &PrefixCondition{
+		set: ps,
+	}
+
+	path := NewPath(nil, bgp.NewIPAddrPrefix(6, "0.0.0.0"), false, []bgp.PathAttributeInterface{}, time.Now(), false)
+	assert.True(t, m.Evaluate(path, nil))
+
+	path = NewPath(nil, bgp.NewIPAddrPrefix(10, "0.0.0.0"), false, []bgp.PathAttributeInterface{}, time.Now(), false)
+	assert.False(t, m.Evaluate(path, nil))
+
+	path = NewPath(nil, bgp.NewIPAddrPrefix(25, "0.0.0.0"), false, []bgp.PathAttributeInterface{}, time.Now(), false)
+	assert.True(t, m.Evaluate(path, nil))
+
+	path = NewPath(nil, bgp.NewIPAddrPrefix(30, "0.0.0.0"), false, []bgp.PathAttributeInterface{}, time.Now(), false)
+	assert.True(t, m.Evaluate(path, nil))
+
+	p3 := config.Prefix{
+		IpPrefix:        "0.0.0.0/0",
+		MasklengthRange: "9..10",
+	}
+	ps2, err := NewPrefixSet(config.PrefixSet{
+		PrefixSetName: "ps2",
+		PrefixList:    []config.Prefix{p3},
+	})
+	assert.Nil(t, err)
+	err = ps.Append(ps2)
+	assert.Nil(t, err)
+
+	path = NewPath(nil, bgp.NewIPAddrPrefix(10, "0.0.0.0"), false, []bgp.PathAttributeInterface{}, time.Now(), false)
+	assert.True(t, m.Evaluate(path, nil))
+
+	ps3, err := NewPrefixSet(config.PrefixSet{
+		PrefixSetName: "ps3",
+		PrefixList:    []config.Prefix{p1},
+	})
+	assert.Nil(t, err)
+	err = ps.Remove(ps3)
+	assert.Nil(t, err)
+
+	path = NewPath(nil, bgp.NewIPAddrPrefix(6, "0.0.0.0"), false, []bgp.PathAttributeInterface{}, time.Now(), false)
+	assert.False(t, m.Evaluate(path, nil))
+
+}