From 8bdf2edc47dc5fcca7a2364c33b8328985e5e087 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Rafa=C5=82=20Sumis=C5=82awski?=
 <rafal.sumislawski@coralogix.com>
Date: Tue, 23 Apr 2024 11:33:30 +0200
Subject: [PATCH] Feature/ngstn 726 sst support (#12)

This PR extends the logic of applying lambda handler instrumentation in two ways.
* Supporting handler defined in *.mjs files
* Adding additinal InstrumentationNodeModuleDefinition matching by filename, which turns out to be what is needed in a module environment with the import-in-the-middle-hook
---
 .../src/instrumentation.ts                    | 36 ++++++++++++++++---
 1 file changed, 31 insertions(+), 5 deletions(-)

diff --git a/plugins/node/opentelemetry-instrumentation-aws-lambda/src/instrumentation.ts b/plugins/node/opentelemetry-instrumentation-aws-lambda/src/instrumentation.ts
index be7a58fa59..450b1bc905 100644
--- a/plugins/node/opentelemetry-instrumentation-aws-lambda/src/instrumentation.ts
+++ b/plugins/node/opentelemetry-instrumentation-aws-lambda/src/instrumentation.ts
@@ -123,15 +123,23 @@ export class AwsLambdaInstrumentation extends InstrumentationBase {
 
     // Lambda loads user function using an absolute path.
     let filename = path.resolve(taskRoot, moduleRoot, module);
-    if (!filename.endsWith('.js')) {
-      // its impossible to know in advance if the user has a cjs or js file.
-      // check that the .js file exists otherwise fallback to next known possibility
+    if (!filename.includes('.')) {
+      // The filename has no extension, we need to check if there is a .js, .cjs or .mjs file
       try {
         fs.statSync(`${filename}.js`);
         filename += '.js';
       } catch (e) {
-        // fallback to .cjs
-        filename += '.cjs';
+        try {
+          fs.statSync(`${filename}.cjs`);
+          filename += '.cjs';
+        } catch (e) {
+          try {
+            fs.statSync(`${filename}.mjs`);
+            filename += '.mjs';
+          } catch (e) {
+            diag.warn(`AwsLambdaInstrumentation couldn't find the handler file: ${filename}`);
+          }
+        }
       }
     }
 
@@ -192,6 +200,24 @@ export class AwsLambdaInstrumentation extends InstrumentationBase {
           this._unwrap(moduleExports, functionName);
         }
       ),
+      // This approach works when the handler is an .mjs file in the function source code
+      new InstrumentationNodeModuleDefinition(
+        filename,
+        ['*'],
+        (moduleExports: LambdaModule) => {
+          diag.debug('Applying patch for lambda handler');
+          if (isWrapped(moduleExports[functionName])) {
+            this._unwrap(moduleExports, functionName);
+          }
+          this._wrap(moduleExports, functionName, this._getHandler());
+          return moduleExports;
+        },
+        (moduleExports?: LambdaModule) => {
+          if (moduleExports === undefined) return;
+          diag.debug('Removing patch for lambda handler');
+          this._unwrap(moduleExports, functionName);
+        }
+      ),
     ];
   }