--- fwsync/driver/fwsync.c	2022/08/04 00:24:39	1.7
+++ fwsync/driver/fwsync.c	2022/08/26 13:33:43	1.12.2.1
@@ -1,7 +1,29 @@
-/*************************************************************************
-* (C) 2022 CloudSigma AG - Sofia/Bulgaria
-*  by Michael Pounov <misho@elwix.org>
-**************************************************************************/
+/*-
+ * SPDX-License-Identifier: BSD-2-Clause-FreeBSD
+ *
+ * Copyright (c) 2022 Michael Pounov <misho@elwix.org>, CloudSigma AG
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions
+ * are met:
+ * 1. Redistributions of source code must retain the above copyright
+ *    notice, this list of conditions and the following disclaimer.
+ * 2. Redistributions in binary form must reproduce the above copyright
+ *    notice, this list of conditions and the following disclaimer in the
+ *    documentation and/or other materials provided with the distribution.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS ``AS IS'' AND
+ * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+ * ARE DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
+ * OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
+ * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
+ * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
+ * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
+ * SUCH DAMAGE.
+ */
 #include "fwsync.h"
 
 
@@ -33,11 +55,10 @@ fwsync_edge_proc(void *arg)
 		if (m_length(m, NULL) != sizeof(struct fws_proto)) {
 			printf("FWSync packet length=%d isn't match expected %lu\n", 
 					m_length(m, NULL), sizeof(struct fws_proto));
+			m_freem(m);
 			return;
 		}
 
-		printf("%u...-> %s version=%c\n", m->m_len, m->m_data, pkt->fws_version);
-
 		switch (pkt->fws_version) {
 			case FWS_PKTVER_STATE:
 				fwsync_add_state(pkt);
@@ -47,8 +68,10 @@ fwsync_edge_proc(void *arg)
 				break;
 			default:
 				printf("FWSync packet was discarded due to wrong version\n");
-				return;
+				break;
 		}
+
+		m_freem(m);
 	}
 }
 
@@ -161,16 +184,20 @@ fwsync_destroy(struct ip_fw_chain *ch, ip_fw3_opheader
 
 	n = (int*) (oh + 1);
 	if (*n & CFG_SYNC_EDGE) {
+		if (fws_ctx.config & CTX_EDGE_ONLINE) {
+			ipfw_unregister_state_sync();
+			ipfw_unregister_alias_sync();
+		}
+
+		callout_drain(&fws_co);
+
 		fws_cfg.cfg.on &= ~CFG_SYNC_EDGE;
 		fws_cfg.cfg.edge = 0;
 		fws_cfg.cfg.addrs = 0;
 		memset(fws_cfg.cfg_addr, 0, sizeof fws_cfg.cfg_addr[0]);
 
-		if (fws_ctx.config & CTX_EDGE_ONLINE) {
-			callout_drain(&fws_co);
-			soshutdown(fws_ctx.sockz[CFG_SYNC_ADDR_EDGE], SHUT_RD);
-			soclose(fws_ctx.sockz[CFG_SYNC_ADDR_EDGE]);
-		}
+		soshutdown(fws_ctx.sockz[CFG_SYNC_ADDR_EDGE], SHUT_RD);
+		soclose(fws_ctx.sockz[CFG_SYNC_ADDR_EDGE]);
 	}
 	if (*n & CFG_SYNC_COLLECTOR) {
 		if (fws_ctx.config & (CTX_COLLECTOR_1_ONLINE | CTX_COLLECTOR_2_ONLINE)) {
@@ -178,6 +205,8 @@ fwsync_destroy(struct ip_fw_chain *ch, ip_fw3_opheader
 			ipfw_unregister_alias_hook();
 		}
 
+		taskqueue_drain(fws_tq, &fws_sndpkt_task);
+
 		fws_cfg.cfg.on &= ~CFG_SYNC_COLLECTOR;
 		fws_cfg.cfg.collector = 0;
 		fws_cfg.cfg.addrs = 0;
@@ -191,6 +220,8 @@ fwsync_destroy(struct ip_fw_chain *ch, ip_fw3_opheader
 
 	fws_ctx.config ^= fws_ctx.config;
 
+	memset(&fws_acct, 0, sizeof fws_acct);
+
 	return 0;
 }
 
@@ -264,9 +295,15 @@ fwsync_start(struct ip_fw_chain *ch, ip_fw3_opheader *
 
 	if ((*n & CFG_SYNC_EDGE) && (fws_ctx.config & CTX_EDGE_READY)) {
 		fws_cfg.cfg.on |= CFG_SYNC_EDGE;
-		fws_ctx.config |= CTX_EDGE_ONLINE;
 
 		callout_reset(&fws_co, hz, fwsync_edge_proc, NULL);
+
+		if (!(fws_ctx.config & CTX_EDGE_ONLINE)) {
+			ipfw_register_state_sync(fwsync_state_sync);
+			ipfw_register_alias_sync(fwsync_alias_sync);
+		}
+
+		fws_ctx.config |= CTX_EDGE_ONLINE;
 	}
 
 	if ((*n & CFG_SYNC_COLLECTOR) && (fws_ctx.config & CTX_COLLECTOR_1_READY)) {
@@ -321,6 +358,9 @@ fwsync_stop(struct ip_fw_chain *ch, ip_fw3_opheader *o
 		fws_ctx.config &= ~CTX_EDGE_ONLINE;
 
 		callout_drain(&fws_co);
+
+		ipfw_unregister_state_sync();
+		ipfw_unregister_alias_sync();
 	}
 
 	if ((*n & CFG_SYNC_COLLECTOR) && (fws_ctx.config & CTX_COLLECTOR_2_ONLINE))
@@ -334,6 +374,8 @@ fwsync_stop(struct ip_fw_chain *ch, ip_fw3_opheader *o
 
 		ipfw_unregister_state_hook();
 		ipfw_unregister_alias_hook();
+
+		taskqueue_drain(fws_tq, &fws_sndpkt_task);
 	}
 
 	return 0;