fix: api.CallContext alias conflicts with hooked package (alibaba#235)

pkg/rules/test/fmt7/hook.go

@@ -17,13 +17,13 @@ package fmt7
 import (
 	_ "fmt"
 
-	"github.com/alibaba/opentelemetry-go-auto-instrumentation/pkg/api"
+	aliasapi "github.com/alibaba/opentelemetry-go-auto-instrumentation/pkg/api" // both alias api and instrumented package(fmt) are imported
 )
 
-func onEnterSprintf3(call api.CallContext, format string, arg ...any) {
+func onEnterSprintf3(call aliasapi.CallContext, format string, arg ...any) {
 	println("a3")
 }
 
-func onExitSprintf3(call api.CallContext, s string) {
+func onExitSprintf3(call aliasapi.CallContext, s string) {
 	print("b3")
 }

test/helloworld_test.go

@@ -25,6 +25,7 @@ const HelloworldAppName = "helloworld"
 func TestRunHelloworld(t *testing.T) {
 	UseApp(HelloworldAppName)
 
+	RunSet(t, "-rule=")
 	RunGoBuild(t, "go", "build") // no test rules, build as usual
 	stdout, _ := RunApp(t, HelloworldAppName)
 	ExpectContains(t, stdout, "helloworld")

tool/preprocess/match.go

@@ -106,11 +106,6 @@ func findAvailableRules() []resource.InstRule {
 		return nil
 	}
 
-	// If rule file is not set, we will use the default rules
-	if config.GetConf().RuleJsonFiles == "" {
-		return loadDefaultRules()
-	}
-
 	rules := make([]resource.InstRule, 0)
 
 	// Load default rules unless explicitly disabled
@@ -119,27 +114,29 @@ func findAvailableRules() []resource.InstRule {
 		rules = append(rules, defaultRules...)
 	}
 
-	// Load multiple rule files if provided
-	if strings.Contains(config.GetConf().RuleJsonFiles, ",") {
-		ruleFiles := strings.Split(config.GetConf().RuleJsonFiles, ",")
-		for _, ruleFile := range ruleFiles {
-			r, err := loadRuleFile(ruleFile)
-			if err != nil {
-				log.Printf("Failed to load rules: %v", err)
-				continue
+	// If rule files are provided, load them
+	if config.GetConf().RuleJsonFiles != "" {
+		// Load multiple rule files
+		if strings.Contains(config.GetConf().RuleJsonFiles, ",") {
+			ruleFiles := strings.Split(config.GetConf().RuleJsonFiles, ",")
+			for _, ruleFile := range ruleFiles {
+				r, err := loadRuleFile(ruleFile)
+				if err != nil {
+					log.Printf("Failed to load rules: %v", err)
+					continue
+				}
+				rules = append(rules, r...)
 			}
-			rules = append(rules, r...)
+			return rules
 		}
-		return rules
-	}
-
-	// Load the one rule file if provided
-	rs, err := loadRuleFile(config.GetConf().RuleJsonFiles)
-	if err != nil {
-		log.Printf("Failed to load rules: %v", err)
-		return nil
+		// Load the one rule file
+		rs, err := loadRuleFile(config.GetConf().RuleJsonFiles)
+		if err != nil {
+			log.Printf("Failed to load rules: %v", err)
+			return nil
+		}
+		rules = append(rules, rs...)
 	}
-	rules = append(rules, rs...)
 	return rules
 }
 

tool/preprocess/preprocess.go

@@ -59,7 +59,6 @@ const (
 	ReorderInitFile    = "reorder_init.go"
 	StdRulesPrefix     = "github.com/alibaba/opentelemetry-go-auto-instrumentation/pkg/"
 	StdRulesPath       = "github.com/alibaba/opentelemetry-go-auto-instrumentation/pkg/rules"
-	ApiPath            = "github.com/alibaba/opentelemetry-go-auto-instrumentation/pkg/api"
 )
 
 // @@ Change should sync with trampoline template

tool/preprocess/setup.go

@@ -29,6 +29,12 @@ import (
 	"github.com/dave/dst"
 )
 
+const (
+	ApiPackage     = "api"
+	ApiImportPath  = "github.com/alibaba/opentelemetry-go-auto-instrumentation/pkg/api"
+	ApiCallContext = "CallContext"
+)
+
 func initRuleDir() (err error) {
 	if exist, _ := util.PathExists(OtelRules); exist {
 		err = os.RemoveAll(OtelRules)
@@ -105,38 +111,47 @@ func (dp *DepProcessor) copyRules(pkgName string) (err error) {
 
 	return nil
 }
-
-func renameCallContext(astRoot *dst.File, bundle *resource.RuleBundle) {
-	pkgName := bundle.PackageName
-	// Find out if the target import path is aliased to another name
-	// if so, we need to rename api.CallContext to the alias name
-	// instead of the package name
-	for _, spec := range astRoot.Imports {
-		// Same import path and alias name is not null?
-		// One exception is the alias name is "_", we should ignore it
-		if shared.IsStringLit(spec.Path, bundle.ImportPath) &&
-			spec.Name != nil &&
-			spec.Name.Name != shared.IdentIgnore {
-			pkgName = spec.Name.Name
-			break
+func rectifyCallContext(astRoot *dst.File, bundle *resource.RuleBundle) {
+	// We write hook code by using api.CallContext as the first parameter, but
+	// the actual package name is not api. Given net/http package, the actual
+	// package name is http, so we should rectify the package name in the hook
+	// code to the correct package name. We did this by renaming the import path
+	// of api to the correct package name, and add an alias name for "api", this
+	// is required because CallContext is defined in the api package, and we can
+	// omit the package name before, but we can't do that now because of renaming
+	newAliasName := bundle.PackageName + util.RandomString(5)
+	alias := ApiPackage
+	spec := shared.FindImport(astRoot, ApiImportPath)
+	if spec != nil {
+		if spec.Name != nil {
+			alias = spec.Name.Name
 		}
 	}
+	// Check if the function has api.CallContext as the first parameter
+	// If so, rename it to the correct package name
 	for _, decl := range astRoot.Decls {
 		if f, ok := decl.(*dst.FuncDecl); ok {
+			foundCallContext := false
+
 			params := f.Type.Params.List
 			for _, param := range params {
 				if sele, ok := param.Type.(*dst.SelectorExpr); ok {
 					if x, ok := sele.X.(*dst.Ident); ok {
-						if x.Name == "api" && sele.Sel.Name == "CallContext" {
-							x.Name = pkgName
+						if x.Name == alias && sele.Sel.Name == ApiCallContext {
+							foundCallContext = true
+							x.Name = newAliasName
+							break
 						}
 					}
 				}
 			}
+			if foundCallContext {
+				spec.Path.Value = fmt.Sprintf("%q", bundle.ImportPath)
+				spec.Name = &dst.Ident{Name: newAliasName}
+			}
 		}
 	}
 }
-
 func makeHookPublic(astRoot *dst.File, bundle *resource.RuleBundle) {
 	// Only make hook public, keep it as it is if it's not a hook
 	hooks := make(map[string]bool)
@@ -168,45 +183,6 @@ func makeHookPublic(astRoot *dst.File, bundle *resource.RuleBundle) {
 	}
 }
 
-func renameImport(root *dst.File, oldPath, newPath string) bool {
-	// Find out if the old import and replace it with new one. Why we dont
-	// remove old import and add new one? Because we are not sure if the
-	// new import will be used, it's a compilation error if we import it
-	// but never use it.
-	for _, decl := range root.Decls {
-		if genDecl, ok := decl.(*dst.GenDecl); ok &&
-			genDecl.Tok == token.IMPORT {
-			for _, spec := range genDecl.Specs {
-				if importSpec, ok := spec.(*dst.ImportSpec); ok {
-					if importSpec.Path.Value == fmt.Sprintf("%q", oldPath) {
-						// In case the new import is already present, try to
-						// remove it first
-						oldSpec := shared.RemoveImport(root, newPath)
-						// Replace old with new one
-						importSpec.Path.Value = fmt.Sprintf("%q", newPath)
-						// Respect alias name of old import, if any
-						if oldSpec != nil {
-							importSpec.Name = oldSpec.Name
-
-							// Unless the alias name is "_", we should keep it
-							// For "_" alias, we should add additional normal
-							// variant for CallContext usage, i.e. keep both
-							// imports, one for existing usages, one for
-							// CallContext usage
-							if oldSpec.Name != nil &&
-								oldSpec.Name.Name == shared.IdentIgnore {
-								shared.AddImport(root, newPath)
-							}
-						}
-						return true
-					}
-				}
-			}
-		}
-	}
-	return false
-}
-
 func (dp *DepProcessor) copyRule(path, target string,
 	bundle *resource.RuleBundle) error {
 	text, err := util.ReadFile(path)
@@ -222,14 +198,11 @@ func (dp *DepProcessor) copyRule(path, target string,
 	astRoot.Name.Name = filepath.Base(filepath.Dir(target))
 
 	// Rename api.CallContext to correct package name if present
-	renameCallContext(astRoot, bundle)
+	rectifyCallContext(astRoot, bundle)
 
 	// Make hook functions public
 	makeHookPublic(astRoot, bundle)
 
-	// Rename "api" import to the correct package prefix
-	renameImport(astRoot, ApiPath, bundle.ImportPath)
-
 	// Copy used rule into project
 	_, err = shared.WriteAstToFile(astRoot, target)
 	if err != nil {

tool/shared/ast.go

@@ -290,6 +290,22 @@ func RemoveImport(root *dst.File, path string) *dst.ImportSpec {
 	return nil
 }
 
+func FindImport(root *dst.File, path string) *dst.ImportSpec {
+	for _, decl := range root.Decls {
+		if genDecl, ok := decl.(*dst.GenDecl); ok &&
+			genDecl.Tok == token.IMPORT {
+			for _, spec := range genDecl.Specs {
+				if importSpec, ok := spec.(*dst.ImportSpec); ok {
+					if importSpec.Path.Value == fmt.Sprintf("%q", path) {
+						return importSpec
+					}
+				}
+			}
+		}
+	}
+	return nil
+}
+
 func NewVarDecl(name string, paramTypes *dst.FieldList) *dst.GenDecl {
 	return &dst.GenDecl{
 		Tok: token.VAR,