3 files changed, 59 insertions, 13 deletions

diff --git a/sys/netinet6/frag6.c b/sys/netinet6/frag6.c
index ca90a9865e32..e975b8a0cb31 100644
--- a/sys/netinet6/frag6.c
+++ b/sys/netinet6/frag6.c
@@ -100,6 +100,12 @@ struct	ip6asfrag {
 
 static MALLOC_DEFINE(M_FRAG6, "frag6", "IPv6 fragment reassembly header");
 
+#ifdef VIMAGE
+/* A flag to indicate if IPv6 fragmentation is initialized. */
+VNET_DEFINE_STATIC(bool,		frag6_on);
+#define	V_frag6_on			VNET(frag6_on)
+#endif
+
 /* System wide (global) maximum and count of packets in reassembly queues. */ 
 static int ip6_maxfrags;
 static volatile u_int frag6_nfrags = 0;
@@ -289,6 +295,15 @@ frag6_cleanup(void *arg __unused, struct ifnet *ifp)
 
 	KASSERT(ifp != NULL, ("%s: ifp is NULL", __func__));
 
+#ifdef VIMAGE
+	/*
+	 * Skip processing if IPv6 reassembly is not initialised or
+	 * torn down by frag6_destroy().
+	 */
+	if (!V_frag6_on)
+		return;
+#endif
+
 	CURVNET_SET_QUIET(ifp->if_vnet);
 	for (i = 0; i < IP6REASS_NHASH; i++) {
 		IP6QB_LOCK(i);
@@ -929,6 +944,9 @@ frag6_init(void)
 	}
 	V_ip6qb_hashseed = arc4random();
 	V_ip6_maxfragsperpacket = 64;
+#ifdef VIMAGE
+	V_frag6_on = true;
+#endif
 	if (!IS_DEFAULT_VNET(curvnet))
 		return;
 
@@ -940,32 +958,58 @@ frag6_init(void)
 /*
  * Drain off all datagram fragments.
  */
+static void
+frag6_drain_one(void)
+{
+	struct ip6q *head;
+	uint32_t bucket;
+
+	for (bucket = 0; bucket < IP6REASS_NHASH; bucket++) {
+		IP6QB_LOCK(bucket);
+		head = IP6QB_HEAD(bucket);
+		while (head->ip6q_next != head) {
+			IP6STAT_INC(ip6s_fragdropped);
+			/* XXX in6_ifstat_inc(ifp, ifs6_reass_fail) */
+			frag6_freef(head->ip6q_next, bucket);
+		}
+		IP6QB_UNLOCK(bucket);
+	}
+}
+
 void
 frag6_drain(void)
 {
 	VNET_ITERATOR_DECL(vnet_iter);
-	struct ip6q *head;
-	uint32_t bucket;
 
 	VNET_LIST_RLOCK_NOSLEEP();
 	VNET_FOREACH(vnet_iter) {
 		CURVNET_SET(vnet_iter);
-		for (bucket = 0; bucket < IP6REASS_NHASH; bucket++) {
-			if (IP6QB_TRYLOCK(bucket) == 0)
-				continue;
-			head = IP6QB_HEAD(bucket);
-			while (head->ip6q_next != head) {
-				IP6STAT_INC(ip6s_fragdropped);
-				/* XXX in6_ifstat_inc(ifp, ifs6_reass_fail) */
-				frag6_freef(head->ip6q_next, bucket);
-			}
-			IP6QB_UNLOCK(bucket);
-		}
+		frag6_drain_one();
 		CURVNET_RESTORE();
 	}
 	VNET_LIST_RUNLOCK_NOSLEEP();
 }
 
+#ifdef VIMAGE
+/*
+ * Clear up IPv6 reassembly structures.
+ */
+void
+frag6_destroy(void)
+{
+	uint32_t bucket;
+
+	frag6_drain_one();
+	V_frag6_on = false;
+	for (bucket = 0; bucket < IP6REASS_NHASH; bucket++) {
+		KASSERT(V_ip6qb[bucket].count == 0,
+		    ("%s: V_ip6qb[%d] (%p) count not 0 (%d)", __func__,
+		    bucket, &V_ip6qb[bucket], V_ip6qb[bucket].count));
+		mtx_destroy(&V_ip6qb[bucket].lock);
+	}
+}
+#endif
+
 /*
  * Put an ip fragment on a reassembly chain.
  * Like insque, but pointers in middle of structure.
diff --git a/sys/netinet6/ip6_input.c b/sys/netinet6/ip6_input.c
index 60a6c501158f..13a4b332f041 100644
--- a/sys/netinet6/ip6_input.c
+++ b/sys/netinet6/ip6_input.c
@@ -393,6 +393,7 @@ ip6_destroy(void *unused __unused)
 	}
 	IFNET_RUNLOCK();
 
+	frag6_destroy();
 	nd6_destroy();
 	in6_ifattach_destroy();
 
diff --git a/sys/netinet6/ip6_var.h b/sys/netinet6/ip6_var.h
index be748b31ceab..f698e8eb5cf1 100644
--- a/sys/netinet6/ip6_var.h
+++ b/sys/netinet6/ip6_var.h
@@ -392,6 +392,7 @@ int	ip6_fragment(struct ifnet *, struct mbuf *, int, u_char, int,
 int	route6_input(struct mbuf **, int *, int);
 
 void	frag6_init(void);
+void	frag6_destroy(void);
 int	frag6_input(struct mbuf **, int *, int);
 void	frag6_slowtimo(void);
 void	frag6_drain(void);