fix: add custodian key support to stake instructions

web3.js/module.d.ts

@@ -772,6 +772,7 @@ declare module '@solana/web3.js' {
     authorizedPubkey: PublicKey;
     newAuthorizedPubkey: PublicKey;
     stakeAuthorizationType: StakeAuthorizationType;
+    custodianPubkey?: PublicKey;
   };
 
   export type AuthorizeWithSeedStakeParams = {
@@ -781,6 +782,7 @@ declare module '@solana/web3.js' {
     authorityOwner: PublicKey;
     newAuthorizedPubkey: PublicKey;
     stakeAuthorizationType: StakeAuthorizationType;
+    custodianPubkey?: PublicKey;
   };
 
   export type SplitStakeParams = {
@@ -795,6 +797,7 @@ declare module '@solana/web3.js' {
     authorizedPubkey: PublicKey;
     toPubkey: PublicKey;
     lamports: number;
+    custodianPubkey?: PublicKey;
   };
 
   export type DeactivateStakeParams = {

web3.js/module.flow.js

@@ -780,6 +780,7 @@ declare module '@solana/web3.js' {
     authorizedPubkey: PublicKey,
     newAuthorizedPubkey: PublicKey,
     stakeAuthorizationType: StakeAuthorizationType,
+    custodianPubkey?: PublicKey,
   |};
 
   declare export type AuthorizeWithSeedStakeParams = {|
@@ -789,6 +790,7 @@ declare module '@solana/web3.js' {
     authorityOwner: PublicKey,
     newAuthorizedPubkey: PublicKey,
     stakeAuthorizationType: StakeAuthorizationType,
+    custodianPubkey?: PublicKey,
   |};
 
   declare export type SplitStakeParams = {|
@@ -803,6 +805,7 @@ declare module '@solana/web3.js' {
     authorizedPubkey: PublicKey,
     toPubkey: PublicKey,
     lamports: number,
+    custodianPubkey?: PublicKey,
   |};
 
   declare export type DeactivateStakeParams = {|

web3.js/src/stake-program.js

@@ -116,12 +116,14 @@ export type DelegateStakeParams = {|
  * @property {PublicKey} authorizedPubkey
  * @property {PublicKey} newAuthorizedPubkey
  * @property {StakeAuthorizationType} stakeAuthorizationType
+ * @property {PublicKey} custodianPubkey
  */
 export type AuthorizeStakeParams = {|
   stakePubkey: PublicKey,
   authorizedPubkey: PublicKey,
   newAuthorizedPubkey: PublicKey,
   stakeAuthorizationType: StakeAuthorizationType,
+  custodianPubkey?: PublicKey,
 |};
 
 /**
@@ -133,6 +135,7 @@ export type AuthorizeStakeParams = {|
  * @property {PublicKey} authorityOwner
  * @property {PublicKey} newAuthorizedPubkey
  * @property {StakeAuthorizationType} stakeAuthorizationType
+ * @property {PublicKey} custodianPubkey
  */
 export type AuthorizeWithSeedStakeParams = {|
   stakePubkey: PublicKey,
@@ -141,6 +144,7 @@ export type AuthorizeWithSeedStakeParams = {|
   authorityOwner: PublicKey,
   newAuthorizedPubkey: PublicKey,
   stakeAuthorizationType: StakeAuthorizationType,
+  custodianPubkey?: PublicKey,
 |};
 
 /**
@@ -165,12 +169,14 @@ export type SplitStakeParams = {|
  * @property {PublicKey} authorizedPubkey
  * @property {PublicKey} toPubkey
  * @property {number} lamports
+ * @property {PublicKey} custodianPubkey
  */
 export type WithdrawStakeParams = {|
   stakePubkey: PublicKey,
   authorizedPubkey: PublicKey,
   toPubkey: PublicKey,
   lamports: number,
+  custodianPubkey?: PublicKey,
 |};
 
 /**
@@ -271,7 +277,7 @@ export class StakeInstruction {
       instruction.data,
     );
 
-    return {
+    const o: AuthorizeStakeParams = {
       stakePubkey: instruction.keys[0].pubkey,
       authorizedPubkey: instruction.keys[2].pubkey,
       newAuthorizedPubkey: new PublicKey(newAuthorized),
@@ -279,6 +285,10 @@ export class StakeInstruction {
         index: stakeAuthorizationType,
       },
     };
+    if (instruction.keys.length > 3) {
+      o.custodianPubkey = instruction.keys[3].pubkey;
+    }
+    return o;
   }
 
   /**
@@ -289,6 +299,7 @@ export class StakeInstruction {
   ): AuthorizeWithSeedStakeParams {
     this.checkProgramId(instruction.programId);
     this.checkKeyLength(instruction.keys, 2);
+
     const {
       newAuthorized,
       stakeAuthorizationType,
@@ -299,7 +310,7 @@ export class StakeInstruction {
       instruction.data,
     );
 
-    return {
+    const o: AuthorizeWithSeedStakeParams = {
       stakePubkey: instruction.keys[0].pubkey,
       authorityBase: instruction.keys[1].pubkey,
       authoritySeed: authoritySeed,
@@ -309,6 +320,10 @@ export class StakeInstruction {
         index: stakeAuthorizationType,
       },
     };
+    if (instruction.keys.length > 3) {
+      o.custodianPubkey = instruction.keys[3].pubkey;
+    }
+    return o;
   }
 
   /**
@@ -343,12 +358,16 @@ export class StakeInstruction {
       instruction.data,
     );
 
-    return {
+    const o: WithdrawStakeParams = {
       stakePubkey: instruction.keys[0].pubkey,
       toPubkey: instruction.keys[1].pubkey,
       authorizedPubkey: instruction.keys[4].pubkey,
       lamports,
     };
+    if (instruction.keys.length > 5) {
+      o.custodianPubkey = instruction.keys[5].pubkey;
+    }
+    return o;
   }
 
   /**
@@ -601,6 +620,7 @@ export class StakeProgram {
       authorizedPubkey,
       newAuthorizedPubkey,
       stakeAuthorizationType,
+      custodianPubkey,
     } = params;
 
     const type = STAKE_INSTRUCTION_LAYOUTS.Authorize;
@@ -609,12 +629,16 @@ export class StakeProgram {
       stakeAuthorizationType: stakeAuthorizationType.index,
     });
 
-    return new Transaction().add({
+    const keys = [
-      keys: [
       {pubkey: stakePubkey, isSigner: false, isWritable: true},
       {pubkey: SYSVAR_CLOCK_PUBKEY, isSigner: false, isWritable: true},
       {pubkey: authorizedPubkey, isSigner: true, isWritable: false},
-      ],
+    ];
+    if (custodianPubkey) {
+      keys.push({pubkey: custodianPubkey, isSigner: false, isWritable: false});
+    }
+    return new Transaction().add({
+      keys,
       programId: this.programId,
       data,
     });
@@ -632,6 +656,7 @@ export class StakeProgram {
       authorityOwner,
       newAuthorizedPubkey,
       stakeAuthorizationType,
+      custodianPubkey,
     } = params;
 
     const type = STAKE_INSTRUCTION_LAYOUTS.AuthorizeWithSeed;
@@ -642,12 +667,16 @@ export class StakeProgram {
       authorityOwner: authorityOwner.toBuffer(),
     });
 
-    return new Transaction().add({
+    const keys = [
-      keys: [
       {pubkey: stakePubkey, isSigner: false, isWritable: true},
       {pubkey: authorityBase, isSigner: true, isWritable: false},
       {pubkey: SYSVAR_CLOCK_PUBKEY, isSigner: false, isWritable: false},
-      ],
+    ];
+    if (custodianPubkey) {
+      keys.push({pubkey: custodianPubkey, isSigner: false, isWritable: false});
+    }
+    return new Transaction().add({
+      keys,
       programId: this.programId,
       data,
     });
@@ -687,12 +716,17 @@ export class StakeProgram {
    * Generate a Transaction that withdraws deactivated Stake tokens.
    */
   static withdraw(params: WithdrawStakeParams): Transaction {
-    const {stakePubkey, authorizedPubkey, toPubkey, lamports} = params;
+    const {
+      stakePubkey,
+      authorizedPubkey,
+      toPubkey,
+      lamports,
+      custodianPubkey,
+    } = params;
     const type = STAKE_INSTRUCTION_LAYOUTS.Withdraw;
     const data = encodeData(type, {lamports});
 
-    return new Transaction().add({
+    const keys = [
-      keys: [
       {pubkey: stakePubkey, isSigner: false, isWritable: true},
       {pubkey: toPubkey, isSigner: false, isWritable: true},
       {pubkey: SYSVAR_CLOCK_PUBKEY, isSigner: false, isWritable: false},
@@ -702,7 +736,12 @@ export class StakeProgram {
         isWritable: false,
       },
       {pubkey: authorizedPubkey, isSigner: true, isWritable: false},
-      ],
+    ];
+    if (custodianPubkey) {
+      keys.push({pubkey: custodianPubkey, isSigner: false, isWritable: false});
+    }
+    return new Transaction().add({
+      keys,
       programId: this.programId,
       data,
     });

web3.js/test/stake-program.test.js

@@ -129,6 +129,25 @@ test('authorize', () => {
   expect(params).toEqual(StakeInstruction.decodeAuthorize(stakeInstruction));
 });
 
+test('authorize with custodian', () => {
+  const stakePubkey = new Account().publicKey;
+  const authorizedPubkey = new Account().publicKey;
+  const newAuthorizedPubkey = new Account().publicKey;
+  const stakeAuthorizationType = StakeAuthorizationLayout.Withdrawer;
+  const custodianPubkey = new Account().publicKey;
+  const params = {
+    stakePubkey,
+    authorizedPubkey,
+    newAuthorizedPubkey,
+    stakeAuthorizationType,
+    custodianPubkey,
+  };
+  const transaction = StakeProgram.authorize(params);
+  expect(transaction.instructions).toHaveLength(1);
+  const [stakeInstruction] = transaction.instructions;
+  expect(params).toEqual(StakeInstruction.decodeAuthorize(stakeInstruction));
+});
+
 test('authorizeWithSeed', () => {
   const stakePubkey = new Account().publicKey;
   const authorityBase = new Account().publicKey;
@@ -152,6 +171,31 @@ test('authorizeWithSeed', () => {
   );
 });
 
+test('authorizeWithSeed with custodian', () => {
+  const stakePubkey = new Account().publicKey;
+  const authorityBase = new Account().publicKey;
+  const authoritySeed = 'test string';
+  const authorityOwner = new Account().publicKey;
+  const newAuthorizedPubkey = new Account().publicKey;
+  const stakeAuthorizationType = StakeAuthorizationLayout.Staker;
+  const custodianPubkey = new Account().publicKey;
+  const params = {
+    stakePubkey,
+    authorityBase,
+    authoritySeed,
+    authorityOwner,
+    newAuthorizedPubkey,
+    stakeAuthorizationType,
+    custodianPubkey,
+  };
+  const transaction = StakeProgram.authorizeWithSeed(params);
+  expect(transaction.instructions).toHaveLength(1);
+  const [stakeInstruction] = transaction.instructions;
+  expect(params).toEqual(
+    StakeInstruction.decodeAuthorizeWithSeed(stakeInstruction),
+  );
+});
+
 test('split', () => {
   const stakePubkey = new Account().publicKey;
   const authorizedPubkey = new Account().publicKey;
@@ -194,6 +238,24 @@ test('withdraw', () => {
   expect(params).toEqual(StakeInstruction.decodeWithdraw(stakeInstruction));
 });
 
+test('withdraw with custodian', () => {
+  const stakePubkey = new Account().publicKey;
+  const authorizedPubkey = new Account().publicKey;
+  const toPubkey = new Account().publicKey;
+  const custodianPubkey = new Account().publicKey;
+  const params = {
+    stakePubkey,
+    authorizedPubkey,
+    toPubkey,
+    lamports: 123,
+    custodianPubkey,
+  };
+  const transaction = StakeProgram.withdraw(params);
+  expect(transaction.instructions).toHaveLength(1);
+  const [stakeInstruction] = transaction.instructions;
+  expect(params).toEqual(StakeInstruction.decodeWithdraw(stakeInstruction));
+});
+
 test('deactivate', () => {
   const stakePubkey = new Account().publicKey;
   const authorizedPubkey = new Account().publicKey;