-
Notifications
You must be signed in to change notification settings - Fork 96
/
epoch_state.rs
409 lines (376 loc) · 15 KB
/
epoch_state.rs
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
use std::collections::btree_map::Entry;
use std::collections::{BTreeMap, BTreeSet};
use std::fmt::{self, Display};
use std::marker::PhantomData;
use std::mem::replace;
use std::result;
use std::sync::Arc;
use crate::crypto::Ciphertext;
use bincode;
use log::error;
use rand::Rng;
use serde::{de::DeserializeOwned, Deserialize, Serialize};
use super::{Batch, Error, FaultKind, FaultLog, MessageContent, Result, Step};
use crate::fault_log::Fault;
use crate::subset::{self as cs, Subset, SubsetOutput};
use crate::threshold_decrypt::{self as td, ThresholdDecrypt};
use crate::{Contribution, NetworkInfo, NodeIdT};
type CsStep<N> = cs::Step<N>;
/// The status of an encrypted contribution.
#[derive(Debug)]
enum DecryptionState<N> {
/// Decryption is still ongoing; we are waiting for decryption shares and/or ciphertext.
Ongoing(Box<ThresholdDecrypt<N>>),
/// Decryption is complete. This contains the plaintext.
Complete(Vec<u8>),
}
impl<N: NodeIdT> DecryptionState<N> {
/// Creates a new `ThresholdDecrypt` instance, waiting for shares and a ciphertext.
fn new(netinfo: Arc<NetworkInfo<N>>) -> Self {
DecryptionState::Ongoing(Box::new(ThresholdDecrypt::new(netinfo)))
}
/// Handles a message containing a decryption share.
fn handle_message(&mut self, sender_id: &N, msg: td::Message) -> td::Result<td::Step<N>> {
match self {
DecryptionState::Ongoing(ref mut td) => td.handle_message(sender_id, msg),
DecryptionState::Complete(_) => Ok(td::Step::default()),
}
}
/// Handles a ciphertext input.
fn set_ciphertext(&mut self, ciphertext: Ciphertext) -> td::Result<td::Step<N>> {
match self {
DecryptionState::Ongoing(ref mut td) => {
td.set_ciphertext(ciphertext)?;
td.start_decryption()
}
DecryptionState::Complete(_) => Ok(td::Step::default()),
}
}
}
/// The status of the subset algorithm.
#[derive(Debug)]
enum SubsetState<N> {
/// The algorithm is ongoing: the set of accepted contributions is still undecided.
Ongoing(Subset<N, EpochId>),
/// The algorithm is complete. This contains the set of accepted proposers.
Complete(BTreeSet<N>),
}
impl<N: NodeIdT> SubsetState<N> {
/// Provides input to the Subset instance, unless it has already completed.
fn handle_input(&mut self, proposal: Vec<u8>) -> Result<CsStep<N>> {
match self {
SubsetState::Ongoing(ref mut cs) => cs.propose(proposal),
SubsetState::Complete(_) => return Ok(cs::Step::default()),
}
.map_err(Error::InputSubset)
}
/// Handles a message in the Subset instance, unless it has already completed.
fn handle_message(&mut self, sender_id: &N, msg: cs::Message<N>) -> Result<CsStep<N>> {
match self {
SubsetState::Ongoing(ref mut cs) => cs.handle_message(sender_id, msg),
SubsetState::Complete(_) => return Ok(cs::Step::default()),
}
.map_err(Error::HandleSubsetMessage)
}
/// Returns the number of contributions that we have already received or, after completion, how
/// many have been accepted.
pub fn received_proposals(&self) -> usize {
match self {
SubsetState::Ongoing(ref cs) => cs.received_proposals(),
SubsetState::Complete(ref proposer_ids) => proposer_ids.len(),
}
}
/// Returns the IDs of the accepted proposers, if that has already been decided.
pub fn accepted_ids(&self) -> Option<&BTreeSet<N>> {
match self {
SubsetState::Ongoing(_) => None,
SubsetState::Complete(ref ids) => Some(ids),
}
}
}
/// A flag used when constructing an `EpochState` to determine which behavior to use when receiving
/// proposals from a `Subset` instance.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub enum SubsetHandlingStrategy {
/// Sets the `EpochState` to return proposals as they are contributed.
Incremental,
/// Sets the `EpochState` to return all received proposals once consensus has been finalized.
AllAtEnd,
}
/// Used in an `EpochState` to encapsulate the state necessary to maintain each
/// `SubsetHandlingStrategy`.
#[derive(Debug, Clone)]
enum SubsetHandler<N> {
Incremental,
AllAtEnd(Vec<(N, Vec<u8>)>),
}
/// The result of a call to `SubsetHandler::handle(...)`.
struct SubsetHandleData<N> {
/// The number of contributions propagated from the handler.
contributions: Vec<(N, Vec<u8>)>,
/// Indicates whether the underlying `Subset` algorithm has achieved consensus and whether
/// there may be more contributions or not.
is_done: bool,
}
impl<N> SubsetHandler<N> {
fn handle(&mut self, o: SubsetOutput<N>) -> SubsetHandleData<N> {
use self::SubsetHandler::*;
use self::SubsetOutput::*;
let contributions;
let is_done;
match o {
Contribution(proposer_id, data) => {
let proposal = (proposer_id, data);
contributions = match self {
Incremental => vec![proposal],
AllAtEnd(cs) => {
cs.push(proposal);
vec![]
}
};
is_done = false;
}
Done => {
contributions = match self {
Incremental => vec![],
AllAtEnd(cs) => replace(cs, vec![]),
};
is_done = true;
}
}
SubsetHandleData {
contributions,
is_done,
}
}
}
impl<N> From<SubsetHandlingStrategy> for SubsetHandler<N> {
fn from(s: SubsetHandlingStrategy) -> Self {
use self::SubsetHandlingStrategy::*;
match s {
Incremental => SubsetHandler::Incremental,
AllAtEnd => SubsetHandler::AllAtEnd(Vec::new()),
}
}
}
/// The sub-algorithms and their intermediate results for a single epoch.
#[derive(Debug)]
pub struct EpochState<C, N> {
/// Our epoch number.
epoch: u64,
/// Shared network data.
netinfo: Arc<NetworkInfo<N>>,
/// The status of the subset algorithm.
subset: SubsetState<N>,
/// The status of threshold decryption, by proposer.
decryption: BTreeMap<N, DecryptionState<N>>,
/// Nodes found so far in `Subset` output.
accepted_proposers: BTreeSet<N>,
/// Determines the behavior upon receiving proposals from `subset`.
subset_handler: SubsetHandler<N>,
/// Whether contributions should be encrypted in this epoch.
require_decryption: bool,
_phantom: PhantomData<C>,
}
impl<C, N> EpochState<C, N>
where
C: Contribution + Serialize + DeserializeOwned,
N: NodeIdT,
{
/// Creates a new `Subset` instance.
pub fn new(
netinfo: Arc<NetworkInfo<N>>,
hb_id: u64,
epoch: u64,
subset_handling_strategy: SubsetHandlingStrategy,
require_decryption: bool,
) -> Result<Self> {
let epoch_id = EpochId { hb_id, epoch };
let cs = Subset::new(netinfo.clone(), epoch_id).map_err(Error::CreateSubset)?;
Ok(EpochState {
epoch,
netinfo,
subset: SubsetState::Ongoing(cs),
decryption: BTreeMap::default(),
accepted_proposers: Default::default(),
subset_handler: subset_handling_strategy.into(),
require_decryption,
_phantom: PhantomData,
})
}
/// If the instance hasn't terminated yet, inputs our encrypted contribution.
pub fn propose<R: Rng>(&mut self, proposal: &C, rng: &mut R) -> Result<Step<C, N>> {
let ser_prop = bincode::serialize(&proposal).map_err(|err| Error::ProposeBincode(*err))?;
let cs_step = self.subset.handle_input(if self.require_decryption {
let ciphertext = self
.netinfo
.public_key_set()
.public_key()
.encrypt_with_rng(rng, ser_prop);
bincode::serialize(&ciphertext).map_err(|err| Error::ProposeBincode(*err))?
} else {
ser_prop
})?;
self.process_subset(cs_step)
}
/// Returns the number of contributions that we have already received or, after completion, how
/// many have been accepted.
pub fn received_proposals(&self) -> usize {
self.subset.received_proposals()
}
/// Handles a message for the Subset or a Threshold Decrypt instance.
pub fn handle_message_content(
&mut self,
sender_id: &N,
content: MessageContent<N>,
) -> Result<Step<C, N>> {
match content {
MessageContent::Subset(cs_msg) => {
let cs_step = self.subset.handle_message(sender_id, cs_msg)?;
self.process_subset(cs_step)
}
MessageContent::DecryptionShare { proposer_id, share } => {
if let Some(ref ids) = self.subset.accepted_ids() {
if !ids.contains(&proposer_id) {
let fault_kind = FaultKind::UnexpectedDecryptionShare;
return Ok(Fault::new(sender_id.clone(), fault_kind).into());
}
}
let td_step = match self.decryption.entry(proposer_id.clone()) {
Entry::Occupied(entry) => entry.into_mut(),
Entry::Vacant(entry) => {
entry.insert(DecryptionState::new(self.netinfo.clone()))
}
}
.handle_message(sender_id, share)
.map_err(Error::ThresholdDecrypt)?;
self.process_decryption(proposer_id, td_step)
}
}
}
/// When contributions of transactions have been decrypted for all valid proposers in this
/// epoch, moves those contributions into a batch, outputs the batch and updates the epoch.
pub fn try_output_batch(&self) -> Option<(Batch<C, N>, FaultLog<N>)> {
let proposer_ids = self.subset.accepted_ids()?;
let mut plaintexts = Vec::new();
// Collect accepted plaintexts. Return if some are not decrypted yet.
for id in proposer_ids {
match self.decryption.get(id) {
None | Some(DecryptionState::Ongoing(_)) => return None,
Some(DecryptionState::Complete(ref pt)) => plaintexts.push((id.clone(), pt)),
}
}
let mut fault_log = FaultLog::default();
let mut batch = Batch {
epoch: self.epoch,
contributions: BTreeMap::new(),
};
// Deserialize the output. If it fails, the proposer of that item is faulty.
for (id, plaintext) in plaintexts {
match bincode::deserialize::<C>(plaintext) {
Ok(contrib) => {
batch.contributions.insert(id, contrib);
}
Err(_) => fault_log.append(id, FaultKind::BatchDeserializationFailed),
}
}
Some((batch, fault_log))
}
/// Checks whether the subset has output, and if it does, sends out our decryption shares.
fn process_subset(&mut self, cs_step: CsStep<N>) -> Result<Step<C, N>> {
let mut step = Step::default();
let cs_outputs = step.extend_with(cs_step, FaultKind::SubsetFault, |cs_msg| {
MessageContent::Subset(cs_msg).with_epoch(self.epoch)
});
let mut has_seen_done = false;
for cs_output in cs_outputs {
if has_seen_done {
// TODO: Is there any way we can statically guarantee that?
error!("`SubsetOutput::Done` was not the last `SubsetOutput`");
}
let SubsetHandleData {
contributions,
is_done,
} = self.subset_handler.handle(cs_output);
for (k, v) in contributions {
step.extend(if self.require_decryption {
self.send_decryption_share(k.clone(), &v)?
} else {
self.decryption
.insert(k.clone(), DecryptionState::Complete(v));
Step::default()
});
self.accepted_proposers.insert(k);
}
if is_done {
self.subset = SubsetState::Complete(self.accepted_proposers.clone());
let faulty_shares: Vec<_> = self
.decryption
.keys()
.filter(|id| !self.accepted_proposers.contains(id))
.cloned()
.collect();
for id in faulty_shares {
if let Some(DecryptionState::Ongoing(td)) = self.decryption.remove(&id) {
for id in td.sender_ids() {
let fault_kind = FaultKind::UnexpectedDecryptionShare;
step.fault_log.append(id.clone(), fault_kind);
}
}
}
has_seen_done = true;
}
}
Ok(step)
}
/// Processes a Threshold Decrypt step.
fn process_decryption(&mut self, proposer_id: N, td_step: td::Step<N>) -> Result<Step<C, N>> {
let mut step = Step::default();
let opt_output = step.extend_with(td_step, FaultKind::DecryptionFault, |share| {
MessageContent::DecryptionShare {
proposer_id: proposer_id.clone(),
share,
}
.with_epoch(self.epoch)
});
if let Some(output) = opt_output.into_iter().next() {
self.decryption
.insert(proposer_id, DecryptionState::Complete(output));
}
Ok(step)
}
/// Given the output of the Subset algorithm, inputs the ciphertexts into the Threshold
/// Decrypt instances and sends our own decryption shares.
fn send_decryption_share(&mut self, proposer_id: N, v: &[u8]) -> Result<Step<C, N>> {
let ciphertext: Ciphertext = match bincode::deserialize(v) {
Ok(ciphertext) => ciphertext,
Err(_) => {
return Ok(Fault::new(proposer_id, FaultKind::DeserializeCiphertext).into());
}
};
let td_result = match self.decryption.entry(proposer_id.clone()) {
Entry::Occupied(entry) => entry.into_mut(),
Entry::Vacant(entry) => entry.insert(DecryptionState::new(self.netinfo.clone())),
}
.set_ciphertext(ciphertext);
match td_result {
Ok(td_step) => self.process_decryption(proposer_id, td_step),
Err(td::Error::InvalidCiphertext(_)) => {
Ok(Fault::new(proposer_id, FaultKind::InvalidCiphertext).into())
}
Err(err) => Err(Error::ThresholdDecrypt(err)),
}
}
}
/// A session identifier for a `Subset` sub-algorithm run within an epoch. It consists of the epoch
/// number, and an optional `HoneyBadger` session identifier.
#[derive(Clone, Debug, Serialize)]
struct EpochId {
hb_id: u64,
epoch: u64,
}
impl Display for EpochId {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> result::Result<(), fmt::Error> {
write!(f, "{}/{}", self.hb_id, self.epoch)
}
}