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

package main

import (
	"bytes"
	"encoding/hex"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/pkg/errors"
	"github.com/sdboyer/gps"
)

const ensureShortHelp = `Ensure a dependency is safely vendored in the project`
const ensureLongHelp = `
Ensure is used to fetch project dependencies into the vendor folder, as well as
to set version constraints for specific dependencies. It takes user input,
solves the updated dependency graph of the project, writes any changes to the
manifest and lock file, and places dependencies in the vendor folder.

Package spec:

  <path>[:alt location][@<version specifier>]

Examples:

  dep ensure                                   Populate vendor from existing manifest and lock
  dep ensure github.com/heroku/rollrus@^0.9.1  Update a specific dependency to a specific version

For more detailed usage examples, see dep ensure -examples.
`
const ensureExamples = `
dep ensure

    Solve the project's dependency graph, and place all dependencies in the
    vendor folder. If a dependency is in the lock file, use the version
    specified there. Otherwise, use the most recent version that can satisfy the
    constraints in the manifest file.

dep ensure github.com/heroku/rollrus@^0.9.1

    Same as above, but choose any release >= 0.9.1, < 1.0.0. This form of
    constraint strikes a good balance of safety and flexibility, and should be
    preferred for libraries.

dep ensure github.com/heroku/rollrus@~0.9.0

    Same as above, but choose any release matching 0.9.x, preferring latest. If
    a constraint was previously set in the manifest, this resets it.

dep ensure github.com/heroku/rollrus:git.internal.com/foo/bar

    Fetch the dependency from a different location.

dep ensure -override github.com/heroku/rollrus@^0.9.1

    Forcefully and transitively override any constraint for this dependency.
    This can inadvertantly make your dependency graph unsolvable; use sparingly.
`

func (cmd *ensureCommand) Name() string      { return "ensure" }
func (cmd *ensureCommand) Args() string      { return "[spec...]" }
func (cmd *ensureCommand) ShortHelp() string { return ensureShortHelp }
func (cmd *ensureCommand) LongHelp() string  { return ensureLongHelp }

func (cmd *ensureCommand) Register(fs *flag.FlagSet) {
	fs.BoolVar(&cmd.examples, "examples", false, "print detailed usage examples")
	fs.BoolVar(&cmd.update, "update", false, "ensure all dependencies are at the latest version allowed by the manifest")
	fs.BoolVar(&cmd.dryRun, "n", false, "dry run, don't actually ensure anything")
	fs.Var(&cmd.overrides, "override", "specify an override constraint spec (repeatable)")
}

type ensureCommand struct {
	examples  bool
	update    bool
	dryRun    bool
	overrides stringSlice
}

func (cmd *ensureCommand) Run(args []string) error {
	if cmd.examples {
		fmt.Fprintln(os.Stderr, strings.TrimSpace(ensureExamples))
		return nil
	}

	p, err := depContext.loadProject("")
	if err != nil {
		return err
	}

	sm, err := depContext.sourceManager()
	if err != nil {
		return err
	}
	defer sm.Release()

	var errs []error
	for _, arg := range args {
		// default persist to manifest
		pc, err := getProjectConstraint(arg, sm)
		if err != nil {
			errs = append(errs, err)
		}

		if gps.IsAny(pc.Constraint) && pc.Ident.NetworkName == "" {
			// If the input specified neither a network name nor a constraint,
			// then the strict thing to do would be to remove the entry
			// entirely. But that would probably be quite surprising for users,
			// and it's what rm is for, so just ignore the input.
			//
			// TODO(sdboyer): for this case - or just in general - do we want to
			// add project args to the requires list temporarily for this run?
			if _, has := p.m.Dependencies[pc.Ident.ProjectRoot]; !has {
				logf("No constraint or alternate source specified for %q, omitting from manifest", pc.Ident.ProjectRoot)
			}
			// If it's already in the manifest, no need to log
			continue
		}

		p.m.Dependencies[pc.Ident.ProjectRoot] = gps.ProjectProperties{
			NetworkName: pc.Ident.NetworkName,
			Constraint:  pc.Constraint,
		}

		for i, lp := range p.l.P {
			if lp.Ident() == pc.Ident {
				p.l.P = append(p.l.P[:i], p.l.P[i+1:]...)
				break
			}
		}
	}

	for _, ovr := range cmd.overrides {
		pc, err := getProjectConstraint(ovr, sm)
		if err != nil {
			errs = append(errs, err)
		}

		// Empty overrides are fine (in contrast to deps), because they actually
		// carry meaning - they force the constraints entirely open for a given
		// project. Inadvisable, but meaningful.

		p.m.Ovr[pc.Ident.ProjectRoot] = gps.ProjectProperties{
			NetworkName: pc.Ident.NetworkName,
			Constraint:  pc.Constraint,
		}

		for i, lp := range p.l.P {
			if lp.Ident() == pc.Ident {
				p.l.P = append(p.l.P[:i], p.l.P[i+1:]...)
				break
			}
		}
	}

	if len(errs) > 0 {
		var buf bytes.Buffer
		for err := range errs {
			fmt.Fprintln(&buf, err)
		}

		return errors.New(buf.String())
	}

	params := gps.SolveParameters{
		RootDir:  p.absroot,
		Manifest: p.m,
		Lock:     p.l,
	}
	if *verbose {
		params.Trace = true
		params.TraceLogger = log.New(os.Stderr, "", 0)
	}

	params.RootPackageTree, err = gps.ListPackages(p.absroot, string(p.importroot))
	if err != nil {
		return errors.Wrap(err, "ensure ListPackage for project")
	}
	solver, err := gps.Prepare(params, sm)
	if err != nil {
		return errors.Wrap(err, "ensure Prepare")
	}
	solution, err := solver.Solve()
	if err != nil {
		handleAllTheFailuresOfTheWorld(err)
		return errors.Wrap(err, "ensure Solve()")
	}

	p.l.P = solution.Projects()
	p.l.Memo = solution.InputHash()

	tv, err := ioutil.TempDir("", "vendor")
	if err != nil {
		return errors.Wrap(err, "ensure making temporary vendor")
	}
	defer os.RemoveAll(tv)

	tm, err := ioutil.TempFile("", "manifest")
	if err != nil {
		return errors.Wrap(err, "ensure making temporary manifest")
	}
	tm.Close()
	defer os.Remove(tm.Name())

	tl, err := ioutil.TempFile("", "lock")
	if err != nil {
		return errors.Wrap(err, "ensure making temporary lock file")
	}
	tl.Close()
	defer os.Remove(tl.Name())

	if err := gps.WriteDepTree(tv, p.l, sm, true); err != nil {
		return errors.Wrap(err, "ensure gps.WriteDepTree")
	}

	if err := writeFile(tm.Name(), p.m); err != nil {
		return errors.Wrap(err, "ensure writeFile for manifest")
	}

	if err := writeFile(tl.Name(), p.l); err != nil {
		return errors.Wrap(err, "ensure writeFile for lock")
	}

	if err := copyFile(tm.Name(), filepath.Join(p.absroot, manifestName)); err != nil {
		return errors.Wrap(err, "ensure moving temp manifest into place!")
	}
	os.Remove(tm.Name())

	if err := copyFile(tl.Name(), filepath.Join(p.absroot, lockName)); err != nil {
		return errors.Wrap(err, "ensure moving temp manifest into place!")
	}
	os.Remove(tl.Name())

	os.RemoveAll(filepath.Join(p.absroot, "vendor"))
	if err := copyFolder(tv, filepath.Join(p.absroot, "vendor")); err != nil {
		return errors.Wrap(err, "ensure moving temp vendor")
	}

	return nil
}

type stringSlice []string

func (s *stringSlice) String() string {
	if len(*s) == 0 {
		return "<none>"
	}
	return strings.Join(*s, ", ")
}

func (s *stringSlice) Set(value string) error {
	*s = append(*s, value)
	return nil
}

func getProjectConstraint(arg string, sm *gps.SourceMgr) (gps.ProjectConstraint, error) {
	constraint := gps.ProjectConstraint{
		Constraint: gps.Any(), // default to any; avoids panics later
	}

	// try to split on '@'
	atIndex := strings.Index(arg, "@")
	if atIndex > 0 {
		parts := strings.SplitN(arg, "@", 2)
		constraint.Constraint = deduceConstraint(parts[1])
		arg = parts[0]
	}
	// TODO: What if there is no @, assume default branch (which may not be master) ?
	// TODO: if we decide to keep equals.....

	// split on colon if there is a network location
	colonIndex := strings.Index(arg, ":")
	if colonIndex > 0 {
		parts := strings.SplitN(arg, ":", 2)
		arg = parts[0]
		constraint.Ident.NetworkName = parts[1]
	}

	pr, err := sm.DeduceProjectRoot(arg)
	if err != nil {
		return constraint, errors.Wrapf(err, "could not infer project root from dependency path: %s", arg) // this should go through to the user
	}

	if string(pr) != arg {
		return constraint, errors.Wrapf(err, "dependency path %s is not a project root", arg)
	}
	constraint.Ident.ProjectRoot = gps.ProjectRoot(arg)

	return constraint, nil
}

// deduceConstraint tries to puzzle out what kind of version is given in a string -
// semver, a revision, or as a fallback, a plain tag
func deduceConstraint(s string) gps.Constraint {
	// always semver if we can
	c, err := gps.NewSemverConstraint(s)
	if err == nil {
		return c
	}

	slen := len(s)
	if slen == 40 {
		if _, err = hex.DecodeString(s); err == nil {
			// Whether or not it's intended to be a SHA1 digest, this is a
			// valid byte sequence for that, so go with Revision. This
			// covers git and hg
			return gps.Revision(s)
		}
	}
	// Next, try for bzr, which has a three-component GUID separated by
	// dashes. There should be two, but the email part could contain
	// internal dashes
	if strings.Count(s, "-") >= 2 {
		// Work from the back to avoid potential confusion from the email
		i3 := strings.LastIndex(s, "-")
		// Skip if - is last char, otherwise this would panic on bounds err
		if slen == i3+1 {
			return gps.NewVersion(s)
		}

		i2 := strings.LastIndex(s[:i3], "-")
		if _, err = strconv.ParseUint(s[i2+1:i3], 10, 64); err == nil {
			// Getting this far means it'd pretty much be nuts if it's not a
			// bzr rev, so don't bother parsing the email.
			return gps.Revision(s)
		}
	}

	// If not a plain SHA1 or bzr custom GUID, assume a plain version.
	// TODO: if there is amgibuity here, then prompt the user?
	return gps.NewVersion(s)
}

// copyFolder takes in a directory and copies its contents to the destination.
// It preserves the file mode on files as well.
func copyFolder(src string, dest string) error {
	fi, err := os.Lstat(src)
	if err != nil {
		return err
	}

	err = os.MkdirAll(dest, fi.Mode())
	if err != nil {
		return err
	}

	dir, err := os.Open(src)
	if err != nil {
		return err
	}
	defer dir.Close()

	objects, err := dir.Readdir(-1)
	if err != nil {
		return err
	}

	for _, obj := range objects {
		if obj.Mode()&os.ModeSymlink != 0 {
			continue
		}

		srcfile := filepath.Join(src, obj.Name())
		destfile := filepath.Join(dest, obj.Name())

		if obj.IsDir() {
			err = copyFolder(srcfile, destfile)
			if err != nil {
				return err
			}
			continue
		}

		if err := copyFile(srcfile, destfile); err != nil {
			return err
		}
	}

	return nil
}

// copyFile copies a file from one place to another with the permission bits
// perserved as well.
func copyFile(src string, dest string) error {
	srcfile, err := os.Open(src)
	if err != nil {
		return err
	}
	defer srcfile.Close()

	destfile, err := os.Create(dest)
	if err != nil {
		return err
	}
	defer destfile.Close()

	if _, err := io.Copy(destfile, srcfile); err != nil {
		return err
	}

	srcinfo, err := os.Stat(src)
	if err != nil {
		return err
	}

	return os.Chmod(dest, srcinfo.Mode())
}
<commit_msg>Further review feedback<commit_after>// Copyright 2016 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"bytes"
	"encoding/hex"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/pkg/errors"
	"github.com/sdboyer/gps"
)

const ensureShortHelp = `Ensure a dependency is safely vendored in the project`
const ensureLongHelp = `
Ensure is used to fetch project dependencies into the vendor folder, as well as
to set version constraints for specific dependencies. It takes user input,
solves the updated dependency graph of the project, writes any changes to the
manifest and lock file, and places dependencies in the vendor folder.

Package spec:

  <path>[:alt location][@<version specifier>]

Examples:

  dep ensure                            Populate vendor from existing manifest and lock
  dep ensure github.com/pkg/foo@^1.0.1  Update a specific dependency to a specific version

For more detailed usage examples, see dep ensure -examples.
`
const ensureExamples = `
dep ensure

    Solve the project's dependency graph, and place all dependencies in the
    vendor folder. If a dependency is in the lock file, use the version
    specified there. Otherwise, use the most recent version that can satisfy the
    constraints in the manifest file.

dep ensure github.com/pkg/foo@^1.0.1

    Same as above, but choose any release >= 1.0.1, < 2.0.0. If a constraint was
    previously set in the manifest, this resets it. This form of constraint
    strikes a good balance of safety and flexibility, and should be preferred
    for libraries.

dep ensure github.com/pkg/foo@~1.0.1

    Same as above, but choose any release matching 1.0.x, preferring latest.

dep ensure github.com/pkg/foo:git.internal.com/alt/foo

    Fetch the dependency from a different location.

dep ensure -override github.com/pkg/foo@^1.0.1

    Forcefully and transitively override any constraint for this dependency.
    This can inadvertantly make your dependency graph unsolvable; use sparingly.
`

func (cmd *ensureCommand) Name() string      { return "ensure" }
func (cmd *ensureCommand) Args() string      { return "[spec...]" }
func (cmd *ensureCommand) ShortHelp() string { return ensureShortHelp }
func (cmd *ensureCommand) LongHelp() string  { return ensureLongHelp }

func (cmd *ensureCommand) Register(fs *flag.FlagSet) {
	fs.BoolVar(&cmd.examples, "examples", false, "print detailed usage examples")
	fs.BoolVar(&cmd.update, "update", false, "ensure all dependencies are at the latest version allowed by the manifest")
	fs.BoolVar(&cmd.dryRun, "n", false, "dry run, don't actually ensure anything")
	fs.Var(&cmd.overrides, "override", "specify an override constraint spec (repeatable)")
}

type ensureCommand struct {
	examples  bool
	update    bool
	dryRun    bool
	overrides stringSlice
}

func (cmd *ensureCommand) Run(args []string) error {
	if cmd.examples {
		fmt.Fprintln(os.Stderr, strings.TrimSpace(ensureExamples))
		return nil
	}

	p, err := depContext.loadProject("")
	if err != nil {
		return err
	}

	sm, err := depContext.sourceManager()
	if err != nil {
		return err
	}
	defer sm.Release()

	var errs []error
	for _, arg := range args {
		// default persist to manifest
		pc, err := getProjectConstraint(arg, sm)
		if err != nil {
			errs = append(errs, err)
		}

		if gps.IsAny(pc.Constraint) && pc.Ident.NetworkName == "" {
			// If the input specified neither a network name nor a constraint,
			// then the strict thing to do would be to remove the entry
			// entirely. But that would probably be quite surprising for users,
			// and it's what rm is for, so just ignore the input.
			//
			// TODO(sdboyer): for this case - or just in general - do we want to
			// add project args to the requires list temporarily for this run?
			if _, has := p.m.Dependencies[pc.Ident.ProjectRoot]; !has {
				logf("No constraint or alternate source specified for %q, omitting from manifest", pc.Ident.ProjectRoot)
			}
			// If it's already in the manifest, no need to log
			continue
		}

		p.m.Dependencies[pc.Ident.ProjectRoot] = gps.ProjectProperties{
			NetworkName: pc.Ident.NetworkName,
			Constraint:  pc.Constraint,
		}

		for i, lp := range p.l.P {
			if lp.Ident() == pc.Ident {
				p.l.P = append(p.l.P[:i], p.l.P[i+1:]...)
				break
			}
		}
	}

	for _, ovr := range cmd.overrides {
		pc, err := getProjectConstraint(ovr, sm)
		if err != nil {
			errs = append(errs, err)
		}

		// Empty overrides are fine (in contrast to deps), because they actually
		// carry meaning - they force the constraints entirely open for a given
		// project. Inadvisable, but meaningful.

		p.m.Ovr[pc.Ident.ProjectRoot] = gps.ProjectProperties{
			NetworkName: pc.Ident.NetworkName,
			Constraint:  pc.Constraint,
		}

		for i, lp := range p.l.P {
			if lp.Ident() == pc.Ident {
				p.l.P = append(p.l.P[:i], p.l.P[i+1:]...)
				break
			}
		}
	}

	if len(errs) > 0 {
		var buf bytes.Buffer
		for err := range errs {
			fmt.Fprintln(&buf, err)
		}

		return errors.New(buf.String())
	}

	params := gps.SolveParameters{
		RootDir:  p.absroot,
		Manifest: p.m,
		Lock:     p.l,
	}
	if *verbose {
		params.Trace = true
		params.TraceLogger = log.New(os.Stderr, "", 0)
	}

	params.RootPackageTree, err = gps.ListPackages(p.absroot, string(p.importroot))
	if err != nil {
		return errors.Wrap(err, "ensure ListPackage for project")
	}
	solver, err := gps.Prepare(params, sm)
	if err != nil {
		return errors.Wrap(err, "ensure Prepare")
	}
	solution, err := solver.Solve()
	if err != nil {
		handleAllTheFailuresOfTheWorld(err)
		return errors.Wrap(err, "ensure Solve()")
	}

	p.l.P = solution.Projects()
	p.l.Memo = solution.InputHash()

	tv, err := ioutil.TempDir("", "vendor")
	if err != nil {
		return errors.Wrap(err, "ensure making temporary vendor")
	}
	defer os.RemoveAll(tv)

	tm, err := ioutil.TempFile("", "manifest")
	if err != nil {
		return errors.Wrap(err, "ensure making temporary manifest")
	}
	tm.Close()
	defer os.Remove(tm.Name())

	tl, err := ioutil.TempFile("", "lock")
	if err != nil {
		return errors.Wrap(err, "ensure making temporary lock file")
	}
	tl.Close()
	defer os.Remove(tl.Name())

	if err := gps.WriteDepTree(tv, p.l, sm, true); err != nil {
		return errors.Wrap(err, "ensure gps.WriteDepTree")
	}

	if err := writeFile(tm.Name(), p.m); err != nil {
		return errors.Wrap(err, "ensure writeFile for manifest")
	}

	if err := writeFile(tl.Name(), p.l); err != nil {
		return errors.Wrap(err, "ensure writeFile for lock")
	}

	if err := copyFile(tm.Name(), filepath.Join(p.absroot, manifestName)); err != nil {
		return errors.Wrap(err, "ensure moving temp manifest into place!")
	}
	os.Remove(tm.Name())

	if err := copyFile(tl.Name(), filepath.Join(p.absroot, lockName)); err != nil {
		return errors.Wrap(err, "ensure moving temp manifest into place!")
	}
	os.Remove(tl.Name())

	os.RemoveAll(filepath.Join(p.absroot, "vendor"))
	if err := copyFolder(tv, filepath.Join(p.absroot, "vendor")); err != nil {
		return errors.Wrap(err, "ensure moving temp vendor")
	}

	return nil
}

type stringSlice []string

func (s *stringSlice) String() string {
	if len(*s) == 0 {
		return "<none>"
	}
	return strings.Join(*s, ", ")
}

func (s *stringSlice) Set(value string) error {
	*s = append(*s, value)
	return nil
}

func getProjectConstraint(arg string, sm *gps.SourceMgr) (gps.ProjectConstraint, error) {
	constraint := gps.ProjectConstraint{
		Constraint: gps.Any(), // default to any; avoids panics later
	}

	// try to split on '@'
	atIndex := strings.Index(arg, "@")
	if atIndex > 0 {
		parts := strings.SplitN(arg, "@", 2)
		constraint.Constraint = deduceConstraint(parts[1])
		arg = parts[0]
	}
	// TODO: What if there is no @, assume default branch (which may not be master) ?
	// TODO: if we decide to keep equals.....

	// split on colon if there is a network location
	colonIndex := strings.Index(arg, ":")
	if colonIndex > 0 {
		parts := strings.SplitN(arg, ":", 2)
		arg = parts[0]
		constraint.Ident.NetworkName = parts[1]
	}

	pr, err := sm.DeduceProjectRoot(arg)
	if err != nil {
		return constraint, errors.Wrapf(err, "could not infer project root from dependency path: %s", arg) // this should go through to the user
	}

	if string(pr) != arg {
		return constraint, errors.Wrapf(err, "dependency path %s is not a project root", arg)
	}
	constraint.Ident.ProjectRoot = gps.ProjectRoot(arg)

	return constraint, nil
}

// deduceConstraint tries to puzzle out what kind of version is given in a string -
// semver, a revision, or as a fallback, a plain tag
func deduceConstraint(s string) gps.Constraint {
	// always semver if we can
	c, err := gps.NewSemverConstraint(s)
	if err == nil {
		return c
	}

	slen := len(s)
	if slen == 40 {
		if _, err = hex.DecodeString(s); err == nil {
			// Whether or not it's intended to be a SHA1 digest, this is a
			// valid byte sequence for that, so go with Revision. This
			// covers git and hg
			return gps.Revision(s)
		}
	}
	// Next, try for bzr, which has a three-component GUID separated by
	// dashes. There should be two, but the email part could contain
	// internal dashes
	if strings.Count(s, "-") >= 2 {
		// Work from the back to avoid potential confusion from the email
		i3 := strings.LastIndex(s, "-")
		// Skip if - is last char, otherwise this would panic on bounds err
		if slen == i3+1 {
			return gps.NewVersion(s)
		}

		i2 := strings.LastIndex(s[:i3], "-")
		if _, err = strconv.ParseUint(s[i2+1:i3], 10, 64); err == nil {
			// Getting this far means it'd pretty much be nuts if it's not a
			// bzr rev, so don't bother parsing the email.
			return gps.Revision(s)
		}
	}

	// If not a plain SHA1 or bzr custom GUID, assume a plain version.
	// TODO: if there is amgibuity here, then prompt the user?
	return gps.NewVersion(s)
}

// copyFolder takes in a directory and copies its contents to the destination.
// It preserves the file mode on files as well.
func copyFolder(src string, dest string) error {
	fi, err := os.Lstat(src)
	if err != nil {
		return err
	}

	err = os.MkdirAll(dest, fi.Mode())
	if err != nil {
		return err
	}

	dir, err := os.Open(src)
	if err != nil {
		return err
	}
	defer dir.Close()

	objects, err := dir.Readdir(-1)
	if err != nil {
		return err
	}

	for _, obj := range objects {
		if obj.Mode()&os.ModeSymlink != 0 {
			continue
		}

		srcfile := filepath.Join(src, obj.Name())
		destfile := filepath.Join(dest, obj.Name())

		if obj.IsDir() {
			err = copyFolder(srcfile, destfile)
			if err != nil {
				return err
			}
			continue
		}

		if err := copyFile(srcfile, destfile); err != nil {
			return err
		}
	}

	return nil
}

// copyFile copies a file from one place to another with the permission bits
// perserved as well.
func copyFile(src string, dest string) error {
	srcfile, err := os.Open(src)
	if err != nil {
		return err
	}
	defer srcfile.Close()

	destfile, err := os.Create(dest)
	if err != nil {
		return err
	}
	defer destfile.Close()

	if _, err := io.Copy(destfile, srcfile); err != nil {
		return err
	}

	srcinfo, err := os.Stat(src)
	if err != nil {
		return err
	}

	return os.Chmod(dest, srcinfo.Mode())
}
<|endoftext|>


<commit_before>package program

// Version can be requested through the command line with:
//
//     c2go -v
//
// See https://github.com/elliotchance/c2go/wiki/Release-Process
const Version = "v0.21.6 Zinc 2018-02-07"
<commit_msg>Bump version: v0.21.7 Zinc 2018-02-08<commit_after>package program

// Version can be requested through the command line with:
//
//     c2go -v
//
// See https://github.com/elliotchance/c2go/wiki/Release-Process
const Version = "v0.21.7 Zinc 2018-02-08"
<|endoftext|>


<commit_before>package binder

import (
	"net/http"
	"reflect"

	"github.com/SoCloz/binder/response"
)

// Action Parameters
type ActionArgs struct {
	Name string
	Type reflect.Type
}

// A controller action wrapper
type Wrapper struct {
	Call       reflect.Value
	IsVariadic bool
	Args       []*ActionArgs
}

// Wraps a controller action
//
// Example :
//   func MyAction(id int, param string) {}
//   http.Handle("/my_action", binder.Wrap(MyAction, "id", "param"))
func Wrap(call interface{}, params ...string) *Wrapper {
	w := new(Wrapper)
	w.Call = reflect.ValueOf(call)
	callType := w.Call.Type()
	w.IsVariadic = callType.IsVariadic()
	w.Args = make([]*ActionArgs, callType.NumIn())

	if callType.NumIn() < len(params) {
		panic("Wrong number of params")
	}
	for i := 0; i < callType.NumIn(); i++ {
		typ := callType.In(i)
		var paramName string
		if i < len(params) {
			paramName = params[i]
		} else {
			paramName = "*"
		}
		w.Args[i] = &ActionArgs{Name: paramName, Type: typ}
		if paramName == "*" {
			RegisterStructAttributes(typ)
		}
	}
	return w
}

// The http handler for the wrapped action
func (wr *Wrapper) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	values := r.URL.Query()

	methodArgs := make([]reflect.Value, len(wr.Args))
	for i, a := range wr.Args {
		methodArgs[i] = GetBoundValue(values, a.Name, a.Type)
	}
	var resultValue reflect.Value
	if wr.IsVariadic {
		resultValue = wr.Call.CallSlice(methodArgs)[0]
	} else {
		resultValue = wr.Call.Call(methodArgs)[0]
	}
	resp := resultValue.Interface().(response.Response)
	resp.ApplyTo(w)
}
<commit_msg>comment formatting<commit_after>package binder

import (
	"net/http"
	"reflect"

	"github.com/SoCloz/binder/response"
)

// Action Parameters
type ActionArgs struct {
	Name string
	Type reflect.Type
}

// A controller action wrapper
type Wrapper struct {
	Call       reflect.Value
	IsVariadic bool
	Args       []*ActionArgs
}

// Wraps a controller action
//
// Example :
//   func MyAction(id int, param string) {}
//   http.Handle("/my_action", binder.Wrap(MyAction, "id", "param")
func Wrap(call interface{}, params ...string) *Wrapper {
	w := new(Wrapper)
	w.Call = reflect.ValueOf(call)
	callType := w.Call.Type()
	w.IsVariadic = callType.IsVariadic()
	w.Args = make([]*ActionArgs, callType.NumIn())

	if callType.NumIn() < len(params) {
		panic("Wrong number of params")
	}
	for i := 0; i < callType.NumIn(); i++ {
		typ := callType.In(i)
		var paramName string
		if i < len(params) {
			paramName = params[i]
		} else {
			paramName = "*"
		}
		w.Args[i] = &ActionArgs{Name: paramName, Type: typ}
		if paramName == "*" {
			RegisterStructAttributes(typ)
		}
	}
	return w
}

// The http handler for the wrapped action
func (wr *Wrapper) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	values := r.URL.Query()

	methodArgs := make([]reflect.Value, len(wr.Args))
	for i, a := range wr.Args {
		methodArgs[i] = GetBoundValue(values, a.Name, a.Type)
	}
	var resultValue reflect.Value
	if wr.IsVariadic {
		resultValue = wr.Call.CallSlice(methodArgs)[0]
	} else {
		resultValue = wr.Call.Call(methodArgs)[0]
	}
	resp := resultValue.Interface().(response.Response)
	resp.ApplyTo(w)
}
<|endoftext|>


<commit_before>package main

import (
	"strings"

	"github.com/juju/errors"

	"github.com/rancher/go-rancher-metadata/metadata"
)

//////////////////////////////////////////////////////////////////////////////////
type ServiceWrap struct {
	metadata.Service
}

//////////////////////////////////////////////////////////////////////////////////
type ContainerWrap struct {
	metadata.Container
}

//////////////////////////////////////////////////////////////////////////////////
func (p ContainerWrap) extractPort(idx, partx int) string {
	part := p.Ports[idx]
	if part == "" {
		return part
	}

	items := strings.Split(part, ":")
	if len(items) == 2 {
		part = items[partx]
	} else {
		part = items[0]
	}

	if strings.Contains(part, "/") {
		part = strings.Split(part, "/")[0]
	}

	return part
}

//////////////////////////////////////////////////////////////////////////////////
func (p ContainerWrap) PortInternal(idx int) (string, error) {
	if len(p.Ports) == 0 {
		return "", nil
	}
	if idx >= len(p.Ports) {
		return "", errors.New("PortInternal: index %d out of range", idx)
	}

	return p.extractPort(idx, 1), nil
}

//////////////////////////////////////////////////////////////////////////////////
func (p ContainerWrap) PortExternal(idx int) (string, error) {
	if len(p.Ports) == 0 {
		return "", nil
	}
	if idx >= len(p.Ports) {
		return "", errors.New("PortExternal: index %d out of range", idx)
	}

	return p.extractPort(idx, 0), nil
}

//////////////////////////////////////////////////////////////////////////////////
func (p ContainerWrap) LabelByKey(key string) (string, error) {
	if val, ok := p.Labels[key]; ok {
		return val, nil
	}

	return "", errors.Errorf("LabelByKey: label %s not found", key)
}

//////////////////////////////////////////////////////////////////////////////////
type HostWrap struct {
	metadata.Host
}
<commit_msg>autopush@1446292909<commit_after>package main

import (
	"strings"

	"github.com/juju/errors"

	"github.com/rancher/go-rancher-metadata/metadata"
)

//////////////////////////////////////////////////////////////////////////////////
type ServiceWrap struct {
	metadata.Service
}

//////////////////////////////////////////////////////////////////////////////////
type ContainerWrap struct {
	metadata.Container
}

//////////////////////////////////////////////////////////////////////////////////
func (p ContainerWrap) extractPort(idx, partx int) string {
	part := p.Ports[idx]
	if part == "" {
		return part
	}

	items := strings.Split(part, ":")
	if len(items) == 2 {
		part = items[partx]
	} else {
		part = items[0]
	}

	if strings.Contains(part, "/") {
		part = strings.Split(part, "/")[0]
	}

	return part
}

//////////////////////////////////////////////////////////////////////////////////
func (p ContainerWrap) PortInternal(idx int) (string, error) {
	if len(p.Ports) == 0 {
		return "", nil
	}
	if idx >= len(p.Ports) {
		return "", errors.Errorf("PortInternal: index %d out of range", idx)
	}

	return p.extractPort(idx, 1), nil
}

//////////////////////////////////////////////////////////////////////////////////
func (p ContainerWrap) PortExternal(idx int) (string, error) {
	if len(p.Ports) == 0 {
		return "", nil
	}
	if idx >= len(p.Ports) {
		return "", errors.Errorf("PortExternal: index %d out of range", idx)
	}

	return p.extractPort(idx, 0), nil
}

//////////////////////////////////////////////////////////////////////////////////
func (p ContainerWrap) LabelByKey(key string) (string, error) {
	if val, ok := p.Labels[key]; ok {
		return val, nil
	}

	return "", errors.Errorf("LabelByKey: label %s not found", key)
}

//////////////////////////////////////////////////////////////////////////////////
type HostWrap struct {
	metadata.Host
}
<|endoftext|>


<commit_before>package gocql

import (
	"net"
	"sync"
	"time"
)

type eventDebouncer struct {
	name   string
	timer  *time.Timer
	mu     sync.Mutex
	events []frame

	callback func([]frame)
	quit     chan struct{}
}

func newEventDebouncer(name string, eventHandler func([]frame)) *eventDebouncer {
	e := &eventDebouncer{
		name:     name,
		quit:     make(chan struct{}),
		timer:    time.NewTimer(eventDebounceTime),
		callback: eventHandler,
	}
	e.timer.Stop()
	go e.flusher()

	return e
}

func (e *eventDebouncer) stop() {
	e.quit <- struct{}{} // sync with flusher
	close(e.quit)
}

func (e *eventDebouncer) flusher() {
	for {
		select {
		case <-e.timer.C:
			e.mu.Lock()
			e.flush()
			e.mu.Unlock()
		case <-e.quit:
			return
		}
	}
}

const (
	eventBufferSize   = 1000
	eventDebounceTime = 1 * time.Second
)

// flush must be called with mu locked
func (e *eventDebouncer) flush() {
	if len(e.events) == 0 {
		return
	}

	// if the flush interval is faster than the callback then we will end up calling
	// the callback multiple times, probably a bad idea. In this case we could drop
	// frames?
	go e.callback(e.events)
	e.events = make([]frame, 0, eventBufferSize)
}

func (e *eventDebouncer) debounce(frame frame) {
	e.mu.Lock()
	e.timer.Reset(eventDebounceTime)

	// TODO: probably need a warning to track if this threshold is too low
	if len(e.events) < eventBufferSize {
		e.events = append(e.events, frame)
	} else {
		Logger.Printf("%s: buffer full, dropping event frame: %s", e.name, frame)
	}

	e.mu.Unlock()
}

func (s *Session) handleEvent(framer *framer) {
	// TODO(zariel): need to debounce events frames, and possible also events
	defer framerPool.Put(framer)

	frame, err := framer.parseFrame()
	if err != nil {
		// TODO: logger
		Logger.Printf("gocql: unable to parse event frame: %v\n", err)
		return
	}

	if gocqlDebug {
		Logger.Printf("gocql: handling frame: %v\n", frame)
	}

	// TODO: handle medatadata events
	switch f := frame.(type) {
	case *schemaChangeKeyspace, *schemaChangeFunction, *schemaChangeTable:
		s.schemaEvents.debounce(frame)
	case *topologyChangeEventFrame, *statusChangeEventFrame:
		s.nodeEvents.debounce(frame)
	default:
		Logger.Printf("gocql: invalid event frame (%T): %v\n", f, f)
	}
}

func (s *Session) handleSchemaEvent(frames []frame) {
	if s.schemaDescriber == nil {
		return
	}
	for _, frame := range frames {
		switch f := frame.(type) {
		case *schemaChangeKeyspace:
			s.schemaDescriber.clearSchema(f.keyspace)
		case *schemaChangeTable:
			s.schemaDescriber.clearSchema(f.keyspace)
		}
	}
}

func (s *Session) handleNodeEvent(frames []frame) {
	type nodeEvent struct {
		change string
		host   net.IP
		port   int
	}

	events := make(map[string]*nodeEvent)

	for _, frame := range frames {
		// TODO: can we be sure the order of events in the buffer is correct?
		switch f := frame.(type) {
		case *topologyChangeEventFrame:
			event, ok := events[f.host.String()]
			if !ok {
				event = &nodeEvent{change: f.change, host: f.host, port: f.port}
				events[f.host.String()] = event
			}
			event.change = f.change

		case *statusChangeEventFrame:
			event, ok := events[f.host.String()]
			if !ok {
				event = &nodeEvent{change: f.change, host: f.host, port: f.port}
				events[f.host.String()] = event
			}
			event.change = f.change
		}
	}

	for _, f := range events {
		if gocqlDebug {
			Logger.Printf("gocql: dispatching event: %+v\n", f)
		}

		switch f.change {
		case "NEW_NODE":
			s.handleNewNode(f.host, f.port, true)
		case "REMOVED_NODE":
			s.handleRemovedNode(f.host, f.port)
		case "MOVED_NODE":
		// java-driver handles this, not mentioned in the spec
		// TODO(zariel): refresh token map
		case "UP":
			s.handleNodeUp(f.host, f.port, true)
		case "DOWN":
			s.handleNodeDown(f.host, f.port)
		}
	}
}

func (s *Session) handleNewNode(ip net.IP, port int, waitForBinary bool) {
	var hostInfo *HostInfo
	if s.control != nil && !s.cfg.IgnorePeerAddr {
		var err error
		hostInfo, err = s.control.fetchHostInfo(ip, port)
		if err != nil {
			Logger.Printf("gocql: events: unable to fetch host info for (%s:%d): %v\n", ip, port, err)
			return
		}
	} else {
		hostInfo = &HostInfo{peer: ip, port: port}
	}

	if s.cfg.IgnorePeerAddr && hostInfo.Peer().Equal(ip) {
		hostInfo.setPeer(ip)
	}

	if s.cfg.HostFilter != nil && !s.cfg.HostFilter.Accept(hostInfo) {
		return
	}

	if t := hostInfo.Version().nodeUpDelay(); t > 0 && waitForBinary {
		time.Sleep(t)
	}

	// should this handle token moving?
	if existing, ok := s.ring.addHostIfMissing(hostInfo); ok {
		existing.update(hostInfo)
		hostInfo = existing
	}

	s.pool.addHost(hostInfo)
	s.policy.AddHost(hostInfo)
	hostInfo.setState(NodeUp)
	if s.control != nil && !s.cfg.IgnorePeerAddr {
		s.hostSource.refreshRing()
	}
}

func (s *Session) handleRemovedNode(ip net.IP, port int) {
	// we remove all nodes but only add ones which pass the filter
	host := s.ring.getHost(ip)
	if host == nil {
		host = &HostInfo{peer: ip, port: port}
	}

	if s.cfg.HostFilter != nil && !s.cfg.HostFilter.Accept(host) {
		return
	}

	host.setState(NodeDown)
	s.policy.RemoveHost(host)
	s.pool.removeHost(ip)
	s.ring.removeHost(ip)

	if !s.cfg.IgnorePeerAddr {
		s.hostSource.refreshRing()
	}
}

func (s *Session) handleNodeUp(ip net.IP, port int, waitForBinary bool) {
	if gocqlDebug {
		Logger.Printf("gocql: Session.handleNodeUp: %s:%d\n", ip.String(), port)
	}

	host := s.ring.getHost(ip)
	if host != nil {
		if s.cfg.IgnorePeerAddr && host.Peer().Equal(ip) {
			// TODO: how can this ever be true?
			host.setPeer(ip)
		}

		if s.cfg.HostFilter != nil && !s.cfg.HostFilter.Accept(host) {
			return
		}

		if t := host.Version().nodeUpDelay(); t > 0 && waitForBinary {
			time.Sleep(t)
		}

		s.pool.hostUp(host)
		s.policy.HostUp(host)
		host.setState(NodeUp)
		return
	}

	s.handleNewNode(ip, port, waitForBinary)
}

func (s *Session) handleNodeDown(ip net.IP, port int) {
	if gocqlDebug {
		Logger.Printf("gocql: Session.handleNodeDown: %s:%d\n", ip.String(), port)
	}

	host := s.ring.getHost(ip)
	if host == nil {
		host = &HostInfo{peer: ip, port: port}
	}

	if s.cfg.HostFilter != nil && !s.cfg.HostFilter.Accept(host) {
		return
	}

	host.setState(NodeDown)
	s.policy.HostDown(host)
	s.pool.hostDown(ip)
}
<commit_msg>Fix data race between event processing and session when concurrent goroutines are potentially accessing schema details or creating, altering keyspaces and tables.<commit_after>package gocql

import (
	"net"
	"sync"
	"time"
)

type eventDebouncer struct {
	name   string
	timer  *time.Timer
	mu     sync.Mutex
	events []frame

	callback func([]frame)
	quit     chan struct{}
}

func newEventDebouncer(name string, eventHandler func([]frame)) *eventDebouncer {
	e := &eventDebouncer{
		name:     name,
		quit:     make(chan struct{}),
		timer:    time.NewTimer(eventDebounceTime),
		callback: eventHandler,
	}
	e.timer.Stop()
	go e.flusher()

	return e
}

func (e *eventDebouncer) stop() {
	e.quit <- struct{}{} // sync with flusher
	close(e.quit)
}

func (e *eventDebouncer) flusher() {
	for {
		select {
		case <-e.timer.C:
			e.mu.Lock()
			e.flush()
			e.mu.Unlock()
		case <-e.quit:
			return
		}
	}
}

const (
	eventBufferSize   = 1000
	eventDebounceTime = 1 * time.Second
)

// flush must be called with mu locked
func (e *eventDebouncer) flush() {
	if len(e.events) == 0 {
		return
	}

	// if the flush interval is faster than the callback then we will end up calling
	// the callback multiple times, probably a bad idea. In this case we could drop
	// frames?
	go e.callback(e.events)
	e.events = make([]frame, 0, eventBufferSize)
}

func (e *eventDebouncer) debounce(frame frame) {
	e.mu.Lock()
	e.timer.Reset(eventDebounceTime)

	// TODO: probably need a warning to track if this threshold is too low
	if len(e.events) < eventBufferSize {
		e.events = append(e.events, frame)
	} else {
		Logger.Printf("%s: buffer full, dropping event frame: %s", e.name, frame)
	}

	e.mu.Unlock()
}

func (s *Session) handleEvent(framer *framer) {
	// TODO(zariel): need to debounce events frames, and possible also events
	defer framerPool.Put(framer)

	frame, err := framer.parseFrame()
	if err != nil {
		// TODO: logger
		Logger.Printf("gocql: unable to parse event frame: %v\n", err)
		return
	}

	if gocqlDebug {
		Logger.Printf("gocql: handling frame: %v\n", frame)
	}

	// TODO: handle medatadata events
	switch f := frame.(type) {
	case *schemaChangeKeyspace, *schemaChangeFunction, *schemaChangeTable:
		s.schemaEvents.debounce(frame)
	case *topologyChangeEventFrame, *statusChangeEventFrame:
		s.nodeEvents.debounce(frame)
	default:
		Logger.Printf("gocql: invalid event frame (%T): %v\n", f, f)
	}
}

func (s *Session) handleSchemaEvent(frames []frame) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	if s.schemaDescriber == nil {
		return
	}
	for _, frame := range frames {
		switch f := frame.(type) {
		case *schemaChangeKeyspace:
			s.schemaDescriber.clearSchema(f.keyspace)
		case *schemaChangeTable:
			s.schemaDescriber.clearSchema(f.keyspace)
		}
	}
}

func (s *Session) handleNodeEvent(frames []frame) {
	type nodeEvent struct {
		change string
		host   net.IP
		port   int
	}

	events := make(map[string]*nodeEvent)

	for _, frame := range frames {
		// TODO: can we be sure the order of events in the buffer is correct?
		switch f := frame.(type) {
		case *topologyChangeEventFrame:
			event, ok := events[f.host.String()]
			if !ok {
				event = &nodeEvent{change: f.change, host: f.host, port: f.port}
				events[f.host.String()] = event
			}
			event.change = f.change

		case *statusChangeEventFrame:
			event, ok := events[f.host.String()]
			if !ok {
				event = &nodeEvent{change: f.change, host: f.host, port: f.port}
				events[f.host.String()] = event
			}
			event.change = f.change
		}
	}

	for _, f := range events {
		if gocqlDebug {
			Logger.Printf("gocql: dispatching event: %+v\n", f)
		}

		switch f.change {
		case "NEW_NODE":
			s.handleNewNode(f.host, f.port, true)
		case "REMOVED_NODE":
			s.handleRemovedNode(f.host, f.port)
		case "MOVED_NODE":
		// java-driver handles this, not mentioned in the spec
		// TODO(zariel): refresh token map
		case "UP":
			s.handleNodeUp(f.host, f.port, true)
		case "DOWN":
			s.handleNodeDown(f.host, f.port)
		}
	}
}

func (s *Session) handleNewNode(ip net.IP, port int, waitForBinary bool) {
	var hostInfo *HostInfo
	if s.control != nil && !s.cfg.IgnorePeerAddr {
		var err error
		hostInfo, err = s.control.fetchHostInfo(ip, port)
		if err != nil {
			Logger.Printf("gocql: events: unable to fetch host info for (%s:%d): %v\n", ip, port, err)
			return
		}
	} else {
		hostInfo = &HostInfo{peer: ip, port: port}
	}

	if s.cfg.IgnorePeerAddr && hostInfo.Peer().Equal(ip) {
		hostInfo.setPeer(ip)
	}

	if s.cfg.HostFilter != nil && !s.cfg.HostFilter.Accept(hostInfo) {
		return
	}

	if t := hostInfo.Version().nodeUpDelay(); t > 0 && waitForBinary {
		time.Sleep(t)
	}

	// should this handle token moving?
	if existing, ok := s.ring.addHostIfMissing(hostInfo); ok {
		existing.update(hostInfo)
		hostInfo = existing
	}

	s.pool.addHost(hostInfo)
	s.policy.AddHost(hostInfo)
	hostInfo.setState(NodeUp)
	if s.control != nil && !s.cfg.IgnorePeerAddr {
		s.hostSource.refreshRing()
	}
}

func (s *Session) handleRemovedNode(ip net.IP, port int) {
	// we remove all nodes but only add ones which pass the filter
	host := s.ring.getHost(ip)
	if host == nil {
		host = &HostInfo{peer: ip, port: port}
	}

	if s.cfg.HostFilter != nil && !s.cfg.HostFilter.Accept(host) {
		return
	}

	host.setState(NodeDown)
	s.policy.RemoveHost(host)
	s.pool.removeHost(ip)
	s.ring.removeHost(ip)

	if !s.cfg.IgnorePeerAddr {
		s.hostSource.refreshRing()
	}
}

func (s *Session) handleNodeUp(ip net.IP, port int, waitForBinary bool) {
	if gocqlDebug {
		Logger.Printf("gocql: Session.handleNodeUp: %s:%d\n", ip.String(), port)
	}

	host := s.ring.getHost(ip)
	if host != nil {
		if s.cfg.IgnorePeerAddr && host.Peer().Equal(ip) {
			// TODO: how can this ever be true?
			host.setPeer(ip)
		}

		if s.cfg.HostFilter != nil && !s.cfg.HostFilter.Accept(host) {
			return
		}

		if t := host.Version().nodeUpDelay(); t > 0 && waitForBinary {
			time.Sleep(t)
		}

		s.pool.hostUp(host)
		s.policy.HostUp(host)
		host.setState(NodeUp)
		return
	}

	s.handleNewNode(ip, port, waitForBinary)
}

func (s *Session) handleNodeDown(ip net.IP, port int) {
	if gocqlDebug {
		Logger.Printf("gocql: Session.handleNodeDown: %s:%d\n", ip.String(), port)
	}

	host := s.ring.getHost(ip)
	if host == nil {
		host = &HostInfo{peer: ip, port: port}
	}

	if s.cfg.HostFilter != nil && !s.cfg.HostFilter.Accept(host) {
		return
	}

	host.setState(NodeDown)
	s.policy.HostDown(host)
	s.pool.hostDown(ip)
}
<|endoftext|>


<commit_before>/*
Copyright 2018 Google Inc. All rights reserved.

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 websockets

import (
	"fmt"
	"net/http"
	"sync"
	"time"

	"context"
	"github.com/gorilla/websocket"
)

// Connection implements a websocket client connection.
//
// This wraps a websocket.Conn connection as defined by the gorilla/websocket library,
// and encapsulates it in an API that is a little more amenable to how the server side
// of our websocket shim is implemented.
type Connection struct {
	ctx            context.Context
	cancel         context.CancelFunc
	clientMessages chan string
	serverMessages chan string
}

// This map defines the set of headers that should be stripped from the WS request, as they
// are being verified by the websocket.Client.
var stripHeaderNames = map[string]bool{
	"Upgrade":                  true,
	"Connection":               true,
	"Sec-Websocket-Key":        true,
	"Sec-Websocket-Version":    true,
	"Sec-Websocket-Extensions": true,
	"Sec-Websocket-Protocol":   true,
}

// stripWSHeader strips request headers that are not allowed by the websocket.Client library,
// while preserving the rest.
func stripWSHeader(header http.Header) http.Header {
	result := http.Header{}
	for k, vals := range header {
		if !stripHeaderNames[k] {
			for _, v := range vals {
				result.Add(k, v)
			}
		}
	}
	return result
}

// NewConnection creates and returns a new Connection.
func NewConnection(ctx context.Context, targetURL string, header http.Header, errCallback func(err error)) (*Connection, error) {
	ctx, cancel := context.WithCancel(ctx)
	serverConn, _, err := websocket.DefaultDialer.Dial(targetURL, stripWSHeader(header))
	if err != nil {
		cancel()
		return nil, err
	}

	// The underlying websocket library assumes that clients are written to read from the connection
	// in a tight loop. We prefer to instead have the ability to poll for the next message with a timeout.
	//
	// The most direct way to bridge those two approaches is to spawn a goroutine that reads in a
	// tight loop, and then sends the read messages in to an open channel.
	serverMessages := make(chan string, 10)

	// Since we are using a channel to pull messages from the connection, we will also use one to
	// push messages. That way our handling of reads and writes are consistent.
	clientMessages := make(chan string, 10)
	var wg sync.WaitGroup
	wg.Add(2)
	go func() {
		defer wg.Done()
		defer func() {
			close(serverMessages)
		}()
		for {
			select {
			case <-ctx.Done():
				return
			default:
				_, msgBytes, err := serverConn.ReadMessage()
				if err != nil {
					errCallback(fmt.Errorf("failed to read a websocket message from the server: %v", err))
					// Errors from the server connection are terminal; once an error is returned
					// no subsequent calls will succeed.
					return
				}
				serverMessages <- string(msgBytes)
			}
		}
	}()
	go func() {
		defer wg.Done()
		for {
			select {
			case <-ctx.Done():
				return
			case clientMsg := <-clientMessages:
				if err := serverConn.WriteMessage(websocket.TextMessage, []byte(clientMsg)); err != nil {
					errCallback(fmt.Errorf("failed to forward websocket data to the server: %v", err))
					// Errors writing to the server connection are terminal; once an error is returned
					// no subsequent calls will succeed.
					return
				}
			}
		}
	}()
	go func() {
		wg.Wait()
		if err := serverConn.WriteMessage(
			websocket.CloseMessage,
			websocket.FormatCloseMessage(websocket.CloseNormalClosure, "")); err != nil {
			errCallback(fmt.Errorf("failure issuing a close message on a server websocket connection: %v", err))
		}
		if err := serverConn.Close(); err != nil {
			errCallback(fmt.Errorf("failure closing a server websocket connection: %v", err))
		}
	}()
	return &Connection{
		ctx:            ctx,
		cancel:         cancel,
		clientMessages: clientMessages,
		serverMessages: serverMessages,
	}, nil
}

// Close closes the websocket client connection.
func (conn *Connection) Close() {
	conn.cancel()
	close(conn.clientMessages)
}

// SendClientMessage sends the given message to the websocket server.
//
// The returned error value is non-nill if the connection has been closed.
func (conn *Connection) SendClientMessage(msg string) error {
	select {
	case <-conn.ctx.Done():
		return fmt.Errorf("attempt to send a client message on a closed websocket connection")
	default:
		conn.clientMessages <- msg
	}
	return nil
}

// ReadServerMessages reads the next messages from the websocket server.
//
// The returned error value is non-nill if the connection has been closed.
//
// The returned []string value is nil if the error is non-nil, or if the method
// times out while waiting for a server message.
func (conn *Connection) ReadServerMessages() ([]string, error) {
	var msgs []string
	select {
	case serverMsg, ok := <-conn.serverMessages:
		if !ok {
			// The server messages channel has been closed.
			return nil, fmt.Errorf("attempt to read a server message from a closed websocket connection")
		}
		msgs = append(msgs, serverMsg)
		for {
			select {
			case serverMsg, ok := <-conn.serverMessages:
				if ok {
					msgs = append(msgs, serverMsg)
				} else {
					return msgs, nil
				}
			default:
				return msgs, nil
			}
		}
	case <-time.After(time.Second * 20):
		return nil, nil
	}
}
<commit_msg>Address review comment<commit_after>/*
Copyright 2018 Google Inc. All rights reserved.

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 websockets

import (
	"fmt"
	"net/http"
	"sync"
	"time"

	"context"
	"github.com/gorilla/websocket"
)

// Connection implements a websocket client connection.
//
// This wraps a websocket.Conn connection as defined by the gorilla/websocket library,
// and encapsulates it in an API that is a little more amenable to how the server side
// of our websocket shim is implemented.
type Connection struct {
	ctx            context.Context
	cancel         context.CancelFunc
	clientMessages chan string
	serverMessages chan string
}

// This map defines the set of headers that should be stripped from the WS request, as they
// are being verified by the websocket.Client.
var stripHeaderNames = map[string]bool{
	"Upgrade":                  true,
	"Connection":               true,
	"Sec-Websocket-Key":        true,
	"Sec-Websocket-Version":    true,
	"Sec-Websocket-Extensions": true,
	"Sec-Websocket-Protocol":   true,
}

// stripWSHeader strips request headers that are not allowed by the websocket.Client library,
// while preserving the rest.
func stripWSHeader(header http.Header) http.Header {
	result := http.Header{}
	for k, v := range header {
		if !stripHeaderNames[k] {
			result[k] = v
		}
	}
	return result
}

// NewConnection creates and returns a new Connection.
func NewConnection(ctx context.Context, targetURL string, header http.Header, errCallback func(err error)) (*Connection, error) {
	ctx, cancel := context.WithCancel(ctx)
	serverConn, _, err := websocket.DefaultDialer.Dial(targetURL, stripWSHeader(header))
	if err != nil {
		cancel()
		return nil, err
	}

	// The underlying websocket library assumes that clients are written to read from the connection
	// in a tight loop. We prefer to instead have the ability to poll for the next message with a timeout.
	//
	// The most direct way to bridge those two approaches is to spawn a goroutine that reads in a
	// tight loop, and then sends the read messages in to an open channel.
	serverMessages := make(chan string, 10)

	// Since we are using a channel to pull messages from the connection, we will also use one to
	// push messages. That way our handling of reads and writes are consistent.
	clientMessages := make(chan string, 10)
	var wg sync.WaitGroup
	wg.Add(2)
	go func() {
		defer wg.Done()
		defer func() {
			close(serverMessages)
		}()
		for {
			select {
			case <-ctx.Done():
				return
			default:
				_, msgBytes, err := serverConn.ReadMessage()
				if err != nil {
					errCallback(fmt.Errorf("failed to read a websocket message from the server: %v", err))
					// Errors from the server connection are terminal; once an error is returned
					// no subsequent calls will succeed.
					return
				}
				serverMessages <- string(msgBytes)
			}
		}
	}()
	go func() {
		defer wg.Done()
		for {
			select {
			case <-ctx.Done():
				return
			case clientMsg := <-clientMessages:
				if err := serverConn.WriteMessage(websocket.TextMessage, []byte(clientMsg)); err != nil {
					errCallback(fmt.Errorf("failed to forward websocket data to the server: %v", err))
					// Errors writing to the server connection are terminal; once an error is returned
					// no subsequent calls will succeed.
					return
				}
			}
		}
	}()
	go func() {
		wg.Wait()
		if err := serverConn.WriteMessage(
			websocket.CloseMessage,
			websocket.FormatCloseMessage(websocket.CloseNormalClosure, "")); err != nil {
			errCallback(fmt.Errorf("failure issuing a close message on a server websocket connection: %v", err))
		}
		if err := serverConn.Close(); err != nil {
			errCallback(fmt.Errorf("failure closing a server websocket connection: %v", err))
		}
	}()
	return &Connection{
		ctx:            ctx,
		cancel:         cancel,
		clientMessages: clientMessages,
		serverMessages: serverMessages,
	}, nil
}

// Close closes the websocket client connection.
func (conn *Connection) Close() {
	conn.cancel()
	close(conn.clientMessages)
}

// SendClientMessage sends the given message to the websocket server.
//
// The returned error value is non-nill if the connection has been closed.
func (conn *Connection) SendClientMessage(msg string) error {
	select {
	case <-conn.ctx.Done():
		return fmt.Errorf("attempt to send a client message on a closed websocket connection")
	default:
		conn.clientMessages <- msg
	}
	return nil
}

// ReadServerMessages reads the next messages from the websocket server.
//
// The returned error value is non-nill if the connection has been closed.
//
// The returned []string value is nil if the error is non-nil, or if the method
// times out while waiting for a server message.
func (conn *Connection) ReadServerMessages() ([]string, error) {
	var msgs []string
	select {
	case serverMsg, ok := <-conn.serverMessages:
		if !ok {
			// The server messages channel has been closed.
			return nil, fmt.Errorf("attempt to read a server message from a closed websocket connection")
		}
		msgs = append(msgs, serverMsg)
		for {
			select {
			case serverMsg, ok := <-conn.serverMessages:
				if ok {
					msgs = append(msgs, serverMsg)
				} else {
					return msgs, nil
				}
			default:
				return msgs, nil
			}
		}
	case <-time.After(time.Second * 20):
		return nil, nil
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"
	"path/filepath"
	"testing"

	"github.com/docker/docker/cliconfig"
	"github.com/docker/docker/pkg/reexec"
	"github.com/go-check/check"
)

func Test(t *testing.T) {
	reexec.Init() // This is required for external graphdriver tests

	if !isLocalDaemon {
		fmt.Println("INFO: Testing against a remote daemon")
	} else {
		fmt.Println("INFO: Testing against a local daemon")
	}

	check.TestingT(t)
}

func init() {
	check.Suite(&DockerSuite{})
}

type DockerSuite struct {
}

func (s *DockerSuite) TearDownTest(c *check.C) {
	unpauseAllContainers()
	deleteAllContainers()
	deleteAllImages()
	deleteAllVolumes()
	deleteAllNetworks()
}

func init() {
	check.Suite(&DockerRegistrySuite{
		ds: &DockerSuite{},
	})
}

type DockerRegistrySuite struct {
	ds  *DockerSuite
	reg *testRegistryV2
	d   *Daemon
}

func (s *DockerRegistrySuite) SetUpTest(c *check.C) {
	testRequires(c, DaemonIsLinux, RegistryHosting)
	s.reg = setupRegistry(c, false, "", "")
	s.d = NewDaemon(c)
}

func (s *DockerRegistrySuite) TearDownTest(c *check.C) {
	if s.reg != nil {
		s.reg.Close()
	}
	if s.d != nil {
		s.d.Stop()
	}
	s.ds.TearDownTest(c)
}

func init() {
	check.Suite(&DockerSchema1RegistrySuite{
		ds: &DockerSuite{},
	})
}

type DockerSchema1RegistrySuite struct {
	ds  *DockerSuite
	reg *testRegistryV2
	d   *Daemon
}

func (s *DockerSchema1RegistrySuite) SetUpTest(c *check.C) {
	testRequires(c, DaemonIsLinux, RegistryHosting)
	s.reg = setupRegistry(c, true, "", "")
	s.d = NewDaemon(c)
}

func (s *DockerSchema1RegistrySuite) TearDownTest(c *check.C) {
	if s.reg != nil {
		s.reg.Close()
	}
	if s.d != nil {
		s.d.Stop()
	}
	s.ds.TearDownTest(c)
}

func init() {
	check.Suite(&DockerRegistryAuthHtpasswdSuite{
		ds: &DockerSuite{},
	})
}

type DockerRegistryAuthHtpasswdSuite struct {
	ds  *DockerSuite
	reg *testRegistryV2
	d   *Daemon
}

func (s *DockerRegistryAuthHtpasswdSuite) SetUpTest(c *check.C) {
	testRequires(c, DaemonIsLinux, RegistryHosting)
	s.reg = setupRegistry(c, false, "htpasswd", "")
	s.d = NewDaemon(c)
}

func (s *DockerRegistryAuthHtpasswdSuite) TearDownTest(c *check.C) {
	if s.reg != nil {
		out, err := s.d.Cmd("logout", privateRegistryURL)
		c.Assert(err, check.IsNil, check.Commentf(out))
		s.reg.Close()
	}
	if s.d != nil {
		s.d.Stop()
	}
	s.ds.TearDownTest(c)
}

func init() {
	check.Suite(&DockerRegistryAuthTokenSuite{
		ds: &DockerSuite{},
	})
}

type DockerRegistryAuthTokenSuite struct {
	ds  *DockerSuite
	reg *testRegistryV2
	d   *Daemon
}

func (s *DockerRegistryAuthTokenSuite) SetUpTest(c *check.C) {
	testRequires(c, DaemonIsLinux, RegistryHosting)
	s.d = NewDaemon(c)
}

func (s *DockerRegistryAuthTokenSuite) TearDownTest(c *check.C) {
	if s.reg != nil {
		out, err := s.d.Cmd("logout", privateRegistryURL)
		c.Assert(err, check.IsNil, check.Commentf(out))
		s.reg.Close()
	}
	if s.d != nil {
		s.d.Stop()
	}
	s.ds.TearDownTest(c)
}

func (s *DockerRegistryAuthTokenSuite) setupRegistryWithTokenService(c *check.C, tokenURL string) {
	if s == nil {
		c.Fatal("registry suite isn't initialized")
	}
	s.reg = setupRegistry(c, false, "token", tokenURL)
}

func init() {
	check.Suite(&DockerDaemonSuite{
		ds: &DockerSuite{},
	})
}

type DockerDaemonSuite struct {
	ds *DockerSuite
	d  *Daemon
}

func (s *DockerDaemonSuite) SetUpTest(c *check.C) {
	testRequires(c, DaemonIsLinux)
	s.d = NewDaemon(c)
}

func (s *DockerDaemonSuite) TearDownTest(c *check.C) {
	testRequires(c, DaemonIsLinux)
	if s.d != nil {
		s.d.Stop()
	}
	s.ds.TearDownTest(c)
}

const defaultSwarmPort = 2477

func init() {
	check.Suite(&DockerSwarmSuite{
		ds: &DockerSuite{},
	})
}

type DockerSwarmSuite struct {
	ds        *DockerSuite
	daemons   []*SwarmDaemon
	portIndex int
}

func (s *DockerSwarmSuite) SetUpTest(c *check.C) {
	testRequires(c, DaemonIsLinux)
}

func (s *DockerSwarmSuite) AddDaemon(c *check.C, joinSwarm, manager bool) *SwarmDaemon {
	d := &SwarmDaemon{
		Daemon: NewDaemon(c),
		port:   defaultSwarmPort + s.portIndex,
	}
	d.listenAddr = fmt.Sprintf("0.0.0.0:%d", d.port)
	err := d.StartWithBusybox()
	c.Assert(err, check.IsNil)

	if joinSwarm == true {
		if len(s.daemons) > 0 {
			c.Assert(d.Join(s.daemons[0].listenAddr, "", "", manager), check.IsNil)
		} else {
			aa := make(map[string]bool)
			aa["worker"] = true
			aa["manager"] = true
			c.Assert(d.Init(aa, ""), check.IsNil)
		}
	}

	s.portIndex++
	s.daemons = append(s.daemons, d)

	return d
}

func (s *DockerSwarmSuite) TearDownTest(c *check.C) {
	testRequires(c, DaemonIsLinux)
	for _, d := range s.daemons {
		d.Stop()
	}
	s.daemons = nil
	s.portIndex = 0

	s.ds.TearDownTest(c)
}

func init() {
	check.Suite(&DockerTrustSuite{
		ds: &DockerSuite{},
	})
}

type DockerTrustSuite struct {
	ds  *DockerSuite
	reg *testRegistryV2
	not *testNotary
}

func (s *DockerTrustSuite) SetUpTest(c *check.C) {
	testRequires(c, RegistryHosting, NotaryServerHosting)
	s.reg = setupRegistry(c, false, "", "")
	s.not = setupNotary(c)
}

func (s *DockerTrustSuite) TearDownTest(c *check.C) {
	if s.reg != nil {
		s.reg.Close()
	}
	if s.not != nil {
		s.not.Close()
	}

	// Remove trusted keys and metadata after test
	os.RemoveAll(filepath.Join(cliconfig.ConfigDir(), "trust"))
	s.ds.TearDownTest(c)
}
<commit_msg>Disable iptables in swarm test daemons<commit_after>package main

import (
	"fmt"
	"os"
	"path/filepath"
	"testing"

	"github.com/docker/docker/cliconfig"
	"github.com/docker/docker/pkg/reexec"
	"github.com/go-check/check"
)

func Test(t *testing.T) {
	reexec.Init() // This is required for external graphdriver tests

	if !isLocalDaemon {
		fmt.Println("INFO: Testing against a remote daemon")
	} else {
		fmt.Println("INFO: Testing against a local daemon")
	}

	check.TestingT(t)
}

func init() {
	check.Suite(&DockerSuite{})
}

type DockerSuite struct {
}

func (s *DockerSuite) TearDownTest(c *check.C) {
	unpauseAllContainers()
	deleteAllContainers()
	deleteAllImages()
	deleteAllVolumes()
	deleteAllNetworks()
}

func init() {
	check.Suite(&DockerRegistrySuite{
		ds: &DockerSuite{},
	})
}

type DockerRegistrySuite struct {
	ds  *DockerSuite
	reg *testRegistryV2
	d   *Daemon
}

func (s *DockerRegistrySuite) SetUpTest(c *check.C) {
	testRequires(c, DaemonIsLinux, RegistryHosting)
	s.reg = setupRegistry(c, false, "", "")
	s.d = NewDaemon(c)
}

func (s *DockerRegistrySuite) TearDownTest(c *check.C) {
	if s.reg != nil {
		s.reg.Close()
	}
	if s.d != nil {
		s.d.Stop()
	}
	s.ds.TearDownTest(c)
}

func init() {
	check.Suite(&DockerSchema1RegistrySuite{
		ds: &DockerSuite{},
	})
}

type DockerSchema1RegistrySuite struct {
	ds  *DockerSuite
	reg *testRegistryV2
	d   *Daemon
}

func (s *DockerSchema1RegistrySuite) SetUpTest(c *check.C) {
	testRequires(c, DaemonIsLinux, RegistryHosting)
	s.reg = setupRegistry(c, true, "", "")
	s.d = NewDaemon(c)
}

func (s *DockerSchema1RegistrySuite) TearDownTest(c *check.C) {
	if s.reg != nil {
		s.reg.Close()
	}
	if s.d != nil {
		s.d.Stop()
	}
	s.ds.TearDownTest(c)
}

func init() {
	check.Suite(&DockerRegistryAuthHtpasswdSuite{
		ds: &DockerSuite{},
	})
}

type DockerRegistryAuthHtpasswdSuite struct {
	ds  *DockerSuite
	reg *testRegistryV2
	d   *Daemon
}

func (s *DockerRegistryAuthHtpasswdSuite) SetUpTest(c *check.C) {
	testRequires(c, DaemonIsLinux, RegistryHosting)
	s.reg = setupRegistry(c, false, "htpasswd", "")
	s.d = NewDaemon(c)
}

func (s *DockerRegistryAuthHtpasswdSuite) TearDownTest(c *check.C) {
	if s.reg != nil {
		out, err := s.d.Cmd("logout", privateRegistryURL)
		c.Assert(err, check.IsNil, check.Commentf(out))
		s.reg.Close()
	}
	if s.d != nil {
		s.d.Stop()
	}
	s.ds.TearDownTest(c)
}

func init() {
	check.Suite(&DockerRegistryAuthTokenSuite{
		ds: &DockerSuite{},
	})
}

type DockerRegistryAuthTokenSuite struct {
	ds  *DockerSuite
	reg *testRegistryV2
	d   *Daemon
}

func (s *DockerRegistryAuthTokenSuite) SetUpTest(c *check.C) {
	testRequires(c, DaemonIsLinux, RegistryHosting)
	s.d = NewDaemon(c)
}

func (s *DockerRegistryAuthTokenSuite) TearDownTest(c *check.C) {
	if s.reg != nil {
		out, err := s.d.Cmd("logout", privateRegistryURL)
		c.Assert(err, check.IsNil, check.Commentf(out))
		s.reg.Close()
	}
	if s.d != nil {
		s.d.Stop()
	}
	s.ds.TearDownTest(c)
}

func (s *DockerRegistryAuthTokenSuite) setupRegistryWithTokenService(c *check.C, tokenURL string) {
	if s == nil {
		c.Fatal("registry suite isn't initialized")
	}
	s.reg = setupRegistry(c, false, "token", tokenURL)
}

func init() {
	check.Suite(&DockerDaemonSuite{
		ds: &DockerSuite{},
	})
}

type DockerDaemonSuite struct {
	ds *DockerSuite
	d  *Daemon
}

func (s *DockerDaemonSuite) SetUpTest(c *check.C) {
	testRequires(c, DaemonIsLinux)
	s.d = NewDaemon(c)
}

func (s *DockerDaemonSuite) TearDownTest(c *check.C) {
	testRequires(c, DaemonIsLinux)
	if s.d != nil {
		s.d.Stop()
	}
	s.ds.TearDownTest(c)
}

const defaultSwarmPort = 2477

func init() {
	check.Suite(&DockerSwarmSuite{
		ds: &DockerSuite{},
	})
}

type DockerSwarmSuite struct {
	ds        *DockerSuite
	daemons   []*SwarmDaemon
	portIndex int
}

func (s *DockerSwarmSuite) SetUpTest(c *check.C) {
	testRequires(c, DaemonIsLinux)
}

func (s *DockerSwarmSuite) AddDaemon(c *check.C, joinSwarm, manager bool) *SwarmDaemon {
	d := &SwarmDaemon{
		Daemon: NewDaemon(c),
		port:   defaultSwarmPort + s.portIndex,
	}
	d.listenAddr = fmt.Sprintf("0.0.0.0:%d", d.port)
	err := d.StartWithBusybox("--iptables=false") // avoid networking conflicts
	c.Assert(err, check.IsNil)

	if joinSwarm == true {
		if len(s.daemons) > 0 {
			c.Assert(d.Join(s.daemons[0].listenAddr, "", "", manager), check.IsNil)
		} else {
			aa := make(map[string]bool)
			aa["worker"] = true
			aa["manager"] = true
			c.Assert(d.Init(aa, ""), check.IsNil)
		}
	}

	s.portIndex++
	s.daemons = append(s.daemons, d)

	return d
}

func (s *DockerSwarmSuite) TearDownTest(c *check.C) {
	testRequires(c, DaemonIsLinux)
	for _, d := range s.daemons {
		d.Stop()
	}
	s.daemons = nil
	s.portIndex = 0

	s.ds.TearDownTest(c)
}

func init() {
	check.Suite(&DockerTrustSuite{
		ds: &DockerSuite{},
	})
}

type DockerTrustSuite struct {
	ds  *DockerSuite
	reg *testRegistryV2
	not *testNotary
}

func (s *DockerTrustSuite) SetUpTest(c *check.C) {
	testRequires(c, RegistryHosting, NotaryServerHosting)
	s.reg = setupRegistry(c, false, "", "")
	s.not = setupNotary(c)
}

func (s *DockerTrustSuite) TearDownTest(c *check.C) {
	if s.reg != nil {
		s.reg.Close()
	}
	if s.not != nil {
		s.not.Close()
	}

	// Remove trusted keys and metadata after test
	os.RemoveAll(filepath.Join(cliconfig.ConfigDir(), "trust"))
	s.ds.TearDownTest(c)
}
<|endoftext|>


<commit_before>package main

//go:generate go run generate.go protocol/input/json.json ../../aws/protocol/jsonrpc/jsonrpc_input_test.go

import (
	"bytes"
	"encoding/json"
	"fmt"
	"net/url"
	"os"
	"regexp"
	"strconv"
	"strings"
	"text/template"

	"github.com/awslabs/aws-sdk-go/internal/model/api"
	"github.com/awslabs/aws-sdk-go/internal/util"
)

type TestSuite struct {
	*api.API
	Description string
	Cases       []TestCase
	title       string
}

type TestCase struct {
	*TestSuite
	Given      *api.Operation
	Params     interface{}     `json:",omitempty"`
	Data       interface{}     `json:"result,omitempty"`
	InputTest  TestExpectation `json:"serialized"`
	OutputTest TestExpectation `json:"response"`
}

type TestExpectation struct {
	Body       string
	URI        string
	Headers    map[string]string
	StatusCode uint `json:"status_code"`
}

const preamble = `
var _ bytes.Buffer // always import bytes
var _ http.Request
`

var reStripSpace = regexp.MustCompile(`\s(\w)`)

var reImportRemoval = regexp.MustCompile(`(?s:import \((.+?)\))`)

func removeImports(code string) string {
	return reImportRemoval.ReplaceAllString(code, "")
}

var extraImports = []string{
	"bytes",
	"encoding/json",
	"io/ioutil",
	"net/http",
	"testing",
	"github.com/awslabs/aws-sdk-go/internal/util",
	"github.com/stretchr/testify/assert",
}

func addImports(code string) string {
	importNames := make([]string, len(extraImports))
	for i, n := range extraImports {
		importNames[i] = fmt.Sprintf("%q", n)
	}
	str := reImportRemoval.ReplaceAllString(code, "import (\n$1\n"+strings.Join(importNames, "\n")+")")
	return str
}

func (t *TestSuite) TestSuite() string {
	var buf bytes.Buffer

	t.title = reStripSpace.ReplaceAllStringFunc(t.Description, func(x string) string {
		return strings.ToUpper(x[1:])
	})

	for idx, c := range t.Cases {
		c.TestSuite = t
		buf.WriteString(c.TestCase(idx) + "\n")
	}
	return util.GoFmt(buf.String())
}

var tplInputTestCase = template.Must(template.New("inputcase").Parse(`
func Test{{ .OpName }}(t *testing.T) {
	svc := New{{ .TestCase.TestSuite.API.StructName }}(nil)

	var input {{ .Given.InputRef.ShapeName }}
	json.Unmarshal([]byte({{ .ParamsString }}), &input)
	req := svc.{{ .Given.ExportedName }}Request(&input)
	r := req.HTTPRequest

	// build request
	{{ .TestCase.TestSuite.API.ProtocolPackage }}.Build(req)
	assert.NoError(t, req.Error)

	// assert body
	body, _ := ioutil.ReadAll(r.Body)
	assert.Equal(t, util.Trim({{ .Body }}), util.Trim(string(body)))

	// assert URL
	assert.Equal(t, "{{ .TestCase.InputTest.URI }}", r.URL.Path)

	// assert headers
{{ range $k, $v := .TestCase.InputTest.Headers }}assert.Equal(t, "{{ $v }}", r.Header.Get("{{ $k }}"))
{{ end }}
}
`))

type tplInputTestCaseData struct {
	*TestCase
	Body, OpName, ParamsString string
}

var tplOutputTestCase = template.Must(template.New("outputcase").Parse(`
func Test{{ .OpName }}(t *testing.T) {
	svc := New{{ .TestCase.TestSuite.API.StructName }}(nil)

	buf := bytes.NewReader([]byte({{ .Body }}))
	req, _ := svc.{{ .Given.ExportedName }}Request()
	req.HTTPResponse = &http.Response{StatusCode: 200, Body: ioutil.NopCloser(buf)}

	// unmarshal response
	{{ .TestCase.TestSuite.API.ProtocolPackage }}.Unmarshal(req)
	assert.NoError(t, req.Error)

	// assert response
	buf, _ := json.Marshal(req.Data)
	assert.Equal(t, util.Trim({{ .ResponseString }}), util.Trim(string(buf)))

	// assert headers
{{ range $k, $v := .TestCase.OutputTest.Headers }}assert.Equal(t, "{{ $v }}", req.HTTPResponse.Header.Get("{{ $k }}"))
{{ end }}
}
`))

type tplOutputTestCaseData struct {
	*TestCase
	Body, OpName, ResponseString string
}

func (i *TestCase) TestCase(idx int) string {
	var buf bytes.Buffer

	opName := i.API.StructName() + i.TestSuite.title + "Case" + strconv.Itoa(idx+1)

	if i.Params != nil { // input test
		// query test should sort body as form encoded values
		switch i.API.ProtocolPackage() {
		case "query", "ec2":
			m, _ := url.ParseQuery(i.InputTest.Body)
			i.InputTest.Body = m.Encode()
		}

		pBuf, _ := json.Marshal(i.Params)
		input := tplInputTestCaseData{
			TestCase:     i,
			Body:         fmt.Sprintf("%q", i.InputTest.Body),
			OpName:       strings.ToUpper(opName[0:1]) + opName[1:],
			ParamsString: fmt.Sprintf("%q", pBuf),
		}

		if err := tplInputTestCase.Execute(&buf, input); err != nil {
			panic(err)
		}
	} else {
		pBuf, _ := json.Marshal(i.Data)
		output := tplOutputTestCaseData{
			TestCase:       i,
			Body:           fmt.Sprintf("%q", i.OutputTest.Body),
			OpName:         strings.ToUpper(opName[0:1]) + opName[1:],
			ResponseString: fmt.Sprintf("%q", pBuf),
		}

		if err := tplOutputTestCase.Execute(&buf, output); err != nil {
			panic(err)
		}
	}

	return util.GoFmt(buf.String())
}

func GenerateTestSuite(filename string) string {
	inout := "Input"
	if strings.Contains(filename, "/output/") {
		inout = "Output"
	}

	var suites []TestSuite
	f, err := os.Open(filename)
	if err != nil {
		panic(err)
	}

	err = json.NewDecoder(f).Decode(&suites)
	if err != nil {
		panic(err)
	}

	var buf bytes.Buffer
	buf.WriteString("package " + suites[0].ProtocolPackage() + "_test\n\n")

	var innerBuf bytes.Buffer
	innerBuf.WriteString("//\n// Tests begin here\n//\n\n\n")

	for i, suite := range suites {
		svcPrefix := inout + "Service" + strconv.Itoa(i+1)
		suite.API.Metadata.ServiceAbbreviation = svcPrefix + "ProtocolTest"
		suite.API.Operations = map[string]*api.Operation{}
		for idx, c := range suite.Cases {
			c.Given.ExportedName = svcPrefix + "TestCaseOperation" + strconv.Itoa(idx+1)
			suite.API.Operations[c.Given.ExportedName] = c.Given
		}

		suite.API.Setup()

		for n, s := range suite.API.Shapes {
			s.Rename(svcPrefix + "TestShape" + n)
		}

		svcCode := addImports(suite.API.ServiceGoCode())
		if i == 0 {
			importMatch := reImportRemoval.FindStringSubmatch(svcCode)
			buf.WriteString(importMatch[0] + "\n\n")
			buf.WriteString(preamble + "\n\n")
		}
		svcCode = removeImports(svcCode)
		svcCode = strings.Replace(svcCode, "func New(", "func New"+suite.API.StructName()+"(", -1)

		buf.WriteString(svcCode + "\n\n")
		buf.WriteString(removeImports(suite.API.APIGoCode()) + "\n\n")
		innerBuf.WriteString(suite.TestSuite() + "\n")
	}

	return util.GoFmt(buf.String() + innerBuf.String())
}

func main() {
	out := GenerateTestSuite(os.Args[1])
	if len(os.Args) == 3 {
		f, err := os.Create(os.Args[2])
		defer f.Close()
		if err != nil {
			panic(err)
		}
		f.WriteString(out + "\n")
	} else {
		fmt.Println(out)
	}
}
<commit_msg>Add nice endpoint prefix for prettier generated test code<commit_after>package main

//go:generate go run generate.go protocol/input/json.json ../../aws/protocol/jsonrpc/jsonrpc_input_test.go

import (
	"bytes"
	"encoding/json"
	"fmt"
	"net/url"
	"os"
	"regexp"
	"strconv"
	"strings"
	"text/template"

	"github.com/awslabs/aws-sdk-go/internal/model/api"
	"github.com/awslabs/aws-sdk-go/internal/util"
)

type TestSuite struct {
	*api.API
	Description string
	Cases       []TestCase
	title       string
}

type TestCase struct {
	*TestSuite
	Given      *api.Operation
	Params     interface{}     `json:",omitempty"`
	Data       interface{}     `json:"result,omitempty"`
	InputTest  TestExpectation `json:"serialized"`
	OutputTest TestExpectation `json:"response"`
}

type TestExpectation struct {
	Body       string
	URI        string
	Headers    map[string]string
	StatusCode uint `json:"status_code"`
}

const preamble = `
var _ bytes.Buffer // always import bytes
var _ http.Request
`

var reStripSpace = regexp.MustCompile(`\s(\w)`)

var reImportRemoval = regexp.MustCompile(`(?s:import \((.+?)\))`)

func removeImports(code string) string {
	return reImportRemoval.ReplaceAllString(code, "")
}

var extraImports = []string{
	"bytes",
	"encoding/json",
	"io/ioutil",
	"net/http",
	"testing",
	"github.com/awslabs/aws-sdk-go/internal/util",
	"github.com/stretchr/testify/assert",
}

func addImports(code string) string {
	importNames := make([]string, len(extraImports))
	for i, n := range extraImports {
		importNames[i] = fmt.Sprintf("%q", n)
	}
	str := reImportRemoval.ReplaceAllString(code, "import (\n$1\n"+strings.Join(importNames, "\n")+")")
	return str
}

func (t *TestSuite) TestSuite() string {
	var buf bytes.Buffer

	t.title = reStripSpace.ReplaceAllStringFunc(t.Description, func(x string) string {
		return strings.ToUpper(x[1:])
	})

	for idx, c := range t.Cases {
		c.TestSuite = t
		buf.WriteString(c.TestCase(idx) + "\n")
	}
	return util.GoFmt(buf.String())
}

var tplInputTestCase = template.Must(template.New("inputcase").Parse(`
func Test{{ .OpName }}(t *testing.T) {
	svc := New{{ .TestCase.TestSuite.API.StructName }}(nil)

	var input {{ .Given.InputRef.ShapeName }}
	json.Unmarshal([]byte({{ .ParamsString }}), &input)
	req := svc.{{ .Given.ExportedName }}Request(&input)
	r := req.HTTPRequest

	// build request
	{{ .TestCase.TestSuite.API.ProtocolPackage }}.Build(req)
	assert.NoError(t, req.Error)

	// assert body
	body, _ := ioutil.ReadAll(r.Body)
	assert.Equal(t, util.Trim({{ .Body }}), util.Trim(string(body)))

	// assert URL
	assert.Equal(t, "{{ .TestCase.InputTest.URI }}", r.URL.Path)

	// assert headers
{{ range $k, $v := .TestCase.InputTest.Headers }}assert.Equal(t, "{{ $v }}", r.Header.Get("{{ $k }}"))
{{ end }}
}
`))

type tplInputTestCaseData struct {
	*TestCase
	Body, OpName, ParamsString string
}

var tplOutputTestCase = template.Must(template.New("outputcase").Parse(`
func Test{{ .OpName }}(t *testing.T) {
	svc := New{{ .TestCase.TestSuite.API.StructName }}(nil)

	buf := bytes.NewReader([]byte({{ .Body }}))
	req, _ := svc.{{ .Given.ExportedName }}Request()
	req.HTTPResponse = &http.Response{StatusCode: 200, Body: ioutil.NopCloser(buf)}

	// unmarshal response
	{{ .TestCase.TestSuite.API.ProtocolPackage }}.Unmarshal(req)
	assert.NoError(t, req.Error)

	// assert response
	buf, _ := json.Marshal(req.Data)
	assert.Equal(t, util.Trim({{ .ResponseString }}), util.Trim(string(buf)))

	// assert headers
{{ range $k, $v := .TestCase.OutputTest.Headers }}assert.Equal(t, "{{ $v }}", req.HTTPResponse.Header.Get("{{ $k }}"))
{{ end }}
}
`))

type tplOutputTestCaseData struct {
	*TestCase
	Body, OpName, ResponseString string
}

func (i *TestCase) TestCase(idx int) string {
	var buf bytes.Buffer

	opName := i.API.StructName() + i.TestSuite.title + "Case" + strconv.Itoa(idx+1)

	if i.Params != nil { // input test
		// query test should sort body as form encoded values
		switch i.API.ProtocolPackage() {
		case "query", "ec2":
			m, _ := url.ParseQuery(i.InputTest.Body)
			i.InputTest.Body = m.Encode()
		}

		pBuf, _ := json.Marshal(i.Params)
		input := tplInputTestCaseData{
			TestCase:     i,
			Body:         fmt.Sprintf("%q", i.InputTest.Body),
			OpName:       strings.ToUpper(opName[0:1]) + opName[1:],
			ParamsString: fmt.Sprintf("%q", pBuf),
		}

		if err := tplInputTestCase.Execute(&buf, input); err != nil {
			panic(err)
		}
	} else {
		pBuf, _ := json.Marshal(i.Data)
		output := tplOutputTestCaseData{
			TestCase:       i,
			Body:           fmt.Sprintf("%q", i.OutputTest.Body),
			OpName:         strings.ToUpper(opName[0:1]) + opName[1:],
			ResponseString: fmt.Sprintf("%q", pBuf),
		}

		if err := tplOutputTestCase.Execute(&buf, output); err != nil {
			panic(err)
		}
	}

	return util.GoFmt(buf.String())
}

func GenerateTestSuite(filename string) string {
	inout := "Input"
	if strings.Contains(filename, "/output/") {
		inout = "Output"
	}

	var suites []TestSuite
	f, err := os.Open(filename)
	if err != nil {
		panic(err)
	}

	err = json.NewDecoder(f).Decode(&suites)
	if err != nil {
		panic(err)
	}

	var buf bytes.Buffer
	buf.WriteString("package " + suites[0].ProtocolPackage() + "_test\n\n")

	var innerBuf bytes.Buffer
	innerBuf.WriteString("//\n// Tests begin here\n//\n\n\n")

	for i, suite := range suites {
		svcPrefix := inout + "Service" + strconv.Itoa(i+1)
		suite.API.Metadata.ServiceAbbreviation = svcPrefix + "ProtocolTest"
		suite.API.Operations = map[string]*api.Operation{}
		for idx, c := range suite.Cases {
			c.Given.ExportedName = svcPrefix + "TestCaseOperation" + strconv.Itoa(idx+1)
			suite.API.Operations[c.Given.ExportedName] = c.Given
		}

		suite.API.Setup()
		suite.API.Metadata.EndpointPrefix = suite.API.PackageName()

		for n, s := range suite.API.Shapes {
			s.Rename(svcPrefix + "TestShape" + n)
		}

		svcCode := addImports(suite.API.ServiceGoCode())
		if i == 0 {
			importMatch := reImportRemoval.FindStringSubmatch(svcCode)
			buf.WriteString(importMatch[0] + "\n\n")
			buf.WriteString(preamble + "\n\n")
		}
		svcCode = removeImports(svcCode)
		svcCode = strings.Replace(svcCode, "func New(", "func New"+suite.API.StructName()+"(", -1)

		buf.WriteString(svcCode + "\n\n")
		buf.WriteString(removeImports(suite.API.APIGoCode()) + "\n\n")
		innerBuf.WriteString(suite.TestSuite() + "\n")
	}

	return util.GoFmt(buf.String() + innerBuf.String())
}

func main() {
	out := GenerateTestSuite(os.Args[1])
	if len(os.Args) == 3 {
		f, err := os.Create(os.Args[2])
		defer f.Close()
		if err != nil {
			panic(err)
		}
		f.WriteString(out + "\n")
	} else {
		fmt.Println(out)
	}
}
<|endoftext|>


<commit_before>package operator

import (
	"fmt"
	"log"
	"time"

	"k8s.io/client-go/1.5/kubernetes"
	v1core "k8s.io/client-go/1.5/kubernetes/typed/core/v1"
	"k8s.io/client-go/1.5/pkg/api"
	v1api "k8s.io/client-go/1.5/pkg/api/v1"
	"k8s.io/client-go/1.5/pkg/fields"
	"k8s.io/client-go/1.5/pkg/util/flowcontrol"
	"k8s.io/client-go/1.5/pkg/watch"
	"k8s.io/client-go/1.5/tools/record"

	"github.com/coreos-inc/container-linux-update-operator/internal/constants"
	"github.com/coreos-inc/container-linux-update-operator/internal/k8sutil"
)

const (
	eventReasonRebootFailed = "RebootFailed"
	eventSourceComponent    = "update-operator"
)

var (
	// justRebootedSelector is a selector for combination of annotations
	// expected to be on a node after it has completed a reboot.
	//
	// The update-operator sets constants.AnnotationOkToReboot to true to
	// trigger a reboot, and the update-agent sets
	// constants.AnnotationRebootNeeded and
	// constants.AnnotationRebootInProgress to false when it has finished.
	justRebootedSelector = fields.Set(map[string]string{
		constants.AnnotationOkToReboot:       constants.True,
		constants.AnnotationRebootNeeded:     constants.False,
		constants.AnnotationRebootInProgress: constants.False,
	}).AsSelector()

	// wantsRebootSelector is a selector for the annotation expected to be on a node when it wants to be rebooted.
	//
	// The update-agent sets constants.AnnotationRebootNeeded to true when
	// it would like to reboot, and false when it starts up.
	//
	// If constants.AnnotationRebootPaused is set to "true", the update-agent will not consider it for rebooting.
	wantsRebootSelector = fields.ParseSelectorOrDie(constants.AnnotationRebootNeeded + "==" + constants.True + "," + constants.AnnotationRebootPaused + "!=" + constants.True)
)

type Kontroller struct {
	kc *kubernetes.Clientset
	nc v1core.NodeInterface
	er record.EventRecorder
}

func New() (*Kontroller, error) {
	// set up kubernetes in-cluster client
	kc, err := k8sutil.InClusterClient()
	if err != nil {
		return nil, fmt.Errorf("error creating Kubernetes client: %v", err)
	}

	// node interface
	nc := kc.Nodes()

	// create event emitter
	broadcaster := record.NewBroadcaster()
	broadcaster.StartRecordingToSink(&v1core.EventSinkImpl{Interface: kc.Events("")})
	er := broadcaster.NewRecorder(v1api.EventSource{Component: eventSourceComponent})

	return &Kontroller{kc, nc, er}, nil
}

func (k *Kontroller) Run() error {
	rl := flowcontrol.NewTokenBucketRateLimiter(0.2, 1)
	for {
		rl.Accept()

		nodelist, err := k.nc.List(api.ListOptions{})
		if err != nil {
			log.Printf("Failed listing nodes %v", err)
			continue
		}

		nodes := k8sutil.FilterNodesByAnnotation(nodelist.Items, justRebootedSelector)

		if len(nodes) > 0 {
			log.Printf("Found %d rebooted nodes, setting annotation %q to false", len(nodes), constants.AnnotationOkToReboot)
		}

		for _, n := range nodes {
			if err := k8sutil.SetNodeAnnotations(k.nc, n.Name, map[string]string{
				constants.AnnotationOkToReboot: constants.False,
			}); err != nil {
				log.Printf("Failed setting annotation %q on node %q to false: %v", constants.AnnotationOkToReboot, n.Name, err)
			}
		}

		nodelist, err = k.nc.List(api.ListOptions{})
		if err != nil {
			log.Printf("Failed listing nodes: %v", err)
			continue
		}

		nodes = k8sutil.FilterNodesByAnnotation(nodelist.Items, wantsRebootSelector)

		// pick N of these machines
		// TODO: for now, synchronous with N == 1. might be async w/ a channel in the future to handle N > 1
		if len(nodes) == 0 {
			continue
		}

		n := nodes[0]

		log.Printf("Found %d nodes that need a reboot, rebooting %q", len(nodes), n.Name)

		k.handleReboot(&n)
	}
}

func (k *Kontroller) handleReboot(n *v1api.Node) {
	// node wants to reboot, so let it.
	if err := k8sutil.SetNodeAnnotations(k.nc, n.Name, map[string]string{
		constants.AnnotationOkToReboot: constants.True,
	}); err != nil {
		log.Printf("Failed to set annotation %q on node %q: %v", constants.AnnotationOkToReboot, n.Name, err)
		return
	}

	// wait for it to come back...
	watcher, err := k.nc.Watch(api.ListOptions{
		FieldSelector:   fields.OneTermEqualSelector("metadata.name", n.Name),
		ResourceVersion: n.ResourceVersion,
	})

	conds := []watch.ConditionFunc{
		k8sutil.NodeAnnotationCondition(constants.AnnotationOkToReboot, constants.True),
		k8sutil.NodeAnnotationCondition(constants.AnnotationRebootNeeded, constants.False),
		k8sutil.NodeAnnotationCondition(constants.AnnotationRebootInProgress, constants.False),
	}
	_, err = watch.Until(time.Hour*1, watcher, conds...)
	if err != nil {
		log.Printf("Waiting for label %q on node %q failed: %v", constants.AnnotationOkToReboot, n.Name, err)
		log.Printf("Failed to wait for successful reboot of node %q", n.Name)

		k.er.Eventf(n, api.EventTypeWarning, eventReasonRebootFailed, "Timed out waiting for node to return after a reboot")
	}

	// node rebooted successfully, or at least set the labels we expected from klocksmith after a reboot.
}
<commit_msg>internal/operator: switch to glog<commit_after>package operator

import (
	"fmt"
	"time"

	"github.com/golang/glog"
	"k8s.io/client-go/1.5/kubernetes"
	v1core "k8s.io/client-go/1.5/kubernetes/typed/core/v1"
	"k8s.io/client-go/1.5/pkg/api"
	v1api "k8s.io/client-go/1.5/pkg/api/v1"
	"k8s.io/client-go/1.5/pkg/fields"
	"k8s.io/client-go/1.5/pkg/util/flowcontrol"
	"k8s.io/client-go/1.5/pkg/watch"
	"k8s.io/client-go/1.5/tools/record"

	"github.com/coreos-inc/container-linux-update-operator/internal/constants"
	"github.com/coreos-inc/container-linux-update-operator/internal/k8sutil"
)

const (
	eventReasonRebootFailed = "RebootFailed"
	eventSourceComponent    = "update-operator"
)

var (
	// justRebootedSelector is a selector for combination of annotations
	// expected to be on a node after it has completed a reboot.
	//
	// The update-operator sets constants.AnnotationOkToReboot to true to
	// trigger a reboot, and the update-agent sets
	// constants.AnnotationRebootNeeded and
	// constants.AnnotationRebootInProgress to false when it has finished.
	justRebootedSelector = fields.Set(map[string]string{
		constants.AnnotationOkToReboot:       constants.True,
		constants.AnnotationRebootNeeded:     constants.False,
		constants.AnnotationRebootInProgress: constants.False,
	}).AsSelector()

	// wantsRebootSelector is a selector for the annotation expected to be on a node when it wants to be rebooted.
	//
	// The update-agent sets constants.AnnotationRebootNeeded to true when
	// it would like to reboot, and false when it starts up.
	//
	// If constants.AnnotationRebootPaused is set to "true", the update-agent will not consider it for rebooting.
	wantsRebootSelector = fields.ParseSelectorOrDie(constants.AnnotationRebootNeeded + "==" + constants.True + "," + constants.AnnotationRebootPaused + "!=" + constants.True)
)

type Kontroller struct {
	kc *kubernetes.Clientset
	nc v1core.NodeInterface
	er record.EventRecorder
}

func New() (*Kontroller, error) {
	// set up kubernetes in-cluster client
	kc, err := k8sutil.InClusterClient()
	if err != nil {
		return nil, fmt.Errorf("error creating Kubernetes client: %v", err)
	}

	// node interface
	nc := kc.Nodes()

	// create event emitter
	broadcaster := record.NewBroadcaster()
	broadcaster.StartRecordingToSink(&v1core.EventSinkImpl{Interface: kc.Events("")})
	er := broadcaster.NewRecorder(v1api.EventSource{Component: eventSourceComponent})

	return &Kontroller{kc, nc, er}, nil
}

func (k *Kontroller) Run() error {
	rl := flowcontrol.NewTokenBucketRateLimiter(0.2, 1)
	for {
		rl.Accept()

		nodelist, err := k.nc.List(api.ListOptions{})
		if err != nil {
			glog.Infof("Failed listing nodes %v", err)
			continue
		}

		nodes := k8sutil.FilterNodesByAnnotation(nodelist.Items, justRebootedSelector)

		if len(nodes) > 0 {
			glog.Infof("Found %d rebooted nodes, setting annotation %q to false", len(nodes), constants.AnnotationOkToReboot)
		}

		for _, n := range nodes {
			if err := k8sutil.SetNodeAnnotations(k.nc, n.Name, map[string]string{
				constants.AnnotationOkToReboot: constants.False,
			}); err != nil {
				glog.Infof("Failed setting annotation %q on node %q to false: %v", constants.AnnotationOkToReboot, n.Name, err)
			}
		}

		nodelist, err = k.nc.List(api.ListOptions{})
		if err != nil {
			glog.Infof("Failed listing nodes: %v", err)
			continue
		}

		nodes = k8sutil.FilterNodesByAnnotation(nodelist.Items, wantsRebootSelector)

		// pick N of these machines
		// TODO: for now, synchronous with N == 1. might be async w/ a channel in the future to handle N > 1
		if len(nodes) == 0 {
			continue
		}

		n := nodes[0]

		glog.Infof("Found %d nodes that need a reboot, rebooting %q", len(nodes), n.Name)

		k.handleReboot(&n)
	}
}

func (k *Kontroller) handleReboot(n *v1api.Node) {
	// node wants to reboot, so let it.
	if err := k8sutil.SetNodeAnnotations(k.nc, n.Name, map[string]string{
		constants.AnnotationOkToReboot: constants.True,
	}); err != nil {
		glog.Infof("Failed to set annotation %q on node %q: %v", constants.AnnotationOkToReboot, n.Name, err)
		return
	}

	// wait for it to come back...
	watcher, err := k.nc.Watch(api.ListOptions{
		FieldSelector:   fields.OneTermEqualSelector("metadata.name", n.Name),
		ResourceVersion: n.ResourceVersion,
	})

	conds := []watch.ConditionFunc{
		k8sutil.NodeAnnotationCondition(constants.AnnotationOkToReboot, constants.True),
		k8sutil.NodeAnnotationCondition(constants.AnnotationRebootNeeded, constants.False),
		k8sutil.NodeAnnotationCondition(constants.AnnotationRebootInProgress, constants.False),
	}
	_, err = watch.Until(time.Hour*1, watcher, conds...)
	if err != nil {
		glog.Infof("Waiting for label %q on node %q failed: %v", constants.AnnotationOkToReboot, n.Name, err)
		glog.Infof("Failed to wait for successful reboot of node %q", n.Name)

		k.er.Eventf(n, api.EventTypeWarning, eventReasonRebootFailed, "Timed out waiting for node to return after a reboot")
	}

	// node rebooted successfully, or at least set the labels we expected from klocksmith after a reboot.
}
<|endoftext|>


<commit_before>// Copyright 2015 CoreOS, Inc.
//
// 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 ec2

import (
	"bufio"
	"fmt"
	"net"
	"strings"
	"time"

	"github.com/coreos/coreos-metadata/internal/providers"
	"github.com/coreos/coreos-metadata/internal/retry"
)

func FetchMetadata() (providers.Metadata, error) {
	public, err := fetchIP("public-ipv4")
	if err != nil {
		return providers.Metadata{}, err
	}
	local, err := fetchIP("local-ipv4")
	if err != nil {
		return providers.Metadata{}, err
	}
	hostname, err := fetchString("hostname")
	if err != nil {
		return providers.Metadata{}, err
	}

	sshKeys, err := fetchSshKeys()
	if err != nil {
		return providers.Metadata{}, err
	}

	return providers.Metadata{
		Attributes: map[string]string{
			"EC2_IPV4_LOCAL":  local.String(),
			"EC2_IPV4_PUBLIC": public.String(),
			"EC2_HOSTNAME":    hostname,
		},
		SshKeys: sshKeys,
	}, nil
}

func fetchString(key string) (string, error) {
	body, err := retry.Client{
		InitialBackoff: time.Second,
		MaxBackoff:     time.Second * 5,
		MaxAttempts:    10,
	}.Get("http://169.254.169.254/2009-04-04/meta-data/" + key)
	return string(body), err
}

func fetchIP(key string) (net.IP, error) {
	str, err := fetchString(key)
	if err != nil {
		return nil, err
	}
	if ip := net.ParseIP(str); ip != nil {
		return ip, nil
	} else {
		return nil, fmt.Errorf("couldn't parse %q as IP address", str)
	}
}

func fetchSshKeys() ([]string, error) {
	keydata, err := fetchString("public-keys")
	if err != nil {
		return nil, fmt.Errorf("error reading keys: %v", err)
	}

	scanner := bufio.NewScanner(strings.NewReader(keydata))
	keynames := []string{}
	for scanner.Scan() {
		keynames = append(keynames, scanner.Text())
	}
	if err := scanner.Err(); err != nil {
		return nil, fmt.Errorf("error parsing keys: %v", err)
	}

	keyIDs := make(map[string]string)
	for _, keyname := range keynames {
		tokens := strings.SplitN(keyname, "=", 2)
		if len(tokens) != 2 {
			return nil, fmt.Errorf("malformed public key: %q", keyname)
		}
		keyIDs[tokens[1]] = tokens[0]
	}

	keys := []string{}
	for _, id := range keyIDs {
		sshkey, err := fetchString(fmt.Sprintf("public-keys/%s/openssh-key", id))
		if err != nil {
			return nil, err
		}
		keys = append(keys, sshkey)
	}

	return keys, nil
}
<commit_msg>providers/ec2: Fetch EC2 Instance ID.<commit_after>// Copyright 2015 CoreOS, Inc.
//
// 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 ec2

import (
	"bufio"
	"fmt"
	"net"
	"strings"
	"time"

	"github.com/coreos/coreos-metadata/internal/providers"
	"github.com/coreos/coreos-metadata/internal/retry"
)

func FetchMetadata() (providers.Metadata, error) {
	instanceId, err := fetchString("instance-id")
	if err != nil {
		return providers.Metadata{}, err
	}

	public, err := fetchIP("public-ipv4")
	if err != nil {
		return providers.Metadata{}, err
	}
	local, err := fetchIP("local-ipv4")
	if err != nil {
		return providers.Metadata{}, err
	}
	hostname, err := fetchString("hostname")
	if err != nil {
		return providers.Metadata{}, err
	}

	sshKeys, err := fetchSshKeys()
	if err != nil {
		return providers.Metadata{}, err
	}

	return providers.Metadata{
		Attributes: map[string]string{
			"EC2_INSTANCE_ID": instanceId,
			"EC2_IPV4_LOCAL":  local.String(),
			"EC2_IPV4_PUBLIC": public.String(),
			"EC2_HOSTNAME":    hostname,
		},
		SshKeys: sshKeys,
	}, nil
}

func fetchString(key string) (string, error) {
	body, err := retry.Client{
		InitialBackoff: time.Second,
		MaxBackoff:     time.Second * 5,
		MaxAttempts:    10,
	}.Get("http://169.254.169.254/2009-04-04/meta-data/" + key)
	return string(body), err
}

func fetchIP(key string) (net.IP, error) {
	str, err := fetchString(key)
	if err != nil {
		return nil, err
	}
	if ip := net.ParseIP(str); ip != nil {
		return ip, nil
	} else {
		return nil, fmt.Errorf("couldn't parse %q as IP address", str)
	}
}

func fetchSshKeys() ([]string, error) {
	keydata, err := fetchString("public-keys")
	if err != nil {
		return nil, fmt.Errorf("error reading keys: %v", err)
	}

	scanner := bufio.NewScanner(strings.NewReader(keydata))
	keynames := []string{}
	for scanner.Scan() {
		keynames = append(keynames, scanner.Text())
	}
	if err := scanner.Err(); err != nil {
		return nil, fmt.Errorf("error parsing keys: %v", err)
	}

	keyIDs := make(map[string]string)
	for _, keyname := range keynames {
		tokens := strings.SplitN(keyname, "=", 2)
		if len(tokens) != 2 {
			return nil, fmt.Errorf("malformed public key: %q", keyname)
		}
		keyIDs[tokens[1]] = tokens[0]
	}

	keys := []string{}
	for _, id := range keyIDs {
		sshkey, err := fetchString(fmt.Sprintf("public-keys/%s/openssh-key", id))
		if err != nil {
			return nil, err
		}
		keys = append(keys, sshkey)
	}

	return keys, nil
}
<|endoftext|>