Penguinjanator
diff --git a/‎packages/shared-frontend-utils/src/networkUtil.test.ts‎
Lines changed: 169 additions & 0 deletions b/‎packages/shared-frontend-utils/src/networkUtil.test.ts‎
Lines changed: 169 additions & 0 deletions
diff --git a/‎packages/shared-frontend-utils/src/networkUtil.ts‎
Lines changed: 94 additions & 34 deletions b/‎packages/shared-frontend-utils/src/networkUtil.ts‎
Lines changed: 94 additions & 34 deletions
@@ -0,0 +1,169 @@
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
+
+import { getClientCountry, isInChina } from './networkUtil'
+
+const traceResponse = (body: string, ok = true) =>
+  ({ ok, text: () => Promise.resolve(body) }) as Response
+
+const TRACE_BODY = (loc: string) =>
+  `fl=1187f27\nh=cloud.comfy.org\nloc=${loc}\ntls=TLSv1.3\n`
+
+/** Blackholes: never settles and ignores the abort signal. */
+const neverSettles = () => new Promise<Response>(() => {})
+
+/** Resolves to the settled value, or to `PENDING` if the promise has not. */
+const settlementOf = <T>(promise: Promise<T>) =>
+  Promise.race([promise, Promise.resolve().then(() => 'PENDING' as const)])
+
+const fetchMock = vi.fn()
+
+beforeEach(() => {
+  vi.stubGlobal('fetch', fetchMock)
+  vi.stubGlobal('navigator', { language: 'en-US' })
+  fetchMock.mockReset()
+})
+
+afterEach(() => {
+  vi.unstubAllGlobals()
+  vi.useRealTimers()
+})
+
+describe('getClientCountry', () => {
+  it('returns the uppercased loc value from the edge trace', async () => {
+    fetchMock.mockResolvedValue(traceResponse(TRACE_BODY('cn')))
+
+    await expect(getClientCountry()).resolves.toBe('CN')
+  })
+
+  it.for([
+    [
+      'a non-200 response',
+      () => Promise.resolve(traceResponse('loc=CN', false))
+    ],
+    [
+      'a body with no loc line',
+      () => Promise.resolve(traceResponse('h=x\nts=1'))
+    ],
+    ['an empty loc value', () => Promise.resolve(traceResponse('loc=\n'))],
+    [
+      'an unknown-country sentinel',
+      () => Promise.resolve(traceResponse(TRACE_BODY('XX')))
+    ],
+    ['a Tor sentinel', () => Promise.resolve(traceResponse(TRACE_BODY('T1')))],
+    ['a network error', () => Promise.reject(new Error('offline'))]
+  ] as const)('returns undefined for %s', async ([, respond]) => {
+    fetchMock.mockImplementation(respond)
+
+    await expect(getClientCountry()).resolves.toBeUndefined()
+  })
+
+  it('gives up rather than hanging when the edge never answers', async () => {
+    vi.useFakeTimers()
+    fetchMock.mockImplementation(neverSettles)
+
+    const country = getClientCountry()
+    expect(await settlementOf(country)).toBe('PENDING')
+
+    await vi.advanceTimersByTimeAsync(2000)
+
+    expect(await settlementOf(country)).toBeUndefined()
+  })
+
+  it('gives up when the edge sends headers and then stalls the body', async () => {
+    vi.useFakeTimers()
+    fetchMock.mockResolvedValue({
+      ok: true,
+      text: () => new Promise<string>(() => {})
+    } as Response)
+
+    const country = getClientCountry()
+    expect(await settlementOf(country)).toBe('PENDING')
+
+    await vi.advanceTimersByTimeAsync(2000)
+
+    expect(await settlementOf(country)).toBeUndefined()
+  })
+})
+
+describe('isInChina', () => {
+  it('trusts the edge answer over the reachability heuristic', async () => {
+    fetchMock.mockResolvedValue(traceResponse(TRACE_BODY('CN')))
+
+    await expect(isInChina()).resolves.toBe(true)
+    expect(fetchMock).toHaveBeenCalledTimes(1)
+  })
+
+  it('reports outside China when the edge names another country', async () => {
+    fetchMock.mockResolvedValue(traceResponse(TRACE_BODY('IN')))
+
+    await expect(isInChina()).resolves.toBe(false)
+  })
+
+  it('reports inside China when only Baidu answers, and fast', async () => {
+    fetchMock.mockImplementation((url: string) =>
+      url.includes('cdn-cgi') || url.includes('google')
+        ? Promise.reject(new Error('blocked'))
+        : Promise.resolve({ ok: true } as Response)
+    )
+
+    await expect(
+      isInChina(),
+      'a China-routed client on a non-zh locale is only detectable by Baidu latency'
+    ).resolves.toBe(true)
+  })
+
+  it.for(['XX', 'T1'] as const)(
+    'falls back to the heuristic when the edge answers %s',
+    async (sentinel) => {
+      vi.stubGlobal('navigator', { language: 'zh-CN' })
+      fetchMock.mockImplementation((url: string) =>
+        url.includes('cdn-cgi')
+          ? Promise.resolve(traceResponse(TRACE_BODY(sentinel)))
+          : Promise.reject(new Error('blocked'))
+      )
+
+      await expect(
+        isInChina(),
+        'a sentinel names no country, so treating it as "not China" would wave through Tor and unknown-IP clients'
+      ).resolves.toBe(true)
+    }
+  )
+
+  it('reports outside China when the edge is down but Google answers', async () => {
+    fetchMock.mockImplementation((url: string) =>
+      url.includes('cdn-cgi')
+        ? Promise.reject(new Error('edge down'))
+        : Promise.resolve({ ok: true } as Response)
+    )
+
+    await expect(isInChina()).resolves.toBe(false)
+  })
+
+  it('falls back to the locale when neither probe answers', async () => {
+    fetchMock.mockRejectedValue(new Error('blocked'))
+
+    await expect(isInChina()).resolves.toBe(false)
+  })
+
+  it('falls back to the reachability heuristic when the edge cannot answer', async () => {
+    vi.stubGlobal('navigator', { language: 'zh-CN' })
+    fetchMock.mockImplementation((url: string) =>
+      url.includes('cdn-cgi')
+        ? Promise.reject(new Error('edge down'))
+        : Promise.reject(new Error('blocked'))
+    )
+
+    await expect(isInChina()).resolves.toBe(true)
+  })
+
+  it('settles even when every probe hangs forever', async () => {
+    vi.useFakeTimers()
+    fetchMock.mockImplementation(neverSettles)
+
+    const verdict = isInChina()
+    for (const _ of [0, 1, 2]) await vi.advanceTimersByTimeAsync(2000)
+
+    expect(await settlementOf(verdict)).toBe(false)
+    expect(fetchMock).toHaveBeenCalledTimes(3)
+  })
+})
@@ -4,59 +4,119 @@ const VALID_STATUS_CODES = [200, 201, 301, 302, 307, 308]
 export const checkUrlReachable = async (url: string): Promise<boolean> => {
   try {
     const response = await axios.head(url)
-    // Additional check for successful response
     return VALID_STATUS_CODES.includes(response.status)
   } catch {
     return false
   }
 }
 
 /**
- * Checks if the user is likely in mainland China by:
- * 1. Checking navigator.language
- * 2. Testing connectivity to commonly blocked services
- * 3. Testing latency to China-specific domains
+ * A CDN implementation detail, not a contract we own. The durable form is a
+ * first-party endpoint echoing `CF-IPCountry`.
  */
-export async function isInChina(): Promise<boolean> {
-  // Quick check based on language/locale
-  const isChineseLocale = navigator.language.toLowerCase().startsWith('zh-cn')
+const CLIENT_COUNTRY_URL = 'https://cloud.comfy.org/cdn-cgi/trace'
+
+/** Bounds every probe leg; two previously had none. */
+const PROBE_TIMEOUT_MS = 2000
+
+/** Baidu answering this fast implies a China route. */
+const CHINA_LATENCY_MS = 150
+
+/** `XX` is an unknown country and `T1` is Tor: neither names where the client is. */
+const UNRESOLVED_COUNTRIES = new Set(['XX', 'T1'])
+
+const parseTraceCountry = (body: string): string | undefined => {
+  const country =
+    body
+      .split('\n')
+      .find((line) => line.startsWith('loc='))
+      ?.slice('loc='.length)
+      .trim()
+      .toUpperCase() || undefined
+
+  return country && UNRESOLVED_COUNTRIES.has(country) ? undefined : country
+}
+
+/**
+ * The abort signal alone is not enough: a request the browser never resolves
+ * nor rejects would hang forever, so the deadline rejects independently.
+ *
+ * `read` runs inside the deadline so a response that sends headers and then
+ * stalls its body is bounded too.
+ */
+async function fetchWithin<T>(
+  url: string,
+  init: RequestInit,
+  read: (response: Response) => Promise<T>
+): Promise<T> {
+  const controller = new AbortController()
+  let expire: ReturnType<typeof setTimeout> | undefined
+
+  const deadline = new Promise<never>((_, reject) => {
+    expire = setTimeout(() => {
+      controller.abort()
+      reject(new Error(`Timed out after ${PROBE_TIMEOUT_MS}ms: ${url}`))
+    }, PROBE_TIMEOUT_MS)
+  })
 
   try {
-    // Test connectivity to Google - commonly blocked in China
-    const googleTest = await Promise.race([
-      fetch('https://www.google.com', {
-        mode: 'no-cors',
-        cache: 'no-cache'
-      }),
-      new Promise((_, reject) => setTimeout(() => reject(), 2000))
+    return await Promise.race([
+      fetch(url, { ...init, signal: controller.signal }).then(read),
+      deadline
     ])
+  } finally {
+    clearTimeout(expire)
+  }
+}
 
-    // If Google is accessible, user is likely not in China
-    if (googleTest) {
-      return false
-    }
+/** ISO country from the CDN edge, or `undefined` when it cannot answer. */
+export async function getClientCountry(): Promise<string | undefined> {
+  try {
+    const body = await fetchWithin(
+      CLIENT_COUNTRY_URL,
+      { cache: 'no-store' },
+      async (response) => (response.ok ? response.text() : undefined)
+    )
+    return body === undefined ? undefined : parseTraceCountry(body)
   } catch {
-    // Google is not accessible - potential indicator of being in China
-    if (isChineseLocale) {
-      return true
-    }
+    return undefined
+  }
+}
+
+const probe = (url: string) =>
+  fetchWithin(url, { mode: 'no-cors', cache: 'no-cache' }, async () => {})
+
+/**
+ * Fallback for when the edge cannot answer. Unsound both ways: a VPN user in
+ * China reaches Google, and a `zh-CN` user anywhere is blocked whenever Google
+ * is briefly unreachable.
+ */
+async function isInChinaByProbe(): Promise<boolean> {
+  const isChineseLocale = navigator.language.toLowerCase().startsWith('zh-cn')
+
+  try {
+    await probe('https://www.google.com')
+    return false
+  } catch {
+    if (isChineseLocale) return true
 
-    // Additional check - test latency to a reliable Chinese domain
     try {
       const start = performance.now()
-      await fetch('https://www.baidu.com', {
-        mode: 'no-cors',
-        cache: 'no-cache'
-      })
-      const latency = performance.now() - start
-
-      // If Baidu responds quickly (<150ms), user is likely in China
-      return latency < 150
+      await probe('https://www.baidu.com')
+      return performance.now() - start < CHINA_LATENCY_MS
     } catch {
-      // If both tests fail, default to locale check
       return isChineseLocale
     }
   }
+}
+
+/**
+ * Prefers the edge's geo-IP, falling back to the heuristic. Always settles, so
+ * callers must not add a timeout that could pre-empt a slow but real answer.
+ */
+export async function isInChina(): Promise<boolean> {
+  const country = await getClientCountry()
+  if (country !== undefined) return country === 'CN'
 
-  return false
+  return isInChinaByProbe()
 }