Skip to main content

lenso_service/
direct_grpc.rs

1use crate::{
2    AuthenticatedServiceContext, AuthenticatedServicePrincipal, AuthenticatedTransportBinding,
3    CallPolicyDeclaration, CallPolicyEvent, CallPolicyEvidence, CallPolicyFailure,
4    CallPolicyPermit, CallPolicyRuntime, CallPolicyTerminalOutcome, DelegatedActorContext,
5    DelegatedContextVerifier, EndpointResolver, IdentityDecisionRecorder, RetryDecision,
6    ServiceContext, ServiceContextAdmission, ServiceContextPolicy, ServiceReference, TenantContext,
7    WorkloadIdentityProvider, WorkloadIdentityVerification,
8    support_grpc_v1::{
9        GetSlaRequest, ProbeSlaRequest, SlaResponse, UpdateSlaRequest,
10        support_service_client::SupportServiceClient,
11    },
12};
13use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD};
14use prost::Message;
15use prost_types::FileDescriptorSet;
16use serde::{Deserialize, Serialize};
17use std::{
18    collections::BTreeMap,
19    sync::Arc,
20    time::{Duration, SystemTime, UNIX_EPOCH},
21};
22use tonic::{Code, Request, Status, metadata::MetadataValue, transport::Channel};
23
24const DEADLINE_METADATA: &str = "x-lenso-deadline-unix-ms";
25const IDEMPOTENCY_METADATA: &str = "idempotency-key";
26const AUTHORIZATION_METADATA: &str = "authorization";
27const DELEGATED_ACTOR_METADATA: &str = "x-lenso-delegated-actor";
28const TENANT_CONTEXT_METADATA: &str = "x-lenso-tenant-context";
29
30#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
31#[serde(rename_all = "snake_case")]
32pub enum GrpcIdempotency {
33    Unknown,
34    Idempotent,
35    RequiresKey,
36}
37
38#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
39#[serde(rename_all = "camelCase")]
40pub struct DirectGrpcOperation {
41    pub operation_id: String,
42    pub path: String,
43    pub idempotency: GrpcIdempotency,
44    pub call_policy: CallPolicyDeclaration,
45    pub request_type: String,
46    pub response_type: String,
47}
48
49impl DirectGrpcOperation {
50    fn retry_decision(&self, code: Code, attempt: u32, key: Option<&str>) -> RetryDecision {
51        if self.idempotency == GrpcIdempotency::Unknown {
52            return RetryDecision::no("operation_retry_safety_unknown");
53        }
54        if attempt >= self.call_policy.max_attempts {
55            return RetryDecision::no("initial_policy_attempt_limit");
56        }
57        if !matches!(code, Code::Unavailable | Code::ResourceExhausted) {
58            return RetryDecision::no("failure_not_retryable");
59        }
60        match self.idempotency {
61            GrpcIdempotency::Idempotent => RetryDecision::yes(),
62            GrpcIdempotency::RequiresKey if key.is_some_and(|value| !value.is_empty()) => {
63                RetryDecision::yes()
64            }
65            GrpcIdempotency::RequiresKey => RetryDecision::no("idempotency_key_required"),
66            GrpcIdempotency::Unknown => unreachable!("unknown safety returns before matching"),
67        }
68    }
69}
70
71#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
72#[serde(rename_all = "camelCase")]
73pub struct DirectGrpcBindings {
74    pub contract_id: String,
75    pub version: String,
76    pub service_name: String,
77    pub operations: Vec<DirectGrpcOperation>,
78}
79
80impl DirectGrpcBindings {
81    #[must_use]
82    pub fn operation(&self, id: &str) -> Option<&DirectGrpcOperation> {
83        self.operations
84            .iter()
85            .find(|operation| operation.operation_id == id)
86    }
87}
88
89#[derive(Debug, Clone)]
90pub struct DirectGrpcServerPolicy {
91    bindings: DirectGrpcBindings,
92    runtime: CallPolicyRuntime,
93    workload_identity: Option<DirectGrpcWorkloadIdentity>,
94    service_context: Option<DirectGrpcServiceContext>,
95}
96
97pub struct DirectGrpcAdmission {
98    pub call_policy_permit: CallPolicyPermit,
99    pub authenticated_service_principal: Option<AuthenticatedServicePrincipal>,
100    pub authenticated_service_context: Option<AuthenticatedServiceContext>,
101}
102
103impl std::fmt::Debug for DirectGrpcAdmission {
104    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
105        formatter
106            .debug_struct("DirectGrpcAdmission")
107            .field(
108                "authenticated_service_principal",
109                &self.authenticated_service_principal,
110            )
111            .field(
112                "authenticated_service_context",
113                &self.authenticated_service_context,
114            )
115            .finish_non_exhaustive()
116    }
117}
118
119#[derive(Debug, Clone)]
120struct DirectGrpcWorkloadIdentity {
121    provider: Arc<dyn WorkloadIdentityProvider>,
122    audience: String,
123}
124
125#[derive(Debug, Clone)]
126struct DirectGrpcServiceContext {
127    admission: ServiceContextAdmission,
128    recorder: Arc<dyn IdentityDecisionRecorder>,
129}
130
131impl DirectGrpcServerPolicy {
132    #[must_use]
133    pub fn new(
134        bindings: DirectGrpcBindings,
135        provider: Arc<dyn WorkloadIdentityProvider>,
136        audience: impl Into<String>,
137        context_provider: Arc<dyn DelegatedContextVerifier>,
138        context_policies: impl IntoIterator<Item = (String, ServiceContextPolicy)>,
139        evidence_recorder: Arc<dyn IdentityDecisionRecorder>,
140    ) -> Self {
141        Self::new_unchecked(bindings)
142            .with_workload_identity(provider, audience)
143            .with_service_context(context_provider, context_policies, evidence_recorder)
144    }
145
146    #[must_use]
147    #[cfg(debug_assertions)]
148    pub fn new_without_workload_identity(bindings: DirectGrpcBindings) -> Self {
149        Self::new_unchecked(bindings)
150    }
151
152    fn new_unchecked(bindings: DirectGrpcBindings) -> Self {
153        Self {
154            bindings,
155            runtime: CallPolicyRuntime::default(),
156            workload_identity: None,
157            service_context: None,
158        }
159    }
160
161    #[must_use]
162    pub fn with_policy_runtime(mut self, runtime: CallPolicyRuntime) -> Self {
163        self.runtime = runtime;
164        self
165    }
166
167    #[must_use]
168    pub fn with_workload_identity(
169        mut self,
170        provider: Arc<dyn WorkloadIdentityProvider>,
171        audience: impl Into<String>,
172    ) -> Self {
173        self.workload_identity = Some(DirectGrpcWorkloadIdentity {
174            provider,
175            audience: audience.into(),
176        });
177        self
178    }
179
180    #[must_use]
181    pub fn with_service_context<I, S>(
182        mut self,
183        provider: Arc<dyn DelegatedContextVerifier>,
184        policies: I,
185        recorder: Arc<dyn IdentityDecisionRecorder>,
186    ) -> Self
187    where
188        I: IntoIterator<Item = (S, ServiceContextPolicy)>,
189        S: Into<String>,
190    {
191        self.service_context = Some(DirectGrpcServiceContext {
192            admission: ServiceContextAdmission::new(provider, policies),
193            recorder,
194        });
195        self
196    }
197
198    pub fn admit<T>(
199        &self,
200        operation_id: &str,
201        request: &Request<T>,
202    ) -> Result<DirectGrpcAdmission, DirectGrpcAdmissionError> {
203        let operation = self.bindings.operation(operation_id).ok_or_else(|| {
204            DirectGrpcAdmissionError::Contract("operation_not_declared".to_owned())
205        })?;
206        let authenticated_service_principal = if let Some(identity) = &self.workload_identity {
207            let credential = request
208                .metadata()
209                .get(AUTHORIZATION_METADATA)
210                .and_then(|value| value.to_str().ok())
211                .and_then(|value| value.strip_prefix("Bearer "))
212                .ok_or_else(|| {
213                    DirectGrpcAdmissionError::Unauthenticated(
214                        "workload_identity_required".to_owned(),
215                    )
216                })?;
217            let binding = request
218                .extensions()
219                .get::<AuthenticatedTransportBinding>()
220                .ok_or_else(|| {
221                    DirectGrpcAdmissionError::Unauthenticated(
222                        "authenticated_transport_binding_required".to_owned(),
223                    )
224                })?;
225            let principal = identity
226                .provider
227                .verify(
228                    credential,
229                    &WorkloadIdentityVerification::new(
230                        &identity.audience,
231                        &binding.0,
232                        self.runtime.now_ms(),
233                    ),
234                )
235                .map_err(|error| {
236                    DirectGrpcAdmissionError::Unauthenticated(error.evidence.outcome)
237                })?;
238            Some(principal)
239        } else {
240            None
241        };
242        let authenticated_service_context = if let Some(context) = &self.service_context {
243            let actor = decode_metadata::<DelegatedActorContext>(
244                request.metadata(),
245                DELEGATED_ACTOR_METADATA,
246            );
247            let tenant =
248                decode_metadata::<TenantContext>(request.metadata(), TENANT_CONTEXT_METADATA);
249            let decision = match (actor, tenant) {
250                (Ok(actor), Ok(tenant)) => {
251                    let service_context = actor.map(|actor| ServiceContext::new(actor, tenant));
252                    context.admission.admit(
253                        operation_id,
254                        service_context.as_ref(),
255                        self.runtime.now_ms(),
256                    )
257                }
258                _ => Err(context.admission.invalid_proof(operation_id)),
259            };
260            match decision {
261                Ok(authenticated) => {
262                    context
263                        .recorder
264                        .record(&authenticated.evidence)
265                        .map_err(|_| {
266                            DirectGrpcAdmissionError::PermissionDenied(
267                                "identity_evidence_persistence_failed".to_owned(),
268                            )
269                        })?;
270                    Some(authenticated)
271                }
272                Err(error) => {
273                    context.recorder.record(&error.evidence).map_err(|_| {
274                        DirectGrpcAdmissionError::PermissionDenied(
275                            "identity_evidence_persistence_failed".to_owned(),
276                        )
277                    })?;
278                    return Err(DirectGrpcAdmissionError::PermissionDenied(
279                        error.evidence.outcome,
280                    ));
281                }
282            }
283        } else {
284            None
285        };
286        let deadline = request
287            .metadata()
288            .get(DEADLINE_METADATA)
289            .and_then(|value| value.to_str().ok())
290            .and_then(|value| value.parse::<u64>().ok());
291        if deadline.is_none_or(|deadline| deadline <= self.runtime.now_ms()) {
292            return Err(DirectGrpcAdmissionError::DeadlineExpired {
293                evidence: rejected_evidence(CallPolicyEvent::DeadlineExpired),
294            });
295        }
296        let idempotency_key = request
297            .metadata()
298            .get(IDEMPOTENCY_METADATA)
299            .and_then(|value| value.to_str().ok());
300        if operation.idempotency == GrpcIdempotency::RequiresKey
301            && idempotency_key.is_none_or(str::is_empty)
302        {
303            return Err(DirectGrpcAdmissionError::IdempotencyKeyRequired {
304                evidence: rejected_evidence(CallPolicyEvent::CallFailed),
305            });
306        }
307        let operation_key = format!("{}:{}", self.bindings.contract_id, operation.operation_id);
308        let call_policy_permit = self
309            .runtime
310            .admit(operation_key, &operation.call_policy)
311            .map_err(|event| DirectGrpcAdmissionError::Overloaded {
312                evidence: rejected_evidence(event),
313            })?;
314        Ok(DirectGrpcAdmission {
315            call_policy_permit,
316            authenticated_service_principal,
317            authenticated_service_context,
318        })
319    }
320}
321
322#[derive(Debug, Clone, PartialEq, Eq)]
323pub enum DirectGrpcAdmissionError {
324    Contract(String),
325    Unauthenticated(String),
326    PermissionDenied(String),
327    DeadlineExpired { evidence: CallPolicyEvidence },
328    IdempotencyKeyRequired { evidence: CallPolicyEvidence },
329    Overloaded { evidence: CallPolicyEvidence },
330}
331
332impl DirectGrpcAdmissionError {
333    #[must_use]
334    pub fn status(&self) -> Status {
335        match self {
336            Self::Contract(message) => Status::unimplemented(message.clone()),
337            Self::Unauthenticated(reason) => Status::unauthenticated(reason.clone()),
338            Self::PermissionDenied(reason) => Status::permission_denied(reason.clone()),
339            Self::DeadlineExpired { .. } => Status::deadline_exceeded("deadline_expired"),
340            Self::IdempotencyKeyRequired { .. } => {
341                Status::invalid_argument("idempotency_key_required")
342            }
343            Self::Overloaded { .. } => Status::resource_exhausted("overload_rejected"),
344        }
345    }
346}
347
348fn rejected_evidence(event: CallPolicyEvent) -> CallPolicyEvidence {
349    CallPolicyEvidence {
350        events: vec![event],
351        attempts: 0,
352        terminal_outcome: CallPolicyTerminalOutcome::Rejected,
353        fallback_handler: None,
354    }
355}
356
357pub fn generate_direct_grpc_bindings(
358    contract_id: impl Into<String>,
359    version: impl Into<String>,
360    proto_source: &str,
361    descriptor_bytes: &[u8],
362) -> Result<DirectGrpcBindings, String> {
363    let descriptor =
364        FileDescriptorSet::decode(descriptor_bytes).map_err(|error| error.to_string())?;
365    let file = descriptor
366        .file
367        .first()
368        .ok_or("Protobuf descriptor requires a file")?;
369    let service = file
370        .service
371        .first()
372        .ok_or("Protobuf descriptor requires a service")?;
373    let package = file
374        .package
375        .as_deref()
376        .ok_or("Protobuf descriptor requires a package")?;
377    let service_name = service
378        .name
379        .as_deref()
380        .ok_or("Protobuf service requires a name")?;
381    let idempotency = protobuf_idempotency_annotations(proto_source)?;
382    let call_policies = parse_protobuf_call_policies(
383        proto_source,
384        service.method.iter().map(|method| {
385            let name = method.name.as_deref().unwrap_or_default();
386            (
387                name,
388                idempotency
389                    .get(name)
390                    .copied()
391                    .unwrap_or(GrpcIdempotency::Unknown),
392            )
393        }),
394    )?;
395    let mut operations = service
396        .method
397        .iter()
398        .map(|method| {
399            let name = method
400                .name
401                .clone()
402                .ok_or("Protobuf method requires a name")?;
403            let request_type = method
404                .input_type
405                .clone()
406                .ok_or("Protobuf method requires an input type")?;
407            let idempotency = idempotency
408                .get(&name)
409                .copied()
410                .unwrap_or(GrpcIdempotency::Unknown);
411            Ok(DirectGrpcOperation {
412                path: format!("/{package}.{service_name}/{name}"),
413                idempotency,
414                call_policy: call_policies
415                    .get(&name)
416                    .cloned()
417                    .expect("every descriptor method receives a call policy"),
418                operation_id: name,
419                request_type,
420                response_type: method
421                    .output_type
422                    .clone()
423                    .ok_or("Protobuf method requires an output type")?,
424            })
425        })
426        .collect::<Result<Vec<_>, &str>>()?;
427    operations.sort_by(|left, right| left.operation_id.cmp(&right.operation_id));
428    Ok(DirectGrpcBindings {
429        contract_id: contract_id.into(),
430        version: version.into(),
431        service_name: format!("{package}.{service_name}"),
432        operations,
433    })
434}
435
436pub fn parse_protobuf_call_policies<I, S>(
437    source: &str,
438    operations: I,
439) -> Result<std::collections::BTreeMap<String, CallPolicyDeclaration>, String>
440where
441    I: IntoIterator<Item = (S, GrpcIdempotency)>,
442    S: AsRef<str>,
443{
444    let operations = operations
445        .into_iter()
446        .map(|(name, idempotency)| (name.as_ref().to_owned(), idempotency))
447        .collect::<std::collections::BTreeMap<_, _>>();
448    let mut declared = std::collections::BTreeMap::new();
449    let mut pending: Option<&str> = None;
450    for line in source.lines().map(str::trim) {
451        if let Some(value) = line.strip_prefix("// lenso-call-policy: ") {
452            pending = Some(value);
453        } else if let Some(method) = line
454            .strip_prefix("rpc ")
455            .and_then(|line| line.split('(').next())
456        {
457            let method = method.trim();
458            if let Some(value) = pending.take() {
459                let policy: CallPolicyDeclaration =
460                    serde_json::from_str(value).map_err(|error| {
461                        format!("rpc {method} has invalid lenso-call-policy: {error}")
462                    })?;
463                declared.insert(method.to_owned(), policy);
464            }
465        } else if !line.starts_with("//") && !line.is_empty() {
466            pending = None;
467        }
468    }
469    let mut policies = std::collections::BTreeMap::new();
470    for (method, idempotency) in operations {
471        let retry_safe = idempotency != GrpcIdempotency::Unknown;
472        let policy = declared
473            .remove(&method)
474            .ok_or_else(|| format!("rpc {method} requires lenso-call-policy"))?;
475        if let Some(issue) = policy.validate(retry_safe).into_iter().next() {
476            return Err(format!(
477                "rpc {method} lenso-call-policy.{}: {}",
478                issue.path, issue.code
479            ));
480        }
481        policies.insert(method, policy);
482    }
483    Ok(policies)
484}
485
486fn protobuf_idempotency_annotations(
487    source: &str,
488) -> Result<std::collections::BTreeMap<String, GrpcIdempotency>, String> {
489    let mut annotations = std::collections::BTreeMap::new();
490    let mut pending = None;
491    for line in source.lines().map(str::trim) {
492        if let Some(value) = line.strip_prefix("// lenso-idempotency: ") {
493            pending = Some(match value {
494                "idempotent" => GrpcIdempotency::Idempotent,
495                "requires_key" => GrpcIdempotency::RequiresKey,
496                other => return Err(format!("unsupported lenso-idempotency `{other}`")),
497            });
498        } else if let Some(method) = line
499            .strip_prefix("rpc ")
500            .and_then(|line| line.split('(').next())
501        {
502            if let Some(value) = pending.take() {
503                annotations.insert(method.trim().to_owned(), value);
504            }
505        } else if !line.starts_with("//") && !line.is_empty() {
506            pending = None;
507        }
508    }
509    Ok(annotations)
510}
511
512#[derive(Debug, Clone, PartialEq, Eq)]
513pub struct DirectGrpcEvidence {
514    pub operation_id: String,
515    pub decision: String,
516    pub attempts: u32,
517    pub call_policy: CallPolicyEvidence,
518    pub grpc_code: Option<Code>,
519    pub grpc_message: Option<String>,
520}
521
522#[derive(Debug)]
523pub struct DirectGrpcResponse {
524    pub payload: Vec<u8>,
525    pub metadata: tonic::metadata::MetadataMap,
526    pub evidence: DirectGrpcEvidence,
527}
528
529pub struct DirectGrpcClient<R> {
530    resolver: R,
531    bindings: DirectGrpcBindings,
532    policy_runtime: CallPolicyRuntime,
533    fallbacks: BTreeMap<String, Arc<GrpcFallback>>,
534    workload_credential: Option<String>,
535    delegated_actor_context: Option<DelegatedActorContext>,
536    tenant_context: Option<TenantContext>,
537}
538type GrpcFallback = dyn Fn(CallPolicyFailure) -> Vec<u8> + Send + Sync;
539
540impl<R> std::fmt::Debug for DirectGrpcClient<R> {
541    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
542        formatter
543            .debug_struct("DirectGrpcClient")
544            .field("bindings", &self.bindings)
545            .field("policy_runtime", &self.policy_runtime)
546            .field("fallbacks", &self.fallbacks.keys().collect::<Vec<_>>())
547            .finish_non_exhaustive()
548    }
549}
550
551enum GrpcPolicyStart {
552    Permit(CallPolicyPermit),
553    Fallback(DirectGrpcResponse),
554}
555
556impl<R: EndpointResolver> DirectGrpcClient<R> {
557    #[must_use]
558    pub fn new(resolver: R, bindings: DirectGrpcBindings) -> Self {
559        Self {
560            resolver,
561            bindings,
562            policy_runtime: CallPolicyRuntime::default(),
563            fallbacks: BTreeMap::new(),
564            workload_credential: None,
565            delegated_actor_context: None,
566            tenant_context: None,
567        }
568    }
569
570    #[must_use]
571    pub fn with_policy_runtime(mut self, policy_runtime: CallPolicyRuntime) -> Self {
572        self.policy_runtime = policy_runtime;
573        self
574    }
575
576    #[must_use]
577    pub fn with_fallback<F>(mut self, handler: impl Into<String>, fallback: F) -> Self
578    where
579        F: Fn(CallPolicyFailure) -> Vec<u8> + Send + Sync + 'static,
580    {
581        self.fallbacks.insert(handler.into(), Arc::new(fallback));
582        self
583    }
584
585    #[must_use]
586    pub fn with_workload_credential(mut self, credential: impl Into<String>) -> Self {
587        self.workload_credential = Some(credential.into());
588        self
589    }
590
591    #[must_use]
592    pub fn with_service_context(
593        mut self,
594        actor: DelegatedActorContext,
595        tenant: Option<TenantContext>,
596    ) -> Self {
597        self.delegated_actor_context = Some(actor);
598        self.tenant_context = tenant;
599        self
600    }
601
602    pub async fn get_sla(
603        &self,
604        service: &ServiceReference,
605        payload: Vec<u8>,
606        deadline: u64,
607    ) -> Result<DirectGrpcResponse, DirectGrpcCallError> {
608        let operation = self.operation("GetSla")?;
609        if deadline <= now_ms() {
610            return self.deadline_failure(operation, 0, Vec::new());
611        }
612        let mut permit = match self.start_policy(service, operation)? {
613            GrpcPolicyStart::Permit(permit) => Some(permit),
614            GrpcPolicyStart::Fallback(response) => return Ok(response),
615        };
616        let (mut client, _) = match self.prepare(service, "GetSla", deadline).await {
617            Ok(prepared) => prepared,
618            Err(error) => {
619                return self.finish_prepare_error(
620                    operation,
621                    permit.take().expect("policy permit exists"),
622                    error,
623                );
624            }
625        };
626        let mut retry_events = Vec::new();
627        for attempt in 1..=operation.call_policy.max_attempts {
628            let request = match request(
629                GetSlaRequest {
630                    payload: payload.clone(),
631                },
632                deadline,
633                None,
634                self.workload_credential.as_deref(),
635                self.delegated_actor_context.as_ref(),
636                self.tenant_context.as_ref(),
637            ) {
638                Ok(request) => request,
639                Err(DirectGrpcCallError::Contract(message)) if message == "deadline_expired" => {
640                    retry_events.push(CallPolicyEvent::DeadlineExpired);
641                    let events = permit
642                        .take()
643                        .expect("policy permit exists")
644                        .failure_after(retry_events);
645                    return self.deadline_failure(operation, attempt - 1, events);
646                }
647                Err(error) => return Err(error),
648            };
649            match client.get_sla(request).await {
650                Ok(response) => {
651                    return Ok(success(
652                        operation,
653                        response,
654                        attempt,
655                        permit.take().expect("policy permit exists"),
656                        retry_events,
657                    ));
658                }
659                Err(status) => {
660                    let decision = operation.retry_decision(status.code(), attempt, None);
661                    if decision.should_retry {
662                        retry_events.push(CallPolicyEvent::RetryScheduled);
663                        continue;
664                    }
665                    return self.finish_status(
666                        operation,
667                        status,
668                        attempt,
669                        permit.take().expect("policy permit exists"),
670                        retry_events,
671                        decision,
672                    );
673                }
674            }
675        }
676        unreachable!()
677    }
678
679    pub async fn update_sla(
680        &self,
681        service: &ServiceReference,
682        payload: Vec<u8>,
683        deadline: u64,
684        idempotency_key: &str,
685    ) -> Result<DirectGrpcResponse, DirectGrpcCallError> {
686        if idempotency_key.is_empty() {
687            return Err(DirectGrpcCallError::Contract(
688                "idempotency_key_required".to_owned(),
689            ));
690        }
691        let operation = self.operation("UpdateSla")?;
692        if deadline <= now_ms() {
693            return self.deadline_failure(operation, 0, Vec::new());
694        }
695        let mut permit = match self.start_policy(service, operation)? {
696            GrpcPolicyStart::Permit(permit) => Some(permit),
697            GrpcPolicyStart::Fallback(response) => return Ok(response),
698        };
699        let (mut client, _) = match self.prepare(service, "UpdateSla", deadline).await {
700            Ok(prepared) => prepared,
701            Err(error) => {
702                return self.finish_prepare_error(
703                    operation,
704                    permit.take().expect("policy permit exists"),
705                    error,
706                );
707            }
708        };
709        let mut retry_events = Vec::new();
710        for attempt in 1..=operation.call_policy.max_attempts {
711            let request = match request(
712                UpdateSlaRequest {
713                    payload: payload.clone(),
714                },
715                deadline,
716                Some(idempotency_key),
717                self.workload_credential.as_deref(),
718                self.delegated_actor_context.as_ref(),
719                self.tenant_context.as_ref(),
720            ) {
721                Ok(request) => request,
722                Err(DirectGrpcCallError::Contract(message)) if message == "deadline_expired" => {
723                    retry_events.push(CallPolicyEvent::DeadlineExpired);
724                    let events = permit
725                        .take()
726                        .expect("policy permit exists")
727                        .failure_after(retry_events);
728                    return self.deadline_failure(operation, attempt - 1, events);
729                }
730                Err(error) => return Err(error),
731            };
732            match client.update_sla(request).await {
733                Ok(response) => {
734                    return Ok(success(
735                        operation,
736                        response,
737                        attempt,
738                        permit.take().expect("policy permit exists"),
739                        retry_events,
740                    ));
741                }
742                Err(status) => {
743                    let decision =
744                        operation.retry_decision(status.code(), attempt, Some(idempotency_key));
745                    if decision.should_retry {
746                        retry_events.push(CallPolicyEvent::RetryScheduled);
747                        continue;
748                    }
749                    return self.finish_status(
750                        operation,
751                        status,
752                        attempt,
753                        permit.take().expect("policy permit exists"),
754                        retry_events,
755                        decision,
756                    );
757                }
758            }
759        }
760        unreachable!()
761    }
762
763    pub async fn probe_sla(
764        &self,
765        service: &ServiceReference,
766        payload: Vec<u8>,
767        deadline: u64,
768    ) -> Result<DirectGrpcResponse, DirectGrpcCallError> {
769        let operation = self.operation("ProbeSla")?;
770        if deadline <= now_ms() {
771            return self.deadline_failure(operation, 0, Vec::new());
772        }
773        let mut permit = match self.start_policy(service, operation)? {
774            GrpcPolicyStart::Permit(permit) => Some(permit),
775            GrpcPolicyStart::Fallback(response) => return Ok(response),
776        };
777        let (mut client, _) = match self.prepare(service, "ProbeSla", deadline).await {
778            Ok(prepared) => prepared,
779            Err(error) => {
780                return self.finish_prepare_error(
781                    operation,
782                    permit.take().expect("policy permit exists"),
783                    error,
784                );
785            }
786        };
787        let request = match request(
788            ProbeSlaRequest { payload },
789            deadline,
790            None,
791            self.workload_credential.as_deref(),
792            self.delegated_actor_context.as_ref(),
793            self.tenant_context.as_ref(),
794        ) {
795            Ok(request) => request,
796            Err(DirectGrpcCallError::Contract(message)) if message == "deadline_expired" => {
797                let events = permit
798                    .take()
799                    .expect("policy permit exists")
800                    .failure_after(vec![CallPolicyEvent::DeadlineExpired]);
801                return self.deadline_failure(operation, 0, events);
802            }
803            Err(error) => return Err(error),
804        };
805        match client.probe_sla(request).await {
806            Ok(response) => Ok(success(
807                operation,
808                response,
809                1,
810                permit.take().expect("policy permit exists"),
811                Vec::new(),
812            )),
813            Err(status) => {
814                let decision = operation.retry_decision(status.code(), 1, None);
815                self.finish_status(
816                    operation,
817                    status,
818                    1,
819                    permit.take().expect("policy permit exists"),
820                    Vec::new(),
821                    decision,
822                )
823            }
824        }
825    }
826
827    fn start_policy(
828        &self,
829        service: &ServiceReference,
830        operation: &DirectGrpcOperation,
831    ) -> Result<GrpcPolicyStart, DirectGrpcCallError> {
832        let operation_key = format!(
833            "{}:{}:{}",
834            service.as_str(),
835            self.bindings.contract_id,
836            operation.operation_id
837        );
838        match self
839            .policy_runtime
840            .begin_call(operation_key, &operation.call_policy)
841        {
842            Ok(permit) => Ok(GrpcPolicyStart::Permit(permit)),
843            Err(event) => {
844                let failure = match event {
845                    CallPolicyEvent::CircuitOpen => CallPolicyFailure::CircuitOpen,
846                    CallPolicyEvent::BulkheadSaturated => CallPolicyFailure::BulkheadSaturated,
847                    _ => CallPolicyFailure::NonRetryableFailure,
848                };
849                let evidence = CallPolicyEvidence {
850                    events: vec![event],
851                    attempts: 0,
852                    terminal_outcome: CallPolicyTerminalOutcome::Rejected,
853                    fallback_handler: None,
854                };
855                if let Some(response) =
856                    self.fallback_response(operation, failure, 0, evidence.events.clone(), None)
857                {
858                    return Ok(GrpcPolicyStart::Fallback(response));
859                }
860                Err(DirectGrpcCallError::Policy { failure, evidence })
861            }
862        }
863    }
864
865    fn operation(&self, operation_id: &str) -> Result<&DirectGrpcOperation, DirectGrpcCallError> {
866        self.bindings
867            .operation(operation_id)
868            .ok_or_else(|| DirectGrpcCallError::Contract("operation_not_declared".to_owned()))
869    }
870
871    fn finish_status(
872        &self,
873        operation: &DirectGrpcOperation,
874        status: Status,
875        attempt: u32,
876        permit: CallPolicyPermit,
877        mut retry_events: Vec<CallPolicyEvent>,
878        decision: RetryDecision,
879    ) -> Result<DirectGrpcResponse, DirectGrpcCallError> {
880        let retryable = matches!(status.code(), Code::Unavailable | Code::ResourceExhausted);
881        let deadline_expired = status.code() == Code::DeadlineExceeded;
882        if deadline_expired {
883            retry_events.push(CallPolicyEvent::DeadlineExpired);
884        } else if status.code() == Code::ResourceExhausted {
885            retry_events.push(CallPolicyEvent::OverloadRejected);
886        } else {
887            retry_events.push(CallPolicyEvent::CallFailed);
888        }
889        let events = if retryable || deadline_expired {
890            permit.failure_after(retry_events)
891        } else {
892            permit.success_after(retry_events)
893        };
894        let failure = if deadline_expired {
895            CallPolicyFailure::DeadlineExpired
896        } else if status.code() == Code::ResourceExhausted {
897            CallPolicyFailure::OverloadRejected
898        } else if retryable {
899            CallPolicyFailure::RetryableFailure
900        } else {
901            CallPolicyFailure::NonRetryableFailure
902        };
903        if let Some(response) =
904            self.fallback_response(operation, failure, attempt, events.clone(), Some(&status))
905        {
906            return Ok(response);
907        }
908        let evidence = DirectGrpcEvidence {
909            operation_id: operation.operation_id.clone(),
910            decision: if deadline_expired {
911                "deadline_expired".to_owned()
912            } else {
913                decision.reason.to_owned()
914            },
915            attempts: attempt,
916            call_policy: CallPolicyEvidence {
917                events,
918                attempts: attempt,
919                terminal_outcome: CallPolicyTerminalOutcome::Failed,
920                fallback_handler: None,
921            },
922            grpc_code: Some(status.code()),
923            grpc_message: Some(status.message().to_owned()),
924        };
925        Err(DirectGrpcCallError::Status { status, evidence })
926    }
927
928    fn fallback_response(
929        &self,
930        operation: &DirectGrpcOperation,
931        failure: CallPolicyFailure,
932        attempts: u32,
933        mut events: Vec<CallPolicyEvent>,
934        native_status: Option<&Status>,
935    ) -> Option<DirectGrpcResponse> {
936        let declaration = operation.call_policy.fallback_for(failure)?;
937        let fallback = self.fallbacks.get(&declaration.handler)?;
938        events.push(CallPolicyEvent::FallbackApplied);
939        Some(DirectGrpcResponse {
940            payload: fallback(failure),
941            metadata: tonic::metadata::MetadataMap::new(),
942            evidence: DirectGrpcEvidence {
943                operation_id: operation.operation_id.clone(),
944                decision: "fallback_applied".to_owned(),
945                attempts,
946                call_policy: CallPolicyEvidence {
947                    events,
948                    attempts,
949                    terminal_outcome: CallPolicyTerminalOutcome::Fallback,
950                    fallback_handler: Some(declaration.handler.clone()),
951                },
952                grpc_code: native_status.map(Status::code),
953                grpc_message: native_status.map(|status| status.message().to_owned()),
954            },
955        })
956    }
957
958    fn deadline_failure(
959        &self,
960        operation: &DirectGrpcOperation,
961        attempts: u32,
962        mut events: Vec<CallPolicyEvent>,
963    ) -> Result<DirectGrpcResponse, DirectGrpcCallError> {
964        if !events.contains(&CallPolicyEvent::DeadlineExpired) {
965            events.push(CallPolicyEvent::DeadlineExpired);
966        }
967        if let Some(response) = self.fallback_response(
968            operation,
969            CallPolicyFailure::DeadlineExpired,
970            attempts,
971            events.clone(),
972            None,
973        ) {
974            return Ok(response);
975        }
976        Err(DirectGrpcCallError::Policy {
977            failure: CallPolicyFailure::DeadlineExpired,
978            evidence: CallPolicyEvidence {
979                events,
980                attempts,
981                terminal_outcome: CallPolicyTerminalOutcome::Failed,
982                fallback_handler: None,
983            },
984        })
985    }
986
987    fn finish_prepare_error(
988        &self,
989        operation: &DirectGrpcOperation,
990        permit: CallPolicyPermit,
991        error: DirectGrpcCallError,
992    ) -> Result<DirectGrpcResponse, DirectGrpcCallError> {
993        match error {
994            DirectGrpcCallError::Contract(message) if message == "deadline_expired" => {
995                let events = permit.failure_after(vec![CallPolicyEvent::DeadlineExpired]);
996                self.deadline_failure(operation, 0, events)
997            }
998            DirectGrpcCallError::Transport {
999                source,
1000                mut evidence,
1001            } => {
1002                let events = permit.failure_after(vec![CallPolicyEvent::CallFailed]);
1003                if let Some(response) = self.fallback_response(
1004                    operation,
1005                    CallPolicyFailure::TransportFailure,
1006                    0,
1007                    events.clone(),
1008                    None,
1009                ) {
1010                    return Ok(response);
1011                }
1012                evidence.call_policy = CallPolicyEvidence {
1013                    events,
1014                    attempts: 0,
1015                    terminal_outcome: CallPolicyTerminalOutcome::Failed,
1016                    fallback_handler: None,
1017                };
1018                Err(DirectGrpcCallError::Transport { source, evidence })
1019            }
1020            error => Err(error),
1021        }
1022    }
1023
1024    async fn prepare<'a>(
1025        &'a self,
1026        service: &ServiceReference,
1027        operation_id: &str,
1028        deadline: u64,
1029    ) -> Result<(SupportServiceClient<Channel>, &'a DirectGrpcOperation), DirectGrpcCallError> {
1030        if deadline <= now_ms() {
1031            return Err(DirectGrpcCallError::Contract("deadline_expired".to_owned()));
1032        }
1033        let operation = self
1034            .bindings
1035            .operation(operation_id)
1036            .ok_or_else(|| DirectGrpcCallError::Contract("operation_not_declared".to_owned()))?;
1037        let endpoint = self
1038            .resolver
1039            .resolve(service)
1040            .map_err(|error| DirectGrpcCallError::Resolution(error.to_string()))?
1041            .endpoints
1042            .into_iter()
1043            .next()
1044            .ok_or_else(|| DirectGrpcCallError::Resolution("no usable endpoint".to_owned()))?;
1045        let remaining = deadline.saturating_sub(now_ms());
1046        if remaining == 0 {
1047            return Err(DirectGrpcCallError::Contract("deadline_expired".to_owned()));
1048        }
1049        let channel = Channel::from_shared(endpoint.address)
1050            .map_err(DirectGrpcCallError::InvalidEndpoint)?
1051            .connect_timeout(Duration::from_millis(remaining))
1052            .connect()
1053            .await
1054            .map_err(|source| DirectGrpcCallError::Transport {
1055                source,
1056                evidence: transport_evidence(operation_id),
1057            })?;
1058        Ok((SupportServiceClient::new(channel), operation))
1059    }
1060}
1061
1062fn request<T>(
1063    message: T,
1064    deadline: u64,
1065    key: Option<&str>,
1066    workload_credential: Option<&str>,
1067    delegated_actor_context: Option<&DelegatedActorContext>,
1068    tenant_context: Option<&TenantContext>,
1069) -> Result<Request<T>, DirectGrpcCallError> {
1070    let remaining = deadline.saturating_sub(now_ms());
1071    if remaining == 0 {
1072        return Err(DirectGrpcCallError::Contract("deadline_expired".to_owned()));
1073    }
1074    let mut request = Request::new(message);
1075    request.set_timeout(Duration::from_millis(remaining));
1076    request.metadata_mut().insert(
1077        DEADLINE_METADATA,
1078        MetadataValue::try_from(deadline.to_string())
1079            .map_err(|error| DirectGrpcCallError::Contract(error.to_string()))?,
1080    );
1081    if let Some(key) = key {
1082        request.metadata_mut().insert(
1083            IDEMPOTENCY_METADATA,
1084            MetadataValue::try_from(key)
1085                .map_err(|error| DirectGrpcCallError::Contract(error.to_string()))?,
1086        );
1087    }
1088    if let Some(credential) = workload_credential {
1089        request.metadata_mut().insert(
1090            AUTHORIZATION_METADATA,
1091            MetadataValue::try_from(format!("Bearer {credential}"))
1092                .map_err(|error| DirectGrpcCallError::Contract(error.to_string()))?,
1093        );
1094    }
1095    if let Some(actor) = delegated_actor_context {
1096        insert_context_metadata(request.metadata_mut(), DELEGATED_ACTOR_METADATA, actor)?;
1097    }
1098    if let Some(tenant) = tenant_context {
1099        insert_context_metadata(request.metadata_mut(), TENANT_CONTEXT_METADATA, tenant)?;
1100    }
1101    Ok(request)
1102}
1103
1104fn insert_context_metadata<T: Serialize>(
1105    metadata: &mut tonic::metadata::MetadataMap,
1106    name: &'static str,
1107    value: &T,
1108) -> Result<(), DirectGrpcCallError> {
1109    let encoded = serde_json::to_vec(value)
1110        .map(|json| URL_SAFE_NO_PAD.encode(json))
1111        .map_err(|error| DirectGrpcCallError::Contract(error.to_string()))?;
1112    metadata.insert(
1113        name,
1114        MetadataValue::try_from(encoded)
1115            .map_err(|error| DirectGrpcCallError::Contract(error.to_string()))?,
1116    );
1117    Ok(())
1118}
1119
1120fn decode_metadata<T: for<'de> Deserialize<'de>>(
1121    metadata: &tonic::metadata::MetadataMap,
1122    name: &'static str,
1123) -> Result<Option<T>, ()> {
1124    let Some(value) = metadata.get(name) else {
1125        return Ok(None);
1126    };
1127    value
1128        .to_str()
1129        .map_err(|_| ())
1130        .and_then(|value| URL_SAFE_NO_PAD.decode(value).map_err(|_| ()))
1131        .and_then(|value| serde_json::from_slice(&value).map_err(|_| ()))
1132        .map(Some)
1133}
1134
1135fn success(
1136    operation: &DirectGrpcOperation,
1137    response: tonic::Response<SlaResponse>,
1138    attempts: u32,
1139    permit: CallPolicyPermit,
1140    mut retry_events: Vec<CallPolicyEvent>,
1141) -> DirectGrpcResponse {
1142    let (metadata, value, _) = response.into_parts();
1143    retry_events.push(CallPolicyEvent::CallCompleted);
1144    let events = permit.success_after(retry_events);
1145    DirectGrpcResponse {
1146        payload: value.payload,
1147        metadata,
1148        evidence: DirectGrpcEvidence {
1149            operation_id: operation.operation_id.clone(),
1150            decision: "call_completed".to_owned(),
1151            attempts,
1152            call_policy: CallPolicyEvidence {
1153                events,
1154                attempts,
1155                terminal_outcome: CallPolicyTerminalOutcome::Completed,
1156                fallback_handler: None,
1157            },
1158            grpc_code: Some(Code::Ok),
1159            grpc_message: None,
1160        },
1161    }
1162}
1163
1164fn transport_evidence(operation_id: &str) -> DirectGrpcEvidence {
1165    DirectGrpcEvidence {
1166        operation_id: operation_id.to_owned(),
1167        decision: "transport_failure_no_retry".to_owned(),
1168        attempts: 0,
1169        call_policy: CallPolicyEvidence {
1170            events: vec![CallPolicyEvent::CallFailed],
1171            attempts: 0,
1172            terminal_outcome: CallPolicyTerminalOutcome::Failed,
1173            fallback_handler: None,
1174        },
1175        grpc_code: None,
1176        grpc_message: None,
1177    }
1178}
1179
1180#[derive(Debug)]
1181pub enum DirectGrpcCallError {
1182    Contract(String),
1183    Resolution(String),
1184    InvalidEndpoint(http::uri::InvalidUri),
1185    Transport {
1186        source: tonic::transport::Error,
1187        evidence: DirectGrpcEvidence,
1188    },
1189    Status {
1190        status: Status,
1191        evidence: DirectGrpcEvidence,
1192    },
1193    Policy {
1194        failure: CallPolicyFailure,
1195        evidence: CallPolicyEvidence,
1196    },
1197}
1198impl std::fmt::Display for DirectGrpcCallError {
1199    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1200        match self {
1201            Self::Contract(value) | Self::Resolution(value) => formatter.write_str(value),
1202            Self::InvalidEndpoint(error) => error.fmt(formatter),
1203            Self::Transport { source, .. } => source.fmt(formatter),
1204            Self::Status { status, .. } => status.fmt(formatter),
1205            Self::Policy { failure, .. } => write!(formatter, "call policy rejected: {failure:?}"),
1206        }
1207    }
1208}
1209impl std::error::Error for DirectGrpcCallError {
1210    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
1211        match self {
1212            Self::InvalidEndpoint(error) => Some(error),
1213            Self::Transport { source, .. } => Some(source),
1214            Self::Status { status, .. } => Some(status),
1215            Self::Contract(_) | Self::Resolution(_) | Self::Policy { .. } => None,
1216        }
1217    }
1218}
1219
1220fn now_ms() -> u64 {
1221    SystemTime::now()
1222        .duration_since(UNIX_EPOCH)
1223        .unwrap_or_default()
1224        .as_millis() as u64
1225}