Skip to content

Commit

Permalink
Update NFT contracts to use custom errors
Browse files Browse the repository at this point in the history
  • Loading branch information
kyledewy committed Feb 9, 2023
1 parent 18024f0 commit 9e2a171
Show file tree
Hide file tree
Showing 6 changed files with 141 additions and 138 deletions.
10 changes: 10 additions & 0 deletions contracts/interfaces/ICoverNFT.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,4 +13,14 @@ interface ICoverNFT is IERC721 {
function burn(uint tokenId) external;

function operatorTransferFrom(address from, address to, uint256 tokenId) external;

error NotOperator();
error NotMinted();
error ZeroAddress();
error WrongFrom();
error InvalidRecipient();
error InvalidNewOperatorAddress();
error NotAuthorized();
error UnsafeRecipient();
error AlreadyMinted();
}
10 changes: 10 additions & 0 deletions contracts/interfaces/IStakingNFT.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,4 +20,14 @@ interface IStakingNFT is IERC721 {

function stakingPoolOf(uint tokenId) external view returns (uint poolId);

error NotOperator();
error NotMinted();
error ZeroAddress();
error WrongFrom();
error InvalidRecipient();
error InvalidNewOperatorAddress();
error NotAuthorized();
error UnsafeRecipient();
error AlreadyMinted();
error NotStakingPool();
}
76 changes: 23 additions & 53 deletions contracts/modules/cover/CoverNFT.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@ contract CoverNFT is ICoverNFT {
address public operator;

modifier onlyOperator {
require(msg.sender == operator, "CoverNFT: Not operator");
if (msg.sender != operator) revert NotOperator();
_;
}

Expand All @@ -47,11 +47,11 @@ contract CoverNFT is ICoverNFT {
}

function ownerOf(uint256 id) public view virtual returns (address owner) {
require((owner = _ownerOf[id]) != address(0), "NOT_MINTED");
if ((owner = _ownerOf[id]) == address(0)) revert NotMinted();
}

function balanceOf(address owner) public view virtual returns (uint256) {
require(owner != address(0), "ZERO_ADDRESS");
if (owner == address(0)) revert ZeroAddress();

return _balanceOf[owner];
}
Expand All @@ -68,8 +68,8 @@ contract CoverNFT is ICoverNFT {

function operatorTransferFrom(address from, address to, uint256 tokenId) external onlyOperator {

require(from == _ownerOf[tokenId], "WRONG_FROM");
require(to != address(0), "INVALID_RECIPIENT");
if (from != _ownerOf[tokenId]) revert WrongFrom();
if (to == address(0)) revert InvalidRecipient();

// Underflow of the sender's balance is impossible because we check for
// ownership above and the recipient's balance can't realistically overflow.
Expand All @@ -85,7 +85,7 @@ contract CoverNFT is ICoverNFT {
}

function changeOperator(address _newOperator) public onlyOperator returns (bool) {
require(_newOperator != address(0), "CoverNFT: Invalid newOperator address");
if (_newOperator == address(0)) revert InvalidNewOperatorAddress();

operator = _newOperator;
return true;
Expand All @@ -95,7 +95,7 @@ contract CoverNFT is ICoverNFT {
function approve(address spender, uint256 id) public virtual {
address owner = _ownerOf[id];

require(msg.sender == owner || isApprovedForAll[owner][msg.sender], "NOT_AUTHORIZED");
if (msg.sender != owner && !isApprovedForAll[owner][msg.sender]) revert NotAuthorized();

getApproved[id] = spender;

Expand All @@ -113,14 +113,13 @@ contract CoverNFT is ICoverNFT {
address to,
uint256 id
) public virtual {
require(from == _ownerOf[id], "WRONG_FROM");
if (from != _ownerOf[id]) revert WrongFrom();

require(to != address(0), "INVALID_RECIPIENT");
if (to == address(0)) revert InvalidRecipient();

require(
msg.sender == from || isApprovedForAll[from][msg.sender] || msg.sender == getApproved[id],
"NOT_AUTHORIZED"
);
if (msg.sender != from && !isApprovedForAll[from][msg.sender] && msg.sender != getApproved[id]) {
revert NotAuthorized();
}

// Underflow of the sender's balance is impossible because we check for
// ownership above and the recipient's balance can't realistically overflow.
Expand All @@ -144,12 +143,10 @@ contract CoverNFT is ICoverNFT {
) public virtual {
transferFrom(from, to, id);

require(
to.code.length == 0 ||
ERC721TokenReceiver(to).onERC721Received(msg.sender, from, id, "") ==
ERC721TokenReceiver.onERC721Received.selector,
"UNSAFE_RECIPIENT"
);
if (to.code.length != 0
&& ERC721TokenReceiver(to).onERC721Received(msg.sender, from, id, "") != ERC721TokenReceiver.onERC721Received.selector) {
revert UnsafeRecipient();
}
}

function safeTransferFrom(
Expand All @@ -160,12 +157,10 @@ contract CoverNFT is ICoverNFT {
) public virtual {
transferFrom(from, to, id);

require(
to.code.length == 0 ||
ERC721TokenReceiver(to).onERC721Received(msg.sender, from, id, data) ==
ERC721TokenReceiver.onERC721Received.selector,
"UNSAFE_RECIPIENT"
);
if (to.code.length != 0
&& ERC721TokenReceiver(to).onERC721Received(msg.sender, from, id, data) != ERC721TokenReceiver.onERC721Received.selector) {
revert UnsafeRecipient();
}
}


Expand All @@ -180,9 +175,9 @@ contract CoverNFT is ICoverNFT {
// Internal functions

function _mint(address to, uint256 id) internal virtual {
require(to != address(0), "INVALID_RECIPIENT");
if (to == address(0)) revert InvalidRecipient();

require(_ownerOf[id] == address(0), "ALREADY_MINTED");
if (_ownerOf[id] != address(0)) revert AlreadyMinted();

// Counter overflow is incredibly unrealistic.
unchecked {
Expand All @@ -197,7 +192,7 @@ contract CoverNFT is ICoverNFT {
function _burn(uint256 id) internal virtual {
address owner = _ownerOf[id];

require(owner != address(0), "NOT_MINTED");
if (owner == address(0)) revert NotMinted();

// Ownership check above ensures no underflow.
unchecked {
Expand All @@ -211,31 +206,6 @@ contract CoverNFT is ICoverNFT {
emit Transfer(owner, address(0), id);
}

function _safeMint(address to, uint256 id) internal virtual {
_mint(to, id);

require(
to.code.length == 0 ||
ERC721TokenReceiver(to).onERC721Received(msg.sender, address(0), id, "") ==
ERC721TokenReceiver.onERC721Received.selector,
"UNSAFE_RECIPIENT"
);
}

function _safeMint(
address to,
uint256 id,
bytes memory data
) internal virtual {
_mint(to, id);

require(
to.code.length == 0 ||
ERC721TokenReceiver(to).onERC721Received(msg.sender, address(0), id, data) ==
ERC721TokenReceiver.onERC721Received.selector,
"UNSAFE_RECIPIENT"
);
}
}

/// @notice A generic interface for a contract which properly accepts ERC721 tokens.
Expand Down
90 changes: 42 additions & 48 deletions contracts/modules/staking/StakingNFT.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -40,16 +40,16 @@ contract StakingNFT is IStakingNFT {
// operator functions

function operatorTransferFrom(address from, address to, uint id) external {
require(msg.sender == operator, "NOT_OPERATOR");
require(from == ownerOf(id), "WRONG_FROM");
require(to != address(0), "INVALID_RECIPIENT");
if (msg.sender != operator) revert NotOperator();
if (from != ownerOf(id)) revert WrongFrom();
if (to == address(0)) revert InvalidRecipient();

// underflow of the sender's balance is impossible because we check for
// ownership above and the recipient's balance can't realistically overflow
unchecked {
_balanceOf[from]--;
_balanceOf[to]++;
}
unchecked {
_balanceOf[from]--;
_balanceOf[to]++;
}

_tokenInfo[id].owner = to;
delete getApproved[id];
Expand All @@ -58,27 +58,26 @@ contract StakingNFT is IStakingNFT {
}

function changeOperator(address newOperator) public {
require(msg.sender == operator, "NOT_OPERATOR");
require(newOperator != address(0), "INVALID_OPERATOR");
if (msg.sender != operator) revert NotOperator();
if (newOperator == address(0)) revert InvalidNewOperatorAddress();
operator = newOperator;
}

// minting and supply

function mint(uint poolId, address to) public returns (uint id) {

require(
msg.sender == StakingPoolLibrary.getAddress(stakingPoolFactory, poolId),
'NOT_STAKING_POOL'
);
if (msg.sender != StakingPoolLibrary.getAddress(stakingPoolFactory, poolId)) {
revert NotStakingPool();
}

require(to != address(0), "INVALID_RECIPIENT");
if (to == address(0)) revert InvalidRecipient();

// counter overflow is incredibly unrealistic
unchecked {
id = _totalSupply++;
_balanceOf[to]++;
}
unchecked {
id = _totalSupply++;
_balanceOf[to]++;
}

_tokenInfo[id].owner = to;
_tokenInfo[id].poolId = uint96(poolId);
Expand All @@ -95,21 +94,22 @@ contract StakingNFT is IStakingNFT {
function tokenInfo(uint tokenId) public view returns (uint poolId, address owner) {
poolId = _tokenInfo[tokenId].poolId;
owner = _tokenInfo[tokenId].owner;
require(owner != address(0), "NOT_MINTED");
if (owner == address(0)) revert NotMinted();
}

function stakingPoolOf(uint tokenId) public view returns (uint poolId) {
poolId = _tokenInfo[tokenId].poolId;
require(poolId != 0, "NOT_MINTED");
if (poolId == 0) revert NotMinted();
}

// ERC165

function supportsInterface(bytes4 interfaceId) public pure returns (bool) {
return
interfaceId == 0x01ffc9a7 || // ERC165 Interface ID for ERC165
interfaceId == 0x80ac58cd || // ERC165 Interface ID for ERC721
interfaceId == 0x5b5e139f; // ERC165 Interface ID for ERC721Metadata
interfaceId == 0x01ffc9a7 || // ERC165 Interface ID for ERC165
interfaceId == 0x80ac58cd || // ERC165 Interface ID for ERC721
interfaceId == 0x5b5e139f;
// ERC165 Interface ID for ERC721Metadata
}

// ERC721
Expand All @@ -122,17 +122,17 @@ contract StakingNFT is IStakingNFT {

function ownerOf(uint id) public view returns (address owner) {
owner = _tokenInfo[id].owner;
require(owner != address(0), "NOT_MINTED");
if (owner == address(0)) revert NotMinted();
}

function balanceOf(address owner) public view returns (uint) {
require(owner != address(0), "ZERO_ADDRESS");
if (owner == address(0)) revert ZeroAddress();
return _balanceOf[owner];
}

function approve(address spender, uint id) public {
address owner = ownerOf(id);
require(msg.sender == owner || isApprovedForAll[owner][msg.sender], "NOT_AUTHORIZED");
if (msg.sender != owner && !isApprovedForAll[owner][msg.sender]) revert NotAuthorized();
getApproved[id] = spender;
emit Approval(owner, spender, id);
}
Expand All @@ -151,20 +151,19 @@ contract StakingNFT is IStakingNFT {

function transferFrom(address from, address to, uint id) public {

require(from == ownerOf(id), "WRONG_FROM");
require(to != address(0), "INVALID_RECIPIENT");
if (from != ownerOf(id)) revert WrongFrom();
if (to == address(0)) revert InvalidRecipient();

require(
msg.sender == from || isApprovedForAll[from][msg.sender] || msg.sender == getApproved[id],
"NOT_AUTHORIZED"
);
if (msg.sender != from && !isApprovedForAll[from][msg.sender] && msg.sender != getApproved[id]) {
revert NotAuthorized();
}

// underflow of the sender's balance is impossible because we check for
// ownership above and the recipient's balance can't realistically overflow
unchecked {
_balanceOf[from]--;
_balanceOf[to]++;
}
unchecked {
_balanceOf[from]--;
_balanceOf[to]++;
}

_tokenInfo[id].owner = to;
delete getApproved[id];
Expand All @@ -174,12 +173,9 @@ contract StakingNFT is IStakingNFT {

function safeTransferFrom(address from, address to, uint id) public {
transferFrom(from, to, id);
require(
to.code.length == 0 ||
ERC721TokenReceiver(to).onERC721Received(msg.sender, from, id, "")
== ERC721TokenReceiver.onERC721Received.selector,
"UNSAFE_RECIPIENT"
);
if (to.code.length != 0 && ERC721TokenReceiver(to).onERC721Received(msg.sender, from, id, "") != ERC721TokenReceiver.onERC721Received.selector) {
revert UnsafeRecipient();
}
}

function safeTransferFrom(
Expand All @@ -189,12 +185,10 @@ contract StakingNFT is IStakingNFT {
bytes calldata data
) public {
transferFrom(from, to, id);
require(
to.code.length == 0 ||
ERC721TokenReceiver(to).onERC721Received(msg.sender, from, id, data)
== ERC721TokenReceiver.onERC721Received.selector,
"UNSAFE_RECIPIENT"
);

if (to.code.length != 0 && ERC721TokenReceiver(to).onERC721Received(msg.sender, from, id, data) != ERC721TokenReceiver.onERC721Received.selector) {
revert UnsafeRecipient();
}
}
}

Expand Down

0 comments on commit 9e2a171

Please sign in to comment.