botnet_core/
database.rs

1use crate::{AsBytes, BotnetResult, Bytes, Key};
2use async_std::sync::{Arc, Mutex};
3use async_trait::async_trait;
4use std::collections::HashMap;
5
6#[cfg(feature = "redisdb")]
7use redis::{aio::Connection as RedisConnection, AsyncCommands, Client as RedisClient};
8
9pub trait DatabaseKey:
10    Key + AsBytes + std::cmp::Eq + std::hash::Hash + Send + Sync
11{
12}
13
14#[async_trait]
15pub trait Database {
16    async fn set_key(&mut self, k: impl DatabaseKey, v: Bytes) -> BotnetResult<()>;
17    async fn get_key(&self, k: impl DatabaseKey) -> BotnetResult<Option<Bytes>>;
18    async fn set_bytes(&self, b: Bytes, v: Bytes) -> BotnetResult<()>;
19    async fn get_bytes(&self, k: &Bytes) -> BotnetResult<Option<Bytes>>;
20}
21
22enum DbType {
23    InMemory,
24    #[allow(unused)]
25    Redis,
26}
27
28pub struct InMemory {
29    #[allow(unused)]
30    db_type: DbType,
31    items: Arc<Mutex<HashMap<Bytes, Bytes>>>,
32}
33
34impl InMemory {
35    pub fn new() -> Self {
36        Self {
37            db_type: DbType::InMemory,
38            items: Arc::new(Mutex::new(HashMap::default())),
39        }
40    }
41}
42
43impl Default for InMemory {
44    fn default() -> Self {
45        Self::new()
46    }
47}
48
49#[async_trait]
50impl Database for InMemory {
51    async fn set_key(&mut self, k: impl DatabaseKey, v: Bytes) -> BotnetResult<()> {
52        self.items.lock().await.insert(k.flatten(), v);
53        Ok(())
54    }
55
56    async fn get_key(&self, k: impl DatabaseKey) -> BotnetResult<Option<Bytes>> {
57        Ok(self.items.lock().await.remove(&k.flatten()))
58    }
59
60    async fn set_bytes(&self, k: Bytes, v: Bytes) -> BotnetResult<()> {
61        self.items.lock().await.insert(k, v);
62        Ok(())
63    }
64
65    async fn get_bytes(&self, k: &Bytes) -> BotnetResult<Option<Bytes>> {
66        Ok(self.items.lock().await.remove(k))
67    }
68}
69
70#[cfg(feature = "redisdb")]
71pub struct Redis {
72    conn: Arc<Mutex<RedisConnection>>,
73}
74
75#[cfg(feature = "redisdb")]
76impl Redis {
77    pub async fn new(url: &str) -> BotnetResult<Self> {
78        let client = RedisClient::open(url)?;
79        let conn = client.get_tokio_connection().await?;
80        Ok(Self {
81            conn: Arc::new(Mutex::new(conn)),
82        })
83    }
84}
85
86#[cfg(feature = "redisdb")]
87#[async_trait]
88impl Database for Redis {
89    async fn set_key(&mut self, k: impl DatabaseKey, v: Bytes) -> BotnetResult<()> {
90        unimplemented!()
91    }
92
93    async fn get_key(&self, k: impl DatabaseKey) -> BotnetResult<Option<Bytes>> {
94        unimplemented!()
95    }
96
97    async fn set_bytes(&self, k: Bytes, v: Bytes) -> BotnetResult<()> {
98        unimplemented!()
99    }
100
101    async fn get_bytes(&self, k: &Bytes) -> BotnetResult<Option<Bytes>> {
102        unimplemented!()
103    }
104}
botnet_core/database.rs

botnet_core/
database.rs
