<commit_before>//go:generate struct-markdown
//go:generate mapstructure-to-hcl2 -type MachineImageFilter

package triton

import (
	"fmt"

	"github.com/hashicorp/packer/hcl2template"
	"github.com/hashicorp/packer/template/interpolate"
)

// SourceMachineConfig represents the configuration to run a machine using
// the SDC API in order for provisioning to take place.
type SourceMachineConfig struct {
	// Name of the VM used for building the
	// image. Does not affect (and does not have to be the same) as the name for a
	// VM instance running this image. Maximum 512 characters but should in
	// practice be much shorter (think between 5 and 20 characters). For example
	// mysql-64-server-image-builder. When omitted defaults to
	// packer-builder-[image_name].
	MachineName string `mapstructure:"source_machine_name" required:"false"`
	// The Triton package to use while
	// building the image. Does not affect (and does not have to be the same) as
	// the package which will be used for a VM instance running this image. On the
	// Joyent public cloud this could for example be g3-standard-0.5-smartos.
	MachinePackage string `mapstructure:"source_machine_package" required:"true"`
	// The UUID of the image to base the new
	// image on. Triton supports multiple types of images, called 'brands' in
	// Triton / Joyent lingo, for contains and VM's. See the chapter Containers
	// and virtual machines in
	// the Joyent Triton documentation for detailed information. The following
	// brands are currently supported by this builder:joyent andkvm. The
	// choice of base image automatically decides the brand. On the Joyent public
	// cloud a valid source_machine_image could for example be
	// 70e3ae72-96b6-11e6-9056-9737fd4d0764 for version 16.3.1 of the 64bit
	// SmartOS base image (a 'joyent' brand image). source_machine_image_filter
	// can be used to populate this UUID.
	MachineImage string `mapstructure:"source_machine_image" required:"true"`
	// The UUID's of Triton
	// networks added to the source machine used for creating the image. For
	// example if any of the provisioners which are run need Internet access you
	// will need to add the UUID's of the appropriate networks here. If this is
	// not specified, instances will be placed into the default Triton public and
	// internal networks.
	MachineNetworks []string `mapstructure:"source_machine_networks" required:"false"`
	// Triton metadata
	// applied to the VM used to create the image. Metadata can be used to pass
	// configuration information to the VM without the need for networking. See
	// Using the metadata
	// API in the
	// Joyent documentation for more information. This can for example be used to
	// set the user-script metadata key to have Triton start a user supplied
	// script after the VM has booted.
	MachineMetadata map[string]string `mapstructure:"source_machine_metadata" required:"false"`
	// Tags applied to the VM used to create the image.
	MachineTags map[string]string `mapstructure:"source_machine_tags" required:"false"`
	// Same as [`source_machine_tags`](#source_machine_tags) but defined as a
	// singular block containing a key and a value field. In HCL2 mode the
	// [`dynamic_block`](https://packer.io/docs/configuration/from-1.5/expressions.html#dynamic-blocks)
	// will allow you to create those programatically.
	MachineTag hcl2template.KeyValues `mapstructure:"source_machine_tag" required:"false"`
	// Whether or not the firewall
	// of the VM used to create an image of is enabled. The Triton firewall only
	// filters inbound traffic to the VM. All outbound traffic is always allowed.
	// Currently this builder does not provide an interface to add specific
	// firewall rules. Unless you have a global rule defined in Triton which
	// allows SSH traffic enabling the firewall will interfere with the SSH
	// provisioner. The default is false.
	MachineFirewallEnabled bool `mapstructure:"source_machine_firewall_enabled" required:"false"`
	// Filters used to populate the
	// source_machine_image field. Example:
	MachineImageFilters MachineImageFilter `mapstructure:"source_machine_image_filter" required:"false"`
}

type MachineImageFilter struct {
	MostRecent bool `mapstructure:"most_recent"`
	Name       string
	OS         string
	Version    string
	Public     bool
	State      string
	Owner      string
	Type       string
}

func (m *MachineImageFilter) Empty() bool {
	return m.Name == "" && m.OS == "" && m.Version == "" && m.State == "" && m.Owner == "" && m.Type == ""
}

// Prepare performs basic validation on a SourceMachineConfig struct.
func (c *SourceMachineConfig) Prepare(ctx *interpolate.Context) []error {
	var errs []error

	if c.MachinePackage == "" {
		errs = append(errs, fmt.Errorf("A source_machine_package must be specified"))
	}

	if c.MachineImage != "" && c.MachineImageFilters.Name != "" {
		errs = append(errs, fmt.Errorf("You cannot specify a Machine Image and also Machine Name filter"))
	}

	if c.MachineNetworks == nil {
		c.MachineNetworks = []string{}
	}

	if c.MachineMetadata == nil {
		c.MachineMetadata = make(map[string]string)
	}

	if c.MachineTags == nil {
		c.MachineTags = make(map[string]string)
	}

	if len(errs) > 0 {
		return errs
	}

	return nil
}
<commit_msg>Update source_machine_config.go<commit_after>//go:generate struct-markdown
//go:generate mapstructure-to-hcl2 -type MachineImageFilter

package triton

import (
	"fmt"

	"github.com/hashicorp/packer/hcl2template"
	"github.com/hashicorp/packer/template/interpolate"
)

// SourceMachineConfig represents the configuration to run a machine using
// the SDC API in order for provisioning to take place.
type SourceMachineConfig struct {
	// Name of the VM used for building the
	// image. Does not affect (and does not have to be the same) as the name for a
	// VM instance running this image. Maximum 512 characters but should in
	// practice be much shorter (think between 5 and 20 characters). For example
	// mysql-64-server-image-builder. When omitted defaults to
	// packer-builder-[image_name].
	MachineName string `mapstructure:"source_machine_name" required:"false"`
	// The Triton package to use while
	// building the image. Does not affect (and does not have to be the same) as
	// the package which will be used for a VM instance running this image. On the
	// Joyent public cloud this could for example be g3-standard-0.5-smartos.
	MachinePackage string `mapstructure:"source_machine_package" required:"true"`
	// The UUID of the image to base the new
	// image on. Triton supports multiple types of images, called 'brands' in
	// Triton / Joyent lingo, for contains and VM's. See the chapter Containers
	// and virtual machines in
	// the Joyent Triton documentation for detailed information. The following
	// brands are currently supported by this builder:joyent andkvm. The
	// choice of base image automatically decides the brand. On the Joyent public
	// cloud a valid source_machine_image could for example be
	// 70e3ae72-96b6-11e6-9056-9737fd4d0764 for version 16.3.1 of the 64bit
	// SmartOS base image (a 'joyent' brand image). source_machine_image_filter
	// can be used to populate this UUID.
	MachineImage string `mapstructure:"source_machine_image" required:"true"`
	// The UUID's of Triton
	// networks added to the source machine used for creating the image. For
	// example if any of the provisioners which are run need Internet access you
	// will need to add the UUID's of the appropriate networks here. If this is
	// not specified, instances will be placed into the default Triton public and
	// internal networks.
	MachineNetworks []string `mapstructure:"source_machine_networks" required:"false"`
	// Triton metadata
	// applied to the VM used to create the image. Metadata can be used to pass
	// configuration information to the VM without the need for networking. See
	// Using the metadata
	// API in the
	// Joyent documentation for more information. This can for example be used to
	// set the user-script metadata key to have Triton start a user supplied
	// script after the VM has booted.
	MachineMetadata map[string]string `mapstructure:"source_machine_metadata" required:"false"`
	// Tags applied to the VM used to create the image.
	MachineTags map[string]string `mapstructure:"source_machine_tags" required:"false"`
	// Same as [`source_machine_tags`](#source_machine_tags) but defined as a
	// singular block containing a key and a value field. In HCL2 mode the
	// [`dynamic_block`](https://packer.io/docs/configuration/from-1.5/expressions.html#dynamic-blocks)
	// will allow you to create those programatically.
	MachineTag hcl2template.KeyValues `mapstructure:"source_machine_tag" required:"false"`
	// Whether or not the firewall
	// of the VM used to create an image of is enabled. The Triton firewall only
	// filters inbound traffic to the VM. All outbound traffic is always allowed.
	// Currently this builder does not provide an interface to add specific
	// firewall rules. Unless you have a global rule defined in Triton which
	// allows SSH traffic enabling the firewall will interfere with the SSH
	// provisioner. The default is false.
	MachineFirewallEnabled bool `mapstructure:"source_machine_firewall_enabled" required:"false"`
	// Filters used to populate the
	// source_machine_image field. Example:
	MachineImageFilters MachineImageFilter `mapstructure:"source_machine_image_filter" required:"false"`
}

type MachineImageFilter struct {
	MostRecent bool `mapstructure:"most_recent"`
	Name       string
	OS         string
	Version    string
	Public     bool
	State      string
	Owner      string
	Type       string
}

func (m *MachineImageFilter) Empty() bool {
	return m.Name == "" && m.OS == "" && m.Version == "" && m.State == "" && m.Owner == "" && m.Type == ""
}

// Prepare performs basic validation on a SourceMachineConfig struct.
func (c *SourceMachineConfig) Prepare(ctx *interpolate.Context) []error {
	var errs []error

	if c.MachinePackage == "" {
		errs = append(errs, fmt.Errorf("A source_machine_package must be specified"))
	}

	if c.MachineImage != "" && c.MachineImageFilters.Name != "" {
		errs = append(errs, fmt.Errorf("You cannot specify a Machine Image and also Machine Name filter"))
	}

	if c.MachineNetworks == nil {
		c.MachineNetworks = []string{}
	}

	if c.MachineMetadata == nil {
		c.MachineMetadata = make(map[string]string)
	}

	if c.MachineTags == nil {
		c.MachineTags = make(map[string]string)
	}

	errs = append(errs, c.MachineTag.CopyOn(c.MachineTags)...)

	return errs
}
<|endoftext|>


<commit_before>package operators

import (
	"math"
)

var (
	add = &Operator{
		Name:          "+",
		Precedence:    1,
		Associativity: L,
		Args:          2,
		Operation: func(args []float64) float64 {
			return args[0] + args[1]
		},
	}
	sub = &Operator{
		Name:          "-",
		Precedence:    1,
		Associativity: L,
		Args:          2,
		Operation: func(args []float64) float64 {
			return args[0] - args[1]
		},
	}
	neg = &Operator{
		Name:          "neg",
		Precedence:    2,
		Associativity: L,
		Args:          1,
		Operation: func(args []float64) float64 {
			return 0 - args[0]
		},
	}
	mul = &Operator{
		Name:          "*",
		Precedence:    2,
		Associativity: L,
		Args:          2,
		Operation: func(args []float64) float64 {
			return args[0] * args[1]
		},
	}
	div = &Operator{
		Name:          "/",
		Precedence:    2,
		Associativity: L,
		Args:          2,
		Operation: func(args []float64) float64 {
			return args[0] / args[1]
		},
	}
	mod = &Operator{
		Name:          "%",
		Precedence:    2,
		Associativity: L,
		Args:          2,
		Operation: func(args []float64) float64 {
			return math.Mod(args[0], args[1])
		},
	}
	pow = &Operator{
		Name:          "^",
		Precedence:    3,
		Associativity: R,
		Args:          2,
		Operation: func(args []float64) float64 {
			return math.Pow(args[0], args[1])
		},
	}

)

func init() {
	Register(add)
	Register(sub)
	Register(neg)
	Register(pow)
	Register(mul)
    Register(mod)
	Register(div)
}
<commit_msg>Fix up formatting for modulo operator<commit_after>package operators

import (
	"math"
)

var (
	add = &Operator{
		Name:          "+",
		Precedence:    1,
		Associativity: L,
		Args:          2,
		Operation: func(args []float64) float64 {
			return args[0] + args[1]
		},
	}
	sub = &Operator{
		Name:          "-",
		Precedence:    1,
		Associativity: L,
		Args:          2,
		Operation: func(args []float64) float64 {
			return args[0] - args[1]
		},
	}
	neg = &Operator{
		Name:          "neg",
		Precedence:    2,
		Associativity: L,
		Args:          1,
		Operation: func(args []float64) float64 {
			return 0 - args[0]
		},
	}
	mul = &Operator{
		Name:          "*",
		Precedence:    2,
		Associativity: L,
		Args:          2,
		Operation: func(args []float64) float64 {
			return args[0] * args[1]
		},
	}
	div = &Operator{
		Name:          "/",
		Precedence:    2,
		Associativity: L,
		Args:          2,
		Operation: func(args []float64) float64 {
			return args[0] / args[1]
		},
	}
	mod = &Operator{
		Name:          "%",
		Precedence:    2,
		Associativity: L,
		Args:          2,
		Operation: func(args []float64) float64 {
			return math.Mod(args[0], args[1])
		},
	}
	pow = &Operator{
		Name:          "^",
		Precedence:    3,
		Associativity: R,
		Args:          2,
		Operation: func(args []float64) float64 {
			return math.Pow(args[0], args[1])
		},
	}
)

func init() {
	Register(add)
	Register(sub)
	Register(neg)
	Register(pow)
	Register(mul)
	Register(mod)
	Register(div)
}
<|endoftext|>


<commit_before>package application_test

import (
	"errors"

	"code.cloudfoundry.org/cli/cf/api/applications/applicationsfakes"
	"code.cloudfoundry.org/cli/cf/commandregistry"
	"code.cloudfoundry.org/cli/cf/configuration/coreconfig"
	"code.cloudfoundry.org/cli/cf/models"
	"code.cloudfoundry.org/cli/cf/requirements"
	"code.cloudfoundry.org/cli/cf/requirements/requirementsfakes"
	testcmd "code.cloudfoundry.org/cli/cf/util/testhelpers/commands"
	testconfig "code.cloudfoundry.org/cli/cf/util/testhelpers/configuration"
	testterm "code.cloudfoundry.org/cli/cf/util/testhelpers/terminal"

	. "code.cloudfoundry.org/cli/cf/util/testhelpers/matchers"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("set-env command", func() {
	var (
		ui                  *testterm.FakeUI
		configRepo          coreconfig.Repository
		app                 models.Application
		appRepo             *applicationsfakes.FakeRepository
		requirementsFactory *requirementsfakes.FakeFactory
		deps                commandregistry.Dependency
	)

	updateCommandDependency := func(pluginCall bool) {
		deps.UI = ui
		deps.Config = configRepo
		deps.RepoLocator = deps.RepoLocator.SetApplicationRepository(appRepo)
		commandregistry.Commands.SetCommand(commandregistry.Commands.FindCommand("set-env").SetDependency(deps, pluginCall))
	}

	BeforeEach(func() {
		ui = &testterm.FakeUI{}
		app = models.Application{}
		app.Name = "my-app"
		app.GUID = "my-app-guid"
		appRepo = new(applicationsfakes.FakeRepository)
		requirementsFactory = new(requirementsfakes.FakeFactory)
		configRepo = testconfig.NewRepositoryWithDefaults()
	})

	runCommand := func(args ...string) bool {
		return testcmd.RunCLICommand("set-env", args, requirementsFactory, updateCommandDependency, false, ui)
	}

	Describe("requirements", func() {
		BeforeEach(func() {
			applicationReq := new(requirementsfakes.FakeApplicationRequirement)
			applicationReq.GetApplicationReturns(app)
			requirementsFactory.NewApplicationRequirementReturns(applicationReq)
		})

		It("fails when login is not successful", func() {
			requirementsFactory.NewLoginRequirementReturns(requirements.Failing{Message: "not logged in"})

			Expect(runCommand("hey", "gabba", "gabba")).To(BeFalse())
		})

		It("fails when a space is not targeted", func() {
			requirementsFactory.NewLoginRequirementReturns(requirements.Passing{})
			requirementsFactory.NewTargetedSpaceRequirementReturns(requirements.Failing{Message: "not targeting space"})

			Expect(runCommand("hey", "gabba", "gabba")).To(BeFalse())
		})

		It("fails with usage when not provided with exactly three args", func() {
			requirementsFactory.NewLoginRequirementReturns(requirements.Passing{})
			requirementsFactory.NewTargetedSpaceRequirementReturns(requirements.Passing{})

			runCommand("zomg", "too", "many", "args")
			Expect(ui.Outputs()).To(ContainSubstrings(
				[]string{"Incorrect Usage", "Requires", "arguments"},
			))
		})
	})

	Context("when logged in, a space is targeted and given enough args", func() {
		BeforeEach(func() {
			app.EnvironmentVars = map[string]interface{}{"foo": "bar"}
			applicationReq := new(requirementsfakes.FakeApplicationRequirement)
			applicationReq.GetApplicationReturns(app)
			requirementsFactory.NewApplicationRequirementReturns(applicationReq)
			requirementsFactory.NewLoginRequirementReturns(requirements.Passing{})
			requirementsFactory.NewTargetedSpaceRequirementReturns(requirements.Passing{})
		})

		Context("when it is new", func() {
			It("is created", func() {
				runCommand("my-app", "DATABASE_URL", "mysql://new-example.com/my-db")

				Expect(ui.Outputs()).To(ContainSubstrings(
					[]string{
						"Setting env variable",
						"DATABASE_URL",
						"my-app",
						"my-org",
						"my-space",
						"my-user",
					},
					[]string{"OK"},
					[]string{"TIP"},
				))

				Expect(ui.Outputs()).ToNot(ContainSubstrings([]string{"mysql://new-example.com/my-db"}))

				appGUID, params := appRepo.UpdateArgsForCall(0)
				Expect(appGUID).To(Equal(app.GUID))
				Expect(*params.EnvironmentVars).To(Equal(map[string]interface{}{
					"DATABASE_URL": "mysql://new-example.com/my-db",
					"foo":          "bar",
				}))
			})
		})

		Context("when it already exists", func() {
			BeforeEach(func() {
				app.EnvironmentVars["DATABASE_URL"] = "mysql://old-example.com/my-db"
			})

			It("is updated", func() {
				runCommand("my-app", "DATABASE_URL", "mysql://new-example.com/my-db")

				Expect(ui.Outputs()).To(ContainSubstrings(
					[]string{
						"Setting env variable",
						"DATABASE_URL",
						"mysql://new-example.com/my-db",
						"my-app",
						"my-org",
						"my-space",
						"my-user",
					},
					[]string{"OK"},
					[]string{"TIP"},
				))
			})
		})

		It("allows the variable value to begin with a hyphen", func() {
			runCommand("my-app", "MY_VAR", "--has-a-cool-value")

			Expect(ui.Outputs()).To(ContainSubstrings(
				[]string{
					"Setting env variable",
					"MY_VAR",
					"--has-a-cool-value",
				},
				[]string{"OK"},
				[]string{"TIP"},
			))
			_, params := appRepo.UpdateArgsForCall(0)
			Expect(*params.EnvironmentVars).To(Equal(map[string]interface{}{
				"MY_VAR": "--has-a-cool-value",
				"foo":    "bar",
			}))
		})

		Context("when setting fails", func() {
			BeforeEach(func() {
				appRepo.UpdateReturns(models.Application{}, errors.New("Error updating app."))
			})

			It("tells the user", func() {
				runCommand("please", "dont", "fail")

				Expect(ui.Outputs()).To(ContainSubstrings(
					[]string{"Setting env variable"},
					[]string{"FAILED"},
					[]string{"Error updating app."},
				))
			})
		})

		It("gives the appropriate tip", func() {
			runCommand("my-app", "DATABASE_URL", "mysql://new-example.com/my-db")
			Expect(ui.Outputs()).To(ContainSubstrings(
				[]string{"TIP: Use 'cf restage my-app' to ensure your env variable changes take effect"},
			))
		})
	})
})
<commit_msg>Fixed failing unit tests on legacy set-env<commit_after>package application_test

import (
	"errors"

	"code.cloudfoundry.org/cli/cf/api/applications/applicationsfakes"
	"code.cloudfoundry.org/cli/cf/commandregistry"
	"code.cloudfoundry.org/cli/cf/configuration/coreconfig"
	"code.cloudfoundry.org/cli/cf/models"
	"code.cloudfoundry.org/cli/cf/requirements"
	"code.cloudfoundry.org/cli/cf/requirements/requirementsfakes"
	testcmd "code.cloudfoundry.org/cli/cf/util/testhelpers/commands"
	testconfig "code.cloudfoundry.org/cli/cf/util/testhelpers/configuration"
	testterm "code.cloudfoundry.org/cli/cf/util/testhelpers/terminal"

	. "code.cloudfoundry.org/cli/cf/util/testhelpers/matchers"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("set-env command", func() {
	var (
		ui                  *testterm.FakeUI
		configRepo          coreconfig.Repository
		app                 models.Application
		appRepo             *applicationsfakes.FakeRepository
		requirementsFactory *requirementsfakes.FakeFactory
		deps                commandregistry.Dependency
	)

	updateCommandDependency := func(pluginCall bool) {
		deps.UI = ui
		deps.Config = configRepo
		deps.RepoLocator = deps.RepoLocator.SetApplicationRepository(appRepo)
		commandregistry.Commands.SetCommand(commandregistry.Commands.FindCommand("set-env").SetDependency(deps, pluginCall))
	}

	BeforeEach(func() {
		ui = &testterm.FakeUI{}
		app = models.Application{}
		app.Name = "my-app"
		app.GUID = "my-app-guid"
		appRepo = new(applicationsfakes.FakeRepository)
		requirementsFactory = new(requirementsfakes.FakeFactory)
		configRepo = testconfig.NewRepositoryWithDefaults()
	})

	runCommand := func(args ...string) bool {
		return testcmd.RunCLICommand("set-env", args, requirementsFactory, updateCommandDependency, false, ui)
	}

	Describe("requirements", func() {
		BeforeEach(func() {
			applicationReq := new(requirementsfakes.FakeApplicationRequirement)
			applicationReq.GetApplicationReturns(app)
			requirementsFactory.NewApplicationRequirementReturns(applicationReq)
		})

		It("fails when login is not successful", func() {
			requirementsFactory.NewLoginRequirementReturns(requirements.Failing{Message: "not logged in"})

			Expect(runCommand("hey", "gabba", "gabba")).To(BeFalse())
		})

		It("fails when a space is not targeted", func() {
			requirementsFactory.NewLoginRequirementReturns(requirements.Passing{})
			requirementsFactory.NewTargetedSpaceRequirementReturns(requirements.Failing{Message: "not targeting space"})

			Expect(runCommand("hey", "gabba", "gabba")).To(BeFalse())
		})

		It("fails with usage when not provided with exactly three args", func() {
			requirementsFactory.NewLoginRequirementReturns(requirements.Passing{})
			requirementsFactory.NewTargetedSpaceRequirementReturns(requirements.Passing{})

			runCommand("zomg", "too", "many", "args")
			Expect(ui.Outputs()).To(ContainSubstrings(
				[]string{"Incorrect Usage", "Requires", "arguments"},
			))
		})
	})

	Context("when logged in, a space is targeted and given enough args", func() {
		BeforeEach(func() {
			app.EnvironmentVars = map[string]interface{}{"foo": "bar"}
			applicationReq := new(requirementsfakes.FakeApplicationRequirement)
			applicationReq.GetApplicationReturns(app)
			requirementsFactory.NewApplicationRequirementReturns(applicationReq)
			requirementsFactory.NewLoginRequirementReturns(requirements.Passing{})
			requirementsFactory.NewTargetedSpaceRequirementReturns(requirements.Passing{})
		})

		Context("when it is new", func() {
			It("is created", func() {
				runCommand("my-app", "DATABASE_URL", "mysql://new-example.com/my-db")

				Expect(ui.Outputs()).To(ContainSubstrings(
					[]string{
						"Setting env variable",
						"DATABASE_URL",
						"my-app",
						"my-org",
						"my-space",
						"my-user",
					},
					[]string{"OK"},
					[]string{"TIP"},
				))

				Expect(ui.Outputs()).ToNot(ContainSubstrings([]string{"mysql://new-example.com/my-db"}))

				appGUID, params := appRepo.UpdateArgsForCall(0)
				Expect(appGUID).To(Equal(app.GUID))
				Expect(*params.EnvironmentVars).To(Equal(map[string]interface{}{
					"DATABASE_URL": "mysql://new-example.com/my-db",
					"foo":          "bar",
				}))
			})
		})

		Context("when it already exists", func() {
			BeforeEach(func() {
				app.EnvironmentVars["DATABASE_URL"] = "mysql://old-example.com/my-db"
			})

			It("is updated", func() {
				runCommand("my-app", "DATABASE_URL", "mysql://new-example.com/my-db")

				Expect(ui.Outputs()).To(ContainSubstrings(
					[]string{
						"Setting env variable",
						"DATABASE_URL",
						"my-app",
						"my-org",
						"my-space",
						"my-user",
					},
					[]string{"OK"},
					[]string{"TIP"},
				))

				Expect(ui.Outputs()).ToNot(ContainSubstrings([]string{"mysql://new-example.com/my-db"}))
			})
		})

		It("allows the variable value to begin with a hyphen", func() {
			runCommand("my-app", "MY_VAR", "--has-a-cool-value")

			Expect(ui.Outputs()).To(ContainSubstrings(
				[]string{
					"Setting env variable",
					"MY_VAR",
				},
				[]string{"OK"},
				[]string{"TIP"},
			))
			_, params := appRepo.UpdateArgsForCall(0)
			Expect(*params.EnvironmentVars).To(Equal(map[string]interface{}{
				"MY_VAR": "--has-a-cool-value",
				"foo":    "bar",
			}))
		})

		Context("when setting fails", func() {
			BeforeEach(func() {
				appRepo.UpdateReturns(models.Application{}, errors.New("Error updating app."))
			})

			It("tells the user", func() {
				runCommand("please", "dont", "fail")

				Expect(ui.Outputs()).To(ContainSubstrings(
					[]string{"Setting env variable"},
					[]string{"FAILED"},
					[]string{"Error updating app."},
				))
			})
		})

		It("gives the appropriate tip", func() {
			runCommand("my-app", "DATABASE_URL", "mysql://new-example.com/my-db")
			Expect(ui.Outputs()).To(ContainSubstrings(
				[]string{"TIP: Use 'cf restage my-app' to ensure your env variable changes take effect"},
			))
		})
	})
})
<|endoftext|>


<commit_before>/*
                       WARNING WARNING WARNING

                Attention all potential contributors

   This testfile is not in the best state. We've been slowly transitioning
   from the built in "testing" package to using Ginkgo. As you can see, we've
   changed the format, but a lot of the setup, test body, descriptions, etc
   are either hardcoded, completely lacking, or misleading.

   For example:

   Describe("Testing with ginkgo"...)      // This is not a great description
   It("TestDoesSoemthing"...)              // This is a horrible description

   Describe("create-user command"...       // Describe the actual object under test
   It("creates a user when provided ..."   // this is more descriptive

   For good examples of writing Ginkgo tests for the cli, refer to

   src/github.com/cloudfoundry/cli/cf/commands/application/delete_app_test.go
   src/github.com/cloudfoundry/cli/cf/terminal/ui_test.go
   src/github.com/cloudfoundry/loggregator_consumer/consumer_test.go
*/

package user_test

import (
	. "github.com/cloudfoundry/cli/cf/commands/user"
	"github.com/cloudfoundry/cli/cf/configuration"
	"github.com/cloudfoundry/cli/cf/models"
	testapi "github.com/cloudfoundry/cli/testhelpers/api"
	testcmd "github.com/cloudfoundry/cli/testhelpers/commands"
	testconfig "github.com/cloudfoundry/cli/testhelpers/configuration"
	testreq "github.com/cloudfoundry/cli/testhelpers/requirements"
	testterm "github.com/cloudfoundry/cli/testhelpers/terminal"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	. "github.com/cloudfoundry/cli/testhelpers/matchers"
)

func getSetSpaceRoleDeps() (requirementsFactory *testreq.FakeReqFactory, spaceRepo *testapi.FakeSpaceRepository, userRepo *testapi.FakeUserRepository) {
	requirementsFactory = &testreq.FakeReqFactory{}
	spaceRepo = &testapi.FakeSpaceRepository{}
	userRepo = &testapi.FakeUserRepository{}
	return
}

func callSetSpaceRole(args []string, requirementsFactory *testreq.FakeReqFactory, spaceRepo *testapi.FakeSpaceRepository, userRepo *testapi.FakeUserRepository) (ui *testterm.FakeUI) {
	ui = new(testterm.FakeUI)
	configRepo := testconfig.NewRepositoryWithDefaults()
	accessToken, err := testconfig.EncodeAccessToken(configuration.TokenInfo{
		Username: "current-user",
	})
	Expect(err).NotTo(HaveOccurred())
	configRepo.SetAccessToken(accessToken)

	cmd := NewSetSpaceRole(ui, configRepo, spaceRepo, userRepo)
	testcmd.RunCommand(cmd, args, requirementsFactory)
	return
}

var _ = Describe("Testing with ginkgo", func() {
	It("TestSetSpaceRoleFailsWithUsage", func() {
		requirementsFactory, spaceRepo, userRepo := getSetSpaceRoleDeps()

		ui := callSetSpaceRole([]string{}, requirementsFactory, spaceRepo, userRepo)
		Expect(ui.FailedWithUsage).To(BeTrue())

		ui = callSetSpaceRole([]string{"my-user"}, requirementsFactory, spaceRepo, userRepo)
		Expect(ui.FailedWithUsage).To(BeTrue())

		ui = callSetSpaceRole([]string{"my-user", "my-org"}, requirementsFactory, spaceRepo, userRepo)
		Expect(ui.FailedWithUsage).To(BeTrue())

		ui = callSetSpaceRole([]string{"my-user", "my-org", "my-space"}, requirementsFactory, spaceRepo, userRepo)
		Expect(ui.FailedWithUsage).To(BeTrue())

		ui = callSetSpaceRole([]string{"my-user", "my-org", "my-space", "my-role"}, requirementsFactory, spaceRepo, userRepo)

		Expect(ui.FailedWithUsage).To(BeFalse())
	})

	It("TestSetSpaceRoleRequirements", func() {
		args := []string{"username", "org", "space", "role"}
		requirementsFactory, spaceRepo, userRepo := getSetSpaceRoleDeps()

		requirementsFactory.LoginSuccess = false
		callSetSpaceRole(args, requirementsFactory, spaceRepo, userRepo)
		Expect(testcmd.CommandDidPassRequirements).To(BeFalse())

		requirementsFactory.LoginSuccess = true
		callSetSpaceRole(args, requirementsFactory, spaceRepo, userRepo)
		Expect(testcmd.CommandDidPassRequirements).To(BeTrue())

		Expect(requirementsFactory.UserUsername).To(Equal("username"))
		Expect(requirementsFactory.OrganizationName).To(Equal("org"))
	})

	It("TestSetSpaceRole", func() {
		org := models.Organization{}
		org.Guid = "my-org-guid"
		org.Name = "my-org"

		args := []string{"some-user", "some-org", "some-space", "SpaceManager"}

		requirementsFactory, spaceRepo, userRepo := getSetSpaceRoleDeps()
		requirementsFactory.LoginSuccess = true
		requirementsFactory.UserFields = models.UserFields{}
		requirementsFactory.UserFields.Guid = "my-user-guid"
		requirementsFactory.UserFields.Username = "my-user"
		requirementsFactory.Organization = org

		spaceRepo.FindByNameInOrgSpace = models.Space{}
		spaceRepo.FindByNameInOrgSpace.Guid = "my-space-guid"
		spaceRepo.FindByNameInOrgSpace.Name = "my-space"
		spaceRepo.FindByNameInOrgSpace.Organization = org.OrganizationFields

		ui := callSetSpaceRole(args, requirementsFactory, spaceRepo, userRepo)

		Expect(ui.Outputs).To(ContainSubstrings(
			[]string{"Assigning role ", "SpaceManager", "my-user", "my-org", "my-space", "current-user"},
			[]string{"OK"},
		))

		Expect(spaceRepo.FindByNameInOrgName).To(Equal("some-space"))
		Expect(spaceRepo.FindByNameInOrgOrgGuid).To(Equal("my-org-guid"))

		Expect(userRepo.SetSpaceRoleUserGuid).To(Equal("my-user-guid"))
		Expect(userRepo.SetSpaceRoleSpaceGuid).To(Equal("my-space-guid"))
		Expect(userRepo.SetSpaceRoleRole).To(Equal(models.SPACE_MANAGER))
	})
})
<commit_msg>Cleanup set space role test<commit_after>package user_test

import (
	. "github.com/cloudfoundry/cli/cf/commands/user"
	"github.com/cloudfoundry/cli/cf/configuration"
	"github.com/cloudfoundry/cli/cf/models"
	testapi "github.com/cloudfoundry/cli/testhelpers/api"
	testcmd "github.com/cloudfoundry/cli/testhelpers/commands"
	testconfig "github.com/cloudfoundry/cli/testhelpers/configuration"
	testreq "github.com/cloudfoundry/cli/testhelpers/requirements"
	testterm "github.com/cloudfoundry/cli/testhelpers/terminal"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	. "github.com/cloudfoundry/cli/testhelpers/matchers"
)

var _ = Describe("set-space-role command", func() {
	var (
		ui                  *testterm.FakeUI
		requirementsFactory *testreq.FakeReqFactory
		spaceRepo           *testapi.FakeSpaceRepository
		userRepo            *testapi.FakeUserRepository
		configRepo          configuration.ReadWriter
	)

	BeforeEach(func() {
		configRepo = testconfig.NewRepositoryWithDefaults()
		accessToken, err := testconfig.EncodeAccessToken(configuration.TokenInfo{Username: "current-user"})
		Expect(err).NotTo(HaveOccurred())
		configRepo.SetAccessToken(accessToken)

		ui = &testterm.FakeUI{}
		requirementsFactory = &testreq.FakeReqFactory{}
		spaceRepo = &testapi.FakeSpaceRepository{}
		userRepo = &testapi.FakeUserRepository{}
	})

	runCommand := func(args ...string) {
		testcmd.RunCommand(NewSetSpaceRole(ui, configRepo, spaceRepo, userRepo), args, requirementsFactory)
	}

	It("fails with usage when not provided exactly four args", func() {
		runCommand("foo", "bar", "baz")
		Expect(ui.FailedWithUsage).To(BeTrue())
	})

	It("does not fail with usage when provided four args", func() {
		runCommand("whatever", "these", "are", "args")
		Expect(ui.FailedWithUsage).To(BeFalse())
	})

	Describe("requirements", func() {
		It("fails when not logged in", func() {
			runCommand("username", "org", "space", "role")

			Expect(testcmd.CommandDidPassRequirements).To(BeFalse())
		})

		It("succeeds when logged in", func() {
			requirementsFactory.LoginSuccess = true
			runCommand("username", "org", "space", "role")

			Expect(testcmd.CommandDidPassRequirements).To(BeTrue())
			Expect(requirementsFactory.UserUsername).To(Equal("username"))
			Expect(requirementsFactory.OrganizationName).To(Equal("org"))
		})
	})

	Context("when logged in", func() {
		BeforeEach(func() {
			requirementsFactory.LoginSuccess = true

			org := models.Organization{}
			org.Guid = "my-org-guid"
			org.Name = "my-org"

			requirementsFactory.UserFields = models.UserFields{Guid: "my-user-guid", Username: "my-user"}
			requirementsFactory.Organization = org

			spaceRepo.FindByNameInOrgSpace = models.Space{}
			spaceRepo.FindByNameInOrgSpace.Guid = "my-space-guid"
			spaceRepo.FindByNameInOrgSpace.Name = "my-space"
			spaceRepo.FindByNameInOrgSpace.Organization = org.OrganizationFields
		})

		It("sets the given space role on the given user", func() {
			runCommand("some-user", "some-org", "some-space", "SpaceManager")

			Expect(ui.Outputs).To(ContainSubstrings(
				[]string{"Assigning role ", "SpaceManager", "my-user", "my-org", "my-space", "current-user"},
				[]string{"OK"},
			))

			Expect(spaceRepo.FindByNameInOrgName).To(Equal("some-space"))
			Expect(spaceRepo.FindByNameInOrgOrgGuid).To(Equal("my-org-guid"))

			Expect(userRepo.SetSpaceRoleUserGuid).To(Equal("my-user-guid"))
			Expect(userRepo.SetSpaceRoleSpaceGuid).To(Equal("my-space-guid"))
			Expect(userRepo.SetSpaceRoleRole).To(Equal(models.SPACE_MANAGER))
		})
	})
})
<|endoftext|>


<commit_before>package packet

import (
	"bytes"
	"strconv"
	"strings"
	"testing"
)

func TestPacketEncoding(t *testing.T) {
	cases := []struct {
		decoded []byte // excluding CRC byte
		encoded []byte
	}{
		{parseBytes("A7 12 89 86 5D 00"), parseBytes("A9 6C 72 69 96 A6 94 D5 55 2C E5")},
		{parseBytes("A7 12 89 86 06 00"), parseBytes("A9 6C 72 69 96 A6 56 65 55 C6 55")},
		{parseBytes("A7 12 89 86 15 09"), parseBytes("A9 6C 72 69 96 A6 C6 55 59 96 65")},
		{parseBytes("A7 12 89 86 8D 00"), parseBytes("A9 6C 72 69 96 A6 68 D5 55 2D 55")},
		{parseBytes("A7 12 89 86 8D 09 03 37 32 32 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00"), parseBytes("A9 6C 72 69 96 A6 68 D5 59 56 38 D6 8F 28 F2 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 8D 95")},
		{parseBytes("A8 0F 25 C1 23 0D 19 1C 50 00 8F 00 90 00 34 34 99 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00"), parseBytes("A9 A5 5C CA 5B 31 CA 35 4D C5 9C 6C 95 55 55 69 C5 55 65 55 55 8F 48 F4 65 95 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 34 E3 96")},
	}
	for _, c := range cases {
		t.Run("encode", func(t *testing.T) {
			result := Encode(c.decoded)
			if !bytes.Equal(result, c.encoded) {
				t.Errorf("Encode(% X) == % X, want % X", c.decoded, result, c.encoded)
			}
		})
		t.Run("decode", func(t *testing.T) {
			result, err := Decode(c.encoded)
			if err != nil {
				t.Errorf("Decode(% X) == %v, want % X", c.encoded, err, c.decoded)
				return
			}
			if !bytes.Equal(result, c.decoded) {
				t.Errorf("Decode(% X) == % X, want % X", c.encoded, result, c.decoded)
			}
		})
	}
}

func parseBytes(hex string) []byte {
	fields := strings.Fields(hex)
	data := make([]byte, len(fields))
	for i, s := range fields {
		b, err := strconv.ParseUint(string(s), 16, 8)
		if err != nil {
			panic(err)
		}
		data[i] = byte(b)
	}
	return data
}
<commit_msg>Simplify parseBytes function and make spaces optional<commit_after>package packet

import (
	"bytes"
	"fmt"
	"strings"
	"testing"
)

func TestPacketEncoding(t *testing.T) {
	cases := []struct {
		decoded []byte // excluding CRC byte
		encoded []byte
	}{
		{parseBytes("A7 12 89 86 5D 00"), parseBytes("A9 6C 72 69 96 A6 94 D5 55 2C E5")},
		{parseBytes("A7 12 89 86 06 00"), parseBytes("A9 6C 72 69 96 A6 56 65 55 C6 55")},
		{parseBytes("A7 12 89 86 15 09"), parseBytes("A9 6C 72 69 96 A6 C6 55 59 96 65")},
		{parseBytes("A7 12 89 86 8D 00"), parseBytes("A9 6C 72 69 96 A6 68 D5 55 2D 55")},
		{parseBytes("A7 12 89 86 8D 09 03 37 32 32 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00"), parseBytes("A9 6C 72 69 96 A6 68 D5 59 56 38 D6 8F 28 F2 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 8D 95")},
		{parseBytes("A8 0F 25 C1 23 0D 19 1C 50 00 8F 00 90 00 34 34 99 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00"), parseBytes("A9 A5 5C CA 5B 31 CA 35 4D C5 9C 6C 95 55 55 69 C5 55 65 55 55 8F 48 F4 65 95 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 55 34 E3 96")},
	}
	for _, c := range cases {
		t.Run("encode", func(t *testing.T) {
			result := Encode(c.decoded)
			if !bytes.Equal(result, c.encoded) {
				t.Errorf("Encode(% X) == % X, want % X", c.decoded, result, c.encoded)
			}
		})
		t.Run("decode", func(t *testing.T) {
			result, err := Decode(c.encoded)
			if err != nil {
				t.Errorf("Decode(% X) == %v, want % X", c.encoded, err, c.decoded)
				return
			}
			if !bytes.Equal(result, c.decoded) {
				t.Errorf("Decode(% X) == % X, want % X", c.encoded, result, c.decoded)
			}
		})
	}
}

func parseBytes(hex string) []byte {
	r := strings.NewReader(hex)
	var data []byte
	for {
		var b byte
		n, err := fmt.Fscanf(r, "%02x", &b)
		if n == 0 {
			break
		}
		if err != nil {
			panic(err)
		}
		data = append(data, b)
	}
	return data
}
<|endoftext|>


<commit_before>package quic

import (
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha256"
	"fmt"
	"hash"
	"log"
	"net"
	"sync"
	"time"

	"github.com/lucas-clemente/quic-go/internal/protocol"
	"github.com/lucas-clemente/quic-go/internal/utils"
	"github.com/lucas-clemente/quic-go/internal/wire"
	"github.com/lucas-clemente/quic-go/logging"
)

type statelessResetErr struct {
	token protocol.StatelessResetToken
}

func (e statelessResetErr) Error() string {
	return fmt.Sprintf("received a stateless reset with token %x", e.token)
}

// The packetHandlerMap stores packetHandlers, identified by connection ID.
// It is used:
// * by the server to store sessions
// * when multiplexing outgoing connections to store clients
type packetHandlerMap struct {
	mutex sync.RWMutex

	conn      connection
	connIDLen int

	handlers    map[string] /* string(ConnectionID)*/ packetHandler
	resetTokens map[protocol.StatelessResetToken] /* stateless reset token */ packetHandler
	server      unknownPacketHandler

	listening chan struct{} // is closed when listen returns
	closed    bool

	deleteRetiredSessionsAfter time.Duration

	statelessResetEnabled bool
	statelessResetMutex   sync.Mutex
	statelessResetHasher  hash.Hash

	tracer logging.Tracer
	logger utils.Logger
}

var _ packetHandlerManager = &packetHandlerMap{}

func setReceiveBuffer(c net.PacketConn, logger utils.Logger) {
	conn, ok := c.(interface{ SetReadBuffer(int) error })
	if !ok {
		logger.Debugf("Connection doesn't allow setting of receive buffer size")
		return
	}
	size, err := inspectReadBuffer(c)
	if err != nil {
		log.Printf("Failed to determine receive buffer size: %s", err)
		return
	}
	if size >= protocol.DesiredReceiveBufferSize {
		logger.Debugf("Conn has receive buffer of %d kiB (wanted: at least %d kiB)", size/1024, protocol.DesiredReceiveBufferSize/1024)
	}
	if err := conn.SetReadBuffer(protocol.DesiredReceiveBufferSize); err != nil {
		log.Printf("Failed to increase receive buffer size: %s\n", err)
		return
	}
	newSize, err := inspectReadBuffer(c)
	if err != nil {
		log.Printf("Failed to determine receive buffer size: %s", err)
		return
	}
	if newSize == size {
		log.Printf("Failed to determine receive buffer size: %s", err)
		return
	}
	if newSize < protocol.DesiredReceiveBufferSize {
		log.Printf("Failed to sufficiently increase receive buffer size. Was: %d kiB, wanted: %d kiB, got: %d kiB.", size/1024, protocol.DesiredReceiveBufferSize/1024, newSize/1024)
		return
	}
	logger.Debugf("Increased receive buffer size to %d kiB", newSize/1024)
}

func newPacketHandlerMap(
	c net.PacketConn,
	connIDLen int,
	statelessResetKey []byte,
	tracer logging.Tracer,
	logger utils.Logger,
) (packetHandlerManager, error) {
	setReceiveBuffer(c, logger)
	conn, err := wrapConn(c)
	if err != nil {
		return nil, err
	}
	m := &packetHandlerMap{
		conn:                       conn,
		connIDLen:                  connIDLen,
		listening:                  make(chan struct{}),
		handlers:                   make(map[string]packetHandler),
		resetTokens:                make(map[protocol.StatelessResetToken]packetHandler),
		deleteRetiredSessionsAfter: protocol.RetiredConnectionIDDeleteTimeout,
		statelessResetEnabled:      len(statelessResetKey) > 0,
		statelessResetHasher:       hmac.New(sha256.New, statelessResetKey),
		tracer:                     tracer,
		logger:                     logger,
	}
	go m.listen()

	if logger.Debug() {
		go m.logUsage()
	}
	return m, nil
}

func (h *packetHandlerMap) logUsage() {
	ticker := time.NewTicker(2 * time.Second)
	var printedZero bool
	for {
		select {
		case <-h.listening:
			return
		case <-ticker.C:
		}

		h.mutex.Lock()
		numHandlers := len(h.handlers)
		numTokens := len(h.resetTokens)
		h.mutex.Unlock()
		// If the number tracked handlers and tokens is zero, only print it a single time.
		hasZero := numHandlers == 0 && numTokens == 0
		if !hasZero || (hasZero && !printedZero) {
			h.logger.Debugf("Tracking %d connection IDs and %d reset tokens.\n", numHandlers, numTokens)
			printedZero = false
			if hasZero {
				printedZero = true
			}
		}
	}
}

func (h *packetHandlerMap) Add(id protocol.ConnectionID, handler packetHandler) bool /* was added */ {
	sid := string(id)

	h.mutex.Lock()
	defer h.mutex.Unlock()

	if _, ok := h.handlers[sid]; ok {
		h.logger.Debugf("Not adding connection ID %s, as it already exists.", id)
		return false
	}
	h.handlers[sid] = handler
	h.logger.Debugf("Adding connection ID %s.", id)
	return true
}

func (h *packetHandlerMap) AddWithConnID(clientDestConnID, newConnID protocol.ConnectionID, fn func() packetHandler) bool {
	sid := string(clientDestConnID)
	h.mutex.Lock()
	defer h.mutex.Unlock()

	if _, ok := h.handlers[sid]; ok {
		h.logger.Debugf("Not adding connection ID %s for a new session, as it already exists.", clientDestConnID)
		return false
	}

	sess := fn()
	h.handlers[sid] = sess
	h.handlers[string(newConnID)] = sess
	h.logger.Debugf("Adding connection IDs %s and %s for a new session.", clientDestConnID, newConnID)
	return true
}

func (h *packetHandlerMap) Remove(id protocol.ConnectionID) {
	h.mutex.Lock()
	delete(h.handlers, string(id))
	h.mutex.Unlock()
	h.logger.Debugf("Removing connection ID %s.", id)
}

func (h *packetHandlerMap) Retire(id protocol.ConnectionID) {
	h.logger.Debugf("Retiring connection ID %s in %s.", id, h.deleteRetiredSessionsAfter)
	time.AfterFunc(h.deleteRetiredSessionsAfter, func() {
		h.mutex.Lock()
		delete(h.handlers, string(id))
		h.mutex.Unlock()
		h.logger.Debugf("Removing connection ID %s after it has been retired.", id)
	})
}

func (h *packetHandlerMap) ReplaceWithClosed(id protocol.ConnectionID, handler packetHandler) {
	h.mutex.Lock()
	h.handlers[string(id)] = handler
	h.mutex.Unlock()
	h.logger.Debugf("Replacing session for connection ID %s with a closed session.", id)

	time.AfterFunc(h.deleteRetiredSessionsAfter, func() {
		h.mutex.Lock()
		handler.shutdown()
		delete(h.handlers, string(id))
		h.mutex.Unlock()
		h.logger.Debugf("Removing connection ID %s for a closed session after it has been retired.", id)
	})
}

func (h *packetHandlerMap) AddResetToken(token protocol.StatelessResetToken, handler packetHandler) {
	h.mutex.Lock()
	h.resetTokens[token] = handler
	h.mutex.Unlock()
}

func (h *packetHandlerMap) RemoveResetToken(token protocol.StatelessResetToken) {
	h.mutex.Lock()
	delete(h.resetTokens, token)
	h.mutex.Unlock()
}

func (h *packetHandlerMap) RetireResetToken(token protocol.StatelessResetToken) {
	time.AfterFunc(h.deleteRetiredSessionsAfter, func() {
		h.mutex.Lock()
		delete(h.resetTokens, token)
		h.mutex.Unlock()
	})
}

func (h *packetHandlerMap) SetServer(s unknownPacketHandler) {
	h.mutex.Lock()
	h.server = s
	h.mutex.Unlock()
}

func (h *packetHandlerMap) CloseServer() {
	h.mutex.Lock()
	if h.server == nil {
		h.mutex.Unlock()
		return
	}
	h.server = nil
	var wg sync.WaitGroup
	for _, handler := range h.handlers {
		if handler.getPerspective() == protocol.PerspectiveServer {
			wg.Add(1)
			go func(handler packetHandler) {
				// blocks until the CONNECTION_CLOSE has been sent and the run-loop has stopped
				handler.shutdown()
				wg.Done()
			}(handler)
		}
	}
	h.mutex.Unlock()
	wg.Wait()
}

// Destroy the underlying connection and wait until listen() has returned.
// It does not close active sessions.
func (h *packetHandlerMap) Destroy() error {
	if err := h.conn.Close(); err != nil {
		return err
	}
	<-h.listening // wait until listening returns
	return nil
}

func (h *packetHandlerMap) close(e error) error {
	h.mutex.Lock()
	if h.closed {
		h.mutex.Unlock()
		return nil
	}

	var wg sync.WaitGroup
	for _, handler := range h.handlers {
		wg.Add(1)
		go func(handler packetHandler) {
			handler.destroy(e)
			wg.Done()
		}(handler)
	}

	if h.server != nil {
		h.server.setCloseError(e)
	}
	h.closed = true
	h.mutex.Unlock()
	wg.Wait()
	return getMultiplexer().RemoveConn(h.conn)
}

func (h *packetHandlerMap) listen() {
	defer close(h.listening)
	for {
		p, err := h.conn.ReadPacket()
		if nerr, ok := err.(net.Error); ok && nerr.Temporary() {
			h.logger.Debugf("Temporary error reading from conn: %w", err)
			continue
		}
		if err != nil {
			h.close(err)
			return
		}
		h.handlePacket(p)
	}
}

func (h *packetHandlerMap) handlePacket(p *receivedPacket) {
	connID, err := wire.ParseConnectionID(p.data, h.connIDLen)
	if err != nil {
		h.logger.Debugf("error parsing connection ID on packet from %s: %s", p.remoteAddr, err)
		if h.tracer != nil {
			h.tracer.DroppedPacket(p.remoteAddr, logging.PacketTypeNotDetermined, p.Size(), logging.PacketDropHeaderParseError)
		}
		p.buffer.MaybeRelease()
		return
	}

	h.mutex.RLock()
	defer h.mutex.RUnlock()

	if isStatelessReset := h.maybeHandleStatelessReset(p.data); isStatelessReset {
		return
	}

	handler, handlerFound := h.handlers[string(connID)]

	if handlerFound { // existing session
		handler.handlePacket(p)
		return
	}
	if p.data[0]&0x80 == 0 {
		go h.maybeSendStatelessReset(p, connID)
		return
	}
	if h.server == nil { // no server set
		h.logger.Debugf("received a packet with an unexpected connection ID %s", connID)
		return
	}
	h.server.handlePacket(p)
}

func (h *packetHandlerMap) maybeHandleStatelessReset(data []byte) bool {
	// stateless resets are always short header packets
	if data[0]&0x80 != 0 {
		return false
	}
	if len(data) < 17 /* type byte + 16 bytes for the reset token */ {
		return false
	}

	var token protocol.StatelessResetToken
	copy(token[:], data[len(data)-16:])
	if sess, ok := h.resetTokens[token]; ok {
		h.logger.Debugf("Received a stateless reset with token %#x. Closing session.", token)
		go sess.destroy(statelessResetErr{token: token})
		return true
	}
	return false
}

func (h *packetHandlerMap) GetStatelessResetToken(connID protocol.ConnectionID) protocol.StatelessResetToken {
	var token protocol.StatelessResetToken
	if !h.statelessResetEnabled {
		// Return a random stateless reset token.
		// This token will be sent in the server's transport parameters.
		// By using a random token, an off-path attacker won't be able to disrupt the connection.
		rand.Read(token[:])
		return token
	}
	h.statelessResetMutex.Lock()
	h.statelessResetHasher.Write(connID.Bytes())
	copy(token[:], h.statelessResetHasher.Sum(nil))
	h.statelessResetHasher.Reset()
	h.statelessResetMutex.Unlock()
	return token
}

func (h *packetHandlerMap) maybeSendStatelessReset(p *receivedPacket, connID protocol.ConnectionID) {
	defer p.buffer.Release()
	if !h.statelessResetEnabled {
		return
	}
	// Don't send a stateless reset in response to very small packets.
	// This includes packets that could be stateless resets.
	if len(p.data) <= protocol.MinStatelessResetSize {
		return
	}
	token := h.GetStatelessResetToken(connID)
	h.logger.Debugf("Sending stateless reset to %s (connection ID: %s). Token: %#x", p.remoteAddr, connID, token)
	data := make([]byte, protocol.MinStatelessResetSize-16, protocol.MinStatelessResetSize)
	rand.Read(data)
	data[0] = (data[0] & 0x7f) | 0x40
	data = append(data, token[:]...)
	if _, err := h.conn.WriteTo(data, p.remoteAddr); err != nil {
		h.logger.Debugf("Error sending Stateless Reset: %s", err)
	}
}
<commit_msg>replace the RWMutex with a Mutex in the packet handler map<commit_after>package quic

import (
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha256"
	"fmt"
	"hash"
	"log"
	"net"
	"sync"
	"time"

	"github.com/lucas-clemente/quic-go/internal/protocol"
	"github.com/lucas-clemente/quic-go/internal/utils"
	"github.com/lucas-clemente/quic-go/internal/wire"
	"github.com/lucas-clemente/quic-go/logging"
)

type statelessResetErr struct {
	token protocol.StatelessResetToken
}

func (e statelessResetErr) Error() string {
	return fmt.Sprintf("received a stateless reset with token %x", e.token)
}

// The packetHandlerMap stores packetHandlers, identified by connection ID.
// It is used:
// * by the server to store sessions
// * when multiplexing outgoing connections to store clients
type packetHandlerMap struct {
	mutex sync.Mutex

	conn      connection
	connIDLen int

	handlers    map[string] /* string(ConnectionID)*/ packetHandler
	resetTokens map[protocol.StatelessResetToken] /* stateless reset token */ packetHandler
	server      unknownPacketHandler

	listening chan struct{} // is closed when listen returns
	closed    bool

	deleteRetiredSessionsAfter time.Duration

	statelessResetEnabled bool
	statelessResetMutex   sync.Mutex
	statelessResetHasher  hash.Hash

	tracer logging.Tracer
	logger utils.Logger
}

var _ packetHandlerManager = &packetHandlerMap{}

func setReceiveBuffer(c net.PacketConn, logger utils.Logger) {
	conn, ok := c.(interface{ SetReadBuffer(int) error })
	if !ok {
		logger.Debugf("Connection doesn't allow setting of receive buffer size")
		return
	}
	size, err := inspectReadBuffer(c)
	if err != nil {
		log.Printf("Failed to determine receive buffer size: %s", err)
		return
	}
	if size >= protocol.DesiredReceiveBufferSize {
		logger.Debugf("Conn has receive buffer of %d kiB (wanted: at least %d kiB)", size/1024, protocol.DesiredReceiveBufferSize/1024)
	}
	if err := conn.SetReadBuffer(protocol.DesiredReceiveBufferSize); err != nil {
		log.Printf("Failed to increase receive buffer size: %s\n", err)
		return
	}
	newSize, err := inspectReadBuffer(c)
	if err != nil {
		log.Printf("Failed to determine receive buffer size: %s", err)
		return
	}
	if newSize == size {
		log.Printf("Failed to determine receive buffer size: %s", err)
		return
	}
	if newSize < protocol.DesiredReceiveBufferSize {
		log.Printf("Failed to sufficiently increase receive buffer size. Was: %d kiB, wanted: %d kiB, got: %d kiB.", size/1024, protocol.DesiredReceiveBufferSize/1024, newSize/1024)
		return
	}
	logger.Debugf("Increased receive buffer size to %d kiB", newSize/1024)
}

func newPacketHandlerMap(
	c net.PacketConn,
	connIDLen int,
	statelessResetKey []byte,
	tracer logging.Tracer,
	logger utils.Logger,
) (packetHandlerManager, error) {
	setReceiveBuffer(c, logger)
	conn, err := wrapConn(c)
	if err != nil {
		return nil, err
	}
	m := &packetHandlerMap{
		conn:                       conn,
		connIDLen:                  connIDLen,
		listening:                  make(chan struct{}),
		handlers:                   make(map[string]packetHandler),
		resetTokens:                make(map[protocol.StatelessResetToken]packetHandler),
		deleteRetiredSessionsAfter: protocol.RetiredConnectionIDDeleteTimeout,
		statelessResetEnabled:      len(statelessResetKey) > 0,
		statelessResetHasher:       hmac.New(sha256.New, statelessResetKey),
		tracer:                     tracer,
		logger:                     logger,
	}
	go m.listen()

	if logger.Debug() {
		go m.logUsage()
	}
	return m, nil
}

func (h *packetHandlerMap) logUsage() {
	ticker := time.NewTicker(2 * time.Second)
	var printedZero bool
	for {
		select {
		case <-h.listening:
			return
		case <-ticker.C:
		}

		h.mutex.Lock()
		numHandlers := len(h.handlers)
		numTokens := len(h.resetTokens)
		h.mutex.Unlock()
		// If the number tracked handlers and tokens is zero, only print it a single time.
		hasZero := numHandlers == 0 && numTokens == 0
		if !hasZero || (hasZero && !printedZero) {
			h.logger.Debugf("Tracking %d connection IDs and %d reset tokens.\n", numHandlers, numTokens)
			printedZero = false
			if hasZero {
				printedZero = true
			}
		}
	}
}

func (h *packetHandlerMap) Add(id protocol.ConnectionID, handler packetHandler) bool /* was added */ {
	sid := string(id)

	h.mutex.Lock()
	defer h.mutex.Unlock()

	if _, ok := h.handlers[sid]; ok {
		h.logger.Debugf("Not adding connection ID %s, as it already exists.", id)
		return false
	}
	h.handlers[sid] = handler
	h.logger.Debugf("Adding connection ID %s.", id)
	return true
}

func (h *packetHandlerMap) AddWithConnID(clientDestConnID, newConnID protocol.ConnectionID, fn func() packetHandler) bool {
	sid := string(clientDestConnID)
	h.mutex.Lock()
	defer h.mutex.Unlock()

	if _, ok := h.handlers[sid]; ok {
		h.logger.Debugf("Not adding connection ID %s for a new session, as it already exists.", clientDestConnID)
		return false
	}

	sess := fn()
	h.handlers[sid] = sess
	h.handlers[string(newConnID)] = sess
	h.logger.Debugf("Adding connection IDs %s and %s for a new session.", clientDestConnID, newConnID)
	return true
}

func (h *packetHandlerMap) Remove(id protocol.ConnectionID) {
	h.mutex.Lock()
	delete(h.handlers, string(id))
	h.mutex.Unlock()
	h.logger.Debugf("Removing connection ID %s.", id)
}

func (h *packetHandlerMap) Retire(id protocol.ConnectionID) {
	h.logger.Debugf("Retiring connection ID %s in %s.", id, h.deleteRetiredSessionsAfter)
	time.AfterFunc(h.deleteRetiredSessionsAfter, func() {
		h.mutex.Lock()
		delete(h.handlers, string(id))
		h.mutex.Unlock()
		h.logger.Debugf("Removing connection ID %s after it has been retired.", id)
	})
}

func (h *packetHandlerMap) ReplaceWithClosed(id protocol.ConnectionID, handler packetHandler) {
	h.mutex.Lock()
	h.handlers[string(id)] = handler
	h.mutex.Unlock()
	h.logger.Debugf("Replacing session for connection ID %s with a closed session.", id)

	time.AfterFunc(h.deleteRetiredSessionsAfter, func() {
		h.mutex.Lock()
		handler.shutdown()
		delete(h.handlers, string(id))
		h.mutex.Unlock()
		h.logger.Debugf("Removing connection ID %s for a closed session after it has been retired.", id)
	})
}

func (h *packetHandlerMap) AddResetToken(token protocol.StatelessResetToken, handler packetHandler) {
	h.mutex.Lock()
	h.resetTokens[token] = handler
	h.mutex.Unlock()
}

func (h *packetHandlerMap) RemoveResetToken(token protocol.StatelessResetToken) {
	h.mutex.Lock()
	delete(h.resetTokens, token)
	h.mutex.Unlock()
}

func (h *packetHandlerMap) RetireResetToken(token protocol.StatelessResetToken) {
	time.AfterFunc(h.deleteRetiredSessionsAfter, func() {
		h.mutex.Lock()
		delete(h.resetTokens, token)
		h.mutex.Unlock()
	})
}

func (h *packetHandlerMap) SetServer(s unknownPacketHandler) {
	h.mutex.Lock()
	h.server = s
	h.mutex.Unlock()
}

func (h *packetHandlerMap) CloseServer() {
	h.mutex.Lock()
	if h.server == nil {
		h.mutex.Unlock()
		return
	}
	h.server = nil
	var wg sync.WaitGroup
	for _, handler := range h.handlers {
		if handler.getPerspective() == protocol.PerspectiveServer {
			wg.Add(1)
			go func(handler packetHandler) {
				// blocks until the CONNECTION_CLOSE has been sent and the run-loop has stopped
				handler.shutdown()
				wg.Done()
			}(handler)
		}
	}
	h.mutex.Unlock()
	wg.Wait()
}

// Destroy the underlying connection and wait until listen() has returned.
// It does not close active sessions.
func (h *packetHandlerMap) Destroy() error {
	if err := h.conn.Close(); err != nil {
		return err
	}
	<-h.listening // wait until listening returns
	return nil
}

func (h *packetHandlerMap) close(e error) error {
	h.mutex.Lock()
	if h.closed {
		h.mutex.Unlock()
		return nil
	}

	var wg sync.WaitGroup
	for _, handler := range h.handlers {
		wg.Add(1)
		go func(handler packetHandler) {
			handler.destroy(e)
			wg.Done()
		}(handler)
	}

	if h.server != nil {
		h.server.setCloseError(e)
	}
	h.closed = true
	h.mutex.Unlock()
	wg.Wait()
	return getMultiplexer().RemoveConn(h.conn)
}

func (h *packetHandlerMap) listen() {
	defer close(h.listening)
	for {
		p, err := h.conn.ReadPacket()
		if nerr, ok := err.(net.Error); ok && nerr.Temporary() {
			h.logger.Debugf("Temporary error reading from conn: %w", err)
			continue
		}
		if err != nil {
			h.close(err)
			return
		}
		h.handlePacket(p)
	}
}

func (h *packetHandlerMap) handlePacket(p *receivedPacket) {
	connID, err := wire.ParseConnectionID(p.data, h.connIDLen)
	if err != nil {
		h.logger.Debugf("error parsing connection ID on packet from %s: %s", p.remoteAddr, err)
		if h.tracer != nil {
			h.tracer.DroppedPacket(p.remoteAddr, logging.PacketTypeNotDetermined, p.Size(), logging.PacketDropHeaderParseError)
		}
		p.buffer.MaybeRelease()
		return
	}

	h.mutex.Lock()
	defer h.mutex.Unlock()

	if isStatelessReset := h.maybeHandleStatelessReset(p.data); isStatelessReset {
		return
	}

	if handler, ok := h.handlers[string(connID)]; ok { // existing session
		handler.handlePacket(p)
		return
	}
	if p.data[0]&0x80 == 0 {
		go h.maybeSendStatelessReset(p, connID)
		return
	}
	if h.server == nil { // no server set
		h.logger.Debugf("received a packet with an unexpected connection ID %s", connID)
		return
	}
	h.server.handlePacket(p)
}

func (h *packetHandlerMap) maybeHandleStatelessReset(data []byte) bool {
	// stateless resets are always short header packets
	if data[0]&0x80 != 0 {
		return false
	}
	if len(data) < 17 /* type byte + 16 bytes for the reset token */ {
		return false
	}

	var token protocol.StatelessResetToken
	copy(token[:], data[len(data)-16:])
	if sess, ok := h.resetTokens[token]; ok {
		h.logger.Debugf("Received a stateless reset with token %#x. Closing session.", token)
		go sess.destroy(statelessResetErr{token: token})
		return true
	}
	return false
}

func (h *packetHandlerMap) GetStatelessResetToken(connID protocol.ConnectionID) protocol.StatelessResetToken {
	var token protocol.StatelessResetToken
	if !h.statelessResetEnabled {
		// Return a random stateless reset token.
		// This token will be sent in the server's transport parameters.
		// By using a random token, an off-path attacker won't be able to disrupt the connection.
		rand.Read(token[:])
		return token
	}
	h.statelessResetMutex.Lock()
	h.statelessResetHasher.Write(connID.Bytes())
	copy(token[:], h.statelessResetHasher.Sum(nil))
	h.statelessResetHasher.Reset()
	h.statelessResetMutex.Unlock()
	return token
}

func (h *packetHandlerMap) maybeSendStatelessReset(p *receivedPacket, connID protocol.ConnectionID) {
	defer p.buffer.Release()
	if !h.statelessResetEnabled {
		return
	}
	// Don't send a stateless reset in response to very small packets.
	// This includes packets that could be stateless resets.
	if len(p.data) <= protocol.MinStatelessResetSize {
		return
	}
	token := h.GetStatelessResetToken(connID)
	h.logger.Debugf("Sending stateless reset to %s (connection ID: %s). Token: %#x", p.remoteAddr, connID, token)
	data := make([]byte, protocol.MinStatelessResetSize-16, protocol.MinStatelessResetSize)
	rand.Read(data)
	data[0] = (data[0] & 0x7f) | 0x40
	data = append(data, token[:]...)
	if _, err := h.conn.WriteTo(data, p.remoteAddr); err != nil {
		h.logger.Debugf("Error sending Stateless Reset: %s", err)
	}
}
<|endoftext|>


<commit_before>// Copyright 2012-2015 Samuel Stauffer. All rights reserved.
// Use of this source code is governed by a 3-clause BSD
// license that can be found in the LICENSE file.

package parser

import (
	"bytes"
	"encoding/json"
	"reflect"
	"strconv"
	"testing"
)

func TestServiceParsing(t *testing.T) {
	parser := &Parser{}
	thrift, err := parser.Parse(bytes.NewBuffer([]byte(`
		include "other.thrift"

		namespace go somepkg
		namespace python some.module123
		namespace python.py-twisted another

		const map<string,string> M1 = {"hello": "world", "goodnight": "moon"}
		const string S1 = "foo\"\tbar"
		const string S2 = 'foo\'\tbar'
		const list<i64> L = [1, 2, 3];

		union myUnion
		{
			1: double dbl = 1.1;
			2: string str = "2";
			3: i32 int32 = 3;
			4: i64 int64
				= 5;
		}

		enum Operation
		{
			ADD = 1,
			SUBTRACT = 2
		}

		enum NoNewLineBeforeBrace {
			ADD = 1,
			SUBTRACT = 2
		}

		service ServiceNAME extends SomeBase
		{
			# authenticate method
			// comment2
			/* some other
			   comments */
			string login(1:string password) throws (1:AuthenticationException authex),
			oneway void explode();
			blah something()
		}

		struct SomeStruct {
			1: double dbl = 1.2,
			2: optional string abc
		}

		struct NewLineBeforeBrace
		{
			1: double dbl = 1.2,
			2: optional string abc
		}`)))

	if err != nil {
		t.Fatalf("Service parsing failed with error %s", err.Error())
	}

	if thrift.Includes["other"] != "other.thrift" {
		t.Errorf("Include not parsed: %+v", thrift.Includes)
	}

	if c := thrift.Constants["M1"]; c == nil {
		t.Errorf("M1 constant missing")
	} else if c.Name != "M1" {
		t.Errorf("M1 name not M1, got '%s'", c.Name)
	} else if v, e := c.Type.String(), "map<string,string>"; v != e {
		t.Errorf("Expected type '%s' for M1, got '%s'", e, v)
	} else if _, ok := c.Value.([]KeyValue); !ok {
		t.Errorf("Expected []KeyValue value for M1, got %T", c.Value)
	}

	if c := thrift.Constants["S1"]; c == nil {
		t.Errorf("S1 constant missing")
	} else if v, e := c.Value.(string), "foo\"\tbar"; e != v {
		t.Errorf("Excepted %s for constnat S1, got %s", strconv.Quote(e), strconv.Quote(v))
	}
	if c := thrift.Constants["S2"]; c == nil {
		t.Errorf("S2 constant missing")
	} else if v, e := c.Value.(string), "foo'\tbar"; e != v {
		t.Errorf("Excepted %s for constnat S2, got %s", strconv.Quote(e), strconv.Quote(v))
	}

	expConst := &Constant{
		Name: "L",
		Type: &Type{
			Name:      "list",
			ValueType: &Type{Name: "i64"},
		},
		Value: []interface{}{int64(1), int64(2), int64(3)},
	}
	if c := thrift.Constants["L"]; c == nil {
		t.Errorf("L constant missing")
	} else if !reflect.DeepEqual(c, expConst) {
		t.Errorf("Expected for L:\n%s\ngot\n%s", pprint(expConst), pprint(c))
	}

	expectedStruct := &Struct{
		Name: "SomeStruct",
		Fields: []*Field{
			{
				ID:      1,
				Name:    "dbl",
				Default: 1.2,
				Type: &Type{
					Name: "double",
				},
			},
			{
				ID:       2,
				Name:     "abc",
				Optional: true,
				Type: &Type{
					Name: "string",
				},
			},
		},
	}
	if s := thrift.Structs["SomeStruct"]; s == nil {
		t.Errorf("SomeStruct missing")
	} else if !reflect.DeepEqual(s, expectedStruct) {
		t.Errorf("Expected\n%s\ngot\n%s", pprint(expectedStruct), pprint(s))
	}

	expectedUnion := &Struct{
		Name: "myUnion",
		Fields: []*Field{
			{
				ID:       1,
				Name:     "dbl",
				Default:  1.1,
				Optional: true,
				Type: &Type{
					Name: "double",
				},
			},
			{
				ID:       2,
				Name:     "str",
				Default:  "2",
				Optional: true,
				Type: &Type{
					Name: "string",
				},
			},
			{
				ID:       3,
				Name:     "int32",
				Default:  int64(3),
				Optional: true,
				Type: &Type{
					Name: "i32",
				},
			},
			{
				ID:       4,
				Name:     "int64",
				Default:  int64(5),
				Optional: true,
				Type: &Type{
					Name: "i64",
				},
			},
		},
	}
	if u := thrift.Unions["myUnion"]; u == nil {
		t.Errorf("myUnion missing")
	} else if !reflect.DeepEqual(u, expectedUnion) {
		t.Errorf("Expected\n%s\ngot\n%s", pprint(expectedUnion), pprint(u))
	}

	expectedEnum := &Enum{
		Name: "Operation",
		Values: map[string]*EnumValue{
			"ADD": &EnumValue{
				Name:  "ADD",
				Value: 1,
			},
			"SUBTRACT": &EnumValue{
				Name:  "SUBTRACT",
				Value: 2,
			},
		},
	}
	if e := thrift.Enums["Operation"]; e == nil {
		t.Errorf("enum Operation missing")
	} else if !reflect.DeepEqual(e, expectedEnum) {
		t.Errorf("Expected\n%s\ngot\n%s", pprint(expectedEnum), pprint(e))
	}

	if len(thrift.Services) != 1 {
		t.Fatalf("Parsing service returned %d services rather than 1 as expected", len(thrift.Services))
	}
	svc := thrift.Services["ServiceNAME"]
	if svc == nil || svc.Name != "ServiceNAME" {
		t.Fatalf("Parsing service expected to find 'ServiceNAME' rather than '%+v'", thrift.Services)
	} else if svc.Extends != "SomeBase" {
		t.Errorf("Expected extends 'SomeBase' got '%s'", svc.Extends)
	}

	expected := map[string]*Service{
		"ServiceNAME": &Service{
			Name:    "ServiceNAME",
			Extends: "SomeBase",
			Methods: map[string]*Method{
				"login": &Method{
					Name: "login",
					ReturnType: &Type{
						Name: "string",
					},
					Arguments: []*Field{
						&Field{
							ID:       1,
							Name:     "password",
							Optional: false,
							Type: &Type{
								Name: "string",
							},
						},
					},
					Exceptions: []*Field{
						&Field{
							ID:       1,
							Name:     "authex",
							Optional: true,
							Type: &Type{
								Name: "AuthenticationException",
							},
						},
					},
				},
				"explode": &Method{
					Name:       "explode",
					ReturnType: nil,
					Oneway:     true,
					Arguments:  []*Field{},
				},
			},
		},
	}
	for n, m := range expected["ServiceNAME"].Methods {
		if !reflect.DeepEqual(svc.Methods[n], m) {
			t.Fatalf("Parsing service returned method\n%s\ninstead of\n%s", pprint(svc.Methods[n]), pprint(m))
		}
	}
}

func TestParseConstant(t *testing.T) {
	parser := &Parser{}
	thrift, err := parser.Parse(bytes.NewBuffer([]byte(`
		const string C1 = "test"
		const string C2 = C1
		`)))
	if err != nil {
		t.Fatalf("Service parsing failed with error %s", err.Error())
	}

	expected := map[string]*Constant{
		"C1": &Constant{
			Name:  "C1",
			Type:  &Type{Name: "string"},
			Value: "test",
		},
		"C2": &Constant{
			Name:  "C2",
			Type:  &Type{Name: "string"},
			Value: Identifier("C1"),
		},
	}
	if got := thrift.Constants; !reflect.DeepEqual(expected, got) {
		t.Errorf("Unexpected constant parsing got\n%s\ninstead of\n%s", pprint(expected), pprint(got))
	}
}

// func TestParseFile(t *testing.T) {
// 	th, err := ParseFile("../testfiles/full.thrift")
// 	if err != nil {
// 		t.Fatal(err)
// 	}
// 	b, err := json.MarshalIndent(th, "", "    ")
// 	if err != nil {
// 		t.Fatal(err)
// 	}
// 	_ = b
// }

func pprint(v interface{}) string {
	b, err := json.MarshalIndent(v, "", "    ")
	if err != nil {
		panic(err)
	}
	return string(b)
}
<commit_msg>Enable TestParseFile for cassandra and HBase files<commit_after>// Copyright 2012-2015 Samuel Stauffer. All rights reserved.
// Use of this source code is governed by a 3-clause BSD
// license that can be found in the LICENSE file.

package parser

import (
	"bytes"
	"encoding/json"
	"path/filepath"
	"reflect"
	"strconv"
	"testing"
)

func TestServiceParsing(t *testing.T) {
	parser := &Parser{}
	thrift, err := parser.Parse(bytes.NewBuffer([]byte(`
		include "other.thrift"

		namespace go somepkg
		namespace python some.module123
		namespace python.py-twisted another

		const map<string,string> M1 = {"hello": "world", "goodnight": "moon"}
		const string S1 = "foo\"\tbar"
		const string S2 = 'foo\'\tbar'
		const list<i64> L = [1, 2, 3];

		union myUnion
		{
			1: double dbl = 1.1;
			2: string str = "2";
			3: i32 int32 = 3;
			4: i64 int64
				= 5;
		}

		enum Operation
		{
			ADD = 1,
			SUBTRACT = 2
		}

		enum NoNewLineBeforeBrace {
			ADD = 1,
			SUBTRACT = 2
		}

		service ServiceNAME extends SomeBase
		{
			# authenticate method
			// comment2
			/* some other
			   comments */
			string login(1:string password) throws (1:AuthenticationException authex),
			oneway void explode();
			blah something()
		}

		struct SomeStruct {
			1: double dbl = 1.2,
			2: optional string abc
		}

		struct NewLineBeforeBrace
		{
			1: double dbl = 1.2,
			2: optional string abc
		}`)))

	if err != nil {
		t.Fatalf("Service parsing failed with error %s", err.Error())
	}

	if thrift.Includes["other"] != "other.thrift" {
		t.Errorf("Include not parsed: %+v", thrift.Includes)
	}

	if c := thrift.Constants["M1"]; c == nil {
		t.Errorf("M1 constant missing")
	} else if c.Name != "M1" {
		t.Errorf("M1 name not M1, got '%s'", c.Name)
	} else if v, e := c.Type.String(), "map<string,string>"; v != e {
		t.Errorf("Expected type '%s' for M1, got '%s'", e, v)
	} else if _, ok := c.Value.([]KeyValue); !ok {
		t.Errorf("Expected []KeyValue value for M1, got %T", c.Value)
	}

	if c := thrift.Constants["S1"]; c == nil {
		t.Errorf("S1 constant missing")
	} else if v, e := c.Value.(string), "foo\"\tbar"; e != v {
		t.Errorf("Excepted %s for constnat S1, got %s", strconv.Quote(e), strconv.Quote(v))
	}
	if c := thrift.Constants["S2"]; c == nil {
		t.Errorf("S2 constant missing")
	} else if v, e := c.Value.(string), "foo'\tbar"; e != v {
		t.Errorf("Excepted %s for constnat S2, got %s", strconv.Quote(e), strconv.Quote(v))
	}

	expConst := &Constant{
		Name: "L",
		Type: &Type{
			Name:      "list",
			ValueType: &Type{Name: "i64"},
		},
		Value: []interface{}{int64(1), int64(2), int64(3)},
	}
	if c := thrift.Constants["L"]; c == nil {
		t.Errorf("L constant missing")
	} else if !reflect.DeepEqual(c, expConst) {
		t.Errorf("Expected for L:\n%s\ngot\n%s", pprint(expConst), pprint(c))
	}

	expectedStruct := &Struct{
		Name: "SomeStruct",
		Fields: []*Field{
			{
				ID:      1,
				Name:    "dbl",
				Default: 1.2,
				Type: &Type{
					Name: "double",
				},
			},
			{
				ID:       2,
				Name:     "abc",
				Optional: true,
				Type: &Type{
					Name: "string",
				},
			},
		},
	}
	if s := thrift.Structs["SomeStruct"]; s == nil {
		t.Errorf("SomeStruct missing")
	} else if !reflect.DeepEqual(s, expectedStruct) {
		t.Errorf("Expected\n%s\ngot\n%s", pprint(expectedStruct), pprint(s))
	}

	expectedUnion := &Struct{
		Name: "myUnion",
		Fields: []*Field{
			{
				ID:       1,
				Name:     "dbl",
				Default:  1.1,
				Optional: true,
				Type: &Type{
					Name: "double",
				},
			},
			{
				ID:       2,
				Name:     "str",
				Default:  "2",
				Optional: true,
				Type: &Type{
					Name: "string",
				},
			},
			{
				ID:       3,
				Name:     "int32",
				Default:  int64(3),
				Optional: true,
				Type: &Type{
					Name: "i32",
				},
			},
			{
				ID:       4,
				Name:     "int64",
				Default:  int64(5),
				Optional: true,
				Type: &Type{
					Name: "i64",
				},
			},
		},
	}
	if u := thrift.Unions["myUnion"]; u == nil {
		t.Errorf("myUnion missing")
	} else if !reflect.DeepEqual(u, expectedUnion) {
		t.Errorf("Expected\n%s\ngot\n%s", pprint(expectedUnion), pprint(u))
	}

	expectedEnum := &Enum{
		Name: "Operation",
		Values: map[string]*EnumValue{
			"ADD": &EnumValue{
				Name:  "ADD",
				Value: 1,
			},
			"SUBTRACT": &EnumValue{
				Name:  "SUBTRACT",
				Value: 2,
			},
		},
	}
	if e := thrift.Enums["Operation"]; e == nil {
		t.Errorf("enum Operation missing")
	} else if !reflect.DeepEqual(e, expectedEnum) {
		t.Errorf("Expected\n%s\ngot\n%s", pprint(expectedEnum), pprint(e))
	}

	if len(thrift.Services) != 1 {
		t.Fatalf("Parsing service returned %d services rather than 1 as expected", len(thrift.Services))
	}
	svc := thrift.Services["ServiceNAME"]
	if svc == nil || svc.Name != "ServiceNAME" {
		t.Fatalf("Parsing service expected to find 'ServiceNAME' rather than '%+v'", thrift.Services)
	} else if svc.Extends != "SomeBase" {
		t.Errorf("Expected extends 'SomeBase' got '%s'", svc.Extends)
	}

	expected := map[string]*Service{
		"ServiceNAME": &Service{
			Name:    "ServiceNAME",
			Extends: "SomeBase",
			Methods: map[string]*Method{
				"login": &Method{
					Name: "login",
					ReturnType: &Type{
						Name: "string",
					},
					Arguments: []*Field{
						&Field{
							ID:       1,
							Name:     "password",
							Optional: false,
							Type: &Type{
								Name: "string",
							},
						},
					},
					Exceptions: []*Field{
						&Field{
							ID:       1,
							Name:     "authex",
							Optional: true,
							Type: &Type{
								Name: "AuthenticationException",
							},
						},
					},
				},
				"explode": &Method{
					Name:       "explode",
					ReturnType: nil,
					Oneway:     true,
					Arguments:  []*Field{},
				},
			},
		},
	}
	for n, m := range expected["ServiceNAME"].Methods {
		if !reflect.DeepEqual(svc.Methods[n], m) {
			t.Fatalf("Parsing service returned method\n%s\ninstead of\n%s", pprint(svc.Methods[n]), pprint(m))
		}
	}
}

func TestParseConstant(t *testing.T) {
	parser := &Parser{}
	thrift, err := parser.Parse(bytes.NewBuffer([]byte(`
		const string C1 = "test"
		const string C2 = C1
		`)))
	if err != nil {
		t.Fatalf("Service parsing failed with error %s", err.Error())
	}

	expected := map[string]*Constant{
		"C1": &Constant{
			Name:  "C1",
			Type:  &Type{Name: "string"},
			Value: "test",
		},
		"C2": &Constant{
			Name:  "C2",
			Type:  &Type{Name: "string"},
			Value: Identifier("C1"),
		},
	}
	if got := thrift.Constants; !reflect.DeepEqual(expected, got) {
		t.Errorf("Unexpected constant parsing got\n%s\ninstead of\n%s", pprint(expected), pprint(got))
	}
}

func TestParseFiles(t *testing.T) {
	files := []string{
		"cassandra.thrift",
		"Hbase.thrift",
	}

	for _, f := range files {
		_, err := ParseFile(filepath.Join("../testfiles", f))
		if err != nil {
			t.Errorf("Failed to parse file %q: %v", f, err)
		}
	}
}

func pprint(v interface{}) string {
	b, err := json.MarshalIndent(v, "", "    ")
	if err != nil {
		panic(err)
	}
	return string(b)
}
<|endoftext|>


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

package zap

import (
	"fmt"
	"testing"

	"github.com/couchbase/vellum"
)

type enumTestEntry struct {
	key string
	val uint64
}

type enumTestWant struct {
	key string
	idx int
	val uint64
}

func TestEnumerator(t *testing.T) {
	tests := []struct {
		desc string
		in   [][]enumTestEntry
		want []enumTestWant
	}{
		{
			desc: "two non-empty enumerators with no duplicate keys",
			in: [][]enumTestEntry{
				[]enumTestEntry{
					{"a", 1},
					{"c", 3},
					{"e", 5},
				},
				[]enumTestEntry{
					{"b", 2},
					{"d", 4},
					{"f", 6},
				},
			},
			want: []enumTestWant{
				{"a", 0, 1},
				{"b", 1, 2},
				{"c", 0, 3},
				{"d", 1, 4},
				{"e", 0, 5},
				{"f", 1, 6},
			},
		},
		{
			desc: "two non-empty enumerators with duplicate keys",
			in: [][]enumTestEntry{
				[]enumTestEntry{
					{"a", 1},
					{"c", 3},
					{"e", 5},
				},
				[]enumTestEntry{
					{"a", 2},
					{"c", 4},
					{"e", 6},
				},
			},
			want: []enumTestWant{
				{"a", 0, 1},
				{"a", 1, 2},
				{"c", 0, 3},
				{"c", 1, 4},
				{"e", 0, 5},
				{"e", 1, 6},
			},
		},
		{
			desc: "first iterator is empty",
			in: [][]enumTestEntry{
				[]enumTestEntry{},
				[]enumTestEntry{
					{"a", 2},
					{"c", 4},
					{"e", 6},
				},
			},
			want: []enumTestWant{
				{"a", 1, 2},
				{"c", 1, 4},
				{"e", 1, 6},
			},
		},
		{
			desc: "last iterator is empty",
			in: [][]enumTestEntry{
				[]enumTestEntry{
					{"a", 1},
					{"c", 3},
					{"e", 5},
				},
				[]enumTestEntry{},
			},
			want: []enumTestWant{
				{"a", 0, 1},
				{"c", 0, 3},
				{"e", 0, 5},
			},
		},
		{
			desc: "two different length enumerators with duplicate keys",
			in: [][]enumTestEntry{
				[]enumTestEntry{
					{"a", 1},
					{"c", 3},
					{"e", 5},
				},
				[]enumTestEntry{
					{"a", 2},
					{"b", 4},
					{"d", 1000},
					{"e", 6},
				},
			},
			want: []enumTestWant{
				{"a", 0, 1},
				{"a", 1, 2},
				{"b", 1, 4},
				{"c", 0, 3},
				{"d", 1, 1000},
				{"e", 0, 5},
				{"e", 1, 6},
			},
		},
	}

	for _, test := range tests {
		var itrs []vellum.Iterator
		for _, entries := range test.in {
			itrs = append(itrs, &testIterator{entries: entries})
		}

		enumerator, err := newEnumerator(itrs)
		if err != nil {
			t.Fatalf("%s - expected no err on newNumerator, got: %v", test.desc, err)
		}

		wanti := 0
		for wanti < len(test.want) {
			if err != nil {
				t.Fatalf("%s - wanted no err, got: %v", test.desc, err)
			}

			currK, currIdx, currV := enumerator.Current()

			want := test.want[wanti]
			if want.key != string(currK) {
				t.Fatalf("%s - wrong key, wanted: %#v, got: %q, %d, %d", test.desc,
					want, currK, currIdx, currV)
			}
			if want.idx != currIdx {
				t.Fatalf("%s - wrong idx, wanted: %#v, got: %q, %d, %d", test.desc,
					want, currK, currIdx, currV)
			}
			if want.val != currV {
				t.Fatalf("%s - wrong val, wanted: %#v, got: %q, %d, %d", test.desc,
					want, currK, currIdx, currV)
			}

			wanti += 1

			err = enumerator.Next()
		}

		if err != vellum.ErrIteratorDone {
			t.Fatalf("%s - expected ErrIteratorDone, got: %v", test.desc, err)
		}

		err = enumerator.Close()
		if err != nil {
			t.Fatalf("%s - expected nil err on close, got: %v", test.desc, err)
		}

		for _, itr := range itrs {
			if itr.(*testIterator).curr != 654321 {
				t.Fatalf("%s - expected child iter to be closed", test.desc)
			}
		}
	}
}

type testIterator struct {
	entries []enumTestEntry
	curr    int
}

func (m *testIterator) Current() ([]byte, uint64) {
	if m.curr >= len(m.entries) {
		return nil, 0
	}
	return []byte(m.entries[m.curr].key), m.entries[m.curr].val
}

func (m *testIterator) Next() error {
	m.curr++
	if m.curr >= len(m.entries) {
		return vellum.ErrIteratorDone
	}
	return nil
}

func (m *testIterator) Seek(key []byte) error {
	return fmt.Errorf("not implemented for enumerator unit tests")
}

func (m *testIterator) Reset(f *vellum.FST,
	startKeyInclusive, endKeyExclusive []byte, aut vellum.Automaton) error {
	return fmt.Errorf("not implemented for enumerator unit tests")
}

func (m *testIterator) Close() error {
	m.curr = 654321
	return nil
}
<commit_msg>EnumeratorTests adopting Exists method<commit_after>//  Copyright (c) 2018 Couchbase, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// 		http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
// implied.  See the License for the specific language governing
// permissions and limitations under the License.

package zap

import (
	"fmt"
	"testing"

	"github.com/couchbase/vellum"
)

type enumTestEntry struct {
	key string
	val uint64
}

type enumTestWant struct {
	key string
	idx int
	val uint64
}

func TestEnumerator(t *testing.T) {
	tests := []struct {
		desc string
		in   [][]enumTestEntry
		want []enumTestWant
	}{
		{
			desc: "two non-empty enumerators with no duplicate keys",
			in: [][]enumTestEntry{
				[]enumTestEntry{
					{"a", 1},
					{"c", 3},
					{"e", 5},
				},
				[]enumTestEntry{
					{"b", 2},
					{"d", 4},
					{"f", 6},
				},
			},
			want: []enumTestWant{
				{"a", 0, 1},
				{"b", 1, 2},
				{"c", 0, 3},
				{"d", 1, 4},
				{"e", 0, 5},
				{"f", 1, 6},
			},
		},
		{
			desc: "two non-empty enumerators with duplicate keys",
			in: [][]enumTestEntry{
				[]enumTestEntry{
					{"a", 1},
					{"c", 3},
					{"e", 5},
				},
				[]enumTestEntry{
					{"a", 2},
					{"c", 4},
					{"e", 6},
				},
			},
			want: []enumTestWant{
				{"a", 0, 1},
				{"a", 1, 2},
				{"c", 0, 3},
				{"c", 1, 4},
				{"e", 0, 5},
				{"e", 1, 6},
			},
		},
		{
			desc: "first iterator is empty",
			in: [][]enumTestEntry{
				[]enumTestEntry{},
				[]enumTestEntry{
					{"a", 2},
					{"c", 4},
					{"e", 6},
				},
			},
			want: []enumTestWant{
				{"a", 1, 2},
				{"c", 1, 4},
				{"e", 1, 6},
			},
		},
		{
			desc: "last iterator is empty",
			in: [][]enumTestEntry{
				[]enumTestEntry{
					{"a", 1},
					{"c", 3},
					{"e", 5},
				},
				[]enumTestEntry{},
			},
			want: []enumTestWant{
				{"a", 0, 1},
				{"c", 0, 3},
				{"e", 0, 5},
			},
		},
		{
			desc: "two different length enumerators with duplicate keys",
			in: [][]enumTestEntry{
				[]enumTestEntry{
					{"a", 1},
					{"c", 3},
					{"e", 5},
				},
				[]enumTestEntry{
					{"a", 2},
					{"b", 4},
					{"d", 1000},
					{"e", 6},
				},
			},
			want: []enumTestWant{
				{"a", 0, 1},
				{"a", 1, 2},
				{"b", 1, 4},
				{"c", 0, 3},
				{"d", 1, 1000},
				{"e", 0, 5},
				{"e", 1, 6},
			},
		},
	}

	for _, test := range tests {
		var itrs []vellum.Iterator
		for _, entries := range test.in {
			itrs = append(itrs, &testIterator{entries: entries})
		}

		enumerator, err := newEnumerator(itrs)
		if err != nil {
			t.Fatalf("%s - expected no err on newNumerator, got: %v", test.desc, err)
		}

		wanti := 0
		for wanti < len(test.want) {
			if err != nil {
				t.Fatalf("%s - wanted no err, got: %v", test.desc, err)
			}

			currK, currIdx, currV := enumerator.Current()

			want := test.want[wanti]
			if want.key != string(currK) {
				t.Fatalf("%s - wrong key, wanted: %#v, got: %q, %d, %d", test.desc,
					want, currK, currIdx, currV)
			}
			if want.idx != currIdx {
				t.Fatalf("%s - wrong idx, wanted: %#v, got: %q, %d, %d", test.desc,
					want, currK, currIdx, currV)
			}
			if want.val != currV {
				t.Fatalf("%s - wrong val, wanted: %#v, got: %q, %d, %d", test.desc,
					want, currK, currIdx, currV)
			}

			wanti += 1

			err = enumerator.Next()
		}

		if err != vellum.ErrIteratorDone {
			t.Fatalf("%s - expected ErrIteratorDone, got: %v", test.desc, err)
		}

		err = enumerator.Close()
		if err != nil {
			t.Fatalf("%s - expected nil err on close, got: %v", test.desc, err)
		}

		for _, itr := range itrs {
			if itr.(*testIterator).curr != 654321 {
				t.Fatalf("%s - expected child iter to be closed", test.desc)
			}
		}
	}
}

type testIterator struct {
	entries []enumTestEntry
	curr    int
}

func (m *testIterator) Current() ([]byte, uint64) {
	if m.curr >= len(m.entries) {
		return nil, 0
	}
	return []byte(m.entries[m.curr].key), m.entries[m.curr].val
}

func (m *testIterator) Next() error {
	m.curr++
	if m.curr >= len(m.entries) {
		return vellum.ErrIteratorDone
	}
	return nil
}

func (m *testIterator) Seek(key []byte) error {
	return fmt.Errorf("not implemented for enumerator unit tests")
}

func (m *testIterator) Reset(f *vellum.FST,
	startKeyInclusive, endKeyExclusive []byte, aut vellum.Automaton) error {
	return fmt.Errorf("not implemented for enumerator unit tests")
}

func (m *testIterator) Close() error {
	m.curr = 654321
	return nil
}

func (m *testIterator) Exists(key []byte) (bool, error) {
	return false, fmt.Errorf("not implemented for enumerator unit tests")
}
<|endoftext|>


<commit_before>// Package parser parses Reddit's JSON responses into protobuffers.
package parser

import (
	"encoding/json"
	"fmt"
	"io"

	"github.com/turnage/redditproto"
)

// ParseLinkListing returns a slice of Links which hold the same data the JSON
// link listing provided contains.
func ParseLinkListing(content io.ReadCloser) ([]*redditproto.Link, error) {
	if content == nil {
		return nil, fmt.Errorf("no content provided")
	}

	listing := &redditproto.LinkListing{}
	decoder := json.NewDecoder(content)
	if err := decoder.Decode(listing); err != nil {
		return nil, err
	}

	return unpackLinkListing(listing)
}

// ParseThread parses a combination link listing and comment listing, which
// Reddit returns when asked for the JSON digest of a thread. This contains the
// submission's information, and all of its comments. The returned link will
// have the Comments field filled, and the comments will have their ReplyTree
// field filled.
func ParseThread(content io.ReadCloser) (*redditproto.Link, error) {
	if content == nil {
		return nil, fmt.Errorf("no content provided")
	}

	listings := []interface{}{
		&redditproto.LinkListing{},
		&redditproto.CommentListing{},
	}
	decoder := json.NewDecoder(content)
	if err := decoder.Decode(&listings); err != nil {
		return nil, err
	}

	if len(listings) != 2 {
		return nil, fmt.Errorf(
			"json decoding malformed the listings: %v",
			listings)
	}

	linkListing := listings[0].(*redditproto.LinkListing)
	commentListing := listings[1].(*redditproto.CommentListing)

	unpackedLinks, err := unpackLinkListing(linkListing)
	if err != nil {
		return nil, err
	}

	if len(unpackedLinks) != 1 {
		return nil, fmt.Errorf(
			"unexpected amount of links (%d)",
			len(unpackedLinks))
	}

	link := unpackedLinks[0]
	link.Comments = unpackCommentListing(commentListing)

	return link, nil
}

// ParseInbox returns a slice of messages in the inbox JSON digest.
func ParseInbox(content io.ReadCloser) ([]*redditproto.Message, error) {
	if content == nil {
		return nil, fmt.Errorf("no content provided")
	}

	messageListing := &redditproto.MessageListing{}
	decoder := json.NewDecoder(content)
	if err := decoder.Decode(messageListing); err != nil {
		return nil, err
	}

	if messageListing.GetData() == nil {
		return nil, fmt.Errorf("no data in listing")
	}

	if messageListing.GetData().GetChildren() == nil {
		return nil, fmt.Errorf("the listing was infertile")
	}

	messages := make(
		[]*redditproto.Message,
		len(messageListing.GetData().GetChildren()))
	for i, child := range messageListing.GetData().GetChildren() {
		messages[i] = child.GetData()
	}

	return messages, nil
}

// unpackLinkListing returns a slice of the links contained in a link listing.
func unpackLinkListing(
	listing *redditproto.LinkListing,
) ([]*redditproto.Link, error) {
	if listing.GetData() == nil {
		return nil, fmt.Errorf("no data field; got %v", listing)
	}

	if listing.GetData().GetChildren() == nil {
		return nil, fmt.Errorf("data has no children; got %v", listing)
	}

	links := make([]*redditproto.Link, len(listing.GetData().GetChildren()))
	for i, child := range listing.GetData().GetChildren() {
		links[i] = child.GetData()
	}
	return links, nil
}

// unpackCommentListing returns a slice of the comments contained in a comment
// listing.
func unpackCommentListing(
	listing *redditproto.CommentListing,
) []*redditproto.Comment {
	if listing.GetData() == nil {
		return nil
	}

	if listing.GetData().GetChildren() == nil {
		return nil
	}

	comments := make(
		[]*redditproto.Comment,
		len(listing.GetData().GetChildren()))
	for i, child := range listing.GetData().GetChildren() {
		comments[i] = child.GetData()
		if comments[i].Replies != nil {
			comments[i].ReplyTree = unpackCommentListing(comments[i].Replies)
			comments[i].Replies = nil
		}
	}
	return comments
}
<commit_msg>Close all connections.<commit_after>// Package parser parses Reddit's JSON responses into protobuffers.
package parser

import (
	"encoding/json"
	"fmt"
	"io"

	"github.com/turnage/redditproto"
)

// ParseLinkListing returns a slice of Links which hold the same data the JSON
// link listing provided contains.
func ParseLinkListing(content io.ReadCloser) ([]*redditproto.Link, error) {
	if content == nil {
		return nil, fmt.Errorf("no content provided")
	}
	defer content.Close()

	listing := &redditproto.LinkListing{}
	decoder := json.NewDecoder(content)
	if err := decoder.Decode(listing); err != nil {
		return nil, err
	}

	return unpackLinkListing(listing)
}

// ParseThread parses a combination link listing and comment listing, which
// Reddit returns when asked for the JSON digest of a thread. This contains the
// submission's information, and all of its comments. The returned link will
// have the Comments field filled, and the comments will have their ReplyTree
// field filled.
func ParseThread(content io.ReadCloser) (*redditproto.Link, error) {
	if content == nil {
		return nil, fmt.Errorf("no content provided")
	}
	defer content.Close()

	listings := []interface{}{
		&redditproto.LinkListing{},
		&redditproto.CommentListing{},
	}
	decoder := json.NewDecoder(content)
	if err := decoder.Decode(&listings); err != nil {
		return nil, err
	}

	if len(listings) != 2 {
		return nil, fmt.Errorf(
			"json decoding malformed the listings: %v",
			listings)
	}

	linkListing := listings[0].(*redditproto.LinkListing)
	commentListing := listings[1].(*redditproto.CommentListing)

	unpackedLinks, err := unpackLinkListing(linkListing)
	if err != nil {
		return nil, err
	}

	if len(unpackedLinks) != 1 {
		return nil, fmt.Errorf(
			"unexpected amount of links (%d)",
			len(unpackedLinks))
	}

	link := unpackedLinks[0]
	link.Comments = unpackCommentListing(commentListing)

	return link, nil
}

// ParseInbox returns a slice of messages in the inbox JSON digest.
func ParseInbox(content io.ReadCloser) ([]*redditproto.Message, error) {
	if content == nil {
		return nil, fmt.Errorf("no content provided")
	}
	defer content.Close()

	messageListing := &redditproto.MessageListing{}
	decoder := json.NewDecoder(content)
	if err := decoder.Decode(messageListing); err != nil {
		return nil, err
	}

	if messageListing.GetData() == nil {
		return nil, fmt.Errorf("no data in listing")
	}

	if messageListing.GetData().GetChildren() == nil {
		return nil, fmt.Errorf("the listing was infertile")
	}

	messages := make(
		[]*redditproto.Message,
		len(messageListing.GetData().GetChildren()))
	for i, child := range messageListing.GetData().GetChildren() {
		messages[i] = child.GetData()
	}

	return messages, nil
}

// unpackLinkListing returns a slice of the links contained in a link listing.
func unpackLinkListing(
	listing *redditproto.LinkListing,
) ([]*redditproto.Link, error) {
	if listing.GetData() == nil {
		return nil, fmt.Errorf("no data field; got %v", listing)
	}

	if listing.GetData().GetChildren() == nil {
		return nil, fmt.Errorf("data has no children; got %v", listing)
	}

	links := make([]*redditproto.Link, len(listing.GetData().GetChildren()))
	for i, child := range listing.GetData().GetChildren() {
		links[i] = child.GetData()
	}
	return links, nil
}

// unpackCommentListing returns a slice of the comments contained in a comment
// listing.
func unpackCommentListing(
	listing *redditproto.CommentListing,
) []*redditproto.Comment {
	if listing.GetData() == nil {
		return nil
	}

	if listing.GetData().GetChildren() == nil {
		return nil
	}

	comments := make(
		[]*redditproto.Comment,
		len(listing.GetData().GetChildren()))
	for i, child := range listing.GetData().GetChildren() {
		comments[i] = child.GetData()
		if comments[i].Replies != nil {
			comments[i].ReplyTree = unpackCommentListing(comments[i].Replies)
			comments[i].Replies = nil
		}
	}
	return comments
}
<|endoftext|>


<commit_before>package ec2

import (
	"fmt"
	"strings"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/pchico83/i2kit/cli/schemas/environment"
)

//CreateSG creates the project security group
func CreateSG(e *environment.Environment, config *aws.Config) error {
	vpc, err := GetVPC(e, config)
	if err != nil {
		return err
	}
	svc := ec2.New(session.New(), config)
	description := fmt.Sprintf("Security Group for environment %s", e.Name)
	name := fmt.Sprintf("%s-i2kit", e.Name)
	csgi := &ec2.CreateSecurityGroupInput{
		Description: &description,
		GroupName:   &name,
		VpcId:       &vpc,
	}
	sg, err := svc.CreateSecurityGroup(csgi)
	if err != nil {
		if !strings.Contains(err.Error(), "InvalidGroup.Duplicate") {
			return err
		}
		dsgi := &ec2.DescribeSecurityGroupsInput{
			GroupNames: []*string{&name},
		}
		sgs, err2 := svc.DescribeSecurityGroups(dsgi)
		if err2 != nil {
			return err2
		}
		for _, i := range sgs.SecurityGroups {
			if *i.VpcId == vpc {
				e.Provider.SecurityGroup = *i.GroupId
			}
		}
	} else {
		e.Provider.SecurityGroup = *sg.GroupId
	}

	if e.Provider.SecurityGroup == "" {
		return fmt.Errorf("Error retrieving SG '%s'", name)
	}

	asgii := &ec2.AuthorizeSecurityGroupIngressInput{
		GroupName:               &name,
		SourceSecurityGroupName: &name,
	}
	protocol := "-1"
	if _, err = svc.AuthorizeSecurityGroupIngress(asgii); err != nil {
		if strings.Contains(err.Error(), "InvalidParameterValue") {
			asgii = &ec2.AuthorizeSecurityGroupIngressInput{
				GroupId: &e.Provider.SecurityGroup,
				IpPermissions: []*ec2.IpPermission{
					&ec2.IpPermission{
						IpProtocol: &protocol,
						UserIdGroupPairs: []*ec2.UserIdGroupPair{
							&ec2.UserIdGroupPair{
								GroupId: &e.Provider.SecurityGroup,
							},
						},
					},
				},
			}
			_, err = svc.AuthorizeSecurityGroupIngress(asgii)
		}
	}
	if err != nil && !strings.Contains(err.Error(), "InvalidPermission.Duplicate") {
		return err
	}
	var ports int64
	ports = 22
	cidrIP := "0.0.0.0/0"
	asgii = &ec2.AuthorizeSecurityGroupIngressInput{
		FromPort:   &ports,
		ToPort:     &ports,
		GroupName:  &name,
		IpProtocol: &protocol,
		CidrIp:     &cidrIP,
	}
	if _, err = svc.AuthorizeSecurityGroupIngress(asgii); err != nil {
		if strings.Contains(err.Error(), "InvalidParameterValue") {
			asgii.SetGroupName("")
			asgii.SetGroupId(e.Provider.SecurityGroup)
			_, err = svc.AuthorizeSecurityGroupIngress(asgii)
		}
	}
	if err != nil && !strings.Contains(err.Error(), "InvalidPermission.Duplicate") {
		return err
	}
	return nil
}
<commit_msg>use a filter to find security groups by name (#72)<commit_after>package ec2

import (
	"fmt"
	"strings"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/pchico83/i2kit/cli/schemas/environment"
)

var groupNameFilter = "group-name"

//CreateSG creates the project security group
func CreateSG(e *environment.Environment, config *aws.Config) error {
	vpc, err := GetVPC(e, config)
	if err != nil {
		return err
	}
	svc := ec2.New(session.New(), config)
	description := fmt.Sprintf("Security Group for environment %s", e.Name)
	name := fmt.Sprintf("%s-i2kit", e.Name)
	csgi := &ec2.CreateSecurityGroupInput{
		Description: &description,
		GroupName:   &name,
		VpcId:       &vpc,
	}
	sg, err := svc.CreateSecurityGroup(csgi)
	if err != nil {
		if !strings.Contains(err.Error(), "InvalidGroup.Duplicate") {
			return err
		}

		dsgi := &ec2.DescribeSecurityGroupsInput{
			Filters: []*ec2.Filter{&ec2.Filter{Name: &groupNameFilter, Values: []*string{&name}}},
		}
		sgs, err2 := svc.DescribeSecurityGroups(dsgi)

		if err2 != nil {
			return err2
		}

		for _, i := range sgs.SecurityGroups {
			if *i.VpcId == vpc {
				e.Provider.SecurityGroup = *i.GroupId
			}
		}
	} else {
		e.Provider.SecurityGroup = *sg.GroupId
	}

	if e.Provider.SecurityGroup == "" {
		return fmt.Errorf("Error retrieving SG '%s'", name)
	}

	asgii := &ec2.AuthorizeSecurityGroupIngressInput{
		GroupName:               &name,
		SourceSecurityGroupName: &name,
	}
	protocol := "-1"
	_, err = svc.AuthorizeSecurityGroupIngress(asgii)
	if err != nil {
		if strings.Contains(err.Error(), "InvalidParameterValue") {
			asgii = &ec2.AuthorizeSecurityGroupIngressInput{
				GroupId: &e.Provider.SecurityGroup,
				IpPermissions: []*ec2.IpPermission{
					&ec2.IpPermission{
						IpProtocol: &protocol,
						UserIdGroupPairs: []*ec2.UserIdGroupPair{
							&ec2.UserIdGroupPair{
								GroupId: &e.Provider.SecurityGroup,
							},
						},
					},
				},
			}
			_, err = svc.AuthorizeSecurityGroupIngress(asgii)
		}
	}

	if err != nil && !strings.Contains(err.Error(), "InvalidPermission.Duplicate") {
		return err
	}
	var ports int64
	ports = 22
	cidrIP := "0.0.0.0/0"
	asgii = &ec2.AuthorizeSecurityGroupIngressInput{
		FromPort:   &ports,
		ToPort:     &ports,
		GroupName:  &name,
		IpProtocol: &protocol,
		CidrIp:     &cidrIP,
	}
	if _, err = svc.AuthorizeSecurityGroupIngress(asgii); err != nil {
		if strings.Contains(err.Error(), "InvalidParameterValue") {
			asgii.SetGroupName("")
			asgii.SetGroupId(e.Provider.SecurityGroup)
			_, err = svc.AuthorizeSecurityGroupIngress(asgii)
		}
	}
	if err != nil && !strings.Contains(err.Error(), "InvalidPermission.Duplicate") {
		return err
	}
	return nil
}
<|endoftext|>