feat(internal/injector): add method-call join point

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))
+}

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)
+		})
+	}
+}