Add peer id inlining for small public keys (#1237)

* Add peer id inlining for small public keys

* Apply @twittner suggestions

* Restore hashing

Originally from libp2p/rust-libp2p@c6e8e63.

Author: tomaka

src/errors.rs

@@ -5,13 +5,19 @@ use std::{error, fmt};
 pub enum EncodeError {
     /// The requested hash algorithm isn't supported by this library.
     UnsupportedType,
+    /// The input length is too large for the hash algorithm.
+    UnsupportedInputLength,
 }
 
 impl fmt::Display for EncodeError {
     #[inline]
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         match *self {
             EncodeError::UnsupportedType => write!(f, "This type is not supported yet"),
+            EncodeError::UnsupportedInputLength => write!(
+                f,
+                "The length of the input for the given hash is not yet supported"
+            ),
         }
     }
 }

src/hashes.rs

@@ -3,6 +3,8 @@
 /// Not all hash types are supported by this library.
 #[derive(PartialEq, Eq, Clone, Debug, Copy, Hash)]
 pub enum Hash {
+    /// Identity (Raw binary )
+    Identity,
     /// SHA-1 (20-byte hash size)
     SHA1,
     /// SHA-256 (32-byte hash size)
@@ -39,6 +41,7 @@ impl Hash {
     /// Get the corresponding hash code.
     pub fn code(self) -> u16 {
         match self {
+            Hash::Identity => 0x00,
             Hash::SHA1 => 0x11,
             Hash::SHA2256 => 0x12,
             Hash::SHA2512 => 0x13,
@@ -60,6 +63,7 @@ impl Hash {
     /// Get the hash length in bytes.
     pub fn size(self) -> u8 {
         match self {
+            Hash::Identity => 42,
             Hash::SHA1 => 20,
             Hash::SHA2256 => 32,
             Hash::SHA2512 => 64,
@@ -81,6 +85,7 @@ impl Hash {
     /// Returns the algorithm corresponding to a code, or `None` if no algorithm is matching.
     pub fn from_code(code: u16) -> Option<Hash> {
         Some(match code {
+            0x00 => Hash::Identity,
             0x11 => Hash::SHA1,
             0x12 => Hash::SHA2256,
             0x13 => Hash::SHA2512,

src/lib.rs

@@ -79,26 +79,47 @@ macro_rules! match_encoder {
 /// ```
 ///
 pub fn encode(hash: Hash, input: &[u8]) -> Result<Multihash, EncodeError> {
-    let (offset, mut output) = encode_hash(hash);
-    match_encoder!(hash for (input, &mut output[offset ..]) {
-        SHA1 => sha1::Sha1,
-        SHA2256 => sha2::Sha256,
-        SHA2512 => sha2::Sha512,
-        SHA3224 => tiny::new_sha3_224,
-        SHA3256 => tiny::new_sha3_256,
-        SHA3384 => tiny::new_sha3_384,
-        SHA3512 => tiny::new_sha3_512,
-        Keccak224 => tiny::new_keccak224,
-        Keccak256 => tiny::new_keccak256,
-        Keccak384 => tiny::new_keccak384,
-        Keccak512 => tiny::new_keccak512,
-        Blake2b512 => blake2::blake2b,
-        Blake2s256 => blake2::blake2s,
-    });
-
-    Ok(Multihash {
-        bytes: output.freeze(),
-    })
+    // Custom length encoding for the identity multihash
+    if let Hash::Identity = hash {
+        if u64::from(std::u32::MAX) < as_u64(input.len()) {
+            return Err(EncodeError::UnsupportedInputLength);
+        }
+        let mut buf = encode::u16_buffer();
+        let code = encode::u16(hash.code(), &mut buf);
+        let mut len_buf = encode::u32_buffer();
+        let size = encode::u32(input.len() as u32, &mut len_buf);
+
+        let total_len = code.len() + size.len() + input.len();
+
+        let mut output = BytesMut::with_capacity(total_len);
+        output.put_slice(code);
+        output.put_slice(size);
+        output.put_slice(input);
+        Ok(Multihash {
+            bytes: output.freeze(),
+        })
+    } else {
+        let (offset, mut output) = encode_hash(hash);
+        match_encoder!(hash for (input, &mut output[offset ..]) {
+            SHA1 => sha1::Sha1,
+            SHA2256 => sha2::Sha256,
+            SHA2512 => sha2::Sha512,
+            SHA3224 => tiny::new_sha3_224,
+            SHA3256 => tiny::new_sha3_256,
+            SHA3384 => tiny::new_sha3_384,
+            SHA3512 => tiny::new_sha3_512,
+            Keccak224 => tiny::new_keccak224,
+            Keccak256 => tiny::new_keccak256,
+            Keccak384 => tiny::new_keccak384,
+            Keccak512 => tiny::new_keccak512,
+            Blake2b512 => blake2::blake2b,
+            Blake2s256 => blake2::blake2s,
+        });
+
+        Ok(Multihash {
+            bytes: output.freeze(),
+        })
+    }
 }
 
 // Encode the given [`Hash`] value and ensure the returned [`BytesMut`]
@@ -206,6 +227,16 @@ impl<'a> MultihashRef<'a> {
         let (code, bytes) = decode::u16(&input).map_err(|_| DecodeError::BadInputLength)?;
 
         let alg = Hash::from_code(code).ok_or(DecodeError::UnknownCode)?;
+
+        // handle the identity case
+        if alg == Hash::Identity {
+            let (hash_len, bytes) = decode::u32(&bytes).map_err(|_| DecodeError::BadInputLength)?;
+            if as_u64(bytes.len()) != u64::from(hash_len) {
+                return Err(DecodeError::BadInputLength);
+            }
+            return Ok(MultihashRef { bytes: input });
+        }
+
         let hash_len = usize::from(alg.size());
 
         // Length of input after hash code should be exactly hash_len + 1
@@ -256,3 +287,8 @@ impl<'a> PartialEq<Multihash> for MultihashRef<'a> {
         self.bytes == &*other.bytes
     }
 }
+
+#[cfg(any(target_pointer_width = "32", target_pointer_width = "64"))]
+fn as_u64(a: usize) -> u64 {
+    a as u64
+}