fix(sandbox): don't answer a DNS query after the proxy has closed

packages/core/src/sandbox/dns-proxy.test.ts

@@ -101,3 +101,128 @@ describe('startDnsProxy', () => {
     proxy = null;
   });
 });
+
+// Regression: closing the proxy while a forward is awaiting its upstream reply
+// used to call send() on the already-closed server socket. dgram throws
+// ERR_SOCKET_DGRAM_NOT_RUNNING *synchronously* from inside the upstream
+// 'message' handler, i.e. after the enclosing Promise executor's synchronous
+// phase, so the promise never catches it and it surfaces as an uncaught
+// exception that fails the whole vitest run.
+describe('startDnsProxy shutdown race', () => {
+  /** A stub upstream resolver that answers `delayMs` after the query arrives. */
+  function startStubUpstream(delayMs: number): Promise<{
+    port: number;
+    close: () => Promise<void>;
+  }> {
+    const sock = createSocket('udp4');
+    const timers = new Set<ReturnType<typeof setTimeout>>();
+    sock.on('message', (msg, rinfo) => {
+      const timer = setTimeout(() => {
+        timers.delete(timer);
+        try {
+          sock.send(buildNxDomain(msg), rinfo.port, rinfo.address);
+        } catch {
+          // Stub already torn down.
+        }
+      }, delayMs);
+      timers.add(timer);
+    });
+    return new Promise((resolve, reject) => {
+      sock.once('error', reject);
+      sock.bind(0, '127.0.0.1', () => {
+        sock.removeListener('error', reject);
+        resolve({
+          port: sock.address().port,
+          close: () =>
+            new Promise<void>((done) => {
+              for (const timer of timers) clearTimeout(timer);
+              timers.clear();
+              sock.close(() => done());
+            }),
+        });
+      });
+    });
+  }
+
+  /**
+   * Run `body` with vitest's own uncaughtException handlers detached, and
+   * report whatever escaped. Without that swap the process-level handler turns
+   * a reproduction into a run-level failure instead of a clean assertion.
+   */
+  async function captureUncaught(body: () => Promise<void>): Promise<unknown[]> {
+    const escaped: unknown[] = [];
+    const capture = (err: unknown): void => {
+      escaped.push(err);
+    };
+    const prior = process.listeners('uncaughtException');
+    process.removeAllListeners('uncaughtException');
+    process.on('uncaughtException', capture);
+    try {
+      await body();
+      // Let the late upstream reply land while our handler is still installed.
+      await new Promise((r) => setTimeout(r, 150));
+    } finally {
+      process.removeListener('uncaughtException', capture);
+      for (const listener of prior) {
+        process.on('uncaughtException', listener as (err: Error) => void);
+      }
+    }
+    return escaped;
+  }
+
+  it('does not throw when the upstream reply lands after close()', async () => {
+    const upstream = await startStubUpstream(120);
+    try {
+      const escaped = await captureUncaught(async () => {
+        const proxy = await startDnsProxy({
+          allowedDomains: ['github.com'],
+          upstream: '127.0.0.1',
+          upstreamPort: upstream.port,
+          log: () => {},
+        });
+        const client = createSocket('udp4');
+        await new Promise<void>((resolve, reject) => {
+          client.send(buildQuery('github.com'), proxy.port, '127.0.0.1', (err) =>
+            err ? reject(err) : resolve(),
+          );
+        });
+        // Give the forward time to reach the stub, then close mid-flight.
+        await new Promise((r) => setTimeout(r, 40));
+        await proxy.close();
+        client.close();
+      });
+      expect(escaped).toEqual([]);
+    } finally {
+      await upstream.close();
+    }
+  });
+
+  it('close() drops the pending upstream socket and its timeout timer', async () => {
+    // A stub that never answers in time: without cleanup the 5s upstream timer
+    // (and its socket) outlive close() and keep the event loop alive.
+    const upstream = await startStubUpstream(60_000);
+    try {
+      const proxy = await startDnsProxy({
+        allowedDomains: ['github.com'],
+        upstream: '127.0.0.1',
+        upstreamPort: upstream.port,
+        log: () => {},
+      });
+      const client = createSocket('udp4');
+      await new Promise<void>((resolve, reject) => {
+        client.send(buildQuery('github.com'), proxy.port, '127.0.0.1', (err) =>
+          err ? reject(err) : resolve(),
+        );
+      });
+      await new Promise((r) => setTimeout(r, 40));
+      const before = process.getActiveResourcesInfo().length;
+      await proxy.close();
+      client.close();
+      const after = process.getActiveResourcesInfo().length;
+      // Server socket + upstream socket + the 5s timer all released.
+      expect(after).toBeLessThan(before);
+    } finally {
+      await upstream.close();
+    }
+  });
+});

packages/core/src/sandbox/dns-proxy.ts

@@ -17,6 +17,8 @@ export interface DnsProxyOpts {
   allowedDomains: string[];
   /** Upstream DNS server for allowed lookups (default 1.1.1.1). */
   upstream?: string;
+  /** Upstream DNS port; default 53. Overridable so tests can run a local stub. */
+  upstreamPort?: number;
   /** Bind address; default 127.0.0.1. */
   bindAddr?: string;
   /** Bind port; default 0 (random). */
@@ -32,11 +34,21 @@ export interface DnsProxyHandle {
   close: () => Promise<void>;
 }
 
+/** Shutdown bookkeeping shared between the server socket and in-flight forwards. */
+interface ProxyState {
+  /** Set synchronously by close(), before the server socket is torn down. */
+  closed: boolean;
+  /** Abandon callbacks, keyed by upstream socket, for forwards still in flight. */
+  pending: Map<Socket, () => void>;
+}
+
 export async function startDnsProxy(opts: DnsProxyOpts): Promise<DnsProxyHandle> {
   const allowed = new Set(opts.allowedDomains.map((d) => d.toLowerCase()));
   const upstream = opts.upstream ?? '1.1.1.1';
+  const upstreamPort = opts.upstreamPort ?? 53;
   const log = opts.log ?? (() => {});
   const sock = createSocket('udp4');
+  const state: ProxyState = { closed: false, pending: new Map() };
 
   sock.on('message', (msg, rinfo) => {
     const domain = parseQName(msg);
@@ -51,8 +63,11 @@ export async function startDnsProxy(opts: DnsProxyOpts): Promise<DnsProxyHandle>
       return;
     }
     log(`[dns-proxy] ALLOW ${norm} → ${upstream}`);
-    forward(sock, msg, rinfo, upstream).catch((err: Error) => {
+    forward(sock, msg, rinfo, upstream, upstreamPort, state).catch((err: Error) => {
       log(`[dns-proxy] forward error: ${err.message}`);
+      // Same shutdown hazard as inside forward(): a rejection can land after
+      // close(), and send() on a closed socket throws synchronously here.
+      if (state.closed) return;
       sock.send(buildNxDomain(msg), rinfo.port, rinfo.address);
     });
   });
@@ -70,6 +85,11 @@ export async function startDnsProxy(opts: DnsProxyOpts): Promise<DnsProxyHandle>
     port,
     close: () =>
       new Promise<void>((resolve) => {
+        // Flag first: forwards check this before touching `sock`, and both run
+        // on the same thread, so nothing can slip between the check and a send.
+        state.closed = true;
+        // Snapshot — abandoning a forward removes it from the map.
+        for (const abandon of [...state.pending.values()]) abandon();
         try {
           sock.close(() => resolve());
         } catch {
@@ -119,26 +139,56 @@ function forward(
   query: Buffer,
   reply: { address: string; port: number },
   upstream: string,
+  upstreamPort: number,
+  state: ProxyState,
 ): Promise<void> {
   return new Promise((resolve, reject) => {
     const upSock = createSocket('udp4');
     const timer = setTimeout(() => {
-      upSock.close();
+      release();
       reject(new Error('upstream timeout'));
     }, 5000);
-    upSock.once('message', (msg) => {
+
+    /** Drop the upstream socket and its timer. Safe to call more than once. */
+    function release(): void {
       clearTimeout(timer);
-      serverSock.send(msg, reply.port, reply.address, (err) => {
+      state.pending.delete(upSock);
+      try {
         upSock.close();
+      } catch {
+        // Already closed.
+      }
+    }
+
+    // close() calls this for every forward still waiting: the proxy socket is
+    // gone, so nobody is left to answer, and leaving the timer armed would
+    // hold the event loop open for another 5s.
+    state.pending.set(upSock, () => {
+      release();
+      resolve();
+    });
+
+    upSock.once('message', (msg) => {
+      release();
+      // The proxy can be closed while the reply is in flight. send() on a
+      // closed socket throws ERR_SOCKET_DGRAM_NOT_RUNNING synchronously from
+      // inside this handler — past the executor's synchronous phase, so the
+      // promise never sees it and it escapes as an uncaught exception.
+      if (state.closed) {
+        resolve();
+        return;
+      }
+      serverSock.send(msg, reply.port, reply.address, (err) => {
         if (err) reject(err);
         else resolve();
       });
     });
+
     upSock.once('error', (err) => {
-      clearTimeout(timer);
-      upSock.close();
+      release();
       reject(err);
     });
-    upSock.send(query, 53, upstream);
+
+    upSock.send(query, upstreamPort, upstream);
   });
 }