feat: support private policy

src/paymasterclient.ts

@@ -9,6 +9,15 @@ export type IsSponsorableResponse = {
   SponsorWebsite: string
 }
 
+export type IsSponsorableOptions = {
+  PrivatePolicyUUID?: string
+}
+
+export type SendRawTransactionOptions = {
+  PrivatePolicyUUID?: string
+  UserAgent?: string
+}
+
 export enum GaslessTransactionStatus { New = 0, Pending = 1, Confirmed = 2, Failed = 3, Invalid = 4}
 
 export type GaslessTransaction = {
@@ -48,36 +57,99 @@ export type Bundle = {
   readonly ChainID: number
 }
 
-export class PaymasterClient extends ethers.JsonRpcProvider {
-  constructor(url?: string | FetchRequest, network?: Networkish, options?: JsonRpcApiProviderOptions) {
-    super(url, network, options)
+export class PaymasterClient {
+  private sponsorClient: ethers.JsonRpcProvider
+  private userClient: ethers.JsonRpcProvider
+
+  constructor(
+    userUrl: string | FetchRequest,
+    sponsorUrl: string | FetchRequest,
+    network?: Networkish,
+    options?: JsonRpcApiProviderOptions
+  ) {
+    this.userClient = new ethers.JsonRpcProvider(userUrl, network, options)
+    this.sponsorClient = new ethers.JsonRpcProvider(sponsorUrl, network, options)
   }
 
   async chainID(): Promise<string> {
-    return await this.send('eth_chainId', [])
+    return await this.userClient.send('eth_chainId', [])
   }
 
-  async isSponsorable(tx: TransactionRequest): Promise<IsSponsorableResponse> {
-    return await this.send('pm_isSponsorable', [tx])
+  async isSponsorable(tx: TransactionRequest, opts: IsSponsorableOptions = {}): Promise<IsSponsorableResponse> {
+    if (opts.PrivatePolicyUUID) {
+      // Create a new provider with the updated header
+      const newConnection = this.sponsorClient._getConnection();
+      newConnection.setHeader("X-MegaFuel-Policy-Uuid", opts.PrivatePolicyUUID);
+
+      // Create a new provider with the modified connection
+      const sponsorProviderWithHeader = new ethers.JsonRpcProvider(
+        newConnection,
+        (this.sponsorClient as any)._network,
+        {
+          staticNetwork: (this.sponsorClient as any)._network,
+          batchMaxCount: (this.sponsorClient as any).batchMaxCount,
+          polling: (this.sponsorClient as any).polling
+        }
+      );
+
+      return await sponsorProviderWithHeader.send('pm_isSponsorable', [tx]);
+    }
+    return await this.userClient.send('pm_isSponsorable', [tx]);
   }
 
-  async sendRawTransaction(signedTx: string): Promise<string> {
-    return await this.send('eth_sendRawTransaction', [signedTx])
+  async sendRawTransaction(signedTx: string, opts: SendRawTransactionOptions = {}): Promise<string> {
+    let sponsorProvider = this.sponsorClient;
+
+    if (opts.UserAgent || opts.PrivatePolicyUUID) {
+      // Create a new provider with the updated headers
+      const newConnection = this.sponsorClient._getConnection();
+
+      if (opts.UserAgent) {
+        newConnection.setHeader("User-Agent", opts.UserAgent);
+      }
+      if (opts.PrivatePolicyUUID) {
+        newConnection.setHeader("X-MegaFuel-Policy-Uuid", opts.PrivatePolicyUUID);
+      }
+
+      // Create a new provider with the modified connection
+      sponsorProvider = new ethers.JsonRpcProvider(
+        newConnection,
+        (this.sponsorClient as any)._network,
+        {
+          staticNetwork: (this.sponsorClient as any)._network,
+          batchMaxCount: (this.sponsorClient as any).batchMaxCount,
+          polling: (this.sponsorClient as any).polling
+        }
+      );
+    }
+
+    if (opts.PrivatePolicyUUID) {
+      return await sponsorProvider.send('eth_sendRawTransaction', [signedTx]);
+    }
+    return await this.userClient.send('eth_sendRawTransaction', [signedTx]);
   }
 
   async getGaslessTransactionByHash(hash: string): Promise<GaslessTransaction> {
-    return await this.send('eth_getGaslessTransactionByHash', [hash])
+    return await this.userClient.send('eth_getGaslessTransactionByHash', [hash])
   }
 
   async getSponsorTxByTxHash(hash: string): Promise<SponsorTx> {
-    return await this.send('pm_getSponsorTxByTxHash', [hash])
+    return await this.userClient.send('pm_getSponsorTxByTxHash', [hash])
   }
 
   async getSponsorTxByBundleUuid(bundleUuid: string): Promise<SponsorTx> {
-    return await this.send('pm_getSponsorTxByBundleUuid', [bundleUuid])
+    return await this.userClient.send('pm_getSponsorTxByBundleUuid', [bundleUuid])
   }
 
   async getBundleByUuid(bundleUuid: string): Promise<Bundle> {
-    return await this.send('pm_getBundleByUuid', [bundleUuid])
+    return await this.userClient.send('pm_getBundleByUuid', [bundleUuid])
+  }
+
+  getSponsorProvider(): ethers.JsonRpcProvider {
+    return this.sponsorClient
+  }
+
+  getUserProvider(): ethers.JsonRpcProvider {
+    return this.userClient
   }
 }

tests/env.ts

@@ -16,3 +16,4 @@ export const ACCOUNT_ADDRESS = '0xF9A8db17431DD8563747D6FC770297E438Aa12eB'
 export const CONTRACT_METHOD = '0xa9059cbb'
 export const TOKEN_CONTRACT_ADDRESS = '0xeD24FC36d5Ee211Ea25A80239Fb8C4Cfd80f12Ee'
 export const RECIPIENT_ADDRESS = '0xDE08B1Fd79b7016F8DD3Df11f7fa0FbfdF07c941'
+export const PRIVATE_POLICY_UUID = "90f1ba4c-1f93-4759-b8a9-da4d59c668b4"

tests/paymaster.spec.ts

@@ -7,7 +7,8 @@ import {
   transformToGaslessTransaction,
   delay, transformSponsorTxResponse, transformBundleResponse,
 } from './utils'
-import {TOKEN_CONTRACT_ADDRESS, CHAIN_ID, RECIPIENT_ADDRESS} from './env'
+import {IsSponsorableOptions, SendRawTransactionOptions} from '../src/paymasterclient'
+import {TOKEN_CONTRACT_ADDRESS, CHAIN_ID, RECIPIENT_ADDRESS, PRIVATE_POLICY_UUID} from './env'
 import {ethers} from 'ethers'
 
 let TX_HASH = ''
@@ -34,7 +35,7 @@ describe('paymasterQuery', () => {
     test('should successfully determine if transaction is sponsorable', async () => {
       const tokenContract = new ethers.Contract(TOKEN_CONTRACT_ADDRESS, tokenAbi, wallet)
       const tokenAmount = ethers.parseUnits('0', 18)
-      const nonce = await paymasterClient.getTransactionCount(wallet.address, 'pending')
+      const nonce = await paymasterClient.getUserProvider().getTransactionCount(wallet.address, 'pending')
 
       const transaction = await tokenContract.transfer.populateTransaction(RECIPIENT_ADDRESS.toLowerCase(), tokenAmount)
       transaction.from = wallet.address
@@ -96,4 +97,54 @@ describe('paymasterQuery', () => {
       expect(sponsorTx.TxHash).toEqual(tx.TxHash)
     }, 13000)
   })
+
+
+    /**
+   * Test for checking if a private policy transaction is sponsorable.
+   */
+    describe('isSponsorable', () => {
+      test('should successfully determine if transaction is sponsorable', async () => {
+        const tokenContract = new ethers.Contract(TOKEN_CONTRACT_ADDRESS, tokenAbi, wallet)
+        const tokenAmount = ethers.parseUnits('0', 18)
+        const nonce = await paymasterClient.getUserProvider().getTransactionCount(wallet.address, 'pending')
+
+        const transaction = await tokenContract.transfer.populateTransaction(RECIPIENT_ADDRESS.toLowerCase(), tokenAmount)
+        transaction.from = wallet.address
+        transaction.nonce = nonce
+        transaction.gasLimit = BigInt(100000)
+        transaction.chainId = BigInt(CHAIN_ID)
+        transaction.gasPrice = BigInt(0)
+
+        const safeTransaction = {
+          ...transaction,
+          gasLimit: transaction.gasLimit.toString(),
+          chainId: transaction.chainId.toString(),
+          gasPrice: transaction.gasPrice.toString(),
+        }
+
+        console.log('Prepared transaction:', safeTransaction)
+
+        const opt: IsSponsorableOptions = {
+          PrivatePolicyUUID: PRIVATE_POLICY_UUID
+        };
+
+        const resRaw = await paymasterClient.isSponsorable(safeTransaction, opt)
+        const res = transformIsSponsorableResponse(resRaw)
+        expect(res.Sponsorable).toEqual(true)
+
+        const txOpt: SendRawTransactionOptions = {
+          PrivatePolicyUUID: PRIVATE_POLICY_UUID,
+          UserAgent: "TEST USER AGENT"
+        };
+
+        const signedTx = await wallet.signTransaction(safeTransaction)
+        try {
+          const tx = await paymasterClient.sendRawTransaction(signedTx,txOpt)
+          TX_HASH = tx
+          console.log('Transaction hash received:', TX_HASH)
+        } catch (error) {
+          console.error('Transaction failed:', error)
+        }
+      }, 100000) // Extends the default timeout as this test involves network calls
+    })
 })

tests/sponsor.spec.ts

@@ -158,10 +158,10 @@ describe('sponsorQuery', () => {
    * Tests retrieving user spend data.
    */
   describe('getUserSpendData', () => {
-    test('should return null for spend data when user has none', async () => {
+    test('should return not null for user spend data', async () => {
       const res = await sponsorClient.getUserSpendData(ACCOUNT_ADDRESS, POLICY_UUID)
 
-      expect(res).toBeNull()
+      expect(res).not.toBeNull()
       console.log('User spend data:', res)
     })
   })

tests/utils.ts

@@ -16,7 +16,7 @@ export const sponsorClient = new SponsorClient(SPONSOR_URL, undefined, {staticNe
 export const assemblyProvider = new ethers.JsonRpcProvider(CHAIN_URL)
 
 // Provider for sending the transaction (e.g., could be a different network or provider)
-export const paymasterClient = new PaymasterClient(PAYMASTER_URL)
+export const paymasterClient = new PaymasterClient(PAYMASTER_URL,SPONSOR_URL+"/"+CHAIN_ID, undefined, {staticNetwork: ethers.Network.from(Number(CHAIN_ID))})
 
 export const wallet = new ethers.Wallet(PRIVATE_KEY, assemblyProvider)
 // ERC20 token ABI (only including the transfer function)