diff --git a/backend/internal/pkg/xai/sso_device.go b/backend/internal/pkg/xai/sso_device.go index e533e394d2..825616b93c 100644 --- a/backend/internal/pkg/xai/sso_device.go +++ b/backend/internal/pkg/xai/sso_device.go @@ -8,6 +8,7 @@ import ( "fmt" "io" "net/http" + "net/http/cookiejar" "net/url" "sort" "strconv" @@ -25,6 +26,7 @@ const ( SSOConversionTimeout = 90 * time.Second ssoMaxAuthBody = 2 << 20 + ssoMaxTokenLength = 16 << 10 ssoDefaultUA = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36" ssoDefaultTokenTTL = 6 * time.Hour ) @@ -51,7 +53,7 @@ type SSODeviceOptions struct { type ssoDeviceFlow struct { client SSODeviceHTTPClient userAgent string - cookies map[string]string + cookieJar http.CookieJar sleep func(context.Context, time.Duration) error } @@ -80,11 +82,16 @@ func ConvertSSOToBuild(ctx context.Context, ssoToken string, opts *SSODeviceOpti if sleep == nil { sleep = sleepContext } + jar, err := cookiejar.New(nil) + if err != nil { + return nil, err + } + seedSSOCookies(jar, ssoToken) flow := &ssoDeviceFlow{ client: client, userAgent: userAgent, - cookies: map[string]string{"sso": ssoToken, "sso-rw": ssoToken}, + cookieJar: jar, sleep: sleep, } return flow.convert(ctx) @@ -253,7 +260,7 @@ 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 cookie := f.cookieHeader(); cookie != "" { + if cookie := f.cookieHeader(request.URL); cookie != "" { request.Header.Set("Cookie", cookie) } if currentForm != nil { @@ -264,7 +271,7 @@ func (f *ssoDeviceFlow) do(ctx context.Context, method, endpoint string, form ur if err != nil { return 0, currentURL, nil, err } - f.captureCookies(response) + f.captureCookies(request.URL, response) data, readErr := io.ReadAll(io.LimitReader(response.Body, ssoMaxAuthBody+1)) _ = response.Body.Close() if readErr != nil { @@ -298,30 +305,49 @@ func (f *ssoDeviceFlow) do(ctx context.Context, method, endpoint string, form ur return 0, currentURL, nil, errors.New("xAI OAuth redirected too many times") } -func (f *ssoDeviceFlow) captureCookies(response *http.Response) { +func seedSSOCookies(jar http.CookieJar, token string) { + if jar == nil { + return + } + for _, rawURL := range []string{SSOAccountsURL, OAuthIssuer + "/"} { + target, err := url.Parse(rawURL) + if err != nil { + continue + } + jar.SetCookies(target, []*http.Cookie{ + {Name: "sso", Value: token, Path: "/", Secure: true, HttpOnly: true}, + {Name: "sso-rw", Value: token, Path: "/", Secure: true, HttpOnly: true}, + }) + } +} + +func (f *ssoDeviceFlow) captureCookies(requestURL *url.URL, response *http.Response) { + if f == nil || f.cookieJar == nil || requestURL == nil || response == nil { + return + } + cookies := make([]*http.Cookie, 0) for _, cookie := range response.Cookies() { name := strings.TrimSpace(cookie.Name) value := strings.TrimSpace(cookie.Value) if name == "" || len(name) > 128 || len(value) > 16384 || strings.ContainsAny(name+value, "\r\n\x00") { continue } - if cookie.MaxAge < 0 { - delete(f.cookies, name) - continue - } - f.cookies[name] = value + cookie.Name = name + cookie.Value = value + cookies = append(cookies, cookie) } + f.cookieJar.SetCookies(requestURL, cookies) } -func (f *ssoDeviceFlow) cookieHeader() string { - keys := make([]string, 0, len(f.cookies)) - for key := range f.cookies { - keys = append(keys, key) +func (f *ssoDeviceFlow) cookieHeader(requestURL *url.URL) string { + if f == nil || f.cookieJar == nil || requestURL == nil { + return "" } - sort.Strings(keys) - parts := make([]string, 0, len(keys)) - for _, key := range keys { - parts = append(parts, key+"="+f.cookies[key]) + cookies := f.cookieJar.Cookies(requestURL) + sort.Slice(cookies, func(i, j int) bool { return cookies[i].Name < cookies[j].Name }) + parts := make([]string, 0, len(cookies)) + for _, cookie := range cookies { + parts = append(parts, cookie.Name+"="+cookie.Value) } return strings.Join(parts, "; ") } @@ -363,7 +389,11 @@ func NormalizeSSOToken(value string) string { } func sanitizeSSOToken(value string) string { - return strings.NewReplacer("\r", "", "\n", "", "\x00", "").Replace(strings.TrimSpace(value)) + value = strings.NewReplacer("\r", "", "\n", "", "\x00", "").Replace(strings.TrimSpace(value)) + if len(value) > ssoMaxTokenLength { + return "" + } + return value } func DecodeJWTClaims(token string) map[string]any { diff --git a/backend/internal/pkg/xai/sso_device_test.go b/backend/internal/pkg/xai/sso_device_test.go index 27dff15f4d..cd3bae3aa3 100644 --- a/backend/internal/pkg/xai/sso_device_test.go +++ b/backend/internal/pkg/xai/sso_device_test.go @@ -6,6 +6,7 @@ import ( "context" "io" "net/http" + "net/http/cookiejar" "net/url" "strings" "testing" @@ -25,7 +26,7 @@ func (c *ssoDeviceFakeClient) Do(req *http.Request) (*http.Response, error) { switch req.URL.String() { case SSOAccountsURL: require.Equal(c.t, http.MethodGet, req.Method) - return ssoDeviceResponse(http.StatusOK, http.Header{"Set-Cookie": {"session=web-session; Path=/"}}, `{}`), nil + return ssoDeviceResponse(http.StatusOK, http.Header{"Set-Cookie": {"session=web-session; Domain=x.ai; Path=/"}}, `{}`), nil case SSODeviceURL: require.Equal(c.t, http.MethodPost, req.Method) values := readSSODeviceForm(c.t, req) @@ -92,6 +93,33 @@ func TestNormalizeSSOTokenAcceptsCookieHeader(t *testing.T) { require.Equal(t, "token-1", NormalizeSSOToken("Cookie: foo=bar; sso=token-1; sso-rw=token-2")) require.Equal(t, "token-2", NormalizeSSOToken("sso-rw=token-2; foo=bar")) require.Equal(t, "raw-token", NormalizeSSOToken(" raw-token ; ignored=1")) + require.Empty(t, NormalizeSSOToken(strings.Repeat("x", ssoMaxTokenLength+1))) +} + +func TestSSODeviceCookieJarHonorsDomainAndPath(t *testing.T) { + jar, err := cookiejar.New(nil) + require.NoError(t, err) + flow := &ssoDeviceFlow{cookieJar: jar} + accountsURL, err := url.Parse("https://accounts.x.ai/") + require.NoError(t, err) + authURL, err := url.Parse("https://auth.x.ai/oauth2/device/verify") + require.NoError(t, err) + + flow.captureCookies(accountsURL, ssoDeviceResponse(http.StatusOK, http.Header{"Set-Cookie": { + "host-only=accounts; Path=/", + "shared=all-xai; Domain=x.ai; Path=/", + "narrow=oauth-only; Domain=x.ai; Path=/oauth2", + }}, "")) + + authCookies := flow.cookieHeader(authURL) + require.NotContains(t, authCookies, "host-only=accounts") + require.Contains(t, authCookies, "shared=all-xai") + require.Contains(t, authCookies, "narrow=oauth-only") + + accountsCookies := flow.cookieHeader(accountsURL) + require.Contains(t, accountsCookies, "host-only=accounts") + require.Contains(t, accountsCookies, "shared=all-xai") + require.NotContains(t, accountsCookies, "narrow=oauth-only") } func ssoDeviceResponse(status int, header http.Header, body string) *http.Response {