github.com/ratrocket/u-root@v0.0.0-20180201221235-1cf9f48ee2cf/cmds/installcommand/installcommand.go

// Copyright 2012-2017 the u-root Authors. All rights reserved
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

// Install command from a go source file.
//
// Synopsis:
//     SYMLINK [ARGS...]
//     installcommand [INSTALLCOMMAND_ARGS...] COMMAND [ARGS...]
//
// Description:
//     u-root commands are lazily compiled. Uncompiled commands in the /bin
//     directory are symbolic links to installcommand. When executed through
//     the symbolic link, installcommand will build the command from source and
//     exec it.
//
//     The second form allows commands to be installed and exec'ed without a
//     symbolic link. In this form additional arguments such as `-v` and
//     `-ludicrous` can be passed into installcommand.
//
// Options:
//     -lowpri:    the scheduler priority to lowered before starting
//     -ludicrous: print out ALL the output from the go build commands
//     -onlybuild:     just build, do not exec
//     -noforce:   do not build if a file already exists at the destination
//     -v:         print all build commands
import (
	"flag"
	"fmt"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"syscall"

	"github.com/u-root/u-root/pkg/uroot/util"
)

var (
	urpath    = "/go/bin:/ubin:/buildbin:/bbin:/usr/local/bin:"
	lowpri    = flag.Bool("lowpri", false, "the scheduler priority is lowered before starting")
	ludicrous = flag.Bool("ludicrous", false, "print out ALL the output from the go build commands")
	onlybuild = flag.Bool("onlybuild", false, "only build, i.e. do not execute the process after building")
	noforce   = flag.Bool("noforce", false, "do not build if a file already exists at the destination")
	verbose   = flag.Bool("v", false, "print all build commands")
	debug     = func(string, ...interface{}) {}
)

type form struct {
	// Name of the command, ex: "ls"
	cmdName string
	// Args passed to the command, ex: {"-l", "-R"}
	cmdArgs []string
	// Args intended for installcommand
	lowPri    bool
	ludicrous bool
	onlybuild bool
	noForce   bool
	verbose   bool
}

func usage() {
	fmt.Fprintf(os.Stderr, "Usage: installcommand [INSTALLCOMMAND_ARGS...] COMMAND [ARGS...]\n")
	os.Exit(2)
}

// Parse the command line to determine the form.
func parseCommandLine() form {
	// First form:
	//     SYMLINK [ARGS...]
	if !strings.HasSuffix(os.Args[0], "installcommand") {
		return form{
			cmdName: filepath.Base(os.Args[0]),
			cmdArgs: os.Args[1:],
		}
	}

	// Second form:
	//     installcommand [INSTALLCOMMAND_ARGS...] COMMAND [ARGS...]
	flag.Parse()
	if flag.NArg() < 1 {
		log.Println("Second form requires a COMMAND argument")
		usage()
	}
	return form{
		cmdName:   flag.Arg(0),
		cmdArgs:   flag.Args()[1:],
		lowPri:    *lowpri,
		ludicrous: *ludicrous,
		onlybuild: *onlybuild,
		noForce:   *noforce,
		verbose:   *verbose,
	}
}

// run runs the command with the information from form.
// Since run can potentially never return, since it can use Exec,
// it should never return in any other case. Hence, if all goes well
// at the end, we os.Exit(0)
func run(n string, form form) {
	if os.Getenv("INSTALLCOMMAND_NOEXEC") == "" {
		err := syscall.Exec(n, append([]string{form.cmdName}, form.cmdArgs...), os.Environ())
		// Regardless of whether err is nil, if Exec returns at all, it failed
		// at its job. Print an error and then let's see if a normal run can succeed.
		log.Printf("Failed to exec %s: %v", form.cmdName, err)
	}

	cmd := exec.Command(n)
	cmd.Args = append([]string{form.cmdName}, form.cmdArgs...)
	cmd.Stdin = os.Stdin
	cmd.Stderr = os.Stderr
	cmd.Stdout = os.Stdout
	if err := cmd.Run(); err != nil {
		exitErr, ok := err.(*exec.ExitError)
		if !ok {
			log.Fatal(err)
		}
		exitWithStatus(exitErr)
	}
	os.Exit(0)
}

func main() {
	form := parseCommandLine()

	if form.lowPri {
		if err := syscall.Setpriority(syscall.PRIO_PROCESS, 0, 20); err != nil {
			log.Printf("Cannot set low priority: %v", err)
		}
	}

	a := []string{"install"}

	if form.verbose {
		debug = log.Printf
		a = append(a, "-x")
	}

	debug("Command name: %v\n", form.cmdName)
	destDir := "/ubin"
	destFile := filepath.Join(destDir, form.cmdName)

	// Is the command there? This covers a race condition
	// in that some other process may have caused it to be
	// built.
	_, err := os.Stat(destFile)

	// The file exists.
	if err == nil {
		if form.onlybuild {
			os.Exit(0)
		}
		run(destFile, form)
	}

	// If we are here, things did not go so well. We have to build
	// the command.  Which means we have to find the source.  Now
	// that we can add new commands, we need to find out where the
	// command is.  The current rule is that command source is
	// either in /src/github.com/u-root/u-root/cmds/cmdName or
	// /src/*/*/cmdName.  We stat that first name and, if it's there,
	// we're good to go; if not, we do the glob.
	// We considered building these paths into install command via an
	// init() function in a file generated by initramfs. We went with this
	// simpler approach as it allows people to add new commands even after boot:
	// it suffices to, e.g.
	// go get github.com/new/fancyshell
	// ln -s /buildbin/installcommand /buildbin/fancyshell
	// and the glob will find it. If we get to the point that the glob
	// no longer works we can go with filepath.Walk (which glob uses anyway)
	// but for now this works.
	var src string
	for _, p := range util.CmdsGlob {
		g := filepath.Join("/src", p, form.cmdName)
		debug("Check %v", g)
		l, err := filepath.Glob(g)
		debug("glob is %v %v", l, err)
		if err != nil {
			log.Printf("Glob %v fails: %v", g, err)
			continue
		}
		if len(l) == 0 {
			continue
		}

		if len(l) > 1 {
			log.Printf("Found these paths for %v: %v; going with %v", form.cmdName, l, l[0])
		}
		src = l[0]
		break
	}

	if src == "" {
		log.Fatalf("Can not find source code for %v", form.cmdName)
	}

	r, err := filepath.Rel("/src", src)
	if err != nil {
		log.Fatalf("Can't take rel of %v: %v", src, err)
	}
	cmd := exec.Command("go", append(a, r)...)

	// Set GOGC if unset. The best value is determined empirically and
	// depends on the machine and Go version. For the workload of compiling
	// a small Go program, values larger than the default perform better.
	// See: /scripts/build_perf.sh
	if _, ok := os.LookupEnv("GOGC"); !ok {
		cmd.Env = append(os.Environ(), "GOGC=400")
	}

	cmd.Dir = "/"

	debug("Run %v", cmd)
	out, err := cmd.CombinedOutput()
	debug("installcommand: go build returned")

	if err != nil {
		p := os.Getenv("PATH")
		log.Fatalf("installcommand: trying to build {cmdName: %v, PATH %s, err %v, out %s}", form.cmdName, p, err, out)
	}

	if *ludicrous {
		debug(string(out))
	}
	if !form.onlybuild {
		run(destFile, form)
	}
}