第十五章：一个最小但完整的控制塔示例

#[macro_use] extern crate rocket;

use std::sync::atomic::{AtomicU64, Ordering};

use rocket::fairing::{Fairing, Info, Kind};
use rocket::http::Status;
use rocket::request::{self, FromRequest, Outcome, Request};
use rocket::response::status;
use rocket::serde::{Deserialize, Serialize, json::Json};
use rocket::{Data, Response, State};

struct AppState {
    service_name: String,
    launch_count: AtomicU64,
}

#[derive(Debug, Deserialize)]
#[serde(crate = "rocket::serde")]
struct CreateLaunch {
    mission: String,
    rocket: String,
}

#[derive(Debug, Serialize)]
#[serde(crate = "rocket::serde")]
struct LaunchView {
    id: u64,
    mission: String,
    rocket: String,
}

#[derive(Debug, Serialize)]
#[serde(crate = "rocket::serde")]
struct ErrorMessage {
    error: String,
}

struct ApiKey;

#[rocket::async_trait]
impl<'r> FromRequest<'r> for ApiKey {
    type Error = ();

    async fn from_request(req: &'r Request<'_>) -> request::Outcome<Self, Self::Error> {
        match req.headers().get_one("x-api-key") {
            Some("launch-secret") => Outcome::Success(ApiKey),
            _ => Outcome::Error((Status::Unauthorized, ())),
        }
    }
}

struct TowerAudit;

#[rocket::async_trait]
impl Fairing for TowerAudit {
    fn info(&self) -> Info {
        Info {
            name: "Tower Audit",
            kind: Kind::Request | Kind::Response,
        }
    }

    async fn on_request(&self, req: &mut Request<'_>, _data: &mut Data<'_>) {
        req.local_cache(|| std::time::SystemTime::now());
    }

    async fn on_response<'r>(&self, req: &'r Request<'_>, res: &mut Response<'r>) {
        let started = req.local_cache(|| std::time::SystemTime::now());
        if let Ok(duration) = started.elapsed() {
            res.set_raw_header("x-tower-time-ms", duration.as_millis().to_string());
        }
    }
}

#[get("/health")]
fn health(state: &State<AppState>) -> String {
    format!(
        "{} alive, launches={}",
        state.service_name,
        state.launch_count.load(Ordering::Relaxed)
    )
}

#[post("/launches", data = "<input>")]
fn create_launch(
    _key: ApiKey,
    state: &State<AppState>,
    input: Json<CreateLaunch>,
) -> Result<Json<LaunchView>, status::Custom<Json<ErrorMessage>>> {
    if input.mission.trim().is_empty() {
        return Err(status::Custom(
            Status::BadRequest,
            Json(ErrorMessage {
                error: "mission cannot be empty".into(),
            }),
        ));
    }

    let id = state.launch_count.fetch_add(1, Ordering::Relaxed) + 1;

    Ok(Json(LaunchView {
        id,
        mission: input.mission.clone(),
        rocket: input.rocket.clone(),
    }))
}

#[launch]
fn rocket() -> _ {
    rocket::build()
        .manage(AppState {
            service_name: "rocket-control-tower".into(),
            launch_count: AtomicU64::new(0),
        })
        .attach(TowerAudit)
        .mount("/", routes![health, create_launch])
}