diff --git a/backend/internal/pkg/xai/sso_device.go b/backend/internal/pkg/xai/sso_device.go index 825616b93c..5a06d24f90 100644 --- a/backend/internal/pkg/xai/sso_device.go +++ b/backend/internal/pkg/xai/sso_device.go @@ -1,6 +1,7 @@ package xai import ( + "bytes" "context" "encoding/base64" "encoding/json" @@ -14,6 +15,8 @@ import ( "strconv" "strings" "time" + + "golang.org/x/net/html" ) const ( @@ -147,7 +150,7 @@ func (f *ssoDeviceFlow) convert(ctx context.Context) (*TokenResponse, error) { return nil, fmt.Errorf("open xAI device verification page: %w", SSOHTTPError{Status: status}) } - status, finalURL, _, err = f.do(ctx, http.MethodPost, SSOVerifyURL, url.Values{"user_code": {device.UserCode}}) + status, finalURL, body, err = f.do(ctx, http.MethodPost, SSOVerifyURL, url.Values{"user_code": {device.UserCode}}) if err != nil { return nil, err } @@ -158,11 +161,19 @@ func (f *ssoDeviceFlow) convert(ctx context.Context) (*TokenResponse, error) { return nil, errors.New("xAI device verification did not reach consent page") } - status, finalURL, _, err = f.do(ctx, http.MethodPost, SSOApproveURL, url.Values{ + approval := url.Values{ "user_code": {device.UserCode}, "action": {"allow"}, "principal_type": {"User"}, "principal_id": {""}, + } + if consentToken := ssoConsentToken(body); consentToken != "" { + approval.Set("consent_token", consentToken) + } + consentURL, _ := url.Parse(finalURL) // f.do only returns validated xAI URLs. + status, finalURL, _, err = f.do(ctx, http.MethodPost, SSOApproveURL, approval, http.Header{ + "Origin": {consentURL.Scheme + "://" + consentURL.Host}, + "Referer": {finalURL}, }) if err != nil { return nil, err @@ -177,6 +188,37 @@ func (f *ssoDeviceFlow) convert(ctx context.Context) (*TokenResponse, error) { return f.pollToken(ctx, device.DeviceCode, time.Duration(device.Interval)*time.Second, time.Duration(device.ExpiresIn)*time.Second) } +// xAI binds device approval to a hidden token on the consent page. Keep it in +// this flow only; older consent pages without the field remain supported. +func ssoConsentToken(body []byte) string { + tokenizer := html.NewTokenizer(bytes.NewReader(body)) + for { + switch tokenizer.Next() { + case html.ErrorToken: + return "" + case html.StartTagToken, html.SelfClosingTagToken: + token := tokenizer.Token() + if token.Data != "input" { + continue + } + var name, value, inputType string + for _, attr := range token.Attr { + switch attr.Key { + case "name": + name = attr.Val + case "value": + value = attr.Val + case "type": + inputType = attr.Val + } + } + if name == "consent_token" && strings.EqualFold(inputType, "hidden") { + return value + } + } + } +} + func (f *ssoDeviceFlow) pollToken(ctx context.Context, deviceCode string, interval, expiresIn time.Duration) (*TokenResponse, error) { if interval < time.Second { interval = time.Second @@ -241,7 +283,7 @@ func (f *ssoDeviceFlow) pollToken(ctx context.Context, deviceCode string, interv return nil, errors.New("xAI device flow token polling timed out") } -func (f *ssoDeviceFlow) do(ctx context.Context, method, endpoint string, form url.Values) (int, string, []byte, error) { +func (f *ssoDeviceFlow) do(ctx context.Context, method, endpoint string, form url.Values, headers ...http.Header) (int, string, []byte, error) { if !safeXAIAuthURL(endpoint) { return 0, "", nil, errors.New("xAI OAuth URL is not trusted") } @@ -260,6 +302,15 @@ func (f *ssoDeviceFlow) do(ctx context.Context, method, endpoint string, form ur request.Header.Set("Accept", "application/json, text/html;q=0.9, */*;q=0.8") request.Header.Set("Accept-Language", "zh-CN,zh;q=0.9,en;q=0.8") request.Header.Set("User-Agent", f.userAgent) + if redirects == 0 { + for _, header := range headers { + for name, values := range header { + for _, value := range values { + request.Header.Add(name, value) + } + } + } + } if cookie := f.cookieHeader(request.URL); cookie != "" { request.Header.Set("Cookie", cookie) } diff --git a/backend/internal/pkg/xai/sso_device_test.go b/backend/internal/pkg/xai/sso_device_test.go index cd3bae3aa3..4b59dd4a57 100644 --- a/backend/internal/pkg/xai/sso_device_test.go +++ b/backend/internal/pkg/xai/sso_device_test.go @@ -23,6 +23,10 @@ type ssoDeviceFakeClient struct { func (c *ssoDeviceFakeClient) Do(req *http.Request) (*http.Response, error) { c.cookieHeaders = append(c.cookieHeaders, req.Header.Get("Cookie")) + if req.URL.String() != SSOApproveURL { + require.Empty(c.t, req.Header.Get("Origin")) + require.Empty(c.t, req.Header.Get("Referer")) + } switch req.URL.String() { case SSOAccountsURL: require.Equal(c.t, http.MethodGet, req.Method) @@ -43,13 +47,16 @@ func (c *ssoDeviceFakeClient) Do(req *http.Request) (*http.Response, error) { return ssoDeviceResponse(http.StatusFound, http.Header{"Location": {"/oauth2/device/consent"}}, ``), nil case "https://auth.x.ai/oauth2/device/consent": require.Equal(c.t, http.MethodGet, req.Method) - return ssoDeviceResponse(http.StatusOK, nil, `consent`), nil + return ssoDeviceResponse(http.StatusOK, nil, `
`), nil case SSOApproveURL: require.Equal(c.t, http.MethodPost, req.Method) + require.Equal(c.t, "https://auth.x.ai", req.Header.Get("Origin")) + require.Equal(c.t, "https://auth.x.ai/oauth2/device/consent", req.Header.Get("Referer")) values := readSSODeviceForm(c.t, req) require.Equal(c.t, "USER-1", values.Get("user_code")) require.Equal(c.t, "allow", values.Get("action")) require.Equal(c.t, "User", values.Get("principal_type")) + require.Equal(c.t, "consent&token", values.Get("consent_token")) return ssoDeviceResponse(http.StatusSeeOther, http.Header{"Location": {"/oauth2/device/done"}}, ``), nil case "https://auth.x.ai/oauth2/device/done": require.Equal(c.t, http.MethodGet, req.Method) @@ -87,6 +94,24 @@ func TestConvertSSOToBuildCompletesDeviceFlow(t *testing.T) { require.Contains(t, client.cookieHeaders[0], "sso-rw=sso-token") require.Contains(t, client.cookieHeaders[len(client.cookieHeaders)-1], "session=web-session") require.Contains(t, client.cookieHeaders[len(client.cookieHeaders)-1], "csrf=csrf-token") + for _, cookie := range client.cookieHeaders { + require.NotContains(t, cookie, "consent&token") + } +} + +func TestSSOConsentToken(t *testing.T) { + for _, tc := range []struct { + name, body, want string + }{ + {"legacy page", `consent`, ""}, + {"encoded hidden value", ``, "a&b+c"}, + {"case insensitive HTML", ``, "token"}, + {"ignore script and visible fields", ``, "real"}, + } { + t.Run(tc.name, func(t *testing.T) { + require.Equal(t, tc.want, ssoConsentToken([]byte(tc.body))) + }) + } } func TestNormalizeSSOTokenAcceptsCookieHeader(t *testing.T) {