Refactor: improve check domain alive

This commit is contained in:
SukkaW 2024-11-06 21:03:26 +08:00
parent 35ab5ce23d
commit ad5b2ffdd5

View File

@ -46,7 +46,16 @@ const dohServers: Array<[string, DNS2.DnsResolver]> = ([
}) })
] as const); ] as const);
const queue = newQueue(8); const queue = newQueue(20);
const mutex = new Map<string, Promise<unknown>>();
function keyedAsyncMutexWithQueue<T>(key: string, fn: () => Promise<T>) {
if (mutex.has(key)) {
return mutex.get(key) as Promise<T>;
}
const promise = queue.add(() => fn()).finally(() => mutex.delete(key));
mutex.set(key, promise);
return promise;
}
class DnsError extends Error { class DnsError extends Error {
name = 'DnsError'; name = 'DnsError';
@ -128,7 +137,7 @@ async function isApexDomainAlive(apexDomain: string): Promise<[string, boolean]>
if (Object.keys(whois).length > 0) { if (Object.keys(whois).length > 0) {
// TODO: this is a workaround for https://github.com/LayeredStudio/whoiser/issues/117 // TODO: this is a workaround for https://github.com/LayeredStudio/whoiser/issues/117
if ('text' in whois && Array.isArray(whois.text) && whois.text.some(value => whoisNotFoundKeywordTest(value.toLowerCase()))) { if ('text' in whois && Array.isArray(whois.text) && whois.text.some(value => whoisNotFoundKeywordTest(value.toLowerCase()))) {
console.log(picocolors.red('[domain dead]'), 'whois no match', { domain: apexDomain }); console.log(picocolors.red('[domain dead]'), 'whois not found', { domain: apexDomain });
domainAliveMap.set(apexDomain, false); domainAliveMap.set(apexDomain, false);
return [apexDomain, false]; return [apexDomain, false];
} }
@ -140,35 +149,24 @@ async function isApexDomainAlive(apexDomain: string): Promise<[string, boolean]>
console.log({ whois }); console.log({ whois });
} }
console.log(picocolors.red('[domain dead]'), 'whois no match', { domain: apexDomain }); console.log(picocolors.red('[domain dead]'), 'whois not found', { domain: apexDomain });
domainAliveMap.set(apexDomain, false); domainAliveMap.set(apexDomain, false);
return [apexDomain, false]; return [apexDomain, false];
} }
const domainMutex = new Map<string, Promise<[string, boolean]>>();
export async function isDomainAlive(domain: string, isSuffix: boolean): Promise<[string, boolean]> { export async function isDomainAlive(domain: string, isSuffix: boolean): Promise<[string, boolean]> {
if (domain[0] === '.') { if (domainAliveMap.has(domain)) {
domain = domain.slice(1); return [domain, domainAliveMap.get(domain)!];
} }
const apexDomain = tldts.getDomain(domain, looseTldtsOpt); const apexDomain = tldts.getDomain(domain, looseTldtsOpt);
if (!apexDomain) { if (!apexDomain) {
console.log('[domain invalid]', 'no apex domain', { domain }); console.log('[domain invalid]', 'no apex domain', { domain });
domainAliveMap.set(domain, true);
return [domain, true] as const; return [domain, true] as const;
} }
let apexDomainAlivePromise; const apexDomainAlive = await keyedAsyncMutexWithQueue(apexDomain, () => isApexDomainAlive(apexDomain));
if (domainMutex.has(domain)) {
apexDomainAlivePromise = domainMutex.get(domain)!;
} else {
apexDomainAlivePromise = queue.add(() => isApexDomainAlive(apexDomain).then(res => {
domainMutex.delete(domain);
return res;
}));
domainMutex.set(domain, apexDomainAlivePromise);
}
const apexDomainAlive = await apexDomainAlivePromise;
if (!apexDomainAlive[1]) { if (!apexDomainAlive[1]) {
domainAliveMap.set(domain, false); domainAliveMap.set(domain, false);
@ -176,12 +174,14 @@ export async function isDomainAlive(domain: string, isSuffix: boolean): Promise<
} }
if (!isSuffix) { if (!isSuffix) {
const aRecords = (await resolve(domain, 'A')); const $domain = domain[0] === '.' ? domain.slice(1) : domain;
const aRecords = (await resolve($domain, 'A'));
if (aRecords.answers.length === 0) { if (aRecords.answers.length === 0) {
const aaaaRecords = (await resolve(domain, 'AAAA')); const aaaaRecords = (await resolve($domain, 'AAAA'));
if (aaaaRecords.answers.length === 0) { if (aaaaRecords.answers.length === 0) {
console.log(picocolors.red('[domain dead]'), 'no A/AAAA records', { domain, a: aRecords.dns, aaaa: aaaaRecords.dns }); console.log(picocolors.red('[domain dead]'), 'no A/AAAA records', { domain, a: aRecords.dns, aaaa: aaaaRecords.dns });
domainAliveMap.set(domain, false); domainAliveMap.set($domain, false);
return [domain, false] as const; return [domain, false] as const;
} }
} }
@ -201,14 +201,8 @@ export async function runAgainstRuleset(filepath: string) {
switch (type) { switch (type) {
case 'DOMAIN-SUFFIX': case 'DOMAIN-SUFFIX':
case 'DOMAIN': { case 'DOMAIN': {
if (!domainMutex.has(domain)) { promises.push(keyedAsyncMutexWithQueue(domain, () => isDomainAlive(domain, type === 'DOMAIN-SUFFIX')));
const promise = queue.add(() => isDomainAlive(domain, type === 'DOMAIN-SUFFIX')).then(res => {
domainMutex.delete(domain);
return res;
});
domainMutex.set(domain, promise);
promises.push(promise);
}
break; break;
} }
// no default // no default