Initial Commit

client/src/config.rs

@@ -0,0 +1,20 @@
+use common::updater::SyncContext;
+use communicator::ConnectConfig;
+use serde::{Deserialize, Serialize};
+use std::collections::HashMap;
+
+#[derive(Debug, Clone, Default, Serialize, Deserialize)]
+pub struct ClientConfig {
+    pub log_level: Option<String>,
+    #[serde(flatten, default)]
+    pub connect: ConnectConfig,
+    pub profiles: HashMap<String, Profile>,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct Profile {
+    #[serde(flatten)]
+    pub sync_context: SyncContext,
+    pub path: String,
+    pub interval: Option<u64>
+}

client/src/http.rs

@@ -0,0 +1,58 @@
+use crate::config::Profile;
+use anyhow::anyhow;
+use bytes::Bytes;
+use common::http::{CloseRequest, DownloadRequest, SyncResponse};
+use common::stream::client_receive;
+use communicator::http::{json_from_response, request, response_to_async_read, text_from_response};
+use h2::client;
+use std::path::Path;
+use std::pin::pin;
+
+pub async fn sync(
+    client: &mut client::SendRequest<Bytes>,
+    profile: &Profile,
+) -> anyhow::Result<SyncResponse> {
+    let mut response = request(
+        client,
+        "/sync",
+        serde_json::to_string(&profile.sync_context)?,
+    )
+    .await?;
+    if !response.status().is_success() {
+        let body = text_from_response(&mut response).await?;
+        Err(anyhow!("Failed to request '/sync': {}", body))?;
+    }
+
+    json_from_response(&mut response).await
+}
+
+pub async fn download(
+    client: &mut client::SendRequest<Bytes>,
+    req: &DownloadRequest,
+    profile: &Profile,
+) -> anyhow::Result<()> {
+    let mut response = request(client, "/download", serde_json::to_string(req)?).await?;
+    if !response.status().is_success() {
+        let body = text_from_response(&mut response).await?;
+        Err(anyhow!("Failed to request '/download': {}", body))?;
+    }
+
+    client_receive(
+        pin!(response_to_async_read(response)),
+        Path::new(&profile.path),
+    )
+    .await
+}
+
+pub async fn close(
+    client: &mut client::SendRequest<Bytes>,
+    req: &CloseRequest,
+) -> anyhow::Result<()> {
+    let mut response = request(client, "/close", serde_json::to_string(req)?).await?;
+    if response.status().is_success() {
+        return Ok(());
+    }
+
+    let body = text_from_response(&mut response).await?;
+    Err(anyhow!("Failed to close request: {}", body))
+}

client/src/main.rs

@@ -0,0 +1,175 @@
+use crate::config::{ClientConfig, Profile};
+use crate::task::run;
+use communicator::{Identity, TunnelEndpoint, TunnelListener, connect_tunnel};
+use lazy_static::lazy_static;
+use log::{LevelFilter, error, info};
+use simplelog::{ColorChoice, Config, TermLogger, TerminalMode};
+use std::fs;
+use std::str::FromStr;
+use std::sync::Arc;
+use std::time::Duration;
+use structopt::StructOpt;
+use tokio::sync::{OwnedSemaphorePermit, Semaphore};
+use tokio::time::sleep;
+use tokio_util::sync::CancellationToken;
+
+mod config;
+mod http;
+mod task;
+
+#[derive(StructOpt)]
+struct CommandOpt {
+    #[structopt(short = "p", long)]
+    pub profile: Vec<String>,
+}
+
+#[tokio::main]
+async fn main() -> anyhow::Result<()> {
+    let log_level = LevelFilter::from_str(
+        CONFIG
+            .log_level
+            .clone()
+            .unwrap_or("INFO".to_string())
+            .as_str(),
+    )
+    .unwrap_or(LevelFilter::Info);
+
+    TermLogger::init(
+        log_level,
+        Config::default(),
+        TerminalMode::Mixed,
+        ColorChoice::Auto,
+    )?;
+
+    let command_opts = CommandOpt::from_args();
+
+    if command_opts.profile.is_empty() {
+        error!("No profile specified. Use -p or --profile to specify at least one profile.");
+        std::process::exit(1);
+    }
+
+    let profiles = command_opts
+        .profile
+        .iter()
+        .map(|p| {
+            let profile = CONFIG.profiles.get(p.as_str());
+            match profile {
+                Some(profile) => Ok((p.to_string(), profile.clone())),
+                None => Err(anyhow::anyhow!("Profile `{}` not found in config", p)),
+            }
+        })
+        .collect::<Result<Vec<_>, _>>()?;
+
+    let mut tasks = vec![];
+    for profile in profiles {
+        let profile = Arc::new(profile.clone());
+        let semaphore = Arc::new(Semaphore::new(1));
+        let cancel_token = CancellationToken::new();
+        let post_task = {
+            async |profile: Arc<(String, Profile)>,
+                   permit: OwnedSemaphorePermit,
+                   cancel_token: CancellationToken| {
+                match profile.1.interval {
+                    None => {
+                        cancel_token.cancel();
+                    }
+                    Some(interval) => {
+                        sleep(Duration::from_secs(interval)).await;
+                    }
+                }
+                drop(permit);
+            }
+        };
+        for server_conf in &CONFIG.connect.server {
+            let server_conf = server_conf.clone().into_tunnel_config(Identity::Client)?;
+            let url = server_conf.url.clone();
+            let server = TunnelListener::bind(server_conf).await?;
+            let semaphore = semaphore.clone();
+            let cancel_token = cancel_token.clone();
+            let profile = profile.clone();
+            info!("tcp server started on {}", url);
+            tasks.push(tokio::task::spawn(async move {
+                loop {
+                    let endpoint = server
+                        .accept()
+                        .await
+                        .map_err(|e| error!("Failed to accept connection: {}", e));
+                    let endpoint = if let Ok(endpoint) = endpoint {
+                        endpoint
+                    } else {
+                        continue;
+                    };
+                    if let TunnelEndpoint::Client(client) = endpoint {
+                        if cancel_token.is_cancelled() {
+                            return;
+                        }
+                        let permit = semaphore.clone().acquire_owned().await.unwrap();
+                        let result = run(client, profile.clone()).await;
+                        match result {
+                            Ok(true) => {
+                                post_task(profile.clone(), permit, cancel_token.clone()).await;
+                            }
+                            Err(error) => {
+                                error!("{}", error);
+                            }
+                            _ => {}
+                        }
+                    }
+                }
+            }));
+        }
+
+        for client_conf in &CONFIG.connect.client {
+            let client_conf = client_conf.clone().into_tunnel_config(Identity::Client);
+            let semaphore = semaphore.clone();
+            let cancel_token = cancel_token.clone();
+            let profile = profile.clone();
+            info!("tcp client started for {}", client_conf.url);
+            tasks.push(tokio::task::spawn(async move {
+                loop {
+                    if cancel_token.is_cancelled() {
+                        return;
+                    }
+                    let endpoint = connect_tunnel(client_conf.clone())
+                        .await
+                        .map_err(|e| error!("Failed to accept connection: {}", e));
+                    let endpoint = if let Ok(endpoint) = endpoint {
+                        endpoint
+                    } else {
+                        continue;
+                    };
+                    if let TunnelEndpoint::Client(client) = endpoint {
+                        if cancel_token.is_cancelled() {
+                            return;
+                        }
+                        let permit = semaphore.clone().acquire_owned().await.unwrap();
+                        let result = run(client, profile.clone()).await;
+                        match result {
+                            Ok(true) => {
+                                post_task(profile.clone(), permit, cancel_token.clone()).await;
+                            }
+                            Err(error) => {
+                                error!("{}", error);
+                            }
+                            _ => {}
+                        }
+                    }
+                    sleep(Duration::from_secs(10)).await;
+                }
+            }));
+        }
+    }
+
+    for task in tasks {
+        let _ = task.await.map_err(|e| error!("{}", e));
+    }
+
+    Ok(())
+}
+lazy_static! {
+    pub static ref CONFIG: ClientConfig = {
+        let raw = fs::read_to_string("sekai-unpacker-client.yaml").unwrap();
+        let config: ClientConfig = yaml_serde::from_str(raw.as_str()).unwrap();
+        config
+    };
+}

client/src/task.rs

@@ -0,0 +1,146 @@
+use crate::config::Profile;
+use crate::http::{close, download, sync};
+use common::http::{CloseRequest, DownloadRequest};
+use communicator::ClientManager;
+use log::{error, info};
+use serde::{Deserialize, Serialize};
+use std::collections::HashMap;
+use std::path::{Path, PathBuf};
+use std::sync::Arc;
+use std::sync::atomic::{AtomicUsize, Ordering};
+use tokio::sync::{RwLock, Semaphore};
+
+pub async fn run(
+    client: Arc<ClientManager>,
+    profile: Arc<(String, Profile)>,
+) -> anyhow::Result<bool> {
+    info!("[{}]: Starting sync", profile.0);
+    let sync_resp = sync(&mut client.get_client().await?, &profile.1).await?;
+    let id = sync_resp.id;
+    let local_manifest = Arc::new(
+        AutoSaveManifest::new(
+            5,
+            Path::new(&profile.1.path)
+                .join("manifest.json")
+                .to_path_buf(),
+        )
+        .await?,
+    );
+    let manifest_snapshot = { local_manifest.manifest.read().await.clone() };
+    let tasks = sync_resp
+        .tasks
+        .into_iter()
+        .filter(|task| {
+            let bundle_name = &task.bundle_path;
+            match manifest_snapshot.bundles.get(bundle_name) {
+                Some(local_hash) => local_hash != &task.bundle_hash,
+                None => true,
+            }
+        })
+        .collect::<Vec<_>>();
+    info!("[{}]: Collected {} tasks", profile.0, tasks.len());
+    let n = 5;
+    let semaphore = Arc::new(Semaphore::new(n));
+    let mut handles = Vec::new();
+    for task in tasks {
+        let permit = semaphore.clone().acquire_owned().await?;
+        let client = client.clone();
+        let id = id.clone();
+        let local_manifest = local_manifest.clone();
+        let profile = profile.clone();
+
+        handles.push(tokio::task::spawn(async move {
+            let req = DownloadRequest {
+                id: id.clone(),
+                task: task.clone(),
+            };
+            let result = download(&mut client.get_client().await.unwrap(), &req, &profile.1).await;
+            if let Err(e) = result
+                && let Some(_) = e.downcast_ref::<h2::Error>()
+            {
+                download(&mut client.get_client().await.unwrap(), &req, &profile.1)
+                    .await
+                    .unwrap();
+            }
+
+            local_manifest
+                .add_bundle(task.bundle_path.clone(), task.bundle_hash.clone())
+                .await
+                .unwrap();
+            drop(permit);
+        }));
+    }
+    let mut succeed = 0;
+    let mut failed = 0;
+    for handle in handles {
+        let r = handle.await;
+        if let Err(e) = r {
+            error!("{}", e);
+            failed += 1;
+        } else {
+            succeed += 1;
+        }
+    }
+    local_manifest.save().await?;
+    info!(
+        "[{}]: Sync finished with {} succeed, {} failed",
+        profile.0, succeed, failed
+    );
+    let req = CloseRequest { id: id.clone() };
+    close(&mut client.get_client().await?, &req).await?;
+
+    Ok(failed == 0)
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+struct Manifest {
+    #[serde(default)]
+    bundles: HashMap<String, String>,
+}
+
+pub struct AutoSaveManifest {
+    manifest: Arc<RwLock<Manifest>>,
+    counter: AtomicUsize,
+    save_interval: usize,
+    storage_path: PathBuf,
+}
+
+impl AutoSaveManifest {
+    pub async fn new(interval: usize, path: PathBuf) -> anyhow::Result<Self> {
+        Ok(Self {
+            manifest: Arc::new(RwLock::new(serde_json::from_str(
+                &tokio::fs::read_to_string(&path)
+                    .await
+                    .unwrap_or("{}".to_owned()),
+            )?)),
+            counter: AtomicUsize::new(0),
+            save_interval: interval,
+            storage_path: path,
+        })
+    }
+
+    pub async fn add_bundle(&self, key: String, value: String) -> anyhow::Result<()> {
+        {
+            let mut w = self.manifest.write().await;
+            w.bundles.insert(key, value);
+        }
+
+        let current_count = self.counter.fetch_add(1, Ordering::SeqCst) + 1;
+
+        if current_count.is_multiple_of(self.save_interval) {
+            self.save().await?;
+        }
+
+        Ok(())
+    }
+
+    pub async fn save(&self) -> anyhow::Result<()> {
+        let data = {
+            let r = self.manifest.read().await;
+            serde_json::to_vec(&*r)?
+        };
+
+        tokio::fs::write(&self.storage_path, data).await?;
+        Ok(())
+    }
+}