Skip to content

Commit c5c63fe

Browse files
mindrunnermindrunner
committed
add stake program functionality
1 parent 87e61cf commit c5c63fe

File tree

2 files changed

+338
-0
lines changed

2 files changed

+338
-0
lines changed

‎src/solana/_layouts/stake_instructions.py

+54
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
from enum import IntEnum
2+
3+
from construct import Switch # type: ignore
4+
from construct import Int32ul, Int64ul, Int64sl, Pass # type: ignore
5+
from construct import Struct as cStruct
6+
from spl.token._layouts import PUBLIC_KEY_LAYOUT
7+
8+
9+
class StakeInstructionType(IntEnum):
10+
"""Instruction types for staking program."""
11+
12+
INITIALIZE_STAKE_ACCOUNT = 0
13+
DELEGATE_STAKE_ACCOUNT = 2
14+
DEACTIVATE_STAKE_ACCOUNT = 5
15+
WITHDRAW_STAKE_ACCOUNT = 4
16+
17+
18+
_AUTHORIZED_LAYOUT = cStruct(
19+
"staker" / PUBLIC_KEY_LAYOUT,
20+
"withdrawer" / PUBLIC_KEY_LAYOUT,
21+
)
22+
23+
_LOCKUP_LAYOUT = cStruct(
24+
"unix_timestamp" / Int64sl,
25+
"epoch" / Int64ul,
26+
"custodian" / PUBLIC_KEY_LAYOUT,
27+
)
28+
29+
INITIALIZE_STAKE_ACCOUNT_LAYOUT = cStruct(
30+
"authorized" / _AUTHORIZED_LAYOUT,
31+
"lockup" / _LOCKUP_LAYOUT,
32+
)
33+
34+
WITHDRAW_STAKE_ACCOUNT_LAYOUT = cStruct(
35+
"lamports" / Int64ul,
36+
)
37+
38+
DELEGATE_STAKE_INSTRUCTIONS_LAYOUT = cStruct(
39+
"instruction_type" / Int32ul,
40+
)
41+
42+
STAKE_INSTRUCTIONS_LAYOUT = cStruct(
43+
"instruction_type" / Int32ul,
44+
"args"
45+
/ Switch(
46+
lambda this: this.instruction_type,
47+
{
48+
StakeInstructionType.INITIALIZE_STAKE_ACCOUNT: INITIALIZE_STAKE_ACCOUNT_LAYOUT,
49+
StakeInstructionType.DELEGATE_STAKE_ACCOUNT: cStruct(),
50+
StakeInstructionType.DEACTIVATE_STAKE_ACCOUNT: Pass,
51+
StakeInstructionType.WITHDRAW_STAKE_ACCOUNT: WITHDRAW_STAKE_ACCOUNT_LAYOUT,
52+
},
53+
),
54+
)

‎src/solana/stake_program.py

+284
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,284 @@
1+
from typing import NamedTuple, Union
2+
3+
from solders import sysvar
4+
from solders.instruction import AccountMeta, Instruction
5+
from solders.pubkey import Pubkey
6+
from solders.system_program import create_account, CreateAccountParams, create_account_with_seed, \
7+
CreateAccountWithSeedParams
8+
9+
from solana._layouts.stake_instructions import STAKE_INSTRUCTIONS_LAYOUT, StakeInstructionType
10+
from solana.transaction import Transaction
11+
12+
STAKE_CONFIG_PUBKEY: Pubkey = Pubkey.from_string("StakeConfig11111111111111111111111111111111")
13+
STAKE_PUBKEY: Pubkey = Pubkey.from_string("Stake11111111111111111111111111111111111111")
14+
15+
16+
class Authorized(NamedTuple):
17+
"""Staking account authority info."""
18+
staker: Pubkey
19+
""""""
20+
withdrawer: Pubkey
21+
""""""
22+
23+
24+
class Lockup(NamedTuple):
25+
"""Stake account lockup info."""
26+
unix_timestamp: int
27+
""""""
28+
epoch: int
29+
""""""
30+
custodian: Pubkey
31+
32+
33+
class InitializeStakeParams(NamedTuple):
34+
"""Initialize Staking params"""
35+
stake_pubkey: Pubkey
36+
""""""
37+
authorized: Authorized
38+
""""""
39+
lockup: Lockup
40+
""""""
41+
42+
43+
class CreateStakeAccountParams(NamedTuple):
44+
"""Create stake account transaction params."""
45+
46+
from_pubkey: Pubkey
47+
""""""
48+
stake_pubkey: Pubkey
49+
""""""
50+
authorized: Authorized
51+
""""""
52+
lockup: Lockup
53+
""""""
54+
lamports: int
55+
""""""
56+
57+
58+
class CreateStakeAccountWithSeedParams(NamedTuple):
59+
"""Create stake account with seed transaction params."""
60+
61+
from_pubkey: Pubkey
62+
""""""
63+
stake_pubkey: Pubkey
64+
""""""
65+
base_pubkey: Pubkey
66+
""""""
67+
seed: str
68+
""""""
69+
authorized: Authorized
70+
""""""
71+
lockup: Lockup
72+
""""""
73+
lamports: int
74+
""""""
75+
76+
77+
class DelegateStakeParams(NamedTuple):
78+
"""Create delegate stake account transaction params."""
79+
80+
stake_pubkey: Pubkey
81+
""""""
82+
authorized_pubkey: Pubkey
83+
""""""
84+
vote_pubkey: Pubkey
85+
""""""
86+
87+
88+
class CreateAccountAndDelegateStakeParams(NamedTuple):
89+
"""Create and delegate a stake account transaction params"""
90+
91+
from_pubkey: Pubkey
92+
""""""
93+
stake_pubkey: Pubkey
94+
""""""
95+
vote_pubkey: Pubkey
96+
""""""
97+
authorized: Authorized
98+
""""""
99+
lockup: Lockup
100+
""""""
101+
lamports: int
102+
""""""
103+
104+
105+
class CreateAccountWithSeedAndDelegateStakeParams(NamedTuple):
106+
"""Create and delegate stake account with seed transaction params."""
107+
108+
from_pubkey: Pubkey
109+
""""""
110+
stake_pubkey: Pubkey
111+
""""""
112+
base_pubkey: Pubkey
113+
""""""
114+
seed: str
115+
""""""
116+
vote_pubkey: Pubkey
117+
""""""
118+
authorized: Authorized
119+
""""""
120+
lockup: Lockup
121+
""""""
122+
lamports: int
123+
""""""
124+
125+
126+
class WithdrawStakeParams(NamedTuple):
127+
"""Withdraw stake account params"""
128+
129+
stake_pubkey: Pubkey
130+
""""""
131+
withdrawer_pubkey: Pubkey
132+
""""""
133+
to_pubkey: Pubkey
134+
""""""
135+
lamports: int
136+
""""""
137+
custodian_pubkey: Pubkey
138+
""""""
139+
140+
141+
def withdraw_stake(params: WithdrawStakeParams) -> Transaction:
142+
data = STAKE_INSTRUCTIONS_LAYOUT.build(
143+
{
144+
"instruction_type:": StakeInstructionType.WITHDRAW_STAKE_ACCOUNT,
145+
"args": {
146+
"lamports": params.lamports
147+
}
148+
}
149+
)
150+
151+
withdraw_instruction = Instruction(
152+
accounts=[
153+
AccountMeta(pubkey=params.stake_pubkey, is_signer=True, is_writable=True),
154+
AccountMeta(pubkey=params.to_pubkey, is_signer=False, is_writable=True),
155+
AccountMeta(pubkey=params.to_pubkey, is_signer=False, is_writable=True),
156+
AccountMeta(pubkey=params.to_pubkey, is_signer=False, is_writable=True),
157+
AccountMeta(pubkey=params.to_pubkey, is_signer=False, is_writable=True),
158+
],
159+
program_id=Pubkey.default(),
160+
data=data,
161+
)
162+
163+
return Transaction(fee_payer=params.stake_pubkey).add(withdraw_instruction)
164+
165+
166+
def create_account_and_delegate_stake(
167+
params: Union[CreateAccountAndDelegateStakeParams, CreateAccountWithSeedAndDelegateStakeParams]) -> Transaction:
168+
"""Generate a transaction to crate and delegate a stake account"""
169+
170+
initialize_stake_instruction = initialize_stake(
171+
InitializeStakeParams(
172+
stake_pubkey=params.stake_pubkey,
173+
authorized=params.authorized,
174+
lockup=params.lockup,
175+
)
176+
)
177+
178+
create_account_instruction = _create_stake_account_instruction(params=params)
179+
180+
delegate_stake_instruction = delegate_stake(
181+
DelegateStakeParams(
182+
stake_pubkey=params.stake_pubkey,
183+
authorized_pubkey=params.authorized.staker,
184+
vote_pubkey=params.vote_pubkey,
185+
)
186+
)
187+
188+
return Transaction(fee_payer=params.from_pubkey).add(
189+
create_account_instruction,
190+
initialize_stake_instruction,
191+
delegate_stake_instruction)
192+
193+
194+
def delegate_stake(params: DelegateStakeParams) -> Instruction:
195+
"""Generate an instruction to delete a Stake account"""
196+
197+
data = STAKE_INSTRUCTIONS_LAYOUT.build(
198+
{
199+
"instruction_type": StakeInstructionType.DELEGATE_STAKE_ACCOUNT,
200+
"args": {}
201+
}
202+
)
203+
return Instruction(
204+
accounts=[
205+
AccountMeta(pubkey=params.stake_pubkey, is_signer=False, is_writable=True),
206+
AccountMeta(pubkey=params.vote_pubkey, is_signer=False, is_writable=False),
207+
AccountMeta(pubkey=sysvar.CLOCK, is_signer=False, is_writable=False),
208+
AccountMeta(pubkey=sysvar.STAKE_HISTORY, is_signer=False, is_writable=False),
209+
AccountMeta(pubkey=STAKE_CONFIG_PUBKEY, is_signer=False, is_writable=False),
210+
AccountMeta(pubkey=params.authorized_pubkey, is_signer=True, is_writable=False),
211+
],
212+
program_id=STAKE_PUBKEY,
213+
data=data,
214+
)
215+
216+
217+
def initialize_stake(params: InitializeStakeParams) -> Instruction:
218+
data = STAKE_INSTRUCTIONS_LAYOUT.build(
219+
{
220+
"instruction_type": StakeInstructionType.INITIALIZE_STAKE_ACCOUNT,
221+
"args": {
222+
"authorized": {
223+
"staker": params.authorized.staker.__bytes__(),
224+
"withdrawer": params.authorized.withdrawer.__bytes__()
225+
},
226+
"lockup": {
227+
"unix_timestamp": params.lockup.unix_timestamp,
228+
"epoch": params.lockup.epoch,
229+
"custodian": params.lockup.custodian.__bytes__()
230+
},
231+
}
232+
}
233+
)
234+
235+
return Instruction(
236+
accounts=[
237+
AccountMeta(pubkey=params.stake_pubkey, is_signer=False, is_writable=True),
238+
AccountMeta(pubkey=sysvar.RENT, is_signer=False, is_writable=False),
239+
],
240+
program_id=STAKE_PUBKEY,
241+
data=data,
242+
)
243+
244+
245+
def _create_stake_account_instruction(params: Union[
246+
CreateStakeAccountParams, CreateStakeAccountWithSeedParams, CreateAccountAndDelegateStakeParams, CreateAccountWithSeedAndDelegateStakeParams]) -> Instruction:
247+
if isinstance(params, CreateStakeAccountParams) or isinstance(params, CreateAccountAndDelegateStakeParams):
248+
return create_account(
249+
CreateAccountParams(
250+
from_pubkey=params.from_pubkey,
251+
to_pubkey=params.stake_pubkey,
252+
lamports=params.lamports,
253+
space=200,
254+
owner=STAKE_PUBKEY
255+
)
256+
)
257+
return create_account_with_seed(
258+
CreateAccountWithSeedParams(
259+
from_pubkey=params.from_pubkey,
260+
to_pubkey=params.stake_pubkey,
261+
base=params.base_pubkey,
262+
seed=params.seed,
263+
lamports=params.lamports,
264+
space=200,
265+
owner=STAKE_PUBKEY,
266+
267+
)
268+
)
269+
270+
271+
def create_stake_account(params: Union[CreateStakeAccountParams, CreateStakeAccountWithSeedParams]) -> Transaction:
272+
"""Generate a Transaction that creates a new Staking Account"""
273+
274+
initialize_stake_instruction = initialize_stake(
275+
InitializeStakeParams(
276+
stake_pubkey=params.stake_pubkey,
277+
authorized=params.authorized,
278+
lockup=params.lockup,
279+
)
280+
)
281+
282+
create_account_instruction = _create_stake_account_instruction(params=params)
283+
284+
return Transaction(fee_payer=params.from_pubkey).add(create_account_instruction, initialize_stake_instruction)

0 commit comments

Comments
 (0)