<commit_before>package ec2

import (
	"fmt"
	"launchpad.net/goamz/ec2"
	"launchpad.net/goamz/s3"
	"launchpad.net/juju/go/environs"
	"launchpad.net/juju/go/state"
	"sync"
)

const zkPort = 2181

var zkPortSuffix = fmt.Sprintf(":%d", zkPort)

func init() {
	environs.RegisterProvider("ec2", environProvider{})
}

type environProvider struct{}

var _ environs.EnvironProvider = environProvider{}

type environ struct {
	name             string
	config           *providerConfig
	ec2              *ec2.EC2
	s3               *s3.S3
	checkBucket      sync.Once
	checkBucketError error
}

var _ environs.Environ = (*environ)(nil)

type instance struct {
	*ec2.Instance
}

func (inst *instance) String() string {
	return inst.Id()
}

var _ environs.Instance = (*instance)(nil)

func (inst *instance) Id() string {
	return inst.InstanceId
}

func (inst *instance) DNSName() string {
	return inst.Instance.DNSName
}

func (environProvider) Open(name string, config interface{}) (e environs.Environ, err error) {
	cfg := config.(*providerConfig)
	if Regions[cfg.region].EC2Endpoint == "" {
		return nil, fmt.Errorf("no ec2 endpoint found for region %q, opening %q", cfg.region, name)
	}
	return &environ{
		name:   name,
		config: cfg,
		ec2:    ec2.New(cfg.auth, Regions[cfg.region]),
		s3:     s3.New(cfg.auth, Regions[cfg.region]),
	}, nil
}

func (e *environ) Bootstrap() error {
	_, err := e.loadState()
	if err == nil {
		return fmt.Errorf("environment is already bootstrapped")
	}
	if s3err, _ := err.(*s3.Error); s3err != nil && s3err.StatusCode != 404 {
		return err
	}
	inst, err := e.startInstance(0, nil, true)
	if err != nil {
		return fmt.Errorf("cannot start bootstrap instance: %v", err)
	}
	err = e.saveState(&bootstrapState{
		ZookeeperInstances: []string{inst.Id()},
	})
	if err != nil {
		// ignore error on StopInstance because the previous error is
		// more important.
		e.StopInstances([]environs.Instance{inst})
		return err
	}
	// TODO make safe in the case of racing Bootstraps
	// If two Bootstraps are called concurrently, there's
	// no way to use S3 to make sure that only one succeeds.
	// Perhaps consider using SimpleDB for state storage
	// which would enable that possibility.

	return nil
}

func (e *environ) StateInfo() (*state.Info, error) {
	st, err := e.loadState()
	if err != nil {
		return nil, err
	}
	resp, err := e.ec2.Instances(st.ZookeeperInstances, nil)
	if err != nil {
		return nil, fmt.Errorf("cannot list instances: %v", err)
	}
	var insts []environs.Instance
	for i := range resp.Reservations {
		r := &resp.Reservations[i]
		for j := range r.Instances {
			insts = append(insts, &instance{&r.Instances[j]})
		}
	}

	addrs := make([]string, len(insts))
	for i, inst := range insts {
		// TODO make this poll until the DNSName becomes available.
		addr := inst.DNSName()
		if addr == "" {
			return nil, fmt.Errorf("zookeeper instance %q does not yet have a DNS address", inst.Id())
		}
		addrs[i] = addr + zkPortSuffix
	}
	return &state.Info{Addrs: addrs}, nil
}

func (e *environ) StartInstance(machineId int, info *state.Info) (environs.Instance, error) {
	return e.startInstance(machineId, info, false)
}

// startInstance is the internal version of StartInstance, used by Bootstrap
// as well as via StartInstance itself. If master is true, a bootstrap
// instance will be started.
func (e *environ) startInstance(machineId int, _ *state.Info, master bool) (environs.Instance, error) {
	image, err := FindImageSpec(DefaultImageConstraint)
	if err != nil {
		return nil, fmt.Errorf("cannot find image: %v", err)
	}
	groups, err := e.setUpGroups(machineId)
	if err != nil {
		return nil, fmt.Errorf("cannot set up groups: %v", err)
	}
	instances, err := e.ec2.RunInstances(&ec2.RunInstances{
		ImageId:        image.ImageId,
		MinCount:       1,
		MaxCount:       1,
		UserData:       nil,
		InstanceType:   "m1.small",
		SecurityGroups: groups,
	})
	if err != nil {
		return nil, fmt.Errorf("cannot run instances: %v", err)
	}
	if len(instances.Instances) != 1 {
		return nil, fmt.Errorf("expected 1 started instance, got %d", len(instances.Instances))
	}
	return &instance{&instances.Instances[0]}, nil
}

func (e *environ) StopInstances(insts []environs.Instance) error {
	ids := make([]string, len(insts))
	for i, inst := range insts {
		ids[i] = inst.(*instance).InstanceId
	}
	return e.terminateInstances(ids)
}

func (e *environ) Instances(ids []string) ([]environs.Instance, error) {
	if len(ids) == 0 {
		return nil, nil
	}
	insts := make([]environs.Instance, len(ids))

	// TODO make a series of requests to cope with eventual consistency.
	filter := ec2.NewFilter()
	filter.Add("instance-state-name", "pending", "running")
	filter.Add("group-name", e.groupName())
	filter.Add("instance-id", ids...)
	resp, err := e.ec2.Instances(nil, filter)
	if err != nil {
		return nil, err
	}
	// For each requested id, add it to the returned instances
	// if we find it in the response.
	n := 0
	for i, id := range ids {
		if insts[i] != nil {
			continue
		}
		for j := range resp.Reservations {
			r := &resp.Reservations[j]
			for k := range r.Instances {
				inst := &r.Instances[k]
				if inst.InstanceId == id {
					insts[i] = &instance{inst}
					n++
				}
			}
		}
	}
	if n == 0 {
		return nil, environs.ErrMissingInstance
	}
	if n < len(ids) {
		return insts, environs.ErrMissingInstance
	}
	return insts, err
}

func (e *environ) Destroy(insts []environs.Instance) error {
	// Try to find all the instances in the environ's group.
	filter := ec2.NewFilter()
	filter.Add("instance-state-name", "pending", "running")
	filter.Add("group-name", e.groupName())
	resp, err := e.ec2.Instances(nil, filter)
	if err != nil {
		return fmt.Errorf("cannot get instances: %v", err)
	}
	var ids []string
	found := make(map[string]bool)
	for _, r := range resp.Reservations {
		for _, inst := range r.Instances {
			ids = append(ids, inst.InstanceId)
			found[inst.InstanceId] = true
		}
	}

	// Then add any instances we've been told about
	// but haven't yet shown up in the instance list.
	for _, inst := range insts {
		id := inst.(*instance).InstanceId
		if !found[id] {
			ids = append(ids, id)
			found[id] = true
		}
	}
	err = e.terminateInstances(ids)
	if err != nil {
		return err
	}
	err = e.deleteState()
	if err != nil {
		return err
	}
	return nil
}

func (e *environ) terminateInstances(ids []string) error {
	if len(ids) == 0 {
		return nil
	}
	_, err := e.ec2.TerminateInstances(ids)
	if err == nil || ec2ErrCode(err) != "InvalidInstance.NotFound" {
		return err
	}
	var firstErr error
	// If we get a NotFound error, it means that no instances have been
	// terminated, so try them one by one, ignoring NotFound errors.
	for _, id := range ids {
		_, err = e.ec2.TerminateInstances([]string{id})
		if ec2ErrCode(err) == "InvalidInstance.NotFound" {
			err = nil
		}
		if err != nil && firstErr == nil {
			firstErr = err
		}
	}
	return firstErr
}

func (e *environ) machineGroupName(machineId int) string {
	return fmt.Sprintf("%s-%d", e.groupName(), machineId)
}

func (e *environ) groupName() string {
	return "juju-" + e.name
}

// setUpGroups creates the security groups for the new machine, and
// returns them.
//
// Instances are tagged with a group so they can be distinguished from
// other instances that might be running on the same EC2 account.  In
// addition, a specific machine security group is created for each
// machine, so that its firewall rules can be configured per machine.
func (e *environ) setUpGroups(machineId int) ([]ec2.SecurityGroup, error) {
	jujuGroup, err := e.ensureGroup(e.groupName(), "juju group for "+e.name,
		[]ec2.IPPerm{
			// TODO delete this authorization when we can do
			// the zookeeper ssh tunnelling.
			{
				Protocol:  "tcp",
				FromPort:  zkPort,
				ToPort:    zkPort,
				SourceIPs: []string{"0.0.0.0/0"},
			},
			{
				Protocol:  "tcp",
				FromPort:  22,
				ToPort:    22,
				SourceIPs: []string{"0.0.0.0/0"},
			},
			// TODO authorize internal traffic
		})
	if err != nil {
		return nil, err
	}
	descr := fmt.Sprintf("juju group for %s machine %d", e.name, machineId)
	jujuMachineGroup, err := e.ensureGroup(e.machineGroupName(machineId), descr, nil)
	if err != nil {
		return nil, err
	}
	return []ec2.SecurityGroup{jujuGroup, jujuMachineGroup}, nil
}

// zg holds the zero security group.
var zg ec2.SecurityGroup

// ensureGroup tries to ensure that a security group exists with the given
// name and permissions. If the group already exists, its permissions
// will be changed accordingly. If the group does not exist, it will be created
// with the given description. It returns the group.
func (e *environ) ensureGroup(name, descr string, perms []ec2.IPPerm) (g ec2.SecurityGroup, err error) {
	resp, err := e.ec2.CreateSecurityGroup(name, descr)
	if err != nil && ec2ErrCode(err) != "InvalidGroup.Duplicate" {
		return zg, err
	}

	want := newPermSet(perms)
	var have permSet
	if err == nil {
		g = resp.SecurityGroup
	} else {
		resp, err := e.ec2.SecurityGroups(ec2.SecurityGroupNames(name), nil)
		if err != nil {
			return zg, err
		}
		// TODO do we mind if the old group has the wrong description?
		have = newPermSet(resp.Groups[0].IPPerms)
		g = resp.Groups[0].SecurityGroup
	}
	revoke := make(permSet)
	for p := range have {
		if !want[p] {
			revoke[p] = true
		}
	}
	if len(revoke) > 0 {
		_, err := e.ec2.RevokeSecurityGroup(g, revoke.ipPerms())
		if err != nil {
			return zg, fmt.Errorf("cannot revoke security group: %v", err)
		}
	}

	add := make(map[ipPerm]bool)
	for p := range want {
		if !have[p] {
			add[p] = true
		}
	}
	if len(add) > 0 {
		_, err := e.ec2.AuthorizeSecurityGroup(g, perms)
		if err != nil {
			return zg, fmt.Errorf("cannot authorize securityGroup: %v", err)
		}
	}
	return g, nil
}

// ipPerm represents a permission for a group or an ip address range
// to access the given range of ports. Only one of groupId or ipAddr
// should be non-empty.
type ipPerm struct {
	protocol string
	fromPort int
	toPort   int
	groupId  string
	ipAddr   string
}

type permSet map[ipPerm]bool

// newPermSet returns a set of all the permissions in the
// given slice of IPPerms. It ignores the name and owner
// id in source groups, using group ids only.
func newPermSet(ps []ec2.IPPerm) permSet {
	m := make(permSet)
	for _, p := range ps {
		ipp := ipPerm{
			protocol: p.Protocol,
			fromPort: p.FromPort,
			toPort:   p.ToPort,
		}
		for _, g := range p.SourceGroups {
			ipp.groupId = g.Id
			m[ipp] = true
		}
		ipp.groupId = ""
		for _, ip := range p.SourceIPs {
			ipp.ipAddr = ip
			m[ipp] = true
		}
	}
	return m
}

// ipPerms returns the given set of permissions
// as a slice of IPPerms.
func (m permSet) ipPerms() (ps []ec2.IPPerm) {
	// We could compact the permissions, but it
	// hardly seems worth it.
	for p := range m {
		ipp := ec2.IPPerm{
			Protocol: p.protocol,
			FromPort: p.fromPort,
			ToPort:   p.toPort,
		}
		if p.ipAddr != "" {
			ipp.SourceIPs = []string{p.ipAddr}
		} else {
			ipp.SourceGroups = []ec2.UserSecurityGroup{{Id: p.groupId}}
		}
		ps = append(ps, ipp)
	}
	return
}

// If the err is of type *ec2.Error, ec2ErrCode returns
// its code, otherwise it returns the empty string.
func ec2ErrCode(err error) string {
	ec2err, _ := err.(*ec2.Error)
	if ec2err == nil {
		return ""
	}
	return ec2err.Code
}
<commit_msg>ec2: fix error code bug. fix add permission bug<commit_after>package ec2

import (
	"fmt"
	"launchpad.net/goamz/ec2"
	"launchpad.net/goamz/s3"
	"launchpad.net/juju/go/environs"
	"launchpad.net/juju/go/state"
	"sync"
)

const zkPort = 2181

var zkPortSuffix = fmt.Sprintf(":%d", zkPort)

func init() {
	environs.RegisterProvider("ec2", environProvider{})
}

type environProvider struct{}

var _ environs.EnvironProvider = environProvider{}

type environ struct {
	name             string
	config           *providerConfig
	ec2              *ec2.EC2
	s3               *s3.S3
	checkBucket      sync.Once
	checkBucketError error
}

var _ environs.Environ = (*environ)(nil)

type instance struct {
	*ec2.Instance
}

func (inst *instance) String() string {
	return inst.Id()
}

var _ environs.Instance = (*instance)(nil)

func (inst *instance) Id() string {
	return inst.InstanceId
}

func (inst *instance) DNSName() string {
	return inst.Instance.DNSName
}

func (environProvider) Open(name string, config interface{}) (e environs.Environ, err error) {
	cfg := config.(*providerConfig)
	if Regions[cfg.region].EC2Endpoint == "" {
		return nil, fmt.Errorf("no ec2 endpoint found for region %q, opening %q", cfg.region, name)
	}
	return &environ{
		name:   name,
		config: cfg,
		ec2:    ec2.New(cfg.auth, Regions[cfg.region]),
		s3:     s3.New(cfg.auth, Regions[cfg.region]),
	}, nil
}

func (e *environ) Bootstrap() error {
	_, err := e.loadState()
	if err == nil {
		return fmt.Errorf("environment is already bootstrapped")
	}
	if s3err, _ := err.(*s3.Error); s3err != nil && s3err.StatusCode != 404 {
		return err
	}
	inst, err := e.startInstance(0, nil, true)
	if err != nil {
		return fmt.Errorf("cannot start bootstrap instance: %v", err)
	}
	err = e.saveState(&bootstrapState{
		ZookeeperInstances: []string{inst.Id()},
	})
	if err != nil {
		// ignore error on StopInstance because the previous error is
		// more important.
		e.StopInstances([]environs.Instance{inst})
		return err
	}
	// TODO make safe in the case of racing Bootstraps
	// If two Bootstraps are called concurrently, there's
	// no way to use S3 to make sure that only one succeeds.
	// Perhaps consider using SimpleDB for state storage
	// which would enable that possibility.

	return nil
}

func (e *environ) StateInfo() (*state.Info, error) {
	st, err := e.loadState()
	if err != nil {
		return nil, err
	}
	resp, err := e.ec2.Instances(st.ZookeeperInstances, nil)
	if err != nil {
		return nil, fmt.Errorf("cannot list instances: %v", err)
	}
	var insts []environs.Instance
	for i := range resp.Reservations {
		r := &resp.Reservations[i]
		for j := range r.Instances {
			insts = append(insts, &instance{&r.Instances[j]})
		}
	}

	addrs := make([]string, len(insts))
	for i, inst := range insts {
		// TODO make this poll until the DNSName becomes available.
		addr := inst.DNSName()
		if addr == "" {
			return nil, fmt.Errorf("zookeeper instance %q does not yet have a DNS address", inst.Id())
		}
		addrs[i] = addr + zkPortSuffix
	}
	return &state.Info{Addrs: addrs}, nil
}

func (e *environ) StartInstance(machineId int, info *state.Info) (environs.Instance, error) {
	return e.startInstance(machineId, info, false)
}

// startInstance is the internal version of StartInstance, used by Bootstrap
// as well as via StartInstance itself. If master is true, a bootstrap
// instance will be started.
func (e *environ) startInstance(machineId int, _ *state.Info, master bool) (environs.Instance, error) {
	image, err := FindImageSpec(DefaultImageConstraint)
	if err != nil {
		return nil, fmt.Errorf("cannot find image: %v", err)
	}
	groups, err := e.setUpGroups(machineId)
	if err != nil {
		return nil, fmt.Errorf("cannot set up groups: %v", err)
	}
	instances, err := e.ec2.RunInstances(&ec2.RunInstances{
		ImageId:        image.ImageId,
		MinCount:       1,
		MaxCount:       1,
		UserData:       nil,
		InstanceType:   "m1.small",
		SecurityGroups: groups,
	})
	if err != nil {
		return nil, fmt.Errorf("cannot run instances: %v", err)
	}
	if len(instances.Instances) != 1 {
		return nil, fmt.Errorf("expected 1 started instance, got %d", len(instances.Instances))
	}
	return &instance{&instances.Instances[0]}, nil
}

func (e *environ) StopInstances(insts []environs.Instance) error {
	ids := make([]string, len(insts))
	for i, inst := range insts {
		ids[i] = inst.(*instance).InstanceId
	}
	return e.terminateInstances(ids)
}

func (e *environ) Instances(ids []string) ([]environs.Instance, error) {
	if len(ids) == 0 {
		return nil, nil
	}
	insts := make([]environs.Instance, len(ids))

	// TODO make a series of requests to cope with eventual consistency.
	filter := ec2.NewFilter()
	filter.Add("instance-state-name", "pending", "running")
	filter.Add("group-name", e.groupName())
	filter.Add("instance-id", ids...)
	resp, err := e.ec2.Instances(nil, filter)
	if err != nil {
		return nil, err
	}
	// For each requested id, add it to the returned instances
	// if we find it in the response.
	n := 0
	for i, id := range ids {
		if insts[i] != nil {
			continue
		}
		for j := range resp.Reservations {
			r := &resp.Reservations[j]
			for k := range r.Instances {
				inst := &r.Instances[k]
				if inst.InstanceId == id {
					insts[i] = &instance{inst}
					n++
				}
			}
		}
	}
	if n == 0 {
		return nil, environs.ErrMissingInstance
	}
	if n < len(ids) {
		return insts, environs.ErrMissingInstance
	}
	return insts, err
}

func (e *environ) Destroy(insts []environs.Instance) error {
	// Try to find all the instances in the environ's group.
	filter := ec2.NewFilter()
	filter.Add("instance-state-name", "pending", "running")
	filter.Add("group-name", e.groupName())
	resp, err := e.ec2.Instances(nil, filter)
	if err != nil {
		return fmt.Errorf("cannot get instances: %v", err)
	}
	var ids []string
	found := make(map[string]bool)
	for _, r := range resp.Reservations {
		for _, inst := range r.Instances {
			ids = append(ids, inst.InstanceId)
			found[inst.InstanceId] = true
		}
	}

	// Then add any instances we've been told about
	// but haven't yet shown up in the instance list.
	for _, inst := range insts {
		id := inst.(*instance).InstanceId
		if !found[id] {
			ids = append(ids, id)
			found[id] = true
		}
	}
	err = e.terminateInstances(ids)
	if err != nil {
		return err
	}
	err = e.deleteState()
	if err != nil {
		return err
	}
	return nil
}

func (e *environ) terminateInstances(ids []string) error {
	if len(ids) == 0 {
		return nil
	}
	_, err := e.ec2.TerminateInstances(ids)
	if err == nil || ec2ErrCode(err) != "InvalidInstanceID.NotFound" {
		return err
	}
	var firstErr error
	// If we get a NotFound error, it means that no instances have been
	// terminated, so try them one by one, ignoring NotFound errors.
	for _, id := range ids {
		_, err = e.ec2.TerminateInstances([]string{id})
		if ec2ErrCode(err) == "InvalidInstanceID.NotFound" {
			err = nil
		}
		if err != nil && firstErr == nil {
			firstErr = err
		}
	}
	return firstErr
}

func (e *environ) machineGroupName(machineId int) string {
	return fmt.Sprintf("%s-%d", e.groupName(), machineId)
}

func (e *environ) groupName() string {
	return "juju-" + e.name
}

// setUpGroups creates the security groups for the new machine, and
// returns them.
//
// Instances are tagged with a group so they can be distinguished from
// other instances that might be running on the same EC2 account.  In
// addition, a specific machine security group is created for each
// machine, so that its firewall rules can be configured per machine.
func (e *environ) setUpGroups(machineId int) ([]ec2.SecurityGroup, error) {
	jujuGroup, err := e.ensureGroup(e.groupName(), "juju group for "+e.name,
		[]ec2.IPPerm{
			// TODO delete this authorization when we can do
			// the zookeeper ssh tunnelling.
			{
				Protocol:  "tcp",
				FromPort:  zkPort,
				ToPort:    zkPort,
				SourceIPs: []string{"0.0.0.0/0"},
			},
			{
				Protocol:  "tcp",
				FromPort:  22,
				ToPort:    22,
				SourceIPs: []string{"0.0.0.0/0"},
			},
			// TODO authorize internal traffic
		})
	if err != nil {
		return nil, err
	}
	descr := fmt.Sprintf("juju group for %s machine %d", e.name, machineId)
	jujuMachineGroup, err := e.ensureGroup(e.machineGroupName(machineId), descr, nil)
	if err != nil {
		return nil, err
	}
	return []ec2.SecurityGroup{jujuGroup, jujuMachineGroup}, nil
}

// zg holds the zero security group.
var zg ec2.SecurityGroup

// ensureGroup tries to ensure that a security group exists with the given
// name and permissions. If the group already exists, its permissions
// will be changed accordingly. If the group does not exist, it will be created
// with the given description. It returns the group.
func (e *environ) ensureGroup(name, descr string, perms []ec2.IPPerm) (g ec2.SecurityGroup, err error) {
	resp, err := e.ec2.CreateSecurityGroup(name, descr)
	if err != nil && ec2ErrCode(err) != "InvalidGroup.Duplicate" {
		return zg, err
	}

	want := newPermSet(perms)
	var have permSet
	if err == nil {
		g = resp.SecurityGroup
	} else {
		resp, err := e.ec2.SecurityGroups(ec2.SecurityGroupNames(name), nil)
		if err != nil {
			return zg, err
		}
		// TODO do we mind if the old group has the wrong description?
		have = newPermSet(resp.Groups[0].IPPerms)
		g = resp.Groups[0].SecurityGroup
	}
	revoke := make(permSet)
	for p := range have {
		if !want[p] {
			revoke[p] = true
		}
	}
	if len(revoke) > 0 {
		_, err := e.ec2.RevokeSecurityGroup(g, revoke.ipPerms())
		if err != nil {
			return zg, fmt.Errorf("cannot revoke security group: %v", err)
		}
	}

	add := make(permSet)
	for p := range want {
		if !have[p] {
			add[p] = true
		}
	}
	if len(add) > 0 {
		_, err := e.ec2.AuthorizeSecurityGroup(g, add.ipPerms())
		if err != nil {
			return zg, fmt.Errorf("cannot authorize securityGroup: %v", err)
		}
	}
	return g, nil
}

// ipPerm represents a permission for a group or an ip address range
// to access the given range of ports. Only one of groupId or ipAddr
// should be non-empty.
type ipPerm struct {
	protocol string
	fromPort int
	toPort   int
	groupId  string
	ipAddr   string
}

type permSet map[ipPerm]bool

// newPermSet returns a set of all the permissions in the
// given slice of IPPerms. It ignores the name and owner
// id in source groups, using group ids only.
func newPermSet(ps []ec2.IPPerm) permSet {
	m := make(permSet)
	for _, p := range ps {
		ipp := ipPerm{
			protocol: p.Protocol,
			fromPort: p.FromPort,
			toPort:   p.ToPort,
		}
		for _, g := range p.SourceGroups {
			ipp.groupId = g.Id
			m[ipp] = true
		}
		ipp.groupId = ""
		for _, ip := range p.SourceIPs {
			ipp.ipAddr = ip
			m[ipp] = true
		}
	}
	return m
}

// ipPerms returns the given set of permissions
// as a slice of IPPerms.
func (m permSet) ipPerms() (ps []ec2.IPPerm) {
	// We could compact the permissions, but it
	// hardly seems worth it.
	for p := range m {
		ipp := ec2.IPPerm{
			Protocol: p.protocol,
			FromPort: p.fromPort,
			ToPort:   p.toPort,
		}
		if p.ipAddr != "" {
			ipp.SourceIPs = []string{p.ipAddr}
		} else {
			ipp.SourceGroups = []ec2.UserSecurityGroup{{Id: p.groupId}}
		}
		ps = append(ps, ipp)
	}
	return
}

// If the err is of type *ec2.Error, ec2ErrCode returns
// its code, otherwise it returns the empty string.
func ec2ErrCode(err error) string {
	ec2err, _ := err.(*ec2.Error)
	if ec2err == nil {
		return ""
	}
	return ec2err.Code
}
<|endoftext|>


<commit_before>package consul

import (
	"fmt"
	"sync"

	log "github.com/hashicorp/go-hclog"

	"github.com/hashicorp/nomad/command/agent/consul"
	"github.com/mitchellh/go-testing-interface"
)

// MockConsulOp represents the register/deregister operations.
type MockConsulOp struct {
	Op      string // add, remove, or update
	AllocID string
	Task    string
}

func NewMockConsulOp(op, allocID, task string) MockConsulOp {
	if op != "add" && op != "remove" && op != "update" && op != "alloc_registrations" {
		panic(fmt.Errorf("invalid consul op: %s", op))
	}
	return MockConsulOp{
		Op:      op,
		AllocID: allocID,
		Task:    task,
	}
}

// MockConsulServiceClient implements the ConsulServiceAPI interface to record
// and log task registration/deregistration.
type MockConsulServiceClient struct {
	ops []MockConsulOp
	mu  sync.Mutex

	logger log.Logger

	// AllocRegistrationsFn allows injecting return values for the
	// AllocRegistrations function.
	AllocRegistrationsFn func(allocID string) (*consul.AllocRegistration, error)
}

func NewMockConsulServiceClient(t testing.T, logger log.Logger) *MockConsulServiceClient {
	logger = logger.Named("mock_consul")
	m := MockConsulServiceClient{
		ops:    make([]MockConsulOp, 0, 20),
		logger: logger,
	}
	return &m
}

func (m *MockConsulServiceClient) UpdateTask(old, new *consul.TaskServices) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.logger.Trace("UpdateTask", "alloc_id", new.AllocID, "task", new.Name)
	m.ops = append(m.ops, NewMockConsulOp("update", new.AllocID, new.Name))
	return nil
}

func (m *MockConsulServiceClient) RegisterTask(task *consul.TaskServices) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.logger.Trace("RegisterTask", "alloc_id", task.AllocID, "task", task.Name)
	m.ops = append(m.ops, NewMockConsulOp("add", task.AllocID, task.Name))
	return nil
}

func (m *MockConsulServiceClient) RemoveTask(task *consul.TaskServices) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.logger.Trace("RemoveTask", "alloc_id", task.AllocID, "task", task.Name)
	m.ops = append(m.ops, NewMockConsulOp("remove", task.AllocID, task.Name))
}

func (m *MockConsulServiceClient) AllocRegistrations(allocID string) (*consul.AllocRegistration, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.logger.Trace("AllocRegistrations", "alloc_id", allocID)
	m.ops = append(m.ops, NewMockConsulOp("alloc_registrations", allocID, ""))

	if m.AllocRegistrationsFn != nil {
		return m.AllocRegistrationsFn(allocID)
	}

	return nil, nil
}

func (m *MockConsulServiceClient) GetOps() []MockConsulOp {
	m.mu.Lock()
	defer m.mu.Unlock()
	return m.ops
}
<commit_msg>test: add some extra logging<commit_after>package consul

import (
	"fmt"
	"sync"

	log "github.com/hashicorp/go-hclog"

	"github.com/hashicorp/nomad/command/agent/consul"
	testing "github.com/mitchellh/go-testing-interface"
)

// MockConsulOp represents the register/deregister operations.
type MockConsulOp struct {
	Op      string // add, remove, or update
	AllocID string
	Task    string
}

func NewMockConsulOp(op, allocID, task string) MockConsulOp {
	if op != "add" && op != "remove" && op != "update" && op != "alloc_registrations" {
		panic(fmt.Errorf("invalid consul op: %s", op))
	}
	return MockConsulOp{
		Op:      op,
		AllocID: allocID,
		Task:    task,
	}
}

// MockConsulServiceClient implements the ConsulServiceAPI interface to record
// and log task registration/deregistration.
type MockConsulServiceClient struct {
	ops []MockConsulOp
	mu  sync.Mutex

	logger log.Logger

	// AllocRegistrationsFn allows injecting return values for the
	// AllocRegistrations function.
	AllocRegistrationsFn func(allocID string) (*consul.AllocRegistration, error)
}

func NewMockConsulServiceClient(t testing.T, logger log.Logger) *MockConsulServiceClient {
	logger = logger.Named("mock_consul")
	m := MockConsulServiceClient{
		ops:    make([]MockConsulOp, 0, 20),
		logger: logger,
	}
	return &m
}

func (m *MockConsulServiceClient) UpdateTask(old, newSvcs *consul.TaskServices) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.logger.Trace("UpdateTask", "alloc_id", newSvcs.AllocID, "task", newSvcs.Name,
		"old_services", len(old.Services), "new_services", len(newSvcs.Services),
	)
	m.ops = append(m.ops, NewMockConsulOp("update", newSvcs.AllocID, newSvcs.Name))
	return nil
}

func (m *MockConsulServiceClient) RegisterTask(task *consul.TaskServices) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.logger.Trace("RegisterTask", "alloc_id", task.AllocID, "task", task.Name,
		"services", len(task.Services),
	)
	m.ops = append(m.ops, NewMockConsulOp("add", task.AllocID, task.Name))
	return nil
}

func (m *MockConsulServiceClient) RemoveTask(task *consul.TaskServices) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.logger.Trace("RemoveTask", "alloc_id", task.AllocID, "task", task.Name,
		"services", len(task.Services),
	)
	m.ops = append(m.ops, NewMockConsulOp("remove", task.AllocID, task.Name))
}

func (m *MockConsulServiceClient) AllocRegistrations(allocID string) (*consul.AllocRegistration, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.logger.Trace("AllocRegistrations", "alloc_id", allocID)
	m.ops = append(m.ops, NewMockConsulOp("alloc_registrations", allocID, ""))

	if m.AllocRegistrationsFn != nil {
		return m.AllocRegistrationsFn(allocID)
	}

	return nil, nil
}

func (m *MockConsulServiceClient) GetOps() []MockConsulOp {
	m.mu.Lock()
	defer m.mu.Unlock()
	return m.ops
}
<|endoftext|>


<commit_before>package ripo

import (
	"fmt"
	"log"
	"strings"
)

var errorDispatcher = func(request ExtendedRequest, rpcErr RPCError) {
	parts := []string{
		fmt.Sprintf("Code=%v", rpcErr.Code()),
		fmt.Sprintf("Message=%#v", rpcErr.Message()),
	}
	if rpcErr.Cause() != nil {
		parts = append(parts, fmt.Sprintf("Original=%#v", rpcErr.Cause().Error()))
	}
	if len(rpcErr.Details()) > 0 {
		parts = append(parts, fmt.Sprintf("Details=%#v", rpcErr.Details()))
	}
	tbLines := []string{
		"Traceback:",
	}
	for _, record := range rpcErr.Traceback(request.HandlerName()).Records() {
		tbLines = append(tbLines, fmt.Sprintf(
			"\tFile %v, Func %v, Line %v",
			record.File(),
			record.FunctionLocal(),
			record.Line(),
		))
	}
	parts = append(parts, strings.Join(tbLines, "\n"))
	log.Printf("RPCError: %v, \n", strings.Join(parts, ", "))

}

func SetErrorDispatcher(dispatcher func(request ExtendedRequest, rpcErr RPCError)) {
	errorDispatcher = dispatcher
}
<commit_msg>print Cause= instead of Original= in default errorDispatcher<commit_after>package ripo

import (
	"fmt"
	"log"
	"strings"
)

var errorDispatcher = func(request ExtendedRequest, rpcErr RPCError) {
	parts := []string{
		fmt.Sprintf("Code=%v", rpcErr.Code()),
		fmt.Sprintf("Message=%#v", rpcErr.Message()),
	}
	if rpcErr.Cause() != nil {
		parts = append(parts, fmt.Sprintf("Cause=%#v", rpcErr.Cause().Error()))
	}
	if len(rpcErr.Details()) > 0 {
		parts = append(parts, fmt.Sprintf("Details=%#v", rpcErr.Details()))
	}
	tbLines := []string{
		"Traceback:",
	}
	for _, record := range rpcErr.Traceback(request.HandlerName()).Records() {
		tbLines = append(tbLines, fmt.Sprintf(
			"\tFile %v, Func %v, Line %v",
			record.File(),
			record.FunctionLocal(),
			record.Line(),
		))
	}
	parts = append(parts, strings.Join(tbLines, "\n"))
	log.Printf("RPCError: %v, \n", strings.Join(parts, ", "))

}

func SetErrorDispatcher(dispatcher func(request ExtendedRequest, rpcErr RPCError)) {
	errorDispatcher = dispatcher
}
<|endoftext|>


<commit_before>// Copyright (c) 2013 The AUTHORS
//
// Use of this source code is governed by The MIT License
// that can be found in the LICENSE file.

package archiver

import (
	"archive/tar"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
)

type tarArchiver struct {
	opts Options
}

func newTarArchiver(opts Options) *tarArchiver {
	return &tarArchiver{opts}
}

func (archiver *tarArchiver) Archive(srcDir string) (archive *os.File, err error) {
	// Make sure the artifacts source directory exists and is not empty.
	dir, err := os.Open(srcDir)
	if err != nil {
		return nil, err
	}
	defer dir.Close()

	info, err := dir.Readdir(1)
	if err != nil {
		return nil, err
	}

	if len(info) == 0 {
		return nil, ErrNoArtifacts
	}

	// Pack the artifacts directory.
	wd, err := os.Getwd()
	if err != nil {
		return nil, err
	}

	ar, err := ioutil.TempFile(wd, "artifacts_archive_")
	if err != nil {
		return nil, err
	}

	if archiver.opts.Verbose() {
		fmt.Println("Packing artifacts")
	}

	aw := tar.NewWriter(ar)

	err = filepath.Walk(srcDir, func(path string, info os.FileInfo, err error) error {
		// Stop on error.
		if err != nil {
			return err
		}

		relative := path[len(srcDir):]

		// Skip root.
		if len(relative) == 0 {
			return nil
		}

		// Drop the leading slash.
		relative = relative[1:]

		// Prepare tar header.
		header, err := tar.FileInfoHeader(info, "")
		if err != nil {
			return err
		}

		// Append a trailing slash if this is a directory.
		if info.IsDir() {
			relative = fmt.Sprintf("%v%c", filepath.Clean(relative), os.PathSeparator)
		}

		if archiver.opts.Verbose() {
			fmt.Println("   ", relative)
		}

		// Tar always uses '/' as the separator.
		header.Name = filepath.ToSlash(relative)

		if archiver.opts.Dry() {
			header.Size = 0
		}

		// Write tar header.
		if err := aw.WriteHeader(header); err != nil {
			return err
		}

		// Do not write directories.
		if info.IsDir() {
			return nil
		}

		// Open the artifacts file.
		file, err := os.Open(path)
		if err != nil {
			return err
		}
		defer file.Close()

		// Copy the file into the archive.
		if !archiver.opts.Dry() {
			_, err = io.Copy(aw, file)
		}
		return err
	})
	if err != nil {
		aw.Close()
		ar.Close()
		return nil, err
	}

	if archiver.opts.Verbose() {
		fmt.Println("Archive created")
	}

	// Make sure we close the archive writer properly.
	if err := aw.Close(); err != nil {
		ar.Close()
		return nil, err
	}

	// Rewind to the beginning of the archive, otherwise the following reads
	// will return no data at all.
	if _, err := ar.Seek(0, os.SEEK_SET); err != nil {
		ar.Close()
		return nil, err
	}

	// Return the archive file, open and set to offset 0.
	return ar, nil
}
<commit_msg>utils/archiver/tar: Remove tmp archive on error<commit_after>// Copyright (c) 2013 The AUTHORS
//
// Use of this source code is governed by The MIT License
// that can be found in the LICENSE file.

package archiver

import (
	"archive/tar"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
)

type tarArchiver struct {
	opts Options
}

func newTarArchiver(opts Options) *tarArchiver {
	return &tarArchiver{opts}
}

func (archiver *tarArchiver) Archive(srcDir string) (archive *os.File, err error) {
	// Make sure the artifacts source directory exists and is not empty.
	dir, err := os.Open(srcDir)
	if err != nil {
		return nil, err
	}
	defer dir.Close()

	info, err := dir.Readdir(1)
	if err != nil {
		return nil, err
	}

	if len(info) == 0 {
		return nil, ErrNoArtifacts
	}

	// Pack the artifacts directory.
	wd, err := os.Getwd()
	if err != nil {
		return nil, err
	}

	ar, err := ioutil.TempFile(wd, "artifacts_archive_")
	if err != nil {
		return nil, err
	}

	if archiver.opts.Verbose() {
		fmt.Println("Packing artifacts")
	}

	aw := tar.NewWriter(ar)

	err = filepath.Walk(srcDir, func(path string, info os.FileInfo, err error) error {
		// Stop on error.
		if err != nil {
			return err
		}

		relative := path[len(srcDir):]

		// Skip root.
		if len(relative) == 0 {
			return nil
		}

		// Drop the leading slash.
		relative = relative[1:]

		// Prepare tar header.
		header, err := tar.FileInfoHeader(info, "")
		if err != nil {
			return err
		}

		// Append a trailing slash if this is a directory.
		if info.IsDir() {
			relative = fmt.Sprintf("%v%c", filepath.Clean(relative), os.PathSeparator)
		}

		if archiver.opts.Verbose() {
			fmt.Println("   ", relative)
		}

		// Tar always uses '/' as the separator.
		header.Name = filepath.ToSlash(relative)

		if archiver.opts.Dry() {
			header.Size = 0
		}

		// Write tar header.
		if err := aw.WriteHeader(header); err != nil {
			return err
		}

		// Do not write directories.
		if info.IsDir() {
			return nil
		}

		// Open the artifacts file.
		file, err := os.Open(path)
		if err != nil {
			return err
		}
		defer file.Close()

		// Copy the file into the archive.
		if !archiver.opts.Dry() {
			_, err = io.Copy(aw, file)
		}
		return err
	})
	if err != nil {
		aw.Close()
		ar.Close()
		os.Remove(ar.Name())
		return nil, err
	}

	if archiver.opts.Verbose() {
		fmt.Println("Archive created")
	}

	// Make sure we close the archive writer properly.
	if err := aw.Close(); err != nil {
		ar.Close()
		os.Remove(ar.Name())
		return nil, err
	}

	// Rewind to the beginning of the archive, otherwise the following reads
	// will return no data at all.
	if _, err := ar.Seek(0, os.SEEK_SET); err != nil {
		ar.Close()
		os.Remove(ar.Name())
		return nil, err
	}

	// Return the archive file, open and set to offset 0.
	return ar, nil
}
<|endoftext|>


<commit_before>package testutil

import (
	"fmt"
	"strings"
	"testing"

	"github.com/prometheus/client_golang/prometheus"
	prommodel "github.com/prometheus/client_model/go"
)

func AssertMetricValue(t *testing.T, g prometheus.Gatherer, n string, exp float64) {
	t.Helper()

	mf, err := g.Gather()
	if err != nil {
		t.Fatal(err)
	}

	name := strings.TrimSuffix(n, "_sum")
	name = strings.TrimSuffix(name, "_count")

	for _, m := range mf {
		if m.GetName() == name {
			var got float64
			for _, metric := range m.GetMetric() {
				switch m.GetType() {
				case prommodel.MetricType_COUNTER:
					got += metric.GetCounter().GetValue()
				case prommodel.MetricType_GAUGE:
					got += metric.GetGauge().GetValue()
				case prommodel.MetricType_HISTOGRAM:
					switch {
					case strings.HasSuffix(n, "_sum"):
						got += float64(metric.GetHistogram().GetSampleSum())
					case strings.HasSuffix(n, "_count"):
						got += float64(metric.GetHistogram().GetSampleCount())
					}
				}
			}
			if got != exp {
				t.Errorf("metric %s value do not match, expected %g but got %g", n, exp, got)
				return
			} else {
				//t.Logf("metric %s has expected value %g", n, exp)
				return
			}
		}
	}

	t.Errorf("metric %s does not exists", n)
}

func AssertMetricDimensions(t *testing.T, g prometheus.Gatherer, n string, dimensions map[string]string) {
	t.Helper()

	mf, err := g.Gather()
	if err != nil {
		t.Fatal(err)
	}

	name := strings.TrimSuffix(n, "_sum")
	name = strings.TrimSuffix(name, "_count")

	var sb strings.Builder
	for _, m := range mf {
		if m.GetName() == name {
			//var got float64
			for _, metric := range m.GetMetric() {
				var match int
			GivenLabels:
				for _, given := range metric.GetLabel() {
					for expectedKey, expectedValue := range dimensions {
						if given.GetName() != expectedKey {
							//t.Log("dont match 1", m.GetName(), given.GetName(), given.GetValue(), expectedKey, expectedValue)
							continue
						}
						if given.GetValue() != expectedValue {
							//t.Log("dont match 2", m.GetName(), given.GetName(), given.GetValue(), expectedKey, expectedValue)
							continue
						}
						match++
						//t.Log("match", m.GetName(), given.GetName(), given.GetValue(), expectedKey, expectedValue)
						continue GivenLabels
					}
				}
				if match == len(metric.GetLabel()) {
					return
				}
				sb.WriteString(fmt.Sprintf("metric checked, but does not match: %s\n	%v\n	%v\n", m.GetName(), metric.GetLabel(), dimensions))
			}
		}
	}

	t.Errorf("metric %s with dimensions %v does not exists", n, dimensions)
	t.Error(sb.String())
}
<commit_msg>v4/internal/assert - commented out code removed<commit_after>package testutil

import (
	"fmt"
	"strings"
	"testing"

	"github.com/prometheus/client_golang/prometheus"
	prommodel "github.com/prometheus/client_model/go"
)

func AssertMetricValue(t *testing.T, g prometheus.Gatherer, n string, exp float64) {
	t.Helper()

	mf, err := g.Gather()
	if err != nil {
		t.Fatal(err)
	}

	name := strings.TrimSuffix(n, "_sum")
	name = strings.TrimSuffix(name, "_count")

	for _, m := range mf {
		if m.GetName() == name {
			var got float64
			for _, metric := range m.GetMetric() {
				switch m.GetType() {
				case prommodel.MetricType_COUNTER:
					got += metric.GetCounter().GetValue()
				case prommodel.MetricType_GAUGE:
					got += metric.GetGauge().GetValue()
				case prommodel.MetricType_HISTOGRAM:
					switch {
					case strings.HasSuffix(n, "_sum"):
						got += float64(metric.GetHistogram().GetSampleSum())
					case strings.HasSuffix(n, "_count"):
						got += float64(metric.GetHistogram().GetSampleCount())
					}
				}
			}
			if got != exp {
				t.Errorf("metric %s value do not match, expected %g but got %g", n, exp, got)
				return
			} else {
				return
			}
		}
	}

	t.Errorf("metric %s does not exists", n)
}

func AssertMetricDimensions(t *testing.T, g prometheus.Gatherer, n string, dimensions map[string]string) {
	t.Helper()

	mf, err := g.Gather()
	if err != nil {
		t.Fatal(err)
	}

	name := strings.TrimSuffix(n, "_sum")
	name = strings.TrimSuffix(name, "_count")

	var sb strings.Builder
	for _, m := range mf {
		if m.GetName() == name {
			//var got float64
			for _, metric := range m.GetMetric() {
				var match int
			GivenLabels:
				for _, given := range metric.GetLabel() {
					for expectedKey, expectedValue := range dimensions {
						if given.GetName() != expectedKey {
							continue
						}
						if given.GetValue() != expectedValue {
							continue
						}
						match++
						continue GivenLabels
					}
				}
				if match == len(metric.GetLabel()) {
					return
				}
				sb.WriteString(fmt.Sprintf("metric checked, but does not match: %s\n	%v\n	%v\n", m.GetName(), metric.GetLabel(), dimensions))
			}
		}
	}

	t.Errorf("metric %s with dimensions %v does not exists", n, dimensions)
	t.Error(sb.String())
}
<|endoftext|>


<commit_before>package kubernetes

import (
	"bytes"
	"fmt"
	"strings"

	"github.com/go-kit/kit/log"
	"gopkg.in/yaml.v2"
	"k8s.io/apiextensions-apiserver/pkg/apis/apiextensions/v1beta1"
	"k8s.io/apimachinery/pkg/runtime/schema"

	kresource "github.com/weaveworks/flux/cluster/kubernetes/resource"
	"github.com/weaveworks/flux/image"
	"github.com/weaveworks/flux/resource"
)

// ResourceScopes maps resource definitions (GroupVersionKind) to whether they are namespaced or not
type ResourceScopes map[schema.GroupVersionKind]v1beta1.ResourceScope

// namespacer assigns namespaces to manifests that need it (or "" if
// the manifest should not have a namespace.
type namespacer interface {
	// EffectiveNamespace gives the namespace that would be used were
	// the manifest to be applied. This may be "", indicating that it
	// should not have a namespace (i.e., it's a cluster-level
	// resource).
	EffectiveNamespace(manifest kresource.KubeManifest, knownScopes ResourceScopes) (string, error)
}

// manifests is an implementation of cluster.Manifests, particular to
// Kubernetes. Aside from loading manifests from files, it does some
// "post-processsing" to make sure the view of the manifests is what
// would be applied; in particular, it fills in the namespace of
// manifests that would be given a default namespace when applied.
type manifests struct {
	namespacer       namespacer
	logger           log.Logger
	resourceWarnings map[string]struct{}
}

func NewManifests(ns namespacer, logger log.Logger) *manifests {
	return &manifests{
		namespacer:       ns,
		logger:           logger,
		resourceWarnings: map[string]struct{}{},
	}
}

func getCRDScopes(manifests map[string]kresource.KubeManifest) ResourceScopes {
	result := ResourceScopes{}
	for _, km := range manifests {
		if km.GetKind() == "CustomResourceDefinition" {
			var crd v1beta1.CustomResourceDefinition
			if err := yaml.Unmarshal(km.Bytes(), &crd); err != nil {
				// The CRD can't be parsed, so we (intentionally) ignore it and
				// just hope for EffectiveNamespace() to find its scope in the cluster if needed.
				continue
			}
			crdVersions := crd.Spec.Versions
			if len(crdVersions) == 0 {
				crdVersions = []v1beta1.CustomResourceDefinitionVersion{{Name: crd.Spec.Version}}
			}
			for _, crdVersion := range crdVersions {
				gvk := schema.GroupVersionKind{
					Group:   crd.Spec.Group,
					Version: crdVersion.Name,
					Kind:    crd.Spec.Names.Kind,
				}
				result[gvk] = crd.Spec.Scope
			}
		}
	}
	return result
}

func (m *manifests) setEffectiveNamespaces(manifests map[string]kresource.KubeManifest) (map[string]resource.Resource, error) {
	knownScopes := getCRDScopes(manifests)
	result := map[string]resource.Resource{}
	for _, km := range manifests {
		resID := km.ResourceID()
		resIDStr := resID.String()
		ns, err := m.namespacer.EffectiveNamespace(km, knownScopes)
		if err != nil {
			if strings.Contains(err.Error(), "not found") {
				// discard the resource and keep going after making sure we logged about it
				if _, warningLogged := m.resourceWarnings[resIDStr]; !warningLogged {
					_, kind, name := resID.Components()
					partialResIDStr := kind + "/" + name
					m.logger.Log(
						"warn", fmt.Sprintf("cannot find scope of resource %s: %s", partialResIDStr, err),
						"impact", fmt.Sprintf("resource %s will be excluded until its scope is available", partialResIDStr))
					m.resourceWarnings[resIDStr] = struct{}{}
				}
				continue
			}
			return nil, err
		}
		km.SetNamespace(ns)
		if _, warningLogged := m.resourceWarnings[resIDStr]; warningLogged {
			// indicate that we found the resource's scope and allow logging a warning again
			m.logger.Log("info", fmt.Sprintf("found scope of resource %s, back in bussiness!", km.ResourceID().String()))
			delete(m.resourceWarnings, resIDStr)
		}
		result[km.ResourceID().String()] = km
	}
	return result, nil
}

func (m *manifests) LoadManifests(baseDir string, paths []string) (map[string]resource.Resource, error) {
	manifests, err := kresource.Load(baseDir, paths)
	if err != nil {
		return nil, err
	}
	return m.setEffectiveNamespaces(manifests)
}

func (m *manifests) ParseManifest(def []byte, source string) (map[string]resource.Resource, error) {
	resources, err := kresource.ParseMultidoc(def, source)
	if err != nil {
		return nil, err
	}
	// Note: setEffectiveNamespaces() won't work for CRD instances whose CRD is yet to be created
	// (due to the CRD not being present in kresources).
	// We could get out of our way to fix this (or give a better error) but:
	// 1. With the exception of HelmReleases CRD instances are not workloads anyways.
	// 2. The problem is eventually fixed by the first successful sync.
	result, err := m.setEffectiveNamespaces(resources)
	if err != nil {
		return nil, err
	}
	return result, nil
}

func (m *manifests) SetWorkloadContainerImage(def []byte, id resource.ID, container string, image image.Ref) ([]byte, error) {
	return updateWorkload(def, id, container, image)
}

func (m *manifests) CreateManifestPatch(originalManifests, modifiedManifests []byte, originalSource, modifiedSource string) ([]byte, error) {
	return createManifestPatch(originalManifests, modifiedManifests, originalSource, modifiedSource)
}

func (m *manifests) ApplyManifestPatch(originalManifests, patchManifests []byte, originalSource, patchSource string) ([]byte, error) {
	return applyManifestPatch(originalManifests, patchManifests, originalSource, patchSource)
}

func (m *manifests) AppendManifestToBuffer(manifest []byte, buffer *bytes.Buffer) error {
	return appendYAMLToBuffer(manifest, buffer)
}

func appendYAMLToBuffer(manifest []byte, buffer *bytes.Buffer) error {
	separator := "---\n"
	bytes := buffer.Bytes()
	if len(bytes) > 0 && bytes[len(bytes)-1] != '\n' {
		separator = "\n---\n"
	}
	if _, err := buffer.WriteString(separator); err != nil {
		return fmt.Errorf("cannot write to internal buffer: %s", err)
	}
	if _, err := buffer.Write(manifest); err != nil {
		return fmt.Errorf("cannot write to internal buffer: %s", err)
	}
	return nil
}

// UpdateWorkloadPolicies in policies.go
<commit_msg>Correct bussiness to business<commit_after>package kubernetes

import (
	"bytes"
	"fmt"
	"strings"

	"github.com/go-kit/kit/log"
	"gopkg.in/yaml.v2"
	"k8s.io/apiextensions-apiserver/pkg/apis/apiextensions/v1beta1"
	"k8s.io/apimachinery/pkg/runtime/schema"

	kresource "github.com/weaveworks/flux/cluster/kubernetes/resource"
	"github.com/weaveworks/flux/image"
	"github.com/weaveworks/flux/resource"
)

// ResourceScopes maps resource definitions (GroupVersionKind) to whether they are namespaced or not
type ResourceScopes map[schema.GroupVersionKind]v1beta1.ResourceScope

// namespacer assigns namespaces to manifests that need it (or "" if
// the manifest should not have a namespace.
type namespacer interface {
	// EffectiveNamespace gives the namespace that would be used were
	// the manifest to be applied. This may be "", indicating that it
	// should not have a namespace (i.e., it's a cluster-level
	// resource).
	EffectiveNamespace(manifest kresource.KubeManifest, knownScopes ResourceScopes) (string, error)
}

// manifests is an implementation of cluster.Manifests, particular to
// Kubernetes. Aside from loading manifests from files, it does some
// "post-processsing" to make sure the view of the manifests is what
// would be applied; in particular, it fills in the namespace of
// manifests that would be given a default namespace when applied.
type manifests struct {
	namespacer       namespacer
	logger           log.Logger
	resourceWarnings map[string]struct{}
}

func NewManifests(ns namespacer, logger log.Logger) *manifests {
	return &manifests{
		namespacer:       ns,
		logger:           logger,
		resourceWarnings: map[string]struct{}{},
	}
}

func getCRDScopes(manifests map[string]kresource.KubeManifest) ResourceScopes {
	result := ResourceScopes{}
	for _, km := range manifests {
		if km.GetKind() == "CustomResourceDefinition" {
			var crd v1beta1.CustomResourceDefinition
			if err := yaml.Unmarshal(km.Bytes(), &crd); err != nil {
				// The CRD can't be parsed, so we (intentionally) ignore it and
				// just hope for EffectiveNamespace() to find its scope in the cluster if needed.
				continue
			}
			crdVersions := crd.Spec.Versions
			if len(crdVersions) == 0 {
				crdVersions = []v1beta1.CustomResourceDefinitionVersion{{Name: crd.Spec.Version}}
			}
			for _, crdVersion := range crdVersions {
				gvk := schema.GroupVersionKind{
					Group:   crd.Spec.Group,
					Version: crdVersion.Name,
					Kind:    crd.Spec.Names.Kind,
				}
				result[gvk] = crd.Spec.Scope
			}
		}
	}
	return result
}

func (m *manifests) setEffectiveNamespaces(manifests map[string]kresource.KubeManifest) (map[string]resource.Resource, error) {
	knownScopes := getCRDScopes(manifests)
	result := map[string]resource.Resource{}
	for _, km := range manifests {
		resID := km.ResourceID()
		resIDStr := resID.String()
		ns, err := m.namespacer.EffectiveNamespace(km, knownScopes)
		if err != nil {
			if strings.Contains(err.Error(), "not found") {
				// discard the resource and keep going after making sure we logged about it
				if _, warningLogged := m.resourceWarnings[resIDStr]; !warningLogged {
					_, kind, name := resID.Components()
					partialResIDStr := kind + "/" + name
					m.logger.Log(
						"warn", fmt.Sprintf("cannot find scope of resource %s: %s", partialResIDStr, err),
						"impact", fmt.Sprintf("resource %s will be excluded until its scope is available", partialResIDStr))
					m.resourceWarnings[resIDStr] = struct{}{}
				}
				continue
			}
			return nil, err
		}
		km.SetNamespace(ns)
		if _, warningLogged := m.resourceWarnings[resIDStr]; warningLogged {
			// indicate that we found the resource's scope and allow logging a warning again
			m.logger.Log("info", fmt.Sprintf("found scope of resource %s, back in business!", km.ResourceID().String()))
			delete(m.resourceWarnings, resIDStr)
		}
		result[km.ResourceID().String()] = km
	}
	return result, nil
}

func (m *manifests) LoadManifests(baseDir string, paths []string) (map[string]resource.Resource, error) {
	manifests, err := kresource.Load(baseDir, paths)
	if err != nil {
		return nil, err
	}
	return m.setEffectiveNamespaces(manifests)
}

func (m *manifests) ParseManifest(def []byte, source string) (map[string]resource.Resource, error) {
	resources, err := kresource.ParseMultidoc(def, source)
	if err != nil {
		return nil, err
	}
	// Note: setEffectiveNamespaces() won't work for CRD instances whose CRD is yet to be created
	// (due to the CRD not being present in kresources).
	// We could get out of our way to fix this (or give a better error) but:
	// 1. With the exception of HelmReleases CRD instances are not workloads anyways.
	// 2. The problem is eventually fixed by the first successful sync.
	result, err := m.setEffectiveNamespaces(resources)
	if err != nil {
		return nil, err
	}
	return result, nil
}

func (m *manifests) SetWorkloadContainerImage(def []byte, id resource.ID, container string, image image.Ref) ([]byte, error) {
	return updateWorkload(def, id, container, image)
}

func (m *manifests) CreateManifestPatch(originalManifests, modifiedManifests []byte, originalSource, modifiedSource string) ([]byte, error) {
	return createManifestPatch(originalManifests, modifiedManifests, originalSource, modifiedSource)
}

func (m *manifests) ApplyManifestPatch(originalManifests, patchManifests []byte, originalSource, patchSource string) ([]byte, error) {
	return applyManifestPatch(originalManifests, patchManifests, originalSource, patchSource)
}

func (m *manifests) AppendManifestToBuffer(manifest []byte, buffer *bytes.Buffer) error {
	return appendYAMLToBuffer(manifest, buffer)
}

func appendYAMLToBuffer(manifest []byte, buffer *bytes.Buffer) error {
	separator := "---\n"
	bytes := buffer.Bytes()
	if len(bytes) > 0 && bytes[len(bytes)-1] != '\n' {
		separator = "\n---\n"
	}
	if _, err := buffer.WriteString(separator); err != nil {
		return fmt.Errorf("cannot write to internal buffer: %s", err)
	}
	if _, err := buffer.Write(manifest); err != nil {
		return fmt.Errorf("cannot write to internal buffer: %s", err)
	}
	return nil
}

// UpdateWorkloadPolicies in policies.go
<|endoftext|>


<commit_before>package provider

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"strings"

	"github.com/NetSys/quilt/util"
	homedir "github.com/mitchellh/go-homedir"
)

var vagrantCmd = "vagrant"
var shCmd = "sh"

type vagrantAPI struct{}

func newVagrantAPI() vagrantAPI {
	vagrant := vagrantAPI{}
	return vagrant
}

func (api vagrantAPI) Init(cloudConfig string, size string, id string) error {
	vdir, err := api.VagrantDir()
	if err != nil {
		return err
	}
	if _, err := os.Stat(vdir); os.IsNotExist(err) {
		os.Mkdir(vdir, os.ModeDir|os.ModePerm)
	}
	path := vdir + id
	os.Mkdir(path, os.ModeDir|os.ModePerm)

	_, err = api.Shell(id, `vagrant --machine-readable init coreos-beta`)
	if err != nil {
		api.Destroy(id)
		return err
	}

	err = util.WriteFile(path+"/user-data", []byte(cloudConfig), 0644)
	if err != nil {
		api.Destroy(id)
		return err
	}

	vagrant := vagrantFile()
	err = util.WriteFile(path+"/vagrantFile", []byte(vagrant), 0644)
	if err != nil {
		api.Destroy(id)
		return err
	}

	err = util.WriteFile(path+"/size", []byte(size), 0644)
	if err != nil {
		api.Destroy(id)
		return err
	}

	return nil
}

func (api vagrantAPI) Up(id string) error {
	_, err := api.Shell(id, `vagrant --machine-readable up`)
	if err != nil {
		return err
	}
	return nil
}

func (api vagrantAPI) Destroy(id string) error {
	_, err := api.Shell(id,
		`vagrant --machine-readable destroy -f; cd ../; rm -rf %s`)
	if err != nil {
		return err
	}
	return nil
}

func (api vagrantAPI) PublicIP(id string) (string, error) {
	ip, err := api.Shell(id,
		`vagrant ssh -c "ip address show eth1 | grep 'inet ' | " +
		"sed -e 's/^.*inet //' -e 's/\/.*$//' | tr -d '\n'"`)
	if err != nil {
		return "", err
	}
	return string(ip[:]), nil
}

func (api vagrantAPI) Status(id string) (string, error) {
	output, err := api.Shell(id, `vagrant --machine-readable status`)
	if err != nil {
		return "", err
	}
	lines := bytes.Split(output, []byte("\n"))
	for _, line := range lines {
		words := strings.Split(string(line[:]), ",")
		if len(words) >= 4 {
			if strings.Compare(words[2], "state") == 0 {
				return words[3], nil
			}
		}
	}
	return "", nil
}

func (api vagrantAPI) List() ([]string, error) {
	subdirs := []string{}
	vdir, err := api.VagrantDir()
	if err != nil {
		return nil, err
	}
	if _, err := os.Stat(vdir); os.IsNotExist(err) {
		return subdirs, nil
	}

	files, err := ioutil.ReadDir(vdir)
	if err != nil {
		return subdirs, err
	}
	for _, file := range files {
		subdirs = append(subdirs, file.Name())
	}
	return subdirs, nil
}

func (api vagrantAPI) AddBox(name string, provider string) error {
	/* Adding a box fails if it already exists, hence the check. */
	exists, err := api.ContainsBox(name)
	if err != nil {
		return err
	}
	if exists {
		return nil
	}
	err = exec.Command(vagrantCmd, []string{"--machine-readable", "box", "add",
		"--provider", provider, name}...).Run()
	if err != nil {
		return err
	}
	return nil
}

func (api vagrantAPI) ContainsBox(name string) (bool, error) {
	output, err := exec.Command(vagrantCmd, []string{"--machine-readable", "box",
		"list"}...).Output()
	if err != nil {
		return false, err
	}
	lines := bytes.Split(output, []byte("\n"))
	for _, line := range lines {
		words := strings.Split(string(line[:]), ",")
		if words[len(words)-1] == name {
			return true, nil
		}
	}
	return false, nil
}

func (api vagrantAPI) Shell(id string, commands string) ([]byte, error) {
	chdir := `(cd %s; `
	vdir, err := api.VagrantDir()
	if err != nil {
		return nil, err
	}
	chdir = fmt.Sprintf(chdir, vdir+id)
	shellCommand := chdir + strings.Replace(commands, "%s", id, -1) + ")"
	output, err := exec.Command(shCmd, []string{"-c", shellCommand}...).Output()
	return output, err
}

func (api vagrantAPI) VagrantDir() (string, error) {
	dir, err := homedir.Dir()
	if err != nil {
		return "", err
	}
	vagrantDir := dir + "/.vagrant/"
	return vagrantDir, nil
}

func (api vagrantAPI) Size(id string) string {
	size, err := api.Shell(id, "cat size")
	if err != nil {
		return ""
	}
	return string(size)
}

func vagrantFile() string {
	vagrantfile := `CLOUD_CONFIG_PATH = File.join(File.dirname(__FILE__), "user-data")
SIZE_PATH = File.join(File.dirname(__FILE__), "size")
Vagrant.require_version ">= 1.6.0"

size = File.open(SIZE_PATH).read.strip.split(",")
Vagrant.configure(2) do |config|
  config.vm.box = "boxcutter/ubuntu1504"

  config.vm.network "private_network", type: "dhcp"

  ram=(size[0].to_f*1024).to_i
  cpus=size[1]
  config.vm.provider "virtualbox" do |v|
    v.memory = ram
    v.cpus = cpus
  end

  if File.exist?(CLOUD_CONFIG_PATH)
    config.vm.provision "shell", path: "#{CLOUD_CONFIG_PATH}"
  end
end
`
	return vagrantfile
}

// VagrantCreateSize creates an encoded string representing the amount of RAM
// and number of CPUs for an instance.
func (api vagrantAPI) CreateSize(ram, cpu float64) string {
	if ram < 1 {
		ram = 1
	}
	if cpu < 1 {
		cpu = 1
	}
	return fmt.Sprintf("%g,%g", ram, cpu)
}
<commit_msg>vagrant: Error with more context<commit_after>package provider

import (
	"bytes"
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"strings"

	"github.com/NetSys/quilt/util"
	homedir "github.com/mitchellh/go-homedir"
)

var vagrantCmd = "vagrant"
var shCmd = "sh"

type vagrantAPI struct{}

func newVagrantAPI() vagrantAPI {
	vagrant := vagrantAPI{}
	return vagrant
}

func (api vagrantAPI) Init(cloudConfig string, size string, id string) error {
	vdir, err := api.VagrantDir()
	if err != nil {
		return err
	}
	if _, err := os.Stat(vdir); os.IsNotExist(err) {
		os.Mkdir(vdir, os.ModeDir|os.ModePerm)
	}
	path := vdir + id
	os.Mkdir(path, os.ModeDir|os.ModePerm)

	_, err = api.Shell(id, `vagrant --machine-readable init coreos-beta`)
	if err != nil {
		api.Destroy(id)
		return errors.New("unable to init machine")
	}

	err = util.WriteFile(path+"/user-data", []byte(cloudConfig), 0644)
	if err != nil {
		api.Destroy(id)
		return err
	}

	vagrant := vagrantFile()
	err = util.WriteFile(path+"/vagrantFile", []byte(vagrant), 0644)
	if err != nil {
		api.Destroy(id)
		return err
	}

	err = util.WriteFile(path+"/size", []byte(size), 0644)
	if err != nil {
		api.Destroy(id)
		return err
	}

	return nil
}

func (api vagrantAPI) Up(id string) error {
	_, err := api.Shell(id, `vagrant --machine-readable up`)
	if err != nil {
		return errors.New("unable to check machine status")
	}
	return nil
}

func (api vagrantAPI) Destroy(id string) error {
	_, err := api.Shell(id,
		`vagrant --machine-readable destroy -f; cd ../; rm -rf %s`)
	if err != nil {
		return errors.New("unable to destroy machine")
	}
	return nil
}

func (api vagrantAPI) PublicIP(id string) (string, error) {
	ip, err := api.Shell(id,
		`vagrant ssh -c "ip address show eth1 | grep 'inet ' | " +
		"sed -e 's/^.*inet //' -e 's/\/.*$//' | tr -d '\n'"`)
	if err != nil {
		return "", err
	}
	return string(ip[:]), nil
}

func (api vagrantAPI) Status(id string) (string, error) {
	output, err := api.Shell(id, `vagrant --machine-readable status`)
	if err != nil {
		return "", errors.New("unable to retrieve machine status")
	}
	lines := bytes.Split(output, []byte("\n"))
	for _, line := range lines {
		words := strings.Split(string(line[:]), ",")
		if len(words) >= 4 {
			if strings.Compare(words[2], "state") == 0 {
				return words[3], nil
			}
		}
	}
	return "", nil
}

func (api vagrantAPI) List() ([]string, error) {
	subdirs := []string{}
	vdir, err := api.VagrantDir()
	if err != nil {
		return nil, err
	}
	if _, err := os.Stat(vdir); os.IsNotExist(err) {
		return subdirs, nil
	}

	files, err := ioutil.ReadDir(vdir)
	if err != nil {
		return subdirs, err
	}
	for _, file := range files {
		subdirs = append(subdirs, file.Name())
	}
	return subdirs, nil
}

func (api vagrantAPI) AddBox(name string, provider string) error {
	/* Adding a box fails if it already exists, hence the check. */
	exists, err := api.ContainsBox(name)
	if err == nil && exists {
		return nil
	}
	err = exec.Command(vagrantCmd, []string{"--machine-readable", "box", "add",
		"--provider", provider, name}...).Run()
	if err != nil {
		return errors.New("unable to add box")
	}
	return nil
}

func (api vagrantAPI) ContainsBox(name string) (bool, error) {
	output, err := exec.Command(vagrantCmd, []string{"--machine-readable", "box",
		"list"}...).Output()
	if err != nil {
		return false, errors.New("unable to list machines")
	}
	lines := bytes.Split(output, []byte("\n"))
	for _, line := range lines {
		words := strings.Split(string(line[:]), ",")
		if words[len(words)-1] == name {
			return true, nil
		}
	}
	return false, nil
}

func (api vagrantAPI) Shell(id string, commands string) ([]byte, error) {
	chdir := `(cd %s; `
	vdir, err := api.VagrantDir()
	if err != nil {
		return nil, err
	}
	chdir = fmt.Sprintf(chdir, vdir+id)
	shellCommand := chdir + strings.Replace(commands, "%s", id, -1) + ")"
	output, err := exec.Command(shCmd, []string{"-c", shellCommand}...).Output()
	return output, err
}

func (api vagrantAPI) VagrantDir() (string, error) {
	dir, err := homedir.Dir()
	if err != nil {
		return "", err
	}
	vagrantDir := dir + "/.vagrant/"
	return vagrantDir, nil
}

func (api vagrantAPI) Size(id string) string {
	size, err := api.Shell(id, "cat size")
	if err != nil {
		return ""
	}
	return string(size)
}

func vagrantFile() string {
	vagrantfile := `CLOUD_CONFIG_PATH = File.join(File.dirname(__FILE__), "user-data")
SIZE_PATH = File.join(File.dirname(__FILE__), "size")
Vagrant.require_version ">= 1.6.0"

size = File.open(SIZE_PATH).read.strip.split(",")
Vagrant.configure(2) do |config|
  config.vm.box = "boxcutter/ubuntu1504"

  config.vm.network "private_network", type: "dhcp"

  ram=(size[0].to_f*1024).to_i
  cpus=size[1]
  config.vm.provider "virtualbox" do |v|
    v.memory = ram
    v.cpus = cpus
  end

  if File.exist?(CLOUD_CONFIG_PATH)
    config.vm.provision "shell", path: "#{CLOUD_CONFIG_PATH}"
  end
end
`
	return vagrantfile
}

// VagrantCreateSize creates an encoded string representing the amount of RAM
// and number of CPUs for an instance.
func (api vagrantAPI) CreateSize(ram, cpu float64) string {
	if ram < 1 {
		ram = 1
	}
	if cpu < 1 {
		cpu = 1
	}
	return fmt.Sprintf("%g,%g", ram, cpu)
}
<|endoftext|>


<commit_before>/*
Copyright 2018 The Kubernetes Authors.

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 util

import (
	"bytes"
	"encoding/json"
	"fmt"
	"time"
)

// GetString tries to return value from map cast to string type. If value doesn't exist, error is returned.
func GetString(dict map[string]interface{}, key string) (string, error) {
	return getString(dict, key, "", fmt.Errorf("key %s not found", key))
}

// GetInt tries to return value from map cast to int type. If value doesn't exist, error is returned.
func GetInt(dict map[string]interface{}, key string) (int, error) {
	return getInt(dict, key, 0, fmt.Errorf("key %s not found", key))
}

// GetFloat64 tries to return value from map cast to float64 type. If value doesn't exist, error is returned.
func GetFloat64(dict map[string]interface{}, key string) (float64, error) {
	return getFloat64(dict, key, 0, fmt.Errorf("key %s not found", key))
}

// GetDuration tries to return value from map cast to duration type. If value doesn't exist, error is returned.
func GetDuration(dict map[string]interface{}, key string) (time.Duration, error) {
	return getDuration(dict, key, 0, fmt.Errorf("key %s not found", key))
}

// GetStringOrDefault tries to return value from map cast to string type. If value doesn't exist default value is used.
func GetStringOrDefault(dict map[string]interface{}, key string, defaultValue string) (string, error) {
	return getString(dict, key, defaultValue, nil)
}

// GetIntOrDefault tries to return value from map cast to int type. If value doesn't exist default value is used.
func GetIntOrDefault(dict map[string]interface{}, key string, defaultValue int) (int, error) {
	return getInt(dict, key, defaultValue, nil)
}

// GetFloat64OrDefault tries to return value from map cast to float64 type. If value doesn't exist default value is used.
func GetFloat64OrDefault(dict map[string]interface{}, key string, defaultValue float64) (float64, error) {
	return getFloat64(dict, key, defaultValue, nil)
}

// GetDurationOrDefault tries to return value from map cast to duration type. If value doesn't exist default value is used.
func GetDurationOrDefault(dict map[string]interface{}, key string, defaultValue time.Duration) (time.Duration, error) {
	return getDuration(dict, key, defaultValue, nil)
}

func getString(dict map[string]interface{}, key string, defaultValue string, defaultError error) (string, error) {
	value, exists := dict[key]
	if !exists {
		return defaultValue, defaultError
	}

	stringValue, ok := value.(string)
	if !ok {
		return "", fmt.Errorf("type assertion error: %v is not a string", value)
	}
	return stringValue, nil
}

func getInt(dict map[string]interface{}, key string, defaultValue int, defaultError error) (int, error) {
	value, exists := dict[key]
	if !exists {
		return defaultValue, defaultError
	}

	// Types from interface{} create from json cannot be cast directly to int.
	floatValue, ok := value.(float64)
	if !ok {
		return 0, fmt.Errorf("type assertion error: %v of is not an int", value)
	}
	return int(floatValue), nil
}

func getFloat64(dict map[string]interface{}, key string, defaultValue float64, defaultError error) (float64, error) {
	value, exists := dict[key]
	if !exists {
		return defaultValue, defaultError
	}

	floatValue, ok := value.(float64)
	if !ok {
		return 0, fmt.Errorf("type assertion error: %v of is not an int", value)
	}
	return floatValue, nil
}

func getDuration(dict map[string]interface{}, key string, defaultValue time.Duration, defaultError error) (time.Duration, error) {
	defaultErr := fmt.Errorf("key %s not found", key)
	durationString, err := getString(dict, key, "", defaultErr)
	if err != nil {
		if err == defaultErr {
			return defaultValue, defaultError
		}
		return 0, err
	}

	duration, err := time.ParseDuration(durationString)
	if err != nil {
		return 0, fmt.Errorf("parsing duration error: %v", err)
	}
	return duration, nil
}

// PrettyPrintJSON converts given data into formatted json.
func PrettyPrintJSON(data interface{}) (string, error) {
	output := &bytes.Buffer{}
	if err := json.NewEncoder(output).Encode(data); err != nil {
		return "", fmt.Errorf("building encoder error: %v", err)
	}
	formatted := &bytes.Buffer{}
	if err := json.Indent(formatted, output.Bytes(), "", "  "); err != nil {
		return "", fmt.Errorf("indenting error: %v", err)
	}
	return string(formatted.Bytes()), nil
}
<commit_msg>int and float parsing fix<commit_after>/*
Copyright 2018 The Kubernetes Authors.

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 util

import (
	"bytes"
	"encoding/json"
	"fmt"
	"strconv"
	"time"
)

// GetString tries to return value from map cast to string type. If value doesn't exist, error is returned.
func GetString(dict map[string]interface{}, key string) (string, error) {
	return getString(dict, key, "", fmt.Errorf("key %s not found", key))
}

// GetInt tries to return value from map cast to int type. If value doesn't exist, error is returned.
func GetInt(dict map[string]interface{}, key string) (int, error) {
	return getInt(dict, key, 0, fmt.Errorf("key %s not found", key))
}

// GetFloat64 tries to return value from map cast to float64 type. If value doesn't exist, error is returned.
func GetFloat64(dict map[string]interface{}, key string) (float64, error) {
	return getFloat64(dict, key, 0, fmt.Errorf("key %s not found", key))
}

// GetDuration tries to return value from map cast to duration type. If value doesn't exist, error is returned.
func GetDuration(dict map[string]interface{}, key string) (time.Duration, error) {
	return getDuration(dict, key, 0, fmt.Errorf("key %s not found", key))
}

// GetStringOrDefault tries to return value from map cast to string type. If value doesn't exist default value is used.
func GetStringOrDefault(dict map[string]interface{}, key string, defaultValue string) (string, error) {
	return getString(dict, key, defaultValue, nil)
}

// GetIntOrDefault tries to return value from map cast to int type. If value doesn't exist default value is used.
func GetIntOrDefault(dict map[string]interface{}, key string, defaultValue int) (int, error) {
	return getInt(dict, key, defaultValue, nil)
}

// GetFloat64OrDefault tries to return value from map cast to float64 type. If value doesn't exist default value is used.
func GetFloat64OrDefault(dict map[string]interface{}, key string, defaultValue float64) (float64, error) {
	return getFloat64(dict, key, defaultValue, nil)
}

// GetDurationOrDefault tries to return value from map cast to duration type. If value doesn't exist default value is used.
func GetDurationOrDefault(dict map[string]interface{}, key string, defaultValue time.Duration) (time.Duration, error) {
	return getDuration(dict, key, defaultValue, nil)
}

func getString(dict map[string]interface{}, key string, defaultValue string, defaultError error) (string, error) {
	value, exists := dict[key]
	if !exists {
		return defaultValue, defaultError
	}

	stringValue, ok := value.(string)
	if !ok {
		return "", fmt.Errorf("type assertion error: %v is not a string", value)
	}
	return stringValue, nil
}

func getInt(dict map[string]interface{}, key string, defaultValue int, defaultError error) (int, error) {
	value, exists := dict[key]
	if !exists {
		return defaultValue, defaultError
	}

	intValue, ok := value.(int)
	if ok {
		return intValue, nil
	}
	// Types from interface{} create from json cannot be cast directly to int.
	floatValue, ok := value.(float64)
	if ok {
		return int(floatValue), nil
	}
	stringValue, ok := value.(string)
	if ok {
		if i, err := strconv.Atoi(stringValue); err != nil {
			return i, nil
		}
	}
	return 0, fmt.Errorf("type assertion error: %v of is not an int", value)
}

func getFloat64(dict map[string]interface{}, key string, defaultValue float64, defaultError error) (float64, error) {
	value, exists := dict[key]
	if !exists {
		return defaultValue, defaultError
	}

	floatValue, ok := value.(float64)
	if ok {
		return floatValue, nil
	}
	stringValue, ok := value.(string)
	if ok {
		if f, err := strconv.ParseFloat(stringValue, 64); err != nil {
			return f, nil
		}
	}
	return 0, fmt.Errorf("type assertion error: %v of is not a float", value)
}

func getDuration(dict map[string]interface{}, key string, defaultValue time.Duration, defaultError error) (time.Duration, error) {
	defaultErr := fmt.Errorf("key %s not found", key)
	durationString, err := getString(dict, key, "", defaultErr)
	if err != nil {
		if err == defaultErr {
			return defaultValue, defaultError
		}
		return 0, err
	}

	duration, err := time.ParseDuration(durationString)
	if err != nil {
		return 0, fmt.Errorf("parsing duration error: %v", err)
	}
	return duration, nil
}

// PrettyPrintJSON converts given data into formatted json.
func PrettyPrintJSON(data interface{}) (string, error) {
	output := &bytes.Buffer{}
	if err := json.NewEncoder(output).Encode(data); err != nil {
		return "", fmt.Errorf("building encoder error: %v", err)
	}
	formatted := &bytes.Buffer{}
	if err := json.Indent(formatted, output.Bytes(), "", "  "); err != nil {
		return "", fmt.Errorf("indenting error: %v", err)
	}
	return string(formatted.Bytes()), nil
}
<|endoftext|>


<commit_before>// Copyright 2017 Jeff Foley. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

package sources

import (
	"fmt"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/OWASP/Amass/config"
	eb "github.com/OWASP/Amass/eventbus"
	"github.com/OWASP/Amass/resolvers"
	"github.com/OWASP/Amass/services"
	"github.com/OWASP/Amass/utils"
	"github.com/PuerkitoBio/goquery"
	"github.com/geziyor/geziyor"
	"github.com/geziyor/geziyor/client"
)

var (
	nameStripRE = regexp.MustCompile("^((20)|(25)|(2b)|(2f)|(3d)|(3a)|(40))+")
	maxCrawlSem = utils.NewSimpleSemaphore(50)
)

// GetAllSources returns a slice of all data source services, initialized and ready.
func GetAllSources(cfg *config.Config, bus *eb.EventBus, pool *resolvers.ResolverPool) []services.Service {
	return []services.Service{
		NewAlienVault(cfg, bus, pool),
		NewArchiveIt(cfg, bus, pool),
		NewArchiveToday(cfg, bus, pool),
		NewArquivo(cfg, bus, pool),
		NewAsk(cfg, bus, pool),
		NewBaidu(cfg, bus, pool),
		NewBinaryEdge(cfg, bus, pool),
		NewBing(cfg, bus, pool),
		NewBufferOver(cfg, bus, pool),
		NewCensys(cfg, bus, pool),
		NewCertSpotter(cfg, bus, pool),
		NewCIRCL(cfg, bus, pool),
		NewCommonCrawl(cfg, bus, pool),
		NewCrtsh(cfg, bus, pool),
		NewDNSDB(cfg, bus, pool),
		NewDNSDumpster(cfg, bus, pool),
		NewDNSTable(cfg, bus, pool),
		NewDogpile(cfg, bus, pool),
		NewEntrust(cfg, bus, pool),
		NewExalead(cfg, bus, pool),
		NewGoogle(cfg, bus, pool),
		NewHackerOne(cfg, bus, pool),
		NewHackerTarget(cfg, bus, pool),
		NewIPv4Info(cfg, bus, pool),
		NewLoCArchive(cfg, bus, pool),
		NewMnemonic(cfg, bus, pool),
		NewNetcraft(cfg, bus, pool),
		NewNetworksDB(cfg, bus, pool),
		NewOpenUKArchive(cfg, bus, pool),
		NewPassiveTotal(cfg, bus, pool),
		NewPTRArchive(cfg, bus, pool),
		NewRADb(cfg, bus, pool),
		NewRiddler(cfg, bus, pool),
		NewRobtex(cfg, bus, pool),
		NewSiteDossier(cfg, bus, pool),
		NewSecurityTrails(cfg, bus, pool),
		NewShadowServer(cfg, bus, pool),
		NewShodan(cfg, bus, pool),
		NewSpyse(cfg, bus, pool),
		NewSublist3rAPI(cfg, bus, pool),
		NewTeamCymru(cfg, bus, pool),
		NewThreatCrowd(cfg, bus, pool),
		NewTwitter(cfg, bus, pool),
		NewUKGovArchive(cfg, bus, pool),
		NewUmbrella(cfg, bus, pool),
		NewURLScan(cfg, bus, pool),
		NewViewDNS(cfg, bus, pool),
		NewVirusTotal(cfg, bus, pool),
		NewWayback(cfg, bus, pool),
		NewYahoo(cfg, bus, pool),
	}
}

// Clean up the names scraped from the web.
func cleanName(name string) string {
	name = strings.TrimSpace(strings.ToLower(name))

	for {
		if i := nameStripRE.FindStringIndex(name); i != nil {
			name = name[i[1]:]
		} else {
			break
		}
	}

	name = strings.Trim(name, "-")
	// Remove dots at the beginning of names
	if len(name) > 1 && name[0] == '.' {
		name = name[1:]
	}
	return name
}

func crawl(service services.Service, baseURL, baseDomain, subdomain, domain string) ([]string, error) {
	var results []string

	maxCrawlSem.Acquire(1)
	defer maxCrawlSem.Release(1)

	re := service.Config().DomainRegex(domain)
	if re == nil {
		return results, fmt.Errorf("crawler error: Failed to obtain regex object for: %s", domain)
	}

	start := fmt.Sprintf("%s/%s/%s", baseURL, strconv.Itoa(time.Now().Year()), subdomain)
	geziyor.NewGeziyor(&geziyor.Options{
		AllowedDomains:              []string{baseDomain},
		StartURLs:                   []string{start},
		Timeout:                     30 * time.Second,
		UserAgent:                   utils.UserAgent,
		RequestDelay:                time.Second,
		RequestDelayRandomize:       true,
		LogDisabled:                 true,
		ConcurrentRequests:          3,
		ConcurrentRequestsPerDomain: 3,
		ParseFunc: func(g *geziyor.Geziyor, r *client.Response) {
			r.HTMLDoc.Find("a").Each(func(i int, s *goquery.Selection) {
				if href, ok := s.Attr("href"); ok {
					if sub := re.FindString(r.JoinURL(href)); sub != "" {
						results = utils.UniqueAppend(results, cleanName(sub))
					}
				}
			})
		},
	}).Start()

	return results, nil
}
<commit_msg>added Google certificate transparency as a data source<commit_after>// Copyright 2017 Jeff Foley. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

package sources

import (
	"fmt"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/OWASP/Amass/config"
	eb "github.com/OWASP/Amass/eventbus"
	"github.com/OWASP/Amass/resolvers"
	"github.com/OWASP/Amass/services"
	"github.com/OWASP/Amass/utils"
	"github.com/PuerkitoBio/goquery"
	"github.com/geziyor/geziyor"
	"github.com/geziyor/geziyor/client"
)

var (
	nameStripRE = regexp.MustCompile("^((20)|(25)|(2b)|(2f)|(3d)|(3a)|(40))+")
	maxCrawlSem = utils.NewSimpleSemaphore(50)
)

// GetAllSources returns a slice of all data source services, initialized and ready.
func GetAllSources(cfg *config.Config, bus *eb.EventBus, pool *resolvers.ResolverPool) []services.Service {
	return []services.Service{
		NewAlienVault(cfg, bus, pool),
		NewArchiveIt(cfg, bus, pool),
		NewArchiveToday(cfg, bus, pool),
		NewArquivo(cfg, bus, pool),
		NewAsk(cfg, bus, pool),
		NewBaidu(cfg, bus, pool),
		NewBinaryEdge(cfg, bus, pool),
		NewBing(cfg, bus, pool),
		NewBufferOver(cfg, bus, pool),
		NewCensys(cfg, bus, pool),
		NewCertSpotter(cfg, bus, pool),
		NewCIRCL(cfg, bus, pool),
		NewCommonCrawl(cfg, bus, pool),
		NewCrtsh(cfg, bus, pool),
		NewDNSDB(cfg, bus, pool),
		NewDNSDumpster(cfg, bus, pool),
		NewDNSTable(cfg, bus, pool),
		NewDogpile(cfg, bus, pool),
		NewEntrust(cfg, bus, pool),
		NewExalead(cfg, bus, pool),
		NewGoogle(cfg, bus, pool),
		NewGoogleCT(cfg, bus, pool),
		NewHackerOne(cfg, bus, pool),
		NewHackerTarget(cfg, bus, pool),
		NewIPv4Info(cfg, bus, pool),
		NewLoCArchive(cfg, bus, pool),
		NewMnemonic(cfg, bus, pool),
		NewNetcraft(cfg, bus, pool),
		NewNetworksDB(cfg, bus, pool),
		NewOpenUKArchive(cfg, bus, pool),
		NewPassiveTotal(cfg, bus, pool),
		NewPTRArchive(cfg, bus, pool),
		NewRADb(cfg, bus, pool),
		NewRiddler(cfg, bus, pool),
		NewRobtex(cfg, bus, pool),
		NewSiteDossier(cfg, bus, pool),
		NewSecurityTrails(cfg, bus, pool),
		NewShadowServer(cfg, bus, pool),
		NewShodan(cfg, bus, pool),
		NewSpyse(cfg, bus, pool),
		NewSublist3rAPI(cfg, bus, pool),
		NewTeamCymru(cfg, bus, pool),
		NewThreatCrowd(cfg, bus, pool),
		NewTwitter(cfg, bus, pool),
		NewUKGovArchive(cfg, bus, pool),
		NewUmbrella(cfg, bus, pool),
		NewURLScan(cfg, bus, pool),
		NewViewDNS(cfg, bus, pool),
		NewVirusTotal(cfg, bus, pool),
		NewWayback(cfg, bus, pool),
		NewYahoo(cfg, bus, pool),
	}
}

// Clean up the names scraped from the web.
func cleanName(name string) string {
	name = strings.TrimSpace(strings.ToLower(name))

	for {
		if i := nameStripRE.FindStringIndex(name); i != nil {
			name = name[i[1]:]
		} else {
			break
		}
	}

	name = strings.Trim(name, "-")
	// Remove dots at the beginning of names
	if len(name) > 1 && name[0] == '.' {
		name = name[1:]
	}
	return name
}

func crawl(service services.Service, baseURL, baseDomain, subdomain, domain string) ([]string, error) {
	var results []string

	maxCrawlSem.Acquire(1)
	defer maxCrawlSem.Release(1)

	re := service.Config().DomainRegex(domain)
	if re == nil {
		return results, fmt.Errorf("crawler error: Failed to obtain regex object for: %s", domain)
	}

	start := fmt.Sprintf("%s/%s/%s", baseURL, strconv.Itoa(time.Now().Year()), subdomain)
	geziyor.NewGeziyor(&geziyor.Options{
		AllowedDomains:              []string{baseDomain},
		StartURLs:                   []string{start},
		Timeout:                     30 * time.Second,
		UserAgent:                   utils.UserAgent,
		RequestDelay:                time.Second,
		RequestDelayRandomize:       true,
		LogDisabled:                 true,
		ConcurrentRequests:          3,
		ConcurrentRequestsPerDomain: 3,
		ParseFunc: func(g *geziyor.Geziyor, r *client.Response) {
			r.HTMLDoc.Find("a").Each(func(i int, s *goquery.Selection) {
				if href, ok := s.Attr("href"); ok {
					if sub := re.FindString(r.JoinURL(href)); sub != "" {
						results = utils.UniqueAppend(results, cleanName(sub))
					}
				}
			})
		},
	}).Start()

	return results, nil
}
<|endoftext|>


<commit_before>package webhook

import (
	"encoding/json"

	"github.com/plimble/errors"
	"github.com/plimble/validator"
)

type ServiceInterface interface {
	ParseGithubWebhook(payload string) (*Webhook, error)
	ParseGitlabWebhook(payload string) (*Webhook, error)
}

type Service struct{}

func New() *Service {
	return &Service{}
}

func (s *Service) ParseGithubWebhook(payload string) (*Webhook, error) {
	v := validator.New()
	v.RequiredString(payload, "payload")
	if err := v.GetError(); err != nil {
		return nil, errors.BadRequestErr(err)
	}

	gwh := &GithubWebHookResult{}
	err := json.Unmarshal([]byte(payload), gwh)

	return gwh.MapToWebHook(), errors.WithStack(err)
}

func (s *Service) ParseGitlabWebhook(payload string) (*Webhook, error) {
	v := validator.New()
	v.RequiredString(payload, "payload")
	if err := v.GetError(); err != nil {
		return nil, errors.BadRequestErr(err)
	}

	gwh := &GitlabWebhookResult{}
	err := json.Unmarshal([]byte(payload), gwh)

	return gwh.MapToWebHook(), errors.WithStack(err)
}
<commit_msg>remove validator<commit_after>package webhook

import (
	"encoding/json"

	"github.com/plimble/errors"
)

type ServiceInterface interface {
	ParseGithubWebhook(payload string) (*Webhook, error)
	ParseGitlabWebhook(payload string) (*Webhook, error)
}

type Service struct{}

func New() *Service {
	return &Service{}
}

func (s *Service) ParseGithubWebhook(payload string) (*Webhook, error) {
	if payload == "" {
		return nil, errors.BadRequest("payload is missing")
	}

	gwh := &GithubWebHookResult{}
	err := json.Unmarshal([]byte(payload), gwh)

	return gwh.MapToWebHook(), errors.WithStack(err)
}

func (s *Service) ParseGitlabWebhook(payload string) (*Webhook, error) {
	if payload == "" {
		return nil, errors.BadRequest("payload is missing")
	}

	gwh := &GitlabWebhookResult{}
	err := json.Unmarshal([]byte(payload), gwh)

	return gwh.MapToWebHook(), errors.WithStack(err)
}
<|endoftext|>


<commit_before>package node

import (
	"github.com/MG-RAST/Shock/shock-server/logger"
	"github.com/MG-RAST/golib/mgo/bson"
	"time"
)

type NodeReaper struct{}

func NewNodeReaper() *NodeReaper {
	return &NodeReaper{}
}

func (nr *NodeReaper) Handle() {
	for {
		// sleep
		time.Sleep(10 * time.Second)
		// query to get expired nodes
		nodes := Nodes{}
		query := nr.getQuery()
		nodes.GetAll(query)
		// delete expired nodes
		for _, n := range nodes {
			if err := n.Delete(); err != nil {
				err_msg := "err:@node_delete: " + err.Error()
				logger.Error(err_msg)
			}
		}
	}
}

func (nr *NodeReaper) getQuery() (query bson.M) {
	hasExpire := bson.M{"expiration": bson.M{"$exists": true}}   // has the field
	toExpire := bson.M{"expiration": bson.M{"$ne": time.Time{}}} // value has been set, not default
	isExpired := bson.M{"expiration": bson.M{"$gt": time.Now()}} // value is too old
	query = bson.M{"$and": []bson.M{hasExpire, toExpire, isExpired}}
	return
}
<commit_msg>log deletion<commit_after>package node

import (
	"github.com/MG-RAST/Shock/shock-server/logger"
	"github.com/MG-RAST/golib/mgo/bson"
	"time"
)

type NodeReaper struct{}

func NewNodeReaper() *NodeReaper {
	return &NodeReaper{}
}

func (nr *NodeReaper) Handle() {
	for {
		// sleep
		time.Sleep(10 * time.Second)
		// query to get expired nodes
		nodes := Nodes{}
		query := nr.getQuery()
		nodes.GetAll(query)
		// delete expired nodes
		for _, n := range nodes {
			logger.Info("access", "Deleting expired node: "+n.Id)
			if err := n.Delete(); err != nil {
				err_msg := "err:@node_delete: " + err.Error()
				logger.Error(err_msg)
			}
		}
	}
}

func (nr *NodeReaper) getQuery() (query bson.M) {
	hasExpire := bson.M{"expiration": bson.M{"$exists": true}}   // has the field
	toExpire := bson.M{"expiration": bson.M{"$ne": time.Time{}}} // value has been set, not default
	isExpired := bson.M{"expiration": bson.M{"$gt": time.Now()}} // value is too old
	query = bson.M{"$and": []bson.M{hasExpire, toExpire, isExpired}}
	return
}
<|endoftext|>


<commit_before>// Copyright 2016 IBM Corporation
//
//   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 discovery

import (
	"errors"
	"net"
	"net/http"

	log "github.com/Sirupsen/logrus"
	"github.com/amalgam8/amalgam8/registry/utils/logging"
	"github.com/ant0ine/go-json-rest/rest"
)

const (
	module = "SIDECARREGISTER"
)

// Server defines an interface for controlling server lifecycle
type Server interface {
	Start() error
	Stop()
}

type server struct {
	config   *Config
	listener net.Listener
	logger   *log.Entry
}

// NewDiscoveryServer creates a new server based on the provided configuration options.
// Returns a valid Server interface on success or an error on failure
func NewDiscoveryServer(conf *Config) (Server, error) {
	if conf == nil {
		return nil, errors.New("null Discovery configuration provided")
	}

	s := &server{
		config: conf,
		logger: logging.GetLogger(module),
	}

	s.logger.Infof("Creating Discovery REST API on %s", s.config.HTTPAddressSpec)

	return s, nil
}

func (s *server) Start() error {
	handler, err := s.setup()
	if err != nil {
		return err
	}

	go func() {
		err = s.serve(handler)
		if err != nil {
			s.logger.WithFields(log.Fields{
				"error": err,
			}).Warn("Failed to start the server")
		}
	}()

	return nil
}

func (s *server) Stop() {
	s.logger.Info("Stopping rest server")

	if s.listener != nil {
		if err := s.listener.Close(); err != nil {
			s.logger.WithFields(log.Fields{
				"error": err,
			}).Warn("Failed to close listener")
		}
	}
}

func (s *server) setup() (http.Handler, error) {
	restAPI := rest.NewApi()

	restAPI.Use(
		&rest.RecoverMiddleware{},
		&rest.TimerMiddleware{},
		&rest.RecorderMiddleware{},
		&rest.ContentTypeCheckerMiddleware{},
	)

	discoveryAPI := NewDiscovery(s.config.Discovery)

	var routes []*rest.Route
	routes = append(routes, discoveryAPI.Routes()...)

	router, err := rest.MakeRouter(routes...)
	if err != nil {
		return nil, err
	}

	restAPI.SetApp(router)
	return restAPI.MakeHandler(), nil
}

func (s *server) serve(h http.Handler) error {
	s.logger.Infof("Starting service discovery REST API on %s", s.config.HTTPAddressSpec)

	listener, err := net.Listen("tcp", s.config.HTTPAddressSpec)
	if err != nil {
		s.logger.WithFields(log.Fields{
			"error": err,
		}).Error("Failed to start the server")
		return err
	}

	s.listener = listener
	if err := http.Serve(listener, h); err != nil {
		s.logger.WithFields(log.Fields{
			"error": err,
		}).Warn("Server has aborted")
		return err
	}

	return nil
}
<commit_msg>code review comments for discovery<commit_after>// Copyright 2016 IBM Corporation
//
//   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 discovery

import (
	"errors"
	"net"
	"net/http"

	"sync"

	log "github.com/Sirupsen/logrus"
	"github.com/amalgam8/amalgam8/registry/utils/logging"
	"github.com/ant0ine/go-json-rest/rest"
)

const (
	module = "SIDECARREGISTER"
)

// Server defines an interface for controlling server lifecycle
type Server interface {
	Start() error
	Stop()
}

type server struct {
	config   *Config
	listener net.Listener
	logger   *log.Entry
	mutex    sync.Mutex
}

// NewDiscoveryServer creates a new server based on the provided configuration options.
// Returns a valid Server interface on success or an error on failure
func NewDiscoveryServer(conf *Config) (Server, error) {
	if conf == nil {
		return nil, errors.New("null Discovery configuration provided")
	}

	s := &server{
		config: conf,
		logger: logging.GetLogger(module),
	}

	s.logger.Infof("Creating Discovery REST API on %s", s.config.HTTPAddressSpec)

	return s, nil
}

func (s *server) Start() error {
	handler, err := s.setup()
	if err != nil {
		return err
	}

	go func() {
		err = s.serve(handler)
		if err != nil {
			s.logger.WithFields(log.Fields{
				"error": err,
			}).Warn("Failed to start the server")
		}
	}()

	return nil
}

func (s *server) Stop() {
	s.logger.Info("Stopping rest server")
	s.mutex.Lock()
	defer s.mutex.Unlock()
	if s.listener != nil {
		if err := s.listener.Close(); err != nil {
			s.logger.WithFields(log.Fields{
				"error": err,
			}).Warn("Failed to close listener")
		}
	}
}

func (s *server) setup() (http.Handler, error) {
	restAPI := rest.NewApi()

	restAPI.Use(
		&rest.RecoverMiddleware{},
		&rest.TimerMiddleware{},
		&rest.RecorderMiddleware{},
		&rest.ContentTypeCheckerMiddleware{},
	)

	discoveryAPI := NewDiscovery(s.config.Discovery)

	router, err := rest.MakeRouter(discoveryAPI.Routes()...)
	if err != nil {
		return nil, err
	}

	restAPI.SetApp(router)
	return restAPI.MakeHandler(), nil
}

func (s *server) serve(h http.Handler) error {
	s.logger.Infof("Starting service discovery REST API on %s", s.config.HTTPAddressSpec)

	listener, err := net.Listen("tcp", s.config.HTTPAddressSpec)
	if err != nil {
		s.logger.WithFields(log.Fields{
			"error": err,
		}).Error("Failed to start the server")
		return err
	}

	s.mutex.Lock()
	defer s.mutex.Unlock()
	s.listener = listener
	if err := http.Serve(listener, h); err != nil {
		s.logger.WithFields(log.Fields{
			"error": err,
		}).Warn("Server has aborted")
		return err
	}

	return nil
}
<|endoftext|>