<commit_before>package main

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"strings"
)

const (
	ADD_SERVICE    = "http://localhost:3000/api/v1/service/add"
	RUN_CONTAINER  = "http://localhost:3000/api/v1/container/run"
	LIST_CONTAINER = "http://localhost:3000/api/v1/container/list"
	EXEC_CONTAINER = "http://localhost:3000/api/v1/container/exec"
	RM_CONTAINER   = "http://localhost:3000/api/v1/container/rm"
)

func call(url string) {
	if len(os.Args) != 4 {
		help()
		return
	}

	filename := os.Args[3]

	raw, err := ioutil.ReadFile(filename)
	if err != nil {
		fmt.Println(err)
		return
	}

	req, err := http.NewRequest("POST", url, bytes.NewBuffer(raw))
	req.Header.Set("Content-Type", "application/json")
	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		panic(err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != 201 {
		body, _ := ioutil.ReadAll(resp.Body)
		fmt.Println(string(body))
	}
}

func list_containers() {
	resp, err := http.Get(LIST_CONTAINER)
	if err != nil {
		fmt.Println(err)
		return
	}

	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		fmt.Println(err)
		return
	}

	fmt.Println(string(body))
}

func exec_container() {
	if len(os.Args) <= 4 {
		help()
		return
	}

	raw := []byte(fmt.Sprintf("{\"name\":\"%s\"}", os.Args[3]))
	fmt.Println(string(raw))
	req, err := http.NewRequest("POST", EXEC_CONTAINER, bytes.NewBuffer(raw))
	req.Header.Set("Content-Type", "application/json")
	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		panic(err)
	}
	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		fmt.Println(err)
		return
	}

	pid := string(body)
	cmd := strings.Join(os.Args[4:], " ")
	path, err := exec.LookPath("nsenter")
	if err != nil {
		fmt.Println("Cannot find nsenter")
		return
	}

	command := strings.Split(fmt.Sprintf("%s --target %s --pid --net --mount %s", path, pid, cmd), " ")
	run := &exec.Cmd{
		Path: command[0],
		Args: command,
	}

	run.Stdin = os.Stdin
	run.Stdout = os.Stdout
	run.Stderr = os.Stderr
	run.Start()
	run.Wait()
}

func rm_container() {
	if len(os.Args) != 4 {
		help()
		return
	}

	raw := []byte(fmt.Sprintf("{\"name\":\"%s\"}", os.Args[3]))
	fmt.Println(string(raw))
	req, err := http.NewRequest("POST", RM_CONTAINER, bytes.NewBuffer(raw))
	req.Header.Set("Content-Type", "application/json")
	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		panic(err)
	}
	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		fmt.Println(err)
		return
	}
	fmt.Println(string(body))
}

func help() {
	fmt.Println("Commands:")
	fmt.Println("\tcontainer")
	fmt.Println("\tservice")
	fmt.Println()

	fmt.Println("container:")
	fmt.Println("\trun filename.json - Runs command specified in filename.json")
	fmt.Println("\texec container_hash command... - Runs specified command in container")
	fmt.Println("\tlist - Lists all running containers")
	fmt.Println("\trm container_hash - Stops container")
	fmt.Println()

	fmt.Println("service:")
	fmt.Println("\tadd filename.json - Creates a service with details from filename.json")
}

func main() {
	switch os.Args[1] {
	case "service":
		switch os.Args[2] {
		case "add":
			call(ADD_SERVICE)
		}
	case "container":
		switch os.Args[2] {
		case "run":
			call(RUN_CONTAINER)
		case "list":
			list_containers()
		case "exec":
			exec_container()
		case "rm":
			rm_container()
		}
	}
}
<commit_msg>Added more help<commit_after>package main

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"strings"
)

const (
	ADD_SERVICE    = "http://localhost:3000/api/v1/service/add"
	RUN_CONTAINER  = "http://localhost:3000/api/v1/container/run"
	LIST_CONTAINER = "http://localhost:3000/api/v1/container/list"
	EXEC_CONTAINER = "http://localhost:3000/api/v1/container/exec"
	RM_CONTAINER   = "http://localhost:3000/api/v1/container/rm"
)

func call(url string) {
	if len(os.Args) != 4 {
		help()
		return
	}

	filename := os.Args[3]

	raw, err := ioutil.ReadFile(filename)
	if err != nil {
		fmt.Println(err)
		return
	}

	req, err := http.NewRequest("POST", url, bytes.NewBuffer(raw))
	req.Header.Set("Content-Type", "application/json")
	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		panic(err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != 201 {
		body, _ := ioutil.ReadAll(resp.Body)
		fmt.Println(string(body))
	}
}

func list_containers() {
	resp, err := http.Get(LIST_CONTAINER)
	if err != nil {
		fmt.Println(err)
		return
	}

	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		fmt.Println(err)
		return
	}

	fmt.Println(string(body))
}

func exec_container() {
	if len(os.Args) <= 4 {
		help()
		return
	}

	raw := []byte(fmt.Sprintf("{\"name\":\"%s\"}", os.Args[3]))
	fmt.Println(string(raw))
	req, err := http.NewRequest("POST", EXEC_CONTAINER, bytes.NewBuffer(raw))
	req.Header.Set("Content-Type", "application/json")
	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		panic(err)
	}
	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		fmt.Println(err)
		return
	}

	pid := string(body)
	cmd := strings.Join(os.Args[4:], " ")
	path, err := exec.LookPath("nsenter")
	if err != nil {
		fmt.Println("Cannot find nsenter")
		return
	}

	command := strings.Split(fmt.Sprintf("%s --target %s --pid --net --mount %s", path, pid, cmd), " ")
	run := &exec.Cmd{
		Path: command[0],
		Args: command,
	}

	run.Stdin = os.Stdin
	run.Stdout = os.Stdout
	run.Stderr = os.Stderr
	run.Start()
	run.Wait()
}

func rm_container() {
	if len(os.Args) != 4 {
		help()
		return
	}

	raw := []byte(fmt.Sprintf("{\"name\":\"%s\"}", os.Args[3]))
	fmt.Println(string(raw))
	req, err := http.NewRequest("POST", RM_CONTAINER, bytes.NewBuffer(raw))
	req.Header.Set("Content-Type", "application/json")
	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		panic(err)
	}
	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		fmt.Println(err)
		return
	}
	fmt.Println(string(body))
}

func help() {
	fmt.Println("Commands:")
	fmt.Println("\tcontainer")
	fmt.Println("\tservice")
	fmt.Println()

	fmt.Println("container:")
	fmt.Println("\trun filename.json - Runs command specified in filename.json")
	fmt.Println("\texec container_hash command... - Runs specified command in container")
	fmt.Println("\tlist - Lists all running containers")
	fmt.Println("\trm container_hash - Stops container")
	fmt.Println()

	fmt.Println("service:")
	fmt.Println("\tadd filename.json - Creates a service with details from filename.json")
}

func main() {
	if len(os.Args) < 3 {
		help()
		return
	}

	switch os.Args[1] {
	case "service":
		switch os.Args[2] {
		case "add":
			call(ADD_SERVICE)
		default:
			help()
		}
	case "container":
		switch os.Args[2] {
		case "run":
			call(RUN_CONTAINER)
		case "list":
			list_containers()
		case "exec":
			exec_container()
		case "rm":
			rm_container()
		default:
			help()
		}
	default:
		help()
	}
}
<|endoftext|>


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

package cmd

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

	"github.com/spf13/cobra"
	"github.com/spf13/viper"
)

var initCmd = &cobra.Command{
	Use:     "init [name]",
	Aliases: []string{"initialize", "initialise", "create"},
	Short:   "Initialize a Cobra Application",
	Long: `Initialize (cobra init) will create a new application, with a license
and the appropriate structure for a Cobra-based CLI application.

  * If a name is provided, it will be created in the current directory;
  * If no name is provided, the current directory will be assumed;
  * If a relative path is provided, it will be created inside $GOPATH
    (e.g. github.com/spf13/hugo);
  * If an absolute path is provided, it will be created;
  * If the directory already exists but is empty, it will be used.

Init will not use an existing directory with contents.`,

	Run: func(cmd *cobra.Command, args []string) {
		var project *Project
		if len(args) == 0 {
			wd, err := os.Getwd()
			if err != nil {
				er(err)
			}
			project = NewProjectFromPath(wd)
		} else if len(args) == 1 {
			arg := args[0]
			if filepath.IsAbs(arg) {
				project = NewProjectFromPath(arg)
			} else {
				project = NewProject(arg)
			}
		} else {
			er("please enter the name")
		}

		initializeProject(project)

		fmt.Fprintln(cmd.OutOrStdout(), `Your Cobra application is ready at
`+project.AbsPath()+`.

Give it a try by going there and running `+"`go run main.go`."+`
Add commands to it by running `+"`cobra add [cmdname]`.")
	},
}

func initializeProject(project *Project) {
	if !exists(project.AbsPath()) { // If path doesn't yet exist, create it
		err := os.MkdirAll(project.AbsPath(), os.ModePerm)
		if err != nil {
			er(err)
		}
	} else if !isEmpty(project.AbsPath()) { // If path exists and is not empty don't use it
		er("Cobra will not create a new project in a non empty directory: " + project.AbsPath())
	}

	// We have a directory and it's empty. Time to initialize it.
	createLicenseFile(project.License(), project.AbsPath())
	createMainFile(project)
	createRootCmdFile(project)
}

func createLicenseFile(license License, path string) {
	data := make(map[string]interface{})
	data["copyright"] = copyrightLine()

	// Generate license template from text and data.
	text, err := executeTemplate(license.Text, data)
	if err != nil {
		er(err)
	}

	// Write license text to LICENSE file.
	err = writeStringToFile(filepath.Join(path, "LICENSE"), text)
	if err != nil {
		er(err)
	}
}

func createMainFile(project *Project) {
	mainTemplate := `{{ comment .copyright }}
{{if .license}}{{ comment .license }}{{end}}

package main

import "{{ .importpath }}"

func main() {
	cmd.Execute()
}
`
	data := make(map[string]interface{})
	data["copyright"] = copyrightLine()
	data["license"] = project.License().Header
	data["importpath"] = path.Join(project.Name(), filepath.Base(project.CmdPath()))

	mainScript, err := executeTemplate(mainTemplate, data)
	if err != nil {
		er(err)
	}

	err = writeStringToFile(filepath.Join(project.AbsPath(), "main.go"), mainScript)
	if err != nil {
		er(err)
	}
}

func createRootCmdFile(project *Project) {
	template := `{{comment .copyright}}
{{if .license}}{{comment .license}}{{end}}

package cmd

import (
	"fmt"
	"os"

	homedir "github.com/mitchellh/go-homedir"
	"github.com/spf13/cobra"
{{if .viper}}	"github.com/spf13/viper"{{end}}
)

{{if .viper}}var cfgFile string{{end}}

// RootCmd represents the base command when called without any subcommands
var RootCmd = &cobra.Command{
	Use:   "{{.appName}}",
	Short: "A brief description of your application",
	Long: ` + "`" + `A longer description that spans multiple lines and likely contains
examples and usage of using your application. For example:

Cobra is a CLI library for Go that empowers applications.
This application is a tool to generate the needed files
to quickly create a Cobra application.` + "`" + `,
	// Uncomment the following line if your bare application
	// has an action associated with it:
	//	Run: func(cmd *cobra.Command, args []string) { },
}

// Execute adds all child commands to the root command sets flags appropriately.
// This is called by main.main(). It only needs to happen once to the rootCmd.
func Execute() {
	if err := RootCmd.Execute(); err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
}

func init() {
{{if .viper}}	cobra.OnInitialize(initConfig){{end}}

	// Here you will define your flags and configuration settings.
	// Cobra supports persistent flags, which, if defined here,
	// will be global for your application.{{ if .viper }}
	RootCmd.PersistentFlags().StringVar(&cfgFile, "config", "", "config file (default is $HOME/.{{ .appName }}.yaml)"){{ else }}
	// RootCmd.PersistentFlags().StringVar(&cfgFile, "config", "", "config file (default is $HOME/.{{ .appName }}.yaml)"){{ end }}

	// Cobra also supports local flags, which will only run
	// when this action is called directly.
	RootCmd.Flags().BoolP("toggle", "t", false, "Help message for toggle")
}{{ if .viper }}

// initConfig reads in config file and ENV variables if set.
func initConfig() {
	if cfgFile != "" {
		// Use config file from the flag.
		viper.SetConfigFile(cfgFile)
	} else {
		// Find home directory.
		home, err := homedir.Dir()
		if err != nil {
			fmt.Println(home)
			os.Exit(1)
		}

		// Search config in home directory with name ".cobra" (without extension).
		viper.AddConfigPath(home)
		viper.SetConfigName(".cobra")
	}

	viper.AutomaticEnv() // read in environment variables that match

	// If a config file is found, read it in.
	if err := viper.ReadInConfig(); err == nil {
		fmt.Println("Using config file:", viper.ConfigFileUsed())
	}
}{{ end }}
`

	data := make(map[string]interface{})
	data["copyright"] = copyrightLine()
	data["viper"] = viper.GetBool("useViper")
	data["license"] = project.License().Header
	data["appName"] = path.Base(project.Name())

	rootCmdScript, err := executeTemplate(template, data)
	if err != nil {
		er(err)
	}

	err = writeStringToFile(filepath.Join(project.CmdPath(), "root.go"), rootCmdScript)
	if err != nil {
		er(err)
	}

}
<commit_msg>cmd: Fix incorrect initialzation by `cobra init .`<commit_after>// Copyright © 2015 Steve Francia <spf@spf13.com>.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package cmd

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

	"github.com/spf13/cobra"
	"github.com/spf13/viper"
)

var initCmd = &cobra.Command{
	Use:     "init [name]",
	Aliases: []string{"initialize", "initialise", "create"},
	Short:   "Initialize a Cobra Application",
	Long: `Initialize (cobra init) will create a new application, with a license
and the appropriate structure for a Cobra-based CLI application.

  * If a name is provided, it will be created in the current directory;
  * If no name is provided, the current directory will be assumed;
  * If a relative path is provided, it will be created inside $GOPATH
    (e.g. github.com/spf13/hugo);
  * If an absolute path is provided, it will be created;
  * If the directory already exists but is empty, it will be used.

Init will not use an existing directory with contents.`,

	Run: func(cmd *cobra.Command, args []string) {
		wd, err := os.Getwd()
		if err != nil {
			er(err)
		}

		var project *Project
		if len(args) == 0 {
			project = NewProjectFromPath(wd)
		} else if len(args) == 1 {
			arg := args[0]
			if arg[0] == '.' {
				arg = filepath.Join(wd, arg)
			}
			if filepath.IsAbs(arg) {
				project = NewProjectFromPath(arg)
			} else {
				project = NewProject(arg)
			}
		} else {
			er("please enter the name")
		}

		initializeProject(project)

		fmt.Fprintln(cmd.OutOrStdout(), `Your Cobra application is ready at
`+project.AbsPath()+`.

Give it a try by going there and running `+"`go run main.go`."+`
Add commands to it by running `+"`cobra add [cmdname]`.")
	},
}

func initializeProject(project *Project) {
	if !exists(project.AbsPath()) { // If path doesn't yet exist, create it
		err := os.MkdirAll(project.AbsPath(), os.ModePerm)
		if err != nil {
			er(err)
		}
	} else if !isEmpty(project.AbsPath()) { // If path exists and is not empty don't use it
		er("Cobra will not create a new project in a non empty directory: " + project.AbsPath())
	}

	// We have a directory and it's empty. Time to initialize it.
	createLicenseFile(project.License(), project.AbsPath())
	createMainFile(project)
	createRootCmdFile(project)
}

func createLicenseFile(license License, path string) {
	data := make(map[string]interface{})
	data["copyright"] = copyrightLine()

	// Generate license template from text and data.
	text, err := executeTemplate(license.Text, data)
	if err != nil {
		er(err)
	}

	// Write license text to LICENSE file.
	err = writeStringToFile(filepath.Join(path, "LICENSE"), text)
	if err != nil {
		er(err)
	}
}

func createMainFile(project *Project) {
	mainTemplate := `{{ comment .copyright }}
{{if .license}}{{ comment .license }}{{end}}

package main

import "{{ .importpath }}"

func main() {
	cmd.Execute()
}
`
	data := make(map[string]interface{})
	data["copyright"] = copyrightLine()
	data["license"] = project.License().Header
	data["importpath"] = path.Join(project.Name(), filepath.Base(project.CmdPath()))

	mainScript, err := executeTemplate(mainTemplate, data)
	if err != nil {
		er(err)
	}

	err = writeStringToFile(filepath.Join(project.AbsPath(), "main.go"), mainScript)
	if err != nil {
		er(err)
	}
}

func createRootCmdFile(project *Project) {
	template := `{{comment .copyright}}
{{if .license}}{{comment .license}}{{end}}

package cmd

import (
	"fmt"
	"os"

	homedir "github.com/mitchellh/go-homedir"
	"github.com/spf13/cobra"
{{if .viper}}	"github.com/spf13/viper"{{end}}
)

{{if .viper}}var cfgFile string{{end}}

// RootCmd represents the base command when called without any subcommands
var RootCmd = &cobra.Command{
	Use:   "{{.appName}}",
	Short: "A brief description of your application",
	Long: ` + "`" + `A longer description that spans multiple lines and likely contains
examples and usage of using your application. For example:

Cobra is a CLI library for Go that empowers applications.
This application is a tool to generate the needed files
to quickly create a Cobra application.` + "`" + `,
	// Uncomment the following line if your bare application
	// has an action associated with it:
	//	Run: func(cmd *cobra.Command, args []string) { },
}

// Execute adds all child commands to the root command sets flags appropriately.
// This is called by main.main(). It only needs to happen once to the rootCmd.
func Execute() {
	if err := RootCmd.Execute(); err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
}

func init() {
{{if .viper}}	cobra.OnInitialize(initConfig){{end}}

	// Here you will define your flags and configuration settings.
	// Cobra supports persistent flags, which, if defined here,
	// will be global for your application.{{ if .viper }}
	RootCmd.PersistentFlags().StringVar(&cfgFile, "config", "", "config file (default is $HOME/.{{ .appName }}.yaml)"){{ else }}
	// RootCmd.PersistentFlags().StringVar(&cfgFile, "config", "", "config file (default is $HOME/.{{ .appName }}.yaml)"){{ end }}

	// Cobra also supports local flags, which will only run
	// when this action is called directly.
	RootCmd.Flags().BoolP("toggle", "t", false, "Help message for toggle")
}{{ if .viper }}

// initConfig reads in config file and ENV variables if set.
func initConfig() {
	if cfgFile != "" {
		// Use config file from the flag.
		viper.SetConfigFile(cfgFile)
	} else {
		// Find home directory.
		home, err := homedir.Dir()
		if err != nil {
			fmt.Println(home)
			os.Exit(1)
		}

		// Search config in home directory with name ".cobra" (without extension).
		viper.AddConfigPath(home)
		viper.SetConfigName(".cobra")
	}

	viper.AutomaticEnv() // read in environment variables that match

	// If a config file is found, read it in.
	if err := viper.ReadInConfig(); err == nil {
		fmt.Println("Using config file:", viper.ConfigFileUsed())
	}
}{{ end }}
`

	data := make(map[string]interface{})
	data["copyright"] = copyrightLine()
	data["viper"] = viper.GetBool("useViper")
	data["license"] = project.License().Header
	data["appName"] = path.Base(project.Name())

	rootCmdScript, err := executeTemplate(template, data)
	if err != nil {
		er(err)
	}

	err = writeStringToFile(filepath.Join(project.CmdPath(), "root.go"), rootCmdScript)
	if err != nil {
		er(err)
	}

}
<|endoftext|>


<commit_before>/*
Copyright 2017 Gravitational, Inc.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// package multiplexer implements SSH and TLS multiplexing
// on the same listener
//
// mux, _ := multiplexer.New(Config{Listener: listener})
// mux.SSH() // returns listener getting SSH connections
// mux.TLS() // returns listener getting TLS connections
//
package multiplexer

import (
	"bufio"
	"bytes"
	"context"
	"net"
	"sync"
	"time"

	"github.com/gravitational/teleport/lib/defaults"

	"github.com/gravitational/trace"
	"github.com/jonboulle/clockwork"
	log "github.com/sirupsen/logrus"
)

// Config is a multiplexer config
type Config struct {
	// Listener is listener to multiplex connection on
	Listener net.Listener
	// Context is a context to signal stops, cancellations
	Context context.Context
	// ReadDeadline is a connection read deadline,
	// set to defaults.ReadHeadersTimeout if unspecified
	ReadDeadline time.Duration
	// Clock is a clock to override in tests, set to real time clock
	// by default
	Clock clockwork.Clock
}

// CheckAndSetDefaults verifies configuration and sets defaults
func (c *Config) CheckAndSetDefaults() error {
	if c.Listener == nil {
		return trace.BadParameter("missing parameter Listener")
	}
	if c.Context == nil {
		c.Context = context.TODO()
	}
	if c.ReadDeadline == 0 {
		c.ReadDeadline = defaults.ReadHeadersTimeout
	}
	if c.Clock == nil {
		c.Clock = clockwork.NewRealClock()
	}
	return nil
}

// New returns a new instance of multiplexer
func New(cfg Config) (*Mux, error) {
	if err := cfg.CheckAndSetDefaults(); err != nil {
		return nil, trace.Wrap(err)
	}
	ctx, cancel := context.WithCancel(cfg.Context)
	waitContext, waitCancel := context.WithCancel(context.TODO())
	return &Mux{
		Entry: log.WithFields(log.Fields{
			trace.Component: "mux",
		}),
		Config:      cfg,
		context:     ctx,
		cancel:      cancel,
		sshListener: newListener(ctx, cfg.Listener.Addr()),
		tlsListener: newListener(ctx, cfg.Listener.Addr()),
		waitContext: waitContext,
		waitCancel:  waitCancel,
	}, nil
}

// Mux supports having both SSH and TLS on the same listener socket
type Mux struct {
	sync.RWMutex
	*log.Entry
	Config
	listenerClosed bool
	sshListener    *Listener
	tlsListener    *Listener
	context        context.Context
	cancel         context.CancelFunc
	waitContext    context.Context
	waitCancel     context.CancelFunc
}

// SSH returns listener that receives SSH connections
func (m *Mux) SSH() net.Listener {
	return m.sshListener
}

// TLS returns listener that receives TLS connections
func (m *Mux) TLS() net.Listener {
	return m.tlsListener
}

func (m *Mux) isClosed() bool {
	m.RLock()
	defer m.RUnlock()
	return m.listenerClosed
}

func (m *Mux) closeListener() {
	m.Lock()
	defer m.Unlock()
	// propagate close signal to other listeners
	m.cancel()
	if m.Listener == nil {
		return
	}
	if m.listenerClosed {
		return
	}
	m.listenerClosed = true
	m.Listener.Close()
}

// Close closes listener
func (m *Mux) Close() error {
	m.closeListener()
	return nil
}

// Wait waits until listener shuts down and stops accepting new connections
// this is to workaround issue https://github.com/golang/go/issues/10527
// in tests
func (m *Mux) Wait() {
	<-m.waitContext.Done()
}

// Serve is a blocking function that serves on the listening socket
// and accepts requests. Every request is served in a separate goroutine
func (m *Mux) Serve() error {
	defer m.waitCancel()
	backoffTimer := time.NewTicker(5 * time.Second)
	defer backoffTimer.Stop()
	for {
		conn, err := m.Listener.Accept()
		if err == nil {
			go m.detectAndForward(conn)
			continue
		}
		if m.isClosed() {
			return nil
		}
		select {
		case <-backoffTimer.C:
			m.Debugf("backoff on accept error: %v", trace.DebugReport(err))
		case <-m.context.Done():
			return nil
		}
	}
}

func (m *Mux) detectAndForward(conn net.Conn) {
	err := conn.SetReadDeadline(m.Clock.Now().Add(m.ReadDeadline))
	if err != nil {
		m.Warning(err.Error())
		conn.Close()
		return
	}
	connWrapper, err := detect(conn)
	if err != nil {
		m.Warning(err.Error())
		conn.Close()
		return
	}

	err = conn.SetReadDeadline(time.Time{})
	if err != nil {
		m.Warning(err.Error())
		conn.Close()
		return
	}

	switch connWrapper.protocol {
	case ProtoTLS:
		select {
		case m.tlsListener.connC <- connWrapper:
		case <-m.context.Done():
			connWrapper.Close()
			return
		}
	case ProtoSSH:
		select {
		case m.sshListener.connC <- connWrapper:
		case <-m.context.Done():
			connWrapper.Close()
			return
		}
	default:
		// should not get here, handle this just in case
		connWrapper.Close()
		m.Errorf("detected but unsupported protocol: %v", connWrapper.protocol)
	}
}

func detect(conn net.Conn) (*Conn, error) {
	reader := bufio.NewReader(conn)

	var proxyLine *ProxyLine
	for i := 0; i < 2; i++ {
		bytes, err := reader.Peek(3)
		if err != nil {
			return nil, trace.Wrap(err, "failed to peek connection")
		}

		proto, err := detectProto(bytes)
		if err != nil {
			return nil, trace.Wrap(err)
		}

		switch proto {
		case ProtoProxy:
			if proxyLine != nil {
				return nil, trace.BadParameter("duplicate proxy line")
			}
			proxyLine, err = ReadProxyLine(reader)
			if err != nil {
				return nil, trace.Wrap(err)
			}
			// repeat the cycle to detect the protocol
		case ProtoTLS, ProtoSSH:
			return &Conn{
				protocol:  proto,
				Conn:      conn,
				reader:    reader,
				proxyLine: proxyLine,
			}, nil
		}
	}
	// if code ended here after two attempts, something is wrong
	return nil, trace.BadParameter("unknown protocol")
}

const (
	// ProtoUnknown is for unknown protocol
	ProtoUnknown = iota
	// ProtoTLS is TLS protocol
	ProtoTLS
	// ProtoSSH is SSH protocol
	ProtoSSH
	// ProtoProxy is a HAProxy proxy line protocol
	ProtoProxy
)

var (
	proxyPrefix = []byte{'P', 'R', 'O', 'X', 'Y'}
	sshPrefix   = []byte{'S', 'S', 'H'}
	tlsPrefix   = []byte{0x16}
)

func detectProto(in []byte) (int, error) {
	switch {
	// reader peeks only 3 bytes, slice the longer proxy prefix
	case bytes.HasPrefix(in, proxyPrefix[:3]):
		return ProtoProxy, nil
	case bytes.HasPrefix(in, sshPrefix):
		return ProtoSSH, nil
	case bytes.HasPrefix(in, tlsPrefix):
		return ProtoTLS, nil
	default:
		return ProtoUnknown, trace.BadParameter("failed to detect protocol")
	}
}
<commit_msg>add comment to explain two attempts<commit_after>/*
Copyright 2017 Gravitational, Inc.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// package multiplexer implements SSH and TLS multiplexing
// on the same listener
//
// mux, _ := multiplexer.New(Config{Listener: listener})
// mux.SSH() // returns listener getting SSH connections
// mux.TLS() // returns listener getting TLS connections
//
package multiplexer

import (
	"bufio"
	"bytes"
	"context"
	"net"
	"sync"
	"time"

	"github.com/gravitational/teleport/lib/defaults"

	"github.com/gravitational/trace"
	"github.com/jonboulle/clockwork"
	log "github.com/sirupsen/logrus"
)

// Config is a multiplexer config
type Config struct {
	// Listener is listener to multiplex connection on
	Listener net.Listener
	// Context is a context to signal stops, cancellations
	Context context.Context
	// ReadDeadline is a connection read deadline,
	// set to defaults.ReadHeadersTimeout if unspecified
	ReadDeadline time.Duration
	// Clock is a clock to override in tests, set to real time clock
	// by default
	Clock clockwork.Clock
}

// CheckAndSetDefaults verifies configuration and sets defaults
func (c *Config) CheckAndSetDefaults() error {
	if c.Listener == nil {
		return trace.BadParameter("missing parameter Listener")
	}
	if c.Context == nil {
		c.Context = context.TODO()
	}
	if c.ReadDeadline == 0 {
		c.ReadDeadline = defaults.ReadHeadersTimeout
	}
	if c.Clock == nil {
		c.Clock = clockwork.NewRealClock()
	}
	return nil
}

// New returns a new instance of multiplexer
func New(cfg Config) (*Mux, error) {
	if err := cfg.CheckAndSetDefaults(); err != nil {
		return nil, trace.Wrap(err)
	}
	ctx, cancel := context.WithCancel(cfg.Context)
	waitContext, waitCancel := context.WithCancel(context.TODO())
	return &Mux{
		Entry: log.WithFields(log.Fields{
			trace.Component: "mux",
		}),
		Config:      cfg,
		context:     ctx,
		cancel:      cancel,
		sshListener: newListener(ctx, cfg.Listener.Addr()),
		tlsListener: newListener(ctx, cfg.Listener.Addr()),
		waitContext: waitContext,
		waitCancel:  waitCancel,
	}, nil
}

// Mux supports having both SSH and TLS on the same listener socket
type Mux struct {
	sync.RWMutex
	*log.Entry
	Config
	listenerClosed bool
	sshListener    *Listener
	tlsListener    *Listener
	context        context.Context
	cancel         context.CancelFunc
	waitContext    context.Context
	waitCancel     context.CancelFunc
}

// SSH returns listener that receives SSH connections
func (m *Mux) SSH() net.Listener {
	return m.sshListener
}

// TLS returns listener that receives TLS connections
func (m *Mux) TLS() net.Listener {
	return m.tlsListener
}

func (m *Mux) isClosed() bool {
	m.RLock()
	defer m.RUnlock()
	return m.listenerClosed
}

func (m *Mux) closeListener() {
	m.Lock()
	defer m.Unlock()
	// propagate close signal to other listeners
	m.cancel()
	if m.Listener == nil {
		return
	}
	if m.listenerClosed {
		return
	}
	m.listenerClosed = true
	m.Listener.Close()
}

// Close closes listener
func (m *Mux) Close() error {
	m.closeListener()
	return nil
}

// Wait waits until listener shuts down and stops accepting new connections
// this is to workaround issue https://github.com/golang/go/issues/10527
// in tests
func (m *Mux) Wait() {
	<-m.waitContext.Done()
}

// Serve is a blocking function that serves on the listening socket
// and accepts requests. Every request is served in a separate goroutine
func (m *Mux) Serve() error {
	defer m.waitCancel()
	backoffTimer := time.NewTicker(5 * time.Second)
	defer backoffTimer.Stop()
	for {
		conn, err := m.Listener.Accept()
		if err == nil {
			go m.detectAndForward(conn)
			continue
		}
		if m.isClosed() {
			return nil
		}
		select {
		case <-backoffTimer.C:
			m.Debugf("backoff on accept error: %v", trace.DebugReport(err))
		case <-m.context.Done():
			return nil
		}
	}
}

func (m *Mux) detectAndForward(conn net.Conn) {
	err := conn.SetReadDeadline(m.Clock.Now().Add(m.ReadDeadline))
	if err != nil {
		m.Warning(err.Error())
		conn.Close()
		return
	}
	connWrapper, err := detect(conn)
	if err != nil {
		m.Warning(err.Error())
		conn.Close()
		return
	}

	err = conn.SetReadDeadline(time.Time{})
	if err != nil {
		m.Warning(err.Error())
		conn.Close()
		return
	}

	switch connWrapper.protocol {
	case ProtoTLS:
		select {
		case m.tlsListener.connC <- connWrapper:
		case <-m.context.Done():
			connWrapper.Close()
			return
		}
	case ProtoSSH:
		select {
		case m.sshListener.connC <- connWrapper:
		case <-m.context.Done():
			connWrapper.Close()
			return
		}
	default:
		// should not get here, handle this just in case
		connWrapper.Close()
		m.Errorf("detected but unsupported protocol: %v", connWrapper.protocol)
	}
}

func detect(conn net.Conn) (*Conn, error) {
	reader := bufio.NewReader(conn)

	// the first attempt is to parse optional proxy
	// protocol line that is injected by load balancers
	// before actual protocol traffic flows.
	// if the first attempt encounters proxy it
	// goes to the second pass to do protocol detection
	var proxyLine *ProxyLine
	for i := 0; i < 2; i++ {
		bytes, err := reader.Peek(3)
		if err != nil {
			return nil, trace.Wrap(err, "failed to peek connection")
		}

		proto, err := detectProto(bytes)
		if err != nil {
			return nil, trace.Wrap(err)
		}

		switch proto {
		case ProtoProxy:
			if proxyLine != nil {
				return nil, trace.BadParameter("duplicate proxy line")
			}
			proxyLine, err = ReadProxyLine(reader)
			if err != nil {
				return nil, trace.Wrap(err)
			}
			// repeat the cycle to detect the protocol
		case ProtoTLS, ProtoSSH:
			return &Conn{
				protocol:  proto,
				Conn:      conn,
				reader:    reader,
				proxyLine: proxyLine,
			}, nil
		}
	}
	// if code ended here after two attempts, something is wrong
	return nil, trace.BadParameter("unknown protocol")
}

const (
	// ProtoUnknown is for unknown protocol
	ProtoUnknown = iota
	// ProtoTLS is TLS protocol
	ProtoTLS
	// ProtoSSH is SSH protocol
	ProtoSSH
	// ProtoProxy is a HAProxy proxy line protocol
	ProtoProxy
)

var (
	proxyPrefix = []byte{'P', 'R', 'O', 'X', 'Y'}
	sshPrefix   = []byte{'S', 'S', 'H'}
	tlsPrefix   = []byte{0x16}
)

func detectProto(in []byte) (int, error) {
	switch {
	// reader peeks only 3 bytes, slice the longer proxy prefix
	case bytes.HasPrefix(in, proxyPrefix[:3]):
		return ProtoProxy, nil
	case bytes.HasPrefix(in, sshPrefix):
		return ProtoSSH, nil
	case bytes.HasPrefix(in, tlsPrefix):
		return ProtoTLS, nil
	default:
		return ProtoUnknown, trace.BadParameter("failed to detect protocol")
	}
}
<|endoftext|>


<commit_before>package metakv

import (
	"encoding/binary"
	"encoding/json"
	"io/ioutil"
	"log"
	"net/http"
	"net/http/httptest"
	"net/url"
	"sort"
	"strings"
	"sync"
	"testing"
)

type entry struct {
	v []byte
	r []byte
}

type mockKV struct {
	l           sync.Mutex
	counter     uint64
	data        map[string]entry
	subscribers map[uint64]chan KVEntry
	srv         *httptest.Server
}

func (kv *mockKV) runMock() func() {
	srv := httptest.NewServer(http.HandlerFunc(kv.Handle))
	if kv.data == nil {
		kv.data = make(map[string]entry)
		kv.subscribers = make(map[uint64]chan KVEntry)
	}
	kv.srv = srv
	return func() {
		srv.Close()
	}
}

func replyJSON(w http.ResponseWriter, value interface{}) {
	json.NewEncoder(w).Encode(value)
}

func (kv *mockKV) broadcast(kve KVEntry) {
	for _, s := range kv.subscribers {
		s <- kve
	}
}

func (kv *mockKV) setLocked(path string, value string) {
	rev := make([]byte, 8)
	binary.LittleEndian.PutUint64(rev, kv.counter)
	kv.counter++
	v := []byte(value)
	e := entry{v, rev}
	kv.data[path] = e

	kv.broadcast(KVEntry{Path: path, Value: v, Rev: rev})
}

func (kv *mockKV) subscribeLocked(ch chan KVEntry) func() {
	id := kv.counter
	kv.counter++
	kv.subscribers[id] = ch
	return func() {
		kv.l.Lock()
		defer kv.l.Unlock()
		delete(kv.subscribers, id)
	}
}

type entriesSlice []KVEntry

func (p entriesSlice) Len() int {
	return len(p)
}
func (p entriesSlice) Less(i, j int) bool {
	return string(p[i].Path) < string(p[j].Path)
}
func (p entriesSlice) Swap(i, j int) {
	p[i], p[j] = p[j], p[i]
}

func (kv *mockKV) checkRevision(rev, path string) bool {
	if rev == "" {
		return true
	}
	e := kv.data[path]
	return string(e.r) == rev
}

func (kv *mockKV) Handle(w http.ResponseWriter, req *http.Request) {
	path := strings.TrimPrefix(req.URL.Path, "/_metakv")
	if path == req.URL.Path {
		panic("Prefix /_metakv is not found")
	}
	isDir := strings.HasSuffix(path, "/")

	if req.Method == "GET" && isDir {
		kv.handleIterate(w, req)
		return
	}

	kv.l.Lock()
	defer kv.l.Unlock()

	switch req.Method {
	case "GET":
		e, exists := kv.data[path]
		if !exists {
			w.Write([]byte("{}"))
			return
		}
		replyJSON(w, map[string][]byte{"value": e.v, "rev": e.r})
	case "PUT":
		req.ParseForm()

		form := req.PostForm
		create := form.Get("create") != ""
		value := form.Get("value")
		rev := form.Get("rev")

		if !kv.checkRevision(rev, path) {
			w.WriteHeader(409)
			return
		}

		if create {
			if _, exists := kv.data[path]; exists {
				w.WriteHeader(409)
				return
			}
		}
		kv.setLocked(path, value)
	case "DELETE":
		rev := req.URL.Query().Get("rev")

		if !kv.checkRevision(rev, path) {
			w.WriteHeader(409)
			return
		}

		kv.broadcast(KVEntry{path, nil, nil})
		delete(kv.data, path)
	default:
		w.WriteHeader(404)
	}
}

func (kv *mockKV) handleIterate(w http.ResponseWriter, req *http.Request) {
	kv.l.Lock()
	locked := true
	defer func() {
		if locked {
			kv.l.Unlock()
		}
	}()

	continuous := req.URL.Query().Get("feed") == "continuous"
	entries := make([]KVEntry, 0, len(kv.data))
	for k, e := range kv.data {
		entries = append(entries, KVEntry{Path: k, Value: e.v, Rev: e.r})
	}
	sort.Sort(entriesSlice(entries))
	enc := json.NewEncoder(w)
	for _, e := range entries {
		err := enc.Encode(e)
		if err != nil {
			log.Fatal(err)
		}
	}
	if continuous {
		w.(http.Flusher).Flush()

		ch := make(chan KVEntry, 16)
		defer kv.subscribeLocked(ch)()

		kv.l.Unlock()
		locked = false

		log.Print("Waiting for rows")
		closed := w.(http.CloseNotifier).CloseNotify()

		for {
			select {
			case e := <-ch:
				log.Printf("Observed {%s, %s, %s}", e.Path, e.Value, e.Rev)
				err := enc.Encode(e)
				if err != nil {
					log.Printf("Got error in subscribe path: %v", err)
					return
				}
				w.(http.Flusher).Flush()
			case <-closed:
				log.Print("receiver is dead")
				return
			}
		}
	}
}

type myT struct{ *testing.T }

func (t *myT) okStatus(statusCode int, err error) {
	if err != nil {
		t.Fatalf("Got error from http call: %v", err)
	}
	if statusCode != 200 {
		t.Fatalf("Expected code 200. Got: %d", statusCode)
	}
}

func (t *myT) emptyBody(resp *http.Response, err error) {
	defer resp.Body.Close()
	t.okStatus(resp.StatusCode, err)
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		t.Fatalf("Got error trying to read body: %v", err)
	}
	if len(body) != 0 {
		t.Fatalf("Expected empty body. Got: `%s'", string(body))
	}
}

func must(t *testing.T) *myT { return &myT{t} }

func (kv *mockKV) fullPath(path string) string {
	return kv.srv.URL + "/_metakv" + path
}

func (kv *mockKV) doGet(path string, response interface{}) (statusCode int, err error) {
	resp, err := http.Get(kv.fullPath(path))
	if err != nil {
		return 0, err
	}
	defer resp.Body.Close()
	return resp.StatusCode, json.NewDecoder(resp.Body).Decode(response)
}

func (kv *mockKV) doPut(path, value string) (resp *http.Response, err error) {
	values := url.Values{"value": {"foobar"}}
	body := strings.NewReader(values.Encode())
	req, err := http.NewRequest("PUT", kv.fullPath(path), body)
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	if err != nil {
		return nil, err
	}
	client := &http.Client{}
	return client.Do(req)
}

func TestMock(t *testing.T) {
	kv := &mockKV{}
	defer kv.runMock()()

	var m map[string]interface{}
	must(t).okStatus(kv.doGet("/test", &m))
	if len(m) != 0 {
		t.Fatalf("Expected get against empty kv to return {}")
	}

	must(t).emptyBody(kv.doPut("/test", "foobar"))

	var kve kvEntry
	must(t).okStatus(kv.doGet("/test", &kve))
	if string(kve.Value) != "foobar" {
		t.Fatalf("failed to get expected value (foobar). Got: %s", kve.Value)
	}
}

func TestSanity(t *testing.T) {
	kv := &mockKV{}
	defer kv.runMock()()

	mockStore := &store{
		url:    kv.srv.URL + "/_metakv",
		client: http.DefaultClient,
	}

	if err := mockStore.add("/_sanity/garbage", []byte("v"), false); err != nil {
		t.Logf("add failed with: %v", err)
	}
	doExecuteBasicSanityTest(t.Log, mockStore)
}
<commit_msg>fix compilation error<commit_after>package metakv

import (
	"encoding/binary"
	"encoding/json"
	"io/ioutil"
	"log"
	"net/http"
	"net/http/httptest"
	"net/url"
	"sort"
	"strings"
	"sync"
	"testing"
)

type entry struct {
	v []byte
	r []byte
}

type mockKV struct {
	l           sync.Mutex
	counter     uint64
	data        map[string]entry
	subscribers map[uint64]chan KVEntry
	srv         *httptest.Server
}

func (kv *mockKV) runMock() func() {
	srv := httptest.NewServer(http.HandlerFunc(kv.Handle))
	if kv.data == nil {
		kv.data = make(map[string]entry)
		kv.subscribers = make(map[uint64]chan KVEntry)
	}
	kv.srv = srv
	return func() {
		srv.Close()
	}
}

func replyJSON(w http.ResponseWriter, value interface{}) {
	json.NewEncoder(w).Encode(value)
}

func (kv *mockKV) broadcast(kve KVEntry) {
	for _, s := range kv.subscribers {
		s <- kve
	}
}

func (kv *mockKV) setLocked(path string, value string) {
	rev := make([]byte, 8)
	binary.LittleEndian.PutUint64(rev, kv.counter)
	kv.counter++
	v := []byte(value)
	e := entry{v, rev}
	kv.data[path] = e

	kv.broadcast(KVEntry{Path: path, Value: v, Rev: rev})
}

func (kv *mockKV) subscribeLocked(ch chan KVEntry) func() {
	id := kv.counter
	kv.counter++
	kv.subscribers[id] = ch
	return func() {
		kv.l.Lock()
		defer kv.l.Unlock()
		delete(kv.subscribers, id)
	}
}

type entriesSlice []KVEntry

func (p entriesSlice) Len() int {
	return len(p)
}
func (p entriesSlice) Less(i, j int) bool {
	return string(p[i].Path) < string(p[j].Path)
}
func (p entriesSlice) Swap(i, j int) {
	p[i], p[j] = p[j], p[i]
}

func (kv *mockKV) checkRevision(rev, path string) bool {
	if rev == "" {
		return true
	}
	e := kv.data[path]
	return string(e.r) == rev
}

func (kv *mockKV) Handle(w http.ResponseWriter, req *http.Request) {
	path := strings.TrimPrefix(req.URL.Path, "/_metakv")
	if path == req.URL.Path {
		panic("Prefix /_metakv is not found")
	}
	isDir := strings.HasSuffix(path, "/")

	if req.Method == "GET" && isDir {
		kv.handleIterate(w, req)
		return
	}

	kv.l.Lock()
	defer kv.l.Unlock()

	switch req.Method {
	case "GET":
		e, exists := kv.data[path]
		if !exists {
			w.Write([]byte("{}"))
			return
		}
		replyJSON(w, map[string][]byte{"value": e.v, "rev": e.r})
	case "PUT":
		req.ParseForm()

		form := req.PostForm
		create := form.Get("create") != ""
		value := form.Get("value")
		rev := form.Get("rev")

		if !kv.checkRevision(rev, path) {
			w.WriteHeader(409)
			return
		}

		if create {
			if _, exists := kv.data[path]; exists {
				w.WriteHeader(409)
				return
			}
		}
		kv.setLocked(path, value)
	case "DELETE":
		rev := req.URL.Query().Get("rev")

		if !kv.checkRevision(rev, path) {
			w.WriteHeader(409)
			return
		}

		kv.broadcast(KVEntry{path, nil, nil})
		delete(kv.data, path)
	default:
		w.WriteHeader(404)
	}
}

func (kv *mockKV) handleIterate(w http.ResponseWriter, req *http.Request) {
	kv.l.Lock()
	locked := true
	defer func() {
		if locked {
			kv.l.Unlock()
		}
	}()

	continuous := req.URL.Query().Get("feed") == "continuous"
	entries := make([]KVEntry, 0, len(kv.data))
	for k, e := range kv.data {
		entries = append(entries, KVEntry{Path: k, Value: e.v, Rev: e.r})
	}
	sort.Sort(entriesSlice(entries))
	enc := json.NewEncoder(w)
	for _, e := range entries {
		err := enc.Encode(e)
		if err != nil {
			log.Fatal(err)
		}
	}
	if continuous {
		w.(http.Flusher).Flush()

		ch := make(chan KVEntry, 16)
		defer kv.subscribeLocked(ch)()

		kv.l.Unlock()
		locked = false

		log.Print("Waiting for rows")
		closed := w.(http.CloseNotifier).CloseNotify()

		for {
			select {
			case e := <-ch:
				log.Printf("Observed {%s, %s, %s}", e.Path, e.Value, e.Rev)
				err := enc.Encode(e)
				if err != nil {
					log.Printf("Got error in subscribe path: %v", err)
					return
				}
				w.(http.Flusher).Flush()
			case <-closed:
				log.Print("receiver is dead")
				return
			}
		}
	}
}

type myT struct{ *testing.T }

func (t *myT) okStatus(statusCode int, err error) {
	if err != nil {
		t.Fatalf("Got error from http call: %v", err)
	}
	if statusCode != 200 {
		t.Fatalf("Expected code 200. Got: %d", statusCode)
	}
}

func (t *myT) emptyBody(resp *http.Response, err error) {
	defer resp.Body.Close()
	t.okStatus(resp.StatusCode, err)
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		t.Fatalf("Got error trying to read body: %v", err)
	}
	if len(body) != 0 {
		t.Fatalf("Expected empty body. Got: `%s'", string(body))
	}
}

func must(t *testing.T) *myT { return &myT{t} }

func (kv *mockKV) fullPath(path string) string {
	return kv.srv.URL + "/_metakv" + path
}

func (kv *mockKV) doGet(path string, response interface{}) (statusCode int, err error) {
	resp, err := http.Get(kv.fullPath(path))
	if err != nil {
		return 0, err
	}
	defer resp.Body.Close()
	return resp.StatusCode, json.NewDecoder(resp.Body).Decode(response)
}

func (kv *mockKV) doPut(path, value string) (resp *http.Response, err error) {
	values := url.Values{"value": {"foobar"}}
	body := strings.NewReader(values.Encode())
	req, err := http.NewRequest("PUT", kv.fullPath(path), body)
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	if err != nil {
		return nil, err
	}
	client := &http.Client{}
	return client.Do(req)
}

func TestMock(t *testing.T) {
	kv := &mockKV{}
	defer kv.runMock()()

	var m map[string]interface{}
	must(t).okStatus(kv.doGet("/test", &m))
	if len(m) != 0 {
		t.Fatalf("Expected get against empty kv to return {}")
	}

	must(t).emptyBody(kv.doPut("/test", "foobar"))

	var kve kvEntry
	must(t).okStatus(kv.doGet("/test", &kve))
	if string(kve.Value) != "foobar" {
		t.Fatalf("failed to get expected value (foobar). Got: %s", kve.Value)
	}
}

func TestSanity(t *testing.T) {
	kv := &mockKV{}
	defer kv.runMock()()

	url, err := url.Parse(kv.srv.URL + "/_metakv")
	if err != nil {
		panic(err)
	}

	mockStore := &store{
		url:    url,
		client: http.DefaultClient,
	}

	if err := mockStore.add("/_sanity/garbage", []byte("v"), false); err != nil {
		t.Logf("add failed with: %v", err)
	}
	doExecuteBasicSanityTest(t.Log, mockStore)
}
<|endoftext|>


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

package v1

import (
	restful "github.com/emicklei/go-restful"

	"k8s.io/heapster/metrics/api/v1/types"
	"k8s.io/heapster/metrics/core"
	"k8s.io/heapster/metrics/sinks/metric"
)

type Api struct {
	runningInKubernetes bool
	metricSink          *metricsink.MetricSink
	gkeMetrics          map[string]core.MetricDescriptor
	gkeLabels           map[string]core.LabelDescriptor
}

// Create a new Api to serve from the specified cache.
func NewApi(runningInKubernetes bool, metricSink *metricsink.MetricSink) *Api {
	gkeMetrics := make(map[string]core.MetricDescriptor)
	gkeLabels := make(map[string]core.LabelDescriptor)
	for _, val := range core.StandardMetrics {
		gkeMetrics[val.Name] = val.MetricDescriptor
	}
	for _, val := range core.CommonLabels() {
		gkeLabels[val.Key] = val
	}
	for _, val := range core.ContainerLabels() {
		gkeLabels[val.Key] = val
	}
	for _, val := range core.PodLabels() {
		gkeLabels[val.Key] = val
	}

	return &Api{
		runningInKubernetes: runningInKubernetes,
		metricSink:          metricSink,
		gkeMetrics:          gkeMetrics,
		gkeLabels:           gkeLabels,
	}
}

// Register the mainApi on the specified endpoint.
func (a *Api) Register(container *restful.Container) {
	ws := new(restful.WebService)
	ws.Path("/api/v1/metric-export").
		Doc("Exports the latest point for all Heapster metrics").
		Produces(restful.MIME_JSON)
	ws.Route(ws.GET("").
		To(a.exportMetrics).
		Doc("export the latest data point for all metrics").
		Operation("exportMetrics").
		Writes([]*types.Timeseries{}))
	container.Add(ws)
	ws = new(restful.WebService)
	ws.Path("/api/v1/metric-export-schema").
		Doc("Schema for metrics exported by heapster").
		Produces(restful.MIME_JSON)
	ws.Route(ws.GET("").
		To(a.exportMetricsSchema).
		Doc("export the schema for all metrics").
		Operation("exportmetricsSchema").
		Writes(types.TimeseriesSchema{}))
	container.Add(ws)

	if a.metricSink != nil {
		a.RegisterModel(container)
	}
}

func convertLabelDescriptor(ld core.LabelDescriptor) types.LabelDescriptor {
	return types.LabelDescriptor{
		Key:         ld.Key,
		Description: ld.Description,
	}
}

func convertMetricDescriptor(md core.MetricDescriptor) types.MetricDescriptor {
	result := types.MetricDescriptor{
		Name:        md.Name,
		Description: md.Description,
		Labels:      make([]types.LabelDescriptor, 0, len(md.Labels)),
	}
	for _, label := range md.Labels {
		result.Labels = append(result.Labels, convertLabelDescriptor(label))
	}

	switch md.Type {
	case core.MetricCumulative:
		result.Type = "cumulative"
	case core.MetricGauge:
		result.Type = "gauge"
	case core.MetricDelta:
		result.Type = "delta"
	}

	switch md.ValueType {
	case core.ValueInt64:
		result.ValueType = "int64"
	case core.ValueFloat:
		result.ValueType = "double"
	}

	switch md.Units {
	case core.UnitsBytes:
		result.Units = "bytes"
	case core.UnitsMilliseconds:
		result.Units = "ms"
	case core.UnitsNanoseconds:
		result.Units = "ns"
	case core.UnitsMillicores:
		result.Units = "millicores"
	}
	return result
}

func (a *Api) exportMetricsSchema(request *restful.Request, response *restful.Response) {
	result := types.TimeseriesSchema{
		Metrics:      make([]types.MetricDescriptor, 0),
		CommonLabels: make([]types.LabelDescriptor, 0),
		PodLabels:    make([]types.LabelDescriptor, 0),
	}
	for _, metric := range core.StandardMetrics {
		if _, found := a.gkeMetrics[metric.Name]; found {
			result.Metrics = append(result.Metrics, convertMetricDescriptor(metric.MetricDescriptor))
		}
	}
	for _, label := range core.CommonLabels() {
		if _, found := a.gkeLabels[label.Key]; found {
			result.PodLabels = append(result.PodLabels, convertLabelDescriptor(label))
		}
	}
	for _, label := range core.PodLabels() {
		if _, found := a.gkeLabels[label.Key]; found {
			result.PodLabels = append(result.PodLabels, convertLabelDescriptor(label))
		}
	}
	response.WriteEntity(result)
}

func (a *Api) exportMetrics(request *restful.Request, response *restful.Response) {
	shortStorage := a.metricSink.GetShortStore()
	tsmap := make(map[string]*types.Timeseries)

	for _, batch := range shortStorage {
		for key, ms := range batch.MetricSets {
			ts := tsmap[key]

			msType := ms.Labels[core.LabelMetricSetType.Key]

			if msType != core.MetricSetTypeNode &&
				msType != core.MetricSetTypePodContainer &&
				msType != core.MetricSetTypeSystemContainer {
				continue
			}

			if ts == nil {
				ts = &types.Timeseries{
					Metrics: make(map[string][]types.Point),
					Labels:  make(map[string]string),
				}
				for labelName, labelValue := range ms.Labels {
					if _, ok := a.gkeLabels[labelName]; ok {
						ts.Labels[labelName] = labelValue
					}
				}
				if msType == core.MetricSetTypeNode {
					ts.Labels[core.LabelContainerName.Key] = "machine"
				}
				tsmap[key] = ts
			}
			for metricName, metricVal := range ms.MetricValues {
				if _, ok := a.gkeMetrics[metricName]; ok {
					points := ts.Metrics[metricName]
					if points == nil {
						points = make([]types.Point, 0, len(shortStorage))
					}
					point := types.Point{
						Start: batch.Timestamp,
						End:   batch.Timestamp,
					}
					if metricVal.ValueType == core.ValueInt64 {
						point.Value = &metricVal.IntValue
					} else if metricVal.ValueType == core.ValueFloat {
						point.Value = &metricVal.FloatValue
					} else {
						continue
					}
					points = append(points, point)
					ts.Metrics[metricName] = points
				}
			}
		}
	}
	timeseries := make([]*types.Timeseries, 0, len(tsmap))
	for _, ts := range tsmap {
		timeseries = append(timeseries, ts)
	}

	response.WriteEntity(timeseries)
}
<commit_msg>Fix metric value setting in GKE<commit_after>// Copyright 2015 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package v1

import (
	restful "github.com/emicklei/go-restful"

	"k8s.io/heapster/metrics/api/v1/types"
	"k8s.io/heapster/metrics/core"
	"k8s.io/heapster/metrics/sinks/metric"
)

type Api struct {
	runningInKubernetes bool
	metricSink          *metricsink.MetricSink
	gkeMetrics          map[string]core.MetricDescriptor
	gkeLabels           map[string]core.LabelDescriptor
}

// Create a new Api to serve from the specified cache.
func NewApi(runningInKubernetes bool, metricSink *metricsink.MetricSink) *Api {
	gkeMetrics := make(map[string]core.MetricDescriptor)
	gkeLabels := make(map[string]core.LabelDescriptor)
	for _, val := range core.StandardMetrics {
		gkeMetrics[val.Name] = val.MetricDescriptor
	}
	for _, val := range core.CommonLabels() {
		gkeLabels[val.Key] = val
	}
	for _, val := range core.ContainerLabels() {
		gkeLabels[val.Key] = val
	}
	for _, val := range core.PodLabels() {
		gkeLabels[val.Key] = val
	}

	return &Api{
		runningInKubernetes: runningInKubernetes,
		metricSink:          metricSink,
		gkeMetrics:          gkeMetrics,
		gkeLabels:           gkeLabels,
	}
}

// Register the mainApi on the specified endpoint.
func (a *Api) Register(container *restful.Container) {
	ws := new(restful.WebService)
	ws.Path("/api/v1/metric-export").
		Doc("Exports the latest point for all Heapster metrics").
		Produces(restful.MIME_JSON)
	ws.Route(ws.GET("").
		To(a.exportMetrics).
		Doc("export the latest data point for all metrics").
		Operation("exportMetrics").
		Writes([]*types.Timeseries{}))
	container.Add(ws)
	ws = new(restful.WebService)
	ws.Path("/api/v1/metric-export-schema").
		Doc("Schema for metrics exported by heapster").
		Produces(restful.MIME_JSON)
	ws.Route(ws.GET("").
		To(a.exportMetricsSchema).
		Doc("export the schema for all metrics").
		Operation("exportmetricsSchema").
		Writes(types.TimeseriesSchema{}))
	container.Add(ws)

	if a.metricSink != nil {
		a.RegisterModel(container)
	}
}

func convertLabelDescriptor(ld core.LabelDescriptor) types.LabelDescriptor {
	return types.LabelDescriptor{
		Key:         ld.Key,
		Description: ld.Description,
	}
}

func convertMetricDescriptor(md core.MetricDescriptor) types.MetricDescriptor {
	result := types.MetricDescriptor{
		Name:        md.Name,
		Description: md.Description,
		Labels:      make([]types.LabelDescriptor, 0, len(md.Labels)),
	}
	for _, label := range md.Labels {
		result.Labels = append(result.Labels, convertLabelDescriptor(label))
	}

	switch md.Type {
	case core.MetricCumulative:
		result.Type = "cumulative"
	case core.MetricGauge:
		result.Type = "gauge"
	case core.MetricDelta:
		result.Type = "delta"
	}

	switch md.ValueType {
	case core.ValueInt64:
		result.ValueType = "int64"
	case core.ValueFloat:
		result.ValueType = "double"
	}

	switch md.Units {
	case core.UnitsBytes:
		result.Units = "bytes"
	case core.UnitsMilliseconds:
		result.Units = "ms"
	case core.UnitsNanoseconds:
		result.Units = "ns"
	case core.UnitsMillicores:
		result.Units = "millicores"
	}
	return result
}

func (a *Api) exportMetricsSchema(request *restful.Request, response *restful.Response) {
	result := types.TimeseriesSchema{
		Metrics:      make([]types.MetricDescriptor, 0),
		CommonLabels: make([]types.LabelDescriptor, 0),
		PodLabels:    make([]types.LabelDescriptor, 0),
	}
	for _, metric := range core.StandardMetrics {
		if _, found := a.gkeMetrics[metric.Name]; found {
			result.Metrics = append(result.Metrics, convertMetricDescriptor(metric.MetricDescriptor))
		}
	}
	for _, label := range core.CommonLabels() {
		if _, found := a.gkeLabels[label.Key]; found {
			result.PodLabels = append(result.PodLabels, convertLabelDescriptor(label))
		}
	}
	for _, label := range core.PodLabels() {
		if _, found := a.gkeLabels[label.Key]; found {
			result.PodLabels = append(result.PodLabels, convertLabelDescriptor(label))
		}
	}
	response.WriteEntity(result)
}

func (a *Api) exportMetrics(request *restful.Request, response *restful.Response) {
	shortStorage := a.metricSink.GetShortStore()
	tsmap := make(map[string]*types.Timeseries)

	for _, batch := range shortStorage {
		for key, ms := range batch.MetricSets {
			ts := tsmap[key]

			msType := ms.Labels[core.LabelMetricSetType.Key]

			if msType != core.MetricSetTypeNode &&
				msType != core.MetricSetTypePodContainer &&
				msType != core.MetricSetTypeSystemContainer {
				continue
			}

			if ts == nil {
				ts = &types.Timeseries{
					Metrics: make(map[string][]types.Point),
					Labels:  make(map[string]string),
				}
				for labelName, labelValue := range ms.Labels {
					if _, ok := a.gkeLabels[labelName]; ok {
						ts.Labels[labelName] = labelValue
					}
				}
				if msType == core.MetricSetTypeNode {
					ts.Labels[core.LabelContainerName.Key] = "machine"
				}
				tsmap[key] = ts
			}
			for metricName, metricVal := range ms.MetricValues {
				if _, ok := a.gkeMetrics[metricName]; ok {
					points := ts.Metrics[metricName]
					if points == nil {
						points = make([]types.Point, 0, len(shortStorage))
					}
					point := types.Point{
						Start: batch.Timestamp,
						End:   batch.Timestamp,
					}
					var value interface{}
					if metricVal.ValueType == core.ValueInt64 {
						value = metricVal.IntValue
					} else if metricVal.ValueType == core.ValueFloat {
						value = metricVal.FloatValue
					} else {
						continue
					}
					point.Value = value
					points = append(points, point)
					ts.Metrics[metricName] = points
				}
			}
		}
	}
	timeseries := make([]*types.Timeseries, 0, len(tsmap))
	for _, ts := range tsmap {
		timeseries = append(timeseries, ts)
	}

	response.WriteEntity(timeseries)
}
<|endoftext|>


<commit_before>/*
   Copyright 2014 Outbrain Inc.

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
*/

package main

import (
	"flag"
	"fmt"
	"github.com/outbrain/golib/log"
	"github.com/outbrain/zookeepercli/output"
	"github.com/outbrain/zookeepercli/zk"
	"io/ioutil"
	"math/rand"
	"os"
	"strings"
	"time"
)

// main is the application's entry point.
func main() {
	servers := flag.String("servers", "", "srv1[:port1][,srv2[:port2]...]")
	command := flag.String("c", "", "command, required (exists|get|ls|lsr|create|creater|set|delete|rm|deleter|rmr|getacl|setacl)")
	force := flag.Bool("force", false, "force operation")
	format := flag.String("format", "txt", "output format (txt|json)")
	verbose := flag.Bool("verbose", false, "verbose")
	debug := flag.Bool("debug", false, "debug mode (very verbose)")
	stack := flag.Bool("stack", false, "add stack trace upon error")
	authUser := flag.String("auth_usr", "", "optional, digest scheme, user")
	authPwd := flag.String("auth_pwd", "", "optional, digest scheme, pwd")
	acls := flag.String("acls", "31", "optional, csv list [1|,2|,4|,8|,16|,31]")
	flag.Parse()

	log.SetLevel(log.ERROR)
	if *verbose {
		log.SetLevel(log.INFO)
	}
	if *debug {
		log.SetLevel(log.DEBUG)
	}
	if *stack {
		log.SetPrintStackTrace(*stack)
	}

	log.Info("starting")

	if *servers == "" {
		log.Fatal("Expected comma delimited list of servers via --servers")
	}
	serversArray := strings.Split(*servers, ",")
	if len(serversArray) == 0 {
		log.Fatal("Expected comma delimited list of servers via --servers")
	}

	if len(*command) == 0 {
		log.Fatal("Expected command (-c) (exists|get|ls|lsr|create|creater|set|delete|getacl|setacl)")
	}

	if len(flag.Args()) < 1 {
		log.Fatal("Expected path argument")
	}
	path := flag.Arg(0)
	if *command == "ls" {
	} else if strings.HasSuffix(path, "/") {
		log.Fatal("Path must not end with '/'")
	}

	rand.Seed(time.Now().UnixNano())
	zk.SetServers(serversArray)

	if *authUser != "" && *authPwd != "" {
		authExp := fmt.Sprint(*authUser, ":", *authPwd)
		zk.SetAuth("digest", []byte(authExp))
	}

	if *command == "creater" {
		*command = "create"
		*force = true
	}
	switch *command {
	case "exists":
		{
			if exists, err := zk.Exists(path); err == nil && exists {
				output.PrintString([]byte("true"), *format)
			} else {
				log.Fatale(err)
			}
		}
	case "get":
		{
			if result, err := zk.Get(path); err == nil {
				output.PrintString(result, *format)
			} else {
				log.Fatale(err)
			}
		}
	case "getacl":
		{
			if result, err := zk.GetACL(path); err == nil {
				output.PrintStringArray(result, *format)
			} else {
				log.Fatale(err)
			}
		}
	case "ls":
		{
			if result, err := zk.Children(path); err == nil {
				output.PrintStringArray(result, *format)
			} else {
				log.Fatale(err)
			}
		}
	case "lsr":
		{
			if result, err := zk.ChildrenRecursive(path); err == nil {
				output.PrintStringArray(result, *format)
			} else {
				log.Fatale(err)
			}
		}
	case "create":
		{
			if len(flag.Args()) < 2 {
				log.Fatal("Expected data argument")
			}
			if *authUser != "" && *authPwd != "" {
				perms, err := zk.BuildACL("digest", *authUser, *authPwd, *acls)
				if err != nil {
					log.Fatale(err)
				}
				if result, err := zk.CreateWithACL(path, []byte(flag.Arg(1)), *force, perms); err == nil {
					log.Infof("Created %+v", result)
				} else {
					log.Fatale(err)
				}
			} else {
				if result, err := zk.Create(path, []byte(flag.Arg(1)), *force); err == nil {
					log.Infof("Created %+v", result)
				} else {
					log.Fatale(err)
				}
			}
		}
	case "set":
		{
			var info []byte
			if len(flag.Args()) > 1 {
				info = []byte(flag.Arg(1))
			} else {
				var err error
				info, err = ioutil.ReadAll(os.Stdin)
				if err != nil {
					log.Fatale(err)
				}
			}
			if result, err := zk.Set(path, info); err == nil {
				log.Infof("Set %+v", result)
			} else {
				log.Fatale(err)
			}
		}
	case "setacl":
		{
			var aclstr string
			if len(flag.Args()) > 1 {
				aclstr = flag.Arg(1)
			} else {
				var err error
				data, err := ioutil.ReadAll(os.Stdin)
				aclstr = string(data)
				if err != nil {
					log.Fatale(err)
				}
			}
			if result, err := zk.SetACL(path, aclstr); err == nil {
				log.Infof("Set %+v", result)
			} else {
				log.Fatale(err)
			}
		}
	case "delete", "rm":
		{
			if err := zk.Delete(path); err != nil {
				log.Fatale(err)
			}
		}
	case "deleter", "rmr":
		{
			if !(*force) {
				log.Fatal("deleter (recursive) command requires --force for safety measure")
			}
			if err := zk.DeleteRecursive(path); err != nil {
				log.Fatale(err)
			}
		}
	default:
		log.Fatalf("Unknown command: %s", *command)
	}
}
<commit_msg>updated error message on missing command<commit_after>/*
   Copyright 2014 Outbrain Inc.

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
*/

package main

import (
	"flag"
	"fmt"
	"github.com/outbrain/golib/log"
	"github.com/outbrain/zookeepercli/output"
	"github.com/outbrain/zookeepercli/zk"
	"io/ioutil"
	"math/rand"
	"os"
	"strings"
	"time"
)

// main is the application's entry point.
func main() {
	servers := flag.String("servers", "", "srv1[:port1][,srv2[:port2]...]")
	command := flag.String("c", "", "command, required (exists|get|ls|lsr|create|creater|set|delete|rm|deleter|rmr|getacl|setacl)")
	force := flag.Bool("force", false, "force operation")
	format := flag.String("format", "txt", "output format (txt|json)")
	verbose := flag.Bool("verbose", false, "verbose")
	debug := flag.Bool("debug", false, "debug mode (very verbose)")
	stack := flag.Bool("stack", false, "add stack trace upon error")
	authUser := flag.String("auth_usr", "", "optional, digest scheme, user")
	authPwd := flag.String("auth_pwd", "", "optional, digest scheme, pwd")
	acls := flag.String("acls", "31", "optional, csv list [1|,2|,4|,8|,16|,31]")
	flag.Parse()

	log.SetLevel(log.ERROR)
	if *verbose {
		log.SetLevel(log.INFO)
	}
	if *debug {
		log.SetLevel(log.DEBUG)
	}
	if *stack {
		log.SetPrintStackTrace(*stack)
	}

	log.Info("starting")

	if *servers == "" {
		log.Fatal("Expected comma delimited list of servers via --servers")
	}
	serversArray := strings.Split(*servers, ",")
	if len(serversArray) == 0 {
		log.Fatal("Expected comma delimited list of servers via --servers")
	}

	if len(*command) == 0 {
		log.Fatal("Expected command (-c) (exists|get|ls|lsr|create|creater|set|delete|rm|deleter|rmr|getacl|setacl)")
	}

	if len(flag.Args()) < 1 {
		log.Fatal("Expected path argument")
	}
	path := flag.Arg(0)
	if *command == "ls" {
	} else if strings.HasSuffix(path, "/") {
		log.Fatal("Path must not end with '/'")
	}

	rand.Seed(time.Now().UnixNano())
	zk.SetServers(serversArray)

	if *authUser != "" && *authPwd != "" {
		authExp := fmt.Sprint(*authUser, ":", *authPwd)
		zk.SetAuth("digest", []byte(authExp))
	}

	if *command == "creater" {
		*command = "create"
		*force = true
	}
	switch *command {
	case "exists":
		{
			if exists, err := zk.Exists(path); err == nil && exists {
				output.PrintString([]byte("true"), *format)
			} else {
				log.Fatale(err)
			}
		}
	case "get":
		{
			if result, err := zk.Get(path); err == nil {
				output.PrintString(result, *format)
			} else {
				log.Fatale(err)
			}
		}
	case "getacl":
		{
			if result, err := zk.GetACL(path); err == nil {
				output.PrintStringArray(result, *format)
			} else {
				log.Fatale(err)
			}
		}
	case "ls":
		{
			if result, err := zk.Children(path); err == nil {
				output.PrintStringArray(result, *format)
			} else {
				log.Fatale(err)
			}
		}
	case "lsr":
		{
			if result, err := zk.ChildrenRecursive(path); err == nil {
				output.PrintStringArray(result, *format)
			} else {
				log.Fatale(err)
			}
		}
	case "create":
		{
			if len(flag.Args()) < 2 {
				log.Fatal("Expected data argument")
			}
			if *authUser != "" && *authPwd != "" {
				perms, err := zk.BuildACL("digest", *authUser, *authPwd, *acls)
				if err != nil {
					log.Fatale(err)
				}
				if result, err := zk.CreateWithACL(path, []byte(flag.Arg(1)), *force, perms); err == nil {
					log.Infof("Created %+v", result)
				} else {
					log.Fatale(err)
				}
			} else {
				if result, err := zk.Create(path, []byte(flag.Arg(1)), *force); err == nil {
					log.Infof("Created %+v", result)
				} else {
					log.Fatale(err)
				}
			}
		}
	case "set":
		{
			var info []byte
			if len(flag.Args()) > 1 {
				info = []byte(flag.Arg(1))
			} else {
				var err error
				info, err = ioutil.ReadAll(os.Stdin)
				if err != nil {
					log.Fatale(err)
				}
			}
			if result, err := zk.Set(path, info); err == nil {
				log.Infof("Set %+v", result)
			} else {
				log.Fatale(err)
			}
		}
	case "setacl":
		{
			var aclstr string
			if len(flag.Args()) > 1 {
				aclstr = flag.Arg(1)
			} else {
				var err error
				data, err := ioutil.ReadAll(os.Stdin)
				aclstr = string(data)
				if err != nil {
					log.Fatale(err)
				}
			}
			if result, err := zk.SetACL(path, aclstr); err == nil {
				log.Infof("Set %+v", result)
			} else {
				log.Fatale(err)
			}
		}
	case "delete", "rm":
		{
			if err := zk.Delete(path); err != nil {
				log.Fatale(err)
			}
		}
	case "deleter", "rmr":
		{
			if !(*force) {
				log.Fatal("deleter (recursive) command requires --force for safety measure")
			}
			if err := zk.DeleteRecursive(path); err != nil {
				log.Fatale(err)
			}
		}
	default:
		log.Fatalf("Unknown command: %s", *command)
	}
}
<|endoftext|>


<commit_before>package commands

import (
	"bytes"
	"crypto/rand"
	"errors"
	"fmt"
	"io"
	"sort"
	"strings"
	"text/tabwriter"

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

	ci "gx/ipfs/QmP1DfoUjiWH2ZBo1PBH6FupdBucbDepx3HpWmEY6JMUpY/go-libp2p-crypto"
	peer "gx/ipfs/QmdS9KpbDyPrieswibZhkod1oXqRwZJrUPzxCofAMWpFGq/go-libp2p-peer"
)

var KeyCmd = &cmds.Command{
	Helptext: cmds.HelpText{
		Tagline: "Create and list IPNS name keypairs",
		ShortDescription: `
'ipfs key gen' generates a new keypair for usage with IPNS and 'ipfs name publish'.

  > ipfs key gen --type=rsa --size=2048 mykey
  > ipfs name publish --key=mykey QmSomeHash

'ipfs key list' lists the available keys.

  > ipfs key list
  self
  mykey
		`,
	},
	Subcommands: map[string]*cmds.Command{
		"gen":  keyGenCmd,
		"list": keyListCmd,
		"rm":   keyRmCmd,
	},
}

type KeyOutput struct {
	Name string
	Id   string
}

type KeyOutputList struct {
	Keys []KeyOutput
}

var keyGenCmd = &cmds.Command{
	Helptext: cmds.HelpText{
		Tagline: "Create a new keypair",
	},
	Options: []cmds.Option{
		cmds.StringOption("type", "t", "type of the key to create [rsa, ed25519]"),
		cmds.IntOption("size", "s", "size of the key to generate"),
	},
	Arguments: []cmds.Argument{
		cmds.StringArg("name", true, false, "name of key to create"),
	},
	Run: func(req cmds.Request, res cmds.Response) {
		n, err := req.InvocContext().GetNode()
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		typ, f, err := req.Option("type").String()
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		if !f {
			res.SetError(fmt.Errorf("please specify a key type with --type"), cmds.ErrNormal)
			return
		}

		size, sizefound, err := req.Option("size").Int()
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		name := req.Arguments()[0]
		if name == "self" {
			res.SetError(fmt.Errorf("cannot create key with name 'self'"), cmds.ErrNormal)
			return
		}

		var sk ci.PrivKey
		var pk ci.PubKey

		switch typ {
		case "rsa":
			if !sizefound {
				res.SetError(fmt.Errorf("please specify a key size with --size"), cmds.ErrNormal)
				return
			}

			priv, pub, err := ci.GenerateKeyPairWithReader(ci.RSA, size, rand.Reader)
			if err != nil {
				res.SetError(err, cmds.ErrNormal)
				return
			}

			sk = priv
			pk = pub
		case "ed25519":
			priv, pub, err := ci.GenerateEd25519Key(rand.Reader)
			if err != nil {
				res.SetError(err, cmds.ErrNormal)
				return
			}

			sk = priv
			pk = pub
		default:
			res.SetError(fmt.Errorf("unrecognized key type: %s", typ), cmds.ErrNormal)
			return
		}

		err = n.Repo.Keystore().Put(name, sk)
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		pid, err := peer.IDFromPublicKey(pk)
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		res.SetOutput(&KeyOutput{
			Name: name,
			Id:   pid.Pretty(),
		})
	},
	Marshalers: cmds.MarshalerMap{
		cmds.Text: func(res cmds.Response) (io.Reader, error) {
			k, ok := res.Output().(*KeyOutput)
			if !ok {
				return nil, fmt.Errorf("expected a KeyOutput as command result")
			}

			return strings.NewReader(k.Id + "\n"), nil
		},
	},
	Type: KeyOutput{},
}

var keyListCmd = &cmds.Command{
	Helptext: cmds.HelpText{
		Tagline: "List all local keypairs",
	},
	Options: []cmds.Option{
		cmds.BoolOption("l", "Show extra information about keys."),
	},
	Run: func(req cmds.Request, res cmds.Response) {
		n, err := req.InvocContext().GetNode()
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		keys, err := n.Repo.Keystore().List()
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		sort.Strings(keys)

		list := make([]KeyOutput, 0, len(keys)+1)

		list = append(list, KeyOutput{Name: "self", Id: n.Identity.Pretty()})

		for _, key := range keys {
			privKey, err := n.Repo.Keystore().Get(key)
			if err != nil {
				res.SetError(err, cmds.ErrNormal)
				return
			}

			pubKey := privKey.GetPublic()

			pid, err := peer.IDFromPublicKey(pubKey)
			if err != nil {
				res.SetError(err, cmds.ErrNormal)
				return
			}

			list = append(list, KeyOutput{Name: key, Id: pid.Pretty()})
		}

		res.SetOutput(&KeyOutputList{list})
	},
	Marshalers: cmds.MarshalerMap{
		cmds.Text: keyOutputListMarshaler,
	},
	Type: KeyOutputList{},
}

var keyRmCmd = &cmds.Command{
	Helptext: cmds.HelpText{
		Tagline: "Remove a keypair",
	},
	Arguments: []cmds.Argument{
		cmds.StringArg("name", true, false, "name of key to remove"),
	},
	Run: func(req cmds.Request, res cmds.Response) {
		n, err := req.InvocContext().GetNode()
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		name := req.Arguments()[0]
		if name == "self" {
			res.SetError(fmt.Errorf("cannot remove key with name 'self'"), cmds.ErrNormal)
			return
		}

		removed, err := n.Repo.Keystore().Get(name)
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		err = n.Repo.Keystore().Delete(name)
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		pubKey := removed.GetPublic()

		pid, err := peer.IDFromPublicKey(pubKey)
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		res.SetOutput(&KeyOutput{
			Name: name,
			Id:   pid.Pretty(),
		})
	},
	Marshalers: cmds.MarshalerMap{
		cmds.Text: func(res cmds.Response) (io.Reader, error) {
			v := res.Output().(*KeyOutput)
			s := fmt.Sprintf("Removed key %s with Id: %s\n", v.Name, v.Id)
			return strings.NewReader(s), nil
		},
	},
	Type: KeyOutput{},
}

func keyOutputListMarshaler(res cmds.Response) (io.Reader, error) {
	withId, _, _ := res.Request().Option("l").Bool()

	list, ok := res.Output().(*KeyOutputList)
	if !ok {
		return nil, errors.New("failed to cast []KeyOutput")
	}

	buf := new(bytes.Buffer)
	w := tabwriter.NewWriter(buf, 1, 2, 1, ' ', 0)
	for _, s := range list.Keys {
		if withId {
			fmt.Fprintf(w, "%s\t%s\t\n", s.Id, s.Name)
		} else {
			fmt.Fprintf(w, "%s\n", s.Name)
		}
	}
	w.Flush()
	return buf, nil
}
<commit_msg>Make ipfs key rm variadic<commit_after>package commands

import (
	"bytes"
	"crypto/rand"
	"errors"
	"fmt"
	"io"
	"sort"
	"strings"
	"text/tabwriter"

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

	ci "gx/ipfs/QmP1DfoUjiWH2ZBo1PBH6FupdBucbDepx3HpWmEY6JMUpY/go-libp2p-crypto"
	peer "gx/ipfs/QmdS9KpbDyPrieswibZhkod1oXqRwZJrUPzxCofAMWpFGq/go-libp2p-peer"
)

var KeyCmd = &cmds.Command{
	Helptext: cmds.HelpText{
		Tagline: "Create and list IPNS name keypairs",
		ShortDescription: `
'ipfs key gen' generates a new keypair for usage with IPNS and 'ipfs name publish'.

  > ipfs key gen --type=rsa --size=2048 mykey
  > ipfs name publish --key=mykey QmSomeHash

'ipfs key list' lists the available keys.

  > ipfs key list
  self
  mykey
		`,
	},
	Subcommands: map[string]*cmds.Command{
		"gen":  keyGenCmd,
		"list": keyListCmd,
		"rm":   keyRmCmd,
	},
}

type KeyOutput struct {
	Name string
	Id   string
}

type KeyOutputList struct {
	Keys []KeyOutput
}

var keyGenCmd = &cmds.Command{
	Helptext: cmds.HelpText{
		Tagline: "Create a new keypair",
	},
	Options: []cmds.Option{
		cmds.StringOption("type", "t", "type of the key to create [rsa, ed25519]"),
		cmds.IntOption("size", "s", "size of the key to generate"),
	},
	Arguments: []cmds.Argument{
		cmds.StringArg("name", true, false, "name of key to create"),
	},
	Run: func(req cmds.Request, res cmds.Response) {
		n, err := req.InvocContext().GetNode()
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		typ, f, err := req.Option("type").String()
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		if !f {
			res.SetError(fmt.Errorf("please specify a key type with --type"), cmds.ErrNormal)
			return
		}

		size, sizefound, err := req.Option("size").Int()
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		name := req.Arguments()[0]
		if name == "self" {
			res.SetError(fmt.Errorf("cannot create key with name 'self'"), cmds.ErrNormal)
			return
		}

		var sk ci.PrivKey
		var pk ci.PubKey

		switch typ {
		case "rsa":
			if !sizefound {
				res.SetError(fmt.Errorf("please specify a key size with --size"), cmds.ErrNormal)
				return
			}

			priv, pub, err := ci.GenerateKeyPairWithReader(ci.RSA, size, rand.Reader)
			if err != nil {
				res.SetError(err, cmds.ErrNormal)
				return
			}

			sk = priv
			pk = pub
		case "ed25519":
			priv, pub, err := ci.GenerateEd25519Key(rand.Reader)
			if err != nil {
				res.SetError(err, cmds.ErrNormal)
				return
			}

			sk = priv
			pk = pub
		default:
			res.SetError(fmt.Errorf("unrecognized key type: %s", typ), cmds.ErrNormal)
			return
		}

		err = n.Repo.Keystore().Put(name, sk)
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		pid, err := peer.IDFromPublicKey(pk)
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		res.SetOutput(&KeyOutput{
			Name: name,
			Id:   pid.Pretty(),
		})
	},
	Marshalers: cmds.MarshalerMap{
		cmds.Text: func(res cmds.Response) (io.Reader, error) {
			k, ok := res.Output().(*KeyOutput)
			if !ok {
				return nil, fmt.Errorf("expected a KeyOutput as command result")
			}

			return strings.NewReader(k.Id + "\n"), nil
		},
	},
	Type: KeyOutput{},
}

var keyListCmd = &cmds.Command{
	Helptext: cmds.HelpText{
		Tagline: "List all local keypairs",
	},
	Options: []cmds.Option{
		cmds.BoolOption("l", "Show extra information about keys."),
	},
	Run: func(req cmds.Request, res cmds.Response) {
		n, err := req.InvocContext().GetNode()
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		keys, err := n.Repo.Keystore().List()
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		sort.Strings(keys)

		list := make([]KeyOutput, 0, len(keys)+1)

		list = append(list, KeyOutput{Name: "self", Id: n.Identity.Pretty()})

		for _, key := range keys {
			privKey, err := n.Repo.Keystore().Get(key)
			if err != nil {
				res.SetError(err, cmds.ErrNormal)
				return
			}

			pubKey := privKey.GetPublic()

			pid, err := peer.IDFromPublicKey(pubKey)
			if err != nil {
				res.SetError(err, cmds.ErrNormal)
				return
			}

			list = append(list, KeyOutput{Name: key, Id: pid.Pretty()})
		}

		res.SetOutput(&KeyOutputList{list})
	},
	Marshalers: cmds.MarshalerMap{
		cmds.Text: keyOutputListMarshaler,
	},
	Type: KeyOutputList{},
}

var keyRmCmd = &cmds.Command{
	Helptext: cmds.HelpText{
		Tagline: "Remove a keypair",
	},
	Arguments: []cmds.Argument{
		cmds.StringArg("name", true, true, "names of keys to remove"),
	},
	Options: []cmds.Option{
		cmds.BoolOption("l", "Show extra information about keys."),
	},
	Run: func(req cmds.Request, res cmds.Response) {
		n, err := req.InvocContext().GetNode()
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		names := req.Arguments()

		list := make([]KeyOutput, 0, len(names))
		for _, name := range names {
			if name == "self" {
				res.SetError(fmt.Errorf("cannot remove key with name 'self'"), cmds.ErrNormal)
				return
			}

			removed, err := n.Repo.Keystore().Get(name)
			if err != nil {
				res.SetError(fmt.Errorf("no key named %s was found", name), cmds.ErrNormal)
				return
			}

			pubKey := removed.GetPublic()

			pid, err := peer.IDFromPublicKey(pubKey)
			if err != nil {
				res.SetError(err, cmds.ErrNormal)
				return
			}

			list = append(list, KeyOutput{Name: name, Id: pid.Pretty()})
		}

		for _, name := range names {
			err = n.Repo.Keystore().Delete(name)
			if err != nil {
				res.SetError(err, cmds.ErrNormal)
				return
			}
		}

		res.SetOutput(&KeyOutputList{list})
	},
	Marshalers: cmds.MarshalerMap{
		cmds.Text: keyOutputListMarshaler,
	},
	Type: KeyOutputList{},
}

func keyOutputListMarshaler(res cmds.Response) (io.Reader, error) {
	withId, _, _ := res.Request().Option("l").Bool()

	list, ok := res.Output().(*KeyOutputList)
	if !ok {
		return nil, errors.New("failed to cast []KeyOutput")
	}

	buf := new(bytes.Buffer)
	w := tabwriter.NewWriter(buf, 1, 2, 1, ' ', 0)
	for _, s := range list.Keys {
		if withId {
			fmt.Fprintf(w, "%s\t%s\t\n", s.Id, s.Name)
		} else {
			fmt.Fprintf(w, "%s\n", s.Name)
		}
	}
	w.Flush()
	return buf, nil
}
<|endoftext|>