From f5eff9e0f9035bc27addea914a7afad0a2190fe6 Mon Sep 17 00:00:00 2001
From: Kira Sotnikov <s8tn1kv@gmail.com>
Date: Wed, 8 Jan 2025 00:54:49 +0300
Subject: [PATCH] Fix mode reg and install

---
 .github/workflows/release.yml |  2 ++
 .gitignore                    |  1 +
 deploy/install                | 30 ++++++++++++++---------------
 src/bin/agent.rs              |  7 +++++--
 src/jobs/agent.rs             | 36 ++++++++++++++++++++++++-----------
 5 files changed, 48 insertions(+), 28 deletions(-)

diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml
index c06f5d8..b8d6329 100644
--- a/.github/workflows/release.yml
+++ b/.github/workflows/release.yml
@@ -86,6 +86,8 @@ jobs:
           files: |
             target/release/api-${{ env.ARCH }}
             target/release/agent-${{ env.ARCH }}
+            deploy/install
+            docs/
             README.md
             config-agent-example.toml
             config-api-example.toml
diff --git a/.gitignore b/.gitignore
index 48515f6..81cbd87 100644
--- a/.gitignore
+++ b/.gitignore
@@ -15,3 +15,4 @@ users.json
 xray-config-mk2nl.json
 connection-lines.txt
 dev/data
+deploy/env
diff --git a/deploy/install b/deploy/install
index 103f51f..6ba1114 100755
--- a/deploy/install
+++ b/deploy/install
@@ -50,18 +50,18 @@ PG_USERNAME="${PG_USERNAME:-postgres}"
 PG_PASSWORD="${PG_PASSWORD:-password}"
 API_ENDPOINT=${API_ENDPOINT:-http://localhost:3005}
 
-sudo mkdir -p "$INSTALL_DIR"
+mkdir -p "$INSTALL_DIR"
 cd "$INSTALL_DIR"
 
 echo "Installing xray-core version $XRAY_VERSION..."
 curl -L -o xray.zip "$XRAY_URL"
-sudo unzip -o xray.zip -d "$INSTALL_DIR"
-sudo rm xray.zip
+unzip -o xray.zip -d "$INSTALL_DIR"
+rm xray.zip
 
 echo "XRAY Core Version"
 $INSTALL_DIR/xray --version
 
-cat <<EOF | sudo tee /etc/systemd/system/xray.service
+cat <<EOF | tee /etc/systemd/system/xray.service
 [Unit]
 Description=Xray Core Service
 After=network.target
@@ -74,14 +74,14 @@ Restart=on-failure
 WantedBy=multi-user.target
 EOF
 
-sudo systemctl daemon-reload
-sudo systemctl enable xray
+systemctl daemon-reload
+systemctl enable xray
 
 echo "Installing agent version $PONY_VERSION..."
 curl -L -o agent "$AGENT_URL"
-sudo chmod +x agent
+chmod +x agent
 
-cat <<EOF | sudo tee /etc/systemd/system/agent.service
+cat <<EOF | tee /etc/systemd/system/agent.service
 [Unit]
 Description=Pony Agent Service
 After=xray.service
@@ -97,8 +97,8 @@ RestartSec=5
 WantedBy=multi-user.target
 EOF
 
-sudo systemctl daemon-reload
-sudo systemctl enable agent
+systemctl daemon-reload
+systemctl enable agent
 
 echo "Generating configurations..."
 
@@ -124,7 +124,7 @@ fi
 if [[ -f "$XRAY_CONFIG_PATH" ]]; then
   echo "File $XRAY_CONFIG_PATH already exist. skip."
 else
-    cat <<EOF | sudo tee "$XRAY_CONFIG_PATH"
+    cat <<EOF | tee "$XRAY_CONFIG_PATH"
 {
   "log": {
     "loglevel": "debug"
@@ -314,7 +314,7 @@ fi
 if [[ -f "$PONY_CONFIG_PATH" ]]; then
   echo "File $PONY_CONFIG_PATH already exist. skip."
 else
-    cat <<EOF | sudo tee "$PONY_CONFIG_PATH"
+    cat <<EOF | tee "$PONY_CONFIG_PATH"
 [debug]
 enabled = true
 web_server = "127.0.0.1"
@@ -344,17 +344,17 @@ sub_endpoint = "$ZMQ_ENDPOINT"
 
 [node]
 env = "$ENV"
-uuid = $(./xray uuid)
+uuid = "$(./xray uuid)"
 hostname = "$HOSTNAME"
 default_interface = "$default_interface"
 ipv4 = "$default_ip"
 
 [pg]
 host = "${PG_HOST}"
-port = "${PG_PORT}"
+port = ${PG_PORT}
 db = "${PG_DB}"
 username = "${PG_USERNAME}"
-password = "${PG_PASSWORD}
+password = "${PG_PASSWORD}"
 
 [api]
 endpoint = "$API_ENDPOINT"
diff --git a/src/bin/agent.rs b/src/bin/agent.rs
index 18f5804..971e3cf 100644
--- a/src/bin/agent.rs
+++ b/src/bin/agent.rs
@@ -173,8 +173,11 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     let _ = {
         let settings = settings.clone();
         debug!("----->>>>> Register node");
-        let _ =
-            agent::register_node(state.clone(), settings.clone(), settings.node.env.clone()).await;
+        if let Err(e) =
+            agent::register_node(state.clone(), settings.clone(), settings.node.env.clone()).await
+        {
+            panic!("Cannot register node {:?}", e);
+        }
     };
 
     if debug {
diff --git a/src/jobs/agent.rs b/src/jobs/agent.rs
index b9cc163..33142df 100644
--- a/src/jobs/agent.rs
+++ b/src/jobs/agent.rs
@@ -1,6 +1,7 @@
 use chrono::{Duration, Utc};
 use log::{debug, error};
-use reqwest::Client;
+use reqwest::Url;
+use reqwest::{Client, StatusCode};
 use std::{error::Error, sync::Arc};
 use tokio::sync::Mutex;
 use uuid::Uuid;
@@ -168,28 +169,41 @@ pub async fn register_node(
     let node_state = state.lock().await;
     let node = node_state.nodes.get(&env).clone();
 
-    let client = Client::new();
+    let mut endpoint = Url::parse(&settings.api.endpoint)?;
+    endpoint
+        .path_segments_mut()
+        .map_err(|_| "Invalid API endpoint")?
+        .push("node")
+        .push("register");
+    let endpoint = endpoint.to_string();
 
-    let endpoint = format!("{}/node/register", settings.api.endpoint);
     debug!("ENDPOINT: {}", endpoint);
 
     match serde_json::to_string_pretty(&node) {
-        Ok(json) => debug!("NODE {}", json),
-        Err(e) => error!("Error serializing to JSON: {}", e),
+        Ok(json) => debug!("Serialized node for environment '{}': {}", env, json),
+        Err(e) => error!("Error serializing node for environment '{}': {}", env, e),
     }
 
     if let Some(env) = node {
         if let Some(node) = env.first() {
-            let res = client.post(endpoint).json(&node).send().await?;
-            if res.status().is_success() {
-                debug!("Req success!");
+            let res = Client::new()
+                .post(&endpoint)
+                .header("Content-Type", "application/json")
+                .json(&node)
+                .send()
+                .await?;
+
+            if res.status().is_success() || res.status() == StatusCode::NOT_MODIFIED {
+                return Ok(());
             } else {
-                error!("Req error: {} {:?}", res.status(), res);
+                return Err(format!("Req error: {} {:?}", res.status(), res).into());
             }
+        } else {
+            return Err("No nodes found in environment".into());
         }
+    } else {
+        return Err("No data found for the given environment".into());
     }
-
-    Ok(())
 }
 
 pub async fn init_state(