Refresh expired websocket sessions (stripe#468)

Author: richardm-stripe

pkg/ansi/ansi.go

@@ -105,24 +105,30 @@ func Linkify(text, url string, w io.Writer) string {
 	return fmt.Sprintf("\x1b]8;;%s\x1b\\%s\x1b]8;;\x1b\\", url, text)
 }
 
-// StartSpinner starts a spinner with the given message. If the writer doesn't
-// support colors, it simply prints the message.
-func StartSpinner(msg string, w io.Writer) *spinner.Spinner {
-	if !shouldUseColors(w) {
-		fmt.Fprintln(w, msg)
-		return nil
-	}
+type charset = []string
 
+func getCharset() charset {
 	// See https://github.com/briandowns/spinner#available-character-sets for
 	// list of available charsets
-	charSetIdx := 11
 	if runtime.GOOS == "windows" {
 		// Less fancy, but uses ASCII characters so works with Windows default
 		// console.
-		charSetIdx = 8
+		return spinner.CharSets[8]
+	}
+	return spinner.CharSets[11]
+}
+
+const duration = time.Duration(100) * time.Millisecond
+
+// StartNewSpinner starts a new spinner with the given message. If the writer doesn't
+// support colors, it simply prints the message.
+func StartNewSpinner(msg string, w io.Writer) *spinner.Spinner {
+	if !shouldUseColors(w) {
+		fmt.Fprintln(w, msg)
+		return nil
 	}
 
-	s := spinner.New(spinner.CharSets[charSetIdx], time.Duration(100)*time.Millisecond)
+	s := spinner.New(getCharset(), duration)
 	s.Writer = w
 
 	if msg != "" {
@@ -134,6 +140,20 @@ func StartSpinner(msg string, w io.Writer) *spinner.Spinner {
 	return s
 }
 
+// StartSpinner updates an existing spinner's message, and starts it if it was stopped
+func StartSpinner(s *spinner.Spinner, msg string, w io.Writer) {
+	if s == nil {
+		fmt.Fprintln(w, msg)
+		return
+	}
+	if msg != "" {
+		s.Suffix = " " + msg
+	}
+	if !s.Active() {
+		s.Start()
+	}
+}
+
 // StopSpinner stops a spinner with the given message. If the writer doesn't
 // support colors, it simply prints the message.
 func StopSpinner(s *spinner.Spinner, msg string, w io.Writer) {

pkg/cmd/listen.go

@@ -1,6 +1,7 @@
 package cmd
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"net/url"
@@ -163,7 +164,7 @@ func (lc *listenCmd) runListenCmd(cmd *cobra.Command, args []string) error {
 		NoWSS:               lc.noWSS,
 	}, lc.events)
 
-	err = p.Run()
+	err = p.Run(context.Background())
 	if err != nil {
 		return err
 	}

pkg/cmd/logs/tail.go

@@ -6,6 +6,8 @@ import (
 	log "github.com/sirupsen/logrus"
 	"github.com/spf13/cobra"
 
+	"context"
+
 	"github.com/stripe/stripe-cli/pkg/config"
 	logTailing "github.com/stripe/stripe-cli/pkg/logtailing"
 	"github.com/stripe/stripe-cli/pkg/validators"
@@ -158,7 +160,7 @@ func (tailCmd *TailCmd) runTailCmd(cmd *cobra.Command, args []string) error {
 		WebSocketFeature: requestLogsWebSocketFeature,
 	})
 
-	err = tailer.Run()
+	err = tailer.Run(context.Background())
 	if err != nil {
 		return err
 	}

pkg/cmd/samples/create.go

@@ -84,7 +84,7 @@ To see supported samples, run 'stripe samples list'`, args[0])
 		return fmt.Errorf("Path already exists for: %s", destination)
 	}
 
-	spinner := ansi.StartSpinner(fmt.Sprintf("Downloading %s", selectedSample), os.Stdout)
+	spinner := ansi.StartNewSpinner(fmt.Sprintf("Downloading %s", selectedSample), os.Stdout)
 
 	if cc.forceRefresh {
 		err := sample.DeleteCache(selectedSample)
@@ -145,7 +145,7 @@ To see supported samples, run 'stripe samples list'`, args[0])
 		os.Exit(1)
 	}()
 
-	spinner = ansi.StartSpinner(fmt.Sprintf("Copying files over... %s", selectedSample), os.Stdout)
+	spinner = ansi.StartNewSpinner(fmt.Sprintf("Copying files over... %s", selectedSample), os.Stdout)
 	// Create the target folder to copy the sample in to. We do
 	// this here in case any of the steps above fail, minimizing
 	// the change that we create a dangling empty folder
@@ -164,7 +164,7 @@ To see supported samples, run 'stripe samples list'`, args[0])
 	ansi.StopSpinner(spinner, "", os.Stdout)
 	fmt.Printf("%s %s\n", color.Green("✔"), ansi.Faint("Files copied"))
 
-	spinner = ansi.StartSpinner(fmt.Sprintf("Configuring your code... %s", selectedSample), os.Stdout)
+	spinner = ansi.StartNewSpinner(fmt.Sprintf("Configuring your code... %s", selectedSample), os.Stdout)
 
 	err = sample.ConfigureDotEnv(targetPath)
 	if err != nil {

pkg/cmd/status.go

@@ -77,7 +77,7 @@ func (sc *statusCmd) runStatusCmd(cmd *cobra.Command, args []string) error {
 		if sc.hideSpinner {
 			time.Sleep(time.Duration(sc.pollRate) * time.Second)
 		} else {
-			spinner := ansi.StartSpinner("", os.Stderr)
+			spinner := ansi.StartNewSpinner("", os.Stderr)
 			time.Sleep(time.Duration(sc.pollRate) * time.Second)
 			ansi.StopSpinner(spinner, "", os.Stderr)
 		}

pkg/login/client_login.go

@@ -54,18 +54,18 @@ func Login(baseURL string, config *config.Config, input io.Reader) error {
 	if isSSH() {
 		fmt.Printf("To authenticate with Stripe, please go to: %s\n", links.BrowserURL)
 
-		s = ansi.StartSpinner("Waiting for confirmation...", os.Stdout)
+		s = ansi.StartNewSpinner("Waiting for confirmation...", os.Stdout)
 	} else {
 		fmt.Printf("Press Enter to open the browser (^C to quit)")
 		fmt.Fscanln(input)
 
-		s = ansi.StartSpinner("Waiting for confirmation...", os.Stdout)
+		s = ansi.StartNewSpinner("Waiting for confirmation...", os.Stdout)
 
 		err = openBrowser(links.BrowserURL)
 		if err != nil {
 			msg := fmt.Sprintf("Failed to open browser, please go to %s manually.", links.BrowserURL)
 			ansi.StopSpinner(s, msg, os.Stdout)
-			s = ansi.StartSpinner("Waiting for confirmation...", os.Stdout)
+			s = ansi.StartNewSpinner("Waiting for confirmation...", os.Stdout)
 		}
 	}
 

pkg/logtailing/tailer.go

@@ -66,8 +66,6 @@ type Tailer struct {
 	webSocketClient  *websocket.Client
 
 	interruptCh chan os.Signal
-
-	ctx context.Context
 }
 
 // EventPayload is the mapping for fields in event payloads from request log tailing
@@ -107,60 +105,84 @@ func New(cfg *Config) *Tailer {
 	}
 }
 
-// Run sets the websocket connection
-func (t *Tailer) Run() error {
-	s := ansi.StartSpinner("Getting ready...", t.cfg.Log.Out)
-
+func withSIGTERMCancel(ctx context.Context, onCancel func()) context.Context {
 	// Create a context that will be canceled when Ctrl+C is pressed
-	ctx, cancel := context.WithCancel(context.Background())
-	t.ctx = ctx
-	signal.Notify(t.interruptCh, os.Interrupt, syscall.SIGTERM)
+	ctx, cancel := context.WithCancel(ctx)
+
+	interruptCh := make(chan os.Signal, 1)
+	signal.Notify(interruptCh, os.Interrupt, syscall.SIGTERM)
 
 	go func() {
+		<-interruptCh
+		onCancel()
+		cancel()
+	}()
+	return ctx
+}
+
+const maxConnectAttempts = 3
+
+// Run sets the websocket connection
+func (t *Tailer) Run(ctx context.Context) error {
+	s := ansi.StartNewSpinner("Getting ready...", t.cfg.Log.Out)
+
+	ctx = withSIGTERMCancel(ctx, func() {
 		log.WithFields(log.Fields{
 			"prefix": "logtailing.Tailer.Run",
 		}).Debug("Ctrl+C received, cleaning up...")
+	})
 
-		<-t.interruptCh
-		cancel()
-	}()
+	var warned = false
+	var nAttempts int = 0
 
-	// Create the CLI session
-	session, err := t.createSession()
-	if err != nil {
-		ansi.StopSpinner(s, "", t.cfg.Log.Out)
-		t.cfg.Log.Fatalf("Error while authenticating with Stripe: %v", err)
-	}
+	for nAttempts < maxConnectAttempts {
+		session, err := t.createSession(ctx)
 
-	// Create and start the websocket client
-	t.webSocketClient = websocket.NewClient(
-		session.WebSocketURL,
-		session.WebSocketID,
-		session.WebSocketAuthorizedFeature,
-		&websocket.Config{
-			EventHandler:      websocket.EventHandlerFunc(t.processRequestLogEvent),
-			Log:               t.cfg.Log,
-			NoWSS:             t.cfg.NoWSS,
-			ReconnectInterval: time.Duration(session.ReconnectDelay) * time.Second,
-		},
-	)
-	go t.webSocketClient.Run()
-
-	select {
-	case <-t.webSocketClient.Connected():
-		ansi.StopSpinner(s, "Ready! You're now waiting to receive API request logs (^C to quit)", t.cfg.Log.Out)
-	case <-t.ctx.Done():
-		ansi.StopSpinner(s, "", t.cfg.Log.Out)
-		t.cfg.Log.Fatalf("Aborting")
-	}
+		if err != nil {
+			ansi.StopSpinner(s, "", t.cfg.Log.Out)
+			t.cfg.Log.Fatalf("Error while authenticating with Stripe: %v", err)
+		}
 
-	if session.DisplayConnectFilterWarning {
-		color := ansi.Color(os.Stdout)
-		fmt.Printf("%s you specified the 'account' filter for Connect accounts but are not a Connect user, so the filter will not be applied.\n", color.Yellow("Warning"))
-	}
+		if session.DisplayConnectFilterWarning && !warned {
+			color := ansi.Color(os.Stdout)
+			fmt.Printf("%s you specified the 'account' filter for Connect accounts but are not a Connect user, so the filter will not be applied.\n", color.Yellow("Warning"))
+			// Only display this warning once
+			warned = true
+		}
 
-	// Block until context is done (i.e. Ctrl+C is pressed)
-	<-t.ctx.Done()
+		t.webSocketClient = websocket.NewClient(
+			session.WebSocketURL,
+			session.WebSocketID,
+			session.WebSocketAuthorizedFeature,
+			&websocket.Config{
+				EventHandler:      websocket.EventHandlerFunc(t.processRequestLogEvent),
+				Log:               t.cfg.Log,
+				NoWSS:             t.cfg.NoWSS,
+				ReconnectInterval: time.Duration(session.ReconnectDelay) * time.Second,
+			},
+		)
+
+		go func() {
+			<-t.webSocketClient.Connected()
+			nAttempts = 0
+			ansi.StopSpinner(s, "Ready! You're now waiting to receive API request logs (^C to quit)", t.cfg.Log.Out)
+		}()
+
+		go t.webSocketClient.Run(ctx)
+		nAttempts++
+
+		select {
+		case <-ctx.Done():
+			ansi.StopSpinner(s, "", t.cfg.Log.Out)
+			t.cfg.Log.Fatalf("Aborting")
+		case <-t.webSocketClient.NotifyExpired:
+			if nAttempts < maxConnectAttempts {
+				ansi.StartSpinner(s, "Session expired, reconnecting...", t.cfg.Log.Out)
+			} else {
+				t.cfg.Log.Fatalf("Session expired. Terminating after %d failed attempts to reauthorize", nAttempts)
+			}
+		}
+	}
 
 	if t.webSocketClient != nil {
 		t.webSocketClient.Stop()
@@ -173,7 +195,7 @@ func (t *Tailer) Run() error {
 	return nil
 }
 
-func (t *Tailer) createSession() (*stripeauth.StripeCLISession, error) {
+func (t *Tailer) createSession(ctx context.Context) (*stripeauth.StripeCLISession, error) {
 	var session *stripeauth.StripeCLISession
 
 	var err error
@@ -189,15 +211,15 @@ func (t *Tailer) createSession() (*stripeauth.StripeCLISession, error) {
 		// Try to authorize at least 5 times before failing. Sometimes we have random
 		// transient errors that we just need to retry for.
 		for i := 0; i <= 5; i++ {
-			session, err = t.stripeAuthClient.Authorize(t.ctx, t.cfg.DeviceName, t.cfg.WebSocketFeature, &filters)
+			session, err = t.stripeAuthClient.Authorize(ctx, t.cfg.DeviceName, t.cfg.WebSocketFeature, &filters)
 
 			if err == nil {
 				exitCh <- struct{}{}
 				return
 			}
 
 			select {
-			case <-t.ctx.Done():
+			case <-ctx.Done():
 				exitCh <- struct{}{}
 				return
 			case <-time.After(1 * time.Second):