Add restore subscription/connections

Author: the2pizza

dev/dev.sql

@@ -158,7 +158,7 @@ ALTER TABLE connections RENAME COLUMN user_id TO subscription_id;
 
 CREATE TABLE subscriptions (
     id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
-    expires_at TIMESTAMP WITH TIME ZONE NOT NULL,
+    expires_at TIMESTAMP WITH TIME ZONE ,
     referred_by CHAR(13),  
 
     created_at TIMESTAMP WITH TIME ZONE DEFAULT now(),

src/bin/api/core/http/handlers/sub.rs

@@ -545,11 +545,6 @@ where
                 "Subscription {}  are expired",
                 sub_param.id
             ))));
-        } else {
-            return Ok(Box::new(http::not_found(&format!(
-                "Subscription {}  are not found",
-                sub_param.id
-            ))));
         }
     }
 

src/bin/api/core/tasks.rs

@@ -3,6 +3,7 @@ use chrono::NaiveTime;
 use chrono::TimeZone;
 use chrono::Utc;
 use futures::future::join_all;
+use pony::http::requests::ConnUpdateRequest;
 use rand::Rng;
 use std::collections::HashMap;
 use std::time::Duration;
@@ -37,6 +38,7 @@ pub trait Tasks {
     async fn collect_conn_stat(&self) -> Result<()>;
     async fn cleanup_expired_connections(&self, interval_sec: u64, publisher: ZmqPublisher);
     async fn cleanup_expired_subscriptions(&self, interval_sec: u64, publisher: ZmqPublisher);
+    async fn restore_subscriptions(&self, interval_sec: u64, publisher: ZmqPublisher);
 }
 
 #[async_trait]
@@ -168,6 +170,77 @@ impl Tasks for Api<HashMap<String, Vec<Node>>, Connection, Subscription> {
         }
     }
 
+    async fn restore_subscriptions(&self, interval_sec: u64, publisher: ZmqPublisher) {
+        let mut interval = tokio::time::interval(Duration::from_secs(interval_sec));
+
+        loop {
+            interval.tick().await;
+            log::debug!("Run restore subscriptions task");
+
+            let expired_subs: Vec<uuid::Uuid> = {
+                let mem = self.sync.memory.read().await;
+                mem.subscriptions
+                    .iter()
+                    .filter_map(|(id, sub)| if sub.is_active() { Some(*id) } else { None })
+                    .collect()
+            };
+
+            for sub_id in expired_subs {
+                let conns_to_restore: Vec<(uuid::Uuid, Connection)> = {
+                    let mem = self.sync.memory.read().await;
+                    mem.connections
+                        .get_by_subscription_id(&sub_id)
+                        .map(|conns| {
+                            conns
+                                .iter()
+                                .filter(|(_id, c)| c.get_deleted())
+                                .filter_map(|(id, c)| Some((*id, c.clone().into())))
+                                .collect()
+                        })
+                        .unwrap_or_default()
+                };
+
+                for (conn_id, conn) in conns_to_restore {
+                    let msg = conn.as_update_message(&conn_id);
+                    if let Ok(bytes) = rkyv::to_bytes::<_, 1024>(&msg) {
+                        let key = conn
+                            .node_id
+                            .map(|id| id.to_string())
+                            .unwrap_or_else(|| conn.get_env());
+                        let _ = publisher.send_binary(&key, bytes.as_ref()).await;
+                    }
+
+                    let conn_upd = ConnUpdateRequest {
+                        env: Some(conn.get_env()),
+                        is_deleted: Some(false),
+                        password: conn.get_password(),
+                        days: None,
+                    };
+
+                    match SyncOp::update_conn(&self.sync, &conn_id, conn_upd).await {
+                        Ok(StorageOperationStatus::Updated(_)) => {
+                            log::info!("Expired connection {} restored", conn_id);
+                        }
+                        Ok(status) => {
+                            log::warn!(
+                                "Connection {} could not be restored: {:?}",
+                                conn_id,
+                                status
+                            );
+                        }
+                        Err(e) => {
+                            log::error!(
+                                "Failed to restore expired connection {}: {:?}",
+                                conn_id,
+                                e
+                            );
+                        }
+                    }
+                }
+            }
+        }
+    }
+
     async fn periodic_db_sync(&self, interval_sec: u64) {
         let base = Duration::from_secs(interval_sec);
 

src/bin/api/main.rs

@@ -175,7 +175,7 @@ async fn main() -> Result<()> {
     let _ = tokio::spawn({
         let api = api.clone();
         let publisher = publisher.clone();
-        let job_interval = Duration::from_secs(60);
+        let job_interval = Duration::from_secs(settings.api.subscription_expire_interval);
         log::info!("cleanup_expired_subscriptions task started");
 
         async move {
@@ -184,6 +184,18 @@ async fn main() -> Result<()> {
         }
     });
 
+    let _ = tokio::spawn({
+        let api = api.clone();
+        let publisher = publisher.clone();
+        let job_interval = Duration::from_secs(settings.api.subscription_restore_interval);
+        log::info!("restore_subscriptions task started");
+
+        async move {
+            api.restore_subscriptions(job_interval.as_secs(), publisher)
+                .await;
+        }
+    });
+
     let api = api.clone();
     let api_handle = tokio::spawn(async move {
         if let Err(e) = api.run(settings.api.hostname).await {

src/config/settings.rs

@@ -80,12 +80,6 @@ fn default_api_token() -> String {
 fn default_label() -> String {
     "🏴‍☠️🏴‍☠️🏴‍☠️ dev".to_string()
 }
-fn default_node_healthcheck_timeout() -> i16 {
-    60
-}
-fn default_conn_limit_check_interval() -> u64 {
-    60
-}
 
 fn default_stat_job_interval() -> u64 {
     60
@@ -117,12 +111,24 @@ fn default_db_sync_interval_sec() -> u64 {
     300
 }
 
+fn default_subscription_restore_interval_sec() -> u64 {
+    60
+}
+
+fn default_subscription_expire_interval_sec() -> u64 {
+    60
+}
+
+fn default_node_healthcheck_timeout() -> i16 {
+    60
+}
+
 fn default_max_bandwidth_bps() -> i64 {
     100_000_000
 }
 
 fn default_hostname() -> String {
-    "http://localhost".to_string()
+    "http://localhost:5005".to_string()
 }
 
 #[derive(Clone, Debug, Deserialize, Default)]
@@ -131,10 +137,6 @@ pub struct ApiServiceConfig {
     pub address: Option<Ipv4Addr>,
     #[serde(default = "default_api_web_port")]
     pub port: u16,
-    #[serde(default = "default_node_healthcheck_timeout")]
-    pub node_health_check_timeout: i16,
-    #[serde(default = "default_conn_limit_check_interval")]
-    pub conn_limit_check_interval: u64,
     #[serde(default = "default_collect_conn_stat_interval")]
     pub collect_conn_stat_interval: u64,
     #[serde(default = "default_healthcheck_interval")]
@@ -151,6 +153,12 @@ pub struct ApiServiceConfig {
     pub db_sync_interval_sec: u64,
     #[serde(default = "default_hostname")]
     pub hostname: String,
+    #[serde(default = "default_subscription_restore_interval_sec")]
+    pub subscription_restore_interval: u64,
+    #[serde(default = "default_subscription_expire_interval_sec")]
+    pub subscription_expire_interval: u64,
+    #[serde(default = "default_node_healthcheck_timeout")]
+    pub node_health_check_timeout: i16,
 }
 
 #[derive(Clone, Debug, Deserialize, Default)]

src/memory/storage/subscription.rs

@@ -22,13 +22,6 @@ where
 {
     fn count_invited_by(&self, referral_code: &str) -> usize {
         self.values()
-            .inspect(|s| {
-                log::debug!(
-                    "cmp: referred_by={:?}  target={:?}",
-                    s.referred_by(),
-                    referral_code
-                );
-            })
             .filter(|s| s.referred_by().as_deref() == Some(referral_code))
             .count()
     }

src/xray_op/client.rs

@@ -145,7 +145,9 @@ impl HandlerActions for Arc<Mutex<HandlerClient>> {
                     "Remove SS user error, password not provided".to_string(),
                 ))
             }
-            Tag::Wireguard => todo!(),
+            Tag::Wireguard => Err(crate::PonyError::Custom(
+                "Removing Wireguard is not implemented".to_string(),
+            )),
         }
     }
 }