refactor addressbook sync

Signed-off-by: 21pages <pages21@163.com>

libs/hbb_common/src/config.rs

@@ -1,6 +1,7 @@
 use std::{
-    collections::HashMap,
+    collections::{HashMap, HashSet},
     fs,
+    io::{Read, Write},
     net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr},
     ops::{Deref, DerefMut},
     path::{Path, PathBuf},
@@ -13,10 +14,12 @@ use rand::Rng;
 use regex::Regex;
 use serde as de;
 use serde_derive::{Deserialize, Serialize};
+use serde_json;
 use sodiumoxide::base64;
 use sodiumoxide::crypto::sign;
 
 use crate::{
+    compress::{compress, decompress},
     log,
     password_security::{
         decrypt_str_or_original, decrypt_vec_or_original, encrypt_str_or_original,
@@ -31,6 +34,7 @@ pub const REG_INTERVAL: i64 = 12_000;
 pub const COMPRESS_LEVEL: i32 = 3;
 const SERIAL: i32 = 3;
 const PASSWORD_ENC_VERSION: &str = "00";
+const ENCRYPT_MAX_LEN: usize = 128;
 
 // config2 options
 #[cfg(target_os = "linux")]
@@ -57,6 +61,7 @@ lazy_static::lazy_static! {
     pub static ref APP_NAME: Arc<RwLock<String>> = Arc::new(RwLock::new("RustDesk".to_owned()));
     static ref KEY_PAIR: Arc<Mutex<Option<KeyPair>>> = Default::default();
     static ref USER_DEFAULT_CONFIG: Arc<RwLock<(UserDefaultConfig, Instant)>> = Arc::new(RwLock::new((UserDefaultConfig::load(), Instant::now())));
+    pub static ref NEW_STORED_PEER_CONFIG: Arc<Mutex<HashSet<String>>> = Default::default();
 }
 
 lazy_static::lazy_static! {
@@ -376,7 +381,8 @@ impl Config2 {
     fn store(&self) {
         let mut config = self.clone();
         if let Some(mut socks) = config.socks {
-            socks.password = encrypt_str_or_original(&socks.password, PASSWORD_ENC_VERSION);
+            socks.password =
+                encrypt_str_or_original(&socks.password, PASSWORD_ENC_VERSION, ENCRYPT_MAX_LEN);
             config.socks = Some(socks);
         }
         Config::store_(&config, "2");
@@ -485,8 +491,9 @@ impl Config {
 
     fn store(&self) {
         let mut config = self.clone();
-        config.password = encrypt_str_or_original(&config.password, PASSWORD_ENC_VERSION);
-        config.enc_id = encrypt_str_or_original(&config.id, PASSWORD_ENC_VERSION);
+        config.password =
+            encrypt_str_or_original(&config.password, PASSWORD_ENC_VERSION, ENCRYPT_MAX_LEN);
+        config.enc_id = encrypt_str_or_original(&config.id, PASSWORD_ENC_VERSION, ENCRYPT_MAX_LEN);
         config.id = "".to_owned();
         Config::store_(&config, "");
     }
@@ -980,15 +987,17 @@ impl PeerConfig {
     pub fn store(&self, id: &str) {
         let _lock = CONFIG.read().unwrap();
         let mut config = self.clone();
-        config.password = encrypt_vec_or_original(&config.password, PASSWORD_ENC_VERSION);
+        config.password =
+            encrypt_vec_or_original(&config.password, PASSWORD_ENC_VERSION, ENCRYPT_MAX_LEN);
         for opt in ["rdp_password", "os-username", "os-password"] {
             if let Some(v) = config.options.get_mut(opt) {
-                *v = encrypt_str_or_original(v, PASSWORD_ENC_VERSION)
+                *v = encrypt_str_or_original(v, PASSWORD_ENC_VERSION, ENCRYPT_MAX_LEN)
             }
         }
         if let Err(err) = store_path(Self::path(id), config) {
             log::error!("Failed to store config: {}", err);
         }
+        NEW_STORED_PEER_CONFIG.lock().unwrap().insert(id.to_owned());
     }
 
     pub fn remove(id: &str) {
@@ -1014,7 +1023,7 @@ impl PeerConfig {
         Config::with_extension(Config::path(path))
     }
 
-    pub fn peers() -> Vec<(String, SystemTime, PeerConfig)> {
+    pub fn peers(id_filters: Option<Vec<String>>) -> Vec<(String, SystemTime, PeerConfig)> {
         if let Ok(peers) = Config::path(PEERS).read_dir() {
             if let Ok(peers) = peers
                 .map(|res| res.map(|e| e.path()))
@@ -1027,7 +1036,6 @@ impl PeerConfig {
                             && p.extension().map(|p| p.to_str().unwrap_or("")) == Some("toml")
                     })
                     .map(|p| {
-                        let t = crate::get_modified_time(p);
                         let id = p
                             .file_stem()
                             .map(|p| p.to_str().unwrap_or(""))
@@ -1041,12 +1049,21 @@ impl PeerConfig {
                         } else {
                             id
                         };
-
-                        let c = PeerConfig::load(&id_decoded_string);
+                        (id_decoded_string, p)
+                    })
+                    .filter(|(id, _)| {
+                        let Some(filters) = &id_filters else {
+                            return true;
+                        };
+                        filters.contains(id)
+                    })
+                    .map(|(id, p)| {
+                        let t = crate::get_modified_time(p);
+                        let c = PeerConfig::load(&id);
                         if c.info.platform.is_empty() {
                             fs::remove_file(p).ok();
                         }
-                        (id_decoded_string, t, c)
+                        (id, t, c)
                     })
                     .filter(|p| !p.2.info.platform.is_empty())
                     .collect();
@@ -1445,6 +1462,74 @@ impl UserDefaultConfig {
     }
 }
 
+#[derive(Debug, Default, Serialize, Deserialize, Clone)]
+pub struct AbPeer {
+    #[serde(
+        default,
+        deserialize_with = "deserialize_string",
+        skip_serializing_if = "String::is_empty"
+    )]
+    pub id: String,
+    #[serde(
+        default,
+        deserialize_with = "deserialize_string",
+        skip_serializing_if = "String::is_empty"
+    )]
+    pub hash: String,
+}
+
+#[derive(Debug, Default, Serialize, Deserialize, Clone)]
+pub struct Ab {
+    #[serde(
+        default,
+        deserialize_with = "deserialize_string",
+        skip_serializing_if = "String::is_empty"
+    )]
+    pub access_token: String,
+    #[serde(default, deserialize_with = "deserialize_vec_abpeer")]
+    pub peers: Vec<AbPeer>,
+}
+
+impl Ab {
+    fn path() -> PathBuf {
+        let filename = format!("{}_ab", APP_NAME.read().unwrap().clone());
+        Config::path(filename)
+    }
+
+    pub fn store(json: String) {
+        if let Ok(mut file) = std::fs::File::create(Self::path()) {
+            let data = compress(json.as_bytes());
+            let max_len = 32 * 1024 * 1024;
+            if data.len() > max_len {
+                // not store original
+                return;
+            }
+            let data = encrypt_vec_or_original(&data, PASSWORD_ENC_VERSION, max_len);
+            file.write_all(&data).ok();
+        };
+    }
+
+    pub fn load() -> Ab {
+        if let Ok(mut file) = std::fs::File::open(Self::path()) {
+            let mut data = vec![];
+            if file.read_to_end(&mut data).is_ok() {
+                let (data, succ, _) = decrypt_vec_or_original(&data, PASSWORD_ENC_VERSION);
+                if succ {
+                    let data = decompress(&data);
+                    if let Ok(ab) = serde_json::from_str::<Ab>(&String::from_utf8_lossy(&data)) {
+                        return ab;
+                    }
+                }
+            }
+        };
+        Ab::default()
+    }
+
+    pub fn remove() {
+        std::fs::remove_file(Self::path()).ok();
+    }
+}
+
 // use default value when field type is wrong
 macro_rules! deserialize_default {
     ($func_name:ident, $return_type:ty) => {
@@ -1464,6 +1549,7 @@ deserialize_default!(deserialize_vec_u8, Vec<u8>);
 deserialize_default!(deserialize_vec_string, Vec<String>);
 deserialize_default!(deserialize_vec_i32_string_i32, Vec<(i32, String, i32)>);
 deserialize_default!(deserialize_vec_discoverypeer, Vec<DiscoveryPeer>);
+deserialize_default!(deserialize_vec_abpeer, Vec<AbPeer>);
 deserialize_default!(deserialize_keypair, KeyPair);
 deserialize_default!(deserialize_size, Size);
 deserialize_default!(deserialize_hashmap_string_string, HashMap<String, String>);