Repository: cedar-policy/cedar (rust)

Issue:
Consider options for `LosslessPolicy` when locations are not available
              This bit. Looks like if `loc` is `None` the `LosslessPolicy` will contain `<undefined>`

https://github.com/cedar-policy/cedar/blob/40be99c42cd6e102515b28db6bfb4c0b1dfc9421/cedar-policy/src/api.rs#L2234-L2243

I think what we'll need is to add a `None` variant to `LosslessPolicy` for the fast parsing path. So, you won't get lossless roundtripping if you used fast parsing.

_Originally posted by @john-h-kastner-aws in https://github.com/cedar-policy/cedar/pull/1626#discussion_r2093676167_
          
---

As John points out, we need to figure out a solution for when we don't have locations. At the moment, we agreed on `expect`ing a location.

Pull request description:
Handle missing source code info
## Description of changes

Extend our existing APIs to handle the (upcoming) case where source location information won't be available after text-based parsing. The `LosslessPolicy` struct in particular requires a new variant (`Empty`) for when we don't pass text from the parsing function that returns the text associated to each policy (source code locations are required for that).

## Issue #, if available

Resolves #1627 

## Checklist for requesting a review

The change in this PR is (choose one, and delete the other options):

- [ ] A breaking change requiring a major version bump to `cedar-policy` (e.g., changes to the signature of an existing API).
- [ ] A backwards-compatible change requiring a minor version bump to `cedar-policy` (e.g., addition of a new API).
- [ ] A bug fix or other functionality change requiring a patch to `cedar-policy`.
- [x] A change "invisible" to users (e.g., documentation, changes to "internal" crates like `cedar-policy-core`, `cedar-validator`, etc.)
- [ ] A change (breaking or otherwise) that only impacts unreleased or experimental code.

I confirm that this PR (choose one, and delete the other options):

- [ ] Updates the "Unreleased" section of the CHANGELOG with a description of my change (required for major/minor version bumps).
- [x] Does not update the CHANGELOG because my change does not significantly impact released code.

I confirm that [`cedar-spec`](https://github.com/cedar-policy/cedar-spec) (choose one, and delete the other options):

- [x] Does not require updates because my change does not impact the Cedar formal model or DRT infrastructure.
- [ ] Requires updates, and I have made / will make these updates myself. (Please include in your description a timeline or link to the relevant PR in `cedar-spec`, and how you have tested that your updates are correct.)
- [ ] Requires updates, but I do not plan to make them in the near future. (Make sure that your changes are hidden behind a feature flag to mark them as experimental.)
- [ ] I'm not sure how my change impacts `cedar-spec`. (Post your PR anyways, and we'll discuss in the comments.)

I confirm that [`docs.cedarpolicy.com`](https://docs.cedarpolicy.com/) (choose one, and delete the other options):

- [x] Does not require updates because my change does not impact the Cedar language specification.
- [ ] Requires updates, and I have made / will make these updates myself. (Please include in your description a timeline or link to the relevant PR in [`cedar-docs`](https://github.com/cedar-policy/cedar-docs). PRs should be targeted at a `staging-X.Y` branch, not `main`.)
- [ ] I'm not sure how my change impacts the documentation. (Post your PR anyways, and we'll discuss in the comments.)


Patch:
diff --git a/cedar-policy-core/src/parser.rs b/cedar-policy-core/src/parser.rs
index 0f6b2c7c..24b9429e 100644
--- a/cedar-policy-core/src/parser.rs
+++ b/cedar-policy-core/src/parser.rs
@@ -61,7 +61,7 @@ pub fn parse_policyset(text: &str) -> Result<ast::PolicySet, err::ParseErrors> {
 /// appear as a key in the returned map.
 pub fn parse_policyset_and_also_return_policy_text(
     text: &str,
-) -> Result<(HashMap<ast::PolicyID, &str>, ast::PolicySet), err::ParseErrors> {
+) -> Result<(HashMap<ast::PolicyID, Option<&str>>, ast::PolicySet), err::ParseErrors> {
     let cst = text_to_cst::parse_policies(text)?;
     let pset = cst.to_policyset()?;
     // PANIC SAFETY Shouldn't be `none` since `parse_policies()` and `to_policyset()` didn't return `Err`
@@ -77,13 +77,13 @@ pub fn parse_policyset_and_also_return_policy_text(
         .with_generated_policyids()
         .expect("shouldn't be `None` since `parse_policies` and `to_policyset` didn't return `Err`")
         .map(|(id, policy)| {
-            let loc = policy
-                .loc
-                .as_ref()
-                .expect("shouldn't be `None` since we parse with locations");
-            (id, &text[loc.start()..loc.end()])
+            if let Some(loc) = &policy.loc {
+                (id, Some(&text[loc.start()..loc.end()]))
+            } else {
+                (id, None)
+            }
         })
-        .collect::<HashMap<ast::PolicyID, &str>>();
+        .collect::<HashMap<ast::PolicyID, Option<&str>>>();
     Ok((texts, pset))
 }
 
@@ -724,17 +724,17 @@ mod tests {
         assert_eq!(texts.len(), 2);
         assert_eq!(
             texts.get(&PolicyID::from_string("policy0")),
-            Some(
-                &r#"permit(principal, action, resource)
+            Some(&Some(
+                r#"permit(principal, action, resource)
             when { principal == resource.owner };"#
-            )
+            ))
         );
         assert_eq!(
             texts.get(&PolicyID::from_string("policy1")),
-            Some(
-                &r#"forbid(principal, action == Action::"modify", resource) // a comment
+            Some(&Some(
+                r#"forbid(principal, action == Action::"modify", resource) // a comment
             when { resource . highSecurity };"#
-            )
+            ))
         );
     }
 
diff --git a/cedar-policy/src/api.rs b/cedar-policy/src/api.rs
index 97ecda94..6ce44da0 100644
--- a/cedar-policy/src/api.rs
+++ b/cedar-policy/src/api.rs
@@ -2354,7 +2354,12 @@ impl PolicySet {
         let templates = self
             .templates
             .into_iter()
-            .map(|(id, template)| template.lossless.est().map(|est| (id.into(), est)))
+            .map(|(id, template)| {
+                template
+                    .lossless
+                    .est(|| template.ast.clone().into())
+                    .map(|est| (id.into(), est))
+            })
             .collect::<Result<HashMap<_, _>, _>>()?;
         let est = est::PolicySet {
             templates,
@@ -2797,7 +2802,14 @@ impl PolicySet {
 impl std::fmt::Display for PolicySet {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         // prefer to display the lossless format
-        write!(f, "{}", self.policies().map(|p| &p.lossless).join("\n"))
+        let mut policies = self.policies().peekable();
+        while let Some(policy) = policies.next() {
+            policy.lossless.fmt(|| policy.ast.clone().into(), f)?;
+            if policies.peek().is_some() {
+                writeln!(f)?;
+            }
+        }
+        Ok(())
     }
 }
 
@@ -2822,7 +2834,7 @@ fn is_static_or_link(
         }
         None => policy
             .lossless
-            .est()
+            .est(|| policy.ast.clone().into())
             .map(|est| Either::Left((id.into(), est))),
     }
 }
@@ -2972,7 +2984,7 @@ impl Template {
         let ast = parser::parse_template(id.map(Into::into), src.as_ref())?;
         Ok(Self {
             ast,
-            lossless: LosslessPolicy::policy_or_template_text(src.as_ref()),
+            lossless: LosslessPolicy::policy_or_template_text(Some(src.as_ref())),
         })
     }
 
@@ -3133,7 +3145,7 @@ impl Template {
 
     /// Get the JSON representation of this `Template`.
     pub fn to_json(&self) -> Result<serde_json::Value, PolicyToJsonError> {
-        let est = self.lossless.est()?;
+        let est = self.lossless.est(|| self.ast.clone().into())?;
         serde_json::to_value(est).map_err(Into::into)
     }
 
@@ -3147,6 +3159,7 @@ impl Template {
     /// the `cedar-policy-formatter` crate.
     pub fn to_cedar(&self) -> String {
         match &self.lossless {
+            LosslessPolicy::Empty => self.ast.to_string(),
             LosslessPolicy::Est(_) => self.ast.to_string(),
             LosslessPolicy::Text { text, .. } => text.clone(),
         }
@@ -3164,7 +3177,7 @@ impl Template {
 impl std::fmt::Display for Template {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         // prefer to display the lossless format
-        self.lossless.fmt(f)
+        self.lossless.fmt(|| self.ast.clone().into(), f)
     }
 }
 
@@ -3477,7 +3490,7 @@ impl Policy {
         let (_, ast) = ast::Template::link_static_policy(inline_ast);
         Ok(Self {
             ast,
-            lossless: LosslessPolicy::policy_or_template_text(policy_src.as_ref()),
+            lossless: LosslessPolicy::policy_or_template_text(Some(policy_src.as_ref())),
         })
     }
 
@@ -3588,7 +3601,7 @@ impl Policy {
         #[allow(clippy::expect_used)]
         let cloned_est = self
             .lossless
-            .est()
+            .est(|| self.ast.clone().into())
             .expect("Internal error, failed to construct est.");
 
         let mapping = mapping.into_iter().map(|(k, v)| (k.0, v.0)).collect();
@@ -3636,7 +3649,7 @@ impl Policy {
     /// assert_eq!(json, Policy::from_json(None, json.clone()).unwrap().to_json().unwrap());
     /// ```
     pub fn to_json(&self) -> Result<serde_json::Value, PolicyToJsonError> {
-        let est = self.lossless.est()?;
+        let est = self.lossless.est(|| self.ast.clone().into())?;
         serde_json::to_value(est).map_err(Into::into)
     }
 
@@ -3656,6 +3669,7 @@ impl Policy {
     /// the `cedar-policy-formatter` crate.
     pub fn to_cedar(&self) -> Option<String> {
         match &self.lossless {
+            LosslessPolicy::Empty => Some(self.ast.to_string()),
             LosslessPolicy::Est(_) => Some(self.ast.to_string()),
             LosslessPolicy::Text { text, slots } => {
                 if slots.is_empty() {
@@ -3691,7 +3705,7 @@ impl Policy {
         let text = ast.to_string(); // assume that pretty-printing is faster than `est::Policy::from(ast.clone())`; is that true?
         Self {
             ast,
-            lossless: LosslessPolicy::policy_or_template_text(text),
+            lossless: LosslessPolicy::policy_or_template_text(Some(text)),
         }
     }
 }
@@ -3699,7 +3713,7 @@ impl Policy {
 impl std::fmt::Display for Policy {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         // prefer to display the lossless format
-        self.lossless.fmt(f)
+        self.lossless.fmt(|| self.ast.clone().into(), f)
     }
 }
 
@@ -3722,6 +3736,8 @@ impl FromStr for Policy {
 /// This structure can be used for static policies, linked policies, and templates.
 #[derive(Debug, Clone)]
 pub(crate) enum LosslessPolicy {
+    /// An empty representation
+    Empty,
     /// EST representation
     Est(est::Policy),
     /// Text representation
@@ -3737,16 +3753,24 @@ pub(crate) enum LosslessPolicy {
 
 impl LosslessPolicy {
     /// Create a new `LosslessPolicy` from the text of a policy or template.
-    fn policy_or_template_text(text: impl Into<String>) -> Self {
-        Self::Text {
-            text: text.into(),
-            slots: HashMap::new(),
+    fn policy_or_template_text(text: Option<impl Into<String>>) -> Self {
+        match text {
+            None => Self::Empty,
+            Some(text) => Self::Text {
+                text: text.into(),
+                slots: HashMap::new(),
+            },
         }
     }
 
-    /// Get the EST representation of this static policy, linked policy, or template
-    fn est(&self) -> Result<est::Policy, PolicyToJsonError> {
+    /// Get the EST representation of this static policy, linked policy, or template.
+    fn est(
+        &self,
+        fallback_est: impl FnOnce() -> est::Policy,
+    ) -> Result<est::Policy, PolicyToJsonError> {
         match self {
+            // Fall back to the `policy` AST if the lossless representation is empty
+            Self::Empty => Ok(fallback_est()),
             Self::Est(est) => Ok(est.clone()),
             Self::Text { text, slots } => {
                 let est =
@@ -3766,6 +3790,7 @@ impl LosslessPolicy {
         vals: impl IntoIterator<Item = (ast::SlotId, &'a ast::EntityUID)>,
     ) -> Result<Self, est::LinkingError> {
         match self {
+            Self::Empty => Ok(Self::Empty),
             Self::Est(est) => {
                 let unwrapped_est_vals: HashMap<
                     ast::SlotId,
@@ -3783,11 +3808,17 @@ impl LosslessPolicy {
             }
         }
     }
-}
 
-impl std::fmt::Display for LosslessPolicy {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+    fn fmt(
+        &self,
+        fallback_est: impl FnOnce() -> est::Policy,
+        f: &mut std::fmt::Formatter<'_>,
+    ) -> std::fmt::Result {
         match self {
+            Self::Empty => match self.est(fallback_est) {
+                Ok(est) => write!(f, "{est}"),
+                Err(e) => write!(f, "<invalid policy: {e}>"),
+            },
             Self::Est(est) => write!(f, "{est}"),
             Self::Text { text, slots } => {
                 if slots.is_empty() {
@@ -3798,7 +3829,7 @@ impl std::fmt::Display for LosslessPolicy {
                     // want to use the actual parser; right now we reuse
                     // another implementation by just converting to EST and
                     // printing that
-                    match self.est() {
+                    match self.est(fallback_est) {
                         Ok(est) => write!(f, "{est}"),
                         Err(e) => write!(f, "<invalid linked policy: {e}>"),
                     }
@@ -5341,6 +5372,59 @@ action CreateList in Create appliesTo {
     }
 }
 
+#[cfg(test)]
+mod test_lossless_empty {
+    use super::{LosslessPolicy, Policy, PolicyId, Template};
+
+    #[test]
+    fn test_lossless_empty_policy() {
+        const STATIC_POLICY_TEXT: &str = "permit(principal,action,resource);";
+        let policy0 = Policy::parse(Some(PolicyId::new("policy0")), STATIC_POLICY_TEXT)
+            .expect("Failed to parse");
+        let lossy_policy0 = Policy {
+            ast: policy0.ast.clone(),
+            lossless: LosslessPolicy::policy_or_template_text(None::<&str>),
+        };
+        // The `to_cedar` representation becomes lossy since we didn't provide text
+        assert_eq!(
+            lossy_policy0.to_cedar(),
+            Some(String::from(
+                "permit(\n  principal,\n  action,\n  resource\n) when {\n  true\n};"
+            ))
+        );
+        // The EST representation is obtained from the AST
+        let lossy_policy0_est = lossy_policy0
+            .lossless
+            .est(|| policy0.ast.clone().into())
+            .unwrap();
+        assert_eq!(lossy_policy0_est, policy0.ast.into());
+    }
+
+    #[test]
+    fn test_lossless_empty_template() {
+        const TEMPLATE_TEXT: &str = "permit(principal == ?principal,action,resource);";
+        let template0 = Template::parse(Some(PolicyId::new("template0")), TEMPLATE_TEXT)
+            .expect("Failed to parse");
+        let lossy_template0 = Template {
+            ast: template0.ast.clone(),
+            lossless: LosslessPolicy::policy_or_template_text(None::<&str>),
+        };
+        // The `to_cedar` representation becomes lossy since we didn't provide text
+        assert_eq!(
+            lossy_template0.to_cedar(),
+            String::from(
+                "permit(\n  principal == ?principal,\n  action,\n  resource\n) when {\n  true\n};"
+            )
+        );
+        // The EST representation is obtained from the AST
+        let lossy_template0_est = lossy_template0
+            .lossless
+            .est(|| template0.ast.clone().into())
+            .unwrap();
+        assert_eq!(lossy_template0_est, template0.ast.into());
+    }
+}
+
 /// Given a schema and policy set, compute an entity manifest.
 ///
 /// The policies must validate against the schema in strict mode,


Test patch:
diff --git a/cedar-policy/src/test/test.rs b/cedar-policy/src/test/test.rs
index 6b79a071..69ecb599 100644
--- a/cedar-policy/src/test/test.rs
+++ b/cedar-policy/src/test/test.rs
@@ -787,6 +787,49 @@ mod policy_set_tests {
         );
     }
 
+    #[test]
+    fn policyset_fmt_static() {
+        const STATIC_POLICY_TEXT: &str = "permit(principal,action,resource);";
+        let mut pset = PolicySet::new();
+        let policy0 = Policy::parse(Some(PolicyId::new("policy0")), STATIC_POLICY_TEXT)
+            .expect("Failed to parse");
+        let policy1 = Policy::parse(Some(PolicyId::new("policy1")), STATIC_POLICY_TEXT)
+            .expect("Failed to parse");
+        pset.add(policy0).unwrap();
+        pset.add(policy1).unwrap();
+        let policy_fmt = format!("{}", pset);
+        let mut expected_fmt = String::from(STATIC_POLICY_TEXT);
+        expected_fmt.push('\n');
+        expected_fmt.push_str(STATIC_POLICY_TEXT);
+        assert_eq!(expected_fmt, policy_fmt);
+    }
+
+    #[test]
+    fn policyset_fmt_template() {
+        const TEMPLATE_TEXT: &str = "permit(principal == ?principal,action,resource);";
+        const LINKED_POLICY_TEXT: &str = "permit(principal == Test::\"test\", action, resource);";
+        let mut pset = PolicySet::new();
+        let template0 = Template::parse(Some(PolicyId::new("template0")), TEMPLATE_TEXT)
+            .expect("Failed to parse");
+        let template1 = Template::parse(Some(PolicyId::new("template1")), TEMPLATE_TEXT)
+            .expect("Failed to parse");
+        pset.add_template(template0).unwrap();
+        pset.add_template(template1).unwrap();
+        let env0: HashMap<SlotId, EntityUid> =
+            HashMap::from([(SlotId::principal(), EntityUid::from_strs("Test", "test"))]);
+        pset.link(PolicyId::new("template0"), PolicyId::new("linked0"), env0)
+            .expect("Failed to link");
+        let env1: HashMap<SlotId, EntityUid> =
+            HashMap::from([(SlotId::principal(), EntityUid::from_strs("Test", "test"))]);
+        pset.link(PolicyId::new("template1"), PolicyId::new("linked1"), env1)
+            .expect("Failed to link");
+        let policy_fmt = format!("{}", pset);
+        let mut expected_fmt = String::from(LINKED_POLICY_TEXT);
+        expected_fmt.push('\n');
+        expected_fmt.push_str(LINKED_POLICY_TEXT);
+        assert_eq!(expected_fmt, policy_fmt);
+    }
+
     #[cfg(feature = "partial-eval")]
     #[test]
     fn unknown_entities() {