keytrans: Prefer expect() to unwrap() and more code deduplication

Oyelowo · Aug 13, 2024 · 9a6898a · 9a6898a
1 parent 360b335
commit 9a6898a
Show file tree

Hide file tree

Showing 4 changed files with 31 additions and 33 deletions.
diff --git a/rust/keytrans/src/commitments.rs b/rust/keytrans/src/commitments.rs
@@ -19,7 +19,7 @@ pub fn commit(search_key: &[u8], data: &[u8], nonce: &[u8; 16]) -> Vec<u8> {
     // a serialized public key. Neither should reach 2^32 bound.
     let data_len: u32 = data.len().try_into().expect("data too large");
 
-    let mut mac = HmacSha256::new_from_slice(FIXED_KEY).unwrap();
+    let mut mac = HmacSha256::new_from_slice(FIXED_KEY).expect("can create hmac from fixed key");
     mac.update(nonce);
     mac.update(&key_len.to_be_bytes());
     mac.update(search_key);

diff --git a/rust/keytrans/src/lib.rs b/rust/keytrans/src/lib.rs
@@ -2,6 +2,9 @@
 // Copyright 2024 Signal Messenger, LLC.
 // SPDX-License-Identifier: AGPL-3.0-only
 //
+
+#![cfg_attr(not(test), warn(clippy::unwrap_used))]
+
 mod commitments;
 mod guide;
 mod implicit;

diff --git a/rust/keytrans/src/verify.rs b/rust/keytrans/src/verify.rs
@@ -302,7 +302,7 @@ fn verify_timestamp(
     } = allowed_range;
     let now = SystemTime::now()
         .duration_since(UNIX_EPOCH)
-        .unwrap()
+        .expect("valid system time")
         .as_millis() as i128;
     let delta = now - timestamp as i128;
     let format_message = |s: &str| match description {
@@ -447,10 +447,7 @@ async fn verify_search_internal(
     }
 
     // Evaluate the inclusion proof to get a candidate root value.
-    let mut ids: Vec<u64> = leaves.keys().cloned().collect();
-    ids.sort();
-
-    let values: Vec<[u8; 32]> = ids.iter().map(|id| *leaves.get(id).unwrap()).collect();
+    let (ids, values) = into_sorted_pairs(leaves);
 
     let inclusion_proof = get_hash_proof(&search_proof.inclusion)?;
     let root = evaluate_batch_proof(&ids, tree_size, &values, &inclusion_proof)?;
@@ -563,10 +560,7 @@ pub async fn verify_monitor(
             }
         }
     } else {
-        let mut ids: Vec<u64> = mpa.leaves.keys().cloned().collect();
-        ids.sort();
-
-        let values: Vec<[u8; 32]> = ids.iter().map(|id| *mpa.leaves.get(id).unwrap()).collect();
+        let (ids, values) = into_sorted_pairs(mpa.leaves);
 
         evaluate_batch_proof(&ids, tree_size, &values, &inclusion_proof)?
     };
@@ -829,7 +823,9 @@ impl MonitoringDataWrapper {
 
     async fn save(self, storage: &mut dyn LogStore, search_key: &str) -> Result<()> {
         if self.changed {
-            storage.set_data(search_key, self.inner.unwrap()).await?
+            if let Some(data) = self.inner {
+                storage.set_data(search_key, data).await?
+            }
         }
         Ok(())
     }
@@ -886,10 +882,7 @@ pub async fn truncate_search_response(
         result.expect("truncate_search_response called with search response that is not verified");
 
     // Evaluate the inclusion proof to get root value.
-    let mut ids: Vec<u64> = leaves.keys().cloned().collect();
-    ids.sort();
-
-    let values: Vec<[u8; 32]> = ids.iter().map(|id| *leaves.get(id).unwrap()).collect();
+    let (ids, values) = into_sorted_pairs(leaves);
 
     let inclusion_proof = get_hash_proof(&search_proof.inclusion)?;
 
@@ -902,6 +895,12 @@ pub async fn truncate_search_response(
     Ok((result_id + 1, root))
 }
 
+fn into_sorted_pairs<K: Ord + Copy, V>(map: HashMap<K, V>) -> (Vec<K>, Vec<V>) {
+    let mut pairs = map.into_iter().collect::<Vec<_>>();
+    pairs.sort_by_key(|pair| pair.0);
+    pairs.into_iter().unzip()
+}
+
 #[cfg(test)]
 mod test {
     use assert_matches::assert_matches;

diff --git a/rust/keytrans/src/vrf.rs b/rust/keytrans/src/vrf.rs
@@ -36,7 +36,7 @@ fn encode_to_curve_try_and_increment(salt: &[u8], data: &[u8]) -> EdwardsPoint {
         hasher.update([i, DOMAIN_SEPARATOR_BACK]);
 
         let r = hasher.finalize_reset();
-        match CompressedEdwardsY(r[..32].try_into().unwrap()).decompress() {
+        match CompressedEdwardsY(r[..32].try_into().expect("hash has enough bytes")).decompress() {
             Some(pt) => return pt.mul_by_cofactor(),
             None => i += 1,
         }
@@ -52,7 +52,7 @@ fn generate_challenge(pts: [&[u8; 32]; 5]) -> [u8; 16] {
     hasher.update([DOMAIN_SEPARATOR_BACK]);
     let c = hasher.finalize();
 
-    c[..16].try_into().unwrap()
+    c[..16].try_into().expect("hash has enough bytes")
 }
 
 fn proof_to_hash(gamma: &EdwardsPoint) -> [u8; 32] {
@@ -62,7 +62,7 @@ fn proof_to_hash(gamma: &EdwardsPoint) -> [u8; 32] {
     hasher.update([DOMAIN_SEPARATOR_BACK]);
     let index = hasher.finalize();
 
-    index[..32].try_into().unwrap()
+    index[..32].try_into().expect("hash has enough bytes")
 }
 
 /// PublicKey holds a VRF public key.
@@ -91,24 +91,20 @@ impl PublicKey {
     /// the index if so.
     pub fn proof_to_hash(&self, m: &[u8], proof: &[u8; 80]) -> Result<[u8; 32]> {
         // Decode proof into its component parts: gamma, c, and s.
-        let gamma = match CompressedEdwardsY(proof[..32].try_into().unwrap()).decompress() {
-            Some(pt) => pt,
-            None => return Err(Error::InvalidProof),
-        };
+        let gamma = CompressedEdwardsY(proof[..32].try_into().expect("proof has enough bytes"))
+            .decompress()
+            .ok_or(Error::InvalidProof)?;
 
         let mut c_bytes = [0u8; 32];
         c_bytes[..16].copy_from_slice(&proof[32..48]);
-        let c = Scalar::from_canonical_bytes(c_bytes);
-        if c.is_none().into() {
-            return Err(Error::InvalidProof);
-        }
-        let c = -(c.unwrap());
+        let c = -Scalar::from_canonical_bytes(c_bytes)
+            .into_option()
+            .ok_or(Error::InvalidProof)?;
 
-        let s = Scalar::from_canonical_bytes(proof[48..80].try_into().unwrap());
-        if s.is_none().into() {
-            return Err(Error::InvalidProof);
-        }
-        let s = s.unwrap();
+        let s =
+            Scalar::from_canonical_bytes(proof[48..80].try_into().expect("proof has enough bytes"))
+                .into_option()
+                .ok_or(Error::InvalidProof)?;
 
         // H = encode_to_curve_try_and_increment(pk, m)
         // U = [s]B - [c]Y
@@ -122,7 +118,7 @@ impl PublicKey {
         let c_prime = generate_challenge([
             &self.compressed,
             &h.compress().0,
-            proof[..32].try_into().unwrap(),
+            proof[..32].try_into().expect("proof has enough bytes"),
             &u.compress().0,
             &v.compress().0,
         ]);