From 2e564d81a0f1212c7b66cbfc4a9986a0dc9b1d49 Mon Sep 17 00:00:00 2001
From: David Wolinsky <isaac.wolinsky@gmail.com>
Date: Tue, 14 Feb 2023 13:12:15 -0800
Subject: [PATCH] [framework] fix object linear transfer ref

Make sure that linear transfers can only be used if the owner field is
correct.
---
 .../framework/aptos-framework/doc/object.md   |  4 ++
 .../aptos-framework/sources/object.move       | 48 +++++++++++++++----
 2 files changed, 42 insertions(+), 10 deletions(-)

diff --git a/aptos-move/framework/aptos-framework/doc/object.md b/aptos-move/framework/aptos-framework/doc/object.md
index 189cc30237d9c..98c90699158be 100644
--- a/aptos-move/framework/aptos-framework/doc/object.md
+++ b/aptos-move/framework/aptos-framework/doc/object.md
@@ -1109,6 +1109,10 @@ Transfer to the destination address using a LinearTransferRef.
 
 <pre><code><b>public</b> <b>fun</b> <a href="object.md#0x1_object_transfer_with_ref">transfer_with_ref</a>(ref: <a href="object.md#0x1_object_LinearTransferRef">LinearTransferRef</a>, <b>to</b>: <b>address</b>) <b>acquires</b> <a href="object.md#0x1_object_ObjectCore">ObjectCore</a> {
     <b>let</b> <a href="object.md#0x1_object">object</a> = <b>borrow_global_mut</b>&lt;<a href="object.md#0x1_object_ObjectCore">ObjectCore</a>&gt;(ref.self);
+    <b>assert</b>!(
+        <a href="object.md#0x1_object">object</a>.owner == ref.owner,
+        <a href="../../aptos-stdlib/../move-stdlib/doc/error.md#0x1_error_permission_denied">error::permission_denied</a>(<a href="object.md#0x1_object_ENOT_OBJECT_OWNER">ENOT_OBJECT_OWNER</a>),
+    );
     <a href="event.md#0x1_event_emit_event">event::emit_event</a>(
         &<b>mut</b> <a href="object.md#0x1_object">object</a>.transfer_events,
         <a href="object.md#0x1_object_TransferEvent">TransferEvent</a> {
diff --git a/aptos-move/framework/aptos-framework/sources/object.move b/aptos-move/framework/aptos-framework/sources/object.move
index 229bc985f2fbd..95d3ab777d027 100644
--- a/aptos-move/framework/aptos-framework/sources/object.move
+++ b/aptos-move/framework/aptos-framework/sources/object.move
@@ -296,6 +296,10 @@ module aptos_framework::object {
     /// Transfer to the destination address using a LinearTransferRef.
     public fun transfer_with_ref(ref: LinearTransferRef, to: address) acquires ObjectCore {
         let object = borrow_global_mut<ObjectCore>(ref.self);
+        assert!(
+            object.owner == ref.owner,
+            error::permission_denied(ENOT_OBJECT_OWNER),
+        );
         event::emit_event(
             &mut object.transfer_events,
             TransferEvent {
@@ -441,9 +445,9 @@ module aptos_framework::object {
     struct Weapon has key { }
 
     #[test_only]
-    public fun create_hero(creator: &signer): Object<Hero> acquires ObjectCore {
-        let hero_creator_ref = create_named_object(creator, b"hero");
-        let hero_signer = generate_signer(&hero_creator_ref);
+    public fun create_hero(creator: &signer): (ConstructorRef, Object<Hero>) acquires ObjectCore {
+        let hero_constructor_ref = create_named_object(creator, b"hero");
+        let hero_signer = generate_signer(&hero_constructor_ref);
         let guid_for_equip_events = create_guid(&hero_signer);
         move_to(
             &hero_signer,
@@ -453,15 +457,17 @@ module aptos_framework::object {
             },
         );
 
-        object_from_constructor_ref<Hero>(&hero_creator_ref)
+        let hero = object_from_constructor_ref<Hero>(&hero_constructor_ref);
+        (hero_constructor_ref, hero)
     }
 
     #[test_only]
-    public fun create_weapon(creator: &signer): Object<Weapon> {
-        let weapon_creator_ref = create_named_object(creator, b"weapon");
-        let weapon_signer = generate_signer(&weapon_creator_ref);
+    public fun create_weapon(creator: &signer): (ConstructorRef, Object<Weapon>) {
+        let weapon_constructor_ref = create_named_object(creator, b"weapon");
+        let weapon_signer = generate_signer(&weapon_constructor_ref);
         move_to(&weapon_signer, Weapon { });
-        object_from_constructor_ref<Weapon>(&weapon_creator_ref)
+        let weapon = object_from_constructor_ref<Weapon>(&weapon_constructor_ref);
+        (weapon_constructor_ref, weapon)
     }
 
     #[test_only]
@@ -496,10 +502,32 @@ module aptos_framework::object {
 
     #[test(creator = @0x123)]
     fun test_object(creator: &signer) acquires Hero, ObjectCore {
-        let hero = create_hero(creator);
-        let weapon = create_weapon(creator);
+        let (_, hero) = create_hero(creator);
+        let (_, weapon) = create_weapon(creator);
 
         hero_equip(creator, hero, weapon);
         hero_unequip(creator, hero, weapon);
     }
+
+    #[test(creator = @0x123)]
+    fun test_linear_transfer(creator: &signer) acquires ObjectCore {
+        let (hero_constructor, hero) = create_hero(creator);
+        let transfer_ref = generate_transfer_ref(&hero_constructor);
+        let linear_transfer_ref = generate_linear_transfer_ref(&transfer_ref);
+        transfer_with_ref(linear_transfer_ref, @0x456);
+        assert!(owner(hero) == @0x456, 0);
+    }
+
+    #[test(creator = @0x123)]
+    #[expected_failure(abort_code = 0x50004, location = Self)]
+    fun test_bad_linear_transfer(creator: &signer) acquires ObjectCore {
+        let (hero_constructor, hero) = create_hero(creator);
+        let transfer_ref = generate_transfer_ref(&hero_constructor);
+        let linear_transfer_ref_good = generate_linear_transfer_ref(&transfer_ref);
+        // This will contain the address of the creator
+        let linear_transfer_ref_bad = generate_linear_transfer_ref(&transfer_ref);
+        transfer_with_ref(linear_transfer_ref_good, @0x456);
+        assert!(owner(hero) == @0x456, 0);
+        transfer_with_ref(linear_transfer_ref_bad, @0x789);
+    }
 }