Fixed some issue of local dns.

Author: masterking32

internal/client/dns_inflight.go

@@ -13,18 +13,25 @@ import (
 )
 
 type dnsInflightManager struct {
-	timeout time.Duration
-	mu      sync.Mutex
-	items   map[string]time.Time
+	timeout       time.Duration
+	cleanupWindow time.Duration
+	nextCleanupAt time.Time
+	mu            sync.Mutex
+	items         map[string]time.Time
 }
 
 func newDNSInflightManager(timeout time.Duration) *dnsInflightManager {
 	if timeout <= 0 {
 		timeout = 30 * time.Second
 	}
+	cleanupWindow := timeout / 4
+	if cleanupWindow < time.Second {
+		cleanupWindow = time.Second
+	}
 	return &dnsInflightManager{
-		timeout: timeout,
-		items:   make(map[string]time.Time),
+		timeout:       timeout,
+		cleanupWindow: cleanupWindow,
+		items:         make(map[string]time.Time),
 	}
 }
 
@@ -37,10 +44,13 @@ func (m *dnsInflightManager) Begin(cacheKey []byte, now time.Time) bool {
 	defer m.mu.Unlock()
 
 	key := string(cacheKey)
-	for existingKey, createdAt := range m.items {
-		if now.Sub(createdAt) >= m.timeout {
-			delete(m.items, existingKey)
+	if m.nextCleanupAt.IsZero() || !now.Before(m.nextCleanupAt) {
+		for existingKey, createdAt := range m.items {
+			if now.Sub(createdAt) >= m.timeout {
+				delete(m.items, existingKey)
+			}
 		}
+		m.nextCleanupAt = now.Add(m.cleanupWindow)
 	}
 
 	if createdAt, ok := m.items[key]; ok && now.Sub(createdAt) < m.timeout {

internal/client/dns_inflight_test.go

@@ -0,0 +1,43 @@
+// ==============================================================================
+// MasterDnsVPN
+// Author: MasterkinG32
+// Github: https://github.com/masterking32
+// Year: 2026
+// ==============================================================================
+
+package client
+
+import (
+	"testing"
+	"time"
+)
+
+func TestDNSInflightBeginDedupesWithinTimeout(t *testing.T) {
+	manager := newDNSInflightManager(10 * time.Second)
+	now := time.Unix(1700000000, 0)
+	key := []byte("example")
+
+	if !manager.Begin(key, now) {
+		t.Fatal("first begin should succeed")
+	}
+	if manager.Begin(key, now.Add(5*time.Second)) {
+		t.Fatal("second begin within timeout should be deduped")
+	}
+	if !manager.Begin(key, now.Add(11*time.Second)) {
+		t.Fatal("begin after timeout should succeed again")
+	}
+}
+
+func TestDNSInflightCompleteReleasesKey(t *testing.T) {
+	manager := newDNSInflightManager(10 * time.Second)
+	now := time.Unix(1700000000, 0)
+	key := []byte("example")
+
+	if !manager.Begin(key, now) {
+		t.Fatal("first begin should succeed")
+	}
+	manager.Complete(key)
+	if !manager.Begin(key, now.Add(time.Second)) {
+		t.Fatal("begin after complete should succeed")
+	}
+}

internal/client/local_dns.go

@@ -16,7 +16,8 @@ import (
 )
 
 type localDNSRequest struct {
-	packet []byte
+	buffer []byte
+	size   int
 	addr   *net.UDPAddr
 }
 
@@ -44,12 +45,18 @@ func (c *Client) RunLocalDNSListener(ctx context.Context) error {
 	)
 
 	queue := make(chan localDNSRequest, c.cfg.LocalDNSQueueSize)
+	packetPool := sync.Pool{
+		New: func() any {
+			return make([]byte, EDnsSafeUDPSize)
+		},
+	}
+
 	var workerWG sync.WaitGroup
 	for range c.cfg.LocalDNSWorkers {
 		workerWG.Add(1)
 		go func() {
 			defer workerWG.Done()
-			c.localDNSWorker(ctx, conn, queue)
+			c.localDNSWorker(ctx, conn, queue, &packetPool)
 		}()
 	}
 
@@ -59,28 +66,30 @@ func (c *Client) RunLocalDNSListener(ctx context.Context) error {
 	}()
 	go c.runLocalDNSCacheFlushLoop(ctx)
 
-	buffer := make([]byte, EDnsSafeUDPSize)
 	for {
+		buffer := packetPool.Get().([]byte)
 		n, addr, err := conn.ReadFromUDP(buffer)
 		if err != nil {
+			packetPool.Put(buffer)
 			if ctx.Err() != nil {
 				break
 			}
 			return err
 		}
 
-		packet := append([]byte(nil), buffer[:n]...)
 		select {
-		case queue <- localDNSRequest{packet: packet, addr: addr}:
+		case queue <- localDNSRequest{buffer: buffer, size: n, addr: addr}:
 		case <-ctx.Done():
+			packetPool.Put(buffer)
 			close(queue)
 			workerWG.Wait()
 			return nil
 		default:
-			response, _ := DnsParser.BuildServerFailureResponse(packet)
+			response, _ := DnsParser.BuildServerFailureResponse(buffer[:n])
 			if len(response) != 0 {
 				_, _ = conn.WriteToUDP(response, addr)
 			}
+			packetPool.Put(buffer)
 		}
 	}
 
@@ -89,7 +98,7 @@ func (c *Client) RunLocalDNSListener(ctx context.Context) error {
 	return nil
 }
 
-func (c *Client) localDNSWorker(ctx context.Context, conn *net.UDPConn, queue <-chan localDNSRequest) {
+func (c *Client) localDNSWorker(ctx context.Context, conn *net.UDPConn, queue <-chan localDNSRequest, packetPool *sync.Pool) {
 	for {
 		select {
 		case <-ctx.Done():
@@ -100,11 +109,18 @@ func (c *Client) localDNSWorker(ctx context.Context, conn *net.UDPConn, queue <-
 			}
 			func() {
 				defer func() {
-					if recover() != nil {
-						return
+					if packetPool != nil && req.buffer != nil {
+						packetPool.Put(req.buffer)
+					}
+					if recovered := recover(); recovered != nil && c != nil && c.log != nil {
+						c.log.Errorf(
+							"💥 <red>Local DNS Handler Panic Recovered</red> <magenta>|</magenta> <yellow>%v</yellow>",
+							recovered,
+						)
 					}
 				}()
-				response, _ := c.handleDNSQueryPacket(req.packet)
+
+				response, _ := c.handleDNSQueryPacket(req.buffer[:req.size])
 				if len(response) != 0 {
 					_, _ = conn.WriteToUDP(response, req.addr)
 				}