Skip to content
This repository was archived by the owner on Mar 11, 2025. It is now read-only.

Commit 95f6c4f

Browse files
joncinquejoncinque
authored
token-2022: Add harvest from accounts to mint (#2820)
* token-2022: Add harvest from accounts to mint * Address feedback
1 parent 46d27c7 commit 95f6c4f

File tree

3 files changed

+302
-8
lines changed

3 files changed

+302
-8
lines changed

‎token/program-2022/src/extension/transfer_fee/processor.rs

Lines changed: 57 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -4,13 +4,13 @@ use {
44
error::TokenError,
55
extension::{
66
transfer_fee::{
7-
instruction::TransferFeeInstruction, TransferFee, TransferFeeConfig,
8-
MAX_FEE_BASIS_POINTS,
7+
instruction::TransferFeeInstruction, TransferFee, TransferFeeAmount,
8+
TransferFeeConfig, MAX_FEE_BASIS_POINTS,
99
},
1010
StateWithExtensionsMut,
1111
},
1212
processor::Processor,
13-
state::Mint,
13+
state::{Account, Mint},
1414
},
1515
solana_program::{
1616
account_info::{next_account_info, AccountInfo},
@@ -108,6 +108,53 @@ fn process_set_transfer_fee(
108108
Ok(())
109109
}
110110

111+
fn harvest_from_account<'a, 'b>(
112+
mint_key: &'b Pubkey,
113+
mint_extension: &'b mut TransferFeeConfig,
114+
token_account_info: &'b AccountInfo<'a>,
115+
) -> Result<(), TokenError> {
116+
let mut token_account_data = token_account_info.data.borrow_mut();
117+
let mut token_account = StateWithExtensionsMut::<Account>::unpack(&mut token_account_data)
118+
.map_err(|_| TokenError::InvalidState)?;
119+
if token_account.base.mint != *mint_key {
120+
return Err(TokenError::MintMismatch);
121+
}
122+
check_program_account(token_account_info.owner).map_err(|_| TokenError::InvalidState)?;
123+
let token_account_extension = token_account
124+
.get_extension_mut::<TransferFeeAmount>()
125+
.map_err(|_| TokenError::InvalidState)?;
126+
let account_withheld_amount = u64::from(token_account_extension.withheld_amount);
127+
let mint_withheld_amount = u64::from(mint_extension.withheld_amount);
128+
mint_extension.withheld_amount = mint_withheld_amount
129+
.checked_add(account_withheld_amount)
130+
.ok_or(TokenError::Overflow)?
131+
.into();
132+
token_account_extension.withheld_amount = 0.into();
133+
Ok(())
134+
}
135+
136+
fn process_harvest_withheld_tokens_to_mint(accounts: &[AccountInfo]) -> ProgramResult {
137+
let account_info_iter = &mut accounts.iter();
138+
let mint_account_info = next_account_info(account_info_iter)?;
139+
let token_account_infos = account_info_iter.as_slice();
140+
141+
let mut mint_data = mint_account_info.data.borrow_mut();
142+
let mut mint = StateWithExtensionsMut::<Mint>::unpack(&mut mint_data)?;
143+
let mint_extension = mint.get_extension_mut::<TransferFeeConfig>()?;
144+
145+
for token_account_info in token_account_infos {
146+
match harvest_from_account(mint_account_info.key, mint_extension, token_account_info) {
147+
// Shouldn't ever happen, but if it does, we don't want to propagate any half-done changes
148+
Err(TokenError::Overflow) => return Err(TokenError::Overflow.into()),
149+
Err(e) => {
150+
msg!("Error harvesting from {}: {}", token_account_info.key, e);
151+
}
152+
Ok(_) => {}
153+
}
154+
}
155+
Ok(())
156+
}
157+
111158
pub(crate) fn process_instruction(
112159
program_id: &Pubkey,
113160
accounts: &[AccountInfo],
@@ -133,7 +180,7 @@ pub(crate) fn process_instruction(
133180
decimals,
134181
fee,
135182
} => {
136-
msg!("Instruction: TransferCheckedWithFee");
183+
msg!("TransferFeeInstruction: TransferCheckedWithFee");
137184
Processor::process_transfer(program_id, accounts, amount, Some(decimals), Some(fee))
138185
}
139186
TransferFeeInstruction::WithdrawWithheldTokensFromMint => {
@@ -143,11 +190,15 @@ pub(crate) fn process_instruction(
143190
unimplemented!();
144191
}
145192
TransferFeeInstruction::HarvestWithheldTokensToMint => {
146-
unimplemented!();
193+
msg!("TransferFeeInstruction: HarvestWithheldTokensToMint");
194+
process_harvest_withheld_tokens_to_mint(accounts)
147195
}
148196
TransferFeeInstruction::SetTransferFee {
149197
transfer_fee_basis_points,
150198
maximum_fee,
151-
} => process_set_transfer_fee(program_id, accounts, transfer_fee_basis_points, maximum_fee),
199+
} => {
200+
msg!("TransferFeeInstruction: SetTransferFee");
201+
process_set_transfer_fee(program_id, accounts, transfer_fee_basis_points, maximum_fee)
202+
}
152203
}
153204
}

‎token/program-2022/tests/transfer_fee.rs

Lines changed: 229 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,10 @@ use {
1515
},
1616
instruction,
1717
},
18-
spl_token_client::token::{ExtensionInitializationParams, TokenError as TokenClientError},
18+
spl_token_client::{
19+
client::ProgramBanksClientProcessTransaction,
20+
token::{ExtensionInitializationParams, Token, TokenError as TokenClientError},
21+
},
1922
std::convert::TryInto,
2023
};
2124

@@ -302,7 +305,7 @@ async fn set_fee() {
302305
}
303306

304307
#[tokio::test]
305-
async fn fail_set_fee_unsupported_mint() {
308+
async fn fail_unsupported_mint() {
306309
let mut context = TestContext::new().await;
307310
context.init_token_with_mint(vec![]).await.unwrap();
308311
let TokenContext {
@@ -323,6 +326,17 @@ async fn fail_set_fee_unsupported_mint() {
323326
TransactionError::InstructionError(0, InstructionError::InvalidAccountData)
324327
)))
325328
);
329+
let error = token
330+
.harvest_withheld_tokens_to_mint(&[])
331+
.await
332+
.err()
333+
.unwrap();
334+
assert_eq!(
335+
error,
336+
TokenClientError::Client(Box::new(TransportError::TransactionError(
337+
TransactionError::InstructionError(0, InstructionError::InvalidAccountData)
338+
)))
339+
);
326340
}
327341

328342
#[tokio::test]
@@ -993,3 +1007,216 @@ async fn no_fees_from_self_transfer() {
9931007
let extension = alice_state.get_extension::<TransferFeeAmount>().unwrap();
9941008
assert_eq!(extension.withheld_amount, 0.into());
9951009
}
1010+
1011+
async fn create_and_transfer_to_account(
1012+
token: &Token<ProgramBanksClientProcessTransaction, Keypair>,
1013+
source: &Pubkey,
1014+
authority: &Keypair,
1015+
owner: &Pubkey,
1016+
amount: u64,
1017+
decimals: u8,
1018+
) -> Pubkey {
1019+
let account = token
1020+
.create_auxiliary_token_account(&Keypair::new(), owner)
1021+
.await
1022+
.unwrap();
1023+
token
1024+
.transfer_checked(source, &account, authority, amount, decimals)
1025+
.await
1026+
.unwrap();
1027+
account
1028+
}
1029+
1030+
#[tokio::test]
1031+
async fn harvest_withheld_tokens_to_mint() {
1032+
let TransferFeeConfigWithKeypairs {
1033+
transfer_fee_config_authority,
1034+
withdraw_withheld_authority,
1035+
transfer_fee_config,
1036+
..
1037+
} = test_transfer_fee_config_with_keypairs();
1038+
let mut context = TestContext::new().await;
1039+
let transfer_fee_basis_points = u16::from(
1040+
transfer_fee_config
1041+
.newer_transfer_fee
1042+
.transfer_fee_basis_points,
1043+
);
1044+
let maximum_fee = u64::from(transfer_fee_config.newer_transfer_fee.maximum_fee);
1045+
context
1046+
.init_token_with_mint(vec![ExtensionInitializationParams::TransferFeeConfig {
1047+
transfer_fee_config_authority: transfer_fee_config_authority.pubkey().into(),
1048+
withdraw_withheld_authority: withdraw_withheld_authority.pubkey().into(),
1049+
transfer_fee_basis_points,
1050+
maximum_fee,
1051+
}])
1052+
.await
1053+
.unwrap();
1054+
let TokenContext {
1055+
decimals,
1056+
mint_authority,
1057+
token,
1058+
alice,
1059+
..
1060+
} = context.token_context.as_ref().unwrap();
1061+
1062+
let alice_account = Keypair::new();
1063+
let alice_account = token
1064+
.create_auxiliary_token_account(&alice_account, &alice.pubkey())
1065+
.await
1066+
.unwrap();
1067+
1068+
// mint a lot of tokens
1069+
let amount = maximum_fee;
1070+
let alice_amount = amount * 100;
1071+
token
1072+
.mint_to(&alice_account, mint_authority, alice_amount)
1073+
.await
1074+
.unwrap();
1075+
1076+
// harvest from zero accounts
1077+
token.harvest_withheld_tokens_to_mint(&[]).await.unwrap();
1078+
let state = token.get_mint_info().await.unwrap();
1079+
let extension = state.get_extension::<TransferFeeConfig>().unwrap();
1080+
assert_eq!(extension.withheld_amount, 0.into());
1081+
1082+
// harvest from one account
1083+
let accumulated_fees = transfer_fee_config.calculate_epoch_fee(0, amount).unwrap();
1084+
let account = create_and_transfer_to_account(
1085+
token,
1086+
&alice_account,
1087+
alice,
1088+
&alice.pubkey(),
1089+
amount,
1090+
*decimals,
1091+
)
1092+
.await;
1093+
token
1094+
.harvest_withheld_tokens_to_mint(&[&account])
1095+
.await
1096+
.unwrap();
1097+
let state = token.get_account_info(&account).await.unwrap();
1098+
let extension = state.get_extension::<TransferFeeAmount>().unwrap();
1099+
assert_eq!(extension.withheld_amount, 0.into());
1100+
let state = token.get_mint_info().await.unwrap();
1101+
let extension = state.get_extension::<TransferFeeConfig>().unwrap();
1102+
assert_eq!(extension.withheld_amount, accumulated_fees.into());
1103+
1104+
// harvest again from the same account
1105+
token
1106+
.harvest_withheld_tokens_to_mint(&[&account])
1107+
.await
1108+
.unwrap();
1109+
let state = token.get_account_info(&account).await.unwrap();
1110+
let extension = state.get_extension::<TransferFeeAmount>().unwrap();
1111+
assert_eq!(extension.withheld_amount, 0.into());
1112+
let state = token.get_mint_info().await.unwrap();
1113+
let extension = state.get_extension::<TransferFeeConfig>().unwrap();
1114+
assert_eq!(extension.withheld_amount, accumulated_fees.into());
1115+
1116+
// no fail harvesting from account belonging to different mint, but nothing
1117+
// happens
1118+
let account = create_and_transfer_to_account(
1119+
token,
1120+
&alice_account,
1121+
alice,
1122+
&alice.pubkey(),
1123+
amount,
1124+
*decimals,
1125+
)
1126+
.await;
1127+
context
1128+
.init_token_with_mint(vec![ExtensionInitializationParams::TransferFeeConfig {
1129+
transfer_fee_config_authority: transfer_fee_config_authority.pubkey().into(),
1130+
withdraw_withheld_authority: withdraw_withheld_authority.pubkey().into(),
1131+
transfer_fee_basis_points,
1132+
maximum_fee,
1133+
}])
1134+
.await
1135+
.unwrap();
1136+
let TokenContext { token, .. } = context.token_context.unwrap();
1137+
token
1138+
.harvest_withheld_tokens_to_mint(&[&account])
1139+
.await
1140+
.unwrap();
1141+
let state = token.get_mint_info().await.unwrap();
1142+
let extension = state.get_extension::<TransferFeeConfig>().unwrap();
1143+
assert_eq!(extension.withheld_amount, 0.into());
1144+
}
1145+
1146+
#[tokio::test]
1147+
async fn max_harvest_withheld_tokens_to_mint() {
1148+
let TransferFeeConfigWithKeypairs {
1149+
transfer_fee_config_authority,
1150+
withdraw_withheld_authority,
1151+
transfer_fee_config,
1152+
..
1153+
} = test_transfer_fee_config_with_keypairs();
1154+
let mut context = TestContext::new().await;
1155+
let transfer_fee_basis_points = u16::from(
1156+
transfer_fee_config
1157+
.newer_transfer_fee
1158+
.transfer_fee_basis_points,
1159+
);
1160+
let maximum_fee = u64::from(transfer_fee_config.newer_transfer_fee.maximum_fee);
1161+
context
1162+
.init_token_with_mint(vec![ExtensionInitializationParams::TransferFeeConfig {
1163+
transfer_fee_config_authority: transfer_fee_config_authority.pubkey().into(),
1164+
withdraw_withheld_authority: withdraw_withheld_authority.pubkey().into(),
1165+
transfer_fee_basis_points,
1166+
maximum_fee,
1167+
}])
1168+
.await
1169+
.unwrap();
1170+
let TokenContext {
1171+
decimals,
1172+
mint_authority,
1173+
token,
1174+
alice,
1175+
..
1176+
} = context.token_context.unwrap();
1177+
1178+
let alice_account = Keypair::new();
1179+
let alice_account = token
1180+
.create_auxiliary_token_account(&alice_account, &alice.pubkey())
1181+
.await
1182+
.unwrap();
1183+
1184+
// mint a lot of tokens
1185+
let amount = maximum_fee;
1186+
let alice_amount = amount * 100;
1187+
token
1188+
.mint_to(&alice_account, &mint_authority, alice_amount)
1189+
.await
1190+
.unwrap();
1191+
// harvest from max accounts, which is around 35, AKA 34 accounts + 1 mint
1192+
// see https://docs.solana.com/proposals/transactions-v2#problem
1193+
let mut accounts = vec![];
1194+
let max_accounts = 34;
1195+
for _ in 0..max_accounts {
1196+
let account = create_and_transfer_to_account(
1197+
&token,
1198+
&alice_account,
1199+
&alice,
1200+
&alice.pubkey(),
1201+
amount,
1202+
decimals,
1203+
)
1204+
.await;
1205+
accounts.push(account);
1206+
}
1207+
let accounts: Vec<_> = accounts.iter().collect();
1208+
let accumulated_fees =
1209+
max_accounts * transfer_fee_config.calculate_epoch_fee(0, amount).unwrap();
1210+
token
1211+
.harvest_withheld_tokens_to_mint(&accounts)
1212+
.await
1213+
.unwrap();
1214+
for account in accounts {
1215+
let state = token.get_account_info(account).await.unwrap();
1216+
let extension = state.get_extension::<TransferFeeAmount>().unwrap();
1217+
assert_eq!(extension.withheld_amount, 0.into());
1218+
}
1219+
let state = token.get_mint_info().await.unwrap();
1220+
let extension = state.get_extension::<TransferFeeConfig>().unwrap();
1221+
assert_eq!(extension.withheld_amount, accumulated_fees.into());
1222+
}

‎token/rust/src/token.rs

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -522,4 +522,20 @@ where
522522
)
523523
.await
524524
}
525+
526+
/// Harvest withheld tokens to mint
527+
pub async fn harvest_withheld_tokens_to_mint(
528+
&self,
529+
sources: &[&Pubkey],
530+
) -> TokenResult<T::Output> {
531+
self.process_ixs(
532+
&[transfer_fee::instruction::harvest_withheld_tokens_to_mint(
533+
&self.program_id,
534+
&self.pubkey,
535+
sources,
536+
)?],
537+
&[&self.payer],
538+
)
539+
.await
540+
}
525541
}

0 commit comments

Comments
 (0)