3 files changed, 82 insertions, 0 deletions

diff --git a/zebra/zebra_rib.c b/zebra/zebra_rib.c
index 0035eb6e..8da6c84a 100644
--- a/zebra/zebra_rib.c
+++ b/zebra/zebra_rib.c
@@ -2905,6 +2905,61 @@ rib_sweep_route (void)
   rib_sweep_table (vrf_table (AFI_IP6, SAFI_UNICAST, 0));
 }
 
+/* Delete routes learned from a given client.  */
+/* TODO(wsun) May need to split the sweep process into multiple batches,
+ * so that the process won't take too long if the table is large. */
+static void
+rib_sweep_client_table (struct route_table *table, int rib_type)
+{
+  struct route_node *rn;
+  struct rib *rib;
+  struct rib *next;
+  int ret = 0;
+
+  if (table)
+    for (rn = route_top (table); rn; rn = route_next (rn))
+      for (rib = rn->info; rib; rib = next)
+	{
+	  next = rib->next;
+
+	  if (CHECK_FLAG (rib->status, RIB_ENTRY_REMOVED))
+	    continue;
+
+	  if (rib->type == rib_type)
+            if (CHECK_FLAG (rib->flags, ZEBRA_FLAG_SELECTED))
+	      {
+                /* TODO(wsun) Is this mandatory? What about graceful restart/
+                 * non-stop forwarding */
+	        ret = rib_uninstall_kernel (rn, rib);
+	        if (! ret)
+                  rib_delnode (rn, rib);
+                else
+                  zlog_err ("%s: could not delete routes from kernel!",
+                            __func__);
+	      }
+            else
+              {
+                /* Always delete the node. */
+                rib_delnode (rn, rib);
+              }
+	}
+}
+
+/* Sweep all routes learned from a given client from RIB tables.  */
+void
+rib_sweep_client_route (struct zserv *client)
+{
+  assert(client);
+  int route_type = client->route_type;
+  if (route_type != ZEBRA_ROUTE_MAX)
+    {
+      zlog_debug ("%s: Removing existing routes from client type %d",
+                  __func__, route_type);
+      rib_sweep_client_table (vrf_table (AFI_IP, SAFI_UNICAST, 0), route_type);
+      rib_sweep_client_table (vrf_table (AFI_IP6, SAFI_UNICAST, 0), route_type);
+    }
+}
+
 /* Close RIB and clean up kernel routes. */
 static void
 rib_close_table (struct route_table *table)
diff --git a/zebra/zserv.c b/zebra/zserv.c
index 9e6f6253..2330135a 100644
--- a/zebra/zserv.c
+++ b/zebra/zserv.c
@@ -750,6 +750,13 @@ zread_ipv4_add (struct zserv *client, u_short length)
   
   /* Type, flags, message. */
   rib->type = stream_getc (s);
+  /* Update client's route type if it is not done yet. */
+  /* It is done here since only zread_ipv4/6_add() and
+   * zread_ipv4/6_delete() decode Zebra messages and retrieve
+   * route types. */
+  if (client->route_type == ZEBRA_ROUTE_MAX)
+    client->route_type = rib->type;
+
   rib->flags = stream_getc (s);
   message = stream_getc (s); 
   rib->uptime = time (NULL);
@@ -924,6 +931,11 @@ zread_ipv6_add (struct zserv *client, u_short length)
 
   /* Type, flags, message. */
   api.type = stream_getc (s);
+  /* Update the route type of the client. 
+   * Same as in zread_ipv4_add(). */
+  if (client->route_type == ZEBRA_ROUTE_MAX)
+    client->route_type = api.type;
+
   api.flags = stream_getc (s);
   api.message = stream_getc (s);
 
@@ -1077,6 +1089,14 @@ zread_router_id_delete (struct zserv *client, u_short length)
 static void
 zebra_client_close (struct zserv *client)
 {
+  struct stream *s;
+
+  /* Sweep all routes learned from the client first. */
+  rib_sweep_client_route(client);
+  /* Reset the route type. It may not be necessary since the
+   * whole client will be freed. */
+  client->route_type = ZEBRA_ROUTE_MAX;
+
   /* Close file descriptor. */
   if (client->sock)
     {
@@ -1115,6 +1135,9 @@ zebra_client_create (int sock)
 
   /* Make client input/output buffer. */
   client->sock = sock;
+  /* Set the default route type to ZEBRA_ROUTE_MAX; it will be updated
+   * once new routes are received. */
+  client->route_type = ZEBRA_ROUTE_MAX;
   client->ibuf = stream_new (ZEBRA_MAX_PACKET_SIZ);
   client->obuf = stream_new (ZEBRA_MAX_PACKET_SIZ);
   client->wb = buffer_new(0);
diff --git a/zebra/zserv.h b/zebra/zserv.h
index a7371830..e37041f8 100644
--- a/zebra/zserv.h
+++ b/zebra/zserv.h
@@ -38,6 +38,10 @@ struct zserv
   /* Client file descriptor. */
   int sock;
 
+  /* Client route type. */
+  /* Assuming each client contains only one type of route. */
+  int route_type;
+
   /* Input/output buffer to the client. */
   struct stream *ibuf;
   struct stream *obuf;