From 56879517f1cb731b237a7420fc8e458b2d9ea9a7 Mon Sep 17 00:00:00 2001
From: Zoltan Papp <zoltan.pmail@gmail.com>
Date: Wed, 16 Aug 2023 09:56:46 +0200
Subject: [PATCH] Fix eBPF apply logic

---
 client/internal/dns/service_listener.go | 35 ++++++++++++++++++-------
 1 file changed, 26 insertions(+), 9 deletions(-)

diff --git a/client/internal/dns/service_listener.go b/client/internal/dns/service_listener.go
index 228981504..cc9eb3e00 100644
--- a/client/internal/dns/service_listener.go
+++ b/client/internal/dns/service_listener.go
@@ -47,9 +47,6 @@ func newServiceViaListener(wgIface WGIface, customAddr *netip.AddrPort) *service
 		},
 	}
 
-	if runtime.GOOS == "linux" && customAddr != nil {
-		s.ebpfService = ebpf.GetEbpfManagerInstance()
-	}
 	return s
 }
 
@@ -69,6 +66,14 @@ func (s *serviceViaListener) Listen() error {
 	}
 	s.server.Addr = fmt.Sprintf("%s:%d", s.listenIP, s.listenPort)
 
+	if s.shouldApplyPortFwd() {
+		s.ebpfService = ebpf.GetEbpfManagerInstance()
+		err = s.ebpfService.LoadDNSFwd(s.listenIP, s.listenPort)
+		if err != nil {
+			log.Warnf("failed to load DNS port fwd, custom port may not support well: %s", err)
+			s.ebpfService = nil
+		}
+	}
 	log.Debugf("starting dns on %s", s.server.Addr)
 	go func() {
 		s.setListenerStatus(true)
@@ -80,12 +85,6 @@ func (s *serviceViaListener) Listen() error {
 		}
 	}()
 
-	if s.ebpfService != nil && s.listenPort != defaultPort {
-		err = s.ebpfService.LoadDNSFwd(s.listenIP, s.listenPort)
-		if err != nil {
-			return err
-		}
-	}
 	return nil
 }
 
@@ -122,6 +121,9 @@ func (s *serviceViaListener) DeregisterMux(pattern string) {
 }
 
 func (s *serviceViaListener) RuntimePort() int {
+	s.listenerFlagLock.Lock()
+	defer s.listenerFlagLock.Unlock()
+
 	if s.ebpfService != nil {
 		return defaultPort
 	} else {
@@ -168,3 +170,18 @@ func (s *serviceViaListener) evalRuntimeAddress() (string, int, error) {
 
 	return s.getFirstListenerAvailable()
 }
+
+func (s *serviceViaListener) shouldApplyPortFwd() bool {
+	if runtime.GOOS != "linux" {
+		return false
+	}
+
+	if s.customAddr != nil {
+		return false
+	}
+
+	if s.listenPort == defaultPort {
+		return false
+	}
+	return true
+}