1use alloy::{
2 network::Ethereum,
3 primitives::{Address, U256},
4 providers::{PendingTransactionBuilder, Provider},
5};
6use anyhow::{bail, Context as _, Result};
7use hotshot_contract_adapter::{
8 evm::DecodeRevert as _,
9 reward::RewardClaimInput,
10 sol_types::{
11 EspTokenV2, LightClientV3,
12 RewardClaim::{self, RewardClaimErrors},
13 StakeTableV2::{self, StakeTableV2Errors},
14 },
15};
16use url::Url;
17
18pub async fn claim_withdrawal(
19 provider: impl Provider,
20 stake_table: Address,
21 validator_address: Address,
22) -> Result<PendingTransactionBuilder<Ethereum>> {
23 let st = StakeTableV2::new(stake_table, provider);
24 st.claimWithdrawal(validator_address)
25 .send()
26 .await
27 .maybe_decode_revert::<StakeTableV2Errors>()
28}
29
30pub async fn claim_validator_exit(
31 provider: impl Provider,
32 stake_table: Address,
33 validator_address: Address,
34) -> Result<PendingTransactionBuilder<Ethereum>> {
35 let st = StakeTableV2::new(stake_table, provider);
36 st.claimValidatorExit(validator_address)
37 .send()
38 .await
39 .maybe_decode_revert::<StakeTableV2Errors>()
40}
41
42struct RewardClaimData {
43 reward_claim_address: Address,
44 claim_input: RewardClaimInput,
45}
46
47async fn try_fetch_reward_claim_data(
48 provider: impl Provider + Clone,
49 stake_table_address: Address,
50 espresso_url: &Url,
51 claimer_address: Address,
52) -> Result<Option<RewardClaimData>> {
53 let stake_table = StakeTableV2::new(stake_table_address, &provider);
54 let token_address = stake_table
55 .token()
56 .call()
57 .await
58 .context("Failed to get token address from stake table")?;
59
60 let esp_token = EspTokenV2::new(token_address, &provider);
61 let reward_claim_address = esp_token
62 .rewardClaim()
63 .call()
64 .await
65 .context("Failed to get reward claim address from token contract")?;
66
67 if reward_claim_address == Address::ZERO {
68 bail!("Reward claim contract not set on ESP token");
69 }
70
71 let light_client_address = stake_table
72 .lightClient()
73 .call()
74 .await
75 .context("Failed to get light client address from stake table")?;
76 let light_client = LightClientV3::new(light_client_address, &provider);
77
78 let auth_root = light_client
79 .authRoot()
80 .call()
81 .await
82 .context("Failed to get auth root from light client")?;
83
84 if auth_root == U256::ZERO {
85 return Ok(None);
86 }
87
88 let finalized_state = light_client
89 .finalizedState()
90 .call()
91 .await
92 .context("Failed to get finalized state from light client")?;
93 let block_height = finalized_state.blockHeight;
94
95 let reward_claim_url = format!(
96 "{}reward-state-v2/reward-claim-input/{}/{}",
97 espresso_url, block_height, claimer_address
98 );
99
100 let http_client = reqwest::Client::new();
101 let response = http_client
102 .get(&reward_claim_url)
103 .header("Accept", "application/json")
104 .send()
105 .await
106 .context("Failed to fetch reward claim input from Espresso API")?;
107
108 if response.status() == reqwest::StatusCode::NOT_FOUND {
109 return Ok(None);
110 }
111
112 let response = response
113 .error_for_status()
114 .context("Espresso API returned error status")?;
115
116 let claim_input: RewardClaimInput = response
117 .json()
118 .await
119 .context("Failed to parse reward claim input from API response")?;
120
121 Ok(Some(RewardClaimData {
122 reward_claim_address,
123 claim_input,
124 }))
125}
126
127pub async fn claim_reward(
128 provider: impl Provider + Clone,
129 stake_table_address: Address,
130 espresso_url: Url,
131 claimer_address: Address,
132) -> Result<PendingTransactionBuilder<Ethereum>> {
133 let data = try_fetch_reward_claim_data(
134 &provider,
135 stake_table_address,
136 &espresso_url,
137 claimer_address,
138 )
139 .await?
140 .context("No reward claim data found for address")?;
141
142 let reward_claim = RewardClaim::new(data.reward_claim_address, provider);
143 reward_claim
144 .claimRewards(
145 data.claim_input.lifetime_rewards,
146 data.claim_input.auth_data.into(),
147 )
148 .send()
149 .await
150 .maybe_decode_revert::<RewardClaimErrors>()
151}
152
153pub async fn unclaimed_rewards(
154 provider: impl Provider + Clone,
155 stake_table_address: Address,
156 espresso_url: Url,
157 claimer_address: Address,
158) -> Result<U256> {
159 let Some(data) = try_fetch_reward_claim_data(
160 &provider,
161 stake_table_address,
162 &espresso_url,
163 claimer_address,
164 )
165 .await?
166 else {
167 return Ok(U256::ZERO);
168 };
169
170 let reward_claim = RewardClaim::new(data.reward_claim_address, &provider);
171 let already_claimed = reward_claim.claimedRewards(claimer_address).call().await?;
172
173 let unclaimed = data
174 .claim_input
175 .lifetime_rewards
176 .checked_sub(already_claimed)
177 .unwrap_or(U256::ZERO);
178
179 Ok(unclaimed)
180}
181
182#[cfg(test)]
183mod test {
184 use alloy::primitives::{utils::parse_ether, U256};
185 use warp::Filter as _;
186
187 use super::*;
188 use crate::{deploy::TestSystem, receipt::ReceiptExt};
189
190 #[tokio::test]
191 async fn test_claim_withdrawal() -> Result<()> {
192 let system = TestSystem::deploy().await?;
193 let amount = parse_ether("1.23")?;
194 system.register_validator().await?;
195 system.delegate(amount).await?;
196 system.undelegate(amount).await?;
197 system.warp_to_unlock_time().await?;
198
199 let validator_address = system.deployer_address;
200 let receipt = claim_withdrawal(&system.provider, system.stake_table, validator_address)
201 .await?
202 .assert_success()
203 .await?;
204
205 let event = receipt
206 .decoded_log::<StakeTableV2::WithdrawalClaimed>()
207 .unwrap();
208 assert_eq!(event.amount, amount);
209
210 Ok(())
211 }
212
213 #[tokio::test]
214 async fn test_claim_validator_exit() -> Result<()> {
215 let system = TestSystem::deploy().await?;
216 let amount = parse_ether("1.23")?;
217 system.register_validator().await?;
218 system.delegate(amount).await?;
219 system.deregister_validator().await?;
220 system.warp_to_unlock_time().await?;
221
222 let validator_address = system.deployer_address;
223 let receipt = claim_validator_exit(&system.provider, system.stake_table, validator_address)
224 .await?
225 .assert_success()
226 .await?;
227
228 let event = receipt
229 .decoded_log::<StakeTableV2::ValidatorExitClaimed>()
230 .unwrap();
231 assert_eq!(event.amount, amount);
232
233 Ok(())
234 }
235
236 #[tokio::test]
237 async fn test_claim_reward() -> Result<()> {
238 let system = TestSystem::deploy().await?;
239 let reward_balance = U256::from(1000000);
240
241 let espresso_url = system.setup_reward_claim_mock(reward_balance).await?;
242
243 tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
244
245 let balance_before = system.balance(system.deployer_address).await?;
246
247 let receipt = claim_reward(
248 &system.provider,
249 system.stake_table,
250 espresso_url,
251 system.deployer_address,
252 )
253 .await?
254 .assert_success()
255 .await?;
256
257 let event = receipt
258 .decoded_log::<RewardClaim::RewardsClaimed>()
259 .unwrap();
260 assert_eq!(event.amount, reward_balance);
261
262 let balance_after = system.balance(system.deployer_address).await?;
263 assert_eq!(balance_after, balance_before + reward_balance);
264
265 Ok(())
266 }
267
268 #[tokio::test]
269 async fn test_unclaimed_rewards_not_found() -> Result<()> {
270 let system = TestSystem::deploy().await?;
271
272 let port = portpicker::pick_unused_port().expect("No ports available");
273
274 let route = warp::path!("reward-state-v2" / "reward-claim-input" / u64 / String)
275 .map(|_, _| warp::reply::with_status(warp::reply(), warp::http::StatusCode::NOT_FOUND));
276
277 tokio::spawn(warp::serve(route).run(([127, 0, 0, 1], port)));
278
279 tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
280
281 let espresso_url = format!("http://localhost:{}/", port).parse()?;
282
283 let unclaimed = unclaimed_rewards(
284 &system.provider,
285 system.stake_table,
286 espresso_url,
287 system.deployer_address,
288 )
289 .await?;
290
291 assert_eq!(unclaimed, U256::ZERO);
292
293 Ok(())
294 }
295}