@absol-labs/agent 0.6.0 → 0.7.0

package/src/wallet/privy-session-broker.ts

@@ -0,0 +1,634 @@
+import { createHash, randomBytes, timingSafeEqual } from "node:crypto";
+
+import { streamEscrowV2Abi } from "@absol-labs/shared";
+import {
+  decodeFunctionData,
+  getAddress,
+  isAddress,
+  type Address,
+  type Hex,
+} from "viem";
+import { z } from "zod";
+
+import {
+  checkMandate,
+  signedSpendMandateSchema,
+  verifyMandateOwnerSignature,
+  verifySignedMandate,
+  type SignedSpendMandate,
+} from "../mandates/mandate.js";
+import {
+  invocationCapabilityDomainName,
+  invocationCapabilityDomainVersion,
+  invocationCapabilitySchema,
+} from "../capability/invocation-capability.js";
+
+export const PRIVY_BROKER_CHAIN_ID = 84532;
+export const PRIVY_BROKER_CAIP2 = "eip155:84532" as const;
+export const PRIVY_BROKER_ESCROW =
+  "0x0f09f36Ccc05A7c9882F438721C08De314dFd46C" as const;
+export const PRIVY_BROKER_USDC =
+  "0x036CbD53842c5426634e7929541eC2318f3dCF7e" as const;
+export const PRIVY_POLICY_MAX_TOTAL_USDC = 1_000_000n;
+export const PRIVY_POLICY_MAX_RATE_USDC = 1_000n;
+export const PRIVY_POLICY_MAX_DURATION_SECONDS = 3_600;
+export const PRIVY_MAX_SESSION_SECONDS = 86_400;
+
+const wireMandateSchema = z.object({
+  mandateId: z.string(),
+  owner: z.string(),
+  chainId: z.number().int(),
+  issuedAt: z.number().int(),
+  mandate: z.object({
+    maxPerStreamUsdc: z.coerce.bigint(),
+    maxTotalUsdc: z.coerce.bigint(),
+    maxRatePerSecondUsdc: z.coerce.bigint(),
+    maxDurationSeconds: z.number().int(),
+    allowedOperators: z.array(z.string()).optional(),
+    expiresAt: z.number().int(),
+  }),
+  signature: z.string(),
+});
+
+const transactionSchema = z.object({
+  from: z.string().refine(isAddress),
+  to: z.string().refine(isAddress),
+  data: z.string().regex(/^0x[0-9a-fA-F]*$/),
+  value: z.union([z.string(), z.number(), z.bigint()]).optional(),
+  chainId: z.union([z.string(), z.number()]).optional(),
+});
+
+export interface PrivyLinkedWallet {
+  readonly id: string | null;
+  readonly address: string;
+  readonly chain_type: string;
+  readonly connector_type?: string;
+  readonly wallet_client_type?: string;
+}
+
+export interface PrivyBrokerProvider {
+  verifyAccessToken(token: string): Promise<{ readonly userId: string }>;
+  getUserWallets(userId: string): Promise<readonly PrivyLinkedWallet[]>;
+  sendTransaction(input: {
+    readonly walletId: string;
+    readonly transaction: Record<string, unknown>;
+    readonly authorizationPrivateKey: string;
+    readonly sponsor: boolean;
+  }): Promise<Hex>;
+  signTypedData(input: {
+    readonly walletId: string;
+    readonly typedData: Record<string, unknown>;
+    readonly authorizationPrivateKey: string;
+  }): Promise<Hex>;
+}
+
+export interface BrokerStreamReader {
+  getStream(streamId: Hex): Promise<{
+    readonly buyer: Address;
+    readonly serviceRef: Hex;
+    readonly status: "active" | "closed";
+  }>;
+}
+
+export interface PrivyBrokerSession {
+  readonly id: string;
+  readonly tokenHash: string;
+  readonly userId: string;
+  readonly walletId: string;
+  readonly address: Address;
+  readonly signedMandate: SignedSpendMandate;
+  readonly expiresAt: number;
+  readonly createdAt: number;
+  readonly revokedAt?: number;
+  readonly spentSoFarUsdc: bigint;
+  /** True after the broker has requested gas sponsorship for one write. */
+  readonly sponsorshipUsed?: boolean;
+}
+
+export interface PrivyBrokerSessionStore {
+  findByTokenHash(tokenHash: string): Promise<PrivyBrokerSession | null>;
+  get(id: string): Promise<PrivyBrokerSession | null>;
+  put(session: PrivyBrokerSession): Promise<void>;
+}
+
+export interface CreatePrivyBrokerSessionInput {
+  readonly accessToken: string;
+  readonly walletId: string;
+  readonly address: Address;
+  readonly signedMandate: unknown;
+}
+
+export interface PrivyBrokerRpcInput {
+  readonly token: string;
+  readonly method: "eth_sendTransaction" | "eth_signTypedData_v4";
+  readonly params: readonly unknown[];
+}
+
+export class PrivySessionBrokerError extends Error {
+  constructor(
+    readonly code:
+      | "unauthorized"
+      | "invalid-request"
+      | "expired"
+      | "revoked"
+      | "policy-denied"
+      | "provider-error",
+    message: string,
+    options?: { readonly cause?: unknown },
+  ) {
+    super(message, options);
+    this.name = "PrivySessionBrokerError";
+  }
+}
+
+export class PrivySessionBroker {
+  private readonly locks = new Map<string, Promise<void>>();
+
+  constructor(
+    private readonly dependencies: {
+      readonly provider: PrivyBrokerProvider;
+      readonly store: PrivyBrokerSessionStore;
+      readonly streamReader: BrokerStreamReader;
+      readonly authorizationPrivateKey: string;
+      readonly nowSeconds?: () => number;
+    },
+  ) {}
+
+  async createSession(input: CreatePrivyBrokerSessionInput): Promise<{
+    readonly sessionId: string;
+    readonly token: string;
+    readonly address: Address;
+    readonly expiresAt: number;
+  }> {
+    const now = this.now();
+    const auth = await this.dependencies.provider
+      .verifyAccessToken(input.accessToken)
+      .catch((error: unknown) => {
+        throw new PrivySessionBrokerError(
+          "unauthorized",
+          "Privy access token is invalid",
+          { cause: error },
+        );
+      });
+    const address = requireAddress(input.address);
+    const wallets = await this.dependencies.provider.getUserWallets(
+      auth.userId,
+    );
+    const wallet = wallets.find(
+      (candidate) =>
+        candidate.id === input.walletId &&
+        candidate.address.toLowerCase() === address.toLowerCase() &&
+        candidate.chain_type === "ethereum" &&
+        candidate.connector_type === "embedded" &&
+        candidate.wallet_client_type === "privy",
+    );
+    if (!wallet?.id) {
+      throw new PrivySessionBrokerError(
+        "unauthorized",
+        "wallet is not a user-owned Privy embedded Ethereum wallet",
+      );
+    }
+
+    const signedMandate = parseWireMandate(input.signedMandate);
+    if (
+      signedMandate.owner.toLowerCase() !== address.toLowerCase() ||
+      signedMandate.chainId !== PRIVY_BROKER_CHAIN_ID
+    ) {
+      throw new PrivySessionBrokerError(
+        "policy-denied",
+        "mandate owner or chain does not match the embedded wallet",
+      );
+    }
+    const verified = await verifySignedMandate(signedMandate, {
+      nowSeconds: now,
+    });
+    if (!verified.allowed) {
+      throw new PrivySessionBrokerError(
+        "policy-denied",
+        `mandate denied: ${verified.reason}`,
+      );
+    }
+    assertMandateWithinPrivyPolicy(signedMandate);
+
+    const expiresAt = Math.min(
+      signedMandate.mandate.expiresAt,
+      now + PRIVY_MAX_SESSION_SECONDS,
+    );
+    const token = randomBytes(32).toString("base64url");
+    const session: PrivyBrokerSession = {
+      id: randomBytes(16).toString("hex"),
+      tokenHash: hashToken(token),
+      userId: auth.userId,
+      walletId: wallet.id,
+      address,
+      signedMandate,
+      expiresAt,
+      createdAt: now,
+      spentSoFarUsdc: 0n,
+    };
+    await this.dependencies.store.put(session);
+    return { sessionId: session.id, token, address, expiresAt };
+  }
+
+  async revokeSession(accessToken: string, sessionId: string): Promise<void> {
+    const auth = await this.dependencies.provider
+      .verifyAccessToken(accessToken)
+      .catch(() => {
+        throw new PrivySessionBrokerError(
+          "unauthorized",
+          "Privy access token is invalid",
+        );
+      });
+    const session = await this.dependencies.store.get(sessionId);
+    if (session === null || session.userId !== auth.userId) {
+      throw new PrivySessionBrokerError("unauthorized", "session not found");
+    }
+    await this.dependencies.store.put({
+      ...session,
+      revokedAt: this.now(),
+    });
+  }
+
+  async rpc(input: PrivyBrokerRpcInput): Promise<Hex> {
+    const session = await this.requireSession(input.token);
+    return this.withSessionLock(session.id, async () => {
+      const current = await this.requireSession(input.token);
+      if (input.method === "eth_sendTransaction") {
+        return this.sendTransaction(current, input.params);
+      }
+      return this.signTypedData(current, input.params);
+    });
+  }
+
+  private async sendTransaction(
+    session: PrivyBrokerSession,
+    params: readonly unknown[],
+  ): Promise<Hex> {
+    const transaction = transactionSchema.parse(params[0]);
+    if (transaction.from.toLowerCase() !== session.address.toLowerCase()) {
+      throw new PrivySessionBrokerError(
+        "policy-denied",
+        "transaction sender does not match session wallet",
+      );
+    }
+    assertChainId(transaction.chainId);
+    assertZeroValue(transaction.value);
+    const to = getAddress(transaction.to);
+    const data = transaction.data as Hex;
+
+    if (to.toLowerCase() === PRIVY_BROKER_USDC.toLowerCase()) {
+      assertBoundedApproval(data, session.signedMandate);
+    } else if (to.toLowerCase() === PRIVY_BROKER_ESCROW.toLowerCase()) {
+      const decoded = decodeFunctionData({ abi: streamEscrowV2Abi, data });
+      if (decoded.functionName === "openStream") {
+        const [operator, , deposit, rate, duration] = decoded.args;
+        const decision = await checkMandate(
+          session.signedMandate,
+          {
+            operator,
+            budgetUsdc: deposit,
+            ratePerSecondUsdc: rate,
+            maxDurationSeconds: Number(duration),
+          },
+          session.spentSoFarUsdc,
+          { nowSeconds: this.now() },
+        );
+        if (!decision.allowed) {
+          throw new PrivySessionBrokerError(
+            "policy-denied",
+            `mandate denied: ${decision.reason}`,
+          );
+        }
+        const hash = await this.providerSend(session, transaction);
+        await this.dependencies.store.put({
+          ...session,
+          sponsorshipUsed: true,
+          spentSoFarUsdc: session.spentSoFarUsdc + deposit,
+        });
+        return hash;
+      }
+      if (
+        decoded.functionName !== "close" &&
+        decoded.functionName !== "reclaim" &&
+        decoded.functionName !== "reclaimUnverified"
+      ) {
+        throw new PrivySessionBrokerError(
+          "policy-denied",
+          "escrow method is not allowed for an agent session",
+        );
+      }
+      const owner = await verifyMandateOwnerSignature(session.signedMandate);
+      if (!owner.allowed) {
+        throw new PrivySessionBrokerError(
+          "policy-denied",
+          "mandate owner signature is invalid",
+        );
+      }
+      const streamId = decoded.args[0];
+      const stream = await this.dependencies.streamReader.getStream(streamId);
+      if (stream.buyer.toLowerCase() !== session.address.toLowerCase()) {
+        throw new PrivySessionBrokerError(
+          "policy-denied",
+          "session wallet is not the stream buyer",
+        );
+      }
+    } else {
+      throw new PrivySessionBrokerError(
+        "policy-denied",
+        "transaction target is not allowlisted",
+      );
+    }
+    return this.providerSend(session, transaction);
+  }
+
+  private async signTypedData(
+    session: PrivyBrokerSession,
+    params: readonly unknown[],
+  ): Promise<Hex> {
+    const raw = params[1];
+    const typedData =
+      typeof raw === "string"
+        ? (JSON.parse(raw) as Record<string, unknown>)
+        : raw;
+    if (typedData === null || typeof typedData !== "object") {
+      throw new PrivySessionBrokerError(
+        "invalid-request",
+        "typed data must be an object",
+      );
+    }
+    const typedDataRecord = typedData as Record<string, unknown>;
+    if (
+      typeof params[0] !== "string" ||
+      params[0].toLowerCase() !== session.address.toLowerCase()
+    ) {
+      throw new PrivySessionBrokerError(
+        "policy-denied",
+        "typed-data signer does not match the session wallet",
+      );
+    }
+    const domain = typedDataRecord.domain as
+      | Record<string, unknown>
+      | undefined;
+    if (
+      domain?.name !== invocationCapabilityDomainName ||
+      domain.version !== invocationCapabilityDomainVersion ||
+      Number(domain.chainId) !== PRIVY_BROKER_CHAIN_ID ||
+      typeof domain.verifyingContract !== "string" ||
+      domain.verifyingContract.toLowerCase() !==
+        PRIVY_BROKER_ESCROW.toLowerCase() ||
+      typedDataRecord.primaryType !== "InvocationCapability"
+    ) {
+      throw new PrivySessionBrokerError(
+        "policy-denied",
+        "only Metrik invocation capabilities may be signed",
+      );
+    }
+    const rawMessage = typedDataRecord.message as
+      | Record<string, unknown>
+      | undefined;
+    const capability = invocationCapabilitySchema.parse({
+      ...rawMessage,
+      expiry:
+        typeof rawMessage?.expiry === "string"
+          ? Number(rawMessage.expiry)
+          : rawMessage?.expiry,
+    });
+    if (
+      capability.buyer.toLowerCase() !== session.address.toLowerCase() ||
+      capability.expiry <= this.now() ||
+      capability.expiry > Math.min(session.expiresAt, this.now() + 120)
+    ) {
+      throw new PrivySessionBrokerError(
+        "policy-denied",
+        "capability buyer or expiry is outside the session",
+      );
+    }
+    const stream = await this.dependencies.streamReader.getStream(
+      capability.streamId as Hex,
+    );
+    if (
+      stream.status !== "active" ||
+      stream.buyer.toLowerCase() !== session.address.toLowerCase() ||
+      stream.serviceRef.toLowerCase() !== capability.serviceRef.toLowerCase()
+    ) {
+      throw new PrivySessionBrokerError(
+        "policy-denied",
+        "capability does not match an active buyer stream",
+      );
+    }
+    try {
+      return await this.dependencies.provider.signTypedData({
+        walletId: session.walletId,
+        typedData: typedDataRecord,
+        authorizationPrivateKey: this.dependencies.authorizationPrivateKey,
+      });
+    } catch (error) {
+      throw new PrivySessionBrokerError(
+        "provider-error",
+        "Privy rejected the typed-data signature request",
+        { cause: error },
+      );
+    }
+  }
+
+  private async providerSend(
+    session: PrivyBrokerSession,
+    transaction: z.infer<typeof transactionSchema>,
+  ): Promise<Hex> {
+    const { chainId: _chainId, ...privyTransaction } = transaction;
+    const sponsor = session.sponsorshipUsed !== true;
+    let hash: Hex;
+    try {
+      hash = await this.dependencies.provider.sendTransaction({
+        walletId: session.walletId,
+        transaction: privyTransaction,
+        authorizationPrivateKey: this.dependencies.authorizationPrivateKey,
+        sponsor,
+      });
+    } catch (error) {
+      throw new PrivySessionBrokerError(
+        "provider-error",
+        "Privy rejected the transaction request",
+        { cause: error },
+      );
+    }
+    if (sponsor) {
+      await this.dependencies.store.put({
+        ...session,
+        sponsorshipUsed: true,
+      });
+    }
+    return hash;
+  }
+
+  private async requireSession(token: string): Promise<PrivyBrokerSession> {
+    if (token.length < 32) {
+      throw new PrivySessionBrokerError(
+        "unauthorized",
+        "invalid session token",
+      );
+    }
+    const session = await this.dependencies.store.findByTokenHash(
+      hashToken(token),
+    );
+    if (
+      session === null ||
+      !safeTokenHashEqual(session.tokenHash, hashToken(token))
+    ) {
+      throw new PrivySessionBrokerError(
+        "unauthorized",
+        "invalid session token",
+      );
+    }
+    if (session.revokedAt !== undefined) {
+      throw new PrivySessionBrokerError("revoked", "session is revoked");
+    }
+    if (this.now() >= session.expiresAt) {
+      throw new PrivySessionBrokerError("expired", "session is expired");
+    }
+    return session;
+  }
+
+  private async withSessionLock<T>(
+    id: string,
+    operation: () => Promise<T>,
+  ): Promise<T> {
+    const previous = this.locks.get(id) ?? Promise.resolve();
+    let release!: () => void;
+    const next = new Promise<void>((resolve) => {
+      release = resolve;
+    });
+    const queued = previous.then(() => next);
+    this.locks.set(id, queued);
+    await previous;
+    try {
+      return await operation();
+    } finally {
+      release();
+      if (this.locks.get(id) === queued) this.locks.delete(id);
+    }
+  }
+
+  private now(): number {
+    return this.dependencies.nowSeconds?.() ?? Math.floor(Date.now() / 1_000);
+  }
+}
+
+export class InMemoryPrivyBrokerSessionStore implements PrivyBrokerSessionStore {
+  private readonly sessions = new Map<string, PrivyBrokerSession>();
+
+  async findByTokenHash(tokenHash: string): Promise<PrivyBrokerSession | null> {
+    return (
+      [...this.sessions.values()].find((session) =>
+        safeTokenHashEqual(session.tokenHash, tokenHash),
+      ) ?? null
+    );
+  }
+
+  async get(id: string): Promise<PrivyBrokerSession | null> {
+    return this.sessions.get(id) ?? null;
+  }
+
+  async put(session: PrivyBrokerSession): Promise<void> {
+    this.sessions.set(session.id, session);
+  }
+}
+
+function parseWireMandate(value: unknown): SignedSpendMandate {
+  return signedSpendMandateSchema.parse(wireMandateSchema.parse(value));
+}
+
+function requireAddress(value: string): Address {
+  if (!isAddress(value)) {
+    throw new PrivySessionBrokerError(
+      "invalid-request",
+      "invalid wallet address",
+    );
+  }
+  return getAddress(value);
+}
+
+function assertMandateWithinPrivyPolicy(signed: SignedSpendMandate): void {
+  const mandate = signed.mandate;
+  if (
+    mandate.maxPerStreamUsdc > PRIVY_POLICY_MAX_TOTAL_USDC ||
+    mandate.maxTotalUsdc > PRIVY_POLICY_MAX_TOTAL_USDC ||
+    mandate.maxRatePerSecondUsdc > PRIVY_POLICY_MAX_RATE_USDC ||
+    mandate.maxDurationSeconds > PRIVY_POLICY_MAX_DURATION_SECONDS
+  ) {
+    throw new PrivySessionBrokerError(
+      "policy-denied",
+      "mandate exceeds the installed Privy signer policy",
+    );
+  }
+}
+
+function assertBoundedApproval(data: Hex, signed: SignedSpendMandate): void {
+  const decoded = decodeFunctionData({
+    abi: [
+      {
+        type: "function",
+        name: "approve",
+        stateMutability: "nonpayable",
+        inputs: [
+          { name: "spender", type: "address" },
+          { name: "amount", type: "uint256" },
+        ],
+        outputs: [{ name: "", type: "bool" }],
+      },
+    ] as const,
+    data,
+  });
+  const [spender, amount] = decoded.args;
+  if (
+    spender.toLowerCase() !== PRIVY_BROKER_ESCROW.toLowerCase() ||
+    amount > signed.mandate.maxPerStreamUsdc
+  ) {
+    throw new PrivySessionBrokerError(
+      "policy-denied",
+      "USDC approval exceeds the mandate or targets another spender",
+    );
+  }
+}
+
+function assertChainId(value: string | number | undefined): void {
+  if (value === undefined) return;
+  const parsed =
+    typeof value === "number"
+      ? value
+      : value.startsWith("0x")
+        ? Number.parseInt(value.slice(2), 16)
+        : Number(value);
+  if (parsed !== PRIVY_BROKER_CHAIN_ID) {
+    throw new PrivySessionBrokerError(
+      "policy-denied",
+      "transaction is not for Base Sepolia",
+    );
+  }
+}
+
+function assertZeroValue(value: string | number | bigint | undefined): void {
+  if (value === undefined) return;
+  const parsed =
+    typeof value === "bigint"
+      ? value
+      : typeof value === "number"
+        ? BigInt(value)
+        : BigInt(value);
+  if (parsed !== 0n) {
+    throw new PrivySessionBrokerError(
+      "policy-denied",
+      "native-token transfers are not allowed",
+    );
+  }
+}
+
+function hashToken(token: string): string {
+  return createHash("sha256").update(token).digest("hex");
+}
+
+function safeTokenHashEqual(left: string, right: string): boolean {
+  const a = Buffer.from(left, "hex");
+  const b = Buffer.from(right, "hex");
+  return a.length === b.length && timingSafeEqual(a, b);
+}