From 18e93461a4d343255d92fd4b2fa3e84b31e2fa13 Mon Sep 17 00:00:00 2001
From: Purva-Chaudhari <48947123+Purva-Chaudhari@users.noreply.github.com>
Date: Thu, 19 Dec 2024 17:51:00 +0530
Subject: [PATCH] [FEA-792] added multiplier (#127)

* added multiplier for flightclass
---
 plume/src/Faucet.sol | 67 +++++++++++++++++++++++++++++++++++++++++---
 1 file changed, 63 insertions(+), 4 deletions(-)

diff --git a/plume/src/Faucet.sol b/plume/src/Faucet.sol
index 67633c5..f97d2ca 100644
--- a/plume/src/Faucet.sol
+++ b/plume/src/Faucet.sol
@@ -86,6 +86,9 @@ contract Faucet is Initializable, UUPSUpgradeable {
     /// @notice Indicates a failure because the address is invalid
     error InvalidAddress();
 
+    /// @notice Indicates a failure because the flightClass is invalid
+    error InvalidFlightClass(uint256 flightClass);
+
     /**
      * @notice Indicates a failure because the sender is not authorized to perform the action
      * @param sender Address of the sender that is not authorized
@@ -187,6 +190,7 @@ contract Faucet is Initializable, UUPSUpgradeable {
     /**
      * @notice Get tokens from the faucet
      * @param token Name of the token requested
+     * @param flightClass User's flight class
      * @param salt Random value to prevent replay attacks
      * @param signature Signature of the message signed by the owner
      */
@@ -198,11 +202,14 @@ contract Faucet is Initializable, UUPSUpgradeable {
     ) external onlySignedByOwner(token, flightClass, salt, signature) {
         FaucetStorage storage $ = _getFaucetStorage();
         address tokenAddress = $.tokens[token];
-        uint256 amount = $.dripAmounts[tokenAddress];
-        if (tokenAddress == address(0) || amount == 0) {
+        uint256 baseAmount = $.dripAmounts[tokenAddress];
+
+        if (tokenAddress == address(0) || baseAmount == 0) {
             revert InvalidToken();
         }
 
+        uint256 amount = _calculateDripAmount(baseAmount, flightClass);
+
         if (tokenAddress == ETH_ADDRESS) {
             if (address(this).balance < amount) {
                 revert InsufficientBalance(amount, token);
@@ -253,6 +260,40 @@ contract Faucet is Initializable, UUPSUpgradeable {
         emit Withdrawn(recipient, amount, token);
     }
 
+    /**
+     * @notice Calculate the amount of tokens to mint based on the base amount and flight class
+     * @dev Internal function to calculate the drip amount based on the flight class multiplier
+     *      Flight classes correspond to:
+     *        - Class 1: Economy
+     *        - Class 2: Plus
+     *        - Class 3: Premium
+     *        - Class 4: Business
+     *        - Class 5: First
+     *        - Class 6: Private
+     * @param baseAmount Base amount of tokens to mint
+     * @param flightClass User flight class
+     */
+    function _calculateDripAmount(uint256 baseAmount, uint256 flightClass) internal pure returns (uint256) {
+        uint256 multiplier;
+        if (flightClass == 1) {
+            multiplier = 1; // 1x
+        } else if (flightClass == 2) {
+            multiplier = 11; // 1.1x (scaled by 10)
+        } else if (flightClass == 3) {
+            multiplier = 125; // 1.25x (scaled by 100)
+        } else if (flightClass == 4) {
+            multiplier = 200; // 2x (scaled by 100)
+        } else if (flightClass == 5) {
+            multiplier = 300; // 3x (scaled by 100)
+        } else if (flightClass == 6) {
+            multiplier = 500; // 5x (scaled by 100)
+        } else {
+            revert InvalidFlightClass(flightClass);
+        }
+
+        return (baseAmount * multiplier) / 100; // Normalize for scaling
+    }
+
     /**
      * @notice Set ownership of the faucet contract to the given address
      * @dev Only the owner can call this function
@@ -302,9 +343,9 @@ contract Faucet is Initializable, UUPSUpgradeable {
     }
 
     /**
-     * @notice Get the amount of tokens to mint per faucet call for the given token
+     * @notice Get the base amount of tokens to mint per faucet call for the given token
      * @param token Name of the token to get the amount for
-     * @return dripAmount Amount of tokens to mint per faucet call
+     * @return dripAmount Base amount of tokens to mint per faucet call
      */
     function getDripAmount(
         string calldata token
@@ -313,6 +354,24 @@ contract Faucet is Initializable, UUPSUpgradeable {
         return $.dripAmounts[$.tokens[token]];
     }
 
+    /**
+     * @notice Get the amount of tokens to mint per user call for the given token
+     * @param token Name of the token to get the amount for
+     * @param flightClass User's flight class
+     * @return dripAmount Amount of tokens to mint per faucet call
+     */
+    function getDripAmount(string calldata token, uint256 flightClass) public view returns (uint256 dripAmount) {
+        FaucetStorage storage $ = _getFaucetStorage();
+        address tokenAddress = $.tokens[token];
+        uint256 baseAmount = $.dripAmounts[tokenAddress];
+
+        if (tokenAddress == address(0) || baseAmount == 0) {
+            revert InvalidToken();
+        }
+
+        return _calculateDripAmount(baseAmount, flightClass);
+    }
+
     /**
      * @notice Get the address of the given token
      * @param token Name of the token to get the address for