@unionlabs/payments 0.1.0 → 0.2.0

package/src/Payment.ts

@@ -0,0 +1,523 @@
+/**
+ * - for deposits you only need to sign on the source chain
+ * - for redemptions you need wallet on dest chain
+ */
+import { Effect, pipe } from "effect";
+import * as A from "effect/Array";
+import {
+  type Address,
+  encodePacked,
+  hexToBigInt,
+  hexToBytes,
+  keccak256,
+  padHex,
+  sha256,
+  toHex,
+} from "viem";
+import * as Attestor from "./Attestor.js";
+import * as Constants from "./Constants.js";
+import * as ZAssetRegistry from "./constants/z-asset-registry.js";
+import type { Erc20Address } from "./Domain.js";
+import * as Domain from "./Domain.js";
+import * as Error from "./Error.js";
+import * as Prover from "./Prover.js";
+import * as PublicClient from "./PublicClient.js";
+import * as Poseidon2 from "./poseidon2.js";
+import type * as PS from "./Schema.js";
+import type { LightClientData } from "./types.js";
+import * as Utils from "./Utils.js";
+import * as WalletClient from "./WalletClient.js";
+
+// TODO: into poseidon2 module
+const PRIME =
+  0x30644e72e131a029b85045b68181585d2833e84879b9709143e1f593f0000001n; // BN254 Fr order
+
+/**
+ * @public
+ */
+export interface DepositOptions {
+  /** The 32-byte secret payment key as a hex string */
+  readonly paymentKey: PS.Hex;
+  /** Array of 0-4 beneficiary addresses*/
+  readonly beneficiaries: ReadonlyArray<Address>;
+  /** Amount to deposit (in underlying token's smallest unit) */
+  readonly amount: bigint;
+  readonly destinationChainId: bigint;
+  readonly sourceChainId: bigint;
+  readonly srcErc20Address: Domain.Erc20Address;
+}
+
+/**
+ * @public
+ */
+export interface RedeemOptions {
+  /** The 32-byte secret payment key as a hex string */
+  paymentKey: PS.Hex;
+  /** Array of 0-4 beneficiary addresses */
+  beneficiaries: Address[];
+  /** The beneficiary address to redeem to */
+  beneficiary: Address;
+  /** Amount to redeem */
+  amount: bigint;
+  /** Light client IDs to include in the proof (for anonymity set) */
+  clientIds: number[];
+  /** The specific light client ID to use for the proof */
+  selectedClientId: number;
+  /** Auto unwrap */
+  unwrap?: boolean | undefined;
+}
+
+export const getDepositAddress = Effect.fn("getDepositAddress")(
+  (
+    secret: PS.Hex,
+    beneficiaries: ReadonlyArray<PS.Address>,
+    destinationChainId: bigint,
+  ) =>
+    Effect.gen(function* () {
+      const paddedBeneficiaries = yield* padBeneficiaries(beneficiaries);
+      const unspendableAddress = Poseidon2.computeUnspendableAddress(
+        secret,
+        destinationChainId,
+        paddedBeneficiaries,
+      );
+
+      return Domain.ZAssetAddress(unspendableAddress);
+    }),
+);
+const padBeneficiaries = Effect.fn("padBeneficiaries")(
+  (beneficiaries: ReadonlyArray<PS.Address>) =>
+    Effect.gen(function* () {
+      if (beneficiaries.length > 4) {
+        // throw new Error("Maximum 4 beneficiaries allowed");
+        return yield* new Error.SystemError({
+          method: "padBeneficiaries",
+          module: "PaymentClient",
+          reason: "InvalidData",
+          description: "Maximum 4 beneficiaries allowed",
+        });
+      }
+      const padded: [PS.Address, PS.Address, PS.Address, PS.Address] = [
+        beneficiaries[0] ?? Constants.ZERO_ADDRESS,
+        beneficiaries[1] ?? Constants.ZERO_ADDRESS,
+        beneficiaries[2] ?? Constants.ZERO_ADDRESS,
+        beneficiaries[3] ?? Constants.ZERO_ADDRESS,
+      ];
+      return padded;
+    }),
+);
+
+export const getNullifier = Effect.fn("getNullifier")(
+  (secret: PS.Hex, destinationChainId: bigint) =>
+    Effect.try({
+      try: () => Poseidon2.computeNullifier(secret, destinationChainId),
+      catch: (cause) =>
+        new Error.SystemError({
+          method: "getNullifier",
+          module: "PaymentClient",
+          reason: "InvalidData",
+          cause,
+        }),
+    }),
+);
+/**
+ * Deterministically shuffle light clients using the secret as seed.
+ * Same secret always produces the same ordering for privacy.
+ * @public
+ */
+export function deterministicShuffleClients(
+  clients: LightClientData[],
+  secret: PS.Hex,
+): LightClientData[] {
+  if (clients.length <= 1) {
+    return [...clients];
+  }
+
+  // Seed = SHA256(secret || clientHeights as big-endian uint64s)
+  const secretBytes = hexToBytes(secret);
+  const heightsBytes = new Uint8Array(clients.length * 8);
+  for (let i = 0; i < clients.length; i++) {
+    const view = new DataView(heightsBytes.buffer, i * 8, 8);
+    view.setBigUint64(0, clients[i]!.height, false);
+  }
+
+  const combined = new Uint8Array(secretBytes.length + heightsBytes.length);
+  combined.set(secretBytes, 0);
+  combined.set(heightsBytes, secretBytes.length);
+  const seedHex = sha256(combined);
+  const seed = hexToBytes(seedHex);
+
+  // Fisher-Yates shuffle with deterministic PRNG
+  const shuffled = [...clients];
+  let seedIndex = 0;
+  for (let i = shuffled.length - 1; i > 0; i--) {
+    const randomByte = seed[seedIndex % seed.length]!;
+    seedIndex++;
+    const j = randomByte % (i + 1);
+    [shuffled[i], shuffled[j]] = [shuffled[j]!, shuffled[i]!];
+  }
+
+  return shuffled;
+}
+
+export const computeStorageSlot = Effect.fn("computeStorageSlot")(function* (
+  address: Address,
+  mappingSlot: bigint,
+) {
+  const paddedAddress = padHex(address, { size: 32 });
+  const paddedSlot = padHex(toHex(mappingSlot), { size: 32 });
+  return yield* Effect.succeed(
+    keccak256(
+      encodePacked(["bytes32", "bytes32"], [paddedAddress, paddedSlot]),
+    ),
+  );
+});
+
+const fetchMptProof = Effect.fn(function* (options: {
+  tokenAddress: `0x${string}`;
+  storageSlot: `0x${string}`;
+  blockNumber: bigint;
+}) {
+  yield* Effect.log(options);
+
+  const srcClient = yield* PublicClient.SourcePublicClient;
+  const proof = yield* srcClient.getProof({
+    address: options.tokenAddress,
+    storageKeys: [options.storageSlot],
+    blockNumber: options.blockNumber,
+  });
+
+  let storageProof: `0x${string}`[] = [];
+  let storageValue: `0x${string}` = "0x0";
+
+  if (proof.storageProof.length > 0 && proof.storageProof[0]) {
+    storageProof = proof.storageProof[0].proof as `0x${string}`[];
+    storageValue = toHex(proof.storageProof[0].value);
+  }
+
+  const obj = {
+    accountProof: proof.accountProof,
+    storageProof,
+    storageValue,
+    storageRoot: proof.storageHash,
+  };
+
+  yield* Effect.log("[fetchMptProof]", obj);
+
+  return obj;
+});
+
+export const generateProof = Effect.fn(function* (options: {
+  secret: PS.Hex;
+  beneficiaries: Address[];
+  beneficiary: Address;
+  amount: bigint;
+  clientIds: number[];
+  selectedClientId: number;
+  srcChainId: bigint;
+  dstChainId: bigint;
+  srcErc20Address: Erc20Address;
+  dstErc20Address: Erc20Address;
+}) {
+  if (!options.clientIds.includes(options.selectedClientId)) {
+    return yield* Effect.fail({
+      proof: {
+        success: false,
+        error: `selectedClientId ${options.selectedClientId} not in clientIds`,
+      },
+    });
+  }
+
+  if (options.beneficiary === Constants.ZERO_ADDRESS) {
+    return yield* Effect.fail({
+      proof: {
+        success: false,
+        error: "Beneficiary address cannot be zero",
+      },
+    });
+  }
+
+  const srcZAssetAddress = yield* ZAssetRegistry.getZAssetOrError(
+    options.srcChainId,
+    options.srcErc20Address,
+  );
+  const dstZAssetAddress = yield* ZAssetRegistry.getZAssetOrError(
+    options.dstChainId,
+    options.dstErc20Address,
+  );
+
+  const prover = yield* Prover.Prover;
+  const destinationPublicClient = yield* PublicClient.DestinationPublicClient;
+  const paddedBeneficiaries = yield* padBeneficiaries(options.beneficiaries);
+  const depositAddress = yield* getDepositAddress(
+    options.secret,
+    options.beneficiaries,
+    options.dstChainId,
+  );
+  const lightClients = yield* fetchLightClients({
+    zAssetAddress: dstZAssetAddress,
+    clientIds: options.clientIds,
+  });
+
+  if (!A.isNonEmptyReadonlyArray(lightClients)) {
+    return yield* new Error.SystemError({
+      module: "PaymentClient",
+      method: "generateProof",
+      reason: "InvalidData",
+      description: "No valid light clients found",
+    });
+  }
+
+  // Deterministic shuffle: same secret always produces the same ordering for privacy
+  const shuffled = deterministicShuffleClients(lightClients, options.secret);
+  const selectedClientIndex = yield* pipe(
+    A.findFirstIndex(shuffled, (c) => c.clientId === options.selectedClientId),
+    Effect.mapError(
+      (cause) =>
+        new Error.SystemError({
+          module: "PaymentClient",
+          method: "generateProof",
+          reason: "InvalidData",
+          cause,
+          description: `Client not found after fetching`,
+        }),
+    ),
+  );
+
+  const selectedClient = A.unsafeGet(shuffled, selectedClientIndex);
+
+  const { tokenAddressKey, balanceSlot } =
+    yield* destinationPublicClient.getCounterparty(
+      dstZAssetAddress,
+      options.selectedClientId,
+    );
+
+  // Check if counterparty is configured for this light client
+  const ZERO_BYTES32 =
+    "0x0000000000000000000000000000000000000000000000000000000000000000";
+
+  yield* Effect.log({ balanceSlot });
+
+  if (balanceSlot === ZERO_BYTES32 || tokenAddressKey === ZERO_BYTES32) {
+    return yield* new Error.SystemError({
+      module: "PaymentClient",
+      method: "generateProof",
+      reason: "InvalidData",
+      description:
+        `Light client is not configured as a counterparty on the destination ZAsset. ` +
+        `Please call setCounterparty() on the ZAsset contract first.`,
+    });
+  }
+
+  const mappingSlot = hexToBigInt(balanceSlot);
+
+  yield* Effect.log({ mappingSlot });
+
+  const nullifier = yield* getNullifier(options.secret, options.dstChainId);
+
+  // Check balance at the selected light client's height
+  const balance = yield* destinationPublicClient.readBalanceAtHeight({
+    srcErc20Address: options.srcErc20Address,
+    dstErc20Address: options.dstErc20Address,
+    srcChainId: options.srcChainId,
+    dstChainId: options.dstChainId,
+    depositAddress,
+    nullifier,
+    height: selectedClient.height,
+  });
+
+  if (options.amount > balance.available) {
+    return yield* new Error.SystemError({
+      module: "PaymentClient",
+      method: "generateProof",
+      reason: "InvalidData",
+      description:
+        `Insufficient available balance. Requested: ${options.amount}, Available: ${balance.available} ` +
+        `(Confirmed at height ${selectedClient.height}: ${balance.confirmed}, Already redeemed: ${balance.redeemed})`,
+    });
+  }
+
+  const storageSlot = yield* computeStorageSlot(depositAddress, mappingSlot);
+
+  const mptProof = yield* fetchMptProof({
+    tokenAddress: srcZAssetAddress,
+    storageSlot,
+    blockNumber: selectedClient.height,
+  });
+
+  const witness: Prover.WitnessData = {
+    secret: options.secret,
+    dstChainId: options.dstChainId,
+    beneficiaries: paddedBeneficiaries,
+    beneficiary: options.beneficiary,
+    redeemAmount: options.amount,
+    alreadyRedeemed: balance.redeemed,
+    lightClients: shuffled,
+    selectedClientIndex,
+    mptProof,
+    srcZAssetAddress: srcZAssetAddress,
+    mappingSlot: `0x${mappingSlot.toString(16)}` as `0x${string}`,
+  };
+
+  const proofResult = yield* prover.generateProof(witness);
+
+  yield* Effect.log({ proofResult });
+
+  return {
+    proof: proofResult,
+    metadata: {
+      depositAddress,
+      beneficiary: options.beneficiary,
+      value: options.amount,
+      lightClients: shuffled,
+      nullifier,
+    },
+  } as const;
+});
+
+const fetchLightClients = Effect.fn(function* (options: {
+  zAssetAddress: Domain.ZAssetAddress;
+  clientIds: number[];
+}) {
+  const { zAssetAddress, clientIds } = options;
+  const dstClient = yield* PublicClient.PublicClient;
+  const ibcHandlerAddress =
+    yield* dstClient.getIbcHandlerAddress(zAssetAddress);
+  const results: LightClientData[] = [];
+  for (const clientId of clientIds) {
+    const lightClientAddress = yield* dstClient.getLightClientAddress(
+      ibcHandlerAddress,
+      clientId,
+    );
+    const height = yield* dstClient.getLatestHeight(
+      lightClientAddress,
+      clientId,
+    );
+    const stateRootIndex = yield* dstClient.getStateRootIndex(
+      zAssetAddress,
+      clientId,
+    );
+    const { stateRoot } = yield* dstClient.getConsensusState(
+      lightClientAddress,
+      clientId,
+      height,
+      stateRootIndex,
+    );
+
+    results.push({ clientId, height, stateRoot });
+  }
+
+  return results;
+});
+
+export const prepareRedemption = Effect.fn(function* (options: {
+  /** The 32-byte secret payment key as a hex string */
+  paymentKey: PS.Hex;
+  /** Array of 0-4 beneficiary addresses */
+  beneficiaries: Address[];
+  /** The beneficiary address to redeem to */
+  beneficiary: Address;
+  /** Amount to redeem */
+  amount: bigint;
+  /** Light client IDs to include in the proof (for anonymity set) */
+  clientIds: number[];
+  /** The specific light client ID to use for the proof */
+  selectedClientId: number;
+  srcChainId: bigint;
+  dstChainId: bigint;
+  srcErc20Address: Domain.Erc20Address;
+  dstErc20Address: Domain.Erc20Address;
+  /** Auto unwrap */
+  unwrap?: boolean | undefined;
+}) {
+  const dstWalletClient = yield* WalletClient.DestinationWalletClient;
+  const attestor = yield* Attestor.Attestor;
+
+  // 1. Generate proof
+  const proofResult = yield* generateProof({
+    secret: options.paymentKey,
+    beneficiaries: options.beneficiaries as unknown as any,
+    beneficiary: options.beneficiary,
+    amount: options.amount,
+    clientIds: options.clientIds,
+    selectedClientId: options.selectedClientId,
+    srcChainId: options.srcChainId,
+    dstChainId: options.dstChainId,
+    srcErc20Address: options.srcErc20Address,
+    dstErc20Address: options.dstErc20Address,
+  });
+
+  yield* Effect.log({ proofResult });
+
+  const depositAddress = yield* getDepositAddress(
+    options.paymentKey,
+    options.beneficiaries,
+    options.dstChainId,
+  );
+
+  yield* Effect.log({ depositAddress });
+
+  // 2. Get attestation
+  const attestationResponse = yield* attestor.get({
+    unspendableAddress: depositAddress,
+    beneficiary: options.beneficiary,
+  });
+
+  yield* Effect.log({ attestationResponse });
+
+  // 3. Build redeem params
+  const redeemParams = Utils.proofJsonToRedeemParams(proofResult.proof, {
+    lightClients: proofResult.metadata.lightClients,
+    nullifier: proofResult.metadata.nullifier,
+    value: proofResult.metadata.value,
+    beneficiary: proofResult.metadata.beneficiary,
+    attestedMessage: attestationResponse.attestedMessage,
+    signature: attestationResponse.signature,
+  });
+
+  yield* Effect.log({ redeemParams });
+
+  return yield* dstWalletClient.prepareRedemption({
+    ...redeemParams,
+    dstErc20Address: options.dstErc20Address,
+  });
+});
+
+export const generateKey = Effect.gen(function* () {
+  while (true) {
+    const bytes = new Uint8Array(32);
+    crypto.getRandomValues(bytes);
+    const value = BigInt(
+      "0x" +
+        Array.from(bytes)
+          .map((b) => b.toString(16).padStart(2, "0"))
+          .join(""),
+    );
+    if (value < PRIME) {
+      return ("0x" +
+        Array.from(bytes)
+          .map((b) => b.toString(16).padStart(2, "0"))
+          .join("")) as PS.Hex;
+    }
+  }
+});
+
+export const prepareDeposit = Effect.fn("deposit")((options: DepositOptions) =>
+  Effect.gen(function* () {
+    const srcWalletClient = yield* WalletClient.SourceWalletClient;
+
+    const depositAddress = yield* getDepositAddress(
+      options.paymentKey,
+      options.beneficiaries,
+      options.destinationChainId,
+    );
+
+    return yield* srcWalletClient.prepareDeposit({
+      amount: options.amount,
+      depositAddress,
+      sourceChainId: options.sourceChainId,
+      srcErc20Address: options.srcErc20Address,
+      beneficiaries: options.beneficiaries,
+      destinationChainId: options.destinationChainId,
+    });
+  }),
+);

package/src/PaymentError.ts

@@ -0,0 +1,39 @@
+/**
+ * @since 1.0.0
+ */
+import { TypeIdError } from "./Error.js";
+
+/**
+ * @since 1.0.0
+ */
+export const PaymentErrorTypeId: unique symbol = Symbol.for(
+  "@unionlabs/payments/PaymentError",
+);
+
+/**
+ * @since 1.0.0
+ */
+export type PaymentErrorTypeId = typeof PaymentErrorTypeId;
+
+/**
+ * @since 1.0.0
+ */
+export class PaymentError extends TypeIdError(
+  PaymentErrorTypeId,
+  "PaymentError",
+)<{
+  cause: unknown;
+  message?: string;
+}> {}
+
+// /**
+//  * @since 1.0.0
+//  */
+// export class ResultLengthMismatch extends TypeIdError(PaymentErrorTypeId, "ResultLengthMismatch")<{
+//   readonly expected: number
+//   readonly actual: number
+// }> {
+//   get message() {
+//     return `Expected ${this.expected} results but got ${this.actual}`
+//   }
+// }