|
| 1 | +# Copilot Proxy - Authentication Design |
| 2 | + |
| 3 | +This document provides a detailed implementation plan for the authentication components of the Copilot Proxy. It refines the initial `design_doc.md` by specifying the architecture required to handle the full, dynamic lifecycle of a Copilot API token. |
| 4 | + |
| 5 | +## 1. Overview |
| 6 | + |
| 7 | +The authentication system will be responsible for: |
| 8 | +1. Using a long-lived GitHub OAuth token. |
| 9 | +2. Exchanging it for a short-lived Copilot API token. |
| 10 | +3. Managing the Copilot token in memory. |
| 11 | +4. Automatically refreshing the token before it expires. |
| 12 | +5. Providing a valid token to the part of the proxy that forwards requests. |
| 13 | + |
| 14 | +To achieve this, the `pkg/copilot/` directory will be structured to separate the stateless API calls from the stateful token management. |
| 15 | + |
| 16 | +## 2. Proposed File Structure |
| 17 | + |
| 18 | +The `pkg/copilot/` directory will be organized as follows: |
| 19 | + |
| 20 | +``` |
| 21 | +/pkg/copilot/ |
| 22 | +├─── auth.go # Stateless functions for the token exchange API call. |
| 23 | +├─── token_manager.go # Stateful, thread-safe token lifecycle management. |
| 24 | +└─── client.go # (Updated) The proxy client, which uses the TokenManager. |
| 25 | +``` |
| 26 | + |
| 27 | +## 3. Component Deep Dive |
| 28 | + |
| 29 | +### `pkg/copilot/auth.go` |
| 30 | + |
| 31 | +**Responsibility:** Contains the low-level, stateless function for performing the GitHub-to-Copilot token exchange. |
| 32 | + |
| 33 | +**Implementation Details:** |
| 34 | + |
| 35 | +```go |
| 36 | +package copilot |
| 37 | + |
| 38 | +import ( |
| 39 | + "context" |
| 40 | + "encoding/json" |
| 41 | + "fmt" |
| 42 | + "net/http" |
| 43 | +) |
| 44 | + |
| 45 | +// ExchangeTokenResponse defines the structure of the JSON response from the token exchange endpoint. |
| 46 | +type ExchangeTokenResponse struct { |
| 47 | + Token string `json:"token"` |
| 48 | + ExpiresAt int64 `json:"expires_at"` |
| 49 | + RefreshIn int64 `json:"refresh_in"` |
| 50 | +} |
| 51 | + |
| 52 | +// ExchangeGitHubToken takes a GitHub OAuth token and exchanges it for a short-lived Copilot token. |
| 53 | +func ExchangeGitHubToken(ctx context.Context, githubToken string) (*ExchangeTokenResponse, error) { |
| 54 | + // 1. Create a new GET request to the exchange endpoint. |
| 55 | + url := "https://api.github.com/copilot_internal/v2/token" |
| 56 | + req, err := http.NewRequestWithContext(ctx, "GET", url, nil) |
| 57 | + if err != nil { |
| 58 | + return nil, fmt.Errorf("failed to create token exchange request: %w", err) |
| 59 | + } |
| 60 | + |
| 61 | + // 2. Add the required headers. |
| 62 | + req.Header.Set("Authorization", "Bearer "+githubToken) |
| 63 | + req.Header.Set("Accept", "application/json") |
| 64 | + |
| 65 | + // 3. Execute the request. |
| 66 | + resp, err := http.DefaultClient.Do(req) |
| 67 | + if err != nil { |
| 68 | + return nil, fmt.Errorf("failed to execute token exchange request: %w", err) |
| 69 | + } |
| 70 | + defer resp.Body.Close() |
| 71 | + |
| 72 | + // 4. Handle non-successful status codes. |
| 73 | + if resp.StatusCode != http.StatusOK { |
| 74 | + return nil, fmt.Errorf("token exchange request failed with status: %s", resp.Status) |
| 75 | + } |
| 76 | + |
| 77 | + // 5. Unmarshal the JSON response. |
| 78 | + var tokenResponse ExchangeTokenResponse |
| 79 | + if err := json.NewDecoder(resp.Body).Decode(&tokenResponse); err != nil { |
| 80 | + return nil, fmt.Errorf("failed to decode token exchange response: %w", err) |
| 81 | + } |
| 82 | + |
| 83 | + // 6. Return the response. |
| 84 | + return &tokenResponse, nil |
| 85 | +} |
| 86 | +``` |
| 87 | + |
| 88 | +### `pkg/copilot/token_manager.go` |
| 89 | + |
| 90 | +**Responsibility:** A stateful, thread-safe manager that holds the Copilot token and runs a background process to refresh it automatically. |
| 91 | + |
| 92 | +**Implementation Details:** |
| 93 | + |
| 94 | +```go |
| 95 | +package copilot |
| 96 | + |
| 97 | +import ( |
| 98 | + "context" |
| 99 | + "log/slog" |
| 100 | + "sync" |
| 101 | + "time" |
| 102 | +) |
| 103 | + |
| 104 | +// TokenManager handles the Copilot token and its refresh cycle. |
| 105 | +type TokenManager struct { |
| 106 | + mu sync.RWMutex |
| 107 | + githubToken string |
| 108 | + copilotToken string |
| 109 | + refreshesAt time.Time |
| 110 | + logger *slog.Logger |
| 111 | + stopCh chan struct{} |
| 112 | +} |
| 113 | + |
| 114 | +// NewTokenManager creates a manager, gets the initial token, and starts the refresh loop. |
| 115 | +func NewTokenManager(ctx context.Context, githubToken string, logger *slog.Logger) (*TokenManager, error) { |
| 116 | + tm := &TokenManager{ |
| 117 | + githubToken: githubToken, |
| 118 | + logger: logger, |
| 119 | + stopCh: make(chan struct{}), |
| 120 | + } |
| 121 | + |
| 122 | + if err := tm.refresh(ctx); err != nil { |
| 123 | + return nil, fmt.Errorf("initial token refresh failed: %w", err) |
| 124 | + } |
| 125 | + |
| 126 | + go tm.refreshTokenLoop() |
| 127 | + return tm, nil |
| 128 | +} |
| 129 | + |
| 130 | +// GetToken returns the current, valid Copilot token in a thread-safe way. |
| 131 | +func (tm *TokenManager) GetToken() string { |
| 132 | + tm.mu.RLock() |
| 133 | + defer tm.mu.RUnlock() |
| 134 | + return tm.copilotToken |
| 135 | +} |
| 136 | + |
| 137 | +// Close gracefully stops the background refresh loop. |
| 138 | +func (tm *TokenManager) Close() { |
| 139 | + close(tm.stopCh) |
| 140 | +} |
| 141 | + |
| 142 | +// refreshTokenLoop runs in the background, refreshing the token before it expires. |
| 143 | +func (tm *TokenManager) refreshTokenLoop() { |
| 144 | + for { |
| 145 | + tm.mu.RLock() |
| 146 | + duration := time.Until(tm.refreshesAt) |
| 147 | + tm.mu.RUnlock() |
| 148 | + |
| 149 | + select { |
| 150 | + case <-time.After(duration): |
| 151 | + tm.logger.Info("Refreshing Copilot token") |
| 152 | + if err := tm.refresh(context.Background()); err != nil { |
| 153 | + tm.logger.Error("Failed to refresh token", "error", err) |
| 154 | + } |
| 155 | + case <-tm.stopCh: |
| 156 | + tm.logger.Info("Token refresh loop stopped") |
| 157 | + return |
| 158 | + } |
| 159 | + } |
| 160 | +} |
| 161 | + |
| 162 | +// refresh executes the token exchange and updates the manager's state. |
| 163 | +func (tm *TokenManager) refresh(ctx context.Context) error { |
| 164 | + resp, err := ExchangeGitHubToken(ctx, tm.githubToken) |
| 165 | + if err != nil { |
| 166 | + return err |
| 167 | + } |
| 168 | + |
| 169 | + tm.mu.Lock() |
| 170 | + defer tm.mu.Unlock() |
| 171 | + |
| 172 | + tm.copilotToken = resp.Token |
| 173 | + // Refresh 60 seconds before the official refresh_in time as a buffer. |
| 174 | + refreshDuration := time.Duration(resp.RefreshIn-60) * time.Second |
| 175 | + tm.refreshesAt = time.Now().Add(refreshDuration) |
| 176 | + |
| 177 | + tm.logger.Info("Successfully refreshed Copilot token", "expires_at", resp.ExpiresAt) |
| 178 | + return nil |
| 179 | +} |
| 180 | +``` |
| 181 | + |
| 182 | +### `pkg/copilot/client.go` (Updated) |
| 183 | + |
| 184 | +**Responsibility:** The proxy client, updated to use the `TokenManager`. |
| 185 | + |
| 186 | +**Implementation Details:** |
| 187 | + |
| 188 | +```go |
| 189 | +package copilot |
| 190 | + |
| 191 | +import ( |
| 192 | + "context" |
| 193 | + "net/http" |
| 194 | + "time" |
| 195 | +) |
| 196 | + |
| 197 | +// Client now holds a reference to the TokenManager. |
| 198 | +type Client struct { |
| 199 | + httpClient *http.Client |
| 200 | + tokenManager *TokenManager |
| 201 | +} |
| 202 | + |
| 203 | +// NewClient is updated to accept the TokenManager. |
| 204 | +func NewClient(tokenManager *TokenManager, timeout time.Duration) *Client { |
| 205 | + return &Client{ |
| 206 | + httpClient: &http.Client{Timeout: timeout}, |
| 207 | + tokenManager: tokenManager, |
| 208 | + } |
| 209 | +} |
| 210 | + |
| 211 | +// ForwardRequest gets the latest token from the manager before each request. |
| 212 | +func (c *Client) ForwardRequest(ctx context.Context, incomingReq *http.Request) (*http.Response, error) { |
| 213 | + // ... (request creation logic remains the same) |
| 214 | + |
| 215 | + // Get the latest valid token for this specific request. |
| 216 | + token := c.tokenManager.GetToken() |
| 217 | + upstreamReq.Header.Set("Authorization", "Bearer "+token) |
| 218 | + |
| 219 | + // ... (request execution logic remains the same) |
| 220 | +} |
| 221 | +``` |
0 commit comments