refactor!: add ws endpoint, storage trait

registry/storage/mod.rs

@@ -0,0 +1,38 @@
+use crate::{
+    error::{Error, Result},
+    utils::Peer,
+};
+
+mod valkey;
+
+pub use valkey::RegisterRequest;
+
+pub enum Storage {
+    Valkey(valkey::ValkeyStorage),
+}
+
+pub trait StorageImpl {
+    async fn register_device(&self, request: &RegisterRequest) -> Result<()>;
+    async fn get_peers(&self) -> Result<Vec<Peer>>;
+}
+
+impl StorageImpl for Storage {
+    async fn register_device(&self, request: &RegisterRequest) -> Result<()> {
+        match self {
+            Self::Valkey(storage) => storage.register_device(request).await,
+        }
+    }
+
+    async fn get_peers(&self) -> Result<Vec<Peer>> {
+        match self {
+            Self::Valkey(storage) => storage.get_peers().await,
+        }
+    }
+}
+
+pub fn get_storage_from_env() -> Result<Storage> {
+    Ok(Storage::Valkey(valkey::ValkeyStorage {
+        valkey_client: redis::Client::open("redis://127.0.0.1:6379/")
+            .map_err(|e| Error::valkey_connect(e, "127.0.0.1:6379/".to_string()))?,
+    }))
+}

registry/storage/valkey.rs

@@ -0,0 +1,108 @@
+use std::collections::{HashMap, HashSet};
+use std::net::IpAddr;
+
+use ipnetwork::IpNetwork;
+use redis::AsyncTypedCommands;
+use serde::Deserialize;
+
+use crate::error::Result;
+use crate::utils::{Peer, WireguardPublicKey};
+use crate::{error::Error, storage::StorageImpl};
+
+pub struct ValkeyStorage {
+    pub valkey_client: redis::Client,
+}
+
+#[derive(Deserialize, Clone)]
+pub struct RegisterRequest {
+    pub public_ip: IpAddr,
+    pub public_key: WireguardPublicKey,
+    pub port: String,
+    pub allowed_ips: Vec<IpNetwork>,
+}
+
+impl StorageImpl for ValkeyStorage {
+    async fn register_device(&self, request: &RegisterRequest) -> Result<()> {
+        let mut conn = self
+            .valkey_client
+            .get_multiplexed_async_connection()
+            .await
+            .map_err(|e| Error::valkey_get_connection(e))?;
+        conn.hset_multiple::<_, _, _>(
+            format!("peer:{}", request.public_key.as_str()),
+            &[
+                ("public_ip", &request.public_ip.to_string()),
+                (
+                    "allowed_ips",
+                    &serde_json::to_string(&request.allowed_ips)
+                        .map_err(|e| Error::serialize_json(e, "serializing allowed_ips"))?,
+                ),
+                ("port", &request.port),
+            ],
+        )
+        .await
+        .map_err(|e| {
+            Error::add_peer(
+                e,
+                request.public_key.as_str(),
+                request.public_ip.to_string(),
+            )
+        })?;
+        conn.sadd("peers", &request.public_key.as_str())
+            .await
+            .map_err(|e| {
+                Error::add_peer(
+                    e,
+                    request.public_key.as_str(),
+                    request.public_ip.to_string(),
+                )
+            })?;
+
+        Ok(())
+    }
+
+    async fn get_peers(&self) -> Result<Vec<Peer>> {
+        let mut conn = self
+            .valkey_client
+            .get_multiplexed_async_connection()
+            .await
+            .map_err(|e| Error::valkey_get_connection(e))?;
+        let keys: HashSet<String> = conn
+            .smembers("peers")
+            .await
+            .map_err(|e| Error::get_peer(e))?;
+
+        if keys.is_empty() {
+            return Ok(vec![]);
+        }
+
+        let keys: Vec<String> = keys.into_iter().collect();
+        let mut pipe = redis::pipe();
+        for key in keys.iter() {
+            pipe.hgetall(format!("peer:{key}"));
+        }
+
+        let peers: Vec<Peer> = pipe
+            .query_async::<Vec<HashMap<String, String>>>(&mut conn)
+            .await
+            .map_err(|e| Error::get_peer(e))?
+            .into_iter()
+            .zip(keys.iter())
+            .map(|(peer, key): (HashMap<String, String>, &String)| {
+                let allowed_ips: Vec<IpNetwork> = peer
+                    .get("allowed_ips")
+                    .map(|s| serde_json::from_str(s).unwrap_or_default())
+                    .unwrap_or_default();
+
+                Peer {
+                    public_key: key.clone(),
+                    public_ip: peer.get("public_ip").unwrap().to_string(),
+                    port: peer.get("port").unwrap().to_string(),
+                    allowed_ips,
+                }
+            })
+            .collect();
+
+        Ok(peers)
+    }
+}