IpSubnet and SdlanSock's sqlx try_from

src/packet/register_super.rs

@@ -47,8 +47,8 @@ mod test {
                 port: 1,
                 has_v6: true,
                 v6_port: 2345,
-                v4: vec![0; 4],
-                v6: vec![1; 16],
+                v4: [0; 4],
+                v6: [1; 16],
             }),
             dev_addr: peer::IpSubnet::new(192, 24),
             pub_key: "public key",

src/peer.rs

@@ -1,11 +1,10 @@
 #![allow(unused)]
-use std::default::Default;
-use std::os::unix::net;
-use std::sync::atomic::{AtomicU32, AtomicU64, AtomicU8, Ordering};
-use std::sync::Mutex;
-
+use dashmap::DashMap;
 use serde::{Deserialize, Serialize};
 use sqlx::prelude::FromRow;
+use std::default::Default;
+use std::sync::atomic::{AtomicU32, AtomicU64, AtomicU8, Ordering};
+use std::sync::Mutex;
 
 #[derive(Debug)]
 pub struct Peer {
@@ -38,8 +37,8 @@ impl Peer {
                 port: 0,
                 has_v6: false,
                 v6_port: 0,
-                v4: vec![0; 4],
-                v6: vec![0; 16],
+                v4: [0; 4],
+                v6: [0; 16],
             }),
             pub_key: Mutex::new(vec![]),
             timeout: 0,
@@ -52,6 +51,8 @@ impl Peer {
 }
 
 /// IpSubnet, 对端ipv4信息
+///
+/*
 #[derive(Debug, Serialize, Deserialize, FromRow)]
 pub struct IpSubnet {
     net_addr: u32,
@@ -74,29 +75,47 @@ impl IpSubnet {
         self.net_bit_len
     }
 }
-pub struct IpSubnetAtomic {
+*/
+
+#[derive(Debug, Serialize, Deserialize, FromRow)]
+pub struct IpSubnet {
+    #[sqlx(try_from = "u32")]
     pub net_addr: AtomicU32,
+    #[sqlx(try_from = "u8")]
     pub net_bit_len: AtomicU8,
 }
 
-impl From<IpSubnet> for IpSubnetAtomic {
-    fn from(value: IpSubnet) -> Self {
-        Self {
-            net_addr: AtomicU32::new(value.net_addr()),
-            net_bit_len: AtomicU8::new(value.net_bit_len()),
-        }
-    }
-}
-
-impl IpSubnetAtomic {
+impl IpSubnet {
     pub fn new(net_addr: u32, net_bit_len: u8) -> Self {
         Self {
             net_addr: AtomicU32::new(net_addr),
             net_bit_len: AtomicU8::new(net_bit_len),
         }
     }
+    pub fn net_addr(&self) -> u32 {
+        self.net_addr.load(Ordering::Relaxed)
+    }
+    pub fn net_bit_len(&self) -> u8 {
+        self.net_bit_len.load(Ordering::Relaxed)
+    }
 }
 
+/// SdlanSock: 对端对外的ip信息，包括ipv4和ipv6
+#[derive(Debug, Serialize, Deserialize, PartialEq, sqlx::FromRow)]
+pub struct SdlanSock {
+    pub family: u8,
+    pub port: u32,
+    pub has_v6: bool,
+    pub v6_port: u32,
+    // pub v4: Vec<u8>,
+    // pub v6: Vec<u8>,
+    #[sqlx(try_from = "Vec<u8>")]
+    pub v4: [u8; 4],
+    #[sqlx(try_from = "Vec<u8>")]
+    pub v6: [u8; 16],
+}
+
+/*
 /// SdlanSock: 对端对外的ip信息，包括ipv4和ipv6
 #[derive(Debug, Serialize, Deserialize, PartialEq, sqlx::FromRow)]
 pub struct SdlanSock {
@@ -109,3 +128,111 @@ pub struct SdlanSock {
     // pub v4: [u8; 4],
     // pub v6: [u8; 16],
 }
+*/
+
+use std::borrow::Cow;
+use std::sync::Arc;
+pub struct PeerMap {
+    pub peers: DashMap<String, Arc<Peer>>,
+    // ip to peer's uuid, the ip is the logical ip
+    // of tun
+    pub peer_ip_map: DashMap<u32, String>,
+    /// peer's read only sdlansock
+    /// if a peer's sock changes, just remove
+    /// and re-insert it.
+    pub peer_sock: DashMap<String, SdlanSock>,
+}
+
+impl PeerMap {
+    pub fn new() -> Self {
+        Self {
+            peers: DashMap::new(),
+            peer_ip_map: DashMap::new(),
+            peer_sock: DashMap::new(),
+        }
+    }
+
+    pub fn insert_peer(&self, data: Arc<Peer>) {
+        let ip = data.dev_addr.net_addr.load(Ordering::Relaxed);
+        if ip != 0 {
+            let id = data.id.clone();
+            self.peer_ip_map.insert(ip, id);
+        }
+        self.peers.insert(data.id.clone(), data);
+    }
+
+    pub fn peer_match<F: Fn(&Arc<Peer>) -> bool>(&self, id: &str, f: F) -> Option<Arc<Peer>> {
+        if let Some(v) = self.peers.get(id) {
+            if f(&v) {
+                return Some(v.clone());
+            }
+        }
+        None
+    }
+
+    pub fn find_peer_by_id(&self, id: &str) -> Option<Arc<Peer>> {
+        if let Some(v) = self.peers.get(id) {
+            return Some(v.clone());
+        }
+        None
+    }
+
+    pub fn find_peer_by_sock(&self, sock: &SdlanSock) -> Option<Arc<Peer>> {
+        for k in self.peers.iter() {
+            let itersock = k.value().sock.lock().unwrap();
+            if itersock.family == sock.family
+                && itersock.port == sock.port
+                && itersock.v4 == sock.v4
+            {
+                return Some(k.value().clone());
+            }
+        }
+
+        None
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use crate::utils::gen_uuid;
+
+    use super::*;
+
+    #[test]
+    fn test_peer_map() {
+        let id = gen_uuid();
+        let pm = PeerMap::new();
+        let mut p = Peer::new(&id);
+        {
+            let mut sl = p.sock.lock().unwrap();
+            sl.family = 0;
+            sl.v4 = [1; 4];
+            sl.port = 20;
+        }
+        pm.insert_peer(Arc::new(p));
+        if let Some(p2) = pm.find_peer_by_id(&id) {
+            let pl = p2.sock.lock().unwrap();
+            assert_eq!(pl.family, 0);
+            assert_eq!(pl.port, 20);
+            assert_eq!(pl.v4, [1; 4]);
+        } else {
+            panic!("cannot find p2 with id");
+        }
+
+        if let Some(p3) = pm.find_peer_by_sock(&SdlanSock {
+            family: 0,
+            port: 20,
+            has_v6: true,
+            v6_port: 30,
+            v4: [1; 4],
+            v6: [1; 16],
+        }) {
+            let pl = p3.sock.lock().unwrap();
+            assert_eq!(pl.family, 0);
+            assert_eq!(pl.port, 20);
+            assert_eq!(pl.v4, [1; 4]);
+        } else {
+            panic!("failed to get peer with sock");
+        }
+    }
+}