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,18 +23,54 @@ 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.
+// It uses double-checked locking pattern for thread-safe lazy initialization
+// and is safe for concurrent use. It also handles the case where Validate()
+// was not called by lazily initializing the atomic.Pointer.
+func (d *Dynamic) getOnce() *sync.Once {
+	// Handle case where Validate() was not called
+	if d.once == nil {
+		d.once = &atomic.Pointer[*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 atomically replaces the sync.Once with a new instance,
+// allowing the fetch operation to be retried. This is called when
+// fetch fails to enable retry on the next call.
+func (d *Dynamic) resetOnce() {
+	// Atomically swap with a new sync.Once
+	once := &sync.Once{}
+	d.once.Store(&once)
+}
+
+// GetDomainAndDomainRegex returns all domains and domain regexes from the dynamic
+// secret and its embedded secrets. It deduplicates the results before returning.
 func (d *Dynamic) GetDomainAndDomainRegex() ([]string, []string) {
 	var domains []string
 	var domainRegex []string
@@ -51,6 +87,8 @@ func (d *Dynamic) GetDomainAndDomainRegex() ([]string, []string) {
 	return uniqueDomains, uniqueDomainRegex
 }
 
+// UnmarshalJSON implements json.Unmarshaler for Dynamic.
+// It handles the inline Secret embedding correctly during JSON unmarshalling.
 func (d *Dynamic) UnmarshalJSON(data []byte) error {
 	if d == nil {
 		return errkit.New("cannot unmarshal into nil Dynamic struct")
@@ -70,8 +108,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")
 	}
@@ -94,32 +134,16 @@ func (d *Dynamic) Validate() error {
 	return nil
 }
 
-// SetLazyFetchCallback sets the lazy fetch callback for the dynamic secret
+// SetLazyFetchCallback sets the lazy fetch callback for the dynamic secret.
+// The callback will be invoked when Fetch() or GetStrategies() is first called.
 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
 }
 
+// applyValuesToSecret replaces template variables (e.g., {{token}}) in the
+// secret's headers, cookies, params, username, password, and token fields
+// with the corresponding values from the Dynamic's Extracted map.
+// It also parses raw cookies after template replacement.
 func (d *Dynamic) applyValuesToSecret(secret *Secret) error {
 	// evaluate headers
 	for i, header := range secret.Headers {
@@ -181,20 +205,79 @@ 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
+// the extracted values in a single atomic operation. This method MUST be called
+// under sync.Once guard to ensure thread-safe fetch-and-hydrate semantics.
+// On error, the once guard is reset to allow retry on the next call.
+func (d *Dynamic) fetchAndHydrate() {
+	d.mu.Lock()
+	// Check if fetchCallback is nil before calling
+	if d.fetchCallback == nil {
+		d.err = fmt.Errorf("fetchCallback is not set for dynamic secret")
+		d.mu.Unlock()
+		// Reset once to allow retry on next call
+		d.resetOnce()
+		return
+	}
+	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.err = err
+			d.mu.Unlock()
+			// 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.err = err
+			d.mu.Unlock()
+			// Reset once to allow retry on next call
+			d.resetOnce()
+			return
+		}
+	}
+	d.mu.Unlock()
+
+	// Mark as fetched successfully (only after successful fetch and hydration)
+	// Check for nil to handle case where Validate() was not called
+	if d.fetched != nil {
+		d.fetched.Store(true)
+	}
+}
+
+// GetStrategies returns the auth strategies for the dynamic secret.
+// It ensures that fetch and hydrate are called exactly once, and all concurrent
+// callers block until the operation completes. If the fetch fails, it returns nil.
+// The once guard is reset on failure to allow retry on the 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())
@@ -205,33 +288,46 @@ func (d *Dynamic) GetStrategies() []AuthStrategy {
 	return strategies
 }
 
-// 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
-	}
-
-	// Try to set fetching flag atomically
-	if !d.fetching.CompareAndSwap(false, true) {
-		// Already fetching, return current error
-		return d.error
-	}
+// Fetch triggers the lazy fetch of the dynamic secret and returns any error.
+// It ensures fetch and hydrate are called exactly once, and all concurrent
+// callers block until the operation completes. On error, the once guard is
+// reset to allow retry on the 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)
 
-	// 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 from the last fetch operation, if any.
+// It is safe for concurrent use.
+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.Extracted = nil
+	d.mu.Unlock()
 }
 
-// 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()
 }