From 00890ef9e020fdcb202842763f9a06085fcd1629 Mon Sep 17 00:00:00 2001
From: Michael Vines <mvines@gmail.com>
Date: Sat, 20 Oct 2018 20:21:25 -0500
Subject: [PATCH] fix: add Token::setOwner

---
 web3.js/module.flow.js             |  9 ++++--
 web3.js/src/token-program.js       | 47 ++++++++++++++++++++++++++----
 web3.js/test/token-program.test.js | 31 ++++++++++++++++++--
 3 files changed, 76 insertions(+), 11 deletions(-)

diff --git a/web3.js/module.flow.js b/web3.js/module.flow.js
index 677a19d5fe..df005b2392 100644
--- a/web3.js/module.flow.js
+++ b/web3.js/module.flow.js
@@ -145,15 +145,20 @@ declare module '@solana/web3.js' {
     ): Promise<void>;
     approve(
       owner: Account,
-      source: PublicKey,
+      account: PublicKey,
       delegate: PublicKey,
       amount: number | TokenAmount
     ): Promise<void>;
     revoke(
       owner: Account,
-      source: PublicKey,
+      account: PublicKey,
       delegate: PublicKey
     ): Promise<void>;
+    setOwner(
+      owner: Account,
+      account: PublicKey,
+      newOwner: PublicKey
+    ): Promise<void>;
   }
 
   // === src/loader.js ===
diff --git a/web3.js/src/token-program.js b/web3.js/src/token-program.js
index cd35200c56..de7c7d8fc3 100644
--- a/web3.js/src/token-program.js
+++ b/web3.js/src/token-program.js
@@ -402,13 +402,13 @@ export class Token {
    * Grant a third-party permission to transfer up the specified number of tokens from an account
    *
    * @param owner Owner of the source token account
-   * @param source Source token account
+   * @param account Public key of the token account
    * @param delegate Token account authorized to perform a transfer tokens from the source account
    * @param amount Maximum number of tokens the delegate may transfer
    */
   async approve(
     owner: Account,
-    source: PublicKey,
+    account: PublicKey,
     delegate: PublicKey,
     amount: number | TokenAmount
   ): Promise<void> {
@@ -429,7 +429,7 @@ export class Token {
 
     const transaction = new Transaction({
       fee: 0,
-      keys: [owner.publicKey, source, delegate],
+      keys: [owner.publicKey, account, delegate],
       programId: this.programId,
       userdata,
     });
@@ -440,15 +440,50 @@ export class Token {
    * Remove approval for the transfer of any remaining tokens
    *
    * @param owner Owner of the source token account
-   * @param source Source token account
+   * @param account Public key of the token account
    * @param delegate Token account to revoke authorization from
    */
   revoke(
     owner: Account,
-    source: PublicKey,
+    account: PublicKey,
     delegate: PublicKey
   ): Promise<void> {
-    return this.approve(owner, source, delegate, 0);
+    return this.approve(owner, account, delegate, 0);
+  }
+
+  /**
+   * Assign a new owner to the account
+   *
+   * @param owner Owner of the token account
+   * @param account Public key of the token account
+   * @param newOwner New owner of the token account
+   */
+  async setOwner(
+    owner: Account,
+    account: PublicKey,
+    newOwner: PublicKey,
+  ): Promise<void> {
+
+    const userdataLayout = BufferLayout.struct([
+      BufferLayout.u32('instruction'),
+    ]);
+
+    const userdata = Buffer.alloc(userdataLayout.span);
+    userdataLayout.encode(
+      {
+        instruction: 4, // SetOwner instruction
+      },
+      userdata,
+    );
+
+    const keys = [owner.publicKey, account,newOwner];
+    const transaction = new Transaction({
+      fee: 0,
+      keys,
+      programId: this.programId,
+      userdata,
+    });
+    await sendAndConfirmTransaction(this.connection, owner, transaction);
   }
 }
 
diff --git a/web3.js/test/token-program.test.js b/web3.js/test/token-program.test.js
index eedae795e4..26958ff886 100644
--- a/web3.js/test/token-program.test.js
+++ b/web3.js/test/token-program.test.js
@@ -473,7 +473,7 @@ test('invalid approve', async () => {
   const account2 = await testToken.newAccount(owner);
 
   // account2 is not a delegate account of account1
-  expect(
+  await expect(
     testToken.approve(
       owner,
       account1,
@@ -483,7 +483,7 @@ test('invalid approve', async () => {
   ).rejects.toThrow();
 
   // account1Delegate is not a delegate account of account2
-  expect(
+  await expect(
     testToken.approve(
       owner,
       account2,
@@ -548,7 +548,7 @@ test.skip('fail on approve overspend', async () => {
   expect(delegateAccountInfo.amount.toNumber()).toBe(0);
   expect(delegateAccountInfo.originalAmount.toNumber()).toBe(2);
 
-  expect(
+  await expect(
     testToken.transfer(
       owner,
       account1Delegate,
@@ -557,3 +557,28 @@ test.skip('fail on approve overspend', async () => {
     )
   ).rejects.toThrow();
 });
+
+
+test('set owner', async () => {
+  if (mockRpcEnabled) {
+    console.log('non-live test skipped');
+    return;
+  }
+
+  const connection = new Connection(url);
+  const owner = await newAccountWithTokens(connection);
+  const newOwner = await newAccountWithTokens(connection);
+
+  const account = await testToken.newAccount(owner);
+
+  await testToken.setOwner(owner, account, newOwner.publicKey);
+  await expect(
+    testToken.setOwner(owner, account, newOwner.publicKey)
+  ).rejects.toThrow();
+
+  await testToken.setOwner(newOwner, account, owner.publicKey);
+  await expect(
+    testToken.setOwner(newOwner, account, owner.publicKey)
+  ).rejects.toThrow();
+
+});