From 47a98d1dbedd36e6edefda37a7e71187a5162001 Mon Sep 17 00:00:00 2001
From: dovholuknf <46322585+dovholuknf@users.noreply.github.com>
Date: Fri, 9 Jul 2021 12:59:59 -0400
Subject: [PATCH] allow 'adding back' removed intercepts

---
 service/cziti/ctun.go              | 31 +++++++++++++++++++++++-------
 service/cziti/sdk.go               |  2 ++
 service/ziti-tunnel/service/ipc.go | 13 ++++++++++++-
 3 files changed, 38 insertions(+), 8 deletions(-)

diff --git a/service/cziti/ctun.go b/service/cziti/ctun.go
index 09c082cd..62b1508b 100644
--- a/service/cziti/ctun.go
+++ b/service/cziti/ctun.go
@@ -45,6 +45,7 @@ void readIdle(uv_prepare_t *idler);
 
 void call_on_packet(void *packet, ssize_t len, packet_cb cb, void *ctx);
 void remove_intercepts(uv_async_t *handle);
+void add_intercepts(uv_async_t *handle);
 void free_async(uv_handle_t* timer);
 
 dns_manager* get_dns_mgr_from_c();
@@ -242,26 +243,42 @@ func (t *tunnel) runWriteLoop() {
 	}
 }
 
-type RemoveWG struct {
+type TunnelerActionWaitGroup struct {
 	Wg    *sync.WaitGroup
 	Czsvc *ZService
 }
 
-func RemoveIntercept(rwg *RemoveWG) {
+func RemoveIntercept(removeWaitGroup *TunnelerActionWaitGroup) {
 	async := (*C.uv_async_t)(C.malloc(C.sizeof_uv_async_t))
-	async.data = unsafe.Pointer(rwg)
+	async.data = unsafe.Pointer(removeWaitGroup)
 	C.uv_async_init(_impl.libuvCtx.l, async, C.uv_async_cb(C.remove_intercepts))
 	C.uv_async_send((*C.uv_async_t)(unsafe.Pointer(async)))
 }
 
+func AddIntercept(addWaitGroup *TunnelerActionWaitGroup) {
+	async := (*C.uv_async_t)(C.malloc(C.sizeof_uv_async_t))
+	async.data = unsafe.Pointer(addWaitGroup)
+	C.uv_async_init(_impl.libuvCtx.l, async, C.uv_async_cb(C.add_intercepts))
+	C.uv_async_send((*C.uv_async_t)(unsafe.Pointer(async)))
+}
+
 //export remove_intercepts
 func remove_intercepts(async *C.uv_async_t) {
-	rwg := (*RemoveWG)(async.data)
-	cSvcName := C.CString(rwg.Czsvc.Name)
+	removeWaitGroup := (*TunnelerActionWaitGroup)(async.data)
+	cSvcName := C.CString(removeWaitGroup.Czsvc.Name)
 	defer C.free(unsafe.Pointer(cSvcName))
-	C.ziti_tunneler_stop_intercepting(theTun.tunCtx, unsafe.Pointer(rwg.Czsvc.Czctx), cSvcName)
+	C.ziti_tunneler_stop_intercepting(theTun.tunCtx, unsafe.Pointer(removeWaitGroup.Czsvc.Czctx), cSvcName)
 	C.uv_close((*C.uv_handle_t)(unsafe.Pointer(async)), C.uv_close_cb(C.free_async))
-	rwg.Wg.Done()
+	removeWaitGroup.Wg.Done()
+}
+
+//export add_intercepts
+func add_intercepts(async *C.uv_async_t) {
+	addWaitGroup := (*TunnelerActionWaitGroup)(async.data)
+
+	C.ziti_sdk_c_on_service(addWaitGroup.Czsvc.Czctx, addWaitGroup.Czsvc.Czsvc, C.ZITI_OK, unsafe.Pointer(theTun.tunCtx))
+	C.uv_close((*C.uv_handle_t)(unsafe.Pointer(async)), C.uv_close_cb(C.free_async))
+	addWaitGroup.Wg.Done()
 }
 
 //export apply_dns_go
diff --git a/service/cziti/sdk.go b/service/cziti/sdk.go
index 2f116a56..1d057093 100644
--- a/service/cziti/sdk.go
+++ b/service/cziti/sdk.go
@@ -113,6 +113,7 @@ type ZService struct {
 	Name    string
 	Id      string
 	Service *dto.Service
+	Czsvc   *C.ziti_service
 	Czctx   C.ziti_context
 }
 
@@ -375,6 +376,7 @@ func serviceCB(ziti_ctx C.ziti_context, service *C.ziti_service, status C.int, z
 			Name:    name,
 			Id:      svcId,
 			Service: svc,
+			Czsvc:   service,
 			Czctx:   ziti_ctx,
 		}
 		zid.Services.Store(svcId, &added)
diff --git a/service/ziti-tunnel/service/ipc.go b/service/ziti-tunnel/service/ipc.go
index 2b7cd53f..da965f39 100644
--- a/service/ziti-tunnel/service/ipc.go
+++ b/service/ziti-tunnel/service/ipc.go
@@ -940,6 +940,17 @@ func connectIdentity(id *Id) {
 
 		id.CId.Services.Range(func(key interface{}, value interface{}) bool {
 			id.Services = append(id.Services, nil)
+
+			val := value.(*cziti.ZService)
+			var wg sync.WaitGroup
+			wg.Add(1)
+			rwg := &cziti.TunnelerActionWaitGroup{
+				Wg:    &wg,
+				Czsvc: val,
+			}
+			cziti.AddIntercept(rwg)
+			wg.Wait()
+
 			return true
 		})
 
@@ -964,7 +975,7 @@ func disconnectIdentity(id *Id) error {
 				val := value.(*cziti.ZService)
 				var wg sync.WaitGroup
 				wg.Add(1)
-				rwg := &cziti.RemoveWG{
+				rwg := &cziti.TunnelerActionWaitGroup{
 					Wg:    &wg,
 					Czsvc: val,
 				}