Fix error due to missing header in the response in resolveSampling me…

…thod (#706) * Fix error where header is missing in the response in resolveSampling method * update var naming
aws · Jan 10, 2025 · d27bfd2 · d27bfd2
1 parent 72d523d
commit d27bfd2
Show file tree

Hide file tree

Showing 2 changed files with 81 additions and 9 deletions.
diff --git a/packages/core/lib/middleware/mw_utils.js b/packages/core/lib/middleware/mw_utils.js
@@ -106,8 +106,8 @@ var utils = {
       }
     }
 
-    if (amznTraceHeader.sampled === '?') {
-      res.header[XRAY_HEADER] = 'Root=' + amznTraceHeader.root + ';Sampled=' + (isSampled ? '1' : '0');
+    if (amznTraceHeader.sampled === '?' && res.header) {
+      res.header(XRAY_HEADER, 'Root=' + amznTraceHeader.root + ';Sampled=' + (isSampled ? '1' : '0'));
     }
 
     if (!isSampled) {
@@ -172,8 +172,10 @@ var utils = {
     var name = this.resolveName(req.headers.host);
     var segment = new Segment(name, amznTraceHeader.root, amznTraceHeader.parent);
 
-    var responseWithEmbeddedRequest = Object.assign({}, res, { req: req });
-    this.resolveSampling(amznTraceHeader, segment, responseWithEmbeddedRequest);
+    if (!res.req) {
+      res.req = req;
+    }
+    this.resolveSampling(amznTraceHeader, segment, res);
 
     segment.addIncomingRequestData(new IncomingRequestData(req));
 

diff --git a/packages/core/test/unit/middleware/mw_utils.test.js b/packages/core/test/unit/middleware/mw_utils.test.js
@@ -14,6 +14,7 @@ describe('Middleware utils', function() {
   var envVarName = 'envDefaultName';
   var hostName = 'www.myhost.com';
   var traceId = '1-f9194208-2c7ad569f5d6ff149137be86';
+  var parentId = '74051af127d2bcba';
 
   function reloadMWUtils() {
     var path = '../../../lib/logger';
@@ -45,8 +46,6 @@ describe('Middleware utils', function() {
   });
 
   describe('#processHeaders', function() {
-    var parentId = '74051af127d2bcba';
-
     it('should return an empty array on an undefined request', function() {
       var headers = MWUtils.processHeaders();
 
@@ -149,7 +148,7 @@ describe('Middleware utils', function() {
   });
 
   describe('#resolveSampling', function() {
-    var res, sandbox, segment, shouldSampleStub;
+    var res, sandbox, segment, responseHeaders, shouldSampleStub;
 
     beforeEach(function() {
       sandbox = sinon.createSandbox();
@@ -158,18 +157,22 @@ describe('Middleware utils', function() {
       shouldSampleStub = sandbox.stub(MWUtils.sampler, 'shouldSample').returns(true);
 
       segment = {};
+      responseHeaders = {};
       res = {
         req: {
           headers: { host: 'moop.hello.com' },
           url: '/evergreen',
           method: 'GET',
         },
-        header: {}
+        header: (headerKey, headerValue) => {
+          responseHeaders[headerKey] = headerValue;
+        }
       };
     });
 
     afterEach(function() {
       sandbox.restore();
+      responseHeaders = {};
     });
 
     it('should not mark segment as not traced if the sampled header is set to "1"', function() {
@@ -209,7 +212,7 @@ describe('Middleware utils', function() {
       MWUtils.resolveSampling(headers, segment, res);
 
       var expected = new RegExp('^Root=' + traceId + ';Sampled=1$');
-      assert.match(res.header[XRAY_HEADER], expected);
+      assert.match(responseHeaders[XRAY_HEADER], expected);
     });
 
     it('should mark segment as not traced if the sampling rules check returns false', function() {
@@ -220,6 +223,73 @@ describe('Middleware utils', function() {
 
       assert.equal(segment.notTraced, true);
     });
+
+    it('should not throw error when res.header is undefined and Sampled=?', function() {
+      var resWithoutHeader = {
+        req: {
+          headers: {},
+          url: '/api/move/up',
+          method: 'GET',
+        }
+      };
+      shouldSampleStub.returns(false);
+      var headers = { root: traceId, sampled: '?' };
+
+      assert.doesNotThrow(
+        () => {
+          MWUtils.resolveSampling(headers, segment, resWithoutHeader);
+        }
+      );
+
+      assert.equal(segment.notTraced, true);
+    });
+  });
+
+
+  describe('#traceRequestResponseCycle', function() {
+    var sandbox, shouldSampleStub;
+
+    beforeEach(function() {
+      sandbox = sinon.createSandbox();
+      MWUtils.sampler = localSampler;
+      MWUtils.setDefaultName(defaultName);
+
+      shouldSampleStub = sandbox.stub(MWUtils.sampler, 'shouldSample').returns(true);
+    });
+
+    afterEach(function() {
+      sandbox.restore();
+    });
+
+    it('should not throw error when Sampled=?', function() {
+      var req = {
+        headers: {
+          url: '/api/move/up',
+          [XRAY_HEADER]: 'Root=' + traceId + ';Parent=' + parentId + ';Sampled=?'
+        },
+        host: hostName,
+        method: 'GET',
+      };
+      var segment;
+      var responseHeaders = {};
+
+      var res = {
+        req: req,
+        on: (name, callback) => {},
+        header: (headerKey, headerValue) => {
+          responseHeaders[headerKey] = headerValue;
+        }
+      };
+      shouldSampleStub.returns(false);
+
+      assert.doesNotThrow(
+        () => {
+          segment = MWUtils.traceRequestResponseCycle(req, res);
+        }
+      );
+      assert.equal(segment.notTraced, true);
+      assert.equal(responseHeaders[XRAY_HEADER], 'Root=' + traceId + ';Sampled=0');
+    });
   });
 
   describe('#samplingWildcardMatch', function() {