From 3be19ce659411b0f721eea75a7409d396913ecb3 Mon Sep 17 00:00:00 2001
From: lza_menace <lza_menace@protonmail.com>
Date: Thu, 17 Nov 2022 15:34:32 -0800
Subject: [PATCH] add some admin funcs

---
 foundry.toml      |  2 +-
 src/SendIt.sol    | 32 ++++++++++++++++++++++++++------
 test/SendIt.t.sol | 20 +++++++++++++-------
 3 files changed, 40 insertions(+), 14 deletions(-)

diff --git a/foundry.toml b/foundry.toml
index 4cb5efa..c7dd4a8 100644
--- a/foundry.toml
+++ b/foundry.toml
@@ -2,7 +2,7 @@
 src = 'src'
 out = 'out'
 libs = ['lib']
-gas_reports = ["*"]
+gas_reports = ["SendIt"]
 
 [rpc_endpoints]
 goerli = "${GOERLI_RPC_URL}"
diff --git a/src/SendIt.sol b/src/SendIt.sol
index 13e3b02..1766287 100644
--- a/src/SendIt.sol
+++ b/src/SendIt.sol
@@ -1,12 +1,14 @@
 // SPDX-License-Identifier: UNLICENSED
 pragma solidity ^0.8.13;
 
-import "solmate/tokens/ERC721.sol";
-import "solmate/tokens/ERC1155.sol";
+import {ERC721} from "openzeppelin-contracts/token/ERC721/ERC721.sol";
+import {ERC1155} from "openzeppelin-contracts/token/ERC1155/ERC1155.sol";
+import {Ownable} from "openzeppelin-contracts/access/Ownable.sol";
 
-contract SendIt {
+contract SendIt is Ownable {
 
-    mapping(address => address) public addressVault;
+    mapping(address => address) public addressVault; // users can store their personal vaults for ease of use
+    uint256 public usageFee = .0001 ether;           // charge a small fee for the cost savings it provides
 
     event TokenTransfer(address indexed contractAddress, uint256 tokenIndex, address indexed from, address indexed to);
     
@@ -27,6 +29,23 @@ contract SendIt {
         _;
     }
 
+    /*************************
+    Admin
+    **************************/
+
+    function updateFee(uint256 amount) external onlyOwner {
+        usageFee = amount;
+    }
+
+    function withdraw() external onlyOwner {
+        uint256 balance = address(this).balance;
+        payable(msg.sender).transfer(balance);
+    }
+
+    /*************************
+    User
+    **************************/
+
     function updateVault(address vaultAddress) external {
         addressVault[msg.sender] = vaultAddress;
     }
@@ -36,7 +55,7 @@ contract SendIt {
         uint256 tokenIndex,
         address recipient,
         bool isERC1155
-    ) public {
+    ) private {
         if (isERC1155) {
             require(ERC1155(contractAddress).balanceOf(msg.sender, tokenIndex) > 0, "Sender is not the token owner, cannot proceed with transfer.");
             require(ERC1155(contractAddress).isApprovedForAll(msg.sender, address(this)), "Contract not approved to send tokens on Sender behalf.");
@@ -54,8 +73,9 @@ contract SendIt {
         uint256[] calldata tokenIndexes,
         address[] calldata recipients,
         bool isERC1155
-    ) external {
+    ) external payable {
         require(tokenIndexes.length == recipients.length, "Array lengths must match.");
+        require(msg.value >= tokenIndexes.length * usageFee, "Invalid usage fee sent.");
         for(uint256 i; i < tokenIndexes.length; i++) {
             contractTransfer(contractAddress, tokenIndexes[i], recipients[i], isERC1155);
         }
diff --git a/test/SendIt.t.sol b/test/SendIt.t.sol
index 8d59345..4192d13 100644
--- a/test/SendIt.t.sol
+++ b/test/SendIt.t.sol
@@ -17,16 +17,19 @@ contract SendItTest is Test {
 
     function testBulkTransferSucceeds() public {
         uint256 amt = 20;
+        uint256 fee = sendit.usageFee();
+        uint256 val = fee * amt;
         uint256[] memory tokenIndexes = new uint256[](amt);
         address[] memory recipients = new address[](amt);
         for (uint256 i; i < amt; i++) {
             tokenIndexes[i] = i + 1;
             recipients[i] = address(1);
         }
+        vm.deal(address(5), 1 ether);
         vm.startPrank(address(5));
         nft.mint(address(5), amt);
         nft.setApprovalForAll(address(sendit), true);
-        sendit.contractBulkTransfer(
+        sendit.contractBulkTransfer{value: val}(
             address(nft),
             tokenIndexes,
             recipients,
@@ -44,10 +47,13 @@ contract SendItTest is Test {
         vm.stopPrank();
     }
 
-    // function testUpdateVault() public {
-    //     sendit.updateVault(address(3));
-    //     console.log(address(sendit));
-    //     // console.log(sendit);
-    //     assertEq(sendit.addressVault(address(2)), address(0));
-    // }
+    function testUpdateVault() public {
+        sendit.updateVault(address(3));
+        assertEq(sendit.addressVault(address(2)), address(0));
+    }
+
+    // only token owner can bulk transfer
+    // can only proceed after approval set
+    // only owner can updateFee
+    // only owner can withdraw
 }