From 299519dda6049840690747404dcbc947738dbfe1 Mon Sep 17 00:00:00 2001
From: dovholuknf <46322585+dovholuknf@users.noreply.github.com>
Date: Tue, 11 Aug 2020 17:07:23 -0400
Subject: [PATCH] closes issue #94 - add hosted services

also handle DNS requests that are both without a . and with.

as well as when the connection-specific local domain is on the DNS request
---
 service/cziti/dns-proxy.go     | 21 ++++++++++++++++++++-
 service/cziti/dns.go           | 17 +++++++++--------
 service/cziti/windns/windns.go | 28 ++++++++++++++++++++++++++++
 3 files changed, 57 insertions(+), 9 deletions(-)

diff --git a/service/cziti/dns-proxy.go b/service/cziti/dns-proxy.go
index ddd6c838..dc6fca05 100644
--- a/service/cziti/dns-proxy.go
+++ b/service/cziti/dns-proxy.go
@@ -28,6 +28,8 @@ import (
 	"github.com/openziti/desktop-edge-win/service/cziti/windns"
 )
 
+var domains []string // get any connection-specific local domains
+
 func processDNSquery(packet []byte, p *net.UDPAddr, s *net.UDPConn) {
 	q := &dns.Msg{}
 	if err := q.Unpack(packet); err != nil {
@@ -46,8 +48,24 @@ func processDNSquery(packet []byte, p *net.UDPAddr, s *net.UDPConn) {
 
 	var ip net.IP
 	// fmt.Printf("query: Type(%d) name(%s)\n", query.Qtype, query.Name)
+	dnsName := strings.TrimSpace(query.Name)
+	ip = DNS.Resolve(dnsName)
 
-	ip = DNS.Resolve(query.Name)
+	// never proxy hostnames that we know about regardless of type
+	if ip == nil {
+		// no direct hit. need to now check to see if the dns query used a connection-specific local domain
+		for _, d := range domains {			
+			if strings.HasSuffix(dnsName, d) {
+				dnsNameTrimmed := strings.Trim(dnsName, d)
+				// dns request has domain appended - removing and resolving
+				ip = DNS.Resolve(dnsNameTrimmed)
+				if ip != nil {
+					log.Debugf("DNS %s failed at first but succeeded by looking up %s", dnsName, dnsNameTrimmed)
+					break
+				}
+			}
+		}
+	}
 
 	// never proxy hostnames that we know about regardless of type
 	if ip != nil {
@@ -167,6 +185,7 @@ func dnsPanicRecover() {
 }
 
 func runDNSproxy(dnsServers []string) {
+	domains = windns.GetConnectionSpecificDomains()
 	defer func() {
 		if err := recover(); err != nil {
 			log.Errorf("Recovered in f. %v", err)
diff --git a/service/cziti/dns.go b/service/cziti/dns.go
index 447172c8..5fa0c674 100644
--- a/service/cziti/dns.go
+++ b/service/cziti/dns.go
@@ -43,7 +43,7 @@ var DNS = &dnsImpl{}
 
 type dnsImpl struct {
 	cidr    uint32
-	counter uint32
+	ipCount uint32
 
 	serviceMap map[string]ctxService
 
@@ -67,12 +67,13 @@ type ctxService struct {
 }
 
 func normalizeDnsName(dnsName string) string {
-	if strings.HasSuffix(dnsName, ".") {
-		return dnsName
-	} else {
+	dnsName = strings.TrimSpace(dnsName)
+	if !strings.HasSuffix(dnsName, ".") {
 		// append a period to the dnsName - forcing it to be a FQDN
-		return dnsName + "."
+		dnsName += "."
 	}
+
+	return strings.ToLower(dnsName) //normalize the casing
 }
 
 // RegisterService will return the next ip address in the configured range. If the ip address is not
@@ -115,7 +116,7 @@ func (dns *dnsImpl) RegisterService(svcId string, dnsNameToReg string, port uint
 		}
 	} else {
 		// if not used at all - map it
-		nextAddr := dns.cidr | atomic.AddUint32(&dns.counter, 1)
+		nextAddr := dns.cidr | atomic.AddUint32(&dns.ipCount, 1)
 		ip = make(net.IP, 4)
 		binary.BigEndian.PutUint32(ip, nextAddr)
 
@@ -145,7 +146,7 @@ func (dns *dnsImpl) DeregisterService(ctx *CZitiCtx, name string) {
 	for k, sc := range dns.serviceMap {
 		if sc.ctx == ctx && sc.name == name {
 			if sc.count < 1 {
-				log.Infof("removing %s from DNS mapping", name)
+				log.Infof("removing service named %s from DNS mapping known as %s", name, k)
 				delete(dns.serviceMap, k)
 			} else {
 				// another service is using the mapping - can't remove it yet so decrement
@@ -182,6 +183,6 @@ func DnsInit(ip string, maskBits int) {
 		i := net.ParseIP(ip).To4()
 		mask := net.CIDRMask(maskBits, 32)
 		DNS.cidr = binary.BigEndian.Uint32(i) & binary.BigEndian.Uint32(mask)
-		DNS.counter = 2
+		DNS.ipCount = 2
 	})
 }
diff --git a/service/cziti/windns/windns.go b/service/cziti/windns/windns.go
index 8072f6f4..c51fd80f 100644
--- a/service/cziti/windns/windns.go
+++ b/service/cziti/windns/windns.go
@@ -43,6 +43,34 @@ func ResetDNS() {
 	}
 }
 
+func GetConnectionSpecificDomains() []string {
+	script := `Get-DnsClient | Select-Object ConnectionSpecificSuffix | Get-Unique | ForEach-Object { $_.ConnectionSpecificSuffix }`
+
+	cmd := exec.Command("powershell", "-Command", script)
+	cmd.Stderr = os.Stdout
+	output := new(bytes.Buffer)
+	cmd.Stdout = output
+
+	err := cmd.Run()
+
+	if err != nil {
+		panic(err)
+	}
+
+	var names []string
+	for {
+		domain, err := output.ReadString('\n')
+		if err != nil {
+			break
+		}
+		domain = strings.TrimSpace(domain)
+		if !strings.HasSuffix(domain, ".") {
+			names = append(names, domain + ".")
+		}
+	}
+	return names
+}
+
 func GetUpstreamDNS() []string {
 	script := `Get-DnsClientServerAddress | ForEach-Object { $_.ServerAddresses } | Sort-Object | Get-Unique`