From d95f1488cc4363e42659ea1b8114e6ad71b99197 Mon Sep 17 00:00:00 2001
From: Krasimir Raykov <raykov@weichain.tech>
Date: Tue, 15 Feb 2022 09:59:04 +0200
Subject: [PATCH] feat: allow registering a referral through the proxy

---
 contracts/proxy/LiqualityRouter.sol | 36 +++++++++++++++++++++++++----
 1 file changed, 32 insertions(+), 4 deletions(-)

diff --git a/contracts/proxy/LiqualityRouter.sol b/contracts/proxy/LiqualityRouter.sol
index c78bafc..d12d97d 100644
--- a/contracts/proxy/LiqualityRouter.sol
+++ b/contracts/proxy/LiqualityRouter.sol
@@ -5,27 +5,55 @@ import "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
 import "./interfaces/ILiqualityRouter.sol";
 import "./interfaces/ILiqualityProxy.sol";
 import "./LibTransfer.sol";
+import "../referrals/interfaces/IReferralRegistry.sol";
 
 contract LiqualityRouter is ILiqualityRouter {
     using LibTransfer for address;
     using SafeERC20 for IERC20;
 
     ILiqualityProxy public liqualityProxy;
+    IReferralRegistry public referralRegistry;
 
     constructor(address _liqualityProxy) {
         liqualityProxy = ILiqualityProxy(_liqualityProxy);
     }
 
-    function route(
+    function routeWithReferral(
         address target,
         address tokenFrom,
+        address referrer,
         uint256 amount,
         bytes calldata data,
         FeeData[] calldata fees
     ) external payable {
-        uint256 totalFee = 0;
         address user = msg.sender;
+        // handle referrals
+        Referral memory referral = referralRegistry.getReferral(user);
+        if (referral.referrer != address(0x0)) {
+            referralRegistry.registerReferral(referrer, user);
+        }
+        _route(target, user, tokenFrom, amount, data, fees);
+    }
 
+    function route(
+        address target,
+        address tokenFrom,
+        uint256 amount,
+        bytes calldata data,
+        FeeData[] calldata fees
+    ) external payable {
+        _route(target, msg.sender, tokenFrom, amount, data, fees);
+    }
+
+    function _route(
+        address target,
+        address from,
+        address tokenFrom,
+        uint256 amount,
+        bytes calldata data,
+        FeeData[] calldata fees
+    ) internal {
+        uint256 totalFee = 0;
         // handle ETH fees
         if (tokenFrom == address(0)) {
             for (uint256 i = 0; i < fees.length; i++) {
@@ -35,9 +63,9 @@ contract LiqualityRouter is ILiqualityRouter {
         }
         // handle ERC20 fees
         else {
-            IERC20(tokenFrom).safeTransferFrom(user, address(liqualityProxy), amount);
+            IERC20(tokenFrom).safeTransferFrom(from, address(liqualityProxy), amount);
             for (uint256 i = 0; i < fees.length; i++) {
-                IERC20(tokenFrom).safeTransferFrom(user, fees[i].account, fees[i].fee);
+                IERC20(tokenFrom).safeTransferFrom(from, fees[i].account, fees[i].fee);
             }
         }