1use std::{convert::TryFrom, future::ready, sync::Arc};
2
3use wasmtime::Caller;
4
5use crate::{
6 drivers::Capability,
7 guest_data::{GuestError, GuestResult},
8 operation::{Contract, Operation},
9 registry::{InstanceRegistry, ResourceId, ResourceType},
10 session::Session,
11};
12use selium_abi::{SessionCreate, SessionEntitlement, SessionRemove, SessionResource};
13
14type SessionOps<C> = (
15 Arc<Operation<SessionCreateDriver<C>>>,
16 Arc<Operation<SessionRemoveDriver<C>>>,
17 Arc<Operation<SessionAddEntitlementDriver<C>>>,
18 Arc<Operation<SessionRemoveEntitlementDriver<C>>>,
19 Arc<Operation<SessionAddResourceDriver<C>>>,
20 Arc<Operation<SessionRemoveResourceDriver<C>>>,
21);
22
23pub trait SessionLifecycleCapability {
25 type Error: Into<GuestError>;
26
27 fn create(&self, parent: &Session, pubkey: [u8; 32]) -> Result<Session, Self::Error>;
29 fn add_entitlement(
31 &self,
32 target: &mut Session,
33 entitlement: Capability,
34 ) -> Result<(), Self::Error>;
35 fn rm_entitlement(
37 &self,
38 target: &mut Session,
39 entitlement: Capability,
40 ) -> Result<(), Self::Error>;
41 fn add_resource(
43 &self,
44 target: &mut Session,
45 entitlement: Capability,
46 resource: ResourceId,
47 ) -> Result<bool, Self::Error>;
48 fn rm_resource(
50 &self,
51 target: &mut Session,
52 entitlement: Capability,
53 resource: ResourceId,
54 ) -> Result<bool, Self::Error>;
55 fn remove(&self, target: &Session) -> Result<(), Self::Error>;
57}
58
59impl<T> SessionLifecycleCapability for Arc<T>
60where
61 T: SessionLifecycleCapability,
62{
63 type Error = T::Error;
64
65 fn create(&self, parent: &Session, pubkey: [u8; 32]) -> Result<Session, Self::Error> {
66 self.as_ref().create(parent, pubkey)
67 }
68
69 fn add_entitlement(
70 &self,
71 target: &mut Session,
72 entitlement: Capability,
73 ) -> Result<(), Self::Error> {
74 self.as_ref().add_entitlement(target, entitlement)
75 }
76
77 fn rm_entitlement(
78 &self,
79 target: &mut Session,
80 entitlement: Capability,
81 ) -> Result<(), Self::Error> {
82 self.as_ref().rm_entitlement(target, entitlement)
83 }
84
85 fn add_resource(
86 &self,
87 target: &mut Session,
88 entitlement: Capability,
89 resource: ResourceId,
90 ) -> Result<bool, Self::Error> {
91 self.as_ref().add_resource(target, entitlement, resource)
92 }
93
94 fn rm_resource(
95 &self,
96 target: &mut Session,
97 entitlement: Capability,
98 resource: ResourceId,
99 ) -> Result<bool, Self::Error> {
100 self.as_ref().rm_resource(target, entitlement, resource)
101 }
102
103 fn remove(&self, target: &Session) -> Result<(), Self::Error> {
104 self.as_ref().remove(target)
105 }
106}
107
108pub struct SessionCreateDriver<Impl>(Impl);
109pub struct SessionAddEntitlementDriver<Impl>(Impl);
110pub struct SessionRemoveEntitlementDriver<Impl>(Impl);
111pub struct SessionAddResourceDriver<Impl>(Impl);
112pub struct SessionRemoveResourceDriver<Impl>(Impl);
113pub struct SessionRemoveDriver<Impl>(Impl);
114
115impl<Impl> Contract for SessionCreateDriver<Impl>
116where
117 Impl: SessionLifecycleCapability + Clone + Send + 'static,
118{
119 type Input = SessionCreate;
120 type Output = u32;
121
122 fn to_future(
123 &self,
124 caller: &mut Caller<'_, InstanceRegistry>,
125 input: Self::Input,
126 ) -> impl Future<Output = GuestResult<Self::Output>> + 'static {
127 let inner = self.0.clone();
128 let SessionCreate { session_id, pubkey } = input;
129
130 let result = (|| -> GuestResult<u32> {
131 let parent_slot = session_id as usize;
132 let new_session = match caller
133 .data()
134 .with::<Session, _>(parent_slot, |session| inner.clone().create(session, pubkey))
135 {
136 Some(Ok(session)) => session,
137 Some(Err(err)) => return Err(err.into()),
138 None => return Err(GuestError::NotFound),
139 };
140
141 let slot = {
142 caller
143 .data_mut()
144 .insert(new_session, None, ResourceType::Session)
145 .map_err(GuestError::from)?
146 };
147
148 {
149 let granted = caller
150 .data()
151 .with::<Session, _>(parent_slot, |session| {
152 session.grant_resource(Capability::SessionLifecycle, slot)
153 })
154 .ok_or(GuestError::NotFound)?;
155
156 if !granted {
157 return Err(GuestError::PermissionDenied);
158 }
159 }
160
161 let handle = u32::try_from(slot).map_err(|_| GuestError::InvalidArgument)?;
162 Ok(handle)
163 })();
164
165 ready(result)
166 }
167}
168
169impl<Impl> Contract for SessionAddEntitlementDriver<Impl>
170where
171 Impl: SessionLifecycleCapability + Clone + Send + 'static,
172{
173 type Input = SessionEntitlement;
174 type Output = ();
175
176 fn to_future(
177 &self,
178 caller: &mut Caller<'_, InstanceRegistry>,
179 input: Self::Input,
180 ) -> impl Future<Output = GuestResult<Self::Output>> + 'static {
181 let inner = self.0.clone();
182 let SessionEntitlement {
183 session_id,
184 target_id,
185 capability,
186 } = input;
187
188 let result = (|| -> GuestResult<()> {
189 let session_slot = session_id as usize;
190 let target_slot = target_id as usize;
191
192 let authorised = caller
193 .data()
194 .with::<Session, _>(session_slot, |parent| {
195 parent.authorise(Capability::SessionLifecycle, target_slot)
196 })
197 .ok_or(GuestError::NotFound)?;
198
199 if !authorised {
200 return Err(GuestError::PermissionDenied);
201 }
202
203 match caller
204 .data_mut()
205 .with::<Session, _>(target_slot, move |target| {
206 inner.clone().add_entitlement(target, capability)
207 }) {
208 Some(Ok(())) => Ok(()),
209 Some(Err(err)) => Err(err.into()),
210 None => Err(GuestError::NotFound),
211 }
212 })();
213
214 ready(result)
215 }
216}
217
218impl<Impl> Contract for SessionRemoveEntitlementDriver<Impl>
219where
220 Impl: SessionLifecycleCapability + Clone + Send + 'static,
221{
222 type Input = SessionEntitlement;
223 type Output = ();
224
225 fn to_future(
226 &self,
227 caller: &mut Caller<'_, InstanceRegistry>,
228 input: Self::Input,
229 ) -> impl Future<Output = GuestResult<Self::Output>> + 'static {
230 let inner = self.0.clone();
231 let SessionEntitlement {
232 session_id,
233 target_id,
234 capability,
235 } = input;
236
237 let result = (|| -> GuestResult<()> {
238 let session_slot = session_id as usize;
239 let target_slot = target_id as usize;
240
241 let authorised = caller
242 .data()
243 .with::<Session, _>(session_slot, |parent| {
244 parent.authorise(Capability::SessionLifecycle, target_slot)
245 })
246 .ok_or(GuestError::NotFound)?;
247
248 if !authorised {
249 return Err(GuestError::PermissionDenied);
250 }
251
252 match caller
253 .data_mut()
254 .with::<Session, _>(target_slot, move |target| {
255 inner.clone().rm_entitlement(target, capability)
256 }) {
257 Some(Ok(())) => Ok(()),
258 Some(Err(err)) => Err(err.into()),
259 None => Err(GuestError::NotFound),
260 }
261 })();
262
263 ready(result)
264 }
265}
266
267impl<Impl> Contract for SessionAddResourceDriver<Impl>
268where
269 Impl: SessionLifecycleCapability + Clone + Send + 'static,
270{
271 type Input = SessionResource;
272 type Output = u32;
273
274 fn to_future(
275 &self,
276 caller: &mut Caller<'_, InstanceRegistry>,
277 input: Self::Input,
278 ) -> impl Future<Output = GuestResult<Self::Output>> + 'static {
279 let inner = self.0.clone();
280 let SessionResource {
281 session_id,
282 target_id,
283 capability,
284 resource_id,
285 } = input;
286
287 let result = (|| -> GuestResult<u32> {
288 let session_slot = session_id as usize;
289 let target_slot = target_id as usize;
290 let resource_slot =
291 ResourceId::try_from(resource_id).map_err(|_| GuestError::InvalidArgument)?;
292
293 let authorised = caller
294 .data()
295 .with::<Session, _>(session_slot, |parent| {
296 parent.authorise(Capability::SessionLifecycle, target_slot)
297 })
298 .ok_or(GuestError::NotFound)?;
299
300 if !authorised {
301 return Err(GuestError::PermissionDenied);
302 }
303
304 match caller
305 .data_mut()
306 .with::<Session, _>(target_slot, move |target| {
307 inner
308 .clone()
309 .add_resource(target, capability, resource_slot)
310 }) {
311 Some(Ok(true)) => Ok(1),
312 Some(Ok(false)) => Ok(0),
313 Some(Err(err)) => Err(err.into()),
314 None => Err(GuestError::NotFound),
315 }
316 })();
317
318 ready(result)
319 }
320}
321
322impl<Impl> Contract for SessionRemoveResourceDriver<Impl>
323where
324 Impl: SessionLifecycleCapability + Clone + Send + 'static,
325{
326 type Input = SessionResource;
327 type Output = u32;
328
329 fn to_future(
330 &self,
331 caller: &mut Caller<'_, InstanceRegistry>,
332 input: Self::Input,
333 ) -> impl Future<Output = GuestResult<Self::Output>> + 'static {
334 let inner = self.0.clone();
335 let SessionResource {
336 session_id,
337 target_id,
338 capability,
339 resource_id,
340 } = input;
341
342 let result = (|| -> GuestResult<u32> {
343 let session_slot = session_id as usize;
344 let target_slot = target_id as usize;
345 let resource_slot =
346 ResourceId::try_from(resource_id).map_err(|_| GuestError::InvalidArgument)?;
347
348 let authorised = caller
349 .data()
350 .with::<Session, _>(session_slot, |parent| {
351 parent.authorise(Capability::SessionLifecycle, target_slot)
352 })
353 .ok_or(GuestError::NotFound)?;
354
355 if !authorised {
356 return Err(GuestError::PermissionDenied);
357 }
358
359 match caller
360 .data_mut()
361 .with::<Session, _>(target_slot, move |target| {
362 inner.clone().rm_resource(target, capability, resource_slot)
363 }) {
364 Some(Ok(removed)) => Ok(if removed { 1 } else { 0 }),
365 Some(Err(err)) => Err(err.into()),
366 None => Err(GuestError::NotFound),
367 }
368 })();
369
370 ready(result)
371 }
372}
373
374impl<Impl> Contract for SessionRemoveDriver<Impl>
375where
376 Impl: SessionLifecycleCapability + Clone + Send + 'static,
377{
378 type Input = SessionRemove;
379 type Output = ();
380
381 fn to_future(
382 &self,
383 caller: &mut Caller<'_, InstanceRegistry>,
384 input: Self::Input,
385 ) -> impl Future<Output = GuestResult<Self::Output>> + 'static {
386 let inner = self.0.clone();
387 let SessionRemove {
388 session_id,
389 target_id,
390 } = input;
391
392 let result = (|| -> GuestResult<()> {
393 let session_slot = session_id as usize;
394 let target_slot = target_id as usize;
395
396 let authorised = caller
397 .data()
398 .with::<Session, _>(session_slot, |parent| {
399 parent.authorise(Capability::SessionLifecycle, target_slot)
400 })
401 .ok_or(GuestError::NotFound)?;
402
403 if !authorised {
404 return Err(GuestError::PermissionDenied);
405 }
406
407 if let Some(Err(err)) = caller
408 .data()
409 .with::<Session, _>(target_slot, |target| inner.clone().remove(target))
410 {
411 return Err(err.into());
412 }
413
414 caller.data_mut().remove::<Session>(target_slot);
415
416 match caller.data().with::<Session, _>(session_slot, |session| {
417 session.revoke_resource(Capability::SessionLifecycle, target_slot)
418 }) {
419 Some(Ok(_)) => Ok(()),
420 Some(Err(err)) => Err(err.into()),
421 None => Err(GuestError::NotFound),
422 }
423 })();
424
425 ready(result)
426 }
427}
428
429pub fn operations<C>(cap: C) -> SessionOps<C>
430where
431 C: SessionLifecycleCapability + Clone + Send + 'static,
432{
433 (
434 Operation::from_hostcall(
435 SessionCreateDriver(cap.clone()),
436 selium_abi::hostcall_contract!(SESSION_CREATE),
437 ),
438 Operation::from_hostcall(
439 SessionRemoveDriver(cap.clone()),
440 selium_abi::hostcall_contract!(SESSION_REMOVE),
441 ),
442 Operation::from_hostcall(
443 SessionAddEntitlementDriver(cap.clone()),
444 selium_abi::hostcall_contract!(SESSION_ADD_ENTITLEMENT),
445 ),
446 Operation::from_hostcall(
447 SessionRemoveEntitlementDriver(cap.clone()),
448 selium_abi::hostcall_contract!(SESSION_RM_ENTITLEMENT),
449 ),
450 Operation::from_hostcall(
451 SessionAddResourceDriver(cap.clone()),
452 selium_abi::hostcall_contract!(SESSION_ADD_RESOURCE),
453 ),
454 Operation::from_hostcall(
455 SessionRemoveResourceDriver(cap),
456 selium_abi::hostcall_contract!(SESSION_RM_RESOURCE),
457 ),
458 )
459}