From 205148ad1fbb397d325d3aa984e7da8c94a59baf Mon Sep 17 00:00:00 2001
From: Mariano Cano <mariano.cano@gmail.com>
Date: Thu, 7 Oct 2021 12:43:24 -0700
Subject: [PATCH] Fix exit after defer.

---
 cmd/step-awskms-init/main.go   |  5 ++++-
 cmd/step-ca/main.go            | 10 ++++++++--
 cmd/step-cloudkms-init/main.go |  5 ++++-
 cmd/step-pkcs11-init/main.go   |  5 ++++-
 cmd/step-yubikey-init/main.go  |  5 ++++-
 5 files changed, 24 insertions(+), 6 deletions(-)

diff --git a/cmd/step-awskms-init/main.go b/cmd/step-awskms-init/main.go
index 7b9dc7a3..0678ef39 100644
--- a/cmd/step-awskms-init/main.go
+++ b/cmd/step-awskms-init/main.go
@@ -33,7 +33,6 @@ func main() {
 
 	// Initialize windows terminal
 	ui.Init()
-	defer ui.Reset()
 
 	c, err := awskms.New(context.Background(), apiv1.Options{
 		Type:            string(apiv1.AmazonKMS),
@@ -54,10 +53,14 @@ func main() {
 			fatal(err)
 		}
 	}
+
+	// Reset windows terminal
+	ui.Reset()
 }
 
 func fatal(err error) {
 	fmt.Fprintln(os.Stderr, err)
+	ui.Reset()
 	os.Exit(1)
 }
 
diff --git a/cmd/step-ca/main.go b/cmd/step-ca/main.go
index d9d17fed..e0123678 100644
--- a/cmd/step-ca/main.go
+++ b/cmd/step-ca/main.go
@@ -53,6 +53,11 @@ func init() {
 	rand.Seed(time.Now().UnixNano())
 }
 
+func exit(code int) {
+	ui.Reset()
+	os.Exit(code)
+}
+
 // appHelpTemplate contains the modified template for the main app
 var appHelpTemplate = `## NAME
 **{{.HelpName}}** -- {{.Usage}}
@@ -93,7 +98,6 @@ Please send us a sentence or two, good or bad: **feedback@smallstep.com** or htt
 func main() {
 	// Initialize windows terminal
 	ui.Init()
-	defer ui.Reset()
 
 	// Override global framework components
 	cli.VersionPrinter = func(c *cli.Context) {
@@ -169,8 +173,10 @@ $ step-ca $STEPPATH/config/ca.json --password-file ./password.txt
 		} else {
 			fmt.Fprintln(os.Stderr, err)
 		}
-		os.Exit(1)
+		exit(1)
 	}
+
+	exit(0)
 }
 
 func flagValue(f cli.Flag) reflect.Value {
diff --git a/cmd/step-cloudkms-init/main.go b/cmd/step-cloudkms-init/main.go
index b924f1a1..14bf50f1 100644
--- a/cmd/step-cloudkms-init/main.go
+++ b/cmd/step-cloudkms-init/main.go
@@ -64,7 +64,6 @@ func main() {
 
 	// Initialize windows terminal
 	ui.Init()
-	defer ui.Reset()
 
 	c, err := cloudkms.New(context.Background(), apiv1.Options{
 		Type:            string(apiv1.CloudKMS),
@@ -84,10 +83,14 @@ func main() {
 			fatal(err)
 		}
 	}
+
+	// Reset windows terminal
+	ui.Reset()
 }
 
 func fatal(err error) {
 	fmt.Fprintln(os.Stderr, err)
+	ui.Reset()
 	os.Exit(1)
 }
 
diff --git a/cmd/step-pkcs11-init/main.go b/cmd/step-pkcs11-init/main.go
index 5d9ba3e1..78c531c6 100644
--- a/cmd/step-pkcs11-init/main.go
+++ b/cmd/step-pkcs11-init/main.go
@@ -131,7 +131,6 @@ func main() {
 
 	// Initialize windows terminal
 	ui.Init()
-	defer ui.Reset()
 
 	if u.Get("pin-value") == "" && u.Get("pin-source") == "" && c.Pin == "" {
 		pin, err := ui.PromptPassword("What is the PKCS#11 PIN?")
@@ -205,6 +204,9 @@ func main() {
 	if err := createPKI(k, c); err != nil {
 		fatalClose(err, k)
 	}
+
+	// Reset windows terminal
+	ui.Reset()
 }
 
 func fatal(err error) {
@@ -213,6 +215,7 @@ func fatal(err error) {
 	} else {
 		fmt.Fprintln(os.Stderr, err)
 	}
+	ui.Reset()
 	os.Exit(1)
 }
 
diff --git a/cmd/step-yubikey-init/main.go b/cmd/step-yubikey-init/main.go
index cc55614a..163d0fcb 100644
--- a/cmd/step-yubikey-init/main.go
+++ b/cmd/step-yubikey-init/main.go
@@ -89,7 +89,6 @@ func main() {
 
 	// Initialize windows terminal
 	ui.Init()
-	defer ui.Reset()
 
 	pin, err := ui.PromptPassword("What is the YubiKey PIN?")
 	if err != nil {
@@ -123,6 +122,9 @@ func main() {
 	defer func() {
 		_ = k.Close()
 	}()
+
+	// Reset windows terminal
+	ui.Reset()
 }
 
 func fatal(err error) {
@@ -131,6 +133,7 @@ func fatal(err error) {
 	} else {
 		fmt.Fprintln(os.Stderr, err)
 	}
+	ui.Reset()
 	os.Exit(1)
 }