lambda_microvm_hook_server/
hook_server.rs

1use crate::MicroVmError;
2use crate::error::{ApiError, error_response};
3use crate::request::RunHookRequest;
4use crate::spawn_command::spawn_command;
5use crate::state::HookServerState;
6use axum::Router;
7use axum::extract::{Json, State, rejection::JsonRejection};
8use axum::http::StatusCode;
9use axum::response::{IntoResponse, Response};
10use axum::routing::post;
11use serde::Serialize;
12use std::sync::Arc;
13use tokio::net::TcpListener;
14
15pub const BASE_PATH: &str = "/aws/lambda-microvms/runtime/v1";
16
17#[derive(Serialize)]
18struct StatusBody {
19    status: &'static str,
20}
21
22pub struct HookServer {
23    state: Arc<HookServerState>,
24}
25
26impl HookServer {
27    pub fn new() -> Self {
28        Self { state: HookServerState::new() }
29    }
30
31    pub async fn serve(self, listener: TcpListener) -> Result<(), MicroVmError> {
32        let router = self.router();
33        let state = Arc::clone(&self.state);
34        axum::serve(listener, router)
35            .with_graceful_shutdown(async move {
36                state.wait_for_completion().await;
37            })
38            .await
39            .map_err(MicroVmError::Server)?;
40
41        self.state.take_result()
42    }
43
44    fn router(&self) -> Router {
45        Router::new()
46            .route(&format!("{BASE_PATH}/ready"), post(ready))
47            .route(&format!("{BASE_PATH}/run"), post(run))
48            .route(&format!("{BASE_PATH}/terminate"), post(terminate))
49            .fallback(not_found)
50            .method_not_allowed_fallback(not_found)
51            .with_state(Arc::clone(&self.state))
52    }
53}
54
55impl Default for HookServer {
56    fn default() -> Self {
57        Self::new()
58    }
59}
60
61async fn ready() -> impl IntoResponse {
62    (StatusCode::OK, Json(StatusBody { status: "ready" }))
63}
64
65async fn run(
66    State(state): State<Arc<HookServerState>>,
67    request: Result<Json<RunHookRequest>, JsonRejection>,
68) -> Result<Json<StatusBody>, ApiError> {
69    let Json(request) = request.map_err(|error| ApiError::bad_request(format!("invalid request JSON: {error}")))?;
70    if request.microvm_id.trim().is_empty() {
71        return Err(ApiError::bad_request("microvmId must not be empty"));
72    }
73
74    let mut payload = request.run_hook_payload;
75    if payload.command.trim().is_empty() {
76        return Err(ApiError::bad_request("command must not be blank"));
77    }
78
79    if !state.claim_run() {
80        return Err(ApiError::conflict("run already started"));
81    }
82
83    payload.environment.insert("AWS_LAMBDA_MICROVM_ID".to_string(), request.microvm_id);
84    let command_result = spawn_command(payload.command, payload.args, payload.environment, state.cancellation_token())
85        .map_err(|error| ApiError::initialization(&state, error))?;
86    state.track_command(command_result);
87    Ok(Json(StatusBody { status: "accepted" }))
88}
89
90async fn terminate(State(state): State<Arc<HookServerState>>) -> Json<StatusBody> {
91    state.cancel();
92    if state.claim_run() {
93        state.finish(Ok(()));
94    }
95    state.wait_for_completion().await;
96    Json(StatusBody { status: "terminating" })
97}
98
99async fn not_found() -> Response {
100    error_response(StatusCode::NOT_FOUND, "not found")
101}
lambda_microvm_hook_server/hook_server.rs

lambda_microvm_hook_server/
hook_server.rs
