fix: Support multiple internal extensions (#74)

* fix: Support multiple internal extensions

* chore: Add comment

Author: ServerlessLife

.vscode/launch.json

@@ -24,6 +24,16 @@
       "type": "node",
       "cwd": "${workspaceRoot}/test/cdk-basic"
     },
+    {
+      "name": "LLDebugger - CDK basic - remove",
+      "program": "${workspaceRoot}/node_modules/tsx/dist/cli.mjs",
+      "args": ["../../src/lldebugger.ts", "-c environment=test", "-r"],
+      "request": "launch",
+      "skipFiles": ["<node_internals>/**"],
+      "console": "integratedTerminal",
+      "type": "node",
+      "cwd": "${workspaceRoot}/test/cdk-basic"
+    },
     {
       "name": "LLDebugger - CDK basic - monorepo",
       "program": "${workspaceRoot}/node_modules/tsx/dist/cli.mjs",

src/extension/interceptor.ts

@@ -6,6 +6,7 @@ const workerId = crypto.randomBytes(16).toString('hex');
 const topic = `${process.env.LLD_DEBUGGER_ID}/events/${workerId}`;
 
 const ORIGINAL_HANDLER_KEY = 'ORIGINAL_HANDLER';
+const originalHandlerName = process.env[ORIGINAL_HANDLER_KEY];
 const observableInterval = process.env.LLD_OBSERVABLE_INTERVAL
   ? parseInt(process.env.LLD_OBSERVABLE_INTERVAL!)
   : 0;
@@ -106,10 +107,65 @@ async function regularMode(context: any, event: any) {
  * Observable mode, which sends the event to the IoT service and doesn't wait for a response. It executes the original handler.
  */
 async function observableMode(context: any, event: any) {
-  const regularHandler = async () => {
-    const handler = await getOriginalHandler();
-    return await handler(event, context);
-  };
+  let regularHandler: undefined | (() => Promise<any>) = undefined;
+
+  if (process.env.LLD_INITIAL_AWS_LAMBDA_EXEC_WRAPPER) {
+    try {
+      Logger.log(
+        `Another extensions exists ${process.env.LLD_INITIAL_AWS_LAMBDA_EXEC_WRAPPER}.`,
+      );
+
+      const { promisify } = require('util');
+      const exec = require('child_process').exec;
+      const execAsync = promisify(exec);
+
+      // read the content of the script
+      const fs = require('fs/promises');
+      const originalScript = await fs.readFile(
+        process.env.LLD_INITIAL_AWS_LAMBDA_EXEC_WRAPPER,
+        'utf8',
+      );
+
+      Logger.verbose('Original script', originalScript);
+
+      // - set original handler
+      // - run second extension script
+      // - print environment variables
+      const script = `export _HANDLER=${process.env.ORIGINAL_HANDLER}
+                      ${originalScript}
+                      echo _HANDLER=$_HANDLER`;
+
+      Logger.verbose('Execute script', script);
+
+      const response = await execAsync(script);
+
+      Logger.verbose(`Output of the script: ${response.stdout}`);
+      // parse environment variables I got from the script
+      const handlerLine = response.stdout
+        .split('\n')
+        .find((line: string) => line.startsWith('_HANDLER'));
+      const oldHandler = handlerLine.split('=')[1];
+
+      Logger.verbose(`Getting handler "${oldHandler}" for another extension`);
+
+      regularHandler = async () => {
+        const handler = await getOriginalHandler(oldHandler);
+        return await handler(event, context);
+      };
+    } catch (e: any) {
+      Logger.error(
+        `Error while running the initial AWS_LAMBDA_EXEC_WRAPPER: ${e.message}`,
+        e,
+      );
+    }
+  }
+
+  if (!regularHandler) {
+    regularHandler = async () => {
+      const handler = await getOriginalHandler(originalHandlerName);
+      return await handler(event, context);
+    };
+  }
 
   const observableHandler = async () => {
     // prevent sending too many events
@@ -155,14 +211,18 @@ async function observableMode(context: any, event: any) {
   return response;
 }
 
-async function getOriginalHandler(): Promise<any> {
+async function getOriginalHandler(
+  originalHandlerName: string | undefined,
+): Promise<any> {
+  Logger.verbose('Original handler:', originalHandlerName);
+
   // @ts-ignore
   const { load } = await import('./aws/UserFunction');
 
-  if (process.env[ORIGINAL_HANDLER_KEY] === undefined)
+  if (originalHandlerName === undefined)
     throw Error('Missing original handler');
   return load(
     process.env.LAMBDA_TASK_ROOT!,
-    process.env[ORIGINAL_HANDLER_KEY],
+    originalHandlerName,
   ) as Promise<any>;
 }

src/infraDeploy.ts

@@ -26,6 +26,7 @@ let iamClient: IAMClient | undefined;
 
 const inlinePolicyName = 'LambdaLiveDebuggerPolicy';
 const layerName = 'LambdaLiveDebugger';
+const lldWrapperPath = '/opt/lld-wrapper';
 
 /**
  * Policy document to attach to the Lambda role
@@ -300,7 +301,16 @@ async function removeLayerFromLambda(functionName: string) {
       );
     }
 
-    const ddlEnvironmentVariables = getEnvironmentVarablesForDebugger('xxx', 0);
+    const initalExecWraper =
+      environmentVariables.LLD_INITIAL_AWS_LAMBDA_EXEC_WRAPPER;
+
+    const ddlEnvironmentVariables = getEnvironmentVarablesForDebugger({
+      // set dummy data, so we just get the list of environment variables
+      functionId: 'xxx',
+      timeout: 0,
+      verbose: true,
+      initalExecWraper: 'test',
+    });
 
     // check if environment variables are set for each property
     for (const [key] of Object.entries(ddlEnvironmentVariables)) {
@@ -323,10 +333,22 @@ async function removeLayerFromLambda(functionName: string) {
       //remove environment variables
       for (const [key] of Object.entries(ddlEnvironmentVariables)) {
         if (environmentVariables && environmentVariables[key]) {
-          delete environmentVariables[key];
+          if (key === 'AWS_LAMBDA_EXEC_WRAPPER') {
+            if (environmentVariables[key] === lldWrapperPath) {
+              delete environmentVariables[key];
+            } else {
+              // do not remove the original AWS_LAMBDA_EXEC_WRAPPER that was set before LLD
+            }
+          } else {
+            delete environmentVariables[key];
+          }
         }
       }
 
+      if (initalExecWraper) {
+        environmentVariables.AWS_LAMBDA_EXEC_WRAPPER = initalExecWraper;
+      }
+
       Logger.verbose(
         'New environment variables',
         JSON.stringify(environmentVariables, null, 2),
@@ -445,10 +467,24 @@ async function attachLayerToLambda(
     Logger.verbose('Layer with the wrong version attached to the function');
   }
 
-  const ddlEnvironmentVariables = getEnvironmentVarablesForDebugger(
+  // support for multiple internal Lambda extensions
+  const initalExecWraper =
+    environmentVariables.AWS_LAMBDA_EXEC_WRAPPER !== lldWrapperPath
+      ? environmentVariables.AWS_LAMBDA_EXEC_WRAPPER
+      : undefined;
+
+  if (initalExecWraper) {
+    Logger.warn(
+      `[Function ${functionName}] Another internal Lambda extension is already attached to the function, which might cause unpredictable behavior.`,
+    );
+  }
+
+  const ddlEnvironmentVariables = getEnvironmentVarablesForDebugger({
     functionId,
-    initialTimeout,
-  );
+    timeout: initialTimeout,
+    verbose: Configuration.config.verbose,
+    initalExecWraper,
+  });
 
   // check if environment variables are already set for each property
   for (const [key, value] of Object.entries(ddlEnvironmentVariables)) {
@@ -551,22 +587,36 @@ async function addPolicyToLambdaRole(functionName: string) {
 
 /**
  * Get the environment variables for the Lambda function
- * @param functionId
- * @param timeout
- * @returns
  */
-function getEnvironmentVarablesForDebugger(
-  functionId: string,
-  timeout: number | undefined,
-): Record<string, string> {
-  return {
+function getEnvironmentVarablesForDebugger({
+  functionId,
+  timeout,
+  verbose,
+  initalExecWraper,
+}: {
+  functionId: string;
+  timeout: number | undefined;
+  verbose: boolean | undefined;
+  initalExecWraper: string | undefined;
+}): Record<string, string> {
+  const env: Record<string, string> = {
     LLD_FUNCTION_ID: functionId,
-    AWS_LAMBDA_EXEC_WRAPPER: '/opt/lld-wrapper',
+    AWS_LAMBDA_EXEC_WRAPPER: lldWrapperPath,
     LLD_DEBUGGER_ID: Configuration.config.debuggerId,
     LLD_INITIAL_TIMEOUT: timeout ? timeout.toString() : '-1', // should never be negative
     LLD_OBSERVABLE_MODE: Configuration.config.observable ? 'true' : 'false',
     LLD_OBSERVABLE_INTERVAL: Configuration.config.interval.toString(),
   };
+
+  if (initalExecWraper) {
+    env.LLD_INITIAL_AWS_LAMBDA_EXEC_WRAPPER = initalExecWraper;
+  }
+
+  if (verbose) {
+    env.LLD_VERBOSE = 'true';
+  }
+
+  return env;
 }
 
 /**

src/logger.ts

@@ -63,11 +63,20 @@ function setVerbose(enabled: boolean) {
   verboseEnabled = enabled;
 }
 
+/**
+ * Check if verbose logging is enabled
+ * @returns Whether verbose logging is enabled
+ */
+function isVerbose() {
+  return verboseEnabled;
+}
+
 export const Logger = {
   log,
   error,
   warn,
   important,
   verbose,
   setVerbose,
+  isVerbose,
 };