zebra: add iterator for walking all tables in RIB

* lib/zebra.h

    Add macro ZEBRA_NUM_OF, which returns the number of elements in a
    static array.

  * zebra/rib.h

    Add the rib_tables_iter_t structure and associated functions,
    which allow one to walk all tables in the rib.

  * zebra/zebra_rib.c

    - Add vrf_id_get_next() to retrieve the first VRF id (if any) that
      is greater than a given VRF id.

    - Add rib_tables_iter_next().

Signed-off-by: Avneesh Sachdev <avneesh@opensourcerouting.org>
Signed-off-by: David Lamparter <equinox@opensourcerouting.org>

3 files changed, 157 insertions, 0 deletions

diff --git a/lib/zebra.h b/lib/zebra.h
index f8a6be30..404b832b 100644
--- a/lib/zebra.h
+++ b/lib/zebra.h
@@ -387,6 +387,8 @@ struct in_pktinfo
 #define MIN(a, b) ((a) < (b) ? (a) : (b))
 #endif
 
+#define ZEBRA_NUM_OF(x) (sizeof (x) / sizeof (x[0]))
+
 /* For old definition. */
 #ifndef IN6_ARE_ADDR_EQUAL
 #define IN6_ARE_ADDR_EQUAL IN6_IS_ADDR_EQUAL
diff --git a/zebra/rib.h b/zebra/rib.h
index c98d99a4..4f99d714 100644
--- a/zebra/rib.h
+++ b/zebra/rib.h
@@ -286,6 +286,25 @@ typedef struct rib_table_info_t_
 
 } rib_table_info_t;
 
+typedef enum
+{
+  RIB_TABLES_ITER_S_INIT,
+  RIB_TABLES_ITER_S_ITERATING,
+  RIB_TABLES_ITER_S_DONE
+} rib_tables_iter_state_t;
+
+/*
+ * Structure that holds state for iterating over all tables in the
+ * Routing Information Base.
+ */
+typedef struct rib_tables_iter_t_
+{
+  uint32_t vrf_id;
+  int afi_safi_ix;
+
+  rib_tables_iter_state_t state;
+} rib_tables_iter_t;
+
 extern struct nexthop *nexthop_ifindex_add (struct rib *, unsigned int);
 extern struct nexthop *nexthop_ifname_add (struct rib *, char *);
 extern struct nexthop *nexthop_blackhole_add (struct rib *);
@@ -374,6 +393,7 @@ static_delete_ipv6 (struct prefix *p, u_char type, struct in6_addr *gate,
 #endif /* HAVE_IPV6 */
 
 extern int rib_gc_dest (struct route_node *rn);
+extern struct route_table *rib_tables_iter_next (rib_tables_iter_t *iter);
 
 /*
  * Inline functions.
@@ -453,4 +473,36 @@ rib_dest_vrf (rib_dest_t *dest)
   return rib_table_info (rib_dest_table (dest))->vrf;
 }
 
+/*
+ * rib_tables_iter_init
+ */
+static inline void
+rib_tables_iter_init (rib_tables_iter_t *iter)
+
+{
+  memset (iter, 0, sizeof (*iter));
+  iter->state = RIB_TABLES_ITER_S_INIT;
+}
+
+/*
+ * rib_tables_iter_started
+ *
+ * Returns TRUE if this iterator has started iterating over the set of
+ * tables.
+ */
+static inline int
+rib_tables_iter_started (rib_tables_iter_t *iter)
+{
+  return iter->state != RIB_TABLES_ITER_S_INIT;
+}
+
+/*
+ * rib_tables_iter_cleanup
+ */
+static inline void
+rib_tables_iter_cleanup (rib_tables_iter_t *iter)
+{
+  iter->state = RIB_TABLES_ITER_S_DONE;
+}
+
 #endif /*_ZEBRA_RIB_H */
diff --git a/zebra/zebra_rib.c b/zebra/zebra_rib.c
index 2cbee935..5c75b909 100644
--- a/zebra/zebra_rib.c
+++ b/zebra/zebra_rib.c
@@ -3077,3 +3077,106 @@ rib_init (void)
   /* VRF initialization.  */
   vrf_init ();
 }
+
+/*
+ * vrf_id_get_next
+ *
+ * Get the first vrf id that is greater than the given vrf id if any.
+ *
+ * Returns TRUE if a vrf id was found, FALSE otherwise.
+ */
+static inline int
+vrf_id_get_next (uint32_t id, uint32_t *next_id_p)
+{
+  while (++id < vector_active (vrf_vector))
+    {
+      if (vrf_lookup (id))
+	{
+	  *next_id_p = id;
+	  return 1;
+	}
+    }
+
+  return 0;
+}
+
+/*
+ * rib_tables_iter_next
+ *
+ * Returns the next table in the iteration.
+ */
+struct route_table *
+rib_tables_iter_next (rib_tables_iter_t *iter)
+{
+  struct route_table *table;
+
+  /*
+   * Array that helps us go over all AFI/SAFI combinations via one
+   * index.
+   */
+  static struct {
+    afi_t afi;
+    safi_t safi;
+  } afi_safis[] = {
+    { AFI_IP, SAFI_UNICAST },
+    { AFI_IP, SAFI_MULTICAST },
+    { AFI_IP6, SAFI_UNICAST },
+    { AFI_IP6, SAFI_MULTICAST },
+  };
+
+  table = NULL;
+
+  switch (iter->state)
+    {
+
+    case RIB_TABLES_ITER_S_INIT:
+      iter->vrf_id = 0;
+      iter->afi_safi_ix = -1;
+
+      /* Fall through */
+
+    case RIB_TABLES_ITER_S_ITERATING:
+      iter->afi_safi_ix++;
+      while (1)
+	{
+
+	  while (iter->afi_safi_ix < (int) ZEBRA_NUM_OF (afi_safis))
+	    {
+	      table = vrf_table (afi_safis[iter->afi_safi_ix].afi,
+				 afi_safis[iter->afi_safi_ix].safi,
+				 iter->vrf_id);
+	      if (table)
+		break;
+
+	      iter->afi_safi_ix++;
+	    }
+
+	  /*
+	   * Found another table in this vrf.
+	   */
+	  if (table)
+	    break;
+
+	  /*
+	   * Done with all tables in the current vrf, go to the next
+	   * one.
+	   */
+	  if (!vrf_id_get_next (iter->vrf_id, &iter->vrf_id))
+	    break;
+
+	  iter->afi_safi_ix = 0;
+	}
+
+      break;
+
+    case RIB_TABLES_ITER_S_DONE:
+      return NULL;
+    }
+
+  if (table)
+    iter->state = RIB_TABLES_ITER_S_ITERATING;
+  else
+    iter->state = RIB_TABLES_ITER_S_DONE;
+
+  return table;
+}