1use std::collections::HashMap;
2
3use num_bigint::BigUint;
4use tycho_client::feed::{synchronizer::ComponentWithState, BlockHeader};
5use tycho_common::{models::token::Token, Bytes};
6
7use crate::{
8 evm::protocol::lido::state::{LidoPoolType, LidoState, StakeLimitState, StakingStatus},
9 protocol::{
10 errors::InvalidSnapshotError,
11 models::{DecoderContext, TryFromWithBlock},
12 },
13};
14
15pub const ETH_ADDRESS: &str = "0x0000000000000000000000000000000000000000";
16
17impl TryFromWithBlock<ComponentWithState, BlockHeader> for LidoState {
18 type Error = InvalidSnapshotError;
19
20 async fn try_from_with_header(
22 snapshot: ComponentWithState,
23 _block: BlockHeader,
24 _account_balances: &HashMap<Bytes, HashMap<Bytes, Bytes>>,
25 _all_tokens: &HashMap<Bytes, Token>,
26 _decoder_context: &DecoderContext,
27 ) -> Result<Self, Self::Error> {
28 let id = snapshot.component.id.as_str();
29
30 let pool_type = match snapshot
31 .component
32 .static_attributes
33 .get("protocol_type_name")
34 .and_then(|bytes| std::str::from_utf8(bytes).ok())
35 .ok_or(InvalidSnapshotError::MissingAttribute(
36 "protocol_type_name is missing".to_owned(),
37 ))? {
38 "stETH" => LidoPoolType::StEth,
39 "wstETH" => LidoPoolType::WStEth,
40 _ => {
41 return Err(InvalidSnapshotError::ValueError(format!(
42 "Unknown protocol type name: {:?}",
43 snapshot.component.protocol_type_name
44 )))
45 }
46 };
47
48 let token_to_track_total_pooled_eth = snapshot
49 .component
50 .static_attributes
51 .get("token_to_track_total_pooled_eth")
52 .ok_or(InvalidSnapshotError::MissingAttribute(
53 "token_to_track_total_pooled_eth is missing".to_owned(),
54 ))?
55 .clone();
56
57 let tokens: [Bytes; 2] =
58 [snapshot.component.tokens[0].clone(), snapshot.component.tokens[1].clone()];
59
60 let total_shares = snapshot
61 .state
62 .attributes
63 .get("total_shares")
64 .ok_or(InvalidSnapshotError::MissingAttribute(
65 "Total shares field is missing".to_owned(),
66 ))?;
67
68 let total_pooled_eth = snapshot
69 .state
70 .balances
71 .get(&token_to_track_total_pooled_eth)
72 .ok_or(InvalidSnapshotError::MissingAttribute(
73 "Total shares field is missing".to_owned(),
74 ))?;
75
76 let (staking_status_parsed, staking_limit) = if pool_type == LidoPoolType::StEth {
77 let staking_status = snapshot
78 .state
79 .attributes
80 .get("staking_status")
81 .ok_or(InvalidSnapshotError::MissingAttribute(
82 "Staking_status field is missing".to_owned(),
83 ))?;
84
85 let staking_status_parsed =
86 if let Ok(status_as_str) = std::str::from_utf8(staking_status) {
87 match status_as_str {
88 "Limited" => StakingStatus::Limited,
89 "Paused" => StakingStatus::Paused,
90 "Unlimited" => StakingStatus::Unlimited,
91 _ => {
92 return Err(InvalidSnapshotError::ValueError(
93 "status_as_str parsed to invalid status".to_owned(),
94 ))
95 }
96 }
97 } else {
98 return Err(InvalidSnapshotError::ValueError(
99 "status_as_str cannot be parsed".to_owned(),
100 ))
101 };
102
103 let staking_limit = snapshot
104 .state
105 .attributes
106 .get("staking_limit")
107 .ok_or(InvalidSnapshotError::MissingAttribute(
108 "Staking_limit field is missing".to_owned(),
109 ))?;
110 (staking_status_parsed, staking_limit)
111 } else {
112 (StakingStatus::Limited, &Bytes::from(vec![0; 32]))
113 };
114
115 let total_wrapped_st_eth = if pool_type == LidoPoolType::StEth {
116 None
117 } else {
118 Some(BigUint::from_bytes_be(
119 snapshot
120 .state
121 .attributes
122 .get("total_wstETH")
123 .ok_or(InvalidSnapshotError::MissingAttribute(
124 "Total pooled eth field is missing".to_owned(),
125 ))?,
126 ))
127 };
128
129 Ok(Self {
130 pool_type,
131 total_shares: BigUint::from_bytes_be(total_shares),
132 total_pooled_eth: BigUint::from_bytes_be(total_pooled_eth),
133 total_wrapped_st_eth,
134 id: id.into(),
135 native_address: ETH_ADDRESS.into(),
136 stake_limits_state: StakeLimitState {
137 staking_status: staking_status_parsed,
138 staking_limit: BigUint::from_bytes_be(staking_limit),
139 },
140 tokens,
141 token_to_track_total_pooled_eth,
142 })
143 }
144}
145
146#[cfg(test)]
147mod tests {
148 use std::{collections::HashMap, str::FromStr};
149
150 use chrono::NaiveDateTime;
151 use num_bigint::BigUint;
152 use num_traits::Zero;
153 use rstest::rstest;
154 use tycho_client::feed::{synchronizer::ComponentWithState, BlockHeader};
155 use tycho_common::{
156 dto::{Chain, ChangeType, ProtocolComponent, ResponseProtocolState},
157 Bytes,
158 };
159
160 use crate::{
161 evm::protocol::lido::{
162 decoder::ETH_ADDRESS,
163 state::{LidoPoolType, LidoState, StakeLimitState},
164 },
165 protocol::{
166 errors::InvalidSnapshotError,
167 models::{DecoderContext, TryFromWithBlock},
168 },
169 };
170
171 const ST_ETH_ADDRESS_PROXY: &str = "0xae7ab96520de3a18e5e111b5eaab095312d7fe84";
172 const WST_ETH_ADDRESS: &str = "0x7f39c581f595b53c5cb19bd0b3f8da6c935e2ca0";
173
174 fn header() -> BlockHeader {
175 BlockHeader {
176 number: 1,
177 hash: Bytes::from(vec![0; 32]),
178 parent_hash: Bytes::from(vec![0; 32]),
179 revert: false,
180 timestamp: 1,
181 }
182 }
183
184 #[tokio::test]
185 async fn test_lido_steth_try_from() {
186 let mut static_attr = HashMap::new();
187 static_attr.insert(
188 "token_to_track_total_pooled_eth".to_string(),
189 "0x0000000000000000000000000000000000000000"
190 .as_bytes()
191 .to_vec()
192 .into(),
193 );
194 static_attr.insert(
195 "token_to_track_total_pooled_eth".to_string(),
196 Bytes::from_str(ETH_ADDRESS).unwrap(),
197 );
198 static_attr.insert("protocol_type_name".to_string(), "stETH".as_bytes().to_vec().into());
199
200 let pc = ProtocolComponent {
201 id: ST_ETH_ADDRESS_PROXY.to_string(),
202 protocol_system: "protocol_system".to_owned(),
203 protocol_type_name: "protocol_type_name".to_owned(),
204 chain: Chain::Ethereum,
205 tokens: vec![
206 Bytes::from("0x0000000000000000000000000000000000000000"),
207 Bytes::from("0xae7ab96520de3a18e5e111b5eaab095312d7fe84"),
208 ],
209 contract_ids: vec![],
210 static_attributes: static_attr,
211 change: ChangeType::Creation,
212 creation_tx: Bytes::from(vec![0; 32]),
213 created_at: NaiveDateTime::default(),
214 };
215
216 let snapshot = ComponentWithState {
217 state: ResponseProtocolState {
218 component_id: ST_ETH_ADDRESS_PROXY.to_owned(),
219 attributes: HashMap::from([
220 ("total_shares".to_string(), Bytes::from(vec![0; 32])),
221 ("staking_status".to_string(), "Limited".as_bytes().to_vec().into()),
222 ("staking_limit".to_string(), Bytes::from(vec![0; 32])),
223 ]),
224 balances: HashMap::from([(Bytes::from(ETH_ADDRESS), Bytes::from(vec![0; 32]))]),
225 },
226 component: pc,
227 component_tvl: None,
228 entrypoints: Vec::new(),
229 };
230
231 let decoder_context = DecoderContext::new();
232
233 let result = LidoState::try_from_with_header(
234 snapshot,
235 header(),
236 &HashMap::new(),
237 &HashMap::new(),
238 &decoder_context,
239 )
240 .await;
241
242 assert!(result.is_ok());
243 assert_eq!(
244 result.unwrap(),
245 LidoState {
246 pool_type: LidoPoolType::StEth,
247 total_shares: BigUint::zero(),
248 total_pooled_eth: BigUint::zero(),
249 total_wrapped_st_eth: None,
250 id: ST_ETH_ADDRESS_PROXY.into(),
251 native_address: ETH_ADDRESS.into(),
252 stake_limits_state: StakeLimitState {
253 staking_status: crate::evm::protocol::lido::state::StakingStatus::Limited,
254 staking_limit: BigUint::zero(),
255 },
256 tokens: [
257 Bytes::from("0x0000000000000000000000000000000000000000"),
258 Bytes::from("0xae7ab96520de3a18e5e111b5eaab095312d7fe84"),
259 ],
260 token_to_track_total_pooled_eth: Bytes::from(ETH_ADDRESS)
261 }
262 );
263 }
264
265 #[tokio::test]
266 #[rstest]
267 #[case::missing_total_shares("total_shares")]
268 #[case::missing_staking_status("staking_status")]
269 #[case::missing_staking_limit("staking_limit")]
270 async fn test_lido_try_from_missing_attribute(#[case] missing_attribute: &str) {
271 let mut static_attr = HashMap::new();
272 static_attr.insert(
273 "token_to_track_total_pooled_eth".to_string(),
274 Bytes::from_str(ETH_ADDRESS).unwrap(),
275 );
276
277 let pc = ProtocolComponent {
278 id: ST_ETH_ADDRESS_PROXY.to_string(),
279 protocol_system: "protocol_system".to_owned(),
280 protocol_type_name: "protocol_type_name".to_owned(),
281 chain: Chain::Ethereum,
282 tokens: vec![
283 Bytes::from("0x0000000000000000000000000000000000000000"),
284 Bytes::from("0xae7ab96520de3a18e5e111b5eaab095312d7fe84"),
285 ],
286 contract_ids: vec![],
287 static_attributes: static_attr,
288 change: ChangeType::Creation,
289 creation_tx: Bytes::from(vec![0; 32]),
290 created_at: NaiveDateTime::default(),
291 };
292
293 let mut snapshot = ComponentWithState {
294 state: ResponseProtocolState {
295 component_id: ST_ETH_ADDRESS_PROXY.to_owned(),
296 attributes: HashMap::from([
297 ("total_shares".to_string(), Bytes::from(vec![0; 32])),
298 ("staking_status".to_string(), "Limited".as_bytes().to_vec().into()),
299 ("staking_limit".to_string(), Bytes::from(vec![0; 32])),
300 ]),
301 balances: HashMap::from([(
302 Bytes::from_str(ETH_ADDRESS).unwrap(),
303 Bytes::from(vec![0; 32]),
304 )]),
305 },
306 component: pc,
307 component_tvl: None,
308 entrypoints: Vec::new(),
309 };
310 snapshot
311 .state
312 .attributes
313 .remove(missing_attribute);
314
315 let decoder_context = DecoderContext::new();
316
317 let result = LidoState::try_from_with_header(
318 snapshot,
319 header(),
320 &HashMap::new(),
321 &HashMap::new(),
322 &decoder_context,
323 )
324 .await;
325
326 assert!(result.is_err());
327 assert!(matches!(result.unwrap_err(), InvalidSnapshotError::MissingAttribute(_)));
328 }
329
330 #[tokio::test]
331 async fn test_lido_wst_eth_try_from() {
332 let mut static_attr = HashMap::new();
333 static_attr.insert(
334 "token_to_track_total_pooled_eth".to_string(),
335 "0xae7ab96520de3a18e5e111b5eaab095312d7fe84"
336 .as_bytes()
337 .to_vec()
338 .into(),
339 );
340 static_attr.insert(
341 "token_to_track_total_pooled_eth".to_string(),
342 Bytes::from_str(ST_ETH_ADDRESS_PROXY).unwrap(),
343 );
344 static_attr.insert("protocol_type_name".to_string(), "wstETH".as_bytes().to_vec().into());
345
346 let pc = ProtocolComponent {
347 id: WST_ETH_ADDRESS.to_string(),
348 protocol_system: "protocol_system".to_owned(),
349 protocol_type_name: "protocol_type_name".to_owned(),
350 chain: Chain::Ethereum,
351 tokens: vec![
352 Bytes::from("0x7f39c581f595b53c5cb19bd0b3f8da6c935e2ca0"),
353 Bytes::from("0xae7ab96520de3a18e5e111b5eaab095312d7fe84"),
354 ],
355 contract_ids: vec![],
356 static_attributes: static_attr,
357 change: ChangeType::Creation,
358 creation_tx: Bytes::from(vec![0; 32]),
359 created_at: NaiveDateTime::default(),
360 };
361
362 let snapshot = ComponentWithState {
363 state: ResponseProtocolState {
364 component_id: ST_ETH_ADDRESS_PROXY.to_owned(),
365 attributes: HashMap::from([
366 ("total_shares".to_string(), Bytes::from(vec![0; 32])),
367 ("total_wstETH".to_string(), Bytes::from(vec![0; 32])),
368 ]),
369 balances: HashMap::from([(
370 Bytes::from_str(ST_ETH_ADDRESS_PROXY).unwrap(),
371 Bytes::from(vec![0; 32]),
372 )]),
373 },
374 component: pc,
375 component_tvl: None,
376 entrypoints: Vec::new(),
377 };
378
379 let decoder_context = DecoderContext::new();
380
381 let result = LidoState::try_from_with_header(
382 snapshot,
383 header(),
384 &HashMap::new(),
385 &HashMap::new(),
386 &decoder_context,
387 )
388 .await;
389
390 assert!(result.is_ok());
391 assert_eq!(
392 result.unwrap(),
393 LidoState {
394 pool_type: LidoPoolType::WStEth,
395 total_shares: BigUint::zero(),
396 total_pooled_eth: BigUint::zero(),
397 total_wrapped_st_eth: Some(BigUint::zero()),
398 id: WST_ETH_ADDRESS.into(),
399 native_address: ETH_ADDRESS.into(),
400 stake_limits_state: StakeLimitState {
401 staking_status: crate::evm::protocol::lido::state::StakingStatus::Limited,
402 staking_limit: BigUint::zero(),
403 },
404 tokens: [
405 Bytes::from("0x7f39c581f595b53c5cb19bd0b3f8da6c935e2ca0"),
406 Bytes::from("0xae7ab96520de3a18e5e111b5eaab095312d7fe84"),
407 ],
408 token_to_track_total_pooled_eth: Bytes::from(
409 "0xae7ab96520de3a18e5e111b5eaab095312d7fe84"
410 )
411 }
412 );
413 }
414
415 #[tokio::test]
416 #[rstest]
417 #[case::missing_total_shares("total_shares")]
418 #[case::missing_total_wst_eth("total_wstETH")]
419 async fn test_lido_wst_try_from_missing_attribute(#[case] missing_attribute: &str) {
420 let pc = ProtocolComponent {
421 id: WST_ETH_ADDRESS.to_string(),
422 protocol_system: "protocol_system".to_owned(),
423 protocol_type_name: "protocol_type_name".to_owned(),
424 chain: Chain::Ethereum,
425 tokens: vec![
426 Bytes::from("0x7f39c581f595b53c5cb19bd0b3f8da6c935e2ca0"),
427 Bytes::from("0xae7ab96520de3a18e5e111b5eaab095312d7fe84"),
428 ],
429 contract_ids: vec![],
430 static_attributes: HashMap::new(),
431 change: ChangeType::Creation,
432 creation_tx: Bytes::from(vec![0; 32]),
433 created_at: NaiveDateTime::default(),
434 };
435
436 let mut snapshot = ComponentWithState {
437 state: ResponseProtocolState {
438 component_id: ST_ETH_ADDRESS_PROXY.to_owned(),
439 attributes: HashMap::from([
440 ("total_shares".to_string(), Bytes::from(vec![0; 32])),
441 ("total_wstETH".to_string(), Bytes::from(vec![0; 32])),
442 ]),
443 balances: HashMap::from([(
444 Bytes::from(ST_ETH_ADDRESS_PROXY),
445 Bytes::from(vec![0; 32]),
446 )]),
447 },
448 component: pc,
449 component_tvl: None,
450 entrypoints: Vec::new(),
451 };
452 snapshot
453 .state
454 .attributes
455 .remove(missing_attribute);
456
457 let decoder_context = DecoderContext::new();
458
459 let result = LidoState::try_from_with_header(
460 snapshot,
461 header(),
462 &HashMap::new(),
463 &HashMap::new(),
464 &decoder_context,
465 )
466 .await;
467
468 assert!(result.is_err());
469 }
470}