github/sub2api
4
0
Fork 0
mirror of https://github.com/Wei-Shaw/sub2api.git synced 2026-10-19 16:17:56 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(grok): include consent token and origin in SSO device approval

Browse Source
This commit is contained in:
whhphd
2026-10-06 22:53:10 +08:00
parent b8dece9000
commit 54810f8c82
2 changed files with 80 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/internal/pkg/xai/sso_device.go
+54 -3
View File
@@ -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)
}
backend/internal/pkg/xai/sso_device_test.go
+26 -1
View File
@@ -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, `<html>consent</html>`), nil
return ssoDeviceResponse(http.StatusOK, nil, `<html><form><input type="hidden" name="consent_token" value="consent&amp;token"/></form></html>`), 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", `<html>consent</html>`, ""},
{"encoded hidden value", `<input value='a&amp;b&#43;c' name='consent_token' type='hidden'>`, "a&b+c"},
{"case insensitive HTML", `<INPUT TYPE="HIDDEN" NAME="consent_token" VALUE="token"/>`, "token"},
{"ignore script and visible fields", `<script>"<input type='hidden' name='consent_token' value='fake'>"</script><input name='consent_token' value='visible'><input type='hidden' name='user_code' value='user'><input type='hidden' name='consent_token' value='real'>`, "real"},
} {
t.Run(tc.name, func(t *testing.T) {
require.Equal(t, tc.want, ssoConsentToken([]byte(tc.body)))
})
}
}
func TestNormalizeSSOTokenAcceptsCookieHeader(t *testing.T) {