fix: modified addExpense to turn ethAddress to be optional

Some queries like getAccountInfo, getBalanceInfo, FileContentsQuery, etc. also add expense to remainingBalance but don't necessarily need to have an originalCaller (ethAddress). Therefore, addExpense can accept nullable ethAddress value. Only when ethAddress or ipAddress is valid, utilize spendingPlan logic. Otherwise, skip completely. Signed-off-by: Logan Nguyen <logan.nguyen@swirldslabs.com>
hashgraph · Sep 30, 2024 · 955111f · 955111f
1 parent f8c554e
commit 955111f
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 41 deletions.
diff --git a/packages/relay/src/lib/services/hbarLimitService/index.ts b/packages/relay/src/lib/services/hbarLimitService/index.ts
@@ -209,41 +209,43 @@ export class HbarLimitService implements IHbarLimitService {
   /**
    * Add expense to the remaining budget.
    * @param {number} cost - The cost of the expense.
-   * @param {string} ethAddress - The Ethereum address to add the expense to.
+   * @param {string} [ethAddress] - The optional Ethereum address to add the expense to.
    * @param {string} [ipAddress] - The optional IP address to add the expense to.
    * @param {string} [requestId] - An optional unique request ID for tracking the request.
    * @returns {Promise<void>} - A promise that resolves when the expense has been added.
    */
-  async addExpense(cost: number, ethAddress: string, ipAddress?: string, requestId?: string): Promise<void> {
-    if (!ethAddress && !ipAddress) {
-      throw new Error('Cannot add expense without an eth address or ip address');
-    }
+  async addExpense(cost: number, ethAddress?: string, ipAddress?: string, requestId?: string): Promise<void> {
+    const requestIdPrefix = formatRequestIdMessage(requestId);
 
-    let spendingPlan = await this.getSpendingPlan(ethAddress, ipAddress);
-    if (!spendingPlan) {
-      // Create a basic spending plan if none exists for the eth address or ip address
-      spendingPlan = await this.createBasicSpendingPlan(ethAddress, ipAddress);
-    }
+    if (ethAddress || ipAddress) {
+      let spendingPlan = await this.getSpendingPlan(ethAddress, ipAddress);
 
-    const requestIdPrefix = formatRequestIdMessage(requestId);
-    this.logger.trace(
-      `${requestIdPrefix} Adding expense of ${cost} to spending plan with ID ${spendingPlan.id}, new spentToday=${
-        spendingPlan.spentToday + cost
-      }`,
-    );
+      if (!spendingPlan) {
+        // Create a basic spending plan if none exists for the eth address or ip address
+        spendingPlan = await this.createBasicSpendingPlan(ethAddress, ipAddress);
+      }
+
+      this.logger.trace(
+        `${requestIdPrefix} Spending plan expense update: planID=${spendingPlan.id}, subscriptionType=${
+          spendingPlan.subscriptionType
+        }, cost=${cost}, originalSpentToday=${spendingPlan.spentToday}, updatedSpentToday=${
+          spendingPlan.spentToday + cost
+        }`,
+      );
 
-    // Check if the spending plan is being used for the first time today
-    if (spendingPlan.spentToday === 0) {
-      this.dailyUniqueSpendingPlansCounter[spendingPlan.subscriptionType].inc(1);
+      // Check if the spending plan is being used for the first time today
+      if (spendingPlan.spentToday === 0) {
+        this.dailyUniqueSpendingPlansCounter[spendingPlan.subscriptionType].inc(1);
+      }
+      await this.hbarSpendingPlanRepository.addAmountToSpentToday(spendingPlan.id, cost);
+
+      // Done asynchronously in the background
+      this.updateAverageDailyUsagePerSubscriptionType(spendingPlan.subscriptionType).then();
     }
 
-    await this.hbarSpendingPlanRepository.addAmountToSpentToday(spendingPlan.id, cost);
     this.remainingBudget -= cost;
     this.hbarLimitRemainingGauge.set(this.remainingBudget);
 
-    // Done asynchronously in the background
-    this.updateAverageDailyUsagePerSubscriptionType(spendingPlan.subscriptionType).then();
-
     this.logger.trace(
       `${requestIdPrefix} HBAR rate limit expense update: cost=${cost}, remainingBudget=${this.remainingBudget}`,
     );
@@ -334,12 +336,12 @@ export class HbarLimitService implements IHbarLimitService {
 
   /**
    * Gets the spending plan for the given eth address or ip address.
-   * @param {string} ethAddress - The eth address to get the spending plan for.
+   * @param {string} [ethAddress] - The eth address to get the spending plan for.
    * @param {string} [ipAddress] - The ip address to get the spending plan for.
    * @returns {Promise<IDetailedHbarSpendingPlan | null>} - A promise that resolves with the spending plan or null if none exists.
    * @private
    */
-  private async getSpendingPlan(ethAddress: string, ipAddress?: string): Promise<IDetailedHbarSpendingPlan | null> {
+  private async getSpendingPlan(ethAddress?: string, ipAddress?: string): Promise<IDetailedHbarSpendingPlan | null> {
     if (ethAddress) {
       try {
         return await this.getSpendingPlanByEthAddress(ethAddress);
@@ -381,12 +383,12 @@ export class HbarLimitService implements IHbarLimitService {
 
   /**
    * Creates a basic spending plan for the given eth address.
-   * @param {string} ethAddress - The eth address to create the spending plan for.
+   * @param {string} [ethAddress] - The eth address to create the spending plan for.
    * @param {string} [ipAddress] - The ip address to create the spending plan for.
    * @returns {Promise<IDetailedHbarSpendingPlan>} - A promise that resolves with the created spending plan.
    * @private
    */
-  private async createBasicSpendingPlan(ethAddress: string, ipAddress?: string): Promise<IDetailedHbarSpendingPlan> {
+  private async createBasicSpendingPlan(ethAddress?: string, ipAddress?: string): Promise<IDetailedHbarSpendingPlan> {
     if (!ethAddress && !ipAddress) {
       throw new Error('Cannot create a spending plan without an associated eth address or ip address');
     }

diff --git a/packages/relay/tests/lib/services/hbarLimitService/hbarLimitService.spec.ts b/packages/relay/tests/lib/services/hbarLimitService/hbarLimitService.spec.ts
@@ -628,20 +628,6 @@ describe('HbarLimitService', function () {
       sinon.assert.calledOnceWithExactly(incDailyUniqueSpendingPlansCounterSpy, 1);
     };
 
-    it('should throw an error if neither ethAddress nor ipAddress is provided', async function () {
-      const ipAddresses = ['', null, undefined];
-      const ethAddresses = ['', null, undefined];
-      const testCases = ethAddresses.flatMap((ethAddress) =>
-        ipAddresses.map((ipAddress) => ({ ethAddress, ipAddress })),
-      );
-      for (const { ethAddress, ipAddress } of testCases) {
-        // @ts-ignore
-        await expect(hbarLimitService.addExpense(100, ethAddress, ipAddress)).to.be.eventually.rejectedWith(
-          'Cannot add expense without an eth address or ip address',
-        );
-      }
-    });
-
     it('should create a basic spending plan if none exists', async function () {
       const newSpendingPlan = createSpendingPlan(mockPlanId);
       hbarSpendingPlanRepositoryStub.create.resolves(newSpendingPlan);