aios/plugins/aios-tenant-boundary/index.js

import crypto from "node:crypto";
import path from "node:path";
import { canonicalize, loadStoredRecord, readPublicKey, validatePayload, verifyRecord } from "./provenance.js";
import { classifyRun, decideToolCall, runKey, validateWriterParams } from "./policy.js";
import { executeEvidenceWrite } from "./writer.js";
import { createRuntimeApi } from "./runtime.js";

const activeRuns = new Map();
const permits = new Map();
const MAX_RUN_AGE_MS = 60 * 60 * 1000;
const DEFAULT_PERMIT_TTL_MS = 2 * 60 * 1000;

function pluginConfig(event, api) {
  return event?.context?.pluginConfig ?? api.pluginConfig ?? {};
}

function prune(now = Date.now()) {
  for (const [key, value] of activeRuns) if (now - value.recordedAt > MAX_RUN_AGE_MS) activeRuns.delete(key);
  for (const [key, value] of permits) if (now - value.recordedAt > value.ttlMs) permits.delete(key);
}

function toolCallId(event) {
  return typeof event?.toolCallId === "string" && event.toolCallId.trim() ? event.toolCallId.trim() : undefined;
}

function eventParams(event) {
  return event?.params ?? event?.parameters;
}

function requireAbsoluteConfig(config, key) {
  const value = config[key];
  if (typeof value !== "string" || !path.isAbsolute(value)) throw new Error(`config_${key}_missing`);
  return value;
}

export function authorizeEvidenceCall({ event, trust, config, now = Date.now() }) {
  const callId = toolCallId(event);
  if (!callId) return { ok: false, code: "tool_call_id_missing" };
  if (!trust?.trusted) return { ok: false, code: trust?.reason || "run_provenance_missing" };
  const params = eventParams(event);
  if (!params || typeof params !== "object" || Array.isArray(params)) return { ok: false, code: "tool_params_missing" };
  let record;
  try { record = loadStoredRecord(requireAbsoluteConfig(config, "provenanceStoreDir"), params.provenance_id); }
  catch (error) {
    const code = error instanceof Error && ["unsafe_provenance_file", "invalid_provenance_location", "provenance_record_too_large"].includes(error.message)
      ? error.message : "provenance_store_rejected";
    return { ok: false, code };
  }
  const validation = validatePayload(record.payload, now);
  if (!validation.ok) return validation;
  let publicKey;
  try { publicKey = readPublicKey(requireAbsoluteConfig(config, "provenancePublicKeyPath")); }
  catch { return { ok: false, code: "public_key_rejected" }; }
  const signature = verifyRecord(record, publicKey, config.provenanceKeyId);
  if (!signature.ok) return signature;
  const payload = record.payload;
  const runId = runKey(event, {});
  if (!runId || payload.run_id !== runId) return { ok: false, code: "run_id_mismatch" };
  if (payload.sender_id !== trust.senderId) return { ok: false, code: "sender_id_mismatch" };
  if (config.writerTaskBranch && payload.task_branch !== config.writerTaskBranch) return { ok: false, code: "task_branch_mismatch" };
  if (config.writerTaskId && payload.task_id !== config.writerTaskId) return { ok: false, code: "task_id_mismatch" };
  if (config.provenanceIssuer && payload.issued_by !== config.provenanceIssuer) return { ok: false, code: "issuer_mismatch" };
  if (config.executor && payload.executor !== config.executor) return { ok: false, code: "executor_mismatch" };
  if (config.executorInstance && payload.executor_instance !== config.executorInstance) return { ok: false, code: "executor_instance_mismatch" };
  const paramsValidation = validateWriterParams(params, payload);
  if (!paramsValidation.ok) return paramsValidation;
  const digest = crypto.createHash("sha256").update(canonicalize(params)).digest("hex");
  if (permits.has(callId)) return { ok: false, code: "permit_replay" };
  const ttlMs = Number.isSafeInteger(config.permitTtlMs) && config.permitTtlMs > 0
    ? Math.min(config.permitTtlMs, DEFAULT_PERMIT_TTL_MS)
    : DEFAULT_PERMIT_TTL_MS;
  permits.set(callId, { payload, paramsDigest: digest, recordedAt: now, ttlMs });
  return { ok: true };
}

export function consumePermit(toolCallIdValue, params, now = Date.now()) {
  const permit = permits.get(toolCallIdValue);
  permits.delete(toolCallIdValue);
  if (!permit) throw new Error("permit_missing");
  if (now - permit.recordedAt > permit.ttlMs) throw new Error("permit_expired");
  const digest = crypto.createHash("sha256").update(canonicalize(params)).digest("hex");
  if (!crypto.timingSafeEqual(Buffer.from(digest), Buffer.from(permit.paramsDigest))) throw new Error("permit_params_mismatch");
  return permit;
}

function toolResult(details) {
  return { content: [{ type: "text", text: JSON.stringify(details) }], details };
}

function registerWriterTool(api) {
  api.registerTool((toolContext) => ({
    name: "aios_evidence_write",
    label: "AIOS Evidence Writer",
    description: "Write one Owner-authorized JSON evidence file and normally push the current task branch.",
    parameters: {
      type: "object",
      additionalProperties: false,
      required: ["provenance_id", "path", "operation", "document"],
      properties: {
        provenance_id: { type: "string" },
        path: { type: "string" },
        operation: { type: "string", enum: ["create", "update"] },
        document: { type: "object" },
      },
    },
    executionMode: "sequential",
    async execute(callId, params) {
      const permit = consumePermit(callId, params);
      const config = toolContext?.getRuntimeConfig?.()?.plugins?.entries?.["aios-tenant-boundary"]?.config
        ?? api.pluginConfig ?? {};
      const result = await executeEvidenceWrite({
        repoRoot: requireAbsoluteConfig(config, "writerRepoRoot"),
        evidencePath: params.path,
        operation: params.operation,
        document: params.document,
        payload: permit.payload,
        pushRemote: config.writerPushRemote ?? "gitea-http",
        askPassPath: config.writerAskPassPath ?? "",
        auditLogPath: requireAbsoluteConfig(config, "writerAuditLogPath"),
        gitBinary: config.gitBinary ?? "/usr/bin/git",
      });
      return toolResult(result);
    },
  }), { name: "aios_evidence_write" });
}

export default {
  id: "aios-tenant-boundary",
  name: "AIOS Tenant Boundary",
  description: "Fail-closed Owner provenance and dedicated evidence-writer boundary.",
  register(api) {
    const runtimeApi = createRuntimeApi(api);
    api.registerHttpRoute({ path: "/v1/health", auth: "plugin", match: "exact", handler: runtimeApi.handleHealth });
    api.registerHttpRoute({ path: "/v1/execute", auth: "plugin", match: "exact", handler: runtimeApi.handleExecute });
    api.registerHttpRoute({ path: "/v1/commands/", auth: "plugin", match: "prefix", handler: runtimeApi.commandRequest });
    registerWriterTool(api);
    api.on("before_agent_run", async (event, ctx) => {
      prune();
      const config = pluginConfig(event, api);
      const trust = classifyRun(event, ctx, config);
      const key = runKey(event, ctx);
      if (key) activeRuns.set(key, { ...trust, recordedAt: Date.now() });
      if (config.blockUnknownAgentRuns === true && !trust.trusted && trust.reason !== "job_id_present") {
        return { outcome: "block", reason: trust.reason, message: "Authenticated Owner provenance is required." };
      }
      return { outcome: "pass" };
    }, { priority: 100 });

    api.on("before_tool_call", async (event, ctx) => {
      prune();
      const config = pluginConfig(event, api);
      const trust = activeRuns.get(runKey(event, ctx));
      const decision = decideToolCall(event.toolName, trust, config);
      if (decision.action === "block") return { block: true, blockReason: `AIOS V2R1 policy denied: ${decision.code}` };
      if (decision.action === "verify_provenance") {
        const authorized = authorizeEvidenceCall({ event, trust, config });
        if (!authorized.ok) return { block: true, blockReason: `AIOS V2R1 provenance denied: ${authorized.code}` };
        return undefined;
      }
      if (decision.action === "approval") {
        return { requireApproval: {
          title: `Approve ${String(event.toolName)}`,
          description: "Authenticated Owner requested a non-evidence side effect.",
          severity: "warning", timeoutMs: Number(config.approvalTimeoutMs ?? 60000),
          allowedDecisions: ["allow-once", "deny"], pluginId: "aios-tenant-boundary",
        } };
      }
      return undefined;
    }, { priority: 100 });

    api.on("agent_end", async (event, ctx) => {
      const key = runKey(event, ctx);
      if (key) activeRuns.delete(key);
    });
  },
};

export const __state = { activeRuns, permits };