1use std::sync::Arc;
4
5use tracing::{debug, info};
6use typesec_core::{
7 Capability, Permission, Resource,
8 policy::{CapabilityError, MintOptions, PolicyEngine, mint_capability_for_id_async},
9 typestate::{Agent, AgentError, AgentState, Authenticated, Credentials, Unauthenticated},
10};
11
12pub struct SecureAgent<S: AgentState> {
22 inner: Agent<S>,
23}
24
25impl SecureAgent<Unauthenticated> {
26 pub fn new(engine: Arc<dyn PolicyEngine>) -> Self {
28 Self {
29 inner: Agent::new(engine),
30 }
31 }
32
33 pub fn authenticate_with(
40 self,
41 credentials: Credentials,
42 authenticator: &dyn typesec_core::typestate::Authenticator,
43 ) -> Result<SecureAgent<Authenticated>, AgentError> {
44 let inner = self.inner.authenticate_with(credentials, authenticator)?;
45 Ok(SecureAgent { inner })
46 }
47
48 pub fn authenticate_unverified(
52 self,
53 credentials: Credentials,
54 ) -> Result<SecureAgent<Authenticated>, AgentError> {
55 let inner = self.inner.authenticate_unverified(credentials)?;
56 Ok(SecureAgent { inner })
57 }
58}
59
60impl SecureAgent<Authenticated> {
61 pub fn subject(&self) -> &str {
63 self.inner.subject()
64 }
65
66 pub fn engine(&self) -> Arc<dyn PolicyEngine> {
70 self.inner.engine().clone()
71 }
72
73 pub async fn request_capability<P: Permission, R: Resource>(
84 &self,
85 resource: &R,
86 ) -> Result<Capability<P, R>, CapabilityError> {
87 self.request_capability_with(resource, MintOptions::default())
88 .await
89 }
90
91 pub async fn request_capability_with<P: Permission, R: Resource>(
96 &self,
97 resource: &R,
98 options: MintOptions,
99 ) -> Result<Capability<P, R>, CapabilityError> {
100 let subject = self.subject().to_owned();
101 let action = P::name();
102 let resource_id = resource.resource_id().to_owned();
103 let engine = self.inner.engine().clone();
104
105 debug!(%subject, action, %resource_id, "requesting capability");
106
107 let cap = mint_capability_for_id_async::<P, R>(
108 engine.as_ref(),
109 subject.as_str(),
110 resource_id.as_str(),
111 &options,
112 )
113 .await?;
114
115 info!(%subject, action, %resource_id, "capability granted");
116
117 Ok(cap)
118 }
119
120 pub async fn execute<P, R, F, Fut>(
142 &self,
143 cap: &Capability<P, R>,
144 resource: &R,
145 action: F,
146 ) -> Result<(), crate::executor::TaskError>
147 where
148 P: Permission,
149 R: Resource,
150 F: FnOnce(&R) -> Fut,
151 Fut: std::future::Future<Output = Result<(), crate::executor::TaskError>>,
152 {
153 if cap.subject() != self.subject() {
154 return Err(crate::executor::TaskError::CapabilityMismatch(format!(
155 "capability was minted for subject '{}', not '{}'",
156 cap.subject(),
157 self.subject()
158 )));
159 }
160 if cap.resource_id() != resource.resource_id() {
161 return Err(crate::executor::TaskError::CapabilityMismatch(format!(
162 "capability covers resource '{}', not '{}'",
163 cap.resource_id(),
164 resource.resource_id()
165 )));
166 }
167 cap.ensure_active()?;
168
169 info!(
170 subject = %self.subject(),
171 permission = %Capability::<P, R>::permission_name(),
172 resource = %cap.resource_id(),
173 "executing with capability"
174 );
175
176 action(resource).await
177 }
178}
179
180impl<S: AgentState> std::fmt::Debug for SecureAgent<S> {
181 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
182 write!(f, "SecureAgent({:?})", self.inner)
183 }
184}
185
186pub struct AgentBuilder {
188 engine: Option<Arc<dyn PolicyEngine>>,
189}
190
191impl AgentBuilder {
192 pub fn new() -> Self {
194 Self { engine: None }
195 }
196
197 pub fn with_engine(mut self, engine: Arc<dyn PolicyEngine>) -> Self {
199 self.engine = Some(engine);
200 self
201 }
202
203 pub fn with_composed_engine(
212 self,
213 primary: Arc<dyn PolicyEngine>,
214 fallback: Arc<dyn PolicyEngine>,
215 ) -> Self {
216 self.with_composed_engine_strategy(
217 primary,
218 fallback,
219 typesec_core::combinator::CombineStrategy::PriorityOrder,
220 )
221 }
222
223 pub fn with_composed_engine_strategy(
225 mut self,
226 primary: Arc<dyn PolicyEngine>,
227 fallback: Arc<dyn PolicyEngine>,
228 strategy: typesec_core::combinator::CombineStrategy,
229 ) -> Self {
230 use typesec_core::combinator::PolicyEngineBuilder;
231 let engine = PolicyEngineBuilder::new()
232 .add_engine(primary)
233 .add_engine(fallback)
234 .strategy(strategy)
235 .build();
236 self.engine = Some(Arc::new(engine));
237 self
238 }
239
240 pub fn build(self) -> Result<SecureAgent<Unauthenticated>, String> {
242 let engine = self.engine.ok_or("no policy engine configured")?;
243 Ok(SecureAgent::new(engine))
244 }
245}
246
247impl Default for AgentBuilder {
248 fn default() -> Self {
249 Self::new()
250 }
251}
252
253#[cfg(test)]
254mod tests {
255 use super::*;
256 use std::sync::Arc;
257 use typesec_core::{
258 ResourceId, SubjectId, permissions::CanRead, policy::PolicyResult,
259 resource::GenericResource,
260 };
261
262 struct AllowAll;
263 impl PolicyEngine for AllowAll {
264 fn check(&self, _: &SubjectId, _: &str, _: &ResourceId) -> PolicyResult {
265 PolicyResult::Allow
266 }
267 }
268
269 struct DenyAll;
270 impl PolicyEngine for DenyAll {
271 fn check(&self, _: &SubjectId, _: &str, _: &ResourceId) -> PolicyResult {
272 PolicyResult::Deny("DenyAll".into())
273 }
274 }
275
276 #[tokio::test]
277 async fn full_flow_allow() {
278 let agent = SecureAgent::new(Arc::new(AllowAll));
279 let agent = agent
280 .authenticate_unverified(Credentials::new("agent:test", "tok"))
281 .expect("auth ok");
282 let resource = GenericResource::new("reports/q1", "report");
283 let cap: Capability<CanRead, GenericResource> = agent
284 .request_capability(&resource)
285 .await
286 .expect("should get cap");
287 assert_eq!(cap.subject(), "agent:test");
288 }
289
290 #[tokio::test]
291 async fn denied_request_returns_error() {
292 let agent = SecureAgent::new(Arc::new(DenyAll));
293 let agent = agent
294 .authenticate_unverified(Credentials::new("agent:test", "tok"))
295 .expect("auth ok");
296 let resource = GenericResource::new("reports/q1", "report");
297 let result: Result<Capability<CanRead, GenericResource>, _> =
298 agent.request_capability(&resource).await;
299 assert!(result.is_err());
300 }
301
302 #[tokio::test]
303 async fn execute_requires_capability() {
304 let agent = SecureAgent::new(Arc::new(AllowAll));
305 let agent = agent
306 .authenticate_unverified(Credentials::new("agent:test", "tok"))
307 .expect("auth ok");
308 let resource = GenericResource::new("reports/q1", "report");
309 let cap: Capability<CanRead, GenericResource> =
310 agent.request_capability(&resource).await.expect("cap ok");
311
312 let executed = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false));
313 let executed_clone = executed.clone();
314
315 agent
316 .execute(&cap, &resource, |_r| {
317 let flag = executed_clone.clone();
318 Box::pin(async move {
319 flag.store(true, std::sync::atomic::Ordering::SeqCst);
320 Ok(())
321 })
322 })
323 .await
324 .expect("execute ok");
325
326 assert!(executed.load(std::sync::atomic::Ordering::SeqCst));
327 }
328
329 #[tokio::test]
330 async fn execute_rejects_capability_for_other_resource() {
331 let agent = SecureAgent::new(Arc::new(AllowAll));
332 let agent = agent
333 .authenticate_unverified(Credentials::new("agent:test", "tok"))
334 .expect("auth ok");
335 let q1 = GenericResource::new("reports/q1", "report");
336 let q2 = GenericResource::new("reports/q2", "report");
337 let cap: Capability<CanRead, GenericResource> =
338 agent.request_capability(&q1).await.expect("cap ok");
339
340 let result = agent
342 .execute(&cap, &q2, |_r| Box::pin(async { Ok(()) }))
343 .await;
344 assert!(matches!(
345 result,
346 Err(crate::executor::TaskError::CapabilityMismatch(_))
347 ));
348 }
349
350 #[tokio::test]
351 async fn execute_rejects_capability_for_other_subject() {
352 let resource = GenericResource::new("reports/q1", "report");
353
354 let other = SecureAgent::new(Arc::new(AllowAll))
356 .authenticate_unverified(Credentials::new("agent:other", "tok"))
357 .expect("auth ok");
358 let cap: Capability<CanRead, GenericResource> =
359 other.request_capability(&resource).await.expect("cap ok");
360
361 let agent = SecureAgent::new(Arc::new(AllowAll))
363 .authenticate_unverified(Credentials::new("agent:test", "tok"))
364 .expect("auth ok");
365 let result = agent
366 .execute(&cap, &resource, |_r| Box::pin(async { Ok(()) }))
367 .await;
368 assert!(matches!(
369 result,
370 Err(crate::executor::TaskError::CapabilityMismatch(_))
371 ));
372 }
373}