<commit_before>// Copyright © 2017 G. Hussain Chinoy <ghchinoy@gmail.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 (
	"bufio"
	"encoding/csv"
	"fmt"
	"log"
	"os"
	"sort"
	"strings"

	"github.com/olekukonko/tablewriter"
	toml "github.com/pelletier/go-toml"
	"github.com/spf13/afero"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
)

// profilesCmd represents the profile command
var profilesCmd = &cobra.Command{
	Use:   "profiles",
	Short: "Manage profiles",
	Long:  `Add, remove, list profiles to manage Cloud Elements access`,
}

// addProfileCmd represents the addProfile command
var addProfileCmd = &cobra.Command{
	Use:   "add <profile>",
	Short: "add a new profile",
	Long:  `Adds a new profile to the available profiles. Provide a name to get started.`,
	Run: func(cmd *cobra.Command, args []string) {
		// check for args, if arg, then list the details for that particular profile
		if len(args) > 0 {
			profile = args[0]
		} else {
			fmt.Println("please provide a profile name to add, profile add <name>")
			os.Exit(1)
		}
		fmt.Printf("%7s: %s\n", "profile", profile)
		if viper.IsSet(profile) {
			fmt.Printf("Profile %s exists.\n", profile)
			os.Exit(1)
		}
		reader := bufio.NewReader(os.Stdin)
		fmt.Print("base URI: ")
		base, _ := reader.ReadString('\n')
		base = strings.Replace(base, "\n", "", -1)
		fmt.Print("user token: ")
		user, _ := reader.ReadString('\n')
		user = strings.Replace(user, "\n", "", -1)
		fmt.Print("org token: ")
		org, _ := reader.ReadString('\n')
		org = strings.Replace(org, "\n", "", -1)

		viper.Set(profile+".base", base)
		viper.Set(profile+".org", org)
		viper.Set(profile+".user", user)

		err := writeConfigAs(viper.ConfigFileUsed(), true)
		if err != nil {
			fmt.Println("Unable to write config file", err.Error())
			fmt.Printf("Config file %s unchanged.\n", viper.ConfigFileUsed())
		}
		fmt.Printf("Added profile %s\n", profile)
	},
}

var longProfile bool

// listProfilesCmd represents the listProfiles command
var listProfilesCmd = &cobra.Command{
	Use:   "list",
	Short: "lists available profiles",
	Long:  `Lists available profiles`,
	Run: func(cmd *cobra.Command, args []string) {
		// check for args, if arg, then list the details for that particular profile
		if len(args) > 0 {
			profile = args[0]
		}
		fmt.Printf("%7s: %s", "profile", profile)
		if viper.IsSet(profile) {
			p := viper.GetStringMap(profile)

			if val, ok := p["label"]; ok {
				fmt.Printf(" (%s)\n", val)
			} else {
				fmt.Println()
			}

			for k, v := range p {
				if k == "base" {
					fmt.Printf("%7s: %s\n", k, v)
				}
			}
		} else {
			fmt.Printf("No %s profile exists in config file %s.", profile, cfgFile)
		}
		fmt.Println()
		settings := viper.AllSettings()
		var profiles []string
		for k := range settings {
			profiles = append(profiles, k)
		}
		sort.Strings(profiles)
		if longProfile {
			data := [][]string{}
			fmt.Printf("%v profiles\n", len(profiles))
			for _, k := range profiles {
				if k == "profile" {
					break
				}
				v := settings[k].(map[string]interface{})
				if v["base"] == nil {
					break
				}
				data = append(data, []string{
					k,
					fmt.Sprintf("%s", v["base"]),
				})
			}
			if outputCSV {
				w := csv.NewWriter(os.Stdout)
				for _, record := range data {
					if err := w.Write(record); err != nil {
						log.Fatalln("error writing record to csv:", err)
					}
				}
				w.Flush()
				if err := w.Error(); err != nil {
					log.Fatal(err)
				}
			} else {
				table := tablewriter.NewWriter(os.Stdout)
				table.SetHeader([]string{"profile", "base url"})
				table.SetBorder(false)
				table.SetAlignment(tablewriter.ALIGN_LEFT)
				table.AppendBulk(data)
				table.Render()
			}
			os.Exit(0)
		}
		posn := sort.SearchStrings(profiles, "profile")
		profiles = append(profiles[:posn], profiles[posn+1:]...)
		fmt.Println("Valid profiles:", strings.Join(profiles, ", "))
	},
}

// setProfileCmd represents the setProfile command
var setProfileCmd = &cobra.Command{
	Use:   "set <profile>",
	Short: "sets a profile to be the default profile",
	Long:  `Sets given profile name as the default profile`,
	Run: func(cmd *cobra.Command, args []string) {

		// check for args, if arg, then list the details for that particular profile
		if len(args) > 0 {
			profile = args[0]
		}
		fmt.Printf("%7s: %s\n", "profile", profile)
		// check if specified profile exists
		if viper.IsSet(profile) {
			p := viper.GetStringMap(profile)

			// copy profile data to default, update label
			// indicate change has occurred

			viper.Set("default.base", p["base"])
			viper.Set("default.org", p["org"])
			viper.Set("default.user", p["user"])
			viper.Set("default.label", profile)

			// writing back has a PR to make this more formal: https://github.com/spf13/viper/pull/287
			// Once that PR is merged, replace writeConfigFile with
			/*
				if err := viper.WriteConfigAs(viper.ConfigFileUsed()); err != nil {
					fmt.Fatal(err)
				}
			*/
			err := writeConfigAs(viper.ConfigFileUsed(), true)
			if err != nil {
				fmt.Println("Unable to write config file", err.Error())
				fmt.Printf("Config file %s unchanged.\n", viper.ConfigFileUsed())
			}
			fmt.Printf("Default profile set to %s\n", profile)

			for k, v := range p {
				if k == "base" {
					fmt.Printf("%7s: %s\n", k, v)
				}
				if k == "label" {
					fmt.Printf("%7s: %s\n", k, v)
				}
			}
		} else { // if not, end
			fmt.Printf("No %s profile exists in config file %s.\n", profile, cfgFile)
			fmt.Printf("Cannot set %s as default profile.\n", profile)
			fmt.Println()
			settings := viper.AllSettings()
			var profiles []string
			for k := range settings {
				profiles = append(profiles, k)
			}
			sort.Strings(profiles)
			posn := sort.SearchStrings(profiles, "profile")
			profiles = append(profiles[:posn], profiles[posn+1:]...)
			fmt.Println("Valid profiles:", strings.Join(profiles, ", "))
		}
	},
}

func writeConfigAs(filename string, force bool) error {

	t, err := toml.TreeFromMap(viper.AllSettings())
	if err != nil {
		return err
	}
	s := t.String()

	var flags int
	if force == true {
		flags = os.O_CREATE | os.O_TRUNC | os.O_WRONLY
	} else {
		if _, err := os.Stat(filename); os.IsNotExist(err) {
			flags = os.O_WRONLY
		} else {
			return fmt.Errorf("file: %s exists - cannot overwrite, use force option", filename)
		}
	}

	var AppFs = afero.NewOsFs()
	f, err := AppFs.OpenFile(filename, flags, os.FileMode(0644))
	if err != nil {
		return err
	}

	_, err = f.WriteString(s)

	return nil
}

func init() {
	RootCmd.AddCommand(profilesCmd)

	profilesCmd.PersistentFlags().StringVar(&cfgFile, "config", "", cfgHelp)
	profilesCmd.PersistentFlags().StringVar(&profile, "profile", "default", "profile name")
	// Set bash-completion
	validConfigFilenames := []string{"toml", ""}
	profilesCmd.PersistentFlags().SetAnnotation("config", cobra.BashCompFilenameExt, validConfigFilenames)

	profilesCmd.AddCommand(listProfilesCmd)
	listProfilesCmd.PersistentFlags().BoolVarP(&longProfile, "long", "l", false, "show long profile")
	listProfilesCmd.PersistentFlags().BoolVarP(&outputCSV, "csv", "", false, "output as CSV")
	profilesCmd.AddCommand(addProfileCmd)
	profilesCmd.AddCommand(setProfileCmd)

}
<commit_msg>Corrects csv flag outputting only profile data<commit_after>// Copyright © 2017 G. Hussain Chinoy <ghchinoy@gmail.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 (
	"bufio"
	"encoding/csv"
	"fmt"
	"log"
	"os"
	"sort"
	"strings"

	"github.com/olekukonko/tablewriter"
	toml "github.com/pelletier/go-toml"
	"github.com/spf13/afero"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
)

// profilesCmd represents the profile command
var profilesCmd = &cobra.Command{
	Use:   "profiles",
	Short: "Manage profiles",
	Long:  `Add, remove, list profiles to manage Cloud Elements access`,
}

// addProfileCmd represents the addProfile command
var addProfileCmd = &cobra.Command{
	Use:   "add <profile>",
	Short: "add a new profile",
	Long:  `Adds a new profile to the available profiles. Provide a name to get started.`,
	Run: func(cmd *cobra.Command, args []string) {
		// check for args, if arg, then list the details for that particular profile
		if len(args) > 0 {
			profile = args[0]
		} else {
			fmt.Println("please provide a profile name to add, profile add <name>")
			os.Exit(1)
		}
		fmt.Printf("%7s: %s\n", "profile", profile)
		if viper.IsSet(profile) {
			fmt.Printf("Profile %s exists.\n", profile)
			os.Exit(1)
		}
		reader := bufio.NewReader(os.Stdin)
		fmt.Print("base URI: ")
		base, _ := reader.ReadString('\n')
		base = strings.Replace(base, "\n", "", -1)
		fmt.Print("user token: ")
		user, _ := reader.ReadString('\n')
		user = strings.Replace(user, "\n", "", -1)
		fmt.Print("org token: ")
		org, _ := reader.ReadString('\n')
		org = strings.Replace(org, "\n", "", -1)

		viper.Set(profile+".base", base)
		viper.Set(profile+".org", org)
		viper.Set(profile+".user", user)

		err := writeConfigAs(viper.ConfigFileUsed(), true)
		if err != nil {
			fmt.Println("Unable to write config file", err.Error())
			fmt.Printf("Config file %s unchanged.\n", viper.ConfigFileUsed())
		}
		fmt.Printf("Added profile %s\n", profile)
	},
}

var longProfile bool

// listProfilesCmd represents the listProfiles command
var listProfilesCmd = &cobra.Command{
	Use:   "list",
	Short: "lists available profiles",
	Long:  `Lists available profiles`,
	Run: func(cmd *cobra.Command, args []string) {
		// check for args, if arg, then list the details for that particular profile
		if len(args) > 0 {
			profile = args[0]
		}

		if !(outputCSV) {
			fmt.Printf("%7s: %s", "profile", profile)
			if viper.IsSet(profile) {
				p := viper.GetStringMap(profile)

				if val, ok := p["label"]; ok {
					fmt.Printf(" (%s)\n", val)
				} else {
					fmt.Println()
				}

				for k, v := range p {
					if k == "base" {
						fmt.Printf("%7s: %s\n", k, v)
					}
				}
			} else {
				fmt.Printf("No %s profile exists in config file %s.", profile, cfgFile)
			}
			fmt.Println()
		}

		settings := viper.AllSettings()
		var profiles []string
		for k := range settings {
			profiles = append(profiles, k)
		}
		sort.Strings(profiles)
		if longProfile {
			data := [][]string{}
			if !(outputCSV) {
				fmt.Printf("%v profiles\n", len(profiles))
			}
			for _, k := range profiles {
				if k == "profile" {
					break
				}
				v := settings[k].(map[string]interface{})
				if v["base"] == nil {
					break
				}
				data = append(data, []string{
					k,
					fmt.Sprintf("%s", v["base"]),
				})
			}
			if outputCSV {
				w := csv.NewWriter(os.Stdout)
				for _, record := range data {
					if err := w.Write(record); err != nil {
						log.Fatalln("error writing record to csv:", err)
					}
				}
				w.Flush()
				if err := w.Error(); err != nil {
					log.Fatal(err)
				}
			} else {
				table := tablewriter.NewWriter(os.Stdout)
				table.SetHeader([]string{"profile", "base url"})
				table.SetBorder(false)
				table.SetAlignment(tablewriter.ALIGN_LEFT)
				table.AppendBulk(data)
				table.Render()
			}
			os.Exit(0)
		}
		posn := sort.SearchStrings(profiles, "profile")
		profiles = append(profiles[:posn], profiles[posn+1:]...)
		fmt.Println("Valid profiles:", strings.Join(profiles, ", "))
	},
}

// setProfileCmd represents the setProfile command
var setProfileCmd = &cobra.Command{
	Use:   "set <profile>",
	Short: "sets a profile to be the default profile",
	Long:  `Sets given profile name as the default profile`,
	Run: func(cmd *cobra.Command, args []string) {

		// check for args, if arg, then list the details for that particular profile
		if len(args) > 0 {
			profile = args[0]
		}
		fmt.Printf("%7s: %s\n", "profile", profile)
		// check if specified profile exists
		if viper.IsSet(profile) {
			p := viper.GetStringMap(profile)

			// copy profile data to default, update label
			// indicate change has occurred

			viper.Set("default.base", p["base"])
			viper.Set("default.org", p["org"])
			viper.Set("default.user", p["user"])
			viper.Set("default.label", profile)

			// writing back has a PR to make this more formal: https://github.com/spf13/viper/pull/287
			// Once that PR is merged, replace writeConfigFile with
			/*
				if err := viper.WriteConfigAs(viper.ConfigFileUsed()); err != nil {
					fmt.Fatal(err)
				}
			*/
			err := writeConfigAs(viper.ConfigFileUsed(), true)
			if err != nil {
				fmt.Println("Unable to write config file", err.Error())
				fmt.Printf("Config file %s unchanged.\n", viper.ConfigFileUsed())
			}
			fmt.Printf("Default profile set to %s\n", profile)

			for k, v := range p {
				if k == "base" {
					fmt.Printf("%7s: %s\n", k, v)
				}
				if k == "label" {
					fmt.Printf("%7s: %s\n", k, v)
				}
			}
		} else { // if not, end
			fmt.Printf("No %s profile exists in config file %s.\n", profile, cfgFile)
			fmt.Printf("Cannot set %s as default profile.\n", profile)
			fmt.Println()
			settings := viper.AllSettings()
			var profiles []string
			for k := range settings {
				profiles = append(profiles, k)
			}
			sort.Strings(profiles)
			posn := sort.SearchStrings(profiles, "profile")
			profiles = append(profiles[:posn], profiles[posn+1:]...)
			fmt.Println("Valid profiles:", strings.Join(profiles, ", "))
		}
	},
}

func writeConfigAs(filename string, force bool) error {

	t, err := toml.TreeFromMap(viper.AllSettings())
	if err != nil {
		return err
	}
	s := t.String()

	var flags int
	if force == true {
		flags = os.O_CREATE | os.O_TRUNC | os.O_WRONLY
	} else {
		if _, err := os.Stat(filename); os.IsNotExist(err) {
			flags = os.O_WRONLY
		} else {
			return fmt.Errorf("file: %s exists - cannot overwrite, use force option", filename)
		}
	}

	var AppFs = afero.NewOsFs()
	f, err := AppFs.OpenFile(filename, flags, os.FileMode(0644))
	if err != nil {
		return err
	}

	_, err = f.WriteString(s)

	return nil
}

func init() {
	RootCmd.AddCommand(profilesCmd)

	profilesCmd.PersistentFlags().StringVar(&cfgFile, "config", "", cfgHelp)
	profilesCmd.PersistentFlags().StringVar(&profile, "profile", "default", "profile name")
	// Set bash-completion
	validConfigFilenames := []string{"toml", ""}
	profilesCmd.PersistentFlags().SetAnnotation("config", cobra.BashCompFilenameExt, validConfigFilenames)

	profilesCmd.AddCommand(listProfilesCmd)
	listProfilesCmd.PersistentFlags().BoolVarP(&longProfile, "long", "l", false, "show long profile")
	listProfilesCmd.PersistentFlags().BoolVarP(&outputCSV, "csv", "", false, "output as CSV")
	profilesCmd.AddCommand(addProfileCmd)
	profilesCmd.AddCommand(setProfileCmd)

}
<|endoftext|>


<commit_before>package main

import (
	"io"
	"io/ioutil"
	"log"
	"os"
	"strings"

	"github.com/pressly/stackup/client"
	"github.com/pressly/stackup/config"
	"github.com/pressly/stackup/terminal"

	"github.com/pressly/prefixer"

	gossh "golang.org/x/crypto/ssh"
	"gopkg.in/yaml.v2"
)

// usage prints help for an arg and exits.
func usage(conf *config.Config, arg int) {
	log.Println("Usage: sup <network> <target/command>\n")
	switch arg {
	case 1:
		// <network> missing, print available hosts.
		log.Println("Available networks (from Supfile):")
		for name, network := range conf.Networks {
			log.Printf("- %v\n", name)
			for _, host := range network.Hosts {
				log.Printf("   - %v\n", host)
			}
		}
	case 2:
		// <target/command> not found or missing,
		// print available targets/commands.
		log.Println("Available targets (from Supfile):")
		for name, commands := range conf.Targets {
			log.Printf("- %v", name)
			for _, cmd := range commands {
				log.Printf("\t%v\n", cmd)
			}
		}
		log.Println()
		log.Println("Available commands (from Supfile):")
		for name, cmd := range conf.Commands {
			log.Printf("- %v\t%v", name, cmd.Desc)
		}
	}
	os.Exit(1)
}

// parseArgs parses os.Args for network and commands to be run.
func parseArgs(conf *config.Config) (config.Network, []config.Command) {
	var commands []config.Command

	// Check for the first argument first
	if len(os.Args) < 2 {
		usage(conf, len(os.Args))
	}
	// Does the <network> exist?
	network, ok := conf.Networks[os.Args[1]]
	if !ok {
		log.Printf("Unknown network \"%v\"\n\n", os.Args[1])
		usage(conf, 1)
	}

	// Does <network> have any hosts?
	if len(network.Hosts) == 0 {
		log.Printf("No hosts specified for network \"%v\"", os.Args[1])
		usage(conf, 1)
	}

	// Check for the second argument
	if len(os.Args) < 3 {
		usage(conf, len(os.Args))
	}
	// Does the <target/command> exist?
	target, isTarget := conf.Targets[os.Args[2]]
	if isTarget {
		// It's the target. Loop over its commands.
		for _, cmd := range target {
			// Does the target's command exist?
			command, isCommand := conf.Commands[cmd]
			if !isCommand {
				log.Printf("Unknown command \"%v\" (from target \"%v\": %v)\n\n", cmd, os.Args[2], target)
				usage(conf, 2)
			}
			command.Name = cmd
			commands = append(commands, command)
		}
	} else {
		// It's probably a command. Does it exist?
		command, isCommand := conf.Commands[os.Args[2]]
		if !isCommand {
			// Not a target, nor command.
			log.Printf("Unknown target/command \"%v\"\n\n", os.Args[2])
			usage(conf, 2)
		}
		command.Name = os.Args[2]
		commands = append(commands, command)
	}

	// Check for extra arguments
	if len(os.Args) != 3 {
		usage(conf, len(os.Args))
	}

	return network, commands
}

func main() {
	var (
		conf       config.Config
		paddingLen int
	)

	// Read configuration file.
	data, _ := ioutil.ReadFile("./Supfile")
	if err := yaml.Unmarshal(data, &conf); err != nil {
		log.Fatal(err)
	}

	// Parse network and commands to be run from os.Args.
	network, commands := parseArgs(&conf)

	// Process all ENVs into a string of form
	// `export FOO="bar"; export BAR="baz";`.
	env := ``
	for name, value := range conf.Env {
		env += `export ` + name + `="` + value + `";`
	}
	for name, value := range network.Env {
		env += `export ` + name + `="` + value + `";`
	}

	// Create clients for every host (either SSH or Localhost).
	var clients []client.Client
	for _, host := range network.Hosts {
		var c client.Client

		if host == "localhost" { // LocalhostClient

			localhostClient := &client.LocalhostClient{
				Env: env,
			}
			if err := localhostClient.Connect(host); err != nil {
				log.Fatal(err)
			}

			c = localhostClient

		} else { // SSHClient

			sshClient := &client.SSHClient{
				Env: env,
			}
			if err := sshClient.Connect(host); err != nil {
				log.Fatal(err)
			}
			defer sshClient.Close()

			c = sshClient
		}

		len := len(c.Prefix())
		if len > paddingLen {
			paddingLen = len
		}

		clients = append(clients, c)
	}

	// Run command or run multiple commands defined by target sequentally.
	for _, cmd := range commands {
		// Translate command into task(s).
		tasks, err := client.TasksFromConfigCommand(cmd)
		if err != nil {
			log.Fatalf("TasksFromConfigCommand(): ", err)
		}

		// Run tasks sequentally.
		for _, task := range tasks {

			// Run task in parallel.
			for i, c := range clients {
				padding := strings.Repeat(" ", paddingLen-(len(c.Prefix())))
				color := terminal.Colors[i%len(terminal.Colors)]

				prefix := color + padding + c.Prefix() + " | "
				err := c.Run(task)
				if err != nil {
					log.Fatalf("%sexit %v", prefix, err)
				}

				// Copy over tasks's STDOUT.
				go func(c client.Client) {
					switch t := c.(type) {
					case *client.SSHClient:
						if _, err := io.Copy(os.Stdout, prefixer.New(t.RemoteStdout, prefix)); err != nil {
							log.Printf("%sSTDOUT: %v", t.Prefix(), err)
						}
					case *client.LocalhostClient:
						if _, err := io.Copy(os.Stdout, prefixer.New(t.Stdout, prefix)); err != nil {
							log.Printf("%sSTDOUT: %v", t.Prefix(), err)
						}
					}
				}(c)

				// Copy over tasks's STDERR.
				go func(c client.Client) {
					switch t := c.(type) {
					case *client.SSHClient:
						if _, err := io.Copy(os.Stderr, prefixer.New(t.RemoteStderr, prefix)); err != nil {
							log.Printf("%sSTDERR: %v", t.Prefix(), err)
						}
					case *client.LocalhostClient:
						if _, err := io.Copy(os.Stderr, prefixer.New(t.Stderr, prefix)); err != nil {
							log.Printf("%sSTDERR: %v", t.Prefix(), err)
						}
					}
				}(c)
			}

			// Wait for every client to finish the task.
			for _, c := range clients {
				if err := c.Wait(); err != nil {
					//TODO: Handle the SSH ExitError in ssh pkg
					e, ok := err.(*gossh.ExitError)
					if ok && e.ExitStatus() != 15 {
						log.Fatalf("%sexit %v", c.Prefix(), e.ExitStatus())
					}
					log.Fatalf("%s: %v", c.Prefix(), err)
				}
			}
		}
	}

	//TODO: We should wait for all io.Copy() goroutines.
	//TODO: We should not exit 0, if there was an error.
}
<commit_msg>Error out early, if no Supfile is provided<commit_after>package main

import (
	"io"
	"io/ioutil"
	"log"
	"os"
	"strings"

	"github.com/pressly/stackup/client"
	"github.com/pressly/stackup/config"
	"github.com/pressly/stackup/terminal"

	"github.com/pressly/prefixer"

	gossh "golang.org/x/crypto/ssh"
	"gopkg.in/yaml.v2"
)

// usage prints help for an arg and exits.
func usage(conf *config.Config, arg int) {
	log.Println("Usage: sup <network> <target/command>\n")
	switch arg {
	case 1:
		// <network> missing, print available hosts.
		log.Println("Available networks (from Supfile):")
		for name, network := range conf.Networks {
			log.Printf("- %v\n", name)
			for _, host := range network.Hosts {
				log.Printf("   - %v\n", host)
			}
		}
	case 2:
		// <target/command> not found or missing,
		// print available targets/commands.
		log.Println("Available targets (from Supfile):")
		for name, commands := range conf.Targets {
			log.Printf("- %v", name)
			for _, cmd := range commands {
				log.Printf("\t%v\n", cmd)
			}
		}
		log.Println()
		log.Println("Available commands (from Supfile):")
		for name, cmd := range conf.Commands {
			log.Printf("- %v\t%v", name, cmd.Desc)
		}
	}
	os.Exit(1)
}

// parseArgs parses os.Args for network and commands to be run.
func parseArgsOrDie(conf *config.Config) (config.Network, []config.Command) {
	var commands []config.Command

	// Check for the first argument first
	if len(os.Args) < 2 {
		usage(conf, len(os.Args))
	}
	// Does the <network> exist?
	network, ok := conf.Networks[os.Args[1]]
	if !ok {
		log.Printf("Unknown network \"%v\"\n\n", os.Args[1])
		usage(conf, 1)
	}

	// Does <network> have any hosts?
	if len(network.Hosts) == 0 {
		log.Printf("No hosts specified for network \"%v\"", os.Args[1])
		usage(conf, 1)
	}

	// Check for the second argument
	if len(os.Args) < 3 {
		usage(conf, len(os.Args))
	}
	// Does the <target/command> exist?
	target, isTarget := conf.Targets[os.Args[2]]
	if isTarget {
		// It's the target. Loop over its commands.
		for _, cmd := range target {
			// Does the target's command exist?
			command, isCommand := conf.Commands[cmd]
			if !isCommand {
				log.Printf("Unknown command \"%v\" (from target \"%v\": %v)\n\n", cmd, os.Args[2], target)
				usage(conf, 2)
			}
			command.Name = cmd
			commands = append(commands, command)
		}
	} else {
		// It's probably a command. Does it exist?
		command, isCommand := conf.Commands[os.Args[2]]
		if !isCommand {
			// Not a target, nor command.
			log.Printf("Unknown target/command \"%v\"\n\n", os.Args[2])
			usage(conf, 2)
		}
		command.Name = os.Args[2]
		commands = append(commands, command)
	}

	// Check for extra arguments
	if len(os.Args) != 3 {
		usage(conf, len(os.Args))
	}

	return network, commands
}

func main() {
	var (
		conf       config.Config
		paddingLen int
	)

	// Read configuration file.
	data, err := ioutil.ReadFile("./Supfile")
	if err != nil {
		log.Fatal(err)
	}
	err = yaml.Unmarshal(data, &conf)
	if err != nil {
		log.Fatal(err)
	}

	// Parse network and commands to be run from os.Args.
	network, commands := parseArgsOrDie(&conf)

	// Process all ENVs into a string of form
	// `export FOO="bar"; export BAR="baz";`.
	env := ``
	for name, value := range conf.Env {
		env += `export ` + name + `="` + value + `";`
	}
	for name, value := range network.Env {
		env += `export ` + name + `="` + value + `";`
	}

	// Create clients for every host (either SSH or Localhost).
	var clients []client.Client
	for _, host := range network.Hosts {
		var c client.Client

		if host == "localhost" { // LocalhostClient

			localhostClient := &client.LocalhostClient{
				Env: env,
			}
			if err := localhostClient.Connect(host); err != nil {
				log.Fatal(err)
			}

			c = localhostClient

		} else { // SSHClient

			sshClient := &client.SSHClient{
				Env: env,
			}
			if err := sshClient.Connect(host); err != nil {
				log.Fatal(err)
			}
			defer sshClient.Close()

			c = sshClient
		}

		len := len(c.Prefix())
		if len > paddingLen {
			paddingLen = len
		}

		clients = append(clients, c)
	}

	// Run command or run multiple commands defined by target sequentally.
	for _, cmd := range commands {
		// Translate command into task(s).
		tasks, err := client.TasksFromConfigCommand(cmd)
		if err != nil {
			log.Fatalf("TasksFromConfigCommand(): ", err)
		}

		// Run tasks sequentally.
		for _, task := range tasks {

			// Run task in parallel.
			for i, c := range clients {
				padding := strings.Repeat(" ", paddingLen-(len(c.Prefix())))
				color := terminal.Colors[i%len(terminal.Colors)]

				prefix := color + padding + c.Prefix() + " | "
				err := c.Run(task)
				if err != nil {
					log.Fatalf("%sexit %v", prefix, err)
				}

				// Copy over tasks's STDOUT.
				go func(c client.Client) {
					switch t := c.(type) {
					case *client.SSHClient:
						if _, err := io.Copy(os.Stdout, prefixer.New(t.RemoteStdout, prefix)); err != nil {
							log.Printf("%sSTDOUT: %v", t.Prefix(), err)
						}
					case *client.LocalhostClient:
						if _, err := io.Copy(os.Stdout, prefixer.New(t.Stdout, prefix)); err != nil {
							log.Printf("%sSTDOUT: %v", t.Prefix(), err)
						}
					}
				}(c)

				// Copy over tasks's STDERR.
				go func(c client.Client) {
					switch t := c.(type) {
					case *client.SSHClient:
						if _, err := io.Copy(os.Stderr, prefixer.New(t.RemoteStderr, prefix)); err != nil {
							log.Printf("%sSTDERR: %v", t.Prefix(), err)
						}
					case *client.LocalhostClient:
						if _, err := io.Copy(os.Stderr, prefixer.New(t.Stderr, prefix)); err != nil {
							log.Printf("%sSTDERR: %v", t.Prefix(), err)
						}
					}
				}(c)
			}

			// Wait for every client to finish the task.
			for _, c := range clients {
				if err := c.Wait(); err != nil {
					//TODO: Handle the SSH ExitError in ssh pkg
					e, ok := err.(*gossh.ExitError)
					if ok && e.ExitStatus() != 15 {
						log.Fatalf("%sexit %v", c.Prefix(), e.ExitStatus())
					}
					log.Fatalf("%s: %v", c.Prefix(), err)
				}
			}
		}
	}

	//TODO: We should wait for all io.Copy() goroutines.
	//TODO: We should not exit 0, if there was an error.
}
<|endoftext|>


<commit_before>package cmd

import (
	"fmt"
	"io/ioutil"
	"path/filepath"
	"strings"

	"github.com/dtan4/valec/aws"
	"github.com/dtan4/valec/secret"
	"github.com/fatih/color"
	"github.com/pkg/errors"
	"github.com/spf13/cobra"
)

// validateCmd represents the validate command
var validateCmd = &cobra.Command{
	Use:   "validate SECRETDIR",
	Short: "Validate secrets in local files",
	RunE: func(cmd *cobra.Command, args []string) error {
		if len(args) < 1 {
			return errors.New("Please specify secret directory.")
		}
		dirname := args[0]

		files, err := ioutil.ReadDir(dirname)
		if err != nil {
			return errors.Wrapf(err, "Failed to read directory. dirname=%s", dirname)
		}

		for _, file := range files {
			if strings.HasPrefix(file.Name(), ".") || !yamlExtRegexp.Match([]byte(file.Name())) {
				continue
			}

			filename := filepath.Join(dirname, file.Name())

			if err := validateFile(filename); err != nil {
				return errors.Wrapf(err, "Failed to validate secrets. filename=%s", filename)
			}
		}

		fmt.Println("All secrets are valid.")

		return nil
	},
}

func validateFile(filename string) error {
	fmt.Println(filename)

	secrets, err := secret.LoadFromYAML(filename)
	if err != nil {
		return errors.Wrapf(err, "Failed to load secrets. filename=%s", filename)
	}

	hasError := false
	red := color.New(color.FgRed)

	for _, secret := range secrets {
		if _, err := aws.KMS.DecryptBase64(secret.Key, secret.Value); err != nil {
			red.Printf("  Secret value is invalid. Please try `valec encrypt`. key=%s\n", secret.Key)
			hasError = true
		}
	}

	if hasError {
		return errors.New("Some secrets are invalid.")
	}

	return nil
}

func init() {
	RootCmd.AddCommand(validateCmd)
}
<commit_msg>Validate secret files recursively<commit_after>package cmd

import (
	"fmt"

	"github.com/dtan4/valec/aws"
	"github.com/dtan4/valec/secret"
	"github.com/dtan4/valec/util"
	"github.com/fatih/color"
	"github.com/pkg/errors"
	"github.com/spf13/cobra"
)

// validateCmd represents the validate command
var validateCmd = &cobra.Command{
	Use:   "validate SECRETDIR",
	Short: "Validate secrets in local files",
	RunE: func(cmd *cobra.Command, args []string) error {
		if len(args) < 1 {
			return errors.New("Please specify secret directory.")
		}
		dirname := args[0]

		files, err := util.ListYAMLFiles(dirname)
		if err != nil {
			return errors.Wrapf(err, "Failed to read directory. dirname=%s", dirname)
		}

		for _, file := range files {
			if err := validateFile(file); err != nil {
				return errors.Wrapf(err, "Failed to validate secrets. filename=%s", file)
			}
		}

		fmt.Println("All secrets are valid.")

		return nil
	},
}

func validateFile(filename string) error {
	fmt.Println(filename)

	secrets, err := secret.LoadFromYAML(filename)
	if err != nil {
		return errors.Wrapf(err, "Failed to load secrets. filename=%s", filename)
	}

	hasError := false
	red := color.New(color.FgRed)

	for _, secret := range secrets {
		if _, err := aws.KMS.DecryptBase64(secret.Key, secret.Value); err != nil {
			red.Printf("  Secret value is invalid. Please try `valec encrypt`. key=%s\n", secret.Key)
			hasError = true
		}
	}

	if hasError {
		return errors.New("Some secrets are invalid.")
	}

	return nil
}

func init() {
	RootCmd.AddCommand(validateCmd)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"

	"github.com/manifoldco/go-manifold"
	"github.com/urfave/cli"

	"github.com/manifoldco/manifold-cli/errs"
)

func validateName(cliCtx *cli.Context, option string) (string, error) {
	val := cliCtx.String(option)
	if val == "" {
		return val, nil
	}

	name := manifold.Name(val)
	if err := name.Validate(nil); err != nil {
		return "", errs.NewUsageExitError(cliCtx, cli.NewExitError(
			fmt.Sprintf("You've provided an invalid %s name!", option), -1,
		))
	}

	return val, nil
}

func validateLabel(cliCtx *cli.Context, option string) (string, error) {
	val := cliCtx.String(option)
	if val == "" {
		return val, nil
	}

	label := manifold.Label(val)
	if err := label.Validate(nil); err != nil {
		return "", errs.NewUsageExitError(cliCtx, cli.NewExitError(
			fmt.Sprintf("You've provided an invalid %s!", option), -1,
		))
	}

	return val, nil
}

func requiredLabel(cliCtx *cli.Context, option string) (string, error) {
	val := cliCtx.String(option)
	if val == "" {
		return "", errs.NewUsageExitError(cliCtx, cli.NewExitError(
			fmt.Sprintf("--%s is required", option), -1,
		))
	}

	return validateLabel(cliCtx, option)
}

func optionalArgLabel(cliCtx *cli.Context, idx int, name string) (string, error) {
	args := cliCtx.Args()

	if len(args) < idx+1 {
		return "", nil
	}

	val := args[idx]
	l := manifold.Label(val)
	if err := l.Validate(nil); err != nil {
		return "", errs.NewUsageExitError(cliCtx, cli.NewExitError(
			fmt.Sprintf("You've provided an invalid %s!", name), -1,
		))
	}

	return val, nil
}

func optionalArgName(cliCtx *cli.Context, idx int, name string) (string, error) {
	args := cliCtx.Args()

	if len(args) < idx+1 {
		return "", nil
	}

	val := args[idx]
	n := manifold.Name(val)
	if err := n.Validate(nil); err != nil {
		return "", errs.NewUsageExitError(cliCtx, cli.NewExitError(
			fmt.Sprintf("You've provided an invalid %s name", name), -1,
		))
	}

	return val, nil
}

func maxOptionalArgsLength(cliCtx *cli.Context, size int) error {
	args := cliCtx.Args()

	if len(args) > size {
		return errs.ErrTooManyArgs
	}

	return nil
}
<commit_msg>Add optional typeName for validateName<commit_after>package main

import (
	"fmt"

	"github.com/manifoldco/go-manifold"
	"github.com/urfave/cli"

	"github.com/manifoldco/manifold-cli/errs"
)

func validateName(cliCtx *cli.Context, option string, typeName ...string) (string, error) {
	val := cliCtx.String(option)
	if val == "" {
		return val, nil
	}

	if len(typeName) > 1 {
		panic("Only one typeName can be provided")
	} else if len(typeName) == 0 {
		typeName = make([]string, 1)
		typeName[0] = option
	}

	name := manifold.Name(val)
	if err := name.Validate(nil); err != nil {
		return "", errs.NewUsageExitError(cliCtx, cli.NewExitError(
			fmt.Sprintf("You've provided an invalid %s name!", typeName[0]), -1,
		))
	}

	return val, nil
}

func validateLabel(cliCtx *cli.Context, option string) (string, error) {
	val := cliCtx.String(option)
	if val == "" {
		return val, nil
	}

	label := manifold.Label(val)
	if err := label.Validate(nil); err != nil {
		return "", errs.NewUsageExitError(cliCtx, cli.NewExitError(
			fmt.Sprintf("You've provided an invalid %s!", option), -1,
		))
	}

	return val, nil
}

func requiredLabel(cliCtx *cli.Context, option string) (string, error) {
	val := cliCtx.String(option)
	if val == "" {
		return "", errs.NewUsageExitError(cliCtx, cli.NewExitError(
			fmt.Sprintf("--%s is required", option), -1,
		))
	}

	return validateLabel(cliCtx, option)
}

func optionalArgLabel(cliCtx *cli.Context, idx int, name string) (string, error) {
	args := cliCtx.Args()

	if len(args) < idx+1 {
		return "", nil
	}

	val := args[idx]
	l := manifold.Label(val)
	if err := l.Validate(nil); err != nil {
		return "", errs.NewUsageExitError(cliCtx, cli.NewExitError(
			fmt.Sprintf("You've provided an invalid %s!", name), -1,
		))
	}

	return val, nil
}

func optionalArgName(cliCtx *cli.Context, idx int, name string) (string, error) {
	args := cliCtx.Args()

	if len(args) < idx+1 {
		return "", nil
	}

	val := args[idx]
	n := manifold.Name(val)
	if err := n.Validate(nil); err != nil {
		return "", errs.NewUsageExitError(cliCtx, cli.NewExitError(
			fmt.Sprintf("You've provided an invalid %s name", name), -1,
		))
	}

	return val, nil
}

func maxOptionalArgsLength(cliCtx *cli.Context, size int) error {
	args := cliCtx.Args()

	if len(args) > size {
		return errs.ErrTooManyArgs
	}

	return nil
}
<|endoftext|>


<commit_before>package torrent

import (
	"bytes"
	"io"
	"time"

	"github.com/anacrolix/chansync"
	"github.com/anacrolix/log"
	"github.com/anacrolix/sync"

	pp "github.com/anacrolix/torrent/peer_protocol"
)

func (pc *PeerConn) startWriter() {
	w := &pc.messageWriter
	*w = peerConnMsgWriter{
		fillWriteBuffer: func() {
			pc.locker().Lock()
			defer pc.locker().Unlock()
			if pc.closed.IsSet() {
				return
			}
			pc.fillWriteBuffer()
		},
		closed: &pc.closed,
		logger: pc.logger,
		w:      pc.w,
		keepAlive: func() bool {
			pc.locker().Lock()
			defer pc.locker().Unlock()
			return pc.useful()
		},
		writeBuffer: new(bytes.Buffer),
	}
	go func() {
		defer pc.locker().Unlock()
		defer pc.close()
		defer pc.locker().Lock()
		pc.messageWriter.run(pc.t.cl.config.KeepAliveTimeout)
	}()
}

type peerConnMsgWriter struct {
	// Must not be called with the local mutex held, as it will call back into the write method.
	fillWriteBuffer func()
	closed          *chansync.SetOnce
	logger          log.Logger
	w               io.Writer
	keepAlive       func() bool

	mu        sync.Mutex
	writeCond chansync.BroadcastCond
	// Pointer so we can swap with the "front buffer".
	writeBuffer *bytes.Buffer
}

// Routine that writes to the peer. Some of what to write is buffered by
// activity elsewhere in the Client, and some is determined locally when the
// connection is writable.
func (cn *peerConnMsgWriter) run(keepAliveTimeout time.Duration) {
	var (
		lastWrite      time.Time = time.Now()
		keepAliveTimer *time.Timer
	)
	cn.mu.Lock()
	defer cn.mu.Unlock()
	keepAliveTimer = time.AfterFunc(keepAliveTimeout, func() {
		cn.mu.Lock()
		defer cn.mu.Unlock()
		if time.Since(lastWrite) >= keepAliveTimeout {
			cn.writeCond.Broadcast()
		}
		keepAliveTimer.Reset(keepAliveTimeout)
	})
	defer keepAliveTimer.Stop()
	frontBuf := new(bytes.Buffer)
	for {
		if cn.closed.IsSet() {
			return
		}
		keepAlive := false
		if cn.writeBuffer.Len() == 0 {
			func() {
				cn.mu.Unlock()
				defer cn.mu.Lock()
				cn.fillWriteBuffer()
				keepAlive = cn.keepAlive()
			}()
		}
		if cn.writeBuffer.Len() == 0 && time.Since(lastWrite) >= keepAliveTimeout && keepAlive {
			cn.writeBuffer.Write(pp.Message{Keepalive: true}.MustMarshalBinary())
			torrent.Add("written keepalives", 1)
		}
		if cn.writeBuffer.Len() == 0 {
			writeCond := cn.writeCond.Signaled()
			cn.mu.Unlock()
			select {
			case <-cn.closed.Done():
			case <-writeCond:
			}
			cn.mu.Lock()
			continue
		}
		// Flip the buffers.
		frontBuf, cn.writeBuffer = cn.writeBuffer, frontBuf
		cn.mu.Unlock()
		n, err := cn.w.Write(frontBuf.Bytes())
		cn.mu.Lock()
		if n != 0 {
			lastWrite = time.Now()
			keepAliveTimer.Reset(keepAliveTimeout)
		}
		if err != nil {
			cn.logger.WithDefaultLevel(log.Debug).Printf("error writing: %v", err)
			return
		}
		if n != frontBuf.Len() {
			panic("short write")
		}
		frontBuf.Reset()
	}
}

func (cn *peerConnMsgWriter) write(msg pp.Message) bool {
	cn.mu.Lock()
	defer cn.mu.Unlock()
	cn.writeBuffer.Write(msg.MustMarshalBinary())
	cn.writeCond.Broadcast()
	return !cn.writeBufferFull()
}

func (cn *peerConnMsgWriter) writeBufferFull() bool {
	return cn.writeBuffer.Len() >= writeBufferHighWaterLen
}
<commit_msg>Rewrite peerConnMsgWriter.run<commit_after>package torrent

import (
	"bytes"
	"io"
	"time"

	"github.com/anacrolix/chansync"
	"github.com/anacrolix/log"
	"github.com/anacrolix/sync"

	pp "github.com/anacrolix/torrent/peer_protocol"
)

func (pc *PeerConn) startWriter() {
	w := &pc.messageWriter
	*w = peerConnMsgWriter{
		fillWriteBuffer: func() {
			pc.locker().Lock()
			defer pc.locker().Unlock()
			if pc.closed.IsSet() {
				return
			}
			pc.fillWriteBuffer()
		},
		closed: &pc.closed,
		logger: pc.logger,
		w:      pc.w,
		keepAlive: func() bool {
			pc.locker().Lock()
			defer pc.locker().Unlock()
			return pc.useful()
		},
		writeBuffer: new(bytes.Buffer),
	}
	go func() {
		defer pc.locker().Unlock()
		defer pc.close()
		defer pc.locker().Lock()
		pc.messageWriter.run(pc.t.cl.config.KeepAliveTimeout)
	}()
}

type peerConnMsgWriter struct {
	// Must not be called with the local mutex held, as it will call back into the write method.
	fillWriteBuffer func()
	closed          *chansync.SetOnce
	logger          log.Logger
	w               io.Writer
	keepAlive       func() bool

	mu        sync.Mutex
	writeCond chansync.BroadcastCond
	// Pointer so we can swap with the "front buffer".
	writeBuffer *bytes.Buffer
}

// Routine that writes to the peer. Some of what to write is buffered by
// activity elsewhere in the Client, and some is determined locally when the
// connection is writable.
func (cn *peerConnMsgWriter) run(keepAliveTimeout time.Duration) {
	lastWrite := time.Now()
	frontBuf := new(bytes.Buffer)
	for {
		if cn.closed.IsSet() {
			return
		}
		cn.fillWriteBuffer()
		keepAlive := cn.keepAlive()
		cn.mu.Lock()
		if cn.writeBuffer.Len() == 0 && time.Since(lastWrite) >= keepAliveTimeout && keepAlive {
			cn.writeBuffer.Write(pp.Message{Keepalive: true}.MustMarshalBinary())
			torrent.Add("written keepalives", 1)
		}
		if cn.writeBuffer.Len() == 0 {
			writeCond := cn.writeCond.Signaled()
			cn.mu.Unlock()
			select {
			case <-cn.closed.Done():
			case <-writeCond:
			case <-time.After(time.Until(lastWrite.Add(keepAliveTimeout))):
			}
			continue
		}
		// Flip the buffers.
		frontBuf, cn.writeBuffer = cn.writeBuffer, frontBuf
		cn.mu.Unlock()
		n, err := cn.w.Write(frontBuf.Bytes())
		if n != 0 {
			lastWrite = time.Now()
		}
		if err != nil {
			cn.logger.WithDefaultLevel(log.Debug).Printf("error writing: %v", err)
			return
		}
		if n != frontBuf.Len() {
			panic("short write")
		}
		frontBuf.Reset()
	}
}

func (cn *peerConnMsgWriter) write(msg pp.Message) bool {
	cn.mu.Lock()
	defer cn.mu.Unlock()
	cn.writeBuffer.Write(msg.MustMarshalBinary())
	cn.writeCond.Broadcast()
	return !cn.writeBufferFull()
}

func (cn *peerConnMsgWriter) writeBufferFull() bool {
	return cn.writeBuffer.Len() >= writeBufferHighWaterLen
}
<|endoftext|>


<commit_before>/*
Copyright IBM Corp. 2017 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 channel

import (
	"fmt"
	"strings"

	"github.com/hyperledger/fabric/common/flogging"
	"github.com/hyperledger/fabric/msp"
	"github.com/hyperledger/fabric/peer/common"
	ab "github.com/hyperledger/fabric/protos/orderer"
	pb "github.com/hyperledger/fabric/protos/peer"
	"github.com/spf13/cobra"
	"github.com/spf13/pflag"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
)

const (
	channelFuncName = "channel"
	shortDes        = "Operate a channel: create|fetch|join|list|update."
	longDes         = "Operate a channel: create|fetch|join|list|update."
)

var logger = flogging.MustGetLogger("channelCmd")

type OrdererRequirement bool
type EndorserRequirement bool

const (
	EndorserRequired    EndorserRequirement = true
	EndorserNotRequired EndorserRequirement = false
	OrdererRequired     OrdererRequirement  = true
	OrdererNotRequired  OrdererRequirement  = false
)

var (
	// join related variables.
	genesisBlockPath string

	// create related variables
	chainID          string
	channelTxFile    string
	orderingEndpoint string
	tls              bool
	caFile           string
	timeout          int
)

// Cmd returns the cobra command for Node
func Cmd(cf *ChannelCmdFactory) *cobra.Command {
	AddFlags(channelCmd)

	channelCmd.AddCommand(createCmd(cf))
	channelCmd.AddCommand(fetchCmd(cf))
	channelCmd.AddCommand(joinCmd(cf))
	channelCmd.AddCommand(listCmd(cf))
	channelCmd.AddCommand(updateCmd(cf))

	return channelCmd
}

// AddFlags adds flags for create and join
func AddFlags(cmd *cobra.Command) {
	flags := cmd.PersistentFlags()

	flags.StringVarP(&orderingEndpoint, "orderer", "o", "", "Ordering service endpoint")
	flags.BoolVarP(&tls, "tls", "", false, "Use TLS when communicating with the orderer endpoint")
	flags.StringVarP(&caFile, "cafile", "", "", "Path to file containing PEM-encoded trusted certificate(s) for the ordering endpoint")
}

var flags *pflag.FlagSet

func init() {
	resetFlags()
}

// Explicitly define a method to facilitate tests
func resetFlags() {
	flags = &pflag.FlagSet{}

	flags.StringVarP(&genesisBlockPath, "blockpath", "b", common.UndefinedParamValue, "Path to file containing genesis block")
	flags.StringVarP(&chainID, "channelID", "c", common.UndefinedParamValue, "In case of a newChain command, the channel ID to create.")
	flags.StringVarP(&channelTxFile, "file", "f", "", "Configuration transaction file generated by a tool such as configtxgen for submitting to orderer")
	flags.IntVarP(&timeout, "timeout", "t", 5, "Channel creation timeout")
}

func attachFlags(cmd *cobra.Command, names []string) {
	cmdFlags := cmd.Flags()
	for _, name := range names {
		if flag := flags.Lookup(name); flag != nil {
			cmdFlags.AddFlag(flag)
		} else {
			logger.Fatalf("Could not find flag '%s' to attach to commond '%s'", name, cmd.Name())
		}
	}
}

var channelCmd = &cobra.Command{
	Use:   channelFuncName,
	Short: fmt.Sprint(shortDes),
	Long:  fmt.Sprint(longDes),
}

type BroadcastClientFactory func() (common.BroadcastClient, error)

// ChannelCmdFactory holds the clients used by ChannelCmdFactory
type ChannelCmdFactory struct {
	EndorserClient   pb.EndorserClient
	Signer           msp.SigningIdentity
	BroadcastClient  common.BroadcastClient
	DeliverClient    deliverClientIntf
	BroadcastFactory BroadcastClientFactory
}

// InitCmdFactory init the ChannelCmdFactory with clients to endorser and orderer according to params
func InitCmdFactory(isEndorserRequired EndorserRequirement, isOrdererRequired OrdererRequirement) (*ChannelCmdFactory, error) {
	var err error

	cmdFact := &ChannelCmdFactory{}

	cmdFact.Signer, err = common.GetDefaultSignerFnc()
	if err != nil {
		return nil, fmt.Errorf("Error getting default signer: %s", err)
	}

	cmdFact.BroadcastFactory = func() (common.BroadcastClient, error) {
		return common.GetBroadcastClientFnc(orderingEndpoint, tls, caFile)
	}

	//for join and list, we need the endorser as well
	if isEndorserRequired {
		cmdFact.EndorserClient, err = common.GetEndorserClientFnc()
		if err != nil {
			return nil, fmt.Errorf("Error getting endorser client %s: %s", channelFuncName, err)
		}
	}

	//for create and fetch, we need the orderer as well
	if isOrdererRequired {
		if len(strings.Split(orderingEndpoint, ":")) != 2 {
			return nil, fmt.Errorf("Ordering service endpoint %s is not valid or missing", orderingEndpoint)
		}

		var opts []grpc.DialOption
		// check for TLS
		if tls {
			if caFile != "" {
				creds, err := credentials.NewClientTLSFromFile(caFile, "")
				if err != nil {
					return nil, fmt.Errorf("Error connecting to %s due to %s", orderingEndpoint, err)
				}
				opts = append(opts, grpc.WithTransportCredentials(creds))
			}
		} else {
			opts = append(opts, grpc.WithInsecure())
		}
		conn, err := grpc.Dial(orderingEndpoint, opts...)
		if err != nil {
			return nil, err
		}

		client, err := ab.NewAtomicBroadcastClient(conn).Deliver(context.TODO())
		if err != nil {
			return nil, fmt.Errorf("Error connecting due to  %s", err)
		}

		cmdFact.DeliverClient = newDeliverClient(conn, client, chainID)
	}
	logger.Infof("Endorser and orderer connections initialized")
	return cmdFact, nil
}
<commit_msg>[FAB-5319] peer cli orderer TLS hostname override<commit_after>/*
Copyright IBM Corp. 2017 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 channel

import (
	"fmt"
	"strings"

	"github.com/hyperledger/fabric/common/flogging"
	"github.com/hyperledger/fabric/msp"
	"github.com/hyperledger/fabric/peer/common"
	ab "github.com/hyperledger/fabric/protos/orderer"
	pb "github.com/hyperledger/fabric/protos/peer"
	"github.com/spf13/cobra"
	"github.com/spf13/pflag"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
)

const (
	channelFuncName = "channel"
	shortDes        = "Operate a channel: create|fetch|join|list|update."
	longDes         = "Operate a channel: create|fetch|join|list|update."
)

var logger = flogging.MustGetLogger("channelCmd")

type OrdererRequirement bool
type EndorserRequirement bool

const (
	EndorserRequired    EndorserRequirement = true
	EndorserNotRequired EndorserRequirement = false
	OrdererRequired     OrdererRequirement  = true
	OrdererNotRequired  OrdererRequirement  = false
)

var (
	// join related variables.
	genesisBlockPath string

	// create related variables
	chainID                    string
	channelTxFile              string
	orderingEndpoint           string
	tls                        bool
	caFile                     string
	ordererTLSHostnameOverride string
	timeout                    int
)

// Cmd returns the cobra command for Node
func Cmd(cf *ChannelCmdFactory) *cobra.Command {
	AddFlags(channelCmd)

	channelCmd.AddCommand(createCmd(cf))
	channelCmd.AddCommand(fetchCmd(cf))
	channelCmd.AddCommand(joinCmd(cf))
	channelCmd.AddCommand(listCmd(cf))
	channelCmd.AddCommand(updateCmd(cf))

	return channelCmd
}

// AddFlags adds flags for create and join
func AddFlags(cmd *cobra.Command) {
	flags := cmd.PersistentFlags()

	flags.StringVarP(&orderingEndpoint, "orderer", "o", "", "Ordering service endpoint")
	flags.BoolVarP(&tls, "tls", "", false, "Use TLS when communicating with the orderer endpoint")
	flags.StringVarP(&caFile, "cafile", "", "", "Path to file containing PEM-encoded trusted certificate(s) for the ordering endpoint")
	flags.StringVarP(&ordererTLSHostnameOverride, "ordererTLSHostnameOverride", "", "", "The hostname override to use when validating the TLS connection to the orderer.")
}

var flags *pflag.FlagSet

func init() {
	resetFlags()
}

// Explicitly define a method to facilitate tests
func resetFlags() {
	flags = &pflag.FlagSet{}

	flags.StringVarP(&genesisBlockPath, "blockpath", "b", common.UndefinedParamValue, "Path to file containing genesis block")
	flags.StringVarP(&chainID, "channelID", "c", common.UndefinedParamValue, "In case of a newChain command, the channel ID to create.")
	flags.StringVarP(&channelTxFile, "file", "f", "", "Configuration transaction file generated by a tool such as configtxgen for submitting to orderer")
	flags.IntVarP(&timeout, "timeout", "t", 5, "Channel creation timeout")
}

func attachFlags(cmd *cobra.Command, names []string) {
	cmdFlags := cmd.Flags()
	for _, name := range names {
		if flag := flags.Lookup(name); flag != nil {
			cmdFlags.AddFlag(flag)
		} else {
			logger.Fatalf("Could not find flag '%s' to attach to commond '%s'", name, cmd.Name())
		}
	}
}

var channelCmd = &cobra.Command{
	Use:   channelFuncName,
	Short: fmt.Sprint(shortDes),
	Long:  fmt.Sprint(longDes),
}

type BroadcastClientFactory func() (common.BroadcastClient, error)

// ChannelCmdFactory holds the clients used by ChannelCmdFactory
type ChannelCmdFactory struct {
	EndorserClient   pb.EndorserClient
	Signer           msp.SigningIdentity
	BroadcastClient  common.BroadcastClient
	DeliverClient    deliverClientIntf
	BroadcastFactory BroadcastClientFactory
}

// InitCmdFactory init the ChannelCmdFactory with clients to endorser and orderer according to params
func InitCmdFactory(isEndorserRequired EndorserRequirement, isOrdererRequired OrdererRequirement) (*ChannelCmdFactory, error) {
	var err error

	cmdFact := &ChannelCmdFactory{}

	cmdFact.Signer, err = common.GetDefaultSignerFnc()
	if err != nil {
		return nil, fmt.Errorf("Error getting default signer: %s", err)
	}

	cmdFact.BroadcastFactory = func() (common.BroadcastClient, error) {
		return common.GetBroadcastClientFnc(orderingEndpoint, tls, caFile)
	}

	//for join and list, we need the endorser as well
	if isEndorserRequired {
		cmdFact.EndorserClient, err = common.GetEndorserClientFnc()
		if err != nil {
			return nil, fmt.Errorf("Error getting endorser client %s: %s", channelFuncName, err)
		}
	}

	//for create and fetch, we need the orderer as well
	if isOrdererRequired {
		if len(strings.Split(orderingEndpoint, ":")) != 2 {
			return nil, fmt.Errorf("Ordering service endpoint %s is not valid or missing", orderingEndpoint)
		}

		var opts []grpc.DialOption
		// check for TLS
		if tls {
			if caFile != "" {
				creds, err := credentials.NewClientTLSFromFile(caFile, ordererTLSHostnameOverride)
				if err != nil {
					return nil, fmt.Errorf("Error connecting to %s due to %s", orderingEndpoint, err)
				}
				opts = append(opts, grpc.WithTransportCredentials(creds))
			}
		} else {
			opts = append(opts, grpc.WithInsecure())
		}
		conn, err := grpc.Dial(orderingEndpoint, opts...)
		if err != nil {
			return nil, err
		}

		client, err := ab.NewAtomicBroadcastClient(conn).Deliver(context.TODO())
		if err != nil {
			return nil, fmt.Errorf("Error connecting due to  %s", err)
		}

		cmdFact.DeliverClient = newDeliverClient(conn, client, chainID)
	}
	logger.Infof("Endorser and orderer connections initialized")
	return cmdFact, nil
}
<|endoftext|>


<commit_before>// Package object is an abstract data record which tracks state changes.
// It's meant to make it easier to map key-value records into ORM / RDBMS
// structures.
// The state change tracking can be useful when the values of primary keys
// need to change. (Changing a foreign key on a table with a composite key, for example)
package object

import (
	"fmt"
	"strconv"
)

// Array is our object array container to facilitate a couple of instances
// where slices are needed.
type Array []*Object

// Object struct encapsulates our key-value pairs and a single-item per-key history
// of the previous value stored for a given key.
type Object struct {
	Type          string
	KV            map[string]interface{} `json:"KV"`
	ChangedFields map[string]interface{} `json:"ChangedFields"`
	Children      map[string]Array
	saved         bool
}

// New is an empty constructor
func New(objType string) *Object {
	return &Object{Type: objType, KV: makeEmptyMap(), ChangedFields: makeEmptyMap(), Children: makeEmptyChildrenMap(), saved: false}
}

// MakeArray will construct an array of length 'size'
func MakeArray(size int) Array {
	objAry := make(Array, size)
	return objAry
}

// NewArray is our single-object array constructor
func NewArray(val *Object) Array {
	objAry := make(Array, 1)
	objAry[0] = val
	return objAry
}

// NewEmptyArray is an empty array constructor
func NewEmptyArray() Array {
	objAry := make(Array, 0)
	return objAry
}

func makeEmptyMap() map[string]interface{} {
	return make(map[string]interface{})
}

func makeEmptyChildrenMap() map[string]Array {
	return make(map[string]Array)
}

// Get is our accessor
func (o Object) Get(k string) interface{} {
	return o.KV[k]
}

// GetString is a safe, typed string accessor
func (o Object) GetString(k string) (string, bool) {
	v, ok := o.KV[k].(string)
	return v, ok
}

// GetInt is a safe, typed int64 accessor
func (o Object) GetInt(k string) (int64, bool) {
	v, ok := o.KV[k].(int64)
	return v, ok
}

// GetFloat is a safe, typed int64 accessor
func (o Object) GetFloat(k string) (float64, bool) {
	v, ok := o.KV[k].(float64)
	return v, ok
}

// GetIntAlways is a safe, typed int64 accessor. It will force conversion away
// from float64 and uint64 values.
func (o Object) GetIntAlways(k string) (int64, error) {
	switch v := o.KV[k].(type) {
	case float64:
		fl := o.KV[k].(float64)
		return int64(fl), nil
	case int64:
		fl := o.KV[k].(int64)
		return fl, nil
	case uint64:
		fl := o.KV[k].(uint64)
		return int64(fl), nil
	case string:
		fl := o.KV[k].(string)
		return strconv.ParseInt(fl, 10, 64)
	default:
		return 0, fmt.Errorf("GetIntAlways: unrecognized type %v", v)
	}
}

// GetUintAlways is a safe, typed uint64 accessor. It will force conversion away
// from float64 and int64 values.
func (o Object) GetUintAlways(k string) (uint64, error) {
	switch v := o.KV[k].(type) {
	case float64:
		fl := o.KV[k].(float64)
		return uint64(fl), nil
	case int64:
		fl := o.KV[k].(int64)
		return uint64(fl), nil
	case uint64:
		fl := o.KV[k].(uint64)
		return fl, nil
	case string:
		fl := o.KV[k].(string)
		return strconv.ParseUint(fl, 10, 64)
	default:
		return 0, fmt.Errorf("GetIntAlways: unrecognized type %v", v)
	}
}

// Set is our setter
func (o *Object) Set(k string, v interface{}) {
	oldVal := o.Get(k)

	if oldVal != nil {
		// Avoid redundant Set()s
		if oldVal == v {
			return
		}
		o.FieldChanged(k, oldVal)
	}
	if o.GetSaved() {
		o.SetSaved(false)
	}
	o.rawSet(k, v)
}

// FieldChanged records the previous value for something that is about to be set
func (o Object) FieldChanged(k string, oldVal interface{}) {
	o.ChangedFields[k] = oldVal
}

func (o Object) rawSet(k string, v interface{}) {
	o.KV[k] = v
}

// ResetChangedFields can be used in conjunction with an ORM... For instance,
// once a Save() method is invoked
func (o *Object) ResetChangedFields() {
	o.ChangedFields = make(map[string]interface{})
}

// SetSaved is our setter for the 'object is saved' status field
func (o *Object) SetSaved(status bool) {
	o.saved = status
}

// GetSaved is our getter for the 'object is saved' bool field
func (o Object) GetSaved() bool {
	return o.saved
}
<commit_msg>minor handling for nils, not sure this is ideal<commit_after>// Package object is an abstract data record which tracks state changes.
// It's meant to make it easier to map key-value records into ORM / RDBMS
// structures.
// The state change tracking can be useful when the values of primary keys
// need to change. (Changing a foreign key on a table with a composite key, for example)
package object

import (
	"fmt"
	"strconv"
)

// Array is our object array container to facilitate a couple of instances
// where slices are needed.
type Array []*Object

// Object struct encapsulates our key-value pairs and a single-item per-key history
// of the previous value stored for a given key.
type Object struct {
	Type          string
	KV            map[string]interface{} `json:"KV"`
	ChangedFields map[string]interface{} `json:"ChangedFields"`
	Children      map[string]Array
	saved         bool
}

// New is an empty constructor
func New(objType string) *Object {
	return &Object{Type: objType, KV: makeEmptyMap(), ChangedFields: makeEmptyMap(), Children: makeEmptyChildrenMap(), saved: false}
}

// MakeArray will construct an array of length 'size'
func MakeArray(size int) Array {
	objAry := make(Array, size)
	return objAry
}

// NewArray is our single-object array constructor
func NewArray(val *Object) Array {
	objAry := make(Array, 1)
	objAry[0] = val
	return objAry
}

// NewEmptyArray is an empty array constructor
func NewEmptyArray() Array {
	objAry := make(Array, 0)
	return objAry
}

func makeEmptyMap() map[string]interface{} {
	return make(map[string]interface{})
}

func makeEmptyChildrenMap() map[string]Array {
	return make(map[string]Array)
}

// Get is our accessor
func (o Object) Get(k string) interface{} {
	return o.KV[k]
}

// GetString is a safe, typed string accessor
func (o Object) GetString(k string) (string, bool) {
	v, ok := o.KV[k].(string)
	return v, ok
}

// GetInt is a safe, typed int64 accessor
func (o Object) GetInt(k string) (int64, bool) {
	v, ok := o.KV[k].(int64)
	return v, ok
}

// GetFloat is a safe, typed int64 accessor
func (o Object) GetFloat(k string) (float64, bool) {
	v, ok := o.KV[k].(float64)
	return v, ok
}

// GetIntAlways is a safe, typed int64 accessor. It will force conversion away
// from float64 and uint64 values.
func (o Object) GetIntAlways(k string) (int64, error) {
	switch v := o.KV[k].(type) {
	case float64:
		fl := o.KV[k].(float64)
		return int64(fl), nil
	case int64:
		fl := o.KV[k].(int64)
		return fl, nil
	case uint64:
		fl := o.KV[k].(uint64)
		return int64(fl), nil
	case string:
		fl := o.KV[k].(string)
		return strconv.ParseInt(fl, 10, 64)
	case nil:
		return 0, nil
	default:
		return 0, fmt.Errorf("GetIntAlways: unrecognized type %v", v)
	}
}

// GetUintAlways is a safe, typed uint64 accessor. It will force conversion away
// from float64 and int64 values.
func (o Object) GetUintAlways(k string) (uint64, error) {
	switch v := o.KV[k].(type) {
	case float64:
		fl := o.KV[k].(float64)
		return uint64(fl), nil
	case int64:
		fl := o.KV[k].(int64)
		return uint64(fl), nil
	case uint64:
		fl := o.KV[k].(uint64)
		return fl, nil
	case string:
		fl := o.KV[k].(string)
		return strconv.ParseUint(fl, 10, 64)
	default:
		return 0, fmt.Errorf("GetIntAlways: unrecognized type %v", v)
	}
}

// Set is our setter
func (o *Object) Set(k string, v interface{}) {
	oldVal := o.Get(k)

	if oldVal != nil {
		// Avoid redundant Set()s
		if oldVal == v {
			return
		}
		o.FieldChanged(k, oldVal)
	}
	if o.GetSaved() {
		o.SetSaved(false)
	}
	o.rawSet(k, v)
}

// FieldChanged records the previous value for something that is about to be set
func (o Object) FieldChanged(k string, oldVal interface{}) {
	o.ChangedFields[k] = oldVal
}

func (o Object) rawSet(k string, v interface{}) {
	o.KV[k] = v
}

// ResetChangedFields can be used in conjunction with an ORM... For instance,
// once a Save() method is invoked
func (o *Object) ResetChangedFields() {
	o.ChangedFields = make(map[string]interface{})
}

// SetSaved is our setter for the 'object is saved' status field
func (o *Object) SetSaved(status bool) {
	o.saved = status
}

// GetSaved is our getter for the 'object is saved' bool field
func (o Object) GetSaved() bool {
	return o.saved
}
<|endoftext|>


<commit_before>package index
<commit_msg>Test snapshot package<commit_after>package index

import (
	"os"
	"reflect"
	"strconv"
	"testing"
)

var (
	tmpdirsnap = "/tmp/test-snap/"
)

func setupro(t testing.TB) func() {
	if err := os.RemoveAll(tmpdirsnap); err != nil {
		t.Fatal(err)
	}

	if err := os.MkdirAll(tmpdirsnap, 0777); err != nil {
		t.Fatal(err)
	}

	return func() {
		if err := os.RemoveAll(tmpdirsnap); err != nil {
			t.Fatal(err)
		}
	}
}

func TestSnapshot(t *testing.T) {
	defer setupro(t)()

	tree := WrapNode(nil)
	for i := 0; i < 3; i++ {
		istr := strconv.Itoa(i)
		flds := []string{"r" + istr, "b" + istr}
		tree.Ensure(flds).Node.RecordID = int64(i)
	}

	s, err := writeSnapshot(tmpdirsnap, tree)
	if err != nil {
		t.Fatal(err)
	}

	for i := 0; i < 3; i++ {
		istr := strconv.Itoa(i)
		flds := []string{"r" + istr, "b" + istr}

		res, err := s.RootNode.FindOne(flds)
		if err != nil {
			t.Fatal(err)
		} else if res == nil {
			t.Fatal("missing res")
		}

		if !reflect.DeepEqual(flds, res.Fields) || res.RecordID != int64(i) {
			t.Fatal("wrong value")
		}
	}

	if err := s.Close(); err != nil {
		t.Fatal(err)
	}

	for i := 0; i < 3; i++ {
		if s, err = LoadSnap(tmpdirsnap); err != nil {
			t.Fatal(err)
		}

		for i := 0; i < 3; i++ {
			istr := strconv.Itoa(i)
			name := "r" + istr

			res, ok := s.RootNode.Children[name]
			if !ok {
				t.Fatal("should have entry")
			} else if res != nil {
				t.Fatal("should be nil")
			}
		}

		for i := 0; i < 3; i++ {
			istr := strconv.Itoa(i)

			br, err := s.LoadBranch("r" + istr)
			if err != nil {
				t.Fatal(err)
			}

			bb, ok := br.Children["b"+istr]
			if !ok || bb == nil {
				t.Fatal("invalid child node")
			}

			ex := tree.Children["r"+istr].Children["b"+istr]
			if bb.Node.RecordID != ex.Node.RecordID {
				t.Fatal("wrong node value")
			}
		}

		if err := s.Close(); err != nil {
			t.Fatal(err)
		}
	}
}
<|endoftext|>


<commit_before>// Copyright (c) 2015 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package gen

import (
	"math/rand"
	"reflect"
	"testing"
	"testing/quick"
	"time"

	tc "go.uber.org/thriftrw/gen/internal/tests/containers"
	te "go.uber.org/thriftrw/gen/internal/tests/enums"
	tx "go.uber.org/thriftrw/gen/internal/tests/exceptions"
	tf "go.uber.org/thriftrw/gen/internal/tests/services"
	ts "go.uber.org/thriftrw/gen/internal/tests/structs"
	td "go.uber.org/thriftrw/gen/internal/tests/typedefs"
	tu "go.uber.org/thriftrw/gen/internal/tests/unions"
	"go.uber.org/thriftrw/wire"
	"go.uber.org/zap/zapcore"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func defaultValueGenerator(typ reflect.Type) func(*testing.T, *rand.Rand) thriftType {
	return func(t *testing.T, rand *rand.Rand) thriftType {
		for {
			// We will keep trying to generate a value until a valid one
			// is found.

			v, ok := quick.Value(typ, rand)
			require.True(t, ok, "failed to generate a value")

			tval := v.Addr().Interface().(thriftType)

			// TODO(abg): ToWire + EvaluateValue to validate here means we end
			// up serializing this value twice. We may want to include a
			// Validate method on generated types.

			w, err := tval.ToWire()
			if err != nil {
				// Value fails validity check. Try again.
				continue
			}

			// Because we evaluate collections lazily, validation issues
			// with items in them won't be known until we try to serialize
			// it or explicitly evaluate the lazy lists with
			// wire.EvaluateValue.
			if err := wire.EvaluateValue(w); err != nil {
				// Value fails validity check. Try again.
				continue
			}

			return tval
		}
	}
}

// enumValueGenerator builds a generator for random enum values given the
// `*_Values` function for that enum.
func enumValueGenerator(valuesFunc interface{}) func(*testing.T, *rand.Rand) thriftType {
	vfunc := reflect.ValueOf(valuesFunc)
	typ := vfunc.Type().Out(0).Elem() // Foo_Values() []Foo -> Foo
	return func(t *testing.T, rand *rand.Rand) thriftType {
		knownValues := vfunc.Call(nil)[0]

		var giveV reflect.Value
		// Flip a coin to decide whether we're evaluating a known or
		// unknown value.
		if rand.Int()%2 == 0 && knownValues.Len() > 0 {
			// Pick a known value at random
			giveV = knownValues.Index(rand.Intn(knownValues.Len()))
		} else {
			// give = MyEnum($randomValue)
			giveV = reflect.New(typ).Elem()
			giveV.Set(reflect.ValueOf(rand.Int31()).Convert(typ))
		}

		return giveV.Addr().Interface().(thriftType)
	}
}

func TestQuickRoundTrip(t *testing.T) {
	type testCase struct {
		// Sample value of the type to be tested.
		Sample interface{}

		// Specifies how we generate valid values of this type. Defaults to
		// defaultValueGenerator(Type) if unspecified.
		Generator func(*testing.T, *rand.Rand) thriftType

		// If set, logging for this type will not be tested. This is needed
		// for typedefs of primitives which can't implement ArrayMarshaler or
		// ObjectMarshaler.
		NoLog bool
	}

	// The following types from our tests have been skipped.
	// - unions.ArbitraryValue: Self-reference causes testing/quick to loop
	//   for too long
	// - services.KeyValue_SetValue_Args{}: Accepts an ArbitraryValue
	// - services.KeyValue_SetValueV2_Args: Accepts an ArbitraryValue
	// - services.KeyValue_GetManyValues_Result{}: Produces an ArbitraryValue
	// - services.KeyValue_GetValue_Result{}: Produces an ArbitraryValue

	// TODO(abg): ^Use custom generators to make this not-a-problem.

	tests := []testCase{
		// structs, unions, and exceptions
		{Sample: tc.ContainersOfContainers{}},
		{Sample: tc.EnumContainers{}},
		{Sample: tc.ListOfConflictingEnums{}},
		{Sample: tc.ListOfConflictingUUIDs{}},
		{Sample: tc.MapOfBinaryAndString{}},
		{Sample: tc.PrimitiveContainersRequired{}},
		{Sample: tc.PrimitiveContainers{}},
		{Sample: td.DefaultPrimitiveTypedef{}},
		{Sample: td.Event{}},
		{Sample: td.I128{}},
		{Sample: td.Transition{}},
		{Sample: te.StructWithOptionalEnum{}},
		{Sample: tf.Cache_Clear_Args{}},
		{Sample: tf.Cache_ClearAfter_Args{}},
		{Sample: tf.ConflictingNames_SetValue_Args{}},
		{Sample: tf.ConflictingNames_SetValue_Result{}},
		{Sample: tf.ConflictingNamesSetValueArgs{}},
		{Sample: tf.InternalError{}},
		{Sample: tf.KeyValue_DeleteValue_Args{}},
		{Sample: tf.KeyValue_DeleteValue_Result{}},
		{Sample: tf.KeyValue_GetManyValues_Args{}},
		{Sample: tf.KeyValue_GetValue_Args{}},
		{Sample: tf.KeyValue_SetValue_Result{}},
		{Sample: tf.KeyValue_SetValueV2_Result{}},
		{Sample: tf.KeyValue_Size_Args{}},
		{Sample: tf.KeyValue_Size_Result{}},
		{Sample: tf.NonStandardServiceName_NonStandardFunctionName_Args{}},
		{Sample: tf.NonStandardServiceName_NonStandardFunctionName_Result{}},
		{Sample: ts.ContactInfo{}},
		{Sample: ts.DefaultsStruct{}},
		{Sample: ts.Edge{}},
		{Sample: ts.EmptyStruct{}},
		{Sample: ts.Frame{}},
		{Sample: ts.GoTags{}},
		{Sample: ts.Graph{}},
		{Sample: ts.Node{}},
		{Sample: ts.Omit{}},
		{Sample: ts.Point{}},
		{Sample: ts.PrimitiveOptionalStruct{}},
		{Sample: ts.PrimitiveRequiredStruct{}},
		{Sample: ts.Rename{}},
		{Sample: ts.Size{}},
		{Sample: ts.StructLabels{}},
		{Sample: ts.User{}},
		{Sample: ts.ZapOptOutStruct{}},
		{Sample: tu.Document{}},
		{Sample: tu.EmptyUnion{}},
		{Sample: tx.DoesNotExistException{}},
		{Sample: tx.EmptyException{}},

		// typedefs
		{Sample: td.BinarySet{}},
		{Sample: td.EdgeMap{}},
		{Sample: td.FrameGroup{}},
		{Sample: td.MyEnum(0)},
		{Sample: td.PDF{}, NoLog: true},
		{Sample: td.PointMap{}},
		{Sample: td.State(""), NoLog: true},
		{Sample: td.StateMap{}},
		{Sample: td.Timestamp(0), NoLog: true},
		{Sample: td.UUID{}},

		// enums
		{
			Sample:    te.EmptyEnum(0),
			Generator: enumValueGenerator(te.EmptyEnum_Values),
		},
		{
			Sample:    te.EnumDefault(0),
			Generator: enumValueGenerator(te.EnumDefault_Values),
		},
		{
			Sample:    te.EnumWithDuplicateName(0),
			Generator: enumValueGenerator(te.EnumWithDuplicateName_Values),
		},
		{
			Sample:    te.EnumWithDuplicateValues(0),
			Generator: enumValueGenerator(te.EnumWithDuplicateValues_Values),
		},
		{
			Sample:    te.EnumWithLabel(0),
			Generator: enumValueGenerator(te.EnumWithLabel_Values),
		},
		{
			Sample:    te.EnumWithValues(0),
			Generator: enumValueGenerator(te.EnumWithValues_Values),
		},
		{
			Sample:    te.LowerCaseEnum(0),
			Generator: enumValueGenerator(te.LowerCaseEnum_Values),
		},
		{
			Sample:    te.RecordType(0),
			Generator: enumValueGenerator(te.RecordType_Values),
		},
		{
			Sample:    te.RecordTypeValues(0),
			Generator: enumValueGenerator(te.RecordTypeValues_Values),
		},
	}

	// Log the seed so that we can reproduce this if it ever fails.
	seed := time.Now().UnixNano()
	rand := rand.New(rand.NewSource(seed))
	t.Logf("Using seed %v for testing/quick", seed)

	const numValues = 1000 // number of values to test against
	for _, tt := range tests {
		typ := reflect.TypeOf(tt.Sample)
		t.Run(typ.Name(), func(t *testing.T) {
			generator := tt.Generator
			if generator == nil {
				generator = defaultValueGenerator(typ)
			}

			values := make([]thriftType, numValues)
			for i := range values {
				values[i] = generator(t, rand)
			}

			t.Run("Thrift", func(t *testing.T) {
				for _, give := range values {
					w, err := give.ToWire()
					require.NoError(t, err, "failed to Thrift encode %v", give)

					got := reflect.New(typ).Interface().(thriftType)
					require.NoError(t, got.FromWire(w), "failed to Thrift decode from %v", w)

					assert.Equal(t, got, give)
				}
			})

			t.Run("String", func(t *testing.T) {
				for _, give := range values {
					assert.NotPanics(t, func() {
						_ = give.String()
					}, "failed to String %#v", give)
				}
			})

			if !tt.NoLog {
				t.Run("Zap", func(t *testing.T) {
					for _, give := range values {
						assert.NotPanics(t, func() {
							enc := zapcore.NewMapObjectEncoder()

							if obj, ok := give.(zapcore.ObjectMarshaler); ok {
								assert.NoErrorf(t, obj.MarshalLogObject(enc), "failed to log %v", give)
								return
							}

							if arr, ok := give.(zapcore.ArrayMarshaler); ok {
								assert.NoErrorf(t, enc.AddArray("values", arr), "failed to log %v", give)
								return
							}

							t.Fatal(
								"Type does not implement zapcore.ObjectMarshaler or zapcore.ArrayMarshaler. "+
									"Did you mean to add NoLog?", typ)
						}, "failed to log %v", give)
					}
				})
			}
		})
	}
}
<commit_msg>gen/test/quick: JSON round-tripping<commit_after>// Copyright (c) 2015 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package gen

import (
	"encoding/json"
	"math/rand"
	"reflect"
	"testing"
	"testing/quick"
	"time"

	tc "go.uber.org/thriftrw/gen/internal/tests/containers"
	te "go.uber.org/thriftrw/gen/internal/tests/enums"
	tx "go.uber.org/thriftrw/gen/internal/tests/exceptions"
	tf "go.uber.org/thriftrw/gen/internal/tests/services"
	ts "go.uber.org/thriftrw/gen/internal/tests/structs"
	td "go.uber.org/thriftrw/gen/internal/tests/typedefs"
	tu "go.uber.org/thriftrw/gen/internal/tests/unions"
	"go.uber.org/thriftrw/wire"
	"go.uber.org/zap/zapcore"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func defaultValueGenerator(typ reflect.Type) func(*testing.T, *rand.Rand) thriftType {
	return func(t *testing.T, rand *rand.Rand) thriftType {
		for {
			// We will keep trying to generate a value until a valid one
			// is found.

			v, ok := quick.Value(typ, rand)
			require.True(t, ok, "failed to generate a value")

			tval := v.Addr().Interface().(thriftType)

			// TODO(abg): ToWire + EvaluateValue to validate here means we end
			// up serializing this value twice. We may want to include a
			// Validate method on generated types.

			w, err := tval.ToWire()
			if err != nil {
				// Value fails validity check. Try again.
				continue
			}

			// Because we evaluate collections lazily, validation issues
			// with items in them won't be known until we try to serialize
			// it or explicitly evaluate the lazy lists with
			// wire.EvaluateValue.
			if err := wire.EvaluateValue(w); err != nil {
				// Value fails validity check. Try again.
				continue
			}

			return tval
		}
	}
}

// enumValueGenerator builds a generator for random enum values given the
// `*_Values` function for that enum.
func enumValueGenerator(valuesFunc interface{}) func(*testing.T, *rand.Rand) thriftType {
	vfunc := reflect.ValueOf(valuesFunc)
	typ := vfunc.Type().Out(0).Elem() // Foo_Values() []Foo -> Foo
	return func(t *testing.T, rand *rand.Rand) thriftType {
		knownValues := vfunc.Call(nil)[0]

		var giveV reflect.Value
		// Flip a coin to decide whether we're evaluating a known or
		// unknown value.
		if rand.Int()%2 == 0 && knownValues.Len() > 0 {
			// Pick a known value at random
			giveV = knownValues.Index(rand.Intn(knownValues.Len()))
		} else {
			// give = MyEnum($randomValue)
			giveV = reflect.New(typ).Elem()
			giveV.Set(reflect.ValueOf(rand.Int31()).Convert(typ))
		}

		return giveV.Addr().Interface().(thriftType)
	}
}

func TestQuickRoundTrip(t *testing.T) {
	type testCase struct {
		// Sample value of the type to be tested.
		Sample interface{}

		// Specifies how we generate valid values of this type. Defaults to
		// defaultValueGenerator(Type) if unspecified.
		Generator func(*testing.T, *rand.Rand) thriftType

		// If set, logging for this type will not be tested. This is needed
		// for typedefs of primitives which can't implement ArrayMarshaler or
		// ObjectMarshaler.
		NoLog bool

		// Whether we should evaluate JSON round-tripping. This is opt-in
		// rather than opt-out because struct types that use collections won't
		// round-trip with JSON successfully due to nil versus empty
		// collection differences.
		JSON bool
	}

	// The following types from our tests have been skipped.
	// - unions.ArbitraryValue: Self-reference causes testing/quick to loop
	//   for too long
	// - services.KeyValue_SetValue_Args{}: Accepts an ArbitraryValue
	// - services.KeyValue_SetValueV2_Args: Accepts an ArbitraryValue
	// - services.KeyValue_GetManyValues_Result{}: Produces an ArbitraryValue
	// - services.KeyValue_GetValue_Result{}: Produces an ArbitraryValue

	// TODO(abg): ^Use custom generators to make this not-a-problem.

	tests := []testCase{
		// structs, unions, and exceptions
		{Sample: tc.ContainersOfContainers{}},
		{Sample: tc.EnumContainers{}},
		{Sample: tc.ListOfConflictingEnums{}},
		{Sample: tc.ListOfConflictingUUIDs{}},
		{Sample: tc.MapOfBinaryAndString{}},
		{Sample: tc.PrimitiveContainersRequired{}},
		{Sample: tc.PrimitiveContainers{}},
		{Sample: td.DefaultPrimitiveTypedef{}},
		{Sample: td.Event{}},
		{Sample: td.I128{}},
		{Sample: td.Transition{}},
		{Sample: te.StructWithOptionalEnum{}},
		{Sample: tf.Cache_Clear_Args{}},
		{Sample: tf.Cache_ClearAfter_Args{}},
		{Sample: tf.ConflictingNames_SetValue_Args{}},
		{Sample: tf.ConflictingNames_SetValue_Result{}},
		{Sample: tf.ConflictingNamesSetValueArgs{}},
		{Sample: tf.InternalError{}},
		{Sample: tf.KeyValue_DeleteValue_Args{}},
		{Sample: tf.KeyValue_DeleteValue_Result{}},
		{Sample: tf.KeyValue_GetManyValues_Args{}},
		{Sample: tf.KeyValue_GetValue_Args{}},
		{Sample: tf.KeyValue_SetValue_Result{}},
		{Sample: tf.KeyValue_SetValueV2_Result{}},
		{Sample: tf.KeyValue_Size_Args{}},
		{Sample: tf.KeyValue_Size_Result{}},
		{Sample: tf.NonStandardServiceName_NonStandardFunctionName_Args{}},
		{Sample: tf.NonStandardServiceName_NonStandardFunctionName_Result{}},
		{Sample: ts.ContactInfo{}},
		{Sample: ts.DefaultsStruct{}},
		{Sample: ts.Edge{}},
		{Sample: ts.EmptyStruct{}},
		{Sample: ts.Frame{}},
		{Sample: ts.GoTags{}},
		{Sample: ts.Graph{}},
		{Sample: ts.Node{}},
		{Sample: ts.Omit{}},
		{Sample: ts.Point{}},
		{Sample: ts.PrimitiveOptionalStruct{}},
		{Sample: ts.PrimitiveRequiredStruct{}},
		{Sample: ts.Rename{}},
		{Sample: ts.Size{}},
		{Sample: ts.StructLabels{}},
		{Sample: ts.User{}},
		{Sample: ts.ZapOptOutStruct{}},
		{Sample: tu.Document{}},
		{Sample: tu.EmptyUnion{}},
		{Sample: tx.DoesNotExistException{}},
		{Sample: tx.EmptyException{}},

		// typedefs
		{Sample: td.BinarySet{}},
		{Sample: td.EdgeMap{}},
		{Sample: td.FrameGroup{}},
		{Sample: td.MyEnum(0)},
		{Sample: td.PDF{}, NoLog: true},
		{Sample: td.PointMap{}},
		{Sample: td.State(""), NoLog: true},
		{Sample: td.StateMap{}},
		{Sample: td.Timestamp(0), NoLog: true},
		{Sample: td.UUID{}},

		// enums
		{
			Sample:    te.EmptyEnum(0),
			Generator: enumValueGenerator(te.EmptyEnum_Values),
			JSON:      true,
		},
		{
			Sample:    te.EnumDefault(0),
			Generator: enumValueGenerator(te.EnumDefault_Values),
			JSON:      true,
		},
		{
			Sample:    te.EnumWithDuplicateName(0),
			Generator: enumValueGenerator(te.EnumWithDuplicateName_Values),
			JSON:      true,
		},
		{
			Sample:    te.EnumWithDuplicateValues(0),
			Generator: enumValueGenerator(te.EnumWithDuplicateValues_Values),
			JSON:      true,
		},
		{
			Sample:    te.EnumWithLabel(0),
			Generator: enumValueGenerator(te.EnumWithLabel_Values),
			JSON:      true,
		},
		{
			Sample:    te.EnumWithValues(0),
			Generator: enumValueGenerator(te.EnumWithValues_Values),
			JSON:      true,
		},
		{
			Sample:    te.LowerCaseEnum(0),
			Generator: enumValueGenerator(te.LowerCaseEnum_Values),
			JSON:      true,
		},
		{
			Sample:    te.RecordType(0),
			Generator: enumValueGenerator(te.RecordType_Values),
			JSON:      true,
		},
		{
			Sample:    te.RecordTypeValues(0),
			Generator: enumValueGenerator(te.RecordTypeValues_Values),
			JSON:      true,
		},
	}

	// Log the seed so that we can reproduce this if it ever fails.
	seed := time.Now().UnixNano()
	rand := rand.New(rand.NewSource(seed))
	t.Logf("Using seed %v for testing/quick", seed)

	const numValues = 1000 // number of values to test against
	for _, tt := range tests {
		typ := reflect.TypeOf(tt.Sample)
		t.Run(typ.Name(), func(t *testing.T) {
			generator := tt.Generator
			if generator == nil {
				generator = defaultValueGenerator(typ)
			}

			values := make([]thriftType, numValues)
			for i := range values {
				values[i] = generator(t, rand)
			}

			t.Run("Thrift", func(t *testing.T) {
				for _, give := range values {
					w, err := give.ToWire()
					require.NoError(t, err, "failed to Thrift encode %v", give)

					got := reflect.New(typ).Interface().(thriftType)
					require.NoError(t, got.FromWire(w), "failed to Thrift decode from %v", w)

					assert.Equal(t, got, give)
				}
			})

			t.Run("String", func(t *testing.T) {
				for _, give := range values {
					assert.NotPanics(t, func() {
						_ = give.String()
					}, "failed to String %#v", give)
				}
			})

			if tt.JSON {
				t.Run("JSON", func(t *testing.T) {
					for _, giveValue := range values {
						give, ok := giveValue.(json.Marshaler)
						require.True(t, ok, "Type does not implement json.Marshaler")

						bs, err := give.MarshalJSON()
						require.NoError(t, err, "failed to encode %v", give)

						got, ok := reflect.New(typ).Interface().(json.Unmarshaler)
						require.True(t, ok, "Type does not implement json.Unmarshaler")

						require.NoError(t, got.UnmarshalJSON(bs), "failed to decode from %q", bs)
						assert.Equal(t, got, give, "could not round-trip")
					}
				})
			}

			if !tt.NoLog {
				t.Run("Zap", func(t *testing.T) {
					for _, give := range values {
						assert.NotPanics(t, func() {
							enc := zapcore.NewMapObjectEncoder()

							if obj, ok := give.(zapcore.ObjectMarshaler); ok {
								assert.NoErrorf(t, obj.MarshalLogObject(enc), "failed to log %v", give)
								return
							}

							if arr, ok := give.(zapcore.ArrayMarshaler); ok {
								assert.NoErrorf(t, enc.AddArray("values", arr), "failed to log %v", give)
								return
							}

							t.Fatal(
								"Type does not implement zapcore.ObjectMarshaler or zapcore.ArrayMarshaler. "+
									"Did you mean to add NoLog?", typ)
						}, "failed to log %v", give)
					}
				})
			}
		})
	}
}
<|endoftext|>


<commit_before>/*
Copyright (c) 2019 VMware, 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 cns

import (
	"context"
	"errors"

	cnstypes "github.com/vmware/govmomi/cns/types"
	"github.com/vmware/govmomi/object"
	vim25types "github.com/vmware/govmomi/vim25/types"
)

// GetTaskInfo gets the task info given a task
func GetTaskInfo(ctx context.Context, task *object.Task) (*vim25types.TaskInfo, error) {
	taskInfo, err := task.WaitForResult(ctx, nil)
	if err != nil {
		return nil, err
	}
	return taskInfo, nil
}

// GetTaskResult gets the task result given a task info
func GetTaskResult(ctx context.Context, taskInfo *vim25types.TaskInfo) (cnstypes.BaseCnsVolumeOperationResult, error) {
	if taskInfo == nil {
		return nil, errors.New("TaskInfo is empty")
	}
	if taskInfo.Result != nil {
		volumeOperationBatchResult := taskInfo.Result.(cnstypes.CnsVolumeOperationBatchResult)
		if &volumeOperationBatchResult == nil ||
			volumeOperationBatchResult.VolumeResults == nil ||
			len(volumeOperationBatchResult.VolumeResults) == 0 {
			return nil, errors.New("Cannot get VolumeOperationResult")
		}
		return volumeOperationBatchResult.VolumeResults[0], nil
	}
	return nil, errors.New("TaskInfo result is empty")
}

// GetTaskResultArray gets the task result array for a specified task info
func GetTaskResultArray(ctx context.Context, taskInfo *vim25types.TaskInfo) ([]cnstypes.BaseCnsVolumeOperationResult, error) {
	if taskInfo == nil {
		return nil, errors.New("TaskInfo is empty")
	}
	if taskInfo.Result != nil {
		volumeOperationBatchResult := taskInfo.Result.(cnstypes.CnsVolumeOperationBatchResult)
		if &volumeOperationBatchResult == nil ||
			volumeOperationBatchResult.VolumeResults == nil ||
			len(volumeOperationBatchResult.VolumeResults) == 0 {
			return nil, errors.New("Cannot get VolumeOperationResult")
		}
		return volumeOperationBatchResult.VolumeResults, nil
	}
	return nil, errors.New("TaskInfo result is empty")
}

// dropUnknownCreateSpecElements helps drop newly added elements in the CnsVolumeCreateSpec, which are not known to the prior vSphere releases
func dropUnknownCreateSpecElements(c *Client, createSpecList []cnstypes.CnsVolumeCreateSpec) []cnstypes.CnsVolumeCreateSpec {
	if c.serviceClient.Version == ReleaseVSAN67u3 {
		// Dropping optional fields not known to vSAN 6.7U3
		updatedcreateSpecList := make([]cnstypes.CnsVolumeCreateSpec, 0, len(createSpecList))
		for _, createSpec := range createSpecList {
			createSpec.Metadata.ContainerCluster.ClusterFlavor = ""
			createSpec.Metadata.ContainerCluster.ClusterDistribution = ""
			createSpec.Metadata.ContainerClusterArray = nil
			var updatedEntityMetadata []cnstypes.BaseCnsEntityMetadata
			for _, entityMetadata := range createSpec.Metadata.EntityMetadata {
				k8sEntityMetadata := interface{}(entityMetadata).(*cnstypes.CnsKubernetesEntityMetadata)
				k8sEntityMetadata.ClusterID = ""
				k8sEntityMetadata.ReferredEntity = nil
				updatedEntityMetadata = append(updatedEntityMetadata, cnstypes.BaseCnsEntityMetadata(k8sEntityMetadata))
			}
			createSpec.Metadata.EntityMetadata = updatedEntityMetadata
			_, ok := createSpec.BackingObjectDetails.(*cnstypes.CnsBlockBackingDetails)
			if ok {
				createSpec.BackingObjectDetails.(*cnstypes.CnsBlockBackingDetails).BackingDiskUrlPath = ""
			}
			updatedcreateSpecList = append(updatedcreateSpecList, createSpec)
		}
		createSpecList = updatedcreateSpecList
	} else if c.serviceClient.Version == ReleaseVSAN70 {
		updatedcreateSpecList := make([]cnstypes.CnsVolumeCreateSpec, 0, len(createSpecList))
		for _, createSpec := range createSpecList {
			createSpec.Metadata.ContainerCluster.ClusterDistribution = ""
			_, ok := createSpec.BackingObjectDetails.(*cnstypes.CnsBlockBackingDetails)
			if ok {
				createSpec.BackingObjectDetails.(*cnstypes.CnsBlockBackingDetails).BackingDiskUrlPath = ""
			}
			updatedcreateSpecList = append(updatedcreateSpecList, createSpec)
		}
		createSpecList = updatedcreateSpecList
	}
	return createSpecList
}

// dropUnknownVolumeMetadataUpdateSpecElements helps drop newly added elements in the CnsVolumeMetadataUpdateSpec, which are not known to the prior vSphere releases
func dropUnknownVolumeMetadataUpdateSpecElements(c *Client, updateSpecList []cnstypes.CnsVolumeMetadataUpdateSpec) []cnstypes.CnsVolumeMetadataUpdateSpec {
	// Dropping optional fields not known to vSAN 6.7U3
	if c.serviceClient.Version == ReleaseVSAN67u3 {
		updatedUpdateSpecList := make([]cnstypes.CnsVolumeMetadataUpdateSpec, 0, len(updateSpecList))
		for _, updateSpec := range updateSpecList {
			updateSpec.Metadata.ContainerCluster.ClusterFlavor = ""
			updateSpec.Metadata.ContainerCluster.ClusterDistribution = ""
			var updatedEntityMetadata []cnstypes.BaseCnsEntityMetadata
			for _, entityMetadata := range updateSpec.Metadata.EntityMetadata {
				k8sEntityMetadata := interface{}(entityMetadata).(*cnstypes.CnsKubernetesEntityMetadata)
				k8sEntityMetadata.ClusterID = ""
				k8sEntityMetadata.ReferredEntity = nil
				updatedEntityMetadata = append(updatedEntityMetadata, cnstypes.BaseCnsEntityMetadata(k8sEntityMetadata))
			}
			updateSpec.Metadata.ContainerClusterArray = nil
			updateSpec.Metadata.EntityMetadata = updatedEntityMetadata
			updatedUpdateSpecList = append(updatedUpdateSpecList, updateSpec)
		}
		updateSpecList = updatedUpdateSpecList
	} else if c.serviceClient.Version == ReleaseVSAN70 {
		updatedUpdateSpecList := make([]cnstypes.CnsVolumeMetadataUpdateSpec, 0, len(updateSpecList))
		for _, updateSpec := range updateSpecList {
			updateSpec.Metadata.ContainerCluster.ClusterDistribution = ""
			var updatedContainerClusterArray []cnstypes.CnsContainerCluster
			for _, containerCluster := range updateSpec.Metadata.ContainerClusterArray {
				containerCluster.ClusterDistribution = ""
				updatedContainerClusterArray = append(updatedContainerClusterArray, containerCluster)
			}
			updateSpec.Metadata.ContainerClusterArray = updatedContainerClusterArray
			updatedUpdateSpecList = append(updatedUpdateSpecList, updateSpec)
		}
		updateSpecList = updatedUpdateSpecList
	}
	return updateSpecList
}
<commit_msg>Drop clusterDistribution from vSAN 7.0 update and create spec elements<commit_after>/*
Copyright (c) 2019 VMware, 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 cns

import (
	"context"
	"errors"

	cnstypes "github.com/vmware/govmomi/cns/types"
	"github.com/vmware/govmomi/object"
	vim25types "github.com/vmware/govmomi/vim25/types"
)

// GetTaskInfo gets the task info given a task
func GetTaskInfo(ctx context.Context, task *object.Task) (*vim25types.TaskInfo, error) {
	taskInfo, err := task.WaitForResult(ctx, nil)
	if err != nil {
		return nil, err
	}
	return taskInfo, nil
}

// GetTaskResult gets the task result given a task info
func GetTaskResult(ctx context.Context, taskInfo *vim25types.TaskInfo) (cnstypes.BaseCnsVolumeOperationResult, error) {
	if taskInfo == nil {
		return nil, errors.New("TaskInfo is empty")
	}
	if taskInfo.Result != nil {
		volumeOperationBatchResult := taskInfo.Result.(cnstypes.CnsVolumeOperationBatchResult)
		if &volumeOperationBatchResult == nil ||
			volumeOperationBatchResult.VolumeResults == nil ||
			len(volumeOperationBatchResult.VolumeResults) == 0 {
			return nil, errors.New("Cannot get VolumeOperationResult")
		}
		return volumeOperationBatchResult.VolumeResults[0], nil
	}
	return nil, errors.New("TaskInfo result is empty")
}

// GetTaskResultArray gets the task result array for a specified task info
func GetTaskResultArray(ctx context.Context, taskInfo *vim25types.TaskInfo) ([]cnstypes.BaseCnsVolumeOperationResult, error) {
	if taskInfo == nil {
		return nil, errors.New("TaskInfo is empty")
	}
	if taskInfo.Result != nil {
		volumeOperationBatchResult := taskInfo.Result.(cnstypes.CnsVolumeOperationBatchResult)
		if &volumeOperationBatchResult == nil ||
			volumeOperationBatchResult.VolumeResults == nil ||
			len(volumeOperationBatchResult.VolumeResults) == 0 {
			return nil, errors.New("Cannot get VolumeOperationResult")
		}
		return volumeOperationBatchResult.VolumeResults, nil
	}
	return nil, errors.New("TaskInfo result is empty")
}

// dropUnknownCreateSpecElements helps drop newly added elements in the CnsVolumeCreateSpec, which are not known to the prior vSphere releases
func dropUnknownCreateSpecElements(c *Client, createSpecList []cnstypes.CnsVolumeCreateSpec) []cnstypes.CnsVolumeCreateSpec {
	updatedcreateSpecList := make([]cnstypes.CnsVolumeCreateSpec, 0, len(createSpecList))
	switch c.serviceClient.Version {
	case ReleaseVSAN67u3:
		// Dropping optional fields not known to vSAN 6.7U3
		for _, createSpec := range createSpecList {
			createSpec.Metadata.ContainerCluster.ClusterFlavor = ""
			createSpec.Metadata.ContainerCluster.ClusterDistribution = ""
			createSpec.Metadata.ContainerClusterArray = nil
			var updatedEntityMetadata []cnstypes.BaseCnsEntityMetadata
			for _, entityMetadata := range createSpec.Metadata.EntityMetadata {
				k8sEntityMetadata := interface{}(entityMetadata).(*cnstypes.CnsKubernetesEntityMetadata)
				k8sEntityMetadata.ClusterID = ""
				k8sEntityMetadata.ReferredEntity = nil
				updatedEntityMetadata = append(updatedEntityMetadata, cnstypes.BaseCnsEntityMetadata(k8sEntityMetadata))
			}
			createSpec.Metadata.EntityMetadata = updatedEntityMetadata
			_, ok := createSpec.BackingObjectDetails.(*cnstypes.CnsBlockBackingDetails)
			if ok {
				createSpec.BackingObjectDetails.(*cnstypes.CnsBlockBackingDetails).BackingDiskUrlPath = ""
			}
			updatedcreateSpecList = append(updatedcreateSpecList, createSpec)
		}
		createSpecList = updatedcreateSpecList
	case ReleaseVSAN70:
		// Dropping optional fields not known to vSAN 7.0
		for _, createSpec := range createSpecList {
			createSpec.Metadata.ContainerCluster.ClusterDistribution = ""
			var updatedContainerClusterArray []cnstypes.CnsContainerCluster
			for _, containerCluster := range createSpec.Metadata.ContainerClusterArray {
				containerCluster.ClusterDistribution = ""
				updatedContainerClusterArray = append(updatedContainerClusterArray, containerCluster)
			}
			createSpec.Metadata.ContainerClusterArray = updatedContainerClusterArray
			_, ok := createSpec.BackingObjectDetails.(*cnstypes.CnsBlockBackingDetails)
			if ok {
				createSpec.BackingObjectDetails.(*cnstypes.CnsBlockBackingDetails).BackingDiskUrlPath = ""
			}
			updatedcreateSpecList = append(updatedcreateSpecList, createSpec)
		}
		createSpecList = updatedcreateSpecList
	case ReleaseVSAN70u1:
		// Dropping optional fields not known to vSAN 7.0U1
		for _, createSpec := range createSpecList {
			createSpec.Metadata.ContainerCluster.ClusterDistribution = ""
			var updatedContainerClusterArray []cnstypes.CnsContainerCluster
			for _, containerCluster := range createSpec.Metadata.ContainerClusterArray {
				containerCluster.ClusterDistribution = ""
				updatedContainerClusterArray = append(updatedContainerClusterArray, containerCluster)
			}
			createSpec.Metadata.ContainerClusterArray = updatedContainerClusterArray
			updatedcreateSpecList = append(updatedcreateSpecList, createSpec)
		}
		createSpecList = updatedcreateSpecList
	}
	return createSpecList
}

// dropUnknownVolumeMetadataUpdateSpecElements helps drop newly added elements in the CnsVolumeMetadataUpdateSpec, which are not known to the prior vSphere releases
func dropUnknownVolumeMetadataUpdateSpecElements(c *Client, updateSpecList []cnstypes.CnsVolumeMetadataUpdateSpec) []cnstypes.CnsVolumeMetadataUpdateSpec {
	// Dropping optional fields not known to vSAN 6.7U3
	if c.serviceClient.Version == ReleaseVSAN67u3 {
		updatedUpdateSpecList := make([]cnstypes.CnsVolumeMetadataUpdateSpec, 0, len(updateSpecList))
		for _, updateSpec := range updateSpecList {
			updateSpec.Metadata.ContainerCluster.ClusterFlavor = ""
			updateSpec.Metadata.ContainerCluster.ClusterDistribution = ""
			var updatedEntityMetadata []cnstypes.BaseCnsEntityMetadata
			for _, entityMetadata := range updateSpec.Metadata.EntityMetadata {
				k8sEntityMetadata := interface{}(entityMetadata).(*cnstypes.CnsKubernetesEntityMetadata)
				k8sEntityMetadata.ClusterID = ""
				k8sEntityMetadata.ReferredEntity = nil
				updatedEntityMetadata = append(updatedEntityMetadata, cnstypes.BaseCnsEntityMetadata(k8sEntityMetadata))
			}
			updateSpec.Metadata.ContainerClusterArray = nil
			updateSpec.Metadata.EntityMetadata = updatedEntityMetadata
			updatedUpdateSpecList = append(updatedUpdateSpecList, updateSpec)
		}
		updateSpecList = updatedUpdateSpecList
	} else if c.serviceClient.Version == ReleaseVSAN70 || c.serviceClient.Version == ReleaseVSAN70u1 {
		updatedUpdateSpecList := make([]cnstypes.CnsVolumeMetadataUpdateSpec, 0, len(updateSpecList))
		for _, updateSpec := range updateSpecList {
			updateSpec.Metadata.ContainerCluster.ClusterDistribution = ""
			var updatedContainerClusterArray []cnstypes.CnsContainerCluster
			for _, containerCluster := range updateSpec.Metadata.ContainerClusterArray {
				containerCluster.ClusterDistribution = ""
				updatedContainerClusterArray = append(updatedContainerClusterArray, containerCluster)
			}
			updateSpec.Metadata.ContainerClusterArray = updatedContainerClusterArray
			updatedUpdateSpecList = append(updatedUpdateSpecList, updateSpec)
		}
		updateSpecList = updatedUpdateSpecList
	}
	return updateSpecList
}
<|endoftext|>