1use std::collections::{BTreeMap, BTreeSet, HashMap};
10use std::convert::Infallible;
11use std::io;
12use std::net::SocketAddr;
13use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering};
14use std::sync::{Arc, Mutex};
15use std::time::Duration;
16
17use http_body_util::{BodyExt, Full, Limited};
18use hyper::body::{Bytes, Incoming};
19use hyper::header::{AUTHORIZATION, CONTENT_TYPE};
20use hyper::service::service_fn;
21use hyper::{Method, Request, Response, StatusCode};
22use hyper_util::rt::TokioIo;
23use rustev_contract::canonical::canonical_value_bytes;
24use rustev_contract::descriptor::BackendDescriptor;
25use rustev_contract::ids::DescriptorId;
26use rustev_contract::remote::{
27 Batching, CancelItem, CancelRequest, CancelResponse, CancelSupport, DescribeRequest,
28 DescribeResponse, Disclosure, HostedBackend, InferRequest, InferResponse, ItemAnswer,
29 ItemResult, MAX_ITEMS, MAX_TEXT_BYTES, PROTOCOL_MAJOR, RemoteCancelAnswer, RemoteCode,
30 RemoteTerms, ServedIdentity, ServedTerms, cut_to,
31};
32use rustev_contract::run::{Charge, CostBound, CostModel};
33use rustev_contract::time::DurationMs;
34use rustev_contract::{Document, Identified, schema};
35use rustev_core::seams::{
36 AdapterFailure, AttemptCall, AttemptReport, CallContext, CancelAck, CancelSignal,
37 DecisionBackend, RemoteEnd,
38};
39use tokio::net::TcpListener;
40use tokio::sync::watch;
41use tokio::task::JoinHandle;
42use tokio::time::Instant;
43
44use crate::cost::sum_charges;
45use crate::http::Secret;
46use crate::taxonomy::{code_of_detail, failure_class, is_local_only};
47use crate::wire::digest_bytes;
48
49pub const DESCRIBE_PATH: &str = "/rustev/remote/1/describe";
51pub const INFER_PATH: &str = "/rustev/remote/1/infer";
52pub const CANCEL_PATH: &str = "/rustev/remote/1/cancel";
53
54#[derive(Clone)]
56pub struct ServerConfig {
57 pub bearer: Option<Secret>,
59 pub idempotency_window_ms: u64,
60 pub max_request_bytes: u64,
61 pub max_response_bytes: u64,
62 pub cancellation: CancelSupport,
63 pub batching: Batching,
65 pub refusals_charged: bool,
67 pub max_in_flight: usize,
69 pub cancel_wait_ms: u64,
71 pub max_remembered: usize,
74 pub versions: Vec<u32>,
78}
79
80impl Default for ServerConfig {
81 fn default() -> Self {
82 ServerConfig {
83 bearer: None,
84 idempotency_window_ms: 10 * 60 * 1000,
85 max_request_bytes: 1024 * 1024,
86 max_response_bytes: 1024 * 1024,
87 cancellation: CancelSupport::Confirmed,
88 batching: Batching::None,
89 refusals_charged: false,
90 max_in_flight: 64,
91 cancel_wait_ms: 1_000,
92 max_remembered: 10_000,
93 versions: vec![PROTOCOL_MAJOR],
94 }
95 }
96}
97
98#[derive(Clone)]
100pub struct Hosted {
101 pub backend: Arc<dyn DecisionBackend>,
102 pub served: ServedTerms,
105 pub cost: Option<CostBound>,
107 pub privacy: Vec<String>,
108}
109
110impl Hosted {
111 pub fn new(backend: Arc<dyn DecisionBackend>) -> Self {
112 let identity = backend.descriptor().artifact.as_str().to_string();
113 Hosted {
114 backend,
115 served: ServedTerms {
116 disclosure: Disclosure::Pinned,
117 identity: Some(identity),
118 },
119 cost: None,
120 privacy: vec![],
121 }
122 }
123
124 fn cost(&self) -> CostBound {
125 if let Some(c) = self.cost {
126 return c;
127 }
128 let bound = self.backend.cost_bound(b"{}");
129 match (self.backend.cost_model(), bound) {
130 (CostModel::Bounded, b @ CostBound::Bounded { .. }) => b,
131 (CostModel::Unknown, _) => CostBound::Unknown,
132 (_, CostBound::Bounded { max_units }) => CostBound::Estimated { units: max_units },
133 (_, b) => b,
134 }
135 }
136}
137
138#[derive(Debug, Clone)]
140struct Recorded {
141 result: ItemResult,
142 charge: Charge,
143 ack: CancelAck,
144 remote: RemoteEnd,
145}
146
147struct Entry {
148 backend_id: String,
149 content: String,
150 signal: CancelSignal,
151 cancel_requested: Arc<AtomicBool>,
152 done: watch::Receiver<Option<Recorded>>,
153 finished_at: Option<Instant>,
154}
155
156struct State {
157 config: ServerConfig,
158 hosted: Mutex<BTreeMap<String, Hosted>>,
159 attempts: Mutex<HashMap<String, Entry>>,
160 in_flight: AtomicUsize,
161 dispatched: AtomicU64,
162}
163
164#[derive(Clone)]
166pub struct RemoteServer {
167 state: Arc<State>,
168}
169
170pub struct ServerHandle {
172 pub addr: SocketAddr,
173 task: JoinHandle<()>,
174}
175
176impl Drop for ServerHandle {
177 fn drop(&mut self) {
178 self.task.abort();
179 }
180}
181
182impl ServerHandle {
183 pub fn url(&self) -> String {
185 format!("http://{}", self.addr)
186 }
187}
188
189fn lock<T>(m: &Mutex<T>) -> std::sync::MutexGuard<'_, T> {
190 m.lock().unwrap_or_else(|e| e.into_inner())
191}
192
193impl RemoteServer {
194 pub fn new(config: ServerConfig) -> Self {
195 RemoteServer {
196 state: Arc::new(State {
197 config,
198 hosted: Mutex::new(BTreeMap::new()),
199 attempts: Mutex::new(HashMap::new()),
200 in_flight: AtomicUsize::new(0),
201 dispatched: AtomicU64::new(0),
202 }),
203 }
204 }
205
206 pub fn host(&self, hosted: Hosted) {
210 let id = hosted.backend.descriptor().backend_id.clone();
211 lock(&self.state.hosted).insert(id, hosted);
212 }
213
214 pub fn dispatched(&self) -> u64 {
216 self.state.dispatched.load(Ordering::SeqCst)
217 }
218
219 pub fn serve(
222 &self,
223 listener: TcpListener,
224 tls: Option<Arc<rustls::ServerConfig>>,
225 ) -> io::Result<ServerHandle> {
226 let addr = listener.local_addr()?;
227 if tls.is_none() && !addr.ip().is_loopback() {
228 return Err(io::Error::new(
229 io::ErrorKind::InvalidInput,
230 "TLS is required except for loopback",
231 ));
232 }
233 let acceptor = tls.map(tokio_rustls::TlsAcceptor::from);
234 let state = self.state.clone();
235 let task = tokio::spawn(async move {
236 loop {
237 let Ok((stream, _)) = listener.accept().await else {
238 continue;
239 };
240 let _ = stream.set_nodelay(true);
241 let state = state.clone();
242 let acceptor = acceptor.clone();
243 tokio::spawn(async move {
244 match acceptor {
245 None => serve_conn(state, stream).await,
246 Some(a) => {
247 if let Ok(tls) = a.accept(stream).await {
248 serve_conn(state, tls).await
249 }
250 }
251 }
252 });
253 }
254 });
255 Ok(ServerHandle { addr, task })
256 }
257
258 pub async fn serve_loopback(&self) -> io::Result<ServerHandle> {
261 let listener = TcpListener::bind(("127.0.0.1", 0)).await?;
262 self.serve(listener, None)
263 }
264}
265
266async fn serve_conn<S>(state: Arc<State>, stream: S)
267where
268 S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin + Send + 'static,
269{
270 let svc = service_fn(move |req| {
271 let state = state.clone();
272 async move { Ok::<_, Infallible>(handle(state, req).await) }
273 });
274 let _ = hyper::server::conn::http1::Builder::new()
275 .serve_connection(TokioIo::new(stream), svc)
276 .await;
277}
278
279fn constant_time_eq(a: &[u8], b: &[u8]) -> bool {
282 if a.len() != b.len() {
283 return false;
284 }
285 a.iter().zip(b).fold(0u8, |acc, (x, y)| acc | (x ^ y)) == 0
286}
287
288fn reply(status: StatusCode, body: Vec<u8>) -> Response<Full<Bytes>> {
289 let mut r = Response::new(Full::new(Bytes::from(body)));
290 *r.status_mut() = status;
291 r.headers_mut()
292 .insert(CONTENT_TYPE, "application/json".parse().expect("static"));
293 r
294}
295
296fn refuse(status: StatusCode, why: &str) -> Response<Full<Bytes>> {
297 let body = serde_json::json!({ "error": cut_to(why, MAX_TEXT_BYTES) });
298 reply(status, serde_json::to_vec(&body).unwrap_or_default())
299}
300
301struct InFlight<'a>(&'a AtomicUsize);
302
303impl Drop for InFlight<'_> {
304 fn drop(&mut self) {
305 self.0.fetch_sub(1, Ordering::SeqCst);
306 }
307}
308
309async fn handle(state: Arc<State>, req: Request<Incoming>) -> Response<Full<Bytes>> {
310 if let Some(b) = &state.config.bearer {
311 let expected = format!("Bearer {}", b.expose());
312 let ok = req
313 .headers()
314 .get(AUTHORIZATION)
315 .is_some_and(|v| constant_time_eq(v.as_bytes(), expected.as_bytes()));
316 if !ok {
317 return refuse(StatusCode::UNAUTHORIZED, "credentials refused");
318 }
319 }
320 if req.method() != Method::POST {
321 return refuse(StatusCode::METHOD_NOT_ALLOWED, "POST only");
322 }
323 let n = state.in_flight.fetch_add(1, Ordering::SeqCst) + 1;
324 let _guard = InFlight(&state.in_flight);
325 if n > state.config.max_in_flight {
326 return refuse(StatusCode::SERVICE_UNAVAILABLE, "overloaded");
327 }
328 let path = req.uri().path().to_string();
329 let max = state.config.max_request_bytes as usize;
330 let body = match Limited::new(req.into_body(), max).collect().await {
331 Ok(b) => b.to_bytes(),
332 Err(_) => return refuse(StatusCode::PAYLOAD_TOO_LARGE, "request too large"),
333 };
334 let r = match path.as_str() {
335 DESCRIBE_PATH => describe(&state, &body),
336 INFER_PATH => infer(&state, &body).await,
337 CANCEL_PATH => cancel(&state, &body).await,
338 _ => return refuse(StatusCode::NOT_FOUND, "no such document"),
339 };
340 match r {
341 Ok(bytes) if bytes.len() as u64 <= state.config.max_response_bytes => {
342 reply(StatusCode::OK, bytes)
343 }
344 Ok(_) => refuse(StatusCode::INTERNAL_SERVER_ERROR, "response too large"),
345 Err((status, why)) => refuse(status, &why),
346 }
347}
348
349type Refusal = (StatusCode, String);
350
351fn bad(why: impl Into<String>) -> Refusal {
352 (StatusCode::BAD_REQUEST, why.into())
353}
354
355fn terms(state: &State, h: &Hosted) -> RemoteTerms {
356 RemoteTerms {
357 cost: h.cost(),
358 refusals_charged: state.config.refusals_charged,
359 cancellation: state.config.cancellation,
360 served: h.served.clone(),
361 batching: state.config.batching,
362 max_request_bytes: state.config.max_request_bytes,
363 max_response_bytes: state.config.max_response_bytes,
364 idempotency_window_ms: state.config.idempotency_window_ms,
365 privacy: h.privacy.clone(),
366 }
367}
368
369fn describe(state: &State, body: &[u8]) -> Result<Vec<u8>, Refusal> {
370 let req = DescribeRequest::parse(body).map_err(|e| bad(e.to_string()))?;
371 let Some(version) = req
372 .versions
373 .iter()
374 .filter(|v| state.config.versions.contains(v))
375 .max()
376 .copied()
377 else {
378 return Err(bad("no common protocol version"));
379 };
380 let hosted = lock(&state.hosted).clone();
381 let mut backends = vec![];
382 for h in hosted.values() {
383 let descriptor = h.backend.descriptor().clone();
384 let descriptor_id = descriptor
385 .id()
386 .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
387 backends.push(HostedBackend {
388 descriptor,
389 descriptor_id,
390 terms: terms(state, h),
391 });
392 }
393 DescribeResponse {
394 schema: schema::REMOTE_DESCRIBE.into(),
395 version,
396 backends,
397 }
398 .record_canonical()
399 .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))
400}
401
402fn failure(code: RemoteCode, detail: &str) -> ItemResult {
403 ItemResult::Failure {
404 code,
405 detail: cut_to(detail, MAX_TEXT_BYTES),
406 }
407}
408
409fn refused_all(req: &InferRequest, code: RemoteCode, detail: &str) -> Vec<(String, Recorded)> {
410 req.items
411 .iter()
412 .map(|i| {
413 (
414 i.attempt_id.clone(),
415 Recorded {
416 result: failure(code, detail),
417 charge: Charge::Observed { units: 0 },
418 ack: CancelAck::NotRequested,
419 remote: RemoteEnd::Finished,
420 },
421 )
422 })
423 .collect()
424}
425
426fn state_bytes(projection: &str) -> Option<Vec<u8>> {
428 let v: serde_json::Value = serde_json::from_str(projection).ok()?;
429 let s = serde_json::json!({ "instance": v.get("instance")?, "values": v.get("values")? });
430 canonical_value_bytes(&s).ok()
431}
432
433fn content_key(
434 backend_id: &str,
435 descriptor: &DescriptorId,
436 projection: &str,
437 principal: &str,
438) -> String {
439 let mut b = Vec::new();
440 for part in [
441 backend_id.as_bytes(),
442 descriptor.as_str().as_bytes(),
443 projection.as_bytes(),
444 principal.as_bytes(),
445 ] {
446 b.extend_from_slice(&(part.len() as u64).to_be_bytes());
447 b.extend_from_slice(part);
448 }
449 digest_bytes(&b)
450}
451
452async fn infer(state: &Arc<State>, body: &[u8]) -> Result<Vec<u8>, Refusal> {
453 let req = InferRequest::parse(body).map_err(|e| bad(e.to_string()))?;
454 if req.items.is_empty() {
455 return Err(bad("no items"));
456 }
457 let ids: BTreeSet<&str> = req.items.iter().map(|i| i.attempt_id.as_str()).collect();
458 if ids.len() != req.items.len() {
459 return Err(bad("an attempt id appears twice"));
460 }
461 if !req.members_fit_version() {
463 return Err(bad("principal reference not allowed or not canonical"));
464 }
465 let hosted = lock(&state.hosted).get(&req.backend_id).cloned();
466 let answers: Vec<(String, Recorded)> = match hosted {
467 _ if !state.config.versions.contains(&req.version) => {
468 refused_all(&req, RemoteCode::Capability, "unsupported protocol version")
469 }
470 None => refused_all(&req, RemoteCode::Capability, "no such backend"),
471 Some(h) => run_items(state, &req, &h).await?,
472 };
473 let hosted = lock(&state.hosted).get(&req.backend_id).cloned();
474 let (artifact, served) = match &hosted {
475 Some(h) => (
476 h.backend.descriptor().artifact.clone(),
477 match h.served.disclosure {
478 Disclosure::Unknown => ServedIdentity::unknown(),
479 d => ServedIdentity {
480 identity: h.served.identity.clone(),
481 disclosure: d,
482 },
483 },
484 ),
485 None => (req.artifact.clone(), ServedIdentity::unknown()),
486 };
487 let charge = sum_charges(&answers.iter().map(|(_, r)| r.charge).collect::<Vec<_>>());
488 let items: Vec<ItemAnswer> = answers
489 .into_iter()
490 .map(|(attempt_id, r)| ItemAnswer {
491 attempt_id,
492 result: r.result,
493 })
494 .collect();
495 let mut resp = InferResponse {
496 schema: schema::REMOTE_INFER.into(),
497 version: req.version,
498 artifact,
499 served,
500 route: vec![],
501 generation_id: None,
502 items,
503 charge,
504 extras: BTreeMap::new(),
505 };
506 if let Ok(b) = resp.record_canonical() {
507 return Ok(b);
508 }
509 for i in &mut resp.items {
511 if let ItemResult::Output(o) = &i.result
512 && rustev_contract::canonical::record_canonical_bytes(o).is_err()
513 {
514 i.result = failure(RemoteCode::RemoteError, "output has no JSON form");
515 }
516 }
517 resp.record_canonical()
518 .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))
519}
520
521enum Slot {
522 Answer(Recorded),
523 Wait(watch::Receiver<Option<Recorded>>),
524}
525
526async fn run_items(
527 state: &Arc<State>,
528 req: &InferRequest,
529 h: &Hosted,
530) -> Result<Vec<(String, Recorded)>, Refusal> {
531 let descriptor: &BackendDescriptor = h.backend.descriptor();
532 let current = descriptor
533 .id()
534 .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
535 if current != req.descriptor || descriptor.artifact != req.artifact {
536 return Ok(refused_all(
537 req,
538 RemoteCode::IdentityMismatch,
539 "the hosted descriptor or artifact differs from the negotiated one",
540 ));
541 }
542 if req.items.len() > 1
549 && descriptor.operations.iter().any(|o| {
550 o.screen.is_some()
551 && o.isolation == Some(rustev_contract::descriptor::Isolation::PerItem)
552 })
553 {
554 return Ok(refused_all(
555 req,
556 RemoteCode::Capability,
557 "the screening capability requires per-item isolation",
558 ));
559 }
560 let max_items = match state.config.batching {
561 Batching::None => 1,
562 Batching::SharedState { max_items } => (max_items as usize).min(MAX_ITEMS),
563 };
564 if req.items.len() > max_items {
565 return Ok(refused_all(req, RemoteCode::Capability, "too many items"));
566 }
567 if req.items.len() > 1 {
568 let first = state_bytes(&req.items[0].projection);
569 if first.is_none()
570 || req
571 .items
572 .iter()
573 .any(|i| state_bytes(&i.projection) != first)
574 {
575 return Err((
576 StatusCode::UNPROCESSABLE_ENTITY,
577 "batched items must share their values and instance".into(),
578 ));
579 }
580 }
581 let window = Duration::from_millis(state.config.idempotency_window_ms);
582 let mut slots = Vec::with_capacity(req.items.len());
583 {
584 let mut attempts = lock(&state.attempts);
585 let now = Instant::now();
586 attempts.retain(|_, e| e.finished_at.is_none_or(|t| now.duration_since(t) < window));
587 let new = req
588 .items
589 .iter()
590 .filter(|i| !attempts.contains_key(&i.attempt_id))
591 .count();
592 if attempts.len() + new > state.config.max_remembered {
593 return Err((StatusCode::SERVICE_UNAVAILABLE, "overloaded".into()));
594 }
595 let principal_digest = req
596 .principal
597 .as_ref()
598 .map(|p| p.evidence().digest)
599 .unwrap_or_default();
600 let principal_handle = req
601 .principal
602 .as_ref()
603 .and_then(|p| p.reference_bytes())
604 .unwrap_or_default();
605 for item in &req.items {
606 let content = content_key(
607 &req.backend_id,
608 &req.descriptor,
609 &item.projection,
610 &principal_digest,
611 );
612 if let Some(e) = attempts.get(&item.attempt_id) {
613 let recorded = e.done.borrow().clone();
614 slots.push(Slot::Answer(if e.content != content {
615 Recorded {
616 result: failure(RemoteCode::MalformedRequest, "attempt_id_conflict"),
617 charge: Charge::Observed { units: 0 },
618 ack: CancelAck::NotRequested,
619 remote: RemoteEnd::Finished,
620 }
621 } else {
622 match recorded {
623 Some(r) => r,
624 None => Recorded {
625 result: ItemResult::InProgress,
626 charge: Charge::Unknown,
627 ack: CancelAck::NotRequested,
628 remote: RemoteEnd::PossiblyContinuing,
629 },
630 }
631 }));
632 continue;
633 }
634 if item.projection.len() as u64 > descriptor.input_limit.max_bytes {
635 slots.push(Slot::Answer(Recorded {
636 result: failure(RemoteCode::Capability, "input exceeds the declared limit"),
637 charge: Charge::Observed { units: 0 },
638 ack: CancelAck::NotRequested,
639 remote: RemoteEnd::Finished,
640 }));
641 continue;
642 }
643 let (tx, rx) = watch::channel(None);
644 let signal = CancelSignal::new();
645 let cancel_requested = Arc::new(AtomicBool::new(false));
646 attempts.insert(
647 item.attempt_id.clone(),
648 Entry {
649 backend_id: req.backend_id.clone(),
650 content,
651 signal: signal.clone(),
652 cancel_requested: cancel_requested.clone(),
653 done: rx.clone(),
654 finished_at: None,
655 },
656 );
657 spawn_attempt(
658 state.clone(),
659 h.backend.clone(),
660 item.attempt_id.clone(),
661 item.projection.clone().into_bytes(),
662 req.budget_ms,
663 req.trace_id.clone(),
664 principal_handle.clone(),
665 signal,
666 cancel_requested,
667 tx,
668 );
669 slots.push(Slot::Wait(rx));
670 }
671 }
672 let mut out = Vec::with_capacity(slots.len());
673 for (item, slot) in req.items.iter().zip(slots) {
674 let r = match slot {
675 Slot::Answer(r) => r,
676 Slot::Wait(mut rx) => match rx.wait_for(Option::is_some).await {
677 Ok(r) => r.clone().expect("waited for Some"),
678 Err(_) => Recorded {
679 result: failure(RemoteCode::RemoteError, "the attempt vanished"),
680 charge: Charge::Unknown,
681 ack: CancelAck::NotRequested,
682 remote: RemoteEnd::PossiblyContinuing,
683 },
684 },
685 };
686 out.push((item.attempt_id.clone(), r));
687 }
688 Ok(out)
689}
690
691#[allow(clippy::too_many_arguments)]
692fn spawn_attempt(
693 state: Arc<State>,
694 backend: Arc<dyn DecisionBackend>,
695 attempt_id: String,
696 projection: Vec<u8>,
697 budget_ms: u64,
698 trace_id: String,
699 principal_handle: Vec<u8>,
700 signal: CancelSignal,
701 cancel_requested: Arc<AtomicBool>,
702 tx: watch::Sender<Option<Recorded>>,
703) {
704 let wait = Duration::from_millis(state.config.cancel_wait_ms);
705 let deadline_hit = Arc::new(AtomicBool::new(false));
706 let hit = deadline_hit.clone();
707 let requested = cancel_requested.clone();
708 let id = attempt_id.clone();
709 let outer_state = state.clone();
710 let inner = tokio::spawn(async move {
711 if budget_ms == 0 {
712 return Recorded {
713 result: failure(RemoteCode::RemoteDeadline, "no budget left; not started"),
714 charge: Charge::Observed { units: 0 },
715 ack: CancelAck::NotRequested,
716 remote: RemoteEnd::Finished,
717 };
718 }
719 state.dispatched.fetch_add(1, Ordering::SeqCst);
720 let cx = CallContext {
721 remaining_ms: DurationMs(budget_ms),
722 trace_id,
723 principal_handle,
727 };
728 let fut = backend.infer(AttemptCall {
729 projection: &projection,
730 attempt_id: &attempt_id,
731 cancel: signal.clone(),
732 cx: &cx,
733 });
734 tokio::pin!(fut);
735 let sleep = tokio::time::sleep(Duration::from_millis(budget_ms));
736 tokio::pin!(sleep);
737 let report = tokio::select! {
738 r = &mut fut => Some(r),
739 _ = &mut sleep => {
740 hit.store(true, Ordering::SeqCst);
741 signal.raise();
742 tokio::time::timeout(wait, &mut fut).await.ok()
743 }
744 };
745 match report {
746 Some(r) => map_report(
747 r,
748 hit.load(Ordering::SeqCst),
749 requested.load(Ordering::SeqCst),
750 ),
751 None => unconfirmed_deadline(),
752 }
753 });
754 tokio::spawn(async move {
755 let r = match inner.await {
756 Ok(r) => r,
757 Err(_) => Recorded {
758 result: failure(RemoteCode::RemoteError, "the hosted backend panicked"),
759 charge: Charge::Unknown,
760 ack: CancelAck::NotRequested,
761 remote: RemoteEnd::PossiblyContinuing,
762 },
763 };
764 let mut attempts = lock(&outer_state.attempts);
766 if let Some(e) = attempts.get_mut(&id) {
767 e.finished_at = Some(Instant::now());
768 }
769 let _ = tx.send(Some(r));
770 });
771}
772
773fn unconfirmed_deadline() -> Recorded {
778 Recorded {
779 result: failure(
780 RemoteCode::RemoteError,
781 "the budget ran out and the hosted backend did not confirm a stop",
782 ),
783 charge: Charge::Unknown,
784 ack: CancelAck::Unconfirmed,
785 remote: RemoteEnd::PossiblyContinuing,
786 }
787}
788
789fn map_report(r: AttemptReport, deadline_hit: bool, cancel_requested: bool) -> Recorded {
794 let result = match r.result {
795 Ok(o) => ItemResult::Output(o),
796 Err((AdapterFailure::Cancelled, d)) if deadline_hit && r.cancel == CancelAck::Stopped => {
798 failure(RemoteCode::RemoteDeadline, &d)
799 }
800 Err((AdapterFailure::Cancelled, _)) if deadline_hit => return unconfirmed_deadline(),
801 Err((AdapterFailure::Cancelled, d)) if cancel_requested => {
802 failure(RemoteCode::Cancelled, &d)
803 }
804 Err((AdapterFailure::Cancelled, d)) => failure(
805 RemoteCode::RemoteError,
806 &format!("reported cancelled without a request: {d}"),
807 ),
808 Err((class, d)) => match code_of_detail(&d) {
809 Some(c) if failure_class(c) == class && !is_local_only(c) => {
810 let rest = d
811 .strip_prefix(&format!("remote:{}", c.as_str()))
812 .unwrap_or("")
813 .trim_start_matches(": ");
814 failure(c, rest)
815 }
816 _ => failure(
817 match class {
818 AdapterFailure::Transient => RemoteCode::RemoteError,
819 AdapterFailure::Overloaded => RemoteCode::Overloaded,
820 _ => RemoteCode::MalformedRequest,
821 },
822 &d,
823 ),
824 },
825 };
826 Recorded {
827 result,
828 charge: r.charge,
829 ack: r.cancel,
830 remote: r.remote,
831 }
832}
833
834fn cancel_answer(support: CancelSupport, r: &Recorded) -> RemoteCancelAnswer {
835 let finished = r.remote == RemoteEnd::Finished && r.ack != CancelAck::Unconfirmed;
836 match (support, finished, r.charge) {
837 (CancelSupport::Confirmed, true, c) if c != Charge::Unknown => {
838 RemoteCancelAnswer::Stopped { charge: c }
839 }
840 _ => RemoteCancelAnswer::Unconfirmed,
841 }
842}
843
844async fn cancel(state: &Arc<State>, body: &[u8]) -> Result<Vec<u8>, Refusal> {
845 let req = CancelRequest::parse(body).map_err(|e| bad(e.to_string()))?;
846 let support = state.config.cancellation;
847 if !state.config.versions.contains(&req.version) || support == CancelSupport::None {
848 return Err((
849 StatusCode::UNPROCESSABLE_ENTITY,
850 "cancellation is not supported".into(),
851 ));
852 }
853 let mut waits = vec![];
854 {
855 let attempts = lock(&state.attempts);
856 for id in &req.attempt_ids {
857 match attempts.get(id) {
858 Some(e) if e.backend_id == req.backend_id => {
859 e.cancel_requested.store(true, Ordering::SeqCst);
860 e.signal.raise();
861 waits.push((id.clone(), Some(e.done.clone())));
862 }
863 _ => waits.push((id.clone(), None)),
864 }
865 }
866 }
867 let wait = Duration::from_millis(state.config.cancel_wait_ms);
868 let mut answers = vec![];
869 for (attempt_id, rx) in waits {
870 let answer = match rx {
871 None => RemoteCancelAnswer::UnknownAttempt,
872 Some(mut rx) => match tokio::time::timeout(wait, rx.wait_for(Option::is_some)).await {
873 Ok(Ok(r)) => cancel_answer(support, r.as_ref().expect("waited for Some")),
874 _ => RemoteCancelAnswer::Unconfirmed,
875 },
876 };
877 answers.push(CancelItem { attempt_id, answer });
878 }
879 CancelResponse {
880 schema: schema::REMOTE_CANCEL.into(),
881 version: req.version,
882 answers,
883 }
884 .record_canonical()
885 .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))
886}
887
888#[cfg(test)]
889mod tests {
890 use super::*;
891
892 #[test]
893 fn bearer_comparison_is_exact_and_length_checked() {
894 assert!(constant_time_eq(b"Bearer abc", b"Bearer abc"));
895 assert!(!constant_time_eq(b"Bearer abc", b"Bearer abd"));
896 assert!(!constant_time_eq(b"Bearer abc", b"Bearer ab"));
897 assert!(!constant_time_eq(b"", b"x"));
898 assert!(constant_time_eq(b"", b""));
899 }
900
901 fn cancelled(ack: CancelAck, charge: Charge) -> AttemptReport {
902 AttemptReport {
903 result: Err((AdapterFailure::Cancelled, "SYNTHETIC".into())),
904 charge,
905 cancel: ack,
906 remote: RemoteEnd::Finished,
907 }
908 }
909
910 #[test]
911 fn remote_deadline_only_for_a_confirmed_stop() {
912 let stopped = map_report(
913 cancelled(CancelAck::Stopped, Charge::Observed { units: 1 }),
914 true,
915 false,
916 );
917 assert!(matches!(
918 stopped.result,
919 ItemResult::Failure {
920 code: RemoteCode::RemoteDeadline,
921 ..
922 }
923 ));
924 assert_eq!(stopped.charge, Charge::Observed { units: 1 });
925 let unconfirmed = map_report(
926 cancelled(CancelAck::Unconfirmed, Charge::Observed { units: 1 }),
927 true,
928 false,
929 );
930 assert!(matches!(
931 unconfirmed.result,
932 ItemResult::Failure {
933 code: RemoteCode::RemoteError,
934 ..
935 }
936 ));
937 assert_eq!(unconfirmed.charge, Charge::Unknown);
938 }
939
940 #[test]
941 fn a_hosted_transport_code_is_never_sent_on() {
942 let r = map_report(
943 AttemptReport {
944 result: Err((
945 AdapterFailure::Transient,
946 "remote:transport_unsent: SYNTHETIC".into(),
947 )),
948 charge: Charge::Observed { units: 3 },
949 cancel: CancelAck::NotRequested,
950 remote: RemoteEnd::Finished,
951 },
952 false,
953 false,
954 );
955 assert!(matches!(
956 r.result,
957 ItemResult::Failure {
958 code: RemoteCode::RemoteError,
959 ..
960 }
961 ));
962 assert_eq!(r.charge, Charge::Observed { units: 3 });
963 }
964}