From 7326aa76b0fe8582e62a8bc145300158c00f7e62 Mon Sep 17 00:00:00 2001
From: fengmk2 <fengmk2@gmail.com>
Date: Tue, 3 Dec 2024 21:01:29 +0800
Subject: [PATCH] fix: publish undici:client:sendHeaders message on H2

Signed-off-by: fengmk2 <fengmk2@gmail.com>
---
 lib/dispatcher/client-h2.js                  |   9 ++
 test/node-test/diagnostics-channel/get-h2.js | 101 +++++++++++++++++++
 2 files changed, 110 insertions(+)
 create mode 100644 test/node-test/diagnostics-channel/get-h2.js

diff --git a/lib/dispatcher/client-h2.js b/lib/dispatcher/client-h2.js
index 1276d6b2e05..d34c7cab989 100644
--- a/lib/dispatcher/client-h2.js
+++ b/lib/dispatcher/client-h2.js
@@ -30,6 +30,7 @@ const {
   kClosed,
   kBodyTimeout
 } = require('../core/symbols.js')
+const { channels } = require('../core/diagnostics.js')
 
 const kOpenStreams = Symbol('open streams')
 
@@ -448,6 +449,14 @@ function writeH2 (client, request) {
 
   session.ref()
 
+  if (channels.sendHeaders.hasSubscribers) {
+    let header = ''
+    for (const key in headers) {
+      header += `${key}: ${headers[key]}\r\n`
+    }
+    channels.sendHeaders.publish({ request, headers: header, socket: session[kSocket] })
+  }
+
   // TODO(metcoder95): add support for sending trailers
   const shouldEndStream = method === 'GET' || method === 'HEAD' || body === null
   if (expectContinue) {
diff --git a/test/node-test/diagnostics-channel/get-h2.js b/test/node-test/diagnostics-channel/get-h2.js
new file mode 100644
index 00000000000..90d16287087
--- /dev/null
+++ b/test/node-test/diagnostics-channel/get-h2.js
@@ -0,0 +1,101 @@
+'use strict'
+
+const { tspl } = require('@matteo.collina/tspl')
+const { test, after } = require('node:test')
+const { createSecureServer } = require('node:http2')
+const diagnosticsChannel = require('node:diagnostics_channel')
+const { once } = require('node:events')
+const pem = require('https-pem')
+const { Client } = require('../../..')
+
+test('Diagnostics channel - get support H2', async t => {
+  const server = createSecureServer(pem)
+
+  server.on('stream', (stream, headers, _flags, rawHeaders) => {
+    t.strictEqual(headers['x-my-header'], 'foo')
+    t.strictEqual(headers[':method'], 'GET')
+    stream.respond({
+      'content-type': 'text/plain; charset=utf-8',
+      'x-custom-h2': 'hello',
+      ':status': 200
+    })
+    stream.end('hello h2!')
+  })
+
+  server.listen(0)
+  await once(server, 'listening')
+
+  diagnosticsChannel.channel('undici:request:create').subscribe(({ request }) => {
+    t.strictEqual(request.origin, `https://localhost:${server.address().port}`)
+    t.strictEqual(request.completed, false)
+    t.strictEqual(request.method, 'GET')
+    t.strictEqual(request.path, '/')
+  })
+
+  let _socket
+  diagnosticsChannel.channel('undici:client:connected').subscribe(({ socket }) => {
+    _socket = socket
+  })
+
+  diagnosticsChannel.channel('undici:client:sendHeaders').subscribe(({ headers, socket }) => {
+    t.strictEqual(_socket, socket)
+    const expectedHeaders = [
+      'x-my-header: foo',
+      `:authority: localhost:${server.address().port}`,
+      ':method: GET',
+      ':path: /',
+      ':scheme: https'
+    ]
+    t.strictEqual(headers, expectedHeaders.join('\r\n') + '\r\n')
+  })
+
+  const client = new Client(`https://localhost:${server.address().port}`, {
+    connect: {
+      rejectUnauthorized: false
+    },
+    allowH2: true
+  })
+
+  t = tspl(t, { plan: 24 })
+  after(() => server.close())
+  after(() => client.close())
+
+  let body = []
+  let response = await client.request({
+    path: '/',
+    method: 'GET',
+    headers: {
+      'x-my-header': 'foo'
+    }
+  })
+
+  response.body.on('data', chunk => {
+    body.push(chunk)
+  })
+
+  await once(response.body, 'end')
+  t.strictEqual(response.statusCode, 200)
+  t.strictEqual(response.headers['content-type'], 'text/plain; charset=utf-8')
+  t.strictEqual(response.headers['x-custom-h2'], 'hello')
+  t.strictEqual(Buffer.concat(body).toString('utf8'), 'hello h2!')
+
+  // request again
+  body = []
+  response = await client.request({
+    path: '/',
+    method: 'GET',
+    headers: {
+      'x-my-header': 'foo'
+    }
+  })
+
+  response.body.on('data', chunk => {
+    body.push(chunk)
+  })
+
+  await once(response.body, 'end')
+  t.strictEqual(response.statusCode, 200)
+  t.strictEqual(response.headers['content-type'], 'text/plain; charset=utf-8')
+  t.strictEqual(response.headers['x-custom-h2'], 'hello')
+  t.strictEqual(Buffer.concat(body).toString('utf8'), 'hello h2!')
+})