botnet-core 0.2.0

botnet
Documentation 
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
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104

use crate::{AsBytes, BotnetResult, Bytes, Key};
use async_std::sync::{Arc, Mutex};
use async_trait::async_trait;
use std::collections::HashMap;

#[cfg(feature = "redisdb")]
use redis::{aio::Connection as RedisConnection, AsyncCommands, Client as RedisClient};

pub trait DatabaseKey:
    Key + AsBytes + std::cmp::Eq + std::hash::Hash + Send + Sync
{
}

#[async_trait]
pub trait Database {
    async fn set_key(&mut self, k: impl DatabaseKey, v: Bytes) -> BotnetResult<()>;
    async fn get_key(&self, k: impl DatabaseKey) -> BotnetResult<Option<Bytes>>;
    async fn set_bytes(&self, b: Bytes, v: Bytes) -> BotnetResult<()>;
    async fn get_bytes(&self, k: &Bytes) -> BotnetResult<Option<Bytes>>;
}

enum DbType {
    InMemory,
    #[allow(unused)]
    Redis,
}

pub struct InMemory {
    #[allow(unused)]
    db_type: DbType,
    items: Arc<Mutex<HashMap<Bytes, Bytes>>>,
}

impl InMemory {
    pub fn new() -> Self {
        Self {
            db_type: DbType::InMemory,
            items: Arc::new(Mutex::new(HashMap::default())),
        }
    }
}

impl Default for InMemory {
    fn default() -> Self {
        Self::new()
    }
}

#[async_trait]
impl Database for InMemory {
    async fn set_key(&mut self, k: impl DatabaseKey, v: Bytes) -> BotnetResult<()> {
        self.items.lock().await.insert(k.flatten(), v);
        Ok(())
    }

    async fn get_key(&self, k: impl DatabaseKey) -> BotnetResult<Option<Bytes>> {
        Ok(self.items.lock().await.remove(&k.flatten()))
    }

    async fn set_bytes(&self, k: Bytes, v: Bytes) -> BotnetResult<()> {
        self.items.lock().await.insert(k, v);
        Ok(())
    }

    async fn get_bytes(&self, k: &Bytes) -> BotnetResult<Option<Bytes>> {
        Ok(self.items.lock().await.remove(k))
    }
}

#[cfg(feature = "redisdb")]
pub struct Redis {
    conn: Arc<Mutex<RedisConnection>>,
}

#[cfg(feature = "redisdb")]
impl Redis {
    pub async fn new(url: &str) -> BotnetResult<Self> {
        let client = RedisClient::open(url)?;
        let conn = client.get_tokio_connection().await?;
        Ok(Self {
            conn: Arc::new(Mutex::new(conn)),
        })
    }
}

#[cfg(feature = "redisdb")]
#[async_trait]
impl Database for Redis {
    async fn set_key(&mut self, k: impl DatabaseKey, v: Bytes) -> BotnetResult<()> {
        unimplemented!()
    }

    async fn get_key(&self, k: impl DatabaseKey) -> BotnetResult<Option<Bytes>> {
        unimplemented!()
    }

    async fn set_bytes(&self, k: Bytes, v: Bytes) -> BotnetResult<()> {
        unimplemented!()
    }

    async fn get_bytes(&self, k: &Bytes) -> BotnetResult<Option<Bytes>> {
        unimplemented!()
    }
}