lit_rust_sdk/client/
connect.rs

1// use crate::blockchain::Staking;
2use crate::{
3    blockchain::Staking,
4    types::{HandshakeRequest, HandshakeResponse, NodeConnectionInfo},
5};
6use eyre::Result;
7use rand::Rng;
8use tracing::{info, warn};
9
10impl<P: alloy::providers::Provider> super::LitNodeClient<P> {
11    pub async fn connect(&mut self) -> Result<()> {
12        info!(
13            "Starting connection to Lit Network: {:?}",
14            self.config.lit_network
15        );
16
17        // let _epoch = self.current_epoch_state().await?;
18        // TODO: initialize the listener
19
20        let network_info = self.get_network_info().await?;
21        info!("Found network info: {:?}", network_info);
22
23        let validators = network_info._2;
24        let mut bootstrap_urls = Vec::with_capacity(validators.len());
25        for validator in validators {
26            // Convert validator.ip (u32) to standard IPv4 string
27            let ip_addr = std::net::Ipv4Addr::from(validator.ip);
28            bootstrap_urls.push(format!("https://{}:{}", ip_addr, validator.port));
29        }
30
31        let min_node_count = network_info._1.to::<usize>();
32        self.min_node_count = Some(min_node_count);
33
34        let epoch = network_info._0;
35        self.epoch = Some(epoch);
36
37        self.handshake_with_nodes(&bootstrap_urls).await?;
38
39        self.update_network_state_from_consensus();
40        self.ready = true;
41        info!("Successfully connected to Lit Network");
42        Ok(())
43    }
44
45    async fn get_network_info(
46        &self,
47    ) -> Result<Staking::getActiveUnkickedValidatorStructsAndCountsReturn> {
48        let network_info = self
49            .staking
50            .getActiveUnkickedValidatorStructsAndCounts()
51            .call()
52            .await?;
53        Ok(network_info)
54    }
55
56    async fn handshake_with_nodes(&mut self, urls: &[String]) -> Result<()> {
57        let mut successful_connections = 0;
58        for url in urls {
59            match self.handshake_with_node(url).await {
60                Ok(response) => {
61                    info!("Successfully connected to node: {}", url);
62                    self.connection_state.insert(
63                        url.clone(),
64                        NodeConnectionInfo {
65                            url: url.clone(),
66                            handshake_response: response,
67                        },
68                    );
69                    successful_connections += 1;
70                }
71                Err(e) => {
72                    warn!("Failed to connect to node {}: {}", url, e);
73                }
74            }
75        }
76
77        if successful_connections < urls.len() {
78            return Err(eyre::eyre!(format!(
79                "Not enough nodes connected. Connected: {}, Required: {}",
80                successful_connections,
81                urls.len()
82            )));
83        }
84        Ok(())
85    }
86
87    async fn handshake_with_node(&self, url: &str) -> Result<HandshakeResponse> {
88        let challenge = self.generate_challenge();
89        let request = HandshakeRequest {
90            client_public_key: "test".to_string(),
91            challenge: challenge.clone(),
92        };
93        let handshake_url = format!("{}/web/handshake", url);
94        let request_id = self.generate_request_id();
95
96        info!("Sending handshake to {}: {:?}", handshake_url, request);
97
98        let response = self
99            .http_client
100            .post(&handshake_url)
101            .header("X-Request-Id", request_id)
102            .json(&request)
103            .send()
104            .await?;
105
106        if !response.status().is_success() {
107            let status = response.status();
108            let body = response
109                .text()
110                .await
111                .unwrap_or_else(|_| "Unable to read body".to_string());
112            warn!("Handshake failed with status {}: {}", status, body);
113            return Err(eyre::eyre!(format!(
114                "Handshake failed with status {}: {}",
115                status, body
116            )));
117        }
118
119        let body_text = response.text().await?;
120        info!("Handshake response body: {}", body_text);
121
122        let handshake_response: HandshakeResponse =
123            serde_json::from_str(&body_text).map_err(|e| {
124                warn!("Failed to parse handshake response: {}", e);
125                eyre::eyre!(e)
126            })?;
127        Ok(handshake_response)
128    }
129
130    // async fn current_epoch_state(&self) -> Result<Staking::Epoch> {
131    //     let epoch = self.staking.epoch().call().await?;
132    //     epoch.
133    //     Ok(epoch)
134    // }
135
136    fn generate_challenge(&self) -> String {
137        let mut rng = rand::thread_rng();
138        let bytes: Vec<u8> = (0..32).map(|_| rng.gen()).collect();
139        hex::encode(bytes)
140    }
141
142    pub(crate) fn generate_request_id(&self) -> String {
143        let mut rng = rand::thread_rng();
144        let bytes: Vec<u8> = (0..16).map(|_| rng.gen()).collect();
145        hex::encode(bytes)
146    }
147
148    pub(crate) fn update_network_state_from_consensus(&mut self) {
149        let responses: Vec<HandshakeResponse> = self
150            .connection_state
151            .iter()
152            .map(|entry| entry.handshake_response.clone())
153            .collect();
154        if responses.is_empty() {
155            return;
156        }
157        let first = &responses[0];
158        self.subnet_pub_key = Some(first.subnet_pub_key.clone());
159        self.network_pub_key = Some(first.network_pub_key.clone());
160        self.network_pub_key_set = Some(first.network_pub_key_set.clone());
161        self.hd_root_pubkeys = Some(first.hd_root_pubkeys.clone());
162        self.latest_blockhash = Some(first.latest_blockhash.clone());
163    }
164}