From 085d64fbdaaf627f29ffbfe0b34a444a4db33ebc Mon Sep 17 00:00:00 2001
From: lyyyuna <lyyyuna@outlook.com>
Date: Fri, 15 Nov 2024 15:08:53 +0800
Subject: [PATCH] feat: register extra information during build

---
 cmd/build.go        |  7 +++++--
 cmd/install.go      |  2 ++
 pkg/build/agent.tpl |  7 ++++++-
 pkg/build/build.go  |  2 +-
 pkg/build/config.go |  8 ++++++++
 pkg/build/inject.go | 18 ++++++++++--------
 6 files changed, 32 insertions(+), 12 deletions(-)

diff --git a/cmd/build.go b/cmd/build.go
index 0b0d72df..4a3f2289 100644
--- a/cmd/build.go
+++ b/cmd/build.go
@@ -26,13 +26,15 @@ var buildCmd = &cobra.Command{
 }
 
 var (
-	gocmode string
-	gochost string
+	gocmode  string
+	gochost  string
+	gocextra string
 )
 
 func init() {
 	buildCmd.Flags().StringVarP(&gocmode, "gocmode", "", "count", "coverage mode: set, count, atomic, watch")
 	buildCmd.Flags().StringVarP(&gochost, "gochost", "", "127.0.0.1:7777", "specify the host of the goc sever")
+	buildCmd.Flags().StringVarP(&gocextra, "gocextra", "", "", "specify the extra information injected into the build")
 	rootCmd.AddCommand(buildCmd)
 }
 
@@ -47,6 +49,7 @@ func buildAction(cmd *cobra.Command, args []string) {
 		build.WithArgs(args),
 		build.WithBuild(),
 		build.WithDebug(globalDebug),
+		build.WithExtra(gocextra),
 	)
 	b.Build()
 
diff --git a/cmd/install.go b/cmd/install.go
index 662d744a..ebbfac0d 100644
--- a/cmd/install.go
+++ b/cmd/install.go
@@ -28,6 +28,7 @@ var installCmd = &cobra.Command{
 func init() {
 	installCmd.Flags().StringVarP(&gocmode, "gocmode", "", "count", "coverage mode: set, count, atomic, watch")
 	installCmd.Flags().StringVarP(&gochost, "gochost", "", "127.0.0.1:7777", "specify the host of the goc sever")
+	installCmd.Flags().StringVarP(&gocextra, "gocextra", "", "", "specify the extra information injected into the build")
 	rootCmd.AddCommand(installCmd)
 }
 
@@ -42,6 +43,7 @@ func installAction(cmd *cobra.Command, args []string) {
 		build.WithArgs(args),
 		build.WithInstall(),
 		build.WithDebug(globalDebug),
+		build.WithExtra(gocextra),
 	)
 	b.Install()
 
diff --git a/pkg/build/agent.tpl b/pkg/build/agent.tpl
index 08069292..4d20bb53 100644
--- a/pkg/build/agent.tpl
+++ b/pkg/build/agent.tpl
@@ -45,7 +45,7 @@ var (
 	token string
 	id string
 	cond = sync.NewCond(&sync.Mutex{})
-	register_extra = os.Getenv("GOC_REGISTER_EXTRA")
+	register_extra = "{{.Extra}}"
 )
 
 func init() {
@@ -55,6 +55,11 @@ func init() {
 		host = host_env
 	}
 
+	// init extra information
+	if os.Getenv("GOC_REGISTER_EXTRA") != "" {
+		register_extra = os.Getenv("GOC_REGISTER_EXTRA")
+	}
+
 	var dialer = websocket.DefaultDialer
 
 	go func() {
diff --git a/pkg/build/build.go b/pkg/build/build.go
index eece0d2a..326c175a 100644
--- a/pkg/build/build.go
+++ b/pkg/build/build.go
@@ -31,6 +31,7 @@ type Build struct {
 	Debug bool
 	Host  string
 	Mode  string // cover mode
+	Extra string
 
 	GOPATH           string
 	GOBIN            string
@@ -53,7 +54,6 @@ type Build struct {
 }
 
 // NewBuild creates a Build struct
-//
 func NewBuild(opts ...gocOption) *Build {
 	b := &Build{}
 
diff --git a/pkg/build/config.go b/pkg/build/config.go
index bb052a85..91344e25 100644
--- a/pkg/build/config.go
+++ b/pkg/build/config.go
@@ -14,6 +14,8 @@
 package build
 
 import (
+	"strings"
+
 	"github.com/spf13/pflag"
 )
 
@@ -60,3 +62,9 @@ func WithDebug(enable bool) gocOption {
 		b.Debug = enable
 	}
 }
+
+func WithExtra(extra string) gocOption {
+	return func(b *Build) {
+		b.Extra = strings.TrimSpace(extra)
+	}
+}
diff --git a/pkg/build/inject.go b/pkg/build/inject.go
index 8c00b30c..5a6d59dc 100644
--- a/pkg/build/inject.go
+++ b/pkg/build/inject.go
@@ -127,14 +127,14 @@ func (b *Build) getPkgTmpDir(pkgDir string) string {
 
 // injectGocAgent inject handlers like following
 //
-// - xxx.go
-// - yyy_package
-// - main.go
-// - goc-cover-agent-apis-auto-generated-11111-22222-bridge.go
-// - goc-cover-agent-apis-auto-generated-11111-22222-package
-//  |
-//  -- rpcagent.go
-//  -- watchagent.go
+//   - xxx.go
+//   - yyy_package
+//   - main.go
+//   - goc-cover-agent-apis-auto-generated-11111-22222-bridge.go
+//   - goc-cover-agent-apis-auto-generated-11111-22222-package
+//     |
+//     -- rpcagent.go
+//     -- watchagent.go
 //
 // 11111_22222_bridge.go 仅仅用于引用 11111_22222_package, where package contains ws agent main logic.
 // 使用 bridge.go 文件是为了避免插桩逻辑中的变量名污染 main 包
@@ -195,12 +195,14 @@ func (b *Build) injectGocAgent(where string, covers []*PackageCover) {
 		Package                  string
 		Host                     string
 		Mode                     string
+		Extra                    string
 	}{
 		Covers:                   covers,
 		GlobalCoverVarImportPath: b.GlobalCoverVarImportPath,
 		Package:                  injectPkgName,
 		Host:                     b.Host,
 		Mode:                     _coverMode,
+		Extra:                    b.Extra,
 	}
 
 	if err := coverMainTmpl.Execute(f2, tmplData); err != nil {