token.go

package main

import (
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"gopkg.in/yaml.v3"
)

const (
	copilotTokenURL = "https://api.github.com/copilot_internal/v2/token"
	// Refresh the token this many seconds before it actually expires.
	tokenRefreshBuffer = 60
)

// ghHostsFile is the default path for gh CLI OAuth token storage.
func ghHostsFile() string {
	if p := os.Getenv("GH_CONFIG_DIR"); p != "" {
		return filepath.Join(p, "hosts.yml")
	}
	home, _ := os.UserHomeDir()
	return filepath.Join(home, ".config", "gh", "hosts.yml")
}

// ghHosts is the minimal structure we need from hosts.yml.
type ghHosts map[string]struct {
	OAuthToken string `yaml:"oauth_token"`
}

// readGitHubPAT tries to find a GitHub PAT in order:
//  1. GITHUB_TOKEN env var
//  2. GH_TOKEN env var
//  3. gh CLI config (~/.config/gh/hosts.yml)
func readGitHubPAT() (string, error) {
	if t := os.Getenv("GITHUB_TOKEN"); t != "" {
		return t, nil
	}
	if t := os.Getenv("GH_TOKEN"); t != "" {
		return t, nil
	}

	hostsPath := ghHostsFile()
	data, err := os.ReadFile(hostsPath)
	if err != nil {
		return "", fmt.Errorf("no GITHUB_TOKEN set and cannot read %s: %w", hostsPath, err)
	}
	var hosts ghHosts
	if err := yaml.Unmarshal(data, &hosts); err != nil {
		return "", fmt.Errorf("parsing %s: %w", hostsPath, err)
	}
	for _, host := range []string{"github.com", "api.github.com"} {
		if h, ok := hosts[host]; ok && h.OAuthToken != "" {
			return h.OAuthToken, nil
		}
	}
	return "", fmt.Errorf("no oauth_token found in %s", hostsPath)
}

// copilotTokenResponse is the JSON returned by the Copilot token endpoint.
type copilotTokenResponse struct {
	Token     string `json:"token"`
	ExpiresAt int64  `json:"expires_at"`
}

// TokenManager caches a short-lived Copilot API token and refreshes it
// transparently before expiry.
type TokenManager struct {
	mu        sync.Mutex
	token     string
	expiresAt time.Time
	client    *http.Client
}

// NewTokenManager creates a TokenManager and eagerly fetches the first token.
func NewTokenManager() (*TokenManager, error) {
	tm := &TokenManager{
		client: &http.Client{Timeout: 15 * time.Second},
	}
	if err := tm.refresh(); err != nil {
		return nil, err
	}
	return tm, nil
}

// Token returns a valid Copilot API token, refreshing if necessary.
func (tm *TokenManager) Token() (string, error) {
	tm.mu.Lock()
	defer tm.mu.Unlock()

	if time.Now().Before(tm.expiresAt.Add(-tokenRefreshBuffer * time.Second)) {
		return tm.token, nil
	}
	if err := tm.refresh(); err != nil {
		return "", err
	}
	return tm.token, nil
}

// refresh fetches a new short-lived token from the GitHub API.
// Caller must hold tm.mu (or be in NewTokenManager before any goroutine sees tm).
func (tm *TokenManager) refresh() error {
	pat, err := readGitHubPAT()
	if err != nil {
		return fmt.Errorf("reading GitHub PAT: %w", err)
	}

	req, err := http.NewRequest(http.MethodGet, copilotTokenURL, nil)
	if err != nil {
		return err
	}
	req.Header.Set("Authorization", "token "+pat)
	req.Header.Set("Accept", "application/json")
	req.Header.Set("User-Agent", "GitHubCopilotChat/0.26.7")

	resp, err := tm.client.Do(req)
	if err != nil {
		return fmt.Errorf("fetching Copilot token: %w", err)
	}
	defer resp.Body.Close()

	body, _ := io.ReadAll(resp.Body)
	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("Copilot token endpoint returned %d: %s", resp.StatusCode, strings.TrimSpace(string(body)))
	}

	var tr copilotTokenResponse
	if err := json.Unmarshal(body, &tr); err != nil {
		return fmt.Errorf("parsing Copilot token response: %w", err)
	}
	if tr.Token == "" {
		return fmt.Errorf("empty token in Copilot token response")
	}

	tm.token = tr.Token
	tm.expiresAt = time.Unix(tr.ExpiresAt, 0)
	logger.Printf("Copilot token refreshed, expires at %s", tm.expiresAt.Format(time.RFC3339))
	return nil
}