Merge pull request #655 from hahwul/improve/server-testcodes

Enhance scan error handling and add tests for scan functionality

Author: hahwul

codecov.yml

@@ -1,2 +1,3 @@
 ignore:  
-  - cmd
+  - cmd
+  - pkg/server/docs/docs.go

go.sum

@@ -16,8 +16,8 @@ github.com/briandowns/spinner v1.23.2 h1:Zc6ecUnI+YzLmJniCfDNaMbW0Wid1d5+qcTq4L2
 github.com/briandowns/spinner v1.23.2/go.mod h1:LaZeM4wm2Ywy6vO571mvhQNRcWfRUnXOs0RcKV0wYKM=
 github.com/chromedp/cdproto v0.0.0-20250222051814-50c6cb17f10a h1:EnkQjhmp/MxhDB4KOTssv6xC20aQ9rhFRCfGHTsTqmE=
 github.com/chromedp/cdproto v0.0.0-20250222051814-50c6cb17f10a/go.mod h1:NItd7aLkcfOA/dcMXvl8p1u+lQqioRMq/SqDp71Pb/k=
-github.com/chromedp/chromedp v0.13.1 h1:FDh9CfaAt0w70gl69Hb69M/xgZrWuppH9AW22aGa+iU=
-github.com/chromedp/chromedp v0.13.1/go.mod h1:O3nO4Lno7iLoVX+7GdqQkehhKG7DtLf/zFRyJo0AhXY=
+github.com/chromedp/chromedp v0.13.0 h1:ydOqt7Y9LkwgutrX5C8bx49D+o63L6WcGUDyIoE0A5M=
+github.com/chromedp/chromedp v0.13.0/go.mod h1:O3nO4Lno7iLoVX+7GdqQkehhKG7DtLf/zFRyJo0AhXY=
 github.com/chromedp/sysutil v1.1.0 h1:PUFNv5EcprjqXZD9nJb9b/c9ibAbxiYo4exNWZyipwM=
 github.com/chromedp/sysutil v1.1.0/go.mod h1:WiThHUdltqCNKGc4gaU50XgYjwjYIhKWoHGPTUfWTJ8=
 github.com/cpuguy83/go-md2man/v2 v2.0.6/go.mod h1:oOW0eioCTA6cOiMLiUPZOpcVxMig6NIQQ7OS05n1F4g=

pkg/server/scan.go

@@ -34,7 +34,12 @@ func ScanFromAPI(url string, rqOptions model.Options, options model.Options, sid
 	escapedURL := cleanURL(url)
 	vLog.WithField("data1", sid).Debug(escapedURL)
 	vLog.WithField("data1", sid).Debug(newOptions)
-	_, _ = scan.Scan(url, newOptions, sid)
+	_, err := scan.Scan(url, newOptions, sid)
+	if err != nil {
+vLog.WithField("data1", sid).Error("Scan failed for URL:", url, ": ", err)
+		return
+	}
+	vLog.WithField("data1", sid).Info("Scan completed successfully")
 }
 
 // GetScan is get scan information

pkg/server/server_test.go

@@ -1,9 +1,12 @@
 package server
 
 import (
+	"bytes"
+	"encoding/json"
 	"net/http"
 	"net/http/httptest"
 	"testing"
+	"time"
 
 	"github.com/hahwul/dalfox/v2/pkg/model"
 	_ "github.com/hahwul/dalfox/v2/pkg/server/docs"
@@ -77,3 +80,99 @@ func Test_healthHandler(t *testing.T) {
 		assert.Contains(t, rec.Body.String(), "ok")
 	}
 }
+
+func Test_postScanHandler(t *testing.T) {
+	e := echo.New()
+	rq := Req{
+		URL: "http://example.com",
+		Options: model.Options{
+			Method: "GET",
+		},
+	}
+	body, _ := json.Marshal(rq)
+	req := httptest.NewRequest(http.MethodPost, "/scan", bytes.NewReader(body))
+	req.Header.Set(echo.HeaderContentType, echo.MIMEApplicationJSON)
+	rec := httptest.NewRecorder()
+	c := e.NewContext(req, rec)
+
+	scans := []string{}
+	options := model.Options{
+		Scan: map[string]model.Scan{},
+	}
+
+	if assert.NoError(t, postScanHandler(c, &scans, options)) {
+		assert.Equal(t, http.StatusOK, rec.Code)
+		assert.Contains(t, rec.Body.String(), "code")
+		assert.Contains(t, rec.Body.String(), "msg")
+		assert.Contains(t, rec.Body.String(), "data")
+	}
+}
+
+func Test_GetScan(t *testing.T) {
+	options := model.Options{
+		Scan: map[string]model.Scan{
+			"test-scan": {URL: "http://example.com", Results: []model.PoC{{Type: "finish"}}},
+		},
+	}
+	scan := GetScan("test-scan", options)
+	assert.Equal(t, "http://example.com", scan.URL)
+	assert.Equal(t, "finish", scan.Results[0].Type)
+}
+
+func Test_GetScans(t *testing.T) {
+	options := model.Options{
+		Scan: map[string]model.Scan{
+			"test-scan1": {URL: "http://example1.com"},
+			"test-scan2": {URL: "http://example2.com"},
+		},
+	}
+	scans := GetScans(options)
+	assert.Contains(t, scans, "test-scan1")
+	assert.Contains(t, scans, "test-scan2")
+}
+
+func Test_ScanFromAPI(t *testing.T) {
+	options := model.Options{
+		Debug: true,
+		Scan:  map[string]model.Scan{},
+	}
+	rqOptions := model.Options{
+		Method: "GET",
+	}
+	sid := "test-scan-id"
+
+	t.Run("Successful Scan", func(t *testing.T) {
+		ScanFromAPI("http://example.com", rqOptions, options, sid)
+		// Add assertions to verify the scan was successful
+	})
+
+	t.Run("Scan with Error", func(t *testing.T) {
+		ScanFromAPI("http://invalid-url", rqOptions, options, sid)
+		// Add assertions to verify error handling
+	})
+}
+
+func Test_setupEchoServer(t *testing.T) {
+	options := model.Options{
+		ServerHost: "localhost",
+		ServerPort: 6664,
+	}
+	scans := []string{}
+	e := setupEchoServer(options, &scans)
+
+	assert.NotNil(t, e)
+	assert.Equal(t, "localhost:6664", e.Server.Addr)
+}
+
+func Test_RunAPIServer(t *testing.T) {
+	options := model.Options{
+		ServerHost: "localhost",
+		ServerPort: 6664,
+	}
+	go RunAPIServer(options)
+	time.Sleep(1 * time.Second)
+
+	resp, err := http.Get("http://localhost:6664/health")
+	assert.NoError(t, err)
+	assert.Equal(t, http.StatusOK, resp.StatusCode)
+}