From 05c6364d90b9c291957d3e78c5dfb00a4fe80265 Mon Sep 17 00:00:00 2001 From: Mzack9999 Date: Tue, 18 Nov 2025 16:46:15 +0400 Subject: [PATCH] better error handling --- common/httpx/httpx.go | 2 +- .../pagetypeclassifier/pagetypeclassifier.go | 25 ++++++------------- .../pagetypeclassifier_test.go | 19 ++++++++------ runner/runner.go | 6 ++++- 4 files changed, 25 insertions(+), 27 deletions(-) diff --git a/common/httpx/httpx.go b/common/httpx/httpx.go index 1b75ee8..39d1ce3 100644 --- a/common/httpx/httpx.go +++ b/common/httpx/httpx.go @@ -1,6 +1,7 @@ package httpx import ( + "context" "crypto/tls" "fmt" "io" @@ -25,7 +26,6 @@ import ( pdhttputil "github.com/projectdiscovery/utils/http" stringsutil "github.com/projectdiscovery/utils/strings" urlutil "github.com/projectdiscovery/utils/url" - "golang.org/x/net/context" "golang.org/x/net/http2" ) diff --git a/common/pagetypeclassifier/pagetypeclassifier.go b/common/pagetypeclassifier/pagetypeclassifier.go index 5e76747..6e62ea2 100644 --- a/common/pagetypeclassifier/pagetypeclassifier.go +++ b/common/pagetypeclassifier/pagetypeclassifier.go @@ -14,34 +14,23 @@ type PageTypeClassifier struct { classifier *naive_bayes.NaiveBayesClassifier } -func New() *PageTypeClassifier { +func New() (*PageTypeClassifier, error) { classifier, err := naive_bayes.NewClassifierFromFileData(classifierData) if err != nil { - panic(err) + return nil, err } - return &PageTypeClassifier{classifier: classifier} + return &PageTypeClassifier{classifier: classifier}, nil } func (n *PageTypeClassifier) Classify(html string) string { - text := htmlToText(html) - if text == "" { + text, err := htmlToText(html) + if err != nil || text == "" { return "other" } return n.classifier.Classify(text) } // htmlToText safely converts HTML to text and protects against panics from Go's HTML parser. -func htmlToText(html string) (text string) { - defer func() { - if r := recover(); r != nil { - // Optionally log this event, e.g., log.Printf("Recovered in htmlToText: %v", r) - text = "" - } - }() - var err error - text, err = htmltomarkdown.ConvertString(html) - if err != nil { - text = "" - } - return +func htmlToText(html string) (string, error) { + return htmltomarkdown.ConvertString(html) } diff --git a/common/pagetypeclassifier/pagetypeclassifier_test.go b/common/pagetypeclassifier/pagetypeclassifier_test.go index 30571f0..4d0e9a7 100644 --- a/common/pagetypeclassifier/pagetypeclassifier_test.go +++ b/common/pagetypeclassifier/pagetypeclassifier_test.go @@ -3,19 +3,22 @@ package pagetypeclassifier import ( "testing" - "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" ) func TestPageTypeClassifier(t *testing.T) { t.Run("test creation of new PageTypeClassifier", func(t *testing.T) { - epc := New() - assert.NotNil(t, epc) + epc, err := New() + require.NoError(t, err) + require.NotNil(t, epc) }) t.Run("test classification non error page text", func(t *testing.T) { - epc := New() - assert.Equal(t, "nonerror", epc.Classify(` + epc, err := New() + require.NoError(t, err) + require.NotNil(t, epc) + require.Equal(t, "nonerror", epc.Classify(` @@ -30,8 +33,10 @@ func TestPageTypeClassifier(t *testing.T) { }) t.Run("test classification on error page text", func(t *testing.T) { - epc := New() - assert.Equal(t, "error", epc.Classify(` + epc, err := New() + require.NoError(t, err) + require.NotNil(t, epc) + require.Equal(t, "error", epc.Classify(` Error 403: Forbidden diff --git a/runner/runner.go b/runner/runner.go index 82df4fa..2035199 100644 --- a/runner/runner.go +++ b/runner/runner.go @@ -385,7 +385,11 @@ func New(options *Options) (*Runner, error) { } runner.simHashes = gcache.New[uint64, struct{}](1000).ARC().Build() - runner.pageTypeClassifier = pagetypeclassifier.New() + pageTypeClassifier, err := pagetypeclassifier.New() + if err != nil { + return nil, err + } + runner.pageTypeClassifier = pageTypeClassifier if options.HttpApiEndpoint != "" { apiServer := NewServer(options.HttpApiEndpoint, options)