refactor: optimize az descendant queries BED-4867 (#2892)

* refactor: optimize az descendant queries BED-4867

* test: integration tests

* chore: cleanup

* fix: rabbit feedback for disconnected paths

* Add test for nested management group

* fix: nested management group descendents

---------

Co-authored-by: Elijah Brown <ebrown@specterops.io>

Author: urangel

cmd/api/src/api/v2/azure.go

@@ -180,7 +180,7 @@ func TestResources_GetAZRelatedEntities(t *testing.T) {
 			},
 			expected: expected{
 				responseCode:   http.StatusInternalServerError,
-				responseBody:   `{"errors":[{"context":"","message":"error fetching primary display kinds: database error"}],"http_status":500,"request_id":"","timestamp":"0001-01-01T00:00:00Z"}`,
+				responseBody:   `{"errors":[{"context":"","message":"an internal error has occurred that is preventing the service from servicing this request"}],"http_status":500,"request_id":"","timestamp":"0001-01-01T00:00:00Z"}`,
 				responseHeader: http.Header{"Content-Type": []string{"application/json"}},
 			},
 		},

cmd/api/src/api/v2/azure_test.go

@@ -949,3 +949,222 @@ func TestEntityDetails(t *testing.T) {
 		})
 	})
 }
+
+// TestFetchEntityDescendentPaths_DirectPathsToRoot verifies that each terminal node with a direct
+// azure.Contains edge to the root tenant is included in the returned path set.
+func TestFetchEntityDescendentPaths_DirectPathsToRoot(t *testing.T) {
+	t.Parallel()
+
+	suite := setupIntegrationTestSuite(t)
+	defer teardownIntegrationTestSuite(t, &suite)
+
+	var (
+		tenantID   = integration.RandomObjectID(t)
+		tenantNode = NewAzureTenant(t, &suite, tenantID)
+		app1       = NewAzureApplication(t, &suite, "App1", integration.RandomObjectID(t), tenantID)
+		app2       = NewAzureApplication(t, &suite, "App2", integration.RandomObjectID(t), tenantID)
+	)
+
+	NewRelationship(t, &suite, tenantNode, app1, graphAzure.Contains)
+	NewRelationship(t, &suite, tenantNode, app2, graphAzure.Contains)
+
+	err := suite.GraphDB.ReadTransaction(suite.Context, func(tx graph.Transaction) error {
+		paths, err := azure.FetchEntityDescendentPaths(tx, tenantNode, graphAzure.App)
+		require.NoError(t, err)
+
+		nodeIDs := paths.AllNodes().IDs()
+		require.Len(t, nodeIDs, 3)
+		require.Contains(t, nodeIDs, tenantNode.ID)
+		require.Contains(t, nodeIDs, app1.ID)
+		require.Contains(t, nodeIDs, app2.ID)
+
+		return nil
+	})
+	require.NoError(t, err)
+}
+
+// TestFetchEntityDescendentPaths_MultiHopPathToRoot verifies that a terminal node connected to the
+// root through one or more intermediate azure.Contains hops is fully represented in the path set,
+// with every intermediary node included.
+func TestFetchEntityDescendentPaths_MultiHopPathToRoot(t *testing.T) {
+	t.Parallel()
+
+	suite := setupIntegrationTestSuite(t)
+	defer teardownIntegrationTestSuite(t, &suite)
+
+	var (
+		tenantID         = integration.RandomObjectID(t)
+		tenantNode       = NewAzureTenant(t, &suite, tenantID)
+		subscriptionNode = NewNode(t, &suite, graph.AsProperties(graph.PropertyMap{
+			common.Name:         "TestSubscription",
+			common.ObjectID:     integration.RandomObjectID(t),
+			graphAzure.TenantID: tenantID,
+		}), graphAzure.Entity, graphAzure.Subscription)
+		appNode = NewAzureApplication(t, &suite, "App", integration.RandomObjectID(t), tenantID)
+	)
+
+	// tenant --Contains--> subscription --Contains--> app
+	NewRelationship(t, &suite, tenantNode, subscriptionNode, graphAzure.Contains)
+	NewRelationship(t, &suite, subscriptionNode, appNode, graphAzure.Contains)
+
+	err := suite.GraphDB.ReadTransaction(suite.Context, func(tx graph.Transaction) error {
+		paths, err := azure.FetchEntityDescendentPaths(tx, tenantNode, graphAzure.App)
+		require.NoError(t, err)
+
+		nodeIDs := paths.AllNodes().IDs()
+		require.Len(t, nodeIDs, 3)
+		require.Contains(t, nodeIDs, tenantNode.ID)
+		require.Contains(t, nodeIDs, subscriptionNode.ID)
+		require.Contains(t, nodeIDs, appNode.ID)
+
+		return nil
+	})
+	require.NoError(t, err)
+}
+
+// TestFetchEntityDescendentPaths_NoTerminalsReturnsEmpty verifies that an empty path set is
+// returned when no nodes of the requested kind exist within the tenant.
+func TestFetchEntityDescendentPaths_NoTerminalsReturnsEmpty(t *testing.T) {
+	t.Parallel()
+
+	suite := setupIntegrationTestSuite(t)
+	defer teardownIntegrationTestSuite(t, &suite)
+
+	var (
+		tenantID   = integration.RandomObjectID(t)
+		tenantNode = NewAzureTenant(t, &suite, tenantID)
+	)
+
+	err := suite.GraphDB.ReadTransaction(suite.Context, func(tx graph.Transaction) error {
+		paths, err := azure.FetchEntityDescendentPaths(tx, tenantNode, graphAzure.App)
+		require.NoError(t, err)
+		require.Equal(t, 0, paths.Len())
+
+		return nil
+	})
+	require.NoError(t, err)
+}
+
+// TestFetchEntityDescendentPaths_SharedIntermediateNode verifies that two terminal nodes sharing
+// an intermediate container node both produce complete paths reaching root. Each traversal
+// independently reaches root through the shared intermediate without being truncated.
+func TestFetchEntityDescendentPaths_SharedIntermediateNode(t *testing.T) {
+	t.Parallel()
+
+	suite := setupIntegrationTestSuite(t)
+	defer teardownIntegrationTestSuite(t, &suite)
+
+	var (
+		tenantID         = integration.RandomObjectID(t)
+		tenantNode       = NewAzureTenant(t, &suite, tenantID)
+		subscriptionNode = NewNode(t, &suite, graph.AsProperties(graph.PropertyMap{
+			common.Name:         "SharedSubscription",
+			common.ObjectID:     integration.RandomObjectID(t),
+			graphAzure.TenantID: tenantID,
+		}), graphAzure.Entity, graphAzure.Subscription)
+		app1 = NewAzureApplication(t, &suite, "App1", integration.RandomObjectID(t), tenantID)
+		app2 = NewAzureApplication(t, &suite, "App2", integration.RandomObjectID(t), tenantID)
+	)
+
+	// tenant --Contains--> subscription --Contains--> app1
+	//                                   --Contains--> app2
+	NewRelationship(t, &suite, tenantNode, subscriptionNode, graphAzure.Contains)
+	NewRelationship(t, &suite, subscriptionNode, app1, graphAzure.Contains)
+	NewRelationship(t, &suite, subscriptionNode, app2, graphAzure.Contains)
+
+	err := suite.GraphDB.ReadTransaction(suite.Context, func(tx graph.Transaction) error {
+		paths, err := azure.FetchEntityDescendentPaths(tx, tenantNode, graphAzure.App)
+		require.NoError(t, err)
+
+		nodeIDs := paths.AllNodes().IDs()
+		require.Contains(t, nodeIDs, tenantNode.ID)
+		require.Contains(t, nodeIDs, subscriptionNode.ID)
+		require.Contains(t, nodeIDs, app1.ID)
+		require.Contains(t, nodeIDs, app2.ID)
+
+		return nil
+	})
+	require.NoError(t, err)
+}
+
+// TestFetchEntityDescendentPaths_TerminalNotConnectedToRoot verifies that a terminal node
+// which exists in the same tenant but has no azure.Contains path leading back to root
+// is excluded from the returned path set.
+func TestFetchEntityDescendentPaths_TerminalNotConnectedToRoot(t *testing.T) {
+	t.Parallel()
+
+	suite := setupIntegrationTestSuite(t)
+	defer teardownIntegrationTestSuite(t, &suite)
+
+	var (
+		tenantID        = integration.RandomObjectID(t)
+		tenantNode      = NewAzureTenant(t, &suite, tenantID)
+		connectedApp    = NewAzureApplication(t, &suite, "ConnectedApp", integration.RandomObjectID(t), tenantID)
+		disconnectedApp = NewAzureApplication(t, &suite, "DisconnectedApp", integration.RandomObjectID(t), tenantID)
+	)
+
+	// Only connectedApp has a Contains edge to root; disconnectedApp shares the tenantID
+	// but has no path leading back to tenantNode.
+	NewRelationship(t, &suite, tenantNode, connectedApp, graphAzure.Contains)
+
+	err := suite.GraphDB.ReadTransaction(suite.Context, func(tx graph.Transaction) error {
+		paths, err := azure.FetchEntityDescendentPaths(tx, tenantNode, graphAzure.App)
+		require.NoError(t, err)
+
+		nodeIDs := paths.AllNodes().IDs()
+		require.Contains(t, nodeIDs, tenantNode.ID)
+		require.Contains(t, nodeIDs, connectedApp.ID)
+		require.NotContains(t, nodeIDs, disconnectedApp.ID)
+
+		return nil
+	})
+	require.NoError(t, err)
+}
+
+// TestListEntityDescendents_NestedManagementGroupToSubscription verifies that a Subscription
+// nested under one or more intermediate ManagementGroups is returned when listing the
+// DescendentSubscriptions of an ancestor ManagementGroup.
+func TestListEntityDescendents_NestedManagementGroupToSubscription(t *testing.T) {
+	t.Parallel()
+
+	suite := setupIntegrationTestSuite(t)
+	defer teardownIntegrationTestSuite(t, &suite)
+
+	var (
+		tenantID        = integration.RandomObjectID(t)
+		mgRootObjectID  = integration.RandomObjectID(t)
+		mgChildObjectID = integration.RandomObjectID(t)
+		subObjectID     = integration.RandomObjectID(t)
+		tenantNode      = NewAzureTenant(t, &suite, tenantID)
+		mgRootNode      = NewNode(t, &suite, graph.AsProperties(graph.PropertyMap{
+			common.Name:         "MGRoot",
+			common.ObjectID:     mgRootObjectID,
+			graphAzure.TenantID: tenantID,
+		}), graphAzure.Entity, graphAzure.ManagementGroup)
+		mgChildNode = NewNode(t, &suite, graph.AsProperties(graph.PropertyMap{
+			common.Name:         "MGChild",
+			common.ObjectID:     mgChildObjectID,
+			graphAzure.TenantID: tenantID,
+		}), graphAzure.Entity, graphAzure.ManagementGroup)
+		nestedSubNode = NewNode(t, &suite, graph.AsProperties(graph.PropertyMap{
+			common.Name:         "NestedSubscription",
+			common.ObjectID:     subObjectID,
+			graphAzure.TenantID: tenantID,
+		}), graphAzure.Entity, graphAzure.Subscription)
+	)
+
+	// tenant --Contains--> mgRoot --Contains--> mgChild --Contains--> nestedSub
+	NewRelationship(t, &suite, tenantNode, mgRootNode, graphAzure.Contains)
+	NewRelationship(t, &suite, mgRootNode, mgChildNode, graphAzure.Contains)
+	NewRelationship(t, &suite, mgChildNode, nestedSubNode, graphAzure.Contains)
+
+	nodes, err := azure.ListEntityDescendents(
+		suite.Context, suite.GraphDB,
+		azure.RelatedEntityTypeDescendentSubscriptions, graphAzure.ManagementGroup,
+		mgRootObjectID, 0, 0,
+	)
+	require.NoError(t, err)
+
+	require.Equal(t, 1, len(nodes), "expected nested subscription to be returned as a descendent of mgRoot")
+	require.Contains(t, nodes.IDs(), nestedSubNode.ID)
+}

packages/go/analysis/azure/azure_integration_test.go

@@ -25,7 +25,7 @@ import (
 	"github.com/specterops/dawgs/ops"
 )
 
-func ListEntityDescendentPaths(ctx context.Context, db graph.Database, relatedEntityType RelatedEntityType, objectID string) (graph.PathSet, error) {
+func ListEntityDescendentPaths(ctx context.Context, db graph.Database, relatedEntityType RelatedEntityType, sourceKind graph.Kind, objectID string) (graph.PathSet, error) {
 	var paths graph.PathSet
 
 	return paths, db.ReadTransaction(ctx, func(tx graph.Transaction) error {
@@ -73,13 +73,17 @@ func ListEntityDescendentPaths(ctx context.Context, db graph.Database, relatedEn
 				return ErrInvalidRelatedEntityType
 			}
 
-			paths, err = FetchEntityDescendentPaths(tx, node, targetKind)
+			if isDirectDescendent(sourceKind, targetKind) {
+				paths, err = FetchDirectDescendentPaths(tx, node, targetKind)
+			} else {
+				paths, err = FetchEntityDescendentPaths(tx, node, targetKind)
+			}
 			return err
 		}
 	})
 }
 
-func ListEntityDescendents(ctx context.Context, db graph.Database, relatedEntityType RelatedEntityType, objectID string, skip, limit int) (graph.NodeSet, error) {
+func ListEntityDescendents(ctx context.Context, db graph.Database, relatedEntityType RelatedEntityType, sourceKind graph.Kind, objectID string, skip, limit int) (graph.NodeSet, error) {
 	var nodes graph.NodeSet
 
 	return nodes, db.ReadTransaction(ctx, func(tx graph.Transaction) error {
@@ -127,7 +131,11 @@ func ListEntityDescendents(ctx context.Context, db graph.Database, relatedEntity
 				return ErrInvalidRelatedEntityType
 			}
 
-			nodes, err = FetchEntityDescendents(tx, node, skip, limit, targetKind)
+			if isDirectDescendent(sourceKind, targetKind) {
+				nodes, err = FetchDirectEntityDescendents(tx, node, targetKind)
+			} else {
+				nodes, err = FetchEntityDescendents(tx, node, targetKind)
+			}
 			return err
 		}
 	})

packages/go/analysis/azure/db_ops.go

@@ -45,7 +45,7 @@ func ManagementGroupEntityDetails(ctx context.Context, db graph.Database, primar
 func PopulateManagementGroupEntityDetailsCounts(tx graph.Transaction, node *graph.Node, details ManagementGroupDetails) (ManagementGroupDetails, error) {
 	var descendentKinds = GetDescendentKinds(azure.ManagementGroup)
 
-	if descendents, err := FetchEntityDescendentCounts(tx, node, 0, 0, descendentKinds...); err != nil {
+	if descendents, err := FetchEntityDescendentCounts(tx, node, descendentKinds...); err != nil {
 		return details, err
 	} else {
 		details.Descendents = descendents