Add more support for fs.FS in template parsing

pkg/catalog/disk/find.go

@@ -216,12 +216,17 @@ func (c *DiskCatalog) findGlobPathMatches(absPath string, processed map[string]s
 // is a file, it returns true otherwise false with no errors.
 func (c *DiskCatalog) findFileMatches(absPath string, processed map[string]struct{}) (match string, matched bool, err error) {
 	if c.templatesFS != nil {
-		absPath = strings.TrimPrefix(absPath, "/")
+		absPath = strings.Trim(absPath, "/")
 	}
 	var info fs.File
 	if c.templatesFS == nil {
 		info, err = os.Open(absPath)
 	} else {
+		// If we were given no path, then it's not a file, it's the root, and we can quietly return.
+		if absPath == "" {
+			return "", false, nil
+		}
+
 		info, err = c.templatesFS.Open(absPath)
 	}
 	if err != nil {
@@ -263,6 +268,12 @@ func (c *DiskCatalog) findDirectoryMatches(absPath string, processed map[string]
 			},
 		)
 	} else {
+		// For the special case of the root directory, we need to pass "." to `fs.WalkDir`.
+		if absPath == "" {
+			absPath = "."
+		}
+		absPath = strings.TrimSuffix(absPath, "/")
+
 		err = fs.WalkDir(
 			c.templatesFS,
 			absPath,

pkg/catalog/disk/path.go

@@ -2,6 +2,7 @@ package disk
 
 import (
 	"fmt"
+	"io/fs"
 	"os"
 	"path/filepath"
 	"strings"
@@ -21,24 +22,31 @@ func (c *DiskCatalog) ResolvePath(templateName, second string) (string, error) {
 	if filepath.IsAbs(templateName) {
 		return templateName, nil
 	}
+	if c.templatesFS != nil {
+		if potentialPath, err := c.tryResolve(templateName); err != errNoValidCombination {
+			return potentialPath, nil
+		}
+	}
 	if second != "" {
 		secondBasePath := filepath.Join(filepath.Dir(second), templateName)
 		if potentialPath, err := c.tryResolve(secondBasePath); err != errNoValidCombination {
 			return potentialPath, nil
 		}
 	}
 
-	curDirectory, err := os.Getwd()
-	if err != nil {
-		return "", err
-	}
+	if c.templatesFS == nil {
+		curDirectory, err := os.Getwd()
+		if err != nil {
+			return "", err
+		}
 
-	templatePath := filepath.Join(curDirectory, templateName)
-	if potentialPath, err := c.tryResolve(templatePath); err != errNoValidCombination {
-		return potentialPath, nil
+		templatePath := filepath.Join(curDirectory, templateName)
+		if potentialPath, err := c.tryResolve(templatePath); err != errNoValidCombination {
+			return potentialPath, nil
+		}
 	}
 
-	templatePath = filepath.Join(config.DefaultConfig.GetTemplateDir(), templateName)
+	templatePath := filepath.Join(config.DefaultConfig.GetTemplateDir(), templateName)
 	if potentialPath, err := c.tryResolve(templatePath); err != errNoValidCombination {
 		return potentialPath, nil
 	}
@@ -50,8 +58,14 @@ var errNoValidCombination = errors.New("no valid combination found")
 
 // tryResolve attempts to load locate the target by iterating across all the folders tree
 func (c *DiskCatalog) tryResolve(fullPath string) (string, error) {
-	if fileutil.FileOrFolderExists(fullPath) {
-		return fullPath, nil
+	if c.templatesFS == nil {
+		if fileutil.FileOrFolderExists(fullPath) {
+			return fullPath, nil
+		}
+	} else {
+		if _, err := fs.Stat(c.templatesFS, fullPath); err == nil {
+			return fullPath, nil
+		}
 	}
 	return "", errNoValidCombination
 }

pkg/protocols/common/generators/validate.go

@@ -22,6 +22,13 @@ func (g *PayloadGenerator) validate(payloads map[string]interface{}, templatePat
 				return errors.New("invalid number of lines in payload")
 			}
 
+			// For historical reasons, "validate" checks to see if the payload file exist.
+			// If we're using a custom helper function, then we need to skip any validation beyond just checking the string syntax.
+			// Actually attempting to load the file will determine whether or not it exists.
+			if g.options.LoadHelperFileFunction != nil {
+				return nil
+			}
+
 			// check if it's a file and try to load it
 			if fileutil.FileExists(payloadType) {
 				continue

pkg/types/types.go

@@ -23,6 +23,9 @@ var (
 	ErrNoMoreRequests = io.EOF
 )
 
+// LoadHelperFileFunction can be used to load a helper file.
+type LoadHelperFileFunction func(helperFile, templatePath string, catalog catalog.Catalog) (io.ReadCloser, error)
+
 // Options contains the configuration options for nuclei scanner.
 type Options struct {
 	// Tags contains a list of tags to execute templates for. Multiple paths
@@ -408,6 +411,9 @@ type Options struct {
 	HttpApiEndpoint string
 	// ListTemplateProfiles lists all available template profiles
 	ListTemplateProfiles bool
+	// LoadHelperFileFunction is a function that will be used to execute LoadHelperFile.
+	// If none is provided, then the default implementation will be used.
+	LoadHelperFileFunction LoadHelperFileFunction
 	// timeouts contains various types of timeouts used in nuclei
 	// these timeouts are derived from dial-timeout (-timeout) with known multipliers
 	// This is internally managed and does not need to be set by user by explicitly setting
@@ -540,10 +546,21 @@ func (options *Options) ParseHeadlessOptionalArguments() map[string]string {
 	return optionalArguments
 }
 
-// LoadHelperFile loads a helper file needed for the template
+// LoadHelperFile loads a helper file needed for the template.
+//
+// If LoadHelperFileFunction is set, then that function will be used.
+// Otherwise, the default implementation will be used, which respects the sandbox rules and only loads files from allowed directories.
+func (options *Options) LoadHelperFile(helperFile, templatePath string, catalog catalog.Catalog) (io.ReadCloser, error) {
+	if options.LoadHelperFileFunction != nil {
+		return options.LoadHelperFileFunction(helperFile, templatePath, catalog)
+	}
+	return options.defaultLoadHelperFile(helperFile, templatePath, catalog)
+}
+
+// defaultLoadHelperFile loads a helper file needed for the template
 // this respects the sandbox rules and only loads files from
 // allowed directories
-func (options *Options) LoadHelperFile(helperFile, templatePath string, catalog catalog.Catalog) (io.ReadCloser, error) {
+func (options *Options) defaultLoadHelperFile(helperFile, templatePath string, catalog catalog.Catalog) (io.ReadCloser, error) {
 	if !options.AllowLocalFileAccess {
 		// if global file access is disabled try loading with restrictions
 		absPath, err := options.GetValidAbsPath(helperFile, templatePath)