1package auth
2
3import (
4 "context"
5 "encoding/json"
6 "fmt"
7 "io"
8 "net/http"
9 "net/url"
10 "strings"
11 "time"
12)
13
14const (
15 xaiClientID = "b1a00492-073a-47ea-816f-4c329264a828"
16 xaiScope = "openid profile email offline_access grok-cli:access api:access"
17)
18
19var (
20 xaiDeviceCodeURL = "https://auth.x.ai/oauth2/device/code"
21 xaiTokenURL = "https://auth.x.ai/oauth2/token"
22)
23
24type XAIDeviceAuthorization struct {
25 DeviceCode string `json:"device_code"`
26 UserCode string `json:"user_code"`
27 VerificationURI string `json:"verification_uri"`
28 VerificationURIComplete string `json:"verification_uri_complete"`
29 ExpiresIn int `json:"expires_in"`
30 Interval int `json:"interval"`
31}
32
33func RequestXAIDeviceAuthorization(ctx context.Context) (XAIDeviceAuthorization, error) {
34 values := url.Values{"client_id": {xaiClientID}, "scope": {xaiScope}}
35 var out XAIDeviceAuthorization
36 if err := postXAIForm(ctx, xaiDeviceCodeURL, values, &out); err != nil {
37 return out, fmt.Errorf("xai device authorization: %w", err)
38 }
39 if out.DeviceCode == "" || out.UserCode == "" || out.VerificationURI == "" || out.ExpiresIn <= 0 {
40 return out, fmt.Errorf("xai device authorization: incomplete response")
41 }
42 verificationURL := out.VerificationURIComplete
43 if verificationURL == "" {
44 verificationURL = out.VerificationURI
45 }
46 u, err := url.Parse(verificationURL)
47 if err != nil || u.Scheme != "https" {
48 return out, fmt.Errorf("xai device authorization: untrusted verification uri")
49 }
50 if out.VerificationURIComplete == "" {
51 q := u.Query()
52 q.Set("user_code", out.UserCode)
53 u.RawQuery = q.Encode()
54 out.VerificationURIComplete = u.String()
55 }
56 return out, nil
57}
58
59func PollXAIDeviceToken(ctx context.Context, dev XAIDeviceAuthorization) (*OAuthToken, error) {
60 interval := time.Duration(dev.Interval) * time.Second
61 if interval <= 0 {
62 interval = 5 * time.Second
63 }
64 deadline := time.NewTimer(time.Duration(dev.ExpiresIn) * time.Second)
65 defer deadline.Stop()
66
67 for {
68 select {
69 case <-ctx.Done():
70 return nil, ctx.Err()
71 case <-deadline.C:
72 return nil, fmt.Errorf("xai device code expired")
73 case <-time.After(interval):
74 }
75
76 values := url.Values{
77 "grant_type": {"urn:ietf:params:oauth:grant-type:device_code"},
78 "client_id": {xaiClientID},
79 "device_code": {dev.DeviceCode},
80 }
81 tok, retry, nextInterval, err := requestXAIToken(ctx, values, "")
82 if err != nil {
83 return nil, err
84 }
85 if tok != nil {
86 return tok, nil
87 }
88 if !retry {
89 return nil, fmt.Errorf("xai device authorization failed")
90 }
91 if nextInterval > 0 {
92 interval = nextInterval
93 }
94 }
95}
96
97func RefreshXAIToken(ctx context.Context, refreshToken string) (*OAuthToken, error) {
98 values := url.Values{
99 "grant_type": {"refresh_token"},
100 "client_id": {xaiClientID},
101 "refresh_token": {refreshToken},
102 }
103 tok, _, _, err := requestXAIToken(ctx, values, refreshToken)
104 return tok, err
105}
106
107func requestXAIToken(ctx context.Context, values url.Values, previousRefresh string) (*OAuthToken, bool, time.Duration, error) {
108 var body struct {
109 AccessToken string `json:"access_token"`
110 RefreshToken string `json:"refresh_token"`
111 TokenType string `json:"token_type"`
112 Scope string `json:"scope"`
113 ExpiresIn int `json:"expires_in"`
114 Error string `json:"error"`
115 ErrorDescription string `json:"error_description"`
116 Interval int `json:"interval"`
117 }
118 status, err := postXAIFormStatus(ctx, xaiTokenURL, values, &body)
119 if err != nil {
120 return nil, false, 0, fmt.Errorf("xai token request: %w", err)
121 }
122 if status >= 200 && status < 300 && body.AccessToken != "" {
123 if body.RefreshToken == "" {
124 body.RefreshToken = previousRefresh
125 }
126 expires := body.ExpiresIn
127 if expires <= 0 {
128 expires = 3600
129 }
130 return &OAuthToken{
131 AccessToken: body.AccessToken, RefreshToken: body.RefreshToken,
132 TokenType: body.TokenType, Scope: body.Scope,
133 Expiry: time.Now().Add(time.Duration(expires)*time.Second - 5*time.Minute),
134 ClientID: xaiClientID,
135 }, false, 0, nil
136 }
137 switch body.Error {
138 case "authorization_pending":
139 return nil, true, 0, nil
140 case "slow_down":
141 return nil, true, time.Duration(body.Interval) * time.Second, nil
142 case "access_denied", "authorization_denied":
143 return nil, false, 0, fmt.Errorf("xai device authorization was denied")
144 case "expired_token":
145 return nil, false, 0, fmt.Errorf("xai device code expired")
146 default:
147 detail := strings.TrimSpace(strings.Join([]string{body.Error, body.ErrorDescription}, ": "))
148 if detail == "" {
149 detail = http.StatusText(status)
150 }
151 return nil, false, 0, fmt.Errorf("xai oauth http %d: %s", status, detail)
152 }
153}
154
155func postXAIForm(ctx context.Context, endpoint string, values url.Values, out any) error {
156 status, err := postXAIFormStatus(ctx, endpoint, values, out)
157 if err != nil {
158 return err
159 }
160 if status < 200 || status >= 300 {
161 return fmt.Errorf("http %d", status)
162 }
163 return nil
164}
165
166func postXAIFormStatus(ctx context.Context, endpoint string, values url.Values, out any) (int, error) {
167 req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, strings.NewReader(values.Encode()))
168 if err != nil {
169 return 0, err
170 }
171 req.Header.Set("Accept", "application/json")
172 req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
173 resp, err := http.DefaultClient.Do(req)
174 if err != nil {
175 return 0, err
176 }
177 defer resp.Body.Close()
178 body, err := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
179 if err != nil {
180 return resp.StatusCode, err
181 }
182 if err := json.Unmarshal(body, out); err != nil {
183 return resp.StatusCode, fmt.Errorf("invalid json: %w", err)
184 }
185 return resp.StatusCode, nil
186}