lambda-microvm-hook-server 0.1.0

use crate::MicroVmError;
use crate::error::{ApiError, error_response};
use crate::request::RunHookRequest;
use crate::spawn_command::spawn_command;
use crate::state::HookServerState;
use axum::Router;
use axum::extract::{Json, State, rejection::JsonRejection};
use axum::http::StatusCode;
use axum::response::{IntoResponse, Response};
use axum::routing::post;
use serde::Serialize;
use std::sync::Arc;
use tokio::net::TcpListener;

pub const BASE_PATH: &str = "/aws/lambda-microvms/runtime/v1";

#[derive(Serialize)]
struct StatusBody {
    status: &'static str,
}

pub struct HookServer {
    state: Arc<HookServerState>,
}

impl HookServer {
    pub fn new() -> Self {
        Self { state: HookServerState::new() }
    }

    pub async fn serve(self, listener: TcpListener) -> Result<(), MicroVmError> {
        let router = self.router();
        let state = Arc::clone(&self.state);
        axum::serve(listener, router)
            .with_graceful_shutdown(async move {
                state.wait_for_completion().await;
            })
            .await
            .map_err(MicroVmError::Server)?;

        self.state.take_result()
    }

    fn router(&self) -> Router {
        Router::new()
            .route(&format!("{BASE_PATH}/ready"), post(ready))
            .route(&format!("{BASE_PATH}/run"), post(run))
            .route(&format!("{BASE_PATH}/terminate"), post(terminate))
            .fallback(not_found)
            .method_not_allowed_fallback(not_found)
            .with_state(Arc::clone(&self.state))
    }
}

impl Default for HookServer {
    fn default() -> Self {
        Self::new()
    }
}

async fn ready() -> impl IntoResponse {
    (StatusCode::OK, Json(StatusBody { status: "ready" }))
}

async fn run(
    State(state): State<Arc<HookServerState>>,
    request: Result<Json<RunHookRequest>, JsonRejection>,
) -> Result<Json<StatusBody>, ApiError> {
    let Json(request) = request.map_err(|error| ApiError::bad_request(format!("invalid request JSON: {error}")))?;
    if request.microvm_id.trim().is_empty() {
        return Err(ApiError::bad_request("microvmId must not be empty"));
    }

    let mut payload = request.run_hook_payload;
    if payload.command.trim().is_empty() {
        return Err(ApiError::bad_request("command must not be blank"));
    }

    if !state.claim_run() {
        return Err(ApiError::conflict("run already started"));
    }

    payload.environment.insert("AWS_LAMBDA_MICROVM_ID".to_string(), request.microvm_id);
    let command_result = spawn_command(payload.command, payload.args, payload.environment, state.cancellation_token())
        .map_err(|error| ApiError::initialization(&state, error))?;
    state.track_command(command_result);
    Ok(Json(StatusBody { status: "accepted" }))
}

async fn terminate(State(state): State<Arc<HookServerState>>) -> Json<StatusBody> {
    state.cancel();
    if state.claim_run() {
        state.finish(Ok(()));
    }
    state.wait_for_completion().await;
    Json(StatusBody { status: "terminating" })
}

async fn not_found() -> Response {
    error_response(StatusCode::NOT_FOUND, "not found")
}