From e4265bbdf8598a4887a5414b9f62a1719fb17a99 Mon Sep 17 00:00:00 2001
From: Nedyalko Andreev <n@andreev.sh>
Date: Fri, 31 Mar 2023 11:54:26 +0300
Subject: [PATCH] Fix panic handling and add a new exit code and a test for
 panics

Because we started calling os.Exit() in a defer statement, we circumvented the regular handling of panics that bubble up to the top by the Go runtime. I could have removed the os.Exit() call from the defer, but handling panics ourselves in this way allows us to handle them exactly how we want, assign them their own unique exit code, and test the whole thing with an e2e test.
---
 cmd/root.go               | 13 +++++++++++++
 cmd/root_test.go          | 24 ++++++++++++++++++++++++
 errext/exitcodes/codes.go |  1 +
 3 files changed, 38 insertions(+)

diff --git a/cmd/root.go b/cmd/root.go
index 9628f0a764d..9de1d625c98 100644
--- a/cmd/root.go
+++ b/cmd/root.go
@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"io"
 	stdlog "log"
+	"runtime/debug"
 	"strconv"
 	"strings"
 	"sync"
@@ -17,6 +18,7 @@ import (
 
 	"go.k6.io/k6/cmd/state"
 	"go.k6.io/k6/errext"
+	"go.k6.io/k6/errext/exitcodes"
 	"go.k6.io/k6/lib/consts"
 	"go.k6.io/k6/log"
 )
@@ -88,6 +90,17 @@ func (c *rootCommand) execute() {
 		c.globalState.OSExit(exitCode)
 	}()
 
+	defer func() {
+		if r := recover(); r != nil {
+			exitCode = int(exitcodes.GoPanic)
+			err := fmt.Errorf("unexpected k6 panic: %s\n%s", r, debug.Stack())
+			if c.loggerIsRemote {
+				c.globalState.FallbackLogger.Error(err)
+			}
+			c.globalState.Logger.Error(err)
+		}
+	}()
+
 	err := c.cmd.Execute()
 	if err == nil {
 		exitCode = 0
diff --git a/cmd/root_test.go b/cmd/root_test.go
index a1b3d3d0283..d1ccf9f54a7 100644
--- a/cmd/root_test.go
+++ b/cmd/root_test.go
@@ -5,8 +5,10 @@ import (
 	"testing"
 
 	"github.com/sirupsen/logrus"
+	"github.com/spf13/cobra"
 	"github.com/stretchr/testify/assert"
 	"go.k6.io/k6/cmd/tests"
+	"go.k6.io/k6/errext/exitcodes"
 	"go.k6.io/k6/lib/testutils"
 )
 
@@ -37,3 +39,25 @@ func TestDeprecatedOptionWarning(t *testing.T) {
 	assert.False(t, testutils.LogContains(logMsgs, logrus.InfoLevel, "logformat"))
 	assert.Contains(t, ts.Stdout.String(), `--logformat has been deprecated`)
 }
+
+func TestPanicHandling(t *testing.T) {
+	t.Parallel()
+
+	ts := tests.NewGlobalTestState(t)
+	ts.CmdArgs = []string{"k6", "panic"}
+	ts.ExpectedExitCode = int(exitcodes.GoPanic)
+
+	rootCmd := newRootCommand(ts.GlobalState)
+	rootCmd.cmd.AddCommand(&cobra.Command{
+		Use: "panic",
+		RunE: func(cmd *cobra.Command, args []string) error {
+			panic("oh no, oh no, oh no,no,no,no,no")
+		},
+	})
+	rootCmd.execute()
+
+	t.Log(ts.Stderr.String())
+	logMsgs := ts.LoggerHook.Drain()
+	assert.True(t, testutils.LogContains(logMsgs, logrus.ErrorLevel, "unexpected k6 panic: oh no"))
+	assert.True(t, testutils.LogContains(logMsgs, logrus.ErrorLevel, "cmd.TestPanicHandling")) // check stacktrace
+}
diff --git a/errext/exitcodes/codes.go b/errext/exitcodes/codes.go
index 8efc0578c8c..69050b571b6 100644
--- a/errext/exitcodes/codes.go
+++ b/errext/exitcodes/codes.go
@@ -20,4 +20,5 @@ const (
 	CannotStartRESTAPI       ExitCode = 106
 	ScriptException          ExitCode = 107
 	ScriptAborted            ExitCode = 108
+	GoPanic                  ExitCode = 109
 )