fix(injector): nil-dereference on a specific import order

internal/injector/imports.go

@@ -9,6 +9,7 @@
 	"go/token"
 	"strconv"
 
+	"github.com/DataDog/orchestrion/internal/log"
 	"github.com/dave/dst"
 	"github.com/dave/dst/dstutil"
 )
@@ -20,7 +21,7 @@
 // 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)
+	specsByPath := importSpecsByImportPath(file.Imports)
 
 	retain := filterExtraneousImports(specsByPath)
 
@@ -32,16 +33,16 @@
 	filterDecls(file, retain)
 }
 
-func importSpecsByImportPath(file *dst.File) map[string][]*dst.ImportSpec {
-	byPath := make(map[string][]*dst.ImportSpec, len(file.Imports))
+func importSpecsByImportPath(imports []*dst.ImportSpec) map[string][]*dst.ImportSpec {
+	byPath := make(map[string][]*dst.ImportSpec, len(imports))
 
-	for _, imp := range file.Imports {
+	for _, imp := range imports {
 		path, err := strconv.Unquote(imp.Path.Value)
 		if err != nil {
+			log.Debugf("failed to unquote import path %q: %v\n", imp.Path.Value, err)
 			continue
 		}
-		list := append(byPath[path], imp)
-		byPath[path] = list
+		byPath[path] = append(byPath[path], imp)
 	}
 
 	return byPath
@@ -53,10 +54,22 @@
 	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 == "_" {
+			// Preference order is: spec.Name == nil < spec.Name.Name == "_" < spec.Name.Name != "_"
+			if spec.Name == nil {
 				continue
 			}
-			retain = spec
+
+			if retain.Name == nil && spec.Name != nil {
+				retain = spec
+				continue
+			}
+
+			if retain.Name.Name == "_" && spec.Name.Name != "_" {
+				retain = spec
+				continue
+			}
+
+			// We found a non-empty alias, no need to look further.
 			break
 		}
 		result[retain] = struct{}{}

internal/injector/imports_test.go

@@ -0,0 +1,131 @@
+// 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 2024 Datadog, Inc.
+
+package injector
+
+import (
+	"testing"
+
+	"github.com/dave/dst"
+	"github.com/stretchr/testify/require"
+)
+
+func TestFilterExtraneousImports(t *testing.T) {
+	qualifedImport := func(name string) *dst.ImportSpec {
+		return &dst.ImportSpec{Name: &dst.Ident{Name: name}}
+	}
+
+	for _, tc := range []struct {
+		name     string
+		in       []*dst.ImportSpec
+		expected *dst.ImportSpec
+	}{
+		{
+			name: "simple",
+			in: []*dst.ImportSpec{
+				qualifedImport("test"),
+			},
+			expected: qualifedImport("test"),
+		},
+		{
+			name: "multiple-qualified",
+			in: []*dst.ImportSpec{
+				qualifedImport("test1"),
+				qualifedImport("test2"),
+			},
+			expected: qualifedImport("test1"),
+		},
+		{
+			name: "one-qualified-first",
+			in: []*dst.ImportSpec{
+				qualifedImport("test1"),
+				{Name: &dst.Ident{Name: "_"}},
+			},
+			expected: qualifedImport("test1"),
+		},
+		{
+			name: "one-qualified-last",
+			in: []*dst.ImportSpec{
+				qualifedImport("test1"),
+			},
+			expected: qualifedImport("test1"),
+		},
+		{
+			name: "one-qualified-first-with-nil",
+			in: []*dst.ImportSpec{
+				qualifedImport("test1"),
+				{Name: nil},
+			},
+			expected: qualifedImport("test1"),
+		},
+		{
+			name: "one-qualified-last-with-nil",
+			in: []*dst.ImportSpec{
+				{Name: nil},
+				qualifedImport("test1"),
+			},
+			expected: qualifedImport("test1"),
+		},
+		{
+			name: "complex-1",
+			in: []*dst.ImportSpec{
+				{Name: nil},
+				{Name: &dst.Ident{Name: "_"}},
+				qualifedImport("test1"),
+				{Name: nil},
+			},
+			expected: qualifedImport("test1"),
+		},
+		{
+			name: "complex-2",
+			in: []*dst.ImportSpec{
+				{Name: &dst.Ident{Name: "_"}},
+				{Name: nil},
+				{Name: nil},
+			},
+			expected: &dst.ImportSpec{Name: &dst.Ident{Name: "_"}},
+		},
+		{
+			name: "complex-3",
+			in: []*dst.ImportSpec{
+				{Name: &dst.Ident{Name: "_"}},
+				{Name: nil},
+				{Name: nil},
+			},
+			expected: &dst.ImportSpec{Name: &dst.Ident{Name: "_"}},
+		},
+		{
+			name: "complex-4",
+			in: []*dst.ImportSpec{
+				{Name: nil},
+				{Name: &dst.Ident{Name: "_"}},
+				{Name: nil},
+			},
+			expected: &dst.ImportSpec{Name: &dst.Ident{Name: "_"}},
+		},
+		{
+			name: "complex-5",
+			in: []*dst.ImportSpec{
+				{Name: nil},
+				qualifedImport("test1"),
+				{Name: &dst.Ident{Name: "_"}},
+				{Name: nil},
+				qualifedImport("test2"),
+			},
+			expected: qualifedImport("test1"),
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			out := filterExtraneousImports(map[string][]*dst.ImportSpec{
+				"test": tc.in,
+			})
+
+			require.Len(t, out, 1)
+			for v := range out {
+				require.EqualValues(t, tc.expected, v)
+			}
+		})
+	}
+}

internal/injector/injector.go

@@ -96,7 +96,7 @@ func (i *Injector) InjectFiles(files []string) (map[string]InjectedFile, error)
 
 	wg.Add(len(astFiles))
 	for idx, astFile := range astFiles {
-		go func(astFile *ast.File) {
+		go func(idx int, astFile *ast.File) {
 			defer wg.Done()
 
 			decorator := decorator.NewDecoratorWithImports(fset, i.ImportPath, gotypes.New(uses))
@@ -123,7 +123,7 @@ func (i *Injector) InjectFiles(files []string) (map[string]InjectedFile, error)
 			resultMu.Lock()
 			defer resultMu.Unlock()
 			result[files[idx]] = res.InjectedFile
-		}(astFile)
+		}(idx, astFile)
 	}
 	wg.Wait()
 

internal/injector/write.go

@@ -13,13 +13,15 @@ import (
 	"path/filepath"
 
 	"github.com/DataDog/orchestrion/internal/injector/lineinfo"
+	"github.com/DataDog/orchestrion/internal/log"
 	"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) {
+	log.Tracef("writing modified file %q\n", decorator.Filenames[file])
 	canonicalizeImports(file)
 
 	filename := decorator.Filenames[file]