refactoring the relay startup logic

common/src/enums.rs

@@ -4,4 +4,22 @@ use clap::ValueEnum;
 pub enum DataSource {
     Mysql,
     Postgres,
-}
+}
+
+#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, ValueEnum, Debug)]
+pub enum ZtmType {
+    Agent,
+    Relay,
+}
+
+impl std::str::FromStr for ZtmType {
+    type Err = String;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        match s.to_lowercase().as_str() {
+            "agent" => Ok(ZtmType::Agent),
+            "relay" => Ok(ZtmType::Relay),
+            _ => Err(format!("'{}' is not a valid ztm type", s)),
+        }
+    }
+}

common/src/model.rs

@@ -1,13 +1,19 @@
 use clap::Args;
 use serde::Deserialize;
 
+use crate::enums::ZtmType;
 
 #[derive(Args, Clone, Debug)]
 pub struct CommonOptions {
     #[arg(long, default_value_t = String::from("127.0.0.1"))]
     pub host: String,
-}
 
+    #[arg(long)]
+    pub ztm: Option<ZtmType>,
+
+    #[arg(long, default_value_t = 8001)]
+    pub relay_port: u16,
+}
 
 #[derive(Deserialize, Debug)]
 pub struct GetParams {

gateway/src/https_server.rs

@@ -13,6 +13,7 @@ use axum::routing::get;
 use axum::Router;
 use axum_server::tls_rustls::RustlsConfig;
 use clap::Args;
+use common::enums::ZtmType;
 use regex::Regex;
 use tower::ServiceBuilder;
 use tower_http::cors::{Any, CorsLayer};
@@ -27,6 +28,7 @@ use jupiter::context::Context;
 use jupiter::raw_storage::local_storage::LocalStorage;
 
 use crate::api_service::router::ApiServiceState;
+use crate::relay_server::run_relay_server;
 use crate::{api_service, lfs};
 
 #[derive(Args, Clone, Debug)]
@@ -80,11 +82,13 @@ pub fn remove_git_suffix(uri: Uri, git_suffix: &str) -> PathBuf {
 
 pub async fn https_server(config: Config, options: HttpsOptions) {
     let HttpsOptions {
-        common: CommonOptions { host },
+        common: CommonOptions { host, .. },
         https_key_path,
         https_cert_path,
         https_port,
-    } = options;
+    } = options.clone();
+
+    check_run_with_ztm(config.clone(), options.common);
 
     let app = app(config, host.clone(), https_port).await;
 
@@ -101,9 +105,11 @@ pub async fn https_server(config: Config, options: HttpsOptions) {
 
 pub async fn http_server(config: Config, options: HttpOptions) {
     let HttpOptions {
-        common: CommonOptions { host },
+        common: CommonOptions { host, .. },
         http_port,
-    } = options;
+    } = options.clone();
+
+    check_run_with_ztm(config.clone(), options.common);
 
     let app = app(config, host.clone(), http_port).await;
 
@@ -239,5 +245,28 @@ async fn put_method_router(
     }
 }
 
+pub fn check_run_with_ztm(config: Config, common: CommonOptions) {
+    let ztm_type = match common.ztm {
+        Some(z) => z,
+        None => {
+            return;
+        }
+    };
+    match ztm_type {
+        ZtmType::Agent => {
+            //TODO Mega server join a ztm mesh
+        }
+        ZtmType::Relay => {
+            //Start a sub thread to run relay server
+            let config_clone = config.clone();
+            let host_clone = common.host.clone();
+            let relay_port = common.relay_port;
+            tokio::spawn(
+                async move { run_relay_server(config_clone, host_clone, relay_port).await },
+            );
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {}

gateway/src/relay_server.rs

@@ -6,9 +6,10 @@ use axum::extract::{Query, State};
 use axum::http::{Request, Response, StatusCode, Uri};
 use axum::routing::get;
 use axum::Router;
-use clap::Args;
-use common::config::Config;
-use common::model::CommonOptions;
+use common::config::{Config, ZTMConfig};
+use gemini::ztm::{
+    connect_ztm_hub, create_ztm_certificate, create_ztm_service, delete_ztm_certificate,
+};
 use gemini::RelayGetParams;
 use jupiter::context::Context;
 use regex::Regex;
@@ -20,25 +21,22 @@ use tower_http::trace::TraceLayer;
 use crate::api_service;
 use crate::api_service::router::ApiServiceState;
 
-#[derive(Args, Clone, Debug)]
-pub struct RelayOptions {
-    #[clap(flatten)]
-    pub common: CommonOptions,
+pub async fn run_relay_server(config: Config, host: String, port: u16) {
+    let app = app(config.clone(), host.clone(), port).await;
 
-    #[arg(long, default_value_t = 8001)]
-    pub http_port: u16,
-}
-
-pub async fn http_server(config: Config, options: RelayOptions) {
-    let RelayOptions {
-        common: CommonOptions { host },
-        http_port,
-    } = options;
-
-    let app = app(config, host.clone(), http_port).await;
-
-    let server_url = format!("{}:{}", host, http_port);
+    let ztm_config = config.ztm;
+    match relay_connect_ztm(ztm_config, port).await {
+        Ok(s) => {
+            tracing::info!("relay connect ztm success: {s}");
+        }
+        Err(e) => {
+            tracing::info!("relay connect ztm failed : {e}");
+            return;
+        }
+    }
 
+    let server_url = format!("{}:{}", host, port);
+    tracing::info!("start relay server: {server_url}");
     let addr = SocketAddr::from_str(&server_url).unwrap();
     let listener = tokio::net::TcpListener::bind(addr).await.unwrap();
     axum::serve(listener, app.into_make_service())
@@ -90,9 +88,9 @@ async fn get_method_router(
 ) -> Result<Response<Body>, (StatusCode, String)> {
     let ztm_config = state.context.config.ztm.clone();
     if Regex::new(r"/hello$").unwrap().is_match(uri.path()) {
-        return gemini::http::handler::hello_gemini(params).await;
+        return hello_relay(params).await;
     } else if Regex::new(r"/certificate$").unwrap().is_match(uri.path()) {
-        return gemini::http::handler::certificate(ztm_config, params).await;
+        return certificate(ztm_config, params).await;
     }
     Err((
         StatusCode::NOT_FOUND,
@@ -102,18 +100,79 @@ async fn get_method_router(
 
 async fn post_method_router(
     state: State<AppState>,
-    uri: Uri,
-    req: Request<Body>,
+    _uri: Uri,
+    _req: Request<Body>,
 ) -> Result<Response<Body>, (StatusCode, String)> {
-    let ztm_config = state.context.config.ztm.clone();
-    if Regex::new(r"/relay_init$").unwrap().is_match(uri.path()) {
-        return gemini::http::handler::init(ztm_config, req, state.0.port).await;
-    }
+    let _ztm_config = state.context.config.ztm.clone();
     Err((
         StatusCode::NOT_FOUND,
         String::from("Operation not supported\n"),
     ))
 }
 
+pub async fn hello_relay(_params: RelayGetParams) -> Result<Response<Body>, (StatusCode, String)> {
+    Ok(Response::builder().body(Body::from("hello relay")).unwrap())
+}
+
+pub async fn certificate(
+    config: ZTMConfig,
+    params: RelayGetParams,
+) -> Result<Response<Body>, (StatusCode, String)> {
+    if params.name.is_none() {
+        return Err((StatusCode::BAD_REQUEST, "not enough paras".to_string()));
+    }
+    let name = params.name.unwrap();
+    let permit = match create_ztm_certificate(config, name.clone()).await {
+        Ok(p) => p,
+        Err(e) => {
+            return Err((StatusCode::INTERNAL_SERVER_ERROR, e));
+        }
+    };
+
+    let permit_json = serde_json::to_string(&permit).unwrap();
+    tracing::info!("new permit [{name}]: {permit_json}");
+
+    Ok(Response::builder()
+        .header("Content-Type", "application/json")
+        .body(Body::from(permit_json))
+        .unwrap())
+}
+
+pub async fn relay_connect_ztm(config: ZTMConfig, relay_port: u16) -> Result<String, String> {
+    // 1. generate a permit for relay
+    let name = "relay".to_string();
+    match delete_ztm_certificate(config.clone(), name.clone()).await {
+        Ok(_s) => (),
+        Err(e) => {
+            return Err(e);
+        }
+    }
+    let permit = match create_ztm_certificate(config.clone(), name).await {
+        Ok(p) => p,
+        Err(e) => {
+            return Err(e);
+        }
+    };
+
+    // 2. connect to ZTM hub (join a mesh)
+    let mesh = match connect_ztm_hub(config.clone(), permit).await {
+        Ok(m) => m,
+        Err(e) => {
+            return Err(e);
+        }
+    };
+
+    // 3. create a ZTM service
+    let response_text =
+        match create_ztm_service(config, mesh.agent.id, "relay".to_string(), relay_port).await {
+            Ok(m) => m,
+            Err(e) => {
+                return Err(e);
+            }
+        };
+
+    Ok(response_text)
+}
+
 #[cfg(test)]
 mod tests {}

gateway/src/ssh_server.rs

@@ -62,7 +62,7 @@ pub async fn start_server(config: Config, command: &SshOptions) {
     let ru_config = Arc::new(ru_config);
 
     let SshOptions {
-        common: CommonOptions { host },
+        common: CommonOptions { host, .. },
         custom:
             SshCustom {
                 ssh_port,

gemini/src/http/handler.rs

@@ -1,88 +1,2 @@
-use axum::{
-    body::{to_bytes, Body},
-    http::{Request, Response, StatusCode},
-};
-use common::config::ZTMConfig;
-
-use crate::{
-    ztm::{
-        ZTMUserPermit, {connect_ztm_hub, create_ztm_certificate, create_ztm_service},
-    },
-    RelayGetParams,
-};
-
-pub async fn hello_gemini(_params: RelayGetParams) -> Result<Response<Body>, (StatusCode, String)> {
-    Ok(Response::builder()
-        .body(Body::from("hello gemini"))
-        .unwrap())
-}
-
-pub async fn certificate(
-    config: ZTMConfig,
-    params: RelayGetParams,
-) -> Result<Response<Body>, (StatusCode, String)> {
-    if params.name.is_none() {
-        return Err((StatusCode::BAD_REQUEST, "not enough paras".to_string()));
-    }
-    let name = params.name.unwrap();
-    let permit = match create_ztm_certificate(config, name.clone()).await {
-        Ok(p) => p,
-        Err(e) => {
-            return Err((StatusCode::INTERNAL_SERVER_ERROR, e));
-        }
-    };
-
-    let permit_json = serde_json::to_string(&permit).unwrap();
-    tracing::info!("new permit [{name}]: {permit_json}");
-
-    Ok(Response::builder()
-        .header("Content-Type", "application/json")
-        .body(Body::from(permit_json))
-        .unwrap())
-}
-
-pub async fn init(
-    config: ZTMConfig,
-    req: Request<Body>,
-    relay_port: u16,
-) -> Result<Response<Body>, (StatusCode, String)> {
-    // transfer body to json
-    let body_bytes = match to_bytes(req.into_body(), usize::MAX).await {
-        Ok(b) => b,
-        Err(e) => {
-            return Err((StatusCode::BAD_REQUEST, e.to_string()));
-        }
-    };
-
-    let permit: ZTMUserPermit = match serde_json::from_slice(&body_bytes) {
-        Ok(p) => p,
-        Err(e) => {
-            return Err((StatusCode::BAD_REQUEST, e.to_string()));
-        }
-    };
-
-    // 1. connect to ZTM hub (join a mesh)
-    let mesh = match connect_ztm_hub(config.clone(), permit).await {
-        Ok(m) => m,
-        Err(e) => {
-            return Err((StatusCode::INTERNAL_SERVER_ERROR, e));
-        }
-    };
-
-    // 2. create a ZTM service
-    let response_text =
-        match create_ztm_service(config, mesh.agent.id, "relay".to_string(), relay_port).await {
-            Ok(m) => m,
-            Err(e) => {
-                return Err((StatusCode::INTERNAL_SERVER_ERROR, e));
-            }
-        };
-
-    Ok(Response::builder()
-        .header("Content-Type", "application/json")
-        .body(Body::from(response_text))
-        .unwrap())
-}
-
 #[cfg(test)]
 mod tests {}