extend provider tests to cover delayed discovery

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

Hide file tree

Showing 4 changed files with 135 additions and 8 deletions.
diff --git a/providers-sdk/v1/plugin/service_test.go b/providers-sdk/v1/plugin/service_test.go
@@ -133,6 +133,20 @@ func TestDeprecatedAddRuntime(t *testing.T) {
 	assert.Equal(t, s.lastConnectionID, uint32(200))
 }
 
+func TestDeprecatedAddRuntime_DisableDelayedDiscovery(t *testing.T) {
+	s := NewService()
+	inv := &inventory.Config{}
+	_, err := s.AddRuntime(inv, func(connId uint32) (*Runtime, error) {
+		c := newTestConnection(connId)
+		return &Runtime{
+			Connection: c,
+		}, nil
+	})
+	require.NoError(t, err)
+	require.Contains(t, inv.Options, DISABLE_DELAYED_DISCOVERY_OPTION)
+	assert.Equal(t, "true", inv.Options[DISABLE_DELAYED_DISCOVERY_OPTION])
+}
+
 func TestAddRuntime_ParentNotExist(t *testing.T) {
 	s := NewService()
 	parentId := uint32(10)

diff --git a/providers/os/connection/container/image_connection_test.go b/providers/os/connection/container/image_connection_test.go
@@ -17,7 +17,9 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"go.mondoo.com/cnquery/v10/providers-sdk/v1/inventory"
+	"go.mondoo.com/cnquery/v10/providers-sdk/v1/plugin"
 	"go.mondoo.com/cnquery/v10/providers/os/connection/container"
+	"go.mondoo.com/cnquery/v10/providers/os/connection/container/image"
 	"go.mondoo.com/cnquery/v10/providers/os/connection/tar"
 )
 
@@ -68,6 +70,32 @@ type dockerConnTest struct {
 	testfile string
 }
 
+func TestNewImageConnection_DelayDiscovery(t *testing.T) {
+	ref, err := name.ParseReference(alpineImage, name.WeakValidation)
+	require.NoError(t, err)
+
+	img, err := image.LoadImageFromRegistry(ref)
+	require.NoError(t, err)
+
+	inv := &inventory.Config{Options: map[string]string{}}
+	_, err = container.NewImageConnection(1, inv, &inventory.Asset{}, img)
+	require.NoError(t, err)
+	assert.True(t, inv.DelayDiscovery)
+}
+
+func TestNewImageConnection_DisableDelayDiscovery(t *testing.T) {
+	ref, err := name.ParseReference(alpineImage, name.WeakValidation)
+	require.NoError(t, err)
+
+	img, err := image.LoadImageFromRegistry(ref)
+	require.NoError(t, err)
+
+	inv := &inventory.Config{Options: map[string]string{plugin.DISABLE_DELAYED_DISCOVERY_OPTION: "true"}}
+	_, err = container.NewImageConnection(1, inv, &inventory.Asset{}, img)
+	require.NoError(t, err)
+	assert.False(t, inv.DelayDiscovery)
+}
+
 func TestImageConnections(t *testing.T) {
 	var testConnections []dockerConnTest
 
@@ -214,7 +242,6 @@ func TestImageConnections(t *testing.T) {
 			})
 		})
 	}
-
 }
 
 func TestTarSymlinkFile(t *testing.T) {

diff --git a/providers/os/connection/docker/container_connection_test.go b/providers/os/connection/docker/container_connection_test.go
@@ -6,11 +6,12 @@ package docker
 import (
 	"context"
 	"fmt"
-	"github.com/docker/docker/client"
 	"io"
 	"os"
 	"testing"
 
+	"github.com/docker/docker/client"
+
 	"github.com/docker/docker/api/types"
 	"github.com/docker/docker/api/types/container"
 	"github.com/docker/docker/api/types/network"
@@ -19,6 +20,7 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"go.mondoo.com/cnquery/v10/providers-sdk/v1/inventory"
+	"go.mondoo.com/cnquery/v10/providers-sdk/v1/plugin"
 	"go.mondoo.com/cnquery/v10/providers/os/connection/tar"
 )
 
@@ -31,14 +33,50 @@ func TestAssetNameForRemoteImages(t *testing.T) {
 	retries := 3
 	counter := 0
 
+	config := &inventory.Config{
+		Type: "docker-image",
+		Host: "gcr.io/google-containers/busybox:1.27.2",
+	}
+	asset = &inventory.Asset{
+		Connections: []*inventory.Config{config},
+	}
+
 	for {
-		config := &inventory.Config{
-			Type: "docker-image",
-			Host: "gcr.io/google-containers/busybox:1.27.2",
-		}
-		asset = &inventory.Asset{
-			Connections: []*inventory.Config{config},
+		conn, err = NewContainerImageConnection(0, config, asset)
+		if counter > retries || (err == nil && conn != nil) {
+			break
 		}
+		counter++
+	}
+	require.NoError(t, err)
+	require.NotNil(t, conn)
+
+	assert.True(t, config.DelayDiscovery)
+	assert.Equal(t, "gcr.io/google-containers/busybox@545e6a6310a2", asset.Name)
+	assert.Contains(t, asset.PlatformIds, "//platformid.api.mondoo.app/runtime/docker/images/545e6a6310a27636260920bc07b994a299b6708a1b26910cfefd335fdfb60d2b")
+}
+
+// This test has an external dependency on the gcr.io registry
+// To test this specific case, we cannot use a stored image, we need to call remote.Get
+func TestAssetNameForRemoteImages_DisableDelayedDiscovery(t *testing.T) {
+	var err error
+	var conn *tar.Connection
+	var asset *inventory.Asset
+	retries := 3
+	counter := 0
+
+	config := &inventory.Config{
+		Type: "docker-image",
+		Host: "gcr.io/google-containers/busybox:1.27.2",
+		Options: map[string]string{
+			plugin.DISABLE_DELAYED_DISCOVERY_OPTION: "true",
+		},
+	}
+	asset = &inventory.Asset{
+		Connections: []*inventory.Config{config},
+	}
+
+	for {
 		conn, err = NewContainerImageConnection(0, config, asset)
 		if counter > retries || (err == nil && conn != nil) {
 			break
@@ -48,6 +86,7 @@ func TestAssetNameForRemoteImages(t *testing.T) {
 	require.NoError(t, err)
 	require.NotNil(t, conn)
 
+	assert.False(t, config.DelayDiscovery)
 	assert.Equal(t, "gcr.io/google-containers/busybox@545e6a6310a2", asset.Name)
 	assert.Contains(t, asset.PlatformIds, "//platformid.api.mondoo.app/runtime/docker/images/545e6a6310a27636260920bc07b994a299b6708a1b26910cfefd335fdfb60d2b")
 }

diff --git a/providers/os/provider/provider_test.go b/providers/os/provider/provider_test.go
@@ -69,6 +69,53 @@ func TestLocalConnectionIdDetectors(t *testing.T) {
 	require.NotNil(t, shutdownconnectResp)
 }
 
+func TestLocalConnectionIdDetectors_DelayedDiscovery(t *testing.T) {
+	srv := &Service{
+		Service: plugin.NewService(),
+	}
+
+	connectResp, err := srv.Connect(&plugin.ConnectReq{
+		Asset: &inventory.Asset{
+			Connections: []*inventory.Config{
+				{
+					Type:           "local",
+					DelayDiscovery: true,
+				},
+			},
+		},
+	}, nil)
+	require.NoError(t, err)
+	require.NotNil(t, connectResp)
+
+	require.Len(t, connectResp.Asset.IdDetector, 2)
+	require.Contains(t, connectResp.Asset.IdDetector, ids.IdDetector_Hostname)
+	require.Contains(t, connectResp.Asset.IdDetector, ids.IdDetector_CloudDetect)
+	require.NotContains(t, connectResp.Asset.IdDetector, ids.IdDetector_SshHostkey)
+	require.Len(t, connectResp.Asset.PlatformIds, 1)
+	require.Nil(t, connectResp.Asset.Platform)
+
+	// Disable delayed discovery and reconnect
+	connectResp.Asset.Connections[0].DelayDiscovery = false
+	connectResp, err = srv.Connect(&plugin.ConnectReq{
+		Asset: connectResp.Asset,
+	}, nil)
+	require.NoError(t, err)
+	require.NotNil(t, connectResp)
+
+	require.Len(t, connectResp.Asset.IdDetector, 2)
+	require.Contains(t, connectResp.Asset.IdDetector, ids.IdDetector_Hostname)
+	require.Contains(t, connectResp.Asset.IdDetector, ids.IdDetector_CloudDetect)
+	require.NotContains(t, connectResp.Asset.IdDetector, ids.IdDetector_SshHostkey)
+	// Now the platformIDs are cleaned up
+	require.Len(t, connectResp.Asset.PlatformIds, 2)
+	// Verify the platform is set
+	require.NotNil(t, connectResp.Asset.Platform)
+
+	shutdownconnectResp, err := srv.Shutdown(&plugin.ShutdownReq{})
+	require.NoError(t, err)
+	require.NotNil(t, shutdownconnectResp)
+}
+
 func TestIdentifyDockerString(t *testing.T) {
 	tests := []struct {
 		input string