From 8b62f248c2ddae5e934e8dc02583f23cc07f0ded Mon Sep 17 00:00:00 2001
From: Yann Normand <yann.normand@gmail.com>
Date: Wed, 16 Oct 2024 22:10:37 +1000
Subject: [PATCH] fix: relax ContextMap value constraint

---
 src/context.ts                                 |  6 +++---
 src/request.ts                                 | 12 +++++-------
 src/tests/index.spec.ts                        | 17 +++++++++++++++++
 src/types.ts                                   |  2 +-
 tap-snapshots/src/tests/index.spec.ts.test.cjs |  4 ++++
 5 files changed, 30 insertions(+), 11 deletions(-)

diff --git a/src/context.ts b/src/context.ts
index 56518af..0f87f99 100644
--- a/src/context.ts
+++ b/src/context.ts
@@ -3,19 +3,19 @@ const CONTEXT_SYMBOL = Symbol.for('aws.lambda.runtime.context');
 
 export interface ContextMap {
   awsRequestId: string;
-  [key: string]: string;
+  [key: string]: unknown;
 }
 
 export interface ContextStorageProvider {
   getContext: () => ContextMap;
   setContext: (map: ContextMap) => void;
-  updateContext: (values: Record<string, string>) => void;
+  updateContext: (values: Record<string, unknown>) => void;
 }
 
 export const GlobalContextStorageProvider: ContextStorageProvider = {
   getContext: () => (global as any)[CONTEXT_SYMBOL] as ContextMap,
   setContext: (map: ContextMap) => ((global as any)[CONTEXT_SYMBOL] = map),
-  updateContext: (values: Record<string, string>) => {
+  updateContext: (values: Record<string, unknown>) => {
     const ctx = GlobalContextStorageProvider.getContext();
     (global as any)[CONTEXT_SYMBOL] = {
       ...ctx,
diff --git a/src/request.ts b/src/request.ts
index 767f07d..e17be94 100644
--- a/src/request.ts
+++ b/src/request.ts
@@ -26,17 +26,16 @@ export const lambdaRequestTracker = (
 
   // capture any correlation headers sent from upstream callers
   if (event.headers) {
-    Object.keys(event.headers).forEach((header) => {
+    for (const [header, value] of Object.entries(event.headers)) {
       if (header.toLowerCase().startsWith(CORRELATION_HEADER)) {
-        // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
-        ctx[header] = event.headers![header] as string;
+        ctx[header] = value;
       }
-    });
+    }
   }
 
   // capture the xray trace id if its enabled
   if (process.env[AMAZON_TRACE_ID]) {
-    ctx[CORRELATION_TRACE_ID] = process.env[AMAZON_TRACE_ID] as string;
+    ctx[CORRELATION_TRACE_ID] = process.env[AMAZON_TRACE_ID];
   }
 
   // set the correlation id if not already set by upstream callers
@@ -49,9 +48,8 @@ export const lambdaRequestTracker = (
   if (options.requestMixin) {
     const result = options.requestMixin(event, context);
     for (const key in result) {
-      // Cast this to string for typescript
       // when the JSON serializer runs, by default it omits undefined properties
-      ctx[key] = result[key] as string;
+      ctx[key] = result[key];
     }
   }
 
diff --git a/src/tests/index.spec.ts b/src/tests/index.spec.ts
index 4a2e1ad..08a4e57 100644
--- a/src/tests/index.spec.ts
+++ b/src/tests/index.spec.ts
@@ -180,6 +180,23 @@ tap.test('should allow removing default request data', (t) => {
   t.end();
 });
 
+tap.test('should allow different types of context values', (t) => {
+  const { log, output, withRequest } = createLogger(undefined, {
+    requestMixin: () => ({
+      'number': 12,
+      'boolean': true,
+      'object': {
+        foo: "bar"
+      }
+    }),
+  });
+
+  withRequest({}, { awsRequestId: '431234' });
+  log.info('Message with trace ID');
+  t.matchSnapshot(output.buffer);
+  t.end();
+});
+
 tap.test('should allow default pino formatter', (t) => {
   const { log, output, withRequest } = createLogger(undefined, {
     formatter: new PinoLogFormatter(),
diff --git a/src/types.ts b/src/types.ts
index ec6db53..b59733f 100644
--- a/src/types.ts
+++ b/src/types.ts
@@ -70,7 +70,7 @@ export interface LambdaRequestTrackerOptions {
   requestMixin?: (
     event: LambdaEvent,
     context: LambdaContext,
-  ) => { [key: string]: string | undefined };
+  ) => { [key: string]: unknown };
   /**
    * Custom storage provider.
    * If not supplied, defaults to global context storage.
diff --git a/tap-snapshots/src/tests/index.spec.ts.test.cjs b/tap-snapshots/src/tests/index.spec.ts.test.cjs
index d64e1bb..bdb5517 100644
--- a/tap-snapshots/src/tests/index.spec.ts.test.cjs
+++ b/tap-snapshots/src/tests/index.spec.ts.test.cjs
@@ -17,6 +17,10 @@ exports[`src/tests/index.spec.ts TAP should allow default pino formatter > must
 {"awsRequestId":"431234","x-correlation-id":"431234","level":30,"time":1480572000000,"msg":"Message with pino formatter"}
 `
 
+exports[`src/tests/index.spec.ts TAP should allow different types of context values > must match snapshot 1`] = `
+2016-12-01T06:00:00.000Z	431234	INFO	Message with trace ID	{"awsRequestId":"431234","x-correlation-id":"431234","number":12,"boolean":true,"object":{"foo":"bar"},"level":30,"time":1480572000000,"msg":"Message with trace ID"}
+`
+
 exports[`src/tests/index.spec.ts TAP should allow modifying the global context > must match snapshot 1`] = `
 2016-12-01T06:00:00.000Z	98875	ERROR	Context updates	{"awsRequestId":"98875","x-correlation-data":"abbb","x-correlation-id":"98875","userId":"12","level":50,"time":1480572000000,"msg":"Context updates"}
 `