From 43cb48231ac5adc5404868ff61d044aeb34e0608 Mon Sep 17 00:00:00 2001
From: wwqgtxx <wwqgtxx@gmail.com>
Date: Thu, 26 Sep 2024 22:21:59 +0800
Subject: [PATCH] cache: add dns cache in udp packet sender

reduce the cost of re-resolving DNS for each packet received and prevent the target IP from jumping between multiple resolution results
---
 constant/adapters.go |  4 ++++
 tunnel/connection.go | 29 +++++++++++++++++++++++++----
 tunnel/tunnel.go     | 14 +-------------
 3 files changed, 30 insertions(+), 17 deletions(-)

diff --git a/constant/adapters.go b/constant/adapters.go
index 3731cd607..cb47f8716 100644
--- a/constant/adapters.go
+++ b/constant/adapters.go
@@ -298,7 +298,11 @@ type PacketSender interface {
 	// Send will send PacketAdapter nonblocking
 	// the implement must call UDPPacket.Drop() inside Send
 	Send(PacketAdapter)
+	// Process is a blocking loop to send PacketAdapter to PacketConn and update the WriteBackProxy
 	Process(PacketConn, WriteBackProxy)
+	// ResolveUDP do a local resolve UDP dns blocking if metadata is not resolved
+	ResolveUDP(*Metadata) error
+	// Close stop the Process loop
 	Close()
 }
 
diff --git a/tunnel/connection.go b/tunnel/connection.go
index 17e4efd04..1ea0678cd 100644
--- a/tunnel/connection.go
+++ b/tunnel/connection.go
@@ -7,7 +7,9 @@ import (
 	"net/netip"
 	"time"
 
+	"github.com/metacubex/mihomo/common/lru"
 	N "github.com/metacubex/mihomo/common/net"
+	"github.com/metacubex/mihomo/component/resolver"
 	C "github.com/metacubex/mihomo/constant"
 	"github.com/metacubex/mihomo/log"
 )
@@ -16,6 +18,7 @@ type packetSender struct {
 	ctx    context.Context
 	cancel context.CancelFunc
 	ch     chan C.PacketAdapter
+	cache  *lru.LruCache[string, netip.Addr]
 }
 
 // newPacketSender return a chan based C.PacketSender
@@ -27,6 +30,7 @@ func newPacketSender() C.PacketSender {
 		ctx:    ctx,
 		cancel: cancel,
 		ch:     ch,
+		cache:  lru.New[string, netip.Addr](lru.WithSize[string, netip.Addr](senderCapacity)),
 	}
 }
 
@@ -39,7 +43,11 @@ func (s *packetSender) Process(pc C.PacketConn, proxy C.WriteBackProxy) {
 			if proxy != nil {
 				proxy.UpdateWriteBack(packet)
 			}
-			_ = handleUDPToRemote(packet, pc, packet.Metadata())
+			if err := s.ResolveUDP(packet.Metadata()); err != nil {
+				log.Warnln("[UDP] Resolve Ip error: %s", err)
+			} else {
+				_ = handleUDPToRemote(packet, pc, packet.Metadata())
+			}
 			packet.Drop()
 		}
 	}
@@ -79,11 +87,24 @@ func (s *packetSender) Close() {
 	s.dropAll()
 }
 
-func handleUDPToRemote(packet C.UDPPacket, pc C.PacketConn, metadata *C.Metadata) error {
-	if err := resolveUDP(metadata); err != nil {
-		return err
+func (s *packetSender) ResolveUDP(metadata *C.Metadata) (err error) {
+	// local resolve UDP dns
+	if !metadata.Resolved() {
+		ip, ok := s.cache.Get(metadata.Host)
+		if !ok {
+			ip, err = resolver.ResolveIP(s.ctx, metadata.Host)
+			if err != nil {
+				return err
+			}
+			s.cache.Set(metadata.Host, ip)
+		}
+
+		metadata.DstIP = ip
 	}
+	return nil
+}
 
+func handleUDPToRemote(packet C.UDPPacket, pc C.PacketConn, metadata *C.Metadata) error {
 	addr := metadata.UDPAddr()
 	if addr == nil {
 		return errors.New("udp addr invalid")
diff --git a/tunnel/tunnel.go b/tunnel/tunnel.go
index af16e4aef..5c136eb24 100644
--- a/tunnel/tunnel.go
+++ b/tunnel/tunnel.go
@@ -346,18 +346,6 @@ func resolveMetadata(metadata *C.Metadata) (proxy C.Proxy, rule C.Rule, err erro
 	return
 }
 
-func resolveUDP(metadata *C.Metadata) error {
-	// local resolve UDP dns
-	if !metadata.Resolved() {
-		ip, err := resolver.ResolveIP(context.Background(), metadata.Host)
-		if err != nil {
-			return err
-		}
-		metadata.DstIP = ip
-	}
-	return nil
-}
-
 // processUDP starts a loop to handle udp packet
 func processUDP(queue chan C.PacketAdapter) {
 	for conn := range queue {
@@ -398,7 +386,7 @@ func handleUDPConn(packet C.PacketAdapter) {
 	sender, loaded := natTable.GetOrCreate(key, newPacketSender)
 	if !loaded {
 		dial := func() (C.PacketConn, C.WriteBackProxy, error) {
-			if err := resolveUDP(metadata); err != nil {
+			if err := sender.ResolveUDP(metadata); err != nil {
 				log.Warnln("[UDP] Resolve Ip error: %s", err)
 				return nil, nil, err
 			}