Fix overflow/underflow issue in math calc

Signed-off-by: Penghui Liao <[email protected]>
giftedchen · Apr 9, 2021 · 29ae8b3 · 29ae8b3
1 parent ce0f2d8
commit 29ae8b3
Show file tree

Hide file tree

Showing 9 changed files with 140 additions and 55 deletions.
diff --git a/.prettierrc b/.prettierrc
@@ -1,15 +1,4 @@
 {
-  "overrides": [
-    {
-      "files": "*.sol",
-      "options": {
-        "printWidth": 120,
-        "tabWidth": 4,
-        "useTabs": false,
-        "singleQuote": true,
-        "bracketSpacing": false,
-        "explicitTypes": "always"
-      }
-    }
-  ]
+  "singleQuote": true,
+  "printWidth": 120
 }
diff --git a/contracts/FlashBot.sol b/contracts/FlashBot.sol
@@ -236,8 +236,44 @@ contract FlashBot is Ownable {
 
     /// @dev calculate the maximum base asset amount to borrow in order to get maximum profit during arbitrage
     function calcBorrowAmount(OrderedReserves memory reserves) internal pure returns (uint256 amount) {
+        // we can't use a1,b1,a2,b2 directly, because it will result overflow/underflow of the intermediate result
+        // so we:
+        //    1. divide all the numbers by d to prevent from overflow/underflow
+        //    2. calculate the result use above numbers
+        //    3. multiple d with the result to get the final result
+        // note this workaround is only suitable for ERC20 token with 18 decimals, which I believe most tokens do
+
+        uint256 min1 = reserves.a1 < reserves.b1 ? reserves.a1 : reserves.b1;
+        uint256 min2 = reserves.a2 < reserves.b2 ? reserves.a2 : reserves.b2;
+        uint256 min = min1 < min2 ? min1 : min2;
+
+        uint256 d;
+        if (min > 1e24) {
+            d = 1e20;
+        } else if (min > 1e23) {
+            d = 1e19;
+        } else if (min > 1e22) {
+            d = 1e18;
+        } else if (min > 1e21) {
+            d = 1e17;
+        } else if (min > 1e20) {
+            d = 1e16;
+        } else if (min > 1e19) {
+            d = 1e15;
+        } else if (min > 1e18) {
+            d = 1e14;
+        } else if (min > 1e17) {
+            d = 1e13;
+        } else if (min > 1e16) {
+            d = 1e12;
+        } else if (min > 1e15) {
+            d = 1e11;
+        } else {
+            d = 1e10;
+        }
+
         (int256 a1, int256 a2, int256 b1, int256 b2) =
-            (int256(reserves.a1), int256(reserves.a2), int256(reserves.b1), int256(reserves.b2));
+            (int256(reserves.a1 / d), int256(reserves.a2 / d), int256(reserves.b1 / d), int256(reserves.b2 / d));
 
         int256 a = a1 * b1 - a2 * b1;
         int256 b = 2 * b1 * b2 * (a1 + a2);
@@ -246,8 +282,8 @@ contract FlashBot is Ownable {
         (int256 x1, int256 x2) = calcSolutionForQuadratic(a, b, c);
 
         // 0 < x < b1 and 0 < x < b2
-        assert((x1 > 0 && x1 < b1 && x1 < b2) || (x1 > 0 && x1 < b1 && x1 < b2));
-        amount = (x1 > 0 && x1 < b1 && x1 < b2) ? uint256(x1) : uint256(x2);
+        require((x1 > 0 && x1 < b1 && x1 < b2) || (x1 > 0 && x1 < b1 && x1 < b2), "Wrong input order");
+        amount = (x1 > 0 && x1 < b1 && x1 < b2) ? uint256(x1) * d : uint256(x2) * d;
     }
 
     /// @dev find solution of quadratic equation: ax^2 + bx + c = 0, only return the positive solution
@@ -270,7 +306,7 @@ contract FlashBot is Ownable {
         assert(n > 1);
 
         // The scale factor is a crude way to turn everything into integer calcs.
-        // Actually do (n * 10 ^ 6) ^ (1/2)
+        // Actually do (n * 10 ^ 4) ^ (1/2)
         uint256 _n = n * 10**6;
         uint256 c = _n;
         res = _n;

diff --git a/contracts/test/TestMath.sol b/contracts/test/TestMath.sol
@@ -6,7 +6,7 @@ import '../FlashBot.sol';
 contract TestMath is FlashBot {
     constructor() FlashBot(address(1)) {}
 
-    function _calcBorrowAmount(OrderedReserves memory reserves) internal pure returns (uint256) {
+    function _calcBorrowAmount(OrderedReserves memory reserves) public pure returns (uint256) {
         return calcBorrowAmount(reserves);
     }
 

diff --git a/package.json b/package.json
@@ -9,12 +9,14 @@
     "@typechain/ethers-v5": "^6.0.5",
     "@typechain/hardhat": "^1.0.1",
     "@types/chai": "^4.2.16",
+    "@types/lodash": "^4.14.168",
     "@types/mocha": "^8.2.2",
     "@types/node": "^14.14.37",
     "chai": "^4.3.4",
     "ethereum-waffle": "^3.0.0",
     "ethers": "^5.0.0",
     "hardhat": "^2.1.2",
+    "lodash": "^4.17.21",
     "prettier": "^2.2.1",
     "prettier-plugin-solidity": "^1.0.0-beta.7",
     "ts-generator": "^0.1.1",

diff --git a/test/AccessControlTest.ts b/test/AccessControlTest.ts
@@ -1,28 +1,28 @@
-import { expect } from "chai";
-import { ethers, waffle } from "hardhat";
-import { TestERC20 } from "../typechain/TestERC20";
-import { FlashBot } from "../typechain/FlashBot";
+import { expect } from 'chai';
+import { ethers, waffle } from 'hardhat';
+import { TestERC20 } from '../typechain/TestERC20';
+import { FlashBot } from '../typechain/FlashBot';
 
-import { flashBotFixture } from "./shared/fixtures";
+import { flashBotFixture } from './shared/fixtures';
 
-describe("FlashBot access control", () => {
+describe('FlashBot access control', () => {
   let weth: TestERC20;
   let flashBot: FlashBot;
 
   beforeEach(async () => {
     ({ weth, flashBot } = await waffle.loadFixture(flashBotFixture));
   });
 
-  it("Should set owner to deployer", async () => {
+  it('Should set owner to deployer', async () => {
     const [owner] = await ethers.getSigners();
     const fbOwner = await flashBot.owner();
     expect(fbOwner).to.be.equal(owner.address);
   });
 
-  it("Should be receivable", async () => {
+  it('Should be receivable', async () => {
     const [owner] = await ethers.getSigners();
 
-    const amount = ethers.utils.parseEther("5.1");
+    const amount = ethers.utils.parseEther('5.1');
     await owner.sendTransaction({
       to: flashBot.address,
       value: amount,
@@ -32,24 +32,24 @@ describe("FlashBot access control", () => {
     expect(balance).to.be.eq(amount);
   });
 
-  it("Should be withdrawable", async () => {
+  it('Should be withdrawable', async () => {
     const [owner, addr1] = await ethers.getSigners();
 
-    const amount = ethers.utils.parseEther("5.1");
+    const amount = ethers.utils.parseEther('5.1');
     await addr1.sendTransaction({
       to: flashBot.address,
       value: amount,
     });
 
-    const wethAmount = ethers.utils.parseEther("100.1");
+    const wethAmount = ethers.utils.parseEther('100.1');
     await weth.mint(flashBot.address, wethAmount);
 
     const balanceBefore = await owner.getBalance();
     const wethBlanceBefore = await weth.balanceOf(owner.address);
 
     // let addr1 withdraw so the gas not spend on owner
     expect(await flashBot.connect(addr1).withdraw())
-      .to.emit(flashBot, "Withdrawn")
+      .to.emit(flashBot, 'Withdrawn')
       .withArgs(owner.address, amount);
 
     const balanceAfter = await owner.getBalance();

diff --git a/test/MathTest.ts b/test/MathTest.ts
@@ -1,22 +1,81 @@
-import { ethers } from "hardhat";
-import { TestMath } from "../typechain/TestMath";
-
-import { expect } from "./shared/expect";
+import lodash from 'lodash';
+import { expect } from 'chai';
+import { ethers } from 'hardhat';
+import { TestMath } from '../typechain/TestMath';
 
 const { BigNumber } = ethers;
 
-describe("MathTest", () => {
+describe('MathTest', () => {
   let math: TestMath;
 
   beforeEach(async () => {
-    const mathTestFactory = await ethers.getContractFactory("TestMath");
+    const mathTestFactory = await ethers.getContractFactory('TestMath');
     math = (await mathTestFactory.deploy()) as TestMath;
   });
 
-  it("Should calculate square root corectly", async () => {
-    const input = BigNumber.from("100");
-    const res = await math._sqrt(input);
-    console.log(res.toString());
-    expect(res).to.be.eq(BigNumber.from(10));
+  describe('#sqrt', () => {
+    it('calculate square root correctly with small input', async () => {
+      const input = BigNumber.from('100');
+      const res = await math._sqrt(input);
+      expect(res).to.be.eq(BigNumber.from(10));
+    });
+
+    it('calculate square root correctly with large input', async () => {
+      const input = ethers.utils.parseEther('10000');
+      const res = await math._sqrt(input);
+      expect(res).to.be.eq(BigNumber.from('100000000000'));
+    });
+  });
+
+  describe('#calcSolutionForQuadratic', () => {
+    it('calculate right solution for quadratic', async () => {
+      const [a, b, c] = ['59995000000', '120100000000000', '59500000000000000'].map((v) => ethers.utils.parseEther(v));
+      const [x1, x2] = await math._calcSolutionForQuadratic(a, b, c);
+
+      expect(x1).to.be.eq(BigNumber.from('-900'));
+      expect(x2).to.be.eq(BigNumber.from('-1101'));
+    });
+
+    it('calculate right solution for quadratic with negative number', async () => {
+      const [a, b, c] = ['-10000000000', '2200000000000000', '-1000000000000000000'].map((v) =>
+        ethers.utils.parseEther(v)
+      );
+      const [x1, x2] = await math._calcSolutionForQuadratic(a, b, c);
+
+      expect(x1).to.be.eq(BigNumber.from('455'));
+      expect(x2).to.be.eq(BigNumber.from('219544'));
+    });
+  });
+
+  describe('#calcBorrowAmount', () => {
+    it('returns right amount with small liquidity pairs', async () => {
+      const reserves = { a1: '5000', b1: '10', a2: '6000', b2: '10' };
+      const input = lodash.mapValues(reserves, (v) => ethers.utils.parseEther(v));
+      const res = await math._calcBorrowAmount(input);
+      // @ts-ignore
+      expect(res).to.be.closeTo(ethers.utils.parseEther('0.45'), ethers.utils.parseEther('0.1'));
+    });
+
+    it('returns right amount with large liquidity pairs', async () => {
+      const reserves = { a1: '1200000000', b1: '600000', a2: '1000000000', b2: '300000' };
+      const input = lodash.mapValues(reserves, (v) => ethers.utils.parseEther(v));
+      const res = await math._calcBorrowAmount(input);
+      // @ts-ignore
+      expect(res).to.be.closeTo(ethers.utils.parseEther('53052.8604'), ethers.utils.parseEther('1500'));
+    });
+
+    it('returns right amount with high difference liquidity pairs', async () => {
+      const reserves = { a1: '1200000000', b1: '600000', a2: '100000', b2: '30' };
+      const input = lodash.mapValues(reserves, (v) => ethers.utils.parseEther(v));
+      const res = await math._calcBorrowAmount(input);
+      // @ts-ignore
+      expect(res).to.be.closeTo(ethers.utils.parseEther('8.729'), ethers.utils.parseEther('0.5'));
+    });
+
+    it('revert with wrong order input', async () => {
+      const reserves = { a1: '1000000000', b1: '300000', a2: '1200000000', b2: '600000' };
+      const input = lodash.mapValues(reserves, (v) => ethers.utils.parseEther(v));
+      await expect(math._calcBorrowAmount(input)).to.be.revertedWith('Wrong input order');
+    });
   });
 });
diff --git a/test/shared/expect.ts b/test/shared/expect.ts
diff --git a/test/shared/fixtures.ts b/test/shared/fixtures.ts
@@ -1,17 +1,17 @@
-import { ethers } from "hardhat";
-import { TestERC20 } from "../../typechain/TestERC20";
-import { FlashBot } from "../../typechain/FlashBot";
+import { ethers } from 'hardhat';
+import { TestERC20 } from '../../typechain/TestERC20';
+import { FlashBot } from '../../typechain/FlashBot';
 
 interface FlashBotFixture {
   weth: TestERC20;
   flashBot: FlashBot;
 }
 
 export const flashBotFixture = async (): Promise<FlashBotFixture> => {
-  const flashBotFactory = await ethers.getContractFactory("FlashBot");
-  const tokenFactory = await ethers.getContractFactory("TestERC20");
+  const flashBotFactory = await ethers.getContractFactory('FlashBot');
+  const tokenFactory = await ethers.getContractFactory('TestERC20');
 
-  const weth = (await tokenFactory.deploy("Weth", "WETH")) as TestERC20;
+  const weth = (await tokenFactory.deploy('Weth', 'WETH')) as TestERC20;
   const flashBot = (await flashBotFactory.deploy(weth.address)) as FlashBot;
 
   return { weth, flashBot };

diff --git a/yarn.lock b/yarn.lock
@@ -626,6 +626,11 @@
   resolved "https://registry.npm.taobao.org/@types/chai/download/@types/chai-4.2.16.tgz?cache=0&sync_timestamp=1617390997822&other_urls=https%3A%2F%2Fregistry.npm.taobao.org%2F%40types%2Fchai%2Fdownload%2F%40types%2Fchai-4.2.16.tgz#f09cc36e18d28274f942e7201147cce34d97e8c8"
   integrity sha1-8JzDbhjSgnT5QucgEUfM402X6Mg=
 
+"@types/lodash@^4.14.168":
+  version "4.14.168"
+  resolved "https://registry.npm.taobao.org/@types/lodash/download/@types/lodash-4.14.168.tgz#fe24632e79b7ade3f132891afff86caa5e5ce008"
+  integrity sha1-/iRjLnm3rePxMoka//hsql5c4Ag=
+
 "@types/lru-cache@^5.1.0":
   version "5.1.0"
   resolved "https://registry.npm.taobao.org/@types/lru-cache/download/@types/lru-cache-5.1.0.tgz#57f228f2b80c046b4a1bd5cac031f81f207f4f03"
@@ -4619,7 +4624,7 @@ [email protected]:
   resolved "https://registry.npm.taobao.org/lodash/download/lodash-4.17.20.tgz#b44a9b6297bcb698f1c51a3545a2b3b368d59c52"
   integrity sha1-tEqbYpe8tpjxxRo1RaKzs2jVnFI=
 
-lodash@^4.17.11, lodash@^4.17.14, lodash@^4.17.15, lodash@^4.17.4:
+lodash@^4.17.11, lodash@^4.17.14, lodash@^4.17.15, lodash@^4.17.21, lodash@^4.17.4:
   version "4.17.21"
   resolved "https://registry.npm.taobao.org/lodash/download/lodash-4.17.21.tgz#679591c564c3bffaae8454cf0b3df370c3d6911c"
   integrity sha1-Z5WRxWTDv/quhFTPCz3zcMPWkRw=