Updating opaque interfaces to include ids from the internet draft (#56)

src/opaque.rs

@@ -7,19 +7,25 @@
 
 use crate::{
     ciphersuite::CipherSuite,
-    envelope::{Envelope, ExportKeySize},
-    errors::{utils::check_slice_size, InternalPakeError, PakeError, ProtocolError},
+    envelope::{Envelope, EnvelopeCredentialsFormat, ExportKeySize},
+    errors::{
+        utils::{check_slice_size, check_slice_size_atleast},
+        InternalPakeError, PakeError, ProtocolError,
+    },
     group::Group,
     hash::Hash,
     key_exchange::traits::{KeyExchange, ToBytes},
     keypair::{KeyPair, SizedBytes},
     oprf,
     oprf::OprfClientBytes,
-    serialization::{serialize, tokenize, CredentialType, ProtocolMessageType},
+    serialization::{
+        serialize, tokenize, u8_to_credential_type, CredentialType, ProtocolMessageType,
+    },
     slow_hash::SlowHash,
 };
 use generic_array::{typenum::Unsigned, GenericArray};
 use rand_core::{CryptoRng, RngCore};
+use std::collections::HashMap;
 use std::{convert::TryFrom, marker::PhantomData};
 use zeroize::Zeroize;
 
@@ -28,6 +34,8 @@ use zeroize::Zeroize;
 
 /// The message sent by the client to the server, to initiate registration
 pub struct RegisterFirstMessage<Grp> {
+    /// User identity
+    id_u: Vec<u8>,
     /// blinded password information
     alpha: Grp,
 }
@@ -35,42 +43,44 @@ pub struct RegisterFirstMessage<Grp> {
 impl<Grp: Group> TryFrom<&[u8]> for RegisterFirstMessage<Grp> {
     type Error = ProtocolError;
     fn try_from(first_message_bytes: &[u8]) -> Result<Self, Self::Error> {
-        let checked_slice = check_slice_size(
-            first_message_bytes,
-            Grp::ElemLen::to_usize(),
-            "first_message_bytes",
-        )?;
+        let elem_len = Grp::ElemLen::to_usize();
+        let checked_slice =
+            check_slice_size_atleast(first_message_bytes, elem_len, "first_message_bytes")?;
+
+        let id_u = checked_slice[..checked_slice.len() - elem_len].to_vec();
+
         // Check that the message is actually containing an element of the
         // correct subgroup
-        let arr = GenericArray::from_slice(checked_slice);
+        let arr = GenericArray::from_slice(&checked_slice[checked_slice.len() - elem_len..]);
         let alpha = Grp::from_element_slice(arr)?;
-        Ok(Self { alpha })
+        Ok(Self { id_u, alpha })
     }
 }
 
 impl<Grp: Group> RegisterFirstMessage<Grp> {
-    /// byte representation for the registration request
-    fn to_bytes(&self) -> GenericArray<u8, Grp::ElemLen> {
-        self.alpha.to_arr()
+    /// Byte representation for the registration request
+    pub fn to_bytes(&self) -> Vec<u8> {
+        [&self.id_u[..], &self.alpha.to_arr().to_vec()[..]].concat()
     }
-}
 
-impl<Grp: Group> RegisterFirstMessage<Grp> {
     /// Serialization into bytes
     pub fn serialize(&self) -> Vec<u8> {
         let mut registration_request: Vec<u8> = Vec::new();
-        registration_request.extend_from_slice(&serialize(Vec::new(), 2));
-        registration_request.extend_from_slice(&serialize((&self.to_bytes()).to_vec(), 2));
+        registration_request.extend_from_slice(&serialize(&self.id_u, 2));
+        registration_request.extend_from_slice(&serialize(&self.alpha.to_arr(), 2));
 
         let mut output: Vec<u8> = Vec::new();
         output.push(ProtocolMessageType::from(self) as u8 + 1);
-        output.extend_from_slice(&serialize(registration_request, 3));
+        output.extend_from_slice(&serialize(&registration_request, 3));
         output
     }
 
     /// Deserialization from bytes
     pub fn deserialize(input: &[u8]) -> Result<Self, ProtocolError> {
-        if input[0] != ProtocolMessageType::RegistrationRequest as u8 + 1 {
+        if input.is_empty()
+            || input.is_empty()
+            || input[0] != ProtocolMessageType::RegistrationRequest as u8 + 1
+        {
             return Err(PakeError::SerializationError.into());
         }
 
@@ -79,14 +89,23 @@ impl<Grp: Group> RegisterFirstMessage<Grp> {
             return Err(PakeError::SerializationError.into());
         }
 
-        let (_, remainder) = tokenize(data, 2)?;
+        let (id_u, remainder) = tokenize(data, 2)?;
         let (alpha_bytes, remainder) = tokenize(remainder, 2)?;
 
         if !remainder.is_empty() {
             return Err(PakeError::SerializationError.into());
         }
 
-        Self::try_from(&alpha_bytes[..])
+        let checked_slice = check_slice_size(
+            &alpha_bytes,
+            Grp::ElemLen::to_usize(),
+            "first_message_bytes",
+        )?;
+        // Check that the message is actually containing an element of the
+        // correct subgroup
+        let arr = GenericArray::from_slice(checked_slice);
+        let alpha = Grp::from_element_slice(arr)?;
+        Ok(Self { id_u, alpha })
     }
 }
 
@@ -95,6 +114,10 @@ impl<Grp: Group> RegisterFirstMessage<Grp> {
 pub struct RegisterSecondMessage<Grp> {
     /// The server's oprf output
     beta: Grp,
+    /// Server's static public key
+    server_s_pk: Vec<u8>,
+    /// Envelope credentials format
+    ecf: EnvelopeCredentialsFormat,
 }
 
 impl<Grp> TryFrom<&[u8]> for RegisterSecondMessage<Grp>
@@ -103,17 +126,25 @@ where
 {
     type Error = ProtocolError;
 
-    fn try_from(second_message_bytes: &[u8]) -> Result<Self, Self::Error> {
-        let checked_slice = check_slice_size(
-            second_message_bytes,
-            Grp::ElemLen::to_usize(),
-            "second_message_bytes",
-        )?;
+    fn try_from(bytes: &[u8]) -> Result<Self, Self::Error> {
+        let elem_len = Grp::ElemLen::to_usize();
+        let checked_slice = check_slice_size_atleast(bytes, elem_len, "second_message_bytes")?;
+
         // Check that the message is actually containing an element of the
         // correct subgroup
-        let arr = GenericArray::from_slice(&checked_slice);
+        let arr = GenericArray::from_slice(&checked_slice[..elem_len]);
         let beta = Grp::from_element_slice(arr)?;
-        Ok(Self { beta })
+
+        let server_s_pk = checked_slice[elem_len..].to_vec();
+
+        // Note that we use a default envelope credentials format here, since it
+        // is not included in the byte representation
+        let ecf = EnvelopeCredentialsFormat::default()?;
+        Ok(Self {
+            beta,
+            server_s_pk,
+            ecf,
+        })
     }
 }
 
@@ -121,30 +152,47 @@ impl<Grp> RegisterSecondMessage<Grp>
 where
     Grp: Group,
 {
-    /// byte representation for the registration response message
-    fn to_bytes(&self) -> Vec<u8> {
-        self.beta.to_arr().to_vec()
+    /// Byte representation for the registration response message. This does not
+    /// include the envelope credentials format
+    pub fn to_bytes(&self) -> Vec<u8> {
+        [&self.beta.to_arr().to_vec()[..], &self.server_s_pk[..]].concat()
     }
 
     /// Serialization into bytes
     pub fn serialize(&self) -> Vec<u8> {
         let mut registration_response: Vec<u8> = Vec::new();
-        registration_response.extend_from_slice(&serialize((&self.to_bytes()).to_vec(), 2));
-        registration_response.extend_from_slice(&serialize(Vec::new(), 2));
+        registration_response.extend_from_slice(&serialize(&self.beta.to_arr(), 2));
+        registration_response.extend_from_slice(&serialize(&self.server_s_pk, 2));
 
-        // TODO: The following should not be hardcoded, but instead be customizable
-        registration_response.extend_from_slice(&[1u8, CredentialType::SkU as u8 + 1]);
-        registration_response.extend_from_slice(&[1u8, CredentialType::PkS as u8 + 1]);
+        // Handle ecf serialization
+        let secret_credentials: Vec<u8> = self
+            .ecf
+            .secret_credentials
+            .iter()
+            .map(|&x| x as u8 + 1)
+            .collect();
+        let cleartext_credentials: Vec<u8> = self
+            .ecf
+            .cleartext_credentials
+            .iter()
+            .map(|&x| x as u8 + 1)
+            .collect();
+        let ecf_serialized = [
+            serialize(&secret_credentials, 1),
+            serialize(&cleartext_credentials, 1),
+        ]
+        .concat();
+        registration_response.extend_from_slice(&ecf_serialized);
 
         let mut output: Vec<u8> = Vec::new();
         output.push(ProtocolMessageType::from(self) as u8 + 1);
-        output.extend_from_slice(&serialize(registration_response, 3));
+        output.extend_from_slice(&serialize(&registration_response, 3));
         output
     }
 
     /// Deserialization from bytes
     pub fn deserialize(input: &[u8]) -> Result<Self, ProtocolError> {
-        if input[0] != ProtocolMessageType::RegistrationResponse as u8 + 1 {
+        if input.is_empty() || input[0] != ProtocolMessageType::RegistrationResponse as u8 + 1 {
             return Err(PakeError::SerializationError.into());
         }
 
@@ -154,18 +202,39 @@ where
         }
 
         let (beta_bytes, remainder) = tokenize(data, 2)?;
-        let (_, remainder) = tokenize(remainder, 2)?;
+        let (server_s_pk, remainder) = tokenize(remainder, 2)?;
 
-        // TODO: The following should affect what is placed in the envelope rather than
-        // being ignored
-        let (_, remainder) = tokenize(remainder, 1)?;
-        let (_, remainder) = tokenize(remainder, 1)?;
+        // Handle ecf deserialization
+        let (secret_credentials, remainder) = tokenize(remainder, 1)?;
+        let (cleartext_credentials, remainder) = tokenize(remainder, 1)?;
+        let sc = secret_credentials
+            .iter()
+            .map(|x| u8_to_credential_type(*x).ok_or(PakeError::SerializationError))
+            .collect::<Result<Vec<CredentialType>, _>>()?;
+        let cc = cleartext_credentials
+            .iter()
+            .map(|x| u8_to_credential_type(*x).ok_or(PakeError::SerializationError))
+            .collect::<Result<Vec<CredentialType>, _>>()?;
+        let ecf = EnvelopeCredentialsFormat::new(sc, cc)?;
 
         if !remainder.is_empty() {
             return Err(PakeError::SerializationError.into());
         }
 
-        Self::try_from(&beta_bytes[..])
+        let checked_slice = check_slice_size(
+            &beta_bytes,
+            Grp::ElemLen::to_usize(),
+            "second_message_bytes",
+        )?;
+        // Check that the message is actually containing an element of the
+        // correct subgroup
+        let arr = GenericArray::from_slice(&checked_slice);
+        let beta = Grp::from_element_slice(arr)?;
+        Ok(Self {
+            ecf,
+            server_s_pk,
+            beta,
+        })
     }
 }
 
@@ -213,17 +282,17 @@ where
     pub fn serialize(&self) -> Vec<u8> {
         let mut registration_upload: Vec<u8> = Vec::new();
         registration_upload.extend_from_slice(&self.envelope.serialize());
-        registration_upload.extend_from_slice(&serialize(self.client_s_pk.to_arr().to_vec(), 2));
+        registration_upload.extend_from_slice(&serialize(&self.client_s_pk.to_arr(), 2));
 
         let mut output: Vec<u8> = Vec::new();
         output.push(ProtocolMessageType::from(self) as u8 + 1);
-        output.extend_from_slice(&serialize(registration_upload, 3));
+        output.extend_from_slice(&serialize(&registration_upload, 3));
         output
     }
 
     /// Deserialization from bytes
     pub fn deserialize(input: &[u8]) -> Result<Self, ProtocolError> {
-        if input[0] != ProtocolMessageType::RegistrationUpload as u8 + 1 {
+        if input.is_empty() || input[0] != ProtocolMessageType::RegistrationUpload as u8 + 1 {
             return Err(PakeError::SerializationError.into());
         }
 
@@ -248,6 +317,8 @@ where
 
 /// The message sent by the user to the server, to initiate registration
 pub struct LoginFirstMessage<CS: CipherSuite> {
+    /// User identity
+    id_u: Vec<u8>,
     /// blinded password information
     alpha: CS::Group,
     ke1_message: <CS::KeyExchange as KeyExchange<CS::Hash, CS::KeyFormat>>::KE1Message,
@@ -256,27 +327,7 @@ pub struct LoginFirstMessage<CS: CipherSuite> {
 impl<CS: CipherSuite> TryFrom<&[u8]> for LoginFirstMessage<CS> {
     type Error = ProtocolError;
     fn try_from(first_message_bytes: &[u8]) -> Result<Self, Self::Error> {
-        let min_expected_len = <CS::Group as Group>::ElemLen::to_usize();
-        let checked_slice = (if first_message_bytes.len() <= min_expected_len {
-            Err(InternalPakeError::SizeError {
-                name: "first_message_bytes",
-                len: min_expected_len,
-                actual_len: first_message_bytes.len(),
-            })
-        } else {
-            Ok(first_message_bytes)
-        })?;
-        // Check that the message is actually containing an element of the
-        // correct subgroup
-        let elem_len = <CS::Group as Group>::ElemLen::to_usize();
-        let arr = GenericArray::from_slice(&checked_slice[..elem_len]);
-        let alpha = CS::Group::from_element_slice(arr)?;
-
-        let ke1_message =
-            <CS::KeyExchange as KeyExchange<CS::Hash, CS::KeyFormat>>::KE1Message::try_from(
-                &checked_slice[elem_len..],
-            )?;
-        Ok(Self { alpha, ke1_message })
+        Self::deserialize(first_message_bytes)
     }
 }
 
@@ -289,33 +340,46 @@ impl<CS: CipherSuite> LoginFirstMessage<CS> {
     /// Serialization into bytes
     pub fn serialize(&self) -> Vec<u8> {
         let mut credential_request: Vec<u8> = Vec::new();
-        credential_request.extend_from_slice(&serialize(Vec::new(), 2));
-        credential_request.extend_from_slice(&serialize((&self.alpha.to_arr()).to_vec(), 2));
+        credential_request.extend_from_slice(&serialize(&self.id_u, 2));
+        credential_request.extend_from_slice(&serialize(&self.alpha.to_arr(), 2));
 
         let mut output: Vec<u8> = Vec::new();
         output.push(ProtocolMessageType::from(self) as u8 + 1);
-        output.extend_from_slice(&serialize(credential_request, 3));
+        output.extend_from_slice(&serialize(&credential_request, 3));
         output.extend_from_slice(&self.ke1_message.to_bytes());
         output
     }
 
     /// Deserialization from bytes
     pub fn deserialize(input: &[u8]) -> Result<Self, ProtocolError> {
-        if input[0] != ProtocolMessageType::CredentialRequest as u8 + 1 {
+        if input.is_empty() || input[0] != ProtocolMessageType::CredentialRequest as u8 + 1 {
             return Err(PakeError::SerializationError.into());
         }
 
         let (data, ke1m) = tokenize(input[1..].to_vec(), 3)?;
 
-        let (_, remainder) = tokenize(data, 2)?;
+        let (id_u, remainder) = tokenize(data, 2)?;
         let (alpha_bytes, remainder) = tokenize(remainder, 2)?;
 
         if !remainder.is_empty() {
             return Err(PakeError::SerializationError.into());
         }
 
-        let concatenated = [&alpha_bytes[..], &ke1m[..]].concat();
-        Self::try_from(&concatenated[..])
+        let elem_len = <CS::Group as Group>::ElemLen::to_usize();
+        let checked_slice = check_slice_size(&alpha_bytes, elem_len, "login_first_message_bytes")?;
+        let arr = GenericArray::from_slice(&checked_slice[..elem_len]);
+        let alpha = <CS::Group as Group>::from_element_slice(arr)?;
+
+        let ke1_message =
+            <CS::KeyExchange as KeyExchange<CS::Hash, CS::KeyFormat>>::KE1Message::try_from(
+                &ke1m[..],
+            )?;
+
+        Ok(Self {
+            id_u,
+            alpha,
+            ke1_message,
+        })
     }
 }
 
@@ -333,28 +397,25 @@ impl<CS: CipherSuite> LoginSecondMessage<CS> {
     /// Serialization into bytes
     pub fn serialize(&self) -> Vec<u8> {
         let mut credential_response: Vec<u8> = Vec::new();
-        credential_response.extend_from_slice(&serialize((&self.beta.to_arr()).to_vec(), 2));
-        credential_response.extend_from_slice(&serialize((&self.envelope.to_bytes()).to_vec(), 2));
-        credential_response.extend_from_slice(&serialize(Vec::new(), 2));
+        credential_response.extend_from_slice(&serialize(&self.beta.to_arr(), 2));
+        credential_response.extend_from_slice(&serialize(&self.envelope.to_bytes(), 2));
 
         let mut output: Vec<u8> = Vec::new();
         output.push(ProtocolMessageType::from(self) as u8 + 1);
-        output.extend_from_slice(&serialize(credential_response, 3));
+        output.extend_from_slice(&serialize(&credential_response, 3));
         output.extend_from_slice(&self.ke2_message.to_bytes());
         output
     }
 
     /// Deserialization from bytes
     pub fn deserialize(input: &[u8]) -> Result<Self, ProtocolError> {
-        if input[0] != ProtocolMessageType::CredentialResponse as u8 + 1 {
+        if input.is_empty() || input[0] != ProtocolMessageType::CredentialResponse as u8 + 1 {
             return Err(PakeError::SerializationError.into());
         }
 
         let (data, ke2m) = tokenize(input[1..].to_vec(), 3)?;
-
         let (beta_bytes, remainder) = tokenize(data, 2)?;
         let (envelope_bytes, remainder) = tokenize(remainder, 2)?;
-        let (_, remainder) = tokenize(remainder, 2)?;
 
         if !remainder.is_empty() {
             return Err(PakeError::SerializationError.into());
@@ -368,15 +429,9 @@ impl<CS: CipherSuite> LoginSecondMessage<CS> {
 impl<CS: CipherSuite> TryFrom<&[u8]> for LoginSecondMessage<CS> {
     type Error = ProtocolError;
     fn try_from(second_message_bytes: &[u8]) -> Result<Self, Self::Error> {
-        let key_len = <<CS::KeyFormat as KeyPair>::Repr as SizedBytes>::Len::to_usize();
-        let envelope_size = key_len + Envelope::<CS::Hash>::additional_size();
         let elem_len = <CS::Group as Group>::ElemLen::to_usize();
-        let ke2_message_size = CS::KeyExchange::ke2_message_size();
-        let checked_slice = check_slice_size(
-            second_message_bytes,
-            elem_len + envelope_size + ke2_message_size,
-            "login_second_message_bytes",
-        )?;
+        let checked_slice =
+            check_slice_size_atleast(second_message_bytes, elem_len, "login_second_message_bytes")?;
 
         // Check that the message is actually containing an element of the
         // correct subgroup
@@ -384,12 +439,14 @@ impl<CS: CipherSuite> TryFrom<&[u8]> for LoginSecondMessage<CS> {
         let arr = GenericArray::from_slice(beta_bytes);
         let beta = CS::Group::from_element_slice(arr)?;
 
-        let envelope =
-            Envelope::<CS::Hash>::from_bytes(&checked_slice[elem_len..elem_len + envelope_size])?;
+        let (envelope, remainder) = Envelope::<CS::Hash>::deserialize(&checked_slice[elem_len..])?;
 
+        let ke2_message_size = CS::KeyExchange::ke2_message_size();
+        let checked_remainder =
+            check_slice_size_atleast(&remainder, ke2_message_size, "login_second_message_bytes")?;
         let ke2_message =
             <CS::KeyExchange as KeyExchange<CS::Hash, CS::KeyFormat>>::KE2Message::try_from(
-                &checked_slice[elem_len + envelope_size..],
+                &checked_remainder,
             )?;
 
         Ok(Self {
@@ -417,6 +474,14 @@ impl<CS: CipherSuite> TryFrom<&[u8]> for LoginThirdMessage<CS> {
 }
 
 impl<CS: CipherSuite> LoginThirdMessage<CS> {
+    /// Serialization into bytes
+    pub fn serialize(&self) -> Vec<u8> {
+        let mut output: Vec<u8> = Vec::new();
+        output.push(ProtocolMessageType::from(self) as u8 + 1);
+        output.extend_from_slice(&self.ke3_message.to_bytes());
+        output
+    }
+
     /// byte representation for the login finalization
     pub fn to_bytes(&self) -> Vec<u8> {
         self.ke3_message.to_bytes()
@@ -428,6 +493,10 @@ impl<CS: CipherSuite> LoginThirdMessage<CS> {
 
 /// The state elements the client holds to register itself
 pub struct ClientRegistration<CS: CipherSuite> {
+    /// User identity
+    id_u: Vec<u8>,
+    /// Server identity
+    id_s: Vec<u8>,
     /// a blinding factor
     pub(crate) blinding_factor: <CS::Group as Group>::Scalar,
     /// the client's password
@@ -436,7 +505,10 @@ pub struct ClientRegistration<CS: CipherSuite> {
 
 impl<CS: CipherSuite> TryFrom<&[u8]> for ClientRegistration<CS> {
     type Error = ProtocolError;
-    fn try_from(bytes: &[u8]) -> Result<Self, Self::Error> {
+    fn try_from(input: &[u8]) -> Result<Self, Self::Error> {
+        let (id_u, bytes) = tokenize(input.to_vec(), 2)?;
+        let (id_s, bytes) = tokenize(bytes.to_vec(), 2)?;
+
         let min_expected_len = <CS::Group as Group>::ScalarLen::to_usize();
         let checked_slice = (if bytes.len() <= min_expected_len {
             Err(InternalPakeError::SizeError {
@@ -455,6 +527,8 @@ impl<CS: CipherSuite> TryFrom<&[u8]> for ClientRegistration<CS> {
         let blinding_factor = CS::Group::from_scalar_slice(blinding_factor_bytes)?;
         let password = checked_slice[scalar_len..].to_vec();
         Ok(Self {
+            id_u,
+            id_s,
             blinding_factor,
             password,
         })
@@ -465,6 +539,8 @@ impl<CS: CipherSuite> ClientRegistration<CS> {
     /// byte representation for the client's registration state
     pub fn to_bytes(&self) -> Vec<u8> {
         let output: Vec<u8> = [
+            &serialize(&self.id_u, 2),
+            &serialize(&self.id_s, 2),
             &CS::Group::scalar_as_bytes(&self.blinding_factor)[..],
             &self.password,
         ]
@@ -502,6 +578,23 @@ impl<CS: CipherSuite> ClientRegistration<CS> {
         password: &[u8],
         pepper: Option<&[u8]>,
         blinding_factor_rng: &mut R,
+    ) -> Result<(RegisterFirstMessage<CS::Group>, Self), ProtocolError> {
+        Self::start_with_user_and_server_name(
+            &Vec::new(),
+            &Vec::new(),
+            password,
+            pepper,
+            blinding_factor_rng,
+        )
+    }
+
+    /// Same as ClientRegistration::start, but also accepts a username and server name as input
+    pub fn start_with_user_and_server_name<R: RngCore + CryptoRng>(
+        user_name: &[u8],
+        server_name: &[u8],
+        password: &[u8],
+        pepper: Option<&[u8]>,
+        blinding_factor_rng: &mut R,
     ) -> Result<(RegisterFirstMessage<CS::Group>, Self), ProtocolError> {
         let OprfClientBytes {
             alpha,
@@ -509,8 +602,13 @@ impl<CS: CipherSuite> ClientRegistration<CS> {
         } = oprf::generate_oprf1::<R, CS::Group>(&password, pepper, blinding_factor_rng)?;
 
         Ok((
-            RegisterFirstMessage::<CS::Group> { alpha },
+            RegisterFirstMessage::<CS::Group> {
+                id_u: user_name.to_vec(),
+                alpha,
+            },
             Self {
+                id_u: user_name.to_vec(),
+                id_s: server_name.to_vec(),
                 blinding_factor,
                 password: password.to_vec(),
             },
@@ -561,6 +659,17 @@ impl<CS: CipherSuite> ClientRegistration<CS> {
         r2: RegisterSecondMessage<CS::Group>,
         server_s_pk: &<CS::KeyFormat as KeyPair>::Repr,
         rng: &mut R,
+    ) -> Result<ClientRegistrationFinishResult<CS::KeyFormat, CS::Hash>, ProtocolError> {
+        let mut r2_cloned = r2;
+        r2_cloned.server_s_pk = server_s_pk.to_arr().to_vec();
+        self.finish_using_transmitted_server_public_key(r2_cloned, rng)
+    }
+
+    /// Same as finish, but without the server public key check
+    pub fn finish_using_transmitted_server_public_key<R: CryptoRng + RngCore>(
+        self,
+        r2: RegisterSecondMessage<CS::Group>,
+        rng: &mut R,
     ) -> Result<ClientRegistrationFinishResult<CS::KeyFormat, CS::Hash>, ProtocolError> {
         let client_static_keypair = CS::KeyFormat::generate_random(rng)?;
 
@@ -570,12 +679,21 @@ impl<CS: CipherSuite> ClientRegistration<CS> {
             &self.blinding_factor,
         )?;
 
-        let (envelope, export_key) = Envelope::<CS::Hash>::seal(
-            &password_derived_key,
-            &client_static_keypair.private().to_arr(),
-            &server_s_pk.to_arr(),
-            rng,
-        )?;
+        let mut credentials_map: HashMap<CredentialType, Vec<u8>> = HashMap::new();
+        credentials_map.insert(
+            CredentialType::SkU,
+            client_static_keypair.private().to_arr().to_vec(),
+        );
+        credentials_map.insert(
+            CredentialType::PkU,
+            client_static_keypair.public().to_arr().to_vec(),
+        );
+        credentials_map.insert(CredentialType::PkS, r2.server_s_pk);
+        credentials_map.insert(CredentialType::IdU, self.id_u.clone());
+        credentials_map.insert(CredentialType::IdS, self.id_s.clone());
+
+        let (envelope, export_key) =
+            Envelope::<CS::Hash>::seal(&password_derived_key, r2.ecf, credentials_map, rng)?;
 
         Ok((
             RegisterThirdMessage {
@@ -632,36 +750,36 @@ where
     >: generic_array::ArrayLength<u8>,
 {
     type Error = ProtocolError;
-    fn try_from(server_registration_bytes: &[u8]) -> Result<Self, Self::Error> {
-        let key_len = <<CS::KeyFormat as KeyPair>::Repr as SizedBytes>::Len::to_usize();
-        let scalar_len = <CS::Group as Group>::ScalarLen::to_usize();
-        let envelope_size = key_len + Envelope::<CS::Hash>::additional_size();
 
-        if server_registration_bytes.len() == scalar_len {
+    /// The format of a serialized ServerRegistration object:
+    /// oprf_key | client_s_pk | envelope
+    fn try_from(input: &[u8]) -> Result<Self, Self::Error> {
+        let scalar_len = <CS::Group as Group>::ScalarLen::to_usize();
+        if input.len() == scalar_len {
             return Ok(Self {
-                oprf_key: CS::Group::from_scalar_slice(GenericArray::from_slice(
-                    server_registration_bytes,
-                ))?,
+                oprf_key: CS::Group::from_scalar_slice(GenericArray::from_slice(input))?,
                 client_s_pk: None,
                 envelope: None,
             });
         }
 
-        let checked_bytes = check_slice_size(
-            server_registration_bytes,
-            envelope_size + key_len + scalar_len,
-            "server_registration_bytes",
-        )?;
+        // Need to do this check manually because envelope is variable-size
+        let key_len = <<CS::KeyFormat as KeyPair>::Repr as SizedBytes>::Len::to_usize();
+
+        let checked_bytes =
+            check_slice_size_atleast(&input, scalar_len + key_len, "server_registration_bytes")?;
+
         let oprf_key_bytes = GenericArray::from_slice(&checked_bytes[..scalar_len]);
         let oprf_key = CS::Group::from_scalar_slice(oprf_key_bytes)?;
         let unchecked_client_s_pk = <CS::KeyFormat as KeyPair>::Repr::from_bytes(
             &checked_bytes[scalar_len..scalar_len + key_len],
         )?;
         let client_s_pk = CS::KeyFormat::check_public_key(unchecked_client_s_pk)?;
+
+        let envelope = Envelope::<CS::Hash>::from_bytes(&checked_bytes[scalar_len + key_len..])?;
+
         Ok(Self {
-            envelope: Some(Envelope::<CS::Hash>::from_bytes(
-                &checked_bytes[checked_bytes.len() - envelope_size..],
-            )?),
+            envelope: Some(envelope),
             client_s_pk: Some(client_s_pk),
             oprf_key,
         })
@@ -720,6 +838,30 @@ where
     pub fn start<R: RngCore + CryptoRng>(
         message: RegisterFirstMessage<CS::Group>,
         rng: &mut R,
+    ) -> Result<(RegisterSecondMessage<CS::Group>, Self), ProtocolError> {
+        Self::start_with_server_pk(message, &Vec::new(), rng)
+    }
+
+    /// Same as start, but with the ability to supply a server_s_pk as input
+    pub fn start_with_server_pk<R: RngCore + CryptoRng>(
+        message: RegisterFirstMessage<CS::Group>,
+        server_s_pk: &[u8],
+        rng: &mut R,
+    ) -> Result<(RegisterSecondMessage<CS::Group>, Self), ProtocolError> {
+        Self::start_with_server_pk_and_ecf(
+            message,
+            server_s_pk,
+            EnvelopeCredentialsFormat::default()?,
+            rng,
+        )
+    }
+
+    /// Same as start, but with the ability to supply a server_s_pk as input and envelope credentials format
+    pub fn start_with_server_pk_and_ecf<R: RngCore + CryptoRng>(
+        message: RegisterFirstMessage<CS::Group>,
+        server_s_pk: &[u8],
+        ecf: EnvelopeCredentialsFormat,
+        rng: &mut R,
     ) -> Result<(RegisterSecondMessage<CS::Group>, Self), ProtocolError> {
         // RFC: generate oprf_key (salt) and v_u = g^oprf_key
         let oprf_key = CS::Group::random_scalar(rng);
@@ -728,7 +870,11 @@ where
         let beta = oprf::generate_oprf2::<CS::Group>(message.alpha, &oprf_key)?;
 
         Ok((
-            RegisterSecondMessage { beta },
+            RegisterSecondMessage {
+                beta,
+                server_s_pk: server_s_pk.to_vec(),
+                ecf,
+            },
             Self {
                 envelope: None,
                 client_s_pk: None,
@@ -786,6 +932,10 @@ where
 
 /// The state elements the client holds to perform a login
 pub struct ClientLogin<CS: CipherSuite> {
+    /// User identity
+    id_u: Vec<u8>,
+    /// Server identity
+    id_s: Vec<u8>,
     /// A blinding factor, which is used to mask (and unmask) secret
     /// information before transmission
     blinding_factor: <CS::Group as Group>::Scalar,
@@ -796,7 +946,10 @@ pub struct ClientLogin<CS: CipherSuite> {
 
 impl<CS: CipherSuite> TryFrom<&[u8]> for ClientLogin<CS> {
     type Error = ProtocolError;
-    fn try_from(bytes: &[u8]) -> Result<Self, Self::Error> {
+    fn try_from(input: &[u8]) -> Result<Self, Self::Error> {
+        let (id_u, bytes) = tokenize(input.to_vec(), 2)?;
+        let (id_s, bytes) = tokenize(bytes.to_vec(), 2)?;
+
         let scalar_len = <CS::Group as Group>::ScalarLen::to_usize();
         let ke1_state_size =
             <CS::KeyExchange as KeyExchange<CS::Hash, CS::KeyFormat>>::ke1_state_size();
@@ -809,7 +962,7 @@ impl<CS: CipherSuite> TryFrom<&[u8]> for ClientLogin<CS> {
                 actual_len: bytes.len(),
             })
         } else {
-            Ok(bytes)
+            Ok(bytes.clone())
         })?;
 
         let blinding_factor_bytes = GenericArray::from_slice(&checked_slice[..scalar_len]);
@@ -820,6 +973,8 @@ impl<CS: CipherSuite> TryFrom<&[u8]> for ClientLogin<CS> {
             )?;
         let password = bytes[scalar_len + ke1_state_size..].to_vec();
         Ok(Self {
+            id_u,
+            id_s,
             blinding_factor,
             password,
             ke1_state,
@@ -831,6 +986,8 @@ impl<CS: CipherSuite> ClientLogin<CS> {
     /// byte representation for the client's login state
     pub fn to_bytes(&self) -> Vec<u8> {
         let output: Vec<u8> = [
+            &serialize(&self.id_u, 2),
+            &serialize(&self.id_s, 2),
             &CS::Group::scalar_as_bytes(&self.blinding_factor)[..],
             &self.ke1_state.to_bytes(),
             &self.password,
@@ -875,6 +1032,17 @@ impl<CS: CipherSuite> ClientLogin<CS> {
         password: &[u8],
         pepper: Option<&[u8]>,
         rng: &mut R,
+    ) -> Result<(LoginFirstMessage<CS>, Self), ProtocolError> {
+        Self::start_with_user_and_server_name(&Vec::new(), &Vec::new(), password, pepper, rng)
+    }
+
+    /// Same as start, but allows the user to supply a username and server name
+    pub fn start_with_user_and_server_name<R: RngCore + CryptoRng>(
+        user_name: &[u8],
+        server_name: &[u8],
+        password: &[u8],
+        pepper: Option<&[u8]>,
+        rng: &mut R,
     ) -> Result<(LoginFirstMessage<CS>, Self), ProtocolError> {
         let OprfClientBytes {
             alpha,
@@ -883,11 +1051,17 @@ impl<CS: CipherSuite> ClientLogin<CS> {
 
         let (ke1_state, ke1_message) = CS::KeyExchange::generate_ke1(alpha.to_arr().to_vec(), rng)?;
 
-        let l1 = LoginFirstMessage { alpha, ke1_message };
+        let l1 = LoginFirstMessage {
+            id_u: user_name.to_vec(),
+            alpha,
+            ke1_message,
+        };
 
         Ok((
             l1,
             Self {
+                id_u: user_name.to_vec(),
+                id_s: server_name.to_vec(),
                 blinding_factor,
                 password: password.to_vec(),
                 ke1_state,
@@ -933,7 +1107,7 @@ impl<CS: CipherSuite> ClientLogin<CS> {
     pub fn finish<R: RngCore + CryptoRng>(
         self,
         l2: LoginSecondMessage<CS>,
-        server_s_pk: &<<CS as CipherSuite>::KeyFormat as KeyPair>::Repr,
+        _server_s_pk: &<<CS as CipherSuite>::KeyFormat as KeyPair>::Repr,
         _client_e_sk_rng: &mut R,
     ) -> Result<ClientLoginFinishResult<CS>, ProtocolError> {
         let l2_bytes: Vec<u8> = [&l2.beta.to_arr()[..], &l2.envelope.to_bytes()].concat();
@@ -946,7 +1120,7 @@ impl<CS: CipherSuite> ClientLogin<CS> {
 
         let opened_envelope = &l2
             .envelope
-            .open(&password_derived_key, &server_s_pk.to_arr())
+            .open(&password_derived_key)
             .map_err(|e| match e {
                 InternalPakeError::SealOpenHmacError => PakeError::InvalidLoginError,
                 err => PakeError::from(err),
@@ -956,8 +1130,12 @@ impl<CS: CipherSuite> ClientLogin<CS> {
             l2_bytes,
             l2.ke2_message,
             &self.ke1_state,
-            server_s_pk.clone(),
-            <CS::KeyFormat as KeyPair>::Repr::from_bytes(&opened_envelope.plaintext)?,
+            <CS::KeyFormat as KeyPair>::Repr::from_bytes(
+                &opened_envelope.credentials_map[&CredentialType::PkS],
+            )?,
+            <CS::KeyFormat as KeyPair>::Repr::from_bytes(
+                &opened_envelope.credentials_map[&CredentialType::SkU],
+            )?,
         )?;
 
         Ok((