From 50421a5ad9ee8cb49e629dc7a9a8caddee959401 Mon Sep 17 00:00:00 2001
From: Kelvin Jin <kelvinjin@google.com>
Date: Wed, 3 Apr 2019 14:13:11 -0700
Subject: [PATCH] fix: don't let trace context injection throw (#989)

---
 src/plugins/plugin-http.ts      | 26 +++++++++++++++++++++++---
 test/plugins/test-trace-http.ts | 29 +++++++++++++++--------------
 2 files changed, 38 insertions(+), 17 deletions(-)

diff --git a/src/plugins/plugin-http.ts b/src/plugins/plugin-http.ts
index 8ad2ee3fb..25da9afe5 100644
--- a/src/plugins/plugin-http.ts
+++ b/src/plugins/plugin-http.ts
@@ -28,6 +28,9 @@ type HttpModule = typeof httpModule;
 type HttpsModule = typeof httpsModule;
 type RequestFunction = typeof request;
 
+const ERR_HTTP_HEADERS_SENT = 'ERR_HTTP_HEADERS_SENT';
+const ERR_HTTP_HEADERS_SENT_MSG = 'Can\'t set headers after they are sent.';
+
 // tslint:disable:no-any
 const isString = is.string as (value: any) => value is string;
 // url.URL is used for type checking, but doesn't exist in Node <7.
@@ -44,12 +47,16 @@ function getSpanName(options: ClientRequestArgs|url.URL) {
 
 /**
  * Returns whether the Expect header is on the given options object.
+ * Assumes only that the header key is either capitalized, lowercase, or
+ * all-caps for simplicity purposes.
  * @param options Options for http.request.
  */
 function hasExpectHeader(options: ClientRequestArgs|url.URL): boolean {
   return !!(
       (options as ClientRequestArgs).headers &&
-      (options as ClientRequestArgs).headers!.Expect);
+      ((options as ClientRequestArgs).headers!.Expect ||
+       (options as ClientRequestArgs).headers!.expect ||
+       (options as ClientRequestArgs).headers!.EXPECT));
 }
 
 function extractUrl(
@@ -126,6 +133,8 @@ function makeRequestTrace(
     // inject the trace context header instead of using ClientRequest#setHeader.
     // (We don't do this generally because cloning the options object is an
     // expensive operation.)
+    // See https://github.com/googleapis/cloud-trace-nodejs/pull/766 for a full
+    // explanation.
     let traceHeaderPreinjected = false;
     if (hasExpectHeader(options)) {
       traceHeaderPreinjected = true;
@@ -179,8 +188,19 @@ function makeRequestTrace(
     // Inject the trace context header, but only if it wasn't already injected
     // earlier.
     if (!traceHeaderPreinjected) {
-      req.setHeader(
-          api.constants.TRACE_CONTEXT_HEADER_NAME, span.getTraceContext());
+      try {
+        req.setHeader(
+            api.constants.TRACE_CONTEXT_HEADER_NAME, span.getTraceContext());
+      } catch (e) {
+        if (e.code === ERR_HTTP_HEADERS_SENT ||
+            e.message === ERR_HTTP_HEADERS_SENT_MSG) {
+          // Swallow the error.
+          // This would happen in the pathological case where the Expect header
+          // exists but is not detected by hasExpectHeader.
+        } else {
+          throw e;
+        }
+      }
     }
     return req;
   };
diff --git a/test/plugins/test-trace-http.ts b/test/plugins/test-trace-http.ts
index 6982e09d3..df0554140 100644
--- a/test/plugins/test-trace-http.ts
+++ b/test/plugins/test-trace-http.ts
@@ -191,20 +191,6 @@ for (const nodule of Object.keys(servers) as Array<keyof typeof servers>) {
             await waitForResponse.done;
           }
         },
-        {
-          description: 'calling http.get with Expect header',
-          fn: async () => {
-            const waitForResponse = new WaitForResponse();
-            const req = http.get(
-                {
-                  port,
-                  rejectUnauthorized: false,
-                  headers: {Expect: '100-continue'}
-                },
-                waitForResponse.handleResponse);
-            await waitForResponse.done;
-          }
-        },
         {
           description: 'calling http.get, but timing out and emitting an error',
           fn: async () => {
@@ -219,6 +205,21 @@ for (const nodule of Object.keys(servers) as Array<keyof typeof servers>) {
             await waitForResponse.done;
           }
         },
+        ...['Expect', 'expect', 'EXPECT', 'eXpEcT'].map(
+            key => ({
+              description: `calling http.get with ${key} header`,
+              fn: async () => {
+                const waitForResponse = new WaitForResponse();
+                http.get(
+                    {
+                      port,
+                      rejectUnauthorized: false,
+                      headers: {[key]: '100-continue'}
+                    },
+                    waitForResponse.handleResponse);
+                await waitForResponse.done;
+              }
+            }))
       ];
 
       beforeEach(async () => {