From dd12415098a5cb4d9846683cae9508e857fa8835 Mon Sep 17 00:00:00 2001 From: Mahad Janjua Date: Thu, 6 Jun 2024 17:19:07 -0700 Subject: [PATCH] Allow lambda passthrough trace header --- packages/core/lib/patchers/aws3_p.ts | 4 +- packages/core/lib/patchers/aws_p.js | 3 + packages/core/lib/patchers/http_p.js | 6 +- .../lib/segments/attributes/subsegment.js | 1 + packages/core/lib/segments/segment.js | 1 + packages/core/lib/utils.js | 4 + .../core/test/unit/patchers/aws_p.test.js | 88 +++++++++++++++++++ 7 files changed, 104 insertions(+), 3 deletions(-) diff --git a/packages/core/lib/patchers/aws3_p.ts b/packages/core/lib/patchers/aws3_p.ts index 782317e4..c72d17a6 100644 --- a/packages/core/lib/patchers/aws3_p.ts +++ b/packages/core/lib/patchers/aws3_p.ts @@ -140,7 +140,9 @@ const getXRayMiddleware = (config: RegionResolvedConfig, manualSegment?: Segment } } - args.request.headers['X-Amzn-Trace-Id'] = traceHeader; + if (!segment.noOp) { + args.request.headers['X-Amzn-Trace-Id'] = traceHeader; + } let res; try { diff --git a/packages/core/lib/patchers/aws_p.js b/packages/core/lib/patchers/aws_p.js index 776d33ef..52c1f021 100644 --- a/packages/core/lib/patchers/aws_p.js +++ b/packages/core/lib/patchers/aws_p.js @@ -88,6 +88,9 @@ function captureAWSRequest(req) { const data = parent.segment ? parent.segment.additionalTraceData : parent.additionalTraceData; var buildListener = function(req) { + if (parent.noOp) { + return; + } let traceHeader = 'Root=' + traceId + ';Parent=' + subsegment.id + ';Sampled=' + (subsegment.notTraced ? '0' : '1'); if (data != null) { diff --git a/packages/core/lib/patchers/http_p.js b/packages/core/lib/patchers/http_p.js index fcf2f16f..d84563e1 100644 --- a/packages/core/lib/patchers/http_p.js +++ b/packages/core/lib/patchers/http_p.js @@ -130,8 +130,10 @@ function enableCapture(module, downstreamXRayEnabled, subsegmentCallback) { options.headers = {}; } - options.headers['X-Amzn-Trace-Id'] = 'Root=' + root.trace_id + ';Parent=' + subsegment.id + - ';Sampled=' + (subsegment.notTraced ? '0' : '1'); + if (!parent.noOp) { + options.headers['X-Amzn-Trace-Id'] = 'Root=' + root.trace_id + ';Parent=' + subsegment.id + + ';Sampled=' + (subsegment.notTraced ? '0' : '1'); + } const errorCapturer = function errorCapturer(e) { if (subsegmentCallback) { diff --git a/packages/core/lib/segments/attributes/subsegment.js b/packages/core/lib/segments/attributes/subsegment.js index b46a52f6..0ead6272 100644 --- a/packages/core/lib/segments/attributes/subsegment.js +++ b/packages/core/lib/segments/attributes/subsegment.js @@ -74,6 +74,7 @@ Subsegment.prototype.addSubsegment = function(subsegment) { subsegment.parent = this; subsegment.notTraced = subsegment.parent.notTraced; + subsegment.noOp = subsegment.parent.noOp; if (subsegment.end_time === undefined) { this.incrementCounter(subsegment.counter); diff --git a/packages/core/lib/segments/segment.js b/packages/core/lib/segments/segment.js index b8f24533..52bf3141 100644 --- a/packages/core/lib/segments/segment.js +++ b/packages/core/lib/segments/segment.js @@ -248,6 +248,7 @@ Segment.prototype.addSubsegment = function addSubsegment(subsegment) { subsegment.parent = this; subsegment.notTraced = subsegment.parent.notTraced; + subsegment.noOp = subsegment.parent.noOp; this.subsegments.push(subsegment); if (!subsegment.end_time) { diff --git a/packages/core/lib/utils.js b/packages/core/lib/utils.js index 93efcd3b..e047b71a 100644 --- a/packages/core/lib/utils.js +++ b/packages/core/lib/utils.js @@ -168,6 +168,10 @@ var utils = { if (!traceData) { traceData = {}; logger.getLogger().error('_X_AMZN_TRACE_ID is empty or has an invalid format'); + } else if (traceData.root && !traceData.parent && !traceData.sampled) { + // Lambda PassThrough only has root, treat as valid in this case and mark the segment + segment.noOp = true; + valid = true; } else if (!traceData.root || !traceData.parent || !traceData.sampled) { logger.getLogger().error('_X_AMZN_TRACE_ID is missing required information'); } else { diff --git a/packages/core/test/unit/patchers/aws_p.test.js b/packages/core/test/unit/patchers/aws_p.test.js index fdd48d93..c4220818 100644 --- a/packages/core/test/unit/patchers/aws_p.test.js +++ b/packages/core/test/unit/patchers/aws_p.test.js @@ -348,4 +348,92 @@ describe('AWS patcher', function() { }); }); + + + describe('#captureAWSRequest-PassThrough-Header', function() { + var awsClient, awsRequest, MyEmitter, sandbox, segment, stubResolve, addNewSubsegmentStub, sub, addNewServiceSubsegmentStub, service; + + before(function() { + MyEmitter = function() { + EventEmitter.call(this); + }; + + awsClient = { + customizeRequests: function customizeRequests(captureAWSRequest) { + this.call = captureAWSRequest; + }, + throttledError: function throttledError() {} + }; + awsClient = awsPatcher.captureAWSClient(awsClient); + + util.inherits(MyEmitter, EventEmitter); + }); + + beforeEach(function() { + sandbox = sinon.createSandbox(); + + awsRequest = { + httpRequest: { + method: 'GET', + url: '/', + connection: { + remoteAddress: 'localhost' + }, + headers: {} + }, + response: {} + }; + + awsRequest.on = function(event, fcn) { + if (event === 'complete') { + this.emitter.on(event, fcn.bind(this, this.response)); + } else { + this.emitter.on(event, fcn.bind(this, this)); + } + return this; + }; + + awsRequest.emitter = new MyEmitter(); + + segment = new Segment('testSegment', traceId); + segment.noOp = true; // enforce passthrough behaviour + segment.additionalTraceData = {'Foo': 'bar'}; + sub = segment.addNewSubsegmentWithoutSampling('subseg'); + service = sub.addNewSubsegmentWithoutSampling('service'); + + stubResolve = sandbox.stub(contextUtils, 'resolveSegment').returns(sub); + addNewSubsegmentStub = sandbox.stub(segment, 'addNewSubsegmentWithoutSampling').returns(sub); + addNewServiceSubsegmentStub = sandbox.stub(sub, 'addNewSubsegmentWithoutSampling').returns(service); + }); + + afterEach(function() { + sandbox.restore(); + }); + + it('should log an info statement and exit if parent is not found on the context or on the call params', function(done) { + stubResolve.returns(); + var logStub = sandbox.stub(logger, 'info'); + + awsClient.call(awsRequest); + + setTimeout(function() { + logStub.should.have.been.calledOnce; + done(); + }, 50); + }); + + it('should not inject the tracing headers if passthrough mode', function(done) { + sandbox.stub(contextUtils, 'isAutomaticMode').returns(true); + + awsClient.call(awsRequest); + + awsRequest.emitter.emit('build'); + + setTimeout(function() { + assert.equal(awsRequest.httpRequest.headers['X-Amzn-Trace-Id'], undefined); + done(); + }, 50); + }); + + }); });