Kademlia Records (#1144)

* initial implementation of the records

* move to multihash keys

* correctly process query results

* comments and formatting

* correctly return closer_peers in query

* checking wrong peer id in test

* Apply suggestions from code review

Co-Authored-By: Roman Borschel <romanb@users.noreply.github.com>

* Fix changes from suggestions

* Send responses to PUT_VALUE requests

* Shortcut in get_value

* Update protocols/kad/src/behaviour.rs

Co-Authored-By: Roman Borschel <romanb@users.noreply.github.com>

* Revert "Update protocols/kad/src/behaviour.rs"

This reverts commit 579ce742a7f4c94587f1e1f0866d2a3a37418efb.

* Remove duplicate insertion

* Adds a record to a PUT_VALUE response

* Fix a racy put_value test

* Store value ourselves only if we are in K closest

* Abstract over storage

* Revert "Abstract over storage": bad take

This reverts commit eaebf5b6d915712eaf3b05929577fdf697f204d8.

* Abstract over records storage using hashmap as default

* Constructor for custom records

* New Record type and its traits

* Fix outdated storage name

* Fixes returning an event

* Change FindNodeReq key type to Multihash

* WriteState for a second stage of a PUT_VALUE request

* GET_VALUE should not have a record

* Refactor a match arm

* Add successes and failures counters to PutValueRes

* If value is found no need to return closer peers

* Remove a custo storage from tests

* Rename a test to get_value_not_found

* Adds a TODO to change FindNode request key to Multihash

Co-Authored-By: Roman Borschel <romanb@users.noreply.github.com>

* Move MemoryRecordStorage to record.rs

* Return a Cow-ed Record from get

* Fix incorrect GET_VALUE parsing

* Various fixes with review

* Fixes get_value_not_found

* Fix peerids names in test

* another fix

* PutValue correctly distributes values

* Simplify the test

* Check that results are actually the closest

* Reverts changes to tests

* Fix the test topology and checking the results

* Run put_value test ten times

* Adds a get_value test

* Apply suggestions from code review

Co-Authored-By: Roman Borschel <romanb@users.noreply.github.com>

* Make Record fields public

* Moves WriteState to write.rs

* A couple of minor fixes

* Another few fixes of review

* Simplify the put_value test

* Dont synchronously return an error from put_value

* Formatting fixes and comments

* Collect a bunch of results

* Take exactly as much elements as neede

* Check if the peer is still connected

* Adds a multiple GetValueResults results number test

* Unnecessary mut iterators in put_value

* Ask for num_results in get_value

* Dont allocate twice in get_value

* Dont count same errored peer multiple times

* Apply suggestions from code review

Co-Authored-By: Roman Borschel <romanb@users.noreply.github.com>

* Fix another review

* Apply suggestions from code review

Co-Authored-By: Pierre Krieger <pierre.krieger1708@gmail.com>

* Bring back FromIterator and improve a panic message

* Update protocols/kad/src/behaviour.rs

Co-Authored-By: Pierre Krieger <pierre.krieger1708@gmail.com>

protocols/kad/src/handler.rs

@@ -22,6 +22,7 @@ use crate::protocol::{
     KadInStreamSink, KadOutStreamSink, KadPeer, KadRequestMsg, KadResponseMsg,
     KademliaProtocolConfig,
 };
+use crate::record::Record;
 use futures::prelude::*;
 use libp2p_core::protocols_handler::{
     KeepAlive,
@@ -182,6 +183,42 @@ pub enum KademliaHandlerEvent<TUserData> {
         /// Known provider for this key.
         provider_peer: KadPeer,
     },
+
+    /// Request to get a value from the dht records
+    GetValue {
+        /// Key for which we should look in the dht
+        key: Multihash,
+        /// Identifier of the request. Needs to be passed back when answering.
+        request_id: KademliaRequestId,
+    },
+
+    /// Response to a `KademliaHandlerIn::GetValue`.
+    GetValueRes {
+        /// The result is present if the key has been found
+        result: Option<Record>,
+        /// Nodes closest to the key.
+        closer_peers: Vec<KadPeer>,
+        /// The user data passed to the `GetValue`.
+        user_data: TUserData,
+    },
+
+    /// Request to put a value in the dht records
+    PutValue {
+        /// The key of the record
+        key: Multihash,
+        /// The value of the record
+        value: Vec<u8>,
+        /// Identifier of the request. Needs to be passed back when answering.
+        request_id: KademliaRequestId,
+    },
+
+    /// Response to a request to put a value
+    PutValueRes {
+        /// The key we were putting in
+        key: Multihash,
+        /// The user data passed to the `GetValue`.
+        user_data: TUserData,
+    }
 }
 
 /// Error that can happen when requesting an RPC query.
@@ -229,12 +266,13 @@ impl From<ProtocolsHandlerUpgrErr<io::Error>> for KademliaHandlerQueryErr {
 }
 
 /// Event to send to the handler.
+#[derive(Debug)]
 pub enum KademliaHandlerIn<TUserData> {
     /// Request for the list of nodes whose IDs are the closest to `key`. The number of nodes
     /// returned is not specified, but should be around 20.
     FindNodeReq {
         /// Identifier of the node.
-        key: PeerId,
+        key: Multihash,
         /// Custom user data. Passed back in the out event when the results arrive.
         user_data: TUserData,
     },
@@ -280,6 +318,44 @@ pub enum KademliaHandlerIn<TUserData> {
         /// Known provider for this key.
         provider_peer: KadPeer,
     },
+
+    /// Request to get a node from the dht
+    GetValue {
+        /// The key of the value we are looking for
+        key: Multihash,
+        /// Custom data. Passed back in the out event when the results arrive.
+        user_data: TUserData,
+    },
+
+    /// Response to a `GetValue`.
+    GetValueRes {
+        /// The value that might have been found in our storage.
+        result: Option<Record>,
+        /// Nodes that are closer to the key we were searching for.
+        closer_peers: Vec<KadPeer>,
+        /// Identifier of the request that was made by the remote.
+        request_id: KademliaRequestId,
+    },
+
+    /// Put a value into the dht records.
+    PutValue {
+        /// The key of the record.
+        key: Multihash,
+        /// The value of the record.
+        value: Vec<u8>,
+        /// Custom data. Passed back in the out event when the results arrive.
+        user_data: TUserData,
+    },
+
+    /// Response to a `PutValue`.
+    PutValueRes {
+        /// Key of the value that was put.
+        key: Multihash,
+        /// Value that was put.
+        value: Vec<u8>,
+        /// Identifier of the request that was made by the remote.
+        request_id: KademliaRequestId,
+    }
 }
 
 /// Unique identifier for a request. Must be passed back in order to answer a request from
@@ -397,9 +473,15 @@ where
     fn inject_event(&mut self, message: KademliaHandlerIn<TUserData>) {
         match message {
             KademliaHandlerIn::FindNodeReq { key, user_data } => {
-                let msg = KadRequestMsg::FindNode { key: key.clone() };
-                self.substreams
-                    .push(SubstreamState::OutPendingOpen(msg, Some(user_data.clone())));
+                // FIXME: Change `KadRequestMsg::FindNode::key` to be a `Multihash`.
+                match PeerId::from_multihash(key.clone()) {
+                    Ok(key) => {
+                        let msg = KadRequestMsg::FindNode { key };
+                        self.substreams
+                            .push(SubstreamState::OutPendingOpen(msg, Some(user_data.clone())));
+                    },
+                    Err(_) => (),
+                }
             }
             KademliaHandlerIn::FindNodeRes {
                 closer_peers,
@@ -465,6 +547,74 @@ where
                 self.substreams
                     .push(SubstreamState::OutPendingOpen(msg, None));
             }
+            KademliaHandlerIn::GetValue { key, user_data } => {
+                let msg = KadRequestMsg::GetValue { key };
+                self.substreams
+                    .push(SubstreamState::OutPendingOpen(msg, Some(user_data)));
+
+            }
+            KademliaHandlerIn::PutValue { key, value, user_data } => {
+                let msg = KadRequestMsg::PutValue {
+                    key,
+                    value,
+                };
+
+                self.substreams
+                    .push(SubstreamState::OutPendingOpen(msg, Some(user_data)));
+            }
+            KademliaHandlerIn::GetValueRes {
+                result,
+                closer_peers,
+                request_id,
+            } => {
+                let pos = self.substreams.iter().position(|state| match state {
+                    SubstreamState::InWaitingUser(ref conn_id, _)
+                        => conn_id == &request_id.connec_unique_id,
+                    _ => false,
+                });
+
+                if let Some(pos) = pos {
+                    let (conn_id, substream) = match self.substreams.remove(pos) {
+                        SubstreamState::InWaitingUser(conn_id, substream) => (conn_id, substream),
+                        _ => unreachable!(),
+                    };
+
+                    let msg = KadResponseMsg::GetValue {
+                        result,
+                        closer_peers: closer_peers.clone(),
+                    };
+                    self.substreams
+                        .push(SubstreamState::InPendingSend(conn_id, substream, msg));
+                }
+            }
+            KademliaHandlerIn::PutValueRes {
+                key,
+                request_id,
+                value,
+            } => {
+                let pos = self.substreams.iter().position(|state| match state {
+                    SubstreamState::InWaitingUser(ref conn_id, _)
+                        if conn_id == &request_id.connec_unique_id =>
+                        {
+                            true
+                        }
+                    _ => false,
+                });
+
+                if let Some(pos) = pos {
+                    let (conn_id, substream) = match self.substreams.remove(pos) {
+                        SubstreamState::InWaitingUser(conn_id, substream) => (conn_id, substream),
+                        _ => unreachable!(),
+                    };
+
+                    let msg = KadResponseMsg::PutValue {
+                        key,
+                        value,
+                    };
+                    self.substreams
+                        .push(SubstreamState::InPendingSend(conn_id, substream, msg));
+                }
+            }
         }
     }
 
@@ -740,6 +890,15 @@ fn process_kad_request<TUserData>(
         KadRequestMsg::AddProvider { key, provider_peer } => {
             Ok(KademliaHandlerEvent::AddProvider { key, provider_peer })
         }
+        KadRequestMsg::GetValue { key } => Ok(KademliaHandlerEvent::GetValue {
+            key,
+            request_id: KademliaRequestId { connec_unique_id },
+        }),
+        KadRequestMsg::PutValue { key, value } => Ok(KademliaHandlerEvent::PutValue {
+            key,
+            value,
+            request_id: KademliaRequestId { connec_unique_id },
+        })
     }
 }
 
@@ -771,5 +930,19 @@ fn process_kad_response<TUserData>(
             provider_peers,
             user_data,
         },
+        KadResponseMsg::GetValue {
+            result,
+            closer_peers,
+        } => KademliaHandlerEvent::GetValueRes {
+            result,
+            closer_peers,
+            user_data,
+        },
+        KadResponseMsg::PutValue { key, .. } => {
+            KademliaHandlerEvent::PutValueRes {
+                key,
+                user_data,
+            }
+        }
     }
 }