<commit_before>package server

import (
	"crypto/rand"
	"encoding/hex"
	"errors"
	"fmt"
	"net"
	"strconv"

	"github.com/g8os/blockstor/redisstub"
	"github.com/g8os/blockstor/tlog"
	log "github.com/glendc/go-mini-log"
)

// DefaultConfig creates a new config, using sane defaults
func DefaultConfig() *Config {
	privKey := make([]byte, tlog.KeySize)
	rand.Read(privKey)

	nonce := make([]byte, tlog.NonceSize)
	rand.Read(nonce)
	hexNonce := hex.EncodeToString(nonce)

	return &Config{
		K:          4,
		M:          2,
		ListenAddr: "", // allocate dynamically
		FlushSize:  25,
		FlushTime:  25,
		PrivKey:    string(privKey),
		HexNonce:   hexNonce,
	}
}

// Config used for creating the tlogserver
type Config struct {
	K                 int
	M                 int
	ListenAddr        string
	FlushSize         int
	FlushTime         int
	PrivKey           string
	HexNonce          string
	ObjStoreAddresses []string
}

// ValidateAndCreateObjStoreAddresses checks if the config
// contains exactly as many object store server addresses as expected.
// If more are given, an error is returned.
// If none are given, memorystubs can be created if `allowInMemoryStubs` is true, otherwise an error is returned.
// If less are given, the missing addresses are generated, using consecutive ports,
// starting with the port following the port of the last given object store server address
func (cfg *Config) ValidateAndCreateObjStoreAddresses(allowInMemoryStubs bool) error {
	length := len(cfg.ObjStoreAddresses)
	expectedLength := cfg.ObjStorServerCount()

	if length == 0 {
		if !allowInMemoryStubs {
			return errors.New("no objstore servers available, and inmemory stubs aren't allowed")
		}

		// create in-memory ledisdb servers
		// if no object store address is given
		cfg.ObjStoreAddresses = make([]string, expectedLength)
		for i := range cfg.ObjStoreAddresses {
			// create in-memory redis
			objstor := redisstub.NewMemoryRedis()
			cfg.ObjStoreAddresses[i] = objstor.Address()

			// listen to in-memory redis on new goroutine
			//defer objstor.Close()
			go objstor.Listen()
		}
	} else if length < expectedLength {
		// parse last given address's host and port
		lastAddr := cfg.ObjStoreAddresses[length-1]
		host, rawPort, err := net.SplitHostPort(lastAddr)
		if err != nil {
			return fmt.Errorf("objstor address %s is illegal: %v", lastAddr, err)
		}
		port, err := strconv.Atoi(rawPort)
		if err != nil {
			return fmt.Errorf("objstor address %s is illegal: %v", lastAddr, err)
		}
		// add the missing servers dynamically
		// (it is assumed that the missing servers are live on the ports
		//  following the last server's port)
		port++
		portBound := port + (expectedLength - length)
		for ; port < portBound; port++ {
			addr := net.JoinHostPort(host, strconv.Itoa(port))
			log.Debug("add missing objstor server address:", addr)
			cfg.ObjStoreAddresses = append(cfg.ObjStoreAddresses, addr)
		}
	} else if length > expectedLength {
		return fmt.Errorf(
			"too many objstor servers given, only %d are required",
			expectedLength)
	}

	return nil
}

// ObjStorServerCount returns the amount of
// Object Store Servers are required/available/expected.
func (cfg *Config) ObjStorServerCount() int {
	return cfg.K + cfg.M + 1
}

// ObjStoreServerAddresses returns the objstore server addresses,
// iff the config defines exactly as many of the expected server addresses
func (cfg *Config) ObjStoreServerAddresses() ([]string, error) {
	if expected := cfg.ObjStorServerCount(); len(cfg.ObjStoreAddresses) != expected {
		return nil, fmt.Errorf(
			"expected %d objstor servers to be available, while %d servers are available",
			expected, len(cfg.ObjStoreAddresses))
	}

	return cfg.ObjStoreAddresses, nil
}
<commit_msg>tlog : make the default config to not return random values.<commit_after>package server

import (
	"errors"
	"fmt"
	"net"
	"strconv"

	"github.com/g8os/blockstor/redisstub"
	log "github.com/glendc/go-mini-log"
)

// DefaultConfig creates a new config, using sane defaults
func DefaultConfig() *Config {
	return &Config{
		K:          4,
		M:          2,
		ListenAddr: "0.0.0.0:11211",
		FlushSize:  25,
		FlushTime:  25,
		PrivKey:    "12345678901234567890123456789012",
		HexNonce:   "37b8e8a308c354048d245f6d",
	}
}

// Config used for creating the tlogserver
type Config struct {
	K                 int
	M                 int
	ListenAddr        string
	FlushSize         int
	FlushTime         int
	PrivKey           string
	HexNonce          string
	ObjStoreAddresses []string
}

// ValidateAndCreateObjStoreAddresses checks if the config
// contains exactly as many object store server addresses as expected.
// If more are given, an error is returned.
// If none are given, memorystubs can be created if `allowInMemoryStubs` is true, otherwise an error is returned.
// If less are given, the missing addresses are generated, using consecutive ports,
// starting with the port following the port of the last given object store server address
func (cfg *Config) ValidateAndCreateObjStoreAddresses(allowInMemoryStubs bool) error {
	length := len(cfg.ObjStoreAddresses)
	expectedLength := cfg.ObjStorServerCount()

	if length == 0 {
		if !allowInMemoryStubs {
			return errors.New("no objstore servers available, and inmemory stubs aren't allowed")
		}

		// create in-memory ledisdb servers
		// if no object store address is given
		cfg.ObjStoreAddresses = make([]string, expectedLength)
		for i := range cfg.ObjStoreAddresses {
			// create in-memory redis
			objstor := redisstub.NewMemoryRedis()
			cfg.ObjStoreAddresses[i] = objstor.Address()

			// listen to in-memory redis on new goroutine
			//defer objstor.Close()
			go objstor.Listen()
		}
	} else if length < expectedLength {
		// parse last given address's host and port
		lastAddr := cfg.ObjStoreAddresses[length-1]
		host, rawPort, err := net.SplitHostPort(lastAddr)
		if err != nil {
			return fmt.Errorf("objstor address %s is illegal: %v", lastAddr, err)
		}
		port, err := strconv.Atoi(rawPort)
		if err != nil {
			return fmt.Errorf("objstor address %s is illegal: %v", lastAddr, err)
		}
		// add the missing servers dynamically
		// (it is assumed that the missing servers are live on the ports
		//  following the last server's port)
		port++
		portBound := port + (expectedLength - length)
		for ; port < portBound; port++ {
			addr := net.JoinHostPort(host, strconv.Itoa(port))
			log.Debug("add missing objstor server address:", addr)
			cfg.ObjStoreAddresses = append(cfg.ObjStoreAddresses, addr)
		}
	} else if length > expectedLength {
		return fmt.Errorf(
			"too many objstor servers given, only %d are required",
			expectedLength)
	}

	return nil
}

// ObjStorServerCount returns the amount of
// Object Store Servers are required/available/expected.
func (cfg *Config) ObjStorServerCount() int {
	return cfg.K + cfg.M + 1
}

// ObjStoreServerAddresses returns the objstore server addresses,
// iff the config defines exactly as many of the expected server addresses
func (cfg *Config) ObjStoreServerAddresses() ([]string, error) {
	if expected := cfg.ObjStorServerCount(); len(cfg.ObjStoreAddresses) != expected {
		return nil, fmt.Errorf(
			"expected %d objstor servers to be available, while %d servers are available",
			expected, len(cfg.ObjStoreAddresses))
	}

	return cfg.ObjStoreAddresses, nil
}
<|endoftext|>


<commit_before>package cmd

import (
	"fmt"
	"io"
	"io/ioutil"
	"os"

	"github.com/MakeNowJust/heredoc"
	"github.com/spf13/cobra"

	kcmdutil "k8s.io/kubernetes/pkg/kubectl/cmd/util"

	buildapi "github.com/openshift/origin/pkg/build/api"
	"github.com/openshift/origin/pkg/cmd/templates"
	cmdutil "github.com/openshift/origin/pkg/cmd/util"
	"github.com/openshift/origin/pkg/cmd/util/clientcmd"
	configcmd "github.com/openshift/origin/pkg/config/cmd"
	newapp "github.com/openshift/origin/pkg/generate/app"
	newcmd "github.com/openshift/origin/pkg/generate/app/cmd"
)

// NewBuildRecommendedCommandName is the recommended command name.
const NewBuildRecommendedCommandName = "new-build"

var (
	newBuildLong = templates.LongDesc(`
		Create a new build by specifying source code

		This command will try to create a build configuration for your application using images and
		code that has a public repository. It will lookup the images on the local Docker installation
		(if available), a Docker registry, or an image stream.

		If you specify a source code URL, it will set up a build that takes your source code and converts
		it into an image that can run inside of a pod. Local source must be in a git repository that has a
		remote repository that the server can see.

		Once the build configuration is created a new build will be automatically triggered.
		You can use '%[1]s status' to check the progress.`)

	newBuildExample = templates.Examples(`
	  # Create a build config based on the source code in the current git repository (with a public
	  # remote) and a Docker image
	  %[1]s %[2]s . --docker-image=repo/langimage

	  # Create a NodeJS build config based on the provided [image]~[source code] combination
	  %[1]s %[2]s openshift/nodejs-010-centos7~https://github.com/openshift/nodejs-ex.git

	  # Create a build config from a remote repository using its beta2 branch
	  %[1]s %[2]s https://github.com/openshift/ruby-hello-world#beta2

	  # Create a build config using a Dockerfile specified as an argument
	  %[1]s %[2]s -D $'FROM centos:7\nRUN yum install -y httpd'

	  # Create a build config from a remote repository and add custom environment variables
	  %[1]s %[2]s https://github.com/openshift/ruby-hello-world RACK_ENV=development

	  # Create a build config from a remote repository and inject the npmrc into a build
	  %[1]s %[2]s https://github.com/openshift/ruby-hello-world --build-secret npmrc:.npmrc

	  # Create a build config that gets its input from a remote repository and another Docker image
	  %[1]s %[2]s https://github.com/openshift/ruby-hello-world --source-image=openshift/jenkins-1-centos7 --source-image-path=/var/lib/jenkins:tmp`)

	newBuildNoInput = `You must specify one or more images, image streams, or source code locations to create a build.

To build from an existing image stream tag or Docker image, provide the name of the image and
the source code location:

  %[1]s %[2]s openshift/nodejs-010-centos7~https://github.com/openshift/nodejs-ex.git

If you only specify the source repository location (local or remote), the command will look at
the repo to determine the type, and then look for a matching image on your server or on the
default Docker registry.

  %[1]s %[2]s https://github.com/openshift/nodejs-ex.git

will look for an image called "nodejs" in your current project, the 'openshift' project, or
on the Docker Hub.
`
)

type NewBuildOptions struct {
	*ObjectGeneratorOptions
}

// NewCmdNewBuild implements the OpenShift cli new-build command
func NewCmdNewBuild(name, baseName string, f *clientcmd.Factory, in io.Reader, out, errout io.Writer) *cobra.Command {
	config := newcmd.NewAppConfig()
	config.ExpectToBuild = true
	o := &NewBuildOptions{ObjectGeneratorOptions: &ObjectGeneratorOptions{Config: config}}

	cmd := &cobra.Command{
		Use:        fmt.Sprintf("%s (IMAGE | IMAGESTREAM | PATH | URL ...)", name),
		Short:      "Create a new build configuration",
		Long:       fmt.Sprintf(newBuildLong, baseName, name),
		Example:    fmt.Sprintf(newBuildExample, baseName, name),
		SuggestFor: []string{"build", "builds"},
		Run: func(c *cobra.Command, args []string) {
			kcmdutil.CheckErr(o.Complete(baseName, name, f, c, args, in, out, errout))
			err := o.RunNewBuild()
			if err == cmdutil.ErrExit {
				os.Exit(1)
			}
			kcmdutil.CheckErr(err)
		},
	}

	cmd.Flags().StringSliceVar(&config.SourceRepositories, "code", config.SourceRepositories, "Source code in the build configuration.")
	cmd.Flags().StringSliceVarP(&config.ImageStreams, "image", "", config.ImageStreams, "Name of an image stream to to use as a builder. (deprecated)")
	cmd.Flags().MarkDeprecated("image", "use --image-stream instead")
	cmd.Flags().StringSliceVarP(&config.ImageStreams, "image-stream", "i", config.ImageStreams, "Name of an image stream to to use as a builder.")
	cmd.Flags().StringSliceVar(&config.DockerImages, "docker-image", config.DockerImages, "Name of a Docker image to use as a builder.")
	cmd.Flags().StringSliceVar(&config.Secrets, "build-secret", config.Secrets, "Secret and destination to use as an input for the build.")
	cmd.Flags().StringVar(&config.Name, "name", "", "Set name to use for generated build artifacts.")
	cmd.Flags().StringVar(&config.To, "to", "", "Push built images to this image stream tag (or Docker image repository if --to-docker is set).")
	cmd.Flags().BoolVar(&config.OutputDocker, "to-docker", false, "If true, have the build output push to a Docker repository.")
	cmd.Flags().StringArrayVar(&config.BuildEnvironment, "build-env", config.BuildEnvironment, "Specify a key-value pair for an environment variable to set into resulting image.")
	cmd.Flags().MarkHidden("build-env")
	cmd.Flags().StringArrayVarP(&config.BuildEnvironment, "env", "e", config.BuildEnvironment, "Specify a key-value pair for an environment variable to set into resulting image.")
	cmd.Flags().StringArrayVar(&config.BuildEnvironmentFiles, "build-env-file", config.BuildEnvironmentFiles, "File containing key-value pairs of environment variables to set into each container.")
	cmd.MarkFlagFilename("build-env-file")
	cmd.Flags().MarkHidden("build-env-file")
	cmd.Flags().StringArrayVar(&config.BuildEnvironmentFiles, "env-file", config.BuildEnvironmentFiles, "File containing key-value pairs of environment variables to set into each container.")
	cmd.MarkFlagFilename("env-file")
	cmd.Flags().Var(&config.Strategy, "strategy", "Specify the build strategy to use if you don't want to detect (docker|pipeline|source).")
	cmd.Flags().StringVarP(&config.Dockerfile, "dockerfile", "D", "", "Specify the contents of a Dockerfile to build directly, implies --strategy=docker. Pass '-' to read from STDIN.")
	cmd.Flags().StringArrayVar(&config.BuildArgs, "build-arg", config.BuildArgs, "Specify a key-value pair to pass to Docker during the build.")
	cmd.Flags().BoolVar(&config.BinaryBuild, "binary", false, "Instead of expecting a source URL, set the build to expect binary contents. Will disable triggers.")
	cmd.Flags().StringP("labels", "l", "", "Label to set in all generated resources.")
	cmd.Flags().BoolVar(&config.AllowMissingImages, "allow-missing-images", false, "If true, indicates that referenced Docker images that cannot be found locally or in a registry should still be used.")
	cmd.Flags().BoolVar(&config.AllowMissingImageStreamTags, "allow-missing-imagestream-tags", false, "If true, indicates that image stream tags that don't exist should still be used.")
	cmd.Flags().StringVar(&config.ContextDir, "context-dir", "", "Context directory to be used for the build.")
	cmd.Flags().BoolVar(&config.NoOutput, "no-output", false, "If true, the build output will not be pushed anywhere.")
	cmd.Flags().StringVar(&config.SourceImage, "source-image", "", "Specify an image to use as source for the build.  You must also specify --source-image-path.")
	cmd.Flags().StringVar(&config.SourceImagePath, "source-image-path", "", "Specify the file or directory to copy from the source image and its destination in the build directory. Format: [source]:[destination-dir].")

	o.Action.BindForOutput(cmd.Flags())
	cmd.Flags().String("output-version", "", "The preferred API versions of the output objects")

	return cmd
}

// Complete sets any default behavior for the command
func (o *NewBuildOptions) Complete(baseName, commandName string, f *clientcmd.Factory, c *cobra.Command, args []string, in io.Reader, out, errout io.Writer) error {
	bo := o.ObjectGeneratorOptions
	err := bo.Complete(baseName, commandName, f, c, args, in, out, errout)
	if err != nil {
		return err
	}

	bo.Config.AllowNonNumericExposedPorts = true
	if bo.Config.Dockerfile == "-" {
		data, err := ioutil.ReadAll(in)
		if err != nil {
			return err
		}
		bo.Config.Dockerfile = string(data)
	}

	return nil
}

// RunNewBuild contains all the necessary functionality for the OpenShift cli new-build command
func (o *NewBuildOptions) RunNewBuild() error {
	config := o.Config
	out := o.Action.Out

	checkGitInstalled(out)

	result, err := config.Run()
	if err != nil {
		return handleError(err, o.BaseName, o.CommandName, o.CommandPath, config, transformBuildError)
	}

	if len(config.Labels) == 0 && len(result.Name) > 0 {
		config.Labels = map[string]string{"build": result.Name}
	}

	if err := setLabels(config.Labels, result); err != nil {
		return err
	}
	if err := setAnnotations(map[string]string{newcmd.GeneratedByNamespace: newcmd.GeneratedByNewBuild}, result); err != nil {
		return err
	}

	if o.Action.ShouldPrint() {
		return o.PrintObject(result.List)
	}

	if errs := o.Action.WithMessage(configcmd.CreateMessage(config.Labels), "created").Run(result.List, result.Namespace); len(errs) > 0 {
		return cmdutil.ErrExit
	}

	if !o.Action.Verbose() || o.Action.DryRun {
		return nil
	}

	indent := o.Action.DefaultIndent()
	for _, item := range result.List.Items {
		switch t := item.(type) {
		case *buildapi.BuildConfig:
			if len(t.Spec.Triggers) > 0 && t.Spec.Source.Binary == nil {
				fmt.Fprintf(out, "%sBuild configuration %q created and build triggered.\n", indent, t.Name)
				fmt.Fprintf(out, "%sRun '%s logs -f bc/%s' to stream the build progress.\n", indent, o.BaseName, t.Name)
			}
		}
	}

	return nil
}

func transformBuildError(err error, baseName, commandName, commandPath string, groups errorGroups) {
	switch t := err.(type) {
	case newapp.ErrNoMatch:
		groups.Add(
			"no-matches",
			heredoc.Docf(`
				The '%[1]s' command will match arguments to the following types:

				  1. Images tagged into image streams in the current project or the 'openshift' project
				     - if you don't specify a tag, we'll add ':latest'
				  2. Images in the Docker Hub, on remote registries, or on the local Docker engine
				  3. Git repository URLs or local paths that point to Git repositories

				--allow-missing-images can be used to force the use of an image that was not matched

				See '%[1]s -h' for examples.`, commandPath,
			),
			t,
			t.Errs...,
		)
		return
	}
	switch err {
	case newcmd.ErrNoInputs:
		groups.Add("", "", usageError(commandPath, newBuildNoInput, baseName, commandName))
		return
	}
	transformRunError(err, baseName, commandName, commandPath, groups)
}
<commit_msg>add env flag to  new-build example<commit_after>package cmd

import (
	"fmt"
	"io"
	"io/ioutil"
	"os"

	"github.com/MakeNowJust/heredoc"
	"github.com/spf13/cobra"

	kcmdutil "k8s.io/kubernetes/pkg/kubectl/cmd/util"

	buildapi "github.com/openshift/origin/pkg/build/api"
	"github.com/openshift/origin/pkg/cmd/templates"
	cmdutil "github.com/openshift/origin/pkg/cmd/util"
	"github.com/openshift/origin/pkg/cmd/util/clientcmd"
	configcmd "github.com/openshift/origin/pkg/config/cmd"
	newapp "github.com/openshift/origin/pkg/generate/app"
	newcmd "github.com/openshift/origin/pkg/generate/app/cmd"
)

// NewBuildRecommendedCommandName is the recommended command name.
const NewBuildRecommendedCommandName = "new-build"

var (
	newBuildLong = templates.LongDesc(`
		Create a new build by specifying source code

		This command will try to create a build configuration for your application using images and
		code that has a public repository. It will lookup the images on the local Docker installation
		(if available), a Docker registry, or an image stream.

		If you specify a source code URL, it will set up a build that takes your source code and converts
		it into an image that can run inside of a pod. Local source must be in a git repository that has a
		remote repository that the server can see.

		Once the build configuration is created a new build will be automatically triggered.
		You can use '%[1]s status' to check the progress.`)

	newBuildExample = templates.Examples(`
	  # Create a build config based on the source code in the current git repository (with a public
	  # remote) and a Docker image
	  %[1]s %[2]s . --docker-image=repo/langimage

	  # Create a NodeJS build config based on the provided [image]~[source code] combination
	  %[1]s %[2]s openshift/nodejs-010-centos7~https://github.com/openshift/nodejs-ex.git

	  # Create a build config from a remote repository using its beta2 branch
	  %[1]s %[2]s https://github.com/openshift/ruby-hello-world#beta2

	  # Create a build config using a Dockerfile specified as an argument
	  %[1]s %[2]s -D $'FROM centos:7\nRUN yum install -y httpd'

	  # Create a build config from a remote repository and add custom environment variables
	  %[1]s %[2]s https://github.com/openshift/ruby-hello-world -e RACK_ENV=development

	  # Create a build config from a remote repository and inject the npmrc into a build
	  %[1]s %[2]s https://github.com/openshift/ruby-hello-world --build-secret npmrc:.npmrc

	  # Create a build config that gets its input from a remote repository and another Docker image
	  %[1]s %[2]s https://github.com/openshift/ruby-hello-world --source-image=openshift/jenkins-1-centos7 --source-image-path=/var/lib/jenkins:tmp`)

	newBuildNoInput = `You must specify one or more images, image streams, or source code locations to create a build.

To build from an existing image stream tag or Docker image, provide the name of the image and
the source code location:

  %[1]s %[2]s openshift/nodejs-010-centos7~https://github.com/openshift/nodejs-ex.git

If you only specify the source repository location (local or remote), the command will look at
the repo to determine the type, and then look for a matching image on your server or on the
default Docker registry.

  %[1]s %[2]s https://github.com/openshift/nodejs-ex.git

will look for an image called "nodejs" in your current project, the 'openshift' project, or
on the Docker Hub.
`
)

type NewBuildOptions struct {
	*ObjectGeneratorOptions
}

// NewCmdNewBuild implements the OpenShift cli new-build command
func NewCmdNewBuild(name, baseName string, f *clientcmd.Factory, in io.Reader, out, errout io.Writer) *cobra.Command {
	config := newcmd.NewAppConfig()
	config.ExpectToBuild = true
	o := &NewBuildOptions{ObjectGeneratorOptions: &ObjectGeneratorOptions{Config: config}}

	cmd := &cobra.Command{
		Use:        fmt.Sprintf("%s (IMAGE | IMAGESTREAM | PATH | URL ...)", name),
		Short:      "Create a new build configuration",
		Long:       fmt.Sprintf(newBuildLong, baseName, name),
		Example:    fmt.Sprintf(newBuildExample, baseName, name),
		SuggestFor: []string{"build", "builds"},
		Run: func(c *cobra.Command, args []string) {
			kcmdutil.CheckErr(o.Complete(baseName, name, f, c, args, in, out, errout))
			err := o.RunNewBuild()
			if err == cmdutil.ErrExit {
				os.Exit(1)
			}
			kcmdutil.CheckErr(err)
		},
	}

	cmd.Flags().StringSliceVar(&config.SourceRepositories, "code", config.SourceRepositories, "Source code in the build configuration.")
	cmd.Flags().StringSliceVarP(&config.ImageStreams, "image", "", config.ImageStreams, "Name of an image stream to to use as a builder. (deprecated)")
	cmd.Flags().MarkDeprecated("image", "use --image-stream instead")
	cmd.Flags().StringSliceVarP(&config.ImageStreams, "image-stream", "i", config.ImageStreams, "Name of an image stream to to use as a builder.")
	cmd.Flags().StringSliceVar(&config.DockerImages, "docker-image", config.DockerImages, "Name of a Docker image to use as a builder.")
	cmd.Flags().StringSliceVar(&config.Secrets, "build-secret", config.Secrets, "Secret and destination to use as an input for the build.")
	cmd.Flags().StringVar(&config.Name, "name", "", "Set name to use for generated build artifacts.")
	cmd.Flags().StringVar(&config.To, "to", "", "Push built images to this image stream tag (or Docker image repository if --to-docker is set).")
	cmd.Flags().BoolVar(&config.OutputDocker, "to-docker", false, "If true, have the build output push to a Docker repository.")
	cmd.Flags().StringArrayVar(&config.BuildEnvironment, "build-env", config.BuildEnvironment, "Specify a key-value pair for an environment variable to set into resulting image.")
	cmd.Flags().MarkHidden("build-env")
	cmd.Flags().StringArrayVarP(&config.BuildEnvironment, "env", "e", config.BuildEnvironment, "Specify a key-value pair for an environment variable to set into resulting image.")
	cmd.Flags().StringArrayVar(&config.BuildEnvironmentFiles, "build-env-file", config.BuildEnvironmentFiles, "File containing key-value pairs of environment variables to set into each container.")
	cmd.MarkFlagFilename("build-env-file")
	cmd.Flags().MarkHidden("build-env-file")
	cmd.Flags().StringArrayVar(&config.BuildEnvironmentFiles, "env-file", config.BuildEnvironmentFiles, "File containing key-value pairs of environment variables to set into each container.")
	cmd.MarkFlagFilename("env-file")
	cmd.Flags().Var(&config.Strategy, "strategy", "Specify the build strategy to use if you don't want to detect (docker|pipeline|source).")
	cmd.Flags().StringVarP(&config.Dockerfile, "dockerfile", "D", "", "Specify the contents of a Dockerfile to build directly, implies --strategy=docker. Pass '-' to read from STDIN.")
	cmd.Flags().StringArrayVar(&config.BuildArgs, "build-arg", config.BuildArgs, "Specify a key-value pair to pass to Docker during the build.")
	cmd.Flags().BoolVar(&config.BinaryBuild, "binary", false, "Instead of expecting a source URL, set the build to expect binary contents. Will disable triggers.")
	cmd.Flags().StringP("labels", "l", "", "Label to set in all generated resources.")
	cmd.Flags().BoolVar(&config.AllowMissingImages, "allow-missing-images", false, "If true, indicates that referenced Docker images that cannot be found locally or in a registry should still be used.")
	cmd.Flags().BoolVar(&config.AllowMissingImageStreamTags, "allow-missing-imagestream-tags", false, "If true, indicates that image stream tags that don't exist should still be used.")
	cmd.Flags().StringVar(&config.ContextDir, "context-dir", "", "Context directory to be used for the build.")
	cmd.Flags().BoolVar(&config.NoOutput, "no-output", false, "If true, the build output will not be pushed anywhere.")
	cmd.Flags().StringVar(&config.SourceImage, "source-image", "", "Specify an image to use as source for the build.  You must also specify --source-image-path.")
	cmd.Flags().StringVar(&config.SourceImagePath, "source-image-path", "", "Specify the file or directory to copy from the source image and its destination in the build directory. Format: [source]:[destination-dir].")

	o.Action.BindForOutput(cmd.Flags())
	cmd.Flags().String("output-version", "", "The preferred API versions of the output objects")

	return cmd
}

// Complete sets any default behavior for the command
func (o *NewBuildOptions) Complete(baseName, commandName string, f *clientcmd.Factory, c *cobra.Command, args []string, in io.Reader, out, errout io.Writer) error {
	bo := o.ObjectGeneratorOptions
	err := bo.Complete(baseName, commandName, f, c, args, in, out, errout)
	if err != nil {
		return err
	}

	bo.Config.AllowNonNumericExposedPorts = true
	if bo.Config.Dockerfile == "-" {
		data, err := ioutil.ReadAll(in)
		if err != nil {
			return err
		}
		bo.Config.Dockerfile = string(data)
	}

	return nil
}

// RunNewBuild contains all the necessary functionality for the OpenShift cli new-build command
func (o *NewBuildOptions) RunNewBuild() error {
	config := o.Config
	out := o.Action.Out

	checkGitInstalled(out)

	result, err := config.Run()
	if err != nil {
		return handleError(err, o.BaseName, o.CommandName, o.CommandPath, config, transformBuildError)
	}

	if len(config.Labels) == 0 && len(result.Name) > 0 {
		config.Labels = map[string]string{"build": result.Name}
	}

	if err := setLabels(config.Labels, result); err != nil {
		return err
	}
	if err := setAnnotations(map[string]string{newcmd.GeneratedByNamespace: newcmd.GeneratedByNewBuild}, result); err != nil {
		return err
	}

	if o.Action.ShouldPrint() {
		return o.PrintObject(result.List)
	}

	if errs := o.Action.WithMessage(configcmd.CreateMessage(config.Labels), "created").Run(result.List, result.Namespace); len(errs) > 0 {
		return cmdutil.ErrExit
	}

	if !o.Action.Verbose() || o.Action.DryRun {
		return nil
	}

	indent := o.Action.DefaultIndent()
	for _, item := range result.List.Items {
		switch t := item.(type) {
		case *buildapi.BuildConfig:
			if len(t.Spec.Triggers) > 0 && t.Spec.Source.Binary == nil {
				fmt.Fprintf(out, "%sBuild configuration %q created and build triggered.\n", indent, t.Name)
				fmt.Fprintf(out, "%sRun '%s logs -f bc/%s' to stream the build progress.\n", indent, o.BaseName, t.Name)
			}
		}
	}

	return nil
}

func transformBuildError(err error, baseName, commandName, commandPath string, groups errorGroups) {
	switch t := err.(type) {
	case newapp.ErrNoMatch:
		groups.Add(
			"no-matches",
			heredoc.Docf(`
				The '%[1]s' command will match arguments to the following types:

				  1. Images tagged into image streams in the current project or the 'openshift' project
				     - if you don't specify a tag, we'll add ':latest'
				  2. Images in the Docker Hub, on remote registries, or on the local Docker engine
				  3. Git repository URLs or local paths that point to Git repositories

				--allow-missing-images can be used to force the use of an image that was not matched

				See '%[1]s -h' for examples.`, commandPath,
			),
			t,
			t.Errs...,
		)
		return
	}
	switch err {
	case newcmd.ErrNoInputs:
		groups.Add("", "", usageError(commandPath, newBuildNoInput, baseName, commandName))
		return
	}
	transformRunError(err, baseName, commandName, commandPath, groups)
}
<|endoftext|>


<commit_before>package cmd

import (
	"fmt"

	"github.com/dcos/dcos-cli/pkg/config"
	"github.com/spf13/cobra"
)

// configShowCmd represents the `dcos config show` subcommand.
var configShowCmd = &cobra.Command{
	Use:  "show",
	Args: cobra.MaximumNArgs(1),
	RunE: runConfigShowCmd,
}

var maskedKeys = []string{
	"core.dcos_acs_token",
}

func runConfigShowCmd(cmd *cobra.Command, args []string) error {
	store := attachedCluster().Config.Store()

	if len(args) == 1 {
		return configShow(store, args[0])
	}

	configShowAll(store)
	return nil
}

func configShowAll(store *config.Store) {
	for _, key := range store.Keys() {
		if val := store.Get(key); val != nil {
			for _, maskedKey := range maskedKeys {
				if maskedKey == key {
					val = "********"
				}
			}
			fmt.Printf("%s %v\n", key, val)
		}
	}
}

func configShow(store *config.Store, key string) error {
	if val := store.Get(key); val != nil {
		fmt.Printf("%v\n", val)
		return nil
	}
	return fmt.Errorf("unknown key \"%s\"", key)
}

func init() {
	configCmd.AddCommand(configShowCmd)
}
<commit_msg>Break once a masked key is found in `config show`<commit_after>package cmd

import (
	"fmt"

	"github.com/dcos/dcos-cli/pkg/config"
	"github.com/spf13/cobra"
)

// configShowCmd represents the `dcos config show` subcommand.
var configShowCmd = &cobra.Command{
	Use:  "show",
	Args: cobra.MaximumNArgs(1),
	RunE: runConfigShowCmd,
}

var maskedKeys = []string{
	"core.dcos_acs_token",
}

func runConfigShowCmd(cmd *cobra.Command, args []string) error {
	store := attachedCluster().Config.Store()

	if len(args) == 1 {
		return configShow(store, args[0])
	}

	configShowAll(store)
	return nil
}

func configShowAll(store *config.Store) {
	for _, key := range store.Keys() {
		if val := store.Get(key); val != nil {
			for _, maskedKey := range maskedKeys {
				if maskedKey == key {
					val = "********"
					break
				}
			}
			fmt.Printf("%s %v\n", key, val)
		}
	}
}

func configShow(store *config.Store, key string) error {
	if val := store.Get(key); val != nil {
		fmt.Printf("%v\n", val)
		return nil
	}
	return fmt.Errorf("unknown key \"%s\"", key)
}

func init() {
	configCmd.AddCommand(configShowCmd)
}
<|endoftext|>


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

package server

import (
	"fmt"
	"sync"
	"time"

	healthModels "github.com/cilium/cilium/api/v1/health/models"
	healthApi "github.com/cilium/cilium/api/v1/health/server"
	"github.com/cilium/cilium/api/v1/health/server/restapi"
	ciliumPkg "github.com/cilium/cilium/pkg/client"
	"github.com/cilium/cilium/pkg/health/defaults"
	"github.com/cilium/cilium/pkg/lock"
	"github.com/cilium/cilium/pkg/logging"
	"github.com/cilium/cilium/pkg/logging/logfields"

	"github.com/go-openapi/loads"
	"github.com/jessevdk/go-flags"
)

// AdminOption is an option for determining over which protocols the APIs are
// exposed.
type AdminOption string

const (
	// AdminOptionAny exposes every API over both Unix and HTTP sockets.
	AdminOptionAny AdminOption = "any"

	// AdminOptionUnix restricts most APIs to hosting over Unix sockets.
	AdminOptionUnix AdminOption = "unix"
)

var (
	log = logging.DefaultLogger.WithField(logfields.LogSubsys, "health-server")

	// PortToPaths is a convenience map for access to the ports and their
	// common string representations
	PortToPaths = map[int]string{
		defaults.HTTPPathPort: "Via L3",
	}

	// AdminOptions is the slice of all valid AdminOption values.
	AdminOptions = []AdminOption{
		AdminOptionAny,
		AdminOptionUnix,
	}
)

// Config stores the configuration data for a cilium-health server.
type Config struct {
	Debug         bool
	Passive       bool
	Admin         AdminOption
	CiliumURI     string
	ProbeInterval time.Duration
	ProbeDeadline time.Duration
}

// ipString is an IP address used as a more descriptive type name in maps.
type ipString string

// nodeMap maps IP addresses to healthNode objectss for convenient access to
// node information.
type nodeMap map[ipString]healthNode

// Server is the cilium-health daemon that is in charge of performing health
// and connectivity checks periodically, and serving the cilium-health API.
type Server struct {
	healthApi.Server  // Server to provide cilium-health API
	*ciliumPkg.Client // Client to "GET /healthz" on cilium daemon
	Config

	waitgroup  sync.WaitGroup      // Used to synchronize all goroutines
	tcpServers []*healthApi.Server // Servers for external pings
	startTime  time.Time

	// The lock protects against read and write access to the IP->Node map,
	// the list of statuses as most recently seen, and the last time a
	// probe was conducted.
	lock.RWMutex
	connectivity *healthReport
	localStatus  *healthModels.SelfStatus
}

// DumpUptime returns the time that this server has been running.
func (s *Server) DumpUptime() string {
	return time.Since(s.startTime).String()
}

// getNodes fetches the latest set of nodes in the cluster from the Cilium
// daemon, and updates the Server's 'nodes' map.
func (s *Server) getNodes() (nodeMap, error) {
	scopedLog := log
	if s.CiliumURI != "" {
		scopedLog = log.WithField("URI", s.CiliumURI)
	}
	scopedLog.Debug("Sending request for /healthz ...")

	resp, err := s.Daemon.GetHealthz(nil)
	if err != nil {
		return nil, fmt.Errorf("unable to get agent health: %s", err)
	}
	log.Debug("Got cilium /healthz")

	if resp == nil || resp.Payload == nil || resp.Payload.Cluster == nil {
		return nil, fmt.Errorf("received nil health response")
	}

	if resp.Payload.Cluster.Self != "" {
		s.localStatus = &healthModels.SelfStatus{
			Name: resp.Payload.Cluster.Self,
		}
	}

	nodes := make(nodeMap)
	for _, n := range resp.Payload.Cluster.Nodes {
		if n.PrimaryAddress != nil {
			if n.PrimaryAddress.IPV4 != nil {
				nodes[ipString(n.PrimaryAddress.IPV4.IP)] = NewHealthNode(n)
			}
			if n.PrimaryAddress.IPV6 != nil {
				nodes[ipString(n.PrimaryAddress.IPV6.IP)] = NewHealthNode(n)
			}
		}
		for _, addr := range n.SecondaryAddresses {
			nodes[ipString(addr.IP)] = NewHealthNode(n)
		}
		if n.HealthEndpointAddress != nil {
			if n.HealthEndpointAddress.IPV4 != nil {
				nodes[ipString(n.HealthEndpointAddress.IPV4.IP)] = NewHealthNode(n)
			}
			if n.HealthEndpointAddress.IPV6 != nil {
				nodes[ipString(n.HealthEndpointAddress.IPV6.IP)] = NewHealthNode(n)
			}
		}
	}
	return nodes, nil
}

// updateCluster makes the specified health report visible to the API.
//
// It only updates the server's API-visible health report if the provided
// report started after the current report.
func (s *Server) updateCluster(report *healthReport) {
	s.Lock()
	defer s.Unlock()

	if s.connectivity.startTime.Before(report.startTime) {
		s.connectivity = report
	}
}

// GetStatusResponse returns the most recent cluster connectivity status.
func (s *Server) GetStatusResponse() *healthModels.HealthStatusResponse {
	s.RLock()
	defer s.RUnlock()

	var name string
	// Check if localStatus is populated already. If not, the name is empty
	if s.localStatus != nil {
		name = s.localStatus.Name
	}

	return &healthModels.HealthStatusResponse{
		Local: &healthModels.SelfStatus{
			Name: name,
		},
		Nodes:     s.connectivity.nodes,
		Timestamp: s.connectivity.startTime.Format(time.RFC3339),
	}
}

// FetchStatusResponse updates the cluster with the latest set of nodes,
// runs a synchronous probe across the cluster, updates the connectivity cache
// and returns the results.
func (s *Server) FetchStatusResponse() (*healthModels.HealthStatusResponse, error) {
	nodes, err := s.getNodes()
	if err != nil {
		return nil, err
	}

	prober := newProber(s, nodes)
	if err := prober.Run(); err != nil {
		log.WithError(err).Info("Failed to run ping")
		return nil, err
	}
	log.Debug("Run complete")
	s.updateCluster(prober.getResults())

	return s.GetStatusResponse(), nil
}

// Run services that are not considered 'Passive': Actively probing other
// hosts and endpoints over ICMP and HTTP, and hosting a local Unix socket.
// Blocks indefinitely, or returns any errors that occur hosting the Unix
// socket API server.
func (s *Server) runActiveServices() error {
	// Run it once at the start so we get some initial status
	s.FetchStatusResponse()

	nodes, _ := s.getNodes()
	prober := newProber(s, nodes)
	prober.MaxRTT = s.ProbeInterval
	prober.OnIdle = func() {
		// Fetch results and update set of nodes to probe every
		// ProbeInterval
		s.updateCluster(prober.getResults())
		if nodes, err := s.getNodes(); err == nil {
			prober.setNodes(nodes)
		}
	}
	prober.RunLoop()

	err := s.Server.Serve()
	prober.Stop()

	return err
}

// Serve spins up the following goroutines:
// * TCP API Server: Responders to the health API "/hello" message, one per path
//
// Also, if "Passive" is not set in s.Config:
// * Prober: Periodically run pings across the cluster at a configured interval
//   and update the server's connectivity status cache.
// * Unix API Server: Handle all health API requests over a unix socket.
func (s *Server) Serve() (err error) {
	for i := range s.tcpServers {
		s.waitgroup.Add(1)
		srv := s.tcpServers[i]
		go func() {
			defer s.waitgroup.Done()
			srv.Serve()
		}()
	}

	if !s.Config.Passive {
		err = s.runActiveServices()
	}
	s.waitgroup.Wait()

	return err
}

// Shutdown server and clean up resources
func (s *Server) Shutdown() {
	for i := range s.tcpServers {
		s.tcpServers[i].Shutdown()
	}
	if !s.Config.Passive {
		s.Server.Shutdown()
	}
}

func enableAPI(opt AdminOption, tcpPort int) bool {
	switch opt {
	case AdminOptionAny:
		return true
	case AdminOptionUnix:
		return tcpPort == 0
	default:
		return false
	}
}

// newServer instantiates a new instance of the API that serves the health
// API on the specified port. If tcpPort is 0, then a unix socket is opened
// which serves the entire API. If a tcpPort is specified, then it returns
// a server which only answers get requests for the root URL "/".
func (s *Server) newServer(spec *loads.Document, tcpPort int) *healthApi.Server {
	api := restapi.NewCiliumHealthAPI(spec)
	api.Logger = log.Printf

	// /hello
	api.GetHelloHandler = NewGetHelloHandler(s)

	if enableAPI(s.Config.Admin, tcpPort) {
		api.GetHealthzHandler = NewGetHealthzHandler(s)
		api.ConnectivityGetStatusHandler = NewGetStatusHandler(s)
		api.ConnectivityPutStatusProbeHandler = NewPutStatusProbeHandler(s)
	}

	srv := healthApi.NewServer(api)
	if tcpPort == 0 {
		srv.EnabledListeners = []string{"unix"}
		srv.SocketPath = flags.Filename(defaults.SockPath)
	} else {
		srv.EnabledListeners = []string{"http"}
		srv.Port = tcpPort
		srv.Host = "" // FIXME: Allow binding to specific IPs
	}
	srv.ConfigureAPI()

	return srv
}

// NewServer creates a server to handle health requests.
func NewServer(config Config) (*Server, error) {
	server := &Server{
		startTime:    time.Now(),
		Config:       config,
		tcpServers:   []*healthApi.Server{},
		connectivity: &healthReport{},
	}

	swaggerSpec, err := loads.Analyzed(healthApi.SwaggerJSON, "")
	if err != nil {
		return nil, err
	}

	if !config.Passive {
		cl, err := ciliumPkg.NewClient(config.CiliumURI)
		if err != nil {
			return nil, err
		}

		server.Client = cl
		server.Server = *server.newServer(swaggerSpec, 0)
	}
	for port := range PortToPaths {
		srv := server.newServer(swaggerSpec, port)
		server.tcpServers = append(server.tcpServers, srv)
	}

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

package server

import (
	"fmt"
	"sync"
	"time"

	healthModels "github.com/cilium/cilium/api/v1/health/models"
	healthApi "github.com/cilium/cilium/api/v1/health/server"
	"github.com/cilium/cilium/api/v1/health/server/restapi"
	ciliumPkg "github.com/cilium/cilium/pkg/client"
	"github.com/cilium/cilium/pkg/health/defaults"
	"github.com/cilium/cilium/pkg/lock"
	"github.com/cilium/cilium/pkg/logging"
	"github.com/cilium/cilium/pkg/logging/logfields"

	"github.com/go-openapi/loads"
	"github.com/jessevdk/go-flags"
)

// AdminOption is an option for determining over which protocols the APIs are
// exposed.
type AdminOption string

const (
	// AdminOptionAny exposes every API over both Unix and HTTP sockets.
	AdminOptionAny AdminOption = "any"

	// AdminOptionUnix restricts most APIs to hosting over Unix sockets.
	AdminOptionUnix AdminOption = "unix"
)

var (
	log = logging.DefaultLogger.WithField(logfields.LogSubsys, "health-server")

	// PortToPaths is a convenience map for access to the ports and their
	// common string representations
	PortToPaths = map[int]string{
		defaults.HTTPPathPort: "Via L3",
	}

	// AdminOptions is the slice of all valid AdminOption values.
	AdminOptions = []AdminOption{
		AdminOptionAny,
		AdminOptionUnix,
	}
)

// Config stores the configuration data for a cilium-health server.
type Config struct {
	Debug         bool
	Passive       bool
	Admin         AdminOption
	CiliumURI     string
	ProbeInterval time.Duration
	ProbeDeadline time.Duration
}

// ipString is an IP address used as a more descriptive type name in maps.
type ipString string

// nodeMap maps IP addresses to healthNode objectss for convenient access to
// node information.
type nodeMap map[ipString]healthNode

// Server is the cilium-health daemon that is in charge of performing health
// and connectivity checks periodically, and serving the cilium-health API.
type Server struct {
	healthApi.Server  // Server to provide cilium-health API
	*ciliumPkg.Client // Client to "GET /healthz" on cilium daemon
	Config

	waitgroup  sync.WaitGroup      // Used to synchronize all goroutines
	tcpServers []*healthApi.Server // Servers for external pings
	startTime  time.Time

	// The lock protects against read and write access to the IP->Node map,
	// the list of statuses as most recently seen, and the last time a
	// probe was conducted.
	lock.RWMutex
	connectivity *healthReport
	localStatus  *healthModels.SelfStatus
}

// DumpUptime returns the time that this server has been running.
func (s *Server) DumpUptime() string {
	return time.Since(s.startTime).String()
}

// getNodes fetches the latest set of nodes in the cluster from the Cilium
// daemon, and updates the Server's 'nodes' map.
func (s *Server) getNodes() (nodeMap, error) {
	scopedLog := log
	if s.CiliumURI != "" {
		scopedLog = log.WithField("URI", s.CiliumURI)
	}
	scopedLog.Debug("Sending request for /healthz ...")

	resp, err := s.Daemon.GetHealthz(nil)
	if err != nil {
		return nil, fmt.Errorf("unable to get agent health: %s", err)
	}
	log.Debug("Got cilium /healthz")

	if resp == nil || resp.Payload == nil || resp.Payload.Cluster == nil {
		return nil, fmt.Errorf("received nil health response")
	}

	if resp.Payload.Cluster.Self != "" {
		s.localStatus = &healthModels.SelfStatus{
			Name: resp.Payload.Cluster.Self,
		}
	}

	nodes := make(nodeMap)
	for _, n := range resp.Payload.Cluster.Nodes {
		if n.PrimaryAddress != nil {
			if n.PrimaryAddress.IPV4 != nil {
				nodes[ipString(n.PrimaryAddress.IPV4.IP)] = NewHealthNode(n)
			}
			if n.PrimaryAddress.IPV6 != nil {
				nodes[ipString(n.PrimaryAddress.IPV6.IP)] = NewHealthNode(n)
			}
		}
		for _, addr := range n.SecondaryAddresses {
			nodes[ipString(addr.IP)] = NewHealthNode(n)
		}
		if n.HealthEndpointAddress != nil {
			if n.HealthEndpointAddress.IPV4 != nil {
				nodes[ipString(n.HealthEndpointAddress.IPV4.IP)] = NewHealthNode(n)
			}
			if n.HealthEndpointAddress.IPV6 != nil {
				nodes[ipString(n.HealthEndpointAddress.IPV6.IP)] = NewHealthNode(n)
			}
		}
	}
	return nodes, nil
}

// updateCluster makes the specified health report visible to the API.
//
// It only updates the server's API-visible health report if the provided
// report started after the current report.
func (s *Server) updateCluster(report *healthReport) {
	s.Lock()
	defer s.Unlock()

	if s.connectivity.startTime.Before(report.startTime) {
		s.connectivity = report
	}
}

// GetStatusResponse returns the most recent cluster connectivity status.
func (s *Server) GetStatusResponse() *healthModels.HealthStatusResponse {
	s.RLock()
	defer s.RUnlock()

	var name string
	// Check if localStatus is populated already. If not, the name is empty
	if s.localStatus != nil {
		name = s.localStatus.Name
	}

	return &healthModels.HealthStatusResponse{
		Local: &healthModels.SelfStatus{
			Name: name,
		},
		Nodes:     s.connectivity.nodes,
		Timestamp: s.connectivity.startTime.Format(time.RFC3339),
	}
}

// FetchStatusResponse updates the cluster with the latest set of nodes,
// runs a synchronous probe across the cluster, updates the connectivity cache
// and returns the results.
func (s *Server) FetchStatusResponse() (*healthModels.HealthStatusResponse, error) {
	nodes, err := s.getNodes()
	if err != nil {
		return nil, err
	}

	prober := newProber(s, nodes)
	if err := prober.Run(); err != nil {
		log.WithError(err).Info("Failed to run ping")
		return nil, err
	}
	log.Debug("Run complete")
	s.updateCluster(prober.getResults())

	return s.GetStatusResponse(), nil
}

// Run services that are not considered 'Passive': Actively probing other
// hosts and endpoints over ICMP and HTTP, and hosting a local Unix socket.
// Blocks indefinitely, or returns any errors that occur hosting the Unix
// socket API server.
func (s *Server) runActiveServices() error {
	// Run it once at the start so we get some initial status
	s.FetchStatusResponse()

	nodes, _ := s.getNodes()
	prober := newProber(s, nodes)
	prober.MaxRTT = s.ProbeInterval
	prober.OnIdle = func() {
		// Fetch results and update set of nodes to probe every
		// ProbeInterval
		s.updateCluster(prober.getResults())
		if nodes, err := s.getNodes(); err == nil {
			prober.setNodes(nodes)
		}
	}
	prober.RunLoop()
	defer prober.Stop()

	return s.Server.Serve()
}

// Serve spins up the following goroutines:
// * TCP API Server: Responders to the health API "/hello" message, one per path
//
// Also, if "Passive" is not set in s.Config:
// * Prober: Periodically run pings across the cluster at a configured interval
//   and update the server's connectivity status cache.
// * Unix API Server: Handle all health API requests over a unix socket.
func (s *Server) Serve() (err error) {
	for i := range s.tcpServers {
		s.waitgroup.Add(1)
		srv := s.tcpServers[i]
		go func() {
			defer s.waitgroup.Done()
			srv.Serve()
		}()
	}

	if !s.Config.Passive {
		err = s.runActiveServices()
	}
	s.waitgroup.Wait()

	return err
}

// Shutdown server and clean up resources
func (s *Server) Shutdown() {
	for i := range s.tcpServers {
		s.tcpServers[i].Shutdown()
	}
	if !s.Config.Passive {
		s.Server.Shutdown()
	}
}

func enableAPI(opt AdminOption, tcpPort int) bool {
	switch opt {
	case AdminOptionAny:
		return true
	case AdminOptionUnix:
		return tcpPort == 0
	default:
		return false
	}
}

// newServer instantiates a new instance of the API that serves the health
// API on the specified port. If tcpPort is 0, then a unix socket is opened
// which serves the entire API. If a tcpPort is specified, then it returns
// a server which only answers get requests for the root URL "/".
func (s *Server) newServer(spec *loads.Document, tcpPort int) *healthApi.Server {
	api := restapi.NewCiliumHealthAPI(spec)
	api.Logger = log.Printf

	// /hello
	api.GetHelloHandler = NewGetHelloHandler(s)

	if enableAPI(s.Config.Admin, tcpPort) {
		api.GetHealthzHandler = NewGetHealthzHandler(s)
		api.ConnectivityGetStatusHandler = NewGetStatusHandler(s)
		api.ConnectivityPutStatusProbeHandler = NewPutStatusProbeHandler(s)
	}

	srv := healthApi.NewServer(api)
	if tcpPort == 0 {
		srv.EnabledListeners = []string{"unix"}
		srv.SocketPath = flags.Filename(defaults.SockPath)
	} else {
		srv.EnabledListeners = []string{"http"}
		srv.Port = tcpPort
		srv.Host = "" // FIXME: Allow binding to specific IPs
	}
	srv.ConfigureAPI()

	return srv
}

// NewServer creates a server to handle health requests.
func NewServer(config Config) (*Server, error) {
	server := &Server{
		startTime:    time.Now(),
		Config:       config,
		tcpServers:   []*healthApi.Server{},
		connectivity: &healthReport{},
	}

	swaggerSpec, err := loads.Analyzed(healthApi.SwaggerJSON, "")
	if err != nil {
		return nil, err
	}

	if !config.Passive {
		cl, err := ciliumPkg.NewClient(config.CiliumURI)
		if err != nil {
			return nil, err
		}

		server.Client = cl
		server.Server = *server.newServer(swaggerSpec, 0)
	}
	for port := range PortToPaths {
		srv := server.newServer(swaggerSpec, port)
		server.tcpServers = append(server.tcpServers, srv)
	}

	return server, nil
}
<|endoftext|>


<commit_before>package model

import (
	"net/http"
	"testing"
	"time"

	"github.com/labstack/echo"
	sd "github.com/labstack/echo/engine/standard"
)

const (
	etcdAddr   = "http://192.168.70.13:2379"
	etcdPrefix = "/gateway2"
)

var (
	serverAddr    = "127.0.0.1:12345"
	apiURL        = "/api/test"
	apiMethod     = "GET"
	checkDuration = 3
	checkTimeout  = 2
	clusterName   = "app"
	lbName        = "ROUNDROBIN"
	sleep         = false
)

var rt *RouteTable

func createRouteTable(t *testing.T) {
	store, err := NewEtcdStore([]string{etcdAddr}, etcdPrefix)

	if nil != err {
		t.Fatalf("create etcd store err.addr:<%s>", err)
	}

	store.Clean()

	rt = NewRouteTable(store)
	time.Sleep(time.Second * 1)
}

func createLocalServer() {
	e := echo.New()

	e.Get("/check", func() echo.HandlerFunc {
		return func(c echo.Context) error {
			if sleep {
				time.Sleep(time.Second * time.Duration(checkTimeout+1))
			}

			return c.String(http.StatusOK, "OK")
		}
	}())

	e.Run(sd.New(serverAddr))
}

func waitNotify() {
	time.Sleep(time.Second * 1)
}

func TestCreateRouteTable(t *testing.T) {
	createRouteTable(t)
}

func TestEtcdWatchNewServer(t *testing.T) {
	go createLocalServer()

	server := &Server{
		Schema:          "http",
		Addr:            serverAddr,
		CheckPath:       "/check",
		CheckDuration:   checkDuration,
		CheckTimeout:    checkTimeout,
		MaxQPS:          1500,
		HalfToOpen:      10,
		HalfTrafficRate: 10,
		CloseCount:      100,
	}

	err := rt.store.SaveServer(server)

	if nil != err {
		t.Error("add server err.")
		return
	}

	waitNotify()

	if len(rt.svrs) != 1 {
		t.Errorf("expect:<1>, acture:<%d>", len(rt.svrs))
		return
	}

	if rt.svrs[serverAddr].lock == nil {
		t.Error("server init error.")
		return
	}
}

func TestServerCheckOk(t *testing.T) {
	time.Sleep(time.Second * time.Duration(checkDuration))

	if rt.svrs[serverAddr].Status == Down {
		t.Errorf("status check ok err.expect:<UP>, acture:<%v>", Down)
	}
}

func TestServerCheckTimeout(t *testing.T) {
	defer func() {
		sleep = false
	}()

	sleep = true
	time.Sleep(time.Second * time.Duration(checkDuration*2+1)) // 等待两个周期

	if rt.svrs[serverAddr].Status == Up {
		t.Errorf("status check timeout err.expect:<DOWN>, acture:<%v>", Up)
		return
	}
}

func TestServerCheckTimeoutRecovery(t *testing.T) {
	time.Sleep(time.Second * time.Duration(checkDuration*2+1)) // 等待两个周期

	if rt.svrs[serverAddr].Status == Down {
		t.Errorf("status check timeout recovery err.expect:<UP>, acture:<%v>", Up)
		return
	}
}

func TestEtcdWatchNewCluster(t *testing.T) {
	cluster := &Cluster{
		Name:   clusterName,
		LbName: lbName,
	}

	err := rt.store.SaveCluster(cluster)

	if nil != err {
		t.Error("add cluster err.")
		return
	}

	waitNotify()

	if len(rt.clusters) == 1 {
		return
	}

	t.Errorf("expect:<1>, acture:<%d>", len(rt.clusters))
}

func TestEtcdWatchNewBind(t *testing.T) {
	bind := &Bind{
		ClusterName: clusterName,
		ServerAddr:  serverAddr,
	}

	err := rt.store.SaveBind(bind)

	if nil != err {
		t.Error("add cluster err.")
		return
	}

	waitNotify()

	if len(rt.mapping) == 1 {
		return
	}

	t.Errorf("expect:<1>, acture:<%d>. %+v", len(rt.mapping), rt.mapping)
}

func TestEtcdWatchNewAPI(t *testing.T) {
	n := &Node{
		AttrName:    "test",
		ClusterName: clusterName,
	}

	err := rt.store.SaveAPI(&API{
		URL:    apiURL,
		Method: apiMethod,
		Nodes:  []*Node{n},
	})

	if nil != err {
		t.Error("add api err.")
		return
	}

	waitNotify()

	if len(rt.apis) == 1 {
		return
	}

	t.Errorf("expect:<1>, acture:<%d>", len(rt.apis))
}

func TestEtcdWatchUpdateServer(t *testing.T) {
	server := &Server{
		Schema:          "http",
		Addr:            serverAddr,
		CheckPath:       "/check",
		CheckDuration:   checkDuration,
		CheckTimeout:    checkTimeout * 2,
		MaxQPS:          3000,
		HalfToOpen:      100,
		HalfTrafficRate: 30,
		CloseCount:      200,
	}

	err := rt.store.UpdateServer(server)

	if nil != err {
		t.Error("update server err.")
		return
	}

	waitNotify()

	svr := rt.svrs[serverAddr]

	if svr.MaxQPS != server.MaxQPS {
		t.Errorf("MaxQPS expect:<%d>, acture:<%d>. ", server.MaxQPS, svr.MaxQPS)
		return
	}

	if svr.HalfToOpen != server.HalfToOpen {
		t.Errorf("HalfToOpen expect:<%d>, acture:<%d>. ", server.HalfToOpen, svr.HalfToOpen)
		return
	}

	if svr.HalfTrafficRate != server.HalfTrafficRate {
		t.Errorf("HalfTrafficRate expect:<%d>, acture:<%d>. ", server.HalfTrafficRate, svr.HalfTrafficRate)
		return
	}

	if svr.CloseCount != server.CloseCount {
		t.Errorf("CloseCount expect:<%d>, acture:<%d>. ", server.CloseCount, svr.CloseCount)
		return
	}

	if svr.CheckTimeout == server.CheckTimeout {
		t.Errorf("CheckTimeout expect:<%d>, acture:<%d>. ", svr.CheckTimeout, server.CheckTimeout)
		return
	}
}

func TestEtcdWatchUpdateCluster(t *testing.T) {
	cluster := &Cluster{
		Name:   clusterName,
		LbName: lbName,
	}

	err := rt.store.UpdateCluster(cluster)

	if nil != err {
		t.Error("update cluster err.")
		return
	}

	waitNotify()

	existCluster := rt.clusters[clusterName]

	if existCluster.LbName != cluster.LbName {
		t.Errorf("LbName expect:<%s>, acture:<%s>. ", cluster.LbName, existCluster.LbName)
		return
	}
}

func TestEtcdWatchUpdateAPI(t *testing.T) {
	n := &Node{
		AttrName:    "test",
		ClusterName: clusterName,
	}

	n2 := &Node{
		AttrName:    "tes2t",
		ClusterName: clusterName,
	}

	api := &API{
		URL:    apiURL,
		Method: apiMethod,
		Nodes:  []*Node{n, n2},
	}

	err := rt.store.UpdateAPI(api)

	if nil != err {
		t.Error("update api err.")
		return
	}

	waitNotify()

	existAPI, _ := rt.apis[getAPIKey(api.URL, api.Method)]

	if len(existAPI.Nodes) != len(api.Nodes) {
		t.Errorf("Nodes expect:<%s>, acture:<%s>. ", len(existAPI.Nodes), len(api.Nodes))
		return
	}
}

func TestEtcdWatchDeleteCluster(t *testing.T) {
	err := rt.store.DeleteCluster(clusterName)

	if nil != err {
		t.Error("delete cluster err.")
		return
	}

	waitNotify()

	if len(rt.clusters) != 0 {
		t.Errorf("clusters expect:<0>, acture:<%d>", len(rt.clusters))
		return
	}

	banded, _ := rt.mapping[serverAddr]

	if len(banded) != 0 {
		t.Errorf("banded expect:<0>, acture:<%d>", len(banded))
		return
	}
}

func TestEtcdWatchDeleteServer(t *testing.T) {
	err := rt.store.DeleteServer(serverAddr)

	if nil != err {
		t.Error("delete server err.")
		return
	}

	waitNotify()

	if len(rt.svrs) != 0 {
		t.Errorf("svrs expect:<0>, acture:<%d>", len(rt.svrs))
		return
	}

	if len(rt.mapping) != 0 {
		t.Errorf("mapping expect:<0>, acture:<%d>", len(rt.mapping))
		return
	}
}

func TestEtcdWatchDeleteAPI(t *testing.T) {
	err := rt.store.DeleteAPI(apiURL, apiMethod)

	if nil != err {
		t.Error("delete api err.")
		return
	}

	waitNotify()

	if len(rt.apis) != 0 {
		t.Errorf("apis expect:<0>, acture:<%d>", len(rt.apis))
		return
	}
}

func TestEtcdWatchNewRouting(t *testing.T) {
	r, err := NewRouting(`desc = "test"; deadline = 100; rule = ["$query_abc == 10", "$query_123 == 20"];`, clusterName, "")

	if nil != err {
		t.Error("add routing err.")
		return
	}

	err = rt.store.SaveRouting(r)

	if nil != err {
		t.Error("add routing err.")
		return
	}

	waitNotify()

	if len(rt.routings) == 1 {
		delete(rt.routings, r.ID)
		return
	}

	t.Errorf("expect:<1>, acture:<%d>", len(rt.routings))
}

func TestEtcdWatchDeleteRouting(t *testing.T) {
	r, err := NewRouting(`desc = "test"; deadline = 3; rule = ["$query_abc == 10", "$query_123 == 20"];`, clusterName, "")

	if nil != err {
		t.Error("add routing err.")
		return
	}

	err = rt.store.SaveRouting(r)

	if nil != err {
		t.Error("add routing err.")
		return
	}

	time.Sleep(time.Second * 30)

	if len(rt.routings) == 0 {
		return
	}

	t.Errorf("expect:<0>, acture:<%d>", len(rt.routings))
}
<commit_msg>fix warn about go_vet<commit_after>package model

import (
	"net/http"
	"testing"
	"time"

	"github.com/labstack/echo"
	sd "github.com/labstack/echo/engine/standard"
)

const (
	etcdAddr   = "http://192.168.70.13:2379"
	etcdPrefix = "/gateway2"
)

var (
	serverAddr    = "127.0.0.1:12345"
	apiURL        = "/api/test"
	apiMethod     = "GET"
	checkDuration = 3
	checkTimeout  = 2
	clusterName   = "app"
	lbName        = "ROUNDROBIN"
	sleep         = false
)

var rt *RouteTable

func createRouteTable(t *testing.T) {
	store, err := NewEtcdStore([]string{etcdAddr}, etcdPrefix)

	if nil != err {
		t.Fatalf("create etcd store err.addr:<%s>", err)
	}

	store.Clean()

	rt = NewRouteTable(store)
	time.Sleep(time.Second * 1)
}

func createLocalServer() {
	e := echo.New()

	e.Get("/check", func() echo.HandlerFunc {
		return func(c echo.Context) error {
			if sleep {
				time.Sleep(time.Second * time.Duration(checkTimeout+1))
			}

			return c.String(http.StatusOK, "OK")
		}
	}())

	e.Run(sd.New(serverAddr))
}

func waitNotify() {
	time.Sleep(time.Second * 1)
}

func TestCreateRouteTable(t *testing.T) {
	createRouteTable(t)
}

func TestEtcdWatchNewServer(t *testing.T) {
	go createLocalServer()

	server := &Server{
		Schema:          "http",
		Addr:            serverAddr,
		CheckPath:       "/check",
		CheckDuration:   checkDuration,
		CheckTimeout:    checkTimeout,
		MaxQPS:          1500,
		HalfToOpen:      10,
		HalfTrafficRate: 10,
		CloseCount:      100,
	}

	err := rt.store.SaveServer(server)

	if nil != err {
		t.Error("add server err.")
		return
	}

	waitNotify()

	if len(rt.svrs) != 1 {
		t.Errorf("expect:<1>, acture:<%d>", len(rt.svrs))
		return
	}

	if rt.svrs[serverAddr].lock == nil {
		t.Error("server init error.")
		return
	}
}

func TestServerCheckOk(t *testing.T) {
	time.Sleep(time.Second * time.Duration(checkDuration))

	if rt.svrs[serverAddr].Status == Down {
		t.Errorf("status check ok err.expect:<UP>, acture:<%v>", Down)
	}
}

func TestServerCheckTimeout(t *testing.T) {
	defer func() {
		sleep = false
	}()

	sleep = true
	time.Sleep(time.Second * time.Duration(checkDuration*2+1)) // 等待两个周期

	if rt.svrs[serverAddr].Status == Up {
		t.Errorf("status check timeout err.expect:<DOWN>, acture:<%v>", Up)
		return
	}
}

func TestServerCheckTimeoutRecovery(t *testing.T) {
	time.Sleep(time.Second * time.Duration(checkDuration*2+1)) // 等待两个周期

	if rt.svrs[serverAddr].Status == Down {
		t.Errorf("status check timeout recovery err.expect:<UP>, acture:<%v>", Up)
		return
	}
}

func TestEtcdWatchNewCluster(t *testing.T) {
	cluster := &Cluster{
		Name:   clusterName,
		LbName: lbName,
	}

	err := rt.store.SaveCluster(cluster)

	if nil != err {
		t.Error("add cluster err.")
		return
	}

	waitNotify()

	if len(rt.clusters) == 1 {
		return
	}

	t.Errorf("expect:<1>, acture:<%d>", len(rt.clusters))
}

func TestEtcdWatchNewBind(t *testing.T) {
	bind := &Bind{
		ClusterName: clusterName,
		ServerAddr:  serverAddr,
	}

	err := rt.store.SaveBind(bind)

	if nil != err {
		t.Error("add cluster err.")
		return
	}

	waitNotify()

	if len(rt.mapping) == 1 {
		return
	}

	t.Errorf("expect:<1>, acture:<%d>. %+v", len(rt.mapping), rt.mapping)
}

func TestEtcdWatchNewAPI(t *testing.T) {
	n := &Node{
		AttrName:    "test",
		ClusterName: clusterName,
	}

	err := rt.store.SaveAPI(&API{
		URL:    apiURL,
		Method: apiMethod,
		Nodes:  []*Node{n},
	})

	if nil != err {
		t.Error("add api err.")
		return
	}

	waitNotify()

	if len(rt.apis) == 1 {
		return
	}

	t.Errorf("expect:<1>, acture:<%d>", len(rt.apis))
}

func TestEtcdWatchUpdateServer(t *testing.T) {
	server := &Server{
		Schema:          "http",
		Addr:            serverAddr,
		CheckPath:       "/check",
		CheckDuration:   checkDuration,
		CheckTimeout:    checkTimeout * 2,
		MaxQPS:          3000,
		HalfToOpen:      100,
		HalfTrafficRate: 30,
		CloseCount:      200,
	}

	err := rt.store.UpdateServer(server)

	if nil != err {
		t.Error("update server err.")
		return
	}

	waitNotify()

	svr := rt.svrs[serverAddr]

	if svr.MaxQPS != server.MaxQPS {
		t.Errorf("MaxQPS expect:<%d>, acture:<%d>. ", server.MaxQPS, svr.MaxQPS)
		return
	}

	if svr.HalfToOpen != server.HalfToOpen {
		t.Errorf("HalfToOpen expect:<%d>, acture:<%d>. ", server.HalfToOpen, svr.HalfToOpen)
		return
	}

	if svr.HalfTrafficRate != server.HalfTrafficRate {
		t.Errorf("HalfTrafficRate expect:<%d>, acture:<%d>. ", server.HalfTrafficRate, svr.HalfTrafficRate)
		return
	}

	if svr.CloseCount != server.CloseCount {
		t.Errorf("CloseCount expect:<%d>, acture:<%d>. ", server.CloseCount, svr.CloseCount)
		return
	}

	if svr.CheckTimeout == server.CheckTimeout {
		t.Errorf("CheckTimeout expect:<%d>, acture:<%d>. ", svr.CheckTimeout, server.CheckTimeout)
		return
	}
}

func TestEtcdWatchUpdateCluster(t *testing.T) {
	cluster := &Cluster{
		Name:   clusterName,
		LbName: lbName,
	}

	err := rt.store.UpdateCluster(cluster)

	if nil != err {
		t.Error("update cluster err.")
		return
	}

	waitNotify()

	existCluster := rt.clusters[clusterName]

	if existCluster.LbName != cluster.LbName {
		t.Errorf("LbName expect:<%s>, acture:<%s>. ", cluster.LbName, existCluster.LbName)
		return
	}
}

func TestEtcdWatchUpdateAPI(t *testing.T) {
	n := &Node{
		AttrName:    "test",
		ClusterName: clusterName,
	}

	n2 := &Node{
		AttrName:    "tes2t",
		ClusterName: clusterName,
	}

	api := &API{
		URL:    apiURL,
		Method: apiMethod,
		Nodes:  []*Node{n, n2},
	}

	err := rt.store.UpdateAPI(api)

	if nil != err {
		t.Error("update api err.")
		return
	}

	waitNotify()

	existAPI, _ := rt.apis[getAPIKey(api.URL, api.Method)]

	if len(existAPI.Nodes) != len(api.Nodes) {
		t.Errorf("Nodes expect:<%s>, acture:<%d>. ", len(existAPI.Nodes), len(api.Nodes))
		return
	}
}

func TestEtcdWatchDeleteCluster(t *testing.T) {
	err := rt.store.DeleteCluster(clusterName)

	if nil != err {
		t.Error("delete cluster err.")
		return
	}

	waitNotify()

	if len(rt.clusters) != 0 {
		t.Errorf("clusters expect:<0>, acture:<%d>", len(rt.clusters))
		return
	}

	banded, _ := rt.mapping[serverAddr]

	if len(banded) != 0 {
		t.Errorf("banded expect:<0>, acture:<%d>", len(banded))
		return
	}
}

func TestEtcdWatchDeleteServer(t *testing.T) {
	err := rt.store.DeleteServer(serverAddr)

	if nil != err {
		t.Error("delete server err.")
		return
	}

	waitNotify()

	if len(rt.svrs) != 0 {
		t.Errorf("svrs expect:<0>, acture:<%d>", len(rt.svrs))
		return
	}

	if len(rt.mapping) != 0 {
		t.Errorf("mapping expect:<0>, acture:<%d>", len(rt.mapping))
		return
	}
}

func TestEtcdWatchDeleteAPI(t *testing.T) {
	err := rt.store.DeleteAPI(apiURL, apiMethod)

	if nil != err {
		t.Error("delete api err.")
		return
	}

	waitNotify()

	if len(rt.apis) != 0 {
		t.Errorf("apis expect:<0>, acture:<%d>", len(rt.apis))
		return
	}
}

func TestEtcdWatchNewRouting(t *testing.T) {
	r, err := NewRouting(`desc = "test"; deadline = 100; rule = ["$query_abc == 10", "$query_123 == 20"];`, clusterName, "")

	if nil != err {
		t.Error("add routing err.")
		return
	}

	err = rt.store.SaveRouting(r)

	if nil != err {
		t.Error("add routing err.")
		return
	}

	waitNotify()

	if len(rt.routings) == 1 {
		delete(rt.routings, r.ID)
		return
	}

	t.Errorf("expect:<1>, acture:<%d>", len(rt.routings))
}

func TestEtcdWatchDeleteRouting(t *testing.T) {
	r, err := NewRouting(`desc = "test"; deadline = 3; rule = ["$query_abc == 10", "$query_123 == 20"];`, clusterName, "")

	if nil != err {
		t.Error("add routing err.")
		return
	}

	err = rt.store.SaveRouting(r)

	if nil != err {
		t.Error("add routing err.")
		return
	}

	time.Sleep(time.Second * 30)

	if len(rt.routings) == 0 {
		return
	}

	t.Errorf("expect:<0>, acture:<%d>", len(rt.routings))
}
<|endoftext|>


<commit_before>// +build !windows

package ping

import (
	"errors"
	"fmt"
	"os/exec"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/influxdata/telegraf"
	"github.com/influxdata/telegraf/internal"
	"github.com/influxdata/telegraf/plugins/inputs"
)

// HostPinger is a function that runs the "ping" function using a list of
// passed arguments. This can be easily switched with a mocked ping function
// for unit test purposes (see ping_test.go)
type HostPinger func(timeout float64, args ...string) (string, error)

type Ping struct {
	// Interval at which to ping (ping -i <INTERVAL>)
	PingInterval float64 `toml:"ping_interval"`

	// Number of pings to send (ping -c <COUNT>)
	Count int

	// Ping timeout, in seconds. 0 means no timeout (ping -W <TIMEOUT>)
	Timeout float64

	// Interface to send ping from (ping -I <INTERFACE>)
	Interface string

	// URLs to ping
	Urls []string

	// host ping function
	pingHost HostPinger
}

func (_ *Ping) Description() string {
	return "Ping given url(s) and return statistics"
}

const sampleConfig = `
  ## NOTE: this plugin forks the ping command. You may need to set capabilities
  ## via setcap cap_net_raw+p /bin/ping
  #
  ## List of urls to ping
  urls = ["www.google.com"] # required
  ## number of pings to send per collection (ping -c <COUNT>)
  # count = 1
  ## interval, in s, at which to ping. 0 == default (ping -i <PING_INTERVAL>)
  # ping_interval = 1.0
  ## per-ping timeout, in s. 0 == no timeout (ping -W <TIMEOUT>)
  # timeout = 1.0
  ## interface to send ping from (ping -I <INTERFACE>)
  # interface = ""
`

func (_ *Ping) SampleConfig() string {
	return sampleConfig
}

func (p *Ping) Gather(acc telegraf.Accumulator) error {

	var wg sync.WaitGroup

	// Spin off a go routine for each url to ping
	for _, url := range p.Urls {
		wg.Add(1)
		go func(u string) {
			defer wg.Done()
			args := p.args(u)
			totalTimeout := float64(p.Count)*p.Timeout + float64(p.Count-1)*p.PingInterval

			out, err := p.pingHost(totalTimeout, args...)
			if err != nil {
				// Some implementations of ping return a 1 exit code on
				// timeout, if this occurs we will not exit and try to parse
				// the output.
				status := -1
				if exitError, ok := err.(*exec.ExitError); ok {
					if ws, ok := exitError.Sys().(syscall.WaitStatus); ok {
						status = ws.ExitStatus()
					}
				}

				if status != 1 {
					// Combine go err + stderr output
					out = strings.TrimSpace(out)
					if len(out) > 0 {
						acc.AddError(fmt.Errorf("%s, %s", out, err))
					} else {
						acc.AddError(err)
					}
					return
				}
			}

			tags := map[string]string{"url": u}
			trans, rec, min, avg, max, stddev, err := processPingOutput(out)
			if err != nil {
				// fatal error
				acc.AddError(fmt.Errorf("%s: %s", err, u))
				return
			}
			// Calculate packet loss percentage
			loss := float64(trans-rec) / float64(trans) * 100.0
			fields := map[string]interface{}{
				"packets_transmitted": trans,
				"packets_received":    rec,
				"percent_packet_loss": loss,
			}
			if min > 0 {
				fields["minimum_response_ms"] = min
			}
			if avg > 0 {
				fields["average_response_ms"] = avg
			}
			if max > 0 {
				fields["maximum_response_ms"] = max
			}
			if stddev > 0 {
				fields["standard_deviation_ms"] = stddev
			}
			acc.AddFields("ping", fields, tags)
		}(url)
	}

	wg.Wait()

	return nil
}

func hostPinger(timeout float64, args ...string) (string, error) {
	bin, err := exec.LookPath("ping")
	if err != nil {
		return "", err
	}
	c := exec.Command(bin, args...)
	out, err := internal.CombinedOutputTimeout(c,
		time.Second*time.Duration(timeout+1))
	return string(out), err
}

// args returns the arguments for the 'ping' executable
func (p *Ping) args(url string) []string {
	// Build the ping command args based on toml config
	args := []string{"-c", strconv.Itoa(p.Count), "-n", "-s", "16"}
	if p.PingInterval > 0 {
		args = append(args, "-i", strconv.FormatFloat(p.PingInterval, 'f', 1, 64))
	}
	if p.Timeout > 0 {
		switch runtime.GOOS {
		case "darwin":
			args = append(args, "-W", strconv.FormatFloat(p.Timeout*1000, 'f', 1, 64))
		case "linux":
			args = append(args, "-W", strconv.FormatFloat(p.Timeout, 'f', 1, 64))
		default:
			// Not sure the best option here, just assume GNU ping?
			args = append(args, "-W", strconv.FormatFloat(p.Timeout, 'f', 1, 64))
		}
	}
	if p.Interface != "" {
		args = append(args, "-I", p.Interface)
	}
	args = append(args, url)
	return args
}

// processPingOutput takes in a string output from the ping command, like:
//
//     PING www.google.com (173.194.115.84): 56 data bytes
//     64 bytes from 173.194.115.84: icmp_seq=0 ttl=54 time=52.172 ms
//     64 bytes from 173.194.115.84: icmp_seq=1 ttl=54 time=34.843 ms
//
//     --- www.google.com ping statistics ---
//     2 packets transmitted, 2 packets received, 0.0% packet loss
//     round-trip min/avg/max/stddev = 34.843/43.508/52.172/8.664 ms
//
// It returns (<transmitted packets>, <received packets>, <average response>)
func processPingOutput(out string) (int, int, float64, float64, float64, float64, error) {
	var trans, recv int
	var min, avg, max, stddev float64
	// Set this error to nil if we find a 'transmitted' line
	err := errors.New("Fatal error processing ping output")
	lines := strings.Split(out, "\n")
	for _, line := range lines {
		if strings.Contains(line, "transmitted") &&
			strings.Contains(line, "received") {
			err = nil
			stats := strings.Split(line, ", ")
			// Transmitted packets
			trans, err = strconv.Atoi(strings.Split(stats[0], " ")[0])
			if err != nil {
				return trans, recv, min, avg, max, stddev, err
			}
			// Received packets
			recv, err = strconv.Atoi(strings.Split(stats[1], " ")[0])
			if err != nil {
				return trans, recv, min, avg, max, stddev, err
			}
		} else if strings.Contains(line, "min/avg/max") {
			stats := strings.Split(line, " ")[3]
			min, err = strconv.ParseFloat(strings.Split(stats, "/")[0], 64)
			avg, err = strconv.ParseFloat(strings.Split(stats, "/")[1], 64)
			max, err = strconv.ParseFloat(strings.Split(stats, "/")[2], 64)
			stddev, err = strconv.ParseFloat(strings.Split(stats, "/")[3], 64)
			if err != nil {
				return trans, recv, min, avg, max, stddev, err
			}
		}
	}
	return trans, recv, min, avg, max, stddev, err
}

func init() {
	inputs.Add("ping", func() telegraf.Input {
		return &Ping{
			pingHost:     hostPinger,
			PingInterval: 1.0,
			Count:        1,
			Timeout:      1.0,
		}
	})
}
<commit_msg>Use 5 second timeout overhead when waiting for ping to complete<commit_after>// +build !windows

package ping

import (
	"errors"
	"fmt"
	"os/exec"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/influxdata/telegraf"
	"github.com/influxdata/telegraf/internal"
	"github.com/influxdata/telegraf/plugins/inputs"
)

// HostPinger is a function that runs the "ping" function using a list of
// passed arguments. This can be easily switched with a mocked ping function
// for unit test purposes (see ping_test.go)
type HostPinger func(timeout float64, args ...string) (string, error)

type Ping struct {
	// Interval at which to ping (ping -i <INTERVAL>)
	PingInterval float64 `toml:"ping_interval"`

	// Number of pings to send (ping -c <COUNT>)
	Count int

	// Ping timeout, in seconds. 0 means no timeout (ping -W <TIMEOUT>)
	Timeout float64

	// Interface to send ping from (ping -I <INTERFACE>)
	Interface string

	// URLs to ping
	Urls []string

	// host ping function
	pingHost HostPinger
}

func (_ *Ping) Description() string {
	return "Ping given url(s) and return statistics"
}

const sampleConfig = `
  ## NOTE: this plugin forks the ping command. You may need to set capabilities
  ## via setcap cap_net_raw+p /bin/ping
  #
  ## List of urls to ping
  urls = ["www.google.com"] # required
  ## number of pings to send per collection (ping -c <COUNT>)
  # count = 1
  ## interval, in s, at which to ping. 0 == default (ping -i <PING_INTERVAL>)
  # ping_interval = 1.0
  ## per-ping timeout, in s. 0 == no timeout (ping -W <TIMEOUT>)
  # timeout = 1.0
  ## interface to send ping from (ping -I <INTERFACE>)
  # interface = ""
`

func (_ *Ping) SampleConfig() string {
	return sampleConfig
}

func (p *Ping) Gather(acc telegraf.Accumulator) error {

	var wg sync.WaitGroup

	// Spin off a go routine for each url to ping
	for _, url := range p.Urls {
		wg.Add(1)
		go func(u string) {
			defer wg.Done()
			args := p.args(u)
			totalTimeout := float64(p.Count)*p.Timeout + float64(p.Count-1)*p.PingInterval

			out, err := p.pingHost(totalTimeout, args...)
			if err != nil {
				// Some implementations of ping return a 1 exit code on
				// timeout, if this occurs we will not exit and try to parse
				// the output.
				status := -1
				if exitError, ok := err.(*exec.ExitError); ok {
					if ws, ok := exitError.Sys().(syscall.WaitStatus); ok {
						status = ws.ExitStatus()
					}
				}

				if status != 1 {
					// Combine go err + stderr output
					out = strings.TrimSpace(out)
					if len(out) > 0 {
						acc.AddError(fmt.Errorf("%s, %s", out, err))
					} else {
						acc.AddError(err)
					}
					return
				}
			}

			tags := map[string]string{"url": u}
			trans, rec, min, avg, max, stddev, err := processPingOutput(out)
			if err != nil {
				// fatal error
				acc.AddError(fmt.Errorf("%s: %s", err, u))
				return
			}
			// Calculate packet loss percentage
			loss := float64(trans-rec) / float64(trans) * 100.0
			fields := map[string]interface{}{
				"packets_transmitted": trans,
				"packets_received":    rec,
				"percent_packet_loss": loss,
			}
			if min > 0 {
				fields["minimum_response_ms"] = min
			}
			if avg > 0 {
				fields["average_response_ms"] = avg
			}
			if max > 0 {
				fields["maximum_response_ms"] = max
			}
			if stddev > 0 {
				fields["standard_deviation_ms"] = stddev
			}
			acc.AddFields("ping", fields, tags)
		}(url)
	}

	wg.Wait()

	return nil
}

func hostPinger(timeout float64, args ...string) (string, error) {
	bin, err := exec.LookPath("ping")
	if err != nil {
		return "", err
	}
	c := exec.Command(bin, args...)
	out, err := internal.CombinedOutputTimeout(c,
		time.Second*time.Duration(timeout+5))
	return string(out), err
}

// args returns the arguments for the 'ping' executable
func (p *Ping) args(url string) []string {
	// Build the ping command args based on toml config
	args := []string{"-c", strconv.Itoa(p.Count), "-n", "-s", "16"}
	if p.PingInterval > 0 {
		args = append(args, "-i", strconv.FormatFloat(p.PingInterval, 'f', 1, 64))
	}
	if p.Timeout > 0 {
		switch runtime.GOOS {
		case "darwin":
			args = append(args, "-W", strconv.FormatFloat(p.Timeout*1000, 'f', 1, 64))
		case "linux":
			args = append(args, "-W", strconv.FormatFloat(p.Timeout, 'f', 1, 64))
		default:
			// Not sure the best option here, just assume GNU ping?
			args = append(args, "-W", strconv.FormatFloat(p.Timeout, 'f', 1, 64))
		}
	}
	if p.Interface != "" {
		args = append(args, "-I", p.Interface)
	}
	args = append(args, url)
	return args
}

// processPingOutput takes in a string output from the ping command, like:
//
//     PING www.google.com (173.194.115.84): 56 data bytes
//     64 bytes from 173.194.115.84: icmp_seq=0 ttl=54 time=52.172 ms
//     64 bytes from 173.194.115.84: icmp_seq=1 ttl=54 time=34.843 ms
//
//     --- www.google.com ping statistics ---
//     2 packets transmitted, 2 packets received, 0.0% packet loss
//     round-trip min/avg/max/stddev = 34.843/43.508/52.172/8.664 ms
//
// It returns (<transmitted packets>, <received packets>, <average response>)
func processPingOutput(out string) (int, int, float64, float64, float64, float64, error) {
	var trans, recv int
	var min, avg, max, stddev float64
	// Set this error to nil if we find a 'transmitted' line
	err := errors.New("Fatal error processing ping output")
	lines := strings.Split(out, "\n")
	for _, line := range lines {
		if strings.Contains(line, "transmitted") &&
			strings.Contains(line, "received") {
			err = nil
			stats := strings.Split(line, ", ")
			// Transmitted packets
			trans, err = strconv.Atoi(strings.Split(stats[0], " ")[0])
			if err != nil {
				return trans, recv, min, avg, max, stddev, err
			}
			// Received packets
			recv, err = strconv.Atoi(strings.Split(stats[1], " ")[0])
			if err != nil {
				return trans, recv, min, avg, max, stddev, err
			}
		} else if strings.Contains(line, "min/avg/max") {
			stats := strings.Split(line, " ")[3]
			min, err = strconv.ParseFloat(strings.Split(stats, "/")[0], 64)
			avg, err = strconv.ParseFloat(strings.Split(stats, "/")[1], 64)
			max, err = strconv.ParseFloat(strings.Split(stats, "/")[2], 64)
			stddev, err = strconv.ParseFloat(strings.Split(stats, "/")[3], 64)
			if err != nil {
				return trans, recv, min, avg, max, stddev, err
			}
		}
	}
	return trans, recv, min, avg, max, stddev, err
}

func init() {
	inputs.Add("ping", func() telegraf.Input {
		return &Ping{
			pingHost:     hostPinger,
			PingInterval: 1.0,
			Count:        1,
			Timeout:      1.0,
		}
	})
}
<|endoftext|>


<commit_before>package middleware

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

	"github.com/labstack/echo/v4"
	"golang.org/x/time/rate"
)

type (
	// RateLimiterStore is the interface to be implemented by custom stores.
	RateLimiterStore interface {
		// Stores for the rate limiter have to implement the Allow method
		Allow(identifier string) (bool, error)
	}
)

type (
	// RateLimiterConfig defines the configuration for the rate limiter
	RateLimiterConfig struct {
		Skipper    Skipper
		BeforeFunc BeforeFunc
		// IdentifierExtractor uses echo.Context to extract the identifier for a visitor
		IdentifierExtractor Extractor
		// Store defines a store for the rate limiter
		Store RateLimiterStore
		// ErrorHandler provides a handler to be called when IdentifierExtractor returns an error
		ErrorHandler func(context echo.Context, err error) error
		// DenyHandler provides a handler to be called when RateLimiter denies access
		DenyHandler func(context echo.Context, identifier string, err error) error
	}
	// Extractor is used to extract data from echo.Context
	Extractor func(context echo.Context) (string, error)
)

// errors
var (
	// ErrRateLimitExceeded denotes an error raised when rate limit is exceeded
	ErrRateLimitExceeded = echo.NewHTTPError(http.StatusTooManyRequests, "rate limit exceeded")
	// ErrExtractorError denotes an error raised when extractor function is unsuccessful
	ErrExtractorError = echo.NewHTTPError(http.StatusForbidden, "error while extracting identifier")
)

// DefaultRateLimiterConfig defines default values for RateLimiterConfig
var DefaultRateLimiterConfig = RateLimiterConfig{
	Skipper: DefaultSkipper,
	IdentifierExtractor: func(ctx echo.Context) (string, error) {
		id := ctx.RealIP()
		return id, nil
	},
	ErrorHandler: func(context echo.Context, err error) error {
		return &echo.HTTPError{
			Code:     ErrExtractorError.Code,
			Message:  ErrExtractorError.Message,
			Internal: err,
		}
	},
	DenyHandler: func(context echo.Context, identifier string, err error) error {
		return &echo.HTTPError{
			Code:     ErrRateLimitExceeded.Code,
			Message:  ErrRateLimitExceeded.Message,
			Internal: err,
		}
	},
}

/*
RateLimiter returns a rate limiting middleware

	e := echo.New()

	limiterStore := middleware.NewRateLimiterMemoryStore(20)

	e.GET("/rate-limited", func(c echo.Context) error {
		return c.String(http.StatusOK, "test")
	}, RateLimiter(limiterStore))
*/
func RateLimiter(store RateLimiterStore) echo.MiddlewareFunc {
	config := DefaultRateLimiterConfig
	config.Store = store

	return RateLimiterWithConfig(config)
}

/*
RateLimiterWithConfig returns a rate limiting middleware

	e := echo.New()

	config := middleware.RateLimiterConfig{
		Skipper: DefaultSkipper,
		Store: middleware.NewRateLimiterMemoryStore(
			middleware.RateLimiterMemoryStoreConfig{Rate: 10, Burst: 30, ExpiresIn: 3 * time.Minute}
		)
		IdentifierExtractor: func(ctx echo.Context) (string, error) {
			id := ctx.RealIP()
			return id, nil
		},
		ErrorHandler: func(context echo.Context, err error) error {
			return context.JSON(http.StatusTooManyRequests, nil)
		},
		DenyHandler: func(context echo.Context, identifier string) error {
			return context.JSON(http.StatusForbidden, nil)
		},
	}

	e.GET("/rate-limited", func(c echo.Context) error {
		return c.String(http.StatusOK, "test")
	}, middleware.RateLimiterWithConfig(config))
*/
func RateLimiterWithConfig(config RateLimiterConfig) echo.MiddlewareFunc {
	if config.Skipper == nil {
		config.Skipper = DefaultRateLimiterConfig.Skipper
	}
	if config.IdentifierExtractor == nil {
		config.IdentifierExtractor = DefaultRateLimiterConfig.IdentifierExtractor
	}
	if config.ErrorHandler == nil {
		config.ErrorHandler = DefaultRateLimiterConfig.ErrorHandler
	}
	if config.DenyHandler == nil {
		config.DenyHandler = DefaultRateLimiterConfig.DenyHandler
	}
	if config.Store == nil {
		panic("Store configuration must be provided")
	}
	return func(next echo.HandlerFunc) echo.HandlerFunc {
		return func(c echo.Context) error {
			if config.Skipper(c) {
				return next(c)
			}
			if config.BeforeFunc != nil {
				config.BeforeFunc(c)
			}

			identifier, err := config.IdentifierExtractor(c)
			if err != nil {
				c.Error(config.ErrorHandler(c, err))
				return nil
			}

			if allow, err := config.Store.Allow(identifier); !allow {
				c.Error(config.DenyHandler(c, identifier, err))
				return nil
			}
			return next(c)
		}
	}
}

type (
	// RateLimiterMemoryStore is the built-in store implementation for RateLimiter
	RateLimiterMemoryStore struct {
		visitors    map[string]*Visitor
		mutex       sync.Mutex
		rate        rate.Limit
		burst       int
		expiresIn   time.Duration
		lastCleanup time.Time
	}
	// Visitor signifies a unique user's limiter details
	Visitor struct {
		*rate.Limiter
		lastSeen time.Time
	}
)

/*
NewRateLimiterMemoryStore returns an instance of RateLimiterMemoryStore with
the provided rate (as req/s). The provided rate less than 1 will be treated as zero.
Also rate between 0 and 1 will be treated as zero.
for more info check out Limiter docs - https://pkg.go.dev/golang.org/x/time/rate#Limit.

Burst and ExpiresIn will be set to default values.

Example (with 20 requests/sec):

	limiterStore := middleware.NewRateLimiterMemoryStore(20)

*/
func NewRateLimiterMemoryStore(rate rate.Limit) (store *RateLimiterMemoryStore) {
	return NewRateLimiterMemoryStoreWithConfig(RateLimiterMemoryStoreConfig{
		Rate: rate,
	})
}

/*
NewRateLimiterMemoryStoreWithConfig returns an instance of RateLimiterMemoryStore
with the provided configuration. Rate must be provided. Burst will be set to the value of
the configured rate if not provided or set to 0.

The build-in memory store is usually capable for modest loads. For higher loads other
store implementations should be considered.

Characteristics:
* Concurrency above 100 parallel requests may causes measurable lock contention
* A high number of different IP addresses (above 16000) may be impacted by the internally used Go map
* A high number of requests from a single IP address may cause lock contention

Example:

	limiterStore := middleware.NewRateLimiterMemoryStoreWithConfig(
		middleware.RateLimiterMemoryStoreConfig{Rate: 50, Burst: 200, ExpiresIn: 5 * time.Minute},
	)
*/
func NewRateLimiterMemoryStoreWithConfig(config RateLimiterMemoryStoreConfig) (store *RateLimiterMemoryStore) {
	store = &RateLimiterMemoryStore{}

	store.rate = config.Rate
	store.burst = config.Burst
	store.expiresIn = config.ExpiresIn
	if config.ExpiresIn == 0 {
		store.expiresIn = DefaultRateLimiterMemoryStoreConfig.ExpiresIn
	}
	if config.Burst == 0 {
		store.burst = int(config.Rate)
	}
	store.visitors = make(map[string]*Visitor)
	store.lastCleanup = now()
	return
}

// RateLimiterMemoryStoreConfig represents configuration for RateLimiterMemoryStore
type RateLimiterMemoryStoreConfig struct {
	Rate      rate.Limit    // Rate of requests allowed to pass as req/s
	Burst     int           // Burst additionally allows a number of requests to pass when rate limit is reached
	ExpiresIn time.Duration // ExpiresIn is the duration after that a rate limiter is cleaned up
}

// DefaultRateLimiterMemoryStoreConfig provides default configuration values for RateLimiterMemoryStore
var DefaultRateLimiterMemoryStoreConfig = RateLimiterMemoryStoreConfig{
	ExpiresIn: 3 * time.Minute,
}

// Allow implements RateLimiterStore.Allow
func (store *RateLimiterMemoryStore) Allow(identifier string) (bool, error) {
	store.mutex.Lock()
	limiter, exists := store.visitors[identifier]
	if !exists {
		limiter = new(Visitor)
		limiter.Limiter = rate.NewLimiter(store.rate, store.burst)
		store.visitors[identifier] = limiter
	}
	limiter.lastSeen = now()
	if now().Sub(store.lastCleanup) > store.expiresIn {
		store.cleanupStaleVisitors()
	}
	store.mutex.Unlock()
	return limiter.AllowN(now(), 1), nil
}

/*
cleanupStaleVisitors helps manage the size of the visitors map by removing stale records
of users who haven't visited again after the configured expiry time has elapsed
*/
func (store *RateLimiterMemoryStore) cleanupStaleVisitors() {
	for id, visitor := range store.visitors {
		if now().Sub(visitor.lastSeen) > store.expiresIn {
			delete(store.visitors, id)
		}
	}
	store.lastCleanup = now()
}

/*
actual time method which is mocked in test file
*/
var now = time.Now
<commit_msg>removed unnecessary comments<commit_after>package middleware

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

	"github.com/labstack/echo/v4"
	"golang.org/x/time/rate"
)

type (
	// RateLimiterStore is the interface to be implemented by custom stores.
	RateLimiterStore interface {
		// Stores for the rate limiter have to implement the Allow method
		Allow(identifier string) (bool, error)
	}
)

type (
	// RateLimiterConfig defines the configuration for the rate limiter
	RateLimiterConfig struct {
		Skipper    Skipper
		BeforeFunc BeforeFunc
		// IdentifierExtractor uses echo.Context to extract the identifier for a visitor
		IdentifierExtractor Extractor
		// Store defines a store for the rate limiter
		Store RateLimiterStore
		// ErrorHandler provides a handler to be called when IdentifierExtractor returns an error
		ErrorHandler func(context echo.Context, err error) error
		// DenyHandler provides a handler to be called when RateLimiter denies access
		DenyHandler func(context echo.Context, identifier string, err error) error
	}
	// Extractor is used to extract data from echo.Context
	Extractor func(context echo.Context) (string, error)
)

// errors
var (
	// ErrRateLimitExceeded denotes an error raised when rate limit is exceeded
	ErrRateLimitExceeded = echo.NewHTTPError(http.StatusTooManyRequests, "rate limit exceeded")
	// ErrExtractorError denotes an error raised when extractor function is unsuccessful
	ErrExtractorError = echo.NewHTTPError(http.StatusForbidden, "error while extracting identifier")
)

// DefaultRateLimiterConfig defines default values for RateLimiterConfig
var DefaultRateLimiterConfig = RateLimiterConfig{
	Skipper: DefaultSkipper,
	IdentifierExtractor: func(ctx echo.Context) (string, error) {
		id := ctx.RealIP()
		return id, nil
	},
	ErrorHandler: func(context echo.Context, err error) error {
		return &echo.HTTPError{
			Code:     ErrExtractorError.Code,
			Message:  ErrExtractorError.Message,
			Internal: err,
		}
	},
	DenyHandler: func(context echo.Context, identifier string, err error) error {
		return &echo.HTTPError{
			Code:     ErrRateLimitExceeded.Code,
			Message:  ErrRateLimitExceeded.Message,
			Internal: err,
		}
	},
}

/*
RateLimiter returns a rate limiting middleware

	e := echo.New()

	limiterStore := middleware.NewRateLimiterMemoryStore(20)

	e.GET("/rate-limited", func(c echo.Context) error {
		return c.String(http.StatusOK, "test")
	}, RateLimiter(limiterStore))
*/
func RateLimiter(store RateLimiterStore) echo.MiddlewareFunc {
	config := DefaultRateLimiterConfig
	config.Store = store

	return RateLimiterWithConfig(config)
}

/*
RateLimiterWithConfig returns a rate limiting middleware

	e := echo.New()

	config := middleware.RateLimiterConfig{
		Skipper: DefaultSkipper,
		Store: middleware.NewRateLimiterMemoryStore(
			middleware.RateLimiterMemoryStoreConfig{Rate: 10, Burst: 30, ExpiresIn: 3 * time.Minute}
		)
		IdentifierExtractor: func(ctx echo.Context) (string, error) {
			id := ctx.RealIP()
			return id, nil
		},
		ErrorHandler: func(context echo.Context, err error) error {
			return context.JSON(http.StatusTooManyRequests, nil)
		},
		DenyHandler: func(context echo.Context, identifier string) error {
			return context.JSON(http.StatusForbidden, nil)
		},
	}

	e.GET("/rate-limited", func(c echo.Context) error {
		return c.String(http.StatusOK, "test")
	}, middleware.RateLimiterWithConfig(config))
*/
func RateLimiterWithConfig(config RateLimiterConfig) echo.MiddlewareFunc {
	if config.Skipper == nil {
		config.Skipper = DefaultRateLimiterConfig.Skipper
	}
	if config.IdentifierExtractor == nil {
		config.IdentifierExtractor = DefaultRateLimiterConfig.IdentifierExtractor
	}
	if config.ErrorHandler == nil {
		config.ErrorHandler = DefaultRateLimiterConfig.ErrorHandler
	}
	if config.DenyHandler == nil {
		config.DenyHandler = DefaultRateLimiterConfig.DenyHandler
	}
	if config.Store == nil {
		panic("Store configuration must be provided")
	}
	return func(next echo.HandlerFunc) echo.HandlerFunc {
		return func(c echo.Context) error {
			if config.Skipper(c) {
				return next(c)
			}
			if config.BeforeFunc != nil {
				config.BeforeFunc(c)
			}

			identifier, err := config.IdentifierExtractor(c)
			if err != nil {
				c.Error(config.ErrorHandler(c, err))
				return nil
			}

			if allow, err := config.Store.Allow(identifier); !allow {
				c.Error(config.DenyHandler(c, identifier, err))
				return nil
			}
			return next(c)
		}
	}
}

type (
	// RateLimiterMemoryStore is the built-in store implementation for RateLimiter
	RateLimiterMemoryStore struct {
		visitors    map[string]*Visitor
		mutex       sync.Mutex
		rate        rate.Limit
		burst       int
		expiresIn   time.Duration
		lastCleanup time.Time
	}
	// Visitor signifies a unique user's limiter details
	Visitor struct {
		*rate.Limiter
		lastSeen time.Time
	}
)

/*
NewRateLimiterMemoryStore returns an instance of RateLimiterMemoryStore with
the provided rate (as req/s). The provided rate less than 1 will be treated as zero.
for more info check out Limiter docs - https://pkg.go.dev/golang.org/x/time/rate#Limit.

Burst and ExpiresIn will be set to default values.

Example (with 20 requests/sec):

	limiterStore := middleware.NewRateLimiterMemoryStore(20)

*/
func NewRateLimiterMemoryStore(rate rate.Limit) (store *RateLimiterMemoryStore) {
	return NewRateLimiterMemoryStoreWithConfig(RateLimiterMemoryStoreConfig{
		Rate: rate,
	})
}

/*
NewRateLimiterMemoryStoreWithConfig returns an instance of RateLimiterMemoryStore
with the provided configuration. Rate must be provided. Burst will be set to the value of
the configured rate if not provided or set to 0.

The build-in memory store is usually capable for modest loads. For higher loads other
store implementations should be considered.

Characteristics:
* Concurrency above 100 parallel requests may causes measurable lock contention
* A high number of different IP addresses (above 16000) may be impacted by the internally used Go map
* A high number of requests from a single IP address may cause lock contention

Example:

	limiterStore := middleware.NewRateLimiterMemoryStoreWithConfig(
		middleware.RateLimiterMemoryStoreConfig{Rate: 50, Burst: 200, ExpiresIn: 5 * time.Minute},
	)
*/
func NewRateLimiterMemoryStoreWithConfig(config RateLimiterMemoryStoreConfig) (store *RateLimiterMemoryStore) {
	store = &RateLimiterMemoryStore{}

	store.rate = config.Rate
	store.burst = config.Burst
	store.expiresIn = config.ExpiresIn
	if config.ExpiresIn == 0 {
		store.expiresIn = DefaultRateLimiterMemoryStoreConfig.ExpiresIn
	}
	if config.Burst == 0 {
		store.burst = int(config.Rate)
	}
	store.visitors = make(map[string]*Visitor)
	store.lastCleanup = now()
	return
}

// RateLimiterMemoryStoreConfig represents configuration for RateLimiterMemoryStore
type RateLimiterMemoryStoreConfig struct {
	Rate      rate.Limit    // Rate of requests allowed to pass as req/s
	Burst     int           // Burst additionally allows a number of requests to pass when rate limit is reached
	ExpiresIn time.Duration // ExpiresIn is the duration after that a rate limiter is cleaned up
}

// DefaultRateLimiterMemoryStoreConfig provides default configuration values for RateLimiterMemoryStore
var DefaultRateLimiterMemoryStoreConfig = RateLimiterMemoryStoreConfig{
	ExpiresIn: 3 * time.Minute,
}

// Allow implements RateLimiterStore.Allow
func (store *RateLimiterMemoryStore) Allow(identifier string) (bool, error) {
	store.mutex.Lock()
	limiter, exists := store.visitors[identifier]
	if !exists {
		limiter = new(Visitor)
		limiter.Limiter = rate.NewLimiter(store.rate, store.burst)
		store.visitors[identifier] = limiter
	}
	limiter.lastSeen = now()
	if now().Sub(store.lastCleanup) > store.expiresIn {
		store.cleanupStaleVisitors()
	}
	store.mutex.Unlock()
	return limiter.AllowN(now(), 1), nil
}

/*
cleanupStaleVisitors helps manage the size of the visitors map by removing stale records
of users who haven't visited again after the configured expiry time has elapsed
*/
func (store *RateLimiterMemoryStore) cleanupStaleVisitors() {
	for id, visitor := range store.visitors {
		if now().Sub(visitor.lastSeen) > store.expiresIn {
			delete(store.visitors, id)
		}
	}
	store.lastCleanup = now()
}

/*
actual time method which is mocked in test file
*/
var now = time.Now
<|endoftext|>