From 06f93ff029fc3d6226421b6f15e9483dfd203464 Mon Sep 17 00:00:00 2001
From: Shivaansh Kapoor <shivaanshkapoor02@gmail.com>
Date: Mon, 9 Sep 2024 12:07:16 +0400
Subject: [PATCH] tokenDirectory to support multiple tokens

---
 .../common/BiconomyTokenPaymasterErrors.sol   |  2 +-
 .../interfaces/IBiconomyTokenPaymaster.sol    | 11 ++++++
 contracts/token/BiconomyTokenPaymaster.sol    | 38 +++++++++++++------
 3 files changed, 38 insertions(+), 13 deletions(-)

diff --git a/contracts/common/BiconomyTokenPaymasterErrors.sol b/contracts/common/BiconomyTokenPaymasterErrors.sol
index 9b0c365..48de59d 100644
--- a/contracts/common/BiconomyTokenPaymasterErrors.sol
+++ b/contracts/common/BiconomyTokenPaymasterErrors.sol
@@ -35,7 +35,7 @@ contract BiconomyTokenPaymasterErrors {
     /**
      * @notice Throws when each token doesnt have a corresponding oracle
      */
-    error TokensAndOraclesLengthMismatch();
+    error TokensAndInfoLengthMismatch();
 
     /**
      * @notice Throws when oracle returns invalid price
diff --git a/contracts/interfaces/IBiconomyTokenPaymaster.sol b/contracts/interfaces/IBiconomyTokenPaymaster.sol
index 58157f9..724e2b4 100644
--- a/contracts/interfaces/IBiconomyTokenPaymaster.sol
+++ b/contracts/interfaces/IBiconomyTokenPaymaster.sol
@@ -1,7 +1,15 @@
 // SPDX-License-Identifier: GPL-3.0
 pragma solidity ^0.8.26;
 
+import { IOracle } from "./oracles/IOracle.sol";
+
 interface IBiconomyTokenPaymaster {
+    // Struct for storing information about the token
+    struct TokenInfo {
+        IOracle oracle;
+        uint8 decimals;
+    }
+
     event UnaccountedGasChanged(uint256 indexed oldValue, uint256 indexed newValue);
     event FixedDynamicAdjustmentChanged(uint256 indexed oldValue, uint256 indexed newValue);
     event FeeCollectorChanged(address indexed oldFeeCollector, address indexed newFeeCollector, address indexed actor);
@@ -11,10 +19,13 @@ interface IBiconomyTokenPaymaster {
     event DynamicAdjustmentCollected(address indexed paymasterId, uint256 indexed dynamicAdjustment);
     event Received(address indexed sender, uint256 value);
     event TokensWithdrawn(address indexed token, address indexed to, uint256 indexed amount, address actor);
+    event UpdatedTokenDirectory(address indexed tokenAddress, IOracle indexed oracle, uint8 decimals);
 
     function setFeeCollector(address _newFeeCollector) external payable;
 
     function setUnaccountedGas(uint256 value) external payable;
 
     function setDynamicAdjustment(uint256 _newUnaccountedGas) external payable;
+
+    function setTokenInfo(address _tokenAddress, IOracle _oracle, uint8 _decimals) external payable;
 }
diff --git a/contracts/token/BiconomyTokenPaymaster.sol b/contracts/token/BiconomyTokenPaymaster.sol
index df36782..7a62cb0 100644
--- a/contracts/token/BiconomyTokenPaymaster.sol
+++ b/contracts/token/BiconomyTokenPaymaster.sol
@@ -30,11 +30,6 @@ contract BiconomyTokenPaymaster is
 {
     using UserOperationLib for PackedUserOperation;
 
-    struct TokenInfo {
-        IOracle oracle;
-        uint8 decimals;
-    }
-
     // State variables
     address public feeCollector;
     uint256 public unaccountedGas;
@@ -54,6 +49,7 @@ contract BiconomyTokenPaymaster is
         uint256 _dynamicAdjustment,
         IOracle _nativeOracle,
         address[] memory _tokens, // Array of token addresses
+        uint8[] memory _decimals, // Array of corresponding token decimals
         IOracle[] memory _oracles // Array of corresponding oracle addresses
     )
         BasePaymaster(_owner, _entryPoint)
@@ -62,8 +58,8 @@ contract BiconomyTokenPaymaster is
             revert UnaccountedGasTooHigh();
         } else if (_dynamicAdjustment > MAX_DYNAMIC_ADJUSTMENT || _dynamicAdjustment == 0) {
             revert InvalidDynamicAdjustment();
-        } else if (_tokens.length != _oracles.length) {
-            revert TokensAndOraclesLengthMismatch();
+        } else if (_tokens.length != _oracles.length || _tokens.length != _decimals.length) {
+            revert TokensAndInfoLengthMismatch();
         }
         assembly ("memory-safe") {
             sstore(feeCollector.slot, address()) // initialize fee collector to this contract
@@ -74,7 +70,7 @@ contract BiconomyTokenPaymaster is
 
         // Populate the tokenToOracle mapping
         for (uint256 i = 0; i < _tokens.length; i++) {
-            tokenDirectory[_tokens[i]] = TokenInfo(_oracles[i], ERC20(_tokens[i]).decimals());
+            tokenDirectory[_tokens[i]] = TokenInfo(_oracles[i], _decimals[i]);
         }
     }
 
@@ -200,7 +196,7 @@ contract BiconomyTokenPaymaster is
     }
 
     /**
-     * @dev Set a new unaccountedEPGasOverhead value.
+     * @dev Set a new dynamicAdjustment value.
      * @param _newDynamicAdjustment The new value to be set as the unaccounted gas value
      * @notice only to be called by the owner of the contract.
      */
@@ -215,6 +211,26 @@ contract BiconomyTokenPaymaster is
         emit FixedDynamicAdjustmentChanged(oldDynamicAdjustment, _newDynamicAdjustment);
     }
 
+    /**
+     * @dev Set or update a TokenInfo entry in the tokenDirectory mapping.
+     * @param _tokenAddress The new value to be set as the unaccounted gas value
+     * @param _oracle The new value to be set as the unaccounted gas value
+     * @param _decimals The new value to be set as the unaccounted gas value
+     * @notice only to be called by the owner of the contract.
+     */
+    function setTokenInfo(
+        address _tokenAddress,
+        IOracle _oracle,
+        uint8 _decimals
+    )
+        external
+        payable
+        override
+        onlyOwner
+    {
+        tokenDirectory[_tokenAddress] = TokenInfo(_oracle, _decimals);
+    }
+
     /**
      * @dev Validate a user operation.
      * This method is abstract in BasePaymaster and must be implemented in derived contracts.
@@ -230,9 +246,7 @@ contract BiconomyTokenPaymaster is
         internal
         override
         returns (bytes memory context, uint256 validationData)
-    {
-        
-    }
+    { }
 
     /**
      * @dev Post-operation handler.