Support context cancellation and cleanup

This should avoid leaving temporary build directories lingering after
Ctrl+C.

builder.go

@@ -15,6 +15,7 @@
 package xcaddy
 
 import (
+	"context"
 	"fmt"
 	"io/ioutil"
 	"log"
@@ -35,11 +36,12 @@ import (
 type Builder struct {
 	CaddyVersion string       `json:"caddy_version,omitempty"`
 	Plugins      []Dependency `json:"plugins,omitempty"`
+	Replacements []Replace    `json:"replacements,omitempty"`
 }
 
 // Build builds Caddy at the configured version with the
 // configured plugins and plops down a binary at outputFile.
-func (b Builder) Build(outputFile string) error {
+func (b Builder) Build(ctx context.Context, outputFile string) error {
 	if b.CaddyVersion == "" {
 		return fmt.Errorf("CaddyVersion must be set")
 	}
@@ -55,7 +57,7 @@ func (b Builder) Build(outputFile string) error {
 		return err
 	}
 
-	env, err := b.newEnvironment()
+	env, err := b.newEnvironment(ctx)
 	if err != nil {
 		return err
 	}
@@ -69,7 +71,7 @@ func (b Builder) Build(outputFile string) error {
 		"-trimpath",
 	)
 	cmd.Env = append(os.Environ(), "CGO_ENABLED=0")
-	err = env.runCommand(cmd, 5*time.Minute)
+	err = env.runCommand(ctx, cmd, 5*time.Minute)
 	if err != nil {
 		return err
 	}
@@ -88,10 +90,15 @@ type Dependency struct {
 
 	// The version of the Go module, like used with `go get`.
 	Version string
+}
 
-	// Optional path to a replacement module. Equivalent to
+// Replace represents a Go module replacement.
-	// a `replace` directive in go.mod.
+type Replace struct {
-	Replace string
+	// The import path of the module being replaced.
+	Old string
+
+	// The path to the replacement module.
+	New string
 }
 
 // newTempFolder creates a new folder in a temporary location.

cmd/xcaddy/main.go

@@ -15,10 +15,12 @@
 package main
 
 import (
+	"context"
 	"fmt"
 	"log"
 	"os"
 	"os/exec"
+	"os/signal"
 	"path/filepath"
 	"runtime"
 	"strings"
@@ -28,20 +30,24 @@ import (
 )
 
 func main() {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	go trapSignals(ctx, cancel)
+
 	if len(os.Args) > 1 && os.Args[1] == "build" {
-		if err := runBuild(os.Args[2:]); err != nil {
+		if err := runBuild(ctx, os.Args[2:]); err != nil {
 			log.Fatalf("[ERROR] %v", err)
 		}
 		return
 	}
 
 	// TODO: the caddy version needs to be settable by the user... maybe an env var?
-	if err := runDev("v2.0.0-rc.1", os.Args[1:]); err != nil {
+	if err := runDev(ctx, "v2.0.0-rc.3", os.Args[1:]); err != nil {
 		log.Fatalf("[ERROR] %v", err)
 	}
 }
 
-func runBuild(args []string) error {
+func runBuild(ctx context.Context, args []string) error {
 	// parse the command line args... rather primitively
 	var caddyVersion, output string
 	var plugins []xcaddy.Dependency
@@ -93,7 +99,7 @@ func runBuild(args []string) error {
 		CaddyVersion: caddyVersion,
 		Plugins:      plugins,
 	}
-	err := builder.Build(output)
+	err := builder.Build(ctx, output)
 	if err != nil {
 		log.Fatalf("[FATAL] %v", err)
 	}
@@ -115,7 +121,7 @@ func runBuild(args []string) error {
 	return nil
 }
 
-func runDev(caddyVersion string, args []string) error {
+func runDev(ctx context.Context, caddyVersion string, args []string) error {
 	const binOutput = "./caddy"
 
 	// get current/main module name
@@ -134,17 +140,44 @@ func runDev(caddyVersion string, args []string) error {
 	}
 	moduleDir := strings.TrimSpace(string(out))
 
+	// make sure the module being developed is replaced
+	// so that the local copy is used
+	replacements := []xcaddy.Replace{
+		{
+			Old: currentModule,
+			New: moduleDir,
+		},
+	}
+
+	// replace directives only apply to the top-level/main go.mod,
+	// and since this tool is a carry-through for the user's actual
+	// go.mod, we need to transfer their replace directives through
+	// to the one we're making
+	cmd = exec.Command("go", "list", "-m", "-f={{if .Replace}}{{.Path}} => {{.Replace}}{{end}}", "all")
+	out, err = cmd.Output()
+	if err != nil {
+		return err
+	}
+	for _, line := range strings.Split(string(out), "\n") {
+		parts := strings.Split(line, "=>")
+		if len(parts) != 2 || parts[0] == "" || parts[1] == "" {
+			continue
+		}
+		replacements = append(replacements, xcaddy.Replace{
+			Old: strings.TrimSpace(parts[0]),
+			New: strings.TrimSpace(parts[1]),
+		})
+	}
+
 	// build caddy with this module plugged in
 	builder := xcaddy.Builder{
 		CaddyVersion: caddyVersion,
 		Plugins: []xcaddy.Dependency{
-			{
+			{ModulePath: currentModule},
-				ModulePath: currentModule,
-				Replace:    moduleDir,
-			},
 		},
+		Replacements: replacements,
 	}
-	err = builder.Build(binOutput)
+	err = builder.Build(ctx, binOutput)
 	if err != nil {
 		return err
 	}
@@ -168,9 +201,22 @@ func runDev(caddyVersion string, args []string) error {
 	}
 	defer cleanup()
 	go func() {
-		time.Sleep(2 * time.Second)
+		time.Sleep(5 * time.Second)
 		cleanup()
 	}()
 
 	return cmd.Wait()
 }
+
+func trapSignals(ctx context.Context, cancel context.CancelFunc) {
+	sig := make(chan os.Signal, 1)
+	signal.Notify(sig, os.Interrupt)
+
+	select {
+	case <-sig:
+		log.Printf("[INFO] SIGINT: Shutting down")
+		cancel()
+	case <-ctx.Done():
+		return
+	}
+}

environment.go

@@ -16,6 +16,7 @@ package xcaddy
 
 import (
 	"bytes"
+	"context"
 	"fmt"
 	"html/template"
 	"io/ioutil"
@@ -27,7 +28,7 @@ import (
 	"time"
 )
 
-func (b Builder) newEnvironment() (*environment, error) {
+func (b Builder) newEnvironment(ctx context.Context) (*environment, error) {
 	// assume Caddy v2 if no semantic version is provided
 	caddyModulePath := defaultCaddyModulePath
 	if !strings.HasPrefix(b.CaddyVersion, "v") || !strings.Contains(b.CaddyVersion, ".") {
@@ -47,11 +48,11 @@ func (b Builder) newEnvironment() (*environment, error) {
 	}
 
 	// create the context for the main module template
-	ctx := moduleTemplateContext{
+	tplCtx := goModTemplateContext{
 		CaddyModule: caddyModulePath,
 	}
 	for _, p := range b.Plugins {
-		ctx.Plugins = append(ctx.Plugins, p.ModulePath)
+		tplCtx.Plugins = append(tplCtx.Plugins, p.ModulePath)
 	}
 
 	// evaluate the template for the main module
@@ -60,7 +61,7 @@ func (b Builder) newEnvironment() (*environment, error) {
 	if err != nil {
 		return nil, err
 	}
-	err = tpl.Execute(&buf, ctx)
+	err = tpl.Execute(&buf, tplCtx)
 	if err != nil {
 		return nil, err
 	}
@@ -98,39 +99,52 @@ func (b Builder) newEnvironment() (*environment, error) {
 	// initialize the go module
 	log.Println("[INFO] Initializing Go module")
 	cmd := env.newCommand("go", "mod", "init", "caddy")
-	err = env.runCommand(cmd, 10*time.Second)
+	err = env.runCommand(ctx, cmd, 10*time.Second)
 	if err != nil {
 		return nil, err
 	}
 
 	// specify module replacements before pinning versions
-	for _, p := range b.Plugins {
+	replaced := make(map[string]string)
-		if p.Replace == "" {
+	for _, r := range b.Replacements {
-			continue
+		log.Printf("[INFO] Replace %s => %s", r.Old, r.New)
-		}
-		log.Printf("[INFO] Replace %s => %s", p.ModulePath, p.Replace)
 		cmd := env.newCommand("go", "mod", "edit",
-			"-replace", fmt.Sprintf("%s=%s", p.ModulePath, p.Replace))
+			"-replace", fmt.Sprintf("%s=%s", r.Old, r.New))
-		err := env.runCommand(cmd, 10*time.Second)
+		err := env.runCommand(ctx, cmd, 10*time.Second)
 		if err != nil {
 			return nil, err
 		}
+		replaced[r.Old] = r.New
+	}
+
+	// check for early abort
+	select {
+	case <-ctx.Done():
+		return nil, ctx.Err()
+	default:
 	}
 
 	// pin versions by populating go.mod, first for Caddy itself and then plugins
 	log.Println("[INFO] Pinning versions")
-	err = env.execGoGet(caddyModulePath, b.CaddyVersion)
+	err = env.execGoGet(ctx, caddyModulePath, b.CaddyVersion)
 	if err != nil {
 		return nil, err
 	}
 	for _, p := range b.Plugins {
-		if p.Replace != "" {
+		// if module is locally available; do not "go get" it
+		if replaced[p.ModulePath] != "" {
 			continue
 		}
-		err = env.execGoGet(p.ModulePath, p.Version)
+		err = env.execGoGet(ctx, p.ModulePath, p.Version)
 		if err != nil {
 			return nil, err
 		}
+		// check for early abort
+		select {
+		case <-ctx.Done():
+			return nil, ctx.Err()
+		default:
+		}
 	}
 
 	log.Println("[INFO] Build environment ready")
@@ -160,38 +174,59 @@ func (env environment) newCommand(command string, args ...string) *exec.Cmd {
 	return cmd
 }
 
-func (env environment) runCommand(cmd *exec.Cmd, timeout time.Duration) error {
+func (env environment) runCommand(ctx context.Context, cmd *exec.Cmd, timeout time.Duration) error {
 	log.Printf("[INFO] exec (timeout=%s): %+v ", timeout, cmd)
 
-	// no timeout? this is easy; just run it
+	if timeout > 0 {
-	if timeout == 0 {
+		var cancel context.CancelFunc
-		return cmd.Run()
+		ctx, cancel = context.WithTimeout(ctx, timeout)
+		defer cancel()
 	}
 
-	// otherwise start it and use a timer
+	// start the command; if it fails to start, report error immediately
 	err := cmd.Start()
 	if err != nil {
 		return err
 	}
-	timer := time.AfterFunc(timeout, func() {
+
-		err = fmt.Errorf("timed out (builder-enforced)")
+	// wait for the command in a goroutine; the reason for this is
-		cmd.Process.Kill()
+	// very subtle: if, in our select, we do `case cmdErr := <-cmd.Wait()`,
-	})
+	// then that case would be chosen immediately, because cmd.Wait() is
-	waitErr := cmd.Wait()
+	// immediately available (even though it blocks for potentially a long
-	timer.Stop()
+	// time, it can be evaluated immediately). So we have to remove that
-	if err != nil {
+	// evaluation from the `case` statement.
-		return err
+	cmdErrChan := make(chan error)
+	go func() {
+		cmdErrChan <- cmd.Wait()
+	}()
+
+	// unblock either when the command finishes, or when the done
+	// channel is closed -- whichever comes first
+	select {
+	case cmdErr := <-cmdErrChan:
+		// process ended; report any error immediately
+		return cmdErr
+	case <-ctx.Done():
+		// context was canceled, either due to timeout or
+		// maybe a signal from higher up canceled the parent
+		// context; presumably, the OS also sent the signal
+		// to the child process, so wait for it to die
+		select {
+		case <-time.After(15 * time.Second):
+			cmd.Process.Kill()
+		case <-cmdErrChan:
+		}
+		return ctx.Err()
 	}
-	return waitErr
 }
 
-func (env environment) execGoGet(modulePath, moduleVersion string) error {
+func (env environment) execGoGet(ctx context.Context, modulePath, moduleVersion string) error {
 	mod := modulePath + "@" + moduleVersion
 	cmd := env.newCommand("go", "get", "-d", "-v", mod)
-	return env.runCommand(cmd, 60*time.Second)
+	return env.runCommand(ctx, cmd, 5*time.Minute)
 }
 
-type moduleTemplateContext struct {
+type goModTemplateContext struct {
 	CaddyModule string
 	Plugins     []string
 }