Support vrfs in zapi multipath

3 files changed, 49 insertions, 30 deletions

diff --git a/pkg/server/server.go b/pkg/server/server.go
index c5dd18f3..4fbb76ff 100644
--- a/pkg/server/server.go
+++ b/pkg/server/server.go
@@ -698,26 +698,39 @@ func clonePathList(pathList []*table.Path) []*table.Path {
 	return l
 }
 
+func (s *BgpServer) setPathVrfIdMap(paths []*table.Path, m map[uint32]bool) {
+	for _, p := range paths {
+		switch p.GetRouteFamily() {
+		case bgp.RF_IPv4_VPN, bgp.RF_IPv6_VPN:
+			for _, vrf := range s.globalRib.Vrfs {
+				if vrf.Id != 0 && table.CanImportToVrf(vrf, p) {
+					m[uint32(vrf.Id)] = true
+				}
+			}
+		default:
+			m[zebra.VRF_DEFAULT] = true
+		}
+	}
+}
+
+// Note: the destination would be the same for all the paths passed here
+// The wather (only zapi) needs a unique list of vrf IDs
 func (s *BgpServer) notifyBestWatcher(best []*table.Path, multipath [][]*table.Path) {
 	if table.SelectionOptions.DisableBestPathSelection {
 		// Note: If best path selection disabled, no best path to notify.
 		return
 	}
+	m := make(map[uint32]bool)
 	clonedM := make([][]*table.Path, len(multipath))
 	for i, pathList := range multipath {
 		clonedM[i] = clonePathList(pathList)
+		if table.UseMultiplePaths.Enabled {
+			s.setPathVrfIdMap(clonedM[i], m)
+		}
 	}
 	clonedB := clonePathList(best)
-	m := make(map[uint32]bool)
-	for _, p := range clonedB {
-		switch p.GetRouteFamily() {
-		case bgp.RF_IPv4_VPN, bgp.RF_IPv6_VPN:
-			for _, vrf := range s.globalRib.Vrfs {
-				if vrf.Id != 0 && table.CanImportToVrf(vrf, p) {
-					m[uint32(vrf.Id)] = true
-				}
-			}
-		}
+	if !table.UseMultiplePaths.Enabled {
+		s.setPathVrfIdMap(clonedB, m)
 	}
 	w := &watchEventBestPath{PathList: clonedB, MultiPathList: clonedM}
 	if len(m) > 0 {
diff --git a/pkg/server/server_test.go b/pkg/server/server_test.go
index a390ce58..3cbf53a9 100644
--- a/pkg/server/server_test.go
+++ b/pkg/server/server_test.go
@@ -533,7 +533,8 @@ func TestMonitor(test *testing.T) {
 	assert.Equal(1, len(b.PathList))
 	assert.Equal("10.0.0.0/24", b.PathList[0].GetNlri().String())
 	assert.False(b.PathList[0].IsWithdraw)
-	assert.Equal(0, len(b.Vrf))
+	assert.Equal(1, len(b.Vrf))
+	assert.True(b.Vrf[0])
 
 	// Withdraws the previous route.
 	// NOTE: Withdraw should not require any path attribute.
@@ -545,7 +546,8 @@ func TestMonitor(test *testing.T) {
 	assert.Equal(1, len(b.PathList))
 	assert.Equal("10.0.0.0/24", b.PathList[0].GetNlri().String())
 	assert.True(b.PathList[0].IsWithdraw)
-	assert.Equal(0, len(b.Vrf))
+	assert.Equal(1, len(b.Vrf))
+	assert.True(b.Vrf[0])
 
 	// Stops the watcher still having an item.
 	w.Stop()
@@ -590,7 +592,7 @@ func TestMonitor(test *testing.T) {
 	assert.False(u.PathList[0].IsWithdraw)
 
 	// Withdraws the previous route.
-	// NOTE: Withdow should not require any path attribute.
+	// NOTE: Withdraw should not require any path attribute.
 	if err := t.addPathList("", []*table.Path{table.NewPath(nil, bgp.NewIPAddrPrefix(24, "10.2.0.0"), true, nil, time.Now(), false)}); err != nil {
 		log.Fatal(err)
 	}
@@ -620,7 +622,19 @@ func TestMonitor(test *testing.T) {
 	assert.True(b.Vrf[1])
 	assert.True(b.Vrf[2])
 
-	// Stops the watcher still having an item.
+	// Withdraw the route
+	if err := s.addPathList("vrf1", []*table.Path{table.NewPath(nil, bgp.NewIPAddrPrefix(24, "10.0.0.0"), true, attrs, time.Now(), false)}); err != nil {
+		log.Fatal(err)
+	}
+	ev = <-w.Event()
+	b = ev.(*watchEventBestPath)
+	assert.Equal(1, len(b.PathList))
+	assert.Equal("111:111:10.0.0.0/24", b.PathList[0].GetNlri().String())
+	assert.True(b.PathList[0].IsWithdraw)
+	assert.Equal(2, len(b.Vrf))
+	assert.True(b.Vrf[1])
+	assert.True(b.Vrf[2])
+
 	w.Stop()
 }
 
diff --git a/pkg/server/zclient.go b/pkg/server/zclient.go
index 390628f0..c002851a 100644
--- a/pkg/server/zclient.go
+++ b/pkg/server/zclient.go
@@ -421,27 +421,19 @@ func (z *zebraClient) loop() {
 				if table.UseMultiplePaths.Enabled {
 					for _, paths := range msg.MultiPathList {
 						z.updatePathByNexthopCache(paths)
-						if body, isWithdraw := newIPRouteBody(paths, 0, z); body != nil {
-							z.client.SendIPRoute(0, body, isWithdraw)
-						}
-						if body := newNexthopRegisterBody(paths, z.nexthopCache); body != nil {
-							z.client.SendNexthopRegister(0, body, false)
+						for i := range msg.Vrf {
+							if body, isWithdraw := newIPRouteBody(paths, i, z); body != nil {
+								z.client.SendIPRoute(i, body, isWithdraw)
+							}
+							if body := newNexthopRegisterBody(paths, z.nexthopCache); body != nil {
+								z.client.SendNexthopRegister(i, body, false)
+							}
 						}
 					}
 				} else {
 					z.updatePathByNexthopCache(msg.PathList)
 					for _, path := range msg.PathList {
-						vrfs := []uint32{0}
-						if msg.Vrf != nil {
-							for vrfId := range msg.Vrf {
-								vrfs = append(vrfs, vrfId)
-							}
-						}
-						for _, i := range vrfs {
-							routeFamily := path.GetRouteFamily()
-							if i == zebra.VRF_DEFAULT && (routeFamily == bgp.RF_IPv4_VPN || routeFamily == bgp.RF_IPv6_VPN) {
-								continue
-							}
+						for i := range msg.Vrf {
 							if body, isWithdraw := newIPRouteBody([]*table.Path{path}, i, z); body != nil {
 								err := z.client.SendIPRoute(i, body, isWithdraw)
 								if err != nil {