Assignment: fix vulnerabilities/loopholes in ERC20 and ERC721 token contracts

contracts/ERC20.sol

@@ -11,62 +11,84 @@ contract ERC20 is IERC20 {
 
     uint256 public totalSupply;
     mapping(address => uint256) public balanceOf;
-    mapping(address => mapping(address => uint256)) public allowance;
+    mapping(address => mapping(address => uint256)) public allowance; // 1stAddr - owner, 2ndAddr - spender
     string public name;
     string public symbol;
+    uint8 public decimals;
+
+    address public owner;
 
     constructor(string memory _name, string memory _symbol, uint8 _decimals) {
         name = _name;
         symbol = _symbol;
         decimals = _decimals;
+
+        owner = msg.sender;
+    }
+
+    modifier onlyOwner(){
+        require(msg.sender == owner, "You are not authorized");
+        _;
     }
 
     function transfer(address recipient, uint256 amount)
         external
         returns (bool)
     {
+        require(recipient != address(0), "Recipient cannot be address zero");
+        require(balanceOf[msg.sender] > amount, "Insufficient balance");
+
         balanceOf[msg.sender] -= amount;
         balanceOf[recipient] += amount;
         emit Transfer(msg.sender, recipient, amount);
-
         return true;
     }
 
-    function approve(address spender, uint256 amount) external returns (bool) {
+    mapping(address => bool) private approvalMap;
+
+    function approve(address spender, uint256 amount) external onlyOwner returns (bool) {
+        require(balanceOf[msg.sender] >= amount, "Insufficient balance");
         allowance[msg.sender][spender] = amount;
+        approvalMap[spender] = true;
         emit Approval(msg.sender, spender, amount);
         return true;
     }
 
+
     function transferFrom(address sender, address recipient, uint256 amount)
         external
         returns (bool)
     {
-        require(msg.sender != recipient, "cannot transfer to self");
+        require(recipient != address(0), "Recipient cannot be address zero");
+        require(approvalMap[msg.sender] == true, "You have not been approved by sender");
+        require(allowance[sender][msg.sender] >= amount, "Allowance Exceeded");
+        require(msg.sender != recipient, "You cannot be the recipoent");
+
         allowance[sender][msg.sender] -= amount;
         balanceOf[sender] -= amount;
         balanceOf[recipient] += amount;
         emit Transfer(sender, recipient, amount);
         return true;
     }
 
-    function _mint(address to, uint256 amount) internal {
+    function _mint(address to, uint256 amount) internal onlyOwner {
         balanceOf[to] += amount;
         totalSupply += amount;
         emit Transfer(address(0), to, amount);
     }
 
     function _burn(address from, uint256 amount) internal {
+        require(balanceOf[from] >= amount, "Excess Burn Amount");
         balanceOf[from] -= amount;
         totalSupply -= amount;
         emit Transfer(from, address(0), amount);
     }
 
-    function mint(address to, uint256 amount) external {
+    function mint(address to, uint256 amount) external onlyOwner {
         _mint(to, amount);
     }
 
-    function burn(address from, uint256 amount) external {
+    function burn(address from, uint256 amount) external onlyOwner {
         _burn(from, amount);
     }
 }

contracts/ERC721.sol

@@ -52,6 +52,17 @@ contract ERC721 is IERC721 {
         address indexed owner, address indexed operator, bool approved
     );
 
+    address private _owner;
+
+    constructor() {
+        _owner = msg.sender; 
+    }
+
+    modifier onlyOwner() {
+        require(_owner == msg.sender, "You are not anthorized");
+        _;
+    }
+
     // Mapping from token ID to owner address
     mapping(uint256 => address) internal _ownerOf;
 
@@ -84,11 +95,13 @@ contract ERC721 is IERC721 {
     }
 
     function setApprovalForAll(address operator, bool approved) external {
+        require(operator != address(0), "Operator cannot be address zero");
         isApprovedForAll[msg.sender][operator] = approved;
         emit ApprovalForAll(msg.sender, operator, approved);
     }
 
     function approve(address spender, uint256 id) external {
+        require(spender != address(0), "Spender cannot be address zero");
         address owner = _ownerOf[id];
         require(
             msg.sender == owner || isApprovedForAll[owner][msg.sender],
@@ -117,6 +130,7 @@ contract ERC721 is IERC721 {
     }
 
     function transferFrom(address from, address to, uint256 id) public {
+        require(from != address(0), "from address cannot be address zero");
         require(from == _ownerOf[id], "from != owner");
         require(to != address(0), "transfer to zero address");
 
@@ -158,7 +172,7 @@ contract ERC721 is IERC721 {
         );
     }
 
-    function _mint(address to, uint256 id) internal {
+    function _mint(address to, uint256 id) internal onlyOwner {
         require(to != address(0), "mint to zero address");
         require(_ownerOf[id] == address(0), "already minted");
 
@@ -182,7 +196,7 @@ contract ERC721 is IERC721 {
 }
 
 contract MyNFT is ERC721 {
-    function mint(address to, uint256 id) external {
+    function mint(address to, uint256 id) external onlyOwner{
         _mint(to, id);
     }