⭐️ support re-using MQL cache between provider connections (#3274)

* ⭐️ support re-using MQL cache between provider connections Signed-off-by: Ivan Milchev <[email protected]> * remove go-memoize from k8s provider Signed-off-by: Ivan Milchev <[email protected]> * make k8s discovery reuses mql cache Signed-off-by: Ivan Milchev <[email protected]> * migrate arista provider Signed-off-by: Ivan Milchev <[email protected]> * migrate atlassian provider Signed-off-by: Ivan Milchev <[email protected]> * migrate aws provider Signed-off-by: Ivan Milchev <[email protected]> * migrate os provider Signed-off-by: Ivan Milchev <[email protected]> * migrate azure provider Signed-off-by: Ivan Milchev <[email protected]> * migrate equinix provider Signed-off-by: Ivan Milchev <[email protected]> * re-use plugin.Connection instead of copying functions Signed-off-by: Ivan Milchev <[email protected]> * migrate gcp provider Signed-off-by: Ivan Milchev <[email protected]> * migrate github provider Signed-off-by: Ivan Milchev <[email protected]> * migrate gitlab provider Signed-off-by: Ivan Milchev <[email protected]> * make sure cache is reused for discovered gcp assets Signed-off-by: Ivan Milchev <[email protected]> * migrate google-workspace provider Signed-off-by: Ivan Milchev <[email protected]> * migrate ipmi provider Signed-off-by: Ivan Milchev <[email protected]> * migrate ms365 provider Signed-off-by: Ivan Milchev <[email protected]> * migrate network provider Signed-off-by: Ivan Milchev <[email protected]> * define plugin.Connection and reuse in all providers Signed-off-by: Ivan Milchev <[email protected]> * migrate okta provider Signed-off-by: Ivan Milchev <[email protected]> * migrate opcua provider Signed-off-by: Ivan Milchev <[email protected]> * migrate slack provider Signed-off-by: Ivan Milchev <[email protected]> * migrate terraform provider Signed-off-by: Ivan Milchev <[email protected]> * migrate vcd provider Signed-off-by: Ivan Milchev <[email protected]> * migrate vsphere provider Signed-off-by: Ivan Milchev <[email protected]> * fix test build Signed-off-by: Ivan Milchev <[email protected]> * fix linter error Signed-off-by: Ivan Milchev <[email protected]> * fix mock provider Signed-off-by: Ivan Milchev <[email protected]> * fix dns tests for network provider Signed-off-by: Ivan Milchev <[email protected]> * fix more tests Signed-off-by: Ivan Milchev <[email protected]> * re-use plugin.Connection in os provider Signed-off-by: Ivan Milchev <[email protected]> * fix tests Signed-off-by: Ivan Milchev <[email protected]> * add tests for connection sharing Signed-off-by: Ivan Milchev <[email protected]> * do not use pointer for ParentID Signed-off-by: Ivan Milchev <[email protected]> * use plugin.Connection in winrm Signed-off-by: Ivan Milchev <[email protected]> * fix tests Signed-off-by: Ivan Milchev <[email protected]> --------- Signed-off-by: Ivan Milchev <[email protected]>
mondoohq · Feb 14, 2024 · 054b1dc · 054b1dc
1 parent 8ea1370
commit 054b1dc
Show file tree

Hide file tree

Showing 125 changed files with 1,046 additions and 962 deletions.
diff --git a/providers-sdk/v1/inventory/inventory.go b/providers-sdk/v1/inventory/inventory.go
@@ -378,7 +378,8 @@ func (p *Platform) PrettyTitle() string {
 }
 
 type cloneSettings struct {
-	noDiscovery bool
+	noDiscovery        bool
+	parentConnectionId *uint32
 }
 
 type CloneOption interface {
@@ -394,6 +395,17 @@ type withoutDiscovery struct{}
 
 func (w withoutDiscovery) Apply(o *cloneSettings) { o.noDiscovery = true }
 
+// WithoutDiscovery removes the discovery flags in the opts to ensure the same discovery does not run again
+func WithParentConnectionId(parentId uint32) CloneOption {
+	return withParentConnectionId{parentId: parentId}
+}
+
+type withParentConnectionId struct {
+	parentId uint32
+}
+
+func (w withParentConnectionId) Apply(o *cloneSettings) { o.parentConnectionId = &w.parentId }
+
 func (cfg *Config) Clone(opts ...CloneOption) *Config {
 	if cfg == nil {
 		return nil
@@ -409,6 +421,9 @@ func (cfg *Config) Clone(opts ...CloneOption) *Config {
 	if cloneSettings.noDiscovery {
 		clonedObject.Discover = &Discovery{}
 	}
+	if cloneSettings.parentConnectionId != nil {
+		clonedObject.ParentConnectionId = *cloneSettings.parentConnectionId
+	}
 
 	return clonedObject
 }

diff --git a/providers-sdk/v1/inventory/inventory.pb.go b/providers-sdk/v1/inventory/inventory.pb.go
diff --git a/providers-sdk/v1/inventory/inventory.proto b/providers-sdk/v1/inventory/inventory.proto
@@ -177,6 +177,7 @@ message Config {
   int32 port = 3;
   string path = 4;
   uint32 id = 5;
+  uint32 parent_connection_id = 30;
   string type = 12;
 
   // credentials available for this provider configuration

diff --git a/providers-sdk/v1/plugin/connection.go b/providers-sdk/v1/plugin/connection.go
@@ -0,0 +1,37 @@
+// Copyright (c) Mondoo, Inc.
+// SPDX-License-Identifier: BUSL-1.1
+
+package plugin
+
+import inventory "go.mondoo.com/cnquery/v10/providers-sdk/v1/inventory"
+
+type Connection interface {
+	ID() uint32
+
+	// ParentID returns the ID of the parent connection. If this returns >0,
+	// the connection with that ID will be used to store and get data.
+	ParentID() uint32
+}
+
+type connection struct {
+	id       uint32
+	parentId uint32
+}
+
+func NewConnection(id uint32, asset *inventory.Asset) Connection {
+	conn := &connection{
+		id: id,
+	}
+	if len(asset.Connections) > 0 && asset.Connections[0].ParentConnectionId > 0 {
+		conn.parentId = asset.Connections[0].ParentConnectionId
+	}
+	return conn
+}
+
+func (c *connection) ID() uint32 {
+	return c.id
+}
+
+func (c *connection) ParentID() uint32 {
+	return c.parentId
+}
diff --git a/providers-sdk/v1/plugin/connection_test.go b/providers-sdk/v1/plugin/connection_test.go
@@ -0,0 +1,48 @@
+// Copyright (c) Mondoo, Inc.
+// SPDX-License-Identifier: BUSL-1.1
+
+package plugin
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	inventory "go.mondoo.com/cnquery/v10/providers-sdk/v1/inventory"
+)
+
+func TestConnection_ID(t *testing.T) {
+	c := NewConnection(1, &inventory.Asset{})
+	require.NotNil(t, c)
+	assert.Equal(t, uint32(1), c.ID())
+}
+
+func TestConnection_ParentID_Nil(t *testing.T) {
+	c := NewConnection(1, &inventory.Asset{})
+	require.NotNil(t, c)
+	assert.Equal(t, uint32(0), c.ParentID())
+}
+
+func TestConnection_ParentID(t *testing.T) {
+	c := NewConnection(1, &inventory.Asset{
+		Connections: []*inventory.Config{
+			{
+				ParentConnectionId: 2,
+			},
+		},
+	})
+	require.NotNil(t, c)
+	assert.Equal(t, uint32(2), c.ParentID())
+}
+
+func TestConnection_ParentID_0(t *testing.T) {
+	c := NewConnection(1, &inventory.Asset{
+		Connections: []*inventory.Config{
+			{
+				ParentConnectionId: 0,
+			},
+		},
+	})
+	require.NotNil(t, c)
+	assert.Equal(t, uint32(0), c.ParentID())
+}
diff --git a/providers-sdk/v1/plugin/runtime.go b/providers-sdk/v1/plugin/runtime.go
@@ -14,7 +14,7 @@ import (
 
 type Runtime struct {
 	Connection     Connection
-	Resources      syncx.Map[Resource]
+	Resources      *syncx.Map[Resource]
 	Callback       ProviderCallback
 	HasRecording   bool
 	CreateResource CreateNamedResource
@@ -24,8 +24,27 @@ type Runtime struct {
 	Upstream       *upstream.UpstreamClient
 }
 
-type Connection interface {
-	ID() uint32
+func NewRuntime(
+	conn Connection,
+	callback ProviderCallback,
+	hasRecording bool,
+	createResource CreateNamedResource,
+	newResource NewResource,
+	getData GetData,
+	setData SetData,
+	upstream *upstream.UpstreamClient,
+) *Runtime {
+	return &Runtime{
+		Connection:     conn,
+		Resources:      &syncx.Map[Resource]{},
+		Callback:       callback,
+		HasRecording:   hasRecording,
+		CreateResource: createResource,
+		NewResource:    newResource,
+		GetData:        getData,
+		SetData:        setData,
+		Upstream:       upstream,
+	}
 }
 
 type (

diff --git a/providers-sdk/v1/plugin/service.go b/providers-sdk/v1/plugin/service.go
@@ -42,13 +42,30 @@ func (s *Service) AddRuntime(createRuntime func(connId uint32) (*Runtime, error)
 		s.lastConnectionID--
 		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[s.lastConnectionID] = runtime
 	return runtime, nil
 }
 
 func (s *Service) GetRuntime(id uint32) (*Runtime, error) {
 	s.runtimesLock.Lock()
 	defer s.runtimesLock.Unlock()
+	return s.doGetRuntime(id)
+}
+
+// doGetRuntime is a helper function to get a runtime by its ID. It MUST be called
+// with a lock on s.runtimesLock.
+func (s *Service) doGetRuntime(id uint32) (*Runtime, error) {
 	if runtime, ok := s.runtimes[id]; ok {
 		return runtime, nil
 	}

diff --git a/providers-sdk/v1/plugin/service_test.go b/providers-sdk/v1/plugin/service_test.go
@@ -9,10 +9,12 @@ import (
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"go.mondoo.com/cnquery/v10/utils/syncx"
 )
 
 type TestConnection struct {
-	id uint32
+	id       uint32
+	parentId uint32
 }
 
 func newTestConnection(id uint32) *TestConnection {
@@ -23,6 +25,10 @@ func (c *TestConnection) ID() uint32 {
 	return c.id
 }
 
+func (c *TestConnection) ParentID() uint32 {
+	return c.parentId
+}
+
 type TestConnectionWithClose struct {
 	*TestConnection
 	closed bool
@@ -36,6 +42,16 @@ func (c *TestConnectionWithClose) Close() {
 	c.closed = true
 }
 
+type TestResource struct{}
+
+func (r *TestResource) MqlID() string {
+	return "test.resource"
+}
+
+func (r *TestResource) MqlName() string {
+	return "Test Resource"
+}
+
 func TestAddRuntime(t *testing.T) {
 	s := NewService()
 	wg := sync.WaitGroup{}
@@ -63,6 +79,52 @@ func TestAddRuntime(t *testing.T) {
 	assert.Equal(t, s.lastConnectionID, uint32(200))
 }
 
+func TestAddRuntime_ParentNotExist(t *testing.T) {
+	s := NewService()
+	parentId := uint32(10)
+	_, err := s.AddRuntime(func(connId uint32) (*Runtime, error) {
+		c := newTestConnection(connId)
+		c.parentId = parentId
+		return &Runtime{
+			Connection: c,
+		}, nil
+	})
+	require.Error(t, err)
+	assert.Equal(t, "parent connection 10 not found", err.Error())
+}
+
+func TestAddRuntime_Parent(t *testing.T) {
+	s := NewService()
+
+	parent, err := s.AddRuntime(func(connId uint32) (*Runtime, error) {
+		resMap := &syncx.Map[Resource]{}
+		resMap.Set("test.resource", &TestResource{})
+
+		return &Runtime{
+			Resources:  resMap,
+			Connection: newTestConnection(connId),
+		}, nil
+	})
+	require.NoError(t, err)
+
+	parentId := parent.Connection.ID()
+	child, err := s.AddRuntime(func(connId uint32) (*Runtime, error) {
+		c := newTestConnection(connId)
+		c.parentId = parentId
+		return &Runtime{
+			Connection: c,
+		}, nil
+	})
+	require.NoError(t, err)
+
+	// Check that the resources for the parent and the child are the same
+	assert.Equal(t, parent.Resources, child.Resources)
+
+	// Add another resource and check that it appears in the child runtime
+	parent.Resources.Set("test.resource2", &TestResource{})
+	assert.Equal(t, parent.Resources, child.Resources)
+}
+
 func TestGetRuntime(t *testing.T) {
 	s := NewService()
 

diff --git a/providers-sdk/v1/testutils/mockprovider/mockprovider.go b/providers-sdk/v1/testutils/mockprovider/mockprovider.go
@@ -11,6 +11,7 @@ import (
 	"go.mondoo.com/cnquery/v10/llx"
 	"go.mondoo.com/cnquery/v10/providers-sdk/v1/plugin"
 	"go.mondoo.com/cnquery/v10/providers-sdk/v1/testutils/mockprovider/resources"
+	"go.mondoo.com/cnquery/v10/utils/syncx"
 )
 
 var Config = plugin.Provider{
@@ -44,6 +45,7 @@ func (s *Service) Connect(req *plugin.ConnectReq, callback plugin.ProviderCallba
 	s.lastConnectionID++
 	connID := s.lastConnectionID
 	runtime := &plugin.Runtime{
+		Resources:    &syncx.Map[plugin.Resource]{},
 		Callback:     callback,
 		HasRecording: req.HasRecording,
 	}

diff --git a/providers/arista/connection/connection.go b/providers/arista/connection/connection.go
@@ -8,11 +8,12 @@ import (
 
 	"github.com/aristanetworks/goeapi"
 	"go.mondoo.com/cnquery/v10/providers-sdk/v1/inventory"
+	"go.mondoo.com/cnquery/v10/providers-sdk/v1/plugin"
 	"go.mondoo.com/cnquery/v10/providers-sdk/v1/vault"
 )
 
 type AristaConnection struct {
-	id    uint32
+	plugin.Connection
 	Conf  *inventory.Config
 	asset *inventory.Asset
 	// custom connection fields
@@ -21,9 +22,9 @@ type AristaConnection struct {
 
 func NewAristaConnection(id uint32, asset *inventory.Asset, conf *inventory.Config) (*AristaConnection, error) {
 	conn := &AristaConnection{
-		Conf:  conf,
-		id:    id,
-		asset: asset,
+		Connection: plugin.NewConnection(id, asset),
+		Conf:       conf,
+		asset:      asset,
 	}
 
 	// initialize connection
@@ -59,10 +60,6 @@ func (c *AristaConnection) Name() string {
 	return "arista"
 }
 
-func (c *AristaConnection) ID() uint32 {
-	return c.id
-}
-
 func (c *AristaConnection) Asset() *inventory.Asset {
 	return c.asset
 }

diff --git a/providers/arista/provider/provider.go b/providers/arista/provider/provider.go
@@ -139,16 +139,15 @@ func (s *Service) connect(req *plugin.ConnectReq, callback plugin.ProviderCallba
 		}
 		asset.Connections[0].Id = connId
 
-		return &plugin.Runtime{
-			Connection:     conn,
-			Callback:       callback,
-			HasRecording:   req.HasRecording,
-			CreateResource: resources.CreateResource,
-			NewResource:    resources.NewResource,
-			GetData:        resources.GetData,
-			SetData:        resources.SetData,
-			Upstream:       upstream,
-		}, nil
+		return plugin.NewRuntime(
+			conn,
+			callback,
+			req.HasRecording,
+			resources.CreateResource,
+			resources.NewResource,
+			resources.GetData,
+			resources.SetData,
+			upstream), nil
 	})
 	if err != nil {
 		return nil, err