diff --git a/pkg/listen/proxy/proxy.go b/pkg/listen/proxy/proxy.go index 6d8367e2..84ee8450 100644 --- a/pkg/listen/proxy/proxy.go +++ b/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 diff --git a/pkg/websocket/client.go b/pkg/websocket/client.go index 675a7dff..1e235410 100644 --- a/pkg/websocket/client.go +++ b/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 diff --git a/pkg/websocket/client_test.go b/pkg/websocket/client_test.go new file mode 100644 index 00000000..b48a0325 --- /dev/null +++ b/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") + } +}