feat: IntelliSense and hover help updates

Signed-off-by: Kevin Hakanson <kevhak@amazon.com>

src/completion.ts

@@ -282,7 +282,12 @@ const provideCedarPeriodTriggerItems = async (
     const properties = splitPropertyChain(found[0]);
     const schemaDoc = await getSchemaTextDocument(undefined, document);
     if (schemaDoc) {
-      let entities = narrowEntityTypes(schemaDoc, properties[0]);
+      let entities = narrowEntityTypes(
+        schemaDoc,
+        properties[0],
+        document,
+        position
+      );
 
       if (properties.length === 1) {
         return createEntityTypesAttributeItems(position, schemaDoc, entities);

src/hover.ts

@@ -35,12 +35,13 @@ const getPrevNextCharacters = (
 
 const createVariableHover = async (
   document: vscode.TextDocument,
+  position: vscode.Position,
   word: string
 ): Promise<vscode.Hover | undefined> => {
   let mdarray: vscode.MarkdownString[] = [];
   const schemaDoc = await getSchemaTextDocument(undefined, document);
   if (schemaDoc) {
-    let entities = narrowEntityTypes(schemaDoc, word);
+    let entities = narrowEntityTypes(schemaDoc, word, document, position);
     entities.forEach((entityType) => {
       const md = new vscode.MarkdownString();
       md.appendCodeblock(entityType, 'cedar');
@@ -58,13 +59,19 @@ const createVariableHover = async (
 
 const createPropertyHover = async (
   document: vscode.TextDocument,
+  position: vscode.Position,
   properties: string[],
   range: vscode.Range | undefined
 ): Promise<vscode.Hover | undefined> => {
   let mdarray: vscode.MarkdownString[] = [];
   const schemaDoc = await getSchemaTextDocument(undefined, document);
   if (schemaDoc) {
-    let entities = narrowEntityTypes(schemaDoc, properties[0]);
+    let entities = narrowEntityTypes(
+      schemaDoc,
+      properties[0],
+      document,
+      position
+    );
     const completions = parseCedarSchemaDoc(schemaDoc).completions;
     let word = properties[properties.length - 1];
     entities.forEach((entityType) => {
@@ -106,7 +113,7 @@ export class CedarHoverProvider implements vscode.HoverProvider {
         ['principal', 'resource', 'context', 'action'].includes(word)
       ) {
         return new Promise(async (resolve) => {
-          let result = await createVariableHover(document, word);
+          let result = await createVariableHover(document, position, word);
           resolve(result);
         });
       }
@@ -177,7 +184,12 @@ export class CedarHoverProvider implements vscode.HoverProvider {
           const properties = splitPropertyChain(found[0]);
           return new Promise(async (resolve) => {
             let result = undefined;
-            result = await createPropertyHover(document, properties, range);
+            result = await createPropertyHover(
+              document,
+              position,
+              properties,
+              range
+            );
             resolve(result);
           });
         }

src/parser.ts

@@ -61,10 +61,17 @@ export type ReferencedRange = {
  * Cedar policies
  */
 
+type EntityTypes = {
+  principals: string[];
+  resources: string[];
+  actions: string[];
+};
+
 export type PolicyRange = {
   id: string;
   range: vscode.Range;
   effectRange: vscode.Range;
+  entityTypes: EntityTypes | undefined;
 };
 
 type PolicyCacheItem = {
@@ -176,14 +183,20 @@ export const parseCedarPoliciesDoc = (
       }
     });
 
-    if (linePreComment.trim().endsWith(';')) {
+    if (
+      // end of policy or
+      linePreComment.trim().endsWith(';') ||
+      // end of file
+      (i === cedarDoc.lineCount - 1 && !effectRange.isEqual(DEFAULT_RANGE))
+    ) {
       const policyRange = {
         id: id || `policy${count}`,
         range: new vscode.Range(
           new vscode.Position(startLine, 0),
           new vscode.Position(i, textLine.length)
         ),
         effectRange: effectRange,
+        entityTypes: undefined,
       };
       policies.push(policyRange);
 

src/test/suite/validation.test.ts

@@ -22,6 +22,7 @@ import {
 } from '../../regex';
 import * as fs from 'fs';
 import * as path from 'path';
+import { determineEntityTypes } from '../../validate';
 
 const readTestDataFile = (dirname: string, filename: string): string => {
   const filepath = path.join(process.cwd(), 'testdata', dirname, filename);
@@ -410,3 +411,59 @@ suite('Validation RegEx Test Suite', () => {
     result.free();
   });
 });
+
+suite('Validate Policy Entities Test Suite', () => {
+  const fetchEntityTypes = async (
+    head: string
+  ): Promise<{
+    principals: string[];
+    resources: string[];
+    actions: string[];
+  }> => {
+    const schemaDoc = await vscode.workspace.openTextDocument(
+      path.join(process.cwd(), 'testdata', 'narrow', 'cedarschema.json')
+    );
+
+    const principalTypes = determineEntityTypes(schemaDoc, 'principal', head);
+    const resourceTypes = determineEntityTypes(schemaDoc, 'resource', head);
+    const actionIds = determineEntityTypes(schemaDoc, 'action', head);
+
+    return Promise.resolve({
+      principals: principalTypes,
+      resources: resourceTypes,
+      actions: actionIds,
+    });
+  };
+
+  test('validate unspecified', async () => {
+    const head = `permit(principal, action, resource)`;
+    const e = await fetchEntityTypes(head);
+
+    assert.equal(e.principals.length, 2);
+    assert.equal(e.principals[0], 'NS::E1');
+    assert.equal(e.principals[1], 'NS::E2');
+
+    assert.equal(e.actions.length, 3);
+    assert.equal(e.actions[0], `NS::Action::"a"`);
+    assert.equal(e.actions[1], `NS::Action::"a1"`);
+    assert.equal(e.actions[2], `NS::Action::"a2"`);
+
+    assert.equal(e.resources.length, 2);
+    assert.equal(e.resources[0], 'NS::R1');
+    assert.equal(e.resources[1], 'NS::R2');
+  });
+
+  test('validate E1 a1 R1', async () => {
+    const head = `permit (principal in NS::E::"id", action == NS::Action::"a1", resource)`;
+    const e = await fetchEntityTypes(head);
+
+    assert.equal(e.principals.length, 1);
+    assert.equal(e.principals[0], 'NS::E1');
+
+    assert.equal(e.actions.length, 1);
+    assert.equal(e.actions[0], `NS::Action::"a1"`);
+
+    assert.equal(e.resources.length, 1);
+    assert.equal(e.resources[0], 'NS::R1');
+  });
+});

src/validate.ts

@@ -125,17 +125,61 @@ class ValidationCache {
     }
   }
 }
+const HEAD_REGEX = /(?:(permit|forbid)\s*\((.|\n)*?\))(?<!\s*(;|when|unless))/;
+export const fetchEntityTypes = (
+  schemaDoc: vscode.TextDocument,
+  cedarDoc: vscode.TextDocument,
+  position: vscode.Position
+): {
+  principals: string[];
+  resources: string[];
+  actions: string[];
+} => {
+  let head: string | undefined = undefined;
+  for (let policy of parseCedarPoliciesDoc(cedarDoc).policies) {
+    if (policy.range.contains(position)) {
+      if (policy.entityTypes === undefined) {
+        const match = cedarDoc.getText(policy.range).match(HEAD_REGEX);
+        if (match) {
+          head = match[0];
+        }
+        const principalTypes = determineEntityTypes(
+          schemaDoc,
+          'principal',
+          head
+        );
+        const resourceTypes = determineEntityTypes(schemaDoc, 'resource', head);
+        const actionIds = determineEntityTypes(schemaDoc, 'action', head);
+        policy.entityTypes = {
+          principals: principalTypes,
+          resources: resourceTypes,
+          actions: actionIds,
+        };
 
+        return policy.entityTypes;
+      } else {
+        return policy.entityTypes;
+      }
+    }
+  }
+
+  return validationCache.fetchEntityTypes(schemaDoc);
+};
 const validationCache = new ValidationCache();
 export const clearValidationCache = () => {
   validationCache.clear();
 };
 export const narrowEntityTypes = (
   schemaDoc: vscode.TextDocument,
-  scope: string
+  scope: string,
+  cedarDoc: vscode.TextDocument,
+  position: vscode.Position
 ): string[] => {
-  const { principals, resources, actions } =
-    validationCache.fetchEntityTypes(schemaDoc);
+  const { principals, resources, actions } = fetchEntityTypes(
+    schemaDoc,
+    cedarDoc,
+    position
+  );
   let entities: string[] = [];
   if (scope === 'principal') {
     entities = principals;
@@ -278,14 +322,14 @@ const ATTRIBUTE_REGEX =
   /attribute `__vscode__` in context for (?<suggestion>.+) not found/;
 export const determineEntityTypes = (
   schemaDoc: vscode.TextDocument,
-  scope: 'principal' | 'resource' | 'action'
+  scope: 'principal' | 'resource' | 'action',
+  head: string = 'permit (principal, action, resource)'
 ): string[] => {
   const types: string[] = [];
   const expr = scope === 'action' ? 'context.__vscode__' : scope;
-  const policyResult: cedar.ValidatePolicyResult = cedar.validatePolicy(
-    schemaDoc.getText(),
-    `permit (principal, action, resource) when { ${expr} };`
-  );
+  const tmpPolicy = `${head} when { ${expr} };`;
+  let policyResult: cedar.ValidatePolicyResult;
+  policyResult = cedar.validatePolicy(schemaDoc.getText(), tmpPolicy);
   if (policyResult.success === false && policyResult.errors) {
     policyResult.errors.forEach((e) => {
       let found =
@@ -299,7 +343,7 @@ export const determineEntityTypes = (
     });
   }
   policyResult.free();
-  return types;
+  return types.sort();
 };
 
 export const validateEntitiesDoc = async (

testdata/narrow/cedarschema.json

@@ -0,0 +1,105 @@
+{
+  "NS": {
+    "entityTypes": {
+      "E": {},
+      "E1": {
+        "memberOfTypes": ["E"],
+        "shape": {
+          "type": "Record",
+          "attributes": {
+            "p1": {
+              "type": "String"
+            }
+          }
+        }
+      },
+      "E2": {
+        "memberOfTypes": ["E"],
+        "shape": {
+          "type": "Record",
+          "attributes": {
+            "p1": {
+              "type": "Long"
+            }
+          }
+        }
+      },
+      "R": {},
+      "R1": {
+        "memberOfTypes": ["R"],
+        "shape": {
+          "type": "Record",
+          "attributes": {
+            "p1": {
+              "type": "String"
+            }
+          }
+        }
+      },
+      "R2": {
+        "memberOfTypes": ["R"],
+        "shape": {
+          "type": "Record",
+          "attributes": {
+            "p1": {
+              "type": "Long"
+            }
+          }
+        }
+      }
+    },
+    "actions": {
+      "as": {
+        "appliesTo": {
+          "principalTypes": [],
+          "resourceTypes": []
+        }
+      },
+      "a": {
+        "memberOf": [{ "id": "as" }],
+        "appliesTo": {
+          "principalTypes": ["E1", "E2"],
+          "resourceTypes": ["R1", "R2"],
+          "context": {
+            "type": "Record",
+            "attributes": {
+              "c1": {
+                "type": "Long"
+              }
+            }
+          }
+        }
+      },
+      "a1": {
+        "memberOf": [{ "id": "as" }],
+        "appliesTo": {
+          "principalTypes": ["E1"],
+          "resourceTypes": ["R1"],
+          "context": {
+            "type": "Record",
+            "attributes": {
+              "c1": {
+                "type": "Long"
+              }
+            }
+          }
+        }
+      },
+      "a2": {
+        "memberOf": [{ "id": "as" }],
+        "appliesTo": {
+          "principalTypes": ["E2"],
+          "resourceTypes": ["R2"],
+          "context": {
+            "type": "Record",
+            "attributes": {
+              "c1": {
+                "type": "Long"
+              }
+            }
+          }
+        }
+      }
+    }
+  }
+}