Update SubscriptionManager to latest version

See https://github.com/nucypher/nucypher-contracts/pull/2

nucypher/blockchain/eth/sol/source/contracts/SubscriptionManager.sol

@@ -2,38 +2,56 @@
 
 pragma solidity ^0.8.0;
 
-contract SubscriptionManager {
+import "@openzeppelin-upgradeable/contracts/access/AccessControlUpgradeable.sol";
+import "@openzeppelin-upgradeable/contracts/proxy/utils/Initializable.sol";
 
-    uint256 private constant RATE_PER_DAY = 50 gwei;
-    uint256 public constant RATE_PER_SECOND = RATE_PER_DAY / 1 days;
+contract SubscriptionManager is Initializable, AccessControlUpgradeable {
 
-    struct Policy { // TODO: Optimize struct layout
+    bytes32 public constant SET_RATE_ROLE =
+        keccak256("Power to set the fee rate");
+    bytes32 public constant WITHDRAW_ROLE =
+        keccak256("Power to withdraw funds from SubscriptionManager");
+
+    // The layout of policy struct is optimized, so sponsor, timestamps and size
+    // fit in a single 256-word.
+    struct Policy {
         address payable sponsor;
+        uint32 startTimestamp;
+        uint32 endTimestamp;
+        uint16 size; // also known as `N`
         address owner;
-        uint64 startTimestamp;
-        uint64 endTimestamp;
     }
 
     event PolicyCreated(
         bytes16 indexed policyId,
         address indexed sponsor,
         address indexed owner,
-        uint64 startTimestamp,
-        uint64 endTimestamp
+        uint16 size,
+        uint32 startTimestamp,
+        uint32 endTimestamp
     );
-    
-    address payable public owner;
-    mapping (bytes16 => Policy) public policies;
 
-    constructor(){
-        owner = payable(msg.sender);
+    event FeeRateUpdated(uint256 oldFeeRate, uint256 newFeeRate);
+
+    // Per-second service fee rate
+    uint256 public feeRate;
+
+    // Mapping that stores policy structs, keyed by policy ID
+    mapping (bytes16 => Policy) internal _policies;
+
+    function initialize(uint256 _feeRate) public initializer {
+        _setFeeRate(_feeRate);
+        _setupRole(SET_RATE_ROLE, msg.sender);
+        _setupRole(WITHDRAW_ROLE, msg.sender);
+        _setupRole(DEFAULT_ADMIN_ROLE, msg.sender);
     }
 
     function createPolicy(
         bytes16 _policyId,
         address _policyOwner,
-        uint64 _startTimestamp,
-        uint64 _endTimestamp
+        uint16 _size,
+        uint32 _startTimestamp,
+        uint32 _endTimestamp
     )
         external payable
     {
@@ -42,31 +60,33 @@ contract SubscriptionManager {
             block.timestamp < _endTimestamp,
             "Invalid timestamps"
         );
-        uint64 duration = _endTimestamp - _startTimestamp;
+        uint32 duration = _endTimestamp - _startTimestamp;
         require(
-            duration > 0 &&
-            msg.value == RATE_PER_SECOND * uint64(duration)
+            duration > 0 && _size > 0 &&
+            msg.value == feeRate * _size * uint32(duration)
         );
-        //Policy storage policy = 
-        _createPolicy(_policyId, _policyOwner, _startTimestamp, _endTimestamp);
+
+        _createPolicy(_policyId, _policyOwner, _size, _startTimestamp, _endTimestamp);
     }
 
     /**
     * @notice Create policy
     * @param _policyId Policy id
     * @param _policyOwner Policy owner. Zero address means sender is owner
+    * @param _size Number of nodes involved in the policy
     * @param _startTimestamp Start timestamp of the policy in seconds
     * @param _endTimestamp End timestamp of the policy in seconds
     */
     function _createPolicy(
         bytes16 _policyId,
         address _policyOwner,
-        uint64 _startTimestamp,
-        uint64 _endTimestamp
+        uint16 _size,
+        uint32 _startTimestamp,
+        uint32 _endTimestamp
     )
         internal returns (Policy storage policy)
     {
-        policy = policies[_policyId];
+        policy = _policies[_policyId];
         require(
             policy.endTimestamp < block.timestamp,
             "Policy is currently active"
@@ -75,6 +95,7 @@ contract SubscriptionManager {
         policy.sponsor = payable(msg.sender);
         policy.startTimestamp = _startTimestamp;
         policy.endTimestamp = _endTimestamp;
+        policy.size = _size;
 
         if (_policyOwner != msg.sender && _policyOwner != address(0)) {
             policy.owner = _policyOwner;
@@ -84,20 +105,34 @@ contract SubscriptionManager {
             _policyId,
             msg.sender,
             _policyOwner == address(0) ? msg.sender : _policyOwner,
+            _size,
             _startTimestamp,
             _endTimestamp
         );
     }
 
-    function isPolicyActive(bytes16 _policyID) public view returns(bool){
-        return policies[_policyID].endTimestamp > block.timestamp;
+    function getPolicy(bytes16 _policyID) public view returns(Policy memory){
+        return _policies[_policyID];
     }
 
-    function sweep(address payable recipient) external {
-        require(msg.sender == owner);
+    function isPolicyActive(bytes16 _policyID) public view returns(bool){
+        return _policies[_policyID].endTimestamp > block.timestamp;
+    }
+
+    function _setFeeRate(uint256 newFee) internal {
+        uint256 oldFee = feeRate;
+        feeRate = newFee;
+        emit FeeRateUpdated(oldFee, newFee);
+    }
+
+    function setFeeRate(uint256 _rate_per_second) onlyRole(SET_RATE_ROLE) external {
+        _setFeeRate(_rate_per_second);
+    }
+
+    function sweep(address payable recipient) onlyRole(WITHDRAW_ROLE) external {
         uint256 balance = address(this).balance;
         (bool sent, ) = recipient.call{value: balance}("");
         require(sent, "Failed transfer");
     }
 
-}
+}