feat: Implement ttl on ArcjetDecision & ArcjetRuleResults (#7)

Co-authored-by: David Mytton <[email protected]>
arcjet · Dec 14, 2023 · c0529ce · c0529ce
1 parent f822c9f
commit c0529ce
Show file tree

Hide file tree

Showing 23 changed files with 524 additions and 27 deletions.
diff --git a/arcjet/index.ts b/arcjet/index.ts
@@ -60,8 +60,8 @@ class Cache<T> {
   }
 
   get(key: string) {
-    const expiresAt = this.expires.get(key);
-    if (typeof expiresAt !== "undefined" && expiresAt > Date.now()) {
+    const ttl = this.ttl(key);
+    if (ttl > 0) {
       return this.data.get(key);
     } else {
       // Cleanup if expired
@@ -74,6 +74,12 @@ class Cache<T> {
     this.expires.set(key, Date.now() + ttl);
     this.data.set(key, value);
   }
+
+  ttl(key: string): number {
+    const now = Date.now();
+    const expiresAt = this.expires.get(key) ?? now;
+    return expiresAt - now;
+  }
 }
 
 function errorMessage(err: unknown): string {
@@ -548,12 +554,14 @@ export function validateEmail(
       ): Promise<ArcjetRuleResult> {
         if (await analyze.isValidEmail(email, analyzeOpts)) {
           return new ArcjetRuleResult({
+            ttl: 0,
             state: "RUN",
             conclusion: "ALLOW",
             reason: new ArcjetEmailReason({ emailTypes: [] }),
           });
         } else {
           return new ArcjetRuleResult({
+            ttl: 0,
             state: "RUN",
             conclusion: "DENY",
             reason: new ArcjetEmailReason({
@@ -640,6 +648,7 @@ export function detectBot(
           block.includes(BotType[botResult.bot_type] as ArcjetBotType)
         ) {
           return new ArcjetRuleResult({
+            ttl: 60000,
             state: "RUN",
             conclusion: "DENY",
             reason: new ArcjetBotReason({
@@ -651,6 +660,7 @@ export function detectBot(
           });
         } else {
           return new ArcjetRuleResult({
+            ttl: 60000,
             state: "RUN",
             conclusion: "ALLOW",
             reason: new ArcjetBotReason({
@@ -763,9 +773,6 @@ export default function arcjet<
   // TODO(#132): Support configurable caching
   const blockCache = new Cache<ArcjetReason>();
 
-  // TTL for in-memory caching of decisions (in milliseconds)
-  const blockCacheTimeoutInMs = 60000;
-
   const flatSortedRules = rules.flat(1).sort((a, b) => a.priority - b.priority);
 
   return Object.freeze({
@@ -802,6 +809,7 @@ export default function arcjet<
         log.error("Failure running rules. Only 10 rules may be specified.");
 
         const decision = new ArcjetErrorDecision({
+          ttl: 0,
           reason: new ArcjetErrorReason("Only 10 rules may be specified"),
           // No results because the sorted rules were too long and we don't want
           // to instantiate a ton of NOT_RUN results
@@ -824,6 +832,7 @@ export default function arcjet<
       // Default all rules to NOT_RUN/ALLOW before doing anything
       for (let idx = 0; idx < flatSortedRules.length; idx++) {
         results[idx] = new ArcjetRuleResult({
+          ttl: 0,
           state: "NOT_RUN",
           conclusion: "ALLOW",
           reason: new ArcjetReason(),
@@ -835,7 +844,7 @@ export default function arcjet<
       // some instances where the instance is not recycled immediately. If so, we
       // can take advantage of that.
       log.time("cache");
-      const existingBlockReason = await blockCache.get(fingerprint);
+      const existingBlockReason = blockCache.get(fingerprint);
       log.timeEnd("cache");
 
       // If already blocked then we can async log to the API and return the
@@ -847,6 +856,7 @@ export default function arcjet<
           existingBlockReason,
         });
         const decision = new ArcjetDenyDecision({
+          ttl: blockCache.ttl(fingerprint),
           reason: existingBlockReason,
           // All results will be NOT_RUN because we used a cached decision
           results,
@@ -903,6 +913,7 @@ export default function arcjet<
           );
 
           results[idx] = new ArcjetRuleResult({
+            ttl: 0,
             state: "RUN",
             conclusion: "ERROR",
             reason: new ArcjetErrorReason(err),
@@ -915,8 +926,9 @@ export default function arcjet<
           log.timeEnd("local");
 
           const decision = new ArcjetDenyDecision({
-            results,
+            ttl: results[idx].ttl,
             reason: results[idx].reason,
+            results,
           });
 
           // Only a DENY decision is reported to avoid creating 2 entries for a
@@ -929,20 +941,22 @@ export default function arcjet<
             flatSortedRules,
           );
 
-          // If we're not in DRY_RUN mode, we want to cache and return this DENY
-          // decision.
+          // If we're not in DRY_RUN mode, we want to cache non-zero TTL results
+          // and return this DENY decision.
           if (rule.mode !== "DRY_RUN") {
-            log.debug("Caching decision", {
-              fingerprint,
-              conclusion: decision.conclusion,
-              reason: decision.reason,
-            });
-
-            await blockCache.set(
-              fingerprint,
-              decision.reason,
-              blockCacheTimeoutInMs,
-            );
+            if (results[idx].ttl > 0) {
+              log.debug("Caching decision for %d milliseconds", decision.ttl, {
+                fingerprint,
+                conclusion: decision.conclusion,
+                reason: decision.reason,
+              });
+
+              blockCache.set(
+                fingerprint,
+                decision.reason,
+                decision.ttl,
+              );
+            }
 
             return decision;
           }
@@ -974,19 +988,21 @@ export default function arcjet<
           fingerprint,
           path: details.path,
           runtime: runtime(),
+          ttl: decision.ttl,
           conclusion: decision.conclusion,
           reason: decision.reason,
           ruleResults: decision.results,
         });
 
-        // If the decision is to block, we cache the block locally
-        if (decision.isDenied()) {
-          log.debug("decide: Caching block locally");
+        // If the decision is to block and we have a non-zero TTL, we cache the
+        // block locally
+        if (decision.isDenied() && decision.ttl > 0) {
+          log.debug("decide: Caching block locally for %d milliseconds", decision.ttl);
 
-          await blockCache.set(
+          blockCache.set(
             fingerprint,
             decision.reason,
-            blockCacheTimeoutInMs,
+            decision.ttl,
           );
         }
 
@@ -997,6 +1013,7 @@ export default function arcjet<
           errorMessage(err),
         );
         const decision = new ArcjetErrorDecision({
+          ttl: 0,
           reason: new ArcjetErrorReason(err),
           results,
         });

diff --git a/arcjet/test/index.edge.test.ts b/arcjet/test/index.edge.test.ts
@@ -18,6 +18,7 @@ describe("Arcjet: Env = Edge runtime", () => {
     const client = {
       decide: jest.fn(async () => {
         return new ArcjetAllowDecision({
+          ttl: 0,
           reason: new ArcjetTestReason(),
           results: [],
         });