feat: add collision resistance for truncated key id (#45)

raphaelrobert · web-flow · commit c6be7b0097e5 · 2025-11-30T23:36:07.000+01:00
* allow truncated key id collisions for public tokens

* avoid collisions when creating keys

* address review comments
diff --git a/src/amortized_tokens/server.rs b/src/amortized_tokens/server.rs
@@ -7,7 +7,7 @@ use typenum::Unsigned;
 use voprf::{BlindedElement, Group, Result, VoprfServer, VoprfServerBatchEvaluateFinishResult};
 
 use crate::{
-    NonceStore, TokenInput,
+    COLLISION_AVOIDANCE_ATTEMPTS, NonceStore, TokenInput,
     common::{
         errors::{CreateKeypairError, IssueTokenResponseError, RedeemTokenError},
         private::{PrivateCipherSuite, PublicKey, public_key_to_token_key_id},
@@ -25,6 +25,14 @@ pub struct Server<CS: PrivateCipherSuite> {
 }
 
 impl<CS: PrivateCipherSuite> Server<CS> {
+    fn server_from_seed(seed: &[u8], info: &[u8]) -> Result<VoprfServer<CS>, CreateKeypairError>
+    where
+        <CS::Group as Group>::Scalar: Send + Sync,
+        <CS::Group as Group>::Elem: Send + Sync,
+    {
+        VoprfServer::<CS>::new_from_seed(seed, info).map_err(|_| CreateKeypairError::SeedError)
+    }
+
     /// Create a new server. The new server does not contain any key material.
     #[must_use]
     pub const fn new() -> Self {
@@ -45,14 +53,29 @@ impl<CS: PrivateCipherSuite> Server<CS> {
         <CS::Group as Group>::Scalar: Send + Sync,
         <CS::Group as Group>::Elem: Send + Sync,
     {
-        let mut seed = GenericArray::<_, <CS::Group as Group>::ScalarLen>::default();
-        OsRng.fill_bytes(&mut seed);
-        self.create_keypair_internal(key_store, &seed, b"PrivacyPass")
-            .await
+        for _ in 0..COLLISION_AVOIDANCE_ATTEMPTS {
+            let mut seed = GenericArray::<_, <CS::Group as Group>::ScalarLen>::default();
+            OsRng.fill_bytes(&mut seed);
+            let server = Self::server_from_seed(&seed, b"PrivacyPass")?;
+            let public_key = server.get_public_key();
+            let truncated_token_key_id =
+                truncate_token_key_id(&public_key_to_token_key_id::<CS>(&public_key));
+
+            if key_store.get(&truncated_token_key_id).await.is_some() {
+                continue;
+            }
+
+            if key_store.insert(truncated_token_key_id, server).await {
+                return Ok(public_key);
+            }
+        }
+        Err(CreateKeypairError::CollisionExhausted)
     }
 
-    /// Creates a new keypair and inserts it into the key store.
-    async fn create_keypair_internal<BKS: PrivateKeyStore<CS = CS>>(
+    /// Creates a new keypair with explicit parameters and inserts it into the
+    /// key store.
+    #[cfg(feature = "kat")]
+    pub async fn create_keypair_with_params<BKS: PrivateKeyStore<CS = CS>>(
         &self,
         key_store: &BKS,
         seed: &[u8],
@@ -62,31 +85,14 @@ impl<CS: PrivateCipherSuite> Server<CS> {
         <CS::Group as Group>::Scalar: Send + Sync,
         <CS::Group as Group>::Elem: Send + Sync,
     {
-        let server = VoprfServer::<CS>::new_from_seed(seed, info)
-            .map_err(|_| CreateKeypairError::SeedError)?;
+        let server = Self::server_from_seed(seed, info)?;
         let public_key = server.get_public_key();
         let truncated_token_key_id =
             truncate_token_key_id(&public_key_to_token_key_id::<CS>(&server.get_public_key()));
         key_store.insert(truncated_token_key_id, server).await;
         Ok(public_key)
     }
 
-    /// Creates a new keypair with explicit parameters and inserts it into the
-    /// key store.
-    #[cfg(feature = "kat")]
-    pub async fn create_keypair_with_params<BKS: PrivateKeyStore<CS = CS>>(
-        &self,
-        key_store: &BKS,
-        seed: &[u8],
-        info: &[u8],
-    ) -> Result<PublicKey<CS>, CreateKeypairError>
-    where
-        <CS::Group as Group>::Scalar: Send + Sync,
-        <CS::Group as Group>::Elem: Send + Sync,
-    {
-        self.create_keypair_internal(key_store, seed, info).await
-    }
-
     /// Issues a token response.
     ///
     /// # Errors
diff --git a/src/common/errors.rs b/src/common/errors.rs
@@ -13,9 +13,12 @@ pub enum SerializationError {
 /// Errors that can occur when creating a keypair.
 #[derive(Error, Debug, PartialEq, Eq)]
 pub enum CreateKeypairError {
-    #[error("Seed is too long")]
-    /// Error when the seed is too long.
+    #[error("Seed does not have the right length")]
+    /// Error when the seed does not have the right length
     SeedError,
+    #[error("Collision exhausted")]
+    /// Error when collision attempts are exhausted
+    CollisionExhausted,
 }
 
 /// Errors that can occur when issuing token requests.
diff --git a/src/common/store.rs b/src/common/store.rs
@@ -12,12 +12,16 @@ use super::private::PrivateCipherSuite;
 pub trait PrivateKeyStore {
     /// The cipher suite used for the key store.
     type CS: PrivateCipherSuite;
-    /// Inserts a keypair with a given `truncated_token_key_id` into the key store.
+    /// Inserts a keypair with a given `truncated_token_key_id` into the key
+    /// store, only if it does not collide with an existing
+    /// `truncated_token_key_id`.
+    ///
+    /// Returns `true` if the key was inserted, `false` if a collision occurred.
     async fn insert(
         &self,
         truncated_token_key_id: TruncatedTokenKeyId,
         server: VoprfServer<Self::CS>,
-    );
+    ) -> bool;
     /// Returns a keypair with a given `truncated_token_key_id` from the key store.
     async fn get(
         &self,
diff --git a/src/lib.rs b/src/lib.rs
@@ -101,3 +101,5 @@ impl TokenInput {
         token_input
     }
 }
+
+pub(crate) const COLLISION_AVOIDANCE_ATTEMPTS: usize = 100;
diff --git a/src/private_tokens/server.rs b/src/private_tokens/server.rs
@@ -1,13 +1,13 @@
 //! Server-side implementation of Privately Verifiable Token protocol.
 
-use generic_array::ArrayLength;
+use generic_array::{ArrayLength, GenericArray};
 use rand::{RngCore, rngs::OsRng};
 use sha2::digest::OutputSizeUser;
 use typenum::Unsigned;
 use voprf::{BlindedElement, Group, Result, VoprfServer};
 
 use crate::{
-    NonceStore, TokenInput,
+    COLLISION_AVOIDANCE_ATTEMPTS, NonceStore, TokenInput,
     auth::authorize::Token,
     common::{
         errors::{CreateKeypairError, IssueTokenResponseError, RedeemTokenError},
@@ -26,6 +26,10 @@ pub struct Server<CS: PrivateCipherSuite> {
 }
 
 impl<CS: PrivateCipherSuite> Server<CS> {
+    fn server_from_seed(seed: &[u8], info: &[u8]) -> Result<VoprfServer<CS>, CreateKeypairError> {
+        VoprfServer::<CS>::new_from_seed(seed, info).map_err(|_| CreateKeypairError::SeedError)
+    }
+
     /// Creates a new server.
     #[must_use]
     pub const fn new() -> Self {
@@ -42,26 +46,23 @@ impl<CS: PrivateCipherSuite> Server<CS> {
         &self,
         key_store: &PKS,
     ) -> Result<PublicKey<CS>, CreateKeypairError> {
-        let mut seed = vec![0u8; <<CS::Group as Group>::ScalarLen as Unsigned>::USIZE];
-        OsRng.fill_bytes(&mut seed);
-        self.create_keypair_internal(key_store, &seed, b"PrivacyPass")
-            .await
-    }
+        for _ in 0..COLLISION_AVOIDANCE_ATTEMPTS {
+            let mut seed = GenericArray::<_, <CS::Group as Group>::ScalarLen>::default();
+            OsRng.fill_bytes(&mut seed);
+            let server = Self::server_from_seed(&seed, b"PrivacyPass")?;
+            let public_key = server.get_public_key();
+            let truncated_token_key_id =
+                truncate_token_key_id(&public_key_to_token_key_id::<CS>(&public_key));
 
-    /// Creates a new keypair and inserts it into the key store.
-    async fn create_keypair_internal<PKS: PrivateKeyStore<CS = CS>>(
-        &self,
-        key_store: &PKS,
-        seed: &[u8],
-        info: &[u8],
-    ) -> Result<PublicKey<CS>, CreateKeypairError> {
-        let server = VoprfServer::<CS>::new_from_seed(seed, info)
-            .map_err(|_| CreateKeypairError::SeedError)?;
-        let public_key = server.get_public_key();
-        let truncated_token_key_id =
-            truncate_token_key_id(&public_key_to_token_key_id::<CS>(&server.get_public_key()));
-        key_store.insert(truncated_token_key_id, server).await;
-        Ok(public_key)
+            if key_store.get(&truncated_token_key_id).await.is_some() {
+                continue;
+            }
+
+            if key_store.insert(truncated_token_key_id, server).await {
+                return Ok(public_key);
+            }
+        }
+        Err(CreateKeypairError::CollisionExhausted)
     }
 
     /// Creates a new keypair with explicit parameters and inserts it into the
@@ -73,7 +74,12 @@ impl<CS: PrivateCipherSuite> Server<CS> {
         seed: &[u8],
         info: &[u8],
     ) -> Result<PublicKey<CS>, CreateKeypairError> {
-        self.create_keypair_internal(key_store, seed, info).await
+        let server = Self::server_from_seed(seed, info)?;
+        let public_key = server.get_public_key();
+        let truncated_token_key_id =
+            truncate_token_key_id(&public_key_to_token_key_id::<CS>(&server.get_public_key()));
+        key_store.insert(truncated_token_key_id, server).await;
+        Ok(public_key)
     }
 
     /// Issues a token response.
diff --git a/src/public_tokens/server.rs b/src/public_tokens/server.rs
@@ -6,7 +6,7 @@ use generic_array::ArrayLength;
 use rand::{CryptoRng, RngCore, rngs::OsRng};
 
 use crate::{
-    NonceStore, TokenInput, TokenType, TruncatedTokenKeyId,
+    COLLISION_AVOIDANCE_ATTEMPTS, NonceStore, TokenInput, TokenType, TruncatedTokenKeyId,
     auth::authorize::Token,
     common::errors::{CreateKeypairError, IssueTokenResponseError, RedeemTokenError},
 };
@@ -18,8 +18,12 @@ use super::{NK, TokenRequest, TokenResponse, public_key_to_token_key_id, truncat
 #[async_trait]
 
 pub trait IssuerKeyStore: Send + Sync {
-    /// Inserts a keypair with a given `truncated_token_key_id` into the key store.
-    async fn insert(&self, truncated_token_key_id: TruncatedTokenKeyId, server: KeyPair);
+    /// Inserts a keypair with a given `truncated_token_key_id` into the key
+    /// store, only if it does not collide with an existing
+    /// `truncated_token_key_id`.
+    ///
+    /// Returns `true` if the key was inserted, `false` if a collision occurred.
+    async fn insert(&self, truncated_token_key_id: TruncatedTokenKeyId, server: KeyPair) -> bool;
     /// Returns a keypair with a given `truncated_token_key_id` from the key store.
     async fn get(&self, truncated_token_key_id: &TruncatedTokenKeyId) -> Option<KeyPair>;
 }
@@ -30,8 +34,8 @@ pub trait IssuerKeyStore: Send + Sync {
 pub trait OriginKeyStore {
     /// Inserts a keypair with a given `truncated_token_key_id` into the key store.
     async fn insert(&self, truncated_token_key_id: TruncatedTokenKeyId, server: PublicKey);
-    /// Returns a keypair with a given `truncated_token_key_id` from the key store.
-    async fn get(&self, truncated_token_key_id: &TruncatedTokenKeyId) -> Option<PublicKey>;
+    /// Returns all public keys with a given `truncated_token_key_id` from the key store.
+    async fn get(&self, truncated_token_key_id: &TruncatedTokenKeyId) -> Vec<PublicKey>;
 }
 
 /// Serializes a keypair into a DER-encoded PKCS#8 document.
@@ -64,14 +68,23 @@ impl IssuerServer {
         rng: &mut R,
         key_store: &IKS,
     ) -> Result<PublicKey, CreateKeypairError> {
-        let key_pair =
-            KeyPair::generate(rng, KEYSIZE_IN_BITS).map_err(|_| CreateKeypairError::SeedError)?;
-        let truncated_token_key_id =
-            truncate_token_key_id(&public_key_to_token_key_id(&key_pair.pk));
-        key_store
-            .insert(truncated_token_key_id, key_pair.clone())
-            .await;
-        Ok(key_pair.pk)
+        for _ in 0..COLLISION_AVOIDANCE_ATTEMPTS {
+            let key_pair = KeyPair::generate(rng, KEYSIZE_IN_BITS)
+                .map_err(|_| CreateKeypairError::SeedError)?;
+            let truncated_token_key_id =
+                truncate_token_key_id(&public_key_to_token_key_id(&key_pair.pk));
+
+            if key_store.get(&truncated_token_key_id).await.is_some() {
+                continue;
+            }
+
+            let public_key = key_pair.pk.clone();
+
+            if key_store.insert(truncated_token_key_id, key_pair).await {
+                return Ok(public_key);
+            }
+        }
+        Err(CreateKeypairError::CollisionExhausted)
     }
 
     /// Issues a new token response.
@@ -152,17 +165,26 @@ impl OriginServer {
             *token.token_key_id(),
         );
 
-        let public_key = key_store
-            .get(&truncate_token_key_id(token.token_key_id()))
-            .await
-            .ok_or(RedeemTokenError::KeyIdNotFound)?;
+        let truncated_token_key_id = truncate_token_key_id(token.token_key_id());
+        let public_keys = key_store.get(&truncated_token_key_id).await;
+        if public_keys.is_empty() {
+            return Err(RedeemTokenError::KeyIdNotFound);
+        }
 
         let options = Options::default();
         let signature = Signature(token.authenticator().to_vec());
+        let token_input_bytes = token_input.serialize();
+
+        let verified = public_keys.iter().any(|public_key| {
+            signature
+                .verify(public_key, None, &token_input_bytes, &options)
+                .is_ok()
+        });
+
+        if !verified {
+            return Err(RedeemTokenError::InvalidToken);
+        }
 
-        signature
-            .verify(&public_key, None, token_input.serialize(), &options)
-            .map_err(|_| RedeemTokenError::InvalidToken)?;
         nonce_store.insert(token.nonce()).await;
         Ok(())
     }
diff --git a/src/test_utils/private_memory_store.rs b/src/test_utils/private_memory_store.rs
@@ -1,6 +1,9 @@
 //! This module contains in-memory implementations of the `PrivateKeyStore` trait.
 use async_trait::async_trait;
-use std::{collections::HashMap, fmt::Debug};
+use std::{
+    collections::{HashMap, hash_map::Entry},
+    fmt::Debug,
+};
 use tokio::sync::Mutex;
 use voprf::*;
 
@@ -18,9 +21,18 @@ pub struct MemoryKeyStoreVoprf<CS: PrivateCipherSuite> {
 impl<C: PrivateCipherSuite> PrivateKeyStore for MemoryKeyStoreVoprf<C> {
     type CS = C;
 
-    async fn insert(&self, truncated_token_key_id: TruncatedTokenKeyId, server: VoprfServer<C>) {
+    async fn insert(
+        &self,
+        truncated_token_key_id: TruncatedTokenKeyId,
+        server: VoprfServer<C>,
+    ) -> bool {
         let mut keys = self.keys.lock().await;
-        keys.insert(truncated_token_key_id, server.clone());
+        if let Entry::Vacant(e) = keys.entry(truncated_token_key_id) {
+            e.insert(server);
+            true
+        } else {
+            false
+        }
     }
 
     async fn get(&self, truncated_token_key_id: &TruncatedTokenKeyId) -> Option<VoprfServer<C>> {
diff --git a/src/test_utils/public_memory_store.rs b/src/test_utils/public_memory_store.rs
diff --git a/tests/public_tokens.rs b/tests/public_tokens.rs

Original file line number	Diff line number	Diff line change
`@@ -101,3 +101,5 @@ impl TokenInput {`
`101`	`101`	`token_input`
`102`	`102`	`}`
`103`	`103`	`}`
	`104`	`+`
	`105`	`+pub(crate) const COLLISION_AVOIDANCE_ATTEMPTS: usize = 100;`