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