api: handle redirect during websocket upgrade (#18903)

When attempting a WebSocket connection upgrade the client may receive a
redirect request from the server, in which case the request should be
reattempted using the new address present in the `Location` header.

api/api.go

@@ -986,28 +986,42 @@ func (c *Client) websocket(endpoint string, q *QueryOptions) (*websocket.Conn, *
 	conn, resp, err := dialer.Dial(rhttp.URL.String(), rhttp.Header)
 
 	// check resp status code, as it's more informative than handshake error we get from ws library
-	if resp != nil && resp.StatusCode != http.StatusSwitchingProtocols {
-		var buf bytes.Buffer
+	if resp != nil {
+		switch resp.StatusCode {
+		case http.StatusSwitchingProtocols:
+			// Connection upgrade was successful.
 
-		if resp.Header.Get("Content-Encoding") == "gzip" {
-			greader, err := gzip.NewReader(resp.Body)
+		case http.StatusPermanentRedirect, http.StatusTemporaryRedirect, http.StatusMovedPermanently:
+			loc := resp.Header.Get("Location")
+			u, err := url.Parse(loc)
 			if err != nil {
-				return nil, nil, newUnexpectedResponseError(
-					fromStatusCode(resp.StatusCode),
-					withExpectedStatuses([]int{http.StatusSwitchingProtocols}),
-					withError(err))
+				return nil, nil, fmt.Errorf("invalid redirect location %q: %w", loc, err)
 			}
-			io.Copy(&buf, greader)
-		} else {
-			io.Copy(&buf, resp.Body)
-		}
-		resp.Body.Close()
+			return c.websocket(u.Path, q)
 
-		return nil, nil, newUnexpectedResponseError(
-			fromStatusCode(resp.StatusCode),
-			withExpectedStatuses([]int{http.StatusSwitchingProtocols}),
-			withBody(fmt.Sprint(buf.Bytes())),
-		)
+		default:
+			var buf bytes.Buffer
+
+			if resp.Header.Get("Content-Encoding") == "gzip" {
+				greader, err := gzip.NewReader(resp.Body)
+				if err != nil {
+					return nil, nil, newUnexpectedResponseError(
+						fromStatusCode(resp.StatusCode),
+						withExpectedStatuses([]int{http.StatusSwitchingProtocols}),
+						withError(err))
+				}
+				_, _ = io.Copy(&buf, greader)
+			} else {
+				_, _ = io.Copy(&buf, resp.Body)
+			}
+			_ = resp.Body.Close()
+
+			return nil, nil, newUnexpectedResponseError(
+				fromStatusCode(resp.StatusCode),
+				withExpectedStatuses([]int{http.StatusSwitchingProtocols}),
+				withBody(buf.String()),
+			)
+		}
 	}
 
 	return conn, resp, err