Fix: 🎯 Fix Race Condition in GitHub App Installation Token Cache under Concurrent Requests

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()
+}

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)
+}