From cd8a59bf0225a9a771c3d01559f4bb6d2b29bb07 Mon Sep 17 00:00:00 2001
From: Jeff Huang <jeff@dinngo.co>
Date: Tue, 13 Aug 2024 14:55:55 +0800
Subject: [PATCH] fix: check dest token

---
 src/logics/stargate-v2/configs.ts          | 58 ++++++++++++----------
 src/logics/stargate-v2/logic.swap-token.ts | 34 +++++++++++--
 2 files changed, 62 insertions(+), 30 deletions(-)

diff --git a/src/logics/stargate-v2/configs.ts b/src/logics/stargate-v2/configs.ts
index 466f88a..50923ca 100644
--- a/src/logics/stargate-v2/configs.ts
+++ b/src/logics/stargate-v2/configs.ts
@@ -5,9 +5,9 @@ export enum EndpointId {
   ETHEREUM = 30101,
   BNB = 30102,
   AVALANCHE = 30106,
+  POLYGON = 30109,
   ARBITRUM = 30110,
   OPTIMISM = 30111,
-  POLYGON = 30109,
   METIS = 30151,
   BASE = 30184,
   IOTA = 30284,
@@ -380,32 +380,29 @@ export const configs: Config[] = [
   },
 ];
 
-export const [supportedChainIds, configMap, configMapById, tokensMap, poolConfigMapById, poolConfigMapByToken] =
-  configs.reduce(
-    (accumulator, config) => {
-      accumulator[0].push(config.chainId);
-      accumulator[1][config.chainId] = config;
-      accumulator[2][config.eid] = config;
-      accumulator[3][config.chainId] = new Set();
-      accumulator[4][config.chainId] = {};
-      accumulator[5][config.chainId] = {};
-      for (const pool of config.pools) {
-        accumulator[3][config.chainId].add(pool.token);
-        accumulator[4][config.chainId][pool.id] = { chainId: config.chainId, ...pool };
-        accumulator[5][config.chainId][pool.token.address] = { chainId: config.chainId, ...pool };
-      }
+export const [supportedChainIds, configMap, tokensMap, poolConfigMapById, poolConfigMapByToken] = configs.reduce(
+  (accumulator, config) => {
+    accumulator[0].push(config.chainId);
+    accumulator[1][config.chainId] = config;
+    accumulator[2][config.chainId] = new Set();
+    accumulator[3][config.chainId] = {};
+    accumulator[4][config.chainId] = {};
+    for (const pool of config.pools) {
+      accumulator[2][config.chainId].add(pool.token);
+      accumulator[3][config.chainId][pool.id] = { chainId: config.chainId, ...pool };
+      accumulator[4][config.chainId][pool.token.address] = { chainId: config.chainId, ...pool };
+    }
 
-      return accumulator;
-    },
-    [[], {}, {}, {}, {}, {}] as [
-      number[],
-      Record<number, Config>,
-      Record<number, Config>,
-      Record<number, Set<common.Token>>,
-      Record<number, Record<number, PoolConfig>>,
-      Record<number, Record<string, PoolConfig>>
-    ]
-  );
+    return accumulator;
+  },
+  [[], {}, {}, {}, {}] as [
+    number[],
+    Record<number, Config>,
+    Record<number, Set<common.Token>>,
+    Record<number, Record<number, PoolConfig>>,
+    Record<number, Record<string, PoolConfig>>
+  ]
+);
 
 export function getMarkets(chainId: number) {
   return configMap[chainId].pools;
@@ -419,6 +416,10 @@ export function getTokens(chainId: number) {
   return [...tokensMap[chainId]];
 }
 
+export function getTokenByPoolId(chainId: number, poolId: number) {
+  return poolConfigMapById[chainId][poolId].token;
+}
+
 export function getPoolByTokenAddress(chainId: number, tokenAddress: string) {
   return poolConfigMapByToken[chainId][tokenAddress].address;
 }
@@ -431,3 +432,8 @@ export function getDestChainIds(srcChainId: number, srcToken: common.Token) {
   }
   return [...destChainIds];
 }
+
+export function getDestToken(srcToken: common.Token, destChainId: number) {
+  const srcPoolConfig = poolConfigMapByToken[srcToken.chainId][srcToken.address] ?? [];
+  return getTokenByPoolId(destChainId, srcPoolConfig.id);
+}
diff --git a/src/logics/stargate-v2/logic.swap-token.ts b/src/logics/stargate-v2/logic.swap-token.ts
index b325dc6..a9b5379 100644
--- a/src/logics/stargate-v2/logic.swap-token.ts
+++ b/src/logics/stargate-v2/logic.swap-token.ts
@@ -1,13 +1,20 @@
 import { StargatePool__factory } from './contracts';
 import * as common from '@protocolink/common';
 import * as core from '@protocolink/core';
-import { getDestChainIds, getEndpointId, getMarkets, getPoolByTokenAddress, supportedChainIds } from './configs';
+import {
+  getDestChainIds,
+  getDestToken,
+  getEndpointId,
+  getMarkets,
+  getPoolByTokenAddress,
+  supportedChainIds,
+} from './configs';
 import { getNativeToken } from '@protocolink/common';
 import { utils } from 'ethers';
 
 export type SwapTokenLogicTokenList = {
   srcToken: common.Token;
-  destChainIds: number[];
+  destTokens: common.Token[];
 }[];
 
 export type SwapTokenLogicParams = core.TokenToTokenExactInParams<{
@@ -36,10 +43,16 @@ export class SwapTokenLogic extends core.Logic implements core.LogicBuilderInter
       srcTokens.push(market.token);
     }
 
-    // find destination chain ids
+    // find destination tokens
     for (const srcToken of srcTokens) {
+      const destTokens: common.Token[] = [];
+
       const destChainIds = getDestChainIds(this.chainId, srcToken);
-      tokenList.push({ srcToken, destChainIds });
+      for (const destChainId of destChainIds) {
+        const destToken = getDestToken(srcToken, destChainId);
+        destTokens.push(destToken);
+      }
+      tokenList.push({ srcToken, destTokens });
     }
 
     return tokenList;
@@ -47,6 +60,19 @@ export class SwapTokenLogic extends core.Logic implements core.LogicBuilderInter
 
   public async quote(params: SwapTokenLogicParams) {
     const { input, tokenOut, receiver } = params;
+
+    // check if tokenOut is legit
+    const destToken = getDestToken(input.token, tokenOut.chainId);
+    if (!tokenOut.is(destToken)) {
+      return {
+        input,
+        output: new common.TokenAmount(tokenOut),
+        fee: '0',
+        lzTokenFee: '0',
+        receiver,
+      };
+    }
+
     const dstEid = getEndpointId(tokenOut.chainId);
     const to = utils.hexZeroPad(utils.solidityPack(['address'], [receiver]), 32);
     const poolAddress = getPoolByTokenAddress(this.chainId, input.token.address);