Adds annotation support for Policy

CedarJava/src/main/java/com/cedarpolicy/model/policy/Policy.java

@@ -23,7 +23,8 @@
 import com.cedarpolicy.model.Effect;
 
 import java.util.concurrent.atomic.AtomicInteger;
-
+import java.util.HashMap;
+import java.util.Map;
 
 /** Policies in the Cedar language. */
 public class Policy {
@@ -36,18 +37,18 @@ public class Policy {
     public final String policySrc;
     /** Policy ID. */
     public final String policyID;
+    /** Annotations */
+    private Map<String, String> annotations;
 
     /**
      * Creates a Cedar policy object.
      *
-     * @param policy String containing the source code of a Cedar policy in the Cedar policy
-     *     language.
-     * @param policyID The id of this policy. Must be unique. Note: We may flip the order of the
-     *     arguments here for idiomatic reasons.
+     * @param policy   String containing the source code of a Cedar policy in the Cedar policy language.
+     * @param policyID The id of this policy. Must be unique. Note: We may flip the order of the arguments here for
+     *                 idiomatic reasons.
      */
     @SuppressFBWarnings("CT_CONSTRUCTOR_THROW")
-    public Policy(
-            @JsonProperty("policySrc") String policy, @JsonProperty("policyID") String policyID)
+    public Policy(@JsonProperty("policySrc") String policy, @JsonProperty("policyID") String policyID)
             throws NullPointerException {
 
         if (policy == null) {
@@ -82,8 +83,8 @@ public String toString() {
     /**
      * Returns the effect of a policy.
      *
-     * Determines the policy effect by attempting static policy first, then template.
-     * In future, it will only support static policies once new class is introduced for Template.
+     * Determines the policy effect by attempting static policy first, then template. In future, it will only support
+     * static policies once new class is introduced for Template.
      *
      * @return The effect of the policy, either "permit" or "forbid"
      * @throws InternalException
@@ -114,17 +115,69 @@ public static Policy parseStaticPolicy(String policyStr) throws InternalExceptio
         return new Policy(policyText, null);
     }
 
-    public static Policy parsePolicyTemplate(String templateStr)  throws InternalException, NullPointerException {
+    public static Policy parsePolicyTemplate(String templateStr) throws InternalException, NullPointerException {
         String templateText = parsePolicyTemplateJni(templateStr);
         return new Policy(templateText, null);
     }
 
+    /**
+     * Gets a copy of the policy annotations map. Annotations are loaded lazily when this method is first called. Works
+     * for both static policies and templates.
+     *
+     * @return A new HashMap containing the policy's annotations. For annotations without explicit values, an empty
+     *         string ("") is used as the value
+     */
+    public Map<String, String> getAnnotations() throws InternalException {
+        ensureAnnotationsLoaded();
+        return new HashMap<>(this.annotations);
+    }
+
+    /**
+     * Gets the value of a specific annotation by its key.
+     *
+     * @param key The annotation key to look up
+     * @return The value associated with the annotation key, or null if the key doesn't exist
+     * @throws InternalException if there is an error loading or parsing the annotations
+     */
+    public String getAnnotation(String key) throws InternalException {
+        ensureAnnotationsLoaded();
+        return this.annotations.getOrDefault(key, null);
+    }
+
+    /**
+     * Ensures that the annotations map is loaded for this policy. If annotations haven't been loaded yet, attempts to
+     * load them first from static policy, then falls back to template if needed.
+     *
+     * @throws InternalException if there is an error loading or parsing the annotations
+     */
+    private void ensureAnnotationsLoaded() throws InternalException {
+        if (annotations == null) {
+            try {
+                this.annotations = getPolicyAnnotationsJni(this.policySrc);
+            } catch (InternalException e) {
+                if (e.getMessage().contains("expected a static policy")) {
+                    this.annotations = getTemplateAnnotationsJni(this.policySrc);
+                } else {
+                    throw e;
+                }
+            }
+        }
+    }
+
     private static native String parsePolicyJni(String policyStr) throws InternalException, NullPointerException;
+
     private static native String parsePolicyTemplateJni(String policyTemplateStr)
             throws InternalException, NullPointerException;
 
     private native String toJsonJni(String policyStr) throws InternalException, NullPointerException;
+
     private static native String fromJsonJni(String policyJsonStr) throws InternalException, NullPointerException;
+
     private native String policyEffectJni(String policyStr) throws InternalException, NullPointerException;
+
     private native String templateEffectJni(String policyStr) throws InternalException, NullPointerException;
+
+    private static native Map<String, String> getPolicyAnnotationsJni(String policyStr) throws InternalException;
+
+    private static native Map<String, String> getTemplateAnnotationsJni(String policyStr) throws InternalException;
 }

CedarJava/src/test/java/com/cedarpolicy/PolicyTests.java

@@ -19,15 +19,18 @@
 import com.cedarpolicy.model.exception.InternalException;
 import com.cedarpolicy.model.policy.Policy;
 import com.cedarpolicy.model.Effect;
-import org.junit.jupiter.api.Test;
 
+import org.junit.jupiter.api.Test;
 import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertNotEquals;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.junit.jupiter.api.Assertions.fail;
 
+import java.util.Map;
+import java.util.HashMap;
+
 public class PolicyTests {
     @Test
     public void parseStaticPolicyTests() {
@@ -140,4 +143,88 @@ public void policyEffectTest() throws InternalException {
         assertEquals(forbidTemplate.effect(), Effect.FORBID);
 
     }
+
+    @Test
+    public void givenStaticPolicyGetAnnotationsReturns() throws InternalException {
+        Policy staticPolicy = Policy.parseStaticPolicy("""
+            @id("policyID1")
+            @annotation("myAnnotation")
+            @emptyAnnotation
+            permit(principal, action, resource);
+            """);
+
+        Map<String, String> expectedMap = new HashMap<>();
+        expectedMap.put("id", "policyID1");
+        expectedMap.put("annotation", "myAnnotation");
+        expectedMap.put("emptyAnnotation", "");
+
+        assertEquals(staticPolicy.getAnnotations(), expectedMap);
+
+        Policy staticPolicyNoAnnotations = Policy.parseStaticPolicy("""
+            permit(principal, action, resource);
+            """);
+
+        assertEquals(staticPolicyNoAnnotations.getAnnotations(), new HashMap<>());
+    }
+
+    @Test
+    public void givenStaticPolicyGetAnnotationReturns() throws InternalException {
+        Policy staticPolicy = Policy.parseStaticPolicy("""
+            @id("policyID1")
+            @annotation("myAnnotation")
+            @emptyAnnotation
+            permit(principal, action, resource);
+            """);
+
+        assertEquals(staticPolicy.getAnnotation("annotation"), "myAnnotation");
+        assertEquals(staticPolicy.getAnnotation("emptyAnnotation"), "");
+
+        Policy staticPolicyNoAnnotations = Policy.parseStaticPolicy("""
+            permit(principal, action, resource);
+            """);
+
+        assertEquals(staticPolicyNoAnnotations.getAnnotation("invalidAnnotation"), null);
+    }
+
+    @Test
+    public void givenTemplatePolicyGetAnnotationsReturns() throws InternalException {
+        Policy templatePolicy = Policy.parsePolicyTemplate("""
+            @id("policyID1")
+            @annotation("myAnnotation")
+            @emptyAnnotation
+            permit(principal == ?principal, action, resource);
+            """);
+
+        Map<String, String> expectedMap = new HashMap<>();
+        expectedMap.put("id", "policyID1");
+        expectedMap.put("annotation", "myAnnotation");
+        expectedMap.put("emptyAnnotation", "");
+
+        assertEquals(templatePolicy.getAnnotations(), expectedMap);
+
+        Policy templatePolicyNoAnnotations = Policy.parsePolicyTemplate("""
+            permit(principal == ?principal, action, resource);
+            """);
+
+        assertEquals(templatePolicyNoAnnotations.getAnnotations(), new HashMap<>());
+    }
+
+    @Test
+    public void givenTemplatePolicyGetAnnotationReturns() throws InternalException {
+        Policy templatePolicy = Policy.parsePolicyTemplate("""
+            @id("policyID1")
+            @annotation("myAnnotation")
+            @emptyAnnotation
+            permit(principal == ?principal, action, resource);
+            """);
+
+        assertEquals(templatePolicy.getAnnotation("annotation"), "myAnnotation");
+        assertEquals(templatePolicy.getAnnotation("emptyAnnotation"), "");
+
+        Policy templatePolicyNoAnnotations = Policy.parsePolicyTemplate("""
+            permit(principal == ?principal, action, resource);
+            """);
+
+        assertEquals(templatePolicyNoAnnotations.getAnnotation("invalidAnnotation"), null);
+    }
 }