Repository: novifinancial/opaque-ke
Pull request #80 (2020-11-06 10:47:30)

Adding expand_message_xmd and RFC compliance with VOPRF and H2C
Addresses #72, #73, and #74 

Note that x25519 is not in the VOPRF spec, and so we actually don't have anything to match against... but this will handle ristretto255 at least.

huitseeker: @kevinlewi See also the run results on #81 which show we already have issues in this area
https://github.com/novifinancial/opaque-ke/pull/81/checks?check_run_id=1374926279

Patch:
diff --git a/src/errors.rs b/src/errors.rs
index 7fb161f..ac2b642 100644
--- a/src/errors.rs
+++ b/src/errors.rs
@@ -27,6 +27,8 @@ pub enum InternalPakeError {
     SubGroupError,
     /// hashing to a key failed
     HashingFailure,
+    /// Computing the hash-to-curve function failed
+    HashToCurveError,
     /// Computing HKDF failed while deriving subkeys
     HkdfError,
     /// Computing HMAC failed while supplying a secret key
diff --git a/src/key_exchange/tripledh.rs b/src/key_exchange/tripledh.rs
index f1df66f..553be9a 100644
--- a/src/key_exchange/tripledh.rs
+++ b/src/key_exchange/tripledh.rs
@@ -287,7 +287,6 @@ impl<KeyFormat: KeyPair> TryFrom<&[u8]> for KE1Message<KeyFormat> {
         })
     }
 }
-
 /// The server state produced after the second key exchange message
 pub struct KE2State<HashLen: ArrayLength<u8>> {
     km3: GenericArray<u8, HashLen>,
diff --git a/src/map_to_curve.rs b/src/map_to_curve.rs
index 64ffabe..41e8e54 100644
--- a/src/map_to_curve.rs
+++ b/src/map_to_curve.rs
@@ -6,31 +6,231 @@
 //! Defines the GroupWithMapToCurve trait to specify how to map a password to a
 //! curve point
 
+use crate::errors::InternalPakeError;
 use crate::group::Group;
+use crate::hash::Hash;
+use crate::serialization::i2osp;
 use curve25519_dalek::{edwards::EdwardsPoint, ristretto::RistrettoPoint};
-
+use digest::{BlockInput, Digest};
+use generic_array::typenum::Unsigned;
+use generic_array::GenericArray;
 use hkdf::Hkdf;
-use sha2::{Sha256, Sha512};
+use sha2::Sha256;
 
 /// A subtrait of Group specifying how to hash a password into a point
 pub trait GroupWithMapToCurve: Group {
-    /// transforms a password and optional pepper into a curve point
-    fn map_to_curve(password: &[u8], dst: Option<&[u8]>) -> Self;
-}
+    /// The ciphersuite identifier as dictated by
+    /// https://www.ietf.org/archive/id/draft-irtf-cfrg-voprf-05.txt
+    const SUITE_ID: usize;
 
-// TODO: incorporate expand_message_xmd from https://www.ietf.org/archive/id/draft-irtf-cfrg-hash-to-curve-10.txt
-// instead of using HKDF-extract here
+    /// transforms a password and domain separation tag (DST) into a curve point
+    fn map_to_curve(msg: &[u8], dst: &[u8]) -> Result<Self, InternalPakeError>;
+
+    /// Generates the contextString parameter as defined in
+    /// https://www.ietf.org/archive/id/draft-irtf-cfrg-voprf-05.txt
+    fn get_context_string(mode: u8) -> Vec<u8> {
+        [i2osp(mode as usize, 1), i2osp(Self::SUITE_ID, 2)].concat()
+    }
+}
 
 impl GroupWithMapToCurve for RistrettoPoint {
-    fn map_to_curve(password: &[u8], dst: Option<&[u8]>) -> Self {
-        let (hashed_input, _) = Hkdf::<Sha512>::extract(dst, password);
-        <Self as Group>::hash_to_curve(&hashed_input)
+    const SUITE_ID: usize = 0x01;
+
+    // Implements the hash_to_ristretto255() function from
+    // https://www.ietf.org/archive/id/draft-irtf-cfrg-hash-to-curve-10.txt
+    fn map_to_curve(msg: &[u8], dst: &[u8]) -> Result<Self, InternalPakeError> {
+        let uniform_bytes = expand_message_xmd::<Sha256>(msg, dst, 64)?;
+        Ok(<Self as Group>::hash_to_curve(
+            &GenericArray::clone_from_slice(&uniform_bytes[..]),
+        ))
     }
 }
 
 impl GroupWithMapToCurve for EdwardsPoint {
-    fn map_to_curve(password: &[u8], dst: Option<&[u8]>) -> Self {
-        let (hashed_input, _) = Hkdf::<Sha256>::extract(dst, password);
-        <Self as Group>::hash_to_curve(&hashed_input)
+    const SUITE_ID: usize = 0x09; // FIXME, seemingly unsupported by VOPRF RFC?
+    fn map_to_curve(msg: &[u8], dst: &[u8]) -> Result<Self, InternalPakeError> {
+        let (hashed_input, _) = Hkdf::<Sha256>::extract(Some(dst), msg);
+        Ok(<Self as Group>::hash_to_curve(&hashed_input))
+    }
+}
+
+// Computes ceil(x / y)
+fn div_ceil(x: usize, y: usize) -> usize {
+    let additive = (x % y != 0) as usize;
+    x / y + additive
+}
+
+fn xor(x: &[u8], y: &[u8]) -> Result<Vec<u8>, InternalPakeError> {
+    if x.len() != y.len() {
+        return Err(InternalPakeError::HashToCurveError);
+    }
+
+    Ok(x.iter().zip(y).map(|(&x1, &x2)| x1 ^ x2).collect())
+}
+
+// Corresponds to the expand_message_xmd() function defined in
+// https://www.ietf.org/archive/id/draft-irtf-cfrg-hash-to-curve-10.txt
+pub(crate) fn expand_message_xmd<H: Hash>(
+    msg: &[u8],
+    dst: &[u8],
+    len_in_bytes: usize,
+) -> Result<Vec<u8>, InternalPakeError> {
+    let b_in_bytes = <H as Digest>::OutputSize::to_usize();
+    let r_in_bytes = <H as BlockInput>::BlockSize::to_usize();
+
+    let ell = div_ceil(len_in_bytes, b_in_bytes);
+    if ell > 255 {
+        return Err(InternalPakeError::HashToCurveError);
+    }
+    let dst_prime = [dst, &i2osp(dst.len(), 1)].concat();
+    let z_pad = i2osp(0, r_in_bytes);
+    let l_i_b_str = i2osp(len_in_bytes, 2);
+    let msg_prime = [&z_pad, msg, &l_i_b_str, &i2osp(0, 1), &dst_prime].concat();
+
+    let mut b: Vec<Vec<u8>> = Vec::new();
+    b.push(H::digest(&msg_prime).to_vec()); // b[0]
+
+    let mut h = H::new();
+    h.update(&b[0]);
+    h.update(&i2osp(1, 1));
+    h.update(&dst_prime);
+    b.push(h.finalize().to_vec()); // b[1]
+
+    let mut uniform_bytes: Vec<u8> = Vec::new();
+    uniform_bytes.extend_from_slice(&b[1]);
+
+    for i in 2..(ell + 1) {
+        let mut h = H::new();
+        h.update(xor(&b[0], &b[i - 1])?);
+        h.update(&i2osp(i, 1));
+        h.update(&dst_prime);
+        b.push(h.finalize().to_vec()); // b[i]
+        uniform_bytes.extend_from_slice(&b[i]);
+    }
+
+    Ok(uniform_bytes[..len_in_bytes].to_vec())
+}
+
+#[cfg(test)]
+mod tests {
+
+    struct Params {
+        msg: &'static str,
+        len_in_bytes: usize,
+        uniform_bytes: &'static str,
+    }
+
+    #[test]
+    fn test_expand_message_xmd() {
+        // Test vectors taken from Section K.1 of https://www.ietf.org/archive/id/draft-irtf-cfrg-hash-to-curve-10.txt
+        let test_vectors: Vec<Params> = vec![
+            Params {
+                msg: "",
+                len_in_bytes: 0x20,
+                uniform_bytes: "f659819a6473c1835b25ea59e3d38914c98b374f0970b7e4c\
+                92181df928fca88",
+            },
+            Params {
+                msg: "abc",
+                len_in_bytes: 0x20,
+                uniform_bytes: "1c38f7c211ef233367b2420d04798fa4698080a8901021a79\
+                5a1151775fe4da7",
+            },
+            Params {
+                msg: "abcdef0123456789",
+                len_in_bytes: 0x20,
+                uniform_bytes: "8f7e7b66791f0da0dbb5ec7c22ec637f79758c0a48170bfb7c4611bd304ece89",
+            },
+            Params {
+                msg: "q128_qqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqq\
+                qqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqq\
+                qqqqqqqqqqqqqqqqqqqqqqqqq",
+                len_in_bytes: 0x20,
+                uniform_bytes: "72d5aa5ec810370d1f0013c0df2f1d65699494ee2a39f72e\
+                1716b1b964e1c642",
+            },
+            Params {
+                msg: "a512_aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa\
+                aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa\
+                aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa\
+                aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa\
+                aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa\
+                aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa\
+                aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa\
+                aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa\
+                aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa\
+                aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",
+                len_in_bytes: 0x20,
+                uniform_bytes: "3b8e704fc48336aca4c2a12195b720882f2162a4b7b13a9c\
+                350db46f429b771b",
+            },
+            Params {
+                msg: "",
+                len_in_bytes: 0x80,
+                uniform_bytes: "8bcffd1a3cae24cf9cd7ab85628fd111bb17e3739d3b53f8\
+                9580d217aa79526f1708354a76a402d3569d6a9d19ef3de4d0b991\
+                e4f54b9f20dcde9b95a66824cbdf6c1a963a1913d43fd7ac443a02\
+                fc5d9d8d77e2071b86ab114a9f34150954a7531da568a1ea8c7608\
+                61c0cde2005afc2c114042ee7b5848f5303f0611cf297f",
+            },
+            Params {
+                msg: "abc",
+                len_in_bytes: 0x80,
+                uniform_bytes: "fe994ec51bdaa821598047b3121c149b364b178606d5e72b\
+                fbb713933acc29c186f316baecf7ea22212f2496ef3f785a27e84a\
+                40d8b299cec56032763eceeff4c61bd1fe65ed81decafff4a31d01\
+                98619c0aa0c6c51fca15520789925e813dcfd318b542f879944127\
+                1f4db9ee3b8092a7a2e8d5b75b73e28fb1ab6b4573c192",
+            },
+            Params {
+                msg: "abcdef0123456789",
+                len_in_bytes: 0x80,
+                uniform_bytes: "c9ec7941811b1e19ce98e21db28d22259354d4d0643e3011\
+                75e2f474e030d32694e9dd5520dde93f3600d8edad94e5c3649030\
+                88a7228cc9eff685d7eaac50d5a5a8229d083b51de4ccc3733917f\
+                4b9535a819b445814890b7029b5de805bf62b33a4dc7e24acdf2c9\
+                24e9fe50d55a6b832c8c84c7f82474b34e48c6d43867be",
+            },
+            Params {
+                msg: "q128_qqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqq\
+                qqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqq\
+                qqqqqqqqqqqqqqqqqqqqqqqqq",
+                len_in_bytes: 0x80,
+                uniform_bytes: "48e256ddba722053ba462b2b93351fc966026e6d6db49318\
+                9798181c5f3feea377b5a6f1d8368d7453faef715f9aecb078cd40\
+                2cbd548c0e179c4ed1e4c7e5b048e0a39d31817b5b24f50db58bb3\
+                720fe96ba53db947842120a068816ac05c159bb5266c63658b4f00\
+                0cbf87b1209a225def8ef1dca917bcda79a1e42acd8069",
+            },
+            Params {
+                msg: "a512_aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa\
+                aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa\
+                aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa\
+                aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa\
+                aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa\
+                aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa\
+                aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa\
+                aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa\
+                aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa\
+                aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",
+                len_in_bytes: 0x80,
+                uniform_bytes: "396962db47f749ec3b5042ce2452b619607f27fd3939ece2\
+                746a7614fb83a1d097f554df3927b084e55de92c7871430d6b95c2\
+                a13896d8a33bc48587b1f66d21b128a1a8240d5b0c26dfe795a1a8\
+                42a0807bb148b77c2ef82ed4b6c9f7fcb732e7f94466c8b51e52bf\
+                378fba044a31f5cb44583a892f5969dcd73b3fa128816e",
+            },
+        ];
+        let dst = "QUUX-V01-CS02-with-expander";
+
+        for tv in test_vectors {
+            let uniform_bytes = super::expand_message_xmd::<sha2::Sha256>(
+                tv.msg.as_bytes(),
+                dst.as_bytes(),
+                tv.len_in_bytes,
+            )
+            .unwrap();
+            assert_eq!(tv.uniform_bytes, hex::encode(uniform_bytes));
+        }
     }
 }
diff --git a/src/opaque.rs b/src/opaque.rs
index e16c97d..e45211b 100644
--- a/src/opaque.rs
+++ b/src/opaque.rs
@@ -16,6 +16,7 @@ use crate::{
     hash::Hash,
     key_exchange::traits::{KeyExchange, ToBytes},
     keypair::{KeyPair, SizedBytesExt},
+    map_to_curve::GroupWithMapToCurve,
     oprf,
     serialization::{
         serialize, tokenize, u8_to_credential_type, CredentialType, ProtocolMessageType,
@@ -29,6 +30,8 @@ use std::collections::HashMap;
 use std::{convert::TryFrom, marker::PhantomData};
 use zeroize::Zeroize;
 
+static STR_OPAQUE_VERSION: &[u8] = b"OPAQUE00";
+
 // Messages
 // =========
 
@@ -1275,10 +1278,14 @@ impl<CS: CipherSuite> ServerLogin<CS> {
 }
 
 // Helper functions
-fn get_password_derived_key<G: Group, SH: SlowHash<D>, D: Hash>(
+fn get_password_derived_key<G: GroupWithMapToCurve, SH: SlowHash<D>, D: Hash>(
     token: &oprf::Token<G>,
     beta: G,
 ) -> Result<Vec<u8>, InternalPakeError> {
-    let oprf_output = oprf::unblind_and_finalize::<G, D>(token, beta)?;
+    let oprf_output = oprf::finalize::<G, D>(
+        &token.data,
+        &oprf::unblind::<G>(token, beta),
+        STR_OPAQUE_VERSION,
+    );
     SH::hash(oprf_output)
 }
diff --git a/src/oprf.rs b/src/oprf.rs
index 6dae4a5..7c41975 100644
--- a/src/oprf.rs
+++ b/src/oprf.rs
@@ -5,10 +5,10 @@
 
 use crate::{
     errors::InternalPakeError, group::Group, hash::Hash, map_to_curve::GroupWithMapToCurve,
+    serialization::serialize,
 };
 use digest::Digest;
 use generic_array::GenericArray;
-use hkdf::Hkdf;
 use rand_core::{CryptoRng, RngCore};
 
 /// Used to store the OPRF input and blinding factor
@@ -18,6 +18,8 @@ pub struct Token<Grp: Group> {
 }
 
 static STR_VOPRF: &[u8] = b"VOPRF05";
+static STR_VOPRF_FINALIZE: &[u8] = b"VOPRF05-Finalize-";
+static MODE_BASE: u8 = 0x01;
 
 /// Computes the first step for the multiplicative blinding version of DH-OPRF. This
 /// message is sent from the client (who holds the input) to the server (who holds the OPRF key).
@@ -28,7 +30,8 @@ pub(crate) fn blind<R: RngCore + CryptoRng, G: GroupWithMapToCurve>(
     blinding_factor_rng: &mut R,
     #[cfg(test)] postprocess: fn(G::Scalar) -> G::Scalar,
 ) -> Result<(Token<G>, G), InternalPakeError> {
-    let mapped_point = G::map_to_curve(input, Some(STR_VOPRF)); // TODO: add contextString from RFC
+    let dst = [STR_VOPRF, &G::get_context_string(MODE_BASE)].concat();
+    let mapped_point = G::map_to_curve(input, &dst)?;
     let blinding_factor = G::random_scalar(blinding_factor_rng);
     #[cfg(test)]
     let blind = postprocess(blinding_factor);
@@ -53,15 +56,25 @@ pub(crate) fn evaluate<G: Group>(point: G, oprf_key: &G::Scalar) -> Result<G, In
 
 /// Computes the third step for the multiplicative blinding version of DH-OPRF, in which
 /// the client unblinds the server's message.
-pub(crate) fn unblind_and_finalize<G: Group, H: Hash>(
-    token: &Token<G>,
-    point: G,
-) -> Result<GenericArray<u8, <H as Digest>::OutputSize>, InternalPakeError> {
+pub(crate) fn unblind<G: Group>(token: &Token<G>, point: G) -> Vec<u8> {
     let unblinded = point * &G::scalar_invert(&token.blind);
-    let ikm: Vec<u8> = [&unblinded.to_arr()[..], &token.data].concat();
-    // TODO: implement proper finalizing code here
-    let (prk, _) = Hkdf::<H>::extract(None, &ikm);
-    Ok(prk)
+    unblinded.to_arr().to_vec()
+}
+
+pub(crate) fn finalize<G: GroupWithMapToCurve, H: Hash>(
+    token_data: &[u8],
+    issued_token: &[u8],
+    info: &[u8],
+) -> GenericArray<u8, <H as Digest>::OutputSize> {
+    let finalize_dst = [STR_VOPRF_FINALIZE, &G::get_context_string(MODE_BASE)].concat();
+    let hash_input = [
+        serialize(token_data, 2),
+        serialize(issued_token, 2),
+        serialize(info, 2),
+        serialize(&finalize_dst, 2),
+    ]
+    .concat();
+    <H as Digest>::digest(&hash_input)
 }
 
 ////////////////////////
@@ -93,11 +106,15 @@ pub fn evaluate_shim<G: Group>(point: G, oprf_key: &G::Scalar) -> Result<G, Inte
 #[cfg(feature = "bench")]
 #[doc(hidden)]
 #[inline]
-pub fn unblind_and_finalize_shim<G: Group, H: Hash>(
+pub fn unblind_and_finalize_shim<G: GroupWithMapToCurve, H: Hash>(
     token: &Token<G>,
     point: G,
 ) -> Result<GenericArray<u8, <H as Digest>::OutputSize>, InternalPakeError> {
-    unblind_and_finalize::<G, H>(token, point)
+    Ok(finalize::<G, H>(
+        &token.data,
+        &unblind::<G>(token, point),
+        b"",
+    ))
 }
 
 ///////////
@@ -111,23 +128,19 @@ mod tests {
     use crate::group::Group;
     use curve25519_dalek::ristretto::RistrettoPoint;
     use generic_array::{arr, GenericArray};
-    use hkdf::Hkdf;
     use rand_core::OsRng;
-    use sha2::{Sha256, Sha512};
 
     fn prf(
         input: &[u8],
         oprf_key: &[u8; 32],
     ) -> GenericArray<u8, <RistrettoPoint as Group>::ElemLen> {
-        let (hashed_input, _) = Hkdf::<Sha512>::extract(Some(STR_VOPRF), &input);
-        let point = RistrettoPoint::hash_to_curve(GenericArray::from_slice(&hashed_input));
+        let dst = [STR_VOPRF, &RistrettoPoint::get_context_string(MODE_BASE)].concat();
+        let point = RistrettoPoint::map_to_curve(input, &dst).unwrap();
         let scalar =
             RistrettoPoint::from_scalar_slice(GenericArray::from_slice(&oprf_key[..])).unwrap();
         let res = point * scalar;
-        let ikm: Vec<u8> = [&res.to_arr()[..], &input].concat();
 
-        let (prk, _) = Hkdf::<Sha256>::extract(None, &ikm);
-        prk
+        finalize::<RistrettoPoint, sha2::Sha256>(&input, &res.to_arr().to_vec(), b"")
     }
 
     #[test]
@@ -142,7 +155,8 @@ mod tests {
         ];
         let oprf_key = RistrettoPoint::from_scalar_slice(&oprf_key_bytes)?;
         let beta = evaluate::<RistrettoPoint>(alpha, &oprf_key)?;
-        let res = unblind_and_finalize::<RistrettoPoint, sha2::Sha256>(&token, beta)?;
+        let res =
+            finalize::<RistrettoPoint, sha2::Sha256>(&token.data, &unblind(&token, beta), b"");
         let res2 = prf(&input[..], &oprf_key.as_bytes());
         assert_eq!(res, res2);
         Ok(())
@@ -155,18 +169,13 @@ mod tests {
         rng.fill_bytes(&mut input);
         let (token, alpha) =
             blind::<_, RistrettoPoint>(&input, &mut rng, std::convert::identity).unwrap();
-        let res = unblind_and_finalize::<RistrettoPoint, sha2::Sha256>(&token, alpha).unwrap();
+        let res =
+            finalize::<RistrettoPoint, sha2::Sha256>(&token.data, &unblind(&token, alpha), b"");
 
-        let (hashed_input, _) = Hkdf::<Sha512>::extract(Some(STR_VOPRF), &input);
-        let mut bits = [0u8; 64];
-        bits.copy_from_slice(&hashed_input);
+        let dst = [STR_VOPRF, &RistrettoPoint::get_context_string(MODE_BASE)].concat();
+        let point = RistrettoPoint::map_to_curve(&input, &dst).unwrap();
+        let res2 = finalize::<RistrettoPoint, sha2::Sha256>(&input, &point.to_arr().to_vec(), b"");
 
-        let point = RistrettoPoint::from_uniform_bytes(&bits);
-        let mut ikm: Vec<u8> = Vec::new();
-        ikm.extend_from_slice(&point.to_arr());
-        ikm.extend_from_slice(&input);
-        let (prk, _) = Hkdf::<Sha256>::extract(None, &ikm);
-
-        assert_eq!(res, prk);
+        assert_eq!(res, res2);
     }
 }
diff --git a/src/serialization/mod.rs b/src/serialization/mod.rs
index 58b71f8..bc67f9c 100644
--- a/src/serialization/mod.rs
+++ b/src/serialization/mod.rs
@@ -80,33 +80,45 @@ impl<T: CipherSuite> From<&LoginThirdMessage<T>> for ProtocolMessageType {
     }
 }
 
-pub(crate) fn serialize(input: &[u8], max_bytes: usize) -> Vec<u8> {
-    let mut output: Vec<u8> = Vec::new();
-    output
-        .extend_from_slice(&input.len().to_be_bytes()[std::mem::size_of::<usize>() - max_bytes..]);
-    output.extend_from_slice(&input[..]);
+// Corresponds to the I2OSP() function from RFC8017
+pub(crate) fn i2osp(input: usize, length: usize) -> Vec<u8> {
+    if length <= std::mem::size_of::<usize>() {
+        return (&input.to_be_bytes()[std::mem::size_of::<usize>() - length..]).to_vec();
+    }
+
+    let mut output = vec![0u8; length];
+    output.splice(
+        length - std::mem::size_of::<usize>()..length,
+        input.to_be_bytes().iter().cloned(),
+    );
     output
 }
 
-pub(crate) fn tokenize(input: Vec<u8>, size_bytes: usize) -> Result<(Vec<u8>, Vec<u8>), PakeError> {
-    if size_bytes > 8 || input.len() < size_bytes {
+// Corresponds to the OS2IP() function from RFC8017
+pub(crate) fn os2ip(input: &[u8]) -> Result<usize, PakeError> {
+    if input.len() > std::mem::size_of::<usize>() {
+        // TODO:: check RFC compliance in refusing this
         return Err(PakeError::SerializationError);
     }
 
-    let mut size_array = [0u8; 8];
-    for i in 0..size_bytes {
-        size_array[8 - size_bytes + i] = input[i];
-    }
+    let mut output_array = [0u8; std::mem::size_of::<usize>()];
+    output_array[std::mem::size_of::<usize>() - input.len()..].copy_from_slice(input);
+    Ok(usize::from_be_bytes(output_array))
+}
 
-    let big_size = u64::from_be_bytes(size_array);
+// Computes I2OSP(len(input), max_bytes) || input
+pub(crate) fn serialize(input: &[u8], max_bytes: usize) -> Vec<u8> {
+    [&i2osp(input.len(), max_bytes), &input[..]].concat()
+}
 
-    // TODO:: check RFC compliance in refusing this
-    if big_size >= u32::MAX as u64 {
+// Tokenizes an input of the format I2OSP(len(input), max_bytes) || input, outputting
+// (input, remainder)
+pub(crate) fn tokenize(input: Vec<u8>, size_bytes: usize) -> Result<(Vec<u8>, Vec<u8>), PakeError> {
+    if size_bytes > std::mem::size_of::<usize>() || input.len() < size_bytes {
         return Err(PakeError::SerializationError);
     }
 
-    let size = big_size as usize;
-
+    let size = os2ip(&input[..size_bytes])?;
     if size_bytes + size > input.len() {
         return Err(PakeError::SerializationError);
     }


Test patch:
diff --git a/src/serialization/tests.rs b/src/serialization/tests.rs
index 0023912..b8fc8be 100644
--- a/src/serialization/tests.rs
+++ b/src/serialization/tests.rs
@@ -13,7 +13,7 @@ use crate::{
     },
     keypair::{KeyPair, X25519KeyPair},
     opaque::*,
-    serialization::{serialize, ProtocolMessageType},
+    serialization::{i2osp, os2ip, serialize, ProtocolMessageType},
 };
 
 use curve25519_dalek::ristretto::RistrettoPoint;
@@ -240,7 +240,7 @@ fn login_second_message_roundtrip() {
 
     let ke2m: Vec<u8> = [&server_nonce[..], &server_e_kp.public(), &mac[..]].concat();
 
-    let total_length_without_ke2m: usize = pt_bytes.len() + envelope.to_bytes().len() + 2;
+    let total_length_without_ke2m = pt_bytes.len() + envelope.to_bytes().len() + 2;
 
     let mut input = Vec::new();
     input.extend_from_slice(&[ProtocolMessageType::CredentialResponse as u8 + 1]);
@@ -312,6 +312,11 @@ fn ke1_message_roundtrip() {
 
 proptest! {
 
+#[test]
+fn test_i2osp_os2ip(bytes in vec(any::<u8>(), 0..std::mem::size_of::<usize>())) {
+    assert_eq!(i2osp(os2ip(&bytes)?, bytes.len()), bytes);
+}
+
 #[test]
 fn test_nocrash_register_first_message(bytes in vec(any::<u8>(), 0..200)) {
     RegisterFirstMessage::<RistrettoPoint>::try_from(&bytes[..]).map_or(true, |_| true);
diff --git a/src/tests/opaque_ke_test.rs b/src/tests/opaque_ke_test.rs
index 873515f..b1f4280 100644
--- a/src/tests/opaque_ke_test.rs
+++ b/src/tests/opaque_ke_test.rs
@@ -66,35 +66,35 @@ pub struct TestVectorParameters {
 
 static TEST_VECTOR: &str = r#"
 {
-    "client_s_pk": "e7dce23c92339bf0719c2faaa25d9b32997501316552241661331ddb2d621d0f",
-    "client_s_sk": "88a85d87f7a6afbcd26c3a44e9c6d6e8b2e66f5f8bc50ab2043e499088bb6c72",
-    "client_e_pk": "a05e21d85c342f5c53f9a048d63a9a04f9d1cef364011963dcdb66dc76c5a752",
-    "client_e_sk": "80de26e3c0fd87c860030bb600b0a8b33eba54c02e243b4584137da648c5847a",
-    "server_s_pk": "1766e2000c0914af7431de8c7ae8757eb2c974006dc7d9fe037cbb1569665a64",
-    "server_s_sk": "d86b6db7d70be0a6155e5f4f1a1047b55a0fde382115e4ad524b34edf3bf4b60",
-    "server_e_pk": "0a7b46933a6a86942f89eadac7bda1e513ce97776f4e7e65127f2685bc3f250b",
-    "server_e_sk": "483b9af7dbce61f734bc824acc892ff5a3c655f29e0ba129eb24ae3ae944a56d",
+    "client_s_pk": "bf958dfae87e144b9ef7e5599991ef5b3c057ce9bd0120d9812d56997ace575f",
+    "client_s_sk": "d028843af0a5d410ee5a560ddda7e7c9c313f7793e54a41708a530e67de26677",
+    "client_e_pk": "d3704416b573b0a3605a88971e144a1edcb5d981855fe54a77cdff02bf6e6457",
+    "client_e_sk": "7833dbb6853ff6ee3f9a69dad738c42de6cdf7c1ae7491248122b8169383ae4c",
+    "server_s_pk": "c82312c965f9e08fda7aca33287beb2f1c39a0afbc6033ab22fa4ba3d4fb2d3f",
+    "server_s_sk": "3829629dc84a4a0fd3d04e225e108d2947668564326c3c45e80bc52c5fd8b16b",
+    "server_e_pk": "415c2c4b1769268b1c7fc703cf652ce2df8249709d6d1b8b2608fff1e744e22b",
+    "server_e_sk": "a02800350e22a75a294d2365cb54f5c6ae6461514ce66b2d2b418965949b514d",
     "id_u": "696455",
     "id_s": "696453",
     "password": "70617373776f7264",
-    "blinding_factor": "803739feafc45969f9be5a6ab6ea0cbc8c6ee5144bc7ce83b6f8acb77e160c04",
-    "oprf_key": "288b3f25582a4d5f1a775442805a44b77266430e90124e288176d56eaadabd05",
-    "envelope_nonce": "6ed65068854948760f5e23042e1f64614382828206ac24389addb709beb2df5d",
-    "client_nonce": "1642c14481e4d3822ca5e038696f86b133c2a3c28bc0c4ee2f6e10305443c5dc",
-    "server_nonce": "d7d2f3b5257e67742d8449cf66eca9b9ea89dc9aa811e71806ebf309e6e27a87",
-    "r1": "01000022002074b8033d642dfcf0cec75df539801d3fe9262919bd8fd2cd742c3ec211522ffe",
-    "r2": "02000028002048e74c2b17750b3875d80688c9217c10df421dbcec56f35f79cccc1085f3f2a8000001010103",
-    "r3": "030000ae6ed65068854948760f5e23042e1f64614382828206ac24389addb709beb2df5d00235c8cbe660578c0a07a0be2ffad5c638a9d80dfab9323c3f2f51b15100bdf6a5116506b00230300201766e2000c0914af7431de8c7ae8757eb2c974006dc7d9fe037cbb1569665a64002008f76f5e11b480be387407b920479273fc7e039edba805b55aa1e2d72965796a0020e7dce23c92339bf0719c2faaa25d9b32997501316552241661331ddb2d621d0f",
-    "l1": "04000022002074b8033d642dfcf0cec75df539801d3fe9262919bd8fd2cd742c3ec211522ffe1642c14481e4d3822ca5e038696f86b133c2a3c28bc0c4ee2f6e10305443c5dca05e21d85c342f5c53f9a048d63a9a04f9d1cef364011963dcdb66dc76c5a752",
-    "l2": "050000ae002048e74c2b17750b3875d80688c9217c10df421dbcec56f35f79cccc1085f3f2a86ed65068854948760f5e23042e1f64614382828206ac24389addb709beb2df5d00235c8cbe660578c0a07a0be2ffad5c638a9d80dfab9323c3f2f51b15100bdf6a5116506b00230300201766e2000c0914af7431de8c7ae8757eb2c974006dc7d9fe037cbb1569665a64002008f76f5e11b480be387407b920479273fc7e039edba805b55aa1e2d72965796a483b9af7dbce61f734bc824acc892ff5a3c655f29e0ba129eb24ae3ae944a56d0a7b46933a6a86942f89eadac7bda1e513ce97776f4e7e65127f2685bc3f250b4968264f94572d261ce8a3df0f83ecd0e022c6ca0c3ee6ba54a8fc5f283ca698",
-    "l3": "30b74db9252e8e2219bb2a0d0708a10b6eac6fee0c2c8ae754062e993ecd18d4",
-    "client_registration_state": "00036964550003696453803739feafc45969f9be5a6ab6ea0cbc8c6ee5144bc7ce83b6f8acb77e160c0470617373776f7264",
-    "client_login_state": "00036964550003696453803739feafc45969f9be5a6ab6ea0cbc8c6ee5144bc7ce83b6f8acb77e160c0480de26e3c0fd87c860030bb600b0a8b33eba54c02e243b4584137da648c5847a1642c14481e4d3822ca5e038696f86b133c2a3c28bc0c4ee2f6e10305443c5dcbafb8dcb23578e52a981d66816d55659b49caf268dce519d190b42571ff031b570617373776f7264",
-    "server_registration_state": "288b3f25582a4d5f1a775442805a44b77266430e90124e288176d56eaadabd05",
-    "server_login_state": "94a3839f257cac066c706a4086516f9bf3f99108309f0d7e3d8423cb716ad6544682fc2f3d3e719931c16f6b4befe0438b8daaff99bd21ce2306ea23882f208d7e482d1a199d47861c103f6d0859e4b9a13375a97901f9f1c9d75bbaf52f43b1",
-    "password_file": "288b3f25582a4d5f1a775442805a44b77266430e90124e288176d56eaadabd05e7dce23c92339bf0719c2faaa25d9b32997501316552241661331ddb2d621d0f6ed65068854948760f5e23042e1f64614382828206ac24389addb709beb2df5d00235c8cbe660578c0a07a0be2ffad5c638a9d80dfab9323c3f2f51b15100bdf6a5116506b00230300201766e2000c0914af7431de8c7ae8757eb2c974006dc7d9fe037cbb1569665a64002008f76f5e11b480be387407b920479273fc7e039edba805b55aa1e2d72965796a",
-    "export_key": "e7a3827bf6f1404ecb210fde8bab2367daa0ae206b959f981816cef02f931d81",
-    "shared_secret": "7e482d1a199d47861c103f6d0859e4b9a13375a97901f9f1c9d75bbaf52f43b1"
+    "blinding_factor": "5feb5851ccacfe610bef7bb970c05fb2f47381f6423fea4436248c6a77ac7a08",
+    "oprf_key": "305500023999e82e0ac5f34c7d13dea3f88b9897cd91b6206f1df8ee2d856f0f",
+    "envelope_nonce": "c2d0720e9987bd4380068a6a23480bb1fa0eea01d84c71afb580f23fb09acae8",
+    "client_nonce": "ffaf9633c8dadc578524fa0993098fb0d30bb9ed092b0d5b2a75bf24381553c8",
+    "server_nonce": "441df2d948b11acbbe6e3ffcbc5fd6e2ff115176f5b16878d7c198371cff0927",
+    "r1": "010000220020b54de93fd9c9cadbfba797cba1e8a1469cbd673ed0fdfd2a6fda644d85664753",
+    "r2": "020000280020a08e6edb2dfb4cc05c9c5530e7ef07532db317b959825546c49eed4a5b7c40a6000001010103",
+    "r3": "030000aec2d0720e9987bd4380068a6a23480bb1fa0eea01d84c71afb580f23fb09acae80023306c65ae3f96ac0a892ea5843e0debe597cde1dbf2561c18c495839e2bc9e887fd10750023030020c82312c965f9e08fda7aca33287beb2f1c39a0afbc6033ab22fa4ba3d4fb2d3f0020251cd488033f5c5a7a85548e0b5c2424a7e0c331e11f51c6de67b697c3f9321c0020bf958dfae87e144b9ef7e5599991ef5b3c057ce9bd0120d9812d56997ace575f",
+    "l1": "040000220020b54de93fd9c9cadbfba797cba1e8a1469cbd673ed0fdfd2a6fda644d85664753ffaf9633c8dadc578524fa0993098fb0d30bb9ed092b0d5b2a75bf24381553c8d3704416b573b0a3605a88971e144a1edcb5d981855fe54a77cdff02bf6e6457",
+    "l2": "050000ae0020a08e6edb2dfb4cc05c9c5530e7ef07532db317b959825546c49eed4a5b7c40a6c2d0720e9987bd4380068a6a23480bb1fa0eea01d84c71afb580f23fb09acae80023306c65ae3f96ac0a892ea5843e0debe597cde1dbf2561c18c495839e2bc9e887fd10750023030020c82312c965f9e08fda7aca33287beb2f1c39a0afbc6033ab22fa4ba3d4fb2d3f0020251cd488033f5c5a7a85548e0b5c2424a7e0c331e11f51c6de67b697c3f9321ca02800350e22a75a294d2365cb54f5c6ae6461514ce66b2d2b418965949b514d415c2c4b1769268b1c7fc703cf652ce2df8249709d6d1b8b2608fff1e744e22b79264b111c4ab164204f16d3f2da5a2cba3c736607f01906e60b3f70c887c1be",
+    "l3": "ce39e86b59b96ae8497909a8b4a67f9cade11badbe7c95392a44ba9ca31362af",
+    "client_registration_state": "000369645500036964535feb5851ccacfe610bef7bb970c05fb2f47381f6423fea4436248c6a77ac7a0870617373776f7264",
+    "client_login_state": "000369645500036964535feb5851ccacfe610bef7bb970c05fb2f47381f6423fea4436248c6a77ac7a087833dbb6853ff6ee3f9a69dad738c42de6cdf7c1ae7491248122b8169383ae4cffaf9633c8dadc578524fa0993098fb0d30bb9ed092b0d5b2a75bf24381553c812d4a730af1022f304f21cfe4cdcf98c0494da4c7c86f45e94b6449713990a9870617373776f7264",
+    "server_registration_state": "305500023999e82e0ac5f34c7d13dea3f88b9897cd91b6206f1df8ee2d856f0f",
+    "server_login_state": "c5b50039ee9c716d4bbbbae949322bdeaa23ab8e8b176832884b467f990ce55af2ee66df1a0b63ee7a3ae343231af7f2c64339238ec3cd0e261cc95727e47fd849bb17adafce3161fa1787d622f27bfeffa947d2bdd128c90e11c6aa4c04c390",
+    "password_file": "305500023999e82e0ac5f34c7d13dea3f88b9897cd91b6206f1df8ee2d856f0fbf958dfae87e144b9ef7e5599991ef5b3c057ce9bd0120d9812d56997ace575fc2d0720e9987bd4380068a6a23480bb1fa0eea01d84c71afb580f23fb09acae80023306c65ae3f96ac0a892ea5843e0debe597cde1dbf2561c18c495839e2bc9e887fd10750023030020c82312c965f9e08fda7aca33287beb2f1c39a0afbc6033ab22fa4ba3d4fb2d3f0020251cd488033f5c5a7a85548e0b5c2424a7e0c331e11f51c6de67b697c3f9321c",
+    "export_key": "2ae428e07164253279df1193a9904bb0b9419750142109818e594f537bd0b854",
+    "shared_secret": "49bb17adafce3161fa1787d622f27bfeffa947d2bdd128c90e11c6aa4c04c390"
 }
 "#;