From 2b4741fbc71c1cc8ed1e4e94b65288337daf880b Mon Sep 17 00:00:00 2001
From: wwqgtxx <wwqgtxx@gmail.com>
Date: Wed, 12 Jun 2024 17:09:21 +0800
Subject: [PATCH] chore: add inner dns proxied connection statistic to restful
 api

---
 dns/util.go | 56 ++++++++++++++++++++++++++++++++++++-----------------
 1 file changed, 38 insertions(+), 18 deletions(-)

diff --git a/dns/util.go b/dns/util.go
index 9b84a5075..f13279fac 100644
--- a/dns/util.go
+++ b/dns/util.go
@@ -18,6 +18,7 @@ import (
 	C "github.com/metacubex/mihomo/constant"
 	"github.com/metacubex/mihomo/log"
 	"github.com/metacubex/mihomo/tunnel"
+	"github.com/metacubex/mihomo/tunnel/statistic"
 
 	D "github.com/miekg/dns"
 	"github.com/samber/lo"
@@ -204,6 +205,7 @@ func getDialHandler(r *Resolver, proxyAdapter C.ProxyAdapter, proxyName string,
 				}
 			}
 
+			var rule C.Rule
 			if proxyAdapter == nil {
 				if proxyName == RespectRules {
 					if !metadata.Resolved() {
@@ -214,7 +216,7 @@ func getDialHandler(r *Resolver, proxyAdapter C.ProxyAdapter, proxyName string,
 						}
 						metadata.DstIP = dstIP
 					}
-					proxyAdapter, _, err = tunnel.ResolveMetadata(metadata)
+					proxyAdapter, rule, err = tunnel.ResolveMetadata(metadata)
 					if err != nil {
 						return nil, err
 					}
@@ -228,22 +230,30 @@ func getDialHandler(r *Resolver, proxyAdapter C.ProxyAdapter, proxyName string,
 			}
 
 			if strings.Contains(network, "tcp") {
-				if proxyAdapter != nil {
-					if proxyAdapter.IsL3Protocol(metadata) { // L3 proxy should resolve domain before to avoid loopback
-						if !metadata.Resolved() {
-							dstIP, err := resolver.ResolveIPWithResolver(ctx, metadata.Host, r)
-							if err != nil {
-								return nil, err
-							}
-							metadata.DstIP = dstIP
-						}
-						metadata.Host = "" // clear host to avoid double resolve in proxy
-					}
-					log.Debugln("%s", metadata.RemoteAddress())
-					return proxyAdapter.DialContext(ctx, metadata, opts...)
+				if proxyAdapter == nil {
+					opts = append(opts, dialer.WithResolver(r))
+					return dialer.DialContext(ctx, network, addr, opts...)
 				}
-				opts = append(opts, dialer.WithResolver(r))
-				return dialer.DialContext(ctx, network, addr, opts...)
+
+				if proxyAdapter.IsL3Protocol(metadata) { // L3 proxy should resolve domain before to avoid loopback
+					if !metadata.Resolved() {
+						dstIP, err := resolver.ResolveIPWithResolver(ctx, metadata.Host, r)
+						if err != nil {
+							return nil, err
+						}
+						metadata.DstIP = dstIP
+					}
+					metadata.Host = "" // clear host to avoid double resolve in proxy
+				}
+
+				conn, err := proxyAdapter.DialContext(ctx, metadata, opts...)
+				if err != nil {
+					return nil, err
+				}
+
+				conn = statistic.NewTCPTracker(conn, statistic.DefaultManager, metadata, rule, 0, 0, false)
+
+				return conn, nil
 			} else {
 				if proxyAdapter == nil {
 					return dialer.DialContext(ctx, network, addr, opts...)
@@ -258,6 +268,8 @@ func getDialHandler(r *Resolver, proxyAdapter C.ProxyAdapter, proxyName string,
 					return nil, err
 				}
 
+				packetConn = statistic.NewUDPTracker(packetConn, statistic.DefaultManager, metadata, rule, 0, 0, false)
+
 				return N.NewBindPacketConn(packetConn, metadata.UDPAddr()), nil
 			}
 		}
@@ -281,9 +293,10 @@ func listenPacket(ctx context.Context, proxyAdapter C.ProxyAdapter, proxyName st
 		metadata.DstIP = dstIP
 	}
 
+	var rule C.Rule
 	if proxyAdapter == nil {
 		if proxyName == RespectRules {
-			proxyAdapter, _, err = tunnel.ResolveMetadata(metadata)
+			proxyAdapter, rule, err = tunnel.ResolveMetadata(metadata)
 			if err != nil {
 				return nil, err
 			}
@@ -304,7 +317,14 @@ func listenPacket(ctx context.Context, proxyAdapter C.ProxyAdapter, proxyName st
 		return nil, fmt.Errorf("proxy adapter [%s] UDP is not supported", proxyAdapter)
 	}
 
-	return proxyAdapter.ListenPacketContext(ctx, metadata, opts...)
+	packetConn, err := proxyAdapter.ListenPacketContext(ctx, metadata, opts...)
+	if err != nil {
+		return nil, err
+	}
+
+	packetConn = statistic.NewUDPTracker(packetConn, statistic.DefaultManager, metadata, rule, 0, 0, false)
+
+	return packetConn, nil
 }
 
 func batchExchange(ctx context.Context, clients []dnsClient, m *D.Msg) (msg *D.Msg, cache bool, err error) {