ckb_light_client_protocol_server/
lib.rs

1//! Server-side implementation for CKB light client protocol.
2//!
3//! TODO(light-client) More documentation.
4
5use std::sync::Arc;
6
7use ckb_logger::{debug, error, info, trace, warn};
8use ckb_network::{CKBProtocolContext, CKBProtocolHandler, PeerIndex, async_trait, bytes::Bytes};
9use ckb_shared::Shared;
10use ckb_store::ChainStore;
11use ckb_types::{core, packed, prelude::*};
12
13use crate::prelude::*;
14
15mod components;
16mod constant;
17mod prelude;
18mod status;
19
20#[cfg(test)]
21mod tests;
22
23pub use status::{Status, StatusCode};
24
25/// Light client protocol handler.
26pub struct LightClientProtocol {
27    /// Sync shared state.
28    pub shared: Shared,
29}
30
31impl LightClientProtocol {
32    /// Create a new light client protocol handler.
33    pub fn new(shared: Shared) -> Self {
34        Self { shared }
35    }
36}
37
38#[async_trait]
39impl CKBProtocolHandler for LightClientProtocol {
40    async fn init(&mut self, _nc: Arc<dyn CKBProtocolContext + Sync>) {}
41
42    async fn connected(
43        &mut self,
44        _nc: Arc<dyn CKBProtocolContext + Sync>,
45        peer: PeerIndex,
46        version: &str,
47    ) {
48        info!("LightClient({}).connected peer={}", version, peer);
49    }
50
51    async fn disconnected(&mut self, _nc: Arc<dyn CKBProtocolContext + Sync>, peer: PeerIndex) {
52        info!("LightClient.disconnected peer={}", peer);
53    }
54
55    async fn received(
56        &mut self,
57        nc: Arc<dyn CKBProtocolContext + Sync>,
58        peer: PeerIndex,
59        data: Bytes,
60    ) {
61        trace!("LightClient.received peer={}", peer);
62
63        let msg = match packed::LightClientMessageReader::from_slice(&data) {
64            Ok(msg) => msg.to_enum(),
65            _ => {
66                warn!(
67                    "LightClient.received a malformed message from Peer({})",
68                    peer
69                );
70                nc.ban_peer(
71                    peer,
72                    constant::BAD_MESSAGE_BAN_TIME,
73                    String::from("send us a malformed message"),
74                );
75                return;
76            }
77        };
78
79        let item_name = msg.item_name();
80        let status = self.try_process(&nc, peer, msg).await;
81        if let Some(ban_time) = status.should_ban() {
82            error!(
83                "process {} from {}; ban {:?} since result is {}",
84                item_name, peer, ban_time, status
85            );
86            nc.ban_peer(peer, ban_time, status.to_string());
87        } else if status.should_warn() {
88            warn!("process {} from {}; result is {}", item_name, peer, status);
89        } else if !status.is_ok() {
90            debug!("process {} from {}; result is {}", item_name, peer, status);
91        }
92    }
93}
94
95impl LightClientProtocol {
96    async fn try_process(
97        &mut self,
98        nc: &Arc<dyn CKBProtocolContext + Sync>,
99        peer_index: PeerIndex,
100        message: packed::LightClientMessageUnionReader<'_>,
101    ) -> Status {
102        match message {
103            packed::LightClientMessageUnionReader::GetLastState(reader) => {
104                components::GetLastStateProcess::new(reader, self, peer_index, nc)
105                    .execute()
106                    .await
107            }
108            packed::LightClientMessageUnionReader::GetLastStateProof(reader) => {
109                components::GetLastStateProofProcess::new(reader, self, peer_index, nc)
110                    .execute()
111                    .await
112            }
113            packed::LightClientMessageUnionReader::GetBlocksProof(reader) => {
114                components::GetBlocksProofProcess::new(reader, self, peer_index, nc)
115                    .execute()
116                    .await
117            }
118            packed::LightClientMessageUnionReader::GetTransactionsProof(reader) => {
119                components::GetTransactionsProofProcess::new(reader, self, peer_index, nc)
120                    .execute()
121                    .await
122            }
123            _ => StatusCode::UnexpectedProtocolMessage.into(),
124        }
125    }
126
127    pub(crate) fn get_verifiable_tip_header(&self) -> Result<packed::VerifiableHeader, String> {
128        let snapshot = self.shared.snapshot();
129
130        let tip_hash = snapshot.tip_hash();
131        let tip_block = snapshot
132            .get_block(&tip_hash)
133            .expect("checked: tip block should be existed");
134        let parent_chain_root = if tip_block.is_genesis() {
135            Default::default()
136        } else {
137            let mmr = snapshot.chain_root_mmr(tip_block.number() - 1);
138            match mmr.get_root() {
139                Ok(root) => root,
140                Err(err) => {
141                    let errmsg = format!("failed to generate a root since {err:?}");
142                    return Err(errmsg);
143                }
144            }
145        };
146
147        let tip_header = packed::VerifiableHeader::new_builder()
148            .header(tip_block.header().data())
149            .uncles_hash(tip_block.calc_uncles_hash())
150            .extension(Pack::pack(&tip_block.extension()))
151            .parent_chain_root(parent_chain_root)
152            .build();
153
154        Ok(tip_header)
155    }
156
157    pub(crate) async fn reply_tip_state<T>(
158        &self,
159        peer: PeerIndex,
160        nc: &Arc<dyn CKBProtocolContext + Sync>,
161    ) -> Status
162    where
163        T: Entity,
164        <T as Entity>::Builder: ProverMessageBuilder,
165        <<T as Entity>::Builder as Builder>::Entity: Into<packed::LightClientMessageUnion>,
166    {
167        let tip_header = match self.get_verifiable_tip_header() {
168            Ok(tip_state) => tip_state,
169            Err(errmsg) => {
170                return StatusCode::InternalError.with_context(errmsg);
171            }
172        };
173        let content = T::new_builder().set_last_header(tip_header).build();
174        let message = packed::LightClientMessage::new_builder()
175            .set(content)
176            .build();
177        nc.reply(peer, &message).await;
178        Status::ok()
179    }
180
181    pub(crate) async fn reply_proof<T>(
182        &self,
183        peer: PeerIndex,
184        nc: &Arc<dyn CKBProtocolContext + Sync>,
185        last_block: &core::BlockView,
186        items_positions: Vec<u64>,
187        proved_items: <<T as Entity>::Builder as ProverMessageBuilder>::ProvedItems,
188        missing_items: <<T as Entity>::Builder as ProverMessageBuilder>::MissingItems,
189    ) -> Status
190    where
191        T: Entity,
192        <T as Entity>::Builder: ProverMessageBuilder,
193        <<T as Entity>::Builder as Builder>::Entity: Into<packed::LightClientMessageUnion>,
194    {
195        let (parent_chain_root, proof) = {
196            let snapshot = self.shared.snapshot();
197            let mmr = snapshot.chain_root_mmr(last_block.number() - 1);
198            let parent_chain_root = match mmr.get_root() {
199                Ok(root) => root,
200                Err(err) => {
201                    let errmsg = format!("failed to generate a root since {err:?}");
202                    return StatusCode::InternalError.with_context(errmsg);
203                }
204            };
205            let proof = if items_positions.is_empty() {
206                Default::default()
207            } else {
208                match mmr.gen_proof(items_positions) {
209                    Ok(proof) => proof.proof_items().to_owned(),
210                    Err(err) => {
211                        let errmsg = format!("failed to generate a proof since {err:?}");
212                        return StatusCode::InternalError.with_context(errmsg);
213                    }
214                }
215            };
216            (parent_chain_root, proof)
217        };
218        let verifiable_last_header = packed::VerifiableHeader::new_builder()
219            .header(last_block.data().header())
220            .uncles_hash(last_block.calc_uncles_hash())
221            .extension(Pack::pack(&last_block.extension()))
222            .parent_chain_root(parent_chain_root)
223            .build();
224        let content = T::new_builder()
225            .set_last_header(verifiable_last_header)
226            .set_proof(proof.into())
227            .set_proved_items(proved_items)
228            .set_missing_items(missing_items)
229            .build();
230        let message = packed::LightClientMessage::new_builder()
231            .set(content)
232            .build();
233        nc.reply(peer, &message).await;
234        Status::ok()
235    }
236}