<commit_before>package cli

import (
	"bufio"
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"os"
	"strconv"

	"errors"

	"os/exec"

	"io/ioutil"

	"github.com/tucnak/climax"
	"github.com/vsco/dcdr/config"
	"github.com/vsco/dcdr/kv"
	"github.com/vsco/dcdr/models"
	"github.com/vsco/dcdr/ui"
)

var (
	InvalidPercentileFormat = errors.New("invalid -value format. use -value=[0.0-1.0]")
	InvalidBoolFormat       = errors.New("invalid -value format. use -value=[true,false]")
	InvalidFeatureType      = errors.New("invalid -type. use -type=[boolean|percentile]")
)

type Controller struct {
	Config *config.Config
	Client kv.ClientIFace
}

func NewController(cfg *config.Config, kv kv.ClientIFace) (cc *Controller) {
	cc = &Controller{
		Config: cfg,
		Client: kv,
	}

	return
}

func (cc *Controller) Watch(ctx climax.Context) int {
	cmd := exec.Command(
		"consul",
		"watch",
		"-type",
		"keyprefix",
		"-prefix",
		cc.Config.Namespace,
		"cat")

	pr, pw := io.Pipe()
	cmd.Stdout = pw
	b := &bytes.Buffer{}
	cmd.Stderr = b

	scanner := bufio.NewScanner(pr)
	scanner.Split(bufio.ScanLines)

	go func() {
		for scanner.Scan() {
			fts, err := models.KVsToFeatures(scanner.Bytes())

			if err != nil {
				fmt.Println(err)
				return
			}

			bts, err := fts.ToJSON()

			if err != nil {
				log.Println(err)
				return
			}

			err = ioutil.WriteFile(cc.Config.FilePath, bts, 0644)

			if err != nil {
				log.Println(err)
				return
			}

			log.Printf("%s wrote changes to %s\n", cc.Config.Username, cc.Config.FilePath)
		}

		if scanner.Err() != nil {
			fmt.Println(scanner.Err())
		}
	}()

	log.Printf("watching namespace: %s\n", cc.Config.Namespace)

	if err := cmd.Run(); err != nil {
		log.Println(err, b)
	}

	return 0
}

func (cc *Controller) List(ctx climax.Context) int {
	pf, _ := ctx.Get("prefix")
	scope, _ := ctx.Get("scope")

	if pf != "" && scope == "" {
		scope = models.DefaultScope
	}

	features, err := cc.Client.List(pf, scope)

	if err != nil {
		fmt.Println(err)
		return 1
	}

	if len(features) == 0 {
		fmt.Printf("No feature flags found in namespace: %s.\n", cc.Client.Namespace())
		return 1
	}

	ui.New().DrawTable(features)

	return 0
}

func (cc *Controller) ParseContext(ctx climax.Context) (*kv.SetRequest, error) {
	name, _ := ctx.Get("name")
	val, _ := ctx.Get("value")
	typ, _ := ctx.Get("type")
	cmt, _ := ctx.Get("comment")
	scp, _ := ctx.Get("scope")
	ft := models.GetFeatureType(typ)

	var v interface{}
	var err error

	switch ft {
	case models.Percentile:
		v, err = strconv.ParseFloat(val, 64)

		if err != nil {
			return nil, InvalidPercentileFormat
		}
	case models.Boolean:
		v, err = strconv.ParseBool(val)

		if err != nil {
			return nil, InvalidBoolFormat
		}
	case models.Invalid:
		return nil, InvalidFeatureType
	}

	return &kv.SetRequest{
		Key:       name,
		Value:     v,
		Scope:     scp,
		Namespace: cc.Config.Namespace,
		Comment:   cmt,
		User:      cc.Config.Username,
	}, nil
}

func (cc *Controller) Set(ctx climax.Context) int {
	sr, err := cc.ParseContext(ctx)

	if err != nil {
		fmt.Println(err)
		return 1
	}

	err = cc.Client.Set(sr)

	if err != nil {
		fmt.Println(err)
		return 1
	}

	fmt.Printf("set flag '%s'\n", sr.Key)

	return 0
}

func (cc *Controller) Delete(ctx climax.Context) int {
	name, _ := ctx.Get("name")
	scope, _ := ctx.Get("scope")

	if name == "" {
		fmt.Println("name is required")
		return 1
	}

	if scope == "" {
		scope = models.DefaultScope
	}

	err := cc.Client.Delete(name, scope)

	if err != nil {
		fmt.Println(err)
		return 1
	}

	fmt.Printf("deleted flag %s/%s\n", scope, name)

	return 0
}

func (cc *Controller) Init(ctx climax.Context) int {
	_, create := ctx.Get("create")

	err := cc.Client.InitRepo(create)

	if err != nil {
		fmt.Println(err)
		return 1
	}

	if create {
		fmt.Printf("initialized new repo in %s and pushed to %s", cc.Config.Git.RepoPath, cc.Config.Git.RepoURL)
	} else {
		fmt.Printf("cloned %s into %s", cc.Config.Git.RepoURL, cc.Config.Git.RepoPath)
	}

	return 0
}

func (cc *Controller) Import(ctx climax.Context) int {
	bts, err := ioutil.ReadAll(os.Stdin)

	if err != nil {
		fmt.Println(err)
		return 1
	}

	var kvs map[string]interface{}
	err = json.Unmarshal(bts, &kvs)

	if err != nil {
		fmt.Println(err)
		return 1
	}

	scope, _ := ctx.Get("scope")

	if scope == "" {
		scope = models.DefaultScope
	}

	for k, v := range kvs {
		sr := &kv.SetRequest{
			Key:       k,
			Value:     v,
			Namespace: cc.Config.Namespace,
			Scope:     scope,
		}

		err = cc.Client.Set(sr)

		if err != nil {
			fmt.Println(err)
			return 1
		}

		fmt.Printf("set %s to %+v\n", k, v)
	}

	return 1
}
<commit_msg>add newlines to repo messages<commit_after>package cli

import (
	"bufio"
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"os"
	"strconv"

	"errors"

	"os/exec"

	"io/ioutil"

	"github.com/tucnak/climax"
	"github.com/vsco/dcdr/config"
	"github.com/vsco/dcdr/kv"
	"github.com/vsco/dcdr/models"
	"github.com/vsco/dcdr/ui"
)

var (
	InvalidPercentileFormat = errors.New("invalid -value format. use -value=[0.0-1.0]")
	InvalidBoolFormat       = errors.New("invalid -value format. use -value=[true,false]")
	InvalidFeatureType      = errors.New("invalid -type. use -type=[boolean|percentile]")
)

type Controller struct {
	Config *config.Config
	Client kv.ClientIFace
}

func NewController(cfg *config.Config, kv kv.ClientIFace) (cc *Controller) {
	cc = &Controller{
		Config: cfg,
		Client: kv,
	}

	return
}

func (cc *Controller) Watch(ctx climax.Context) int {
	cmd := exec.Command(
		"consul",
		"watch",
		"-type",
		"keyprefix",
		"-prefix",
		cc.Config.Namespace,
		"cat")

	pr, pw := io.Pipe()
	cmd.Stdout = pw
	b := &bytes.Buffer{}
	cmd.Stderr = b

	scanner := bufio.NewScanner(pr)
	scanner.Split(bufio.ScanLines)

	go func() {
		for scanner.Scan() {
			fts, err := models.KVsToFeatures(scanner.Bytes())

			if err != nil {
				fmt.Println(err)
				return
			}

			bts, err := fts.ToJSON()

			if err != nil {
				log.Println(err)
				return
			}

			err = ioutil.WriteFile(cc.Config.FilePath, bts, 0644)

			if err != nil {
				log.Println(err)
				return
			}

			log.Printf("%s wrote changes to %s\n", cc.Config.Username, cc.Config.FilePath)
		}

		if scanner.Err() != nil {
			fmt.Println(scanner.Err())
		}
	}()

	log.Printf("watching namespace: %s\n", cc.Config.Namespace)

	if err := cmd.Run(); err != nil {
		log.Println(err, b)
	}

	return 0
}

func (cc *Controller) List(ctx climax.Context) int {
	pf, _ := ctx.Get("prefix")
	scope, _ := ctx.Get("scope")

	if pf != "" && scope == "" {
		scope = models.DefaultScope
	}

	features, err := cc.Client.List(pf, scope)

	if err != nil {
		fmt.Println(err)
		return 1
	}

	if len(features) == 0 {
		fmt.Printf("No feature flags found in namespace: %s.\n", cc.Client.Namespace())
		return 1
	}

	ui.New().DrawTable(features)

	return 0
}

func (cc *Controller) ParseContext(ctx climax.Context) (*kv.SetRequest, error) {
	name, _ := ctx.Get("name")
	val, _ := ctx.Get("value")
	typ, _ := ctx.Get("type")
	cmt, _ := ctx.Get("comment")
	scp, _ := ctx.Get("scope")
	ft := models.GetFeatureType(typ)

	var v interface{}
	var err error

	switch ft {
	case models.Percentile:
		v, err = strconv.ParseFloat(val, 64)

		if err != nil {
			return nil, InvalidPercentileFormat
		}
	case models.Boolean:
		v, err = strconv.ParseBool(val)

		if err != nil {
			return nil, InvalidBoolFormat
		}
	case models.Invalid:
		return nil, InvalidFeatureType
	}

	return &kv.SetRequest{
		Key:       name,
		Value:     v,
		Scope:     scp,
		Namespace: cc.Config.Namespace,
		Comment:   cmt,
		User:      cc.Config.Username,
	}, nil
}

func (cc *Controller) Set(ctx climax.Context) int {
	sr, err := cc.ParseContext(ctx)

	if err != nil {
		fmt.Println(err)
		return 1
	}

	err = cc.Client.Set(sr)

	if err != nil {
		fmt.Println(err)
		return 1
	}

	fmt.Printf("set flag '%s'\n", sr.Key)

	return 0
}

func (cc *Controller) Delete(ctx climax.Context) int {
	name, _ := ctx.Get("name")
	scope, _ := ctx.Get("scope")

	if name == "" {
		fmt.Println("name is required")
		return 1
	}

	if scope == "" {
		scope = models.DefaultScope
	}

	err := cc.Client.Delete(name, scope)

	if err != nil {
		fmt.Println(err)
		return 1
	}

	fmt.Printf("deleted flag %s/%s\n", scope, name)

	return 0
}

func (cc *Controller) Init(ctx climax.Context) int {
	_, create := ctx.Get("create")

	err := cc.Client.InitRepo(create)

	if err != nil {
		fmt.Println(err)
		return 1
	}

	if create {
		fmt.Printf("initialized new repo in %s and pushed to %s\n", cc.Config.Git.RepoPath, cc.Config.Git.RepoURL)
	} else {
		fmt.Printf("cloned %s into %s\n", cc.Config.Git.RepoURL, cc.Config.Git.RepoPath)
	}

	return 0
}

func (cc *Controller) Import(ctx climax.Context) int {
	bts, err := ioutil.ReadAll(os.Stdin)

	if err != nil {
		fmt.Println(err)
		return 1
	}

	var kvs map[string]interface{}
	err = json.Unmarshal(bts, &kvs)

	if err != nil {
		fmt.Println(err)
		return 1
	}

	scope, _ := ctx.Get("scope")

	if scope == "" {
		scope = models.DefaultScope
	}

	for k, v := range kvs {
		sr := &kv.SetRequest{
			Key:       k,
			Value:     v,
			Namespace: cc.Config.Namespace,
			Scope:     scope,
		}

		err = cc.Client.Set(sr)

		if err != nil {
			fmt.Println(err)
			return 1
		}

		fmt.Printf("set %s to %+v\n", k, v)
	}

	return 1
}
<|endoftext|>


<commit_before>package main

import (
	"log"
	"net"
	"os"
	"os/signal"
	"strconv"

	"github.com/coreos/etcd/pkg/flags"
	"github.com/privacylab/talek/common"
	"github.com/privacylab/talek/libtalek"
	"github.com/privacylab/talek/server"
	"github.com/spf13/pflag"
)

// Starts a talek frontend operating with configuration from talekutil
func main() {
	log.Println("----------------------")
	log.Println("--- Talek Frontend ---")
	log.Println("----------------------")

	configPath := pflag.String("client", "talek.conf", "Talek Client Configuration")
	commonPath := pflag.String("common", "common.conf", "Talek Common Configuration")
	systemPath := pflag.String("server", "server.conf", "Talek Server Configuration")
	verbose := pflag.Bool("verbose", false, "Verbose output")
	err := flags.SetPflagsFromEnv(common.EnvPrefix, pflag.CommandLine)
	if err != nil {
		log.Printf("Error reading environment variables, %v\n", err)
		return
	}
	pflag.Parse()

	config := libtalek.ClientConfigFromFile(*configPath)
	config.Config = common.ConfigFromFile(*commonPath)
	serverConfig := server.ConfigFromFile(*systemPath, config.Config)

	replicas := make([]common.ReplicaInterface, len(config.TrustDomains))
	for i, td := range config.TrustDomains {
		replicas[i] = common.NewReplicaRPC(td.Name, td)
	}

	f := server.NewFrontend("Talek Frontend", serverConfig, replicas)
	f.Verbose = *verbose
	_, port, _ := net.SplitHostPort(config.FrontendAddr)
	pnum, _ := strconv.Atoi(port)
	_ = server.NewNetworkRPC(common.FrontendInterface(f), pnum)

	log.Println("Running.")

	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)
	<-c
	f.Close()
}
<commit_msg>help on error<commit_after>package main

import (
	"log"
	"net"
	"os"
	"os/signal"
	"strconv"

	"github.com/coreos/etcd/pkg/flags"
	"github.com/privacylab/talek/common"
	"github.com/privacylab/talek/libtalek"
	"github.com/privacylab/talek/server"
	"github.com/spf13/pflag"
)

// Starts a talek frontend operating with configuration from talekutil
func main() {
	log.Println("----------------------")
	log.Println("--- Talek Frontend ---")
	log.Println("----------------------")

	configPath := pflag.String("client", "talek.conf", "Talek Client Configuration")
	commonPath := pflag.String("common", "common.conf", "Talek Common Configuration")
	systemPath := pflag.String("server", "server.conf", "Talek Server Configuration")
	verbose := pflag.Bool("verbose", false, "Verbose output")
	err := flags.SetPflagsFromEnv(common.EnvPrefix, pflag.CommandLine)
	if err != nil {
		log.Printf("Error reading environment variables, %v\n", err)
		return
	}
	pflag.Parse()

	config := libtalek.ClientConfigFromFile(*configPath)
	if config == nil {
		pflag.Usage()
		return
	}
	config.Config = common.ConfigFromFile(*commonPath)
	serverConfig := server.ConfigFromFile(*systemPath, config.Config)

	replicas := make([]common.ReplicaInterface, len(config.TrustDomains))
	for i, td := range config.TrustDomains {
		replicas[i] = common.NewReplicaRPC(td.Name, td)
	}

	f := server.NewFrontend("Talek Frontend", serverConfig, replicas)
	f.Verbose = *verbose
	_, port, _ := net.SplitHostPort(config.FrontendAddr)
	pnum, _ := strconv.Atoi(port)
	_ = server.NewNetworkRPC(common.FrontendInterface(f), pnum)

	log.Println("Running.")

	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)
	<-c
	f.Close()
}
<|endoftext|>


<commit_before>package clicommand

import (
	"os"
	"path/filepath"
	"runtime"

	"github.com/buildkite/agent/agent"
	"github.com/buildkite/agent/cliconfig"
	"github.com/buildkite/agent/logger"
	"github.com/codegangsta/cli"
)

var StartDescription = `Usage:

   buildkite-agent start [arguments...]

Description:

   When a job is ready to run it will call the "bootstrap-script"
   and pass it all the environment variables required for the job to run.
   This script is responsible for checking out the code, and running the
   actual build script defined in the project.

   The agent will run any jobs within a PTY (pseudo terminal) if available.

Example:

   $ buildkite-agent start --token xxx`

type AgentStartConfig struct {
	Config                           string   `cli:"config"`
	Token                            string   `cli:"token" validate:"required"`
	Name                             string   `cli:"name"`
	Priority                         string   `cli:"priority"`
	BootstrapScript                  string   `cli:"bootstrap-script" normalize:"filepath" validate:"required,file-exists"`
	BuildPath                        string   `cli:"build-path" normalize:"filepath" validate:"required"`
	HooksPath                        string   `cli:"hooks-path" normalize:"filepath"`
	MetaData                         []string `cli:"meta-data"`
	MetaDataEC2Tags                  bool     `cli:"meta-data-ec2-tags"`
	NoColor                          bool     `cli:"no-color"`
	NoAutoSSHFingerprintVerification bool     `cli:"no-automatic-ssh-fingerprint-verification"`
	NoCommandEval                    bool     `cli:"no-command-eval"`
	NoPTY                            bool     `cli:"no-pty"`
	Endpoint                         string   `cli:"endpoint" validate:"required"`
	Debug                            bool     `cli:"debug"`
	DebugHTTP                        bool     `cli:"debug-http"`
}

func DefaultConfigFilePaths() (paths []string) {
	// Toggle beetwen windows an *nix paths
	if runtime.GOOS == "windows" {
		paths = []string{
			"$USERPROFILE\\AppData\\Local\\BuildkiteAgent\\buildkite-agent.cfg",
		}
	} else {
		paths = []string{
			"$HOME/.buildkite-agent/buildkite-agent.cfg",
			"/usr/local/etc/buildkite-agent/buildkite-agent.cfg",
			"/etc/buildkite-agent/buildkite-agent.cfg",
		}
	}

	// Also check to see if there's a buildkite-agent.cfg in the folder
	// that the binary is running in.
	pathToBinary, err := filepath.Abs(filepath.Dir(os.Args[0]))
	if err == nil {
		pathToRelativeConfig := filepath.Join(pathToBinary, "buildkite-agent.cfg")
		paths = append([]string{pathToRelativeConfig}, paths...)
	}

	return
}

var AgentStartCommand = cli.Command{
	Name:        "start",
	Usage:       "Starts a Buildkite agent",
	Description: StartDescription,
	Flags: []cli.Flag{
		cli.StringFlag{
			Name:   "config",
			Value:  "",
			Usage:  "Path to a configuration file",
			EnvVar: "BUILDKITE_AGENT_CONFIG",
		},
		cli.StringFlag{
			Name:   "token",
			Value:  "",
			Usage:  "Your account agent token",
			EnvVar: "BUILDKITE_AGENT_TOKEN",
		},
		cli.StringFlag{
			Name:   "name",
			Value:  "",
			Usage:  "The name of the agent",
			EnvVar: "BUILDKITE_AGENT_NAME",
		},
		cli.StringFlag{
			Name:   "priority",
			Value:  "",
			Usage:  "The priority of the agent (higher priorities are assigned work first)",
			EnvVar: "BUILDKITE_AGENT_PRIORITY",
		},
		cli.StringSliceFlag{
			Name:   "meta-data",
			Value:  &cli.StringSlice{},
			Usage:  "Meta data for the agent (default is \"queue=default\")",
			EnvVar: "BUILDKITE_AGENT_META_DATA",
		},
		cli.BoolFlag{
			Name:  "meta-data-ec2-tags",
			Usage: "Populate the meta data from the current instances EC2 Tags",
		},
		cli.StringFlag{
			Name:   "bootstrap-script",
			Value:  "",
			Usage:  "Path to the bootstrap script",
			EnvVar: "BUILDKITE_BOOTSTRAP_SCRIPT_PATH",
		},
		cli.StringFlag{
			Name:   "build-path",
			Value:  "",
			Usage:  "Path to where the builds will run from",
			EnvVar: "BUILDKITE_BUILD_PATH",
		},
		cli.StringFlag{
			Name:   "hooks-path",
			Value:  "",
			Usage:  "Directory where the hook scripts are found",
			EnvVar: "BUILDKITE_HOOKS_PATH",
		},
		cli.BoolFlag{
			Name:   "no-pty",
			Usage:  "Do not run jobs within a pseudo terminal",
			EnvVar: "BUILDKITE_NO_PTY",
		},
		cli.BoolFlag{
			Name:   "no-automatic-ssh-fingerprint-verification",
			Usage:  "Don't automatically verify SSH fingerprints",
			EnvVar: "BUILDKITE_NO_AUTOMATIC_SSH_FINGERPRINT_VERIFICATION",
		},
		cli.BoolFlag{
			Name:   "no-command-eval",
			Usage:  "Don't allow this agent to run arbitrary console commands",
			EnvVar: "BUILDKITE_NO_COMMAND_EVAL",
		},
		EndpointFlag,
		NoColorFlag,
		DebugFlag,
		DebugHTTPFlag,
	},
	Action: func(c *cli.Context) {
		// The configuration will be loaded into this struct
		cfg := AgentStartConfig{}

		// Setup the config loader. You'll see that we also path paths to
		// potential config files. The loader will use the first one it finds.
		loader := cliconfig.Loader{
			CLI:                    c,
			Config:                 &cfg,
			DefaultConfigFilePaths: DefaultConfigFilePaths(),
		}

		// Load the configuration
		if err := loader.Load(); err != nil {
			logger.Fatal("%s", err)
		}

		// Setup the any global configuration options
		HandleGlobalFlags(cfg)

		// Force some settings if on Windows (these aren't supported yet)
		if runtime.GOOS == "windows" {
			cfg.NoAutoSSHFingerprintVerification = true
			cfg.NoPTY = true
			cfg.NoCommandEval = true
		}

		// Setup the agent
		pool := agent.AgentPool{
			Token:           cfg.Token,
			Name:            cfg.Name,
			Priority:        cfg.Priority,
			MetaData:        cfg.MetaData,
			MetaDataEC2Tags: cfg.MetaDataEC2Tags,
			Endpoint:        cfg.Endpoint,
			AgentConfiguration: &agent.AgentConfiguration{
				BootstrapScript:                cfg.BootstrapScript,
				BuildPath:                      cfg.BuildPath,
				HooksPath:                      cfg.HooksPath,
				AutoSSHFingerprintVerification: !cfg.NoAutoSSHFingerprintVerification,
				CommandEval:                    !cfg.NoCommandEval,
				RunInPty:                       !cfg.NoPTY,
			},
		}

		// Store the loaded config file path on the pool so we can
		// show it when the agent starts
		if loader.File != nil {
			pool.ConfigFilePath = loader.File.Path
		}

		// Start the agent pool
		if err := pool.Start(); err != nil {
			logger.Fatal("%s", err)
		}
	},
}
<commit_msg>Don't validate the existance of the `bootstrap-script`<commit_after>package clicommand

import (
	"os"
	"path/filepath"
	"runtime"

	"github.com/buildkite/agent/agent"
	"github.com/buildkite/agent/cliconfig"
	"github.com/buildkite/agent/logger"
	"github.com/codegangsta/cli"
)

var StartDescription = `Usage:

   buildkite-agent start [arguments...]

Description:

   When a job is ready to run it will call the "bootstrap-script"
   and pass it all the environment variables required for the job to run.
   This script is responsible for checking out the code, and running the
   actual build script defined in the project.

   The agent will run any jobs within a PTY (pseudo terminal) if available.

Example:

   $ buildkite-agent start --token xxx`

type AgentStartConfig struct {
	Config                           string   `cli:"config"`
	Token                            string   `cli:"token" validate:"required"`
	Name                             string   `cli:"name"`
	Priority                         string   `cli:"priority"`
	BootstrapScript                  string   `cli:"bootstrap-script" normalize:"filepath" validate:"required"`
	BuildPath                        string   `cli:"build-path" normalize:"filepath" validate:"required"`
	HooksPath                        string   `cli:"hooks-path" normalize:"filepath"`
	MetaData                         []string `cli:"meta-data"`
	MetaDataEC2Tags                  bool     `cli:"meta-data-ec2-tags"`
	NoColor                          bool     `cli:"no-color"`
	NoAutoSSHFingerprintVerification bool     `cli:"no-automatic-ssh-fingerprint-verification"`
	NoCommandEval                    bool     `cli:"no-command-eval"`
	NoPTY                            bool     `cli:"no-pty"`
	Endpoint                         string   `cli:"endpoint" validate:"required"`
	Debug                            bool     `cli:"debug"`
	DebugHTTP                        bool     `cli:"debug-http"`
}

func DefaultConfigFilePaths() (paths []string) {
	// Toggle beetwen windows an *nix paths
	if runtime.GOOS == "windows" {
		paths = []string{
			"$USERPROFILE\\AppData\\Local\\BuildkiteAgent\\buildkite-agent.cfg",
		}
	} else {
		paths = []string{
			"$HOME/.buildkite-agent/buildkite-agent.cfg",
			"/usr/local/etc/buildkite-agent/buildkite-agent.cfg",
			"/etc/buildkite-agent/buildkite-agent.cfg",
		}
	}

	// Also check to see if there's a buildkite-agent.cfg in the folder
	// that the binary is running in.
	pathToBinary, err := filepath.Abs(filepath.Dir(os.Args[0]))
	if err == nil {
		pathToRelativeConfig := filepath.Join(pathToBinary, "buildkite-agent.cfg")
		paths = append([]string{pathToRelativeConfig}, paths...)
	}

	return
}

var AgentStartCommand = cli.Command{
	Name:        "start",
	Usage:       "Starts a Buildkite agent",
	Description: StartDescription,
	Flags: []cli.Flag{
		cli.StringFlag{
			Name:   "config",
			Value:  "",
			Usage:  "Path to a configuration file",
			EnvVar: "BUILDKITE_AGENT_CONFIG",
		},
		cli.StringFlag{
			Name:   "token",
			Value:  "",
			Usage:  "Your account agent token",
			EnvVar: "BUILDKITE_AGENT_TOKEN",
		},
		cli.StringFlag{
			Name:   "name",
			Value:  "",
			Usage:  "The name of the agent",
			EnvVar: "BUILDKITE_AGENT_NAME",
		},
		cli.StringFlag{
			Name:   "priority",
			Value:  "",
			Usage:  "The priority of the agent (higher priorities are assigned work first)",
			EnvVar: "BUILDKITE_AGENT_PRIORITY",
		},
		cli.StringSliceFlag{
			Name:   "meta-data",
			Value:  &cli.StringSlice{},
			Usage:  "Meta data for the agent (default is \"queue=default\")",
			EnvVar: "BUILDKITE_AGENT_META_DATA",
		},
		cli.BoolFlag{
			Name:  "meta-data-ec2-tags",
			Usage: "Populate the meta data from the current instances EC2 Tags",
		},
		cli.StringFlag{
			Name:   "bootstrap-script",
			Value:  "",
			Usage:  "Path to the bootstrap script",
			EnvVar: "BUILDKITE_BOOTSTRAP_SCRIPT_PATH",
		},
		cli.StringFlag{
			Name:   "build-path",
			Value:  "",
			Usage:  "Path to where the builds will run from",
			EnvVar: "BUILDKITE_BUILD_PATH",
		},
		cli.StringFlag{
			Name:   "hooks-path",
			Value:  "",
			Usage:  "Directory where the hook scripts are found",
			EnvVar: "BUILDKITE_HOOKS_PATH",
		},
		cli.BoolFlag{
			Name:   "no-pty",
			Usage:  "Do not run jobs within a pseudo terminal",
			EnvVar: "BUILDKITE_NO_PTY",
		},
		cli.BoolFlag{
			Name:   "no-automatic-ssh-fingerprint-verification",
			Usage:  "Don't automatically verify SSH fingerprints",
			EnvVar: "BUILDKITE_NO_AUTOMATIC_SSH_FINGERPRINT_VERIFICATION",
		},
		cli.BoolFlag{
			Name:   "no-command-eval",
			Usage:  "Don't allow this agent to run arbitrary console commands",
			EnvVar: "BUILDKITE_NO_COMMAND_EVAL",
		},
		EndpointFlag,
		NoColorFlag,
		DebugFlag,
		DebugHTTPFlag,
	},
	Action: func(c *cli.Context) {
		// The configuration will be loaded into this struct
		cfg := AgentStartConfig{}

		// Setup the config loader. You'll see that we also path paths to
		// potential config files. The loader will use the first one it finds.
		loader := cliconfig.Loader{
			CLI:                    c,
			Config:                 &cfg,
			DefaultConfigFilePaths: DefaultConfigFilePaths(),
		}

		// Load the configuration
		if err := loader.Load(); err != nil {
			logger.Fatal("%s", err)
		}

		// Setup the any global configuration options
		HandleGlobalFlags(cfg)

		// Force some settings if on Windows (these aren't supported yet)
		if runtime.GOOS == "windows" {
			cfg.NoAutoSSHFingerprintVerification = true
			cfg.NoPTY = true
			cfg.NoCommandEval = true
		}

		// Setup the agent
		pool := agent.AgentPool{
			Token:           cfg.Token,
			Name:            cfg.Name,
			Priority:        cfg.Priority,
			MetaData:        cfg.MetaData,
			MetaDataEC2Tags: cfg.MetaDataEC2Tags,
			Endpoint:        cfg.Endpoint,
			AgentConfiguration: &agent.AgentConfiguration{
				BootstrapScript:                cfg.BootstrapScript,
				BuildPath:                      cfg.BuildPath,
				HooksPath:                      cfg.HooksPath,
				AutoSSHFingerprintVerification: !cfg.NoAutoSSHFingerprintVerification,
				CommandEval:                    !cfg.NoCommandEval,
				RunInPty:                       !cfg.NoPTY,
			},
		}

		// Store the loaded config file path on the pool so we can
		// show it when the agent starts
		if loader.File != nil {
			pool.ConfigFilePath = loader.File.Path
		}

		// Start the agent pool
		if err := pool.Start(); err != nil {
			logger.Fatal("%s", err)
		}
	},
}
<|endoftext|>


<commit_before>package gatekeeper

import (
	"bytes"
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"net/url"
	"os"
	"strconv"
	"strings"
)

var HttpClient = &http.Client{}
var VaultAddress = os.Getenv("VAULT_ADDR")
var GatekeeperAddr = os.Getenv("GATEKEEPER_ADDR")

var ErrNoTaskId = errors.New("No task id provided.")

type VaultError struct {
	Code   int      `json:"-"`
	Errors []string `json:"errors"`
}

func (e VaultError) Error() string {
	return fmt.Sprintf("%d: %s", e.Code, strings.Join(e.Errors, ", "))
}

func init() {
	tr := &http.Transport{
		TLSClientConfig: &tls.Config{},
	}
	capath := os.Getenv("VAULT_CAPATH")
	cacert := os.Getenv("VAULT_CACERT")

	if b, err := strconv.ParseBool(os.Getenv("VAULT_SKIP_VERIFY")); err == nil && b {
		tr.TLSClientConfig.InsecureSkipVerify = true
	}

	if capath != "" || cacert != "" {
		LoadCA := func() (*x509.CertPool, error) {
			if capath != "" {
				return LoadCAPath(capath)
			} else if cacert != "" {
				return LoadCACert(cacert)
			}
			panic("invariant violation")
		}
		if certs, err := LoadCA(); err == nil {
			tr.TLSClientConfig.RootCAs = certs
		} else {
			fmt.Fprintf(os.Stderr, "Gatekeeper: Failed to read client certs. Error: %v\n", err)
		}
	}
	HttpClient = &http.Client{Transport: tr}
}

func RequestVaultToken(taskId string) (string, error) {
	if taskId == "" {
		return "", ErrNoTaskId
	}
	var gkPath string
	var vaultPath string
	if u, err := url.Parse(GatekeeperAddr); err == nil {
		u.Path = "/token"
		gkPath = u.String()
	} else {
		return "", err
	}

	if u, err := url.Parse(VaultAddress); err == nil {
		u.Path = "/v1/cubbyhole/vault-token"
		vaultPath = u.String()
	} else {
		return "", err
	}

	gkb := struct {
		TaskId string `json:"task_id"`
	}{taskId}
	payload, _ := json.Marshal(gkb)

	var gkResp struct {
		Status string `json:"status"`
		Ok     bool   `json:"ok"`
		Error  string `json:"error"`
		Token  string `json:"token"`
	}

	if resp, err := HttpClient.Post(gkPath, "application/json", bytes.NewReader(payload)); err == nil {
		decoder := json.NewDecoder(resp.Body)
		if err := decoder.Decode(&gkResp); err == nil {
			if gkResp.Ok {
				req, _ := http.NewRequest("GET", vaultPath, nil)
				req.Header.Add("X-Vault-Token", gkResp.Token)
				if resp, err := HttpClient.Do(req); err == nil {
					decoder := json.NewDecoder(resp.Body)
					vaultResp := struct {
						Data struct {
							Token string `json:"token"`
						} `json:"data"`
					}{}
					if resp.StatusCode == 200 {
						if err := decoder.Decode(&vaultResp); err == nil {
							return vaultResp.Data.Token, nil
						} else {
							return "", err
						}
					} else {
						var e VaultError
						e.Code = resp.StatusCode
						if err := decoder.Decode(&e); err == nil {
							return "", e
						} else {
							e.Errors = []string{"communication error."}
							return "", e
						}
					}
				} else {
					return "", err
				}
			} else {
				return "", errors.New(gkResp.Error)
			}
		} else {
			return "", err
		}
	} else {
		return "", err
	}
}

func EnvRequestVaultToken() (string, error) {
	return RequestVaultToken(os.Getenv("MESOS_TASK_ID"))
}
<commit_msg>Refactor RequestVaultToken for readability, to reduce complexity, and to be more idiomatic go<commit_after>package gatekeeper

import (
	"bytes"
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"net/url"
	"os"
	"strconv"
	"strings"
)

var HttpClient = &http.Client{}
var VaultAddress = os.Getenv("VAULT_ADDR")
var GatekeeperAddr = os.Getenv("GATEKEEPER_ADDR")

var ErrNoTaskId = errors.New("No task id provided.")

type VaultError struct {
	Code   int      `json:"-"`
	Errors []string `json:"errors"`
}

func (e VaultError) Error() string {
	return fmt.Sprintf("%d: %s", e.Code, strings.Join(e.Errors, ", "))
}

func init() {
	tr := &http.Transport{
		TLSClientConfig: &tls.Config{},
	}
	capath := os.Getenv("VAULT_CAPATH")
	cacert := os.Getenv("VAULT_CACERT")

	if b, err := strconv.ParseBool(os.Getenv("VAULT_SKIP_VERIFY")); err == nil && b {
		tr.TLSClientConfig.InsecureSkipVerify = true
	}

	if capath != "" || cacert != "" {
		LoadCA := func() (*x509.CertPool, error) {
			if capath != "" {
				return LoadCAPath(capath)
			} else if cacert != "" {
				return LoadCACert(cacert)
			}
			panic("invariant violation")
		}
		if certs, err := LoadCA(); err == nil {
			tr.TLSClientConfig.RootCAs = certs
		} else {
			fmt.Fprintf(os.Stderr, "Gatekeeper: Failed to read client certs. Error: %v\n", err)
		}
	}
	HttpClient = &http.Client{Transport: tr}
}

func RequestVaultToken(taskId string) (string, error) {
	if taskId == "" {
		return "", ErrNoTaskId
	}

	gkAddr, err := url.Parse(GatekeeperAddr)
	if err != nil {
		return "", err
	}
	gkAddr.Path = "/token"

	gkTaskID := struct {
		TaskId string `json:"task_id"`
	}{taskId}

	gkReq, err := json.Marshal(gkTaskID)
	if err != nil {
		return "", err
	}

	gkResp, err := HttpClient.Post(gkAddr.String(), "application/json", bytes.NewReader(gkReq))
	if err != nil {
		return "", err
	}
	defer gkResp.Body.Close()

	var gkTokResp struct {
		OK     bool   `json:"ok"`
		Token  string `json:"token"`
		Status string `json:"status"`
		Error  string `json:"error"`
	}

	if err := json.NewDecoder(gkResp.Body).Decode(&gkTokResp); err != nil {
		return "", err
	}

	if !gkTokResp.OK {
		return "", errors.New(gkTokResp.Error)
	}

	vaultAddr, err := url.Parse(VaultAddress)
	if err != nil {
		return "", err
	}
	vaultAddr.Path = "/v1/cubbyhole/vault-token"

	req, err := http.NewRequest("GET", vaultAddr.String(), nil)
	if err != nil {
		return "", err
	}
	req.Header.Add("X-Vault-Token", gkTokResp.Token)

	vaultResp, err := HttpClient.Do(req)
	if err != nil {
		return "", err
	}
	defer vaultResp.Body.Close()

	bodyDecoder := json.NewDecoder(vaultResp.Body)

	if vaultResp.StatusCode != 200 {
		var vaultErr VaultError
		vaultErr.Code = vaultResp.StatusCode
		if err := bodyDecoder.Decode(&vaultErr); err != nil {
			vaultErr.Errors = []string{"communication error."}
			return "", err
		}

		return "", vaultErr
	}

	vaultSecret := struct {
		Data struct {
			Token string `json:"token"`
		} `json:"data"`
	}{}

	if err := bodyDecoder.Decode(&vaultSecret); err != nil {
		return "", err
	}

	return vaultSecret.Data.Token, nil
}

func EnvRequestVaultToken() (string, error) {
	return RequestVaultToken(os.Getenv("MESOS_TASK_ID"))
}
<|endoftext|>


<commit_before>/**********************************************************\
|                                                          |
|                          hprose                          |
|                                                          |
| Official WebSite: http://www.hprose.com/                 |
|                   http://www.hprose.org/                 |
|                                                          |
\**********************************************************/
/**********************************************************\
 *                                                        *
 * hprose/http_client.go                                  *
 *                                                        *
 * hprose http client for Go.                             *
 *                                                        *
 * LastModified: Feb 8, 2015                              *
 * Author: Ma Bingyao <andot@hprose.com>                  *
 *                                                        *
\**********************************************************/

package hprose

import (
	"crypto/tls"
	"io"
	"io/ioutil"
	"net/http"
	"net/http/cookiejar"
	"net/url"
)

var cookieJar, _ = cookiejar.New(nil)

var DisableGlobalCookie = false

type HttpClient struct {
	*BaseClient
}

type HttpTransporter struct {
	*http.Client
}

func NewHttpClient(uri string) Client {
	client := &HttpClient{NewBaseClient(newHttpTransporter())}
	client.Client = client
	client.SetUri(uri)
	client.SetKeepAlive(true)
	return client
}

func (client *HttpClient) Close() {}

func (client *HttpClient) SetUri(uri string) {
	if u, err := url.Parse(uri); err == nil {
		if u.Scheme != "http" && u.Scheme != "https" {
			panic("This client desn't support " + u.Scheme + " scheme.")
		}
		if u.Scheme == "https" {
			client.SetTLSClientConfig(&tls.Config{InsecureSkipVerify: true})
		}
	}
	client.BaseClient.SetUri(uri)
}

func (client *HttpClient) Http() *http.Client {
	return client.Transporter.(*HttpTransporter).Client
}

func (client *HttpClient) transport() *http.Transport {
	return client.Http().Transport.(*http.Transport)
}

func (client *HttpClient) TLSClientConfig() *tls.Config {
	return client.transport().TLSClientConfig
}

func (client *HttpClient) SetTLSClientConfig(config *tls.Config) {
	client.transport().TLSClientConfig = config
}

func (client *HttpClient) KeepAlive() bool {
	return !client.transport().DisableKeepAlives
}

func (client *HttpClient) SetKeepAlive(enable bool) {
	client.transport().DisableKeepAlives = !enable
}

func (client *HttpClient) Compression() bool {
	return !client.transport().DisableCompression
}

func (client *HttpClient) SetCompression(enable bool) {
	client.transport().DisableCompression = !enable
}

func (client *HttpClient) MaxIdleConnsPerHost() int {
	return client.transport().MaxIdleConnsPerHost
}

func (client *HttpClient) SetMaxIdleConnsPerHost(value int) {
	client.transport().MaxIdleConnsPerHost = value
}

func newHttpTransporter() *HttpTransporter {
	tr := &http.Transport{
		DisableCompression:  true,
		DisableKeepAlives:   false,
		MaxIdleConnsPerHost: 4}
	jar := cookieJar
	if DisableGlobalCookie {
		jar, _ = cookiejar.New(nil)
	}
	return &HttpTransporter{&http.Client{Jar: jar, Transport: tr}}
}

func (h *HttpTransporter) readAll(response *http.Response) (data []byte, err error) {
	if response.ContentLength > 0 {
		data = make([]byte, response.ContentLength)
		_, err = io.ReadFull(response.Body, data)
		return data, err
	}
	if response.ContentLength < 0 {
		return ioutil.ReadAll(response.Body)
	}
	return make([]byte, 0), nil
}

func (h *HttpTransporter) SendAndReceive(uri string, data []byte) ([]byte, error) {
	req, err := http.NewRequest("POST", uri, NewBytesReader(data))
	if err != nil {
		return nil, err
	}
	req.ContentLength = int64(len(data))
	req.Header.Set("Content-Type", "application/hprose")
	resp, err := h.Do(req)
	if err != nil {
		return nil, err
	}
	data, err = ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}
	return data, resp.Body.Close()
}
<commit_msg>Added Header method for HttpClient<commit_after>/**********************************************************\
|                                                          |
|                          hprose                          |
|                                                          |
| Official WebSite: http://www.hprose.com/                 |
|                   http://www.hprose.org/                 |
|                                                          |
\**********************************************************/
/**********************************************************\
 *                                                        *
 * hprose/http_client.go                                  *
 *                                                        *
 * hprose http client for Go.                             *
 *                                                        *
 * LastModified: May 13, 2015                             *
 * Author: Ma Bingyao <andot@hprose.com>                  *
 *                                                        *
\**********************************************************/

package hprose

import (
	"crypto/tls"
	"io"
	"io/ioutil"
	"net/http"
	"net/http/cookiejar"
	"net/url"
)

var cookieJar, _ = cookiejar.New(nil)

var DisableGlobalCookie = false

type HttpClient struct {
	*BaseClient
}

type HttpTransporter struct {
	*http.Client
	Header *http.Header
}

func NewHttpClient(uri string) Client {
	client := &HttpClient{NewBaseClient(newHttpTransporter())}
	client.Client = client
	client.SetUri(uri)
	client.SetKeepAlive(true)
	return client
}

func (client *HttpClient) Close() {}

func (client *HttpClient) SetUri(uri string) {
	if u, err := url.Parse(uri); err == nil {
		if u.Scheme != "http" && u.Scheme != "https" {
			panic("This client desn't support " + u.Scheme + " scheme.")
		}
		if u.Scheme == "https" {
			client.SetTLSClientConfig(&tls.Config{InsecureSkipVerify: true})
		}
	}
	client.BaseClient.SetUri(uri)
}

func (client *HttpClient) Http() *http.Client {
	return client.Transporter.(*HttpTransporter).Client
}

func (client *HttpClient) Header() *http.Header {
	return client.Transporter.(*HttpTransporter).Header
}

func (client *HttpClient) transport() *http.Transport {
	return client.Http().Transport.(*http.Transport)
}

func (client *HttpClient) TLSClientConfig() *tls.Config {
	return client.transport().TLSClientConfig
}

func (client *HttpClient) SetTLSClientConfig(config *tls.Config) {
	client.transport().TLSClientConfig = config
}

func (client *HttpClient) KeepAlive() bool {
	return !client.transport().DisableKeepAlives
}

func (client *HttpClient) SetKeepAlive(enable bool) {
	client.transport().DisableKeepAlives = !enable
}

func (client *HttpClient) Compression() bool {
	return !client.transport().DisableCompression
}

func (client *HttpClient) SetCompression(enable bool) {
	client.transport().DisableCompression = !enable
}

func (client *HttpClient) MaxIdleConnsPerHost() int {
	return client.transport().MaxIdleConnsPerHost
}

func (client *HttpClient) SetMaxIdleConnsPerHost(value int) {
	client.transport().MaxIdleConnsPerHost = value
}

func newHttpTransporter() *HttpTransporter {
	tr := &http.Transport{
		DisableCompression:  true,
		DisableKeepAlives:   false,
		MaxIdleConnsPerHost: 4}
	jar := cookieJar
	if DisableGlobalCookie {
		jar, _ = cookiejar.New(nil)
	}
	return &HttpTransporter{&http.Client{Jar: jar, Transport: tr}, &http.Header{}}
}

func (h *HttpTransporter) readAll(response *http.Response) (data []byte, err error) {
	if response.ContentLength > 0 {
		data = make([]byte, response.ContentLength)
		_, err = io.ReadFull(response.Body, data)
		return data, err
	}
	if response.ContentLength < 0 {
		return ioutil.ReadAll(response.Body)
	}
	return make([]byte, 0), nil
}

func (h *HttpTransporter) SendAndReceive(uri string, data []byte) ([]byte, error) {
	req, err := http.NewRequest("POST", uri, NewBytesReader(data))
	if err != nil {
		return nil, err
	}
	for key, values := range *h.Header {
		for _, value := range values {
			req.Header.Add(key, value)
		}
	}
	req.ContentLength = int64(len(data))
	req.Header.Set("Content-Type", "application/hprose")
	resp, err := h.Do(req)
	if err != nil {
		return nil, err
	}
	data, err = ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}
	return data, resp.Body.Close()
}
<|endoftext|>


<commit_before>package sudoku

import (
	"fmt"
	"strconv"
)

//TODO: investigate bumping this back up when #100 lands
const _MAX_IMPLICATION_STEPS = 5

type forcingChainsTechnique struct {
	*basicSolveTechnique
}

func (self *forcingChainsTechnique) numImplicationSteps(step *SolveStep) int {

	if step == nil {
		return 0
	}

	//Verify that the information we're unpacking is what we expect
	numImplicationSteps, ok := step.extra.(int)

	if !ok {
		numImplicationSteps = 0
	}
	return numImplicationSteps
}

func (self *forcingChainsTechnique) Variants() []string {
	var result []string
	for i := 1; i <= _MAX_IMPLICATION_STEPS+1; i++ {
		result = append(result, self.Name()+" ("+strconv.Itoa(i)+" steps)")
	}
	return result
}

func (self *forcingChainsTechnique) variant(step *SolveStep) string {
	return self.basicSolveTechnique.variant(step) + " (" + strconv.Itoa(self.numImplicationSteps(step)) + " steps)"
}

func (self *forcingChainsTechnique) humanLikelihood(step *SolveStep) float64 {
	//TODO: figure out what the baseDifficulty should be, this might be higher than
	//it's actually in practice

	//Note that this number has to be pretty high because it's competing against
	//HiddenSIZEGROUP, which has the k exponential in its favor.
	return self.difficultyHelper(20000.0)
}

func (self *forcingChainsTechnique) Description(step *SolveStep) string {
	return fmt.Sprintf("cell %s only has two options, %s, and if you put either one in and see the chain of implications it leads to, both ones end up with %s in cell %s, so we can just fill that number in", step.PointerCells.Description(), step.PointerNums.Description(), step.TargetNums.Description(), step.TargetCells.Description())
}

func (self *forcingChainsTechnique) Find(grid *Grid, results chan *SolveStep, done chan bool) {
	//TODO: test that this will find multiple if they exist.

	/*
	 * Conceptually this techinque chooses a cell with two possibilities
	 * and explores forward along two branches, seeing what would happen
	 * if it followed the simple implication chains forward to see if any
	 * cells end up set to the same number on both branches, meaning
	 * that no matter what, the cell will end up that value so you can set it
	 * that way now. In some ways it's like a very easy form of guessing.
	 *
	 * This techinque will do a BFS forward from the chosen cell, and won't
	 * explore more than _MAX_IMPLICATION_STEPS steps out from that. It will
	 * stop exploring if it finds one of two types of contradictions:
	 * 1) It notes that down this branch a single cell has had two different numbers
	 * implicated into it, which implies that somewhere earlier we ran into some inconsistency
	 * or
	 * 2) As soon as we note an inconsistency (a cell with no legal values).
	 *
	 * It is important to note that for every sudoku with one solution (that is, all
	 * legal puzzles), one of the two branches MUST lead to an inconsistency somewhere
	 * it's just a matter of how forward you have to go before you find it. That means
	 * that this technique is sensitive to the order in which you explore the frontiers
	 * of implications and when you choose to bail.
	 *
	 */

	getter := grid.queue().DefaultGetter()

	for {

		//Check if it's time to stop.
		select {
		case <-done:
			return
		default:
		}

		candidate := getter.GetSmallerThan(3)

		if candidate == nil {
			break
		}

		candidateCell := candidate.(*Cell)

		if len(candidateCell.Possibilities()) != 2 {
			//We found one with 1 possibility, which isn't interesting for us--nakedSingle should do that one.
			continue
		}

		firstPossibilityNum := candidateCell.Possibilities()[0]
		secondPossibilityNum := candidateCell.Possibilities()[1]

		firstGrid := grid.Copy()
		secondGrid := grid.Copy()

		//Check that the neighbor isn't just already having a single possibility, because then this technique is overkill.

		firstAccumulator := &chainSearcherAccumulator{make(map[cellRef]IntSlice), make(map[cellRef]IntSlice)}
		secondAccumulator := &chainSearcherAccumulator{make(map[cellRef]IntSlice), make(map[cellRef]IntSlice)}

		chainSearcher(0, _MAX_IMPLICATION_STEPS,
			candidateCell.InGrid(firstGrid),
			firstPossibilityNum, firstAccumulator)

		chainSearcher(0, _MAX_IMPLICATION_STEPS,
			candidateCell.InGrid(secondGrid),
			secondPossibilityNum, secondAccumulator)

		firstAccumulator.reduce()
		secondAccumulator.reduce()

		//See if either branch, at some generation, has the same cell forced to the same number in either generation.

		//We're just going to look at the last generation for each and compare
		//when each cell was setœ instead of doing (expensive!) pairwise
		//comparison across all of them

		for cell, numSlice := range firstAccumulator.numbers {
			if secondNumSlice, ok := secondAccumulator.numbers[cell]; ok {

				//Found two cells that overlap in terms of both being affected.
				//We're only interested in them if they are both set to exactly one item, which is the
				//same number.
				if len(numSlice) != 1 || len(secondNumSlice) != 1 || !numSlice.SameContentAs(secondNumSlice) {
					continue
				}

				numImplicationSteps := firstAccumulator.firstGeneration[cell][0] + secondAccumulator.firstGeneration[cell][0]

				//Is their combined generation count lower than _MAX_IMPLICATION_STEPS?
				if numImplicationSteps > _MAX_IMPLICATION_STEPS+1 {
					//Too many implication steps. :-(
					continue
				}

				//Okay, we have a candidate step. Is it useful?
				step := &SolveStep{self,
					CellSlice{cell.Cell(grid)},
					IntSlice{numSlice[0]},
					CellSlice{candidateCell},
					candidateCell.Possibilities(),
					numImplicationSteps,
				}

				if step.IsUseful(grid) {
					select {
					case results <- step:
					case <-done:
						return
					}
				}

			}
		}

		//TODO: figure out why the tests are coming back with different answers, even when only looking at the key cell
		//that should work from the example.
		//TODO: figure out a way to only compute a generation if required on each branch (don't compute all the way to _MAX_IMPLICATIONS to start)

		//TODO: ideally steps with a higher generation + generation score
		//would be scored as higher diffiuclty maybe include a
		//difficultyMultiplier in SolveStep that we can fill in? Hmmm, but
		//ideally it would factor  in at humanLikelihood level. Having a
		//million different ForcingChainLength techniques would be a
		//nightmare, peformance wise... unless there was a way to pass the
		//work done in one technique to another.

	}
}

type chainSearcherAccumulator struct {
	numbers         map[cellRef]IntSlice
	firstGeneration map[cellRef]IntSlice
}

func (c *chainSearcherAccumulator) String() string {
	result := "Begin map (length " + strconv.Itoa(len(c.numbers)) + ")\n"
	for cell, numSlice := range c.numbers {
		result += "\t" + cell.String() + " : " + numSlice.Description() + " : " + c.firstGeneration[cell].Description() + "\n"
	}
	result += "End map\n"
	return result
}

//Goes through each item in the map and removes duplicates, keeping the smallest generation seen for each unique number.
func (c *chainSearcherAccumulator) reduce() {
	for cell, numList := range c.numbers {
		output := make(map[int]int)
		generationList, ok := c.firstGeneration[cell]
		if !ok {
			panic("numbers and firstGeneration were out of sync")
		}
		for i, num := range numList {
			generation := generationList[i]
			currentGeneration := output[num]
			if currentGeneration == 0 || generation < currentGeneration {
				output[num] = generation
			}
		}
		var resultNumbers IntSlice
		var resultGenerations IntSlice
		for num, generation := range output {
			resultNumbers = append(resultNumbers, num)
			resultGenerations = append(resultGenerations, generation)
		}
		c.numbers[cell] = resultNumbers
		c.firstGeneration[cell] = resultGenerations
	}
}

func chainSearcher(generation int, maxGeneration int, cell *Cell, numToApply int, accum *chainSearcherAccumulator) {

	/*
	 * chainSearcher implements a DFS to search forward through implication chains to
	 * fill out accum with details about cells it sees and sets.
	 * The reason a DFS and not a BFS is called for is because with forcing chains, we
	 * KNOW that either the left or right branch will lead to an inconsistency at some point
	 * (as long as the sudoku has only one valid solution). We want to IGNORE that
	 * inconsistency for as long as possible to follow the implication chains as deep as we can go.
	 * By definition, the end of the DFS will be the farthest a given implication chain can go
	 * towards setting that specific cell to the forced value. This means that we have the maximum
	 * density of implication chain results to sift through to find cells forced to the same value.
	 */

	if generation > maxGeneration {
		//base case
		return
	}

	//Becuase this is a DFS, if we see an invalidity in this grid, it's a meaningful invalidity
	//and we should avoid it.
	if cell.grid.Invalid() {
		return
	}

	cellsToVisit := cell.Neighbors().FilterByPossible(numToApply).FilterByNumPossibilities(2)

	cell.SetNumber(numToApply)

	//Accumulate information about this cell being set. We'll reduce out duplicates later.
	accum.numbers[cell.ref()] = append(accum.numbers[cell.ref()], numToApply)
	accum.firstGeneration[cell.ref()] = append(accum.firstGeneration[cell.ref()], generation)

	for _, cellToVisit := range cellsToVisit {
		possibilities := cellToVisit.Possibilities()

		if len(possibilities) != 1 {
			panic("Expected the cell to have one possibility")
		}

		forcedNum := possibilities[0]

		//recurse
		chainSearcher(generation+1, maxGeneration, cellToVisit, forcedNum, accum)
	}

	//Undo this number and return
	cell.SetNumber(0)
}
<commit_msg>Made the default for a nil step for numImplicationSteps 1, so we don't get a humanLikelihood of 0 on a nil step. And also use numImplicationSteps to tweak humanLikelihood.<commit_after>package sudoku

import (
	"fmt"
	"strconv"
)

//TODO: investigate bumping this back up when #100 lands
const _MAX_IMPLICATION_STEPS = 5

type forcingChainsTechnique struct {
	*basicSolveTechnique
}

func (self *forcingChainsTechnique) numImplicationSteps(step *SolveStep) int {

	if step == nil {
		return 1
	}

	//Verify that the information we're unpacking is what we expect
	numImplicationSteps, ok := step.extra.(int)

	if !ok {
		numImplicationSteps = 1
	}
	return numImplicationSteps
}

func (self *forcingChainsTechnique) Variants() []string {
	var result []string
	for i := 1; i <= _MAX_IMPLICATION_STEPS+1; i++ {
		result = append(result, self.Name()+" ("+strconv.Itoa(i)+" steps)")
	}
	return result
}

func (self *forcingChainsTechnique) variant(step *SolveStep) string {
	return self.basicSolveTechnique.variant(step) + " (" + strconv.Itoa(self.numImplicationSteps(step)) + " steps)"
}

func (self *forcingChainsTechnique) humanLikelihood(step *SolveStep) float64 {
	//TODO: figure out what the baseDifficulty should be, this might be higher than
	//it's actually in practice

	//Note that this number has to be pretty high because it's competing against
	//HiddenSIZEGROUP, which has the k exponential in its favor.
	return float64(self.numImplicationSteps(step)) * self.difficultyHelper(20000.0)
}

func (self *forcingChainsTechnique) Description(step *SolveStep) string {
	return fmt.Sprintf("cell %s only has two options, %s, and if you put either one in and see the chain of implications it leads to, both ones end up with %s in cell %s, so we can just fill that number in", step.PointerCells.Description(), step.PointerNums.Description(), step.TargetNums.Description(), step.TargetCells.Description())
}

func (self *forcingChainsTechnique) Find(grid *Grid, results chan *SolveStep, done chan bool) {
	//TODO: test that this will find multiple if they exist.

	/*
	 * Conceptually this techinque chooses a cell with two possibilities
	 * and explores forward along two branches, seeing what would happen
	 * if it followed the simple implication chains forward to see if any
	 * cells end up set to the same number on both branches, meaning
	 * that no matter what, the cell will end up that value so you can set it
	 * that way now. In some ways it's like a very easy form of guessing.
	 *
	 * This techinque will do a BFS forward from the chosen cell, and won't
	 * explore more than _MAX_IMPLICATION_STEPS steps out from that. It will
	 * stop exploring if it finds one of two types of contradictions:
	 * 1) It notes that down this branch a single cell has had two different numbers
	 * implicated into it, which implies that somewhere earlier we ran into some inconsistency
	 * or
	 * 2) As soon as we note an inconsistency (a cell with no legal values).
	 *
	 * It is important to note that for every sudoku with one solution (that is, all
	 * legal puzzles), one of the two branches MUST lead to an inconsistency somewhere
	 * it's just a matter of how forward you have to go before you find it. That means
	 * that this technique is sensitive to the order in which you explore the frontiers
	 * of implications and when you choose to bail.
	 *
	 */

	getter := grid.queue().DefaultGetter()

	for {

		//Check if it's time to stop.
		select {
		case <-done:
			return
		default:
		}

		candidate := getter.GetSmallerThan(3)

		if candidate == nil {
			break
		}

		candidateCell := candidate.(*Cell)

		if len(candidateCell.Possibilities()) != 2 {
			//We found one with 1 possibility, which isn't interesting for us--nakedSingle should do that one.
			continue
		}

		firstPossibilityNum := candidateCell.Possibilities()[0]
		secondPossibilityNum := candidateCell.Possibilities()[1]

		firstGrid := grid.Copy()
		secondGrid := grid.Copy()

		//Check that the neighbor isn't just already having a single possibility, because then this technique is overkill.

		firstAccumulator := &chainSearcherAccumulator{make(map[cellRef]IntSlice), make(map[cellRef]IntSlice)}
		secondAccumulator := &chainSearcherAccumulator{make(map[cellRef]IntSlice), make(map[cellRef]IntSlice)}

		chainSearcher(0, _MAX_IMPLICATION_STEPS,
			candidateCell.InGrid(firstGrid),
			firstPossibilityNum, firstAccumulator)

		chainSearcher(0, _MAX_IMPLICATION_STEPS,
			candidateCell.InGrid(secondGrid),
			secondPossibilityNum, secondAccumulator)

		firstAccumulator.reduce()
		secondAccumulator.reduce()

		//See if either branch, at some generation, has the same cell forced to the same number in either generation.

		//We're just going to look at the last generation for each and compare
		//when each cell was setœ instead of doing (expensive!) pairwise
		//comparison across all of them

		for cell, numSlice := range firstAccumulator.numbers {
			if secondNumSlice, ok := secondAccumulator.numbers[cell]; ok {

				//Found two cells that overlap in terms of both being affected.
				//We're only interested in them if they are both set to exactly one item, which is the
				//same number.
				if len(numSlice) != 1 || len(secondNumSlice) != 1 || !numSlice.SameContentAs(secondNumSlice) {
					continue
				}

				numImplicationSteps := firstAccumulator.firstGeneration[cell][0] + secondAccumulator.firstGeneration[cell][0]

				//Is their combined generation count lower than _MAX_IMPLICATION_STEPS?
				if numImplicationSteps > _MAX_IMPLICATION_STEPS+1 {
					//Too many implication steps. :-(
					continue
				}

				//Okay, we have a candidate step. Is it useful?
				step := &SolveStep{self,
					CellSlice{cell.Cell(grid)},
					IntSlice{numSlice[0]},
					CellSlice{candidateCell},
					candidateCell.Possibilities(),
					numImplicationSteps,
				}

				if step.IsUseful(grid) {
					select {
					case results <- step:
					case <-done:
						return
					}
				}

			}
		}

		//TODO: figure out why the tests are coming back with different answers, even when only looking at the key cell
		//that should work from the example.
		//TODO: figure out a way to only compute a generation if required on each branch (don't compute all the way to _MAX_IMPLICATIONS to start)

		//TODO: ideally steps with a higher generation + generation score
		//would be scored as higher diffiuclty maybe include a
		//difficultyMultiplier in SolveStep that we can fill in? Hmmm, but
		//ideally it would factor  in at humanLikelihood level. Having a
		//million different ForcingChainLength techniques would be a
		//nightmare, peformance wise... unless there was a way to pass the
		//work done in one technique to another.

	}
}

type chainSearcherAccumulator struct {
	numbers         map[cellRef]IntSlice
	firstGeneration map[cellRef]IntSlice
}

func (c *chainSearcherAccumulator) String() string {
	result := "Begin map (length " + strconv.Itoa(len(c.numbers)) + ")\n"
	for cell, numSlice := range c.numbers {
		result += "\t" + cell.String() + " : " + numSlice.Description() + " : " + c.firstGeneration[cell].Description() + "\n"
	}
	result += "End map\n"
	return result
}

//Goes through each item in the map and removes duplicates, keeping the smallest generation seen for each unique number.
func (c *chainSearcherAccumulator) reduce() {
	for cell, numList := range c.numbers {
		output := make(map[int]int)
		generationList, ok := c.firstGeneration[cell]
		if !ok {
			panic("numbers and firstGeneration were out of sync")
		}
		for i, num := range numList {
			generation := generationList[i]
			currentGeneration := output[num]
			if currentGeneration == 0 || generation < currentGeneration {
				output[num] = generation
			}
		}
		var resultNumbers IntSlice
		var resultGenerations IntSlice
		for num, generation := range output {
			resultNumbers = append(resultNumbers, num)
			resultGenerations = append(resultGenerations, generation)
		}
		c.numbers[cell] = resultNumbers
		c.firstGeneration[cell] = resultGenerations
	}
}

func chainSearcher(generation int, maxGeneration int, cell *Cell, numToApply int, accum *chainSearcherAccumulator) {

	/*
	 * chainSearcher implements a DFS to search forward through implication chains to
	 * fill out accum with details about cells it sees and sets.
	 * The reason a DFS and not a BFS is called for is because with forcing chains, we
	 * KNOW that either the left or right branch will lead to an inconsistency at some point
	 * (as long as the sudoku has only one valid solution). We want to IGNORE that
	 * inconsistency for as long as possible to follow the implication chains as deep as we can go.
	 * By definition, the end of the DFS will be the farthest a given implication chain can go
	 * towards setting that specific cell to the forced value. This means that we have the maximum
	 * density of implication chain results to sift through to find cells forced to the same value.
	 */

	if generation > maxGeneration {
		//base case
		return
	}

	//Becuase this is a DFS, if we see an invalidity in this grid, it's a meaningful invalidity
	//and we should avoid it.
	if cell.grid.Invalid() {
		return
	}

	cellsToVisit := cell.Neighbors().FilterByPossible(numToApply).FilterByNumPossibilities(2)

	cell.SetNumber(numToApply)

	//Accumulate information about this cell being set. We'll reduce out duplicates later.
	accum.numbers[cell.ref()] = append(accum.numbers[cell.ref()], numToApply)
	accum.firstGeneration[cell.ref()] = append(accum.firstGeneration[cell.ref()], generation)

	for _, cellToVisit := range cellsToVisit {
		possibilities := cellToVisit.Possibilities()

		if len(possibilities) != 1 {
			panic("Expected the cell to have one possibility")
		}

		forcedNum := possibilities[0]

		//recurse
		chainSearcher(generation+1, maxGeneration, cellToVisit, forcedNum, accum)
	}

	//Undo this number and return
	cell.SetNumber(0)
}
<|endoftext|>


<commit_before>package integration_test

import (
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"os/user"
	"path/filepath"
	"runtime"

	"github.com/concourse/fly/version"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	"github.com/onsi/gomega/gbytes"
	"github.com/onsi/gomega/gexec"
	"github.com/onsi/gomega/ghttp"
)

var _ = Describe("Syncing", func() {
	var (
		flyVersion string
		flyPath    string
	)

	cliHandler := func() http.HandlerFunc {
		return ghttp.CombineHandlers(
			ghttp.VerifyRequest("GET", "/api/v1/cli"),
			func(w http.ResponseWriter, r *http.Request) {
				arch := r.URL.Query().Get("arch")
				platform := r.URL.Query().Get("platform")

				if arch != "amd64" && platform != runtime.GOOS {
					http.Error(w, "bad params", 500)
					return
				}

				w.WriteHeader(http.StatusOK)
				fmt.Fprint(w, "this will totally execute")
			},
		)
	}

	JustBeforeEach(func() {
		var err error
		flyPath, err = gexec.Build(
			"github.com/concourse/fly",
			"-ldflags", fmt.Sprintf("-X github.com/concourse/fly/version.Version=%s", flyVersion),
		)
		Expect(err).NotTo(HaveOccurred())

		atcServer.AppendHandlers(cliHandler())
	})

	Context("When versions mismatch between fly + atc", func() {
		BeforeEach(func() {
			major, minor, patch, err := version.GetSemver(atcVersion)
			Expect(err).NotTo(HaveOccurred())

			flyVersion = fmt.Sprintf("%d.%d.%d", major, minor, patch+1)
		})
		It("downloads and replaces the currently running executable", func() {
			flyCmd := exec.Command(flyPath, "-t", targetName, "sync")

			sess, err := gexec.Start(flyCmd, GinkgoWriter, GinkgoWriter)
			Expect(err).NotTo(HaveOccurred())

			<-sess.Exited
			Expect(sess.ExitCode()).To(Equal(0))

			expected := []byte("this will totally execute")
			expectBinaryToMatch(flyPath, expected[:8])
		})

		Context("When the user running sync doesn't have write permissions for the target directory", func() {
			It("returns an error, and doesn't download/replace the executable", func() {
				me, err := user.Current()
				Expect(err).ToNot(HaveOccurred())

				if me.Uid == "0" {
					Skip("root can always write; not worth testing")
					return
				}

				os.Chmod(filepath.Dir(flyPath), 0500)

				expectedBinary := readBinary(flyPath)

				flyCmd := exec.Command(flyPath, "-t", targetName, "sync")

				sess, err := gexec.Start(flyCmd, GinkgoWriter, GinkgoWriter)
				Expect(err).NotTo(HaveOccurred())

				<-sess.Exited
				Expect(sess.ExitCode()).To(Equal(1))
				Expect(sess.Err).To(gbytes.Say(".*permission denied"))

				expectBinaryToMatch(flyPath, expectedBinary)
			})
		})
	})

	Context("When versions match between fly + atc", func() {
		BeforeEach(func() {
			flyVersion = atcVersion
		})
		It("informs the user, and doesn't download/replace the executable", func() {
			expectedBinary := readBinary(flyPath)

			flyCmd := exec.Command(flyPath, "-t", targetName, "sync")

			sess, err := gexec.Start(flyCmd, GinkgoWriter, GinkgoWriter)
			Expect(err).NotTo(HaveOccurred())

			<-sess.Exited
			Expect(sess.ExitCode()).To(Equal(0))
			Expect(sess.Out).To(gbytes.Say(`version already matches; skipping`))

			expectBinaryToMatch(flyPath, expectedBinary)
		})
	})
})

func readBinary(path string) []byte {
	expectedBinary, err := ioutil.ReadFile(flyPath)
	Expect(err).NotTo(HaveOccurred())
	return expectedBinary[:8]
}

func expectBinaryToMatch(path string, expectedBinary []byte) {
	contents, err := ioutil.ReadFile(path)
	Expect(err).NotTo(HaveOccurred())

	// don't let ginkgo try and output the entire binary as ascii
	//
	// that is the way to the dark side
	contents = contents[:8]
	Expect(contents).To(Equal(expectedBinary))
}
<commit_msg>skip windows<commit_after>package integration_test

import (
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"os/user"
	"path/filepath"
	"runtime"

	"github.com/concourse/fly/version"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	"github.com/onsi/gomega/gbytes"
	"github.com/onsi/gomega/gexec"
	"github.com/onsi/gomega/ghttp"
)

var _ = Describe("Syncing", func() {
	var (
		flyVersion string
		flyPath    string
	)

	cliHandler := func() http.HandlerFunc {
		return ghttp.CombineHandlers(
			ghttp.VerifyRequest("GET", "/api/v1/cli"),
			func(w http.ResponseWriter, r *http.Request) {
				arch := r.URL.Query().Get("arch")
				platform := r.URL.Query().Get("platform")

				if arch != "amd64" && platform != runtime.GOOS {
					http.Error(w, "bad params", 500)
					return
				}

				w.WriteHeader(http.StatusOK)
				fmt.Fprint(w, "this will totally execute")
			},
		)
	}

	JustBeforeEach(func() {
		var err error
		flyPath, err = gexec.Build(
			"github.com/concourse/fly",
			"-ldflags", fmt.Sprintf("-X github.com/concourse/fly/version.Version=%s", flyVersion),
		)
		Expect(err).NotTo(HaveOccurred())

		atcServer.AppendHandlers(cliHandler())
	})

	Context("When versions mismatch between fly + atc", func() {
		BeforeEach(func() {
			major, minor, patch, err := version.GetSemver(atcVersion)
			Expect(err).NotTo(HaveOccurred())

			flyVersion = fmt.Sprintf("%d.%d.%d", major, minor, patch+1)
		})
		It("downloads and replaces the currently running executable", func() {
			flyCmd := exec.Command(flyPath, "-t", targetName, "sync")

			sess, err := gexec.Start(flyCmd, GinkgoWriter, GinkgoWriter)
			Expect(err).NotTo(HaveOccurred())

			<-sess.Exited
			Expect(sess.ExitCode()).To(Equal(0))

			expected := []byte("this will totally execute")
			expectBinaryToMatch(flyPath, expected[:8])
		})

		Context("When the user running sync doesn't have write permissions for the target directory", func() {
			It("returns an error, and doesn't download/replace the executable", func() {
				me, err := user.Current()
				Expect(err).ToNot(HaveOccurred())

				if me.Uid == "0" {
					Skip("root can always write; not worth testing")
					return
				}

				if runtime.GOOS == "windows" {
					Skip("who knows how windows works; not worth testing")
					return
				}

				os.Chmod(filepath.Dir(flyPath), 0500)

				expectedBinary := readBinary(flyPath)

				flyCmd := exec.Command(flyPath, "-t", targetName, "sync")

				sess, err := gexec.Start(flyCmd, GinkgoWriter, GinkgoWriter)
				Expect(err).NotTo(HaveOccurred())

				<-sess.Exited
				Expect(sess.ExitCode()).To(Equal(1))
				Expect(sess.Err).To(gbytes.Say(".*permission denied"))

				expectBinaryToMatch(flyPath, expectedBinary)
			})
		})
	})

	Context("When versions match between fly + atc", func() {
		BeforeEach(func() {
			flyVersion = atcVersion
		})
		It("informs the user, and doesn't download/replace the executable", func() {
			expectedBinary := readBinary(flyPath)

			flyCmd := exec.Command(flyPath, "-t", targetName, "sync")

			sess, err := gexec.Start(flyCmd, GinkgoWriter, GinkgoWriter)
			Expect(err).NotTo(HaveOccurred())

			<-sess.Exited
			Expect(sess.ExitCode()).To(Equal(0))
			Expect(sess.Out).To(gbytes.Say(`version already matches; skipping`))

			expectBinaryToMatch(flyPath, expectedBinary)
		})
	})
})

func readBinary(path string) []byte {
	expectedBinary, err := ioutil.ReadFile(flyPath)
	Expect(err).NotTo(HaveOccurred())
	return expectedBinary[:8]
}

func expectBinaryToMatch(path string, expectedBinary []byte) {
	contents, err := ioutil.ReadFile(path)
	Expect(err).NotTo(HaveOccurred())

	// don't let ginkgo try and output the entire binary as ascii
	//
	// that is the way to the dark side
	contents = contents[:8]
	Expect(contents).To(Equal(expectedBinary))
}
<|endoftext|>


<commit_before>package secretcrypt

import (
	"flag"
	"fmt"
	"os"
	"testing"

	"github.com/Zemanta/go-secretcrypt/internal"
	"github.com/stretchr/testify/assert"
)

func TestMain(m *testing.M) {
	internal.CryptersMap["plain"] = internal.PlainCrypter{}
	flag.Parse()
	os.Exit(m.Run())
}

func assertStrictSecretValid(t *testing.T, secret StrictSecret) {
	assert.Equal(t, "plain", secret.crypter.Name())
	assert.Equal(t, "my-abc", string(secret.ciphertext))
	assert.Equal(t, internal.DecryptParams{
		"k1": "v1",
		"k2": "v2",
	}, secret.decryptParams)
}

func TestUnmarshalText(t *testing.T) {
	var secret StrictSecret
	err := secret.UnmarshalText([]byte("plain:k1=v1&k2=v2:my-abc"))
	assert.NoError(t, err)
	assertStrictSecretValid(t, secret)
}

func TestNewStrictSecret(t *testing.T) {
	secret, err := LoadStrictSecret("plain:k1=v1&k2=v2:my-abc")
	assert.NoError(t, err)
	assertStrictSecretValid(t, secret)
}

func TestDecrypt(t *testing.T) {
	mockCrypter := &internal.MockCrypter{}
	internal.CryptersMap["mock"] = mockCrypter
	mockCrypter.On(
		"Decrypt",
		internal.Ciphertext("my-abc"),
		internal.DecryptParams{
			"k1": "v1",
			"k2": "v2",
		}).Return("myplaintext", nil)

	secret, err := LoadStrictSecret("mock:k1=v1&k2=v2:my-abc")
	assert.NoError(t, err)

	plaintext, err := secret.Decrypt()
	assert.NoError(t, err)
	assert.Equal(t, plaintext, "myplaintext")
	mockCrypter.AssertExpectations(t)
}

func TestNoCaching(t *testing.T) {
	mockCrypter := &internal.MockCrypter{}
	internal.CryptersMap["mock"] = mockCrypter
	mockCrypter.On(
		"Decrypt",
		internal.Ciphertext("my-abc"),
		internal.DecryptParams{
			"k1": "v1",
			"k2": "v2",
		}).Return("myplaintext", nil)

	secret, err := LoadStrictSecret("mock:k1=v1&k2=v2:my-abc")
	assert.NoError(t, err)

	plaintext, err := secret.Decrypt()
	assert.NoError(t, err)
	assert.Equal(t, plaintext, "myplaintext")
	plaintext2, err := secret.Decrypt()
	assert.NoError(t, err)
	assert.Equal(t, plaintext2, "myplaintext")

	mockCrypter.AssertExpectations(t)
	mockCrypter.AssertNumberOfCalls(t, "Decrypt", 2)
}

func TestEmptyStrictSecret(t *testing.T) {
	zero := StrictSecret{}
	emptyStr, err := LoadStrictSecret("")
	assert.NoError(t, err)
	for _, secret := range []StrictSecret{zero, emptyStr} {
		plaintext, err := secret.Decrypt()
		assert.NoError(t, err)
		assert.Equal(t, plaintext, "")
	}
}

func TestSecret(t *testing.T) {
	var secret Secret
	err := secret.UnmarshalText([]byte("plain:k1=v1&k2=v2:my-abc"))
	assert.NoError(t, err)
	assert.Equal(t, "my-abc", secret.Get())
}

func TestStrictSecretMarshalText(t *testing.T) {
	var ssecret StrictSecret
	err := ssecret.UnmarshalText([]byte("plain:k1=v1&k2=v2:my-abc"))
	assert.NoError(t, err)

	d, err := ssecret.Decrypt()
	assert.NoError(t, err)
	assert.Equal(t, "my-abc", d)

	extraParams := internal.DecryptParams{
		"k3": "v3",
	}
	ssecret.AppendParameters(extraParams)

	text, err := ssecret.MarshalText()
	assert.NoError(t, err)
	assert.Equal(t, "plain:k1=v1&k2=v2&k3=v3:my-abc", string(text))
}

func TestSecretMarshalText(t *testing.T) {
	secret, err := LoadSecret("invalid")
	assert.Error(t, err)

	secret, err = LoadSecret("plain:k1=v1&k2=v2:my-abc")
	assert.NoError(t, err)
	assert.Equal(t, "my-abc", secret.Get())

	text, err := secret.MarshalText()
	assert.NoError(t, err)
	assert.NotNil(t, text)
	assert.Equal(t, "", string(text))
	// assert.Equal(t, "plain:k1=v1&k2=v2:my-abc", string(text)) // TODO: check why it gets "" here
}

func TestStrictSecretUnmarshalTextError(t *testing.T) {
	var ssecret StrictSecret
	err := ssecret.UnmarshalText([]byte("plain:k1=v1&k2=v2Missing3rdComponent"))
	assert.Error(t, err, "missing 3rd component (ciphertext)")

	err = ssecret.UnmarshalText([]byte("invalid:k1=v1&k2=v2:my-abc"))
	assert.Error(t, err, "should be invalid crypter")
}

func TestSecretRedaction(t *testing.T) {
	var s Secret
	err := s.UnmarshalText([]byte("plain:k1=v1&k2=v2:my-abc"))
	assert.NoError(t, err)
	assert.Equal(t, "my-abc", s.Get())
	assert.Equal(t, "<redacted>", s.String())
	assert.Equal(t, "Secret: <redacted>", fmt.Sprintf("Secret: %s", s))
	assert.Equal(t, "Secret: <redacted>", fmt.Sprintf("Secret: %s", &s))
	assert.Equal(t, "<redacted>", s.GoString())
	assert.Equal(t, "Go secret: <redacted>", fmt.Sprintf("Go secret: %#v", s))
	assert.Equal(t, "Go secret: <redacted>", fmt.Sprintf("Go secret: %#v", &s))
}

func TestStrictSecretPlainRedaction(t *testing.T) {
	var ss StrictSecret
	err := ss.UnmarshalText([]byte("plain:k1=v1&k2=v2:my-abc"))
	assert.NoError(t, err)
	d, err := ss.Decrypt()
	assert.NoError(t, err)
	assert.Equal(t, "my-abc", d)

	// note: ciphertext of plain is same as decrypted, not actutally redacted!
	assert.Equal(t, "my-abc", ss.String())
	assert.Equal(t, "Secret: my-abc", fmt.Sprintf("Secret: %s", ss))
	assert.Equal(t, "Secret: my-abc", fmt.Sprintf("Secret: %s", &ss))
	assert.Equal(t, "my-abc", ss.GoString())
	assert.Equal(t, "Go secret: my-abc", fmt.Sprintf("Go secret: %#v", ss))
	assert.Equal(t, "Go secret: my-abc", fmt.Sprintf("Go secret: %#v", &ss))
}
<commit_msg>Fix test<commit_after>package secretcrypt

import (
	"flag"
	"fmt"
	"os"
	"testing"

	"github.com/Zemanta/go-secretcrypt/internal"
	"github.com/stretchr/testify/assert"
)

func TestMain(m *testing.M) {
	internal.CryptersMap["plain"] = internal.PlainCrypter{}
	flag.Parse()
	os.Exit(m.Run())
}

func assertStrictSecretValid(t *testing.T, secret StrictSecret) {
	assert.Equal(t, "plain", secret.crypter.Name())
	assert.Equal(t, "my-abc", string(secret.ciphertext))
	assert.Equal(t, internal.DecryptParams{
		"k1": "v1",
		"k2": "v2",
	}, secret.decryptParams)
}

func TestUnmarshalText(t *testing.T) {
	var secret StrictSecret
	err := secret.UnmarshalText([]byte("plain:k1=v1&k2=v2:my-abc"))
	assert.NoError(t, err)
	assertStrictSecretValid(t, secret)
}

func TestNewStrictSecret(t *testing.T) {
	secret, err := LoadStrictSecret("plain:k1=v1&k2=v2:my-abc")
	assert.NoError(t, err)
	assertStrictSecretValid(t, secret)
}

func TestDecrypt(t *testing.T) {
	mockCrypter := &internal.MockCrypter{}
	internal.CryptersMap["mock"] = mockCrypter
	mockCrypter.On(
		"Decrypt",
		internal.Ciphertext("my-abc"),
		internal.DecryptParams{
			"k1": "v1",
			"k2": "v2",
		}).Return("myplaintext", nil)

	secret, err := LoadStrictSecret("mock:k1=v1&k2=v2:my-abc")
	assert.NoError(t, err)

	plaintext, err := secret.Decrypt()
	assert.NoError(t, err)
	assert.Equal(t, plaintext, "myplaintext")
	mockCrypter.AssertExpectations(t)
}

func TestNoCaching(t *testing.T) {
	mockCrypter := &internal.MockCrypter{}
	internal.CryptersMap["mock"] = mockCrypter
	mockCrypter.On(
		"Decrypt",
		internal.Ciphertext("my-abc"),
		internal.DecryptParams{
			"k1": "v1",
			"k2": "v2",
		}).Return("myplaintext", nil)

	secret, err := LoadStrictSecret("mock:k1=v1&k2=v2:my-abc")
	assert.NoError(t, err)

	plaintext, err := secret.Decrypt()
	assert.NoError(t, err)
	assert.Equal(t, plaintext, "myplaintext")
	plaintext2, err := secret.Decrypt()
	assert.NoError(t, err)
	assert.Equal(t, plaintext2, "myplaintext")

	mockCrypter.AssertExpectations(t)
	mockCrypter.AssertNumberOfCalls(t, "Decrypt", 2)
}

func TestEmptyStrictSecret(t *testing.T) {
	zero := StrictSecret{}
	emptyStr, err := LoadStrictSecret("")
	assert.NoError(t, err)
	for _, secret := range []StrictSecret{zero, emptyStr} {
		plaintext, err := secret.Decrypt()
		assert.NoError(t, err)
		assert.Equal(t, plaintext, "")
	}
}

func TestSecret(t *testing.T) {
	var secret Secret
	err := secret.UnmarshalText([]byte("plain:k1=v1&k2=v2:my-abc"))
	assert.NoError(t, err)
	assert.Equal(t, "my-abc", secret.Get())
}

func TestStrictSecretMarshalText(t *testing.T) {
	var ssecret StrictSecret
	err := ssecret.UnmarshalText([]byte("plain:k1=v1&k2=v2:my-abc"))
	assert.NoError(t, err)

	d, err := ssecret.Decrypt()
	assert.NoError(t, err)
	assert.Equal(t, "my-abc", d)

	extraParams := internal.DecryptParams{
		"k3": "v3",
	}
	ssecret.AppendParameters(extraParams)

	text, err := ssecret.MarshalText()
	assert.NoError(t, err)
	assert.Equal(t, "plain:k1=v1&k2=v2&k3=v3:my-abc", string(text))
}

func TestSecretMarshalText(t *testing.T) {
	secret, err := LoadSecret("invalid")
	assert.Error(t, err)

	secret, err = LoadSecret("plain:k1=v1&k2=v2:my-abc")
	assert.NoError(t, err)
	assert.Equal(t, "my-abc", secret.Get())

	text, err := secret.MarshalText()
	assert.NoError(t, err)
	assert.NotNil(t, text)
	assert.Equal(t, "plain:k1=v1&k2=v2:my-abc", string(text))
}

func TestStrictSecretUnmarshalTextError(t *testing.T) {
	var ssecret StrictSecret
	err := ssecret.UnmarshalText([]byte("plain:k1=v1&k2=v2Missing3rdComponent"))
	assert.Error(t, err, "missing 3rd component (ciphertext)")

	err = ssecret.UnmarshalText([]byte("invalid:k1=v1&k2=v2:my-abc"))
	assert.Error(t, err, "should be invalid crypter")
}

func TestSecretRedaction(t *testing.T) {
	var s Secret
	err := s.UnmarshalText([]byte("plain:k1=v1&k2=v2:my-abc"))
	assert.NoError(t, err)
	assert.Equal(t, "my-abc", s.Get())
	assert.Equal(t, "<redacted>", s.String())
	assert.Equal(t, "Secret: <redacted>", fmt.Sprintf("Secret: %s", s))
	assert.Equal(t, "Secret: <redacted>", fmt.Sprintf("Secret: %s", &s))
	assert.Equal(t, "<redacted>", s.GoString())
	assert.Equal(t, "Go secret: <redacted>", fmt.Sprintf("Go secret: %#v", s))
	assert.Equal(t, "Go secret: <redacted>", fmt.Sprintf("Go secret: %#v", &s))
}

func TestStrictSecretPlainRedaction(t *testing.T) {
	var ss StrictSecret
	err := ss.UnmarshalText([]byte("plain:k1=v1&k2=v2:my-abc"))
	assert.NoError(t, err)
	d, err := ss.Decrypt()
	assert.NoError(t, err)
	assert.Equal(t, "my-abc", d)

	// note: ciphertext of plain is same as decrypted, not actutally redacted!
	assert.Equal(t, "my-abc", ss.String())
	assert.Equal(t, "Secret: my-abc", fmt.Sprintf("Secret: %s", ss))
	assert.Equal(t, "Secret: my-abc", fmt.Sprintf("Secret: %s", &ss))
	assert.Equal(t, "my-abc", ss.GoString())
	assert.Equal(t, "Go secret: my-abc", fmt.Sprintf("Go secret: %#v", ss))
	assert.Equal(t, "Go secret: my-abc", fmt.Sprintf("Go secret: %#v", &ss))
}
<|endoftext|>


<commit_before>// Copyright © 2017 Circonus, Inc. <support@circonus.com>
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
//

package check

import (
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"io/ioutil"
	"math/rand"
	"net"
	"net/url"
	"reflect"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/circonus-labs/circonus-agent/internal/config"
	"github.com/circonus-labs/circonus-gometrics/api"
	"github.com/pkg/errors"
	"github.com/spf13/viper"
)

func (c *Check) setReverseConfig() error {
	c.Lock()
	defer c.Unlock()

	if len(c.bundle.ReverseConnectURLs) == 0 {
		return errors.New("no reverse URLs found in check bundle")
	}
	rURL := c.bundle.ReverseConnectURLs[0]
	rSecret := c.bundle.Config["reverse:secret_key"]

	if rSecret != "" {
		rURL += "#" + rSecret
	}

	// Replace protocol, url.Parse does not understand 'mtev_reverse'.
	// Important part is validating what's after 'proto://'.
	// Using raw tls connections, the url protocol is not germane.
	reverseURL, err := url.Parse(strings.Replace(rURL, "mtev_reverse", "http", -1))
	if err != nil {
		return errors.Wrapf(err, "parsing check bundle reverse URL (%s)", rURL)
	}

	brokerAddr, err := net.ResolveTCPAddr("tcp", reverseURL.Host)
	if err != nil {
		return errors.Wrapf(err, "invalid reverse service address", rURL)
	}

	if len(c.bundle.Brokers) == 0 {
		return errors.New("no brokers found in check bundle")
	}
	brokerID := c.bundle.Brokers[0]

	tlsConfig, err := c.brokerTLSConfig(brokerID, reverseURL)
	if err != nil {
		return errors.Wrapf(err, "creating TLS config for (%s - %s)", brokerID, rURL)
	}

	c.revConfig = &ReverseConfig{
		ReverseURL: reverseURL,
		BrokerID:   brokerID,
		BrokerAddr: brokerAddr,
		TLSConfig:  tlsConfig,
	}

	return nil
}

// brokerTLSConfig returns the correct TLS configuration for the broker
func (c *Check) brokerTLSConfig(cid string, reverseURL *url.URL) (*tls.Config, error) {
	if cid == "" {
		return nil, errors.New("invalid broker cid (empty)")
	}

	bcid := cid

	if ok, _ := regexp.MatchString(`^[0-9]+$`, bcid); ok {
		bcid = "/broker/" + cid
	}

	if ok, _ := regexp.MatchString(`^/broker/[0-9]+$`, bcid); !ok {
		return nil, errors.Errorf("invalid broker cid (%s)", cid)
	}

	broker, err := c.client.FetchBroker(api.CIDType(&bcid))
	if err != nil {
		return nil, errors.Wrapf(err, "unable to retrieve broker (%s)", cid)
	}

	cn, err := c.getBrokerCN(broker, reverseURL)
	if err != nil {
		return nil, err
	}
	cert, err := c.fetchBrokerCA()
	if err != nil {
		return nil, err
	}
	cp := x509.NewCertPool()
	if !cp.AppendCertsFromPEM(cert) {
		return nil, errors.New("unable to add Broker CA Certificate to x509 cert pool")
	}

	tlsConfig := &tls.Config{
		RootCAs:    cp,
		ServerName: cn,
	}

	c.logger.Debug().Str("CN", cn).Msg("setting tls CN")

	return tlsConfig, nil
}

func (c *Check) getBrokerCN(broker *api.Broker, reverseURL *url.URL) (string, error) {
	host := reverseURL.Hostname()

	// OK...
	//
	// mtev_reverse can have an IP or an FQDN for the host portion
	// it used to be that when it was an IP, the CN was needed in order to verify TLS connections
	// otherwise, the FQDN was valid. now, the FQDN may be valid for the cert or it may not be...

	cn := ""

	for _, detail := range broker.Details {
		// certs are generated against the CN (in theory)
		// 1. find the right broker instance with matching IP or external hostname
		// 2. set the tls.Config.ServerName to whatever that instance's CN is currently
		// 3. cert will be valid for TLS conns (in theory)
		if detail.IP != nil && *detail.IP == host {
			cn = detail.CN
			break
		}
		if detail.ExternalHost != nil && *detail.ExternalHost == host {
			cn = detail.CN
			break
		}
	}

	if cn == "" {
		return "", errors.Errorf("unable to match reverse URL host (%s) to broker", host)
	}

	return cn, nil
}

func (c *Check) fetchBrokerCA() ([]byte, error) {
	// use local file if specified
	file := viper.GetString(config.KeyReverseBrokerCAFile)
	if file != "" {
		cert, err := ioutil.ReadFile(file)
		if err != nil {
			return nil, errors.Wrapf(err, "reading specified broker-ca-file (%s)", file)
		}
		return cert, nil
	}

	// otherwise, try the api
	data, err := c.client.Get("/pki/ca.crt")
	if err != nil {
		return nil, errors.Wrap(err, "fetching Broker CA certificate")
	}

	type cacert struct {
		Contents string `json:"contents"`
	}

	var cadata cacert

	if err := json.Unmarshal(data, &cadata); err != nil {
		return nil, errors.Wrap(err, "parsing Broker CA certificate")
	}

	if cadata.Contents == "" {
		return nil, errors.Errorf("no Broker CA certificate in response (%#v)", string(data))
	}

	return []byte(cadata.Contents), nil
}

// Select a broker for use when creating a check, if a specific broker
// was not specified.
func (c *Check) selectBroker(checkType string) (*api.Broker, error) {
	brokerList, err := c.client.FetchBrokers()
	if err != nil {
		return nil, errors.Wrap(err, "select broker")
	}

	if len(*brokerList) == 0 {
		return nil, errors.New("no brokers returned from API")
	}

	validBrokers := make(map[string]api.Broker)
	haveEnterprise := false
	threshold := 10 * time.Second

	for _, broker := range *brokerList {
		broker := broker
		dur, ok := c.isValidBroker(&broker, checkType)
		if ok {
			if dur > threshold {
				continue
			} else if dur == threshold {
				validBrokers[broker.CID] = broker
			} else if dur < threshold {
				validBrokers = make(map[string]api.Broker)
				haveEnterprise = false
				threshold = dur
				validBrokers[broker.CID] = broker
			}
			if broker.Type == "enterprise" {
				haveEnterprise = true
			}
		}
	}

	if haveEnterprise { // eliminate non-enterprise brokers from valid brokers
		for k, v := range validBrokers {
			if v.Type != "enterprise" {
				delete(validBrokers, k)
			}
		}
	}

	if len(validBrokers) == 0 {
		return nil, errors.Errorf("found %d broker(s), zero are valid", len(*brokerList))
	}

	var selectedBroker api.Broker
	validBrokerKeys := reflect.ValueOf(validBrokers).MapKeys()
	if len(validBrokerKeys) == 1 {
		selectedBroker = validBrokers[validBrokerKeys[0].String()]
	} else {
		selectedBroker = validBrokers[validBrokerKeys[rand.Intn(len(validBrokerKeys))].String()]
	}

	c.logger.Debug().Str("broker", selectedBroker.Name).Msg("selected")

	return &selectedBroker, nil
}

// Is the broker valid (active, supports check type, and reachable)
func (c *Check) isValidBroker(broker *api.Broker, checkType string) (time.Duration, bool) {
	var brokerHost string
	var brokerPort string
	var connDuration time.Duration
	valid := false

	for _, detail := range broker.Details {
		detail := detail

		// broker must be active
		if detail.Status != brokerActiveStatus {
			c.logger.Debug().Str("broker", broker.Name).Msg("not active, skipping")
			continue
		}

		// broker must have module loaded for the check type to be used
		if !brokerSupportsCheckType(checkType, &detail) {
			c.logger.Debug().Str("broker", broker.Name).Str("type", checkType).Msg("unsupported check type, skipping")
			continue
		}

		if detail.ExternalPort != 0 {
			brokerPort = strconv.Itoa(int(detail.ExternalPort))
		} else {
			if *detail.Port != 0 {
				brokerPort = strconv.Itoa(int(*detail.Port))
			} else {
				brokerPort = "43191"
			}
		}

		if detail.ExternalHost != nil && *detail.ExternalHost != "" {
			brokerHost = *detail.ExternalHost
		} else {
			brokerHost = *detail.IP
		}

		if brokerHost == "trap.noit.circonus.net" && brokerPort != "443" {
			brokerPort = "443"
		}

		minDelay := int(200 * time.Millisecond)
		maxDelay := int(2 * time.Second)

		for attempt := 1; attempt <= brokerMaxRetries; attempt++ {
			start := time.Now()
			// broker must be reachable and respond within designated time
			conn, err := net.DialTimeout("tcp", net.JoinHostPort(brokerHost, brokerPort), brokerMaxResponseTime)
			if err == nil {
				connDuration = time.Since(start)
				conn.Close()
				valid = true
				break
			}

			delay := time.Duration(rand.Intn(maxDelay-minDelay) + minDelay)

			c.logger.Warn().
				Err(err).
				Str("delay", delay.String()).
				Str("broker", broker.Name).
				Int("attempt", attempt).
				Int("retries", brokerMaxRetries).
				Msg("unable to connect, retrying")

			time.Sleep(delay)
		}

		if valid {
			c.logger.Debug().Str("broker", broker.Name).Msg("valid")
			break
		}
	}

	return connDuration, valid
}

// brokerSupportsCheckType verifies a broker supports the check type to be used
func brokerSupportsCheckType(checkType string, details *api.BrokerDetail) bool {
	baseType := string(checkType)

	if idx := strings.Index(baseType, ":"); idx > 0 {
		baseType = baseType[0:idx]
	}

	for _, module := range details.Modules {
		if module == baseType {
			return true
		}
	}

	return false
}
<commit_msg>upd: move pkg consts to struct<commit_after>// Copyright © 2017 Circonus, Inc. <support@circonus.com>
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
//

package check

import (
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"io/ioutil"
	"math/rand"
	"net"
	"net/url"
	"reflect"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/circonus-labs/circonus-agent/internal/config"
	"github.com/circonus-labs/circonus-gometrics/api"
	"github.com/pkg/errors"
	"github.com/spf13/viper"
)

func (c *Check) setReverseConfig() error {
	c.Lock()
	defer c.Unlock()

	if len(c.bundle.ReverseConnectURLs) == 0 {
		return errors.New("no reverse URLs found in check bundle")
	}
	rURL := c.bundle.ReverseConnectURLs[0]
	rSecret := c.bundle.Config["reverse:secret_key"]

	if rSecret != "" {
		rURL += "#" + rSecret
	}

	// Replace protocol, url.Parse does not understand 'mtev_reverse'.
	// Important part is validating what's after 'proto://'.
	// Using raw tls connections, the url protocol is not germane.
	reverseURL, err := url.Parse(strings.Replace(rURL, "mtev_reverse", "http", -1))
	if err != nil {
		return errors.Wrapf(err, "parsing check bundle reverse URL (%s)", rURL)
	}

	brokerAddr, err := net.ResolveTCPAddr("tcp", reverseURL.Host)
	if err != nil {
		return errors.Wrapf(err, "invalid reverse service address", rURL)
	}

	if len(c.bundle.Brokers) == 0 {
		return errors.New("no brokers found in check bundle")
	}
	brokerID := c.bundle.Brokers[0]

	tlsConfig, err := c.brokerTLSConfig(brokerID, reverseURL)
	if err != nil {
		return errors.Wrapf(err, "creating TLS config for (%s - %s)", brokerID, rURL)
	}

	c.revConfig = &ReverseConfig{
		ReverseURL: reverseURL,
		BrokerID:   brokerID,
		BrokerAddr: brokerAddr,
		TLSConfig:  tlsConfig,
	}

	return nil
}

// brokerTLSConfig returns the correct TLS configuration for the broker
func (c *Check) brokerTLSConfig(cid string, reverseURL *url.URL) (*tls.Config, error) {
	if cid == "" {
		return nil, errors.New("invalid broker cid (empty)")
	}

	bcid := cid

	if ok, _ := regexp.MatchString(`^[0-9]+$`, bcid); ok {
		bcid = "/broker/" + cid
	}

	if ok, _ := regexp.MatchString(`^/broker/[0-9]+$`, bcid); !ok {
		return nil, errors.Errorf("invalid broker cid (%s)", cid)
	}

	broker, err := c.client.FetchBroker(api.CIDType(&bcid))
	if err != nil {
		return nil, errors.Wrapf(err, "unable to retrieve broker (%s)", cid)
	}

	cn, err := c.getBrokerCN(broker, reverseURL)
	if err != nil {
		return nil, err
	}
	cert, err := c.fetchBrokerCA()
	if err != nil {
		return nil, err
	}
	cp := x509.NewCertPool()
	if !cp.AppendCertsFromPEM(cert) {
		return nil, errors.New("unable to add Broker CA Certificate to x509 cert pool")
	}

	tlsConfig := &tls.Config{
		RootCAs:    cp,
		ServerName: cn,
	}

	c.logger.Debug().Str("CN", cn).Msg("setting tls CN")

	return tlsConfig, nil
}

func (c *Check) getBrokerCN(broker *api.Broker, reverseURL *url.URL) (string, error) {
	host := reverseURL.Hostname()

	// OK...
	//
	// mtev_reverse can have an IP or an FQDN for the host portion
	// it used to be that when it was an IP, the CN was needed in order to verify TLS connections
	// otherwise, the FQDN was valid. now, the FQDN may be valid for the cert or it may not be...

	cn := ""

	for _, detail := range broker.Details {
		// certs are generated against the CN (in theory)
		// 1. find the right broker instance with matching IP or external hostname
		// 2. set the tls.Config.ServerName to whatever that instance's CN is currently
		// 3. cert will be valid for TLS conns (in theory)
		if detail.IP != nil && *detail.IP == host {
			cn = detail.CN
			break
		}
		if detail.ExternalHost != nil && *detail.ExternalHost == host {
			cn = detail.CN
			break
		}
	}

	if cn == "" {
		return "", errors.Errorf("unable to match reverse URL host (%s) to broker", host)
	}

	return cn, nil
}

func (c *Check) fetchBrokerCA() ([]byte, error) {
	// use local file if specified
	file := viper.GetString(config.KeyReverseBrokerCAFile)
	if file != "" {
		cert, err := ioutil.ReadFile(file)
		if err != nil {
			return nil, errors.Wrapf(err, "reading specified broker-ca-file (%s)", file)
		}
		return cert, nil
	}

	// otherwise, try the api
	data, err := c.client.Get("/pki/ca.crt")
	if err != nil {
		return nil, errors.Wrap(err, "fetching Broker CA certificate")
	}

	type cacert struct {
		Contents string `json:"contents"`
	}

	var cadata cacert

	if err := json.Unmarshal(data, &cadata); err != nil {
		return nil, errors.Wrap(err, "parsing Broker CA certificate")
	}

	if cadata.Contents == "" {
		return nil, errors.Errorf("no Broker CA certificate in response (%#v)", string(data))
	}

	return []byte(cadata.Contents), nil
}

// Select a broker for use when creating a check, if a specific broker
// was not specified.
func (c *Check) selectBroker(checkType string) (*api.Broker, error) {
	brokerList, err := c.client.FetchBrokers()
	if err != nil {
		return nil, errors.Wrap(err, "select broker")
	}

	if len(*brokerList) == 0 {
		return nil, errors.New("no brokers returned from API")
	}

	validBrokers := make(map[string]api.Broker)
	haveEnterprise := false
	threshold := 10 * time.Second

	for _, broker := range *brokerList {
		broker := broker
		dur, ok := c.isValidBroker(&broker, checkType)
		if ok {
			if dur > threshold {
				continue
			} else if dur == threshold {
				validBrokers[broker.CID] = broker
			} else if dur < threshold {
				validBrokers = make(map[string]api.Broker)
				haveEnterprise = false
				threshold = dur
				validBrokers[broker.CID] = broker
			}
			if broker.Type == "enterprise" {
				haveEnterprise = true
			}
		}
	}

	if haveEnterprise { // eliminate non-enterprise brokers from valid brokers
		for k, v := range validBrokers {
			if v.Type != "enterprise" {
				delete(validBrokers, k)
			}
		}
	}

	if len(validBrokers) == 0 {
		return nil, errors.Errorf("found %d broker(s), zero are valid", len(*brokerList))
	}

	var selectedBroker api.Broker
	validBrokerKeys := reflect.ValueOf(validBrokers).MapKeys()
	if len(validBrokerKeys) == 1 {
		selectedBroker = validBrokers[validBrokerKeys[0].String()]
	} else {
		selectedBroker = validBrokers[validBrokerKeys[rand.Intn(len(validBrokerKeys))].String()]
	}

	c.logger.Debug().Str("broker", selectedBroker.Name).Msg("selected")

	return &selectedBroker, nil
}

// Is the broker valid (active, supports check type, and reachable)
func (c *Check) isValidBroker(broker *api.Broker, checkType string) (time.Duration, bool) {
	var brokerHost string
	var brokerPort string
	var connDuration time.Duration
	valid := false

	for _, detail := range broker.Details {
		detail := detail

		// broker must be active
		if detail.Status != c.statusActiveBroker {
			c.logger.Debug().Str("broker", broker.Name).Msg("not active, skipping")
			continue
		}

		// broker must have module loaded for the check type to be used
		if !brokerSupportsCheckType(checkType, &detail) {
			c.logger.Debug().Str("broker", broker.Name).Str("type", checkType).Msg("unsupported check type, skipping")
			continue
		}

		if detail.ExternalPort != 0 {
			brokerPort = strconv.Itoa(int(detail.ExternalPort))
		} else {
			if *detail.Port != 0 {
				brokerPort = strconv.Itoa(int(*detail.Port))
			} else {
				brokerPort = "43191"
			}
		}

		if detail.ExternalHost != nil && *detail.ExternalHost != "" {
			brokerHost = *detail.ExternalHost
		} else {
			brokerHost = *detail.IP
		}

		if brokerHost == "trap.noit.circonus.net" && brokerPort != "443" {
			brokerPort = "443"
		}

		minDelay := int(200 * time.Millisecond)
		maxDelay := int(2 * time.Second)

		for attempt := 1; attempt <= c.brokerMaxRetries; attempt++ {
			start := time.Now()
			// broker must be reachable and respond within designated time
			conn, err := net.DialTimeout("tcp", net.JoinHostPort(brokerHost, brokerPort), c.brokerMaxResponseTime)
			if err == nil {
				connDuration = time.Since(start)
				conn.Close()
				valid = true
				break
			}

			delay := time.Duration(rand.Intn(maxDelay-minDelay) + minDelay)

			c.logger.Warn().
				Err(err).
				Str("delay", delay.String()).
				Str("broker", broker.Name).
				Int("attempt", attempt).
				Int("retries", c.brokerMaxRetries).
				Msg("unable to connect, retrying")

			time.Sleep(delay)
		}

		if valid {
			c.logger.Debug().Str("broker", broker.Name).Msg("valid")
			break
		}
	}

	return connDuration, valid
}

// brokerSupportsCheckType verifies a broker supports the check type to be used
func brokerSupportsCheckType(checkType string, details *api.BrokerDetail) bool {
	baseType := string(checkType)

	if idx := strings.Index(baseType, ":"); idx > 0 {
		baseType = baseType[0:idx]
	}

	for _, module := range details.Modules {
		if module == baseType {
			return true
		}
	}

	return false
}
<|endoftext|>