1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
//! This module provides route table abstraction for `Kad`.

use std::collections::HashMap;

use async_trait::async_trait;
use futures::lock::Mutex;
use xstack::driver_wrapper;

use crate::KBucketKey;

/// A record store must implement the `Driver-*` traits in this module.
///
/// A `KadStore` is the kademlia's key/value pair storage service,
/// developers can call [`set_store`](crate::KademliaOptions::set_store)
/// to set the customise implementation, the default implementation used by `KademliaRouter` is the [`KadMemoryStore`](crate::KadMemoryStore).
pub mod syscall {
    use async_trait::async_trait;

    use crate::KBucketKey;

    /// A trait that provides functions to access peer informations.
    #[async_trait]
    pub trait DriverKadStore: Sync + Send {
        /// insert a new kad record.
        async fn insert(&self, key: KBucketKey, record: Vec<u8>) -> std::io::Result<()>;

        /// remove record by key.
        async fn remove(&self, key: &KBucketKey) -> std::io::Result<Option<Vec<u8>>>;

        /// Get the record by key.
        async fn get(&self, key: &KBucketKey) -> std::io::Result<Option<Vec<u8>>>;
    }
}

driver_wrapper!(
    ["A type wrapper of [`DriverKadStore`](syscall::DriverKadStore)"]
    KadStore[syscall::DriverKadStore]
);

pub struct KadMemoryStore(Mutex<HashMap<KBucketKey, Vec<u8>>>);

impl KadMemoryStore {
    pub fn new() -> Self {
        KadMemoryStore(Default::default())
    }
}

#[async_trait]
impl syscall::DriverKadStore for KadMemoryStore {
    /// insert a new kad record.
    async fn insert(&self, key: KBucketKey, record: Vec<u8>) -> std::io::Result<()> {
        self.0.lock().await.insert(key, record);

        Ok(())
    }

    /// remove record by key.
    async fn remove(&self, key: &KBucketKey) -> std::io::Result<Option<Vec<u8>>> {
        Ok(self.0.lock().await.remove(key))
    }

    /// Get the record by key.
    async fn get(&self, key: &KBucketKey) -> std::io::Result<Option<Vec<u8>>> {
        Ok(self.0.lock().await.get(key).map(|record| record.to_owned()))
    }
}