Refactor a bit more

celo-org · Jan 10, 2024 · 9c78d2e · 9c78d2e
1 parent 6e703c3
commit 9c78d2e
Showing 1 changed file with 30 additions and 26 deletions.
diff --git a/contracts/fee_currencies.go b/contracts/fee_currencies.go
@@ -97,49 +97,53 @@ func CreditFees(
 		return err
 	}
 
-	leftoverGas, err = transferErc20(evm, abi, caller, feeCurrency, &baseFeeReceiver, baseFee, leftoverGas)
+	leftoverGas, err = transferErc20(evm, abi, caller, feeCurrency, baseFeeReceiver, baseFee, leftoverGas)
 	if err != nil {
 		return fmt.Errorf("base fee: %w", err)
 	}
 
 	// If the tip is non-zero, but the coinbase of the miner is not set, send the tip back to the tx sender
-	if feeTip.Cmp(common.Big0) == 1 {
-		if tipReceiver != common.ZeroAddress {
-			leftoverGas, err = transferErc20(evm, abi, caller, feeCurrency, &tipReceiver, feeTip, leftoverGas)
-			if err != nil {
-				return fmt.Errorf("fee tip: %w", err)
-			}
-		} else {
-			refund.Add(refund, feeTip)
-		}
+	unusedFee, leftoverGas, err := creditFee(evm, abi, caller, feeCurrency, tipReceiver, feeTip, leftoverGas, "fee tip")
+	if err != nil {
+		return err
 	}
+	refund.Add(refund, unusedFee)
 
-	// The data fee is `nil` when the node is not running in rollup mode
 	// If the data fee is non-zero, but the data fee receiver is not set, send the tip back to the tx sender
-	if l1DataFee != nil && l1DataFee.Cmp(common.Big0) == 1 {
-		if l1DataFeeReceiver != common.ZeroAddress {
-			leftoverGas, err = transferErc20(evm, abi, caller, feeCurrency, &l1DataFeeReceiver, l1DataFee, leftoverGas)
-			if err != nil {
-				return fmt.Errorf("l1 data fee: %w", err)
-			}
-		} else {
-			refund.Add(refund, l1DataFee)
-		}
+	unusedFee, leftoverGas, err = creditFee(evm, abi, caller, feeCurrency, l1DataFeeReceiver, l1DataFee, leftoverGas, "l1 data fee")
+	if err != nil {
+		return err
 	}
+	refund.Add(refund, unusedFee)
 
-	if refund.Cmp(common.Big0) == 1 {
-		leftoverGas, err = transferErc20(evm, abi, caller, feeCurrency, &txSender, refund, leftoverGas)
-		if err != nil {
-			return fmt.Errorf("refund: %w", err)
-		}
+	unusedFee, leftoverGas, err = creditFee(evm, abi, caller, feeCurrency, txSender, refund, leftoverGas, "refund")
+	if err != nil {
+		return err
+	}
+	if unusedFee == nil || unusedFee.Cmp(common.Big0) != 0 {
+		return errors.New("could not refund remaining fees to sender")
 	}
 
 	gasUsed := maxGasForCreditGasFeesTransactions - leftoverGas
 	log.Trace("creditFees called", "feeCurrency", *feeCurrency, "gasUsed", gasUsed)
 	return nil
 }
 
-func transferErc20(evm *vm.EVM, abi *abi.ABI, caller vm.AccountRef, feeCurrency, receiver *common.Address, amount *big.Int, gasLimit uint64) (uint64, error) {
+func creditFee(evm *vm.EVM, abi *abi.ABI, caller vm.AccountRef, feeCurrency *common.Address, receiver common.Address, amount *big.Int, gasLimit uint64, action string) (*big.Int, uint64, error) {
+	if amount != nil && amount.Cmp(common.Big0) == 1 {
+		if receiver != common.ZeroAddress {
+			leftoverGas, err := transferErc20(evm, abi, caller, feeCurrency, receiver, amount, gasLimit)
+			if err != nil {
+				return common.Big0, leftoverGas, fmt.Errorf("%s: %w", action, err)
+			}
+		} else {
+			return amount, gasLimit, nil
+		}
+	}
+	return common.Big0, gasLimit, nil
+}
+
+func transferErc20(evm *vm.EVM, abi *abi.ABI, caller vm.AccountRef, feeCurrency *common.Address, receiver common.Address, amount *big.Int, gasLimit uint64) (uint64, error) {
 	// Solidity: function transfer(address to, uint256 amount) returns(bool)
 	transferData, err := abi.Pack("transfer", receiver, amount)
 	if err != nil {