Skip to main content

orbit_runtime/
server.rs

1//! HTTP API (RFC-0001 §11.2) and the HTTP/JSON binding of the worker
2//! protocol (§8). See `protocol/README.md` for the wire format.
3
4use crate::engine::{Engine, EngineError, RegisterWorker, TaskAssignment};
5use crate::identity::{Identities, WorkerIdentity};
6use crate::model::TaskError;
7use axum::extract::{DefaultBodyLimit, Path, Query, Request, State};
8use axum::http::{HeaderMap, HeaderValue, StatusCode};
9use axum::middleware::{self, Next};
10use axum::response::{IntoResponse, Response};
11use axum::routing::{get, post};
12use axum::{Extension, Json, Router};
13use serde::Deserialize;
14use serde_json::{json, Value};
15use std::net::SocketAddr;
16use std::sync::{Arc, Mutex};
17use std::time::Duration;
18use tokio::sync::Notify;
19
20pub const API_VERSION: &str = "orbit-api/1";
21
22pub struct AppState {
23    pub engine: Mutex<Engine>,
24    pub changed: Arc<Notify>,
25    pub auth_token: Option<String>,
26    /// Optional credential restricted to the worker protocol.
27    pub worker_token: Option<String>,
28    /// Per-worker identities (`--workers-file`).
29    pub identities: Identities,
30}
31
32pub type Shared = Arc<AppState>;
33
34pub struct ApiError(pub(crate) EngineError);
35
36impl From<EngineError> for ApiError {
37    fn from(e: EngineError) -> Self {
38        ApiError(e)
39    }
40}
41
42impl IntoResponse for ApiError {
43    fn into_response(self) -> Response {
44        let status =
45            StatusCode::from_u16(self.0.status).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
46        (
47            status,
48            Json(json!({ "error": { "code": self.0.code, "message": self.0.message } })),
49        )
50            .into_response()
51    }
52}
53
54fn bad_request(msg: impl Into<String>) -> ApiError {
55    ApiError(EngineError::new(400, "INVALID_REQUEST", msg))
56}
57
58type ApiResult<T> = Result<T, ApiError>;
59
60pub(crate) fn engine(state: &Shared) -> std::sync::MutexGuard<'_, Engine> {
61    state.engine.lock().unwrap_or_else(|p| p.into_inner())
62}
63
64pub fn router(state: Shared, body_limit: usize) -> Router {
65    Router::new()
66        .route("/v1/health", get(health))
67        .route("/v1/metrics", get(metrics))
68        .route("/v1/workflows", post(register_workflow).get(list_workflows))
69        .route("/v1/workflows/{name}", get(get_workflow_latest))
70        .route("/v1/workflows/{name}/{version}", get(get_workflow))
71        .route("/v1/executions", post(start_execution).get(list_executions))
72        .route("/v1/executions/{id}", get(get_execution))
73        .route("/v1/executions/{id}/history", get(get_history))
74        .route("/v1/executions/{id}/cancel", post(cancel_execution))
75        .route("/v1/executions/{id}/signals/{name}", post(send_signal))
76        .route("/v1/executions/{id}/queries/{name}", get(run_query))
77        .route("/v1/executions/{id}/wait", get(wait_execution))
78        .route("/v1/workers", get(list_workers))
79        .route("/v1/workers/register", post(register_worker))
80        .route("/v1/workers/heartbeat", post(heartbeat))
81        .route("/v1/workers/poll", post(poll))
82        .route("/v1/workers/unregister", post(unregister))
83        .route("/v1/tasks/complete", post(complete_task))
84        .route("/v1/tasks/fail", post(fail_task))
85        .layer(middleware::from_fn_with_state(state.clone(), auth))
86        .layer(DefaultBodyLimit::max(body_limit))
87        .with_state(state)
88}
89
90/// Endpoints a worker credential may call.
91fn is_worker_path(path: &str) -> bool {
92    path.starts_with("/v1/workers/") || path.starts_with("/v1/tasks/")
93}
94
95/// Who is calling, as established by [`auth`].
96#[derive(Clone)]
97pub enum Caller {
98    /// No authentication is configured (loopback development).
99    Open,
100    Admin,
101    /// The shared, restricted worker token (`--worker-token`).
102    SharedWorker,
103    /// A worker from the workers file, by token or client certificate.
104    Worker(Arc<WorkerIdentity>),
105}
106
107impl Caller {
108    fn identity(&self) -> Option<&WorkerIdentity> {
109        match self {
110            Caller::Worker(w) => Some(w),
111            _ => None,
112        }
113    }
114}
115
116fn forbidden(msg: impl Into<String>) -> ApiError {
117    ApiError(EngineError::new(403, "FORBIDDEN", msg))
118}
119
120/// Authentication.
121///
122/// When any credential is configured, every endpoint except `/v1/health`
123/// requires one. The admin token can call anything. Worker credentials (the
124/// shared `--worker-token`, a per-worker token from the workers file, or a
125/// client certificate under mutual TLS) can only call the worker protocol
126/// endpoints, so a compromised worker host cannot start, cancel or read
127/// executions. Per-worker identities are further limited by the handlers to
128/// their own id, tasks, sessions and task tokens.
129/// Establish the caller from a bearer token and/or a verified client
130/// certificate. Shared by the HTTP API and the gRPC binding. `Ok(None)` means
131/// no credential was presented while authentication is required.
132pub(crate) fn identify(
133    state: &AppState,
134    bearer: Option<&str>,
135    cert_der: Option<&[u8]>,
136) -> Result<Option<Caller>, ApiError> {
137    let secured = state.auth_token.is_some() || !state.identities.is_empty();
138    if !secured {
139        return Ok(Some(Caller::Open));
140    }
141    if let Some(t) = bearer {
142        if state
143            .auth_token
144            .as_ref()
145            .is_some_and(|a| constant_time_eq(t.as_bytes(), a.as_bytes()))
146        {
147            return Ok(Some(Caller::Admin));
148        }
149        if let Some(w) = state.identities.by_token(t) {
150            return Ok(Some(Caller::Worker(Arc::new(w.clone()))));
151        }
152        if state
153            .worker_token
154            .as_ref()
155            .is_some_and(|w| constant_time_eq(t.as_bytes(), w.as_bytes()))
156        {
157            return Ok(Some(Caller::SharedWorker));
158        }
159    }
160    if let Some(der) = cert_der {
161        return match state.identities.by_certificate(der) {
162            Some(w) => Ok(Some(Caller::Worker(Arc::new(w.clone())))),
163            None => Err(forbidden(
164                "the client certificate is valid but matches no worker in the workers file",
165            )),
166        };
167    }
168    Ok(None)
169}
170
171async fn auth(State(state): State<Shared>, mut req: Request, next: Next) -> Response {
172    let path = req.uri().path().to_string();
173    let bearer = req
174        .headers()
175        .get("authorization")
176        .and_then(|v| v.to_str().ok())
177        .and_then(|v| v.strip_prefix("Bearer "))
178        .map(str::to_string);
179    let cert = req
180        .extensions()
181        .get::<Arc<axum_server::tls::TlsInfo>>()
182        .and_then(|info| info.peer_certificates.first().cloned());
183    let caller = match identify(&state, bearer.as_deref(), cert.as_deref()) {
184        Ok(Some(c)) => c,
185        Ok(None) if path == "/v1/health" => Caller::Open,
186        Ok(None) => {
187            return ApiError(EngineError::new(
188                401,
189                "UNAUTHENTICATED",
190                "missing or invalid credentials",
191            ))
192            .into_response()
193        }
194        Err(e) => return e.into_response(),
195    };
196    if matches!(caller, Caller::SharedWorker | Caller::Worker(_))
197        && !is_worker_path(&path)
198        && path != "/v1/health"
199    {
200        return forbidden("worker credentials may only call /v1/workers/* and /v1/tasks/*")
201            .into_response();
202    }
203    req.extensions_mut().insert(caller);
204    let mut resp = next.run(req).await;
205    resp.headers_mut()
206        .insert("x-orbit-api-version", HeaderValue::from_static(API_VERSION));
207    resp
208}
209
210/// A per-worker identity may only use sessions it opened.
211pub(crate) fn authorize_session(state: &Shared, caller: &Caller, session: &str) -> ApiResult<()> {
212    if let Some(id) = caller.identity() {
213        if let Some(owner) = engine(state).session_worker(session) {
214            if owner != id.id {
215                return Err(forbidden(format!(
216                    "session belongs to worker `{owner}`, not `{}`",
217                    id.id
218                )));
219            }
220        }
221    }
222    Ok(())
223}
224
225/// A per-worker identity may only report results for tokens issued to it.
226pub(crate) fn authorize_token(state: &Shared, caller: &Caller, token: &str) -> ApiResult<()> {
227    if let Some(id) = caller.identity() {
228        if let Some(owner) = engine(state).token_worker(token) {
229            if owner != id.id {
230                return Err(forbidden(format!(
231                    "task token was issued to worker `{owner}`, not `{}`",
232                    id.id
233                )));
234            }
235        }
236    }
237    Ok(())
238}
239
240fn constant_time_eq(a: &[u8], b: &[u8]) -> bool {
241    a.len() == b.len() && a.iter().zip(b).fold(0u8, |acc, (x, y)| acc | (x ^ y)) == 0
242}
243
244async fn health(State(state): State<Shared>) -> Json<Value> {
245    let e = engine(&state);
246    Json(json!({
247        "status": "ok",
248        "api_version": API_VERSION,
249        "runtime_version": env!("CARGO_PKG_VERSION"),
250        "ir_versions": orbit_ir::SUPPORTED_IR_VERSIONS,
251        "active_executions": e.active_count(),
252        "worker_sessions": e.session_count(),
253    }))
254}
255
256async fn metrics(State(state): State<Shared>) -> Response {
257    let e = engine(&state);
258    let body = e.metrics.render(e.active_count(), e.session_count());
259    let mut headers = HeaderMap::new();
260    headers.insert(
261        "content-type",
262        HeaderValue::from_static("text/plain; version=0.0.4"),
263    );
264    (headers, body).into_response()
265}
266
267async fn register_workflow(State(state): State<Shared>, body: String) -> ApiResult<Response> {
268    let ir = orbit_ir::WorkflowIr::from_json(&body)
269        .map_err(|e| ApiError(EngineError::new(400, "INVALID_IR", e.0)))?;
270    let (name, version, hash) = (
271        ir.name.clone(),
272        ir.version.clone(),
273        ir.definition_hash.clone(),
274    );
275    let created = engine(&state).register_definition(ir)?;
276    let status = if created {
277        StatusCode::CREATED
278    } else {
279        StatusCode::OK
280    };
281    Ok((status, Json(json!({ "name": name, "version": version, "definition_hash": hash, "created": created }))).into_response())
282}
283
284fn summary(d: &orbit_ir::WorkflowIr) -> Value {
285    json!({
286        "name": d.name,
287        "version": d.version,
288        "description": d.description,
289        "definition_hash": d.definition_hash,
290        "steps": d.steps.iter().map(|s| s.id.clone()).collect::<Vec<_>>(),
291        "tasks": d.contracts.keys().collect::<Vec<_>>(),
292    })
293}
294
295async fn list_workflows(State(state): State<Shared>) -> Json<Value> {
296    let defs = engine(&state).definitions();
297    Json(json!({ "workflows": defs.iter().map(|d| summary(d)).collect::<Vec<_>>() }))
298}
299
300async fn get_workflow(
301    State(state): State<Shared>,
302    Path((name, version)): Path<(String, String)>,
303) -> ApiResult<Json<Value>> {
304    let d = engine(&state)
305        .definition(&name, Some(&version))
306        .ok_or_else(|| {
307            ApiError(EngineError::new(
308                404,
309                "WORKFLOW_NOT_FOUND",
310                "workflow not found",
311            ))
312        })?;
313    Ok(Json(serde_json::to_value(&*d).unwrap_or_default()))
314}
315
316async fn get_workflow_latest(
317    State(state): State<Shared>,
318    Path(name): Path<String>,
319) -> ApiResult<Json<Value>> {
320    let d = engine(&state).definition(&name, None).ok_or_else(|| {
321        ApiError(EngineError::new(
322            404,
323            "WORKFLOW_NOT_FOUND",
324            "workflow not found",
325        ))
326    })?;
327    Ok(Json(serde_json::to_value(&*d).unwrap_or_default()))
328}
329
330#[derive(Deserialize)]
331struct StartRequest {
332    workflow: String,
333    #[serde(default)]
334    version: Option<String>,
335    #[serde(default)]
336    input: Value,
337    #[serde(default)]
338    execution_id: Option<String>,
339}
340
341async fn start_execution(
342    State(state): State<Shared>,
343    Json(req): Json<StartRequest>,
344) -> ApiResult<Response> {
345    let r = engine(&state).start(
346        &req.workflow,
347        req.version.as_deref(),
348        req.input,
349        req.execution_id,
350    )?;
351    let status = if r.created {
352        StatusCode::CREATED
353    } else {
354        StatusCode::OK
355    };
356    Ok((status, Json(serde_json::to_value(r).unwrap_or_default())).into_response())
357}
358
359#[derive(Deserialize)]
360struct ListQuery {
361    status: Option<String>,
362    workflow: Option<String>,
363    limit: Option<u32>,
364}
365
366async fn list_executions(
367    State(state): State<Shared>,
368    Query(q): Query<ListQuery>,
369) -> ApiResult<Json<Value>> {
370    let list = engine(&state).list_executions(
371        q.status.as_deref().map(str::to_uppercase).as_deref(),
372        q.workflow.as_deref(),
373        q.limit.unwrap_or(50),
374    )?;
375    Ok(Json(json!({ "executions": list })))
376}
377
378async fn get_execution(
379    State(state): State<Shared>,
380    Path(id): Path<String>,
381) -> ApiResult<Json<Value>> {
382    let s = engine(&state).execution(&id)?;
383    Ok(Json(serde_json::to_value(s).unwrap_or_default()))
384}
385
386async fn get_history(
387    State(state): State<Shared>,
388    Path(id): Path<String>,
389) -> ApiResult<Json<Value>> {
390    let h = engine(&state).history(&id)?;
391    Ok(Json(json!({ "execution_id": id, "events": h })))
392}
393
394#[derive(Deserialize, Default)]
395struct CancelRequest {
396    #[serde(default)]
397    reason: Option<String>,
398}
399
400async fn cancel_execution(
401    State(state): State<Shared>,
402    Path(id): Path<String>,
403    body: Option<Json<CancelRequest>>,
404) -> ApiResult<Json<Value>> {
405    let reason = body
406        .and_then(|b| b.0.reason)
407        .unwrap_or_else(|| "requested via API".into());
408    let s = engine(&state).cancel(&id, &reason)?;
409    Ok(Json(serde_json::to_value(s).unwrap_or_default()))
410}
411
412/// `POST /v1/executions/{id}/signals/{name}` with the payload as the body
413/// (an empty body sends `{}`).
414async fn send_signal(
415    State(state): State<Shared>,
416    Path((id, name)): Path<(String, String)>,
417    body: Option<Json<Value>>,
418) -> ApiResult<Json<Value>> {
419    let payload = body.map(|b| b.0).unwrap_or_else(|| json!({}));
420    let seq = engine(&state).signal(&id, &name, payload)?;
421    Ok(Json(json!({
422        "execution_id": id,
423        "signal": name,
424        "accepted": true,
425        "seq": seq,
426    })))
427}
428
429/// `GET /v1/executions/{id}/queries/{name}`: evaluate a declared query.
430async fn run_query(
431    State(state): State<Shared>,
432    Path((id, name)): Path<(String, String)>,
433) -> ApiResult<Json<Value>> {
434    let value = engine(&state).query(&id, &name)?;
435    Ok(Json(
436        json!({ "execution_id": id, "query": name, "value": value }),
437    ))
438}
439
440#[derive(Deserialize)]
441struct WaitQuery {
442    timeout_ms: Option<u64>,
443}
444
445/// Long-poll until the execution reaches a terminal status (or the timeout).
446async fn wait_execution(
447    State(state): State<Shared>,
448    Path(id): Path<String>,
449    Query(q): Query<WaitQuery>,
450) -> ApiResult<Json<Value>> {
451    let deadline = tokio::time::Instant::now()
452        + Duration::from_millis(q.timeout_ms.unwrap_or(30_000).min(300_000));
453    loop {
454        let notified = state.changed.notified();
455        tokio::pin!(notified);
456        notified.as_mut().enable();
457        let s = engine(&state).execution(&id)?;
458        if s.status.is_terminal() || tokio::time::Instant::now() >= deadline {
459            return Ok(Json(serde_json::to_value(s).unwrap_or_default()));
460        }
461        let _ = tokio::time::timeout_at(deadline, notified).await;
462    }
463}
464
465async fn list_workers(State(state): State<Shared>) -> ApiResult<Json<Value>> {
466    Ok(Json(json!({ "workers": engine(&state).workers()? })))
467}
468
469async fn register_worker(
470    State(state): State<Shared>,
471    Extension(caller): Extension<Caller>,
472    Json(req): Json<RegisterWorker>,
473) -> ApiResult<Json<Value>> {
474    if let Some(id) = caller.identity() {
475        if req.worker_id != id.id {
476            return Err(forbidden(format!(
477                "this credential belongs to worker `{}`, not `{}`",
478                id.id, req.worker_id
479            )));
480        }
481        if let Some(rt) = &id.runtime {
482            if !rt.eq_ignore_ascii_case(&req.runtime) {
483                return Err(forbidden(format!(
484                    "worker `{}` must announce runtime `{rt}`",
485                    id.id
486                )));
487            }
488        }
489        let denied: Vec<&str> = req
490            .contracts
491            .iter()
492            .map(|c| c.task_name.as_str())
493            .filter(|t| !id.may_serve(t))
494            .collect();
495        if !denied.is_empty() {
496            return Err(forbidden(format!(
497                "worker `{}` may not serve: {} (allowed: {})",
498                id.id,
499                denied.join(", "),
500                id.tasks.join(", ")
501            )));
502        }
503    }
504    Ok(Json(
505        serde_json::to_value(engine(&state).register_worker(req)?).unwrap_or_default(),
506    ))
507}
508
509#[derive(Deserialize)]
510struct SessionRequest {
511    worker_session_id: String,
512    #[serde(default)]
513    wait_ms: Option<u64>,
514}
515
516async fn heartbeat(
517    State(state): State<Shared>,
518    Extension(caller): Extension<Caller>,
519    Json(req): Json<SessionRequest>,
520) -> ApiResult<Json<Value>> {
521    authorize_session(&state, &caller, &req.worker_session_id)?;
522    Ok(Json(
523        serde_json::to_value(engine(&state).heartbeat(&req.worker_session_id)?).unwrap_or_default(),
524    ))
525}
526
527async fn unregister(
528    State(state): State<Shared>,
529    Extension(caller): Extension<Caller>,
530    Json(req): Json<SessionRequest>,
531) -> ApiResult<Json<Value>> {
532    authorize_session(&state, &caller, &req.worker_session_id)?;
533    engine(&state).unregister_worker(&req.worker_session_id)?;
534    Ok(Json(json!({ "accepted": true })))
535}
536
537/// Long-poll for a task. Returns 204 when nothing arrived within `wait_ms`.
538async fn poll(
539    State(state): State<Shared>,
540    Extension(caller): Extension<Caller>,
541    Json(req): Json<SessionRequest>,
542) -> ApiResult<Response> {
543    authorize_session(&state, &caller, &req.worker_session_id)?;
544    let wait = Duration::from_millis(req.wait_ms.unwrap_or(20_000).min(60_000));
545    let deadline = tokio::time::Instant::now() + wait;
546    loop {
547        let notified = state.changed.notified();
548        tokio::pin!(notified);
549        notified.as_mut().enable();
550        let task: Option<TaskAssignment> = engine(&state).poll(&req.worker_session_id)?;
551        if let Some(t) = task {
552            return Ok(Json(serde_json::to_value(t).unwrap_or_default()).into_response());
553        }
554        if tokio::time::Instant::now() >= deadline {
555            return Ok(StatusCode::NO_CONTENT.into_response());
556        }
557        let _ = tokio::time::timeout_at(deadline, notified).await;
558    }
559}
560
561#[derive(Deserialize)]
562struct CompleteRequest {
563    token: String,
564    #[serde(default)]
565    output: Value,
566}
567
568async fn complete_task(
569    State(state): State<Shared>,
570    Extension(caller): Extension<Caller>,
571    Json(req): Json<CompleteRequest>,
572) -> ApiResult<Json<Value>> {
573    if req.token.is_empty() {
574        return Err(bad_request("token is required"));
575    }
576    authorize_token(&state, &caller, &req.token)?;
577    Ok(Json(
578        serde_json::to_value(engine(&state).complete(&req.token, req.output)?).unwrap_or_default(),
579    ))
580}
581
582#[derive(Deserialize)]
583struct FailRequest {
584    token: String,
585    error_code: String,
586    #[serde(default)]
587    error_message: String,
588    #[serde(default = "default_true")]
589    retryable: bool,
590    #[serde(default)]
591    details: Option<Value>,
592}
593
594fn default_true() -> bool {
595    true
596}
597
598async fn fail_task(
599    State(state): State<Shared>,
600    Extension(caller): Extension<Caller>,
601    Json(req): Json<FailRequest>,
602) -> ApiResult<Json<Value>> {
603    if req.error_code.is_empty() {
604        return Err(bad_request("error_code is required"));
605    }
606    authorize_token(&state, &caller, &req.token)?;
607    let error = TaskError {
608        code: req.error_code,
609        message: req.error_message,
610        retryable: req.retryable,
611        details: req.details,
612    };
613    Ok(Json(
614        serde_json::to_value(engine(&state).fail(&req.token, error)?).unwrap_or_default(),
615    ))
616}
617
618/// Background loop: timers, retries, deadlines, liveness and built-in tasks.
619pub async fn run_ticker(state: Shared, interval: Duration) {
620    let mut ticker = tokio::time::interval(interval);
621    ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
622    loop {
623        ticker.tick().await;
624        let jobs = match engine(&state).tick() {
625            Ok(j) => j,
626            Err(e) => {
627                tracing::error!(error = %e, "tick failed");
628                continue;
629            }
630        };
631        for job in jobs {
632            let st = state.clone();
633            tokio::spawn(async move {
634                let input = job.input.clone();
635                let timeout = job.timeout_ms;
636                let result = tokio::task::spawn_blocking(move || {
637                    crate::http_adapter::execute(&input, timeout)
638                })
639                .await
640                .unwrap_or_else(|e| Err(TaskError::new("INTERNAL", e.to_string(), true)));
641                let mut e = engine(&st);
642                let r = match result {
643                    Ok(out) => e.complete(&job.token, out),
644                    Err(err) => e.fail(&job.token, err),
645                };
646                if let Err(err) = r {
647                    tracing::error!(error = %err, "recording built-in task result failed");
648                }
649            });
650        }
651    }
652}
653
654/// PEM files for serving HTTPS directly.
655pub struct TlsConfig {
656    pub cert: std::path::PathBuf,
657    pub key: std::path::PathBuf,
658    /// Enables mutual TLS: client certificates signed by this CA are verified
659    /// and identify workers. Clients without a certificate can still connect
660    /// and authenticate with a token.
661    pub client_ca: Option<std::path::PathBuf>,
662}
663
664fn io_err(msg: String) -> std::io::Error {
665    std::io::Error::new(std::io::ErrorKind::InvalidInput, msg)
666}
667
668fn read_pem_certs(
669    path: &std::path::Path,
670) -> std::io::Result<Vec<rustls::pki_types::CertificateDer<'static>>> {
671    let pem =
672        std::fs::read(path).map_err(|e| io_err(format!("cannot read {}: {e}", path.display())))?;
673    let certs = rustls_pemfile::certs(&mut pem.as_slice())
674        .collect::<Result<Vec<_>, _>>()
675        .map_err(|e| io_err(format!("{}: {e}", path.display())))?;
676    if certs.is_empty() {
677        return Err(io_err(format!(
678            "{} contains no certificates",
679            path.display()
680        )));
681    }
682    Ok(certs)
683}
684
685/// Build the rustls server configuration (with client verification for mTLS).
686fn rustls_config(tls: &TlsConfig) -> std::io::Result<rustls::ServerConfig> {
687    let chain = read_pem_certs(&tls.cert)?;
688    let key_pem = std::fs::read(&tls.key)
689        .map_err(|e| io_err(format!("cannot read {}: {e}", tls.key.display())))?;
690    let key = rustls_pemfile::private_key(&mut key_pem.as_slice())
691        .map_err(|e| io_err(format!("{}: {e}", tls.key.display())))?
692        .ok_or_else(|| io_err(format!("{} contains no private key", tls.key.display())))?;
693    let provider = Arc::new(rustls::crypto::aws_lc_rs::default_provider());
694    let builder = rustls::ServerConfig::builder_with_provider(provider.clone())
695        .with_safe_default_protocol_versions()
696        .map_err(|e| io_err(e.to_string()))?;
697    let builder = match &tls.client_ca {
698        Some(ca) => {
699            let mut roots = rustls::RootCertStore::empty();
700            for c in read_pem_certs(ca)? {
701                roots
702                    .add(c)
703                    .map_err(|e| io_err(format!("{}: {e}", ca.display())))?;
704            }
705            let verifier = rustls::server::WebPkiClientVerifier::builder_with_provider(
706                Arc::new(roots),
707                provider,
708            )
709            .allow_unauthenticated()
710            .build()
711            .map_err(|e| io_err(format!("client CA: {e}")))?;
712            builder.with_client_cert_verifier(verifier)
713        }
714        None => builder.with_no_client_auth(),
715    };
716    let mut config = builder
717        .with_single_cert(chain, key)
718        .map_err(|e| io_err(format!("TLS certificate/key: {e}")))?;
719    config.alpn_protocols = vec![b"h2".to_vec(), b"http/1.1".to_vec()];
720    Ok(config)
721}
722
723pub struct ServerConfig {
724    pub addr: SocketAddr,
725    pub auth_token: Option<String>,
726    pub worker_token: Option<String>,
727    pub identities: Identities,
728    pub tls: Option<TlsConfig>,
729    /// Also serve the gRPC binding of the worker protocol on this address.
730    pub grpc_addr: Option<SocketAddr>,
731    /// Pre-bound listeners (socket activation, tests): used instead of
732    /// binding `addr` / `grpc_addr`.
733    pub listener: Option<std::net::TcpListener>,
734    pub grpc_listener: Option<std::net::TcpListener>,
735    pub body_limit: usize,
736    pub tick_interval: Duration,
737}
738
739pub async fn serve(
740    engine: Engine,
741    config: ServerConfig,
742    shutdown: impl std::future::Future<Output = ()> + Send + 'static,
743) -> std::io::Result<()> {
744    // Several rustls crypto backends are compiled in (via different
745    // dependencies); choose one before any TLS configuration is built.
746    let _ = rustls::crypto::aws_lc_rs::default_provider().install_default();
747    let changed = engine.changed.clone();
748    let state = Arc::new(AppState {
749        engine: Mutex::new(engine),
750        changed,
751        auth_token: config.auth_token,
752        worker_token: config.worker_token,
753        identities: config.identities,
754    });
755    let app = router(state.clone(), config.body_limit);
756    let ticker = tokio::spawn(run_ticker(state.clone(), config.tick_interval));
757    // One shutdown signal fanned out to every listener.
758    let (stop_tx, stop_rx) = tokio::sync::watch::channel(false);
759    tokio::spawn(async move {
760        shutdown.await;
761        let _ = stop_tx.send(true);
762    });
763    let stopped = |mut rx: tokio::sync::watch::Receiver<bool>| async move {
764        let _ = rx.wait_for(|v| *v).await;
765    };
766    let bind = |addr: SocketAddr| {
767        std::net::TcpListener::bind(addr)
768            .map_err(|e| std::io::Error::new(e.kind(), format!("cannot listen on {addr}: {e}")))
769    };
770    let grpc_listener = match (config.grpc_listener, config.grpc_addr) {
771        (Some(l), _) => Some(l),
772        (None, Some(addr)) => Some(bind(addr)?),
773        (None, None) => None,
774    };
775    let http_listener = match config.listener {
776        Some(l) => l,
777        None => bind(config.addr)?,
778    };
779    http_listener.set_nonblocking(true)?;
780    let http_addr = http_listener.local_addr()?;
781    let grpc = match grpc_listener {
782        Some(std_listener) => {
783            std_listener.set_nonblocking(true)?;
784            let addr = std_listener.local_addr()?;
785            let listener = tokio::net::TcpListener::from_std(std_listener)?;
786            let incoming =
787                tonic::transport::server::TcpIncoming::from_listener(listener, true, None)
788                    .map_err(|e| io_err(format!("gRPC listener: {e}")))?;
789            let mut builder = tonic::transport::Server::builder();
790            if let Some(tls) = &config.tls {
791                let _ = rustls::crypto::aws_lc_rs::default_provider().install_default();
792                let read = |p: &std::path::Path| {
793                    std::fs::read(p)
794                        .map_err(|e| io_err(format!("cannot read {}: {e}", p.display())))
795                };
796                let mut t = tonic::transport::ServerTlsConfig::new().identity(
797                    tonic::transport::Identity::from_pem(read(&tls.cert)?, read(&tls.key)?),
798                );
799                if let Some(ca) = &tls.client_ca {
800                    t = t
801                        .client_ca_root(tonic::transport::Certificate::from_pem(read(ca)?))
802                        .client_auth_optional(true);
803                }
804                builder = builder
805                    .tls_config(t)
806                    .map_err(|e| io_err(format!("gRPC TLS: {e}")))?;
807            }
808            let service = crate::grpc::GrpcWorkers::new(state.clone());
809            let rx = stop_rx.clone();
810            tracing::info!(addr = %addr, tls = config.tls.is_some(), "gRPC worker protocol listening");
811            Some(tokio::spawn(async move {
812                if let Err(e) = builder
813                    .add_service(service)
814                    .serve_with_incoming_shutdown(incoming, stopped(rx))
815                    .await
816                {
817                    tracing::error!(error = %e, "gRPC server failed");
818                }
819            }))
820        }
821        None => None,
822    };
823    let shutdown = stopped(stop_rx);
824    let result = match config.tls {
825        None => {
826            let listener = tokio::net::TcpListener::from_std(http_listener)?;
827            tracing::info!(addr = %http_addr, tls = false, "ORBIT runtime listening");
828            axum::serve(listener, app)
829                .with_graceful_shutdown(shutdown)
830                .await
831        }
832        Some(tls) => {
833            // Several rustls crypto backends are compiled in (via different
834            // dependencies); pick one explicitly so the choice is not ambiguous.
835            let _ = rustls::crypto::aws_lc_rs::default_provider().install_default();
836            let rustls =
837                axum_server::tls_rustls::RustlsConfig::from_config(Arc::new(rustls_config(&tls)?));
838            let handle = axum_server::Handle::new();
839            let h = handle.clone();
840            tokio::spawn(async move {
841                shutdown.await;
842                h.graceful_shutdown(Some(Duration::from_secs(10)));
843            });
844            tracing::info!(addr = %http_addr, tls = true, mtls = tls.client_ca.is_some(), "ORBIT runtime listening");
845            axum_server::from_tcp_rustls(http_listener, rustls)?
846                .handle(handle)
847                .serve(app.into_make_service())
848                .await
849        }
850    };
851    ticker.abort();
852    if let Some(g) = grpc {
853        let _ = g.await;
854    }
855    result
856}