churn/crates/churn-agent/src/main.rs

mod agent;

use std::{net::SocketAddr, path::PathBuf};

use agent::AgentService;
use anyhow::Error;
use axum::{
    extract::State,
    http::StatusCode,
    response::{IntoResponse, Response},
    routing::{get, post},
    Json, Router,
};
use churn_domain::AgentEnrollReq;
use clap::{Parser, Subcommand};
use serde_json::json;
use wasmtime::{Caller, Engine, Linker, Module, Store};

#[derive(Parser)]
#[command(author, version, about, long_about = None, subcommand_required = true)]
struct Command {
    #[command(subcommand)]
    command: Option<Commands>,
}

#[derive(Subcommand)]
enum Commands {
    Daemon {
        #[arg(env = "CHURN_ADDR", long)]
        host: SocketAddr,
    },

    Connect {
        /// agent name is the hostname which other agents or servers can resolve and connect via. It should be unique
        #[arg(env = "CHURN_AGENT_NAME", long)]
        agent_name: String,

        #[arg(env = "CHURN_ADDR", long)]
        host: SocketAddr,

        #[arg(env = "CHURN_TOKEN", long)]
        token: String,
    },

    Execute {
        #[arg(env = "CHURN_AGENT_EXE", long)]
        exe: PathBuf,

        #[command(subcommand)]
        commands: Option<ExecuteCommands>,
    },
}

#[derive(Subcommand)]
enum ExecuteCommands {
    Source,
}

#[derive(Clone, Default)]
struct AppState {
    agent: AgentService,
}

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    dotenv::dotenv().ok();
    tracing_subscriber::fmt::init();

    let cli = Command::parse();

    handle_command(cli).await?;

    Ok(())
}

async fn handle_command(cmd: Command) -> anyhow::Result<()> {
    match cmd.command.unwrap() {
        Commands::Daemon { host } => {
            tracing::info!("Starting churn server");

            let app = Router::new()
                .route("/enroll", post(enroll))
                .route("/ping", get(ping))
                .with_state(AppState::default());

            tracing::info!("churn server listening on {}", host);
            axum::Server::bind(&host)
                .serve(app.into_make_service())
                .await
                .unwrap();

            Ok(())
        }
        Commands::Connect {
            host: _,
            token: _,
            agent_name: _,
        } => todo!(),
        Commands::Execute { exe, commands } => match commands {
            Some(ExecuteCommands::Source) => Ok(()),
            None => {
                let engine = Engine::default();
                let module = Module::from_file(&engine, exe)?;

                // Create a `Linker` which will be later used to instantiate this module.
                // Host functionality is defined by name within the `Linker`.
                let mut linker = Linker::new(&engine);
                linker.func_wrap("print", "print", |caller: Caller<'_, u32>, param: i32| {
                    println!("Got {} from WebAssembly", param);
                    println!("my host state is: {}", caller.data());
                })?;

                // All wasm objects operate within the context of a "store". Each
                // `Store` has a type parameter to store host-specific data, which in
                // this case we're using `4` for.
                let mut store = Store::new(&engine, 4);
                let instance = linker.instantiate(&mut store, &module)?;
                let hello = instance.get_typed_func::<(), ()>(&mut store, "hello")?;

                // And finally we can call the wasm!
                hello.call(&mut store, ())?;

                Ok(())
            }
        },
    }
}

enum AppError {
    Internal(Error),
}

impl IntoResponse for AppError {
    fn into_response(self) -> Response {
        let (status, error_message) = match self {
            AppError::Internal(e) => {
                tracing::error!("failed with error: {}", e);

                (
                    StatusCode::INTERNAL_SERVER_ERROR,
                    "failed with internal error",
                )
            }
        };

        let body = Json(json!({
            "error": error_message,
        }));

        (status, body).into_response()
    }
}

async fn ping() -> impl IntoResponse {
    "pong!"
}

async fn enroll(
    State(state): State<AppState>,
    Json(req): Json<AgentEnrollReq>,
) -> Result<(), AppError> {
    state
        .agent
        .enroll(&req.agent_name, &req.server, &req.lease)
        .await
        .map_err(AppError::Internal)?;

    Ok(())
}