diff --git a/modules/auth/github_app_token.go b/modules/auth/github_app_token.go new file mode 100644 index 0000000..171cfff --- /dev/null +++ b/modules/auth/github_app_token.go @@ -0,0 +1,94 @@ +// Copyright 2024 The Gitea Authors. All rights reserved. +// SPDX-License-Identifier: MIT + +package auth + +import ( + "context" + "fmt" + "sync" + "time" + + "golang.org/x/sync/singleflight" +) + +// TokenExpiryBuffer is the time before actual expiration when we consider a token expired +const TokenExpiryBuffer = 2 * time.Minute + +type cachedToken struct { + Token string + Expiry time.Time +} + +// GitHubAppTokenCache provides thread-safe caching of GitHub App installation tokens +type GitHubAppTokenCache struct { + mu sync.RWMutex + tokens map[int64]*cachedToken // keyed by installation ID + sfg singleflight.Group +} + +// NewGitHubAppTokenCache creates a new token cache +func NewGitHubAppTokenCache() *GitHubAppTokenCache { + return &GitHubAppTokenCache{ + tokens: make(map[int64]*cachedToken), + } +} + +// GetToken retrieves a valid token for the given installation ID. +// If the cached token is expired or near expiration, it triggers a refresh. +// Concurrent requests for the same installation ID are coalesced using singleflight. +func (c *GitHubAppTokenCache) GetToken(ctx context.Context, installationID int64, refreshFunc func(context.Context, int64) (string, time.Time, error)) (string, error) { + // Fast path: check if we have a valid cached token + c.mu.RLock() + if cached, exists := c.tokens[installationID]; exists { + if time.Now().Add(TokenExpiryBuffer).Before(cached.Expiry) { + token := cached.Token + c.mu.RUnlock() + return token, nil + } + } + c.mu.RUnlock() + + // Token is expired or near expiration - use singleflight to ensure only one refresh + key := fmt.Sprintf("github-token-refresh-%d", installationID) + val, err, _ := c.sfg.Do(key, func() (interface{}, error) { + // Double-check pattern: another goroutine might have refreshed while we waited + c.mu.RLock() + if cached, exists := c.tokens[installationID]; exists { + if time.Now().Add(TokenExpiryBuffer).Before(cached.Expiry) { + token := cached.Token + c.mu.RUnlock() + return token, nil + } + } + c.mu.RUnlock() + + // Actually refresh the token + token, expiry, err := refreshFunc(ctx, installationID) + if err != nil { + return "", err + } + + // Update cache with new token + c.mu.Lock() + c.tokens[installationID] = &cachedToken{ + Token: token, + Expiry: expiry, + } + c.mu.Unlock() + + return token, nil + }) + + if err != nil { + return "", err + } + return val.(string), nil +} + +// InvalidateToken removes a token from the cache, forcing a refresh on next access +func (c *GitHubAppTokenCache) InvalidateToken(installationID int64) { + c.mu.Lock() + delete(c.tokens, installationID) + c.mu.Unlock() +} diff --git a/modules/auth/github_app_token_test.go b/modules/auth/github_app_token_test.go new file mode 100644 index 0000000..507dcc5 --- /dev/null +++ b/modules/auth/github_app_token_test.go @@ -0,0 +1,196 @@ +// Copyright 2024 The Gitea Authors. All rights reserved. +// SPDX-License-Identifier: MIT + +package auth + +import ( + "context" + "fmt" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestGitHubAppTokenCache_ConcurrentAccess(t *testing.T) { + cache := NewGitHubAppTokenCache() + installationID := int64(12345) + + var refreshCount atomic.Int32 + var tokenCounter atomic.Int32 + + // Mock refresh function that simulates GitHub API call + refreshFunc := func(ctx context.Context, id int64) (string, time.Time, error) { + refreshCount.Add(1) + time.Sleep(50 * time.Millisecond) // Simulate network latency + tokenNum := tokenCounter.Add(1) + return fmt.Sprintf("ghs_token_%d", tokenNum), time.Now().Add(5 * time.Minute), nil + } + + // Spawn 50 concurrent goroutines requesting the token + const numGoroutines = 50 + var wg sync.WaitGroup + wg.Add(numGoroutines) + + tokens := make([]string, numGoroutines) + errors := make([]error, numGoroutines) + + for i := 0; i < numGoroutines; i++ { + go func(idx int) { + defer wg.Done() + token, err := cache.GetToken(context.Background(), installationID, refreshFunc) + tokens[idx] = token + errors[idx] = err + }(i) + } + + wg.Wait() + + // Verify all goroutines got a token without error + for i := 0; i < numGoroutines; i++ { + require.NoError(t, errors[i], "goroutine %d got error", i) + require.NotEmpty(t, tokens[i], "goroutine %d got empty token", i) + } + + // Verify all goroutines got the same token (singleflight coalescing) + firstToken := tokens[0] + for i := 1; i < numGoroutines; i++ { + assert.Equal(t, firstToken, tokens[i], "token mismatch at index %d", i) + } + + // Verify refresh was called exactly once (singleflight working) + assert.Equal(t, int32(1), refreshCount.Load(), "refresh should be called exactly once") +} + +func TestGitHubAppTokenCache_ExpirationBuffer(t *testing.T) { + cache := NewGitHubAppTokenCache() + installationID := int64(67890) + + var refreshCount atomic.Int32 + + // First refresh: token expires in 1 minute (within buffer) + refreshFunc := func(ctx context.Context, id int64) (string, time.Time, error) { + count := refreshCount.Add(1) + if count == 1 { + // First call: token expires in 1 minute (should trigger refresh due to buffer) + return "ghs_token_1", time.Now().Add(1 * time.Minute), nil + } + // Second call: token expires in 10 minutes (safe) + return "ghs_token_2", time.Now().Add(10 * time.Minute), nil + } + + // First call should cache the token + token1, err := cache.GetToken(context.Background(), installationID, refreshFunc) + require.NoError(t, err) + assert.Equal(t, "ghs_token_1", token1) + assert.Equal(t, int32(1), refreshCount.Load()) + + // Second call should trigger refresh because token expires within buffer + token2, err := cache.GetToken(context.Background(), installationID, refreshFunc) + require.NoError(t, err) + assert.Equal(t, "ghs_token_2", token2) + assert.Equal(t, int32(2), refreshCount.Load()) + + // Third call should use cached token (no refresh) + token3, err := cache.GetToken(context.Background(), installationID, refreshFunc) + require.NoError(t, err) + assert.Equal(t, "ghs_token_2", token3) + assert.Equal(t, int32(2), refreshCount.Load(), "should not refresh again") +} + +func TestGitHubAppTokenCache_NoExpiredTokenReturned(t *testing.T) { + cache := NewGitHubAppTokenCache() + installationID := int64(11111) + + var refreshCount atomic.Int32 + + // Refresh function that returns very short-lived tokens + refreshFunc := func(ctx context.Context, id int64) (string, time.Time, error) { + count := refreshCount.Add(1) + // Token expires in 100ms + return fmt.Sprintf("ghs_token_%d", count), time.Now().Add(100 * time.Millisecond), nil + } + + // Get initial token + token1, err := cache.GetToken(context.Background(), installationID, refreshFunc) + require.NoError(t, err) + assert.Equal(t, "ghs_token_1", token1) + + // Wait for token to expire + time.Sleep(150 * time.Millisecond) + + // Concurrent requests after expiration + const numGoroutines = 20 + var wg sync.WaitGroup + wg.Add(numGoroutines) + + tokens := make([]string, numGoroutines) + for i := 0; i < numGoroutines; i++ { + go func(idx int) { + defer wg.Done() + token, err := cache.GetToken(context.Background(), installationID, refreshFunc) + require.NoError(t, err) + tokens[idx] = token + }(i) + } + + wg.Wait() + + // All goroutines should get the new token, not the expired one + for i := 0; i < numGoroutines; i++ { + assert.NotEqual(t, "ghs_token_1", tokens[i], "goroutine %d got expired token", i) + assert.Equal(t, "ghs_token_2", tokens[i], "goroutine %d should get refreshed token", i) + } + + // Refresh should be called exactly once for the expired token + assert.Equal(t, int32(2), refreshCount.Load()) +} + +func TestGitHubAppTokenCache_InvalidateToken(t *testing.T) { + cache := NewGitHubAppTokenCache() + installationID := int64(22222) + + var refreshCount atomic.Int32 + + refreshFunc := func(ctx context.Context, id int64) (string, time.Time, error) { + count := refreshCount.Add(1) + return fmt.Sprintf("ghs_token_%d", count), time.Now().Add(10 * time.Minute), nil + } + + // Get initial token + token1, err := cache.GetToken(context.Background(), installationID, refreshFunc) + require.NoError(t, err) + assert.Equal(t, "ghs_token_1", token1) + + // Invalidate the token + cache.InvalidateToken(installationID) + + // Next call should refresh + token2, err := cache.GetToken(context.Background(), installationID, refreshFunc) + require.NoError(t, err) + assert.Equal(t, "ghs_token_2", token2) + assert.Equal(t, int32(2), refreshCount.Load()) +} + +func TestGitHubAppTokenCache_MultipleInstallations(t *testing.T) { + cache := NewGitHubAppTokenCache() + + refreshFunc := func(ctx context.Context, id int64) (string, time.Time, error) { + return fmt.Sprintf("ghs_token_install_%d", id), time.Now().Add(10 * time.Minute), nil + } + + // Get tokens for different installations + token1, err := cache.GetToken(context.Background(), 100, refreshFunc) + require.NoError(t, err) + + token2, err := cache.GetToken(context.Background(), 200, refreshFunc) + require.NoError(t, err) + + // Tokens should be different and installation-specific + assert.Equal(t, "ghs_token_install_100", token1) + assert.Equal(t, "ghs_token_install_200", token2) + assert.NotEqual(t, token1, token2) +}