From ec345d9e530effda7c9630d1fc2eb1a214653686 Mon Sep 17 00:00:00 2001
From: debugtalk <mail@debugtalk.com>
Date: Tue, 14 Jun 2022 23:17:34 +0800
Subject: [PATCH] fix: check if debugtalk.go contains main() function

---
 .../plugin/debugtalk_gen.go                   | 16 +++++++++
 hrp/build.go                                  | 23 ++++++++++---
 hrp/build_test.go                             | 33 ++++++++++++++-----
 3 files changed, 58 insertions(+), 14 deletions(-)
 create mode 100644 examples/demo-with-go-plugin/plugin/debugtalk_gen.go

diff --git a/examples/demo-with-go-plugin/plugin/debugtalk_gen.go b/examples/demo-with-go-plugin/plugin/debugtalk_gen.go
new file mode 100644
index 00000000..df7727c3
--- /dev/null
+++ b/examples/demo-with-go-plugin/plugin/debugtalk_gen.go
@@ -0,0 +1,16 @@
+// NOTE: Generated By hrp v4.1.3, DO NOT EDIT!
+package main
+
+import (
+	"github.com/httprunner/funplugin/fungo"
+)
+
+func main() {
+	fungo.Register("SumTwoInt", SumTwoInt)
+	fungo.Register("SumInts", SumInts)
+	fungo.Register("Sum", Sum)
+	fungo.Register("SetupHookExample", SetupHookExample)
+	fungo.Register("TeardownHookExample", TeardownHookExample)
+	fungo.Register("GetUserAgent", GetUserAgent)
+	fungo.Serve()
+}
diff --git a/hrp/build.go b/hrp/build.go
index 81682dfd..43cdaf82 100644
--- a/hrp/build.go
+++ b/hrp/build.go
@@ -35,7 +35,7 @@ var (
 	regexGoFunctionName = regexFunctions{regexp.MustCompile(`(?m)^func ([a-zA-Z_]\w*)\(.*\)`)}
 )
 
-func (r *regexFunctions) findAllFunctionNames(content string) []string {
+func (r *regexFunctions) findAllFunctionNames(content string) ([]string, error) {
 	var functionNames []string
 	// find all function names
 	functionNameSlice := r.FindAllStringSubmatch(content, -1)
@@ -56,7 +56,11 @@ func (r *regexFunctions) findAllFunctionNames(content string) []string {
 	} else if r == &regexGoFunctionName {
 		// filter main and init function
 		for _, name := range functionNames {
-			if name == "main" || name == "init" {
+			if name == "main" {
+				log.Warn().Msg("plugin debugtalk.go should not define main() function !!!")
+				return nil, errors.New("debugtalk.go should not contain main() function")
+			}
+			if name == "init" {
 				continue
 			}
 			filteredFunctionNames = append(filteredFunctionNames, name)
@@ -64,7 +68,7 @@ func (r *regexFunctions) findAllFunctionNames(content string) []string {
 	}
 
 	log.Info().Strs("functionNames", filteredFunctionNames).Msg("find all function names")
-	return filteredFunctionNames
+	return filteredFunctionNames, nil
 }
 
 type pluginTemplate struct {
@@ -167,14 +171,19 @@ func (pt *pluginTemplate) generateGo(output string) error {
 
 // buildGo builds debugtalk.go to debugtalk.bin
 func buildGo(path string, output string) error {
+	log.Info().Str("path", path).Str("output", output).Msg("start to build go plugin")
 	content, err := os.ReadFile(path)
 	if err != nil {
 		log.Error().Err(err).Msg("failed to read file")
 		return errors.Wrap(err, "read file failed")
 	}
-	functionNames := regexGoFunctionName.findAllFunctionNames(string(content))
+	functionNames, err := regexGoFunctionName.findAllFunctionNames(string(content))
+	if err != nil {
+		return err
+	}
 
 	templateContent := &pluginTemplate{
+		path:          path,
 		Version:       version.VERSION,
 		FunctionNames: functionNames,
 	}
@@ -183,6 +192,7 @@ func buildGo(path string, output string) error {
 
 // buildPy completes funppy information in debugtalk.py
 func buildPy(path string, output string) error {
+	log.Info().Str("path", path).Str("output", output).Msg("start to prepare python plugin")
 	// check the syntax of debugtalk.py
 	err := builtin.ExecCommand("python3", "-m", "py_compile", path)
 	if err != nil {
@@ -194,7 +204,10 @@ func buildPy(path string, output string) error {
 		log.Error().Err(err).Msg("failed to read file")
 		return errors.Wrap(err, "read file failed")
 	}
-	functionNames := regexPyFunctionName.findAllFunctionNames(string(content))
+	functionNames, err := regexPyFunctionName.findAllFunctionNames(string(content))
+	if err != nil {
+		return err
+	}
 
 	templateContent := &pluginTemplate{
 		path:          path,
diff --git a/hrp/build_test.go b/hrp/build_test.go
index c8932301..ab5d88e1 100644
--- a/hrp/build_test.go
+++ b/hrp/build_test.go
@@ -60,7 +60,10 @@ def __test_3():	# private function
 def Test5():	# exported function
     pass
 `
-	names := regexPyFunctionName.findAllFunctionNames(content)
+	names, err := regexPyFunctionName.findAllFunctionNames(content)
+	if !assert.Nil(t, err) {
+		t.FailNow()
+	}
 	if !assert.Contains(t, names, "test_1") {
 		t.FailNow()
 	}
@@ -89,10 +92,6 @@ func testFunc2() {	// exported function
 	return
 }
 
-func main() {	// private function
-	return
-}
-
 func init() {	// private function
 	return
 }
@@ -105,16 +104,16 @@ func _Test3() { // exported function
 // 	return
 // }
 `
-	names := regexGoFunctionName.findAllFunctionNames(content)
+	names, err := regexGoFunctionName.findAllFunctionNames(content)
+	if !assert.Nil(t, err) {
+		t.FailNow()
+	}
 	if !assert.Contains(t, names, "Test1") {
 		t.FailNow()
 	}
 	if !assert.Contains(t, names, "testFunc2") {
 		t.FailNow()
 	}
-	if !assert.NotContains(t, names, "main") {
-		t.FailNow()
-	}
 	if !assert.NotContains(t, names, "init") {
 		t.FailNow()
 	}
@@ -126,3 +125,19 @@ func _Test3() { // exported function
 		t.FailNow()
 	}
 }
+
+func TestFindAllGoFunctionNamesAbnormal(t *testing.T) {
+	content := `
+func init() {	// private function
+	return
+}
+
+func main() {	// should not define main() function
+	return
+}
+`
+	_, err := regexGoFunctionName.findAllFunctionNames(content)
+	if !assert.NotNil(t, err) {
+		t.FailNow()
+	}
+}