1use std::collections::BTreeMap;
2use std::net::SocketAddr;
3use std::sync::Arc;
4use std::time::{SystemTime, UNIX_EPOCH};
5
6use frost_ed25519 as frost;
7use rand::RngCore;
8use tokio::io::{AsyncReadExt, AsyncWriteExt};
9use tokio::net::{TcpListener, TcpStream};
10use tokio::sync::mpsc;
11use tokio::sync::Mutex;
12
13use crate::roast_coordinator::RoastCoordinator;
14use crate::rpc::{
15 verify_enrollment_pop, CaWireMessage, ClientRpcRequest, ClientRpcResponse, QuorumRpcMessage,
16 SignCertRequestPayload,
17};
18use crate::signer_daemon::{SignerPolicy, SignerState};
19use yip_membership::cert::Cert;
20
21pub async fn read_framed<T: serde::de::DeserializeOwned>(
23 stream: &mut TcpStream,
24) -> Result<T, std::io::Error> {
25 let mut len_buf = [0u8; 4];
26 stream.read_exact(&mut len_buf).await?;
27 let len = u32::from_be_bytes(len_buf) as usize;
28 if len > 10 * 1024 * 1024 {
29 return Err(std::io::Error::new(
30 std::io::ErrorKind::InvalidData,
31 "message frame too large",
32 ));
33 }
34 let mut buf = vec![0u8; len];
35 stream.read_exact(&mut buf).await?;
36 postcard::from_bytes(&buf)
37 .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e.to_string()))
38}
39
40pub async fn write_framed<T: serde::Serialize>(
42 stream: &mut TcpStream,
43 msg: &T,
44) -> Result<(), std::io::Error> {
45 let bytes = postcard::to_allocvec(msg)
46 .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e.to_string()))?;
47 let len = bytes.len() as u32;
48 stream.write_all(&len.to_be_bytes()).await?;
49 stream.write_all(&bytes).await?;
50 stream.flush().await?;
51 Ok(())
52}
53
54pub struct CoordinatorState {
56 pub threshold: usize,
57 pub pubkey_package: frost::keys::PublicKeyPackage,
58 pub signer_txs: BTreeMap<frost::Identifier, mpsc::Sender<QuorumRpcMessage>>,
59 pub roast: RoastCoordinator,
60}
61
62impl CoordinatorState {
63 pub fn new(threshold: usize, pubkey_package: frost::keys::PublicKeyPackage) -> Self {
64 Self {
65 threshold,
66 pubkey_package: pubkey_package.clone(),
67 signer_txs: BTreeMap::new(),
68 roast: RoastCoordinator::new(threshold, pubkey_package, Vec::new()),
69 }
70 }
71}
72
73pub async fn run_coordinator(
75 listener: TcpListener,
76 threshold: usize,
77 pubkey_package: frost::keys::PublicKeyPackage,
78) -> Result<(), std::io::Error> {
79 let state = Arc::new(Mutex::new(CoordinatorState::new(threshold, pubkey_package)));
80
81 loop {
82 let (mut socket, _) = listener.accept().await?;
83 let state = Arc::clone(&state);
84
85 tokio::spawn(async move {
86 let first_msg: Result<CaWireMessage, _> = read_framed(&mut socket).await;
87 match first_msg {
88 Ok(CaWireMessage::Quorum(QuorumRpcMessage::PublishCommitments {
89 signer_id,
90 commitments,
91 ..
92 })) => {
93 let (tx, mut rx) = mpsc::channel::<QuorumRpcMessage>(32);
95 {
96 let mut s = state.lock().await;
97 s.signer_txs.insert(signer_id, tx);
98 s.roast.register_commitments(signer_id, commitments);
99 }
100
101 loop {
102 tokio::select! {
103 Some(out_msg) = rx.recv() => {
104 let wire_msg = CaWireMessage::Quorum(out_msg);
105 if write_framed(&mut socket, &wire_msg).await.is_err() {
106 break;
107 }
108 }
109 in_msg = read_framed::<CaWireMessage>(&mut socket) => {
110 match in_msg {
111 Ok(CaWireMessage::Quorum(QuorumRpcMessage::SignCertResponse(payload))) => {
112 let mut s = state.lock().await;
113 let _ = s.roast.receive_share(
114 payload.session_id,
115 payload.signer_id,
116 payload.signature_share,
117 payload.replenish_commitment,
118 );
119 }
120 _ => break,
121 }
122 }
123 }
124 }
125
126 let mut s = state.lock().await;
127 s.signer_txs.remove(&signer_id);
128 }
129 Ok(CaWireMessage::ClientReq(ClientRpcRequest::GetChallenge)) => {
130 let mut challenge_nonce = [0u8; 32];
132 rand::rngs::OsRng.fill_bytes(&mut challenge_nonce);
133 let expires_at_unix = SystemTime::now()
134 .duration_since(UNIX_EPOCH)
135 .unwrap()
136 .as_secs()
137 + 300;
138
139 let resp = CaWireMessage::ClientResp(ClientRpcResponse::Challenge {
140 challenge_nonce,
141 expires_at_unix,
142 });
143 if write_framed(&mut socket, &resp).await.is_err() {
144 return;
145 }
146
147 let submit_req: Result<CaWireMessage, _> = read_framed(&mut socket).await;
149 let payload = match submit_req {
150 Ok(CaWireMessage::ClientReq(ClientRpcRequest::SubmitEnrollment(p))) => *p,
151 _ => return,
152 };
153
154 if payload.challenge_nonce != challenge_nonce {
155 let rej = CaWireMessage::ClientResp(ClientRpcResponse::Rejected {
156 reason: "challenge nonce mismatch".into(),
157 });
158 let _ = write_framed(&mut socket, &rej).await;
159 return;
160 }
161
162 if let Err(e) = verify_enrollment_pop(
163 &payload.pop_signature,
164 &payload.csr_tbs.member_sign_pubkey,
165 &challenge_nonce,
166 &payload.csr_tbs,
167 ) {
168 let rej = CaWireMessage::ClientResp(ClientRpcResponse::Rejected {
169 reason: format!("PoP verification failed: {e}"),
170 });
171 let _ = write_framed(&mut socket, &rej).await;
172 return;
173 }
174
175 let msg = payload.csr_tbs.to_canonical_bytes();
176
177 let mut session_info = None;
179 let mut signers_to_dispatch = Vec::new();
180
181 {
182 let mut s = state.lock().await;
183 s.roast.message = msg.clone();
184 if let Some((session_id, pkg, slot_map)) = s.roast.try_spawn_session() {
185 session_info = Some((session_id, pkg.clone()));
186 for (id, slot) in slot_map {
187 if let Some(tx) = s.signer_txs.get(&id) {
188 let req = SignCertRequestPayload {
189 session_id,
190 slot_id: slot,
191 signing_package: pkg.clone(),
192 cert_tbs: payload.csr_tbs.clone(),
193 credential: payload.credential.clone(),
194 challenge_nonce,
195 pop_signature: payload.pop_signature,
196 };
197 signers_to_dispatch.push((tx.clone(), req));
198 }
199 }
200 }
201 }
202
203 if let Some((session_id, _)) = session_info {
204 for (tx, req) in signers_to_dispatch {
205 let _ = tx
206 .send(QuorumRpcMessage::SignCertRequest(Box::new(req)))
207 .await;
208 }
209
210 let mut final_sig = None;
212 for _ in 0..100 {
213 tokio::time::sleep(std::time::Duration::from_millis(15)).await;
214 let s = state.lock().await;
215 if let Some(session) = s.roast.active_sessions.get(&session_id) {
216 if session.collected_shares.len() >= s.threshold {
217 if let Ok(sig) = frost::aggregate(
218 &session.signing_package,
219 &session.collected_shares,
220 &s.pubkey_package,
221 ) {
222 final_sig = Some(sig);
223 break;
224 }
225 }
226 }
227 }
228
229 if let Some(sig) = final_sig {
230 let sig_bytes = sig.serialize().unwrap();
231 let mut sig_arr = [0u8; 64];
232 sig_arr.copy_from_slice(&sig_bytes);
233
234 let cert = Cert {
235 version: payload.csr_tbs.version,
236 member_pubkey: payload.csr_tbs.member_pubkey,
237 member_sign_pubkey: payload.csr_tbs.member_sign_pubkey,
238 network_id: payload.csr_tbs.network_id,
239 not_before: payload.csr_tbs.not_before,
240 not_after: payload.csr_tbs.not_after,
241 tags: payload.csr_tbs.tags.clone(),
242 ca_sig: sig_arr,
243 };
244
245 let mut cert_bytes = Vec::new();
246 cert.encode(&mut cert_bytes);
247
248 let root_vk = {
249 let s = state.lock().await;
250 let vk_b = s.pubkey_package.verifying_key().serialize().unwrap();
251 let mut arr = [0u8; 32];
252 arr.copy_from_slice(&vk_b);
253 arr
254 };
255
256 let enrolled = CaWireMessage::ClientResp(ClientRpcResponse::Enrolled {
257 cert_bytes,
258 ca_public_key: root_vk,
259 });
260 let _ = write_framed(&mut socket, &enrolled).await;
261 return;
262 }
263 }
264
265 let rej = CaWireMessage::ClientResp(ClientRpcResponse::Rejected {
266 reason: "threshold signing quorum unreachable or stalled".into(),
267 });
268 let _ = write_framed(&mut socket, &rej).await;
269 }
270 _ => {}
271 }
272 });
273 }
274}
275
276pub async fn run_signer(
278 coord_addr: SocketAddr,
279 key_package: frost::keys::KeyPackage,
280 pubkey_package: frost::keys::PublicKeyPackage,
281 policy: SignerPolicy,
282) -> Result<(), std::io::Error> {
283 let mut stream = TcpStream::connect(coord_addr).await?;
284 let signer = Arc::new(SignerState::new(key_package, pubkey_package, policy));
285
286 let commitments = {
288 let mut vault = signer.vault.lock().unwrap();
289 vault.preprocess_batch(signer.key_package.signing_share(), 16)
290 };
291
292 let pub_msg = CaWireMessage::Quorum(QuorumRpcMessage::PublishCommitments {
293 signer_id: *signer.key_package.identifier(),
294 boot_epoch_id: signer.vault.lock().unwrap().boot_epoch_id(),
295 commitments,
296 });
297 write_framed(&mut stream, &pub_msg).await?;
298
299 loop {
300 let msg: CaWireMessage = match read_framed(&mut stream).await {
301 Ok(m) => m,
302 Err(_) => break,
303 };
304
305 if let CaWireMessage::Quorum(QuorumRpcMessage::SignCertRequest(req)) = msg {
306 let resp = signer.handle_sign_request(&req);
307 write_framed(&mut stream, &CaWireMessage::Quorum(resp)).await?;
308 }
309 }
310
311 Ok(())
312}