Repository: stripe/stripe-cli (go)

Issue:
stripe listen does not log when it cannot create a websocket connection
### Issue
I'm not sure why `stripe listen` isn't working - maybe the endpoint is down or I've made a mistake (it was working earlier with no changes on my end). 

Regardless, it's super confusing that the CLI says "Ready! Your webhook signing secret..." when in fact the CLI is silently dialing the websocket endpoint and repeatedly failing.

### Expected Behavior
I expected the `stripe listen` to log its connection failures / retries (without having to run the command with `--log-level=debug`).

### Steps to reproduce
Run `stripe listen` in such a way that it authenticates properly but can't create a websocket connection to stripe-cli.stripe.com.

### Traceback
```
stripe listen --forward-to localhost:4000/webhooks/stripeaccountupdated --print-json --log-level debug
⣾ Getting ready... [Tue, 10 Sep 2019 20:42:17 EDT] DEBUG stripeauth.client.Authorize: Authenticating with Stripe...
⡿ Getting ready... [Tue, 10 Sep 2019 20:42:27 EDT] DEBUG stripeauth.Client.Authorize: Got successful response from Stripe display_connect_filter_warning=false reconnect_delay=36 websocket_authorized_feature=webhook-payloads websocket_id={ws-id}websocket_url=wss://stripe-cli.stripe.com/subscribe/{account}
[Tue, 10 Sep 2019 20:42:27 EDT] DEBUG websocket.client.Run: Attempting to connect to Stripe
> Ready! Your webhook signing secret is {secret} (^C to quit)
[Tue, 10 Sep 2019 20:42:27 EDT] DEBUG websocket.Client.connect: Dialing websocket url=wss://stripe-cli.stripe.com/subscribe/{account}?websocket_feature=webhook-payloads
[Tue, 10 Sep 2019 20:42:37 EDT] DEBUG websocket.Client.connect: Websocket connection error error=dial tcp: i/o timeout
[Tue, 10 Sep 2019 20:42:37 EDT] DEBUG websocket.client.Run: Failed to connect to Stripe. Retrying...
[Tue, 10 Sep 2019 20:42:47 EDT] DEBUG websocket.Client.connect: Dialing websocket url=wss://stripe-cli.stripe.com/subscribe/{account}?websocket_feature=webhook-payloads
[Tue, 10 Sep 2019 20:42:57 EDT] DEBUG websocket.Client.connect: Websocket connection error error=dial tcp: i/o timeout
[Tue, 10 Sep 2019 20:42:57 EDT] DEBUG websocket.client.Run: Failed to connect to Stripe. Retrying...
^C[Tue, 10 Sep 2019 20:43:00 EDT] DEBUG proxy.Proxy.Run: Ctrl+C received, cleaning up...
[Tue, 10 Sep 2019 20:43:00 EDT] DEBUG proxy.Proxy.Run: Bye!
```

### Environment
macOS


Pull request description:
Better synchronization for websocket commands
 ### Reviewers
r? @tomer-stripe 
cc @stripe/dev-platform

 ### Summary
Improves the synchronization of all goroutines in websocket commands:
- create a cancelable [context](https://blog.golang.org/context)
- start a goroutine who will be blocked until Ctrl+C is pressed. When Ctrl+C is pressed, the goroutine cancels the context
- create a goroutine for the session creation request and pass the context to the HTTP client
- block until the goroutine returns
- if the session creation succeeds, pass the context to the websocket client
- block until either the context is canceled or the client is connected
- at this point only, display the success message (previously we would start the websocket client goroutine and immediately display the success message despite the fact that the websocket client was still connecting)

Fixes #156.

Future improvements:
- there's a lot of duplicated code between the proxy and logtailing packages, we should factorize it
- the websocket client still attempts to reconnect endlessly, even when the websocket ID is no longer valid


Patch:
diff --git a/Makefile b/Makefile
index c9d8fab..2109dcd 100644
--- a/Makefile
+++ b/Makefile
@@ -6,10 +6,11 @@ export GO111MODULE := on
 export GOBIN := $(shell pwd)/bin
 export PATH := $(GOBIN):$(PATH)
 export GOPROXY := https://gocenter.io
+export GOLANGCI_LINT_VERSION := v1.21.0
 
 # Install all the build and lint dependencies
 setup:
-	curl -sfL https://install.goreleaser.com/github.com/golangci/golangci-lint.sh | sh
+	curl -sfL https://install.goreleaser.com/github.com/golangci/golangci-lint.sh | sh -s $(GOLANGCI_LINT_VERSION)
 	curl -L https://git.io/misspell | sh
 	go mod download
 .PHONY: setup
diff --git a/pkg/login/client_login.go b/pkg/login/client_login.go
index 78f5334..0128ec0 100644
--- a/pkg/login/client_login.go
+++ b/pkg/login/client_login.go
@@ -1,6 +1,7 @@
 package login
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"io"
@@ -113,7 +114,7 @@ func getLinks(baseURL string, deviceName string) (*Links, error) {
 	data := url.Values{}
 	data.Set("device_name", deviceName)
 
-	res, err := client.PerformRequest(http.MethodPost, stripeCLIAuthPath, data.Encode(), nil)
+	res, err := client.PerformRequest(context.TODO(), http.MethodPost, stripeCLIAuthPath, data.Encode(), nil)
 	if err != nil {
 		return nil, err
 	}
diff --git a/pkg/login/login_message.go b/pkg/login/login_message.go
index 7607bf2..b68c417 100644
--- a/pkg/login/login_message.go
+++ b/pkg/login/login_message.go
@@ -1,6 +1,7 @@
 package login
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"net/url"
@@ -72,7 +73,7 @@ func getUserAccount(baseURL string, apiKey string) (*Account, error) {
 		APIKey:  apiKey,
 	}
 
-	resp, err := client.PerformRequest("GET", "/v1/account", "", nil)
+	resp, err := client.PerformRequest(context.TODO(), "GET", "/v1/account", "", nil)
 
 	if err != nil {
 		return nil, err
diff --git a/pkg/login/poll.go b/pkg/login/poll.go
index ab078cd..2f6ffa4 100644
--- a/pkg/login/poll.go
+++ b/pkg/login/poll.go
@@ -1,6 +1,7 @@
 package login
 
 import (
+	"context"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -51,7 +52,7 @@ func PollForKey(pollURL string, interval time.Duration, maxAttempts int) (*PollA
 
 	var count = 0
 	for count < maxAttempts {
-		res, err := client.PerformRequest(http.MethodGet, parsedURL.Path, parsedURL.Query().Encode(), nil)
+		res, err := client.PerformRequest(context.TODO(), http.MethodGet, parsedURL.Path, parsedURL.Query().Encode(), nil)
 		if err != nil {
 			return nil, nil, err
 		}
diff --git a/pkg/logtailing/tailer.go b/pkg/logtailing/tailer.go
index b43bf3a..e908b2e 100644
--- a/pkg/logtailing/tailer.go
+++ b/pkg/logtailing/tailer.go
@@ -1,6 +1,7 @@
 package logtailing
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"io/ioutil"
@@ -65,6 +66,8 @@ 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
@@ -105,85 +108,134 @@ func New(cfg *Config) *Tailer {
 }
 
 // Run sets the websocket connection
-func (tailer *Tailer) Run() error {
-	s := ansi.StartSpinner("Getting ready...", tailer.cfg.Log.Out)
+func (t *Tailer) Run() error {
+	s := ansi.StartSpinner("Getting ready...", t.cfg.Log.Out)
 
-	// Intercept Ctrl+c so we can do some clean up
-	signal.Notify(tailer.interruptCh, os.Interrupt, syscall.SIGTERM)
+	// 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)
 
-	filters, err := jsonifyFilters(tailer.cfg.Filters)
-	if err != nil {
-		tailer.cfg.Log.Fatalf("Error while converting log filters to JSON encoding: %v", err)
-	}
+	go func() {
+		log.WithFields(log.Fields{
+			"prefix": "logtailing.Tailer.Run",
+		}).Debug("Ctrl+C received, cleaning up...")
 
-	session, err := tailer.stripeAuthClient.Authorize(tailer.cfg.DeviceName, tailer.cfg.WebSocketFeature, &filters)
+		<-t.interruptCh
+		cancel()
+	}()
+
+	// Create the CLI session
+	session, err := t.createSession()
 	if err != nil {
-		tailer.cfg.Log.Fatalf("Error while authenticating with Stripe: %v", err)
+		ansi.StopSpinner(s, "", t.cfg.Log.Out)
+		t.cfg.Log.Fatalf("Error while authenticating with Stripe: %v", err)
 	}
 
-	tailer.webSocketClient = websocket.NewClient(
+	// Create and start the websocket client
+	t.webSocketClient = websocket.NewClient(
 		session.WebSocketURL,
 		session.WebSocketID,
 		session.WebSocketAuthorizedFeature,
 		&websocket.Config{
-			EventHandler:      websocket.EventHandlerFunc(tailer.processRequestLogEvent),
-			Log:               tailer.cfg.Log,
-			NoWSS:             tailer.cfg.NoWSS,
+			EventHandler:      websocket.EventHandlerFunc(t.processRequestLogEvent),
+			Log:               t.cfg.Log,
+			NoWSS:             t.cfg.NoWSS,
 			ReconnectInterval: time.Duration(session.ReconnectDelay) * time.Second,
 		},
 	)
-	go tailer.webSocketClient.Run()
-
-	ansi.StopSpinner(s, "Ready! You're now waiting to receive API request logs (^C to quit)", tailer.cfg.Log.Out)
+	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 session.DisplayConnectFilterWarning {
 		color := ansi.Color(os.Stdout)
-		fmt.Println(fmt.Sprintf("%s you specified the 'account' filter for connect accounts but are not a connect merchant, so the filter will not be applied.", color.Yellow("Warning")))
+		fmt.Println(fmt.Sprintf("%s you specified the 'account' filter for Connect accounts but are not a Connect user, so the filter will not be applied.", color.Yellow("Warning")))
 	}
 
-	// Block until Ctrl+C is received
-	<-tailer.interruptCh
-
-	log.WithFields(log.Fields{
-		"prefix": "logs.Tailer.Run",
-	}).Debug("Ctrl+C received, cleaning up...")
+	// Block until context is done (i.e. Ctrl+C is pressed)
+	<-t.ctx.Done()
 
-	if tailer.webSocketClient != nil {
-		tailer.webSocketClient.Stop()
+	if t.webSocketClient != nil {
+		t.webSocketClient.Stop()
 	}
 
 	log.WithFields(log.Fields{
-		"prefix": "logs.Tailer.Run",
+		"prefix": "logtailing.Tailer.Run",
 	}).Debug("Bye!")
 
 	return nil
 }
 
-func (tailer *Tailer) processRequestLogEvent(msg websocket.IncomingMessage) {
+func (t *Tailer) createSession() (*stripeauth.StripeCLISession, error) {
+	var session *stripeauth.StripeCLISession
+
+	var err error
+
+	exitCh := make(chan struct{})
+
+	filters, err := jsonifyFilters(t.cfg.Filters)
+	if err != nil {
+		t.cfg.Log.Fatalf("Error while converting log filters to JSON encoding: %v", err)
+	}
+
+	go func() {
+		// 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)
+
+			if err == nil {
+				exitCh <- struct{}{}
+				return
+			}
+
+			select {
+			case <-t.ctx.Done():
+				exitCh <- struct{}{}
+				return
+			case <-time.After(1 * time.Second):
+			}
+		}
+
+		exitCh <- struct{}{}
+	}()
+	<-exitCh
+
+	return session, err
+}
+
+func (t *Tailer) processRequestLogEvent(msg websocket.IncomingMessage) {
 	if msg.RequestLogEvent == nil {
-		tailer.cfg.Log.Debug("WebSocket specified for request logs received non-request-logs event")
+		t.cfg.Log.Debug("WebSocket specified for request logs received non-request-logs event")
 		return
 	}
 
 	requestLogEvent := msg.RequestLogEvent
 
-	tailer.cfg.Log.WithFields(log.Fields{
-		"prefix":     "logs.Tailer.processRequestLogEvent",
+	t.cfg.Log.WithFields(log.Fields{
+		"prefix":     "logtailing.Tailer.processRequestLogEvent",
 		"webhook_id": requestLogEvent.RequestLogID,
 	}).Debugf("Processing request log event")
 
 	var payload EventPayload
 	if err := json.Unmarshal([]byte(requestLogEvent.EventPayload), &payload); err != nil {
-		tailer.cfg.Log.Debug("Received malformed payload: ", err)
+		t.cfg.Log.Debug("Received malformed payload: ", err)
 	}
 
 	// Don't show stripecli/sessions logs since they're generated by the CLI
 	if payload.URL == "/v1/stripecli/sessions" {
-		tailer.cfg.Log.Debug("Filtering out /v1/stripecli/sessions from logs")
+		t.cfg.Log.Debug("Filtering out /v1/stripecli/sessions from logs")
 		return
 	}
 
-	if tailer.cfg.OutputFormat == outputFormatJSON {
+	if t.cfg.OutputFormat == outputFormatJSON {
 		fmt.Println(ansi.ColorizeJSON(requestLogEvent.EventPayload, false, os.Stdout))
 		return
 	}
diff --git a/pkg/proxy/proxy.go b/pkg/proxy/proxy.go
index e7d91c9..0957f56 100644
--- a/pkg/proxy/proxy.go
+++ b/pkg/proxy/proxy.go
@@ -1,6 +1,7 @@
 package proxy
 
 import (
+	"context"
 	"crypto/tls"
 	"encoding/json"
 	"fmt"
@@ -84,6 +85,8 @@ type Proxy struct {
 	events map[string]bool
 
 	interruptCh chan os.Signal
+
+	ctx context.Context
 }
 
 // Run sets the websocket connection and starts the Goroutines to forward
@@ -91,30 +94,28 @@ type Proxy struct {
 func (p *Proxy) Run() error {
 	s := ansi.StartSpinner("Getting ready...", p.cfg.Log.Out)
 
-	// Intercept Ctrl+c so we can do some clean up
+	// Create a context that will be canceled when Ctrl+C is pressed
+	ctx, cancel := context.WithCancel(context.Background())
+	p.ctx = ctx
 	signal.Notify(p.interruptCh, os.Interrupt, syscall.SIGTERM)
 
-	var session *stripeauth.StripeCLISession
-
-	var err 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 = p.stripeAuthClient.Authorize(p.cfg.DeviceName, p.cfg.WebSocketFeature, nil)
-
-		if err == nil {
-			break
-		}
+	go func() {
+		log.WithFields(log.Fields{
+			"prefix": "proxy.Proxy.Run",
+		}).Debug("Ctrl+C received, cleaning up...")
 
-		time.Sleep(1 * time.Second)
-	}
+		<-p.interruptCh
+		cancel()
+	}()
 
+	// Create the CLI session
+	session, err := p.createSession()
 	if err != nil {
 		ansi.StopSpinner(s, "", p.cfg.Log.Out)
 		p.cfg.Log.Fatalf("Error while authenticating with Stripe: %v", err)
 	}
 
+	// Create and start the websocket client
 	p.webSocketClient = websocket.NewClient(
 		session.WebSocketURL,
 		session.WebSocketID,
@@ -124,18 +125,21 @@ func (p *Proxy) Run() error {
 			NoWSS:             p.cfg.NoWSS,
 			ReconnectInterval: time.Duration(session.ReconnectDelay) * time.Second,
 			EventHandler:      websocket.EventHandlerFunc(p.processWebhookEvent),
+			Ctx:               p.ctx,
 		},
 	)
 	go p.webSocketClient.Run()
 
-	ansi.StopSpinner(s, fmt.Sprintf("Ready! Your webhook signing secret is %s (^C to quit)", ansi.Bold(session.Secret)), p.cfg.Log.Out)
-
-	// Block until Ctrl+C is received
-	<-p.interruptCh
+	select {
+	case <-p.webSocketClient.Connected():
+		ansi.StopSpinner(s, fmt.Sprintf("Ready! Your webhook signing secret is %s (^C to quit)", ansi.Bold(session.Secret)), p.cfg.Log.Out)
+	case <-p.ctx.Done():
+		ansi.StopSpinner(s, "", p.cfg.Log.Out)
+		p.cfg.Log.Fatalf("Aborting")
+	}
 
-	log.WithFields(log.Fields{
-		"prefix": "proxy.Proxy.Run",
-	}).Debug("Ctrl+C received, cleaning up...")
+	// Block until context is done (i.e. Ctrl+C is pressed)
+	<-p.ctx.Done()
 
 	if p.webSocketClient != nil {
 		p.webSocketClient.Stop()
@@ -148,6 +152,39 @@ func (p *Proxy) Run() error {
 	return nil
 }
 
+func (p *Proxy) createSession() (*stripeauth.StripeCLISession, error) {
+	var session *stripeauth.StripeCLISession
+
+	var err error
+
+	exitCh := make(chan struct{})
+
+	go func() {
+		// 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 = p.stripeAuthClient.Authorize(p.ctx, p.cfg.DeviceName, p.cfg.WebSocketFeature, nil)
+
+			if err == nil {
+				exitCh <- struct{}{}
+				return
+			}
+
+			select {
+			case <-p.ctx.Done():
+				exitCh <- struct{}{}
+				return
+			case <-time.After(1 * time.Second):
+			}
+		}
+
+		exitCh <- struct{}{}
+	}()
+	<-exitCh
+
+	return session, err
+}
+
 func (p *Proxy) filterWebhookEvent(msg *websocket.WebhookEvent) bool {
 	if msg.Endpoint.APIVersion != nil && !p.cfg.UseLatestAPIVersion {
 		p.cfg.Log.WithFields(log.Fields{
diff --git a/pkg/requests/base.go b/pkg/requests/base.go
index 5e90a96..62888c9 100644
--- a/pkg/requests/base.go
+++ b/pkg/requests/base.go
@@ -2,6 +2,7 @@ package requests
 
 import (
 	"bufio"
+	"context"
 	"fmt"
 	"io/ioutil"
 	"net/http"
@@ -134,7 +135,7 @@ func (rb *Base) MakeRequest(apiKey, path string, params *RequestParameters, errO
 		rb.setVersionHeader(req, params)
 	}
 
-	resp, err := client.PerformRequest(rb.Method, path, data, configureReq)
+	resp, err := client.PerformRequest(context.TODO(), rb.Method, path, data, configureReq)
 	if err != nil {
 		return []byte{}, err
 	}
diff --git a/pkg/samples/samples.go b/pkg/samples/samples.go
index 7fd3151..34f082a 100644
--- a/pkg/samples/samples.go
+++ b/pkg/samples/samples.go
@@ -1,6 +1,7 @@
 package samples
 
 import (
+	"context"
 	"fmt"
 	"os"
 	"path/filepath"
@@ -396,7 +397,7 @@ func (s *Samples) ConfigureDotEnv(sampleLocation string) error {
 
 	authClient := stripeauth.NewClient(apiKey, nil)
 
-	authSession, err := authClient.Authorize(deviceName, "webhooks", nil)
+	authSession, err := authClient.Authorize(context.TODO(), deviceName, "webhooks", nil)
 	if err != nil {
 		return err
 	}
diff --git a/pkg/stripe/client.go b/pkg/stripe/client.go
index 3983f0f..4f0744c 100644
--- a/pkg/stripe/client.go
+++ b/pkg/stripe/client.go
@@ -39,7 +39,7 @@ type Client struct {
 }
 
 // PerformRequest sends a request to Stripe and returns the response.
-func (c *Client) PerformRequest(method, path string, params string, configure func(*http.Request)) (*http.Response, error) {
+func (c *Client) PerformRequest(ctx context.Context, method, path string, params string, configure func(*http.Request)) (*http.Response, error) {
 	url, err := url.Parse(path)
 	if err != nil {
 		return nil, err
@@ -76,6 +76,10 @@ func (c *Client) PerformRequest(method, path string, params string, configure fu
 		c.httpClient = newHTTPClient(c.Verbose, os.Getenv("STRIPE_CLI_UNIX_SOCKET"))
 	}
 
+	if ctx != nil {
+		req = req.WithContext(ctx)
+	}
+
 	resp, err := c.httpClient.Do(req)
 	if err != nil {
 		return nil, err
diff --git a/pkg/stripeauth/client.go b/pkg/stripeauth/client.go
index 5158548..b6701f4 100644
--- a/pkg/stripeauth/client.go
+++ b/pkg/stripeauth/client.go
@@ -1,6 +1,7 @@
 package stripeauth
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"io/ioutil"
@@ -36,7 +37,7 @@ type Client struct {
 }
 
 // Authorize sends a request to Stripe to initiate a new CLI session.
-func (c *Client) Authorize(deviceName string, websocketFeature string, filters *string) (*StripeCLISession, error) {
+func (c *Client) Authorize(ctx context.Context, deviceName string, websocketFeature string, filters *string) (*StripeCLISession, error) {
 	c.cfg.Log.WithFields(log.Fields{
 		"prefix": "stripeauth.client.Authorize",
 	}).Debug("Authenticating with Stripe...")
@@ -59,7 +60,7 @@ func (c *Client) Authorize(deviceName string, websocketFeature string, filters *
 		APIKey:  c.apiKey,
 	}
 
-	resp, err := client.PerformRequest(http.MethodPost, stripeCLISessionPath, form.Encode(), nil)
+	resp, err := client.PerformRequest(ctx, http.MethodPost, stripeCLISessionPath, form.Encode(), nil)
 	if err != nil {
 		return nil, err
 	}
diff --git a/pkg/websocket/client.go b/pkg/websocket/client.go
index 36bcd77..12d950f 100644
--- a/pkg/websocket/client.go
+++ b/pkg/websocket/client.go
@@ -1,6 +1,7 @@
 package websocket
 
 import (
+	"context"
 	"encoding/json"
 	"io/ioutil"
 	"net"
@@ -45,6 +46,8 @@ type Config struct {
 	WriteWait time.Duration
 
 	EventHandler EventHandler
+
+	Ctx context.Context
 }
 
 // EventHandler handles an event.
@@ -88,6 +91,21 @@ type Client struct {
 	wg            *sync.WaitGroup
 }
 
+// Connected returns a channel that's closed when the client has finished
+// establishing the websocket connection.
+func (c *Client) Connected() <-chan struct{} {
+	d := make(chan struct{})
+
+	go func() {
+		for !c.isConnected {
+			time.Sleep(100 * time.Millisecond)
+		}
+		close(d)
+	}()
+
+	return d
+}
+
 // Run starts listening for incoming webhook requests from Stripe.
 func (c *Client) Run() {
 	for {
@@ -100,9 +118,21 @@ func (c *Client) Run() {
 			c.cfg.Log.WithFields(log.Fields{
 				"prefix": "websocket.client.Run",
 			}).Debug("Failed to connect to Stripe. Retrying...")
-			time.Sleep(c.cfg.ConnectAttemptWait)
+
+			select {
+			case <-c.cfg.Ctx.Done():
+				return
+			case <-time.After(c.cfg.ConnectAttemptWait):
+			}
 		}
+
 		select {
+		case <-c.cfg.Ctx.Done():
+			close(c.send)
+			close(c.stopReadPump)
+			close(c.stopWritePump)
+
+			return
 		case <-c.done:
 			close(c.send)
 			close(c.stopReadPump)
@@ -164,7 +194,7 @@ func (c *Client) connect() bool {
 		"url":    url,
 	}).Debug("Dialing websocket")
 
-	conn, resp, err := c.cfg.Dialer.Dial(url, header)
+	conn, resp, err := c.cfg.Dialer.DialContext(c.cfg.Ctx, url, header)
 	if err != nil {
 		c.cfg.Log.WithFields(log.Fields{
 			"prefix": "websocket.Client.connect",
@@ -384,6 +414,10 @@ func NewClient(url string, webSocketID string, websocketAuthorizedFeature string
 		cfg.EventHandler = nullEventHandler
 	}
 
+	if cfg.Ctx == nil {
+		cfg.Ctx = context.Background()
+	}
+
 	return &Client{
 		URL:                        url,
 		WebSocketID:                webSocketID,


Test patch:
diff --git a/pkg/stripe/client_test.go b/pkg/stripe/client_test.go
index 7ab9a07..da2d45f 100644
--- a/pkg/stripe/client_test.go
+++ b/pkg/stripe/client_test.go
@@ -1,6 +1,7 @@
 package stripe
 
 import (
+	"context"
 	"io/ioutil"
 	"net/http"
 	"net/http/httptest"
@@ -30,7 +31,7 @@ func TestPerformRequest_ParamsEncoding_Delete(t *testing.T) {
 	params.Add("key_a", "value_a")
 	params.Add("key_b", "value_b")
 
-	resp, err := client.PerformRequest(http.MethodDelete, "/delete", params.Encode(), nil)
+	resp, err := client.PerformRequest(context.TODO(), http.MethodDelete, "/delete", params.Encode(), nil)
 	require.NoError(t, err)
 
 	defer resp.Body.Close()
@@ -56,7 +57,7 @@ func TestPerformRequest_ParamsEncoding_Get(t *testing.T) {
 	params.Add("key_a", "value_a")
 	params.Add("key_b", "value_b")
 
-	resp, err := client.PerformRequest(http.MethodGet, "/get", params.Encode(), nil)
+	resp, err := client.PerformRequest(context.TODO(), http.MethodGet, "/get", params.Encode(), nil)
 	require.NoError(t, err)
 
 	defer resp.Body.Close()
@@ -82,7 +83,7 @@ func TestPerformRequest_ParamsEncoding_Post(t *testing.T) {
 	params.Add("key_a", "value_a")
 	params.Add("key_b", "value_b")
 
-	resp, err := client.PerformRequest(http.MethodPost, "/post", params.Encode(), nil)
+	resp, err := client.PerformRequest(context.TODO(), http.MethodPost, "/post", params.Encode(), nil)
 	require.NoError(t, err)
 
 	defer resp.Body.Close()
@@ -100,7 +101,7 @@ func TestPerformRequest_ApiKey_Provided(t *testing.T) {
 		APIKey:  "sk_test_1234",
 	}
 
-	resp, err := client.PerformRequest(http.MethodGet, "/get", "", nil)
+	resp, err := client.PerformRequest(context.TODO(), http.MethodGet, "/get", "", nil)
 	require.NoError(t, err)
 
 	defer resp.Body.Close()
@@ -117,7 +118,7 @@ func TestPerformRequest_ApiKey_Omitted(t *testing.T) {
 		BaseURL: baseURL,
 	}
 
-	resp, err := client.PerformRequest(http.MethodGet, "/get", "", nil)
+	resp, err := client.PerformRequest(context.TODO(), http.MethodGet, "/get", "", nil)
 	require.NoError(t, err)
 
 	defer resp.Body.Close()
@@ -134,7 +135,7 @@ func TestPerformRequest_ConfigureFunc(t *testing.T) {
 		BaseURL: baseURL,
 	}
 
-	resp, err := client.PerformRequest(http.MethodGet, "/get", "", func(r *http.Request) {
+	resp, err := client.PerformRequest(context.TODO(), http.MethodGet, "/get", "", func(r *http.Request) {
 		r.Header.Add("Stripe-Version", "2019-07-10")
 	})
 	require.NoError(t, err)
diff --git a/pkg/stripeauth/client_test.go b/pkg/stripeauth/client_test.go
index 121c009..1d99955 100644
--- a/pkg/stripeauth/client_test.go
+++ b/pkg/stripeauth/client_test.go
@@ -1,6 +1,7 @@
 package stripeauth
 
 import (
+	"context"
 	"encoding/json"
 	"io/ioutil"
 	"net/http"
@@ -36,7 +37,7 @@ func TestAuthorize(t *testing.T) {
 	client := NewClient("sk_test_123", &Config{
 		APIBaseURL: ts.URL,
 	})
-	session, err := client.Authorize("my-device", "webhooks", nil)
+	session, err := client.Authorize(context.TODO(), "my-device", "webhooks", nil)
 	require.NoError(t, err)
 	require.Equal(t, "some-id", session.WebSocketID)
 	require.Equal(t, "wss://example.com/subscribe/acct_123", session.WebSocketURL)
@@ -54,7 +55,7 @@ func TestUserAgent(t *testing.T) {
 	client := NewClient("sk_test_123", &Config{
 		APIBaseURL: ts.URL,
 	})
-	client.Authorize("my-device", "webhooks", nil)
+	client.Authorize(context.TODO(), "my-device", "webhooks", nil)
 }
 
 func TestStripeClientUserAgent(t *testing.T) {
@@ -77,5 +78,5 @@ func TestStripeClientUserAgent(t *testing.T) {
 	client := NewClient("sk_test_123", &Config{
 		APIBaseURL: ts.URL,
 	})
-	client.Authorize("my-device", "webhooks", nil)
+	client.Authorize(context.TODO(), "my-device", "webhooks", nil)
 }