1 #![allow(missing_docs)]
2 // Copyright 2023 Google LLC
3 //
4 // Licensed under the Apache License, Version 2.0 (the "License");
5 // you may not use this file except in compliance with the License.
6 // You may obtain a copy of the License at
7 //
8 //     http://www.apache.org/licenses/LICENSE-2.0
9 //
10 // Unless required by applicable law or agreed to in writing, software
11 // distributed under the License is distributed on an "AS IS" BASIS,
12 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 // See the License for the specific language governing permissions and
14 // limitations under the License.
15 #![allow(clippy::expect_used)]
16 // TODO: remove this and convert all unwraps to expects
17 #![allow(clippy::unwrap_used)]
18 
19 use std::{
20     collections::HashSet,
21     fmt::{self, Formatter},
22     marker::PhantomData,
23 };
24 
25 use crypto_provider::elliptic_curve::EphemeralSecret;
26 use crypto_provider::p256::{P256EcdhProvider, P256PublicKey, P256};
27 use crypto_provider::x25519::X25519;
28 use crypto_provider::CryptoProvider;
29 use crypto_provider::{
30     elliptic_curve::{EcdhProvider, PublicKey},
31     hkdf::Hkdf,
32     sha2::{Sha256, Sha512},
33     CryptoRng,
34 };
35 use ukey2_proto::protobuf::Message;
36 use ukey2_proto::ukey2_all_proto::{securemessage, ukey};
37 
38 use crate::proto_adapter::NextProtocol;
39 pub(crate) use crate::proto_adapter::{
40     CipherCommitment, ClientFinished, ClientInit, GenericPublicKey, HandshakeCipher,
41     IntoAdapter as _, ServerInit, ToWrappedMessage as _,
42 };
43 
44 pub trait WireCompatibilityLayer {
encode_public_key<C: CryptoProvider>( &self, key: Vec<u8>, cipher: HandshakeCipher, ) -> Option<Vec<u8>>45     fn encode_public_key<C: CryptoProvider>(
46         &self,
47         key: Vec<u8>,
48         cipher: HandshakeCipher,
49     ) -> Option<Vec<u8>>;
decode_public_key<C: CryptoProvider>( &self, key: Vec<u8>, cipher: HandshakeCipher, ) -> Option<Vec<u8>>50     fn decode_public_key<C: CryptoProvider>(
51         &self,
52         key: Vec<u8>,
53         cipher: HandshakeCipher,
54     ) -> Option<Vec<u8>>;
55 }
56 
57 #[derive(Clone)]
58 pub enum HandshakeImplementation {
59     /// Implementation of ukey2 exchange handshake according to the specs in
60     /// <https://github.com/google/ukey2/blob/master/README.md>.
61     ///
62     /// In particular, when encoding for the P256 public key, this uses the standardized encoding
63     /// described in [SEC 1](https://www.secg.org/sec1-v2.pdf).
64     ///
65     /// For X25519, the public key is the x-coordinate in little endian per RFC 7748.
66     Spec,
67     /// Implementation of ukey2 exchange handshake that matches
68     /// [the Java implementation](https://github.com/google/ukey2/blob/master/src/main/java/com/google/security/cryptauth/lib/securegcm/Ukey2Handshake.java),
69     /// but different from what the specs says.
70     ///
71     /// In particular, when encoding for the P256 curve, the public key is represented as serialized
72     /// bytes of the following proto:
73     /// ```text
74     /// message EcP256PublicKey {
75     ///     // x and y are encoded in big-endian two's complement (slightly wasteful)
76     ///     // Client MUST verify (x,y) is a valid point on NIST P256
77     ///     required bytes x = 1;
78     ///     required bytes y = 2;
79     /// }
80     /// ```
81     ///
82     /// Encoding for X25519 is not supported in this mode.
83     PublicKeyInProtobuf,
84 }
85 
86 impl WireCompatibilityLayer for HandshakeImplementation {
encode_public_key<C: CryptoProvider>( &self, key: Vec<u8>, cipher: HandshakeCipher, ) -> Option<Vec<u8>>87     fn encode_public_key<C: CryptoProvider>(
88         &self,
89         key: Vec<u8>,
90         cipher: HandshakeCipher,
91     ) -> Option<Vec<u8>> {
92         match self {
93             HandshakeImplementation::Spec => Some(key),
94             HandshakeImplementation::PublicKeyInProtobuf => match cipher {
95                 HandshakeCipher::P256Sha512 => {
96                     let p256_key =
97                         <C::P256 as P256EcdhProvider>::PublicKey::from_bytes(key.as_slice())
98                             .expect("");
99                     let (x, y) = p256_key.to_affine_coordinates().unwrap();
100                     let bigboi_x = num_bigint::BigInt::from_biguint(
101                         num_bigint::Sign::Plus,
102                         num_bigint::BigUint::from_bytes_be(x.to_vec().as_slice()),
103                     );
104                     let bigboi_y = num_bigint::BigInt::from_biguint(
105                         num_bigint::Sign::Plus,
106                         num_bigint::BigUint::from_bytes_be(y.to_vec().as_slice()),
107                     );
108                     let proto_key = securemessage::EcP256PublicKey {
109                         x: Some(bigboi_x.to_signed_bytes_be()),
110                         y: Some(bigboi_y.to_signed_bytes_be()),
111                         ..Default::default()
112                     };
113                     let key = securemessage::GenericPublicKey {
114                         type_: Some(securemessage::PublicKeyType::EC_P256.into()),
115                         ec_p256_public_key: Some(proto_key).into(),
116                         ..Default::default()
117                     };
118                     key.write_to_bytes().ok()
119                 }
120                 HandshakeCipher::Curve25519Sha512 => None,
121             },
122         }
123     }
124 
decode_public_key<C: CryptoProvider>( &self, key: Vec<u8>, cipher: HandshakeCipher, ) -> Option<Vec<u8>>125     fn decode_public_key<C: CryptoProvider>(
126         &self,
127         key: Vec<u8>,
128         cipher: HandshakeCipher,
129     ) -> Option<Vec<u8>> {
130         match self {
131             HandshakeImplementation::Spec => Some(key),
132             HandshakeImplementation::PublicKeyInProtobuf => {
133                 // key will be wrapped in a genericpublickey
134                 let public_key: GenericPublicKey<C> =
135                     securemessage::GenericPublicKey::parse_from_bytes(key.as_slice())
136                         .ok()?
137                         .into_adapter()
138                         .ok()?;
139                 match public_key {
140                     GenericPublicKey::Ec256(key) => {
141                         debug_assert_eq!(cipher, HandshakeCipher::P256Sha512);
142                         Some(key.to_bytes().to_vec())
143                     }
144                 }
145             }
146         }
147     }
148 }
149 
150 pub struct Ukey2ServerStage1<C: CryptoProvider> {
151     pub(crate) next_protocols: HashSet<NextProtocol>,
152     pub(crate) handshake_impl: HandshakeImplementation,
153     _marker: PhantomData<C>,
154 }
155 
156 impl<C: CryptoProvider> fmt::Debug for Ukey2ServerStage1<C> {
fmt(&self, f: &mut Formatter<'_>) -> fmt::Result157     fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
158         write!(f, "Ukey2ServerS1")
159     }
160 }
161 
162 impl<C: CryptoProvider> Ukey2ServerStage1<C> {
from(next_protocols: HashSet<String>, handshake_impl: HandshakeImplementation) -> Self163     pub fn from(next_protocols: HashSet<String>, handshake_impl: HandshakeImplementation) -> Self {
164         Self {
165             next_protocols: next_protocols.iter().filter_map(|p| p.try_into().ok()).collect(),
166             handshake_impl,
167             _marker: PhantomData,
168         }
169     }
170 
handle_client_init<R: rand::Rng + rand::CryptoRng>( self, rng: &mut R, client_init: ClientInit, client_init_msg_bytes: Vec<u8>, ) -> Result<Ukey2ServerStage2<C>, ClientInitError>171     pub(crate) fn handle_client_init<R: rand::Rng + rand::CryptoRng>(
172         self,
173         rng: &mut R,
174         client_init: ClientInit,
175         client_init_msg_bytes: Vec<u8>,
176     ) -> Result<Ukey2ServerStage2<C>, ClientInitError> {
177         if client_init.version() != 1 {
178             return Err(ClientInitError::BadVersion);
179         }
180 
181         let next_protocols = client_init.next_protocols();
182         let Some(selected_protocol) = next_protocols.intersection(&self.next_protocols).min()
183         else {
184             return Err(ClientInitError::BadNextProtocol);
185         };
186 
187         // nothing to check here about client_init.random -- already been validated as 32 bytes
188 
189         // all cipher types are supported, so no BAD_HANDSHAKE_CIPHER case
190         let commitment = client_init
191             .commitments()
192             .iter()
193             // we want to get the first matching cipher, but max_by_key returns the last max,
194             // so iterate in reverse direction
195             .rev()
196             // proto enum uses the priority as the numeric value
197             .max_by_key(|c| c.cipher().as_proto() as i32)
198             .ok_or(ClientInitError::BadHandshakeCipher)?;
199         match commitment.cipher() {
200             // pick in priority order
201             HandshakeCipher::Curve25519Sha512 => {
202                 let secret = ServerKeyPair::Curve25519(
203                     <C::X25519 as EcdhProvider<X25519>>::EphemeralSecret::generate_random(&mut
204                         <<<C::X25519 as EcdhProvider<X25519>>::EphemeralSecret as EphemeralSecret<
205                             X25519,
206                         >>::Rng as CryptoRng>::new(),
207                     ),
208                 );
209                 Ok(Ukey2ServerStage2::from(
210                     &mut *rng,
211                     client_init_msg_bytes,
212                     commitment.clone(),
213                     secret,
214                     self.handshake_impl,
215                     *selected_protocol,
216                 ))
217             }
218             HandshakeCipher::P256Sha512 => {
219                 let secret = ServerKeyPair::P256(
220                     <C::P256 as EcdhProvider<P256>>::EphemeralSecret::generate_random(
221                         &mut<<<C::P256 as EcdhProvider<P256>>::EphemeralSecret as EphemeralSecret<
222                             P256,
223                         >>::Rng as CryptoRng>::new(),
224                     ),
225                 );
226                 Ok(Ukey2ServerStage2::from(
227                     &mut *rng,
228                     client_init_msg_bytes,
229                     commitment.clone(),
230                     secret,
231                     self.handshake_impl,
232                     *selected_protocol,
233                 ))
234             }
235         }
236     }
237 }
238 
239 enum ServerKeyPair<C: CryptoProvider> {
240     Curve25519(<C::X25519 as EcdhProvider<X25519>>::EphemeralSecret),
241     P256(<C::P256 as EcdhProvider<P256>>::EphemeralSecret),
242 }
243 
244 pub struct Ukey2ServerStage2<C: CryptoProvider> {
245     client_init_msg: Vec<u8>,
246     server_init_msg: Vec<u8>,
247     commitment: CipherCommitment,
248     key_pair: ServerKeyPair<C>,
249     pub(crate) handshake_impl: HandshakeImplementation,
250     next_protocol: NextProtocol,
251     _marker: PhantomData<C>,
252 }
253 
254 impl<C: CryptoProvider> fmt::Debug for Ukey2ServerStage2<C> {
fmt(&self, f: &mut Formatter<'_>) -> fmt::Result255     fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
256         write!(f, "Ukey2ServerS2")
257     }
258 }
259 
260 const HKDF_SALT_AUTH: &[u8] = b"UKEY2 v1 auth";
261 const HKDF_SALT_NEXT: &[u8] = b"UKEY2 v1 next";
262 
263 impl<C: CryptoProvider> Ukey2ServerStage2<C> {
from<R: rand::Rng + rand::CryptoRng>( rng: &mut R, client_init_msg: Vec<u8>, commitment: CipherCommitment, key_pair: ServerKeyPair<C>, handshake_impl: HandshakeImplementation, next_protocol: NextProtocol, ) -> Self264     fn from<R: rand::Rng + rand::CryptoRng>(
265         rng: &mut R,
266         client_init_msg: Vec<u8>,
267         commitment: CipherCommitment,
268         key_pair: ServerKeyPair<C>,
269         handshake_impl: HandshakeImplementation,
270         next_protocol: NextProtocol,
271     ) -> Self {
272         let random: [u8; 32] = rng.gen();
273         let mut server_init = ukey::Ukey2ServerInit::default();
274         server_init.set_version(1);
275         server_init.set_random(random.to_vec());
276         server_init.set_handshake_cipher(commitment.cipher().as_proto());
277         server_init.set_public_key(match &key_pair {
278             ServerKeyPair::Curve25519(es) => es.public_key_bytes().as_ref().to_vec(),
279             ServerKeyPair::P256(es) => handshake_impl
280                 .encode_public_key::<C>(
281                     es.public_key_bytes().as_ref().to_vec(),
282                     HandshakeCipher::P256Sha512,
283                 )
284                 .unwrap(),
285         });
286         server_init.set_selected_next_protocol(next_protocol.to_string());
287 
288         Self {
289             client_init_msg,
290             server_init_msg: server_init.to_wrapped_msg().write_to_bytes().unwrap(),
291             commitment,
292             key_pair,
293             handshake_impl,
294             next_protocol,
295             _marker: PhantomData,
296         }
297     }
298 
server_init_msg(&self) -> &[u8]299     pub fn server_init_msg(&self) -> &[u8] {
300         &self.server_init_msg
301     }
302 
handle_client_finished_msg( self, msg: ClientFinished, client_finished_msg_bytes: &[u8], ) -> Result<Ukey2Server, ClientFinishedError>303     pub(crate) fn handle_client_finished_msg(
304         self,
305         msg: ClientFinished,
306         client_finished_msg_bytes: &[u8],
307     ) -> Result<Ukey2Server, ClientFinishedError> {
308         let hash_bytes = C::Sha512::sha512(client_finished_msg_bytes);
309         // must be constant time to avoid timing attack on hash equality
310         if C::constant_time_eq(hash_bytes.as_slice(), self.commitment.commitment()) {
311             // handshake is complete
312             // independently derive shared DH key
313             let shared_secret_bytes = match self.key_pair {
314                 ServerKeyPair::Curve25519(es) => {
315                     let buf = msg.public_key.into_iter().collect::<Vec<u8>>();
316                     let public_key: [u8; 32] =
317                         (&buf[..]).try_into().map_err(|_| ClientFinishedError::BadEd25519Key)?;
318                     es.diffie_hellman(
319                         &<C::X25519 as EcdhProvider<X25519>>::PublicKey::from_bytes(&public_key)
320                             .map_err(|_| ClientFinishedError::BadEd25519Key)?,
321                     )
322                     .map_err(|_| ClientFinishedError::BadKeyExchange)?
323                     .into()
324                 }
325                 ServerKeyPair::P256(es) => {
326                     let other_public_key =
327                         &<C::P256 as P256EcdhProvider>::PublicKey::from_sec1_bytes(
328                             self.handshake_impl
329                                 .decode_public_key::<C>(msg.public_key, HandshakeCipher::P256Sha512)
330                                 .ok_or(ClientFinishedError::BadP256Key)?
331                                 .as_slice(),
332                         )
333                         .map_err(|_| ClientFinishedError::BadP256Key)?;
334                     es.diffie_hellman(other_public_key)
335                         .map_err(|_| ClientFinishedError::BadKeyExchange)?
336                         .into()
337                 }
338             };
339             let shared_secret_sha256 = C::Sha256::sha256(&shared_secret_bytes).to_vec();
340             Ok(Ukey2Server {
341                 completed_handshake: CompletedHandshake::new(
342                     self.client_init_msg,
343                     self.server_init_msg,
344                     shared_secret_sha256,
345                     self.next_protocol,
346                 ),
347             })
348         } else {
349             Err(ClientFinishedError::UnknownCommitment)
350         }
351     }
352 }
353 
354 /// Representation of the UKEY2 server information after the handshake has been completed. An
355 /// instance of this can be created by going through the handshake state machine (starting from
356 /// [`Ukey2ServerStage1`]).
357 pub struct Ukey2Server {
358     completed_handshake: CompletedHandshake,
359 }
360 
361 impl fmt::Debug for Ukey2Server {
fmt(&self, f: &mut Formatter<'_>) -> fmt::Result362     fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
363         write!(f, "Ukey2Server")
364     }
365 }
366 
367 impl Ukey2Server {
completed_handshake(&self) -> &CompletedHandshake368     pub fn completed_handshake(&self) -> &CompletedHandshake {
369         &self.completed_handshake
370     }
371 }
372 
373 pub struct Ukey2ClientStage1<C: CryptoProvider> {
374     curve25519_secret: <C::X25519 as EcdhProvider<X25519>>::EphemeralSecret,
375     p256_secret: <C::P256 as EcdhProvider<P256>>::EphemeralSecret,
376     curve25519_client_finished_bytes: Vec<u8>,
377     p256_client_finished_bytes: Vec<u8>,
378     client_init_bytes: Vec<u8>,
379     commitment_ciphers: Vec<HandshakeCipher>,
380     handshake_impl: HandshakeImplementation,
381     next_protocols: Vec<NextProtocol>,
382     _marker: PhantomData<C>,
383 }
384 
385 impl<C: CryptoProvider> fmt::Debug for Ukey2ClientStage1<C> {
fmt(&self, f: &mut Formatter<'_>) -> fmt::Result386     fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
387         write!(f, "Ukey2Client1")
388     }
389 }
390 
391 impl<C: CryptoProvider> Ukey2ClientStage1<C> {
392     // Clippy: we assert that there must be at least one element in `next_protocols`, so indexing
393     // [0] is safe.
394     #[allow(clippy::indexing_slicing)]
from<R: rand::Rng + rand::SeedableRng + rand::CryptoRng>( rng: &mut R, next_protocols: Vec<NextProtocol>, handshake_impl: HandshakeImplementation, ) -> Self395     pub fn from<R: rand::Rng + rand::SeedableRng + rand::CryptoRng>(
396         rng: &mut R,
397         next_protocols: Vec<NextProtocol>,
398         handshake_impl: HandshakeImplementation,
399     ) -> Self {
400         assert!(!next_protocols.is_empty());
401         let random = rng.gen::<[u8; 32]>().to_vec();
402         // Curve25519 ClientFinished Message
403         let curve25519_secret =
404             <C::X25519 as EcdhProvider<X25519>>::EphemeralSecret::generate_random(
405                 &mut <<<C::X25519 as EcdhProvider<X25519>>::EphemeralSecret as EphemeralSecret<
406                     X25519,
407                 >>::Rng as CryptoRng>::new(),
408             );
409         let curve25519_client_finished_bytes = {
410             let client_finished = ukey::Ukey2ClientFinished {
411                 public_key: Some(curve25519_secret.public_key_bytes().as_ref().to_vec()),
412                 ..Default::default()
413             };
414             client_finished.to_wrapped_msg().write_to_bytes().unwrap()
415         };
416         let curve25519_client_finished_hash =
417             C::Sha512::sha512(&curve25519_client_finished_bytes).to_vec();
418 
419         // P256 ClientFinished Message
420         let p256_secret = <C::P256 as EcdhProvider<P256>>::EphemeralSecret::generate_random(
421             &mut <<<C::P256 as EcdhProvider<P256>>::EphemeralSecret as EphemeralSecret<
422                 P256,
423             >>::Rng as CryptoRng>::new(),
424         );
425         let p256_client_finished_bytes = {
426             let client_finished = ukey::Ukey2ClientFinished {
427                 public_key: Some(
428                     handshake_impl
429                         .encode_public_key::<C>(
430                             p256_secret.public_key_bytes().as_ref().to_vec(),
431                             HandshakeCipher::P256Sha512,
432                         )
433                         .expect("Output of p256_secret.public_key_bytes should always be valid input for encode_public_key"),
434                 ),
435                 ..Default::default()
436             };
437             client_finished.to_wrapped_msg().write_to_bytes().unwrap()
438         };
439         let p256_client_finished_hash = C::Sha512::sha512(&p256_client_finished_bytes).to_vec();
440 
441         // ClientInit Message
442         let client_init_bytes = {
443             let curve25519_commitment = ukey::ukey2client_init::CipherCommitment {
444                 handshake_cipher: Some(HandshakeCipher::Curve25519Sha512.as_proto().into()),
445                 commitment: Some(curve25519_client_finished_hash),
446                 ..Default::default()
447             };
448 
449             let p256_commitment = ukey::ukey2client_init::CipherCommitment {
450                 handshake_cipher: Some(HandshakeCipher::P256Sha512.as_proto().into()),
451                 commitment: Some(p256_client_finished_hash),
452                 ..Default::default()
453             };
454 
455             let client_init = ukey::Ukey2ClientInit {
456                 version: Some(1),
457                 random: Some(random),
458                 cipher_commitments: vec![curve25519_commitment, p256_commitment],
459                 next_protocol: Some(next_protocols[0].to_string()),
460                 next_protocols: next_protocols.iter().map(|x| x.to_string()).collect(),
461                 ..Default::default()
462             };
463             client_init.to_wrapped_msg().write_to_bytes().unwrap()
464         };
465 
466         Self {
467             curve25519_secret,
468             p256_secret,
469             curve25519_client_finished_bytes,
470             p256_client_finished_bytes,
471             client_init_bytes,
472             commitment_ciphers: vec![
473                 HandshakeCipher::Curve25519Sha512,
474                 HandshakeCipher::P256Sha512,
475             ],
476             handshake_impl,
477             next_protocols,
478             _marker: PhantomData,
479         }
480     }
481 
client_init_msg(&self) -> &[u8]482     pub fn client_init_msg(&self) -> &[u8] {
483         &self.client_init_bytes
484     }
485 
handle_server_init( self, server_init: ServerInit, server_init_bytes: Vec<u8>, ) -> Result<Ukey2Client, ServerInitError>486     pub(crate) fn handle_server_init(
487         self,
488         server_init: ServerInit,
489         server_init_bytes: Vec<u8>,
490     ) -> Result<Ukey2Client, ServerInitError> {
491         if server_init.version() != 1 {
492             return Err(ServerInitError::BadVersion);
493         }
494 
495         if !self.next_protocols.contains(&server_init.selected_next_protocol()) {
496             return Err(ServerInitError::BadNextProtocol);
497         }
498         let next_protocol = server_init.selected_next_protocol();
499 
500         // loop over all commitments every time for a semblance of constant time-ness
501         let server_cipher = self
502             .commitment_ciphers
503             .iter()
504             .fold(None, |accum, c| {
505                 if server_init.handshake_cipher() == *c {
506                     match accum {
507                         None => Some(c),
508                         Some(_) => accum,
509                     }
510                 } else {
511                     accum
512                 }
513             })
514             .ok_or(ServerInitError::BadHandshakeCipher)?;
515         let (server_shared_secret, client_finished_bytes) = match server_cipher {
516             HandshakeCipher::P256Sha512 => {
517                 let other_public_key = &<C::P256 as P256EcdhProvider>::PublicKey::from_sec1_bytes(
518                     self.handshake_impl
519                         .decode_public_key::<C>(
520                             server_init.public_key.to_vec(),
521                             HandshakeCipher::P256Sha512,
522                         )
523                         .ok_or(ServerInitError::BadPublicKey)?
524                         .as_slice(),
525                 )
526                 .map_err(|_| ServerInitError::BadPublicKey)?;
527                 let shared_secret = self
528                     .p256_secret
529                     .diffie_hellman(other_public_key)
530                     .map_err(|_| ServerInitError::BadKeyExchange)?;
531                 let shared_secret_bytes: [u8; 32] = shared_secret.into();
532                 (shared_secret_bytes, self.p256_client_finished_bytes)
533             }
534             HandshakeCipher::Curve25519Sha512 => {
535                 let pub_key: [u8; 32] =
536                     server_init.public_key.try_into().map_err(|_| ServerInitError::BadPublicKey)?;
537                 (
538                     self.curve25519_secret
539                         .diffie_hellman(
540                             &<C::X25519 as EcdhProvider<X25519>>::PublicKey::from_bytes(&pub_key)
541                                 .map_err(|_| ServerInitError::BadPublicKey)?,
542                         )
543                         .map_err(|_| ServerInitError::BadKeyExchange)?
544                         .into(),
545                     self.curve25519_client_finished_bytes,
546                 )
547             }
548         };
549         let shared_secret_bytes = C::Sha256::sha256(&server_shared_secret).to_vec();
550         Ok(Ukey2Client {
551             client_finished_bytes,
552             completed_handshake: CompletedHandshake::new(
553                 self.client_init_bytes,
554                 server_init_bytes.to_vec(),
555                 shared_secret_bytes,
556                 next_protocol,
557             ),
558         })
559     }
560 }
561 
562 #[derive(Debug)]
563 #[allow(clippy::enum_variant_names)]
564 pub(crate) enum ServerInitError {
565     BadVersion,
566     BadHandshakeCipher,
567     BadPublicKey,
568     /// The diffie-hellman key exchange failed to generate a shared secret
569     BadKeyExchange,
570     /// The server sent an invalid next protocol that is not available to the client.
571     BadNextProtocol,
572 }
573 
574 #[derive(Clone)]
575 pub struct Ukey2Client {
576     completed_handshake: CompletedHandshake,
577     client_finished_bytes: Vec<u8>,
578 }
579 
580 impl fmt::Debug for Ukey2Client {
fmt(&self, f: &mut Formatter<'_>) -> fmt::Result581     fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
582         write!(f, "Ukey2Client")
583     }
584 }
585 
586 impl Ukey2Client {
client_finished_msg(&self) -> &[u8]587     pub fn client_finished_msg(&self) -> &[u8] {
588         &self.client_finished_bytes
589     }
590 
completed_handshake(&self) -> &CompletedHandshake591     pub fn completed_handshake(&self) -> &CompletedHandshake {
592         &self.completed_handshake
593     }
594 }
595 
596 #[allow(clippy::enum_variant_names)]
597 pub enum ClientInitError {
598     BadVersion,
599     BadHandshakeCipher,
600     BadNextProtocol,
601 }
602 
603 pub enum ClientFinishedError {
604     BadEd25519Key,
605     BadP256Key,
606     UnknownCommitment,
607     /// The diffie-hellman key exchange failed to generate a shared secret
608     BadKeyExchange,
609 }
610 
611 /// The result of completing the UKEY2 handshake.
612 #[derive(Clone)]
613 pub struct CompletedHandshake {
614     client_init_bytes: Vec<u8>,
615     server_init_bytes: Vec<u8>,
616     shared_secret: Vec<u8>,
617     pub next_protocol: NextProtocol,
618 }
619 
620 impl CompletedHandshake {
new( client_init_bytes: Vec<u8>, server_init_bytes: Vec<u8>, shared_secret: Vec<u8>, next_protocol: NextProtocol, ) -> Self621     fn new(
622         client_init_bytes: Vec<u8>,
623         server_init_bytes: Vec<u8>,
624         shared_secret: Vec<u8>,
625         next_protocol: NextProtocol,
626     ) -> Self {
627         Self { client_init_bytes, server_init_bytes, shared_secret, next_protocol }
628     }
629 
630     /// Returns an HKDF for the UKEY2 `AUTH_STRING`.
auth_string<C: CryptoProvider>(&self) -> HandshakeHkdf<C>631     pub fn auth_string<C: CryptoProvider>(&self) -> HandshakeHkdf<C> {
632         HandshakeHkdf::new(
633             &self.client_init_bytes,
634             &self.server_init_bytes,
635             C::HkdfSha256::new(Some(HKDF_SALT_AUTH), &self.shared_secret),
636         )
637     }
638 
639     /// Returns an HKDF for the UKEY2 `NEXT_SECRET`.
next_protocol_secret<C: CryptoProvider>(&self) -> HandshakeHkdf<C>640     pub fn next_protocol_secret<C: CryptoProvider>(&self) -> HandshakeHkdf<C> {
641         HandshakeHkdf::new(
642             &self.client_init_bytes,
643             &self.server_init_bytes,
644             C::HkdfSha256::new(Some(HKDF_SALT_NEXT), &self.shared_secret),
645         )
646     }
647 }
648 
649 /// A UKEY2 handshake secret that can derive output at the caller's preferred length.
650 pub struct HandshakeHkdf<'a, C: CryptoProvider> {
651     client_init_bytes: &'a [u8],
652     server_init_bytes: &'a [u8],
653     hkdf: C::HkdfSha256,
654 }
655 
656 impl<'a, C: CryptoProvider> HandshakeHkdf<'a, C> {
657     /// Returns `None` if the requested size > 255 * 512 bytes.
derive_array<const N: usize>(&self) -> Option<[u8; N]>658     pub fn derive_array<const N: usize>(&self) -> Option<[u8; N]> {
659         let mut buf = [0; N];
660         self.derive_slice(&mut buf).map(|_| buf)
661     }
662 
663     /// Returns `None` if the requested `length` > 255 * 512 bytes.
derive_vec(&self, length: usize) -> Option<Vec<u8>>664     pub fn derive_vec(&self, length: usize) -> Option<Vec<u8>> {
665         let mut buf = vec![0; length];
666         self.derive_slice(&mut buf).map(|_| buf)
667     }
668 
669     /// Returns `None` if the provided `buf` has size > 255 * 512 bytes.
derive_slice(&self, buf: &mut [u8]) -> Option<()>670     pub fn derive_slice(&self, buf: &mut [u8]) -> Option<()> {
671         self.hkdf.expand_multi_info(&[self.client_init_bytes, self.server_init_bytes], buf).ok()
672     }
673 
new(client_init_bytes: &'a [u8], server_init_bytes: &'a [u8], hkdf: C::HkdfSha256) -> Self674     fn new(client_init_bytes: &'a [u8], server_init_bytes: &'a [u8], hkdf: C::HkdfSha256) -> Self {
675         Self { client_init_bytes, server_init_bytes, hkdf }
676     }
677 }
678