From 778fc592b900eccd929ff70d119a4c0bb290a0ce Mon Sep 17 00:00:00 2001
From: dovholuknf <46322585+dovholuknf@users.noreply.github.com>
Date: Thu, 18 Jun 2020 16:14:46 -0400
Subject: [PATCH] rework dns addition to take all existing dns into account so
 if the service fails the network won't end up looking like it's broken

---
 service/cziti/ctun.go              |  2 +-
 service/cziti/windns/windns.go     | 77 ++++++++++++++++++++++++++++--
 service/ziti-tunnel/service/ipc.go |  3 +-
 3 files changed, 77 insertions(+), 5 deletions(-)

diff --git a/service/cziti/ctun.go b/service/cziti/ctun.go
index ed882d4f..3ea04c7c 100644
--- a/service/cziti/ctun.go
+++ b/service/cziti/ctun.go
@@ -263,6 +263,6 @@ func AddIntercept(svcId string, service string, host string, port uint16, ctx *C
 	for _, t := range devMap {
 		log.Debug("adding intercept for: %s, %s, %d", service, host, port)
 		res := C.ziti_tunneler_intercept_v1(t.tunCtx, unsafe.Pointer(ctx.nf), C.CString(svcId), C.CString(service), C.CString(host), C.int(port))
-		log.Debug("intercept added", res)
+		log.Debugf("intercept added", res)
 	}
 }
\ No newline at end of file
diff --git a/service/cziti/windns/windns.go b/service/cziti/windns/windns.go
index 2900a400..1a27a84b 100644
--- a/service/cziti/windns/windns.go
+++ b/service/cziti/windns/windns.go
@@ -19,18 +19,18 @@ package windns
 
 import (
 	"bytes"
-	"fmt"
+	"github.com/michaelquigley/pfxlog"
 	"net"
 	"os"
 	"os/exec"
 	"strings"
-
-	"github.com/michaelquigley/pfxlog"
 )
 
 var log = pfxlog.Logger()
 
 func ResetDNS() {
+	log.Infof("skipping ResetDNS for now")
+
 	log.Infof("resetting dns...")
 	log.Info("restoring dns to original-ish state")
 
@@ -77,14 +77,85 @@ func GetUpstreamDNS() []string {
 }
 
 func ReplaceDNS(ips []net.IP) {
+
+	/*
+	ifaces, err := net.Interfaces()
+	if err != nil {
+		fmt.Print(fmt.Errorf("localAddresses: %+v\n", err.Error()))
+		return
+	}
+	for _, i := range ifaces {
+		addrs, err := i.Addrs()
+		if err != nil {
+			fmt.Print(fmt.Errorf("localAddresses: %+v\n", err.Error()))
+			continue
+		}
+		for _, a := range addrs {
+			switch v := a.(type) {
+			case *net.IPAddr:
+				fmt.Printf("%v : %s (%s)\n", i.Name, v, v.IP.DefaultMask())
+
+			case *net.IPNet:
+				fmt.Printf("%v : %s [%v/%v]\n", i.Name, v, v.IP, v.Mask)
+			}
+
+		}
+	}
+
 	var names []string
 	for _, i := range ips {
 		names = append(names, i.String())
 	}
 	addresses := strings.Join(names, ",")
+
+	*/
+	/*
 	script := fmt.Sprintf(
 		`Get-NetIPInterface | ForEach-Object { Set-DnsClientServerAddress -InterfaceIndex $_.ifIndex -ServerAddresses %s }`,
 		addresses)
+	*/
+
+
+
+	script := `$dnsinfo=Get-DnsClientServerAddress
+
+# see https://docs.microsoft.com/en-us/dotnet/api/system.net.sockets.addressfamily
+$IPv4=2
+$IPv6=23
+
+$dnsUpdates = @{}
+
+foreach ($dns in $dnsinfo)
+{
+    if($dnsUpdates[$dns.InterfaceIndex] -eq $null) { $dnsUpdates[$dns.InterfaceIndex]=[System.Collections.ArrayList]@() }
+    if($dns.AddressFamily -eq $IPv6) {
+        $dnsServers=$dns.ServerAddresses
+        $ArrList=[System.Collections.ArrayList]@($dnsServers)
+        if(($dnsServers -ne $null) -and ($dnsServers.Contains("::1")) ) {
+            echo ($dns.InterfaceAlias + " IPv6 already contains ::1")
+        } else {
+            $ArrList.Insert(0,"::1")
+        }
+        $dnsUpdates[$dns.InterfaceIndex].AddRange($ArrList)
+    }
+    elseif($dns.AddressFamily -eq $IPv4){
+        $dnsServers=$dns.ServerAddresses
+        $ArrList=[System.Collections.ArrayList]@($dnsServers)
+        if(($dnsServers -ne $null) -and ($dnsServers.Contains("127.0.0.1")) ) {
+            echo ($dns.InterfaceAlias + " IPv4 already contains 127.0.0.1")
+        } else {
+            $ArrList.Insert(0,"127.0.0.1")
+        }
+        $dnsUpdates[$dns.InterfaceIndex].AddRange($ArrList)
+    }
+}
+
+foreach ($key in $dnsUpdates.Keys)
+{
+    $dnsServers=$dnsUpdates[$key]
+    Set-DnsClientServerAddress -InterfaceIndex $key -ServerAddresses ($dnsServers)
+}`
+
 	cmd := exec.Command("powershell", "-Command", script)
 	cmd.Stderr = os.Stderr
 	cmd.Stdout = os.Stdout
diff --git a/service/ziti-tunnel/service/ipc.go b/service/ziti-tunnel/service/ipc.go
index 5c52c3ea..e80512b5 100644
--- a/service/ziti-tunnel/service/ipc.go
+++ b/service/ziti-tunnel/service/ipc.go
@@ -681,6 +681,7 @@ func disconnectIdentity(id *dto.Identity) error {
 
 	id.Active = false
 	if id.Connected {
+		log.Debugf("ranging over services all services to remove intercept and deregister the service")
 		for _, s := range id.Services {
 			cziti.RemoveIntercept(s.Id)
 			cziti.DNS.DeregisterService(id.NFContext, s.Name)
@@ -777,7 +778,7 @@ func acceptServices() {
 							Fingerprint: id.FingerPrint,
 							Service:     svc,
 						}
-						log.Debug(" dispatched added service change event")
+						log.Debug("dispatched added service change event")
 					case cziti.REMOVED:
 						for idx, svc := range id.Services {
 							if svc.Name == c.Service.Name {