refactor DNS resolution and IP handling; remove unused utility

entrypoints/background.ts

@@ -4,51 +4,87 @@ let currentTabUrl: string | null = null
 
 async function resolveARecord(hostname: string): Promise<string | null> {
   try {
-    const dnsResponse = await fetch(
+    let dnsResponse = await fetch(
       `https://cloudflare-dns.com/dns-query?name=${hostname}&type=A`,
-      {
-        headers: { Accept: 'application/dns-json' },
-      }
+      { headers: { Accept: 'application/dns-json' } }
     )
-
-    if (!dnsResponse.ok) {
-      console.error(`DNS query failed: ${dnsResponse.status}`)
-      return null
+    if (dnsResponse.ok) {
+      const dnsData = await dnsResponse.json()
+      const aRecord = dnsData.Answer?.find(
+        (entry: DNSEntry) => entry.type === 1
+      )?.data
+      if (aRecord) return aRecord
     }
 
-    const dnsData = await dnsResponse.json()
-    return (
-      dnsData.Answer?.find((entry: DNSEntry) => entry.type === 1)?.data || null
+    dnsResponse = await fetch(
+      `https://cloudflare-dns.com/dns-query?name=${hostname}&type=AAAA`,
+      { headers: { Accept: 'application/dns-json' } }
     )
+    if (dnsResponse.ok) {
+      const dnsData = await dnsResponse.json()
+      const aaaaRecord = dnsData.Answer?.find(
+        (entry: DNSEntry) => entry.type === 28
+      )?.data
+      if (aaaaRecord) return aaaaRecord
+    }
+
+    return null
   } catch (error) {
     console.error('Failed to fetch DNS data:', error)
     return null
   }
 }
 
+function isIP(host: string): boolean {
+  const cleanedHost = host.replace(/^\[|\]$/g, '')
+
+  const ipv4Regex =
+    /^(25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.(25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.(25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.(25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)$/
+
+  const ipv6Regex =
+    /(([0-9a-fA-F]{1,4}:){7,7}[0-9a-fA-F]{1,4}|([0-9a-fA-F]{1,4}:){1,7}:|([0-9a-fA-F]{1,4}:){1,6}:[0-9a-fA-F]{1,4}|([0-9a-fA-F]{1,4}:){1,5}(:[0-9a-fA-F]{1,4}){1,2}|([0-9a-fA-F]{1,4}:){1,4}(:[0-9a-fA-F]{1,4}){1,3}|([0-9a-fA-F]{1,4}:){1,3}(:[0-9a-fA-F]{1,4}){1,4}|([0-9a-fA-F]{1,4}:){1,2}(:[0-9a-fA-F]{1,4}){1,5}|[0-9a-fA-F]{1,4}:((:[0-9a-fA-F]{1,4}){1,6})|:((:[0-9a-fA-F]{1,4}){1,7}|:)|fe80:(:[0-9a-fA-F]{0,4}){0,4}%[0-9a-zA-Z]{1,}|::(ffff(:0{1,4}){0,1}:){0,1}((25[0-5]|(2[0-4]|1{0,1}[0-9]){0,1}[0-9])\.){3,3}(25[0-5]|(2[0-4]|1{0,1}[0-9]){0,1}[0-9])|([0-9a-fA-F]{1,4}:){1,4}:((25[0-5]|(2[0-4]|1{0,1}[0-9]){0,1}[0-9])\.){3,3}(25[0-5]|(2[0-4]|1{0,1}[0-9]){0,1}[0-9]))/
+
+  return ipv4Regex.test(cleanedHost) || ipv6Regex.test(cleanedHost)
+}
+
+async function getIPInfo(host: string): Promise<any | null> {
+  const cleanedHost = host.replace(/^\[|\]$/g, '')
+
+  if (isIP(cleanedHost)) {
+    const response = await fetch(`https://ip.albert.lol/${cleanedHost}`)
+    const data = await response.json()
+    return data.ip ? data : null
+  }
+
+  const resolvedIp = await resolveARecord(cleanedHost)
+  if (!resolvedIp) return null
+
+  const response = await fetch(`https://ip.albert.lol/${resolvedIp}`)
+  return await response.json()
+}
+
 async function handleTabUpdate(url: string) {
   if (url === currentTabUrl) return
   currentTabUrl = url
 
   try {
     const hostname = new URL(url).hostname
-    const ip = await resolveARecord(hostname)
-
-    if (!ip) {
+    const apiData = await getIPInfo(hostname)
+    if (!apiData || !apiData.ip) {
       await updateIcon(null)
       return
     }
 
-    const apiResponse = await fetch(`https://ip.albert.lol/${ip}`)
-    const apiData = await apiResponse.json()
-
+    if (apiData.bogon === true) {
+      await updateIcon(null)
+      return
+    }
     const country = apiData.country || null
     const asn = apiData.org?.split(' ')[0]
     let iconCode = country
     if (!iconCode && asn === 'AS13335') {
       iconCode = 'cloudflare'
     }
-
     await updateIcon(iconCode)
   } catch (error) {
     console.error('Failed to handle tab update:', error)
@@ -72,18 +108,37 @@ export default defineBackground({
         if (request.type === 'FETCH_SERVER_INFO') {
           ;(async () => {
             try {
-              const ip = await resolveARecord(request.hostname)
-              if (!ip) {
+              const cleanHostname =
+                request.hostname.startsWith('[') &&
+                request.hostname.endsWith(']')
+                  ? request.hostname.slice(1, -1)
+                  : request.hostname
+
+              const apiData = await getIPInfo(cleanHostname)
+              if (!apiData || !apiData.ip) {
                 sendResponse({ error: 'DNS resolution failed', data: null })
                 return
               }
+              if (apiData.bogon === true) {
+                await updateIcon(null)
+                sendResponse({
+                  error: null,
+                  data: {
+                    origin: '',
+                    ip: cleanHostname,
+                    hostname: '',
+                    country: '',
+                    city: '',
+                    org: '',
+                    isLocal: true,
+                    isBrowserResource: false,
+                  },
+                })
+                return
+              }
 
-              const apiResponse = await fetch(`https://ip.albert.lol/${ip}`)
-              const apiData = await apiResponse.json()
-
-              const parsed = psl.parse(request.hostname)
+              const parsed = psl.parse(cleanHostname)
               const origin = 'domain' in parsed ? parsed.domain : null
-
               const country = apiData.country || null
               const asn = apiData.org?.split(' ')[0]
               let iconCode = country
@@ -91,7 +146,6 @@ export default defineBackground({
                 iconCode = 'cloudflare'
               }
               await updateIcon(iconCode)
-
               sendResponse({
                 error: null,
                 data: {
@@ -101,6 +155,8 @@ export default defineBackground({
                   country: apiData.country || null,
                   city: apiData.city || null,
                   org: apiData.org,
+                  isLocal: false,
+                  isBrowserResource: false,
                 },
               })
             } catch (error) {
@@ -113,7 +169,6 @@ export default defineBackground({
           })()
           return true
         }
-
         sendResponse({ error: 'Unknown request type', data: null })
         return true
       }

hooks/useTabData.ts

@@ -1,12 +1,3 @@
-import { useState, useEffect } from 'react'
-import { isPrivateIP } from '@/utils'
-import {
-  FetchServerInfoRequest,
-  FetchServerInfoResponse,
-  ServerData,
-} from '@/utils/model'
-import browser from 'webextension-polyfill'
-
 export function useTabData() {
   const [data, setData] = useState<ServerData | null>(null)
   const [loading, setLoading] = useState(true)
@@ -38,20 +29,6 @@ export function useTabData() {
           })
         }
 
-        const isInternal = isPrivateIP(hostname)
-        if (isInternal) {
-          return setData({
-            origin: '',
-            ip: hostname,
-            hostname: '',
-            country: '',
-            city: '',
-            org: '',
-            isLocal: true,
-            isBrowserResource: false,
-          })
-        }
-
         const response = await browser.runtime.sendMessage<
           FetchServerInfoRequest,
           FetchServerInfoResponse

utils/index.ts

@@ -1,36 +0,0 @@
-export function isPrivateIP(host: string): boolean {
-  try {
-    const rawIp = host.startsWith('[') ? host.slice(1, -1) : host
-
-    if (rawIp === 'localhost') return true
-
-    if (rawIp.includes('.')) {
-      const parts = rawIp.split('.').map(Number)
-      if (parts.length !== 4 || parts.some(isNaN)) return false
-
-      return (
-        parts[0] === 10 ||
-        (parts[0] === 172 && parts[1] >= 16 && parts[1] <= 31) ||
-        (parts[0] === 192 && parts[1] === 168) ||
-        (parts[0] === 169 && parts[1] === 254)
-      )
-    }
-
-    if (rawIp.includes(':')) {
-      const ip = rawIp.split('%')[0]
-
-      if (ip === '::1') return true
-
-      const firstHextet = parseInt(ip.split(':')[0], 16)
-      if ((firstHextet & 0xfe00) === 0xfc00) return true
-
-      if ((firstHextet & 0xffc0) === 0xfe80) return true
-
-      return false
-    }
-
-    return false
-  } catch {
-    return false
-  }
-}

utils/model.ts

@@ -15,11 +15,11 @@ export interface DNSEntry {
 }
 
 export interface FetchServerInfoRequest {
-  type: 'FETCH_SERVER_INFO';
-  hostname: string;
+  type: 'FETCH_SERVER_INFO'
+  hostname: string
 }
 
 export interface FetchServerInfoResponse {
-  error?: string;
-  data?: ServerData;
+  error?: string
+  data?: ServerData
 }