fix(listen): recover when the CLI session expires server-side

pkg/listen/proxy/proxy.go

@@ -50,7 +50,7 @@ type Config struct {
 	// Disable periodic health checks of the local server
 	NoHealthcheck bool
 	// Output mode: interactive, compact, quiet
-	Output string
+	Output   string
 	GuestURL string
 	// MaxConnections allows tuning the maximum concurrent connections per host.
 	// Default: 50 concurrent connections
@@ -141,6 +141,25 @@ func (p *Proxy) Run(parentCtx context.Context) error {
 		return fmt.Errorf("error while starting a new session")
 	}
 
+	// Session-recreation data sent on every websocket connect so the server
+	// can recreate the session if it has expired server-side while the CLI
+	// is still running.
+	var connectionIDs []string
+	for _, connection := range p.connections {
+		connectionIDs = append(connectionIDs, connection.ID)
+	}
+	var sessionFiltersJSON []byte
+	if p.cfg.Filters != nil {
+		var err error
+		sessionFiltersJSON, err = json.Marshal(p.cfg.Filters)
+		if err != nil {
+			log.WithFields(log.Fields{
+				"prefix": "proxy.Proxy.Run",
+			}).Debug("Failed to serialize session filters for reconnect headers: ", err)
+			sessionFiltersJSON = nil
+		}
+	}
+
 	// Main loop to keep attempting to connect to Hookdeck once
 	// we have created a session.
 	for canConnect() {
@@ -150,9 +169,11 @@ func (p *Proxy) Run(parentCtx context.Context) error {
 			p.cfg.Key,
 			p.cfg.ProjectID,
 			&websocket.Config{
-				Log:          p.cfg.Log,
-				NoWSS:        p.cfg.NoWSS,
-				EventHandler: websocket.EventHandlerFunc(p.processAttempt),
+				Log:                p.cfg.Log,
+				NoWSS:              p.cfg.NoWSS,
+				EventHandler:       websocket.EventHandlerFunc(p.processAttempt),
+				WebhookIDs:         connectionIDs,
+				SessionFiltersJSON: sessionFiltersJSON,
 			},
 		)
 
@@ -205,6 +226,23 @@ func (p *Proxy) Run(parentCtx context.Context) error {
 				p.renderer.Cleanup()
 				return fmt.Errorf("Could not connect. Terminating after %d failed attempts to establish a connection.", nAttempts)
 			}
+			// The server reported that the session no longer exists (e.g. it
+			// expired server-side). Reconnecting with the same websocket ID
+			// can never succeed, so create a fresh session before retrying.
+			if p.webSocketClient.SessionExpired() {
+				log.WithFields(log.Fields{
+					"prefix": "proxy.Proxy.Run",
+				}).Debug("Session expired server-side, creating a new session before reconnecting")
+
+				newSession, err := p.createSession(signalCtx)
+				if err != nil || newSession.Id == "" {
+					log.WithFields(log.Fields{
+						"prefix": "proxy.Proxy.Run",
+					}).Debug("Failed to create a replacement session, retrying with the existing session: ", err)
+				} else {
+					session = newSession
+				}
+			}
 		}
 
 		// Add backoff delay between all retry attempts

pkg/websocket/client.go

@@ -11,6 +11,7 @@ import (
 	"os"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	ws "github.com/gorilla/websocket"
@@ -41,6 +42,16 @@ type Config struct {
 	WriteWait time.Duration
 
 	EventHandler EventHandler
+
+	// WebhookIDs are the connection IDs the session listens to. They are sent
+	// on every connect via the X-Webhook-Ids header so the server can recreate
+	// the session if it has expired server-side.
+	WebhookIDs []string
+
+	// SessionFiltersJSON is the JSON-encoded session filters, sent on every
+	// connect via the X-Session-Filters header (base64-encoded). Nil when the
+	// session has no filters.
+	SessionFiltersJSON []byte
 }
 
 // EventHandler handles an event.
@@ -82,6 +93,12 @@ type Client struct {
 	done        chan struct{}
 	isConnected bool
 
+	// sessionExpired is set when the server indicates the session no longer
+	// exists (close code 4001, or an "Unknown WebSocket ID." rejection during
+	// the upgrade). The owner of the client should create a new session
+	// instead of reconnecting with the same websocket ID.
+	sessionExpired atomic.Bool
+
 	NotifyExpired chan struct{}
 	notifyClose   chan error
 	send          chan *OutgoingMessage
@@ -90,6 +107,13 @@ type Client struct {
 	wg            *sync.WaitGroup
 }
 
+// SessionExpired reports whether the server indicated that the session tied
+// to this client's WebSocketID no longer exists. When true, reconnecting with
+// the same WebSocketID cannot succeed; a new session must be created first.
+func (c *Client) SessionExpired() bool {
+	return c.sessionExpired.Load()
+}
+
 // Connected returns a channel that's closed when the client has finished
 // establishing the websocket connection.
 func (c *Client) Connected() <-chan struct{} {
@@ -123,6 +147,7 @@ func (c *Client) Run(ctx context.Context) {
 		}).Debug("Failed to connect to Hookdeck. Retrying...")
 
 		if err == ErrUnknownID {
+			c.sessionExpired.Store(true)
 			c.cfg.Log.WithFields(log.Fields{
 				"prefix": "websocket.client.Run",
 			}).Debug("Websocket session is expired.")
@@ -205,14 +230,20 @@ var unknownIDMessage string = "Unknown WebSocket ID."
 // ErrUnknownID can occur when the websocket session is expired or invalid
 var ErrUnknownID error = errors.New(unknownIDMessage)
 
+// closeSessionExpired is the close code the server sends when the session
+// tied to the Websocket-Id header no longer exists (SESSION_EXPIRED) or the
+// session-recreation headers could not be parsed (INVALID_SESSION_DATA).
+const closeSessionExpired = 4001
+
 func basicAuth(username, password string) string {
 	auth := username + ":" + password
 	return base64.StdEncoding.EncodeToString([]byte(auth))
 }
 
-// connect makes a single attempt to connect to the websocket URL. It returns
-// the success of the attempt.
-func (c *Client) connect(ctx context.Context) error {
+// connectHeaders builds the headers sent with the websocket upgrade request.
+// X-Webhook-Ids and X-Session-Filters let the server recreate the session if
+// it has expired server-side (e.g. Redis TTL) while the CLI is still running.
+func (c *Client) connectHeaders() http.Header {
 	header := http.Header{}
 	// Disable compression by requiring "identity"
 	header.Set("Accept-Encoding", "identity")
@@ -222,6 +253,23 @@ func (c *Client) connect(ctx context.Context) error {
 	header.Set("X-Team-Id", c.TeamID)
 	header.Set("Authorization", "Basic "+basicAuth(c.CLIKey, ""))
 
+	if len(c.cfg.WebhookIDs) > 0 {
+		header.Set("X-Webhook-Ids", strings.Join(c.cfg.WebhookIDs, ","))
+	}
+	// Base64 keeps non-ASCII filter values intact; the server decodes the
+	// header as base64-encoded JSON.
+	if len(c.cfg.SessionFiltersJSON) > 0 {
+		header.Set("X-Session-Filters", base64.StdEncoding.EncodeToString(c.cfg.SessionFiltersJSON))
+	}
+
+	return header
+}
+
+// connect makes a single attempt to connect to the websocket URL. It returns
+// the success of the attempt.
+func (c *Client) connect(ctx context.Context) error {
+	header := c.connectHeaders()
+
 	url := c.URL
 	if c.cfg.NoWSS && strings.HasPrefix(url, "wss") {
 		url = "ws" + strings.TrimPrefix(c.URL, "wss")
@@ -309,6 +357,12 @@ func (c *Client) readPump() {
 					"prefix": "websocket.Client.readPump",
 				}).Debug("stopReadPump")
 			default:
+				if ws.IsCloseError(err, closeSessionExpired) {
+					c.sessionExpired.Store(true)
+					c.cfg.Log.WithFields(log.Fields{
+						"prefix": "websocket.Client.readPump",
+					}).Debug("Server closed the connection because the session expired: ", err)
+				}
 				switch {
 				case !ws.IsCloseError(err):
 					// read errors do not prevent websocket reconnects in the CLI so we should

pkg/websocket/client_test.go

@@ -0,0 +1,124 @@
+package websocket
+
+import (
+	"context"
+	"encoding/base64"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+	"time"
+
+	ws "github.com/gorilla/websocket"
+)
+
+func TestConnectHeadersIncludeSessionRecreationData(t *testing.T) {
+	filtersJSON := []byte(`{"body":{"event_type":"PAYOUT_NORMAL_FAILED"}}`)
+
+	client := NewClient("wss://example.com", "cses_123", "key_123", "tm_123", &Config{
+		WebhookIDs:         []string{"web_1", "web_2"},
+		SessionFiltersJSON: filtersJSON,
+	})
+
+	header := client.connectHeaders()
+
+	if got := header.Get("Websocket-Id"); got != "cses_123" {
+		t.Errorf("Websocket-Id = %q, want %q", got, "cses_123")
+	}
+
+	if got := header.Get("X-Webhook-Ids"); got != "web_1,web_2" {
+		t.Errorf("X-Webhook-Ids = %q, want %q", got, "web_1,web_2")
+	}
+
+	encoded := header.Get("X-Session-Filters")
+	if encoded == "" {
+		t.Fatal("X-Session-Filters header is missing")
+	}
+	decoded, err := base64.StdEncoding.DecodeString(encoded)
+	if err != nil {
+		t.Fatalf("X-Session-Filters is not valid base64: %v", err)
+	}
+	if string(decoded) != string(filtersJSON) {
+		t.Errorf("X-Session-Filters decodes to %q, want %q", decoded, filtersJSON)
+	}
+}
+
+func TestConnectHeadersOmitEmptySessionRecreationData(t *testing.T) {
+	client := NewClient("wss://example.com", "cses_123", "key_123", "tm_123", &Config{})
+
+	header := client.connectHeaders()
+
+	if got := header.Get("X-Webhook-Ids"); got != "" {
+		t.Errorf("X-Webhook-Ids = %q, want empty", got)
+	}
+	if got := header.Get("X-Session-Filters"); got != "" {
+		t.Errorf("X-Session-Filters = %q, want empty", got)
+	}
+}
+
+// runClientAgainstServer connects a client to a test websocket server whose
+// handler is given the upgraded connection, then waits for the client to
+// report the connection as lost. It returns the client for assertions.
+func runClientAgainstServer(t *testing.T, handler func(conn *ws.Conn)) *Client {
+	t.Helper()
+
+	upgrader := ws.Upgrader{}
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		conn, err := upgrader.Upgrade(w, r, nil)
+		if err != nil {
+			t.Errorf("failed to upgrade connection: %v", err)
+			return
+		}
+		defer conn.Close()
+		handler(conn)
+	}))
+	defer server.Close()
+
+	url := "ws" + strings.TrimPrefix(server.URL, "http")
+	client := NewClient(url, "cses_123", "key_123", "tm_123", &Config{
+		WebhookIDs: []string{"web_1"},
+	})
+
+	go client.Run(context.Background())
+
+	select {
+	case <-client.NotifyExpired:
+	case <-time.After(5 * time.Second):
+		t.Fatal("timed out waiting for the client to report the connection as lost")
+	}
+
+	return client
+}
+
+func TestSessionExpiredCloseSetsSessionExpired(t *testing.T) {
+	client := runClientAgainstServer(t, func(conn *ws.Conn) {
+		deadline := time.Now().Add(time.Second)
+		if err := conn.WriteControl(ws.CloseMessage, ws.FormatCloseMessage(closeSessionExpired, "SESSION_EXPIRED"), deadline); err != nil {
+			t.Errorf("failed to send close message: %v", err)
+			return
+		}
+		// Wait for the client to echo the close frame before tearing down.
+		conn.SetReadDeadline(time.Now().Add(time.Second))
+		conn.ReadMessage()
+	})
+
+	if !client.SessionExpired() {
+		t.Error("SessionExpired() = false after a 4001 close, want true")
+	}
+}
+
+func TestNormalCloseDoesNotSetSessionExpired(t *testing.T) {
+	client := runClientAgainstServer(t, func(conn *ws.Conn) {
+		deadline := time.Now().Add(time.Second)
+		if err := conn.WriteControl(ws.CloseMessage, ws.FormatCloseMessage(ws.CloseGoingAway, "SERVER_SHUTDOWN"), deadline); err != nil {
+			t.Errorf("failed to send close message: %v", err)
+			return
+		}
+		conn.SetReadDeadline(time.Now().Add(time.Second))
+		conn.ReadMessage()
+	})
+
+	if client.SessionExpired() {
+		t.Error("SessionExpired() = true after a 1001 close, want false")
+	}
+}