Fix error where header is missing in the response in resolveSampling …

…method
aws · Jan 9, 2025 · eb8ae16 · eb8ae16
1 parent 72d523d
commit eb8ae16
Show file tree

Hide file tree

Showing 2 changed files with 67 additions and 5 deletions.
diff --git a/packages/core/lib/middleware/mw_utils.js b/packages/core/lib/middleware/mw_utils.js
@@ -106,7 +106,7 @@ var utils = {
       }
     }
 
-    if (amznTraceHeader.sampled === '?') {
+    if (amznTraceHeader.sampled === '?' && res.header) {
       res.header[XRAY_HEADER] = 'Root=' + amznTraceHeader.root + ';Sampled=' + (isSampled ? '1' : '0');
     }
 
@@ -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();
 
@@ -220,6 +219,67 @@ 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',
+        host: hostName,
+        method: 'GET',
+        [XRAY_HEADER]: 'Root=' + traceId + '; Parent=' + parentId + '; Sampled=?'
+      };
+      var segment;
+
+      var resWithoutHeader = {
+        req: req,
+        on: (name, callback) => {}
+      };
+      shouldSampleStub.returns(false);
+
+      assert.doesNotThrow(
+        () => {
+          segment = MWUtils.traceRequestResponseCycle(req, resWithoutHeader);
+        }
+      );
+      assert.equal(segment.notTraced, true);
+    });
   });
 
   describe('#samplingWildcardMatch', function() {