make sure resource overriding and extension works

Signed-off-by: Ivan Milchev <[email protected]>

providers-sdk/v1/resources/resources.proto

@@ -40,6 +40,7 @@ message ResourceInfo {
   string min_mondoo_version = 25;
   string defaults = 26;
   string provider = 27;
+  repeated ResourceInfo others = 29;
 }
 
 message Field {
@@ -54,4 +55,5 @@ message Field {
   bool is_implicit_resource = 24;
   bool is_embedded = 25;
   string provider = 27;
+  repeated Field others = 29;
 }

providers-sdk/v1/resources/schema.go

@@ -12,13 +12,19 @@ type ResourcesSchema interface {
 // Add another schema and return yourself. other may be nil.
 // The other schema overrides specifications in this schema, unless
 // it is trying to extend a resource whose base is already defined.
-func (s *Schema) Add(other ResourcesSchema) *Schema {
+func (s *Schema) Add(other ResourcesSchema) ResourcesSchema {
 	if other == nil {
 		return s
 	}
 
 	for k, v := range other.AllResources() {
 		if existing, ok := s.Resources[k]; ok {
+			// If neither resource is an extension, we can't merge them. We store both references.
+			if !v.IsExtension && !existing.IsExtension {
+				existing.Others = append(existing.Others, v)
+				continue
+			}
+
 			// We will merge resources into it until we find one that is not extending.
 			// Technically, this should only happen with one resource and one only,
 			// i.e. the root resource. In case they are incorrectly specified, the
@@ -56,7 +62,11 @@ func (s *Schema) Add(other ResourcesSchema) *Schema {
 				existing.Fields = map[string]*Field{}
 			}
 			for fk, fv := range v.Fields {
-				existing.Fields[fk] = fv
+				if fExisting, ok := existing.Fields[fk]; ok {
+					fExisting.Others = append(fExisting.Others, fv)
+				} else {
+					existing.Fields[fk] = fv
+				}
 			}
 		} else {
 			ri := &ResourceInfo{
@@ -89,9 +99,19 @@ func (s *Schema) Lookup(name string) *ResourceInfo {
 
 func (s *Schema) LookupField(resource string, field string) (*ResourceInfo, *Field) {
 	res := s.Lookup(resource)
-	if res == nil || res.Fields == nil {
+	if res == nil {
 		return res, nil
 	}
+
+	// If the fields don't exist in the current resource, check the other instances of it
+	if res.Fields == nil {
+		for _, o := range res.Others {
+			if o.Fields != nil && o.Fields[field] != nil {
+				res = o
+				break
+			}
+		}
+	}
 	return res, res.Fields[field]
 }
 

providers/coordinator.go

@@ -43,9 +43,9 @@ type ProvidersCoordinator interface {
 
 var BuiltinCoreID = coreconf.Config.ID
 
-var Coordinator ProvidersCoordinator = newCoordinator()
+var Coordinator ProvidersCoordinator
 
-func newCoordinator() ProvidersCoordinator {
+func newCoordinator() *coordinator {
 	c := &coordinator{
 		runningByID: map[string]*RunningProvider{},
 		runtimes:    map[string]*Runtime{},
@@ -306,6 +306,7 @@ func (c *coordinator) unsafeStartProvider(id string, update UpdateProvidersConfi
 			mp := x.Runtime.Plugin.(*mockProviderService)
 			mp.Init(x.Runtime)
 		}
+		c.schema.Add(id, x.Runtime.Schema)
 		return x.Runtime, nil
 	}
 
@@ -385,8 +386,7 @@ func (c *coordinator) unsafeStartProvider(id string, update UpdateProvidersConfi
 		gracePeriod: 3 * time.Second,
 	}
 
-	c.schema.Add(provider.Name, provider.Schema)
-	c.schema.prioritizeIDs(BuiltinCoreID)
+	c.schema.Add(provider.ID, provider.Schema)
 
 	if err := res.heartbeat(); err != nil {
 		return nil, err

providers/extensible_schema.go

@@ -8,7 +8,6 @@ import (
 
 	"github.com/rs/zerolog/log"
 	"go.mondoo.com/cnquery/v10/providers-sdk/v1/resources"
-	"golang.org/x/exp/slices"
 )
 
 type ExtensibleSchema interface {
@@ -39,7 +38,7 @@ func newExtensibleSchema() extensibleSchema {
 	}
 }
 
-func (x *extensibleSchema) Add(name string, schema *resources.Schema) {
+func (x *extensibleSchema) Add(name string, schema resources.ResourcesSchema) {
 	x.sync.Lock()
 	x.unsafeAdd(name, schema)
 	x.unsafeRefresh()
@@ -101,31 +100,19 @@ func (x *extensibleSchema) LookupField(resource string, field string) (*resource
 	x.sync.Lock()
 	defer x.sync.Unlock()
 
-	found, ok := x.roAggregate.Resources[resource]
-	if !ok {
-		if x.lastRefreshed >= LastProviderInstall {
-			return nil, nil
-		}
-
-		found, ok = x.roAggregate.Resources[resource]
-		if !ok {
-			return nil, nil
-		}
-		return found, found.Fields[field]
+	res, f := x.roAggregate.LookupField(resource, field)
+	if res != nil && f != nil {
+		return res, f
 	}
 
-	fieldObj, ok := found.Fields[field]
-	if ok {
-		return found, fieldObj
-	}
 	if x.lastRefreshed >= LastProviderInstall {
-		return found, nil
+		return nil, nil
 	}
 
 	x.unsafeLoadAll()
 	x.unsafeRefresh()
 
-	return found, found.Fields[field]
+	return x.roAggregate.LookupField(resource, field)
 }
 
 // Prioritize the provider IDs in the order that is provided. Any other
@@ -185,18 +172,16 @@ func (x *extensibleSchema) unsafeRefresh() {
 	res := resources.Schema{
 		Resources: map[string]*resources.ResourceInfo{},
 	}
-	for id, schema := range x.loaded {
-		if !slices.Contains(x.prioritization, id) {
-			res.Add(schema)
-		}
+	for _, schema := range x.loaded {
+		res.Add(schema)
 	}
 
-	for i := len(x.prioritization) - 1; i >= 0; i-- {
-		id := x.prioritization[i]
-		if s := x.loaded[id]; s != nil {
-			res.Add(s)
-		}
-	}
+	// for i := len(x.prioritization) - 1; i >= 0; i-- {
+	// 	id := x.prioritization[i]
+	// 	if s := x.loaded[id]; s != nil {
+	// 		res.Add(s)
+	// 	}
+	// }
 
 	// Note: This object is read-only and thus must be re-created to
 	// prevent concurrency issues with access outside this struct

providers/providers.go

@@ -58,6 +58,26 @@ func init() {
 	}
 
 	LastProviderInstall = time.Now().Unix()
+
+	// Initialize the global coordinator instance
+	coordinator := newCoordinator()
+	Coordinator = coordinator
+
+	// Load the schema for all the active providers
+	providers, err := ListActive()
+	if err != nil {
+		log.Error().Err(err).Msg("failed to list active providers")
+		return
+	}
+
+	for name := range providers {
+		schema, err := Coordinator.LoadSchema(name)
+		if err != nil {
+			log.Error().Err(err).Str("provider", name).Msg("failed to load schema")
+			continue
+		}
+		coordinator.schema.Add(name, schema)
+	}
 }
 
 type ProviderLookup struct {

providers/runtime.go

@@ -536,6 +536,22 @@ func (r *Runtime) lookupResourceProvider(resource string) (*ConnectedProvider, *
 		return nil, nil, errors.New("cannot find resource '" + resource + "' in schema")
 	}
 
+	// prioritize ids
+	resourcesPerProvider := map[string]*resources.ResourceInfo{
+		info.Provider: info,
+	}
+	for _, other := range info.Others {
+		resourcesPerProvider[other.Provider] = other
+	}
+
+	priority := []string{BuiltinCoreID, r.Provider.Instance.ID}
+	for i := len(priority) - 1; i >= 0; i-- {
+		id := priority[i]
+		if s := resourcesPerProvider[id]; s != nil {
+			info = s
+		}
+	}
+
 	if info.Provider == "" {
 		// This case happens when the resource is only bridging a resource chain,
 		// i.e. it is extending in nature (which we only test for the warning).
@@ -597,6 +613,14 @@ func (r *Runtime) lookupFieldProvider(resource string, field string) (*Connected
 		return nil, nil, nil, errors.New("cannot find field '" + field + "' in resource '" + resource + "'")
 	}
 
+	// Make sure we grab the field that matches the provider of this runtime (if possible).
+	for _, f := range fieldInfo.Others {
+		if f.Provider == r.Provider.Instance.ID {
+			fieldInfo = f
+			break
+		}
+	}
+
 	if provider := r.providers[fieldInfo.Provider]; provider != nil {
 		return provider, resourceInfo, fieldInfo, provider.ConnectionError
 	}