From 50ac22b665b10cf4b33c57dbb9f3327a7df4acb5 Mon Sep 17 00:00:00 2001 From: Fabio Date: Tue, 3 Mar 2026 22:30:39 +0100 Subject: [PATCH] fix: refresh GitLab OAuth2 token on 401 during long runs Long-running analyses on large GitLab groups fail mid-run because the glab CLI OAuth2 token expires (~2h). Add a token refresh callback to the GitLab provider that re-invokes glab CLI on 401 responses, retrying the request once with the fresh token. Co-Authored-By: Claude Opus 4.6 --- cmd/codemium/main.go | 17 ++++++- internal/provider/gitlab.go | 47 ++++++++++++++++-- internal/provider/gitlab_test.go | 81 ++++++++++++++++++++++++++++++++ 3 files changed, 138 insertions(+), 7 deletions(-) diff --git a/cmd/codemium/main.go b/cmd/codemium/main.go index 5024827..f45c5a6 100644 --- a/cmd/codemium/main.go +++ b/cmd/codemium/main.go @@ -349,7 +349,14 @@ func runAnalyze(cmd *cobra.Command, args []string) error { return fmt.Errorf("--group is required for gitlab") } baseURL := os.Getenv("CODEMIUM_GITLAB_URL") - prov = provider.NewGitLab(cred.AccessToken, baseURL, httpClient) + // Enable token refresh when token was sourced from glab CLI + var refreshFn func() (string, bool) + if os.Getenv("CODEMIUM_GITLAB_TOKEN") == "" { + if _, err := store.Load("gitlab"); err != nil { + refreshFn = auth.GlabCLIToken + } + } + prov = provider.NewGitLabWithRefresh(cred.AccessToken, baseURL, httpClient, refreshFn) default: return fmt.Errorf("unsupported provider: %s", providerName) } @@ -1093,7 +1100,13 @@ func runTrends(cmd *cobra.Command, args []string) error { return fmt.Errorf("--group is required for gitlab") } baseURL := os.Getenv("CODEMIUM_GITLAB_URL") - prov = provider.NewGitLab(cred.AccessToken, baseURL, httpClient) + var refreshFn func() (string, bool) + if os.Getenv("CODEMIUM_GITLAB_TOKEN") == "" { + if _, err := store.Load("gitlab"); err != nil { + refreshFn = auth.GlabCLIToken + } + } + prov = provider.NewGitLabWithRefresh(cred.AccessToken, baseURL, httpClient, refreshFn) default: return fmt.Errorf("unsupported provider: %s", providerName) } diff --git a/internal/provider/gitlab.go b/internal/provider/gitlab.go index 5c69154..f86752d 100644 --- a/internal/provider/gitlab.go +++ b/internal/provider/gitlab.go @@ -8,6 +8,7 @@ import ( "net/http" "net/url" "strings" + "sync" "time" "github.com/dsablic/codemium/internal/model" @@ -17,9 +18,11 @@ const gitlabAPIBase = "https://gitlab.com" // GitLab implements Provider and CommitLister for GitLab. type GitLab struct { - token string - baseURL string - client *http.Client + token string + baseURL string + client *http.Client + tokenRefresh func() (string, bool) + mu sync.Mutex } // NewGitLab creates a new GitLab provider. If baseURL is empty, @@ -38,6 +41,14 @@ func NewGitLab(token string, baseURL string, client *http.Client) *GitLab { } } +// NewGitLabWithRefresh creates a GitLab provider with a token refresh callback. +// The refreshFn is called on 401 responses to obtain a fresh token (e.g. via glab CLI). +func NewGitLabWithRefresh(token string, baseURL string, client *http.Client, refreshFn func() (string, bool)) *GitLab { + g := NewGitLab(token, baseURL, client) + g.tokenRefresh = refreshFn + return g +} + type gitlabProject struct { ID int `json:"id"` Path string `json:"path"` @@ -96,12 +107,38 @@ func (g *GitLab) ListRepos(ctx context.Context, opts ListOpts) ([]model.Repo, er } func (g *GitLab) doGet(ctx context.Context, reqURL string) (*http.Response, error) { + g.mu.Lock() + token := g.token + g.mu.Unlock() + req, err := http.NewRequestWithContext(ctx, http.MethodGet, reqURL, nil) if err != nil { return nil, err } - req.Header.Set("Authorization", "Bearer "+g.token) - return g.client.Do(req) + req.Header.Set("Authorization", "Bearer "+token) + + resp, err := g.client.Do(req) + if err != nil { + return nil, err + } + + if resp.StatusCode == http.StatusUnauthorized && g.tokenRefresh != nil { + resp.Body.Close() + if newToken, ok := g.tokenRefresh(); ok && newToken != token { + g.mu.Lock() + g.token = newToken + g.mu.Unlock() + + req2, err := http.NewRequestWithContext(ctx, http.MethodGet, reqURL, nil) + if err != nil { + return nil, err + } + req2.Header.Set("Authorization", "Bearer "+newToken) + return g.client.Do(req2) + } + } + + return resp, nil } func (g *GitLab) fetchPage(ctx context.Context, pageURL string) ([]model.Repo, string, error) { diff --git a/internal/provider/gitlab_test.go b/internal/provider/gitlab_test.go index 6e94547..362780d 100644 --- a/internal/provider/gitlab_test.go +++ b/internal/provider/gitlab_test.go @@ -262,6 +262,87 @@ func TestGitLabCommitStats(t *testing.T) { } } +func TestGitLabTokenRefreshOn401(t *testing.T) { + calls := 0 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + calls++ + auth := r.Header.Get("Authorization") + if auth == "Bearer expired-token" { + w.WriteHeader(http.StatusUnauthorized) + return + } + if auth == "Bearer fresh-token" { + json.NewEncoder(w).Encode([]map[string]any{ + { + "id": 1, "path": "repo-1", "path_with_namespace": "g/repo-1", + "name": "Repo 1", "web_url": "h", "http_url_to_repo": "c", + "default_branch": "main", "archived": false, + "forked_from_project": nil, + "namespace": map[string]any{"full_path": "g"}, + }, + }) + return + } + t.Errorf("unexpected auth header: %s", auth) + w.WriteHeader(http.StatusUnauthorized) + })) + defer server.Close() + + refreshCalls := 0 + refreshFn := func() (string, bool) { + refreshCalls++ + return "fresh-token", true + } + + gl := provider.NewGitLabWithRefresh("expired-token", server.URL, nil, refreshFn) + repos, err := gl.ListRepos(context.Background(), provider.ListOpts{Organization: "g"}) + if err != nil { + t.Fatalf("expected successful retry, got error: %v", err) + } + if len(repos) != 1 { + t.Fatalf("expected 1 repo, got %d", len(repos)) + } + if refreshCalls != 1 { + t.Errorf("expected 1 refresh call, got %d", refreshCalls) + } + if calls != 2 { + t.Errorf("expected 2 server calls (401 + retry), got %d", calls) + } +} + +func TestGitLabTokenRefreshNotCalledWithoutCallback(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusUnauthorized) + })) + defer server.Close() + + gl := provider.NewGitLab("expired-token", server.URL, nil) + _, err := gl.ListRepos(context.Background(), provider.ListOpts{Organization: "g"}) + if err == nil { + t.Fatal("expected error on 401 without refresh callback") + } + if !strings.Contains(err.Error(), "401") { + t.Errorf("expected 401 in error, got: %v", err) + } +} + +func TestGitLabTokenRefreshFailure(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusUnauthorized) + })) + defer server.Close() + + refreshFn := func() (string, bool) { + return "", false + } + + gl := provider.NewGitLabWithRefresh("expired-token", server.URL, nil, refreshFn) + _, err := gl.ListRepos(context.Background(), provider.ListOpts{Organization: "g"}) + if err == nil { + t.Fatal("expected error when refresh fails") + } +} + func TestGitLabListCommitsLimit(t *testing.T) { page := 0 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {