From b9a53442b27e5f830b8ce687f72147f00f7487ff Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E7=8E=8B=E8=89=AF?= <841369634@qq.com>
Date: Sat, 14 Sep 2024 17:29:07 +0800
Subject: [PATCH] =?UTF-8?q?optimize:=20=E4=BB=A3=E7=90=86=E8=AF=B7?=
 =?UTF-8?q?=E6=B1=82=E4=BD=BF=E7=94=A8=E4=BA=86DNS=E8=8E=B7=E5=8F=96?=
 =?UTF-8?q?=E5=88=B0=E7=9A=84IP=E6=97=B6=EF=BC=8C=E6=B7=BB=E5=8A=A0?=
 =?UTF-8?q?=E5=93=8D=E5=BA=94=E5=A4=B4=20`DS-Lookup`=EF=BC=8C=E8=BF=94?=
 =?UTF-8?q?=E5=9B=9E=E5=BD=93=E5=89=8D=E6=89=80=E4=BD=BF=E7=94=A8=E7=9A=84?=
 =?UTF-8?q?IP=E3=80=82?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../mitmproxy/src/lib/proxy/mitmproxy/createConnectHandler.js | 2 +-
 .../mitmproxy/src/lib/proxy/mitmproxy/createRequestHandler.js | 2 +-
 packages/mitmproxy/src/lib/proxy/mitmproxy/dnsLookup.js       | 4 +++-
 3 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/packages/mitmproxy/src/lib/proxy/mitmproxy/createConnectHandler.js b/packages/mitmproxy/src/lib/proxy/mitmproxy/createConnectHandler.js
index 05963b0..469256c 100644
--- a/packages/mitmproxy/src/lib/proxy/mitmproxy/createConnectHandler.js
+++ b/packages/mitmproxy/src/lib/proxy/mitmproxy/createConnectHandler.js
@@ -62,7 +62,7 @@ function connect (req, cltSocket, head, hostname, port, dnsConfig) {
     if (dnsConfig && dnsConfig.providers) {
       const dns = DnsUtil.hasDnsLookup(dnsConfig, hostname)
       if (dns) {
-        options.lookup = dnsLookup.createLookupFunc(dns, 'connect', hostport, isDnsIntercept)
+        options.lookup = dnsLookup.createLookupFunc(null, dns, 'connect', hostport, isDnsIntercept)
       }
     }
     const proxySocket = net.connect(options, () => {
diff --git a/packages/mitmproxy/src/lib/proxy/mitmproxy/createRequestHandler.js b/packages/mitmproxy/src/lib/proxy/mitmproxy/createRequestHandler.js
index 1a25321..492e52b 100644
--- a/packages/mitmproxy/src/lib/proxy/mitmproxy/createRequestHandler.js
+++ b/packages/mitmproxy/src/lib/proxy/mitmproxy/createRequestHandler.js
@@ -122,7 +122,7 @@ module.exports = function createRequestHandler (createIntercepts, middlewares, e
               }
             }
             if (dns) {
-              rOptions.lookup = dnsLookup.createLookupFunc(dns, 'request url', url, isDnsIntercept)
+              rOptions.lookup = dnsLookup.createLookupFunc(res, dns, 'request url', url, isDnsIntercept)
             }
           }
 
diff --git a/packages/mitmproxy/src/lib/proxy/mitmproxy/dnsLookup.js b/packages/mitmproxy/src/lib/proxy/mitmproxy/dnsLookup.js
index e434824..9efe2f4 100644
--- a/packages/mitmproxy/src/lib/proxy/mitmproxy/dnsLookup.js
+++ b/packages/mitmproxy/src/lib/proxy/mitmproxy/dnsLookup.js
@@ -3,7 +3,7 @@ const log = require('../../../utils/util.log')
 const defaultDns = require('dns')
 
 module.exports = {
-  createLookupFunc: function (dns, action, target, isDnsIntercept) {
+  createLookupFunc: function (res, dns, action, target, isDnsIntercept) {
     target = target ? (', target: ' + target) : ''
 
     return (hostname, options, callback) => {
@@ -12,6 +12,7 @@ module.exports = {
         const aliveIpObj = tester.pickFastAliveIpObj()
         if (aliveIpObj) {
           log.info(`----- ${action}: ${hostname}, use alive ip from dns '${aliveIpObj.dns}': ${aliveIpObj.host}${target} -----`)
+          if (res) res.setHeader('DS-Lookup', `IpTester: ${aliveIpObj.host}(${aliveIpObj.dns})`)
           callback(null, aliveIpObj.host, 4)
           return
         } else {
@@ -41,6 +42,7 @@ module.exports = {
           }
           if (isTestFailedIp === false) {
             log.info(`----- ${action}: ${hostname}, use ip from dns '${dns.name}': ${ip}${target} -----`)
+            if (res) res.setHeader('DS-Lookup', `DNS: ${ip.host}(${dns.name})`)
             callback(null, ip, 4)
             return
           } else {