Repository: permitio/cedar-agent (language: rust)

Issue:
Passing additional entities to `POST /is_authorized` endpoint
Our system has a substantial number of entities, making it impractical to store them all in the memory of the Cedar Agent. To overcome this limitation, we propose enhancing the functionality of the `POST /is_authorized` endpoint by allowing to pass in additional entities.  These entities would be considered when evaluating policies for the current request.

### How I would expect it to work:

Let's say that the data stored in the memory of Cedar Agent is the following:
```json
[
  {
    "attrs": {},
    "parents": [],
    "uid": {
      "id": "fraud",
      "type": "ResourceGroup"
    }
  },
  {
    "attrs": {},
    "parents": [
      {
        "id": "Admin",
        "type": "Role"
      }
    ],
    "uid": {
      "id": "admin@domain.com",
      "type": "User"
    }
  },
  {
    "attrs": {},
    "parents": [],
    "uid": {
      "id": "get",
      "type": "Action"
    }
  },
  {
    "attrs": {},
    "parents": [],
    "uid": {
      "id": "Admin",
      "type": "Role"
    }
  }
]
```

Additionally, let's consider the following policies:
```json
[
  {
    "id": "fraud-docs-policy",
    "content": "permit(principal in Role::\"Admin\",action in [Action::\"get\"],resource in ResourceGroup::\"fraud\");"
  }
]
```

Consider the authorization request:
```json
{
  "principal": "User::\"admin@domain.com\"",
  "action": "Action::\"get\"",
  "resource": "Document::\"fraud-doc\"",
  "additional_entities": [
    {
      "attrs": {},
      "parents": [
        {
          "id": "fraud",
          "type": "ResourceGroup"
        }
      ],
      "uid": {
        "id": "fraud-doc",
        "type": "Document"
      }
    }
  ]
}
```

The expected decision for this authorization request is "Allow" because the principal `User:"admin@domain.com"` 
is permitted to perform action `get` on resources belonging to `ResourceGroup:"fraud",`and the resource `Document:"fraud-doc"` belongs to that resource group.


Of course, if you have any other ideas how to solve this, then feel free to let us know.


Discussion / hints:

Fix patch:
diff --git a/src/routes/authorization.rs b/src/routes/authorization.rs
index f014772..f65f00f 100644
--- a/src/routes/authorization.rs
+++ b/src/routes/authorization.rs
@@ -19,7 +19,7 @@ pub async fn is_authorized(
     data_store: &State<Box<dyn DataStore>>,
     authorizer: &State<Authorizer>,
     authorization_call: Json<AuthorizationCall>,
-) -> Result<Json<AuthorizationAnswer>, AgentError> {    
+) -> Result<Json<AuthorizationAnswer>, AgentError> {
     let policies = policy_store.policy_set().await;
     let query: AuthorizationRequest = match authorization_call.into_inner().try_into() {
         Ok(query) => query,
@@ -30,16 +30,19 @@ pub async fn is_authorized(
         }
     };
 
-    // Temporary solution to override fetching entities from the datastore by directly passing it to the REST body. 
-    // Eventually this logic will be replaced in favor of performing live patch updates  
-    let (request, entities) = &query.get_request_entities();    
+    // Temporary solution to override fetching entities from the datastore by directly passing it to the REST body.
+    // Eventually this logic will be replaced in favor of performing live patch updates
+    let stored_entities = data_store.entities().await;
+    let (request, entities) = match query.get_request_entities(stored_entities) {
+        Ok(result) => result,
+        Err(err)=> {
+            return Err(AgentError::BadRequest {
+                reason: err.to_string(),
+            })
+        }
+    };
 
-    let request_entities = match entities {
-        None => data_store.entities().await,
-        Some(ents) => ents.clone()
-    };    
-    
     info!("Querying cedar using {:?}", &request);
-    let answer = authorizer.is_authorized(&request, &policies, &request_entities);
+    let answer = authorizer.is_authorized(&request, &policies, &entities);
     Ok(Json::from(AuthorizationAnswer::from(answer)))
 }
diff --git a/src/schemas/authorization.rs b/src/schemas/authorization.rs
index 531a3c3..61155c6 100644
--- a/src/schemas/authorization.rs
+++ b/src/schemas/authorization.rs
@@ -6,6 +6,7 @@ use std::str::FromStr;
 use cedar_policy::{Context, EntityUid, EvaluationError, Request, Response, Entities};
 use cedar_policy_core::authorizer::Decision;
 use cedar_policy_core::parser::err::ParseErrors;
+use cedar_policy_core::entities::EntitiesError;
 
 use rocket::serde::json::serde_json;
 use rocket_okapi::okapi::schemars;
@@ -19,34 +20,54 @@ pub struct AuthorizationCall {
     resource: Option<String>,
     context: Option<serde_json::Value>,
     entities: Option<serde_json::Value>,
+    additional_entities: Option<serde_json::Value>,
     /// Optional schema in JSON format.
     /// If present, this will inform the parsing: for instance, it will allow
     /// `__entity` and `__extn` escapes to be implicit, and it will error if
     /// attributes have the wrong types (e.g., string instead of integer).
     /// currently unsupported
     #[schemars(skip)]
-    policies: Option<String>,  
-    
+    policies: Option<String>,
 }
 
 pub struct AuthorizationRequest {
     request: Request,
-    entities: Option<Entities>
+    entities: Option<Entities>,
+    additional_entities: Option<Entities>,
 }
 
 impl AuthorizationRequest {
-
-    pub fn new(request: Request, entities: Option<Entities>) -> AuthorizationRequest {
-        AuthorizationRequest { request, entities }
+    pub fn new(
+        request: Request,
+        entities: Option<Entities>,
+        additional_entities: Option<Entities>,
+    ) -> AuthorizationRequest {
+        AuthorizationRequest {
+            request,
+            entities,
+            additional_entities,
+        }
     }
 
     pub fn get_entities(self) -> Option<Entities> {
         self.entities
     }
 
-    pub fn get_request_entities(self) -> (Request, Option<Entities>) {        
-
-        (self.request, self.entities)
+    pub fn get_request_entities(self, stored_entities: Entities) -> Result<(Request, Entities), EntitiesError> {
+        let request_entities = match self.entities {
+            None => stored_entities,
+            Some(ents) => ents.clone()
+        };
+        let patched_entities = match self.additional_entities {
+            None => request_entities,
+            Some(ents) => {
+                match Entities::from_entities(request_entities.iter().chain(ents.iter()).cloned()) {
+                    Ok(entities) => entities,
+                    Err(err) => return Err(err)
+                }
+            }
+        };
+        Ok((self.request, patched_entities))
     }
 }
 
@@ -61,11 +82,25 @@ fn string_to_euid(optional_str: Option<String>) -> Result<Option<EntityUid>, Par
 }
 
 impl AuthorizationCall {
-
-    pub fn new(principal: Option<String>, action: Option<String>, resource: Option<String>, context:Option<rocket::serde::json::Value>, entities:Option<rocket::serde::json::Value>, policies: Option<String>) -> AuthorizationCall {
-        AuthorizationCall { principal, action, resource, context, entities, policies }
+    pub fn new(
+        principal: Option<String>,
+        action: Option<String>,
+        resource: Option<String>,
+        context: Option<rocket::serde::json::Value>,
+        entities: Option<rocket::serde::json::Value>,
+        additional_entities: Option<rocket::serde::json::Value>,
+        policies: Option<String>,
+    ) -> AuthorizationCall {
+        AuthorizationCall {
+            principal,
+            action,
+            resource,
+            context,
+            entities,
+            additional_entities,
+            policies,
+        }
     }
-
 }
 
 impl TryInto<AuthorizationRequest> for AuthorizationCall {
@@ -84,7 +119,7 @@ impl TryInto<AuthorizationRequest> for AuthorizationCall {
             Ok(r) => r,
             Err(e) => return Err(e.into()),
         };
-        let entities = match self.entities {            
+        let entities = match self.entities {
             Some(et) => match Entities::from_json_value(et, None) {
                 Ok(et) => {
                     Some(et)
@@ -93,6 +128,13 @@ impl TryInto<AuthorizationRequest> for AuthorizationCall {
             },
             None => None,
         };
+        let additional_entities = match self.additional_entities {
+            Some(et) => match Entities::from_json_value(et, None) {
+                Ok(et) => Some(et),
+                Err(e) => return Err(e.into()),
+            },
+            None => None,
+        };
         let context = match self.context {
             Some(c) => match Context::from_json_value(c, None) {
                 Ok(c) => c,
@@ -100,7 +142,11 @@ impl TryInto<AuthorizationRequest> for AuthorizationCall {
             },
             None => Context::empty(),
         };
-        Ok(AuthorizationRequest::new(Request::new(principal, action, resource, context), entities))
+        Ok(AuthorizationRequest::new(
+            Request::new(principal, action, resource, context),
+            entities,
+            additional_entities,
+        ))
     }
 }
 


Test patch:
diff --git a/tests/services/data_tests.rs b/tests/services/data_tests.rs
index 217a613..d95f953 100644
--- a/tests/services/data_tests.rs
+++ b/tests/services/data_tests.rs
@@ -35,11 +35,11 @@ async fn test_load_entities_from_file() {
 }
 
 #[tokio::test]
-async fn test_load_empty_entities_from_authz_call() {   
+async fn test_load_empty_entities_from_authz_call() {
 
     let entities: String = String::from("[]");
 
-    let query = make_authz_call(entities);        
+    let query = make_authz_call(entities);
 
     match query {
         Ok(req) => assert_eq!(req.get_entities().unwrap(), Entities::empty()),
@@ -48,7 +48,7 @@ async fn test_load_empty_entities_from_authz_call() {
 }
 
 #[tokio::test]
-async fn test_load_no_entities_from_authz_call() {       
+async fn test_load_no_entities_from_authz_call() {
 
     let query = make_authz_call_no_entities();
 
@@ -60,7 +60,7 @@ async fn test_load_no_entities_from_authz_call() {
 
 
 #[tokio::test]
-async fn test_load_entities_from_authz_call() {   
+async fn test_load_entities_from_authz_call() {
 
     let entities: String = r#"
     [
@@ -89,13 +89,104 @@ async fn test_load_entities_from_authz_call() {
     "#
     .to_string();
 
-    let query = make_authz_call(entities);    
+    let query = make_authz_call(entities);
 
     match query {
-        Ok(req) => {                        
-            assert_ne!(req.get_entities(), None);            
+        Ok(req) => {
+            assert_ne!(req.get_entities(), None);
         },
-        _ => assert!(false)        
+        _ => assert!(false)
+    };
+}
+
+#[tokio::test]
+async fn test_combine_entities_with_additional_entities(){
+    let stored_entities: String = r#"
+    [
+        {
+            "attrs": {},
+            "parents": [
+                {
+                    "id": "Admin",
+                    "type": "Role"
+                }
+            ],
+            "uid": {
+                "id": "admin.1@domain.com",
+                "type": "User"
+            }
+        },
+        {
+            "attrs": {},
+            "parents": [],
+            "uid": {
+                "id": "delete",
+                "type": "Action"
+            }
+        }
+    ]
+    "#
+    .to_string();
+
+    let additional_entities: String = r#"
+    [
+        {
+            "attrs": {},
+            "parents": [],
+            "uid": {
+                "id": "Admin",
+                "type": "Role"
+            }
+        }
+    ]
+    "#.to_string();
+
+    let expected_result: String = r#"
+    [
+        {
+            "attrs": {},
+            "parents": [
+                {
+                    "id": "Admin",
+                    "type": "Role"
+                }
+            ],
+            "uid": {
+                "id": "admin.1@domain.com",
+                "type": "User"
+            }
+        },
+        {
+            "attrs": {},
+            "parents": [],
+            "uid": {
+                "id": "delete",
+                "type": "Action"
+            }
+        },
+        {
+            "attrs": {},
+            "parents": [],
+            "uid": {
+                "id": "Admin",
+                "type": "Role"
+            }
+        }
+    ]
+    "#.to_string();
+
+    let query = make_authz_call_with_additional_entities(additional_entities);
+
+    match query {
+        Ok(req) => {
+            match req.get_request_entities(Entities::from_json_str(&stored_entities, None).unwrap()) {
+                Ok((_request, entities)) => {
+                    assert_eq!(entities, Entities::from_json_str(&expected_result, None).unwrap())
+                },
+                _ => assert!(false)
+            };
+        },
+        _ => assert!(false)
     };
 }
 
@@ -111,6 +202,7 @@ fn make_authz_call_no_entities() -> Result<AuthorizationRequest, Box<dyn Error>>
         None,
         None,
         None,
+        None,
     );
     return authorization_call.try_into();
 }
@@ -127,7 +219,26 @@ fn make_authz_call(entities: String) -> Result<AuthorizationRequest, Box<dyn Err
         None,
         rocket::serde::json::from_str(&entities).unwrap(),
         None,
+        None,
     );
     return authorization_call.try_into();
+}
+
+fn make_authz_call_with_additional_entities(
+    additional_entities: String
+) -> Result<AuthorizationRequest, Box<dyn Error>> {
+    let principal: Option<String> = Some("User::\"Test\"".to_string());
+    let action: Option<String> = Some("Action::\"Delete\"".to_string());
+    let resource: Option<String> = Some("Document::\"cedar-agent.pdf\"".to_string());
 
-}
\ No newline at end of file
+    let authorization_call = AuthorizationCall::new(
+        principal,
+        action,
+        resource,
+        None,
+        None,
+        rocket::serde::json::from_str(&additional_entities).unwrap(),
+        None,
+    );
+    return authorization_call.try_into();
+}