Bed 4262 az panic recovery

cmd/list-app-owners.go

@@ -44,19 +44,21 @@ var listAppOwnersCmd = &cobra.Command{
 
 func listAppOwnersCmdImpl(cmd *cobra.Command, args []string) {
 	ctx, stop := signal.NotifyContext(cmd.Context(), os.Interrupt, os.Kill)
+	panicChan := panicChan()
 	defer gracefulShutdown(stop)
 
 	log.V(1).Info("testing connections")
 	azClient := connectAndCreateClient()
 	log.Info("collecting azure app owners...")
 	start := time.Now()
-	stream := listAppOwners(ctx, azClient, listApps(ctx, azClient))
+	stream := listAppOwners(ctx, azClient, panicChan, listApps(ctx, azClient, panicChan))
+	handleBubbledPanic(ctx, panicChan, stop)
 	outputStream(ctx, stream)
 	duration := time.Since(start)
 	log.Info("collection completed", "duration", duration.String())
 }
 
-func listAppOwners(ctx context.Context, client client.AzureClient, apps <-chan azureWrapper[models.App]) <-chan azureWrapper[models.AppOwners] {
+func listAppOwners(ctx context.Context, client client.AzureClient, panicChan chan error, apps <-chan azureWrapper[models.App]) <-chan azureWrapper[models.AppOwners] {
 	var (
 		out     = make(chan azureWrapper[models.AppOwners])
 		streams = pipeline.Demux(ctx.Done(), apps, 25)
@@ -67,6 +69,7 @@ func listAppOwners(ctx context.Context, client client.AzureClient, apps <-chan a
 	for i := range streams {
 		stream := streams[i]
 		go func() {
+			defer panicRecovery(panicChan)
 			defer wg.Done()
 			for app := range stream {
 				var (

cmd/list-app-owners_test.go

@@ -50,7 +50,7 @@ func TestListAppOwners(t *testing.T) {
 	mockClient.EXPECT().TenantInfo().Return(mockTenant).AnyTimes()
 	mockClient.EXPECT().ListAzureADAppOwners(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(mockAppOwnerChannel).Times(1)
 	mockClient.EXPECT().ListAzureADAppOwners(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(mockAppOwnerChannel2).Times(1)
-	channel := listAppOwners(ctx, mockClient, mockAppsChannel)
+	channel := listAppOwners(ctx, mockClient, panicChan(), mockAppsChannel)
 
 	go func() {
 		defer close(mockAppsChannel)

cmd/list-app-role-assignments.go

@@ -51,14 +51,16 @@ func listAppRoleAssignmentsCmdImpl(cmd *cobra.Command, args []string) {
 	azClient := connectAndCreateClient()
 	log.Info("collecting azure active directory app role assignments...")
 	start := time.Now()
-	servicePrincipals := listServicePrincipals(ctx, azClient)
-	stream := listAppRoleAssignments(ctx, azClient, servicePrincipals)
+	panicChan := panicChan()
+	servicePrincipals := listServicePrincipals(ctx, azClient, panicChan)
+	stream := listAppRoleAssignments(ctx, azClient, panicChan, servicePrincipals)
+	handleBubbledPanic(ctx, panicChan, stop)
 	outputStream(ctx, stream)
 	duration := time.Since(start)
 	log.Info("collection completed", "duration", duration.String())
 }
 
-func listAppRoleAssignments(ctx context.Context, client client.AzureClient, servicePrincipals <-chan interface{}) <-chan interface{} {
+func listAppRoleAssignments(ctx context.Context, client client.AzureClient, panicChan chan error, servicePrincipals <-chan interface{}) <-chan interface{} {
 	var (
 		out         = make(chan interface{})
 		filteredSPs = make(chan models.ServicePrincipal)
@@ -67,6 +69,7 @@ func listAppRoleAssignments(ctx context.Context, client client.AzureClient, serv
 	)
 
 	go func() {
+		defer panicRecovery(panicChan)
 		defer close(filteredSPs)
 
 		for result := range pipeline.OrDone(ctx.Done(), servicePrincipals) {
@@ -87,6 +90,7 @@ func listAppRoleAssignments(ctx context.Context, client client.AzureClient, serv
 	for i := range streams {
 		stream := streams[i]
 		go func() {
+			defer panicRecovery(panicChan)
 			defer wg.Done()
 			for servicePrincipal := range stream {
 				var (

cmd/list-apps.go

@@ -49,16 +49,19 @@ func listAppsCmdImpl(cmd *cobra.Command, args []string) {
 	azClient := connectAndCreateClient()
 	log.Info("collecting azure active directory applications...")
 	start := time.Now()
-	stream := listApps(ctx, azClient)
+	panicChan := panicChan()
+	stream := listApps(ctx, azClient, panicChan)
+	handleBubbledPanic(ctx, panicChan, stop)
 	outputStream(ctx, stream)
 	duration := time.Since(start)
 	log.Info("collection completed", "duration", duration.String())
 }
 
-func listApps(ctx context.Context, client client.AzureClient) <-chan azureWrapper[models.App] {
+func listApps(ctx context.Context, client client.AzureClient, panicChan chan error) <-chan azureWrapper[models.App] {
 	out := make(chan azureWrapper[models.App])
 
 	go func() {
+		defer panicRecovery(panicChan)
 		defer close(out)
 		count := 0
 		for item := range client.ListAzureADApps(ctx, "", "", "", "", nil) {

cmd/list-apps_test.go

@@ -56,7 +56,7 @@ func TestListApps(t *testing.T) {
 		}
 	}()
 
-	channel := listApps(ctx, mockClient)
+	channel := listApps(ctx, mockClient, panicChan())
 	<-channel
 	if _, ok := <-channel; ok {
 		t.Error("expected channel to close from an error result but it did not")

cmd/list-automation-account-role-assignments.go

@@ -52,14 +52,16 @@ func listAutomationAccountRoleAssignmentImpl(cmd *cobra.Command, args []string)
 	azClient := connectAndCreateClient()
 	log.Info("collecting azure automation account role assignments...")
 	start := time.Now()
-	subscriptions := listSubscriptions(ctx, azClient)
-	stream := listAutomationAccountRoleAssignments(ctx, azClient, listAutomationAccounts(ctx, azClient, subscriptions))
+	panicChan := panicChan()
+	subscriptions := listSubscriptions(ctx, azClient, panicChan)
+	stream := listAutomationAccountRoleAssignments(ctx, azClient, panicChan, listAutomationAccounts(ctx, azClient, panicChan, subscriptions))
+	handleBubbledPanic(ctx, panicChan, stop)
 	outputStream(ctx, stream)
 	duration := time.Since(start)
 	log.Info("collection completed", "duration", duration.String())
 }
 
-func listAutomationAccountRoleAssignments(ctx context.Context, client client.AzureClient, automationAccounts <-chan interface{}) <-chan interface{} {
+func listAutomationAccountRoleAssignments(ctx context.Context, client client.AzureClient, panicChan chan error, automationAccounts <-chan interface{}) <-chan interface{} {
 	var (
 		out     = make(chan interface{})
 		ids     = make(chan string)
@@ -68,6 +70,7 @@ func listAutomationAccountRoleAssignments(ctx context.Context, client client.Azu
 	)
 
 	go func() {
+		defer panicRecovery(panicChan)
 		defer close(ids)
 
 		for result := range pipeline.OrDone(ctx.Done(), automationAccounts) {
@@ -86,6 +89,7 @@ func listAutomationAccountRoleAssignments(ctx context.Context, client client.Azu
 	for i := range streams {
 		stream := streams[i]
 		go func() {
+			defer panicRecovery(panicChan)
 			defer wg.Done()
 			for id := range stream {
 				var (

cmd/list-automation-accounts.go

@@ -51,13 +51,15 @@ func listAutomationAccountsCmdImpl(cmd *cobra.Command, args []string) {
 	azClient := connectAndCreateClient()
 	log.Info("collecting azure automation accounts...")
 	start := time.Now()
-	stream := listAutomationAccounts(ctx, azClient, listSubscriptions(ctx, azClient))
+	panicChan := panicChan()
+	stream := listAutomationAccounts(ctx, azClient, panicChan, listSubscriptions(ctx, azClient, panicChan))
+	handleBubbledPanic(ctx, panicChan, stop)
 	outputStream(ctx, stream)
 	duration := time.Since(start)
 	log.Info("collection completed", "duration", duration.String())
 }
 
-func listAutomationAccounts(ctx context.Context, client client.AzureClient, subscriptions <-chan interface{}) <-chan interface{} {
+func listAutomationAccounts(ctx context.Context, client client.AzureClient, panicChan chan error, subscriptions <-chan interface{}) <-chan interface{} {
 	var (
 		out     = make(chan interface{})
 		ids     = make(chan string)
@@ -66,6 +68,7 @@ func listAutomationAccounts(ctx context.Context, client client.AzureClient, subs
 	)
 
 	go func() {
+		defer panicRecovery(panicChan)
 		defer close(ids)
 		for result := range pipeline.OrDone(ctx.Done(), subscriptions) {
 			if subscription, ok := result.(AzureWrapper).Data.(models.Subscription); !ok {
@@ -83,6 +86,7 @@ func listAutomationAccounts(ctx context.Context, client client.AzureClient, subs
 	for i := range streams {
 		stream := streams[i]
 		go func() {
+			defer panicRecovery(panicChan)
 			defer wg.Done()
 			for id := range stream {
 				count := 0

cmd/list-azure-ad.go

@@ -53,13 +53,15 @@ func listAzureADCmdImpl(cmd *cobra.Command, args []string) {
 	azClient := connectAndCreateClient()
 	log.Info("collecting azure ad objects...")
 	start := time.Now()
-	stream := listAllAD(ctx, azClient)
+	panicChan := panicChan()
+	stream := listAllAD(ctx, azClient, panicChan)
+	handleBubbledPanic(ctx, panicChan, stop)
 	outputStream(ctx, stream)
 	duration := time.Since(start)
 	log.Info("collection completed", "duration", duration.String())
 }
 
-func listAllAD(ctx context.Context, client client.AzureClient) <-chan interface{} {
+func listAllAD(ctx context.Context, client client.AzureClient, panicChan chan error) <-chan interface{} {
 	var (
 		devices  = make(chan interface{})
 		devices2 = make(chan interface{})
@@ -79,35 +81,35 @@ func listAllAD(ctx context.Context, client client.AzureClient) <-chan interface{
 	)
 
 	// Enumerate Apps, AppOwners and AppMembers
-	appChans := pipeline.TeeFixed(ctx.Done(), listApps(ctx, client), 2)
+	appChans := pipeline.TeeFixed(ctx.Done(), listApps(ctx, client, panicChan), 2)
 	apps := pipeline.ToAny(ctx.Done(), appChans[0])
-	appOwners := pipeline.ToAny(ctx.Done(), listAppOwners(ctx, client, appChans[1]))
+	appOwners := pipeline.ToAny(ctx.Done(), listAppOwners(ctx, client, panicChan, appChans[1]))
 
 	// Enumerate Devices and DeviceOwners
-	pipeline.Tee(ctx.Done(), listDevices(ctx, client), devices, devices2)
-	deviceOwners := listDeviceOwners(ctx, client, devices2)
+	pipeline.Tee(ctx.Done(), listDevices(ctx, client, panicChan), devices, devices2)
+	deviceOwners := listDeviceOwners(ctx, client, panicChan, devices2)
 
 	// Enumerate Groups, GroupOwners and GroupMembers
-	pipeline.Tee(ctx.Done(), listGroups(ctx, client), groups, groups2, groups3)
-	groupOwners := listGroupOwners(ctx, client, groups2)
-	groupMembers := listGroupMembers(ctx, client, groups3)
+	pipeline.Tee(ctx.Done(), listGroups(ctx, client, panicChan), groups, groups2, groups3)
+	groupOwners := listGroupOwners(ctx, client, panicChan, groups2)
+	groupMembers := listGroupMembers(ctx, client, panicChan, groups3)
 
 	// Enumerate ServicePrincipals and ServicePrincipalOwners
-	pipeline.Tee(ctx.Done(), listServicePrincipals(ctx, client), servicePrincipals, servicePrincipals2, servicePrincipals3)
-	servicePrincipalOwners := listServicePrincipalOwners(ctx, client, servicePrincipals2)
+	pipeline.Tee(ctx.Done(), listServicePrincipals(ctx, client, panicChan), servicePrincipals, servicePrincipals2, servicePrincipals3)
+	servicePrincipalOwners := listServicePrincipalOwners(ctx, client, panicChan, servicePrincipals2)
 
 	// Enumerate Tenants
-	pipeline.Tee(ctx.Done(), listTenants(ctx, client), tenants)
+	pipeline.Tee(ctx.Done(), listTenants(ctx, client, panicChan), tenants)
 
 	// Enumerate Users
-	users := listUsers(ctx, client)
+	users := listUsers(ctx, client, panicChan)
 
 	// Enumerate Roles and RoleAssignments
-	pipeline.Tee(ctx.Done(), listRoles(ctx, client), roles, roles2)
-	roleAssignments := listRoleAssignments(ctx, client, roles2)
+	pipeline.Tee(ctx.Done(), listRoles(ctx, client, panicChan), roles, roles2)
+	roleAssignments := listRoleAssignments(ctx, client, panicChan, roles2)
 
 	// Enumerate AppRoleAssignments
-	appRoleAssignments := listAppRoleAssignments(ctx, client, servicePrincipals3)
+	appRoleAssignments := listAppRoleAssignments(ctx, client, panicChan, servicePrincipals3)
 
 	return pipeline.Mux(ctx.Done(),
 		appOwners,