From 8204bba79f49f83a9e404f69d0c453a36911b607 Mon Sep 17 00:00:00 2001
From: decolua <decoluadt@example.com>
Date: Tue, 28 Apr 2026 10:08:57 +0700
Subject: [PATCH] Refactor MitmServerCard to use input field for API key
 selection and enhance shutdown process to remove DNS entries synchronously.
 Added removeAllDNSEntriesSync function for safe cleanup during shutdown.

---
 .../cli-tools/components/MitmServerCard.js    | 26 +++++++++----------
 src/mitm/dns/dnsConfig.js                     | 24 +++++++++++++++++
 src/mitm/server.js                            | 11 +++++++-
 3 files changed, 46 insertions(+), 15 deletions(-)

diff --git a/src/app/(dashboard)/dashboard/cli-tools/components/MitmServerCard.js b/src/app/(dashboard)/dashboard/cli-tools/components/MitmServerCard.js
index 63fa42b8..8766f7b3 100644
--- a/src/app/(dashboard)/dashboard/cli-tools/components/MitmServerCard.js
+++ b/src/app/(dashboard)/dashboard/cli-tools/components/MitmServerCard.js
@@ -176,22 +176,20 @@ export default function MitmServerCard({ apiKeys, cloudEnabled, onStatusChange }
               <div className="flex items-center gap-2">
                 <span className="w-32 shrink-0 text-sm font-semibold text-text-main text-right">API Key</span>
                 <span className="material-symbols-outlined text-text-muted text-[14px]">arrow_forward</span>
-                {apiKeys?.length > 0 ? (
-                  <select
-                    value={selectedApiKey}
-                    onChange={(e) => setSelectedApiKey(e.target.value)}
-                    className="flex-1 min-w-0 px-2 py-1.5 bg-surface rounded text-xs border border-border text-text-main focus:outline-none focus:ring-1 focus:ring-primary/50"
-                  >
+                <input
+                  type="text"
+                  list="mitm-api-keys"
+                  value={selectedApiKey}
+                  onChange={(e) => setSelectedApiKey(e.target.value)}
+                  placeholder={cloudEnabled ? "Enter or pick API key" : "sk_9router (default)"}
+                  className="flex-1 min-w-0 px-2 py-1.5 bg-surface rounded border border-border text-xs text-text-main focus:outline-none focus:ring-1 focus:ring-primary/50"
+                />
+                {apiKeys?.length > 0 && (
+                  <datalist id="mitm-api-keys">
                     {apiKeys.map((key) => (
-                      <option key={key.id} value={key.key}>
-                        {key.key}
-                      </option>
+                      <option key={key.id} value={key.key}>{key.name || key.key}</option>
                     ))}
-                  </select>
-                ) : (
-                  <span className="flex-1 px-2 py-1.5 text-xs text-text-muted">
-                    {cloudEnabled ? "No API keys — create one in Keys page" : "sk_9router (default)"}
-                  </span>
+                  </datalist>
                 )}
               </div>
             )}
diff --git a/src/mitm/dns/dnsConfig.js b/src/mitm/dns/dnsConfig.js
index 70d00487..39af3364 100644
--- a/src/mitm/dns/dnsConfig.js
+++ b/src/mitm/dns/dnsConfig.js
@@ -214,11 +214,35 @@ async function removeAllDNSEntries(sudoPassword) {
   }
 }
 
+/**
+ * Sync removal of ALL tool DNS entries — for use during process shutdown
+ * when async ops aren't safe. Assumes caller already has root/admin rights.
+ */
+function removeAllDNSEntriesSync() {
+  try {
+    if (!fs.existsSync(HOSTS_FILE)) return;
+    const allHosts = Object.values(TOOL_HOSTS).flat();
+    const content = fs.readFileSync(HOSTS_FILE, "utf8");
+    const eol = IS_WIN ? "\r\n" : "\n";
+    const filtered = content.split(/\r?\n/).filter(l => !allHosts.some(h => l.includes(h))).join(eol);
+    if (filtered === content) return;
+    fs.writeFileSync(HOSTS_FILE, filtered, "utf8");
+    if (IS_WIN) {
+      try { execSync("ipconfig /flushdns", { windowsHide: true, stdio: "ignore" }); } catch { /* ignore */ }
+    } else if (IS_MAC) {
+      try { execSync("dscacheutil -flushcache && killall -HUP mDNSResponder", { stdio: "ignore" }); } catch { /* ignore */ }
+    } else {
+      try { execSync("resolvectl flush-caches 2>/dev/null || true", { stdio: "ignore" }); } catch { /* ignore */ }
+    }
+  } catch { /* best effort during shutdown */ }
+}
+
 module.exports = {
   TOOL_HOSTS,
   addDNSEntry,
   removeDNSEntry,
   removeAllDNSEntries,
+  removeAllDNSEntriesSync,
   execWithPassword,
   isSudoAvailable,
   executeElevatedPowerShell,
diff --git a/src/mitm/server.js b/src/mitm/server.js
index c444dacf..78e3362d 100644
--- a/src/mitm/server.js
+++ b/src/mitm/server.js
@@ -266,7 +266,16 @@ server.on("error", (e) => {
   process.exit(1);
 });
 
-const shutdown = () => server.close(() => process.exit(0));
+const { removeAllDNSEntriesSync } = require("./dns/dnsConfig");
+let isShuttingDown = false;
+const shutdown = () => {
+  if (isShuttingDown) return;
+  isShuttingDown = true;
+  // Strip tool hosts from /etc/hosts so other apps aren't broken after exit
+  removeAllDNSEntriesSync();
+  const forceExit = setTimeout(() => process.exit(0), 1500);
+  server.close(() => { clearTimeout(forceExit); process.exit(0); });
+};
 process.on("SIGTERM", shutdown);
 process.on("SIGINT", shutdown);
 if (process.platform === "win32") process.on("SIGBREAK", shutdown);