From 1f266cec60c3ee550b956a8e9b2dc1ad68431bf0 Mon Sep 17 00:00:00 2001
From: fengmk2 <suqian.yf@antgroup.com>
Date: Sat, 1 Jun 2024 10:33:46 +0800
Subject: [PATCH] fix: adpater http/2 agent on diagnosticsChannel

https://github.com/node-modules/urllib/issues/510
---
 src/diagnosticsChannel.ts       | 17 ++++++++++++-----
 test/options.dispatcher.test.ts | 17 ++++++++++++++++-
 2 files changed, 28 insertions(+), 6 deletions(-)

diff --git a/src/diagnosticsChannel.ts b/src/diagnosticsChannel.ts
index 51fe618f..97780e00 100644
--- a/src/diagnosticsChannel.ts
+++ b/src/diagnosticsChannel.ts
@@ -184,10 +184,15 @@ export function initDiagnosticsChannel() {
 
     // get socket from opaque
     const socket = opaque[symbols.kRequestSocket];
-    socket[symbols.kHandledResponses]++;
-    debug('[%s] Request#%d get %s response headers on Socket#%d (handled %d responses, sock: %o)',
-      name, opaque[symbols.kRequestId], response.statusCode, socket[symbols.kSocketId], socket[symbols.kHandledResponses],
-      formatSocket(socket));
+    if (socket) {
+      socket[symbols.kHandledResponses]++;
+      debug('[%s] Request#%d get %s response headers on Socket#%d (handled %d responses, sock: %o)',
+        name, opaque[symbols.kRequestId], response.statusCode, socket[symbols.kSocketId], socket[symbols.kHandledResponses],
+        formatSocket(socket));
+    } else {
+      debug('[%s] Request#%d get %s response headers on Unknown Socket',
+        name, opaque[symbols.kRequestId], response.statusCode);
+    }
 
     if (!opaque[symbols.kEnableRequestTiming]) return;
     opaque[symbols.kRequestTiming].waiting = performanceTime(opaque[symbols.kRequestStartTime]);
@@ -197,7 +202,9 @@ export function initDiagnosticsChannel() {
   subscribe('undici:request:trailers', (message, name) => {
     const { request } = message as DiagnosticsChannel.RequestTrailersMessage;
     const opaque = getRequestOpaque(request, kHandler);
-    if (!opaque || !opaque[symbols.kRequestId]) return;
+    if (!opaque || !opaque[symbols.kRequestId]) {
+      return;
+    }
 
     debug('[%s] Request#%d get response body and trailers', name, opaque[symbols.kRequestId]);
 
diff --git a/test/options.dispatcher.test.ts b/test/options.dispatcher.test.ts
index 0571f870..bde1e00e 100644
--- a/test/options.dispatcher.test.ts
+++ b/test/options.dispatcher.test.ts
@@ -1,7 +1,7 @@
 import { strict as assert } from 'node:assert';
 import { describe, it, beforeAll, afterAll } from 'vitest';
 import setup from 'proxy';
-import { request, ProxyAgent, getGlobalDispatcher, setGlobalDispatcher } from '../src';
+import { request, ProxyAgent, getGlobalDispatcher, setGlobalDispatcher, Agent } from '../src';
 import { startServer } from './fixtures/server';
 
 describe('options.dispatcher.test.ts', () => {
@@ -62,4 +62,19 @@ describe('options.dispatcher.test.ts', () => {
     assert.equal(response.data, '<h1>hello</h1>');
     setGlobalDispatcher(agent);
   });
+
+  it('should work with http/2 dispatcher', async () => {
+    // https://github.com/nodejs/undici/issues/2750#issuecomment-1941009554
+    const agent = new Agent({
+      allowH2: true,
+    });
+    assert(agent);
+    const response = await request('https://registry.npmmirror.com', {
+      dataType: 'json',
+      timing: true,
+      dispatcher: agent,
+    });
+    assert.equal(response.status, 200);
+    assert.equal(response.headers['content-type'], 'application/json; charset=utf-8');
+  });
 });