From 62d73330f61c2fff564e43f51b5a368f1f60f202 Mon Sep 17 00:00:00 2001 From: Romain Marcadier Date: Tue, 27 Aug 2024 10:00:00 +0200 Subject: [PATCH] refactor(injector): new API with better performance (#234) Create a new `injector.Injector` API that does not rely on `decorator.Load` (internally using `packages.Load`), instead using the basic `go/types` API to type-check the AST nodes in order to obtain the `Uses` map that is sufficient to build an import-managing `decorator.Decorator` instance. The package name resolution is done by parsing type information from the archives mentioned in the `importcfg` file. This change removes unnecessary compilation of un-instrumented archives which are not useful, and hence saves time and disk space. --- The new API also does not consider `PreserveLineInfo` to be optional (this was never exposed to end-users anyway), so a bunch of test reference files have changed to now include line directives. --------- Co-authored-by: Eliott Bouhana <47679741+eliottness@users.noreply.github.com> --- internal/injector/builtin/generated.go | 25 -- internal/injector/builtin/generator/main.go | 38 -- internal/injector/builtin/integration_test.go | 88 +++-- internal/injector/check.go | 34 ++ internal/injector/imports.go | 115 ++++++ internal/injector/injector.go | 356 ++++++------------ internal/injector/injector_test.go | 73 ++-- internal/injector/parse.go | 137 +++++++ internal/injector/parse_test.go | 72 ++++ internal/injector/restorer.go | 68 ++++ .../access-return-value/expected.diff | 9 +- .../chi5-newroute-dotimport/config.yml | 2 - .../injector/chi5-newroute/config.yml | 2 - .../testdata/injector/database-sql/config.yml | 2 - .../testdata/injector/directive/config.yml | 2 - .../testdata/injector/grpc/config.yml | 2 - .../http-anonymous-handler/expected.diff | 16 +- .../injector/http-default/expected.diff | 19 +- .../injector/http-line-info/config.yml | 2 - .../injector/http-no-arg-name/expected.diff | 18 +- .../testdata/injector/http-server/config.yml | 2 - .../injector/name-conflict/config.yml | 3 - .../testdata/injector/redigo/config.yml | 2 - internal/injector/write.go | 55 +++ internal/toolexec/aspect/oncompile.go | 63 ++-- internal/toolexec/importcfg/lookup.go | 33 ++ internal/toolexec/proxy/compile.go | 27 +- internal/toolexec/proxy/compile_test.go | 29 +- 28 files changed, 824 insertions(+), 470 deletions(-) create mode 100644 internal/injector/check.go create mode 100644 internal/injector/imports.go create mode 100644 internal/injector/parse.go create mode 100644 internal/injector/parse_test.go create mode 100644 internal/injector/restorer.go create mode 100644 internal/injector/write.go create mode 100644 internal/toolexec/importcfg/lookup.go diff --git a/internal/injector/builtin/generated.go b/internal/injector/builtin/generated.go index dcb66ae22..ae815a6ba 100644 --- a/internal/injector/builtin/generated.go +++ b/internal/injector/builtin/generated.go @@ -684,31 +684,6 @@ var Aspects = [...]aspect.Aspect{ }, } -// RestorerMap is a set of import path to name mappings for packages that would be incorrectly named by restorer.Guess -var RestorerMap = map[string]string{ - "gopkg.in/DataDog/dd-trace-go.v1/contrib/IBM/sarama.v1": "sarama", - "gopkg.in/DataDog/dd-trace-go.v1/contrib/cloud.google.com/go/pubsub.v1": "pubsub", - "gopkg.in/DataDog/dd-trace-go.v1/contrib/confluentinc/confluent-kafka-go/kafka.v2": "kafka", - "gopkg.in/DataDog/dd-trace-go.v1/contrib/dimfeld/httptreemux.v5": "httptreemux", - "gopkg.in/DataDog/dd-trace-go.v1/contrib/elastic/go-elasticsearch.v6": "elastic", - "gopkg.in/DataDog/dd-trace-go.v1/contrib/emicklei/go-restful": "restful", - "gopkg.in/DataDog/dd-trace-go.v1/contrib/emicklei/go-restful.v3": "restful", - "gopkg.in/DataDog/dd-trace-go.v1/contrib/go-chi/chi.v5": "chi", - "gopkg.in/DataDog/dd-trace-go.v1/contrib/go-pg/pg.v10": "pg", - "gopkg.in/DataDog/dd-trace-go.v1/contrib/go-redis/redis.v7": "redis", - "gopkg.in/DataDog/dd-trace-go.v1/contrib/go-redis/redis.v8": "redis", - "gopkg.in/DataDog/dd-trace-go.v1/contrib/gofiber/fiber.v2": "fiber", - "gopkg.in/DataDog/dd-trace-go.v1/contrib/google.golang.org/grpc.v12": "grpc", - "gopkg.in/DataDog/dd-trace-go.v1/contrib/gopkg.in/jinzhu/gorm.v1": "gorm", - "gopkg.in/DataDog/dd-trace-go.v1/contrib/gorm.io/gorm.v1": "gorm", - "gopkg.in/DataDog/dd-trace-go.v1/contrib/graph-gophers/graphql-go": "graphql", - "gopkg.in/DataDog/dd-trace-go.v1/contrib/jackc/pgx.v5": "pgx", - "gopkg.in/DataDog/dd-trace-go.v1/contrib/labstack/echo.v4": "echo", - "gopkg.in/DataDog/dd-trace-go.v1/contrib/redis/go-redis.v9": "redis", - "gopkg.in/DataDog/dd-trace-go.v1/contrib/segmentio/kafka.go.v0": "kafka", - "gopkg.in/DataDog/dd-trace-go.v1/contrib/valyala/fasthttp.v1": "fasthttp", -} - // InjectedPaths is a set of import paths that may be injected by built-in aspects. This list is used to ensure proper // invalidation of cached artifacts when injected dependencies change. var InjectedPaths = [...]string{ diff --git a/internal/injector/builtin/generator/main.go b/internal/injector/builtin/generator/main.go index 2b2c6b374..aa1427b4e 100644 --- a/internal/injector/builtin/generator/main.go +++ b/internal/injector/builtin/generator/main.go @@ -19,7 +19,6 @@ import ( "strings" "github.com/dave/jennifer/jen" - "golang.org/x/tools/go/packages" "gopkg.in/yaml.v3" ) @@ -162,29 +161,6 @@ func main() { g.Empty().Line() }) - file.Comment("RestorerMap is a set of import path to name mappings for packages that would be incorrectly named by restorer.Guess") - file.Var().Id("RestorerMap").Op("=").Map(jen.String()).String().ValuesFunc(func(g *jen.Group) { - pkgs, err := packages.Load(&packages.Config{}, "gopkg.in/DataDog/dd-trace-go.v1/...") - if err != nil { - log.Fatalf("Failed to load packages: %v\n", err) - } - sort.Sort(sortable(pkgs)) - - for _, pkg := range pkgs { - if strings.HasSuffix(pkg.PkgPath, "/internal") || strings.Contains(pkg.PkgPath, "/internal/") { - // We don't care about internal packages here (at least for now) - continue - } - if strings.HasSuffix(pkg.PkgPath, "/"+pkg.Name) { - // We don't care about packages for which `restorer.Guess` would infer the right name. - continue - } - g.Line().Lit(pkg.PkgPath).Op(":").Lit(pkg.Name) - } - - g.Line().Empty() - }) - file.Comment("InjectedPaths is a set of import paths that may be injected by built-in aspects. This list is used to ensure proper") file.Comment("invalidation of cached artifacts when injected dependencies change.") file.Var().Id("InjectedPaths").Op("=").Index(jen.Op("...")).String().ValuesFunc(func(g *jen.Group) { @@ -213,20 +189,6 @@ func main() { } } -type sortable []*packages.Package - -func (s sortable) Len() int { - return len(s) -} - -func (s sortable) Less(i, j int) bool { - return s[i].PkgPath < s[j].PkgPath -} - -func (s sortable) Swap(i, j int) { - s[i], s[j] = s[j], s[i] -} - func writeAll(w io.Writer, data []byte) { for len(data) > 0 { n, err := w.Write(data) diff --git a/internal/injector/builtin/integration_test.go b/internal/injector/builtin/integration_test.go index 5f702bbbc..c4337802b 100644 --- a/internal/injector/builtin/integration_test.go +++ b/internal/injector/builtin/integration_test.go @@ -10,6 +10,7 @@ import ( "errors" "flag" "fmt" + "io" "os" "path/filepath" "runtime" @@ -20,6 +21,7 @@ import ( "github.com/datadog/orchestrion/internal/version" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "golang.org/x/tools/go/packages" ) var ( @@ -44,51 +46,71 @@ func Test(t *testing.T) { pkgDir := filepath.Join(samplesDir, dir) + testLookup := func(path string) (io.ReadCloser, error) { + pkgs, err := packages.Load( + &packages.Config{ + Mode: packages.NeedExportFile, + Dir: pkgDir, + Logf: t.Logf, + }, + path, + ) + if err != nil { + return nil, err + } + file := pkgs[0].ExportFile + if file == "" { + return nil, fmt.Errorf("no export data for %s", path) + } + return os.Open(file) + } + tmp := t.TempDir() - inj, err := injector.New(pkgDir, injector.Options{ + inj := injector.Injector{ Aspects: builtin.Aspects[:], - Dir: pkgDir, ModifiedFile: func(filename string) string { return filepath.Join(tmp, filepath.Base(filename)) }, - PreserveLineInfo: true, - }) - require.NoError(t, err) + RootConfig: map[string]string{"httpmode": "wrap"}, + ImportPath: fmt.Sprintf("github.com/datadog/orchestrion/samples/%s", dir), + Lookup: testLookup, + } files, err := filepath.Glob(filepath.Join(pkgDir, "*.go")) require.NoError(t, err) - for _, filename := range files { - t.Run(filepath.Base(filename), func(t *testing.T) { - res, err := inj.InjectFile(filename, map[string]string{"httpmode": "wrap"}) - require.NoError(t, err) - - referenceFile := filepath.Join(referenceDir, dir, filepath.Base(filename)) + ".snap" - if !res.Modified { - _, err := os.Stat(referenceFile) - if updateSnapshots && err == nil { - require.NoError(t, os.Remove(referenceFile)) - } - require.ErrorIs(t, err, os.ErrNotExist) - return - } - - data, err := os.ReadFile(res.Filename) - require.NoError(t, err) - data = bytes.ReplaceAll(data, []byte(samplesDir), []byte("samples")) - data = bytes.ReplaceAll(data, []byte(fmt.Sprintf("%q", version.Tag)), []byte("\"\"")) + results, err := inj.InjectFiles(files) + require.NoError(t, err) - reference, err := os.ReadFile(referenceFile) - if updateSnapshots && errors.Is(err, os.ErrNotExist) { - require.NoError(t, os.MkdirAll(filepath.Dir(referenceFile), 0o755)) - require.NoError(t, os.WriteFile(referenceFile, data, 0o644)) - } - require.NoError(t, err) + for _, filename := range files { + referenceFile := filepath.Join(referenceDir, dir, filepath.Base(filename)) + ".snap" - if !assert.Equal(t, string(reference), string(data)) && updateSnapshots { - require.NoError(t, os.WriteFile(referenceFile, data, 0o644)) + res, modified := results[filename] + if !modified { + _, err := os.Stat(referenceFile) + if updateSnapshots && err == nil { + require.NoError(t, os.Remove(referenceFile)) } - }) + require.ErrorIs(t, err, os.ErrNotExist) + return + } + + data, err := os.ReadFile(res.Filename) + require.NoError(t, err) + + data = bytes.ReplaceAll(data, []byte(samplesDir), []byte("samples")) + data = bytes.ReplaceAll(data, []byte(fmt.Sprintf("%q", version.Tag)), []byte("\"\"")) + + reference, err := os.ReadFile(referenceFile) + if updateSnapshots && errors.Is(err, os.ErrNotExist) { + require.NoError(t, os.MkdirAll(filepath.Dir(referenceFile), 0o755)) + require.NoError(t, os.WriteFile(referenceFile, data, 0o644)) + } + require.NoError(t, err) + + if !assert.Equal(t, string(reference), string(data)) && updateSnapshots { + require.NoError(t, os.WriteFile(referenceFile, data, 0o644)) + } } }) } diff --git a/internal/injector/check.go b/internal/injector/check.go new file mode 100644 index 000000000..759e46474 --- /dev/null +++ b/internal/injector/check.go @@ -0,0 +1,34 @@ +// 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 injector + +import ( + "fmt" + "go/ast" + "go/importer" + "go/token" + "go/types" + "runtime" +) + +// typeCheck runs the Go type checker on the provided files, and returns the +// Uses type information map that is built in the process. +func (i *Injector) typeCheck(fset *token.FileSet, files []*ast.File) (map[*ast.Ident]types.Object, error) { + pkg := types.NewPackage(i.ImportPath, i.Name) + typeInfo := types.Info{Uses: make(map[*ast.Ident]types.Object)} + + checkerCfg := types.Config{ + GoVersion: i.GoVersion, + Importer: importer.ForCompiler(fset, runtime.Compiler, i.Lookup), + } + checker := types.NewChecker(&checkerCfg, fset, pkg, &typeInfo) + + if err := checker.Files(files); err != nil { + return nil, fmt.Errorf("type-checking files: %w", err) + } + + return typeInfo.Uses, nil +} diff --git a/internal/injector/imports.go b/internal/injector/imports.go new file mode 100644 index 000000000..3c8764442 --- /dev/null +++ b/internal/injector/imports.go @@ -0,0 +1,115 @@ +// 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 injector + +import ( + "go/token" + "strconv" + + "github.com/dave/dst" + "github.com/dave/dst/dstutil" +) + +// canonicalizeImports works around the issue detailed in https://github.com/dave/dst/issues/45 +// where dave/dst improperly handles multiple imports of the same package with different aliases, +// resulting in invalid output source code. +// +// To do so, it modifies the AST file so that it only includes a single import per path, using the +// first non-empty alias found. +func canonicalizeImports(file *dst.File) { + specsByPath := importSpecsByImportPath(file) + + retain := filterExtraneousImports(specsByPath) + + file.Imports = file.Imports[:0] // Re-use the backing store, we'll keep <= what was there. + for spec := range retain { + file.Imports = append(file.Imports, spec) + } + + filterDecls(file, retain) +} + +func importSpecsByImportPath(file *dst.File) map[string][]*dst.ImportSpec { + byPath := make(map[string][]*dst.ImportSpec, len(file.Imports)) + + for _, imp := range file.Imports { + path, err := strconv.Unquote(imp.Path.Value) + if err != nil { + continue + } + list := append(byPath[path], imp) + byPath[path] = list + } + + return byPath +} + +func filterExtraneousImports(byPath map[string][]*dst.ImportSpec) map[*dst.ImportSpec]struct{} { + result := make(map[*dst.ImportSpec]struct{}, len(byPath)) + + for _, specs := range byPath { + retain := specs[0] + for _, spec := range specs[1:] { + if (spec.Name == nil && (retain.Name == nil || retain.Name.Name != "_")) || spec.Name.Name == "_" { + continue + } + retain = spec + break + } + result[retain] = struct{}{} + } + + return result +} + +func filterDecls(file *dst.File, retain map[*dst.ImportSpec]struct{}) { + dstutil.Apply( + file, + func(csor *dstutil.Cursor) bool { + switch node := csor.Node().(type) { + case *dst.GenDecl: + // Only visit the children of `import` declarations. + return node.Tok == token.IMPORT + case *dst.ImportSpec: + // Filter out ImportSpec entries to keep only those in retain + if _, ret := retain[node]; !ret { + csor.Delete() + } + // No need to traverse children. + return false + case dst.Decl: + // No need to visit any other kind of declaration + return false + default: + // Visit other node types (e.g, the *ast.File) + return true + } + }, + func(csor *dstutil.Cursor) bool { + switch node := csor.Node().(type) { + case *dst.GenDecl: + if node.Tok != token.IMPORT { + // Imports are before any other kind of declaration, we can abort traversal as soon as we + // find a declaration that is not an `import` declaration. + return false + } + + if len(node.Specs) == 0 { + csor.Delete() + } + // Proceed with the rest of the nodes (there may be more imports). + return true + case dst.Decl: + // Imports are before any other kind of declaration, we can abort traversal as soon as we + // find a declaration that is not an `import` declaration. + return false + default: + // Proceed with the rest of the nodes (there may be imports down there). + return true + } + }, + ) +} diff --git a/internal/injector/injector.go b/internal/injector/injector.go index 049c675ce..ce19b6980 100644 --- a/internal/injector/injector.go +++ b/internal/injector/injector.go @@ -9,294 +9,189 @@ package injector import ( - "bytes" "errors" "fmt" - "go/parser" + "go/importer" "go/token" - "os" - "path/filepath" - "sync" - "github.com/datadog/orchestrion/internal/goflags" "github.com/datadog/orchestrion/internal/injector/aspect" "github.com/datadog/orchestrion/internal/injector/aspect/context" - "github.com/datadog/orchestrion/internal/injector/builtin" - "github.com/datadog/orchestrion/internal/injector/lineinfo" "github.com/datadog/orchestrion/internal/injector/typed" - "github.com/datadog/orchestrion/internal/log" "github.com/dave/dst" "github.com/dave/dst/decorator" - "github.com/dave/dst/decorator/resolver/guess" + "github.com/dave/dst/decorator/resolver" + "github.com/dave/dst/decorator/resolver/gotypes" "github.com/dave/dst/dstutil" - "golang.org/x/tools/go/packages" ) type ( - // Injector injects go code into a program. + // Injector injects go code into a specific Go package. Injector struct { - fileset *token.FileSet - decorators []*decorator.Decorator - newRestorer func(string) *decorator.FileRestorer - opts Options - mutex sync.Mutex // Guards access to InjectFile - } - - // ModifiedFileFn is called with the original file and must return the path to use when writing a modified version. - ModifiedFileFn func(string) string - - Options struct { - // Aspects is the set of configured injections to attempt. + // Aspects is the set of configured aspects to use. Aspects []aspect.Aspect - // Dir is the working directory to use for resolving dependencies, etc... If blank, the current working directory is - // used. - Dir string - // IncludeTests requests test files to be prepared for injection, too. - IncludeTests bool - // ModifiedFile is called to obtain the file name to use when writing a modified file. If nil, the original file is - // overwritten in-place. - ModifiedFile ModifiedFileFn - // PreserveLineInfo enables emission of //line directives to preserve line information from the original file, so - // that stack traces resolve to the original source code. This is strongly recommended when performing compile-time - // injection. - PreserveLineInfo bool - } -) - -// New creates a new injector with the specified options. -func New(pkgDir string, opts Options) (*Injector, error) { - fileset := token.NewFileSet() - cfg := &packages.Config{ - // Explicitly disable toolexec for this, as if provided via $GOFLAGS, it - // would be honored by the go/packages loader and that'd cause this call to - // become a fork-bomb. - BuildFlags: []string{"-toolexec="}, - Dir: opts.Dir, - Fset: fileset, - Mode: packages.NeedName | - packages.NeedFiles | - packages.NeedImports | - packages.NeedTypes | - packages.NeedTypesSizes | - packages.NeedSyntax | - packages.NeedTypesInfo, - Tests: opts.IncludeTests, - Logf: func(format string, args ...any) { log.Tracef(format+"\n", args...) }, - } - if flags, err := goflags.Flags(); err == nil { - // Honor any `-tags` flags provided by the user, as these may affect what - // is getting compiled or not. - if tags, hasTags := flags.Get("-tags"); hasTags { - cfg.BuildFlags = append(cfg.BuildFlags, fmt.Sprintf("-tags=%s", tags)) - } - } - - var ( - pkgPath string - decorators []*decorator.Decorator - restorerMap map[string]string - ) - if pkgs, err := decorator.Load(cfg, pkgDir); err != nil { - return nil, err - } else { - decorators = make([]*decorator.Decorator, 0, len(pkgs)) - restorerMap = make(map[string]string, len(builtin.RestorerMap)) - for path, name := range builtin.RestorerMap { - restorerMap[path] = name - } - - for _, pkg := range pkgs { - if len(pkg.Errors) > 0 { - errs := make([]error, len(pkg.Errors)) - for i := range pkg.Errors { - errs[i] = pkg.Errors[i] - } - return nil, errors.Join(errs...) - } - if pkgPath == "" { - // The first non-blank package path is the "top level" one (the one we care about). - pkgPath = pkg.PkgPath - } + // ImportPath is the import path of the package that will be injected. + ImportPath string + // Name is the name of the package that will be injected. If blank, it will be determined from parsing source files. + Name string + // GoVersion is the go runtime version required by this package. If blank, no go runtime compatibility will be + // asserted. + GoVersion string - for _, imp := range pkg.Imports { - if imp.Name == "" { - // Happens when there is an error while processing the import, typically inability to resolve the name due to - // a typo or something. If we allow blank names in the map, the restorer just removes the qualifiers, which is - // obviously undesirable. - continue - } - restorerMap[imp.PkgPath] = imp.Name - } + // ModifiedFile is called to determine the output file name for a modified file. If nil, the input file is modified + // in-place. + ModifiedFile func(string) string + // Lookup is a function that resolves and imported package's archive file. + Lookup importer.Lookup + // RootConfig is the root configuration value to use. + RootConfig map[string]string - // pkg.Decorator is nil in case the package in question does not include any go source file. This can be the case - // when building test packages that do not include any non-test source file; in which case the "package under - // test" is empty. This is because the loader returns three different entries when processing tests: - // 1. The package under test (which may be empty) - // 2. The test functions - // 3. The test binary/main - if pkg.Decorator != nil { - decorators = append(decorators, pkg.Decorator) - } - } + // restorerResolver is used to restore modified files. It's created on-demand then re-used. + restorerResolver resolver.RestorerResolver } - if len(decorators) == 0 { - return nil, errors.New("no decorators could be created") + // InjectedFile contains information about a modified file. It can be used to update compilation instructions. + InjectedFile struct { + // References holds new references created while injecting the package, if any. + References typed.ReferenceMap + // Filename is the name of the file that needs to be compiled in place of the original one. It may be identical to + // the input file if the Injector.ModifiedFile function is nil or returns identity. + Filename string } - return &Injector{ - fileset: fileset, - decorators: decorators, - newRestorer: func(filename string) *decorator.FileRestorer { - res := decorator.NewRestorerWithImports(pkgPath, guess.WithMap(restorerMap)).FileRestorer() - res.Name = filename - return res - }, - opts: opts, - }, nil -} - -type ( - // Result describes the result of an injection operation. - Result struct { - References typed.ReferenceMap // New package references injected into the file and what kind of reference they are - Filename string // The file name of the output file (may be different from the input file) - Modified bool // Whether the file was modified + result struct { + InjectedFile + Modified bool } ) -// Injects code in the specified file. This method can be called concurrently by multiple goroutines, -// as is guarded by a sync.Mutex. -func (i *Injector) InjectFile(filename string, rootConfig map[string]string) (res Result, err error) { - i.mutex.Lock() - defer i.mutex.Unlock() - - res.Filename = filename +// InjectFiles performs injections on the specified files. All provided file paths must belong to the import path set on +// the receiving Injector. The method returns a map that associates the original source file path to the modified file +// information. It does not contain entries for unmodified files. +func (i *Injector) InjectFiles(files []string) (map[string]InjectedFile, error) { + if err := i.validate(); err != nil { + return nil, err + } - file, decorator, err := i.lookupDecoratedFile(filename) + fset := token.NewFileSet() + astFiles, err := parseFiles(fset, files) if err != nil { - return res, err + return nil, err } - if res.Modified, res.References, err = i.inject(file, decorator, rootConfig); err != nil { - return res, err + uses, err := i.typeCheck(fset, astFiles) + if err != nil { + return nil, err } - if res.Modified { - buf := bytes.NewBuffer(nil) - if err = i.newRestorer(filename).Fprint(buf, file); err != nil { - return res, err + decorator := decorator.NewDecoratorWithImports(fset, i.ImportPath, gotypes.New(uses)) + dstFiles := make([]*dst.File, len(astFiles)) + for idx, astFile := range astFiles { + dstFiles[idx], err = decorator.DecorateFile(astFile) + if err != nil { + return nil, err } + } - res.Filename = i.outputFileFor(filename) - if err := os.MkdirAll(filepath.Dir(res.Filename), 0o755); err != nil { - return res, err + result := make(map[string]InjectedFile, len(files)) + for idx, dstFile := range dstFiles { + res, err := i.injectFile(decorator, dstFile) + if err != nil { + return nil, err + } + if res.Modified { + result[files[idx]] = res.InjectedFile } - err = os.WriteFile(res.Filename, postProcess(buf.Bytes()), 0o644) } - - return res, err + return result, nil } -func (i *Injector) lookupDecoratedFile(filename string) (*dst.File, *decorator.Decorator, error) { - stat, err := os.Stat(filename) - if err != nil { - return nil, nil, err +func (i *Injector) validate() error { + var err error + if i.ImportPath == "" { + err = errors.Join(err, fmt.Errorf("invalid %T: missing ImportPath", i)) } - - for _, dec := range i.decorators { - for node, name := range dec.Filenames { - if name == "" { - // A bunch of synthetic nodes won't have a file name. - continue - } - s, err := os.Stat(name) - if err != nil { - continue - } - if os.SameFile(stat, s) { - return node, dec, nil - } - } + if i.Lookup == nil { + err = errors.Join(err, fmt.Errorf("invalid %T: missing Lookup", i)) } + return err +} + +// injectFile injects code in the specified file. This method can be called concurrently by multiple goroutines, +// as is guarded by a sync.Mutex. +func (i *Injector) injectFile(decorator *decorator.Decorator, file *dst.File) (result, error) { + result := result{InjectedFile: InjectedFile{Filename: decorator.Filenames[file]}} - src, err := os.ReadFile(filename) + var err error + result.Modified, result.References, err = i.applyAspects(decorator, file, i.RootConfig) if err != nil { - return nil, nil, err + return result, err } - decorator := i.decorators[0] - file, err := decorator.ParseFile(filename, src, parser.ParseComments) - if err != nil { - return nil, nil, err + if result.Modified { + result.Filename, err = i.writeModifiedFile(decorator, file) + if err != nil { + return result, err + } } - return file, decorator, nil + return result, nil } -// inject performs all configured injections on the specified file. It returns whether the file was -// modified, any import references introduced by modifications. In case of an error, the -// trasnformation aborts as quickly as possible and returns the error. -func (i *Injector) inject(file *dst.File, decorator *decorator.Decorator, rootConfig map[string]string) (mod bool, refs typed.ReferenceMap, err error) { - var chain *context.NodeChain +func (i *Injector) applyAspects(decorator *decorator.Decorator, file *dst.File, rootConfig map[string]string) (bool, typed.ReferenceMap, error) { + var ( + chain *context.NodeChain + modified bool + references typed.ReferenceMap + err error + ) - dstutil.Apply( - file, - func(csor *dstutil.Cursor) bool { - if err != nil || csor.Node() == nil || ddIgnored(csor.Node()) { - return false - } - root := chain == nil - chain = chain.Child(csor) - if root { - chain.SetConfig(rootConfig) - } - return true - }, - func(csor *dstutil.Cursor) bool { - if err != nil || csor.Node() == nil || ddIgnored(csor.Node()) { - return false - } + pre := func(csor *dstutil.Cursor) bool { + if err != nil || csor.Node() == nil || ddIgnored(csor.Node()) { + return false + } + root := chain == nil + chain = chain.Child(csor) + if root { + chain.SetConfig(rootConfig) + } + return true + } - // Pop the ancestry stack now that we're done with this node. - defer func() { chain = chain.Parent() }() + post := func(csor *dstutil.Cursor) bool { + // Pop the ancestry stack now that we're done with this node. + defer func() { chain = chain.Parent() }() - var changed bool - changed, err = i.injectNode(chain.Context( - csor, - decorator.Path, - file, - &refs, - decorator, - )) - mod = mod || changed + var changed bool + changed, err = i.injectNode(chain.Context( + csor, + decorator.Path, + file, + &references, + decorator, + )) + modified = modified || changed - return err == nil - }, - ) + return err == nil + } + + dstutil.Apply(file, pre, post) // We only inject synthetic imports here because it may offset declarations by one position in // case a new import declaration is necessary, which causes dstutil.Apply to re-traverse the // current declaration. - if refs.AddSyntheticImports(file) { - mod = true - } - - if mod && i.opts.PreserveLineInfo { - if err := lineinfo.AnnotateMovedNodes(decorator, file, i.newRestorer); err != nil { - return mod, refs, err - } + if references.AddSyntheticImports(file) { + modified = true } - return + return modified, references, err } -// injectNode assesses all configured injections agaisnt the current node, and performs any AST +// injectNode assesses all configured aspects agaisnt the current node, and performs any AST // transformations. It returns whether the AST was indeed modified. In case of an error, the // injector aborts immediately and returns the error. func (i *Injector) injectNode(ctx context.AdviceContext) (mod bool, err error) { - for _, inj := range i.opts.Aspects { + for _, inj := range i.Aspects { if !inj.JoinPoint.Matches(ctx) { continue } @@ -314,12 +209,3 @@ func (i *Injector) injectNode(ctx context.AdviceContext) (mod bool, err error) { return } - -// outputFileFor returns the file name to be used when writing a modified file. It uses the options -// specified when building this Injector. -func (i *Injector) outputFileFor(filename string) string { - if i.opts.ModifiedFile == nil { - return filename - } - return i.opts.ModifiedFile(filename) -} diff --git a/internal/injector/injector_test.go b/internal/injector/injector_test.go index 2f63c0607..d0b176321 100644 --- a/internal/injector/injector_test.go +++ b/internal/injector/injector_test.go @@ -7,6 +7,7 @@ package injector_test import ( "fmt" + "io" "os" "os/exec" "path/filepath" @@ -23,17 +24,19 @@ import ( "github.com/sergi/go-diff/diffmatchpatch" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "golang.org/x/tools/go/packages" "gopkg.in/yaml.v3" "gotest.tools/v3/golden" ) type testConfig struct { Aspects []aspect.Aspect `yaml:"aspects"` - PreserveLineInfo bool `yaml:"preserveLineInfo"` SyntheticReferences map[string]typed.ReferenceKind `yaml:"syntheticReferences"` Code string `yaml:"code"` } +const testModuleName = "dummy/test/module" + func Test(t *testing.T) { t.Parallel() @@ -58,13 +61,33 @@ func Test(t *testing.T) { t.Run(testName, func(t *testing.T) { t.Parallel() + tmp := t.TempDir() + + testLookup := func(path string) (io.ReadCloser, error) { + pkgs, err := packages.Load( + &packages.Config{ + Mode: packages.NeedExportFile, + Dir: tmp, + Logf: t.Logf, + }, + path, + ) + if err != nil { + return nil, err + } + file := pkgs[0].ExportFile + if file == "" { + return nil, fmt.Errorf("no export file found for %q", path) + } + return os.Open(file) + } + data, err := os.ReadFile(filepath.Join(testPath, "config.yml")) require.NoError(t, err, "failed to read test configuration") var config testConfig require.NoError(t, yaml.Unmarshal(data, &config), "failed to parse test configuration") - tmp := t.TempDir() - runGo(t, tmp, "mod", "init", "dummy/test/module") + runGo(t, tmp, "mod", "init", testModuleName) runGo(t, tmp, "mod", "edit", "-replace", fmt.Sprintf("github.com/datadog/orchestrion=%s", rootDir)) runGo(t, tmp, "mod", "edit", "-replace", fmt.Sprintf("orchestrion/integration=%s", integDir)) @@ -73,39 +96,37 @@ func Test(t *testing.T) { require.NoError(t, os.WriteFile(inputFile, []byte(original), 0o644), "failed to create main.go") runGo(t, tmp, "mod", "tidy") - inj, err := injector.New(tmp, injector.Options{ - Aspects: config.Aspects, - Dir: tmp, - PreserveLineInfo: config.PreserveLineInfo, - ModifiedFile: func(path string) string { return path + ".edited.go" }, - }) - require.NoError(t, err, "failed to create injector") + inj := injector.Injector{ + Aspects: config.Aspects, + ModifiedFile: func(path string) string { return filepath.Join(tmp, filepath.Base(path)+".edited.go") }, + ImportPath: testModuleName, + Lookup: testLookup, + } - res, err := inj.InjectFile(inputFile, nil) + res, err := inj.InjectFiles([]string{inputFile}) require.NoError(t, err, "failed to inject file") - if res.Modified { - assert.Equal(t, inputFile+".edited.go", res.Filename) - } else { - assert.Equal(t, inputFile, res.Filename) + resFile, modified := res[inputFile] + if !modified { + golden.Assert(t, "", filepath.Join(dirName, testName, "expected.diff")) + return } - modifiedData, err := os.ReadFile(res.Filename) - require.NoError(t, err, "failed to read modified file") - modified := normalize(modifiedData, inputFile) + assert.Equal(t, filepath.Join(tmp, filepath.Base(inputFile)+".edited.go"), resFile.Filename) + assert.Equal(t, config.SyntheticReferences, resFile.References.Map()) - assert.Equal(t, config.SyntheticReferences, res.References.Map()) + modifiedData, err := os.ReadFile(resFile.Filename) + require.NoError(t, err, "failed to read modified file") + normalized := normalize(modifiedData, inputFile) - edits := myers.ComputeEdits(span.URIFromPath("input.go"), original, modified) + edits := myers.ComputeEdits(span.URIFromPath("input.go"), original, normalized) diff := gotextdiff.ToUnified("input.go", "output.go", original, edits) golden.Assert(t, fmt.Sprint(diff), filepath.Join(dirName, testName, "expected.diff")) - if res.Modified { - // Verify that the modified code compiles... - os.Rename(res.Filename, inputFile) - runGo(t, tmp, "mod", "tidy") - runGo(t, tmp, "build", inputFile) - } + // Verify that the modified code compiles... + os.Rename(resFile.Filename, inputFile) + runGo(t, tmp, "mod", "tidy") + runGo(t, tmp, "build", inputFile) }) } } diff --git a/internal/injector/parse.go b/internal/injector/parse.go new file mode 100644 index 000000000..09e7a8997 --- /dev/null +++ b/internal/injector/parse.go @@ -0,0 +1,137 @@ +// 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 injector + +import ( + "bytes" + "fmt" + "go/ast" + "go/parser" + "go/token" + "io" + "os" +) + +// parseFiles parses all provided files, after having applied any leading +// "//line" directives if present. +func parseFiles(fset *token.FileSet, files []string) ([]*ast.File, error) { + result := make([]*ast.File, len(files)) + for idx, file := range files { + var err error + result[idx], err = parseFile(fset, file) + if err != nil { + return nil, err + } + } + return result, nil +} + +// parseFile parses the provided filename, after having applied a leading +// "//line" directive if one is present. +func parseFile(fset *token.FileSet, filename string) (*ast.File, error) { + file, err := os.Open(filename) + if err != nil { + return nil, fmt.Errorf("open %q: %w", filename, err) + } + defer file.Close() + + // If the file begins with a "//line :1:1" directive, we consume it and + // then pretend the "" was our filename all along. This simplifies + // handling of line offsets further down the line and removes some duplicated + // effort to do it early. + mappedFilename := filename + if mapped, err := consumeLineDirective(file); err != nil { + return nil, fmt.Errorf("peeking at first line of %q: %w", filename, err) + } else if mapped != "" { + mappedFilename = mapped + } + + astFile, err := parser.ParseFile(fset, mappedFilename, file, parser.ParseComments) + if err != nil { + return nil, fmt.Errorf("parsing %q: %w", filename, err) + } + + return astFile, nil +} + +// consumeLineDirective consumes the first line from r if it's a "//line" +// directive that either does not have line/column information or has it set to +// line 1 (and column 1). If the directive is consumed, the filename it refers +// to is returned. Otherwise, the reader is rewinded to its original position. +func consumeLineDirective(r io.ReadSeeker) (string, error) { + var buf [7]byte + n, err := r.Read(buf[:]) + if err != nil { + return "", err + } + if string(buf[:n]) != "//line " { + _, err := r.Seek(0, io.SeekStart) + return "", err + } + + buffer := make([]byte, 0, 128) + var wasCR, done bool + for !done { + if n, err := r.Read(buf[:1]); err != nil { + return "", err + } else if n == 0 { + // Reached EOF + break + } + switch c := buf[0]; c { + case '\n': + done = true + case '\r': + wasCR = true + continue + default: + if wasCR { + // We saw a CR and this is not an LF, so we rewind one byte and bail out. + if _, err := r.Seek(-1, io.SeekCurrent); err != nil { + return "", err + } + done = true + } else { + buffer = append(buffer, c) + } + } + } + + // Remove any leading or trailing white space + if rest, pos, hadPos := cutPositionSuffix(bytes.TrimSpace(buffer)); !hadPos { + return string(rest), nil + } else if pos != 1 { + // It was not at position 1, so it's not the directive we're looking for. + _, err := r.Seek(0, io.SeekStart) + return "", err + } else if rest, pos, hadPos := cutPositionSuffix(rest); !hadPos || pos == 1 { + return string(rest), nil + } + // It was not at position 1, so it's not the directive we're looking for. + _, err = r.Seek(0, io.SeekStart) + return "", err +} + +// cutPositionSuffix removes a trailing ":" from the provided buffer, if present. +func cutPositionSuffix(buf []byte) ([]byte, int, bool) { + cutOff := len(buf) - 1 + + // First, consume the integer at the end of the buffer. + pos := 0 + pow := 1 + for buf[cutOff] >= '0' && buf[cutOff] <= '9' { + pos += pow * int(buf[cutOff]-'0') + pow *= 10 + cutOff-- + } + + // If there's no ":" before the integer, or there was no digit at all, it was not a position... + if buf[cutOff] != ':' || pow == 1 { + return buf, 0, false + } + + return buf[:cutOff], pos, true +} diff --git a/internal/injector/parse_test.go b/internal/injector/parse_test.go new file mode 100644 index 000000000..f325dac4f --- /dev/null +++ b/internal/injector/parse_test.go @@ -0,0 +1,72 @@ +// 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 injector + +import ( + "bytes" + "io" + "strings" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestConsumeLineDirective(t *testing.T) { + source := []string{ + "package main", + "", + "func main() {", + "\t/* ... */", + "}", + "", + } + + lineSeparators := map[string]string{ + "CR": "\r", + "LF": "\n", + "CRLF": "\r\n", + } + + for sepStyle, sep := range lineSeparators { + t.Run(sepStyle, func(t *testing.T) { + sourceBytes := []byte(strings.Join(source, sep)) + + cases := map[string]string{ + "// no directive": "", + // Bad directives (ignored) + "//line path/to/file.go:1:42": "", + "//line path/to/file.go:1337": "", + // Valid directives + "//line path/to/file.go:1:1": "path/to/file.go", + "//line path/to/file.go:1": "path/to/file.go", + "//line path/to/file.go": "path/to/file.go", + } + for directive, expectedOutcome := range cases { + t.Run(directive, func(t *testing.T) { + var buffer bytes.Buffer + buffer.WriteString(directive) + buffer.WriteString(sep) + buffer.Write(sourceBytes) + + reader := bytes.NewReader(buffer.Bytes()) + filename, err := consumeLineDirective(reader) + require.NoError(t, err) + require.Equal(t, expectedOutcome, filename) + + rest, err := io.ReadAll(reader) + require.NoError(t, err) + + expected := sourceBytes + if expectedOutcome == "" { + expected = buffer.Bytes() + } + + require.Equal(t, string(expected), string(rest)) + }) + } + }) + } +} diff --git a/internal/injector/restorer.go b/internal/injector/restorer.go new file mode 100644 index 000000000..2dac42dd3 --- /dev/null +++ b/internal/injector/restorer.go @@ -0,0 +1,68 @@ +// 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 injector + +import ( + "fmt" + "go/importer" + "go/token" + "go/types" + + "github.com/dave/dst/decorator" + "golang.org/x/tools/go/gcexportdata" +) + +type lookupResolver struct { + lookup importer.Lookup + + fset *token.FileSet + imports map[string]*types.Package +} + +func (i *Injector) newRestorer(filename string) *decorator.FileRestorer { + if i.restorerResolver == nil { + i.restorerResolver = &lookupResolver{lookup: i.Lookup} + } + + return &decorator.FileRestorer{ + Restorer: decorator.NewRestorerWithImports(i.ImportPath, i.restorerResolver), + Name: filename, + } +} + +func (r *lookupResolver) ResolvePackage(path string) (string, error) { + // The "unsafe" package does not have an archive, so it's hard-coded here. + if path == "unsafe" { + return "unsafe", nil + } + + // If this is present in "cache", we can return right away! + if pkg, ok := r.imports[path]; ok { + return pkg.Name(), nil + } + + if r.fset == nil { + r.fset = token.NewFileSet() + } + if r.imports == nil { + r.imports = make(map[string]*types.Package) + } + + rd, err := r.lookup(path) + if err != nil { + return "", fmt.Errorf("lookup %q: %w", path, err) + } + gcr, err := gcexportdata.NewReader(rd) + if err != nil { + return "", fmt.Errorf("locating gc export data: %w", err) + } + pkg, err := gcexportdata.Read(gcr, r.fset, r.imports, path) + if err != nil { + return "", fmt.Errorf("reading gc export data: %w", err) + } + + return pkg.Name(), err +} diff --git a/internal/injector/testdata/injector/access-return-value/expected.diff b/internal/injector/testdata/injector/access-return-value/expected.diff index d371251f6..901a8d23d 100644 --- a/internal/injector/testdata/injector/access-return-value/expected.diff +++ b/internal/injector/testdata/injector/access-return-value/expected.diff @@ -1,11 +1,17 @@ --- input.go +++ output.go -@@ -5,7 +5,14 @@ +@@ -1,3 +1,4 @@ ++//line input.go:1:1 + package test + + import ( +@@ -5,7 +6,16 @@ "log" ) -func test() (interface{}, error) { +func test() (_ interface{}, __result__1 error) { ++//line + { + defer func() { + if __result__1 != nil { @@ -13,6 +19,7 @@ + } + }() + } ++//line input.go:9 log.Println("Running test function...") return nil, errors.ErrUnsupported } diff --git a/internal/injector/testdata/injector/chi5-newroute-dotimport/config.yml b/internal/injector/testdata/injector/chi5-newroute-dotimport/config.yml index 70098019b..c20f16753 100644 --- a/internal/injector/testdata/injector/chi5-newroute-dotimport/config.yml +++ b/internal/injector/testdata/injector/chi5-newroute-dotimport/config.yml @@ -15,8 +15,6 @@ aspects: return mux }() -preserveLineInfo: true - syntheticReferences: gopkg.in/DataDog/dd-trace-go.v1/contrib/go-chi/chi.v5: true diff --git a/internal/injector/testdata/injector/chi5-newroute/config.yml b/internal/injector/testdata/injector/chi5-newroute/config.yml index 1a485ecea..80a29db6c 100644 --- a/internal/injector/testdata/injector/chi5-newroute/config.yml +++ b/internal/injector/testdata/injector/chi5-newroute/config.yml @@ -15,8 +15,6 @@ aspects: return mux }() -preserveLineInfo: true - syntheticReferences: gopkg.in/DataDog/dd-trace-go.v1/contrib/go-chi/chi.v5: true diff --git a/internal/injector/testdata/injector/database-sql/config.yml b/internal/injector/testdata/injector/database-sql/config.yml index efe962b18..56b85c0a9 100644 --- a/internal/injector/testdata/injector/database-sql/config.yml +++ b/internal/injector/testdata/injector/database-sql/config.yml @@ -22,8 +22,6 @@ aspects: {{range .AST.Args}}{{.}}, {{end}}) -preserveLineInfo: true - syntheticReferences: gopkg.in/DataDog/dd-trace-go.v1/contrib/database/sql: true diff --git a/internal/injector/testdata/injector/directive/config.yml b/internal/injector/testdata/injector/directive/config.yml index 0da9a7681..8c2551cff 100644 --- a/internal/injector/testdata/injector/directive/config.yml +++ b/internal/injector/testdata/injector/directive/config.yml @@ -24,8 +24,6 @@ aspects: , {{printf "%q" .Key}}, {{printf "%q" .Value}} {{- end -}}) -preserveLineInfo: true - syntheticReferences: github.com/datadog/orchestrion/instrument: true github.com/datadog/orchestrion/instrument/event: true diff --git a/internal/injector/testdata/injector/grpc/config.yml b/internal/injector/testdata/injector/grpc/config.yml index 19b244c24..cf8070564 100644 --- a/internal/injector/testdata/injector/grpc/config.yml +++ b/internal/injector/testdata/injector/grpc/config.yml @@ -27,8 +27,6 @@ aspects: - imports: *imports template: grpc.UnaryInterceptor(grpctrace.UnaryServerInterceptor()) -preserveLineInfo: true - syntheticReferences: gopkg.in/DataDog/dd-trace-go.v1/contrib/google.golang.org/grpc: true diff --git a/internal/injector/testdata/injector/http-anonymous-handler/expected.diff b/internal/injector/testdata/injector/http-anonymous-handler/expected.diff index 2e7d50733..03d2d3f7a 100644 --- a/internal/injector/testdata/injector/http-anonymous-handler/expected.diff +++ b/internal/injector/testdata/injector/http-anonymous-handler/expected.diff @@ -1,18 +1,25 @@ --- input.go +++ output.go -@@ -6,6 +6,8 @@ - "log" +@@ -1,3 +1,4 @@ ++//line input.go:1:1 + package main + + import ( +@@ -7,12 +8,36 @@ "net/http" + "orchestrion/integration" ++//line + __orchestrion_instrument "github.com/datadog/orchestrion/instrument" + __orchestrion_event "github.com/datadog/orchestrion/instrument/event" - "orchestrion/integration" ) -@@ -13,6 +15,24 @@ ++//line input.go:12 + func main() { s := &http.Server{ Addr: ":8085", Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { ++//line + { + r = r.WithContext(__orchestrion_instrument.Report( + r.Context(), @@ -31,6 +38,7 @@ + "http.useragent", r.Header.Get("User-Agent"), + ) + } ++//line input.go:16 data, err := io.ReadAll(r.Body) if err != nil { w.WriteHeader(http.StatusBadRequest) diff --git a/internal/injector/testdata/injector/http-default/expected.diff b/internal/injector/testdata/injector/http-default/expected.diff index d7bcc802c..660c81e63 100644 --- a/internal/injector/testdata/injector/http-default/expected.diff +++ b/internal/injector/testdata/injector/http-default/expected.diff @@ -1,18 +1,28 @@ --- input.go +++ output.go -@@ -6,6 +6,8 @@ - "log" +@@ -1,3 +1,4 @@ ++//line input.go:1:1 + package main + + import ( +@@ -7,8 +8,12 @@ "net/http" + "orchestrion/integration" ++//line + __orchestrion_instrument "github.com/datadog/orchestrion/instrument" + __orchestrion_event "github.com/datadog/orchestrion/instrument/event" - "orchestrion/integration" ) -@@ -21,6 +23,26 @@ ++//line input.go:12 + func main() { + s := &http.Server{ + Addr: ":8085", +@@ -21,6 +24,28 @@ } func handle(w http.ResponseWriter, r *http.Request) { ++//line + { + r = r.WithContext(__orchestrion_instrument.Report( + r.Context(), @@ -33,6 +43,7 @@ + "http.useragent", r.Header.Get("User-Agent"), + ) + } ++//line input.go:24 data, err := io.ReadAll(r.Body) if err != nil { w.WriteHeader(http.StatusBadRequest) diff --git a/internal/injector/testdata/injector/http-line-info/config.yml b/internal/injector/testdata/injector/http-line-info/config.yml index 17bef760b..1aa613144 100644 --- a/internal/injector/testdata/injector/http-line-info/config.yml +++ b/internal/injector/testdata/injector/http-line-info/config.yml @@ -37,8 +37,6 @@ aspects: {{ range .DirectiveArgs "dd:span" -}}{{printf "%q, %q," .Key .Value}}{{- end }} ) -preserveLineInfo: true - syntheticReferences: github.com/datadog/orchestrion/instrument/event: true github.com/datadog/orchestrion/instrument: true diff --git a/internal/injector/testdata/injector/http-no-arg-name/expected.diff b/internal/injector/testdata/injector/http-no-arg-name/expected.diff index fdf29b39f..3377cf530 100644 --- a/internal/injector/testdata/injector/http-no-arg-name/expected.diff +++ b/internal/injector/testdata/injector/http-no-arg-name/expected.diff @@ -1,20 +1,30 @@ --- input.go +++ output.go -@@ -5,6 +5,8 @@ - "log" +@@ -1,3 +1,4 @@ ++//line input.go:1:1 + package main + + import ( +@@ -6,8 +7,12 @@ "net/http" + "orchestrion/integration" ++//line + __orchestrion_instrument "github.com/datadog/orchestrion/instrument" + __orchestrion_event "github.com/datadog/orchestrion/instrument/event" - "orchestrion/integration" ) -@@ -19,5 +21,25 @@ ++//line input.go:11 + func main() { + s := &http.Server{ + Addr: ":8085", +@@ -19,5 +22,26 @@ log.Printf("Server shut down: %v", s.ListenAndServe()) } -func handle(http.ResponseWriter, *http.Request) { +func handle(_ http.ResponseWriter, __argument__1 *http.Request) { ++//line + { + __argument__1 = __argument__1.WithContext(__orchestrion_instrument.Report( + __argument__1.Context(), diff --git a/internal/injector/testdata/injector/http-server/config.yml b/internal/injector/testdata/injector/http-server/config.yml index 55f623af7..1963cafb6 100644 --- a/internal/injector/testdata/injector/http-server/config.yml +++ b/internal/injector/testdata/injector/http-server/config.yml @@ -14,8 +14,6 @@ aspects: instrument.WrapHandler({{ . }}) //dd:endwrap -preserveLineInfo: true - syntheticReferences: github.com/datadog/orchestrion/instrument: true diff --git a/internal/injector/testdata/injector/name-conflict/config.yml b/internal/injector/testdata/injector/name-conflict/config.yml index a43d3c66c..66fd1f4bf 100644 --- a/internal/injector/testdata/injector/name-conflict/config.yml +++ b/internal/injector/testdata/injector/name-conflict/config.yml @@ -12,9 +12,6 @@ aspects: template: |- fmt.Println("hello world") -preserveLineInfo: true -mustCompile: true - syntheticReferences: fmt: true diff --git a/internal/injector/testdata/injector/redigo/config.yml b/internal/injector/testdata/injector/redigo/config.yml index 707eb3638..189102157 100644 --- a/internal/injector/testdata/injector/redigo/config.yml +++ b/internal/injector/testdata/injector/redigo/config.yml @@ -17,8 +17,6 @@ aspects: - replace-function: gopkg.in/DataDog/dd-trace-go.v1/contrib/gomodule/redigo.DialURL -preserveLineInfo: true - syntheticReferences: gopkg.in/DataDog/dd-trace-go.v1/contrib/gomodule/redigo: true diff --git a/internal/injector/write.go b/internal/injector/write.go new file mode 100644 index 000000000..132912675 --- /dev/null +++ b/internal/injector/write.go @@ -0,0 +1,55 @@ +// 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 injector + +import ( + "bytes" + "fmt" + "go/format" + "os" + "path/filepath" + + "github.com/datadog/orchestrion/internal/injector/lineinfo" + "github.com/dave/dst" + "github.com/dave/dst/decorator" +) + +// writeModifiedFile writes the modified file to disk after having restored it to Go source code, +// and returns the path to the modified file. +func (i *Injector) writeModifiedFile(decorator *decorator.Decorator, file *dst.File) (string, error) { + canonicalizeImports(file) + + filename := decorator.Filenames[file] + + if err := lineinfo.AnnotateMovedNodes(decorator, file, i.newRestorer); err != nil { + return filename, fmt.Errorf("annotating moved nodes in %q: %w", filename, err) + } + + restorer := i.newRestorer(filename) + astFile, err := restorer.RestoreFile(file) + if err != nil { + return filename, fmt.Errorf("restoring %q: %w", filename, err) + } + + var buf bytes.Buffer + if err := format.Node(&buf, restorer.Fset, astFile); err != nil { + return filename, fmt.Errorf("formatting %q: %w", filename, err) + } + + if i.ModifiedFile != nil { + filename = i.ModifiedFile(filename) + dir := filepath.Dir(filename) + if err := os.MkdirAll(dir, 0o755); err != nil { + return filename, fmt.Errorf("mkdir %q: %w", dir, err) + } + } + + if err := os.WriteFile(filename, postProcess(buf.Bytes()), 0o644); err != nil { + return filename, fmt.Errorf("writing %q: %w", filename, err) + } + + return filename, nil +} diff --git a/internal/toolexec/aspect/oncompile.go b/internal/toolexec/aspect/oncompile.go index e1c76ea64..303196c97 100644 --- a/internal/toolexec/aspect/oncompile.go +++ b/internal/toolexec/aspect/oncompile.go @@ -11,8 +11,6 @@ import ( "os/exec" "path/filepath" "regexp" - "slices" - "strings" "github.com/datadog/orchestrion/internal/injector" "github.com/datadog/orchestrion/internal/injector/aspect" @@ -63,54 +61,43 @@ func (w Weaver) OnCompile(cmd *proxy.CompileCommand) error { } } - if cmd.SourceDir == "" { - // This may be the case for packages that only have CGO components, and all the Go source is generated by the `cgo` - // tools. In those cases, we won't weave anything (we couldn't properly anchor the source tree anyway, so we could - // not create an injector). - log.Debugf("No source directory found, nothing to weave.\n") - return nil + imports, err := importcfg.ParseFile(cmd.Flags.ImportCfg) + if err != nil { + return fmt.Errorf("parsing %q: %w", cmd.Flags.ImportCfg, err) } orchestrionDir := filepath.Join(filepath.Dir(cmd.Flags.Output), "orchestrion") - injector, err := injector.New(cmd.SourceDir, injector.Options{ - Aspects: aspects, - // Include test files if any of the input Go files has a _test.go suffix. - IncludeTests: slices.ContainsFunc(cmd.GoFiles(), func(name string) bool { return strings.HasSuffix(strings.ToLower(name), "_test.go") }), + injector := injector.Injector{ + Aspects: aspects, + RootConfig: map[string]string{"httpmode": "wrap"}, + Lookup: imports.Lookup, + ImportPath: w.ImportPath, + GoVersion: cmd.Flags.GoVersion, ModifiedFile: func(file string) string { return filepath.Join(orchestrionDir, "src", cmd.Flags.Package, filepath.Base(file)) }, - PreserveLineInfo: true, - }) + } + + goFiles := cmd.GoFiles() + results, err := injector.InjectFiles(goFiles) if err != nil { - return fmt.Errorf("creating injector for %s (in %q): %w", w.ImportPath, cmd.SourceDir, err) + return err } references := typed.ReferenceMap{} - for _, gofile := range cmd.GoFiles() { - res, err := injector.InjectFile(gofile, map[string]string{"httpmode": "wrap"}) - if err != nil { - return fmt.Errorf("weaving aspects in %q: %w", gofile, err) + for gofile, modFile := range results { + log.Debugf("Modified source code: %q => %q\n", gofile, modFile.Filename) + if err := cmd.ReplaceParam(gofile, modFile.Filename); err != nil { + return fmt.Errorf("replacing %q with %q: %w", gofile, modFile.Filename, err) } - if res.Modified { - log.Debugf("Modified source code: %q => %q\n", gofile, res.Filename) - if err := cmd.ReplaceParam(gofile, res.Filename); err != nil { - return fmt.Errorf("replacing %q with %q: %w", gofile, res.Filename, err) - } - } - - references.Merge(res.References) + references.Merge(modFile.References) } if references.Count() == 0 { return nil } - reg, err := importcfg.ParseFile(cmd.Flags.ImportCfg) - if err != nil { - return fmt.Errorf("parsing %q: %w", cmd.Flags.ImportCfg, err) - } - var ( linkDeps linkdeps.LinkDeps regUpdated bool @@ -121,14 +108,14 @@ func (w Weaver) OnCompile(cmd *proxy.CompileCommand) error { continue } - if archive, ok := reg.PackageFile[depImportPath]; ok { + if archive, ok := imports.PackageFile[depImportPath]; ok { deps, err := linkdeps.FromArchive(archive) if err != nil { return fmt.Errorf("reading %s from %q: %w", linkdeps.LinkDepsFilename, depImportPath, err) } log.Debugf("Processing %s dependencies from %s[%s]...", linkdeps.LinkDepsFilename, depImportPath, archive) for _, tDep := range deps.Dependencies() { - if _, found := reg.PackageFile[tDep]; !found { + if _, found := imports.PackageFile[tDep]; !found { log.Debugf("Copying %s dependency on %q inherited from %q\n", linkdeps.LinkDepsFilename, tDep, depImportPath) linkDeps.Add(tDep) } @@ -154,18 +141,18 @@ func (w Weaver) OnCompile(cmd *proxy.CompileCommand) error { } log.Debugf("Processing %s dependencies from %s...\n", linkdeps.LinkDepsFilename, dep) for _, tDep := range deps.Dependencies() { - if _, found := reg.PackageFile[tDep]; !found { + if _, found := imports.PackageFile[tDep]; !found { log.Debugf("Copying transitive %s dependency on %q inherited from %q via %q\n", linkdeps.LinkDepsFilename, tDep, depImportPath, dep) linkDeps.Add(tDep) } } - if _, ok := reg.PackageFile[dep]; ok { + if _, ok := imports.PackageFile[dep]; ok { // Already part of natural dependencies, nothing to do... continue } log.Debugf("Recording transitive dependency of %q: %q => %q\n", depImportPath, dep, archive) - reg.PackageFile[dep] = archive + imports.PackageFile[dep] = archive regUpdated = true } } @@ -179,7 +166,7 @@ func (w Weaver) OnCompile(cmd *proxy.CompileCommand) error { if regUpdated { // Creating updated version of the importcfg file, with new dependencies - if err := writeUpdatedImportConfig(reg, cmd.Flags.ImportCfg); err != nil { + if err := writeUpdatedImportConfig(imports, cmd.Flags.ImportCfg); err != nil { return fmt.Errorf("writing updated %q: %w", cmd.Flags.ImportCfg, err) } } diff --git a/internal/toolexec/importcfg/lookup.go b/internal/toolexec/importcfg/lookup.go new file mode 100644 index 000000000..430f3b5ab --- /dev/null +++ b/internal/toolexec/importcfg/lookup.go @@ -0,0 +1,33 @@ +// 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 importcfg + +import ( + "fmt" + "io" + "os" +) + +// Lookup opens the archive file for the provided import path. This allows the +// ImportConfig object to serve as a package information resolver's Lookup +// function. +func (r *ImportConfig) Lookup(path string) (io.ReadCloser, error) { + lookup := path + if p, mapped := r.ImportMap[lookup]; mapped { + lookup = p + } + + filename := r.PackageFile[lookup] + if filename == "" { + err := fmt.Errorf("no package file found for %q", lookup) + if lookup != path { + err = fmt.Errorf("mapped from %q: %w", path, err) + } + return nil, err + } + + return os.Open(filename) +} diff --git a/internal/toolexec/proxy/compile.go b/internal/toolexec/proxy/compile.go index 65c6965b4..82c8c88c8 100644 --- a/internal/toolexec/proxy/compile.go +++ b/internal/toolexec/proxy/compile.go @@ -7,7 +7,6 @@ package proxy import ( "errors" - "fmt" "path/filepath" "strings" ) @@ -17,14 +16,13 @@ type compileFlagSet struct { ImportCfg string `ddflag:"-importcfg"` Output string `ddflag:"-o"` TrimPath string `ddflag:"-trimpath"` + GoVersion string `ddflag:"-goversion"` } // CompileCommand represents a go tool `compile` invocation type CompileCommand struct { command Flags compileFlagSet - // SourceDir is the directory containing source files that are in the command's inputs. - SourceDir string // WorkDir is the $WORK directory managed by the go toolchain. WorkDir string } @@ -54,32 +52,17 @@ func (cmd *CompileCommand) AddFiles(files []string) { } } -func (f *compileFlagSet) Valid() bool { - return f.Package != "" && f.Output != "" && f.ImportCfg != "" && f.TrimPath != "" -} - -func (f *compileFlagSet) String() string { - return fmt.Sprintf("-p=%q -o=%q -importcfg=%q", f.Package, f.Output, f.ImportCfg) -} - func parseCompileCommand(args []string) (Command, error) { if len(args) == 0 { return nil, errors.New("unexpected number of command arguments") } cmd := CompileCommand{command: NewCommand(args)} parseFlags(&cmd.Flags, args[1:]) - files := cmd.GoFiles() - // The ImportCfg file is rooted in the stage directory - stageDir := filepath.Dir(cmd.Flags.ImportCfg) - // The WorkDir is the parent of the stage directory - cmd.WorkDir = filepath.Dir(stageDir) - // NOTE: Some commands just print the tool version, in which case no go file will be provided as arg - for _, f := range files { - if dir := filepath.Dir(f); dir != stageDir { - cmd.SourceDir = dir - break - } + if cmd.Flags.ImportCfg != "" { + // The WorkDir is the parent of the stage directory, which is where the importcfg file is located. + cmd.WorkDir = filepath.Dir(filepath.Dir(cmd.Flags.ImportCfg)) } + return &cmd, nil } diff --git a/internal/toolexec/proxy/compile_test.go b/internal/toolexec/proxy/compile_test.go index 27b9565a6..5fbd11d9c 100644 --- a/internal/toolexec/proxy/compile_test.go +++ b/internal/toolexec/proxy/compile_test.go @@ -14,36 +14,24 @@ import ( func TestParseCompile(t *testing.T) { for name, tc := range map[string]struct { - input []string - stage string - sourceDir string - goFiles []string - flags compileFlagSet + input []string + stage string + goFiles []string + flags compileFlagSet }{ "version_print": { input: []string{"/path/compile", "-V=full"}, stage: ".", }, "compile": { - input: []string{"/path/compile", "-o", "/buildDir/b002/a.out", "-p", "mypackage", "-importcfg", "/buildDir/b002/importcfg", "/source/dir/main.go", "/source/dir/file1.go"}, - stage: "b002", - sourceDir: "/source/dir", - goFiles: []string{"/source/dir/main.go", "/source/dir/file1.go"}, - flags: compileFlagSet{ - Package: "mypackage", - ImportCfg: "/buildDir/b002/importcfg", - Output: "/buildDir/b002/a.out", - }, - }, - "compile-no-source-dir": { - input: []string{"/path/compile", "-o", "/buildDir/b002/a.out", "-p", "mypackage", "-importcfg", "/buildDir/b002/importcfg", "/buildDir/b002/main.go", "/buildDir/b002/file1.go"}, - stage: "b002", - sourceDir: "", - goFiles: []string{"/source/dir/main.go", "/source/dir/file1.go"}, + input: []string{"/path/compile", "-o", "/buildDir/b002/a.out", "-p", "mypackage", "-goversion", "go1.42.1337", "-importcfg", "/buildDir/b002/importcfg", "/source/dir/main.go", "/source/dir/file1.go"}, + stage: "b002", + goFiles: []string{"/source/dir/main.go", "/source/dir/file1.go"}, flags: compileFlagSet{ Package: "mypackage", ImportCfg: "/buildDir/b002/importcfg", Output: "/buildDir/b002/a.out", + GoVersion: "go1.42.1337", }, }, } { @@ -54,7 +42,6 @@ func TestParseCompile(t *testing.T) { require.Equal(t, tc.stage, cmd.Stage()) c := cmd.(*CompileCommand) require.True(t, reflect.DeepEqual(tc.flags, c.Flags)) - require.Equal(t, tc.sourceDir, c.SourceDir) }) } }