[vpython] Implement smart probing.

vpython/application/probe.go

+// Copyright 2017 The LUCI Authors. All rights reserved.
+// Use of this source code is governed under the Apache License, Version 2.0
+// that can be found in the LICENSE file.
+
+package application
+
+import (
+        "fmt"
+        "os"
+        "os/exec"
+        "strings"
+
+        "golang.org/x/net/context"
+
+        "github.com/luci/luci-go/vpython/python"
+
+        "github.com/luci/luci-go/common/errors"
+        "github.com/luci/luci-go/common/logging"
+        "github.com/luci/luci-go/common/system/environ"
+        "github.com/luci/luci-go/common/system/exitcode"
+        "github.com/luci/luci-go/common/system/prober"
+)
+
+const (
+        // CheckWrapperENV is an environment variable that is used to detect if a
+        // given binary is a "vpython" wrapper instance. The "vpython" Main function
+        // will immediately exit with a non-zero return code if this environment
+        // variable is set.
+        //
+        // See checkWrapper for more details.
+        checkWrapperENV = "VPYTHON_CHECK_WRAPPER"
+
+        // checkWrapperSentinel is text that will be output by a "vpython" instance if
+        // checkWrapperENV is set. If an application exits with a non-zero error code
+        // and outputs this text to STDOUT, then it is considered a "vpython" wrapper.
+        checkWrapperSentinel = checkWrapperENV + ": vpython wrapper check, exiting with non-zero return code"
+)
+
+func wrapperCheck(env environ.Env) bool {
+        if env.GetEmpty(checkWrapperENV) != "" {
+                fmt.Fprintln(os.Stdout, checkWrapperSentinel)
+                return true
+        }
+        return false
+}
+
+type lookPath struct {
+        // probeBase is a base prober whose Self and SelfStat fields will be used
+        // during a lookPathImpl call. This ensures that we only have to call
+        // ResolveSelf once for any given application.
+        probeBase prober.Probe
+
+        env        environ.Env
+        lastPath   string
+        lastOutput string
+}
+
+func (lp *lookPath) look(c context.Context, target string) (*python.LookPathResult, error) {
+        // Construct our probe, derived it from our "probeBase" and tailored to the
+        // specific "target" we're looking for.
+        probe := lp.probeBase
+        probe.Target = target
+        probe.CheckWrapper = lp.checkWrapper
+
+        var result python.LookPathResult
+        var err error
+        if result.Path, err = probe.Locate(c, "", lp.env.Clone()); err != nil {
+                return nil, err
+        }
+
+        // During probing, we try and capture the output of "--version". If we do,
+        // and if we confirm that the "path" returned by the probe matches "lastPath",
+        // return this version along with the resolution.
+        if result.Path == lp.lastPath {
+                var err error
+                if result.Version, err = python.ParseVersionOutput(lp.lastOutput); err == nil {
+                        logging.Debugf(c, "Detected Python version %q from probe candidate [%s]", result.Version, result.Path)
+                } else {
+                        logging.Debugf(c, "Failed to parse version from probe candidate [%s]: %s", result.Path, err)
+                }
+        }
+
+        return &result, nil
+}
+
+// checkWrapper is a prober.CheckWrapperFunc that detects if a given path is a
+// "vpython" instance.
+//
+// It does this by running it "path --version" with CheckWrapperENV set. If
+// the target is a "vpython" wrapper, it will immediately exit with a non-zero
+// value (see Main). If it is not a valid Python program, its behavior is
+// undefined. If it is a valid Pytohn instance, it will exit with a Python
+// version string.
+//
+// As a side effect, the text output of the last version probe will be stored
+// in lp.lastOutput. The "look" function can pass this value on to
+// the LookPathResult.
+func (lp *lookPath) checkWrapper(c context.Context, path string, env environ.Env) (isWrapper bool, err error) {
+        env.Set(checkWrapperENV, "1")
+
+        cmd := exec.CommandContext(c, path, "--version")
+        cmd.Env = env.Sorted()
+
+        output, err := cmd.CombinedOutput()
+        rc, ok := exitcode.Get(err)
+        if !ok {
+                err = errors.Annotate(err).Reason("failed to check if %(path)q is a wrapper").
+                        D("path", path).
+                        Err()
+                return
+        }
+
+        // Retain the last output, so our "look" can reference it to extract the
+        // version string.
+        lp.lastPath = path
+        lp.lastOutput = string(output)
+
+        if rc != 0 {
+                // The target exited with a non-zero error code. It is a wrapper if it
+                // printed the sentinel text.
+                isWrapper = strings.TrimSpace(lp.lastOutput) == checkWrapperSentinel
+                if isWrapper {
+                        // The target was successfully identified as a wrapper. Clear "err", which
+                        // is non-nil b/c of the non-zero exit code, to indicate a successful
+                        // wrapper check.
+                        err = nil
+                        return
+                }
+
+                // The target returned non-zero, but didn't identify as a wrapper. It is
+                // likely something that happens to be named the same thing as the target,
+                // which is an error.
+                err = errors.Annotate(err).Reason("wrapper check returned non-zero").Err()
+        }
+
+        return
+}