Merge pull request #718 from openziti/dns-cleanup

Cleanup DNS entries when services are stopped

tunnel/dns/file.go

@@ -19,6 +19,7 @@ package dns
 import (
 	"bufio"
 	"fmt"
+	"io"
 	"net"
 	"os"
 	"sync"
@@ -59,7 +60,7 @@ func (h *hostFile) AddHostname(hostname string, ip net.IP) error {
 	defer f.Close()
 
 	if !isHostPresent(hostname, ip, f) {
-		_, err := f.Seek(0, os.SEEK_END)
+		_, err := f.Seek(0, io.SeekEnd)
 		if err != nil {
 			return err
 		}
@@ -72,7 +73,7 @@ func (h *hostFile) AddHostname(hostname string, ip net.IP) error {
 	return nil
 }
 
-func (h *hostFile) RemoveHostname(s string, ip net.IP) error {
+func (h *hostFile) RemoveHostname(s string) error {
 	h.mutex.Lock()
 	defer h.mutex.Unlock()
 	return nil

tunnel/dns/refcount.go

@@ -0,0 +1,50 @@
+package dns
+
+import (
+	cmap "github.com/orcaman/concurrent-map"
+	"net"
+)
+
+func NewRefCountingResolver(resolver Resolver) Resolver {
+	return &RefCountingResolver{
+		names:   cmap.New(),
+		wrapped: resolver,
+	}
+}
+
+type RefCountingResolver struct {
+	names   cmap.ConcurrentMap
+	wrapped Resolver
+}
+
+func (self *RefCountingResolver) AddHostname(s string, ip net.IP) error {
+	err := self.wrapped.AddHostname(s, ip)
+	if err != nil {
+		self.names.Upsert(s, 1, func(exist bool, valueInMap interface{}, newValue interface{}) interface{} {
+			if exist {
+				return valueInMap.(int) + 1
+			}
+			return 1
+		})
+	}
+	return err
+}
+
+func (self *RefCountingResolver) RemoveHostname(s string) error {
+	val := self.names.Upsert(s, 1, func(exist bool, valueInMap interface{}, newValue interface{}) interface{} {
+		if exist {
+			return valueInMap.(int) - 1
+		}
+		return 0
+	})
+
+	if result := val.(int); result == 0 {
+		self.names.Remove(s)
+		return self.wrapped.RemoveHostname(s)
+	}
+	return nil
+}
+
+func (self *RefCountingResolver) Cleanup() error {
+	return self.wrapped.Cleanup()
+}

tunnel/dns/resolver.go

@@ -20,6 +20,6 @@ import "net"
 
 type Resolver interface {
 	AddHostname(string, net.IP) error
-	RemoveHostname(string, net.IP) error
+	RemoveHostname(string) error
 	Cleanup() error
 }

tunnel/dns/server.go

@@ -65,9 +65,9 @@ func NewResolver(config string) Resolver {
 
 	switch resolverURL.Scheme {
 	case "", "file":
-		return NewHostFile(resolverURL.Path)
+		return NewRefCountingResolver(NewHostFile(resolverURL.Path))
 	case "udp":
-		return NewDnsServer(resolverURL.Host)
+		return NewRefCountingResolver(NewDnsServer(resolverURL.Host))
 	}
 
 	log.Fatalf("invalid resolver configuration '%s'. must be 'file://' or 'udp://' URL", config)
@@ -136,7 +136,7 @@ func (r *resolver) testSystemResolver() error {
 		return fmt.Errorf("unexpected resolved address %s", resolved.IP.String())
 	}
 
-	_ = r.RemoveHostname(resolverTestHostname, resolverTestIP)
+	_ = r.RemoveHostname(resolverTestHostname)
 	return nil
 }
 
@@ -179,7 +179,7 @@ func (r *resolver) AddHostname(hostname string, ip net.IP) error {
 	return nil
 }
 
-func (r *resolver) RemoveHostname(hostname string, ip net.IP) error {
+func (r *resolver) RemoveHostname(hostname string) error {
 	r.namesMtx.Lock()
 	if _, ok := r.names[strings.ToLower(hostname)+"."]; ok {
 		log.Infof("removing %s from resolver", hostname)

tunnel/entities/service.go

@@ -11,6 +11,7 @@ import (
 	"net"
 	"reflect"
 	"strconv"
+	"sync"
 	"time"
 )
 
@@ -277,7 +278,23 @@ type Service struct {
 	HostV2Config         *HostV2Config
 	DialIdentityProvider TemplateFunc
 	SourceAddrProvider   TemplateFunc
-	StopHostHook         func()
+	cleanupActions       []func()
+	lock                 sync.Mutex
+}
+
+func (self *Service) AddCleanupAction(f func()) {
+	self.lock.Lock()
+	defer self.lock.Unlock()
+	self.cleanupActions = append(self.cleanupActions, f)
+}
+
+func (self *Service) RunCleanupActions() {
+	self.lock.Lock()
+	defer self.lock.Unlock()
+
+	for _, action := range self.cleanupActions {
+		action()
+	}
 }
 
 func (self *Service) GetSourceAddr(sourceAddr net.Addr, destAddr net.Addr) string {

tunnel/intercept/interceptor.go

@@ -81,7 +81,7 @@ func (addr *InterceptAddress) String() string {
 func GetInterceptAddresses(service *entities.Service, protocol string, resolver dns.Resolver) ([]*InterceptAddress, error) {
 	var result []*InterceptAddress
 	for _, addr := range service.InterceptV1Config.Addresses {
-		_, cidr, err := getInterceptIP(addr, resolver)
+		_, cidr, err := getInterceptIP(service, addr, resolver)
 		if err != nil {
 			return nil, errors.Wrapf(err, "failed to get intercept IP address for %v", addr)
 		}

tunnel/intercept/iputils.go

@@ -20,6 +20,7 @@ import (
 	"fmt"
 	"github.com/michaelquigley/pfxlog"
 	"github.com/openziti/edge/tunnel/dns"
+	"github.com/openziti/edge/tunnel/entities"
 	"github.com/openziti/edge/tunnel/utils"
 	"net"
 )
@@ -51,7 +52,7 @@ func SetDnsInterceptIpRange(cidr string) error {
 	return nil
 }
 
-func getInterceptIP(hostname string, resolver dns.Resolver) (net.IP, *net.IPNet, error) {
+func getInterceptIP(svc *entities.Service, hostname string, resolver dns.Resolver) (net.IP, *net.IPNet, error) {
 	log := pfxlog.Logger()
 
 	ip, ipNet, err := utils.GetDialIP(hostname)
@@ -66,7 +67,14 @@ func getInterceptIP(hostname string, resolver dns.Resolver) (net.IP, *net.IPNet,
 	if err == nil {
 		if len(addrs) > 0 {
 			ip := addrs[0].To4()
-			_ = resolver.AddHostname(hostname, ip)
+			if err = resolver.AddHostname(hostname, ip); err != nil {
+				log.WithError(err).Errorf("failed to add host/ip mapping to resolver: %v -> %v", hostname, ip)
+			}
+			svc.AddCleanupAction(func() {
+				if err = resolver.RemoveHostname(hostname); err != nil {
+					log.WithError(err).Errorf("failed to remove host mapping from resolver: %v ", hostname)
+				}
+			})
 			prefixLen := utils.AddrBits(ip)
 			ipNet := &net.IPNet{IP: ip, Mask: net.CIDRMask(prefixLen, prefixLen)}
 			return addrs[0], ipNet, nil
@@ -79,7 +87,16 @@ func getInterceptIP(hostname string, resolver dns.Resolver) (net.IP, *net.IPNet,
 	if ip == nil {
 		return nil, nil, fmt.Errorf("invalid IP address or unresolvable hostname: %s", hostname)
 	}
-	_ = resolver.AddHostname(hostname, ip)
+	if err = resolver.AddHostname(hostname, ip); err != nil {
+		log.WithError(err).Errorf("failed to add host/ip mapping to resolver: %v -> %v", hostname, ip)
+	}
+
+	svc.AddCleanupAction(func() {
+		if err = resolver.RemoveHostname(hostname); err != nil {
+			log.WithError(err).Errorf("failed to remove host mapping from resolver: %v ", hostname)
+		}
+	})
+
 	prefixLen := utils.AddrBits(ip)
 	ipNet = &net.IPNet{IP: ip, Mask: net.CIDRMask(prefixLen, prefixLen)}
 	return ip, ipNet, nil

tunnel/intercept/svcpoll.go

@@ -204,9 +204,7 @@ func (self *ServiceListener) removeService(svc *entities.Service) {
 			}
 		}
 
-		if previousService.StopHostHook != nil {
-			previousService.StopHostHook()
-		}
+		previousService.RunCleanupActions()
 
 		delete(self.services, svc.Id)
 	}
@@ -231,7 +229,7 @@ func (self *ServiceListener) host(svc *entities.Service, tracker AddressTracker)
 		}
 		self.healthCheckMgr.UnregisterServiceChecks(svc.Id)
 	}
-	svc.StopHostHook = stopHook
+	svc.AddCleanupAction(stopHook)
 
 	for idx, hostContext := range hostContexts {
 		hostControl, err := self.provider.HostService(hostContext)