feat: add flag to dump pprof profiles to stderr

cmd/utils.go

@@ -34,8 +34,12 @@ import (
 	"os"
 	"path"
 	"path/filepath"
+	"runtime/pprof"
 	"time"
 
+	"github.com/spf13/cobra"
+	"golang.org/x/net/proxy"
+
 	"github.com/bloodhoundad/azurehound/v2/client"
 	client_config "github.com/bloodhoundad/azurehound/v2/client/config"
 	"github.com/bloodhoundad/azurehound/v2/client/rest"
@@ -45,8 +49,6 @@ import (
 	"github.com/bloodhoundad/azurehound/v2/models"
 	"github.com/bloodhoundad/azurehound/v2/pipeline"
 	"github.com/bloodhoundad/azurehound/v2/sinks"
-	"github.com/spf13/cobra"
-	"golang.org/x/net/proxy"
 )
 
 func exit(err error) {
@@ -86,7 +88,9 @@ func persistentPreRunE(cmd *cobra.Command, args []string) error {
 func gracefulShutdown(stop context.CancelFunc) {
 	stop()
 	fmt.Fprintln(os.Stderr, "\nshutting down gracefully, press ctrl+c again to force")
-	// TODO timeout context
+	if profile := pprof.Lookup(config.Pprof.Value().(string)); profile != nil {
+		profile.WriteTo(os.Stderr, 1)
+	}
 }
 
 func testConnections() error {
@@ -224,7 +228,6 @@ func newAzureClient() (client.AzureClient, error) {
 		} else {
 			clientCert = string(content)
 		}
-
 	}
 
 	if file, ok := keyFile.(string); ok && file != "" {

config/config.go

@@ -128,6 +128,12 @@ var (
 		Persistent: true,
 		Default:    "",
 	}
+	Pprof = Config{
+		Name:       "pprof",
+		Usage:      "During graceful shutdown, prints the pprof profile with the provided name to stderr",
+		Persistent: true,
+		Default:    "",
+	}
 
 	// Azure Configurations
 	AzAppId = Config{
@@ -283,6 +289,7 @@ var (
 		LogFile,
 		Proxy,
 		RefreshToken,
+		Pprof,
 	}
 
 	AzureConfig = []Config{