diff --git a/.github/workflows/generate-pgo.yaml b/.github/workflows/generate-pgo.yaml index b0c47d56ec..39fc7e6a10 100644 --- a/.github/workflows/generate-pgo.yaml +++ b/.github/workflows/generate-pgo.yaml @@ -31,6 +31,7 @@ jobs: - uses: actions/checkout@v6 - uses: projectdiscovery/actions/setup/git@v1 - uses: projectdiscovery/actions/setup/go@v1 + - uses: projectdiscovery/actions/cache/nuclei@v1 - name: Generate list run: for i in {1..${{ matrix.targets }}}; do echo "https://honey.scanme.sh/?_=${i}" >> "${LIST_FILE}"; done # NOTE(dwisiswant0): use `-no-mhe` flag to get better samples. diff --git a/.github/workflows/perf-regression.yaml b/.github/workflows/perf-regression.yaml index 85650a09f4..7856067243 100644 --- a/.github/workflows/perf-regression.yaml +++ b/.github/workflows/perf-regression.yaml @@ -14,6 +14,7 @@ jobs: - uses: actions/checkout@v6 - uses: projectdiscovery/actions/setup/go@v1 - uses: projectdiscovery/actions/cache/go-rod-browser@v1 + - uses: projectdiscovery/actions/cache/nuclei@v1 - run: make build-test - run: ./bin/nuclei.test -test.run - -test.bench=. -test.benchmem ./cmd/nuclei/ | tee $BENCH_OUT env: diff --git a/.github/workflows/tests.yaml b/.github/workflows/tests.yaml index e244b2001d..42a46a67db 100644 --- a/.github/workflows/tests.yaml +++ b/.github/workflows/tests.yaml @@ -39,6 +39,7 @@ jobs: - uses: actions/checkout@v6 - uses: projectdiscovery/actions/setup/go@v1 - uses: projectdiscovery/actions/cache/go-rod-browser@v1 + - uses: projectdiscovery/actions/cache/nuclei@v1 - uses: projectdiscovery/actions/free-disk-space@v1 with: llvm: 'false' @@ -66,6 +67,7 @@ jobs: - uses: actions/checkout@v6 - uses: projectdiscovery/actions/setup/go@v1 - uses: projectdiscovery/actions/cache/go-rod-browser@v1 + - uses: projectdiscovery/actions/cache/nuclei@v1 - name: "Simple" run: go run . working-directory: examples/simple/ @@ -88,6 +90,7 @@ jobs: steps: - uses: actions/checkout@v6 - uses: projectdiscovery/actions/setup/go@v1 + - uses: projectdiscovery/actions/cache/nuclei@v1 - uses: projectdiscovery/actions/setup/python@v1 - uses: projectdiscovery/actions/cache/go-rod-browser@v1 - run: bash run.sh "${{ matrix.os }}" @@ -108,6 +111,7 @@ jobs: steps: - uses: actions/checkout@v6 - uses: projectdiscovery/actions/setup/go@v1 + - uses: projectdiscovery/actions/cache/nuclei@v1 - uses: projectdiscovery/actions/setup/python@v1 - uses: projectdiscovery/actions/cache/go-rod-browser@v1 - run: bash run.sh diff --git a/cmd/integration-test/library.go b/cmd/integration-test/library.go index 3513b1d043..2c4cda5764 100644 --- a/cmd/integration-test/library.go +++ b/cmd/integration-test/library.go @@ -15,6 +15,7 @@ import ( "github.com/logrusorgru/aurora" "github.com/pkg/errors" "github.com/projectdiscovery/goflags" + "github.com/projectdiscovery/gologger" "github.com/projectdiscovery/nuclei/v3/pkg/catalog/config" "github.com/projectdiscovery/nuclei/v3/pkg/catalog/disk" "github.com/projectdiscovery/nuclei/v3/pkg/catalog/loader" @@ -70,6 +71,7 @@ func executeNucleiAsLibrary(templatePath, templateURL string) ([]string, error) defaultOpts := types.DefaultOptions() defaultOpts.ExecutionId = "test" + defaultOpts.Logger = gologger.DefaultLogger mockProgress := &testutils.MockProgressClient{} reportingClient, err := reporting.New(&reporting.Options{ExecutionId: defaultOpts.ExecutionId}, "", false) diff --git a/go.mod b/go.mod index 6043a9b47c..18504364ad 100644 --- a/go.mod +++ b/go.mod @@ -87,6 +87,7 @@ require ( github.com/leslie-qiwa/flat v0.0.0-20230424180412-f9d1cf014baa github.com/lib/pq v1.10.9 github.com/mattn/go-sqlite3 v1.14.28 + github.com/maypok86/otter/v2 v2.2.1 github.com/mholt/archives v0.1.5 github.com/microsoft/go-mssqldb v1.9.2 github.com/ory/dockertest/v3 v3.12.0 diff --git a/go.sum b/go.sum index 798cbfc736..8178f4e11f 100644 --- a/go.sum +++ b/go.sum @@ -703,6 +703,8 @@ github.com/mattn/go-sqlite3 v1.14.28 h1:ThEiQrnbtumT+QMknw63Befp/ce/nUPgBPMlRFEu github.com/mattn/go-sqlite3 v1.14.28/go.mod h1:Uh1q+B4BYcTPb+yiD3kU8Ct7aC0hY9fxUwlHK0RXw+Y= github.com/matttproud/golang_protobuf_extensions v1.0.1/go.mod h1:D8He9yQNgCq6Z5Ld7szi9bcBfOoFv/3dc6xSMkL2PC0= github.com/matttproud/golang_protobuf_extensions v1.0.4/go.mod h1:BSXmuO+STAnVfrANrmjBb36TMTDstsz7MSK+HVaYKv4= +github.com/maypok86/otter/v2 v2.2.1 h1:hnGssisMFkdisYcvQ8L019zpYQcdtPse+g0ps2i7cfI= +github.com/maypok86/otter/v2 v2.2.1/go.mod h1:1NKY9bY+kB5jwCXBJfE59u+zAwOt6C7ni1FTlFFMqVs= github.com/mholt/acmez v1.2.0 h1:1hhLxSgY5FvH5HCnGUuwbKY2VQVo8IU7rxXKSnZ7F30= github.com/mholt/acmez v1.2.0/go.mod h1:VT9YwH1xgNX1kmYY89gY8xPJC84BFAisjo8Egigt4kE= github.com/mholt/archives v0.1.5 h1:Fh2hl1j7VEhc6DZs2DLMgiBNChUux154a1G+2esNvzQ= diff --git a/pkg/catalog/index/filter.go b/pkg/catalog/index/filter.go new file mode 100644 index 0000000000..ac4959a531 --- /dev/null +++ b/pkg/catalog/index/filter.go @@ -0,0 +1,344 @@ +package index + +import ( + "path/filepath" + "slices" + "strings" + + "github.com/projectdiscovery/nuclei/v3/pkg/model/types/severity" + "github.com/projectdiscovery/nuclei/v3/pkg/templates/types" +) + +// Filter represents filtering criteria for template metadata. +// +// Inclusion fields (e.g., Authors, Tags, IDs, Severities, ProtocolTypes) use +// AND logic across different filter types and OR logic within each type. +// Exclusion fields (e.g., ExcludeTags, ExcludeIDs, ExcludeSeverities, +// ExcludeProtocolTypes) take precedence over inclusion fields. Additionally, +// IncludeTemplates and IncludeTags can force inclusion of templates even if +// they match exclusion criteria. +type Filter struct { + // Authors to include. + Authors []string + + // Tags to include. + Tags []string + + // ExcludeTags to exclude (takes precedence over Tags). + ExcludeTags []string + + // IncludeTags to force include even if excluded. + IncludeTags []string + + // IDs to include (supports wildcards, OR logic). + IDs []string + + // ExcludeIDs to exclude (supports wildcards). + ExcludeIDs []string + + // IncludeTemplates paths to force include even if excluded. + IncludeTemplates []string + + // ExcludeTemplates paths to exclude. + ExcludeTemplates []string + + // Severities to include. + Severities []severity.Severity + + // ExcludeSeverities to exclude. + ExcludeSeverities []severity.Severity + + // ProtocolTypes to include. + ProtocolTypes []types.ProtocolType + + // ExcludeProtocolTypes to exclude. + ExcludeProtocolTypes []types.ProtocolType +} + +// Matches checks if metadata matches the filter criteria. +func (f *Filter) Matches(m *Metadata) bool { + if f.isForcedInclude(m) { + return true + } + + if f.isExcluded(m) { + return false + } + + if !f.matchesIncludes(m) { + return false + } + + return true +} + +// isForcedInclude checks if template is forced to be included. +func (f *Filter) isForcedInclude(m *Metadata) bool { + if len(f.IncludeTemplates) > 0 { + for _, includePath := range f.IncludeTemplates { + if matchesPath(m.FilePath, includePath) { + return true + } + } + } + + if len(f.IncludeTags) > 0 { + if slices.ContainsFunc(f.IncludeTags, m.HasTag) { + return true + } + } + + return false +} + +// isExcluded checks if template should be excluded. +func (f *Filter) isExcluded(m *Metadata) bool { + if len(f.ExcludeTemplates) > 0 { + for _, excludePath := range f.ExcludeTemplates { + if matchesPath(m.FilePath, excludePath) { + return true + } + } + } + + if len(f.ExcludeTags) > 0 { + if slices.ContainsFunc(f.ExcludeTags, m.HasTag) { + return true + } + } + + if len(f.ExcludeIDs) > 0 { + for _, excludeID := range f.ExcludeIDs { + if matchesID(m.ID, excludeID) { + return true + } + } + } + + if len(f.ExcludeSeverities) > 0 { + if slices.ContainsFunc(f.ExcludeSeverities, m.MatchesSeverity) { + return true + } + } + + if len(f.ExcludeProtocolTypes) > 0 { + if slices.ContainsFunc(f.ExcludeProtocolTypes, m.MatchesProtocol) { + return true + } + } + + return false +} + +// matchesIncludes checks if metadata matches include filters. +// +// Returns true if no include filters are specified, or if all specified filter +// types match. +func (f *Filter) matchesIncludes(m *Metadata) bool { + if len(f.Authors) > 0 { + if !slices.ContainsFunc(f.Authors, m.HasAuthor) { + return false + } + } + + if len(f.Tags) > 0 { + if !slices.ContainsFunc(f.Tags, m.HasTag) { + return false + } + } + + if len(f.IDs) > 0 { + matched := false + for _, id := range f.IDs { + if matchesID(m.ID, id) { + matched = true + break + } + } + if !matched { + return false + } + } + + if len(f.Severities) > 0 { + if !slices.ContainsFunc(f.Severities, m.MatchesSeverity) { + return false + } + } + + if len(f.ProtocolTypes) > 0 { + if !slices.ContainsFunc(f.ProtocolTypes, m.MatchesProtocol) { + return false + } + } + + return true +} + +// matchesID checks if template ID matches pattern (supports wildcards). +func matchesID(templateID, pattern string) bool { + // Convert to lowercase for case-insensitive matching + templateID = strings.ToLower(templateID) + pattern = strings.ToLower(pattern) + + if templateID == pattern { + return true + } + + matched, _ := filepath.Match(pattern, templateID) + + return matched +} + +// matchesPath checks if template path matches pattern. +func matchesPath(templatePath, pattern string) bool { + templatePath = filepath.Clean(templatePath) + pattern = filepath.Clean(pattern) + + if templatePath == pattern { + return true + } + + if strings.HasPrefix(templatePath, pattern+string(filepath.Separator)) { + return true + } + + matched, _ := filepath.Match(pattern, templatePath) + + return matched +} + +// FilterFunc is a function that filters metadata. +type FilterFunc func(*Metadata) bool + +// UnmarshalFilter creates a Filter from nuclei options. +func UnmarshalFilter( + authors, tags, excludeTags, includeTags []string, + ids, excludeIDs []string, + includeTemplates, excludeTemplates []string, + severities, excludeSeverities []string, + protocolTypes, excludeProtocolTypes []string, +) (*Filter, error) { + filter := &Filter{ + Authors: authors, + Tags: tags, + ExcludeTags: excludeTags, + IncludeTags: includeTags, + IDs: ids, + ExcludeIDs: excludeIDs, + IncludeTemplates: includeTemplates, + ExcludeTemplates: excludeTemplates, + } + + for _, sev := range severities { + holder := &severity.Holder{} + if err := holder.UnmarshalYAML(func(v interface{}) error { + *v.(*string) = sev + return nil + }); err == nil { + filter.Severities = append(filter.Severities, holder.Severity) + } + } + + for _, sev := range excludeSeverities { + holder := &severity.Holder{} + if err := holder.UnmarshalYAML(func(v interface{}) error { + *v.(*string) = sev + return nil + }); err == nil { + filter.ExcludeSeverities = append(filter.ExcludeSeverities, holder.Severity) + } + } + + for _, pt := range protocolTypes { + holder := &types.TypeHolder{} + if err := holder.UnmarshalYAML(func(v interface{}) error { + *v.(*string) = pt + return nil + }); err == nil && holder.ProtocolType != types.InvalidProtocol { + filter.ProtocolTypes = append(filter.ProtocolTypes, holder.ProtocolType) + } + } + + for _, pt := range excludeProtocolTypes { + holder := &types.TypeHolder{} + if err := holder.UnmarshalYAML(func(v interface{}) error { + *v.(*string) = pt + return nil + }); err == nil && holder.ProtocolType != types.InvalidProtocol { + filter.ExcludeProtocolTypes = append(filter.ExcludeProtocolTypes, holder.ProtocolType) + } + } + + return filter, nil +} + +// UnmarshalFilterFunc creates a FilterFunc from filter criteria. +func UnmarshalFilterFunc(filter *Filter) FilterFunc { + if filter == nil { + return func(*Metadata) bool { return true } + } + + return filter.Matches +} + +// IsEmpty returns true if filter has no criteria set. +func (f *Filter) IsEmpty() bool { + return len(f.Authors) == 0 && + len(f.Tags) == 0 && + len(f.ExcludeTags) == 0 && + len(f.IncludeTags) == 0 && + len(f.IDs) == 0 && + len(f.ExcludeIDs) == 0 && + len(f.IncludeTemplates) == 0 && + len(f.ExcludeTemplates) == 0 && + len(f.Severities) == 0 && + len(f.ExcludeSeverities) == 0 && + len(f.ProtocolTypes) == 0 && + len(f.ExcludeProtocolTypes) == 0 +} + +// String returns a human-readable representation of the filter. +func (f *Filter) String() string { + var parts []string + + if len(f.Authors) > 0 { + parts = append(parts, "authors="+strings.Join(f.Authors, ",")) + } + + if len(f.Tags) > 0 { + parts = append(parts, "tags="+strings.Join(f.Tags, ",")) + } + + if len(f.ExcludeTags) > 0 { + parts = append(parts, "exclude-tags="+strings.Join(f.ExcludeTags, ",")) + } + + if len(f.IDs) > 0 { + parts = append(parts, "ids="+strings.Join(f.IDs, ",")) + } + + if len(f.Severities) > 0 { + sevs := make([]string, len(f.Severities)) + for i, s := range f.Severities { + sevs[i] = s.String() + } + + parts = append(parts, "severities="+strings.Join(sevs, ",")) + } + + if len(f.ProtocolTypes) > 0 { + pts := make([]string, len(f.ProtocolTypes)) + for i, p := range f.ProtocolTypes { + pts[i] = p.String() + } + + parts = append(parts, "types="+strings.Join(pts, ",")) + } + + if len(parts) == 0 { + return "filter=" + } + + return "filter(" + strings.Join(parts, ", ") + ")" +} diff --git a/pkg/catalog/index/filter_test.go b/pkg/catalog/index/filter_test.go new file mode 100644 index 0000000000..2bc4735e6c --- /dev/null +++ b/pkg/catalog/index/filter_test.go @@ -0,0 +1,407 @@ +package index + +import ( + "os" + "path/filepath" + "testing" + + "github.com/projectdiscovery/nuclei/v3/pkg/model/types/severity" + "github.com/projectdiscovery/nuclei/v3/pkg/templates/types" + "github.com/stretchr/testify/require" +) + +func TestFilterMatches(t *testing.T) { + metadata := &Metadata{ + ID: "test-template-1", + FilePath: "/templates/cves/2021/CVE-2021-1234.yaml", + Name: "Test CVE Template", + Authors: []string{"pdteam", "geeknik"}, + Tags: []string{"cve", "rce", "apache"}, + Severity: "critical", + ProtocolType: "http", + } + + t.Run("Empty filter matches all", func(t *testing.T) { + filter := &Filter{} + require.True(t, filter.Matches(metadata)) + require.True(t, filter.IsEmpty()) + }) + + t.Run("Author filter - match", func(t *testing.T) { + filter := &Filter{Authors: []string{"pdteam"}} + require.True(t, filter.Matches(metadata)) + }) + + t.Run("Author filter - no match", func(t *testing.T) { + filter := &Filter{Authors: []string{"unknown"}} + require.False(t, filter.Matches(metadata)) + }) + + t.Run("Multiple authors - OR logic", func(t *testing.T) { + filter := &Filter{Authors: []string{"unknown", "geeknik"}} + require.True(t, filter.Matches(metadata)) + }) + + t.Run("Tag filter - match", func(t *testing.T) { + filter := &Filter{Tags: []string{"cve"}} + require.True(t, filter.Matches(metadata)) + }) + + t.Run("Tag filter - no match", func(t *testing.T) { + filter := &Filter{Tags: []string{"xss"}} + require.False(t, filter.Matches(metadata)) + }) + + t.Run("Exclude tags - match", func(t *testing.T) { + filter := &Filter{ExcludeTags: []string{"rce"}} + require.False(t, filter.Matches(metadata)) + }) + + t.Run("Include tags overrides exclude", func(t *testing.T) { + filter := &Filter{ + ExcludeTags: []string{"rce"}, + IncludeTags: []string{"cve"}, + } + require.True(t, filter.Matches(metadata)) + }) + + t.Run("ID filter - exact match", func(t *testing.T) { + filter := &Filter{IDs: []string{"test-template-1"}} + require.True(t, filter.Matches(metadata)) + }) + + t.Run("ID filter - wildcard match", func(t *testing.T) { + filter := &Filter{IDs: []string{"test-*"}} + require.True(t, filter.Matches(metadata)) + }) + + t.Run("ID filter - no match", func(t *testing.T) { + filter := &Filter{IDs: []string{"other-*"}} + require.False(t, filter.Matches(metadata)) + }) + + t.Run("Exclude ID - exact match", func(t *testing.T) { + filter := &Filter{ExcludeIDs: []string{"test-template-1"}} + require.False(t, filter.Matches(metadata)) + }) + + t.Run("Exclude ID - wildcard match", func(t *testing.T) { + filter := &Filter{ExcludeIDs: []string{"test-*"}} + require.False(t, filter.Matches(metadata)) + }) + + t.Run("Severity filter - match", func(t *testing.T) { + filter := &Filter{Severities: []severity.Severity{severity.Critical}} + require.True(t, filter.Matches(metadata)) + }) + + t.Run("Severity filter - no match", func(t *testing.T) { + filter := &Filter{Severities: []severity.Severity{severity.High, severity.Medium}} + require.False(t, filter.Matches(metadata)) + }) + + t.Run("Exclude severity - match", func(t *testing.T) { + filter := &Filter{ExcludeSeverities: []severity.Severity{severity.Critical}} + require.False(t, filter.Matches(metadata)) + }) + + t.Run("Protocol type filter - match", func(t *testing.T) { + filter := &Filter{ProtocolTypes: []types.ProtocolType{types.HTTPProtocol}} + require.True(t, filter.Matches(metadata)) + }) + + t.Run("Protocol type filter - no match", func(t *testing.T) { + filter := &Filter{ProtocolTypes: []types.ProtocolType{types.DNSProtocol}} + require.False(t, filter.Matches(metadata)) + }) + + t.Run("Exclude protocol type - match", func(t *testing.T) { + filter := &Filter{ExcludeProtocolTypes: []types.ProtocolType{types.HTTPProtocol}} + require.False(t, filter.Matches(metadata)) + }) + + t.Run("Include templates - path match", func(t *testing.T) { + filter := &Filter{ + ExcludeTags: []string{"cve"}, + IncludeTemplates: []string{"/templates/cves/"}, + } + require.True(t, filter.Matches(metadata)) + }) + + t.Run("Exclude templates - path match", func(t *testing.T) { + filter := &Filter{ + ExcludeTemplates: []string{"/templates/cves/"}, + } + require.False(t, filter.Matches(metadata)) + }) + + t.Run("Complex filter - all match", func(t *testing.T) { + filter := &Filter{ + Authors: []string{"pdteam"}, + Tags: []string{"cve"}, + Severities: []severity.Severity{severity.Critical}, + ProtocolTypes: []types.ProtocolType{types.HTTPProtocol}, + } + require.True(t, filter.Matches(metadata)) + }) + + t.Run("Complex filter - AND logic across types", func(t *testing.T) { + filter := &Filter{ + Authors: []string{"pdteam"}, // matches + Tags: []string{"xss"}, // doesn't match + Severities: []severity.Severity{severity.Critical}, // matches + } + // With AND logic across filter types, doesn't match because tags don't match + // even though author and severity match + require.False(t, filter.Matches(metadata)) + }) + + t.Run("Complex filter - no match at all", func(t *testing.T) { + filter := &Filter{ + Authors: []string{"unknown"}, // doesn't match + Tags: []string{"xss"}, // doesn't match + Severities: []severity.Severity{severity.Low}, // doesn't match + } + require.False(t, filter.Matches(metadata)) + }) +} + +func TestMatchesPath(t *testing.T) { + tests := []struct { + name string + path string + pattern string + expected bool + }{ + {"exact match", "/templates/cves/2021/test.yaml", "/templates/cves/2021/test.yaml", true}, + {"directory prefix", "/templates/cves/2021/test.yaml", "/templates/cves", true}, + {"directory with slash", "/templates/cves/2021/test.yaml", "/templates/cves/", true}, + {"no match", "/templates/cves/2021/test.yaml", "/templates/exploits", false}, + {"wildcard match", "/templates/cves/2021/test.yaml", "/templates/*/2021/*.yaml", true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := matchesPath(tt.path, tt.pattern) + require.Equal(t, tt.expected, result) + }) + } +} + +func TestMatchesID(t *testing.T) { + tests := []struct { + name string + id string + pattern string + expected bool + }{ + {"exact match", "CVE-2021-1234", "CVE-2021-1234", true}, + {"wildcard prefix", "CVE-2021-1234", "CVE-*", true}, + {"wildcard suffix", "CVE-2021-1234", "*-1234", true}, + {"wildcard middle", "CVE-2021-1234", "CVE-*-1234", true}, + {"no match", "CVE-2021-1234", "CVE-2022-*", false}, + {"partial no match", "CVE-2021-1234", "CVE-2021-12", false}, + {"case insensitive exact", "cve-2021-1234", "CVE-2021-1234", true}, + {"case insensitive wildcard", "CVE-2021-1234", "cve-*", true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := matchesID(tt.id, tt.pattern) + require.Equal(t, tt.expected, result) + }) + } +} + +func TestUnmarshalFilter(t *testing.T) { + filter, err := UnmarshalFilter( + []string{"author1", "author2"}, + []string{"tag1", "tag2"}, + []string{"exclude-tag"}, + []string{"include-tag"}, + []string{"id1", "id2*"}, + []string{"exclude-id*"}, + []string{"/include/path"}, + []string{"/exclude/path"}, + []string{"critical", "high"}, + []string{"info"}, + []string{"http", "dns"}, + []string{"file"}, + ) + + require.NoError(t, err) + require.NotNil(t, filter) + + require.Equal(t, []string{"author1", "author2"}, filter.Authors) + require.Equal(t, []string{"tag1", "tag2"}, filter.Tags) + require.Equal(t, []string{"exclude-tag"}, filter.ExcludeTags) + require.Equal(t, []string{"include-tag"}, filter.IncludeTags) + require.Equal(t, []string{"id1", "id2*"}, filter.IDs) + require.Equal(t, []string{"exclude-id*"}, filter.ExcludeIDs) + require.Equal(t, []string{"/include/path"}, filter.IncludeTemplates) + require.Equal(t, []string{"/exclude/path"}, filter.ExcludeTemplates) + + require.Len(t, filter.Severities, 2) + require.Contains(t, filter.Severities, severity.Critical) + require.Contains(t, filter.Severities, severity.High) + + require.Len(t, filter.ExcludeSeverities, 1) + require.Contains(t, filter.ExcludeSeverities, severity.Info) + + require.Len(t, filter.ProtocolTypes, 2) + require.Contains(t, filter.ProtocolTypes, types.HTTPProtocol) + require.Contains(t, filter.ProtocolTypes, types.DNSProtocol) + + require.Len(t, filter.ExcludeProtocolTypes, 1) + require.Contains(t, filter.ExcludeProtocolTypes, types.FileProtocol) +} + +func TestIndexFilter(t *testing.T) { + tmpDir := t.TempDir() + idx, err := NewIndex(tmpDir) + require.NoError(t, err) + + // Create test templates and metadata + templates := []struct { + id string + path string + authors []string + tags []string + severity string + protocol string + }{ + {"cve-2021-1", "/templates/cves/CVE-2021-1.yaml", []string{"pdteam"}, []string{"cve", "rce"}, "critical", "http"}, + {"cve-2021-2", "/templates/cves/CVE-2021-2.yaml", []string{"pdteam"}, []string{"cve", "xss"}, "high", "http"}, + {"exploit-1", "/templates/exploits/exploit-1.yaml", []string{"geeknik"}, []string{"exploit"}, "medium", "dns"}, + {"info-1", "/templates/info/info-1.yaml", []string{"author1"}, []string{"info"}, "info", "http"}, + } + + for _, tmpl := range templates { + tmpFile := filepath.Join(tmpDir, filepath.Base(tmpl.path)) + err := os.WriteFile(tmpFile, []byte("id: "+tmpl.id), 0644) + require.NoError(t, err) + + metadata := &Metadata{ + ID: tmpl.id, + FilePath: tmpFile, + Authors: tmpl.authors, + Tags: tmpl.tags, + Severity: tmpl.severity, + ProtocolType: tmpl.protocol, + } + idx.Set(tmpl.path, metadata) + } + + t.Run("No filter returns all", func(t *testing.T) { + results := idx.Filter(nil) + require.Len(t, results, 4) + }) + + t.Run("Filter by author", func(t *testing.T) { + filter := &Filter{Authors: []string{"pdteam"}} + results := idx.Filter(filter) + require.Len(t, results, 2) + }) + + t.Run("Filter by tag", func(t *testing.T) { + filter := &Filter{Tags: []string{"cve"}} + results := idx.Filter(filter) + require.Len(t, results, 2) + }) + + t.Run("Filter by severity", func(t *testing.T) { + filter := &Filter{Severities: []severity.Severity{severity.Critical}} + results := idx.Filter(filter) + require.Len(t, results, 1) + }) + + t.Run("Filter by protocol type", func(t *testing.T) { + filter := &Filter{ProtocolTypes: []types.ProtocolType{types.HTTPProtocol}} + results := idx.Filter(filter) + require.Len(t, results, 3) + }) + + t.Run("Exclude by severity", func(t *testing.T) { + filter := &Filter{ExcludeSeverities: []severity.Severity{severity.Info}} + results := idx.Filter(filter) + require.Len(t, results, 3) + }) + + t.Run("Exclude by tag", func(t *testing.T) { + filter := &Filter{ExcludeTags: []string{"info"}} + results := idx.Filter(filter) + require.Len(t, results, 3) + }) + + t.Run("Complex filter", func(t *testing.T) { + filter := &Filter{ + Tags: []string{"cve"}, + Severities: []severity.Severity{severity.Critical, severity.High}, + ExcludeSeverities: []severity.Severity{severity.Info}, + } + results := idx.Filter(filter) + require.Len(t, results, 2) + }) + + t.Run("Count with filter", func(t *testing.T) { + filter := &Filter{Tags: []string{"cve"}} + count := idx.Count(filter) + require.Equal(t, 2, count) + }) + + t.Run("Count without filter", func(t *testing.T) { + count := idx.Count(nil) + require.Equal(t, 4, count) + }) +} + +func TestIndexFilterFunc(t *testing.T) { + tmpDir := t.TempDir() + idx, err := NewIndex(tmpDir) + require.NoError(t, err) + + // Add test metadata + for i := 0; i < 5; i++ { + metadata := &Metadata{ + ID: "test-" + string(rune('a'+i)), + FilePath: "/tmp/test.yaml", + Severity: "high", + } + if i%2 == 0 { + metadata.Tags = []string{"even"} + } else { + metadata.Tags = []string{"odd"} + } + idx.Set("/tmp/test-"+string(rune('a'+i))+".yaml", metadata) + } + + t.Run("Custom filter function", func(t *testing.T) { + results := idx.FilterFunc(func(m *Metadata) bool { + return m.HasTag("even") + }) + require.Len(t, results, 3) // 0, 2, 4 + }) + + t.Run("Nil filter function returns all", func(t *testing.T) { + results := idx.FilterFunc(nil) + require.Len(t, results, 5) + }) +} + +func TestFilterString(t *testing.T) { + filter := &Filter{ + Authors: []string{"author1", "author2"}, + Tags: []string{"tag1"}, + Severities: []severity.Severity{severity.Critical, severity.High}, + ProtocolTypes: []types.ProtocolType{types.HTTPProtocol}, + } + + str := filter.String() + require.Contains(t, str, "authors=") + require.Contains(t, str, "tags=") + require.Contains(t, str, "severities=") + require.Contains(t, str, "types=") + + emptyFilter := &Filter{} + require.Equal(t, "filter=", emptyFilter.String()) +} diff --git a/pkg/catalog/index/index.go b/pkg/catalog/index/index.go new file mode 100644 index 0000000000..3024851b72 --- /dev/null +++ b/pkg/catalog/index/index.go @@ -0,0 +1,352 @@ +package index + +import ( + "encoding/gob" + "maps" + "os" + "path/filepath" + "sync" + + "github.com/maypok86/otter/v2" + "github.com/projectdiscovery/nuclei/v3/pkg/catalog/config" + "github.com/projectdiscovery/nuclei/v3/pkg/templates" + folderutil "github.com/projectdiscovery/utils/folder" +) + +const ( + // IndexFileName is the name of the persistent cache file. + IndexFileName = "index.gob" + + // IndexVersion is the schema version for cache invalidation on breaking + // changes. + IndexVersion = 1 + + // DefaultMaxSize is the default maximum number of templates to cache. + DefaultMaxSize = 50000 + + // DefaultMaxWeight is the default maximum weight of the cache. + DefaultMaxWeight = DefaultMaxSize * 800 // ~40MB assuming ~800B/entry +) + +// Index represents a cache for template metadata. +type Index struct { + cache *otter.Cache[string, *Metadata] + cacheFile string + mu sync.RWMutex + version int +} + +// cacheSnapshot represents the serialized cache structure. +type cacheSnapshot struct { + Version int `gob:"version"` + Data map[string]*Metadata `gob:"data"` +} + +// NewIndex creates a new template metadata cache with the given options. +func NewIndex(cacheDir string) (*Index, error) { + if cacheDir == "" { + cacheDir = folderutil.AppCacheDirOrDefault(".nuclei-cache", config.BinaryName) + } + + if err := os.MkdirAll(cacheDir, 0755); err != nil { + return nil, err + } + + cacheFile := filepath.Join(cacheDir, IndexFileName) + + // NOTE(dwisiswant0): Build cache with adaptive sizing based on memory cost. + opts := &otter.Options[string, *Metadata]{ + MaximumWeight: uint64(DefaultMaxWeight), + Weigher: func(key string, value *Metadata) uint32 { + if value == nil { + return uint32(len(key)) + } + + weight := len(key) + weight += len(value.ID) + weight += len(value.FilePath) + weight += 24 // ModTime is time.Time (24B) + weight += len(value.Name) + weight += len(value.Severity) + weight += len(value.ProtocolType) + weight += len(value.TemplateVerifier) + + for _, author := range value.Authors { + weight += len(author) + } + for _, tag := range value.Tags { + weight += len(tag) + } + + return uint32(weight) + }, + } + + cache, err := otter.New(opts) + if err != nil { + return nil, err + } + + c := &Index{ + cache: cache, + cacheFile: cacheFile, + version: IndexVersion, + } + + return c, nil +} + +// NewDefaultIndex creates a index with default settings in the default cache +// directory. +func NewDefaultIndex() (*Index, error) { + return NewIndex("") +} + +// Get retrieves metadata for a template path, validating freshness via mtime. +func (i *Index) Get(path string) (*Metadata, bool) { + i.mu.RLock() + defer i.mu.RUnlock() + + metadata, found := i.cache.GetIfPresent(path) + if !found { + return nil, false + } + + if !metadata.IsValid() { + go i.Delete(path) + + return nil, false + } + + return metadata, true +} + +// Set stores metadata for a template path. +// +// The caller is responsible for ensuring the metadata is valid and contains +// the correct checksum before calling this method. +// Use [SetFromTemplate] for automatic extraction and checksum computation. +// +// Returns the metadata and whether it was successfully cached (false if evicted). +func (i *Index) Set(path string, metadata *Metadata) (*Metadata, bool) { + i.mu.Lock() + defer i.mu.Unlock() + + return i.cache.Set(path, metadata) +} + +// SetFromTemplate extracts metadata from a parsed template and stores it. +// +// Returns the metadata and whether it was successfully cached. The metadata is +// always returned (even on checksum failure) for immediate filtering use. +// Returns false if the metadata was not cached (e.g., set, evicted). +func (i *Index) SetFromTemplate(path string, tpl *templates.Template) (*Metadata, bool) { + metadata := NewMetadataFromTemplate(path, tpl) + + info, err := os.Stat(path) + if err != nil { + return metadata, false + } + metadata.ModTime = info.ModTime() + + if i.cache == nil { + return metadata, false + } + + return i.Set(path, metadata) +} + +// Has checks if metadata exists for a path without validation. +func (i *Index) Has(path string) bool { + i.mu.RLock() + defer i.mu.RUnlock() + + _, found := i.cache.GetIfPresent(path) + + return found +} + +// Delete removes metadata for a path. +func (i *Index) Delete(path string) { + i.mu.Lock() + defer i.mu.Unlock() + + i.cache.Invalidate(path) +} + +// Size returns the number of cached entries. +func (i *Index) Size() int { + i.mu.RLock() + defer i.mu.RUnlock() + + return i.cache.EstimatedSize() +} + +// Clear removes all cached entries. +func (i *Index) Clear() { + i.mu.Lock() + defer i.mu.Unlock() + + i.cache.InvalidateAll() +} + +// Save persists the cache to disk using gob encoding. +func (i *Index) Save() error { + i.mu.RLock() + defer i.mu.RUnlock() + + snapshot := &cacheSnapshot{ + Version: i.version, + Data: make(map[string]*Metadata), + } + + maps.Insert(snapshot.Data, i.cache.All()) + + // NOTE(dwisiswant0): write to temp for atomic op. + tmpFile := i.cacheFile + ".tmp" + file, err := os.Create(tmpFile) + if err != nil { + return err + } + + encoder := gob.NewEncoder(file) + if err := encoder.Encode(snapshot); err != nil { + _ = file.Close() + _ = os.Remove(tmpFile) + + return err + } + + if err := file.Close(); err != nil { + _ = os.Remove(tmpFile) + + return err + } + + if err := os.Rename(tmpFile, i.cacheFile); err != nil { + _ = os.Remove(tmpFile) + + return err + } + + return nil +} + +// Load loads the cache from disk using gob decoding. +func (i *Index) Load() error { + file, err := os.Open(i.cacheFile) + if err != nil { + if os.IsNotExist(err) { + return nil + } + + return err + } + defer func() { _ = file.Close() }() + + var snapshot cacheSnapshot + + decoder := gob.NewDecoder(file) + if err := decoder.Decode(&snapshot); err != nil { + _ = file.Close() + _ = os.Remove(i.cacheFile) + + return nil + } + + if snapshot.Version != i.version { + _ = file.Close() + _ = os.Remove(i.cacheFile) + + return nil + } + + i.mu.Lock() + defer i.mu.Unlock() + + for key, value := range snapshot.Data { + i.cache.Set(key, value) + } + + return nil +} + +// Filter returns all template paths that match the given filter criteria. +func (i *Index) Filter(filter *Filter) []string { + if filter == nil || filter.IsEmpty() { + return i.All() + } + + i.mu.RLock() + defer i.mu.RUnlock() + + var matched []string + for path, metadata := range i.cache.All() { + if filter.Matches(metadata) { + matched = append(matched, path) + } + } + + return matched +} + +// FilterFunc returns all template paths that match the given filter function. +func (i *Index) FilterFunc(fn FilterFunc) []string { + if fn == nil { + return i.All() + } + + i.mu.RLock() + defer i.mu.RUnlock() + + var matched []string + for path, metadata := range i.cache.All() { + if fn(metadata) { + matched = append(matched, path) + } + } + + return matched +} + +// All returns all template paths in the index. +func (i *Index) All() []string { + i.mu.RLock() + defer i.mu.RUnlock() + + paths := make([]string, 0, i.cache.EstimatedSize()) + for path := range i.cache.All() { + paths = append(paths, path) + } + + return paths +} + +// GetAll returns all metadata entries in the index. +func (i *Index) GetAll() map[string]*Metadata { + i.mu.RLock() + defer i.mu.RUnlock() + + result := maps.Collect(i.cache.All()) + + return result +} + +// Count returns the number of templates matching the filter. +func (i *Index) Count(filter *Filter) int { + if filter == nil || filter.IsEmpty() { + return i.Size() + } + + i.mu.RLock() + defer i.mu.RUnlock() + + count := 0 + for _, metadata := range i.cache.All() { + if filter.Matches(metadata) { + count++ + } + } + + return count +} diff --git a/pkg/catalog/index/index_test.go b/pkg/catalog/index/index_test.go new file mode 100644 index 0000000000..de76dd7592 --- /dev/null +++ b/pkg/catalog/index/index_test.go @@ -0,0 +1,737 @@ +package index + +import ( + "fmt" + "os" + "path/filepath" + "testing" + "time" + + "github.com/projectdiscovery/nuclei/v3/pkg/model" + "github.com/projectdiscovery/nuclei/v3/pkg/model/types/severity" + "github.com/projectdiscovery/nuclei/v3/pkg/model/types/stringslice" + "github.com/projectdiscovery/nuclei/v3/pkg/protocols/code" + "github.com/projectdiscovery/nuclei/v3/pkg/protocols/headless" + "github.com/projectdiscovery/nuclei/v3/pkg/protocols/http" + "github.com/projectdiscovery/nuclei/v3/pkg/templates" + "github.com/projectdiscovery/nuclei/v3/pkg/templates/types" + "github.com/stretchr/testify/require" +) + +func TestNewIndex(t *testing.T) { + t.Run("with custom directory", func(t *testing.T) { + tmpDir := t.TempDir() + cache, err := NewIndex(tmpDir) + require.NoError(t, err, "Failed to create cache with custom directory") + require.NotNil(t, cache, "Cache should not be nil") + require.Equal(t, filepath.Join(tmpDir, IndexFileName), cache.cacheFile) + require.Equal(t, IndexVersion, cache.version) + }) + + t.Run("with default directory", func(t *testing.T) { + cache, err := NewDefaultIndex() + require.NoError(t, err, "Failed to create cache with default directory") + require.NotNil(t, cache, "Cache should not be nil") + }) +} + +func TestCacheBasicOperations(t *testing.T) { + tmpDir := t.TempDir() + cache, err := NewIndex(tmpDir) + require.NoError(t, err) + + metadata := &Metadata{ + ID: "concurrent-test", + FilePath: "/tmp/concurrent.yaml", + } + + t.Run("Set and Has", func(t *testing.T) { + cache.Set(metadata.FilePath, metadata) + require.Equal(t, 1, cache.Size(), "Cache size should be 1 after Set") + require.True(t, cache.Has(metadata.FilePath), "Cache should contain the path after Set") + require.False(t, cache.Has("/nonexistent"), "Cache should not contain nonexistent path") + }) + + t.Run("Get with validation", func(t *testing.T) { + // Get should fail validation for nonexistent file + retrieved, found := cache.Get(metadata.FilePath) + require.False(t, found, "Get should fail validation for nonexistent file") + require.Nil(t, retrieved, "Retrieved metadata should be nil for invalid entry") + }) + + t.Run("Delete", func(t *testing.T) { + cache.Set(metadata.FilePath, metadata) + require.True(t, cache.Has(metadata.FilePath), "Cache should contain path before Delete") + + cache.Delete(metadata.FilePath) + require.False(t, cache.Has(metadata.FilePath), "Cache should not contain path after Delete") + }) + + t.Run("Clear", func(t *testing.T) { + cache.Set(metadata.FilePath, metadata) + cache.Set("/tmp/test2.yaml", &Metadata{ID: "test2", FilePath: "/tmp/test2.yaml"}) + require.True(t, cache.Size() > 0, "Cache should have entries before Clear") + + cache.Clear() + require.Equal(t, 0, cache.Size(), "Cache should be empty after Clear") + }) +} + +func TestCachePersistence(t *testing.T) { + tmpDir := t.TempDir() + + metadata1 := &Metadata{ + ID: "persist-test-1", + FilePath: "/tmp/persist1.yaml", + Name: "Persistence Test 1", + Authors: []string{"tester"}, + Tags: []string{"test"}, + Severity: "medium", + ProtocolType: "dns", + } + + metadata2 := &Metadata{ + ID: "persist-test-2", + FilePath: "/tmp/persist2.yaml", + Name: "Persistence Test 2", + Authors: []string{"tester2"}, + Tags: []string{"cve"}, + Severity: "critical", + ProtocolType: "http", + } + + t.Run("Save and Load", func(t *testing.T) { + // Create cache and add entries + cache1, err := NewIndex(tmpDir) + require.NoError(t, err) + + cache1.Set(metadata1.FilePath, metadata1) + cache1.Set(metadata2.FilePath, metadata2) + require.Equal(t, 2, cache1.Size()) + + // Save to disk + err = cache1.Save() + require.NoError(t, err, "Failed to save cache") + + // Verify cache file exists + cacheFile := filepath.Join(tmpDir, IndexFileName) + stat, err := os.Stat(cacheFile) + require.NoError(t, err, "Cache file should exist") + require.Greater(t, stat.Size(), int64(0), "Cache file should not be empty") + + // Create new cache and load + cache2, err := NewIndex(tmpDir) + require.NoError(t, err) + require.Equal(t, 0, cache2.Size(), "New cache should be empty before Load") + + err = cache2.Load() + require.NoError(t, err, "Failed to load cache") + + // Verify data was loaded + require.Equal(t, 2, cache2.Size(), "Loaded cache should have 2 entries") + require.True(t, cache2.Has(metadata1.FilePath), "Loaded cache should contain first entry") + require.True(t, cache2.Has(metadata2.FilePath), "Loaded cache should contain second entry") + }) + + t.Run("Load non-existent cache", func(t *testing.T) { + emptyDir := t.TempDir() + cache, err := NewIndex(emptyDir) + require.NoError(t, err) + + // Loading non-existent cache should not error + err = cache.Load() + require.NoError(t, err, "Loading non-existent cache should not error") + require.Equal(t, 0, cache.Size(), "Cache should be empty after loading non-existent file") + }) + + t.Run("Atomic save", func(t *testing.T) { + cache, err := NewIndex(tmpDir) + require.NoError(t, err) + + cache.Set(metadata1.FilePath, metadata1) + err = cache.Save() + require.NoError(t, err) + + // Verify no .tmp file left behind + tmpFile := filepath.Join(tmpDir, IndexFileName+".tmp") + _, err = os.Stat(tmpFile) + require.True(t, os.IsNotExist(err), "Temporary file should not exist after save") + + // Verify actual cache file exists + cacheFile := filepath.Join(tmpDir, IndexFileName) + _, err = os.Stat(cacheFile) + require.NoError(t, err, "Cache file should exist") + }) +} + +func TestIndexVersionMismatch(t *testing.T) { + tmpDir := t.TempDir() + + // Create cache with current version + cache1, err := NewIndex(tmpDir) + require.NoError(t, err) + + metadata := &Metadata{ + ID: "version-test", + FilePath: "/tmp/version.yaml", + } + cache1.Set(metadata.FilePath, metadata) + + // Save with current version + err = cache1.Save() + require.NoError(t, err) + + // Manually modify version and save again + cache1.version = 999 + err = cache1.Save() + require.NoError(t, err) + + // Try to load with different version + cache2, err := NewIndex(tmpDir) + require.NoError(t, err) + + // Load should succeed but cache should be empty (version mismatch) + err = cache2.Load() + require.NoError(t, err, "Load should not error on version mismatch") + require.Equal(t, 0, cache2.Size(), "Cache should be empty after version mismatch") +} + +func TestCacheCorruptedFile(t *testing.T) { + tmpDir := t.TempDir() + cacheFile := filepath.Join(tmpDir, IndexFileName) + + // Create corrupted cache file + err := os.WriteFile(cacheFile, []byte("corrupted data that is not valid gob"), 0644) + require.NoError(t, err) + + // Try to load corrupted cache + cache, err := NewIndex(tmpDir) + require.NoError(t, err) + + err = cache.Load() + require.NoError(t, err, "Load should not error on corrupted cache") + require.Equal(t, 0, cache.Size(), "Cache should be empty after loading corrupted file") + + // Corrupted file should be removed + _, err = os.Stat(cacheFile) + require.True(t, os.IsNotExist(err), "Corrupted cache file should be removed") +} + +func TestMetadataValidation(t *testing.T) { + tmpDir := t.TempDir() + tmpFile := filepath.Join(tmpDir, "test.yaml") + + t.Run("Valid metadata", func(t *testing.T) { + // Create a test file + err := os.WriteFile(tmpFile, []byte("id: test\ninfo:\n name: Test"), 0644) + require.NoError(t, err) + + info, err := os.Stat(tmpFile) + require.NoError(t, err) + + // Create metadata with correct checksum + metadata := &Metadata{ + ID: "test", + FilePath: tmpFile, + ModTime: info.ModTime(), + } + + // Should be valid + require.True(t, metadata.IsValid(), "Metadata should be valid for unchanged file") + }) + + t.Run("Invalid metadata after file modification", func(t *testing.T) { + // Create the test file first to ensure it exists in this subtest + err := os.WriteFile(tmpFile, []byte("id: test\ninfo:\n name: Test"), 0644) + require.NoError(t, err) + + // Set file ModTime to past to ensure modification is detectable + oldTime := time.Now().Add(-2 * time.Second) + err = os.Chtimes(tmpFile, oldTime, oldTime) + require.NoError(t, err) + + info, err := os.Stat(tmpFile) + require.NoError(t, err) + + metadata := &Metadata{ + ID: "test", + FilePath: tmpFile, + ModTime: info.ModTime(), + } + + // Modify file + err = os.WriteFile(tmpFile, []byte("id: test\ninfo:\n name: Modified"), 0644) + require.NoError(t, err) + + // Should now be invalid + require.False(t, metadata.IsValid(), "Metadata should be invalid after file modification") + }) + + t.Run("Invalid metadata for deleted file", func(t *testing.T) { + // Create the test file first to ensure it exists in this subtest + err := os.WriteFile(tmpFile, []byte("id: test\ninfo:\n name: Test"), 0644) + require.NoError(t, err) + + info, err := os.Stat(tmpFile) + require.NoError(t, err) + + metadata := &Metadata{ + ID: "test", + FilePath: tmpFile, + ModTime: info.ModTime(), + } + + // Delete file + err = os.Remove(tmpFile) + require.NoError(t, err) + + // Should be invalid + require.False(t, metadata.IsValid(), "Metadata should be invalid for deleted file") + }) +} + +func TestSetFromTemplate(t *testing.T) { + tmpDir := t.TempDir() + tmpFile := filepath.Join(tmpDir, "extract.yaml") + + // Create a test file + err := os.WriteFile(tmpFile, []byte("id: extract-test"), 0644) + require.NoError(t, err) + + cache, err := NewIndex(tmpDir) + require.NoError(t, err) + + t.Run("Basic metadata extraction", func(t *testing.T) { + template := &templates.Template{ + ID: "extract-test", + Info: model.Info{ + Name: "Extract Test Template", + Authors: stringslice.StringSlice{Value: "author1,author2"}, + Tags: stringslice.StringSlice{Value: "tag1,tag2"}, + Description: "Test description", + SeverityHolder: severity.Holder{ + Severity: severity.High, + }, + }, + SelfContained: true, + Verified: true, + TemplateVerifier: "test-verifier", + } + + metadata, ok := cache.SetFromTemplate(tmpFile, template) + require.True(t, ok, "Failed to set metadata from template") + require.NotNil(t, metadata, "Metadata should not be nil") + + // Verify core fields + require.Equal(t, "extract-test", metadata.ID) + require.Equal(t, tmpFile, metadata.FilePath) + + // Verify Info fields + require.Equal(t, "Extract Test Template", metadata.Name) + require.Equal(t, []string{"author1,author2"}, metadata.Authors) + require.Equal(t, []string{"tag1,tag2"}, metadata.Tags) + require.Equal(t, "high", metadata.Severity) + + // Verify flags + require.True(t, metadata.Verified) + require.Equal(t, "test-verifier", metadata.TemplateVerifier) + }) + + t.Run("HTTP protocol detection", func(t *testing.T) { + // Create a separate test file for this test + httpFile := filepath.Join(tmpDir, "http-test.yaml") + err := os.WriteFile(httpFile, []byte("id: http-test"), 0644) + require.NoError(t, err) + + template := &templates.Template{ + ID: "http-test", + Info: model.Info{ + Name: "HTTP Test", + Authors: stringslice.StringSlice{Value: "tester"}, + SeverityHolder: severity.Holder{ + Severity: severity.Medium, + }, + }, + RequestsHTTP: []*http.Request{{Method: http.HTTPMethodTypeHolder{MethodType: http.HTTPGet}}}, + } + + metadata, ok := cache.SetFromTemplate(httpFile, template) + require.True(t, ok) + require.NotNil(t, metadata) + require.Equal(t, "http", metadata.ProtocolType) + }) + + t.Run("Extract with missing file", func(t *testing.T) { + template := &templates.Template{ + ID: "missing-test", + Info: model.Info{ + Name: "Missing File Test", + Authors: stringslice.StringSlice{Value: "tester"}, + SeverityHolder: severity.Holder{ + Severity: severity.Low, + }, + }, + } + + metadata, ok := cache.SetFromTemplate("/nonexistent/file.yaml", template) + require.False(t, ok, "Should return false for nonexistent file") + require.NotNil(t, metadata, "Metadata should still be returned") + }) +} + +func TestMetadataMatchingHelpers(t *testing.T) { + metadata := &Metadata{ + Tags: []string{"cve", "rce", "apache"}, + Authors: []string{"pdteam", "geeknik"}, + Severity: "critical", + ProtocolType: "http", + } + + t.Run("HasTag", func(t *testing.T) { + require.True(t, metadata.HasTag("cve")) + require.True(t, metadata.HasTag("rce")) + require.True(t, metadata.HasTag("apache")) + require.False(t, metadata.HasTag("xxe")) + require.False(t, metadata.HasTag("")) + }) + + t.Run("HasAuthor", func(t *testing.T) { + require.True(t, metadata.HasAuthor("pdteam")) + require.True(t, metadata.HasAuthor("geeknik")) + require.False(t, metadata.HasAuthor("unknown")) + require.False(t, metadata.HasAuthor("")) + }) + + t.Run("MatchesSeverity", func(t *testing.T) { + require.True(t, metadata.MatchesSeverity(severity.Critical)) + require.False(t, metadata.MatchesSeverity(severity.High)) + require.False(t, metadata.MatchesSeverity(severity.Medium)) + require.False(t, metadata.MatchesSeverity(severity.Low)) + require.False(t, metadata.MatchesSeverity(severity.Info)) + }) + + t.Run("MatchesProtocol", func(t *testing.T) { + require.True(t, metadata.MatchesProtocol(types.HTTPProtocol)) + require.False(t, metadata.MatchesProtocol(types.DNSProtocol)) + require.False(t, metadata.MatchesProtocol(types.FileProtocol)) + require.False(t, metadata.MatchesProtocol(types.NetworkProtocol)) + }) + + t.Run("Empty metadata", func(t *testing.T) { + emptyMetadata := &Metadata{} + require.False(t, emptyMetadata.HasTag("any")) + require.False(t, emptyMetadata.HasAuthor("any")) + }) +} + +func TestCacheConcurrency(t *testing.T) { + tmpDir := t.TempDir() + cache, err := NewIndex(tmpDir) + require.NoError(t, err) + + // Test concurrent writes + t.Run("Concurrent Set", func(t *testing.T) { + done := make(chan bool) + for i := 0; i < 10; i++ { + go func(id int) { + metadata := &Metadata{ + ID: string(rune('a' + id)), + FilePath: filepath.Join("/tmp", string(rune('a'+id))+".yaml"), + } + cache.Set(metadata.FilePath, metadata) + done <- true + }(i) + } + + // Wait for all goroutines + for i := 0; i < 10; i++ { + <-done + } + + require.Equal(t, 10, cache.Size(), "All concurrent writes should succeed") + }) + + // Test concurrent reads + t.Run("Concurrent Has", func(t *testing.T) { + metadata := &Metadata{ + ID: "concurrent-test", + FilePath: "/tmp/concurrent.yaml", + } + cache.Set(metadata.FilePath, metadata) + + done := make(chan bool) + for i := 0; i < 20; i++ { + go func() { + _ = cache.Has(metadata.FilePath) + done <- true + }() + } + + // Wait for all goroutines + for i := 0; i < 20; i++ { + <-done + } + }) +} + +func TestCacheSize(t *testing.T) { + tmpDir := t.TempDir() + cache, err := NewIndex(tmpDir) + require.NoError(t, err) + + require.Equal(t, 0, cache.Size(), "New cache should have size 0") + + // Add entries + for i := 0; i < 5; i++ { + metadata := &Metadata{ + ID: string(rune('a' + i)), + FilePath: filepath.Join("/tmp", string(rune('a'+i))+".yaml"), + } + cache.Set(metadata.FilePath, metadata) + } + + require.Equal(t, 5, cache.Size(), "Cache should have size 5 after adding 5 entries") + + // Delete entries + cache.Delete(filepath.Join("/tmp", "a.yaml")) + cache.Delete(filepath.Join("/tmp", "b.yaml")) + + require.Equal(t, 3, cache.Size(), "Cache should have size 3 after deleting 2 entries") + + // Clear cache + cache.Clear() + require.Equal(t, 0, cache.Size(), "Cache should have size 0 after Clear") +} + +func TestCacheGetWithValidFile(t *testing.T) { + tmpDir := t.TempDir() + cache, err := NewIndex(tmpDir) + require.NoError(t, err) + + // Create a real file for testing validation + tmpFile := filepath.Join(tmpDir, "test.yaml") + err = os.WriteFile(tmpFile, []byte("id: test"), 0644) + require.NoError(t, err) + + info, err := os.Stat(tmpFile) + require.NoError(t, err) + + metadata := &Metadata{ + ID: "test", + FilePath: tmpFile, + ModTime: info.ModTime(), + Name: "Test Template", + } + + // Set and get should work with valid file + cache.Set(metadata.FilePath, metadata) + retrieved, found := cache.Get(metadata.FilePath) + require.True(t, found, "Should find entry with valid file") + require.NotNil(t, retrieved, "Retrieved metadata should not be nil") + require.Equal(t, metadata.ID, retrieved.ID) +} + +func TestCacheSaveErrorHandling(t *testing.T) { + tmpDir := t.TempDir() + cache, err := NewIndex(tmpDir) + require.NoError(t, err) + + metadata := &Metadata{ + ID: "test", + FilePath: filepath.Join("/tmp", "test.yaml"), + } + cache.Set(metadata.FilePath, metadata) + + // Create a directory where the temp file would be created to force an error + // The Save method creates a file at cacheFile + ".tmp" + conflictPath := filepath.Join(tmpDir, IndexFileName+".tmp") + err = os.Mkdir(conflictPath, 0755) + require.NoError(t, err) + + err = cache.Save() + require.Error(t, err, "Save should fail when temp file cannot be created") +} + +func TestNewCacheWithInvalidDirectory(t *testing.T) { + // Try to create cache in a file path (should fail) + tmpFile := filepath.Join(t.TempDir(), "file.txt") + err := os.WriteFile(tmpFile, []byte("test"), 0644) + require.NoError(t, err) + + cache, err := NewIndex(tmpFile) + require.Error(t, err, "NewCache should fail when path is a file") + require.Nil(t, cache, "Cache should be nil on error") +} + +func TestCacheLoadCorruptedRemoval(t *testing.T) { + tmpDir := t.TempDir() + cacheFile := filepath.Join(tmpDir, IndexFileName) + + // Create corrupted cache file with invalid gob data + err := os.WriteFile(cacheFile, []byte("this is not valid gob encoding at all!"), 0644) + require.NoError(t, err) + + // Verify file exists before Load + _, err = os.Stat(cacheFile) + require.NoError(t, err, "Corrupted file should exist") + + // Load should not error but should remove corrupted file + cache, err := NewIndex(tmpDir) + require.NoError(t, err) + + err = cache.Load() + require.NoError(t, err, "Load should not return error for corrupted file") + + // Verify corrupted file was removed + _, err = os.Stat(cacheFile) + require.True(t, os.IsNotExist(err), "Corrupted file should be removed") + require.Equal(t, 0, cache.Size(), "Cache should be empty after loading corrupted file") +} + +func TestMetadataExtractionWithNilClassification(t *testing.T) { + tmpDir := t.TempDir() + tmpFile := filepath.Join(tmpDir, "test.yaml") + err := os.WriteFile(tmpFile, []byte("id: test"), 0644) + require.NoError(t, err) + + template := &templates.Template{ + ID: "nil-classification", + Info: model.Info{ + Name: "Template without classification", + Authors: stringslice.StringSlice{Value: "tester"}, + SeverityHolder: severity.Holder{ + Severity: severity.Medium, + }, + Classification: nil, // Explicitly nil + }, + } + + cache, err := NewIndex(tmpDir) + require.NoError(t, err) + + metadata, ok := cache.SetFromTemplate(tmpFile, template) + require.True(t, ok) + require.NotNil(t, metadata) +} + +func TestCachePersistenceWithLargeDataset(t *testing.T) { + tmpDir := t.TempDir() + cache, err := NewIndex(tmpDir) + require.NoError(t, err) + + // Add 100 entries to test bulk operations + for i := 0; i < 100; i++ { + metadata := &Metadata{ + ID: fmt.Sprintf("template-%d", i), + FilePath: filepath.Join("/tmp", fmt.Sprintf("template-%d.yaml", i)), + Name: fmt.Sprintf("Template %d", i), + Authors: []string{fmt.Sprintf("author%d", i)}, + Tags: []string{"tag1", "tag2", "tag3"}, + Severity: "high", + } + cache.Set(metadata.FilePath, metadata) + } + + require.Equal(t, 100, cache.Size(), "Cache should contain 100 entries") + + // Save to disk + err = cache.Save() + require.NoError(t, err) + + // Load into new cache + cache2, err := NewIndex(tmpDir) + require.NoError(t, err) + err = cache2.Load() + require.NoError(t, err) + + require.Equal(t, 100, cache2.Size(), "Loaded cache should contain 100 entries") + + // Verify a sample entry + found := cache2.Has(filepath.Join("/tmp", "template-50.yaml")) + require.True(t, found, "Should find sample entry") +} + +func TestMetadataHelperMethods(t *testing.T) { + metadata := &Metadata{ + ID: "helper-test", + Tags: []string{}, + Authors: []string{}, + Severity: "", + ProtocolType: "", + } + + t.Run("Empty tags", func(t *testing.T) { + require.False(t, metadata.HasTag("anytag")) + }) + + t.Run("Empty authors", func(t *testing.T) { + require.False(t, metadata.HasAuthor("anyauthor")) + }) + + t.Run("Empty severity", func(t *testing.T) { + require.False(t, metadata.MatchesSeverity(severity.Critical)) + }) + + t.Run("Empty protocol", func(t *testing.T) { + require.False(t, metadata.MatchesProtocol(types.HTTPProtocol)) + }) +} + +func TestMultipleProtocolsDetection(t *testing.T) { + tmpDir := t.TempDir() + tmpFile := filepath.Join(tmpDir, "multi.yaml") + err := os.WriteFile(tmpFile, []byte("id: multi"), 0644) + require.NoError(t, err) + + // Template with multiple protocol types + template := &templates.Template{ + ID: "multi-protocol", + Info: model.Info{ + Name: "Multi Protocol Template", + Authors: stringslice.StringSlice{Value: "tester"}, + SeverityHolder: severity.Holder{ + Severity: severity.High, + }, + }, + RequestsHTTP: []*http.Request{{Method: http.HTTPMethodTypeHolder{MethodType: http.HTTPGet}}}, + RequestsHeadless: []*headless.Request{{}}, + RequestsCode: []*code.Request{{}}, + } + + cache, err := NewIndex(tmpDir) + require.NoError(t, err) + + metadata, ok := cache.SetFromTemplate(tmpFile, template) + require.True(t, ok) + require.NotNil(t, metadata) + require.Equal(t, "http", metadata.ProtocolType, "Primary protocol should be http") +} + +func TestNewMetadataFromTemplate(t *testing.T) { + tmpl := &templates.Template{ + ID: "test-template", + Info: model.Info{ + Name: "Test Template", + Authors: stringslice.StringSlice{Value: []string{"author"}}, + Tags: stringslice.StringSlice{Value: []string{"tag"}}, + SeverityHolder: severity.Holder{ + Severity: severity.Low, + }, + }, + Verified: true, + TemplateVerifier: "verifier", + } + + path := "/tmp/test.yaml" + metadata := NewMetadataFromTemplate(path, tmpl) + + require.Equal(t, tmpl.ID, metadata.ID) + require.Equal(t, path, metadata.FilePath) + require.Equal(t, tmpl.Info.Name, metadata.Name) + require.Equal(t, tmpl.Info.Authors.ToSlice(), metadata.Authors) + require.Equal(t, tmpl.Info.Tags.ToSlice(), metadata.Tags) + require.Equal(t, tmpl.Info.SeverityHolder.Severity.String(), metadata.Severity) + require.Equal(t, tmpl.Type().String(), metadata.ProtocolType) + require.Equal(t, tmpl.Verified, metadata.Verified) + require.Equal(t, tmpl.TemplateVerifier, metadata.TemplateVerifier) +} diff --git a/pkg/catalog/index/metadata.go b/pkg/catalog/index/metadata.go new file mode 100644 index 0000000000..013ab439c1 --- /dev/null +++ b/pkg/catalog/index/metadata.go @@ -0,0 +1,104 @@ +package index + +import ( + "os" + "slices" + "time" + + "github.com/projectdiscovery/nuclei/v3/pkg/model/types/severity" + "github.com/projectdiscovery/nuclei/v3/pkg/templates" + "github.com/projectdiscovery/nuclei/v3/pkg/templates/types" +) + +// Metadata contains lightweight metadata extracted from a template. +type Metadata struct { + // ID is the unique identifier of the template. + ID string `gob:"id"` + + // FilePath is the path to the template file. + FilePath string `gob:"file_path"` + + // ModTime is the modification time of the template file. + ModTime time.Time `gob:"mod_time"` + + // Name is the name of the template. + Name string `gob:"name"` + + // Authors are the authors of the template. + Authors []string `gob:"authors"` + + // Tags are the tags associated with the template. + Tags []string `gob:"tags"` + + // Severity is the severity level of the template. + Severity string `gob:"severity"` + + // ProtocolType is the primary protocol type of the template. + ProtocolType string `gob:"protocol_type"` + + // Verified indicates whether the template is verified. + Verified bool `gob:"verified"` + + // TemplateVerifier is the verifier used for the template. + TemplateVerifier string `gob:"verifier,omitempty"` + + // NOTE(dwisiswant0): Consider adding more fields here in the future to + // enhance filtering caps w/o loading full templates, such as: + // `has_{code,headless,file}` to indicate presence of protocol-based + // requests, and/or classification fields (CVE, CWE, CVSS, EPSS), if needed. + // + // For maintainers: when adding new fields, don't forget to update the + // Weigher logic in [NewIndex] to account for the new fields in cache weight + // calculation, because it affects cache eviction behavior. Also, consider + // the impact on existing cached data and whether a [IndexVersion] bump is + // needed. +} + +// NewMetadataFromTemplate creates a new metadata object from a template. +func NewMetadataFromTemplate(path string, tpl *templates.Template) *Metadata { + return &Metadata{ + ID: tpl.ID, + FilePath: path, + + Name: tpl.Info.Name, + Authors: tpl.Info.Authors.ToSlice(), + Tags: tpl.Info.Tags.ToSlice(), + Severity: tpl.Info.SeverityHolder.Severity.String(), + + ProtocolType: tpl.Type().String(), + + Verified: tpl.Verified, + TemplateVerifier: tpl.TemplateVerifier, + } +} + +// IsValid checks if the cached metadata is still valid by comparing the file +// modification time. +func (m *Metadata) IsValid() bool { + info, err := os.Stat(m.FilePath) + if err != nil { + return false + } + + return m.ModTime.Equal(info.ModTime()) +} + +// MatchesSeverity checks if the metadata matches the given severity. +func (m *Metadata) MatchesSeverity(sev severity.Severity) bool { + return m.Severity == sev.String() +} + +// MatchesProtocol checks if the metadata matches the given protocol type. +func (m *Metadata) MatchesProtocol(protocolType types.ProtocolType) bool { + return m.ProtocolType == protocolType.String() +} + +// HasTag checks if the metadata contains the given tag. +func (m *Metadata) HasTag(tag string) bool { + return slices.Contains(m.Tags, tag) +} + +// HasAuthor checks if the metadata contains the given author. +func (m *Metadata) HasAuthor(author string) bool { + return slices.Contains(m.Authors, author) +} diff --git a/pkg/catalog/loader/loader.go b/pkg/catalog/loader/loader.go index a68153d875..00bc970b9b 100644 --- a/pkg/catalog/loader/loader.go +++ b/pkg/catalog/loader/loader.go @@ -14,7 +14,7 @@ import ( "github.com/projectdiscovery/gologger" "github.com/projectdiscovery/nuclei/v3/pkg/catalog" "github.com/projectdiscovery/nuclei/v3/pkg/catalog/config" - "github.com/projectdiscovery/nuclei/v3/pkg/catalog/loader/filter" + "github.com/projectdiscovery/nuclei/v3/pkg/catalog/index" "github.com/projectdiscovery/nuclei/v3/pkg/keys" "github.com/projectdiscovery/nuclei/v3/pkg/model/types/severity" "github.com/projectdiscovery/nuclei/v3/pkg/protocols" @@ -77,7 +77,6 @@ type Config struct { type Store struct { id string // id of the store (optional) tagFilter *templates.TagFilter - pathFilter *filter.PathFilter config *Config finalTemplates []string finalWorkflows []string @@ -92,6 +91,16 @@ type Store struct { // parserCacheOnce is used to cache the parser cache result parserCacheOnce func() *templates.Cache + // metadataIndex is the template metadata cache + metadataIndex *index.Index + + // indexFilter is the cached filter for metadata matching + indexFilter *index.Filter + + // saveTemplatesIndexOnce is used to ensure we only save the metadata index + // once + saveMetadataIndexOnce func() + // NotFoundCallback is called for each not found template // This overrides error handling for not found templates NotFoundCallback func(template string) bool @@ -129,17 +138,10 @@ func NewConfig(options *types.Options, catalog catalog.Catalog, executerOpts *pr // New creates a new template store based on provided configuration func New(cfg *Config) (*Store, error) { + // tagFilter only for IncludeConditions (advanced filtering). + // All other filtering (tags, authors, severities, IDs, protocols, paths) is + // handled by [index.Filter]. tagFilter, err := templates.NewTagFilter(&templates.TagFilterConfig{ - Tags: cfg.Tags, - ExcludeTags: cfg.ExcludeTags, - Authors: cfg.Authors, - Severities: cfg.Severities, - ExcludeSeverities: cfg.ExcludeSeverities, - IncludeTags: cfg.IncludeTags, - IncludeIds: cfg.IncludeIds, - ExcludeIds: cfg.ExcludeIds, - Protocols: cfg.Protocols, - ExcludeProtocols: cfg.ExcludeProtocols, IncludeConditions: cfg.IncludeConditions, }) if err != nil { @@ -147,13 +149,9 @@ func New(cfg *Config) (*Store, error) { } store := &Store{ - id: cfg.StoreId, - config: cfg, - tagFilter: tagFilter, - pathFilter: filter.NewPathFilter(&filter.PathFilterConfig{ - IncludedTemplates: cfg.IncludeTemplates, - ExcludedTemplates: cfg.ExcludeTemplates, - }, cfg.Catalog), + id: cfg.StoreId, + config: cfg, + tagFilter: tagFilter, finalTemplates: cfg.Templates, finalWorkflows: cfg.Workflows, logger: cfg.Logger, @@ -171,6 +169,21 @@ func New(cfg *Config) (*Store, error) { return nil }) + // Initialize metadata index and filter (load from disk & cache for reuse) + store.metadataIndex = store.loadTemplatesIndex() + store.indexFilter = store.buildIndexFilter() + store.saveMetadataIndexOnce = sync.OnceFunc(func() { + if store.metadataIndex == nil { + return + } + + if err := store.metadataIndex.Save(); err != nil { + store.logger.Warning().Msgf("Could not save metadata cache: %v", err) + } else { + store.logger.Verbose().Msgf("Saved %d templates to metadata cache", store.metadataIndex.Size()) + } + }) + // Do a check to see if we have URLs in templates flag, if so // we need to processs them separately and remove them from the initial list var templatesFinal []string @@ -302,17 +315,102 @@ func init() { templateIDPathMap = make(map[string]string) } +// buildIndexFilter creates an [index.Filter] from the store configuration. +// This filter handles all basic filtering (paths, tags, authors, severities, +// IDs, protocols). Advanced IncludeConditions filtering is handled separately +// by tagFilter. +func (store *Store) buildIndexFilter() *index.Filter { + includeTemplates, _ := store.config.Catalog.GetTemplatesPath(store.config.IncludeTemplates) + excludeTemplates, _ := store.config.Catalog.GetTemplatesPath(store.config.ExcludeTemplates) + + return &index.Filter{ + Authors: store.config.Authors, + Tags: store.config.Tags, + ExcludeTags: store.config.ExcludeTags, + IncludeTags: store.config.IncludeTags, + IDs: store.config.IncludeIds, + ExcludeIDs: store.config.ExcludeIds, + IncludeTemplates: includeTemplates, + ExcludeTemplates: excludeTemplates, + Severities: []severity.Severity(store.config.Severities), + ExcludeSeverities: []severity.Severity(store.config.ExcludeSeverities), + ProtocolTypes: []templateTypes.ProtocolType(store.config.Protocols), + ExcludeProtocolTypes: []templateTypes.ProtocolType(store.config.ExcludeProtocols), + } +} + +func (store *Store) loadTemplatesIndex() *index.Index { + var metadataIdx *index.Index + + idx, err := index.NewDefaultIndex() + if err != nil { + store.logger.Warning().Msgf("Could not create metadata cache: %v", err) + } else { + metadataIdx = idx + if err := metadataIdx.Load(); err != nil { + store.logger.Warning().Msgf("Could not load metadata cache: %v", err) + } + } + + return metadataIdx +} + // LoadTemplatesOnlyMetadata loads only the metadata of the templates func (store *Store) LoadTemplatesOnlyMetadata() error { + defer store.saveMetadataIndexOnce() + templatePaths, errs := store.config.Catalog.GetTemplatesPath(store.finalTemplates) store.logErroredTemplates(errs) - filteredTemplatePaths := store.pathFilter.Match(templatePaths) - + indexFilter := store.indexFilter validPaths := make(map[string]struct{}) - for templatePath := range filteredTemplatePaths { + + for _, templatePath := range templatePaths { + if store.metadataIndex != nil { + if metadata, found := store.metadataIndex.Get(templatePath); found { + if !indexFilter.Matches(metadata) { + continue + } + + if store.tagFilter != nil { + loaded, err := store.config.ExecutorOptions.Parser.LoadTemplate(templatePath, store.tagFilter, nil, store.config.Catalog) + if !loaded { + if err != nil && strings.Contains(err.Error(), templates.ErrExcluded.Error()) { + stats.Increment(templates.TemplatesExcludedStats) + if config.DefaultConfig.LogAllEvents { + store.logger.Print().Msgf("[%v] %v\n", aurora.Yellow("WRN").String(), err.Error()) + } + } + continue + } + } + + validPaths[templatePath] = struct{}{} + continue + } + } + loaded, err := store.config.ExecutorOptions.Parser.LoadTemplate(templatePath, store.tagFilter, nil, store.config.Catalog) - if loaded || store.pathFilter.MatchIncluded(templatePath) { + if loaded { + templatesCache := store.parserCacheOnce() + if templatesCache != nil { + if template, _, _ := templatesCache.Has(templatePath); template != nil { + var metadata *index.Metadata + if store.metadataIndex != nil { + metadata, _ = store.metadataIndex.SetFromTemplate(templatePath, template) + } else { + metadata = index.NewMetadataFromTemplate(templatePath, template) + } + + if !indexFilter.Matches(metadata) { + continue + } + + validPaths[templatePath] = struct{}{} + continue + } + } + validPaths[templatePath] = struct{}{} } if err != nil { @@ -376,15 +474,24 @@ func (store *Store) LoadTemplatesOnlyMetadata() error { func (store *Store) ValidateTemplates() error { templatePaths, errs := store.config.Catalog.GetTemplatesPath(store.finalTemplates) store.logErroredTemplates(errs) + workflowPaths, errs := store.config.Catalog.GetTemplatesPath(store.finalWorkflows) store.logErroredTemplates(errs) - filteredTemplatePaths := store.pathFilter.Match(templatePaths) - filteredWorkflowPaths := store.pathFilter.Match(workflowPaths) + templatePathsMap := make(map[string]struct{}, len(templatePaths)) + for _, path := range templatePaths { + templatePathsMap[path] = struct{}{} + } + + workflowPathsMap := make(map[string]struct{}, len(workflowPaths)) + for _, path := range workflowPaths { + workflowPathsMap[path] = struct{}{} + } - if store.areTemplatesValid(filteredTemplatePaths) && store.areWorkflowsValid(filteredWorkflowPaths) { + if store.areTemplatesValid(templatePathsMap) && store.areWorkflowsValid(workflowPathsMap) { return nil } + return errors.New("errors occurred during template validation") } @@ -503,10 +610,9 @@ func (store *Store) LoadTemplates(templatesList []string) []*templates.Template func (store *Store) LoadWorkflows(workflowsList []string) []*templates.Template { includedWorkflows, errs := store.config.Catalog.GetTemplatesPath(workflowsList) store.logErroredTemplates(errs) - workflowPathMap := store.pathFilter.Match(includedWorkflows) - loadedWorkflows := make([]*templates.Template, 0, len(workflowPathMap)) - for workflowPath := range workflowPathMap { + loadedWorkflows := make([]*templates.Template, 0, len(includedWorkflows)) + for _, workflowPath := range includedWorkflows { loaded, err := store.config.ExecutorOptions.Parser.LoadWorkflow(workflowPath, store.config.Catalog) if err != nil { store.logger.Warning().Msgf("Could not load workflow %s: %s\n", workflowPath, err) @@ -526,9 +632,12 @@ func (store *Store) LoadWorkflows(workflowsList []string) []*templates.Template // LoadTemplatesWithTags takes a list of templates and extra tags // returning templates that match. func (store *Store) LoadTemplatesWithTags(templatesList, tags []string) []*templates.Template { + defer store.saveMetadataIndexOnce() + + indexFilter := store.indexFilter + includedTemplates, errs := store.config.Catalog.GetTemplatesPath(templatesList) store.logErroredTemplates(errs) - templatePathMap := store.pathFilter.Match(includedTemplates) loadedTemplates := sliceutil.NewSyncSlice[*templates.Template]() loadedTemplateIDs := mapsutil.NewSyncLockMap[string, struct{}]() @@ -572,14 +681,46 @@ func (store *Store) LoadTemplatesWithTags(templatesList, tags []string) []*templ panic("dialers with executionId " + store.config.ExecutorOptions.Options.ExecutionId + " not found") } - for templatePath := range templatePathMap { + for _, templatePath := range includedTemplates { wgLoadTemplates.Add() go func(templatePath string) { defer wgLoadTemplates.Done() + var ( + metadata *index.Metadata + metadataCached bool + ) + + if store.metadataIndex != nil { + if cachedMetadata, found := store.metadataIndex.Get(templatePath); found { + metadata = cachedMetadata + if !indexFilter.Matches(metadata) { + return + } + // NOTE(dwisiswant0): else, tagFilter probably exists (for + // IncludeConditions), which still need to check via + // LoadTemplate. + + metadataCached = true + } + } + loaded, err := store.config.ExecutorOptions.Parser.LoadTemplate(templatePath, store.tagFilter, tags, store.config.Catalog) - if loaded || store.pathFilter.MatchIncluded(templatePath) { + if loaded { parsed, err := templates.Parse(templatePath, store.preprocessor, store.config.ExecutorOptions) + + if parsed != nil && !metadataCached { + if store.metadataIndex != nil { + metadata, _ = store.metadataIndex.SetFromTemplate(templatePath, parsed) + } else { + metadata = index.NewMetadataFromTemplate(templatePath, parsed) + } + + if metadata != nil && !indexFilter.Matches(metadata) { + return + } + } + if err != nil { // exclude templates not compatible with offline matching from total runtime warning stats if !errors.Is(err, templates.ErrIncompatibleWithOfflineMatching) { diff --git a/pkg/catalog/loader/loader_bench_test.go b/pkg/catalog/loader/loader_bench_test.go index 079e928ad5..32ed506e8b 100644 --- a/pkg/catalog/loader/loader_bench_test.go +++ b/pkg/catalog/loader/loader_bench_test.go @@ -8,7 +8,9 @@ import ( "github.com/projectdiscovery/nuclei/v3/pkg/catalog/disk" "github.com/projectdiscovery/nuclei/v3/pkg/catalog/loader" "github.com/projectdiscovery/nuclei/v3/pkg/loader/workflow" + "github.com/projectdiscovery/nuclei/v3/pkg/model/types/severity" "github.com/projectdiscovery/nuclei/v3/pkg/templates" + templateTypes "github.com/projectdiscovery/nuclei/v3/pkg/templates/types" "github.com/projectdiscovery/nuclei/v3/pkg/testutils" ) @@ -41,3 +43,201 @@ func BenchmarkStoreValidateTemplates(b *testing.B) { _ = store.ValidateTemplates() } } + +func BenchmarkLoadTemplates(b *testing.B) { + options := testutils.DefaultOptions.Copy() + options.Logger = &gologger.Logger{} + options.ExecutionId = "bench-load-templates" + testutils.Init(options) + + catalog := disk.NewCatalog(config.DefaultConfig.TemplatesDirectory) + executerOpts := testutils.NewMockExecuterOptions(options, nil) + executerOpts.Parser = templates.NewParser() + + workflowLoader, err := workflow.NewLoader(executerOpts) + if err != nil { + b.Fatalf("could not create workflow loader: %s", err) + } + executerOpts.WorkflowLoader = workflowLoader + + b.Run("NoFilter", func(b *testing.B) { + loaderCfg := loader.NewConfig(options, catalog, executerOpts) + store, err := loader.New(loaderCfg) + if err != nil { + b.Fatalf("could not create store: %s", err) + } + + b.ResetTimer() + b.ReportAllocs() + + for b.Loop() { + _ = store.LoadTemplates([]string{config.DefaultConfig.TemplatesDirectory}) + } + }) + + b.Run("FilterBySeverityCritical", func(b *testing.B) { + opts := options.Copy() + opts.Severities = severity.Severities{severity.Critical} + loaderCfg := loader.NewConfig(opts, catalog, executerOpts) + + store, err := loader.New(loaderCfg) + if err != nil { + b.Fatalf("could not create store: %s", err) + } + + b.ResetTimer() + b.ReportAllocs() + + for b.Loop() { + _ = store.LoadTemplates([]string{config.DefaultConfig.TemplatesDirectory}) + } + }) + + b.Run("FilterBySeverityHighCritical", func(b *testing.B) { + opts := options.Copy() + opts.Severities = severity.Severities{severity.High, severity.Critical} + loaderCfg := loader.NewConfig(opts, catalog, executerOpts) + + store, err := loader.New(loaderCfg) + if err != nil { + b.Fatalf("could not create store: %s", err) + } + + b.ResetTimer() + b.ReportAllocs() + + for b.Loop() { + _ = store.LoadTemplates([]string{config.DefaultConfig.TemplatesDirectory}) + } + }) + + b.Run("FilterByAuthor", func(b *testing.B) { + opts := options.Copy() + opts.Authors = []string{"pdteam"} + loaderCfg := loader.NewConfig(opts, catalog, executerOpts) + + store, err := loader.New(loaderCfg) + if err != nil { + b.Fatalf("could not create store: %s", err) + } + + b.ResetTimer() + b.ReportAllocs() + + for b.Loop() { + _ = store.LoadTemplates([]string{config.DefaultConfig.TemplatesDirectory}) + } + }) + + b.Run("FilterByTags", func(b *testing.B) { + opts := options.Copy() + opts.Tags = []string{"cve", "rce"} + loaderCfg := loader.NewConfig(opts, catalog, executerOpts) + + store, err := loader.New(loaderCfg) + if err != nil { + b.Fatalf("could not create store: %s", err) + } + + b.ResetTimer() + b.ReportAllocs() + + for b.Loop() { + _ = store.LoadTemplates([]string{config.DefaultConfig.TemplatesDirectory}) + } + }) + + b.Run("FilterByProtocol", func(b *testing.B) { + opts := options.Copy() + opts.Protocols = templateTypes.ProtocolTypes{templateTypes.HTTPProtocol} + loaderCfg := loader.NewConfig(opts, catalog, executerOpts) + + store, err := loader.New(loaderCfg) + if err != nil { + b.Fatalf("could not create store: %s", err) + } + + b.ResetTimer() + b.ReportAllocs() + + for b.Loop() { + _ = store.LoadTemplates([]string{config.DefaultConfig.TemplatesDirectory}) + } + }) + + b.Run("ComplexFilter", func(b *testing.B) { + opts := options.Copy() + opts.Severities = severity.Severities{severity.High, severity.Critical} + opts.Authors = []string{"pdteam"} + opts.Tags = []string{"cve"} + loaderCfg := loader.NewConfig(opts, catalog, executerOpts) + + store, err := loader.New(loaderCfg) + if err != nil { + b.Fatalf("could not create store: %s", err) + } + + b.ResetTimer() + b.ReportAllocs() + + for b.Loop() { + _ = store.LoadTemplates([]string{config.DefaultConfig.TemplatesDirectory}) + } + }) +} + +func BenchmarkLoadTemplatesOnlyMetadata(b *testing.B) { + options := testutils.DefaultOptions.Copy() + options.Logger = &gologger.Logger{} + options.ExecutionId = "bench-metadata" + testutils.Init(options) + + catalog := disk.NewCatalog(config.DefaultConfig.TemplatesDirectory) + executerOpts := testutils.NewMockExecuterOptions(options, nil) + executerOpts.Parser = templates.NewParser() + + workflowLoader, err := workflow.NewLoader(executerOpts) + if err != nil { + b.Fatalf("could not create workflow loader: %s", err) + } + executerOpts.WorkflowLoader = workflowLoader + + b.Run("WithoutFilter", func(b *testing.B) { + loaderCfg := loader.NewConfig(options, catalog, executerOpts) + store, err := loader.New(loaderCfg) + if err != nil { + b.Fatalf("could not create store: %s", err) + } + + // Pre-warm the cache + _ = store.LoadTemplatesOnlyMetadata() + + b.ResetTimer() + b.ReportAllocs() + + for b.Loop() { + _ = store.LoadTemplatesOnlyMetadata() + } + }) + + b.Run("WithSeverityFilter", func(b *testing.B) { + opts := options.Copy() + opts.Severities = severity.Severities{severity.Critical} + loaderCfg := loader.NewConfig(opts, catalog, executerOpts) + + store, err := loader.New(loaderCfg) + if err != nil { + b.Fatalf("could not create store: %s", err) + } + + // Pre-warm the cache + _ = store.LoadTemplatesOnlyMetadata() + + b.ResetTimer() + b.ReportAllocs() + + for b.Loop() { + _ = store.LoadTemplatesOnlyMetadata() + } + }) +}