Port to bytes, rlp 0.5 (#25)

Author: vorot93

Cargo.toml

@@ -17,10 +17,11 @@ exclude = [
 [dependencies]
 base64 = "0.12.1"
 bs58 = "0.3.1"
+bytes = "1"
 hex = "0.4.2"
 log = "0.4.8"
 rand = "0.7.3"
-rlp = "0.4.5"
+rlp = "0.5"
 zeroize = "1.1.0"
 libsecp256k1 = { version = "0.3.5", optional = true }
 sha3 = "0.9"

src/builder.rs

@@ -1,4 +1,5 @@
 use crate::{Enr, EnrError, EnrKey, EnrPublicKey, Key, NodeId, MAX_ENR_SIZE};
+use bytes::{Bytes, BytesMut};
 use rlp::RlpStream;
 use std::{collections::BTreeMap, marker::PhantomData, net::IpAddr};
 
@@ -12,7 +13,7 @@ pub struct EnrBuilder<K: EnrKey> {
 
     /// The key-value pairs for the ENR record.
     /// Values are stored as RLP encoded bytes.
-    content: BTreeMap<Key, Vec<u8>>,
+    content: BTreeMap<Key, Bytes>,
 
     /// Pins the generic key types.
     phantom: PhantomData<K>,
@@ -39,11 +40,11 @@ impl<K: EnrKey> EnrBuilder<K> {
 
     /// Adds an arbitrary key-value to the `ENRBuilder`.
     pub fn add_value(&mut self, key: impl AsRef<[u8]>, value: &[u8]) -> &mut Self {
-        self.add_value_rlp(key, rlp::encode(&value))
+        self.add_value_rlp(key, rlp::encode(&value).freeze())
     }
 
     /// Adds an arbitrary key-value where the value is raw RLP encoded bytes.
-    pub fn add_value_rlp(&mut self, key: impl AsRef<[u8]>, rlp: Vec<u8>) -> &mut Self {
+    pub fn add_value_rlp(&mut self, key: impl AsRef<[u8]>, rlp: Bytes) -> &mut Self {
         self.content.insert(key.as_ref().to_vec(), rlp);
         self
     }
@@ -97,16 +98,16 @@ impl<K: EnrKey> EnrBuilder<K> {
     }
 
     /// Generates the rlp-encoded form of the ENR specified by the builder config.
-    fn rlp_content(&self) -> Vec<u8> {
-        let mut stream = RlpStream::new();
+    fn rlp_content(&self) -> BytesMut {
+        let mut stream = RlpStream::new_with_buffer(BytesMut::with_capacity(MAX_ENR_SIZE));
         stream.begin_list(self.content.len() * 2 + 1);
         stream.append(&self.seq);
         for (k, v) in &self.content {
             stream.append(k);
             // The values are stored as raw RLP encoded bytes
             stream.append_raw(v, 1);
         }
-        stream.drain()
+        stream.out()
     }
 
     /// Signs record based on the identity scheme. Currently only "v4" is supported.
@@ -144,7 +145,7 @@ impl<K: EnrKey> EnrBuilder<K> {
             }
         }
 
-        self.add_value_rlp("id", rlp::encode(&self.id.as_bytes()));
+        self.add_value_rlp("id", rlp::encode(&self.id.as_bytes()).freeze());
 
         self.add_public_key(&key.public());
         let rlp_content = self.rlp_content();

src/keys/combined.rs

@@ -4,6 +4,7 @@
 //! Currently only `secp256k1` and `ed25519` key types are supported.
 
 use super::{ed25519_dalek as ed25519, EnrKey, EnrPublicKey, SigningError};
+use bytes::Bytes;
 pub use k256;
 use rand::RngCore;
 use rlp::DecoderError;
@@ -66,7 +67,7 @@ impl EnrKey for CombinedKey {
     }
 
     /// Decodes the raw bytes of an ENR's content into a public key if possible.
-    fn enr_to_public(content: &BTreeMap<Key, Vec<u8>>) -> Result<Self::PublicKey, DecoderError> {
+    fn enr_to_public(content: &BTreeMap<Key, Bytes>) -> Result<Self::PublicKey, DecoderError> {
         k256::ecdsa::SigningKey::enr_to_public(content)
             .map(CombinedPublicKey::Secp256k1)
             .or_else(|_| ed25519::Keypair::enr_to_public(content).map(CombinedPublicKey::from))

src/keys/ed25519.rs

@@ -3,6 +3,7 @@ use super::{
     EnrKey, EnrKeyUnambiguous, EnrPublicKey, SigningError,
 };
 use crate::Key;
+use bytes::Bytes;
 use rlp::DecoderError;
 use std::{collections::BTreeMap, convert::TryFrom};
 
@@ -26,7 +27,7 @@ impl EnrKey for ed25519::Keypair {
     }
 
     /// Decodes the raw bytes of an ENR's content into a public key if possible.
-    fn enr_to_public(content: &BTreeMap<Key, Vec<u8>>) -> Result<Self::PublicKey, DecoderError> {
+    fn enr_to_public(content: &BTreeMap<Key, Bytes>) -> Result<Self::PublicKey, DecoderError> {
         let pubkey_bytes = content
             .get(ENR_KEY.as_bytes())
             .ok_or(DecoderError::Custom("Unknown signature"))?;

src/keys/k256_key.rs

@@ -2,6 +2,7 @@
 
 use super::{EnrKey, EnrKeyUnambiguous, EnrPublicKey, SigningError};
 use crate::Key;
+use bytes::Bytes;
 use k256::{
     ecdsa::{
         signature::{DigestVerifier, RandomizedDigestSigner, Signature as _},
@@ -35,7 +36,7 @@ impl EnrKey for SigningKey {
         self.verify_key()
     }
 
-    fn enr_to_public(content: &BTreeMap<Key, Vec<u8>>) -> Result<Self::PublicKey, DecoderError> {
+    fn enr_to_public(content: &BTreeMap<Key, Bytes>) -> Result<Self::PublicKey, DecoderError> {
         let pubkey_bytes = content
             .get(ENR_KEY.as_bytes())
             .ok_or(DecoderError::Custom("Unknown signature"))?;

src/keys/libsecp256k1.rs

@@ -2,6 +2,7 @@
 
 use super::{secp256k1, EnrKey, EnrKeyUnambiguous, EnrPublicKey, SigningError};
 use crate::{digest, Key};
+use bytes::Bytes;
 use rlp::DecoderError;
 use std::collections::BTreeMap;
 
@@ -30,7 +31,7 @@ impl EnrKey for secp256k1::SecretKey {
     }
 
     /// Decodes the raw bytes of an ENR's content into a public key if possible.
-    fn enr_to_public(content: &BTreeMap<Key, Vec<u8>>) -> Result<Self::PublicKey, DecoderError> {
+    fn enr_to_public(content: &BTreeMap<Key, Bytes>) -> Result<Self::PublicKey, DecoderError> {
         let pubkey_bytes = content
             .get(ENR_KEY.as_bytes())
             .ok_or(DecoderError::Custom("Unknown signature"))?;

src/keys/mod.rs

@@ -31,6 +31,7 @@ pub use k256;
 pub use secp256k1;
 
 use crate::Key;
+use bytes::Bytes;
 use rlp::DecoderError;
 use std::{
     collections::BTreeMap,
@@ -54,7 +55,7 @@ pub trait EnrKey: Send + Sync + Unpin + 'static {
     /// `EnrPublicKey`. It takes the ENR's `BTreeMap` and returns a public key.
     ///
     /// Note: This specifies the supported key schemes for an ENR.
-    fn enr_to_public(content: &BTreeMap<Key, Vec<u8>>) -> Result<Self::PublicKey, DecoderError>;
+    fn enr_to_public(content: &BTreeMap<Key, Bytes>) -> Result<Self::PublicKey, DecoderError>;
 }
 
 /// Trait for keys that are uniquely represented

src/keys/rust_secp256k1.rs

@@ -1,5 +1,6 @@
 use super::{EnrKey, EnrKeyUnambiguous, EnrPublicKey, SigningError};
 use crate::{digest, Key};
+use bytes::Bytes;
 use c_secp256k1::SECP256K1;
 use rlp::DecoderError;
 use std::collections::BTreeMap;
@@ -23,7 +24,7 @@ impl EnrKey for c_secp256k1::SecretKey {
         Self::PublicKey::from_secret_key(SECP256K1, self)
     }
 
-    fn enr_to_public(content: &BTreeMap<Key, Vec<u8>>) -> Result<Self::PublicKey, DecoderError> {
+    fn enr_to_public(content: &BTreeMap<Key, Bytes>) -> Result<Self::PublicKey, DecoderError> {
         let pubkey_bytes = content
             .get(ENR_KEY.as_bytes())
             .ok_or(DecoderError::Custom("Unknown signature"))?;

src/lib.rs

@@ -180,6 +180,7 @@ mod builder;
 mod keys;
 mod node_id;
 
+use bytes::{Bytes, BytesMut};
 use log::debug;
 use rlp::{DecoderError, Rlp, RlpStream};
 use std::collections::BTreeMap;
@@ -227,7 +228,7 @@ pub struct Enr<K: EnrKey> {
     /// Key-value contents of the ENR. A BTreeMap is used to get the keys in sorted order, which is
     /// important for verifying the signature of the ENR.
     /// Everything is stored as raw RLP bytes.
-    content: BTreeMap<Key, Vec<u8>>,
+    content: BTreeMap<Key, Bytes>,
 
     /// The signature of the ENR record, stored as bytes.
     signature: Vec<u8>,
@@ -261,13 +262,13 @@ impl<K: EnrKey> Enr<K> {
     }
 
     /// Reads a custom key from the record if it exists as raw RLP bytes.
-    pub fn get_raw_rlp(&self, key: impl AsRef<[u8]>) -> Option<&Vec<u8>> {
-        self.content.get(key.as_ref())
+    pub fn get_raw_rlp(&self, key: impl AsRef<[u8]>) -> Option<&[u8]> {
+        self.content.get(key.as_ref()).map(AsRef::as_ref)
     }
 
     /// Returns an iterator over all key/value pairs in the ENR.
-    pub fn iter(&self) -> impl Iterator<Item = (&Key, &Vec<u8>)> {
-        self.content.iter()
+    pub fn iter(&self) -> impl Iterator<Item = (&Key, &[u8])> {
+        self.content.iter().map(|(k, v)| (k, v.as_ref()))
     }
 
     /// Returns the IPv4 address of the ENR record if it is defined.
@@ -430,18 +431,10 @@ impl<K: EnrKey> Enr<K> {
         }
     }
 
-    /// RLP encodes the ENR into a byte array.
-    #[must_use]
-    pub fn encode(&self) -> Vec<u8> {
-        let mut s = RlpStream::new();
-        s.append(self);
-        s.drain()
-    }
-
     /// Provides the URL-safe base64 encoded "text" version of the ENR prefixed by "enr:".
     #[must_use]
     pub fn to_base64(&self) -> String {
-        let hex = base64::encode_config(&self.encode(), base64::URL_SAFE_NO_PAD);
+        let hex = base64::encode_config(&rlp::encode(self), base64::URL_SAFE_NO_PAD);
         format!("enr:{}", hex)
     }
 
@@ -480,8 +473,8 @@ impl<K: EnrKey> Enr<K> {
         key: impl AsRef<[u8]>,
         value: &[u8],
         enr_key: &K,
-    ) -> Result<Option<Vec<u8>>, EnrError> {
-        self.insert_raw_rlp(key, rlp::encode(&value), enr_key)
+    ) -> Result<Option<Bytes>, EnrError> {
+        self.insert_raw_rlp(key, rlp::encode(&value).freeze(), enr_key)
     }
 
     /// Adds or modifies a key/value to the ENR record. A `EnrKey` is required to re-sign the record once
@@ -491,11 +484,11 @@ impl<K: EnrKey> Enr<K> {
     pub fn insert_raw_rlp(
         &mut self,
         key: impl AsRef<[u8]>,
-        value: Vec<u8>,
+        value: Bytes,
         enr_key: &K,
-    ) -> Result<Option<Vec<u8>>, EnrError> {
+    ) -> Result<Option<Bytes>, EnrError> {
         // currently only support "v4" identity schemes
-        if key.as_ref() == b"id" && value != b"v4" {
+        if key.as_ref() == b"id" && &*value != b"v4" {
             return Err(EnrError::UnsupportedIdentityScheme);
         }
 
@@ -504,7 +497,7 @@ impl<K: EnrKey> Enr<K> {
         let public_key = enr_key.public();
         let previous_key = self.content.insert(
             public_key.enr_key(),
-            rlp::encode(&public_key.encode().as_ref()),
+            rlp::encode(&public_key.encode().as_ref()).freeze(),
         );
 
         // check the size of the record
@@ -640,27 +633,31 @@ impl<K: EnrKey> Enr<K> {
 
         let (prev_ip, prev_port) = match socket.ip() {
             IpAddr::V4(addr) => (
-                self.content
-                    .insert("ip".into(), rlp::encode(&(&addr.octets() as &[u8]))),
+                self.content.insert(
+                    "ip".into(),
+                    rlp::encode(&(&addr.octets() as &[u8])).freeze(),
+                ),
                 self.content.insert(
                     port_string.clone(),
-                    rlp::encode(&(&socket.port().to_be_bytes() as &[u8])),
+                    rlp::encode(&(&socket.port().to_be_bytes() as &[u8])).freeze(),
                 ),
             ),
             IpAddr::V6(addr) => (
-                self.content
-                    .insert("ip6".into(), rlp::encode(&(&addr.octets() as &[u8]))),
+                self.content.insert(
+                    "ip6".into(),
+                    rlp::encode(&(&addr.octets() as &[u8])).freeze(),
+                ),
                 self.content.insert(
                     port_v6_string.clone(),
-                    rlp::encode(&(&socket.port().to_be_bytes() as &[u8])),
+                    rlp::encode(&(&socket.port().to_be_bytes() as &[u8])).freeze(),
                 ),
             ),
         };
 
         let public_key = key.public();
         let previous_key = self.content.insert(
             public_key.enr_key(),
-            rlp::encode(&public_key.encode().as_ref()),
+            rlp::encode(&public_key.encode().as_ref()).freeze(),
         );
 
         // check the size and revert on failure
@@ -726,8 +723,8 @@ impl<K: EnrKey> Enr<K> {
     // Private Functions //
 
     /// Evaluates the RLP-encoding of the content of the ENR record.
-    fn rlp_content(&self) -> Vec<u8> {
-        let mut stream = RlpStream::new();
+    fn rlp_content(&self) -> BytesMut {
+        let mut stream = RlpStream::new_with_buffer(BytesMut::with_capacity(MAX_ENR_SIZE));
         stream.begin_list(self.content.len() * 2 + 1);
         stream.append(&self.seq);
         for (k, v) in &self.content {
@@ -736,7 +733,7 @@ impl<K: EnrKey> Enr<K> {
             // Values are raw RLP encoded data
             stream.append_raw(v, 1);
         }
-        stream.drain()
+        stream.out()
     }
 
     /// Signs the ENR record based on the identity scheme. Currently only "v4" is supported.
@@ -892,7 +889,7 @@ impl<K: EnrKey> rlp::Decodable for Enr<K> {
                 return Err(DecoderError::Custom("Unsorted keys"));
             }
             prev = Some(key.clone());
-            content.insert(key, value.into());
+            content.insert(key, Bytes::copy_from_slice(value));
         }
 
         // verify we know the signature type