mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-19 16:17:56 +00:00
Merge pull request #7887 from whhphd/fix/grok-sso-consent
fix(grok): include consent token and origin in SSO device approval
This commit is contained in:
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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&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&b+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) {
|
||||
|
||||
Reference in New Issue
Block a user