Skip to content

added support for vault keys usage inside NodeJS code component #88

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Aug 12, 2025
Merged

added support for vault keys usage inside NodeJS code component #88

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
f56bf13
added support for vault keys usage inside NodeJS code component
zubairirfan96 Aug 11, 2025
625c03b
removed misleading comments
zubairirfan96 Aug 11, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
104 changes: 82 additions & 22 deletions packages/core/src/helpers/AWSLambdaCode.helper.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@ import crypto from 'crypto';
import { ConnectorService } from '@sre/Core/ConnectorsService';
import { AccessCandidate } from '@sre/Security/AccessControl/AccessCandidate.class';
import zl from 'zip-lib';
import { InvokeCommand, Runtime, LambdaClient, UpdateFunctionCodeCommand, CreateFunctionCommand, GetFunctionCommand, GetFunctionCommandOutput, InvokeCommandOutput } from '@aws-sdk/client-lambda';
import { InvokeCommand, Runtime, LambdaClient, UpdateFunctionCodeCommand, CreateFunctionCommand, GetFunctionCommand, GetFunctionCommandOutput, InvokeCommandOutput, UpdateFunctionConfigurationCommand } from '@aws-sdk/client-lambda';
import { GetRoleCommand, CreateRoleCommand, IAMClient, GetRoleCommandOutput, CreateRoleCommandOutput } from '@aws-sdk/client-iam';
import fs from 'fs';
import { AWSConfig, AWSCredentials, AWSRegionConfig } from '@sre/types/AWS.types';
Expand All @@ -20,13 +20,15 @@ export function getLambdaFunctionName(agentId: string, componentId: string) {
}


export function generateCodeHash(code_body: string, codeInputs: string[]) {
export function generateCodeHash(code_body: string, codeInputs: string[], envVariables: string[]) {
const bodyHash = getSanitizeCodeHash(code_body);
const inputsHash = getSanitizeCodeHash(JSON.stringify(codeInputs));
return `body-${bodyHash}__inputs-${inputsHash}`;
const envVariablesHash = getSanitizeCodeHash(JSON.stringify(envVariables));
return `body-${bodyHash}__inputs-${inputsHash}__env-${envVariablesHash}`;
}

export function getSanitizeCodeHash(code: string) {
export function getSanitizeCodeHash(rawCode: string) {
const code = replaceVaultKeysTemplateVars(rawCode, {});
let output = '';
let isSingleQuote = false;
let isDoubleQuote = false;
Expand Down Expand Up @@ -82,9 +84,15 @@ export async function setDeployedCodeHash(agentId: string, componentId: string,
.set(`${cachePrefix}_${agentId}-${componentId}`, codeHash, null, null, cacheTTL);
}

export function generateLambdaCode(code: string, parameters: string[]) {
function replaceVaultKeysTemplateVars(code: string, envVariables: Record<string, string>) {
const regex = /\{\{KEY\((.*?)\)\}\}/g;
return code.replaceAll(regex, (match, p1) => `process.env.${p1}`);
}

export function generateLambdaCode(code: string, parameters: string[], envVariables: Record<string, string>) {
const codeWithEnvVariables = envVariables && Object.keys(envVariables).length ? replaceVaultKeysTemplateVars(code, envVariables) : code;
const lambdaCode = `
${code}
${codeWithEnvVariables}
export const handler = async (event, context) => {
try {
context.callbackWaitsForEmptyEventLoop = false;
Expand Down Expand Up @@ -118,7 +126,7 @@ export async function zipCode(directory: string) {
});
}

export async function createOrUpdateLambdaFunction(functionName, zipFilePath, awsConfigs) {
export async function createOrUpdateLambdaFunction(functionName, zipFilePath, awsConfigs, envVariables: Record<string, string>) {
const client = new LambdaClient({
region: awsConfigs.region,
credentials: {
Expand All @@ -142,9 +150,12 @@ export async function createOrUpdateLambdaFunction(functionName, zipFilePath, aw
};
const updateFunctionCodeCommand = new UpdateFunctionCodeCommand(updateCodeParams);
await client.send(updateFunctionCodeCommand);
// Update function configuration to attach layer
await verifyFunctionDeploymentStatus(functionName, client);
// console.log('Lambda function code and configuration updated successfully!');

if (envVariables && Object.keys(envVariables).length) {
await updateLambdaFunctionConfiguration(client, functionName, envVariables);
await verifyFunctionDeploymentStatus(functionName, client);
}
} else {
// Create function if it does not exist
let roleArn = '';
Expand Down Expand Up @@ -188,6 +199,7 @@ export async function createOrUpdateLambdaFunction(functionName, zipFilePath, aw
'auto-delete': 'true',
},
MemorySize: 256,
...(envVariables && Object.keys(envVariables).length ? { Environment: { Variables: envVariables } } : {}),
};

const functionCreateCommand = new CreateFunctionCommand(functionParams);
Expand All @@ -200,6 +212,17 @@ export async function createOrUpdateLambdaFunction(functionName, zipFilePath, aw
}
}

function updateLambdaFunctionConfiguration(client: LambdaClient, functionName: string, envVariables: Record<string, string>) {
const updateFunctionConfigurationParams = {
FunctionName: functionName,
Environment: {
Variables: envVariables,
},
};
const updateFunctionConfigurationCommand = new UpdateFunctionConfigurationCommand(updateFunctionConfigurationParams);
return client.send(updateFunctionConfigurationCommand);
}

export async function waitForRoleDeploymentStatus(roleName, client): Promise<boolean> {
return new Promise((resolve, reject) => {
try {
Expand Down Expand Up @@ -356,10 +379,11 @@ export function reportUsage({ cost, agentId, teamId }: { cost: number; agentId:
});
}

export function validateAsyncMainFunction(code: string): { isValid: boolean; error?: string; parameters?: string[]; dependencies?: string[] } {
export function validateAsyncMainFunction(rawCode: string): { isValid: boolean; error?: string; parameters?: string[]; dependencies?: string[] } {
try {
// Parse the code using acorn
const ast = acorn.parse(code, {
const code = replaceVaultKeysTemplateVars(rawCode, {});
// Parse the code using acorn
const ast = acorn.parse(code, {
ecmaVersion: 'latest',
sourceType: 'module'
});
Expand Down Expand Up @@ -387,8 +411,8 @@ export function validateAsyncMainFunction(code: string): { isValid: boolean; err
}

// Handle CallExpression (require() calls)
if (node.type === 'CallExpression' &&
node.callee.type === 'Identifier' &&
if (node.type === 'CallExpression' &&
node.callee.type === 'Identifier' &&
node.callee.name === 'require' &&
node.arguments.length > 0 &&
node.arguments[0].type === 'Literal') {
Expand All @@ -399,7 +423,7 @@ export function validateAsyncMainFunction(code: string): { isValid: boolean; err
}

// Handle dynamic import() calls
if (node.type === 'CallExpression' &&
if (node.type === 'CallExpression' &&
node.callee.type === 'Import' &&
node.arguments.length > 0 &&
node.arguments[0].type === 'Literal') {
Expand Down Expand Up @@ -473,26 +497,26 @@ export function validateAsyncMainFunction(code: string): { isValid: boolean; err
}

if (!hasMain) {
return {
isValid: false,
return {
isValid: false,
error: 'No main function found at root level',
dependencies
};
}

if (!hasAsyncMain) {
return {
isValid: false,
return {
isValid: false,
error: 'Main function exists but is not async',
dependencies
};
}

return { isValid: true, parameters: mainParameters, dependencies };
} catch (error) {
return {
isValid: false,
error: `Failed to parse code: ${error.message}`
return {
isValid: false,
error: `Failed to parse code: ${error.message}`
};
}
}
Expand Down Expand Up @@ -526,3 +550,39 @@ export function generateCodeFromLegacyComponent(code_body: string, code_imports:
return code;
}

export function extractAllKeyNamesFromTemplateVars(input: string): string[] {
const regex = /\{\{KEY\((.*?)\)\}\}/g;
const matches = [];
let match;
while ((match = regex.exec(input)) !== null) {
if (match[1]) {
matches.push(match[1]);
}
}
return matches;
}
Comment on lines +553 to +563
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🛠️ Refactor suggestion

Add validation for vault key names

Consider validating that extracted key names are valid environment variable names to prevent issues with special characters.

 export function extractAllKeyNamesFromTemplateVars(input: string): string[] {
     const regex = /\{\{KEY\((.*?)\)\}\}/g;
     const matches = [];
     let match;
     while ((match = regex.exec(input)) !== null) {
         if (match[1]) {
+            // Validate key name contains only alphanumeric and underscore
+            if (!/^[a-zA-Z_][a-zA-Z0-9_]*$/.test(match[1])) {
+                throw new Error(`Invalid vault key name: "${match[1]}". Key names must be valid environment variable names.`);
+            }
             matches.push(match[1]);
         }
     }
     return matches;
 }
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
export function extractAllKeyNamesFromTemplateVars(input: string): string[] {
const regex = /\{\{KEY\((.*?)\)\}\}/g;
const matches = [];
let match;
while ((match = regex.exec(input)) !== null) {
if (match[1]) {
matches.push(match[1]);
}
}
return matches;
}
export function extractAllKeyNamesFromTemplateVars(input: string): string[] {
const regex = /\{\{KEY\((.*?)\)\}\}/g;
const matches = [];
let match;
while ((match = regex.exec(input)) !== null) {
if (match[1]) {
// Validate key name contains only alphanumeric characters and underscores, and doesn’t start with a digit
if (!/^[a-zA-Z_][a-zA-Z0-9_]*$/.test(match[1])) {
throw new Error(
`Invalid vault key name: "${match[1]}". Key names must be valid environment variable names.`
);
}
matches.push(match[1]);
}
}
return matches;
}
🤖 Prompt for AI Agents 
In packages/core/src/helpers/AWSLambdaCode.helper.ts around lines 553 to 563,
the extracted vault key names from the template are not validated; update the
function to validate each extracted name against a standard environment-variable
pattern (/^[A-Za-z_][A-Za-z0-9_]*$/), collect any invalid names and either
filter them out or (preferably) throw an Error listing invalid keys so callers
know which entries are bad; ensure the function returns only the valid key names
when successful and include a clear error message naming the offending keys if
validation fails.



async function fetchVaultSecret(keyName: string, agentTeamId: string): Promise<{ value: string, key: string }> {
const vaultSecret = await VaultHelper.getAgentKey(keyName, agentTeamId);
return {
value: vaultSecret,
key: keyName,
};

}
Comment on lines +566 to +573
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue

Add error handling for missing vault keys

The fetchVaultSecret function doesn't handle cases where a vault key might not exist. This could cause the entire deployment to fail if a referenced key is missing.

Consider adding error handling:

 async function fetchVaultSecret(keyName: string, agentTeamId: string): Promise<{ value: string, key: string }> {
-    const vaultSecret = await VaultHelper.getAgentKey(keyName, agentTeamId);
-    return {
-        value: vaultSecret,
-        key: keyName,
-    };
+    try {
+        const vaultSecret = await VaultHelper.getAgentKey(keyName, agentTeamId);
+        return {
+            value: vaultSecret,
+            key: keyName,
+        };
+    } catch (error) {
+        console.error(`Failed to fetch vault key "${keyName}":`, error);
+        throw new Error(`Vault key "${keyName}" not found or inaccessible`);
+    }
 }
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
async function fetchVaultSecret(keyName: string, agentTeamId: string): Promise<{ value: string, key: string }> {
const vaultSecret = await VaultHelper.getAgentKey(keyName, agentTeamId);
return {
value: vaultSecret,
key: keyName,
};
}
async function fetchVaultSecret(keyName: string, agentTeamId: string): Promise<{ value: string, key: string }> {
try {
const vaultSecret = await VaultHelper.getAgentKey(keyName, agentTeamId);
return {
value: vaultSecret,
key: keyName,
};
} catch (error) {
console.error(`Failed to fetch vault key "${keyName}":`, error);
throw new Error(`Vault key "${keyName}" not found or inaccessible`);
}
}
🤖 Prompt for AI Agents 
In packages/core/src/helpers/AWSLambdaCode.helper.ts around lines 566 to 573,
fetchVaultSecret currently assumes VaultHelper.getAgentKey always returns a
value; add handling for missing or empty secrets by checking the returned
vaultSecret (null/undefined/empty string) and, if absent, either throw a
descriptive Error (e.g. `throw new Error(\`Vault secret not found for key
${keyName} and team ${agentTeamId}\`)`) or return a clearly documented fallback
structure indicating the missing key; ensure the caller contract is preserved
(update callers if you choose to throw) and include a short
processLogger.warn/error entry before throwing/returning to aid debugging.


export async function getCurrentEnvironmentVariables(agentTeamId: string, code: string): Promise<Record<string, string>> {
const allKeyNames = extractAllKeyNamesFromTemplateVars(code);
const envVariables: Record<string, string> = {};
const vaultSecrets = await Promise.all(allKeyNames.map((keyName) => fetchVaultSecret(keyName, agentTeamId)));
vaultSecrets.forEach((secret) => {
envVariables[secret.key] = secret.value;
});
return envVariables;
}

export function getSortedObjectValues(obj: Record<string, string>): string[] {
const sortedKeys = Object.keys(obj).sort();
return sortedKeys.map((key) => obj[key]);
}
18 changes: 10 additions & 8 deletions packages/core/src/subsystems/ComputeManager/Code.service/connectors/AWSLambdaCode.class.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@ import { AccessRequest } from '@sre/Security/AccessControl/AccessRequest.class';
import fs from 'fs';
import path from 'path';
import { execSync } from 'child_process';
import { cacheTTL, createOrUpdateLambdaFunction, generateCodeHash, generateLambdaCode, getDeployedCodeHash, getDeployedFunction, getLambdaFunctionName, invokeLambdaFunction, setDeployedCodeHash, updateDeployedCodeTTL, validateAsyncMainFunction, zipCode } from '@sre/helpers/AWSLambdaCode.helper';
import { cacheTTL, createOrUpdateLambdaFunction, generateCodeHash, generateLambdaCode, getCurrentEnvironmentVariables, getDeployedCodeHash, getDeployedFunction, getLambdaFunctionName, getSortedObjectValues, invokeLambdaFunction, setDeployedCodeHash, updateDeployedCodeTTL, validateAsyncMainFunction, zipCode } from '@sre/helpers/AWSLambdaCode.helper';
import { AWSCredentials, AWSRegionConfig } from '@sre/types/AWS.types';
import { Logger } from '@sre/helpers/Log.helper';

Expand All @@ -31,14 +31,16 @@ export class AWSLambdaCode extends CodeConnector {
public async deploy(acRequest: AccessRequest, codeUID: string, input: CodeInput, config: CodeConfig): Promise<CodeDeployment> {
const agentId = acRequest.candidate.id;
const functionName = getLambdaFunctionName(agentId, codeUID);


const [isLambdaExists, exisitingCodeHash] = await Promise.all([
const [isLambdaExists, exisitingCodeHash, currentEnvVariables] = await Promise.all([
getDeployedFunction(functionName, this.awsConfigs),
getDeployedCodeHash(agentId, codeUID),
getCurrentEnvironmentVariables(acRequest.candidate.id, input.code),
]);
Comment on lines +34 to 38
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue

Incorrect parameter passed to getCurrentEnvironmentVariables

The function getCurrentEnvironmentVariables expects agentTeamId as the first parameter, but you're passing acRequest.candidate.id which is the agentId. This could lead to incorrect vault secret retrieval.

Consider retrieving the team ID first:

-            getCurrentEnvironmentVariables(acRequest.candidate.id, input.code),
+            getCurrentEnvironmentVariables(/* need to get teamId from acRequest */, input.code),

You may need to fetch the team ID using the AccountConnector service similar to how it's done in the helper functions.

Committable suggestion skipped: line range outside the PR's diff.

🤖 Prompt for AI Agents 
In
packages/core/src/subsystems/ComputeManager/Code.service/connectors/AWSLambdaCode.class.ts
around lines 34 to 38, getCurrentEnvironmentVariables is being called with
acRequest.candidate.id (agentId) but it expects agentTeamId; fetch the agent's
team ID first (use the AccountConnector or the same helper used elsewhere to
resolve team ID from the agent/account) and then call
getCurrentEnvironmentVariables(agentTeamId, input.code) so the correct vault
secrets are retrieved.


const codeHash = generateCodeHash(input.code, Object.keys(input.inputs));
// sorting values to get the same hash for the same env variables
const envValues = getSortedObjectValues(currentEnvVariables);

const codeHash = generateCodeHash(input.code, Object.keys(input.inputs), envValues);
if (isLambdaExists && exisitingCodeHash === codeHash) {
return {
id: functionName,
Expand All @@ -57,7 +59,7 @@ export class AWSLambdaCode extends CodeConnector {
if (!isValid) {
throw new Error(error || 'Invalid Code')
}
const lambdaCode = generateLambdaCode(input.code, parameters);
const lambdaCode = generateLambdaCode(input.code, parameters, currentEnvVariables);
// create folder
fs.mkdirSync(directory);
// create index.js file
Expand All @@ -68,7 +70,7 @@ export class AWSLambdaCode extends CodeConnector {
execSync(`npm install ${dependencies.join(' ')}`, { cwd: directory });

const zipFilePath = await zipCode(directory);
await createOrUpdateLambdaFunction(functionName, zipFilePath, this.awsConfigs);
await createOrUpdateLambdaFunction(functionName, zipFilePath, this.awsConfigs, currentEnvVariables);
await setDeployedCodeHash(agentId, codeUID, codeHash);
console.debug('Lambda function updated successfully!');
return {
Expand All @@ -83,7 +85,7 @@ export class AWSLambdaCode extends CodeConnector {
try {
fs.rmSync(`${directory}`, { recursive: true, force: true });
fs.unlinkSync(`${directory}.zip`);
} catch (error) {}
} catch (error) { }
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🛠️ Refactor suggestion

Add error logging in catch block

The empty catch block silently swallows errors during cleanup. Consider logging the error for debugging purposes.

-            } catch (error) { }
+            } catch (error) {
+                console.debug('Cleanup error:', error);
+            }
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
} catch (error) { }
} catch (error) {
console.debug('Cleanup error:', error);
}
🤖 Prompt for AI Agents 
In
packages/core/src/subsystems/ComputeManager/Code.service/connectors/AWSLambdaCode.class.ts
around line 88, the catch block is empty and silently swallows errors during
cleanup; update the catch to log the error (including message and stack) with
the class/logger used in this file (e.g., this.logger?.error('Failed during
cleanup in AWSLambdaCode', error) or fall back to console.error if no logger),
ensure the caught value is handled as unknown and included in the log (e.g.,
String(error) and (error as Error).stack) to provide clear context for
debugging, and do not suppress the original error details.

}

}
Expand Down