1use std::time::Duration;
2use std::{collections::BTreeMap, fmt};
3
4use super::{RequestId, lifecycle::CancellationToken};
5
6#[derive(Clone, Eq, PartialEq)]
8pub struct InvocationExtension {
9 key: String,
10 value: Vec<u8>,
11}
12
13impl fmt::Debug for InvocationExtension {
14 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
15 formatter
16 .debug_struct("InvocationExtension")
17 .field("key", &self.key)
18 .field("value", &"<redacted>")
19 .finish()
20 }
21}
22
23impl InvocationExtension {
24 pub fn new(key: impl Into<String>, value: Vec<u8>) -> Self {
26 Self {
27 key: key.into(),
28 value,
29 }
30 }
31
32 pub fn key(&self) -> &str {
34 &self.key
35 }
36
37 pub fn value(&self) -> &[u8] {
39 &self.value
40 }
41}
42
43#[derive(Clone, Eq, PartialEq)]
45pub struct SealedInvocationExtension {
46 key: String,
47 issuer: String,
48 audience: Vec<String>,
49 value: Vec<u8>,
50 proof: String,
51}
52
53impl SealedInvocationExtension {
54 pub fn signed(
59 key: impl Into<String>,
60 issuer: impl Into<String>,
61 audience: impl IntoIterator<Item = impl Into<String>>,
62 value: Vec<u8>,
63 proof: impl Into<String>,
64 ) -> Self {
65 Self {
66 key: key.into(),
67 issuer: issuer.into(),
68 audience: audience.into_iter().map(Into::into).collect(),
69 value,
70 proof: proof.into(),
71 }
72 }
73
74 pub fn key(&self) -> &str {
76 &self.key
77 }
78
79 pub fn issuer(&self) -> &str {
81 &self.issuer
82 }
83
84 pub fn audience(&self) -> &[String] {
86 &self.audience
87 }
88
89 pub fn value(&self) -> &[u8] {
91 &self.value
92 }
93
94 pub fn proof(&self) -> &str {
96 &self.proof
97 }
98
99 pub fn covers(&self, capability_id: &str, operation: &str) -> bool {
101 let target = format!("{capability_id}:{operation}");
102 self.audience.iter().any(|audience| audience == &target)
103 }
104}
105
106impl fmt::Debug for SealedInvocationExtension {
107 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
108 formatter
109 .debug_struct("SealedInvocationExtension")
110 .field("key", &self.key)
111 .field("issuer", &self.issuer)
112 .field("audience", &self.audience)
113 .field("value", &"<redacted>")
114 .field("proof", &"<redacted>")
115 .finish()
116 }
117}
118
119#[derive(Clone, Debug, Eq, PartialEq)]
121pub enum InvocationContextError {
122 EmptyExtensionKey,
124 ExtensionAlreadySet { key: String },
126 SealedExtensionAlreadySet { key: String },
128 InvalidSealedExtension { key: String },
130}
131
132impl std::fmt::Display for InvocationContextError {
133 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
134 match self {
135 Self::EmptyExtensionKey => {
136 formatter.write_str("Invocation Context extension key is empty")
137 }
138 Self::ExtensionAlreadySet { key } => {
139 write!(
140 formatter,
141 "Invocation Context extension `{key}` is already set"
142 )
143 }
144 Self::SealedExtensionAlreadySet { key } => {
145 write!(
146 formatter,
147 "sealed Invocation Context extension `{key}` is already set"
148 )
149 }
150 Self::InvalidSealedExtension { key } => {
151 write!(
152 formatter,
153 "sealed Invocation Context extension `{key}` has invalid provenance"
154 )
155 }
156 }
157 }
158}
159
160#[derive(Clone, Debug)]
162pub struct InvocationContext {
163 pub(super) caller_instance: Option<String>,
164 pub(super) request_id: RequestId,
165 pub(super) deadline: Option<Duration>,
166 pub(super) cancellation: CancellationToken,
167 pub(super) extensions: BTreeMap<String, InvocationExtension>,
168 pub(super) sealed_extensions: BTreeMap<String, SealedInvocationExtension>,
169}
170
171impl InvocationContext {
172 pub fn new(
174 request_id: RequestId,
175 deadline: Option<Duration>,
176 cancellation: CancellationToken,
177 ) -> Self {
178 Self {
179 caller_instance: None,
180 request_id,
181 deadline,
182 cancellation,
183 extensions: BTreeMap::new(),
184 sealed_extensions: BTreeMap::new(),
185 }
186 }
187
188 #[must_use]
190 pub fn with_caller_instance(mut self, caller_instance: impl Into<String>) -> Self {
191 self.caller_instance = Some(caller_instance.into());
192 self
193 }
194
195 pub fn caller_instance(&self) -> Option<&str> {
197 self.caller_instance.as_deref()
198 }
199
200 pub const fn request_id(&self) -> RequestId {
202 self.request_id
203 }
204
205 pub const fn deadline(&self) -> Option<Duration> {
207 self.deadline
208 }
209
210 pub fn cancellation(&self) -> CancellationToken {
212 self.cancellation.clone()
213 }
214
215 pub fn with_extension(
217 mut self,
218 key: impl Into<String>,
219 value: Vec<u8>,
220 ) -> Result<Self, InvocationContextError> {
221 let extension = InvocationExtension::new(key, value);
222 if extension.key().is_empty() {
223 return Err(InvocationContextError::EmptyExtensionKey);
224 }
225 if self.sealed_extensions.contains_key(extension.key()) {
226 return Err(InvocationContextError::SealedExtensionAlreadySet {
227 key: extension.key().to_owned(),
228 });
229 }
230 if self.extensions.contains_key(extension.key()) {
231 return Err(InvocationContextError::ExtensionAlreadySet {
232 key: extension.key().to_owned(),
233 });
234 }
235 self.extensions
236 .insert(extension.key().to_owned(), extension);
237 Ok(self)
238 }
239
240 pub fn with_sealed_extension(
242 mut self,
243 extension: SealedInvocationExtension,
244 ) -> Result<Self, InvocationContextError> {
245 if extension.key().is_empty() {
246 return Err(InvocationContextError::EmptyExtensionKey);
247 }
248 if extension.issuer().is_empty()
249 || extension.audience().is_empty()
250 || extension.proof().is_empty()
251 || extension
252 .audience()
253 .iter()
254 .any(|audience| audience.is_empty())
255 {
256 return Err(InvocationContextError::InvalidSealedExtension {
257 key: extension.key().to_owned(),
258 });
259 }
260 if self.sealed_extensions.contains_key(extension.key())
261 || self.extensions.contains_key(extension.key())
262 {
263 return Err(InvocationContextError::SealedExtensionAlreadySet {
264 key: extension.key().to_owned(),
265 });
266 }
267 self.sealed_extensions
268 .insert(extension.key().to_owned(), extension);
269 Ok(self)
270 }
271
272 pub fn extension(&self, key: &str) -> Option<&[u8]> {
274 self.extensions.get(key).map(InvocationExtension::value)
275 }
276
277 pub fn extensions(&self) -> impl Iterator<Item = &InvocationExtension> {
279 self.extensions.values()
280 }
281
282 pub fn sealed_extension(&self, key: &str) -> Option<&SealedInvocationExtension> {
284 self.sealed_extensions.get(key)
285 }
286
287 pub fn sealed_extensions(&self) -> impl Iterator<Item = &SealedInvocationExtension> {
289 self.sealed_extensions.values()
290 }
291
292 #[must_use]
297 pub fn for_target(mut self, capability_id: &str, operation: &str) -> Self {
298 self.sealed_extensions
299 .retain(|_, extension| extension.covers(capability_id, operation));
300 self
301 }
302
303 pub fn is_cancelled(&self) -> bool {
305 self.cancellation.is_cancelled()
306 }
307
308 pub fn is_expired(&self, now: Duration) -> bool {
310 self.deadline.is_some_and(|deadline| deadline <= now)
311 }
312}