From 15cde2995b32d0d6d83eb13cd5dddc31e00afa31 Mon Sep 17 00:00:00 2001
From: hc-github-team-secure-vault-core
 <82990506+hc-github-team-secure-vault-core@users.noreply.github.com>
Date: Wed, 17 May 2023 12:27:20 -0400
Subject: [PATCH] command/server: add support to write pprof files to the
 filesystem via SIGUSR2 (#20609) (#20620)

* core/server: add support to write pprof files to the filesystem via SIGUSR2

* changelog

* Fix filepath join

* Use core logger

* Simplify logic

* Break on error

Co-authored-by: Jason O'Donnell <2160810+jasonodonnell@users.noreply.github.com>
---
 changelog/20609.txt |  4 +++
 command/server.go   | 73 +++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 77 insertions(+)
 create mode 100644 changelog/20609.txt

diff --git a/changelog/20609.txt b/changelog/20609.txt
new file mode 100644
index 000000000000..fe92833da52d
--- /dev/null
+++ b/changelog/20609.txt
@@ -0,0 +1,4 @@
+```release-note:improvement
+command/server: Add support for dumping pprof files to the filesystem via SIGUSR2 when 
+`VAULT_PPROF_WRITE_TO_FILE=true` is set on the server.
+```
\ No newline at end of file
diff --git a/command/server.go b/command/server.go
index d25b23aed0b4..094178876d92 100644
--- a/command/server.go
+++ b/command/server.go
@@ -1659,6 +1659,79 @@ func (c *ServerCommand) Run(args []string) int {
 		case <-c.SigUSR2Ch:
 			logWriter := c.logger.StandardWriter(&hclog.StandardLoggerOptions{})
 			pprof.Lookup("goroutine").WriteTo(logWriter, 2)
+
+			if os.Getenv("VAULT_STACKTRACE_WRITE_TO_FILE") != "" {
+				c.logger.Info("Writing stacktrace to file")
+
+				dir := ""
+				path := os.Getenv("VAULT_STACKTRACE_FILE_PATH")
+				if path != "" {
+					if _, err := os.Stat(path); err != nil {
+						c.logger.Error("Checking stacktrace path failed", "error", err)
+						continue
+					}
+					dir = path
+				} else {
+					dir, err = os.MkdirTemp("", "vault-stacktrace")
+					if err != nil {
+						c.logger.Error("Could not create temporary directory for stacktrace", "error", err)
+						continue
+					}
+				}
+
+				f, err := os.CreateTemp(dir, "stacktrace")
+				if err != nil {
+					c.logger.Error("Could not create stacktrace file", "error", err)
+					continue
+				}
+
+				if err := pprof.Lookup("goroutine").WriteTo(f, 2); err != nil {
+					f.Close()
+					c.logger.Error("Could not write stacktrace to file", "error", err)
+					continue
+				}
+
+				c.logger.Info(fmt.Sprintf("Wrote stacktrace to: %s", f.Name()))
+				f.Close()
+			}
+
+			// We can only get pprof outputs via the API but sometimes Vault can get
+			// into a state where it cannot process requests so we can get pprof outputs
+			// via SIGUSR2.
+			if os.Getenv("VAULT_PPROF_WRITE_TO_FILE") != "" {
+				dir := ""
+				path := os.Getenv("VAULT_PPROF_FILE_PATH")
+				if path != "" {
+					if _, err := os.Stat(path); err != nil {
+						c.logger.Error("Checking pprof path failed", "error", err)
+						continue
+					}
+					dir = path
+				} else {
+					dir, err = os.MkdirTemp("", "vault-pprof")
+					if err != nil {
+						c.logger.Error("Could not create temporary directory for pprof", "error", err)
+						continue
+					}
+				}
+
+				dumps := []string{"goroutine", "heap", "allocs", "threadcreate"}
+				for _, dump := range dumps {
+					pFile, err := os.Create(filepath.Join(dir, dump))
+					if err != nil {
+						c.logger.Error("error creating pprof file", "name", dump, "error", err)
+						break
+					}
+
+					err = pprof.Lookup(dump).WriteTo(pFile, 0)
+					if err != nil {
+						c.logger.Error("error generating pprof data", "name", dump, "error", err)
+						break
+					}
+				}
+
+				c.logger.Info(fmt.Sprintf("Wrote pprof files to: %s", dir))
+			}
 		}
 	}
 	// Notify systemd that the server is shutting down