Skip to content

Commit 41fa8e2

Browse files
dmitry-markindmitry-markin
authored
kad: Providers part 3: publish provider records (start providing) (#234)
Implement `ADD_PROVIDER` network request and execute `FIND_NODE` query to publish local provider to the target peers.
1 parent d50ec10 commit 41fa8e2

File tree

4 files changed

+246
-28
lines changed

4 files changed

+246
-28
lines changed

‎src/protocol/libp2p/kademlia/handle.rs‎

Lines changed: 35 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -130,6 +130,18 @@ pub(crate) enum KademliaCommand {
130130
query_id: QueryId,
131131
},
132132

133+
/// Register as a content provider for `key`.
134+
StartProviding {
135+
/// Provided key.
136+
key: RecordKey,
137+
138+
/// Our external addresses to publish.
139+
public_addresses: Vec<Multiaddr>,
140+
141+
/// Query ID for the query.
142+
query_id: QueryId,
143+
},
144+
133145
/// Store record locally.
134146
StoreRecord {
135147
// Record.
@@ -175,7 +187,8 @@ pub enum KademliaEvent {
175187
},
176188

177189
/// `PUT_VALUE` query succeeded.
178-
PutRecordSucess {
190+
// TODO: this is never emitted. Implement + add `AddProviderSuccess`.
191+
PutRecordSuccess {
179192
/// Query ID.
180193
query_id: QueryId,
181194

@@ -299,6 +312,27 @@ impl KademliaHandle {
299312
query_id
300313
}
301314

315+
/// Register as a content provider on the DHT.
316+
///
317+
/// Register the local peer ID & its `public_addresses` as a provider for a given `key`.
318+
pub async fn start_providing(
319+
&mut self,
320+
key: RecordKey,
321+
public_addresses: Vec<Multiaddr>,
322+
) -> QueryId {
323+
let query_id = self.next_query_id();
324+
let _ = self
325+
.cmd_tx
326+
.send(KademliaCommand::StartProviding {
327+
key,
328+
public_addresses,
329+
query_id,
330+
})
331+
.await;
332+
333+
query_id
334+
}
335+
302336
/// Store the record in the local store. Used in combination with
303337
/// [`IncomingRecordValidationMode::Manual`].
304338
pub async fn store_record(&mut self, record: Record) {

‎src/protocol/libp2p/kademlia/message.rs‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -172,8 +172,7 @@ impl KademliaMessage {
172172
}
173173

174174
/// Create `ADD_PROVIDER` message with `provider`.
175-
#[allow(unused)]
176-
pub fn add_provider(provider: ProviderRecord) -> Vec<u8> {
175+
pub fn add_provider(provider: ProviderRecord) -> Bytes {
177176
let peer = KademliaPeer::new(
178177
provider.provider,
179178
provider.addresses,
@@ -187,10 +186,10 @@ impl KademliaMessage {
187186
..Default::default()
188187
};
189188

190-
let mut buf = Vec::with_capacity(message.encoded_len());
189+
let mut buf = BytesMut::with_capacity(message.encoded_len());
191190
message.encode(&mut buf).expect("Vec<u8> to provide needed capacity");
192191

193-
buf
192+
buf.freeze()
194193
}
195194

196195
/// Create `GET_PROVIDERS` request for `key`.

‎src/protocol/libp2p/kademlia/mod.rs‎

Lines changed: 137 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -83,13 +83,16 @@ mod schema {
8383
}
8484

8585
/// Peer action.
86-
#[derive(Debug)]
86+
#[derive(Debug, Clone)]
8787
enum PeerAction {
8888
/// Send `FIND_NODE` message to peer.
8989
SendFindNode(QueryId),
9090

9191
/// Send `PUT_VALUE` message to peer.
9292
SendPutValue(Bytes),
93+
94+
/// Send `ADD_PROVIDER` message to peer.
95+
SendAddProvider(Bytes),
9396
}
9497

9598
/// Peer context.
@@ -335,7 +338,12 @@ impl Kademlia {
335338
}
336339
}
337340
Some(PeerAction::SendPutValue(message)) => {
338-
tracing::trace!(target: LOG_TARGET, ?peer, "send `PUT_VALUE` response");
341+
tracing::trace!(target: LOG_TARGET, ?peer, "send `PUT_VALUE` message");
342+
343+
self.executor.send_message(peer, message, substream);
344+
}
345+
Some(PeerAction::SendAddProvider(message)) => {
346+
tracing::trace!(target: LOG_TARGET, ?peer, "send `ADD_PROVIDER` message");
339347

340348
self.executor.send_message(peer, message, substream);
341349
}
@@ -758,6 +766,35 @@ impl Kademlia {
758766

759767
Ok(())
760768
}
769+
QueryAction::AddProviderToFoundNodes { provider, peers } => {
770+
tracing::trace!(
771+
target: LOG_TARGET,
772+
provided_key = ?provider.key,
773+
num_peers = ?peers.len(),
774+
"add provider record to found peers",
775+
);
776+
777+
let provided_key = provider.key.clone();
778+
let message = KademliaMessage::add_provider(provider);
779+
780+
for peer in peers {
781+
if let Err(error) = self.open_substream_or_dial(
782+
peer.peer,
783+
PeerAction::SendAddProvider(message.clone()),
784+
None,
785+
) {
786+
tracing::debug!(
787+
target: LOG_TARGET,
788+
?peer,
789+
?provided_key,
790+
?error,
791+
"failed to add provider record to peer",
792+
)
793+
}
794+
}
795+
796+
Ok(())
797+
}
761798
QueryAction::GetRecordQueryDone { query_id, records } => {
762799
let _ = self
763800
.event_tx
@@ -794,7 +831,11 @@ impl Kademlia {
794831
event = self.service.next() => match event {
795832
Some(TransportEvent::ConnectionEstablished { peer, .. }) => {
796833
if let Err(error) = self.on_connection_established(peer) {
797-
tracing::debug!(target: LOG_TARGET, ?error, "failed to handle established connection");
834+
tracing::debug!(
835+
target: LOG_TARGET,
836+
?error,
837+
"failed to handle established connection",
838+
);
798839
}
799840
}
800841
Some(TransportEvent::ConnectionClosed { peer }) => {
@@ -804,7 +845,10 @@ impl Kademlia {
804845
match direction {
805846
Direction::Inbound => self.on_inbound_substream(peer, substream).await,
806847
Direction::Outbound(substream_id) => {
807-
if let Err(error) = self.on_outbound_substream(peer, substream_id, substream).await {
848+
if let Err(error) = self
849+
.on_outbound_substream(peer, substream_id, substream)
850+
.await
851+
{
808852
tracing::debug!(
809853
target: LOG_TARGET,
810854
?peer,
@@ -819,22 +863,41 @@ impl Kademlia {
819863
Some(TransportEvent::SubstreamOpenFailure { substream, error }) => {
820864
self.on_substream_open_failure(substream, error).await;
821865
}
822-
Some(TransportEvent::DialFailure { peer, address, .. }) => self.on_dial_failure(peer, address),
866+
Some(TransportEvent::DialFailure { peer, address, .. }) =>
867+
self.on_dial_failure(peer, address),
823868
None => return Err(Error::EssentialTaskClosed),
824869
},
825870
context = self.executor.next() => {
826871
let QueryContext { peer, query_id, result } = context.unwrap();
827872

828873
match result {
829874
QueryResult::SendSuccess { substream } => {
830-
tracing::trace!(target: LOG_TARGET, ?peer, query = ?query_id, "message sent to peer");
875+
tracing::trace!(
876+
target: LOG_TARGET,
877+
?peer,
878+
query = ?query_id,
879+
"message sent to peer",
880+
);
831881
let _ = substream.close().await;
832882
}
833883
QueryResult::ReadSuccess { substream, message } => {
834-
tracing::trace!(target: LOG_TARGET, ?peer, query = ?query_id, "message read from peer");
884+
tracing::trace!(target: LOG_TARGET,
885+
?peer,
886+
query = ?query_id,
887+
"message read from peer",
888+
);
835889

836-
if let Err(error) = self.on_message_received(peer, query_id, message, substream).await {
837-
tracing::debug!(target: LOG_TARGET, ?peer, ?error, "failed to process message");
890+
if let Err(error) = self.on_message_received(
891+
peer,
892+
query_id,
893+
message,
894+
substream
895+
).await {
896+
tracing::debug!(target: LOG_TARGET,
897+
?peer,
898+
?error,
899+
"failed to process message",
900+
);
838901
}
839902
}
840903
QueryResult::SubstreamClosed | QueryResult::Timeout => {
@@ -853,22 +916,36 @@ impl Kademlia {
853916
command = self.cmd_rx.recv() => {
854917
match command {
855918
Some(KademliaCommand::FindNode { peer, query_id }) => {
856-
tracing::debug!(target: LOG_TARGET, ?peer, query = ?query_id, "starting `FIND_NODE` query");
919+
tracing::debug!(
920+
target: LOG_TARGET,
921+
?peer,
922+
query = ?query_id,
923+
"starting `FIND_NODE` query",
924+
);
857925

858926
self.engine.start_find_node(
859927
query_id,
860928
peer,
861-
self.routing_table.closest(Key::from(peer), self.replication_factor).into()
929+
self.routing_table
930+
.closest(Key::from(peer), self.replication_factor)
931+
.into()
862932
);
863933
}
864934
Some(KademliaCommand::PutRecord { mut record, query_id }) => {
865-
tracing::debug!(target: LOG_TARGET, query = ?query_id, key = ?record.key, "store record to DHT");
935+
tracing::debug!(
936+
target: LOG_TARGET,
937+
query = ?query_id,
938+
key = ?record.key,
939+
"store record to DHT",
940+
);
866941

867942
// For `PUT_VALUE` requests originating locally we are always the publisher.
868943
record.publisher = Some(self.local_key.clone().into_preimage());
869944

870945
// Make sure TTL is set.
871-
record.expires = record.expires.or_else(|| Some(Instant::now() + self.record_ttl));
946+
record.expires = record
947+
.expires
948+
.or_else(|| Some(Instant::now() + self.record_ttl));
872949

873950
let key = Key::new(record.key.clone());
874951

@@ -880,11 +957,23 @@ impl Kademlia {
880957
self.routing_table.closest(key, self.replication_factor).into(),
881958
);
882959
}
883-
Some(KademliaCommand::PutRecordToPeers { mut record, query_id, peers, update_local_store }) => {
884-
tracing::debug!(target: LOG_TARGET, query = ?query_id, key = ?record.key, "store record to DHT to specified peers");
960+
Some(KademliaCommand::PutRecordToPeers {
961+
mut record,
962+
query_id,
963+
peers,
964+
update_local_store,
965+
}) => {
966+
tracing::debug!(
967+
target: LOG_TARGET,
968+
query = ?query_id,
969+
key = ?record.key,
970+
"store record to DHT to specified peers",
971+
);
885972

886973
// Make sure TTL is set.
887-
record.expires = record.expires.or_else(|| Some(Instant::now() + self.record_ttl));
974+
record.expires = record
975+
.expires
976+
.or_else(|| Some(Instant::now() + self.record_ttl));
888977

889978
if update_local_store {
890979
self.store.put(record.clone());
@@ -898,7 +987,8 @@ impl Kademlia {
898987

899988
match self.routing_table.entry(Key::from(peer)) {
900989
KBucketEntry::Occupied(entry) => Some(entry.clone()),
901-
KBucketEntry::Vacant(entry) if !entry.addresses.is_empty() => Some(entry.clone()),
990+
KBucketEntry::Vacant(entry) if !entry.addresses.is_empty() =>
991+
Some(entry.clone()),
902992
_ => None,
903993
}
904994
}).collect();
@@ -909,6 +999,36 @@ impl Kademlia {
909999
peers,
9101000
);
9111001
}
1002+
Some(KademliaCommand::StartProviding {
1003+
key,
1004+
public_addresses,
1005+
query_id
1006+
}) => {
1007+
tracing::debug!(
1008+
target: LOG_TARGET,
1009+
query = ?query_id,
1010+
?key,
1011+
?public_addresses,
1012+
"register as content provider"
1013+
);
1014+
1015+
let provider = ProviderRecord {
1016+
key: key.clone(),
1017+
provider: self.service.local_peer_id(),
1018+
addresses: public_addresses,
1019+
expires: Instant::now() + self.provider_ttl,
1020+
};
1021+
1022+
self.store.put_provider(provider.clone());
1023+
1024+
self.engine.start_add_provider(
1025+
query_id,
1026+
provider,
1027+
self.routing_table
1028+
.closest(Key::new(key), self.replication_factor)
1029+
.into(),
1030+
);
1031+
}
9121032
Some(KademliaCommand::GetRecord { key, quorum, query_id }) => {
9131033
tracing::debug!(target: LOG_TARGET, ?key, "get record from DHT");
9141034

0 commit comments

Comments
 (0)