Fix Fatal CLI Restore (#345)

* restore cli state on failures (Fatal)

Author: frjcomp

src/pipeleak/cmd/root.go

@@ -19,6 +19,9 @@ import (
 	"github.com/spf13/cobra"
 )
 
+// TerminalRestorer is a function that can be called to restore terminal state
+var TerminalRestorer func()
+
 var (
 	rootCmd = &cobra.Command{
 		Use:     "pipeleak",
@@ -70,6 +73,7 @@ type CustomWriter struct {
 
 func (cw *CustomWriter) Write(p []byte) (n int, err error) {
 	originalLen := len(p)
+
 	if bytes.HasSuffix(p, []byte("\n")) {
 		p = bytes.TrimSuffix(p, []byte("\n"))
 	}
@@ -94,6 +98,17 @@ func (cw *CustomWriter) Write(p []byte) (n int, err error) {
 	return originalLen, nil
 }
 
+// FatalHook is a zerolog hook that restores terminal state before fatal exits
+type FatalHook struct{}
+
+func (h FatalHook) Run(e *zerolog.Event, level zerolog.Level, msg string) {
+	if level == zerolog.FatalLevel {
+		if TerminalRestorer != nil {
+			TerminalRestorer()
+		}
+	}
+}
+
 func initLogger(cmd *cobra.Command) {
 	defaultOut := &CustomWriter{Writer: os.Stdout}
 	colorEnabled := LogColor
@@ -115,12 +130,15 @@ func initLogger(cmd *cobra.Command) {
 		}
 	}
 
+	// Create the fatal hook to restore terminal state
+	fatalHook := FatalHook{}
+
 	if JsonLogoutput {
 		// For JSON output, wrap with HitLevelWriter to transform level field
 		hitWriter := &logging.HitLevelWriter{}
 		hitWriter.SetOutput(defaultOut)
 		logging.SetGlobalHitWriter(hitWriter)
-		log.Logger = zerolog.New(hitWriter).With().Timestamp().Logger()
+		log.Logger = zerolog.New(hitWriter).With().Timestamp().Logger().Hook(fatalHook)
 	} else {
 		// For console output, use custom FormatLevel to color the hit level
 		output := zerolog.ConsoleWriter{
@@ -133,7 +151,7 @@ func initLogger(cmd *cobra.Command) {
 		hitWriter := &logging.HitLevelWriter{}
 		hitWriter.SetOutput(&output)
 		logging.SetGlobalHitWriter(hitWriter)
-		log.Logger = zerolog.New(hitWriter).With().Timestamp().Logger()
+		log.Logger = zerolog.New(hitWriter).With().Timestamp().Logger().Hook(fatalHook)
 	}
 }
 

src/pipeleak/cmd/root_terminal_test.go

@@ -0,0 +1,124 @@
+package cmd
+
+import (
+	"bytes"
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/rs/zerolog"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestTerminalRestorer(t *testing.T) {
+	t.Run("TerminalRestorer_can_be_set", func(t *testing.T) {
+		originalRestorer := TerminalRestorer
+		defer func() { TerminalRestorer = originalRestorer }()
+
+		called := false
+		TerminalRestorer = func() { called = true }
+
+		TerminalRestorer()
+		assert.True(t, called, "TerminalRestorer should be callable")
+	})
+
+	t.Run("TerminalRestorer_nil_safe", func(t *testing.T) {
+		originalRestorer := TerminalRestorer
+		defer func() { TerminalRestorer = originalRestorer }()
+
+		TerminalRestorer = nil
+		assert.NotPanics(t, func() {
+			if TerminalRestorer != nil {
+				TerminalRestorer()
+			}
+		})
+	})
+}
+
+func TestFatalHook(t *testing.T) {
+	t.Run("fatal_level_calls_TerminalRestorer", func(t *testing.T) {
+		originalRestorer := TerminalRestorer
+		defer func() { TerminalRestorer = originalRestorer }()
+
+		called := false
+		TerminalRestorer = func() { called = true }
+
+		hook := FatalHook{}
+
+		// Create a dummy event and run the hook with fatal level
+		var buf bytes.Buffer
+		logger := zerolog.New(&buf)
+		event := logger.Fatal()
+
+		hook.Run(event, zerolog.FatalLevel, "test")
+
+		assert.True(t, called, "TerminalRestorer should be called for fatal level")
+	})
+
+	t.Run("non_fatal_level_does_not_call_TerminalRestorer", func(t *testing.T) {
+		originalRestorer := TerminalRestorer
+		defer func() { TerminalRestorer = originalRestorer }()
+
+		called := false
+		TerminalRestorer = func() { called = true }
+
+		hook := FatalHook{}
+
+		levels := []zerolog.Level{
+			zerolog.TraceLevel,
+			zerolog.DebugLevel,
+			zerolog.InfoLevel,
+			zerolog.WarnLevel,
+			zerolog.ErrorLevel,
+		}
+
+		var buf bytes.Buffer
+		logger := zerolog.New(&buf)
+		for _, lvl := range levels {
+			event := logger.WithLevel(lvl)
+			hook.Run(event, lvl, "test")
+			assert.False(t, called, "TerminalRestorer should not be called for non-fatal levels")
+		}
+	})
+
+	t.Run("nil_TerminalRestorer_does_not_panic", func(t *testing.T) {
+		originalRestorer := TerminalRestorer
+		defer func() { TerminalRestorer = originalRestorer }()
+
+		TerminalRestorer = nil
+		hook := FatalHook{}
+
+		assert.NotPanics(t, func() {
+			var buf bytes.Buffer
+			logger := zerolog.New(&buf)
+			event := logger.Fatal()
+			hook.Run(event, zerolog.FatalLevel, "test")
+		})
+	})
+}
+
+func TestCustomWriter_WritesCorrectly(t *testing.T) {
+	t.Run("Writes_log_to_file", func(t *testing.T) {
+		tmpDir := t.TempDir()
+		logFile := filepath.Join(tmpDir, "test.log")
+
+	f, err := os.OpenFile(logFile, os.O_CREATE|os.O_WRONLY, 0644)
+	require.NoError(t, err)
+	defer func() { _ = f.Close() }()
+
+		writer := &CustomWriter{Writer: f}
+
+		testLog := []byte(`{"level":"info","message":"test"}` + "\n")
+		n, err := writer.Write(testLog)
+
+		assert.NoError(t, err)
+		assert.Equal(t, len(testLog), n, "Should return original length")
+
+	// Read back and verify
+	_ = f.Close()
+		content, err := os.ReadFile(logFile)
+		require.NoError(t, err)
+		assert.Contains(t, string(content), "test", "Log content should be written")
+	})
+}

src/pipeleak/main.go

@@ -1,18 +1,39 @@
 package main
 
 import (
+	"encoding/json"
+	"io"
 	_ "net/http/pprof"
 	"os"
 
 	"github.com/CompassSecurity/pipeleak/cmd"
 	"golang.org/x/term"
 )
 
-var originalTerminalState *term.State
+var originalTermState *term.State
+
+type TerminalRestoringWriter struct {
+	underlying io.Writer
+}
+
+func (w *TerminalRestoringWriter) Write(p []byte) (n int, err error) {
+	var logEntry map[string]interface{}
+	if err := json.Unmarshal(p, &logEntry); err == nil {
+		if level, ok := logEntry["level"].(string); ok && level == "fatal" {
+			_, _ = w.underlying.Write(p)
+			restoreTerminalState()
+			os.Exit(1)
+		}
+	}
+	return w.underlying.Write(p)
+}
 
 func main() {
 	saveTerminalState()
 	defer restoreTerminalState()
+
+	cmd.TerminalRestorer = restoreTerminalState
+
 	err := cmd.Execute()
 	if err != nil {
 		os.Exit(1)
@@ -23,13 +44,18 @@ func saveTerminalState() {
 	if term.IsTerminal(int(os.Stdin.Fd())) {
 		state, err := term.GetState(int(os.Stdin.Fd()))
 		if err == nil {
-			originalTerminalState = state
+			originalTermState = state
 		}
 	}
 }
 
 func restoreTerminalState() {
-	if originalTerminalState != nil {
-		_ = term.Restore(int(os.Stdin.Fd()), originalTerminalState)
+	if originalTermState != nil {
+		_ = term.Restore(int(os.Stdin.Fd()), originalTermState)
 	}
 }
+
+// GetTerminalRestoringWriter wraps the given writer to restore terminal state on fatal logs
+func GetTerminalRestoringWriter(w io.Writer) io.Writer {
+	return &TerminalRestoringWriter{underlying: w}
+}