Kademlia Records

protocols/kad/src/behaviour.rs

@@ -89,6 +89,9 @@ pub struct Kademlia<TSubstream> {
 
     /// Marker to pin the generics.
     marker: PhantomData<TSubstream>,
+
+    /// The records that we keep.
+    records: FnvHashMap<Multihash, Vec<u8>>
 }
 
 /// Opaque type. Each query that we start gets a unique number.
@@ -131,6 +134,22 @@ enum QueryInfoInner {
         /// Which hash we're targetting.
         target: Multihash,
     },
+
+    /// Put the value to the dht records
+    PutValue {
+        /// The key of the record being inserted
+        key: Multihash,
+        /// The value of the record being inserted
+        value: Vec<u8>,
+    },
+
+    /// Get value from the dht record
+    GetValue {
+        /// The key we're looking for
+        key: Multihash,
+        /// The results from peers are stored here
+        results: Vec<(Multihash, Vec<u8>)>,
+    },
 }
 
 impl Into<kbucket::Key<QueryInfo>> for QueryInfo {
@@ -146,6 +165,8 @@ impl AsRef<[u8]> for QueryInfo {
             QueryInfoInner::FindPeer(peer) => peer.as_ref(),
             QueryInfoInner::GetProviders { target, .. } => target.as_bytes(),
             QueryInfoInner::AddProvider { target } => target.as_bytes(),
+            QueryInfoInner::GetValue { key, .. } => key.as_bytes(),
+            QueryInfoInner::PutValue { key, .. } => key.as_bytes(),
         }
     }
 }
@@ -170,6 +191,17 @@ impl QueryInfo {
                 key: unimplemented!(), // TODO: target.clone(),
                 user_data,
             },
+            QueryInfoInner::GetValue { key, .. } => {
+                KademliaHandlerIn::GetValue {
+                    key: key.clone(),
+                    user_data,
+                }
+            },
+            QueryInfoInner::PutValue { key, value } => KademliaHandlerIn::PutValue {
+                key: key.clone(),
+                value: value.clone(),
+                user_data,
+            }
         }
     }
 }
@@ -261,6 +293,7 @@ impl<TSubstream> Kademlia<TSubstream> {
             rpc_timeout: Duration::from_secs(8),
             add_provider: SmallVec::new(),
             marker: PhantomData,
+            records: FnvHashMap::default(),
         }
     }
 
@@ -283,6 +316,26 @@ impl<TSubstream> Kademlia<TSubstream> {
         self.start_query(QueryInfoInner::GetProviders { target, pending_results: Vec::new() });
     }
 
+    /// Starts an iterative `GET_VALUE` request.
+    pub fn get_value(&mut self, key: &Multihash) {
+        if let Some(value) = self.records.get(key) {
+            self.queued_events.push(NetworkBehaviourAction::GenerateEvent(
+                KademliaOut::GetValueResult {
+                    result: Some((key.clone(), value.clone())),
+                    closer_peers: None,
+                }
+            ));
+        }
+        self.start_query(QueryInfoInner::GetValue { key: key.clone(), results: vec![] });
+    }
+
+    /// Starts an iterative `PUT_VALUE` request
+    pub fn put_value(&mut self, key: Multihash, value: Vec<u8>) {
+        // TODO: Probably we shouldn't store the value ourselves
+        self.records.insert(key.clone(), value.clone());
+        self.start_query(QueryInfoInner::PutValue{key, value});
+    }
+
     /// Register the local node as the provider for the given key.
     ///
     /// This will periodically send `ADD_PROVIDER` messages to the nodes closest to the key. When
@@ -623,6 +676,66 @@ where
                 self.add_provider.push((key, provider_peer.node_id));
                 return;
             }
+            KademliaHandlerEvent::GetValue { key, request_id } => {
+                let result = match self.records.get(&key) {
+                    Some(value) => Some((key.clone(), value.clone())),
+                    None => None,
+                };
+                let closer_peers = self.find_closest(&kbucket::Key::from(key), &source);
+
+                self.queued_events.push(NetworkBehaviourAction::SendEvent {
+                    peer_id: source,
+                    event: KademliaHandlerIn::GetValueRes {
+                        result,
+                        closer_peers,
+                        request_id,
+                    },
+                });
+            }
+            KademliaHandlerEvent::GetValueRes {
+                result,
+                closer_peers,
+                user_data,
+            } => {
+                if let Some(query) = self.active_queries.get_mut(&user_data) {
+                    if let QueryInfoInner::GetValue {
+                        key: _,
+                        results
+                    } = &mut query.target_mut().inner {
+                        if let Some(result) = result {
+                            results.push(result);
+                        }
+                    }
+                    // A hacky way to end a query
+                    query.inject_rpc_result(&source, vec![source.clone()]);
+                }
+                self.discovered(&user_data, &source, closer_peers.iter());
+            }
+            KademliaHandlerEvent::PutValue {
+                key,
+                value,
+                request_id
+            } => {
+                self.records.insert(key.clone(), value.clone());
+
+                self.queued_events.push(NetworkBehaviourAction::SendEvent {
+                    peer_id: source,
+                    event: KademliaHandlerIn::PutValueRes {
+                        key,
+                        value,
+                        request_id,
+                    },
+                });
+            }
+            KademliaHandlerEvent::PutValueRes {
+                key: _,
+                user_data,
+            } => {
+                if let Some(query) = self.active_queries.get_mut(&user_data) {
+                    // A hacky way to end a query
+                    query.inject_rpc_result(&source, vec![source.clone()]);
+                }
+            }
         };
     }
 
@@ -753,6 +866,27 @@ where
                             self.queued_events.push(event);
                         }
                     },
+                    QueryInfoInner::GetValue { key: _, results } => {
+                        let result = match results.first() {
+                            Some(a) => Some(a.clone()),
+                            None => None,
+                        };
+
+                        let event = KademliaOut::GetValueResult {
+                            result: result.clone(),
+                            closer_peers: {
+                                match result {
+                                None => Some(closer_peers.collect()),
+                                Some(_) => None,
+                            }},
+                        };
+
+                        break Async::Ready(NetworkBehaviourAction::GenerateEvent(event));
+                    },
+                    QueryInfoInner::PutValue { key, .. } => {
+                        let event = KademliaOut::PutValueResult { key };
+                        break Async::Ready(NetworkBehaviourAction::GenerateEvent(event));
+                    },
                 }
             } else {
                 break Async::NotReady;
@@ -802,6 +936,20 @@ pub enum KademliaOut {
         /// List of peers ordered from closest to furthest away.
         closer_peers: Vec<PeerId>,
     },
+
+    /// Result of a `GET_VALUE` query
+    GetValueResult {
+        /// The result that we have probably received
+        result: Option<(Multihash, Vec<u8>)>,
+        /// List of peers ordered from closes to furthest from the key
+        closer_peers: Option<Vec<PeerId>>,
+    },
+
+    /// Result of a `PUT_VALUE` query
+    PutValueResult {
+        /// The key that we were inserting
+        key: Multihash,
+    }
 }
 
 impl From<kbucket::EntryView<PeerId, Addresses>> for KadPeer {

protocols/kad/src/behaviour/test.rs

@@ -38,6 +38,7 @@ use libp2p_yamux as yamux;
 use rand::random;
 use std::{io, u64};
 use tokio::runtime::Runtime;
+use multihash::Hash;
 
 type TestSwarm = Swarm<
     Boxed<(PeerId, StreamMuxerBox), io::Error>,
@@ -223,3 +224,83 @@ fn unresponsive_not_returned_indirect() {
         }))
         .unwrap();
 }
+
+
+#[test]
+fn search_for_unknown_value() {
+    let (port_base, mut swarms) = build_nodes(3);
+
+    let peer_ids: Vec<_> = swarms.iter()
+        .map(|swarm| Swarm::local_peer_id(&swarm).clone()).collect();
+
+    swarms[0].add_address(&peer_ids[1], Protocol::Memory(port_base + 1).into());
+    swarms[1].add_address(&peer_ids[2], Protocol::Memory(port_base + 2).into());
+
+    let target_key = multihash::encode(Hash::SHA2256, &vec![1,2,3]).unwrap();
+    swarms[0].get_value(&target_key);
+
+    Runtime::new().unwrap().block_on(
+        future::poll_fn(move || -> Result<_, io::Error> {
+            for swarm in &mut swarms {
+                loop {
+                    match swarm.poll().unwrap() {
+                        Async::Ready(Some(KademliaOut::GetValueResult { result, .. })) => {
+                            assert_eq!(result, None);
+                            return Ok(Async::Ready(()));
+                        }
+                        Async::Ready(_) => (),
+                        Async::NotReady => break,
+                    }
+                }
+            }
+
+            Ok(Async::NotReady)
+        }))
+        .unwrap()
+}
+
+#[test]
+fn put_value() {
+    let (port_base, mut swarms) = build_nodes(3);
+
+    let peer_ids: Vec<_> = swarms.iter()
+        .map(|swarm| Swarm::local_peer_id(&swarm).clone()).collect();
+
+    swarms[0].add_address(&peer_ids[1], Protocol::Memory(port_base + 1).into());
+    swarms[1].add_address(&peer_ids[2], Protocol::Memory(port_base + 2).into());
+
+    let target_key = multihash::encode(Hash::SHA2256, &vec![1,2,3]).unwrap();
+
+    // TODO: This is racy, need to be sequenced
+    swarms[0].get_value(&target_key);
+    swarms[1].put_value(target_key.clone(), vec![4,5,6]);
+
+    Runtime::new().unwrap().block_on(
+        future::poll_fn(move || -> Result<_, io::Error> {
+            for swarm in &mut swarms {
+                loop {
+                    match swarm.poll().unwrap() {
+                        Async::Ready(Some(KademliaOut::GetValueResult {
+                            result,
+                            closer_peers,
+                        })) => {
+                            assert_ne!(result, None);
+                            assert_eq!(closer_peers, None);
+                            let value = result.unwrap();
+                            assert_eq!(value.0, target_key);
+                            assert_eq!(value.1, vec![4,5,6]);
+                            return Ok(Async::Ready(()));
+                        }
+                        Async::Ready(Some(KademliaOut::PutValueResult { key })) => {
+                            assert_eq!(key, target_key);
+                        },
+                        Async::Ready(_) => (),
+                        Async::NotReady => break,
+                    }
+                }
+            }
+
+            Ok(Async::NotReady)
+        }))
+        .unwrap()
+}