mirror of
https://github.com/house-of-vanity/yggman.git
synced 2025-12-18 03:57:54 +00:00
first commit
This commit is contained in:
Ultradesu
40
src/modules/example.rs
Normal file
40
src/modules/example.rs
Normal file
@@ -0,0 +1,40 @@
|
||||
use async_trait::async_trait;
|
||||
use std::sync::Arc;
|
||||
use crate::core::context::AppContext;
|
||||
use crate::core::module::Module;
|
||||
use crate::error::Result;
|
||||
|
||||
pub struct ExampleModule {
|
||||
name: String,
|
||||
context: Option<Arc<AppContext>>,
|
||||
}
|
||||
|
||||
|
||||
#[async_trait]
|
||||
impl Module for ExampleModule {
|
||||
fn name(&self) -> &str {
|
||||
&self.name
|
||||
}
|
||||
|
||||
async fn init(&mut self, context: Arc<AppContext>) -> Result<()> {
|
||||
self.context = Some(context);
|
||||
tracing::debug!("Example module initialized");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn start(&self) -> Result<()> {
|
||||
tracing::debug!("Example module started");
|
||||
|
||||
if let Some(ctx) = &self.context {
|
||||
let config = ctx.config_manager.get();
|
||||
tracing::debug!("Current config: {:?}", config);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn stop(&self) -> Result<()> {
|
||||
tracing::debug!("Example module stopped");
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
3
src/modules/mod.rs
Normal file
3
src/modules/mod.rs
Normal file
@@ -0,0 +1,3 @@
|
||||
pub mod example;
|
||||
pub mod web;
|
||||
pub mod websocket;
|
||||
273
src/modules/web.rs
Normal file
273
src/modules/web.rs
Normal file
@@ -0,0 +1,273 @@
|
||||
use async_trait::async_trait;
|
||||
use axum::{
|
||||
extract::{State, Path, WebSocketUpgrade},
|
||||
http::StatusCode,
|
||||
response::{Html, Json, Response},
|
||||
routing::{get, post, put, delete},
|
||||
Router,
|
||||
};
|
||||
use std::sync::Arc;
|
||||
use tower_http::cors::CorsLayer;
|
||||
use sea_orm::DatabaseConnection;
|
||||
|
||||
use crate::core::context::AppContext;
|
||||
use crate::core::module::Module;
|
||||
use crate::error::Result;
|
||||
use crate::node_manager::NodeManager;
|
||||
use crate::yggdrasil::{Node, YggdrasilConfig};
|
||||
|
||||
pub struct WebModule {
|
||||
name: String,
|
||||
context: Option<Arc<AppContext>>,
|
||||
node_manager: Arc<NodeManager>,
|
||||
}
|
||||
|
||||
impl WebModule {
|
||||
pub fn new(db: DatabaseConnection) -> Self {
|
||||
Self {
|
||||
name: "web".to_string(),
|
||||
context: None,
|
||||
node_manager: Arc::new(NodeManager::new(db)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Module for WebModule {
|
||||
fn name(&self) -> &str {
|
||||
&self.name
|
||||
}
|
||||
|
||||
async fn init(&mut self, context: Arc<AppContext>) -> Result<()> {
|
||||
self.context = Some(context);
|
||||
tracing::info!("Web module initialized");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn start(&self) -> Result<()> {
|
||||
let context = self.context.as_ref().unwrap();
|
||||
let config = context.config_manager.get();
|
||||
let port = config.server.port;
|
||||
|
||||
tracing::info!("Starting web server on port {}", port);
|
||||
|
||||
let node_manager = self.node_manager.clone();
|
||||
|
||||
let app = Router::new()
|
||||
.route("/", get(index_handler))
|
||||
.route("/api/nodes", get(get_nodes_handler))
|
||||
.route("/api/nodes", post(add_node_handler))
|
||||
.route("/api/nodes/:id", get(get_node_handler))
|
||||
.route("/api/nodes/:id", put(update_node_handler))
|
||||
.route("/api/nodes/:id", delete(delete_node_handler))
|
||||
.route("/api/configs", get(get_configs_handler))
|
||||
.route("/api/nodes/:id/config", get(get_node_config_handler))
|
||||
.route("/ws/agent", get(ws_agent_handler))
|
||||
.layer(CorsLayer::permissive())
|
||||
.with_state(node_manager);
|
||||
|
||||
let bind_addr = format!("{}:{}", config.server.bind_address, port);
|
||||
let listener = tokio::net::TcpListener::bind(&bind_addr)
|
||||
.await
|
||||
.map_err(|e| crate::error::AppError::Io(e))?;
|
||||
|
||||
tokio::spawn(async move {
|
||||
axum::serve(listener, app)
|
||||
.await
|
||||
.expect("Failed to run web server");
|
||||
});
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn stop(&self) -> Result<()> {
|
||||
tracing::info!("Web module stopped");
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
async fn index_handler() -> Html<&'static str> {
|
||||
Html(include_str!("../../static/index.html"))
|
||||
}
|
||||
|
||||
#[derive(serde::Serialize)]
|
||||
struct NodesResponse {
|
||||
nodes: Vec<Node>,
|
||||
}
|
||||
|
||||
async fn get_nodes_handler(
|
||||
State(node_manager): State<Arc<NodeManager>>,
|
||||
) -> Json<NodesResponse> {
|
||||
let nodes = node_manager.get_all_nodes().await;
|
||||
Json(NodesResponse { nodes })
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
struct AddNodeRequest {
|
||||
name: String,
|
||||
listen: Vec<String>,
|
||||
addresses: Vec<String>,
|
||||
}
|
||||
|
||||
#[derive(serde::Serialize)]
|
||||
struct AddNodeResponse {
|
||||
success: bool,
|
||||
message: String,
|
||||
}
|
||||
|
||||
async fn add_node_handler(
|
||||
State(node_manager): State<Arc<NodeManager>>,
|
||||
Json(payload): Json<AddNodeRequest>,
|
||||
) -> Json<AddNodeResponse> {
|
||||
match node_manager.add_node(payload.name, payload.listen, payload.addresses).await {
|
||||
Ok(_) => {
|
||||
// Broadcast update to all connected agents
|
||||
crate::websocket_state::broadcast_configuration_update(&node_manager).await;
|
||||
|
||||
Json(AddNodeResponse {
|
||||
success: true,
|
||||
message: "Node added successfully".to_string(),
|
||||
})
|
||||
}
|
||||
Err(e) => Json(AddNodeResponse {
|
||||
success: false,
|
||||
message: format!("Failed to add node: {}", e),
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(serde::Serialize)]
|
||||
struct ConfigsResponse {
|
||||
configs: Vec<NodeConfig>,
|
||||
}
|
||||
|
||||
#[derive(serde::Serialize)]
|
||||
struct NodeConfig {
|
||||
node_id: String,
|
||||
node_name: String,
|
||||
node_addresses: Vec<String>,
|
||||
config: YggdrasilConfig,
|
||||
}
|
||||
|
||||
async fn get_configs_handler(
|
||||
State(node_manager): State<Arc<NodeManager>>,
|
||||
) -> Json<ConfigsResponse> {
|
||||
let nodes = node_manager.get_all_nodes().await;
|
||||
let configs_map = node_manager.generate_configs().await;
|
||||
|
||||
let mut configs = Vec::new();
|
||||
for node in nodes {
|
||||
if let Some(config) = configs_map.get(&node.id) {
|
||||
configs.push(NodeConfig {
|
||||
node_id: node.id.clone(),
|
||||
node_name: node.name.clone(),
|
||||
node_addresses: node.addresses.clone(),
|
||||
config: config.clone(),
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
Json(ConfigsResponse { configs })
|
||||
}
|
||||
|
||||
// Get single node handler
|
||||
async fn get_node_handler(
|
||||
State(node_manager): State<Arc<NodeManager>>,
|
||||
Path(node_id): Path<String>,
|
||||
) -> std::result::Result<Json<Node>, StatusCode> {
|
||||
match node_manager.get_node_by_id(&node_id).await {
|
||||
Some(node) => Ok(Json(node)),
|
||||
None => Err(StatusCode::NOT_FOUND),
|
||||
}
|
||||
}
|
||||
|
||||
// Update node handler
|
||||
async fn update_node_handler(
|
||||
State(node_manager): State<Arc<NodeManager>>,
|
||||
Path(node_id): Path<String>,
|
||||
Json(payload): Json<AddNodeRequest>,
|
||||
) -> std::result::Result<Json<AddNodeResponse>, StatusCode> {
|
||||
match node_manager.update_node(&node_id, payload.name, payload.listen, payload.addresses).await {
|
||||
Ok(_) => {
|
||||
// Broadcast update to all connected agents
|
||||
crate::websocket_state::broadcast_configuration_update(&node_manager).await;
|
||||
|
||||
Ok(Json(AddNodeResponse {
|
||||
success: true,
|
||||
message: "Node updated successfully".to_string(),
|
||||
}))
|
||||
}
|
||||
Err(e) => {
|
||||
if e.to_string().contains("Node not found") {
|
||||
Err(StatusCode::NOT_FOUND)
|
||||
} else {
|
||||
Ok(Json(AddNodeResponse {
|
||||
success: false,
|
||||
message: format!("Failed to update node: {}", e),
|
||||
}))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Delete node handler
|
||||
async fn delete_node_handler(
|
||||
State(node_manager): State<Arc<NodeManager>>,
|
||||
Path(node_id): Path<String>,
|
||||
) -> std::result::Result<Json<AddNodeResponse>, StatusCode> {
|
||||
match node_manager.remove_node(&node_id).await {
|
||||
Ok(_) => {
|
||||
// Broadcast update to all connected agents
|
||||
crate::websocket_state::broadcast_configuration_update(&node_manager).await;
|
||||
|
||||
Ok(Json(AddNodeResponse {
|
||||
success: true,
|
||||
message: "Node deleted successfully".to_string(),
|
||||
}))
|
||||
}
|
||||
Err(e) => {
|
||||
if e.to_string().contains("Node not found") {
|
||||
Err(StatusCode::NOT_FOUND)
|
||||
} else {
|
||||
Ok(Json(AddNodeResponse {
|
||||
success: false,
|
||||
message: format!("Failed to delete node: {}", e),
|
||||
}))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Get node configuration for agent
|
||||
async fn get_node_config_handler(
|
||||
State(node_manager): State<Arc<NodeManager>>,
|
||||
Path(node_id): Path<String>,
|
||||
) -> std::result::Result<Json<NodeConfig>, StatusCode> {
|
||||
// Get the node
|
||||
let node = match node_manager.get_node_by_id(&node_id).await {
|
||||
Some(node) => node,
|
||||
None => return Err(StatusCode::NOT_FOUND),
|
||||
};
|
||||
|
||||
// Generate configurations for all nodes
|
||||
let configs_map = node_manager.generate_configs().await;
|
||||
|
||||
// Get config for this specific node
|
||||
match configs_map.get(&node_id) {
|
||||
Some(config) => Ok(Json(NodeConfig {
|
||||
node_id: node.id.clone(),
|
||||
node_name: node.name.clone(),
|
||||
node_addresses: node.addresses.clone(),
|
||||
config: config.clone(),
|
||||
})),
|
||||
None => Err(StatusCode::INTERNAL_SERVER_ERROR),
|
||||
}
|
||||
}
|
||||
|
||||
// WebSocket handler for agents
|
||||
async fn ws_agent_handler(
|
||||
ws: WebSocketUpgrade,
|
||||
State(node_manager): State<Arc<NodeManager>>,
|
||||
) -> Response {
|
||||
ws.on_upgrade(move |socket| crate::modules::websocket::handle_agent_socket(socket, node_manager))
|
||||
}
|
||||
164
src/modules/websocket.rs
Normal file
164
src/modules/websocket.rs
Normal file
@@ -0,0 +1,164 @@
|
||||
use axum::extract::ws::{Message, WebSocket};
|
||||
use futures_util::{SinkExt, StreamExt};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::sync::Arc;
|
||||
use tracing::{debug, error, info, warn};
|
||||
|
||||
use crate::node_manager::NodeManager;
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
#[serde(tag = "type")]
|
||||
pub enum AgentMessage {
|
||||
Register {
|
||||
name: String,
|
||||
addresses: Vec<String>,
|
||||
},
|
||||
Heartbeat,
|
||||
UpdateAddresses {
|
||||
addresses: Vec<String>,
|
||||
},
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
#[serde(tag = "type")]
|
||||
pub enum ServerMessage {
|
||||
Config {
|
||||
node_id: String,
|
||||
private_key: String,
|
||||
listen: Vec<String>,
|
||||
peers: Vec<String>,
|
||||
allowed_public_keys: Vec<String>,
|
||||
},
|
||||
Update {
|
||||
peers: Vec<String>,
|
||||
allowed_public_keys: Vec<String>,
|
||||
},
|
||||
Error {
|
||||
message: String,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
pub async fn handle_agent_socket(
|
||||
socket: WebSocket,
|
||||
node_manager: Arc<NodeManager>,
|
||||
) {
|
||||
let (mut sender, mut receiver) = socket.split();
|
||||
let (tx, mut rx) = tokio::sync::mpsc::channel::<ServerMessage>(100);
|
||||
|
||||
let mut node_id: Option<String> = None;
|
||||
|
||||
// Spawn task to forward messages from channel to WebSocket
|
||||
let send_task = tokio::spawn(async move {
|
||||
while let Some(msg) = rx.recv().await {
|
||||
if let Ok(json) = serde_json::to_string(&msg) {
|
||||
if sender.send(Message::Text(json)).await.is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
// Handle incoming messages
|
||||
while let Some(msg) = receiver.next().await {
|
||||
if let Ok(Message::Text(text)) = msg {
|
||||
match serde_json::from_str::<AgentMessage>(&text) {
|
||||
Ok(agent_msg) => {
|
||||
match agent_msg {
|
||||
AgentMessage::Register { name, addresses } => {
|
||||
info!("Agent registration: {} with addresses {:?}", name, addresses);
|
||||
|
||||
// Get default endpoints from config
|
||||
// TODO: Get from actual config, for now use hardcoded
|
||||
let default_listen = vec!["tcp://0.0.0.0:9001".to_string()];
|
||||
|
||||
// Check if node already exists
|
||||
let node = if let Some(existing_node) = node_manager.get_node_by_name(&name).await {
|
||||
info!("Reusing existing node: {} ({})", existing_node.name, existing_node.id);
|
||||
// Update addresses for existing node
|
||||
match node_manager.update_node(&existing_node.id, name.clone(), default_listen.clone(), addresses).await {
|
||||
Ok(_) => {
|
||||
// Get the updated node
|
||||
node_manager.get_node_by_id(&existing_node.id).await
|
||||
}
|
||||
Err(e) => {
|
||||
warn!("Failed to update existing node addresses: {}", e);
|
||||
Some(existing_node)
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// Create new node
|
||||
info!("Creating new node: {}", name);
|
||||
match node_manager.add_node(name.clone(), default_listen.clone(), addresses).await {
|
||||
Ok(_) => {
|
||||
// Get the newly created node
|
||||
node_manager.get_node_by_name(&name).await
|
||||
}
|
||||
Err(e) => {
|
||||
let error_msg = ServerMessage::Error {
|
||||
message: format!("Failed to register node: {}", e),
|
||||
};
|
||||
let _ = tx.send(error_msg).await;
|
||||
None
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
if let Some(node) = node {
|
||||
node_id = Some(node.id.clone());
|
||||
|
||||
// Register connection
|
||||
crate::websocket_state::register_agent_connection(node.id.clone(), tx.clone()).await;
|
||||
|
||||
// Generate config for this node
|
||||
let configs = node_manager.generate_configs().await;
|
||||
if let Some(config) = configs.get(&node.id) {
|
||||
let peers: Vec<String> = config.peers.clone();
|
||||
let allowed_keys: Vec<String> = config.allowed_public_keys.clone();
|
||||
|
||||
let response = ServerMessage::Config {
|
||||
node_id: node.id.clone(),
|
||||
private_key: node.private_key.clone(),
|
||||
listen: default_listen,
|
||||
peers,
|
||||
allowed_public_keys: allowed_keys,
|
||||
};
|
||||
|
||||
if let Err(e) = tx.send(response).await {
|
||||
error!("Failed to send config to agent: {}", e);
|
||||
}
|
||||
|
||||
// Notify other agents about node connection
|
||||
crate::websocket_state::broadcast_configuration_update(&node_manager).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
AgentMessage::Heartbeat => {
|
||||
debug!("Heartbeat from {:?}", node_id);
|
||||
}
|
||||
AgentMessage::UpdateAddresses { addresses } => {
|
||||
if let Some(id) = &node_id {
|
||||
info!("Address update for {}: {:?}", id, addresses);
|
||||
// TODO: Update node addresses in database
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
warn!("Failed to parse agent message: {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Clean up
|
||||
if let Some(id) = node_id {
|
||||
crate::websocket_state::unregister_agent_connection(&id).await;
|
||||
info!("Agent {} disconnected", id);
|
||||
}
|
||||
|
||||
// Abort send task
|
||||
send_task.abort();
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user