<commit_before>// goldpushk pushes Gold services to production. See go/goldpushk.
//
// Sample usage:
//
//   Deployment of a specific service:
//     $ goldpushk --service diffserver --instance chrome-gpu
//     $ goldpushk -s diffserver -i chrome-gpu
//
//   Deployment of a specific service across multiple instances:
//     $ goldpushk --service diffserver --instance chrome-gpu,skia
//     $ goldpushk -s diffserver -i chrome-gpu,skia
//
//   Deployment of all instances of a given service across all Gold instances:
//     $ goldpushk --service diffserver --instance all
//     $ goldpushk -s diffserver -i all
//
//   Deployment of all services corresponding to a specific Gold instance:
//     $ goldpushk --service all --instance chrome-gpu
//     $ goldpushk -s all -i chrome-gpu
//
//   Deployment of all instances of a given service, designating one of them as the canary:
//     $ goldpushk --service diffserver --instance all --canary skia:diffserver
//
//   Print out all Gold instances and services goldpushk is able to manage:
//     $ goldpushk --list

package main

import (
	"context"
	"errors"
	"fmt"
	"os"
	"sort"
	"strings"
	"text/tabwriter"

	"github.com/spf13/cobra"
	"go.skia.org/infra/go/skerr"
	"go.skia.org/infra/go/sklog"
	"go.skia.org/infra/go/util"
	"go.skia.org/infra/golden/cmd/goldpushk/goldpushk"
)

const (
	// Wildcard value for command line arguments.
	all = "all"

	// Environment variable with path to buildbot repository checkout directory.
	skiaInfraRootEnvVar = "SKIA_INFRA_ROOT"

	// Git repositories.
	skiaPublicConfigRepoUrl = "https://skia.googlesource.com/skia-public-config"
	skiaCorpConfigRepoUrl   = "https://skia.googlesource.com/skia-corp-config"
)

var (
	// Required flags.
	flagInstances []string
	flagServices  []string
	flagCanaries  []string

	// Optional flags.
	flagList                       bool
	flagDryRun                     bool
	flagNoCommit                   bool
	flagMinUptimeSeconds           int
	flagUptimePollFrequencySeconds int

	// Flags for debugging.
	flagLogToStdErr bool
	flagTesting     bool
)

func main() {
	rootCmd := &cobra.Command{
		Use:  "goldpushk",
		Long: "goldpushk pushes Gold services to production.",
		PersistentPreRun: func(cmd *cobra.Command, args []string) {
			logMode := sklog.SLogNone
			if flagLogToStdErr {
				logMode = sklog.SLogStderr
			}
			sklog.SetLogger(sklog.NewStdErrCloudLogger(logMode))
		},
		Run: func(cmd *cobra.Command, args []string) {
			run(cmd)
		},
	}

	rootCmd.Flags().SortFlags = false
	rootCmd.Flags().BoolVar(&flagList, "list", false, "List known Gold instances and services (tip: try combining this flag with --testing).")
	rootCmd.Flags().StringSliceVarP(&flagInstances, "instances", "i", []string{}, "[REQUIRED] Comma-delimited list of Gold instances to target (e.g. \"skia,flutter\"), or \""+all+"\" to target all instances.")
	rootCmd.Flags().StringSliceVarP(&flagServices, "services", "s", []string{}, "[REQUIRED] Comma-delimited list of services to target (e.g. \"skiacorrectness,diffserver\"), or \""+all+"\" to target all services.")
	rootCmd.Flags().StringSliceVarP(&flagCanaries, "canaries", "c", []string{}, "Comma-delimited subset of Gold services to use as canaries, written as instance:service pairs (e.g. \"skia:diffserver,flutter:skiacorrectness\")")
	rootCmd.Flags().BoolVar(&flagDryRun, "dryrun", false, "Do everything except applying the new configuration to Kubernetes and committing changes to Git.")
	rootCmd.Flags().BoolVar(&flagNoCommit, "no-commit", false, "Do not commit configuration changes to the skia-public-config or skia-corp-config Git repositories.")
	rootCmd.Flags().IntVar(&flagMinUptimeSeconds, "min-uptime", 30, "Minimum uptime in seconds required for all services before exiting the monitoring step.")
	rootCmd.Flags().IntVar(&flagUptimePollFrequencySeconds, "poll-freq", 3, "How often to poll Kubernetes for service uptimes, in seconds.")
	rootCmd.Flags().BoolVar(&flagLogToStdErr, "logtostderr", false, "Log debug information to stderr. No logs will be produced if this flag is not set.")
	rootCmd.Flags().BoolVar(&flagTesting, "testing", false, "Do not deploy any production services; use testing services instead.")

	if _, err := rootCmd.ExecuteC(); err != nil {
		sklog.Fatalf("Error while running Cobra command: %s", err)
	}
}

func run(cmd *cobra.Command) {
	// Get set of deployable units. Used as the source of truth across goldpushk.
	var deployableUnitSet goldpushk.DeployableUnitSet
	if flagTesting {
		deployableUnitSet = goldpushk.TestingDeployableUnits()
	} else {
		deployableUnitSet = goldpushk.ProductionDeployableUnits()
	}

	// If --list is passed, print known services and exit. This takes into account flag --testing.
	if flagList {
		if err := listKnownServices(deployableUnitSet); err != nil {
			sklog.Fatalf("Error while printing list of known services: %s", err)
		}
		return
	}

	// If --list was not provided, validate presence of flags --services and --instances.
	if len(flagInstances) == 0 {
		fmt.Println("Error: flag \"instances\" is required.")
		if err := cmd.Usage(); err != nil {
			sklog.Fatalf("Error while printing usage: %s", err)
		}
		os.Exit(1)
	}
	if len(flagServices) == 0 {
		fmt.Println("Error: flag \"services\" is required.")
		if err := cmd.Usage(); err != nil {
			sklog.Fatalf("Error while printing usage: %s", err)
		}
		os.Exit(1)
	}

	// Parse and validate command line flags.
	deployableUnits, canariedDeployableUnits, err := parseAndValidateFlags(deployableUnitSet, flagInstances, flagServices, flagCanaries)
	if err != nil {
		fmt.Printf("Error: %s.\n", err)
		os.Exit(1)
	}

	// Read environment variables.
	skiaInfraRoot, ok := os.LookupEnv(skiaInfraRootEnvVar)
	if !ok {
		fmt.Printf("Error: environment variable %s not set.", skiaInfraRootEnvVar)
		os.Exit(1)
	}

	// Build goldpushk instance.
	gpk := goldpushk.New(deployableUnits, canariedDeployableUnits, skiaInfraRoot, flagDryRun, flagNoCommit, flagMinUptimeSeconds, flagUptimePollFrequencySeconds, skiaPublicConfigRepoUrl, skiaCorpConfigRepoUrl)

	// Run goldpushk.
	if err = gpk.Run(context.Background()); err != nil {
		fmt.Printf("Error: %s.\n", err)
		os.Exit(1)
	}
}

// listKnownServices prints out a table of known services.
func listKnownServices(deployableUnitSet goldpushk.DeployableUnitSet) error {
	mode := "production"
	if flagTesting {
		mode = "testing"
	}
	fmt.Printf("Known Gold instances and services (%s):\n", mode)

	// Print out table header.
	w := tabwriter.NewWriter(os.Stdout, 10, 0, 2, ' ', 0)
	if _, err := fmt.Fprintln(w, "\nINSTANCE\tSERVICE\tCANONICAL NAME"); err != nil {
		return skerr.Wrap(err)
	}

	// Print out table body.
	for _, instance := range deployableUnitSet.KnownInstances() {
		for _, service := range deployableUnitSet.KnownServices() {
			unit, ok := deployableUnitSet.Get(goldpushk.DeployableUnitID{Instance: instance, Service: service})
			if ok {
				if _, err := fmt.Fprintf(w, "%s\t%s\t%s\n", instance, service, unit.CanonicalName()); err != nil {
					return skerr.Wrap(err)
				}
			}
		}
	}

	// Flush output and return.
	if err := w.Flush(); err != nil {
		return skerr.Wrap(err)
	}
	return nil
}

// containsWildcardValue determines whether or not a flag contains the special
// "all" wildcard value.
func containsWildcardValue(flag []string) bool {
	for _, value := range flag {
		if value == all {
			return true
		}
	}
	return false
}

// parseAndValidateFlags validates the given command line flags, retrieves the
// corresponding DeployableUnits from the given DeployableUnitSet and returns
// them as two separate slices according to whether or not they were marked for
// canarying.
func parseAndValidateFlags(deployableUnitSet goldpushk.DeployableUnitSet, instances, services, canaries []string) (deployableUnits, canariedDeployableUnits []goldpushk.DeployableUnit, err error) {
	// Deduplicate inputs.
	instances = util.SSliceDedup(instances)
	services = util.SSliceDedup(services)
	canaries = util.SSliceDedup(canaries)

	// Determine whether --instances or --services are set to "all".
	allInstances := containsWildcardValue(instances)
	allServices := containsWildcardValue(services)

	// If --instances or --services contain the special "all" value, they should
	// not contain any other values.
	if allInstances && len(instances) != 1 {
		return nil, nil, errors.New("flag --instances should contain either \"all\" or a list of Gold instances, but not both")
	}
	if allServices && len(services) != 1 {
		return nil, nil, errors.New("flag --services should contain either \"all\" or a list of Gold services, but not both")
	}
	if allInstances && allServices {
		return nil, nil, errors.New("cannot set both --instances and --services to \"all\"")
	}

	// Validate instances.
	if !allInstances {
		for _, instanceStr := range instances {
			if !deployableUnitSet.IsKnownInstance(goldpushk.Instance(instanceStr)) {
				return nil, nil, fmt.Errorf("unknown Gold instance: \"%s\"", instanceStr)
			}
		}
	}

	// Validate services.
	if !allServices {
		for _, serviceStr := range services {
			if !deployableUnitSet.IsKnownService(goldpushk.Service(serviceStr)) {
				return nil, nil, fmt.Errorf("unknown Gold service: \"%s\"", serviceStr)
			}
		}
	}

	// This slice will be populated with the subset of the cartesian product of
	// flags --instances and --services that is found in the services map.
	var servicesToDeploy []goldpushk.DeployableUnitID

	// Determines whether or not an instance/service pair should be canaried.
	isMarkedForCanarying := map[goldpushk.DeployableUnitID]bool{}

	// Determine the set of instances over which to iterate to compute the
	// cartesian product of flags --instances and --services.
	var instanceIterationSet []goldpushk.Instance
	if containsWildcardValue(instances) {
		// Handle the "all" value.
		instanceIterationSet = deployableUnitSet.KnownInstances()
	} else {
		for _, instanceStr := range instances {
			instanceIterationSet = append(instanceIterationSet, goldpushk.Instance(instanceStr))
		}
	}

	// Determine the set of services over which to iterate to compute the
	// cartesian product of flags --instances and --services.
	var serviceIterationSet []goldpushk.Service
	if containsWildcardValue(services) {
		// Handle the "all" value.
		serviceIterationSet = deployableUnitSet.KnownServices()
	} else {
		for _, serviceStr := range services {
			serviceIterationSet = append(serviceIterationSet, goldpushk.Service(serviceStr))
		}
	}

	// Iterate over the cartesian product of flags --instances and --services.
	for _, instance := range instanceIterationSet {
		for _, service := range serviceIterationSet {
			id := goldpushk.DeployableUnitID{
				Instance: instance,
				Service:  service,
			}

			// Skip if the current instance/service combination is not found in the services map.
			if _, ok := deployableUnitSet.Get(id); !ok {
				continue
			}

			// Save instance/service pair, which is not marked for canarying by default.
			servicesToDeploy = append(servicesToDeploy, id)
			isMarkedForCanarying[id] = false
		}
	}

	// Fail if --instances and --services didn't match any services in the services map.
	if len(servicesToDeploy) == 0 {
		return nil, nil, errors.New("no known Gold services match the values supplied with --instances and --services")
	}

	// Iterate over the --canaries flag.
	for _, canaryStr := range canaries {
		// Validate format and extract substrings.
		canaryStrSplit := strings.Split(canaryStr, ":")
		if len(canaryStrSplit) != 2 || len(canaryStrSplit[0]) == 0 || len(canaryStrSplit[1]) == 0 {
			return nil, nil, fmt.Errorf("invalid canary format: \"%s\"", canaryStr)
		}

		instance := goldpushk.Instance(canaryStrSplit[0])
		service := goldpushk.Service(canaryStrSplit[1])
		instanceServicePair := goldpushk.DeployableUnitID{
			Instance: instance,
			Service:  service,
		}

		// Validate canary subcomponents.
		if !deployableUnitSet.IsKnownInstance(instance) {
			return nil, nil, fmt.Errorf("invalid canary - unknown Gold instance: \"%s\"", canaryStr)
		}
		if !deployableUnitSet.IsKnownService(service) {
			return nil, nil, fmt.Errorf("invalid canary - unknown Gold service: \"%s\"", canaryStr)
		}

		// Canaries should match the services provided with --instances and --services.
		if _, ok := isMarkedForCanarying[instanceServicePair]; !ok {
			return nil, nil, fmt.Errorf("canary does not match any targeted services: \"%s\"", canaryStr)
		}

		// Mark instance/service pair for canarying.
		isMarkedForCanarying[instanceServicePair] = true
	}

	// Sort services to deploy to generate a deterministic output.
	sort.Slice(servicesToDeploy, func(i, j int) bool {
		a := servicesToDeploy[i]
		b := servicesToDeploy[j]
		return a.Instance < b.Instance || (a.Instance == b.Instance && a.Service < b.Service)
	})

	// Build outputs.
	for _, instanceServicePair := range servicesToDeploy {
		deployment, ok := deployableUnitSet.Get(instanceServicePair)
		if !ok {
			sklog.Fatalf("DeployableUnit \"%s\" not found in deployableUnitSet", deployment.CanonicalName())
		}

		if isMarkedForCanarying[instanceServicePair] {
			canariedDeployableUnits = append(canariedDeployableUnits, deployment)
		} else {
			deployableUnits = append(deployableUnits, deployment)
		}
	}

	// If all services to be deployed are marked for canarying, it probably
	// indicates a user error.
	if len(deployableUnits) == 0 {
		return nil, nil, errors.New("all targeted services are marked for canarying")
	}

	return deployableUnits, canariedDeployableUnits, nil
}
<commit_msg>[goldpushk] Don't print out stack trace in the presence of comamnd-line flag errors.<commit_after>// goldpushk pushes Gold services to production. See go/goldpushk.
//
// Sample usage:
//
//   Deployment of a specific service:
//     $ goldpushk --service diffserver --instance chrome-gpu
//     $ goldpushk -s diffserver -i chrome-gpu
//
//   Deployment of a specific service across multiple instances:
//     $ goldpushk --service diffserver --instance chrome-gpu,skia
//     $ goldpushk -s diffserver -i chrome-gpu,skia
//
//   Deployment of all instances of a given service across all Gold instances:
//     $ goldpushk --service diffserver --instance all
//     $ goldpushk -s diffserver -i all
//
//   Deployment of all services corresponding to a specific Gold instance:
//     $ goldpushk --service all --instance chrome-gpu
//     $ goldpushk -s all -i chrome-gpu
//
//   Deployment of all instances of a given service, designating one of them as the canary:
//     $ goldpushk --service diffserver --instance all --canary skia:diffserver
//
//   Print out all Gold instances and services goldpushk is able to manage:
//     $ goldpushk --list

package main

import (
	"context"
	"errors"
	"fmt"
	"os"
	"sort"
	"strings"
	"text/tabwriter"

	"github.com/spf13/cobra"
	"go.skia.org/infra/go/skerr"
	"go.skia.org/infra/go/sklog"
	"go.skia.org/infra/go/util"
	"go.skia.org/infra/golden/cmd/goldpushk/goldpushk"
)

const (
	// Wildcard value for command line arguments.
	all = "all"

	// Environment variable with path to buildbot repository checkout directory.
	skiaInfraRootEnvVar = "SKIA_INFRA_ROOT"

	// Git repositories.
	skiaPublicConfigRepoUrl = "https://skia.googlesource.com/skia-public-config"
	skiaCorpConfigRepoUrl   = "https://skia.googlesource.com/skia-corp-config"
)

var (
	// Required flags.
	flagInstances []string
	flagServices  []string
	flagCanaries  []string

	// Optional flags.
	flagList                       bool
	flagDryRun                     bool
	flagNoCommit                   bool
	flagMinUptimeSeconds           int
	flagUptimePollFrequencySeconds int

	// Flags for debugging.
	flagLogToStdErr bool
	flagTesting     bool
)

func main() {
	// Prevent sklog from using glog.
	sklog.SetLogger(sklog.NewStdErrCloudLogger(sklog.SLogNone))

	rootCmd := &cobra.Command{
		Use:  "goldpushk",
		Long: "goldpushk pushes Gold services to production.",
		PersistentPreRun: func(cmd *cobra.Command, args []string) {
			if flagLogToStdErr {
				sklog.SetLogger(sklog.NewStdErrCloudLogger(sklog.SLogStderr))
			}
		},
		Run: func(cmd *cobra.Command, args []string) {
			run(cmd)
		},
	}

	rootCmd.Flags().SortFlags = false
	rootCmd.Flags().BoolVar(&flagList, "list", false, "List known Gold instances and services (tip: try combining this flag with --testing).")
	rootCmd.Flags().StringSliceVarP(&flagInstances, "instances", "i", []string{}, "[REQUIRED] Comma-delimited list of Gold instances to target (e.g. \"skia,flutter\"), or \""+all+"\" to target all instances.")
	rootCmd.Flags().StringSliceVarP(&flagServices, "services", "s", []string{}, "[REQUIRED] Comma-delimited list of services to target (e.g. \"skiacorrectness,diffserver\"), or \""+all+"\" to target all services.")
	rootCmd.Flags().StringSliceVarP(&flagCanaries, "canaries", "c", []string{}, "Comma-delimited subset of Gold services to use as canaries, written as instance:service pairs (e.g. \"skia:diffserver,flutter:skiacorrectness\")")
	rootCmd.Flags().BoolVar(&flagDryRun, "dryrun", false, "Do everything except applying the new configuration to Kubernetes and committing changes to Git.")
	rootCmd.Flags().BoolVar(&flagNoCommit, "no-commit", false, "Do not commit configuration changes to the skia-public-config or skia-corp-config Git repositories.")
	rootCmd.Flags().IntVar(&flagMinUptimeSeconds, "min-uptime", 30, "Minimum uptime in seconds required for all services before exiting the monitoring step.")
	rootCmd.Flags().IntVar(&flagUptimePollFrequencySeconds, "poll-freq", 3, "How often to poll Kubernetes for service uptimes, in seconds.")
	rootCmd.Flags().BoolVar(&flagLogToStdErr, "logtostderr", false, "Log debug information to stderr. No logs will be produced if this flag is not set.")
	rootCmd.Flags().BoolVar(&flagTesting, "testing", false, "Do not deploy any production services; use testing services instead.")

	// Fail with exit code 1 in the presence of invalid flags.
	if _, err := rootCmd.ExecuteC(); err != nil {
		sklog.Errorf("Failed to execute Cobra command: %s", err)
		os.Exit(1)
	}
}

func run(cmd *cobra.Command) {
	// Get set of deployable units. Used as the source of truth across goldpushk.
	var deployableUnitSet goldpushk.DeployableUnitSet
	if flagTesting {
		deployableUnitSet = goldpushk.TestingDeployableUnits()
	} else {
		deployableUnitSet = goldpushk.ProductionDeployableUnits()
	}

	// If --list is passed, print known services and exit. This takes into account flag --testing.
	if flagList {
		if err := listKnownServices(deployableUnitSet); err != nil {
			sklog.Fatalf("Error while printing list of known services: %s", err)
		}
		return
	}

	// If --list was not provided, validate presence of flags --services and --instances.
	if len(flagInstances) == 0 {
		fmt.Println("Error: flag \"instances\" is required.")
		if err := cmd.Usage(); err != nil {
			sklog.Fatalf("Error while printing usage: %s", err)
		}
		os.Exit(1)
	}
	if len(flagServices) == 0 {
		fmt.Println("Error: flag \"services\" is required.")
		if err := cmd.Usage(); err != nil {
			sklog.Fatalf("Error while printing usage: %s", err)
		}
		os.Exit(1)
	}

	// Parse and validate command line flags.
	deployableUnits, canariedDeployableUnits, err := parseAndValidateFlags(deployableUnitSet, flagInstances, flagServices, flagCanaries)
	if err != nil {
		fmt.Printf("Error: %s.\n", err)
		os.Exit(1)
	}

	// Read environment variables.
	skiaInfraRoot, ok := os.LookupEnv(skiaInfraRootEnvVar)
	if !ok {
		fmt.Printf("Error: environment variable %s not set.", skiaInfraRootEnvVar)
		os.Exit(1)
	}

	// Build goldpushk instance.
	gpk := goldpushk.New(deployableUnits, canariedDeployableUnits, skiaInfraRoot, flagDryRun, flagNoCommit, flagMinUptimeSeconds, flagUptimePollFrequencySeconds, skiaPublicConfigRepoUrl, skiaCorpConfigRepoUrl)

	// Run goldpushk.
	if err = gpk.Run(context.Background()); err != nil {
		fmt.Printf("Error: %s.\n", err)
		os.Exit(1)
	}
}

// listKnownServices prints out a table of known services.
func listKnownServices(deployableUnitSet goldpushk.DeployableUnitSet) error {
	mode := "production"
	if flagTesting {
		mode = "testing"
	}
	fmt.Printf("Known Gold instances and services (%s):\n", mode)

	// Print out table header.
	w := tabwriter.NewWriter(os.Stdout, 10, 0, 2, ' ', 0)
	if _, err := fmt.Fprintln(w, "\nINSTANCE\tSERVICE\tCANONICAL NAME"); err != nil {
		return skerr.Wrap(err)
	}

	// Print out table body.
	for _, instance := range deployableUnitSet.KnownInstances() {
		for _, service := range deployableUnitSet.KnownServices() {
			unit, ok := deployableUnitSet.Get(goldpushk.DeployableUnitID{Instance: instance, Service: service})
			if ok {
				if _, err := fmt.Fprintf(w, "%s\t%s\t%s\n", instance, service, unit.CanonicalName()); err != nil {
					return skerr.Wrap(err)
				}
			}
		}
	}

	// Flush output and return.
	if err := w.Flush(); err != nil {
		return skerr.Wrap(err)
	}
	return nil
}

// containsWildcardValue determines whether or not a flag contains the special
// "all" wildcard value.
func containsWildcardValue(flag []string) bool {
	for _, value := range flag {
		if value == all {
			return true
		}
	}
	return false
}

// parseAndValidateFlags validates the given command line flags, retrieves the
// corresponding DeployableUnits from the given DeployableUnitSet and returns
// them as two separate slices according to whether or not they were marked for
// canarying.
func parseAndValidateFlags(deployableUnitSet goldpushk.DeployableUnitSet, instances, services, canaries []string) (deployableUnits, canariedDeployableUnits []goldpushk.DeployableUnit, err error) {
	// Deduplicate inputs.
	instances = util.SSliceDedup(instances)
	services = util.SSliceDedup(services)
	canaries = util.SSliceDedup(canaries)

	// Determine whether --instances or --services are set to "all".
	allInstances := containsWildcardValue(instances)
	allServices := containsWildcardValue(services)

	// If --instances or --services contain the special "all" value, they should
	// not contain any other values.
	if allInstances && len(instances) != 1 {
		return nil, nil, errors.New("flag --instances should contain either \"all\" or a list of Gold instances, but not both")
	}
	if allServices && len(services) != 1 {
		return nil, nil, errors.New("flag --services should contain either \"all\" or a list of Gold services, but not both")
	}
	if allInstances && allServices {
		return nil, nil, errors.New("cannot set both --instances and --services to \"all\"")
	}

	// Validate instances.
	if !allInstances {
		for _, instanceStr := range instances {
			if !deployableUnitSet.IsKnownInstance(goldpushk.Instance(instanceStr)) {
				return nil, nil, fmt.Errorf("unknown Gold instance: \"%s\"", instanceStr)
			}
		}
	}

	// Validate services.
	if !allServices {
		for _, serviceStr := range services {
			if !deployableUnitSet.IsKnownService(goldpushk.Service(serviceStr)) {
				return nil, nil, fmt.Errorf("unknown Gold service: \"%s\"", serviceStr)
			}
		}
	}

	// This slice will be populated with the subset of the cartesian product of
	// flags --instances and --services that is found in the services map.
	var servicesToDeploy []goldpushk.DeployableUnitID

	// Determines whether or not an instance/service pair should be canaried.
	isMarkedForCanarying := map[goldpushk.DeployableUnitID]bool{}

	// Determine the set of instances over which to iterate to compute the
	// cartesian product of flags --instances and --services.
	var instanceIterationSet []goldpushk.Instance
	if containsWildcardValue(instances) {
		// Handle the "all" value.
		instanceIterationSet = deployableUnitSet.KnownInstances()
	} else {
		for _, instanceStr := range instances {
			instanceIterationSet = append(instanceIterationSet, goldpushk.Instance(instanceStr))
		}
	}

	// Determine the set of services over which to iterate to compute the
	// cartesian product of flags --instances and --services.
	var serviceIterationSet []goldpushk.Service
	if containsWildcardValue(services) {
		// Handle the "all" value.
		serviceIterationSet = deployableUnitSet.KnownServices()
	} else {
		for _, serviceStr := range services {
			serviceIterationSet = append(serviceIterationSet, goldpushk.Service(serviceStr))
		}
	}

	// Iterate over the cartesian product of flags --instances and --services.
	for _, instance := range instanceIterationSet {
		for _, service := range serviceIterationSet {
			id := goldpushk.DeployableUnitID{
				Instance: instance,
				Service:  service,
			}

			// Skip if the current instance/service combination is not found in the services map.
			if _, ok := deployableUnitSet.Get(id); !ok {
				continue
			}

			// Save instance/service pair, which is not marked for canarying by default.
			servicesToDeploy = append(servicesToDeploy, id)
			isMarkedForCanarying[id] = false
		}
	}

	// Fail if --instances and --services didn't match any services in the services map.
	if len(servicesToDeploy) == 0 {
		return nil, nil, errors.New("no known Gold services match the values supplied with --instances and --services")
	}

	// Iterate over the --canaries flag.
	for _, canaryStr := range canaries {
		// Validate format and extract substrings.
		canaryStrSplit := strings.Split(canaryStr, ":")
		if len(canaryStrSplit) != 2 || len(canaryStrSplit[0]) == 0 || len(canaryStrSplit[1]) == 0 {
			return nil, nil, fmt.Errorf("invalid canary format: \"%s\"", canaryStr)
		}

		instance := goldpushk.Instance(canaryStrSplit[0])
		service := goldpushk.Service(canaryStrSplit[1])
		instanceServicePair := goldpushk.DeployableUnitID{
			Instance: instance,
			Service:  service,
		}

		// Validate canary subcomponents.
		if !deployableUnitSet.IsKnownInstance(instance) {
			return nil, nil, fmt.Errorf("invalid canary - unknown Gold instance: \"%s\"", canaryStr)
		}
		if !deployableUnitSet.IsKnownService(service) {
			return nil, nil, fmt.Errorf("invalid canary - unknown Gold service: \"%s\"", canaryStr)
		}

		// Canaries should match the services provided with --instances and --services.
		if _, ok := isMarkedForCanarying[instanceServicePair]; !ok {
			return nil, nil, fmt.Errorf("canary does not match any targeted services: \"%s\"", canaryStr)
		}

		// Mark instance/service pair for canarying.
		isMarkedForCanarying[instanceServicePair] = true
	}

	// Sort services to deploy to generate a deterministic output.
	sort.Slice(servicesToDeploy, func(i, j int) bool {
		a := servicesToDeploy[i]
		b := servicesToDeploy[j]
		return a.Instance < b.Instance || (a.Instance == b.Instance && a.Service < b.Service)
	})

	// Build outputs.
	for _, instanceServicePair := range servicesToDeploy {
		deployment, ok := deployableUnitSet.Get(instanceServicePair)
		if !ok {
			sklog.Fatalf("DeployableUnit \"%s\" not found in deployableUnitSet", deployment.CanonicalName())
		}

		if isMarkedForCanarying[instanceServicePair] {
			canariedDeployableUnits = append(canariedDeployableUnits, deployment)
		} else {
			deployableUnits = append(deployableUnits, deployment)
		}
	}

	// If all services to be deployed are marked for canarying, it probably
	// indicates a user error.
	if len(deployableUnits) == 0 {
		return nil, nil, errors.New("all targeted services are marked for canarying")
	}

	return deployableUnits, canariedDeployableUnits, nil
}
<|endoftext|>


<commit_before>package database

import (
	"crypto/cipher"
	"encoding/binary"
	"fmt"
	"io"
	"os"
	"sync"

	"github.com/golang/protobuf/proto"
	"github.com/golang/snappy"
)

//go:generate protoc --proto_path=. --go_out=. record.proto

type File interface {
	io.ReaderAt
	io.WriterAt
}

type Database struct {
	index, data, rawData File
	indexLen, dataLen    uint64
	mu                   sync.Mutex
	maxSize              uint64
	lru                  *LRU
}

func NewDatabase(index, data *os.File, indexBlock, dataBlock cipher.Block, maxSize, cacheRecords, cacheBytes int) (*Database, error) {
	indexStat, err := index.Stat()
	if err != nil {
		return nil, err
	}
	dataStat, err := data.Stat()
	if err != nil {
		return nil, err
	}
	lru, err := NewLRU(uint64(cacheRecords), uint64(cacheBytes))
	if err != nil {
		return nil, err
	}
	return &Database{
		index:    WrapInCTR(indexBlock, index),
		data:     WrapInCTR(dataBlock, data),
		rawData:  data,
		indexLen: uint64(indexStat.Size()),
		dataLen:  uint64(dataStat.Size()),
		maxSize:  uint64(maxSize),
		lru:      lru,
	}, nil
}

func recordSize(record *Record) uint64 {
	return uint64(len(record.Filename) + len(record.Content))
}

func (d *Database) Lookup(key uint64) (*Record, error) {
	d.mu.Lock()
	r := d.lru.Get(key)
	indexLen := d.indexLen
	dataLen := d.dataLen
	d.mu.Unlock()
	if r != nil {
		return r, nil
	}
	if key >= indexLen/8 {
		return nil, fmt.Errorf("the record does not exist")
	}
	indexBuffer := make([]byte, 8+8)
	indexBuffer2 := indexBuffer
	if key*8 == indexLen-8 {
		// Last element.
		indexBuffer2 = indexBuffer2[0:8]
	}
	if _, err := d.index.ReadAt(indexBuffer2, int64(key*8)); err != nil {
		return nil, err
	}
	dataBegin := binary.LittleEndian.Uint64(indexBuffer[0:8])
	dataEnd := binary.LittleEndian.Uint64(indexBuffer[8:16])
	if key*8 == indexLen-8 {
		// Last element.
		dataEnd = dataLen
	}
	if dataBegin == 0 && key != 0 {
		return nil, fmt.Errorf("the record does not exist in the middle")
	}
	size := dataEnd - dataBegin
	if size > d.maxSize {
		return nil, fmt.Errorf("the record seems to be too long")
	}
	dataBuffer := make([]byte, dataEnd-dataBegin)
	if _, err := d.data.ReadAt(dataBuffer, int64(dataBegin)); err != nil {
		return nil, err
	}
	dataBuffer, err := snappy.Decode(nil, dataBuffer)
	if err != nil {
		return nil, err
	}
	var record Record
	if err := proto.Unmarshal(dataBuffer, &record); err != nil {
		return nil, err
	}
	if record.SelfBurning {
		// Wipe the record.
		dummy := make([]byte, dataEnd-dataBegin)
		if _, err := d.rawData.WriteAt(dummy, int64(dataBegin)); err != nil {
			return nil, err
		}
	} else {
		d.mu.Lock()
		d.lru.Set(key, &record, recordSize(&record))
		d.mu.Unlock()
	}
	return &record, nil
}

func (d *Database) Add(record *Record) (uint64, error) {
	data, err := proto.Marshal(record)
	if err != nil {
		return 0, err
	}
	data = snappy.Encode(nil, data)
	if uint64(len(data)) > d.maxSize {
		return 0, fmt.Errorf("the record seems to be too long")
	}
	d.mu.Lock()
	indexLen := d.indexLen
	d.indexLen += 8
	key := indexLen / 8
	dataLen := d.dataLen
	d.dataLen += uint64(len(data))
	if !record.SelfBurning {
		d.lru.Set(key, record, recordSize(record))
	}
	d.mu.Unlock()
	if _, err := d.data.WriteAt(data, int64(dataLen)); err != nil {
		return 0, err
	}
	indexBuffer := make([]byte, 8)
	binary.LittleEndian.PutUint64(indexBuffer, dataLen)
	if _, err := d.index.WriteAt(indexBuffer, int64(indexLen)); err != nil {
		return 0, err
	}
	return key, nil
}

func (d *Database) RecordsCount() int64 {
	d.mu.Lock()
	indexLen := d.indexLen
	d.mu.Unlock()
	return int64(indexLen) / 8
}
<commit_msg>gopasta: add go:generate for lru.go<commit_after>package database

import (
	"crypto/cipher"
	"encoding/binary"
	"fmt"
	"io"
	"os"
	"sync"

	"github.com/golang/protobuf/proto"
	"github.com/golang/snappy"
)

//go:generate protoc --proto_path=. --go_out=. record.proto

//go:generate go get gitlab.com/starius/lru-gen
//go:generate lru-gen -package database -key "uint64" -value "*Record" -output lru.go

type File interface {
	io.ReaderAt
	io.WriterAt
}

type Database struct {
	index, data, rawData File
	indexLen, dataLen    uint64
	mu                   sync.Mutex
	maxSize              uint64
	lru                  *LRU
}

func NewDatabase(index, data *os.File, indexBlock, dataBlock cipher.Block, maxSize, cacheRecords, cacheBytes int) (*Database, error) {
	indexStat, err := index.Stat()
	if err != nil {
		return nil, err
	}
	dataStat, err := data.Stat()
	if err != nil {
		return nil, err
	}
	lru, err := NewLRU(uint64(cacheRecords), uint64(cacheBytes))
	if err != nil {
		return nil, err
	}
	return &Database{
		index:    WrapInCTR(indexBlock, index),
		data:     WrapInCTR(dataBlock, data),
		rawData:  data,
		indexLen: uint64(indexStat.Size()),
		dataLen:  uint64(dataStat.Size()),
		maxSize:  uint64(maxSize),
		lru:      lru,
	}, nil
}

func recordSize(record *Record) uint64 {
	return uint64(len(record.Filename) + len(record.Content))
}

func (d *Database) Lookup(key uint64) (*Record, error) {
	d.mu.Lock()
	r := d.lru.Get(key)
	indexLen := d.indexLen
	dataLen := d.dataLen
	d.mu.Unlock()
	if r != nil {
		return r, nil
	}
	if key >= indexLen/8 {
		return nil, fmt.Errorf("the record does not exist")
	}
	indexBuffer := make([]byte, 8+8)
	indexBuffer2 := indexBuffer
	if key*8 == indexLen-8 {
		// Last element.
		indexBuffer2 = indexBuffer2[0:8]
	}
	if _, err := d.index.ReadAt(indexBuffer2, int64(key*8)); err != nil {
		return nil, err
	}
	dataBegin := binary.LittleEndian.Uint64(indexBuffer[0:8])
	dataEnd := binary.LittleEndian.Uint64(indexBuffer[8:16])
	if key*8 == indexLen-8 {
		// Last element.
		dataEnd = dataLen
	}
	if dataBegin == 0 && key != 0 {
		return nil, fmt.Errorf("the record does not exist in the middle")
	}
	size := dataEnd - dataBegin
	if size > d.maxSize {
		return nil, fmt.Errorf("the record seems to be too long")
	}
	dataBuffer := make([]byte, dataEnd-dataBegin)
	if _, err := d.data.ReadAt(dataBuffer, int64(dataBegin)); err != nil {
		return nil, err
	}
	dataBuffer, err := snappy.Decode(nil, dataBuffer)
	if err != nil {
		return nil, err
	}
	var record Record
	if err := proto.Unmarshal(dataBuffer, &record); err != nil {
		return nil, err
	}
	if record.SelfBurning {
		// Wipe the record.
		dummy := make([]byte, dataEnd-dataBegin)
		if _, err := d.rawData.WriteAt(dummy, int64(dataBegin)); err != nil {
			return nil, err
		}
	} else {
		d.mu.Lock()
		d.lru.Set(key, &record, recordSize(&record))
		d.mu.Unlock()
	}
	return &record, nil
}

func (d *Database) Add(record *Record) (uint64, error) {
	data, err := proto.Marshal(record)
	if err != nil {
		return 0, err
	}
	data = snappy.Encode(nil, data)
	if uint64(len(data)) > d.maxSize {
		return 0, fmt.Errorf("the record seems to be too long")
	}
	d.mu.Lock()
	indexLen := d.indexLen
	d.indexLen += 8
	key := indexLen / 8
	dataLen := d.dataLen
	d.dataLen += uint64(len(data))
	if !record.SelfBurning {
		d.lru.Set(key, record, recordSize(record))
	}
	d.mu.Unlock()
	if _, err := d.data.WriteAt(data, int64(dataLen)); err != nil {
		return 0, err
	}
	indexBuffer := make([]byte, 8)
	binary.LittleEndian.PutUint64(indexBuffer, dataLen)
	if _, err := d.index.WriteAt(indexBuffer, int64(indexLen)); err != nil {
		return 0, err
	}
	return key, nil
}

func (d *Database) RecordsCount() int64 {
	d.mu.Lock()
	indexLen := d.indexLen
	d.mu.Unlock()
	return int64(indexLen) / 8
}
<|endoftext|>


<commit_before>// Package netlink provides a simple library for netlink. Netlink is
// the interface a user-space program in linux uses to communicate with
// the kernel. It can be used to add and remove interfaces, set up ip
// addresses and routes, and confiugre ipsec. Netlink communication
// requires elevated privileges, so in most cases this code needs to
// be run as root. In addition to dealing with netlink primitives, the
// library attempts to provide an high-level interface that is loosly
// modeled on the iproute2 command line interface.
package netlink

import (
	"net"
	"syscall"
)

const (
	// Family type definitions
	FAMILY_ALL = syscall.AF_UNSPEC
	FAMILY_V4  = syscall.AF_INET
	FAMILY_V6  = syscall.AF_INET6
)

// GetIPFamily returns the family type of a net.IP.
func GetIPFamily(ip net.IP) int {
	if len(ip) <= net.IPv4len {
		return FAMILY_V4
	}
	if ip.To4() != nil {
		return FAMILY_V4
	}
	return FAMILY_V6
}

// ParseIPNet parses a string in ip/net format and returns a net.IPNet.
// This is valuable because addresses in netlink are often IPNets and
// ParseCIDR returns an IPNet with the IP part set to the base IP of the
// range.
func ParseIPNet(s string) (net.IPNet, error) {
	ip, ipNet, err := net.ParseCIDR(s)
	if err != nil {
		return net.IPNet{}, err
	}
	return net.IPNet{ip, ipNet.Mask}, nil
}
<commit_msg>Add helper method for creating an IPNet from an IP<commit_after>// Package netlink provides a simple library for netlink. Netlink is
// the interface a user-space program in linux uses to communicate with
// the kernel. It can be used to add and remove interfaces, set up ip
// addresses and routes, and confiugre ipsec. Netlink communication
// requires elevated privileges, so in most cases this code needs to
// be run as root. In addition to dealing with netlink primitives, the
// library attempts to provide an high-level interface that is loosly
// modeled on the iproute2 command line interface.
package netlink

import (
	"net"
	"syscall"
)

const (
	// Family type definitions
	FAMILY_ALL = syscall.AF_UNSPEC
	FAMILY_V4  = syscall.AF_INET
	FAMILY_V6  = syscall.AF_INET6
)

// GetIPFamily returns the family type of a net.IP.
func GetIPFamily(ip net.IP) int {
	if len(ip) <= net.IPv4len {
		return FAMILY_V4
	}
	if ip.To4() != nil {
		return FAMILY_V4
	}
	return FAMILY_V6
}

// ParseIPNet parses a string in ip/net format and returns a net.IPNet.
// This is valuable because addresses in netlink are often IPNets and
// ParseCIDR returns an IPNet with the IP part set to the base IP of the
// range.
func ParseIPNet(s string) (*net.IPNet, error) {
	ip, ipNet, err := net.ParseCIDR(s)
	if err != nil {
		return nil, err
	}
	return &net.IPNet{ip, ipNet.Mask}, nil
}

// NewIPNet generates an IPNet from an ip address using a netmask of 32.
func NewIPNet(ip net.IP) (*net.IPNet) {
	return &net.IPNet{ip, net.CIDRMask(32, 32)}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"errors"
	"io"
	"net/http"
)

var getworkurl = "http://localhost:9980/miner/headerforwork"
var submitblockurl = "http://localhost:9980/miner/submitheader"

func getHeaderForWork() (target, header []byte, err error) {
	client := &http.Client{}

	req, err := http.NewRequest("GET", getworkurl, nil)
	if err != nil {
		return
	}

	req.Header.Add("User-Agent", "Sia-Agent")
	resp, err := client.Do(req)

	if err != nil {
		return
	}
	defer resp.Body.Close()
	buf := make([]byte, 113)
	n, err := resp.Body.Read(buf)
	if err != nil && err != io.EOF {
		return
	}
	if n < 112 {
		err = errors.New("Invalid response")
	} else {
		err = nil
	}

	target = buf[:32]
	header = buf[32:112]

	return
}

func submitHeader(header []byte) (err error) {
	req, err := http.NewRequest("POST", submitblockurl, bytes.NewReader(header))
	if err != nil {
		return
	}

	req.Header.Add("User-Agent", "Sia-Agent")

	client := &http.Client{}
	_, err = client.Do(req)

	return
}
<commit_msg>Provide some better error reporting since an invalid response can be caused by an unitialized or locked wallet<commit_after>package main

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"net/http"
)

var getworkurl = "http://localhost:9980/miner/headerforwork"
var submitblockurl = "http://localhost:9980/miner/submitheader"

func getHeaderForWork() (target, header []byte, err error) {
	client := &http.Client{}

	req, err := http.NewRequest("GET", getworkurl, nil)
	if err != nil {
		return
	}

	req.Header.Add("User-Agent", "Sia-Agent")
	resp, err := client.Do(req)
	if err != nil {
		return
	}
	defer resp.Body.Close()

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

	if len(buf) < 112 {
		err = fmt.Errorf("Invalid response, only received %d bytes, is your wallet initialized and unlocked?", len(buf))
		return
	}

	target = buf[:32]
	header = buf[32:112]

	return
}

func submitHeader(header []byte) (err error) {
	req, err := http.NewRequest("POST", submitblockurl, bytes.NewReader(header))
	if err != nil {
		return
	}

	req.Header.Add("User-Agent", "Sia-Agent")

	client := &http.Client{}
	_, err = client.Do(req)

	return
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/gob"
	"errors"
	"net"
	"sync"
)

type MsgType int32

const (
	PeerListRequest MsgType = iota
	PeerBroadcast   MsgType = iota

	BlockChainRequest MsgType = iota
	BlockBroadcast    MsgType = iota

	TransactionRequest   MsgType = iota
	TransactionBroadcast MsgType = iota
)

type PeerConn struct {
	base    net.Conn
	encoder *gob.Encoder
	decoder *gob.Decoder
	txLock  sync.Mutex
}

type PeerEvent struct {
	addr  string
	value interface{}
}

type PeerNetwork struct {
	peers    map[string]*PeerConn
	server   net.Listener
	events   chan PeerEvent
	closing  bool
	peerLock sync.RWMutex
}

func NewPeerNetwork(startPeer string) (network *PeerNetwork, err error) {
	tmpAddrs := make([]string, 0)

	if startPeer != "" {
		conn, err := net.Dial("udp", startPeer)
		if err != nil {
			return nil, err
		}
		defer conn.Close()

		encoder := gob.NewEncoder(conn)
		decoder := gob.NewDecoder(conn)

		err = encoder.Encode(PeerListRequest)
		if err != nil {
			return nil, err
		}

		err = decoder.Decode(&tmpAddrs)
		if err != nil {
			return nil, err
		}

		if len(tmpAddrs) == 0 {
			return nil, errors.New("Initial peer returned empty peer list")
		}
	}

	network = &PeerNetwork{
		peers:  make(map[string]*PeerConn, len(tmpAddrs)),
		events: make(chan PeerEvent),
	}
	network.server, err = net.Listen("udp", ":0")
	if err != nil {
		return nil, err
	}

	for _, addr := range tmpAddrs {
		conn, err := net.Dial("udp", addr)
		if err != nil {
			continue
		}

		encoder := gob.NewEncoder(conn)

		err = encoder.Encode(PeerBroadcast)
		if err != nil {
			conn.Close()
			continue
		}

		err = encoder.Encode(network.server.Addr().String())
		if err != nil {
			conn.Close()
			continue
		}

		decoder := gob.NewDecoder(conn)

		network.peers[addr] = &PeerConn{base: conn, encoder: encoder, decoder: decoder}
		go network.ReceiveFromConn(addr)
	}

	go network.AcceptNewConns()
	go network.HandleEvents()

	return network, nil
}

func (network *PeerNetwork) AcceptNewConns() {
	for {
		conn, err := network.server.Accept()

		if err != nil {
			network.events <- PeerEvent{"", err}
			return
		}

		encoder := gob.NewEncoder(conn)
		decoder := gob.NewDecoder(conn)

		var nextType MsgType
		err = decoder.Decode(&nextType)
		if err != nil {
			conn.Close()
			continue
		}

		switch nextType {
		case PeerListRequest:
			peerList := network.PeerAddrList()
			err = encoder.Encode(peerList)
			conn.Close()
		case PeerBroadcast:
			var addr string
			err = decoder.Decode(&addr)
			network.peerLock.Lock()
			if err != nil || network.peers[addr] != nil {
				network.peerLock.Unlock()
				conn.Close()
				continue
			}
			network.peers[addr] = &PeerConn{base: conn, encoder: encoder, decoder: decoder}
			go network.ReceiveFromConn(addr)
			network.peerLock.Unlock()
		default:
			conn.Close()
		}
	}
}

func (network *PeerNetwork) ReceiveFromConn(addr string) {
	peer := network.peers[addr]

	var err error
	var nextType MsgType

	for {
		err = peer.decoder.Decode(&nextType)
		if err != nil {
			network.events <- PeerEvent{addr, err}
			return
		}

		switch nextType {
		default:
			network.events <- PeerEvent{addr, errors.New("Unknown message type received")}
			return
		}
	}
}

func (network *PeerNetwork) HandleEvents() {
	for event := range network.events {
		switch val := event.value.(type) {
		case error:
			if event.addr == "" {
				if network.closing {
					if len(network.peers) == 0 {
						close(network.events)
						return
					}
				} else {
					panic(val)
				}
			} else {
				network.peerLock.Lock()
				delete(network.peers, event.addr)
				network.peerLock.Unlock()
				if len(network.peers) == 0 {
					if network.closing {
						close(network.events)
						return
					} else {
						panic(val)
					}
				}
			}
		}
	}
}

func (network *PeerNetwork) Close() {
	network.closing = true
	network.server.Close()
	for _, peer := range network.peers {
		peer.base.Close()
	}
}

func (network *PeerNetwork) PeerAddrList() []string {
	network.peerLock.RLock()
	defer network.peerLock.RUnlock()

	list := make([]string, 0, len(network.peers))
	for addr, _ := range network.peers {
		list = append(list, addr)
	}
	return list
}
<commit_msg>checkpoint wip<commit_after>package main

import (
	"encoding/gob"
	"errors"
	"net"
	"sync"
)

type MsgType int32

const (
	PeerListRequest MsgType = iota
	PeerBroadcast   MsgType = iota

	BlockChainRequest MsgType = iota
	BlockBroadcast    MsgType = iota

	TransactionRequest   MsgType = iota
	TransactionBroadcast MsgType = iota
)

type PeerConn struct {
	base    net.Conn
	encoder *gob.Encoder
	decoder *gob.Decoder
}

type PeerEvent struct {
	addr  string
	value interface{}
}

type PeerNetwork struct {
	peers    map[string]*PeerConn
	server   net.Listener
	events   chan PeerEvent
	closing  bool
	peerLock sync.RWMutex
}

func NewPeerNetwork(startPeer string) (network *PeerNetwork, err error) {
	tmpAddrs := make([]string, 0)

	if startPeer != "" {
		conn, err := net.Dial("udp", startPeer)
		if err != nil {
			return nil, err
		}
		defer conn.Close()

		encoder := gob.NewEncoder(conn)
		decoder := gob.NewDecoder(conn)

		err = encoder.Encode(PeerListRequest)
		if err != nil {
			return nil, err
		}

		err = decoder.Decode(&tmpAddrs)
		if err != nil {
			return nil, err
		}

		if len(tmpAddrs) == 0 {
			return nil, errors.New("Initial peer returned empty peer list")
		}
	}

	network = &PeerNetwork{
		peers:  make(map[string]*PeerConn, len(tmpAddrs)),
		events: make(chan PeerEvent),
	}
	network.server, err = net.Listen("udp", ":0")
	if err != nil {
		return nil, err
	}

	for _, addr := range tmpAddrs {
		conn, err := net.Dial("udp", addr)
		if err != nil {
			continue
		}

		encoder := gob.NewEncoder(conn)

		err = encoder.Encode(PeerBroadcast)
		if err != nil {
			conn.Close()
			continue
		}

		err = encoder.Encode(network.server.Addr().String())
		if err != nil {
			conn.Close()
			continue
		}

		decoder := gob.NewDecoder(conn)

		network.peers[addr] = &PeerConn{base: conn, encoder: encoder, decoder: decoder}
		go network.ReceiveFromConn(addr)
	}

	go network.AcceptNewConns()
	go network.HandleEvents()

	return network, nil
}

func (network *PeerNetwork) AcceptNewConns() {
	for {
		conn, err := network.server.Accept()

		if err != nil {
			network.events <- PeerEvent{"", err}
			return
		}

		encoder := gob.NewEncoder(conn)
		decoder := gob.NewDecoder(conn)

		var nextType MsgType
		err = decoder.Decode(&nextType)
		if err != nil {
			conn.Close()
			continue
		}

		switch nextType {
		case PeerListRequest:
			peerList := network.PeerAddrList()
			err = encoder.Encode(peerList)
			conn.Close()
		case PeerBroadcast:
			var addr string
			err = decoder.Decode(&addr)
			network.peerLock.Lock()
			if err != nil || network.peers[addr] != nil {
				network.peerLock.Unlock()
				conn.Close()
				continue
			}
			network.peers[addr] = &PeerConn{base: conn, encoder: encoder, decoder: decoder}
			go network.ReceiveFromConn(addr)
			network.peerLock.Unlock()
		default:
			conn.Close()
		}
	}
}

func (network *PeerNetwork) ReceiveFromConn(addr string) {
	peer := network.peers[addr]

	var err error
	var nextType MsgType

	for {
		err = peer.decoder.Decode(&nextType)
		if err != nil {
			network.events <- PeerEvent{addr, err}
			return
		}

		switch nextType {
		default:
			network.events <- PeerEvent{addr, errors.New("Unknown message type received")}
			return
		}
	}
}

func (network *PeerNetwork) HandleEvents() {
	for event := range network.events {
		switch val := event.value.(type) {
		case error:
			if event.addr == "" {
				if network.closing {
					if len(network.peers) == 0 {
						close(network.events)
						return
					}
				} else {
					panic(val)
				}
			} else {
				network.peerLock.Lock()
				delete(network.peers, event.addr)
				network.peerLock.Unlock()
				if len(network.peers) == 0 {
					if network.closing {
						close(network.events)
						return
					} else {
						panic(val)
					}
				}
			}
		}
	}
}

func (network *PeerNetwork) Close() {
	network.closing = true
	network.server.Close()
	for _, peer := range network.peers {
		peer.base.Close()
	}
}

func (network *PeerNetwork) PeerAddrList() []string {
	network.peerLock.RLock()
	defer network.peerLock.RUnlock()

	list := make([]string, 0, len(network.peers))
	for addr, _ := range network.peers {
		list = append(list, addr)
	}
	return list
}
<|endoftext|>


<commit_before>package vat

import (
	"bytes"
	"encoding/xml"
	"errors"
	"io/ioutil"
	"net/http"
	"regexp"
	"strings"
	"time"
)

type viesResponse struct {
	CountryCode string
	VATNumber   string
	RequestDate time.Time
	Valid       bool
	Name        string
	Address     string
}

const serviceURL = "http://ec.europa.eu/taxation_customs/vies/services/checkVatService"

// ErrInvalidVATNumber will be returned when an invalid VAT number is passed to a function that validates existence.
var ErrInvalidVATNumber = errors.New("VAT number is invalid")

// ValidateNumber validates a VAT number by both format and existence.
// The existence check uses the VIES VAT validation SOAP API and will only run when format validation passes.
func ValidateNumber(n string) (bool, error) {
	format, err := ValidateNumberFormat(n)
	existence := false

	if format {
		existence, err = ValidateNumberExistence(n)
	}

	return (format && existence), err
}

// ValidateNumberFormat validates a VAT number by its format.
func ValidateNumberFormat(n string) (bool, error) {
	patterns := map[string]string{
		"AT": "U[A-Z\\d]{8}",
		"BE": "(0\\d{9}|\\d{10})",
		"BG": "\\d{9,10}",
		"CY": "\\d{8}[A-Z]",
		"CZ": "\\d{8,10}",
		"DE": "\\d{9}",
		"DK": "(\\d{2} ?){3}\\d{2}",
		"EE": "\\d{9}",
		"EL": "\\d{9}",
		"ES": "[A-Z]\\d{7}[A-Z]|\\d{8}[A-Z]|[A-Z]\\d{8}",
		"FI": "\\d{8}",
		"FR": "([A-Z]{2}|\\d{2})\\d{9}",
		"GB": "\\d{9}|\\d{12}|(GD|HA)\\d{3}",
		"HR": "\\d{11}",
		"HU": "\\d{8}",
		"IE": "[A-Z\\d]{8}|[A-Z\\d]{9}",
		"IT": "\\d{11}",
		"LT": "(\\d{9}|\\d{12})",
		"LU": "\\d{8}",
		"LV": "\\d{11}",
		"MT": "\\d{8}",
		"NL": "\\d{9}B\\d{2}",
		"PL": "\\d{10}",
		"PT": "\\d{9}",
		"RO": "\\d{2,10}",
		"SE": "\\d{12}",
		"SI": "\\d{8}",
		"SK": "\\d{10}",
	}

	if len(n) < 3 {
		return false, nil
	}

	n = strings.ToUpper(n)
	pattern, ok := patterns[n[0:2]]
	if !ok {
		return false, nil
	}

	matched, err := regexp.MatchString(pattern, n[2:])
	return matched, err
}

// ValidateNumberExistence validates a VAT number by its existence using the VIES VAT API (using SOAP)
func ValidateNumberExistence(n string) (bool, error) {
	r, err := checkVAT(n)
	return r.Valid, err
}

// checkVAT returns *ViesResponse for a VAT number
func checkVAT(vatNumber string) (*viesResponse, error) {
	if len(vatNumber) < 3 {
		return nil, ErrInvalidVATNumber
	}

	e := getEnvelope(vatNumber)
	eb := bytes.NewBufferString(e)
	client := http.Client{
		Timeout: (time.Duration(ServiceTimeout) * time.Second),
	}
	res, err := client.Post(serviceURL, "text/xml;charset=UTF-8", eb)
	if err != nil {
		return nil, ErrServiceUnavailable
	}
	defer res.Body.Close()

	xmlRes, err := ioutil.ReadAll(res.Body)

	// check if response contains "INVALID_INPUT" string
	if bytes.Contains(xmlRes, []byte("INVALID_INPUT")) {
		return nil, ErrInvalidVATNumber
	}

	var rd struct {
		XMLName xml.Name `xml:"Envelope"`
		Soap    struct {
			XMLName xml.Name `xml:"Body"`
			Soap    struct {
				XMLName     xml.Name `xml:"checkVatResponse"`
				CountryCode string   `xml:"countryCode"`
				VATNumber   string   `xml:"vatNumber"`
				RequestDate string   `xml:"requestDate"` // 2015-03-06+01:00
				Valid       bool     `xml:"valid"`
				Name        string   `xml:"name"`
				Address     string   `xml:"address"`
			}
		}
	}
	if err = xml.Unmarshal(xmlRes, &rd); err != nil {
		return nil, err
	}

	pDate, err := time.Parse("2006-01-02-07:00", rd.Soap.Soap.RequestDate)
	if err != nil {
		return nil, err
	}

	r := &viesResponse{
		CountryCode: rd.Soap.Soap.CountryCode,
		VATNumber:   rd.Soap.Soap.VATNumber,
		RequestDate: pDate,
		Valid:       rd.Soap.Soap.Valid,
		Name:        rd.Soap.Soap.Name,
		Address:     rd.Soap.Soap.Address,
	}

	return r, nil
}

// getEnvelope parses envelope template
func getEnvelope(n string) string {
	n = strings.ToUpper(n)
	countryCode := n[0:2]
	vatNumber := n[2:]
	const envelopeTemplate = `
	<soapenv:Envelope xmlns:soapenv="http://schemas.xmlsoap.org/soap/envelope/">
	<soapenv:Header/>
	<soapenv:Body>
	  <checkVat xmlns="urn:ec.europa.eu:taxud:vies:services:checkVat:types">
	    <countryCode>{{.countryCode}}</countryCode>
	    <vatNumber>{{.vatNumber}}</vatNumber>
	  </checkVat>
	</soapenv:Body>
	</soapenv:Envelope>
	`

	e := envelopeTemplate
	e = strings.Replace(e, "{{.countryCode}}", countryCode, 1)
	e = strings.Replace(e, "{{.vatNumber}}", vatNumber, 1)
	return e
}
<commit_msg>check error returned by ioutil.ReadAll<commit_after>package vat

import (
	"bytes"
	"encoding/xml"
	"errors"
	"io/ioutil"
	"net/http"
	"regexp"
	"strings"
	"time"
)

type viesResponse struct {
	CountryCode string
	VATNumber   string
	RequestDate time.Time
	Valid       bool
	Name        string
	Address     string
}

const serviceURL = "http://ec.europa.eu/taxation_customs/vies/services/checkVatService"

// ErrInvalidVATNumber will be returned when an invalid VAT number is passed to a function that validates existence.
var ErrInvalidVATNumber = errors.New("vat: vat number is invalid")

// ValidateNumber validates a VAT number by both format and existence.
// The existence check uses the VIES VAT validation SOAP API and will only run when format validation passes.
func ValidateNumber(n string) (bool, error) {
	format, err := ValidateNumberFormat(n)
	existence := false

	if format {
		existence, err = ValidateNumberExistence(n)
	}

	return (format && existence), err
}

// ValidateNumberFormat validates a VAT number by its format.
func ValidateNumberFormat(n string) (bool, error) {
	patterns := map[string]string{
		"AT": "U[A-Z\\d]{8}",
		"BE": "(0\\d{9}|\\d{10})",
		"BG": "\\d{9,10}",
		"CY": "\\d{8}[A-Z]",
		"CZ": "\\d{8,10}",
		"DE": "\\d{9}",
		"DK": "(\\d{2} ?){3}\\d{2}",
		"EE": "\\d{9}",
		"EL": "\\d{9}",
		"ES": "[A-Z]\\d{7}[A-Z]|\\d{8}[A-Z]|[A-Z]\\d{8}",
		"FI": "\\d{8}",
		"FR": "([A-Z]{2}|\\d{2})\\d{9}",
		"GB": "\\d{9}|\\d{12}|(GD|HA)\\d{3}",
		"HR": "\\d{11}",
		"HU": "\\d{8}",
		"IE": "[A-Z\\d]{8}|[A-Z\\d]{9}",
		"IT": "\\d{11}",
		"LT": "(\\d{9}|\\d{12})",
		"LU": "\\d{8}",
		"LV": "\\d{11}",
		"MT": "\\d{8}",
		"NL": "\\d{9}B\\d{2}",
		"PL": "\\d{10}",
		"PT": "\\d{9}",
		"RO": "\\d{2,10}",
		"SE": "\\d{12}",
		"SI": "\\d{8}",
		"SK": "\\d{10}",
	}

	if len(n) < 3 {
		return false, nil
	}

	n = strings.ToUpper(n)
	pattern, ok := patterns[n[0:2]]
	if !ok {
		return false, nil
	}

	matched, err := regexp.MatchString(pattern, n[2:])
	return matched, err
}

// ValidateNumberExistence validates a VAT number by its existence using the VIES VAT API (using SOAP)
func ValidateNumberExistence(n string) (bool, error) {
	r, err := checkVAT(n)
	return r.Valid, err
}

// checkVAT returns *ViesResponse for a VAT number
func checkVAT(vatNumber string) (*viesResponse, error) {
	if len(vatNumber) < 3 {
		return nil, ErrInvalidVATNumber
	}

	e := getEnvelope(vatNumber)
	eb := bytes.NewBufferString(e)
	client := http.Client{
		Timeout: (time.Duration(ServiceTimeout) * time.Second),
	}
	res, err := client.Post(serviceURL, "text/xml;charset=UTF-8", eb)
	if err != nil {
		return nil, ErrServiceUnavailable
	}
	defer res.Body.Close()

	xmlRes, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return nil, err
	}

	// check if response contains "INVALID_INPUT" string
	if bytes.Contains(xmlRes, []byte("INVALID_INPUT")) {
		return nil, ErrInvalidVATNumber
	}

	var rd struct {
		XMLName xml.Name `xml:"Envelope"`
		Soap    struct {
			XMLName xml.Name `xml:"Body"`
			Soap    struct {
				XMLName     xml.Name `xml:"checkVatResponse"`
				CountryCode string   `xml:"countryCode"`
				VATNumber   string   `xml:"vatNumber"`
				RequestDate string   `xml:"requestDate"` // 2015-03-06+01:00
				Valid       bool     `xml:"valid"`
				Name        string   `xml:"name"`
				Address     string   `xml:"address"`
			}
		}
	}
	if err = xml.Unmarshal(xmlRes, &rd); err != nil {
		return nil, err
	}

	pDate, err := time.Parse("2006-01-02-07:00", rd.Soap.Soap.RequestDate)
	if err != nil {
		return nil, err
	}

	r := &viesResponse{
		CountryCode: rd.Soap.Soap.CountryCode,
		VATNumber:   rd.Soap.Soap.VATNumber,
		RequestDate: pDate,
		Valid:       rd.Soap.Soap.Valid,
		Name:        rd.Soap.Soap.Name,
		Address:     rd.Soap.Soap.Address,
	}

	return r, nil
}

// getEnvelope parses envelope template
func getEnvelope(n string) string {
	n = strings.ToUpper(n)
	countryCode := n[0:2]
	vatNumber := n[2:]
	const envelopeTemplate = `
	<soapenv:Envelope xmlns:soapenv="http://schemas.xmlsoap.org/soap/envelope/">
	<soapenv:Header/>
	<soapenv:Body>
	  <checkVat xmlns="urn:ec.europa.eu:taxud:vies:services:checkVat:types">
	    <countryCode>{{.countryCode}}</countryCode>
	    <vatNumber>{{.vatNumber}}</vatNumber>
	  </checkVat>
	</soapenv:Body>
	</soapenv:Envelope>
	`

	e := envelopeTemplate
	e = strings.Replace(e, "{{.countryCode}}", countryCode, 1)
	e = strings.Replace(e, "{{.vatNumber}}", vatNumber, 1)
	return e
}
<|endoftext|>


<commit_before>package ofutils

import (
	"bufio"
	"crypto/md5"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"math/rand"
	"net/smtp"
	"os"
	"reflect"
	"sort"
	"strconv"
	"strings"
	"time"

	"golang.org/x/text/encoding/simplifiedchinese"
	"golang.org/x/text/transform"

	"github.com/astaxie/beego/orm"
	"github.com/axgle/mahonia"
)

func init() {
	rand.Seed(time.Now().UnixNano())
}

var letterRunes = []rune("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ")

func RandString(n int) string {
	b := make([]rune, n)
	for i := range b {
		b[i] = letterRunes[rand.Intn(len(letterRunes))]
	}
	return string(b)
}
func ChangeJson(data *[]byte, key string, value interface{}) error {
	var m map[string]interface{}
	err := json.Unmarshal(*data, &m)
	if err != nil {
		return err
	}
	m[key] = value
	*data, err = json.Marshal(m)
	return err
}
func MD5(text string) string {
	hasher := md5.New()
	hasher.Write([]byte(text))
	return hex.EncodeToString(hasher.Sum(nil))
}

func ToString(str interface{}) string {
	if str == nil {
		return ""
	}
	return fmt.Sprintf("%v", str)
}

func ToInt(val interface{}) int {
	if val == nil {
		return 0
	}
	s, ok := val.(string)
	if ok {
		i, err := strconv.Atoi(s)
		if err != nil {
			return 0
		}
		return i
	}
	i, ok := val.(int)
	if ok {
		return i
	} else {
		return 0
	}
}

func GetWeekFirstDay() string {
	week := time.Now().Weekday().String()
	var day time.Duration
	switch week {
	case "Sunday":
		day = 6
		break
	case "Monday":
		day = 0
		break
	case "Tuesday":
		day = 1
		break
	case "Wednesday":
		day = 2
		break
	case "Thursday":
		day = 3
		break
	case "Friday":
		day = 4
		break
	case "Saturday":
		day = 5
		break
	}
	date := time.Now().Add(-day * 24 * time.Hour)
	return date.Format("2006-01-02")
}

type ByKey struct {
	Key  string
	List []orm.Params
}

func (a ByKey) Len() int {
	return len(a.List)
}
func (a ByKey) Swap(i, j int) {
	a.List[i], a.List[j] = a.List[j], a.List[i]
}
func (a ByKey) Less(i, j int) bool {
	return ToInt(a.List[i][a.Key]) > ToInt(a.List[j][a.Key])
}

func Sort(list []orm.Params, key string) []orm.Params {
	byKey := ByKey{List: list, Key: key}
	sort.Sort(byKey)
	return byKey.List
}

func SubString(str string, begin, length int) (substr string) {
	// 将字符串的转换成[]rune
	rs := []rune(str)
	lth := len(rs)

	// 简单的越界判断
	if begin < 0 {
		begin = 0
	}
	if begin >= lth {
		begin = lth
	}
	end := begin + length
	if end > lth {
		end = lth
	}
	// 返回子串
	return string(rs[begin:end])
}

func ToFloat(str interface{}) float64 {
	if str == nil {
		return 0
	}
	tf, err := strconv.ParseFloat(ToString(str), 64)
	if err != nil {
		return 0
	}
	return tf
}
func ToJson(datas interface{}) string {
	jsonString, _ := json.Marshal(datas)
	return string(jsonString)
}
func GetEncryptPhone(phone interface{}) string {
	temp := ToString(phone)
	if len(temp) == 11 {
		start := SubString(temp, 0, 3)
		end := SubString(temp, 7, 11)
		temp = start + "****" + end
	}
	return temp
}
func Exist(filename string) bool {
	_, err := os.Stat(filename)
	return err == nil || os.IsExist(err)
}
func SendMail(user, password, host, to, subject, body, mailtype string) error {
	hp := strings.Split(host, ":")
	auth := smtp.PlainAuth("", user, password, hp[0])
	var content_type string
	if mailtype == "html" {
		content_type = "Content-Type: text/" + mailtype + "; charset=UTF-8"
	} else {
		content_type = "Content-Type: text/plain" + "; charset=UTF-8"
	}

	msg := []byte("To: " + to + "\r\nFrom: " + user + "<" + user + ">\r\nSubject: " + subject + "\r\n" + content_type + "\r\n\r\n" + body)
	send_to := strings.Split(to, ";")
	err := smtp.SendMail(host, auth, user, send_to, msg)
	return err
}
func ZeroBefore(i int) string {
	if i < 10 {
		return "0" + ToString(i)
	}
	return ToString(i)
}
func GetTimeStamp() string {
	return time.Now().Format("20060102150405")
}
func Utf8ToGBK(text string) string {
	enc := mahonia.NewEncoder("gbk")
	return enc.ConvertString(text)
}
func GBKFileToUtf8(filePath string) string {
	// Read UTF-8 from a GBK encoded file.
	f, err := os.Open(filePath)
	if err != nil {
		log.Fatal(err)
	}
	r := transform.NewReader(f, simplifiedchinese.GBK.NewDecoder())

	// Read converted UTF-8 from `r` as needed.
	// As an example we'll read line-by-line showing what was read:
	sc := bufio.NewScanner(r)
	result := ""
	for sc.Scan() {
		result += string(sc.Bytes()) + "\n"
	}
	return result
}
func TrimSuffix(s, suffix string) string {
	if strings.HasSuffix(s, suffix) {
		s = s[:len(s)-len(suffix)]
	}
	return s
}
func TrimPrefix(s, suffix string) string {
	if strings.HasPrefix(s, suffix) {
		s = s[len(suffix):]
	}
	return s
}
func Copy(dst, src string) error {
	in, err := os.Open(src)
	if err != nil {
		return err
	}
	defer in.Close()
	out, err := os.Create(dst)
	if err != nil {
		return err
	}
	defer out.Close()
	_, err = io.Copy(out, in)
	cerr := out.Close()
	if err != nil {
		return err
	}
	return cerr
}

func GetStructName(myvar interface{}) string {
	if t := reflect.TypeOf(myvar); t.Kind() == reflect.Ptr {
		return t.Elem().Name()
	} else {
		return t.Name()
	}
}
<commit_msg>update<commit_after>package ofutils

import (
	"bufio"
	"crypto/md5"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"math/rand"
	"net/smtp"
	"os"
	"reflect"
	"sort"
	"strconv"
	"strings"
	"time"

	"golang.org/x/text/encoding/simplifiedchinese"
	"golang.org/x/text/transform"

	"github.com/astaxie/beego/orm"
	"github.com/axgle/mahonia"
)

func init() {
	rand.Seed(time.Now().UnixNano())
}

var letterRunes = []rune("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ")

func RandString(n int) string {
	b := make([]rune, n)
	for i := range b {
		b[i] = letterRunes[rand.Intn(len(letterRunes))]
	}
	return string(b)
}
func ChangeJson(data *[]byte, key string, value interface{}) error {
	var m map[string]interface{}
	err := json.Unmarshal(*data, &m)
	if err != nil {
		return err
	}
	m[key] = value
	*data, err = json.Marshal(m)
	return err
}
func MD5(text string) string {
	hasher := md5.New()
	hasher.Write([]byte(text))
	return hex.EncodeToString(hasher.Sum(nil))
}

func ToString(str interface{}) string {
	if str == nil {
		return ""
	}
	return fmt.Sprintf("%v", str)
}

func ToInt(val interface{}) int {
	if val == nil {
		return 0
	}
	s, ok := val.(string)
	if ok {
		i, err := strconv.Atoi(s)
		if err != nil {
			return 0
		}
		return i
	}
	i, ok := val.(int)
	if ok {
		return i
	} else {
		return 0
	}
}

func GetWeekFirstDay() string {
	week := time.Now().Weekday().String()
	var day time.Duration
	switch week {
	case "Sunday":
		day = 6
		break
	case "Monday":
		day = 0
		break
	case "Tuesday":
		day = 1
		break
	case "Wednesday":
		day = 2
		break
	case "Thursday":
		day = 3
		break
	case "Friday":
		day = 4
		break
	case "Saturday":
		day = 5
		break
	}
	date := time.Now().Add(-day * 24 * time.Hour)
	return date.Format("2006-01-02")
}

type ByKey struct {
	Key  string
	List []orm.Params
}

func (a ByKey) Len() int {
	return len(a.List)
}
func (a ByKey) Swap(i, j int) {
	a.List[i], a.List[j] = a.List[j], a.List[i]
}
func (a ByKey) Less(i, j int) bool {
	return ToInt(a.List[i][a.Key]) > ToInt(a.List[j][a.Key])
}

func Sort(list []orm.Params, key string) []orm.Params {
	byKey := ByKey{List: list, Key: key}
	sort.Sort(byKey)
	return byKey.List
}

func SubString(str string, begin, length int) (substr string) {
	// 将字符串的转换成[]rune
	rs := []rune(str)
	lth := len(rs)

	// 简单的越界判断
	if begin < 0 {
		begin = 0
	}
	if begin >= lth {
		begin = lth
	}
	end := begin + length
	if end > lth {
		end = lth
	}
	// 返回子串
	return string(rs[begin:end])
}

func ToFloat(str interface{}) float64 {
	if str == nil {
		return 0
	}
	tf, err := strconv.ParseFloat(ToString(str), 64)
	if err != nil {
		return 0
	}
	return tf
}
func ToJson(datas interface{}) string {
	jsonString, _ := json.Marshal(datas)
	return string(jsonString)
}
func GetEncryptPhone(phone interface{}) string {
	temp := ToString(phone)
	if len(temp) == 11 {
		start := SubString(temp, 0, 3)
		end := SubString(temp, 7, 11)
		temp = start + "****" + end
	}
	return temp
}
func Exist(filename string) bool {
	_, err := os.Stat(filename)
	return err == nil || os.IsExist(err)
}
func SendMail(user, password, host, to, subject, body, mailtype string) error {
	hp := strings.Split(host, ":")
	auth := smtp.PlainAuth("", user, password, hp[0])
	var content_type string
	if mailtype == "html" {
		content_type = "Content-Type: text/" + mailtype + "; charset=UTF-8"
	} else {
		content_type = "Content-Type: text/plain" + "; charset=UTF-8"
	}

	msg := []byte("To: " + to + "\r\nFrom: " + user + "<" + user + ">\r\nSubject: " + subject + "\r\n" + content_type + "\r\n\r\n" + body)
	send_to := strings.Split(to, ";")
	err := smtp.SendMail(host, auth, user, send_to, msg)
	return err
}
func ZeroBefore(i int) string {
	if i < 10 {
		return "0" + ToString(i)
	}
	return ToString(i)
}
func GetTimeStamp() string {
	return time.Now().Format("20060102150405")
}
func Utf8ToGBK(text string) string {
	enc := mahonia.NewEncoder("gbk")
	return enc.ConvertString(text)
}
func GBKFileToUtf8(filePath string) string {
	// Read UTF-8 from a GBK encoded file.
	f, err := os.Open(filePath)
	if err != nil {
		log.Fatal(err)
	}
	r := transform.NewReader(f, simplifiedchinese.GBK.NewDecoder())

	// Read converted UTF-8 from `r` as needed.
	// As an example we'll read line-by-line showing what was read:
	sc := bufio.NewScanner(r)
	result := ""
	for sc.Scan() {
		result += string(sc.Bytes()) + "\n"
	}
	return result
}
func TrimSuffix(s, suffix string) string {
	if strings.HasSuffix(s, suffix) {
		s = s[:len(s)-len(suffix)]
	}
	return s
}
func TrimPrefix(s, suffix string) string {
	if strings.HasPrefix(s, suffix) {
		s = s[len(suffix):]
	}
	return s
}
func Copy(dst, src string) error {
	in, err := os.Open(src)
	if err != nil {
		return err
	}
	defer in.Close()
	out, err := os.Create(dst)
	if err != nil {
		return err
	}
	defer out.Close()
	_, err = io.Copy(out, in)
	cerr := out.Close()
	if err != nil {
		return err
	}
	return cerr
}

func GetStructName(myvar interface{}) string {
	if t := reflect.TypeOf(myvar); t.Kind() == reflect.Ptr {
		return t.Elem().Name()
	} else {
		return t.Name()
	}
}

func ByteToMapArray(data [][]byte) []map[string]interface{} {
	list := make([]map[string]interface{}, len(data))
	for i, v := range data {
		json.Unmarshal(v, &list[i])
	}
	return list
}
<|endoftext|>


<commit_before>package handling_test

import (
	"bytes"
	"encoding/json"
	"errors"
	"io/ioutil"
	"net/http"
	"reflect"

	. "github.com/gotgo/gokn/handling"
	"github.com/gotgo/gokn/rest"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

type TestResponseWriter struct {
	WriteReturnCount int
	WriteReturnError error
	WriteBytes       []byte
	WriteHeaderCode  int
}

func (trw *TestResponseWriter) Header() http.Header {
	return make(map[string][]string)
}

func (trw *TestResponseWriter) Write(bytes []byte) (int, error) {
	trw.WriteBytes = bytes
	return len(bytes), trw.WriteReturnError
}

func (trw *TestResponseWriter) WriteHeader(code int) {
	trw.WriteHeaderCode = code
}

type TestRouter struct {
	RegisterCount int
	GetCount      int
	PostCount     int
	PutCount      int
	DeleteCount   int
	HeadCount     int
	PatchCount    int
	Handlers      []func(http.ResponseWriter, *http.Request)
}

func NewTestRouter() *TestRouter {
	router := new(TestRouter)
	router.Handlers = []func(http.ResponseWriter, *http.Request){}
	return router
}

func (tr *TestRouter) RequestArgs(req *http.Request) map[string]string {
	return make(map[string]string)
}

func (tr *TestRouter) RegisterRoute(verb, path string, f func(http.ResponseWriter, *http.Request)) {
	tr.RegisterCount++
	switch verb {
	case "GET":
		tr.GetCount++
	case "POST":
		tr.PostCount++
	case "PUT":
		tr.PutCount++
	case "DELETE":
		tr.DeleteCount++
	case "HEAD":
		tr.HeadCount++
	case "PATCH":
		tr.PatchCount++
	}
	tr.Handlers = append(tr.Handlers, f)
}

type TestHandler struct {
	ResponseStatus int
}
type TestStruct struct {
	Message string
}

func NewTestHandler() *TestHandler {
	h := new(TestHandler)
	h.ResponseStatus = 200
	return h
}

func (th *TestHandler) setResponse(resp rest.Responder) {
	if th.ResponseStatus != 200 {
		resp.SetStatus(th.ResponseStatus, "")
	} else {
		resp.SetBody(&TestStruct{"response"})
	}
}

func (th *TestHandler) Get(req *rest.Request, resp rest.Responder) {
	th.setResponse(resp)
}
func (th *TestHandler) Post(req *rest.Request, resp rest.Responder) {
	th.setResponse(resp)
}
func (th *TestHandler) Put(req *rest.Request, resp rest.Responder) {
	th.setResponse(resp)
}
func (th *TestHandler) Delete(req *rest.Request, resp rest.Responder) {
	th.setResponse(resp)
}
func (th *TestHandler) Head(req *rest.Request, resp rest.Responder) {
	th.setResponse(resp)
}
func (th *TestHandler) Patch(req *rest.Request, resp rest.Responder) {
	th.setResponse(resp)
}

func getSpec(resourceT string, verb string) rest.ServerResource {
	def := &rest.ResourceDef{
		ResourceT:    resourceT,
		Verb:         verb,
		RequestBody:  reflect.TypeOf([]byte{}),
		ResponseBody: reflect.TypeOf([]byte{}),
	}
	ct := []string{"application/json"}
	return rest.NewServerResource(def, ct, ct)
}

var _ = Describe("RootHandler", func() {

	var (
		root    *RootHandler
		handler *TestHandler
		router  *TestRouter
		request *http.Request
		writer  *TestResponseWriter
	)

	BeforeEach(func() {
		root = NewRootHandler()
		handler = NewTestHandler()
		router = NewTestRouter()

		tm := &TestStruct{Message: ""}

		buf, _ := json.Marshal(tm)
		request = &http.Request{
			Method:        "POST",
			Body:          ioutil.NopCloser(bytes.NewReader(buf)),
			ContentLength: int64(len(buf)),
		}

		writer = new(TestResponseWriter)
	})

	Context("Test Init", func() {
		It("should have reasonable defaults set after NewRootHandler", func() {
			Expect(root.Log).ToNot(BeNil())
			Expect(root.Binder).ToNot(BeNil())
		})
	})

	Context("Bind", func() {
		It("should bind GET", func() {
			spec := getSpec("/test", "GET")
			root.Bind(router, spec, handler, "")
			Expect(router.RegisterCount).To(Equal(1))
			Expect(router.GetCount).To(Equal(1))
			Expect(len(router.Handlers)).To(Equal(1))
		})
		It("should bind POST", func() {
			spec := getSpec("/test", "POST")
			root.Bind(router, spec, handler, "")
			Expect(router.RegisterCount).To(Equal(1))
			Expect(router.PostCount).To(Equal(1))
			Expect(len(router.Handlers)).To(Equal(1))
		})
		It("should bind PUT", func() {
			spec := getSpec("/test", "PUT")
			root.Bind(router, spec, handler, "")
			Expect(router.RegisterCount).To(Equal(1))
			Expect(router.PutCount).To(Equal(1))
			Expect(len(router.Handlers)).To(Equal(1))
		})
		It("should bind DELETE", func() {
			spec := getSpec("/test", "DELETE")
			root.Bind(router, spec, handler, "")
			Expect(router.RegisterCount).To(Equal(1))
			Expect(router.DeleteCount).To(Equal(1))
			Expect(len(router.Handlers)).To(Equal(1))
		})
		It("should bind HEAD", func() {
			spec := getSpec("/test", "HEAD")
			root.Bind(router, spec, handler, "")
			Expect(router.RegisterCount).To(Equal(1))
			Expect(router.HeadCount).To(Equal(1))
			Expect(len(router.Handlers)).To(Equal(1))
		})
		It("should bind PATCH", func() {
			spec := getSpec("/test", "PATCH")
			root.Bind(router, spec, handler, "")
			Expect(router.RegisterCount).To(Equal(1))
			Expect(router.PatchCount).To(Equal(1))
			Expect(len(router.Handlers)).To(Equal(1))
		})
		It("should bind to all verbs in one go", func() {
			verbs := []string{"GET", "POST", "PUT", "DELETE", "HEAD", "PATCH"}
			for _, verb := range verbs {
				root.Bind(router, getSpec("/test", verb), handler, "")
			}
			Expect(router.RegisterCount).To(Equal(6))
			Expect(router.GetCount).To(Equal(1))
			Expect(router.PostCount).To(Equal(1))
			Expect(router.PutCount).To(Equal(1))
			Expect(router.DeleteCount).To(Equal(1))
			Expect(router.HeadCount).To(Equal(1))
			Expect(router.PatchCount).To(Equal(1))
			Expect(len(router.Handlers)).To(Equal(6))
		})
	})

	Context("BindAll", func() {
		It("should work", func() {
			spec1 := getSpec("/testA", "GET")
			spec2 := getSpec("/testB", "GET")
			spec3 := getSpec("/testC", "POST")
			specs := make(map[rest.ServerResource]rest.Handler)
			specs[spec1] = handler
			specs[spec2] = handler
			specs[spec3] = handler
			root.BindAll(router, specs, "")
			Expect(router.RegisterCount).To(Equal(3))
			Expect(router.GetCount).To(Equal(2))
			Expect(router.PostCount).To(Equal(1))
			Expect(len(router.Handlers)).To(Equal(3))
		})
	})

	Context("Calling wrapped handler", func() {

		It("should write bytes", func() {
			spec := getSpec("/test", "POST")
			root.Bind(router, spec, handler, "")
			Expect(len(router.Handlers)).To(Equal(1))
			wrappedHandler := router.Handlers[0]
			request.Header = make(map[string][]string)
			request.Header["Content-Type"] = []string{"application/json"}
			handler.ResponseStatus = 200
			wrappedHandler(writer, request)
			Expect(writer.WriteBytes).ToNot(BeNil())
			Expect(len(writer.WriteBytes)).Should(BeNumerically(">", 0))
		})

		It("should write header value of error on response error", func() {
			spec := getSpec("/test", "POST")
			root.Bind(router, spec, handler, "")
			wrappedHandler := router.Handlers[0]

			status := 401
			handler.ResponseStatus = status
			wrappedHandler(writer, request)
			Expect(len(writer.WriteBytes)).To(Equal(0))
			Expect(writer.WriteHeaderCode).To(Equal(status))
		})

		It("should return an error code on Encode failure", func() {

			root.Encoders.Set(&ContentTypeEncoder{
				ContentType: "fail",
				Encode: func(v interface{}) ([]byte, error) {
					return nil, errors.New("fail")
				},
			})
			def := &rest.ResourceDef{
				ResourceT: "willfail",
				Verb:      "GET",
			}

			ct := []string{"fail"}
			spec := rest.NewServerResource(def, ct, ct)
			root.Bind(router, spec, handler, "")
			router.Handlers[0](writer, request)
			Expect(writer.WriteHeaderCode).To(Equal(http.StatusInternalServerError))
		})

	})

})
<commit_msg>method signature fix for test<commit_after>package handling_test

import (
	"bytes"
	"encoding/json"
	"errors"
	"io/ioutil"
	"net/http"
	"reflect"

	. "github.com/gotgo/gokn/handling"
	"github.com/gotgo/gokn/rest"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

type TestResponseWriter struct {
	WriteReturnCount int
	WriteReturnError error
	WriteBytes       []byte
	WriteHeaderCode  int
}

func (trw *TestResponseWriter) Header() http.Header {
	return make(map[string][]string)
}

func (trw *TestResponseWriter) Write(bytes []byte) (int, error) {
	trw.WriteBytes = bytes
	return len(bytes), trw.WriteReturnError
}

func (trw *TestResponseWriter) WriteHeader(code int) {
	trw.WriteHeaderCode = code
}

type TestRouter struct {
	RegisterCount int
	GetCount      int
	PostCount     int
	PutCount      int
	DeleteCount   int
	HeadCount     int
	PatchCount    int
	Handlers      []func(http.ResponseWriter, *http.Request)
}

func NewTestRouter() *TestRouter {
	router := new(TestRouter)
	router.Handlers = []func(http.ResponseWriter, *http.Request){}
	return router
}

func (tr *TestRouter) RequestArgs(req *http.Request) map[string]string {
	return make(map[string]string)
}

func (tr *TestRouter) RegisterRoute(verb, path string, f func(http.ResponseWriter, *http.Request)) {
	tr.RegisterCount++
	switch verb {
	case "GET":
		tr.GetCount++
	case "POST":
		tr.PostCount++
	case "PUT":
		tr.PutCount++
	case "DELETE":
		tr.DeleteCount++
	case "HEAD":
		tr.HeadCount++
	case "PATCH":
		tr.PatchCount++
	}
	tr.Handlers = append(tr.Handlers, f)
}

type TestHandler struct {
	ResponseStatus int
}
type TestStruct struct {
	Message string
}

func NewTestHandler() *TestHandler {
	h := new(TestHandler)
	h.ResponseStatus = 200
	return h
}

func (th *TestHandler) setResponse(resp rest.Responder) {
	if th.ResponseStatus != 200 {
		resp.SetStatus(th.ResponseStatus, "", nil)
	} else {
		resp.SetBody(&TestStruct{"response"})
	}
}

func (th *TestHandler) Get(req *rest.Request, resp rest.Responder) {
	th.setResponse(resp)
}
func (th *TestHandler) Post(req *rest.Request, resp rest.Responder) {
	th.setResponse(resp)
}
func (th *TestHandler) Put(req *rest.Request, resp rest.Responder) {
	th.setResponse(resp)
}
func (th *TestHandler) Delete(req *rest.Request, resp rest.Responder) {
	th.setResponse(resp)
}
func (th *TestHandler) Head(req *rest.Request, resp rest.Responder) {
	th.setResponse(resp)
}
func (th *TestHandler) Patch(req *rest.Request, resp rest.Responder) {
	th.setResponse(resp)
}

func getSpec(resourceT string, verb string) rest.ServerResource {
	def := &rest.ResourceDef{
		ResourceT:    resourceT,
		Verb:         verb,
		RequestBody:  reflect.TypeOf([]byte{}),
		ResponseBody: reflect.TypeOf([]byte{}),
	}
	ct := []string{"application/json"}
	return rest.NewServerResource(def, ct, ct)
}

var _ = Describe("RootHandler", func() {

	var (
		root    *RootHandler
		handler *TestHandler
		router  *TestRouter
		request *http.Request
		writer  *TestResponseWriter
	)

	BeforeEach(func() {
		root = NewRootHandler()
		handler = NewTestHandler()
		router = NewTestRouter()

		tm := &TestStruct{Message: ""}

		buf, _ := json.Marshal(tm)
		request = &http.Request{
			Method:        "POST",
			Body:          ioutil.NopCloser(bytes.NewReader(buf)),
			ContentLength: int64(len(buf)),
		}

		writer = new(TestResponseWriter)
	})

	Context("Test Init", func() {
		It("should have reasonable defaults set after NewRootHandler", func() {
			Expect(root.Log).ToNot(BeNil())
			Expect(root.Binder).ToNot(BeNil())
		})
	})

	Context("Bind", func() {
		It("should bind GET", func() {
			spec := getSpec("/test", "GET")
			root.Bind(router, spec, handler, "")
			Expect(router.RegisterCount).To(Equal(1))
			Expect(router.GetCount).To(Equal(1))
			Expect(len(router.Handlers)).To(Equal(1))
		})
		It("should bind POST", func() {
			spec := getSpec("/test", "POST")
			root.Bind(router, spec, handler, "")
			Expect(router.RegisterCount).To(Equal(1))
			Expect(router.PostCount).To(Equal(1))
			Expect(len(router.Handlers)).To(Equal(1))
		})
		It("should bind PUT", func() {
			spec := getSpec("/test", "PUT")
			root.Bind(router, spec, handler, "")
			Expect(router.RegisterCount).To(Equal(1))
			Expect(router.PutCount).To(Equal(1))
			Expect(len(router.Handlers)).To(Equal(1))
		})
		It("should bind DELETE", func() {
			spec := getSpec("/test", "DELETE")
			root.Bind(router, spec, handler, "")
			Expect(router.RegisterCount).To(Equal(1))
			Expect(router.DeleteCount).To(Equal(1))
			Expect(len(router.Handlers)).To(Equal(1))
		})
		It("should bind HEAD", func() {
			spec := getSpec("/test", "HEAD")
			root.Bind(router, spec, handler, "")
			Expect(router.RegisterCount).To(Equal(1))
			Expect(router.HeadCount).To(Equal(1))
			Expect(len(router.Handlers)).To(Equal(1))
		})
		It("should bind PATCH", func() {
			spec := getSpec("/test", "PATCH")
			root.Bind(router, spec, handler, "")
			Expect(router.RegisterCount).To(Equal(1))
			Expect(router.PatchCount).To(Equal(1))
			Expect(len(router.Handlers)).To(Equal(1))
		})
		It("should bind to all verbs in one go", func() {
			verbs := []string{"GET", "POST", "PUT", "DELETE", "HEAD", "PATCH"}
			for _, verb := range verbs {
				root.Bind(router, getSpec("/test", verb), handler, "")
			}
			Expect(router.RegisterCount).To(Equal(6))
			Expect(router.GetCount).To(Equal(1))
			Expect(router.PostCount).To(Equal(1))
			Expect(router.PutCount).To(Equal(1))
			Expect(router.DeleteCount).To(Equal(1))
			Expect(router.HeadCount).To(Equal(1))
			Expect(router.PatchCount).To(Equal(1))
			Expect(len(router.Handlers)).To(Equal(6))
		})
	})

	Context("BindAll", func() {
		It("should work", func() {
			spec1 := getSpec("/testA", "GET")
			spec2 := getSpec("/testB", "GET")
			spec3 := getSpec("/testC", "POST")
			specs := make(map[rest.ServerResource]rest.Handler)
			specs[spec1] = handler
			specs[spec2] = handler
			specs[spec3] = handler
			root.BindAll(router, specs, "")
			Expect(router.RegisterCount).To(Equal(3))
			Expect(router.GetCount).To(Equal(2))
			Expect(router.PostCount).To(Equal(1))
			Expect(len(router.Handlers)).To(Equal(3))
		})
	})

	Context("Calling wrapped handler", func() {

		It("should write bytes", func() {
			spec := getSpec("/test", "POST")
			root.Bind(router, spec, handler, "")
			Expect(len(router.Handlers)).To(Equal(1))
			wrappedHandler := router.Handlers[0]
			request.Header = make(map[string][]string)
			request.Header["Content-Type"] = []string{"application/json"}
			handler.ResponseStatus = 200
			wrappedHandler(writer, request)
			Expect(writer.WriteBytes).ToNot(BeNil())
			Expect(len(writer.WriteBytes)).Should(BeNumerically(">", 0))
		})

		It("should write header value of error on response error", func() {
			spec := getSpec("/test", "POST")
			root.Bind(router, spec, handler, "")
			wrappedHandler := router.Handlers[0]

			status := 401
			handler.ResponseStatus = status
			wrappedHandler(writer, request)
			Expect(len(writer.WriteBytes)).To(Equal(0))
			Expect(writer.WriteHeaderCode).To(Equal(status))
		})

		It("should return an error code on Encode failure", func() {

			root.Encoders.Set(&ContentTypeEncoder{
				ContentType: "fail",
				Encode: func(v interface{}) ([]byte, error) {
					return nil, errors.New("fail")
				},
			})
			def := &rest.ResourceDef{
				ResourceT: "willfail",
				Verb:      "GET",
			}

			ct := []string{"fail"}
			spec := rest.NewServerResource(def, ct, ct)
			root.Bind(router, spec, handler, "")
			router.Handlers[0](writer, request)
			Expect(writer.WriteHeaderCode).To(Equal(http.StatusInternalServerError))
		})

	})

})
<|endoftext|>