feat(whitelist-house): extend house contract to use whitelist feature

primitivefinance · Jul 12, 2021 · 00a5db0 · 00a5db0
1 parent 0cb7242
commit 00a5db0
Show file tree

Hide file tree

Showing 6 changed files with 33 additions and 21 deletions.
diff --git a/contracts/Admin.sol b/contracts/Admin.sol
@@ -13,9 +13,11 @@ contract Admin is IAdmin {
         _;
     }
 
+    constructor() {}
+
     /// @param _admin The address receiving the admin rights
-    constructor(address _admin) {
-        admin = _admin;
+    function initializeAdmin(address _admin) internal {
+      admin = _admin;
     }
 
     /// @inheritdoc IAdmin

diff --git a/contracts/PrimitiveHouse.sol b/contracts/PrimitiveHouse.sol
@@ -2,6 +2,7 @@
 pragma solidity 0.8.0;
 
 import "./interfaces/IPrimitiveHouse.sol";
+import  "./Whitelist.sol";
 import "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
 import "@uniswap/v3-core/contracts/interfaces/IUniswapV3Factory.sol";
 import "@uniswap/v3-core/contracts/interfaces/IUniswapV3Pool.sol";
@@ -10,7 +11,10 @@ import "@primitivefinance/primitive-v2-core/contracts/interfaces/engine/IPrimiti
 import "@primitivefinance/primitive-v2-core/contracts/libraries/Margin.sol";
 import "@primitivefinance/primitive-v2-core/contracts/libraries/Position.sol";
 
-contract PrimitiveHouse is IPrimitiveHouse {
+contract PrimitiveHouse is 
+  Whitelist, 
+  IPrimitiveHouse 
+  {
     using SafeERC20 for IERC20;
     using Margin for mapping(address => Margin.Data);
     using Margin for Margin.Data;
@@ -29,7 +33,7 @@ contract PrimitiveHouse is IPrimitiveHouse {
     mapping(address => Margin.Data) _margins;
     mapping(bytes32 => Position.Data) _positions;
 
-    constructor() {}
+    constructor() Whitelist() {}
 
     modifier lock() {
       require(reentrant != 1, "locked");
@@ -57,7 +61,7 @@ contract PrimitiveHouse is IPrimitiveHouse {
       uint64 sigma, 
       uint32 time, 
       uint riskyPrice
-    ) external override lock  {
+    ) external override lock onlyWhitelisted {
       IPrimitiveEngineActions(engine).create(
         strike, 
         sigma, 
@@ -72,7 +76,7 @@ contract PrimitiveHouse is IPrimitiveHouse {
       address owner, 
       uint256 delRisky, 
       uint256 delStable
-    ) external override lock  {
+    ) external override lock onlyWhitelisted {
       IPrimitiveEngineActions(engine).deposit(
         address(this), 
         delRisky, 
@@ -88,7 +92,7 @@ contract PrimitiveHouse is IPrimitiveHouse {
     function withdraw(
       uint256 delRisky, 
       uint256 delStable
-    ) external override lock  {
+    ) external override lock onlyWhitelisted {
       IPrimitiveEngineActions(engine).withdraw(delRisky, delStable);
 
       _margins.withdraw(delRisky, delStable);
@@ -102,7 +106,7 @@ contract PrimitiveHouse is IPrimitiveHouse {
       address owner,
       uint256 delLiquidity,
       bool fromMargin
-    ) external override lock  {
+    ) external override lock onlyWhitelisted {
       (uint256 delRisky, uint256 delStable) = IPrimitiveEngineActions(engine).allocate(
           poolId, 
           address(this), 
@@ -126,7 +130,7 @@ contract PrimitiveHouse is IPrimitiveHouse {
       address owner, 
       uint256 delLiquidity, 
       uint256 maxPremium
-    ) external override lock  {
+    ) external override lock onlyWhitelisted {
       IPrimitiveEngineActions(engine).borrow(
         poolId, 
         delLiquidity, 
@@ -142,7 +146,7 @@ contract PrimitiveHouse is IPrimitiveHouse {
       address owner,
       uint256 delLiquidity,
       bool fromMargin
-    ) external override lock  {
+    ) external override lock onlyWhitelisted {
       bytes memory data = '0x';
       (uint256 delRisky, uint256 delStable, uint256 premium) = IPrimitiveEngineActions(engine).repay(
           poolId, 
@@ -167,21 +171,21 @@ contract PrimitiveHouse is IPrimitiveHouse {
       uint256 deltaOut, 
       uint256 deltaInMax, 
       bytes calldata data
-    ) external override lock {
+    ) external override lock onlyWhitelisted {
       IPrimitiveEngineActions(engine).swap(poolId, addXRemoveY, deltaOut, deltaInMax, true, data);
     }
 
     function swapXForY(
       bytes32 poolId, 
       uint256 deltaOut
-    ) external override lock {
+    ) external override lock onlyWhitelisted {
       IPrimitiveEngineActions(engine).swap(poolId, true, deltaOut, type(uint256).max, true, new bytes(0));
     }
 
     function swapYForX(
       bytes32 poolId, 
       uint256 deltaOut
-    ) external override lock {
+    ) external override lock onlyWhitelisted {
       IPrimitiveEngineActions(engine).swap(poolId, false, deltaOut, type(uint256).max, true, new bytes(0));
     }
 

diff --git a/contracts/PrimitiveHouseFactory.sol b/contracts/PrimitiveHouseFactory.sol
@@ -4,6 +4,7 @@
 
 import "./interfaces/IPrimitiveHouseFactory.sol";
 import "./PrimitiveHouse.sol";
+import "./interfaces/IAdmin.sol";
 
 contract PrimitiveHouseFactory is IPrimitiveHouseFactory {
     /// @inheritdoc IPrimitiveHouseFactory
@@ -29,6 +30,7 @@ contract PrimitiveHouseFactory is IPrimitiveHouseFactory {
         require(engine != address(0), "Cannot be same token");
         house = deploy(address(this), engine);
         getHouse[engine] = house;
+        IAdmin(house).setAdmin(msg.sender);
         emit Deployed(msg.sender, engine, house);
     }
 

diff --git a/contracts/Whitelist.sol b/contracts/Whitelist.sol
@@ -27,10 +27,10 @@ contract Whitelist is IWhitelist, Admin {
     }
 
 
-    /// EFFECT FUNCTIONS ///
-
-    /// @param _admin The address of the admin
-    constructor(address _admin) Admin(_admin) { }
+    constructor() {
+      initializeAdmin(msg.sender);
+    }
+
 
     /// @inheritdoc IWhitelist
     function addKeys(bytes32[] memory hashes) external override onlyAdmin() {

diff --git a/package.json b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "@primitivefinance/primitive-v2-periphery",
-  "version": "0.0.1",
+  "version": "0.1.1",
   "description": "Primitive V2 periphery smart contracts.",
   "main": "index.js",
   "keywords": [
@@ -12,6 +12,8 @@
   "contracts",
   "artifacts/contracts/PrimitiveEngine.sol/PrimitiveHouse.json",
   "artifacts/contracts/PrimitiveFactory.sol/PrimitiveHouseFactory.json",
+  "artifacts/contracts/Admin.sol/Admin.json",
+  "artifacts/contracts/Whitelist.sol/Whitelist.json",
   "artifacts/contracts/interfaces/**/*.json",
   "!artifacts/contracts/interfaces/**/*.dbg.json"
   ],

diff --git a/test/unit/createTestContracts.ts b/test/unit/createTestContracts.ts
@@ -1,5 +1,5 @@
 import hre, { ethers } from 'hardhat'
-import { Wallet, Contract } from 'ethers'
+import { Wallet, Contract, utils } from 'ethers'
 import { Contracts } from '../../types'
 import { deployContract } from 'ethereum-waffle'
 import * as ContractTypes from '../../typechain'
@@ -34,8 +34,10 @@ async function initializeBaseContracts(deployer: Wallet): Promise<BaseContracts>
   await houseFactory.deploy(engine.address)
   const houseAddr = await houseFactory.getHouse(engine.address)
   const house = (await ethers.getContractAt(PrimitiveHouseAbi, houseAddr)) as unknown as ContractTypes.PrimitiveHouse
-  const admin = (await deploy('Admin', deployer, [deployer.address])) as ContractTypes.Admin
-  const whitelist = (await deploy('Whitelist', deployer, [deployer.address])) as ContractTypes.Whitelist
+  await house.addKeys([utils.solidityKeccak256(['string'], ['wentoken'])])
+  await house.useKey('wentoken', deployer.address)
+  const admin = (await deploy('Admin', deployer)) as ContractTypes.Admin
+  const whitelist = (await deploy('Whitelist', deployer)) as ContractTypes.Whitelist
   return { factory, engine, stable, risky, house, houseFactory, admin, whitelist }
 }