From cc06a33d177bb217b4c9686073242cbf14f16ac3 Mon Sep 17 00:00:00 2001
From: Thrilok Kumar <thrilok2000@gmail.com>
Date: Wed, 16 Mar 2022 16:30:38 +0400
Subject: [PATCH] Fixed polygon Aave v2-to-v3 connector

---
 .../aave/v2-to-v3-import/helpers.sol          | 15 +++++++++++++++
 .../connectors/aave/v2-to-v3-import/main.sol  | 19 +++++++++++--------
 2 files changed, 26 insertions(+), 8 deletions(-)

diff --git a/contracts/polygon/connectors/aave/v2-to-v3-import/helpers.sol b/contracts/polygon/connectors/aave/v2-to-v3-import/helpers.sol
index 5da70943..50b5381c 100644
--- a/contracts/polygon/connectors/aave/v2-to-v3-import/helpers.sol
+++ b/contracts/polygon/connectors/aave/v2-to-v3-import/helpers.sol
@@ -44,7 +44,9 @@ abstract contract Helper is DSMath, Basic {
 
 	struct ImportData {
 		address[] _supplyTokens;
+		address[] _supplyTokensV3;
 		address[] _borrowTokens;
+		address[] _borrowTokensV3;
 		ATokenV2Interface[] aTokens;
 		uint256[] supplyAmts;
 		uint256[] variableBorrowAmts;
@@ -58,7 +60,9 @@ abstract contract Helper is DSMath, Basic {
 
 	struct ImportInputData {
 		address[] supplyTokens;
+		address[] supplyTokensV3;
 		address[] borrowTokens;
+		address[] borrowTokensV3;
 		bool convertStable;
 		uint256[] flashLoanFees;
 	}
@@ -92,6 +96,7 @@ contract _AaveHelper is Helper {
 	) internal returns (ImportData memory) {
 		if (inputData.borrowTokens.length > 0) {
 			data._borrowTokens = new address[](inputData.borrowTokens.length);
+			data._borrowTokensV3 = new address[](inputData.borrowTokensV3.length);
 			data.variableBorrowAmts = new uint256[](
 				inputData.borrowTokens.length
 			);
@@ -114,7 +119,11 @@ contract _AaveHelper is Helper {
 				address _token = inputData.borrowTokens[i] == maticAddr
 					? wmaticAddr
 					: inputData.borrowTokens[i];
+				address _tokenV3 = inputData.borrowTokensV3[i] == maticAddr
+					? wmaticAddr
+					: inputData.borrowTokensV3[i];
 				data._borrowTokens[i] = _token;
+				data._borrowTokensV3[i] = _tokenV3;
 
 				(
 					data.stableBorrowAmts[i],
@@ -159,6 +168,7 @@ contract _AaveHelper is Helper {
 	) internal view returns (ImportData memory) {
 		data.supplyAmts = new uint256[](inputData.supplyTokens.length);
 		data._supplyTokens = new address[](inputData.supplyTokens.length);
+		data._supplyTokensV3 = new address[](inputData.supplyTokensV3.length);
 		data.aTokens = new ATokenV2Interface[](inputData.supplyTokens.length);
 
 		for (uint256 i = 0; i < inputData.supplyTokens.length; i++) {
@@ -175,10 +185,15 @@ contract _AaveHelper is Helper {
 			address _token = inputData.supplyTokens[i] == maticAddr
 				? wmaticAddr
 				: inputData.supplyTokens[i];
+			address _tokenV3 = inputData.supplyTokensV3[i] == maticAddr
+				? wmaticAddr
+				: inputData.supplyTokensV3[i];
+
 			(address _aToken, , ) = aaveV2Data.getReserveTokensAddresses(
 				_token
 			);
 			data._supplyTokens[i] = _token;
+			data._supplyTokensV3[i] = _tokenV3;
 			data.aTokens[i] = ATokenV2Interface(_aToken);
 			data.supplyAmts[i] = data.aTokens[i].balanceOf(userAccount);
 		}
diff --git a/contracts/polygon/connectors/aave/v2-to-v3-import/main.sol b/contracts/polygon/connectors/aave/v2-to-v3-import/main.sol
index 06a7d43b..b5472467 100644
--- a/contracts/polygon/connectors/aave/v2-to-v3-import/main.sol
+++ b/contracts/polygon/connectors/aave/v2-to-v3-import/main.sol
@@ -24,6 +24,9 @@ contract _AaveV2ToV3MigrationResolver is _AaveHelper {
 		ImportInputData memory inputData,
 		bool doImport
 	) internal returns (string memory _eventName, bytes memory _eventParam) {
+		require(inputData.borrowTokens.length == inputData.borrowTokensV3.length, "borrow-tokens-len-not-same");
+		require(inputData.supplyTokens.length == inputData.supplyTokensV3.length, "supply-tokens-len-not-same");
+
 		if (doImport) {
 			// check only when we are importing from user's address
 			require(
@@ -81,31 +84,31 @@ contract _AaveV2ToV3MigrationResolver is _AaveHelper {
 		);
 		// deposit tokens in v3
 		_depositTokensV3(
-			data._supplyTokens.length,
+			data._supplyTokensV3.length,
 			aaveV3,
 			data.supplyAmts,
-			data._supplyTokens
+			data._supplyTokensV3
 		);
 
 		// borrow assets in aave v3 after migrating position
 		if (data.convertStable) {
 			_BorrowVariableV3(
-				data._borrowTokens.length,
+				data._borrowTokensV3.length,
 				aaveV3,
-				data._borrowTokens,
+				data._borrowTokensV3,
 				data.totalBorrowAmtsWithFee
 			);
 		} else {
 			_BorrowStableV3(
-				data._borrowTokens.length,
+				data._borrowTokensV3.length,
 				aaveV3,
-				data._borrowTokens,
+				data._borrowTokensV3,
 				data.stableBorrowAmtsWithFee
 			);
 			_BorrowVariableV3(
-				data._borrowTokens.length,
+				data._borrowTokensV3.length,
 				aaveV3,
-				data._borrowTokens,
+				data._borrowTokensV3,
 				data.variableBorrowAmtsWithFee
 			);
 		}