<commit_before>package cloudca

import (
	"encoding/json"
	"github.com/cloud-ca/go-cloudca/api"
	"github.com/cloud-ca/go-cloudca/services"
)

const (
	VOLUME_TYPE_OS   = "OS"
	VOLUME_TYPE_DATA = "DATA"
)

type Volume struct {
	Id             string `json:"id,omitempty"`
	Name           string `json:"name,omitempty"`
	Type           string `json:"type,omitempty"`
	CreationDate   string `json:"creationDate,omitempty"`
	Size           int    `json:"size,omitempty"`
	DiskOfferingId string `json:"diskOfferingId,omitempty"`
	TemplateId     string `json:"templateId,omitempty"`
	StorageTier    string `json:"storageTier,omitempty"`
	ZoneName       string `json:"zoneName,omitempty"`
	State          string `json:"state,omitempty"`
	InstanceName   string `json:"instanceName,omitempty"`
	InstanceId     string `json:"instanceId,omitempty"`
	InstanceState  string `json:"instanceState,omitempty"`
}

type VolumeService interface {
	Get(id string) (*Volume, error)
	List() ([]Volume, error)
	ListOfType(volumeType string) ([]Volume, error)
	ListWithOptions(options map[string]string) ([]Volume, error)
}

type VolumeApi struct {
	entityService services.EntityService
}

func NewVolumeService(apiClient api.ApiClient, serviceCode string, environmentName string) VolumeService {
	return &VolumeApi{
		entityService: services.NewEntityService(apiClient, serviceCode, environmentName, VOLUME_ENTITY_TYPE),
	}
}

func parseVolume(data []byte) *Volume {
	volume := Volume{}
	json.Unmarshal(data, &volume)
	return &volume
}

func parseVolumeList(data []byte) []Volume {
	volumes := []Volume{}
	json.Unmarshal(data, &volumes)
	return volumes
}

//Get volume with the specified id for the current environment
func (volumeApi *VolumeApi) Get(id string) (*Volume, error) {
	data, err := volumeApi.entityService.Get(id, map[string]string{})
	if err != nil {
		return nil, err
	}
	return parseVolume(data), nil
}

//List all volumes for the current environment
func (volumeApi *VolumeApi) List() ([]Volume, error) {
	return volumeApi.ListWithOptions(map[string]string{})
}

//List all volumes of specified type for the current environment
func (volumeApi *VolumeApi) ListOfType(volumeType string) ([]Volume, error) {
	return volumeApi.ListWithOptions(map[string]string{
		"type": volumeType,
	})
}

//List all volumes for the current environment. Can use options to do sorting and paging.
func (volumeApi *VolumeApi) ListWithOptions(options map[string]string) ([]Volume, error) {
	data, err := volumeApi.entityService.List(options)
	if err != nil {
		return nil, err
	}
	return parseVolumeList(data), nil
}
<commit_msg>Add volume operations<commit_after>package cloudca

import (
	"encoding/json"
	"github.com/cloud-ca/go-cloudca/api"
	"github.com/cloud-ca/go-cloudca/services"
)

const (
	VOLUME_TYPE_OS   = "OS"
	VOLUME_TYPE_DATA = "DATA"
)

type Volume struct {
	Id             string `json:"id,omitempty"`
	Name           string `json:"name,omitempty"`
	Type           string `json:"type,omitempty"`
	CreationDate   string `json:"creationDate,omitempty"`
	Size           int    `json:"size,omitempty"`
	DiskOfferingId string `json:"diskOfferingId,omitempty"`
	TemplateId     string `json:"templateId,omitempty"`
	StorageTier    string `json:"storageTier,omitempty"`
	ZoneName       string `json:"zoneName,omitempty"`
	ZoneId         string `json:"zoneId,omitempty"`
	State          string `json:"state,omitempty"`
	InstanceName   string `json:"instanceName,omitempty"`
	InstanceId     string `json:"instanceId,omitempty"`
	InstanceState  string `json:"instanceState,omitempty"`
}

type VolumeService interface {
	Get(id string) (*Volume, error)
	List() ([]Volume, error)
	ListOfType(volumeType string) ([]Volume, error)
	ListWithOptions(options map[string]string) ([]Volume, error)
	Create(Volume) (*Volume, error)
	Delete(string) error
	AttachToInstance(*Volume, string) error
	DetachFromInstance(*Volume) error
}

type VolumeApi struct {
	entityService services.EntityService
}

func NewVolumeService(apiClient api.ApiClient, serviceCode string, environmentName string) VolumeService {
	return &VolumeApi{
		entityService: services.NewEntityService(apiClient, serviceCode, environmentName, VOLUME_ENTITY_TYPE),
	}
}

func parseVolume(data []byte) *Volume {
	volume := Volume{}
	json.Unmarshal(data, &volume)
	return &volume
}

func parseVolumeList(data []byte) []Volume {
	volumes := []Volume{}
	json.Unmarshal(data, &volumes)
	return volumes
}

//Get volume with the specified id for the current environment
func (volumeApi *VolumeApi) Get(id string) (*Volume, error) {
	data, err := volumeApi.entityService.Get(id, map[string]string{})
	if err != nil {
		return nil, err
	}
	return parseVolume(data), nil
}

//List all volumes for the current environment
func (volumeApi *VolumeApi) List() ([]Volume, error) {
	return volumeApi.ListWithOptions(map[string]string{})
}

//List all volumes of specified type for the current environment
func (volumeApi *VolumeApi) ListOfType(volumeType string) ([]Volume, error) {
	return volumeApi.ListWithOptions(map[string]string{
		"type": volumeType,
	})
}

//List all volumes for the current environment. Can use options to do sorting and paging.
func (volumeApi *VolumeApi) ListWithOptions(options map[string]string) ([]Volume, error) {
	data, err := volumeApi.entityService.List(options)
	if err != nil {
		return nil, err
	}
	return parseVolumeList(data), nil
}

func (api *VolumeApi) Create(volume Volume) (*Volume, error) {
	msg, err := json.Marshal(volume)
	if err != nil {
		return nil, err
	}
	res, err := api.entityService.Create(msg, map[string]string{})
	if err != nil {
		return nil, err
	}
	return parseVolume(res), nil
}

func (api *VolumeApi) Delete(volumeId string) error {
	_, err := api.entityService.Delete(volumeId, []byte{}, map[string]string{})
	return err
}

func (api *VolumeApi) AttachToInstance(volume *Volume, instanceId string) error {
	msg, err := json.Marshal(Volume{
		InstanceId: instanceId,
	})
	if err != nil {
		return err
	}
	_, err = api.entityService.Execute(volume.Id, "attachToInstance", msg, map[string]string{})
	return err
}

func (api *VolumeApi) DetachFromInstance(volume *Volume) error {
	_, err := api.entityService.Execute(volume.Id, "detachFromInstance", []byte{}, map[string]string{})
	return err
}
<|endoftext|>


<commit_before>package netsec

import (
	"atlantis/supervisor/containers/serialize"
	"errors"
	"log"
	"sync"
)

type NetworkSecurity struct {
	sync.Mutex
	SaveFile   string                        // where to save state
	DeniedIPs  map[string]bool               // list of denied IPs. map for easy existence check
	IPGroups   map[string][]string           // group name -> list of infrastructure IPs to blanket deny
	Containers map[string]*ContainerSecurity // container id -> ContainerSecurity
}

func New(saveFile string) *NetworkSecurity {
	return &NetworkSecurity{
		Mutex:      sync.Mutex{},
		SaveFile:   saveFile,
		DeniedIPs:  map[string]bool{},
		IPGroups:   map[string][]string{},
		Containers: map[string]*ContainerSecurity{},
	}
}

func (n *NetworkSecurity) save() {
	// save state
	serialize.SaveAll(serialize.SaveDefinition{
		n.SaveFile,
		n,
	})
}

func (n *NetworkSecurity) UpdateIPGroup(name string, ips []string) error {
	n.Lock()
	defer n.Unlock()

	if err := delConnTrackRule(); err != nil {
		return err
	}

	current, exists := n.IPGroups[name]
	toRemove := []string{}
	if exists {
		// figure out what is being removed
		incomingMap := map[string]bool{}
		for _, ip := range ips {
			incomingMap[ip] = true
		}
		for _, ip := range current {
			if !incomingMap[ip] {
				toRemove = append(toRemove, ip)
			}
		}
	}
	// add blanket deny rule for new IPs
	newIPs := []string{}
	for _, ip := range ips {
		if n.DeniedIPs[ip] {
			// already exists
			continue
		}
		if err := n.rejectIP(ip); err != nil {
			return err
		}
		n.DeniedIPs[ip] = true
		newIPs = append(newIPs, ip)
	}
	// remove blanket deny rule for removed IPs
	for _, ip := range toRemove {
		if !n.DeniedIPs[ip] {
			// already removed
			continue
		}
		n.allowIP(ip)
		delete(n.DeniedIPs, ip)
	}
	// add/remove forward rules for new IPs for everything that uses the name
	for _, contSec := range n.Containers {
		ports, exists := contSec.SecurityGroups[name]
		if !exists {
			// this container does not use this ip group
			continue
		}
		for _, port := range ports {
			// add new ips
			for _, ip := range newIPs {
				contSec.allowPort(ip, port)
			}
			// remove old ips
			for _, ip := range toRemove {
				contSec.rejectPort(ip, port)
			}
		}
	}
	// update ipGroups
	n.IPGroups[name] = ips
	n.save()

	return addConnTrackRule()
}

func (n *NetworkSecurity) DeleteIPGroup(name string) error {
	if err := n.UpdateIPGroup(name, []string{}); err != nil {
		return err
	}
	n.Lock()
	defer n.Unlock()
	delete(n.IPGroups, name)
	n.save()
	return nil
}

func (n *NetworkSecurity) AddContainerSecurity(id string, pid int, sgs map[string][]uint16) error {
	n.Lock()
	defer n.Unlock()
	log.Printf("[netsec] add container security: "+id+", pid: %d, sgs: %#v", pid, sgs)
	if _, exists := n.Containers[id]; exists {
		// we already have security set up for this id. don't do it and return an error.
		log.Println("[netsec] -- not adding, already existed for: " + id)
		return errors.New("Container " + id + " already has Network Security set up.")
	}
	// make sure all groups exist
	for group, _ := range sgs {
		_, exists := n.IPGroups[group]
		if !exists {
			log.Println("[netsec] -- not adding group " + group + " doesn't exist for: " + id)
			return errors.New("IP Group " + group + " does not exist")
		}
	}

	// fetch network info
	contSec, err := NewContainerSecurity(id, pid, sgs)
	if err != nil {
		log.Println("[netsec] -- guano error: " + err.Error())
		return err
	}
	log.Println("[netsec] --> contSec: " + contSec.String())
	contSec.addMark()

	// add forward rules
	for group, ports := range sgs {
		for _, port := range ports {
			ips := n.IPGroups[group]
			for _, ip := range ips {
				if err := contSec.allowPort(ip, port); err != nil {
					defer n.RemoveContainerSecurity(id) // cleanup created references when we error out
					log.Println("[netsec] -- allow port error: " + err.Error())
					return err
				}
			}
		}
	}
	n.Containers[id] = contSec
	n.save()
	log.Println("[netsec] -- added " + id)
	return nil
}

func (n *NetworkSecurity) RemoveContainerSecurity(id string) error {
	n.Lock()
	defer n.Unlock()
	log.Println("[netsec] remove container security: " + id)
	contSec, exists := n.Containers[id]
	if !exists {
		log.Println("[netsec] -- not removing, none existed for: " + id)
		// no container security here, nothing to remove
		return nil
	}

	log.Println("[netsec] --> contSec: " + contSec.String())
	contSec.delMark()
	// remove forward rules
	for group, ports := range contSec.SecurityGroups {
		for _, port := range ports {
			ips := n.IPGroups[group]
			for _, ip := range ips {
				contSec.rejectPort(ip, port)
			}
		}
	}
	delete(n.Containers, id)
	n.save()
	log.Println("[netsec] -- removed " + id)
	return nil
}

func delConnTrackRule() error {
	_, err := executeCommand("iptables", "-D", "FORWARD", "-o", "docker0", "-m", "conntrack", "--ctstate", "RELATED,ESTABLISHED", "-j", "ACCEPT")
	return err
}

func addConnTrackRule() error {
	_, err := executeCommand("iptables", "-I", "FORWARD", "-o", "docker0", "-m", "conntrack", "--ctstate", "RELATED,ESTABLISHED", "-j", "ACCEPT")
	return err
}

func (n *NetworkSecurity) forwardRule(action, ip string) error {
	_, err := executeCommand("iptables", action, "FORWARD", "-i", "docker0", "-d", ip, "-j", "REJECT")
	return err
}

func (n *NetworkSecurity) rejectIP(ip string) error {
	return n.forwardRule("-I", ip)
}

func (n *NetworkSecurity) allowIP(ip string) error {
	return n.forwardRule("-D", ip)
}
<commit_msg>remove docker0 references<commit_after>package netsec

import (
	"atlantis/supervisor/containers/serialize"
	"errors"
	"log"
	"sync"
)

type NetworkSecurity struct {
	sync.Mutex
	SaveFile   string                        // where to save state
	DeniedIPs  map[string]bool               // list of denied IPs. map for easy existence check
	IPGroups   map[string][]string           // group name -> list of infrastructure IPs to blanket deny
	Containers map[string]*ContainerSecurity // container id -> ContainerSecurity
}

func New(saveFile string) *NetworkSecurity {
	return &NetworkSecurity{
		Mutex:      sync.Mutex{},
		SaveFile:   saveFile,
		DeniedIPs:  map[string]bool{},
		IPGroups:   map[string][]string{},
		Containers: map[string]*ContainerSecurity{},
	}
}

func (n *NetworkSecurity) save() {
	// save state
	serialize.SaveAll(serialize.SaveDefinition{
		n.SaveFile,
		n,
	})
}

func (n *NetworkSecurity) UpdateIPGroup(name string, ips []string) error {
	n.Lock()
	defer n.Unlock()

	if err := delConnTrackRule(); err != nil {
		log.Println("[netsec] error deleting track rule: " + err.Error())
		// continue, probably we never added it in the first place. **crosses fingers**
	}

	current, exists := n.IPGroups[name]
	toRemove := []string{}
	if exists {
		// figure out what is being removed
		incomingMap := map[string]bool{}
		for _, ip := range ips {
			incomingMap[ip] = true
		}
		for _, ip := range current {
			if !incomingMap[ip] {
				toRemove = append(toRemove, ip)
			}
		}
	}
	// add blanket deny rule for new IPs
	newIPs := []string{}
	for _, ip := range ips {
		if n.DeniedIPs[ip] {
			// already exists
			continue
		}
		if err := n.rejectIP(ip); err != nil {
			return err
		}
		n.DeniedIPs[ip] = true
		newIPs = append(newIPs, ip)
	}
	// remove blanket deny rule for removed IPs
	for _, ip := range toRemove {
		if !n.DeniedIPs[ip] {
			// already removed
			continue
		}
		n.allowIP(ip)
		delete(n.DeniedIPs, ip)
	}
	// add/remove forward rules for new IPs for everything that uses the name
	for _, contSec := range n.Containers {
		ports, exists := contSec.SecurityGroups[name]
		if !exists {
			// this container does not use this ip group
			continue
		}
		for _, port := range ports {
			// add new ips
			for _, ip := range newIPs {
				contSec.allowPort(ip, port)
			}
			// remove old ips
			for _, ip := range toRemove {
				contSec.rejectPort(ip, port)
			}
		}
	}
	// update ipGroups
	n.IPGroups[name] = ips
	n.save()

	return addConnTrackRule()
}

func (n *NetworkSecurity) DeleteIPGroup(name string) error {
	if err := n.UpdateIPGroup(name, []string{}); err != nil {
		return err
	}
	n.Lock()
	defer n.Unlock()
	delete(n.IPGroups, name)
	n.save()
	return nil
}

func (n *NetworkSecurity) AddContainerSecurity(id string, pid int, sgs map[string][]uint16) error {
	n.Lock()
	defer n.Unlock()
	log.Printf("[netsec] add container security: "+id+", pid: %d, sgs: %#v", pid, sgs)
	if _, exists := n.Containers[id]; exists {
		// we already have security set up for this id. don't do it and return an error.
		log.Println("[netsec] -- not adding, already existed for: " + id)
		return errors.New("Container " + id + " already has Network Security set up.")
	}
	// make sure all groups exist
	for group, _ := range sgs {
		_, exists := n.IPGroups[group]
		if !exists {
			log.Println("[netsec] -- not adding group " + group + " doesn't exist for: " + id)
			return errors.New("IP Group " + group + " does not exist")
		}
	}

	// fetch network info
	contSec, err := NewContainerSecurity(id, pid, sgs)
	if err != nil {
		log.Println("[netsec] -- guano error: " + err.Error())
		return err
	}
	log.Println("[netsec] --> contSec: " + contSec.String())
	contSec.addMark()

	// add forward rules
	for group, ports := range sgs {
		for _, port := range ports {
			ips := n.IPGroups[group]
			for _, ip := range ips {
				if err := contSec.allowPort(ip, port); err != nil {
					defer n.RemoveContainerSecurity(id) // cleanup created references when we error out
					log.Println("[netsec] -- allow port error: " + err.Error())
					return err
				}
			}
		}
	}
	n.Containers[id] = contSec
	n.save()
	log.Println("[netsec] -- added " + id)
	return nil
}

func (n *NetworkSecurity) RemoveContainerSecurity(id string) error {
	n.Lock()
	defer n.Unlock()
	log.Println("[netsec] remove container security: " + id)
	contSec, exists := n.Containers[id]
	if !exists {
		log.Println("[netsec] -- not removing, none existed for: " + id)
		// no container security here, nothing to remove
		return nil
	}

	log.Println("[netsec] --> contSec: " + contSec.String())
	contSec.delMark()
	// remove forward rules
	for group, ports := range contSec.SecurityGroups {
		for _, port := range ports {
			ips := n.IPGroups[group]
			for _, ip := range ips {
				contSec.rejectPort(ip, port)
			}
		}
	}
	delete(n.Containers, id)
	n.save()
	log.Println("[netsec] -- removed " + id)
	return nil
}

func delConnTrackRule() error {
	_, err := executeCommand("iptables", "-D", "FORWARD", "-m", "conntrack", "--ctstate", "RELATED,ESTABLISHED", "-j", "ACCEPT")
	return err
}

func addConnTrackRule() error {
	_, err := executeCommand("iptables", "-I", "FORWARD", "-m", "conntrack", "--ctstate", "RELATED,ESTABLISHED", "-j", "ACCEPT")
	return err
}

func (n *NetworkSecurity) forwardRule(action, ip string) error {
	_, err := executeCommand("iptables", action, "FORWARD", "-d", ip, "-j", "REJECT")
	return err
}

func (n *NetworkSecurity) rejectIP(ip string) error {
	return n.forwardRule("-I", ip)
}

func (n *NetworkSecurity) allowIP(ip string) error {
	return n.forwardRule("-D", ip)
}
<|endoftext|>


<commit_before>// +build windows

package util

import (
	"os"
	"syscall"
	"unsafe"
)

/*
func IsTerminal(file *os.File) (bool, error) {
	s, e := file.Stat()
	if e != nil {
		return true, nil
	}
	m := s.Mode()
	if m&os.ModeDevice != 0 {
		return true, nil
	}
	return false, nil
}
*/

/*
Examples:

    term, err := IsTerminal(os.Stdin)
    term, err := IsTerminal(os.Stdout)
    term, err := IsTerminal(os.Stderr)
*/
func IsTerminal(file *os.File) bool {
	var st uint32
	return getConsoleMode(file.Fd(), &st) == nil
}

func getConsoleMode(hConsoleHandle syscall.Handle, lpMode *uint32) (err error) {
	r1, _, e1 := syscall.Syscall(procGetConsoleMode.Addr(), 2, uintptr(hConsoleHandle), uintptr(unsafe.Pointer(lpMode)), 0)
	if int(r1) == 0 {
		if e1 != 0 {
			err = error(e1)
		} else {
			err = syscall.EINVAL
		}
	}
	return
}
<commit_msg>test vervolg<commit_after>// +build windows

package util

import (
	"os"
	"syscall"
	"unsafe"
)

/*
func IsTerminal(file *os.File) (bool, error) {
	s, e := file.Stat()
	if e != nil {
		return true, nil
	}
	m := s.Mode()
	if m&os.ModeDevice != 0 {
		return true, nil
	}
	return false, nil
}
*/

/*
Examples:

    term, err := IsTerminal(os.Stdin)
    term, err := IsTerminal(os.Stdout)
    term, err := IsTerminal(os.Stderr)
*/
func IsTerminal(file *os.File) bool {
	var st uint32
	return getConsoleMode(syscall.Handle(file.Fd()), &st) == nil
}

func getConsoleMode(hConsoleHandle syscall.Handle, lpMode *uint32) (err error) {
	r1, _, e1 := syscall.Syscall(procGetConsoleMode.Addr(), 2, uintptr(hConsoleHandle), uintptr(unsafe.Pointer(lpMode)), 0)
	if int(r1) == 0 {
		if e1 != 0 {
			err = error(e1)
		} else {
			err = syscall.EINVAL
		}
	}
	return
}
<|endoftext|>


<commit_before>package aws

import (
	"fmt"
	"log"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/service/budgets"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/helper/schema"
)

func resourceAwsBudgetsBudget() *schema.Resource {
	return &schema.Resource{
		Schema: resourceAwsBudgetsBudgetSchema(),
		Create: resourceAwsBudgetsBudgetCreate,
		Read:   resourceAwsBudgetsBudgetRead,
		Update: resourceAwsBudgetsBudgetUpdate,
		Delete: resourceAwsBudgetsBudgetDelete,
	}
}

func resourceAwsBudgetsBudgetSchema() map[string]*schema.Schema {
	return map[string]*schema.Schema{
		"name": {
			Type:          schema.TypeString,
			Optional:      true,
			ConflictsWith: []string{"name_prefix"},
		},
		"name_prefix": {
			Type:     schema.TypeString,
			Optional: true,
			ForceNew: true,
		},
		"budget_type": {
			Type:     schema.TypeString,
			Required: true,
		},
		"limit_amount": {
			Type:     schema.TypeString,
			Required: true,
		},
		"limit_unit": {
			Type:     schema.TypeString,
			Required: true,
		},
		"include_credit": {
			Type:     schema.TypeBool,
			Optional: true,
			Default:  true,
		},
		"include_other_subscription": {
			Type:     schema.TypeBool,
			Optional: true,
			Default:  true,
		},
		"include_recurring": {
			Type:     schema.TypeBool,
			Optional: true,
			Default:  true,
		},
		"include_refund": {
			Type:     schema.TypeBool,
			Optional: true,
			Default:  true,
		},
		"include_subscription": {
			Type:     schema.TypeBool,
			Optional: true,
			Default:  true,
		},
		"include_support": {
			Type:     schema.TypeBool,
			Optional: true,
			Default:  true,
		},
		"include_tax": {
			Type:     schema.TypeBool,
			Optional: true,
			Default:  true,
		},
		"include_upfront": {
			Type:     schema.TypeBool,
			Optional: true,
			Default:  true,
		},
		"use_blended": {
			Type:     schema.TypeBool,
			Optional: true,
			Default:  false,
		},
		"time_period_start": {
			Type:     schema.TypeString,
			Required: true,
		},
		"time_period_end": {
			Type:     schema.TypeString,
			Optional: true,
			Default:  "2087-06-15_00:00",
		},
		"time_unit": {
			Type:     schema.TypeString,
			Required: true,
		},
		"cost_filters": {
			Type:     schema.TypeMap,
			Optional: true,
			Computed: true,
		},
	}
}

func resourceAwsBudgetsBudgetCreate(d *schema.ResourceData, meta interface{}) error {
	client := meta.(*AWSClient).budgetconn
	accountID := meta.(*AWSClient).accountid
	budget, err := newBudgetsBudget(d)
	if err != nil {
		return fmt.Errorf("failed creating budget: %v", err)
	}

	_, err = client.CreateBudget(&budgets.CreateBudgetInput{
		AccountId: &accountID,
		Budget:    budget,
	})
	if err != nil {
		return fmt.Errorf("create budget failed: %v", err)
	}

	d.SetId(*budget.BudgetName)
	return resourceAwsBudgetsBudgetUpdate(d, meta)
}

func resourceAwsBudgetsBudgetRead(d *schema.ResourceData, meta interface{}) error {
	budgetName := d.Id()
	describeBudgetOutput, err := describeBudget(budgetName, meta)
	if isBudgetNotFoundException(err) {
		d.SetId("")
		return nil
	}

	if err != nil {
		return fmt.Errorf("describe budget failed: %v", err)
	}

	flattenedBudget, err := expandBudgetsBudgetFlatten(describeBudgetOutput.Budget)
	if err != nil {
		return fmt.Errorf("failed flattening budget output: %v", err)
	}

	if _, ok := d.GetOk("name"); ok {
		d.Set("name", flattenedBudget.name)
	}

	for k, v := range map[string]interface{}{
		"budget_type":                flattenedBudget.budgetType,
		"time_unit":                  flattenedBudget.timeUnit,
		"cost_filters":               convertCostFiltersToStringMap(flattenedBudget.costFilters),
		"limit_amount":               flattenedBudget.limitAmount,
		"limit_unit":                 flattenedBudget.limitUnit,
		"include_credit":             flattenedBudget.includeCredit,
		"include_other_subscription": flattenedBudget.includeOtherSubscription,
		"include_recurring":          flattenedBudget.includeRecurring,
		"include_refund":             flattenedBudget.includeRefund,
		"include_subscription":       flattenedBudget.includeSubscription,
		"include_support":            flattenedBudget.includeSupport,
		"include_tax":                flattenedBudget.includeTax,
		"include_upfront":            flattenedBudget.includeUpFront,
		"use_blended":                flattenedBudget.useBlended,
		"time_period_start":          flattenedBudget.timePeriodStart.Format("2006-01-02_15:04"),
		"time_period_end":            flattenedBudget.timePeriodEnd.Format("2006-01-02_15:04"),
	} {
		if err := d.Set(k, v); err != nil {
			return err
		}
	}
	return nil
}

func resourceAwsBudgetsBudgetUpdate(d *schema.ResourceData, meta interface{}) error {
	client := meta.(*AWSClient).budgetconn
	accountID := meta.(*AWSClient).accountid
	budget, err := newBudgetsBudget(d)
	if err != nil {
		return fmt.Errorf("could not create budget: %v", err)
	}

	updateBudgetInput := new(budgets.UpdateBudgetInput)
	updateBudgetInput.SetAccountId(accountID)
	updateBudgetInput.SetNewBudget(budget)
	_, err = client.UpdateBudget(updateBudgetInput)
	if err != nil {
		return fmt.Errorf("updaate budget failed: %v", err)
	}

	return resourceAwsBudgetsBudgetRead(d, meta)
}

func resourceAwsBudgetsBudgetDelete(d *schema.ResourceData, meta interface{}) error {
	budgetName := d.Id()
	if !budgetExists(budgetName, meta) {
		log.Printf("[INFO] budget %s could not be found. skipping delete.", d.Id())
		return nil
	}

	client := meta.(*AWSClient).budgetconn
	accountID := meta.(*AWSClient).accountid
	deleteBudgetInput := new(budgets.DeleteBudgetInput)
	deleteBudgetInput.SetBudgetName(budgetName)
	deleteBudgetInput.SetAccountId(accountID)
	_, err := client.DeleteBudget(deleteBudgetInput)
	if err != nil {
		return fmt.Errorf("delete budget failed: %v", err)
	}

	return nil
}

type expandBudgetsBudgetFlattenedBudget struct {
	name                     *string
	budgetType               *string
	timeUnit                 *string
	costFilters              map[string][]*string
	limitAmount              *string
	limitUnit                *string
	includeCredit            *bool
	includeOtherSubscription *bool
	includeRecurring         *bool
	includeRefund            *bool
	includeSubscription      *bool
	includeSupport           *bool
	includeTax               *bool
	includeUpFront           *bool
	useBlended               *bool
	timePeriodStart          *time.Time
	timePeriodEnd            *time.Time
}

func expandBudgetsBudgetFlatten(budget *budgets.Budget) (*expandBudgetsBudgetFlattenedBudget, error) {
	if budget == nil {
		return nil, fmt.Errorf("empty budget returned from budget output: %v", budget)
	}

	budgetLimit := budget.BudgetLimit
	if budgetLimit == nil {
		return nil, fmt.Errorf("empty limit in budget: %v", budget)
	}

	budgetCostTypes := budget.CostTypes
	if budgetCostTypes == nil {
		return nil, fmt.Errorf("empty CostTypes in budget: %v", budget)
	}

	budgetTimePeriod := budget.TimePeriod
	if budgetTimePeriod == nil {
		return nil, fmt.Errorf("empty TimePeriod in budget: %v", budget)
	}

	budgetTimePeriodStart := budgetTimePeriod.Start
	if budgetTimePeriodStart == nil {
		return nil, fmt.Errorf("empty TimePeriodStart in budget: %v", budget)
	}

	budgetTimePeriodEnd := budgetTimePeriod.End
	if budgetTimePeriodEnd == nil {
		return nil, fmt.Errorf("empty TimePeriodEnd in budget: %v", budget)
	}

	return &expandBudgetsBudgetFlattenedBudget{
		name:                     budget.BudgetName,
		budgetType:               budget.BudgetType,
		timeUnit:                 budget.TimeUnit,
		costFilters:              budget.CostFilters,
		limitAmount:              budgetLimit.Amount,
		limitUnit:                budgetLimit.Unit,
		includeCredit:            budgetCostTypes.IncludeCredit,
		includeOtherSubscription: budgetCostTypes.IncludeOtherSubscription,
		includeRecurring:         budgetCostTypes.IncludeRecurring,
		includeRefund:            budgetCostTypes.IncludeRefund,
		includeSubscription:      budgetCostTypes.IncludeSubscription,
		includeSupport:           budgetCostTypes.IncludeSupport,
		includeTax:               budgetCostTypes.IncludeTax,
		includeUpFront:           budgetCostTypes.IncludeUpfront,
		useBlended:               budgetCostTypes.UseBlended,
		timePeriodStart:          budgetTimePeriodStart,
		timePeriodEnd:            budgetTimePeriodEnd,
	}, nil
}

func convertCostFiltersToStringMap(costFilters map[string][]*string) map[string]string {
	convertedCostFilters := make(map[string]string)
	for k, v := range costFilters {
		filterValues := make([]string, 0)
		for _, singleFilterValue := range v {
			filterValues = append(filterValues, *singleFilterValue)
		}

		convertedCostFilters[k] = strings.Join(filterValues, ",")
	}

	return convertedCostFilters
}

func newBudgetsBudget(d *schema.ResourceData) (*budgets.Budget, error) {
	var budgetName string
	if id := d.Id(); id != "" {
		budgetName = id

	} else if v, ok := d.GetOk("name"); ok {
		budgetName = v.(string)

	} else if v, ok := d.GetOk("name_prefix"); ok {
		budgetName = resource.PrefixedUniqueId(v.(string))

	} else {
		budgetName = resource.UniqueId()
	}

	budgetType := d.Get("budget_type").(string)
	budgetLimitAmount := d.Get("limit_amount").(string)
	budgetLimitUnit := d.Get("limit_unit").(string)
	budgetIncludeCredit := d.Get("include_credit").(bool)
	budgetIncludeOtherSubscription := d.Get("include_other_subscription").(bool)
	budgetIncludeRecurring := d.Get("include_recurring").(bool)
	budgetIncludeRefund := d.Get("include_refund").(bool)
	budgetIncludeSubscription := d.Get("include_subscription").(bool)
	budgetIncludeSupport := d.Get("include_support").(bool)
	budgetIncludeTax := d.Get("include_tax").(bool)
	budgetIncludeUpfront := d.Get("include_upfront").(bool)
	budgetUseBlended := d.Get("use_blended").(bool)
	budgetTimeUnit := d.Get("time_unit").(string)
	budgetCostFilters := make(map[string][]*string)
	for k, v := range d.Get("cost_filters").(map[string]interface{}) {
		filterValue := v.(string)
		budgetCostFilters[k] = append(budgetCostFilters[k], &filterValue)
	}

	budgetTimePeriodStart, err := time.Parse("2006-01-02_15:04", d.Get("time_period_start").(string))
	if err != nil {
		return nil, fmt.Errorf("failure parsing time: %v", err)
	}

	budgetTimePeriodEnd, err := time.Parse("2006-01-02_15:04", d.Get("time_period_end").(string))
	if err != nil {
		return nil, fmt.Errorf("failure parsing time: %v", err)
	}

	budget := new(budgets.Budget)
	budget.SetBudgetName(budgetName)
	budget.SetBudgetType(budgetType)
	budget.SetBudgetLimit(&budgets.Spend{
		Amount: &budgetLimitAmount,
		Unit:   &budgetLimitUnit,
	})
	budget.SetCostTypes(&budgets.CostTypes{
		IncludeCredit:            &budgetIncludeCredit,
		IncludeOtherSubscription: &budgetIncludeOtherSubscription,
		IncludeRecurring:         &budgetIncludeRecurring,
		IncludeRefund:            &budgetIncludeRefund,
		IncludeSubscription:      &budgetIncludeSubscription,
		IncludeSupport:           &budgetIncludeSupport,
		IncludeTax:               &budgetIncludeTax,
		IncludeUpfront:           &budgetIncludeUpfront,
		UseBlended:               &budgetUseBlended,
	})
	budget.SetTimePeriod(&budgets.TimePeriod{
		End:   &budgetTimePeriodEnd,
		Start: &budgetTimePeriodStart,
	})
	budget.SetTimeUnit(budgetTimeUnit)
	budget.SetCostFilters(budgetCostFilters)
	return budget, nil
}

func budgetExists(budgetName string, meta interface{}) bool {
	_, err := describeBudget(budgetName, meta)
	if isBudgetNotFoundException(err) {
		return false
	}

	return true
}

func isBudgetNotFoundException(err error) bool {
	if awsErr, ok := err.(awserr.Error); ok && awsErr.Code() == budgets.ErrCodeNotFoundException {
		return true
	}

	return false
}

func describeBudget(budgetName string, meta interface{}) (*budgets.DescribeBudgetOutput, error) {
	client := meta.(*AWSClient).budgetconn
	accountID := meta.(*AWSClient).accountid
	describeBudgetInput := new(budgets.DescribeBudgetInput)
	describeBudgetInput.SetBudgetName(budgetName)
	describeBudgetInput.SetAccountId(accountID)
	return client.DescribeBudget(describeBudgetInput)
}
<commit_msg>swap use of Set for struct init with value<commit_after>package aws

import (
	"fmt"
	"log"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/service/budgets"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/helper/schema"
)

func resourceAwsBudgetsBudget() *schema.Resource {
	return &schema.Resource{
		Schema: resourceAwsBudgetsBudgetSchema(),
		Create: resourceAwsBudgetsBudgetCreate,
		Read:   resourceAwsBudgetsBudgetRead,
		Update: resourceAwsBudgetsBudgetUpdate,
		Delete: resourceAwsBudgetsBudgetDelete,
	}
}

func resourceAwsBudgetsBudgetSchema() map[string]*schema.Schema {
	return map[string]*schema.Schema{
		"name": {
			Type:          schema.TypeString,
			Optional:      true,
			ConflictsWith: []string{"name_prefix"},
		},
		"name_prefix": {
			Type:     schema.TypeString,
			Optional: true,
			ForceNew: true,
		},
		"budget_type": {
			Type:     schema.TypeString,
			Required: true,
		},
		"limit_amount": {
			Type:     schema.TypeString,
			Required: true,
		},
		"limit_unit": {
			Type:     schema.TypeString,
			Required: true,
		},
		"include_credit": {
			Type:     schema.TypeBool,
			Optional: true,
			Default:  true,
		},
		"include_other_subscription": {
			Type:     schema.TypeBool,
			Optional: true,
			Default:  true,
		},
		"include_recurring": {
			Type:     schema.TypeBool,
			Optional: true,
			Default:  true,
		},
		"include_refund": {
			Type:     schema.TypeBool,
			Optional: true,
			Default:  true,
		},
		"include_subscription": {
			Type:     schema.TypeBool,
			Optional: true,
			Default:  true,
		},
		"include_support": {
			Type:     schema.TypeBool,
			Optional: true,
			Default:  true,
		},
		"include_tax": {
			Type:     schema.TypeBool,
			Optional: true,
			Default:  true,
		},
		"include_upfront": {
			Type:     schema.TypeBool,
			Optional: true,
			Default:  true,
		},
		"use_blended": {
			Type:     schema.TypeBool,
			Optional: true,
			Default:  false,
		},
		"time_period_start": {
			Type:     schema.TypeString,
			Required: true,
		},
		"time_period_end": {
			Type:     schema.TypeString,
			Optional: true,
			Default:  "2087-06-15_00:00",
		},
		"time_unit": {
			Type:     schema.TypeString,
			Required: true,
		},
		"cost_filters": {
			Type:     schema.TypeMap,
			Optional: true,
			Computed: true,
		},
	}
}

func resourceAwsBudgetsBudgetCreate(d *schema.ResourceData, meta interface{}) error {
	client := meta.(*AWSClient).budgetconn
	accountID := meta.(*AWSClient).accountid
	budget, err := newBudgetsBudget(d)
	if err != nil {
		return fmt.Errorf("failed creating budget: %v", err)
	}

	_, err = client.CreateBudget(&budgets.CreateBudgetInput{
		AccountId: &accountID,
		Budget:    budget,
	})
	if err != nil {
		return fmt.Errorf("create budget failed: %v", err)
	}

	d.SetId(*budget.BudgetName)
	return resourceAwsBudgetsBudgetUpdate(d, meta)
}

func resourceAwsBudgetsBudgetRead(d *schema.ResourceData, meta interface{}) error {
	budgetName := d.Id()
	describeBudgetOutput, err := describeBudget(budgetName, meta)
	if isBudgetNotFoundException(err) {
		d.SetId("")
		return nil
	}

	if err != nil {
		return fmt.Errorf("describe budget failed: %v", err)
	}

	flattenedBudget, err := expandBudgetsBudgetFlatten(describeBudgetOutput.Budget)
	if err != nil {
		return fmt.Errorf("failed flattening budget output: %v", err)
	}

	if _, ok := d.GetOk("name"); ok {
		d.Set("name", flattenedBudget.name)
	}

	for k, v := range map[string]interface{}{
		"budget_type":                flattenedBudget.budgetType,
		"time_unit":                  flattenedBudget.timeUnit,
		"cost_filters":               convertCostFiltersToStringMap(flattenedBudget.costFilters),
		"limit_amount":               flattenedBudget.limitAmount,
		"limit_unit":                 flattenedBudget.limitUnit,
		"include_credit":             flattenedBudget.includeCredit,
		"include_other_subscription": flattenedBudget.includeOtherSubscription,
		"include_recurring":          flattenedBudget.includeRecurring,
		"include_refund":             flattenedBudget.includeRefund,
		"include_subscription":       flattenedBudget.includeSubscription,
		"include_support":            flattenedBudget.includeSupport,
		"include_tax":                flattenedBudget.includeTax,
		"include_upfront":            flattenedBudget.includeUpFront,
		"use_blended":                flattenedBudget.useBlended,
		"time_period_start":          flattenedBudget.timePeriodStart.Format("2006-01-02_15:04"),
		"time_period_end":            flattenedBudget.timePeriodEnd.Format("2006-01-02_15:04"),
	} {
		if err := d.Set(k, v); err != nil {
			return err
		}
	}
	return nil
}

func resourceAwsBudgetsBudgetUpdate(d *schema.ResourceData, meta interface{}) error {
	client := meta.(*AWSClient).budgetconn
	accountID := meta.(*AWSClient).accountid
	budget, err := newBudgetsBudget(d)
	if err != nil {
		return fmt.Errorf("could not create budget: %v", err)
	}

	_, err = client.UpdateBudget(&budgets.UpdateBudgetInput{
		AccountId: &accountID,
		NewBudget: budget,
	})
	if err != nil {
		return fmt.Errorf("update budget failed: %v", err)
	}

	return resourceAwsBudgetsBudgetRead(d, meta)
}

func resourceAwsBudgetsBudgetDelete(d *schema.ResourceData, meta interface{}) error {
	budgetName := d.Id()
	if !budgetExists(budgetName, meta) {
		log.Printf("[INFO] budget %s could not be found. skipping delete.", d.Id())
		return nil
	}

	client := meta.(*AWSClient).budgetconn
	accountID := meta.(*AWSClient).accountid
	_, err := client.DeleteBudget(&budgets.DeleteBudgetInput{
		BudgetName: &budgetName,
		AccountId:  &accountID,
	})
	if err != nil {
		return fmt.Errorf("delete budget failed: %v", err)
	}

	return nil
}

type expandBudgetsBudgetFlattenedBudget struct {
	name                     *string
	budgetType               *string
	timeUnit                 *string
	costFilters              map[string][]*string
	limitAmount              *string
	limitUnit                *string
	includeCredit            *bool
	includeOtherSubscription *bool
	includeRecurring         *bool
	includeRefund            *bool
	includeSubscription      *bool
	includeSupport           *bool
	includeTax               *bool
	includeUpFront           *bool
	useBlended               *bool
	timePeriodStart          *time.Time
	timePeriodEnd            *time.Time
}

func expandBudgetsBudgetFlatten(budget *budgets.Budget) (*expandBudgetsBudgetFlattenedBudget, error) {
	if budget == nil {
		return nil, fmt.Errorf("empty budget returned from budget output: %v", budget)
	}

	budgetLimit := budget.BudgetLimit
	if budgetLimit == nil {
		return nil, fmt.Errorf("empty limit in budget: %v", budget)
	}

	budgetCostTypes := budget.CostTypes
	if budgetCostTypes == nil {
		return nil, fmt.Errorf("empty CostTypes in budget: %v", budget)
	}

	budgetTimePeriod := budget.TimePeriod
	if budgetTimePeriod == nil {
		return nil, fmt.Errorf("empty TimePeriod in budget: %v", budget)
	}

	budgetTimePeriodStart := budgetTimePeriod.Start
	if budgetTimePeriodStart == nil {
		return nil, fmt.Errorf("empty TimePeriodStart in budget: %v", budget)
	}

	budgetTimePeriodEnd := budgetTimePeriod.End
	if budgetTimePeriodEnd == nil {
		return nil, fmt.Errorf("empty TimePeriodEnd in budget: %v", budget)
	}

	return &expandBudgetsBudgetFlattenedBudget{
		name:                     budget.BudgetName,
		budgetType:               budget.BudgetType,
		timeUnit:                 budget.TimeUnit,
		costFilters:              budget.CostFilters,
		limitAmount:              budgetLimit.Amount,
		limitUnit:                budgetLimit.Unit,
		includeCredit:            budgetCostTypes.IncludeCredit,
		includeOtherSubscription: budgetCostTypes.IncludeOtherSubscription,
		includeRecurring:         budgetCostTypes.IncludeRecurring,
		includeRefund:            budgetCostTypes.IncludeRefund,
		includeSubscription:      budgetCostTypes.IncludeSubscription,
		includeSupport:           budgetCostTypes.IncludeSupport,
		includeTax:               budgetCostTypes.IncludeTax,
		includeUpFront:           budgetCostTypes.IncludeUpfront,
		useBlended:               budgetCostTypes.UseBlended,
		timePeriodStart:          budgetTimePeriodStart,
		timePeriodEnd:            budgetTimePeriodEnd,
	}, nil
}

func convertCostFiltersToStringMap(costFilters map[string][]*string) map[string]string {
	convertedCostFilters := make(map[string]string)
	for k, v := range costFilters {
		filterValues := make([]string, 0)
		for _, singleFilterValue := range v {
			filterValues = append(filterValues, *singleFilterValue)
		}

		convertedCostFilters[k] = strings.Join(filterValues, ",")
	}

	return convertedCostFilters
}

func newBudgetsBudget(d *schema.ResourceData) (*budgets.Budget, error) {
	var budgetName string
	if id := d.Id(); id != "" {
		budgetName = id

	} else if v, ok := d.GetOk("name"); ok {
		budgetName = v.(string)

	} else if v, ok := d.GetOk("name_prefix"); ok {
		budgetName = resource.PrefixedUniqueId(v.(string))

	} else {
		budgetName = resource.UniqueId()
	}

	budgetType := d.Get("budget_type").(string)
	budgetLimitAmount := d.Get("limit_amount").(string)
	budgetLimitUnit := d.Get("limit_unit").(string)
	budgetIncludeCredit := d.Get("include_credit").(bool)
	budgetIncludeOtherSubscription := d.Get("include_other_subscription").(bool)
	budgetIncludeRecurring := d.Get("include_recurring").(bool)
	budgetIncludeRefund := d.Get("include_refund").(bool)
	budgetIncludeSubscription := d.Get("include_subscription").(bool)
	budgetIncludeSupport := d.Get("include_support").(bool)
	budgetIncludeTax := d.Get("include_tax").(bool)
	budgetIncludeUpfront := d.Get("include_upfront").(bool)
	budgetUseBlended := d.Get("use_blended").(bool)
	budgetTimeUnit := d.Get("time_unit").(string)
	budgetCostFilters := make(map[string][]*string)
	for k, v := range d.Get("cost_filters").(map[string]interface{}) {
		filterValue := v.(string)
		budgetCostFilters[k] = append(budgetCostFilters[k], &filterValue)
	}

	budgetTimePeriodStart, err := time.Parse("2006-01-02_15:04", d.Get("time_period_start").(string))
	if err != nil {
		return nil, fmt.Errorf("failure parsing time: %v", err)
	}

	budgetTimePeriodEnd, err := time.Parse("2006-01-02_15:04", d.Get("time_period_end").(string))
	if err != nil {
		return nil, fmt.Errorf("failure parsing time: %v", err)
	}

	budget := &budgets.Budget{
		BudgetName: &budgetName,
		BudgetType: &budgetType,
		BudgetLimit: &budgets.Spend{
			Amount: &budgetLimitAmount,
			Unit:   &budgetLimitUnit,
		},
		CostTypes: costTypes,
		TimePeriod: &budgets.TimePeriod{
			End:   &budgetTimePeriodEnd,
			Start: &budgetTimePeriodStart,
		},
		TimeUnit:    &budgetTimeUnit,
		CostFilters: budgetCostFilters,
	}
	return budget, nil
}

func budgetExists(budgetName string, meta interface{}) bool {
	_, err := describeBudget(budgetName, meta)
	if isBudgetNotFoundException(err) {
		return false
	}

	return true
}

func isBudgetNotFoundException(err error) bool {
	if awsErr, ok := err.(awserr.Error); ok && awsErr.Code() == budgets.ErrCodeNotFoundException {
		return true
	}

	return false
}

func describeBudget(budgetName string, meta interface{}) (*budgets.DescribeBudgetOutput, error) {
	client := meta.(*AWSClient).budgetconn
	accountID := meta.(*AWSClient).accountid
	return client.DescribeBudget(&budgets.DescribeBudgetInput{
		BudgetName: &budgetName,
		AccountId:  &accountID,
	})
}
<|endoftext|>


<commit_before>// TeamSpeak 3 Server Query library
//
// Reference: http://goo.gl/OpJXz
package ts3

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	stdlog "log"
	"net"
	"os"
	"strings"
	"time"
)

const (
	DefaultPort    = "10011"
	VerificationID = "TS3"
)

var (
	// Custom logger
	log = stdlog.New(os.Stdout, "ts3> ", stdlog.LstdFlags)
	// ts3.Dial max timeout
	DialTimeout = 1 * time.Second
)

// Holds incoming, outgoing and notification requests (1) and responses (2 & 3)
type RWChan struct{ In, Out, Not chan string }

type Conn struct {
	conn net.Conn
	rw   RWChan
}

// Dial connects to a local/remote TS3 server. A default port is appended to
// `addr` if user doesn't provide one.
func Dial(addr string) *Conn {
	var (
		err  error
		line string
	)

	// Append DefaultPort if user didn't specify one
	if !strings.Contains(addr, ":") {
		addr += ":" + DefaultPort
	}

	// Try to establish connection
	conn, err := net.DialTimeout("tcp", addr, DialTimeout)
	fatal(err, fmt.Sprintf("Connection error: %v\n", err))

	// Allocate connection object
	ts3conn := &Conn{
		conn: conn,
		rw: RWChan{
			In:  make(chan string),
			Out: make(chan string),
			Not: make(chan string),
		},
	}

	rbuf := bufio.NewReader(conn)

	// Buffer to read from TCP socket; Read first line
	line, err = rbuf.ReadString('\n')
	fatal(err, "Couldn't identify server.")
	fmt.Print(line)

	// Then check if it's a TS3 server
	if !strings.Contains(line, VerificationID) {
		log.Fatal("Not a TeamSpeak 3 server.")
	}

	// Show welcome message
	line, err = rbuf.ReadString('\n')
	fatal(err, "Couldn't recv welcome message.")
	fmt.Print(line)

	// Copy flow: writer (request) -> conn -> reader (response)
	go cp(ts3conn, conn)
	go cp(conn, ts3conn)

	return ts3conn
}

// Read reads data from buffer into p doubling any IAC chars found (0xff), more
// info on RFC 854 (Telnet).  It returns the number of bytes read into p.
func (conn *Conn) Read(p []byte) (int, error) {
	b := []byte(<-conn.rw.In)
	// Double IAC chars
	bytes.Replace(b, []byte{0xff}, []byte{0xff, 0xff}, -1)
	copy(p, b)
	return len(b), nil
}

// Write writes the contents of p into the buffer. It returns the number of
// bytes written.
func (conn *Conn) Write(p []byte) (int, error) {
	s := string(p)
	conn.rw.Out <- s
	return len(p), nil
}

// Close closes underlying TCP Conn to local/remote server.
func (c *Conn) Close() error {
	return c.conn.Close()
}

// Cmd sends a request to a server and waits for its response.
func (c *Conn) Cmd(cmd string) string {
	var s string

	c.rw.In <- cmd + "\n"

	// Some commands output two lines
	var end bool
	for !end {
		select {
		case line := <-c.rw.Out:
			s += line
			if strings.HasPrefix(s, "notify") {
				c.rw.Not <- s
				end = true
			}
			if strings.HasPrefix(s, "error id=") {
				end = true
			}
		case <-time.After(500 * time.Millisecond):
			end = true
		}
	}

	return trimNet(s)
}

// Chans returns a low-level interaction
func (c *Conn) Chans() *RWChan {
	return &c.rw
}

// cp copies from an io.Reader to an io.Writer
func cp(dst io.Writer, src io.Reader) {
	_, err := io.Copy(dst, src)
	fatal(err)
}

// fatal exits application if encounters an error
func fatal(err error, s ...string) {
	if err != nil {
		if len(s) == 0 {
			log.Fatal(err)
		} else {
			log.Fatal(s)
		}
	}
}
<commit_msg>Update ts3.go<commit_after>// TeamSpeak 3 Server Query library
//
// Reference: http://goo.gl/OpJXz
package ts3

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	stdlog "log"
	"net"
	"os"
	"strings"
	"time"
)

const (
	DefaultPort    = "10011"
	VerificationID = "TS3"
)

var (
	// Custom logger
	log = stdlog.New(os.Stdout, "ts3> ", stdlog.LstdFlags)
	// ts3.Dial max timeout
	DialTimeout = 1 * time.Second
)

// Holds incoming, outgoing and notification requests (1) and responses (2 & 3)
type RWChan struct{ In, Out, Not chan string }

type Conn struct {
	conn net.Conn
	rw   RWChan
}

// Dial connects to a local/remote TS3 server. A default port is appended to
// `addr` if user doesn't provide one.
func Dial(addr string) *Conn {
	var (
		err  error
		line string
	)

	// Append DefaultPort if user didn't specify one
	if !strings.Contains(addr, ":") {
		addr += ":" + DefaultPort
	}

	// Try to establish connection
	conn, err := net.DialTimeout("tcp", addr, DialTimeout)
	fatal(err, fmt.Sprintf("Connection error: %v\n", err))

	// Allocate connection object
	ts3conn := &Conn{
		conn: conn,
		rw: RWChan{
			In:  make(chan string),
			Out: make(chan string),
			Not: make(chan string),
		},
	}

	rbuf := bufio.NewReader(conn)

	// Buffer to read from TCP socket; Read first line
	line, err = rbuf.ReadString('\n')
	fatal(err, "Couldn't identify server.")

	// Then check if it's a TS3 server
	if !strings.Contains(line, VerificationID) {
		log.Fatal("Not a TeamSpeak 3 server.")
	}

	// Show welcome message
	line, err = rbuf.ReadString('\n')
	fatal(err, "Couldn't recv welcome message.")

	// Copy flow: writer (request) -> conn -> reader (response)
	go cp(ts3conn, conn)
	go cp(conn, ts3conn)

	return ts3conn
}

// Read reads data from buffer into p doubling any IAC chars found (0xff), more
// info on RFC 854 (Telnet).  It returns the number of bytes read into p.
func (conn *Conn) Read(p []byte) (int, error) {
	b := []byte(<-conn.rw.In)
	// Double IAC chars
	bytes.Replace(b, []byte{0xff}, []byte{0xff, 0xff}, -1)
	copy(p, b)
	return len(b), nil
}

// Write writes the contents of p into the buffer. It returns the number of
// bytes written.
func (conn *Conn) Write(p []byte) (int, error) {
	s := string(p)
	conn.rw.Out <- s
	return len(p), nil
}

// Close closes underlying TCP Conn to local/remote server.
func (c *Conn) Close() error {
	return c.conn.Close()
}

// Cmd sends a request to a server and waits for its response.
func (c *Conn) Cmd(cmd string) string {
	var s string

	c.rw.In <- cmd + "\n"

	// Some commands output two lines
	var end bool
	for !end {
		select {
		case line := <-c.rw.Out:
			s += line
			if strings.HasPrefix(s, "notify") {
				c.rw.Not <- s
				end = true
			}
			if strings.HasPrefix(s, "error id=") {
				end = true
			}
		case <-time.After(500 * time.Millisecond):
			end = true
		}
	}

	return trimNet(s)
}

// Chans returns a low-level interaction
func (c *Conn) Chans() *RWChan {
	return &c.rw
}

// cp copies from an io.Reader to an io.Writer
func cp(dst io.Writer, src io.Reader) {
	_, err := io.Copy(dst, src)
	fatal(err)
}

// fatal exits application if encounters an error
func fatal(err error, s ...string) {
	if err != nil {
		if len(s) == 0 {
			log.Fatal(err)
		} else {
			log.Fatal(s)
		}
	}
}
<|endoftext|>


<commit_before>//go:build linux

// Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License"). You may
// not use this file except in compliance with the License. A copy of the
// License is located at
//
//	http://aws.amazon.com/apache2.0/
//
// or in the "license" file accompanying this file. This file 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 stats

import (
	"context"
	"fmt"
	"time"

	apitaskstatus "github.com/aws/amazon-ecs-agent/agent/api/task/status"
	"github.com/aws/amazon-ecs-agent/agent/config"
	"github.com/aws/amazon-ecs-agent/agent/ecscni"
	"github.com/aws/amazon-ecs-agent/agent/eni/netlinkwrapper"
	"github.com/aws/amazon-ecs-agent/agent/stats/resolver"
	"github.com/aws/amazon-ecs-agent/agent/utils/nswrapper"
	"github.com/aws/amazon-ecs-agent/agent/utils/retry"
	"github.com/containernetworking/plugins/pkg/ns"

	"github.com/cihub/seelog"
	"github.com/docker/docker/api/types"
	dockerstats "github.com/docker/docker/api/types"
	netlinklib "github.com/vishvananda/netlink"
)

const (
	// linkTypeDevice defines the string that's expected to be the output of
	// netlink.Link.Type() method for netlink.Device type.
	linkTypeDevice = "device"
	linkTypeVlan   = "vlan"
	// encapTypeLoopback defines the string that's set for the link.Attrs.EncapType
	// field for localhost devices. The EncapType field defines the link
	// encapsulation method. For localhost, it's set to "loopback".
	encapTypeLoopback = "loopback"
)

// StatsTask abstracts methods to gather and aggregate network data for a task. Used only for AWSVPC mode.
type StatsTask struct {
	StatsQueue            *Queue
	TaskMetadata          *TaskMetadata
	Ctx                   context.Context
	Cancel                context.CancelFunc
	Resolver              resolver.ContainerMetadataResolver
	nswrapperinterface    nswrapper.NS
	netlinkinterface      netlinkwrapper.NetLink
	metricPublishInterval time.Duration
}

func newStatsTaskContainer(taskARN string, containerPID string, numberOfContainers int,
	resolver resolver.ContainerMetadataResolver, publishInterval time.Duration) (*StatsTask, error) {
	nsAgent := nswrapper.NewNS()
	netlinkclient := netlinkwrapper.New()

	ctx, cancel := context.WithCancel(context.Background())
	return &StatsTask{
		TaskMetadata: &TaskMetadata{
			TaskArn:          taskARN,
			ContainerPID:     containerPID,
			NumberContainers: numberOfContainers,
		},
		Ctx:                   ctx,
		Cancel:                cancel,
		Resolver:              resolver,
		netlinkinterface:      netlinkclient,
		nswrapperinterface:    nsAgent,
		metricPublishInterval: publishInterval,
	}, nil
}

func (task *StatsTask) StartStatsCollection() {
	queueSize := int(config.DefaultContainerMetricsPublishInterval.Seconds() * 4)
	task.StatsQueue = NewQueue(queueSize)
	task.StatsQueue.Reset()
	go task.collect()
}

func (task *StatsTask) StopStatsCollection() {
	task.Cancel()
}

func (taskStat *StatsTask) collect() {
	taskArn := taskStat.TaskMetadata.TaskArn
	backoff := retry.NewExponentialBackoff(time.Second*1, time.Second*10, 0.5, 2)

	for {
		err := taskStat.processStatsStream()
		select {
		case <-taskStat.Ctx.Done():
			seelog.Debugf("Stopping stats collection for taskStat %s", taskArn)
			return
		default:
			if err != nil {
				d := backoff.Duration()
				time.Sleep(d)
				seelog.Debugf("Error querying stats for task %s: %v", taskArn, err)
			}
			// We were disconnected from the stats stream.
			// Check if the task is terminal. If it is, stop collecting metrics.
			terminal, err := taskStat.terminal()
			if err != nil {
				// Error determining if the task is terminal. clean-up anyway.
				seelog.Warnf("Error determining if the task %s is terminal, stopping stats collection: %v",
					taskArn, err)
				taskStat.StopStatsCollection()
			} else if terminal {
				seelog.Infof("Task %s is terminal, stopping stats collection", taskArn)
				taskStat.StopStatsCollection()
			}
		}
	}
}

func (taskStat *StatsTask) processStatsStream() error {
	taskArn := taskStat.TaskMetadata.TaskArn
	awsvpcNetworkStats, errC := taskStat.getAWSVPCNetworkStats()

	returnError := false
	for {
		select {
		case <-taskStat.Ctx.Done():
			seelog.Info("task context is done")
			return nil
		case err := <-errC:
			seelog.Warnf("Error encountered processing metrics stream from host, this may affect "+
				"cloudwatch metric accuracy: %s", err)
			returnError = true
		case rawStat, ok := <-awsvpcNetworkStats:
			if !ok {
				if returnError {
					return fmt.Errorf("error encountered processing metrics stream from host")
				}
				return nil
			}
			if err := taskStat.StatsQueue.Add(rawStat); err != nil {
				seelog.Warnf("Task [%s]: error converting stats: %v", taskArn, err)
			}
		}

	}
}

func (taskStat *StatsTask) terminal() (bool, error) {
	resolvedTask, err := taskStat.Resolver.ResolveTaskByARN(taskStat.TaskMetadata.TaskArn)
	if err != nil {
		return false, err
	}
	return resolvedTask.GetKnownStatus() == apitaskstatus.TaskStopped, nil
}

func getDevicesList(linkList []netlinklib.Link) []string {
	var deviceNames []string
	for _, link := range linkList {
		if link.Type() != linkTypeDevice && link.Type() != linkTypeVlan {
			// We only care about netlink.Device/netlink.Vlan types. Ignore other link types.
			continue
		}
		if link.Attrs().EncapType == encapTypeLoopback {
			// Ignore localhost
			continue
		}
		deviceNames = append(deviceNames, link.Attrs().Name)
	}
	return deviceNames
}

func (taskStat *StatsTask) populateNIDeviceList(containerPID string) ([]string, error) {
	var err error
	var deviceList []string
	netNSPath := fmt.Sprintf(ecscni.NetnsFormat, containerPID)
	err = taskStat.nswrapperinterface.WithNetNSPath(netNSPath, func(ns.NetNS) error {
		linksInTaskNetNS, linkErr := taskStat.netlinkinterface.LinkList()
		deviceNames := getDevicesList(linksInTaskNetNS)
		deviceList = append(deviceList, deviceNames...)
		return linkErr
	})
	return deviceList, err
}

func linkStatsToDockerStats(netLinkStats *netlinklib.LinkStatistics, numberOfContainers uint64) dockerstats.NetworkStats {
	networkStats := dockerstats.NetworkStats{
		RxBytes:   netLinkStats.RxBytes / numberOfContainers,
		RxPackets: netLinkStats.RxPackets / numberOfContainers,
		RxErrors:  netLinkStats.RxErrors / numberOfContainers,
		RxDropped: netLinkStats.RxDropped / numberOfContainers,
		TxBytes:   netLinkStats.TxBytes / numberOfContainers,
		TxPackets: netLinkStats.TxPackets / numberOfContainers,
		TxErrors:  netLinkStats.TxErrors / numberOfContainers,
		TxDropped: netLinkStats.TxDropped / numberOfContainers,
	}
	return networkStats
}

func (taskStat *StatsTask) getAWSVPCNetworkStats() (<-chan *types.StatsJSON, <-chan error) {

	errC := make(chan error)
	statsC := make(chan *dockerstats.StatsJSON)
	if taskStat.TaskMetadata.NumberContainers > 0 {
		go func() {
			defer close(statsC)
			statPollTicker := time.NewTicker(taskStat.metricPublishInterval)
			defer statPollTicker.Stop()
			for range statPollTicker.C {
				if len(taskStat.TaskMetadata.DeviceName) == 0 {
					var err error
					taskStat.TaskMetadata.DeviceName, err = taskStat.populateNIDeviceList(taskStat.TaskMetadata.ContainerPID)
					if err != nil {
						errC <- err
						return
					}
				}
				networkStats := make(map[string]dockerstats.NetworkStats, len(taskStat.TaskMetadata.DeviceName))
				for _, device := range taskStat.TaskMetadata.DeviceName {
					var link netlinklib.Link
					err := taskStat.nswrapperinterface.WithNetNSPath(fmt.Sprintf(ecscni.NetnsFormat,
						taskStat.TaskMetadata.ContainerPID),
						func(ns.NetNS) error {
							var linkErr error
							if link, linkErr = taskStat.netlinkinterface.LinkByName(device); linkErr != nil {
								return linkErr
							}
							return nil
						})

					if err != nil {
						errC <- err
						return
					}
					netLinkStats := link.Attrs().Statistics
					networkStats[link.Attrs().Name] = linkStatsToDockerStats(netLinkStats,
						uint64(taskStat.TaskMetadata.NumberContainers))
				}

				dockerStats := &types.StatsJSON{
					Networks: networkStats,
					Stats: types.Stats{
						Read: time.Now(),
					},
				}
				statsC <- dockerStats
			}
		}()
	}

	return statsC, errC
}
<commit_msg>Fix memory leak in task stats collector<commit_after>//go:build linux

// Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License"). You may
// not use this file except in compliance with the License. A copy of the
// License is located at
//
//	http://aws.amazon.com/apache2.0/
//
// or in the "license" file accompanying this file. This file 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 stats

import (
	"context"
	"fmt"
	"time"

	apitaskstatus "github.com/aws/amazon-ecs-agent/agent/api/task/status"
	"github.com/aws/amazon-ecs-agent/agent/config"
	"github.com/aws/amazon-ecs-agent/agent/ecscni"
	"github.com/aws/amazon-ecs-agent/agent/eni/netlinkwrapper"
	"github.com/aws/amazon-ecs-agent/agent/stats/resolver"
	"github.com/aws/amazon-ecs-agent/agent/utils/nswrapper"
	"github.com/aws/amazon-ecs-agent/agent/utils/retry"
	"github.com/containernetworking/plugins/pkg/ns"

	"github.com/cihub/seelog"
	"github.com/docker/docker/api/types"
	dockerstats "github.com/docker/docker/api/types"
	netlinklib "github.com/vishvananda/netlink"
)

const (
	// linkTypeDevice defines the string that's expected to be the output of
	// netlink.Link.Type() method for netlink.Device type.
	linkTypeDevice = "device"
	linkTypeVlan   = "vlan"
	// encapTypeLoopback defines the string that's set for the link.Attrs.EncapType
	// field for localhost devices. The EncapType field defines the link
	// encapsulation method. For localhost, it's set to "loopback".
	encapTypeLoopback = "loopback"
)

// StatsTask abstracts methods to gather and aggregate network data for a task. Used only for AWSVPC mode.
type StatsTask struct {
	StatsQueue            *Queue
	TaskMetadata          *TaskMetadata
	Ctx                   context.Context
	Cancel                context.CancelFunc
	Resolver              resolver.ContainerMetadataResolver
	nswrapperinterface    nswrapper.NS
	netlinkinterface      netlinkwrapper.NetLink
	metricPublishInterval time.Duration
}

func newStatsTaskContainer(taskARN string, containerPID string, numberOfContainers int,
	resolver resolver.ContainerMetadataResolver, publishInterval time.Duration) (*StatsTask, error) {
	nsAgent := nswrapper.NewNS()
	netlinkclient := netlinkwrapper.New()

	ctx, cancel := context.WithCancel(context.Background())
	return &StatsTask{
		TaskMetadata: &TaskMetadata{
			TaskArn:          taskARN,
			ContainerPID:     containerPID,
			NumberContainers: numberOfContainers,
		},
		Ctx:                   ctx,
		Cancel:                cancel,
		Resolver:              resolver,
		netlinkinterface:      netlinkclient,
		nswrapperinterface:    nsAgent,
		metricPublishInterval: publishInterval,
	}, nil
}

func (taskStat *StatsTask) StartStatsCollection() {
	queueSize := int(config.DefaultContainerMetricsPublishInterval.Seconds() * 4)
	taskStat.StatsQueue = NewQueue(queueSize)
	taskStat.StatsQueue.Reset()
	go taskStat.collect()
}

func (taskStat *StatsTask) StopStatsCollection() {
	taskStat.Cancel()
}

func (taskStat *StatsTask) collect() {
	taskArn := taskStat.TaskMetadata.TaskArn
	backoff := retry.NewExponentialBackoff(time.Second*1, time.Second*10, 0.5, 2)

	for {
		err := taskStat.processStatsStream()
		select {
		case <-taskStat.Ctx.Done():
			seelog.Debugf("Stopping stats collection for taskStat %s", taskArn)
			return
		default:
			if err != nil {
				d := backoff.Duration()
				time.Sleep(d)
				seelog.Debugf("Error querying stats for task %s: %v", taskArn, err)
			}
			// We were disconnected from the stats stream.
			// Check if the task is terminal. If it is, stop collecting metrics.
			terminal, err := taskStat.terminal()
			if err != nil {
				// Error determining if the task is terminal. clean-up anyway.
				seelog.Warnf("Error determining if the task %s is terminal, stopping stats collection: %v",
					taskArn, err)
				taskStat.StopStatsCollection()
			} else if terminal {
				seelog.Infof("Task %s is terminal, stopping stats collection", taskArn)
				taskStat.StopStatsCollection()
			}
		}
	}
}

func (taskStat *StatsTask) processStatsStream() error {
	taskArn := taskStat.TaskMetadata.TaskArn
	awsvpcNetworkStats, errC := taskStat.getAWSVPCNetworkStats()

	returnError := false
	for {
		select {
		case <-taskStat.Ctx.Done():
			seelog.Info("task context is done")
			return nil
		case err := <-errC:
			seelog.Warnf("Error encountered processing metrics stream from host, this may affect "+
				"cloudwatch metric accuracy: %s", err)
			returnError = true
		case rawStat, ok := <-awsvpcNetworkStats:
			if !ok {
				if returnError {
					return fmt.Errorf("error encountered processing metrics stream from host")
				}
				return nil
			}
			if err := taskStat.StatsQueue.Add(rawStat); err != nil {
				seelog.Warnf("Task [%s]: error converting stats: %v", taskArn, err)
			}
		}

	}
}

func (taskStat *StatsTask) terminal() (bool, error) {
	resolvedTask, err := taskStat.Resolver.ResolveTaskByARN(taskStat.TaskMetadata.TaskArn)
	if err != nil {
		return false, err
	}
	return resolvedTask.GetKnownStatus() == apitaskstatus.TaskStopped, nil
}

func getDevicesList(linkList []netlinklib.Link) []string {
	var deviceNames []string
	for _, link := range linkList {
		if link.Type() != linkTypeDevice && link.Type() != linkTypeVlan {
			// We only care about netlink.Device/netlink.Vlan types. Ignore other link types.
			continue
		}
		if link.Attrs().EncapType == encapTypeLoopback {
			// Ignore localhost
			continue
		}
		deviceNames = append(deviceNames, link.Attrs().Name)
	}
	return deviceNames
}

func (taskStat *StatsTask) populateNIDeviceList(containerPID string) ([]string, error) {
	var err error
	var deviceList []string
	netNSPath := fmt.Sprintf(ecscni.NetnsFormat, containerPID)
	err = taskStat.nswrapperinterface.WithNetNSPath(netNSPath, func(ns.NetNS) error {
		linksInTaskNetNS, linkErr := taskStat.netlinkinterface.LinkList()
		deviceNames := getDevicesList(linksInTaskNetNS)
		deviceList = append(deviceList, deviceNames...)
		return linkErr
	})
	return deviceList, err
}

func linkStatsToDockerStats(netLinkStats *netlinklib.LinkStatistics, numberOfContainers uint64) dockerstats.NetworkStats {
	networkStats := dockerstats.NetworkStats{
		RxBytes:   netLinkStats.RxBytes / numberOfContainers,
		RxPackets: netLinkStats.RxPackets / numberOfContainers,
		RxErrors:  netLinkStats.RxErrors / numberOfContainers,
		RxDropped: netLinkStats.RxDropped / numberOfContainers,
		TxBytes:   netLinkStats.TxBytes / numberOfContainers,
		TxPackets: netLinkStats.TxPackets / numberOfContainers,
		TxErrors:  netLinkStats.TxErrors / numberOfContainers,
		TxDropped: netLinkStats.TxDropped / numberOfContainers,
	}
	return networkStats
}

func (taskStat *StatsTask) getAWSVPCNetworkStats() (<-chan *types.StatsJSON, <-chan error) {

	errC := make(chan error, 1)
	statsC := make(chan *dockerstats.StatsJSON)
	if taskStat.TaskMetadata.NumberContainers > 0 {
		go func() {
			defer close(statsC)
			statPollTicker := time.NewTicker(taskStat.metricPublishInterval)
			defer statPollTicker.Stop()
			for range statPollTicker.C {
				if len(taskStat.TaskMetadata.DeviceName) == 0 {
					var err error
					taskStat.TaskMetadata.DeviceName, err = taskStat.populateNIDeviceList(taskStat.TaskMetadata.ContainerPID)
					if err != nil {
						errC <- err
						return
					}
				}
				networkStats := make(map[string]dockerstats.NetworkStats, len(taskStat.TaskMetadata.DeviceName))
				for _, device := range taskStat.TaskMetadata.DeviceName {
					var link netlinklib.Link
					err := taskStat.nswrapperinterface.WithNetNSPath(fmt.Sprintf(ecscni.NetnsFormat,
						taskStat.TaskMetadata.ContainerPID),
						func(ns.NetNS) error {
							var linkErr error
							if link, linkErr = taskStat.netlinkinterface.LinkByName(device); linkErr != nil {
								return linkErr
							}
							return nil
						})

					if err != nil {
						errC <- err
						return
					}
					netLinkStats := link.Attrs().Statistics
					networkStats[link.Attrs().Name] = linkStatsToDockerStats(netLinkStats,
						uint64(taskStat.TaskMetadata.NumberContainers))
				}

				dockerStats := &types.StatsJSON{
					Networks: networkStats,
					Stats: types.Stats{
						Read: time.Now(),
					},
				}
				select {
				case <-taskStat.Ctx.Done():
					return
				case statsC <- dockerStats:
				}
			}
		}()
	}

	return statsC, errC
}
<|endoftext|>


<commit_before>// Copyright 2017 Istio 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 kube

import (
	"fmt"
	"sync"

	v1 "k8s.io/api/core/v1"
	"k8s.io/client-go/tools/cache"

	"istio.io/istio/pilot/pkg/model"
	"istio.io/istio/pkg/log"
)

// PodCache is an eventually consistent pod cache
type PodCache struct {
	cacheHandler

	sync.RWMutex
	// keys maintains stable pod IP to name key mapping
	// this allows us to retrieve the latest status by pod IP.
	// This should only contain RUNNING or PENDING pods with an allocated IP.
	keys map[string]string

	c *Controller
}

func newPodCache(ch cacheHandler, c *Controller) *PodCache {
	out := &PodCache{
		cacheHandler: ch,
		c:            c,
		keys:         make(map[string]string),
	}

	ch.handler.Append(func(obj interface{}, ev model.Event) error {
		return out.event(obj, ev)
	})
	return out
}

// event updates the IP-based index (pc.keys).
func (pc *PodCache) event(obj interface{}, ev model.Event) error {
	pc.Lock()
	defer pc.Unlock()

	// When a pod is deleted obj could be an *v1.Pod or a DeletionFinalStateUnknown marker item.
	pod, ok := obj.(*v1.Pod)
	if !ok {
		tombstone, ok := obj.(cache.DeletedFinalStateUnknown)
		if !ok {
			return fmt.Errorf("couldn't get object from tombstone %+v", obj)
		}
		pod, ok = tombstone.Obj.(*v1.Pod)
		if !ok {
			return fmt.Errorf("tombstone contained object that is not a pod %#v", obj)
		}
	}

	ip := pod.Status.PodIP
	// PodIP will be empty when pod is just created, but before the IP is assigned
	// via UpdateStatus.

	if len(ip) > 0 {
		log.Infof("Handling event %s for pod %s in namespace %s -> %v", ev, pod.Name, pod.Namespace, ip)
		key := KeyFunc(pod.Name, pod.Namespace)
		switch ev {
		case model.EventAdd:
			switch pod.Status.Phase {
			case v1.PodPending, v1.PodRunning:
				// add to cache if the pod is running or pending
				pc.keys[ip] = key
				if pc.c.XDSUpdater != nil {
					pc.c.XDSUpdater.WorkloadUpdate(ip, pod.ObjectMeta.Labels, pod.ObjectMeta.Annotations)
				}
			}
		case model.EventUpdate:
			switch pod.Status.Phase {
			case v1.PodPending, v1.PodRunning:
				// add to cache if the pod is running or pending
				pc.keys[ip] = key
				if pc.c.XDSUpdater != nil {
					pc.c.XDSUpdater.WorkloadUpdate(ip, pod.ObjectMeta.Labels, pod.ObjectMeta.Annotations)
				}
			default:
				// delete if the pod switched to other states and is in the cache
				if pc.keys[ip] == key {
					delete(pc.keys, ip)
					if pc.c.XDSUpdater != nil {
						pc.c.XDSUpdater.WorkloadUpdate(ip, nil, nil)
					}
				}
			}
		case model.EventDelete:
			// delete only if this pod was in the cache
			if pc.keys[ip] == key {
				delete(pc.keys, ip)
				if pc.c != nil {
					pc.c.XDSUpdater.WorkloadUpdate(ip, nil, nil)
				}
			}
		}
	}
	return nil
}

// nolint: unparam
func (pc *PodCache) getPodKey(addr string) (string, bool) {
	pc.RLock()
	defer pc.RUnlock()
	key, exists := pc.keys[addr]
	return key, exists
}

// getPodByIp returns the pod or nil if pod not found or an error occurred
func (pc *PodCache) getPodByIP(addr string) *v1.Pod {
	pc.RLock()
	defer pc.RUnlock()

	key, exists := pc.keys[addr]
	if !exists {
		return nil
	}
	item, exists, err := pc.informer.GetStore().GetByKey(key)
	if !exists || err != nil {
		return nil
	}
	return item.(*v1.Pod)
}

// labelsByIP returns pod labels or nil if pod not found or an error occurred
func (pc *PodCache) labelsByIP(addr string) (model.Labels, bool) {
	pod := pc.getPodByIP(addr)
	if pod == nil {
		return nil, false
	}
	return convertLabels(pod.ObjectMeta), true
}
<commit_msg>bug: check pc.c.XDSUpdater instead of pc.c (#11102)<commit_after>// Copyright 2017 Istio 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 kube

import (
	"fmt"
	"sync"

	v1 "k8s.io/api/core/v1"
	"k8s.io/client-go/tools/cache"

	"istio.io/istio/pilot/pkg/model"
	"istio.io/istio/pkg/log"
)

// PodCache is an eventually consistent pod cache
type PodCache struct {
	cacheHandler

	sync.RWMutex
	// keys maintains stable pod IP to name key mapping
	// this allows us to retrieve the latest status by pod IP.
	// This should only contain RUNNING or PENDING pods with an allocated IP.
	keys map[string]string

	c *Controller
}

func newPodCache(ch cacheHandler, c *Controller) *PodCache {
	out := &PodCache{
		cacheHandler: ch,
		c:            c,
		keys:         make(map[string]string),
	}

	ch.handler.Append(func(obj interface{}, ev model.Event) error {
		return out.event(obj, ev)
	})
	return out
}

// event updates the IP-based index (pc.keys).
func (pc *PodCache) event(obj interface{}, ev model.Event) error {
	pc.Lock()
	defer pc.Unlock()

	// When a pod is deleted obj could be an *v1.Pod or a DeletionFinalStateUnknown marker item.
	pod, ok := obj.(*v1.Pod)
	if !ok {
		tombstone, ok := obj.(cache.DeletedFinalStateUnknown)
		if !ok {
			return fmt.Errorf("couldn't get object from tombstone %+v", obj)
		}
		pod, ok = tombstone.Obj.(*v1.Pod)
		if !ok {
			return fmt.Errorf("tombstone contained object that is not a pod %#v", obj)
		}
	}

	ip := pod.Status.PodIP
	// PodIP will be empty when pod is just created, but before the IP is assigned
	// via UpdateStatus.

	if len(ip) > 0 {
		log.Infof("Handling event %s for pod %s in namespace %s -> %v", ev, pod.Name, pod.Namespace, ip)
		key := KeyFunc(pod.Name, pod.Namespace)
		switch ev {
		case model.EventAdd:
			switch pod.Status.Phase {
			case v1.PodPending, v1.PodRunning:
				// add to cache if the pod is running or pending
				pc.keys[ip] = key
				if pc.c != nil && pc.c.XDSUpdater != nil {
					pc.c.XDSUpdater.WorkloadUpdate(ip, pod.ObjectMeta.Labels, pod.ObjectMeta.Annotations)
				}
			}
		case model.EventUpdate:
			switch pod.Status.Phase {
			case v1.PodPending, v1.PodRunning:
				// add to cache if the pod is running or pending
				pc.keys[ip] = key
				if pc.c != nil && pc.c.XDSUpdater != nil {
					pc.c.XDSUpdater.WorkloadUpdate(ip, pod.ObjectMeta.Labels, pod.ObjectMeta.Annotations)
				}
			default:
				// delete if the pod switched to other states and is in the cache
				if pc.keys[ip] == key {
					delete(pc.keys, ip)
					if pc.c != nil && pc.c.XDSUpdater != nil {
						pc.c.XDSUpdater.WorkloadUpdate(ip, nil, nil)
					}
				}
			}
		case model.EventDelete:
			// delete only if this pod was in the cache
			if pc.keys[ip] == key {
				delete(pc.keys, ip)
				if pc.c != nil && pc.c.XDSUpdater != nil {
					pc.c.XDSUpdater.WorkloadUpdate(ip, nil, nil)
				}
			}
		}
	}
	return nil
}

// nolint: unparam
func (pc *PodCache) getPodKey(addr string) (string, bool) {
	pc.RLock()
	defer pc.RUnlock()
	key, exists := pc.keys[addr]
	return key, exists
}

// getPodByIp returns the pod or nil if pod not found or an error occurred
func (pc *PodCache) getPodByIP(addr string) *v1.Pod {
	pc.RLock()
	defer pc.RUnlock()

	key, exists := pc.keys[addr]
	if !exists {
		return nil
	}
	item, exists, err := pc.informer.GetStore().GetByKey(key)
	if !exists || err != nil {
		return nil
	}
	return item.(*v1.Pod)
}

// labelsByIP returns pod labels or nil if pod not found or an error occurred
func (pc *PodCache) labelsByIP(addr string) (model.Labels, bool) {
	pod := pc.getPodByIP(addr)
	if pod == nil {
		return nil, false
	}
	return convertLabels(pod.ObjectMeta), true
}
<|endoftext|>


<commit_before>package s3

import (
	"context"
	"errors"
	"strings"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/dynamodb"
	"github.com/aws/aws-sdk-go/service/s3"
	awsbase "github.com/hashicorp/aws-sdk-go-base"
	"github.com/hashicorp/terraform/backend"
	"github.com/hashicorp/terraform/helper/schema"
)

// New creates a new backend for S3 remote state.
func New() backend.Backend {
	s := &schema.Backend{
		Schema: map[string]*schema.Schema{
			"bucket": {
				Type:        schema.TypeString,
				Required:    true,
				Description: "The name of the S3 bucket",
			},

			"key": {
				Type:        schema.TypeString,
				Required:    true,
				Description: "The path to the state file inside the bucket",
				ValidateFunc: func(v interface{}, s string) ([]string, []error) {
					// s3 will strip leading slashes from an object, so while this will
					// technically be accepted by s3, it will break our workspace hierarchy.
					if strings.HasPrefix(v.(string), "/") {
						return nil, []error{errors.New("key must not start with '/'")}
					}
					return nil, nil
				},
			},

			"region": {
				Type:        schema.TypeString,
				Required:    true,
				Description: "The region of the S3 bucket.",
				DefaultFunc: schema.MultiEnvDefaultFunc([]string{
					"AWS_REGION",
					"AWS_DEFAULT_REGION",
				}, nil),
			},

			"dynamodb_endpoint": {
				Type:        schema.TypeString,
				Optional:    true,
				Description: "A custom endpoint for the DynamoDB API",
				DefaultFunc: schema.EnvDefaultFunc("AWS_DYNAMODB_ENDPOINT", ""),
			},

			"endpoint": {
				Type:        schema.TypeString,
				Optional:    true,
				Description: "A custom endpoint for the S3 API",
				DefaultFunc: schema.EnvDefaultFunc("AWS_S3_ENDPOINT", ""),
			},

			"iam_endpoint": {
				Type:        schema.TypeString,
				Optional:    true,
				Description: "A custom endpoint for the IAM API",
				DefaultFunc: schema.EnvDefaultFunc("AWS_IAM_ENDPOINT", ""),
			},

			"sts_endpoint": {
				Type:        schema.TypeString,
				Optional:    true,
				Description: "A custom endpoint for the STS API",
				DefaultFunc: schema.EnvDefaultFunc("AWS_STS_ENDPOINT", ""),
			},

			"encrypt": {
				Type:        schema.TypeBool,
				Optional:    true,
				Description: "Whether to enable server side encryption of the state file",
				Default:     false,
			},

			"acl": {
				Type:        schema.TypeString,
				Optional:    true,
				Description: "Canned ACL to be applied to the state file",
				Default:     "",
			},

			"access_key": {
				Type:        schema.TypeString,
				Optional:    true,
				Description: "AWS access key",
				Default:     "",
			},

			"secret_key": {
				Type:        schema.TypeString,
				Optional:    true,
				Description: "AWS secret key",
				Default:     "",
			},

			"kms_key_id": {
				Type:        schema.TypeString,
				Optional:    true,
				Description: "The ARN of a KMS Key to use for encrypting the state",
				Default:     "",
			},

			"lock_table": {
				Type:        schema.TypeString,
				Optional:    true,
				Description: "DynamoDB table for state locking",
				Default:     "",
				Deprecated:  "please use the dynamodb_table attribute",
			},

			"dynamodb_table": {
				Type:        schema.TypeString,
				Optional:    true,
				Description: "DynamoDB table for state locking and consistency",
				Default:     "",
			},

			"profile": {
				Type:        schema.TypeString,
				Optional:    true,
				Description: "AWS profile name",
				Default:     "",
			},

			"shared_credentials_file": {
				Type:        schema.TypeString,
				Optional:    true,
				Description: "Path to a shared credentials file",
				Default:     "",
			},

			"token": {
				Type:        schema.TypeString,
				Optional:    true,
				Description: "MFA token",
				Default:     "",
			},

			"skip_credentials_validation": {
				Type:        schema.TypeBool,
				Optional:    true,
				Description: "Skip the credentials validation via STS API.",
				Default:     false,
			},

			"skip_get_ec2_platforms": {
				Type:        schema.TypeBool,
				Optional:    true,
				Description: "Skip getting the supported EC2 platforms.",
				Default:     false,
				Deprecated:  "The S3 Backend does not require EC2 functionality and this attribute is no longer used.",
			},

			"skip_region_validation": {
				Type:        schema.TypeBool,
				Optional:    true,
				Description: "Skip static validation of region name.",
				Default:     false,
			},

			"skip_requesting_account_id": {
				Type:        schema.TypeBool,
				Optional:    true,
				Description: "Skip requesting the account ID.",
				Default:     false,
				Deprecated:  "The S3 Backend no longer automatically looks up the AWS Account ID and this attribute is no longer used.",
			},

			"skip_metadata_api_check": {
				Type:        schema.TypeBool,
				Optional:    true,
				Description: "Skip the AWS Metadata API check.",
				Default:     false,
			},

			"role_arn": {
				Type:        schema.TypeString,
				Optional:    true,
				Description: "The role to be assumed",
				Default:     "",
			},

			"session_name": {
				Type:        schema.TypeString,
				Optional:    true,
				Description: "The session name to use when assuming the role.",
				Default:     "",
			},

			"external_id": {
				Type:        schema.TypeString,
				Optional:    true,
				Description: "The external ID to use when assuming the role",
				Default:     "",
			},

			"assume_role_policy": {
				Type:        schema.TypeString,
				Optional:    true,
				Description: "The permissions applied when assuming a role.",
				Default:     "",
			},

			"workspace_key_prefix": {
				Type:        schema.TypeString,
				Optional:    true,
				Description: "The prefix applied to the non-default state path inside the bucket.",
				Default:     "env:",
				ValidateFunc: func(v interface{}, s string) ([]string, []error) {
					prefix := v.(string)
					if strings.HasPrefix(prefix, "/") || strings.HasSuffix(prefix, "/") {
						return nil, []error{errors.New("workspace_key_prefix must not start or end with '/'")}
					}
					return nil, nil
				},
			},

			"force_path_style": {
				Type:        schema.TypeBool,
				Optional:    true,
				Description: "Force s3 to use path style api.",
				Default:     false,
			},

			"max_retries": {
				Type:        schema.TypeInt,
				Optional:    true,
				Description: "The maximum number of times an AWS API request is retried on retryable failure.",
				Default:     5,
			},
		},
	}

	result := &Backend{Backend: s}
	result.Backend.ConfigureFunc = result.configure
	return result
}

type Backend struct {
	*schema.Backend

	// The fields below are set from configure
	s3Client  *s3.S3
	dynClient *dynamodb.DynamoDB

	bucketName           string
	keyName              string
	serverSideEncryption bool
	acl                  string
	kmsKeyID             string
	ddbTable             string
	workspaceKeyPrefix   string
}

func (b *Backend) configure(ctx context.Context) error {
	if b.s3Client != nil {
		return nil
	}

	// Grab the resource data
	data := schema.FromContextBackendConfig(ctx)

	if !data.Get("skip_region_validation").(bool) {
		if err := awsbase.ValidateRegion(data.Get("region").(string)); err != nil {
			return err
		}
	}

	b.bucketName = data.Get("bucket").(string)
	b.keyName = data.Get("key").(string)
	b.serverSideEncryption = data.Get("encrypt").(bool)
	b.acl = data.Get("acl").(string)
	b.kmsKeyID = data.Get("kms_key_id").(string)
	b.workspaceKeyPrefix = data.Get("workspace_key_prefix").(string)

	b.ddbTable = data.Get("dynamodb_table").(string)
	if b.ddbTable == "" {
		// try the deprecated field
		b.ddbTable = data.Get("lock_table").(string)
	}

	cfg := &awsbase.Config{
		AccessKey:             data.Get("access_key").(string),
		AssumeRoleARN:         data.Get("role_arn").(string),
		AssumeRoleExternalID:  data.Get("external_id").(string),
		AssumeRolePolicy:      data.Get("assume_role_policy").(string),
		AssumeRoleSessionName: data.Get("session_name").(string),
		CredsFilename:         data.Get("shared_credentials_file").(string),
		IamEndpoint:           data.Get("iam_endpoint").(string),
		MaxRetries:            data.Get("max_retries").(int),
		Profile:               data.Get("profile").(string),
		Region:                data.Get("region").(string),
		SecretKey:             data.Get("secret_key").(string),
		SkipCredsValidation:   data.Get("skip_credentials_validation").(bool),
		SkipMetadataApiCheck:  data.Get("skip_metadata_api_check").(bool),
		StsEndpoint:           data.Get("sts_endpoint").(string),
		Token:                 data.Get("token").(string),
	}

	sess, err := awsbase.GetSession(cfg)
	if err != nil {
		return err
	}

	b.dynClient = dynamodb.New(sess.Copy(&aws.Config{
		Endpoint: aws.String(data.Get("dynamodb_endpoint").(string)),
	}))
	b.s3Client = s3.New(sess.Copy(&aws.Config{
		Endpoint:         aws.String(data.Get("endpoint").(string)),
		S3ForcePathStyle: aws.Bool(data.Get("force_path_style").(bool)),
	}))

	return nil
}
<commit_msg>backend/s3: Add debug logging and user agent<commit_after>package s3

import (
	"context"
	"errors"
	"strings"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/dynamodb"
	"github.com/aws/aws-sdk-go/service/s3"
	awsbase "github.com/hashicorp/aws-sdk-go-base"
	"github.com/hashicorp/terraform/backend"
	"github.com/hashicorp/terraform/helper/logging"
	"github.com/hashicorp/terraform/helper/schema"
	"github.com/hashicorp/terraform/version"
)

// New creates a new backend for S3 remote state.
func New() backend.Backend {
	s := &schema.Backend{
		Schema: map[string]*schema.Schema{
			"bucket": {
				Type:        schema.TypeString,
				Required:    true,
				Description: "The name of the S3 bucket",
			},

			"key": {
				Type:        schema.TypeString,
				Required:    true,
				Description: "The path to the state file inside the bucket",
				ValidateFunc: func(v interface{}, s string) ([]string, []error) {
					// s3 will strip leading slashes from an object, so while this will
					// technically be accepted by s3, it will break our workspace hierarchy.
					if strings.HasPrefix(v.(string), "/") {
						return nil, []error{errors.New("key must not start with '/'")}
					}
					return nil, nil
				},
			},

			"region": {
				Type:        schema.TypeString,
				Required:    true,
				Description: "The region of the S3 bucket.",
				DefaultFunc: schema.MultiEnvDefaultFunc([]string{
					"AWS_REGION",
					"AWS_DEFAULT_REGION",
				}, nil),
			},

			"dynamodb_endpoint": {
				Type:        schema.TypeString,
				Optional:    true,
				Description: "A custom endpoint for the DynamoDB API",
				DefaultFunc: schema.EnvDefaultFunc("AWS_DYNAMODB_ENDPOINT", ""),
			},

			"endpoint": {
				Type:        schema.TypeString,
				Optional:    true,
				Description: "A custom endpoint for the S3 API",
				DefaultFunc: schema.EnvDefaultFunc("AWS_S3_ENDPOINT", ""),
			},

			"iam_endpoint": {
				Type:        schema.TypeString,
				Optional:    true,
				Description: "A custom endpoint for the IAM API",
				DefaultFunc: schema.EnvDefaultFunc("AWS_IAM_ENDPOINT", ""),
			},

			"sts_endpoint": {
				Type:        schema.TypeString,
				Optional:    true,
				Description: "A custom endpoint for the STS API",
				DefaultFunc: schema.EnvDefaultFunc("AWS_STS_ENDPOINT", ""),
			},

			"encrypt": {
				Type:        schema.TypeBool,
				Optional:    true,
				Description: "Whether to enable server side encryption of the state file",
				Default:     false,
			},

			"acl": {
				Type:        schema.TypeString,
				Optional:    true,
				Description: "Canned ACL to be applied to the state file",
				Default:     "",
			},

			"access_key": {
				Type:        schema.TypeString,
				Optional:    true,
				Description: "AWS access key",
				Default:     "",
			},

			"secret_key": {
				Type:        schema.TypeString,
				Optional:    true,
				Description: "AWS secret key",
				Default:     "",
			},

			"kms_key_id": {
				Type:        schema.TypeString,
				Optional:    true,
				Description: "The ARN of a KMS Key to use for encrypting the state",
				Default:     "",
			},

			"lock_table": {
				Type:        schema.TypeString,
				Optional:    true,
				Description: "DynamoDB table for state locking",
				Default:     "",
				Deprecated:  "please use the dynamodb_table attribute",
			},

			"dynamodb_table": {
				Type:        schema.TypeString,
				Optional:    true,
				Description: "DynamoDB table for state locking and consistency",
				Default:     "",
			},

			"profile": {
				Type:        schema.TypeString,
				Optional:    true,
				Description: "AWS profile name",
				Default:     "",
			},

			"shared_credentials_file": {
				Type:        schema.TypeString,
				Optional:    true,
				Description: "Path to a shared credentials file",
				Default:     "",
			},

			"token": {
				Type:        schema.TypeString,
				Optional:    true,
				Description: "MFA token",
				Default:     "",
			},

			"skip_credentials_validation": {
				Type:        schema.TypeBool,
				Optional:    true,
				Description: "Skip the credentials validation via STS API.",
				Default:     false,
			},

			"skip_get_ec2_platforms": {
				Type:        schema.TypeBool,
				Optional:    true,
				Description: "Skip getting the supported EC2 platforms.",
				Default:     false,
				Deprecated:  "The S3 Backend does not require EC2 functionality and this attribute is no longer used.",
			},

			"skip_region_validation": {
				Type:        schema.TypeBool,
				Optional:    true,
				Description: "Skip static validation of region name.",
				Default:     false,
			},

			"skip_requesting_account_id": {
				Type:        schema.TypeBool,
				Optional:    true,
				Description: "Skip requesting the account ID.",
				Default:     false,
				Deprecated:  "The S3 Backend no longer automatically looks up the AWS Account ID and this attribute is no longer used.",
			},

			"skip_metadata_api_check": {
				Type:        schema.TypeBool,
				Optional:    true,
				Description: "Skip the AWS Metadata API check.",
				Default:     false,
			},

			"role_arn": {
				Type:        schema.TypeString,
				Optional:    true,
				Description: "The role to be assumed",
				Default:     "",
			},

			"session_name": {
				Type:        schema.TypeString,
				Optional:    true,
				Description: "The session name to use when assuming the role.",
				Default:     "",
			},

			"external_id": {
				Type:        schema.TypeString,
				Optional:    true,
				Description: "The external ID to use when assuming the role",
				Default:     "",
			},

			"assume_role_policy": {
				Type:        schema.TypeString,
				Optional:    true,
				Description: "The permissions applied when assuming a role.",
				Default:     "",
			},

			"workspace_key_prefix": {
				Type:        schema.TypeString,
				Optional:    true,
				Description: "The prefix applied to the non-default state path inside the bucket.",
				Default:     "env:",
				ValidateFunc: func(v interface{}, s string) ([]string, []error) {
					prefix := v.(string)
					if strings.HasPrefix(prefix, "/") || strings.HasSuffix(prefix, "/") {
						return nil, []error{errors.New("workspace_key_prefix must not start or end with '/'")}
					}
					return nil, nil
				},
			},

			"force_path_style": {
				Type:        schema.TypeBool,
				Optional:    true,
				Description: "Force s3 to use path style api.",
				Default:     false,
			},

			"max_retries": {
				Type:        schema.TypeInt,
				Optional:    true,
				Description: "The maximum number of times an AWS API request is retried on retryable failure.",
				Default:     5,
			},
		},
	}

	result := &Backend{Backend: s}
	result.Backend.ConfigureFunc = result.configure
	return result
}

type Backend struct {
	*schema.Backend

	// The fields below are set from configure
	s3Client  *s3.S3
	dynClient *dynamodb.DynamoDB

	bucketName           string
	keyName              string
	serverSideEncryption bool
	acl                  string
	kmsKeyID             string
	ddbTable             string
	workspaceKeyPrefix   string
}

func (b *Backend) configure(ctx context.Context) error {
	if b.s3Client != nil {
		return nil
	}

	// Grab the resource data
	data := schema.FromContextBackendConfig(ctx)

	if !data.Get("skip_region_validation").(bool) {
		if err := awsbase.ValidateRegion(data.Get("region").(string)); err != nil {
			return err
		}
	}

	b.bucketName = data.Get("bucket").(string)
	b.keyName = data.Get("key").(string)
	b.serverSideEncryption = data.Get("encrypt").(bool)
	b.acl = data.Get("acl").(string)
	b.kmsKeyID = data.Get("kms_key_id").(string)
	b.workspaceKeyPrefix = data.Get("workspace_key_prefix").(string)

	b.ddbTable = data.Get("dynamodb_table").(string)
	if b.ddbTable == "" {
		// try the deprecated field
		b.ddbTable = data.Get("lock_table").(string)
	}

	cfg := &awsbase.Config{
		AccessKey:             data.Get("access_key").(string),
		AssumeRoleARN:         data.Get("role_arn").(string),
		AssumeRoleExternalID:  data.Get("external_id").(string),
		AssumeRolePolicy:      data.Get("assume_role_policy").(string),
		AssumeRoleSessionName: data.Get("session_name").(string),
		CredsFilename:         data.Get("shared_credentials_file").(string),
		DebugLogging:          logging.IsDebugOrHigher(),
		IamEndpoint:           data.Get("iam_endpoint").(string),
		MaxRetries:            data.Get("max_retries").(int),
		Profile:               data.Get("profile").(string),
		Region:                data.Get("region").(string),
		SecretKey:             data.Get("secret_key").(string),
		SkipCredsValidation:   data.Get("skip_credentials_validation").(bool),
		SkipMetadataApiCheck:  data.Get("skip_metadata_api_check").(bool),
		StsEndpoint:           data.Get("sts_endpoint").(string),
		Token:                 data.Get("token").(string),
		UserAgentProducts: []*awsbase.UserAgentProduct{
			{Name: "APN", Version: "1.0"},
			{Name: "HashiCorp", Version: "1.0"},
			{Name: "Terraform", Version: version.String()},
		},
	}

	sess, err := awsbase.GetSession(cfg)
	if err != nil {
		return err
	}

	b.dynClient = dynamodb.New(sess.Copy(&aws.Config{
		Endpoint: aws.String(data.Get("dynamodb_endpoint").(string)),
	}))
	b.s3Client = s3.New(sess.Copy(&aws.Config{
		Endpoint:         aws.String(data.Get("endpoint").(string)),
		S3ForcePathStyle: aws.Bool(data.Get("force_path_style").(bool)),
	}))

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"database/sql"
	"flag"
	"fmt"
	"math/rand"
	"runtime"
	"sync"

	_ "github.com/c4pt0r/mysql"

	"github.com/coopernurse/gorp"
)

func NewDbMap() *gorp.DbMap {
	dsn := "root:@tcp(127.0.0.1:4000)/benchmark"
	dbType := "mysql"
	db, err := sql.Open(dbType, dsn)
	if err != nil {
		panic(err.Error())
	}

	dbmap := &gorp.DbMap{Db: db, Dialect: gorp.MySQLDialect{"InnoDB", "UTF8"}}
	dbmap.AddTableWithName(TestData{}, "autoincr_test").SetKeys(true, "Id")

	if err != nil {
		panic(err.Error())
	}
	return dbmap
}

var letters = []rune("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ")

func randSeq(n int) string {
	b := make([]rune, n)
	for i := range b {
		b[i] = letters[rand.Intn(len(letters))]
	}
	return string(b)
}

type TestData struct {
	Id       int    `db:"id"`
	DateTime string `db:"datetime"`
	Data     string `db:"data"`
}

var n = flag.Int("n", 20000, "n")
var testType = flag.String("t", "read", "read|write")
var numWorkers = 50

func WriteTest() {
	wg := sync.WaitGroup{}
	var chans []chan *TestData
	for i := 0; i < numWorkers; i++ {
		c := make(chan *TestData, 100)
		chans = append(chans, c)
		go func(chan *TestData) {
			m := NewDbMap()
			for t := range c {
				err := m.Insert(t)
				if err != nil {
					panic(err)
				}
				wg.Done()
			}
		}(c)
	}
	for i := 0; i < *n; i++ {
		wg.Add(1)
		chans[i%numWorkers] <- &TestData{
			Data: randSeq(1024),
		}
	}
	wg.Wait()
}

func ReadTest() {
	wg := sync.WaitGroup{}
	c := make(chan int)
	for i := 0; i < numWorkers; i++ {
		go func(c chan int) {
			m := NewDbMap()
			t := TestData{}
			for _ = range c {
				x := rand.Intn(2000) + 1
				err := m.SelectOne(&t, "select * from autoincr_test where id = ?", x)
				if err != nil {
					fmt.Println(x)
					panic(err)
				}
				wg.Done()
			}
		}(c)
	}

	for i := 1; i < *n; i++ {
		wg.Add(1)
		c <- i
	}

	wg.Wait()

}

func main() {
	runtime.GOMAXPROCS(runtime.NumCPU())
	flag.Parse()
	if *testType == "read" {
		ReadTest()
	} else {
		WriteTest()
	}
}
<commit_msg>fix bug<commit_after>package main

import (
	"database/sql"
	"flag"
	"fmt"
	"math/rand"
	"runtime"
	"sync"

	_ "github.com/c4pt0r/mysql"

	"github.com/coopernurse/gorp"
)

var letters = []rune("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ")

func randSeq(n int) string {
	b := make([]rune, n)
	for i := range b {
		b[i] = letters[rand.Intn(len(letters))]
	}
	return string(b)
}

func NewDbMap() *gorp.DbMap {
	dsn := "root:@tcp(127.0.0.1:4000)/benchmark"
	dbType := "mysql"
	db, err := sql.Open(dbType, dsn)

	db.SetMaxIdleConns(100)
	db.SetMaxOpenConns(100)

	if err != nil {
		panic(err.Error())
	}

	dbmap := &gorp.DbMap{Db: db, Dialect: gorp.MySQLDialect{"InnoDB", "UTF8"}}
	dbmap.AddTableWithName(TestData{}, "autoincr_test").SetKeys(true, "Id")

	if err != nil {
		panic(err.Error())
	}
	return dbmap
}

type TestData struct {
	Id       int    `db:"id"`
	DateTime string `db:"datetime"`
	Data     string `db:"data"`
}

var n = flag.Int("n", 20000, "n")
var testType = flag.String("t", "read", "read|write")
var numWorkers = 50

func WriteTest() {
	wg := sync.WaitGroup{}
	var chans []chan *TestData
	for i := 0; i < numWorkers; i++ {
		c := make(chan *TestData, 100)
		chans = append(chans, c)
		go func(chan *TestData) {
			m := NewDbMap()
			for t := range c {
				err := m.Insert(t)
				if err != nil {
					panic(err)
				}
				wg.Done()
			}
		}(c)
	}
	for i := 0; i < *n; i++ {
		wg.Add(1)
		chans[i%numWorkers] <- &TestData{
			Data: randSeq(1024),
		}
	}
	wg.Wait()
}

func ReadTest() {
	wg := sync.WaitGroup{}
	c := make(chan int)
	for i := 0; i < numWorkers; i++ {
		go func(c chan int) {
			m := NewDbMap()
			t := TestData{}
			for _ = range c {
				x := rand.Intn(2000) + 1
				err := m.SelectOne(&t, "select * from autoincr_test where id = ?", x)
				if err != nil {
					fmt.Println(x)
					panic(err)
				}
				wg.Done()
			}
		}(c)
	}

	for i := 1; i < *n; i++ {
		wg.Add(1)
		c <- i
	}

	wg.Wait()
}

func main() {
	runtime.GOMAXPROCS(runtime.NumCPU())
	flag.Parse()
	if *testType == "read" {
		ReadTest()
	} else {
		WriteTest()
	}
}
<|endoftext|>