diff --git a/integration/helpers/instance.go b/integration/helpers/instance.go index fc53e5b7d46..a1da8640289 100644 --- a/integration/helpers/instance.go +++ b/integration/helpers/instance.go @@ -33,11 +33,11 @@ import ( "os" "path/filepath" "runtime/debug" + "strings" "testing" "time" "github.com/gorilla/websocket" - "github.com/gravitational/roundtrip" "github.com/gravitational/trace" "github.com/jonboulle/clockwork" log "github.com/sirupsen/logrus" @@ -1532,6 +1532,27 @@ func CreateWebSession(proxyHost, user, password string) (*web.CreateSessionRespo return csResp, resp.Cookies(), nil } +func makeAuthReqOverWS(ws *websocket.Conn, token string) error { + authReq, err := json.Marshal(struct { + Token string `json:"token"` + }{Token: token}) + if err != nil { + return trace.Wrap(err) + } + + if err := ws.WriteMessage(websocket.TextMessage, authReq); err != nil { + return trace.Wrap(err) + } + _, authRes, err := ws.ReadMessage() + if err != nil { + return trace.Wrap(err) + } + if !strings.Contains(string(authRes), `"status":"ok"`) { + return trace.AccessDenied("unexpected response") + } + return nil +} + // SSH establishes an SSH connection via the web api in the same manner that // the web UI does. The returned [web.TerminalStream] should be used as stdin/stdout // for the session. @@ -1539,7 +1560,7 @@ func (w *WebClient) SSH(termReq web.TerminalRequest) (*web.TerminalStream, error u := url.URL{ Host: w.i.Web, Scheme: client.WSS, - Path: fmt.Sprintf("/v1/webapi/sites/%v/connect", w.tc.SiteName), + Path: fmt.Sprintf("/v1/webapi/sites/%v/connect/ws", w.tc.SiteName), } data, err := json.Marshal(termReq) if err != nil { @@ -1548,7 +1569,6 @@ func (w *WebClient) SSH(termReq web.TerminalRequest) (*web.TerminalStream, error q := u.Query() q.Set("params", string(data)) - q.Set(roundtrip.AccessTokenQueryParam, w.token) u.RawQuery = q.Encode() header := http.Header{} @@ -1567,6 +1587,10 @@ func (w *WebClient) SSH(termReq web.TerminalRequest) (*web.TerminalStream, error return nil, trace.Wrap(err) } + if err := makeAuthReqOverWS(ws, w.token); err != nil { + return nil, trace.Wrap(err) + } + defer resp.Body.Close() return web.NewTerminalStream(context.Background(), web.TerminalStreamConfig{WS: ws}), nil } diff --git a/lib/benchmark/web.go b/lib/benchmark/web.go index e3010431478..1e153982085 100644 --- a/lib/benchmark/web.go +++ b/lib/benchmark/web.go @@ -195,7 +195,7 @@ func connectToHost(ctx context.Context, tc *client.TeleportClient, webSession *w u := url.URL{ Host: tc.WebProxyAddr, Scheme: client.WSS, - Path: fmt.Sprintf("/v1/webapi/sites/%v/connect", tc.SiteName), + Path: fmt.Sprintf("/v1/webapi/sites/%v/connect/ws", tc.SiteName), RawQuery: url.Values{ "params": []string{string(data)}, roundtrip.AccessTokenQueryParam: []string{webSession.getToken()}, diff --git a/lib/web/assistant_test.go b/lib/web/assistant_test.go index c67f94fd6ab..15d925dadb7 100644 --- a/lib/web/assistant_test.go +++ b/lib/web/assistant_test.go @@ -31,7 +31,6 @@ import ( "testing" "github.com/gorilla/websocket" - "github.com/gravitational/roundtrip" "github.com/gravitational/trace" "github.com/sashabaranov/go-openai" "github.com/stretchr/testify/assert" @@ -435,7 +434,7 @@ func (s *WebSuite) makeAssistant(_ *testing.T, pack *authPack, conversationID, a u := url.URL{ Host: s.url().Host, Scheme: client.WSS, - Path: fmt.Sprintf("/v1/webapi/sites/%s/assistant", currentSiteShortcut), + Path: fmt.Sprintf("/v1/webapi/sites/%s/assistant/ws", currentSiteShortcut), } q := u.Query() @@ -447,7 +446,6 @@ func (s *WebSuite) makeAssistant(_ *testing.T, pack *authPack, conversationID, a q.Set("action", action) } - q.Set(roundtrip.AccessTokenQueryParam, pack.session.Token) u.RawQuery = q.Encode() dialer := websocket.Dialer{} @@ -466,6 +464,10 @@ func (s *WebSuite) makeAssistant(_ *testing.T, pack *authPack, conversationID, a return nil, trace.Wrap(err) } + if err := makeAuthReqOverWS(ws, pack.session.Token); err != nil { + return nil, trace.Wrap(err) + } + err = resp.Body.Close() if err != nil { return nil, trace.Wrap(err) diff --git a/lib/web/command_test.go b/lib/web/command_test.go index 1f23a65f161..3dcc1823209 100644 --- a/lib/web/command_test.go +++ b/lib/web/command_test.go @@ -35,7 +35,6 @@ import ( "github.com/gogo/protobuf/proto" "github.com/google/uuid" "github.com/gorilla/websocket" - "github.com/gravitational/roundtrip" "github.com/gravitational/trace" "github.com/sashabaranov/go-openai" "github.com/sirupsen/logrus" @@ -273,7 +272,7 @@ func (s *WebSuite) makeCommand(t *testing.T, pack *authPack, conversationID uuid u := url.URL{ Host: s.url().Host, Scheme: client.WSS, - Path: fmt.Sprintf("/v1/webapi/command/%v/execute", currentSiteShortcut), + Path: fmt.Sprintf("/v1/webapi/command/%v/execute/ws", currentSiteShortcut), } data, err := json.Marshal(req) if err != nil { @@ -282,7 +281,6 @@ func (s *WebSuite) makeCommand(t *testing.T, pack *authPack, conversationID uuid q := u.Query() q.Set("params", string(data)) - q.Set(roundtrip.AccessTokenQueryParam, pack.session.Token) u.RawQuery = q.Encode() dialer := websocket.Dialer{} @@ -301,6 +299,10 @@ func (s *WebSuite) makeCommand(t *testing.T, pack *authPack, conversationID uuid return nil, nil, trace.Wrap(err) } + if err := makeAuthReqOverWS(ws, pack.session.Token); err != nil { + return nil, nil, trace.Wrap(err) + } + ty, raw, err := ws.ReadMessage() if err != nil { return nil, nil, trace.Wrap(err)