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"
@@ -30,11 +30,27 @@ type Dynamic struct {
 	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.Value           `json:"-" yaml:"-"` // stores *sync.Once, allows retry on failure
 	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
 }
 
+// getOnce returns the current sync.Once instance, creating a new one if needed
+func (d *Dynamic) getOnce() *sync.Once {
+	if v := d.once.Load(); v != nil {
+		return v.(*sync.Once)
+	}
+	once := &sync.Once{}
+	d.once.Store(once)
+	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() {
+	d.once.Store(&sync.Once{})
+}
+
 func (d *Dynamic) GetDomainAndDomainRegex() ([]string, []string) {
 	var domains []string
 	var domainRegex []string
@@ -70,8 +86,8 @@ func (d *Dynamic) UnmarshalJSON(data []byte) error {
 
 // Validate validates the dynamic secret
 func (d *Dynamic) Validate() error {
+	d.once.Store(&sync.Once{})
 	d.fetched = &atomic.Bool{}
-	d.fetching = &atomic.Bool{}
 	if d.TemplatePath == "" {
 		return errkit.New(" template-path is required for dynamic secret")
 	}
@@ -96,28 +112,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 +176,57 @@ 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.error = d.fetchCallback(d)
+	if d.error != nil {
+		// Reset once to allow retry on next call
+		d.resetOnce()
+		return
+	}
+	if len(d.Extracted) == 0 {
+		d.error = fmt.Errorf("no extracted values found for dynamic secret")
+		// Reset once to allow retry on next call
+		d.resetOnce()
+		return
+	}
+
+	if d.Secret != nil {
+		if err := d.applyValuesToSecret(d.Secret); err != nil {
+			d.error = err
+			// Reset once to allow retry on next call
+			d.resetOnce()
+			return
 		}
-	} else {
-		// Try to fetch if not already fetched
-		_ = d.Fetch(true)
 	}
 
+	for _, secret := range d.Secrets {
+		if err := d.applyValuesToSecret(secret); err != nil {
+			d.error = err
+			// Reset once to allow retry on next call
+			d.resetOnce()
+			return
+		}
+	}
+
+	// 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
 	if d.error != nil {
 		return nil
 	}
+
 	var strategies []AuthStrategy
 	if d.Secret != nil {
 		strategies = append(strategies, d.GetStrategy())
@@ -206,32 +238,33 @@ 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)
-
-	// Mark as fetched and clear fetching flag
-	d.fetched.Store(true)
-	d.fetching.Store(false)
-
-	if d.error != nil && isFatal {
-		gologger.Fatal().Msgf("Could not fetch dynamic secret: %s\n", d.error)
-	}
 	return d.error
 }
 
 // Error returns the error if any
 func (d *Dynamic) Error() error {
 	return 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() {
+	d.once.Store(&sync.Once{})
+	d.fetched.Store(false)
+	d.error = nil
+	d.Extracted = nil
+}
+
+// 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()
+}