libsignal_protocol/state/
bundle.rs

1//
2// Copyright 2020 Signal Messenger, LLC.
3// SPDX-License-Identifier: AGPL-3.0-only
4//
5
6use std::clone::Clone;
7
8use crate::state::{PreKeyId, SignedPreKeyId};
9use crate::{kem, DeviceId, IdentityKey, KyberPreKeyId, PublicKey, Result, SignalProtocolError};
10
11#[derive(Clone)]
12struct SignedPreKey {
13    id: SignedPreKeyId,
14    public_key: PublicKey,
15    signature: Vec<u8>,
16}
17
18impl SignedPreKey {
19    fn new(id: SignedPreKeyId, public_key: PublicKey, signature: Vec<u8>) -> Self {
20        Self {
21            id,
22            public_key,
23            signature,
24        }
25    }
26}
27
28#[derive(Clone)]
29struct KyberPreKey {
30    id: KyberPreKeyId,
31    public_key: kem::PublicKey,
32    signature: Vec<u8>,
33}
34
35impl KyberPreKey {
36    fn new(id: KyberPreKeyId, public_key: kem::PublicKey, signature: Vec<u8>) -> Self {
37        Self {
38            id,
39            public_key,
40            signature,
41        }
42    }
43}
44
45// Represents the raw contents of the pre-key bundle without any notion of required/optional
46// fields.
47// Can be used as a "builder" for PreKeyBundle, in which case all the validation will happen in
48// PreKeyBundle::new.
49pub struct PreKeyBundleContent {
50    pub registration_id: Option<u32>,
51    pub device_id: Option<DeviceId>,
52    pub pre_key_id: Option<PreKeyId>,
53    pub pre_key_public: Option<PublicKey>,
54    pub ec_pre_key_id: Option<SignedPreKeyId>,
55    pub ec_pre_key_public: Option<PublicKey>,
56    pub ec_pre_key_signature: Option<Vec<u8>>,
57    pub identity_key: Option<IdentityKey>,
58    pub kyber_pre_key_id: Option<KyberPreKeyId>,
59    pub kyber_pre_key_public: Option<kem::PublicKey>,
60    pub kyber_pre_key_signature: Option<Vec<u8>>,
61}
62
63impl From<PreKeyBundle> for PreKeyBundleContent {
64    fn from(bundle: PreKeyBundle) -> Self {
65        Self {
66            registration_id: Some(bundle.registration_id),
67            device_id: Some(bundle.device_id),
68            pre_key_id: bundle.pre_key_id,
69            pre_key_public: bundle.pre_key_public,
70            ec_pre_key_id: Some(bundle.ec_signed_pre_key.id),
71            ec_pre_key_public: Some(bundle.ec_signed_pre_key.public_key),
72            ec_pre_key_signature: Some(bundle.ec_signed_pre_key.signature),
73            identity_key: Some(bundle.identity_key),
74            kyber_pre_key_id: bundle.kyber_pre_key.as_ref().map(|kyber| kyber.id),
75            kyber_pre_key_public: bundle
76                .kyber_pre_key
77                .as_ref()
78                .map(|kyber| kyber.public_key.clone()),
79            kyber_pre_key_signature: bundle
80                .kyber_pre_key
81                .as_ref()
82                .map(|kyber| kyber.signature.clone()),
83        }
84    }
85}
86
87impl TryFrom<PreKeyBundleContent> for PreKeyBundle {
88    type Error = SignalProtocolError;
89
90    fn try_from(content: PreKeyBundleContent) -> Result<Self> {
91        let mut bundle = PreKeyBundle::new(
92            content.registration_id.ok_or_else(|| {
93                SignalProtocolError::InvalidArgument("registration_id is required".to_string())
94            })?,
95            content.device_id.ok_or_else(|| {
96                SignalProtocolError::InvalidArgument("device_id is required".to_string())
97            })?,
98            content
99                .pre_key_id
100                .and_then(|id| content.pre_key_public.map(|public| (id, public))),
101            content.ec_pre_key_id.ok_or_else(|| {
102                SignalProtocolError::InvalidArgument("signed_pre_key_id is required".to_string())
103            })?,
104            content.ec_pre_key_public.ok_or_else(|| {
105                SignalProtocolError::InvalidArgument(
106                    "signed_pre_key_public is required".to_string(),
107                )
108            })?,
109            content.ec_pre_key_signature.ok_or_else(|| {
110                SignalProtocolError::InvalidArgument(
111                    "signed_pre_key_signature is required".to_string(),
112                )
113            })?,
114            content.identity_key.ok_or_else(|| {
115                SignalProtocolError::InvalidArgument("identity_key is required".to_string())
116            })?,
117        )?;
118
119        fn zip3<T, U, V>(x: Option<T>, y: Option<U>, z: Option<V>) -> Option<(T, U, V)> {
120            x.zip(y).zip(z).map(|((x, y), z)| (x, y, z))
121        }
122
123        if let Some((kyber_id, kyber_public, kyber_sig)) = zip3(
124            content.kyber_pre_key_id,
125            content.kyber_pre_key_public,
126            content.kyber_pre_key_signature,
127        ) {
128            bundle = bundle.with_kyber_pre_key(kyber_id, kyber_public, kyber_sig);
129        }
130        Ok(bundle)
131    }
132}
133
134#[derive(Clone)]
135pub struct PreKeyBundle {
136    registration_id: u32,
137    device_id: DeviceId,
138    pre_key_id: Option<PreKeyId>,
139    pre_key_public: Option<PublicKey>,
140    ec_signed_pre_key: SignedPreKey,
141    identity_key: IdentityKey,
142    // Optional to support older clients
143    // TODO: remove optionality once the transition is over
144    kyber_pre_key: Option<KyberPreKey>,
145}
146
147impl PreKeyBundle {
148    pub fn new(
149        registration_id: u32,
150        device_id: DeviceId,
151        pre_key: Option<(PreKeyId, PublicKey)>,
152        signed_pre_key_id: SignedPreKeyId,
153        signed_pre_key_public: PublicKey,
154        signed_pre_key_signature: Vec<u8>,
155        identity_key: IdentityKey,
156    ) -> Result<Self> {
157        let (pre_key_id, pre_key_public) = match pre_key {
158            None => (None, None),
159            Some((id, key)) => (Some(id), Some(key)),
160        };
161
162        let ec_signed_pre_key = SignedPreKey::new(
163            signed_pre_key_id,
164            signed_pre_key_public,
165            signed_pre_key_signature,
166        );
167
168        Ok(Self {
169            registration_id,
170            device_id,
171            pre_key_id,
172            pre_key_public,
173            ec_signed_pre_key,
174            identity_key,
175            kyber_pre_key: None,
176        })
177    }
178
179    pub fn with_kyber_pre_key(
180        mut self,
181        pre_key_id: KyberPreKeyId,
182        public_key: kem::PublicKey,
183        signature: Vec<u8>,
184    ) -> Self {
185        self.kyber_pre_key = Some(KyberPreKey::new(pre_key_id, public_key, signature));
186        self
187    }
188
189    pub fn registration_id(&self) -> Result<u32> {
190        Ok(self.registration_id)
191    }
192
193    pub fn device_id(&self) -> Result<DeviceId> {
194        Ok(self.device_id)
195    }
196
197    pub fn pre_key_id(&self) -> Result<Option<PreKeyId>> {
198        Ok(self.pre_key_id)
199    }
200
201    pub fn pre_key_public(&self) -> Result<Option<PublicKey>> {
202        Ok(self.pre_key_public)
203    }
204
205    pub fn signed_pre_key_id(&self) -> Result<SignedPreKeyId> {
206        Ok(self.ec_signed_pre_key.id)
207    }
208
209    pub fn signed_pre_key_public(&self) -> Result<PublicKey> {
210        Ok(self.ec_signed_pre_key.public_key)
211    }
212
213    pub fn signed_pre_key_signature(&self) -> Result<&[u8]> {
214        Ok(self.ec_signed_pre_key.signature.as_ref())
215    }
216
217    pub fn identity_key(&self) -> Result<&IdentityKey> {
218        Ok(&self.identity_key)
219    }
220
221    pub fn has_kyber_pre_key(&self) -> bool {
222        self.kyber_pre_key.is_some()
223    }
224
225    pub fn kyber_pre_key_id(&self) -> Result<Option<KyberPreKeyId>> {
226        Ok(self.kyber_pre_key.as_ref().map(|pre_key| pre_key.id))
227    }
228
229    pub fn kyber_pre_key_public(&self) -> Result<Option<&kem::PublicKey>> {
230        Ok(self
231            .kyber_pre_key
232            .as_ref()
233            .map(|pre_key| &pre_key.public_key))
234    }
235
236    pub fn kyber_pre_key_signature(&self) -> Result<Option<&[u8]>> {
237        Ok(self
238            .kyber_pre_key
239            .as_ref()
240            .map(|pre_key| pre_key.signature.as_ref()))
241    }
242
243    pub fn modify<F>(self, modify: F) -> Result<Self>
244    where
245        F: FnOnce(&mut PreKeyBundleContent),
246    {
247        let mut content = self.into();
248        modify(&mut content);
249        content.try_into()
250    }
251}
libsignal_protocol/state/bundle.rs

libsignal_protocol/state/
bundle.rs
