Use fs.FS when explicitly given (#5312)

* Use the `templateFS` if it's there when calling `OpenFile`

* Add a new constructor

* More refactoring

* Both of my use cases are working

* Fix for legacy assumptions

* minor update: remove gologger debug stmts

---------

Co-authored-by: Tarun Koyalwar <tarun@projectdiscovery.io>

lib/sdk_private.go

@@ -231,8 +231,11 @@ func (e *NucleiEngine) init(ctx context.Context) error {
 	// and also upgrade templates to latest version if available
 	installer.NucleiSDKVersionCheck()
 
+	if DefaultConfig.CanCheckForUpdates() {
 		return e.processUpdateCheckResults()
 	}
+	return nil
+}
 
 type syncOnce struct {
 	sync.Once

pkg/catalog/disk/catalog.go

@@ -11,17 +11,25 @@ import (
 // DiskCatalog is a template catalog helper implementation based on disk
 type DiskCatalog struct {
 	templatesDirectory string
-	templatesFS        fs.FS // TODO: Refactor to use this
+	templatesFS        fs.FS // Due to issues with how Go has implemented fs.FS, we'll have to also implement normal os operations, as well.  See: https://github.com/golang/go/issues/44279
 }
 
 // NewCatalog creates a new Catalog structure using provided input items
 // using disk based items
 func NewCatalog(directory string) *DiskCatalog {
 	catalog := &DiskCatalog{templatesDirectory: directory}
-	if directory != "" {
-		catalog.templatesFS = os.DirFS(directory)
-	} else {
-		catalog.templatesFS = os.DirFS(config.DefaultConfig.GetTemplateDir())
+	if directory == "" {
+		catalog.templatesDirectory = config.DefaultConfig.GetTemplateDir()
+	}
+	return catalog
+}
+
+// NewFSCatalog creates a new Catalog structure using provided input items
+// using the fs.FS as its filesystem.
+func NewFSCatalog(fs fs.FS, directory string) *DiskCatalog {
+	catalog := &DiskCatalog{
+		templatesDirectory: directory,
+		templatesFS:        fs,
 	}
 	return catalog
 }
@@ -29,6 +37,7 @@ func NewCatalog(directory string) *DiskCatalog {
 // OpenFile opens a file and returns an io.ReadCloser to the file.
 // It is used to read template and payload files based on catalog responses.
 func (d *DiskCatalog) OpenFile(filename string) (io.ReadCloser, error) {
+	if d.templatesFS == nil {
 		file, err := os.Open(filename)
 		if err != nil {
 			if file, errx := os.Open(BackwardsCompatiblePaths(d.templatesDirectory, filename)); errx == nil {
@@ -37,3 +46,6 @@ func (d *DiskCatalog) OpenFile(filename string) (io.ReadCloser, error) {
 		}
 		return file, err
 	}
+
+	return d.templatesFS.Open(filename)
+}

pkg/catalog/disk/find.go

@@ -79,7 +79,9 @@ func (c *DiskCatalog) GetTemplatePath(target string) ([]string, error) {
 	}
 
 	// try to handle deprecated template paths
-	absPath := BackwardsCompatiblePaths(c.templatesDirectory, target)
+	absPath := target
+	if c.templatesFS == nil {
+		absPath = BackwardsCompatiblePaths(c.templatesDirectory, target)
 		if absPath != target && strings.TrimPrefix(absPath, c.templatesDirectory+string(filepath.Separator)) != target {
 			if config.DefaultConfig.LogAllEvents {
 				gologger.DefaultLogger.Print().Msgf("[%v] requested Template path %s is deprecated, please update to %s\n", aurora.Yellow("WRN").String(), target, absPath)
@@ -87,10 +89,12 @@ func (c *DiskCatalog) GetTemplatePath(target string) ([]string, error) {
 			deprecatedPathsCounter++
 		}
 
-	absPath, err := c.convertPathToAbsolute(absPath)
+		var err error
+		absPath, err = c.convertPathToAbsolute(absPath)
 		if err != nil {
 			return nil, errors.Wrapf(err, "could not find template file")
 		}
+	}
 
 	// Template input is either a file or a directory
 	match, file, err := c.findFileMatches(absPath, processed)
@@ -143,17 +147,41 @@ func (c *DiskCatalog) findGlobPathMatches(absPath string, processed map[string]s
 		if c.templatesDirectory == "" {
 			templateDir = "./"
 		}
+
+		if c.templatesFS == nil {
 			matches, _ := fs.Glob(os.DirFS(filepath.Join(templateDir, "http")), inputGlob)
 			if len(matches) != 0 {
 				return matches
 			}
+
 			// condition to support network cve related globs
 			matches, _ = fs.Glob(os.DirFS(filepath.Join(templateDir, "network")), inputGlob)
 			return matches
+		} else {
+			sub, err := fs.Sub(c.templatesFS, filepath.Join(templateDir, "http"))
+			if err != nil {
+				return nil
+			}
+			matches, _ := fs.Glob(sub, inputGlob)
+			if len(matches) != 0 {
+				return matches
+			}
+
+			// condition to support network cve related globs
+			sub, err = fs.Sub(c.templatesFS, filepath.Join(templateDir, "network"))
+			if err != nil {
+				return nil
+			}
+			matches, _ = fs.Glob(sub, inputGlob)
+			return matches
+		}
 	}
 
 	var matched []string
-	matches, err := fs.Glob(c.templatesFS, relPath)
+	var matches []string
+	if c.templatesFS == nil {
+		var err error
+		matches, err = filepath.Glob(relPath)
 		if len(matches) != 0 {
 			matched = append(matched, matches...)
 		} else {
@@ -162,6 +190,18 @@ func (c *DiskCatalog) findGlobPathMatches(absPath string, processed map[string]s
 		if err != nil && len(matched) == 0 {
 			return nil, errors.Errorf("wildcard found, but unable to glob: %s\n", err)
 		}
+	} else {
+		var err error
+		matches, err = fs.Glob(c.templatesFS, relPath)
+		if len(matches) != 0 {
+			matched = append(matched, matches...)
+		} else {
+			matched = append(matched, OldPathsResolver(relPath)...)
+		}
+		if err != nil && len(matched) == 0 {
+			return nil, errors.Errorf("wildcard found, but unable to glob: %s\n", err)
+		}
+	}
 	results := make([]string, 0, len(matches))
 	for _, match := range matches {
 		if _, ok := processed[match]; !ok {
@@ -175,11 +215,23 @@ func (c *DiskCatalog) findGlobPathMatches(absPath string, processed map[string]s
 // findFileMatches finds if a path is an absolute file. If the path
 // 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) {
-	info, err := os.Stat(absPath)
+	if c.templatesFS != nil {
+		absPath = strings.TrimPrefix(absPath, "/")
+	}
+	var info fs.File
+	if c.templatesFS == nil {
+		info, err = os.Open(absPath)
+	} else {
+		info, err = c.templatesFS.Open(absPath)
+	}
 	if err != nil {
 		return "", false, err
 	}
-	if !info.Mode().IsRegular() {
+	stat, err := info.Stat()
+	if err != nil {
+		return "", false, err
+	}
+	if !stat.Mode().IsRegular() {
 		return "", false, nil
 	}
 	if _, ok := processed[absPath]; !ok {
@@ -192,7 +244,9 @@ func (c *DiskCatalog) findFileMatches(absPath string, processed map[string]struc
 // findDirectoryMatches finds matches for templates from a directory
 func (c *DiskCatalog) findDirectoryMatches(absPath string, processed map[string]struct{}) ([]string, error) {
 	var results []string
-	err := filepath.WalkDir(
+	var err error
+	if c.templatesFS == nil {
+		err = filepath.WalkDir(
 			absPath,
 			func(path string, d fs.DirEntry, err error) error {
 				// continue on errors
@@ -208,6 +262,25 @@ func (c *DiskCatalog) findDirectoryMatches(absPath string, processed map[string]
 				return nil
 			},
 		)
+	} else {
+		err = fs.WalkDir(
+			c.templatesFS,
+			absPath,
+			func(path string, d fs.DirEntry, err error) error {
+				// continue on errors
+				if err != nil {
+					return nil
+				}
+				if !d.IsDir() && config.GetTemplateFormatFromExt(path) != config.Unknown {
+					if _, ok := processed[path]; !ok {
+						results = append(results, path)
+						processed[path] = struct{}{}
+					}
+				}
+				return nil
+			},
+		)
+	}
 	return results, err
 }