Toolexec example

main.go

package main

// This function will be automatically called from the injected main with toolexec-example.go.
func greet() string {
	return "Hello, World!"
}

Makefile.toolexec

@PHONY: test

test:
	$(eval TMP := $(shell mktemp -d))
	go build -o $(TMP)/toolexec-example.run toolexec-example.go
	go build -a -toolexec=$(TMP)/toolexec-example.run -o hello main.go
  ./hello && go clean
	rm -rf $(TMP)

toolexec-example.go

package main

import (
	"bytes"
	"errors"
	"go/ast"
	"go/parser"
	"go/printer"
	"go/token"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
)

var (
	ErrNotTarget = errors.New("not a target")
)

func main() {
	if len(os.Args) < 3 {
		log.Fatal("usage: go build -toolexec=/path/to/this/binary")
	}
	tmpDir, err := os.MkdirTemp("", "toolexec-*.go")
	if err != nil {
		log.Fatal(err)
	}
	defer os.RemoveAll(tmpDir)
	cmd := parseCommand(os.Args[1], os.Args[2:])
	if err := cmd.wrap(tmpDir); err != nil {
		log.Print(err)
	}
	if err := cmd.run(); err != nil {
		log.Fatal(err)
	}
}

type command struct {
	tool     string
	toolPath string
	args     []string
}

func parseCommand(toolPath string, args []string) *command {
	tool := filepath.Base(toolPath)
	if ext := filepath.Ext(tool); ext != "" {
		// Remove extension
		tool = tool[:len(tool)-len(ext)]
	}
	return &command{
		tool:     tool,
		toolPath: toolPath,
		args:     args,
	}
}

func (c *command) wrap(tmpDir string) error {
	switch c.tool {
	case "compile":
		return c.compile(tmpDir)
	case "link":
		fallthrough
	default:
		return nil
	}
}

func (c *command) run() error {
	cmd := exec.Command(c.toolPath, c.args...)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	cmd.Stdin = os.Stdin
	return cmd.Run()
}

func (c *command) compile(tmpDir string) error {
	var (
		cfg     string
		files   []string
		filesAt int
	)
	for i, arg := range c.args {
		if strings.HasPrefix(arg, "-") {
			continue
		}
		if strings.Contains(arg, "/b001/importcfg") {
			cfg = arg
		} else if strings.HasSuffix(arg, ".go") {
			files = c.args[i:]
			filesAt = i
			break
		}
	}
	if cfg == "" || len(files) == 0 {
		return ErrNotTarget
	}
	var (
		err    error
		target *ast.File
	)
	fset := token.NewFileSet()
	target, err = parser.ParseFile(fset, files[0], nil, parser.ParseComments)
	if err != nil {
		return err
	}
	appendMain(fset, target)
	var output []byte
	buffer := bytes.NewBuffer(output)
	err = printer.Fprint(buffer, fset, target)
	if err != nil {
		return err
	}
	targetPath := filepath.Join(tmpDir, filepath.Base(files[0]))
	err = os.WriteFile(targetPath, buffer.Bytes(), 0644)
	if err != nil {
		return err
	}
	c.args[filesAt] = targetPath
	return nil
}

var payload = `package main

	func main() {
		println(greet())
	}`

func appendMain(fset *token.FileSet, target *ast.File) {
	payload, _ := parser.ParseFile(fset, "", payload, parser.ParseComments)
	target.Decls = append(target.Decls, payload.Decls...)
}

toolexec.go

// Toolexec example
// It showcases a minimal example using toolexec to inject code.