Use updateContext

AikidoSec · Jul 25, 2024 · 125c103 · 125c103
1 parent cd327ea
commit 125c103
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 6 deletions.
diff --git a/library/ratelimiting/shouldRateLimitRequest.ts b/library/ratelimiting/shouldRateLimitRequest.ts
@@ -1,5 +1,5 @@
 import { Agent } from "../agent/Agent";
-import { Context } from "../agent/Context";
+import { Context, updateContext } from "../agent/Context";
 import { isLocalhostIP } from "../helpers/isLocalhostIP";
 
 type Result =
@@ -16,7 +16,10 @@ type Result =
     };
 
 // eslint-disable-next-line max-lines-per-function
-export function shouldRateLimitRequest(context: Context, agent: Agent): Result {
+export function shouldRateLimitRequest(
+  context: Readonly<Context>,
+  agent: Agent
+): Result {
   const match = agent.getConfig().getEndpoint(context);
 
   if (!match) {
@@ -61,7 +64,7 @@ export function shouldRateLimitRequest(context: Context, agent: Agent): Result {
 
     // This function is executed for every middleware and route handler
     // We want to count the request only once
-    context.consumedRateLimitForIP = true;
+    updateContext(context, "consumedRateLimitForIP", true);
 
     if (!allowed) {
       return { block: true, trigger: "ip" };
@@ -79,7 +82,7 @@ export function shouldRateLimitRequest(context: Context, agent: Agent): Result {
 
     // This function is executed for every middleware and route handler
     // We want to count the request only once
-    context.consumedRateLimitForUser = true;
+    updateContext(context, "consumedRateLimitForUser", true);
 
     if (!allowed) {
       return { block: true, trigger: "user" };

diff --git a/library/sources/Lambda.test.ts b/library/sources/Lambda.test.ts
@@ -5,7 +5,7 @@ import { Agent } from "../agent/Agent";
 import { setInstance } from "../agent/AgentSingleton";
 import { ReportingAPIForTesting } from "../agent/api/ReportingAPIForTesting";
 import { Token } from "../agent/api/Token";
-import { getContext } from "../agent/Context";
+import { getContext, updateContext } from "../agent/Context";
 import { LoggerNoop } from "../agent/logger/LoggerNoop";
 import { createLambdaWrapper, SQSEvent, APIGatewayProxyEvent } from "./Lambda";
 
@@ -427,7 +427,7 @@ t.test("it counts attacks", async () => {
 
   const handler = createLambdaWrapper(async (event, context) => {
     const ctx = getContext();
-    ctx.attackDetected = true;
+    updateContext(ctx, "attackDetected", true);
     return ctx;
   });