1use 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
34pub mod relay_server;
36
37const SIGNATURE_STORAGE_CAPACITY: usize = 100;
39
40#[derive(Debug)]
41pub struct StateSigner<ApiVer: StaticVersionType> {
42 sign_key: StateSignKey,
44
45 ver_key: StateVerKey,
47
48 signatures: RwLock<StateSignatureMemStorage>,
50
51 voting_stake_table_state: StakeTableState,
53
54 voting_stake_table_epoch: Option<EpochNumber>,
56
57 stake_table_capacity: usize,
59
60 should_vote: bool,
62
63 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 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 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 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 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 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#[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}