openai/openai-go
Publicmirrored from https://github.com/openai/openai-goAvailable
auth/workloadidentity.go
260lines · modecode
| 1 | package auth |
| 2 | |
| 3 | import ( |
| 4 | "bytes" |
| 5 | "context" |
| 6 | "encoding/json" |
| 7 | "fmt" |
| 8 | "io" |
| 9 | "net/http" |
| 10 | "sync" |
| 11 | "time" |
| 12 | |
| 13 | "github.com/openai/openai-go/v3/shared" |
| 14 | ) |
| 15 | |
| 16 | const ( |
| 17 | TokenExchangeGrantType = "urn:ietf:params:oauth:grant-type:token-exchange" |
| 18 | JWTTokenType = "urn:ietf:params:oauth:token-type:jwt" |
| 19 | IDTokenType = "urn:ietf:params:oauth:token-type:id_token" |
| 20 | DefaultTokenExpiry = 60 * time.Minute |
| 21 | DefaultRefreshBuffer = 20 * time.Minute |
| 22 | TokenExchangeURL = "https://auth.openai.com/oauth/token" |
| 23 | ) |
| 24 | |
| 25 | type WorkloadIdentityAuth struct { |
| 26 | config WorkloadIdentity |
| 27 | |
| 28 | // Protects cachedToken, tokenExpiry, and refreshInFlight |
| 29 | mu sync.Mutex |
| 30 | cachedToken string |
| 31 | tokenExpiry time.Time |
| 32 | refreshInFlight *tokenRefreshState |
| 33 | } |
| 34 | |
| 35 | type tokenRefreshResult struct { |
| 36 | token string |
| 37 | err error |
| 38 | } |
| 39 | |
| 40 | // Coordinates concurrent access to a single in-flight refresh operation |
| 41 | // done channel signals completion to all waiting goroutines |
| 42 | type tokenRefreshState struct { |
| 43 | done chan struct{} |
| 44 | result tokenRefreshResult |
| 45 | } |
| 46 | |
| 47 | type tokenExchangeRequest struct { |
| 48 | GrantType string `json:"grant_type"` |
| 49 | ClientID string `json:"client_id,omitempty"` |
| 50 | SubjectToken string `json:"subject_token"` |
| 51 | SubjectTokenType string `json:"subject_token_type"` |
| 52 | IdentityProviderID string `json:"identity_provider_id"` |
| 53 | ServiceAccountID string `json:"service_account_id"` |
| 54 | } |
| 55 | |
| 56 | func NewWorkloadIdentityAuth(config WorkloadIdentity) (*WorkloadIdentityAuth, error) { |
| 57 | if config.IdentityProviderID == "" { |
| 58 | return nil, fmt.Errorf("WorkloadIdentity: IdentityProviderID is required") |
| 59 | } |
| 60 | if config.ServiceAccountID == "" { |
| 61 | return nil, fmt.Errorf("WorkloadIdentity: ServiceAccountID is required") |
| 62 | } |
| 63 | if config.Provider == nil { |
| 64 | return nil, fmt.Errorf("WorkloadIdentity: Provider is required") |
| 65 | } |
| 66 | if config.RefreshBufferSeconds < 0 { |
| 67 | return nil, fmt.Errorf("WorkloadIdentity: RefreshBufferSeconds must be non-negative") |
| 68 | } |
| 69 | return &WorkloadIdentityAuth{ |
| 70 | config: config, |
| 71 | }, nil |
| 72 | } |
| 73 | |
| 74 | func (w *WorkloadIdentityAuth) GetToken(ctx context.Context, httpClient HTTPDoer) (string, error) { |
| 75 | if httpClient == nil { |
| 76 | httpClient = http.DefaultClient |
| 77 | } |
| 78 | |
| 79 | // Lock for entire decision: check cache, decide refresh strategy, potentially start background refresh |
| 80 | w.mu.Lock() |
| 81 | |
| 82 | if w.cachedToken == "" { |
| 83 | return w.handleLockedRefresh(ctx, httpClient) |
| 84 | } |
| 85 | |
| 86 | now := time.Now() |
| 87 | if now.After(w.tokenExpiry) { |
| 88 | return w.handleLockedRefresh(ctx, httpClient) |
| 89 | } |
| 90 | |
| 91 | refreshBuffer := w.config.RefreshBufferSeconds |
| 92 | if refreshBuffer == 0 { |
| 93 | refreshBuffer = int(DefaultRefreshBuffer / time.Second) |
| 94 | } |
| 95 | refreshTime := w.tokenExpiry.Add(-time.Duration(refreshBuffer) * time.Second) |
| 96 | |
| 97 | // Proactive background refresh: start if within refresh window and no refresh active |
| 98 | if now.After(refreshTime) && w.refreshInFlight == nil { |
| 99 | state := w.beginRefreshLocked() |
| 100 | // Background goroutine with independent context, lock released before spawn |
| 101 | go func() { |
| 102 | refreshCtx, cancel := context.WithTimeout(context.Background(), 30*time.Second) |
| 103 | defer cancel() |
| 104 | token, err := w.refreshToken(refreshCtx, httpClient) |
| 105 | w.finishRefresh(state, token, err) |
| 106 | }() |
| 107 | } |
| 108 | |
| 109 | token := w.cachedToken |
| 110 | w.mu.Unlock() |
| 111 | return token, nil |
| 112 | } |
| 113 | |
| 114 | // Single-flight pattern: ensures only one refresh runs, others wait for result |
| 115 | func (w *WorkloadIdentityAuth) handleLockedRefresh(ctx context.Context, httpClient HTTPDoer) (string, error) { |
| 116 | if w.refreshInFlight == nil { |
| 117 | // No refresh running: start foreground refresh, unlock before blocking operation |
| 118 | state := w.beginRefreshLocked() |
| 119 | w.mu.Unlock() |
| 120 | return w.completeForegroundRefresh(ctx, state, httpClient) |
| 121 | } |
| 122 | |
| 123 | // Refresh already running: unlock and wait for its completion |
| 124 | state := w.refreshInFlight |
| 125 | w.mu.Unlock() |
| 126 | return w.waitForRefresh(ctx, state) |
| 127 | } |
| 128 | |
| 129 | func (w *WorkloadIdentityAuth) invalidateToken() { |
| 130 | w.mu.Lock() |
| 131 | defer w.mu.Unlock() |
| 132 | w.cachedToken = "" |
| 133 | w.tokenExpiry = time.Time{} |
| 134 | } |
| 135 | |
| 136 | func (w *WorkloadIdentityAuth) beginRefreshLocked() *tokenRefreshState { |
| 137 | w.refreshInFlight = &tokenRefreshState{done: make(chan struct{})} |
| 138 | return w.refreshInFlight |
| 139 | } |
| 140 | |
| 141 | func (w *WorkloadIdentityAuth) completeForegroundRefresh(ctx context.Context, state *tokenRefreshState, httpClient HTTPDoer) (string, error) { |
| 142 | token, err := w.refreshToken(ctx, httpClient) |
| 143 | w.finishRefresh(state, token, err) |
| 144 | return token, err |
| 145 | } |
| 146 | |
| 147 | // Atomically publishes refresh result and signals all waiting goroutines via channel close |
| 148 | func (w *WorkloadIdentityAuth) finishRefresh(state *tokenRefreshState, token string, err error) { |
| 149 | w.mu.Lock() |
| 150 | defer w.mu.Unlock() |
| 151 | if w.refreshInFlight != state { |
| 152 | return |
| 153 | } |
| 154 | state.result = tokenRefreshResult{token: token, err: err} |
| 155 | close(state.done) // Broadcasts completion to all waiters |
| 156 | w.refreshInFlight = nil |
| 157 | } |
| 158 | |
| 159 | // Blocks until refresh completes or context is canceled |
| 160 | func (w *WorkloadIdentityAuth) waitForRefresh(ctx context.Context, state *tokenRefreshState) (string, error) { |
| 161 | select { |
| 162 | case <-state.done: // Refresh completed |
| 163 | return state.result.token, state.result.err |
| 164 | case <-ctx.Done(): // Caller context canceled |
| 165 | return "", ctx.Err() |
| 166 | } |
| 167 | } |
| 168 | |
| 169 | func (w *WorkloadIdentityAuth) refreshToken(ctx context.Context, httpClient HTTPDoer) (string, error) { |
| 170 | if httpClient == nil { |
| 171 | httpClient = http.DefaultClient |
| 172 | } |
| 173 | |
| 174 | subjectToken, err := w.config.Provider.GetToken(ctx, httpClient) |
| 175 | if err != nil { |
| 176 | return "", err |
| 177 | } |
| 178 | |
| 179 | subjectTokenType := w.config.Provider.TokenType() |
| 180 | var subjectTokenTypeURN string |
| 181 | switch subjectTokenType { |
| 182 | case SubjectTokenTypeJWT: |
| 183 | subjectTokenTypeURN = JWTTokenType |
| 184 | case SubjectTokenTypeID: |
| 185 | subjectTokenTypeURN = IDTokenType |
| 186 | default: |
| 187 | return "", fmt.Errorf("unsupported subject token type %q", subjectTokenType) |
| 188 | } |
| 189 | |
| 190 | requestBody := tokenExchangeRequest{ |
| 191 | GrantType: TokenExchangeGrantType, |
| 192 | ClientID: w.config.ClientID, |
| 193 | SubjectToken: subjectToken, |
| 194 | SubjectTokenType: subjectTokenTypeURN, |
| 195 | IdentityProviderID: w.config.IdentityProviderID, |
| 196 | ServiceAccountID: w.config.ServiceAccountID, |
| 197 | } |
| 198 | |
| 199 | jsonBody, err := json.Marshal(requestBody) |
| 200 | if err != nil { |
| 201 | return "", fmt.Errorf("failed to marshal token exchange request: %w", err) |
| 202 | } |
| 203 | |
| 204 | req, err := http.NewRequestWithContext(ctx, "POST", TokenExchangeURL, bytes.NewReader(jsonBody)) |
| 205 | if err != nil { |
| 206 | return "", fmt.Errorf("failed to create token exchange request: %w", err) |
| 207 | } |
| 208 | req.Header.Set("Content-Type", "application/json") |
| 209 | |
| 210 | resp, err := httpClient.Do(req) |
| 211 | if err != nil { |
| 212 | return "", fmt.Errorf("failed to exchange token: %w", err) |
| 213 | } |
| 214 | defer resp.Body.Close() |
| 215 | |
| 216 | body, err := io.ReadAll(resp.Body) |
| 217 | if err != nil { |
| 218 | return "", fmt.Errorf("failed to read token exchange response: %w", err) |
| 219 | } |
| 220 | |
| 221 | if resp.StatusCode == http.StatusBadRequest || resp.StatusCode == http.StatusUnauthorized || resp.StatusCode == http.StatusForbidden { |
| 222 | var oauthErr struct { |
| 223 | Error string `json:"error"` |
| 224 | ErrorDescription string `json:"error_description"` |
| 225 | } |
| 226 | if json.Unmarshal(body, &oauthErr) == nil { |
| 227 | return "", &OAuthError{ |
| 228 | StatusCode: resp.StatusCode, |
| 229 | ErrorCode: shared.OAuthErrorCode(oauthErr.Error), |
| 230 | ErrorDescription: oauthErr.ErrorDescription, |
| 231 | } |
| 232 | } |
| 233 | } |
| 234 | |
| 235 | if resp.StatusCode != http.StatusOK { |
| 236 | return "", fmt.Errorf("token exchange failed with status %d: %s", resp.StatusCode, string(body)) |
| 237 | } |
| 238 | |
| 239 | var tokenResp TokenExchangeResponse |
| 240 | if err := json.Unmarshal(body, &tokenResp); err != nil { |
| 241 | return "", fmt.Errorf("failed to decode token exchange response: %w", err) |
| 242 | } |
| 243 | |
| 244 | if tokenResp.AccessToken == "" { |
| 245 | return "", fmt.Errorf("token exchange response missing 'access_token' field. Response: %s", string(body)) |
| 246 | } |
| 247 | |
| 248 | expiresIn := int(DefaultTokenExpiry / time.Second) |
| 249 | if tokenResp.ExpiresIn != nil { |
| 250 | expiresIn = *tokenResp.ExpiresIn |
| 251 | } |
| 252 | |
| 253 | // Atomically update cached token and expiry |
| 254 | w.mu.Lock() |
| 255 | w.cachedToken = tokenResp.AccessToken |
| 256 | w.tokenExpiry = time.Now().Add(time.Duration(expiresIn) * time.Second) |
| 257 | w.mu.Unlock() |
| 258 | |
| 259 | return tokenResp.AccessToken, nil |
| 260 | } |
| 261 | |