fix(lambda): handle missing context for some lambda functions (#7445)

* fix(lambda): handle missing context for Lambda Authorizers

Lambda Authorizers and some other Lambda handler types do not receive
a context object - they only receive an event parameter. Previously,
the extractContext() function would throw "Could not extract context"
when it couldn't find a context object with getRemainingTimeInMillis.

This change makes extractContext() return undefined instead of throwing,
and updates the datadog() wrapper to skip timeout checking when context
is not available. This allows Lambda Authorizers to be instrumented
without crashing.

Fixes: DataDog/datadog-lambda-js#721

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

* Undefined cleanup.

* Add getRemainingTimeInMillis check back.

* Update packages/dd-trace/src/lambda/handler.js

Co-authored-by: Ruben Bridgewater <ruben@bridgewater.de>

* Move extractContext to separate file and add more tesets.

* Linting.

---------

Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>
Co-authored-by: Rey Abolofia <rey.abolofia@datadoghq.com>
Co-authored-by: Rey Abolofia <purple4reina@gmail.com>
Co-authored-by: Ruben Bridgewater <ruben@bridgewater.de>

Author: 4 people

packages/dd-trace/src/lambda/context.js

@@ -0,0 +1,27 @@
+'use strict'
+
+const log = require('../log')
+
+/**
+ * Extracts the context from the given Lambda handler arguments.
+ *
+ * It is possible for users to define a lambda function without specifying a
+ * context arg. In these cases, this function returns null instead of throwing
+ * an error.
+ *
+ * @param {unknown[]} args any amount of arguments
+ * @returns {object | null}
+ */
+exports.extractContext = function extractContext (args) {
+  let context = null
+  for (let i = 0; i < args.length && i < 3; i++) {
+    if (args[i] && typeof args[i].getRemainingTimeInMillis === 'function') {
+      context = args[i]
+      break
+    }
+  }
+  if (!context) {
+    log.debug('Unable to extract context object from Lambda handler arguments')
+  }
+  return context
+}

packages/dd-trace/src/lambda/handler.js

@@ -5,6 +5,7 @@ const { channel } = require('../../../datadog-instrumentations/src/helpers/instr
 const { ERROR_MESSAGE, ERROR_TYPE } = require('../constants')
 const { getValueFromEnvSources } = require('../config/helper')
 const { ImpendingTimeout } = require('./runtime/errors')
+const { extractContext } = require('./context')
 
 const globalTracer = global._ddtrace
 const tracer = globalTracer._tracer
@@ -60,23 +61,6 @@ function crashFlush () {
   }
 }
 
-/**
- * Extracts the context from the given Lambda handler arguments.
- *
- * @param {unknown[]} args any amount of arguments
- * @returns the context, if extraction was succesful.
- */
-function extractContext (args) {
-  let context = args.length > 1 ? args[1] : undefined
-  if (context === undefined || context.getRemainingTimeInMillis === undefined) {
-    context = args.length > 2 ? args[2] : undefined
-    if (context === undefined || context.getRemainingTimeInMillis === undefined) {
-      throw new Error('Could not extract context')
-    }
-  }
-  return context
-}
-
 /**
  * Patches your AWS Lambda handler function to add some tracing support.
  *
@@ -86,7 +70,10 @@ exports.datadog = function datadog (lambdaHandler) {
   return (...args) => {
     const context = extractContext(args)
 
-    checkTimeout(context)
+    if (context) {
+      checkTimeout(context)
+    }
+
     const result = lambdaHandler.apply(this, args)
     if (result && typeof result.then === 'function') {
       return result.then((res) => {

packages/dd-trace/test/lambda/context.spec.js

@@ -0,0 +1,41 @@
+'use strict'
+
+const assert = require('node:assert/strict')
+const { describe, it } = require('mocha')
+const { extractContext } = require('../../src/lambda/context')
+
+describe('context', () => {
+  describe('extractContext', () => {
+    const assertExtractContext = (args, doesExtract) => {
+      it(`properly extracts context object from args length ${args.length}`, () => {
+        const ctx = extractContext(args)
+        if (doesExtract) {
+          assert.strictEqual(typeof ctx.getRemainingTimeInMillis, 'function')
+          assert.strictEqual(ctx.getRemainingTimeInMillis(), 100)
+        } else {
+          assert.strictEqual(ctx, null)
+        }
+      })
+    }
+
+    const contexts = [
+      [null, false],
+      [[], false],
+      [{}, false],
+      [{ getRemainingTimeInMillis: null }, false],
+      [{ getRemainingTimeInMillis: undefined }, false],
+      [{ getRemainingTimeInMillis: 'not a function' }, false],
+      [{ getRemainingTimeInMillis: () => 100 }, true],
+    ]
+
+    assertExtractContext([], false)
+    assertExtractContext([{}], false)
+    contexts.forEach(([context, doesExtract], index) => {
+      describe(`using context case ${index + 1}`, () => {
+        assertExtractContext([{}, context], doesExtract)
+        assertExtractContext([{}, {}, context], doesExtract)
+        assertExtractContext([{}, {}, {}, context], false)
+      })
+    })
+  })
+})

packages/dd-trace/test/lambda/fixtures/handler.js

@@ -65,11 +65,67 @@ const errorHandler = async (_event, _context) => {
   throw new CustomError('my error')
 }
 
+/**
+ * Lambda Authorizer handler - only receives event, no context.
+ * This is the signature used by API Gateway Lambda Authorizers.
+ */
+const authorizerHandler = async (event) => {
+  // Simulate a simple authorizer that returns an IAM policy
+  return {
+    principalId: 'user123',
+    policyDocument: {
+      Version: '2012-10-17',
+      Statement: [
+        {
+          Action: 'execute-api:Invoke',
+          Effect: 'Allow',
+          Resource: event.methodArn || '*',
+        },
+      ],
+    },
+  }
+}
+
+/**
+ * Synchronous Lambda Authorizer handler - only receives event, no context.
+ */
+const authorizerHandlerSync = (event) => {
+  return {
+    principalId: 'user123',
+    policyDocument: {
+      Version: '2012-10-17',
+      Statement: [
+        {
+          Action: 'execute-api:Invoke',
+          Effect: 'Allow',
+          Resource: event.methodArn || '*',
+        },
+      ],
+    },
+  }
+}
+
+/**
+ * Lambda Authorizer handler that throws an error.
+ */
+const authorizerErrorHandler = async (event) => {
+  class AuthorizationError extends Error {
+    constructor (message) {
+      super(message)
+      Object.defineProperty(this, 'name', { value: 'AuthorizationError' })
+    }
+  }
+  throw new AuthorizationError('Unauthorized')
+}
+
 module.exports = {
   finishSpansEarlyTimeoutHandler,
   handler,
   swappedArgsHandler,
   timeoutHandler,
   errorHandler,
   callbackHandler,
+  authorizerHandler,
+  authorizerHandlerSync,
+  authorizerErrorHandler,
 }

packages/dd-trace/test/lambda/index.spec.js

@@ -218,6 +218,125 @@ describe('lambda', () => {
     })
   })
 
+  describe('lambda authorizers (no context)', () => {
+    beforeEach(setup)
+
+    afterEach(() => {
+      restoreEnv()
+      return closeAgent()
+    })
+
+    it('patches async lambda authorizer correctly (event only, no context)', async () => {
+      // Set the desired handler to patch
+      process.env.DD_LAMBDA_HANDLER = 'handler.authorizerHandler'
+      // Load the agent and re-register hook for patching.
+      await loadAgent()
+
+      // Lambda Authorizers only receive an event, no context
+      const _event = {
+        type: 'REQUEST',
+        methodArn: 'arn:aws:execute-api:us-east-1:123456789012:api-id/stage/GET/resource',
+        headers: {
+          Authorization: 'Bearer token123',
+        },
+      }
+
+      // Mock `datadog-lambda` handler resolve and import.
+      const _handlerPath = path.resolve(__dirname, './fixtures/handler.js')
+      const app = require(_handlerPath)
+      datadog = require('./fixtures/datadog-lambda')
+
+      // Run the function without context - this should NOT throw
+      const result = await datadog(app.authorizerHandler)(_event)
+
+      assert.notStrictEqual(result, undefined)
+      assert.strictEqual(result.principalId, 'user123')
+      assert.strictEqual(result.policyDocument.Statement[0].Effect, 'Allow')
+
+      // Expect traces to be correct.
+      const checkTraces = agent.assertSomeTraces((_traces) => {
+        const traces = _traces[0]
+        assert.strictEqual(traces.length, 1)
+        traces.forEach((trace) => {
+          assert.strictEqual(trace.error, 0)
+        })
+      })
+      await checkTraces
+    })
+
+    it('patches sync lambda authorizer correctly (event only, no context)', async () => {
+      // Set the desired handler to patch
+      process.env.DD_LAMBDA_HANDLER = 'handler.authorizerHandlerSync'
+      // Load the agent and re-register hook for patching.
+      await loadAgent()
+
+      // Lambda Authorizers only receive an event, no context
+      const _event = {
+        type: 'REQUEST',
+        methodArn: 'arn:aws:execute-api:us-east-1:123456789012:api-id/stage/GET/resource',
+      }
+
+      // Mock `datadog-lambda` handler resolve and import.
+      const _handlerPath = path.resolve(__dirname, './fixtures/handler.js')
+      const app = require(_handlerPath)
+      datadog = require('./fixtures/datadog-lambda')
+
+      // Run the function without context - this should NOT throw
+      // Note: datadog-lambda mock wraps in async, so we need to await
+      const result = await datadog(app.authorizerHandlerSync)(_event)
+
+      assert.notStrictEqual(result, undefined)
+      assert.strictEqual(result.principalId, 'user123')
+      assert.strictEqual(result.policyDocument.Statement[0].Effect, 'Allow')
+
+      // Expect traces to be correct.
+      const checkTraces = agent.assertSomeTraces((_traces) => {
+        const traces = _traces[0]
+        assert.strictEqual(traces.length, 1)
+        traces.forEach((trace) => {
+          assert.strictEqual(trace.error, 0)
+        })
+      })
+      await checkTraces
+    })
+
+    it('handles errors in lambda authorizer correctly (event only, no context)', async () => {
+      // Set the desired handler to patch
+      process.env.DD_LAMBDA_HANDLER = 'handler.authorizerErrorHandler'
+      // Load the agent and re-register hook for patching.
+      await loadAgent()
+
+      const _event = {
+        type: 'REQUEST',
+        methodArn: 'arn:aws:execute-api:us-east-1:123456789012:api-id/stage/GET/resource',
+      }
+
+      // Mock `datadog-lambda` handler resolve and import.
+      const _handlerPath = path.resolve(__dirname, './fixtures/handler.js')
+      const app = require(_handlerPath)
+      datadog = require('./fixtures/datadog-lambda')
+
+      // Run the function - should throw but not because of missing context
+      try {
+        await datadog(app.authorizerErrorHandler)(_event)
+        assert.fail('Expected error to be thrown')
+      } catch (e) {
+        assert.strictEqual(e.name, 'AuthorizationError')
+        assert.strictEqual(e.message, 'Unauthorized')
+      }
+
+      // Expect traces to be correct with error.
+      const checkTraces = agent.assertSomeTraces((_traces) => {
+        const traces = _traces[0]
+        assert.strictEqual(traces.length, 1)
+        traces.forEach((trace) => {
+          assert.strictEqual(trace.error, 1)
+        })
+      })
+      await checkTraces
+    })
+  })
+
   describe('timeout spans', () => {
     beforeEach(setup)