mirror of
https://github.com/lukaszraczylo/traefikoidc.git
synced 2026-10-13 04:07:37 +00:00
Complete rebuild of the plugin
* Fix bug affecting Azure OIDC authentication ( and most likely others ) * Fixes issue #51 * Ensure that appended roles are unique. Update the documentation. * Improvements targetting possible memory usage spikes. * Additional fixes and cleanup * Refactoring code to fix the issues identified by the users. * Modernize run * Fieldalignment * Multiple changes to improve performance and reduce complexity. - Optimise the errors and recovery. - Deduplicate code in metadata cache. - Remove unused performance monitoring code. - Simplify session management and settings handling. * Fix claims issue. * Add ability to overwrite the default scopes in the settings file * Well.. that escalated quickly. Completely forgot that Traefik uses outdated Yaegi and requires compatibility with 1.20 ( pre-generic Go code ). * Bugfix #51: Ensures that user provided scopes overrides work. * fixup! Bugfix #51: Ensures that user provided scopes overrides work. * fixup! fixup! Bugfix #51: Ensures that user provided scopes overrides work. * Abstract the provider logic into a separate package. * Additional micro fixes and cleanups. * Simplify all the things. * fixup! Simplify all the things. * fixup! fixup! Simplify all the things. * fixup! fixup! fixup! Simplify all the things. * fixup! fixup! fixup! fixup! Simplify all the things. * ... * Cleanup tests. * fixup! Cleanup tests. * fixup! fixup! fixup! Cleanup tests. * fixup! fixup! fixup! fixup! Cleanup tests. * fixup! fixup! fixup! fixup! fixup! Cleanup tests. * Issue #53: Fix CSRF token handling in reverse proxy 1. ✅ HTTPS Detection Fixed (session.go:723) - Now uses X-Forwarded-Proto header instead of r.URL.Scheme - Properly detects HTTPS in reverse proxy environments 2. ✅ SameSite Cookie Attribute Fixed - Removed automatic SameSiteStrictMode for HTTPS (would break OAuth) - Keeps SameSiteLaxMode to allow OAuth callbacks from external domains - Only uses Strict for AJAX requests which don't involve OAuth redirects 3. ✅ Cookie Domain Handling Fixed - Now respects X-Forwarded-Host header for cookie domain - Ensures cookies are set for the public domain, not internal proxy domain 4. ✅ EnhanceSessionSecurity Properly Integrated - Function is now actually called during session save - Applies security enhancements without breaking OAuth flow Why Issue #53 Failed Before: 1. Cookies were not marked Secure in HTTPS environments (browser wouldn't send them back) 2. If they had been Secure with SameSite=Strict, Azure callbacks would still fail 3. Cookie domain might have been wrong (internal vs public domain) Why It Works Now: 1. Cookies are properly marked Secure for HTTPS 2. Uses SameSite=Lax to allow OAuth provider callbacks 3. Cookie domain uses public domain from X-Forwarded-Host 4. CSRF token persists through the entire OAuth flow * Next set of enhancements together with memory usage improvements. * Memory leak fixes and optimisations. * CSRF and Cookie Domain fixes * fixup! CSRF and Cookie Domain fixes * Metadata cache leak fix + profiling * fixup! Metadata cache leak fix + profiling * Memory leaks hunting, part 1337. * Further pursue of perfection. * fixup! Further pursue of perfection. * fixup! fixup! Further pursue of perfection. * fixup! fixup! fixup! Further pursue of perfection. * fixup! fixup! fixup! fixup! Further pursue of perfection. * fixup! fixup! fixup! fixup! fixup! Further pursue of perfection. * fixup! fixup! fixup! fixup! fixup! fixup! Further pursue of perfection. * fixup! fixup! fixup! fixup! fixup! fixup! fixup! Further pursue of perfection. * fixup! fixup! fixup! fixup! fixup! fixup! fixup! fixup! Further pursue of perfection. * fixup! fixup! fixup! fixup! fixup! fixup! fixup! fixup! fixup! Further pursue of perfection. * Clear race conditions * fixup! Clear race conditions * Weekend fun with memory leaks * Splitting code into multiple files with reasonable testing coverage. ``` ok github.com/lukaszraczylo/traefikoidc 117.017s coverage: 72.6% of statements ok github.com/lukaszraczylo/traefikoidc/auth 0.505s coverage: 87.1% of statements ok github.com/lukaszraczylo/traefikoidc/circuit_breaker 0.283s coverage: 99.0% of statements github.com/lukaszraczylo/traefikoidc/config coverage: 0.0% of statements ok github.com/lukaszraczylo/traefikoidc/handlers 0.349s coverage: 98.2% of statements ok github.com/lukaszraczylo/traefikoidc/internal/providers (cached) coverage: 94.3% of statements ok github.com/lukaszraczylo/traefikoidc/middleware 0.808s coverage: 78.0% of statements ok github.com/lukaszraczylo/traefikoidc/recovery 0.653s coverage: 100.0% of statements ok github.com/lukaszraczylo/traefikoidc/session/chunking (cached) coverage: 87.8% of statements ok github.com/lukaszraczylo/traefikoidc/session/core (cached) coverage: 85.6% of statements ok github.com/lukaszraczylo/traefikoidc/session/crypto (cached) coverage: 81.8% of statements ok github.com/lukaszraczylo/traefikoidc/session/storage (cached) coverage: 93.5% of statements ok github.com/lukaszraczylo/traefikoidc/session/validators (cached) coverage: 98.8% of statements ```` * fixup! Splitting code into multiple files with reasonable testing coverage. * fixup! fixup! Splitting code into multiple files with reasonable testing coverage. * Weekend fun with further optimisations. * fixup! Weekend fun with further optimisations. * fixup! fixup! Weekend fun with further optimisations. * fixup! fixup! fixup! Weekend fun with further optimisations. * fixup! fixup! fixup! fixup! Weekend fun with further optimisations. * fixup! fixup! fixup! fixup! fixup! Weekend fun with further optimisations. * Pre-release cleanup. * Enhance test coverage. * fixup! Enhance test coverage. * fixup! fixup! Enhance test coverage. * fixup! fixup! fixup! Enhance test coverage.
This commit is contained in:
Vendored
+426
@@ -0,0 +1,426 @@
|
||||
package cache
|
||||
|
||||
import (
|
||||
"container/list"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Type defines the type of cache for optimized behavior
|
||||
type Type string
|
||||
|
||||
const (
|
||||
TypeToken Type = "token"
|
||||
TypeMetadata Type = "metadata"
|
||||
TypeJWK Type = "jwk"
|
||||
TypeSession Type = "session"
|
||||
TypeGeneral Type = "general"
|
||||
)
|
||||
|
||||
// Logger interface for cache operations
|
||||
type Logger interface {
|
||||
Debug(msg string)
|
||||
Debugf(format string, args ...interface{})
|
||||
Info(msg string)
|
||||
Infof(format string, args ...interface{})
|
||||
Error(msg string)
|
||||
Errorf(format string, args ...interface{})
|
||||
}
|
||||
|
||||
// Config provides configuration for the cache
|
||||
type Config struct {
|
||||
Type Type
|
||||
MaxSize int
|
||||
MaxMemoryBytes int64
|
||||
DefaultTTL time.Duration
|
||||
CleanupInterval time.Duration
|
||||
EnableCompression bool
|
||||
EnableMetrics bool
|
||||
EnableAutoCleanup bool
|
||||
EnableMemoryLimit bool
|
||||
Logger Logger
|
||||
|
||||
// Type-specific configurations
|
||||
TokenConfig *TokenConfig
|
||||
MetadataConfig *MetadataConfig
|
||||
JWKConfig *JWKConfig
|
||||
}
|
||||
|
||||
// TokenConfig provides token-specific cache configuration
|
||||
type TokenConfig struct {
|
||||
BlacklistTTL time.Duration
|
||||
RefreshTokenTTL time.Duration
|
||||
EnableTokenRotation bool
|
||||
}
|
||||
|
||||
// MetadataConfig provides metadata-specific cache configuration
|
||||
type MetadataConfig struct {
|
||||
GracePeriod time.Duration
|
||||
ExtendedGracePeriod time.Duration
|
||||
MaxGracePeriod time.Duration
|
||||
SecurityCriticalMaxGracePeriod time.Duration
|
||||
SecurityCriticalFields []string
|
||||
}
|
||||
|
||||
// JWKConfig provides JWK-specific cache configuration
|
||||
type JWKConfig struct {
|
||||
RefreshInterval time.Duration
|
||||
MinRefreshTime time.Duration
|
||||
MaxKeyAge time.Duration
|
||||
}
|
||||
|
||||
// Item represents a single cache entry
|
||||
type Item struct {
|
||||
Key string
|
||||
Value interface{}
|
||||
Size int64
|
||||
ExpiresAt time.Time
|
||||
LastAccessed time.Time
|
||||
AccessCount int64
|
||||
CacheType Type
|
||||
|
||||
// Type-specific metadata
|
||||
Metadata map[string]interface{}
|
||||
|
||||
// LRU list element reference
|
||||
element *list.Element
|
||||
}
|
||||
|
||||
// Cache provides a single, unified cache implementation
|
||||
type Cache struct {
|
||||
mu sync.RWMutex
|
||||
items map[string]*Item
|
||||
lruList *list.List
|
||||
config Config
|
||||
logger Logger
|
||||
|
||||
// Memory management
|
||||
currentSize int64
|
||||
currentMemory int64
|
||||
|
||||
// Metrics
|
||||
hits int64
|
||||
misses int64
|
||||
evictions int64
|
||||
sets int64
|
||||
|
||||
// Lifecycle management
|
||||
ctx context.Context
|
||||
cancel context.CancelFunc
|
||||
wg sync.WaitGroup
|
||||
stopCleanup chan bool
|
||||
closed int32
|
||||
}
|
||||
|
||||
// DefaultConfig returns a default cache configuration
|
||||
func DefaultConfig() Config {
|
||||
return Config{
|
||||
Type: TypeGeneral,
|
||||
MaxSize: 1000,
|
||||
MaxMemoryBytes: 64 * 1024 * 1024, // 64MB
|
||||
DefaultTTL: 10 * time.Minute,
|
||||
CleanupInterval: 5 * time.Minute,
|
||||
EnableAutoCleanup: true,
|
||||
EnableMemoryLimit: true,
|
||||
EnableMetrics: true,
|
||||
}
|
||||
}
|
||||
|
||||
// New creates a new cache instance
|
||||
func New(config Config) *Cache {
|
||||
if config.Logger == nil {
|
||||
config.Logger = &noOpLogger{}
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
c := &Cache{
|
||||
items: make(map[string]*Item),
|
||||
lruList: list.New(),
|
||||
config: config,
|
||||
logger: config.Logger,
|
||||
ctx: ctx,
|
||||
cancel: cancel,
|
||||
}
|
||||
|
||||
if config.EnableAutoCleanup && config.CleanupInterval > 0 {
|
||||
c.stopCleanup = make(chan bool)
|
||||
c.startCleanupRoutine()
|
||||
}
|
||||
|
||||
return c
|
||||
}
|
||||
|
||||
// Set stores a value with TTL
|
||||
func (c *Cache) Set(key string, value interface{}, ttl time.Duration) error {
|
||||
if atomic.LoadInt32(&c.closed) == 1 {
|
||||
return fmt.Errorf("cache is closed")
|
||||
}
|
||||
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
|
||||
// Calculate size
|
||||
size := c.estimateSize(value)
|
||||
|
||||
// Check memory limit
|
||||
if c.config.EnableMemoryLimit && c.currentMemory+size > c.config.MaxMemoryBytes {
|
||||
c.evictLRU()
|
||||
}
|
||||
|
||||
// Check size limit
|
||||
if c.config.MaxSize > 0 && len(c.items) >= c.config.MaxSize {
|
||||
c.evictLRU()
|
||||
}
|
||||
|
||||
// Create or update item
|
||||
item := &Item{
|
||||
Key: key,
|
||||
Value: value,
|
||||
Size: size,
|
||||
ExpiresAt: time.Now().Add(ttl),
|
||||
LastAccessed: time.Now(),
|
||||
AccessCount: 0,
|
||||
CacheType: c.config.Type,
|
||||
Metadata: make(map[string]interface{}),
|
||||
}
|
||||
|
||||
// Remove old item if exists
|
||||
if oldItem, exists := c.items[key]; exists {
|
||||
c.lruList.Remove(oldItem.element)
|
||||
c.currentMemory -= oldItem.Size
|
||||
c.currentSize--
|
||||
}
|
||||
|
||||
// Add new item
|
||||
item.element = c.lruList.PushFront(item)
|
||||
c.items[key] = item
|
||||
c.currentMemory += size
|
||||
c.currentSize++
|
||||
atomic.AddInt64(&c.sets, 1)
|
||||
|
||||
c.logger.Debugf("Cache: Set key=%s, size=%d, ttl=%v", key, size, ttl)
|
||||
return nil
|
||||
}
|
||||
|
||||
// Get retrieves a value from cache
|
||||
func (c *Cache) Get(key string) (interface{}, bool) {
|
||||
if atomic.LoadInt32(&c.closed) == 1 {
|
||||
return nil, false
|
||||
}
|
||||
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
|
||||
item, exists := c.items[key]
|
||||
if !exists {
|
||||
atomic.AddInt64(&c.misses, 1)
|
||||
return nil, false
|
||||
}
|
||||
|
||||
// Check expiration
|
||||
if time.Now().After(item.ExpiresAt) {
|
||||
c.removeItem(key, item)
|
||||
atomic.AddInt64(&c.misses, 1)
|
||||
return nil, false
|
||||
}
|
||||
|
||||
// Update LRU
|
||||
c.lruList.MoveToFront(item.element)
|
||||
item.LastAccessed = time.Now()
|
||||
item.AccessCount++
|
||||
atomic.AddInt64(&c.hits, 1)
|
||||
|
||||
return item.Value, true
|
||||
}
|
||||
|
||||
// Delete removes a key from cache
|
||||
func (c *Cache) Delete(key string) {
|
||||
if atomic.LoadInt32(&c.closed) == 1 {
|
||||
return
|
||||
}
|
||||
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
|
||||
if item, exists := c.items[key]; exists {
|
||||
c.removeItem(key, item)
|
||||
}
|
||||
}
|
||||
|
||||
// Clear removes all items from cache
|
||||
func (c *Cache) Clear() {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
|
||||
c.items = make(map[string]*Item)
|
||||
c.lruList.Init()
|
||||
c.currentSize = 0
|
||||
c.currentMemory = 0
|
||||
}
|
||||
|
||||
// Size returns the number of items in cache
|
||||
func (c *Cache) Size() int {
|
||||
c.mu.RLock()
|
||||
defer c.mu.RUnlock()
|
||||
return len(c.items)
|
||||
}
|
||||
|
||||
// SetMaxSize updates the maximum cache size
|
||||
func (c *Cache) SetMaxSize(size int) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
c.config.MaxSize = size
|
||||
|
||||
// Evict items if necessary
|
||||
for len(c.items) > size && c.lruList.Len() > 0 {
|
||||
c.evictLRU()
|
||||
}
|
||||
}
|
||||
|
||||
// GetStats returns cache statistics
|
||||
func (c *Cache) GetStats() map[string]interface{} {
|
||||
c.mu.RLock()
|
||||
defer c.mu.RUnlock()
|
||||
|
||||
return map[string]interface{}{
|
||||
"size": c.currentSize,
|
||||
"memory": c.currentMemory,
|
||||
"hits": atomic.LoadInt64(&c.hits),
|
||||
"misses": atomic.LoadInt64(&c.misses),
|
||||
"evictions": atomic.LoadInt64(&c.evictions),
|
||||
"sets": atomic.LoadInt64(&c.sets),
|
||||
"hit_rate": c.calculateHitRate(),
|
||||
"cache_type": string(c.config.Type),
|
||||
}
|
||||
}
|
||||
|
||||
// Close gracefully shuts down the cache
|
||||
func (c *Cache) Close() error {
|
||||
if !atomic.CompareAndSwapInt32(&c.closed, 0, 1) {
|
||||
return fmt.Errorf("cache already closed")
|
||||
}
|
||||
|
||||
c.cancel()
|
||||
if c.config.EnableAutoCleanup {
|
||||
close(c.stopCleanup)
|
||||
c.wg.Wait()
|
||||
}
|
||||
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
// Clear inline to avoid double locking
|
||||
c.items = make(map[string]*Item)
|
||||
c.lruList.Init()
|
||||
c.currentSize = 0
|
||||
c.currentMemory = 0
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// Cleanup removes expired items
|
||||
func (c *Cache) Cleanup() {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
|
||||
now := time.Now()
|
||||
var toRemove []string
|
||||
|
||||
for key, item := range c.items {
|
||||
if now.After(item.ExpiresAt) {
|
||||
toRemove = append(toRemove, key)
|
||||
}
|
||||
}
|
||||
|
||||
for _, key := range toRemove {
|
||||
if item, exists := c.items[key]; exists {
|
||||
c.removeItem(key, item)
|
||||
}
|
||||
}
|
||||
|
||||
c.logger.Debugf("Cache cleanup: removed %d expired items", len(toRemove))
|
||||
}
|
||||
|
||||
// Private methods
|
||||
|
||||
func (c *Cache) removeItem(key string, item *Item) {
|
||||
c.lruList.Remove(item.element)
|
||||
delete(c.items, key)
|
||||
c.currentMemory -= item.Size
|
||||
c.currentSize--
|
||||
}
|
||||
|
||||
func (c *Cache) evictLRU() {
|
||||
if elem := c.lruList.Back(); elem != nil {
|
||||
item := elem.Value.(*Item)
|
||||
c.removeItem(item.Key, item)
|
||||
atomic.AddInt64(&c.evictions, 1)
|
||||
c.logger.Debugf("Cache: Evicted LRU item key=%s", item.Key)
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Cache) estimateSize(value interface{}) int64 {
|
||||
// Simple size estimation
|
||||
switch v := value.(type) {
|
||||
case string:
|
||||
return int64(len(v))
|
||||
case []byte:
|
||||
return int64(len(v))
|
||||
case map[string]interface{}:
|
||||
// Rough estimation for maps
|
||||
data, _ := json.Marshal(v)
|
||||
return int64(len(data))
|
||||
default:
|
||||
// Default size for unknown types
|
||||
return 256
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Cache) calculateHitRate() float64 {
|
||||
hits := atomic.LoadInt64(&c.hits)
|
||||
misses := atomic.LoadInt64(&c.misses)
|
||||
total := hits + misses
|
||||
if total == 0 {
|
||||
return 0
|
||||
}
|
||||
return float64(hits) / float64(total)
|
||||
}
|
||||
|
||||
func (c *Cache) startCleanupRoutine() {
|
||||
c.wg.Add(1)
|
||||
go func() {
|
||||
defer c.wg.Done()
|
||||
ticker := time.NewTicker(c.config.CleanupInterval)
|
||||
defer ticker.Stop()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-ticker.C:
|
||||
c.Cleanup()
|
||||
case <-c.stopCleanup:
|
||||
return
|
||||
case <-c.ctx.Done():
|
||||
return
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
// noOpLogger provides a no-op logger implementation
|
||||
type noOpLogger struct{}
|
||||
|
||||
func (l *noOpLogger) Debug(msg string) {}
|
||||
func (l *noOpLogger) Debugf(format string, args ...interface{}) {}
|
||||
func (l *noOpLogger) Info(msg string) {}
|
||||
func (l *noOpLogger) Infof(format string, args ...interface{}) {}
|
||||
func (l *noOpLogger) Error(msg string) {}
|
||||
func (l *noOpLogger) Errorf(format string, args ...interface{}) {}
|
||||
func (l *noOpLogger) Warn(msg string) {}
|
||||
func (l *noOpLogger) Warnf(format string, args ...interface{}) {}
|
||||
func (l *noOpLogger) Fatal(msg string) {}
|
||||
func (l *noOpLogger) Fatalf(format string, args ...interface{}) {}
|
||||
func (l *noOpLogger) WithField(key string, value interface{}) Logger { return l }
|
||||
func (l *noOpLogger) WithFields(fields map[string]interface{}) Logger { return l }
|
||||
Vendored
+2040
File diff suppressed because it is too large
Load Diff
Vendored
+278
@@ -0,0 +1,278 @@
|
||||
package cache
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// CompatibilityWrapper provides backward compatibility with existing cache interfaces
|
||||
type CompatibilityWrapper struct {
|
||||
cache *Cache
|
||||
}
|
||||
|
||||
// NewCompatibilityWrapper creates a new compatibility wrapper
|
||||
func NewCompatibilityWrapper(cache *Cache) *CompatibilityWrapper {
|
||||
return &CompatibilityWrapper{cache: cache}
|
||||
}
|
||||
|
||||
// CacheInterface implementation for backward compatibility
|
||||
func (c *CompatibilityWrapper) Set(key string, value interface{}, ttl time.Duration) {
|
||||
_ = c.cache.Set(key, value, ttl)
|
||||
}
|
||||
|
||||
func (c *CompatibilityWrapper) Get(key string) (interface{}, bool) {
|
||||
return c.cache.Get(key)
|
||||
}
|
||||
|
||||
func (c *CompatibilityWrapper) Delete(key string) {
|
||||
c.cache.Delete(key)
|
||||
}
|
||||
|
||||
func (c *CompatibilityWrapper) SetMaxSize(size int) {
|
||||
c.cache.SetMaxSize(size)
|
||||
}
|
||||
|
||||
func (c *CompatibilityWrapper) Size() int {
|
||||
return c.cache.Size()
|
||||
}
|
||||
|
||||
func (c *CompatibilityWrapper) Clear() {
|
||||
c.cache.Clear()
|
||||
}
|
||||
|
||||
func (c *CompatibilityWrapper) Cleanup() {
|
||||
c.cache.Cleanup()
|
||||
}
|
||||
|
||||
func (c *CompatibilityWrapper) Close() {
|
||||
_ = c.cache.Close()
|
||||
}
|
||||
|
||||
func (c *CompatibilityWrapper) GetStats() map[string]interface{} {
|
||||
return c.cache.GetStats()
|
||||
}
|
||||
|
||||
// UniversalCacheCompat provides compatibility with the old UniversalCache
|
||||
type UniversalCacheCompat struct {
|
||||
*Cache
|
||||
}
|
||||
|
||||
// NewUniversalCacheCompat creates a compatibility wrapper for UniversalCache
|
||||
func NewUniversalCacheCompat(config Config) *UniversalCacheCompat {
|
||||
return &UniversalCacheCompat{
|
||||
Cache: New(config),
|
||||
}
|
||||
}
|
||||
|
||||
// Set wraps the cache Set method for compatibility
|
||||
func (u *UniversalCacheCompat) Set(key string, value interface{}, ttl time.Duration) error {
|
||||
return u.Cache.Set(key, value, ttl)
|
||||
}
|
||||
|
||||
// TokenCacheCompat provides compatibility with the old TokenCache
|
||||
type TokenCacheCompat struct {
|
||||
cache *TokenCache
|
||||
}
|
||||
|
||||
// NewTokenCacheCompat creates a compatibility wrapper for TokenCache
|
||||
func NewTokenCacheCompat() *TokenCacheCompat {
|
||||
manager := GetGlobalManager(nil)
|
||||
return &TokenCacheCompat{
|
||||
cache: manager.GetTokenCache(),
|
||||
}
|
||||
}
|
||||
|
||||
// Set stores parsed token claims
|
||||
func (t *TokenCacheCompat) Set(token string, claims map[string]interface{}, expiration time.Duration) {
|
||||
_ = t.cache.Set(token, claims, expiration)
|
||||
}
|
||||
|
||||
// Get retrieves cached claims for a token
|
||||
func (t *TokenCacheCompat) Get(token string) (map[string]interface{}, bool) {
|
||||
return t.cache.Get(token)
|
||||
}
|
||||
|
||||
// Delete removes a token from cache
|
||||
func (t *TokenCacheCompat) Delete(token string) {
|
||||
t.cache.Delete(token)
|
||||
}
|
||||
|
||||
// MetadataCacheCompat provides compatibility with the old MetadataCache
|
||||
type MetadataCacheCompat struct {
|
||||
cache *MetadataCache
|
||||
logger Logger
|
||||
wg *sync.WaitGroup
|
||||
}
|
||||
|
||||
// NewMetadataCacheCompat creates a compatibility wrapper for MetadataCache
|
||||
func NewMetadataCacheCompat(wg *sync.WaitGroup) *MetadataCacheCompat {
|
||||
manager := GetGlobalManager(nil)
|
||||
return &MetadataCacheCompat{
|
||||
cache: manager.GetMetadataCache(),
|
||||
logger: manager.logger,
|
||||
wg: wg,
|
||||
}
|
||||
}
|
||||
|
||||
// NewMetadataCacheCompatWithLogger creates a MetadataCache with specific logger
|
||||
func NewMetadataCacheCompatWithLogger(wg *sync.WaitGroup, logger Logger) *MetadataCacheCompat {
|
||||
manager := GetGlobalManager(logger)
|
||||
return &MetadataCacheCompat{
|
||||
cache: manager.GetMetadataCache(),
|
||||
logger: logger,
|
||||
wg: wg,
|
||||
}
|
||||
}
|
||||
|
||||
// Set stores provider metadata with a TTL
|
||||
func (m *MetadataCacheCompat) Set(providerURL string, metadata *ProviderMetadata, ttl time.Duration) error {
|
||||
return m.cache.Set(providerURL, metadata, ttl)
|
||||
}
|
||||
|
||||
// Get retrieves provider metadata from cache
|
||||
func (m *MetadataCacheCompat) Get(providerURL string) (*ProviderMetadata, bool) {
|
||||
return m.cache.Get(providerURL)
|
||||
}
|
||||
|
||||
// Delete removes provider metadata
|
||||
func (m *MetadataCacheCompat) Delete(providerURL string) {
|
||||
m.cache.Delete(providerURL)
|
||||
}
|
||||
|
||||
// GetWithGracePeriod retrieves metadata with grace period support
|
||||
func (m *MetadataCacheCompat) GetWithGracePeriod(ctx context.Context, providerURL string) (*ProviderMetadata, bool) {
|
||||
// For compatibility, just use regular Get
|
||||
return m.cache.Get(providerURL)
|
||||
}
|
||||
|
||||
// JWKCacheCompat provides compatibility with the old JWKCache
|
||||
type JWKCacheCompat struct {
|
||||
cache *JWKCache
|
||||
}
|
||||
|
||||
// NewJWKCacheCompat creates a compatibility wrapper for JWKCache
|
||||
func NewJWKCacheCompat() *JWKCacheCompat {
|
||||
manager := GetGlobalManager(nil)
|
||||
return &JWKCacheCompat{
|
||||
cache: manager.GetJWKCache(),
|
||||
}
|
||||
}
|
||||
|
||||
// GetJWKS retrieves JWKS from cache or fetches from the remote URL if not cached
|
||||
func (j *JWKCacheCompat) GetJWKS(ctx context.Context, jwksURL string, httpClient *http.Client) (*JWKSet, error) {
|
||||
// Check cache first
|
||||
if jwks, found := j.cache.Get(jwksURL); found {
|
||||
return jwks, nil
|
||||
}
|
||||
|
||||
// For compatibility, we don't fetch from remote - that should be done by the caller
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
// Set stores a JWK set
|
||||
func (j *JWKCacheCompat) Set(jwksURL string, jwks *JWKSet, ttl time.Duration) error {
|
||||
return j.cache.Set(jwksURL, jwks, ttl)
|
||||
}
|
||||
|
||||
// Cleanup is a no-op for compatibility
|
||||
func (j *JWKCacheCompat) Cleanup() {}
|
||||
|
||||
// Close is a no-op for compatibility
|
||||
func (j *JWKCacheCompat) Close() {}
|
||||
|
||||
// CacheManagerCompat provides compatibility with the old CacheManager
|
||||
type CacheManagerCompat struct {
|
||||
manager *Manager
|
||||
mu sync.RWMutex
|
||||
}
|
||||
|
||||
// GetGlobalCacheManagerCompat returns a singleton CacheManager instance
|
||||
func GetGlobalCacheManagerCompat(wg *sync.WaitGroup) *CacheManagerCompat {
|
||||
return &CacheManagerCompat{
|
||||
manager: GetGlobalManager(nil),
|
||||
}
|
||||
}
|
||||
|
||||
// GetSharedTokenBlacklist returns the shared token blacklist cache
|
||||
func (c *CacheManagerCompat) GetSharedTokenBlacklist() *CompatibilityWrapper {
|
||||
c.mu.RLock()
|
||||
defer c.mu.RUnlock()
|
||||
return NewCompatibilityWrapper(c.manager.GetRawTokenCache())
|
||||
}
|
||||
|
||||
// GetSharedTokenCache returns the shared token cache
|
||||
func (c *CacheManagerCompat) GetSharedTokenCache() *TokenCacheCompat {
|
||||
c.mu.RLock()
|
||||
defer c.mu.RUnlock()
|
||||
return NewTokenCacheCompat()
|
||||
}
|
||||
|
||||
// GetSharedMetadataCache returns the shared metadata cache
|
||||
func (c *CacheManagerCompat) GetSharedMetadataCache() *MetadataCacheCompat {
|
||||
c.mu.RLock()
|
||||
defer c.mu.RUnlock()
|
||||
return NewMetadataCacheCompat(nil)
|
||||
}
|
||||
|
||||
// GetSharedJWKCache returns the shared JWK cache
|
||||
func (c *CacheManagerCompat) GetSharedJWKCache() *JWKCacheCompat {
|
||||
c.mu.RLock()
|
||||
defer c.mu.RUnlock()
|
||||
return NewJWKCacheCompat()
|
||||
}
|
||||
|
||||
// Close gracefully shuts down all cache components
|
||||
func (c *CacheManagerCompat) Close() error {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
return c.manager.Close()
|
||||
}
|
||||
|
||||
// UniversalCacheManagerCompat provides compatibility with UniversalCacheManager
|
||||
type UniversalCacheManagerCompat struct {
|
||||
manager *Manager
|
||||
logger Logger
|
||||
}
|
||||
|
||||
// GetUniversalCacheManagerCompat returns the global cache manager
|
||||
func GetUniversalCacheManagerCompat(logger Logger) *UniversalCacheManagerCompat {
|
||||
return &UniversalCacheManagerCompat{
|
||||
manager: GetGlobalManager(logger),
|
||||
logger: logger,
|
||||
}
|
||||
}
|
||||
|
||||
// GetTokenCache returns the token cache
|
||||
func (u *UniversalCacheManagerCompat) GetTokenCache() *UniversalCacheCompat {
|
||||
return &UniversalCacheCompat{
|
||||
Cache: u.manager.GetRawTokenCache(),
|
||||
}
|
||||
}
|
||||
|
||||
// GetMetadataCache returns the metadata cache
|
||||
func (u *UniversalCacheManagerCompat) GetMetadataCache() *UniversalCacheCompat {
|
||||
return &UniversalCacheCompat{
|
||||
Cache: u.manager.GetRawMetadataCache(),
|
||||
}
|
||||
}
|
||||
|
||||
// GetJWKCache returns the JWK cache
|
||||
func (u *UniversalCacheManagerCompat) GetJWKCache() *UniversalCacheCompat {
|
||||
return &UniversalCacheCompat{
|
||||
Cache: u.manager.GetRawJWKCache(),
|
||||
}
|
||||
}
|
||||
|
||||
// GetBlacklistCache returns the blacklist cache (uses token cache)
|
||||
func (u *UniversalCacheManagerCompat) GetBlacklistCache() *UniversalCacheCompat {
|
||||
return &UniversalCacheCompat{
|
||||
Cache: u.manager.GetRawTokenCache(),
|
||||
}
|
||||
}
|
||||
|
||||
// Close shuts down the cache manager
|
||||
func (u *UniversalCacheManagerCompat) Close() error {
|
||||
return u.manager.Close()
|
||||
}
|
||||
Vendored
+284
@@ -0,0 +1,284 @@
|
||||
package cache
|
||||
|
||||
import (
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Manager manages multiple cache instances with singleton pattern
|
||||
type Manager struct {
|
||||
mu sync.RWMutex
|
||||
|
||||
// Core caches
|
||||
tokenCache *Cache
|
||||
metadataCache *Cache
|
||||
jwkCache *Cache
|
||||
sessionCache *Cache
|
||||
generalCache *Cache
|
||||
|
||||
// Typed wrappers
|
||||
typedToken *TokenCache
|
||||
typedMetadata *MetadataCache
|
||||
typedJWK *JWKCache
|
||||
typedSession *SessionCache
|
||||
|
||||
logger Logger
|
||||
}
|
||||
|
||||
var (
|
||||
globalManager *Manager
|
||||
globalManagerOnce sync.Once
|
||||
)
|
||||
|
||||
// GetGlobalManager returns the singleton cache manager instance
|
||||
func GetGlobalManager(logger Logger) *Manager {
|
||||
globalManagerOnce.Do(func() {
|
||||
globalManager = NewManager(logger)
|
||||
})
|
||||
return globalManager
|
||||
}
|
||||
|
||||
// NewManager creates a new cache manager
|
||||
func NewManager(logger Logger) *Manager {
|
||||
if logger == nil {
|
||||
logger = &noOpLogger{}
|
||||
}
|
||||
|
||||
m := &Manager{
|
||||
logger: logger,
|
||||
}
|
||||
|
||||
// Initialize core caches with appropriate configurations
|
||||
m.initializeCaches()
|
||||
|
||||
return m
|
||||
}
|
||||
|
||||
// initializeCaches creates all cache instances with appropriate configurations
|
||||
func (m *Manager) initializeCaches() {
|
||||
// Token cache configuration
|
||||
tokenConfig := Config{
|
||||
Type: TypeToken,
|
||||
MaxSize: 5000,
|
||||
MaxMemoryBytes: 32 * 1024 * 1024, // 32MB
|
||||
DefaultTTL: 1 * time.Hour,
|
||||
CleanupInterval: 5 * time.Minute,
|
||||
EnableAutoCleanup: true,
|
||||
EnableMemoryLimit: true,
|
||||
EnableMetrics: true,
|
||||
Logger: m.logger,
|
||||
TokenConfig: &TokenConfig{
|
||||
BlacklistTTL: 24 * time.Hour,
|
||||
RefreshTokenTTL: 7 * 24 * time.Hour,
|
||||
EnableTokenRotation: true,
|
||||
},
|
||||
}
|
||||
m.tokenCache = New(tokenConfig)
|
||||
m.typedToken = NewTokenCache(m.tokenCache)
|
||||
|
||||
// Metadata cache configuration
|
||||
metadataConfig := Config{
|
||||
Type: TypeMetadata,
|
||||
MaxSize: 100,
|
||||
MaxMemoryBytes: 10 * 1024 * 1024, // 10MB
|
||||
DefaultTTL: 24 * time.Hour,
|
||||
CleanupInterval: 30 * time.Minute,
|
||||
EnableAutoCleanup: true,
|
||||
EnableMemoryLimit: true,
|
||||
EnableMetrics: true,
|
||||
Logger: m.logger,
|
||||
MetadataConfig: &MetadataConfig{
|
||||
GracePeriod: 5 * time.Minute,
|
||||
ExtendedGracePeriod: 15 * time.Minute,
|
||||
MaxGracePeriod: 1 * time.Hour,
|
||||
SecurityCriticalMaxGracePeriod: 30 * time.Minute,
|
||||
SecurityCriticalFields: []string{"issuer", "jwks_uri"},
|
||||
},
|
||||
}
|
||||
m.metadataCache = New(metadataConfig)
|
||||
m.typedMetadata = NewMetadataCache(m.metadataCache, *metadataConfig.MetadataConfig)
|
||||
|
||||
// JWK cache configuration
|
||||
jwkConfig := Config{
|
||||
Type: TypeJWK,
|
||||
MaxSize: 50,
|
||||
MaxMemoryBytes: 5 * 1024 * 1024, // 5MB
|
||||
DefaultTTL: 1 * time.Hour,
|
||||
CleanupInterval: 10 * time.Minute,
|
||||
EnableAutoCleanup: true,
|
||||
EnableMemoryLimit: true,
|
||||
EnableMetrics: true,
|
||||
Logger: m.logger,
|
||||
JWKConfig: &JWKConfig{
|
||||
RefreshInterval: 1 * time.Hour,
|
||||
MinRefreshTime: 5 * time.Minute,
|
||||
MaxKeyAge: 24 * time.Hour,
|
||||
},
|
||||
}
|
||||
m.jwkCache = New(jwkConfig)
|
||||
m.typedJWK = NewJWKCache(m.jwkCache)
|
||||
|
||||
// Session cache configuration
|
||||
sessionConfig := Config{
|
||||
Type: TypeSession,
|
||||
MaxSize: 10000,
|
||||
MaxMemoryBytes: 64 * 1024 * 1024, // 64MB
|
||||
DefaultTTL: 30 * time.Minute,
|
||||
CleanupInterval: 5 * time.Minute,
|
||||
EnableAutoCleanup: true,
|
||||
EnableMemoryLimit: true,
|
||||
EnableMetrics: true,
|
||||
Logger: m.logger,
|
||||
}
|
||||
m.sessionCache = New(sessionConfig)
|
||||
m.typedSession = NewSessionCache(m.sessionCache)
|
||||
|
||||
// General cache configuration
|
||||
generalConfig := Config{
|
||||
Type: TypeGeneral,
|
||||
MaxSize: 1000,
|
||||
MaxMemoryBytes: 16 * 1024 * 1024, // 16MB
|
||||
DefaultTTL: 10 * time.Minute,
|
||||
CleanupInterval: 5 * time.Minute,
|
||||
EnableAutoCleanup: true,
|
||||
EnableMemoryLimit: true,
|
||||
EnableMetrics: true,
|
||||
Logger: m.logger,
|
||||
}
|
||||
m.generalCache = New(generalConfig)
|
||||
}
|
||||
|
||||
// GetTokenCache returns the token cache instance
|
||||
func (m *Manager) GetTokenCache() *TokenCache {
|
||||
m.mu.RLock()
|
||||
defer m.mu.RUnlock()
|
||||
return m.typedToken
|
||||
}
|
||||
|
||||
// GetMetadataCache returns the metadata cache instance
|
||||
func (m *Manager) GetMetadataCache() *MetadataCache {
|
||||
m.mu.RLock()
|
||||
defer m.mu.RUnlock()
|
||||
return m.typedMetadata
|
||||
}
|
||||
|
||||
// GetJWKCache returns the JWK cache instance
|
||||
func (m *Manager) GetJWKCache() *JWKCache {
|
||||
m.mu.RLock()
|
||||
defer m.mu.RUnlock()
|
||||
return m.typedJWK
|
||||
}
|
||||
|
||||
// GetSessionCache returns the session cache instance
|
||||
func (m *Manager) GetSessionCache() *SessionCache {
|
||||
m.mu.RLock()
|
||||
defer m.mu.RUnlock()
|
||||
return m.typedSession
|
||||
}
|
||||
|
||||
// GetGeneralCache returns the general cache instance
|
||||
func (m *Manager) GetGeneralCache() *Cache {
|
||||
m.mu.RLock()
|
||||
defer m.mu.RUnlock()
|
||||
return m.generalCache
|
||||
}
|
||||
|
||||
// GetRawTokenCache returns the raw token cache for compatibility
|
||||
func (m *Manager) GetRawTokenCache() *Cache {
|
||||
m.mu.RLock()
|
||||
defer m.mu.RUnlock()
|
||||
return m.tokenCache
|
||||
}
|
||||
|
||||
// GetRawMetadataCache returns the raw metadata cache for compatibility
|
||||
func (m *Manager) GetRawMetadataCache() *Cache {
|
||||
m.mu.RLock()
|
||||
defer m.mu.RUnlock()
|
||||
return m.metadataCache
|
||||
}
|
||||
|
||||
// GetRawJWKCache returns the raw JWK cache for compatibility
|
||||
func (m *Manager) GetRawJWKCache() *Cache {
|
||||
m.mu.RLock()
|
||||
defer m.mu.RUnlock()
|
||||
return m.jwkCache
|
||||
}
|
||||
|
||||
// GetStats returns statistics for all caches
|
||||
func (m *Manager) GetStats() map[string]map[string]interface{} {
|
||||
m.mu.RLock()
|
||||
defer m.mu.RUnlock()
|
||||
|
||||
return map[string]map[string]interface{}{
|
||||
"token": m.tokenCache.GetStats(),
|
||||
"metadata": m.metadataCache.GetStats(),
|
||||
"jwk": m.jwkCache.GetStats(),
|
||||
"session": m.sessionCache.GetStats(),
|
||||
"general": m.generalCache.GetStats(),
|
||||
}
|
||||
}
|
||||
|
||||
// ClearAll clears all cache instances
|
||||
func (m *Manager) ClearAll() {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
|
||||
m.tokenCache.Clear()
|
||||
m.metadataCache.Clear()
|
||||
m.jwkCache.Clear()
|
||||
m.sessionCache.Clear()
|
||||
m.generalCache.Clear()
|
||||
}
|
||||
|
||||
// Close gracefully shuts down all cache instances
|
||||
func (m *Manager) Close() error {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
|
||||
var firstErr error
|
||||
|
||||
if err := m.tokenCache.Close(); err != nil && firstErr == nil {
|
||||
firstErr = err
|
||||
}
|
||||
if err := m.metadataCache.Close(); err != nil && firstErr == nil {
|
||||
firstErr = err
|
||||
}
|
||||
if err := m.jwkCache.Close(); err != nil && firstErr == nil {
|
||||
firstErr = err
|
||||
}
|
||||
if err := m.sessionCache.Close(); err != nil && firstErr == nil {
|
||||
firstErr = err
|
||||
}
|
||||
if err := m.generalCache.Close(); err != nil && firstErr == nil {
|
||||
firstErr = err
|
||||
}
|
||||
|
||||
return firstErr
|
||||
}
|
||||
|
||||
// CleanupAll runs cleanup on all cache instances
|
||||
func (m *Manager) CleanupAll() {
|
||||
m.mu.RLock()
|
||||
defer m.mu.RUnlock()
|
||||
|
||||
m.tokenCache.Cleanup()
|
||||
m.metadataCache.Cleanup()
|
||||
m.jwkCache.Cleanup()
|
||||
m.sessionCache.Cleanup()
|
||||
m.generalCache.Cleanup()
|
||||
}
|
||||
|
||||
// SetLogger updates the logger for all caches
|
||||
func (m *Manager) SetLogger(logger Logger) {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
|
||||
m.logger = logger
|
||||
if logger != nil {
|
||||
m.tokenCache.logger = logger
|
||||
m.metadataCache.logger = logger
|
||||
m.jwkCache.logger = logger
|
||||
m.sessionCache.logger = logger
|
||||
m.generalCache.logger = logger
|
||||
}
|
||||
}
|
||||
Vendored
+315
@@ -0,0 +1,315 @@
|
||||
package cache
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"time"
|
||||
)
|
||||
|
||||
// TypedCache provides a type-safe wrapper around Cache for specific types
|
||||
type TypedCache[T any] struct {
|
||||
cache *Cache
|
||||
prefix string
|
||||
}
|
||||
|
||||
// NewTypedCache creates a new typed cache wrapper
|
||||
func NewTypedCache[T any](cache *Cache, prefix string) *TypedCache[T] {
|
||||
return &TypedCache[T]{
|
||||
cache: cache,
|
||||
prefix: prefix,
|
||||
}
|
||||
}
|
||||
|
||||
// Set stores a typed value
|
||||
func (tc *TypedCache[T]) Set(key string, value T, ttl time.Duration) error {
|
||||
prefixedKey := tc.prefix + key
|
||||
return tc.cache.Set(prefixedKey, value, ttl)
|
||||
}
|
||||
|
||||
// Get retrieves a typed value
|
||||
func (tc *TypedCache[T]) Get(key string) (T, bool) {
|
||||
var zero T
|
||||
prefixedKey := tc.prefix + key
|
||||
|
||||
value, exists := tc.cache.Get(prefixedKey)
|
||||
if !exists {
|
||||
return zero, false
|
||||
}
|
||||
|
||||
// Try direct type assertion first
|
||||
if typedValue, ok := value.(T); ok {
|
||||
return typedValue, true
|
||||
}
|
||||
|
||||
// If that fails, try JSON marshaling/unmarshaling for complex types
|
||||
data, err := json.Marshal(value)
|
||||
if err != nil {
|
||||
return zero, false
|
||||
}
|
||||
|
||||
var result T
|
||||
if err := json.Unmarshal(data, &result); err != nil {
|
||||
return zero, false
|
||||
}
|
||||
|
||||
return result, true
|
||||
}
|
||||
|
||||
// Delete removes a typed value
|
||||
func (tc *TypedCache[T]) Delete(key string) {
|
||||
prefixedKey := tc.prefix + key
|
||||
tc.cache.Delete(prefixedKey)
|
||||
}
|
||||
|
||||
// Clear removes all items with the prefix
|
||||
func (tc *TypedCache[T]) Clear() {
|
||||
// Note: This clears the entire underlying cache
|
||||
// In a production system, you might want to implement prefix-based clearing
|
||||
tc.cache.Clear()
|
||||
}
|
||||
|
||||
// Size returns the size of the underlying cache
|
||||
func (tc *TypedCache[T]) Size() int {
|
||||
return tc.cache.Size()
|
||||
}
|
||||
|
||||
// TokenCache provides specialized caching for JWT tokens
|
||||
type TokenCache struct {
|
||||
cache *TypedCache[map[string]interface{}]
|
||||
}
|
||||
|
||||
// NewTokenCache creates a new token cache
|
||||
func NewTokenCache(baseCache *Cache) *TokenCache {
|
||||
return &TokenCache{
|
||||
cache: NewTypedCache[map[string]interface{}](baseCache, "token:"),
|
||||
}
|
||||
}
|
||||
|
||||
// Set stores parsed token claims
|
||||
func (tc *TokenCache) Set(token string, claims map[string]interface{}, expiration time.Duration) error {
|
||||
return tc.cache.Set(token, claims, expiration)
|
||||
}
|
||||
|
||||
// Get retrieves cached claims for a token
|
||||
func (tc *TokenCache) Get(token string) (map[string]interface{}, bool) {
|
||||
return tc.cache.Get(token)
|
||||
}
|
||||
|
||||
// Delete removes a token from cache
|
||||
func (tc *TokenCache) Delete(token string) {
|
||||
tc.cache.Delete(token)
|
||||
}
|
||||
|
||||
// SetBlacklisted marks a token as blacklisted
|
||||
func (tc *TokenCache) SetBlacklisted(token string, ttl time.Duration) error {
|
||||
blacklistKey := "blacklist:" + token
|
||||
// Store blacklisted status as a map to match the type
|
||||
blacklistData := map[string]interface{}{"blacklisted": true}
|
||||
return tc.cache.Set(blacklistKey, blacklistData, ttl)
|
||||
}
|
||||
|
||||
// IsBlacklisted checks if a token is blacklisted
|
||||
func (tc *TokenCache) IsBlacklisted(token string) bool {
|
||||
blacklistKey := "blacklist:" + token
|
||||
value, exists := tc.cache.Get(blacklistKey)
|
||||
if !exists {
|
||||
return false
|
||||
}
|
||||
// Check if the blacklist data indicates blacklisted status
|
||||
if data, ok := value["blacklisted"]; ok {
|
||||
blacklisted, _ := data.(bool)
|
||||
return blacklisted
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// MetadataCache provides specialized caching for provider metadata
|
||||
type MetadataCache struct {
|
||||
cache *Cache
|
||||
config MetadataConfig
|
||||
}
|
||||
|
||||
// ProviderMetadata represents OIDC provider metadata
|
||||
type ProviderMetadata struct {
|
||||
Issuer string `json:"issuer"`
|
||||
AuthorizationEndpoint string `json:"authorization_endpoint"`
|
||||
TokenEndpoint string `json:"token_endpoint"`
|
||||
UserInfoEndpoint string `json:"userinfo_endpoint"`
|
||||
JWKSUri string `json:"jwks_uri"`
|
||||
ScopesSupported []string `json:"scopes_supported"`
|
||||
}
|
||||
|
||||
// NewMetadataCache creates a new metadata cache
|
||||
func NewMetadataCache(baseCache *Cache, config MetadataConfig) *MetadataCache {
|
||||
return &MetadataCache{
|
||||
cache: baseCache,
|
||||
config: config,
|
||||
}
|
||||
}
|
||||
|
||||
// Set stores provider metadata with grace period support
|
||||
func (mc *MetadataCache) Set(providerURL string, metadata *ProviderMetadata, ttl time.Duration) error {
|
||||
if metadata == nil {
|
||||
return fmt.Errorf("metadata cannot be nil")
|
||||
}
|
||||
|
||||
key := "metadata:" + providerURL
|
||||
|
||||
// Apply grace period if configured
|
||||
if mc.config.GracePeriod > 0 {
|
||||
ttl += mc.config.GracePeriod
|
||||
}
|
||||
|
||||
// Store as JSON for consistency
|
||||
data, err := json.Marshal(metadata)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to marshal metadata: %w", err)
|
||||
}
|
||||
|
||||
return mc.cache.Set(key, data, ttl)
|
||||
}
|
||||
|
||||
// Get retrieves provider metadata from cache
|
||||
func (mc *MetadataCache) Get(providerURL string) (*ProviderMetadata, bool) {
|
||||
key := "metadata:" + providerURL
|
||||
value, exists := mc.cache.Get(key)
|
||||
if !exists {
|
||||
return nil, false
|
||||
}
|
||||
|
||||
// Handle different value types
|
||||
var data []byte
|
||||
switch v := value.(type) {
|
||||
case []byte:
|
||||
data = v
|
||||
case string:
|
||||
data = []byte(v)
|
||||
default:
|
||||
return nil, false
|
||||
}
|
||||
|
||||
var metadata ProviderMetadata
|
||||
if err := json.Unmarshal(data, &metadata); err != nil {
|
||||
return nil, false
|
||||
}
|
||||
|
||||
return &metadata, true
|
||||
}
|
||||
|
||||
// Delete removes provider metadata
|
||||
func (mc *MetadataCache) Delete(providerURL string) {
|
||||
key := "metadata:" + providerURL
|
||||
mc.cache.Delete(key)
|
||||
}
|
||||
|
||||
// JWKCache provides specialized caching for JWK sets
|
||||
type JWKCache struct {
|
||||
cache *Cache
|
||||
}
|
||||
|
||||
// JWKSet represents a set of JSON Web Keys
|
||||
type JWKSet struct {
|
||||
Keys []JWK `json:"keys"`
|
||||
}
|
||||
|
||||
// JWK represents a JSON Web Key
|
||||
type JWK struct {
|
||||
Kid string `json:"kid"`
|
||||
Kty string `json:"kty"`
|
||||
Use string `json:"use"`
|
||||
N string `json:"n"`
|
||||
E string `json:"e"`
|
||||
X5c []string `json:"x5c,omitempty"`
|
||||
}
|
||||
|
||||
// NewJWKCache creates a new JWK cache
|
||||
func NewJWKCache(baseCache *Cache) *JWKCache {
|
||||
return &JWKCache{
|
||||
cache: baseCache,
|
||||
}
|
||||
}
|
||||
|
||||
// Set stores a JWK set
|
||||
func (jc *JWKCache) Set(jwksURL string, jwks *JWKSet, ttl time.Duration) error {
|
||||
if jwks == nil {
|
||||
return fmt.Errorf("JWK set cannot be nil")
|
||||
}
|
||||
|
||||
key := "jwk:" + jwksURL
|
||||
return jc.cache.Set(key, jwks, ttl)
|
||||
}
|
||||
|
||||
// Get retrieves a JWK set from cache
|
||||
func (jc *JWKCache) Get(jwksURL string) (*JWKSet, bool) {
|
||||
key := "jwk:" + jwksURL
|
||||
value, exists := jc.cache.Get(key)
|
||||
if !exists {
|
||||
return nil, false
|
||||
}
|
||||
|
||||
jwks, ok := value.(*JWKSet)
|
||||
if !ok {
|
||||
// Try JSON conversion
|
||||
data, err := json.Marshal(value)
|
||||
if err != nil {
|
||||
return nil, false
|
||||
}
|
||||
|
||||
var result JWKSet
|
||||
if err := json.Unmarshal(data, &result); err != nil {
|
||||
return nil, false
|
||||
}
|
||||
return &result, true
|
||||
}
|
||||
|
||||
return jwks, true
|
||||
}
|
||||
|
||||
// Delete removes a JWK set from cache
|
||||
func (jc *JWKCache) Delete(jwksURL string) {
|
||||
key := "jwk:" + jwksURL
|
||||
jc.cache.Delete(key)
|
||||
}
|
||||
|
||||
// SessionCache provides specialized caching for sessions
|
||||
type SessionCache struct {
|
||||
cache *TypedCache[SessionData]
|
||||
}
|
||||
|
||||
// SessionData represents session information
|
||||
type SessionData struct {
|
||||
ID string `json:"id"`
|
||||
UserID string `json:"user_id"`
|
||||
AccessToken string `json:"access_token"`
|
||||
RefreshToken string `json:"refresh_token"`
|
||||
ExpiresAt time.Time `json:"expires_at"`
|
||||
Claims map[string]interface{} `json:"claims"`
|
||||
}
|
||||
|
||||
// NewSessionCache creates a new session cache
|
||||
func NewSessionCache(baseCache *Cache) *SessionCache {
|
||||
return &SessionCache{
|
||||
cache: NewTypedCache[SessionData](baseCache, "session:"),
|
||||
}
|
||||
}
|
||||
|
||||
// Set stores session data
|
||||
func (sc *SessionCache) Set(sessionID string, data SessionData, ttl time.Duration) error {
|
||||
return sc.cache.Set(sessionID, data, ttl)
|
||||
}
|
||||
|
||||
// Get retrieves session data
|
||||
func (sc *SessionCache) Get(sessionID string) (SessionData, bool) {
|
||||
return sc.cache.Get(sessionID)
|
||||
}
|
||||
|
||||
// Delete removes a session
|
||||
func (sc *SessionCache) Delete(sessionID string) {
|
||||
sc.cache.Delete(sessionID)
|
||||
}
|
||||
|
||||
// Exists checks if a session exists
|
||||
func (sc *SessionCache) Exists(sessionID string) bool {
|
||||
_, exists := sc.cache.Get(sessionID)
|
||||
return exists
|
||||
}
|
||||
@@ -0,0 +1,545 @@
|
||||
package httpclient
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/tls"
|
||||
"fmt"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/http/cookiejar"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Config provides configuration for creating HTTP clients
|
||||
type Config struct {
|
||||
// Timeout for the entire request
|
||||
Timeout time.Duration
|
||||
// MaxRedirects allowed (0 means follow Go's default of 10)
|
||||
MaxRedirects int
|
||||
// UseCookieJar enables cookie jar for the client
|
||||
UseCookieJar bool
|
||||
// Connection settings
|
||||
DialTimeout time.Duration
|
||||
KeepAlive time.Duration
|
||||
TLSHandshakeTimeout time.Duration
|
||||
ResponseHeaderTimeout time.Duration
|
||||
ExpectContinueTimeout time.Duration
|
||||
IdleConnTimeout time.Duration
|
||||
// Connection pool settings
|
||||
MaxIdleConns int
|
||||
MaxIdleConnsPerHost int
|
||||
MaxConnsPerHost int
|
||||
// Buffer settings
|
||||
WriteBufferSize int
|
||||
ReadBufferSize int
|
||||
// Feature flags
|
||||
ForceHTTP2 bool
|
||||
DisableKeepAlives bool
|
||||
DisableCompression bool
|
||||
// TLS configuration
|
||||
TLSConfig *tls.Config
|
||||
}
|
||||
|
||||
// ClientType defines the type of HTTP client for optimized behavior
|
||||
type ClientType string
|
||||
|
||||
const (
|
||||
ClientTypeDefault ClientType = "default"
|
||||
ClientTypeToken ClientType = "token"
|
||||
ClientTypeAPI ClientType = "api"
|
||||
ClientTypeProxy ClientType = "proxy"
|
||||
)
|
||||
|
||||
// PresetConfigs provides pre-configured settings for different client types
|
||||
var PresetConfigs = map[ClientType]Config{
|
||||
ClientTypeDefault: {
|
||||
Timeout: 10 * time.Second, // Reduced from 30s to prevent slowloris attacks
|
||||
MaxRedirects: 5, // Reduced from 10 to prevent redirect loops
|
||||
UseCookieJar: false,
|
||||
DialTimeout: 3 * time.Second,
|
||||
KeepAlive: 15 * time.Second,
|
||||
TLSHandshakeTimeout: 2 * time.Second,
|
||||
ResponseHeaderTimeout: 3 * time.Second,
|
||||
ExpectContinueTimeout: 1 * time.Second,
|
||||
IdleConnTimeout: 5 * time.Second,
|
||||
MaxIdleConns: 20, // Reduced from 100 to limit resource usage
|
||||
MaxIdleConnsPerHost: 2, // Reduced from 10 to prevent connection exhaustion
|
||||
MaxConnsPerHost: 5, // Reduced from 10 to limit concurrent connections
|
||||
WriteBufferSize: 4096,
|
||||
ReadBufferSize: 4096,
|
||||
ForceHTTP2: true,
|
||||
DisableKeepAlives: false,
|
||||
DisableCompression: false,
|
||||
},
|
||||
ClientTypeToken: {
|
||||
Timeout: 10 * time.Second,
|
||||
MaxRedirects: 50, // Token endpoints may redirect more
|
||||
UseCookieJar: true,
|
||||
DialTimeout: 3 * time.Second,
|
||||
KeepAlive: 15 * time.Second,
|
||||
TLSHandshakeTimeout: 2 * time.Second,
|
||||
ResponseHeaderTimeout: 3 * time.Second,
|
||||
ExpectContinueTimeout: 1 * time.Second,
|
||||
IdleConnTimeout: 5 * time.Second,
|
||||
MaxIdleConns: 10,
|
||||
MaxIdleConnsPerHost: 2,
|
||||
MaxConnsPerHost: 5,
|
||||
WriteBufferSize: 4096,
|
||||
ReadBufferSize: 4096,
|
||||
ForceHTTP2: true,
|
||||
DisableKeepAlives: false,
|
||||
DisableCompression: false,
|
||||
},
|
||||
ClientTypeAPI: {
|
||||
Timeout: 30 * time.Second, // Longer for API operations
|
||||
MaxRedirects: 10,
|
||||
UseCookieJar: false,
|
||||
DialTimeout: 5 * time.Second,
|
||||
KeepAlive: 30 * time.Second,
|
||||
TLSHandshakeTimeout: 5 * time.Second,
|
||||
ResponseHeaderTimeout: 10 * time.Second,
|
||||
ExpectContinueTimeout: 1 * time.Second,
|
||||
IdleConnTimeout: 90 * time.Second,
|
||||
MaxIdleConns: 50,
|
||||
MaxIdleConnsPerHost: 5,
|
||||
MaxConnsPerHost: 10,
|
||||
WriteBufferSize: 8192,
|
||||
ReadBufferSize: 8192,
|
||||
ForceHTTP2: true,
|
||||
DisableKeepAlives: false,
|
||||
DisableCompression: false,
|
||||
},
|
||||
ClientTypeProxy: {
|
||||
Timeout: 60 * time.Second, // Proxy needs longer timeouts
|
||||
MaxRedirects: 0, // Proxy should not follow redirects
|
||||
UseCookieJar: false,
|
||||
DialTimeout: 10 * time.Second,
|
||||
KeepAlive: 30 * time.Second,
|
||||
TLSHandshakeTimeout: 5 * time.Second,
|
||||
ResponseHeaderTimeout: 10 * time.Second,
|
||||
ExpectContinueTimeout: 1 * time.Second,
|
||||
IdleConnTimeout: 90 * time.Second,
|
||||
MaxIdleConns: 100,
|
||||
MaxIdleConnsPerHost: 10,
|
||||
MaxConnsPerHost: 20,
|
||||
WriteBufferSize: 16384,
|
||||
ReadBufferSize: 16384,
|
||||
ForceHTTP2: true,
|
||||
DisableKeepAlives: false,
|
||||
DisableCompression: true, // Proxy should not modify content
|
||||
},
|
||||
}
|
||||
|
||||
// Factory provides methods for creating configured HTTP clients
|
||||
type Factory struct {
|
||||
pool *TransportPool
|
||||
logger Logger
|
||||
}
|
||||
|
||||
// Logger interface for HTTP client operations
|
||||
type Logger interface {
|
||||
Debug(msg string)
|
||||
Debugf(format string, args ...interface{})
|
||||
Info(msg string)
|
||||
Infof(format string, args ...interface{})
|
||||
Error(msg string)
|
||||
Errorf(format string, args ...interface{})
|
||||
}
|
||||
|
||||
var (
|
||||
globalFactory *Factory
|
||||
globalFactoryOnce sync.Once
|
||||
)
|
||||
|
||||
// GetGlobalFactory returns the singleton HTTP client factory
|
||||
func GetGlobalFactory(logger Logger) *Factory {
|
||||
globalFactoryOnce.Do(func() {
|
||||
globalFactory = NewFactory(logger)
|
||||
})
|
||||
return globalFactory
|
||||
}
|
||||
|
||||
// NewFactory creates a new HTTP client factory
|
||||
func NewFactory(logger Logger) *Factory {
|
||||
if logger == nil {
|
||||
logger = &noOpLogger{}
|
||||
}
|
||||
return &Factory{
|
||||
pool: GetGlobalTransportPool(),
|
||||
logger: logger,
|
||||
}
|
||||
}
|
||||
|
||||
// CreateClient creates an HTTP client with the specified configuration
|
||||
func (f *Factory) CreateClient(config Config) (*http.Client, error) {
|
||||
// Validate configuration
|
||||
if err := f.ValidateConfig(&config); err != nil {
|
||||
return nil, fmt.Errorf("invalid configuration: %w", err)
|
||||
}
|
||||
|
||||
// Apply TLS configuration if not provided
|
||||
if config.TLSConfig == nil {
|
||||
config.TLSConfig = f.createSecureTLSConfig()
|
||||
}
|
||||
|
||||
// Get or create transport from pool
|
||||
transport := f.pool.GetOrCreateTransport(config)
|
||||
if transport == nil {
|
||||
return nil, fmt.Errorf("failed to create transport: client limit exceeded")
|
||||
}
|
||||
|
||||
// Create HTTP client
|
||||
client := &http.Client{
|
||||
Transport: transport,
|
||||
Timeout: config.Timeout,
|
||||
}
|
||||
|
||||
// Configure redirect policy
|
||||
if config.MaxRedirects > 0 {
|
||||
client.CheckRedirect = func(req *http.Request, via []*http.Request) error {
|
||||
if len(via) >= config.MaxRedirects {
|
||||
return fmt.Errorf("stopped after %d redirects", config.MaxRedirects)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
// Add cookie jar if requested
|
||||
if config.UseCookieJar {
|
||||
jar, err := cookiejar.New(nil)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create cookie jar: %w", err)
|
||||
}
|
||||
client.Jar = jar
|
||||
}
|
||||
|
||||
f.logger.Debugf("Created HTTP client with config: timeout=%v, maxRedirects=%d", config.Timeout, config.MaxRedirects)
|
||||
return client, nil
|
||||
}
|
||||
|
||||
// CreateClientWithPreset creates an HTTP client using a preset configuration
|
||||
func (f *Factory) CreateClientWithPreset(clientType ClientType) (*http.Client, error) {
|
||||
config, ok := PresetConfigs[clientType]
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("unknown client type: %s", clientType)
|
||||
}
|
||||
return f.CreateClient(config)
|
||||
}
|
||||
|
||||
// CreateDefault creates a default HTTP client
|
||||
func (f *Factory) CreateDefault() (*http.Client, error) {
|
||||
return f.CreateClientWithPreset(ClientTypeDefault)
|
||||
}
|
||||
|
||||
// CreateToken creates an HTTP client optimized for token operations
|
||||
func (f *Factory) CreateToken() (*http.Client, error) {
|
||||
return f.CreateClientWithPreset(ClientTypeToken)
|
||||
}
|
||||
|
||||
// CreateAPI creates an HTTP client optimized for API operations
|
||||
func (f *Factory) CreateAPI() (*http.Client, error) {
|
||||
return f.CreateClientWithPreset(ClientTypeAPI)
|
||||
}
|
||||
|
||||
// CreateProxy creates an HTTP client optimized for proxy operations
|
||||
func (f *Factory) CreateProxy() (*http.Client, error) {
|
||||
return f.CreateClientWithPreset(ClientTypeProxy)
|
||||
}
|
||||
|
||||
// ValidateConfig validates HTTP client configuration parameters
|
||||
func (f *Factory) ValidateConfig(config *Config) error {
|
||||
// Validate connection pool limits
|
||||
if config.MaxIdleConns < 0 {
|
||||
return fmt.Errorf("MaxIdleConns cannot be negative: %d", config.MaxIdleConns)
|
||||
}
|
||||
if config.MaxIdleConns > 1000 {
|
||||
return fmt.Errorf("MaxIdleConns too high (max 1000): %d", config.MaxIdleConns)
|
||||
}
|
||||
|
||||
if config.MaxIdleConnsPerHost < 0 {
|
||||
return fmt.Errorf("MaxIdleConnsPerHost cannot be negative: %d", config.MaxIdleConnsPerHost)
|
||||
}
|
||||
if config.MaxIdleConnsPerHost > 100 {
|
||||
return fmt.Errorf("MaxIdleConnsPerHost too high (max 100): %d", config.MaxIdleConnsPerHost)
|
||||
}
|
||||
|
||||
if config.MaxConnsPerHost < 0 {
|
||||
return fmt.Errorf("MaxConnsPerHost cannot be negative: %d", config.MaxConnsPerHost)
|
||||
}
|
||||
if config.MaxConnsPerHost > 200 {
|
||||
return fmt.Errorf("MaxConnsPerHost too high (max 200): %d", config.MaxConnsPerHost)
|
||||
}
|
||||
|
||||
// Validate timeouts
|
||||
if config.Timeout < 0 {
|
||||
return fmt.Errorf("timeout cannot be negative")
|
||||
}
|
||||
if config.Timeout > 5*time.Minute {
|
||||
return fmt.Errorf("timeout too long (max 5 minutes): %v", config.Timeout)
|
||||
}
|
||||
|
||||
// Validate buffer sizes
|
||||
if config.WriteBufferSize < 0 || config.ReadBufferSize < 0 {
|
||||
return fmt.Errorf("buffer sizes cannot be negative")
|
||||
}
|
||||
if config.WriteBufferSize > 1024*1024 || config.ReadBufferSize > 1024*1024 {
|
||||
return fmt.Errorf("buffer sizes too large (max 1MB)")
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// createSecureTLSConfig creates a secure TLS configuration
|
||||
func (f *Factory) createSecureTLSConfig() *tls.Config {
|
||||
return &tls.Config{
|
||||
MinVersion: tls.VersionTLS12, // SECURITY: Enforce TLS 1.2 minimum
|
||||
MaxVersion: tls.VersionTLS13, // Support up to TLS 1.3
|
||||
CipherSuites: []uint16{
|
||||
// TLS 1.3 cipher suites (automatically selected when TLS 1.3 is negotiated)
|
||||
// TLS 1.2 secure cipher suites
|
||||
tls.TLS_ECDHE_ECDSA_WITH_AES_256_GCM_SHA384,
|
||||
tls.TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384,
|
||||
tls.TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256,
|
||||
tls.TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256,
|
||||
},
|
||||
InsecureSkipVerify: false, // SECURITY: Always verify certificates
|
||||
PreferServerCipherSuites: false, // Let client choose best cipher
|
||||
}
|
||||
}
|
||||
|
||||
// TransportPool manages a pool of shared HTTP transports
|
||||
type TransportPool struct {
|
||||
mu sync.RWMutex
|
||||
transports map[string]*sharedTransport
|
||||
maxConns int
|
||||
ctx context.Context
|
||||
cancel context.CancelFunc
|
||||
|
||||
// Resource limits
|
||||
clientCount int32 // Track total HTTP clients
|
||||
maxClients int32 // Limit total clients
|
||||
}
|
||||
|
||||
type sharedTransport struct {
|
||||
transport *http.Transport
|
||||
refCount int32
|
||||
lastUsed time.Time
|
||||
config Config
|
||||
}
|
||||
|
||||
var (
|
||||
globalTransportPool *TransportPool
|
||||
globalTransportPoolOnce sync.Once
|
||||
)
|
||||
|
||||
// GetGlobalTransportPool returns the singleton transport pool instance
|
||||
func GetGlobalTransportPool() *TransportPool {
|
||||
globalTransportPoolOnce.Do(func() {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
globalTransportPool = &TransportPool{
|
||||
transports: make(map[string]*sharedTransport),
|
||||
maxConns: 20, // Reduced from 100 to prevent resource exhaustion
|
||||
ctx: ctx,
|
||||
cancel: cancel,
|
||||
clientCount: 0,
|
||||
maxClients: 5, // Maximum 5 HTTP clients
|
||||
}
|
||||
// Start cleanup goroutine with context cancellation
|
||||
go globalTransportPool.cleanupIdleTransports(ctx)
|
||||
})
|
||||
return globalTransportPool
|
||||
}
|
||||
|
||||
// GetOrCreateTransport gets or creates a shared transport with the given config
|
||||
func (p *TransportPool) GetOrCreateTransport(config Config) *http.Transport {
|
||||
// Check client limit before creating new transport
|
||||
if atomic.LoadInt32(&p.clientCount) >= p.maxClients {
|
||||
// Try to return existing transport if limit reached
|
||||
p.mu.RLock()
|
||||
defer p.mu.RUnlock()
|
||||
for _, shared := range p.transports {
|
||||
if shared != nil && shared.transport != nil {
|
||||
atomic.AddInt32(&shared.refCount, 1)
|
||||
shared.lastUsed = time.Now()
|
||||
return shared.transport
|
||||
}
|
||||
}
|
||||
// If no transport available, return nil
|
||||
return nil
|
||||
}
|
||||
|
||||
p.mu.Lock()
|
||||
defer p.mu.Unlock()
|
||||
|
||||
key := p.configKey(config)
|
||||
|
||||
if shared, exists := p.transports[key]; exists {
|
||||
atomic.AddInt32(&shared.refCount, 1)
|
||||
shared.lastUsed = time.Now()
|
||||
return shared.transport
|
||||
}
|
||||
|
||||
// Create new transport
|
||||
transport := p.createTransport(config)
|
||||
|
||||
p.transports[key] = &sharedTransport{
|
||||
transport: transport,
|
||||
refCount: 1,
|
||||
lastUsed: time.Now(),
|
||||
config: config,
|
||||
}
|
||||
|
||||
atomic.AddInt32(&p.clientCount, 1)
|
||||
return transport
|
||||
}
|
||||
|
||||
// createTransport creates a new HTTP transport with the given configuration
|
||||
func (p *TransportPool) createTransport(config Config) *http.Transport {
|
||||
// Create secure TLS config if not provided
|
||||
tlsConfig := config.TLSConfig
|
||||
if tlsConfig == nil {
|
||||
tlsConfig = &tls.Config{
|
||||
MinVersion: tls.VersionTLS12,
|
||||
MaxVersion: tls.VersionTLS13,
|
||||
}
|
||||
}
|
||||
|
||||
return &http.Transport{
|
||||
Proxy: http.ProxyFromEnvironment,
|
||||
DialContext: (&net.Dialer{
|
||||
Timeout: config.DialTimeout,
|
||||
KeepAlive: config.KeepAlive,
|
||||
}).DialContext,
|
||||
TLSClientConfig: tlsConfig,
|
||||
TLSHandshakeTimeout: config.TLSHandshakeTimeout,
|
||||
ResponseHeaderTimeout: config.ResponseHeaderTimeout,
|
||||
ExpectContinueTimeout: config.ExpectContinueTimeout,
|
||||
IdleConnTimeout: config.IdleConnTimeout,
|
||||
MaxIdleConns: config.MaxIdleConns,
|
||||
MaxIdleConnsPerHost: config.MaxIdleConnsPerHost,
|
||||
MaxConnsPerHost: config.MaxConnsPerHost,
|
||||
WriteBufferSize: config.WriteBufferSize,
|
||||
ReadBufferSize: config.ReadBufferSize,
|
||||
ForceAttemptHTTP2: config.ForceHTTP2,
|
||||
DisableKeepAlives: config.DisableKeepAlives,
|
||||
DisableCompression: config.DisableCompression,
|
||||
}
|
||||
}
|
||||
|
||||
// configKey generates a unique key for the configuration
|
||||
func (p *TransportPool) configKey(config Config) string {
|
||||
return fmt.Sprintf("%v-%d-%d-%d-%d-%v-%v-%v",
|
||||
config.Timeout,
|
||||
config.MaxIdleConns,
|
||||
config.MaxIdleConnsPerHost,
|
||||
config.MaxConnsPerHost,
|
||||
config.MaxRedirects,
|
||||
config.ForceHTTP2,
|
||||
config.DisableKeepAlives,
|
||||
config.DisableCompression,
|
||||
)
|
||||
}
|
||||
|
||||
// cleanupIdleTransports periodically cleans up idle transports
|
||||
func (p *TransportPool) cleanupIdleTransports(ctx context.Context) {
|
||||
ticker := time.NewTicker(5 * time.Minute)
|
||||
defer ticker.Stop()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-ticker.C:
|
||||
p.cleanupIdle()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// cleanupIdle removes idle transports with zero references
|
||||
func (p *TransportPool) cleanupIdle() {
|
||||
p.mu.Lock()
|
||||
defer p.mu.Unlock()
|
||||
|
||||
now := time.Now()
|
||||
var toRemove []string
|
||||
|
||||
for key, shared := range p.transports {
|
||||
if atomic.LoadInt32(&shared.refCount) == 0 && now.Sub(shared.lastUsed) > 10*time.Minute {
|
||||
if shared.transport != nil {
|
||||
shared.transport.CloseIdleConnections()
|
||||
}
|
||||
toRemove = append(toRemove, key)
|
||||
}
|
||||
}
|
||||
|
||||
for _, key := range toRemove {
|
||||
delete(p.transports, key)
|
||||
atomic.AddInt32(&p.clientCount, -1)
|
||||
}
|
||||
}
|
||||
|
||||
// Release decrements the reference count for a transport
|
||||
func (p *TransportPool) Release(transport *http.Transport) {
|
||||
p.mu.RLock()
|
||||
defer p.mu.RUnlock()
|
||||
|
||||
for _, shared := range p.transports {
|
||||
if shared.transport == transport {
|
||||
atomic.AddInt32(&shared.refCount, -1)
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Close shuts down the transport pool
|
||||
func (p *TransportPool) Close() error {
|
||||
p.cancel()
|
||||
|
||||
p.mu.Lock()
|
||||
defer p.mu.Unlock()
|
||||
|
||||
for key, shared := range p.transports {
|
||||
if shared.transport != nil {
|
||||
shared.transport.CloseIdleConnections()
|
||||
}
|
||||
delete(p.transports, key)
|
||||
}
|
||||
|
||||
atomic.StoreInt32(&p.clientCount, 0)
|
||||
return nil
|
||||
}
|
||||
|
||||
// noOpLogger provides a no-op logger implementation
|
||||
type noOpLogger struct{}
|
||||
|
||||
func (l *noOpLogger) Debug(msg string) {}
|
||||
func (l *noOpLogger) Debugf(format string, args ...interface{}) {}
|
||||
func (l *noOpLogger) Info(msg string) {}
|
||||
func (l *noOpLogger) Infof(format string, args ...interface{}) {}
|
||||
func (l *noOpLogger) Error(msg string) {}
|
||||
func (l *noOpLogger) Errorf(format string, args ...interface{}) {}
|
||||
|
||||
// Compatibility functions for backward compatibility
|
||||
|
||||
// CreateDefaultHTTPClient creates a default HTTP client
|
||||
func CreateDefaultHTTPClient() *http.Client {
|
||||
factory := GetGlobalFactory(nil)
|
||||
client, _ := factory.CreateDefault()
|
||||
return client
|
||||
}
|
||||
|
||||
// CreateTokenHTTPClient creates an HTTP client optimized for token operations
|
||||
func CreateTokenHTTPClient() *http.Client {
|
||||
factory := GetGlobalFactory(nil)
|
||||
client, _ := factory.CreateToken()
|
||||
return client
|
||||
}
|
||||
|
||||
// CreateHTTPClientWithConfig creates an HTTP client with custom configuration
|
||||
func CreateHTTPClientWithConfig(config Config) *http.Client {
|
||||
factory := GetGlobalFactory(nil)
|
||||
client, _ := factory.CreateClient(config)
|
||||
return client
|
||||
}
|
||||
@@ -0,0 +1,299 @@
|
||||
package httpclient
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestFactoryCreateClient(t *testing.T) {
|
||||
factory := NewFactory(nil)
|
||||
|
||||
// Test creating default client
|
||||
client, err := factory.CreateDefault()
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create default client: %v", err)
|
||||
}
|
||||
if client == nil {
|
||||
t.Fatal("Expected non-nil client")
|
||||
}
|
||||
|
||||
// Test creating token client
|
||||
tokenClient, err := factory.CreateToken()
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create token client: %v", err)
|
||||
}
|
||||
if tokenClient == nil {
|
||||
t.Fatal("Expected non-nil token client")
|
||||
}
|
||||
}
|
||||
|
||||
func TestFactoryCreateClientWithPreset(t *testing.T) {
|
||||
factory := NewFactory(nil)
|
||||
|
||||
testCases := []struct {
|
||||
name string
|
||||
clientType ClientType
|
||||
shouldFail bool
|
||||
}{
|
||||
{"Default", ClientTypeDefault, false},
|
||||
{"Token", ClientTypeToken, false},
|
||||
{"API", ClientTypeAPI, false},
|
||||
{"Proxy", ClientTypeProxy, false},
|
||||
{"Invalid", ClientType("invalid"), true},
|
||||
}
|
||||
|
||||
for _, tc := range testCases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
client, err := factory.CreateClientWithPreset(tc.clientType)
|
||||
if tc.shouldFail {
|
||||
if err == nil {
|
||||
t.Fatal("Expected error for invalid client type")
|
||||
}
|
||||
} else {
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create %s client: %v", tc.clientType, err)
|
||||
}
|
||||
if client == nil {
|
||||
t.Fatal("Expected non-nil client")
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestFactoryValidateConfig(t *testing.T) {
|
||||
factory := NewFactory(nil)
|
||||
|
||||
testCases := []struct {
|
||||
name string
|
||||
config Config
|
||||
shouldFail bool
|
||||
}{
|
||||
{
|
||||
name: "Valid config",
|
||||
config: PresetConfigs[ClientTypeDefault],
|
||||
shouldFail: false,
|
||||
},
|
||||
{
|
||||
name: "Negative MaxIdleConns",
|
||||
config: Config{
|
||||
MaxIdleConns: -1,
|
||||
},
|
||||
shouldFail: true,
|
||||
},
|
||||
{
|
||||
name: "Excessive MaxIdleConns",
|
||||
config: Config{
|
||||
MaxIdleConns: 2000,
|
||||
},
|
||||
shouldFail: true,
|
||||
},
|
||||
{
|
||||
name: "Negative timeout",
|
||||
config: Config{
|
||||
Timeout: -1 * time.Second,
|
||||
},
|
||||
shouldFail: true,
|
||||
},
|
||||
{
|
||||
name: "Excessive timeout",
|
||||
config: Config{
|
||||
Timeout: 10 * time.Minute,
|
||||
},
|
||||
shouldFail: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tc := range testCases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
err := factory.ValidateConfig(&tc.config)
|
||||
if tc.shouldFail && err == nil {
|
||||
t.Fatal("Expected validation to fail")
|
||||
}
|
||||
if !tc.shouldFail && err != nil {
|
||||
t.Fatalf("Unexpected validation error: %v", err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestTransportPoolConcurrency(t *testing.T) {
|
||||
pool := &TransportPool{
|
||||
transports: make(map[string]*sharedTransport),
|
||||
maxConns: 20,
|
||||
clientCount: 0,
|
||||
maxClients: 5,
|
||||
}
|
||||
|
||||
config := PresetConfigs[ClientTypeDefault]
|
||||
|
||||
var wg sync.WaitGroup
|
||||
numGoroutines := 10
|
||||
|
||||
// Test concurrent transport creation
|
||||
wg.Add(numGoroutines)
|
||||
for i := 0; i < numGoroutines; i++ {
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
transport := pool.GetOrCreateTransport(config)
|
||||
if transport != nil {
|
||||
// Simulate usage
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
pool.Release(transport)
|
||||
}
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
|
||||
// Verify client count is within limits
|
||||
clientCount := atomic.LoadInt32(&pool.clientCount)
|
||||
if clientCount > pool.maxClients {
|
||||
t.Fatalf("Client count %d exceeds max %d", clientCount, pool.maxClients)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHTTPClientRequests(t *testing.T) {
|
||||
// Create test server
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.WriteHeader(http.StatusOK)
|
||||
w.Write([]byte("test response"))
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
factory := NewFactory(nil)
|
||||
client, err := factory.CreateDefault()
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create client: %v", err)
|
||||
}
|
||||
|
||||
// Make request
|
||||
resp, err := client.Get(server.URL)
|
||||
if err != nil {
|
||||
t.Fatalf("Request failed: %v", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
t.Fatalf("Expected status 200, got %d", resp.StatusCode)
|
||||
}
|
||||
}
|
||||
|
||||
func TestClientWithCookieJar(t *testing.T) {
|
||||
config := PresetConfigs[ClientTypeToken]
|
||||
if !config.UseCookieJar {
|
||||
t.Skip("Token client should have cookie jar enabled")
|
||||
}
|
||||
|
||||
factory := NewFactory(nil)
|
||||
client, err := factory.CreateToken()
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create token client: %v", err)
|
||||
}
|
||||
|
||||
if client.Jar == nil {
|
||||
t.Fatal("Expected cookie jar to be set")
|
||||
}
|
||||
}
|
||||
|
||||
func TestTransportPoolCleanup(t *testing.T) {
|
||||
pool := &TransportPool{
|
||||
transports: make(map[string]*sharedTransport),
|
||||
maxConns: 20,
|
||||
clientCount: 0,
|
||||
maxClients: 5,
|
||||
}
|
||||
|
||||
config := PresetConfigs[ClientTypeDefault]
|
||||
|
||||
// Create transport
|
||||
transport := pool.GetOrCreateTransport(config)
|
||||
if transport == nil {
|
||||
t.Fatal("Failed to create transport")
|
||||
}
|
||||
|
||||
// Release transport
|
||||
pool.Release(transport)
|
||||
|
||||
// Simulate idle time
|
||||
pool.mu.Lock()
|
||||
for _, shared := range pool.transports {
|
||||
shared.lastUsed = time.Now().Add(-11 * time.Minute)
|
||||
atomic.StoreInt32(&shared.refCount, 0)
|
||||
}
|
||||
pool.mu.Unlock()
|
||||
|
||||
// Run cleanup
|
||||
pool.cleanupIdle()
|
||||
|
||||
// Verify transport was removed
|
||||
pool.mu.RLock()
|
||||
count := len(pool.transports)
|
||||
pool.mu.RUnlock()
|
||||
|
||||
if count != 0 {
|
||||
t.Fatalf("Expected 0 transports after cleanup, got %d", count)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGlobalFactorySingleton(t *testing.T) {
|
||||
factory1 := GetGlobalFactory(nil)
|
||||
factory2 := GetGlobalFactory(nil)
|
||||
|
||||
if factory1 != factory2 {
|
||||
t.Fatal("Expected singleton factory instances to be the same")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCompatibilityFunctions(t *testing.T) {
|
||||
// Test CreateDefaultHTTPClient
|
||||
defaultClient := CreateDefaultHTTPClient()
|
||||
if defaultClient == nil {
|
||||
t.Fatal("Expected non-nil default client")
|
||||
}
|
||||
|
||||
// Test CreateTokenHTTPClient
|
||||
tokenClient := CreateTokenHTTPClient()
|
||||
if tokenClient == nil {
|
||||
t.Fatal("Expected non-nil token client")
|
||||
}
|
||||
|
||||
// Test CreateHTTPClientWithConfig
|
||||
config := PresetConfigs[ClientTypeAPI]
|
||||
apiClient := CreateHTTPClientWithConfig(config)
|
||||
if apiClient == nil {
|
||||
t.Fatal("Expected non-nil API client")
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkFactoryCreateClient(b *testing.B) {
|
||||
factory := NewFactory(nil)
|
||||
|
||||
b.ResetTimer()
|
||||
b.RunParallel(func(pb *testing.PB) {
|
||||
for pb.Next() {
|
||||
client, err := factory.CreateDefault()
|
||||
if err != nil || client == nil {
|
||||
b.Fatal("Failed to create client")
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func BenchmarkTransportPoolGetOrCreate(b *testing.B) {
|
||||
pool := GetGlobalTransportPool()
|
||||
config := PresetConfigs[ClientTypeDefault]
|
||||
|
||||
b.ResetTimer()
|
||||
b.RunParallel(func(pb *testing.PB) {
|
||||
for pb.Next() {
|
||||
transport := pool.GetOrCreateTransport(config)
|
||||
if transport != nil {
|
||||
pool.Release(transport)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,83 @@
|
||||
package logger
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"log"
|
||||
)
|
||||
|
||||
// LegacyLoggerAdapter wraps the old Logger struct from the main package
|
||||
// to implement the new unified Logger interface. This allows for gradual
|
||||
// migration of the codebase to the new logger interface.
|
||||
type LegacyLoggerAdapter struct {
|
||||
logError *log.Logger
|
||||
logInfo *log.Logger
|
||||
logDebug *log.Logger
|
||||
}
|
||||
|
||||
// NewLegacyAdapter creates a new adapter from the old logger components
|
||||
func NewLegacyAdapter(logError, logInfo, logDebug *log.Logger) Logger {
|
||||
if logError == nil || logInfo == nil || logDebug == nil {
|
||||
return GetNoOpLogger()
|
||||
}
|
||||
return &LegacyLoggerAdapter{
|
||||
logError: logError,
|
||||
logInfo: logInfo,
|
||||
logDebug: logDebug,
|
||||
}
|
||||
}
|
||||
|
||||
// Debug logs a debug message
|
||||
func (l *LegacyLoggerAdapter) Debug(msg string) {
|
||||
l.logDebug.Print(msg)
|
||||
}
|
||||
|
||||
// Debugf logs a formatted debug message
|
||||
func (l *LegacyLoggerAdapter) Debugf(format string, args ...interface{}) {
|
||||
l.logDebug.Printf(format, args...)
|
||||
}
|
||||
|
||||
// Info logs an info message
|
||||
func (l *LegacyLoggerAdapter) Info(msg string) {
|
||||
l.logInfo.Print(msg)
|
||||
}
|
||||
|
||||
// Infof logs a formatted info message
|
||||
func (l *LegacyLoggerAdapter) Infof(format string, args ...interface{}) {
|
||||
l.logInfo.Printf(format, args...)
|
||||
}
|
||||
|
||||
// Error logs an error message
|
||||
func (l *LegacyLoggerAdapter) Error(msg string) {
|
||||
l.logError.Print(msg)
|
||||
}
|
||||
|
||||
// Errorf logs a formatted error message
|
||||
func (l *LegacyLoggerAdapter) Errorf(format string, args ...interface{}) {
|
||||
l.logError.Printf(format, args...)
|
||||
}
|
||||
|
||||
// Printf logs a formatted message at info level
|
||||
func (l *LegacyLoggerAdapter) Printf(format string, args ...interface{}) {
|
||||
l.logInfo.Printf(format, args...)
|
||||
}
|
||||
|
||||
// Println logs a message at info level
|
||||
func (l *LegacyLoggerAdapter) Println(args ...interface{}) {
|
||||
l.logInfo.Print(args...)
|
||||
}
|
||||
|
||||
// Fatalf logs a formatted error message and panics
|
||||
func (l *LegacyLoggerAdapter) Fatalf(format string, args ...interface{}) {
|
||||
l.logError.Printf(format, args...)
|
||||
panic(fmt.Sprintf(format, args...))
|
||||
}
|
||||
|
||||
// WithField returns the same logger (no structured logging support in legacy adapter)
|
||||
func (l *LegacyLoggerAdapter) WithField(key string, value interface{}) Logger {
|
||||
return l
|
||||
}
|
||||
|
||||
// WithFields returns the same logger (no structured logging support in legacy adapter)
|
||||
func (l *LegacyLoggerAdapter) WithFields(fields map[string]interface{}) Logger {
|
||||
return l
|
||||
}
|
||||
@@ -0,0 +1,182 @@
|
||||
package logger
|
||||
|
||||
import (
|
||||
"io"
|
||||
"os"
|
||||
"sync"
|
||||
)
|
||||
|
||||
// Factory creates and manages logger instances with singleton support
|
||||
// for common logger types to reduce memory allocation.
|
||||
type Factory struct {
|
||||
mu sync.RWMutex
|
||||
defaultLogger Logger
|
||||
noOpLogger Logger
|
||||
loggers map[string]Logger
|
||||
defaultLogLevel string
|
||||
}
|
||||
|
||||
var (
|
||||
// globalFactory is the singleton factory instance
|
||||
globalFactory *Factory
|
||||
// factoryOnce ensures the factory is created only once
|
||||
factoryOnce sync.Once
|
||||
)
|
||||
|
||||
// GetFactory returns the global logger factory instance
|
||||
func GetFactory() *Factory {
|
||||
factoryOnce.Do(func() {
|
||||
globalFactory = &Factory{
|
||||
loggers: make(map[string]Logger),
|
||||
defaultLogLevel: "info",
|
||||
}
|
||||
})
|
||||
return globalFactory
|
||||
}
|
||||
|
||||
// SetDefaultLogLevel sets the default log level for new loggers
|
||||
func (f *Factory) SetDefaultLogLevel(level string) {
|
||||
f.mu.Lock()
|
||||
defer f.mu.Unlock()
|
||||
f.defaultLogLevel = level
|
||||
}
|
||||
|
||||
// GetLogger returns a logger for the given name, creating one if it doesn't exist
|
||||
func (f *Factory) GetLogger(name string) Logger {
|
||||
f.mu.RLock()
|
||||
if logger, exists := f.loggers[name]; exists {
|
||||
f.mu.RUnlock()
|
||||
return logger
|
||||
}
|
||||
f.mu.RUnlock()
|
||||
|
||||
// Create new logger
|
||||
f.mu.Lock()
|
||||
defer f.mu.Unlock()
|
||||
|
||||
// Double check after acquiring write lock
|
||||
if logger, exists := f.loggers[name]; exists {
|
||||
return logger
|
||||
}
|
||||
|
||||
logger := f.createLogger(name)
|
||||
f.loggers[name] = logger
|
||||
return logger
|
||||
}
|
||||
|
||||
// createLogger creates a new logger instance
|
||||
func (f *Factory) createLogger(name string) Logger {
|
||||
if name == "noop" || name == "no-op" || name == "discard" {
|
||||
return GetNoOpLogger()
|
||||
}
|
||||
|
||||
// Create logger with appropriate outputs based on environment
|
||||
var errorOut, infoOut, debugOut io.Writer
|
||||
|
||||
if os.Getenv("OIDC_LOG_TO_FILE") == "true" {
|
||||
// Log to files if configured
|
||||
errorOut = getOrCreateLogFile("error.log")
|
||||
infoOut = getOrCreateLogFile("info.log")
|
||||
debugOut = getOrCreateLogFile("debug.log")
|
||||
} else {
|
||||
// Default to stdout/stderr
|
||||
errorOut = os.Stderr
|
||||
infoOut = os.Stdout
|
||||
debugOut = os.Stdout
|
||||
}
|
||||
|
||||
return NewStandardLogger(f.defaultLogLevel, errorOut, infoOut, debugOut)
|
||||
}
|
||||
|
||||
// GetDefaultLogger returns the default logger instance
|
||||
func (f *Factory) GetDefaultLogger() Logger {
|
||||
f.mu.RLock()
|
||||
if f.defaultLogger != nil {
|
||||
f.mu.RUnlock()
|
||||
return f.defaultLogger
|
||||
}
|
||||
f.mu.RUnlock()
|
||||
|
||||
f.mu.Lock()
|
||||
defer f.mu.Unlock()
|
||||
|
||||
if f.defaultLogger == nil {
|
||||
f.defaultLogger = f.createLogger("default")
|
||||
}
|
||||
|
||||
return f.defaultLogger
|
||||
}
|
||||
|
||||
// GetNoOpLogger returns the singleton no-op logger
|
||||
func (f *Factory) GetNoOpLogger() Logger {
|
||||
f.mu.RLock()
|
||||
if f.noOpLogger != nil {
|
||||
f.mu.RUnlock()
|
||||
return f.noOpLogger
|
||||
}
|
||||
f.mu.RUnlock()
|
||||
|
||||
f.mu.Lock()
|
||||
defer f.mu.Unlock()
|
||||
|
||||
if f.noOpLogger == nil {
|
||||
f.noOpLogger = GetNoOpLogger()
|
||||
}
|
||||
|
||||
return f.noOpLogger
|
||||
}
|
||||
|
||||
// Clear removes all cached loggers (useful for testing)
|
||||
func (f *Factory) Clear() {
|
||||
f.mu.Lock()
|
||||
defer f.mu.Unlock()
|
||||
|
||||
f.loggers = make(map[string]Logger)
|
||||
f.defaultLogger = nil
|
||||
// Don't clear noOpLogger as it's a singleton
|
||||
}
|
||||
|
||||
// getOrCreateLogFile returns a file writer for the given log file
|
||||
func getOrCreateLogFile(filename string) io.Writer {
|
||||
logDir := os.Getenv("OIDC_LOG_DIR")
|
||||
if logDir == "" {
|
||||
logDir = "/var/log/traefik-oidc"
|
||||
}
|
||||
|
||||
// Ensure log directory exists
|
||||
if err := os.MkdirAll(logDir, 0755); err != nil {
|
||||
// Fall back to stderr if we can't create the directory
|
||||
return os.Stderr
|
||||
}
|
||||
|
||||
filepath := logDir + "/" + filename
|
||||
file, err := os.OpenFile(filepath, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0644)
|
||||
if err != nil {
|
||||
// Fall back to stderr if we can't open the file
|
||||
return os.Stderr
|
||||
}
|
||||
|
||||
return file
|
||||
}
|
||||
|
||||
// Global convenience functions
|
||||
|
||||
// New creates a new logger with the specified level
|
||||
func New(level string) Logger {
|
||||
return GetFactory().GetLogger(level)
|
||||
}
|
||||
|
||||
// Default returns the default logger
|
||||
func Default() Logger {
|
||||
return GetFactory().GetDefaultLogger()
|
||||
}
|
||||
|
||||
// NoOp returns a no-op logger
|
||||
func NoOp() Logger {
|
||||
return GetFactory().GetNoOpLogger()
|
||||
}
|
||||
|
||||
// WithLevel creates a new logger with the specified level
|
||||
func WithLevel(level string) Logger {
|
||||
return NewStandardLogger(level, os.Stderr, os.Stdout, os.Stdout)
|
||||
}
|
||||
@@ -0,0 +1,312 @@
|
||||
// Package logger provides a unified logging interface for the entire application.
|
||||
// It consolidates all the duplicate logger interfaces into a single, comprehensive
|
||||
// interface that supports different log levels and structured logging.
|
||||
package logger
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"io"
|
||||
"log"
|
||||
"sync"
|
||||
)
|
||||
|
||||
// Logger is the unified interface for all logging operations in the application.
|
||||
// It combines all the methods from the various logger interfaces that were
|
||||
// previously scattered across different packages.
|
||||
type Logger interface {
|
||||
// Basic logging methods
|
||||
Debug(msg string)
|
||||
Debugf(format string, args ...interface{})
|
||||
Info(msg string)
|
||||
Infof(format string, args ...interface{})
|
||||
Error(msg string)
|
||||
Errorf(format string, args ...interface{})
|
||||
|
||||
// Additional methods for compatibility with existing code
|
||||
Printf(format string, args ...interface{})
|
||||
Println(args ...interface{})
|
||||
Fatalf(format string, args ...interface{})
|
||||
|
||||
// Structured logging support
|
||||
WithField(key string, value interface{}) Logger
|
||||
WithFields(fields map[string]interface{}) Logger
|
||||
}
|
||||
|
||||
// StandardLogger implements the Logger interface using Go's standard log package.
|
||||
// It provides thread-safe logging with different output streams for different log levels.
|
||||
type StandardLogger struct {
|
||||
mu sync.RWMutex
|
||||
logError *log.Logger
|
||||
logInfo *log.Logger
|
||||
logDebug *log.Logger
|
||||
fields map[string]interface{}
|
||||
level LogLevel
|
||||
}
|
||||
|
||||
// LogLevel represents the logging level
|
||||
type LogLevel int
|
||||
|
||||
const (
|
||||
// LogLevelDebug enables all log messages
|
||||
LogLevelDebug LogLevel = iota
|
||||
// LogLevelInfo enables info and error messages
|
||||
LogLevelInfo
|
||||
// LogLevelError enables only error messages
|
||||
LogLevelError
|
||||
// LogLevelNone disables all logging
|
||||
LogLevelNone
|
||||
)
|
||||
|
||||
// ParseLogLevel converts a string log level to LogLevel
|
||||
func ParseLogLevel(level string) LogLevel {
|
||||
switch level {
|
||||
case "debug", "DEBUG":
|
||||
return LogLevelDebug
|
||||
case "info", "INFO":
|
||||
return LogLevelInfo
|
||||
case "error", "ERROR":
|
||||
return LogLevelError
|
||||
case "none", "NONE":
|
||||
return LogLevelNone
|
||||
default:
|
||||
return LogLevelInfo
|
||||
}
|
||||
}
|
||||
|
||||
// NewStandardLogger creates a new StandardLogger with the specified log level
|
||||
func NewStandardLogger(level string, errorOutput, infoOutput, debugOutput io.Writer) *StandardLogger {
|
||||
logLevel := ParseLogLevel(level)
|
||||
|
||||
if errorOutput == nil {
|
||||
errorOutput = io.Discard
|
||||
}
|
||||
if infoOutput == nil {
|
||||
infoOutput = io.Discard
|
||||
}
|
||||
if debugOutput == nil {
|
||||
debugOutput = io.Discard
|
||||
}
|
||||
|
||||
return &StandardLogger{
|
||||
logError: log.New(errorOutput, "ERROR: ", log.Ldate|log.Ltime|log.Lshortfile),
|
||||
logInfo: log.New(infoOutput, "INFO: ", log.Ldate|log.Ltime),
|
||||
logDebug: log.New(debugOutput, "DEBUG: ", log.Ldate|log.Ltime|log.Lshortfile),
|
||||
fields: make(map[string]interface{}),
|
||||
level: logLevel,
|
||||
}
|
||||
}
|
||||
|
||||
// Debug logs a debug message
|
||||
func (l *StandardLogger) Debug(msg string) {
|
||||
if l.level <= LogLevelDebug {
|
||||
l.mu.RLock()
|
||||
defer l.mu.RUnlock()
|
||||
if len(l.fields) > 0 {
|
||||
msg = l.formatWithFields(msg)
|
||||
}
|
||||
l.logDebug.Print(msg)
|
||||
}
|
||||
}
|
||||
|
||||
// Debugf logs a formatted debug message
|
||||
func (l *StandardLogger) Debugf(format string, args ...interface{}) {
|
||||
if l.level <= LogLevelDebug {
|
||||
l.mu.RLock()
|
||||
defer l.mu.RUnlock()
|
||||
msg := fmt.Sprintf(format, args...)
|
||||
if len(l.fields) > 0 {
|
||||
msg = l.formatWithFields(msg)
|
||||
}
|
||||
l.logDebug.Print(msg)
|
||||
}
|
||||
}
|
||||
|
||||
// Info logs an info message
|
||||
func (l *StandardLogger) Info(msg string) {
|
||||
if l.level <= LogLevelInfo {
|
||||
l.mu.RLock()
|
||||
defer l.mu.RUnlock()
|
||||
if len(l.fields) > 0 {
|
||||
msg = l.formatWithFields(msg)
|
||||
}
|
||||
l.logInfo.Print(msg)
|
||||
}
|
||||
}
|
||||
|
||||
// Infof logs a formatted info message
|
||||
func (l *StandardLogger) Infof(format string, args ...interface{}) {
|
||||
if l.level <= LogLevelInfo {
|
||||
l.mu.RLock()
|
||||
defer l.mu.RUnlock()
|
||||
msg := fmt.Sprintf(format, args...)
|
||||
if len(l.fields) > 0 {
|
||||
msg = l.formatWithFields(msg)
|
||||
}
|
||||
l.logInfo.Print(msg)
|
||||
}
|
||||
}
|
||||
|
||||
// Error logs an error message
|
||||
func (l *StandardLogger) Error(msg string) {
|
||||
if l.level <= LogLevelError {
|
||||
l.mu.RLock()
|
||||
defer l.mu.RUnlock()
|
||||
if len(l.fields) > 0 {
|
||||
msg = l.formatWithFields(msg)
|
||||
}
|
||||
l.logError.Print(msg)
|
||||
}
|
||||
}
|
||||
|
||||
// Errorf logs a formatted error message
|
||||
func (l *StandardLogger) Errorf(format string, args ...interface{}) {
|
||||
if l.level <= LogLevelError {
|
||||
l.mu.RLock()
|
||||
defer l.mu.RUnlock()
|
||||
msg := fmt.Sprintf(format, args...)
|
||||
if len(l.fields) > 0 {
|
||||
msg = l.formatWithFields(msg)
|
||||
}
|
||||
l.logError.Print(msg)
|
||||
}
|
||||
}
|
||||
|
||||
// Printf logs a formatted message at info level
|
||||
func (l *StandardLogger) Printf(format string, args ...interface{}) {
|
||||
l.Infof(format, args...)
|
||||
}
|
||||
|
||||
// Println logs a message at info level
|
||||
func (l *StandardLogger) Println(args ...interface{}) {
|
||||
l.Info(fmt.Sprint(args...))
|
||||
}
|
||||
|
||||
// Fatalf logs a formatted error message and exits the program
|
||||
func (l *StandardLogger) Fatalf(format string, args ...interface{}) {
|
||||
l.Errorf(format, args...)
|
||||
panic(fmt.Sprintf(format, args...))
|
||||
}
|
||||
|
||||
// WithField returns a new logger with an additional field
|
||||
func (l *StandardLogger) WithField(key string, value interface{}) Logger {
|
||||
l.mu.Lock()
|
||||
defer l.mu.Unlock()
|
||||
|
||||
newLogger := &StandardLogger{
|
||||
logError: l.logError,
|
||||
logInfo: l.logInfo,
|
||||
logDebug: l.logDebug,
|
||||
fields: make(map[string]interface{}, len(l.fields)+1),
|
||||
level: l.level,
|
||||
}
|
||||
|
||||
for k, v := range l.fields {
|
||||
newLogger.fields[k] = v
|
||||
}
|
||||
newLogger.fields[key] = value
|
||||
|
||||
return newLogger
|
||||
}
|
||||
|
||||
// WithFields returns a new logger with additional fields
|
||||
func (l *StandardLogger) WithFields(fields map[string]interface{}) Logger {
|
||||
l.mu.Lock()
|
||||
defer l.mu.Unlock()
|
||||
|
||||
newLogger := &StandardLogger{
|
||||
logError: l.logError,
|
||||
logInfo: l.logInfo,
|
||||
logDebug: l.logDebug,
|
||||
fields: make(map[string]interface{}, len(l.fields)+len(fields)),
|
||||
level: l.level,
|
||||
}
|
||||
|
||||
for k, v := range l.fields {
|
||||
newLogger.fields[k] = v
|
||||
}
|
||||
for k, v := range fields {
|
||||
newLogger.fields[k] = v
|
||||
}
|
||||
|
||||
return newLogger
|
||||
}
|
||||
|
||||
// formatWithFields formats a message with structured fields
|
||||
func (l *StandardLogger) formatWithFields(msg string) string {
|
||||
if len(l.fields) == 0 {
|
||||
return msg
|
||||
}
|
||||
|
||||
fieldsStr := ""
|
||||
for k, v := range l.fields {
|
||||
if fieldsStr != "" {
|
||||
fieldsStr += " "
|
||||
}
|
||||
fieldsStr += fmt.Sprintf("%s=%v", k, v)
|
||||
}
|
||||
|
||||
return fmt.Sprintf("%s [%s]", msg, fieldsStr)
|
||||
}
|
||||
|
||||
// NoOpLogger is a logger that discards all output.
|
||||
// It's useful for testing and for cases where logging should be disabled.
|
||||
type NoOpLogger struct{}
|
||||
|
||||
// Debug discards the message
|
||||
func (n *NoOpLogger) Debug(msg string) {}
|
||||
|
||||
// Debugf discards the formatted message
|
||||
func (n *NoOpLogger) Debugf(format string, args ...interface{}) {}
|
||||
|
||||
// Info discards the message
|
||||
func (n *NoOpLogger) Info(msg string) {}
|
||||
|
||||
// Infof discards the formatted message
|
||||
func (n *NoOpLogger) Infof(format string, args ...interface{}) {}
|
||||
|
||||
// Error discards the message
|
||||
func (n *NoOpLogger) Error(msg string) {}
|
||||
|
||||
// Errorf discards the formatted message
|
||||
func (n *NoOpLogger) Errorf(format string, args ...interface{}) {}
|
||||
|
||||
// Printf discards the formatted message
|
||||
func (n *NoOpLogger) Printf(format string, args ...interface{}) {}
|
||||
|
||||
// Println discards the message
|
||||
func (n *NoOpLogger) Println(args ...interface{}) {}
|
||||
|
||||
// Fatalf discards the message and does not exit
|
||||
func (n *NoOpLogger) Fatalf(format string, args ...interface{}) {}
|
||||
|
||||
// WithField returns the same NoOpLogger
|
||||
func (n *NoOpLogger) WithField(key string, value interface{}) Logger {
|
||||
return n
|
||||
}
|
||||
|
||||
// WithFields returns the same NoOpLogger
|
||||
func (n *NoOpLogger) WithFields(fields map[string]interface{}) Logger {
|
||||
return n
|
||||
}
|
||||
|
||||
var (
|
||||
// singletonNoOpLogger is the global instance of the no-op logger
|
||||
singletonNoOpLogger *NoOpLogger
|
||||
// noOpLoggerOnce ensures the singleton is created only once
|
||||
noOpLoggerOnce sync.Once
|
||||
)
|
||||
|
||||
// GetNoOpLogger returns the singleton no-op logger instance.
|
||||
// This reduces memory allocation by reusing the same no-op logger
|
||||
// instance across the entire application.
|
||||
func GetNoOpLogger() Logger {
|
||||
noOpLoggerOnce.Do(func() {
|
||||
singletonNoOpLogger = &NoOpLogger{}
|
||||
})
|
||||
return singletonNoOpLogger
|
||||
}
|
||||
|
||||
// DefaultLogger creates a default logger based on the provided configuration
|
||||
func DefaultLogger(level string) Logger {
|
||||
return NewStandardLogger(level, log.Writer(), log.Writer(), log.Writer())
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,473 @@
|
||||
// Package pool provides a unified, centralized memory pool management system
|
||||
// for the entire application. It consolidates all duplicate pool implementations
|
||||
// into a single, efficient, and thread-safe package.
|
||||
package pool
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"compress/gzip"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
)
|
||||
|
||||
// Manager is the centralized pool manager that consolidates all memory pools
|
||||
// used throughout the application. It provides a single entry point for
|
||||
// all pooling operations, reducing duplicate code and improving maintainability.
|
||||
type Manager struct {
|
||||
// Buffer pools
|
||||
smallBufferPool *sync.Pool // 1KB buffers
|
||||
mediumBufferPool *sync.Pool // 4KB buffers
|
||||
largeBufferPool *sync.Pool // 8KB buffers
|
||||
xlBufferPool *sync.Pool // 16KB buffers
|
||||
|
||||
// Compression pools
|
||||
gzipWriterPool *sync.Pool
|
||||
gzipReaderPool *sync.Pool
|
||||
|
||||
// String builder pool
|
||||
stringBuilderPool *sync.Pool
|
||||
|
||||
// JWT parsing buffers
|
||||
jwtBufferPool *sync.Pool
|
||||
|
||||
// HTTP response buffers
|
||||
httpResponsePool *sync.Pool
|
||||
|
||||
// Byte slice pools for various sizes
|
||||
byteSlicePools map[int]*sync.Pool
|
||||
poolMu sync.RWMutex
|
||||
|
||||
// Statistics
|
||||
stats PoolStats
|
||||
}
|
||||
|
||||
// PoolStats tracks pool usage statistics
|
||||
type PoolStats struct {
|
||||
BufferGets uint64
|
||||
BufferPuts uint64
|
||||
GzipGets uint64
|
||||
GzipPuts uint64
|
||||
StringGets uint64
|
||||
StringPuts uint64
|
||||
JWTGets uint64
|
||||
JWTPuts uint64
|
||||
HTTPGets uint64
|
||||
HTTPPuts uint64
|
||||
OversizedRejects uint64
|
||||
}
|
||||
|
||||
// JWTBuffer provides pre-allocated buffers for JWT parsing
|
||||
type JWTBuffer struct {
|
||||
Header []byte
|
||||
Payload []byte
|
||||
Signature []byte
|
||||
}
|
||||
|
||||
var (
|
||||
// globalManager is the singleton pool manager instance
|
||||
globalManager *Manager
|
||||
// managerOnce ensures single initialization
|
||||
managerOnce sync.Once
|
||||
)
|
||||
|
||||
// Get returns the global pool manager instance
|
||||
func Get() *Manager {
|
||||
managerOnce.Do(func() {
|
||||
globalManager = newManager()
|
||||
})
|
||||
return globalManager
|
||||
}
|
||||
|
||||
// newManager creates a new pool manager with all pools initialized
|
||||
func newManager() *Manager {
|
||||
m := &Manager{
|
||||
byteSlicePools: make(map[int]*sync.Pool),
|
||||
}
|
||||
|
||||
// Initialize buffer pools with different sizes
|
||||
m.smallBufferPool = &sync.Pool{
|
||||
New: func() interface{} {
|
||||
return bytes.NewBuffer(make([]byte, 0, 1024))
|
||||
},
|
||||
}
|
||||
|
||||
m.mediumBufferPool = &sync.Pool{
|
||||
New: func() interface{} {
|
||||
return bytes.NewBuffer(make([]byte, 0, 4096))
|
||||
},
|
||||
}
|
||||
|
||||
m.largeBufferPool = &sync.Pool{
|
||||
New: func() interface{} {
|
||||
return bytes.NewBuffer(make([]byte, 0, 8192))
|
||||
},
|
||||
}
|
||||
|
||||
m.xlBufferPool = &sync.Pool{
|
||||
New: func() interface{} {
|
||||
return bytes.NewBuffer(make([]byte, 0, 16384))
|
||||
},
|
||||
}
|
||||
|
||||
// Initialize compression pools
|
||||
m.gzipWriterPool = &sync.Pool{
|
||||
New: func() interface{} {
|
||||
w, _ := gzip.NewWriterLevel(nil, gzip.BestSpeed)
|
||||
return w
|
||||
},
|
||||
}
|
||||
|
||||
m.gzipReaderPool = &sync.Pool{
|
||||
New: func() interface{} {
|
||||
return (*gzip.Reader)(nil)
|
||||
},
|
||||
}
|
||||
|
||||
// Initialize string builder pool
|
||||
m.stringBuilderPool = &sync.Pool{
|
||||
New: func() interface{} {
|
||||
sb := &strings.Builder{}
|
||||
sb.Grow(1024)
|
||||
return sb
|
||||
},
|
||||
}
|
||||
|
||||
// Initialize JWT buffer pool
|
||||
m.jwtBufferPool = &sync.Pool{
|
||||
New: func() interface{} {
|
||||
return &JWTBuffer{
|
||||
Header: make([]byte, 0, 512),
|
||||
Payload: make([]byte, 0, 2048),
|
||||
Signature: make([]byte, 0, 512),
|
||||
}
|
||||
},
|
||||
}
|
||||
|
||||
// Initialize HTTP response buffer pool
|
||||
m.httpResponsePool = &sync.Pool{
|
||||
New: func() interface{} {
|
||||
buf := make([]byte, 0, 8192)
|
||||
return &buf
|
||||
},
|
||||
}
|
||||
|
||||
// Initialize common byte slice pools
|
||||
for _, size := range []int{256, 512, 1024, 2048, 4096, 8192, 16384} {
|
||||
size := size // capture for closure
|
||||
m.byteSlicePools[size] = &sync.Pool{
|
||||
New: func() interface{} {
|
||||
b := make([]byte, size)
|
||||
return &b
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
return m
|
||||
}
|
||||
|
||||
// GetBuffer returns a buffer from the appropriate pool based on size hint
|
||||
func (m *Manager) GetBuffer(sizeHint int) *bytes.Buffer {
|
||||
atomic.AddUint64(&m.stats.BufferGets, 1)
|
||||
|
||||
switch {
|
||||
case sizeHint <= 1024:
|
||||
return m.smallBufferPool.Get().(*bytes.Buffer)
|
||||
case sizeHint <= 4096:
|
||||
return m.mediumBufferPool.Get().(*bytes.Buffer)
|
||||
case sizeHint <= 8192:
|
||||
return m.largeBufferPool.Get().(*bytes.Buffer)
|
||||
case sizeHint <= 16384:
|
||||
return m.xlBufferPool.Get().(*bytes.Buffer)
|
||||
default:
|
||||
// For very large buffers, create new ones
|
||||
return bytes.NewBuffer(make([]byte, 0, sizeHint))
|
||||
}
|
||||
}
|
||||
|
||||
// PutBuffer returns a buffer to the appropriate pool
|
||||
func (m *Manager) PutBuffer(buf *bytes.Buffer) {
|
||||
if buf == nil {
|
||||
return
|
||||
}
|
||||
|
||||
atomic.AddUint64(&m.stats.BufferPuts, 1)
|
||||
|
||||
// Reset buffer before returning to pool
|
||||
capacity := buf.Cap()
|
||||
buf.Reset()
|
||||
|
||||
// Reject oversized buffers to prevent memory bloat
|
||||
if capacity > 32768 {
|
||||
atomic.AddUint64(&m.stats.OversizedRejects, 1)
|
||||
return
|
||||
}
|
||||
|
||||
// Return to appropriate pool based on capacity
|
||||
switch {
|
||||
case capacity <= 1024:
|
||||
m.smallBufferPool.Put(buf)
|
||||
case capacity <= 4096:
|
||||
m.mediumBufferPool.Put(buf)
|
||||
case capacity <= 8192:
|
||||
m.largeBufferPool.Put(buf)
|
||||
case capacity <= 16384:
|
||||
m.xlBufferPool.Put(buf)
|
||||
}
|
||||
}
|
||||
|
||||
// GetGzipWriter returns a gzip writer from the pool
|
||||
func (m *Manager) GetGzipWriter() *gzip.Writer {
|
||||
atomic.AddUint64(&m.stats.GzipGets, 1)
|
||||
return m.gzipWriterPool.Get().(*gzip.Writer)
|
||||
}
|
||||
|
||||
// PutGzipWriter returns a gzip writer to the pool
|
||||
func (m *Manager) PutGzipWriter(w *gzip.Writer) {
|
||||
if w == nil {
|
||||
return
|
||||
}
|
||||
atomic.AddUint64(&m.stats.GzipPuts, 1)
|
||||
w.Reset(nil)
|
||||
m.gzipWriterPool.Put(w)
|
||||
}
|
||||
|
||||
// GetGzipReader returns a gzip reader from the pool
|
||||
func (m *Manager) GetGzipReader() *gzip.Reader {
|
||||
atomic.AddUint64(&m.stats.GzipGets, 1)
|
||||
r := m.gzipReaderPool.Get()
|
||||
if r == nil {
|
||||
return nil
|
||||
}
|
||||
return r.(*gzip.Reader)
|
||||
}
|
||||
|
||||
// PutGzipReader returns a gzip reader to the pool
|
||||
func (m *Manager) PutGzipReader(r *gzip.Reader) {
|
||||
if r == nil {
|
||||
return
|
||||
}
|
||||
atomic.AddUint64(&m.stats.GzipPuts, 1)
|
||||
r.Reset(nil)
|
||||
m.gzipReaderPool.Put(r)
|
||||
}
|
||||
|
||||
// GetStringBuilder returns a string builder from the pool
|
||||
func (m *Manager) GetStringBuilder() *strings.Builder {
|
||||
atomic.AddUint64(&m.stats.StringGets, 1)
|
||||
sb := m.stringBuilderPool.Get().(*strings.Builder)
|
||||
sb.Reset()
|
||||
return sb
|
||||
}
|
||||
|
||||
// PutStringBuilder returns a string builder to the pool
|
||||
func (m *Manager) PutStringBuilder(sb *strings.Builder) {
|
||||
if sb == nil {
|
||||
return
|
||||
}
|
||||
|
||||
atomic.AddUint64(&m.stats.StringPuts, 1)
|
||||
|
||||
// Reject oversized builders
|
||||
if sb.Cap() > 16384 {
|
||||
atomic.AddUint64(&m.stats.OversizedRejects, 1)
|
||||
return
|
||||
}
|
||||
|
||||
sb.Reset()
|
||||
m.stringBuilderPool.Put(sb)
|
||||
}
|
||||
|
||||
// GetJWTBuffer returns JWT parsing buffers from the pool
|
||||
func (m *Manager) GetJWTBuffer() *JWTBuffer {
|
||||
atomic.AddUint64(&m.stats.JWTGets, 1)
|
||||
return m.jwtBufferPool.Get().(*JWTBuffer)
|
||||
}
|
||||
|
||||
// PutJWTBuffer returns JWT parsing buffers to the pool
|
||||
func (m *Manager) PutJWTBuffer(buf *JWTBuffer) {
|
||||
if buf == nil {
|
||||
return
|
||||
}
|
||||
|
||||
atomic.AddUint64(&m.stats.JWTPuts, 1)
|
||||
|
||||
// Check for oversized buffers
|
||||
if cap(buf.Header) > 2048 || cap(buf.Payload) > 8192 || cap(buf.Signature) > 2048 {
|
||||
atomic.AddUint64(&m.stats.OversizedRejects, 1)
|
||||
return
|
||||
}
|
||||
|
||||
// Reset slices to zero length
|
||||
buf.Header = buf.Header[:0]
|
||||
buf.Payload = buf.Payload[:0]
|
||||
buf.Signature = buf.Signature[:0]
|
||||
m.jwtBufferPool.Put(buf)
|
||||
}
|
||||
|
||||
// GetHTTPResponseBuffer returns an HTTP response buffer from the pool
|
||||
func (m *Manager) GetHTTPResponseBuffer() []byte {
|
||||
atomic.AddUint64(&m.stats.HTTPGets, 1)
|
||||
return *m.httpResponsePool.Get().(*[]byte)
|
||||
}
|
||||
|
||||
// PutHTTPResponseBuffer returns an HTTP response buffer to the pool
|
||||
func (m *Manager) PutHTTPResponseBuffer(buf []byte) {
|
||||
if buf == nil {
|
||||
return
|
||||
}
|
||||
|
||||
atomic.AddUint64(&m.stats.HTTPPuts, 1)
|
||||
|
||||
// Reject oversized buffers
|
||||
if cap(buf) > 32768 {
|
||||
atomic.AddUint64(&m.stats.OversizedRejects, 1)
|
||||
return
|
||||
}
|
||||
|
||||
buf = buf[:0]
|
||||
m.httpResponsePool.Put(&buf)
|
||||
}
|
||||
|
||||
// GetByteSlice returns a byte slice of the specified size from the pool
|
||||
func (m *Manager) GetByteSlice(size int) []byte {
|
||||
m.poolMu.RLock()
|
||||
pool, exists := m.byteSlicePools[size]
|
||||
m.poolMu.RUnlock()
|
||||
|
||||
if !exists {
|
||||
// Round up to nearest power of 2
|
||||
poolSize := 1
|
||||
for poolSize < size {
|
||||
poolSize *= 2
|
||||
}
|
||||
|
||||
m.poolMu.Lock()
|
||||
// Double-check after acquiring write lock
|
||||
pool, exists = m.byteSlicePools[poolSize]
|
||||
if !exists {
|
||||
pool = &sync.Pool{
|
||||
New: func() interface{} {
|
||||
b := make([]byte, poolSize)
|
||||
return &b
|
||||
},
|
||||
}
|
||||
m.byteSlicePools[poolSize] = pool
|
||||
}
|
||||
m.poolMu.Unlock()
|
||||
}
|
||||
|
||||
b := pool.Get().(*[]byte)
|
||||
return (*b)[:size]
|
||||
}
|
||||
|
||||
// PutByteSlice returns a byte slice to the pool
|
||||
func (m *Manager) PutByteSlice(b []byte) {
|
||||
if b == nil || cap(b) > 65536 { // Don't pool very large slices
|
||||
return
|
||||
}
|
||||
|
||||
size := cap(b)
|
||||
m.poolMu.RLock()
|
||||
pool, exists := m.byteSlicePools[size]
|
||||
m.poolMu.RUnlock()
|
||||
|
||||
if exists {
|
||||
b = b[:0]
|
||||
pool.Put(&b)
|
||||
}
|
||||
}
|
||||
|
||||
// GetStats returns current pool statistics
|
||||
func (m *Manager) GetStats() PoolStats {
|
||||
return PoolStats{
|
||||
BufferGets: atomic.LoadUint64(&m.stats.BufferGets),
|
||||
BufferPuts: atomic.LoadUint64(&m.stats.BufferPuts),
|
||||
GzipGets: atomic.LoadUint64(&m.stats.GzipGets),
|
||||
GzipPuts: atomic.LoadUint64(&m.stats.GzipPuts),
|
||||
StringGets: atomic.LoadUint64(&m.stats.StringGets),
|
||||
StringPuts: atomic.LoadUint64(&m.stats.StringPuts),
|
||||
JWTGets: atomic.LoadUint64(&m.stats.JWTGets),
|
||||
JWTPuts: atomic.LoadUint64(&m.stats.JWTPuts),
|
||||
HTTPGets: atomic.LoadUint64(&m.stats.HTTPGets),
|
||||
HTTPPuts: atomic.LoadUint64(&m.stats.HTTPPuts),
|
||||
OversizedRejects: atomic.LoadUint64(&m.stats.OversizedRejects),
|
||||
}
|
||||
}
|
||||
|
||||
// ResetStats resets all statistics counters
|
||||
func (m *Manager) ResetStats() {
|
||||
atomic.StoreUint64(&m.stats.BufferGets, 0)
|
||||
atomic.StoreUint64(&m.stats.BufferPuts, 0)
|
||||
atomic.StoreUint64(&m.stats.GzipGets, 0)
|
||||
atomic.StoreUint64(&m.stats.GzipPuts, 0)
|
||||
atomic.StoreUint64(&m.stats.StringGets, 0)
|
||||
atomic.StoreUint64(&m.stats.StringPuts, 0)
|
||||
atomic.StoreUint64(&m.stats.JWTGets, 0)
|
||||
atomic.StoreUint64(&m.stats.JWTPuts, 0)
|
||||
atomic.StoreUint64(&m.stats.HTTPGets, 0)
|
||||
atomic.StoreUint64(&m.stats.HTTPPuts, 0)
|
||||
atomic.StoreUint64(&m.stats.OversizedRejects, 0)
|
||||
}
|
||||
|
||||
// Global convenience functions
|
||||
|
||||
// Buffer returns a buffer from the global pool
|
||||
func Buffer(sizeHint int) *bytes.Buffer {
|
||||
return Get().GetBuffer(sizeHint)
|
||||
}
|
||||
|
||||
// ReturnBuffer returns a buffer to the global pool
|
||||
func ReturnBuffer(buf *bytes.Buffer) {
|
||||
Get().PutBuffer(buf)
|
||||
}
|
||||
|
||||
// GzipWriter returns a gzip writer from the global pool
|
||||
func GzipWriter() *gzip.Writer {
|
||||
return Get().GetGzipWriter()
|
||||
}
|
||||
|
||||
// ReturnGzipWriter returns a gzip writer to the global pool
|
||||
func ReturnGzipWriter(w *gzip.Writer) {
|
||||
Get().PutGzipWriter(w)
|
||||
}
|
||||
|
||||
// StringBuilder returns a string builder from the global pool
|
||||
func StringBuilder() *strings.Builder {
|
||||
return Get().GetStringBuilder()
|
||||
}
|
||||
|
||||
// ReturnStringBuilder returns a string builder to the global pool
|
||||
func ReturnStringBuilder(sb *strings.Builder) {
|
||||
Get().PutStringBuilder(sb)
|
||||
}
|
||||
|
||||
// JWTBuffers returns JWT parsing buffers from the global pool
|
||||
func JWTBuffers() *JWTBuffer {
|
||||
return Get().GetJWTBuffer()
|
||||
}
|
||||
|
||||
// ReturnJWTBuffers returns JWT parsing buffers to the global pool
|
||||
func ReturnJWTBuffers(buf *JWTBuffer) {
|
||||
Get().PutJWTBuffer(buf)
|
||||
}
|
||||
|
||||
// HTTPBuffer returns an HTTP response buffer from the global pool
|
||||
func HTTPBuffer() []byte {
|
||||
return Get().GetHTTPResponseBuffer()
|
||||
}
|
||||
|
||||
// ReturnHTTPBuffer returns an HTTP response buffer to the global pool
|
||||
func ReturnHTTPBuffer(buf []byte) {
|
||||
Get().PutHTTPResponseBuffer(buf)
|
||||
}
|
||||
|
||||
// ByteSlice returns a byte slice from the global pool
|
||||
func ByteSlice(size int) []byte {
|
||||
return Get().GetByteSlice(size)
|
||||
}
|
||||
|
||||
// ReturnByteSlice returns a byte slice to the global pool
|
||||
func ReturnByteSlice(b []byte) {
|
||||
Get().PutByteSlice(b)
|
||||
}
|
||||
@@ -0,0 +1,586 @@
|
||||
package pool
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// TestManager_Singleton tests that Get() returns the same instance
|
||||
func TestManager_Singleton(t *testing.T) {
|
||||
manager1 := Get()
|
||||
manager2 := Get()
|
||||
|
||||
if manager1 != manager2 {
|
||||
t.Error("Get() should return the same instance (singleton)")
|
||||
}
|
||||
|
||||
if manager1 == nil {
|
||||
t.Error("Get() should not return nil")
|
||||
}
|
||||
}
|
||||
|
||||
// TestManager_BufferPools tests buffer pool operations
|
||||
func TestManager_BufferPools(t *testing.T) {
|
||||
manager := Get()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
sizeHint int
|
||||
expected int // expected capacity range
|
||||
}{
|
||||
{"small buffer", 512, 1024},
|
||||
{"medium buffer", 2048, 4096},
|
||||
{"large buffer", 6144, 8192},
|
||||
{"xl buffer", 12288, 16384},
|
||||
{"oversized buffer", 32768, 32768}, // Should create new buffer
|
||||
}
|
||||
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
buf := manager.GetBuffer(test.sizeHint)
|
||||
if buf == nil {
|
||||
t.Error("GetBuffer should not return nil")
|
||||
}
|
||||
|
||||
if buf.Cap() < test.sizeHint {
|
||||
t.Errorf("Buffer capacity %d is less than size hint %d", buf.Cap(), test.sizeHint)
|
||||
}
|
||||
|
||||
// Write some data
|
||||
buf.WriteString("test data")
|
||||
if buf.String() != "test data" {
|
||||
t.Error("Buffer should contain written data")
|
||||
}
|
||||
|
||||
// Return to pool
|
||||
manager.PutBuffer(buf)
|
||||
|
||||
// Buffer should be reset when returned to pool
|
||||
buf2 := manager.GetBuffer(test.sizeHint)
|
||||
if buf2.Len() != 0 {
|
||||
t.Error("Buffer from pool should be reset")
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestManager_PutBuffer_Nil tests putting nil buffer
|
||||
func TestManager_PutBuffer_Nil(t *testing.T) {
|
||||
manager := Get()
|
||||
// Should not panic
|
||||
manager.PutBuffer(nil)
|
||||
}
|
||||
|
||||
// TestManager_PutBuffer_Oversized tests rejection of oversized buffers
|
||||
func TestManager_PutBuffer_Oversized(t *testing.T) {
|
||||
manager := Get()
|
||||
manager.ResetStats()
|
||||
|
||||
// Create oversized buffer
|
||||
buf := bytes.NewBuffer(make([]byte, 0, 40000))
|
||||
manager.PutBuffer(buf)
|
||||
|
||||
stats := manager.GetStats()
|
||||
if stats.OversizedRejects == 0 {
|
||||
t.Error("Oversized buffer should be rejected")
|
||||
}
|
||||
}
|
||||
|
||||
// TestManager_GzipPools tests gzip writer and reader pools
|
||||
func TestManager_GzipPools(t *testing.T) {
|
||||
manager := Get()
|
||||
|
||||
// Test gzip writer
|
||||
writer := manager.GetGzipWriter()
|
||||
if writer == nil {
|
||||
t.Error("GetGzipWriter should not return nil")
|
||||
}
|
||||
|
||||
// Test that we can use it
|
||||
var buf bytes.Buffer
|
||||
writer.Reset(&buf)
|
||||
writer.Write([]byte("test data"))
|
||||
writer.Close()
|
||||
|
||||
if buf.Len() == 0 {
|
||||
t.Error("Gzip writer should have written compressed data")
|
||||
}
|
||||
|
||||
// Return to pool
|
||||
manager.PutGzipWriter(writer)
|
||||
|
||||
// Test gzip reader
|
||||
reader := manager.GetGzipReader()
|
||||
// Reader might be nil from pool initially
|
||||
if reader != nil {
|
||||
manager.PutGzipReader(reader)
|
||||
}
|
||||
}
|
||||
|
||||
// TestManager_GzipPools_Nil tests putting nil gzip objects
|
||||
func TestManager_GzipPools_Nil(t *testing.T) {
|
||||
manager := Get()
|
||||
|
||||
// Should not panic
|
||||
manager.PutGzipWriter(nil)
|
||||
manager.PutGzipReader(nil)
|
||||
}
|
||||
|
||||
// TestManager_StringBuilderPool tests string builder pool
|
||||
func TestManager_StringBuilderPool(t *testing.T) {
|
||||
manager := Get()
|
||||
|
||||
sb := manager.GetStringBuilder()
|
||||
if sb == nil {
|
||||
t.Error("GetStringBuilder should not return nil")
|
||||
}
|
||||
|
||||
// Should be reset
|
||||
if sb.Len() != 0 {
|
||||
t.Error("String builder from pool should be reset")
|
||||
}
|
||||
|
||||
// Test writing
|
||||
sb.WriteString("test")
|
||||
sb.WriteString(" data")
|
||||
if sb.String() != "test data" {
|
||||
t.Error("String builder should contain written data")
|
||||
}
|
||||
|
||||
// Return to pool
|
||||
manager.PutStringBuilder(sb)
|
||||
|
||||
// Get another one - should be reset
|
||||
sb2 := manager.GetStringBuilder()
|
||||
if sb2.Len() != 0 {
|
||||
t.Error("String builder from pool should be reset")
|
||||
}
|
||||
}
|
||||
|
||||
// TestManager_StringBuilderPool_Nil tests putting nil string builder
|
||||
func TestManager_StringBuilderPool_Nil(t *testing.T) {
|
||||
manager := Get()
|
||||
// Should not panic
|
||||
manager.PutStringBuilder(nil)
|
||||
}
|
||||
|
||||
// TestManager_StringBuilderPool_Oversized tests rejection of oversized string builders
|
||||
func TestManager_StringBuilderPool_Oversized(t *testing.T) {
|
||||
manager := Get()
|
||||
manager.ResetStats()
|
||||
|
||||
// Create oversized string builder
|
||||
sb := &strings.Builder{}
|
||||
sb.Grow(20000)
|
||||
sb.WriteString("test")
|
||||
|
||||
manager.PutStringBuilder(sb)
|
||||
|
||||
stats := manager.GetStats()
|
||||
if stats.OversizedRejects == 0 {
|
||||
t.Error("Oversized string builder should be rejected")
|
||||
}
|
||||
}
|
||||
|
||||
// TestManager_JWTBufferPool tests JWT buffer pool
|
||||
func TestManager_JWTBufferPool(t *testing.T) {
|
||||
manager := Get()
|
||||
|
||||
jwtBuf := manager.GetJWTBuffer()
|
||||
if jwtBuf == nil {
|
||||
t.Error("GetJWTBuffer should not return nil")
|
||||
return
|
||||
}
|
||||
|
||||
// Check structure
|
||||
if jwtBuf.Header == nil || jwtBuf.Payload == nil || jwtBuf.Signature == nil {
|
||||
t.Error("JWT buffer should have all fields initialized")
|
||||
}
|
||||
|
||||
// Should be empty initially
|
||||
if len(jwtBuf.Header) != 0 || len(jwtBuf.Payload) != 0 || len(jwtBuf.Signature) != 0 {
|
||||
t.Error("JWT buffer from pool should be reset")
|
||||
}
|
||||
|
||||
// Use the buffer
|
||||
jwtBuf.Header = append(jwtBuf.Header, []byte("header")...)
|
||||
jwtBuf.Payload = append(jwtBuf.Payload, []byte("payload")...)
|
||||
jwtBuf.Signature = append(jwtBuf.Signature, []byte("signature")...)
|
||||
|
||||
// Return to pool
|
||||
manager.PutJWTBuffer(jwtBuf)
|
||||
|
||||
// Get another one - should be reset
|
||||
jwtBuf2 := manager.GetJWTBuffer()
|
||||
if len(jwtBuf2.Header) != 0 || len(jwtBuf2.Payload) != 0 || len(jwtBuf2.Signature) != 0 {
|
||||
t.Error("JWT buffer from pool should be reset")
|
||||
}
|
||||
}
|
||||
|
||||
// TestManager_JWTBufferPool_Nil tests putting nil JWT buffer
|
||||
func TestManager_JWTBufferPool_Nil(t *testing.T) {
|
||||
manager := Get()
|
||||
// Should not panic
|
||||
manager.PutJWTBuffer(nil)
|
||||
}
|
||||
|
||||
// TestManager_JWTBufferPool_Oversized tests rejection of oversized JWT buffers
|
||||
func TestManager_JWTBufferPool_Oversized(t *testing.T) {
|
||||
manager := Get()
|
||||
manager.ResetStats()
|
||||
|
||||
// Create oversized JWT buffer
|
||||
jwtBuf := &JWTBuffer{
|
||||
Header: make([]byte, 0, 3000), // Over 2048 limit
|
||||
Payload: make([]byte, 0, 10000), // Over 8192 limit
|
||||
Signature: make([]byte, 0, 3000), // Over 2048 limit
|
||||
}
|
||||
|
||||
manager.PutJWTBuffer(jwtBuf)
|
||||
|
||||
stats := manager.GetStats()
|
||||
if stats.OversizedRejects == 0 {
|
||||
t.Error("Oversized JWT buffer should be rejected")
|
||||
}
|
||||
}
|
||||
|
||||
// TestManager_HTTPResponsePool tests HTTP response buffer pool
|
||||
func TestManager_HTTPResponsePool(t *testing.T) {
|
||||
manager := Get()
|
||||
|
||||
buf := manager.GetHTTPResponseBuffer()
|
||||
if buf == nil {
|
||||
t.Error("GetHTTPResponseBuffer should not return nil")
|
||||
}
|
||||
|
||||
// Should be empty initially
|
||||
if len(buf) != 0 {
|
||||
t.Error("HTTP buffer from pool should be empty")
|
||||
}
|
||||
|
||||
// Use the buffer
|
||||
buf = append(buf, []byte("HTTP response data")...)
|
||||
|
||||
// Return to pool
|
||||
manager.PutHTTPResponseBuffer(buf)
|
||||
|
||||
// Get another one - should be reset
|
||||
buf2 := manager.GetHTTPResponseBuffer()
|
||||
if len(buf2) != 0 {
|
||||
t.Error("HTTP buffer from pool should be reset")
|
||||
}
|
||||
}
|
||||
|
||||
// TestManager_HTTPResponsePool_Nil tests putting nil HTTP buffer
|
||||
func TestManager_HTTPResponsePool_Nil(t *testing.T) {
|
||||
manager := Get()
|
||||
// Should not panic
|
||||
manager.PutHTTPResponseBuffer(nil)
|
||||
}
|
||||
|
||||
// TestManager_HTTPResponsePool_Oversized tests rejection of oversized HTTP buffers
|
||||
func TestManager_HTTPResponsePool_Oversized(t *testing.T) {
|
||||
manager := Get()
|
||||
manager.ResetStats()
|
||||
|
||||
// Create oversized buffer
|
||||
buf := make([]byte, 0, 40000)
|
||||
manager.PutHTTPResponseBuffer(buf)
|
||||
|
||||
stats := manager.GetStats()
|
||||
if stats.OversizedRejects == 0 {
|
||||
t.Error("Oversized HTTP buffer should be rejected")
|
||||
}
|
||||
}
|
||||
|
||||
// TestManager_ByteSlicePool tests byte slice pool with dynamic sizing
|
||||
func TestManager_ByteSlicePool(t *testing.T) {
|
||||
manager := Get()
|
||||
|
||||
tests := []int{256, 512, 1024, 2048, 4096, 8192, 16384}
|
||||
|
||||
for _, size := range tests {
|
||||
t.Run(strings.Join([]string{"size", string(rune(size))}, "_"), func(t *testing.T) {
|
||||
slice := manager.GetByteSlice(size)
|
||||
if slice == nil {
|
||||
t.Error("GetByteSlice should not return nil")
|
||||
}
|
||||
|
||||
if len(slice) != size {
|
||||
t.Errorf("Byte slice length %d != requested size %d", len(slice), size)
|
||||
}
|
||||
|
||||
if cap(slice) < size {
|
||||
t.Errorf("Byte slice capacity %d < requested size %d", cap(slice), size)
|
||||
}
|
||||
|
||||
// Use the slice
|
||||
copy(slice, []byte("test data"))
|
||||
|
||||
// Return to pool
|
||||
manager.PutByteSlice(slice)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestManager_ByteSlicePool_CustomSize tests byte slice pool with non-standard sizes
|
||||
func TestManager_ByteSlicePool_CustomSize(t *testing.T) {
|
||||
manager := Get()
|
||||
|
||||
// Test custom size (should round up to power of 2)
|
||||
slice := manager.GetByteSlice(300)
|
||||
if slice == nil {
|
||||
t.Error("GetByteSlice should not return nil")
|
||||
}
|
||||
|
||||
if len(slice) != 300 {
|
||||
t.Errorf("Byte slice length %d != requested size 300", len(slice))
|
||||
}
|
||||
|
||||
// Capacity should be >= 300 (likely 512 as next power of 2)
|
||||
if cap(slice) < 300 {
|
||||
t.Error("Byte slice capacity should be at least 300")
|
||||
}
|
||||
|
||||
manager.PutByteSlice(slice)
|
||||
}
|
||||
|
||||
// TestManager_ByteSlicePool_Nil tests putting nil byte slice
|
||||
func TestManager_ByteSlicePool_Nil(t *testing.T) {
|
||||
manager := Get()
|
||||
// Should not panic
|
||||
manager.PutByteSlice(nil)
|
||||
}
|
||||
|
||||
// TestManager_ByteSlicePool_Oversized tests rejection of oversized byte slices
|
||||
func TestManager_ByteSlicePool_Oversized(t *testing.T) {
|
||||
manager := Get()
|
||||
|
||||
// Create oversized slice
|
||||
slice := make([]byte, 100000)
|
||||
|
||||
// Should not panic and should not be pooled
|
||||
manager.PutByteSlice(slice)
|
||||
}
|
||||
|
||||
// TestManager_Stats tests statistics tracking
|
||||
func TestManager_Stats(t *testing.T) {
|
||||
manager := Get()
|
||||
manager.ResetStats()
|
||||
|
||||
initialStats := manager.GetStats()
|
||||
if initialStats.BufferGets != 0 || initialStats.BufferPuts != 0 {
|
||||
t.Error("Stats should be zero after reset")
|
||||
}
|
||||
|
||||
// Perform operations
|
||||
buf := manager.GetBuffer(1024)
|
||||
manager.PutBuffer(buf)
|
||||
|
||||
writer := manager.GetGzipWriter()
|
||||
manager.PutGzipWriter(writer)
|
||||
|
||||
sb := manager.GetStringBuilder()
|
||||
manager.PutStringBuilder(sb)
|
||||
|
||||
jwtBuf := manager.GetJWTBuffer()
|
||||
manager.PutJWTBuffer(jwtBuf)
|
||||
|
||||
httpBuf := manager.GetHTTPResponseBuffer()
|
||||
manager.PutHTTPResponseBuffer(httpBuf)
|
||||
|
||||
// Check stats
|
||||
stats := manager.GetStats()
|
||||
if stats.BufferGets == 0 || stats.BufferPuts == 0 {
|
||||
t.Error("Buffer stats should be incremented")
|
||||
}
|
||||
if stats.GzipGets == 0 || stats.GzipPuts == 0 {
|
||||
t.Error("Gzip stats should be incremented")
|
||||
}
|
||||
if stats.StringGets == 0 || stats.StringPuts == 0 {
|
||||
t.Error("String stats should be incremented")
|
||||
}
|
||||
if stats.JWTGets == 0 || stats.JWTPuts == 0 {
|
||||
t.Error("JWT stats should be incremented")
|
||||
}
|
||||
if stats.HTTPGets == 0 || stats.HTTPPuts == 0 {
|
||||
t.Error("HTTP stats should be incremented")
|
||||
}
|
||||
}
|
||||
|
||||
// TestManager_ResetStats tests statistics reset
|
||||
func TestManager_ResetStats(t *testing.T) {
|
||||
manager := Get()
|
||||
|
||||
// Perform some operations
|
||||
buf := manager.GetBuffer(1024)
|
||||
manager.PutBuffer(buf)
|
||||
|
||||
// Check that stats are non-zero
|
||||
stats := manager.GetStats()
|
||||
if stats.BufferGets == 0 {
|
||||
t.Error("Stats should be non-zero before reset")
|
||||
}
|
||||
|
||||
// Reset stats
|
||||
manager.ResetStats()
|
||||
|
||||
// Check that stats are zero
|
||||
resetStats := manager.GetStats()
|
||||
if resetStats.BufferGets != 0 || resetStats.BufferPuts != 0 {
|
||||
t.Error("Stats should be zero after reset")
|
||||
}
|
||||
}
|
||||
|
||||
// TestManager_ConcurrentAccess tests concurrent access to pools
|
||||
func TestManager_ConcurrentAccess(t *testing.T) {
|
||||
manager := Get()
|
||||
manager.ResetStats()
|
||||
|
||||
var wg sync.WaitGroup
|
||||
numGoroutines := 50
|
||||
operationsPerGoroutine := 10
|
||||
|
||||
wg.Add(numGoroutines)
|
||||
for i := 0; i < numGoroutines; i++ {
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
for j := 0; j < operationsPerGoroutine; j++ {
|
||||
// Test buffer pool
|
||||
buf := manager.GetBuffer(1024)
|
||||
buf.WriteString("test")
|
||||
manager.PutBuffer(buf)
|
||||
|
||||
// Test string builder pool
|
||||
sb := manager.GetStringBuilder()
|
||||
sb.WriteString("test")
|
||||
manager.PutStringBuilder(sb)
|
||||
|
||||
// Test JWT buffer pool
|
||||
jwtBuf := manager.GetJWTBuffer()
|
||||
jwtBuf.Header = append(jwtBuf.Header, byte(j))
|
||||
manager.PutJWTBuffer(jwtBuf)
|
||||
|
||||
// Test byte slice pool
|
||||
slice := manager.GetByteSlice(256)
|
||||
slice[0] = byte(j)
|
||||
manager.PutByteSlice(slice)
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
wg.Wait()
|
||||
|
||||
// Check that operations completed without panic
|
||||
stats := manager.GetStats()
|
||||
expectedOps := uint64(numGoroutines * operationsPerGoroutine)
|
||||
if stats.BufferGets < expectedOps || stats.StringGets < expectedOps || stats.JWTGets < expectedOps {
|
||||
t.Error("Some operations may have failed during concurrent access")
|
||||
}
|
||||
}
|
||||
|
||||
// TestGlobalConvenienceFunctions tests the global convenience functions
|
||||
func TestGlobalConvenienceFunctions(t *testing.T) {
|
||||
// Test buffer functions
|
||||
buf := Buffer(1024)
|
||||
if buf == nil {
|
||||
t.Error("Buffer() should not return nil")
|
||||
}
|
||||
buf.WriteString("test")
|
||||
ReturnBuffer(buf)
|
||||
|
||||
// Test gzip functions
|
||||
writer := GzipWriter()
|
||||
if writer == nil {
|
||||
t.Error("GzipWriter() should not return nil")
|
||||
}
|
||||
ReturnGzipWriter(writer)
|
||||
|
||||
// Test string builder functions
|
||||
sb := StringBuilder()
|
||||
if sb == nil {
|
||||
t.Error("StringBuilder() should not return nil")
|
||||
}
|
||||
sb.WriteString("test")
|
||||
ReturnStringBuilder(sb)
|
||||
|
||||
// Test JWT buffer functions
|
||||
jwtBuf := JWTBuffers()
|
||||
if jwtBuf == nil {
|
||||
t.Error("JWTBuffers() should not return nil")
|
||||
}
|
||||
ReturnJWTBuffers(jwtBuf)
|
||||
|
||||
// Test HTTP buffer functions
|
||||
httpBuf := HTTPBuffer()
|
||||
if httpBuf == nil {
|
||||
t.Error("HTTPBuffer() should not return nil")
|
||||
}
|
||||
ReturnHTTPBuffer(httpBuf)
|
||||
|
||||
// Test byte slice functions
|
||||
slice := ByteSlice(256)
|
||||
if slice == nil {
|
||||
t.Error("ByteSlice() should not return nil")
|
||||
}
|
||||
if len(slice) != 256 {
|
||||
t.Error("ByteSlice() should return correct size")
|
||||
}
|
||||
ReturnByteSlice(slice)
|
||||
}
|
||||
|
||||
// Benchmark tests for performance verification
|
||||
func BenchmarkManager_GetBuffer(b *testing.B) {
|
||||
manager := Get()
|
||||
b.ResetTimer()
|
||||
|
||||
for i := 0; i < b.N; i++ {
|
||||
buf := manager.GetBuffer(1024)
|
||||
manager.PutBuffer(buf)
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkManager_GetStringBuilder(b *testing.B) {
|
||||
manager := Get()
|
||||
b.ResetTimer()
|
||||
|
||||
for i := 0; i < b.N; i++ {
|
||||
sb := manager.GetStringBuilder()
|
||||
manager.PutStringBuilder(sb)
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkManager_GetJWTBuffer(b *testing.B) {
|
||||
manager := Get()
|
||||
b.ResetTimer()
|
||||
|
||||
for i := 0; i < b.N; i++ {
|
||||
jwtBuf := manager.GetJWTBuffer()
|
||||
manager.PutJWTBuffer(jwtBuf)
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkManager_GetByteSlice(b *testing.B) {
|
||||
manager := Get()
|
||||
b.ResetTimer()
|
||||
|
||||
for i := 0; i < b.N; i++ {
|
||||
slice := manager.GetByteSlice(1024)
|
||||
manager.PutByteSlice(slice)
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkManager_ConcurrentAccess(b *testing.B) {
|
||||
manager := Get()
|
||||
b.ResetTimer()
|
||||
|
||||
b.RunParallel(func(pb *testing.PB) {
|
||||
for pb.Next() {
|
||||
buf := manager.GetBuffer(1024)
|
||||
buf.WriteString("test")
|
||||
manager.PutBuffer(buf)
|
||||
}
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,370 @@
|
||||
package pool
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/tls"
|
||||
"net"
|
||||
"net/http"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
)
|
||||
|
||||
// TransportPool manages a pool of shared HTTP transports to prevent connection exhaustion
|
||||
// and resource leaks. It provides centralized management of HTTP client transports with
|
||||
// proper lifecycle management and security controls.
|
||||
type TransportPool struct {
|
||||
mu sync.RWMutex
|
||||
transports map[string]*sharedTransport
|
||||
maxConns int
|
||||
ctx context.Context
|
||||
cancel context.CancelFunc
|
||||
clientCount int32 // Track total HTTP clients
|
||||
maxClients int32 // Limit total clients
|
||||
}
|
||||
|
||||
// sharedTransport wraps an HTTP transport with reference counting
|
||||
type sharedTransport struct {
|
||||
transport *http.Transport
|
||||
refCount int32
|
||||
lastUsed time.Time
|
||||
config TransportConfig
|
||||
}
|
||||
|
||||
// TransportConfig defines configuration for HTTP transports
|
||||
type TransportConfig struct {
|
||||
// Timeouts
|
||||
DialTimeout time.Duration
|
||||
TLSHandshakeTimeout time.Duration
|
||||
ResponseHeaderTimeout time.Duration
|
||||
ExpectContinueTimeout time.Duration
|
||||
IdleConnTimeout time.Duration
|
||||
KeepAlive time.Duration
|
||||
|
||||
// Connection limits
|
||||
MaxIdleConns int
|
||||
MaxIdleConnsPerHost int
|
||||
MaxConnsPerHost int
|
||||
|
||||
// Features
|
||||
ForceHTTP2 bool
|
||||
DisableKeepAlives bool
|
||||
DisableCompression bool
|
||||
|
||||
// Buffer sizes
|
||||
WriteBufferSize int
|
||||
ReadBufferSize int
|
||||
|
||||
// TLS
|
||||
InsecureSkipVerify bool
|
||||
MinTLSVersion uint16
|
||||
}
|
||||
|
||||
var (
|
||||
// globalTransportPool is the singleton transport pool instance
|
||||
globalTransportPool *TransportPool
|
||||
// transportPoolOnce ensures single initialization
|
||||
transportPoolOnce sync.Once
|
||||
)
|
||||
|
||||
// GetTransportPool returns the global transport pool instance
|
||||
func GetTransportPool() *TransportPool {
|
||||
transportPoolOnce.Do(func() {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
globalTransportPool = &TransportPool{
|
||||
transports: make(map[string]*sharedTransport),
|
||||
maxConns: 20,
|
||||
ctx: ctx,
|
||||
cancel: cancel,
|
||||
clientCount: 0,
|
||||
maxClients: 5,
|
||||
}
|
||||
go globalTransportPool.cleanupRoutine(ctx)
|
||||
})
|
||||
return globalTransportPool
|
||||
}
|
||||
|
||||
// DefaultTransportConfig returns a secure default configuration
|
||||
func DefaultTransportConfig() TransportConfig {
|
||||
return TransportConfig{
|
||||
DialTimeout: 30 * time.Second,
|
||||
TLSHandshakeTimeout: 10 * time.Second,
|
||||
ResponseHeaderTimeout: 10 * time.Second,
|
||||
ExpectContinueTimeout: 1 * time.Second,
|
||||
IdleConnTimeout: 30 * time.Second,
|
||||
KeepAlive: 30 * time.Second,
|
||||
MaxIdleConns: 10,
|
||||
MaxIdleConnsPerHost: 2,
|
||||
MaxConnsPerHost: 5,
|
||||
ForceHTTP2: true,
|
||||
DisableKeepAlives: false,
|
||||
DisableCompression: false,
|
||||
WriteBufferSize: 4096,
|
||||
ReadBufferSize: 4096,
|
||||
InsecureSkipVerify: false,
|
||||
MinTLSVersion: tls.VersionTLS12,
|
||||
}
|
||||
}
|
||||
|
||||
// GetTransport gets or creates a shared transport with the given config
|
||||
func (p *TransportPool) GetTransport(config TransportConfig) *http.Transport {
|
||||
// Check client limit
|
||||
if atomic.LoadInt32(&p.clientCount) >= p.maxClients {
|
||||
return p.getExistingTransport()
|
||||
}
|
||||
|
||||
key := p.configKey(config)
|
||||
|
||||
// Fast path: check with read lock
|
||||
p.mu.RLock()
|
||||
if shared, exists := p.transports[key]; exists {
|
||||
atomic.AddInt32(&shared.refCount, 1)
|
||||
shared.lastUsed = time.Now()
|
||||
p.mu.RUnlock()
|
||||
return shared.transport
|
||||
}
|
||||
p.mu.RUnlock()
|
||||
|
||||
// Slow path: create new transport
|
||||
p.mu.Lock()
|
||||
defer p.mu.Unlock()
|
||||
|
||||
// Double-check after acquiring write lock
|
||||
if shared, exists := p.transports[key]; exists {
|
||||
atomic.AddInt32(&shared.refCount, 1)
|
||||
shared.lastUsed = time.Now()
|
||||
return shared.transport
|
||||
}
|
||||
|
||||
// Create new transport
|
||||
transport := p.createTransport(config)
|
||||
shared := &sharedTransport{
|
||||
transport: transport,
|
||||
refCount: 1,
|
||||
lastUsed: time.Now(),
|
||||
config: config,
|
||||
}
|
||||
|
||||
p.transports[key] = shared
|
||||
atomic.AddInt32(&p.clientCount, 1)
|
||||
|
||||
return transport
|
||||
}
|
||||
|
||||
// ReleaseTransport decrements the reference count for a transport
|
||||
func (p *TransportPool) ReleaseTransport(transport *http.Transport) {
|
||||
if transport == nil {
|
||||
return
|
||||
}
|
||||
|
||||
p.mu.RLock()
|
||||
defer p.mu.RUnlock()
|
||||
|
||||
for _, shared := range p.transports {
|
||||
if shared.transport == transport {
|
||||
count := atomic.AddInt32(&shared.refCount, -1)
|
||||
if count <= 0 {
|
||||
shared.lastUsed = time.Now()
|
||||
}
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// getExistingTransport returns any available transport when limit is reached
|
||||
func (p *TransportPool) getExistingTransport() *http.Transport {
|
||||
p.mu.RLock()
|
||||
defer p.mu.RUnlock()
|
||||
|
||||
for _, shared := range p.transports {
|
||||
if shared != nil && shared.transport != nil {
|
||||
atomic.AddInt32(&shared.refCount, 1)
|
||||
shared.lastUsed = time.Now()
|
||||
return shared.transport
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// createTransport creates a new HTTP transport with the given config
|
||||
func (p *TransportPool) createTransport(config TransportConfig) *http.Transport {
|
||||
// Set secure defaults
|
||||
if config.MinTLSVersion == 0 {
|
||||
config.MinTLSVersion = tls.VersionTLS12
|
||||
}
|
||||
|
||||
tlsConfig := &tls.Config{
|
||||
MinVersion: config.MinTLSVersion,
|
||||
MaxVersion: tls.VersionTLS13,
|
||||
CipherSuites: []uint16{
|
||||
tls.TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384,
|
||||
tls.TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256,
|
||||
tls.TLS_ECDHE_ECDSA_WITH_AES_256_GCM_SHA384,
|
||||
tls.TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256,
|
||||
},
|
||||
PreferServerCipherSuites: true,
|
||||
InsecureSkipVerify: config.InsecureSkipVerify,
|
||||
}
|
||||
|
||||
return &http.Transport{
|
||||
Proxy: http.ProxyFromEnvironment,
|
||||
DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
|
||||
dialer := &net.Dialer{
|
||||
Timeout: config.DialTimeout,
|
||||
KeepAlive: config.KeepAlive,
|
||||
}
|
||||
return dialer.DialContext(ctx, network, addr)
|
||||
},
|
||||
TLSClientConfig: tlsConfig,
|
||||
ForceAttemptHTTP2: config.ForceHTTP2,
|
||||
TLSHandshakeTimeout: config.TLSHandshakeTimeout,
|
||||
ExpectContinueTimeout: config.ExpectContinueTimeout,
|
||||
MaxIdleConns: config.MaxIdleConns,
|
||||
MaxIdleConnsPerHost: config.MaxIdleConnsPerHost,
|
||||
IdleConnTimeout: config.IdleConnTimeout,
|
||||
DisableKeepAlives: config.DisableKeepAlives,
|
||||
MaxConnsPerHost: config.MaxConnsPerHost,
|
||||
ResponseHeaderTimeout: config.ResponseHeaderTimeout,
|
||||
DisableCompression: config.DisableCompression,
|
||||
WriteBufferSize: config.WriteBufferSize,
|
||||
ReadBufferSize: config.ReadBufferSize,
|
||||
}
|
||||
}
|
||||
|
||||
// configKey generates a unique key for a transport config
|
||||
func (p *TransportPool) configKey(config TransportConfig) string {
|
||||
// Create a simple key based on critical parameters
|
||||
sb := Get().GetStringBuilder()
|
||||
defer Get().PutStringBuilder(sb)
|
||||
|
||||
sb.WriteByte(byte(config.MaxConnsPerHost))
|
||||
sb.WriteByte(byte(config.MaxIdleConnsPerHost))
|
||||
sb.WriteByte(byte(config.MaxIdleConns))
|
||||
if config.ForceHTTP2 {
|
||||
sb.WriteByte(1)
|
||||
} else {
|
||||
sb.WriteByte(0)
|
||||
}
|
||||
if config.DisableKeepAlives {
|
||||
sb.WriteByte(1)
|
||||
} else {
|
||||
sb.WriteByte(0)
|
||||
}
|
||||
if config.DisableCompression {
|
||||
sb.WriteByte(1)
|
||||
} else {
|
||||
sb.WriteByte(0)
|
||||
}
|
||||
if config.InsecureSkipVerify {
|
||||
sb.WriteByte(1)
|
||||
} else {
|
||||
sb.WriteByte(0)
|
||||
}
|
||||
|
||||
return sb.String()
|
||||
}
|
||||
|
||||
// cleanupRoutine periodically cleans up unused transports
|
||||
func (p *TransportPool) cleanupRoutine(ctx context.Context) {
|
||||
ticker := time.NewTicker(1 * time.Minute)
|
||||
defer ticker.Stop()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
p.cleanup()
|
||||
return
|
||||
case <-ticker.C:
|
||||
p.cleanupIdle()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// cleanupIdle removes idle transports
|
||||
func (p *TransportPool) cleanupIdle() {
|
||||
p.mu.Lock()
|
||||
defer p.mu.Unlock()
|
||||
|
||||
now := time.Now()
|
||||
for key, shared := range p.transports {
|
||||
refCount := atomic.LoadInt32(&shared.refCount)
|
||||
if refCount <= 0 && now.Sub(shared.lastUsed) > 2*time.Minute {
|
||||
shared.transport.CloseIdleConnections()
|
||||
delete(p.transports, key)
|
||||
atomic.AddInt32(&p.clientCount, -1)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// cleanup closes all transports
|
||||
func (p *TransportPool) cleanup() {
|
||||
p.mu.Lock()
|
||||
defer p.mu.Unlock()
|
||||
|
||||
for _, shared := range p.transports {
|
||||
shared.transport.CloseIdleConnections()
|
||||
}
|
||||
p.transports = make(map[string]*sharedTransport)
|
||||
atomic.StoreInt32(&p.clientCount, 0)
|
||||
}
|
||||
|
||||
// Shutdown gracefully shuts down the transport pool
|
||||
func (p *TransportPool) Shutdown() {
|
||||
if p.cancel != nil {
|
||||
p.cancel()
|
||||
}
|
||||
}
|
||||
|
||||
// Stats returns transport pool statistics
|
||||
type TransportPoolStats struct {
|
||||
ActiveTransports int
|
||||
TotalClients int32
|
||||
MaxClients int32
|
||||
}
|
||||
|
||||
// GetStats returns current pool statistics
|
||||
func (p *TransportPool) GetStats() TransportPoolStats {
|
||||
p.mu.RLock()
|
||||
defer p.mu.RUnlock()
|
||||
|
||||
activeCount := 0
|
||||
for _, shared := range p.transports {
|
||||
if atomic.LoadInt32(&shared.refCount) > 0 {
|
||||
activeCount++
|
||||
}
|
||||
}
|
||||
|
||||
return TransportPoolStats{
|
||||
ActiveTransports: activeCount,
|
||||
TotalClients: atomic.LoadInt32(&p.clientCount),
|
||||
MaxClients: p.maxClients,
|
||||
}
|
||||
}
|
||||
|
||||
// CreateHTTPClient creates an HTTP client using the transport pool
|
||||
func CreateHTTPClient(config TransportConfig, timeout time.Duration) *http.Client {
|
||||
pool := GetTransportPool()
|
||||
transport := pool.GetTransport(config)
|
||||
|
||||
if transport == nil {
|
||||
// Fallback to a basic client if pool is exhausted
|
||||
return &http.Client{
|
||||
Timeout: timeout,
|
||||
}
|
||||
}
|
||||
|
||||
client := &http.Client{
|
||||
Transport: transport,
|
||||
Timeout: timeout,
|
||||
}
|
||||
|
||||
// Configure redirect policy
|
||||
client.CheckRedirect = func(req *http.Request, via []*http.Request) error {
|
||||
if len(via) >= 10 {
|
||||
return http.ErrUseLastResponse
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
return client
|
||||
}
|
||||
@@ -0,0 +1,593 @@
|
||||
package pool
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/tls"
|
||||
"net/http"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// TestGetTransportPool_Singleton tests that GetTransportPool returns the same instance
|
||||
func TestGetTransportPool_Singleton(t *testing.T) {
|
||||
pool1 := GetTransportPool()
|
||||
pool2 := GetTransportPool()
|
||||
|
||||
if pool1 != pool2 {
|
||||
t.Error("GetTransportPool() should return the same instance (singleton)")
|
||||
}
|
||||
|
||||
if pool1 == nil {
|
||||
t.Error("GetTransportPool() should not return nil")
|
||||
}
|
||||
}
|
||||
|
||||
// TestDefaultTransportConfig tests the default transport configuration
|
||||
func TestDefaultTransportConfig(t *testing.T) {
|
||||
config := DefaultTransportConfig()
|
||||
|
||||
// Verify security defaults
|
||||
if config.MinTLSVersion != tls.VersionTLS12 {
|
||||
t.Errorf("Default MinTLSVersion should be TLS 1.2, got %d", config.MinTLSVersion)
|
||||
}
|
||||
|
||||
if config.InsecureSkipVerify {
|
||||
t.Error("Default should not skip TLS verification")
|
||||
}
|
||||
|
||||
if !config.ForceHTTP2 {
|
||||
t.Error("Default should force HTTP/2")
|
||||
}
|
||||
|
||||
// Verify reasonable timeouts
|
||||
if config.DialTimeout <= 0 {
|
||||
t.Error("DialTimeout should be positive")
|
||||
}
|
||||
|
||||
if config.TLSHandshakeTimeout <= 0 {
|
||||
t.Error("TLSHandshakeTimeout should be positive")
|
||||
}
|
||||
|
||||
if config.ResponseHeaderTimeout <= 0 {
|
||||
t.Error("ResponseHeaderTimeout should be positive")
|
||||
}
|
||||
|
||||
// Verify connection limits
|
||||
if config.MaxIdleConns <= 0 {
|
||||
t.Error("MaxIdleConns should be positive")
|
||||
}
|
||||
|
||||
if config.MaxIdleConnsPerHost <= 0 {
|
||||
t.Error("MaxIdleConnsPerHost should be positive")
|
||||
}
|
||||
|
||||
if config.MaxConnsPerHost <= 0 {
|
||||
t.Error("MaxConnsPerHost should be positive")
|
||||
}
|
||||
}
|
||||
|
||||
// TestTransportPool_GetTransport tests transport creation and reuse
|
||||
func TestTransportPool_GetTransport(t *testing.T) {
|
||||
pool := &TransportPool{
|
||||
transports: make(map[string]*sharedTransport),
|
||||
maxConns: 20,
|
||||
maxClients: 5,
|
||||
}
|
||||
|
||||
config := DefaultTransportConfig()
|
||||
|
||||
// First call should create new transport
|
||||
transport1 := pool.GetTransport(config)
|
||||
if transport1 == nil {
|
||||
t.Error("GetTransport should not return nil")
|
||||
}
|
||||
|
||||
// Second call with same config should return same transport
|
||||
transport2 := pool.GetTransport(config)
|
||||
if transport2 == nil {
|
||||
t.Error("GetTransport should not return nil")
|
||||
}
|
||||
|
||||
if transport1 != transport2 {
|
||||
t.Error("GetTransport should return same transport for same config")
|
||||
}
|
||||
|
||||
// Verify reference counting
|
||||
pool.mu.RLock()
|
||||
key := pool.configKey(config)
|
||||
shared := pool.transports[key]
|
||||
refCount := shared.refCount
|
||||
pool.mu.RUnlock()
|
||||
|
||||
if refCount != 2 {
|
||||
t.Errorf("Reference count should be 2, got %d", refCount)
|
||||
}
|
||||
}
|
||||
|
||||
// TestTransportPool_GetTransport_DifferentConfigs tests transport creation with different configs
|
||||
func TestTransportPool_GetTransport_DifferentConfigs(t *testing.T) {
|
||||
pool := &TransportPool{
|
||||
transports: make(map[string]*sharedTransport),
|
||||
maxConns: 20,
|
||||
maxClients: 5,
|
||||
}
|
||||
|
||||
config1 := DefaultTransportConfig()
|
||||
config2 := DefaultTransportConfig()
|
||||
config2.MaxConnsPerHost = 10 // Different from default
|
||||
|
||||
transport1 := pool.GetTransport(config1)
|
||||
transport2 := pool.GetTransport(config2)
|
||||
|
||||
if transport1 == transport2 {
|
||||
t.Error("Different configs should produce different transports")
|
||||
}
|
||||
}
|
||||
|
||||
// TestTransportPool_GetTransport_ClientLimit tests client limit enforcement
|
||||
func TestTransportPool_GetTransport_ClientLimit(t *testing.T) {
|
||||
pool := &TransportPool{
|
||||
transports: make(map[string]*sharedTransport),
|
||||
maxConns: 20,
|
||||
maxClients: 2, // Low limit for testing
|
||||
clientCount: 2, // Already at limit
|
||||
}
|
||||
|
||||
config := DefaultTransportConfig()
|
||||
|
||||
// Should return existing transport when limit reached
|
||||
transport := pool.GetTransport(config)
|
||||
// Transport might be nil if no existing transports
|
||||
if transport != nil && pool.clientCount > pool.maxClients {
|
||||
t.Error("Should not exceed client limit")
|
||||
}
|
||||
}
|
||||
|
||||
// TestTransportPool_ReleaseTransport tests transport reference counting
|
||||
func TestTransportPool_ReleaseTransport(t *testing.T) {
|
||||
pool := &TransportPool{
|
||||
transports: make(map[string]*sharedTransport),
|
||||
maxConns: 20,
|
||||
maxClients: 5,
|
||||
}
|
||||
|
||||
config := DefaultTransportConfig()
|
||||
|
||||
// Get transport
|
||||
transport := pool.GetTransport(config)
|
||||
if transport == nil {
|
||||
t.Error("GetTransport should not return nil")
|
||||
}
|
||||
|
||||
// Release transport
|
||||
pool.ReleaseTransport(transport)
|
||||
|
||||
// Verify reference count decreased
|
||||
pool.mu.RLock()
|
||||
key := pool.configKey(config)
|
||||
shared := pool.transports[key]
|
||||
refCount := shared.refCount
|
||||
pool.mu.RUnlock()
|
||||
|
||||
if refCount != 0 {
|
||||
t.Errorf("Reference count should be 0 after release, got %d", refCount)
|
||||
}
|
||||
}
|
||||
|
||||
// TestTransportPool_ReleaseTransport_Nil tests releasing nil transport
|
||||
func TestTransportPool_ReleaseTransport_Nil(t *testing.T) {
|
||||
pool := &TransportPool{
|
||||
transports: make(map[string]*sharedTransport),
|
||||
maxConns: 20,
|
||||
maxClients: 5,
|
||||
}
|
||||
|
||||
// Should not panic
|
||||
pool.ReleaseTransport(nil)
|
||||
}
|
||||
|
||||
// TestTransportPool_ReleaseTransport_Unknown tests releasing unknown transport
|
||||
func TestTransportPool_ReleaseTransport_Unknown(t *testing.T) {
|
||||
pool := &TransportPool{
|
||||
transports: make(map[string]*sharedTransport),
|
||||
maxConns: 20,
|
||||
maxClients: 5,
|
||||
}
|
||||
|
||||
// Create a transport not from the pool
|
||||
transport := &http.Transport{}
|
||||
|
||||
// Should not panic
|
||||
pool.ReleaseTransport(transport)
|
||||
}
|
||||
|
||||
// TestTransportPool_createTransport tests transport creation with different configs
|
||||
func TestTransportPool_createTransport(t *testing.T) {
|
||||
pool := &TransportPool{}
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
config TransportConfig
|
||||
}{
|
||||
{
|
||||
"default config",
|
||||
DefaultTransportConfig(),
|
||||
},
|
||||
{
|
||||
"custom timeouts",
|
||||
TransportConfig{
|
||||
DialTimeout: 10 * time.Second,
|
||||
TLSHandshakeTimeout: 5 * time.Second,
|
||||
MinTLSVersion: tls.VersionTLS13,
|
||||
},
|
||||
},
|
||||
{
|
||||
"insecure config",
|
||||
TransportConfig{
|
||||
InsecureSkipVerify: true,
|
||||
MinTLSVersion: tls.VersionTLS10,
|
||||
},
|
||||
},
|
||||
{
|
||||
"no HTTP/2",
|
||||
TransportConfig{
|
||||
ForceHTTP2: false,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
transport := pool.createTransport(test.config)
|
||||
|
||||
if transport == nil {
|
||||
t.Error("createTransport should not return nil")
|
||||
return
|
||||
}
|
||||
|
||||
// Verify TLS config
|
||||
if transport.TLSClientConfig == nil {
|
||||
t.Error("Transport should have TLS config")
|
||||
return
|
||||
}
|
||||
|
||||
// Verify minimum TLS version
|
||||
expectedMinVersion := test.config.MinTLSVersion
|
||||
if expectedMinVersion == 0 {
|
||||
expectedMinVersion = tls.VersionTLS12 // Default
|
||||
}
|
||||
if transport.TLSClientConfig.MinVersion != expectedMinVersion {
|
||||
t.Errorf("TLS MinVersion should be %d, got %d", expectedMinVersion, transport.TLSClientConfig.MinVersion)
|
||||
}
|
||||
|
||||
// Verify max TLS version
|
||||
if transport.TLSClientConfig.MaxVersion != tls.VersionTLS13 {
|
||||
t.Errorf("TLS MaxVersion should be %d, got %d", tls.VersionTLS13, transport.TLSClientConfig.MaxVersion)
|
||||
}
|
||||
|
||||
// Verify InsecureSkipVerify
|
||||
if transport.TLSClientConfig.InsecureSkipVerify != test.config.InsecureSkipVerify {
|
||||
t.Errorf("InsecureSkipVerify should be %v, got %v", test.config.InsecureSkipVerify, transport.TLSClientConfig.InsecureSkipVerify)
|
||||
}
|
||||
|
||||
// Verify HTTP/2
|
||||
if transport.ForceAttemptHTTP2 != test.config.ForceHTTP2 {
|
||||
t.Errorf("ForceAttemptHTTP2 should be %v, got %v", test.config.ForceHTTP2, transport.ForceAttemptHTTP2)
|
||||
}
|
||||
|
||||
// Verify timeouts
|
||||
if test.config.TLSHandshakeTimeout > 0 && transport.TLSHandshakeTimeout != test.config.TLSHandshakeTimeout {
|
||||
t.Errorf("TLSHandshakeTimeout should be %v, got %v", test.config.TLSHandshakeTimeout, transport.TLSHandshakeTimeout)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestTransportPool_configKey tests configuration key generation
|
||||
func TestTransportPool_configKey(t *testing.T) {
|
||||
pool := &TransportPool{}
|
||||
|
||||
config1 := DefaultTransportConfig()
|
||||
config2 := DefaultTransportConfig()
|
||||
|
||||
key1 := pool.configKey(config1)
|
||||
key2 := pool.configKey(config2)
|
||||
|
||||
if key1 != key2 {
|
||||
t.Error("Same configs should generate same key")
|
||||
}
|
||||
|
||||
// Different config
|
||||
config3 := config1
|
||||
config3.MaxConnsPerHost = 999
|
||||
key3 := pool.configKey(config3)
|
||||
|
||||
if key1 == key3 {
|
||||
t.Error("Different configs should generate different keys")
|
||||
}
|
||||
}
|
||||
|
||||
// TestTransportPool_cleanupIdle tests idle transport cleanup
|
||||
func TestTransportPool_cleanupIdle(t *testing.T) {
|
||||
pool := &TransportPool{
|
||||
transports: make(map[string]*sharedTransport),
|
||||
maxConns: 20,
|
||||
maxClients: 5,
|
||||
}
|
||||
|
||||
config := DefaultTransportConfig()
|
||||
transport := pool.createTransport(config)
|
||||
|
||||
// Add transport to pool with old timestamp
|
||||
shared := &sharedTransport{
|
||||
transport: transport,
|
||||
refCount: 0,
|
||||
lastUsed: time.Now().Add(-5 * time.Minute), // Old
|
||||
config: config,
|
||||
}
|
||||
|
||||
key := pool.configKey(config)
|
||||
pool.transports[key] = shared
|
||||
|
||||
// Run cleanup
|
||||
pool.cleanupIdle()
|
||||
|
||||
// Transport should be removed
|
||||
if _, exists := pool.transports[key]; exists {
|
||||
t.Error("Old idle transport should be cleaned up")
|
||||
}
|
||||
}
|
||||
|
||||
// TestTransportPool_cleanup tests full cleanup
|
||||
func TestTransportPool_cleanup(t *testing.T) {
|
||||
pool := &TransportPool{
|
||||
transports: make(map[string]*sharedTransport),
|
||||
maxConns: 20,
|
||||
maxClients: 5,
|
||||
clientCount: 3,
|
||||
}
|
||||
|
||||
config := DefaultTransportConfig()
|
||||
transport := pool.createTransport(config)
|
||||
|
||||
// Add transport to pool
|
||||
shared := &sharedTransport{
|
||||
transport: transport,
|
||||
refCount: 1,
|
||||
lastUsed: time.Now(),
|
||||
config: config,
|
||||
}
|
||||
|
||||
key := pool.configKey(config)
|
||||
pool.transports[key] = shared
|
||||
|
||||
// Run cleanup
|
||||
pool.cleanup()
|
||||
|
||||
// All transports should be removed
|
||||
if len(pool.transports) != 0 {
|
||||
t.Error("All transports should be cleaned up")
|
||||
}
|
||||
|
||||
// Client count should be reset
|
||||
if pool.clientCount != 0 {
|
||||
t.Error("Client count should be reset")
|
||||
}
|
||||
}
|
||||
|
||||
// TestTransportPool_Shutdown tests graceful shutdown
|
||||
func TestTransportPool_Shutdown(t *testing.T) {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
pool := &TransportPool{
|
||||
transports: make(map[string]*sharedTransport),
|
||||
maxConns: 20,
|
||||
maxClients: 5,
|
||||
ctx: ctx,
|
||||
cancel: cancel,
|
||||
}
|
||||
|
||||
// Should not panic
|
||||
pool.Shutdown()
|
||||
}
|
||||
|
||||
// TestTransportPool_GetStats tests statistics
|
||||
func TestTransportPool_GetStats(t *testing.T) {
|
||||
pool := &TransportPool{
|
||||
transports: make(map[string]*sharedTransport),
|
||||
maxConns: 20,
|
||||
maxClients: 5,
|
||||
clientCount: 3,
|
||||
}
|
||||
|
||||
config := DefaultTransportConfig()
|
||||
|
||||
// Add some transports
|
||||
for i := 0; i < 3; i++ {
|
||||
transport := pool.createTransport(config)
|
||||
shared := &sharedTransport{
|
||||
transport: transport,
|
||||
refCount: int32(i % 2), // Some active, some idle
|
||||
lastUsed: time.Now(),
|
||||
config: config,
|
||||
}
|
||||
pool.transports[string(rune(i))] = shared
|
||||
}
|
||||
|
||||
stats := pool.GetStats()
|
||||
|
||||
if stats.TotalClients != 3 {
|
||||
t.Errorf("TotalClients should be 3, got %d", stats.TotalClients)
|
||||
}
|
||||
|
||||
if stats.MaxClients != 5 {
|
||||
t.Errorf("MaxClients should be 5, got %d", stats.MaxClients)
|
||||
}
|
||||
|
||||
if stats.ActiveTransports < 0 || stats.ActiveTransports > 3 {
|
||||
t.Errorf("ActiveTransports should be between 0 and 3, got %d", stats.ActiveTransports)
|
||||
}
|
||||
}
|
||||
|
||||
// TestCreateHTTPClient tests HTTP client creation
|
||||
func TestCreateHTTPClient(t *testing.T) {
|
||||
config := DefaultTransportConfig()
|
||||
timeout := 30 * time.Second
|
||||
|
||||
client := CreateHTTPClient(config, timeout)
|
||||
|
||||
if client == nil {
|
||||
t.Error("CreateHTTPClient should not return nil")
|
||||
return
|
||||
}
|
||||
|
||||
if client.Timeout != timeout {
|
||||
t.Errorf("Client timeout should be %v, got %v", timeout, client.Timeout)
|
||||
}
|
||||
|
||||
if client.Transport == nil {
|
||||
t.Error("Client should have transport")
|
||||
}
|
||||
|
||||
if client.CheckRedirect == nil {
|
||||
t.Error("Client should have redirect policy")
|
||||
}
|
||||
|
||||
// Test redirect policy
|
||||
req := &http.Request{}
|
||||
var via []*http.Request
|
||||
|
||||
// Should allow up to 9 redirects (10 total requests)
|
||||
for i := 0; i < 9; i++ {
|
||||
via = append(via, &http.Request{})
|
||||
err := client.CheckRedirect(req, via)
|
||||
if err != nil {
|
||||
t.Errorf("Should allow %d redirects, got error: %v", i+1, err)
|
||||
}
|
||||
}
|
||||
|
||||
// Should reject 10th redirect (11th total request)
|
||||
via = append(via, &http.Request{})
|
||||
err := client.CheckRedirect(req, via)
|
||||
if err != http.ErrUseLastResponse {
|
||||
t.Error("Should reject too many redirects")
|
||||
}
|
||||
}
|
||||
|
||||
// TestCreateHTTPClient_Fallback tests fallback when pool is exhausted
|
||||
func TestCreateHTTPClient_Fallback(t *testing.T) {
|
||||
// Override global pool with limited one
|
||||
originalPool := globalTransportPool
|
||||
defer func() {
|
||||
globalTransportPool = originalPool
|
||||
}()
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
|
||||
globalTransportPool = &TransportPool{
|
||||
transports: make(map[string]*sharedTransport),
|
||||
maxConns: 20,
|
||||
ctx: ctx,
|
||||
cancel: cancel,
|
||||
clientCount: 10,
|
||||
maxClients: 1, // Very low limit
|
||||
}
|
||||
|
||||
config := DefaultTransportConfig()
|
||||
timeout := 30 * time.Second
|
||||
|
||||
client := CreateHTTPClient(config, timeout)
|
||||
|
||||
if client == nil {
|
||||
t.Error("CreateHTTPClient should not return nil even when pool is exhausted")
|
||||
return
|
||||
}
|
||||
|
||||
if client.Timeout != timeout {
|
||||
t.Errorf("Client timeout should be %v, got %v", timeout, client.Timeout)
|
||||
}
|
||||
}
|
||||
|
||||
// TestTransportPool_ConcurrentAccess tests concurrent access to transport pool
|
||||
func TestTransportPool_ConcurrentAccess(t *testing.T) {
|
||||
pool := &TransportPool{
|
||||
transports: make(map[string]*sharedTransport),
|
||||
maxConns: 20,
|
||||
maxClients: 50, // High limit for concurrent test
|
||||
}
|
||||
|
||||
// Use different configs to reduce contention on single transport
|
||||
baseConfig := DefaultTransportConfig()
|
||||
configs := make([]TransportConfig, 10)
|
||||
for i := range configs {
|
||||
configs[i] = baseConfig
|
||||
configs[i].MaxConnsPerHost = 5 + i // Make each config unique
|
||||
}
|
||||
|
||||
var wg sync.WaitGroup
|
||||
numGoroutines := 10
|
||||
operationsPerGoroutine := 3
|
||||
|
||||
wg.Add(numGoroutines)
|
||||
for i := 0; i < numGoroutines; i++ {
|
||||
go func(goroutineID int) {
|
||||
defer wg.Done()
|
||||
config := configs[goroutineID%len(configs)]
|
||||
for j := 0; j < operationsPerGoroutine; j++ {
|
||||
transport := pool.GetTransport(config)
|
||||
if transport == nil {
|
||||
continue
|
||||
}
|
||||
// Use transport briefly
|
||||
time.Sleep(time.Millisecond)
|
||||
pool.ReleaseTransport(transport)
|
||||
}
|
||||
}(i)
|
||||
}
|
||||
|
||||
wg.Wait()
|
||||
|
||||
// Should not panic and should have reasonable stats
|
||||
stats := pool.GetStats()
|
||||
if stats.TotalClients < 0 || stats.TotalClients > int32(numGoroutines) {
|
||||
t.Errorf("Unexpected client count: %d", stats.TotalClients)
|
||||
}
|
||||
}
|
||||
|
||||
// Benchmark tests for performance verification
|
||||
func BenchmarkTransportPool_GetTransport(b *testing.B) {
|
||||
pool := &TransportPool{
|
||||
transports: make(map[string]*sharedTransport),
|
||||
maxConns: 20,
|
||||
maxClients: 100,
|
||||
}
|
||||
|
||||
config := DefaultTransportConfig()
|
||||
b.ResetTimer()
|
||||
|
||||
for i := 0; i < b.N; i++ {
|
||||
transport := pool.GetTransport(config)
|
||||
pool.ReleaseTransport(transport)
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkCreateHTTPClient(b *testing.B) {
|
||||
config := DefaultTransportConfig()
|
||||
timeout := 30 * time.Second
|
||||
b.ResetTimer()
|
||||
|
||||
for i := 0; i < b.N; i++ {
|
||||
CreateHTTPClient(config, timeout)
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkTransportPool_configKey(b *testing.B) {
|
||||
pool := &TransportPool{}
|
||||
config := DefaultTransportConfig()
|
||||
b.ResetTimer()
|
||||
|
||||
for i := 0; i < b.N; i++ {
|
||||
pool.configKey(config)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,115 @@
|
||||
package providers
|
||||
|
||||
import (
|
||||
"net/url"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Adapter facilitates communication between the legacy TraefikOIDC struct and the new provider system.
|
||||
type Adapter struct {
|
||||
provider OIDCProvider
|
||||
legacySettings LegacySettings
|
||||
tokenVerifier TokenVerifier
|
||||
tokenCache TokenCache
|
||||
}
|
||||
|
||||
// LegacySettings provides the adapter with access to the original configuration values.
|
||||
type LegacySettings interface {
|
||||
GetIssuerURL() string
|
||||
GetAuthURL() string
|
||||
GetScopes() []string
|
||||
IsPKCEEnabled() bool
|
||||
GetClientID() string
|
||||
GetRefreshGracePeriod() time.Duration
|
||||
IsOverrideScopes() bool
|
||||
}
|
||||
|
||||
// NewAdapter creates a new adapter for a given provider and legacy settings.
|
||||
func NewAdapter(provider OIDCProvider, settings LegacySettings, tokenVerifier TokenVerifier, tokenCache TokenCache) *Adapter {
|
||||
return &Adapter{
|
||||
provider: provider,
|
||||
legacySettings: settings,
|
||||
tokenVerifier: tokenVerifier,
|
||||
tokenCache: tokenCache,
|
||||
}
|
||||
}
|
||||
|
||||
// BuildAuthURL constructs the authentication URL using the adapted provider.
|
||||
func (a *Adapter) BuildAuthURL(redirectURL, state, nonce, codeChallenge string) string {
|
||||
params := url.Values{}
|
||||
params.Set("client_id", a.legacySettings.GetClientID())
|
||||
params.Set("response_type", "code")
|
||||
params.Set("redirect_uri", redirectURL)
|
||||
params.Set("state", state)
|
||||
params.Set("nonce", nonce)
|
||||
|
||||
if a.legacySettings.IsPKCEEnabled() && codeChallenge != "" {
|
||||
params.Set("code_challenge", codeChallenge)
|
||||
params.Set("code_challenge_method", "S256")
|
||||
}
|
||||
|
||||
scopes := a.legacySettings.GetScopes()
|
||||
|
||||
if a.legacySettings.IsOverrideScopes() {
|
||||
finalParams := params
|
||||
finalParams.Set("scope", strings.Join(scopes, " "))
|
||||
|
||||
switch a.provider.GetType() {
|
||||
case ProviderTypeGoogle:
|
||||
finalParams.Set("access_type", "offline")
|
||||
finalParams.Set("prompt", "consent")
|
||||
case ProviderTypeAzure:
|
||||
finalParams.Set("response_mode", "query")
|
||||
}
|
||||
|
||||
return a.buildURLWithParams(a.legacySettings.GetAuthURL(), finalParams)
|
||||
}
|
||||
|
||||
authParams, err := a.provider.BuildAuthParams(params, scopes)
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
|
||||
finalParams := authParams.URLValues
|
||||
finalParams.Set("scope", strings.Join(authParams.Scopes, " "))
|
||||
|
||||
return a.buildURLWithParams(a.legacySettings.GetAuthURL(), finalParams)
|
||||
}
|
||||
|
||||
// from the configured issuerURL.
|
||||
func (a *Adapter) buildURLWithParams(baseURL string, params url.Values) string {
|
||||
if !strings.HasPrefix(baseURL, "http://") && !strings.HasPrefix(baseURL, "https://") {
|
||||
issuerURLParsed, err := url.Parse(a.legacySettings.GetIssuerURL())
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
|
||||
baseURLParsed, err := url.Parse(baseURL)
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
|
||||
resolvedURL := issuerURLParsed.ResolveReference(baseURLParsed)
|
||||
resolvedURL.RawQuery = params.Encode()
|
||||
return resolvedURL.String()
|
||||
}
|
||||
|
||||
u, err := url.Parse(baseURL)
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
|
||||
u.RawQuery = params.Encode()
|
||||
return u.String()
|
||||
}
|
||||
|
||||
// ValidateTokens validates tokens using the adapted provider.
|
||||
func (a *Adapter) ValidateTokens(session Session) (*ValidationResult, error) {
|
||||
return a.provider.ValidateTokens(session, a.tokenVerifier, a.tokenCache, a.legacySettings.GetRefreshGracePeriod())
|
||||
}
|
||||
|
||||
// GetType returns the underlying provider's type.
|
||||
func (a *Adapter) GetType() ProviderType {
|
||||
return a.provider.GetType()
|
||||
}
|
||||
@@ -0,0 +1,106 @@
|
||||
package providers
|
||||
|
||||
import (
|
||||
"net/url"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
// AzureProvider encapsulates Azure AD-specific OIDC logic.
|
||||
type AzureProvider struct {
|
||||
*BaseProvider
|
||||
}
|
||||
|
||||
// NewAzureProvider creates a new instance of the AzureProvider.
|
||||
func NewAzureProvider() *AzureProvider {
|
||||
return &AzureProvider{
|
||||
BaseProvider: NewBaseProvider(),
|
||||
}
|
||||
}
|
||||
|
||||
// GetType returns the provider's type.
|
||||
func (p *AzureProvider) GetType() ProviderType {
|
||||
return ProviderTypeAzure
|
||||
}
|
||||
|
||||
// GetCapabilities returns the specific capabilities of the Azure provider.
|
||||
func (p *AzureProvider) GetCapabilities() ProviderCapabilities {
|
||||
return ProviderCapabilities{
|
||||
SupportsRefreshTokens: true,
|
||||
RequiresOfflineAccessScope: true,
|
||||
PreferredTokenValidation: "access",
|
||||
}
|
||||
}
|
||||
|
||||
// BuildAuthParams configures Azure-specific authentication parameters.
|
||||
func (p *AzureProvider) BuildAuthParams(baseParams url.Values, scopes []string) (*AuthParams, error) {
|
||||
baseParams.Set("response_mode", "query")
|
||||
|
||||
hasOfflineAccess := false
|
||||
for _, scope := range scopes {
|
||||
if scope == "offline_access" {
|
||||
hasOfflineAccess = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !hasOfflineAccess {
|
||||
scopes = append(scopes, "offline_access")
|
||||
}
|
||||
|
||||
return &AuthParams{
|
||||
URLValues: baseParams,
|
||||
Scopes: scopes,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Azure may use access tokens for validation, and this method ensures that behavior is preserved.
|
||||
func (p *AzureProvider) ValidateTokens(session Session, verifier TokenVerifier, tokenCache TokenCache, refreshGracePeriod time.Duration) (*ValidationResult, error) {
|
||||
if !session.GetAuthenticated() {
|
||||
if session.GetRefreshToken() != "" {
|
||||
return &ValidationResult{NeedsRefresh: true}, nil
|
||||
}
|
||||
return &ValidationResult{IsExpired: true}, nil
|
||||
}
|
||||
|
||||
accessToken := session.GetAccessToken()
|
||||
idToken := session.GetIDToken()
|
||||
|
||||
if accessToken != "" {
|
||||
if strings.Count(accessToken, ".") == 2 {
|
||||
if err := verifier.VerifyToken(accessToken); err != nil {
|
||||
if idToken != "" {
|
||||
return p.ValidateTokenExpiry(session, idToken, tokenCache, refreshGracePeriod)
|
||||
}
|
||||
if session.GetRefreshToken() != "" {
|
||||
return &ValidationResult{NeedsRefresh: true}, nil
|
||||
}
|
||||
return &ValidationResult{IsExpired: true}, nil
|
||||
}
|
||||
return p.ValidateTokenExpiry(session, accessToken, tokenCache, refreshGracePeriod)
|
||||
}
|
||||
if idToken != "" {
|
||||
return p.ValidateTokenExpiry(session, idToken, tokenCache, refreshGracePeriod)
|
||||
}
|
||||
return &ValidationResult{Authenticated: true}, nil
|
||||
}
|
||||
|
||||
if idToken != "" {
|
||||
if err := verifier.VerifyToken(idToken); err != nil {
|
||||
if session.GetRefreshToken() != "" {
|
||||
return &ValidationResult{NeedsRefresh: true}, nil
|
||||
}
|
||||
return &ValidationResult{IsExpired: true}, nil
|
||||
}
|
||||
return p.ValidateTokenExpiry(session, idToken, tokenCache, refreshGracePeriod)
|
||||
}
|
||||
|
||||
if session.GetRefreshToken() != "" {
|
||||
return &ValidationResult{NeedsRefresh: true}, nil
|
||||
}
|
||||
return &ValidationResult{IsExpired: true}, nil
|
||||
}
|
||||
|
||||
// Azure requires specific tenant configuration and scope handling.
|
||||
func (p *AzureProvider) ValidateConfig() error {
|
||||
return p.BaseProvider.ValidateConfig()
|
||||
}
|
||||
@@ -0,0 +1,140 @@
|
||||
package providers
|
||||
|
||||
import (
|
||||
"net/url"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
// BaseProvider provides common functionality for all OIDC provider implementations.
|
||||
// It defines default behaviors that can be overridden by specific providers.
|
||||
// It can be embedded in specific provider structs to share common logic.
|
||||
type BaseProvider struct {
|
||||
}
|
||||
|
||||
// GetType returns the default provider type (generic).
|
||||
// This should be overridden by specific provider implementations.
|
||||
func (p *BaseProvider) GetType() ProviderType {
|
||||
return ProviderTypeGeneric
|
||||
}
|
||||
|
||||
// GetCapabilities returns default provider capabilities.
|
||||
// This can be overridden by specific providers to declare their unique features.
|
||||
func (p *BaseProvider) GetCapabilities() ProviderCapabilities {
|
||||
return ProviderCapabilities{
|
||||
SupportsRefreshTokens: true,
|
||||
RequiresOfflineAccessScope: true,
|
||||
PreferredTokenValidation: "id",
|
||||
}
|
||||
}
|
||||
|
||||
// ValidateTokens performs basic token validation logic common to all providers.
|
||||
// It checks authentication state, token presence, and determines if refresh is needed.
|
||||
// This method can be extended or replaced by specific providers.
|
||||
func (p *BaseProvider) ValidateTokens(session Session, verifier TokenVerifier, tokenCache TokenCache, refreshGracePeriod time.Duration) (*ValidationResult, error) {
|
||||
if !session.GetAuthenticated() {
|
||||
if session.GetRefreshToken() != "" {
|
||||
return &ValidationResult{NeedsRefresh: true}, nil
|
||||
}
|
||||
return &ValidationResult{}, nil
|
||||
}
|
||||
|
||||
accessToken := session.GetAccessToken()
|
||||
if accessToken == "" {
|
||||
if session.GetRefreshToken() != "" {
|
||||
return &ValidationResult{NeedsRefresh: true}, nil
|
||||
}
|
||||
return &ValidationResult{IsExpired: true}, nil
|
||||
}
|
||||
|
||||
idToken := session.GetIDToken()
|
||||
if idToken == "" {
|
||||
if session.GetRefreshToken() != "" {
|
||||
return &ValidationResult{Authenticated: true, NeedsRefresh: true}, nil
|
||||
}
|
||||
return &ValidationResult{Authenticated: true}, nil
|
||||
}
|
||||
|
||||
if err := verifier.VerifyToken(idToken); err != nil {
|
||||
if strings.Contains(err.Error(), "token has expired") {
|
||||
if session.GetRefreshToken() != "" {
|
||||
return &ValidationResult{NeedsRefresh: true}, nil
|
||||
}
|
||||
return &ValidationResult{IsExpired: true}, nil
|
||||
}
|
||||
if session.GetRefreshToken() != "" {
|
||||
return &ValidationResult{NeedsRefresh: true}, nil
|
||||
}
|
||||
return &ValidationResult{IsExpired: true}, nil
|
||||
}
|
||||
|
||||
return p.ValidateTokenExpiry(session, idToken, tokenCache, refreshGracePeriod)
|
||||
}
|
||||
|
||||
// ValidateTokenExpiry checks if a token is expired or needs refresh based on cached claims.
|
||||
// This method is now exported so provider implementations can reuse this logic without duplication.
|
||||
func (p *BaseProvider) ValidateTokenExpiry(session Session, token string, tokenCache TokenCache, refreshGracePeriod time.Duration) (*ValidationResult, error) {
|
||||
cachedClaims, found := tokenCache.Get(token)
|
||||
if !found {
|
||||
if session.GetRefreshToken() != "" {
|
||||
return &ValidationResult{NeedsRefresh: true}, nil
|
||||
}
|
||||
return &ValidationResult{IsExpired: true}, nil
|
||||
}
|
||||
|
||||
expClaim, ok := cachedClaims["exp"].(float64)
|
||||
if !ok {
|
||||
if session.GetRefreshToken() != "" {
|
||||
return &ValidationResult{NeedsRefresh: true}, nil
|
||||
}
|
||||
return &ValidationResult{IsExpired: true}, nil
|
||||
}
|
||||
|
||||
expTime := time.Unix(int64(expClaim), 0)
|
||||
if expTime.Before(time.Now().Add(refreshGracePeriod)) {
|
||||
if session.GetRefreshToken() != "" {
|
||||
return &ValidationResult{Authenticated: true, NeedsRefresh: true}, nil
|
||||
}
|
||||
return &ValidationResult{Authenticated: true}, nil
|
||||
}
|
||||
|
||||
return &ValidationResult{Authenticated: true}, nil
|
||||
}
|
||||
|
||||
// BuildAuthParams constructs authorization parameters for the provider.
|
||||
// It includes the "offline_access" scope by default for refresh token support.
|
||||
func (p *BaseProvider) BuildAuthParams(baseParams url.Values, scopes []string) (*AuthParams, error) {
|
||||
hasOfflineAccess := false
|
||||
for _, scope := range scopes {
|
||||
if scope == "offline_access" {
|
||||
hasOfflineAccess = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !hasOfflineAccess {
|
||||
scopes = append(scopes, "offline_access")
|
||||
}
|
||||
|
||||
return &AuthParams{
|
||||
URLValues: baseParams,
|
||||
Scopes: scopes,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// HandleTokenRefresh processes provider-specific token refresh logic.
|
||||
// By default, it does nothing and assumes the standard token response is sufficient.
|
||||
func (p *BaseProvider) HandleTokenRefresh(tokenData *TokenResult) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// ValidateConfig checks provider-specific configuration requirements.
|
||||
// By default, it assumes the configuration is valid.
|
||||
func (p *BaseProvider) ValidateConfig() error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// NewBaseProvider creates a new BaseProvider instance.
|
||||
// This can be used when a generic OIDC provider is sufficient.
|
||||
func NewBaseProvider() *BaseProvider {
|
||||
return &BaseProvider{}
|
||||
}
|
||||
@@ -0,0 +1,115 @@
|
||||
package providers
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net/url"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// ProviderFactory encapsulates the logic for creating and configuring OIDC providers.
|
||||
type ProviderFactory struct {
|
||||
registry *ProviderRegistry
|
||||
}
|
||||
|
||||
// NewProviderFactory creates a new factory with a pre-configured registry.
|
||||
func NewProviderFactory() *ProviderFactory {
|
||||
registry := NewProviderRegistry()
|
||||
|
||||
registry.RegisterProvider(NewGenericProvider())
|
||||
registry.RegisterProvider(NewGoogleProvider())
|
||||
registry.RegisterProvider(NewAzureProvider())
|
||||
|
||||
return &ProviderFactory{
|
||||
registry: registry,
|
||||
}
|
||||
}
|
||||
|
||||
// CreateProvider creates an OIDC provider based on the issuer URL.
|
||||
// It automatically detects the provider type and returns a configured instance.
|
||||
func (f *ProviderFactory) CreateProvider(issuerURL string) (OIDCProvider, error) {
|
||||
if issuerURL == "" {
|
||||
return nil, fmt.Errorf("issuer URL cannot be empty")
|
||||
}
|
||||
|
||||
if _, err := url.Parse(issuerURL); err != nil {
|
||||
return nil, fmt.Errorf("invalid issuer URL format: %w", err)
|
||||
}
|
||||
|
||||
provider := f.registry.DetectProvider(issuerURL)
|
||||
if provider == nil {
|
||||
return nil, fmt.Errorf("unable to detect provider for issuer URL: %s", issuerURL)
|
||||
}
|
||||
|
||||
if err := provider.ValidateConfig(); err != nil {
|
||||
return nil, fmt.Errorf("provider configuration validation failed: %w", err)
|
||||
}
|
||||
|
||||
return provider, nil
|
||||
}
|
||||
|
||||
// CreateProviderByType creates a provider instance of the specified type.
|
||||
// This is useful when you want to force a specific provider type regardless of URL.
|
||||
func (f *ProviderFactory) CreateProviderByType(providerType ProviderType) (OIDCProvider, error) {
|
||||
var provider OIDCProvider
|
||||
|
||||
switch providerType {
|
||||
case ProviderTypeGeneric:
|
||||
provider = NewGenericProvider()
|
||||
case ProviderTypeGoogle:
|
||||
provider = NewGoogleProvider()
|
||||
case ProviderTypeAzure:
|
||||
provider = NewAzureProvider()
|
||||
default:
|
||||
return nil, fmt.Errorf("unsupported provider type: %d", providerType)
|
||||
}
|
||||
|
||||
if err := provider.ValidateConfig(); err != nil {
|
||||
return nil, fmt.Errorf("provider configuration validation failed: %w", err)
|
||||
}
|
||||
|
||||
return provider, nil
|
||||
}
|
||||
|
||||
// GetSupportedProviders returns a list of all supported provider types and their detection patterns.
|
||||
func (f *ProviderFactory) GetSupportedProviders() map[ProviderType][]string {
|
||||
return map[ProviderType][]string{
|
||||
ProviderTypeGeneric: {"*"},
|
||||
ProviderTypeGoogle: {"accounts.google.com"},
|
||||
ProviderTypeAzure: {"login.microsoftonline.com", "sts.windows.net"},
|
||||
}
|
||||
}
|
||||
|
||||
// DetectProviderType determines the provider type for a given issuer URL.
|
||||
// This is useful for diagnostic purposes or UI display.
|
||||
func (f *ProviderFactory) DetectProviderType(issuerURL string) (ProviderType, error) {
|
||||
provider, err := f.CreateProvider(issuerURL)
|
||||
if err != nil {
|
||||
return ProviderTypeGeneric, err
|
||||
}
|
||||
return provider.GetType(), nil
|
||||
}
|
||||
|
||||
// IsProviderSupported checks if a given issuer URL is supported by any registered provider.
|
||||
func (f *ProviderFactory) IsProviderSupported(issuerURL string) bool {
|
||||
if issuerURL == "" {
|
||||
return false
|
||||
}
|
||||
|
||||
normalizedURL, err := url.Parse(issuerURL)
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
|
||||
host := strings.ToLower(normalizedURL.Host)
|
||||
supportedProviders := f.GetSupportedProviders()
|
||||
|
||||
for _, patterns := range supportedProviders {
|
||||
for _, pattern := range patterns {
|
||||
if pattern == "*" || strings.Contains(host, strings.ToLower(pattern)) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return false
|
||||
}
|
||||
@@ -0,0 +1,18 @@
|
||||
package providers
|
||||
|
||||
// GenericProvider encapsulates standard OIDC logic for any compliant provider.
|
||||
type GenericProvider struct {
|
||||
*BaseProvider
|
||||
}
|
||||
|
||||
// NewGenericProvider creates a new instance of the GenericProvider.
|
||||
func NewGenericProvider() *GenericProvider {
|
||||
return &GenericProvider{
|
||||
BaseProvider: NewBaseProvider(),
|
||||
}
|
||||
}
|
||||
|
||||
// GetType returns the provider's type.
|
||||
func (p *GenericProvider) GetType() ProviderType {
|
||||
return ProviderTypeGeneric
|
||||
}
|
||||
@@ -0,0 +1,56 @@
|
||||
package providers
|
||||
|
||||
import (
|
||||
"net/url"
|
||||
)
|
||||
|
||||
// GoogleProvider encapsulates Google-specific OIDC logic.
|
||||
type GoogleProvider struct {
|
||||
*BaseProvider
|
||||
}
|
||||
|
||||
// NewGoogleProvider creates a new instance of the GoogleProvider.
|
||||
func NewGoogleProvider() *GoogleProvider {
|
||||
return &GoogleProvider{
|
||||
BaseProvider: NewBaseProvider(),
|
||||
}
|
||||
}
|
||||
|
||||
// GetType returns the provider's type.
|
||||
func (p *GoogleProvider) GetType() ProviderType {
|
||||
return ProviderTypeGoogle
|
||||
}
|
||||
|
||||
// GetCapabilities returns the specific capabilities of the Google provider.
|
||||
func (p *GoogleProvider) GetCapabilities() ProviderCapabilities {
|
||||
return ProviderCapabilities{
|
||||
SupportsRefreshTokens: true,
|
||||
RequiresOfflineAccessScope: false,
|
||||
RequiresPromptConsent: true,
|
||||
PreferredTokenValidation: "id",
|
||||
}
|
||||
}
|
||||
|
||||
// BuildAuthParams configures Google-specific authentication parameters.
|
||||
func (p *GoogleProvider) BuildAuthParams(baseParams url.Values, scopes []string) (*AuthParams, error) {
|
||||
baseParams.Set("access_type", "offline")
|
||||
baseParams.Set("prompt", "consent")
|
||||
|
||||
// Google does not use the "offline_access" scope, so we remove it if present.
|
||||
var filteredScopes []string
|
||||
for _, scope := range scopes {
|
||||
if scope != "offline_access" {
|
||||
filteredScopes = append(filteredScopes, scope)
|
||||
}
|
||||
}
|
||||
|
||||
return &AuthParams{
|
||||
URLValues: baseParams,
|
||||
Scopes: filteredScopes,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Google requires specific scopes and client configuration for proper operation.
|
||||
func (p *GoogleProvider) ValidateConfig() error {
|
||||
return p.BaseProvider.ValidateConfig()
|
||||
}
|
||||
@@ -0,0 +1,79 @@
|
||||
// Package providers implements a universal OIDC provider abstraction system.
|
||||
// It provides a clean interface for different OIDC providers (Google, Azure, Generic)
|
||||
// with provider-specific logic encapsulated in separate implementations.
|
||||
package providers
|
||||
|
||||
import (
|
||||
"net/url"
|
||||
"time"
|
||||
)
|
||||
|
||||
// TokenVerifier defines the interface for token verification.
|
||||
type TokenVerifier interface {
|
||||
VerifyToken(token string) error
|
||||
}
|
||||
|
||||
// TokenCache defines the interface for a token cache.
|
||||
type TokenCache interface {
|
||||
Get(key string) (map[string]interface{}, bool)
|
||||
}
|
||||
|
||||
// ProviderType is an enumeration for identifying different OIDC providers.
|
||||
type ProviderType int
|
||||
|
||||
const (
|
||||
ProviderTypeGeneric ProviderType = iota
|
||||
ProviderTypeGoogle
|
||||
ProviderTypeAzure
|
||||
)
|
||||
|
||||
// ProviderCapabilities defines the specific features and behaviors of an OIDC provider.
|
||||
type ProviderCapabilities struct {
|
||||
PreferredTokenValidation string
|
||||
SupportsRefreshTokens bool
|
||||
RequiresOfflineAccessScope bool
|
||||
RequiresPromptConsent bool
|
||||
}
|
||||
|
||||
// ValidationResult holds the outcome of a token validation check.
|
||||
type ValidationResult struct {
|
||||
Authenticated bool
|
||||
NeedsRefresh bool
|
||||
IsExpired bool
|
||||
}
|
||||
|
||||
// AuthParams contains the provider-specific parameters for building the authorization URL.
|
||||
type AuthParams struct {
|
||||
URLValues url.Values
|
||||
Scopes []string
|
||||
}
|
||||
|
||||
// TokenResult holds the tokens returned by the provider.
|
||||
type TokenResult struct {
|
||||
IDToken string
|
||||
AccessToken string
|
||||
RefreshToken string
|
||||
}
|
||||
|
||||
// This abstraction allows for provider-specific logic to be encapsulated.
|
||||
type OIDCProvider interface {
|
||||
GetType() ProviderType
|
||||
|
||||
GetCapabilities() ProviderCapabilities
|
||||
|
||||
ValidateTokens(session Session, verifier TokenVerifier, tokenCache TokenCache, refreshGracePeriod time.Duration) (*ValidationResult, error)
|
||||
|
||||
BuildAuthParams(baseParams url.Values, scopes []string) (*AuthParams, error)
|
||||
|
||||
HandleTokenRefresh(tokenData *TokenResult) error
|
||||
|
||||
ValidateConfig() error
|
||||
}
|
||||
|
||||
// This interface decouples the providers from the main session management implementation.
|
||||
type Session interface {
|
||||
GetIDToken() string
|
||||
GetAccessToken() string
|
||||
GetRefreshToken() string
|
||||
GetAuthenticated() bool
|
||||
}
|
||||
@@ -0,0 +1,140 @@
|
||||
package providers
|
||||
|
||||
import (
|
||||
"net/url"
|
||||
"strings"
|
||||
"sync"
|
||||
)
|
||||
|
||||
// ProviderRegistry manages a collection of OIDC provider implementations.
|
||||
// It provides thread-safe access to provider instances and caches detection results.
|
||||
type ProviderRegistry struct {
|
||||
cache map[string]OIDCProvider
|
||||
typeMap map[ProviderType]OIDCProvider
|
||||
providers []OIDCProvider
|
||||
mu sync.RWMutex
|
||||
// Bounded cache configuration to prevent memory leaks
|
||||
maxCacheSize int
|
||||
cacheCount int
|
||||
}
|
||||
|
||||
// NewProviderRegistry creates and initializes a new ProviderRegistry.
|
||||
func NewProviderRegistry() *ProviderRegistry {
|
||||
return &ProviderRegistry{
|
||||
providers: make([]OIDCProvider, 0),
|
||||
cache: make(map[string]OIDCProvider),
|
||||
typeMap: make(map[ProviderType]OIDCProvider),
|
||||
maxCacheSize: 1000, // Prevent unbounded cache growth
|
||||
cacheCount: 0,
|
||||
}
|
||||
}
|
||||
|
||||
// RegisterProvider adds a new provider to the registry.
|
||||
// It maintains both a list of providers and a type-to-provider mapping for efficient lookups.
|
||||
func (r *ProviderRegistry) RegisterProvider(provider OIDCProvider) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.providers = append(r.providers, provider)
|
||||
r.typeMap[provider.GetType()] = provider
|
||||
}
|
||||
|
||||
// GetProviderByType retrieves a provider instance by its type.
|
||||
// Returns nil if the provider type is not registered.
|
||||
func (r *ProviderRegistry) GetProviderByType(providerType ProviderType) OIDCProvider {
|
||||
r.mu.RLock()
|
||||
defer r.mu.RUnlock()
|
||||
return r.typeMap[providerType]
|
||||
}
|
||||
|
||||
// GetRegisteredProviders returns a slice of all registered provider types.
|
||||
func (r *ProviderRegistry) GetRegisteredProviders() []ProviderType {
|
||||
r.mu.RLock()
|
||||
defer r.mu.RUnlock()
|
||||
|
||||
types := make([]ProviderType, 0, len(r.typeMap))
|
||||
for providerType := range r.typeMap {
|
||||
types = append(types, providerType)
|
||||
}
|
||||
return types
|
||||
}
|
||||
|
||||
// ClearCache removes all cached provider detection results.
|
||||
// This can be useful for testing or when provider configuration changes.
|
||||
func (r *ProviderRegistry) ClearCache() {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.cache = make(map[string]OIDCProvider)
|
||||
r.cacheCount = 0
|
||||
}
|
||||
|
||||
// evictOldestCacheEntry removes the first cache entry when cache is full
|
||||
// This is a simple eviction strategy - in production, LRU might be preferred
|
||||
func (r *ProviderRegistry) evictOldestCacheEntry() {
|
||||
// Simple eviction: remove first entry found
|
||||
for key := range r.cache {
|
||||
delete(r.cache, key)
|
||||
r.cacheCount--
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
// DetectProvider identifies the appropriate OIDC provider for an issuer URL.
|
||||
// Uses double-checked locking pattern to avoid race conditions while caching results.
|
||||
func (r *ProviderRegistry) DetectProvider(issuerURL string) OIDCProvider {
|
||||
r.mu.RLock()
|
||||
if provider, found := r.cache[issuerURL]; found {
|
||||
r.mu.RUnlock()
|
||||
return provider
|
||||
}
|
||||
r.mu.RUnlock()
|
||||
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
|
||||
if provider, found := r.cache[issuerURL]; found {
|
||||
return provider
|
||||
}
|
||||
|
||||
detectedProvider := r.detectProviderUnsafe(issuerURL)
|
||||
|
||||
// Check if cache is full and evict if necessary
|
||||
if r.cacheCount >= r.maxCacheSize {
|
||||
r.evictOldestCacheEntry()
|
||||
}
|
||||
|
||||
r.cache[issuerURL] = detectedProvider
|
||||
r.cacheCount++
|
||||
|
||||
return detectedProvider
|
||||
}
|
||||
|
||||
// detectProviderUnsafe performs the actual provider detection logic.
|
||||
// This method assumes the caller holds the appropriate lock and should not be called directly.
|
||||
func (r *ProviderRegistry) detectProviderUnsafe(issuerURL string) OIDCProvider {
|
||||
normalizedURL, err := url.Parse(issuerURL)
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
host := normalizedURL.Host
|
||||
|
||||
for _, p := range r.providers {
|
||||
switch p.GetType() {
|
||||
case ProviderTypeGoogle:
|
||||
if strings.Contains(host, "accounts.google.com") {
|
||||
return p
|
||||
}
|
||||
case ProviderTypeAzure:
|
||||
if strings.Contains(host, "login.microsoftonline.com") || strings.Contains(host, "sts.windows.net") {
|
||||
return p
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for _, p := range r.providers {
|
||||
if p.GetType() == ProviderTypeGeneric {
|
||||
return p
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,151 @@
|
||||
package providers
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net/url"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// ConfigValidator provides common configuration validation utilities for providers.
|
||||
type ConfigValidator struct{}
|
||||
|
||||
// NewConfigValidator creates a new configuration validator.
|
||||
func NewConfigValidator() *ConfigValidator {
|
||||
return &ConfigValidator{}
|
||||
}
|
||||
|
||||
// ValidateIssuerURL validates that an issuer URL is properly formatted and accessible.
|
||||
func (v *ConfigValidator) ValidateIssuerURL(issuerURL string) error {
|
||||
if issuerURL == "" {
|
||||
return fmt.Errorf("issuer URL cannot be empty")
|
||||
}
|
||||
|
||||
parsedURL, err := url.Parse(issuerURL)
|
||||
if err != nil {
|
||||
return fmt.Errorf("invalid issuer URL format: %w", err)
|
||||
}
|
||||
|
||||
if parsedURL.Scheme == "" {
|
||||
return fmt.Errorf("issuer URL must include scheme (http/https)")
|
||||
}
|
||||
|
||||
if parsedURL.Scheme != "http" && parsedURL.Scheme != "https" {
|
||||
return fmt.Errorf("issuer URL scheme must be http or https")
|
||||
}
|
||||
|
||||
if parsedURL.Host == "" {
|
||||
return fmt.Errorf("issuer URL must include host")
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// ValidateClientID validates that a client ID is properly formatted.
|
||||
func (v *ConfigValidator) ValidateClientID(clientID string) error {
|
||||
if clientID == "" {
|
||||
return fmt.Errorf("client ID cannot be empty")
|
||||
}
|
||||
|
||||
if len(clientID) < 3 {
|
||||
return fmt.Errorf("client ID appears to be too short")
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// ValidateScopes validates that the provided scopes are reasonable.
|
||||
func (v *ConfigValidator) ValidateScopes(scopes []string) error {
|
||||
if len(scopes) == 0 {
|
||||
return fmt.Errorf("at least one scope must be provided")
|
||||
}
|
||||
|
||||
hasOpenIDScope := false
|
||||
for _, scope := range scopes {
|
||||
if strings.TrimSpace(scope) == "openid" {
|
||||
hasOpenIDScope = true
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
if !hasOpenIDScope {
|
||||
return fmt.Errorf("'openid' scope is required for OIDC authentication")
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// ValidateRedirectURL validates that a redirect URL is properly formatted.
|
||||
func (v *ConfigValidator) ValidateRedirectURL(redirectURL string) error {
|
||||
if redirectURL == "" {
|
||||
return fmt.Errorf("redirect URL cannot be empty")
|
||||
}
|
||||
|
||||
parsedURL, err := url.Parse(redirectURL)
|
||||
if err != nil {
|
||||
return fmt.Errorf("invalid redirect URL format: %w", err)
|
||||
}
|
||||
|
||||
if parsedURL.Scheme == "" {
|
||||
return fmt.Errorf("redirect URL must include scheme (http/https)")
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// ValidateProviderSpecificConfig performs provider-specific validation.
|
||||
func (v *ConfigValidator) ValidateProviderSpecificConfig(provider OIDCProvider, config map[string]interface{}) error {
|
||||
switch provider.GetType() {
|
||||
case ProviderTypeGoogle:
|
||||
return v.validateGoogleConfig(config)
|
||||
case ProviderTypeAzure:
|
||||
return v.validateAzureConfig(config)
|
||||
case ProviderTypeGeneric:
|
||||
return v.validateGenericConfig(config)
|
||||
default:
|
||||
return fmt.Errorf("unknown provider type: %d", provider.GetType())
|
||||
}
|
||||
}
|
||||
|
||||
// validateGoogleConfig validates Google-specific configuration.
|
||||
func (v *ConfigValidator) validateGoogleConfig(config map[string]interface{}) error {
|
||||
if issuerURL, ok := config["issuer_url"].(string); ok {
|
||||
if !strings.Contains(issuerURL, "accounts.google.com") {
|
||||
return fmt.Errorf("google provider requires issuer URL to contain accounts.google.com")
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// validateAzureConfig validates Azure-specific configuration.
|
||||
func (v *ConfigValidator) validateAzureConfig(config map[string]interface{}) error {
|
||||
if issuerURL, ok := config["issuer_url"].(string); ok {
|
||||
if !strings.Contains(issuerURL, "login.microsoftonline.com") && !strings.Contains(issuerURL, "sts.windows.net") {
|
||||
return fmt.Errorf("azure provider requires issuer URL to contain login.microsoftonline.com or sts.windows.net")
|
||||
}
|
||||
}
|
||||
|
||||
if issuerURL, ok := config["issuer_url"].(string); ok {
|
||||
parsedURL, err := url.Parse(issuerURL)
|
||||
if err == nil {
|
||||
pathParts := strings.Split(parsedURL.Path, "/")
|
||||
hasTenantID := false
|
||||
for _, part := range pathParts {
|
||||
if len(part) == 36 && strings.Count(part, "-") == 4 {
|
||||
hasTenantID = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !hasTenantID {
|
||||
return fmt.Errorf("azure issuer URL should include tenant ID")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// validateGenericConfig validates generic OIDC provider configuration.
|
||||
func (v *ConfigValidator) validateGenericConfig(config map[string]interface{}) error {
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,394 @@
|
||||
// Package singleton provides a centralized, thread-safe singleton management system
|
||||
// that consolidates all singleton patterns used throughout the application.
|
||||
// It ensures proper initialization, lifecycle management, and graceful shutdown.
|
||||
package singleton
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
)
|
||||
|
||||
// Registry is the centralized singleton registry that manages all singleton instances
|
||||
// in the application. It provides thread-safe initialization, access, and cleanup.
|
||||
type Registry struct {
|
||||
mu sync.RWMutex
|
||||
instances map[string]*Instance
|
||||
groups map[string]*Group
|
||||
shutdown int32
|
||||
wg sync.WaitGroup
|
||||
}
|
||||
|
||||
// Instance represents a singleton instance with lifecycle management
|
||||
type Instance struct {
|
||||
name string
|
||||
value interface{}
|
||||
initializer func() interface{}
|
||||
finalizer func(interface{})
|
||||
once sync.Once
|
||||
refCount int32
|
||||
}
|
||||
|
||||
// Group represents a group of related singletons
|
||||
type Group struct {
|
||||
name string
|
||||
instances map[string]*Instance
|
||||
mu sync.RWMutex
|
||||
}
|
||||
|
||||
var (
|
||||
// globalRegistry is the singleton registry instance
|
||||
globalRegistry *Registry
|
||||
// registryOnce ensures single initialization
|
||||
registryOnce sync.Once
|
||||
)
|
||||
|
||||
// Get returns the global singleton registry
|
||||
func Get() *Registry {
|
||||
registryOnce.Do(func() {
|
||||
globalRegistry = &Registry{
|
||||
instances: make(map[string]*Instance),
|
||||
groups: make(map[string]*Group),
|
||||
}
|
||||
})
|
||||
return globalRegistry
|
||||
}
|
||||
|
||||
// Register registers a new singleton with its initializer and optional finalizer
|
||||
func (r *Registry) Register(name string, initializer func() interface{}, finalizer func(interface{})) error {
|
||||
if atomic.LoadInt32(&r.shutdown) == 1 {
|
||||
return fmt.Errorf("registry is shutting down")
|
||||
}
|
||||
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
|
||||
if _, exists := r.instances[name]; exists {
|
||||
return fmt.Errorf("singleton %s already registered", name)
|
||||
}
|
||||
|
||||
r.instances[name] = &Instance{
|
||||
name: name,
|
||||
initializer: initializer,
|
||||
finalizer: finalizer,
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetInstance retrieves or initializes a singleton instance
|
||||
func (r *Registry) GetInstance(name string) (interface{}, error) {
|
||||
if atomic.LoadInt32(&r.shutdown) == 1 {
|
||||
return nil, fmt.Errorf("registry is shutting down")
|
||||
}
|
||||
|
||||
r.mu.RLock()
|
||||
instance, exists := r.instances[name]
|
||||
r.mu.RUnlock()
|
||||
|
||||
if !exists {
|
||||
return nil, fmt.Errorf("singleton %s not registered", name)
|
||||
}
|
||||
|
||||
// Initialize the singleton if needed
|
||||
instance.once.Do(func() {
|
||||
if instance.initializer != nil {
|
||||
instance.value = instance.initializer()
|
||||
atomic.AddInt32(&instance.refCount, 1)
|
||||
}
|
||||
})
|
||||
|
||||
return instance.value, nil
|
||||
}
|
||||
|
||||
// MustGet retrieves a singleton instance, panicking if not found
|
||||
func (r *Registry) MustGet(name string) interface{} {
|
||||
val, err := r.GetInstance(name)
|
||||
if err != nil {
|
||||
panic(fmt.Sprintf("singleton %s: %v", name, err))
|
||||
}
|
||||
return val
|
||||
}
|
||||
|
||||
// RegisterGroup creates a new singleton group
|
||||
func (r *Registry) RegisterGroup(name string) error {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
|
||||
if _, exists := r.groups[name]; exists {
|
||||
return fmt.Errorf("group %s already exists", name)
|
||||
}
|
||||
|
||||
r.groups[name] = &Group{
|
||||
name: name,
|
||||
instances: make(map[string]*Instance),
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// AddToGroup adds a singleton to a group
|
||||
func (r *Registry) AddToGroup(groupName, singletonName string) error {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
|
||||
group, groupExists := r.groups[groupName]
|
||||
if !groupExists {
|
||||
return fmt.Errorf("group %s does not exist", groupName)
|
||||
}
|
||||
|
||||
instance, instanceExists := r.instances[singletonName]
|
||||
if !instanceExists {
|
||||
return fmt.Errorf("singleton %s not registered", singletonName)
|
||||
}
|
||||
|
||||
group.mu.Lock()
|
||||
defer group.mu.Unlock()
|
||||
|
||||
group.instances[singletonName] = instance
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetGroup retrieves all singletons in a group
|
||||
func (r *Registry) GetGroup(name string) (map[string]interface{}, error) {
|
||||
r.mu.RLock()
|
||||
group, exists := r.groups[name]
|
||||
r.mu.RUnlock()
|
||||
|
||||
if !exists {
|
||||
return nil, fmt.Errorf("group %s does not exist", name)
|
||||
}
|
||||
|
||||
group.mu.RLock()
|
||||
defer group.mu.RUnlock()
|
||||
|
||||
result := make(map[string]interface{})
|
||||
for name, instance := range group.instances {
|
||||
if instance.value != nil {
|
||||
result[name] = instance.value
|
||||
}
|
||||
}
|
||||
|
||||
return result, nil
|
||||
}
|
||||
|
||||
// AddReference increments the reference count for a singleton
|
||||
func (r *Registry) AddReference(name string) error {
|
||||
r.mu.RLock()
|
||||
instance, exists := r.instances[name]
|
||||
r.mu.RUnlock()
|
||||
|
||||
if !exists {
|
||||
return fmt.Errorf("singleton %s not registered", name)
|
||||
}
|
||||
|
||||
atomic.AddInt32(&instance.refCount, 1)
|
||||
return nil
|
||||
}
|
||||
|
||||
// ReleaseReference decrements the reference count for a singleton
|
||||
func (r *Registry) ReleaseReference(name string) error {
|
||||
r.mu.RLock()
|
||||
instance, exists := r.instances[name]
|
||||
r.mu.RUnlock()
|
||||
|
||||
if !exists {
|
||||
return fmt.Errorf("singleton %s not registered", name)
|
||||
}
|
||||
|
||||
count := atomic.AddInt32(&instance.refCount, -1)
|
||||
if count == 0 && instance.finalizer != nil && instance.value != nil {
|
||||
// Run finalizer when last reference is released
|
||||
go instance.finalizer(instance.value)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetReferenceCount returns the reference count for a singleton
|
||||
func (r *Registry) GetReferenceCount(name string) (int32, error) {
|
||||
r.mu.RLock()
|
||||
instance, exists := r.instances[name]
|
||||
r.mu.RUnlock()
|
||||
|
||||
if !exists {
|
||||
return 0, fmt.Errorf("singleton %s not registered", name)
|
||||
}
|
||||
|
||||
return atomic.LoadInt32(&instance.refCount), nil
|
||||
}
|
||||
|
||||
// Shutdown gracefully shuts down all singletons
|
||||
func (r *Registry) Shutdown(ctx context.Context) error {
|
||||
if !atomic.CompareAndSwapInt32(&r.shutdown, 0, 1) {
|
||||
return fmt.Errorf("registry already shutting down")
|
||||
}
|
||||
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
|
||||
// Create error channel for collecting shutdown errors
|
||||
errChan := make(chan error, len(r.instances))
|
||||
|
||||
// Run finalizers for all initialized singletons
|
||||
for name, instance := range r.instances {
|
||||
if instance.value != nil && instance.finalizer != nil {
|
||||
r.wg.Add(1)
|
||||
go func(n string, i *Instance) {
|
||||
defer r.wg.Done()
|
||||
|
||||
// Run finalizer with panic recovery
|
||||
func() {
|
||||
defer func() {
|
||||
if r := recover(); r != nil {
|
||||
errChan <- fmt.Errorf("finalizer for %s panicked: %v", n, r)
|
||||
}
|
||||
}()
|
||||
i.finalizer(i.value)
|
||||
}()
|
||||
}(name, instance)
|
||||
}
|
||||
}
|
||||
|
||||
// Wait for all finalizers to complete or timeout
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
r.wg.Wait()
|
||||
close(done)
|
||||
}()
|
||||
|
||||
select {
|
||||
case <-done:
|
||||
// All finalizers completed
|
||||
case <-ctx.Done():
|
||||
return fmt.Errorf("shutdown timeout: %w", ctx.Err())
|
||||
}
|
||||
|
||||
// Collect any errors
|
||||
close(errChan)
|
||||
var errs []error
|
||||
for err := range errChan {
|
||||
if err != nil {
|
||||
errs = append(errs, err)
|
||||
}
|
||||
}
|
||||
|
||||
// Clear all instances
|
||||
r.instances = make(map[string]*Instance)
|
||||
r.groups = make(map[string]*Group)
|
||||
|
||||
if len(errs) > 0 {
|
||||
return fmt.Errorf("shutdown errors: %v", errs)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// Reset resets the registry (mainly for testing)
|
||||
func (r *Registry) Reset() {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
|
||||
r.instances = make(map[string]*Instance)
|
||||
r.groups = make(map[string]*Group)
|
||||
atomic.StoreInt32(&r.shutdown, 0)
|
||||
}
|
||||
|
||||
// Stats returns statistics about the registry
|
||||
type Stats struct {
|
||||
TotalRegistered int
|
||||
TotalInitialized int
|
||||
TotalGroups int
|
||||
TotalReferences int32
|
||||
}
|
||||
|
||||
// GetStats returns current registry statistics
|
||||
func (r *Registry) GetStats() Stats {
|
||||
r.mu.RLock()
|
||||
defer r.mu.RUnlock()
|
||||
|
||||
stats := Stats{
|
||||
TotalRegistered: len(r.instances),
|
||||
TotalGroups: len(r.groups),
|
||||
}
|
||||
|
||||
for _, instance := range r.instances {
|
||||
if instance.value != nil {
|
||||
stats.TotalInitialized++
|
||||
}
|
||||
stats.TotalReferences += atomic.LoadInt32(&instance.refCount)
|
||||
}
|
||||
|
||||
return stats
|
||||
}
|
||||
|
||||
// Builder provides a fluent interface for registering singletons
|
||||
type Builder struct {
|
||||
registry *Registry
|
||||
name string
|
||||
initializer func() interface{}
|
||||
finalizer func(interface{})
|
||||
group string
|
||||
}
|
||||
|
||||
// NewBuilder creates a new singleton builder
|
||||
func NewBuilder(name string) *Builder {
|
||||
return &Builder{
|
||||
registry: Get(),
|
||||
name: name,
|
||||
}
|
||||
}
|
||||
|
||||
// WithInitializer sets the initializer function
|
||||
func (b *Builder) WithInitializer(init func() interface{}) *Builder {
|
||||
b.initializer = init
|
||||
return b
|
||||
}
|
||||
|
||||
// WithFinalizer sets the finalizer function
|
||||
func (b *Builder) WithFinalizer(final func(interface{})) *Builder {
|
||||
b.finalizer = final
|
||||
return b
|
||||
}
|
||||
|
||||
// InGroup adds the singleton to a group
|
||||
func (b *Builder) InGroup(group string) *Builder {
|
||||
b.group = group
|
||||
return b
|
||||
}
|
||||
|
||||
// Register registers the singleton with the configured options
|
||||
func (b *Builder) Register() error {
|
||||
if err := b.registry.Register(b.name, b.initializer, b.finalizer); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if b.group != "" {
|
||||
// Ensure group exists
|
||||
if err := b.registry.RegisterGroup(b.group); err != nil {
|
||||
// Group might already exist, which is ok
|
||||
if !contains(err.Error(), "already exists") {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
return b.registry.AddToGroup(b.group, b.name)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// Helper function to check if string contains substring
|
||||
func contains(s, substr string) bool {
|
||||
return len(s) >= len(substr) && (s == substr || len(s) > len(substr) && containsHelper(s, substr))
|
||||
}
|
||||
|
||||
func containsHelper(s, substr string) bool {
|
||||
for i := 0; i <= len(s)-len(substr); i++ {
|
||||
if s[i:i+len(substr)] == substr {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
@@ -0,0 +1,970 @@
|
||||
package singleton
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// TestGet_Singleton tests that Get() returns the same instance
|
||||
func TestGet_Singleton(t *testing.T) {
|
||||
registry1 := Get()
|
||||
registry2 := Get()
|
||||
|
||||
if registry1 != registry2 {
|
||||
t.Error("Get() should return the same instance (singleton)")
|
||||
}
|
||||
|
||||
if registry1 == nil {
|
||||
t.Error("Get() should not return nil")
|
||||
}
|
||||
}
|
||||
|
||||
// TestRegistry_Register tests singleton registration
|
||||
func TestRegistry_Register(t *testing.T) {
|
||||
registry := &Registry{
|
||||
instances: make(map[string]*Instance),
|
||||
groups: make(map[string]*Group),
|
||||
}
|
||||
|
||||
initializer := func() interface{} {
|
||||
return "test-value"
|
||||
}
|
||||
|
||||
finalizer := func(v interface{}) {
|
||||
// Mock finalizer
|
||||
}
|
||||
|
||||
// Test successful registration
|
||||
err := registry.Register("test-singleton", initializer, finalizer)
|
||||
if err != nil {
|
||||
t.Errorf("Register should succeed, got error: %v", err)
|
||||
}
|
||||
|
||||
// Verify instance was registered
|
||||
if len(registry.instances) != 1 {
|
||||
t.Error("Instance should be registered")
|
||||
}
|
||||
|
||||
instance := registry.instances["test-singleton"]
|
||||
if instance == nil {
|
||||
t.Error("Instance should not be nil")
|
||||
return
|
||||
}
|
||||
|
||||
if instance.name != "test-singleton" {
|
||||
t.Errorf("Instance name should be 'test-singleton', got '%s'", instance.name)
|
||||
}
|
||||
|
||||
if instance.initializer == nil {
|
||||
t.Error("Instance should have initializer")
|
||||
}
|
||||
|
||||
if instance.finalizer == nil {
|
||||
t.Error("Instance should have finalizer")
|
||||
}
|
||||
}
|
||||
|
||||
// TestRegistry_Register_Duplicate tests duplicate registration
|
||||
func TestRegistry_Register_Duplicate(t *testing.T) {
|
||||
registry := &Registry{
|
||||
instances: make(map[string]*Instance),
|
||||
groups: make(map[string]*Group),
|
||||
}
|
||||
|
||||
initializer := func() interface{} {
|
||||
return "test-value"
|
||||
}
|
||||
|
||||
// Register first time
|
||||
err := registry.Register("test-singleton", initializer, nil)
|
||||
if err != nil {
|
||||
t.Errorf("First registration should succeed, got error: %v", err)
|
||||
}
|
||||
|
||||
// Register again - should fail
|
||||
err = registry.Register("test-singleton", initializer, nil)
|
||||
if err == nil {
|
||||
t.Error("Duplicate registration should fail")
|
||||
}
|
||||
|
||||
if !strings.Contains(err.Error(), "already registered") {
|
||||
t.Errorf("Error should mention already registered, got: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestRegistry_Register_DuringShutdown tests registration during shutdown
|
||||
func TestRegistry_Register_DuringShutdown(t *testing.T) {
|
||||
registry := &Registry{
|
||||
instances: make(map[string]*Instance),
|
||||
groups: make(map[string]*Group),
|
||||
shutdown: 1, // Already shutting down
|
||||
}
|
||||
|
||||
initializer := func() interface{} {
|
||||
return "test-value"
|
||||
}
|
||||
|
||||
err := registry.Register("test-singleton", initializer, nil)
|
||||
if err == nil {
|
||||
t.Error("Registration during shutdown should fail")
|
||||
}
|
||||
|
||||
if !strings.Contains(err.Error(), "shutting down") {
|
||||
t.Errorf("Error should mention shutting down, got: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestRegistry_GetInstance tests singleton retrieval and initialization
|
||||
func TestRegistry_GetInstance(t *testing.T) {
|
||||
registry := &Registry{
|
||||
instances: make(map[string]*Instance),
|
||||
groups: make(map[string]*Group),
|
||||
}
|
||||
|
||||
callCount := int32(0)
|
||||
testValue := "test-value"
|
||||
|
||||
initializer := func() interface{} {
|
||||
atomic.AddInt32(&callCount, 1)
|
||||
return testValue
|
||||
}
|
||||
|
||||
// Register singleton
|
||||
err := registry.Register("test-singleton", initializer, nil)
|
||||
if err != nil {
|
||||
t.Errorf("Register should succeed, got error: %v", err)
|
||||
}
|
||||
|
||||
// First get - should initialize
|
||||
value1, err := registry.GetInstance("test-singleton")
|
||||
if err != nil {
|
||||
t.Errorf("GetInstance should succeed, got error: %v", err)
|
||||
}
|
||||
|
||||
if value1 != testValue {
|
||||
t.Errorf("Value should be '%s', got '%v'", testValue, value1)
|
||||
}
|
||||
|
||||
if atomic.LoadInt32(&callCount) != 1 {
|
||||
t.Errorf("Initializer should be called once, called %d times", callCount)
|
||||
}
|
||||
|
||||
// Second get - should return same instance without calling initializer
|
||||
value2, err := registry.GetInstance("test-singleton")
|
||||
if err != nil {
|
||||
t.Errorf("GetInstance should succeed, got error: %v", err)
|
||||
}
|
||||
|
||||
if value2 != testValue {
|
||||
t.Errorf("Value should be '%s', got '%v'", testValue, value2)
|
||||
}
|
||||
|
||||
if atomic.LoadInt32(&callCount) != 1 {
|
||||
t.Errorf("Initializer should still be called only once, called %d times", callCount)
|
||||
}
|
||||
}
|
||||
|
||||
// TestRegistry_GetInstance_NotRegistered tests getting unregistered singleton
|
||||
func TestRegistry_GetInstance_NotRegistered(t *testing.T) {
|
||||
registry := &Registry{
|
||||
instances: make(map[string]*Instance),
|
||||
groups: make(map[string]*Group),
|
||||
}
|
||||
|
||||
value, err := registry.GetInstance("non-existent")
|
||||
if err == nil {
|
||||
t.Error("GetInstance of non-existent singleton should fail")
|
||||
}
|
||||
|
||||
if value != nil {
|
||||
t.Error("Value should be nil for non-existent singleton")
|
||||
}
|
||||
|
||||
if !strings.Contains(err.Error(), "not registered") {
|
||||
t.Errorf("Error should mention not registered, got: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestRegistry_GetInstance_DuringShutdown tests getting instance during shutdown
|
||||
func TestRegistry_GetInstance_DuringShutdown(t *testing.T) {
|
||||
registry := &Registry{
|
||||
instances: make(map[string]*Instance),
|
||||
groups: make(map[string]*Group),
|
||||
shutdown: 1, // Already shutting down
|
||||
}
|
||||
|
||||
value, err := registry.GetInstance("test-singleton")
|
||||
if err == nil {
|
||||
t.Error("GetInstance during shutdown should fail")
|
||||
}
|
||||
|
||||
if value != nil {
|
||||
t.Error("Value should be nil during shutdown")
|
||||
}
|
||||
|
||||
if !strings.Contains(err.Error(), "shutting down") {
|
||||
t.Errorf("Error should mention shutting down, got: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestRegistry_MustGet tests MustGet method
|
||||
func TestRegistry_MustGet(t *testing.T) {
|
||||
registry := &Registry{
|
||||
instances: make(map[string]*Instance),
|
||||
groups: make(map[string]*Group),
|
||||
}
|
||||
|
||||
testValue := "test-value"
|
||||
initializer := func() interface{} {
|
||||
return testValue
|
||||
}
|
||||
|
||||
// Register singleton
|
||||
err := registry.Register("test-singleton", initializer, nil)
|
||||
if err != nil {
|
||||
t.Errorf("Register should succeed, got error: %v", err)
|
||||
}
|
||||
|
||||
// MustGet should succeed
|
||||
value := registry.MustGet("test-singleton")
|
||||
if value != testValue {
|
||||
t.Errorf("Value should be '%s', got '%v'", testValue, value)
|
||||
}
|
||||
|
||||
// MustGet non-existent should panic
|
||||
defer func() {
|
||||
if r := recover(); r == nil {
|
||||
t.Error("MustGet of non-existent singleton should panic")
|
||||
}
|
||||
}()
|
||||
|
||||
registry.MustGet("non-existent")
|
||||
}
|
||||
|
||||
// TestRegistry_RegisterGroup tests group registration
|
||||
func TestRegistry_RegisterGroup(t *testing.T) {
|
||||
registry := &Registry{
|
||||
instances: make(map[string]*Instance),
|
||||
groups: make(map[string]*Group),
|
||||
}
|
||||
|
||||
// Test successful group registration
|
||||
err := registry.RegisterGroup("test-group")
|
||||
if err != nil {
|
||||
t.Errorf("RegisterGroup should succeed, got error: %v", err)
|
||||
}
|
||||
|
||||
// Verify group was registered
|
||||
if len(registry.groups) != 1 {
|
||||
t.Error("Group should be registered")
|
||||
}
|
||||
|
||||
group := registry.groups["test-group"]
|
||||
if group == nil {
|
||||
t.Error("Group should not be nil")
|
||||
return
|
||||
}
|
||||
|
||||
if group.name != "test-group" {
|
||||
t.Errorf("Group name should be 'test-group', got '%s'", group.name)
|
||||
}
|
||||
|
||||
// Test duplicate group registration
|
||||
err = registry.RegisterGroup("test-group")
|
||||
if err == nil {
|
||||
t.Error("Duplicate group registration should fail")
|
||||
}
|
||||
|
||||
if !strings.Contains(err.Error(), "already exists") {
|
||||
t.Errorf("Error should mention already exists, got: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestRegistry_AddToGroup tests adding singletons to groups
|
||||
func TestRegistry_AddToGroup(t *testing.T) {
|
||||
registry := &Registry{
|
||||
instances: make(map[string]*Instance),
|
||||
groups: make(map[string]*Group),
|
||||
}
|
||||
|
||||
// Register a singleton
|
||||
initializer := func() interface{} {
|
||||
return "test-value"
|
||||
}
|
||||
|
||||
err := registry.Register("test-singleton", initializer, nil)
|
||||
if err != nil {
|
||||
t.Errorf("Register should succeed, got error: %v", err)
|
||||
}
|
||||
|
||||
// Register a group
|
||||
err = registry.RegisterGroup("test-group")
|
||||
if err != nil {
|
||||
t.Errorf("RegisterGroup should succeed, got error: %v", err)
|
||||
}
|
||||
|
||||
// Add singleton to group
|
||||
err = registry.AddToGroup("test-group", "test-singleton")
|
||||
if err != nil {
|
||||
t.Errorf("AddToGroup should succeed, got error: %v", err)
|
||||
}
|
||||
|
||||
// Verify singleton is in group
|
||||
group := registry.groups["test-group"]
|
||||
if len(group.instances) != 1 {
|
||||
t.Error("Group should contain one instance")
|
||||
}
|
||||
|
||||
if group.instances["test-singleton"] == nil {
|
||||
t.Error("Singleton should be in group")
|
||||
}
|
||||
|
||||
// Test adding to non-existent group
|
||||
err = registry.AddToGroup("non-existent-group", "test-singleton")
|
||||
if err == nil {
|
||||
t.Error("Adding to non-existent group should fail")
|
||||
}
|
||||
|
||||
// Test adding non-existent singleton to group
|
||||
err = registry.AddToGroup("test-group", "non-existent-singleton")
|
||||
if err == nil {
|
||||
t.Error("Adding non-existent singleton should fail")
|
||||
}
|
||||
}
|
||||
|
||||
// TestRegistry_GetGroup tests retrieving group instances
|
||||
func TestRegistry_GetGroup(t *testing.T) {
|
||||
registry := &Registry{
|
||||
instances: make(map[string]*Instance),
|
||||
groups: make(map[string]*Group),
|
||||
}
|
||||
|
||||
// Register singletons
|
||||
err := registry.Register("test-singleton-1", func() interface{} {
|
||||
return "value-1"
|
||||
}, nil)
|
||||
if err != nil {
|
||||
t.Errorf("Register should succeed, got error: %v", err)
|
||||
}
|
||||
|
||||
err = registry.Register("test-singleton-2", func() interface{} {
|
||||
return "value-2"
|
||||
}, nil)
|
||||
if err != nil {
|
||||
t.Errorf("Register should succeed, got error: %v", err)
|
||||
}
|
||||
|
||||
// Register group and add singletons
|
||||
err = registry.RegisterGroup("test-group")
|
||||
if err != nil {
|
||||
t.Errorf("RegisterGroup should succeed, got error: %v", err)
|
||||
}
|
||||
|
||||
err = registry.AddToGroup("test-group", "test-singleton-1")
|
||||
if err != nil {
|
||||
t.Errorf("AddToGroup should succeed, got error: %v", err)
|
||||
}
|
||||
|
||||
err = registry.AddToGroup("test-group", "test-singleton-2")
|
||||
if err != nil {
|
||||
t.Errorf("AddToGroup should succeed, got error: %v", err)
|
||||
}
|
||||
|
||||
// Initialize singletons
|
||||
_, _ = registry.GetInstance("test-singleton-1")
|
||||
_, _ = registry.GetInstance("test-singleton-2")
|
||||
|
||||
// Get group
|
||||
groupInstances, err := registry.GetGroup("test-group")
|
||||
if err != nil {
|
||||
t.Errorf("GetGroup should succeed, got error: %v", err)
|
||||
}
|
||||
|
||||
if len(groupInstances) != 2 {
|
||||
t.Errorf("Group should contain 2 instances, got %d", len(groupInstances))
|
||||
}
|
||||
|
||||
if groupInstances["test-singleton-1"] != "value-1" {
|
||||
t.Error("Group should contain correct instance values")
|
||||
}
|
||||
|
||||
if groupInstances["test-singleton-2"] != "value-2" {
|
||||
t.Error("Group should contain correct instance values")
|
||||
}
|
||||
|
||||
// Test getting non-existent group
|
||||
_, err = registry.GetGroup("non-existent-group")
|
||||
if err == nil {
|
||||
t.Error("Getting non-existent group should fail")
|
||||
}
|
||||
}
|
||||
|
||||
// TestRegistry_ReferenceCountingv tests reference counting
|
||||
func TestRegistry_ReferenceCountingv(t *testing.T) {
|
||||
registry := &Registry{
|
||||
instances: make(map[string]*Instance),
|
||||
groups: make(map[string]*Group),
|
||||
}
|
||||
|
||||
finalizerCalled := int32(0)
|
||||
finalizer := func(v interface{}) {
|
||||
atomic.AddInt32(&finalizerCalled, 1)
|
||||
}
|
||||
|
||||
// Register singleton
|
||||
err := registry.Register("test-singleton", func() interface{} {
|
||||
return "test-value"
|
||||
}, finalizer)
|
||||
if err != nil {
|
||||
t.Errorf("Register should succeed, got error: %v", err)
|
||||
}
|
||||
|
||||
// Initialize singleton (this adds 1 reference)
|
||||
_, err = registry.GetInstance("test-singleton")
|
||||
if err != nil {
|
||||
t.Errorf("GetInstance should succeed, got error: %v", err)
|
||||
}
|
||||
|
||||
// Check initial reference count
|
||||
count, err := registry.GetReferenceCount("test-singleton")
|
||||
if err != nil {
|
||||
t.Errorf("GetReferenceCount should succeed, got error: %v", err)
|
||||
}
|
||||
|
||||
if count != 1 {
|
||||
t.Errorf("Reference count should be 1, got %d", count)
|
||||
}
|
||||
|
||||
// Add reference
|
||||
err = registry.AddReference("test-singleton")
|
||||
if err != nil {
|
||||
t.Errorf("AddReference should succeed, got error: %v", err)
|
||||
}
|
||||
|
||||
count, _ = registry.GetReferenceCount("test-singleton")
|
||||
if count != 2 {
|
||||
t.Errorf("Reference count should be 2, got %d", count)
|
||||
}
|
||||
|
||||
// Release reference
|
||||
err = registry.ReleaseReference("test-singleton")
|
||||
if err != nil {
|
||||
t.Errorf("ReleaseReference should succeed, got error: %v", err)
|
||||
}
|
||||
|
||||
count, _ = registry.GetReferenceCount("test-singleton")
|
||||
if count != 1 {
|
||||
t.Errorf("Reference count should be 1, got %d", count)
|
||||
}
|
||||
|
||||
// Release last reference - should trigger finalizer
|
||||
err = registry.ReleaseReference("test-singleton")
|
||||
if err != nil {
|
||||
t.Errorf("ReleaseReference should succeed, got error: %v", err)
|
||||
}
|
||||
|
||||
count, _ = registry.GetReferenceCount("test-singleton")
|
||||
if count != 0 {
|
||||
t.Errorf("Reference count should be 0, got %d", count)
|
||||
}
|
||||
|
||||
// Wait for finalizer to run (it runs in goroutine)
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
|
||||
if atomic.LoadInt32(&finalizerCalled) != 1 {
|
||||
t.Errorf("Finalizer should be called once, called %d times", finalizerCalled)
|
||||
}
|
||||
|
||||
// Test reference operations on non-existent singleton
|
||||
err = registry.AddReference("non-existent")
|
||||
if err == nil {
|
||||
t.Error("AddReference on non-existent singleton should fail")
|
||||
}
|
||||
|
||||
err = registry.ReleaseReference("non-existent")
|
||||
if err == nil {
|
||||
t.Error("ReleaseReference on non-existent singleton should fail")
|
||||
}
|
||||
|
||||
_, err = registry.GetReferenceCount("non-existent")
|
||||
if err == nil {
|
||||
t.Error("GetReferenceCount on non-existent singleton should fail")
|
||||
}
|
||||
}
|
||||
|
||||
// TestRegistry_Shutdown tests graceful shutdown
|
||||
func TestRegistry_Shutdown(t *testing.T) {
|
||||
registry := &Registry{
|
||||
instances: make(map[string]*Instance),
|
||||
groups: make(map[string]*Group),
|
||||
}
|
||||
|
||||
finalizerCalled := int32(0)
|
||||
finalizer := func(v interface{}) {
|
||||
atomic.AddInt32(&finalizerCalled, 1)
|
||||
}
|
||||
|
||||
// Register and initialize singletons
|
||||
err := registry.Register("test-singleton-1", func() interface{} {
|
||||
return "value-1"
|
||||
}, finalizer)
|
||||
if err != nil {
|
||||
t.Errorf("Register should succeed, got error: %v", err)
|
||||
}
|
||||
|
||||
err = registry.Register("test-singleton-2", func() interface{} {
|
||||
return "value-2"
|
||||
}, finalizer)
|
||||
if err != nil {
|
||||
t.Errorf("Register should succeed, got error: %v", err)
|
||||
}
|
||||
|
||||
// Initialize singletons
|
||||
_, _ = registry.GetInstance("test-singleton-1")
|
||||
_, _ = registry.GetInstance("test-singleton-2")
|
||||
|
||||
// Shutdown
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
|
||||
err = registry.Shutdown(ctx)
|
||||
if err != nil {
|
||||
t.Errorf("Shutdown should succeed, got error: %v", err)
|
||||
}
|
||||
|
||||
// Verify finalizers were called
|
||||
if atomic.LoadInt32(&finalizerCalled) != 2 {
|
||||
t.Errorf("Finalizers should be called 2 times, called %d times", finalizerCalled)
|
||||
}
|
||||
|
||||
// Verify registry is cleared
|
||||
if len(registry.instances) != 0 {
|
||||
t.Error("Instances should be cleared after shutdown")
|
||||
}
|
||||
|
||||
if len(registry.groups) != 0 {
|
||||
t.Error("Groups should be cleared after shutdown")
|
||||
}
|
||||
|
||||
// Verify shutdown flag is set
|
||||
if atomic.LoadInt32(®istry.shutdown) != 1 {
|
||||
t.Error("Shutdown flag should be set")
|
||||
}
|
||||
|
||||
// Test double shutdown
|
||||
err = registry.Shutdown(ctx)
|
||||
if err == nil {
|
||||
t.Error("Double shutdown should fail")
|
||||
}
|
||||
}
|
||||
|
||||
// TestRegistry_Shutdown_Timeout tests shutdown timeout
|
||||
func TestRegistry_Shutdown_Timeout(t *testing.T) {
|
||||
registry := &Registry{
|
||||
instances: make(map[string]*Instance),
|
||||
groups: make(map[string]*Group),
|
||||
}
|
||||
|
||||
// Register singleton with slow finalizer
|
||||
slowFinalizer := func(v interface{}) {
|
||||
time.Sleep(100 * time.Millisecond)
|
||||
}
|
||||
|
||||
err := registry.Register("slow-singleton", func() interface{} {
|
||||
return "value"
|
||||
}, slowFinalizer)
|
||||
if err != nil {
|
||||
t.Errorf("Register should succeed, got error: %v", err)
|
||||
}
|
||||
|
||||
// Initialize singleton
|
||||
_, _ = registry.GetInstance("slow-singleton")
|
||||
|
||||
// Shutdown with short timeout
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond)
|
||||
defer cancel()
|
||||
|
||||
err = registry.Shutdown(ctx)
|
||||
if err == nil {
|
||||
t.Error("Shutdown should timeout")
|
||||
}
|
||||
|
||||
if !strings.Contains(err.Error(), "timeout") {
|
||||
t.Errorf("Error should mention timeout, got: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestRegistry_Shutdown_PanicRecovery tests panic recovery during shutdown
|
||||
func TestRegistry_Shutdown_PanicRecovery(t *testing.T) {
|
||||
registry := &Registry{
|
||||
instances: make(map[string]*Instance),
|
||||
groups: make(map[string]*Group),
|
||||
}
|
||||
|
||||
// Register singleton with panicking finalizer
|
||||
panicFinalizer := func(v interface{}) {
|
||||
panic("finalizer panic")
|
||||
}
|
||||
|
||||
err := registry.Register("panic-singleton", func() interface{} {
|
||||
return "value"
|
||||
}, panicFinalizer)
|
||||
if err != nil {
|
||||
t.Errorf("Register should succeed, got error: %v", err)
|
||||
}
|
||||
|
||||
// Initialize singleton
|
||||
_, _ = registry.GetInstance("panic-singleton")
|
||||
|
||||
// Shutdown should handle panic
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
|
||||
err = registry.Shutdown(ctx)
|
||||
if err == nil {
|
||||
t.Error("Shutdown should report finalizer panic")
|
||||
}
|
||||
|
||||
if !strings.Contains(err.Error(), "panicked") {
|
||||
t.Errorf("Error should mention panic, got: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestRegistry_Reset tests registry reset
|
||||
func TestRegistry_Reset(t *testing.T) {
|
||||
registry := &Registry{
|
||||
instances: make(map[string]*Instance),
|
||||
groups: make(map[string]*Group),
|
||||
shutdown: 1,
|
||||
}
|
||||
|
||||
// Add some data
|
||||
registry.instances["test"] = &Instance{}
|
||||
registry.groups["test"] = &Group{}
|
||||
|
||||
// Reset
|
||||
registry.Reset()
|
||||
|
||||
// Verify everything is cleared
|
||||
if len(registry.instances) != 0 {
|
||||
t.Error("Instances should be cleared after reset")
|
||||
}
|
||||
|
||||
if len(registry.groups) != 0 {
|
||||
t.Error("Groups should be cleared after reset")
|
||||
}
|
||||
|
||||
if atomic.LoadInt32(®istry.shutdown) != 0 {
|
||||
t.Error("Shutdown flag should be cleared after reset")
|
||||
}
|
||||
}
|
||||
|
||||
// TestRegistry_GetStats tests statistics
|
||||
func TestRegistry_GetStats(t *testing.T) {
|
||||
registry := &Registry{
|
||||
instances: make(map[string]*Instance),
|
||||
groups: make(map[string]*Group),
|
||||
}
|
||||
|
||||
// Register singletons
|
||||
err := registry.Register("test-singleton-1", func() interface{} {
|
||||
return "value-1"
|
||||
}, nil)
|
||||
if err != nil {
|
||||
t.Errorf("Register should succeed, got error: %v", err)
|
||||
}
|
||||
|
||||
err = registry.Register("test-singleton-2", func() interface{} {
|
||||
return "value-2"
|
||||
}, nil)
|
||||
if err != nil {
|
||||
t.Errorf("Register should succeed, got error: %v", err)
|
||||
}
|
||||
|
||||
// Register group
|
||||
err = registry.RegisterGroup("test-group")
|
||||
if err != nil {
|
||||
t.Errorf("RegisterGroup should succeed, got error: %v", err)
|
||||
}
|
||||
|
||||
// Initialize one singleton
|
||||
_, _ = registry.GetInstance("test-singleton-1")
|
||||
|
||||
// Add reference
|
||||
_ = registry.AddReference("test-singleton-1")
|
||||
|
||||
// Get stats
|
||||
stats := registry.GetStats()
|
||||
|
||||
if stats.TotalRegistered != 2 {
|
||||
t.Errorf("TotalRegistered should be 2, got %d", stats.TotalRegistered)
|
||||
}
|
||||
|
||||
if stats.TotalInitialized != 1 {
|
||||
t.Errorf("TotalInitialized should be 1, got %d", stats.TotalInitialized)
|
||||
}
|
||||
|
||||
if stats.TotalGroups != 1 {
|
||||
t.Errorf("TotalGroups should be 1, got %d", stats.TotalGroups)
|
||||
}
|
||||
|
||||
if stats.TotalReferences != 2 { // 1 from initialization + 1 from AddReference
|
||||
t.Errorf("TotalReferences should be 2, got %d", stats.TotalReferences)
|
||||
}
|
||||
}
|
||||
|
||||
// TestBuilder tests the fluent builder interface
|
||||
func TestBuilder(t *testing.T) {
|
||||
// Reset global registry for clean test
|
||||
Get().Reset()
|
||||
|
||||
testValue := "builder-test-value"
|
||||
|
||||
initializer := func() interface{} {
|
||||
return testValue
|
||||
}
|
||||
|
||||
finalizer := func(v interface{}) {
|
||||
// Mock finalizer for builder test
|
||||
}
|
||||
|
||||
// Test builder
|
||||
err := NewBuilder("builder-singleton").
|
||||
WithInitializer(initializer).
|
||||
WithFinalizer(finalizer).
|
||||
InGroup("builder-group").
|
||||
Register()
|
||||
|
||||
if err != nil {
|
||||
t.Errorf("Builder registration should succeed, got error: %v", err)
|
||||
}
|
||||
|
||||
// Verify singleton was registered
|
||||
value, err := Get().GetInstance("builder-singleton")
|
||||
if err != nil {
|
||||
t.Errorf("GetInstance should succeed, got error: %v", err)
|
||||
}
|
||||
|
||||
if value != testValue {
|
||||
t.Errorf("Value should be '%s', got '%v'", testValue, value)
|
||||
}
|
||||
|
||||
// Verify group was created and singleton added
|
||||
groupInstances, err := Get().GetGroup("builder-group")
|
||||
if err != nil {
|
||||
t.Errorf("GetGroup should succeed, got error: %v", err)
|
||||
}
|
||||
|
||||
if len(groupInstances) != 1 {
|
||||
t.Errorf("Group should contain 1 instance, got %d", len(groupInstances))
|
||||
}
|
||||
|
||||
if groupInstances["builder-singleton"] != testValue {
|
||||
t.Error("Group should contain correct instance")
|
||||
}
|
||||
}
|
||||
|
||||
// TestBuilder_WithoutGroup tests builder without group
|
||||
func TestBuilder_WithoutGroup(t *testing.T) {
|
||||
registry := &Registry{
|
||||
instances: make(map[string]*Instance),
|
||||
groups: make(map[string]*Group),
|
||||
}
|
||||
|
||||
builder := &Builder{
|
||||
registry: registry,
|
||||
name: "no-group-singleton",
|
||||
}
|
||||
|
||||
err := builder.WithInitializer(func() interface{} {
|
||||
return "value"
|
||||
}).Register()
|
||||
|
||||
if err != nil {
|
||||
t.Errorf("Registration without group should succeed, got error: %v", err)
|
||||
}
|
||||
|
||||
// Verify singleton was registered
|
||||
if len(registry.instances) != 1 {
|
||||
t.Error("Singleton should be registered")
|
||||
}
|
||||
}
|
||||
|
||||
// TestContainsHelper tests the helper string contains function
|
||||
func TestContainsHelper(t *testing.T) {
|
||||
tests := []struct {
|
||||
s string
|
||||
substr string
|
||||
expect bool
|
||||
}{
|
||||
{"hello world", "world", true},
|
||||
{"hello world", "hello", true},
|
||||
{"hello world", "lo wo", true},
|
||||
{"hello world", "xyz", false},
|
||||
{"hello", "hello world", false},
|
||||
{"", "test", false},
|
||||
{"test", "", true},
|
||||
{"", "", true},
|
||||
}
|
||||
|
||||
for _, test := range tests {
|
||||
result := contains(test.s, test.substr)
|
||||
if result != test.expect {
|
||||
t.Errorf("contains(%q, %q) = %v, want %v", test.s, test.substr, result, test.expect)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestRegistry_ConcurrentAccess tests concurrent access to registry
|
||||
func TestRegistry_ConcurrentAccess(t *testing.T) {
|
||||
registry := &Registry{
|
||||
instances: make(map[string]*Instance),
|
||||
groups: make(map[string]*Group),
|
||||
}
|
||||
|
||||
callCount := int32(0)
|
||||
initializer := func() interface{} {
|
||||
atomic.AddInt32(&callCount, 1)
|
||||
return "concurrent-value"
|
||||
}
|
||||
|
||||
// Register singleton
|
||||
err := registry.Register("concurrent-singleton", initializer, nil)
|
||||
if err != nil {
|
||||
t.Errorf("Register should succeed, got error: %v", err)
|
||||
}
|
||||
|
||||
var wg sync.WaitGroup
|
||||
numGoroutines := 50
|
||||
|
||||
// Concurrent access
|
||||
wg.Add(numGoroutines)
|
||||
for i := 0; i < numGoroutines; i++ {
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
value, err := registry.GetInstance("concurrent-singleton")
|
||||
if err != nil {
|
||||
t.Errorf("GetInstance should succeed, got error: %v", err)
|
||||
return
|
||||
}
|
||||
if value != "concurrent-value" {
|
||||
t.Errorf("Value should be 'concurrent-value', got '%v'", value)
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
wg.Wait()
|
||||
|
||||
// Initializer should be called only once despite concurrent access
|
||||
if atomic.LoadInt32(&callCount) != 1 {
|
||||
t.Errorf("Initializer should be called only once, called %d times", callCount)
|
||||
}
|
||||
}
|
||||
|
||||
// TestRegistry_ConcurrentReferenceOperations tests concurrent reference operations
|
||||
func TestRegistry_ConcurrentReferenceOperations(t *testing.T) {
|
||||
registry := &Registry{
|
||||
instances: make(map[string]*Instance),
|
||||
groups: make(map[string]*Group),
|
||||
}
|
||||
|
||||
// Register singleton
|
||||
err := registry.Register("ref-singleton", func() interface{} {
|
||||
return "ref-value"
|
||||
}, nil)
|
||||
if err != nil {
|
||||
t.Errorf("Register should succeed, got error: %v", err)
|
||||
}
|
||||
|
||||
// Initialize singleton
|
||||
_, _ = registry.GetInstance("ref-singleton")
|
||||
|
||||
var wg sync.WaitGroup
|
||||
numGoroutines := 20
|
||||
|
||||
// Concurrent reference operations
|
||||
wg.Add(numGoroutines * 2)
|
||||
for i := 0; i < numGoroutines; i++ {
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
_ = registry.AddReference("ref-singleton")
|
||||
}()
|
||||
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
_ = registry.ReleaseReference("ref-singleton")
|
||||
}()
|
||||
}
|
||||
|
||||
wg.Wait()
|
||||
|
||||
// Reference count should be consistent (initial 1 + net operations)
|
||||
count, err := registry.GetReferenceCount("ref-singleton")
|
||||
if err != nil {
|
||||
t.Errorf("GetReferenceCount should succeed, got error: %v", err)
|
||||
}
|
||||
|
||||
// Count should be >= 0 due to balanced add/release operations
|
||||
if count < 0 {
|
||||
t.Errorf("Reference count should not be negative, got %d", count)
|
||||
}
|
||||
}
|
||||
|
||||
// Benchmark tests for performance verification
|
||||
func BenchmarkRegistry_GetInstance(b *testing.B) {
|
||||
registry := &Registry{
|
||||
instances: make(map[string]*Instance),
|
||||
groups: make(map[string]*Group),
|
||||
}
|
||||
|
||||
registry.Register("benchmark-singleton", func() interface{} {
|
||||
return "benchmark-value"
|
||||
}, nil)
|
||||
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
registry.GetInstance("benchmark-singleton")
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkRegistry_ConcurrentGetInstance(b *testing.B) {
|
||||
registry := &Registry{
|
||||
instances: make(map[string]*Instance),
|
||||
groups: make(map[string]*Group),
|
||||
}
|
||||
|
||||
registry.Register("concurrent-benchmark", func() interface{} {
|
||||
return "concurrent-value"
|
||||
}, nil)
|
||||
|
||||
b.ResetTimer()
|
||||
b.RunParallel(func(pb *testing.PB) {
|
||||
for pb.Next() {
|
||||
registry.GetInstance("concurrent-benchmark")
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func BenchmarkBuilder_Register(b *testing.B) {
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
registry := &Registry{
|
||||
instances: make(map[string]*Instance),
|
||||
groups: make(map[string]*Group),
|
||||
}
|
||||
|
||||
builder := &Builder{
|
||||
registry: registry,
|
||||
name: fmt.Sprintf("benchmark-%d", i),
|
||||
}
|
||||
|
||||
builder.WithInitializer(func() interface{} {
|
||||
return "value"
|
||||
}).Register()
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user