105 lines
2.4 KiB
Rust
105 lines
2.4 KiB
Rust
pub mod api {
|
|
pub mod auth;
|
|
pub mod catchers;
|
|
pub mod mcp;
|
|
}
|
|
|
|
pub mod core {
|
|
pub mod config;
|
|
pub mod db;
|
|
pub mod validation;
|
|
}
|
|
|
|
pub mod transport {
|
|
pub mod execute;
|
|
pub mod protocol;
|
|
pub mod socket;
|
|
}
|
|
|
|
use core::config::Config;
|
|
use dashmap::DashMap;
|
|
use log::info;
|
|
use rocket::routes;
|
|
use serde::Serialize;
|
|
use sqlx::SqlitePool;
|
|
use sqlx::sqlite::SqliteConnectOptions;
|
|
use std::sync::{Arc, atomic::AtomicUsize};
|
|
use tokio::sync::{mpsc, oneshot};
|
|
|
|
const DATABASE_PATH: &str = "server.db";
|
|
|
|
pub type Database = SqlitePool;
|
|
pub type DeviceCounts = DashMap<String, Arc<AtomicUsize>>;
|
|
pub type Registry = DashMap<(String, String), RegistryEntry>;
|
|
|
|
pub struct RegistryEntry {
|
|
pub sender: mpsc::Sender<PendingExec>,
|
|
pub in_flight: Arc<AtomicUsize>,
|
|
}
|
|
|
|
#[derive(Debug, Serialize)]
|
|
pub struct ExecResult {
|
|
pub exit_code: i32,
|
|
pub stdout: String,
|
|
pub stderr: String,
|
|
}
|
|
|
|
pub struct PendingExec {
|
|
pub exec_id: String,
|
|
pub command: String,
|
|
pub reply: oneshot::Sender<ExecResult>,
|
|
}
|
|
|
|
pub struct AppState {
|
|
pub database: Database,
|
|
pub registry: Arc<Registry>,
|
|
pub device_counts: Arc<DeviceCounts>,
|
|
pub config: Arc<Config>,
|
|
}
|
|
|
|
#[rocket::main]
|
|
async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
|
env_logger::init();
|
|
|
|
let database = connect_database().await?;
|
|
apply_schema(&database).await?;
|
|
|
|
info!("Launching server");
|
|
|
|
build_rocket(database).launch().await?;
|
|
|
|
Ok(())
|
|
}
|
|
|
|
async fn connect_database() -> Result<Database, sqlx::Error> {
|
|
let options = SqliteConnectOptions::new()
|
|
.filename(DATABASE_PATH)
|
|
.create_if_missing(true);
|
|
|
|
SqlitePool::connect_with(options).await
|
|
}
|
|
|
|
async fn apply_schema(database: &Database) -> Result<(), sqlx::Error> {
|
|
for statement in include_str!("../schema.sql")
|
|
.split(';')
|
|
.map(str::trim)
|
|
.filter(|statement| !statement.is_empty())
|
|
{
|
|
sqlx::query(statement).execute(database).await?;
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
fn build_rocket(database: Database) -> rocket::Rocket<rocket::Build> {
|
|
rocket::build()
|
|
.manage(AppState {
|
|
database,
|
|
registry: Arc::new(DashMap::new()),
|
|
device_counts: Arc::new(DashMap::new()),
|
|
config: Arc::new(Config::from_env()),
|
|
})
|
|
.mount("/", routes![transport::socket::connect, api::mcp::route])
|
|
.register("/", rocket::catchers![api::catchers::default_catcher])
|
|
}
|