use globally unique connection IDs for providers

Signed-off-by: Ivan Milchev <[email protected]>
mondoohq · Mar 5, 2024 · ee0903f · ee0903f
1 parent 6cb42a3
commit ee0903f
Show file tree

Hide file tree

Showing 26 changed files with 89 additions and 24 deletions.
diff --git a/providers-sdk/v1/inventory/inventory.go b/providers-sdk/v1/inventory/inventory.go
@@ -417,7 +417,7 @@ func (cfg *Config) Clone(opts ...CloneOption) *Config {
 	}
 
 	clonedObject := proto.Clone(cfg).(*Config)
-
+	clonedObject.Id = 0
 	if cloneSettings.noDiscovery {
 		clonedObject.Discover = &Discovery{}
 	}

diff --git a/providers-sdk/v1/plugin/service.go b/providers-sdk/v1/plugin/service.go
@@ -12,6 +12,7 @@ import (
 	"time"
 
 	llx "go.mondoo.com/cnquery/v10/llx"
+	inventory "go.mondoo.com/cnquery/v10/providers-sdk/v1/inventory"
 )
 
 type Service struct {
@@ -31,7 +32,46 @@ func NewService() *Service {
 
 var heartbeatRes HeartbeatRes
 
-func (s *Service) AddRuntime(createRuntime func(connId uint32) (*Runtime, error)) (*Runtime, error) {
+// FIXME: once we move to v12, remove the conf parametrer and remove the connId from the createRuntime function.
+// The connection ID will always be set before the connection call is done, so we don't need to do anything about it here.
+// The parameters are needed now, only to make sure that old clients can work with new providers.
+func (s *Service) AddRuntime(conf *inventory.Config, createRuntime func(connId uint32) (*Runtime, error)) (*Runtime, error) {
+	// FIXME: DEPRECATED, remove in v12.0 vv
+	// This approach is used only when old clients use new providers. We will throw it away in v12
+	if conf.Id == 0 {
+		return s.deprecatedAddRuntime(createRuntime)
+	}
+	// ^^
+
+	s.runtimesLock.Lock()
+	defer s.runtimesLock.Unlock()
+
+	// If a runtime with this ID already exists, then return that
+	if runtime, ok := s.runtimes[conf.Id]; ok {
+		return runtime, nil
+	}
+
+	runtime, err := createRuntime(conf.Id)
+	if err != nil {
+		return nil, err
+	}
+
+	if runtime.Connection != nil {
+		if parentId := runtime.Connection.ParentID(); parentId > 0 {
+			parentRuntime, err := s.doGetRuntime(parentId)
+			if err != nil {
+				return nil, errors.New("parent connection " + strconv.FormatUint(uint64(parentId), 10) + " not found")
+			}
+			runtime.Resources = parentRuntime.Resources
+
+		}
+	}
+	s.runtimes[conf.Id] = runtime
+	return runtime, nil
+}
+
+// FIXME: DEPRECATED, remove in v12.0 vv
+func (s *Service) deprecatedAddRuntime(createRuntime func(connId uint32) (*Runtime, error)) (*Runtime, error) {
 	s.runtimesLock.Lock()
 	defer s.runtimesLock.Unlock()
 
@@ -57,6 +97,8 @@ func (s *Service) AddRuntime(createRuntime func(connId uint32) (*Runtime, error)
 	return runtime, nil
 }
 
+// ^^
+
 func (s *Service) GetRuntime(id uint32) (*Runtime, error) {
 	s.runtimesLock.Lock()
 	defer s.runtimesLock.Unlock()

diff --git a/providers/arista/provider/provider.go b/providers/arista/provider/provider.go
@@ -124,7 +124,7 @@ func (s *Service) connect(req *plugin.ConnectReq, callback plugin.ProviderCallba
 	asset := req.Asset
 	conf := asset.Connections[0]
 
-	runtime, err := s.AddRuntime(func(connId uint32) (*plugin.Runtime, error) {
+	runtime, err := s.AddRuntime(conf, func(connId uint32) (*plugin.Runtime, error) {
 		conn, err := connection.NewAristaConnection(connId, asset, conf)
 		if err != nil {
 			return nil, err

diff --git a/providers/atlassian/provider/provider.go b/providers/atlassian/provider/provider.go
@@ -137,7 +137,7 @@ func (s *Service) connect(req *plugin.ConnectReq, callback plugin.ProviderCallba
 
 	asset := req.Asset
 	conf := asset.Connections[0]
-	runtime, err := s.AddRuntime(func(connId uint32) (*plugin.Runtime, error) {
+	runtime, err := s.AddRuntime(conf, func(connId uint32) (*plugin.Runtime, error) {
 		conn, err := connection.NewConnection(connId, asset, conf)
 		if err != nil {
 			return nil, err

diff --git a/providers/aws/provider/provider.go b/providers/aws/provider/provider.go
@@ -186,7 +186,7 @@ func (s *Service) connect(req *plugin.ConnectReq, callback plugin.ProviderCallba
 	asset := req.Asset
 	conf := asset.Connections[0]
 
-	runtime, err := s.AddRuntime(func(connId uint32) (*plugin.Runtime, error) {
+	runtime, err := s.AddRuntime(conf, func(connId uint32) (*plugin.Runtime, error) {
 		var conn shared.Connection
 		var err error
 

diff --git a/providers/azure/provider/provider.go b/providers/azure/provider/provider.go
@@ -192,7 +192,7 @@ func (s *Service) connect(req *plugin.ConnectReq, callback plugin.ProviderCallba
 	asset := req.Asset
 	conf := asset.Connections[0]
 
-	runtime, err := s.AddRuntime(func(connId uint32) (*plugin.Runtime, error) {
+	runtime, err := s.AddRuntime(conf, func(connId uint32) (*plugin.Runtime, error) {
 		var conn shared.AzureConnection
 		var err error
 

diff --git a/providers/coordinator.go b/providers/coordinator.go
@@ -29,6 +29,7 @@ import (
 //go:generate mockgen -source=../providers-sdk/v1/resources/schema.go -destination=./mock_schema.go -package=providers
 
 type ProvidersCoordinator interface {
+	NextConnectionId() uint32
 	NewRuntime() *Runtime
 	NewRuntimeFrom(parent *Runtime) *Runtime
 	RuntimeFor(asset *inventory.Asset, parent *Runtime) (*Runtime, error)
@@ -57,6 +58,9 @@ func newCoordinator() *coordinator {
 }
 
 type coordinator struct {
+	lastConnectionID uint32
+	connectionsLock  sync.Mutex
+
 	providers   Providers
 	runningByID map[string]*RunningProvider
 
@@ -88,6 +92,13 @@ type ProviderVersion struct {
 	Version string `json:"version"`
 }
 
+func (c *coordinator) NextConnectionId() uint32 {
+	c.connectionsLock.Lock()
+	defer c.connectionsLock.Unlock()
+	c.lastConnectionID++
+	return c.lastConnectionID
+}
+
 func (c *coordinator) tryProviderUpdate(provider *Provider, update UpdateProvidersConfig) (*Provider, error) {
 	if provider.Path == "" {
 		return nil, errors.New("cannot determine installation path for provider")
@@ -269,6 +280,13 @@ func (c *coordinator) RemoveRuntime(runtime *Runtime) {
 			}
 		}
 	}
+
+	// If all providers have been killed, reset the connection IDs back to 0
+	if len(c.runningByID) == 0 {
+		c.connectionsLock.Lock()
+		defer c.connectionsLock.Unlock()
+		c.lastConnectionID = 0
+	}
 }
 
 func (c *coordinator) GetRunningProvider(id string, update UpdateProvidersConfig) (*RunningProvider, error) {

diff --git a/providers/core/provider/provider.go b/providers/core/provider/provider.go
@@ -35,7 +35,7 @@ func (s *Service) Connect(req *plugin.ConnectReq, callback plugin.ProviderCallba
 	}
 
 	connectionId := defaultConnection
-	runtime, err := s.AddRuntime(func(connId uint32) (*plugin.Runtime, error) {
+	runtime, err := s.AddRuntime(req.Asset.Connections[0], func(connId uint32) (*plugin.Runtime, error) {
 		connectionId = connId
 		var upstream *upstream.UpstreamClient
 		var err error

diff --git a/providers/equinix/provider/provider.go b/providers/equinix/provider/provider.go
@@ -108,7 +108,7 @@ func (s *Service) connect(req *plugin.ConnectReq, callback plugin.ProviderCallba
 
 	asset := req.Asset
 	conf := asset.Connections[0]
-	runtime, err := s.AddRuntime(func(connId uint32) (*plugin.Runtime, error) {
+	runtime, err := s.AddRuntime(conf, func(connId uint32) (*plugin.Runtime, error) {
 		conn, err := connection.NewEquinixConnection(connId, asset, conf)
 		if err != nil {
 			return nil, err

diff --git a/providers/gcp/provider/provider.go b/providers/gcp/provider/provider.go
@@ -209,7 +209,7 @@ func (s *Service) connect(req *plugin.ConnectReq, callback plugin.ProviderCallba
 	asset := req.Asset
 	conf := asset.Connections[0]
 
-	runtime, err := s.AddRuntime(func(connId uint32) (*plugin.Runtime, error) {
+	runtime, err := s.AddRuntime(conf, func(connId uint32) (*plugin.Runtime, error) {
 		var conn shared.GcpConnection
 		var err error
 

diff --git a/providers/github/provider/provider.go b/providers/github/provider/provider.go
@@ -127,7 +127,7 @@ func (s *Service) connect(req *plugin.ConnectReq, callback plugin.ProviderCallba
 	}
 
 	asset := req.Asset
-	runtime, err := s.AddRuntime(func(connId uint32) (*plugin.Runtime, error) {
+	runtime, err := s.AddRuntime(asset.Connections[0], func(connId uint32) (*plugin.Runtime, error) {
 		conn, err := connection.NewGithubConnection(connId, asset)
 		if err != nil {
 			return nil, err

diff --git a/providers/gitlab/provider/provider.go b/providers/gitlab/provider/provider.go
@@ -141,7 +141,7 @@ func (s *Service) connect(req *plugin.ConnectReq, callback plugin.ProviderCallba
 
 	asset := req.Asset
 	conf := asset.Connections[0]
-	runtime, err := s.AddRuntime(func(connId uint32) (*plugin.Runtime, error) {
+	runtime, err := s.AddRuntime(conf, func(connId uint32) (*plugin.Runtime, error) {
 		conn, err := connection.NewGitLabConnection(connId, asset, conf)
 		if err != nil {
 			return nil, err

diff --git a/providers/google-workspace/provider/provider.go b/providers/google-workspace/provider/provider.go
@@ -159,7 +159,7 @@ func (s *Service) connect(req *plugin.ConnectReq, callback plugin.ProviderCallba
 
 	asset := req.Asset
 	conf := asset.Connections[0]
-	runtime, err := s.AddRuntime(func(connId uint32) (*plugin.Runtime, error) {
+	runtime, err := s.AddRuntime(conf, func(connId uint32) (*plugin.Runtime, error) {
 		conn, err := connection.NewGoogleWorkspaceConnection(connId, asset, conf)
 		if err != nil {
 			return nil, err

diff --git a/providers/ipmi/provider/provider.go b/providers/ipmi/provider/provider.go
@@ -118,7 +118,7 @@ func (s *Service) connect(req *plugin.ConnectReq, callback plugin.ProviderCallba
 
 	asset := req.Asset
 	conf := asset.Connections[0]
-	runtime, err := s.AddRuntime(func(connId uint32) (*plugin.Runtime, error) {
+	runtime, err := s.AddRuntime(conf, func(connId uint32) (*plugin.Runtime, error) {
 		conn, err := connection.NewIpmiConnection(connId, asset, conf)
 		if err != nil {
 			return nil, err

diff --git a/providers/k8s/provider/provider.go b/providers/k8s/provider/provider.go
@@ -137,7 +137,7 @@ func (s *Service) connect(req *plugin.ConnectReq, callback plugin.ProviderCallba
 	asset := req.Asset
 	conf := asset.Connections[0]
 
-	runtime, err := s.AddRuntime(func(connId uint32) (*plugin.Runtime, error) {
+	runtime, err := s.AddRuntime(conf, func(connId uint32) (*plugin.Runtime, error) {
 		var conn shared.Connection
 		var err error
 		if manifestContent, ok := conf.Options[shared.OPTION_IMMEMORY_CONTENT]; ok {

diff --git a/providers/ms365/provider/provider.go b/providers/ms365/provider/provider.go
@@ -115,7 +115,7 @@ func (s *Service) connect(req *plugin.ConnectReq, callback plugin.ProviderCallba
 
 	asset := req.Asset
 	conf := asset.Connections[0]
-	runtime, err := s.AddRuntime(func(connId uint32) (*plugin.Runtime, error) {
+	runtime, err := s.AddRuntime(conf, func(connId uint32) (*plugin.Runtime, error) {
 		conn, err := connection.NewMs365Connection(connId, asset, conf)
 		if err != nil {
 			return nil, err

diff --git a/providers/network/provider/provider.go b/providers/network/provider/provider.go
@@ -137,7 +137,7 @@ func (s *Service) connect(req *plugin.ConnectReq, callback plugin.ProviderCallba
 	asset := req.Asset
 	conf := asset.Connections[0]
 
-	runtime, err := s.AddRuntime(func(connId uint32) (*plugin.Runtime, error) {
+	runtime, err := s.AddRuntime(conf, func(connId uint32) (*plugin.Runtime, error) {
 		var conn *connection.HostConnection
 
 		switch conf.Type {

diff --git a/providers/oci/provider/provider.go b/providers/oci/provider/provider.go
@@ -137,7 +137,7 @@ func (s *Service) connect(req *plugin.ConnectReq, callback plugin.ProviderCallba
 	asset := req.Asset
 	conf := asset.Connections[0]
 
-	runtime, err := s.AddRuntime(func(connId uint32) (*plugin.Runtime, error) {
+	runtime, err := s.AddRuntime(conf, func(connId uint32) (*plugin.Runtime, error) {
 		conn, err := connection.NewOciConnection(connId, asset, conf)
 		if err != nil {
 			return nil, err

diff --git a/providers/okta/provider/provider.go b/providers/okta/provider/provider.go
@@ -117,7 +117,7 @@ func (s *Service) connect(req *plugin.ConnectReq, callback plugin.ProviderCallba
 
 	asset := req.Asset
 	conf := asset.Connections[0]
-	runtime, err := s.AddRuntime(func(connId uint32) (*plugin.Runtime, error) {
+	runtime, err := s.AddRuntime(conf, func(connId uint32) (*plugin.Runtime, error) {
 		conn, err := connection.NewOktaConnection(connId, asset, conf)
 		if err != nil {
 			return nil, err

diff --git a/providers/opcua/provider/provider.go b/providers/opcua/provider/provider.go
@@ -91,7 +91,7 @@ func (s *Service) connect(req *plugin.ConnectReq, callback plugin.ProviderCallba
 
 	asset := req.Asset
 	conf := asset.Connections[0]
-	runtime, err := s.AddRuntime(func(connId uint32) (*plugin.Runtime, error) {
+	runtime, err := s.AddRuntime(conf, func(connId uint32) (*plugin.Runtime, error) {
 		conn, err := connection.NewOpcuaConnection(connId, asset, conf)
 		if err != nil {
 			return nil, err

diff --git a/providers/os/provider/provider.go b/providers/os/provider/provider.go
@@ -292,7 +292,7 @@ func (s *Service) connect(req *plugin.ConnectReq, callback plugin.ProviderCallba
 	asset := req.Asset
 	conf := asset.Connections[0]
 
-	runtime, err := s.AddRuntime(func(connId uint32) (*plugin.Runtime, error) {
+	runtime, err := s.AddRuntime(conf, func(connId uint32) (*plugin.Runtime, error) {
 		var conn shared.Connection
 		var err error
 

diff --git a/providers/runtime.go b/providers/runtime.go
@@ -193,6 +193,11 @@ func (r *Runtime) Connect(req *plugin.ConnectReq) error {
 		return errors.New("cannot connect to asset, no connection info provided")
 	}
 
+	// If there is no connection ID set, we need to assign one from the coordinator
+	if asset.Connections[0].Id == 0 {
+		asset.Connections[0].Id = Coordinator.NextConnectionId()
+	}
+
 	r.features = req.Features
 	callbacks := providerCallbacks{
 		runtime: r,

diff --git a/providers/slack/provider/provider.go b/providers/slack/provider/provider.go
@@ -119,7 +119,7 @@ func (s *Service) connect(req *plugin.ConnectReq, callback plugin.ProviderCallba
 
 	asset := req.Asset
 	conf := asset.Connections[0]
-	runtime, err := s.AddRuntime(func(connId uint32) (*plugin.Runtime, error) {
+	runtime, err := s.AddRuntime(conf, func(connId uint32) (*plugin.Runtime, error) {
 		var conn *connection.SlackConnection
 		var err error
 

diff --git a/providers/terraform/provider/provider.go b/providers/terraform/provider/provider.go
@@ -135,7 +135,7 @@ func (s *Service) connect(req *plugin.ConnectReq, callback plugin.ProviderCallba
 	asset := req.Asset
 	conf := asset.Connections[0]
 
-	runtime, err := s.AddRuntime(func(connId uint32) (*plugin.Runtime, error) {
+	runtime, err := s.AddRuntime(conf, func(connId uint32) (*plugin.Runtime, error) {
 		var conn *connection.Connection
 		var err error
 

diff --git a/providers/vcd/provider/provider.go b/providers/vcd/provider/provider.go
@@ -109,7 +109,7 @@ func (s *Service) connect(req *plugin.ConnectReq, callback plugin.ProviderCallba
 	asset := req.Asset
 	conf := asset.Connections[0]
 
-	runtime, err := s.AddRuntime(func(connId uint32) (*plugin.Runtime, error) {
+	runtime, err := s.AddRuntime(conf, func(connId uint32) (*plugin.Runtime, error) {
 		conn, err := connection.NewVcdConnection(connId, asset, conf)
 		if err != nil {
 			return nil, err

diff --git a/providers/vsphere/provider/provider.go b/providers/vsphere/provider/provider.go
@@ -139,7 +139,7 @@ func (s *Service) connect(req *plugin.ConnectReq, callback plugin.ProviderCallba
 	asset := req.Asset
 	conf := asset.Connections[0]
 
-	runtime, err := s.AddRuntime(func(connId uint32) (*plugin.Runtime, error) {
+	runtime, err := s.AddRuntime(conf, func(connId uint32) (*plugin.Runtime, error) {
 		conn, err := connection.NewVsphereConnection(connId, asset, conf)
 		if err != nil {
 			return nil, err