From 55cdee05657d10f542a2167de45dcf4a6d0a3d91 Mon Sep 17 00:00:00 2001
From: Patrick East <east.patrick@gmail.com>
Date: Wed, 4 Nov 2020 13:42:02 -0800
Subject: [PATCH] compile: Add import for removed entrypoints

In addition to removing entrypoint rules we will now inject import
statements into modules in the same package to maintain any usage
of the older rules. Previously any usage of the older rules from
the same package without using the fully qualified path would raise
an error.

Signed-off-by: Patrick East <east.patrick@gmail.com>
---
 compile/compile.go | 29 ++++++++++++++++++++++++++++-
 1 file changed, 28 insertions(+), 1 deletion(-)

diff --git a/compile/compile.go b/compile/compile.go
index 4b0bd84588..2383f1cb13 100644
--- a/compile/compile.go
+++ b/compile/compile.go
@@ -437,6 +437,10 @@ func (c *Compiler) compileWasm(ctx context.Context) error {
 
 	// Each entrypoint needs an entry in the manifest along with the
 	// original rule(s) removed from the remaining rego modules.
+
+	// For each package path keep a list of new imports. They are stored as a
+	// map to remove duplicates for each package.
+	requiredImports := map[string]map[string]*ast.Import{}
 	for i, entrypoint := range c.entrypointrefs {
 		entrypointPath := c.entrypoints[i]
 
@@ -451,8 +455,19 @@ func (c *Compiler) compileWasm(ctx context.Context) error {
 			// Drop any rules that match the entrypoint path.
 			var rules []*ast.Rule
 			for _, rule := range mf.Parsed.Rules {
-				if !rule.Path().Equal(entrypoint.Value) {
+				rulePath := rule.Path()
+				if !rulePath.Equal(entrypoint.Value) {
 					rules = append(rules, rule)
+				} else {
+					pkgPath := rule.Module.Package.Path.String()
+					newImport := &ast.Import{Path: ast.NewTerm(rulePath)}
+					if _, ok := requiredImports[pkgPath]; ok {
+						requiredImports[pkgPath][rulePath.String()] = newImport
+					} else {
+						requiredImports[pkgPath] = map[string]*ast.Import{
+							rulePath.String(): newImport,
+						}
+					}
 				}
 			}
 
@@ -466,6 +481,18 @@ func (c *Compiler) compileWasm(ctx context.Context) error {
 		}
 	}
 
+	// Any packages which had rules removed need an import injected for the
+	// removed rule to keep the policies valid.
+	for i := 0; i < len(c.bundle.Modules); i++ {
+		mf := &c.bundle.Modules[i]
+		pkgPath := mf.Parsed.Package.Path.String()
+		if imports, ok := requiredImports[pkgPath]; ok {
+			mf.Raw = nil
+			for _, newImport := range imports {
+				mf.Parsed.Imports = append(mf.Parsed.Imports, newImport)
+			}
+		}
+	}
 	return nil
 }