<commit_before>/*
Calc is a calculator designed to be run on the command line.
Expressions rarely require quoting and work easily
with other command line tools that produce or require
whitespace-separated text.

Here is a brief overview by demonstration:

	# The command line is revert-polish - operands are
	# pushed onto a stack; operators pop them off and
	# push the result(s).
	% calc 3 4 5 pow pow
	373391848741020043532959754184866588225409776783734007750636931722079040617265251229993688938803977220468765065431475158108727054592160858581351336982809187314191748594262580938807019951956404285571818041046681288797402925517668012340617298396574731619152386723046235125934896058590588284654793540505936202376547807442730582144527058988756251452817793413352141920744623027518729185432862375737063985485319476416926263819972887006907013899256524297198527698749274196276811060702333710356481
	%
	% # The default type for a literal without a decimal point is a big integer.
	% # Conversion operators (float, big, rat) can convert from one to another.
	% calc 3 4 5 pow pow float
	-7.087363323575677e+18
	%
	% # A sequence of operations enclosed in [ ] is executed
	% # repeatedly until the stack is empty.
	% calc 3 6 8 10 [ + ]
	27
	%
	% # x is used for multiplication to avoid the
	% # need to quote it in the shell.
	% calc 2 3 4 [ x ]
	24
	%
	% # If there's more than one value left on the stack,
	% # they're all printed.
	% calc 5 7 8 10 +
	5
	7
	18
	%
	% # A sequence of operations enclosed in [[ ]] is executed
	% # for every item (or set of items, if it uses more
	% # than one) in the stack.
	% calc 150 256 645 [[ 10 / ]]
	15
	25
	64
	%
	% # We can nest sequences. Within a sequence,
	% # a repetition operator only goes back to the
	% # start of the sequence. (BUG here)
	% calc 3 4 5 [[ 2 3 4 5 [ + ] x ]]
	126
	70
	%  calc 5 6 7 8 1 2 [[ + ]]
	11
	15
	3
	%
	% # Operations are on integers by default...
	% calc 5 13 /
	0
	%
	% # but can be converted to rationals. In any
	% # operation, the first operand (deepest in the stack)
	% # determines the result. Other operands are
	% # converted to its type.
	% calc 5 rat 13 /
	5/13
	% calc 5 rat 13 / 2/7 +
	61/91
	%
	% # The % operator can be used to determine an output format.
	% calc 3 5 pow %x
	f3
	%
	% # Any format acceptable to Go's fmt.Print may be used.
	% calc 3 5 pow '%#7.5x'
	0x000f3
	%
	% # The % format operator "taints" the top value on the stack
	% # with the format - any operation on that value
	% # will taint the result too.
	% calc 3 '%6d' 5 pow
	   243
	%
	% # If the % format operator is executed on an empty stack,
	% # it sets the default format for the stack.
	% calc %o 10 11 12
	12
	13
	14
	%
*/
package documentation
<commit_msg>calc: rename "documentation" to "main"<commit_after>/*
Calc is a calculator designed to be run on the command line.
Expressions rarely require quoting and work easily
with other command line tools that produce or require
whitespace-separated text.

Here is a brief overview by demonstration:

	# The command line is revert-polish - operands are
	# pushed onto a stack; operators pop them off and
	# push the result(s).
	% calc 3 4 5 pow pow
	373391848741020043532959754184866588225409776783734007750636931722079040617265251229993688938803977220468765065431475158108727054592160858581351336982809187314191748594262580938807019951956404285571818041046681288797402925517668012340617298396574731619152386723046235125934896058590588284654793540505936202376547807442730582144527058988756251452817793413352141920744623027518729185432862375737063985485319476416926263819972887006907013899256524297198527698749274196276811060702333710356481
	%
	% # The default type for a literal without a decimal point is a big integer.
	% # Conversion operators (float, big, rat) can convert from one to another.
	% calc 3 4 5 pow pow float
	-7.087363323575677e+18
	%
	% # A sequence of operations enclosed in [ ] is executed
	% # repeatedly until the stack is empty.
	% calc 3 6 8 10 [ + ]
	27
	%
	% # x is used for multiplication to avoid the
	% # need to quote it in the shell.
	% calc 2 3 4 [ x ]
	24
	%
	% # If there's more than one value left on the stack,
	% # they're all printed.
	% calc 5 7 8 10 +
	5
	7
	18
	%
	% # A sequence of operations enclosed in [[ ]] is executed
	% # for every item (or set of items, if it uses more
	% # than one) in the stack.
	% calc 150 256 645 [[ 10 / ]]
	15
	25
	64
	%
	% # We can nest sequences. Within a sequence,
	% # a repetition operator only goes back to the
	% # start of the sequence. (BUG here)
	% calc 3 4 5 [[ 2 3 4 5 [ + ] x ]]
	126
	70
	%  calc 5 6 7 8 1 2 [[ + ]]
	11
	15
	3
	%
	% # Operations are on integers by default...
	% calc 5 13 /
	0
	%
	% # but can be converted to rationals. In any
	% # operation, the first operand (deepest in the stack)
	% # determines the result. Other operands are
	% # converted to its type.
	% calc 5 rat 13 /
	5/13
	% calc 5 rat 13 / 2/7 +
	61/91
	%
	% # The % operator can be used to determine an output format.
	% calc 3 5 pow %x
	f3
	%
	% # Any format acceptable to Go's fmt.Print may be used.
	% calc 3 5 pow '%#7.5x'
	0x000f3
	%
	% # The % format operator "taints" the top value on the stack
	% # with the format - any operation on that value
	% # will taint the result too.
	% calc 3 '%6d' 5 pow
	   243
	%
	% # If the % format operator is executed on an empty stack,
	% # it sets the default format for the stack.
	% calc %o 10 11 12
	12
	13
	14
	%
*/
package main
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"strconv"

	"github.com/ihcsim/go-soundex"
	"github.com/jessevdk/go-flags"
)

type opts struct {
	Info   bool `short:"i" long:"info" description:"Show Soundex algorithm information"`
	Sample bool `short:"s" long:"sample" description:"Show the sample Soundex chart"`
}

func main() {
	o := opts{}
	args, _ := flags.Parse(&o)

	if o.Info {
		printInfo()
	}

	if o.Sample {
		printSample()
	}

	if len(args) > 0 {
		for _, arg := range args {
			if _, err := strconv.Atoi(arg); err == nil {
				fmt.Println("Unable to encode a number")
				continue
			}

			if arg[0] == '-' {
				continue
			}

			fmt.Printf("Soundex code of %s:%s\n", arg, soundex.Encode(arg))
		}
	}
}

func printInfo() {
	fmt.Println(`
American Soundex Algorithm
==========================

The Soundex value of a name is determined by:

1. Retain the first letter of the name and drop all other occurrences of a, e, i, o, u, y, h, w.
2. Replace consonants with digits as follows (after the first letter):
  * b, f, p, v → 1
  * c, g, j, k, q, s, x, z → 2
  * d, t → 3
  * l → 4
  * m, n → 5
  * r → 6
3. If two or more letters with the same number are adjacent in the original name (before step 1), only retain the first letter; also two letters with the same number separated by 'h' or 'w' are coded as a single number, whereas such letters separated by a vowel are coded twice. This rule also applies to the first letter.
4. If you have too few letters in your word that you can't assign three numbers, append with zeros until there are three numbers. If you have more than 3 letters, just retain the first 3 numbers.`)
}

func printSample() {
	fmt.Println(`
Soundex Code Sample
===================

The following is a list of sample names with their respective Soundex code.

Soundex Code | Name
------------ | ----
A261         | Ashcraft
A261         | Ashcroft
B620         | Burroughs
B620         | Burrows
C532         | Ciondecks
E460         | Ellery
E460         | Euler
E251         | Example
G200         | Gauss
G200         | Ghosh
H416         | Heilbronn
H416         | Hilbert
K530         | Kant
K530         | Knuth
L300         | Ladd
L222         | Lissajous
L300         | Lloyd
L222         | Lukasiewicz
O600         | O'Hara
R163         | Robert
R150         | Rubin
R163         | Rupert
S532         | Soundex
T522         | Tymczak
W350         | Wheaton
`)
}
<commit_msg>Replace the usage of helper functions with Parser.<commit_after>package main

import (
	"fmt"
	"os"
	"strconv"

	"github.com/ihcsim/go-soundex"
	"github.com/jessevdk/go-flags"
)

type opts struct {
	Info   bool `short:"i" long:"info" description:"Show Soundex algorithm information"`
	Sample bool `short:"s" long:"sample" description:"Show the sample Soundex chart"`
}

func main() {
	o := opts{}
	p := flags.NewParser(&o, flags.Default)
	args, err := p.Parse()
	if err != nil {
		fmt.Println("Failed to parse CLI arguments")
		os.Exit(1)
	}
	p.Usage = "[OPTIONS] <name>"

	if o.Info {
		printInfo()
		os.Exit(0)
	}

	if o.Sample {
		printSample()
		os.Exit(0)
	}

	for _, arg := range args {
		if _, err := strconv.Atoi(arg); err == nil {
			fmt.Println("Unable to encode a number")
			continue
		}

		if arg[0] == '-' {
			continue
		}

		fmt.Printf("Soundex code of %s:%s\n", arg, soundex.Encode(arg))
	}
}

func printInfo() {
	fmt.Println(`
American Soundex Algorithm
==========================

The Soundex value of a name is determined by:

1. Retain the first letter of the name and drop all other occurrences of a, e, i, o, u, y, h, w.
2. Replace consonants with digits as follows (after the first letter):
  * b, f, p, v → 1
  * c, g, j, k, q, s, x, z → 2
  * d, t → 3
  * l → 4
  * m, n → 5
  * r → 6
3. If two or more letters with the same number are adjacent in the original name (before step 1), only retain the first letter; also two letters with the same number separated by 'h' or 'w' are coded as a single number, whereas such letters separated by a vowel are coded twice. This rule also applies to the first letter.
4. If you have too few letters in your word that you can't assign three numbers, append with zeros until there are three numbers. If you have more than 3 letters, just retain the first 3 numbers.`)
}

func printSample() {
	fmt.Println(`
Soundex Code Sample
===================

The following is a list of sample names with their respective Soundex code.

Soundex Code | Name
------------ | ----
A261         | Ashcraft
A261         | Ashcroft
B620         | Burroughs
B620         | Burrows
C532         | Ciondecks
E460         | Ellery
E460         | Euler
E251         | Example
G200         | Gauss
G200         | Ghosh
H416         | Heilbronn
H416         | Hilbert
K530         | Kant
K530         | Knuth
L300         | Ladd
L222         | Lissajous
L300         | Lloyd
L222         | Lukasiewicz
O600         | O'Hara
R163         | Robert
R150         | Rubin
R163         | Rupert
S532         | Soundex
T522         | Tymczak
W350         | Wheaton
`)
}
<|endoftext|>


<commit_before>// Copyright 2016 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 (
	"flag"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"strings"

	"github.com/golang/dep"
	"github.com/golang/dep/internal"
	"github.com/pkg/errors"
	"github.com/sdboyer/gps"
	"github.com/sdboyer/gps/pkgtree"
)

const initShortHelp = `Initialize a new project with manifest and lock files`
const initLongHelp = `
Initialize the project at filepath root by parsing its dependencies and writing
manifest and lock files. If root isn't specified, use the current directory.

The version of each dependency will reflect the current state of the GOPATH. If
a dependency doesn't exist in the GOPATH, it won't be written to the manifest,
but it will be solved-for, and will appear in the lock.

Note: init may use the network to solve the dependency graph.

Note: init does NOT vendor dependencies at the moment. See dep ensure.
`

func (cmd *initCommand) Name() string      { return "init" }
func (cmd *initCommand) Args() string      { return "[root]" }
func (cmd *initCommand) ShortHelp() string { return initShortHelp }
func (cmd *initCommand) LongHelp() string  { return initLongHelp }
func (cmd *initCommand) Hidden() bool      { return false }

func (cmd *initCommand) Register(fs *flag.FlagSet) {
	fs.BoolVar(&cmd.noExamples, "no-examples", false, "don't include example in Gopkg.toml")
}

type initCommand struct {
	noExamples bool
}

func trimPathPrefix(p1, p2 string) string {
	if internal.HasFilepathPrefix(p1, p2) {
		return p1[len(p2):]
	}
	return p1
}

func (cmd *initCommand) Run(ctx *dep.Ctx, args []string) error {
	if len(args) > 1 {
		return errors.Errorf("too many args (%d)", len(args))
	}

	var root string
	if len(args) <= 0 {
		wd, err := os.Getwd()
		if err != nil {
			return err
		}
		root = wd
	} else {
		root = args[0]
	}

	mf := filepath.Join(root, dep.ManifestName)
	lf := filepath.Join(root, dep.LockName)

	mok, err := dep.IsRegular(mf)
	if err != nil {
		return err
	}
	if mok {
		return errors.Errorf("manifest already exists: %s", mf)
	}
	// Manifest file does not exist.

	lok, err := dep.IsRegular(lf)
	if err != nil {
		return err
	}
	if lok {
		return errors.Errorf("invalid state: manifest %q does not exist, but lock %q does", mf, lf)
	}

	cpr, err := ctx.SplitAbsoluteProjectRoot(root)
	if err != nil {
		return errors.Wrap(err, "determineProjectRoot")
	}
	internal.Vlogf("Finding dependencies for %q...", cpr)
	pkgT, err := pkgtree.ListPackages(root, cpr)
	if err != nil {
		return errors.Wrap(err, "gps.ListPackages")
	}
	internal.Vlogf("Found %d dependencies.", len(pkgT.Packages))
	sm, err := ctx.SourceManager()
	if err != nil {
		return errors.Wrap(err, "getSourceManager")
	}
	sm.UseDefaultSignalHandling()
	defer sm.Release()

	pd, err := getProjectData(ctx, pkgT, cpr, sm)
	if err != nil {
		return err
	}
	m := &dep.Manifest{
		Dependencies: pd.constraints,
	}

	// Make an initial lock from what knowledge we've collected about the
	// versions on disk
	l := &dep.Lock{
		P: make([]gps.LockedProject, 0, len(pd.ondisk)),
	}

	for pr, v := range pd.ondisk {
		// That we have to chop off these path prefixes is a symptom of
		// a problem in gps itself
		pkgs := make([]string, 0, len(pd.dependencies[pr]))
		prslash := string(pr) + "/"
		for _, pkg := range pd.dependencies[pr] {
			if pkg == string(pr) {
				pkgs = append(pkgs, ".")
			} else {
				pkgs = append(pkgs, trimPathPrefix(pkg, prslash))
			}
		}

		l.P = append(l.P, gps.NewLockedProject(
			gps.ProjectIdentifier{ProjectRoot: pr}, v, pkgs),
		)
	}

	internal.Vlogf("Solving...")
	params := gps.SolveParameters{
		RootDir:         root,
		RootPackageTree: pkgT,
		Manifest:        m,
		Lock:            l,
		ProjectAnalyzer: dep.Analyzer{},
	}

	if *verbose {
		params.Trace = true
		params.TraceLogger = log.New(os.Stderr, "", 0)
	}
	s, err := gps.Prepare(params, sm)
	if err != nil {
		return errors.Wrap(err, "prepare solver")
	}

	soln, err := s.Solve()
	if err != nil {
		handleAllTheFailuresOfTheWorld(err)
		return err
	}
	l = dep.LockFromInterface(soln)

	internal.Vlogf("Writing manifest and lock files.")

	var sw dep.SafeWriter
	sw.Prepare(m, nil, l, dep.VendorAlways)
	if err := sw.Write(root, sm, cmd.noExamples); err != nil {
		return errors.Wrap(err, "safe write of manifest and lock")
	}

	return nil
}

// contains checks if a array of strings contains a value
func contains(a []string, b string) bool {
	for _, v := range a {
		if b == v {
			return true
		}
	}
	return false
}

// isStdLib reports whether $GOROOT/src/path should be considered
// part of the standard distribution. For historical reasons we allow people to add
// their own code to $GOROOT instead of using $GOPATH, but we assume that
// code will start with a domain name (dot in the first element).
// This was loving taken from src/cmd/go/pkg.go in Go's code (isStandardImportPath).
func isStdLib(path string) bool {
	i := strings.Index(path, "/")
	if i < 0 {
		i = len(path)
	}
	elem := path[:i]
	return !strings.Contains(elem, ".")
}

// TODO solve failures can be really creative - we need to be similarly creative
// in handling them and informing the user appropriately
func handleAllTheFailuresOfTheWorld(err error) {
}

func hasImportPathPrefix(s, prefix string) bool {
	if s == prefix {
		return true
	}
	return strings.HasPrefix(s, prefix+"/")
}

type projectData struct {
	constraints  gps.ProjectConstraints          // constraints that could be found
	dependencies map[gps.ProjectRoot][]string    // all dependencies (imports) found by project root
	notondisk    map[gps.ProjectRoot]bool        // projects that were not found on disk
	ondisk       map[gps.ProjectRoot]gps.Version // projects that were found on disk
}

func getProjectData(ctx *dep.Ctx, pkgT pkgtree.PackageTree, cpr string, sm *gps.SourceMgr) (projectData, error) {
	constraints := make(gps.ProjectConstraints)
	dependencies := make(map[gps.ProjectRoot][]string)
	packages := make(map[string]bool)
	notondisk := make(map[gps.ProjectRoot]bool)
	ondisk := make(map[gps.ProjectRoot]gps.Version)

	syncDep := func(pr gps.ProjectRoot, sm *gps.SourceMgr) {
		message := "Cached"
		if err := sm.SyncSourceFor(gps.ProjectIdentifier{ProjectRoot: pr}); err != nil {
			message = "Unable to cache"
		}
		fmt.Fprintf(os.Stderr, "%s %s\n", message, pr)
	}

	rm, _ := pkgT.ToReachMap(true, true, false, nil)
	if len(rm) == 0 {
		return projectData{}, nil
	}

	internal.Vlogf("Building dependency graph...")
	// Exclude stdlib imports from the list returned from Flatten().
	const omitStdlib = false
	for _, ip := range rm.Flatten(omitStdlib) {
		pr, err := sm.DeduceProjectRoot(ip)
		if err != nil {
			return projectData{}, errors.Wrap(err, "sm.DeduceProjectRoot") // TODO: Skip and report ?
		}

		packages[ip] = true
		if _, has := dependencies[pr]; has {
			dependencies[pr] = append(dependencies[pr], ip)
			continue
		}
		go syncDep(pr, sm)

		internal.Vlogf("Found import of %q, analyzing...", ip)

		dependencies[pr] = []string{ip}
		v, err := ctx.VersionInWorkspace(pr)
		if err != nil {
			notondisk[pr] = true
			internal.Vlogf("Could not determine version for %q, assuming version to be master", pr)
			v = gps.Revision("master")
		}

		ondisk[pr] = v
		pp := gps.ProjectProperties{}
		switch v.Type() {
		case gps.IsBranch, gps.IsVersion, gps.IsRevision:
			pp.Constraint = v
		case gps.IsSemver:
			c, _ := gps.NewSemverConstraint("^" + v.String())
			pp.Constraint = c
		}

		constraints[pr] = pp
	}

	internal.Vlogf("Analyzing transitive imports...")
	// Explore the packages we've found for transitive deps, either
	// completing the lock or identifying (more) missing projects that we'll
	// need to ask gps to solve for us.
	colors := make(map[string]uint8)
	const (
		white uint8 = iota
		grey
		black
	)

	// cache of PackageTrees, so we don't parse projects more than once
	ptrees := make(map[gps.ProjectRoot]pkgtree.PackageTree)

	// depth-first traverser
	var dft func(string) error
	dft = func(pkg string) error {
		switch colors[pkg] {
		case white:
			internal.Vlogf("Analyzing %q...", pkg)
			colors[pkg] = grey

			pr, err := sm.DeduceProjectRoot(pkg)
			if err != nil {
				return errors.Wrap(err, "could not deduce project root for "+pkg)
			}

			// We already visited this project root earlier via some other
			// pkg within it, and made the decision that it's not on disk.
			// Respect that decision, and pop the stack.
			if notondisk[pr] {
				colors[pkg] = black
				return nil
			}

			ptree, has := ptrees[pr]
			if !has {
				// It's fine if the root does not exist - it indicates that this
				// project is not present in the workspace, and so we need to
				// solve to deal with this dep.
				r := filepath.Join(ctx.GOPATH, "src", string(pr))
				fi, err := os.Stat(r)
				if os.IsNotExist(err) || !fi.IsDir() {
					colors[pkg] = black
					notondisk[pr] = true
					return nil
				}

				// We know the project is on disk; the question is whether we're
				// first seeing it here, in the transitive exploration, or if it
				// was found in the initial pass on direct imports. We know it's
				// the former if there's no entry for it in the ondisk map.
				if _, in := ondisk[pr]; !in {
					v, err := ctx.VersionInWorkspace(pr)
					if err != nil {
						// Even if we know it's on disk, errors are still
						// possible when trying to deduce version. If we
						// encounter such an error, just treat the project as
						// not being on disk; the solver will work it out.
						colors[pkg] = black
						notondisk[pr] = true
						return nil
					}
					ondisk[pr] = v
				}

				ptree, err = pkgtree.ListPackages(r, string(pr))
				if err != nil {
					// Any error here other than an a nonexistent dir (which
					// can't happen because we covered that case above) is
					// probably critical, so bail out.
					return errors.Wrap(err, "gps.ListPackages")
				}
				ptrees[pr] = ptree
			}

			// Get a reachmap that includes main pkgs (even though importing
			// them is an error, what we're checking right now is simply whether
			// there's a package with go code present on disk), and does not
			// backpropagate errors (again, because our only concern right now
			// is package existence).
			rm, errmap := ptree.ToReachMap(true, false, false, nil)
			reached, ok := rm[pkg]
			if !ok {
				colors[pkg] = black
				// not on disk...
				notondisk[pr] = true
				return nil
			}
			if _, ok := errmap[pkg]; ok {
				// The package is on disk, but contains some errors.
				colors[pkg] = black
				return nil
			}

			if deps, has := dependencies[pr]; has {
				if !contains(deps, pkg) {
					dependencies[pr] = append(deps, pkg)
				}
			} else {
				dependencies[pr] = []string{pkg}
				go syncDep(pr, sm)
			}

			// recurse
			for _, rpkg := range reached.External {
				if isStdLib(rpkg) {
					continue
				}

				err := dft(rpkg)
				if err != nil {
					// Bubble up any errors we encounter
					return err
				}
			}

			colors[pkg] = black
		case grey:
			return errors.Errorf("Import cycle detected on %s", pkg)
		}
		return nil
	}

	// run the depth-first traversal from the set of immediate external
	// package imports we found in the current project
	for pkg := range packages {
		err := dft(pkg)
		if err != nil {
			return projectData{}, err // already errors.Wrap()'d internally
		}
	}

	pd := projectData{
		constraints:  constraints,
		dependencies: dependencies,
		notondisk:    notondisk,
		ondisk:       ondisk,
	}
	return pd, nil
}
<commit_msg>Partially implements new init algorithm<commit_after>// Copyright 2016 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 (
	"flag"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"strings"

	"github.com/golang/dep"
	"github.com/golang/dep/internal"
	"github.com/pkg/errors"
	"github.com/sdboyer/gps"
	"github.com/sdboyer/gps/pkgtree"
)

const initShortHelp = `Initialize a new project with manifest and lock files`
const initLongHelp = `
Initialize the project at filepath root by parsing its dependencies and writing
manifest and lock files. If root isn't specified, use the current directory.

The version of each dependency will reflect the current state of the GOPATH. If
a dependency doesn't exist in the GOPATH, it won't be written to the manifest,
but it will be solved-for, and will appear in the lock.

Note: init may use the network to solve the dependency graph.

Note: init does NOT vendor dependencies at the moment. See dep ensure.
`

func (cmd *initCommand) Name() string      { return "init" }
func (cmd *initCommand) Args() string      { return "[root]" }
func (cmd *initCommand) ShortHelp() string { return initShortHelp }
func (cmd *initCommand) LongHelp() string  { return initLongHelp }
func (cmd *initCommand) Hidden() bool      { return false }

func (cmd *initCommand) Register(fs *flag.FlagSet) {
	fs.BoolVar(&cmd.noExamples, "no-examples", false, "don't include example in Gopkg.toml")
}

type initCommand struct {
	noExamples bool
}

func trimPathPrefix(p1, p2 string) string {
	if internal.HasFilepathPrefix(p1, p2) {
		return p1[len(p2):]
	}
	return p1
}

func (cmd *initCommand) Run(ctx *dep.Ctx, args []string) error {
	if len(args) > 1 {
		return errors.Errorf("too many args (%d)", len(args))
	}

	var root string
	if len(args) <= 0 {
		wd, err := os.Getwd()
		if err != nil {
			return err
		}
		root = wd
	} else {
		root = args[0]
	}

	mf := filepath.Join(root, dep.ManifestName)
	lf := filepath.Join(root, dep.LockName)

	mok, err := dep.IsRegular(mf)
	if err != nil {
		return err
	}
	if mok {
		return errors.Errorf("manifest already exists: %s", mf)
	}
	// Manifest file does not exist.

	lok, err := dep.IsRegular(lf)
	if err != nil {
		return err
	}
	if lok {
		return errors.Errorf("invalid state: manifest %q does not exist, but lock %q does", mf, lf)
	}

	cpr, err := ctx.SplitAbsoluteProjectRoot(root)
	if err != nil {
		return errors.Wrap(err, "determineProjectRoot")
	}
	internal.Vlogf("Finding dependencies for %q...", cpr)
	pkgT, err := pkgtree.ListPackages(root, cpr)
	if err != nil {
		return errors.Wrap(err, "gps.ListPackages")
	}
	internal.Vlogf("Found %d dependencies.", len(pkgT.Packages))
	sm, err := ctx.SourceManager()
	if err != nil {
		return errors.Wrap(err, "getSourceManager")
	}
	sm.UseDefaultSignalHandling()
	defer sm.Release()

	pd, err := getProjectData(ctx, pkgT, cpr, sm)
	if err != nil {
		return err
	}
	m := &dep.Manifest{
		Dependencies: pd.constraints,
	}

	// Make an initial lock from what knowledge we've collected about the
	// versions on disk
	l := &dep.Lock{
		P: make([]gps.LockedProject, 0, len(pd.ondisk)),
	}

	for pr, v := range pd.ondisk {
		// That we have to chop off these path prefixes is a symptom of
		// a problem in gps itself
		pkgs := make([]string, 0, len(pd.dependencies[pr]))
		prslash := string(pr) + "/"
		for _, pkg := range pd.dependencies[pr] {
			if pkg == string(pr) {
				pkgs = append(pkgs, ".")
			} else {
				pkgs = append(pkgs, trimPathPrefix(pkg, prslash))
			}
		}

		l.P = append(l.P, gps.NewLockedProject(
			gps.ProjectIdentifier{ProjectRoot: pr}, v, pkgs),
		)
	}

	// Run solver with project versions found on disk
	internal.Vlogf("Solving...")
	soln, err := getSolverSolution(root, pkgT, m, l, sm)
	if err != nil {
		handleAllTheFailuresOfTheWorld(err)
		return err
	}
	l = dep.LockFromInterface(soln)

	// Pick notondisk project constraints from solution and add to manifest
	for k, _ := range pd.notondisk {
		for _, x := range l.Projects() {
			if k == x.Ident().ProjectRoot {
				m.Dependencies[k] = getProjectPropertiesFromVersion(x.Version())
				break
			}
		}
	}

	// Run solver again with appropriate constraint solutions from previous run
	// to generate the final lock.
	soln, err = getSolverSolution(root, pkgT, m, l, sm)
	if err != nil {
		handleAllTheFailuresOfTheWorld(err)
		return err
	}
	l = dep.LockFromInterface(soln)

	internal.Vlogf("Writing manifest and lock files.")

	var sw dep.SafeWriter
	sw.Prepare(m, nil, l, dep.VendorAlways)
	if err := sw.Write(root, sm, cmd.noExamples); err != nil {
		return errors.Wrap(err, "safe write of manifest and lock")
	}

	return nil
}

// contains checks if a array of strings contains a value
func contains(a []string, b string) bool {
	for _, v := range a {
		if b == v {
			return true
		}
	}
	return false
}

// isStdLib reports whether $GOROOT/src/path should be considered
// part of the standard distribution. For historical reasons we allow people to add
// their own code to $GOROOT instead of using $GOPATH, but we assume that
// code will start with a domain name (dot in the first element).
// This was loving taken from src/cmd/go/pkg.go in Go's code (isStandardImportPath).
func isStdLib(path string) bool {
	i := strings.Index(path, "/")
	if i < 0 {
		i = len(path)
	}
	elem := path[:i]
	return !strings.Contains(elem, ".")
}

// TODO solve failures can be really creative - we need to be similarly creative
// in handling them and informing the user appropriately
func handleAllTheFailuresOfTheWorld(err error) {
}

func hasImportPathPrefix(s, prefix string) bool {
	if s == prefix {
		return true
	}
	return strings.HasPrefix(s, prefix+"/")
}

// getProjectPropertiesFromVersion takes a gps.Version and returns a proper
// gps.ProjectProperties with Constraint value based on the version type.
func getProjectPropertiesFromVersion(v gps.Version) gps.ProjectProperties {
	pp := gps.ProjectProperties{}
	switch v.Type() {
	case gps.IsBranch, gps.IsVersion, gps.IsRevision:
		pp.Constraint = v
	case gps.IsSemver:
		c, _ := gps.NewSemverConstraint("^" + v.String())
		pp.Constraint = c
	}

	return pp
}

// getSolverSolution runs gps solver and returns a solution.
func getSolverSolution(root string, pkgT pkgtree.PackageTree, m *dep.Manifest, l *dep.Lock, sm *gps.SourceMgr) (gps.Solution, error) {
	params := gps.SolveParameters{
		RootDir:         root,
		RootPackageTree: pkgT,
		Manifest:        m,
		Lock:            l,
		ProjectAnalyzer: dep.Analyzer{},
	}

	if *verbose {
		params.Trace = true
		params.TraceLogger = log.New(os.Stderr, "", 0)
	}
	s, err := gps.Prepare(params, sm)
	if err != nil {
		return nil, errors.Wrap(err, "prepare solver")
	}

	return s.Solve()
}

type projectData struct {
	constraints  gps.ProjectConstraints          // constraints that could be found
	dependencies map[gps.ProjectRoot][]string    // all dependencies (imports) found by project root
	notondisk    map[gps.ProjectRoot]bool        // projects that were not found on disk
	ondisk       map[gps.ProjectRoot]gps.Version // projects that were found on disk
}

func getProjectData(ctx *dep.Ctx, pkgT pkgtree.PackageTree, cpr string, sm *gps.SourceMgr) (projectData, error) {
	constraints := make(gps.ProjectConstraints)
	dependencies := make(map[gps.ProjectRoot][]string)
	packages := make(map[string]bool)
	notondisk := make(map[gps.ProjectRoot]bool)
	ondisk := make(map[gps.ProjectRoot]gps.Version)

	syncDep := func(pr gps.ProjectRoot, sm *gps.SourceMgr) {
		message := "Cached"
		if err := sm.SyncSourceFor(gps.ProjectIdentifier{ProjectRoot: pr}); err != nil {
			message = "Unable to cache"
		}
		fmt.Fprintf(os.Stderr, "%s %s\n", message, pr)
	}

	rm, _ := pkgT.ToReachMap(true, true, false, nil)
	if len(rm) == 0 {
		return projectData{}, nil
	}

	internal.Vlogf("Building dependency graph...")
	// Exclude stdlib imports from the list returned from Flatten().
	const omitStdlib = false
	for _, ip := range rm.Flatten(omitStdlib) {
		pr, err := sm.DeduceProjectRoot(ip)
		if err != nil {
			return projectData{}, errors.Wrap(err, "sm.DeduceProjectRoot") // TODO: Skip and report ?
		}

		packages[ip] = true
		if _, has := dependencies[pr]; has {
			dependencies[pr] = append(dependencies[pr], ip)
			continue
		}
		go syncDep(pr, sm)

		internal.Vlogf("Found import of %q, analyzing...", ip)

		dependencies[pr] = []string{ip}
		v, err := ctx.VersionInWorkspace(pr)
		if err != nil {
			notondisk[pr] = true
			internal.Vlogf("Could not determine version for %q, omitting from generated manifest", pr)
			continue
		}

		ondisk[pr] = v
		constraints[pr] = getProjectPropertiesFromVersion(v)
	}

	internal.Vlogf("Analyzing transitive imports...")
	// Explore the packages we've found for transitive deps, either
	// completing the lock or identifying (more) missing projects that we'll
	// need to ask gps to solve for us.
	colors := make(map[string]uint8)
	const (
		white uint8 = iota
		grey
		black
	)

	// cache of PackageTrees, so we don't parse projects more than once
	ptrees := make(map[gps.ProjectRoot]pkgtree.PackageTree)

	// depth-first traverser
	var dft func(string) error
	dft = func(pkg string) error {
		switch colors[pkg] {
		case white:
			internal.Vlogf("Analyzing %q...", pkg)
			colors[pkg] = grey

			pr, err := sm.DeduceProjectRoot(pkg)
			if err != nil {
				return errors.Wrap(err, "could not deduce project root for "+pkg)
			}

			// We already visited this project root earlier via some other
			// pkg within it, and made the decision that it's not on disk.
			// Respect that decision, and pop the stack.
			if notondisk[pr] {
				colors[pkg] = black
				return nil
			}

			ptree, has := ptrees[pr]
			if !has {
				// It's fine if the root does not exist - it indicates that this
				// project is not present in the workspace, and so we need to
				// solve to deal with this dep.
				r := filepath.Join(ctx.GOPATH, "src", string(pr))
				fi, err := os.Stat(r)
				if os.IsNotExist(err) || !fi.IsDir() {
					colors[pkg] = black
					notondisk[pr] = true
					return nil
				}

				// We know the project is on disk; the question is whether we're
				// first seeing it here, in the transitive exploration, or if it
				// was found in the initial pass on direct imports. We know it's
				// the former if there's no entry for it in the ondisk map.
				if _, in := ondisk[pr]; !in {
					v, err := ctx.VersionInWorkspace(pr)
					if err != nil {
						// Even if we know it's on disk, errors are still
						// possible when trying to deduce version. If we
						// encounter such an error, just treat the project as
						// not being on disk; the solver will work it out.
						colors[pkg] = black
						notondisk[pr] = true
						return nil
					}
					ondisk[pr] = v
				}

				ptree, err = pkgtree.ListPackages(r, string(pr))
				if err != nil {
					// Any error here other than an a nonexistent dir (which
					// can't happen because we covered that case above) is
					// probably critical, so bail out.
					return errors.Wrap(err, "gps.ListPackages")
				}
				ptrees[pr] = ptree
			}

			// Get a reachmap that includes main pkgs (even though importing
			// them is an error, what we're checking right now is simply whether
			// there's a package with go code present on disk), and does not
			// backpropagate errors (again, because our only concern right now
			// is package existence).
			rm, errmap := ptree.ToReachMap(true, false, false, nil)
			reached, ok := rm[pkg]
			if !ok {
				colors[pkg] = black
				// not on disk...
				notondisk[pr] = true
				return nil
			}
			if _, ok := errmap[pkg]; ok {
				// The package is on disk, but contains some errors.
				colors[pkg] = black
				return nil
			}

			if deps, has := dependencies[pr]; has {
				if !contains(deps, pkg) {
					dependencies[pr] = append(deps, pkg)
				}
			} else {
				dependencies[pr] = []string{pkg}
				go syncDep(pr, sm)
			}

			// recurse
			for _, rpkg := range reached.External {
				if isStdLib(rpkg) {
					continue
				}

				err := dft(rpkg)
				if err != nil {
					// Bubble up any errors we encounter
					return err
				}
			}

			colors[pkg] = black
		case grey:
			return errors.Errorf("Import cycle detected on %s", pkg)
		}
		return nil
	}

	// run the depth-first traversal from the set of immediate external
	// package imports we found in the current project
	for pkg := range packages {
		err := dft(pkg)
		if err != nil {
			return projectData{}, err // already errors.Wrap()'d internally
		}
	}

	pd := projectData{
		constraints:  constraints,
		dependencies: dependencies,
		notondisk:    notondisk,
		ondisk:       ondisk,
	}
	return pd, nil
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"os"
	"os/signal"
	"path"
	"path/filepath"
	"runtime/pprof"
	"strings"

	"github.com/akalin/gopar/par1"
	"github.com/akalin/gopar/par2"
	"github.com/akalin/gopar/rsec16"
)

type par1LogEncoderDelegate struct{}

func (par1LogEncoderDelegate) OnDataFileLoad(i, n int, path string, byteCount int, err error) {
	if err != nil {
		fmt.Printf("[%d/%d] Loading data file %q failed: %+v\n", i, n, path, err)
	} else {
		fmt.Printf("[%d/%d] Loaded data file %q (%d bytes)\n", i, n, path, byteCount)
	}
}

func (par1LogEncoderDelegate) OnVolumeFileWrite(i, n int, path string, dataByteCount, byteCount int, err error) {
	if err != nil {
		fmt.Printf("[%d/%d] Writing volume file %q failed: %+v\n", i, n, path, err)
	} else {
		fmt.Printf("[%d/%d] Wrote volume file %q (%d data bytes, %d bytes)\n", i, n, path, dataByteCount, byteCount)
	}
}

type par1LogDecoderDelegate struct{}

func (par1LogDecoderDelegate) OnHeaderLoad(headerInfo string) {
	fmt.Printf("Loaded header: %s\n", headerInfo)
}

func (par1LogDecoderDelegate) OnFileEntryLoad(i, n int, filename, entryInfo string) {
	fmt.Printf("[%d/%d] Loaded entry for %q: %s\n", i, n, filename, entryInfo)
}

func (par1LogDecoderDelegate) OnCommentLoad(comment []byte) {
	fmt.Printf("Comment: %q\n", comment)
}

func (par1LogDecoderDelegate) OnDataFileLoad(i, n int, path string, byteCount int, corrupt bool, err error) {
	if err != nil {
		if corrupt {
			fmt.Printf("[%d/%d] Loading data file %q failed; marking as corrupt and skipping: %+v\n", i, n, path, err)
		} else {
			fmt.Printf("[%d/%d] Loading data file %q failed: %+v\n", i, n, path, err)
		}
	} else {
		fmt.Printf("[%d/%d] Loaded data file %q (%d bytes)\n", i, n, path, byteCount)
	}
}

func (par1LogDecoderDelegate) OnDataFileWrite(i, n int, path string, byteCount int, err error) {
	if err != nil {
		fmt.Printf("[%d/%d] Writing data file %q failed: %+v\n", i, n, path, err)
	} else {
		fmt.Printf("[%d/%d] Wrote data file %q (%d bytes)\n", i, n, path, byteCount)
	}
}

func (par1LogDecoderDelegate) OnVolumeFileLoad(i uint64, path string, storedSetHash, computedSetHash [16]byte, dataByteCount int, err error) {
	if os.IsNotExist(err) {
		// Do nothing.
	} else if err != nil {
		fmt.Printf("[%d] Loading volume file %q failed: %+v\n", i, path, err)
	} else {
		fmt.Printf("[%d] Loaded volume file %q (%d data bytes)\n", i, path, dataByteCount)
		if storedSetHash != computedSetHash {
			fmt.Printf("[%d] Warning: stored set hash in %q %x doesn't match computed set hash %x\n", i, path, storedSetHash, computedSetHash)
		}
	}
}

type par2LogEncoderDelegate struct{}

func (par2LogEncoderDelegate) OnDataFileLoad(i, n int, path string, byteCount int, err error) {
	if err != nil {
		fmt.Printf("[%d/%d] Loading data file %q failed: %+v\n", i, n, path, err)
	} else {
		fmt.Printf("[%d/%d] Loaded data file %q (%d bytes)\n", i, n, path, byteCount)
	}
}

func (par2LogEncoderDelegate) OnIndexFileWrite(path string, byteCount int, err error) {
	if err != nil {
		fmt.Printf("Writing index file %q failed: %+v\n", path, err)
	} else {
		fmt.Printf("Wrote index file %q (%d bytes)\n", path, byteCount)
	}
}

func (par2LogEncoderDelegate) OnRecoveryFileWrite(start, count, total int, path string, dataByteCount, byteCount int, err error) {
	if err != nil {
		fmt.Printf("[%d+%d/%d] Writing recovery file %q failed: %+v\n", start, count, total, path, err)
	} else {
		fmt.Printf("[%d+%d/%d] Wrote recovery file %q (%d data bytes, %d bytes)\n", start, count, total, path, dataByteCount, byteCount)
	}
}

type par2LogDecoderDelegate struct{}

func (par2LogDecoderDelegate) OnCreatorPacketLoad(clientID string) {
	fmt.Printf("Loaded creator packet with client ID %q\n", clientID)
}

func (par2LogDecoderDelegate) OnMainPacketLoad(sliceByteCount, recoverySetCount, nonRecoverySetCount int) {
	fmt.Printf("Loaded main packet: slice byte count=%d, recovery set size=%d, non-recovery set size=%d\n", sliceByteCount, recoverySetCount, nonRecoverySetCount)
}

func (par2LogDecoderDelegate) OnFileDescriptionPacketLoad(fileID [16]byte, filename string, byteCount int) {
	fmt.Printf("Loaded file description packet for %q (ID=%x, %d bytes)\n", filename, fileID, byteCount)
}

func (par2LogDecoderDelegate) OnIFSCPacketLoad(fileID [16]byte) {
	fmt.Printf("Loaded checksums for file with ID %x\n", fileID)
}

func (par2LogDecoderDelegate) OnRecoveryPacketLoad(exponent uint16, byteCount int) {
	fmt.Printf("Loaded recovery packet: exponent=%d, byte count=%d\n", exponent, byteCount)
}

func (par2LogDecoderDelegate) OnUnknownPacketLoad(packetType [16]byte, byteCount int) {
	fmt.Printf("Loaded unknown packet of type %q and byte count %d\n", packetType, byteCount)
}

func (par2LogDecoderDelegate) OnOtherPacketSkip(setID [16]byte, packetType [16]byte, byteCount int) {
	fmt.Printf("Skipped packet with set ID %x of type %q and byte count %d\n", setID, packetType, byteCount)
}

func (par2LogDecoderDelegate) OnDataFileLoad(i, n int, path string, byteCount, hits, misses int, err error) {
	if err != nil {
		fmt.Printf("[%d/%d] Loading data file %q failed: %+v\n", i, n, path, err)
	} else {
		fmt.Printf("[%d/%d] Loaded data file %q (%d bytes, %d hits, %d misses)\n", i, n, path, byteCount, hits, misses)
	}
}

func (par2LogDecoderDelegate) OnParityFileLoad(i int, path string, err error) {
	if err != nil {
		fmt.Printf("[%d] Loading volume file %q failed: %+v\n", i, path, err)
	} else {
		fmt.Printf("[%d] Loaded volume file %q\n", i, path)
	}
}

func (par2LogDecoderDelegate) OnDetectCorruptDataChunk(fileID [16]byte, filename string, startByteOffset, endByteOffset int) {
	fmt.Printf("Corrupt data chunk: %q (ID %x), bytes %d to %d\n", filename, fileID, startByteOffset, endByteOffset-1)
}

func (par2LogDecoderDelegate) OnDetectDataFileHashMismatch(fileID [16]byte, filename string) {
	fmt.Printf("Hash mismatch for %q (ID %x)\n", filename, fileID)
}

func (par2LogDecoderDelegate) OnDetectDataFileWrongByteCount(fileID [16]byte, filename string) {
	fmt.Printf("Wrong byte count for %q (ID %x)\n", filename, fileID)
}

func (par2LogDecoderDelegate) OnDataFileWrite(i, n int, path string, byteCount int, err error) {
	if err != nil {
		fmt.Printf("[%d/%d] Writing data file %q failed: %+v\n", i, n, path, err)
	} else {
		fmt.Printf("[%d/%d] Wrote data file %q (%d bytes)\n", i, n, path, byteCount)
	}
}

func printUsageAndExit(name string, flagSet *flag.FlagSet) {
	name = filepath.Base(name)
	fmt.Printf(`
Usage:
  %s c(reate) [options] <PAR file> [files]
  %s v(erify) [options] <PAR file>
  %s r(epair) [options] <PAR file>

Options:
`, name, name, name)
	flagSet.PrintDefaults()
	fmt.Printf("\n")
	os.Exit(2)
}

type encoder interface {
	LoadFileData() error
	ComputeParityData() error
	Write(string) error
}

type decoder interface {
	LoadFileData() error
	LoadParityData() error
	Verify(checkParity bool) (bool, error)
	Repair(checkParity bool) ([]string, error)
}

func newEncoder(parFile string, filePaths []string, sliceByteCount, numParityShards, numGoroutines int) (encoder, error) {
	// TODO: Detect file type more robustly.
	ext := path.Ext(parFile)
	if ext == ".par2" {
		return par2.NewEncoder(par2LogEncoderDelegate{}, filePaths, sliceByteCount, numParityShards, numGoroutines)
	}
	return par1.NewEncoder(par1LogEncoderDelegate{}, filePaths, numParityShards)
}

func newDecoder(parFile string, numGoroutines int) (decoder, error) {
	// TODO: Detect file type more robustly.
	ext := path.Ext(parFile)
	if ext == ".par2" {
		return par2.NewDecoder(par2LogDecoderDelegate{}, parFile, numGoroutines)
	}
	return par1.NewDecoder(par1LogDecoderDelegate{}, parFile)
}

func main() {
	name := os.Args[0]
	flagSet := flag.NewFlagSet(name, flag.ExitOnError)
	flagSet.SetOutput(os.Stdout)
	usage := flagSet.Bool("h", false, "print usage info")
	cpuprofile := flagSet.String("cpuprofile", "", "if non-empty, where to write the CPU profile")
	checkParity := flagSet.Bool("checkparity", false, "check parity when verifying or repairing")
	sliceByteCount := flagSet.Int("s", 2000, "block size in bytes (must be a multiple of 4)")
	numParityShards := flagSet.Int("c", 3, "number of recovery blocks to create (or files, for PAR1)")
	// TODO: Detect hyperthreading and use only number of physical cores.
	numGoroutines := flagSet.Int("g", rsec16.DefaultNumGoroutines(), "number of goroutines to use for encoding/decoding PAR2")
	flagSet.Parse(os.Args[1:])

	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			panic(err)
		}
		defer f.Close()

		err = pprof.StartCPUProfile(f)
		if err != nil {
			panic(err)
		}

		c := make(chan os.Signal, 1)
		signal.Notify(c, os.Interrupt)
		go func() {
			<-c
			pprof.StopCPUProfile()
			os.Exit(1)
		}()

		defer pprof.StopCPUProfile()
	}

	if flagSet.NArg() < 2 || *usage {
		printUsageAndExit(name, flagSet)
	}

	cmd := flagSet.Arg(0)
	parFile := flagSet.Arg(1)

	switch strings.ToLower(cmd) {
	case "c":
		fallthrough
	case "create":
		if flagSet.NArg() == 2 {
			printUsageAndExit(name, flagSet)
		}

		encoder, err := newEncoder(parFile, flagSet.Args()[2:], *sliceByteCount, *numParityShards, *numGoroutines)
		if err != nil {
			panic(err)
		}

		err = encoder.LoadFileData()
		if err != nil {
			panic(err)
		}

		err = encoder.ComputeParityData()
		if err != nil {
			panic(err)
		}

		err = encoder.Write(parFile)
		if err != nil {
			fmt.Printf("Write parity error: %s", err)
			os.Exit(-1)
		}

	case "v":
		fallthrough
	case "verify":
		decoder, err := newDecoder(parFile, *numGoroutines)
		if err != nil {
			panic(err)
		}

		err = decoder.LoadFileData()
		if err != nil {
			panic(err)
		}

		err = decoder.LoadParityData()
		if err != nil {
			panic(err)
		}

		ok, err := decoder.Verify(*checkParity)
		if err != nil {
			panic(err)
		}

		fmt.Printf("Verify result: %t\n", ok)
		if !ok {
			os.Exit(-1)
		}

	case "r":
		fallthrough
	case "repair":
		decoder, err := newDecoder(parFile, *numGoroutines)
		if err != nil {
			panic(err)
		}

		err = decoder.LoadFileData()
		if err != nil {
			panic(err)
		}

		err = decoder.LoadParityData()
		if err != nil {
			panic(err)
		}

		repairedFiles, err := decoder.Repair(false)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Repair error: %s\n", err)
			os.Exit(-1)
		}

		fmt.Printf("Repaired files: %v\n", repairedFiles)
		if err != nil {
			os.Exit(-1)
		}

	default:
		printUsageAndExit(name, flagSet)
	}
}
<commit_msg>Fix usage statement to match actual behavior<commit_after>package main

import (
	"flag"
	"fmt"
	"os"
	"os/signal"
	"path"
	"path/filepath"
	"runtime/pprof"
	"strings"

	"github.com/akalin/gopar/par1"
	"github.com/akalin/gopar/par2"
	"github.com/akalin/gopar/rsec16"
)

type par1LogEncoderDelegate struct{}

func (par1LogEncoderDelegate) OnDataFileLoad(i, n int, path string, byteCount int, err error) {
	if err != nil {
		fmt.Printf("[%d/%d] Loading data file %q failed: %+v\n", i, n, path, err)
	} else {
		fmt.Printf("[%d/%d] Loaded data file %q (%d bytes)\n", i, n, path, byteCount)
	}
}

func (par1LogEncoderDelegate) OnVolumeFileWrite(i, n int, path string, dataByteCount, byteCount int, err error) {
	if err != nil {
		fmt.Printf("[%d/%d] Writing volume file %q failed: %+v\n", i, n, path, err)
	} else {
		fmt.Printf("[%d/%d] Wrote volume file %q (%d data bytes, %d bytes)\n", i, n, path, dataByteCount, byteCount)
	}
}

type par1LogDecoderDelegate struct{}

func (par1LogDecoderDelegate) OnHeaderLoad(headerInfo string) {
	fmt.Printf("Loaded header: %s\n", headerInfo)
}

func (par1LogDecoderDelegate) OnFileEntryLoad(i, n int, filename, entryInfo string) {
	fmt.Printf("[%d/%d] Loaded entry for %q: %s\n", i, n, filename, entryInfo)
}

func (par1LogDecoderDelegate) OnCommentLoad(comment []byte) {
	fmt.Printf("Comment: %q\n", comment)
}

func (par1LogDecoderDelegate) OnDataFileLoad(i, n int, path string, byteCount int, corrupt bool, err error) {
	if err != nil {
		if corrupt {
			fmt.Printf("[%d/%d] Loading data file %q failed; marking as corrupt and skipping: %+v\n", i, n, path, err)
		} else {
			fmt.Printf("[%d/%d] Loading data file %q failed: %+v\n", i, n, path, err)
		}
	} else {
		fmt.Printf("[%d/%d] Loaded data file %q (%d bytes)\n", i, n, path, byteCount)
	}
}

func (par1LogDecoderDelegate) OnDataFileWrite(i, n int, path string, byteCount int, err error) {
	if err != nil {
		fmt.Printf("[%d/%d] Writing data file %q failed: %+v\n", i, n, path, err)
	} else {
		fmt.Printf("[%d/%d] Wrote data file %q (%d bytes)\n", i, n, path, byteCount)
	}
}

func (par1LogDecoderDelegate) OnVolumeFileLoad(i uint64, path string, storedSetHash, computedSetHash [16]byte, dataByteCount int, err error) {
	if os.IsNotExist(err) {
		// Do nothing.
	} else if err != nil {
		fmt.Printf("[%d] Loading volume file %q failed: %+v\n", i, path, err)
	} else {
		fmt.Printf("[%d] Loaded volume file %q (%d data bytes)\n", i, path, dataByteCount)
		if storedSetHash != computedSetHash {
			fmt.Printf("[%d] Warning: stored set hash in %q %x doesn't match computed set hash %x\n", i, path, storedSetHash, computedSetHash)
		}
	}
}

type par2LogEncoderDelegate struct{}

func (par2LogEncoderDelegate) OnDataFileLoad(i, n int, path string, byteCount int, err error) {
	if err != nil {
		fmt.Printf("[%d/%d] Loading data file %q failed: %+v\n", i, n, path, err)
	} else {
		fmt.Printf("[%d/%d] Loaded data file %q (%d bytes)\n", i, n, path, byteCount)
	}
}

func (par2LogEncoderDelegate) OnIndexFileWrite(path string, byteCount int, err error) {
	if err != nil {
		fmt.Printf("Writing index file %q failed: %+v\n", path, err)
	} else {
		fmt.Printf("Wrote index file %q (%d bytes)\n", path, byteCount)
	}
}

func (par2LogEncoderDelegate) OnRecoveryFileWrite(start, count, total int, path string, dataByteCount, byteCount int, err error) {
	if err != nil {
		fmt.Printf("[%d+%d/%d] Writing recovery file %q failed: %+v\n", start, count, total, path, err)
	} else {
		fmt.Printf("[%d+%d/%d] Wrote recovery file %q (%d data bytes, %d bytes)\n", start, count, total, path, dataByteCount, byteCount)
	}
}

type par2LogDecoderDelegate struct{}

func (par2LogDecoderDelegate) OnCreatorPacketLoad(clientID string) {
	fmt.Printf("Loaded creator packet with client ID %q\n", clientID)
}

func (par2LogDecoderDelegate) OnMainPacketLoad(sliceByteCount, recoverySetCount, nonRecoverySetCount int) {
	fmt.Printf("Loaded main packet: slice byte count=%d, recovery set size=%d, non-recovery set size=%d\n", sliceByteCount, recoverySetCount, nonRecoverySetCount)
}

func (par2LogDecoderDelegate) OnFileDescriptionPacketLoad(fileID [16]byte, filename string, byteCount int) {
	fmt.Printf("Loaded file description packet for %q (ID=%x, %d bytes)\n", filename, fileID, byteCount)
}

func (par2LogDecoderDelegate) OnIFSCPacketLoad(fileID [16]byte) {
	fmt.Printf("Loaded checksums for file with ID %x\n", fileID)
}

func (par2LogDecoderDelegate) OnRecoveryPacketLoad(exponent uint16, byteCount int) {
	fmt.Printf("Loaded recovery packet: exponent=%d, byte count=%d\n", exponent, byteCount)
}

func (par2LogDecoderDelegate) OnUnknownPacketLoad(packetType [16]byte, byteCount int) {
	fmt.Printf("Loaded unknown packet of type %q and byte count %d\n", packetType, byteCount)
}

func (par2LogDecoderDelegate) OnOtherPacketSkip(setID [16]byte, packetType [16]byte, byteCount int) {
	fmt.Printf("Skipped packet with set ID %x of type %q and byte count %d\n", setID, packetType, byteCount)
}

func (par2LogDecoderDelegate) OnDataFileLoad(i, n int, path string, byteCount, hits, misses int, err error) {
	if err != nil {
		fmt.Printf("[%d/%d] Loading data file %q failed: %+v\n", i, n, path, err)
	} else {
		fmt.Printf("[%d/%d] Loaded data file %q (%d bytes, %d hits, %d misses)\n", i, n, path, byteCount, hits, misses)
	}
}

func (par2LogDecoderDelegate) OnParityFileLoad(i int, path string, err error) {
	if err != nil {
		fmt.Printf("[%d] Loading volume file %q failed: %+v\n", i, path, err)
	} else {
		fmt.Printf("[%d] Loaded volume file %q\n", i, path)
	}
}

func (par2LogDecoderDelegate) OnDetectCorruptDataChunk(fileID [16]byte, filename string, startByteOffset, endByteOffset int) {
	fmt.Printf("Corrupt data chunk: %q (ID %x), bytes %d to %d\n", filename, fileID, startByteOffset, endByteOffset-1)
}

func (par2LogDecoderDelegate) OnDetectDataFileHashMismatch(fileID [16]byte, filename string) {
	fmt.Printf("Hash mismatch for %q (ID %x)\n", filename, fileID)
}

func (par2LogDecoderDelegate) OnDetectDataFileWrongByteCount(fileID [16]byte, filename string) {
	fmt.Printf("Wrong byte count for %q (ID %x)\n", filename, fileID)
}

func (par2LogDecoderDelegate) OnDataFileWrite(i, n int, path string, byteCount int, err error) {
	if err != nil {
		fmt.Printf("[%d/%d] Writing data file %q failed: %+v\n", i, n, path, err)
	} else {
		fmt.Printf("[%d/%d] Wrote data file %q (%d bytes)\n", i, n, path, byteCount)
	}
}

func printUsageAndExit(name string, flagSet *flag.FlagSet) {
	name = filepath.Base(name)
	fmt.Printf(`
Usage:
  %s [options] c(reate) <PAR file> [files]
  %s [options] v(erify) <PAR file>
  %s [options] r(epair) <PAR file>

Options:
`, name, name, name)
	flagSet.PrintDefaults()
	fmt.Printf("\n")
	os.Exit(2)
}

type encoder interface {
	LoadFileData() error
	ComputeParityData() error
	Write(string) error
}

type decoder interface {
	LoadFileData() error
	LoadParityData() error
	Verify(checkParity bool) (bool, error)
	Repair(checkParity bool) ([]string, error)
}

func newEncoder(parFile string, filePaths []string, sliceByteCount, numParityShards, numGoroutines int) (encoder, error) {
	// TODO: Detect file type more robustly.
	ext := path.Ext(parFile)
	if ext == ".par2" {
		return par2.NewEncoder(par2LogEncoderDelegate{}, filePaths, sliceByteCount, numParityShards, numGoroutines)
	}
	return par1.NewEncoder(par1LogEncoderDelegate{}, filePaths, numParityShards)
}

func newDecoder(parFile string, numGoroutines int) (decoder, error) {
	// TODO: Detect file type more robustly.
	ext := path.Ext(parFile)
	if ext == ".par2" {
		return par2.NewDecoder(par2LogDecoderDelegate{}, parFile, numGoroutines)
	}
	return par1.NewDecoder(par1LogDecoderDelegate{}, parFile)
}

func main() {
	name := os.Args[0]
	flagSet := flag.NewFlagSet(name, flag.ExitOnError)
	flagSet.SetOutput(os.Stdout)
	usage := flagSet.Bool("h", false, "print usage info")
	cpuprofile := flagSet.String("cpuprofile", "", "if non-empty, where to write the CPU profile")
	checkParity := flagSet.Bool("checkparity", false, "check parity when verifying or repairing")
	sliceByteCount := flagSet.Int("s", 2000, "block size in bytes (must be a multiple of 4)")
	numParityShards := flagSet.Int("c", 3, "number of recovery blocks to create (or files, for PAR1)")
	// TODO: Detect hyperthreading and use only number of physical cores.
	numGoroutines := flagSet.Int("g", rsec16.DefaultNumGoroutines(), "number of goroutines to use for encoding/decoding PAR2")
	flagSet.Parse(os.Args[1:])

	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			panic(err)
		}
		defer f.Close()

		err = pprof.StartCPUProfile(f)
		if err != nil {
			panic(err)
		}

		c := make(chan os.Signal, 1)
		signal.Notify(c, os.Interrupt)
		go func() {
			<-c
			pprof.StopCPUProfile()
			os.Exit(1)
		}()

		defer pprof.StopCPUProfile()
	}

	if flagSet.NArg() < 2 || *usage {
		printUsageAndExit(name, flagSet)
	}

	cmd := flagSet.Arg(0)
	parFile := flagSet.Arg(1)

	switch strings.ToLower(cmd) {
	case "c":
		fallthrough
	case "create":
		if flagSet.NArg() == 2 {
			printUsageAndExit(name, flagSet)
		}

		encoder, err := newEncoder(parFile, flagSet.Args()[2:], *sliceByteCount, *numParityShards, *numGoroutines)
		if err != nil {
			panic(err)
		}

		err = encoder.LoadFileData()
		if err != nil {
			panic(err)
		}

		err = encoder.ComputeParityData()
		if err != nil {
			panic(err)
		}

		err = encoder.Write(parFile)
		if err != nil {
			fmt.Printf("Write parity error: %s", err)
			os.Exit(-1)
		}

	case "v":
		fallthrough
	case "verify":
		decoder, err := newDecoder(parFile, *numGoroutines)
		if err != nil {
			panic(err)
		}

		err = decoder.LoadFileData()
		if err != nil {
			panic(err)
		}

		err = decoder.LoadParityData()
		if err != nil {
			panic(err)
		}

		ok, err := decoder.Verify(*checkParity)
		if err != nil {
			panic(err)
		}

		fmt.Printf("Verify result: %t\n", ok)
		if !ok {
			os.Exit(-1)
		}

	case "r":
		fallthrough
	case "repair":
		decoder, err := newDecoder(parFile, *numGoroutines)
		if err != nil {
			panic(err)
		}

		err = decoder.LoadFileData()
		if err != nil {
			panic(err)
		}

		err = decoder.LoadParityData()
		if err != nil {
			panic(err)
		}

		repairedFiles, err := decoder.Repair(false)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Repair error: %s\n", err)
			os.Exit(-1)
		}

		fmt.Printf("Repaired files: %v\n", repairedFiles)
		if err != nil {
			os.Exit(-1)
		}

	default:
		printUsageAndExit(name, flagSet)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"bytes"
	"flag"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"os/exec"
	"runtime/debug"
	"strings"
	"sync"
	"time"

	"git.apache.org/thrift.git/lib/go/thrift"
	"github.com/daviddengcn/go-colortext"
	"github.com/korovkin/limiter"
	"github.com/korovkin/parallel"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
)

type logger struct {
	ticket   int
	hostname string
	isError  bool
	buf      *bytes.Buffer
	print    bool
}

var (
	loggerMutex     = new(sync.Mutex)
	loggerIndex     = int(0)
	loggerStartTime = time.Now()
	loggerHostname  = ""

	flag_verbose = flag.Bool(
		"v",
		false,
		"verbose level: 1")
)

var loggerColors = []ct.Color{
	ct.Green,
	ct.Cyan,
	ct.Magenta,
	ct.Yellow,
	ct.Blue,
	ct.Red,
}

func (l *logger) Write(p []byte) (int, error) {
	buf := bytes.NewBuffer(p)
	wrote := 0
	for {
		line, err := buf.ReadBytes('\n')
		if len(line) > 1 {
			now := time.Now().Format("15:01:02")
			s := string(line)
			ts := time.Since(loggerStartTime).String()
			e := "I"
			if l.isError {
				e = "E"
			}

			{
				loggerMutex.Lock()
				if l.print {
					ct.ChangeColor(loggerColors[l.ticket%len(loggerColors)], false, ct.None, false)
					fmt.Printf("[%-14s %s %s %03d %s] ", ts, l.hostname, now, l.ticket, e)
					ct.ResetColor()
					fmt.Print(s)
				}
				if l.buf != nil {
					l.buf.Write([]byte(s))
				}
				loggerMutex.Unlock()
			}

			wrote += len(line)
		}
		if err != nil {
			break
		}
	}
	if len(p) > 0 && p[len(p)-1] != '\n' {
		fmt.Println()
	}

	return len(p), nil
}

func newLogger(ticket int, collectLines bool) *logger {
	l := &logger{ticket: ticket, buf: nil}
	if collectLines {
		l.buf = &bytes.Buffer{}
	}
	l.print = true
	return l
}

func CheckFatal(e error) error {
	if e != nil {
		debug.PrintStack()
		log.Println("CHECK: ERROR:", e)
		panic(e)
	}
	return e
}

func CheckNotFatal(e error) error {
	if e != nil {
		debug.PrintStack()
		log.Println("CHECK: ERROR:", e, e.Error())
	}
	return e
}

func executeCommand(p *Parallel, ticket int, cmdLine string) (*parallel.Output, error) {
	p.StatNumCommandsStart.Inc()
	T_START := time.Now()
	var err error
	output := &parallel.Output{}
	loggerOut := newLogger(ticket, true)
	loggerOut.isError = false
	loggerErr := newLogger(ticket, true)
	loggerErr.isError = true

	defer func() {
		dt := time.Since(T_START)
		fmt.Fprintf(
			loggerOut,
			"execute: done: dt: "+dt.String()+"\n",
		)
		if err == nil {
			p.StatNumCommandsDone.Inc()
			p.StatCommandLatency.Observe(dt.Seconds())
		}
	}()

	// execute remotely:
	if len(p.Slaves) > 0 {
		slave := p.Slaves[ticket%len(p.Slaves)]

		var transport thrift.TTransport
		transport, err = thrift.NewTSocket(slave.Address)
		if err != nil {
			log.Fatalln("failed to dial slave:", err.Error())
		}

		transport, err = p.transportFactory.GetTransport(transport)
		if err != nil {
			log.Fatalln("failed to GetTransport:", err.Error())
		}

		err = transport.Open()
		if err != nil {
			log.Fatalln("failed to open:", err.Error())
		}

		defer transport.Close()
		client := parallel.NewParallelClientFactory(transport, p.protocolFactory)

		cmd := parallel.Cmd{
			CmdLine: cmdLine,
			Ticket:  int64(ticket),
		}

		loggerOut.hostname = slave.Address
		fmt.Fprintf(loggerOut, "start: '"+cmdLine+"'\n")

		output, err = client.Execute(&cmd)
		if err != nil {
			log.Fatalln("failed to execute:", err.Error())
		}

		hostname := output.Tags["hostname"]
		loggerOut.hostname = hostname
		loggerErr.hostname = hostname

		if *flag_verbose {
			fmt.Fprintf(loggerOut,
				"execute: remote: host: %s stdout: [%s]\n",
				hostname,
				output.Stdout)

			fmt.Fprintf(loggerErr,
				"execute: remote: host: %s stderr: [%s]\n",
				hostname,
				output.Stderr)
		}

		return output, err
	}

	// execute locally:
	cs := []string{"/bin/sh", "-c", cmdLine}
	cmd := exec.Command(cs[0], cs[1:]...)
	cmd.Stdin = nil
	cmd.Stdout = loggerOut
	cmd.Stderr = loggerErr
	cmd.Env = append(
		os.Environ(),
		fmt.Sprintf("PARALLEL_TICKER=%d", ticket),
	)

	fmt.Fprintf(loggerOut, "start: '"+cmdLine+"'\n")

	loggerOut.print = *flag_verbose
	loggerErr.print = *flag_verbose

	err = cmd.Start()
	if err != nil {
		log.Fatalln("failed to start:", err)
		return output, err
	}

	if err == nil {
		err = cmd.Wait()
	}

	loggerOut.print = true
	loggerErr.print = true

	output.Tags = map[string]string{"hostname": loggerHostname}

	if loggerOut.buf != nil {
		output.Stdout = string(loggerOut.buf.Bytes())
	}

	if loggerErr.buf != nil {
		output.Stderr = string(loggerErr.buf.Bytes())
	}

	return output, err
}

type Slave struct {
	Address string
}

type Parallel struct {
	jobs   int
	logger *logger
	worker *limiter.ConcurrencyLimiter

	// master / slave
	protocolFactory  thrift.TProtocolFactory
	transportFactory thrift.TTransportFactory

	// master:
	Slaves []*Slave

	// slave:
	handler         *ParallelSlaveHandler
	serverTransport thrift.TServerTransport
	slaveAddress    string

	// stats:
	StatNumCommandsStart prometheus.Counter
	StatNumCommandsDone  prometheus.Counter
	StatCommandLatency   prometheus.Summary
}

func (p *Parallel) Close() {
	p.worker.Wait()
}

func mainMaster(p *Parallel) {
	var err error

	// connect to slaves:
	for _, slave := range p.Slaves {
		var transport thrift.TTransport
		transport, err = thrift.NewTSocket(slave.Address)
		if err != nil {
			log.Fatalln("failed to dial slave:", err.Error())
		}

		transport, err = p.transportFactory.GetTransport(transport)
		if err != nil {
			log.Fatalln("failed to open:", err.Error())
		}

		err = transport.Open()
		if err != nil {
			log.Fatalln("failed to open:", err.Error())
		}

		defer transport.Close()

		client := parallel.NewParallelClientFactory(transport, p.protocolFactory)
		ok, err := client.Ping()
		if err != nil {
			log.Fatalln("failed to ping client:", err.Error())
		}

		fmt.Fprintf(
			p.logger,
			"adding slave: %s ok: %s",
			slave.Address,
			ok)
	}

	r := bufio.NewReaderSize(os.Stdin, 1*1024*1024)
	fmt.Fprintf(p.logger, "reading from stdin...\n")
	for {
		line, err := r.ReadString('\n')
		if err == io.EOF {
			break
		}
		line = strings.TrimSpace(line)

		p.worker.ExecuteWithTicket(func(ticket int) {
			executeCommand(p, ticket, line)
		})
	}
}

type ParallelSlaveHandler struct {
	p *Parallel
}

func NewParallelSlaveHandler() *ParallelSlaveHandler {
	return &ParallelSlaveHandler{}
}

func (p *ParallelSlaveHandler) Execute(command *parallel.Cmd) (output *parallel.Output, err error) {
	output = nil
	err = nil
	output, err = executeCommand(p.p, int(command.Ticket), command.CmdLine)

	// TODO:: recover, handle panics

	return output, err
}

func (p *ParallelSlaveHandler) Ping() (r string, err error) {
	return "ping:ok", nil
}

func mainSlave(p *Parallel) {
	var err error

	p.serverTransport, err = thrift.NewTServerSocket(p.slaveAddress)

	if err != nil {
		log.Fatalln("failed to start server:", err.Error())
		return
	}

	p.handler = NewParallelSlaveHandler()
	p.handler.p = p

	server := thrift.NewTSimpleServer4(
		parallel.NewParallelProcessor(p.handler),
		p.serverTransport,
		p.transportFactory,
		p.protocolFactory)

	err = server.Serve()
	if err != nil {
		log.Fatalln("failed to run slave:", err.Error())
	}
}

func metricsServer(serverAddress string) {
	metricsHandler := promhttp.HandlerFor(prometheus.DefaultGatherer, promhttp.HandlerOpts{})
	http.HandleFunc("/metrics", func(c http.ResponseWriter, req *http.Request) {
		metricsHandler.ServeHTTP(c, req)
	})

	http.HandleFunc("/",
		func(c http.ResponseWriter, req *http.Request) {
			io.WriteString(c, fmt.Sprintf("parallel %s", time.Now().String()))
		})

	go http.ListenAndServe(serverAddress, nil)
}

func main() {
	T_START := time.Now()
	logger := newLogger(0, false)
	defer func() {
		fmt.Fprintf(logger, "all done: dt: "+time.Since(T_START).String()+"\n")
	}()

	flag_jobs := flag.Int(
		"j",
		2,
		"num of concurrent jobs")

	flag_slave := flag.Bool(
		"slave",
		false,
		"run as slave")

	flag_slaves := flag.String(
		"slaves",
		"",
		"CSV list of slave addresses")

	flag_slave_address := flag.String(
		"address",
		"localhost:9010",
		"slave address")

	flag_slave_metrics_address := flag.String(
		"metrics_address",
		"localhost:9011",
		"slave metric address")

	loggerHostname, _ = os.Hostname()

	flag.Parse()
	fmt.Fprintf(logger, "concurrency limit: %d", *flag_jobs)
	fmt.Fprintf(logger, "slaves: %s", *flag_slaves)

	p := Parallel{}
	p.jobs = *flag_jobs
	p.logger = logger
	p.worker = limiter.NewConcurrencyLimiter(p.jobs)
	p.slaveAddress = *flag_slave_address
	p.Slaves = []*Slave{}
	for _, slaveAddr := range strings.Split(*flag_slaves, ",") {
		if slaveAddr != "" {
			slave := Slave{Address: slaveAddr}
			p.Slaves = append(p.Slaves, &slave)
		}
	}

	defer p.Close()

	// thrift protocol
	p.protocolFactory = thrift.NewTBinaryProtocolFactoryDefault()

	// thrift transport
	p.transportFactory = thrift.NewTTransportFactory()
	p.transportFactory = thrift.NewTFramedTransportFactory(p.transportFactory)

	// stats:
	p.StatNumCommandsStart = prometheus.NewCounter(
		prometheus.CounterOpts{
			Name: "commands_num_start",
			Help: "num received"})
	err := prometheus.Register(p.StatNumCommandsStart)
	CheckFatal(err)

	p.StatNumCommandsDone = prometheus.NewCounter(
		prometheus.CounterOpts{
			Name: "commands_num_done",
			Help: "num received"})
	err = prometheus.Register(p.StatNumCommandsDone)
	CheckFatal(err)

	p.StatCommandLatency = prometheus.NewSummary(prometheus.SummaryOpts{
		Name: "commands_latency",
		Help: "commands latency stat",
	})
	err = prometheus.Register(p.StatCommandLatency)
	CheckFatal(err)

	// run the metrics server:
	go metricsServer(*flag_slave_metrics_address)

	if *flag_slave == false {
		loggerHostname = p.slaveAddress
		logger.hostname = loggerHostname

		fmt.Fprintf(logger, "running as master\n")
		mainMaster(&p)
	} else {
		loggerHostname = p.slaveAddress
		logger.hostname = loggerHostname

		fmt.Fprintf(logger, "running as slave on: %s\n", p.slaveAddress)
		mainSlave(&p)

		// stop the metrics thread as well:
		os.Exit(0)
	}
}
<commit_msg>check whether the metrics server could be started<commit_after>package main

import (
	"bufio"
	"bytes"
	"flag"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"os/exec"
	"runtime/debug"
	"strings"
	"sync"
	"time"

	"git.apache.org/thrift.git/lib/go/thrift"
	"github.com/daviddengcn/go-colortext"
	"github.com/korovkin/limiter"
	"github.com/korovkin/parallel"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
)

type logger struct {
	ticket   int
	hostname string
	isError  bool
	buf      *bytes.Buffer
	print    bool
}

var (
	loggerMutex     = new(sync.Mutex)
	loggerIndex     = int(0)
	loggerStartTime = time.Now()
	loggerHostname  = ""

	flag_verbose = flag.Bool(
		"v",
		false,
		"verbose level: 1")
)

var loggerColors = []ct.Color{
	ct.Green,
	ct.Cyan,
	ct.Magenta,
	ct.Yellow,
	ct.Blue,
	ct.Red,
}

func (l *logger) Write(p []byte) (int, error) {
	buf := bytes.NewBuffer(p)
	wrote := 0
	for {
		line, err := buf.ReadBytes('\n')
		if len(line) > 1 {
			now := time.Now().Format("15:01:02")
			s := string(line)
			ts := time.Since(loggerStartTime).String()
			e := "I"
			if l.isError {
				e = "E"
			}

			{
				loggerMutex.Lock()
				if l.print {
					ct.ChangeColor(loggerColors[l.ticket%len(loggerColors)], false, ct.None, false)
					fmt.Printf("[%-14s %s %s %03d %s] ", ts, l.hostname, now, l.ticket, e)
					ct.ResetColor()
					fmt.Print(s)
				}
				if l.buf != nil {
					l.buf.Write([]byte(s))
				}
				loggerMutex.Unlock()
			}

			wrote += len(line)
		}
		if err != nil {
			break
		}
	}
	if len(p) > 0 && p[len(p)-1] != '\n' {
		fmt.Println()
	}

	return len(p), nil
}

func newLogger(ticket int, collectLines bool) *logger {
	l := &logger{ticket: ticket, buf: nil}
	if collectLines {
		l.buf = &bytes.Buffer{}
	}
	l.print = true
	return l
}

func CheckFatal(e error) error {
	if e != nil {
		debug.PrintStack()
		log.Println("CHECK: ERROR:", e)
		panic(e)
	}
	return e
}

func CheckNotFatal(e error) error {
	if e != nil {
		debug.PrintStack()
		log.Println("CHECK: ERROR:", e, e.Error())
	}
	return e
}

func executeCommand(p *Parallel, ticket int, cmdLine string) (*parallel.Output, error) {
	p.StatNumCommandsStart.Inc()
	T_START := time.Now()
	var err error
	output := &parallel.Output{}
	loggerOut := newLogger(ticket, true)
	loggerOut.isError = false
	loggerErr := newLogger(ticket, true)
	loggerErr.isError = true

	defer func() {
		dt := time.Since(T_START)
		fmt.Fprintf(
			loggerOut,
			"execute: done: dt: "+dt.String()+"\n",
		)
		if err == nil {
			p.StatNumCommandsDone.Inc()
			p.StatCommandLatency.Observe(dt.Seconds())
		}
	}()

	// execute remotely:
	if len(p.Slaves) > 0 {
		slave := p.Slaves[ticket%len(p.Slaves)]

		var transport thrift.TTransport
		transport, err = thrift.NewTSocket(slave.Address)
		if err != nil {
			log.Fatalln("failed to dial slave:", err.Error())
		}

		transport, err = p.transportFactory.GetTransport(transport)
		if err != nil {
			log.Fatalln("failed to GetTransport:", err.Error())
		}

		err = transport.Open()
		if err != nil {
			log.Fatalln("failed to open:", err.Error())
		}

		defer transport.Close()
		client := parallel.NewParallelClientFactory(transport, p.protocolFactory)

		cmd := parallel.Cmd{
			CmdLine: cmdLine,
			Ticket:  int64(ticket),
		}

		loggerOut.hostname = slave.Address
		fmt.Fprintf(loggerOut, "start: '"+cmdLine+"'\n")

		output, err = client.Execute(&cmd)
		if err != nil {
			log.Fatalln("failed to execute:", err.Error())
		}

		hostname := output.Tags["hostname"]
		loggerOut.hostname = hostname
		loggerErr.hostname = hostname

		if *flag_verbose {
			fmt.Fprintf(loggerOut,
				"execute: remote: host: %s stdout: [%s]\n",
				hostname,
				output.Stdout)

			fmt.Fprintf(loggerErr,
				"execute: remote: host: %s stderr: [%s]\n",
				hostname,
				output.Stderr)
		}

		return output, err
	}

	// execute locally:
	cs := []string{"/bin/sh", "-c", cmdLine}
	cmd := exec.Command(cs[0], cs[1:]...)
	cmd.Stdin = nil
	cmd.Stdout = loggerOut
	cmd.Stderr = loggerErr
	cmd.Env = append(
		os.Environ(),
		fmt.Sprintf("PARALLEL_TICKER=%d", ticket),
	)

	fmt.Fprintf(loggerOut, "start: '"+cmdLine+"'\n")

	loggerOut.print = *flag_verbose
	loggerErr.print = *flag_verbose

	err = cmd.Start()
	if err != nil {
		log.Fatalln("failed to start:", err)
		return output, err
	}

	if err == nil {
		err = cmd.Wait()
	}

	loggerOut.print = true
	loggerErr.print = true

	output.Tags = map[string]string{"hostname": loggerHostname}

	if loggerOut.buf != nil {
		output.Stdout = string(loggerOut.buf.Bytes())
	}

	if loggerErr.buf != nil {
		output.Stderr = string(loggerErr.buf.Bytes())
	}

	return output, err
}

type Slave struct {
	Address string
}

type Parallel struct {
	jobs   int
	logger *logger
	worker *limiter.ConcurrencyLimiter

	// master / slave
	protocolFactory  thrift.TProtocolFactory
	transportFactory thrift.TTransportFactory

	// master:
	Slaves []*Slave

	// slave:
	handler         *ParallelSlaveHandler
	serverTransport thrift.TServerTransport
	slaveAddress    string

	// stats:
	StatNumCommandsStart prometheus.Counter
	StatNumCommandsDone  prometheus.Counter
	StatCommandLatency   prometheus.Summary
}

func (p *Parallel) Close() {
	p.worker.Wait()
}

func mainMaster(p *Parallel) {
	var err error
	log.SetFlags(log.Lmicroseconds | log.Ldate | log.Lshortfile)

	// connect to slaves:
	for _, slave := range p.Slaves {
		var transport thrift.TTransport
		transport, err = thrift.NewTSocket(slave.Address)
		if err != nil {
			log.Fatalln("failed to dial slave:", err.Error())
		}

		transport, err = p.transportFactory.GetTransport(transport)
		if err != nil {
			log.Fatalln("failed to open:", err.Error())
		}

		err = transport.Open()
		if err != nil {
			log.Fatalln("failed to open:", err.Error())
		}

		defer transport.Close()

		client := parallel.NewParallelClientFactory(transport, p.protocolFactory)
		ok, err := client.Ping()
		if err != nil {
			log.Fatalln("failed to ping client:", err.Error())
		}

		fmt.Fprintf(
			p.logger,
			"adding slave: %s ok: %s",
			slave.Address,
			ok)
	}

	r := bufio.NewReaderSize(os.Stdin, 1*1024*1024)
	fmt.Fprintf(p.logger, "reading from stdin...\n")
	for {
		line, err := r.ReadString('\n')
		if err == io.EOF {
			break
		}
		line = strings.TrimSpace(line)

		p.worker.ExecuteWithTicket(func(ticket int) {
			executeCommand(p, ticket, line)
		})
	}
}

type ParallelSlaveHandler struct {
	p *Parallel
}

func NewParallelSlaveHandler() *ParallelSlaveHandler {
	return &ParallelSlaveHandler{}
}

func (p *ParallelSlaveHandler) Execute(command *parallel.Cmd) (output *parallel.Output, err error) {
	output = nil
	err = nil
	output, err = executeCommand(p.p, int(command.Ticket), command.CmdLine)

	// TODO:: recover, handle panics

	return output, err
}

func (p *ParallelSlaveHandler) Ping() (r string, err error) {
	return "ping:ok", nil
}

func mainSlave(p *Parallel) {
	var err error

	p.serverTransport, err = thrift.NewTServerSocket(p.slaveAddress)

	if err != nil {
		log.Fatalln("failed to start server:", err.Error())
		return
	}

	p.handler = NewParallelSlaveHandler()
	p.handler.p = p

	server := thrift.NewTSimpleServer4(
		parallel.NewParallelProcessor(p.handler),
		p.serverTransport,
		p.transportFactory,
		p.protocolFactory)

	err = server.Serve()
	if err != nil {
		log.Fatalln("failed to run slave:", err.Error())
	}
}

func metricsServer(p *Parallel, serverAddress string) {
	metricsHandler := promhttp.HandlerFor(prometheus.DefaultGatherer, promhttp.HandlerOpts{})
	http.HandleFunc("/metrics", func(c http.ResponseWriter, req *http.Request) {
		metricsHandler.ServeHTTP(c, req)
	})

	http.HandleFunc("/",
		func(c http.ResponseWriter, req *http.Request) {
			io.WriteString(c,
				fmt.Sprintf(
					"go parallel time: %d slave_address: %s jobs: %d slaves: %d",
					time.Now().Unix(),
					serverAddress,
					p.jobs,
					len(p.Slaves)),
			)
		})

	err := http.ListenAndServe(serverAddress, nil)
	if err != nil {
		log.Println("WARNING: failed to start the metrics server on:", serverAddress, err.Error)
	}
}

func main() {
	T_START := time.Now()
	logger := newLogger(0, false)
	defer func() {
		fmt.Fprintf(logger, "all done: dt: "+time.Since(T_START).String()+"\n")
	}()

	flag_jobs := flag.Int(
		"j",
		2,
		"num of concurrent jobs")

	flag_slave := flag.Bool(
		"slave",
		false,
		"run as slave")

	flag_slaves := flag.String(
		"slaves",
		"",
		"CSV list of slave addresses")

	flag_slave_address := flag.String(
		"address",
		"localhost:9010",
		"slave address")

	flag_slave_metrics_address := flag.String(
		"metrics_address",
		"localhost:9011",
		"slave metric address")

	loggerHostname, _ = os.Hostname()

	flag.Parse()
	fmt.Fprintf(logger, "concurrency limit: %d", *flag_jobs)
	fmt.Fprintf(logger, "slaves: %s", *flag_slaves)

	p := &Parallel{}
	p.jobs = *flag_jobs
	p.logger = logger
	p.worker = limiter.NewConcurrencyLimiter(p.jobs)
	p.slaveAddress = *flag_slave_address
	p.Slaves = []*Slave{}
	for _, slaveAddr := range strings.Split(*flag_slaves, ",") {
		if slaveAddr != "" {
			slave := Slave{Address: slaveAddr}
			p.Slaves = append(p.Slaves, &slave)
		}
	}

	defer p.Close()

	// thrift protocol
	p.protocolFactory = thrift.NewTBinaryProtocolFactoryDefault()

	// thrift transport
	p.transportFactory = thrift.NewTTransportFactory()
	p.transportFactory = thrift.NewTFramedTransportFactory(p.transportFactory)

	// stats:
	p.StatNumCommandsStart = prometheus.NewCounter(
		prometheus.CounterOpts{
			Name: "commands_num_start",
			Help: "num received"})
	err := prometheus.Register(p.StatNumCommandsStart)
	CheckFatal(err)

	p.StatNumCommandsDone = prometheus.NewCounter(
		prometheus.CounterOpts{
			Name: "commands_num_done",
			Help: "num received"})
	err = prometheus.Register(p.StatNumCommandsDone)
	CheckFatal(err)

	p.StatCommandLatency = prometheus.NewSummary(prometheus.SummaryOpts{
		Name: "commands_latency",
		Help: "commands latency stat",
	})
	err = prometheus.Register(p.StatCommandLatency)
	CheckFatal(err)

	// run the metrics server:
	go metricsServer(p, *flag_slave_metrics_address)

	if *flag_slave == false {
		loggerHostname = p.slaveAddress
		logger.hostname = loggerHostname

		fmt.Fprintf(logger, "running as master\n")
		mainMaster(p)
	} else {
		loggerHostname = p.slaveAddress
		logger.hostname = loggerHostname

		fmt.Fprintf(logger, "running as slave on: %s\n", p.slaveAddress)
		mainSlave(p)

		// stop the metrics thread as well:
		os.Exit(0)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"flag"

	rin "github.com/fujiwara/Rin"
)

func main() {
	var (
		config string
	)
	flag.StringVar(&config, "config", "config.yaml", "config file path")
	flag.BoolVar(&rin.Debug, "debug", false, "enable debug logging")
	flag.Parse()
	err := rin.Run(config)
	if err != nil {
		panic(err)
	}
}
<commit_msg>add -version flag<commit_after>package main

import (
	"flag"
	"fmt"

	rin "github.com/fujiwara/Rin"
)

var (
	version   string
	buildDate string
)

func main() {
	var (
		config      string
		showVersion bool
	)
	flag.StringVar(&config, "config", "config.yaml", "config file path")
	flag.StringVar(&config, "c", "config.yaml", "config file path")
	flag.BoolVar(&rin.Debug, "debug", false, "enable debug logging")
	flag.BoolVar(&rin.Debug, "d", false, "enable debug logging")
	flag.BoolVar(&showVersion, "version", false, "show version")
	flag.BoolVar(&showVersion, "v", false, "show version")
	flag.Parse()

	if showVersion {
		fmt.Println("version:", version)
		fmt.Println("build:", buildDate)
		return
	}

	err := rin.Run(config)
	if err != nil {
		panic(err)
	}
}
<|endoftext|>


<commit_before>package shared

import "math/rand"

var charset = []byte("ABCDEFGHIJKLMNOPQRSTUVWXYZ234567")

func RandomKey() string {
	b := make([]byte, 32)
	for i := range b {
		b[i] = charset[rand.Intn(len(charset))]
	}
	return string(b)
}
<commit_msg>use lower case for key<commit_after>package shared

import "math/rand"

var charset = []byte("abcdefghijklmnopqrstuvwxyz234567")

func RandomKey() string {
	b := make([]byte, 32)
	for i := range b {
		b[i] = charset[rand.Intn(len(charset))]
	}
	return string(b)
}
<|endoftext|>


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

package main

import (
	"bytes"
	"crypto/rand"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"syscall"
	"time"
	"unsafe"

	// We use this ssh because it implements port redirection.
	// It can not, however, unpack password-protected keys yet.

	// TODO: get rid of krpty
	"github.com/u-root/u-root/pkg/termios"
	"github.com/u-root/u-root/pkg/ulog"

	// We use this ssh because it can unpack password-protected private keys.
	ossh "golang.org/x/crypto/ssh"
)

// a nonce is a [32]byte containing only printable characters, suitable for use as a string
type nonce [32]byte

var (
	// For the ssh server part
	bin         = flag.String("bin", "cpud", "path of cpu binary")
	debug       = flag.Bool("d", false, "enable debug prints")
	dbg9p       = flag.Bool("dbg9p", false, "show 9p io")
	dump        = flag.Bool("dump", false, "Dump copious output, including a 9p trace, to a temp file at exit")
	hostKeyFile = flag.String("hk", "" /*"/etc/ssh/ssh_host_rsa_key"*/, "file for host key")
	keyFile     = flag.String("key", filepath.Join(os.Getenv("HOME"), ".ssh/cpu_rsa"), "key file")
	mountopts   = flag.String("mountopts", "", "Extra options to add to the 9p mount")
	msize       = flag.Int("msize", 1048576, "msize to use")
	network     = flag.String("network", "tcp", "network to use")
	port        = flag.String("sp", "23", "cpu default port")
	port9p      = flag.String("port9p", "", "port9p # on remote machine for 9p mount")
	root        = flag.String("root", "/", "9p root")
	timeout9P   = flag.String("timeout9p", "100ms", "time to wait for the 9p mount to happen.")

	v          = func(string, ...interface{}) {}
	pid1       bool
	dumpWriter *os.File
)

func verbose(f string, a ...interface{}) {
	v("\r\n"+f+"\r\n", a...)
}

// generateNonce returns a nonce, or an error if random reader fails.
func generateNonce() (nonce, error) {
	var b [len(nonce{}) / 2]byte
	if _, err := rand.Read(b[:]); err != nil {
		return nonce{}, err
	}
	var n nonce
	copy(n[:], fmt.Sprintf("%02x", b))
	return n, nil
}

// String is a Stringer for nonce.
func (n nonce) String() string {
	return string(n[:])
}

func dial(n, a string, config *ossh.ClientConfig) (*ossh.Client, error) {
	client, err := ossh.Dial(n, a, config)
	if err != nil {
		return nil, fmt.Errorf("Failed to dial: %v", err)
	}
	return client, nil
}

func config(kf string) (*ossh.ClientConfig, error) {
	cb := ossh.InsecureIgnoreHostKey()
	//var hostKey ssh.PublicKey
	// A public key may be used to authenticate against the remote
	// server by using an unencrypted PEM-encoded private key file.
	//
	// If you have an encrypted private key, the crypto/x509 package
	// can be used to decrypt it.
	key, err := ioutil.ReadFile(kf)
	if err != nil {
		return nil, fmt.Errorf("unable to read private key %v: %v", kf, err)
	}

	// Create the Signer for this private key.
	signer, err := ossh.ParsePrivateKey(key)
	if err != nil {
		return nil, fmt.Errorf("ParsePrivateKey %v: %v", kf, err)
	}
	if *hostKeyFile != "" {
		hk, err := ioutil.ReadFile(*hostKeyFile)
		if err != nil {
			return nil, fmt.Errorf("unable to read host key %v: %v", *hostKeyFile, err)
		}
		pk, err := ossh.ParsePublicKey(hk)
		if err != nil {
			return nil, fmt.Errorf("host key %v: %v", string(hk), err)
		}
		cb = ossh.FixedHostKey(pk)
	}
	config := &ossh.ClientConfig{
		User: os.Getenv("USER"),
		Auth: []ossh.AuthMethod{
			// Use the PublicKeys method for remote authentication.
			ossh.PublicKeys(signer),
		},
		HostKeyCallback: cb,
	}
	return config, nil
}

func cmd(client *ossh.Client, s string) ([]byte, error) {
	session, err := client.NewSession()
	if err != nil {
		return nil, fmt.Errorf("Failed to create session: %v", err)
	}
	defer session.Close()

	var b bytes.Buffer
	session.Stdout = &b
	if err := session.Run(s); err != nil {
		return nil, fmt.Errorf("Failed to run %v: %v", s, err.Error())
	}
	return b.Bytes(), nil
}

// To make sure defer gets run and you tty is sane on exit
func runClient(host, a string) error {
	c, err := config(*keyFile)
	if err != nil {
		return err
	}
	cl, err := dial(*network, net.JoinHostPort(host, *port), c)
	if err != nil {
		return err
	}
	// Special case: maybe we don't want a namespace. If so, we don't need
	// to open up the socket.
	wantNameSpace := true
	if n, ok := os.LookupEnv("CPU_NAMESPACE"); ok && len(n) == 0 {
		wantNameSpace = false
	}

	var env []string
	cmd := fmt.Sprintf("%v -remote -bin %v", *bin, *bin)
	if wantNameSpace {
		// From setting up the forward to having the nonce written back to us,
		// we only allow 100ms. This is a lot, considering that at this point,
		// the sshd has forked a server for us and it's waiting to be
		// told what to do. We suggest that making the deadline a flag
		// would be a bad move, since people might be tempted to make it
		// large.
		deadline, err := time.ParseDuration(*timeout9P)
		if err != nil {
			return err
		}
		// Arrange port forwarding from remote ssh to our server.
		// Request the remote side to open port 5640 on all interfaces.
		// Note: cl.Listen returns a TCP listener with network is "tcp"
		// or variants. This lets us use a listen deadline.
		l, err := cl.Listen("tcp", "127.0.0.1:0")
		if err != nil {
			return fmt.Errorf("First cl.Listen %v", err)
		}
		ap := strings.Split(l.Addr().String(), ":")
		if len(ap) == 0 {
			return fmt.Errorf("Can't find a port number in %v", l.Addr().String())
		}
		port9p := ap[len(ap)-1]
		v("listener %T %v addr %v port %v", l, l, l.Addr().String(), port)

		nonce, err := generateNonce()
		if err != nil {
			log.Fatalf("Getting nonce: %v", err)
		}
		go srv(l, *root, nonce, deadline)
		cmd = fmt.Sprintf("%s -port9p %v", cmd, port9p)
		env = append(env, "CPUNONCE="+nonce.String())
	}
	cmd = fmt.Sprintf("%s %q", cmd, a)
	if err := shell(cl, cmd, env...); err != nil {
		return err
	}
	return nil
}

func env(s *ossh.Session, envs ...string) {
	for _, v := range append(os.Environ(), envs...) {
		env := strings.SplitN(v, "=", 2)
		if len(env) == 1 {
			env = append(env, "")
		}
		if err := s.Setenv(env[0], env[1]); err != nil {
			log.Printf("Warning: s.Setenv(%q, %q): %v", v, os.Getenv(v), err)
		}
	}
}

func stdin(s *ossh.Session, w io.WriteCloser, r io.Reader) {
	var newLine, tilde bool
	var t = []byte{'~'}
	var b [1]byte
	for {
		if _, err := r.Read(b[:]); err != nil {
			break
		}
		switch b[0] {
		default:
			newLine = false
			if tilde {
				if _, err := w.Write(t[:]); err != nil {
					return
				}
				tilde = false
			}
			if _, err := w.Write(b[:]); err != nil {
				return
			}
		case '\n', '\r':
			newLine = true
			if _, err := w.Write(b[:]); err != nil {
				return
			}
		case '~':
			if newLine {
				newLine = false
				tilde = true
				break
			}
			if _, err := w.Write(t[:]); err != nil {
				return
			}
		case '.':
			if tilde {
				s.Close()
				return
			}
			if _, err := w.Write(b[:]); err != nil {
				return
			}
		}
	}
}

func shell(client *ossh.Client, cmd string, envs ...string) error {
	t, err := termios.New()
	if err != nil {
		return err
	}
	r, err := t.Raw()
	if err != nil {
		return err
	}
	defer t.Set(r)
	if *bin == "" {
		if *bin, err = exec.LookPath("cpu"); err != nil {
			return err
		}
	}

	v("command is %q", cmd)
	session, err := client.NewSession()
	if err != nil {
		return err
	}
	defer session.Close()
	env(session, envs...)
	// Set up terminal modes
	modes := ossh.TerminalModes{
		ossh.ECHO:          0,     // disable echoing
		ossh.TTY_OP_ISPEED: 14400, // input speed = 14.4kbaud
		ossh.TTY_OP_OSPEED: 14400, // output speed = 14.4kbaud
	}
	// Request pseudo terminal
	if err := session.RequestPty("ansi", 40, 80, modes); err != nil {
		log.Fatal("request for pseudo terminal failed: ", err)
	}
	i, err := session.StdinPipe()
	if err != nil {
		return err
	}
	o, err := session.StdoutPipe()
	if err != nil {
		return err
	}
	e, err := session.StderrPipe()
	if err != nil {
		return err
	}

	v("Start remote with command %q", cmd)
	if err := session.Start(cmd); err != nil {
		return fmt.Errorf("Failed to run %v: %v", cmd, err.Error())
	}
	//env(session, "CPUNONCE="+n.String())
	go stdin(session, i, os.Stdin)
	go io.Copy(os.Stdout, o)
	go io.Copy(os.Stderr, e)
	return session.Wait()
}

// We do flag parsing in init so we can
// Unshare if needed while we are still
// single threaded.
func init() {
	flag.Parse()
	if *dump && *debug {
		log.Fatalf("You can only set either dump OR debug")
	}
	if *debug {
		v = log.Printf
	}
	if *dump {
		var err error
		dumpWriter, err = ioutil.TempFile("", "cpu")
		if err != nil {
			log.Fatal(err)
		}
		log.Printf("Logging to %s", dumpWriter.Name())
		*dbg9p = true
		ulog.Log = log.New(dumpWriter, "", log.Ltime|log.Lmicroseconds)
		v = ulog.Log.Printf
	}
}

func setWinsize(f *os.File, w, h int) {
	syscall.Syscall(syscall.SYS_IOCTL, f.Fd(), uintptr(syscall.TIOCSWINSZ),
		uintptr(unsafe.Pointer(&struct{ h, w, x, y uint16 }{uint16(h), uint16(w), 0, 0})))
}

// TODO: we've been tryinmg to figure out the right way to do usage for years.
// If this is a good way, it belongs in the uroot package.
func usage() {
	var b bytes.Buffer
	flag.CommandLine.SetOutput(&b)
	flag.PrintDefaults()
	log.Fatalf("Usage: cpu [options] host [shell command]:\n%v", b.String())
}

func main() {
	args := flag.Args()
	if len(args) == 0 {
		usage()
	}
	host := args[0]
	a := strings.Join(args[1:], " ")
	verbose("Running as client")
	if a == "" {
		a = os.Getenv("SHELL")
	}
	t, err := termios.GetTermios(0)
	if err != nil {
		log.Fatal("Getting Termios")
	}
	if err := runClient(host, a); err != nil {
		e := 1
		log.Printf("SSH error %s", err)
		if x, ok := err.(*ossh.ExitError); ok {
			e = x.ExitStatus()
		}
		defer os.Exit(e)
	}
	if err := termios.SetTermios(0, t); err != nil {
		// Never make this a log.Fatal, it might
		// interfere with the exit handling
		// for errors from the remote process.
		log.Print(err)
	}
}
<commit_msg>startup debug print: print the 9p mount debug value<commit_after>// Copyright 2018-2019 the u-root Authors. All rights reserved
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"bytes"
	"crypto/rand"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"syscall"
	"time"
	"unsafe"

	// We use this ssh because it implements port redirection.
	// It can not, however, unpack password-protected keys yet.

	// TODO: get rid of krpty
	"github.com/u-root/u-root/pkg/termios"
	"github.com/u-root/u-root/pkg/ulog"

	// We use this ssh because it can unpack password-protected private keys.
	ossh "golang.org/x/crypto/ssh"
)

// a nonce is a [32]byte containing only printable characters, suitable for use as a string
type nonce [32]byte

var (
	// For the ssh server part
	bin         = flag.String("bin", "cpud", "path of cpu binary")
	debug       = flag.Bool("d", false, "enable debug prints")
	dbg9p       = flag.Bool("dbg9p", false, "show 9p io")
	dump        = flag.Bool("dump", false, "Dump copious output, including a 9p trace, to a temp file at exit")
	hostKeyFile = flag.String("hk", "" /*"/etc/ssh/ssh_host_rsa_key"*/, "file for host key")
	keyFile     = flag.String("key", filepath.Join(os.Getenv("HOME"), ".ssh/cpu_rsa"), "key file")
	mountopts   = flag.String("mountopts", "", "Extra options to add to the 9p mount")
	msize       = flag.Int("msize", 1048576, "msize to use")
	network     = flag.String("network", "tcp", "network to use")
	port        = flag.String("sp", "23", "cpu default port")
	port9p      = flag.String("port9p", "", "port9p # on remote machine for 9p mount")
	root        = flag.String("root", "/", "9p root")
	timeout9P   = flag.String("timeout9p", "100ms", "time to wait for the 9p mount to happen.")

	v          = func(string, ...interface{}) {}
	pid1       bool
	dumpWriter *os.File
)

func verbose(f string, a ...interface{}) {
	v("\r\n"+f+"\r\n", a...)
}

// generateNonce returns a nonce, or an error if random reader fails.
func generateNonce() (nonce, error) {
	var b [len(nonce{}) / 2]byte
	if _, err := rand.Read(b[:]); err != nil {
		return nonce{}, err
	}
	var n nonce
	copy(n[:], fmt.Sprintf("%02x", b))
	return n, nil
}

// String is a Stringer for nonce.
func (n nonce) String() string {
	return string(n[:])
}

func dial(n, a string, config *ossh.ClientConfig) (*ossh.Client, error) {
	client, err := ossh.Dial(n, a, config)
	if err != nil {
		return nil, fmt.Errorf("Failed to dial: %v", err)
	}
	return client, nil
}

func config(kf string) (*ossh.ClientConfig, error) {
	cb := ossh.InsecureIgnoreHostKey()
	//var hostKey ssh.PublicKey
	// A public key may be used to authenticate against the remote
	// server by using an unencrypted PEM-encoded private key file.
	//
	// If you have an encrypted private key, the crypto/x509 package
	// can be used to decrypt it.
	key, err := ioutil.ReadFile(kf)
	if err != nil {
		return nil, fmt.Errorf("unable to read private key %v: %v", kf, err)
	}

	// Create the Signer for this private key.
	signer, err := ossh.ParsePrivateKey(key)
	if err != nil {
		return nil, fmt.Errorf("ParsePrivateKey %v: %v", kf, err)
	}
	if *hostKeyFile != "" {
		hk, err := ioutil.ReadFile(*hostKeyFile)
		if err != nil {
			return nil, fmt.Errorf("unable to read host key %v: %v", *hostKeyFile, err)
		}
		pk, err := ossh.ParsePublicKey(hk)
		if err != nil {
			return nil, fmt.Errorf("host key %v: %v", string(hk), err)
		}
		cb = ossh.FixedHostKey(pk)
	}
	config := &ossh.ClientConfig{
		User: os.Getenv("USER"),
		Auth: []ossh.AuthMethod{
			// Use the PublicKeys method for remote authentication.
			ossh.PublicKeys(signer),
		},
		HostKeyCallback: cb,
	}
	return config, nil
}

func cmd(client *ossh.Client, s string) ([]byte, error) {
	session, err := client.NewSession()
	if err != nil {
		return nil, fmt.Errorf("Failed to create session: %v", err)
	}
	defer session.Close()

	var b bytes.Buffer
	session.Stdout = &b
	if err := session.Run(s); err != nil {
		return nil, fmt.Errorf("Failed to run %v: %v", s, err.Error())
	}
	return b.Bytes(), nil
}

// To make sure defer gets run and you tty is sane on exit
func runClient(host, a string) error {
	c, err := config(*keyFile)
	if err != nil {
		return err
	}
	cl, err := dial(*network, net.JoinHostPort(host, *port), c)
	if err != nil {
		return err
	}
	// Special case: maybe we don't want a namespace. If so, we don't need
	// to open up the socket.
	wantNameSpace := true
	if n, ok := os.LookupEnv("CPU_NAMESPACE"); ok && len(n) == 0 {
		wantNameSpace = false
	}

	var env []string
	cmd := fmt.Sprintf("%v -remote -bin %v", *bin, *bin)
	if wantNameSpace {
		// From setting up the forward to having the nonce written back to us,
		// we only allow 100ms. This is a lot, considering that at this point,
		// the sshd has forked a server for us and it's waiting to be
		// told what to do. We suggest that making the deadline a flag
		// would be a bad move, since people might be tempted to make it
		// large.
		deadline, err := time.ParseDuration(*timeout9P)
		if err != nil {
			return err
		}
		// Arrange port forwarding from remote ssh to our server.
		// Request the remote side to open port 5640 on all interfaces.
		// Note: cl.Listen returns a TCP listener with network is "tcp"
		// or variants. This lets us use a listen deadline.
		l, err := cl.Listen("tcp", "127.0.0.1:0")
		if err != nil {
			return fmt.Errorf("First cl.Listen %v", err)
		}
		ap := strings.Split(l.Addr().String(), ":")
		if len(ap) == 0 {
			return fmt.Errorf("Can't find a port number in %v", l.Addr().String())
		}
		port9p := ap[len(ap)-1]
		v("listener %T %v addr %v port %v", l, l, l.Addr().String(), port)

		nonce, err := generateNonce()
		if err != nil {
			log.Fatalf("Getting nonce: %v", err)
		}
		go srv(l, *root, nonce, deadline)
		cmd = fmt.Sprintf("%s -port9p %v", cmd, port9p)
		env = append(env, "CPUNONCE="+nonce.String())
	}
	cmd = fmt.Sprintf("%s %q", cmd, a)
	if err := shell(cl, cmd, env...); err != nil {
		return err
	}
	return nil
}

func env(s *ossh.Session, envs ...string) {
	for _, v := range append(os.Environ(), envs...) {
		env := strings.SplitN(v, "=", 2)
		if len(env) == 1 {
			env = append(env, "")
		}
		if err := s.Setenv(env[0], env[1]); err != nil {
			log.Printf("Warning: s.Setenv(%q, %q): %v", v, os.Getenv(v), err)
		}
	}
}

func stdin(s *ossh.Session, w io.WriteCloser, r io.Reader) {
	var newLine, tilde bool
	var t = []byte{'~'}
	var b [1]byte
	for {
		if _, err := r.Read(b[:]); err != nil {
			break
		}
		switch b[0] {
		default:
			newLine = false
			if tilde {
				if _, err := w.Write(t[:]); err != nil {
					return
				}
				tilde = false
			}
			if _, err := w.Write(b[:]); err != nil {
				return
			}
		case '\n', '\r':
			newLine = true
			if _, err := w.Write(b[:]); err != nil {
				return
			}
		case '~':
			if newLine {
				newLine = false
				tilde = true
				break
			}
			if _, err := w.Write(t[:]); err != nil {
				return
			}
		case '.':
			if tilde {
				s.Close()
				return
			}
			if _, err := w.Write(b[:]); err != nil {
				return
			}
		}
	}
}

func shell(client *ossh.Client, cmd string, envs ...string) error {
	t, err := termios.New()
	if err != nil {
		return err
	}
	r, err := t.Raw()
	if err != nil {
		return err
	}
	defer t.Set(r)
	if *bin == "" {
		if *bin, err = exec.LookPath("cpu"); err != nil {
			return err
		}
	}

	v("command is %q", cmd)
	session, err := client.NewSession()
	if err != nil {
		return err
	}
	defer session.Close()
	env(session, envs...)
	// Set up terminal modes
	modes := ossh.TerminalModes{
		ossh.ECHO:          0,     // disable echoing
		ossh.TTY_OP_ISPEED: 14400, // input speed = 14.4kbaud
		ossh.TTY_OP_OSPEED: 14400, // output speed = 14.4kbaud
	}
	// Request pseudo terminal
	if err := session.RequestPty("ansi", 40, 80, modes); err != nil {
		log.Fatal("request for pseudo terminal failed: ", err)
	}
	i, err := session.StdinPipe()
	if err != nil {
		return err
	}
	o, err := session.StdoutPipe()
	if err != nil {
		return err
	}
	e, err := session.StderrPipe()
	if err != nil {
		return err
	}

	v("Start remote with command %q", cmd)
	if err := session.Start(cmd); err != nil {
		return fmt.Errorf("Failed to run %v: %v", cmd, err.Error())
	}
	//env(session, "CPUNONCE="+n.String())
	go stdin(session, i, os.Stdin)
	go io.Copy(os.Stdout, o)
	go io.Copy(os.Stderr, e)
	return session.Wait()
}

// We do flag parsing in init so we can
// Unshare if needed while we are still
// single threaded.
func init() {
	flag.Parse()
	if *dump && *debug {
		log.Fatalf("You can only set either dump OR debug")
	}
	if *debug {
		v = log.Printf
	}
	if *dump {
		var err error
		dumpWriter, err = ioutil.TempFile("", "cpu")
		if err != nil {
			log.Fatal(err)
		}
		log.Printf("Logging to %s", dumpWriter.Name())
		*dbg9p = true
		ulog.Log = log.New(dumpWriter, "", log.Ltime|log.Lmicroseconds)
		v = ulog.Log.Printf
	}
}

func setWinsize(f *os.File, w, h int) {
	syscall.Syscall(syscall.SYS_IOCTL, f.Fd(), uintptr(syscall.TIOCSWINSZ),
		uintptr(unsafe.Pointer(&struct{ h, w, x, y uint16 }{uint16(h), uint16(w), 0, 0})))
}

// TODO: we've been tryinmg to figure out the right way to do usage for years.
// If this is a good way, it belongs in the uroot package.
func usage() {
	var b bytes.Buffer
	flag.CommandLine.SetOutput(&b)
	flag.PrintDefaults()
	log.Fatalf("Usage: cpu [options] host [shell command]:\n%v", b.String())
}

func main() {
	args := flag.Args()
	if len(args) == 0 {
		usage()
	}
	host := args[0]
	a := strings.Join(args[1:], " ")
	verbose("Running as client with server 9p deadline of %v", *timeout9P)
	if a == "" {
		a = os.Getenv("SHELL")
	}
	t, err := termios.GetTermios(0)
	if err != nil {
		log.Fatal("Getting Termios")
	}
	if err := runClient(host, a); err != nil {
		e := 1
		log.Printf("SSH error %s", err)
		if x, ok := err.(*ossh.ExitError); ok {
			e = x.ExitStatus()
		}
		defer os.Exit(e)
	}
	if err := termios.SetTermios(0, t); err != nil {
		// Never make this a log.Fatal, it might
		// interfere with the exit handling
		// for errors from the remote process.
		log.Print(err)
	}
}
<|endoftext|>