steep/crates/steep-server/src/api.rs

use axum::body::{Bytes, to_bytes};
use axum::extract::{FromRef, FromRequest, Path, State};
use axum::http::Request;
use axum::response::IntoResponse;
use axum::response::sse::{Event, KeepAlive, Sse};
use axum::routing::{get, post};
use axum::{Json, Router};
use futures_util::stream::{self, Stream, StreamExt};
use reqwest::StatusCode;
use ring::hmac;
use sentry::integrations::tower::{NewSentryLayer, SentryHttpLayer};
use serde_json::{Value, json};
use tokio_stream::wrappers::{BroadcastStream, errors::BroadcastStreamRecvError};
use tower::ServiceBuilder;
use tower_http::services::ServeDir;
use tower_http::trace::TraceLayer;
use tracing::{info, instrument, warn};

use crate::consts::{GITEA_EVENT_TYPE_HEADER_NAME, GITEA_SIG_HEADER_NAME, MAX_WEBHOOK_BODY_SIZE};
use crate::errors::AppError;
use crate::event_bus::JobEvent;
use crate::gitea::WebhookType;
use crate::metrics;
use crate::schemas::JobStatus;
use crate::state::AppState;

pub async fn start(app_state: AppState) -> anyhow::Result<()> {
    let http_port = app_state.config.http_port;
    let shutdown = app_state.shutdown.clone();

    // The dashboard is served from a static directory; the API routes below take
    // precedence (the directory is only the fallback).
    let front_dir = std::env::var("FRONT_DIR").unwrap_or_else(|_| String::from("front"));
    if !std::path::Path::new(&front_dir).is_dir() {
        warn!(dir = %front_dir, "front directory not found, the dashboard will 404");
    }
    let dashboard = ServeDir::new(&front_dir);

    let app = Router::new()
        .route("/webhook", post(webhook))
        .route("/reviews", get(reviews_list))
        .route("/jobs", get(jobs_list))
        .route("/events", get(lifecycle_events))
        .route("/jobs/{id}/events", get(job_events))
        .fallback_service(dashboard)
        .layer(
            ServiceBuilder::new()
                .layer(NewSentryLayer::<Request<_>>::new_from_top())
                .layer(SentryHttpLayer::new())
                .layer(TraceLayer::new_for_http()),
        )
        .with_state(app_state);

    let listener = tokio::net::TcpListener::bind(format!("0.0.0.0:{}", http_port)).await?;
    info!("Listening API on port {}", http_port);

    axum::serve(listener, app)
        .with_graceful_shutdown(async move { shutdown.cancelled().await })
        .await
        .map_err(anyhow::Error::from)
        .inspect(|_| info!("API shutting down complete"))?;

    Ok(())
}

async fn reviews_list(State(app_state): State<AppState>) -> Result<impl IntoResponse, AppError> {
    let reviews = app_state.reviews.list().await?;

    Ok(Json(json!({
        "reviews": reviews
    })))
}

async fn jobs_list(State(app_state): State<AppState>) -> Result<impl IntoResponse, AppError> {
    let pending = app_state.queue.list(JobStatus::Pending).await?;
    let processing = app_state.queue.list(JobStatus::Processing).await?;

    Ok(Json(json!({
        "pending": pending,
        "processing": processing,
    })))
}

/// SSE stream of the global lifecycle events (the job list).
///
/// The stream is closed as soon as the graceful shutdown starts: a live SSE
/// connection would otherwise keep the server from stopping.
async fn lifecycle_events(
    State(app_state): State<AppState>,
) -> Sse<impl Stream<Item = Result<Event, axum::Error>>> {
    let rx = app_state.event_bus.subscribe_lifecycle();
    let shutdown = app_state.shutdown.cancelled_owned();

    let stream = BroadcastStream::new(rx).map(|result| match result {
        Ok(event) => sse_json("lifecycle", &event),
        Err(BroadcastStreamRecvError::Lagged(skipped)) => sse_gap(skipped),
    });

    Sse::new(stream.take_until(shutdown)).keep_alive(KeepAlive::default())
}

async fn job_events(
    State(app_state): State<AppState>,
    Path(id): Path<u64>,
) -> Result<Sse<impl Stream<Item = Result<Event, axum::Error>>>, AppError> {
    let subscription = app_state
        .event_bus
        .subscribe_job(id)
        .ok_or(AppError::JobNotFoundErr(id))?;
    let shutdown = app_state.shutdown.cancelled_owned();

    let history = subscription
        .history
        .into_iter()
        .map(|event| sse_json("progress", &event));

    let live = BroadcastStream::new(subscription.live).map(|result| match result {
        Ok(event) => sse_json("progress", &event),
        Err(BroadcastStreamRecvError::Lagged(skipped)) => sse_gap(skipped),
    });

    let stream = stream::iter(history).chain(live).take_until(shutdown);

    Ok(Sse::new(stream).keep_alive(KeepAlive::default()))
}

fn sse_json<T: serde::Serialize>(name: &'static str, payload: &T) -> Result<Event, axum::Error> {
    Event::default().event(name).json_data(payload)
}

fn sse_gap(skipped: u64) -> Result<Event, axum::Error> {
    Ok(Event::default().event("gap").data(skipped.to_string()))
}

#[instrument(skip(app_state), fields(webhook_type), err)]
async fn webhook(
    State(app_state): State<AppState>,
    WebhookExtract(wb): WebhookExtract,
) -> Result<impl IntoResponse, AppError> {
    tracing::Span::current().record("webhook_type", tracing::field::debug(&wb));

    let event_type = wb.event_type_str();
    metrics::webhook_received(event_type);

    let event_id = wb.event_id();
    match is_known(&app_state, event_id).await {
        Ok(true) => {
            metrics::webhook_duplicate(event_type);
            return Err(AppError::AlreadyProcessedErr);
        }
        Ok(false) => {}
        Err(err) => {
            metrics::webhook_enqueue_failed(event_type);
            return Err(AppError::EnqueueErr(anyhow::Error::from(err)));
        }
    }

    let job = match app_state.queue.push(wb).await {
        Ok(job) => job,
        Err(err) => {
            metrics::webhook_enqueue_failed(event_type);
            return Err(AppError::EnqueueErr(anyhow::Error::from(err)));
        }
    };

    app_state.event_bus.publish(JobEvent::Queued {
        id: job.id,
        repo: job.repo,
        pr: job.pr,
    });

    Ok((StatusCode::CREATED, "Task started"))
}

/// Whether an event is already queued, or already reviewed.
async fn is_known(app_state: &AppState, event_id: u64) -> toasty::Result<bool> {
    Ok(app_state.queue.is_queued(event_id).await?
        || app_state.reviews.is_reviewed(event_id).await?)
}

pub struct WebhookExtract(pub WebhookType);

impl<S> FromRequest<S> for WebhookExtract
where
    AppState: FromRef<S>,
    S: Send + Sync,
{
    type Rejection = AppError;

    #[instrument(skip(req, state), err)]
    async fn from_request(req: axum::extract::Request, state: &S) -> Result<Self, Self::Rejection> {
        let app_state = AppState::from_ref(state);
        let headers = req.headers();

        let sig_header = extract_header(GITEA_SIG_HEADER_NAME, headers)?;
        let type_header = extract_header(GITEA_EVENT_TYPE_HEADER_NAME, headers)?;
        let body_bytes = read_body(req.into_body()).await?;

        verify_signature(
            app_state.config.webhook_secret.as_bytes(),
            &sig_header,
            &body_bytes,
        )?;

        let body_str = String::from_utf8_lossy(&body_bytes).into_owned();
        sentry::configure_scope(|scope| {
            scope.add_event_processor(move |mut event| {
                let mut request = event.request.take().unwrap_or_default();
                request.data = Some(body_str.clone());
                event.request = Some(request);
                Some(event)
            });
        });

        let webhook = parse_webhook(&type_header, &app_state.bot_name, &body_bytes)?;
        Ok(WebhookExtract(webhook))
    }
}

fn extract_header(key: &str, headers: &axum::http::HeaderMap) -> Result<String, AppError> {
    let value = headers
        .get(key)
        .ok_or(AppError::WebHookMissingHeaderErr(key.into()))?
        .to_str()
        .map_err(anyhow::Error::from)?;

    Ok(value.to_owned())
}

async fn read_body(body: axum::body::Body) -> Result<Bytes, AppError> {
    to_bytes(body, MAX_WEBHOOK_BODY_SIZE)
        .await
        .map_err(anyhow::Error::from)
        .map_err(AppError::from)
}

fn parse_webhook(header: &str, bot_name: &str, body_bytes: &[u8]) -> Result<WebhookType, AppError> {
    let Json(value) =
        Json::<Value>::from_bytes(body_bytes).map_err(|_| AppError::MalformedJsonErr)?;

    WebhookType::from_event(header, bot_name, value)
}

fn verify_signature(secret_key: &[u8], sig_header: &str, body: &[u8]) -> Result<(), AppError> {
    let sig_header_decoded =
        hex::decode(sig_header).map_err(|_| AppError::WebHookSigHeaderInvalidErr)?;
    let key = hmac::Key::new(hmac::HMAC_SHA256, secret_key);

    hmac::verify(&key, body, &sig_header_decoded).map_err(|_| AppError::WebHookSigHeaderInvalidErr)
}