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

// +build go1.3

package lxdclient

import (
	"fmt"
	"io"
	"os"
	"strings"

	"github.com/gorilla/websocket"
	"github.com/juju/errors"
	"github.com/lxc/lxd"
	"github.com/lxc/lxd/shared"

	"github.com/juju/juju/network"
)

// TODO(ericsnow) We probably need to address some of the things that
// get handled in container/lxc/clonetemplate.go.

type rawInstanceClient interface {
	ListContainers() ([]shared.ContainerInfo, error)
	ContainerInfo(name string) (*shared.ContainerInfo, error)
	Init(name string, imgremote string, image string, profiles *[]string, config map[string]string, ephem bool) (*lxd.Response, error)
	Action(name string, action shared.ContainerAction, timeout int, force, stateful bool) (*lxd.Response, error)
	Exec(name string, cmd []string, env map[string]string, stdin io.ReadCloser, stdout io.WriteCloser, stderr io.WriteCloser, controlHandler func(*lxd.Client, *websocket.Conn)) (int, error)
	Delete(name string) (*lxd.Response, error)

	WaitForSuccess(waitURL string) error
	ContainerState(name string) (*shared.ContainerState, error)
}

type instanceClient struct {
	raw    rawInstanceClient
	remote string
}

func (client *instanceClient) addInstance(spec InstanceSpec) error {
	imageRemote := spec.ImageRemote
	if imageRemote == "" {
		imageRemote = client.remote
	}

	imageAlias := spec.Image
	if imageAlias == "" {
		// TODO(ericsnow) Do not have a default?
		// XXX(jam) I don't think we should do this, imageAlias should
		// be mandatory
		imageAlias = "ubuntu"
	}

	var profiles *[]string
	if len(spec.Profiles) > 0 {
		profiles = &spec.Profiles
	}

	// TODO(ericsnow) Copy the image first?

	config := spec.config()
	resp, err := client.raw.Init(spec.Name, imageRemote, imageAlias, profiles, config, spec.Ephemeral)
	if err != nil {
		return errors.Trace(err)
	}

	// Init is an async operation, since the tar -xvf (or whatever) might
	// take a while; the result is an LXD operation id, which we can just
	// wait on until it is finished.
	if err := client.raw.WaitForSuccess(resp.Operation); err != nil {
		// TODO(ericsnow) Handle different failures (from the async
		// operation) differently?
		return errors.Trace(err)
	}

	return nil
}

type execFailure struct {
	cmd    string
	code   int
	stderr string
}

// Error returns the string representation of the error.
func (err execFailure) Error() string {
	return fmt.Sprintf("got non-zero code from %q: (%d) %s", err.cmd, err.code, err.stderr)
}

func (client *instanceClient) exec(spec InstanceSpec, cmd []string) error {
	var env map[string]string

	cmdStr := strings.Join(cmd, " ")
	fmt.Println("running", cmdStr)

	var input, output closingBuffer
	stdin, stdout, stderr := &input, &output, &output
	rc, err := client.raw.Exec(spec.Name, cmd, env, stdin, stdout, stderr, nil)
	if err != nil {
		return errors.Trace(err)
	} else if rc != 0 {
		msg := output.String()
		if msg == "" {
			msg = "<reason unknown>"
		}
		err := &execFailure{
			cmd:    cmdStr,
			code:   rc,
			stderr: msg,
		}
		return errors.Trace(err)
	}

	return nil
}

func (client *instanceClient) chmod(spec InstanceSpec, filename string, mode os.FileMode) error {
	cmd := []string{
		"/bin/chmod",
		fmt.Sprintf("%s", mode),
		filename,
	}

	if err := client.exec(spec, cmd); err != nil {
		return errors.Trace(err)
	}
	return nil
}

func (client *instanceClient) startInstance(spec InstanceSpec) error {
	timeout := -1
	force := false
	// TODO(jam) I believe stateful means to store the memory state when
	// stopping the instance, and then you have to pass stateful=true to
	// start the instance with its state. We aren't supporting that yet.
	stateful := false
	resp, err := client.raw.Action(spec.Name, shared.Start, timeout, force, stateful)
	if err != nil {
		return errors.Trace(err)
	}

	if err := client.raw.WaitForSuccess(resp.Operation); err != nil {
		// TODO(ericsnow) Handle different failures (from the async
		// operation) differently?
		return errors.Trace(err)
	}

	return nil
}

// AddInstance creates a new instance based on the spec's data and
// returns it. The instance will be created using the client.
func (client *instanceClient) AddInstance(spec InstanceSpec) (*Instance, error) {
	if err := client.addInstance(spec); err != nil {
		return nil, errors.Trace(err)
	}

	if err := client.startInstance(spec); err != nil {
		if err := client.removeInstance(spec.Name); err != nil {
			logger.Errorf("could not remove container %q after starting it failed", spec.Name)
		}
		return nil, errors.Trace(err)
	}

	inst, err := client.Instance(spec.Name)
	if err != nil {
		return nil, errors.Trace(err)
	}
	inst.spec = &spec

	return inst, nil
}

// Instance gets the up-to-date info about the given instance
// and returns it.
func (client *instanceClient) Instance(name string) (*Instance, error) {
	info, err := client.raw.ContainerInfo(name)
	if err != nil {
		return nil, errors.Trace(err)
	}

	inst := newInstance(info, nil)
	return inst, nil
}

func (client *instanceClient) Status(name string) (string, error) {
	info, err := client.raw.ContainerInfo(name)
	if err != nil {
		return "", errors.Trace(err)
	}

	return info.Status, nil
}

// Instances sends a request to the API for a list of all instances
// (in the Client's namespace) for which the name starts with the
// provided prefix. The result is also limited to those instances with
// one of the specified statuses (if any).
func (client *instanceClient) Instances(prefix string, statuses ...string) ([]Instance, error) {
	infos, err := client.raw.ListContainers()
	if err != nil {
		return nil, errors.Trace(err)
	}

	var insts []Instance
	for _, info := range infos {
		name := info.Name
		if prefix != "" && !strings.HasPrefix(name, prefix) {
			continue
		}
		if len(statuses) > 0 && !checkStatus(info, statuses) {
			continue
		}

		inst := newInstance(&info, nil)
		insts = append(insts, *inst)
	}
	return insts, nil
}

func checkStatus(info shared.ContainerInfo, statuses []string) bool {
	for _, status := range statuses {
		statusCode := allStatuses[status]
		if info.StatusCode == statusCode {
			return true
		}
	}
	return false
}

// removeInstance sends a request to the API to remove the instance
// with the provided ID. The call blocks until the instance is removed
// (or the request fails).
func (client *instanceClient) removeInstance(name string) error {
	info, err := client.raw.ContainerInfo(name)
	if err != nil {
		return errors.Trace(err)
	}

	//if info.Status.StatusCode != 0 && info.Status.StatusCode != shared.Stopped {
	if info.StatusCode != shared.Stopped {
		timeout := -1
		force := true
		// TODO(jam) I believe stateful means to store the memory state
		// when stopping the instance, and then you have to pass
		// stateful=true to start the instance with its state. We
		// aren't supporting that yet.
		stateful := false
		resp, err := client.raw.Action(name, shared.Stop, timeout, force, stateful)
		if err != nil {
			return errors.Trace(err)
		}

		if err := client.raw.WaitForSuccess(resp.Operation); err != nil {
			// TODO(ericsnow) Handle different failures (from the async
			// operation) differently?
			return errors.Trace(err)
		}
	}

	resp, err := client.raw.Delete(name)
	if err != nil {
		return errors.Trace(err)
	}

	if err := client.raw.WaitForSuccess(resp.Operation); err != nil {
		// TODO(ericsnow) Handle different failures (from the async
		// operation) differently?
		return errors.Trace(err)
	}

	return nil
}

// RemoveInstances sends a request to the API to terminate all
// instances (in the Client's namespace) that match one of the
// provided IDs. If a prefix is provided, only IDs that start with the
// prefix will be considered. The call blocks until all the instances
// are removed or the request fails.
func (client *instanceClient) RemoveInstances(prefix string, names ...string) error {
	if len(names) == 0 {
		return nil
	}

	instances, err := client.Instances(prefix)
	if err != nil {
		return errors.Annotatef(err, "while removing instances %v", names)
	}

	var failed []string
	for _, name := range names {
		if !checkInstanceName(name, instances) {
			// We ignore unknown instance names.
			continue
		}

		if err := client.removeInstance(name); err != nil {
			failed = append(failed, name)
			logger.Errorf("while removing instance %q: %v", name, err)
		}
	}
	if len(failed) != 0 {
		return errors.Errorf("some instance removals failed: %v", failed)
	}
	return nil
}

func checkInstanceName(name string, instances []Instance) bool {
	for _, inst := range instances {
		if inst.Name == name {
			return true
		}
	}
	return false
}

func lxdAddressFamilyToNetworkType(addrFamily string) (network.AddressType, error) {
	switch addrFamily {
	case "inet":
		return network.IPv4Address, nil
	case "inet6":
		return network.IPv6Address, nil
	default:
		return "", errors.Errorf("invalid LXD address family type %s", addrFamily)
	}
}

// Addresses returns the list of network.Addresses for this instance. It
// converts the information that LXD tracks into the Juju network model.
func (client *instanceClient) Addresses(name string) ([]network.Address, error) {
	state, err := client.raw.ContainerState(name)
	if err != nil {
		return nil, err
	}

	networks := state.Network
	if networks == nil {
		return []network.Address{}, nil
	}

	addrs := []network.Address{}

	for name, net := range networks {
		for _, addr := range net.Addresses {
			type_, err := lxdAddressFamilyToNetworkType(addr.Family)
			if err != nil {
				return nil, err
			}

			// TODO(jam) 2016-02-26 For multi NIC support we'll
			// need to start tracking scope and space information.
			// But it doesn't make sense for lxdclient to be space
			// aware. That information at best would need to be
			// passed in. (only the DB really knows what spaces are
			// available, and lxdclient doesn't talk to the DB.)
			addrs = append(addrs, network.Address{
				Value:       addr.Address,
				Type:        type_,
				NetworkName: name,
			})
		}
	}

	return addrs, nil
}
<commit_msg>We don't need a default image alias<commit_after>// Copyright 2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

// +build go1.3

package lxdclient

import (
	"fmt"
	"io"
	"os"
	"strings"

	"github.com/gorilla/websocket"
	"github.com/juju/errors"
	"github.com/lxc/lxd"
	"github.com/lxc/lxd/shared"

	"github.com/juju/juju/network"
)

// TODO(ericsnow) We probably need to address some of the things that
// get handled in container/lxc/clonetemplate.go.

type rawInstanceClient interface {
	ListContainers() ([]shared.ContainerInfo, error)
	ContainerInfo(name string) (*shared.ContainerInfo, error)
	Init(name string, imgremote string, image string, profiles *[]string, config map[string]string, ephem bool) (*lxd.Response, error)
	Action(name string, action shared.ContainerAction, timeout int, force, stateful bool) (*lxd.Response, error)
	Exec(name string, cmd []string, env map[string]string, stdin io.ReadCloser, stdout io.WriteCloser, stderr io.WriteCloser, controlHandler func(*lxd.Client, *websocket.Conn)) (int, error)
	Delete(name string) (*lxd.Response, error)

	WaitForSuccess(waitURL string) error
	ContainerState(name string) (*shared.ContainerState, error)
}

type instanceClient struct {
	raw    rawInstanceClient
	remote string
}

func (client *instanceClient) addInstance(spec InstanceSpec) error {
	imageRemote := spec.ImageRemote
	if imageRemote == "" {
		imageRemote = client.remote
	}

	imageAlias := spec.Image

	var profiles *[]string
	if len(spec.Profiles) > 0 {
		profiles = &spec.Profiles
	}

	// TODO(ericsnow) Copy the image first?

	config := spec.config()
	resp, err := client.raw.Init(spec.Name, imageRemote, imageAlias, profiles, config, spec.Ephemeral)
	if err != nil {
		return errors.Trace(err)
	}

	// Init is an async operation, since the tar -xvf (or whatever) might
	// take a while; the result is an LXD operation id, which we can just
	// wait on until it is finished.
	if err := client.raw.WaitForSuccess(resp.Operation); err != nil {
		// TODO(ericsnow) Handle different failures (from the async
		// operation) differently?
		return errors.Trace(err)
	}

	return nil
}

type execFailure struct {
	cmd    string
	code   int
	stderr string
}

// Error returns the string representation of the error.
func (err execFailure) Error() string {
	return fmt.Sprintf("got non-zero code from %q: (%d) %s", err.cmd, err.code, err.stderr)
}

func (client *instanceClient) exec(spec InstanceSpec, cmd []string) error {
	var env map[string]string

	cmdStr := strings.Join(cmd, " ")
	fmt.Println("running", cmdStr)

	var input, output closingBuffer
	stdin, stdout, stderr := &input, &output, &output
	rc, err := client.raw.Exec(spec.Name, cmd, env, stdin, stdout, stderr, nil)
	if err != nil {
		return errors.Trace(err)
	} else if rc != 0 {
		msg := output.String()
		if msg == "" {
			msg = "<reason unknown>"
		}
		err := &execFailure{
			cmd:    cmdStr,
			code:   rc,
			stderr: msg,
		}
		return errors.Trace(err)
	}

	return nil
}

func (client *instanceClient) chmod(spec InstanceSpec, filename string, mode os.FileMode) error {
	cmd := []string{
		"/bin/chmod",
		fmt.Sprintf("%s", mode),
		filename,
	}

	if err := client.exec(spec, cmd); err != nil {
		return errors.Trace(err)
	}
	return nil
}

func (client *instanceClient) startInstance(spec InstanceSpec) error {
	timeout := -1
	force := false
	// TODO(jam) I believe stateful means to store the memory state when
	// stopping the instance, and then you have to pass stateful=true to
	// start the instance with its state. We aren't supporting that yet.
	stateful := false
	resp, err := client.raw.Action(spec.Name, shared.Start, timeout, force, stateful)
	if err != nil {
		return errors.Trace(err)
	}

	if err := client.raw.WaitForSuccess(resp.Operation); err != nil {
		// TODO(ericsnow) Handle different failures (from the async
		// operation) differently?
		return errors.Trace(err)
	}

	return nil
}

// AddInstance creates a new instance based on the spec's data and
// returns it. The instance will be created using the client.
func (client *instanceClient) AddInstance(spec InstanceSpec) (*Instance, error) {
	if err := client.addInstance(spec); err != nil {
		return nil, errors.Trace(err)
	}

	if err := client.startInstance(spec); err != nil {
		if err := client.removeInstance(spec.Name); err != nil {
			logger.Errorf("could not remove container %q after starting it failed", spec.Name)
		}
		return nil, errors.Trace(err)
	}

	inst, err := client.Instance(spec.Name)
	if err != nil {
		return nil, errors.Trace(err)
	}
	inst.spec = &spec

	return inst, nil
}

// Instance gets the up-to-date info about the given instance
// and returns it.
func (client *instanceClient) Instance(name string) (*Instance, error) {
	info, err := client.raw.ContainerInfo(name)
	if err != nil {
		return nil, errors.Trace(err)
	}

	inst := newInstance(info, nil)
	return inst, nil
}

func (client *instanceClient) Status(name string) (string, error) {
	info, err := client.raw.ContainerInfo(name)
	if err != nil {
		return "", errors.Trace(err)
	}

	return info.Status, nil
}

// Instances sends a request to the API for a list of all instances
// (in the Client's namespace) for which the name starts with the
// provided prefix. The result is also limited to those instances with
// one of the specified statuses (if any).
func (client *instanceClient) Instances(prefix string, statuses ...string) ([]Instance, error) {
	infos, err := client.raw.ListContainers()
	if err != nil {
		return nil, errors.Trace(err)
	}

	var insts []Instance
	for _, info := range infos {
		name := info.Name
		if prefix != "" && !strings.HasPrefix(name, prefix) {
			continue
		}
		if len(statuses) > 0 && !checkStatus(info, statuses) {
			continue
		}

		inst := newInstance(&info, nil)
		insts = append(insts, *inst)
	}
	return insts, nil
}

func checkStatus(info shared.ContainerInfo, statuses []string) bool {
	for _, status := range statuses {
		statusCode := allStatuses[status]
		if info.StatusCode == statusCode {
			return true
		}
	}
	return false
}

// removeInstance sends a request to the API to remove the instance
// with the provided ID. The call blocks until the instance is removed
// (or the request fails).
func (client *instanceClient) removeInstance(name string) error {
	info, err := client.raw.ContainerInfo(name)
	if err != nil {
		return errors.Trace(err)
	}

	//if info.Status.StatusCode != 0 && info.Status.StatusCode != shared.Stopped {
	if info.StatusCode != shared.Stopped {
		timeout := -1
		force := true
		// TODO(jam) I believe stateful means to store the memory state
		// when stopping the instance, and then you have to pass
		// stateful=true to start the instance with its state. We
		// aren't supporting that yet.
		stateful := false
		resp, err := client.raw.Action(name, shared.Stop, timeout, force, stateful)
		if err != nil {
			return errors.Trace(err)
		}

		if err := client.raw.WaitForSuccess(resp.Operation); err != nil {
			// TODO(ericsnow) Handle different failures (from the async
			// operation) differently?
			return errors.Trace(err)
		}
	}

	resp, err := client.raw.Delete(name)
	if err != nil {
		return errors.Trace(err)
	}

	if err := client.raw.WaitForSuccess(resp.Operation); err != nil {
		// TODO(ericsnow) Handle different failures (from the async
		// operation) differently?
		return errors.Trace(err)
	}

	return nil
}

// RemoveInstances sends a request to the API to terminate all
// instances (in the Client's namespace) that match one of the
// provided IDs. If a prefix is provided, only IDs that start with the
// prefix will be considered. The call blocks until all the instances
// are removed or the request fails.
func (client *instanceClient) RemoveInstances(prefix string, names ...string) error {
	if len(names) == 0 {
		return nil
	}

	instances, err := client.Instances(prefix)
	if err != nil {
		return errors.Annotatef(err, "while removing instances %v", names)
	}

	var failed []string
	for _, name := range names {
		if !checkInstanceName(name, instances) {
			// We ignore unknown instance names.
			continue
		}

		if err := client.removeInstance(name); err != nil {
			failed = append(failed, name)
			logger.Errorf("while removing instance %q: %v", name, err)
		}
	}
	if len(failed) != 0 {
		return errors.Errorf("some instance removals failed: %v", failed)
	}
	return nil
}

func checkInstanceName(name string, instances []Instance) bool {
	for _, inst := range instances {
		if inst.Name == name {
			return true
		}
	}
	return false
}

func lxdAddressFamilyToNetworkType(addrFamily string) (network.AddressType, error) {
	switch addrFamily {
	case "inet":
		return network.IPv4Address, nil
	case "inet6":
		return network.IPv6Address, nil
	default:
		return "", errors.Errorf("invalid LXD address family type %s", addrFamily)
	}
}

// Addresses returns the list of network.Addresses for this instance. It
// converts the information that LXD tracks into the Juju network model.
func (client *instanceClient) Addresses(name string) ([]network.Address, error) {
	state, err := client.raw.ContainerState(name)
	if err != nil {
		return nil, err
	}

	networks := state.Network
	if networks == nil {
		return []network.Address{}, nil
	}

	addrs := []network.Address{}

	for name, net := range networks {
		for _, addr := range net.Addresses {
			type_, err := lxdAddressFamilyToNetworkType(addr.Family)
			if err != nil {
				return nil, err
			}

			// TODO(jam) 2016-02-26 For multi NIC support we'll
			// need to start tracking scope and space information.
			// But it doesn't make sense for lxdclient to be space
			// aware. That information at best would need to be
			// passed in. (only the DB really knows what spaces are
			// available, and lxdclient doesn't talk to the DB.)
			addrs = append(addrs, network.Address{
				Value:       addr.Address,
				Type:        type_,
				NetworkName: name,
			})
		}
	}

	return addrs, nil
}
<|endoftext|>


<commit_before>// Copyright © 2017 Aqua Security Software Ltd. <info@aquasec.com>
//
// 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 check

import (
	"fmt"
	"os"
	"regexp"
	"strconv"
	"strings"
)

// test:
// flag: OPTION
// set: (true|false)
// compare:
//   op: (eq|gt|gte|lt|lte|has)
//   value: val

type binOp string

const (
	and binOp = "and"
	or        = "or"
)

type testItem struct {
	Flag    string
	Output  string
	Value   string
	Set     bool
	Compare compare
}

type compare struct {
	Op    string
	Value string
}

func (t *testItem) execute(s string) (result bool) {
	result = false
	match := strings.Contains(s, t.Flag)

	if t.Set {
		var flagVal string
		isset := match

		if isset && t.Compare.Op != "" {
			// Expects flags in the form;
			// --flag=somevalue
			// --flag
			// somevalue
			//pttn := `(` + t.Flag + `)(=)*([^\s,]*) *`
			pttn := `(` + t.Flag + `)(=)*([^\s]*) *`
			flagRe := regexp.MustCompile(pttn)
			vals := flagRe.FindStringSubmatch(s)

			if len(vals) > 0 {
				if vals[3] != "" {
					flagVal = vals[3]
				} else {
					flagVal = vals[1]
				}
			} else {
				fmt.Fprintf(os.Stderr, "invalid flag in testitem definition")
				os.Exit(1)
			}

			switch t.Compare.Op {
			case "eq":
				result = flagVal == t.Compare.Value

			case "noteq":
				result = !(flagVal == t.Compare.Value)

			case "gt":
				a, b := toNumeric(flagVal, t.Compare.Value)
				result = a > b

			case "gte":
				a, b := toNumeric(flagVal, t.Compare.Value)
				result = a >= b

			case "lt":
				a, b := toNumeric(flagVal, t.Compare.Value)
				result = a < b

			case "lte":
				a, b := toNumeric(flagVal, t.Compare.Value)
				result = a <= b

			case "has":
				result = strings.Contains(flagVal, t.Compare.Value)

			case "nothave":
				result = !strings.Contains(flagVal, t.Compare.Value)
			}
		} else {
			result = isset
		}

	} else {
		notset := !match
		result = notset
	}

	return
}

type tests struct {
	TestItems []*testItem `yaml:"test_items"`
	BinOp     binOp       `yaml:"bin_op"`
}

func (ts *tests) execute(s string) (result bool) {
	res := make([]bool, len(ts.TestItems))

	for i, t := range ts.TestItems {
		res[i] = t.execute(s)
	}

	// If no binary operation is specified, default to AND
	switch ts.BinOp {
	default:
		fmt.Fprintf(os.Stderr, "unknown binary operator for tests %s\n", ts.BinOp)
		os.Exit(1)
	case and, "":
		result = true
		for i := range res {
			result = result && res[i]
		}
	case or:
		result = false
		for i := range res {
			result = result || res[i]
		}
	}

	return
}

func toNumeric(a, b string) (c, d int) {
	var err error
	c, err = strconv.Atoi(a)
	if err != nil {
		fmt.Fprintf(os.Stderr, "error converting %s: %s\n", a, err)
		os.Exit(1)
	}
	d, err = strconv.Atoi(b)
	if err != nil {
		fmt.Fprintf(os.Stderr, "error converting %s: %s\n", b, err)
		os.Exit(1)
	}

	return c, d
}
<commit_msg>Do case insensitive comparaison for booleans - Fix #125<commit_after>// Copyright © 2017 Aqua Security Software Ltd. <info@aquasec.com>
//
// 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 check

import (
	"fmt"
	"os"
	"regexp"
	"strconv"
	"strings"
)

// test:
// flag: OPTION
// set: (true|false)
// compare:
//   op: (eq|gt|gte|lt|lte|has)
//   value: val

type binOp string

const (
	and binOp = "and"
	or        = "or"
)

type testItem struct {
	Flag    string
	Output  string
	Value   string
	Set     bool
	Compare compare
}

type compare struct {
	Op    string
	Value string
}

func (t *testItem) execute(s string) (result bool) {
	result = false
	match := strings.Contains(s, t.Flag)

	if t.Set {
		var flagVal string
		isset := match

		if isset && t.Compare.Op != "" {
			// Expects flags in the form;
			// --flag=somevalue
			// --flag
			// somevalue
			//pttn := `(` + t.Flag + `)(=)*([^\s,]*) *`
			pttn := `(` + t.Flag + `)(=)*([^\s]*) *`
			flagRe := regexp.MustCompile(pttn)
			vals := flagRe.FindStringSubmatch(s)

			if len(vals) > 0 {
				if vals[3] != "" {
					flagVal = vals[3]
				} else {
					flagVal = vals[1]
				}
			} else {
				fmt.Fprintf(os.Stderr, "invalid flag in testitem definition")
				os.Exit(1)
			}

			switch t.Compare.Op {
			case "eq":
				// Do case insensitive comparaison for booleans ...
				if value == "false" || value == "true" {
					result = value == t.Compare.Value
				} else {
					result = flagVal == t.Compare.Value
				}

			case "noteq":
				// Do case insensitive comparaison for booleans ...
				if value == "false" || value == "true" {
					result = !(value == t.Compare.Value)
				} else {
					result = !(flagVal == t.Compare.Value)
				}

			case "gt":
				a, b := toNumeric(flagVal, t.Compare.Value)
				result = a > b

			case "gte":
				a, b := toNumeric(flagVal, t.Compare.Value)
				result = a >= b

			case "lt":
				a, b := toNumeric(flagVal, t.Compare.Value)
				result = a < b

			case "lte":
				a, b := toNumeric(flagVal, t.Compare.Value)
				result = a <= b

			case "has":
				result = strings.Contains(flagVal, t.Compare.Value)

			case "nothave":
				result = !strings.Contains(flagVal, t.Compare.Value)
			}
		} else {
			result = isset
		}

	} else {
		notset := !match
		result = notset
	}

	return
}

type tests struct {
	TestItems []*testItem `yaml:"test_items"`
	BinOp     binOp       `yaml:"bin_op"`
}

func (ts *tests) execute(s string) (result bool) {
	res := make([]bool, len(ts.TestItems))

	for i, t := range ts.TestItems {
		res[i] = t.execute(s)
	}

	// If no binary operation is specified, default to AND
	switch ts.BinOp {
	default:
		fmt.Fprintf(os.Stderr, "unknown binary operator for tests %s\n", ts.BinOp)
		os.Exit(1)
	case and, "":
		result = true
		for i := range res {
			result = result && res[i]
		}
	case or:
		result = false
		for i := range res {
			result = result || res[i]
		}
	}

	return
}

func toNumeric(a, b string) (c, d int) {
	var err error
	c, err = strconv.Atoi(a)
	if err != nil {
		fmt.Fprintf(os.Stderr, "error converting %s: %s\n", a, err)
		os.Exit(1)
	}
	d, err = strconv.Atoi(b)
	if err != nil {
		fmt.Fprintf(os.Stderr, "error converting %s: %s\n", b, err)
		os.Exit(1)
	}

	return c, d
}
<|endoftext|>


<commit_before>package server

import (
	"bytes"
	"compress/gzip"
	"crypto/sha256"
	"encoding/base64"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"path/filepath"
	"strconv"
	"strings"
	"text/template"
	"time"

	"github.com/dsnet/compress/brotli"
	"github.com/khlieng/dispatch/assets"
	"github.com/tdewolff/minify/v2"
	"github.com/tdewolff/minify/v2/html"
)

const longCacheControl = "public, max-age=31536000, immutable"
const disabledCacheControl = "no-cache, no-store, must-revalidate"

type File struct {
	Data         []byte
	Length       string
	GzipData     []byte
	GzipLength   string
	Hash         string
	ContentType  string
	CacheControl string
	Compressed   bool
}

type h2PushAsset struct {
	path string
	hash string
}

func newH2PushAsset(name string) h2PushAsset {
	return h2PushAsset{
		path: name,
		hash: strings.Split(name, ".")[1],
	}
}

var (
	files = map[string]*File{}

	indexPage          []byte
	indexPageLen       string
	inlineScriptSha256 string
	serviceWorker      []byte

	h2PushAssets      []h2PushAsset
	h2PushCookieValue string

	contentTypes = map[string]string{
		".js":    "text/javascript",
		".css":   "text/css",
		".woff2": "font/woff2",
		".woff":  "application/font-woff",
		".ttf":   "application/x-font-ttf",
		".png":   "image/png",
		".ico":   "image/x-icon",
		".json":  "application/json",
	}

	robots = []byte("User-agent: *\nDisallow: /")

	hstsHeader string
	cspEnabled bool
)

func (d *Dispatch) initFileServer() {
	cfg := d.Config()

	if cfg.Dev {
		renderIndexPage(indexTemplateData{
			Scripts: []string{"/boot.js", "/main.js"},
		})
	} else {
		bootloader := decompressedAsset(findAssetName("boot*.js"))
		runtime := decompressedAsset(findAssetName("runtime*.js"))

		inlineScript := string(bootloader) + string(runtime)

		hash := sha256.New()
		hash.Write(bootloader)
		hash.Write(runtime)
		inlineScriptSha256 = base64.StdEncoding.EncodeToString(hash.Sum(nil))

		indexStylesheet := "/" + findAssetName("main*.css")
		indexScripts := []string{
			"/" + findAssetName("vendors*.js"),
			"/" + findAssetName("main*.js"),
		}

		h2PushAssets = []h2PushAsset{
			newH2PushAsset(indexStylesheet),
			newH2PushAsset(indexScripts[0]),
			newH2PushAsset(indexScripts[1]),
		}

		for _, asset := range h2PushAssets {
			h2PushCookieValue += asset.hash
		}

		ignoreAssets := []string{
			findAssetName("runtime*.js"),
			findAssetName("boot*.js"),
			"sw.js",
		}

	outer:
		for _, asset := range assets.AssetNames() {
			assetName := strings.TrimSuffix(asset, ".br")

			for _, ignored := range ignoreAssets {
				if ignored == assetName {
					continue outer
				}
			}

			file := &File{
				ContentType:  contentTypes[filepath.Ext(assetName)],
				CacheControl: longCacheControl,
				Compressed:   strings.HasSuffix(asset, ".br"),
			}

			data, err := assets.Asset(asset)
			fatalErr(err)
			file.Data = data
			file.Length = strconv.Itoa(len(data))

			if file.Compressed {
				file.GzipData = gzipAsset(data)
				file.GzipLength = strconv.Itoa(len(file.GzipData))
			}

			files["/"+assetName] = file
		}

		renderIndexPage(indexTemplateData{
			Stylesheet:   indexStylesheet,
			InlineScript: inlineScript,
			Scripts:      indexScripts,
		})

		serviceWorker = decompressedAsset("sw.js")
		hash.Reset()
		hash.Write(indexPage)
		indexHash := base64.StdEncoding.EncodeToString(hash.Sum(nil))

		serviceWorker = append(serviceWorker, []byte(`
workbox.precaching.precacheAndRoute([{
	revision: '`+indexHash+`',
	url: '/'
}]);
workbox.routing.registerNavigationRoute('/');`)...)

		if cfg.HTTPS.HSTS.Enabled && cfg.HTTPS.Enabled {
			hstsHeader = "max-age=" + cfg.HTTPS.HSTS.MaxAge

			if cfg.HTTPS.HSTS.IncludeSubdomains {
				hstsHeader += "; includeSubDomains"
			}
			if cfg.HTTPS.HSTS.Preload {
				hstsHeader += "; preload"
			}
		}

		cspEnabled = true
	}
}

func renderIndexPage(data indexTemplateData) {
	tmpl, err := template.New("").Parse(indexTemplate)
	fatalErr(err)

	m := minify.New()
	m.AddFunc("text/html", html.Minify)

	buf := &bytes.Buffer{}
	gzw, err := gzip.NewWriterLevel(buf, gzip.BestCompression)
	fatalErr(err)
	mw := m.Writer("text/html", gzw)

	fatalErr(tmpl.Execute(mw, data))

	fatalErr(mw.Close())
	fatalErr(gzw.Close())

	indexPage = buf.Bytes()
	indexPageLen = strconv.Itoa(len(indexPage))
}

func findAssetName(glob string) string {
	for _, assetName := range assets.AssetNames() {
		assetName = strings.TrimSuffix(assetName, ".br")

		if m, _ := filepath.Match(glob, assetName); m {
			return assetName
		}
	}
	return ""
}

func decompressAsset(data []byte) []byte {
	br, err := brotli.NewReader(bytes.NewReader(data), nil)
	fatalErr(err)

	buf := &bytes.Buffer{}
	io.Copy(buf, br)
	return buf.Bytes()
}

func decompressedAsset(name string) []byte {
	asset, err := assets.Asset(name + ".br")
	fatalErr(err)
	return decompressAsset(asset)
}

func gzipAsset(data []byte) []byte {
	br, err := brotli.NewReader(bytes.NewReader(data), nil)
	fatalErr(err)

	buf := &bytes.Buffer{}
	gzw, err := gzip.NewWriterLevel(buf, gzip.BestCompression)
	fatalErr(err)

	io.Copy(gzw, br)
	gzw.Close()
	return buf.Bytes()
}

func (d *Dispatch) serveFiles(w http.ResponseWriter, r *http.Request) {
	if r.URL.Path == "/" {
		d.serveIndex(w, r)
	} else if file, ok := files[r.URL.Path]; ok {
		d.serveFile(w, r, file)
	} else if r.URL.Path == "/sw.js" {
		w.Header().Set("Cache-Control", disabledCacheControl)
		w.Header().Set("Content-Type", "text/javascript")
		w.Header().Set("Content-Length", strconv.Itoa(len(serviceWorker)))
		w.Write(serviceWorker)
	} else if r.URL.Path == "/robots.txt" {
		w.Header().Set("Content-Type", "text/plain")
		w.Header().Set("Content-Length", strconv.Itoa(len(robots)))
		w.Write(robots)
	} else {
		d.serveIndex(w, r)
	}
}

func (d *Dispatch) serveIndex(w http.ResponseWriter, r *http.Request) {
	if pusher, ok := w.(http.Pusher); ok {
		options := &http.PushOptions{
			Header: http.Header{
				"Accept-Encoding": r.Header["Accept-Encoding"],
			},
		}

		cookie, err := r.Cookie("push")
		if err != nil {
			for _, asset := range h2PushAssets {
				pusher.Push(asset.path, options)
			}

			setPushCookie(w, r)
		} else {
			pushed := false

			i := 0
			for _, asset := range h2PushAssets {
				if len(cookie.Value) >= i+len(asset.hash) &&
					asset.hash != cookie.Value[i:i+len(asset.hash)] {
					i += len(asset.hash)
					pusher.Push(asset.path, options)
					pushed = true
				}
			}

			if pushed {
				setPushCookie(w, r)
			}
		}
	}

	if cspEnabled {
		var wsSrc string
		if r.TLS != nil {
			wsSrc = "wss://" + r.Host
		} else {
			wsSrc = "ws://" + r.Host
		}

		csp := []string{
			"default-src 'none'",
			"script-src 'self' 'sha256-" + inlineScriptSha256 + "'",
			"style-src 'self' 'unsafe-inline'",
			"font-src 'self'",
			"img-src 'self'",
			"manifest-src 'self'",
			"connect-src 'self' " + wsSrc,
			"worker-src 'self'",
		}

		w.Header().Set("Content-Security-Policy", strings.Join(csp, "; "))
	}

	w.Header().Set("Content-Type", "text/html")
	w.Header().Set("Cache-Control", disabledCacheControl)
	w.Header().Set("X-Content-Type-Options", "nosniff")
	w.Header().Set("X-Frame-Options", "deny")
	w.Header().Set("X-XSS-Protection", "1; mode=block")
	w.Header().Set("Referrer-Policy", "same-origin")

	if hstsHeader != "" {
		w.Header().Set("Strict-Transport-Security", hstsHeader)
	}

	for k, v := range d.Config().Headers {
		w.Header().Set(k, v)
	}

	if strings.Contains(r.Header.Get("Accept-Encoding"), "gzip") {
		w.Header().Set("Content-Encoding", "gzip")
		w.Header().Set("Content-Length", indexPageLen)
		w.Write(indexPage)
	} else {
		serveDecompressed(w, indexPage)
	}
}

func setPushCookie(w http.ResponseWriter, r *http.Request) {
	http.SetCookie(w, &http.Cookie{
		Name:     "push",
		Value:    h2PushCookieValue,
		Path:     "/",
		Expires:  time.Now().AddDate(1, 0, 0),
		HttpOnly: true,
		Secure:   r.TLS != nil,
	})
}

func (d *Dispatch) serveFile(w http.ResponseWriter, r *http.Request, file *File) {
	w.Header().Set("Cache-Control", file.CacheControl)
	w.Header().Set("Content-Type", file.ContentType)

	if file.Compressed && strings.Contains(r.Header.Get("Accept-Encoding"), "br") {
		w.Header().Set("Content-Encoding", "br")
		w.Header().Set("Content-Length", file.Length)
		w.Write(file.Data)
	} else if file.Compressed && strings.Contains(r.Header.Get("Accept-Encoding"), "gzip") {
		w.Header().Set("Content-Encoding", "gzip")
		w.Header().Set("Content-Length", file.GzipLength)
		w.Write(file.GzipData)
	} else if !file.Compressed {
		w.Header().Set("Content-Length", file.Length)
		w.Write(file.Data)
	} else {
		serveDecompressed(w, file.GzipData)
	}
}

func serveDecompressed(w http.ResponseWriter, asset []byte) {
	gzr, err := gzip.NewReader(bytes.NewReader(asset))
	buf, err := ioutil.ReadAll(gzr)
	if err != nil {
		http.Error(w, "", http.StatusInternalServerError)
		return
	}

	w.Header().Set("Content-Length", strconv.Itoa(len(buf)))
	w.Write(buf)
}

func fatalErr(err error) {
	if err != nil {
		log.Fatal(err)
	}
}
<commit_msg>Fix h2 push hash check<commit_after>package server

import (
	"bytes"
	"compress/gzip"
	"crypto/sha256"
	"encoding/base64"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"path/filepath"
	"strconv"
	"strings"
	"text/template"
	"time"

	"github.com/dsnet/compress/brotli"
	"github.com/khlieng/dispatch/assets"
	"github.com/tdewolff/minify/v2"
	"github.com/tdewolff/minify/v2/html"
)

const longCacheControl = "public, max-age=31536000, immutable"
const disabledCacheControl = "no-cache, no-store, must-revalidate"

type File struct {
	Data         []byte
	Length       string
	GzipData     []byte
	GzipLength   string
	Hash         string
	ContentType  string
	CacheControl string
	Compressed   bool
}

type h2PushAsset struct {
	path string
	hash string
}

func newH2PushAsset(name string) h2PushAsset {
	return h2PushAsset{
		path: name,
		hash: strings.Split(name, ".")[1],
	}
}

var (
	files = map[string]*File{}

	indexPage          []byte
	indexPageLen       string
	inlineScriptSha256 string
	serviceWorker      []byte

	h2PushAssets      []h2PushAsset
	h2PushCookieValue string

	contentTypes = map[string]string{
		".js":    "text/javascript",
		".css":   "text/css",
		".woff2": "font/woff2",
		".woff":  "application/font-woff",
		".ttf":   "application/x-font-ttf",
		".png":   "image/png",
		".ico":   "image/x-icon",
		".json":  "application/json",
	}

	robots = []byte("User-agent: *\nDisallow: /")

	hstsHeader string
	cspEnabled bool
)

func (d *Dispatch) initFileServer() {
	cfg := d.Config()

	if cfg.Dev {
		renderIndexPage(indexTemplateData{
			Scripts: []string{"/boot.js", "/main.js"},
		})
	} else {
		bootloader := decompressedAsset(findAssetName("boot*.js"))
		runtime := decompressedAsset(findAssetName("runtime*.js"))

		inlineScript := string(bootloader) + string(runtime)

		hash := sha256.New()
		hash.Write(bootloader)
		hash.Write(runtime)
		inlineScriptSha256 = base64.StdEncoding.EncodeToString(hash.Sum(nil))

		indexStylesheet := "/" + findAssetName("main*.css")
		indexScripts := []string{
			"/" + findAssetName("vendors*.js"),
			"/" + findAssetName("main*.js"),
		}

		h2PushAssets = []h2PushAsset{
			newH2PushAsset(indexStylesheet),
			newH2PushAsset(indexScripts[0]),
			newH2PushAsset(indexScripts[1]),
		}

		for _, asset := range h2PushAssets {
			h2PushCookieValue += asset.hash
		}

		ignoreAssets := []string{
			findAssetName("runtime*.js"),
			findAssetName("boot*.js"),
			"sw.js",
		}

	outer:
		for _, asset := range assets.AssetNames() {
			assetName := strings.TrimSuffix(asset, ".br")

			for _, ignored := range ignoreAssets {
				if ignored == assetName {
					continue outer
				}
			}

			file := &File{
				ContentType:  contentTypes[filepath.Ext(assetName)],
				CacheControl: longCacheControl,
				Compressed:   strings.HasSuffix(asset, ".br"),
			}

			data, err := assets.Asset(asset)
			fatalErr(err)
			file.Data = data
			file.Length = strconv.Itoa(len(data))

			if file.Compressed {
				file.GzipData = gzipAsset(data)
				file.GzipLength = strconv.Itoa(len(file.GzipData))
			}

			files["/"+assetName] = file
		}

		renderIndexPage(indexTemplateData{
			Stylesheet:   indexStylesheet,
			InlineScript: inlineScript,
			Scripts:      indexScripts,
		})

		serviceWorker = decompressedAsset("sw.js")
		hash.Reset()
		hash.Write(indexPage)
		indexHash := base64.StdEncoding.EncodeToString(hash.Sum(nil))

		serviceWorker = append(serviceWorker, []byte(`
workbox.precaching.precacheAndRoute([{
	revision: '`+indexHash+`',
	url: '/'
}]);
workbox.routing.registerNavigationRoute('/');`)...)

		if cfg.HTTPS.HSTS.Enabled && cfg.HTTPS.Enabled {
			hstsHeader = "max-age=" + cfg.HTTPS.HSTS.MaxAge

			if cfg.HTTPS.HSTS.IncludeSubdomains {
				hstsHeader += "; includeSubDomains"
			}
			if cfg.HTTPS.HSTS.Preload {
				hstsHeader += "; preload"
			}
		}

		cspEnabled = true
	}
}

func renderIndexPage(data indexTemplateData) {
	tmpl, err := template.New("").Parse(indexTemplate)
	fatalErr(err)

	m := minify.New()
	m.AddFunc("text/html", html.Minify)

	buf := &bytes.Buffer{}
	gzw, err := gzip.NewWriterLevel(buf, gzip.BestCompression)
	fatalErr(err)
	mw := m.Writer("text/html", gzw)

	fatalErr(tmpl.Execute(mw, data))

	fatalErr(mw.Close())
	fatalErr(gzw.Close())

	indexPage = buf.Bytes()
	indexPageLen = strconv.Itoa(len(indexPage))
}

func findAssetName(glob string) string {
	for _, assetName := range assets.AssetNames() {
		assetName = strings.TrimSuffix(assetName, ".br")

		if m, _ := filepath.Match(glob, assetName); m {
			return assetName
		}
	}
	return ""
}

func decompressAsset(data []byte) []byte {
	br, err := brotli.NewReader(bytes.NewReader(data), nil)
	fatalErr(err)

	buf := &bytes.Buffer{}
	io.Copy(buf, br)
	return buf.Bytes()
}

func decompressedAsset(name string) []byte {
	asset, err := assets.Asset(name + ".br")
	fatalErr(err)
	return decompressAsset(asset)
}

func gzipAsset(data []byte) []byte {
	br, err := brotli.NewReader(bytes.NewReader(data), nil)
	fatalErr(err)

	buf := &bytes.Buffer{}
	gzw, err := gzip.NewWriterLevel(buf, gzip.BestCompression)
	fatalErr(err)

	io.Copy(gzw, br)
	gzw.Close()
	return buf.Bytes()
}

func (d *Dispatch) serveFiles(w http.ResponseWriter, r *http.Request) {
	if r.URL.Path == "/" {
		d.serveIndex(w, r)
	} else if file, ok := files[r.URL.Path]; ok {
		d.serveFile(w, r, file)
	} else if r.URL.Path == "/sw.js" {
		w.Header().Set("Cache-Control", disabledCacheControl)
		w.Header().Set("Content-Type", "text/javascript")
		w.Header().Set("Content-Length", strconv.Itoa(len(serviceWorker)))
		w.Write(serviceWorker)
	} else if r.URL.Path == "/robots.txt" {
		w.Header().Set("Content-Type", "text/plain")
		w.Header().Set("Content-Length", strconv.Itoa(len(robots)))
		w.Write(robots)
	} else {
		d.serveIndex(w, r)
	}
}

func (d *Dispatch) serveIndex(w http.ResponseWriter, r *http.Request) {
	if pusher, ok := w.(http.Pusher); ok {
		options := &http.PushOptions{
			Header: http.Header{
				"Accept-Encoding": r.Header["Accept-Encoding"],
			},
		}

		cookie, err := r.Cookie("push")
		if err != nil {
			for _, asset := range h2PushAssets {
				pusher.Push(asset.path, options)
			}

			setPushCookie(w, r)
		} else {
			pushed := false

			i := 0
			for _, asset := range h2PushAssets {
				if len(cookie.Value) >= i+len(asset.hash) &&
					asset.hash != cookie.Value[i:i+len(asset.hash)] {
					pusher.Push(asset.path, options)
					pushed = true
				}
				i += len(asset.hash)
			}

			if pushed {
				setPushCookie(w, r)
			}
		}
	}

	if cspEnabled {
		var wsSrc string
		if r.TLS != nil {
			wsSrc = "wss://" + r.Host
		} else {
			wsSrc = "ws://" + r.Host
		}

		csp := []string{
			"default-src 'none'",
			"script-src 'self' 'sha256-" + inlineScriptSha256 + "'",
			"style-src 'self' 'unsafe-inline'",
			"font-src 'self'",
			"img-src 'self'",
			"manifest-src 'self'",
			"connect-src 'self' " + wsSrc,
			"worker-src 'self'",
		}

		w.Header().Set("Content-Security-Policy", strings.Join(csp, "; "))
	}

	w.Header().Set("Content-Type", "text/html")
	w.Header().Set("Cache-Control", disabledCacheControl)
	w.Header().Set("X-Content-Type-Options", "nosniff")
	w.Header().Set("X-Frame-Options", "deny")
	w.Header().Set("X-XSS-Protection", "1; mode=block")
	w.Header().Set("Referrer-Policy", "same-origin")

	if hstsHeader != "" {
		w.Header().Set("Strict-Transport-Security", hstsHeader)
	}

	for k, v := range d.Config().Headers {
		w.Header().Set(k, v)
	}

	if strings.Contains(r.Header.Get("Accept-Encoding"), "gzip") {
		w.Header().Set("Content-Encoding", "gzip")
		w.Header().Set("Content-Length", indexPageLen)
		w.Write(indexPage)
	} else {
		serveDecompressed(w, indexPage)
	}
}

func setPushCookie(w http.ResponseWriter, r *http.Request) {
	http.SetCookie(w, &http.Cookie{
		Name:     "push",
		Value:    h2PushCookieValue,
		Path:     "/",
		Expires:  time.Now().AddDate(1, 0, 0),
		HttpOnly: true,
		Secure:   r.TLS != nil,
	})
}

func (d *Dispatch) serveFile(w http.ResponseWriter, r *http.Request, file *File) {
	w.Header().Set("Cache-Control", file.CacheControl)
	w.Header().Set("Content-Type", file.ContentType)

	if file.Compressed && strings.Contains(r.Header.Get("Accept-Encoding"), "br") {
		w.Header().Set("Content-Encoding", "br")
		w.Header().Set("Content-Length", file.Length)
		w.Write(file.Data)
	} else if file.Compressed && strings.Contains(r.Header.Get("Accept-Encoding"), "gzip") {
		w.Header().Set("Content-Encoding", "gzip")
		w.Header().Set("Content-Length", file.GzipLength)
		w.Write(file.GzipData)
	} else if !file.Compressed {
		w.Header().Set("Content-Length", file.Length)
		w.Write(file.Data)
	} else {
		serveDecompressed(w, file.GzipData)
	}
}

func serveDecompressed(w http.ResponseWriter, asset []byte) {
	gzr, err := gzip.NewReader(bytes.NewReader(asset))
	buf, err := ioutil.ReadAll(gzr)
	if err != nil {
		http.Error(w, "", http.StatusInternalServerError)
		return
	}

	w.Header().Set("Content-Length", strconv.Itoa(len(buf)))
	w.Write(buf)
}

func fatalErr(err error) {
	if err != nil {
		log.Fatal(err)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"html/template"
	"net/http"
	"net/http/httptest"
	"testing"
)

func TestRootHandler(t *testing.T) {
	req, err := http.NewRequest("GET", "/", nil)

	if err != nil {
		t.Fatal(err)
	}

	rr := httptest.NewRecorder()
	handler := http.HandlerFunc(rootHandler)

	handler.ServeHTTP(rr, req)

	if status := rr.Code; status != http.StatusOK {
		t.Errorf("handler returned wrong status code: got %v want %v",
			status, http.StatusOK)
	}

	buf := testTemplate(indexTemplate, nil, t)

	if rr.Body.String() != buf.String() {
		t.Errorf("handler returned unexpected body: got %v want %v",
			rr.Body.String(), buf.String())
	}

}

func TestBunReport(t *testing.T) {
	//test GET
	req, err := http.NewRequest("GET", "/bunReport", nil)
	if err != nil {
		t.Fatal(err)
	}

	rr := httptest.NewRecorder()
	handler := http.HandlerFunc(bunReport)

	handler.ServeHTTP(rr, req)

	if status := rr.Code; status != http.StatusOK {
		t.Errorf("handler returned wrong status code: got %v want %v",
			status, http.StatusOK)
	}

	buf := testTemplate(reportTemplate, nil, t)
	if rr.Body.String() != buf.String() {
		t.Errorf("handler returned unexpected body: got %v want %v",
			rr.Body.String(), buf.String())
	}

	//test POST
	req, err = http.NewRequest("POST", "/bunReport?size=1", nil)
	if err != nil {
		t.Fatal(err)
	}

	rr = httptest.NewRecorder()
	handler.ServeHTTP(rr, req)

	if rr.Code != http.StatusOK {
		t.Errorf("handler returned wrong status code: got %v want %v",
			rr.Code, http.StatusOK)
	}

	req, err = http.NewRequest("POST", "/bunReport", nil)
	if err != nil {
		t.Fatal("Unable to handle bunReport with no size")
	}

	rr = httptest.NewRecorder()
	handler.ServeHTTP(rr, req)

	if rr.Code != http.StatusOK {
		t.Errorf("handler returned wrong status code: got %v want %v",
			rr.Code, http.StatusOK)
	}

}

func TestWriteTemplate(t *testing.T) {
	rr := httptest.NewRecorder()
	err := writeTemplate(rr, "", nil)
	if err == nil {
		t.Fatal("Did not receive expected error no such file or directory:", err)
	}

	err = writeTemplate(rr, indexTemplate, nil)
	if err != nil {
		t.Fatal("Error with template file:", err)
	}

	buf := testTemplate(indexTemplate, nil, t)

	if rr.Body.String() != buf.String() {
		t.Errorf("writeTemplate returned unexpected body: got %v expected %v",
			rr.Body.String(), buf.String())
	}
}

func testTemplate(templateFile string, templateData interface{}, t *testing.T) (buf *bytes.Buffer) {
	tpl, err := template.ParseFiles(templateFile)
	if err != nil {
		t.Fatal("Error with template file:", err)
	}

	buf = new(bytes.Buffer)
	err = tpl.Execute(buf, templateData)
	if err != nil {
		t.Fatal("Error executing template:", err)
	}

	return buf
}

func TestMain(t *testing.T) {
	_, err := http.NewRequest("GET", ":-41", nil)
	if err == nil {
		t.Fatal("No error for GET from port -41 ")
	}
	_, err = http.NewRequest("GET", "/fakeURL", nil)
	if err != nil {
		t.Fatal("Error for GET from fakeURL", err)
	}
}
<commit_msg>Rewrote bunReporter handler tests using table<commit_after>package main

import (
	"bytes"
	"html/template"
	"io"
	"net/http"
	"net/http/httptest"
	"testing"
)

func TestRootHandler(t *testing.T) {
	req, err := http.NewRequest("GET", "/", nil)

	if err != nil {
		t.Fatal(err)
	}

	rr := httptest.NewRecorder()
	handler := http.HandlerFunc(rootHandler)

	handler.ServeHTTP(rr, req)

	if rr.Code != http.StatusOK {
		t.Errorf("handler returned wrong status code: got %v want %v",
			rr.Code, http.StatusOK)
	}

	buf := testTemplate(indexTemplate, nil, t)

	if rr.Body.String() != buf.String() {
		t.Errorf("handler returned unexpected body: got %v want %v",
			rr.Body.String(), buf.String())
	}

}

func TestBunReport(t *testing.T) {
	var handlerTests = []struct {
		method      string
		urlString   string
		requestBody io.Reader
		body        string
		status      int
	}{
		{"GET", "/bunReport", nil, testTemplate(reportTemplate, nil, t).String(), http.StatusOK},
		{"POST", "/bunReport", nil, "", http.StatusOK},
		{"POST", "/bunReport?size=1", nil, "", http.StatusOK},
	}

	for _, tt := range handlerTests {
		req, err := http.NewRequest(tt.method, tt.urlString, tt.requestBody)
		if err != nil {
			t.Fatal(err)
		}

		rr := httptest.NewRecorder()
		handler := http.HandlerFunc(bunReport)
		handler.ServeHTTP(rr, req)

		if rr.Code != tt.status {
			t.Errorf("handler returned wrong status code: got %v want %v",
				rr.Code, tt.status)
		}

		if tt.body != "" && rr.Body.String() != tt.body {
			t.Errorf("handler returned unexpected body: got %v want %v",
				rr.Body.String(), tt.body)
		}
	}
}

func TestWriteTemplate(t *testing.T) {
	rr := httptest.NewRecorder()
	err := writeTemplate(rr, "", nil)
	if err == nil {
		t.Fatal("Did not receive expected error no such file or directory:", err)
	}

	err = writeTemplate(rr, indexTemplate, nil)
	if err != nil {
		t.Fatal("Error with template file:", err)
	}

	buf := testTemplate(indexTemplate, nil, t)

	if rr.Body.String() != buf.String() {
		t.Errorf("writeTemplate returned unexpected body: got %v expected %v",
			rr.Body.String(), buf.String())
	}
}

func testTemplate(templateFile string, templateData interface{}, t *testing.T) (buf *bytes.Buffer) {
	tpl, err := template.ParseFiles(templateFile)
	if err != nil {
		t.Fatal("Error with template file:", err)
	}

	buf = new(bytes.Buffer)
	err = tpl.Execute(buf, templateData)
	if err != nil {
		t.Fatal("Error executing template:", err)
	}

	return buf
}

func TestMain(t *testing.T) {
	_, err := http.NewRequest("GET", ":-41", nil)
	if err == nil {
		t.Fatal("No error for GET from port -41 ")
	}
	_, err = http.NewRequest("GET", "/fakeURL", nil)
	if err != nil {
		t.Fatal("Error for GET from fakeURL", err)
	}
}
<|endoftext|>


<commit_before>// Copyright 2013 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"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"runtime"
	"strings"

	"code.google.com/p/go.tools/go/vcs"
)

// These variables are copied from the gobuilder's environment
// to the envv of its subprocesses.
var extraEnv = []string{
	"GOARM",
	"CGO_ENABLED",

	// For Unix derivatives.
	"CC",
	"PATH",
	"TMPDIR",
	"USER",

	// For Plan 9.
	"objtype",
	"cputype",
	"path",
}

// builderEnv represents the environment that a Builder will run tests in.
type builderEnv interface {
	// setup sets up the builder environment and returns the directory to run the buildCmd in.
	setup(repo *Repo, workpath, hash string, envv []string) (string, error)
}

// goEnv represents the builderEnv for the main Go repo.
type goEnv struct {
	goos, goarch string
}

func (b *Builder) envv() []string {
	if runtime.GOOS == "windows" {
		return b.envvWindows()
	}

	var e []string
	if *buildTool == "go" {
		e = []string{
			"GOOS=" + b.goos,
			"GOHOSTOS=" + b.goos,
			"GOARCH=" + b.goarch,
			"GOHOSTARCH=" + b.goarch,
			"GOROOT_FINAL=/usr/local/go",
		}
	}

	for _, k := range extraEnv {
		if s, ok := getenvOk(k); ok {
			e = append(e, k+"="+s)
		}
	}
	return e
}

func (b *Builder) envvWindows() []string {
	var start map[string]string
	if *buildTool == "go" {
		start = map[string]string{
			"GOOS":         b.goos,
			"GOHOSTOS":     b.goos,
			"GOARCH":       b.goarch,
			"GOHOSTARCH":   b.goarch,
			"GOROOT_FINAL": `c:\go`,
			"GOBUILDEXIT":  "1", // exit all.bat with completion status.
		}
	}

	for _, name := range extraEnv {
		if s, ok := getenvOk(name); ok {
			start[name] = s
		}
	}
	skip := map[string]bool{
		"GOBIN":   true,
		"GOPATH":  true,
		"GOROOT":  true,
		"INCLUDE": true,
		"LIB":     true,
	}
	var e []string
	for name, v := range start {
		e = append(e, name+"="+v)
		skip[name] = true
	}
	for _, kv := range os.Environ() {
		s := strings.SplitN(kv, "=", 2)
		name := strings.ToUpper(s[0])
		switch {
		case name == "":
			// variables, like "=C:=C:\", just copy them
			e = append(e, kv)
		case !skip[name]:
			e = append(e, kv)
			skip[name] = true
		}
	}
	return e
}

// setup for a goEnv clones the main go repo to workpath/go at the provided hash
// and returns the path workpath/go/src, the location of all go build scripts.
func (env *goEnv) setup(repo *Repo, workpath, hash string, envv []string) (string, error) {
	goworkpath := filepath.Join(workpath, "go")
	if err := repo.Export(goworkpath, hash); err != nil {
		return "", fmt.Errorf("error exporting repository: %s", err)
	}
	if err := ioutil.WriteFile(filepath.Join(goworkpath, "VERSION"), []byte(hash), 0644); err != nil {
		return "", fmt.Errorf("error writing VERSION file: %s", err)
	}
	return filepath.Join(goworkpath, "src"), nil
}

// gccgoEnv represents the builderEnv for the gccgo compiler.
type gccgoEnv struct{}

// setup for a gccgoEnv clones the gofrontend repo to workpath/go at the hash
// and clones the latest GCC branch to repo.Path/gcc. The gccgo sources are
// replaced with the updated sources in the gofrontend repo and gcc gets
// gets configured and built in workpath/gcc-objdir. The path to
// workpath/gcc-objdir is returned.
func (env *gccgoEnv) setup(repo *Repo, workpath, hash string, envv []string) (string, error) {
	gccpath := filepath.Join(repo.Path, "gcc")

	// get a handle to Git vcs.Cmd for pulling down GCC from the mirror.
	git := vcs.ByCmd("git")

	// only pull down gcc if we don't have a local copy.
	if _, err := os.Stat(gccpath); err != nil {
		if err := timeout(*cmdTimeout, func() error {
			// pull down a working copy of GCC.
			return git.Create(gccpath, *gccPath)
		}); err != nil {
			return "", err
		}
	}

	if err := git.Download(gccpath); err != nil {
		return "", err
	}

	// get the modified files for this commit.
	statusCmd := []string{
		"hg",
		"status",
		"--no-status",
		"--change",
		hash,
	}

	var buf bytes.Buffer
	if _, err := runOutput(*cmdTimeout, envv, &buf, repo.Path, statusCmd...); err != nil {
		return "", fmt.Errorf("Failed to find the modified files for %s: %s", hash, err)
	}
	modifiedFiles := strings.Split(buf.String(), "\n")
	var isMirrored bool
	for _, f := range modifiedFiles {
		if strings.HasPrefix(f, "go/") || strings.HasPrefix(f, "libgo/") {
			isMirrored = true
			break
		}
	}

	// use git log to find the corresponding commit to sync to in the gcc mirror.
	// If the files modified in the gofrontend are mirrored to gcc, we expect a
	// commit with a similar description in the gcc mirror. If the files modified are
	// not mirrored, e.g. in support/, we can sync to the most recent gcc commit that
	// occurred before those files were modified to verify gccgo's status at that point.
	logCmd := []string{
		"git",
		"log",
		"-1",
		"--format=%H",
	}
	var errMsg string
	if isMirrored {
		commitDesc, err := repo.Master.VCS.LogAtRev(repo.Path, hash, "{desc|escape}")
		if err != nil {
			return "", err
		}

		logCmd = append(logCmd, "--grep", "'"+string(commitDesc)+"'", "--regexp-ignore-case")
		errMsg = fmt.Sprintf("Failed to find a commit with a similar description to '%s'", string(commitDesc))
	} else {
		commitDate, err := repo.Master.VCS.LogAtRev(repo.Path, hash, "{date|rfc3339date}")
		if err != nil {
			return "", err
		}

		logCmd = append(logCmd, "--before", string(commitDate))
		errMsg = fmt.Sprintf("Failed to find a commit before '%s'", string(commitDate))
	}

	buf.Reset()
	if _, err := runOutput(*cmdTimeout, envv, &buf, gccpath, logCmd...); err != nil {
		return "", fmt.Errorf("%s: %s", errMsg, err)
	}
	gccRev := buf.String()
	if gccRev == "" {
		return "", fmt.Errorf(errMsg)
	}

	// checkout gccRev
	// TODO(cmang): Fix this to work in parallel mode.
	checkoutCmd := []string{
		"git",
		"reset",
		"--hard",
		strings.TrimSpace(gccRev),
	}
	if _, err := runOutput(*cmdTimeout, envv, ioutil.Discard, gccpath, checkoutCmd...); err != nil {
		return "", fmt.Errorf("Failed to checkout commit at revision %s: %s", gccRev, err)
	}

	// make objdir to work in
	gccobjdir := filepath.Join(workpath, "gcc-objdir")
	if err := os.Mkdir(gccobjdir, mkdirPerm); err != nil {
		return "", err
	}

	// configure GCC with substituted gofrontend and libgo
	gccConfigCmd := []string{
		filepath.Join(gccpath, "configure"),
		"--enable-languages=c,c++,go",
		"--disable-bootstrap",
		"--disable-multilib",
	}
	if _, err := runOutput(*cmdTimeout, envv, ioutil.Discard, gccobjdir, gccConfigCmd...); err != nil {
		return "", fmt.Errorf("Failed to configure GCC: %s", err)
	}

	// build gcc
	if _, err := runOutput(*buildTimeout, envv, ioutil.Discard, gccobjdir, "make"); err != nil {
		return "", fmt.Errorf("Failed to build GCC: %s", err)
	}

	return gccobjdir, nil
}

func getenvOk(k string) (v string, ok bool) {
	v = os.Getenv(k)
	if v != "" {
		return v, true
	}
	keq := k + "="
	for _, kv := range os.Environ() {
		if kv == keq {
			return "", true
		}
	}
	return "", false
}
<commit_msg>go.tools/dashboard/builder: fix imported by not used error<commit_after>// Copyright 2013 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"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"runtime"
	"strings"

	"code.google.com/p/go.tools/go/vcs"
)

// These variables are copied from the gobuilder's environment
// to the envv of its subprocesses.
var extraEnv = []string{
	"GOARM",
	"CGO_ENABLED",

	// For Unix derivatives.
	"CC",
	"PATH",
	"TMPDIR",
	"USER",

	// For Plan 9.
	"objtype",
	"cputype",
	"path",
}

// builderEnv represents the environment that a Builder will run tests in.
type builderEnv interface {
	// setup sets up the builder environment and returns the directory to run the buildCmd in.
	setup(repo *Repo, workpath, hash string, envv []string) (string, error)
}

// goEnv represents the builderEnv for the main Go repo.
type goEnv struct {
	goos, goarch string
}

func (b *Builder) envv() []string {
	if runtime.GOOS == "windows" {
		return b.envvWindows()
	}

	var e []string
	if *buildTool == "go" {
		e = []string{
			"GOOS=" + b.goos,
			"GOHOSTOS=" + b.goos,
			"GOARCH=" + b.goarch,
			"GOHOSTARCH=" + b.goarch,
			"GOROOT_FINAL=/usr/local/go",
		}
	}

	for _, k := range extraEnv {
		if s, ok := getenvOk(k); ok {
			e = append(e, k+"="+s)
		}
	}
	return e
}

func (b *Builder) envvWindows() []string {
	var start map[string]string
	if *buildTool == "go" {
		start = map[string]string{
			"GOOS":         b.goos,
			"GOHOSTOS":     b.goos,
			"GOARCH":       b.goarch,
			"GOHOSTARCH":   b.goarch,
			"GOROOT_FINAL": `c:\go`,
			"GOBUILDEXIT":  "1", // exit all.bat with completion status.
		}
	}

	for _, name := range extraEnv {
		if s, ok := getenvOk(name); ok {
			start[name] = s
		}
	}
	skip := map[string]bool{
		"GOBIN":   true,
		"GOPATH":  true,
		"GOROOT":  true,
		"INCLUDE": true,
		"LIB":     true,
	}
	var e []string
	for name, v := range start {
		e = append(e, name+"="+v)
		skip[name] = true
	}
	for _, kv := range os.Environ() {
		s := strings.SplitN(kv, "=", 2)
		name := strings.ToUpper(s[0])
		switch {
		case name == "":
			// variables, like "=C:=C:\", just copy them
			e = append(e, kv)
		case !skip[name]:
			e = append(e, kv)
			skip[name] = true
		}
	}
	return e
}

// setup for a goEnv clones the main go repo to workpath/go at the provided hash
// and returns the path workpath/go/src, the location of all go build scripts.
func (env *goEnv) setup(repo *Repo, workpath, hash string, envv []string) (string, error) {
	goworkpath := filepath.Join(workpath, "go")
	if err := repo.Export(goworkpath, hash); err != nil {
		return "", fmt.Errorf("error exporting repository: %s", err)
	}
	if err := ioutil.WriteFile(filepath.Join(goworkpath, "VERSION"), []byte(hash), 0644); err != nil {
		return "", fmt.Errorf("error writing VERSION file: %s", err)
	}
	return filepath.Join(goworkpath, "src"), nil
}

// gccgoEnv represents the builderEnv for the gccgo compiler.
type gccgoEnv struct{}

// setup for a gccgoEnv clones the gofrontend repo to workpath/go at the hash
// and clones the latest GCC branch to repo.Path/gcc. The gccgo sources are
// replaced with the updated sources in the gofrontend repo and gcc gets
// gets configured and built in workpath/gcc-objdir. The path to
// workpath/gcc-objdir is returned.
func (env *gccgoEnv) setup(repo *Repo, workpath, hash string, envv []string) (string, error) {
	gccpath := filepath.Join(repo.Path, "gcc")

	// get a handle to Git vcs.Cmd for pulling down GCC from the mirror.
	git := vcs.ByCmd("git")

	// only pull down gcc if we don't have a local copy.
	if _, err := os.Stat(gccpath); err != nil {
		if err := timeout(*cmdTimeout, func() error {
			// pull down a working copy of GCC.
			return git.Create(gccpath, *gccPath)
		}); err != nil {
			return "", err
		}
	}

	if err := git.Download(gccpath); err != nil {
		return "", err
	}

	// get the modified files for this commit.
	statusCmd := []string{
		"hg",
		"status",
		"--no-status",
		"--change",
		hash,
	}

	var buf bytes.Buffer
	if _, err := runOutput(*cmdTimeout, envv, &buf, repo.Path, statusCmd...); err != nil {
		return "", fmt.Errorf("Failed to find the modified files for %s: %s", hash, err)
	}
	modifiedFiles := strings.Split(buf.String(), "\n")
	var isMirrored bool
	for _, f := range modifiedFiles {
		if strings.HasPrefix(f, "go/") || strings.HasPrefix(f, "libgo/") {
			isMirrored = true
			break
		}
	}

	// use git log to find the corresponding commit to sync to in the gcc mirror.
	// If the files modified in the gofrontend are mirrored to gcc, we expect a
	// commit with a similar description in the gcc mirror. If the files modified are
	// not mirrored, e.g. in support/, we can sync to the most recent gcc commit that
	// occurred before those files were modified to verify gccgo's status at that point.
	logCmd := []string{
		"git",
		"log",
		"-1",
		"--format=%H",
	}
	var errMsg string
	if isMirrored {
		commitDesc, err := repo.Master.VCS.LogAtRev(repo.Path, hash, "{desc|escape}")
		if err != nil {
			return "", err
		}

		logCmd = append(logCmd, "--grep", "'"+string(commitDesc)+"'", "--regexp-ignore-case")
		errMsg = fmt.Sprintf("Failed to find a commit with a similar description to '%s'", string(commitDesc))
	} else {
		commitDate, err := repo.Master.VCS.LogAtRev(repo.Path, hash, "{date|rfc3339date}")
		if err != nil {
			return "", err
		}

		logCmd = append(logCmd, "--before", string(commitDate))
		errMsg = fmt.Sprintf("Failed to find a commit before '%s'", string(commitDate))
	}

	buf.Reset()
	if _, err := runOutput(*cmdTimeout, envv, &buf, gccpath, logCmd...); err != nil {
		return "", fmt.Errorf("%s: %s", errMsg, err)
	}
	gccRev := buf.String()
	if gccRev == "" {
		return "", fmt.Errorf(errMsg)
	}

	// checkout gccRev
	// TODO(cmang): Fix this to work in parallel mode.
	checkoutCmd := []string{
		"git",
		"reset",
		"--hard",
		strings.TrimSpace(gccRev),
	}
	if _, err := runOutput(*cmdTimeout, envv, ioutil.Discard, gccpath, checkoutCmd...); err != nil {
		return "", fmt.Errorf("Failed to checkout commit at revision %s: %s", gccRev, err)
	}

	// make objdir to work in
	gccobjdir := filepath.Join(workpath, "gcc-objdir")
	if err := os.Mkdir(gccobjdir, mkdirPerm); err != nil {
		return "", err
	}

	// configure GCC with substituted gofrontend and libgo
	gccConfigCmd := []string{
		filepath.Join(gccpath, "configure"),
		"--enable-languages=c,c++,go",
		"--disable-bootstrap",
		"--disable-multilib",
	}
	if _, err := runOutput(*cmdTimeout, envv, ioutil.Discard, gccobjdir, gccConfigCmd...); err != nil {
		return "", fmt.Errorf("Failed to configure GCC: %s", err)
	}

	// build gcc
	if _, err := runOutput(*buildTimeout, envv, ioutil.Discard, gccobjdir, "make"); err != nil {
		return "", fmt.Errorf("Failed to build GCC: %s", err)
	}

	return gccobjdir, nil
}

func getenvOk(k string) (v string, ok bool) {
	v = os.Getenv(k)
	if v != "" {
		return v, true
	}
	keq := k + "="
	for _, kv := range os.Environ() {
		if kv == keq {
			return "", true
		}
	}
	return "", false
}
<|endoftext|>


<commit_before>// Copyright © 2016 The Things Network
// Use of this source code is governed by the MIT license that can be found in the LICENSE file.

package util

import (
	"encoding/json"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"testing"

	. "github.com/smartystreets/assertions"
)

var (
	url   = "/foo"
	token = "token"
)

type OKResp struct {
	OK string `json:token`
}

type FooResp struct {
	Foo string `json:"foo" valid:"required"`
}

func OKHandler(a *Assertion, method string) http.HandlerFunc {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		a.So(r.RequestURI, ShouldEqual, url)
		a.So(r.Method, ShouldEqual, method)
		a.So(r.Header.Get("Authorization"), ShouldEqual, "bearer "+token)
		resp := OKResp{
			OK: token,
		}
		w.WriteHeader(http.StatusOK)
		encoder := json.NewEncoder(w)
		encoder.Encode(&resp)
	})
}

func FooHandler(a *Assertion, method string) http.HandlerFunc {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		a.So(r.RequestURI, ShouldEqual, url)
		a.So(r.Method, ShouldEqual, method)
		a.So(r.Header.Get("Authorization"), ShouldEqual, "bearer "+token)
		resp := FooResp{
			Foo: token,
		}
		w.WriteHeader(http.StatusOK)
		encoder := json.NewEncoder(w)
		encoder.Encode(&resp)
	})
}

func RedirectHandler(a *Assertion, method string) http.HandlerFunc {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		a.So(r.Header.Get("Authorization"), ShouldEqual, "bearer "+token)
		if r.RequestURI == url {
			w.Header().Set("Location", "/bar")
			w.WriteHeader(307)
		} else {
			a.So(r.RequestURI, ShouldEqual, "/bar")
			resp := FooResp{
				Foo: token,
			}
			w.WriteHeader(http.StatusOK)
			encoder := json.NewEncoder(w)
			encoder.Encode(&resp)
		}
	})
}

func EchoHandler(a *Assertion, method string) http.HandlerFunc {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		a.So(r.RequestURI, ShouldEqual, url)
		a.So(r.Method, ShouldEqual, method)
		a.So(r.Header.Get("Authorization"), ShouldEqual, "bearer "+token)
		w.WriteHeader(http.StatusOK)
		defer r.Body.Close()
		body, err := ioutil.ReadAll(r.Body)
		a.So(err, ShouldBeNil)
		w.Write(body)
	})
}

func TestGET(t *testing.T) {
	a := New(t)
	server := httptest.NewServer(OKHandler(a, "GET"))
	defer server.Close()

	var resp OKResp
	err := GET(server.URL, token, url, &resp)
	a.So(err, ShouldBeNil)
	a.So(resp.OK, ShouldEqual, token)
}

func TestGETDropResponse(t *testing.T) {
	a := New(t)
	server := httptest.NewServer(OKHandler(a, "GET"))
	defer server.Close()

	err := GET(server.URL, token, url, nil)
	a.So(err, ShouldBeNil)
}

func TestGETIllegalResponse(t *testing.T) {
	a := New(t)
	server := httptest.NewServer(OKHandler(a, "GET"))
	defer server.Close()

	var resp FooResp
	err := GET(server.URL, token, url, &resp)
	a.So(err, ShouldNotBeNil)
}

func TestGETIllegalResponseIgnore(t *testing.T) {
	a := New(t)
	server := httptest.NewServer(FooHandler(a, "GET"))
	defer server.Close()

	var resp OKResp
	err := GET(server.URL, token, url, &resp)
	a.So(err, ShouldBeNil)
}

func TestGETRedirect(t *testing.T) {
	a := New(t)
	server := httptest.NewServer(RedirectHandler(a, "GET"))
	defer server.Close()

	var resp OKResp
	err := GET(server.URL, token, url, &resp)
	a.So(err, ShouldBeNil)
}

func TestPUT(t *testing.T) {
	a := New(t)
	server := httptest.NewServer(EchoHandler(a, "PUT"))
	defer server.Close()

	var resp FooResp
	body := FooResp{
		Foo: token,
	}
	err := PUT(server.URL, token, url, body, &resp)
	a.So(err, ShouldBeNil)
	a.So(resp.Foo, ShouldEqual, body.Foo)
}

func TestPUTIllegalRequest(t *testing.T) {
	a := New(t)
	server := httptest.NewServer(EchoHandler(a, "PUT"))
	defer server.Close()

	var resp FooResp
	body := FooResp{}
	err := PUT(server.URL, token, url, body, &resp)
	a.So(err, ShouldNotBeNil)
}

func TestPUTIllegalResponse(t *testing.T) {
	a := New(t)
	server := httptest.NewServer(OKHandler(a, "PUT"))
	defer server.Close()

	var resp FooResp
	err := PUT(server.URL, token, url, nil, &resp)
	a.So(err, ShouldNotBeNil)
}

func TestPUTRedirect(t *testing.T) {
	a := New(t)
	server := httptest.NewServer(RedirectHandler(a, "PUT"))
	defer server.Close()

	var resp FooResp
	err := PUT(server.URL, token, url, nil, &resp)
	a.So(err, ShouldBeNil)
	a.So(resp.Foo, ShouldEqual, token)
}
<commit_msg>add tests for POST and DELETE<commit_after>// Copyright © 2016 The Things Network
// Use of this source code is governed by the MIT license that can be found in the LICENSE file.

package util

import (
	"encoding/json"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"testing"

	. "github.com/smartystreets/assertions"
)

var (
	url   = "/foo"
	token = "token"
)

type OKResp struct {
	OK string `json:token`
}

type FooResp struct {
	Foo string `json:"foo" valid:"required"`
}

func OKHandler(a *Assertion, method string) http.HandlerFunc {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		a.So(r.RequestURI, ShouldEqual, url)
		a.So(r.Method, ShouldEqual, method)
		a.So(r.Header.Get("Authorization"), ShouldEqual, "bearer "+token)
		resp := OKResp{
			OK: token,
		}
		w.WriteHeader(http.StatusOK)
		encoder := json.NewEncoder(w)
		encoder.Encode(&resp)
	})
}

func FooHandler(a *Assertion, method string) http.HandlerFunc {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		a.So(r.RequestURI, ShouldEqual, url)
		a.So(r.Method, ShouldEqual, method)
		a.So(r.Header.Get("Authorization"), ShouldEqual, "bearer "+token)
		resp := FooResp{
			Foo: token,
		}
		w.WriteHeader(http.StatusOK)
		encoder := json.NewEncoder(w)
		encoder.Encode(&resp)
	})
}

func RedirectHandler(a *Assertion, method string) http.HandlerFunc {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		a.So(r.Header.Get("Authorization"), ShouldEqual, "bearer "+token)
		if r.RequestURI == url {
			w.Header().Set("Location", "/bar")
			w.WriteHeader(307)
		} else {
			a.So(r.RequestURI, ShouldEqual, "/bar")
			resp := FooResp{
				Foo: token,
			}
			w.WriteHeader(http.StatusOK)
			encoder := json.NewEncoder(w)
			encoder.Encode(&resp)
		}
	})
}

func EchoHandler(a *Assertion, method string) http.HandlerFunc {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		a.So(r.RequestURI, ShouldEqual, url)
		a.So(r.Method, ShouldEqual, method)
		a.So(r.Header.Get("Authorization"), ShouldEqual, "bearer "+token)
		w.WriteHeader(http.StatusOK)
		defer r.Body.Close()
		body, err := ioutil.ReadAll(r.Body)
		a.So(err, ShouldBeNil)
		w.Write(body)
	})
}

func TestGET(t *testing.T) {
	a := New(t)
	server := httptest.NewServer(OKHandler(a, "GET"))
	defer server.Close()

	var resp OKResp
	err := GET(server.URL, token, url, &resp)
	a.So(err, ShouldBeNil)
	a.So(resp.OK, ShouldEqual, token)
}

func TestGETDropResponse(t *testing.T) {
	a := New(t)
	server := httptest.NewServer(OKHandler(a, "GET"))
	defer server.Close()

	err := GET(server.URL, token, url, nil)
	a.So(err, ShouldBeNil)
}

func TestGETIllegalResponse(t *testing.T) {
	a := New(t)
	server := httptest.NewServer(OKHandler(a, "GET"))
	defer server.Close()

	var resp FooResp
	err := GET(server.URL, token, url, &resp)
	a.So(err, ShouldNotBeNil)
}

func TestGETIllegalResponseIgnore(t *testing.T) {
	a := New(t)
	server := httptest.NewServer(FooHandler(a, "GET"))
	defer server.Close()

	var resp OKResp
	err := GET(server.URL, token, url, &resp)
	a.So(err, ShouldBeNil)
}

func TestGETRedirect(t *testing.T) {
	a := New(t)
	server := httptest.NewServer(RedirectHandler(a, "GET"))
	defer server.Close()

	var resp OKResp
	err := GET(server.URL, token, url, &resp)
	a.So(err, ShouldBeNil)
}

func TestPUT(t *testing.T) {
	a := New(t)
	server := httptest.NewServer(EchoHandler(a, "PUT"))
	defer server.Close()

	var resp FooResp
	body := FooResp{
		Foo: token,
	}
	err := PUT(server.URL, token, url, body, &resp)
	a.So(err, ShouldBeNil)
	a.So(resp.Foo, ShouldEqual, body.Foo)
}

func TestPUTIllegalRequest(t *testing.T) {
	a := New(t)
	server := httptest.NewServer(EchoHandler(a, "PUT"))
	defer server.Close()

	var resp FooResp
	body := FooResp{}
	err := PUT(server.URL, token, url, body, &resp)
	a.So(err, ShouldNotBeNil)
}

func TestPUTIllegalResponse(t *testing.T) {
	a := New(t)
	server := httptest.NewServer(OKHandler(a, "PUT"))
	defer server.Close()

	var resp FooResp
	err := PUT(server.URL, token, url, nil, &resp)
	a.So(err, ShouldNotBeNil)
}

func TestPUTRedirect(t *testing.T) {
	a := New(t)
	server := httptest.NewServer(RedirectHandler(a, "PUT"))
	defer server.Close()

	var resp FooResp
	err := PUT(server.URL, token, url, nil, &resp)
	a.So(err, ShouldBeNil)
	a.So(resp.Foo, ShouldEqual, token)
}

func TestPOST(t *testing.T) {
	a := New(t)
	server := httptest.NewServer(EchoHandler(a, "POST"))
	defer server.Close()

	var resp FooResp
	body := FooResp{
		Foo: token,
	}
	err := POST(server.URL, token, url, body, &resp)
	a.So(err, ShouldBeNil)
	a.So(resp.Foo, ShouldEqual, body.Foo)
}

func TestPOSTIllegalRequest(t *testing.T) {
	a := New(t)
	server := httptest.NewServer(EchoHandler(a, "POST"))
	defer server.Close()

	var resp FooResp
	body := FooResp{}
	err := POST(server.URL, token, url, body, &resp)
	a.So(err, ShouldNotBeNil)
}

func TestPOSTIllegalResponse(t *testing.T) {
	a := New(t)
	server := httptest.NewServer(OKHandler(a, "POST"))
	defer server.Close()

	var resp FooResp
	err := POST(server.URL, token, url, nil, &resp)
	a.So(err, ShouldNotBeNil)
}

func TestPOSTRedirect(t *testing.T) {
	a := New(t)
	server := httptest.NewServer(RedirectHandler(a, "POST"))
	defer server.Close()

	var resp FooResp
	err := POST(server.URL, token, url, nil, &resp)
	a.So(err, ShouldBeNil)
	a.So(resp.Foo, ShouldEqual, token)
}

func TestDELETE(t *testing.T) {
	a := New(t)
	server := httptest.NewServer(OKHandler(a, "DELETE"))
	defer server.Close()

	err := DELETE(server.URL, token, url)
	a.So(err, ShouldBeNil)
}

func TestDELETERedirect(t *testing.T) {
	a := New(t)
	server := httptest.NewServer(RedirectHandler(a, "DELETE"))
	defer server.Close()

	err := DELETE(server.URL, token, url)
	a.So(err, ShouldBeNil)
}
<|endoftext|>


<commit_before>package cli

import (
	"fmt"
	"github.com/brightbox/gobrightbox"
	"gopkg.in/alecthomas/kingpin.v2"
	"sort"
)

type ImagesCommand struct {
	*CliApp
	Id      string
	IdList  []string
	ShowAll bool
}

type imagesForDisplay []brightbox.Image

func (il imagesForDisplay) Len() int      { return len(il) }
func (il imagesForDisplay) Swap(i, j int) { il[i], il[j] = il[j], il[i] }
func (il imagesForDisplay) Less(i, j int) bool {
	a := il.sortKeys(i)
	b := il.sortKeys(j)
	for e := 0; e < len(a); e++ {
		switch a[e].(type) {
		case string:
			if a[e].(string) != b[e].(string) {
				return a[e].(string) < b[e].(string)
			}
		case int64:
			if a[e].(int64) != b[e].(int64) {
				return a[e].(int64) < b[e].(int64)
			}
		case bool:
			if a[e].(bool) != b[e].(bool) {
				return a[e].(bool) == false && b[e].(bool) == true
			}
		}
	}
	return false
}

func (il imagesForDisplay) sortKeys(i int) []interface{} {
	im := il[i]
	return []interface{}{
		!im.Official,
		im.SourceType == "snapshot",
		im.Name, im.Arch,
		im.Status == "deprecated",
		!im.Public,
		-im.CreatedAt.Unix(),
	}
}

func (l *ImagesCommand) list(pc *kingpin.ParseContext) error {
	err := l.Configure()
	if err != nil {
		return err
	}

	// Get the account id in parallel to everything else, we don't need it until
	// later anyway.
	achan := make(chan string)
	go func() {
		achan <- l.accountId()
	}()

	w := tabWriter()
	defer w.Flush()

	images, err := l.Client.Images()
	if err != nil {
		return err
	}
	sortedImages := make(imagesForDisplay, 0, len(*images))
	for _, i := range *images {
		sortedImages = append(sortedImages, i)
	}
	sort.Sort(sortedImages)

	accountId := <-achan

	listRec(w, "ID", "OWNER", "TYPE", "CREATED", "STATUS", "SIZE", "NAME")
	for _, i := range sortedImages {
		owner := i.Owner
		itype := i.SourceType
		if i.Official {
			owner = "brightbox"
			itype = "official"
		}
		name := i.Name
		if i.CompatibilityMode {
			name += " (compat)"
		}
		name += " (" + i.Arch + ")"
		var status string
		if i.Status != "available" {
			status = i.Status
		} else if i.Public {
			status = "public"
		} else {
			status = "private"
		}
		if l.ShowAll == false {
			if !i.Official && i.Owner != accountId {
				continue
			}
		}
		listRec(
			w, i.Id, owner, itype,
			i.CreatedAt.Format("2006-01-02"),
			status, i.VirtualSize, name)
	}
	return nil
}

func (l *ImagesCommand) show(pc *kingpin.ParseContext) error {

	err := l.Configure()
	if err != nil {
		return err
	}
	w := tabWriterRight()
	defer w.Flush()
	s, err := l.Client.Image(l.Id)
	if err != nil {
		l.Fatalf(err.Error())
	}

	owner := s.Owner
	itype := s.SourceType
	if s.Official {
		owner = "brightbox"
		itype = "official"
	}
	drawShow(w, []interface{}{
		"id", s.Id,
		"type", itype,
		"owner", owner,
		"created_at", s.CreatedAt,
		"status", s.Status,
		"locked", s.Locked,
		"arch", s.Arch,
		"name", s.Name,
		"description", s.Description,
		"username", s.Username,
		"virtual_size", s.VirtualSize,
		"disk_size", s.DiskSize,
		"public", s.Public,
		"compatibility_mode", s.CompatibilityMode,
		"official", s.Official,
		"ancestor_id", s.AncestorId,
		"license_name", s.LicenseName,
	})
	return nil

}

func (l *ImagesCommand) destroy(pc *kingpin.ParseContext) error {
	err := l.Configure()
	if err != nil {
		return err
	}
	returnError := false
	for _, id := range l.IdList {
		fmt.Printf("Destroying image %s\n", id)
		err := l.Client.DestroyImage(id)
		if err != nil {
			l.Errorf("%s: %s", err.Error(), id)
			returnError = true
		}
	}
	if returnError {
		return genericError
	}
	return nil
}

func ConfigureImagesCommand(app *CliApp) {
	cmd := ImagesCommand{CliApp: app}
	images := app.Command("images", "Manage server images")
	list := images.Command("list", "List server images").Default().Action(cmd.list)
	list.Flag("show-all", "Show all public images from all accounts").Default("false").BoolVar(&cmd.ShowAll)
	show := images.Command("show", "View details of a server image").Action(cmd.show)
	show.Arg("identifier", "Identifier of image to show").Required().StringVar(&cmd.Id)
	destroy := images.Command("destroy", "Destroy a server image").Action(cmd.destroy)
	destroy.Arg("identifier", "Identifier of image to destroy").Required().StringsVar(&cmd.IdList)

}
<commit_msg>Add field support to images commands<commit_after>package cli

import (
	"fmt"
	"github.com/brightbox/gobrightbox"
	"gopkg.in/alecthomas/kingpin.v2"
	"sort"
	"strings"
)

const (
	DefaultImageListFields = "id,owner,type,created,status,arch,name"
	DefaultImageShowFields = "id,type,owner,created_at,status,locked,arch,name,description,username,virtual_size,disk_size,public,compatibility_mode,official,ancestor_id,license_name"
)

type ImagesCommand struct {
	*CliApp
	Id      string
	IdList  []string
	ShowAll bool
	Fields  string
}

func ImageFields(i *brightbox.Image) map[string]string {
	owner := i.Owner
	itype := i.SourceType
	if i.Official {
		owner = "brightbox"
		itype = "official"
	}
	name := i.Name
	if i.CompatibilityMode {
		name += " (compat)"
	}
	name += " (" + i.Arch + ")"
	var status string
	if i.Status != "available" {
		status = i.Status
	} else if i.Public {
		status = "public"
	} else {
		status = "private"
	}

	return map[string]string{
		"id":                 i.Id,
		"type":               itype,
		"owner":              owner,
		"created":            i.CreatedAt.Format("2006-01-02"),
		"created_at":         formatTime(&i.CreatedAt),
		"status":             status,
		"locked":             formatBool(i.Locked),
		"arch":               i.Arch,
		"name":               i.Name,
		"description":        i.Description,
		"username":           i.Username,
		"virtual_size":       formatInt(i.VirtualSize),
		"disk_size":          formatInt(i.DiskSize),
		"public":             formatBool(i.Public),
		"compatibility_mode": formatBool(i.CompatibilityMode),
		"official":           formatBool(i.Official),
		"ancestor_id":        i.AncestorId,
		"license_name":       i.LicenseName,
	}
}

type imagesForDisplay []brightbox.Image

func (il imagesForDisplay) Len() int      { return len(il) }
func (il imagesForDisplay) Swap(i, j int) { il[i], il[j] = il[j], il[i] }
func (il imagesForDisplay) Less(i, j int) bool {
	a := il.sortKeys(i)
	b := il.sortKeys(j)
	for e := 0; e < len(a); e++ {
		switch a[e].(type) {
		case string:
			if a[e].(string) != b[e].(string) {
				return a[e].(string) < b[e].(string)
			}
		case int64:
			if a[e].(int64) != b[e].(int64) {
				return a[e].(int64) < b[e].(int64)
			}
		case bool:
			if a[e].(bool) != b[e].(bool) {
				return a[e].(bool) == false && b[e].(bool) == true
			}
		}
	}
	return false
}

func (il imagesForDisplay) sortKeys(i int) []interface{} {
	im := il[i]
	return []interface{}{
		!im.Official,
		im.SourceType == "snapshot",
		im.Name, im.Arch,
		im.Status == "deprecated",
		!im.Public,
		-im.CreatedAt.Unix(),
	}
}

func (l *ImagesCommand) list(pc *kingpin.ParseContext) error {
	err := l.Configure()
	if err != nil {
		return err
	}

	// Get the account id in parallel to everything else, we don't need it until
	// later anyway.
	achan := make(chan string)
	go func() {
		achan <- l.accountId()
	}()

	out := new(RowFieldOutput)
	out.Setup(strings.Split(l.Fields, ","))

	images, err := l.Client.Images()
	if err != nil {
		return err
	}
	sortedImages := make(imagesForDisplay, 0, len(*images))
	for _, i := range *images {
		sortedImages = append(sortedImages, i)
	}
	sort.Sort(sortedImages)

	accountId := <-achan
	out.SendHeader()
	for _, i := range sortedImages {
		if l.ShowAll == false {
			if !i.Official && i.Owner != accountId {
				continue
			}
		}
		if err = out.Write(ImageFields(&i)); err != nil {
			return err
		}
	}
	out.Flush()
	return nil
}

func (l *ImagesCommand) show(pc *kingpin.ParseContext) error {

	err := l.Configure()
	if err != nil {
		return err
	}
	out := new(ShowFieldOutput)
	out.Setup(strings.Split(l.Fields, ","))

	i, err := l.Client.Image(l.Id)
	if err != nil {
		l.Fatalf(err.Error())
	}
	if err = out.Write(ImageFields(i)); err != nil {
		return err
	}
	out.Flush()
	return nil

}

func (l *ImagesCommand) destroy(pc *kingpin.ParseContext) error {
	err := l.Configure()
	if err != nil {
		return err
	}
	returnError := false
	for _, id := range l.IdList {
		fmt.Printf("Destroying image %s\n", id)
		err := l.Client.DestroyImage(id)
		if err != nil {
			l.Errorf("%s: %s", err.Error(), id)
			returnError = true
		}
	}
	if returnError {
		return genericError
	}
	return nil
}

func ConfigureImagesCommand(app *CliApp) {
	cmd := ImagesCommand{CliApp: app}
	images := app.Command("images", "Manage server images")
	list := images.Command("list", "List server images").Default().Action(cmd.list)
	list.Flag("fields", "Which fields to display").
		Default(DefaultImageListFields).
		StringVar(&cmd.Fields)
	list.Flag("show-all", "Show all public images from all accounts").Default("false").BoolVar(&cmd.ShowAll)
	show := images.Command("show", "View details of a server image").Action(cmd.show)
	show.Flag("fields", "Which fields to display").
		Default(DefaultImageShowFields).
		StringVar(&cmd.Fields)
	show.Arg("identifier", "Identifier of image to show").Required().StringVar(&cmd.Id)
	destroy := images.Command("destroy", "Destroy a server image").Action(cmd.destroy)
	destroy.Arg("identifier", "Identifier of image to destroy").Required().StringsVar(&cmd.IdList)

}
<|endoftext|>


<commit_before>package commands

import (
	"errors"

	cmds "github.com/ipfs/go-ipfs/commands"

	cmdkit "gx/ipfs/QmSNbH2A1evCCbJSDC6u3RV3GGDhgu6pRGbXHvrN89tMKf/go-ipfs-cmdkit"
)

var MountCmd = &cmds.Command{
	Helptext: cmdkit.HelpText{
		Tagline:          "Not yet implemented on Windows.",
		ShortDescription: "Not yet implemented on Windows. :(",
	},

	Run: func(req cmds.Request, res cmds.Response) {
		res.SetError(errors.New("Mount isn't compatible with Windows yet"), cmdkit.ErrNormal)
	},
}
<commit_msg>fix builds on windows<commit_after>package commands

import (
	"errors"

	cmds "github.com/ipfs/go-ipfs/commands"

	cmdkit "gx/ipfs/QmUyfy4QSr3NXym4etEiRyxBLqqAeKHJuRdi8AACxg63fZ/go-ipfs-cmdkit"
)

var MountCmd = &cmds.Command{
	Helptext: cmdkit.HelpText{
		Tagline:          "Not yet implemented on Windows.",
		ShortDescription: "Not yet implemented on Windows. :(",
	},

	Run: func(req cmds.Request, res cmds.Response) {
		res.SetError(errors.New("Mount isn't compatible with Windows yet"), cmdkit.ErrNormal)
	},
}
<|endoftext|>


<commit_before>package cli

import (
	"fmt"
	"log"
	"time"

	"github.com/99designs/aws-vault/vault"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/iam"
	"gopkg.in/alecthomas/kingpin.v2"
)

type RotateCommandInput struct {
	NoSession   bool
	ProfileName string
	Keyring     *vault.CredentialKeyring
	Config      vault.Config
}

func ConfigureRotateCommand(app *kingpin.Application) {
	input := RotateCommandInput{}

	cmd := app.Command("rotate", "Rotates credentials")

	cmd.Flag("no-session", "Use master credentials, no session or role used").
		Short('n').
		BoolVar(&input.NoSession)

	cmd.Arg("profile", "Name of the profile").
		Required().
		HintAction(awsConfigFile.ProfileNames).
		StringVar(&input.ProfileName)

	cmd.Action(func(c *kingpin.ParseContext) error {
		input.Config.MfaPromptMethod = GlobalFlags.PromptDriver
		input.Keyring = &vault.CredentialKeyring{Keyring: keyringImpl}
		app.FatalIfError(RotateCommand(input), "rotate")
		return nil
	})
}

func RotateCommand(input RotateCommandInput) error {
	vault.UseSession = !input.NoSession
	vault.UseSessionCache = false

	configLoader.BaseConfig = input.Config
	configLoader.ActiveProfile = input.ProfileName
	config, err := configLoader.LoadFromProfile(input.ProfileName)
	if err != nil {
		return err
	}

	masterCredentialsName, err := vault.MasterCredentialsFor(input.ProfileName, input.Keyring, config)
	if err != nil {
		return err
	}

	if input.NoSession {
		fmt.Printf("Rotating credentials stored for profile '%s' using master credentials (takes 10-20 seconds)\n", masterCredentialsName)
	} else {
		fmt.Printf("Rotating credentials stored for profile '%s' using a session from profile '%s' (takes 10-20 seconds)\n", masterCredentialsName, input.ProfileName)
	}

	// Get the existing credentials access key ID
	oldMasterCreds, err := vault.NewMasterCredentials(input.Keyring, masterCredentialsName).Get()
	if err != nil {
		return err
	}
	oldMasterCredsAccessKeyID := vault.FormatKeyForDisplay(oldMasterCreds.AccessKeyID)
	log.Printf("Rotating access key %s\n", oldMasterCredsAccessKeyID)

	fmt.Println("Creating a new access key")

	// create a session to rotate the credentials
	var sessCreds *credentials.Credentials
	if input.NoSession {
		sessCreds = vault.NewMasterCredentials(input.Keyring, config.ProfileName)
	} else {
		sessCreds, err = vault.NewTempCredentials(config, input.Keyring)
		if err != nil {
			return fmt.Errorf("Error getting temporary credentials: %w", err)
		}
	}

	sess, err := vault.NewSession(sessCreds, config.Region)
	if err != nil {
		return err
	}

	// A username is needed for some IAM calls if the credentials have assumed a role
	iamUserName, err := getUsernameIfAssumingRole(sess, config)
	if err != nil {
		return err
	}

	// Create a new access key
	createOut, err := iam.New(sess).CreateAccessKey(&iam.CreateAccessKeyInput{
		UserName: iamUserName,
	})
	if err != nil {
		return err
	}
	fmt.Printf("Created new access key %s\n", vault.FormatKeyForDisplay(*createOut.AccessKey.AccessKeyId))

	newMasterCreds := credentials.Value{
		AccessKeyID:     *createOut.AccessKey.AccessKeyId,
		SecretAccessKey: *createOut.AccessKey.SecretAccessKey,
	}

	err = input.Keyring.Set(masterCredentialsName, newMasterCreds)
	if err != nil {
		return fmt.Errorf("Error storing new access key %s: %w", vault.FormatKeyForDisplay(newMasterCreds.AccessKeyID), err)
	}

	// Delete old sessions
	sessions := input.Keyring.Sessions()
	profileNames, err := getProfilesInChain(input.ProfileName, configLoader)
	for _, profileName := range profileNames {
		if n, _ := sessions.Delete(profileName); n > 0 {
			fmt.Printf("Deleted %d sessions for %s\n", n, profileName)
		}
	}

	// expire the cached credentials
	sessCreds.Expire()

	// Use new credentials to delete old access key
	fmt.Printf("Deleting old access key %s\n", oldMasterCredsAccessKeyID)
	err = retry(time.Second*20, time.Second*2, func() error {
		_, err = iam.New(sess).DeleteAccessKey(&iam.DeleteAccessKeyInput{
			AccessKeyId: &oldMasterCreds.AccessKeyID,
			UserName:    iamUserName,
		})
		return err
	})
	if err != nil {
		return fmt.Errorf("Can't delete old access key %s: %w", oldMasterCredsAccessKeyID, err)
	}
	fmt.Printf("Deleted old access key %s\n", oldMasterCredsAccessKeyID)

	fmt.Println("Finished rotating access key")

	return nil
}

func retry(maxTime time.Duration, sleep time.Duration, f func() error) (err error) {
	t0 := time.Now()
	i := 0
	for {
		i++

		err = f()
		if err == nil {
			return
		}

		elapsed := time.Since(t0)
		if elapsed > maxTime {
			return fmt.Errorf("After %d attempts, last error: %s", i, err)
		}

		time.Sleep(sleep)
		log.Println("Retrying after error:", err)
	}
}

func getUsernameIfAssumingRole(sess *session.Session, config *vault.Config) (*string, error) {
	if config.RoleARN != "" {
		n, err := vault.GetUsernameFromSession(sess)
		if err != nil {
			return nil, err
		}
		log.Printf("Found IAM username '%s'", n)
		return &n, nil
	}
	return nil, nil
}

func getProfilesInChain(profileName string, configLoader *vault.ConfigLoader) (profileNames []string, err error) {
	profileNames = append(profileNames, profileName)

	config, err := configLoader.LoadFromProfile(profileName)
	if err != nil {
		return profileNames, err
	}

	if config.SourceProfile != nil {
		newProfileNames, err := getProfilesInChain(config.SourceProfileName, configLoader)
		if err != nil {
			return profileNames, err
		}
		profileNames = append(profileNames, newProfileNames...)
	}

	return profileNames, nil
}
<commit_msg>Add comment to explain why we might need sesssions<commit_after>package cli

import (
	"fmt"
	"log"
	"time"

	"github.com/99designs/aws-vault/vault"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/iam"
	"gopkg.in/alecthomas/kingpin.v2"
)

type RotateCommandInput struct {
	NoSession   bool
	ProfileName string
	Keyring     *vault.CredentialKeyring
	Config      vault.Config
}

func ConfigureRotateCommand(app *kingpin.Application) {
	input := RotateCommandInput{}

	cmd := app.Command("rotate", "Rotates credentials")

	cmd.Flag("no-session", "Use master credentials, no session or role used").
		Short('n').
		BoolVar(&input.NoSession)

	cmd.Arg("profile", "Name of the profile").
		Required().
		HintAction(awsConfigFile.ProfileNames).
		StringVar(&input.ProfileName)

	cmd.Action(func(c *kingpin.ParseContext) error {
		input.Config.MfaPromptMethod = GlobalFlags.PromptDriver
		input.Keyring = &vault.CredentialKeyring{Keyring: keyringImpl}
		app.FatalIfError(RotateCommand(input), "rotate")
		return nil
	})
}

func RotateCommand(input RotateCommandInput) error {
	// Can't disable sessions completely, might need to use session for MFA-Protected API Access
	vault.UseSession = !input.NoSession
	vault.UseSessionCache = false

	configLoader.BaseConfig = input.Config
	configLoader.ActiveProfile = input.ProfileName
	config, err := configLoader.LoadFromProfile(input.ProfileName)
	if err != nil {
		return err
	}

	masterCredentialsName, err := vault.MasterCredentialsFor(input.ProfileName, input.Keyring, config)
	if err != nil {
		return err
	}

	if input.NoSession {
		fmt.Printf("Rotating credentials stored for profile '%s' using master credentials (takes 10-20 seconds)\n", masterCredentialsName)
	} else {
		fmt.Printf("Rotating credentials stored for profile '%s' using a session from profile '%s' (takes 10-20 seconds)\n", masterCredentialsName, input.ProfileName)
	}

	// Get the existing credentials access key ID
	oldMasterCreds, err := vault.NewMasterCredentials(input.Keyring, masterCredentialsName).Get()
	if err != nil {
		return err
	}
	oldMasterCredsAccessKeyID := vault.FormatKeyForDisplay(oldMasterCreds.AccessKeyID)
	log.Printf("Rotating access key %s\n", oldMasterCredsAccessKeyID)

	fmt.Println("Creating a new access key")

	// create a session to rotate the credentials
	var sessCreds *credentials.Credentials
	if input.NoSession {
		sessCreds = vault.NewMasterCredentials(input.Keyring, config.ProfileName)
	} else {
		sessCreds, err = vault.NewTempCredentials(config, input.Keyring)
		if err != nil {
			return fmt.Errorf("Error getting temporary credentials: %w", err)
		}
	}

	sess, err := vault.NewSession(sessCreds, config.Region)
	if err != nil {
		return err
	}

	// A username is needed for some IAM calls if the credentials have assumed a role
	iamUserName, err := getUsernameIfAssumingRole(sess, config)
	if err != nil {
		return err
	}

	// Create a new access key
	createOut, err := iam.New(sess).CreateAccessKey(&iam.CreateAccessKeyInput{
		UserName: iamUserName,
	})
	if err != nil {
		return err
	}
	fmt.Printf("Created new access key %s\n", vault.FormatKeyForDisplay(*createOut.AccessKey.AccessKeyId))

	newMasterCreds := credentials.Value{
		AccessKeyID:     *createOut.AccessKey.AccessKeyId,
		SecretAccessKey: *createOut.AccessKey.SecretAccessKey,
	}

	err = input.Keyring.Set(masterCredentialsName, newMasterCreds)
	if err != nil {
		return fmt.Errorf("Error storing new access key %s: %w", vault.FormatKeyForDisplay(newMasterCreds.AccessKeyID), err)
	}

	// Delete old sessions
	sessions := input.Keyring.Sessions()
	profileNames, err := getProfilesInChain(input.ProfileName, configLoader)
	for _, profileName := range profileNames {
		if n, _ := sessions.Delete(profileName); n > 0 {
			fmt.Printf("Deleted %d sessions for %s\n", n, profileName)
		}
	}

	// expire the cached credentials
	sessCreds.Expire()

	// Use new credentials to delete old access key
	fmt.Printf("Deleting old access key %s\n", oldMasterCredsAccessKeyID)
	err = retry(time.Second*20, time.Second*2, func() error {
		_, err = iam.New(sess).DeleteAccessKey(&iam.DeleteAccessKeyInput{
			AccessKeyId: &oldMasterCreds.AccessKeyID,
			UserName:    iamUserName,
		})
		return err
	})
	if err != nil {
		return fmt.Errorf("Can't delete old access key %s: %w", oldMasterCredsAccessKeyID, err)
	}
	fmt.Printf("Deleted old access key %s\n", oldMasterCredsAccessKeyID)

	fmt.Println("Finished rotating access key")

	return nil
}

func retry(maxTime time.Duration, sleep time.Duration, f func() error) (err error) {
	t0 := time.Now()
	i := 0
	for {
		i++

		err = f()
		if err == nil {
			return
		}

		elapsed := time.Since(t0)
		if elapsed > maxTime {
			return fmt.Errorf("After %d attempts, last error: %s", i, err)
		}

		time.Sleep(sleep)
		log.Println("Retrying after error:", err)
	}
}

func getUsernameIfAssumingRole(sess *session.Session, config *vault.Config) (*string, error) {
	if config.RoleARN != "" {
		n, err := vault.GetUsernameFromSession(sess)
		if err != nil {
			return nil, err
		}
		log.Printf("Found IAM username '%s'", n)
		return &n, nil
	}
	return nil, nil
}

func getProfilesInChain(profileName string, configLoader *vault.ConfigLoader) (profileNames []string, err error) {
	profileNames = append(profileNames, profileName)

	config, err := configLoader.LoadFromProfile(profileName)
	if err != nil {
		return profileNames, err
	}

	if config.SourceProfile != nil {
		newProfileNames, err := getProfilesInChain(config.SourceProfileName, configLoader)
		if err != nil {
			return profileNames, err
		}
		profileNames = append(profileNames, newProfileNames...)
	}

	return profileNames, nil
}
<|endoftext|>