add --trap-stdlib to support mock most stdlib functions

xhd2015 · May 18, 2024 · ced25b5 · ced25b5
1 parent 7742d22
commit ced25b5
Show file tree

Hide file tree

Showing 28 changed files with 244 additions and 101 deletions.
diff --git a/cmd/xgo/exec_tool/debug.go b/cmd/xgo/exec_tool/debug.go
@@ -20,6 +20,7 @@ func getDebugEnv(xgoCompilerEnableEnv string) map[string]string {
 		"GOCACHE":                        os.Getenv("GOCACHE"),
 		XGO_MAIN_MODULE:                  os.Getenv(XGO_MAIN_MODULE),
 		XGO_COMPILE_PKG_DATA_DIR:         os.Getenv(XGO_COMPILE_PKG_DATA_DIR),
+		XGO_STD_LIB_TRAP_DEFAULT_ALLOW:   os.Getenv(XGO_STD_LIB_TRAP_DEFAULT_ALLOW),
 		"GOROOT":                         "../..",
 		"PATH":                           "../../bin:${env:PATH}",
 		"XGO_COMPILER_ENABLE":            xgoCompilerEnableEnv,

diff --git a/cmd/xgo/exec_tool/env.go b/cmd/xgo/exec_tool/env.go
@@ -15,3 +15,5 @@ const XGO_TOOLCHAIN_VERSION_NUMBER = "XGO_TOOLCHAIN_VERSION_NUMBER"
 const XGO_MAIN_MODULE = "XGO_MAIN_MODULE"
 
 const XGO_COMPILE_PKG_DATA_DIR = "XGO_COMPILE_PKG_DATA_DIR"
+
+const XGO_STD_LIB_TRAP_DEFAULT_ALLOW = "XGO_STD_LIB_TRAP_DEFAULT_ALLOW"
diff --git a/cmd/xgo/main.go b/cmd/xgo/main.go
@@ -151,6 +151,7 @@ func handleBuild(cmd string, args []string) error {
 	dumpIR := opts.dumpIR
 	dumpAST := opts.dumpAST
 	stackTrace := opts.stackTrace
+	trapStdlib := opts.trapStdlib
 
 	if cmdExec && len(remainArgs) == 0 {
 		return fmt.Errorf("exec requires command")
@@ -241,8 +242,11 @@ func handleBuild(cmd string, args []string) error {
 	// gcflags can cause the build cache to invalidate
 	// so separate them with normal one
 	buildCacheSuffix := ""
+	if trapStdlib {
+		buildCacheSuffix += "-trapstd"
+	}
 	if len(gcflags) > 0 {
-		buildCacheSuffix = "-gcflags"
+		buildCacheSuffix += "-gcflags"
 	}
 	buildCacheDir := filepath.Join(instrumentDir, "build-cache"+buildCacheSuffix)
 	revisionFile := filepath.Join(instrumentDir, "xgo-revision.txt")
@@ -461,20 +465,33 @@ func handleBuild(cmd string, args []string) error {
 		execCmd.Env = append(execCmd.Env, "XGO_TOOLCHAIN_VERSION="+VERSION)
 		execCmd.Env = append(execCmd.Env, "XGO_TOOLCHAIN_REVISION="+REVISION)
 		execCmd.Env = append(execCmd.Env, "XGO_TOOLCHAIN_VERSION_NUMBER="+strconv.FormatInt(NUMBER, 10))
-		if dumpIR != "" {
-			execCmd.Env = append(execCmd.Env, "XGO_DEBUG_DUMP_IR="+dumpIR)
-			execCmd.Env = append(execCmd.Env, "XGO_DEBUG_DUMP_IR_FILE="+tmpIRFile)
-		}
-		if dumpAST != "" {
-			execCmd.Env = append(execCmd.Env, "XGO_DEBUG_DUMP_AST="+dumpAST)
-			execCmd.Env = append(execCmd.Env, "XGO_DEBUG_DUMP_AST_FILE="+tmpASTFile)
-		}
+
+		// IR
+		execCmd.Env = append(execCmd.Env, "XGO_DEBUG_DUMP_IR="+dumpIR)
+		execCmd.Env = append(execCmd.Env, "XGO_DEBUG_DUMP_IR_FILE="+tmpIRFile)
+
+		// AST
+		execCmd.Env = append(execCmd.Env, "XGO_DEBUG_DUMP_AST="+dumpAST)
+		execCmd.Env = append(execCmd.Env, "XGO_DEBUG_DUMP_AST_FILE="+tmpASTFile)
+
+		// vscode debug
+		var xgoDebugVscode string
 		if vscodeDebugFile != "" {
-			execCmd.Env = append(execCmd.Env, "XGO_DEBUG_VSCODE="+vscodeDebugFile+vscodeDebugFileSuffix)
+			xgoDebugVscode = vscodeDebugFile + vscodeDebugFileSuffix
 		}
+		execCmd.Env = append(execCmd.Env, "XGO_DEBUG_VSCODE="+xgoDebugVscode)
+
+		// stack trace
 		if stackTrace != "" {
 			execCmd.Env = append(execCmd.Env, "XGO_STACK_TRACE="+stackTrace)
 		}
+
+		// trap stdlib
+		var trapStdlibEnv string
+		if trapStdlib {
+			trapStdlibEnv = "true"
+		}
+		execCmd.Env = append(execCmd.Env, "XGO_STD_LIB_TRAP_DEFAULT_ALLOW="+trapStdlibEnv)
 	}
 	logDebug("command env: %v", execCmd.Env)
 	execCmd.Stdout = os.Stdout

diff --git a/cmd/xgo/option.go b/cmd/xgo/option.go
@@ -50,6 +50,9 @@ type options struct {
 	overlay string
 	modfile string
 
+	// --trap-stdlib
+	trapStdlib bool
+
 	// xgo test --trace
 
 	// --strace, --strace=on, --strace=off
@@ -97,6 +100,7 @@ func parseOptions(args []string) (*options, error) {
 	var overlay string
 	var modfile string
 	var stackTrace string
+	var trapStdlib bool
 
 	var remainArgs []string
 	nArg := len(args)
@@ -254,6 +258,11 @@ func parseOptions(args []string) (*options, error) {
 			continue
 		}
 
+		if arg == "--trap-stdlib" {
+			trapStdlib = true
+			continue
+		}
+
 		if isDevelopment && arg == "--debug-with-dlv" {
 			debugWithDlv = true
 			continue
@@ -340,6 +349,7 @@ func parseOptions(args []string) (*options, error) {
 		overlay:    overlay,
 		modfile:    modfile,
 		stackTrace: stackTrace,
+		trapStdlib: trapStdlib,
 
 		remainArgs: remainArgs,
 	}, nil

diff --git a/cmd/xgo/runtime_gen/core/func.go b/cmd/xgo/runtime_gen/core/func.go
@@ -51,6 +51,9 @@ type FuncInfo struct {
 	// is this a closure?
 	Closure bool
 
+	// is this function from stdlib
+	Stdlib bool
+
 	// source info
 	File string
 	Line int

diff --git a/cmd/xgo/runtime_gen/core/version.go b/cmd/xgo/runtime_gen/core/version.go
@@ -7,8 +7,8 @@ import (
 )
 
 const VERSION = "1.0.35"
-const REVISION = "3431198dea60ee1a41ed69b93034c5919aba32b5+1"
-const NUMBER = 218
+const REVISION = "7742d22be893d673e6de160b6acc01b39502271e+1"
+const NUMBER = 219
 
 // these fields will be filled by compiler
 const XGO_VERSION = ""

diff --git a/cmd/xgo/runtime_gen/functab/functab.go b/cmd/xgo/runtime_gen/functab/functab.go
@@ -169,6 +169,13 @@ func registerFuncInfo(fnInfo interface{}) {
 	name := rv.FieldByName("Name").String()
 	interface_ := rv.FieldByName("Interface").Bool()
 	generic := rv.FieldByName("Generic").Bool()
+
+	var stdlib bool
+	stdlibField := rv.FieldByName("Stdlib")
+	if stdlibField.IsValid() {
+		stdlib = stdlibField.Bool()
+	}
+
 	f := rv.FieldByName("Fn").Interface()
 
 	var firstArgCtx bool
@@ -222,6 +229,7 @@ func registerFuncInfo(fnInfo interface{}) {
 		Interface: interface_,
 		Generic:   generic,
 		Closure:   closure,
+		Stdlib:    stdlib,
 
 		File: file,
 		Line: line,

diff --git a/cmd/xgo/runtime_gen/trace/stack.go b/cmd/xgo/runtime_gen/trace/stack.go
@@ -192,16 +192,20 @@ func ExportFuncInfo(c *core.FuncInfo, opts *ExportOptions) *FuncInfoExport {
 		return nil
 	}
 	return &FuncInfoExport{
+		Kind:         FuncKind(c.Kind.String()),
 		Pkg:          c.Pkg,
 		IdentityName: c.IdentityName,
 		Name:         c.Name,
 		RecvType:     c.RecvType,
 		RecvPtr:      c.RecvPtr,
 
-		Generic:  c.Generic,
-		RecvName: c.RecvName,
-		ArgNames: c.ArgNames,
-		ResNames: c.ResNames,
+		Interface: c.Interface,
+		Generic:   c.Generic,
+		Closure:   c.Closure,
+		Stdlib:    c.Stdlib,
+		RecvName:  c.RecvName,
+		ArgNames:  c.ArgNames,
+		ResNames:  c.ResNames,
 
 		FirstArgCtx:   c.FirstArgCtx,
 		LastResultErr: c.LastResultErr,

diff --git a/cmd/xgo/runtime_gen/trace/stack_export.go b/cmd/xgo/runtime_gen/trace/stack_export.go
@@ -2,6 +2,15 @@ package trace
 
 import "time"
 
+type FuncKind string
+
+const (
+	FuncKind_Func   FuncKind = "func"
+	FuncKind_Var    FuncKind = "var"
+	FuncKind_VarPtr FuncKind = "var_ptr"
+	FuncKind_Const  FuncKind = "const"
+)
+
 type RootExport struct {
 	// current executed function
 	Begin    time.Time
@@ -24,13 +33,18 @@ type StackExport struct {
 
 type FuncInfoExport struct {
 	// FullName string
+	Kind         FuncKind
 	Pkg          string
 	IdentityName string
 	Name         string
 	RecvType     string
 	RecvPtr      bool
 
-	Generic bool
+	// interface method?
+	Interface bool
+	Generic   bool
+	Closure   bool
+	Stdlib    bool
 
 	File string
 	Line int

diff --git a/cmd/xgo/runtime_gen/trap/trap.go b/cmd/xgo/runtime_gen/trap/trap.go
@@ -184,6 +184,7 @@ func trap(f *core.FuncInfo, pc uintptr, recv interface{}, args []interface{}, re
 
 	var resetFlag bool
 	parent := r.top
+
 	if !r.intercepting {
 		resetFlag = true
 		r.intercepting = true

diff --git a/cmd/xgo/trace/stack_export.go b/cmd/xgo/trace/stack_export.go
diff --git a/cmd/xgo/version.go b/cmd/xgo/version.go
@@ -3,8 +3,8 @@ package main
 import "fmt"
 
 const VERSION = "1.0.35"
-const REVISION = "3431198dea60ee1a41ed69b93034c5919aba32b5+1"
-const NUMBER = 218
+const REVISION = "7742d22be893d673e6de160b6acc01b39502271e+1"
+const NUMBER = 219
 
 func getRevision() string {
 	revSuffix := ""

diff --git a/patch/ctxt/ctx.go b/patch/ctxt/ctx.go
@@ -2,7 +2,6 @@ package ctxt
 
 import (
 	"cmd/compile/internal/base"
-	"os"
 	"strings"
 )
 
@@ -11,17 +10,14 @@ const XgoRuntimePkg = XgoModule + "/runtime"
 const XgoRuntimeCorePkg = XgoModule + "/runtime/core"
 const XgoRuntimeTracePkg = XgoModule + "/runtime/trace"
 
-var XgoMainModule = os.Getenv("XGO_MAIN_MODULE")
-var XgoCompilePkgDataDir = os.Getenv("XGO_COMPILE_PKG_DATA_DIR")
-
 const XgoLinkTrapVarForGenerated = "__xgo_link_trap_var_for_generated"
 
 func SkipPackageTrap() bool {
 	pkgPath := GetPkgPath()
 	if pkgPath == "" {
 		return true
 	}
-	if strings.HasPrefix(pkgPath, "runtime/") || strings.HasPrefix(pkgPath, "internal/") {
+	if pkgPath == "runtime" || strings.HasPrefix(pkgPath, "runtime/") || strings.HasPrefix(pkgPath, "internal/") {
 		return true
 	}
 	if base.Flag.Std {
@@ -37,6 +33,12 @@ func SkipPackageTrap() bool {
 		// because generic instantiation happens in other package, so this
 		// func may be a foreigner.
 
+		if XgoStdTrapDefaultAllow {
+			if _, ok := stdBlocklist[pkgPath]["*"]; ok {
+				return true
+			}
+			return false
+		}
 		// allow http
 		if _, ok := stdWhitelist[pkgPath]; ok {
 			return false

diff --git a/patch/ctxt/env.go b/patch/ctxt/env.go
@@ -0,0 +1,9 @@
+package ctxt
+
+import "os"
+
+var XgoMainModule = os.Getenv("XGO_MAIN_MODULE")
+var XgoCompilePkgDataDir = os.Getenv("XGO_COMPILE_PKG_DATA_DIR")
+
+// enabled via: --trap-stdlib
+var XgoStdTrapDefaultAllow = os.Getenv("XGO_STD_LIB_TRAP_DEFAULT_ALLOW") == "true"
diff --git a/patch/ctxt/stdlib.go b/patch/ctxt/stdlib.go
@@ -54,7 +54,39 @@ var stdWhitelist = map[string]map[string]bool{
 	},
 }
 
+// effective when XgoStdTrapDefaultAllow is true
+//
+//	"net":map[string]bool{
+//	   "*": true -> disable all
+//	}
+var stdBlocklist = map[string]map[string]bool{
+	"syscall": map[string]bool{
+		"*": true,
+	},
+	"reflect": map[string]bool{
+		"*": true,
+	},
+	"sync": map[string]bool{
+		"*": true,
+	},
+	"sync/atomic": map[string]bool{
+		"*": true,
+	},
+	"testing": map[string]bool{
+		"*": true,
+	},
+	"unsafe": map[string]bool{
+		"*": true,
+	},
+}
+
 func allowStdFunc(pkgPath string, funcName string) bool {
+	if XgoStdTrapDefaultAllow {
+		if stdBlocklist[pkgPath]["*"] || stdBlocklist[pkgPath][funcName] {
+			return false
+		}
+		return true
+	}
 	if stdWhitelist[pkgPath][funcName] {
 		return true
 	}

diff --git a/patch/syntax/func_stub_def.go b/patch/syntax/func_stub_def.go
@@ -9,6 +9,7 @@ const expected__xgo_stub_def = `struct {
 	Interface    bool
 	Generic      bool
 	Closure      bool // is the given function a closure
+	Stdlib       bool
 	RecvTypeName string
 	RecvPtr      bool
 	Name         string

diff --git a/patch/syntax/helper_code.go b/patch/syntax/helper_code.go
@@ -14,6 +14,7 @@ type __xgo_local_func_stub struct {
 	Interface    bool
 	Generic      bool
 	Closure      bool // is the given function a closure
+	Stdlib       bool
 	RecvTypeName string
 	RecvPtr      bool
 	Name         string
Original file line number	Diff line number	Diff line change
Expand Up		@@ -15,3 +15,5 @@ const XGO_TOOLCHAIN_VERSION_NUMBER = "XGO_TOOLCHAIN_VERSION_NUMBER"
		const XGO_MAIN_MODULE = "XGO_MAIN_MODULE"

		const XGO_COMPILE_PKG_DATA_DIR = "XGO_COMPILE_PKG_DATA_DIR"

		const XGO_STD_LIB_TRAP_DEFAULT_ALLOW = "XGO_STD_LIB_TRAP_DEFAULT_ALLOW"