<commit_before>package main

import (
	"bytes"
	"crypto/md5"
	"fmt"
	"go/ast"
	"go/format"
	"go/printer"
	"go/token"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"syscall"

	"github.com/jessevdk/go-flags"
	"github.com/zimmski/go-tool/importing"

	"github.com/zimmski/go-mutesting"
	"github.com/zimmski/go-mutesting/mutator"
	_ "github.com/zimmski/go-mutesting/mutator/branch"
	_ "github.com/zimmski/go-mutesting/mutator/expression"
	_ "github.com/zimmski/go-mutesting/mutator/statement"
)

const (
	returnOk = iota
	returnHelp
	returnBashCompletion
	returnError
)

const (
	execPassed  = 0
	execFailed  = 1
	execSkipped = 2
)

type options struct {
	General struct {
		Debug                bool `long:"debug" description:"Debug log output"`
		DoNotRemoveTmpFolder bool `long:"do-not-remove-tmp-folder" description:"Do not remove the tmp folder where all mutations are saved to"`
		Help                 bool `long:"help" description:"Show this help message"`
		Verbose              bool `long:"verbose" description:"Verbose log output"`
	} `group:"General options"`

	Files struct {
		Blacklist []string `long:"blacklist" description:"List of MD5 checksums of mutations which should be ignored. Each checksum must end with a new line character."`
		ListFiles bool     `long:"list-files" description:"List found files"`
		PrintAST  bool     `long:"print-ast" description:"Print the ASTs of all given files and exit"`
	} `group:"File options"`

	Mutator struct {
		DisableMutators []string `long:"disable" description:"Disable mutator by their name or using * as a suffix pattern"`
		ListMutators    bool     `long:"list-mutators" description:"List all available mutators"`
	} `group:"Mutator options"`

	Exec struct {
		Exec    string `long:"exec" description:"Execute this command for every mutation"`
		Timeout uint   `long:"exec-timeout" description:"Sets a timeout for the command execution (in seconds)" default:"10"`
	} `group:"Exec options"`

	Remaining struct {
		Targets []string `description:"Packages, directories and files even with patterns"`
	} `positional-args:"true" required:"true"`
}

func checkArguments(args []string, opts *options) (bool, int) {
	p := flags.NewNamedParser("go-mutesting", flags.None)

	p.ShortDescription = "Mutation testing for Go source code"

	if _, err := p.AddGroup("go-mutesting", "go-mutesting arguments", opts); err != nil {
		return true, exitError(err.Error())
	}

	completion := len(os.Getenv("GO_FLAGS_COMPLETION")) > 0

	_, err := p.ParseArgs(args)
	if (opts.General.Help || len(args) == 0) && !completion {
		p.WriteHelp(os.Stdout)

		return true, returnHelp
	} else if opts.Mutator.ListMutators {
		for _, name := range mutator.List() {
			fmt.Println(name)
		}

		return true, returnOk
	}

	if err != nil {
		return true, exitError(err.Error())
	}

	if completion {
		return true, returnBashCompletion
	}

	if opts.General.Debug {
		opts.General.Verbose = true
	}

	return false, 0
}

func debug(opts *options, format string, args ...interface{}) {
	if opts.General.Debug {
		fmt.Printf(format+"\n", args...)
	}
}

func verbose(opts *options, format string, args ...interface{}) {
	if opts.General.Verbose || opts.General.Debug {
		fmt.Printf(format+"\n", args...)
	}
}

func exitError(format string, args ...interface{}) int {
	fmt.Fprintf(os.Stderr, format+"\n", args...)

	return returnError
}

func mainCmd(args []string) int {
	var opts = new(options)
	var mutationBlackList = make(map[string]struct{})

	if exit, exitCode := checkArguments(args, opts); exit {
		return exitCode
	}

	files := importing.FilesOfArgs(opts.Remaining.Targets)
	if len(files) == 0 {
		return exitError("Could not find any suitable Go source files")
	}

	if opts.Files.ListFiles {
		for _, file := range files {
			fmt.Println(file)
		}

		return returnOk
	} else if opts.Files.PrintAST {
		for _, file := range files {
			fmt.Println(file)

			src, _, err := mutesting.ParseFile(file)
			if err != nil {
				return exitError("Could not open file %q: %v", file, err)
			}

			mutesting.PrintWalk(src)

			fmt.Println()
		}

		return returnOk
	}

	if len(opts.Files.Blacklist) > 0 {
		for _, f := range opts.Files.Blacklist {
			c, err := ioutil.ReadFile(f)
			if err != nil {
				return exitError("Cannot read blacklist file %q: %v", f, err)
			}

			for _, line := range strings.Split(string(c), "\n") {
				if line == "" {
					continue
				}

				if len(line) != 32 {
					return exitError("%q is not a MD5 checksum", line)
				}

				mutationBlackList[line] = struct{}{}
			}
		}
	}

	var mutators []mutator.Mutator

MUTATOR:
	for _, name := range mutator.List() {
		if len(opts.Mutator.DisableMutators) > 0 {
			for _, d := range opts.Mutator.DisableMutators {
				pattern := strings.HasSuffix(d, "*")

				if (pattern && strings.HasPrefix(name, d[:len(d)-2])) || (!pattern && name == d) {
					continue MUTATOR
				}
			}
		}

		debug(opts, "Enable mutator %q", name)

		m, _ := mutator.New(name)
		mutators = append(mutators, m)
	}

	tmpDir, err := ioutil.TempDir("", "go-mutesting-")
	if err != nil {
		panic(err)
	}
	debug(opts, "Save mutations into %q", tmpDir)

	var execs []string
	if opts.Exec.Exec != "" {
		execs = strings.Split(opts.Exec.Exec, " ")
	}

	passed := 0
	failed := 0
	skipped := 0

	for _, file := range files {
		debug(opts, "Mutate %q", file)

		src, fset, err := mutesting.ParseFile(file)
		if err != nil {
			return exitError("Could not open file %q: %v", file, err)
		}

		err = os.MkdirAll(tmpDir+"/"+filepath.Dir(file), 0755)
		if err != nil {
			panic(err)
		}

		tmpFile := tmpDir + "/" + file

		originalFile := fmt.Sprintf("%s.original", tmpFile)
		err = copyFile(file, originalFile)
		if err != nil {
			panic(err)
		}
		debug(opts, "Save original into %q", originalFile)

		i := 0

		for _, m := range mutators {
			debug(opts, "Mutator %s", m)

			changed := mutesting.MutateWalk(src, m)

			for {
				_, ok := <-changed

				if !ok {
					break
				}

				mutationFile := fmt.Sprintf("%s.%d", tmpFile, i)
				checksum, duplicate, err := saveAST(mutationBlackList, mutationFile, fset, src)
				if err != nil {
					panic(err)
				}
				if duplicate {
					debug(opts, "%q is a duplicate, we ignore it", mutationFile)
				} else {
					debug(opts, "Save mutation into %q with checksum %s", mutationFile, checksum)

					if len(execs) > 0 {
						debug(opts, "Execute %q for mutation", opts.Exec.Exec)

						execCommand := exec.Command(execs[0], execs[1:]...)

						execCommand.Stderr = os.Stderr
						execCommand.Stdout = os.Stdout

						execCommand.Env = append(os.Environ(), []string{
							"MUTATE_CHANGED=" + mutationFile,
							fmt.Sprintf("MUTATE_DEBUG=%t", opts.General.Debug),
							"MUTATE_ORIGINAL=" + file,
							fmt.Sprintf("MUTATE_TIMEOUT=%d", opts.Exec.Timeout),
							fmt.Sprintf("MUTATE_VERBOSE=%t", opts.General.Verbose),
						}...)

						err = execCommand.Start()
						if err != nil {
							panic(err)
						}

						// TODO timeout here

						err = execCommand.Wait()

						var execExitCode int
						if err == nil {
							execExitCode = 0
						} else if e, ok := err.(*exec.ExitError); ok {
							execExitCode = e.Sys().(syscall.WaitStatus).ExitStatus()
						} else {
							panic(err)
						}

						debug(opts, "Exited with %d", execExitCode)

						msg := fmt.Sprintf("%q with checksum %s", mutationFile, checksum)

						switch execExitCode {
						case 0:
							fmt.Printf("PASS %s\n", msg)

							passed++
						case 1:
							fmt.Printf("FAIL %s\n", msg)

							failed++
						case 2:
							fmt.Printf("SKIP %s\n", msg)

							skipped++
						default:
							fmt.Printf("UNKOWN exit code for %s\n", msg)
						}
					}
				}

				changed <- true

				// ignore original state
				<-changed
				changed <- true

				i++
			}
		}
	}

	if !opts.General.DoNotRemoveTmpFolder {
		err = os.RemoveAll(tmpDir)
		if err != nil {
			panic(err)
		}
		debug(opts, "Remove %q", tmpDir)
	}

	if len(execs) > 0 {
		fmt.Printf("The mutation score is %f (%d passed, %d failed, %d skipped, total is %d)\n", float64(passed)/float64(passed+failed), passed, failed, skipped, passed+failed+skipped)
	} else {
		fmt.Println("Cannot do a mutation testing summary since no exec command was given.")
	}

	return returnOk
}

func main() {
	os.Exit(mainCmd(os.Args[1:]))
}

func copyFile(src string, dst string) (err error) {
	s, err := os.Open(src)
	if err != nil {
		return err
	}
	defer func() {
		e := s.Close()
		if err == nil {
			err = e
		}
	}()

	d, err := os.Create(dst)
	if err != nil {
		return err
	}
	defer func() {
		e := d.Close()
		if err == nil {
			err = e
		}
	}()

	_, err = io.Copy(d, s)
	if err == nil {
		i, err := os.Stat(src)
		if err != nil {
			return err
		}

		return os.Chmod(dst, i.Mode())
	}

	return nil
}

func saveAST(mutationBlackList map[string]struct{}, file string, fset *token.FileSet, node ast.Node) (string, bool, error) {
	var buf bytes.Buffer

	h := md5.New()

	err := printer.Fprint(io.MultiWriter(h, &buf), fset, node)
	if err != nil {
		return "", false, err
	}

	checksum := fmt.Sprintf("%x", h.Sum(nil))

	if _, ok := mutationBlackList[checksum]; ok {
		return checksum, true, nil
	}

	mutationBlackList[checksum] = struct{}{}

	src, err := format.Source(buf.Bytes())
	if err != nil {
		return "", false, err
	}

	err = ioutil.WriteFile(file, src, 0666)
	if err != nil {
		return "", false, err
	}

	return checksum, false, nil
}
<commit_msg>Put copyFile in its own package so I will not be seduced to write it yet again<commit_after>package main

import (
	"bytes"
	"crypto/md5"
	"fmt"
	"go/ast"
	"go/format"
	"go/printer"
	"go/token"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"syscall"

	"github.com/jessevdk/go-flags"
	"github.com/zimmski/go-tool/importing"
	"github.com/zimmski/osutil"

	"github.com/zimmski/go-mutesting"
	"github.com/zimmski/go-mutesting/mutator"
	_ "github.com/zimmski/go-mutesting/mutator/branch"
	_ "github.com/zimmski/go-mutesting/mutator/expression"
	_ "github.com/zimmski/go-mutesting/mutator/statement"
)

const (
	returnOk = iota
	returnHelp
	returnBashCompletion
	returnError
)

const (
	execPassed  = 0
	execFailed  = 1
	execSkipped = 2
)

type options struct {
	General struct {
		Debug                bool `long:"debug" description:"Debug log output"`
		DoNotRemoveTmpFolder bool `long:"do-not-remove-tmp-folder" description:"Do not remove the tmp folder where all mutations are saved to"`
		Help                 bool `long:"help" description:"Show this help message"`
		Verbose              bool `long:"verbose" description:"Verbose log output"`
	} `group:"General options"`

	Files struct {
		Blacklist []string `long:"blacklist" description:"List of MD5 checksums of mutations which should be ignored. Each checksum must end with a new line character."`
		ListFiles bool     `long:"list-files" description:"List found files"`
		PrintAST  bool     `long:"print-ast" description:"Print the ASTs of all given files and exit"`
	} `group:"File options"`

	Mutator struct {
		DisableMutators []string `long:"disable" description:"Disable mutator by their name or using * as a suffix pattern"`
		ListMutators    bool     `long:"list-mutators" description:"List all available mutators"`
	} `group:"Mutator options"`

	Exec struct {
		Exec    string `long:"exec" description:"Execute this command for every mutation"`
		Timeout uint   `long:"exec-timeout" description:"Sets a timeout for the command execution (in seconds)" default:"10"`
	} `group:"Exec options"`

	Remaining struct {
		Targets []string `description:"Packages, directories and files even with patterns"`
	} `positional-args:"true" required:"true"`
}

func checkArguments(args []string, opts *options) (bool, int) {
	p := flags.NewNamedParser("go-mutesting", flags.None)

	p.ShortDescription = "Mutation testing for Go source code"

	if _, err := p.AddGroup("go-mutesting", "go-mutesting arguments", opts); err != nil {
		return true, exitError(err.Error())
	}

	completion := len(os.Getenv("GO_FLAGS_COMPLETION")) > 0

	_, err := p.ParseArgs(args)
	if (opts.General.Help || len(args) == 0) && !completion {
		p.WriteHelp(os.Stdout)

		return true, returnHelp
	} else if opts.Mutator.ListMutators {
		for _, name := range mutator.List() {
			fmt.Println(name)
		}

		return true, returnOk
	}

	if err != nil {
		return true, exitError(err.Error())
	}

	if completion {
		return true, returnBashCompletion
	}

	if opts.General.Debug {
		opts.General.Verbose = true
	}

	return false, 0
}

func debug(opts *options, format string, args ...interface{}) {
	if opts.General.Debug {
		fmt.Printf(format+"\n", args...)
	}
}

func verbose(opts *options, format string, args ...interface{}) {
	if opts.General.Verbose || opts.General.Debug {
		fmt.Printf(format+"\n", args...)
	}
}

func exitError(format string, args ...interface{}) int {
	fmt.Fprintf(os.Stderr, format+"\n", args...)

	return returnError
}

func mainCmd(args []string) int {
	var opts = new(options)
	var mutationBlackList = make(map[string]struct{})

	if exit, exitCode := checkArguments(args, opts); exit {
		return exitCode
	}

	files := importing.FilesOfArgs(opts.Remaining.Targets)
	if len(files) == 0 {
		return exitError("Could not find any suitable Go source files")
	}

	if opts.Files.ListFiles {
		for _, file := range files {
			fmt.Println(file)
		}

		return returnOk
	} else if opts.Files.PrintAST {
		for _, file := range files {
			fmt.Println(file)

			src, _, err := mutesting.ParseFile(file)
			if err != nil {
				return exitError("Could not open file %q: %v", file, err)
			}

			mutesting.PrintWalk(src)

			fmt.Println()
		}

		return returnOk
	}

	if len(opts.Files.Blacklist) > 0 {
		for _, f := range opts.Files.Blacklist {
			c, err := ioutil.ReadFile(f)
			if err != nil {
				return exitError("Cannot read blacklist file %q: %v", f, err)
			}

			for _, line := range strings.Split(string(c), "\n") {
				if line == "" {
					continue
				}

				if len(line) != 32 {
					return exitError("%q is not a MD5 checksum", line)
				}

				mutationBlackList[line] = struct{}{}
			}
		}
	}

	var mutators []mutator.Mutator

MUTATOR:
	for _, name := range mutator.List() {
		if len(opts.Mutator.DisableMutators) > 0 {
			for _, d := range opts.Mutator.DisableMutators {
				pattern := strings.HasSuffix(d, "*")

				if (pattern && strings.HasPrefix(name, d[:len(d)-2])) || (!pattern && name == d) {
					continue MUTATOR
				}
			}
		}

		debug(opts, "Enable mutator %q", name)

		m, _ := mutator.New(name)
		mutators = append(mutators, m)
	}

	tmpDir, err := ioutil.TempDir("", "go-mutesting-")
	if err != nil {
		panic(err)
	}
	debug(opts, "Save mutations into %q", tmpDir)

	var execs []string
	if opts.Exec.Exec != "" {
		execs = strings.Split(opts.Exec.Exec, " ")
	}

	passed := 0
	failed := 0
	skipped := 0

	for _, file := range files {
		debug(opts, "Mutate %q", file)

		src, fset, err := mutesting.ParseFile(file)
		if err != nil {
			return exitError("Could not open file %q: %v", file, err)
		}

		err = os.MkdirAll(tmpDir+"/"+filepath.Dir(file), 0755)
		if err != nil {
			panic(err)
		}

		tmpFile := tmpDir + "/" + file

		originalFile := fmt.Sprintf("%s.original", tmpFile)
		err = osutil.CopyFile(file, originalFile)
		if err != nil {
			panic(err)
		}
		debug(opts, "Save original into %q", originalFile)

		i := 0

		for _, m := range mutators {
			debug(opts, "Mutator %s", m)

			changed := mutesting.MutateWalk(src, m)

			for {
				_, ok := <-changed

				if !ok {
					break
				}

				mutationFile := fmt.Sprintf("%s.%d", tmpFile, i)
				checksum, duplicate, err := saveAST(mutationBlackList, mutationFile, fset, src)
				if err != nil {
					panic(err)
				}
				if duplicate {
					debug(opts, "%q is a duplicate, we ignore it", mutationFile)
				} else {
					debug(opts, "Save mutation into %q with checksum %s", mutationFile, checksum)

					if len(execs) > 0 {
						debug(opts, "Execute %q for mutation", opts.Exec.Exec)

						execCommand := exec.Command(execs[0], execs[1:]...)

						execCommand.Stderr = os.Stderr
						execCommand.Stdout = os.Stdout

						execCommand.Env = append(os.Environ(), []string{
							"MUTATE_CHANGED=" + mutationFile,
							fmt.Sprintf("MUTATE_DEBUG=%t", opts.General.Debug),
							"MUTATE_ORIGINAL=" + file,
							fmt.Sprintf("MUTATE_TIMEOUT=%d", opts.Exec.Timeout),
							fmt.Sprintf("MUTATE_VERBOSE=%t", opts.General.Verbose),
						}...)

						err = execCommand.Start()
						if err != nil {
							panic(err)
						}

						// TODO timeout here

						err = execCommand.Wait()

						var execExitCode int
						if err == nil {
							execExitCode = 0
						} else if e, ok := err.(*exec.ExitError); ok {
							execExitCode = e.Sys().(syscall.WaitStatus).ExitStatus()
						} else {
							panic(err)
						}

						debug(opts, "Exited with %d", execExitCode)

						msg := fmt.Sprintf("%q with checksum %s", mutationFile, checksum)

						switch execExitCode {
						case 0:
							fmt.Printf("PASS %s\n", msg)

							passed++
						case 1:
							fmt.Printf("FAIL %s\n", msg)

							failed++
						case 2:
							fmt.Printf("SKIP %s\n", msg)

							skipped++
						default:
							fmt.Printf("UNKOWN exit code for %s\n", msg)
						}
					}
				}

				changed <- true

				// ignore original state
				<-changed
				changed <- true

				i++
			}
		}
	}

	if !opts.General.DoNotRemoveTmpFolder {
		err = os.RemoveAll(tmpDir)
		if err != nil {
			panic(err)
		}
		debug(opts, "Remove %q", tmpDir)
	}

	if len(execs) > 0 {
		fmt.Printf("The mutation score is %f (%d passed, %d failed, %d skipped, total is %d)\n", float64(passed)/float64(passed+failed), passed, failed, skipped, passed+failed+skipped)
	} else {
		fmt.Println("Cannot do a mutation testing summary since no exec command was given.")
	}

	return returnOk
}

func main() {
	os.Exit(mainCmd(os.Args[1:]))
}

func saveAST(mutationBlackList map[string]struct{}, file string, fset *token.FileSet, node ast.Node) (string, bool, error) {
	var buf bytes.Buffer

	h := md5.New()

	err := printer.Fprint(io.MultiWriter(h, &buf), fset, node)
	if err != nil {
		return "", false, err
	}

	checksum := fmt.Sprintf("%x", h.Sum(nil))

	if _, ok := mutationBlackList[checksum]; ok {
		return checksum, true, nil
	}

	mutationBlackList[checksum] = struct{}{}

	src, err := format.Source(buf.Bytes())
	if err != nil {
		return "", false, err
	}

	err = ioutil.WriteFile(file, src, 0666)
	if err != nil {
		return "", false, err
	}

	return checksum, false, nil
}
<|endoftext|>


<commit_before>// Copyright 2012, 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package cmd_test

import (
	"fmt"

	. "launchpad.net/gocheck"
	"launchpad.net/juju-core/cmd"
	"launchpad.net/juju-core/log"
	"launchpad.net/juju-core/testing"
)

func initDefenestrate(args []string) (*cmd.SuperCommand, *TestCommand, error) {
	jc := cmd.NewSuperCommand(cmd.SuperCommandParams{Name: "jujutest"})
	tc := &TestCommand{Name: "defenestrate"}
	jc.Register(tc)
	return jc, tc, testing.InitCommand(jc, args)
}

type SuperCommandSuite struct{}

var _ = Suite(&SuperCommandSuite{})

const helpText = "\n    help\\s+- show help on a command or other topic"
const helpCommandsText = "commands:" + helpText

func (s *SuperCommandSuite) TestDispatch(c *C) {
	jc := cmd.NewSuperCommand(cmd.SuperCommandParams{Name: "jujutest"})
	info := jc.Info()
	c.Assert(info.Name, Equals, "jujutest")
	c.Assert(info.Args, Equals, "<command> ...")
	c.Assert(info.Doc, Matches, helpCommandsText)

	jc, _, err := initDefenestrate([]string{"discombobulate"})
	c.Assert(err, ErrorMatches, "unrecognized command: jujutest discombobulate")
	info = jc.Info()
	c.Assert(info.Name, Equals, "jujutest")
	c.Assert(info.Args, Equals, "<command> ...")
	c.Assert(info.Doc, Matches, "commands:\n    defenestrate - defenestrate the juju"+helpText)

	jc, tc, err := initDefenestrate([]string{"defenestrate"})
	c.Assert(err, IsNil)
	c.Assert(tc.Option, Equals, "")
	info = jc.Info()
	c.Assert(info.Name, Equals, "jujutest defenestrate")
	c.Assert(info.Args, Equals, "<something>")
	c.Assert(info.Doc, Equals, "defenestrate-doc")

	_, tc, err = initDefenestrate([]string{"defenestrate", "--option", "firmly"})
	c.Assert(err, IsNil)
	c.Assert(tc.Option, Equals, "firmly")

	_, tc, err = initDefenestrate([]string{"defenestrate", "gibberish"})
	c.Assert(err, ErrorMatches, `unrecognized args: \["gibberish"\]`)
}

func (s *SuperCommandSuite) TestRegister(c *C) {
	jc := cmd.NewSuperCommand(cmd.SuperCommandParams{Name: "jujutest"})
	jc.Register(&TestCommand{Name: "flip"})
	jc.Register(&TestCommand{Name: "flap"})
	badCall := func() { jc.Register(&TestCommand{Name: "flap"}) }
	c.Assert(badCall, PanicMatches, "command already registered: flap")
}

func (s *SuperCommandSuite) TestRegisterAlias(c *C) {
	jc := cmd.NewSuperCommand(cmd.SuperCommandParams{Name: "jujutest"})
	jc.Register(&TestCommand{Name: "flip", Aliases: []string{"flap", "flop"}})

	info := jc.Info()
	c.Assert(info.Doc, Equals, `commands:
    flap - alias for flip
    flip - flip the juju
    flop - alias for flip
    help - show help on a command or other topic`)
}

var commandsDoc = `commands:
    flapbabble - flapbabble the juju
    flip       - flip the juju`

func (s *SuperCommandSuite) TestInfo(c *C) {
	jc := cmd.NewSuperCommand(cmd.SuperCommandParams{
		Name:    "jujutest",
		Purpose: "to be purposeful",
		Doc:     "doc\nblah\ndoc",
	})
	info := jc.Info()
	c.Assert(info.Name, Equals, "jujutest")
	c.Assert(info.Purpose, Equals, "to be purposeful")
	// info doc starts with the jc.Doc and ends with the help command
	c.Assert(info.Doc, Matches, jc.Doc+"(.|\n)*")
	c.Assert(info.Doc, Matches, "(.|\n)*"+helpCommandsText)

	jc.Register(&TestCommand{Name: "flip"})
	jc.Register(&TestCommand{Name: "flapbabble"})
	info = jc.Info()
	c.Assert(info.Doc, Matches, jc.Doc+"\n\n"+commandsDoc+helpText)

	jc.Doc = ""
	info = jc.Info()
	c.Assert(info.Doc, Matches, commandsDoc+helpText)
}

func (s *SuperCommandSuite) TestLogging(c *C) {
	target, debug := log.Target(), log.Debug
	defer func() {
		log.SetTarget(target)
		log.Debug = debug
	}()
	jc := cmd.NewSuperCommand(cmd.SuperCommandParams{Name: "jujutest", Log: &cmd.Log{}})
	jc.Register(&TestCommand{Name: "blah"})
	ctx := testing.Context(c)
	code := cmd.Main(jc, ctx, []string{"blah", "--option", "error", "--debug"})
	c.Assert(code, Equals, 1)
	c.Assert(bufferString(ctx.Stderr), Matches, `^.* ERROR command failed: BAM!
error: BAM!
`)
}

func (s *SuperCommandSuite) TestMissingCallback(c *C) {
	var (
		calledName string
		calledArgs []string
	)

	callback := func(ctx *cmd.Context, subcommand string, args []string) error {
		calledName = subcommand
		calledArgs = args
		return nil
	}

	super := cmd.NewSuperCommand(cmd.SuperCommandParams{
		Name:            "jujutest",
		Log:             &cmd.Log{},
		MissingCallback: callback,
	})

	ctx := testing.Context(c)
	code := cmd.Main(super, ctx, []string{"foo", "bar", "baz", "--debug"})
	c.Assert(code, Equals, 0)
	c.Assert(testing.Stdout(ctx), Equals, "")
	c.Assert(testing.Stderr(ctx), Equals, "")
	c.Assert(calledName, Equals, "foo")
	c.Assert(calledArgs, DeepEquals, []string{"bar", "baz", "--debug"})
}

func (s *SuperCommandSuite) TestMissingCallbackErrors(c *C) {
	var (
		calledName string
		calledArgs []string
	)

	callback := func(ctx *cmd.Context, subcommand string, args []string) error {
		calledName = subcommand
		calledArgs = args
		return fmt.Errorf("command not found %q", subcommand)
	}

	super := cmd.NewSuperCommand(cmd.SuperCommandParams{
		Name:            "jujutest",
		Log:             &cmd.Log{},
		MissingCallback: callback,
	})

	ctx := testing.Context(c)
	code := cmd.Main(super, ctx, []string{"foo", "bar", "baz", "--debug"})
	c.Assert(code, Equals, 1)
	c.Assert(testing.Stdout(ctx), Equals, "")
	c.Assert(testing.Stderr(ctx), Equals, "error: command not found \"foo\"\n")
	c.Assert(calledName, Equals, "foo")
	c.Assert(calledArgs, DeepEquals, []string{"bar", "baz", "--debug"})
}

func (s *SuperCommandSuite) TestMissingCallbackContextWiredIn(c *C) {
	callback := func(ctx *cmd.Context, subcommand string, args []string) error {
		fmt.Fprintf(ctx.Stdout, "this is std out")
		fmt.Fprintf(ctx.Stderr, "this is std err")
		return nil
	}

	super := cmd.NewSuperCommand(cmd.SuperCommandParams{
		Name:            "jujutest",
		Log:             &cmd.Log{},
		MissingCallback: callback,
	})

	ctx := testing.Context(c)
	code := cmd.Main(super, ctx, []string{"foo", "bar", "baz", "--debug"})
	c.Assert(code, Equals, 0)
	c.Assert(testing.Stdout(ctx), Equals, "this is std out")
	c.Assert(testing.Stderr(ctx), Equals, "this is std err")
}
<commit_msg>Refactor tests, and tighten up expectations.<commit_after>// Copyright 2012, 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package cmd_test

import (
	"fmt"

	. "launchpad.net/gocheck"
	"launchpad.net/juju-core/cmd"
	"launchpad.net/juju-core/log"
	"launchpad.net/juju-core/testing"
)

func initDefenestrate(args []string) (*cmd.SuperCommand, *TestCommand, error) {
	jc := cmd.NewSuperCommand(cmd.SuperCommandParams{Name: "jujutest"})
	tc := &TestCommand{Name: "defenestrate"}
	jc.Register(tc)
	return jc, tc, testing.InitCommand(jc, args)
}

type SuperCommandSuite struct{}

var _ = Suite(&SuperCommandSuite{})

const helpText = "\n    help\\s+- show help on a command or other topic"
const helpCommandsText = "commands:" + helpText

func (s *SuperCommandSuite) TestDispatch(c *C) {
	jc := cmd.NewSuperCommand(cmd.SuperCommandParams{Name: "jujutest"})
	info := jc.Info()
	c.Assert(info.Name, Equals, "jujutest")
	c.Assert(info.Args, Equals, "<command> ...")
	c.Assert(info.Doc, Matches, helpCommandsText)

	jc, _, err := initDefenestrate([]string{"discombobulate"})
	c.Assert(err, ErrorMatches, "unrecognized command: jujutest discombobulate")
	info = jc.Info()
	c.Assert(info.Name, Equals, "jujutest")
	c.Assert(info.Args, Equals, "<command> ...")
	c.Assert(info.Doc, Matches, "commands:\n    defenestrate - defenestrate the juju"+helpText)

	jc, tc, err := initDefenestrate([]string{"defenestrate"})
	c.Assert(err, IsNil)
	c.Assert(tc.Option, Equals, "")
	info = jc.Info()
	c.Assert(info.Name, Equals, "jujutest defenestrate")
	c.Assert(info.Args, Equals, "<something>")
	c.Assert(info.Doc, Equals, "defenestrate-doc")

	_, tc, err = initDefenestrate([]string{"defenestrate", "--option", "firmly"})
	c.Assert(err, IsNil)
	c.Assert(tc.Option, Equals, "firmly")

	_, tc, err = initDefenestrate([]string{"defenestrate", "gibberish"})
	c.Assert(err, ErrorMatches, `unrecognized args: \["gibberish"\]`)
}

func (s *SuperCommandSuite) TestRegister(c *C) {
	jc := cmd.NewSuperCommand(cmd.SuperCommandParams{Name: "jujutest"})
	jc.Register(&TestCommand{Name: "flip"})
	jc.Register(&TestCommand{Name: "flap"})
	badCall := func() { jc.Register(&TestCommand{Name: "flap"}) }
	c.Assert(badCall, PanicMatches, "command already registered: flap")
}

func (s *SuperCommandSuite) TestRegisterAlias(c *C) {
	jc := cmd.NewSuperCommand(cmd.SuperCommandParams{Name: "jujutest"})
	jc.Register(&TestCommand{Name: "flip", Aliases: []string{"flap", "flop"}})

	info := jc.Info()
	c.Assert(info.Doc, Equals, `commands:
    flap - alias for flip
    flip - flip the juju
    flop - alias for flip
    help - show help on a command or other topic`)
}

var commandsDoc = `commands:
    flapbabble - flapbabble the juju
    flip       - flip the juju`

func (s *SuperCommandSuite) TestInfo(c *C) {
	jc := cmd.NewSuperCommand(cmd.SuperCommandParams{
		Name:    "jujutest",
		Purpose: "to be purposeful",
		Doc:     "doc\nblah\ndoc",
	})
	info := jc.Info()
	c.Assert(info.Name, Equals, "jujutest")
	c.Assert(info.Purpose, Equals, "to be purposeful")
	// info doc starts with the jc.Doc and ends with the help command
	c.Assert(info.Doc, Matches, jc.Doc+"(.|\n)*")
	c.Assert(info.Doc, Matches, "(.|\n)*"+helpCommandsText)

	jc.Register(&TestCommand{Name: "flip"})
	jc.Register(&TestCommand{Name: "flapbabble"})
	info = jc.Info()
	c.Assert(info.Doc, Matches, jc.Doc+"\n\n"+commandsDoc+helpText)

	jc.Doc = ""
	info = jc.Info()
	c.Assert(info.Doc, Matches, commandsDoc+helpText)
}

func (s *SuperCommandSuite) TestLogging(c *C) {
	target, debug := log.Target(), log.Debug
	defer func() {
		log.SetTarget(target)
		log.Debug = debug
	}()
	jc := cmd.NewSuperCommand(cmd.SuperCommandParams{Name: "jujutest", Log: &cmd.Log{}})
	jc.Register(&TestCommand{Name: "blah"})
	ctx := testing.Context(c)
	code := cmd.Main(jc, ctx, []string{"blah", "--option", "error", "--debug"})
	c.Assert(code, Equals, 1)
	c.Assert(bufferString(ctx.Stderr), Matches, `^.* ERROR command failed: BAM!
error: BAM!
`)
}

func NewSuperWithCallback(callback func(*cmd.Context, string, []string) error) cmd.Command {
	return cmd.NewSuperCommand(cmd.SuperCommandParams{
		Name:            "jujutest",
		Log:             &cmd.Log{},
		MissingCallback: callback,
	})
}

func (s *SuperCommandSuite) TestMissingCallback(c *C) {
	var calledName string
	var calledArgs []string

	callback := func(ctx *cmd.Context, subcommand string, args []string) error {
		calledName = subcommand
		calledArgs = args
		return nil
	}

	code := cmd.Main(
		NewSuperWithCallback(callback),
		testing.Context(c),
		[]string{"foo", "bar", "baz", "--debug"})
	c.Assert(code, Equals, 0)
	c.Assert(calledName, Equals, "foo")
	c.Assert(calledArgs, DeepEquals, []string{"bar", "baz", "--debug"})
}

func (s *SuperCommandSuite) TestMissingCallbackErrors(c *C) {
	callback := func(ctx *cmd.Context, subcommand string, args []string) error {
		return fmt.Errorf("command not found %q", subcommand)
	}

	ctx := testing.Context(c)
	code := cmd.Main(NewSuperWithCallback(callback), ctx, []string{"foo"})
	c.Assert(code, Equals, 1)
	c.Assert(testing.Stdout(ctx), Equals, "")
	c.Assert(testing.Stderr(ctx), Equals, "error: command not found \"foo\"\n")
}

func (s *SuperCommandSuite) TestMissingCallbackContextWiredIn(c *C) {
	callback := func(ctx *cmd.Context, subcommand string, args []string) error {
		fmt.Fprintf(ctx.Stdout, "this is std out")
		fmt.Fprintf(ctx.Stderr, "this is std err")
		return nil
	}

	ctx := testing.Context(c)
	code := cmd.Main(NewSuperWithCallback(callback), ctx, []string{"foo", "bar", "baz", "--debug"})
	c.Assert(code, Equals, 0)
	c.Assert(testing.Stdout(ctx), Equals, "this is std out")
	c.Assert(testing.Stderr(ctx), Equals, "this is std err")
}
<|endoftext|>


<commit_before>package tarexport // import "github.com/docker/docker/image/tarexport"

import (
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"runtime"
	"time"

	"github.com/docker/distribution"
	"github.com/docker/distribution/reference"
	"github.com/docker/docker/image"
	"github.com/docker/docker/image/v1"
	"github.com/docker/docker/layer"
	"github.com/docker/docker/pkg/archive"
	"github.com/docker/docker/pkg/system"
	"github.com/opencontainers/go-digest"
	"github.com/pkg/errors"
)

type imageDescriptor struct {
	refs     []reference.NamedTagged
	layers   []string
	image    *image.Image
	layerRef layer.Layer
}

type saveSession struct {
	*tarexporter
	outDir      string
	images      map[image.ID]*imageDescriptor
	savedLayers map[string]struct{}
	diffIDPaths map[layer.DiffID]string // cache every diffID blob to avoid duplicates
}

func (l *tarexporter) Save(names []string, outStream io.Writer) error {
	images, err := l.parseNames(names)
	if err != nil {
		return err
	}

	// Release all the image top layer references
	defer l.releaseLayerReferences(images)
	return (&saveSession{tarexporter: l, images: images}).save(outStream)
}

// parseNames will parse the image names to a map which contains image.ID to *imageDescriptor.
// Each imageDescriptor holds an image top layer reference named 'layerRef'. It is taken here, should be released later.
func (l *tarexporter) parseNames(names []string) (desc map[image.ID]*imageDescriptor, rErr error) {
	imgDescr := make(map[image.ID]*imageDescriptor)
	defer func() {
		if rErr != nil {
			l.releaseLayerReferences(imgDescr)
		}
	}()

	addAssoc := func(id image.ID, ref reference.Named) error {
		if _, ok := imgDescr[id]; !ok {
			descr := &imageDescriptor{}
			if err := l.takeLayerReference(id, descr); err != nil {
				return err
			}
			imgDescr[id] = descr
		}

		if ref != nil {
			if _, ok := ref.(reference.Canonical); ok {
				return nil
			}
			tagged, ok := reference.TagNameOnly(ref).(reference.NamedTagged)
			if !ok {
				return nil
			}

			for _, t := range imgDescr[id].refs {
				if tagged.String() == t.String() {
					return nil
				}
			}
			imgDescr[id].refs = append(imgDescr[id].refs, tagged)
		}
		return nil
	}

	for _, name := range names {
		ref, err := reference.ParseAnyReference(name)
		if err != nil {
			return nil, err
		}
		namedRef, ok := ref.(reference.Named)
		if !ok {
			// Check if digest ID reference
			if digested, ok := ref.(reference.Digested); ok {
				id := image.IDFromDigest(digested.Digest())
				if err := addAssoc(id, nil); err != nil {
					return nil, err
				}
				continue
			}
			return nil, errors.Errorf("invalid reference: %v", name)
		}

		if reference.FamiliarName(namedRef) == string(digest.Canonical) {
			imgID, err := l.is.Search(name)
			if err != nil {
				return nil, err
			}
			if err := addAssoc(imgID, nil); err != nil {
				return nil, err
			}
			continue
		}
		if reference.IsNameOnly(namedRef) {
			assocs := l.rs.ReferencesByName(namedRef)
			for _, assoc := range assocs {
				if err := addAssoc(image.IDFromDigest(assoc.ID), assoc.Ref); err != nil {
					return nil, err
				}
			}
			if len(assocs) == 0 {
				imgID, err := l.is.Search(name)
				if err != nil {
					return nil, err
				}
				if err := addAssoc(imgID, nil); err != nil {
					return nil, err
				}
			}
			continue
		}
		id, err := l.rs.Get(namedRef)
		if err != nil {
			return nil, err
		}
		if err := addAssoc(image.IDFromDigest(id), namedRef); err != nil {
			return nil, err
		}

	}
	return imgDescr, nil
}

// takeLayerReference will take/Get the image top layer reference
func (l *tarexporter) takeLayerReference(id image.ID, imgDescr *imageDescriptor) error {
	img, err := l.is.Get(id)
	if err != nil {
		return err
	}
	imgDescr.image = img
	topLayerID := img.RootFS.ChainID()
	if topLayerID == "" {
		return nil
	}
	os := img.OS
	if os == "" {
		os = runtime.GOOS
	}
	if !system.IsOSSupported(os) {
		return fmt.Errorf("os %q is not supported", os)
	}
	layer, err := l.lss[os].Get(topLayerID)
	if err != nil {
		return err
	}
	imgDescr.layerRef = layer
	return nil
}

// releaseLayerReferences will release all the image top layer references
func (l *tarexporter) releaseLayerReferences(imgDescr map[image.ID]*imageDescriptor) error {
	for _, descr := range imgDescr {
		if descr.layerRef != nil {
			os := descr.image.OS
			if os == "" {
				os = runtime.GOOS
			}
			l.lss[os].Release(descr.layerRef)
		}
	}
	return nil
}

func (s *saveSession) save(outStream io.Writer) error {
	s.savedLayers = make(map[string]struct{})
	s.diffIDPaths = make(map[layer.DiffID]string)

	// get image json
	tempDir, err := ioutil.TempDir("", "docker-export-")
	if err != nil {
		return err
	}
	defer os.RemoveAll(tempDir)

	s.outDir = tempDir
	reposLegacy := make(map[string]map[string]string)

	var manifest []manifestItem
	var parentLinks []parentLink

	for id, imageDescr := range s.images {
		foreignSrcs, err := s.saveImage(id)
		if err != nil {
			return err
		}

		var repoTags []string
		var layers []string

		for _, ref := range imageDescr.refs {
			familiarName := reference.FamiliarName(ref)
			if _, ok := reposLegacy[familiarName]; !ok {
				reposLegacy[familiarName] = make(map[string]string)
			}
			reposLegacy[familiarName][ref.Tag()] = imageDescr.layers[len(imageDescr.layers)-1]
			repoTags = append(repoTags, reference.FamiliarString(ref))
		}

		for _, l := range imageDescr.layers {
			layers = append(layers, filepath.Join(l, legacyLayerFileName))
		}

		manifest = append(manifest, manifestItem{
			Config:       id.Digest().Hex() + ".json",
			RepoTags:     repoTags,
			Layers:       layers,
			LayerSources: foreignSrcs,
		})

		parentID, _ := s.is.GetParent(id)
		parentLinks = append(parentLinks, parentLink{id, parentID})
		s.tarexporter.loggerImgEvent.LogImageEvent(id.String(), id.String(), "save")
	}

	for i, p := range validatedParentLinks(parentLinks) {
		if p.parentID != "" {
			manifest[i].Parent = p.parentID
		}
	}

	if len(reposLegacy) > 0 {
		reposFile := filepath.Join(tempDir, legacyRepositoriesFileName)
		rf, err := os.OpenFile(reposFile, os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0644)
		if err != nil {
			return err
		}

		if err := json.NewEncoder(rf).Encode(reposLegacy); err != nil {
			rf.Close()
			return err
		}

		rf.Close()

		if err := system.Chtimes(reposFile, time.Unix(0, 0), time.Unix(0, 0)); err != nil {
			return err
		}
	}

	manifestFileName := filepath.Join(tempDir, manifestFileName)
	f, err := os.OpenFile(manifestFileName, os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0644)
	if err != nil {
		return err
	}

	if err := json.NewEncoder(f).Encode(manifest); err != nil {
		f.Close()
		return err
	}

	f.Close()

	if err := system.Chtimes(manifestFileName, time.Unix(0, 0), time.Unix(0, 0)); err != nil {
		return err
	}

	fs, err := archive.Tar(tempDir, archive.Uncompressed)
	if err != nil {
		return err
	}
	defer fs.Close()

	_, err = io.Copy(outStream, fs)
	return err
}

func (s *saveSession) saveImage(id image.ID) (map[layer.DiffID]distribution.Descriptor, error) {
	img := s.images[id].image
	if len(img.RootFS.DiffIDs) == 0 {
		return nil, fmt.Errorf("empty export - not implemented")
	}

	var parent digest.Digest
	var layers []string
	var foreignSrcs map[layer.DiffID]distribution.Descriptor
	for i := range img.RootFS.DiffIDs {
		v1Img := image.V1Image{
			// This is for backward compatibility used for
			// pre v1.9 docker.
			Created: time.Unix(0, 0),
		}
		if i == len(img.RootFS.DiffIDs)-1 {
			v1Img = img.V1Image
		}
		rootFS := *img.RootFS
		rootFS.DiffIDs = rootFS.DiffIDs[:i+1]
		v1ID, err := v1.CreateID(v1Img, rootFS.ChainID(), parent)
		if err != nil {
			return nil, err
		}

		v1Img.ID = v1ID.Hex()
		if parent != "" {
			v1Img.Parent = parent.Hex()
		}

		v1Img.OS = img.OS
		src, err := s.saveLayer(rootFS.ChainID(), v1Img, img.Created)
		if err != nil {
			return nil, err
		}
		layers = append(layers, v1Img.ID)
		parent = v1ID
		if src.Digest != "" {
			if foreignSrcs == nil {
				foreignSrcs = make(map[layer.DiffID]distribution.Descriptor)
			}
			foreignSrcs[img.RootFS.DiffIDs[i]] = src
		}
	}

	configFile := filepath.Join(s.outDir, id.Digest().Hex()+".json")
	if err := ioutil.WriteFile(configFile, img.RawJSON(), 0644); err != nil {
		return nil, err
	}
	if err := system.Chtimes(configFile, img.Created, img.Created); err != nil {
		return nil, err
	}

	s.images[id].layers = layers
	return foreignSrcs, nil
}

func (s *saveSession) saveLayer(id layer.ChainID, legacyImg image.V1Image, createdTime time.Time) (distribution.Descriptor, error) {
	if _, exists := s.savedLayers[legacyImg.ID]; exists {
		return distribution.Descriptor{}, nil
	}

	outDir := filepath.Join(s.outDir, legacyImg.ID)
	if err := os.Mkdir(outDir, 0755); err != nil {
		return distribution.Descriptor{}, err
	}

	// todo: why is this version file here?
	if err := ioutil.WriteFile(filepath.Join(outDir, legacyVersionFileName), []byte("1.0"), 0644); err != nil {
		return distribution.Descriptor{}, err
	}

	imageConfig, err := json.Marshal(legacyImg)
	if err != nil {
		return distribution.Descriptor{}, err
	}

	if err := ioutil.WriteFile(filepath.Join(outDir, legacyConfigFileName), imageConfig, 0644); err != nil {
		return distribution.Descriptor{}, err
	}

	// serialize filesystem
	layerPath := filepath.Join(outDir, legacyLayerFileName)
	operatingSystem := legacyImg.OS
	if operatingSystem == "" {
		operatingSystem = runtime.GOOS
	}
	l, err := s.lss[operatingSystem].Get(id)
	if err != nil {
		return distribution.Descriptor{}, err
	}
	defer layer.ReleaseAndLog(s.lss[operatingSystem], l)

	if oldPath, exists := s.diffIDPaths[l.DiffID()]; exists {
		relPath, err := filepath.Rel(outDir, oldPath)
		if err != nil {
			return distribution.Descriptor{}, err
		}
		if err := os.Symlink(relPath, layerPath); err != nil {
			return distribution.Descriptor{}, errors.Wrap(err, "error creating symlink while saving layer")
		}
	} else {
		// Use system.CreateSequential rather than os.Create. This ensures sequential
		// file access on Windows to avoid eating into MM standby list.
		// On Linux, this equates to a regular os.Create.
		tarFile, err := system.CreateSequential(layerPath)
		if err != nil {
			return distribution.Descriptor{}, err
		}
		defer tarFile.Close()

		arch, err := l.TarStream()
		if err != nil {
			return distribution.Descriptor{}, err
		}
		defer arch.Close()

		if _, err := io.Copy(tarFile, arch); err != nil {
			return distribution.Descriptor{}, err
		}

		for _, fname := range []string{"", legacyVersionFileName, legacyConfigFileName, legacyLayerFileName} {
			// todo: maybe save layer created timestamp?
			if err := system.Chtimes(filepath.Join(outDir, fname), createdTime, createdTime); err != nil {
				return distribution.Descriptor{}, err
			}
		}

		s.diffIDPaths[l.DiffID()] = layerPath
	}
	s.savedLayers[legacyImg.ID] = struct{}{}

	var src distribution.Descriptor
	if fs, ok := l.(distribution.Describable); ok {
		src = fs.Descriptor()
	}
	return src, nil
}
<commit_msg>LCOW: Write saved manifest.json in Unix paths<commit_after>package tarexport // import "github.com/docker/docker/image/tarexport"

import (
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path"
	"path/filepath"
	"runtime"
	"time"

	"github.com/docker/distribution"
	"github.com/docker/distribution/reference"
	"github.com/docker/docker/image"
	"github.com/docker/docker/image/v1"
	"github.com/docker/docker/layer"
	"github.com/docker/docker/pkg/archive"
	"github.com/docker/docker/pkg/system"
	"github.com/opencontainers/go-digest"
	"github.com/pkg/errors"
)

type imageDescriptor struct {
	refs     []reference.NamedTagged
	layers   []string
	image    *image.Image
	layerRef layer.Layer
}

type saveSession struct {
	*tarexporter
	outDir      string
	images      map[image.ID]*imageDescriptor
	savedLayers map[string]struct{}
	diffIDPaths map[layer.DiffID]string // cache every diffID blob to avoid duplicates
}

func (l *tarexporter) Save(names []string, outStream io.Writer) error {
	images, err := l.parseNames(names)
	if err != nil {
		return err
	}

	// Release all the image top layer references
	defer l.releaseLayerReferences(images)
	return (&saveSession{tarexporter: l, images: images}).save(outStream)
}

// parseNames will parse the image names to a map which contains image.ID to *imageDescriptor.
// Each imageDescriptor holds an image top layer reference named 'layerRef'. It is taken here, should be released later.
func (l *tarexporter) parseNames(names []string) (desc map[image.ID]*imageDescriptor, rErr error) {
	imgDescr := make(map[image.ID]*imageDescriptor)
	defer func() {
		if rErr != nil {
			l.releaseLayerReferences(imgDescr)
		}
	}()

	addAssoc := func(id image.ID, ref reference.Named) error {
		if _, ok := imgDescr[id]; !ok {
			descr := &imageDescriptor{}
			if err := l.takeLayerReference(id, descr); err != nil {
				return err
			}
			imgDescr[id] = descr
		}

		if ref != nil {
			if _, ok := ref.(reference.Canonical); ok {
				return nil
			}
			tagged, ok := reference.TagNameOnly(ref).(reference.NamedTagged)
			if !ok {
				return nil
			}

			for _, t := range imgDescr[id].refs {
				if tagged.String() == t.String() {
					return nil
				}
			}
			imgDescr[id].refs = append(imgDescr[id].refs, tagged)
		}
		return nil
	}

	for _, name := range names {
		ref, err := reference.ParseAnyReference(name)
		if err != nil {
			return nil, err
		}
		namedRef, ok := ref.(reference.Named)
		if !ok {
			// Check if digest ID reference
			if digested, ok := ref.(reference.Digested); ok {
				id := image.IDFromDigest(digested.Digest())
				if err := addAssoc(id, nil); err != nil {
					return nil, err
				}
				continue
			}
			return nil, errors.Errorf("invalid reference: %v", name)
		}

		if reference.FamiliarName(namedRef) == string(digest.Canonical) {
			imgID, err := l.is.Search(name)
			if err != nil {
				return nil, err
			}
			if err := addAssoc(imgID, nil); err != nil {
				return nil, err
			}
			continue
		}
		if reference.IsNameOnly(namedRef) {
			assocs := l.rs.ReferencesByName(namedRef)
			for _, assoc := range assocs {
				if err := addAssoc(image.IDFromDigest(assoc.ID), assoc.Ref); err != nil {
					return nil, err
				}
			}
			if len(assocs) == 0 {
				imgID, err := l.is.Search(name)
				if err != nil {
					return nil, err
				}
				if err := addAssoc(imgID, nil); err != nil {
					return nil, err
				}
			}
			continue
		}
		id, err := l.rs.Get(namedRef)
		if err != nil {
			return nil, err
		}
		if err := addAssoc(image.IDFromDigest(id), namedRef); err != nil {
			return nil, err
		}

	}
	return imgDescr, nil
}

// takeLayerReference will take/Get the image top layer reference
func (l *tarexporter) takeLayerReference(id image.ID, imgDescr *imageDescriptor) error {
	img, err := l.is.Get(id)
	if err != nil {
		return err
	}
	imgDescr.image = img
	topLayerID := img.RootFS.ChainID()
	if topLayerID == "" {
		return nil
	}
	os := img.OS
	if os == "" {
		os = runtime.GOOS
	}
	if !system.IsOSSupported(os) {
		return fmt.Errorf("os %q is not supported", os)
	}
	layer, err := l.lss[os].Get(topLayerID)
	if err != nil {
		return err
	}
	imgDescr.layerRef = layer
	return nil
}

// releaseLayerReferences will release all the image top layer references
func (l *tarexporter) releaseLayerReferences(imgDescr map[image.ID]*imageDescriptor) error {
	for _, descr := range imgDescr {
		if descr.layerRef != nil {
			os := descr.image.OS
			if os == "" {
				os = runtime.GOOS
			}
			l.lss[os].Release(descr.layerRef)
		}
	}
	return nil
}

func (s *saveSession) save(outStream io.Writer) error {
	s.savedLayers = make(map[string]struct{})
	s.diffIDPaths = make(map[layer.DiffID]string)

	// get image json
	tempDir, err := ioutil.TempDir("", "docker-export-")
	if err != nil {
		return err
	}
	defer os.RemoveAll(tempDir)

	s.outDir = tempDir
	reposLegacy := make(map[string]map[string]string)

	var manifest []manifestItem
	var parentLinks []parentLink

	for id, imageDescr := range s.images {
		foreignSrcs, err := s.saveImage(id)
		if err != nil {
			return err
		}

		var repoTags []string
		var layers []string

		for _, ref := range imageDescr.refs {
			familiarName := reference.FamiliarName(ref)
			if _, ok := reposLegacy[familiarName]; !ok {
				reposLegacy[familiarName] = make(map[string]string)
			}
			reposLegacy[familiarName][ref.Tag()] = imageDescr.layers[len(imageDescr.layers)-1]
			repoTags = append(repoTags, reference.FamiliarString(ref))
		}

		for _, l := range imageDescr.layers {
			// IMPORTANT: We use path, not filepath here to ensure the layers
			// in the manifest use Unix-style forward-slashes. Otherwise, a
			// Linux image saved from LCOW won't be able to be imported on
			// LCOL.
			layers = append(layers, path.Join(l, legacyLayerFileName))
		}

		manifest = append(manifest, manifestItem{
			Config:       id.Digest().Hex() + ".json",
			RepoTags:     repoTags,
			Layers:       layers,
			LayerSources: foreignSrcs,
		})

		parentID, _ := s.is.GetParent(id)
		parentLinks = append(parentLinks, parentLink{id, parentID})
		s.tarexporter.loggerImgEvent.LogImageEvent(id.String(), id.String(), "save")
	}

	for i, p := range validatedParentLinks(parentLinks) {
		if p.parentID != "" {
			manifest[i].Parent = p.parentID
		}
	}

	if len(reposLegacy) > 0 {
		reposFile := filepath.Join(tempDir, legacyRepositoriesFileName)
		rf, err := os.OpenFile(reposFile, os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0644)
		if err != nil {
			return err
		}

		if err := json.NewEncoder(rf).Encode(reposLegacy); err != nil {
			rf.Close()
			return err
		}

		rf.Close()

		if err := system.Chtimes(reposFile, time.Unix(0, 0), time.Unix(0, 0)); err != nil {
			return err
		}
	}

	manifestFileName := filepath.Join(tempDir, manifestFileName)
	f, err := os.OpenFile(manifestFileName, os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0644)
	if err != nil {
		return err
	}

	if err := json.NewEncoder(f).Encode(manifest); err != nil {
		f.Close()
		return err
	}

	f.Close()

	if err := system.Chtimes(manifestFileName, time.Unix(0, 0), time.Unix(0, 0)); err != nil {
		return err
	}

	fs, err := archive.Tar(tempDir, archive.Uncompressed)
	if err != nil {
		return err
	}
	defer fs.Close()

	_, err = io.Copy(outStream, fs)
	return err
}

func (s *saveSession) saveImage(id image.ID) (map[layer.DiffID]distribution.Descriptor, error) {
	img := s.images[id].image
	if len(img.RootFS.DiffIDs) == 0 {
		return nil, fmt.Errorf("empty export - not implemented")
	}

	var parent digest.Digest
	var layers []string
	var foreignSrcs map[layer.DiffID]distribution.Descriptor
	for i := range img.RootFS.DiffIDs {
		v1Img := image.V1Image{
			// This is for backward compatibility used for
			// pre v1.9 docker.
			Created: time.Unix(0, 0),
		}
		if i == len(img.RootFS.DiffIDs)-1 {
			v1Img = img.V1Image
		}
		rootFS := *img.RootFS
		rootFS.DiffIDs = rootFS.DiffIDs[:i+1]
		v1ID, err := v1.CreateID(v1Img, rootFS.ChainID(), parent)
		if err != nil {
			return nil, err
		}

		v1Img.ID = v1ID.Hex()
		if parent != "" {
			v1Img.Parent = parent.Hex()
		}

		v1Img.OS = img.OS
		src, err := s.saveLayer(rootFS.ChainID(), v1Img, img.Created)
		if err != nil {
			return nil, err
		}
		layers = append(layers, v1Img.ID)
		parent = v1ID
		if src.Digest != "" {
			if foreignSrcs == nil {
				foreignSrcs = make(map[layer.DiffID]distribution.Descriptor)
			}
			foreignSrcs[img.RootFS.DiffIDs[i]] = src
		}
	}

	configFile := filepath.Join(s.outDir, id.Digest().Hex()+".json")
	if err := ioutil.WriteFile(configFile, img.RawJSON(), 0644); err != nil {
		return nil, err
	}
	if err := system.Chtimes(configFile, img.Created, img.Created); err != nil {
		return nil, err
	}

	s.images[id].layers = layers
	return foreignSrcs, nil
}

func (s *saveSession) saveLayer(id layer.ChainID, legacyImg image.V1Image, createdTime time.Time) (distribution.Descriptor, error) {
	if _, exists := s.savedLayers[legacyImg.ID]; exists {
		return distribution.Descriptor{}, nil
	}

	outDir := filepath.Join(s.outDir, legacyImg.ID)
	if err := os.Mkdir(outDir, 0755); err != nil {
		return distribution.Descriptor{}, err
	}

	// todo: why is this version file here?
	if err := ioutil.WriteFile(filepath.Join(outDir, legacyVersionFileName), []byte("1.0"), 0644); err != nil {
		return distribution.Descriptor{}, err
	}

	imageConfig, err := json.Marshal(legacyImg)
	if err != nil {
		return distribution.Descriptor{}, err
	}

	if err := ioutil.WriteFile(filepath.Join(outDir, legacyConfigFileName), imageConfig, 0644); err != nil {
		return distribution.Descriptor{}, err
	}

	// serialize filesystem
	layerPath := filepath.Join(outDir, legacyLayerFileName)
	operatingSystem := legacyImg.OS
	if operatingSystem == "" {
		operatingSystem = runtime.GOOS
	}
	l, err := s.lss[operatingSystem].Get(id)
	if err != nil {
		return distribution.Descriptor{}, err
	}
	defer layer.ReleaseAndLog(s.lss[operatingSystem], l)

	if oldPath, exists := s.diffIDPaths[l.DiffID()]; exists {
		relPath, err := filepath.Rel(outDir, oldPath)
		if err != nil {
			return distribution.Descriptor{}, err
		}
		if err := os.Symlink(relPath, layerPath); err != nil {
			return distribution.Descriptor{}, errors.Wrap(err, "error creating symlink while saving layer")
		}
	} else {
		// Use system.CreateSequential rather than os.Create. This ensures sequential
		// file access on Windows to avoid eating into MM standby list.
		// On Linux, this equates to a regular os.Create.
		tarFile, err := system.CreateSequential(layerPath)
		if err != nil {
			return distribution.Descriptor{}, err
		}
		defer tarFile.Close()

		arch, err := l.TarStream()
		if err != nil {
			return distribution.Descriptor{}, err
		}
		defer arch.Close()

		if _, err := io.Copy(tarFile, arch); err != nil {
			return distribution.Descriptor{}, err
		}

		for _, fname := range []string{"", legacyVersionFileName, legacyConfigFileName, legacyLayerFileName} {
			// todo: maybe save layer created timestamp?
			if err := system.Chtimes(filepath.Join(outDir, fname), createdTime, createdTime); err != nil {
				return distribution.Descriptor{}, err
			}
		}

		s.diffIDPaths[l.DiffID()] = layerPath
	}
	s.savedLayers[legacyImg.ID] = struct{}{}

	var src distribution.Descriptor
	if fs, ok := l.(distribution.Describable); ok {
		src = fs.Descriptor()
	}
	return src, nil
}
<|endoftext|>


<commit_before><commit_msg>Remove extra logging<commit_after><|endoftext|>


<commit_before>package rx

import (
	"time"

	"github.com/reactivego/multicast"
)

//jig:template Subject<Foo>
//jig:embeds Observable<Foo>

// SubjectFoo is a combination of an observer and observable. Subjects are
// special because they are the only reactive constructs that support
// multicasting. The items sent to it through its observer side are
// multicasted to multiple clients subscribed to its observable side.
//
// A SubjectFoo embeds ObservableFoo and FooObserveFunc. This exposes the
// methods and fields of both types on SubjectFoo. Use the ObservableFoo
// methods to subscribe to it. Use the FooObserveFunc Next, Error and Complete
// methods to feed data to it.
//
// After a subject has been terminated by calling either Error or Complete,
// it goes into terminated state. All subsequent calls to its observer side
// will be silently ignored. All subsequent subscriptions to the observable
// side will be handled according to the specific behavior of the subject.
// There are different types of subjects, see the different NewXxxSubjectFoo
// functions for more info.
//
// Important! a subject is a hot observable. This means that subscribing to
// it will block the calling goroutine while it is waiting for items and
// notifications to receive. Unless you have code on a different goroutine
// already feeding into the subject, your subscribe will deadlock.
// Alternatively, you could subscribe on a goroutine as shown in the example.
type SubjectFoo struct {
	ObservableFoo
	FooObserveFunc
}

//jig:template NewSubject<Foo>
//jig:needs Subject<Foo>

// NewSubjectFoo creates a new Subject. After the subject is
// terminated, all subsequent subscriptions to the observable side will be
// terminated immediately with either an Error or Complete notification send to
// the subscribing client
//
// Note that this implementation is blocking. When no subcribers are present
// then the data can flow freely. But when there are subscribers, the observable
// goroutine is blocked until all subscribers have processed the next, error or
// complete notification.
func NewSubjectFoo() SubjectFoo {
	ch := multicast.NewChan(1, 16 /*max endpoints*/)

	observable := Observable(func(observe ObserveFunc, subscribeOn Scheduler, subscriber Subscriber) {
		ep, err := ch.NewEndpoint(0)
		if err != nil {
			observe(nil, err, true)
			return
		}
		observable := Create(func(observer Observer) {
			receive := func(value interface{}, err error, closed bool) bool {
				if !closed {
					observer.Next(value)
				} else {
					observer.Error(err)
				}
				return !observer.Closed()
			}
			ep.Range(receive, 0)
		})
		observable(observe, subscribeOn, subscriber.Add(ep.Cancel))
	})

	observer := func(next foo, err error, done bool) {
		if !ch.Closed() {
			if !done {
				ch.FastSend(next)
			} else {
				ch.Close(err)
			}
		}
	}

	return SubjectFoo{observable.AsObservableFoo(), observer}
}

//jig:template MaxReplayCapacity

// MaxReplayCapacity is the maximum size of a replay buffer. Can be modified.
var MaxReplayCapacity = 16383

//jig:template NewReplaySubject<Foo>
//jig:needs Subject<Foo>, MaxReplayCapacity

// NewReplaySubjectFoo creates a new ReplaySubject. ReplaySubject ensures that
// all observers see the same sequence of emitted items, even if they
// subscribe after. When bufferCapacity argument is 0, then MaxReplayCapacity is
// used (currently 16383). When windowDuration argument is 0, then entries added
// to the buffer will remain fresh forever.
//
// Note that this implementation is non-blocking. When no subscribers are
// present the buffer fills up to bufferCapacity after which new items will
// start overwriting the oldest ones according to the FIFO principle.
// If a subscriber cannot keep up with the data rate of the source observable,
// eventually the buffer for the subscriber will overflow. At that moment the
// subscriber will receive an ErrMissingBackpressure error.
func NewReplaySubjectFoo(bufferCapacity int, windowDuration time.Duration) SubjectFoo {
	if bufferCapacity == 0 {
		bufferCapacity = MaxReplayCapacity
	}
	ch := multicast.NewChan(bufferCapacity, 16 /*max endpoints*/)

	observable := Observable(func(observe ObserveFunc, subscribeOn Scheduler, subscriber Subscriber) {
		ep, err := ch.NewEndpoint(multicast.ReplayAll)
		if err != nil {
			observe(nil, err, true)
			return
		}
		observable := Create(func(observer Observer) {
			receive := func(value interface{}, err error, closed bool) bool {
				if !closed {
					observer.Next(value)
				} else {
					observer.Error(err)
				}
				return !observer.Closed()
			}
			ep.Range(receive, windowDuration)
		})
		observable(observe, subscribeOn, subscriber.Add(ep.Cancel))
	})

	observer := func(next foo, err error, done bool) {
		if !ch.Closed() {
			if !done {
				ch.Send(next)
			} else {
				ch.Close(err)
			}
		}
	}

	return SubjectFoo{observable.AsObservableFoo(), observer}
}
<commit_msg>Update Subject documentation.<commit_after>package rx

import (
	"time"

	"github.com/reactivego/multicast"
)

//jig:template Subject<Foo>
//jig:embeds Observable<Foo>

// SubjectFoo is a combination of an observer and observable. Subjects are
// special because they are the only reactive constructs that support
// multicasting. The items sent to it through its observer side are
// multicasted to multiple clients subscribed to its observable side.
//
// A SubjectFoo embeds ObservableFoo and FooObserveFunc. This exposes the
// methods and fields of both types on SubjectFoo. Use the ObservableFoo
// methods to subscribe to it. Use the FooObserveFunc Next, Error and Complete
// methods to feed data to it.
//
// After a subject has been terminated by calling either Error or Complete,
// it goes into terminated state. All subsequent calls to its observer side
// will be silently ignored. All subsequent subscriptions to the observable
// side will be handled according to the specific behavior of the subject.
// There are different types of subjects, see the different NewXxxSubjectFoo
// functions for more info.
type SubjectFoo struct {
	ObservableFoo
	FooObserveFunc
}

//jig:template NewSubject<Foo>
//jig:needs Subject<Foo>

// NewSubjectFoo creates a new Subject. After the subject is
// terminated, all subsequent subscriptions to the observable side will be
// terminated immediately with either an Error or Complete notification send to
// the subscribing client
//
// Note that this implementation is blocking. When no subcribers are present
// then the data can flow freely. But when there are subscribers, the observable
// goroutine is blocked until all subscribers have processed the next, error or
// complete notification.
func NewSubjectFoo() SubjectFoo {
	ch := multicast.NewChan(1, 16 /*max endpoints*/)

	observable := Observable(func(observe ObserveFunc, subscribeOn Scheduler, subscriber Subscriber) {
		ep, err := ch.NewEndpoint(0)
		if err != nil {
			observe(nil, err, true)
			return
		}
		observable := Create(func(observer Observer) {
			receive := func(value interface{}, err error, closed bool) bool {
				if !closed {
					observer.Next(value)
				} else {
					observer.Error(err)
				}
				return !observer.Closed()
			}
			ep.Range(receive, 0)
		})
		observable(observe, subscribeOn, subscriber.Add(ep.Cancel))
	})

	observer := func(next foo, err error, done bool) {
		if !ch.Closed() {
			if !done {
				ch.FastSend(next)
			} else {
				ch.Close(err)
			}
		}
	}

	return SubjectFoo{observable.AsObservableFoo(), observer}
}

//jig:template MaxReplayCapacity

// MaxReplayCapacity is the maximum size of a replay buffer. Can be modified.
var MaxReplayCapacity = 16383

//jig:template NewReplaySubject<Foo>
//jig:needs Subject<Foo>, MaxReplayCapacity

// NewReplaySubjectFoo creates a new ReplaySubject. ReplaySubject ensures that
// all observers see the same sequence of emitted items, even if they
// subscribe after. When bufferCapacity argument is 0, then MaxReplayCapacity is
// used (currently 16383). When windowDuration argument is 0, then entries added
// to the buffer will remain fresh forever.
func NewReplaySubjectFoo(bufferCapacity int, windowDuration time.Duration) SubjectFoo {
	if bufferCapacity == 0 {
		bufferCapacity = MaxReplayCapacity
	}
	ch := multicast.NewChan(bufferCapacity, 16 /*max endpoints*/)

	observable := Observable(func(observe ObserveFunc, subscribeOn Scheduler, subscriber Subscriber) {
		ep, err := ch.NewEndpoint(multicast.ReplayAll)
		if err != nil {
			observe(nil, err, true)
			return
		}
		observable := Create(func(observer Observer) {
			receive := func(value interface{}, err error, closed bool) bool {
				if !closed {
					observer.Next(value)
				} else {
					observer.Error(err)
				}
				return !observer.Closed()
			}
			ep.Range(receive, windowDuration)
		})
		observable(observe, subscribeOn, subscriber.Add(ep.Cancel))
	})

	observer := func(next foo, err error, done bool) {
		if !ch.Closed() {
			if !done {
				ch.Send(next)
			} else {
				ch.Close(err)
			}
		}
	}

	return SubjectFoo{observable.AsObservableFoo(), observer}
}
<|endoftext|>


<commit_before>// Copyright 2015 Andrew E. Bruno
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package khttp is a transport that authenticates all outgoing requests using
// SPNEGO (negotiate authentication) http://tools.ietf.org/html/rfc4559.
package khttp

import (
    "errors"
    "os"
    "fmt"
    "net/http"
    "strings"

    "github.com/ubccr/kerby"
)

var (
    negotiateHeader       = "Negotiate"
    wwwAuthenticateHeader = "WWW-Authenticate"
    authorizationHeader   = "Authorization"
)

// HTTP client transport that authenticates all outgoing
// requests using SPNEGO. Implements the http.RoundTripper interface
type Transport struct {
    // keytab file to use
    KeyTab  string
    // principal
    Principal  string
    // Next specifies the next transport to be used or http.DefaultTransport if nil.
    Next http.RoundTripper
}

// RoundTrip executes a single HTTP transaction performing SPNEGO negotiate
// authentication.
func (t *Transport) RoundTrip(req *http.Request) (*http.Response, error) {
    if len(t.KeyTab) > 0 {
        os.Setenv("KRB5_CLIENT_KTNAME", t.KeyTab)
    }
    service := fmt.Sprintf("HTTP@%s", req.URL.Host)
    kc := new(kerby.KerbClient)
    err := kc.Init(service, t.Principal)
    if err != nil {
        return nil, err
    }

    err = kc.Step("")
    if err != nil {
        return nil, err
    }

    req.Header.Set(authorizationHeader, negotiateHeader+" "+kc.Response())

    tr := t.Next
    if tr == nil {
        tr = http.DefaultTransport
        if tr == nil {
            return nil, errors.New("khttp: no Next transport or DefaultTransport")
        }
    }

    resp, err := tr.RoundTrip(req)
    if err != nil {
        return nil, err
    }

    authReply := strings.Split(resp.Header.Get(wwwAuthenticateHeader), " ")
    if len(authReply) != 2 || authReply[0] != negotiateHeader {
        return nil, errors.New("khttp: server replied with invalid www-authenticate header")
    }

    // Authenticate the reply from the server
    err = kc.Step(authReply[1])
    if err != nil {
        return nil, err
    }
    kc.Clean()

    return resp, nil
}
<commit_msg>Defer cleanup<commit_after>// Copyright 2015 Andrew E. Bruno
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package khttp is a transport that authenticates all outgoing requests using
// SPNEGO (negotiate authentication) http://tools.ietf.org/html/rfc4559.
package khttp

import (
    "errors"
    "os"
    "fmt"
    "net/http"
    "strings"

    "github.com/ubccr/kerby"
)

var (
    negotiateHeader       = "Negotiate"
    wwwAuthenticateHeader = "WWW-Authenticate"
    authorizationHeader   = "Authorization"
)

// HTTP client transport that authenticates all outgoing
// requests using SPNEGO. Implements the http.RoundTripper interface
type Transport struct {
    // keytab file to use
    KeyTab  string
    // principal
    Principal  string
    // Next specifies the next transport to be used or http.DefaultTransport if nil.
    Next http.RoundTripper
}

// RoundTrip executes a single HTTP transaction performing SPNEGO negotiate
// authentication.
func (t *Transport) RoundTrip(req *http.Request) (*http.Response, error) {
    if len(t.KeyTab) > 0 {
        os.Setenv("KRB5_CLIENT_KTNAME", t.KeyTab)
    }
    service := fmt.Sprintf("HTTP@%s", req.URL.Host)
    kc := new(kerby.KerbClient)
    err := kc.Init(service, t.Principal)
    if err != nil {
        return nil, err
    }
    defer kc.Clean()

    err = kc.Step("")
    if err != nil {
        return nil, err
    }

    req.Header.Set(authorizationHeader, negotiateHeader+" "+kc.Response())

    tr := t.Next
    if tr == nil {
        tr = http.DefaultTransport
        if tr == nil {
            return nil, errors.New("khttp: no Next transport or DefaultTransport")
        }
    }

    resp, err := tr.RoundTrip(req)
    if err != nil {
        return nil, err
    }

    authReply := strings.Split(resp.Header.Get(wwwAuthenticateHeader), " ")
    if len(authReply) != 2 || authReply[0] != negotiateHeader {
        return nil, errors.New("khttp: server replied with invalid www-authenticate header")
    }

    // Authenticate the reply from the server
    err = kc.Step(authReply[1])
    if err != nil {
        return nil, err
    }

    return resp, nil
}
<|endoftext|>


<commit_before>package kit

import (
	"path/filepath"
	"sync"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/suite"
	"gopkg.in/fsnotify.v1"
)

const (
	textFixturePath  = "../fixtures/project/assets/application.js"
	watchFixturePath = "../fixtures/project"
)

type FileWatcherTestSuite struct {
	suite.Suite
	watcher *FileWatcher
}

func (suite *FileWatcherTestSuite) TestNewFileReader() {
	watcher, err := newFileWatcher(ThemeClient{}, watchFixturePath, true, eventFilter{}, func(ThemeClient, Asset, EventType, error) {})
	assert.Nil(suite.T(), err)
	assert.Equal(suite.T(), true, watcher.IsWatching())
	watcher.StopWatching()
}

func (suite *FileWatcherTestSuite) TestConvertFsEvents() {
	assetChan := make(chan Asset, 4)
	eventChan := make(chan fsnotify.Event)
	var wg sync.WaitGroup
	wg.Add(2)

	newWatcher := &FileWatcher{
		done:    make(chan bool),
		watcher: &fsnotify.Watcher{Events: eventChan},
	}

	newWatcher.callback = func(client ThemeClient, asset Asset, event EventType, err error) {
		assert.Nil(suite.T(), err)
		assert.Equal(suite.T(), Update, event)
		assetChan <- asset
		wg.Done()
	}

	go convertFsEvents(newWatcher)

	go func() {
		writes := []fsnotify.Event{
			{Name: watchFixturePath + "/templates/template.liquid", Op: fsnotify.Write},
			{Name: watchFixturePath + "/templates/template.liquid", Op: fsnotify.Write},
			{Name: watchFixturePath + "/templates/template.liquid", Op: fsnotify.Write},
			{Name: watchFixturePath + "/templates/customers/test.liquid", Op: fsnotify.Write},
		}
		for _, fsEvent := range writes {
			eventChan <- fsEvent
		}
		close(eventChan)
	}()

	wg.Wait()
	// test that the events are debounced
	assert.Equal(suite.T(), 2, len(assetChan))
}

func (suite *FileWatcherTestSuite) TestStopWatching() {
	watcher, err := newFileWatcher(ThemeClient{}, watchFixturePath, true, eventFilter{}, func(ThemeClient, Asset, EventType, error) {})
	assert.Nil(suite.T(), err)
	assert.Equal(suite.T(), true, watcher.IsWatching())
	watcher.StopWatching()
	time.Sleep(50 * time.Millisecond)
	assert.Equal(suite.T(), false, watcher.IsWatching())
}

func (suite *FileWatcherTestSuite) TestHandleEvent() {
	writes := []struct {
		Name  string
		Event fsnotify.Op
	}{
		{Name: textFixturePath, Event: fsnotify.Create},
		{Name: textFixturePath, Event: fsnotify.Write},
		{Name: textFixturePath, Event: fsnotify.Remove},
		{Name: "../fixtures/project/whatever.txt", Event: fsnotify.Write},
	}

	var wg sync.WaitGroup
	wg.Add(len(writes))

	watcher := &FileWatcher{callback: func(client ThemeClient, asset Asset, event EventType, err error) {
		if err != nil {
			assert.Equal(suite.T(), "File not in project workspace.", err.Error())
			assert.Equal(suite.T(), "../fixtures/project/whatever.txt", asset.Key)
		} else {
			assert.Equal(suite.T(), extractAssetKey(textFixturePath), asset.Key)
		}
		wg.Done()
	}}

	for _, write := range writes {
		handleEvent(watcher, fsnotify.Event{Name: write.Name, Op: write.Event})
	}

	wg.Wait()
}

func (suite *FileWatcherTestSuite) TestExtractAssetKey() {
	tests := map[string]string{
		"/long/path/to/config.yml":                      "",
		"/long/path/to/assets/logo.png":                 "assets/logo.png",
		"/long/path/to/templates/customers/test.liquid": "templates/customers/test.liquid",
		"/long/path/to/config/test.liquid":              "config/test.liquid",
		"/long/path/to/layout/test.liquid":              "layout/test.liquid",
		"/long/path/to/snippets/test.liquid":            "snippets/test.liquid",
		"/long/path/to/templates/test.liquid":           "templates/test.liquid",
		"/long/path/to/locales/test.liquid":             "locales/test.liquid",
		"/long/path/to/sections/test.liquid":            "sections/test.liquid",
	}
	for input, expected := range tests {
		assert.Equal(suite.T(), expected, extractAssetKey(input))
	}
}

func (suite *FileWatcherTestSuite) TestfindDirectoriesToWatch() {
	expected := []string{
		clean(watchFixturePath),
		clean(watchFixturePath + "/assets"),
		clean(watchFixturePath + "/config"),
		clean(watchFixturePath + "/layout"),
		clean(watchFixturePath + "/locales"),
		clean(watchFixturePath + "/snippets"),
		clean(watchFixturePath + "/templates"),
		clean(watchFixturePath + "/templates/customers"),
	}

	files := findDirectoriesToWatch(watchFixturePath, true, func(string) bool { return false })
	assert.Equal(suite.T(), expected, files)

	files = findDirectoriesToWatch(watchFixturePath, false, func(string) bool { return false })
	assert.Equal(suite.T(), []string{clean(watchFixturePath)}, files)
}

func TestFileWatcherTestSuite(t *testing.T) {
	suite.Run(t, new(FileWatcherTestSuite))
}

func clean(path string) string {
	return filepath.Clean(path)
}
<commit_msg>fixing tests<commit_after>package kit

import (
	"path/filepath"
	"sync"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/suite"
	"gopkg.in/fsnotify.v1"
)

const (
	textFixturePath  = "../fixtures/project/assets/application.js"
	watchFixturePath = "../fixtures/project"
)

type FileWatcherTestSuite struct {
	suite.Suite
	watcher *FileWatcher
}

func (suite *FileWatcherTestSuite) TestNewFileReader() {
	watcher, err := newFileWatcher(ThemeClient{}, watchFixturePath, true, eventFilter{}, func(ThemeClient, Asset, EventType, error) {})
	assert.Nil(suite.T(), err)
	assert.Equal(suite.T(), true, watcher.IsWatching())
	watcher.StopWatching()
}

func (suite *FileWatcherTestSuite) TestConvertFsEvents() {
	assetChan := make(chan Asset, 4)
	eventChan := make(chan fsnotify.Event)
	var wg sync.WaitGroup
	wg.Add(2)

	newWatcher := &FileWatcher{
		done:    make(chan bool),
		watcher: &fsnotify.Watcher{Events: eventChan},
	}

	newWatcher.callback = func(client ThemeClient, asset Asset, event EventType, err error) {
		assert.Nil(suite.T(), err)
		assert.Equal(suite.T(), Update, event)
		assetChan <- asset
		wg.Done()
	}

	go convertFsEvents(newWatcher)

	go func() {
		writes := []fsnotify.Event{
			{Name: watchFixturePath + "/templates/template.liquid", Op: fsnotify.Write},
			{Name: watchFixturePath + "/templates/template.liquid", Op: fsnotify.Write},
			{Name: watchFixturePath + "/templates/template.liquid", Op: fsnotify.Write},
			{Name: watchFixturePath + "/templates/customers/test.liquid", Op: fsnotify.Write},
		}
		for _, fsEvent := range writes {
			eventChan <- fsEvent
		}
		close(eventChan)
	}()

	wg.Wait()
	// test that the events are debounced
	assert.Equal(suite.T(), 2, len(assetChan))
}

func (suite *FileWatcherTestSuite) TestStopWatching() {
	watcher, err := newFileWatcher(ThemeClient{}, watchFixturePath, true, eventFilter{}, func(ThemeClient, Asset, EventType, error) {})
	assert.Nil(suite.T(), err)
	assert.Equal(suite.T(), true, watcher.IsWatching())
	watcher.StopWatching()
	time.Sleep(50 * time.Millisecond)
	assert.Equal(suite.T(), false, watcher.IsWatching())
}

func (suite *FileWatcherTestSuite) TestHandleEvent() {
	writes := []struct {
		Name  string
		Event fsnotify.Op
	}{
		{Name: textFixturePath, Event: fsnotify.Create},
		{Name: textFixturePath, Event: fsnotify.Write},
		{Name: textFixturePath, Event: fsnotify.Remove},
		{Name: "../fixtures/project/whatever.txt", Event: fsnotify.Write},
	}

	var wg sync.WaitGroup
	wg.Add(len(writes))

	watcher := &FileWatcher{callback: func(client ThemeClient, asset Asset, event EventType, err error) {
		if err != nil {
			assert.Equal(suite.T(), "File ../fixtures/project/whatever.txt is not in project workspace.", err.Error())
			assert.Equal(suite.T(), "../fixtures/project/whatever.txt", asset.Key)
		} else {
			assert.Equal(suite.T(), extractAssetKey(textFixturePath), asset.Key)
		}
		wg.Done()
	}}

	for _, write := range writes {
		handleEvent(watcher, fsnotify.Event{Name: write.Name, Op: write.Event})
	}

	wg.Wait()
}

func (suite *FileWatcherTestSuite) TestExtractAssetKey() {
	tests := map[string]string{
		"/long/path/to/config.yml":                      "",
		"/long/path/to/assets/logo.png":                 "assets/logo.png",
		"/long/path/to/templates/customers/test.liquid": "templates/customers/test.liquid",
		"/long/path/to/config/test.liquid":              "config/test.liquid",
		"/long/path/to/layout/test.liquid":              "layout/test.liquid",
		"/long/path/to/snippets/test.liquid":            "snippets/test.liquid",
		"/long/path/to/templates/test.liquid":           "templates/test.liquid",
		"/long/path/to/locales/test.liquid":             "locales/test.liquid",
		"/long/path/to/sections/test.liquid":            "sections/test.liquid",
	}
	for input, expected := range tests {
		assert.Equal(suite.T(), expected, extractAssetKey(input))
	}
}

func (suite *FileWatcherTestSuite) TestfindDirectoriesToWatch() {
	expected := []string{
		clean(watchFixturePath),
		clean(watchFixturePath + "/assets"),
		clean(watchFixturePath + "/config"),
		clean(watchFixturePath + "/layout"),
		clean(watchFixturePath + "/locales"),
		clean(watchFixturePath + "/snippets"),
		clean(watchFixturePath + "/templates"),
		clean(watchFixturePath + "/templates/customers"),
	}

	files := findDirectoriesToWatch(watchFixturePath, true, func(string) bool { return false })
	assert.Equal(suite.T(), expected, files)

	files = findDirectoriesToWatch(watchFixturePath, false, func(string) bool { return false })
	assert.Equal(suite.T(), []string{clean(watchFixturePath)}, files)
}

func TestFileWatcherTestSuite(t *testing.T) {
	suite.Run(t, new(FileWatcherTestSuite))
}

func clean(path string) string {
	return filepath.Clean(path)
}
<|endoftext|>


<commit_before>package ext

import (
    "time"
)

func SleepUntil(ms int64) {
    time.Sleep(time.Unix(ms/ 1000, ms % 1000 * int64(time.Millisecond)).Sub(time.Now()))
}
<commit_msg>format and test ok<commit_after>package ext

import (
	"time"
)

func SleepUntil(ms int64) {
	time.Sleep(time.Unix(ms/1000, ms%1000*int64(time.Millisecond)).Sub(time.Now()))
}
<|endoftext|>


<commit_before>package rss

import (
	"strings"
	"time"
)

// TimeLayouts is contains a list of time.Parse() layouts that are used in
// attempts to convert item.Date and item.PubDate string to time.Time values.
// The layouts are attempted in ascending order until either time.Parse()
// does not return an error or all layouts are attempted.
var TimeLayouts = []string{
	"Mon, _2 Jan 2006 15:04:05 MST",
	"Mon, _2 Jan 2006 15:04:05 -0700",
	"Mon, 2 Jan 2006 15:04 MST",
	time.ANSIC,
	time.UnixDate,
	time.RubyDate,
	time.RFC822,
	time.RFC822Z,
	time.RFC850,
	time.RFC1123,
	time.RFC1123Z,
	time.RFC3339,
	time.RFC3339Nano,
}

func parseTime(s string) (time.Time, error) {
	s = strings.TrimSpace(s)

	var e error
	var t time.Time

	for _, layout := range TimeLayouts {
		t, e = time.Parse(layout, s)
		if e == nil {
			return t, e
		}
	}

	return time.Time{}, e
}
<commit_msg>fix time issue<commit_after>package rss

import (
	"strings"
	"time"
)

// TimeLayouts is contains a list of time.Parse() layouts that are used in
// attempts to convert item.Date and item.PubDate string to time.Time values.
// The layouts are attempted in ascending order until either time.Parse()
// does not return an error or all layouts are attempted.
var TimeLayouts = []string{
	"Mon, _2 Jan 2006 15:04:05 MST",
	"Mon, _2 Jan 2006 15:04:05 -0700",
	"Mon, 2 Jan 2006 15:04 MST",
	"01/02/2006 03:04:05 PM",
	time.ANSIC,
	time.UnixDate,
	time.RubyDate,
	time.RFC822,
	time.RFC822Z,
	time.RFC850,
	time.RFC1123,
	time.RFC1123Z,
	time.RFC3339,
	time.RFC3339Nano,
}

func parseTime(s string) (time.Time, error) {
	s = strings.TrimSpace(s)

	var e error
	var t time.Time

	for _, layout := range TimeLayouts {
		t, e = time.Parse(layout, s)
		if e == nil {
			return t, e
		}
	}

	return time.Time{}, e
}
<|endoftext|>


<commit_before>package gohost

import (
	"errors"
	"math"

	"github.com/grpc-ecosystem/go-grpc-middleware"
	"google.golang.org/grpc"
)

const (
	// DefaultMaxSendMsgSize is the default max send message size, per gRPC
	DefaultMaxSendMsgSize = 1024 * 1024 * 4

	// DefaultMaxRecvMsgSize is the default max receive message size, per gRPC
	DefaultMaxRecvMsgSize = 1024 * 1024 * 4
)

// Hoster is used to serve gRPC and HTTP endpoints.
type Hoster struct {
	Service            GRPCService
	GRPCAddr           string
	HTTPAddr           string
	CertFile           string
	KeyFile            string
	InsecureSkipVerify bool
	EnableCORS         bool
	MaxSendMsgSize     int
	MaxRecvMsgSize     int
	UnaryInterceptors  []grpc.UnaryServerInterceptor
	StreamInterceptors []grpc.StreamServerInterceptor
	Logger             func(format string, v ...interface{})
}

// NewHoster creates a new hoster instance with defaults set. This is the minimum required to host a server.
func NewHoster(service GRPCService, grpcAddr string) *Hoster {
	return &Hoster{
		Service:        service,
		GRPCAddr:       grpcAddr,
		MaxSendMsgSize: DefaultMaxSendMsgSize,
		MaxRecvMsgSize: DefaultMaxRecvMsgSize,
	}
}

// ListenAndServe creates and starts the server.
func (h *Hoster) ListenAndServe() error {
	// validate parameters
	if h.Service == nil {
		return errors.New("gRPC service implementation must be provided")
	}
	if h.GRPCAddr == "" {
		return errors.New("gRPC address must be provided")
	}

	// check if HTTP endpoint is enabled
	if h.HTTPAddr != "" {
		// ensure interface is implemented
		httpService, ok := h.Service.(HTTPService)
		if !ok {
			return errors.New("service does not implement HTTP interface")
		}

		// configure dial options
		dialOpts := []grpc.DialOption{
			grpc.WithDefaultCallOptions(grpc.MaxCallSendMsgSize(math.MaxInt32), grpc.MaxCallRecvMsgSize(math.MaxInt32)),
		}

		// start the HTTP endpoint
		if h.IsTLSEnabled() {
			h.log("Starting HTTP endpoint with TLS enabled: %v", h.HTTPAddr)
			go ServeHTTPWithTLS(httpService, h.HTTPAddr, h.GRPCAddr, h.EnableCORS, dialOpts, h.CertFile, h.KeyFile, h.InsecureSkipVerify)
		} else {
			h.log("Starting insecure HTTP endpoint: %v", h.HTTPAddr)
			go ServeHTTP(httpService, h.HTTPAddr, h.GRPCAddr, h.EnableCORS, dialOpts)
		}
	}

	// configure server options
	serverOpts := []grpc.ServerOption{
		grpc.MaxSendMsgSize(h.MaxSendMsgSize),
		grpc.MaxRecvMsgSize(h.MaxRecvMsgSize),
	}

	// add interceptors
	if len(h.UnaryInterceptors) > 0 {
		unaryInterceptorChain := grpc.UnaryInterceptor(grpc_middleware.ChainUnaryServer(h.UnaryInterceptors...))
		serverOpts = append(serverOpts, unaryInterceptorChain)
	}
	if len(h.StreamInterceptors) > 0 {
		streamInterceptorChain := grpc.StreamInterceptor(grpc_middleware.ChainStreamServer(h.StreamInterceptors...))
		serverOpts = append(serverOpts, streamInterceptorChain)
	}

	// start the gRPC endpoint
	if h.IsTLSEnabled() {
		h.log("Starting gRPC endpoint with TLS enabled: %v", h.GRPCAddr)
		return ServeGRPCWithTLS(h.Service, h.GRPCAddr, serverOpts, h.CertFile, h.KeyFile)
	}

	h.log("Starting insecure gRPC endpoint: %v", h.GRPCAddr)
	return ServeGRPC(h.Service, h.GRPCAddr, serverOpts)
}

// IsTLSEnabled will return true if TLS properties are set and ready to use.
func (h *Hoster) IsTLSEnabled() bool {
	return h.CertFile != "" && h.KeyFile != ""
}

// log will safely call the log function provided.
func (h *Hoster) log(format string, v ...interface{}) {
	if h.Logger != nil {
		h.Logger(format, v...)
	}
}
<commit_msg>Added comments to hoster properties.<commit_after>package gohost

import (
	"errors"
	"math"

	"github.com/grpc-ecosystem/go-grpc-middleware"
	"google.golang.org/grpc"
)

const (
	// DefaultMaxSendMsgSize is the default max send message size, per gRPC
	DefaultMaxSendMsgSize = 1024 * 1024 * 4

	// DefaultMaxRecvMsgSize is the default max receive message size, per gRPC
	DefaultMaxRecvMsgSize = 1024 * 1024 * 4
)

// Hoster is used to serve gRPC and HTTP endpoints.
type Hoster struct {
	// Service contains the actual implementation of the service calls. Additionally implement the HTTPService interface if an HTTP endpoint is desired.
	Service GRPCService

	// GRPCAddr is the endpoint (host and port) on which to host the gRPC service.
	GRPCAddr string

	// HTTPAddr is the endpoint (host and port) on which to host the HTTP service. May be left blank if not using HTTP.
	HTTPAddr string

	// CertFile is the certificate file for use with TLS. May be left blank if using insecure mode.
	CertFile string

	// KeyFile is the private key file for use with TLS. May be left blank if using insecure mode.
	KeyFile string

	// InsecureSkipVerify will cause verification of the host name during a TLS handshake to be skipped if set to true.
	InsecureSkipVerify bool

	// EnableCORS will enable all cross-origin resource sharing if set to true.
	EnableCORS bool

	// MaxSendMsgSize will change the size of the message that can be sent from the service.
	MaxSendMsgSize int

	// MaxRecvMsgSize will change the size of the message that can be received by the service.
	MaxRecvMsgSize int

	// UnaryInterceptors is an array of unary interceptors to be used by the service. They will be executed in order, from first to last.
	UnaryInterceptors []grpc.UnaryServerInterceptor

	// StreamInterceptors is an array of stream interceptors to be used by the service. They will be executed in order, from first to last.
	StreamInterceptors []grpc.StreamServerInterceptor

	// Logger is the logging method to be used for info and error logging by the hoster.
	Logger func(format string, v ...interface{})
}

// NewHoster creates a new hoster instance with defaults set. This is the minimum required to host a server.
func NewHoster(service GRPCService, grpcAddr string) *Hoster {
	return &Hoster{
		Service:        service,
		GRPCAddr:       grpcAddr,
		MaxSendMsgSize: DefaultMaxSendMsgSize,
		MaxRecvMsgSize: DefaultMaxRecvMsgSize,
	}
}

// ListenAndServe creates and starts the server.
func (h *Hoster) ListenAndServe() error {
	// validate parameters
	if h.Service == nil {
		return errors.New("gRPC service implementation must be provided")
	}
	if h.GRPCAddr == "" {
		return errors.New("gRPC address must be provided")
	}

	// check if HTTP endpoint is enabled
	if h.HTTPAddr != "" {
		// ensure interface is implemented
		httpService, ok := h.Service.(HTTPService)
		if !ok {
			return errors.New("service does not implement HTTP interface")
		}

		// configure dial options
		dialOpts := []grpc.DialOption{
			grpc.WithDefaultCallOptions(grpc.MaxCallSendMsgSize(math.MaxInt32), grpc.MaxCallRecvMsgSize(math.MaxInt32)),
		}

		// start the HTTP endpoint
		if h.IsTLSEnabled() {
			h.log("Starting HTTP endpoint with TLS enabled: %v", h.HTTPAddr)
			go ServeHTTPWithTLS(httpService, h.HTTPAddr, h.GRPCAddr, h.EnableCORS, dialOpts, h.CertFile, h.KeyFile, h.InsecureSkipVerify)
		} else {
			h.log("Starting insecure HTTP endpoint: %v", h.HTTPAddr)
			go ServeHTTP(httpService, h.HTTPAddr, h.GRPCAddr, h.EnableCORS, dialOpts)
		}
	}

	// configure server options
	serverOpts := []grpc.ServerOption{
		grpc.MaxSendMsgSize(h.MaxSendMsgSize),
		grpc.MaxRecvMsgSize(h.MaxRecvMsgSize),
	}

	// add interceptors
	if len(h.UnaryInterceptors) > 0 {
		unaryInterceptorChain := grpc.UnaryInterceptor(grpc_middleware.ChainUnaryServer(h.UnaryInterceptors...))
		serverOpts = append(serverOpts, unaryInterceptorChain)
	}
	if len(h.StreamInterceptors) > 0 {
		streamInterceptorChain := grpc.StreamInterceptor(grpc_middleware.ChainStreamServer(h.StreamInterceptors...))
		serverOpts = append(serverOpts, streamInterceptorChain)
	}

	// start the gRPC endpoint
	if h.IsTLSEnabled() {
		h.log("Starting gRPC endpoint with TLS enabled: %v", h.GRPCAddr)
		return ServeGRPCWithTLS(h.Service, h.GRPCAddr, serverOpts, h.CertFile, h.KeyFile)
	}

	h.log("Starting insecure gRPC endpoint: %v", h.GRPCAddr)
	return ServeGRPC(h.Service, h.GRPCAddr, serverOpts)
}

// IsTLSEnabled will return true if TLS properties are set and ready to use.
func (h *Hoster) IsTLSEnabled() bool {
	return h.CertFile != "" && h.KeyFile != ""
}

// log will safely call the log function provided.
func (h *Hoster) log(format string, v ...interface{}) {
	if h.Logger != nil {
		h.Logger(format, v...)
	}
}
<|endoftext|>


<commit_before>package httptracker

import (
	"bytes"
	"context"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"strconv"
	"time"

	"github.com/cenkalti/rain/internal/logger"
	"github.com/cenkalti/rain/torrent/internal/tracker"
	"github.com/zeebo/bencode"
)

// TODO get tracker timeout values from config
var httpTimeout = 30 * time.Second

type HTTPTracker struct {
	url       *url.URL
	log       logger.Logger
	http      *http.Client
	transport *http.Transport
	trackerID string
}

var _ tracker.Tracker = (*HTTPTracker)(nil)

func New(u *url.URL) *HTTPTracker {
	transport := &http.Transport{
		Dial: (&net.Dialer{
			Timeout: httpTimeout,
		}).Dial,
		TLSHandshakeTimeout: httpTimeout,
		DisableKeepAlives:   true,
	}
	return &HTTPTracker{
		url:       u,
		log:       logger.New("tracker " + u.String()),
		transport: transport,
		http: &http.Client{
			Timeout:   httpTimeout,
			Transport: transport,
		},
	}
}

func (t *HTTPTracker) Announce(ctx context.Context, req tracker.AnnounceRequest) (*tracker.AnnounceResponse, error) {
	transfer := req.Transfer
	e := req.Event
	numWant := req.NumWant
	peerID := transfer.PeerID
	infoHash := transfer.InfoHash
	q := url.Values{}
	q.Set("info_hash", string(infoHash[:]))
	q.Set("peer_id", string(peerID[:]))
	q.Set("port", strconv.FormatUint(uint64(transfer.Port), 10))
	q.Set("uploaded", strconv.FormatInt(transfer.BytesUploaded, 10))
	q.Set("downloaded", strconv.FormatInt(transfer.BytesDownloaded, 10))
	q.Set("left", strconv.FormatInt(transfer.BytesLeft, 10))
	q.Set("compact", "1")
	q.Set("no_peer_id", "1")
	q.Set("numwant", strconv.Itoa(numWant))
	if e != tracker.EventNone {
		q.Set("event", e.String())
	}
	if t.trackerID != "" {
		q.Set("trackerid", t.trackerID)
	}

	u := t.url
	u.RawQuery = q.Encode()
	t.log.Debugf("making request to: %q", u.String())

	httpReq := &http.Request{
		Method:     "GET",
		URL:        u,
		Proto:      "HTTP/1.1",
		ProtoMajor: 1,
		ProtoMinor: 1,
		Header:     make(http.Header),
		Host:       u.Host,
	}
	httpReq = httpReq.WithContext(ctx)

	do := func() ([]byte, error) {
		resp, err := t.http.Do(httpReq)
		if err != nil {
			return nil, err
		}
		defer resp.Body.Close()
		if resp.StatusCode != 200 {
			data, _ := ioutil.ReadAll(resp.Body)
			return nil, fmt.Errorf("status not 200 OK (status: %d body: %q)", resp.StatusCode, string(data))
		}
		return ioutil.ReadAll(resp.Body)
	}

	body, err := do()
	if uerr, ok := err.(*url.Error); ok && uerr.Err == context.Canceled {
		return nil, context.Canceled
	}

	var response announceResponse
	err = bencode.DecodeBytes(body, &response)
	if err != nil {
		return nil, err
	}

	if response.WarningMessage != "" {
		t.log.Warning(response.WarningMessage)
	}
	if response.FailureReason != "" {
		return nil, tracker.Error(response.FailureReason)
	}

	if response.TrackerID != "" {
		t.trackerID = response.TrackerID
	}

	// Peers may be in binary or dictionary model.
	var peers []*net.TCPAddr
	if len(response.Peers) > 0 {
		if response.Peers[0] == 'l' {
			peers, err = t.parsePeersDictionary(response.Peers)
		} else {
			var b []byte
			err = bencode.DecodeBytes(response.Peers, &b)
			if err != nil {
				return nil, err
			}
			peers, err = tracker.ParsePeersBinary(bytes.NewReader(b), t.log)
		}
	}
	if err != nil {
		return nil, err
	}

	// Filter external IP
	if len(response.ExternalIP) != 0 {
		for i, p := range peers {
			if bytes.Equal(p.IP[:], response.ExternalIP) {
				peers[i], peers = peers[len(peers)-1], peers[:len(peers)-1]
				break
			}
		}
	}

	return &tracker.AnnounceResponse{
		Interval:   time.Duration(response.Interval) * time.Second,
		Leechers:   response.Incomplete,
		Seeders:    response.Complete,
		Peers:      peers,
		ExternalIP: response.ExternalIP,
	}, nil
}

func (t *HTTPTracker) parsePeersDictionary(b bencode.RawMessage) ([]*net.TCPAddr, error) {
	var peers []struct {
		IP   string `bencode:"ip"`
		Port uint16 `bencode:"port"`
	}
	err := bencode.DecodeBytes(b, &peers)
	if err != nil {
		return nil, err
	}

	addrs := make([]*net.TCPAddr, len(peers))
	for i, p := range peers {
		pe := &net.TCPAddr{IP: net.ParseIP(p.IP), Port: int(p.Port)}
		addrs[i] = pe
	}
	return addrs, err
}

func (t *HTTPTracker) Close() error {
	t.transport.CloseIdleConnections()
	return nil
}
<commit_msg>do not eat tracker query params<commit_after>package httptracker

import (
	"bytes"
	"context"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"strconv"
	"time"

	"github.com/cenkalti/rain/internal/logger"
	"github.com/cenkalti/rain/torrent/internal/tracker"
	"github.com/zeebo/bencode"
)

// TODO get tracker timeout values from config
var httpTimeout = 30 * time.Second

type HTTPTracker struct {
	url       *url.URL
	log       logger.Logger
	http      *http.Client
	transport *http.Transport
	trackerID string
}

var _ tracker.Tracker = (*HTTPTracker)(nil)

func New(u *url.URL) *HTTPTracker {
	transport := &http.Transport{
		Dial: (&net.Dialer{
			Timeout: httpTimeout,
		}).Dial,
		TLSHandshakeTimeout: httpTimeout,
		DisableKeepAlives:   true,
	}
	return &HTTPTracker{
		url:       u,
		log:       logger.New("tracker " + u.String()),
		transport: transport,
		http: &http.Client{
			Timeout:   httpTimeout,
			Transport: transport,
		},
	}
}

func (t *HTTPTracker) Announce(ctx context.Context, req tracker.AnnounceRequest) (*tracker.AnnounceResponse, error) {
	transfer := req.Transfer
	e := req.Event
	numWant := req.NumWant
	peerID := transfer.PeerID
	infoHash := transfer.InfoHash
	q := t.url.Query()
	q.Set("info_hash", string(infoHash[:]))
	q.Set("peer_id", string(peerID[:]))
	q.Set("port", strconv.FormatUint(uint64(transfer.Port), 10))
	q.Set("uploaded", strconv.FormatInt(transfer.BytesUploaded, 10))
	q.Set("downloaded", strconv.FormatInt(transfer.BytesDownloaded, 10))
	q.Set("left", strconv.FormatInt(transfer.BytesLeft, 10))
	q.Set("compact", "1")
	q.Set("no_peer_id", "1")
	q.Set("numwant", strconv.Itoa(numWant))
	if e != tracker.EventNone {
		q.Set("event", e.String())
	}
	if t.trackerID != "" {
		q.Set("trackerid", t.trackerID)
	}

	u := t.url
	u.RawQuery = q.Encode()
	t.log.Debugf("making request to: %q", u.String())

	httpReq := &http.Request{
		Method:     "GET",
		URL:        u,
		Proto:      "HTTP/1.1",
		ProtoMajor: 1,
		ProtoMinor: 1,
		Header:     make(http.Header),
		Host:       u.Host,
	}
	httpReq = httpReq.WithContext(ctx)

	do := func() ([]byte, error) {
		resp, err := t.http.Do(httpReq)
		if err != nil {
			return nil, err
		}
		defer resp.Body.Close()
		if resp.StatusCode != 200 {
			data, _ := ioutil.ReadAll(resp.Body)
			return nil, fmt.Errorf("status not 200 OK (status: %d body: %q)", resp.StatusCode, string(data))
		}
		return ioutil.ReadAll(resp.Body)
	}

	body, err := do()
	if uerr, ok := err.(*url.Error); ok && uerr.Err == context.Canceled {
		return nil, context.Canceled
	}

	var response announceResponse
	err = bencode.DecodeBytes(body, &response)
	if err != nil {
		return nil, err
	}

	if response.WarningMessage != "" {
		t.log.Warning(response.WarningMessage)
	}
	if response.FailureReason != "" {
		return nil, tracker.Error(response.FailureReason)
	}

	if response.TrackerID != "" {
		t.trackerID = response.TrackerID
	}

	// Peers may be in binary or dictionary model.
	var peers []*net.TCPAddr
	if len(response.Peers) > 0 {
		if response.Peers[0] == 'l' {
			peers, err = t.parsePeersDictionary(response.Peers)
		} else {
			var b []byte
			err = bencode.DecodeBytes(response.Peers, &b)
			if err != nil {
				return nil, err
			}
			peers, err = tracker.ParsePeersBinary(bytes.NewReader(b), t.log)
		}
	}
	if err != nil {
		return nil, err
	}

	// Filter external IP
	if len(response.ExternalIP) != 0 {
		for i, p := range peers {
			if bytes.Equal(p.IP[:], response.ExternalIP) {
				peers[i], peers = peers[len(peers)-1], peers[:len(peers)-1]
				break
			}
		}
	}

	return &tracker.AnnounceResponse{
		Interval:   time.Duration(response.Interval) * time.Second,
		Leechers:   response.Incomplete,
		Seeders:    response.Complete,
		Peers:      peers,
		ExternalIP: response.ExternalIP,
	}, nil
}

func (t *HTTPTracker) parsePeersDictionary(b bencode.RawMessage) ([]*net.TCPAddr, error) {
	var peers []struct {
		IP   string `bencode:"ip"`
		Port uint16 `bencode:"port"`
	}
	err := bencode.DecodeBytes(b, &peers)
	if err != nil {
		return nil, err
	}

	addrs := make([]*net.TCPAddr, len(peers))
	for i, p := range peers {
		pe := &net.TCPAddr{IP: net.ParseIP(p.IP), Port: int(p.Port)}
		addrs[i] = pe
	}
	return addrs, err
}

func (t *HTTPTracker) Close() error {
	t.transport.CloseIdleConnections()
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2012 Arne Roomann-Kurrik
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package fauxfile

import (
	"testing"
)

func ExpectCwd(t *testing.T, expected string, mf *MockFilesystem) {
	if mf.cwd != expected {
		t.Fatalf("Expected cwd of '%v', got '%v'", expected, mf.cwd)
	}
}

func TestChdir(t *testing.T) {
	mf := NewMockFilesystem()
	mf.Mkdir("/foo", 0755)
	ExpectCwd(t, "/", mf)
	mf.Chdir("foo")
	ExpectCwd(t, "/foo", mf)
}

<commit_msg>Added a mkdir test<commit_after>// Copyright 2012 Arne Roomann-Kurrik
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package fauxfile

import (
	"testing"
)

func ExpectCwd(t *testing.T, expected string, mf *MockFilesystem) {
	if mf.cwd != expected {
		t.Fatalf("Expected cwd of '%v', got '%v'", expected, mf.cwd)
	}
}

func ExpectDir(t *testing.T, path string, mf *MockFilesystem) {
	mode, exists := mf.dirs[path]
	if !exists {
		t.Fatalf("Expected path of '%s' to be present", path)
	}
	if !mode.IsDir() {
		t.Fatalf("Expected '%v' to be directory", path)
	}
}

func TestChdir(t *testing.T) {
	mf := NewMockFilesystem()
	mf.Mkdir("/foo", 0755)
	ExpectCwd(t, "/", mf)
	mf.Chdir("foo")
	ExpectCwd(t, "/foo", mf)
}

func TestMkdir(t *testing.T) {
	mf := NewMockFilesystem()
	mf.Mkdir("/foo", 0755)
	mf.Mkdir("/bar", 0777)
	mf.Chdir("bar")
	mf.Mkdir("baz", 0777)
	ExpectDir(t, "/foo", mf)
	ExpectDir(t, "/bar", mf)
	ExpectDir(t, "/bar/baz", mf)
}

<|endoftext|>