fix(authx): prevent race condition and enable retry on fetch failure

pkg/authprovider/authx/dynamic.go

@@ -3,9 +3,9 @@ package authx
 import (
 	"fmt"
 	"strings"
+	"sync"
 	"sync/atomic"
 
-	"github.com/projectdiscovery/gologger"
 	"github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/replacer"
 	"github.com/projectdiscovery/nuclei/v3/pkg/utils/json"
 	"github.com/projectdiscovery/utils/errkit"
@@ -23,16 +23,43 @@ var (
 // ex: username and password are dynamic secrets, the actual secret is the token obtained
 // after authenticating with the username and password
 type Dynamic struct {
-	*Secret       `yaml:",inline"`       // this is a static secret that will be generated after the dynamic secret is resolved
-	Secrets       []*Secret              `yaml:"secrets"`
-	TemplatePath  string                 `json:"template" yaml:"template"`
-	Variables     []KV                   `json:"variables" yaml:"variables"`
-	Input         string                 `json:"input" yaml:"input"` // (optional) target for the dynamic secret
-	Extracted     map[string]interface{} `json:"-" yaml:"-"`         // extracted values from the dynamic secret
-	fetchCallback LazyFetchSecret        `json:"-" yaml:"-"`
-	fetched       *atomic.Bool           `json:"-" yaml:"-"` // atomic flag to check if the secret has been fetched
-	fetching      *atomic.Bool           `json:"-" yaml:"-"` // atomic flag to prevent recursive fetch calls
-	error         error                  `json:"-" yaml:"-"` // error if any
+	*Secret       `yaml:",inline"` // this is a static secret that will be generated after the dynamic secret is resolved
+	Secrets       []*Secret        `yaml:"secrets"`
+	TemplatePath  string           `json:"template" yaml:"template"`
+	Variables     []KV             `json:"variables" yaml:"variables"`
+	Input         string           `json:"input" yaml:"input"` // (optional) target for the dynamic secret
+	Extracted     map[string]interface{} `json:"-" yaml:"-"`   // extracted values from the dynamic secret
+	fetchCallback LazyFetchSecret  `json:"-" yaml:"-"`
+	once          *atomic.Pointer[*sync.Once] `json:"-" yaml:"-"` // stores *sync.Once, allows retry on failure
+	fetched       *atomic.Bool                `json:"-" yaml:"-"` // atomic flag to check if the secret has been fetched
+	mu            sync.RWMutex                `json:"-" yaml:"-"` // protects err field
+	err           error                       `json:"-" yaml:"-"` // error if any
+}
+
+// getOnce returns the current sync.Once instance, creating a new one if needed
+// Uses double-checked locking pattern for thread-safe lazy initialization
+func (d *Dynamic) getOnce() *sync.Once {
+	// Fast path - check if already initialized
+	ptr := d.once.Load()
+	if ptr != nil {
+		return *ptr
+	}
+	// Slow path - create new sync.Once
+	once := &sync.Once{}
+	// Try to store - if another goroutine beat us, use theirs
+	if !d.once.CompareAndSwap(nil, &once) {
+		// Someone else stored, use their value
+		return *d.once.Load()
+	}
+	return once
+}
+
+// resetOnce resets the sync.Once, allowing retry on next fetch call
+// this is called when fetch fails to enable retry
+func (d *Dynamic) resetOnce() {
+	// Atomically swap with a new sync.Once
+	once := &sync.Once{}
+	d.once.Store(&once)
 }
 
 func (d *Dynamic) GetDomainAndDomainRegex() ([]string, []string) {
@@ -70,8 +97,10 @@ func (d *Dynamic) UnmarshalJSON(data []byte) error {
 
 // Validate validates the dynamic secret
 func (d *Dynamic) Validate() error {
+	d.once = &atomic.Pointer[*sync.Once]{}
+	once := &sync.Once{}
+	d.once.Store(&once)
 	d.fetched = &atomic.Bool{}
-	d.fetching = &atomic.Bool{}
 	if d.TemplatePath == "" {
 		return errkit.New(" template-path is required for dynamic secret")
 	}
@@ -96,28 +125,7 @@ func (d *Dynamic) Validate() error {
 
 // SetLazyFetchCallback sets the lazy fetch callback for the dynamic secret
 func (d *Dynamic) SetLazyFetchCallback(callback LazyFetchSecret) {
-	d.fetchCallback = func(d *Dynamic) error {
-		err := callback(d)
-		if err != nil {
-			return err
-		}
-		if len(d.Extracted) == 0 {
-			return fmt.Errorf("no extracted values found for dynamic secret")
-		}
-
-		if d.Secret != nil {
-			if err := d.applyValuesToSecret(d.Secret); err != nil {
-				return err
-			}
-		}
-
-		for _, secret := range d.Secrets {
-			if err := d.applyValuesToSecret(secret); err != nil {
-				return err
-			}
-		}
-		return nil
-	}
+	d.fetchCallback = callback
 }
 
 func (d *Dynamic) applyValuesToSecret(secret *Secret) error {
@@ -181,20 +189,66 @@ func (d *Dynamic) applyValuesToSecret(secret *Secret) error {
 	return nil
 }
 
-// GetStrategy returns the auth strategies for the dynamic secret
-func (d *Dynamic) GetStrategies() []AuthStrategy {
-	if d.fetched.Load() {
-		if d.error != nil {
-			return nil
+// fetchAndHydrate executes the fetch callback and hydrates all secrets with extracted values
+// this MUST be called under sync.Once guard to ensure atomic fetch-and-hydrate
+// On error, the once guard is reset to allow retry on next call
+func (d *Dynamic) fetchAndHydrate() {
+	d.mu.Lock()
+	d.err = d.fetchCallback(d)
+	if d.err != nil {
+		d.mu.Unlock()
+		// Reset once to allow retry on next call
+		d.resetOnce()
+		return
+	}
+	if len(d.Extracted) == 0 {
+		d.err = fmt.Errorf("no extracted values found for dynamic secret")
+		d.mu.Unlock()
+		// Reset once to allow retry on next call
+		d.resetOnce()
+		return
+	}
+
+	if d.Secret != nil {
+		if err := d.applyValuesToSecret(d.Secret); err != nil {
+			d.mu.Unlock()
+			d.err = err
+			// Reset once to allow retry on next call
+			d.resetOnce()
+			return
 		}
-	} else {
-		// Try to fetch if not already fetched
-		_ = d.Fetch(true)
 	}
 
-	if d.error != nil {
+	for _, secret := range d.Secrets {
+		if err := d.applyValuesToSecret(secret); err != nil {
+			d.mu.Unlock()
+			d.err = err
+			// Reset once to allow retry on next call
+			d.resetOnce()
+			return
+		}
+	}
+	d.mu.Unlock()
+
+	// Mark as fetched successfully (only after successful fetch and hydration)
+	d.fetched.Store(true)
+}
+
+// GetStrategies returns the auth strategies for the dynamic secret
+// It ensures fetch and hydrate are called exactly once and all callers block until complete
+// If fetch fails, the once guard is reset to allow retry on next call
+func (d *Dynamic) GetStrategies() []AuthStrategy {
+	// Use sync.Once to ensure fetch and hydrate are called exactly once and all callers block until complete
+	d.getOnce().Do(d.fetchAndHydrate)
+
+	// If fetch failed, return nil strategies
+	d.mu.RLock()
+	hasError := d.err != nil
+	d.mu.RUnlock()
+	if hasError {
 		return nil
 	}
+
 	var strategies []AuthStrategy
 	if d.Secret != nil {
 		strategies = append(strategies, d.GetStrategy())
@@ -206,32 +260,42 @@ func (d *Dynamic) GetStrategies() []AuthStrategy {
 }
 
 // Fetch fetches the dynamic secret
-// if isFatal is true, it will stop the execution if the secret could not be fetched
-func (d *Dynamic) Fetch(isFatal bool) error {
-	if d.fetched.Load() {
-		return d.error
-	}
+// If fetch fails, the once guard is reset to allow retry on next call
+func (d *Dynamic) Fetch() error {
+	// Use sync.Once to ensure fetch and hydrate are called exactly once and all callers block until complete
+	d.getOnce().Do(d.fetchAndHydrate)
 
-	// Try to set fetching flag atomically
-	if !d.fetching.CompareAndSwap(false, true) {
-		// Already fetching, return current error
-		return d.error
-	}
-
-	// We're the only one fetching, call the callback
-	d.error = d.fetchCallback(d)
+	d.mu.RLock()
+	defer d.mu.RUnlock()
+	return d.err
+}
 
-	// Mark as fetched and clear fetching flag
-	d.fetched.Store(true)
-	d.fetching.Store(false)
+// Error returns the error if any
+func (d *Dynamic) Error() error {
+	d.mu.RLock()
+	defer d.mu.RUnlock()
+	return d.err
+}
 
-	if d.error != nil && isFatal {
-		gologger.Fatal().Msgf("Could not fetch dynamic secret: %s\n", d.error)
+// Reset resets the fetch state, allowing a fresh fetch on next call
+// This is useful when you want to force a re-fetch of the dynamic secret
+// Call Validate() again after Reset() if you want to ensure the secret is still valid
+func (d *Dynamic) Reset() {
+	once := &sync.Once{}
+	d.once.Store(&once)
+	if d.fetched != nil {
+		d.fetched.Store(false)
 	}
-	return d.error
+	d.mu.Lock()
+	d.err = nil
+	d.mu.Unlock()
+	d.Extracted = nil
 }
 
-// Error returns the error if any
-func (d *Dynamic) Error() error {
-	return d.error
+// IsFetched returns true if the dynamic secret has been successfully fetched
+func (d *Dynamic) IsFetched() bool {
+	if d.fetched == nil {
+		return false
+	}
+	return d.fetched.Load()
 }