diff --git a/contracts/components/threat_oracle/ThreatOracle.sol b/contracts/components/threat_oracle/ThreatOracle.sol index 2c0d915b..389967b9 100644 --- a/contracts/components/threat_oracle/ThreatOracle.sol +++ b/contracts/components/threat_oracle/ThreatOracle.sol @@ -6,8 +6,15 @@ pragma solidity ^0.8.9; import "../BaseComponentUpgradeable.sol"; import "./ThreatOracleCore.sol"; +/** + * Onchain registry for addresses flagged by the Attack Detector bot. Each address registered will be assigned a 'threat level', which will be a value + * between 0-5. Only the account who has been granted the role of THREAT_ORACLE_ADMIN can register addresses. + * Addresses that were initially given a non-zero threat level that were later discovered to have been false positives will be updated to have their + * threat level set to zero. + */ contract ThreatOracle is BaseComponentUpgradeable, ThreatOracleCore { string public constant version = "0.1.0"; + /// @custom:oz-upgrades-unsafe-allow constructor constructor(address forwarder) initializer ForwardedContext(forwarder) {} diff --git a/contracts/components/threat_oracle/ThreatOracleCore.sol b/contracts/components/threat_oracle/ThreatOracleCore.sol index e8fab495..d1778e01 100644 --- a/contracts/components/threat_oracle/ThreatOracleCore.sol +++ b/contracts/components/threat_oracle/ThreatOracleCore.sol @@ -9,15 +9,24 @@ import "../BaseComponentUpgradeable.sol"; abstract contract ThreatOracleCore is BaseComponentUpgradeable { using EnumerableMap for EnumerableMap.AddressToUintMap; + /// Map of addresses to their corresponding threat level. EnumerableMap.AddressToUintMap private _addressThreatLevel; event AddressThreatLevelSet(address indexed _address, uint256 indexed threatLevel); - error IncorrectAmounts(uint256 addressAmount, uint256 threatLevelAmount); + error UnevenAmounts(uint256 addressAmount, uint256 threatLevelAmount); error ThreatLevelAlreadySet(address _address, uint256 threatLevel); + /** + * @notice Method to register addresses and assign them a threat level. + * @dev Only accessible to the account that has admin access. + * Does not allow setting the threat level to its current value. + * Threat level is to be between 0-5. + * @param addresses Array of addresses to register. + * @param threatLevels Array of the addresses' corresponding threat levels. + */ function setThreatLevels(address[] calldata addresses, uint256[] calldata threatLevels) external onlyRole(THREAT_ORACLE_ADMIN_ROLE) { - if (addresses.length != threatLevels.length) revert IncorrectAmounts(addresses.length, threatLevels.length); + if (addresses.length != threatLevels.length) revert UnevenAmounts(addresses.length, threatLevels.length); for (uint256 i = 0; i < addresses.length; i++) { _setThreatLevel(addresses[i], threatLevels[i]); @@ -31,15 +40,32 @@ abstract contract ThreatOracleCore is BaseComponentUpgradeable { emit AddressThreatLevelSet(_address, threatLevel); } + /** + * @notice Get threat level for an address + * @dev A return value of '0' does not mean an address is 'no threat', + * as it could also mean it has not been registered. + * @param _address Address of interest. + * @return threat level of the given address. + */ function getThreatLevel(address _address) public view returns (uint256) { (,uint256 fetchedThreatLevel) = _addressThreatLevel.tryGet(_address); return fetchedThreatLevel; } + /** + * @notice Check if address has been registered. + * @param _address Address of interest. + * @return true if the address has been registered, false otherwise. + */ function isRegistered(address _address) public view returns (bool) { return _addressThreatLevel.contains(_address); } + /** + * @notice Gets the total amount of addresses that have been registered. + * @dev Amount includes addresses deemed no threat. + * @return amount of addresses that have been registered. + */ function totalAddressesRegistered() public view returns (uint256) { return _addressThreatLevel.length(); } diff --git a/test/components/threat.oracle.test.js b/test/components/threat.oracle.test.js index ec7dfe23..319c8435 100644 --- a/test/components/threat.oracle.test.js +++ b/test/components/threat.oracle.test.js @@ -82,7 +82,7 @@ describe('Threat Oracles', async function () { const highThreatLevelsAmount = [...mockThreatLevels, 4, 5]; await expect(this.threatOracle.connect(this.accounts.manager).setThreatLevels(mockAddresses, highThreatLevelsAmount)) - .to.be.revertedWith(`IncorrectAmounts(${mockAddresses.length}, ${highThreatLevelsAmount.length})`); + .to.be.revertedWith(`UnevenAmounts(${mockAddresses.length}, ${highThreatLevelsAmount.length})`); for(let i = 0; i < mockAddresses.length; i++) { const mockAddress = mockAddresses[i];