<commit_before>// Copyright 2015 The Cockroach Authors.
//
// 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. See the AUTHORS file
// for names of contributors.
//
// Author: Peter Mattis (peter@cockroachlabs.com)

package localcluster

import (
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"sync"

	"github.com/cockroachdb/cockroach/util/log"
	"github.com/samalba/dockerclient"
)

const (
	builderImage   = "cockroachdb/builder"
	dockerspyImage = "cockroachdb/docker-spy"
	domain         = "local"
)

var cockroachImage = flag.String("i", builderImage, "the docker image to run")
var cockroachBinary = flag.String("b", defaultBinary(), "the binary to run (if image == "+builderImage+")")
var cockroachEntry = flag.String("e", "", "the entry point for the image")
var testCertsDir = filepath.Clean(os.ExpandEnv("${PWD}/../resource/test_certs"))

func prettyJSON(v interface{}) string {
	pretty, err := json.MarshalIndent(v, "", "  ")
	if err != nil {
		panic(err)
	}
	return string(pretty)
}

func defaultBinary() string {
	gopath := filepath.SplitList(os.Getenv("GOPATH"))
	if len(gopath) == 0 {
		return ""
	}
	return gopath[0] + "/bin/linux_amd64/cockroach"
}

func exists(path string) bool {
	if _, err := os.Stat(path); os.IsNotExist(err) {
		return false
	}
	return true
}

func node(i int) string {
	return fmt.Sprintf("roach%d.%s", i, domain)
}

func data(i int) string {
	return fmt.Sprintf("/data%d", i)
}

// The various event types.
const (
	EventCreate     = "create"
	EventDestroy    = "destroy"
	EventDie        = "die"
	EventExecCreate = "exec_create"
	EventExecStart  = "exec_start"
	EventExport     = "export"
	EventKill       = "kill"
	EventOom        = "oom"
	EventPause      = "pause"
	EventRestart    = "restart"
	EventStart      = "start"
	EventStop       = "stop"
	EventUnpause    = "unpause"
)

// Event for a node containing a node index and the type of event.
type Event struct {
	NodeIndex int
	Status    string
}

// Cluster manages a local cockroach cluster running on docker. The cluster is
// composed of a "dns" container which automatically registers dns entries for
// the cockroach nodes, a "volumes" container which manages the persistent
// volumes used for certs and node data and N cockroach nodes.
type Cluster struct {
	client         dockerclient.Client
	stopper        chan struct{}
	mu             sync.Mutex // Protects the fields below
	dns            *Container
	vols           *Container
	Nodes          []*Container
	Events         chan Event
	CertsDir       string
	UseTestCerts   bool
	monitorStopper chan struct{}
}

// Create creates a new local cockroach cluster. The stopper is used to
// gracefully shutdown the channel (e.g. when a signal arrives). The cluster
// must be started before being used.
func Create(numNodes int, stopper chan struct{}) *Cluster {
	select {
	case <-stopper:
		// The stopper was already closed, exit early.
		os.Exit(1)
	default:
	}

	if *cockroachImage == builderImage && !exists(*cockroachBinary) {
		log.Fatalf("\"%s\": does not exist", *cockroachBinary)
	}

	return &Cluster{
		client:       newDockerClient(),
		stopper:      stopper,
		Nodes:        make([]*Container, numNodes),
		UseTestCerts: true,
	}
}

// stopOnPanic is invoked as a deferred function in Start in order to attempt
// to tear down the cluster if a panic occurs while starting it. If the panic
// was initiated by the stopper being closed (which panicOnStop notices) then
// the process is exited with a failure code.
func (l *Cluster) stopOnPanic() {
	if r := recover(); r != nil {
		l.Stop()
		if r != l {
			panic(r)
		}
		os.Exit(1)
	}
}

// panicOnStop tests whether the stopper channel has been closed and panics if
// it has. This allows polling for whether to stop and avoids nasty locking
// complications with trying to call Stop at arbitrary points such as in the
// middle of creating a container.
func (l *Cluster) panicOnStop() {
	if l.stopper == nil {
		panic(l)
	}

	select {
	case <-l.stopper:
		l.stopper = nil
		panic(l)
	default:
	}
}

func (l *Cluster) runDockerSpy() {
	l.panicOnStop()

	create := func() (*Container, error) {
		return createContainer(l.client, dockerclient.ContainerConfig{
			Image: dockerspyImage,
			Cmd:   []string{"--dns-domain=" + domain},
		})
	}
	c, err := create()
	if err == dockerclient.ErrNotFound {
		log.Infof("pulling %s", dockerspyImage)
		err = l.client.PullImage(dockerspyImage, nil)
		if err == nil {
			c, err = create()
		}
	}
	if err != nil {
		panic(err)
	}
	maybePanic(c.Start([]string{"/var/run/docker.sock:/var/run/docker.sock"}, nil, nil))
	maybePanic(c.Inspect())
	c.Name = "docker-spy"
	l.dns = c
	log.Infof("started %s: %s", c.Name, c.NetworkSettings.IPAddress)
}

// create the volumes container that keeps all of the volumes used by
// the cluster.
func (l *Cluster) initCluster() {
	log.Infof("initializing cluster")
	l.panicOnStop()

	vols := map[string]struct{}{}
	for i := range l.Nodes {
		vols[data(i)] = struct{}{}
	}
	create := func() (*Container, error) {
		var entrypoint []string
		if *cockroachImage == builderImage {
			entrypoint = append(entrypoint, "/"+filepath.Base(*cockroachBinary))
		} else if *cockroachEntry != "" {
			entrypoint = append(entrypoint, *cockroachEntry)
		}
		return createContainer(l.client, dockerclient.ContainerConfig{
			Image:      *cockroachImage,
			Volumes:    vols,
			Entrypoint: entrypoint,
			Cmd:        []string{"init", "--stores=ssd=" + data(0)},
		})
	}
	c, err := create()
	if err == dockerclient.ErrNotFound && *cockroachImage == builderImage {
		log.Infof("pulling %s", *cockroachImage)
		err = l.client.PullImage(*cockroachImage, nil)
		if err == nil {
			c, err = create()
		}
	}
	if err != nil {
		panic(err)
	}

	// Use pre-created test-certs when possible: cert creation takes ~5
	// seconds.
	if l.UseTestCerts {
		l.CertsDir = testCertsDir
	} else {
		// Create the temporary certs directory in the current working
		// directory. Boot2docker's handling of binding local directories
		// into the container is very confusing. If the directory being
		// bound has a parent directory that exists in the boot2docker VM
		// then that directory is bound into the container. In particular,
		// that means that binds of /tmp and /var will be problematic.
		cwd, err := os.Getwd()
		if err != nil {
			panic(err)
		}
		l.CertsDir, err = ioutil.TempDir(cwd, ".localcluster.certs.")
		if err != nil {
			panic(err)
		}
	}

	binds := []string{l.CertsDir + ":/certs"}
	if *cockroachImage == builderImage {
		path, err := filepath.Abs(*cockroachBinary)
		if err != nil {
			panic(err)
		}
		binds = append(binds, path+":/"+filepath.Base(*cockroachBinary))
	}
	maybePanic(c.Start(binds, nil, nil))
	maybePanic(c.Wait())
	c.Name = "volumes"
	l.vols = c
}

func (l *Cluster) createRoach(i int, cmd ...string) *Container {
	l.panicOnStop()

	var hostname string
	if i >= 0 {
		hostname = fmt.Sprintf("roach%d", i)
	}
	var entrypoint []string
	if *cockroachImage == builderImage {
		entrypoint = append(entrypoint, "/"+filepath.Base(*cockroachBinary))
	} else if *cockroachEntry != "" {
		entrypoint = append(entrypoint, *cockroachEntry)
	}
	c, err := createContainer(l.client, dockerclient.ContainerConfig{
		Hostname:     hostname,
		Domainname:   domain,
		Image:        *cockroachImage,
		ExposedPorts: map[string]struct{}{"8080/tcp": {}},
		Entrypoint:   entrypoint,
		Cmd:          cmd,
	})
	if err != nil {
		panic(err)
	}
	return c
}

func (l *Cluster) createCACert() {
	log.Infof("creating ca")
	c := l.createRoach(-1, "cert", "--certs=/certs", "create-ca")
	defer c.mustRemove()
	maybePanic(c.Start(nil, nil, l.vols))
	maybePanic(c.Wait())
}

func (l *Cluster) createNodeCerts() {
	log.Infof("creating node certs")
	var nodes []string
	for i := range l.Nodes {
		nodes = append(nodes, node(i))
	}
	args := []string{"cert", "--certs=/certs", "create-node"}
	args = append(args, nodes...)
	c := l.createRoach(-1, args...)
	defer c.mustRemove()
	maybePanic(c.Start(nil, nil, l.vols))
	maybePanic(c.Wait())
}

func (l *Cluster) startNode(i int) *Container {
	gossipNodes := []string{}
	for i := range l.Nodes {
		gossipNodes = append(gossipNodes, node(i)+":8080")
	}

	cmd := []string{
		"start",
		"--stores=ssd=" + data(i),
		"--certs=/certs",
		"--addr=" + node(i) + ":8080",
		"--gossip=" + strings.Join(gossipNodes, ","),
	}
	c := l.createRoach(i, cmd...)
	maybePanic(c.Start(nil, l.dns, l.vols))
	maybePanic(c.Inspect())
	c.Name = node(i)
	log.Infof("started %s: %s", c.Name, c.Addr(""))
	return c
}

func (l *Cluster) processEvent(e dockerclient.EventOrError, monitorStopper chan struct{}) bool {
	l.mu.Lock()
	defer l.mu.Unlock()

	if e.Error != nil {
		log.Errorf("monitoring error: %s", e.Error)
		if l.Events != nil {
			l.Events <- Event{NodeIndex: -1, Status: EventDie}
		}
		return false
	}

	for i, n := range l.Nodes {
		if n != nil && n.Id == e.Id {
			if log.V(1) {
				log.Errorf("node=%d status=%s", i, e.Status)
			}
			if l.Events != nil {
				l.Events <- Event{NodeIndex: i, Status: e.Status}
			}
			return true
		}
	}

	// TODO(pmattis): When we add the ability to start/stop/restart nodes we'll
	// need to keep around a map of old node container ids in order to ignore
	// events on those containers.

	// An event on any other container is unexpected. Die.
	select {
	case <-l.stopper:
	case <-monitorStopper:
	default:
		// There is a very tiny race here: the signal handler might be closing the
		// stopper simultaneously.
		log.Errorf("stopping due to unexpected event: %+v", e)
		r, err := l.client.ContainerLogs(e.Id, &dockerclient.LogOptions{
			Stdout: true,
			Stderr: true,
		})
		if err == nil {
			defer r.Close()
			_, err = io.Copy(os.Stdout, r)
		}
		close(l.stopper)
	}
	return false
}

func (l *Cluster) monitor(ch <-chan dockerclient.EventOrError, monitorStopper chan struct{}) {
	for e := range ch {
		if !l.processEvent(e, monitorStopper) {
			break
		}
	}
}

// Start starts the cluster.
func (l *Cluster) Start() {
	defer l.stopOnPanic()

	l.mu.Lock()
	defer l.mu.Unlock()

	l.runDockerSpy()
	l.initCluster()
	if l.CertsDir != testCertsDir {
		l.createCACert()
		l.createNodeCerts()
	}

	l.monitorStopper = make(chan struct{})
	ch, err := l.client.MonitorEvents(nil, l.monitorStopper)
	if err != nil {
		panic(err)
	}
	go l.monitor(ch, l.monitorStopper)

	for i := range l.Nodes {
		l.Nodes[i] = l.startNode(i)
	}
}

// Stop stops the clusters. It is safe to stop the cluster multiple times.
func (l *Cluster) Stop() {
	log.Infof("stopping")

	l.mu.Lock()
	defer l.mu.Unlock()

	if l.monitorStopper != nil {
		close(l.monitorStopper)
		l.monitorStopper = nil
	}

	if l.dns != nil {
		maybePanic(l.dns.Kill())
		l.dns = nil
	}
	if l.vols != nil {
		maybePanic(l.vols.Kill())
		l.vols = nil
	}
	if l.CertsDir != "" && l.CertsDir != testCertsDir {
		_ = os.RemoveAll(l.CertsDir)
		l.CertsDir = ""
	}
	for i, n := range l.Nodes {
		if n != nil {
			maybePanic(n.Kill())
			l.Nodes[i] = nil
		}
	}
}
<commit_msg>Set --scan-max-idle-time=200ms in order to speed up range replication.<commit_after>// Copyright 2015 The Cockroach Authors.
//
// 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. See the AUTHORS file
// for names of contributors.
//
// Author: Peter Mattis (peter@cockroachlabs.com)

package localcluster

import (
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"sync"

	"github.com/cockroachdb/cockroach/util/log"
	"github.com/samalba/dockerclient"
)

const (
	builderImage   = "cockroachdb/builder"
	dockerspyImage = "cockroachdb/docker-spy"
	domain         = "local"
)

var cockroachImage = flag.String("i", builderImage, "the docker image to run")
var cockroachBinary = flag.String("b", defaultBinary(), "the binary to run (if image == "+builderImage+")")
var cockroachEntry = flag.String("e", "", "the entry point for the image")
var testCertsDir = filepath.Clean(os.ExpandEnv("${PWD}/../resource/test_certs"))

func prettyJSON(v interface{}) string {
	pretty, err := json.MarshalIndent(v, "", "  ")
	if err != nil {
		panic(err)
	}
	return string(pretty)
}

func defaultBinary() string {
	gopath := filepath.SplitList(os.Getenv("GOPATH"))
	if len(gopath) == 0 {
		return ""
	}
	return gopath[0] + "/bin/linux_amd64/cockroach"
}

func exists(path string) bool {
	if _, err := os.Stat(path); os.IsNotExist(err) {
		return false
	}
	return true
}

func node(i int) string {
	return fmt.Sprintf("roach%d.%s", i, domain)
}

func data(i int) string {
	return fmt.Sprintf("/data%d", i)
}

// The various event types.
const (
	EventCreate     = "create"
	EventDestroy    = "destroy"
	EventDie        = "die"
	EventExecCreate = "exec_create"
	EventExecStart  = "exec_start"
	EventExport     = "export"
	EventKill       = "kill"
	EventOom        = "oom"
	EventPause      = "pause"
	EventRestart    = "restart"
	EventStart      = "start"
	EventStop       = "stop"
	EventUnpause    = "unpause"
)

// Event for a node containing a node index and the type of event.
type Event struct {
	NodeIndex int
	Status    string
}

// Cluster manages a local cockroach cluster running on docker. The cluster is
// composed of a "dns" container which automatically registers dns entries for
// the cockroach nodes, a "volumes" container which manages the persistent
// volumes used for certs and node data and N cockroach nodes.
type Cluster struct {
	client         dockerclient.Client
	stopper        chan struct{}
	mu             sync.Mutex // Protects the fields below
	dns            *Container
	vols           *Container
	Nodes          []*Container
	Events         chan Event
	CertsDir       string
	UseTestCerts   bool
	monitorStopper chan struct{}
}

// Create creates a new local cockroach cluster. The stopper is used to
// gracefully shutdown the channel (e.g. when a signal arrives). The cluster
// must be started before being used.
func Create(numNodes int, stopper chan struct{}) *Cluster {
	select {
	case <-stopper:
		// The stopper was already closed, exit early.
		os.Exit(1)
	default:
	}

	if *cockroachImage == builderImage && !exists(*cockroachBinary) {
		log.Fatalf("\"%s\": does not exist", *cockroachBinary)
	}

	return &Cluster{
		client:       newDockerClient(),
		stopper:      stopper,
		Nodes:        make([]*Container, numNodes),
		UseTestCerts: true,
	}
}

// stopOnPanic is invoked as a deferred function in Start in order to attempt
// to tear down the cluster if a panic occurs while starting it. If the panic
// was initiated by the stopper being closed (which panicOnStop notices) then
// the process is exited with a failure code.
func (l *Cluster) stopOnPanic() {
	if r := recover(); r != nil {
		l.Stop()
		if r != l {
			panic(r)
		}
		os.Exit(1)
	}
}

// panicOnStop tests whether the stopper channel has been closed and panics if
// it has. This allows polling for whether to stop and avoids nasty locking
// complications with trying to call Stop at arbitrary points such as in the
// middle of creating a container.
func (l *Cluster) panicOnStop() {
	if l.stopper == nil {
		panic(l)
	}

	select {
	case <-l.stopper:
		l.stopper = nil
		panic(l)
	default:
	}
}

func (l *Cluster) runDockerSpy() {
	l.panicOnStop()

	create := func() (*Container, error) {
		return createContainer(l.client, dockerclient.ContainerConfig{
			Image: dockerspyImage,
			Cmd:   []string{"--dns-domain=" + domain},
		})
	}
	c, err := create()
	if err == dockerclient.ErrNotFound {
		log.Infof("pulling %s", dockerspyImage)
		err = l.client.PullImage(dockerspyImage, nil)
		if err == nil {
			c, err = create()
		}
	}
	if err != nil {
		panic(err)
	}
	maybePanic(c.Start([]string{"/var/run/docker.sock:/var/run/docker.sock"}, nil, nil))
	maybePanic(c.Inspect())
	c.Name = "docker-spy"
	l.dns = c
	log.Infof("started %s: %s", c.Name, c.NetworkSettings.IPAddress)
}

// create the volumes container that keeps all of the volumes used by
// the cluster.
func (l *Cluster) initCluster() {
	log.Infof("initializing cluster")
	l.panicOnStop()

	vols := map[string]struct{}{}
	for i := range l.Nodes {
		vols[data(i)] = struct{}{}
	}
	create := func() (*Container, error) {
		var entrypoint []string
		if *cockroachImage == builderImage {
			entrypoint = append(entrypoint, "/"+filepath.Base(*cockroachBinary))
		} else if *cockroachEntry != "" {
			entrypoint = append(entrypoint, *cockroachEntry)
		}
		return createContainer(l.client, dockerclient.ContainerConfig{
			Image:      *cockroachImage,
			Volumes:    vols,
			Entrypoint: entrypoint,
			Cmd:        []string{"init", "--stores=ssd=" + data(0)},
		})
	}
	c, err := create()
	if err == dockerclient.ErrNotFound && *cockroachImage == builderImage {
		log.Infof("pulling %s", *cockroachImage)
		err = l.client.PullImage(*cockroachImage, nil)
		if err == nil {
			c, err = create()
		}
	}
	if err != nil {
		panic(err)
	}

	// Use pre-created test-certs when possible: cert creation takes ~5
	// seconds.
	if l.UseTestCerts {
		l.CertsDir = testCertsDir
	} else {
		// Create the temporary certs directory in the current working
		// directory. Boot2docker's handling of binding local directories
		// into the container is very confusing. If the directory being
		// bound has a parent directory that exists in the boot2docker VM
		// then that directory is bound into the container. In particular,
		// that means that binds of /tmp and /var will be problematic.
		cwd, err := os.Getwd()
		if err != nil {
			panic(err)
		}
		l.CertsDir, err = ioutil.TempDir(cwd, ".localcluster.certs.")
		if err != nil {
			panic(err)
		}
	}

	binds := []string{l.CertsDir + ":/certs"}
	if *cockroachImage == builderImage {
		path, err := filepath.Abs(*cockroachBinary)
		if err != nil {
			panic(err)
		}
		binds = append(binds, path+":/"+filepath.Base(*cockroachBinary))
	}
	maybePanic(c.Start(binds, nil, nil))
	maybePanic(c.Wait())
	c.Name = "volumes"
	l.vols = c
}

func (l *Cluster) createRoach(i int, cmd ...string) *Container {
	l.panicOnStop()

	var hostname string
	if i >= 0 {
		hostname = fmt.Sprintf("roach%d", i)
	}
	var entrypoint []string
	if *cockroachImage == builderImage {
		entrypoint = append(entrypoint, "/"+filepath.Base(*cockroachBinary))
	} else if *cockroachEntry != "" {
		entrypoint = append(entrypoint, *cockroachEntry)
	}
	c, err := createContainer(l.client, dockerclient.ContainerConfig{
		Hostname:     hostname,
		Domainname:   domain,
		Image:        *cockroachImage,
		ExposedPorts: map[string]struct{}{"8080/tcp": {}},
		Entrypoint:   entrypoint,
		Cmd:          cmd,
	})
	if err != nil {
		panic(err)
	}
	return c
}

func (l *Cluster) createCACert() {
	log.Infof("creating ca")
	c := l.createRoach(-1, "cert", "--certs=/certs", "create-ca")
	defer c.mustRemove()
	maybePanic(c.Start(nil, nil, l.vols))
	maybePanic(c.Wait())
}

func (l *Cluster) createNodeCerts() {
	log.Infof("creating node certs")
	var nodes []string
	for i := range l.Nodes {
		nodes = append(nodes, node(i))
	}
	args := []string{"cert", "--certs=/certs", "create-node"}
	args = append(args, nodes...)
	c := l.createRoach(-1, args...)
	defer c.mustRemove()
	maybePanic(c.Start(nil, nil, l.vols))
	maybePanic(c.Wait())
}

func (l *Cluster) startNode(i int) *Container {
	gossipNodes := []string{}
	for i := range l.Nodes {
		gossipNodes = append(gossipNodes, node(i)+":8080")
	}

	cmd := []string{
		"start",
		"--stores=ssd=" + data(i),
		"--certs=/certs",
		"--addr=" + node(i) + ":8080",
		"--gossip=" + strings.Join(gossipNodes, ","),
		"--scan-max-idle-time=200ms", // set low to speed up tests
	}
	c := l.createRoach(i, cmd...)
	maybePanic(c.Start(nil, l.dns, l.vols))
	maybePanic(c.Inspect())
	c.Name = node(i)
	log.Infof("started %s: %s", c.Name, c.Addr(""))
	return c
}

func (l *Cluster) processEvent(e dockerclient.EventOrError, monitorStopper chan struct{}) bool {
	l.mu.Lock()
	defer l.mu.Unlock()

	if e.Error != nil {
		log.Errorf("monitoring error: %s", e.Error)
		if l.Events != nil {
			l.Events <- Event{NodeIndex: -1, Status: EventDie}
		}
		return false
	}

	for i, n := range l.Nodes {
		if n != nil && n.Id == e.Id {
			if log.V(1) {
				log.Errorf("node=%d status=%s", i, e.Status)
			}
			if l.Events != nil {
				l.Events <- Event{NodeIndex: i, Status: e.Status}
			}
			return true
		}
	}

	// TODO(pmattis): When we add the ability to start/stop/restart nodes we'll
	// need to keep around a map of old node container ids in order to ignore
	// events on those containers.

	// An event on any other container is unexpected. Die.
	select {
	case <-l.stopper:
	case <-monitorStopper:
	default:
		// There is a very tiny race here: the signal handler might be closing the
		// stopper simultaneously.
		log.Errorf("stopping due to unexpected event: %+v", e)
		r, err := l.client.ContainerLogs(e.Id, &dockerclient.LogOptions{
			Stdout: true,
			Stderr: true,
		})
		if err == nil {
			defer r.Close()
			_, err = io.Copy(os.Stdout, r)
		}
		close(l.stopper)
	}
	return false
}

func (l *Cluster) monitor(ch <-chan dockerclient.EventOrError, monitorStopper chan struct{}) {
	for e := range ch {
		if !l.processEvent(e, monitorStopper) {
			break
		}
	}
}

// Start starts the cluster.
func (l *Cluster) Start() {
	defer l.stopOnPanic()

	l.mu.Lock()
	defer l.mu.Unlock()

	l.runDockerSpy()
	l.initCluster()
	if l.CertsDir != testCertsDir {
		l.createCACert()
		l.createNodeCerts()
	}

	l.monitorStopper = make(chan struct{})
	ch, err := l.client.MonitorEvents(nil, l.monitorStopper)
	if err != nil {
		panic(err)
	}
	go l.monitor(ch, l.monitorStopper)

	for i := range l.Nodes {
		l.Nodes[i] = l.startNode(i)
	}
}

// Stop stops the clusters. It is safe to stop the cluster multiple times.
func (l *Cluster) Stop() {
	log.Infof("stopping")

	l.mu.Lock()
	defer l.mu.Unlock()

	if l.monitorStopper != nil {
		close(l.monitorStopper)
		l.monitorStopper = nil
	}

	if l.dns != nil {
		maybePanic(l.dns.Kill())
		l.dns = nil
	}
	if l.vols != nil {
		maybePanic(l.vols.Kill())
		l.vols = nil
	}
	if l.CertsDir != "" && l.CertsDir != testCertsDir {
		_ = os.RemoveAll(l.CertsDir)
		l.CertsDir = ""
	}
	for i, n := range l.Nodes {
		if n != nil {
			maybePanic(n.Kill())
			l.Nodes[i] = nil
		}
	}
}
<|endoftext|>


<commit_before>package cli

import (
	"fmt"
	"strings"

	"github.com/docker/docker/pkg/term"
	"github.com/pkg/errors"
	"github.com/spf13/cobra"
)

// SetupRootCommand sets default usage, help, and error handling for the
// root command.
func SetupRootCommand(rootCmd *cobra.Command) {
	cobra.AddTemplateFunc("hasSubCommands", hasSubCommands)
	cobra.AddTemplateFunc("hasManagementSubCommands", hasManagementSubCommands)
	cobra.AddTemplateFunc("operationSubCommands", operationSubCommands)
	cobra.AddTemplateFunc("managementSubCommands", managementSubCommands)
	cobra.AddTemplateFunc("wrappedFlagUsages", wrappedFlagUsages)

	rootCmd.SetUsageTemplate(usageTemplate)
	rootCmd.SetHelpTemplate(helpTemplate)
	rootCmd.SetFlagErrorFunc(FlagErrorFunc)
	rootCmd.SetHelpCommand(helpCommand)
	rootCmd.SetVersionTemplate("Docker version {{.Version}}\n")

	rootCmd.PersistentFlags().BoolP("help", "h", false, "Print usage")
	rootCmd.PersistentFlags().MarkShorthandDeprecated("help", "please use --help")
	rootCmd.PersistentFlags().Lookup("help").Hidden = true
}

// FlagErrorFunc prints an error message which matches the format of the
// docker/cli/cli error messages
func FlagErrorFunc(cmd *cobra.Command, err error) error {
	if err == nil {
		return nil
	}

	usage := ""
	if cmd.HasSubCommands() {
		usage = "\n\n" + cmd.UsageString()
	}
	return StatusError{
		Status:     fmt.Sprintf("%s\nSee '%s --help'.%s", err, cmd.CommandPath(), usage),
		StatusCode: 125,
	}
}

// VisitAll will traverse all commands from the root.
// This is different from the VisitAll of cobra.Command where only parents
// are checked.
func VisitAll(root *cobra.Command, fn func(*cobra.Command)) {
	for _, cmd := range root.Commands() {
		VisitAll(cmd, fn)
	}
	fn(root)
}

// DisableFlagsInUseLine sets the DisableFlagsInUseLine flag on all
// commands within the tree rooted at cmd.
func DisableFlagsInUseLine(cmd *cobra.Command) {
	VisitAll(cmd, func(ccmd *cobra.Command) {
		// do not add a `[flags]` to the end of the usage line.
		ccmd.DisableFlagsInUseLine = true
	})
}

var helpCommand = &cobra.Command{
	Use:               "help [command]",
	Short:             "Help about the command",
	PersistentPreRun:  func(cmd *cobra.Command, args []string) {},
	PersistentPostRun: func(cmd *cobra.Command, args []string) {},
	RunE: func(c *cobra.Command, args []string) error {
		cmd, args, e := c.Root().Find(args)
		if cmd == nil || e != nil || len(args) > 0 {
			return errors.Errorf("unknown help topic: %v", strings.Join(args, " "))
		}

		helpFunc := cmd.HelpFunc()
		helpFunc(cmd, args)
		return nil
	},
}

func hasSubCommands(cmd *cobra.Command) bool {
	return len(operationSubCommands(cmd)) > 0
}

func hasManagementSubCommands(cmd *cobra.Command) bool {
	return len(managementSubCommands(cmd)) > 0
}

func operationSubCommands(cmd *cobra.Command) []*cobra.Command {
	cmds := []*cobra.Command{}
	for _, sub := range cmd.Commands() {
		if sub.IsAvailableCommand() && !sub.HasSubCommands() {
			cmds = append(cmds, sub)
		}
	}
	return cmds
}

func wrappedFlagUsages(cmd *cobra.Command) string {
	width := 80
	if ws, err := term.GetWinsize(0); err == nil {
		width = int(ws.Width)
	}
	return cmd.Flags().FlagUsagesWrapped(width - 1)
}

func managementSubCommands(cmd *cobra.Command) []*cobra.Command {
	cmds := []*cobra.Command{}
	for _, sub := range cmd.Commands() {
		if sub.IsAvailableCommand() && sub.HasSubCommands() {
			cmds = append(cmds, sub)
		}
	}
	return cmds
}

var usageTemplate = `Usage:

{{- if not .HasSubCommands}}	{{.UseLine}}{{end}}
{{- if .HasSubCommands}}	{{ .CommandPath}}{{- if .HasAvailableFlags}} [OPTIONS]{{end}} COMMAND{{end}}

{{if ne .Long ""}}{{ .Long | trim }}{{ else }}{{ .Short | trim }}{{end}}

{{- if gt .Aliases 0}}

Aliases:
  {{.NameAndAliases}}

{{- end}}
{{- if .HasExample}}

Examples:
{{ .Example }}

{{- end}}
{{- if .HasAvailableFlags}}

Options:
{{ wrappedFlagUsages . | trimRightSpace}}

{{- end}}
{{- if hasManagementSubCommands . }}

Management Commands:

{{- range managementSubCommands . }}
  {{rpad .Name .NamePadding }} {{.Short}}
{{- end}}

{{- end}}
{{- if hasSubCommands .}}

Commands:

{{- range operationSubCommands . }}
  {{rpad .Name .NamePadding }} {{.Short}}
{{- end}}
{{- end}}

{{- if .HasSubCommands }}

Run '{{.CommandPath}} COMMAND --help' for more information on a command.
{{- end}}
`

var helpTemplate = `
{{if or .Runnable .HasSubCommands}}{{.UsageString}}{{end}}`
<commit_msg>Refactor common bits of `SetupRootCommand`<commit_after>package cli

import (
	"fmt"
	"strings"

	"github.com/docker/docker/pkg/term"
	"github.com/pkg/errors"
	"github.com/spf13/cobra"
)

func setupCommonRootCommand(rootCmd *cobra.Command) {
	cobra.AddTemplateFunc("hasSubCommands", hasSubCommands)
	cobra.AddTemplateFunc("hasManagementSubCommands", hasManagementSubCommands)
	cobra.AddTemplateFunc("operationSubCommands", operationSubCommands)
	cobra.AddTemplateFunc("managementSubCommands", managementSubCommands)
	cobra.AddTemplateFunc("wrappedFlagUsages", wrappedFlagUsages)

	rootCmd.SetUsageTemplate(usageTemplate)
	rootCmd.SetHelpTemplate(helpTemplate)
	rootCmd.SetFlagErrorFunc(FlagErrorFunc)
	rootCmd.SetHelpCommand(helpCommand)
}

// SetupRootCommand sets default usage, help, and error handling for the
// root command.
func SetupRootCommand(rootCmd *cobra.Command) {
	setupCommonRootCommand(rootCmd)

	rootCmd.SetVersionTemplate("Docker version {{.Version}}\n")

	rootCmd.PersistentFlags().BoolP("help", "h", false, "Print usage")
	rootCmd.PersistentFlags().MarkShorthandDeprecated("help", "please use --help")
	rootCmd.PersistentFlags().Lookup("help").Hidden = true
}

// FlagErrorFunc prints an error message which matches the format of the
// docker/cli/cli error messages
func FlagErrorFunc(cmd *cobra.Command, err error) error {
	if err == nil {
		return nil
	}

	usage := ""
	if cmd.HasSubCommands() {
		usage = "\n\n" + cmd.UsageString()
	}
	return StatusError{
		Status:     fmt.Sprintf("%s\nSee '%s --help'.%s", err, cmd.CommandPath(), usage),
		StatusCode: 125,
	}
}

// VisitAll will traverse all commands from the root.
// This is different from the VisitAll of cobra.Command where only parents
// are checked.
func VisitAll(root *cobra.Command, fn func(*cobra.Command)) {
	for _, cmd := range root.Commands() {
		VisitAll(cmd, fn)
	}
	fn(root)
}

// DisableFlagsInUseLine sets the DisableFlagsInUseLine flag on all
// commands within the tree rooted at cmd.
func DisableFlagsInUseLine(cmd *cobra.Command) {
	VisitAll(cmd, func(ccmd *cobra.Command) {
		// do not add a `[flags]` to the end of the usage line.
		ccmd.DisableFlagsInUseLine = true
	})
}

var helpCommand = &cobra.Command{
	Use:               "help [command]",
	Short:             "Help about the command",
	PersistentPreRun:  func(cmd *cobra.Command, args []string) {},
	PersistentPostRun: func(cmd *cobra.Command, args []string) {},
	RunE: func(c *cobra.Command, args []string) error {
		cmd, args, e := c.Root().Find(args)
		if cmd == nil || e != nil || len(args) > 0 {
			return errors.Errorf("unknown help topic: %v", strings.Join(args, " "))
		}

		helpFunc := cmd.HelpFunc()
		helpFunc(cmd, args)
		return nil
	},
}

func hasSubCommands(cmd *cobra.Command) bool {
	return len(operationSubCommands(cmd)) > 0
}

func hasManagementSubCommands(cmd *cobra.Command) bool {
	return len(managementSubCommands(cmd)) > 0
}

func operationSubCommands(cmd *cobra.Command) []*cobra.Command {
	cmds := []*cobra.Command{}
	for _, sub := range cmd.Commands() {
		if sub.IsAvailableCommand() && !sub.HasSubCommands() {
			cmds = append(cmds, sub)
		}
	}
	return cmds
}

func wrappedFlagUsages(cmd *cobra.Command) string {
	width := 80
	if ws, err := term.GetWinsize(0); err == nil {
		width = int(ws.Width)
	}
	return cmd.Flags().FlagUsagesWrapped(width - 1)
}

func managementSubCommands(cmd *cobra.Command) []*cobra.Command {
	cmds := []*cobra.Command{}
	for _, sub := range cmd.Commands() {
		if sub.IsAvailableCommand() && sub.HasSubCommands() {
			cmds = append(cmds, sub)
		}
	}
	return cmds
}

var usageTemplate = `Usage:

{{- if not .HasSubCommands}}	{{.UseLine}}{{end}}
{{- if .HasSubCommands}}	{{ .CommandPath}}{{- if .HasAvailableFlags}} [OPTIONS]{{end}} COMMAND{{end}}

{{if ne .Long ""}}{{ .Long | trim }}{{ else }}{{ .Short | trim }}{{end}}

{{- if gt .Aliases 0}}

Aliases:
  {{.NameAndAliases}}

{{- end}}
{{- if .HasExample}}

Examples:
{{ .Example }}

{{- end}}
{{- if .HasAvailableFlags}}

Options:
{{ wrappedFlagUsages . | trimRightSpace}}

{{- end}}
{{- if hasManagementSubCommands . }}

Management Commands:

{{- range managementSubCommands . }}
  {{rpad .Name .NamePadding }} {{.Short}}
{{- end}}

{{- end}}
{{- if hasSubCommands .}}

Commands:

{{- range operationSubCommands . }}
  {{rpad .Name .NamePadding }} {{.Short}}
{{- end}}
{{- end}}

{{- if .HasSubCommands }}

Run '{{.CommandPath}} COMMAND --help' for more information on a command.
{{- end}}
`

var helpTemplate = `
{{if or .Runnable .HasSubCommands}}{{.UsageString}}{{end}}`
<|endoftext|>


<commit_before>package cmd

import (
	"io/ioutil"
	"os"

	"github.com/Sirupsen/logrus"
	"github.com/spf13/cobra"

	"github.com/emc-advanced-dev/pkg/errors"
	"github.com/emc-advanced-dev/unik/pkg/client"
	unikos "github.com/emc-advanced-dev/unik/pkg/os"
	unikutil "github.com/emc-advanced-dev/unik/pkg/util"
)

var name, path, compiler, provider, runArgs string
var mountPoints []string
var force, noCleanup bool

var buildCmd = &cobra.Command{
	Use:   "build",
	Short: "Build a unikernel image from source code files",
	Long: `Compiles source files into a runnable unikernel image.

Images must be compiled for a specific provider, specified with the --provider flag
To see a list of available providers, run 'unik providers'

A unikernel compiler that is compatible with the provider must be specified with the --compiler flag
To see a list of available compilers, run 'unik compilers'

If you wish to attach volumes to instances of an image, the image must be compiled in advance
with a list of the expected mount points. e.g. for an application that reads from a '/data' folder,
the unikernel should be compiled with the flag -mount /data

Runtime arguments to be passed to your unikernel must also be specified at compile time.
You can specify arguments as a single string passed to the --args flag

Image names must be unique. If an image exists with the same name, you can force overwriting with the
--force flag

Example usage:
	unik build --name myUnikernel --path ./myApp/src --compiler rump-go-xen --provider aws --mountpoint /foo --mountpoint /bar --args '-myParameter MYVALUE' --force

	# will create a unikernel named myUnikernel using the sources found in ./myApp/src,
	# compiled using rumprun for the xen hypervisor, targeting AWS infrastructure,
	# expecting a volume to be mounted at /foo at runtime,
	# expecting another volume to be mounted at /bar at runtime,
	# passing '-myParameter MYVALUE' as arguments to the application when it is run,
	# and deleting any previous existing instances and image for the name myUnikernel before compiling

Another example (using only the required parameters):
	unik build --name anotherUnikernel --path ./anotherApp/src --compiler rump-go-vmware --provider vsphere
`,
	Run: func(cmd *cobra.Command, args []string) {
		if err := func() error {
			if name == "" {
				return errors.New("--name must be set", nil)
			}
			if path == "" {
				return errors.New("--path must be set", nil)
			}
			if compiler == "" {
				return errors.New("--compiler must be set", nil)
			}
			if provider == "" {
				return errors.New("--provider must be set", nil)
			}
			if err := readClientConfig(); err != nil {
				return err
			}
			if host == "" {
				host = clientConfig.Host
			}
			logrus.WithFields(logrus.Fields{
				"name":        name,
				"path":        path,
				"compiler":    compiler,
				"provider":    provider,
				"args":        runArgs,
				"mountPoints": mountPoints,
				"force":       force,
				"host":        host,
			}).Infof("running unik build")
			sourceTar, err := ioutil.TempFile(unikutil.UnikTmpDir(), "")
			if err != nil {
				logrus.WithError(err).Error("failed to create tmp tar file")
			}
			defer os.Remove(sourceTar.Name())
			if err := unikos.Compress(path, sourceTar.Name()); err != nil {
				return errors.New("failed to tar sources", err)
			}
			logrus.Infof("App packaged as tarball: %s\n", sourceTar.Name())
			image, err := client.UnikClient(host).Images().Build(name, sourceTar.Name(), compiler, provider, runArgs, mountPoints, force, noCleanup)
			if err != nil {
				return errors.New("building image failed", err)
			}
			printImages(image)
			return nil
		}(); err != nil {
			logrus.Errorf("build failed: %v", err)
			os.Exit(-1)
		}
	},
}

func init() {
	RootCmd.AddCommand(buildCmd)
	buildCmd.Flags().StringVar(&name, "name", "", "<string,required> name to give the unikernel. must be unique")
	buildCmd.Flags().StringVar(&path, "path", "", "<string,required> path to root application sources folder")
	buildCmd.Flags().StringVar(&compiler, "compiler", "", "<string,required> name of the unikernel compiler to use")
	buildCmd.Flags().StringVar(&provider, "provider", "", "<string,required> name of the target infrastructure to compile for")
	buildCmd.Flags().StringVar(&runArgs, "args", "", "<string,optional> to be passed to the unikernel at runtime")
	buildCmd.Flags().StringSliceVar(&mountPoints, "mountpoint", []string{}, "<string,repeated> specify up to 8 mount points for volumes")
	buildCmd.Flags().BoolVar(&force, "force", false, "<bool, optional> force overwriting a previously existing")
	buildCmd.Flags().BoolVar(&noCleanup, "no-cleanup", false, "<bool, optional> for debugging; do not clean up artifacts for images that fail to build")
}
<commit_msg>example command s,xen,aws,<commit_after>package cmd

import (
	"io/ioutil"
	"os"

	"github.com/Sirupsen/logrus"
	"github.com/spf13/cobra"

	"github.com/emc-advanced-dev/pkg/errors"
	"github.com/emc-advanced-dev/unik/pkg/client"
	unikos "github.com/emc-advanced-dev/unik/pkg/os"
	unikutil "github.com/emc-advanced-dev/unik/pkg/util"
)

var name, path, compiler, provider, runArgs string
var mountPoints []string
var force, noCleanup bool

var buildCmd = &cobra.Command{
	Use:   "build",
	Short: "Build a unikernel image from source code files",
	Long: `Compiles source files into a runnable unikernel image.

Images must be compiled for a specific provider, specified with the --provider flag
To see a list of available providers, run 'unik providers'

A unikernel compiler that is compatible with the provider must be specified with the --compiler flag
To see a list of available compilers, run 'unik compilers'

If you wish to attach volumes to instances of an image, the image must be compiled in advance
with a list of the expected mount points. e.g. for an application that reads from a '/data' folder,
the unikernel should be compiled with the flag -mount /data

Runtime arguments to be passed to your unikernel must also be specified at compile time.
You can specify arguments as a single string passed to the --args flag

Image names must be unique. If an image exists with the same name, you can force overwriting with the
--force flag

Example usage:
	unik build --name myUnikernel --path ./myApp/src --compiler rump-go-aws --provider aws --mountpoint /foo --mountpoint /bar --args '-myParameter MYVALUE' --force

	# will create a unikernel named myUnikernel using the sources found in ./myApp/src,
	# compiled using rumprun for the xen hypervisor, targeting AWS infrastructure,
	# expecting a volume to be mounted at /foo at runtime,
	# expecting another volume to be mounted at /bar at runtime,
	# passing '-myParameter MYVALUE' as arguments to the application when it is run,
	# and deleting any previous existing instances and image for the name myUnikernel before compiling

Another example (using only the required parameters):
	unik build --name anotherUnikernel --path ./anotherApp/src --compiler rump-go-vmware --provider vsphere
`,
	Run: func(cmd *cobra.Command, args []string) {
		if err := func() error {
			if name == "" {
				return errors.New("--name must be set", nil)
			}
			if path == "" {
				return errors.New("--path must be set", nil)
			}
			if compiler == "" {
				return errors.New("--compiler must be set", nil)
			}
			if provider == "" {
				return errors.New("--provider must be set", nil)
			}
			if err := readClientConfig(); err != nil {
				return err
			}
			if host == "" {
				host = clientConfig.Host
			}
			logrus.WithFields(logrus.Fields{
				"name":        name,
				"path":        path,
				"compiler":    compiler,
				"provider":    provider,
				"args":        runArgs,
				"mountPoints": mountPoints,
				"force":       force,
				"host":        host,
			}).Infof("running unik build")
			sourceTar, err := ioutil.TempFile(unikutil.UnikTmpDir(), "")
			if err != nil {
				logrus.WithError(err).Error("failed to create tmp tar file")
			}
			defer os.Remove(sourceTar.Name())
			if err := unikos.Compress(path, sourceTar.Name()); err != nil {
				return errors.New("failed to tar sources", err)
			}
			logrus.Infof("App packaged as tarball: %s\n", sourceTar.Name())
			image, err := client.UnikClient(host).Images().Build(name, sourceTar.Name(), compiler, provider, runArgs, mountPoints, force, noCleanup)
			if err != nil {
				return errors.New("building image failed", err)
			}
			printImages(image)
			return nil
		}(); err != nil {
			logrus.Errorf("build failed: %v", err)
			os.Exit(-1)
		}
	},
}

func init() {
	RootCmd.AddCommand(buildCmd)
	buildCmd.Flags().StringVar(&name, "name", "", "<string,required> name to give the unikernel. must be unique")
	buildCmd.Flags().StringVar(&path, "path", "", "<string,required> path to root application sources folder")
	buildCmd.Flags().StringVar(&compiler, "compiler", "", "<string,required> name of the unikernel compiler to use")
	buildCmd.Flags().StringVar(&provider, "provider", "", "<string,required> name of the target infrastructure to compile for")
	buildCmd.Flags().StringVar(&runArgs, "args", "", "<string,optional> to be passed to the unikernel at runtime")
	buildCmd.Flags().StringSliceVar(&mountPoints, "mountpoint", []string{}, "<string,repeated> specify up to 8 mount points for volumes")
	buildCmd.Flags().BoolVar(&force, "force", false, "<bool, optional> force overwriting a previously existing")
	buildCmd.Flags().BoolVar(&noCleanup, "no-cleanup", false, "<bool, optional> for debugging; do not clean up artifacts for images that fail to build")
}
<|endoftext|>


<commit_before>package cmd

import (
	"fmt"
	"math"
	"os"

	"github.com/phil-mansfield/shellfish/io"
	"github.com/phil-mansfield/shellfish/parse"
	"github.com/phil-mansfield/shellfish/cmd/env"
	"github.com/phil-mansfield/shellfish/cmd/memo"
	"github.com/phil-mansfield/shellfish/los/geom"
)

type CheckConfig struct {
	h0, omegaM, omegaL float64
	boxWidth float64
	particleMasses []float64
	particleCount int64
	exampleHalo []float64
}

var _ Mode = &CheckConfig{}

func (config *CheckConfig) ExampleConfig() string {
	return `[check.config]

# All fields are optional. Unless otherwise stated, checks are very rough and
# only done to 10%, so  report float values to at least two decimal places.

# H0 = 70
# OmegaM = 0.30
# OmegaL = 0.70

# BoxWidth is in units of Mpc/h (not your code units).
# BoxWidth = 125

# The total number of particles in the simulation box. This is checked exactly.
# ParticleCount = 1073741824

# ParticleMasses A list of all the valid DM particle masses in Msun/h (not your
# code units).
# ParticleMasses = 1.7e7, 1.8e8, 1.1e9

# ExampleHalo is a the location and size of one halo from your catalog. It
# is a list of six values:
#    ExampleHalo[0] = snapshot number
#    ExampleHalo[1 - 3] = (X, Y, Z) [Mpc/h]
#    ExampleHalo[4] = Radius [Mpc/h]
#    ExampleHalo[5] = Unbound mass contained within radius [Msun/h]
# The radius used for elements 4 and 5 can correspond to any definition or
# no definition at all, but be careful: most halo catalogs report *bound*,
# not *unbound* masses.
# Report position and radius to the highest accuracy that you know.
# ExampleHalo = 100, 4.68299, 100.552, 80.9536, 2.68893, 1.446e+15
`
}

func (config *CheckConfig) ReadConfig(fname string, flags []string) error {

	vars := parse.NewConfigVars("check.config")

	vars.Float(&config.h0, "H0", -1)
	vars.Float(&config.omegaM, "OmegaM", -1)
	vars.Float(&config.omegaL, "OmegaL", -1)
	vars.Float(&config.boxWidth, "BoxWidth", -1)
	vars.Floats(&config.particleMasses, "ParticleMasses", []float64{})
	vars.Int(&config.particleCount, "ParticleCount", -1)
	vars.Floats(&config.exampleHalo, "ExampleHalo", []float64{})

	if fname == "" {
		if len(flags) == 0 { return nil }
		if err := parse.ReadFlags(flags, vars); err != nil { return err }
	} else {
		if err := parse.ReadConfig(fname, vars); err != nil { return err }
		if err := parse.ReadFlags(flags, vars); err != nil { return err }
	}

	if len(config.exampleHalo) != 0 && len(config.exampleHalo) != 6 {
		return fmt.Errorf(
			"ExampleHalo field must have 6 entries, not %d",
			len(config.exampleHalo),
		)
	}

	return nil
}

func (config *CheckConfig) Run(
	gConfig *GlobalConfig, e *env.Environment, stdin []byte,
) ([]string, error) {

	failedTests := []string{}

	buf, err := getVectorBuffer(
		e.ParticleCatalog(int(gConfig.HSnapMax), 0), gConfig,
	)

	hds, fnames, err := memo.ReadHeaders(int(gConfig.HSnapMax), buf, e)
	if err != nil { return nil, err }
	hd, fname := hds[0], fnames[0]

	failedTests = headerChecks(hd, config, failedTests)
	failedTests, err = particleChecks(buf, fname, config, failedTests)
	if err != nil { return nil, err }
	failedTests, err = haloChecks(hd, buf, config, failedTests, e)
	if err != nil { return nil, err }

	return nil, nil

	if len(failedTests) > 0 {
		if len(failedTests) == 1 {
			fmt.Println("Sanity check failed:")
		} else {
			fmt.Println("Sanity checks failed:")
		}

		for _, test := range failedTests {
			 fmt.Println(test)
		}
		os.Exit(1)
	}

	return nil, nil
}

func checkAlmostEq(x, y float64) bool {
	delta := y / 10
	return math.Abs(x - y) < delta
}

func headerChecks(
	hd io.Header, config *CheckConfig, failedTests []string,
) []string {
	if config.h0 > 0 && !checkAlmostEq(config.h0, hd.Cosmo.H100 * 100) {
		msg := fmt.Sprintf(
			"H0 value in check.config is %g, but read H0 value is %g.",
			config.h0, hd.Cosmo.H100 * 100,
		)
		failedTests = append(failedTests, msg)
	}

	if config.omegaM > 0 && !checkAlmostEq(config.omegaM, hd.Cosmo.OmegaM) {
		msg := fmt.Sprintf(
			"OmegaM value in check.config is %g, but read OmegaM value is %g.",
			config.omegaM, hd.Cosmo.OmegaM,
		)
		failedTests = append(failedTests, msg)
	}

	if config.omegaL > 0 && !checkAlmostEq(config.omegaL, hd.Cosmo.OmegaL) {
		msg := fmt.Sprintf(
			"OmegaL value in check.config is %g, but read OmegaL value is %g.",
			config.omegaL, hd.Cosmo.OmegaL,
		)
		failedTests = append(failedTests, msg)
	}

	if config.boxWidth > 0 && !checkAlmostEq(config.boxWidth, hd.TotalWidth) {
		msg := fmt.Sprintf(
			"BoxWidth value in check.config is %g, but read " +
			"BoxWidth value is %g.",  config.boxWidth, hd.TotalWidth,
		)
		failedTests = append(failedTests, msg)
	}

	return failedTests
}

func particleChecks(
	buf io.VectorBuffer, fname string,
	config *CheckConfig, failedTests []string,
) ([]string, error) {

	if len(config.particleMasses) > 0 {

		_, _, ms, _, err := buf.Read(fname)
		if err != nil {
			return failedTests, err
		}

		failureMass := -1.0

		for _, m := range ms {
			found := false
			for _, mm := range config.particleMasses {
				found = found || checkAlmostEq(float64(m), mm)
			}

			if !found {
				failureMass = float64(m)
				break
			}
		}

		if failureMass != -1 {
			msg := fmt.Sprintf(
				"Allowed masses in check.config are %g, but a particle " +
				"with mass %g was found in a paritcle snapshot.",
				config.particleMasses, failureMass,
			)

			failedTests = append(failedTests, msg)
		}
	}

	if config.particleCount > 0 {
		n, err := buf.TotalParticles(fname)
		if err != nil { return failedTests, err }
		if int64(n) != config.particleCount {
			msg := fmt.Sprintf(
				"ParticleCount value in check.config is %d, but read value " +
				"is %d.", config.particleCount, n,
			)
			failedTests = append(failedTests, msg)
		}
	}

	return failedTests, nil
}

func haloChecks(
	hd io.Header, buf io.VectorBuffer,
	config *CheckConfig, failedTests []string,
	e *env.Environment,
) ([]string, error) {
	if len(config.exampleHalo) == 0 { return failedTests, nil }

	snap := int(config.exampleHalo[0])
	hx,hy,hz:=config.exampleHalo[1],config.exampleHalo[2],config.exampleHalo[3]
	hr, hm := config.exampleHalo[4], config.exampleHalo[5]

	snapCoords := [][]float64{
		[]float64{hx}, []float64{hy}, []float64{hz}, []float64{hr},
	}
	s := geom.Sphere{
		C: [3]float32{float32(hx), float32(hy), float32(hz)},
		R: float32(hr),
	}

	/* The following code is identical to the prof mainloop. */

	hds, files, err := memo.ReadHeaders(snap, buf, e)
	if err != nil {
		return nil, err
	}
	hBounds, err := boundingSpheres(snapCoords, &hds[0], e)
	if err != nil {
		return nil, err
	}
	_, intrIdxs := binSphereIntersections(hds, hBounds)

	haloMass := 0.0

	for i := range hds {
		if len(intrIdxs[i]) == 0 {
			continue
		}

		xs, _, ms, _, err := buf.Read(files[i])
		if err != nil {
			return nil, err
		}

		haloMass += addMass(s, &hds[i], xs, ms)

		buf.Close()
	}

	if !checkAlmostEq(hm, haloMass) {
		msg := fmt.Sprintf(
			"ExampleHalo mass in check.config is %g, but measured " +
			"halo mass is %g.", hm, haloMass,
		)
		failedTests = append(failedTests, msg)
	}

	return failedTests, nil
}

func addMass(
	s geom.Sphere, hd *io.Header, xs [][3]float32, ms []float32,
) float64 {
	rMax2 := s.R*s.R

	x0, y0, z0 := s.C[0], s.C[1], s.C[2]
	tw2 := float32(hd.TotalWidth) / 2

	m := 0.0

	min, max := float32(1e6), float32(-1e6)

	for i, vec := range xs {
		x, y, z := vec[0], vec[1], vec[2]
		dx, dy, dz := x - x0, y - y0, z - z0
		dx = wrap(dx, tw2)
		dy = wrap(dy, tw2)
		dz = wrap(dz, tw2)

		if x > max { max = x }
		if x < min { min = x }

		r2 := dx*dx + dy*dy + dz*dz
		if  r2 <= rMax2 {
			m += float64(ms[i])
		}
	}

	return m
}
<commit_msg>Fixed bug in check reporting.<commit_after>package cmd

import (
	"fmt"
	"math"
	"os"

	"github.com/phil-mansfield/shellfish/io"
	"github.com/phil-mansfield/shellfish/parse"
	"github.com/phil-mansfield/shellfish/cmd/env"
	"github.com/phil-mansfield/shellfish/cmd/memo"
	"github.com/phil-mansfield/shellfish/los/geom"
)

type CheckConfig struct {
	h0, omegaM, omegaL float64
	boxWidth float64
	particleMasses []float64
	particleCount int64
	exampleHalo []float64
}

var _ Mode = &CheckConfig{}

func (config *CheckConfig) ExampleConfig() string {
	return `[check.config]

# All fields are optional. Unless otherwise stated, checks are very rough and
# only done to 10%, so  report float values to at least two decimal places.

# H0 = 70
# OmegaM = 0.30
# OmegaL = 0.70

# BoxWidth is in units of Mpc/h (not your code units).
# BoxWidth = 125

# The total number of particles in the simulation box. This is checked exactly.
# ParticleCount = 1073741824

# ParticleMasses A list of all the valid DM particle masses in Msun/h (not your
# code units).
# ParticleMasses = 1.7e7, 1.8e8, 1.1e9

# ExampleHalo is a the location and size of one halo from your catalog. It
# is a list of six values:
#    ExampleHalo[0] = snapshot number
#    ExampleHalo[1 - 3] = (X, Y, Z) [Mpc/h]
#    ExampleHalo[4] = Radius [Mpc/h]
#    ExampleHalo[5] = Unbound mass contained within radius [Msun/h]
# The radius used for elements 4 and 5 can correspond to any definition or
# no definition at all, but be careful: most halo catalogs report *bound*,
# not *unbound* masses.
# Report position and radius to the highest accuracy that you know.
# ExampleHalo = 100, 4.68299, 100.552, 80.9536, 2.68893, 1.446e+15
`
}

func (config *CheckConfig) ReadConfig(fname string, flags []string) error {

	vars := parse.NewConfigVars("check.config")

	vars.Float(&config.h0, "H0", -1)
	vars.Float(&config.omegaM, "OmegaM", -1)
	vars.Float(&config.omegaL, "OmegaL", -1)
	vars.Float(&config.boxWidth, "BoxWidth", -1)
	vars.Floats(&config.particleMasses, "ParticleMasses", []float64{})
	vars.Int(&config.particleCount, "ParticleCount", -1)
	vars.Floats(&config.exampleHalo, "ExampleHalo", []float64{})

	if fname == "" {
		if len(flags) == 0 { return nil }
		if err := parse.ReadFlags(flags, vars); err != nil { return err }
	} else {
		if err := parse.ReadConfig(fname, vars); err != nil { return err }
		if err := parse.ReadFlags(flags, vars); err != nil { return err }
	}

	if len(config.exampleHalo) != 0 && len(config.exampleHalo) != 6 {
		return fmt.Errorf(
			"ExampleHalo field must have 6 entries, not %d",
			len(config.exampleHalo),
		)
	}

	return nil
}

func (config *CheckConfig) Run(
	gConfig *GlobalConfig, e *env.Environment, stdin []byte,
) ([]string, error) {

	failedTests := []string{}

	buf, err := getVectorBuffer(
		e.ParticleCatalog(int(gConfig.HSnapMax), 0), gConfig,
	)

	hds, fnames, err := memo.ReadHeaders(int(gConfig.HSnapMax), buf, e)
	if err != nil { return nil, err }
	hd, fname := hds[0], fnames[0]

	failedTests = headerChecks(hd, config, failedTests)
	failedTests, err = particleChecks(buf, fname, config, failedTests)
	if err != nil { return nil, err }
	failedTests, err = haloChecks(hd, buf, config, failedTests, e)
	if err != nil { return nil, err }

	if len(failedTests) > 0 {
		if len(failedTests) == 1 {
			fmt.Println("Sanity check failed:")
		} else {
			fmt.Println("Sanity checks failed:")
		}

		for _, test := range failedTests {
			 fmt.Println(test)
		}
		os.Exit(1)
	}

	return nil, nil
}

func checkAlmostEq(x, y float64) bool {
	delta := y / 10
	return math.Abs(x - y) < delta
}

func headerChecks(
	hd io.Header, config *CheckConfig, failedTests []string,
) []string {
	if config.h0 > 0 && !checkAlmostEq(config.h0, hd.Cosmo.H100 * 100) {
		msg := fmt.Sprintf(
			"H0 value in check.config is %g, but read H0 value is %g.",
			config.h0, hd.Cosmo.H100 * 100,
		)
		failedTests = append(failedTests, msg)
	}

	if config.omegaM > 0 && !checkAlmostEq(config.omegaM, hd.Cosmo.OmegaM) {
		msg := fmt.Sprintf(
			"OmegaM value in check.config is %g, but read OmegaM value is %g.",
			config.omegaM, hd.Cosmo.OmegaM,
		)
		failedTests = append(failedTests, msg)
	}

	if config.omegaL > 0 && !checkAlmostEq(config.omegaL, hd.Cosmo.OmegaL) {
		msg := fmt.Sprintf(
			"OmegaL value in check.config is %g, but read OmegaL value is %g.",
			config.omegaL, hd.Cosmo.OmegaL,
		)
		failedTests = append(failedTests, msg)
	}

	if config.boxWidth > 0 && !checkAlmostEq(config.boxWidth, hd.TotalWidth) {
		msg := fmt.Sprintf(
			"BoxWidth value in check.config is %g, but read " +
			"BoxWidth value is %g.",  config.boxWidth, hd.TotalWidth,
		)
		failedTests = append(failedTests, msg)
	}

	return failedTests
}

func particleChecks(
	buf io.VectorBuffer, fname string,
	config *CheckConfig, failedTests []string,
) ([]string, error) {

	if len(config.particleMasses) > 0 {

		_, _, ms, _, err := buf.Read(fname)
		if err != nil {
			return failedTests, err
		}

		failureMass := -1.0

		for _, m := range ms {
			found := false
			for _, mm := range config.particleMasses {
				found = found || checkAlmostEq(float64(m), mm)
			}

			if !found {
				failureMass = float64(m)
				break
			}
		}

		if failureMass != -1 {
			msg := fmt.Sprintf(
				"Allowed masses in check.config are %g, but a particle " +
				"with mass %g was found in a paritcle snapshot.",
				config.particleMasses, failureMass,
			)

			failedTests = append(failedTests, msg)
		}

		buf.Close()
	}

	if config.particleCount > 0 {
		n, err := buf.TotalParticles(fname)
		if err != nil { return failedTests, err }
		if int64(n) != config.particleCount {
			msg := fmt.Sprintf(
				"ParticleCount value in check.config is %d, but read value " +
				"is %d.", config.particleCount, n,
			)
			failedTests = append(failedTests, msg)
		}
	}

	return failedTests, nil
}

func haloChecks(
	hd io.Header, buf io.VectorBuffer,
	config *CheckConfig, failedTests []string,
	e *env.Environment,
) ([]string, error) {
	if len(config.exampleHalo) == 0 { return failedTests, nil }

	snap := int(config.exampleHalo[0])
	hx,hy,hz:=config.exampleHalo[1],config.exampleHalo[2],config.exampleHalo[3]
	hr, hm := config.exampleHalo[4], config.exampleHalo[5]

	snapCoords := [][]float64{
		[]float64{hx}, []float64{hy}, []float64{hz}, []float64{hr},
	}
	s := geom.Sphere{
		C: [3]float32{float32(hx), float32(hy), float32(hz)},
		R: float32(hr),
	}

	/* The following code is identical to the prof mainloop. */

	hds, files, err := memo.ReadHeaders(snap, buf, e)
	if err != nil {
		return nil, err
	}
	hBounds, err := boundingSpheres(snapCoords, &hds[0], e)
	if err != nil {
		return nil, err
	}
	_, intrIdxs := binSphereIntersections(hds, hBounds)

	haloMass := 0.0

	for i := range hds {
		if len(intrIdxs[i]) == 0 {
			continue
		}

		xs, _, ms, _, err := buf.Read(files[i])
		if err != nil {
			return nil, err
		}

		haloMass += addMass(s, &hds[i], xs, ms)

		buf.Close()
	}

	if !checkAlmostEq(hm, haloMass) {
		msg := fmt.Sprintf(
			"ExampleHalo mass in check.config is %g, but measured " +
			"halo mass is %g.", hm, haloMass,
		)
		failedTests = append(failedTests, msg)
	}

	return failedTests, nil
}

func addMass(
	s geom.Sphere, hd *io.Header, xs [][3]float32, ms []float32,
) float64 {
	rMax2 := s.R*s.R

	x0, y0, z0 := s.C[0], s.C[1], s.C[2]
	tw2 := float32(hd.TotalWidth) / 2

	m := 0.0

	min, max := float32(1e6), float32(-1e6)

	for i, vec := range xs {
		x, y, z := vec[0], vec[1], vec[2]
		dx, dy, dz := x - x0, y - y0, z - z0
		dx = wrap(dx, tw2)
		dy = wrap(dy, tw2)
		dz = wrap(dz, tw2)

		if x > max { max = x }
		if x < min { min = x }

		r2 := dx*dx + dy*dy + dz*dz
		if  r2 <= rMax2 {
			m += float64(ms[i])
		}
	}

	return m
}
<|endoftext|>


<commit_before>// Copyright 2014 tsuru 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 cmd

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"runtime"
	"strings"

	"github.com/tsuru/tsuru/exec"
)

var execut exec.Executor

const callbackPage = `<!DOCTYPE html>

	%s

`

const successMarkup = `

Login Successful!

You can close this window now.

`

const errorMarkup = `

Login Failed!

%s

`

func executor() exec.Executor {
	if execut == nil {
		execut = exec.OsExecutor{}
	}
	return execut
}

func port(schemeData map[string]string) string {
	p := schemeData["port"]
	if p != "" {
		return fmt.Sprintf(":%s", p)
	}
	return ":0"
}

func open(url string) error {
	var opts exec.ExecuteOptions
	opts = exec.ExecuteOptions{
		Cmd:  "open",
		Args: []string{url},
	}
	if runtime.GOOS == "linux" {
		opts = exec.ExecuteOptions{
			Cmd:  "xdg-open",
			Args: []string{url},
		}
	}
	return executor().Execute(opts)
}

func convertToken(code, redirectUrl string) (string, error) {
	var token string
	params := map[string]string{"code": code, "redirectUrl": redirectUrl}
	u, err := GetURL("/auth/login")
	if err != nil {
		return token, fmt.Errorf("Error in GetURL: %s", err.Error())
	}
	var buf bytes.Buffer
	err = json.NewEncoder(&buf).Encode(params)
	if err != nil {
		return token, fmt.Errorf("Error encoding params %#v: %s", params, err.Error())
	}
	resp, err := http.Post(u, "application/json", &buf)
	if err != nil {
		return token, fmt.Errorf("Error during login post: %s", err.Error())
	}
	defer resp.Body.Close()
	result, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return token, fmt.Errorf("Error reading body: %s", err.Error())
	}
	data := make(map[string]interface{})
	err = json.Unmarshal(result, &data)
	if err != nil {
		return token, fmt.Errorf("Error parsing response: %s - %s", result, err.Error())
	}
	return data["token"].(string), nil
}

func callback(redirectUrl string, finish chan bool) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		defer func() {
			finish <- true
		}()
		var page string
		token, err := convertToken(r.URL.Query().Get("code"), redirectUrl)
		if err == nil {
			writeToken(token)
			page = fmt.Sprintf(callbackPage, successMarkup)
		} else {
			msg := fmt.Sprintf(errorMarkup, err.Error())
			page = fmt.Sprintf(callbackPage, msg)
		}
		w.Header().Add("Content-Type", "text/html")
		w.Write([]byte(page))
	}
}

func (c *login) oauthLogin(context *Context, client *Client) error {
	schemeData := c.getScheme().Data
	finish := make(chan bool)
	l, err := net.Listen("tcp", port(schemeData))
	if err != nil {
		return err
	}
	_, port, err := net.SplitHostPort(l.Addr().String())
	if err != nil {
		return err
	}
	redirectUrl := fmt.Sprintf("http://localhost:%s", port)
	authUrl := strings.Replace(schemeData["authorizeUrl"], "__redirect_url__", redirectUrl, 1)
	http.HandleFunc("/", callback(redirectUrl, finish))
	server := &http.Server{}
	go server.Serve(l)
	open(authUrl)
	<-finish
	return nil
}
<commit_msg>cmd: add message for sucessfull login.<commit_after>// Copyright 2015 tsuru 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 cmd

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"runtime"
	"strings"

	"github.com/tsuru/tsuru/exec"
)

var execut exec.Executor

const callbackPage = `<!DOCTYPE html>

	%s

`

const successMarkup = `

Login Successful!

You can close this window now.

`

const errorMarkup = `

Login Failed!

%s

`

func executor() exec.Executor {
	if execut == nil {
		execut = exec.OsExecutor{}
	}
	return execut
}

func port(schemeData map[string]string) string {
	p := schemeData["port"]
	if p != "" {
		return fmt.Sprintf(":%s", p)
	}
	return ":0"
}

func open(url string) error {
	var opts exec.ExecuteOptions
	opts = exec.ExecuteOptions{
		Cmd:  "open",
		Args: []string{url},
	}
	if runtime.GOOS == "linux" {
		opts = exec.ExecuteOptions{
			Cmd:  "xdg-open",
			Args: []string{url},
		}
	}
	return executor().Execute(opts)
}

func convertToken(code, redirectUrl string) (string, error) {
	var token string
	params := map[string]string{"code": code, "redirectUrl": redirectUrl}
	u, err := GetURL("/auth/login")
	if err != nil {
		return token, fmt.Errorf("Error in GetURL: %s", err.Error())
	}
	var buf bytes.Buffer
	err = json.NewEncoder(&buf).Encode(params)
	if err != nil {
		return token, fmt.Errorf("Error encoding params %#v: %s", params, err.Error())
	}
	resp, err := http.Post(u, "application/json", &buf)
	if err != nil {
		return token, fmt.Errorf("Error during login post: %s", err.Error())
	}
	defer resp.Body.Close()
	result, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return token, fmt.Errorf("Error reading body: %s", err.Error())
	}
	data := make(map[string]interface{})
	err = json.Unmarshal(result, &data)
	if err != nil {
		return token, fmt.Errorf("Error parsing response: %s - %s", result, err.Error())
	}
	return data["token"].(string), nil
}

func callback(redirectUrl string, finish chan bool) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		defer func() {
			finish <- true
		}()
		var page string
		token, err := convertToken(r.URL.Query().Get("code"), redirectUrl)
		if err == nil {
			writeToken(token)
			page = fmt.Sprintf(callbackPage, successMarkup)
		} else {
			msg := fmt.Sprintf(errorMarkup, err.Error())
			page = fmt.Sprintf(callbackPage, msg)
		}
		w.Header().Add("Content-Type", "text/html")
		w.Write([]byte(page))
	}
}

func (c *login) oauthLogin(context *Context, client *Client) error {
	schemeData := c.getScheme().Data
	finish := make(chan bool)
	l, err := net.Listen("tcp", port(schemeData))
	if err != nil {
		return err
	}
	_, port, err := net.SplitHostPort(l.Addr().String())
	if err != nil {
		return err
	}
	redirectUrl := fmt.Sprintf("http://localhost:%s", port)
	authUrl := strings.Replace(schemeData["authorizeUrl"], "__redirect_url__", redirectUrl, 1)
	http.HandleFunc("/", callback(redirectUrl, finish))
	server := &http.Server{}
	go server.Serve(l)
	open(authUrl)
	<-finish
	fmt.Fprintln(context.Stdout, "Successfully logged in!")
	return nil
}
<|endoftext|>


<commit_before>package cmd

import (
	//"fmt"
	"bufio"
	"os"

	"github.com/spf13/cobra"
	"github.com/spf13/viper"
	"github.com/udzura/collector/collectorlib"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/route53"
)

var hostedZone string
var checkID string
var domains []string
var setID string
var weight int64

// watchCmd represents the watch command
var watchCmd = &cobra.Command{
	Use:   "watch",
	Short: "The IP watcher",
	Long: `The IP watcher.
This subcommand is intended to run under "consul watch".`,
	Run: func(cmd *cobra.Command, args []string) {
		os.Exit(runWatcher())
	},
}

func init() {
	RootCmd.AddCommand(watchCmd)

	watchCmd.Flags().StringVarP(&hostedZone, "hosted-zone", "H", "", "Hosted zone to update")
	watchCmd.Flags().StringVarP(&checkID, "check-id", "C", "", "CheckID to use for IP output")
	watchCmd.Flags().StringSliceVarP(&domains, "domain", "D", []string{}, "Full domain name to keep global IPs")
	watchCmd.Flags().StringVarP(&setID, "set-id", "S", "", "SetIdentifier")
	watchCmd.Flags().Int64VarP(&weight, "weight", "W", 0, "Round Robin Weight")

	viper.SetEnvPrefix("collector")
	viper.BindEnv("hosted_zone")
	viper.BindEnv("check_id")
	viper.BindEnv("domain")
	viper.BindEnv("set_id")
	viper.BindEnv("weight")

	cobra.OnInitialize(initEnviron)
}

func initEnviron() {
	if hostedZone == "" {
		hostedZone = viper.GetString("hosted_zone")
	}
	if checkID == "" {
		checkID = viper.GetString("check_id")
	}
	if len(domains) == 0 {
		domains = viper.GetStringSlice("domain")
	}
	if viper.IsSet("set_id") {
		setID = viper.GetString("set_id")
	}
	if viper.IsSet("weight") {
		weight = viper.GetInt64("weight")
	}

}

func runWatcher() int {
	collectorlib.Logger.Infof("Watch called: hostedZone=%s, domains=%v, checkID=%s, setID=%s, weight=%d", hostedZone, domains, checkID, setID, weight)

	reader := bufio.NewReader(os.Stdin)
	_, err := reader.Peek(1)
	if err != nil {
		collectorlib.Logger.Errorln(err.Error())
		return -1
	}
	defer os.Stdin.Close()

	req, err := collectorlib.ParseRequest(reader)
	if err != nil {
		collectorlib.Logger.Errorln(err.Error())
		return -1
	}
	if checkID != "" {
		req.TargetCheckID = checkID
	}

	svc := route53.New(session.New())
	p1 := &route53.ListHostedZonesByNameInput{
		DNSName: aws.String(hostedZone),
	}
	resp, err := svc.ListHostedZonesByName(p1)

	if err != nil {
		collectorlib.Logger.Errorln(err.Error())
		return -1
	}

	var target *route53.HostedZone
	for _, z := range resp.HostedZones {
		if *z.Name == hostedZone+"." {
			target = z
		}
	}
	if target == nil {
		collectorlib.Logger.Errorf("Hosted zone not found. response: %v", resp.HostedZones)
		return -1
	}

	collectorlib.Logger.Infof("get response: %s(%s)", *target.Name, *target.Id)

	domainModels, err := collectorlib.NewDomains(domains)
	if err != nil {
		collectorlib.Logger.Errorln(err.Error())
		return -1
	}

	for _, domain := range domainModels {
		collectorlib.Logger.Infof("handling domain: %s", domain.FQDN)
		ips := req.IPsByTag(domain.Tag)
		if len(ips) == 0 {
			collectorlib.Logger.Warnln("Hey, no Ip included. Skipping for fail-safe.")
			continue
		}
		collectorlib.Logger.Infof("IPs: %v", ips)

		var oldIPs []string
		p2 := &route53.ListResourceRecordSetsInput{
			HostedZoneId:    target.Id,
			StartRecordName: aws.String(domain.FQDN + "."),
			StartRecordType: aws.String("A"),
		}
		r2, err := svc.ListResourceRecordSets(p2)
		if err != nil {
			collectorlib.Logger.Errorln(err.Error())
			return -1
		}
		if len(r2.ResourceRecordSets) > 0 {
			rrset := r2.ResourceRecordSets[0]
			if *rrset.Type == "A" {
				for _, rr := range rrset.ResourceRecords {
					oldIPs = append(oldIPs, *rr.Value)
				}
			}
		}
		collectorlib.Logger.Infof("existing IPs: %v", oldIPs)
		diff := collectorlib.NewDiff(oldIPs, ips)

		if diff.IsChanged() {
			var rrs []*route53.ResourceRecord
			for _, ip := range ips {
				rrs = append(rrs, &route53.ResourceRecord{
					Value: aws.String(ip),
				})
			}
			p3 := &route53.ChangeResourceRecordSetsInput{
				ChangeBatch: &route53.ChangeBatch{
					Changes: []*route53.Change{
						{
							Action: aws.String("UPSERT"),
							ResourceRecordSet: &route53.ResourceRecordSet{
								Name:            aws.String(domain.FQDN + "."),
								Type:            aws.String("A"),
								ResourceRecords: rrs,
								TTL:             aws.Int64(60),
								SetIdentifier: aws.String(setID),
								Weight: aws.Int64(weight),
							},
						},
					},
					Comment: aws.String("Update via collector"),
				},
				HostedZoneId: target.Id,
			}
			r3, err := svc.ChangeResourceRecordSets(p3)
			if err != nil {
				collectorlib.Logger.Errorln(err.Error())
				return -1
			}

			collectorlib.NotifyToSlack(domain.FQDN, diff)
			collectorlib.Logger.Infof("Success: %v", r3.ChangeInfo)
		} else {
			collectorlib.Logger.Infof("No change, skipping.")
		}
	}

	return 0
}
<commit_msg>set weight only specified<commit_after>package cmd

import (
	//"fmt"
	"bufio"
	"os"

	"github.com/spf13/cobra"
	"github.com/spf13/viper"
	"github.com/udzura/collector/collectorlib"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/route53"
)

var hostedZone string
var checkID string
var domains []string
var setID string
var weight int64

// watchCmd represents the watch command
var watchCmd = &cobra.Command{
	Use:   "watch",
	Short: "The IP watcher",
	Long: `The IP watcher.
This subcommand is intended to run under "consul watch".`,
	Run: func(cmd *cobra.Command, args []string) {
		os.Exit(runWatcher())
	},
}

func init() {
	RootCmd.AddCommand(watchCmd)

	watchCmd.Flags().StringVarP(&hostedZone, "hosted-zone", "H", "", "Hosted zone to update")
	watchCmd.Flags().StringVarP(&checkID, "check-id", "C", "", "CheckID to use for IP output")
	watchCmd.Flags().StringSliceVarP(&domains, "domain", "D", []string{}, "Full domain name to keep global IPs")
	watchCmd.Flags().StringVarP(&setID, "set-id", "S", "", "SetIdentifier")
	watchCmd.Flags().Int64VarP(&weight, "weight", "W", 0, "Round Robin Weight")

	viper.SetEnvPrefix("collector")
	viper.BindEnv("hosted_zone")
	viper.BindEnv("check_id")
	viper.BindEnv("domain")
	viper.BindEnv("set_id")
	viper.BindEnv("weight")

	cobra.OnInitialize(initEnviron)
}

func initEnviron() {
	if hostedZone == "" {
		hostedZone = viper.GetString("hosted_zone")
	}
	if checkID == "" {
		checkID = viper.GetString("check_id")
	}
	if len(domains) == 0 {
		domains = viper.GetStringSlice("domain")
	}
	if viper.IsSet("set_id") {
		setID = viper.GetString("set_id")
	}
	if viper.IsSet("weight") {
		weight = viper.GetInt64("weight")
	}

}

func runWatcher() int {
	collectorlib.Logger.Infof("Watch called: hostedZone=%s, domains=%v, checkID=%s, setID=%s, weight=%d", hostedZone, domains, checkID, setID, weight)

	reader := bufio.NewReader(os.Stdin)
	_, err := reader.Peek(1)
	if err != nil {
		collectorlib.Logger.Errorln(err.Error())
		return -1
	}
	defer os.Stdin.Close()

	req, err := collectorlib.ParseRequest(reader)
	if err != nil {
		collectorlib.Logger.Errorln(err.Error())
		return -1
	}
	if checkID != "" {
		req.TargetCheckID = checkID
	}

	svc := route53.New(session.New())
	p1 := &route53.ListHostedZonesByNameInput{
		DNSName: aws.String(hostedZone),
	}
	resp, err := svc.ListHostedZonesByName(p1)

	if err != nil {
		collectorlib.Logger.Errorln(err.Error())
		return -1
	}

	var target *route53.HostedZone
	for _, z := range resp.HostedZones {
		if *z.Name == hostedZone+"." {
			target = z
		}
	}
	if target == nil {
		collectorlib.Logger.Errorf("Hosted zone not found. response: %v", resp.HostedZones)
		return -1
	}

	collectorlib.Logger.Infof("get response: %s(%s)", *target.Name, *target.Id)

	domainModels, err := collectorlib.NewDomains(domains)
	if err != nil {
		collectorlib.Logger.Errorln(err.Error())
		return -1
	}

	for _, domain := range domainModels {
		collectorlib.Logger.Infof("handling domain: %s", domain.FQDN)
		ips := req.IPsByTag(domain.Tag)
		if len(ips) == 0 {
			collectorlib.Logger.Warnln("Hey, no Ip included. Skipping for fail-safe.")
			continue
		}
		collectorlib.Logger.Infof("IPs: %v", ips)

		var oldIPs []string
		p2 := &route53.ListResourceRecordSetsInput{
			HostedZoneId:    target.Id,
			StartRecordName: aws.String(domain.FQDN + "."),
			StartRecordType: aws.String("A"),
		}
		r2, err := svc.ListResourceRecordSets(p2)
		if err != nil {
			collectorlib.Logger.Errorln(err.Error())
			return -1
		}
		if len(r2.ResourceRecordSets) > 0 {
			rrset := r2.ResourceRecordSets[0]
			if *rrset.Type == "A" {
				for _, rr := range rrset.ResourceRecords {
					oldIPs = append(oldIPs, *rr.Value)
				}
			}
		}
		collectorlib.Logger.Infof("existing IPs: %v", oldIPs)
		diff := collectorlib.NewDiff(oldIPs, ips)

		if diff.IsChanged() {
			var rrs []*route53.ResourceRecord
			for _, ip := range ips {
				rrs = append(rrs, &route53.ResourceRecord{
					Value: aws.String(ip),
				})
			}
			rrset := &route53.ResourceRecordSet{
				Name:            aws.String(domain.FQDN + "."),
				Type:            aws.String("A"),
				ResourceRecords: rrs,
				TTL:             aws.Int64(60),
			}
			if setID != "" {
				rrset.SetSetIdentifier(*aws.String(setID))
				rrset.SetWeight(*aws.Int64(weight))
			}

			p3 := &route53.ChangeResourceRecordSetsInput{
				ChangeBatch: &route53.ChangeBatch{
					Changes: []*route53.Change{
						{
							Action: aws.String("UPSERT"),
							ResourceRecordSet: rrset,
						},
					},
					Comment: aws.String("Update via collector"),
				},
				HostedZoneId: target.Id,
			}
			r3, err := svc.ChangeResourceRecordSets(p3)
			if err != nil {
				collectorlib.Logger.Errorln(err.Error())
				return -1
			}

			collectorlib.NotifyToSlack(domain.FQDN, diff)
			collectorlib.Logger.Infof("Success: %v", r3.ChangeInfo)
		} else {
			collectorlib.Logger.Infof("No change, skipping.")
		}
	}

	return 0
}
<|endoftext|>


<commit_before>package code

import (
	"net/http"
)

type (
	Code         int
	TextHTTPCode struct {
		Text     string
		HTTPCode int
	}
	CodeMap map[Code]TextHTTPCode
)

const (
	RequestTimeout   Code = -9 //提交超时
	AbnormalResponse Code = -8 //响应异常
	OperationTimeout Code = -7 //操作超时
	Unsupported      Code = -6 //不支持的操作
	RepeatOperation  Code = -5 //重复操作
	DataNotFound     Code = -4 //数据未找到
	UserNotFound     Code = -3 //用户未找到
	NonPrivileged    Code = -2 //无权限
	Unauthenticated  Code = -1 //未登录
	Failure          Code = 0  //操作失败
	Success          Code = 1  //操作成功
)

// CodeDict 状态码字典
var CodeDict = CodeMap{
	RequestTimeout:   {"RequestTimeout", http.StatusOK},
	AbnormalResponse: {"AbnormalResponse", http.StatusOK},
	OperationTimeout: {"OperationTimeout", http.StatusOK},
	Unsupported:      {"Unsupported", http.StatusOK},
	RepeatOperation:  {"RepeatOperation", http.StatusOK},
	DataNotFound:     {"DataNotFound", http.StatusOK},
	UserNotFound:     {"UserNotFound", http.StatusOK},
	NonPrivileged:    {"NonPrivileged", http.StatusOK},
	Unauthenticated:  {"Unauthenticated", http.StatusOK},
	Failure:          {"Failure", http.StatusOK},
	Success:          {"Success", http.StatusOK},
}

func (c Code) String() string {
	if v, y := CodeDict[c]; y {
		return v.Text
	}
	return `Undefined`
}

// Int 返回int类型的自定义状态码
func (c Code) Int() int {
	return int(c)
}

// HTTPCode 返回HTTP状态码
func (c Code) HTTPCode() int {
	if v, y := CodeDict[c]; y {
		return v.HTTPCode
	}
	return http.StatusOK
}

func (s CodeMap) Get(code int) TextHTTPCode {
	v, _ := s[Code(code)]
	return v
}
<commit_msg>improved<commit_after>package code

import (
	"net/http"
)

type (
	Code         int
	TextHTTPCode struct {
		Text     string
		HTTPCode int
	}
	CodeMap map[Code]TextHTTPCode
)

const (
	BalanceNoEnough  Code = -10 //余额不足
	RequestTimeout   Code = -9  //提交超时
	AbnormalResponse Code = -8  //响应异常
	OperationTimeout Code = -7  //操作超时
	Unsupported      Code = -6  //不支持的操作
	RepeatOperation  Code = -5  //重复操作
	DataNotFound     Code = -4  //数据未找到
	UserNotFound     Code = -3  //用户未找到
	NonPrivileged    Code = -2  //无权限
	Unauthenticated  Code = -1  //未登录
	Failure          Code = 0   //操作失败
	Success          Code = 1   //操作成功
)

// CodeDict 状态码字典
var CodeDict = CodeMap{
	BalanceNoEnough:  {"BalanceNoEnough", http.StatusOK},
	RequestTimeout:   {"RequestTimeout", http.StatusOK},
	AbnormalResponse: {"AbnormalResponse", http.StatusOK},
	OperationTimeout: {"OperationTimeout", http.StatusOK},
	Unsupported:      {"Unsupported", http.StatusOK},
	RepeatOperation:  {"RepeatOperation", http.StatusOK},
	DataNotFound:     {"DataNotFound", http.StatusOK},
	UserNotFound:     {"UserNotFound", http.StatusOK},
	NonPrivileged:    {"NonPrivileged", http.StatusOK},
	Unauthenticated:  {"Unauthenticated", http.StatusOK},
	Failure:          {"Failure", http.StatusOK},
	Success:          {"Success", http.StatusOK},
}

func (c Code) String() string {
	if v, y := CodeDict[c]; y {
		return v.Text
	}
	return `Undefined`
}

// Int 返回int类型的自定义状态码
func (c Code) Int() int {
	return int(c)
}

// HTTPCode 返回HTTP状态码
func (c Code) HTTPCode() int {
	if v, y := CodeDict[c]; y {
		return v.HTTPCode
	}
	return http.StatusOK
}

func (s CodeMap) Get(code int) TextHTTPCode {
	v, _ := s[Code(code)]
	return v
}
<|endoftext|>


<commit_before>package cloudflare

import (
	"encoding/json"
	"net/url"
	"strconv"
	"time"

	"github.com/pkg/errors"
)

// AccessRule represents a firewall access rule.
type AccessRule struct {
	ID            string                  `json:"id,omitempty"`
	Notes         string                  `json:"notes,omitempty"`
	AllowedModes  []string                `json:"allowed_modes,omitempty"`
	Mode          string                  `json:"mode,omitempty"`
	Configuration AccessRuleConfiguration `json:"configuration,omitempty"`
	Scope         AccessRuleScope         `json:"scope,omitempty"`
	CreatedOn     time.Time               `json:"created_on,omitempty"`
	ModifiedOn    time.Time               `json:"modified_on,omitempty"`
}

// AccessRuleConfiguration represents the configuration of a firewall
// access rule.
type AccessRuleConfiguration struct {
	Target string `json:"target,omitempty"`
	Value  string `json:"value,omitempty"`
}

// AccessRuleScope represents the scope of a firewall access rule.
type AccessRuleScope struct {
	ID    string `json:"id,omitempty"`
	Email string `json:"email,omitempty"`
	Name  string `json:"name,omitempty"`
	Type  string `json:"type,omitempty"`
}

// AccessRuleResponse represents the response from the firewall access
// rule endpoint.
type AccessRuleResponse struct {
	Result AccessRule `json:"result"`
	Response
	ResultInfo `json:"result_info"`
}

// AccessRuleListResponse represents the response from the list access rules
// endpoint.
type AccessRuleListResponse struct {
	Result []AccessRule `json:"result"`
	Response
	ResultInfo `json:"result_info"`
}

// ListUserAccessRules returns a slice of access rules for the logged-in user.
//
// This takes an AccessRule to allow filtering of the results returned.
//
// API reference: https://api.cloudflare.com/#user-level-firewall-access-rule-list-access-rules
func (api *API) ListUserAccessRules(accessRule AccessRule, page int) (*AccessRuleListResponse, error) {
	return api.listAccessRules("/user", accessRule, page)
}

// CreateUserAccessRule creates a firewall access rule for the logged-in user.
//
// API reference: https://api.cloudflare.com/#user-level-firewall-access-rule-create-access-rule
func (api *API) CreateUserAccessRule(accessRule AccessRule) (*AccessRuleResponse, error) {
	return api.createAccessRule("/user", accessRule)
}

// UpdateUserAccessRule updates a single access rule for the logged-in user &
// given access rule identifier.
//
// API reference: https://api.cloudflare.com/#user-level-firewall-access-rule-update-access-rule
func (api *API) UpdateUserAccessRule(accessRuleID string, accessRule AccessRule) (*AccessRuleResponse, error) {
	return api.updateAccessRule("/user", accessRuleID, accessRule)
}

// DeleteUserAccessRule deletes a single access rule for the logged-in user and
// access rule identifiers.
//
// API reference: https://api.cloudflare.com/#user-level-firewall-access-rule-update-access-rule
func (api *API) DeleteUserAccessRule(accessRuleID string) (*AccessRuleResponse, error) {
	return api.deleteAccessRule("/user", accessRuleID)
}

// ListZoneAccessRules returns a slice of access rules for the given zone
// identifier.
//
// This takes an AccessRule to allow filtering of the results returned.
//
// API reference: https://api.cloudflare.com/#firewall-access-rule-for-a-zone-list-access-rules
func (api *API) ListZoneAccessRules(zoneID string, accessRule AccessRule, page int) (*AccessRuleListResponse, error) {
	return api.listAccessRules("/zones/"+zoneID, accessRule, page)
}

// CreateZoneAccessRule creates a firewall access rule for the given zone
// identifier.
//
// API reference: https://api.cloudflare.com/#firewall-access-rule-for-a-zone-create-access-rule
func (api *API) CreateZoneAccessRule(zoneID string, accessRule AccessRule) (*AccessRuleResponse, error) {
	return api.createAccessRule("/zones/"+zoneID, accessRule)
}

// UpdateZoneAccessRule updates a single access rule for the given zone &
// access rule identifiers.
//
// API reference: https://api.cloudflare.com/#firewall-access-rule-for-a-zone-update-access-rule
func (api *API) UpdateZoneAccessRule(zoneID, accessRuleID string, accessRule AccessRule) (*AccessRuleResponse, error) {
	return api.updateAccessRule("/zones/"+zoneID, accessRuleID, accessRule)
}

// DeleteZoneAccessRule deletes a single access rule for the given zone and
// access rule identifiers.
//
// API reference: https://api.cloudflare.com/#firewall-access-rule-for-a-zone-delete-access-rule
func (api *API) DeleteZoneAccessRule(zoneID, accessRuleID string) (*AccessRuleResponse, error) {
	return api.deleteAccessRule("/zones/"+zoneID, accessRuleID)
}

// ListOrganizationAccessRules returns a slice of access rules for the given
// organization identifier.
//
// This takes an AccessRule to allow filtering of the results returned.
//
// API reference: https://api.cloudflare.com/#organization-level-firewall-access-rule-list-access-rules
func (api *API) ListOrganizationAccessRules(organizationID string, accessRule AccessRule, page int) (*AccessRuleListResponse, error) {
	return api.listAccessRules("/accounts/"+organizationID, accessRule, page)
}

// CreateOrganizationAccessRule creates a firewall access rule for the given
// organization identifier.
//
// API reference: https://api.cloudflare.com/#organization-level-firewall-access-rule-create-access-rule
func (api *API) CreateOrganizationAccessRule(organizationID string, accessRule AccessRule) (*AccessRuleResponse, error) {
	return api.createAccessRule("/accounts/"+organizationID, accessRule)
}

// UpdateOrganizationAccessRule updates a single access rule for the given
// organization & access rule identifiers.
//
// API reference: https://api.cloudflare.com/#organization-level-firewall-access-rule-update-access-rule
func (api *API) UpdateOrganizationAccessRule(organizationID, accessRuleID string, accessRule AccessRule) (*AccessRuleResponse, error) {
	return api.updateAccessRule("/accounts/"+organizationID, accessRuleID, accessRule)
}

// DeleteOrganizationAccessRule deletes a single access rule for the given
// organization and access rule identifiers.
//
// API reference: https://api.cloudflare.com/#organization-level-firewall-access-rule-delete-access-rule
func (api *API) DeleteOrganizationAccessRule(organizationID, accessRuleID string) (*AccessRuleResponse, error) {
	return api.deleteAccessRule("/accounts/"+organizationID, accessRuleID)
}

func (api *API) listAccessRules(prefix string, accessRule AccessRule, page int) (*AccessRuleListResponse, error) {
	// Construct a query string
	v := url.Values{}
	if page <= 0 {
		page = 1
	}
	v.Set("page", strconv.Itoa(page))
	// Request as many rules as possible per page - API max is 100
	v.Set("per_page", "100")
	if accessRule.Notes != "" {
		v.Set("notes", accessRule.Notes)
	}
	if accessRule.Mode != "" {
		v.Set("mode", accessRule.Mode)
	}
	if accessRule.Scope.Type != "" {
		v.Set("scope_type", accessRule.Scope.Type)
	}
	if accessRule.Configuration.Value != "" {
		v.Set("configuration_value", accessRule.Configuration.Value)
	}
	if accessRule.Configuration.Target != "" {
		v.Set("configuration_target", accessRule.Configuration.Target)
	}
	v.Set("page", strconv.Itoa(page))
	query := "?" + v.Encode()

	uri := prefix + "/firewall/access_rules/rules" + query
	res, err := api.makeRequest("GET", uri, nil)
	if err != nil {
		return nil, errors.Wrap(err, errMakeRequestError)
	}

	response := &AccessRuleListResponse{}
	err = json.Unmarshal(res, &response)
	if err != nil {
		return nil, errors.Wrap(err, errUnmarshalError)
	}
	return response, nil
}

func (api *API) createAccessRule(prefix string, accessRule AccessRule) (*AccessRuleResponse, error) {
	uri := prefix + "/firewall/access_rules/rules"
	res, err := api.makeRequest("POST", uri, accessRule)
	if err != nil {
		return nil, errors.Wrap(err, errMakeRequestError)
	}

	response := &AccessRuleResponse{}
	err = json.Unmarshal(res, &response)
	if err != nil {
		return nil, errors.Wrap(err, errUnmarshalError)
	}

	return response, nil
}

func (api *API) updateAccessRule(prefix, accessRuleID string, accessRule AccessRule) (*AccessRuleResponse, error) {
	uri := prefix + "/firewall/access_rules/rules/" + accessRuleID
	res, err := api.makeRequest("PATCH", uri, accessRule)
	if err != nil {
		return nil, errors.Wrap(err, errMakeRequestError)
	}

	response := &AccessRuleResponse{}
	err = json.Unmarshal(res, &response)
	if err != nil {
		return nil, errors.Wrap(err, errUnmarshalError)
	}
	return response, nil
}

func (api *API) deleteAccessRule(prefix, accessRuleID string) (*AccessRuleResponse, error) {
	uri := prefix + "/firewall/access_rules/rules/" + accessRuleID
	res, err := api.makeRequest("DELETE", uri, nil)
	if err != nil {
		return nil, errors.Wrap(err, errMakeRequestError)
	}

	response := &AccessRuleResponse{}
	err = json.Unmarshal(res, &response)
	if err != nil {
		return nil, errors.Wrap(err, errUnmarshalError)
	}

	return response, nil
}
<commit_msg>Add functions for retrieving individual access rules<commit_after>package cloudflare

import (
	"encoding/json"
	"net/url"
	"strconv"
	"time"

	"github.com/pkg/errors"
)

// AccessRule represents a firewall access rule.
type AccessRule struct {
	ID            string                  `json:"id,omitempty"`
	Notes         string                  `json:"notes,omitempty"`
	AllowedModes  []string                `json:"allowed_modes,omitempty"`
	Mode          string                  `json:"mode,omitempty"`
	Configuration AccessRuleConfiguration `json:"configuration,omitempty"`
	Scope         AccessRuleScope         `json:"scope,omitempty"`
	CreatedOn     time.Time               `json:"created_on,omitempty"`
	ModifiedOn    time.Time               `json:"modified_on,omitempty"`
}

// AccessRuleConfiguration represents the configuration of a firewall
// access rule.
type AccessRuleConfiguration struct {
	Target string `json:"target,omitempty"`
	Value  string `json:"value,omitempty"`
}

// AccessRuleScope represents the scope of a firewall access rule.
type AccessRuleScope struct {
	ID    string `json:"id,omitempty"`
	Email string `json:"email,omitempty"`
	Name  string `json:"name,omitempty"`
	Type  string `json:"type,omitempty"`
}

// AccessRuleResponse represents the response from the firewall access
// rule endpoint.
type AccessRuleResponse struct {
	Result AccessRule `json:"result"`
	Response
	ResultInfo `json:"result_info"`
}

// AccessRuleListResponse represents the response from the list access rules
// endpoint.
type AccessRuleListResponse struct {
	Result []AccessRule `json:"result"`
	Response
	ResultInfo `json:"result_info"`
}

// ListUserAccessRules returns a slice of access rules for the logged-in user.
//
// This takes an AccessRule to allow filtering of the results returned.
//
// API reference: https://api.cloudflare.com/#user-level-firewall-access-rule-list-access-rules
func (api *API) ListUserAccessRules(accessRule AccessRule, page int) (*AccessRuleListResponse, error) {
	return api.listAccessRules("/user", accessRule, page)
}

// CreateUserAccessRule creates a firewall access rule for the logged-in user.
//
// API reference: https://api.cloudflare.com/#user-level-firewall-access-rule-create-access-rule
func (api *API) CreateUserAccessRule(accessRule AccessRule) (*AccessRuleResponse, error) {
	return api.createAccessRule("/user", accessRule)
}

// UserAccessRule returns the details of a user's account access rule.
//
// API reference: https://api.cloudflare.com/#user-level-firewall-access-rule-list-access-rules
func (api *API) UserAccessRule(accessRuleID string) (*AccessRuleResponse, error) {
	return api.retrieveAccessRule("/user", accessRuleID)
}

// UpdateUserAccessRule updates a single access rule for the logged-in user &
// given access rule identifier.
//
// API reference: https://api.cloudflare.com/#user-level-firewall-access-rule-update-access-rule
func (api *API) UpdateUserAccessRule(accessRuleID string, accessRule AccessRule) (*AccessRuleResponse, error) {
	return api.updateAccessRule("/user", accessRuleID, accessRule)
}

// DeleteUserAccessRule deletes a single access rule for the logged-in user and
// access rule identifiers.
//
// API reference: https://api.cloudflare.com/#user-level-firewall-access-rule-update-access-rule
func (api *API) DeleteUserAccessRule(accessRuleID string) (*AccessRuleResponse, error) {
	return api.deleteAccessRule("/user", accessRuleID)
}

// ListZoneAccessRules returns a slice of access rules for the given zone
// identifier.
//
// This takes an AccessRule to allow filtering of the results returned.
//
// API reference: https://api.cloudflare.com/#firewall-access-rule-for-a-zone-list-access-rules
func (api *API) ListZoneAccessRules(zoneID string, accessRule AccessRule, page int) (*AccessRuleListResponse, error) {
	return api.listAccessRules("/zones/"+zoneID, accessRule, page)
}

// CreateZoneAccessRule creates a firewall access rule for the given zone
// identifier.
//
// API reference: https://api.cloudflare.com/#firewall-access-rule-for-a-zone-create-access-rule
func (api *API) CreateZoneAccessRule(zoneID string, accessRule AccessRule) (*AccessRuleResponse, error) {
	return api.createAccessRule("/zones/"+zoneID, accessRule)
}

// ZoneAccessRule returns the details of a zone's access rule.
//
// API reference: https://api.cloudflare.com/#firewall-access-rule-for-a-zone-list-access-rules
func (api *API) ZoneAccessRule(zoneID string, accessRuleID string) (*AccessRuleResponse, error) {
	return api.retrieveAccessRule("/zones/"+zoneID, accessRuleID)
}

// UpdateZoneAccessRule updates a single access rule for the given zone &
// access rule identifiers.
//
// API reference: https://api.cloudflare.com/#firewall-access-rule-for-a-zone-update-access-rule
func (api *API) UpdateZoneAccessRule(zoneID, accessRuleID string, accessRule AccessRule) (*AccessRuleResponse, error) {
	return api.updateAccessRule("/zones/"+zoneID, accessRuleID, accessRule)
}

// DeleteZoneAccessRule deletes a single access rule for the given zone and
// access rule identifiers.
//
// API reference: https://api.cloudflare.com/#firewall-access-rule-for-a-zone-delete-access-rule
func (api *API) DeleteZoneAccessRule(zoneID, accessRuleID string) (*AccessRuleResponse, error) {
	return api.deleteAccessRule("/zones/"+zoneID, accessRuleID)
}

// ListOrganizationAccessRules returns a slice of access rules for the given
// organization identifier.
//
// This takes an AccessRule to allow filtering of the results returned.
//
// API reference: https://api.cloudflare.com/#organization-level-firewall-access-rule-list-access-rules
func (api *API) ListOrganizationAccessRules(organizationID string, accessRule AccessRule, page int) (*AccessRuleListResponse, error) {
	return api.listAccessRules("/accounts/"+organizationID, accessRule, page)
}

// CreateOrganizationAccessRule creates a firewall access rule for the given
// organization identifier.
//
// API reference: https://api.cloudflare.com/#organization-level-firewall-access-rule-create-access-rule
func (api *API) CreateOrganizationAccessRule(organizationID string, accessRule AccessRule) (*AccessRuleResponse, error) {
	return api.createAccessRule("/accounts/"+organizationID, accessRule)
}

// OrganizationAccessRule returns the details of an organization's access rule.
//
// API reference: https://api.cloudflare.com/#account-level-firewall-access-rule-access-rule-details
func (api *API) OrganizationAccessRule(organizationID string, accessRuleID string) (*AccessRuleResponse, error) {
	return api.retrieveAccessRule("/accounts/"+organizationID, accessRuleID)
}

// UpdateOrganizationAccessRule updates a single access rule for the given
// organization & access rule identifiers.
//
// API reference: https://api.cloudflare.com/#organization-level-firewall-access-rule-update-access-rule
func (api *API) UpdateOrganizationAccessRule(organizationID, accessRuleID string, accessRule AccessRule) (*AccessRuleResponse, error) {
	return api.updateAccessRule("/accounts/"+organizationID, accessRuleID, accessRule)
}

// DeleteOrganizationAccessRule deletes a single access rule for the given
// organization and access rule identifiers.
//
// API reference: https://api.cloudflare.com/#organization-level-firewall-access-rule-delete-access-rule
func (api *API) DeleteOrganizationAccessRule(organizationID, accessRuleID string) (*AccessRuleResponse, error) {
	return api.deleteAccessRule("/accounts/"+organizationID, accessRuleID)
}

func (api *API) listAccessRules(prefix string, accessRule AccessRule, page int) (*AccessRuleListResponse, error) {
	// Construct a query string
	v := url.Values{}
	if page <= 0 {
		page = 1
	}
	v.Set("page", strconv.Itoa(page))
	// Request as many rules as possible per page - API max is 100
	v.Set("per_page", "100")
	if accessRule.Notes != "" {
		v.Set("notes", accessRule.Notes)
	}
	if accessRule.Mode != "" {
		v.Set("mode", accessRule.Mode)
	}
	if accessRule.Scope.Type != "" {
		v.Set("scope_type", accessRule.Scope.Type)
	}
	if accessRule.Configuration.Value != "" {
		v.Set("configuration_value", accessRule.Configuration.Value)
	}
	if accessRule.Configuration.Target != "" {
		v.Set("configuration_target", accessRule.Configuration.Target)
	}
	v.Set("page", strconv.Itoa(page))
	query := "?" + v.Encode()

	uri := prefix + "/firewall/access_rules/rules" + query
	res, err := api.makeRequest("GET", uri, nil)
	if err != nil {
		return nil, errors.Wrap(err, errMakeRequestError)
	}

	response := &AccessRuleListResponse{}
	err = json.Unmarshal(res, &response)
	if err != nil {
		return nil, errors.Wrap(err, errUnmarshalError)
	}
	return response, nil
}

func (api *API) createAccessRule(prefix string, accessRule AccessRule) (*AccessRuleResponse, error) {
	uri := prefix + "/firewall/access_rules/rules"
	res, err := api.makeRequest("POST", uri, accessRule)
	if err != nil {
		return nil, errors.Wrap(err, errMakeRequestError)
	}

	response := &AccessRuleResponse{}
	err = json.Unmarshal(res, &response)
	if err != nil {
		return nil, errors.Wrap(err, errUnmarshalError)
	}

	return response, nil
}

func (api *API) retrieveAccessRule(prefix, accessRuleID string) (*AccessRuleResponse, error) {
	uri := prefix + "/firewall/access_rules/rules/" + accessRuleID

	res, err := api.makeRequest("GET", uri, nil)

	if err != nil {
		return nil, errors.Wrap(err, errMakeRequestError)
	}

	response := &AccessRuleResponse{}
	err = json.Unmarshal(res, &response)
	if err != nil {
		return nil, errors.Wrap(err, errUnmarshalError)
	}

	return response, nil
}

func (api *API) updateAccessRule(prefix, accessRuleID string, accessRule AccessRule) (*AccessRuleResponse, error) {
	uri := prefix + "/firewall/access_rules/rules/" + accessRuleID
	res, err := api.makeRequest("PATCH", uri, accessRule)
	if err != nil {
		return nil, errors.Wrap(err, errMakeRequestError)
	}

	response := &AccessRuleResponse{}
	err = json.Unmarshal(res, &response)
	if err != nil {
		return nil, errors.Wrap(err, errUnmarshalError)
	}
	return response, nil
}

func (api *API) deleteAccessRule(prefix, accessRuleID string) (*AccessRuleResponse, error) {
	uri := prefix + "/firewall/access_rules/rules/" + accessRuleID
	res, err := api.makeRequest("DELETE", uri, nil)
	if err != nil {
		return nil, errors.Wrap(err, errMakeRequestError)
	}

	response := &AccessRuleResponse{}
	err = json.Unmarshal(res, &response)
	if err != nil {
		return nil, errors.Wrap(err, errUnmarshalError)
	}

	return response, nil
}
<|endoftext|>