Skip to main content

espresso_node/
state_signature.rs

1//! Utilities for generating and storing the most recent light client state signatures.
2
3use std::collections::{HashMap, VecDeque};
4
5use alloy::primitives::FixedBytes;
6use async_lock::RwLock;
7use hotshot::types::{Event, EventType, SchnorrPubKey};
8use hotshot_contract_adapter::light_client::derive_signed_state_digest;
9use hotshot_types::{
10    data::{EpochNumber, Leaf2},
11    event::LeafInfo,
12    light_client::{
13        LCV2StateSignatureRequestBody, LCV3StateSignatureRequestBody, LightClientState,
14        StakeTableState, StateSignKey, StateSignature, StateVerKey,
15    },
16    new_protocol::CoordinatorEvent,
17    stake_table::HSStakeTable,
18    traits::{
19        block_contents::BlockHeader,
20        signature_key::{
21            LCV1StateSignatureKey, LCV2StateSignatureKey, LCV3StateSignatureKey,
22            StakeTableEntryType,
23        },
24    },
25    utils::{is_ge_epoch_root, option_epoch_from_block_number},
26};
27use jf_signature::SignatureError;
28use surf_disco::{Client, Url};
29use tide_disco::error::ServerError;
30use vbs::version::StaticVersionType;
31
32use crate::{SeqTypes, consensus_handle::ConsensusHandle};
33
34/// A relay server that's collecting and serving the light client state signatures
35pub mod relay_server;
36
37/// Capacity for the in memory signature storage.
38const SIGNATURE_STORAGE_CAPACITY: usize = 100;
39
40#[derive(Debug)]
41pub struct StateSigner<ApiVer: StaticVersionType> {
42    /// Key for signing a new light client state
43    sign_key: StateSignKey,
44
45    /// Key for verifying a light client state
46    ver_key: StateVerKey,
47
48    /// The most recent light client state signatures
49    signatures: RwLock<StateSignatureMemStorage>,
50
51    /// Commitment for current fixed stake table
52    voting_stake_table_state: StakeTableState,
53
54    /// epoch for the current stake table state
55    voting_stake_table_epoch: Option<EpochNumber>,
56
57    /// Capacity of the stake table
58    stake_table_capacity: usize,
59
60    /// Indicate whether the signer is in the voting stake table and should sign the state
61    should_vote: bool,
62
63    /// The state relay server url
64    relay_server_client: Option<Client<ServerError, ApiVer>>,
65}
66
67impl<ApiVer: StaticVersionType> StateSigner<ApiVer> {
68    pub fn new(
69        sign_key: StateSignKey,
70        ver_key: StateVerKey,
71        voting_stake_table_state: StakeTableState,
72        voting_stake_table_epoch: Option<EpochNumber>,
73        stake_table_capacity: usize,
74        should_vote: bool,
75    ) -> Self {
76        Self {
77            sign_key,
78            ver_key,
79            voting_stake_table_state,
80            voting_stake_table_epoch,
81            stake_table_capacity,
82            signatures: Default::default(),
83            should_vote,
84            relay_server_client: Default::default(),
85        }
86    }
87
88    /// Connect to the given state relay server to send signed HotShot states to.
89    pub fn with_relay_server(mut self, url: Url) -> Self {
90        self.relay_server_client = Some(Client::new(url));
91        self
92    }
93
94    pub(super) async fn handle_event<I>(
95        &mut self,
96        event: &CoordinatorEvent<SeqTypes>,
97        consensus_handle: &ConsensusHandle<SeqTypes, I>,
98    ) where
99        I: hotshot::traits::NodeImplementation<SeqTypes>,
100        I::Storage: hotshot_new_protocol::storage::NewProtocolStorage<SeqTypes>,
101    {
102        let leaf: &Leaf2<SeqTypes> = match event {
103            CoordinatorEvent::LegacyEvent(Event {
104                event: EventType::Decide { leaf_chain, .. },
105                ..
106            }) => match leaf_chain.first() {
107                Some(LeafInfo { leaf, .. }) => leaf,
108                None => return,
109            },
110            CoordinatorEvent::NewDecide { leaf_infos, .. } => match leaf_infos.first() {
111                Some(info) => &info.leaf,
112                None => return,
113            },
114            _ => return,
115        };
116        match leaf
117            .block_header()
118            .get_light_client_state(leaf.view_number())
119        {
120            Ok(state) => {
121                tracing::debug!("New leaves decided. Latest block height: {}", leaf.height(),);
122
123                let cur_block_height = state.block_height;
124                let blocks_per_epoch = *consensus_handle.epoch_height().await;
125
126                let option_state_epoch = option_epoch_from_block_number(
127                    leaf.with_epoch,
128                    cur_block_height,
129                    blocks_per_epoch,
130                );
131
132                if self.voting_stake_table_epoch != option_state_epoch {
133                    let Ok(membership) = consensus_handle
134                        .membership_coordinator()
135                        .await
136                        .stake_table_for_epoch(option_state_epoch)
137                    else {
138                        tracing::error!(
139                            "Failed to get membership for epoch: {:?}",
140                            option_state_epoch
141                        );
142                        return;
143                    };
144                    let stake_table = HSStakeTable::from_iter(membership.stake_table());
145                    match stake_table.commitment(self.stake_table_capacity) {
146                        Ok(stake_table_state) => {
147                            self.should_vote = should_vote(&stake_table, &self.ver_key);
148                            self.voting_stake_table_epoch = option_state_epoch;
149                            self.voting_stake_table_state = stake_table_state;
150                        },
151                        Err(err) => {
152                            tracing::error!("Failed to compute stake table commitment: {:?}", err);
153                            return;
154                        },
155                    }
156                }
157
158                if !self.should_vote {
159                    tracing::debug!(
160                        "Not signing the state at block height {} since not in the voting stake \
161                         table",
162                        state.block_height
163                    );
164                    return;
165                }
166
167                // The last few state updates are handled in the consensus, we do not sign them.
168                if leaf.with_epoch & is_ge_epoch_root(cur_block_height, blocks_per_epoch) {
169                    tracing::debug!("Skipping epoch transition block {cur_block_height}");
170                    return;
171                }
172
173                let Ok(auth_root) = leaf.block_header().auth_root() else {
174                    tracing::error!("Failed to get auth root for light client state");
175                    return;
176                };
177
178                let Ok(request_body) = self
179                    .get_request_body(&state, &self.voting_stake_table_state, auth_root)
180                    .await
181                else {
182                    tracing::error!("Failed to sign new state");
183                    return;
184                };
185
186                if let Some(client) = &self.relay_server_client {
187                    if let Err(error) = client
188                        .post::<()>("api/state")
189                        .body_binary(&request_body)
190                        .unwrap()
191                        .send()
192                        .await
193                    {
194                        tracing::error!("Error posting signature to the relay server: {:?}", error);
195                    }
196
197                    if !leaf.with_epoch {
198                        // Before epoch upgrade, we need to sign the state for the legacy light client
199                        let Ok(legacy_signature) = self.legacy_sign_new_state(&state).await else {
200                            tracing::error!("Failed to sign new state for legacy light client");
201                            return;
202                        };
203                        let legacy_request_body = LCV2StateSignatureRequestBody {
204                            key: self.ver_key.clone(),
205                            state,
206                            next_stake: StakeTableState::default(),
207                            signature: legacy_signature,
208                        };
209                        if let Err(error) = client
210                            .post::<()>("api/legacy-state")
211                            .body_binary(&legacy_request_body)
212                            .unwrap()
213                            .send()
214                            .await
215                        {
216                            tracing::error!(
217                                "Error posting signature for legacy light client to the relay \
218                                 server: {:?}",
219                                error
220                            );
221                        }
222                    }
223                }
224            },
225            Err(err) => {
226                tracing::error!("Error generating light client state: {:?}", err)
227            },
228        }
229    }
230
231    /// Return a signature of a light client state at given height.
232    pub async fn get_state_signature(&self, height: u64) -> Option<LCV3StateSignatureRequestBody> {
233        let pool_guard = self.signatures.read().await;
234        pool_guard.get_signature(height)
235    }
236
237    /// Sign the light client state at given height and store it.
238    async fn get_request_body(
239        &self,
240        state: &LightClientState,
241        next_stake_table: &StakeTableState,
242        auth_root: FixedBytes<32>,
243    ) -> Result<LCV3StateSignatureRequestBody, SignatureError> {
244        let signed_state_digest = derive_signed_state_digest(state, next_stake_table, &auth_root);
245        let signature = <SchnorrPubKey as LCV3StateSignatureKey>::sign_state(
246            &self.sign_key,
247            signed_state_digest,
248        )?;
249        let v2signature = <SchnorrPubKey as LCV2StateSignatureKey>::sign_state(
250            &self.sign_key,
251            state,
252            next_stake_table,
253        )?;
254        let request_body = LCV3StateSignatureRequestBody {
255            key: self.ver_key.clone(),
256            state: *state,
257            next_stake: *next_stake_table,
258            signature,
259            v2_signature: v2signature.clone(),
260            auth_root,
261        };
262        let mut pool_guard = self.signatures.write().await;
263        pool_guard.push(state.block_height, request_body.clone());
264        tracing::debug!(
265            "New signature added for block height {}",
266            state.block_height
267        );
268        Ok(request_body)
269    }
270
271    async fn legacy_sign_new_state(
272        &self,
273        state: &LightClientState,
274    ) -> Result<StateSignature, SignatureError> {
275        <SchnorrPubKey as LCV1StateSignatureKey>::sign_state(&self.sign_key, state)
276    }
277}
278
279/// A rolling in-memory storage for the most recent light client state signatures.
280#[derive(Debug, Default)]
281pub struct StateSignatureMemStorage {
282    pool: HashMap<u64, LCV3StateSignatureRequestBody>,
283    deque: VecDeque<u64>,
284}
285
286impl StateSignatureMemStorage {
287    pub fn push(&mut self, height: u64, signature: LCV3StateSignatureRequestBody) {
288        self.pool.insert(height, signature);
289        self.deque.push_back(height);
290        if self.pool.len() > SIGNATURE_STORAGE_CAPACITY {
291            self.pool.remove(&self.deque.pop_front().unwrap());
292        }
293    }
294
295    pub fn get_signature(&self, height: u64) -> Option<LCV3StateSignatureRequestBody> {
296        self.pool.get(&height).cloned()
297    }
298}
299
300pub(crate) fn should_vote(stake_table: &HSStakeTable<SeqTypes>, ver_key: &SchnorrPubKey) -> bool {
301    stake_table
302        .0
303        .iter()
304        .any(|peer| &peer.state_ver_key == ver_key && !peer.stake_table_entry.stake().is_zero())
305}