🧹 add agent info in upstream requests

Signed-off-by: Ivan Milchev <[email protected]>
mondoohq · Mar 8, 2024 · 3c0ca99 · 3c0ca99
1 parent 067a377
commit 3c0ca99
Show file tree

Hide file tree

Showing 5 changed files with 15 additions and 14 deletions.
diff --git a/apps/cnspec/cmd/serve.go b/apps/cnspec/cmd/serve.go
@@ -73,7 +73,7 @@ var serveCmd = &cobra.Command{
 		ctx := cnquery.SetFeatures(context.Background(), cnquery.DefaultFeatures)
 
 		if scanConf != nil && scanConf.runtime.UpstreamConfig != nil {
-			client, err := scanConf.runtime.UpstreamConfig.InitClient()
+			client, err := scanConf.runtime.UpstreamConfig.InitClient(ctx)
 			if err != nil {
 				return cli_errors.NewCommandError(errors.Wrap(err, "could not initialize upstream client"), 1)
 			}

diff --git a/apps/cnspec/cmd/vuln.go b/apps/cnspec/cmd/vuln.go
@@ -4,6 +4,7 @@
 package cmd
 
 import (
+	"context"
 	"encoding/json"
 	"strings"
 
@@ -123,7 +124,7 @@ var vulnCmdRun = func(cmd *cobra.Command, runtime *providers.Runtime, cliRes *pl
 		return
 	}
 
-	client, err := runtime.UpstreamConfig.InitClient()
+	client, err := runtime.UpstreamConfig.InitClient(context.Background())
 	if err != nil {
 		if status, ok := status.FromError(err); ok {
 			code := status.Code()

diff --git a/go.mod b/go.mod
@@ -34,7 +34,7 @@ require (
 	github.com/spf13/pflag v1.0.6-0.20201009195203-85dd5c8bc61c
 	github.com/spf13/viper v1.18.2
 	github.com/stretchr/testify v1.9.0
-	go.mondoo.com/cnquery/v10 v10.6.1
+	go.mondoo.com/cnquery/v10 v10.6.2-0.20240308154801-22bff171192f
 	go.mondoo.com/mondoo-go v0.0.0-20240307081130-83b805182b6b
 	go.mondoo.com/ranger-rpc v0.6.0
 	go.opentelemetry.io/otel v1.24.0

diff --git a/go.sum b/go.sum
@@ -1189,8 +1189,8 @@ go-simpler.org/sloglint v0.4.0/go.mod h1:v6zJ++j/thFPhefs2wEXoCKwT10yo5nkBDYRCXy
 go.etcd.io/etcd/api/v3 v3.5.1/go.mod h1:cbVKeC6lCfl7j/8jBhAK6aIYO9XOjdptoxU/nLQcPvs=
 go.etcd.io/etcd/client/pkg/v3 v3.5.1/go.mod h1:IJHfcCEKxYu1Os13ZdwCwIUTUVGYTSAM3YSwc9/Ac1g=
 go.etcd.io/etcd/client/v2 v2.305.1/go.mod h1:pMEacxZW7o8pg4CrFE7pquyCJJzZvkvdD2RibOCCCGs=
-go.mondoo.com/cnquery/v10 v10.6.1 h1:2p9ZXgsh5khmXQb8dQ8Ehl6BAqT2raN1gQvVBzqbDWs=
-go.mondoo.com/cnquery/v10 v10.6.1/go.mod h1:cH1ISKlLU+LAVJbMjol4TBOSI+NqR3QPWTrotqs8fbg=
+go.mondoo.com/cnquery/v10 v10.6.2-0.20240308154801-22bff171192f h1:2wGSYK4D1Nhp+llLScoI2jXxhx5P1nJmhRKh7XVlfWQ=
+go.mondoo.com/cnquery/v10 v10.6.2-0.20240308154801-22bff171192f/go.mod h1:cH1ISKlLU+LAVJbMjol4TBOSI+NqR3QPWTrotqs8fbg=
 go.mondoo.com/mondoo-go v0.0.0-20240307081130-83b805182b6b h1:MSUQyChh5C21eThnx0h684PqA3gbumiZ9KG5ymiLGZc=
 go.mondoo.com/mondoo-go v0.0.0-20240307081130-83b805182b6b/go.mod h1:JDTzAPQ3y9bVoDOBrkhcd1zFKHKjmeuJ0jZRTeQ46QA=
 go.mondoo.com/ranger-rpc v0.6.0 h1:u8FY0NGhR1B/LZk9+DdmqksrWc007argv+nMGJGa3OA=

diff --git a/policy/scan/local_scanner.go b/policy/scan/local_scanner.go
@@ -237,7 +237,7 @@ func createReporter(ctx context.Context, job *Job, upstream *upstream.UpstreamCo
 		// if we use upstream with/without incognito, we want to fetch the bundle here to ensure we only fetch it once
 		// for all assets in the same space
 		if upstream != nil && upstream.Creds != nil {
-			client, err := upstream.InitClient()
+			client, err := upstream.InitClient(ctx)
 			if err != nil {
 				return nil, err
 			}
@@ -324,7 +324,7 @@ func (s *LocalScanner) distributeJob(job *Job, ctx context.Context, upstream *up
 	spaceMrn := ""
 	var services *policy.Services
 	if upstream != nil && upstream.ApiEndpoint != "" && !upstream.Incognito {
-		client, err := upstream.InitClient()
+		client, err := upstream.InitClient(ctx)
 		if err != nil {
 			return nil, err
 		}
@@ -478,14 +478,14 @@ func handleDelayedDiscovery(ctx context.Context, asset *inventory.Asset, runtime
 	return asset, nil
 }
 
-func (s *LocalScanner) upstreamServices(conf *upstream.UpstreamConfig) *policy.Services {
+func (s *LocalScanner) upstreamServices(ctx context.Context, conf *upstream.UpstreamConfig) *policy.Services {
 	if conf == nil ||
 		conf.ApiEndpoint == "" ||
 		conf.Incognito {
 		return nil
 	}
 
-	client, err := s.upstreamClient(conf)
+	client, err := s.upstreamClient(ctx, conf)
 	if err != nil {
 		log.Error().Err(err).Msg("could not init upstream client")
 		return nil
@@ -513,7 +513,7 @@ func (s *LocalScanner) RunAssetJob(job *AssetJob) {
 
 	job.Reporter.AddReport(job.Asset, results)
 
-	upstream := s.upstreamServices(job.UpstreamConfig)
+	upstream := s.upstreamServices(job.Ctx, job.UpstreamConfig)
 	if upstream != nil {
 		// get new gql client
 		mondooClient, err := gql.NewClient(job.UpstreamConfig, s._upstreamClient.HttpClient)
@@ -536,12 +536,12 @@ func (s *LocalScanner) RunAssetJob(job *AssetJob) {
 	}
 }
 
-func (s *LocalScanner) upstreamClient(conf *upstream.UpstreamConfig) (*upstream.UpstreamClient, error) {
+func (s *LocalScanner) upstreamClient(ctx context.Context, conf *upstream.UpstreamConfig) (*upstream.UpstreamClient, error) {
 	if s._upstreamClient != nil {
 		return s._upstreamClient, nil
 	}
 
-	client, err := conf.InitClient()
+	client, err := conf.InitClient(ctx)
 	if err != nil {
 		return nil, err
 	}
@@ -557,7 +557,7 @@ func (s *LocalScanner) runMotorizedAsset(job *AssetJob) (*AssetReport, error) {
 	runtimeErr := inmemory.WithDb(s.runtime, s.resolvedPolicyCache, func(db *inmemory.Db, services *policy.LocalServices) error {
 		if job.UpstreamConfig.ApiEndpoint != "" && !job.UpstreamConfig.Incognito {
 			log.Debug().Msg("using API endpoint " + job.UpstreamConfig.ApiEndpoint)
-			client, err := s.upstreamClient(job.UpstreamConfig)
+			client, err := s.upstreamClient(job.Ctx, job.UpstreamConfig)
 			if err != nil {
 				return err
 			}
@@ -631,7 +631,7 @@ func (s *LocalScanner) GarbageCollectAssets(ctx context.Context, garbageCollectO
 		return nil, status.Errorf(codes.Internal, "missing upstream config in service")
 	}
 
-	client, err := s.upstreamClient(s.upstream)
+	client, err := s.upstreamClient(ctx, s.upstream)
 	if err != nil {
 		return nil, err
 	}