[key server] Misc clean ups (#415)

Author: jonas-lj

crates/crypto/src/lib.rs

@@ -132,11 +132,7 @@ pub fn seal_encrypt(
         indices, shares, ..
     } = split(&mut rng, base_key, threshold, number_of_shares)?;
 
-    let services = key_servers
-        .iter()
-        .zip(indices)
-        .map(|(s, i)| (*s, i))
-        .collect::<Vec<_>>();
+    let services = key_servers.iter().cloned().zip(indices).collect::<Vec<_>>();
 
     let encrypted_shares = match public_keys {
         IBEPublicKeys::BonehFranklinBLS12381(pks) => {

crates/crypto/src/tss.rs

@@ -114,37 +114,6 @@ pub fn combine<const N: usize>(shares: &[(u8, [u8; N])]) -> FastCryptoResult<[u8
         .expect("fixed length"))
 }
 
-pub fn split_with_given_shares<const N: usize>(
-    given_shares: &[[u8; N]],
-    number_of_shares: u8,
-) -> FastCryptoResult<SecretSharing<N>> {
-    let threshold = given_shares.len();
-    if threshold > number_of_shares as usize || threshold == 0 {
-        return Err(InvalidInput);
-    }
-
-    let indices = (1..=number_of_shares).collect_vec();
-
-    // Share each byte of the secret individually.
-    let (secret, byte_shares): (Vec<u8>, Vec<Vec<u8>>) = (0..N)
-        .map(|i| {
-            split_byte_with_given_shares(&given_shares.iter().map(|s| s[i]).collect_vec(), &indices)
-        })
-        .collect::<FastCryptoResult<Vec<_>>>()?
-        .into_iter()
-        .unzip();
-
-    // Combine the byte shares into shares.
-    let shares = transpose(&byte_shares)?;
-    let secret = secret.try_into().expect("fixed length");
-
-    Ok(SecretSharing {
-        secret,
-        indices,
-        shares,
-    })
-}
-
 /// Internal function to share a secret.
 /// This is an implementation of Shamir's secret sharing over the Galois field of 256 elements.
 /// See https://dl.acm.org/doi/10.1145/359168.359176.
@@ -175,44 +144,6 @@ fn split_byte<R: AllowedRng>(
         .collect())
 }
 
-/// Create a secret sharing of `num_shares` shares such that at least `threshold` shares are needed
-/// to reconstruct the byte and such that the first `threshold` shares will be the given ones.
-///
-/// The shared secret will be determined by the given shares, and the process is deterministic.
-///
-/// Returns the secret and a vector of the shares.
-fn split_byte_with_given_shares(
-    given_shares: &[u8],
-    indices: &[u8],
-) -> FastCryptoResult<(u8, Vec<u8>)> {
-    let number_of_shares = indices.len();
-    let threshold = given_shares.len() + 1;
-    assert!(threshold <= number_of_shares && number_of_shares <= 255 && threshold > 0);
-    assert!(indices.iter().all(|&i| i != 0) && indices.iter().all_unique());
-
-    // Construct the polynomial that interpolates the given shares and the secret.
-    let polynomial = Polynomial::interpolate(
-        &indices
-            .iter()
-            .zip(given_shares)
-            .map(|(&x, &y)| (x.into(), y.into()))
-            .collect_vec(),
-    );
-
-    // The secret is the constant term of the polynomial.
-    let secret = polynomial.0[0].0;
-
-    // Evaluate the polynomial at the remaining indices to get the remaining shares.
-    let remaining_shares = indices[given_shares.len()..]
-        .iter()
-        .map(|i| polynomial.evaluate(&i.into()).0)
-        .collect();
-
-    let shares = [given_shares.to_vec(), remaining_shares].concat();
-
-    Ok((secret, shares))
-}
-
 /// Internal function to reconstruct a secret.
 /// This is an implementation of Shamir's secret sharing over the Galois field of 256 elements.
 /// See https://dl.acm.org/doi/10.1145/359168.359176.
@@ -324,44 +255,4 @@ mod tests {
 
         assert_ne!(combine(&shares[..1]).unwrap(), expected);
     }
-
-    #[test]
-    fn test_split_byte_with_given_shares() {
-        let given_shares = [5, 19];
-        let indices = [1, 2, 3, 4, 5];
-
-        let (secret, shares) = split_byte_with_given_shares(&given_shares, &indices).unwrap();
-
-        let reconstructed = combine_byte(&[
-            (indices[0], shares[0]),
-            (indices[2], shares[2]),
-            (indices[4], shares[4]),
-        ])
-        .unwrap();
-        assert_eq!(reconstructed, secret);
-    }
-
-    #[test]
-    fn test_with_given_shares() {
-        let given_shares = [
-            *b"AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA",
-            *b"BBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBB",
-            *b"CCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCC",
-        ];
-        let threshold = given_shares.len() as u8;
-        let SecretSharing {
-            secret,
-            indices,
-            shares,
-        } = split_with_given_shares(&given_shares, 5).unwrap();
-
-        assert_eq!(threshold, given_shares.len() as u8);
-        assert_eq!(shares[0], given_shares[0]);
-        assert_eq!(shares[1], given_shares[1]);
-
-        assert_eq!(
-            secret,
-            combine(&(1..4).map(|i| (indices[i], shares[i])).collect_vec()).unwrap()
-        );
-    }
 }

crates/key-server/src/master_keys.rs

@@ -230,6 +230,7 @@ impl MasterKeys {
                 .ok_or(InternalError::InvalidServiceId),
         }
     }
+
     /// Load committee version to determine which master share to use.
     pub(crate) fn get_committee_server_master_share(
         &self,

crates/key-server/src/metrics.rs

@@ -171,11 +171,11 @@ pub(crate) fn call_with_duration<T>(metrics: Option<&Histogram>, closure: impl F
 pub(crate) fn status_callback(metrics: &IntCounterVec) -> impl Fn(bool) + use<> {
     let metrics = metrics.clone();
     move |status: bool| {
-        let value = match status {
+        let label = match status {
             true => "success",
             false => "failure",
         };
-        metrics.with_label_values(&[value]).inc();
+        metrics.with_label_values(&[label]).inc();
     }
 }
 

crates/key-server/src/metrics_push.rs

@@ -108,10 +108,10 @@ pub async fn push_metrics(
 
     if !response.status().is_success() {
         let status = response.status();
-        let body = match response.text().await {
-            Ok(body) => body,
-            Err(error) => format!("couldn't decode response body; {error}"),
-        };
+        let body = response
+            .text()
+            .await
+            .unwrap_or_else(|error| format!("couldn't decode response body; {error}"));
         return Err(anyhow::anyhow!(
             "metrics push failed: [{}]:{}",
             status,

crates/key-server/src/server.rs

@@ -554,7 +554,7 @@ impl Server {
     fn create_response(
         &self,
         first_pkg_id: ObjectID,
-        ids: &[KeyId],
+        ids: Vec<KeyId>,
         enc_key: &ElGamalPublicKey,
     ) -> FetchKeyResponse {
         debug!(
@@ -566,16 +566,13 @@ impl Server {
             .get_key_for_package(&first_pkg_id)
             .expect("checked already");
         let decryption_keys = ids
-            .iter()
+            .into_iter()
             .map(|id| {
                 // Requested key
-                let key = ibe::extract(master_key, id);
+                let key = ibe::extract(master_key, &id);
                 // ElGamal encryption of key under the user's public key
                 let encrypted_key = encrypt(&mut thread_rng(), &key, enc_key);
-                DecryptionKey {
-                    id: id.to_owned(),
-                    encrypted_key,
-                }
+                DecryptionKey { id, encrypted_key }
             })
             .collect();
         FetchKeyResponse { decryption_keys }
@@ -797,7 +794,7 @@ async fn handle_fetch_key(
             Json(
                 app_state
                     .server
-                    .create_response(first_pkg_id, &full_ids, &payload.enc_key),
+                    .create_response(first_pkg_id, full_ids, &payload.enc_key),
             )
         })
 }

crates/key-server/src/tests/externals.rs

@@ -80,7 +80,7 @@ pub(crate) async fn get_key(
         .map(|(pkg_id, ids)| {
             elgamal::decrypt(
                 &sk,
-                &server.create_response(pkg_id, &ids, &pk).decryption_keys[0].encrypted_key,
+                &server.create_response(pkg_id, ids, &pk).decryption_keys[0].encrypted_key,
             )
         })
 }

crates/key-server/src/utils.rs

@@ -29,11 +29,11 @@ macro_rules! git_version {
 use crate::types::IbeMasterKey;
 use anyhow::anyhow;
 use crypto::ibe::MASTER_KEY_LENGTH;
-use fastcrypto::encoding::Encoding;
+use fastcrypto::encoding::{Encoding, Hex};
 use fastcrypto::serde_helpers::ToFromByteArray;
 pub use git_version;
 use std::env;
-use sui_types::base_types::ObjectID;
+use sui_types::base_types::{ObjectID, SUI_ADDRESS_LENGTH};
 
 /// Read a byte array from an environment variable and decode it using the specified encoding.
 pub fn decode_byte_array<E: Encoding, const N: usize>(env_name: &str) -> anyhow::Result<[u8; N]> {
@@ -57,8 +57,5 @@ pub fn decode_master_key<E: Encoding>(env_name: &str) -> anyhow::Result<IbeMaste
 
 /// Read an ObjectID from an environment variable.
 pub fn decode_object_id(env_name: &str) -> anyhow::Result<ObjectID> {
-    let hex_string =
-        env::var(env_name).map_err(|_| anyhow!("Environment variable {} must be set", env_name))?;
-    ObjectID::from_hex_literal(&hex_string)
-        .map_err(|_| anyhow!("Invalid ObjectID for environment variable {env_name}"))
+    decode_byte_array::<Hex, SUI_ADDRESS_LENGTH>(env_name).map(ObjectID::new)
 }