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 a1a870b..6e62ea2 100644
--- a/common/pagetypeclassifier/pagetypeclassifier.go
+++ b/common/pagetypeclassifier/pagetypeclassifier.go
@@ -14,26 +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)
}
-func htmlToText(html string) string {
- text, err := htmltomarkdown.ConvertString(html)
- if err != nil {
- panic(err)
- }
- return text
+// htmlToText safely converts HTML to text and protects against panics from Go's HTML parser.
+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)