@aikidosec/firewall
Version:
Zen by Aikido is an embedded Application Firewall that autonomously protects Node.js apps against common and critical attacks, provides rate limiting, detects malicious traffic (including bots), and more.
156 lines (155 loc) • 6.15 kB
JavaScript
;
Object.defineProperty(exports, "__esModule", { value: true });
exports.AwsSDKVersion3 = void 0;
const wrapExport_1 = require("../agent/hooks/wrapExport");
const isPlainObject_1 = require("../helpers/isPlainObject");
function isUsage(usage) {
return ((0, isPlainObject_1.isPlainObject)(usage) &&
typeof usage.input_tokens === "number" &&
typeof usage.output_tokens === "number");
}
function isInvokeResponse(response) {
return (0, isPlainObject_1.isPlainObject)(response);
}
function isConverseUsage(usage) {
return ((0, isPlainObject_1.isPlainObject)(usage) &&
typeof usage.inputTokens === "number" &&
typeof usage.outputTokens === "number");
}
function isConverseResponse(response) {
return (0, isPlainObject_1.isPlainObject)(response);
}
class AwsSDKVersion3 {
processInvokeModelResponse(response, command, agent) {
if (!isInvokeResponse(response)) {
return;
}
let body;
try {
const json = new TextDecoder("utf-8", {
fatal: true,
}).decode(response.body);
body = JSON.parse(json);
}
catch {
// Ignore errors when parsing the response body
return;
}
// @ts-expect-error We don't know the type of command
if (command && command.input && typeof command.input.modelId === "string") {
let inputTokens = 0;
let outputTokens = 0;
if (isUsage(body.usage)) {
inputTokens = body.usage.input_tokens;
outputTokens = body.usage.output_tokens;
}
let modelId = undefined;
// @ts-expect-error We don't know the type of command
if (!command.input.modelId.startsWith("arn:")) {
// @ts-expect-error We don't know the type of command
modelId = command.input.modelId;
}
if (!modelId && typeof body.model === "string") {
modelId = body.model;
}
if (!modelId) {
return;
}
const aiStats = agent.getAIStatistics();
aiStats.onAICall({
provider: "bedrock",
model: modelId,
inputTokens: inputTokens,
outputTokens: outputTokens,
});
}
}
processConverseResponse(response, command, agent) {
// @ts-expect-error We don't know the type of command
if (!command || !command.input || !command.input.modelId) {
return;
}
if (!isConverseResponse(response)) {
return;
}
// @ts-expect-error We don't know the type of command
const modelId = command.input.modelId;
// Don't report if modelId is an ARN
// There's no way to get the actual model name like we can with InvokeModel
if (modelId.startsWith("arn:")) {
return;
}
let inputTokens = 0;
let outputTokens = 0;
if (isConverseUsage(response.usage)) {
inputTokens = response.usage.inputTokens || 0;
outputTokens = response.usage.outputTokens || 0;
}
const aiStats = agent.getAIStatistics();
aiStats.onAICall({
provider: "bedrock",
model: modelId,
inputTokens: inputTokens,
outputTokens: outputTokens,
});
}
onRequire(exports, pkgInfo) {
if (exports.BedrockRuntimeClient) {
(0, wrapExport_1.wrapExport)(exports.BedrockRuntimeClient.prototype, "send", pkgInfo, {
kind: "ai_op",
modifyReturnValue: (args, returnValue, agent) => {
if (args.length > 0) {
const command = args[0];
if (returnValue instanceof Promise) {
// Inspect the response after the promise resolves, it won't change the original promise
returnValue
.then((response) => {
if (exports.InvokeModelCommand &&
command instanceof exports.InvokeModelCommand) {
this.processInvokeModelResponse(response, command, agent);
}
else if (exports.ConverseCommand &&
command instanceof exports.ConverseCommand) {
this.processConverseResponse(response, command, agent);
}
})
.catch((error) => {
agent.onErrorThrownByInterceptor({
error: error,
method: "send.<promise>",
module: "@aws-sdk/client-bedrock-runtime",
});
});
}
}
return returnValue;
},
});
}
}
wrap(hooks) {
hooks
.addPackage("@aws-sdk/client-bedrock-runtime")
.withVersion("^3.0.0")
.onRequire((exports, pkgInfo) => {
this.onRequire(exports, pkgInfo);
})
// ESM instrumentation not added yet
// because the package.json "main" field points to CJS build
// and "module" is not supported by Node.js:
// "module": "./dist-es/index.js",
.addFileInstrumentation({
path: "dist-cjs/index.js",
functions: [],
accessLocalVariables: {
names: ["module.exports"],
cb: (vars, pkgInfo) => {
if (vars.length > 0 && (0, isPlainObject_1.isPlainObject)(vars[0])) {
this.onRequire(vars[0], pkgInfo);
}
},
},
});
}
}
exports.AwsSDKVersion3 = AwsSDKVersion3;