Skip to content

AI Agent External Tool-Broker Lab

Exercises action-bound approval, replay resistance, kill-switch behavior, and unknown-argument denial in a local tool broker.

Application Security 3 min read

Implementation: Tested

Implementation

broker.js JavaScript · 518 lines
labs/ai-agent-security/broker.js
Open on GitHub View raw 
1#!/usr/bin/env node
2"use strict";
3 
4/**
5 * Tested teaching implementation for labs/ai-agent-security.
6 *
7 * This broker receives authenticated context separately from the model-proposed
8 * call. It demonstrates deterministic authorization; it does not validate
9 * identity tokens, make network requests, or issue real credentials.
10 *
11 * Approval consumption uses an in-memory atomic compare-and-set that models a
12 * durable conditional write. Production systems must use a durable store such
13 * as a PostgreSQL conditional update / unique insert, Redis SET NX, DynamoDB
14 * conditional write, or another transactionally enforced compare-and-set.
15 * A JavaScript Set is not equivalent to a distributed durable store.
16 */
17 
18class BrokerDenied extends Error {
19  constructor(code, message) {
20    super(message);
21    this.name = "BrokerDenied";
22    this.code = code;
23  }
24}
25 
26function deny(code, message) {
27  throw new BrokerDenied(code, message);
28}
29 
30function isRecord(value) {
31  return value !== null && typeof value === "object" && !Array.isArray(value);
32}
33 
34function requireRecord(value, code, label) {
35  if (!isRecord(value)) deny(code, `${label} must be an object`);
36  return value;
37}
38 
39function requireIdentifier(value, code, label) {
40  if (typeof value !== "string" || value.length < 1 || value.length > 128) {
41    deny(code, `${label} must be a nonempty string of at most 128 characters`);
42  }
43  return value;
44}
45 
46function requireSafePositiveInteger(value, code, label) {
47  if (!Number.isSafeInteger(value) || value <= 0) {
48    deny(code, `${label} must be a positive safe integer`);
49  }
50  return value;
51}
52 
53function hasOwn(object, key) {
54  return Object.prototype.hasOwnProperty.call(object, key);
55}
56 
57const {createHash} = require("node:crypto");
58 
59/**
60 * Canonical action binding for approval consumption.
61 * Key order is fixed. The hash is SHA-256 of JSON.stringify of that object —
62 * not a delimiter-joined string (which would be ambiguous for attacker-controlled
63 * identifier characters).
64 */
65function canonicalActionBinding(binding) {
66  return {
67    principalId: binding.principalId,
68    tenantId: binding.tenantId,
69    tool: binding.tool,
70    amountCents: binding.amountCents,
71    destinationRef: binding.destinationRef
72  };
73}
74 
75function hashActionBinding(binding) {
76  const canonical = JSON.stringify(canonicalActionBinding(binding));
77  return createHash("sha256").update(canonical, "utf8").digest("hex");
78}
79 
80/** @deprecated Use hashActionBinding; retained name for ApprovalStore callers. */
81function digestAction(binding) {
82  return hashActionBinding(binding);
83}
84 
85function safeAuditMetadata(authenticatedContext, proposedCall) {
86  const args = isRecord(proposedCall?.arguments) ? proposedCall.arguments : {};
87  return {
88    principalId:
89      typeof authenticatedContext?.principalId === "string"
90        ? authenticatedContext.principalId
91        : "<invalid>",
92    tenantId:
93      typeof proposedCall?.tenantId === "string"
94        ? proposedCall.tenantId
95        : "<invalid>",
96    tool:
97      typeof proposedCall?.tool === "string" ? proposedCall.tool : "<invalid>",
98    amountCents:
99      Number.isSafeInteger(args.amountCents) ? args.amountCents : null,
100    destinationRef:
101      typeof args.destinationRef === "string"
102        ? args.destinationRef
103        : "<invalid>",
104  };
105}
106 
107/**
108 * In-memory approval store that models durable conditional consumption.
109 *
110 * getApprovedAction is read-only. consumeIfUnused serializes consumption so
111 * only one concurrent caller can observe an unused approval and mark it
112 * consumed. Consumed state is not reverted if a later executor fails;
113 * production systems should reconcile by idempotency key before retrying.
114 */
115class ApprovalStore {
116  #approvals = new Map();
117  #consumed = new Map();
118  #mutex = Promise.resolve();
119  #readDelayMs;
120 
121  constructor(approvals = [], {readDelayMs = 0} = {}) {
122    this.#readDelayMs = readDelayMs;
123    for (const approval of approvals) {
124      if (!isRecord(approval) || typeof approval.id !== "string") {
125        throw new TypeError("each approval must be an object with an id");
126      }
127      this.#approvals.set(approval.id, Object.freeze({...approval}));
128    }
129  }
130 
131  async getApprovedAction(approvalId) {
132    if (this.#readDelayMs > 0) {
133      await new Promise((resolve) => setTimeout(resolve, this.#readDelayMs));
134    }
135    return this.#approvals.get(approvalId) ?? null;
136  }
137 
138  /**
139   * Atomically consume an approval if it is unused and still bound to the
140   * supplied action digest. Returns true only for the first valid consumption.
141   */
142  async consumeIfUnused({approvalId, actionDigest, nowMs}) {
143    const run = this.#mutex.then(() => {
144      if (typeof approvalId !== "string" || approvalId.length < 1) {
145        return false;
146      }
147      if (typeof actionDigest !== "string" || actionDigest.length < 1) {
148        return false;
149      }
150      if (!Number.isSafeInteger(nowMs)) {
151        return false;
152      }
153      if (this.#consumed.has(approvalId)) {
154        return false;
155      }
156 
157      const approval = this.#approvals.get(approvalId);
158      if (!isRecord(approval) || approval.decision !== "approved") {
159        return false;
160      }
161      if (!Number.isSafeInteger(approval.expiresAtMs) || approval.expiresAtMs <= nowMs) {
162        return false;
163      }
164 
165      const expectedDigest = digestAction({
166        principalId: approval.principalId,
167        tenantId: approval.tenantId,
168        tool: approval.tool,
169        amountCents: approval.amountCents,
170        destinationRef: approval.destinationRef,
171      });
172      if (expectedDigest !== actionDigest) {
173        return false;
174      }
175 
176      // Durable CAS model: mark consumed before returning success. Executor
177      // failure must not free the approval for a second attempt.
178      this.#consumed.set(approvalId, {
179        actionDigest,
180        consumedAtMs: nowMs,
181      });
182      return true;
183    });
184 
185    this.#mutex = run.then(
186      () => undefined,
187      () => undefined,
188    );
189    return run;
190  }
191}
192 
193class ToolBroker {
194  #approvalStore;
195  #audit;
196  #clock;
197  #executor;
198  #isExecutionEnabled;
199  #policy;
200 
201  constructor({
202    policy,
203    approvalStore,
204    executor,
205    isExecutionEnabled = () => true,
206    audit = () => {},
207    clock = () => Date.now(),
208  }) {
209    if (!isRecord(policy) || !isRecord(policy.principals)) {
210      throw new TypeError("policy.principals must be an object");
211    }
212    if (
213      approvalStore == null ||
214      typeof approvalStore.getApprovedAction !== "function" ||
215      typeof approvalStore.consumeIfUnused !== "function"
216    ) {
217      throw new TypeError(
218        "approvalStore must implement getApprovedAction and consumeIfUnused",
219      );
220    }
221    if (typeof executor !== "function") {
222      throw new TypeError("executor must be a function");
223    }
224    if (typeof isExecutionEnabled !== "function") {
225      throw new TypeError("isExecutionEnabled must be a function");
226    }
227    if (typeof audit !== "function") {
228      throw new TypeError("audit must be a function");
229    }
230    if (typeof clock !== "function") {
231      throw new TypeError("clock must be a function");
232    }
233 
234    this.#policy = policy;
235    this.#approvalStore = approvalStore;
236    this.#executor = executor;
237    this.#isExecutionEnabled = isExecutionEnabled;
238    this.#audit = audit;
239    this.#clock = clock;
240  }
241 
242  async execute(authenticatedContext, proposedCall, approvalEvidence = null) {
243    const auditMetadata = safeAuditMetadata(authenticatedContext, proposedCall);
244 
245    try {
246      const authorizedCall = await this.#authorize(
247        authenticatedContext,
248        proposedCall,
249        approvalEvidence,
250      );
251 
252      // A production high-impact path should fail closed if this durable audit
253      // decision cannot be recorded. This lab's injected audit sink is in-memory.
254      this.#audit({
255        event: "authorization_decision",
256        outcome: "allow",
257        reason: "POLICY_ALLOWED",
258        ...auditMetadata,
259        approvalId: authorizedCall.approvalId,
260      });
261 
262      try {
263        const result = await this.#executor(authorizedCall);
264        this.#audit({
265          event: "execution_outcome",
266          outcome: "success",
267          ...auditMetadata,
268          approvalId: authorizedCall.approvalId,
269        });
270        return result;
271      } catch (error) {
272        this.#audit({
273          event: "execution_outcome",
274          outcome: "failure",
275          reason: "EXECUTOR_FAILED",
276          ...auditMetadata,
277          approvalId: authorizedCall.approvalId,
278        });
279        throw error;
280      }
281    } catch (error) {
282      if (error instanceof BrokerDenied) {
283        this.#audit({
284          event: "authorization_decision",
285          outcome: "deny",
286          reason: error.code,
287          ...auditMetadata,
288          approvalId:
289            isRecord(approvalEvidence) &&
290            typeof approvalEvidence.id === "string"
291              ? approvalEvidence.id
292              : null,
293        });
294      }
295      throw error;
296    }
297  }
298 
299  async #authorize(authenticatedContext, proposedCall, approvalEvidence) {
300    const context = requireRecord(
301      authenticatedContext,
302      "INVALID_CONTEXT",
303      "authenticatedContext",
304    );
305    const principalId = requireIdentifier(
306      context.principalId,
307      "INVALID_CONTEXT",
308      "authenticatedContext.principalId",
309    );
310 
311    const call = requireRecord(
312      proposedCall,
313      "INVALID_CALL",
314      "proposedCall",
315    );
316    const tenantId = requireIdentifier(
317      call.tenantId,
318      "INVALID_CALL",
319      "proposedCall.tenantId",
320    );
321    const tool = requireIdentifier(
322      call.tool,
323      "INVALID_CALL",
324      "proposedCall.tool",
325    );
326    const args = requireRecord(
327      call.arguments,
328      "INVALID_ARGUMENTS",
329      "proposedCall.arguments",
330    );
331 
332    const allowedArgumentNames = new Set([
333      "amountCents",
334      "destinationRef",
335    ]);
336    for (const argumentName of Object.keys(args)) {
337      if (!allowedArgumentNames.has(argumentName)) {
338        deny(
339          "UNKNOWN_ARGUMENT",
340          `argument ${argumentName} is not in the closed schema`,
341        );
342      }
343    }
344 
345    const amountCents = requireSafePositiveInteger(
346      args.amountCents,
347      "INVALID_AMOUNT",
348      "proposedCall.arguments.amountCents",
349    );
350    const destinationRef = requireIdentifier(
351      args.destinationRef,
352      "INVALID_DESTINATION",
353      "proposedCall.arguments.destinationRef",
354    );
355 
356    const executionEnabled = await this.#isExecutionEnabled({
357      principalId,
358      tenantId,
359      tool,
360    });
361    if (executionEnabled !== true) {
362      deny("EXECUTION_DISABLED", "the execution kill switch is active");
363    }
364 
365    if (!hasOwn(this.#policy.principals, principalId)) {
366      deny("UNAUTHORIZED_PRINCIPAL", "principal is not registered");
367    }
368    const principalPolicy = requireRecord(
369      this.#policy.principals[principalId],
370      "INVALID_POLICY",
371      "principal policy",
372    );
373 
374    if (
375      !Array.isArray(principalPolicy.tenants) ||
376      !principalPolicy.tenants.includes(tenantId)
377    ) {
378      deny("UNAUTHORIZED_TENANT", "principal cannot act for this tenant");
379    }
380 
381    if (!isRecord(principalPolicy.tools) || !hasOwn(principalPolicy.tools, tool)) {
382      deny("UNAUTHORIZED_TOOL", "tool is not allowed for this principal");
383    }
384    const toolPolicy = requireRecord(
385      principalPolicy.tools[tool],
386      "INVALID_POLICY",
387      "tool policy",
388    );
389 
390    const maxAmountCents = requireSafePositiveInteger(
391      toolPolicy.maxAmountCents,
392      "INVALID_POLICY",
393      "tool policy maxAmountCents",
394    );
395    const approvalRequiredAtCents = requireSafePositiveInteger(
396      toolPolicy.approvalRequiredAtCents,
397      "INVALID_POLICY",
398      "tool policy approvalRequiredAtCents",
399    );
400    if (approvalRequiredAtCents > maxAmountCents) {
401      deny(
402        "INVALID_POLICY",
403        "approval threshold cannot exceed the maximum amount",
404      );
405    }
406 
407    if (amountCents > maxAmountCents) {
408      deny("AMOUNT_EXCEEDS_POLICY", "amount exceeds the policy maximum");
409    }
410 
411    if (
412      !Array.isArray(toolPolicy.allowedDestinationRefs) ||
413      !toolPolicy.allowedDestinationRefs.includes(destinationRef)
414    ) {
415      deny(
416        "UNAUTHORIZED_DESTINATION",
417        "destination is not allowed by policy",
418      );
419    }
420 
421    let approvalId = null;
422    if (amountCents >= approvalRequiredAtCents) {
423      const evidence = requireRecord(
424        approvalEvidence,
425        "APPROVAL_REQUIRED",
426        "approvalEvidence",
427      );
428      approvalId = requireIdentifier(
429        evidence.id,
430        "INVALID_APPROVAL",
431        "approvalEvidence.id",
432      );
433 
434      const actionDigest = digestAction({
435        principalId,
436        tenantId,
437        tool,
438        amountCents,
439        destinationRef,
440      });
441      const nowMs = this.#clock();
442 
443      // Verification (read) and consumption (atomic write) are separated.
444      let approval;
445      try {
446        approval = await this.#approvalStore.getApprovedAction(approvalId);
447      } catch {
448        deny("APPROVAL_STORE_ERROR", "approval store failed closed");
449      }
450 
451      if (!isRecord(approval) || approval.decision !== "approved") {
452        deny("INVALID_APPROVAL", "approval service did not approve the action");
453      }
454      if (approval.id !== approvalId) {
455        deny("INVALID_APPROVAL", "approval identifier does not match");
456      }
457      if (
458        !Number.isSafeInteger(approval.expiresAtMs) ||
459        approval.expiresAtMs <= nowMs
460      ) {
461        deny("APPROVAL_EXPIRED", "approval has expired");
462      }
463 
464      const expectedBinding = {
465        principalId,
466        tenantId,
467        tool,
468        amountCents,
469        destinationRef,
470      };
471      for (const [field, expectedValue] of Object.entries(expectedBinding)) {
472        if (approval[field] !== expectedValue) {
473          deny(
474            "APPROVAL_SCOPE_MISMATCH",
475            `approval is not bound to the requested ${field}`,
476          );
477        }
478      }
479 
480      // Consume before invoking the executor. Only one concurrent request may
481      // succeed; the loser receives APPROVAL_REPLAYED. Consumed state is not
482      // reverted if the executor later fails.
483      let consumed;
484      try {
485        consumed = await this.#approvalStore.consumeIfUnused({
486          approvalId,
487          actionDigest,
488          nowMs,
489        });
490      } catch {
491        deny("APPROVAL_STORE_ERROR", "approval store failed closed");
492      }
493      if (consumed !== true) {
494        deny("APPROVAL_REPLAYED", "approval has already been consumed");
495      }
496    }
497 
498    return Object.freeze({
499      principalId,
500      tenantId,
501      tool,
502      arguments: Object.freeze({
503        amountCents,
504        destinationRef,
505      }),
506      approvalId,
507    });
508  }
509}
510 
511module.exports = {
512  ApprovalStore,
513  BrokerDenied,
514  ToolBroker,
515  canonicalActionBinding,
516  digestAction,
517  hashActionBinding,
518};

Run it

  • node --test labs/ai-agent-security/tests/broker.test.js

Status: tested teaching implementation. The surrounding production architecture in appsec/ai-agent-security.md is conceptual and partially tested.

This dependency-free Node.js lab demonstrates one narrow security boundary: a model may propose an action, but a broker outside the model decides whether a fake executor receives it. The executor stores an in-memory record and returns status: simulated; it performs no network, filesystem, financial, cloud, or other external action.

Security properties exercised

The broker receives authenticatedContext separately from proposedCall and then:

  • resolves a static server-side policy for the authenticated principal;
  • rejects an unauthorized principal, tenant, tool, destination, or amount;
  • accepts only a closed argument schema;
  • requires approval at or above the configured impact threshold;
  • asks an injected approval verifier for the authoritative approval record;
  • binds approval to principal, tenant, tool, amount, and destination;
  • consumes approval before execution to reject replay;
  • checks an independently injected execution kill switch;
  • emits sanitized decision and outcome events to an in-memory audit sink;
  • invokes a fake executor only after all controls allow the call.

The tests assert that the fake executor's invocation count stays zero for denied actions. This distinction matters: a model may still propose a forbidden action, but it must not acquire authority by doing so.

Run

From the repository root with Node.js 22 or newer:

node --test labs/ai-agent-security/tests/broker.test.js

No package installation, environment variable, credential, service, container, or network connection is required.

Files

  • broker.js contains the broker and stable denial codes.
  • tests/broker.test.js has tests for both what should work and what should get denied, plus the fake policy, approval store, audit sink, kill switch, and executor.

Deliberate limitations

This is not a production authorization library. In particular:

  • authenticatedContext is assumed to come from a trusted gateway. The lab does not validate a session, workload certificate, issuer, signature, audience, freshness, revocation, or tenant membership source.
  • The policy is an in-process fixture. It has no administrative authorization, durable version, distribution, cache, or outage behavior.
  • The approval verifier reads an in-memory map. It does not authenticate an approver, verify a signature, enforce separation of duties, or persist single-use state across processes.
  • Replay state, audit events, and fake execution records are memory-only.
  • The broker validates one teaching schema for a simulated payment-like action. A real broker needs a reviewed, versioned schema and constraints per tool.
  • The lab does not implement MCP, OAuth, token exchange, downstream reauthorization, idempotency reconciliation, a sandbox, secret delivery, egress control, rate limiting, or durable incident evidence.

For production, derive identity and tenant outside the model, use a maintained authorization implementation, bind short-lived credentials to the selected resource, make the downstream service authorize again, and test fail-closed behavior at every boundary.

Research AI-agent authorization External authorization, action-bound approval, concurrent local consumption, and bounded tool execution.