From 50e2020771d2b6b5271ef5857b901efe9f8962c8 Mon Sep 17 00:00:00 2001 From: Ulises Rangel Date: Mon, 15 Jun 2026 10:30:33 -0500 Subject: [PATCH 1/6] refactor: optimize az descendant queries BED-4867 --- cmd/api/src/api/v2/azure.go | 153 ++++++++------- cmd/api/src/queries/graph.go | 1 + packages/go/analysis/azure/db_ops.go | 24 ++- packages/go/analysis/azure/queries.go | 174 ++++++++++++++++-- packages/go/analysis/azure/resource_group.go | 2 +- .../js-client-library/src/client.ts | 2 +- 6 files changed, 262 insertions(+), 94 deletions(-) diff --git a/cmd/api/src/api/v2/azure.go b/cmd/api/src/api/v2/azure.go index 80da05dcfa84..6f2f6184d60c 100644 --- a/cmd/api/src/api/v2/azure.go +++ b/cmd/api/src/api/v2/azure.go @@ -78,15 +78,10 @@ var ( ErrParameterRelatedEntityType = errors.New("invalid related entity type") ) -func graphRelatedEntityType(request *http.Request, db database.Database, graphDb graph.Database, entityType, objectID string) (any, int, *api.ErrorWrapper) { +func graphRelatedEntityType(request *http.Request, graphDb graph.Database, primaryDisplayKinds graphschema.PrimaryDisplayKinds, options relatedEntityTypeOptions) (any, int, *api.ErrorWrapper) { ctx := request.Context() - primaryDisplayKinds, err := db.GetPrimaryDisplayKinds(request.Context()) - if err != nil { - return nil, 0, api.BuildErrorResponse(http.StatusInternalServerError, fmt.Sprintf("error fetching primary display kinds: %v", err), request) - } - - switch relatedEntityType := azure.RelatedEntityType(entityType); relatedEntityType { + switch relatedEntityType := azure.RelatedEntityType(options.RelatedEntityString); relatedEntityType { case azure.RelatedEntityTypeDescendentUsers, azure.RelatedEntityTypeDescendentGroups, azure.RelatedEntityTypeDescendentManagementGroups, azure.RelatedEntityTypeDescendentSubscriptions, azure.RelatedEntityTypeDescendentResourceGroups, azure.RelatedEntityTypeDescendentVirtualMachines, @@ -98,116 +93,116 @@ func graphRelatedEntityType(request *http.Request, db database.Database, graphDb azure.RelatedEntityTypeDescendentWebApps, azure.RelatedEntityTypeDescendentAutomationAccounts, azure.RelatedEntityTypeDescendentLogicApps, azure.RelatedEntityTypeDescendentFunctionApps: - if descendents, err := azure.ListEntityDescendentPaths(ctx, graphDb, relatedEntityType, objectID); err != nil { - return nil, 0, api.BuildErrorResponse(http.StatusInternalServerError, fmt.Sprintf("error fetching related entity type %s: %v", entityType, err), request) + if descendents, err := azure.ListEntityDescendentPaths(ctx, graphDb, relatedEntityType, options.SourceKind, options.SourceObjectID); err != nil { + return nil, 0, api.BuildErrorResponse(http.StatusInternalServerError, fmt.Sprintf("error fetching related entity type %s: %v", options.RelatedEntityString, err), request) } else { return bloodhoundgraph.PathSetToBloodHoundGraph(primaryDisplayKinds, descendents), descendents.Len(), nil } case azure.RelatedEntityTypeActiveAssignments: - if assignments, err := azure.ListEntityActiveAssignmentPaths(ctx, graphDb, objectID); err != nil { - return nil, 0, api.BuildErrorResponse(http.StatusInternalServerError, fmt.Sprintf("error fetching related entity type %s: %v", entityType, err), request) + if assignments, err := azure.ListEntityActiveAssignmentPaths(ctx, graphDb, options.SourceObjectID); err != nil { + return nil, 0, api.BuildErrorResponse(http.StatusInternalServerError, fmt.Sprintf("error fetching related entity type %s: %v", options.RelatedEntityString, err), request) } else { return bloodhoundgraph.PathSetToBloodHoundGraph(primaryDisplayKinds, assignments), assignments.Len(), nil } case azure.RelatedEntityTypePIMAssignments: - if assignments, err := azure.ListEntityPIMAssignmentPaths(ctx, graphDb, objectID); err != nil { - return nil, 0, api.BuildErrorResponse(http.StatusInternalServerError, fmt.Sprintf("error fetching related entity type %s: %v", entityType, err), request) + if assignments, err := azure.ListEntityPIMAssignmentPaths(ctx, graphDb, options.SourceObjectID); err != nil { + return nil, 0, api.BuildErrorResponse(http.StatusInternalServerError, fmt.Sprintf("error fetching related entity type %s: %v", options.RelatedEntityString, err), request) } else { return bloodhoundgraph.PathSetToBloodHoundGraph(primaryDisplayKinds, assignments), assignments.Len(), nil } case azure.RelatedEntityTypeRoleApprovers: - if approvers, err := azure.ListRoleApproverPaths(ctx, graphDb, objectID); err != nil { - return nil, 0, api.BuildErrorResponse(http.StatusInternalServerError, fmt.Sprintf("error fetching related entity type %s: %v", entityType, err), request) + if approvers, err := azure.ListRoleApproverPaths(ctx, graphDb, options.SourceObjectID); err != nil { + return nil, 0, api.BuildErrorResponse(http.StatusInternalServerError, fmt.Sprintf("error fetching related entity type %s: %v", options.RelatedEntityString, err), request) } else { return bloodhoundgraph.PathSetToBloodHoundGraph(primaryDisplayKinds, approvers), approvers.Len(), nil } case azure.RelatedEntityTypeVaultKeyReaders, azure.RelatedEntityTypeVaultSecretReaders, azure.RelatedEntityTypeVaultCertReaders, azure.RelatedEntityTypeVaultAllReaders: - if groupMembers, err := azure.ListKeyVaultReaderPaths(ctx, graphDb, relatedEntityType, objectID); err != nil { - return nil, 0, api.BuildErrorResponse(http.StatusInternalServerError, fmt.Sprintf("error fetching related entity type %s: %v", entityType, err), request) + if groupMembers, err := azure.ListKeyVaultReaderPaths(ctx, graphDb, relatedEntityType, options.SourceObjectID); err != nil { + return nil, 0, api.BuildErrorResponse(http.StatusInternalServerError, fmt.Sprintf("error fetching related entity type %s: %v", options.RelatedEntityString, err), request) } else { return bloodhoundgraph.PathSetToBloodHoundGraph(primaryDisplayKinds, groupMembers), groupMembers.Len(), nil } case azure.RelatedEntityTypeGroupMembers: - if groupMembers, err := azure.ListEntityGroupMemberPaths(ctx, graphDb, objectID); err != nil { - return nil, 0, api.BuildErrorResponse(http.StatusInternalServerError, fmt.Sprintf("error fetching related entity type %s: %v", entityType, err), request) + if groupMembers, err := azure.ListEntityGroupMemberPaths(ctx, graphDb, options.SourceObjectID); err != nil { + return nil, 0, api.BuildErrorResponse(http.StatusInternalServerError, fmt.Sprintf("error fetching related entity type %s: %v", options.RelatedEntityString, err), request) } else { return bloodhoundgraph.PathSetToBloodHoundGraph(primaryDisplayKinds, groupMembers), groupMembers.Len(), nil } case azure.RelatedEntityTypeGroupMembership: - if groupMembership, err := azure.ListEntityGroupMembershipPaths(ctx, graphDb, objectID); err != nil { - return nil, 0, api.BuildErrorResponse(http.StatusInternalServerError, fmt.Sprintf("error fetching related entity type %s: %v", entityType, err), request) + if groupMembership, err := azure.ListEntityGroupMembershipPaths(ctx, graphDb, options.SourceObjectID); err != nil { + return nil, 0, api.BuildErrorResponse(http.StatusInternalServerError, fmt.Sprintf("error fetching related entity type %s: %v", options.RelatedEntityString, err), request) } else { return bloodhoundgraph.PathSetToBloodHoundGraph(primaryDisplayKinds, groupMembership), groupMembership.Len(), nil } case azure.RelatedEntityTypeRoles: - if userRoles, err := azure.ListEntityRolePaths(ctx, graphDb, objectID); err != nil { - return nil, 0, api.BuildErrorResponse(http.StatusInternalServerError, fmt.Sprintf("error fetching related entity type %s: %v", entityType, err), request) + if userRoles, err := azure.ListEntityRolePaths(ctx, graphDb, options.SourceObjectID); err != nil { + return nil, 0, api.BuildErrorResponse(http.StatusInternalServerError, fmt.Sprintf("error fetching related entity type %s: %v", options.RelatedEntityString, err), request) } else { return bloodhoundgraph.PathSetToBloodHoundGraph(primaryDisplayKinds, userRoles), userRoles.Len(), nil } case azure.RelatedEntityTypeEligibleAndApproverRoles: - if eligibleAndApproverRoles, err := azure.ListEntityEligibleAndApproverRolePaths(ctx, graphDb, objectID); err != nil { - return nil, 0, api.BuildErrorResponse(http.StatusInternalServerError, fmt.Sprintf("error fetching related entity type %s: %v", entityType, err), request) + if eligibleAndApproverRoles, err := azure.ListEntityEligibleAndApproverRolePaths(ctx, graphDb, options.SourceObjectID); err != nil { + return nil, 0, api.BuildErrorResponse(http.StatusInternalServerError, fmt.Sprintf("error fetching related entity type %s: %v", options.RelatedEntityString, err), request) } else { return bloodhoundgraph.PathSetToBloodHoundGraph(primaryDisplayKinds, eligibleAndApproverRoles), eligibleAndApproverRoles.Len(), nil } case azure.RelatedEntityTypeOutboundExecutionPrivileges: - if executionPrivileges, err := azure.ListEntityExecutionPrivilegePaths(ctx, graphDb, objectID, graph.DirectionOutbound); err != nil { - return nil, 0, api.BuildErrorResponse(http.StatusInternalServerError, fmt.Sprintf("error fetching related entity type %s: %v", entityType, err), request) + if executionPrivileges, err := azure.ListEntityExecutionPrivilegePaths(ctx, graphDb, options.SourceObjectID, graph.DirectionOutbound); err != nil { + return nil, 0, api.BuildErrorResponse(http.StatusInternalServerError, fmt.Sprintf("error fetching related entity type %s: %v", options.RelatedEntityString, err), request) } else { return bloodhoundgraph.PathSetToBloodHoundGraph(primaryDisplayKinds, executionPrivileges), executionPrivileges.Len(), nil } case azure.RelatedEntityTypeInboundExecutionPrivileges: - if executionPrivileges, err := azure.ListEntityExecutionPrivilegePaths(ctx, graphDb, objectID, graph.DirectionInbound); err != nil { - return nil, 0, api.BuildErrorResponse(http.StatusInternalServerError, fmt.Sprintf("error fetching related entity type %s: %v", entityType, err), request) + if executionPrivileges, err := azure.ListEntityExecutionPrivilegePaths(ctx, graphDb, options.SourceObjectID, graph.DirectionInbound); err != nil { + return nil, 0, api.BuildErrorResponse(http.StatusInternalServerError, fmt.Sprintf("error fetching related entity type %s: %v", options.RelatedEntityString, err), request) } else { return bloodhoundgraph.PathSetToBloodHoundGraph(primaryDisplayKinds, executionPrivileges), executionPrivileges.Len(), nil } case azure.RelatedEntityTypeOutboundAbusableAppRoleAssignments: - if objectControl, err := azure.ListEntityAbusableAppRoleAssignmentsPaths(ctx, graphDb, objectID, graph.DirectionOutbound); err != nil { - return nil, 0, api.BuildErrorResponse(http.StatusInternalServerError, fmt.Sprintf("error fetching related entity type %s: %v", entityType, err), request) + if objectControl, err := azure.ListEntityAbusableAppRoleAssignmentsPaths(ctx, graphDb, options.SourceObjectID, graph.DirectionOutbound); err != nil { + return nil, 0, api.BuildErrorResponse(http.StatusInternalServerError, fmt.Sprintf("error fetching related entity type %s: %v", options.RelatedEntityString, err), request) } else { return bloodhoundgraph.PathSetToBloodHoundGraph(primaryDisplayKinds, objectControl), objectControl.Len(), nil } case azure.RelatedEntityTypeInboundAbusableAppRoleAssignments: - if objectControl, err := azure.ListEntityAbusableAppRoleAssignmentsPaths(ctx, graphDb, objectID, graph.DirectionInbound); err != nil { - return nil, 0, api.BuildErrorResponse(http.StatusInternalServerError, fmt.Sprintf("error fetching related entity type %s: %v", entityType, err), request) + if objectControl, err := azure.ListEntityAbusableAppRoleAssignmentsPaths(ctx, graphDb, options.SourceObjectID, graph.DirectionInbound); err != nil { + return nil, 0, api.BuildErrorResponse(http.StatusInternalServerError, fmt.Sprintf("error fetching related entity type %s: %v", options.RelatedEntityString, err), request) } else { return bloodhoundgraph.PathSetToBloodHoundGraph(primaryDisplayKinds, objectControl), objectControl.Len(), nil } case azure.RelatedEntityTypeOutboundControl: - if objectControl, err := azure.ListEntityObjectControlPaths(ctx, graphDb, objectID, graph.DirectionOutbound); err != nil { - return nil, 0, api.BuildErrorResponse(http.StatusInternalServerError, fmt.Sprintf("error fetching related entity type %s: %v", entityType, err), request) + if objectControl, err := azure.ListEntityObjectControlPaths(ctx, graphDb, options.SourceObjectID, graph.DirectionOutbound); err != nil { + return nil, 0, api.BuildErrorResponse(http.StatusInternalServerError, fmt.Sprintf("error fetching related entity type %s: %v", options.RelatedEntityString, err), request) } else { return bloodhoundgraph.PathSetToBloodHoundGraph(primaryDisplayKinds, objectControl), objectControl.Len(), nil } case azure.RelatedEntityTypeInboundControl: - if objectControl, err := azure.ListEntityObjectControlPaths(ctx, graphDb, objectID, graph.DirectionInbound); err != nil { - return nil, 0, api.BuildErrorResponse(http.StatusInternalServerError, fmt.Sprintf("error fetching related entity type %s: %v", entityType, err), request) + if objectControl, err := azure.ListEntityObjectControlPaths(ctx, graphDb, options.SourceObjectID, graph.DirectionInbound); err != nil { + return nil, 0, api.BuildErrorResponse(http.StatusInternalServerError, fmt.Sprintf("error fetching related entity type %s: %v", options.RelatedEntityString, err), request) } else { return bloodhoundgraph.PathSetToBloodHoundGraph(primaryDisplayKinds, objectControl), objectControl.Len(), nil } case azure.RelatedEntityTypeFederatedIdentityCredentials: - if fics, err := azure.ListAppFederatedIdentityCredentialPaths(ctx, graphDb, objectID); err != nil { - return nil, 0, api.BuildErrorResponse(http.StatusInternalServerError, fmt.Sprintf("error fetching related entity type %s: %v", entityType, err), request) + if fics, err := azure.ListAppFederatedIdentityCredentialPaths(ctx, graphDb, options.SourceObjectID); err != nil { + return nil, 0, api.BuildErrorResponse(http.StatusInternalServerError, fmt.Sprintf("error fetching related entity type %s: %v", options.RelatedEntityString, err), request) } else { return bloodhoundgraph.PathSetToBloodHoundGraph(primaryDisplayKinds, fics), fics.Len(), nil } default: - return nil, 0, api.BuildErrorResponse(http.StatusNotFound, fmt.Sprintf("no matching related entity list type for %s", entityType), request) + return nil, 0, api.BuildErrorResponse(http.StatusNotFound, fmt.Sprintf("no matching related entity list type for %s", options.RelatedEntityString), request) } } @@ -221,13 +216,21 @@ func nodeSetToOrderedSlice(nodeSet graph.NodeSet) []*graph.Node { return nodes } -func listRelatedEntityType(ctx context.Context, db graph.Database, primaryDisplayKinds graphschema.PrimaryDisplayKinds, entityType, objectID string, skip, limit int) ([]azure.Node, int, error) { +type relatedEntityTypeOptions struct { + RelatedEntityString string + SourceKind graph.Kind + SourceObjectID string + skip int + limit int +} + +func listRelatedEntityType(ctx context.Context, db graph.Database, primaryDisplayKinds graphschema.PrimaryDisplayKinds, options relatedEntityTypeOptions) ([]azure.Node, int, error) { var ( nodeSet graph.NodeSet err error ) // NOTE: All skip/limit passed to lower level queries is currently hardcoded to 0 so we can get the full count of the dataset for skip/limit tracking - switch relatedEntityType := azure.RelatedEntityType(entityType); relatedEntityType { + switch relatedEntityType := azure.RelatedEntityType(options.RelatedEntityString); relatedEntityType { case azure.RelatedEntityTypeDescendentUsers, azure.RelatedEntityTypeDescendentGroups, azure.RelatedEntityTypeDescendentManagementGroups, azure.RelatedEntityTypeDescendentSubscriptions, azure.RelatedEntityTypeDescendentResourceGroups, azure.RelatedEntityTypeDescendentVirtualMachines, @@ -239,77 +242,77 @@ func listRelatedEntityType(ctx context.Context, db graph.Database, primaryDispla azure.RelatedEntityTypeDescendentWebApps, azure.RelatedEntityTypeDescendentLogicApps, azure.RelatedEntityTypeDescendentFunctionApps, azure.RelatedEntityTypeDescendentAutomationAccounts: - if nodeSet, err = azure.ListEntityDescendents(ctx, db, relatedEntityType, objectID, 0, 0); err != nil { + if nodeSet, err = azure.ListEntityDescendents(ctx, db, relatedEntityType, options.SourceKind, options.SourceObjectID, 0, 0); err != nil { return nil, 0, err } case azure.RelatedEntityTypeActiveAssignments: - if nodeSet, err = azure.ListEntityActiveAssignments(ctx, db, objectID, 0, 0); err != nil { + if nodeSet, err = azure.ListEntityActiveAssignments(ctx, db, options.SourceObjectID, 0, 0); err != nil { return nil, 0, err } case azure.RelatedEntityTypePIMAssignments: - if nodeSet, err = azure.ListEntityPIMAssignments(ctx, db, objectID, 0, 0); err != nil { + if nodeSet, err = azure.ListEntityPIMAssignments(ctx, db, options.SourceObjectID, 0, 0); err != nil { return nil, 0, err } case azure.RelatedEntityTypeRoleApprovers: - if nodeSet, err = azure.ListRoleApprovers(ctx, db, objectID, 0, 0); err != nil { + if nodeSet, err = azure.ListRoleApprovers(ctx, db, options.SourceObjectID, 0, 0); err != nil { return nil, 0, err } case azure.RelatedEntityTypeVaultKeyReaders, azure.RelatedEntityTypeVaultSecretReaders, azure.RelatedEntityTypeVaultCertReaders, azure.RelatedEntityTypeVaultAllReaders: - if nodeSet, err = azure.ListKeyVaultReaders(ctx, db, relatedEntityType, objectID, 0, 0); err != nil { + if nodeSet, err = azure.ListKeyVaultReaders(ctx, db, relatedEntityType, options.SourceObjectID, 0, 0); err != nil { return nil, 0, err } case azure.RelatedEntityTypeGroupMembers: - if nodeSet, err = azure.ListEntityGroupMembers(ctx, db, objectID, 0, 0); err != nil { + if nodeSet, err = azure.ListEntityGroupMembers(ctx, db, options.SourceObjectID, 0, 0); err != nil { return nil, 0, err } case azure.RelatedEntityTypeGroupMembership: - if nodeSet, err = azure.ListEntityGroupMembership(ctx, db, objectID, 0, 0); err != nil { + if nodeSet, err = azure.ListEntityGroupMembership(ctx, db, options.SourceObjectID, 0, 0); err != nil { return nil, 0, err } case azure.RelatedEntityTypeRoles: - if nodeSet, err = azure.ListEntityRoles(ctx, db, objectID, 0, 0); err != nil { + if nodeSet, err = azure.ListEntityRoles(ctx, db, options.SourceObjectID, 0, 0); err != nil { return nil, 0, err } case azure.RelatedEntityTypeEligibleAndApproverRoles: - if nodeSet, err = azure.ListEntityEligibleAndApproverRoles(ctx, db, objectID); err != nil { + if nodeSet, err = azure.ListEntityEligibleAndApproverRoles(ctx, db, options.SourceObjectID); err != nil { return nil, 0, err } case azure.RelatedEntityTypeOutboundExecutionPrivileges: - if nodeSet, err = azure.ListEntityExecutionPrivileges(ctx, db, objectID, graph.DirectionOutbound, 0, 0); err != nil { + if nodeSet, err = azure.ListEntityExecutionPrivileges(ctx, db, options.SourceObjectID, graph.DirectionOutbound, 0, 0); err != nil { return nil, 0, err } case azure.RelatedEntityTypeInboundExecutionPrivileges: - if nodeSet, err = azure.ListEntityExecutionPrivileges(ctx, db, objectID, graph.DirectionInbound, 0, 0); err != nil { + if nodeSet, err = azure.ListEntityExecutionPrivileges(ctx, db, options.SourceObjectID, graph.DirectionInbound, 0, 0); err != nil { return nil, 0, err } case azure.RelatedEntityTypeOutboundAbusableAppRoleAssignments: - if nodeSet, err = azure.ListEntityAbusableAppRoleAssignments(ctx, db, objectID, graph.DirectionOutbound, 0, 0); err != nil { + if nodeSet, err = azure.ListEntityAbusableAppRoleAssignments(ctx, db, options.SourceObjectID, graph.DirectionOutbound, 0, 0); err != nil { return nil, 0, err } case azure.RelatedEntityTypeInboundAbusableAppRoleAssignments: - if nodeSet, err = azure.ListEntityAbusableAppRoleAssignments(ctx, db, objectID, graph.DirectionInbound, 0, 0); err != nil { + if nodeSet, err = azure.ListEntityAbusableAppRoleAssignments(ctx, db, options.SourceObjectID, graph.DirectionInbound, 0, 0); err != nil { return nil, 0, err } case azure.RelatedEntityTypeOutboundControl: - if nodeSet, err = azure.ListEntityObjectControl(ctx, db, objectID, graph.DirectionOutbound, 0, 0); err != nil { + if nodeSet, err = azure.ListEntityObjectControl(ctx, db, options.SourceObjectID, graph.DirectionOutbound, 0, 0); err != nil { return nil, 0, err } case azure.RelatedEntityTypeInboundControl: - if nodeSet, err = azure.ListEntityObjectControl(ctx, db, objectID, graph.DirectionInbound, 0, 0); err != nil { + if nodeSet, err = azure.ListEntityObjectControl(ctx, db, options.SourceObjectID, graph.DirectionInbound, 0, 0); err != nil { return nil, 0, err } case azure.RelatedEntityTypeFederatedIdentityCredentials: - if nodeSet, err = azure.ListAppFederatedIdentityCredentials(ctx, db, objectID, 0, 0); err != nil { + if nodeSet, err = azure.ListAppFederatedIdentityCredentials(ctx, db, options.SourceObjectID, 0, 0); err != nil { return nil, 0, err } @@ -319,20 +322,20 @@ func listRelatedEntityType(ctx context.Context, db graph.Database, primaryDispla nodeCount := nodeSet.Len() - if skip > nodeCount { + if options.skip > nodeCount { return nil, 0, ErrParameterSkip } - if skip+limit > nodeCount { - limit = nodeSet.Len() - skip + if options.skip+options.limit > nodeCount { + options.limit = nodeSet.Len() - options.skip } - s := nodeSetToOrderedSlice(nodeSet)[skip : skip+limit] + s := nodeSetToOrderedSlice(nodeSet)[options.skip : options.skip+options.limit] return azure.FromGraphNodes(primaryDisplayKinds, s), nodeCount, nil } -func (s *Resources) GetAZRelatedEntities(ctx context.Context, response http.ResponseWriter, request *http.Request, objectID string) { +func (s *Resources) GetAZRelatedEntities(ctx context.Context, response http.ResponseWriter, request *http.Request, objectID string, sourceKind graph.Kind) { var ( queryParams = request.URL.Query() returnType = queryParams.Get(relatedEntityReturnTypeQueryParameterName) @@ -351,16 +354,30 @@ func (s *Resources) GetAZRelatedEntities(ctx context.Context, response http.Resp api.WriteErrorResponse(ctx, ErrBadQueryParameter(request, model.PaginationQueryParameterSkip, err), response) } else if limit, err := ParseLimitQueryParameter(queryParams, 100); err != nil { api.WriteErrorResponse(ctx, ErrBadQueryParameter(request, model.PaginationQueryParameterLimit, err), response) + } else if primaryDisplayKinds, err := s.DB.GetPrimaryDisplayKinds(ctx); err != nil { + api.HandleDatabaseError(request, response, err) } else if returnType == relatedEntityReturnTypeGraph { - if data, _, apiErr := graphRelatedEntityType(request, s.DB, s.Graph, relatedEntityType, objectID); apiErr != nil { + options := relatedEntityTypeOptions{ + RelatedEntityString: relatedEntityType, + SourceObjectID: objectID, + SourceKind: sourceKind, + skip: skip, + limit: limit, + } + if data, _, apiErr := graphRelatedEntityType(request, s.Graph, primaryDisplayKinds, options); apiErr != nil { api.WriteErrorResponse(ctx, apiErr, response) } else { api.WriteJSONResponse(ctx, data, http.StatusOK, response) } - } else if primaryDisplayKinds, err := s.DB.GetPrimaryDisplayKinds(request.Context()); err != nil { - api.HandleDatabaseError(request, response, err) } else { - if nodes, count, err := listRelatedEntityType(ctx, s.Graph, primaryDisplayKinds, relatedEntityType, objectID, skip, limit); err != nil { + options := relatedEntityTypeOptions{ + RelatedEntityString: relatedEntityType, + SourceObjectID: objectID, + SourceKind: sourceKind, + skip: skip, + limit: limit, + } + if nodes, count, err := listRelatedEntityType(ctx, s.Graph, primaryDisplayKinds, options); err != nil { if errors.Is(err, ErrParameterSkip) { api.WriteErrorResponse(ctx, api.BuildErrorResponse(http.StatusBadRequest, fmt.Sprintf(utils.ErrorInvalidSkip, skip), request), response) } else if errors.Is(err, ErrParameterRelatedEntityType) { @@ -457,7 +474,7 @@ func (s *Resources) GetAZEntity(response http.ResponseWriter, request *http.Requ } else if !hasAccess { api.WriteErrorResponse(request.Context(), api.BuildErrorResponse(http.StatusForbidden, api.ErrorResponseDetailsForbidden, request), response) } else if relatedEntityTypeStr := queryVars.Get(relatedEntityTypeQueryParameterName); relatedEntityTypeStr != "" { - s.GetAZRelatedEntities(request.Context(), response, request, objectID) + s.GetAZRelatedEntities(request.Context(), response, request, objectID, azKind) } else if includeCounts, err := api.ParseOptionalBool(queryVars.Get(api.QueryParameterIncludeCounts), true); err != nil { api.WriteErrorResponse(request.Context(), api.BuildErrorResponse(http.StatusBadRequest, api.ErrorResponseDetailsBadQueryParameterFilters, request), response) } else if entityInformation, err := GetAZEntityInformation(request.Context(), s.DB, s.Graph, entityType, objectID, includeCounts); err != nil { diff --git a/cmd/api/src/queries/graph.go b/cmd/api/src/queries/graph.go index d676ecc4b567..92b9c525aa61 100644 --- a/cmd/api/src/queries/graph.go +++ b/cmd/api/src/queries/graph.go @@ -1013,6 +1013,7 @@ func (s *GraphQuery) runListQuery(ctx context.Context, primaryDisplayKinds graph } func (s *GraphQuery) runCountQuery(ctx context.Context, node *graph.Node, params EntityQueryParameters, cacheEnabled bool) (any, int, error) { + params.RequestedType = model.DataTypeList result, err := s.runMaybeCachedEntityQuery(ctx, node, params, cacheEnabled) return nil, result.Len(), err } diff --git a/packages/go/analysis/azure/db_ops.go b/packages/go/analysis/azure/db_ops.go index 466fe69838f6..7f7126d0aeca 100644 --- a/packages/go/analysis/azure/db_ops.go +++ b/packages/go/analysis/azure/db_ops.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,19 @@ func ListEntityDescendentPaths(ctx context.Context, db graph.Database, relatedEn return ErrInvalidRelatedEntityType } - paths, err = FetchEntityDescendentPaths(tx, node, targetKind) - return err + if validEndKinds, ok := containsValidEndKinds[sourceKind]; ok && validEndKinds.ContainsOneOf(targetKind) { + paths, err = FetchDirectDescendentPaths(tx, node, targetKind) + return err + } 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,8 +133,14 @@ func ListEntityDescendents(ctx context.Context, db graph.Database, relatedEntity return ErrInvalidRelatedEntityType } - nodes, err = FetchEntityDescendents(tx, node, skip, limit, targetKind) - return err + if validEndKinds, ok := containsValidEndKinds[sourceKind]; ok && validEndKinds.ContainsOneOf(targetKind) { + nodes, err = FetchDirectEntityDescendents(tx, node, skip, limit, targetKind) + return err + } else { + nodes, err = FetchEntityDescendents(tx, node, skip, limit, targetKind) + return err + } + } }) } diff --git a/packages/go/analysis/azure/queries.go b/packages/go/analysis/azure/queries.go index 6faecec03167..53500bd57c51 100644 --- a/packages/go/analysis/azure/queries.go +++ b/packages/go/analysis/azure/queries.go @@ -18,6 +18,7 @@ package azure import ( "context" + "net/url" "log/slog" "strings" @@ -28,6 +29,7 @@ import ( "github.com/specterops/bloodhound/packages/go/graphschema/ad" "github.com/specterops/bloodhound/packages/go/graphschema/azure" "github.com/specterops/bloodhound/packages/go/graphschema/common" + "github.com/specterops/dawgs/cardinality" "github.com/specterops/dawgs/graph" "github.com/specterops/dawgs/ops" "github.com/specterops/dawgs/query" @@ -788,27 +790,159 @@ func FetchKeyVaultReaderCounts(tx graph.Transaction, keyVault *graph.Node) (KeyV return keyVaultReaders, nil } -func EntityDescendentsTraversal(root *graph.Node, _ ...graph.Kind) ops.TraversalPlan { - return ops.TraversalPlan{ - Root: root, - Direction: graph.DirectionOutbound, - BranchQuery: FilterContains, +// containsValidEndKinds maps each valid AZContains start node kind to the set of valid end node kinds. +// This is derived directly from the ingestion logic in bhce/packages/go/ein/azure.go. +var containsValidEndKinds = map[graph.Kind]graph.Kinds{ + azure.Tenant: { + azure.App, + azure.Device, + azure.Group, + azure.Role, + azure.ServicePrincipal, + azure.Subscription, + azure.User, + }, + azure.ManagementGroup: { + azure.ManagementGroup, + azure.Subscription, + }, + azure.Subscription: { + azure.ResourceGroup, + }, + azure.ResourceGroup: { + azure.VMScaleSet, + azure.FunctionApp, + azure.KeyVault, + azure.LogicApp, + azure.VM, + azure.ManagedCluster, + azure.ContainerRegistry, + azure.WebApp, + azure.AutomationAccount, + }, +} + +// nodeAzureTenantID returns the Azure tenant ID for a node. For Tenant nodes the tenant ID is +// stored in the common.ObjectID property; for all other Azure nodes it is stored in azure.TenantID. +func nodeAzureTenantID(node *graph.Node) (string, error) { + if node.Kinds.ContainsOneOf(azure.Tenant) { + return node.Properties.Get(common.ObjectID.String()).String() } + + return node.Properties.Get(azure.TenantID.String()).String() } -func FetchEntityDescendentPaths(tx graph.Transaction, root *graph.Node, descendentKinds ...graph.Kind) (graph.PathSet, error) { - return ops.TraverseIntermediaryPaths(tx, EntityDescendentsTraversal(root, descendentKinds...), func(node *graph.Node) bool { - return node.Kinds.ContainsOneOf(descendentKinds...) - }) +// FetchDirectDescendentPaths fetches the direct AZContains relationships from the root node. Both +// the valid start/end node kind pairing and the tenant-id match are enforced +func FetchDirectDescendentPaths(tx graph.Transaction, root *graph.Node, descendentKind ...graph.Kind) (graph.PathSet, error) { + if tenantID, err := nodeAzureTenantID(root); err != nil { + return nil, err + } else { + return ops.FetchPathSet(tx.Relationships().Filter(query.And( + query.Equals(query.StartID(), root.ID), + query.Kind(query.Relationship(), azure.Contains), + query.KindIn(query.End(), descendentKind...), + query.Equals(query.EndProperty(azure.TenantID.String()), tenantID), + ))) + } } -func FetchEntityDescendents(tx graph.Transaction, root *graph.Node, skip, limit int, descendentKinds ...graph.Kind) (graph.NodeSet, error) { - if paths, err := FetchEntityDescendentPaths(tx, root, descendentKinds...); err != nil { +func FetchDirectEntityDescendents(tx graph.Transaction, root *graph.Node, skip, limit int, descendentKind graph.Kind) (graph.NodeSet, error) { + if paths, err := FetchDirectDescendentPaths(tx, root, descendentKind); err != nil { return nil, err } else { nodes := paths.AllNodes() nodes.Remove(root.ID) - return nodes.ContainingNodeKinds(descendentKinds...), nil + return nodes, nil + } +} + +func FetchDescendentKindByTenantID(tx graph.Transaction, root *graph.Node, skip, limit int, descendentKind ...graph.Kind) (graph.NodeSet, error) { + if tenantID, err := nodeAzureTenantID(root); err != nil { + return nil, err + } else if nodes, err := ops.FetchNodeSet(tx.Nodes().Filter(query.And( + query.KindIn(query.Node(), descendentKind...), + query.Equals(query.NodeProperty(azure.TenantID.String()), tenantID), + ))); err != nil { + return nil, err + } else { + return nodes, nil + } +} + +// FetchEntityDescendentPaths fetches paths from each terminal node of descendentKind (within the +// same Azure tenant as root) back up to root via azure.Contains relationships. Each traversal +// halts upon reaching root or encountering a node already claimed by a prior traversal, preventing +// duplicate path segments across traversals. +func FetchEntityDescendentPaths(tx graph.Transaction, root *graph.Node, descendentKind ...graph.Kind) (graph.PathSet, error) { + var ( + visitedBitmap = cardinality.NewBitmap64() + pathSet = graph.NewPathSet() + ) + + // Pre-populate the visited bitmap with root so that any traversal halts when it reaches it + visitedBitmap.Add(root.ID.Uint64()) + + terminalNodes, err := FetchDescendentKindByTenantID(tx, root, 0, 0, descendentKind...) + if err != nil { + return nil, err + } + + for _, terminalNode := range terminalNodes { + // Skip this terminal if it was already reached and claimed by a prior traversal + if visitedBitmap.Contains(terminalNode.ID.Uint64()) { + continue + } + + if paths, err := ops.TraversePaths(tx, ops.TraversalPlan{ + Root: terminalNode, + Direction: graph.DirectionInbound, + BranchQuery: func() graph.Criteria { + return query.Kind(query.Relationship(), azure.Contains) + }, + ExpansionFilter: func(segment *graph.PathSegment) bool { + nodeID := segment.Node.ID.Uint64() + if visitedBitmap.Contains(nodeID) { + return false + } + visitedBitmap.Add(nodeID) + return true + }, + }); err != nil { + return nil, err + } else { + pathSet.AddPathSet(paths) + } + } + + return pathSet, nil +} + +func FetchEntityDescendents(tx graph.Transaction, root *graph.Node, skip, limit int, descendentKind graph.Kind) (graph.NodeSet, error) { + if paths, err := FetchEntityDescendentPaths(tx, root, descendentKind); err != nil { + return nil, err + } else { + nodes := paths.AllNodes() + nodes.Remove(root.ID) + return nodes.ContainingNodeKinds(descendentKind), nil + } +} + +func FetchDirectDescendentCounts(tx graph.Transaction, root *graph.Node, skip, limit int, descendentKinds ...graph.Kind) (Descendents, error) { + if paths, err := FetchDirectDescendentPaths(tx, root, descendentKinds...); err != nil { + return Descendents{}, err + } else { + details := Descendents{ + DescendentCounts: map[string]int{}, + } + kindSet := paths.AllNodes().KindSet() + kindSet.RemoveNode(root.ID) + + for _, kind := range descendentKinds { + details.DescendentCounts[kind.String()] = int(kindSet.Count(kind)) + } + + return details, nil } } @@ -846,10 +980,14 @@ func fetchRolesTraversalPlan(root *graph.Node) ops.TraversalPlan { // FetchEntityByObjectID pulls a node by its ObjectID. It requires a kind to perform index lookups against. func FetchEntityByObjectID(tx graph.Transaction, objectID string) (*graph.Node, error) { - return tx.Nodes().Filterf(func() graph.Criteria { - return query.And( - query.Kind(query.Node(), azure.Entity), - query.Equals(query.NodeProperty(common.ObjectID.String()), objectID), - ) - }).First() + if decodedObjectID, err := url.QueryUnescape(objectID); err != nil { + return nil, err + } else { + return tx.Nodes().Filterf(func() graph.Criteria { + return query.And( + query.Kind(query.Node(), azure.Entity), + query.Equals(query.NodeProperty(common.ObjectID.String()), decodedObjectID), + ) + }).First() + } } diff --git a/packages/go/analysis/azure/resource_group.go b/packages/go/analysis/azure/resource_group.go index 4330bb44bfcc..efdb0b5b520c 100644 --- a/packages/go/analysis/azure/resource_group.go +++ b/packages/go/analysis/azure/resource_group.go @@ -43,7 +43,7 @@ func ResourceGroupEntityDetails(ctx context.Context, db graph.Database, primaryD func PopulateResourceGroupEntityDetailsCounts(tx graph.Transaction, node *graph.Node, details ResourceGroupDetails) (ResourceGroupDetails, error) { var descendentKinds = GetDescendentKinds(azure.ResourceGroup) - if descendents, err := FetchEntityDescendentCounts(tx, node, 0, 0, descendentKinds...); err != nil { + if descendents, err := FetchDirectDescendentCounts(tx, node, 0, 0, descendentKinds...); err != nil { return details, err } else { details.Descendents = descendents diff --git a/packages/javascript/js-client-library/src/client.ts b/packages/javascript/js-client-library/src/client.ts index 6c18df122cce..d8dc2680111c 100644 --- a/packages/javascript/js-client-library/src/client.ts +++ b/packages/javascript/js-client-library/src/client.ts @@ -1235,7 +1235,7 @@ class BHEAPIClient { Object.assign( { params: { - object_id: id, + object_id: encodeURIComponent(id), related_entity_type: relatedEntityType, counts, skip, From e623bcc3f3b0f591b6bd2fea0603025a5e8a2139 Mon Sep 17 00:00:00 2001 From: Ulises Rangel Date: Mon, 15 Jun 2026 12:22:23 -0500 Subject: [PATCH 2/6] test: integration tests --- .../analysis/azure/azure_integration_test.go | 128 ++++++++++++++++++ 1 file changed, 128 insertions(+) diff --git a/packages/go/analysis/azure/azure_integration_test.go b/packages/go/analysis/azure/azure_integration_test.go index 74ebeef1b571..873beb5a8b23 100644 --- a/packages/go/analysis/azure/azure_integration_test.go +++ b/packages/go/analysis/azure/azure_integration_test.go @@ -949,3 +949,131 @@ 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_DifferentTenantExcluded verifies that nodes sharing the same kind +// but belonging to a different Azure tenant are not included in the path set. +func TestFetchEntityDescendentPaths_DifferentTenantExcluded(t *testing.T) { + t.Parallel() + + suite := setupIntegrationTestSuite(t) + defer teardownIntegrationTestSuite(t, &suite) + + var ( + tenantID = integration.RandomObjectID(t) + otherTenantID = integration.RandomObjectID(t) + tenantNode = NewAzureTenant(t, &suite, tenantID) + ownApp = NewAzureApplication(t, &suite, "OwnApp", integration.RandomObjectID(t), tenantID) + foreignApp = NewAzureApplication(t, &suite, "ForeignApp", integration.RandomObjectID(t), otherTenantID) + ) + + NewRelationship(t, &suite, tenantNode, ownApp, graphAzure.Contains) + // foreignApp is deliberately not connected; it belongs to a different tenant. + + 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, ownApp.ID) + require.NotContains(t, nodeIDs, foreignApp.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) +} From a965703ecf07a8a9a517c6af854baffc6f76cf16 Mon Sep 17 00:00:00 2001 From: Ulises Rangel Date: Mon, 15 Jun 2026 16:15:10 -0500 Subject: [PATCH 3/6] chore: cleanup --- cmd/api/src/api/v2/azure.go | 142 +++++++++--------- cmd/api/src/api/v2/azure_test.go | 2 +- cmd/api/src/queries/graph.go | 1 - .../analysis/azure/azure_integration_test.go | 33 ---- packages/go/analysis/azure/db_ops.go | 16 +- .../go/analysis/azure/management_group.go | 2 +- packages/go/analysis/azure/queries.go | 54 +++---- packages/go/analysis/azure/resource_group.go | 2 +- packages/go/analysis/azure/subscription.go | 2 +- packages/go/analysis/azure/tenant.go | 2 +- .../EntityInfoDataTableGraphed.test.tsx | 7 +- .../EntityInfoDataTableGraphed.tsx | 42 +++--- .../js-client-library/src/client.ts | 2 +- 13 files changed, 139 insertions(+), 168 deletions(-) diff --git a/cmd/api/src/api/v2/azure.go b/cmd/api/src/api/v2/azure.go index 6f2f6184d60c..21b48164113f 100644 --- a/cmd/api/src/api/v2/azure.go +++ b/cmd/api/src/api/v2/azure.go @@ -79,9 +79,12 @@ var ( ) func graphRelatedEntityType(request *http.Request, graphDb graph.Database, primaryDisplayKinds graphschema.PrimaryDisplayKinds, options relatedEntityTypeOptions) (any, int, *api.ErrorWrapper) { - ctx := request.Context() + var ( + ctx = request.Context() + objectID = options.sourceObjectID + ) - switch relatedEntityType := azure.RelatedEntityType(options.RelatedEntityString); relatedEntityType { + switch relatedEntityType := azure.RelatedEntityType(options.relatedEntityString); relatedEntityType { case azure.RelatedEntityTypeDescendentUsers, azure.RelatedEntityTypeDescendentGroups, azure.RelatedEntityTypeDescendentManagementGroups, azure.RelatedEntityTypeDescendentSubscriptions, azure.RelatedEntityTypeDescendentResourceGroups, azure.RelatedEntityTypeDescendentVirtualMachines, @@ -93,116 +96,116 @@ func graphRelatedEntityType(request *http.Request, graphDb graph.Database, prima azure.RelatedEntityTypeDescendentWebApps, azure.RelatedEntityTypeDescendentAutomationAccounts, azure.RelatedEntityTypeDescendentLogicApps, azure.RelatedEntityTypeDescendentFunctionApps: - if descendents, err := azure.ListEntityDescendentPaths(ctx, graphDb, relatedEntityType, options.SourceKind, options.SourceObjectID); err != nil { - return nil, 0, api.BuildErrorResponse(http.StatusInternalServerError, fmt.Sprintf("error fetching related entity type %s: %v", options.RelatedEntityString, err), request) + if descendents, err := azure.ListEntityDescendentPaths(ctx, graphDb, relatedEntityType, options.sourceKind, objectID); err != nil { + return nil, 0, api.BuildErrorResponse(http.StatusInternalServerError, fmt.Sprintf("error fetching related entity type %s: %v", options.relatedEntityString, err), request) } else { return bloodhoundgraph.PathSetToBloodHoundGraph(primaryDisplayKinds, descendents), descendents.Len(), nil } case azure.RelatedEntityTypeActiveAssignments: - if assignments, err := azure.ListEntityActiveAssignmentPaths(ctx, graphDb, options.SourceObjectID); err != nil { - return nil, 0, api.BuildErrorResponse(http.StatusInternalServerError, fmt.Sprintf("error fetching related entity type %s: %v", options.RelatedEntityString, err), request) + if assignments, err := azure.ListEntityActiveAssignmentPaths(ctx, graphDb, objectID); err != nil { + return nil, 0, api.BuildErrorResponse(http.StatusInternalServerError, fmt.Sprintf("error fetching related entity type %s: %v", options.relatedEntityString, err), request) } else { return bloodhoundgraph.PathSetToBloodHoundGraph(primaryDisplayKinds, assignments), assignments.Len(), nil } case azure.RelatedEntityTypePIMAssignments: - if assignments, err := azure.ListEntityPIMAssignmentPaths(ctx, graphDb, options.SourceObjectID); err != nil { - return nil, 0, api.BuildErrorResponse(http.StatusInternalServerError, fmt.Sprintf("error fetching related entity type %s: %v", options.RelatedEntityString, err), request) + if assignments, err := azure.ListEntityPIMAssignmentPaths(ctx, graphDb, objectID); err != nil { + return nil, 0, api.BuildErrorResponse(http.StatusInternalServerError, fmt.Sprintf("error fetching related entity type %s: %v", options.relatedEntityString, err), request) } else { return bloodhoundgraph.PathSetToBloodHoundGraph(primaryDisplayKinds, assignments), assignments.Len(), nil } case azure.RelatedEntityTypeRoleApprovers: - if approvers, err := azure.ListRoleApproverPaths(ctx, graphDb, options.SourceObjectID); err != nil { - return nil, 0, api.BuildErrorResponse(http.StatusInternalServerError, fmt.Sprintf("error fetching related entity type %s: %v", options.RelatedEntityString, err), request) + if approvers, err := azure.ListRoleApproverPaths(ctx, graphDb, objectID); err != nil { + return nil, 0, api.BuildErrorResponse(http.StatusInternalServerError, fmt.Sprintf("error fetching related entity type %s: %v", options.relatedEntityString, err), request) } else { return bloodhoundgraph.PathSetToBloodHoundGraph(primaryDisplayKinds, approvers), approvers.Len(), nil } case azure.RelatedEntityTypeVaultKeyReaders, azure.RelatedEntityTypeVaultSecretReaders, azure.RelatedEntityTypeVaultCertReaders, azure.RelatedEntityTypeVaultAllReaders: - if groupMembers, err := azure.ListKeyVaultReaderPaths(ctx, graphDb, relatedEntityType, options.SourceObjectID); err != nil { - return nil, 0, api.BuildErrorResponse(http.StatusInternalServerError, fmt.Sprintf("error fetching related entity type %s: %v", options.RelatedEntityString, err), request) + if groupMembers, err := azure.ListKeyVaultReaderPaths(ctx, graphDb, relatedEntityType, objectID); err != nil { + return nil, 0, api.BuildErrorResponse(http.StatusInternalServerError, fmt.Sprintf("error fetching related entity type %s: %v", options.relatedEntityString, err), request) } else { return bloodhoundgraph.PathSetToBloodHoundGraph(primaryDisplayKinds, groupMembers), groupMembers.Len(), nil } case azure.RelatedEntityTypeGroupMembers: - if groupMembers, err := azure.ListEntityGroupMemberPaths(ctx, graphDb, options.SourceObjectID); err != nil { - return nil, 0, api.BuildErrorResponse(http.StatusInternalServerError, fmt.Sprintf("error fetching related entity type %s: %v", options.RelatedEntityString, err), request) + if groupMembers, err := azure.ListEntityGroupMemberPaths(ctx, graphDb, objectID); err != nil { + return nil, 0, api.BuildErrorResponse(http.StatusInternalServerError, fmt.Sprintf("error fetching related entity type %s: %v", options.relatedEntityString, err), request) } else { return bloodhoundgraph.PathSetToBloodHoundGraph(primaryDisplayKinds, groupMembers), groupMembers.Len(), nil } case azure.RelatedEntityTypeGroupMembership: - if groupMembership, err := azure.ListEntityGroupMembershipPaths(ctx, graphDb, options.SourceObjectID); err != nil { - return nil, 0, api.BuildErrorResponse(http.StatusInternalServerError, fmt.Sprintf("error fetching related entity type %s: %v", options.RelatedEntityString, err), request) + if groupMembership, err := azure.ListEntityGroupMembershipPaths(ctx, graphDb, objectID); err != nil { + return nil, 0, api.BuildErrorResponse(http.StatusInternalServerError, fmt.Sprintf("error fetching related entity type %s: %v", options.relatedEntityString, err), request) } else { return bloodhoundgraph.PathSetToBloodHoundGraph(primaryDisplayKinds, groupMembership), groupMembership.Len(), nil } case azure.RelatedEntityTypeRoles: - if userRoles, err := azure.ListEntityRolePaths(ctx, graphDb, options.SourceObjectID); err != nil { - return nil, 0, api.BuildErrorResponse(http.StatusInternalServerError, fmt.Sprintf("error fetching related entity type %s: %v", options.RelatedEntityString, err), request) + if userRoles, err := azure.ListEntityRolePaths(ctx, graphDb, objectID); err != nil { + return nil, 0, api.BuildErrorResponse(http.StatusInternalServerError, fmt.Sprintf("error fetching related entity type %s: %v", options.relatedEntityString, err), request) } else { return bloodhoundgraph.PathSetToBloodHoundGraph(primaryDisplayKinds, userRoles), userRoles.Len(), nil } case azure.RelatedEntityTypeEligibleAndApproverRoles: - if eligibleAndApproverRoles, err := azure.ListEntityEligibleAndApproverRolePaths(ctx, graphDb, options.SourceObjectID); err != nil { - return nil, 0, api.BuildErrorResponse(http.StatusInternalServerError, fmt.Sprintf("error fetching related entity type %s: %v", options.RelatedEntityString, err), request) + if eligibleAndApproverRoles, err := azure.ListEntityEligibleAndApproverRolePaths(ctx, graphDb, objectID); err != nil { + return nil, 0, api.BuildErrorResponse(http.StatusInternalServerError, fmt.Sprintf("error fetching related entity type %s: %v", options.relatedEntityString, err), request) } else { return bloodhoundgraph.PathSetToBloodHoundGraph(primaryDisplayKinds, eligibleAndApproverRoles), eligibleAndApproverRoles.Len(), nil } case azure.RelatedEntityTypeOutboundExecutionPrivileges: - if executionPrivileges, err := azure.ListEntityExecutionPrivilegePaths(ctx, graphDb, options.SourceObjectID, graph.DirectionOutbound); err != nil { - return nil, 0, api.BuildErrorResponse(http.StatusInternalServerError, fmt.Sprintf("error fetching related entity type %s: %v", options.RelatedEntityString, err), request) + if executionPrivileges, err := azure.ListEntityExecutionPrivilegePaths(ctx, graphDb, objectID, graph.DirectionOutbound); err != nil { + return nil, 0, api.BuildErrorResponse(http.StatusInternalServerError, fmt.Sprintf("error fetching related entity type %s: %v", options.relatedEntityString, err), request) } else { return bloodhoundgraph.PathSetToBloodHoundGraph(primaryDisplayKinds, executionPrivileges), executionPrivileges.Len(), nil } case azure.RelatedEntityTypeInboundExecutionPrivileges: - if executionPrivileges, err := azure.ListEntityExecutionPrivilegePaths(ctx, graphDb, options.SourceObjectID, graph.DirectionInbound); err != nil { - return nil, 0, api.BuildErrorResponse(http.StatusInternalServerError, fmt.Sprintf("error fetching related entity type %s: %v", options.RelatedEntityString, err), request) + if executionPrivileges, err := azure.ListEntityExecutionPrivilegePaths(ctx, graphDb, objectID, graph.DirectionInbound); err != nil { + return nil, 0, api.BuildErrorResponse(http.StatusInternalServerError, fmt.Sprintf("error fetching related entity type %s: %v", options.relatedEntityString, err), request) } else { return bloodhoundgraph.PathSetToBloodHoundGraph(primaryDisplayKinds, executionPrivileges), executionPrivileges.Len(), nil } case azure.RelatedEntityTypeOutboundAbusableAppRoleAssignments: - if objectControl, err := azure.ListEntityAbusableAppRoleAssignmentsPaths(ctx, graphDb, options.SourceObjectID, graph.DirectionOutbound); err != nil { - return nil, 0, api.BuildErrorResponse(http.StatusInternalServerError, fmt.Sprintf("error fetching related entity type %s: %v", options.RelatedEntityString, err), request) + if objectControl, err := azure.ListEntityAbusableAppRoleAssignmentsPaths(ctx, graphDb, objectID, graph.DirectionOutbound); err != nil { + return nil, 0, api.BuildErrorResponse(http.StatusInternalServerError, fmt.Sprintf("error fetching related entity type %s: %v", options.relatedEntityString, err), request) } else { return bloodhoundgraph.PathSetToBloodHoundGraph(primaryDisplayKinds, objectControl), objectControl.Len(), nil } case azure.RelatedEntityTypeInboundAbusableAppRoleAssignments: - if objectControl, err := azure.ListEntityAbusableAppRoleAssignmentsPaths(ctx, graphDb, options.SourceObjectID, graph.DirectionInbound); err != nil { - return nil, 0, api.BuildErrorResponse(http.StatusInternalServerError, fmt.Sprintf("error fetching related entity type %s: %v", options.RelatedEntityString, err), request) + if objectControl, err := azure.ListEntityAbusableAppRoleAssignmentsPaths(ctx, graphDb, objectID, graph.DirectionInbound); err != nil { + return nil, 0, api.BuildErrorResponse(http.StatusInternalServerError, fmt.Sprintf("error fetching related entity type %s: %v", options.relatedEntityString, err), request) } else { return bloodhoundgraph.PathSetToBloodHoundGraph(primaryDisplayKinds, objectControl), objectControl.Len(), nil } case azure.RelatedEntityTypeOutboundControl: - if objectControl, err := azure.ListEntityObjectControlPaths(ctx, graphDb, options.SourceObjectID, graph.DirectionOutbound); err != nil { - return nil, 0, api.BuildErrorResponse(http.StatusInternalServerError, fmt.Sprintf("error fetching related entity type %s: %v", options.RelatedEntityString, err), request) + if objectControl, err := azure.ListEntityObjectControlPaths(ctx, graphDb, objectID, graph.DirectionOutbound); err != nil { + return nil, 0, api.BuildErrorResponse(http.StatusInternalServerError, fmt.Sprintf("error fetching related entity type %s: %v", options.relatedEntityString, err), request) } else { return bloodhoundgraph.PathSetToBloodHoundGraph(primaryDisplayKinds, objectControl), objectControl.Len(), nil } case azure.RelatedEntityTypeInboundControl: - if objectControl, err := azure.ListEntityObjectControlPaths(ctx, graphDb, options.SourceObjectID, graph.DirectionInbound); err != nil { - return nil, 0, api.BuildErrorResponse(http.StatusInternalServerError, fmt.Sprintf("error fetching related entity type %s: %v", options.RelatedEntityString, err), request) + if objectControl, err := azure.ListEntityObjectControlPaths(ctx, graphDb, objectID, graph.DirectionInbound); err != nil { + return nil, 0, api.BuildErrorResponse(http.StatusInternalServerError, fmt.Sprintf("error fetching related entity type %s: %v", options.relatedEntityString, err), request) } else { return bloodhoundgraph.PathSetToBloodHoundGraph(primaryDisplayKinds, objectControl), objectControl.Len(), nil } case azure.RelatedEntityTypeFederatedIdentityCredentials: - if fics, err := azure.ListAppFederatedIdentityCredentialPaths(ctx, graphDb, options.SourceObjectID); err != nil { - return nil, 0, api.BuildErrorResponse(http.StatusInternalServerError, fmt.Sprintf("error fetching related entity type %s: %v", options.RelatedEntityString, err), request) + if fics, err := azure.ListAppFederatedIdentityCredentialPaths(ctx, graphDb, objectID); err != nil { + return nil, 0, api.BuildErrorResponse(http.StatusInternalServerError, fmt.Sprintf("error fetching related entity type %s: %v", options.relatedEntityString, err), request) } else { return bloodhoundgraph.PathSetToBloodHoundGraph(primaryDisplayKinds, fics), fics.Len(), nil } default: - return nil, 0, api.BuildErrorResponse(http.StatusNotFound, fmt.Sprintf("no matching related entity list type for %s", options.RelatedEntityString), request) + return nil, 0, api.BuildErrorResponse(http.StatusNotFound, fmt.Sprintf("no matching related entity list type for %s", options.relatedEntityString), request) } } @@ -217,20 +220,23 @@ func nodeSetToOrderedSlice(nodeSet graph.NodeSet) []*graph.Node { } type relatedEntityTypeOptions struct { - RelatedEntityString string - SourceKind graph.Kind - SourceObjectID string + relatedEntityString string + sourceKind graph.Kind + sourceObjectID string skip int limit int } func listRelatedEntityType(ctx context.Context, db graph.Database, primaryDisplayKinds graphschema.PrimaryDisplayKinds, options relatedEntityTypeOptions) ([]azure.Node, int, error) { var ( - nodeSet graph.NodeSet - err error + nodeSet graph.NodeSet + err error + objectID = options.sourceObjectID + skip = options.skip + limit = options.limit ) // NOTE: All skip/limit passed to lower level queries is currently hardcoded to 0 so we can get the full count of the dataset for skip/limit tracking - switch relatedEntityType := azure.RelatedEntityType(options.RelatedEntityString); relatedEntityType { + switch relatedEntityType := azure.RelatedEntityType(options.relatedEntityString); relatedEntityType { case azure.RelatedEntityTypeDescendentUsers, azure.RelatedEntityTypeDescendentGroups, azure.RelatedEntityTypeDescendentManagementGroups, azure.RelatedEntityTypeDescendentSubscriptions, azure.RelatedEntityTypeDescendentResourceGroups, azure.RelatedEntityTypeDescendentVirtualMachines, @@ -242,77 +248,77 @@ func listRelatedEntityType(ctx context.Context, db graph.Database, primaryDispla azure.RelatedEntityTypeDescendentWebApps, azure.RelatedEntityTypeDescendentLogicApps, azure.RelatedEntityTypeDescendentFunctionApps, azure.RelatedEntityTypeDescendentAutomationAccounts: - if nodeSet, err = azure.ListEntityDescendents(ctx, db, relatedEntityType, options.SourceKind, options.SourceObjectID, 0, 0); err != nil { + if nodeSet, err = azure.ListEntityDescendents(ctx, db, relatedEntityType, options.sourceKind, objectID, 0, 0); err != nil { return nil, 0, err } case azure.RelatedEntityTypeActiveAssignments: - if nodeSet, err = azure.ListEntityActiveAssignments(ctx, db, options.SourceObjectID, 0, 0); err != nil { + if nodeSet, err = azure.ListEntityActiveAssignments(ctx, db, objectID, 0, 0); err != nil { return nil, 0, err } case azure.RelatedEntityTypePIMAssignments: - if nodeSet, err = azure.ListEntityPIMAssignments(ctx, db, options.SourceObjectID, 0, 0); err != nil { + if nodeSet, err = azure.ListEntityPIMAssignments(ctx, db, objectID, 0, 0); err != nil { return nil, 0, err } case azure.RelatedEntityTypeRoleApprovers: - if nodeSet, err = azure.ListRoleApprovers(ctx, db, options.SourceObjectID, 0, 0); err != nil { + if nodeSet, err = azure.ListRoleApprovers(ctx, db, objectID, 0, 0); err != nil { return nil, 0, err } case azure.RelatedEntityTypeVaultKeyReaders, azure.RelatedEntityTypeVaultSecretReaders, azure.RelatedEntityTypeVaultCertReaders, azure.RelatedEntityTypeVaultAllReaders: - if nodeSet, err = azure.ListKeyVaultReaders(ctx, db, relatedEntityType, options.SourceObjectID, 0, 0); err != nil { + if nodeSet, err = azure.ListKeyVaultReaders(ctx, db, relatedEntityType, objectID, 0, 0); err != nil { return nil, 0, err } case azure.RelatedEntityTypeGroupMembers: - if nodeSet, err = azure.ListEntityGroupMembers(ctx, db, options.SourceObjectID, 0, 0); err != nil { + if nodeSet, err = azure.ListEntityGroupMembers(ctx, db, objectID, 0, 0); err != nil { return nil, 0, err } case azure.RelatedEntityTypeGroupMembership: - if nodeSet, err = azure.ListEntityGroupMembership(ctx, db, options.SourceObjectID, 0, 0); err != nil { + if nodeSet, err = azure.ListEntityGroupMembership(ctx, db, objectID, 0, 0); err != nil { return nil, 0, err } case azure.RelatedEntityTypeRoles: - if nodeSet, err = azure.ListEntityRoles(ctx, db, options.SourceObjectID, 0, 0); err != nil { + if nodeSet, err = azure.ListEntityRoles(ctx, db, objectID, 0, 0); err != nil { return nil, 0, err } case azure.RelatedEntityTypeEligibleAndApproverRoles: - if nodeSet, err = azure.ListEntityEligibleAndApproverRoles(ctx, db, options.SourceObjectID); err != nil { + if nodeSet, err = azure.ListEntityEligibleAndApproverRoles(ctx, db, objectID); err != nil { return nil, 0, err } case azure.RelatedEntityTypeOutboundExecutionPrivileges: - if nodeSet, err = azure.ListEntityExecutionPrivileges(ctx, db, options.SourceObjectID, graph.DirectionOutbound, 0, 0); err != nil { + if nodeSet, err = azure.ListEntityExecutionPrivileges(ctx, db, objectID, graph.DirectionOutbound, 0, 0); err != nil { return nil, 0, err } case azure.RelatedEntityTypeInboundExecutionPrivileges: - if nodeSet, err = azure.ListEntityExecutionPrivileges(ctx, db, options.SourceObjectID, graph.DirectionInbound, 0, 0); err != nil { + if nodeSet, err = azure.ListEntityExecutionPrivileges(ctx, db, objectID, graph.DirectionInbound, 0, 0); err != nil { return nil, 0, err } case azure.RelatedEntityTypeOutboundAbusableAppRoleAssignments: - if nodeSet, err = azure.ListEntityAbusableAppRoleAssignments(ctx, db, options.SourceObjectID, graph.DirectionOutbound, 0, 0); err != nil { + if nodeSet, err = azure.ListEntityAbusableAppRoleAssignments(ctx, db, objectID, graph.DirectionOutbound, 0, 0); err != nil { return nil, 0, err } case azure.RelatedEntityTypeInboundAbusableAppRoleAssignments: - if nodeSet, err = azure.ListEntityAbusableAppRoleAssignments(ctx, db, options.SourceObjectID, graph.DirectionInbound, 0, 0); err != nil { + if nodeSet, err = azure.ListEntityAbusableAppRoleAssignments(ctx, db, objectID, graph.DirectionInbound, 0, 0); err != nil { return nil, 0, err } case azure.RelatedEntityTypeOutboundControl: - if nodeSet, err = azure.ListEntityObjectControl(ctx, db, options.SourceObjectID, graph.DirectionOutbound, 0, 0); err != nil { + if nodeSet, err = azure.ListEntityObjectControl(ctx, db, objectID, graph.DirectionOutbound, 0, 0); err != nil { return nil, 0, err } case azure.RelatedEntityTypeInboundControl: - if nodeSet, err = azure.ListEntityObjectControl(ctx, db, options.SourceObjectID, graph.DirectionInbound, 0, 0); err != nil { + if nodeSet, err = azure.ListEntityObjectControl(ctx, db, objectID, graph.DirectionInbound, 0, 0); err != nil { return nil, 0, err } case azure.RelatedEntityTypeFederatedIdentityCredentials: - if nodeSet, err = azure.ListAppFederatedIdentityCredentials(ctx, db, options.SourceObjectID, 0, 0); err != nil { + if nodeSet, err = azure.ListAppFederatedIdentityCredentials(ctx, db, objectID, 0, 0); err != nil { return nil, 0, err } @@ -322,15 +328,15 @@ func listRelatedEntityType(ctx context.Context, db graph.Database, primaryDispla nodeCount := nodeSet.Len() - if options.skip > nodeCount { + if skip > nodeCount { return nil, 0, ErrParameterSkip } - if options.skip+options.limit > nodeCount { - options.limit = nodeSet.Len() - options.skip + if skip+limit > nodeCount { + limit = nodeSet.Len() - skip } - s := nodeSetToOrderedSlice(nodeSet)[options.skip : options.skip+options.limit] + s := nodeSetToOrderedSlice(nodeSet)[skip : skip+limit] return azure.FromGraphNodes(primaryDisplayKinds, s), nodeCount, nil } @@ -358,11 +364,9 @@ func (s *Resources) GetAZRelatedEntities(ctx context.Context, response http.Resp api.HandleDatabaseError(request, response, err) } else if returnType == relatedEntityReturnTypeGraph { options := relatedEntityTypeOptions{ - RelatedEntityString: relatedEntityType, - SourceObjectID: objectID, - SourceKind: sourceKind, - skip: skip, - limit: limit, + relatedEntityString: relatedEntityType, + sourceObjectID: objectID, + sourceKind: sourceKind, } if data, _, apiErr := graphRelatedEntityType(request, s.Graph, primaryDisplayKinds, options); apiErr != nil { api.WriteErrorResponse(ctx, apiErr, response) @@ -371,9 +375,9 @@ func (s *Resources) GetAZRelatedEntities(ctx context.Context, response http.Resp } } else { options := relatedEntityTypeOptions{ - RelatedEntityString: relatedEntityType, - SourceObjectID: objectID, - SourceKind: sourceKind, + relatedEntityString: relatedEntityType, + sourceObjectID: objectID, + sourceKind: sourceKind, skip: skip, limit: limit, } diff --git a/cmd/api/src/api/v2/azure_test.go b/cmd/api/src/api/v2/azure_test.go index 4b06521b488c..565e7c17990f 100644 --- a/cmd/api/src/api/v2/azure_test.go +++ b/cmd/api/src/api/v2/azure_test.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"}}, }, }, diff --git a/cmd/api/src/queries/graph.go b/cmd/api/src/queries/graph.go index 92b9c525aa61..d676ecc4b567 100644 --- a/cmd/api/src/queries/graph.go +++ b/cmd/api/src/queries/graph.go @@ -1013,7 +1013,6 @@ func (s *GraphQuery) runListQuery(ctx context.Context, primaryDisplayKinds graph } func (s *GraphQuery) runCountQuery(ctx context.Context, node *graph.Node, params EntityQueryParameters, cacheEnabled bool) (any, int, error) { - params.RequestedType = model.DataTypeList result, err := s.runMaybeCachedEntityQuery(ctx, node, params, cacheEnabled) return nil, result.Len(), err } diff --git a/packages/go/analysis/azure/azure_integration_test.go b/packages/go/analysis/azure/azure_integration_test.go index 873beb5a8b23..fc6ae3bf3eae 100644 --- a/packages/go/analysis/azure/azure_integration_test.go +++ b/packages/go/analysis/azure/azure_integration_test.go @@ -1022,39 +1022,6 @@ func TestFetchEntityDescendentPaths_MultiHopPathToRoot(t *testing.T) { require.NoError(t, err) } -// TestFetchEntityDescendentPaths_DifferentTenantExcluded verifies that nodes sharing the same kind -// but belonging to a different Azure tenant are not included in the path set. -func TestFetchEntityDescendentPaths_DifferentTenantExcluded(t *testing.T) { - t.Parallel() - - suite := setupIntegrationTestSuite(t) - defer teardownIntegrationTestSuite(t, &suite) - - var ( - tenantID = integration.RandomObjectID(t) - otherTenantID = integration.RandomObjectID(t) - tenantNode = NewAzureTenant(t, &suite, tenantID) - ownApp = NewAzureApplication(t, &suite, "OwnApp", integration.RandomObjectID(t), tenantID) - foreignApp = NewAzureApplication(t, &suite, "ForeignApp", integration.RandomObjectID(t), otherTenantID) - ) - - NewRelationship(t, &suite, tenantNode, ownApp, graphAzure.Contains) - // foreignApp is deliberately not connected; it belongs to a different tenant. - - 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, ownApp.ID) - require.NotContains(t, nodeIDs, foreignApp.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) { diff --git a/packages/go/analysis/azure/db_ops.go b/packages/go/analysis/azure/db_ops.go index 7f7126d0aeca..0466ff57943e 100644 --- a/packages/go/analysis/azure/db_ops.go +++ b/packages/go/analysis/azure/db_ops.go @@ -73,14 +73,12 @@ func ListEntityDescendentPaths(ctx context.Context, db graph.Database, relatedEn return ErrInvalidRelatedEntityType } - if validEndKinds, ok := containsValidEndKinds[sourceKind]; ok && validEndKinds.ContainsOneOf(targetKind) { + if isDirectDescendent(sourceKind, targetKind) { paths, err = FetchDirectDescendentPaths(tx, node, targetKind) - return err } else { paths, err = FetchEntityDescendentPaths(tx, node, targetKind) - return err } - + return err } }) } @@ -133,14 +131,12 @@ func ListEntityDescendents(ctx context.Context, db graph.Database, relatedEntity return ErrInvalidRelatedEntityType } - if validEndKinds, ok := containsValidEndKinds[sourceKind]; ok && validEndKinds.ContainsOneOf(targetKind) { - nodes, err = FetchDirectEntityDescendents(tx, node, skip, limit, targetKind) - return err + if isDirectDescendent(sourceKind, targetKind) { + nodes, err = FetchDirectEntityDescendents(tx, node, targetKind) } else { - nodes, err = FetchEntityDescendents(tx, node, skip, limit, targetKind) - return err + nodes, err = FetchEntityDescendents(tx, node, targetKind) } - + return err } }) } diff --git a/packages/go/analysis/azure/management_group.go b/packages/go/analysis/azure/management_group.go index eb99b1af2654..013f4e15400d 100644 --- a/packages/go/analysis/azure/management_group.go +++ b/packages/go/analysis/azure/management_group.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 diff --git a/packages/go/analysis/azure/queries.go b/packages/go/analysis/azure/queries.go index 53500bd57c51..8611e07b5c17 100644 --- a/packages/go/analysis/azure/queries.go +++ b/packages/go/analysis/azure/queries.go @@ -18,7 +18,6 @@ package azure import ( "context" - "net/url" "log/slog" "strings" @@ -822,6 +821,13 @@ var containsValidEndKinds = map[graph.Kind]graph.Kinds{ }, } +func isDirectDescendent(sourceKind, targetKind graph.Kind) bool { + if validEndKinds, ok := containsValidEndKinds[sourceKind]; ok && validEndKinds.ContainsOneOf(targetKind) { + return true + } + return false +} + // nodeAzureTenantID returns the Azure tenant ID for a node. For Tenant nodes the tenant ID is // stored in the common.ObjectID property; for all other Azure nodes it is stored in azure.TenantID. func nodeAzureTenantID(node *graph.Node) (string, error) { @@ -832,22 +838,16 @@ func nodeAzureTenantID(node *graph.Node) (string, error) { return node.Properties.Get(azure.TenantID.String()).String() } -// FetchDirectDescendentPaths fetches the direct AZContains relationships from the root node. Both -// the valid start/end node kind pairing and the tenant-id match are enforced +// FetchDirectDescendentPaths fetches the direct AZContains relationships from the root node func FetchDirectDescendentPaths(tx graph.Transaction, root *graph.Node, descendentKind ...graph.Kind) (graph.PathSet, error) { - if tenantID, err := nodeAzureTenantID(root); err != nil { - return nil, err - } else { - return ops.FetchPathSet(tx.Relationships().Filter(query.And( - query.Equals(query.StartID(), root.ID), - query.Kind(query.Relationship(), azure.Contains), - query.KindIn(query.End(), descendentKind...), - query.Equals(query.EndProperty(azure.TenantID.String()), tenantID), - ))) - } + return ops.FetchPathSet(tx.Relationships().Filter(query.And( + query.Equals(query.StartID(), root.ID), + query.Kind(query.Relationship(), azure.Contains), + query.KindIn(query.End(), descendentKind...), + ))) } -func FetchDirectEntityDescendents(tx graph.Transaction, root *graph.Node, skip, limit int, descendentKind graph.Kind) (graph.NodeSet, error) { +func FetchDirectEntityDescendents(tx graph.Transaction, root *graph.Node, descendentKind graph.Kind) (graph.NodeSet, error) { if paths, err := FetchDirectDescendentPaths(tx, root, descendentKind); err != nil { return nil, err } else { @@ -857,7 +857,7 @@ func FetchDirectEntityDescendents(tx graph.Transaction, root *graph.Node, skip, } } -func FetchDescendentKindByTenantID(tx graph.Transaction, root *graph.Node, skip, limit int, descendentKind ...graph.Kind) (graph.NodeSet, error) { +func FetchDescendentKindByTenantID(tx graph.Transaction, root *graph.Node, descendentKind ...graph.Kind) (graph.NodeSet, error) { if tenantID, err := nodeAzureTenantID(root); err != nil { return nil, err } else if nodes, err := ops.FetchNodeSet(tx.Nodes().Filter(query.And( @@ -883,7 +883,7 @@ func FetchEntityDescendentPaths(tx graph.Transaction, root *graph.Node, descende // Pre-populate the visited bitmap with root so that any traversal halts when it reaches it visitedBitmap.Add(root.ID.Uint64()) - terminalNodes, err := FetchDescendentKindByTenantID(tx, root, 0, 0, descendentKind...) + terminalNodes, err := FetchDescendentKindByTenantID(tx, root, descendentKind...) if err != nil { return nil, err } @@ -918,7 +918,7 @@ func FetchEntityDescendentPaths(tx graph.Transaction, root *graph.Node, descende return pathSet, nil } -func FetchEntityDescendents(tx graph.Transaction, root *graph.Node, skip, limit int, descendentKind graph.Kind) (graph.NodeSet, error) { +func FetchEntityDescendents(tx graph.Transaction, root *graph.Node, descendentKind graph.Kind) (graph.NodeSet, error) { if paths, err := FetchEntityDescendentPaths(tx, root, descendentKind); err != nil { return nil, err } else { @@ -928,7 +928,7 @@ func FetchEntityDescendents(tx graph.Transaction, root *graph.Node, skip, limit } } -func FetchDirectDescendentCounts(tx graph.Transaction, root *graph.Node, skip, limit int, descendentKinds ...graph.Kind) (Descendents, error) { +func FetchDirectDescendentCounts(tx graph.Transaction, root *graph.Node, descendentKinds ...graph.Kind) (Descendents, error) { if paths, err := FetchDirectDescendentPaths(tx, root, descendentKinds...); err != nil { return Descendents{}, err } else { @@ -946,7 +946,7 @@ func FetchDirectDescendentCounts(tx graph.Transaction, root *graph.Node, skip, l } } -func FetchEntityDescendentCounts(tx graph.Transaction, root *graph.Node, skip, limit int, descendentKinds ...graph.Kind) (Descendents, error) { +func FetchEntityDescendentCounts(tx graph.Transaction, root *graph.Node, descendentKinds ...graph.Kind) (Descendents, error) { if paths, err := FetchEntityDescendentPaths(tx, root, descendentKinds...); err != nil { return Descendents{}, err } else { @@ -980,14 +980,10 @@ func fetchRolesTraversalPlan(root *graph.Node) ops.TraversalPlan { // FetchEntityByObjectID pulls a node by its ObjectID. It requires a kind to perform index lookups against. func FetchEntityByObjectID(tx graph.Transaction, objectID string) (*graph.Node, error) { - if decodedObjectID, err := url.QueryUnescape(objectID); err != nil { - return nil, err - } else { - return tx.Nodes().Filterf(func() graph.Criteria { - return query.And( - query.Kind(query.Node(), azure.Entity), - query.Equals(query.NodeProperty(common.ObjectID.String()), decodedObjectID), - ) - }).First() - } + return tx.Nodes().Filterf(func() graph.Criteria { + return query.And( + query.Kind(query.Node(), azure.Entity), + query.Equals(query.NodeProperty(common.ObjectID.String()), objectID), + ) + }).First() } diff --git a/packages/go/analysis/azure/resource_group.go b/packages/go/analysis/azure/resource_group.go index efdb0b5b520c..3cf042faf11a 100644 --- a/packages/go/analysis/azure/resource_group.go +++ b/packages/go/analysis/azure/resource_group.go @@ -43,7 +43,7 @@ func ResourceGroupEntityDetails(ctx context.Context, db graph.Database, primaryD func PopulateResourceGroupEntityDetailsCounts(tx graph.Transaction, node *graph.Node, details ResourceGroupDetails) (ResourceGroupDetails, error) { var descendentKinds = GetDescendentKinds(azure.ResourceGroup) - if descendents, err := FetchDirectDescendentCounts(tx, node, 0, 0, descendentKinds...); err != nil { + if descendents, err := FetchDirectDescendentCounts(tx, node, descendentKinds...); err != nil { return details, err } else { details.Descendents = descendents diff --git a/packages/go/analysis/azure/subscription.go b/packages/go/analysis/azure/subscription.go index ad8106a7f040..7a1c320c494e 100644 --- a/packages/go/analysis/azure/subscription.go +++ b/packages/go/analysis/azure/subscription.go @@ -43,7 +43,7 @@ func SubscriptionEntityDetails(ctx context.Context, db graph.Database, primaryDi func PopulateSubscriptionEntityDetailsCounts(tx graph.Transaction, node *graph.Node, details SubscriptionDetails) (SubscriptionDetails, error) { var descendentKinds = GetDescendentKinds(azure.Subscription) - 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 diff --git a/packages/go/analysis/azure/tenant.go b/packages/go/analysis/azure/tenant.go index cdf94f9e43ea..4e5bb0cba5d0 100644 --- a/packages/go/analysis/azure/tenant.go +++ b/packages/go/analysis/azure/tenant.go @@ -48,7 +48,7 @@ func TenantEntityDetails(ctx context.Context, db graph.Database, primaryDisplayK func PopulateTenantEntityDetailsCounts(tx graph.Transaction, node *graph.Node, details TenantDetails) (TenantDetails, error) { descendentKinds := GetDescendentKinds(azure.Tenant) - 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 diff --git a/packages/javascript/bh-shared-ui/src/components/EntityInfoDataTableGraphed/EntityInfoDataTableGraphed.test.tsx b/packages/javascript/bh-shared-ui/src/components/EntityInfoDataTableGraphed/EntityInfoDataTableGraphed.test.tsx index ad60093e0dcd..cfe0f5a0576c 100644 --- a/packages/javascript/bh-shared-ui/src/components/EntityInfoDataTableGraphed/EntityInfoDataTableGraphed.test.tsx +++ b/packages/javascript/bh-shared-ui/src/components/EntityInfoDataTableGraphed/EntityInfoDataTableGraphed.test.tsx @@ -131,7 +131,7 @@ describe('EntityInfoDataTableGraphed', () => { expect(sum).not.toBeNull(); }); - it('displays ! icon when one of the Affected Object calls fail', async () => { + it('sums the remaining sections when one of the Affected Object calls fails', async () => { console.error = vi.fn(); server.use( rest.get(`api/v2/gpos/${objectId}/ous`, (req, res, ctx) => { @@ -141,9 +141,10 @@ describe('EntityInfoDataTableGraphed', () => { render(); - const errorIcon = await screen.findByTestId('ErrorOutlineIcon'); + const sum = await screen.findByText('5,056'); + expect(sum).not.toBeNull(); - expect(errorIcon).not.toBeNull(); + expect(screen.queryByTestId('ErrorOutlineIcon')).toBeNull(); }); it('displays 0 when a given sections returns empty, and sums the rest of the sections correctly', async () => { diff --git a/packages/javascript/bh-shared-ui/src/components/EntityInfoDataTableGraphed/EntityInfoDataTableGraphed.tsx b/packages/javascript/bh-shared-ui/src/components/EntityInfoDataTableGraphed/EntityInfoDataTableGraphed.tsx index 1a257bc52748..49495852ed14 100644 --- a/packages/javascript/bh-shared-ui/src/components/EntityInfoDataTableGraphed/EntityInfoDataTableGraphed.tsx +++ b/packages/javascript/bh-shared-ui/src/components/EntityInfoDataTableGraphed/EntityInfoDataTableGraphed.tsx @@ -13,6 +13,7 @@ // limitations under the License. // // SPDX-License-Identifier: Apache-2.0 +import { PaginatedResponse } from 'js-client-library'; import { useQuery } from 'react-query'; import { NODE_GRAPH_RENDER_LIMIT } from '../../constants'; import { useExploreParams } from '../../hooks'; @@ -21,6 +22,27 @@ import { EntityInfoDataTableProps, entityRelationshipEndpoints } from '../../uti import EntityInfoCollapsibleSection from '../EntityInfo/EntityInfoCollapsibleSection'; import InfiniteScrollingTable from '../InfiniteScrollingTable'; +function getCount( + queryData: Array>> | PaginatedResponse | undefined, + countLabel: string | undefined +): number | undefined { + if (Array.isArray(queryData)) { + const fulfilledData = queryData.filter((result) => result.status === 'fulfilled').map((result) => result.value); + + if (countLabel !== undefined) { + const labeledSection = fulfilledData.find((sectionData: any) => sectionData?.countLabel === countLabel); + return labeledSection?.count; + } else { + return fulfilledData.reduce((acc, val) => { + const sectionCount = val?.count ?? 0; + return acc + sectionCount; + }, 0); + } + } else if (queryData) { + return queryData?.count ?? 0; + } +} + export const EntityInfoDataTableGraphed: React.FC = ({ id, label, @@ -35,13 +57,13 @@ export const EntityInfoDataTableGraphed: React.FC = ({ const isExpandedPanelSection = (expandedPanelSections as string[]).includes(label); const countQuery = useQuery( - ['relatedCount', label, id], + ['relatedCount', label, id, sections], () => { if (endpoint) { return endpoint({ id, skip: 0, limit: 128 }); } if (sections) - return Promise.all( + return Promise.allSettled( sections.map((section: EntityInfoDataTableProps) => { const endpoint = section.queryType ? entityRelationshipEndpoints[section.queryType] : undefined; return endpoint ? endpoint({ id, skip: 0, limit: 128 }) : Promise.resolve(); @@ -104,21 +126,7 @@ export const EntityInfoDataTableGraphed: React.FC = ({ setNodeSearchParams(item); }; - let count: number | undefined; - if (Array.isArray(countQuery.data)) { - if (countLabel !== undefined) { - countQuery.data.forEach((sectionData: any) => { - if (sectionData.countLabel === countLabel) count = sectionData.count; - }); - } else { - count = countQuery.data.reduce((acc, val) => { - const count = val?.count ?? 0; - return acc + count; - }, 0); - } - } else if (countQuery.data) { - count = countQuery.data?.count ?? 0; - } + const count = getCount(countQuery.data, countLabel); return ( Date: Tue, 16 Jun 2026 11:07:34 -0500 Subject: [PATCH 4/6] fix: rabbit feedback for disconnected paths --- .../analysis/azure/azure_integration_test.go | 76 +++++++++++++++++++ packages/go/analysis/azure/queries.go | 36 +++++---- 2 files changed, 93 insertions(+), 19 deletions(-) diff --git a/packages/go/analysis/azure/azure_integration_test.go b/packages/go/analysis/azure/azure_integration_test.go index fc6ae3bf3eae..4143e463382b 100644 --- a/packages/go/analysis/azure/azure_integration_test.go +++ b/packages/go/analysis/azure/azure_integration_test.go @@ -1044,3 +1044,79 @@ func TestFetchEntityDescendentPaths_NoTerminalsReturnsEmpty(t *testing.T) { }) 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) +} diff --git a/packages/go/analysis/azure/queries.go b/packages/go/analysis/azure/queries.go index 8611e07b5c17..2dc6732d30bd 100644 --- a/packages/go/analysis/azure/queries.go +++ b/packages/go/analysis/azure/queries.go @@ -28,7 +28,6 @@ import ( "github.com/specterops/bloodhound/packages/go/graphschema/ad" "github.com/specterops/bloodhound/packages/go/graphschema/azure" "github.com/specterops/bloodhound/packages/go/graphschema/common" - "github.com/specterops/dawgs/cardinality" "github.com/specterops/dawgs/graph" "github.com/specterops/dawgs/ops" "github.com/specterops/dawgs/query" @@ -821,6 +820,8 @@ var containsValidEndKinds = map[graph.Kind]graph.Kinds{ }, } +// isDirectDescendent reports whether targetKind is a valid direct AZContains child of sourceKind, +// according to the containsValidEndKinds mapping derived from the Azure ingestion logic. func isDirectDescendent(sourceKind, targetKind graph.Kind) bool { if validEndKinds, ok := containsValidEndKinds[sourceKind]; ok && validEndKinds.ContainsOneOf(targetKind) { return true @@ -847,6 +848,8 @@ func FetchDirectDescendentPaths(tx graph.Transaction, root *graph.Node, descende ))) } +// FetchDirectEntityDescendents fetches the set of nodes of descendentKind that are direct AZContains +// children of root, excluding root itself. func FetchDirectEntityDescendents(tx graph.Transaction, root *graph.Node, descendentKind graph.Kind) (graph.NodeSet, error) { if paths, err := FetchDirectDescendentPaths(tx, root, descendentKind); err != nil { return nil, err @@ -857,6 +860,9 @@ func FetchDirectEntityDescendents(tx graph.Transaction, root *graph.Node, descen } } +// FetchDescendentKindByTenantID fetches the set of nodes matching descendentKind that reside within +// the same Azure tenant as root, determined by comparing the azure.TenantID property against root's +// tenant ID. func FetchDescendentKindByTenantID(tx graph.Transaction, root *graph.Node, descendentKind ...graph.Kind) (graph.NodeSet, error) { if tenantID, err := nodeAzureTenantID(root); err != nil { return nil, err @@ -872,16 +878,9 @@ func FetchDescendentKindByTenantID(tx graph.Transaction, root *graph.Node, desce // FetchEntityDescendentPaths fetches paths from each terminal node of descendentKind (within the // same Azure tenant as root) back up to root via azure.Contains relationships. Each traversal -// halts upon reaching root or encountering a node already claimed by a prior traversal, preventing -// duplicate path segments across traversals. +// halts upon reaching root, and only paths that successfully reach root are included in the result. func FetchEntityDescendentPaths(tx graph.Transaction, root *graph.Node, descendentKind ...graph.Kind) (graph.PathSet, error) { - var ( - visitedBitmap = cardinality.NewBitmap64() - pathSet = graph.NewPathSet() - ) - - // Pre-populate the visited bitmap with root so that any traversal halts when it reaches it - visitedBitmap.Add(root.ID.Uint64()) + pathSet := graph.NewPathSet() terminalNodes, err := FetchDescendentKindByTenantID(tx, root, descendentKind...) if err != nil { @@ -889,11 +888,7 @@ func FetchEntityDescendentPaths(tx graph.Transaction, root *graph.Node, descende } for _, terminalNode := range terminalNodes { - // Skip this terminal if it was already reached and claimed by a prior traversal - if visitedBitmap.Contains(terminalNode.ID.Uint64()) { - continue - } - + reachedRoot := false if paths, err := ops.TraversePaths(tx, ops.TraversalPlan{ Root: terminalNode, Direction: graph.DirectionInbound, @@ -901,17 +896,18 @@ func FetchEntityDescendentPaths(tx graph.Transaction, root *graph.Node, descende return query.Kind(query.Relationship(), azure.Contains) }, ExpansionFilter: func(segment *graph.PathSegment) bool { - nodeID := segment.Node.ID.Uint64() - if visitedBitmap.Contains(nodeID) { + if segment.Node.ID == root.ID { + reachedRoot = true return false } - visitedBitmap.Add(nodeID) return true }, }); err != nil { return nil, err } else { - pathSet.AddPathSet(paths) + if reachedRoot { + pathSet.AddPathSet(paths) + } } } @@ -928,6 +924,8 @@ func FetchEntityDescendents(tx graph.Transaction, root *graph.Node, descendentKi } } +// FetchDirectDescendentCounts returns the per-kind counts of nodes that are direct AZContains +// children of root, keyed by kind string and excluding root itself. func FetchDirectDescendentCounts(tx graph.Transaction, root *graph.Node, descendentKinds ...graph.Kind) (Descendents, error) { if paths, err := FetchDirectDescendentPaths(tx, root, descendentKinds...); err != nil { return Descendents{}, err From 8557a8ebaca7747df66bc27fb7131808540f7da4 Mon Sep 17 00:00:00 2001 From: Elijah Brown Date: Tue, 23 Jun 2026 09:51:30 -0700 Subject: [PATCH 5/6] Add test for nested management group --- .../analysis/azure/azure_integration_test.go | 48 +++++++++++++++++++ 1 file changed, 48 insertions(+) diff --git a/packages/go/analysis/azure/azure_integration_test.go b/packages/go/analysis/azure/azure_integration_test.go index 4143e463382b..efe3c323cbbe 100644 --- a/packages/go/analysis/azure/azure_integration_test.go +++ b/packages/go/analysis/azure/azure_integration_test.go @@ -1120,3 +1120,51 @@ func TestFetchEntityDescendentPaths_TerminalNotConnectedToRoot(t *testing.T) { }) 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) +} From a02ff873749fc39bb6d84ee88203b2bba8ba477f Mon Sep 17 00:00:00 2001 From: Ulises Rangel Date: Tue, 23 Jun 2026 13:12:36 -0500 Subject: [PATCH 6/6] fix: nested management group descendents --- packages/go/analysis/azure/queries.go | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/packages/go/analysis/azure/queries.go b/packages/go/analysis/azure/queries.go index 2dc6732d30bd..cb5d61fd0b13 100644 --- a/packages/go/analysis/azure/queries.go +++ b/packages/go/analysis/azure/queries.go @@ -788,9 +788,10 @@ func FetchKeyVaultReaderCounts(tx graph.Transaction, keyVault *graph.Node) (KeyV return keyVaultReaders, nil } -// containsValidEndKinds maps each valid AZContains start node kind to the set of valid end node kinds. +// directRelationshipKinds maps each valid AZContains start node kind to the set of valid end node kinds +// where it is a single hop relationship as opposed to a path of multiple relationships. // This is derived directly from the ingestion logic in bhce/packages/go/ein/azure.go. -var containsValidEndKinds = map[graph.Kind]graph.Kinds{ +var directRelationshipKinds = map[graph.Kind]graph.Kinds{ azure.Tenant: { azure.App, azure.Device, @@ -800,10 +801,6 @@ var containsValidEndKinds = map[graph.Kind]graph.Kinds{ azure.Subscription, azure.User, }, - azure.ManagementGroup: { - azure.ManagementGroup, - azure.Subscription, - }, azure.Subscription: { azure.ResourceGroup, }, @@ -823,7 +820,7 @@ var containsValidEndKinds = map[graph.Kind]graph.Kinds{ // isDirectDescendent reports whether targetKind is a valid direct AZContains child of sourceKind, // according to the containsValidEndKinds mapping derived from the Azure ingestion logic. func isDirectDescendent(sourceKind, targetKind graph.Kind) bool { - if validEndKinds, ok := containsValidEndKinds[sourceKind]; ok && validEndKinds.ContainsOneOf(targetKind) { + if validEndKinds, ok := directRelationshipKinds[sourceKind]; ok && validEndKinds.ContainsOneOf(targetKind) { return true } return false