Skip to content

Commit b96ef75

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

File tree

3 files changed

+348
-1
lines changed

3 files changed

+348
-1
lines changed

pyproject.toml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -86,7 +86,7 @@ dummy-variable-rgx = "^(_+|(_+[a-zA-Z0-9_]*[a-zA-Z0-9]+?))$"
8686
target-version = "py310"
8787

8888
[tool.ruff.pydocstyle]
89-
convention = "google"
89+
convention = "google"
9090

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

src/solana/stake_program.py

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

0 commit comments

Comments
 (0)