<commit_before>package control

import (
	"bufio"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"strings"

	log "github.com/Sirupsen/logrus"
	"gopkg.in/yaml.v2"

	dockerClient "github.com/fsouza/go-dockerclient"

	"github.com/codegangsta/cli"
	"github.com/rancherio/os/cmd/power"
	"github.com/rancherio/os/config"
	"github.com/rancherio/os/docker"
)

type Images struct {
	Current   string   `yaml:"current,omitempty"`
	Available []string `yaml:"available,omitempty"`
}

func osSubcommands() []cli.Command {
	return []cli.Command{
		{
			Name:   "upgrade",
			Usage:  "upgrade to latest version",
			Action: osUpgrade,
			Flags: []cli.Flag{
				cli.BoolFlag{
					Name:  "stage, s",
					Usage: "Only stage the new upgrade, don't apply it",
				},
				cli.StringFlag{
					Name:  "image, i",
					Usage: "upgrade to a certain image",
				},
				cli.BoolFlag{
					Name:  "force, f",
					Usage: "do not prompt for input",
				},
			},
		},
		{
			Name:   "list",
			Usage:  "list the current available versions",
			Action: osMetaDataGet,
		},
		{
			Name:   "version",
			Usage:  "show the currently installed version",
			Action: osVersion,
		},
	}
}

func getImages() (*Images, error) {
	upgradeUrl, err := getUpgradeUrl()
	if err != nil {
		return nil, err
	}

	var body []byte

	if strings.HasPrefix(upgradeUrl, "/") {
		body, err = ioutil.ReadFile(upgradeUrl)
		if err != nil {
			return nil, err
		}
	} else {
		u, err := url.Parse(upgradeUrl)
		if err != nil {
			return nil, err
		}

		q := u.Query()
		q.Set("current", config.VERSION)
		u.RawQuery = q.Encode()
		upgradeUrl = u.String()

		resp, err := http.Get(upgradeUrl)
		if err != nil {
			return nil, err
		}
		defer resp.Body.Close()
		body, err = ioutil.ReadAll(resp.Body)
		if err != nil {
			return nil, err
		}
	}

	return parseBody(body)
}

func osMetaDataGet(c *cli.Context) {
	images, err := getImages()
	if err != nil {
		log.Fatal(err)
	}

	client, err := docker.NewSystemClient()
	if err != nil {
		log.Fatal(err)
	}

	for _, image := range images.Available {
		_, err := client.InspectImage(image)
		if err == dockerClient.ErrNoSuchImage {
			fmt.Println(image, "remote")
		} else {
			fmt.Println(image, "local")
		}
	}
}

func getLatestImage() (string, error) {
	images, err := getImages()
	if err != nil {
		return "", err
	}

	return images.Current, nil
}

func osUpgrade(c *cli.Context) {
	image := c.String("image")

	if image == "" {
		var err error
		image, err = getLatestImage()
		if err != nil {
			log.Fatal(err)
		}
		if image == "" {
			log.Fatal("Failed to find latest image")
		}
	}
	startUpgradeContainer(image, c.Bool("stage"), c.Bool("force"))
}

func osVersion(c *cli.Context) {
	fmt.Println(config.VERSION)
}

func yes(in *bufio.Reader, question string) bool {
	fmt.Printf("%s [y/N]: ", question)
	line, err := in.ReadString('\n')
	if err != nil {
		log.Fatal(err)
	}

	return strings.ToLower(line[0:1]) == "y"
}

func startUpgradeContainer(image string, stage, force bool) {
	in := bufio.NewReader(os.Stdin)

	container := docker.NewContainer(config.DOCKER_SYSTEM_HOST, &config.ContainerConfig{
		Cmd: "--name=os-upgrade " +
			"--log-driver=json-file" +
			"--rm " +
			"--privileged " +
			"--net=host " +
			image + " " +
			"-t rancher-upgrade " +
			"-r " + config.VERSION,
	}).Stage()

	if container.Err != nil {
		log.Fatal(container.Err)
	}

	if !stage {
		fmt.Printf("Upgrading to %s\n", image)

		if !force {
			if !yes(in, "Continue") {
				os.Exit(1)
			}
		}

		container.Start()
		if container.Err != nil {
			log.Fatal(container.Err)
		}

		client, err := docker.NewClient(config.DOCKER_SYSTEM_HOST)
		if err != nil {
			log.Fatal(err)
		}

		go func() {
			client.Logs(dockerClient.LogsOptions{
				Container:    container.Container.ID,
				OutputStream: os.Stdout,
				ErrorStream:  os.Stderr,
				Follow:       true,
				Stdout:       true,
				Stderr:       true,
			})
		}()

		exit, err := client.WaitContainer(container.Container.ID)
		if err != nil {
			log.Fatal(err)
		}

		if container.Err != nil {
			log.Fatal(container.Err)
		}

		if exit == 0 {
			if force || yes(in, "Continue with reboot") {
				log.Info("Rebooting")
				power.Reboot()
			}
		} else {
			log.Error("Upgrade failed")
			os.Exit(exit)
		}
	}
}

func parseBody(body []byte) (*Images, error) {
	update := &Images{}
	err := yaml.Unmarshal(body, update)
	if err != nil {
		return nil, err
	}

	return update, nil
}

func getUpgradeUrl() (string, error) {
	cfg, err := config.LoadConfig()
	if err != nil {
		return "", err
	}

	return cfg.Upgrade.Url, nil
}
<commit_msg>Fix broken OS upgrade<commit_after>package control

import (
	"bufio"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"strings"

	log "github.com/Sirupsen/logrus"
	"gopkg.in/yaml.v2"

	dockerClient "github.com/fsouza/go-dockerclient"

	"github.com/codegangsta/cli"
	"github.com/rancherio/os/cmd/power"
	"github.com/rancherio/os/config"
	"github.com/rancherio/os/docker"
)

type Images struct {
	Current   string   `yaml:"current,omitempty"`
	Available []string `yaml:"available,omitempty"`
}

func osSubcommands() []cli.Command {
	return []cli.Command{
		{
			Name:   "upgrade",
			Usage:  "upgrade to latest version",
			Action: osUpgrade,
			Flags: []cli.Flag{
				cli.BoolFlag{
					Name:  "stage, s",
					Usage: "Only stage the new upgrade, don't apply it",
				},
				cli.StringFlag{
					Name:  "image, i",
					Usage: "upgrade to a certain image",
				},
				cli.BoolFlag{
					Name:  "force, f",
					Usage: "do not prompt for input",
				},
			},
		},
		{
			Name:   "list",
			Usage:  "list the current available versions",
			Action: osMetaDataGet,
		},
		{
			Name:   "version",
			Usage:  "show the currently installed version",
			Action: osVersion,
		},
	}
}

func getImages() (*Images, error) {
	upgradeUrl, err := getUpgradeUrl()
	if err != nil {
		return nil, err
	}

	var body []byte

	if strings.HasPrefix(upgradeUrl, "/") {
		body, err = ioutil.ReadFile(upgradeUrl)
		if err != nil {
			return nil, err
		}
	} else {
		u, err := url.Parse(upgradeUrl)
		if err != nil {
			return nil, err
		}

		q := u.Query()
		q.Set("current", config.VERSION)
		u.RawQuery = q.Encode()
		upgradeUrl = u.String()

		resp, err := http.Get(upgradeUrl)
		if err != nil {
			return nil, err
		}
		defer resp.Body.Close()
		body, err = ioutil.ReadAll(resp.Body)
		if err != nil {
			return nil, err
		}
	}

	return parseBody(body)
}

func osMetaDataGet(c *cli.Context) {
	images, err := getImages()
	if err != nil {
		log.Fatal(err)
	}

	client, err := docker.NewSystemClient()
	if err != nil {
		log.Fatal(err)
	}

	for _, image := range images.Available {
		_, err := client.InspectImage(image)
		if err == dockerClient.ErrNoSuchImage {
			fmt.Println(image, "remote")
		} else {
			fmt.Println(image, "local")
		}
	}
}

func getLatestImage() (string, error) {
	images, err := getImages()
	if err != nil {
		return "", err
	}

	return images.Current, nil
}

func osUpgrade(c *cli.Context) {
	image := c.String("image")

	if image == "" {
		var err error
		image, err = getLatestImage()
		if err != nil {
			log.Fatal(err)
		}
		if image == "" {
			log.Fatal("Failed to find latest image")
		}
	}
	startUpgradeContainer(image, c.Bool("stage"), c.Bool("force"))
}

func osVersion(c *cli.Context) {
	fmt.Println(config.VERSION)
}

func yes(in *bufio.Reader, question string) bool {
	fmt.Printf("%s [y/N]: ", question)
	line, err := in.ReadString('\n')
	if err != nil {
		log.Fatal(err)
	}

	return strings.ToLower(line[0:1]) == "y"
}

func startUpgradeContainer(image string, stage, force bool) {
	in := bufio.NewReader(os.Stdin)

	container := docker.NewContainer(config.DOCKER_SYSTEM_HOST, &config.ContainerConfig{
		Cmd: "--name=os-upgrade " +
			"--log-driver=json-file " +
			"--rm " +
			"--privileged " +
			"--net=host " +
			image + " " +
			"-t rancher-upgrade " +
			"-r " + config.VERSION,
	}).Stage()

	if container.Err != nil {
		log.Fatal(container.Err)
	}

	if !stage {
		fmt.Printf("Upgrading to %s\n", image)

		if !force {
			if !yes(in, "Continue") {
				os.Exit(1)
			}
		}

		container.Start()
		if container.Err != nil {
			log.Fatal(container.Err)
		}

		client, err := docker.NewClient(config.DOCKER_SYSTEM_HOST)
		if err != nil {
			log.Fatal(err)
		}

		go func() {
			client.Logs(dockerClient.LogsOptions{
				Container:    container.Container.ID,
				OutputStream: os.Stdout,
				ErrorStream:  os.Stderr,
				Follow:       true,
				Stdout:       true,
				Stderr:       true,
			})
		}()

		exit, err := client.WaitContainer(container.Container.ID)
		if err != nil {
			log.Fatal(err)
		}

		if container.Err != nil {
			log.Fatal(container.Err)
		}

		if exit == 0 {
			if force || yes(in, "Continue with reboot") {
				log.Info("Rebooting")
				power.Reboot()
			}
		} else {
			log.Error("Upgrade failed")
			os.Exit(exit)
		}
	}
}

func parseBody(body []byte) (*Images, error) {
	update := &Images{}
	err := yaml.Unmarshal(body, update)
	if err != nil {
		return nil, err
	}

	return update, nil
}

func getUpgradeUrl() (string, error) {
	cfg, err := config.LoadConfig()
	if err != nil {
		return "", err
	}

	return cfg.Upgrade.Url, nil
}
<|endoftext|>


<commit_before>// Copyright 2015 The Go 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

import (
	"context"
	"flag"
	"fmt"
	"io"
	"os"
	"strings"

	"golang.org/x/build/buildlet"
	"golang.org/x/build/dashboard"
	"golang.org/x/build/internal/envutil"
	"golang.org/x/build/internal/gomote/protos"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
)

func legacyRun(args []string) error {
	if activeGroup != nil {
		return fmt.Errorf("command does not support groups")
	}

	fs := flag.NewFlagSet("run", flag.ContinueOnError)
	fs.Usage = func() {
		fmt.Fprintln(os.Stderr, "run usage: gomote run [run-opts] <instance> <cmd> [args...]")
		fs.PrintDefaults()
		os.Exit(1)
	}
	var sys bool
	fs.BoolVar(&sys, "system", false, "run inside the system, and not inside the workdir; this is implicit if cmd starts with '/'")
	var debug bool
	fs.BoolVar(&debug, "debug", false, "write debug info about the command's execution before it begins")
	var env stringSlice
	fs.Var(&env, "e", "Environment variable KEY=value. The -e flag may be repeated multiple times to add multiple things to the environment.")
	var firewall bool
	fs.BoolVar(&firewall, "firewall", false, "Enable outbound firewall on machine. This is on by default on many builders (where supported) but disabled by default on gomote for ease of debugging. Once any command has been run with the -firewall flag on, it's on for the lifetime of that gomote instance.")
	var path string
	fs.StringVar(&path, "path", "", "Comma-separated list of ExecOpts.Path elements. The special string 'EMPTY' means to run without any $PATH. The empty string (default) does not modify the $PATH. Otherwise, the following expansions apply: the string '$PATH' expands to the current PATH element(s), the substring '$WORKDIR' expands to the buildlet's temp workdir.")

	var dir string
	fs.StringVar(&dir, "dir", "", "Directory to run from. Defaults to the directory of the command, or the work directory if -system is true.")
	var builderEnv string
	fs.StringVar(&builderEnv, "builderenv", "", "Optional alternate builder to act like. Must share the same underlying buildlet host type, or it's an error. For instance, linux-amd64-race or linux-386-387 are compatible with linux-amd64, but openbsd-amd64 and openbsd-386 are different hosts.")

	fs.Parse(args)
	if fs.NArg() < 2 {
		fs.Usage()
	}
	name, cmd := fs.Arg(0), fs.Arg(1)

	var conf *dashboard.BuildConfig

	bc, conf, err := clientAndConf(name)
	if err != nil {
		return err
	}

	if builderEnv != "" {
		altConf, ok := dashboard.Builders[builderEnv]
		if !ok {
			return fmt.Errorf("unknown --builderenv=%q builder value", builderEnv)
		}
		if altConf.HostType != conf.HostType {
			return fmt.Errorf("--builderEnv=%q has host type %q, which is not compatible with the named buildlet's host type %q",
				builderEnv, altConf.HostType, conf.HostType)
		}
		conf = altConf
	}

	var pathOpt []string
	if path == "EMPTY" {
		pathOpt = []string{} // non-nil
	} else if path != "" {
		pathOpt = strings.Split(path, ",")
	}
	env = append(env, "GO_DISABLE_OUTBOUND_NETWORK="+fmt.Sprint(firewall))

	remoteErr, execErr := bc.Exec(context.Background(), cmd, buildlet.ExecOpts{
		Dir:         dir,
		SystemLevel: sys || strings.HasPrefix(cmd, "/"),
		Output:      os.Stdout,
		Args:        fs.Args()[2:],
		ExtraEnv:    envutil.Dedup(conf.GOOS(), append(conf.Env(), []string(env)...)),
		Debug:       debug,
		Path:        pathOpt,
	})
	if execErr != nil {
		return fmt.Errorf("Error trying to execute %s: %v", cmd, execErr)
	}
	return remoteErr
}

// stringSlice implements flag.Value, specifically for storing environment
// variable key=value pairs.
type stringSlice []string

func (*stringSlice) String() string { return "" } // default value

func (ss *stringSlice) Set(v string) error {
	if v != "" {
		if !strings.Contains(v, "=") {
			return fmt.Errorf("-e argument %q doesn't contains an '=' sign.", v)
		}
		*ss = append(*ss, v)
	}
	return nil
}

func run(args []string) error {
	if activeGroup != nil {
		return fmt.Errorf("command does not yet support groups")
	}

	fs := flag.NewFlagSet("run", flag.ContinueOnError)
	fs.Usage = func() {
		fmt.Fprintln(os.Stderr, "run usage: gomote run [run-opts] <instance> <cmd> [args...]")
		fs.PrintDefaults()
		os.Exit(1)
	}
	var sys bool
	fs.BoolVar(&sys, "system", false, "run inside the system, and not inside the workdir; this is implicit if cmd starts with '/'")
	var debug bool
	fs.BoolVar(&debug, "debug", false, "write debug info about the command's execution before it begins")
	var env stringSlice
	fs.Var(&env, "e", "Environment variable KEY=value. The -e flag may be repeated multiple times to add multiple things to the environment.")
	var firewall bool
	fs.BoolVar(&firewall, "firewall", false, "Enable outbound firewall on machine. This is on by default on many builders (where supported) but disabled by default on gomote for ease of debugging. Once any command has been run with the -firewall flag on, it's on for the lifetime of that gomote instance.")
	var path string
	fs.StringVar(&path, "path", "", "Comma-separated list of ExecOpts.Path elements. The special string 'EMPTY' means to run without any $PATH. The empty string (default) does not modify the $PATH. Otherwise, the following expansions apply: the string '$PATH' expands to the current PATH element(s), the substring '$WORKDIR' expands to the buildlet's temp workdir.")

	var dir string
	fs.StringVar(&dir, "dir", "", "Directory to run from. Defaults to the directory of the command, or the work directory if -system is true.")
	var builderEnv string
	fs.StringVar(&builderEnv, "builderenv", "", "Optional alternate builder to act like. Must share the same underlying buildlet host type, or it's an error. For instance, linux-amd64-race or linux-386-387 are compatible with linux-amd64, but openbsd-amd64 and openbsd-386 are different hosts.")

	fs.Parse(args)
	if fs.NArg() < 2 {
		fs.Usage()
	}
	name, cmd := fs.Arg(0), fs.Arg(1)
	var pathOpt []string
	if path == "EMPTY" {
		pathOpt = []string{} // non-nil
	} else if path != "" {
		pathOpt = strings.Split(path, ",")
	}
	env = append(env, "GO_DISABLE_OUTBOUND_NETWORK="+fmt.Sprint(firewall))

	ctx := context.Background()
	client := gomoteServerClient(ctx)
	stream, err := client.ExecuteCommand(ctx, &protos.ExecuteCommandRequest{
		AppendEnvironment: []string(env),
		Args:              fs.Args()[2:],
		Command:           cmd,
		Debug:             debug,
		Directory:         dir,
		GomoteId:          name,
		Path:              pathOpt,
		SystemLevel:       sys || strings.HasPrefix(cmd, "/"),
		ImitateHostType:   builderEnv,
	})
	if err != nil {
		return fmt.Errorf("unable to execute %s: %s", cmd, statusFromError(err))
	}
	for {
		update, err := stream.Recv()
		if err == io.EOF {
			return nil
		}
		if err != nil {
			// execution error
			if status.Code(err) == codes.Aborted {
				return fmt.Errorf("Error trying to execute %s: %v", cmd, statusFromError(err))
			}
			// remote error
			return fmt.Errorf("unable to execute %s: %s", cmd, statusFromError(err))
		}
		os.Stdout.Write(update.GetOutput())
	}
}
<commit_msg>cmd/gomote: add support for groups to run<commit_after>// Copyright 2015 The Go 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

import (
	"context"
	"flag"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"

	"golang.org/x/build/buildlet"
	"golang.org/x/build/dashboard"
	"golang.org/x/build/internal/envutil"
	"golang.org/x/build/internal/gomote/protos"
	"golang.org/x/sync/errgroup"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
)

func legacyRun(args []string) error {
	if activeGroup != nil {
		return fmt.Errorf("command does not support groups")
	}

	fs := flag.NewFlagSet("run", flag.ContinueOnError)
	fs.Usage = func() {
		fmt.Fprintln(os.Stderr, "run usage: gomote run [run-opts] <instance> <cmd> [args...]")
		fs.PrintDefaults()
		os.Exit(1)
	}
	var sys bool
	fs.BoolVar(&sys, "system", false, "run inside the system, and not inside the workdir; this is implicit if cmd starts with '/'")
	var debug bool
	fs.BoolVar(&debug, "debug", false, "write debug info about the command's execution before it begins")
	var env stringSlice
	fs.Var(&env, "e", "Environment variable KEY=value. The -e flag may be repeated multiple times to add multiple things to the environment.")
	var firewall bool
	fs.BoolVar(&firewall, "firewall", false, "Enable outbound firewall on machine. This is on by default on many builders (where supported) but disabled by default on gomote for ease of debugging. Once any command has been run with the -firewall flag on, it's on for the lifetime of that gomote instance.")
	var path string
	fs.StringVar(&path, "path", "", "Comma-separated list of ExecOpts.Path elements. The special string 'EMPTY' means to run without any $PATH. The empty string (default) does not modify the $PATH. Otherwise, the following expansions apply: the string '$PATH' expands to the current PATH element(s), the substring '$WORKDIR' expands to the buildlet's temp workdir.")

	var dir string
	fs.StringVar(&dir, "dir", "", "Directory to run from. Defaults to the directory of the command, or the work directory if -system is true.")
	var builderEnv string
	fs.StringVar(&builderEnv, "builderenv", "", "Optional alternate builder to act like. Must share the same underlying buildlet host type, or it's an error. For instance, linux-amd64-race or linux-386-387 are compatible with linux-amd64, but openbsd-amd64 and openbsd-386 are different hosts.")

	fs.Parse(args)
	if fs.NArg() < 2 {
		fs.Usage()
	}
	name, cmd := fs.Arg(0), fs.Arg(1)

	var conf *dashboard.BuildConfig

	bc, conf, err := clientAndConf(name)
	if err != nil {
		return err
	}

	if builderEnv != "" {
		altConf, ok := dashboard.Builders[builderEnv]
		if !ok {
			return fmt.Errorf("unknown --builderenv=%q builder value", builderEnv)
		}
		if altConf.HostType != conf.HostType {
			return fmt.Errorf("--builderEnv=%q has host type %q, which is not compatible with the named buildlet's host type %q",
				builderEnv, altConf.HostType, conf.HostType)
		}
		conf = altConf
	}

	var pathOpt []string
	if path == "EMPTY" {
		pathOpt = []string{} // non-nil
	} else if path != "" {
		pathOpt = strings.Split(path, ",")
	}
	env = append(env, "GO_DISABLE_OUTBOUND_NETWORK="+fmt.Sprint(firewall))

	remoteErr, execErr := bc.Exec(context.Background(), cmd, buildlet.ExecOpts{
		Dir:         dir,
		SystemLevel: sys || strings.HasPrefix(cmd, "/"),
		Output:      os.Stdout,
		Args:        fs.Args()[2:],
		ExtraEnv:    envutil.Dedup(conf.GOOS(), append(conf.Env(), []string(env)...)),
		Debug:       debug,
		Path:        pathOpt,
	})
	if execErr != nil {
		return fmt.Errorf("Error trying to execute %s: %v", cmd, execErr)
	}
	return remoteErr
}

// stringSlice implements flag.Value, specifically for storing environment
// variable key=value pairs.
type stringSlice []string

func (*stringSlice) String() string { return "" } // default value

func (ss *stringSlice) Set(v string) error {
	if v != "" {
		if !strings.Contains(v, "=") {
			return fmt.Errorf("-e argument %q doesn't contains an '=' sign.", v)
		}
		*ss = append(*ss, v)
	}
	return nil
}

func run(args []string) error {
	fs := flag.NewFlagSet("run", flag.ContinueOnError)
	fs.Usage = func() {
		fmt.Fprintln(os.Stderr, "run usage: gomote run [run-opts] <instance> <cmd> [args...]")
		fs.PrintDefaults()
		os.Exit(1)
	}
	var sys bool
	fs.BoolVar(&sys, "system", false, "run inside the system, and not inside the workdir; this is implicit if cmd starts with '/'")
	var debug bool
	fs.BoolVar(&debug, "debug", false, "write debug info about the command's execution before it begins")
	var env stringSlice
	fs.Var(&env, "e", "Environment variable KEY=value. The -e flag may be repeated multiple times to add multiple things to the environment.")
	var firewall bool
	fs.BoolVar(&firewall, "firewall", false, "Enable outbound firewall on machine. This is on by default on many builders (where supported) but disabled by default on gomote for ease of debugging. Once any command has been run with the -firewall flag on, it's on for the lifetime of that gomote instance.")
	var path string
	fs.StringVar(&path, "path", "", "Comma-separated list of ExecOpts.Path elements. The special string 'EMPTY' means to run without any $PATH. The empty string (default) does not modify the $PATH. Otherwise, the following expansions apply: the string '$PATH' expands to the current PATH element(s), the substring '$WORKDIR' expands to the buildlet's temp workdir.")

	var dir string
	fs.StringVar(&dir, "dir", "", "Directory to run from. Defaults to the directory of the command, or the work directory if -system is true.")
	var builderEnv string
	fs.StringVar(&builderEnv, "builderenv", "", "Optional alternate builder to act like. Must share the same underlying buildlet host type, or it's an error. For instance, linux-amd64-race or linux-386-387 are compatible with linux-amd64, but openbsd-amd64 and openbsd-386 are different hosts.")

	fs.Parse(args)
	if fs.NArg() == 0 {
		fs.Usage()
	}
	// First check if the instance name refers to a live instance.
	ctx := context.Background()
	client := gomoteServerClient(ctx)
	_, err := client.InstanceAlive(ctx, &protos.InstanceAliveRequest{
		GomoteId: fs.Arg(0),
	})
	var cmd string
	var cmdArgs []string
	var runSet []string
	if err != nil {
		// When there's no active group, this must be an instance name.
		// Given that we got an error, we should surface that.
		if activeGroup == nil {
			return fmt.Errorf("instance %q: %s", fs.Arg(0), statusFromError(err))
		}
		// When there is an active group, this just means that we're going
		// to use the group instead and assume the rest is a command.
		for _, inst := range activeGroup.Instances {
			runSet = append(runSet, inst)
		}
		cmd = fs.Arg(0)
		cmdArgs = fs.Args()[1:]
	} else {
		runSet = append(runSet, fs.Arg(0))
		if fs.NArg() == 1 {
			fmt.Fprintln(os.Stderr, "missing command")
			fs.Usage()
		}
		cmd = fs.Arg(1)
		cmdArgs = fs.Args()[2:]
	}
	var pathOpt []string
	if path == "EMPTY" {
		pathOpt = []string{} // non-nil
	} else if path != "" {
		pathOpt = strings.Split(path, ",")
	}
	env = append(env, "GO_DISABLE_OUTBOUND_NETWORK="+fmt.Sprint(firewall))

	detailedProgress := len(runSet) == 1

	// Create temporary directory for output.
	// This is useful even if we don't have multiple gomotes running, since
	// it's easy to accidentally lose the output.
	tmpOutDir, err := os.MkdirTemp("", "gomote")
	if err != nil {
		return err
	}

	eg, ctx := errgroup.WithContext(context.Background())
	for _, inst := range runSet {
		inst := inst
		if !detailedProgress {
			fmt.Fprintf(os.Stderr, "# Running command on %q...\n", inst)
		}
		eg.Go(func() error {
			outf, err := os.Create(filepath.Join(tmpOutDir, fmt.Sprintf("%s.stdout", inst)))
			if err != nil {
				return err
			}
			defer func() {
				outf.Close()
				fmt.Fprintf(os.Stderr, "# Wrote results from %q to %q.\n", inst, outf.Name())
			}()
			fmt.Fprintf(os.Stderr, "# Streaming results from %q to %q...\n", inst, outf.Name())
			var outWriter io.Writer
			if detailedProgress {
				outWriter = io.MultiWriter(os.Stdout, outf)
			} else {
				outWriter = outf
			}

			client := gomoteServerClient(ctx)
			stream, err := client.ExecuteCommand(ctx, &protos.ExecuteCommandRequest{
				AppendEnvironment: []string(env),
				Args:              cmdArgs,
				Command:           cmd,
				Debug:             debug,
				Directory:         dir,
				GomoteId:          inst,
				Path:              pathOpt,
				SystemLevel:       sys || strings.HasPrefix(cmd, "/"),
				ImitateHostType:   builderEnv,
			})
			if err != nil {
				return fmt.Errorf("unable to execute %s: %s", cmd, statusFromError(err))
			}
			for {
				update, err := stream.Recv()
				if err == io.EOF {
					return nil
				}
				if err != nil {
					// execution error
					if status.Code(err) == codes.Aborted {
						return fmt.Errorf("Error trying to execute %s: %v", cmd, statusFromError(err))
					}
					// remote error
					return fmt.Errorf("unable to execute %s: %s", cmd, statusFromError(err))
				}
				fmt.Fprintf(outWriter, string(update.GetOutput()))
			}
		})
	}
	return eg.Wait()
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"
	"strings"

	"github.com/kubernetes/helm/pkg/chart"
	"github.com/kubernetes/helm/pkg/repo"
	"github.com/spf13/cobra"
)

var untarFile bool

// TODO(vaikas): Is there a well known directory for these??
var untarDir string

func init() {
	RootCommand.AddCommand(fetchCmd)
	fetchCmd.Flags().BoolVar(&untarFile, "untar", false, "If set to true, will untar the chart after downloading it.")
	fetchCmd.Flags().StringVar(&untarDir, "untardir", ".", "If untar is specified, this flag specifies where to untar the chart.")
}

var fetchCmd = &cobra.Command{
	Use:   "fetch [chart URL | repo/chartname]",
	Short: "Download a chart from a repository and (optionally) unpack it in local directory.",
	Long:  "",
	RunE:  fetch,
}

func fetch(cmd *cobra.Command, args []string) error {
	if len(args) == 0 {
		return fmt.Errorf("This command needs at least one argument, url or repo/name of the chart.")
	}

	f, err := repo.LoadRepositoriesFile(repositoriesFile())
	if err != nil {
		return err
	}

	// get download url
	u, err := mapRepoArg(args[0], f.Repositories)
	if err != nil {
		return err
	}

	resp, err := http.Get(u.String())
	if err != nil {
		return err
	}
	if resp.StatusCode != 200 {
		return fmt.Errorf("Failed to fetch %s : %s", u.String(), resp.Status)
	}

	defer resp.Body.Close()
	if untarFile {
		return chart.Expand(untarDir, resp.Body)
	}
	p := strings.Split(u.String(), "/")
	return saveChartFile(p[len(p)-1], resp.Body)
}

// mapRepoArg figures out which format the argument is given, and creates a fetchable
// url from it.
func mapRepoArg(arg string, r map[string]string) (*url.URL, error) {
	// See if it's already a full URL.
	u, err := url.ParseRequestURI(arg)
	if err == nil {
		// If it has a scheme and host and path, it's a full URL
		if u.IsAbs() && len(u.Host) > 0 && len(u.Path) > 0 {
			return u, nil
		}
		return nil, fmt.Errorf("Invalid chart url format: %s", arg)
	}
	// See if it's of the form: repo/path_to_chart
	p := strings.Split(arg, "/")
	if len(p) > 1 {
		if baseURL, ok := r[p[0]]; ok {
			if !strings.HasSuffix(baseURL, "/") {
				baseURL = baseURL + "/"
			}
			return url.ParseRequestURI(baseURL + strings.Join(p[1:], "/"))
		}
		return nil, fmt.Errorf("No such repo: %s", p[0])
	}
	return nil, fmt.Errorf("Invalid chart url format: %s", arg)
}

func saveChartFile(c string, r io.Reader) error {
	// Grab the chart name that we'll use for the name of the file to download to.
	out, err := os.Create(c)
	if err != nil {
		return err
	}
	defer out.Close()

	_, err = io.Copy(out, r)
	return err
}
<commit_msg>remove the todo based on PR comments<commit_after>package main

import (
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"
	"strings"

	"github.com/kubernetes/helm/pkg/chart"
	"github.com/kubernetes/helm/pkg/repo"
	"github.com/spf13/cobra"
)

var untarFile bool
var untarDir string

func init() {
	RootCommand.AddCommand(fetchCmd)
	fetchCmd.Flags().BoolVar(&untarFile, "untar", false, "If set to true, will untar the chart after downloading it.")
	fetchCmd.Flags().StringVar(&untarDir, "untardir", ".", "If untar is specified, this flag specifies where to untar the chart.")
}

var fetchCmd = &cobra.Command{
	Use:   "fetch [chart URL | repo/chartname]",
	Short: "Download a chart from a repository and (optionally) unpack it in local directory.",
	Long:  "",
	RunE:  fetch,
}

func fetch(cmd *cobra.Command, args []string) error {
	if len(args) == 0 {
		return fmt.Errorf("This command needs at least one argument, url or repo/name of the chart.")
	}

	f, err := repo.LoadRepositoriesFile(repositoriesFile())
	if err != nil {
		return err
	}

	// get download url
	u, err := mapRepoArg(args[0], f.Repositories)
	if err != nil {
		return err
	}

	resp, err := http.Get(u.String())
	if err != nil {
		return err
	}
	if resp.StatusCode != 200 {
		return fmt.Errorf("Failed to fetch %s : %s", u.String(), resp.Status)
	}

	defer resp.Body.Close()
	if untarFile {
		return chart.Expand(untarDir, resp.Body)
	}
	p := strings.Split(u.String(), "/")
	return saveChartFile(p[len(p)-1], resp.Body)
}

// mapRepoArg figures out which format the argument is given, and creates a fetchable
// url from it.
func mapRepoArg(arg string, r map[string]string) (*url.URL, error) {
	// See if it's already a full URL.
	u, err := url.ParseRequestURI(arg)
	if err == nil {
		// If it has a scheme and host and path, it's a full URL
		if u.IsAbs() && len(u.Host) > 0 && len(u.Path) > 0 {
			return u, nil
		}
		return nil, fmt.Errorf("Invalid chart url format: %s", arg)
	}
	// See if it's of the form: repo/path_to_chart
	p := strings.Split(arg, "/")
	if len(p) > 1 {
		if baseURL, ok := r[p[0]]; ok {
			if !strings.HasSuffix(baseURL, "/") {
				baseURL = baseURL + "/"
			}
			return url.ParseRequestURI(baseURL + strings.Join(p[1:], "/"))
		}
		return nil, fmt.Errorf("No such repo: %s", p[0])
	}
	return nil, fmt.Errorf("Invalid chart url format: %s", arg)
}

func saveChartFile(c string, r io.Reader) error {
	// Grab the chart name that we'll use for the name of the file to download to.
	out, err := os.Create(c)
	if err != nil {
		return err
	}
	defer out.Close()

	_, err = io.Copy(out, r)
	return err
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"
	"strings"

	"github.com/spf13/cobra"
	"k8s.io/helm/pkg/chartutil"
	"k8s.io/helm/pkg/repo"
)

var untarFile bool
var untarDir string

func init() {
	RootCommand.AddCommand(fetchCmd)
	fetchCmd.Flags().BoolVar(&untarFile, "untar", false, "If set to true, will untar the chart after downloading it.")
	fetchCmd.Flags().StringVar(&untarDir, "untardir", ".", "If untar is specified, this flag specifies where to untar the chart.")
}

var fetchCmd = &cobra.Command{
	Use:   "fetch [chart URL | repo/chartname]",
	Short: "Download a chart from a repository and (optionally) unpack it in local directory.",
	Long:  "",
	RunE:  fetch,
}

func fetch(cmd *cobra.Command, args []string) error {
	if len(args) == 0 {
		return fmt.Errorf("This command needs at least one argument, url or repo/name of the chart.")
	}

	f, err := repo.LoadRepositoriesFile(repositoriesFile())
	if err != nil {
		return err
	}

	// get download url
	u, err := mapRepoArg(args[0], f.Repositories)
	if err != nil {
		return err
	}

	resp, err := http.Get(u.String())
	if err != nil {
		return err
	}
	if resp.StatusCode != 200 {
		return fmt.Errorf("Failed to fetch %s : %s", u.String(), resp.Status)
	}

	defer resp.Body.Close()
	if untarFile {
		return chartutil.Expand(untarDir, resp.Body)
	}
	p := strings.Split(u.String(), "/")
	return saveChartFile(p[len(p)-1], resp.Body)
}

// mapRepoArg figures out which format the argument is given, and creates a fetchable
// url from it.
func mapRepoArg(arg string, r map[string]string) (*url.URL, error) {
	// See if it's already a full URL.
	u, err := url.ParseRequestURI(arg)
	if err == nil {
		// If it has a scheme and host and path, it's a full URL
		if u.IsAbs() && len(u.Host) > 0 && len(u.Path) > 0 {
			return u, nil
		}
		return nil, fmt.Errorf("Invalid chart url format: %s", arg)
	}
	// See if it's of the form: repo/path_to_chart
	p := strings.Split(arg, "/")
	if len(p) > 1 {
		if baseURL, ok := r[p[0]]; ok {
			if !strings.HasSuffix(baseURL, "/") {
				baseURL = baseURL + "/"
			}
			return url.ParseRequestURI(baseURL + strings.Join(p[1:], "/"))
		}
		return nil, fmt.Errorf("No such repo: %s", p[0])
	}
	return nil, fmt.Errorf("Invalid chart url format: %s", arg)
}

func saveChartFile(c string, r io.Reader) error {
	// Grab the chart name that we'll use for the name of the file to download to.
	out, err := os.Create(c)
	if err != nil {
		return err
	}
	defer out.Close()

	_, err = io.Copy(out, r)
	return err
}
<commit_msg>fix(helm): add .tgz to package names on fetch<commit_after>package main

import (
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"strings"

	"github.com/spf13/cobra"
	"k8s.io/helm/pkg/chartutil"
	"k8s.io/helm/pkg/repo"
)

var untarFile bool
var untarDir string

func init() {
	RootCommand.AddCommand(fetchCmd)
	fetchCmd.Flags().BoolVar(&untarFile, "untar", false, "If set to true, will untar the chart after downloading it.")
	fetchCmd.Flags().StringVar(&untarDir, "untardir", ".", "If untar is specified, this flag specifies where to untar the chart.")
}

var fetchCmd = &cobra.Command{
	Use:   "fetch [chart URL | repo/chartname]",
	Short: "Download a chart from a repository and (optionally) unpack it in local directory.",
	Long:  "",
	RunE:  fetch,
}

func fetch(cmd *cobra.Command, args []string) error {
	if len(args) == 0 {
		return fmt.Errorf("This command needs at least one argument, url or repo/name of the chart.")
	}

	pname := args[0]
	if filepath.Ext(pname) != ".tgz" {
		pname += ".tgz"
	}

	f, err := repo.LoadRepositoriesFile(repositoriesFile())
	if err != nil {
		return err
	}

	// get download url
	u, err := mapRepoArg(pname, f.Repositories)
	if err != nil {
		return err
	}

	resp, err := http.Get(u.String())
	if err != nil {
		return err
	}
	if resp.StatusCode != 200 {
		return fmt.Errorf("Failed to fetch %s : %s", u.String(), resp.Status)
	}

	defer resp.Body.Close()
	if untarFile {
		return chartutil.Expand(untarDir, resp.Body)
	}
	p := strings.Split(u.String(), "/")
	return saveChartFile(p[len(p)-1], resp.Body)
}

// mapRepoArg figures out which format the argument is given, and creates a fetchable
// url from it.
func mapRepoArg(arg string, r map[string]string) (*url.URL, error) {
	// See if it's already a full URL.
	u, err := url.ParseRequestURI(arg)
	if err == nil {
		// If it has a scheme and host and path, it's a full URL
		if u.IsAbs() && len(u.Host) > 0 && len(u.Path) > 0 {
			return u, nil
		}
		return nil, fmt.Errorf("Invalid chart url format: %s", arg)
	}
	// See if it's of the form: repo/path_to_chart
	p := strings.Split(arg, "/")
	if len(p) > 1 {
		if baseURL, ok := r[p[0]]; ok {
			if !strings.HasSuffix(baseURL, "/") {
				baseURL = baseURL + "/"
			}
			return url.ParseRequestURI(baseURL + strings.Join(p[1:], "/"))
		}
		return nil, fmt.Errorf("No such repo: %s", p[0])
	}
	return nil, fmt.Errorf("Invalid chart url format: %s", arg)
}

func saveChartFile(c string, r io.Reader) error {
	// Grab the chart name that we'll use for the name of the file to download to.
	out, err := os.Create(c)
	if err != nil {
		return err
	}
	defer out.Close()

	_, err = io.Copy(out, r)
	return err
}
<|endoftext|>


<commit_before>package command

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"os"
	"os/signal"
	"strings"
	"syscall"

	"github.com/docker/docker/pkg/term"
	"github.com/hashicorp/nomad/api"
	"github.com/hashicorp/nomad/api/contexts"
	"github.com/hashicorp/nomad/helper/escapingio"
	"github.com/posener/complete"
)

type AllocExecCommand struct {
	Meta

	Stdin  io.Reader
	Stdout io.WriteCloser
	Stderr io.WriteCloser
}

func (l *AllocExecCommand) Help() string {
	helpText := `
Usage: nomad alloc exec [options] <allocation> <command>

  Run command inside the environment of the given allocation and task.

General Options:

  ` + generalOptionsUsage() + `

Exec Specific Options:

  -task <task-name>
    Sets the task to exec command in

  -job
    Use a random allocation from the specified job ID.

  -i
    Pass stdin to the container, defaults to true.  Pass -i=false to disable.

  -t
    Allocate a pseudo-tty, defaults to true if stdin is detected to be a tty session.
    Pass -t=false to disable explicitly.

  -e <escape_char>
    Sets the escape character for sessions with a pty (default: '~').  The escape
    character is only recognized at the beginning of a line.  The escape character
    followed by a dot ('.') closes the connection.  Setting the character to
    'none' disables any escapes and makes the session fully transparent.
  `
	return strings.TrimSpace(helpText)
}

func (l *AllocExecCommand) Synopsis() string {
	return "Execute commands in task"
}

func (c *AllocExecCommand) AutocompleteFlags() complete.Flags {
	return mergeAutocompleteFlags(c.Meta.AutocompleteFlags(FlagSetClient),
		complete.Flags{
			"--task": complete.PredictAnything,
			"-job":   complete.PredictAnything,
			"-i":     complete.PredictNothing,
			"-t":     complete.PredictNothing,
			"-e":     complete.PredictSet("none", "~"),
		})
}

func (l *AllocExecCommand) AutocompleteArgs() complete.Predictor {
	return complete.PredictFunc(func(a complete.Args) []string {
		client, err := l.Meta.Client()
		if err != nil {
			return nil
		}

		resp, _, err := client.Search().PrefixSearch(a.Last, contexts.Allocs, nil)
		if err != nil {
			return []string{}
		}
		return resp.Matches[contexts.Allocs]
	})
}

func (l *AllocExecCommand) Name() string { return "alloc exec" }

func (l *AllocExecCommand) Run(args []string) int {
	var job, stdinOpt, ttyOpt bool
	var task string
	var escapeChar string

	flags := l.Meta.FlagSet(l.Name(), FlagSetClient)
	flags.Usage = func() { l.Ui.Output(l.Help()) }
	flags.BoolVar(&job, "job", false, "")
	flags.StringVar(&task, "task", "", "")
	flags.StringVar(&escapeChar, "e", "~", "")

	flags.BoolVar(&stdinOpt, "i", true, "")

	stdinTty := isStdinTty()
	flags.BoolVar(&ttyOpt, "t", stdinTty, "")

	if err := flags.Parse(args); err != nil {
		return 1
	}
	args = flags.Args()

	if ttyOpt && !stdinOpt {
		l.Ui.Error("-i must be enabled if running with tty")
		return 1
	}

	if escapeChar == "none" {
		escapeChar = ""
	} else if len(escapeChar) > 1 {
		l.Ui.Error("-e requires 'none' or a single character")
		return 1
	}

	if numArgs := len(args); numArgs < 1 {
		if job {
			l.Ui.Error("A job ID is required")
		} else {
			l.Ui.Error("An allocation ID is required")
		}

		l.Ui.Error(commandErrorText(l))
		return 1
	} else if numArgs < 2 {
		l.Ui.Error("A command is required")
		l.Ui.Error(commandErrorText(l))
		return 1
	}

	command := args[1:]

	client, err := l.Meta.Client()
	if err != nil {
		l.Ui.Error(fmt.Sprintf("Error initializing client: %v", err))
		return 1
	}

	// If -job is specified, use random allocation, otherwise use provided allocation
	allocID := args[0]
	if job {
		allocID, err = getRandomJobAlloc(client, args[0])
		if err != nil {
			l.Ui.Error(fmt.Sprintf("Error fetching allocations: %v", err))
			return 1
		}
	}

	length := shortId

	// Query the allocation info
	if len(allocID) == 1 {
		l.Ui.Error(fmt.Sprintf("Alloc ID must contain at least two characters."))
		return 1
	}

	allocID = sanitizeUUIDPrefix(allocID)
	allocs, _, err := client.Allocations().PrefixList(allocID)
	if err != nil {
		l.Ui.Error(fmt.Sprintf("Error querying allocation: %v", err))
		return 1
	}
	if len(allocs) == 0 {
		l.Ui.Error(fmt.Sprintf("No allocation(s) with prefix or id %q found", allocID))
		return 1
	}
	if len(allocs) > 1 {
		// Format the allocs
		out := formatAllocListStubs(allocs, false, length)
		l.Ui.Error(fmt.Sprintf("Prefix matched multiple allocations\n\n%s", out))
		return 1
	}
	// Prefix lookup matched a single allocation
	alloc, _, err := client.Allocations().Info(allocs[0].ID, nil)
	if err != nil {
		l.Ui.Error(fmt.Sprintf("Error querying allocation: %s", err))
		return 1
	}

	if task == "" {
		task, err = lookupAllocTask(alloc)

		if err != nil {
			l.Ui.Error(err.Error())
			l.Ui.Error("\nPlease specify the task.")
			return 1
		}
	}

	if err := validateTaskExistsInAllocation(task, alloc); err != nil {
		l.Ui.Error(err.Error())
		return 1
	}

	if l.Stdin == nil {
		l.Stdin = os.Stdin
	}
	if l.Stdout == nil {
		l.Stdout = os.Stdout
	}
	if l.Stderr == nil {
		l.Stderr = os.Stderr
	}

	var stdin io.Reader = l.Stdin
	if !stdinOpt {
		stdin = bytes.NewReader(nil)
	}

	code, err := l.execImpl(client, alloc, task, ttyOpt, command, escapeChar, stdin, l.Stdout, l.Stderr)
	if err != nil {
		l.Ui.Error(fmt.Sprintf("failed to exec into task: %v", err))
		return 1
	}

	return code
}

// execImpl invokes the Alloc Exec api call, it also prepares and restores terminal states as necessary.
func (l *AllocExecCommand) execImpl(client *api.Client, alloc *api.Allocation, task string, tty bool,
	command []string, escapeChar string, stdin io.Reader, stdout, stderr io.WriteCloser) (int, error) {

	sizeCh := make(chan api.TerminalSize, 1)

	ctx, cancelFn := context.WithCancel(context.Background())
	defer cancelFn()

	// When tty, ensures we capture all user input and monitor terminal resizes.
	if tty {
		if stdin == nil {
			return -1, fmt.Errorf("stdin is null")
		}

		inCleanup, err := setRawTerminal(stdin)
		if err != nil {
			return -1, err
		}
		defer inCleanup()

		outCleanup, err := setRawTerminalOutput(stdout)
		if err != nil {
			return -1, err
		}
		defer outCleanup()

		sizeCleanup, err := watchTerminalSize(stdout, sizeCh)
		if err != nil {
			return -1, err
		}
		defer sizeCleanup()

		if escapeChar != "" {
			stdin = escapingio.NewReader(stdin, escapeChar[0], func(c byte) bool {
				switch c {
				case '.':
					stderr.Write([]byte("\nConnection closed\n"))
					cancelFn()
					return true
				default:
					return false
				}
			})

		}
	}

	signalCh := make(chan os.Signal, 1)
	signal.Notify(signalCh, os.Interrupt, syscall.SIGTERM)
	go func() {
		for range signalCh {
			cancelFn()
		}
	}()

	return client.Allocations().Exec(ctx,
		alloc, task, tty, command, stdin, stdout, stderr, sizeCh, nil)
}

func isStdinTty() bool {
	_, isTerminal := term.GetFdInfo(os.Stdin)
	return isTerminal
}

// setRawTerminal sets the stream terminal in raw mode, so process captures
// Ctrl+C and other commands to forward to remote process.
// It returns a cleanup function that restores terminal to original mode.
func setRawTerminal(stream interface{}) (cleanup func(), err error) {
	fd, isTerminal := term.GetFdInfo(stream)
	if !isTerminal {
		return nil, errors.New("not a terminal")
	}

	state, err := term.SetRawTerminal(fd)
	if err != nil {
		return nil, err
	}

	return func() { term.RestoreTerminal(fd, state) }, nil
}

// setRawTerminalOutput sets the output stream in Windows to raw mode,
// so it disables LF -> CRLF translation.
// It's basically a no-op on unix.
func setRawTerminalOutput(stream interface{}) (cleanup func(), err error) {
	fd, isTerminal := term.GetFdInfo(stream)
	if !isTerminal {
		return nil, errors.New("not a terminal")
	}

	state, err := term.SetRawTerminalOutput(fd)
	if err != nil {
		return nil, err
	}

	return func() { term.RestoreTerminal(fd, state) }, nil
}

// watchTerminalSize watches terminal size changes to propagate to remote tty.
func watchTerminalSize(out io.Writer, resize chan<- api.TerminalSize) (func(), error) {
	fd, isTerminal := term.GetFdInfo(out)
	if !isTerminal {
		return nil, errors.New("not a terminal")
	}

	ctx, cancel := context.WithCancel(context.Background())

	signalCh := make(chan os.Signal, 1)
	setupWindowNotification(signalCh)

	sendTerminalSize := func() {
		s, err := term.GetWinsize(fd)
		if err != nil {
			return
		}

		resize <- api.TerminalSize{
			Height: int(s.Height),
			Width:  int(s.Width),
		}
	}
	go func() {
		for {
			select {
			case <-ctx.Done():
				return
			case <-signalCh:
				sendTerminalSize()
			}
		}
	}()

	go func() {
		// send initial size
		sendTerminalSize()
	}()

	return cancel, nil
}
<commit_msg>Restore tty start before emitting errors<commit_after>package command

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"os"
	"os/signal"
	"strings"
	"syscall"

	"github.com/docker/docker/pkg/term"
	"github.com/hashicorp/nomad/api"
	"github.com/hashicorp/nomad/api/contexts"
	"github.com/hashicorp/nomad/helper/escapingio"
	"github.com/posener/complete"
)

type AllocExecCommand struct {
	Meta

	Stdin  io.Reader
	Stdout io.WriteCloser
	Stderr io.WriteCloser
}

func (l *AllocExecCommand) Help() string {
	helpText := `
Usage: nomad alloc exec [options] <allocation> <command>

  Run command inside the environment of the given allocation and task.

General Options:

  ` + generalOptionsUsage() + `

Exec Specific Options:

  -task <task-name>
    Sets the task to exec command in

  -job
    Use a random allocation from the specified job ID.

  -i
    Pass stdin to the container, defaults to true.  Pass -i=false to disable.

  -t
    Allocate a pseudo-tty, defaults to true if stdin is detected to be a tty session.
    Pass -t=false to disable explicitly.

  -e <escape_char>
    Sets the escape character for sessions with a pty (default: '~').  The escape
    character is only recognized at the beginning of a line.  The escape character
    followed by a dot ('.') closes the connection.  Setting the character to
    'none' disables any escapes and makes the session fully transparent.
  `
	return strings.TrimSpace(helpText)
}

func (l *AllocExecCommand) Synopsis() string {
	return "Execute commands in task"
}

func (c *AllocExecCommand) AutocompleteFlags() complete.Flags {
	return mergeAutocompleteFlags(c.Meta.AutocompleteFlags(FlagSetClient),
		complete.Flags{
			"--task": complete.PredictAnything,
			"-job":   complete.PredictAnything,
			"-i":     complete.PredictNothing,
			"-t":     complete.PredictNothing,
			"-e":     complete.PredictSet("none", "~"),
		})
}

func (l *AllocExecCommand) AutocompleteArgs() complete.Predictor {
	return complete.PredictFunc(func(a complete.Args) []string {
		client, err := l.Meta.Client()
		if err != nil {
			return nil
		}

		resp, _, err := client.Search().PrefixSearch(a.Last, contexts.Allocs, nil)
		if err != nil {
			return []string{}
		}
		return resp.Matches[contexts.Allocs]
	})
}

func (l *AllocExecCommand) Name() string { return "alloc exec" }

func (l *AllocExecCommand) Run(args []string) int {
	var job, stdinOpt, ttyOpt bool
	var task string
	var escapeChar string

	flags := l.Meta.FlagSet(l.Name(), FlagSetClient)
	flags.Usage = func() { l.Ui.Output(l.Help()) }
	flags.BoolVar(&job, "job", false, "")
	flags.StringVar(&task, "task", "", "")
	flags.StringVar(&escapeChar, "e", "~", "")

	flags.BoolVar(&stdinOpt, "i", true, "")

	stdinTty := isStdinTty()
	flags.BoolVar(&ttyOpt, "t", stdinTty, "")

	if err := flags.Parse(args); err != nil {
		return 1
	}
	args = flags.Args()

	if ttyOpt && !stdinOpt {
		l.Ui.Error("-i must be enabled if running with tty")
		return 1
	}

	if escapeChar == "none" {
		escapeChar = ""
	} else if len(escapeChar) > 1 {
		l.Ui.Error("-e requires 'none' or a single character")
		return 1
	}

	if numArgs := len(args); numArgs < 1 {
		if job {
			l.Ui.Error("A job ID is required")
		} else {
			l.Ui.Error("An allocation ID is required")
		}

		l.Ui.Error(commandErrorText(l))
		return 1
	} else if numArgs < 2 {
		l.Ui.Error("A command is required")
		l.Ui.Error(commandErrorText(l))
		return 1
	}

	command := args[1:]

	client, err := l.Meta.Client()
	if err != nil {
		l.Ui.Error(fmt.Sprintf("Error initializing client: %v", err))
		return 1
	}

	// If -job is specified, use random allocation, otherwise use provided allocation
	allocID := args[0]
	if job {
		allocID, err = getRandomJobAlloc(client, args[0])
		if err != nil {
			l.Ui.Error(fmt.Sprintf("Error fetching allocations: %v", err))
			return 1
		}
	}

	length := shortId

	// Query the allocation info
	if len(allocID) == 1 {
		l.Ui.Error(fmt.Sprintf("Alloc ID must contain at least two characters."))
		return 1
	}

	allocID = sanitizeUUIDPrefix(allocID)
	allocs, _, err := client.Allocations().PrefixList(allocID)
	if err != nil {
		l.Ui.Error(fmt.Sprintf("Error querying allocation: %v", err))
		return 1
	}
	if len(allocs) == 0 {
		l.Ui.Error(fmt.Sprintf("No allocation(s) with prefix or id %q found", allocID))
		return 1
	}
	if len(allocs) > 1 {
		// Format the allocs
		out := formatAllocListStubs(allocs, false, length)
		l.Ui.Error(fmt.Sprintf("Prefix matched multiple allocations\n\n%s", out))
		return 1
	}
	// Prefix lookup matched a single allocation
	alloc, _, err := client.Allocations().Info(allocs[0].ID, nil)
	if err != nil {
		l.Ui.Error(fmt.Sprintf("Error querying allocation: %s", err))
		return 1
	}

	if task == "" {
		task, err = lookupAllocTask(alloc)

		if err != nil {
			l.Ui.Error(err.Error())
			l.Ui.Error("\nPlease specify the task.")
			return 1
		}
	}

	if err := validateTaskExistsInAllocation(task, alloc); err != nil {
		l.Ui.Error(err.Error())
		return 1
	}

	if l.Stdin == nil {
		l.Stdin = os.Stdin
	}
	if l.Stdout == nil {
		l.Stdout = os.Stdout
	}
	if l.Stderr == nil {
		l.Stderr = os.Stderr
	}

	var stdin io.Reader = l.Stdin
	if !stdinOpt {
		stdin = bytes.NewReader(nil)
	}

	code, err := l.execImpl(client, alloc, task, ttyOpt, command, escapeChar, stdin, l.Stdout, l.Stderr)
	if err != nil {
		l.Ui.Error(fmt.Sprintf("failed to exec into task: %v", err))
		return 1
	}

	return code
}

// execImpl invokes the Alloc Exec api call, it also prepares and restores terminal states as necessary.
func (l *AllocExecCommand) execImpl(client *api.Client, alloc *api.Allocation, task string, tty bool,
	command []string, escapeChar string, stdin io.Reader, stdout, stderr io.WriteCloser) (int, error) {

	sizeCh := make(chan api.TerminalSize, 1)

	ctx, cancelFn := context.WithCancel(context.Background())
	defer cancelFn()

	// When tty, ensures we capture all user input and monitor terminal resizes.
	if tty {
		if stdin == nil {
			return -1, fmt.Errorf("stdin is null")
		}

		inCleanup, err := setRawTerminal(stdin)
		if err != nil {
			return -1, err
		}
		defer inCleanup()

		outCleanup, err := setRawTerminalOutput(stdout)
		if err != nil {
			return -1, err
		}
		defer outCleanup()

		sizeCleanup, err := watchTerminalSize(stdout, sizeCh)
		if err != nil {
			return -1, err
		}
		defer sizeCleanup()

		if escapeChar != "" {
			stdin = escapingio.NewReader(stdin, escapeChar[0], func(c byte) bool {
				switch c {
				case '.':
					// need to restore tty state so error reporting here
					// gets emitted at beginning of line
					outCleanup()
					inCleanup()

					stderr.Write([]byte("\nConnection closed\n"))
					cancelFn()
					return true
				default:
					return false
				}
			})
		}
	}

	signalCh := make(chan os.Signal, 1)
	signal.Notify(signalCh, os.Interrupt, syscall.SIGTERM)
	go func() {
		for range signalCh {
			cancelFn()
		}
	}()

	return client.Allocations().Exec(ctx,
		alloc, task, tty, command, stdin, stdout, stderr, sizeCh, nil)
}

func isStdinTty() bool {
	_, isTerminal := term.GetFdInfo(os.Stdin)
	return isTerminal
}

// setRawTerminal sets the stream terminal in raw mode, so process captures
// Ctrl+C and other commands to forward to remote process.
// It returns a cleanup function that restores terminal to original mode.
func setRawTerminal(stream interface{}) (cleanup func(), err error) {
	fd, isTerminal := term.GetFdInfo(stream)
	if !isTerminal {
		return nil, errors.New("not a terminal")
	}

	state, err := term.SetRawTerminal(fd)
	if err != nil {
		return nil, err
	}

	return func() { term.RestoreTerminal(fd, state) }, nil
}

// setRawTerminalOutput sets the output stream in Windows to raw mode,
// so it disables LF -> CRLF translation.
// It's basically a no-op on unix.
func setRawTerminalOutput(stream interface{}) (cleanup func(), err error) {
	fd, isTerminal := term.GetFdInfo(stream)
	if !isTerminal {
		return nil, errors.New("not a terminal")
	}

	state, err := term.SetRawTerminalOutput(fd)
	if err != nil {
		return nil, err
	}

	return func() { term.RestoreTerminal(fd, state) }, nil
}

// watchTerminalSize watches terminal size changes to propagate to remote tty.
func watchTerminalSize(out io.Writer, resize chan<- api.TerminalSize) (func(), error) {
	fd, isTerminal := term.GetFdInfo(out)
	if !isTerminal {
		return nil, errors.New("not a terminal")
	}

	ctx, cancel := context.WithCancel(context.Background())

	signalCh := make(chan os.Signal, 1)
	setupWindowNotification(signalCh)

	sendTerminalSize := func() {
		s, err := term.GetWinsize(fd)
		if err != nil {
			return
		}

		resize <- api.TerminalSize{
			Height: int(s.Height),
			Width:  int(s.Width),
		}
	}
	go func() {
		for {
			select {
			case <-ctx.Done():
				return
			case <-signalCh:
				sendTerminalSize()
			}
		}
	}()

	go func() {
		// send initial size
		sendTerminalSize()
	}()

	return cancel, nil
}
<|endoftext|>


<commit_before>package command

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"os"
	"os/signal"
	"strings"
	"syscall"

	"github.com/docker/docker/pkg/term"
	"github.com/hashicorp/nomad/api"
	"github.com/hashicorp/nomad/api/contexts"
	"github.com/hashicorp/nomad/helper/escapingio"
	"github.com/posener/complete"
)

type AllocExecCommand struct {
	Meta

	Stdin  io.Reader
	Stdout io.WriteCloser
	Stderr io.WriteCloser
}

func (l *AllocExecCommand) Help() string {
	helpText := `
Usage: nomad alloc exec [options] <allocation> <command>

  Run command inside the environment of the given allocation and task.

General Options:

  ` + generalOptionsUsage() + `

Exec Specific Options:

  -task <task-name>
    Sets the task to exec command in

  -job
    Use a random allocation from the specified job ID.

  -i
    Pass stdin to the container, defaults to true.  Pass -i=false to disable.

  -t
    Allocate a pseudo-tty, defaults to true if stdin is detected to be a tty session.
    Pass -t=false to disable explicitly.

  -e <escape_char>
    Sets the escape character for sessions with a pty (default: '~').  The escape
    character is only recognized at the beginning of a line.  The escape character
    followed by a dot ('.') closes the connection.  Setting the character to
    'none' disables any escapes and makes the session fully transparent.
  `
	return strings.TrimSpace(helpText)
}

func (l *AllocExecCommand) Synopsis() string {
	return "Execute commands in task"
}

func (c *AllocExecCommand) AutocompleteFlags() complete.Flags {
	return mergeAutocompleteFlags(c.Meta.AutocompleteFlags(FlagSetClient),
		complete.Flags{
			"--task": complete.PredictAnything,
			"-job":   complete.PredictAnything,
			"-i":     complete.PredictNothing,
			"-t":     complete.PredictNothing,
			"-e":     complete.PredictSet("none", "~"),
		})
}

func (l *AllocExecCommand) AutocompleteArgs() complete.Predictor {
	return complete.PredictFunc(func(a complete.Args) []string {
		client, err := l.Meta.Client()
		if err != nil {
			return nil
		}

		resp, _, err := client.Search().PrefixSearch(a.Last, contexts.Allocs, nil)
		if err != nil {
			return []string{}
		}
		return resp.Matches[contexts.Allocs]
	})
}

func (l *AllocExecCommand) Name() string { return "alloc exec" }

func (l *AllocExecCommand) Run(args []string) int {
	var job, stdinOpt, ttyOpt bool
	var task string
	var escapeChar string

	flags := l.Meta.FlagSet(l.Name(), FlagSetClient)
	flags.Usage = func() { l.Ui.Output(l.Help()) }
	flags.BoolVar(&job, "job", false, "")
	flags.StringVar(&task, "task", "", "")
	flags.StringVar(&escapeChar, "e", "~", "")

	flags.BoolVar(&stdinOpt, "i", true, "")

	stdinTty := isStdinTty()
	flags.BoolVar(&ttyOpt, "t", stdinTty, "")

	if err := flags.Parse(args); err != nil {
		return 1
	}
	args = flags.Args()

	if ttyOpt && !stdinOpt {
		l.Ui.Error("-i must be enabled if running with tty")
		return 1
	}

	if escapeChar == "none" {
		escapeChar = ""
	} else if len(escapeChar) > 1 {
		l.Ui.Error("-e requires 'none' or a single character")
		return 1
	}

	if numArgs := len(args); numArgs < 1 {
		if job {
			l.Ui.Error("A job ID is required")
		} else {
			l.Ui.Error("An allocation ID is required")
		}

		l.Ui.Error(commandErrorText(l))
		return 1
	} else if numArgs < 2 {
		l.Ui.Error("A command is required")
		l.Ui.Error(commandErrorText(l))
		return 1
	}

	command := args[1:]

	client, err := l.Meta.Client()
	if err != nil {
		l.Ui.Error(fmt.Sprintf("Error initializing client: %v", err))
		return 1
	}

	// If -job is specified, use random allocation, otherwise use provided allocation
	allocID := args[0]
	if job {
		allocID, err = getRandomJobAlloc(client, args[0])
		if err != nil {
			l.Ui.Error(fmt.Sprintf("Error fetching allocations: %v", err))
			return 1
		}
	}

	length := shortId

	// Query the allocation info
	if len(allocID) == 1 {
		l.Ui.Error(fmt.Sprintf("Alloc ID must contain at least two characters."))
		return 1
	}

	allocID = sanitizeUUIDPrefix(allocID)
	allocs, _, err := client.Allocations().PrefixList(allocID)
	if err != nil {
		l.Ui.Error(fmt.Sprintf("Error querying allocation: %v", err))
		return 1
	}
	if len(allocs) == 0 {
		l.Ui.Error(fmt.Sprintf("No allocation(s) with prefix or id %q found", allocID))
		return 1
	}
	if len(allocs) > 1 {
		// Format the allocs
		out := formatAllocListStubs(allocs, false, length)
		l.Ui.Error(fmt.Sprintf("Prefix matched multiple allocations\n\n%s", out))
		return 1
	}
	// Prefix lookup matched a single allocation
	alloc, _, err := client.Allocations().Info(allocs[0].ID, nil)
	if err != nil {
		l.Ui.Error(fmt.Sprintf("Error querying allocation: %s", err))
		return 1
	}

	if task == "" {
		task, err = lookupAllocTask(alloc)

		if err != nil {
			l.Ui.Error(err.Error())
			l.Ui.Error("\nPlease specify the task.")
			return 1
		}
	}

	if err := validateTaskExistsInAllocation(task, alloc); err != nil {
		l.Ui.Error(err.Error())
		return 1
	}

	if l.Stdin == nil {
		l.Stdin = os.Stdin
	}
	if l.Stdout == nil {
		l.Stdout = os.Stdout
	}
	if l.Stderr == nil {
		l.Stderr = os.Stderr
	}

	var stdin io.Reader = l.Stdin
	if !stdinOpt {
		stdin = bytes.NewReader(nil)
	}

	code, err := l.execImpl(client, alloc, task, ttyOpt, command, escapeChar, stdin, l.Stdout, l.Stderr)
	if err != nil {
		l.Ui.Error(fmt.Sprintf("failed to exec into task: %v", err))
		return 1
	}

	return code
}

// execImpl invokes the Alloc Exec api call, it also prepares and restores terminal states as necessary.
func (l *AllocExecCommand) execImpl(client *api.Client, alloc *api.Allocation, task string, tty bool,
	command []string, escapeChar string, stdin io.Reader, stdout, stderr io.WriteCloser) (int, error) {

	sizeCh := make(chan api.TerminalSize, 1)

	ctx, cancelFn := context.WithCancel(context.Background())
	defer cancelFn()

	// When tty, ensures we capture all user input and monitor terminal resizes.
	if tty {
		if stdin == nil {
			return -1, fmt.Errorf("stdin is null")
		}

		inCleanup, err := setRawTerminal(stdin)
		if err != nil {
			return -1, err
		}
		defer inCleanup()

		outCleanup, err := setRawTerminalOutput(stdout)
		if err != nil {
			return -1, err
		}
		defer outCleanup()

		sizeCleanup, err := watchTerminalSize(stdout, sizeCh)
		if err != nil {
			return -1, err
		}
		defer sizeCleanup()

		if escapeChar != "" {
			stdin = escapingio.NewReader(stdin, escapeChar[0], func(c byte) bool {
				switch c {
				case '.':
					// need to restore tty state so error reporting here
					// gets emitted at beginning of line
					outCleanup()
					inCleanup()

					stderr.Write([]byte("\nConnection closed\n"))
					cancelFn()
					return true
				default:
					return false
				}
			})
		}
	}

	signalCh := make(chan os.Signal, 1)
	signal.Notify(signalCh, os.Interrupt, syscall.SIGTERM)
	go func() {
		for range signalCh {
			cancelFn()
		}
	}()

	return client.Allocations().Exec(ctx,
		alloc, task, tty, command, stdin, stdout, stderr, sizeCh, nil)
}

func isStdinTty() bool {
	_, isTerminal := term.GetFdInfo(os.Stdin)
	return isTerminal
}

// setRawTerminal sets the stream terminal in raw mode, so process captures
// Ctrl+C and other commands to forward to remote process.
// It returns a cleanup function that restores terminal to original mode.
func setRawTerminal(stream interface{}) (cleanup func(), err error) {
	fd, isTerminal := term.GetFdInfo(stream)
	if !isTerminal {
		return nil, errors.New("not a terminal")
	}

	state, err := term.SetRawTerminal(fd)
	if err != nil {
		return nil, err
	}

	return func() { term.RestoreTerminal(fd, state) }, nil
}

// setRawTerminalOutput sets the output stream in Windows to raw mode,
// so it disables LF -> CRLF translation.
// It's basically a no-op on unix.
func setRawTerminalOutput(stream interface{}) (cleanup func(), err error) {
	fd, isTerminal := term.GetFdInfo(stream)
	if !isTerminal {
		return nil, errors.New("not a terminal")
	}

	state, err := term.SetRawTerminalOutput(fd)
	if err != nil {
		return nil, err
	}

	return func() { term.RestoreTerminal(fd, state) }, nil
}

// watchTerminalSize watches terminal size changes to propagate to remote tty.
func watchTerminalSize(out io.Writer, resize chan<- api.TerminalSize) (func(), error) {
	fd, isTerminal := term.GetFdInfo(out)
	if !isTerminal {
		return nil, errors.New("not a terminal")
	}

	ctx, cancel := context.WithCancel(context.Background())

	signalCh := make(chan os.Signal, 1)
	setupWindowNotification(signalCh)

	sendTerminalSize := func() {
		s, err := term.GetWinsize(fd)
		if err != nil {
			return
		}

		resize <- api.TerminalSize{
			Height: int(s.Height),
			Width:  int(s.Width),
		}
	}
	go func() {
		for {
			select {
			case <-ctx.Done():
				return
			case <-signalCh:
				sendTerminalSize()
			}
		}
	}()

	go func() {
		// send initial size
		sendTerminalSize()
	}()

	return cancel, nil
}
<commit_msg>nomad exec: check stdout for tty as well<commit_after>package command

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"os"
	"os/signal"
	"strings"
	"syscall"

	"github.com/docker/docker/pkg/term"
	"github.com/hashicorp/nomad/api"
	"github.com/hashicorp/nomad/api/contexts"
	"github.com/hashicorp/nomad/helper/escapingio"
	"github.com/posener/complete"
)

type AllocExecCommand struct {
	Meta

	Stdin  io.Reader
	Stdout io.WriteCloser
	Stderr io.WriteCloser
}

func (l *AllocExecCommand) Help() string {
	helpText := `
Usage: nomad alloc exec [options] <allocation> <command>

  Run command inside the environment of the given allocation and task.

General Options:

  ` + generalOptionsUsage() + `

Exec Specific Options:

  -task <task-name>
    Sets the task to exec command in

  -job
    Use a random allocation from the specified job ID.

  -i
    Pass stdin to the container, defaults to true.  Pass -i=false to disable.

  -t
    Allocate a pseudo-tty, defaults to true if stdin is detected to be a tty session.
    Pass -t=false to disable explicitly.

  -e <escape_char>
    Sets the escape character for sessions with a pty (default: '~').  The escape
    character is only recognized at the beginning of a line.  The escape character
    followed by a dot ('.') closes the connection.  Setting the character to
    'none' disables any escapes and makes the session fully transparent.
  `
	return strings.TrimSpace(helpText)
}

func (l *AllocExecCommand) Synopsis() string {
	return "Execute commands in task"
}

func (c *AllocExecCommand) AutocompleteFlags() complete.Flags {
	return mergeAutocompleteFlags(c.Meta.AutocompleteFlags(FlagSetClient),
		complete.Flags{
			"--task": complete.PredictAnything,
			"-job":   complete.PredictAnything,
			"-i":     complete.PredictNothing,
			"-t":     complete.PredictNothing,
			"-e":     complete.PredictSet("none", "~"),
		})
}

func (l *AllocExecCommand) AutocompleteArgs() complete.Predictor {
	return complete.PredictFunc(func(a complete.Args) []string {
		client, err := l.Meta.Client()
		if err != nil {
			return nil
		}

		resp, _, err := client.Search().PrefixSearch(a.Last, contexts.Allocs, nil)
		if err != nil {
			return []string{}
		}
		return resp.Matches[contexts.Allocs]
	})
}

func (l *AllocExecCommand) Name() string { return "alloc exec" }

func (l *AllocExecCommand) Run(args []string) int {
	var job, stdinOpt, ttyOpt bool
	var task string
	var escapeChar string

	flags := l.Meta.FlagSet(l.Name(), FlagSetClient)
	flags.Usage = func() { l.Ui.Output(l.Help()) }
	flags.BoolVar(&job, "job", false, "")
	flags.StringVar(&task, "task", "", "")
	flags.StringVar(&escapeChar, "e", "~", "")

	flags.BoolVar(&stdinOpt, "i", true, "")

	inferredTty := isTty()
	flags.BoolVar(&ttyOpt, "t", inferredTty, "")

	if err := flags.Parse(args); err != nil {
		return 1
	}
	args = flags.Args()

	if ttyOpt && !stdinOpt {
		l.Ui.Error("-i must be enabled if running with tty")
		return 1
	}

	if escapeChar == "none" {
		escapeChar = ""
	} else if len(escapeChar) > 1 {
		l.Ui.Error("-e requires 'none' or a single character")
		return 1
	}

	if numArgs := len(args); numArgs < 1 {
		if job {
			l.Ui.Error("A job ID is required")
		} else {
			l.Ui.Error("An allocation ID is required")
		}

		l.Ui.Error(commandErrorText(l))
		return 1
	} else if numArgs < 2 {
		l.Ui.Error("A command is required")
		l.Ui.Error(commandErrorText(l))
		return 1
	}

	command := args[1:]

	client, err := l.Meta.Client()
	if err != nil {
		l.Ui.Error(fmt.Sprintf("Error initializing client: %v", err))
		return 1
	}

	// If -job is specified, use random allocation, otherwise use provided allocation
	allocID := args[0]
	if job {
		allocID, err = getRandomJobAlloc(client, args[0])
		if err != nil {
			l.Ui.Error(fmt.Sprintf("Error fetching allocations: %v", err))
			return 1
		}
	}

	length := shortId

	// Query the allocation info
	if len(allocID) == 1 {
		l.Ui.Error(fmt.Sprintf("Alloc ID must contain at least two characters."))
		return 1
	}

	allocID = sanitizeUUIDPrefix(allocID)
	allocs, _, err := client.Allocations().PrefixList(allocID)
	if err != nil {
		l.Ui.Error(fmt.Sprintf("Error querying allocation: %v", err))
		return 1
	}
	if len(allocs) == 0 {
		l.Ui.Error(fmt.Sprintf("No allocation(s) with prefix or id %q found", allocID))
		return 1
	}
	if len(allocs) > 1 {
		// Format the allocs
		out := formatAllocListStubs(allocs, false, length)
		l.Ui.Error(fmt.Sprintf("Prefix matched multiple allocations\n\n%s", out))
		return 1
	}
	// Prefix lookup matched a single allocation
	alloc, _, err := client.Allocations().Info(allocs[0].ID, nil)
	if err != nil {
		l.Ui.Error(fmt.Sprintf("Error querying allocation: %s", err))
		return 1
	}

	if task == "" {
		task, err = lookupAllocTask(alloc)

		if err != nil {
			l.Ui.Error(err.Error())
			l.Ui.Error("\nPlease specify the task.")
			return 1
		}
	}

	if err := validateTaskExistsInAllocation(task, alloc); err != nil {
		l.Ui.Error(err.Error())
		return 1
	}

	if l.Stdin == nil {
		l.Stdin = os.Stdin
	}
	if l.Stdout == nil {
		l.Stdout = os.Stdout
	}
	if l.Stderr == nil {
		l.Stderr = os.Stderr
	}

	var stdin io.Reader = l.Stdin
	if !stdinOpt {
		stdin = bytes.NewReader(nil)
	}

	code, err := l.execImpl(client, alloc, task, ttyOpt, command, escapeChar, stdin, l.Stdout, l.Stderr)
	if err != nil {
		l.Ui.Error(fmt.Sprintf("failed to exec into task: %v", err))
		return 1
	}

	return code
}

// execImpl invokes the Alloc Exec api call, it also prepares and restores terminal states as necessary.
func (l *AllocExecCommand) execImpl(client *api.Client, alloc *api.Allocation, task string, tty bool,
	command []string, escapeChar string, stdin io.Reader, stdout, stderr io.WriteCloser) (int, error) {

	sizeCh := make(chan api.TerminalSize, 1)

	ctx, cancelFn := context.WithCancel(context.Background())
	defer cancelFn()

	// When tty, ensures we capture all user input and monitor terminal resizes.
	if tty {
		if stdin == nil {
			return -1, fmt.Errorf("stdin is null")
		}

		inCleanup, err := setRawTerminal(stdin)
		if err != nil {
			return -1, err
		}
		defer inCleanup()

		outCleanup, err := setRawTerminalOutput(stdout)
		if err != nil {
			return -1, err
		}
		defer outCleanup()

		sizeCleanup, err := watchTerminalSize(stdout, sizeCh)
		if err != nil {
			return -1, err
		}
		defer sizeCleanup()

		if escapeChar != "" {
			stdin = escapingio.NewReader(stdin, escapeChar[0], func(c byte) bool {
				switch c {
				case '.':
					// need to restore tty state so error reporting here
					// gets emitted at beginning of line
					outCleanup()
					inCleanup()

					stderr.Write([]byte("\nConnection closed\n"))
					cancelFn()
					return true
				default:
					return false
				}
			})
		}
	}

	signalCh := make(chan os.Signal, 1)
	signal.Notify(signalCh, os.Interrupt, syscall.SIGTERM)
	go func() {
		for range signalCh {
			cancelFn()
		}
	}()

	return client.Allocations().Exec(ctx,
		alloc, task, tty, command, stdin, stdout, stderr, sizeCh, nil)
}

// isTty returns true if both stdin and stdout are a TTY
func isTty() bool {
	_, isStdinTerminal := term.GetFdInfo(os.Stdin)
	_, isStdoutTerminal := term.GetFdInfo(os.Stdout)
	return isStdinTerminal && isStdoutTerminal
}

// setRawTerminal sets the stream terminal in raw mode, so process captures
// Ctrl+C and other commands to forward to remote process.
// It returns a cleanup function that restores terminal to original mode.
func setRawTerminal(stream interface{}) (cleanup func(), err error) {
	fd, isTerminal := term.GetFdInfo(stream)
	if !isTerminal {
		return nil, errors.New("not a terminal")
	}

	state, err := term.SetRawTerminal(fd)
	if err != nil {
		return nil, err
	}

	return func() { term.RestoreTerminal(fd, state) }, nil
}

// setRawTerminalOutput sets the output stream in Windows to raw mode,
// so it disables LF -> CRLF translation.
// It's basically a no-op on unix.
func setRawTerminalOutput(stream interface{}) (cleanup func(), err error) {
	fd, isTerminal := term.GetFdInfo(stream)
	if !isTerminal {
		return nil, errors.New("not a terminal")
	}

	state, err := term.SetRawTerminalOutput(fd)
	if err != nil {
		return nil, err
	}

	return func() { term.RestoreTerminal(fd, state) }, nil
}

// watchTerminalSize watches terminal size changes to propagate to remote tty.
func watchTerminalSize(out io.Writer, resize chan<- api.TerminalSize) (func(), error) {
	fd, isTerminal := term.GetFdInfo(out)
	if !isTerminal {
		return nil, errors.New("not a terminal")
	}

	ctx, cancel := context.WithCancel(context.Background())

	signalCh := make(chan os.Signal, 1)
	setupWindowNotification(signalCh)

	sendTerminalSize := func() {
		s, err := term.GetWinsize(fd)
		if err != nil {
			return
		}

		resize <- api.TerminalSize{
			Height: int(s.Height),
			Width:  int(s.Width),
		}
	}
	go func() {
		for {
			select {
			case <-ctx.Done():
				return
			case <-signalCh:
				sendTerminalSize()
			}
		}
	}()

	go func() {
		// send initial size
		sendTerminalSize()
	}()

	return cancel, nil
}
<|endoftext|>


<commit_before>package main

import (
	// "code.google.com/p/go.net/websocket"
	"encoding/json"
	"flag"
	"fmt"
	"golang.org/x/net/websocket"
	"gopkg.in/redis.v1"
	"io"
	"log"
	"net/http"
	"net/url"
	"os"
	"strings"
)

var (
	logger                    *log.Logger
	logFile                   string
	redisHost                 string
	redisPort                 int
	redisChannel              string
	redisEndpoint             string
	websocketHost             string
	websocketPort             int
	websocketEndpoint         string
	websocketRoute            string
	websocketAllowInsecure    bool
	websocketAllowableOrigins string
	websocketAllowableURLs    []url.URL
	redisClient               *redis.Client
)

// What I'd really like to do is pass in a list of allowable Origin
// URLs but the signature for websocket.Config only allows a single
// url.URL thingy and trying to subclass it resulted in a cascading
// series of "unexpected yak" style errors from the compilers. It
// seems like something that must be possible but my Go-fu is still
// weak... (20140729/straup)

// From Crowley (20140729):
// As for your multiple origins problem:  I recommend you construct
// a map[string]*websocket.Server mapping your full set of origins.
// Do that before you start serving traffic and then you can access
// the map without locks.  Each of those *websocket.Server values
// implements the http.Handler interface so you can in your handler
// look up the right one for the origin and call its ServeHTTP method.

// type pubsocketdConfig struct {
//     websocket.Config
//     Origin []url.URL
// }

func pubsocketdHandler(w http.ResponseWriter, req *http.Request) {

	// See above (20140729/straup)

	if websocketAllowInsecure {

		s := websocket.Server{
			Handler: websocket.Handler(pubSubHandler),
		}

		s.ServeHTTP(w, req)

	} else {

		originURL := websocketAllowableURLs[0]

		pubsocketdConfig := websocket.Config{Origin: &originURL}

		s := websocket.Server{
			Config:    pubsocketdConfig,
			Handshake: pubsocketdHandshake,
			Handler:   websocket.Handler(pubSubHandler),
		}

		s.ServeHTTP(w, req)
	}

}

func pubsocketdHandshake(config *websocket.Config, req *http.Request) (err error) {

	remoteAddr := req.RemoteAddr
	headers := req.Header

	origin := headers.Get("Origin")
	realIP := headers.Get("X-Real-IP")

	if origin == "" {
		logger.Printf("[%s][%s][handshake] missing origin", realIP, remoteAddr)
		return fmt.Errorf("missing origin")
	}

	parsed, err := url.Parse(origin)

	if err != nil {
		logger.Printf("[%s][%s][handshake] failed to parse origin, %v", realIP, remoteAddr, origin)
		return fmt.Errorf("failed to parse origin")
	}

	// See above inre: config.Origin being/becoming a list of url.URLs
	// (20140727/straup)

	if parsed.String() != config.Origin.String() {
		logger.Printf("[%s][%s][handshake] invalid origin, expected %v but got %v", realIP, remoteAddr, config.Origin, parsed)
		return fmt.Errorf("invalid origin")
	}

	logger.Printf("[%s][%s][handshake] OK", realIP, remoteAddr)
	return
}

func pubSubHandler(ws *websocket.Conn) {

	remoteAddr := ws.Request().RemoteAddr
	headers := ws.Request().Header

	realIP := headers.Get("X-Real-IP")

	logger.Printf("[%s][%s][request] OK", realIP, remoteAddr)

	pubsubClient := redisClient.PubSub()
	defer pubsubClient.Close()

	if err := pubsubClient.Subscribe(redisChannel); err != nil {
		logger.Printf("[%s][%s][error] failed to subscribe to pubsub channel %v, because %s", realIP, remoteAddr, redisChannel, err)
		ws.Close()
		return
	}

	log.Printf("[%s][%s][connect] OK", realIP, remoteAddr)

	for ws != nil {

		i, _ := pubsubClient.Receive()

		if msg, _ := i.(*redis.Message); msg != nil {

			// log.Printf("[%s][%s][send] %s", realIP, remoteAddr, msg.Payload)

			var json_blob interface{}
			bytes_blob := []byte(msg.Payload)

			if err := json.Unmarshal(bytes_blob, &json_blob); err != nil {
				logger.Printf("[%s][%s][error] failed to parse JSON %v, because %v", realIP, remoteAddr, msg.Payload, err)
				continue
			}

			if err := websocket.JSON.Send(ws, json_blob); err != nil {
				logger.Printf("[%s][%s][error] failed to send JSON, because %v", realIP, remoteAddr, err)
				ws.Close()
				break
			}

			logger.Printf("[%s][%s][send] OK", realIP, remoteAddr)
		}
	}
}

func main() {

	flag.StringVar(&websocketHost, "ws-host", "127.0.0.1", "WebSocket host")
	flag.IntVar(&websocketPort, "ws-port", 8080, "WebSocket port")
	flag.StringVar(&websocketRoute, "ws-route", "/", "WebSocket route")
	flag.StringVar(&websocketAllowableOrigins, "ws-origin", "", "WebSocket allowable origin(s)")
	flag.BoolVar(&websocketAllowInsecure, "ws-insecure", false, "Allow WebSocket server to run in insecure mode")
	flag.StringVar(&redisHost, "rs-host", "127.0.0.1", "Redis host")
	flag.IntVar(&redisPort, "rs-port", 6379, "Redis port")
	flag.StringVar(&redisChannel, "rs-channel", "pubsocketd", "Redis channel")
	flag.StringVar(&logFile, "ps-log-file", "", "Log requests to this file")

	flag.Parse()

	if logFile != "" {
		file, err := os.OpenFile(logFile, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0660)

		if err != nil {
			panic(err)
		}

		multi := io.MultiWriter(file, os.Stdout)
		logger = log.New(multi, "[pubsocketd] ", log.Ldate|log.Ltime|log.Lshortfile)
	} else {
		multi := io.MultiWriter(os.Stdout)
		logger = log.New(multi, "[pubsocketd] ", log.Ldate|log.Ltime|log.Lshortfile)
	}

	if !websocketAllowInsecure {
		if websocketAllowableOrigins == "" {
			logger.Fatalf("Missing allowable Origin (-ws-origin=http://example.com)")
		}

		allowed := strings.Split(websocketAllowableOrigins, ",")
		count := len(allowed)

		if count > 1 {
			logger.Fatalf("Only one origin server is supported at the moment")
		}

		websocketAllowableURLs = make([]url.URL, 0, count)

		for _, test := range allowed {

			test := strings.TrimSpace(test)

			url, err := url.Parse(test)

			if err != nil {
				logger.Fatalf("Invalid Origin parameter: %v, %v", test, err)
			}

			websocketAllowableURLs = append(websocketAllowableURLs, *url)
		}
	}

	websocketEndpoint = fmt.Sprintf("%s:%d", websocketHost, websocketPort)
	redisEndpoint = fmt.Sprintf("%s:%d", redisHost, redisPort)

	redisClient = redis.NewTCPClient(&redis.Options{
		Addr: redisEndpoint,
	})

	defer redisClient.Close()

	// Normally this is the sort of thing you'd expect to do:
	// http.Handle(websocketRoute, websocket.Handler(pubSubHandler))

	// However since we're going to be aggressively paranoid about checking
	// the Origin headers we're going to set up our own websocket Server
	// thingy complete with custom Config and Handshake directive and
	// pass the whole thing off to HandleFunc (20140727/straup)

	// See also:
	// http://www.christian-schneider.net/CrossSiteWebSocketHijacking.html
	// https://code.google.com/p/go/source/browse/websocket/server.go?repo=net

	http.HandleFunc(websocketRoute, pubsocketdHandler)

	if websocketAllowInsecure {
		logger.Printf("[init] listening for websocket requests on " + websocketEndpoint + websocketRoute + ", in INSECURE MODE which is not advised for production use")

	} else {
		logger.Printf("[init] listening for websocket requests on " + websocketEndpoint + websocketRoute + ", from " + websocketAllowableOrigins)
	}

	logger.Printf("[init] listening for pubsub messages from " + redisEndpoint + " sent to the " + redisChannel + " channel")

	if err := http.ListenAndServe(websocketEndpoint, nil); err != nil {
		logger.Fatalf("Failed to start websocket server, because %v", err)
	}
}
<commit_msg>:lock: pubsocketd optionally speaks TLS<commit_after>package main

import (
	// "code.google.com/p/go.net/websocket"
	"encoding/json"
	"flag"
	"fmt"
	"golang.org/x/net/websocket"
	"gopkg.in/redis.v1"
	"io"
	"log"
	"net/http"
	"net/url"
	"os"
	"strings"
)

var (
	logger                    *log.Logger
	logFile                   string
	redisHost                 string
	redisPort                 int
	redisChannel              string
	redisEndpoint             string
	websocketHost             string
	websocketPort             int
	websocketEndpoint         string
	websocketRoute            string
	websocketAllowInsecure    bool
	websocketAllowableOrigins string
	websocketAllowableURLs    []url.URL
	redisClient               *redis.Client
	tlsCert                   string
	tlsKey                    string
)

// What I'd really like to do is pass in a list of allowable Origin
// URLs but the signature for websocket.Config only allows a single
// url.URL thingy and trying to subclass it resulted in a cascading
// series of "unexpected yak" style errors from the compilers. It
// seems like something that must be possible but my Go-fu is still
// weak... (20140729/straup)

// From Crowley (20140729):
// As for your multiple origins problem:  I recommend you construct
// a map[string]*websocket.Server mapping your full set of origins.
// Do that before you start serving traffic and then you can access
// the map without locks.  Each of those *websocket.Server values
// implements the http.Handler interface so you can in your handler
// look up the right one for the origin and call its ServeHTTP method.

// type pubsocketdConfig struct {
//     websocket.Config
//     Origin []url.URL
// }

func pubsocketdHandler(w http.ResponseWriter, req *http.Request) {

	// See above (20140729/straup)

	if websocketAllowInsecure {

		s := websocket.Server{
			Handler: websocket.Handler(pubSubHandler),
		}

		s.ServeHTTP(w, req)

	} else {

		originURL := websocketAllowableURLs[0]

		pubsocketdConfig := websocket.Config{Origin: &originURL}

		s := websocket.Server{
			Config:    pubsocketdConfig,
			Handshake: pubsocketdHandshake,
			Handler:   websocket.Handler(pubSubHandler),
		}

		s.ServeHTTP(w, req)
	}

}

func pubsocketdHandshake(config *websocket.Config, req *http.Request) (err error) {

	remoteAddr := req.RemoteAddr
	headers := req.Header

	origin := headers.Get("Origin")
	realIP := headers.Get("X-Real-IP")

	if origin == "" {
		logger.Printf("[%s][%s][handshake] missing origin", realIP, remoteAddr)
		return fmt.Errorf("missing origin")
	}

	parsed, err := url.Parse(origin)

	if err != nil {
		logger.Printf("[%s][%s][handshake] failed to parse origin, %v", realIP, remoteAddr, origin)
		return fmt.Errorf("failed to parse origin")
	}

	// See above inre: config.Origin being/becoming a list of url.URLs
	// (20140727/straup)

	if parsed.String() != config.Origin.String() {
		logger.Printf("[%s][%s][handshake] invalid origin, expected %v but got %v", realIP, remoteAddr, config.Origin, parsed)
		return fmt.Errorf("invalid origin")
	}

	logger.Printf("[%s][%s][handshake] OK", realIP, remoteAddr)
	return
}

func pubSubHandler(ws *websocket.Conn) {

	remoteAddr := ws.Request().RemoteAddr
	headers := ws.Request().Header

	realIP := headers.Get("X-Real-IP")

	logger.Printf("[%s][%s][request] OK", realIP, remoteAddr)

	pubsubClient := redisClient.PubSub()
	defer pubsubClient.Close()

	if err := pubsubClient.Subscribe(redisChannel); err != nil {
		logger.Printf("[%s][%s][error] failed to subscribe to pubsub channel %v, because %s", realIP, remoteAddr, redisChannel, err)
		ws.Close()
		return
	}

	log.Printf("[%s][%s][connect] OK", realIP, remoteAddr)

	for ws != nil {

		i, _ := pubsubClient.Receive()

		if msg, _ := i.(*redis.Message); msg != nil {

			// log.Printf("[%s][%s][send] %s", realIP, remoteAddr, msg.Payload)

			var json_blob interface{}
			bytes_blob := []byte(msg.Payload)

			if err := json.Unmarshal(bytes_blob, &json_blob); err != nil {
				logger.Printf("[%s][%s][error] failed to parse JSON %v, because %v", realIP, remoteAddr, msg.Payload, err)
				continue
			}

			if err := websocket.JSON.Send(ws, json_blob); err != nil {
				logger.Printf("[%s][%s][error] failed to send JSON, because %v", realIP, remoteAddr, err)
				ws.Close()
				break
			}

			logger.Printf("[%s][%s][send] OK", realIP, remoteAddr)
		}
	}
}

func main() {

	flag.StringVar(&websocketHost, "ws-host", "127.0.0.1", "WebSocket host")
	flag.IntVar(&websocketPort, "ws-port", 8080, "WebSocket port")
	flag.StringVar(&websocketRoute, "ws-route", "/", "WebSocket route")
	flag.StringVar(&websocketAllowableOrigins, "ws-origin", "", "WebSocket allowable origin(s)")
	flag.BoolVar(&websocketAllowInsecure, "ws-insecure", false, "Allow WebSocket server to run in insecure mode")
	flag.StringVar(&redisHost, "rs-host", "127.0.0.1", "Redis host")
	flag.IntVar(&redisPort, "rs-port", 6379, "Redis port")
	flag.StringVar(&redisChannel, "rs-channel", "pubsocketd", "Redis channel")
	flag.StringVar(&logFile, "ps-log-file", "", "Log requests to this file")
	flag.StringVar(&tlsCert, "tls-cert", "", "Path to TLS certificate")
	flag.StringVar(&tlsKey, "tls-key", "", "Path to TLS key")

	flag.Parse()

	if logFile != "" {
		file, err := os.OpenFile(logFile, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0660)

		if err != nil {
			panic(err)
		}

		multi := io.MultiWriter(file, os.Stdout)
		logger = log.New(multi, "[pubsocketd] ", log.Ldate|log.Ltime|log.Lshortfile)
	} else {
		multi := io.MultiWriter(os.Stdout)
		logger = log.New(multi, "[pubsocketd] ", log.Ldate|log.Ltime|log.Lshortfile)
	}

	if !websocketAllowInsecure {
		if websocketAllowableOrigins == "" {
			logger.Fatalf("Missing allowable Origin (-ws-origin=http://example.com)")
		}

		allowed := strings.Split(websocketAllowableOrigins, ",")
		count := len(allowed)

		if count > 1 {
			logger.Fatalf("Only one origin server is supported at the moment")
		}

		websocketAllowableURLs = make([]url.URL, 0, count)

		for _, test := range allowed {

			test := strings.TrimSpace(test)

			url, err := url.Parse(test)

			if err != nil {
				logger.Fatalf("Invalid Origin parameter: %v, %v", test, err)
			}

			websocketAllowableURLs = append(websocketAllowableURLs, *url)
		}
	}

	websocketEndpoint = fmt.Sprintf("%s:%d", websocketHost, websocketPort)
	redisEndpoint = fmt.Sprintf("%s:%d", redisHost, redisPort)

	redisClient = redis.NewTCPClient(&redis.Options{
		Addr: redisEndpoint,
	})

	defer redisClient.Close()

	// Normally this is the sort of thing you'd expect to do:
	// http.Handle(websocketRoute, websocket.Handler(pubSubHandler))

	// However since we're going to be aggressively paranoid about checking
	// the Origin headers we're going to set up our own websocket Server
	// thingy complete with custom Config and Handshake directive and
	// pass the whole thing off to HandleFunc (20140727/straup)

	// See also:
	// http://www.christian-schneider.net/CrossSiteWebSocketHijacking.html
	// https://code.google.com/p/go/source/browse/websocket/server.go?repo=net

	http.HandleFunc(websocketRoute, pubsocketdHandler)

	if websocketAllowInsecure {
		logger.Printf("[init] listening for websocket requests on " + websocketEndpoint + websocketRoute + ", in INSECURE MODE which is not advised for production use")

	} else {
		logger.Printf("[init] listening for websocket requests on " + websocketEndpoint + websocketRoute + ", from " + websocketAllowableOrigins)
	}

	logger.Printf("[init] listening for pubsub messages from " + redisEndpoint + " sent to the " + redisChannel + " channel")

	if tlsCert != "" && tlsKey != "" {
		if err := http.ListenAndServeTLS(websocketEndpoint, tlsCert, tlsKey, nil); err != nil {
			logger.Fatalf("Failed to start websocket server, because %v", err)
		}
	} else {
		if err := http.ListenAndServe(websocketEndpoint, nil); err != nil {
			logger.Fatalf("Failed to start websocket server, because %v", err)
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"log"

	"github.com/VojtechVitek/rerun"
)

var (
	watch flagStringSlice
)

type flagStringSlice []string

func (f *flagStringSlice) String() string {
	return fmt.Sprintf("%v", *f)
}

func (f *flagStringSlice) Set(value string) error {
	*f = append(*f, value)
	return nil
}

func init() {
	flag.Var(&watch, "watch", "Watch directory/file")
}

type argType int

const (
	argNone argType = iota
	argWatch
	argIgnore
	argRun
)

func main() {
	watcher, err := rerun.NewWatcher()
	if err != nil {
		log.Fatal(err)
	}

	go watcher.Watch()
	defer watcher.Close()

	select {}
}
<commit_msg>CLI parser for -watch, -ignore and -run args<commit_after>package main

import (
	"fmt"
	"log"

	"os"

	"github.com/VojtechVitek/rerun"
)

func main() {
	watcher, err := rerun.NewWatcher()
	if err != nil {
		log.Fatal(err)
	}

	args := os.Args[1:]
	mode := argNone

	// Parse command line arguments.
	// -watch dirs ...
	// -ignore dirs ...
	// -run command ...
loop:
	for i, arg := range args {
		switch mode {
		case argNone, argWatch, argIgnore:
			switch arg {
			case "-watch":
				mode = argWatch
				continue
			case "-ignore":
				mode = argIgnore
				continue
			case "-run":
				mode = argRun
				continue
			}
		}

		switch mode {
		case argWatch:
			watcher.Add(arg)
		case argIgnore:
			watcher.Ignore(arg)
		case argRun:
			args = args[i:]
			break loop
		default:
			break loop
		}
	}

	if mode == argNone {
		log.Fatal("interactive mode")
	}

	fmt.Printf("Run: %+v\n", args)

	go watcher.Watch()
	defer watcher.Close()

	select {}
}
<|endoftext|>


<commit_before>/*
seqls - list directory contents, rolled up into file sequences
*/
package main

import (
	"flag"
	"fmt"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"sync"

	"github.com/MichaelTJones/walk"
	"github.com/justinfx/gofileseq"
)

var (
	numWorkers = 50
	recurse    = flag.Bool("r", false, `recursively scan all sub-directories`)
	allFiles   = flag.Bool("a", false, `list all files, including those without frame patterns`)
)

func init() {
	// If $GOMAXPROCS isn't set, use the full capacity of the machine.
	// For small machines, use at least 4 threads.
	if os.Getenv("GOMAXPROCS") == "" {
		n := runtime.NumCPU()
		if n < 4 {
			n = 4
		}
		runtime.GOMAXPROCS(n)
	}
}

const usage = `seqls: list directory contents, rolled up into file sequences

Usage:  seqls [options] [path [...]]

One or more paths may be provided. If no path is provided, the current
working directory will be scanned.
`

func main() {
	flag.Usage = func() {
		fmt.Fprintln(os.Stderr, usage)
		flag.PrintDefaults()
		fmt.Fprintln(os.Stderr, "\n")
		os.Exit(1)
	}
	flag.Parse()

	roots := flag.Args()
	if len(roots) == 0 {
		roots = []string{"."}
	}

	var wg sync.WaitGroup

	pathChan := make(chan string)
	seqChan := make(chan fileseq.FileSequences)

	// fmt.Printf("launching %d workers\n", numWorkers)

	var listerFn func(string) (fileseq.FileSequences, error)
	if *allFiles {
		listerFn = fileseq.ListFiles
	} else {
		listerFn = fileseq.FindSequencesOnDisk
	}

	// Start the workers to find sequences
	for i := 0; i < numWorkers; i++ {
		wg.Add(1)
		go func() {
			for path := range pathChan {
				seqs, err := listerFn(path)
				if err != nil {
					// Bail out of the app for any path error
					fmt.Fprintf(os.Stderr, "Error finding sequence in dir %q: %s\n", path, err)
					os.Exit(255)
				}
				seqChan <- seqs
			}
			wg.Done()
		}()
	}

	// Load the root paths into the source channel
	unique := cleanDirs(roots)
	go func() {
		if *recurse {
			loadRecursive(unique, pathChan)
		} else {
			for r := range unique {
				pathChan <- r
			}
		}
		close(pathChan)
	}()

	// Will close the output channel when no more source
	// paths are being added
	go func() {
		wg.Wait()
		close(seqChan)
	}()

	// Pull out all processed sequences and print
	for seqs := range seqChan {
		for _, s := range seqs {
			fmt.Println(s)
		}
	}

	os.Exit(0)
}

// Parallel walk the root paths and populate the path
// channel for the worker routines to consume.
func loadRecursive(roots uniquePaths, out chan string) {
	walkFn := func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if info.IsDir() {
			out <- path

		} else if fi, err := os.Stat(path); err == nil && fi.IsDir() {
			out <- path
		}
		return nil
	}
	for r := range roots {
		r := r
		if err := walk.Walk(r, walkFn); err != nil {
			fmt.Fprintf(os.Stderr, "Error: Failed to read dir %q: %s\n", r, err)
		}
	}
}

type uniquePaths map[string]struct{}

// Take a list of paths and reduce them to cleaned
// and unique paths.
func cleanDirs(paths []string) uniquePaths {
	u := make(uniquePaths)
	var (
		fi  os.FileInfo
		err error
	)
	for _, p := range paths {
		p := strings.TrimSpace(filepath.Clean(p))
		if p == "" {
			continue
		}
		if fi, err = os.Stat(p); err != nil {
			fmt.Fprintf(os.Stderr, "Error: Failed to read path %q: %s\n", p, err)
			continue
		}
		if !fi.IsDir() {
			continue
		}
		u[p] = struct{}{}
	}
	return u
}
<commit_msg>seqls: (features) Added listing all files, long listing, abspath, and human units<commit_after>/*
seqls - list directory contents, rolled up into file sequences
*/
package main

import (

	// "flag"

	"fmt"
	"io"
	"os"
	"os/user"
	"path/filepath"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"text/tabwriter"
	"time"

	"github.com/MichaelTJones/walk"
	"github.com/jessevdk/go-flags"
	"github.com/justinfx/gofileseq"
)

// number of goroutines to spawn for processing directories
// and building FileSequence results
var numWorkers = 50

var opts struct {
	Recurse   bool `short:"r" long:"recurse" description:"recursively scan all sub-directories"`
	AllFiles  bool `short:"a" long:"all" description:"list all files, including those without frame patterns"`
	LongList  bool `short:"l" long:"long" description:"long listing; include extra stat information"`
	AbsPath   bool `short:"f" long:"full" description:"show absolute paths"`
	HumanSize bool `short:"H" long:"human" description:"when using long listing, show human-readable file size units"`
}

const DateFmt = `Jan _2 15:04`

func init() {
	// If $GOMAXPROCS isn't set, use the full capacity of the machine.
	// For small machines, use at least 4 threads.
	if os.Getenv("GOMAXPROCS") == "" {
		n := runtime.NumCPU()
		if n < 4 {
			n = 4
		}
		runtime.GOMAXPROCS(n)
	}
}

const usage = `[options] [path [...]]

List directory contents, rolled up into file sequences.

One or more paths may be provided. If no path is provided, the current
working directory will be scanned.

Only files are shown in listing results.
`

func main() {
	parser := flags.NewParser(&opts, flags.HelpFlag|flags.PrintErrors)
	parser.Name = "seqls"
	parser.Usage = usage
	roots, err := parser.Parse()

	if err != nil {
		os.Exit(1)
	}

	if len(roots) == 0 {
		// Use the current directory if specific dirs were not passed
		roots = []string{"."}
	}

	var wg sync.WaitGroup

	pathChan := make(chan string)
	seqChan := make(chan fileseq.FileSequences)

	// fmt.Printf("launching %d workers\n", numWorkers)

	var listerFn func(string) (fileseq.FileSequences, error)
	if opts.AllFiles {
		listerFn = fileseq.ListFiles
	} else {
		listerFn = fileseq.FindSequencesOnDisk
	}

	// Start the workers to find sequences
	for i := 0; i < numWorkers; i++ {
		wg.Add(1)
		go func() {
			for path := range pathChan {
				seqs, err := listerFn(path)
				if err != nil {
					// Bail out of the app for any path error
					fmt.Fprintf(os.Stderr, "Error finding sequence in dir %q: %s\n", path, err)
					os.Exit(255)
				}
				seqChan <- seqs
			}
			wg.Done()
		}()
	}

	// Load the root paths into the source channel
	unique := cleanDirs(roots)
	go func() {
		if opts.Recurse {
			loadRecursive(unique, pathChan)
		} else {
			for r := range unique {
				pathChan <- r
			}
		}
		close(pathChan)
	}()

	// Will close the output channel when no more source
	// paths are being added
	go func() {
		wg.Wait()
		close(seqChan)
	}()

	// Pull out all processed sequences and print
	var w io.Writer
	if opts.LongList {
		w = tabwriter.NewWriter(os.Stdout, 5, 0, 2, ' ', 0)
	} else {
		w = os.Stdout
	}

	var printer func(io.Writer, *fileseq.FileSequence)
	if opts.LongList {
		printer = printLongListing
	} else {
		printer = printShortListing
	}

	// Start processing the results that are feeding in the channel
	for seqs := range seqChan {
		for _, s := range seqs {
			printer(w, s)
		}
	}

	if opts.LongList {
		// Flush the tabwriter buffer
		w.(*tabwriter.Writer).Flush()
	}

	os.Exit(0)
}

// Parallel walk the root paths and populate the path
// channel for the worker routines to consume.
func loadRecursive(roots uniquePaths, out chan string) {
	walkFn := func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if info.IsDir() {
			out <- path

		} else if fi, err := os.Stat(path); err == nil && fi.IsDir() {
			out <- path
		}
		return nil
	}
	for r := range roots {
		r := r
		if err := walk.Walk(r, walkFn); err != nil {
			fmt.Fprintf(os.Stderr, "Error: Failed to read dir %q: %s\n", r, err)
		}
	}
}

type uniquePaths map[string]struct{}

// Take a list of paths and reduce them to cleaned
// and unique paths.
func cleanDirs(paths []string) uniquePaths {
	u := make(uniquePaths)
	var (
		fi  os.FileInfo
		err error
	)
	for _, p := range paths {
		p := strings.TrimSpace(filepath.Clean(p))
		if p == "" {
			continue
		}
		if fi, err = os.Stat(p); err != nil {
			fmt.Fprintf(os.Stderr, "Error: Failed to read path %q: %s\n", p, err)
			continue
		}
		if !fi.IsDir() {
			continue
		}
		u[p] = struct{}{}
	}
	return u
}

func printShortListing(w io.Writer, fs *fileseq.FileSequence) {
	str := fs.String()
	if opts.AbsPath {
		abs, err := filepath.Abs(str)
		if err == nil {
			str = abs
		}
	}
	fmt.Fprintln(w, str)
}

func printLongListing(w io.Writer, fs *fileseq.FileSequence) {
	var (
		err      error
		byteSize int64
		stat     os.FileInfo
		str      string
	)

	count := fs.Len()
	if count == 1 {
		// Only a single file
		str = fs.Index(0)
		stat, err = os.Stat(str)
		if err == nil {
			byteSize = stat.Size()
		}

	} else {
		str = fs.String()
		// For a sequence of files, we want to get the most
		// recent modtime, as well as the total size of the
		// combined files.
		mod := time.Unix(0, 0)
		var thisStat os.FileInfo
		var thisStr string
		for i := 0; i < count; i++ {
			thisStr = fs.Index(i)
			thisStat, err = os.Stat(thisStr)
			if err != nil {
				continue
			}
			byteSize += thisStat.Size()
			if thisStat.ModTime().After(mod) {
				stat = thisStat
				mod = thisStat.ModTime()
			}
		}
	}

	if opts.AbsPath {
		abs, err := filepath.Abs(str)
		if err == nil {
			str = abs
		}
	}

	if err != nil {
		fmt.Println(str)
		fmt.Fprintf(os.Stderr, "Error stating file: %s\n", err.Error())
		return
	}

	stat_t := stat.Sys().(*syscall.Stat_t)
	usr := strconv.Itoa(int(stat_t.Uid))
	gid := "-"

	uObj, err := user.LookupId(usr)
	if err == nil {
		usr = uObj.Username
		gid = uObj.Gid
	}

	var size interface{}
	if opts.HumanSize {
		size = ByteSize(byteSize)
	} else {
		size = byteSize
	}

	fmt.Fprintf(w, "%s\t%s\t%s\t%v\t%s\t%s\n",
		stat.Mode(),
		usr,
		gid,
		size,
		stat.ModTime().Format(DateFmt),
		str)

}

type ByteSize float64

const (
	_           = iota // ignore first value by assigning to blank identifier
	KB ByteSize = 1 << (10 * iota)
	MB
	GB
	TB
)

func (b ByteSize) String() string {
	switch {
	case b >= TB:
		return fmt.Sprintf("%.1fT", b/TB)
	case b >= GB:
		return fmt.Sprintf("%.1fG", b/GB)
	case b >= MB:
		return fmt.Sprintf("%.1fM", b/MB)
	case b >= KB:
		return fmt.Sprintf("%.1fK", b/KB)
	case b == 0:
		return strconv.Itoa(0)
	}
	return fmt.Sprintf("%.0f", b)
}
<|endoftext|>