From 1db9c95d5bd805b682826ca6c46191c23847a3b7 Mon Sep 17 00:00:00 2001
From: Hidetatsu Yaginuma <deetyler@protonmail.com>
Date: Fri, 7 May 2021 15:52:16 +0900
Subject: [PATCH] handle panic

---
 command/runner.go          | 25 ++++++++++++++++++++-----
 command/subcommand.go      |  4 ----
 kubectl/subcommand.go      | 14 +-------------
 kubectl/subcommand_test.go |  2 --
 4 files changed, 21 insertions(+), 24 deletions(-)

diff --git a/command/runner.go b/command/runner.go
index 03a922e..1f317dc 100644
--- a/command/runner.go
+++ b/command/runner.go
@@ -1,7 +1,9 @@
 package command
 
 import (
+	"bytes"
 	"fmt"
+	"io"
 	"os"
 	"os/exec"
 	"strings"
@@ -57,7 +59,7 @@ func Run(args []string, version string) error {
 
 	// when should not colorize, just run command and return
 	// TODO: right now, krew is unsupported by kubecolor but it should be.
-	if !shouldColorize || subcommandInfo.IsKrew {
+	if !shouldColorize {
 		cmd.Stdout = Stdout
 		cmd.Stderr = Stderr
 		if err := cmd.Start(); err != nil {
@@ -72,16 +74,21 @@ func Run(args []string, version string) error {
 	}
 
 	// when colorize, capture stdout and err then colorize it
-	outReader, err := cmd.StdoutPipe()
+	cmdOut, err := cmd.StdoutPipe()
 	if err != nil {
 		return err
 	}
 
-	errReader, err := cmd.StderrPipe()
+	cmdErr, err := cmd.StderrPipe()
 	if err != nil {
 		return err
 	}
 
+	// make buffer to be used in defer recover()
+	buff := new(bytes.Buffer)
+	outReader := io.TeeReader(cmdOut, buff)
+	errReader := io.TeeReader(cmdErr, buff)
+
 	if err := cmd.Start(); err != nil {
 		return err
 	}
@@ -92,14 +99,22 @@ func Run(args []string, version string) error {
 
 	wg.Add(1)
 	go func() {
+		defer wg.Done()
+		defer func() {
+			if r := recover(); r != nil {
+				fmt.Fprintf(os.Stdout, buff.String())
+			}
+		}()
+
+		// This can panic when kubecolor has bug, so recover in defer
 		printers.FullColoredPrinter.Print(outReader, Stdout)
-		wg.Done()
 	}()
 
 	wg.Add(1)
 	go func() {
+		defer wg.Done()
+		// This will unlikely panic
 		printers.ErrorPrinter.Print(errReader, Stderr)
-		wg.Done()
 	}()
 
 	wg.Wait()
diff --git a/command/subcommand.go b/command/subcommand.go
index 7b0c453..7be6329 100644
--- a/command/subcommand.go
+++ b/command/subcommand.go
@@ -16,10 +16,6 @@ func ResolveSubcommand(args []string, config *KubecolorConfig) (bool, *kubectl.S
 	// subcommandFound becomes false when subcommand is not found; e.g. "kubecolor --help"
 	subcommandInfo, subcommandFound := kubectl.InspectSubcommandInfo(args)
 
-	if subcommandInfo.IsKrew {
-		return false, subcommandInfo
-	}
-
 	// if --plain found, it does not colorize
 	if config.Plain {
 		return false, subcommandInfo
diff --git a/kubectl/subcommand.go b/kubectl/subcommand.go
index ac4b678..679e33d 100644
--- a/kubectl/subcommand.go
+++ b/kubectl/subcommand.go
@@ -195,22 +195,10 @@ func CollectCommandlineOptions(args []string, info *SubcommandInfo) {
 	}
 }
 
+// TODO: return shouldColorize = false when the given args is for plugin
 func InspectSubcommandInfo(args []string) (*SubcommandInfo, bool) {
-	// TODO: support krew
-	contains := func(s []string, e string) bool {
-		for _, a := range s {
-			if a == e {
-				return true
-			}
-		}
-		return false
-	}
 	ret := &SubcommandInfo{}
 
-	if contains(args, "krew") {
-		return &SubcommandInfo{IsKrew: true}, false
-	}
-
 	CollectCommandlineOptions(args, ret)
 
 	for i := range args {
diff --git a/kubectl/subcommand_test.go b/kubectl/subcommand_test.go
index a116e09..33153fe 100644
--- a/kubectl/subcommand_test.go
+++ b/kubectl/subcommand_test.go
@@ -62,8 +62,6 @@ func TestInspectSubcommandInfo(t *testing.T) {
 
 		{"apply", &SubcommandInfo{Subcommand: Apply}, true},
 
-		{"krew version", &SubcommandInfo{IsKrew: true}, false},
-
 		{"", &SubcommandInfo{}, false},
 	}
 	for _, tt := range tests {