From 729b78ae435f1745c27ea7d296f8c04f4769565c Mon Sep 17 00:00:00 2001
From: David Pokora <dpokora@gmail.com>
Date: Mon, 26 Feb 2024 12:11:50 -0800
Subject: [PATCH] Added a unique exit code for failed tests, added generic way
 to bubble up errors with exit codes from cmd, organized exit code definitions
 into one package (#301)

Co-authored-by: anishnaik <anish.naik@trailofbits.com>
---
 cmd/exitcodes/error_with_exit_code.go | 37 +++++++++++++++++++++++++++
 cmd/exitcodes/exit_codes.go           | 21 +++++++++++++++
 cmd/fuzz.go                           |  6 +++++
 main.go                               |  6 ++---
 4 files changed, 67 insertions(+), 3 deletions(-)
 create mode 100644 cmd/exitcodes/error_with_exit_code.go
 create mode 100644 cmd/exitcodes/exit_codes.go

diff --git a/cmd/exitcodes/error_with_exit_code.go b/cmd/exitcodes/error_with_exit_code.go
new file mode 100644
index 00000000..eb2367a7
--- /dev/null
+++ b/cmd/exitcodes/error_with_exit_code.go
@@ -0,0 +1,37 @@
+package exitcodes
+
+// ErrorWithExitCode is an `error` type that wraps an existing error and exit code, providing exit codes
+// for a given error if they are bubbled up to the top-level.
+type ErrorWithExitCode struct {
+	err      error
+	exitCode int
+}
+
+// NewErrorWithExitCode creates a new error (ErrorWithExitCode) with the provided internal error and exit code.
+func NewErrorWithExitCode(err error, exitCode int) *ErrorWithExitCode {
+	return &ErrorWithExitCode{
+		err:      err,
+		exitCode: exitCode,
+	}
+}
+
+// Error returns the error message string, implementing the `error` interface.
+func (e *ErrorWithExitCode) Error() string {
+	return e.err.Error()
+}
+
+// GetErrorExitCode checks the given exit code that the application should exit with, if this error is bubbled to
+// the top-level. This will be 0 for a nil error, 1 for a generic error, or arbitrary if the error is of type
+// ErrorWithExitCode.
+// Returns the exit code associated with the error.
+func GetErrorExitCode(err error) int {
+	// If we have no error, return 0, if we have a generic error, return 1, if we have a custom error code, unwrap
+	// and return it.
+	if err == nil {
+		return ExitCodeSuccess
+	} else if unwrappedErr, ok := err.(*ErrorWithExitCode); ok {
+		return unwrappedErr.exitCode
+	} else {
+		return ExitCodeGeneralError
+	}
+}
diff --git a/cmd/exitcodes/exit_codes.go b/cmd/exitcodes/exit_codes.go
new file mode 100644
index 00000000..cb6c7c98
--- /dev/null
+++ b/cmd/exitcodes/exit_codes.go
@@ -0,0 +1,21 @@
+package exitcodes
+
+const (
+	// ================================
+	// Platform-universal exit codes
+	// ================================
+
+	// ExitCodeSuccess indicates no errors or failures had occurred.
+	ExitCodeSuccess = 0
+
+	// ExitCodeGeneralError indicates some type of general error occurred.
+	ExitCodeGeneralError = 1
+
+	// ================================
+	// Application-specific exit codes
+	// ================================
+	// Note: Despite not being standardized, exit codes 2-5 are often used for common use cases, so we avoid them.
+
+	// ExitCodeTestFailed indicates a test case had failed.
+	ExitCodeTestFailed = 7
+)
diff --git a/cmd/fuzz.go b/cmd/fuzz.go
index 3fc22fd3..28453ae8 100644
--- a/cmd/fuzz.go
+++ b/cmd/fuzz.go
@@ -2,6 +2,7 @@ package cmd
 
 import (
 	"fmt"
+	"github.com/crytic/medusa/cmd/exitcodes"
 	"github.com/crytic/medusa/logging/colors"
 	"os"
 	"os/signal"
@@ -160,5 +161,10 @@ func cmdRunFuzz(cmd *cobra.Command, args []string) error {
 	// Start the fuzzing process with our cancellable context.
 	err = fuzzer.Start()
 
+	// If we have no error and failed test cases, we'll want to return a special exit code
+	if err == nil && len(fuzzer.TestCasesWithStatus(fuzzing.TestCaseStatusFailed)) > 0 {
+		return exitcodes.NewErrorWithExitCode(err, exitcodes.ExitCodeTestFailed)
+	}
+
 	return err
 }
diff --git a/main.go b/main.go
index 4902eaeb..ad2537e1 100644
--- a/main.go
+++ b/main.go
@@ -2,6 +2,7 @@ package main
 
 import (
 	"github.com/crytic/medusa/cmd"
+	"github.com/crytic/medusa/cmd/exitcodes"
 	"os"
 )
 
@@ -9,7 +10,6 @@ func main() {
 	// Run our root CLI command, which contains all underlying command logic and will handle parsing/invocation.
 	err := cmd.Execute()
 
-	if err != nil {
-		os.Exit(1)
-	}
+	// Determine the exit code from any potential error and exit out.
+	os.Exit(exitcodes.GetErrorExitCode(err))
 }