diff --git a/internal/injector/aspect/join/method_call.go b/internal/injector/aspect/join/method_call.go new file mode 100644 index 00000000..9061fdbd --- /dev/null +++ b/internal/injector/aspect/join/method_call.go @@ -0,0 +1,180 @@ +// Unless explicitly stated otherwise all files in this repository are licensed +// under the Apache License Version 2.0. +// This product includes software developed at Datadog (https://www.datadoghq.com/). +// Copyright 2023-present Datadog, Inc. + +package join + +import ( + gocontext "context" + "errors" + "fmt" + "go/types" + + "github.com/DataDog/orchestrion/internal/fingerprint" + "github.com/DataDog/orchestrion/internal/injector/aspect/context" + "github.com/DataDog/orchestrion/internal/injector/aspect/may" + "github.com/DataDog/orchestrion/internal/injector/typed" + "github.com/DataDog/orchestrion/internal/yaml" + "github.com/dave/dst" + "github.com/goccy/go-yaml/ast" +) + +type ( + MethodCallMatch int + + methodCall struct { + Receiver typed.TypeName + Name string + Match MethodCallMatch + } +) + +const ( + // MethodCallMatchAny matches calls regardless of whether the receiver is a pointer or value. This is the default. + MethodCallMatchAny MethodCallMatch = iota + // MethodCallMatchPointerOnly matches only calls where the receiver is a pointer type. + MethodCallMatchPointerOnly + // MethodCallMatchValueOnly matches only calls where the receiver is a value type. + MethodCallMatchValueOnly +) + +func MethodCall(receiver typed.TypeName, name string, match MethodCallMatch) *methodCall { + return &methodCall{Receiver: receiver, Name: name, Match: match} +} + +func (m *methodCall) ImpliesImported() []string { + if path := m.Receiver.ImportPath; path != "" { + return []string{path} + } + return nil +} + +func (m *methodCall) PackageMayMatch(ctx *may.PackageContext) may.MatchType { + return ctx.PackageImports(m.Receiver.ImportPath) +} + +func (m *methodCall) FileMayMatch(ctx *may.FileContext) may.MatchType { + return ctx.FileContains(m.Name) +} + +func (m *methodCall) Matches(ctx context.AspectContext) bool { + call, ok := ctx.Node().(*dst.CallExpr) + if !ok { + return false + } + + selector, ok := call.Fun.(*dst.SelectorExpr) + if !ok || selector.Sel.Name != m.Name { + return false + } + + recvType := ctx.ResolveType(selector.X) + return m.matchesType(recvType) +} + +func (m *methodCall) matchesType(t types.Type) bool { + if t == nil { + return false + } + + switch m.Match { + case MethodCallMatchPointerOnly: + ptr, ok := t.(*types.Pointer) + if !ok { + return false + } + return m.matchesNamed(ptr.Elem()) + case MethodCallMatchValueOnly: + return m.matchesNamed(t) + default: // MethodCallMatchAny + if ptr, ok := t.(*types.Pointer); ok { + t = ptr.Elem() + } + return m.matchesNamed(t) + } +} + +func (m *methodCall) matchesNamed(t types.Type) bool { + named, ok := t.(*types.Named) + if !ok { + return false + } + obj := named.Obj() + return obj.Pkg() != nil && + obj.Pkg().Path() == m.Receiver.ImportPath && + obj.Name() == m.Receiver.Name +} + +func (m *methodCall) Hash(h *fingerprint.Hasher) error { + return h.Named("method-call", m.Receiver, fingerprint.String(m.Name), m.Match) +} + +func init() { + unmarshalers["method-call"] = func(ctx gocontext.Context, node ast.Node) (Point, error) { + var spec struct { + Receiver string `yaml:"receiver"` + Name string `yaml:"name"` + Match MethodCallMatch `yaml:"match"` + } + if err := yaml.NodeToValueContext(ctx, node, &spec); err != nil { + return nil, err + } + + if spec.Receiver == "" { + return nil, errors.New("method-call: missing required field 'receiver'") + } + if spec.Name == "" { + return nil, errors.New("method-call: missing required field 'name'") + } + + tn, err := typed.NewTypeName(spec.Receiver) + if err != nil { + return nil, fmt.Errorf("method-call: invalid receiver type %q: %w", spec.Receiver, err) + } + if tn.Pointer { + return nil, fmt.Errorf("method-call: receiver type must not include a pointer sigil (use match: pointer-only instead): %q", spec.Receiver) + } + + return MethodCall(tn, spec.Name, spec.Match), nil + } +} + +var _ yaml.NodeUnmarshalerContext = (*MethodCallMatch)(nil) + +func (m *MethodCallMatch) UnmarshalYAML(ctx gocontext.Context, node ast.Node) error { + var name string + if err := yaml.NodeToValueContext(ctx, node, &name); err != nil { + return err + } + + switch name { + case "any", "": + *m = MethodCallMatchAny + case "pointer-only": + *m = MethodCallMatchPointerOnly + case "value-only": + *m = MethodCallMatchValueOnly + default: + return fmt.Errorf("invalid method-call.match value: %q", name) + } + + return nil +} + +func (m MethodCallMatch) String() string { + switch m { + case MethodCallMatchAny: + return "any" + case MethodCallMatchPointerOnly: + return "pointer-only" + case MethodCallMatchValueOnly: + return "value-only" + default: + panic(fmt.Errorf("invalid MethodCallMatch(%d)", int(m))) + } +} + +func (m MethodCallMatch) Hash(h *fingerprint.Hasher) error { + return h.Named("method-call-match", fingerprint.Int(m)) +} diff --git a/internal/injector/aspect/join/method_call_test.go b/internal/injector/aspect/join/method_call_test.go new file mode 100644 index 00000000..dce34b34 --- /dev/null +++ b/internal/injector/aspect/join/method_call_test.go @@ -0,0 +1,205 @@ +// Unless explicitly stated otherwise all files in this repository are licensed +// under the Apache License Version 2.0. +// This product includes software developed at Datadog (https://www.datadoghq.com/). +// Copyright 2023-present Datadog, Inc. + +package join + +import ( + gocontext "context" + "go/types" + "testing" + + "github.com/goccy/go-yaml" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/DataDog/orchestrion/internal/fingerprint" + "github.com/DataDog/orchestrion/internal/injector/aspect/may" + "github.com/DataDog/orchestrion/internal/injector/typed" +) + +func newNamedType(pkgPath string, pkgName string, typeName string) *types.Named { + pkg := types.NewPackage(pkgPath, pkgName) + obj := types.NewTypeName(0, pkg, typeName, nil) + return types.NewNamed(obj, types.NewStruct(nil, nil), nil) +} + +func TestMethodCallMatchesType(t *testing.T) { + zapLogger := newNamedType("go.uber.org/zap", "zap", "Logger") + zapLoggerPtr := types.NewPointer(zapLogger) + otherLogger := newNamedType("example.com/other", "other", "Logger") + otherLoggerPtr := types.NewPointer(otherLogger) + + tn, err := typed.NewTypeName("go.uber.org/zap.Logger") + require.NoError(t, err) + + tests := []struct { + name string + match MethodCallMatch + typ types.Type + want bool + }{ + {name: "any: matches pointer", match: MethodCallMatchAny, typ: zapLoggerPtr, want: true}, + {name: "any: matches value", match: MethodCallMatchAny, typ: zapLogger, want: true}, + {name: "pointer-only: matches pointer", match: MethodCallMatchPointerOnly, typ: zapLoggerPtr, want: true}, + {name: "pointer-only: rejects value", match: MethodCallMatchPointerOnly, typ: zapLogger, want: false}, + {name: "value-only: matches value", match: MethodCallMatchValueOnly, typ: zapLogger, want: true}, + {name: "value-only: rejects pointer", match: MethodCallMatchValueOnly, typ: zapLoggerPtr, want: false}, + {name: "any: rejects different package", match: MethodCallMatchAny, typ: otherLoggerPtr, want: false}, + {name: "any: rejects nil", match: MethodCallMatchAny, typ: nil, want: false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + m := MethodCall(tn, "Info", tt.match) + assert.Equal(t, tt.want, m.matchesType(tt.typ)) + }) + } +} + +func TestMethodCallPackageMayMatch(t *testing.T) { + tn, err := typed.NewTypeName("go.uber.org/zap.Logger") + require.NoError(t, err) + m := MethodCall(tn, "Info", MethodCallMatchAny) + + importing := &may.PackageContext{ImportMap: map[string]string{"go.uber.org/zap": "zap.a"}} + notImporting := &may.PackageContext{ImportMap: map[string]string{"example.com/other": "other.a"}} + + assert.Equal(t, may.Match, m.PackageMayMatch(importing)) + assert.Equal(t, may.NeverMatch, m.PackageMayMatch(notImporting)) +} + +func TestMethodCallFileMayMatch(t *testing.T) { + tn, err := typed.NewTypeName("go.uber.org/zap.Logger") + require.NoError(t, err) + m := MethodCall(tn, "Info", MethodCallMatchAny) + + withMethod := &may.FileContext{FileContent: []byte(`package main; func f() { logger.Info("hi") }`)} + withoutMethod := &may.FileContext{FileContent: []byte(`package main; func f() { logger.Debug("hi") }`)} + + assert.Equal(t, may.Match, m.FileMayMatch(withMethod)) + assert.Equal(t, may.NeverMatch, m.FileMayMatch(withoutMethod)) +} + +func TestMethodCallImpliesImported(t *testing.T) { + tn, err := typed.NewTypeName("go.uber.org/zap.Logger") + require.NoError(t, err) + m := MethodCall(tn, "Info", MethodCallMatchAny) + assert.Equal(t, []string{"go.uber.org/zap"}, m.ImpliesImported()) +} + +func TestMethodCallHash(t *testing.T) { + tn1, _ := typed.NewTypeName("go.uber.org/zap.Logger") + tn2, _ := typed.NewTypeName("go.uber.org/zap.Logger") + tn3, _ := typed.NewTypeName("example.com/other.Logger") + + m1 := MethodCall(tn1, "Info", MethodCallMatchAny) + m2 := MethodCall(tn2, "Info", MethodCallMatchAny) + m3 := MethodCall(tn3, "Info", MethodCallMatchAny) + m4 := MethodCall(tn1, "Debug", MethodCallMatchAny) + m5 := MethodCall(tn1, "Info", MethodCallMatchPointerOnly) + + hash := func(m *methodCall) string { + h := fingerprint.New() + require.NoError(t, m.Hash(h)) + return h.Finish() + } + + assert.Equal(t, hash(m1), hash(m2), "identical method-calls must hash equally") + assert.NotEqual(t, hash(m1), hash(m3), "different receiver packages must hash differently") + assert.NotEqual(t, hash(m1), hash(m4), "different method names must hash differently") + assert.NotEqual(t, hash(m1), hash(m5), "different match modes must hash differently") +} + +func TestMethodCallUnmarshalYAML(t *testing.T) { + tests := []struct { + name string + yaml string + wantImport string + wantType string + wantMethod string + wantMatch MethodCallMatch + wantErr bool + }{ + { + name: "defaults to any", + yaml: `method-call: + receiver: "go.uber.org/zap.Logger" + name: Info`, + wantImport: "go.uber.org/zap", + wantType: "Logger", + wantMethod: "Info", + wantMatch: MethodCallMatchAny, + }, + { + name: "pointer-only", + yaml: `method-call: + receiver: "go.uber.org/zap.Logger" + name: Info + match: pointer-only`, + wantImport: "go.uber.org/zap", + wantType: "Logger", + wantMethod: "Info", + wantMatch: MethodCallMatchPointerOnly, + }, + { + name: "value-only", + yaml: `method-call: + receiver: "go.uber.org/zap.SugaredLogger" + name: Debugw + match: value-only`, + wantImport: "go.uber.org/zap", + wantType: "SugaredLogger", + wantMethod: "Debugw", + wantMatch: MethodCallMatchValueOnly, + }, + { + name: "pointer sigil in receiver is rejected", + yaml: `method-call: + receiver: "*go.uber.org/zap.Logger" + name: Info`, + wantErr: true, + }, + { + name: "missing receiver is rejected", + yaml: `method-call: + name: Info`, + wantErr: true, + }, + { + name: "missing name is rejected", + yaml: `method-call: + receiver: "go.uber.org/zap.Logger"`, + wantErr: true, + }, + } + + fn, ok := unmarshalers["method-call"] + require.True(t, ok, "method-call unmarshaler must be registered") + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var data map[string]any + err := yaml.Unmarshal([]byte(tt.yaml), &data) + require.NoError(t, err) + + node, err := yaml.ValueToNode(data["method-call"]) + require.NoError(t, err) + + result, err := fn(gocontext.Background(), node) + if tt.wantErr { + assert.Error(t, err) + return + } + require.NoError(t, err) + + m, ok := result.(*methodCall) + require.True(t, ok) + assert.Equal(t, tt.wantImport, m.Receiver.ImportPath) + assert.Equal(t, tt.wantType, m.Receiver.Name) + assert.Equal(t, tt.wantMethod, m.Name) + assert.Equal(t, tt.wantMatch, m.Match) + }) + } +} diff --git a/internal/injector/config/schema.json b/internal/injector/config/schema.json index 98d4a946..549f627e 100644 --- a/internal/injector/config/schema.json +++ b/internal/injector/config/schema.json @@ -372,6 +372,7 @@ { "$ref": "#/$defs/join-point/function" }, { "$ref": "#/$defs/join-point/function-call" }, { "$ref": "#/$defs/join-point/import-path" }, + { "$ref": "#/$defs/join-point/method-call" }, { "$ref": "#/$defs/join-point/not" }, { "$ref": "#/$defs/join-point/one-of" }, { "$ref": "#/$defs/join-point/package-name" }, @@ -628,6 +629,50 @@ { "import-path": "github.com/gorilla/mux" } ] }, + "method-call": { + "required": ["method-call"], + "unevaluatedProperties": false, + "properties": { + "method-call": { + "title": "Target method calls by receiver type", + "markdownDescription": "The `method-call` join point matches call expressions of the form `receiver.Method(args...)` where the receiver's type is a specific fully-qualified named type. Unlike `function-call`, it uses type information to distinguish between different types that expose the same method name.\n\nThe `match` field controls whether to match pointer receivers (`pointer-only`), value receivers (`value-only`), or both (`any`, the default).", + "type": "object", + "required": ["receiver", "name"], + "properties": { + "receiver": { + "description": "The fully qualified type name of the method receiver (without pointer sigil).", + "$ref": "#/$defs/go/qualified-identifier" + }, + "name": { + "description": "The name of the method to match.", + "$ref": "#/$defs/go/identifier" + }, + "match": { + "description": "Whether to match pointer receivers, value receivers, or both (default: any).", + "type": "string", + "enum": ["pointer-only", "value-only", "any"], + "default": "any" + } + }, + "additionalProperties": false + } + }, + "examples": [ + { + "method-call": { + "receiver": "go.uber.org/zap.Logger", + "name": "Info" + } + }, + { + "method-call": { + "receiver": "net/http.Client", + "name": "Do", + "match": "pointer-only" + } + } + ] + }, "not": { "required": ["not"], "unevaluatedProperties": false, diff --git a/internal/injector/testdata/injector/method-call-interface/config.yml b/internal/injector/testdata/injector/method-call-interface/config.yml new file mode 100644 index 00000000..b95bbb49 --- /dev/null +++ b/internal/injector/testdata/injector/method-call-interface/config.yml @@ -0,0 +1,31 @@ +%YAML 1.1 +--- +# Verifies that method-call matches calls on variables typed as an interface, +# but does NOT match calls on concrete types that happen to implement that interface. +aspects: + - id: io-reader-read + join-point: + method-call: + receiver: "io.Reader" + name: Read + advice: + - wrap-expression: + template: |- + io.MultiReader({{ .AST.Fun.X }}).Read( + {{ index .AST.Args 0 }}, + ) + +code: |- + package test + + import ( + "io" + "strings" + ) + + func doRead(r io.Reader, sr *strings.Reader, buf []byte) { + // Should be rewritten: receiver is typed as io.Reader (interface) + _, _ = r.Read(buf) + // Should NOT be rewritten: receiver is *strings.Reader (concrete type that implements io.Reader) + _, _ = sr.Read(buf) + } diff --git a/internal/injector/testdata/injector/method-call-interface/modified.go.snap b/internal/injector/testdata/injector/method-call-interface/modified.go.snap new file mode 100644 index 00000000..e1d9bf1f --- /dev/null +++ b/internal/injector/testdata/injector/method-call-interface/modified.go.snap @@ -0,0 +1,22 @@ +//line input.go:1:1 +package test + +import ( + "io" + "strings" +) + +func doRead(r io.Reader, sr *strings.Reader, buf []byte) { + // Should be rewritten: receiver is typed as io.Reader (interface) + _, _ = +//line :1 + io.MultiReader( +//line input.go:10 + r). +//line :1 + Read( +//line input.go:10 + buf) + // Should NOT be rewritten: receiver is *strings.Reader (concrete type that implements io.Reader) + _, _ = sr.Read(buf) +} diff --git a/internal/injector/testdata/injector/method-call/config.yml b/internal/injector/testdata/injector/method-call/config.yml new file mode 100644 index 00000000..ce546b90 --- /dev/null +++ b/internal/injector/testdata/injector/method-call/config.yml @@ -0,0 +1,37 @@ +%YAML 1.1 +--- +# Verifies that method-call matches calls on the correct receiver type and +# leaves calls on a different type with the same method name untouched. +aspects: + - id: http-client-do + join-point: + method-call: + receiver: "net/http.Client" + name: Do + match: pointer-only + advice: + - wrap-expression: + template: |- + http.DefaultClient.Do( + {{ index .AST.Args 0 }}, + ) + +code: |- + package test + + import ( + "net/http" + ) + + type myClient struct{} + + func (c *myClient) Do(req *http.Request) (*http.Response, error) { + return nil, nil + } + + func doRequest(client *http.Client, other *myClient, req *http.Request) { + // Should be rewritten: receiver is *http.Client (pointer-only match) + _, _ = client.Do(req) + // Should NOT be rewritten: receiver is *myClient, not *http.Client + _, _ = other.Do(req) + } diff --git a/internal/injector/testdata/injector/method-call/modified.go.snap b/internal/injector/testdata/injector/method-call/modified.go.snap new file mode 100644 index 00000000..0a796eb9 --- /dev/null +++ b/internal/injector/testdata/injector/method-call/modified.go.snap @@ -0,0 +1,23 @@ +//line input.go:1:1 +package test + +import ( + "net/http" +) + +type myClient struct{} + +func (c *myClient) Do(req *http.Request) (*http.Response, error) { + return nil, nil +} + +func doRequest(client *http.Client, other *myClient, req *http.Request) { + // Should be rewritten: receiver is *http.Client (pointer-only match) + _, _ = +//line :1 + http.DefaultClient.Do( +//line input.go:15 + req) + // Should NOT be rewritten: receiver is *myClient, not *http.Client + _, _ = other.Do(req) +} diff --git a/main_test.go b/main_test.go index b7186d64..28115c59 100644 --- a/main_test.go +++ b/main_test.go @@ -10,6 +10,7 @@ import ( "encoding/json" "fmt" "io" + "io/fs" "net/http" "os" "os/exec" @@ -20,6 +21,7 @@ import ( "testing" "github.com/stretchr/testify/require" + "golang.org/x/mod/modfile" ) func TestBuildFromModuleSubdirectory(t *testing.T) { @@ -59,13 +61,14 @@ type benchCase interface { var benchCases = map[string]func(b *testing.B) benchCase{ "DataDog:orchestrion": benchmarkOrchestrion, // normal build - "traefik:traefik": benchmarkGithub("traefik", "traefik", "./...", false), - "go-delve:delve": benchmarkGithub("go-delve", "delve", "./...", false), - "jlegrone:tctx": benchmarkGithub("jlegrone", "tctx", "./...", false), - "tinylib:msgp": benchmarkGithub("tinylib", "msgp", "./...", false), + "traefik:traefik": benchmarkGithub("traefik", "traefik", "", "./...", false), + "go-delve:delve": benchmarkGithub("go-delve", "delve", "", "./...", false), + "jlegrone:tctx": benchmarkGithub("jlegrone", "tctx", "", "./...", false), + "tinylib:msgp": benchmarkGithub("tinylib", "msgp", "", "./...", false), + "etcd-io:etcd": benchmarkGithub("etcd-io", "etcd", "server", "./...", false), // test packages - "gin-gonic:gin.test": benchmarkGithub("gin-gonic", "gin", "./...", true), - "jlegrone:tctx.test": benchmarkGithub("jlegrone", "tctx", "./...", true), + "gin-gonic:gin.test": benchmarkGithub("gin-gonic", "gin", "", "./...", true), + "jlegrone:tctx.test": benchmarkGithub("jlegrone", "tctx", "", "./...", true), } func Benchmark(b *testing.B) { @@ -93,7 +96,10 @@ type benchGithub struct { harness } -func benchmarkGithub(owner string, repo string, build string, testbuild bool) func(b *testing.B) benchCase { +// benchmarkGithub builds a benchmark case for a github repo. If subdir is +// non-empty, setup and build run against that sub-module (etcd-style monorepos) +// rather than the repo root. +func benchmarkGithub(owner string, repo string, subdir string, build string, testbuild bool) func(b *testing.B) benchCase { return func(b *testing.B) benchCase { tc := &benchGithub{harness{build: build, testbuild: testbuild}} @@ -101,8 +107,14 @@ func benchmarkGithub(owner string, repo string, build string, testbuild bool) fu b.Logf("Latest release is %s/%s@%s", owner, repo, tag) tc.gitCloneGithub(b, owner, repo, tag) + if subdir != "" { + tc.dir = filepath.Join(tc.dir, subdir) + } tc.exec(b, "go", "mod", "download") tc.exec(b, "go", "mod", "edit", "-replace=github.com/DataDog/orchestrion="+rootDir) + if replace := os.Getenv("DD_TRACE_GO_REPLACE"); replace != "" { + applyDDTraceGoReplaces(b, &tc.runner, replace) + } if stat, err := os.Stat(filepath.Join(tc.dir, "vendor")); err == nil && stat.IsDir() { // If there's a vendor dir, we need to update the `modules.txt` in there to reflect the replacement. tc.exec(b, "go", "mod", "vendor") @@ -248,6 +260,51 @@ func getGithubToken() (string, bool) { return strings.TrimSpace(bytes.String()), true } +// applyDDTraceGoReplaces walks replace, finds every go.mod whose module path +// starts with github.com/DataDog/dd-trace-go, and adds a corresponding replace +// directive to the target's go.mod. This lets the benchmarks build against a +// local checkout of dd-trace-go (including branches where transitive contribs +// reference an unpublished placeholder version like v2.10.0-dev). +func applyDDTraceGoReplaces(b *testing.B, r *runner, replace string) { + b.Helper() + + err := filepath.WalkDir(replace, func(path string, d fs.DirEntry, err error) error { + if err != nil { + return err + } + if d.IsDir() { + if path != replace { + name := d.Name() + if strings.HasPrefix(name, ".") || strings.HasPrefix(name, "_") { + return filepath.SkipDir + } + } + return nil + } + if d.Name() != "go.mod" { + return nil + } + data, err := os.ReadFile(path) + if err != nil { + return err + } + mod, err := modfile.Parse(path, data, nil) + if err != nil { + return err + } + if mod.Module == nil { + return nil + } + modPath := mod.Module.Mod.Path + if !strings.HasPrefix(modPath, "github.com/DataDog/dd-trace-go") { + return nil + } + r.exec(b, "go", "mod", "edit", "-replace="+modPath+"="+filepath.Dir(path)) + return nil + }) + require.NoError(b, err) +} + func (h *harness) gitCloneGithub(b *testing.B, owner string, repo string, tag string) string { b.Helper()