router white list

Swap/flamingo-contract-swap/FlamingoSwapFactory/FlamingoSwapFactoryContract.cs

@@ -220,6 +220,7 @@ private static bool SetFeeTo(byte[] feeTo)
             return true;
         }
 
+
 
         /// <summary>
         /// 获取pair

Swap/flamingo-contract-swap/FlamingoSwapPair/FlamingoSwapPairContract.Admin.cs

@@ -17,12 +17,14 @@ partial class FlamingoSwapPairContract
 
         static readonly byte[] superAdmin = "AZaCs7GwthGy9fku2nFXtbrdKBRmrUQoFP".ToScriptHash();
 
-        ///// <summary>
-        ///// Factory地址
-        ///// </summary>
-        //static readonly byte[] FactoryContract = "32904da4441d544d05074774ade7c891d912f61a".HexToBytes();
+        /// <summary>
+        /// WhiteList 合约地址
+        /// </summary>
+        static readonly byte[] WhiteListContract = "3008f596f4fbdcaf712d6fc0ad2e9a522cc061cf".HexToBytes();
+
 
         const string AdminKey = nameof(superAdmin);
+        private const string WhiteListContractKey = nameof(WhiteListContract);
 
 
         /// <summary>
@@ -49,6 +51,39 @@ public static bool SetAdmin(byte[] admin)
         }
 
 
+        /// <summary>
+        /// 获取WhiteListContract地址
+        /// </summary>
+        /// <returns></returns>
+        public static byte[] GetWhiteListContract()
+        {
+            var whiteList = Storage.Get(WhiteListContractKey);
+            return whiteList.Length == 20 ? whiteList : WhiteListContract;
+        }
+
+        /// <summary>
+        /// 设置WhiteListContract地址
+        /// </summary>
+        /// <param name="whiteList"></param>
+        /// <returns></returns>
+        public static bool SetWhiteListContract(byte[] whiteList)
+        {
+            Assert(whiteList.Length == 20, "WhiteList contract Invalid");
+            Assert(Runtime.CheckWitness(GetAdmin()), "Forbidden");
+            Storage.Put(WhiteListContractKey, whiteList);
+            return true;
+        }
+
+        /// <summary>
+        /// 检查<see cref="callScript"/>是否为router合约
+        /// </summary>
+        /// <param name="callScript"></param>
+        /// <returns></returns>
+        private static bool CheckIsRouter(byte[] callScript)
+        {
+            var whiteList = GetWhiteListContract();
+            return ((Func<string, object[], bool>)whiteList.ToDelegate())("checkRouter", new object[] { callScript });
+        }
 
         #endregion
 

Swap/flamingo-contract-swap/FlamingoSwapPair/FlamingoSwapPairContract.Nep5.cs

@@ -7,6 +7,7 @@
 using System.Threading.Tasks;
 using Neo.SmartContract.Framework;
 using Neo.SmartContract.Framework.Services.Neo;
+using Neo.SmartContract.Framework.Services.System;
 
 namespace FlamingoSwapPair
 {
@@ -70,6 +71,12 @@ private static bool Transfer(byte[] from, byte[] to, BigInteger amount, byte[] c
             Assert(from.Length == 20 && to.Length == 20, "The parameters from and to SHOULD be 20-byte addresses.");
             Assert(amount >= 0, "The parameter amount MUST be greater than 0.");
 
+            var me = ExecutionEngine.ExecutingScriptHash;
+            if (to == me)
+            {
+                Assert(CheckIsRouter(callscript), "Only support transfer to me by Router");
+            }
+
             if (!Runtime.CheckWitness(from) && from.AsBigInteger() != callscript.AsBigInteger())
                 return false;
             StorageMap asset = Storage.CurrentContext.CreateMap(BalanceMapKey);

Swap/flamingo-contract-swap/FlamingoSwapPair/FlamingoSwapPairContract.cs

@@ -10,8 +10,6 @@ namespace FlamingoSwapPair
     partial class FlamingoSwapPairContract : SmartContract
     {
 
-
-
         /// <summary>
         /// https://uniswap.org/docs/v2/protocol-overview/smart-contracts/#minimum-liquidity
         /// </summary>
@@ -124,6 +122,12 @@ public static object Main(string method, object[] args)
 
                 if (method == "setAdmin") return SetAdmin((byte[])args[0]);
 
+                if (method == "getWhiteListContract") return GetWhiteListContract();
+
+                if (method == "setWhiteListContract") return SetWhiteListContract((byte[])args[0]);
+
+                if (method == "checkIsRouter") return CheckIsRouter((byte[])args[0]);
+
                 if (method == "upgrade")
                 {
                     Assert(args.Length == 9, "upgrade: args.Length != 9.");
@@ -144,8 +148,6 @@ public static object Main(string method, object[] args)
         }
 
 
-
-
         #region Swap
 
         /// <summary>
@@ -157,6 +159,8 @@ public static object Main(string method, object[] args)
         /// <param name="toAddress"></param>
         public static bool Swap(byte[] msgSender, BigInteger amount0Out, BigInteger amount1Out, byte[] toAddress)
         {
+            Assert(CheckIsRouter(msgSender), "Only Router Can Swap");
+
             var me = ExecutionEngine.ExecutingScriptHash;
 
             //转出量必需一个为0一个为正数
@@ -218,6 +222,7 @@ public static bool Swap(byte[] msgSender, BigInteger amount0Out, BigInteger amou
         /// <returns></returns>
         public static object Burn(byte[] msgSender, byte[] toAddress)
         {
+            Assert(CheckIsRouter(msgSender),"Only Router Can Burn");
             var me = ExecutionEngine.ExecutingScriptHash;
             var r = ReservePair;
 
@@ -259,13 +264,14 @@ public static object Burn(byte[] msgSender, byte[] toAddress)
 
         /// <summary>
         /// 铸造代币，此方法应该由router在AddLiquidity时调用
-        /// todo:禁止外部直接调用
         /// </summary>
         /// <param name="msgSender"></param>
         /// <param name="toAddress"></param>
         /// <returns>返回本次铸币量</returns>
         public static BigInteger Mint(byte[] msgSender, byte[] toAddress)
         {
+            Assert(CheckIsRouter(msgSender), "Only Router Can Mint");
+
             var me = ExecutionEngine.ExecutingScriptHash;
 
             var r = ReservePair;

Swap/flamingo-contract-swap/FlamingoSwapPairWhiteList/FlamingoSwapPairWhiteList.Admin.cs

@@ -0,0 +1,47 @@
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Text;
+using System.Threading.Tasks;
+using Neo.SmartContract.Framework;
+using Neo.SmartContract.Framework.Services.Neo;
+
+namespace FlamingoSwapPairWhiteList
+{
+    partial class FlamingoSwapPairWhiteList
+    {
+        #region Admin
+
+        static readonly byte[] superAdmin = "AZaCs7GwthGy9fku2nFXtbrdKBRmrUQoFP".ToScriptHash();
+        const string AdminKey = nameof(superAdmin);
+
+
+        /// <summary>
+        /// 获取合约管理员
+        /// </summary>
+        /// <returns></returns>
+        public static byte[] GetAdmin()
+        {
+            var admin = Storage.Get(AdminKey);
+            return admin.Length == 20 ? admin : superAdmin;
+        }
+
+        /// <summary>
+        /// 设置合约管理员
+        /// </summary>
+        /// <param name="admin"></param>
+        /// <returns></returns>
+        public static bool SetAdmin(byte[] admin)
+        {
+            Assert(admin.Length == 20, "NewAdmin Invalid");
+            Assert(Runtime.CheckWitness(GetAdmin()), "Forbidden");
+            Storage.Put(AdminKey, admin);
+            return true;
+        }
+
+
+
+        #endregion
+
+    }
+}

Swap/flamingo-contract-swap/FlamingoSwapPairWhiteList/FlamingoSwapPairWhiteList.Helper.cs

@@ -0,0 +1,40 @@
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Text;
+using System.Threading.Tasks;
+using Neo.SmartContract.Framework;
+using Neo.SmartContract.Framework.Services.Neo;
+
+namespace FlamingoSwapPairWhiteList
+{
+    partial class FlamingoSwapPairWhiteList
+    {
+    /// <summary>
+    /// 断言
+    /// </summary>
+    /// <param name="condition"></param>
+    /// <param name="message"></param>
+    private static void Assert(bool condition, string message)
+    {
+        if (!condition)
+        {
+            Runtime.Notify("Fault:" + message);
+            throw new Exception(message);
+        }
+    }
+
+        ///// <summary>
+        ///// 断言,节约gas
+        ///// </summary>
+        ///// <param name="condition"></param>
+        ///// <param name="message"></param>
+        //[OpCode(OpCode.THROWIFNOT)]
+        //[OpCode(OpCode.DROP)]
+        //private static extern void Assert(bool condition, string message);
+
+
+        [OpCode(OpCode.APPEND)]
+        private static extern void Append<T>(T[] array, T newItem);
+    }
+}

Swap/flamingo-contract-swap/FlamingoSwapPairWhiteList/FlamingoSwapPairWhiteList.cs

@@ -0,0 +1,128 @@
+using System;
+using System.ComponentModel;
+using System.Numerics;
+using Neo.SmartContract.Framework;
+using Neo.SmartContract.Framework.Services.Neo;
+using Neo.SmartContract.Framework.Services.System;
+
+namespace FlamingoSwapPairWhiteList
+{
+    partial class FlamingoSwapPairWhiteList : SmartContract
+    {
+
+        /// <summary>
+        /// 允许向LP合约转入LP token的白名单地址
+        /// </summary>
+        private const string WhiteList = "WhiteList";
+
+
+        #region 通知
+
+        /// <summary>
+        /// params: routerHash
+        /// </summary>
+        [DisplayName("addRouter")]
+        public static event Action<byte[]> onAddRouter;
+
+        /// <summary>
+        /// params: routerHash
+        /// </summary>
+        [DisplayName("removeRouter")]
+        public static event Action<byte[]> onRemoveRouter;
+
+
+        #endregion
+
+        public static object Main(string method, object[] args)
+        {
+            if (Runtime.Trigger == TriggerType.Verification)
+            {
+                return Runtime.CheckWitness(GetAdmin());
+            }
+            else if (Runtime.Trigger == TriggerType.Application)
+            {
+                //合约调用时，等价以太坊的msg.sender
+                //直接调用时，此处为 tx.Script.ToScriptHash();
+                //var msgSender = ExecutionEngine.CallingScriptHash;
+                if (method == "checkRouter") return CheckRouterWhiteList((byte[])args[0]);
+                if (method == "addRouter") return AddRouterWhiteList((byte[])args[0]);
+                if (method == "removeRouter") return RemoveRouterWhiteList((byte[])args[0]);
+                if (method == "getAllRouter") return GetAllRouterWhiteList();
+                if (method == "getAdmin") return GetAdmin();
+                if (method == "setAdmin") return SetAdmin((byte[])args[0]);
+            }
+            return false;
+        }
+
+
+
+        /// <summary>
+        /// 增加router白名单，加过白名单的router才能完成burn
+        /// </summary>
+        /// <param name="router">Nep5 tokenA</param>
+        /// <returns></returns>
+        public static bool AddRouterWhiteList(byte[] router)
+        {
+            Assert(Runtime.CheckWitness(GetAdmin()), "Forbidden");
+            Assert(router.Length == 20, "Invalid Router Address");
+            var key = WhiteList.AsByteArray().Concat(router);
+            Storage.Put(key, 1);
+            onAddRouter(router);
+            return true;
+        }
+
+
+        /// <summary>
+        /// 移除router白名单，加过白名单的router才能完成burn
+        /// </summary>
+        /// <param name="router">Nep5 tokenA</param>
+        /// <returns></returns>
+        public static bool RemoveRouterWhiteList(byte[] router)
+        {
+            Assert(Runtime.CheckWitness(GetAdmin()), "Forbidden");
+            Assert(router.Length == 20, "Invalid Router Address");
+            var key = WhiteList.AsByteArray().Concat(router);
+            Storage.Delete(key);
+            onRemoveRouter(router);
+            return true;
+        }
+
+
+        /// <summary>
+        /// 检查router白名单
+        /// </summary>
+        /// <param name="router">Nep5 tokenA</param>
+        /// <returns></returns>
+        public static bool CheckRouterWhiteList(byte[] router)
+        {
+            if (router.Length != 20)
+            {
+                return false;
+            }
+            var key = WhiteList.AsByteArray().Concat(router);
+            var value = Storage.Get(key).AsBigInteger();
+            return value > 0;
+        }
+
+
+        /// <summary>
+        /// 查询router白名单
+        /// </summary>
+        /// <returns></returns>
+        public static byte[][] GetAllRouterWhiteList()
+        {
+            var iterator = Storage.Find(WhiteList);
+            var result = new byte[0][];
+            while (iterator.Next())
+            {
+                if (iterator.Value.AsBigInteger() > 0)
+                {
+                    var router = iterator.Key.AsByteArray().Last(20);
+                    Append(result, router);
+                }
+            }
+            return result;
+        }
+
+    }
+}

Swap/flamingo-contract-swap/FlamingoSwapPairWhiteList/FlamingoSwapPairWhiteList.csproj

@@ -0,0 +1,15 @@
+<Project Sdk="Microsoft.NET.Sdk">
+
+  <PropertyGroup>
+    <TargetFramework>netcoreapp2.1</TargetFramework>
+    <CopyLocalLockFileAssemblies>true</CopyLocalLockFileAssemblies>
+  </PropertyGroup>
+
+  <ItemGroup>
+    <PackageReference Include="Neo.SmartContract.Framework" Version="2.10.0.0" />
+  </ItemGroup>
+
+  <Target Name="NeoBuild" AfterTargets="AfterBuild">
+    <Exec WorkingDirectory="$(OutputPath)" Command="neon $(AssemblyName).dll --compatible"></Exec>
+  </Target>
+</Project>