[token] add partial withdraw with capability (#5652)

aptos-labs · Dec 1, 2022 · c4f5fb6 · c4f5fb6
1 parent 21660ed
commit c4f5fb6
Show file tree

Hide file tree

Showing 2 changed files with 137 additions and 0 deletions.
diff --git a/aptos-move/framework/aptos-token/doc/token.md b/aptos-move/framework/aptos-token/doc/token.md
@@ -55,6 +55,7 @@ Checkout our developer doc on our token standard https://aptos.dev/concepts/coin
 -  [Function `transfer`](#0x3_token_transfer)
 -  [Function `create_withdraw_capability`](#0x3_token_create_withdraw_capability)
 -  [Function `withdraw_with_capability`](#0x3_token_withdraw_with_capability)
+-  [Function `partial_withdraw_with_capability`](#0x3_token_partial_withdraw_with_capability)
 -  [Function `withdraw_token`](#0x3_token_withdraw_token)
 -  [Function `create_collection`](#0x3_token_create_collection)
 -  [Function `check_collection_exists`](#0x3_token_check_collection_exists)
@@ -1135,6 +1136,16 @@ The field is not mutable
 
 
 
+<a name="0x3_token_EINSUFFICIENT_WITHDRAW_CAPABILITY_AMOUNT"></a>
+
+Withdraw capability doesn't have sufficient amount
+
+
+<pre><code><b>const</b> <a href="token.md#0x3_token_EINSUFFICIENT_WITHDRAW_CAPABILITY_AMOUNT">EINSUFFICIENT_WITHDRAW_CAPABILITY_AMOUNT</a>: u64 = 37;
+</code></pre>
+
+
+
 <a name="0x3_token_EINVALID_MAXIMUM"></a>
 
 Collection or tokendata maximum must be larger than supply
@@ -2657,6 +2668,60 @@ Withdraw the token with a capability
 
 
 
+</details>
+
+<a name="0x3_token_partial_withdraw_with_capability"></a>
+
+## Function `partial_withdraw_with_capability`
+
+Withdraw the token with a capability.
+
+
+<pre><code><b>public</b> <b>fun</b> <a href="token.md#0x3_token_partial_withdraw_with_capability">partial_withdraw_with_capability</a>(withdraw_proof: <a href="token.md#0x3_token_WithdrawCapability">token::WithdrawCapability</a>, withdraw_amount: u64): (<a href="token.md#0x3_token_Token">token::Token</a>, <a href="../../aptos-framework/../aptos-stdlib/../move-stdlib/doc/option.md#0x1_option_Option">option::Option</a>&lt;<a href="token.md#0x3_token_WithdrawCapability">token::WithdrawCapability</a>&gt;)
+</code></pre>
+
+
+
+<details>
+<summary>Implementation</summary>
+
+
+<pre><code><b>public</b> <b>fun</b> <a href="token.md#0x3_token_partial_withdraw_with_capability">partial_withdraw_with_capability</a>(
+    withdraw_proof: <a href="token.md#0x3_token_WithdrawCapability">WithdrawCapability</a>,
+    withdraw_amount: u64,
+): (<a href="token.md#0x3_token_Token">Token</a>, Option&lt;<a href="token.md#0x3_token_WithdrawCapability">WithdrawCapability</a>&gt;) <b>acquires</b> <a href="token.md#0x3_token_TokenStore">TokenStore</a> {
+    // verify the delegation hasn't expired yet
+    <b>assert</b>!(<a href="../../aptos-framework/doc/timestamp.md#0x1_timestamp_now_seconds">timestamp::now_seconds</a>() &lt;= *&withdraw_proof.expiration_sec, <a href="../../aptos-framework/../aptos-stdlib/../move-stdlib/doc/error.md#0x1_error_invalid_argument">error::invalid_argument</a>(<a href="token.md#0x3_token_EWITHDRAW_PROOF_EXPIRES">EWITHDRAW_PROOF_EXPIRES</a>));
+
+    <b>assert</b>!(withdraw_amount &lt;= withdraw_proof.amount, <a href="../../aptos-framework/../aptos-stdlib/../move-stdlib/doc/error.md#0x1_error_invalid_argument">error::invalid_argument</a>(<a href="token.md#0x3_token_EINSUFFICIENT_WITHDRAW_CAPABILITY_AMOUNT">EINSUFFICIENT_WITHDRAW_CAPABILITY_AMOUNT</a>));
+
+    <b>let</b> res: Option&lt;<a href="token.md#0x3_token_WithdrawCapability">WithdrawCapability</a>&gt; = <b>if</b> (withdraw_amount == withdraw_proof.amount) {
+        <a href="../../aptos-framework/../aptos-stdlib/../move-stdlib/doc/option.md#0x1_option_none">option::none</a>&lt;<a href="token.md#0x3_token_WithdrawCapability">WithdrawCapability</a>&gt;()
+    } <b>else</b> {
+        <a href="../../aptos-framework/../aptos-stdlib/../move-stdlib/doc/option.md#0x1_option_some">option::some</a>(
+            <a href="token.md#0x3_token_WithdrawCapability">WithdrawCapability</a> {
+                token_owner: withdraw_proof.token_owner,
+                token_id: withdraw_proof.token_id,
+                amount: withdraw_proof.amount - withdraw_amount,
+                expiration_sec: withdraw_proof.expiration_sec,
+            }
+        )
+    };
+
+    (
+        <a href="token.md#0x3_token_withdraw_with_event_internal">withdraw_with_event_internal</a>(
+            withdraw_proof.token_owner,
+            withdraw_proof.token_id,
+            withdraw_amount,
+        ),
+        res
+    )
+
+}
+</code></pre>
+
+
+
 </details>
 
 <a name="0x3_token_withdraw_token"></a>

diff --git a/aptos-move/framework/aptos-token/sources/token.move b/aptos-move/framework/aptos-token/sources/token.move
@@ -145,6 +145,10 @@ module aptos_token::token {
     /// Token Properties count doesn't match
     const ETOKEN_PROPERTIES_COUNT_NOT_MATCH: u64 = 37;
 
+
+    /// Withdraw capability doesn't have sufficient amount
+    const EINSUFFICIENT_WITHDRAW_CAPABILITY_AMOUNT: u64 = 38;
+
     //
     // Core data structures for holding tokens
     //
@@ -963,6 +967,40 @@ module aptos_token::token {
         )
     }
 
+    /// Withdraw the token with a capability.
+    public fun partial_withdraw_with_capability(
+        withdraw_proof: WithdrawCapability,
+        withdraw_amount: u64,
+    ): (Token, Option<WithdrawCapability>) acquires TokenStore {
+        // verify the delegation hasn't expired yet
+        assert!(timestamp::now_seconds() <= *&withdraw_proof.expiration_sec, error::invalid_argument(EWITHDRAW_PROOF_EXPIRES));
+
+        assert!(withdraw_amount <= withdraw_proof.amount, error::invalid_argument(EINSUFFICIENT_WITHDRAW_CAPABILITY_AMOUNT));
+
+        let res: Option<WithdrawCapability> = if (withdraw_amount == withdraw_proof.amount) {
+            option::none<WithdrawCapability>()
+        } else {
+            option::some(
+                WithdrawCapability {
+                    token_owner: withdraw_proof.token_owner,
+                    token_id: withdraw_proof.token_id,
+                    amount: withdraw_proof.amount - withdraw_amount,
+                    expiration_sec: withdraw_proof.expiration_sec,
+                }
+            )
+        };
+
+        (
+            withdraw_with_event_internal(
+                withdraw_proof.token_owner,
+                withdraw_proof.token_id,
+                withdraw_amount,
+            ),
+            res
+        )
+
+    }
+
     public fun withdraw_token(
         account: &signer,
         id: TokenId,
@@ -2448,6 +2486,40 @@ module aptos_token::token {
         create_royalty(101, 100, @0xcafe);
     }
 
+    #[test(framework = @0x1, creator = @0xcafe)]
+    fun test_partial_withdraw_with_proof(creator: &signer, framework: &signer): Token acquires TokenStore, Collections {
+        timestamp::set_time_has_started_for_testing(framework);
+        account::create_account_for_test(signer::address_of(creator));
+        let token_id = create_collection_and_token(
+            creator,
+            4,
+            4,
+            4,
+            vector<String>[],
+            vector<vector<u8>>[],
+            vector<String>[],
+            vector<bool>[false, false, false],
+            vector<bool>[false, false, false, false, false],
+        );
+
+        timestamp::update_global_time_for_test(1000000);
+
+        // provide the proof to the account
+        let cap = create_withdraw_capability(
+            creator, // ask user to provide address to avoid ambiguity from rotated keys
+            token_id,
+            3,
+            2000000,
+        );
+
+        let (token, capability) = partial_withdraw_with_capability(cap, 1);
+        assert!(option::borrow<WithdrawCapability>(&capability).amount == 2, 1);
+        let (token_1, cap) = partial_withdraw_with_capability(option::extract(&mut capability), 2);
+        assert!(option::is_none(&cap), 1);
+        merge(&mut token, token_1);
+        token
+    }
+
     //
     // Deprecated functions
     //