add support GetTable method for vrfs #2235

3 files changed, 78 insertions, 15 deletions

diff --git a/internal/pkg/table/table.go b/internal/pkg/table/table.go
index bfcf2059..2e0dc94a 100644
--- a/internal/pkg/table/table.go
+++ b/internal/pkg/table/table.go
@@ -434,14 +434,41 @@ type TableInfo struct {
 	NumAccepted    int
 }
 
-func (t *Table) Info(id string, as uint32) *TableInfo {
+type TableInfoOptions struct {
+	ID  string
+	AS  uint32
+	VRF *Vrf
+}
+
+func (t *Table) Info(option ...TableInfoOptions) *TableInfo {
 	var numD, numP int
+
+	id := GLOBAL_RIB_NAME
+	var vrf *Vrf
+	as := uint32(0)
+
+	for _, o := range option {
+		if o.ID != "" {
+			id = o.ID
+		}
+		if o.VRF != nil {
+			vrf = o.VRF
+		}
+		as = o.AS
+	}
+
 	for _, d := range t.destinations {
-		n := 0
-		if id == GLOBAL_RIB_NAME {
-			n = len(d.knownPathList)
-		} else {
-			n = len(d.GetKnownPathList(id, as))
+		paths := d.GetKnownPathList(id, as)
+		n := len(paths)
+
+		if vrf != nil {
+			ps := make([]*Path, 0, len(paths))
+			for _, p := range paths {
+				if CanImportToVrf(vrf, p) {
+					ps = append(ps, p.ToLocal())
+				}
+			}
+			n = len(ps)
 		}
 		if n != 0 {
 			numD++
diff --git a/internal/pkg/table/table_manager.go b/internal/pkg/table/table_manager.go
index aa0cd0fc..b16f135b 100644
--- a/internal/pkg/table/table_manager.go
+++ b/internal/pkg/table/table_manager.go
@@ -354,11 +354,3 @@ func (manager *TableManager) GetDestination(path *Path) *Destination {
 	}
 	return t.GetDestination(path.GetNlri())
 }
-
-func (manager *TableManager) TableInfo(id string, as uint32, family bgp.RouteFamily) (*TableInfo, error) {
-	t, ok := manager.Tables[family]
-	if !ok {
-		return nil, fmt.Errorf("address family %s is not configured", family)
-	}
-	return t.Info(id, as), nil
-}
diff --git a/pkg/server/server.go b/pkg/server/server.go
index d9e14926..a2c3dfbd 100644
--- a/pkg/server/server.go
+++ b/pkg/server/server.go
@@ -2568,7 +2568,15 @@ func (s *BgpServer) getRibInfo(addr string, family bgp.RouteFamily) (info *table
 			as = peer.AS()
 			m = s.rsRib
 		}
-		info, err = m.TableInfo(id, as, family)
+
+		af := bgp.RouteFamily(family)
+		tbl, ok := m.Tables[af]
+		if !ok {
+			return fmt.Errorf("address family: %s not supported", af)
+		}
+
+		info = tbl.Info(table.TableInfoOptions{ID: id, AS: as})
+
 		return err
 	}, true)
 	return
@@ -2595,6 +2603,40 @@ func (s *BgpServer) getAdjRibInfo(addr string, family bgp.RouteFamily, in bool)
 	return
 }
 
+func (s *BgpServer) getVrfRibInfo(name string, family bgp.RouteFamily) (info *table.TableInfo, err error) {
+	err = s.mgmtOperation(func() error {
+		m := s.globalRib
+		vrfs := m.Vrfs
+		if _, ok := vrfs[name]; !ok {
+			return fmt.Errorf("vrf %s not found", name)
+		}
+
+		var af bgp.RouteFamily
+		switch family {
+		case bgp.RF_IPv4_UC:
+			af = bgp.RF_IPv4_VPN
+		case bgp.RF_IPv6_UC:
+			af = bgp.RF_IPv6_VPN
+		case bgp.RF_FS_IPv4_UC:
+			af = bgp.RF_FS_IPv4_VPN
+		case bgp.RF_FS_IPv6_UC:
+			af = bgp.RF_FS_IPv6_VPN
+		case bgp.RF_EVPN:
+			af = bgp.RF_EVPN
+		}
+
+		tbl, ok := m.Tables[af]
+		if !ok {
+			return fmt.Errorf("address family: %s not supported", af)
+		}
+
+		info = tbl.Info(table.TableInfoOptions{VRF: vrfs[name]})
+
+		return err
+	}, true)
+	return
+}
+
 func (s *BgpServer) GetTable(ctx context.Context, r *api.GetTableRequest) (*api.GetTableResponse, error) {
 	if r == nil {
 		return nil, fmt.Errorf("invalid request")
@@ -2614,6 +2656,8 @@ func (s *BgpServer) GetTable(ctx context.Context, r *api.GetTableRequest) (*api.
 		fallthrough
 	case api.TableType_ADJ_OUT:
 		info, err = s.getAdjRibInfo(r.Name, family, in)
+	case api.TableType_VRF:
+		info, err = s.getVrfRibInfo(r.Name, family)
 	default:
 		return nil, fmt.Errorf("unsupported resource type: %s", r.TableType)
 	}