1use 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 pub worker_token: Option<String>,
28 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/ready", get(ready))
71 .route("/v1/metrics", get(metrics))
72 .route("/v1/workflows", post(register_workflow).get(list_workflows))
73 .route("/v1/workflows/{name}", get(get_workflow_latest))
74 .route("/v1/workflows/{name}/{version}", get(get_workflow))
75 .route("/v1/executions", post(start_execution).get(list_executions))
76 .route("/v1/executions/{id}", get(get_execution))
77 .route("/v1/executions/{id}/history", get(get_history))
78 .route("/v1/executions/{id}/cancel", post(cancel_execution))
79 .route("/v1/executions/{id}/signals/{name}", post(send_signal))
80 .route("/v1/executions/{id}/queries/{name}", get(run_query))
81 .route("/v1/executions/{id}/wait", get(wait_execution))
82 .route("/v1/schedules", post(create_schedule).get(list_schedules))
83 .route(
84 "/v1/schedules/{id}",
85 get(get_schedule)
86 .patch(update_schedule)
87 .delete(delete_schedule),
88 )
89 .route("/v1/schedules/{id}/pause", post(pause_schedule))
90 .route("/v1/schedules/{id}/resume", post(resume_schedule))
91 .route("/v1/schedules/{id}/trigger", post(trigger_schedule))
92 .route("/v1/workers", get(list_workers))
93 .route("/v1/workers/register", post(register_worker))
94 .route("/v1/workers/heartbeat", post(heartbeat))
95 .route("/v1/workers/poll", post(poll))
96 .route("/v1/workers/unregister", post(unregister))
97 .route("/v1/tasks/complete", post(complete_task))
98 .route("/v1/tasks/fail", post(fail_task))
99 .layer(middleware::from_fn_with_state(state.clone(), auth))
100 .layer(DefaultBodyLimit::max(body_limit))
101 .with_state(state)
102}
103
104fn is_worker_path(path: &str) -> bool {
106 path.starts_with("/v1/workers/") || path.starts_with("/v1/tasks/")
107}
108
109#[derive(Clone)]
111pub enum Caller {
112 Open,
114 Admin,
115 SharedWorker,
117 Worker(Arc<WorkerIdentity>),
119 Client(Arc<crate::identity::ClientIdentity>),
122}
123
124impl Caller {
125 fn identity(&self) -> Option<&WorkerIdentity> {
126 match self {
127 Caller::Worker(w) => Some(w),
128 _ => None,
129 }
130 }
131}
132
133fn forbidden(msg: impl Into<String>) -> ApiError {
134 ApiError(EngineError::new(403, "FORBIDDEN", msg))
135}
136
137fn allow(caller: &Caller, action: &str, workflow: &str) -> ApiResult<()> {
139 if let Caller::Client(c) = caller {
140 if !c.allows(action, workflow) {
141 return Err(forbidden(format!(
142 "client `{}` may not {action} executions of `{workflow}`",
143 c.id
144 )));
145 }
146 }
147 Ok(())
148}
149
150fn allow_definition(caller: &Caller, workflow: &str) -> ApiResult<()> {
152 if let Caller::Client(c) = caller {
153 if !c.can_read(workflow) {
154 return Err(forbidden(format!(
155 "client `{}` may not read the definition of `{workflow}`",
156 c.id
157 )));
158 }
159 }
160 Ok(())
161}
162
163fn allow_execution(state: &Shared, caller: &Caller, action: &str, id: &str) -> ApiResult<()> {
165 if matches!(caller, Caller::Client(_)) {
166 let workflow = engine(state).execution(id)?.workflow;
167 allow(caller, action, &workflow)?;
168 }
169 Ok(())
170}
171
172pub(crate) fn identify(
185 state: &AppState,
186 bearer: Option<&str>,
187 cert_der: Option<&[u8]>,
188) -> Result<Option<Caller>, ApiError> {
189 let secured = state.auth_token.is_some() || !state.identities.is_empty();
190 if !secured {
191 return Ok(Some(Caller::Open));
192 }
193 if let Some(t) = bearer {
194 if state
195 .auth_token
196 .as_ref()
197 .is_some_and(|a| constant_time_eq(t.as_bytes(), a.as_bytes()))
198 {
199 return Ok(Some(Caller::Admin));
200 }
201 if let Some(w) = state.identities.by_token(t) {
202 return Ok(Some(Caller::Worker(Arc::new(w.clone()))));
203 }
204 if let Some(c) = state.identities.client_by_token(t) {
205 return Ok(Some(Caller::Client(Arc::new(c.clone()))));
206 }
207 if state
208 .worker_token
209 .as_ref()
210 .is_some_and(|w| constant_time_eq(t.as_bytes(), w.as_bytes()))
211 {
212 return Ok(Some(Caller::SharedWorker));
213 }
214 }
215 if let Some(der) = cert_der {
216 if let Some(w) = state.identities.by_certificate(der) {
217 return Ok(Some(Caller::Worker(Arc::new(w.clone()))));
218 }
219 return match state.identities.client_by_certificate(der) {
220 Some(c) => Ok(Some(Caller::Client(Arc::new(c.clone())))),
221 None => Err(forbidden(
222 "the client certificate is valid but matches no worker or client in the workers file",
223 )),
224 };
225 }
226 Ok(None)
227}
228
229async fn auth(State(state): State<Shared>, mut req: Request, next: Next) -> Response {
230 let path = req.uri().path().to_string();
231 let bearer = req
232 .headers()
233 .get("authorization")
234 .and_then(|v| v.to_str().ok())
235 .and_then(|v| v.strip_prefix("Bearer "))
236 .map(str::to_string);
237 let cert = req
238 .extensions()
239 .get::<Arc<axum_server::tls::TlsInfo>>()
240 .and_then(|info| info.peer_certificates.first().cloned());
241 let caller = match identify(&state, bearer.as_deref(), cert.as_deref()) {
242 Ok(Some(c)) => c,
243 Ok(None) if path == "/v1/health" => Caller::Open,
244 Ok(None) => {
245 return ApiError(EngineError::new(
246 401,
247 "UNAUTHENTICATED",
248 "missing or invalid credentials",
249 ))
250 .into_response()
251 }
252 Err(e) => return e.into_response(),
253 };
254 if matches!(caller, Caller::SharedWorker | Caller::Worker(_))
255 && !is_worker_path(&path)
256 && path != "/v1/health"
257 {
258 return forbidden("worker credentials may only call /v1/workers/* and /v1/tasks/*")
259 .into_response();
260 }
261 if let Caller::Client(c) = &caller {
262 let admin_only = (path == "/v1/workflows" && req.method() == axum::http::Method::POST)
263 || path.starts_with("/v1/workers")
264 || path == "/v1/metrics";
265 if is_worker_path(&path) || admin_only {
266 return forbidden(format!(
267 "client `{}` may only use the execution API (start, signal, query, read, cancel) and schedules",
268 c.id
269 ))
270 .into_response();
271 }
272 }
273 req.extensions_mut().insert(caller);
274 let mut resp = next.run(req).await;
275 resp.headers_mut()
276 .insert("x-orbit-api-version", HeaderValue::from_static(API_VERSION));
277 resp
278}
279
280pub(crate) fn authorize_session(state: &Shared, caller: &Caller, session: &str) -> ApiResult<()> {
282 if let Some(id) = caller.identity() {
283 if let Some(owner) = engine(state).session_worker_lookup(session) {
284 if owner != id.id {
285 return Err(forbidden(format!(
286 "session belongs to worker `{owner}`, not `{}`",
287 id.id
288 )));
289 }
290 }
291 }
292 Ok(())
293}
294
295pub(crate) fn authorize_token(state: &Shared, caller: &Caller, token: &str) -> ApiResult<()> {
297 if let Some(id) = caller.identity() {
298 if let Some(owner) = engine(state).token_worker_lookup(token) {
299 if owner != id.id {
300 return Err(forbidden(format!(
301 "task token was issued to worker `{owner}`, not `{}`",
302 id.id
303 )));
304 }
305 }
306 }
307 Ok(())
308}
309
310fn constant_time_eq(a: &[u8], b: &[u8]) -> bool {
311 a.len() == b.len() && a.iter().zip(b).fold(0u8, |acc, (x, y)| acc | (x ^ y)) == 0
312}
313
314async fn ready() -> Json<Value> {
315 Json(json!({ "role": "single", "ready": true }))
316}
317
318async fn health(State(state): State<Shared>) -> Json<Value> {
319 let e = engine(&state);
320 Json(json!({
321 "status": "ok",
322 "api_version": API_VERSION,
323 "runtime_version": env!("CARGO_PKG_VERSION"),
324 "ir_versions": orbit_ir::SUPPORTED_IR_VERSIONS,
325 "active_executions": e.active_count(),
326 "worker_sessions": e.session_count(),
327 }))
328}
329
330async fn metrics(State(state): State<Shared>) -> Response {
331 let e = engine(&state);
332 let body = e.metrics.render(e.active_count(), e.session_count());
333 let mut headers = HeaderMap::new();
334 headers.insert(
335 "content-type",
336 HeaderValue::from_static("text/plain; version=0.0.4"),
337 );
338 (headers, body).into_response()
339}
340
341async fn register_workflow(State(state): State<Shared>, body: String) -> ApiResult<Response> {
342 let ir = orbit_ir::WorkflowIr::from_json(&body)
343 .map_err(|e| ApiError(EngineError::new(400, "INVALID_IR", e.0)))?;
344 let (name, version, hash) = (
345 ir.name.clone(),
346 ir.version.clone(),
347 ir.definition_hash.clone(),
348 );
349 let created = engine(&state).register_definition(ir)?;
350 let status = if created {
351 StatusCode::CREATED
352 } else {
353 StatusCode::OK
354 };
355 Ok((status, Json(json!({ "name": name, "version": version, "definition_hash": hash, "created": created }))).into_response())
356}
357
358fn summary(d: &orbit_ir::WorkflowIr) -> Value {
359 json!({
360 "name": d.name,
361 "version": d.version,
362 "description": d.description,
363 "definition_hash": d.definition_hash,
364 "steps": d.steps.iter().map(|s| s.id.clone()).collect::<Vec<_>>(),
365 "tasks": d.contracts.keys().collect::<Vec<_>>(),
366 })
367}
368
369async fn list_workflows(
370 State(state): State<Shared>,
371 Extension(caller): Extension<Caller>,
372) -> Json<Value> {
373 let defs = engine(&state).definitions();
374 let visible = |name: &str| match &caller {
375 Caller::Client(c) => c.can_read(name),
376 _ => true,
377 };
378 Json(
379 json!({ "workflows": defs.iter().filter(|d| visible(&d.name)).map(|d| summary(d)).collect::<Vec<_>>() }),
380 )
381}
382
383async fn get_workflow(
384 State(state): State<Shared>,
385 Extension(caller): Extension<Caller>,
386 Path((name, version)): Path<(String, String)>,
387) -> ApiResult<Json<Value>> {
388 allow_definition(&caller, &name)?;
389 let d = engine(&state)
390 .definition(&name, Some(&version))
391 .ok_or_else(|| {
392 ApiError(EngineError::new(
393 404,
394 "WORKFLOW_NOT_FOUND",
395 "workflow not found",
396 ))
397 })?;
398 Ok(Json(serde_json::to_value(&*d).unwrap_or_default()))
399}
400
401async fn get_workflow_latest(
402 State(state): State<Shared>,
403 Extension(caller): Extension<Caller>,
404 Path(name): Path<String>,
405) -> ApiResult<Json<Value>> {
406 allow_definition(&caller, &name)?;
407 let d = engine(&state).definition(&name, None).ok_or_else(|| {
408 ApiError(EngineError::new(
409 404,
410 "WORKFLOW_NOT_FOUND",
411 "workflow not found",
412 ))
413 })?;
414 Ok(Json(serde_json::to_value(&*d).unwrap_or_default()))
415}
416
417#[derive(Deserialize)]
418struct StartRequest {
419 workflow: String,
420 #[serde(default)]
421 version: Option<String>,
422 #[serde(default)]
423 input: Value,
424 #[serde(default)]
425 execution_id: Option<String>,
426}
427
428async fn start_execution(
429 State(state): State<Shared>,
430 Extension(caller): Extension<Caller>,
431 Json(req): Json<StartRequest>,
432) -> ApiResult<Response> {
433 allow(&caller, "start", &req.workflow)?;
434 let r = engine(&state).start(
435 &req.workflow,
436 req.version.as_deref(),
437 req.input,
438 req.execution_id,
439 )?;
440 let status = if r.created {
441 StatusCode::CREATED
442 } else {
443 StatusCode::OK
444 };
445 Ok((status, Json(serde_json::to_value(r).unwrap_or_default())).into_response())
446}
447
448#[derive(Deserialize)]
449struct ListQuery {
450 status: Option<String>,
451 workflow: Option<String>,
452 limit: Option<u32>,
453}
454
455async fn list_executions(
456 State(state): State<Shared>,
457 Extension(caller): Extension<Caller>,
458 Query(q): Query<ListQuery>,
459) -> ApiResult<Json<Value>> {
460 let limit = q.limit.unwrap_or(50).clamp(1, 1000);
461 let fetch = if matches!(caller, Caller::Client(_)) {
464 1000
465 } else {
466 limit
467 };
468 let mut list = engine(&state).list_executions(
469 q.status.as_deref().map(str::to_uppercase).as_deref(),
470 q.workflow.as_deref(),
471 fetch,
472 )?;
473 if let Caller::Client(c) = &caller {
474 list.retain(|e| c.can_read(&e.workflow));
475 list.truncate(limit as usize);
476 }
477 Ok(Json(json!({ "executions": list })))
478}
479
480async fn get_execution(
481 State(state): State<Shared>,
482 Extension(caller): Extension<Caller>,
483 Path(id): Path<String>,
484) -> ApiResult<Json<Value>> {
485 allow_execution(&state, &caller, "read", &id)?;
486 let s = engine(&state).execution(&id)?;
487 Ok(Json(serde_json::to_value(s).unwrap_or_default()))
488}
489
490#[derive(Deserialize)]
491struct HistoryQuery {
492 run: Option<u32>,
493}
494
495async fn get_history(
498 State(state): State<Shared>,
499 Extension(caller): Extension<Caller>,
500 Path(id): Path<String>,
501 Query(q): Query<HistoryQuery>,
502) -> ApiResult<Json<Value>> {
503 allow_execution(&state, &caller, "read", &id)?;
504 let e = engine(&state);
505 let h = match q.run {
506 Some(run) => e.run_history(&id, run)?,
507 None => e.history(&id)?,
508 };
509 let run = crate::model::ExecutionState::replay(&id, &h[..h.len().min(1)]).run;
510 let mut body = json!({ "execution_id": id, "events": h });
511 if run > 1 || q.run.is_some() {
512 body["run"] = json!(run);
513 }
514 Ok(Json(body))
515}
516
517fn allow_schedule(caller: &Caller, workflow: &str) -> ApiResult<()> {
521 if let Caller::Client(c) = caller {
522 if !c.allows("schedule", workflow) {
523 return Err(forbidden(format!(
524 "client `{}` may not manage schedules of `{workflow}`",
525 c.id
526 )));
527 }
528 }
529 Ok(())
530}
531
532fn allow_schedule_id(state: &Shared, caller: &Caller, id: &str) -> ApiResult<()> {
533 if matches!(caller, Caller::Client(_)) {
534 let workflow = engine(state).schedule_workflow(id)?;
535 allow_schedule(caller, &workflow)?;
536 }
537 Ok(())
538}
539
540async fn create_schedule(
541 State(state): State<Shared>,
542 Extension(caller): Extension<Caller>,
543 Json(req): Json<crate::schedule::CreateSchedule>,
544) -> ApiResult<Response> {
545 allow_schedule(&caller, req.spec["workflow"].as_str().unwrap_or(""))?;
546 let v = engine(&state).create_schedule(req)?;
547 Ok((StatusCode::CREATED, Json(v)).into_response())
548}
549
550async fn list_schedules(
551 State(state): State<Shared>,
552 Extension(caller): Extension<Caller>,
553) -> ApiResult<Json<Value>> {
554 let mut list = engine(&state).list_schedules()?;
555 if let Caller::Client(c) = &caller {
556 list.retain(|s| c.allows("schedule", s["workflow"].as_str().unwrap_or("")));
557 }
558 Ok(Json(json!({ "schedules": list })))
559}
560
561async fn get_schedule(
562 State(state): State<Shared>,
563 Extension(caller): Extension<Caller>,
564 Path(id): Path<String>,
565) -> ApiResult<Json<Value>> {
566 allow_schedule_id(&state, &caller, &id)?;
567 Ok(Json(engine(&state).get_schedule(&id)?))
568}
569
570async fn update_schedule(
571 State(state): State<Shared>,
572 Extension(caller): Extension<Caller>,
573 Path(id): Path<String>,
574 Json(patch): Json<Value>,
575) -> ApiResult<Json<Value>> {
576 allow_schedule_id(&state, &caller, &id)?;
577 if let Some(w) = patch.get("workflow").and_then(Value::as_str) {
578 allow_schedule(&caller, w)?;
579 }
580 Ok(Json(engine(&state).update_schedule(&id, patch)?))
581}
582
583async fn delete_schedule(
584 State(state): State<Shared>,
585 Extension(caller): Extension<Caller>,
586 Path(id): Path<String>,
587) -> ApiResult<Json<Value>> {
588 allow_schedule_id(&state, &caller, &id)?;
589 engine(&state).delete_schedule(&id)?;
590 Ok(Json(json!({ "schedule_id": id, "deleted": true })))
591}
592
593async fn pause_schedule(
594 State(state): State<Shared>,
595 Extension(caller): Extension<Caller>,
596 Path(id): Path<String>,
597) -> ApiResult<Json<Value>> {
598 allow_schedule_id(&state, &caller, &id)?;
599 Ok(Json(engine(&state).set_schedule_paused(&id, true)?))
600}
601
602async fn resume_schedule(
603 State(state): State<Shared>,
604 Extension(caller): Extension<Caller>,
605 Path(id): Path<String>,
606) -> ApiResult<Json<Value>> {
607 allow_schedule_id(&state, &caller, &id)?;
608 Ok(Json(engine(&state).set_schedule_paused(&id, false)?))
609}
610
611async fn trigger_schedule(
614 State(state): State<Shared>,
615 Extension(caller): Extension<Caller>,
616 Path(id): Path<String>,
617) -> ApiResult<Json<Value>> {
618 allow_schedule_id(&state, &caller, &id)?;
619 let t = engine(&state).trigger_schedule(&id)?;
620 Ok(Json(serde_json::to_value(t).unwrap_or_default()))
621}
622
623#[derive(Deserialize, Default)]
624struct CancelRequest {
625 #[serde(default)]
626 reason: Option<String>,
627}
628
629async fn cancel_execution(
630 State(state): State<Shared>,
631 Extension(caller): Extension<Caller>,
632 Path(id): Path<String>,
633 body: Option<Json<CancelRequest>>,
634) -> ApiResult<Json<Value>> {
635 allow_execution(&state, &caller, "cancel", &id)?;
636 let reason = body
637 .and_then(|b| b.0.reason)
638 .unwrap_or_else(|| "requested via API".into());
639 let s = engine(&state).cancel(&id, &reason)?;
640 Ok(Json(serde_json::to_value(s).unwrap_or_default()))
641}
642
643async fn send_signal(
646 State(state): State<Shared>,
647 Extension(caller): Extension<Caller>,
648 Path((id, name)): Path<(String, String)>,
649 body: Option<Json<Value>>,
650) -> ApiResult<Json<Value>> {
651 allow_execution(&state, &caller, "signal", &id)?;
652 let payload = body.map(|b| b.0).unwrap_or_else(|| json!({}));
653 let seq = engine(&state).signal(&id, &name, payload)?;
654 Ok(Json(json!({
655 "execution_id": id,
656 "signal": name,
657 "accepted": true,
658 "seq": seq,
659 })))
660}
661
662async fn run_query(
664 State(state): State<Shared>,
665 Extension(caller): Extension<Caller>,
666 Path((id, name)): Path<(String, String)>,
667) -> ApiResult<Json<Value>> {
668 allow_execution(&state, &caller, "query", &id)?;
669 let value = engine(&state).query(&id, &name)?;
670 Ok(Json(
671 json!({ "execution_id": id, "query": name, "value": value }),
672 ))
673}
674
675#[derive(Deserialize)]
676struct WaitQuery {
677 timeout_ms: Option<u64>,
678}
679
680async fn wait_execution(
682 State(state): State<Shared>,
683 Extension(caller): Extension<Caller>,
684 Path(id): Path<String>,
685 Query(q): Query<WaitQuery>,
686) -> ApiResult<Json<Value>> {
687 allow_execution(&state, &caller, "read", &id)?;
688 let deadline = tokio::time::Instant::now()
689 + Duration::from_millis(q.timeout_ms.unwrap_or(30_000).min(300_000));
690 loop {
691 let notified = state.changed.notified();
692 tokio::pin!(notified);
693 notified.as_mut().enable();
694 let s = engine(&state).execution(&id)?;
695 if s.status.is_terminal() || tokio::time::Instant::now() >= deadline {
696 return Ok(Json(serde_json::to_value(s).unwrap_or_default()));
697 }
698 let _ = tokio::time::timeout_at(deadline, notified).await;
699 }
700}
701
702async fn list_workers(State(state): State<Shared>) -> ApiResult<Json<Value>> {
703 Ok(Json(json!({ "workers": engine(&state).workers()? })))
704}
705
706async fn register_worker(
707 State(state): State<Shared>,
708 Extension(caller): Extension<Caller>,
709 Json(req): Json<RegisterWorker>,
710) -> ApiResult<Json<Value>> {
711 if let Some(id) = caller.identity() {
712 if req.worker_id != id.id {
713 return Err(forbidden(format!(
714 "this credential belongs to worker `{}`, not `{}`",
715 id.id, req.worker_id
716 )));
717 }
718 if let Some(rt) = &id.runtime {
719 if !rt.eq_ignore_ascii_case(&req.runtime) {
720 return Err(forbidden(format!(
721 "worker `{}` must announce runtime `{rt}`",
722 id.id
723 )));
724 }
725 }
726 let denied: Vec<&str> = req
727 .contracts
728 .iter()
729 .map(|c| c.task_name.as_str())
730 .filter(|t| !id.may_serve(t))
731 .collect();
732 if !denied.is_empty() {
733 return Err(forbidden(format!(
734 "worker `{}` may not serve: {} (allowed: {})",
735 id.id,
736 denied.join(", "),
737 id.tasks.join(", ")
738 )));
739 }
740 }
741 Ok(Json(
742 serde_json::to_value(engine(&state).register_worker(req)?).unwrap_or_default(),
743 ))
744}
745
746#[derive(Deserialize)]
747struct SessionRequest {
748 worker_session_id: String,
749 #[serde(default)]
750 wait_ms: Option<u64>,
751}
752
753async fn heartbeat(
754 State(state): State<Shared>,
755 Extension(caller): Extension<Caller>,
756 Json(req): Json<SessionRequest>,
757) -> ApiResult<Json<Value>> {
758 authorize_session(&state, &caller, &req.worker_session_id)?;
759 Ok(Json(
760 serde_json::to_value(engine(&state).heartbeat(&req.worker_session_id)?).unwrap_or_default(),
761 ))
762}
763
764async fn unregister(
765 State(state): State<Shared>,
766 Extension(caller): Extension<Caller>,
767 Json(req): Json<SessionRequest>,
768) -> ApiResult<Json<Value>> {
769 authorize_session(&state, &caller, &req.worker_session_id)?;
770 engine(&state).unregister_worker(&req.worker_session_id)?;
771 Ok(Json(json!({ "accepted": true })))
772}
773
774async fn poll(
776 State(state): State<Shared>,
777 Extension(caller): Extension<Caller>,
778 Json(req): Json<SessionRequest>,
779) -> ApiResult<Response> {
780 authorize_session(&state, &caller, &req.worker_session_id)?;
781 let wait = Duration::from_millis(req.wait_ms.unwrap_or(20_000).min(60_000));
782 let deadline = tokio::time::Instant::now() + wait;
783 loop {
784 let notified = state.changed.notified();
785 tokio::pin!(notified);
786 notified.as_mut().enable();
787 let task: Option<TaskAssignment> = crate::shared::poll(&state, &req.worker_session_id)?;
788 if let Some(t) = task {
789 return Ok(Json(serde_json::to_value(t).unwrap_or_default()).into_response());
790 }
791 if tokio::time::Instant::now() >= deadline {
792 return Ok(StatusCode::NO_CONTENT.into_response());
793 }
794 let _ = tokio::time::timeout_at(deadline, notified).await;
795 }
796}
797
798#[derive(Deserialize)]
799struct CompleteRequest {
800 token: String,
801 #[serde(default)]
802 output: Value,
803}
804
805async fn complete_task(
806 State(state): State<Shared>,
807 Extension(caller): Extension<Caller>,
808 Json(req): Json<CompleteRequest>,
809) -> ApiResult<Json<Value>> {
810 if req.token.is_empty() {
811 return Err(bad_request("token is required"));
812 }
813 authorize_token(&state, &caller, &req.token)?;
814 Ok(Json(
815 serde_json::to_value(engine(&state).complete(&req.token, req.output)?).unwrap_or_default(),
816 ))
817}
818
819#[derive(Deserialize)]
820struct FailRequest {
821 token: String,
822 error_code: String,
823 #[serde(default)]
824 error_message: String,
825 #[serde(default = "default_true")]
826 retryable: bool,
827 #[serde(default)]
828 details: Option<Value>,
829}
830
831fn default_true() -> bool {
832 true
833}
834
835async fn fail_task(
836 State(state): State<Shared>,
837 Extension(caller): Extension<Caller>,
838 Json(req): Json<FailRequest>,
839) -> ApiResult<Json<Value>> {
840 if req.error_code.is_empty() {
841 return Err(bad_request("error_code is required"));
842 }
843 authorize_token(&state, &caller, &req.token)?;
844 let error = TaskError {
845 code: req.error_code,
846 message: req.error_message,
847 retryable: req.retryable,
848 details: req.details,
849 };
850 Ok(Json(
851 serde_json::to_value(engine(&state).fail(&req.token, error)?).unwrap_or_default(),
852 ))
853}
854
855pub async fn run_ticker(state: Shared, interval: Duration) {
857 let mut ticker = tokio::time::interval(interval);
858 ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
859 loop {
860 ticker.tick().await;
861 let jobs = match engine(&state).tick() {
862 Ok(j) => j,
863 Err(e) => {
864 tracing::error!(error = %e, "tick failed");
865 continue;
866 }
867 };
868 for job in jobs {
869 let st = state.clone();
870 tokio::spawn(async move {
871 let input = job.input.clone();
872 let timeout = job.timeout_ms;
873 let result = tokio::task::spawn_blocking(move || {
874 crate::http_adapter::execute(&input, timeout)
875 })
876 .await
877 .unwrap_or_else(|e| Err(TaskError::new("INTERNAL", e.to_string(), true)));
878 let mut e = engine(&st);
879 let r = match result {
880 Ok(out) => e.complete(&job.token, out),
881 Err(err) => e.fail(&job.token, err),
882 };
883 if let Err(err) = r {
884 tracing::error!(error = %err, "recording built-in task result failed");
885 }
886 });
887 }
888 }
889}
890
891pub struct TlsConfig {
893 pub cert: std::path::PathBuf,
894 pub key: std::path::PathBuf,
895 pub client_ca: Option<std::path::PathBuf>,
899}
900
901fn io_err(msg: String) -> std::io::Error {
902 std::io::Error::new(std::io::ErrorKind::InvalidInput, msg)
903}
904
905fn read_pem_certs(
906 path: &std::path::Path,
907) -> std::io::Result<Vec<rustls::pki_types::CertificateDer<'static>>> {
908 let pem =
909 std::fs::read(path).map_err(|e| io_err(format!("cannot read {}: {e}", path.display())))?;
910 let certs = rustls_pemfile::certs(&mut pem.as_slice())
911 .collect::<Result<Vec<_>, _>>()
912 .map_err(|e| io_err(format!("{}: {e}", path.display())))?;
913 if certs.is_empty() {
914 return Err(io_err(format!(
915 "{} contains no certificates",
916 path.display()
917 )));
918 }
919 Ok(certs)
920}
921
922fn rustls_config(tls: &TlsConfig) -> std::io::Result<rustls::ServerConfig> {
924 let chain = read_pem_certs(&tls.cert)?;
925 let key_pem = std::fs::read(&tls.key)
926 .map_err(|e| io_err(format!("cannot read {}: {e}", tls.key.display())))?;
927 let key = rustls_pemfile::private_key(&mut key_pem.as_slice())
928 .map_err(|e| io_err(format!("{}: {e}", tls.key.display())))?
929 .ok_or_else(|| io_err(format!("{} contains no private key", tls.key.display())))?;
930 let provider = Arc::new(rustls::crypto::aws_lc_rs::default_provider());
931 let builder = rustls::ServerConfig::builder_with_provider(provider.clone())
932 .with_safe_default_protocol_versions()
933 .map_err(|e| io_err(e.to_string()))?;
934 let builder = match &tls.client_ca {
935 Some(ca) => {
936 let mut roots = rustls::RootCertStore::empty();
937 for c in read_pem_certs(ca)? {
938 roots
939 .add(c)
940 .map_err(|e| io_err(format!("{}: {e}", ca.display())))?;
941 }
942 let verifier = rustls::server::WebPkiClientVerifier::builder_with_provider(
943 Arc::new(roots),
944 provider,
945 )
946 .allow_unauthenticated()
947 .build()
948 .map_err(|e| io_err(format!("client CA: {e}")))?;
949 builder.with_client_cert_verifier(verifier)
950 }
951 None => builder.with_no_client_auth(),
952 };
953 let mut config = builder
954 .with_single_cert(chain, key)
955 .map_err(|e| io_err(format!("TLS certificate/key: {e}")))?;
956 config.alpn_protocols = vec![b"h2".to_vec(), b"http/1.1".to_vec()];
957 Ok(config)
958}
959
960pub struct ServerConfig {
961 pub addr: SocketAddr,
962 pub auth_token: Option<String>,
963 pub worker_token: Option<String>,
964 pub identities: Identities,
965 pub tls: Option<TlsConfig>,
966 pub grpc_addr: Option<SocketAddr>,
968 pub listener: Option<std::net::TcpListener>,
971 pub grpc_listener: Option<std::net::TcpListener>,
972 pub body_limit: usize,
973 pub tick_interval: Duration,
974 pub ui: bool,
976}
977
978pub async fn serve(
979 engine: Engine,
980 config: ServerConfig,
981 shutdown: impl std::future::Future<Output = ()> + Send + 'static,
982) -> std::io::Result<()> {
983 serve_node(engine, config, None, shutdown).await
984}
985
986pub async fn serve_node(
990 engine: Engine,
991 config: ServerConfig,
992 cluster: Option<Arc<crate::cluster::Node>>,
993 shutdown: impl std::future::Future<Output = ()> + Send + 'static,
994) -> std::io::Result<()> {
995 let _ = rustls::crypto::aws_lc_rs::default_provider().install_default();
998 let changed = engine.changed.clone();
999 let shared = engine.shared_context();
1001 let state = Arc::new(AppState {
1002 engine: Mutex::new(engine),
1003 changed,
1004 auth_token: config.auth_token,
1005 worker_token: config.worker_token,
1006 identities: config.identities,
1007 });
1008 let app = router(state.clone(), config.body_limit);
1009 let app = match &shared {
1010 Some(ctx) => crate::shared::layer(app, ctx.clone(), state.clone()),
1011 None => crate::cluster::layer(app, cluster.clone(), state.clone()),
1012 };
1013 let app = if config.ui {
1016 app.merge(crate::ui::router())
1017 } else {
1018 app
1019 };
1020 if let Some(node) = &cluster {
1021 crate::cluster::start(node.clone(), state.clone())?;
1022 }
1023 if let Some(ctx) = &shared {
1024 crate::shared::start(ctx.clone(), state.clone())?;
1025 }
1026 let ticker = tokio::spawn(run_ticker(state.clone(), config.tick_interval));
1027 let (stop_tx, stop_rx) = tokio::sync::watch::channel(false);
1029 tokio::spawn(async move {
1030 shutdown.await;
1031 let _ = stop_tx.send(true);
1032 });
1033 let stopped = |mut rx: tokio::sync::watch::Receiver<bool>| async move {
1034 let _ = rx.wait_for(|v| *v).await;
1035 };
1036 let bind = |addr: SocketAddr| {
1037 std::net::TcpListener::bind(addr)
1038 .map_err(|e| std::io::Error::new(e.kind(), format!("cannot listen on {addr}: {e}")))
1039 };
1040 let grpc_listener = match (config.grpc_listener, config.grpc_addr) {
1041 (Some(l), _) => Some(l),
1042 (None, Some(addr)) => Some(bind(addr)?),
1043 (None, None) => None,
1044 };
1045 let http_listener = match config.listener {
1046 Some(l) => l,
1047 None => bind(config.addr)?,
1048 };
1049 http_listener.set_nonblocking(true)?;
1050 let http_addr = http_listener.local_addr()?;
1051 let grpc = match grpc_listener {
1052 Some(std_listener) => {
1053 std_listener.set_nonblocking(true)?;
1054 let addr = std_listener.local_addr()?;
1055 let listener = tokio::net::TcpListener::from_std(std_listener)?;
1056 let incoming =
1057 tonic::transport::server::TcpIncoming::from_listener(listener, true, None)
1058 .map_err(|e| io_err(format!("gRPC listener: {e}")))?;
1059 let mut builder = tonic::transport::Server::builder();
1060 if let Some(tls) = &config.tls {
1061 let _ = rustls::crypto::aws_lc_rs::default_provider().install_default();
1062 let read = |p: &std::path::Path| {
1063 std::fs::read(p)
1064 .map_err(|e| io_err(format!("cannot read {}: {e}", p.display())))
1065 };
1066 let mut t = tonic::transport::ServerTlsConfig::new().identity(
1067 tonic::transport::Identity::from_pem(read(&tls.cert)?, read(&tls.key)?),
1068 );
1069 if let Some(ca) = &tls.client_ca {
1070 t = t
1071 .client_ca_root(tonic::transport::Certificate::from_pem(read(ca)?))
1072 .client_auth_optional(true);
1073 }
1074 builder = builder
1075 .tls_config(t)
1076 .map_err(|e| io_err(format!("gRPC TLS: {e}")))?;
1077 }
1078 let service = crate::grpc::GrpcWorkers::new(state.clone());
1079 let rx = stop_rx.clone();
1080 tracing::info!(addr = %addr, tls = config.tls.is_some(), "gRPC worker protocol listening");
1081 let router = match &cluster {
1082 Some(node) => {
1083 builder.add_service(tonic::service::interceptor::InterceptedService::new(
1084 service,
1085 crate::cluster::grpc_gate(node.clone()),
1086 ))
1087 }
1088 None => builder.add_service(service),
1089 };
1090 Some(tokio::spawn(async move {
1091 if let Err(e) = router
1092 .serve_with_incoming_shutdown(incoming, stopped(rx))
1093 .await
1094 {
1095 tracing::error!(error = %e, "gRPC server failed");
1096 }
1097 }))
1098 }
1099 None => None,
1100 };
1101 let shutdown = stopped(stop_rx);
1102 let result = match config.tls {
1103 None => {
1104 let listener = tokio::net::TcpListener::from_std(http_listener)?;
1105 tracing::info!(addr = %http_addr, tls = false, ui = config.ui, "ORBIT runtime listening");
1106 axum::serve(listener, app)
1107 .with_graceful_shutdown(shutdown)
1108 .await
1109 }
1110 Some(tls) => {
1111 let _ = rustls::crypto::aws_lc_rs::default_provider().install_default();
1114 let rustls =
1115 axum_server::tls_rustls::RustlsConfig::from_config(Arc::new(rustls_config(&tls)?));
1116 let handle = axum_server::Handle::new();
1117 let h = handle.clone();
1118 tokio::spawn(async move {
1119 shutdown.await;
1120 h.graceful_shutdown(Some(Duration::from_secs(10)));
1121 });
1122 tracing::info!(addr = %http_addr, tls = true, mtls = tls.client_ca.is_some(), ui = config.ui, "ORBIT runtime listening");
1123 axum_server::from_tcp_rustls(http_listener, rustls)?
1124 .handle(handle)
1125 .serve(app.into_make_service())
1126 .await
1127 }
1128 };
1129 ticker.abort();
1130 if let Some(ctx) = &shared {
1131 ctx.stop();
1132 }
1133 if let Some(g) = grpc {
1134 let _ = g.await;
1135 }
1136 result
1137}