diff --git a/tunnel/dns/file.go b/tunnel/dns/file.go index 0e648ff59..56e813b43 100644 --- a/tunnel/dns/file.go +++ b/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 diff --git a/tunnel/dns/refcount.go b/tunnel/dns/refcount.go new file mode 100644 index 000000000..9a07cf678 --- /dev/null +++ b/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() +} diff --git a/tunnel/dns/resolver.go b/tunnel/dns/resolver.go index 2897030b9..487ed0d00 100644 --- a/tunnel/dns/resolver.go +++ b/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 } diff --git a/tunnel/dns/server.go b/tunnel/dns/server.go index 2862b4017..42c8a791c 100644 --- a/tunnel/dns/server.go +++ b/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) diff --git a/tunnel/entities/service.go b/tunnel/entities/service.go index 5d1fb2e97..0815bda89 100644 --- a/tunnel/entities/service.go +++ b/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 { diff --git a/tunnel/intercept/interceptor.go b/tunnel/intercept/interceptor.go index 5d4bf7b53..92bc6d1d7 100644 --- a/tunnel/intercept/interceptor.go +++ b/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) } diff --git a/tunnel/intercept/iputils.go b/tunnel/intercept/iputils.go index 0aa21c1c9..9e3891aee 100644 --- a/tunnel/intercept/iputils.go +++ b/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 diff --git a/tunnel/intercept/svcpoll.go b/tunnel/intercept/svcpoll.go index ab6dcc2c8..9c4f87fee 100644 --- a/tunnel/intercept/svcpoll.go +++ b/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)