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

use alloc::boxed::Box;
use alloc::collections::BTreeSet;
use alloc::string::String;

use serde::{Deserialize, Serialize};
use umbral_pre::Capsule;

use crate::address::Address;
use crate::conditions::Conditions;
use crate::message_kit::MessageKit;
use crate::versioning::{
    messagepack_deserialize, messagepack_serialize, ProtocolObject, ProtocolObjectInner,
};

/// An object encapsulating the information necessary for retrieval of cfrags from Ursulas.
/// Contains the capsule and the checksum addresses of Ursulas from which the requester
/// already received cfrags.
#[derive(PartialEq, Debug, Serialize, Deserialize)]
pub struct RetrievalKit {
    /// The ciphertext's capsule.
    pub capsule: Capsule,
    /// The addresses that have already been queried for reencryption.
    pub queried_addresses: BTreeSet<Address>,
    /// A blob of bytes containing decryption conditions for this message.
    pub conditions: Option<Conditions>,
}

impl RetrievalKit {
    /// Creates a new retrival kit from a message kit.
    pub fn from_message_kit(message_kit: &MessageKit) -> Self {
        Self {
            capsule: message_kit.capsule,
            queried_addresses: BTreeSet::<Address>::new(),
            conditions: message_kit.conditions.clone(),
        }
    }

    /// Creates a new retrieval kit recording the addresses already queried for reencryption.
    pub fn new(
        capsule: &Capsule,
        queried_addresses: impl IntoIterator<Item = Address>,
        conditions: Option<&Conditions>,
    ) -> Self {
        // Can store cfrags too, if we're worried about Ursulas supplying duplicate ones.
        Self {
            capsule: *capsule,
            queried_addresses: queried_addresses.into_iter().collect(),
            conditions: conditions.cloned(),
        }
    }
}

#[derive(Deserialize)]
struct RetrievalKit0 {
    capsule: Capsule,
    queried_addresses: BTreeSet<Address>,
}

impl<'a> ProtocolObjectInner<'a> for RetrievalKit {
    fn brand() -> [u8; 4] {
        *b"RKit"
    }

    fn version() -> (u16, u16) {
        (1, 1)
    }

    fn unversioned_to_bytes(&self) -> Box<[u8]> {
        messagepack_serialize(&self)
    }

    fn unversioned_from_bytes(minor_version: u16, bytes: &[u8]) -> Option<Result<Self, String>> {
        match minor_version {
            0 => {
                let rkit =
                    messagepack_deserialize::<RetrievalKit0>(bytes).map(|rkit| RetrievalKit {
                        capsule: rkit.capsule,
                        queried_addresses: rkit.queried_addresses,
                        conditions: None,
                    });
                Some(rkit)
            }
            1 => Some(messagepack_deserialize(bytes)),
            _ => None,
        }
    }
}

impl<'a> ProtocolObject<'a> for RetrievalKit {}