Safe Abstract Session Keys

useBondingCurveBuy.ts

import { usePublicClient } from "./usePublicClient";
import { IToken } from "@/types/token";
import { parseAbi, parseEther } from "viem";
import { ethers } from "ethers";
import { useAbstractClient } from "@abstract-foundation/agw-react";
import { useAbstractSession } from "@/hooks/useCreateAbstractSession";
import { privateKeyToAccount } from "viem/accounts";
import { useSessionClientChain } from "./useSessionClientChain";

export const useBondingCurveBuy = (chain: any) => {
  const { data: client } = useAbstractClient();
  const { getStoredSession } = useAbstractSession(chain);
  const sessionClientChain = useSessionClientChain(chain);

  const buyTokens = async (
    token: IToken,
    ethAmount: string,
    minimumOut: bigint
  ) => {
    try {
      const publicClient = await usePublicClient(chain ?? token.chain, true);
      
      const minimumOutInWei = BigInt(
        ethers.utils.parseUnits(minimumOut.toString(), "wei").toString()
      );

      const sessionData = await getStoredSession();
      
      if (!sessionData) {
        throw new Error("No session data found");
      }

      const sessionSigner = privateKeyToAccount(sessionData.privateKey);

      const sessionClient = client?.toSessionClient(
        sessionSigner,
        sessionData.session
      );
      
      if (!sessionClient) {
        throw new Error("Failed to create session client");
      }

      const tx = await sessionClient.writeContract({
        abi: parseAbi(["function buy(address,uint256) external payable"]),
        account: sessionClient.account,
        chain: sessionClientChain,
        address: chain.g8KeepFactoryAddress ?? token.chain.g8KeepFactoryAddress as `0x${string}`,
        functionName: "buy",
        args: [token.address as `0x${string}`, minimumOutInWei],
        value: BigInt(parseEther(ethAmount)),
      });

      await publicClient.waitForTransactionReceipt({ hash: tx! });
      return tx;
    } catch (error) {
      console.error("Error buying from bonding curve:", error);
      throw error;
    }
  };

  return { buyTokens };
};

useCreateAbstractSession.ts

import { useCreateSession } from "@abstract-foundation/agw-react";
import { generatePrivateKey, privateKeyToAccount } from "viem/accounts";
import { toFunctionSelector } from "viem";
import { LimitType, getSessionHash } from "@abstract-foundation/agw-client/sessions";
import { useAccount } from "wagmi";
import { usePublicClient } from "./usePublicClient";
import { SESSION_KEY_VALIDATOR } from "@/types/constants";
import { G8KEEP_BONDING_CURVE_FACTORY_ADDRESS, DEV_G8KEEP_BONDING_CURVE_FACTORY_ADDRESS } from "@/types/constants";

const LOCAL_STORAGE_KEY_PREFIX = "abstract_session_";
const ENCRYPTION_KEY_PREFIX = "encryption_key_";

const ABI = [
  {
    inputs: [
      { internalType: "address", name: "account", type: "address" },
      { internalType: "bytes32", name: "sessionHash", type: "bytes32" },
    ],
    name: "sessionStatus",
    outputs: [
      { internalType: "enum SessionLib.Status", name: "", type: "uint8" },
    ],
    stateMutability: "view",
    type: "function",
  },
];

enum SessionStatus {
  NotInitialized = 0,
  Active = 1,
  Closed = 2,
  Expired = 3,
}

export const useAbstractSession = (chain: any) => {
  const { address } = useAccount();
  const { createSessionAsync } = useCreateSession();
  const factoryAddress = chain?.name == "Abstract" ? G8KEEP_BONDING_CURVE_FACTORY_ADDRESS : DEV_G8KEEP_BONDING_CURVE_FACTORY_ADDRESS;

  const getStorageKey = (userAddress: string) =>
    `${LOCAL_STORAGE_KEY_PREFIX}${userAddress}`;

  const getEncryptionKey = async (userAddress: string): Promise<CryptoKey> => {
    const storedKey = localStorage.getItem(
      `${ENCRYPTION_KEY_PREFIX}${userAddress}`
    );
    if (storedKey) {
      return crypto.subtle.importKey(
        "raw",
        Buffer.from(storedKey, "hex"),
        { name: "AES-GCM" },
        false,
        ["encrypt", "decrypt"]
      );
    }

    const key = await crypto.subtle.generateKey(
      { name: "AES-GCM", length: 256 },
      true,
      ["encrypt", "decrypt"]
    );

    const exportedKey = await crypto.subtle.exportKey("raw", key);
    localStorage.setItem(
      `${ENCRYPTION_KEY_PREFIX}${userAddress}`,
      Buffer.from(exportedKey).toString("hex")
    );

    return key;
  };

  const encrypt = async (data: string, key: CryptoKey): Promise<string> => {
    const iv = crypto.getRandomValues(new Uint8Array(12));
    const encrypted = await crypto.subtle.encrypt(
      { name: "AES-GCM", iv },
      key,
      new TextEncoder().encode(data)
    );

    return JSON.stringify({
      iv: Buffer.from(iv).toString("hex"),
      data: Buffer.from(encrypted).toString("hex"),
    });
  };

  const decrypt = async (
    encryptedData: string,
    key: CryptoKey
  ): Promise<string> => {
    const { iv, data } = JSON.parse(encryptedData);
    const decrypted = await crypto.subtle.decrypt(
      { name: "AES-GCM", iv: Buffer.from(iv, "hex") },
      key,
      Buffer.from(data, "hex")
    );

    return new TextDecoder().decode(decrypted);
  };

  const getStoredSession = async () => {
    if (!address) return null;
    const encryptedData = localStorage.getItem(getStorageKey(address));
    if (!encryptedData) return null;

    try {
      const key = await getEncryptionKey(address);
      const decryptedData = await decrypt(encryptedData, key);
      const parsedData = JSON.parse(decryptedData);
      const sessionHash = getSessionHash(parsedData.session);
      await validateSession(address, sessionHash);
      return JSON.parse(decryptedData);
    } catch (error) {
      console.error("Failed to decrypt session:", error);
      return null;
    }
  };

  const validateSession = async (
    address: string,
    sessionHash: string
  ) => {
    const publicClient = await usePublicClient(chain, true);

    try {
      const status = (await publicClient.readContract({
        address: SESSION_KEY_VALIDATOR as `0x${string}`,
        abi: ABI,
        functionName: "sessionStatus",
        args: [address as `0x${string}`, sessionHash],
      })) as SessionStatus;

      const isValid = status === SessionStatus.Active;

      if (!isValid) {
        clearStoredSession();
        await createAndStoreSession();
      }
    } catch (error) {
      console.error("Failed to validate session:", error);
      return;
    }
  }

  const createAndStoreSession = async () => {
    if (!address) return null;

    try {
      const sessionPrivateKey = generatePrivateKey();
      const sessionSigner = privateKeyToAccount(sessionPrivateKey);
      const maxBigInt = BigInt(
        "0xffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff"
      );
      const almostMaxBigInt = maxBigInt - BigInt(1);

      const { session } = await createSessionAsync({
        session: {
          signer: sessionSigner.address,
          expiresAt: BigInt(Math.floor(Date.now() / 1000) + 60 * 60 * 24 * 30), // 30 days
          feeLimit: {
            limitType: LimitType.Lifetime,
            limit: almostMaxBigInt,
            period: BigInt(0),
          },
          callPolicies: [
            {
              target: factoryAddress,
              selector: toFunctionSelector(
                "buy(address,uint256) external payable"
              ),
              valueLimit: {
                limitType: LimitType.Unlimited,
                limit: almostMaxBigInt,
                period: BigInt(0),
              },
              maxValuePerUse: almostMaxBigInt,
              constraints: [],
            },
            {
              target: factoryAddress,
              selector: toFunctionSelector(
                "sell(address,uint112,uint112) external payable"
              ),
              valueLimit: {
                limitType: LimitType.Unlimited,
                limit: BigInt(0),
                period: BigInt(0),
              },
              maxValuePerUse: almostMaxBigInt,
              constraints: [],
            },
            {
              target: factoryAddress,
              selector: toFunctionSelector(
                "deployToken(address,address,string,string,uint256,string,string,bytes32) external payable"
              ),
              valueLimit: {
                limitType: LimitType.Unlimited,
                limit: almostMaxBigInt,
                period: BigInt(0),
              },
              maxValuePerUse: almostMaxBigInt,
              constraints: [],
            },
          ],
          transferPolicies: [],
        },
      });

      const sessionData = { session, privateKey: sessionPrivateKey };
      const key = await getEncryptionKey(address);
      const encryptedData = await encrypt(
        JSON.stringify(sessionData, (_, value) =>
          typeof value === "bigint" ? value.toString() : value
        ),
        key
      );

      localStorage.setItem(getStorageKey(address), encryptedData);
      return sessionData;
    } catch (error) {
      console.error("Failed to create session:", error);
      throw new Error("Session creation failed");
    }
  };

  const clearStoredSession = () => {
    if (address) {
      localStorage.removeItem(getStorageKey(address));
      localStorage.removeItem(`${ENCRYPTION_KEY_PREFIX}${address}`);
    }
  };

  return { getStoredSession, validateSession, createAndStoreSession, clearStoredSession };
};