fix(instrumentation-undici): fix several header handling handling bugs

The handling for User-Agent had a bunch of bugs, most notably with the handling of multiple-valued headers, which are rare but legal.

Code similar to the added test "another header with multiple values" caused us errors in production from the user-agent code, and is where I started. Reading the code, writing tests, and improving the types revealed several more bugs in the same code as well as the span-to-attributes handling; the added tests "multiple user-agents", "another header with value user-agent" also fail before this PR (the others are just to actually cover the ordinary cases.

Similarly, I also fixed an incorrect case in undici v5 where it would treat a `user-agent-bogus` header as a user agent, but I couldn't write a test case since the tests run with the newer version.

Author: benjaminjkraft

plugins/node/instrumentation-undici/src/types.ts

@@ -22,9 +22,10 @@ export interface UndiciRequest {
   path: string;
   /**
    * Serialized string of headers in the form `name: value\r\n` for v5
-   * Array of strings v6
+   * Array of strings `[key1, value1, key2, value2]`, where values are
+   * `string | string[]` for v6
    */
-  headers: string | string[];
+  headers: string | (string | string[])[];
   /**
    * Helper method to add headers (from v6)
    */

plugins/node/instrumentation-undici/src/undici.ts

@@ -165,6 +165,49 @@ export class UndiciInstrumentation extends InstrumentationBase<UndiciInstrumenta
     });
   }
 
+  /**
+   * For each header in the request, call the callback. Skips likely-invalid
+   * headers. Multi-valued headers are passed through. The loop exits early if
+   * the callback returns true.
+   */
+  private forEachRequestHeader(
+    request: UndiciRequest,
+    callback: (key: string, value: string | string[]) => boolean | undefined
+  ): void {
+    if (Array.isArray(request.headers)) {
+      // headers are an array [k1, v2, k2, v2] (undici v6+)
+      for (let i = 0; i < request.headers.length; i += 2) {
+        const key = request.headers[i];
+        if (typeof key !== 'string') {
+          // Shouldn't happen, but the types don't know that, and let's be safe
+          continue;
+        }
+        if (callback(key, request.headers[i + 1])) {
+          break;
+        }
+      }
+    } else if (typeof request.headers === 'string') {
+      // headers are a raw string (undici v5)
+      const headers = request.headers.split('\r\n');
+      for (const line of headers) {
+        if (!line) {
+          continue;
+        }
+        const colonIndex = line.indexOf(':');
+        if (colonIndex === -1) {
+          // Invalid header? Probably this can't happen, but again let's be safe.
+          continue;
+        }
+        const key = line.substring(0, colonIndex);
+        const value = line.substring(0, colonIndex + 1);
+
+        if (callback(key, value)) {
+          break;
+        }
+      }
+    }
+  }
+
   // This is the 1st message we receive for each request (fired after request creation). Here we will
   // create the span and populate some atttributes, then link the span to the request for further
   // span processing
@@ -218,26 +261,16 @@ export class UndiciInstrumentation extends InstrumentationBase<UndiciInstrumenta
     }
 
     // Get user agent from headers
-    let userAgent;
-    if (Array.isArray(request.headers)) {
-      const idx = request.headers.findIndex(
-        h => h.toLowerCase() === 'user-agent'
-      );
-      if (idx >= 0) {
-        userAgent = request.headers[idx + 1];
+    this.forEachRequestHeader(request, (key, value) => {
+      if (key.toLowerCase() === 'user-agent') {
+        // user-agent should only appear once per the spec, but the library doesn't
+        // prevent passing it multiple times, so we handle that to be safe.
+        const userAgent = Array.isArray(value) ? value[0] : value;
+        attributes[SemanticAttributes.USER_AGENT_ORIGINAL] = userAgent;
+        return true; // no need to keep iterating
       }
-    } else if (typeof request.headers === 'string') {
-      const headers = request.headers.split('\r\n');
-      const uaHeader = headers.find(h =>
-        h.toLowerCase().startsWith('user-agent')
-      );
-      userAgent =
-        uaHeader && uaHeader.substring(uaHeader.indexOf(':') + 1).trim();
-    }
-
-    if (userAgent) {
-      attributes[SemanticAttributes.USER_AGENT_ORIGINAL] = userAgent;
-    }
+      return false;
+    });
 
     // Get attributes from the hook if present
     const hookAttributes = safeExecuteInTheMiddle(
@@ -330,25 +363,14 @@ export class UndiciInstrumentation extends InstrumentationBase<UndiciInstrumenta
         config.headersToSpanAttributes.requestHeaders.map(n => n.toLowerCase())
       );
 
-      // headers could be in form
-      // ['name: value', ...] for v5
-      // ['name', 'value', ...] for v6
-      const rawHeaders = Array.isArray(request.headers)
-        ? request.headers
-        : request.headers.split('\r\n');
-      rawHeaders.forEach((h, idx) => {
-        const sepIndex = h.indexOf(':');
-        const hasSeparator = sepIndex !== -1;
-        const name = (
-          hasSeparator ? h.substring(0, sepIndex) : h
-        ).toLowerCase();
-        const value = hasSeparator
-          ? h.substring(sepIndex + 1)
-          : rawHeaders[idx + 1];
-
+      this.forEachRequestHeader(request, (key, value) => {
+        const name = key.toLowerCase();
         if (headersToAttribs.has(name)) {
-          spanAttributes[`http.request.header.${name}`] = value.trim();
+          spanAttributes[`http.request.header.${name}`] = value
+            .toString()
+            .trim();
         }
+        return false; // keep iterating always, there may be more
       });
     }
 

plugins/node/instrumentation-undici/test/undici.test.ts

@@ -795,43 +795,87 @@ describe('UndiciInstrumentation `undici` tests', function () {
       });
     });
 
-    it('should not report an user-agent if it was not defined', async function () {
-      let spans = memoryExporter.getFinishedSpans();
-      assert.strictEqual(spans.length, 0);
+    const userAgentRequests: Array<{
+      name: string;
+      headers: Record<string, string | string[] | undefined>;
+      expectedUserAgent: string | undefined;
+    }> = [
+      {
+        name: 'no user-agent',
+        headers: { 'foo-client': 'bar' },
+        expectedUserAgent: undefined,
+      },
+      {
+        name: 'a user-agent',
+        headers: { 'foo-client': 'bar', 'user-agent': 'custom' },
+        expectedUserAgent: 'custom',
+      },
+      {
+        name: 'explicitly-undefined user-agent',
+        headers: { 'foo-client': 'bar', 'user-agent': undefined },
+        expectedUserAgent: undefined,
+      },
+      // contra the spec, but we shouldn't crash
+      {
+        name: 'multiple user-agents',
+        headers: {
+          'foo-client': 'bar',
+          'user-agent': ['agent', 'other-agent'],
+        },
+        expectedUserAgent: 'agent',
+      },
+      {
+        name: 'another header with value user-agent',
+        headers: { 'foo-client': 'user-agent', 'user-agent': 'custom' },
+        expectedUserAgent: 'custom',
+      },
+      {
+        name: 'another header with multiple values',
+        headers: { 'foo-client': ['one', 'two'], 'user-agent': 'custom' },
+        expectedUserAgent: 'custom',
+      },
+      {
+        name: 'another header with explicitly-undefined value',
+        headers: { 'foo-client': undefined, 'user-agent': 'custom' },
+        expectedUserAgent: 'custom',
+      },
+    ];
 
-      // Do some requests
-      const headers = {
-        'foo-client': 'bar',
-      };
+    for (const testCase of userAgentRequests) {
+      it(`should report the correct user-agent when the request has ${testCase.name}`, async function () {
+        let spans = memoryExporter.getFinishedSpans();
+        assert.strictEqual(spans.length, 0);
 
-      const queryRequestUrl = `${protocol}://${hostname}:${mockServer.port}/?query=test`;
-      const queryResponse = await undici.request(queryRequestUrl, { headers });
-      await consumeResponseBody(queryResponse.body);
+        const queryRequestUrl = `${protocol}://${hostname}:${mockServer.port}/?query=test`;
+        const queryResponse = await undici.request(queryRequestUrl, {
+          headers: testCase.headers,
+        });
+        await consumeResponseBody(queryResponse.body);
 
-      assert.ok(
-        queryResponse.headers['propagation-error'] == null,
-        'propagation is set for instrumented requests'
-      );
+        assert.ok(
+          queryResponse.headers['propagation-error'] == null,
+          'propagation is set for instrumented requests'
+        );
 
-      spans = memoryExporter.getFinishedSpans();
-      const span = spans[0];
-      assert.ok(span, 'a span is present');
-      assert.strictEqual(spans.length, 1);
-      assertSpan(span, {
-        hostname: 'localhost',
-        httpStatusCode: queryResponse.statusCode,
-        httpMethod: 'GET',
-        path: '/',
-        query: '?query=test',
-        reqHeaders: headers,
-        resHeaders: queryResponse.headers,
+        spans = memoryExporter.getFinishedSpans();
+        const span = spans[0];
+        assert.ok(span, 'a span is present');
+        assert.strictEqual(spans.length, 1);
+        assertSpan(span, {
+          hostname: 'localhost',
+          httpStatusCode: queryResponse.statusCode,
+          httpMethod: 'GET',
+          path: '/',
+          query: '?query=test',
+          reqHeaders: testCase.headers,
+          resHeaders: queryResponse.headers,
+        });
+        assert.strictEqual(
+          span.attributes['user_agent.original'],
+          testCase.expectedUserAgent
+        );
       });
-      assert.strictEqual(
-        span.attributes['user_agent.original'],
-        undefined,
-        'user-agent is undefined'
-      );
-    });
+    }
 
     it('should create valid span if request.path is a full URL', async function () {
       let spans = memoryExporter.getFinishedSpans();

plugins/node/instrumentation-undici/test/utils/assertSpan.ts

@@ -171,7 +171,7 @@ export const assertSpan = (
     if (userAgent) {
       assert.strictEqual(
         span.attributes[SemanticAttributes.USER_AGENT_ORIGINAL],
-        userAgent
+        Array.isArray(userAgent) ? userAgent[0] : userAgent
       );
     }
   }