bundle: Add info about the correct rego version to parse modules on the store (#7278)

Fixing an issue where the rego-version for individual modules was lost during bundle deactivation (bundle lifecycle) if this version diverged from the active runtime rego-version. This could cause reloading of v0 bundles to fail when OPA was not running with the `--v0-compatible` flag.

Signed-off-by: Ashutosh Narkar <[email protected]>
Co-authored-by: Johan Fylling <[email protected]>
(cherry picked from commit 5d5329d)

Author: ashutosh-narkar

v1/bundle/bundle.go

@@ -67,8 +67,9 @@ type Bundle struct {
 
 // Raw contains raw bytes representing the bundle's content
 type Raw struct {
-	Path  string
-	Value []byte
+	Path   string
+	Value  []byte
+	module *ModuleFile
 }
 
 // Patch contains an array of objects wherein each object represents the patch operation to be
@@ -633,15 +634,6 @@ func (r *Reader) Read() (Bundle, error) {
 			fullPath := r.fullPath(path)
 			bs := buf.Bytes()
 
-			if r.lazyLoadingMode {
-				p := fullPath
-				if r.name != "" {
-					p = modulePathWithPrefix(r.name, fullPath)
-				}
-
-				raw = append(raw, Raw{Path: p, Value: bs})
-			}
-
 			// Modules are parsed after we've had a chance to read the manifest
 			mf := ModuleFile{
 				URL:          f.URL(),
@@ -650,6 +642,15 @@ func (r *Reader) Read() (Bundle, error) {
 				Raw:          bs,
 			}
 			modules = append(modules, mf)
+
+			if r.lazyLoadingMode {
+				p := fullPath
+				if r.name != "" {
+					p = modulePathWithPrefix(r.name, fullPath)
+				}
+
+				raw = append(raw, Raw{Path: p, Value: bs, module: &mf})
+			}
 		} else if filepath.Base(path) == WasmFile {
 			bundle.WasmModules = append(bundle.WasmModules, WasmModuleFile{
 				URL:  f.URL(),
@@ -1220,6 +1221,19 @@ func (b *Bundle) RegoVersionForFile(path string, def ast.RegoVersion) (ast.RegoV
 	return def, fmt.Errorf("unknown bundle rego-version %d for file '%s'", *version, path)
 }
 
+func (m *Manifest) RegoVersionForFile(path string) (ast.RegoVersion, error) {
+	v, err := m.numericRegoVersionForFile(path)
+	if err != nil {
+		return ast.RegoUndefined, err
+	}
+
+	if v == nil {
+		return ast.RegoUndefined, nil
+	}
+
+	return ast.RegoVersionFromInt(*v), nil
+}
+
 func (m *Manifest) numericRegoVersionForFile(path string) (*int, error) {
 	var version *int
 

v1/bundle/file.go

@@ -438,15 +438,11 @@ func (it *iterator) Next() (*storage.Update, error) {
 		for _, item := range it.raw {
 			f := file{name: item.Path}
 
-			fpath := strings.TrimLeft(normalizePath(filepath.Dir(f.name)), "/.")
-			if strings.HasSuffix(f.name, RegoExt) {
-				fpath = strings.Trim(normalizePath(f.name), "/")
+			p, err := getFileStoragePath(f.name)
+			if err != nil {
+				return nil, err
 			}
 
-			p, ok := storage.ParsePathEscaped("/" + fpath)
-			if !ok {
-				return nil, fmt.Errorf("storage path invalid: %v", f.name)
-			}
 			f.path = p
 
 			f.raw = item.Value
@@ -506,3 +502,16 @@ func getdepth(path string, isDir bool) int {
 	basePath := strings.Trim(filepath.Dir(filepath.ToSlash(path)), "/")
 	return len(strings.Split(basePath, "/"))
 }
+
+func getFileStoragePath(path string) (storage.Path, error) {
+	fpath := strings.TrimLeft(normalizePath(filepath.Dir(path)), "/.")
+	if strings.HasSuffix(path, RegoExt) {
+		fpath = strings.Trim(normalizePath(path), "/")
+	}
+
+	p, ok := storage.ParsePathEscaped("/" + fpath)
+	if !ok {
+		return nil, fmt.Errorf("storage path invalid: %v", path)
+	}
+	return p, nil
+}

v1/bundle/store.go

@@ -23,6 +23,8 @@ import (
 // BundlesBasePath is the storage path used for storing bundle metadata
 var BundlesBasePath = storage.MustParsePath("/system/bundles")
 
+var ModulesInfoBasePath = storage.MustParsePath("/system/modules")
+
 // Note: As needed these helpers could be memoized.
 
 // ManifestStoragePath is the storage path used for the given named bundle manifest.
@@ -59,6 +61,14 @@ func metadataPath(name string) storage.Path {
 	return append(BundlesBasePath, name, "manifest", "metadata")
 }
 
+func moduleRegoVersionPath(id string) storage.Path {
+	return append(ModulesInfoBasePath, strings.Trim(id, "/"), "rego_version")
+}
+
+func moduleInfoPath(id string) storage.Path {
+	return append(ModulesInfoBasePath, strings.Trim(id, "/"))
+}
+
 func read(ctx context.Context, store storage.Store, txn storage.Transaction, path storage.Path) (interface{}, error) {
 	value, err := store.Read(ctx, txn, path)
 	if err != nil {
@@ -166,6 +176,16 @@ func eraseWasmModulesFromStore(ctx context.Context, store storage.Store, txn sto
 	return suppressNotFound(err)
 }
 
+func eraseModuleRegoVersionsFromStore(ctx context.Context, store storage.Store, txn storage.Transaction, modules []string) error {
+	for _, module := range modules {
+		err := store.Write(ctx, txn, storage.RemoveOp, moduleInfoPath(module), nil)
+		if err := suppressNotFound(err); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
 // ReadWasmMetadataFromStore will read Wasm module resolver metadata from the store.
 func ReadWasmMetadataFromStore(ctx context.Context, store storage.Store, txn storage.Transaction, name string) ([]WasmResolver, error) {
 	path := wasmEntrypointsPath(name)
@@ -626,7 +646,7 @@ func eraseBundles(ctx context.Context, store storage.Store, txn storage.Transact
 		return nil, err
 	}
 
-	remaining, err := erasePolicies(ctx, store, txn, parserOpts, roots)
+	remaining, removed, err := erasePolicies(ctx, store, txn, parserOpts, roots)
 	if err != nil {
 		return nil, err
 	}
@@ -649,6 +669,11 @@ func eraseBundles(ctx context.Context, store storage.Store, txn storage.Transact
 		}
 	}
 
+	err = eraseModuleRegoVersionsFromStore(ctx, store, txn, removed)
+	if err != nil {
+		return nil, err
+	}
+
 	return remaining, nil
 }
 
@@ -668,44 +693,103 @@ func eraseData(ctx context.Context, store storage.Store, txn storage.Transaction
 	return nil
 }
 
-func erasePolicies(ctx context.Context, store storage.Store, txn storage.Transaction, parserOpts ast.ParserOptions, roots map[string]struct{}) (map[string]*ast.Module, error) {
+type moduleInfo struct {
+	RegoVersion ast.RegoVersion `json:"rego_version"`
+}
+
+func readModuleInfoFromStore(ctx context.Context, store storage.Store, txn storage.Transaction) (map[string]moduleInfo, error) {
+	value, err := read(ctx, store, txn, ModulesInfoBasePath)
+	if suppressNotFound(err) != nil {
+		return nil, err
+	}
+
+	if value == nil {
+		return nil, nil
+	}
+
+	if m, ok := value.(map[string]any); ok {
+		versions := make(map[string]moduleInfo, len(m))
+
+		for k, v := range m {
+			if m0, ok := v.(map[string]any); ok {
+				if ver, ok := m0["rego_version"]; ok {
+					if vs, ok := ver.(json.Number); ok {
+						i, err := vs.Int64()
+						if err != nil {
+							return nil, fmt.Errorf("corrupt rego version")
+						}
+						versions[k] = moduleInfo{RegoVersion: ast.RegoVersionFromInt(int(i))}
+					}
+				}
+			}
+		}
+		return versions, nil
+	}
+
+	return nil, fmt.Errorf("corrupt rego version")
+}
+
+func erasePolicies(ctx context.Context, store storage.Store, txn storage.Transaction, parserOpts ast.ParserOptions, roots map[string]struct{}) (map[string]*ast.Module, []string, error) {
 
 	ids, err := store.ListPolicies(ctx, txn)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
+	}
+
+	modulesInfo, err := readModuleInfoFromStore(ctx, store, txn)
+	if err != nil {
+		return nil, nil, fmt.Errorf("failed to read module info from store: %w", err)
+	}
+
+	getRegoVersion := func(modId string) (ast.RegoVersion, bool) {
+		info, ok := modulesInfo[modId]
+		if !ok {
+			return ast.RegoUndefined, false
+		}
+		return info.RegoVersion, true
 	}
 
 	remaining := map[string]*ast.Module{}
+	var removed []string
 
 	for _, id := range ids {
 		bs, err := store.GetPolicy(ctx, txn, id)
 		if err != nil {
-			return nil, err
+			return nil, nil, err
 		}
-		module, err := ast.ParseModuleWithOpts(id, string(bs), parserOpts)
+
+		parserOptsCpy := parserOpts
+		if regoVersion, ok := getRegoVersion(id); ok {
+			parserOptsCpy.RegoVersion = regoVersion
+		}
+
+		module, err := ast.ParseModuleWithOpts(id, string(bs), parserOptsCpy)
 		if err != nil {
-			return nil, err
+			return nil, nil, err
 		}
 		path, err := module.Package.Path.Ptr()
 		if err != nil {
-			return nil, err
+			return nil, nil, err
 		}
 		deleted := false
 		for root := range roots {
 			if RootPathsContain([]string{root}, path) {
 				if err := store.DeletePolicy(ctx, txn, id); err != nil {
-					return nil, err
+					return nil, nil, err
 				}
 				deleted = true
 				break
 			}
 		}
-		if !deleted {
+
+		if deleted {
+			removed = append(removed, id)
+		} else {
 			remaining[id] = module
 		}
 	}
 
-	return remaining, nil
+	return remaining, removed, nil
 }
 
 func writeManifestToStore(opts *ActivateOpts, name string, manifest Manifest) error {
@@ -758,6 +842,12 @@ func writeDataAndModules(ctx context.Context, store storage.Store, txn storage.T
 				if err := store.UpsertPolicy(ctx, txn, path, mf.Raw); err != nil {
 					return err
 				}
+
+				if regoVersion, err := b.RegoVersionForFile(mf.Path, ast.RegoUndefined); err == nil && regoVersion != ast.RegoUndefined {
+					if err := write(ctx, store, txn, moduleRegoVersionPath(path), regoVersion.Int()); err != nil {
+						return fmt.Errorf("failed to write rego version for '%s' in bundle '%s': %w", mf.Path, name, err)
+					}
+				}
 			}
 		} else {
 			params.BasePaths = *b.Manifest.Roots
@@ -766,6 +856,25 @@ func writeDataAndModules(ctx context.Context, store storage.Store, txn storage.T
 			if err != nil {
 				return fmt.Errorf("store truncate failed for bundle '%s': %v", name, err)
 			}
+
+			for _, f := range b.Raw {
+				if strings.HasSuffix(f.Path, RegoExt) {
+					p, err := getFileStoragePath(f.Path)
+					if err != nil {
+						return fmt.Errorf("failed get storage path for module '%s' in bundle '%s': %w", f.Path, name, err)
+					}
+
+					if m := f.module; m != nil {
+						// 'f.module.Path' contains the module's path as it relates to the bundle root, and can be used for looking up the rego-version.
+						// 'f.Path' can differ, based on how the bundle reader was initialized.
+						if regoVersion, err := b.RegoVersionForFile(f.module.Path, ast.RegoUndefined); err == nil && regoVersion != ast.RegoUndefined {
+							if err := write(ctx, store, txn, moduleRegoVersionPath(p.String()), regoVersion.Int()); err != nil {
+								return fmt.Errorf("failed to write rego version for '%s' in bundle '%s': %w", f.Path, name, err)
+							}
+						}
+					}
+				}
+			}
 		}
 	}