diff --git a/contracts/universalSchemes/VestingScheme.sol b/contracts/universalSchemes/VestingScheme.sol index d1c72fae..cbddae9d 100644 --- a/contracts/universalSchemes/VestingScheme.sol +++ b/contracts/universalSchemes/VestingScheme.sol @@ -39,6 +39,7 @@ contract VestingScheme is UniversalScheme, VotingMachineCallbacks, ProposalExecu uint256 signaturesReqToCancel; uint256 collectedPeriods; uint256 signaturesReceivedCounter; + Avatar avatar; mapping(address=>bool) signers; mapping(address=>bool) signaturesReceived; } @@ -90,6 +91,7 @@ contract VestingScheme is UniversalScheme, VotingMachineCallbacks, ProposalExecu ControllerInterface controller = ControllerInterface(avatar.owner()); uint256 tokensToMint = proposedAgreement.amountPerPeriod.mul(proposedAgreement.numOfAgreedPeriods); require(controller.mintTokens(tokensToMint, address(this), address(avatar))); + proposedAgreement.avatar = avatar; agreements[agreementsCounter] = proposedAgreement; agreementsCounter++; // Log the new agreement: @@ -291,6 +293,10 @@ contract VestingScheme is UniversalScheme, VotingMachineCallbacks, ProposalExecu */ function collect(uint256 _agreementId) public onlyBeneficiary(_agreementId) { Agreement memory agreement = agreements[_agreementId]; + if (agreement.avatar != Avatar(0)) { + require(ControllerInterface(agreement.avatar.owner()) + .isSchemeRegistered(address(this), address(agreement.avatar))); + } uint256 periodsFromStartingBlock = (block.number.sub(agreement.startingBlock)).div(agreement.periodLength); require(periodsFromStartingBlock >= agreement.cliffInPeriods); @@ -319,6 +325,10 @@ contract VestingScheme is UniversalScheme, VotingMachineCallbacks, ProposalExecu function cancelAgreement(uint256 _agreementId) internal { Agreement memory agreement = agreements[_agreementId]; delete agreements[_agreementId]; + if (agreement.avatar != Avatar(0)) { + require(ControllerInterface(agreement.avatar.owner()) + .isSchemeRegistered(address(this), address(agreement.avatar))); + } uint256 periodsLeft = agreement.numOfAgreedPeriods.sub(agreement.collectedPeriods); uint256 tokensLeft = periodsLeft.mul(agreement.amountPerPeriod); address(agreement.token).safeTransfer(agreement.returnOnCancelAddress, tokensLeft);