Fix hanging on non-serving commands

cmd/ecs_credential_provider.go

@@ -38,7 +38,7 @@ func init() {
 
 var ecsCredentialProvider = &cobra.Command{
 	Use:   "ecs_credential_provider",
-	Short: "Run a local ECS Credential Provider endpoint that serves and caches credentials for roles on demand",
+	Short: "RunService a local ECS Credential Provider endpoint that serves and caches credentials for roles on demand",
 	RunE:  runEcsMetadata,
 }
 
@@ -65,8 +65,6 @@ func runEcsMetadata(cmd *cobra.Command, args []string) error {
 	// Check for interrupt signal and exit cleanly
 	<-shutdown
 	log.Print("Shutdown signal received, exiting weep...")
-	// Send a signal to show that we're done shutting down
-	done <- 0
 
 	return nil
 }

cmd/metadata.go

@@ -42,7 +42,7 @@ func init() {
 
 var metadataCmd = &cobra.Command{
 	Use:   "metadata [role_name]",
-	Short: "Run a local Instance Metadata Service (IMDS) endpoint that serves credentials",
+	Short: "RunService a local Instance Metadata Service (IMDS) endpoint that serves credentials",
 	Args:  cobra.ExactArgs(1),
 	RunE:  runMetadata,
 }
@@ -86,8 +86,6 @@ func runMetadata(cmd *cobra.Command, args []string) error {
 	// Check for interrupt signal and exit cleanly
 	<-shutdown
 	log.Print("Shutdown signal received, exiting weep meta-data service...")
-	// Send a signal to show that we're done shutting down
-	done <- 0
 
 	return nil
 }

cmd/root.go

@@ -19,9 +19,11 @@ package cmd
 import (
 	"fmt"
 	"os"
+	"os/signal"
 	"path"
 	"runtime"
 	"strings"
+	"syscall"
 
 	"github.com/netflix/weep/util"
 
@@ -41,8 +43,6 @@ var (
 		Long:              "Weep is a CLI tool that manages AWS access via ConsoleMe for local development.",
 		DisableAutoGenTag: true,
 	}
-	shutdown chan os.Signal
-	done     chan int
 )
 
 func init() {
@@ -53,15 +53,20 @@ func init() {
 	rootCmd.PersistentFlags().StringSliceVarP(&assumeRole, "assume-role", "A", make([]string, 0), "one or more roles to assume after retrieving credentials")
 	rootCmd.PersistentFlags().StringVar(&logFormat, "log-format", "", "log format (json or tty)")
 	rootCmd.PersistentFlags().StringVar(&logLevel, "log-level", "", "log level (debug, info, warn)")
-
+	rootCmd.PersistentFlags().BoolVarP(&runAsService, "svc", "s", false, "run weep as a service")
 }
 
-func Execute(shutdownSignal chan os.Signal, doneSignal chan int) {
-	shutdown = shutdownSignal
-	done = doneSignal
-	if err := rootCmd.Execute(); err != nil {
-		_, _ = fmt.Fprintln(os.Stderr, err)
-		os.Exit(1)
+func Execute() {
+	shutdown := make(chan os.Signal, 1)
+	done = make(chan int, 1)
+	signal.Notify(shutdown, syscall.SIGINT, syscall.SIGTERM)
+	if runAsService {
+		RunService()
+	} else {
+		if err := rootCmd.Execute(); err != nil {
+			_, _ = fmt.Fprintln(os.Stderr, err)
+			os.Exit(1)
+		}
 	}
 }
 

run/run.go → cmd/svc.go

@@ -14,20 +14,14 @@
  * limitations under the License.
  */
 
-package run
+package cmd
 
 import (
-	"os"
-	"os/signal"
-	"syscall"
-
 	"github.com/kardianos/service"
-	"github.com/netflix/weep/cmd"
 	log "github.com/sirupsen/logrus"
 )
 
 var svcLogger service.Logger
-var done chan int
 
 type program struct{}
 
@@ -37,18 +31,16 @@ func (p *program) Start(s service.Service) error {
 }
 
 func (p *program) run() {
-	shutdown := make(chan os.Signal, 1)
-	done = make(chan int, 1)
-	signal.Notify(shutdown, syscall.SIGINT, syscall.SIGTERM)
-	cmd.Execute(shutdown, done)
+	_ = rootCmd.Execute()
+	done <- 0
 }
 
 func (p *program) Stop(s service.Service) error {
 	<-done
 	return nil
 }
 
-func Run() {
+func RunService() {
 	svcConfig := &service.Config{
 		Name:        "Weep",
 		DisplayName: "Weep",
@@ -68,4 +60,5 @@ func Run() {
 	if err != nil {
 		_ = svcLogger.Error(err)
 	}
+	done <- 0
 }

cmd/vars.go

@@ -16,6 +16,8 @@
 
 package cmd
 
+import "os"
+
 var (
 	assumeRole            []string
 	role                  string
@@ -32,4 +34,7 @@ var (
 	cfgFile               string
 	logLevel              string
 	logFormat             string
+	runAsService          bool
+	shutdown              chan os.Signal
+	done                  chan int
 )

main.go

@@ -19,7 +19,8 @@ package main
 import (
 	"os"
 
-	"github.com/netflix/weep/run"
+	"github.com/netflix/weep/cmd"
+
 	log "github.com/sirupsen/logrus"
 )
 
@@ -29,5 +30,5 @@ func init() {
 }
 
 func main() {
-	run.Run()
+	cmd.Execute()
 }