Add functionality to configure ThreatOracle address

forta-network · Dec 4, 2023 · d4ed4af · d4ed4af
1 parent 740b41c
commit d4ed4af
Show file tree

Hide file tree

Showing 2 changed files with 43 additions and 3 deletions.
diff --git a/contracts/components/threat_oracle/ThreatOracleProvider.sol b/contracts/components/threat_oracle/ThreatOracleProvider.sol
@@ -37,7 +37,7 @@ import "./IThreatOracle.sol";
  *          { ... }
  *       ```
  *     - `onlyNonThreatAccounts(address[] memory)` is intended to check arguments
- *     - of type `address[]`. If the function has multiple arguments of this type,
+ *       of type `address[]`. If the function has multiple arguments of this type,
  *       then the modifier can be used multiple times. For example:
  *       ```
  *          function foobar(
@@ -51,6 +51,10 @@ import "./IThreatOracle.sol";
  *       Additionally, this modifier implements a upper bound on `address[]` amount
  *       as a preventative technique for potential DoS attacks
  *
+ * Lastly, this contract implements an `internal` function to update the contract
+ * that is serving as the `ThreatOracle` for testing. This function should be called by another
+ * function with either `public` or `external` visibility WITH access control to only
+ * allow a privelaged account to update the used `ThreatOracle` contract.
  */
 abstract contract ThreatOracleProvider {
     bytes32 constant private EXPLOIT_CATEGORY = keccak256("exploit");
@@ -62,6 +66,8 @@ abstract contract ThreatOracleProvider {
     // (DoS prevention)
     uint8 private _maxAddressArgumentAmount;
 
+    event ThreatOracleContractUpdated(address indexed previousAddress, address indexed newAddress);
+
     error ThreatAccountIdentified(address account, string threatCategory, uint8 confidenceScore);
     error MaxAddressArgumentAmountExceeded(uint8 maxAddressArgumentAmount, uint exceedingAmount);
 
@@ -98,6 +104,10 @@ abstract contract ThreatOracleProvider {
         _maxAddressArgumentAmount = __maxAddressArgumentAmount;
     }
 
+    function getThreatOracleAddress() public view returns (address) {
+        return address(_threatOracle);
+    }
+
     function _confirmNonThreatAccount(address account) internal view {
         (string memory category, uint8 confidenceScore) = _threatOracle.getThreatProperties(account);
         bytes32 categoryHashed = keccak256(abi.encodePacked(category));
@@ -106,4 +116,11 @@ abstract contract ThreatOracleProvider {
             confidenceScore >= _minConfidenceScore
         ) revert ThreatAccountIdentified(account, category, confidenceScore);
     }
+
+    function _updateThreatOracleContractAddress(address newThreatOracleAddress) internal {
+        address prevThreatOracleAddress = address(_threatOracle);
+        _threatOracle = IThreatOracle(newThreatOracleAddress);
+
+        emit ThreatOracleContractUpdated(prevThreatOracleAddress, newThreatOracleAddress);
+    }
 }
diff --git a/contracts/mocks/threat_oracle/MockThreatOracleConsumer.sol b/contracts/mocks/threat_oracle/MockThreatOracleConsumer.sol
@@ -7,9 +7,25 @@ import "../../components/threat_oracle/ThreatOracleProvider.sol";
 
 contract MockThreatOracleConsumer is ThreatOracleProvider {
 
-    constructor(address _threatOracle, uint8 _minConfidenceScore, uint8 __maxAddressArgumentAmount)
+    address public owner;
+
+    error MsgSenderNotOwner(address msgSender);
+
+    modifier onlyOwner(address account) {
+        if (account != owner) revert MsgSenderNotOwner(account);
+        _;
+    }
+
+    constructor(
+        address _threatOracle,
+        uint8 _minConfidenceScore,
+        uint8 __maxAddressArgumentAmount,
+        address __owner
+    )
         ThreatOracleProvider(_threatOracle, _minConfidenceScore, __maxAddressArgumentAmount)
-    {}
+    {
+        owner = __owner;
+    }
 
     function foo() public view onlyNonThreatMsgSenderAndTxOrigin() returns (bool) {
         return true;
@@ -47,4 +63,11 @@ contract MockThreatOracleConsumer is ThreatOracleProvider {
         return true;
     }
 
+    function updateThreatOracleContractAddress(address newThreatOracleAddress)
+        external
+        onlyOwner(msg.sender)
+    {
+        _updateThreatOracleContractAddress(newThreatOracleAddress);
+    }
+
 }