diff --git a/public/samples/VRF/v2-5/SubscriptionManager.sol b/public/samples/VRF/v2-5/SubscriptionManager.sol index 7338af92640..f806eef2e87 100644 --- a/public/samples/VRF/v2-5/SubscriptionManager.sol +++ b/public/samples/VRF/v2-5/SubscriptionManager.sol @@ -19,6 +19,10 @@ import {VRFV2PlusClient} from "@chainlink/contracts/src/v0.8/vrf/dev/libraries/V */ contract VRFv2PlusSubscriptionManager is VRFConsumerBaseV2Plus { + + error InsufficientValue(uint256 required); + error NotEnoughBalance(uint256 available); + LinkTokenInterface LINKTOKEN; // Sepolia coordinator. For other networks, @@ -61,7 +65,7 @@ contract VRFv2PlusSubscriptionManager is VRFConsumerBaseV2Plus { } // Assumes the subscription is funded sufficiently. - function requestRandomWords() external onlyOwner { + function requestRandomWords(bool enableNativePayment) external onlyOwner { // Will revert if subscription is not set and funded. s_requestId = s_vrfCoordinator.requestRandomWords( VRFV2PlusClient.RandomWordsRequest({ @@ -71,7 +75,7 @@ contract VRFv2PlusSubscriptionManager is VRFConsumerBaseV2Plus { callbackGasLimit: callbackGasLimit, numWords: numWords, extraArgs: VRFV2PlusClient._argsToBytes( - VRFV2PlusClient.ExtraArgsV1({nativePayment: false}) + VRFV2PlusClient.ExtraArgsV1({nativePayment: enableNativePayment}) ) }) ); @@ -101,6 +105,18 @@ contract VRFv2PlusSubscriptionManager is VRFConsumerBaseV2Plus { ); } + // 1000000000000000000 = 1 ether / native token + function topUpSubscriptionWithNativeToken(uint256 amount) external payable onlyOwner + { + if (msg.value < amount) { + revert InsufficientValue(amount); + } + + s_vrfCoordinator.fundSubscriptionWithNative{value: amount}( + s_subscriptionId + ); + } + function addConsumer(address consumerAddress) external onlyOwner { // Add a consumer contract to the subscription. s_vrfCoordinator.addConsumer(s_subscriptionId, consumerAddress); @@ -122,4 +138,13 @@ contract VRFv2PlusSubscriptionManager is VRFConsumerBaseV2Plus { function withdraw(uint256 amount, address to) external onlyOwner { LINKTOKEN.transfer(to, amount); } + + // Transfer this contract's native token balance to an address. + // 1000000000000000000 = 1 ether / native token + function withdrawNativeToken(uint256 amount, address payable to) external onlyOwner { + if (address(this).balance < amount) { + revert NotEnoughBalance(address(this).balance); + } + to.transfer(amount); + } }