add functions to reclaim leases and test

contracts/LockedTokens.cdc

@@ -82,8 +82,17 @@ access(all) contract LockedTokens {
         access(UnlockTokens) fun increaseUnlockLimit(delta: UFix64)
     }
 
+    /// Entitlement for the token admin to unlock tokens from
+    /// the token sale and various grants
     access(all) entitlement UnlockTokens
 
+    /// Entitlement for the token admin to use to recover leased tokens
+    /// directly from accounts who have leased tokens to operate nodes
+    /// Since there are no existing capabilities with this entitlement,
+    /// it is only used when authorizing a transaction that already
+    /// has access to the account
+    access(all) entitlement RecoverLease
+
     /// This token manager resource is stored in the shared account to manage access
     /// to the locked token vault and to the staking/delegating resources.
     access(all) resource LockedTokenManager: FungibleToken.Receiver, FungibleToken.Provider, TokenAdmin {
@@ -258,6 +267,19 @@ access(all) contract LockedTokens {
             return delegatorRef
         }
 
+        /// The following two functions are late additions to replicate functionality
+        /// that was lost with the Crescendo upgrade
+        /// They allow the account that owns the TokenManager to borrow a reference to
+        /// the node or delegator directly from its storage
+        /// This is only used by the Flow Foundation to recover leases that it has given to node operators
+        access(RecoverLease) view fun borrowNodeForLease(): auth(FlowIDTableStaking.NodeOperator) &FlowIDTableStaking.NodeStaker? {
+            return self.borrowNode()
+        }
+
+        access(RecoverLease) view fun borrowDelegatorForLease(): auth(FlowIDTableStaking.DelegatorOwner) &FlowIDTableStaking.NodeDelegator? {
+            return self.borrowDelegator()
+        }
+
         access(UnlockTokens) fun removeNode(): @FlowIDTableStaking.NodeStaker? {
             let node <- self.nodeStaker <- nil
 

lib/go/templates/lockedtokens_templates.go

@@ -16,6 +16,7 @@ const (
 	depositAccountCreatorCapabilityFilename        = "lockedTokens/admin/admin_deposit_account_creator.cdc"
 	removeDelegatorFilename                        = "lockedTokens/admin/admin_remove_delegator.cdc"
 	getBadAccountsFilename                         = "lockedTokens/admin/get_unlocking_bad_accounts.cdc"
+	recoverLeaseTokensFilename                     = "lockedTokens/admin/recover_lease_tokens.cdc"
 
 	// Custody Provider / Wallet provider Account creation templates
 	setupCustodyAccountFilename                  = "lockedTokens/admin/custody_setup_account_creator.cdc"
@@ -120,6 +121,12 @@ func GenerateGetBadAccountsScript(env Environment) []byte {
 	return []byte(ReplaceAddresses(code, env))
 }
 
+func GenerateRecoverLeaseTokensScript(env Environment) []byte {
+	code := assets.MustAssetString(recoverLeaseTokensFilename)
+
+	return []byte(ReplaceAddresses(code, env))
+}
+
 /************ Custody Provider Transactions ********************/
 
 func GenerateSetupCustodyAccountScript(env Environment) []byte {

lib/go/test/flow_lockedtokens_test.go

@@ -627,6 +627,27 @@ func TestLockedTokensStaker(t *testing.T) {
 		result = executeScriptAndCheck(t, b, templates.GenerateGetUnlockLimitScript(env), [][]byte{jsoncdc.MustEncode(cadence.Address(joshAddress))})
 		assertEqual(t, CadenceUFix64("8500.0"), result)
 	})
+
+	t.Run("Should be able to claim leased tokens as the admin", func(t *testing.T) {
+
+		script := templates.GenerateRecoverLeaseTokensScript(env)
+
+		tx := createTxWithTemplateAndAuthorizer(b, script, joshSharedAddress)
+
+		_ = tx.AddArgument(CadenceUFix64("10.0"))
+		_ = tx.AddArgument(cadence.NewAddress(adminAddress))
+
+		signAndSubmit(
+			t, b, tx,
+			[]flow.Address{joshSharedAddress},
+			[]crypto.Signer{adminSigner},
+			false,
+		)
+
+		// Check balance of locked account
+		result := executeScriptAndCheck(t, b, templates.GenerateGetFlowBalanceScript(env), [][]byte{jsoncdc.MustEncode(cadence.Address(joshSharedAddress))})
+		assertEqual(t, CadenceUFix64("0.0001"), result)
+	})
 }
 
 func TestLockedTokensDelegator(t *testing.T) {

transactions/lockedTokens/admin/recover_lease_tokens.cdc

@@ -0,0 +1,40 @@
+import "FungibleToken"
+import "FlowToken"
+import "LockedTokens"
+
+transaction(amount: UFix64, to: Address) {
+
+    // The Vault resource that holds the tokens that are being transferred
+    let sentVault: @{FungibleToken.Vault}
+
+    prepare(signer: auth(BorrowValue) &Account) {
+
+        // Get a reference to the signer's locked token manager
+        let tokenManagerRef = signer.storage.borrow<auth(FungibleToken.Withdraw, LockedTokens.RecoverLease) &LockedTokens.LockedTokenManager>(from: LockedTokens.LockedTokenManagerStoragePath)
+            ?? panic("The signer does not store a LockedTokenManager object at the path "
+                    .concat(LockedTokens.LockedTokenManagerStoragePath.toString()))
+
+        let nodeRef = tokenManagerRef.borrowNodeForLease()
+            ?? panic("Could not borrow a reference to a node in the LockedTokenManager of the signer's account")
+
+        // Withdraw enough tokens to pay for fees, assuming there are some rewards in the rewards bucket
+        tokenManagerRef.deposit(from: <-nodeRef.withdrawRewardedTokens(amount: 0.0001)!)
+
+        // Withdraw tokens from the signer's stored vault
+        self.sentVault <- nodeRef.withdrawUnstakedTokens(amount: amount)!
+    }
+
+    execute {
+
+        // Get a reference to the recipient's Receiver
+        let receiverRef =  getAccount(to)
+            .capabilities.borrow<&{FungibleToken.Receiver}>(/public/flowTokenReceiver)
+            ?? panic("Could not borrow a Receiver reference to the FlowToken Vault in account "
+                .concat(to.toString()).concat(" at path /public/flowTokenReceiver")
+                .concat(". Make sure you are sending to an address that has ")
+                .concat("a FlowToken Vault set up properly at the specified path."))
+
+        // Deposit the withdrawn tokens in the recipient's receiver
+        receiverRef.deposit(from: <-self.sentVault)
+    }
+}