fix bugs

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

Author: imilchev

.vscode/launch.json

@@ -159,14 +159,10 @@
       "program": "${workspaceRoot}/apps/cnquery/cnquery.go",
       "cwd": "${workspaceRoot}/",
       "args": [
-        "run",
+        "scan",
         "k8s",
-        "--namespaces",
-        "default",
-        "-c",
-        "k8s.pods",
-        "--log-level",
-        "debug"
+        "--discover", "namespaces,pods",
+        "--namespaces", "default"
       ],
     },
     {
@@ -195,7 +191,7 @@
       "program": "${workspaceRoot}/apps/cnquery/cnquery.go",
       "console": "integratedTerminal",
       "args": [
-        "shell", "ssh", "[email protected]",
+        "shell", "cloudflare"
       ],
     },
     {

apps/cnquery/cmd/plugin.go

@@ -114,7 +114,11 @@ func (c *cnqueryPlugin) RunQuery(conf *run.RunQueryConfig, runtime *providers.Ru
 	}
 
 	ctx := context.Background()
-	discoveredAssets, err := scan.DiscoverAssets(ctx, conf.Inventory, upstreamConfig, runtime.Recording())
+	discoveredAssets, err := scan.DiscoverAssets(ctx, &scan.AssetDiscoveryRequest{
+		Inv:       conf.Inventory,
+		Upstream:  upstreamConfig,
+		Recording: runtime.Recording(),
+	})
 	if err != nil {
 		return err
 	}

apps/cnquery/cmd/shell.go

@@ -104,9 +104,11 @@ func StartShell(runtime *providers.Runtime, conf *ShellConfig) error {
 	// we go through inventory resolution to resolve credentials properly for the passed-in asset
 	ctx := context.Background()
 	discoveredAssets, err := scan.DiscoverAssets(ctx,
-		inventory.New(inventory.WithAssets(conf.Asset)),
-		conf.UpstreamConfig,
-		runtime.Recording())
+		&scan.AssetDiscoveryRequest{
+			Inv:       inventory.New(inventory.WithAssets(conf.Asset)),
+			Upstream:  conf.UpstreamConfig,
+			Recording: runtime.Recording(),
+		})
 	if err != nil {
 		log.Fatal().Err(err).Msg("could not process assets")
 	}

explorer/scan/discovery.go

@@ -79,9 +79,18 @@ func (d *DiscoveredAssets) GetAssetsByPlatformID(platformID string) []*AssetWith
 	return assets
 }
 
+type AssetDiscoveryRequest struct {
+	Inv       *inventory.Inventory
+	Upstream  *upstream.UpstreamConfig
+	Recording llx.Recording
+	// DiscoveredAssets *DiscoveredAssets
+	SkipRoot bool
+	Depth    uint
+}
+
 // DiscoverAssets discovers assets from the given inventory and upstream configuration. Returns only unique assets
-func DiscoverAssets(ctx context.Context, inv *inventory.Inventory, upstream *upstream.UpstreamConfig, recording llx.Recording) (*DiscoveredAssets, error) {
-	im, err := manager.NewManager(manager.WithInventory(inv, providers.DefaultRuntime()))
+func DiscoverAssets(ctx context.Context, req *AssetDiscoveryRequest) (*DiscoveredAssets, error) {
+	im, err := manager.NewManager(manager.WithInventory(req.Inv, providers.DefaultRuntime()))
 	if err != nil {
 		return nil, errors.New("failed to resolve inventory for connection: " + err.Error())
 	}
@@ -96,7 +105,7 @@ func DiscoverAssets(ctx context.Context, inv *inventory.Inventory, upstream *ups
 	// CI/CD scan and we need to apply the runtime labels to the assets
 	if runtimeEnv != nil &&
 		runtimeEnv.IsAutomatedEnv() &&
-		inv.Spec.Assets[0].Category == inventory.AssetCategory_CATEGORY_CICD {
+		req.Inv.Spec.Assets[0].Category == inventory.AssetCategory_CATEGORY_CICD {
 		runtimeLabels = runtimeEnv.Labels()
 	}
 
@@ -110,13 +119,18 @@ func DiscoverAssets(ctx context.Context, inv *inventory.Inventory, upstream *ups
 		}
 
 		// create runtime for root asset
-		rootAssetWithRuntime, err := createRuntimeForAsset(resolvedRootAsset, upstream, recording)
+		rootAssetWithRuntime, err := createRuntimeForAsset(resolvedRootAsset, req.Upstream, req.Recording)
 		if err != nil {
 			log.Error().Err(err).Str("asset", resolvedRootAsset.Name).Msg("unable to create runtime for asset")
 			discoveredAssets.AddError(resolvedRootAsset, err)
 			continue
 		}
 
+		// If the root asset is nil, then we observed an asset we have already discovered
+		if rootAssetWithRuntime == nil {
+			continue
+		}
+
 		resolvedRootAsset = rootAssetWithRuntime.Asset // to ensure we get all the information the connect call gave us
 
 		// Make sure the root runtime is closed at the end of the loop if needed. This will close the runtimes for all
@@ -145,7 +159,7 @@ func DiscoverAssets(ctx context.Context, inv *inventory.Inventory, upstream *ups
 		}
 
 		// for all discovered assets, we apply mondoo-specific labels and annotations that come from the root asset
-		discoverAssets(rootAssetWithRuntime, resolvedRootAsset, discoveredAssets, runtimeLabels, upstream, recording)
+		discoverAssets(rootAssetWithRuntime, resolvedRootAsset, discoveredAssets, runtimeLabels, req.Upstream, req.Recording)
 	}
 
 	// if there is exactly one asset, assure that the --asset-name is used
@@ -204,7 +218,9 @@ func discoverAssets(rootAssetWithRuntime *AssetWithRuntime, resolvedRootAsset *i
 			// If the asset has been already added, we should close its runtime
 			if !discoveredAssets.Add(resolvedAsset, assetWithRuntime.Runtime) {
 				assetWithRuntime.Runtime.Close()
+				continue
 			}
+
 			discoverAssets(assetWithRuntime, resolvedRootAsset, discoveredAssets, runtimeLabels, upstream, recording)
 		} else {
 			discoverAssets(assetWithRuntime, resolvedRootAsset, discoveredAssets, runtimeLabels, upstream, recording)

explorer/scan/discovery_test.go

@@ -142,7 +142,11 @@ func TestDiscoverAssets(t *testing.T) {
 
 	t.Run("normal", func(t *testing.T) {
 		inv := getInventory()
-		discoveredAssets, err := DiscoverAssets(context.Background(), inv, nil, recording.Null{})
+		discoveredAssets, err := DiscoverAssets(context.Background(), &AssetDiscoveryRequest{
+			Inv:              inv,
+			DiscoveredAssets: &DiscoveredAssets{platformIds: map[string]struct{}{}},
+			Recording:        recording.Null{},
+		})
 		require.NoError(t, err)
 		assert.Len(t, discoveredAssets.Assets, 3)
 		assert.Len(t, discoveredAssets.Errors, 0)
@@ -158,7 +162,11 @@ func TestDiscoverAssets(t *testing.T) {
 	t.Run("with duplicate root assets", func(t *testing.T) {
 		inv := getInventory()
 		inv.Spec.Assets = append(inv.Spec.Assets, inv.Spec.Assets[0])
-		discoveredAssets, err := DiscoverAssets(context.Background(), inv, nil, recording.Null{})
+		discoveredAssets, err := DiscoverAssets(context.Background(), &AssetDiscoveryRequest{
+			Inv:              inv,
+			DiscoveredAssets: &DiscoveredAssets{platformIds: map[string]struct{}{}},
+			Recording:        recording.Null{},
+		})
 		require.NoError(t, err)
 
 		// Make sure no duplicates are returned
@@ -173,7 +181,11 @@ func TestDiscoverAssets(t *testing.T) {
 	t.Run("with duplicate discovered assets", func(t *testing.T) {
 		inv := getInventory()
 		inv.Spec.Assets[0].Connections[0].Options["path"] = "./testdata/3pods_with_duplicate.yaml"
-		discoveredAssets, err := DiscoverAssets(context.Background(), inv, nil, recording.Null{})
+		discoveredAssets, err := DiscoverAssets(context.Background(), &AssetDiscoveryRequest{
+			Inv:              inv,
+			DiscoveredAssets: &DiscoveredAssets{platformIds: map[string]struct{}{}},
+			Recording:        recording.Null{},
+		})
 		require.NoError(t, err)
 
 		// Make sure no duplicates are returned
@@ -191,7 +203,11 @@ func TestDiscoverAssets(t *testing.T) {
 			"key1": "value1",
 			"key2": "value2",
 		}
-		discoveredAssets, err := DiscoverAssets(context.Background(), inv, nil, recording.Null{})
+		discoveredAssets, err := DiscoverAssets(context.Background(), &AssetDiscoveryRequest{
+			Inv:              inv,
+			DiscoveredAssets: &DiscoveredAssets{platformIds: map[string]struct{}{}},
+			Recording:        recording.Null{},
+		})
 		require.NoError(t, err)
 
 		for _, asset := range discoveredAssets.Assets {
@@ -209,7 +225,11 @@ func TestDiscoverAssets(t *testing.T) {
 	t.Run("copy root asset managedBy", func(t *testing.T) {
 		inv := getInventory()
 		inv.Spec.Assets[0].ManagedBy = "managed-by-test"
-		discoveredAssets, err := DiscoverAssets(context.Background(), inv, nil, recording.Null{})
+		discoveredAssets, err := DiscoverAssets(context.Background(), &AssetDiscoveryRequest{
+			Inv:              inv,
+			DiscoveredAssets: &DiscoveredAssets{platformIds: map[string]struct{}{}},
+			Recording:        recording.Null{},
+		})
 		require.NoError(t, err)
 
 		for _, asset := range discoveredAssets.Assets {
@@ -234,7 +254,11 @@ func TestDiscoverAssets(t *testing.T) {
 		}()
 		inv.Spec.Assets[0].Category = inventory.AssetCategory_CATEGORY_CICD
 		require.NoError(t, os.Setenv("GITHUB_ACTION", "go-test"))
-		discoveredAssets, err := DiscoverAssets(context.Background(), inv, nil, recording.Null{})
+		discoveredAssets, err := DiscoverAssets(context.Background(), &AssetDiscoveryRequest{
+			Inv:              inv,
+			DiscoveredAssets: &DiscoveredAssets{platformIds: map[string]struct{}{}},
+			Recording:        recording.Null{},
+		})
 		require.NoError(t, err)
 
 		for _, asset := range discoveredAssets.Assets {
@@ -261,7 +285,11 @@ func TestDiscoverAssets(t *testing.T) {
 		}()
 		inv.Spec.Assets[0].Category = inventory.AssetCategory_CATEGORY_CICD
 		require.NoError(t, os.Setenv("GITHUB_ACTION", "go-test"))
-		discoveredAssets, err := DiscoverAssets(context.Background(), inv, nil, recording.Null{})
+		discoveredAssets, err := DiscoverAssets(context.Background(), &AssetDiscoveryRequest{
+			Inv:              inv,
+			DiscoveredAssets: &DiscoveredAssets{platformIds: map[string]struct{}{}},
+			Recording:        recording.Null{},
+		})
 		require.NoError(t, err)
 
 		for _, asset := range discoveredAssets.Assets {
@@ -278,7 +306,11 @@ func TestDiscoverAssets(t *testing.T) {
 		inv := getInventory()
 		inv.Spec.Assets[0].Connections[0].Type = "local"
 
-		discoveredAssets, err := DiscoverAssets(context.Background(), inv, nil, recording.Null{})
+		discoveredAssets, err := DiscoverAssets(context.Background(), &AssetDiscoveryRequest{
+			Inv:              inv,
+			DiscoveredAssets: &DiscoveredAssets{platformIds: map[string]struct{}{}},
+			Recording:        recording.Null{},
+		})
 		require.NoError(t, err)
 		assert.Len(t, discoveredAssets.Assets, 1)
 

explorer/scan/local_scanner.go

@@ -208,23 +208,6 @@ func CreateProgressBar(discoveredAssets *DiscoveredAssets, disableProgressBar bo
 func (s *LocalScanner) distributeJob(job *Job, ctx context.Context, upstream *upstream.UpstreamConfig) (*explorer.ReportCollection, error) {
 	log.Info().Msgf("discover related assets for %d asset(s)", len(job.Inventory.Spec.Assets))
 
-	discoveredAssets, err := DiscoverAssets(ctx, job.Inventory, upstream, s.recording)
-	if err != nil {
-		return nil, err
-	}
-
-	// For each discovered asset, we initialize a new runtime and connect to it.
-	// Within this process, we set up a catch-all deferred function, that shuts
-	// down all runtimes, in case we exit early.
-	defer func() {
-		for _, asset := range discoveredAssets.Assets {
-			// we can call close multiple times and it will only execute once
-			if asset.Runtime != nil {
-				asset.Runtime.Close()
-			}
-		}
-	}()
-
 	// plan scan jobs
 	reporter := NewAggregateReporter()
 	if job.Bundle == nil && upstream != nil && upstream.Creds != nil {
@@ -246,18 +229,79 @@ func (s *LocalScanner) distributeJob(job *Job, ctx context.Context, upstream *up
 		reporter.AddBundle(bundle)
 	}
 
+	processedAssets := map[string]struct{}{}
+	if err := s.loopyLoop(ctx, job, processedAssets, reporter, upstream); err != nil {
+		return nil, err
+	}
+
+	return reporter.Reports(), nil
+}
+
+func (s *LocalScanner) loopyLoop(ctx context.Context, job *Job, processedAssets map[string]struct{}, reporter *AggregateReporter, upstream *upstream.UpstreamConfig) error {
+	discoveryReq := &AssetDiscoveryRequest{
+		Inv:       job.Inventory,
+		Upstream:  upstream,
+		Recording: s.recording,
+	}
+	discoveredAssets, err := DiscoverAssets(ctx, discoveryReq)
+	if err != nil {
+		return err
+	}
+
+	newAssets := &DiscoveredAssets{platformIds: processedAssets}
+	for _, asset := range discoveredAssets.Assets {
+		newAssets.Add(asset.Asset, asset.Runtime)
+	}
+
+	if len(newAssets.Assets) == 0 {
+		return nil
+	}
+
+	if err := s.scanAssets(ctx, newAssets, job, reporter, upstream); err != nil {
+		return err
+	}
+	for pid := range discoveredAssets.platformIds {
+		processedAssets[pid] = struct{}{}
+	}
+
+	for _, asset := range newAssets.Assets {
+		err := s.loopyLoop(ctx, &Job{
+			Inventory: &inventory.Inventory{
+				Spec: &inventory.InventorySpec{
+					Assets: []*inventory.Asset{asset.Asset},
+				},
+			},
+			Bundle:           job.Bundle,
+			DoRecord:         job.DoRecord,
+			QueryPackFilters: job.QueryPackFilters,
+			Props:            job.Props,
+		}, processedAssets, reporter, upstream)
+
+		// Close the runtime for the asset once we are done with it
+		if asset.Runtime != nil {
+			asset.Runtime.Close()
+		}
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+func (s *LocalScanner) scanAssets(ctx context.Context, discoveredAssets *DiscoveredAssets, job *Job, reporter *AggregateReporter, upstream *upstream.UpstreamConfig) error {
 	// if we had asset errors we want to place them into the reporter
 	for i := range discoveredAssets.Errors {
 		reporter.AddScanError(discoveredAssets.Errors[i].Asset, discoveredAssets.Errors[i].Err)
 	}
 
 	if len(discoveredAssets.Assets) == 0 {
-		return reporter.Reports(), nil
+		return nil
 	}
 
 	multiprogress, err := CreateProgressBar(discoveredAssets, s.disableProgressBar)
 	if err != nil {
-		return nil, err
+		return err
 	}
 	// start the progress bar
 	scanGroups := sync.WaitGroup{}
@@ -274,13 +318,13 @@ func (s *LocalScanner) distributeJob(job *Job, ctx context.Context, upstream *up
 	if upstream != nil && upstream.ApiEndpoint != "" && !upstream.Incognito {
 		client, err := upstream.InitClient(ctx)
 		if err != nil {
-			return nil, err
+			return err
 		}
 		spaceMrn = client.SpaceMrn
 
 		services, err = explorer.NewRemoteServices(client.ApiEndpoint, client.Plugins, client.HttpClient)
 		if err != nil {
-			return nil, err
+			return err
 		}
 	}
 
@@ -305,7 +349,7 @@ func (s *LocalScanner) distributeJob(job *Job, ctx context.Context, upstream *up
 				List:     assetsToSync,
 			})
 			if err != nil {
-				return nil, err
+				return err
 			}
 			log.Debug().Int("assets", len(resp.Details)).Msg("got assets details")
 			platformAssetMapping := make(map[string]*explorer.SynchronizeAssetsRespAssetDetail)
@@ -332,7 +376,7 @@ func (s *LocalScanner) distributeJob(job *Job, ctx context.Context, upstream *up
 					randID := "//" + explorer.SERVICE_NAME + "/" + explorer.MRN_RESOURCE_ASSET + "/" + ksuid.New().String()
 					x, err := mrn.NewMRN(randID)
 					if err != nil {
-						return nil, multierr.Wrap(err, "failed to generate a random asset MRN")
+						return multierr.Wrap(err, "failed to generate a random asset MRN")
 					}
 					asset.Mrn = x.String()
 				}
@@ -342,7 +386,7 @@ func (s *LocalScanner) distributeJob(job *Job, ctx context.Context, upstream *up
 		// if a bundle was provided check that it matches the filter, bundles can also be downloaded
 		// later therefore we do not want to stop execution here
 		if job.Bundle != nil && job.Bundle.FilterQueryPacks(job.QueryPackFilters) {
-			return nil, errors.New("all available packs filtered out. nothing to do")
+			return errors.New("all available packs filtered out. nothing to do")
 		}
 
 		wg := sync.WaitGroup{}
@@ -395,7 +439,7 @@ func (s *LocalScanner) distributeJob(job *Job, ctx context.Context, upstream *up
 		wg.Wait()
 	}
 	scanGroups.Wait()
-	return reporter.Reports(), nil
+	return nil
 }
 
 func HandleDelayedDiscovery(ctx context.Context, asset *inventory.Asset, runtime *providers.Runtime, services *explorer.Services, spaceMrn string) (*inventory.Asset, error) {