Skip to main content

yip_ca/
serve.rs

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
21/// Reads a length-prefixed postcard message from an async stream.
22pub 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
40/// Writes a length-prefixed postcard message to an async stream.
41pub 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
54/// Shared coordinator state managing connected signers and ROAST coordination.
55pub 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
73/// Runs the Coordinator daemon loop.
74pub 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                    // Handle persistent Signer connection
94                    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                    // Handle client enrollment challenge request
131                    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                    // Await enrollment submission
148                    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                    // Start ROAST signing round
178                    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                        // Wait for ROAST completion (polling shares)
211                        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
276/// Runs a persistent Signer daemon connecting to the Coordinator.
277pub 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    // Preprocess initial batch of 16 commitments
287    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}