<commit_before>package exoscale

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"path/filepath"
	"strings"
	"text/template"
	"time"

	"github.com/codegangsta/cli"
	"github.com/docker/machine/drivers"
	"github.com/docker/machine/log"
	"github.com/docker/machine/provider"
	"github.com/docker/machine/state"
	"github.com/pyr/egoscale/src/egoscale"
)

type Driver struct {
	URL              string
	ApiKey           string
	ApiSecretKey     string
	InstanceProfile  string
	DiskSize         int
	Image            string
	SecurityGroup    string
	AvailabilityZone string
	MachineName      string
	KeyPair          string
	IPAddress        string
	PublicKey        string
	Id               string
	CaCertPath       string
	PrivateKeyPath   string
	DriverKeyPath    string
	SwarmMaster      bool
	SwarmHost        string
	SwarmDiscovery   string
	storePath        string
}

func init() {
	drivers.Register("exoscale", &drivers.RegisteredDriver{
		New:            NewDriver,
		GetCreateFlags: GetCreateFlags,
	})
}

// RegisterCreateFlags registers the flags this driver adds to
// "docker hosts create"
func GetCreateFlags() []cli.Flag {
	return []cli.Flag{
		cli.StringFlag{
			EnvVar: "EXOSCALE_ENDPOINT",
			Name:   "exoscale-url",
			Usage:  "exoscale API endpoint",
		},
		cli.StringFlag{
			EnvVar: "EXOSCALE_API_KEY",
			Name:   "exoscale-api-key",
			Usage:  "exoscale API key",
		},
		cli.StringFlag{
			EnvVar: "EXOSCALE_API_SECRET",
			Name:   "exoscale-api-secret-key",
			Usage:  "exoscale API secret key",
		},
		cli.StringFlag{
			EnvVar: "EXOSCALE_INSTANCE_PROFILE",
			Name:   "exoscale-instance-profile",
			Value:  "small",
			Usage:  "exoscale instance profile (small, medium, large, ...)",
		},
		cli.IntFlag{
			EnvVar: "EXOSCALE_DISK_SIZE",
			Name:   "exoscale-disk-size",
			Value:  50,
			Usage:  "exoscale disk size (10, 50, 100, 200, 400)",
		},
		cli.StringFlag{
			EnvVar: "EXSOCALE_IMAGE",
			Name:   "exoscale-image",
			Value:  "ubuntu-14.04",
			Usage:  "exoscale image template",
		},
		cli.StringFlag{
			EnvVar: "EXOSCALE_SECURITY_GROUP",
			Name:   "exoscale-security-group",
			Value:  "docker-machine",
			Usage:  "exoscale security group",
		},
		cli.StringFlag{
			EnvVar: "EXOSCALE_AVAILABILITY_ZONE",
			Name:   "exoscale-availability-zone",
			Value:  "ch-gva-2",
			Usage:  "exoscale availibility zone",
		},
		cli.StringFlag{
			EnvVar: "EXOSCALE_KEYPAIR",
			Name:   "exoscale-keypair",
			Usage:  "exoscale keypair name",
		},
	}
}

func NewDriver(machineName string, storePath string, caCert string, privateKey string) (drivers.Driver, error) {
	return &Driver{MachineName: machineName, storePath: storePath, CaCertPath: caCert, PrivateKeyPath: privateKey}, nil
}

func (d *Driver) AuthorizePort(ports []*drivers.Port) error {
	return nil
}

func (d *Driver) DeauthorizePort(ports []*drivers.Port) error {
	return nil
}

func (d *Driver) GetMachineName() string {
	return d.MachineName
}

func (d *Driver) GetSSHHostname() (string, error) {
	return d.GetIP()
}

func (d *Driver) GetSSHKeyPath() string {
	return filepath.Join(d.storePath, "id_rsa")
}

func (d *Driver) GetSSHPort() (int, error) {
	return 22, nil
}

func (d *Driver) GetSSHUsername() string {
	return "ubuntu"
}

func (d *Driver) GetProviderType() provider.ProviderType {
	return provider.Remote
}

func (d *Driver) DriverName() string {
	return "exoscale"
}

func (d *Driver) SetConfigFromFlags(flags drivers.DriverOptions) error {
	d.URL = flags.String("exoscale-endpoint")
	d.ApiKey = flags.String("exoscale-api-key")
	d.ApiSecretKey = flags.String("exoscale-api-secret-key")
	d.InstanceProfile = flags.String("exoscale-instance-profile")
	d.DiskSize = flags.Int("exoscale-disk-size")
	d.Image = flags.String("exoscale-image")
	d.SecurityGroup = flags.String("exoscale-security-group")
	d.AvailabilityZone = flags.String("exoscale-availability-zone")
	d.KeyPair = flags.String("exoscale-keypair")
	d.SwarmMaster = flags.Bool("swarm-master")
	d.SwarmHost = flags.String("swarm-host")
	d.SwarmDiscovery = flags.String("swarm-discovery")

	if d.URL == "" {
		d.URL = "https://api.exoscale.ch/compute"
	}
	if d.ApiKey == "" || d.ApiSecretKey == "" {
		return fmt.Errorf("Please specify an API key (--exoscale-api-key) and an API secret key (--exoscale-api-secret-key).")
	}

	return nil
}

func (d *Driver) GetURL() (string, error) {
	ip, err := d.GetIP()
	if err != nil {
		return "", err
	}
	return fmt.Sprintf("tcp://%s:2376", ip), nil
}

func (d *Driver) GetIP() (string, error) {
	if d.IPAddress == "" {
		return "", fmt.Errorf("IP address is not set")
	}
	return d.IPAddress, nil
}

func (d *Driver) GetState() (state.State, error) {
	client := egoscale.NewClient(d.URL, d.ApiKey, d.ApiSecretKey)
	vm, err := client.GetVirtualMachine(d.Id)
	if err != nil {
		return state.Error, err
	}
	switch vm.State {
	case "Starting":
		return state.Starting, nil
	case "Running":
		return state.Running, nil
	case "Stopping":
		return state.Running, nil
	case "Stopped":
		return state.Stopped, nil
	case "Destroyed":
		return state.Stopped, nil
	case "Expunging":
		return state.Stopped, nil
	case "Migrating":
		return state.Paused, nil
	case "Error":
		return state.Error, nil
	case "Unknown":
		return state.Error, nil
	case "Shutdowned":
		return state.Stopped, nil
	}
	return state.None, nil
}

func (d *Driver) PreCreateCheck() error {
	return nil
}

func (d *Driver) Create() error {
	log.Infof("Querying exoscale for the requested parameters...")
	client := egoscale.NewClient(d.URL, d.ApiKey, d.ApiSecretKey)
	topology, err := client.GetTopology()
	if err != nil {
		return err
	}

	// Availability zone UUID
	zone, ok := topology.Zones[d.AvailabilityZone]
	if !ok {
		return fmt.Errorf("Availability zone %v doesn't exist",
			d.AvailabilityZone)
	}
	log.Debugf("Availability zone %v = %s", d.AvailabilityZone, zone)

	// Image UUID
	var tpl string
	images, ok := topology.Images[strings.ToLower(d.Image)]
	if ok {
		tpl, ok = images[d.DiskSize]
	}
	if !ok {
		return fmt.Errorf("Unable to find image %v with size %d",
			d.Image, d.DiskSize)
	}
	log.Debugf("Image %v(%d) = %s", d.Image, d.DiskSize, tpl)

	// Profile UUID
	profile, ok := topology.Profiles[strings.ToLower(d.InstanceProfile)]
	if !ok {
		return fmt.Errorf("Unable to find the %s profile",
			d.InstanceProfile)
	}
	log.Debugf("Profile %v = %s", d.InstanceProfile, profile)

	// Security group
	sg, ok := topology.SecurityGroups[d.SecurityGroup]
	if !ok {
		log.Infof("Security group %v does not exist, create it",
			d.SecurityGroup)
		rules := []egoscale.SecurityGroupRule{
			{
				SecurityGroupId: "",
				Cidr:            "0.0.0.0/0",
				Protocol:        "TCP",
				Port:            22,
			},
			{
				SecurityGroupId: "",
				Cidr:            "0.0.0.0/0",
				Protocol:        "TCP",
				Port:            2376,
			},
			{
				SecurityGroupId: "",
				Cidr:            "0.0.0.0/0",
				Protocol:        "ICMP",
				IcmpType:        8,
				IcmpCode:        0,
			},
		}
		sgresp, err := client.CreateSecurityGroupWithRules(d.SecurityGroup,
			rules,
			make([]egoscale.SecurityGroupRule, 0, 0))
		if err != nil {
			return err
		}
		sg = sgresp.Id
	}
	log.Debugf("Security group %v = %s", d.SecurityGroup, sg)

	if d.KeyPair == "" {
		log.Infof("Generate an SSH keypair...")
		kpresp, err := client.CreateKeypair(d.MachineName)
		if err != nil {
			return err
		}
		err = ioutil.WriteFile(d.GetSSHKeyPath(), []byte(kpresp.Privatekey), 0600)
		if err != nil {
			return err
		}
		d.KeyPair = d.MachineName
	}

	log.Infof("Spawn exoscale host...")

	userdata, err := d.getCloudInit()
	if err != nil {
		return err
	}
	log.Debugf("Using the following cloud-init file:")
	log.Debugf("%s", userdata)

	machineProfile := egoscale.MachineProfile{
		Template:        tpl,
		ServiceOffering: profile,
		SecurityGroups:  []string{sg},
		Userdata:        userdata,
		Zone:            zone,
		Keypair:         d.KeyPair,
		Name:            d.MachineName,
	}

	cvmresp, err := client.CreateVirtualMachine(machineProfile)
	if err != nil {
		return err
	}

	vm, err := d.waitForVM(client, cvmresp)
	if err != nil {
		return err
	}
	d.IPAddress = vm.Nic[0].Ipaddress
	d.Id = vm.Id

	return nil
}

func (d *Driver) Start() error {
	vmstate, err := d.GetState()
	if err != nil {
		return err
	}
	if vmstate == state.Running || vmstate == state.Starting {
		log.Infof("Host is already running or starting")
		return nil
	}

	client := egoscale.NewClient(d.URL, d.ApiKey, d.ApiSecretKey)
	svmresp, err := client.StartVirtualMachine(d.Id)
	if err != nil {
		return err
	}
	_, err = d.waitForVM(client, svmresp)
	if err != nil {
		return err
	}
	return nil
}

func (d *Driver) Stop() error {
	vmstate, err := d.GetState()
	if err != nil {
		return err
	}
	if vmstate == state.Stopped {
		log.Infof("Host is already stopped")
		return nil
	}

	client := egoscale.NewClient(d.URL, d.ApiKey, d.ApiSecretKey)
	svmresp, err := client.StopVirtualMachine(d.Id)
	if err != nil {
		return err
	}
	_, err = d.waitForVM(client, svmresp)
	if err != nil {
		return err
	}
	return nil
}

func (d *Driver) Remove() error {
	client := egoscale.NewClient(d.URL, d.ApiKey, d.ApiSecretKey)
	dvmresp, err := client.DestroyVirtualMachine(d.Id)
	if err != nil {
		return err
	}
	_, err = d.waitForVM(client, dvmresp)
	if err != nil {
		return err
	}
	return nil
}

func (d *Driver) Restart() error {
	vmstate, err := d.GetState()
	if err != nil {
		return err
	}
	if vmstate == state.Stopped {
		return fmt.Errorf("Host is stopped, use start command to start it")
	}

	client := egoscale.NewClient(d.URL, d.ApiKey, d.ApiSecretKey)
	svmresp, err := client.RebootVirtualMachine(d.Id)
	if err != nil {
		return err
	}
	_, err = d.waitForVM(client, svmresp)
	if err != nil {
		return err
	}

	return nil
}

func (d *Driver) Kill() error {
	return d.Stop()
}

func (d *Driver) waitForVM(client *egoscale.Client, jobid string) (*egoscale.DeployVirtualMachineResponse, error) {
	log.Infof("Waiting for VM...")
	maxRepeats := 60
	i := 0
	var resp *egoscale.QueryAsyncJobResultResponse
	var err error
	for ; i < maxRepeats; i++ {
		resp, err = client.PollAsyncJob(jobid)
		if err != nil {
			return nil, err
		}

		if resp.Jobstatus == 1 {
			break
		}
		time.Sleep(2 * time.Second)
	}
	if i == maxRepeats {
		return nil, fmt.Errorf("Timeout while waiting for VM")
	}
	vm, err := client.AsyncToVirtualMachine(*resp)
	if err != nil {
		return nil, err
	}

	return vm, nil
}

// Build a cloud-init user data string that will install and run
// docker.
func (d *Driver) getCloudInit() (string, error) {
	const tpl = `#cloud-config
manage_etc_hosts: true
fqdn: {{ .MachineName }}
resize_rootfs: true
`
	var buffer bytes.Buffer

	tmpl, err := template.New("cloud-init").Parse(tpl)
	if err != nil {
		return "", err
	}
	err = tmpl.Execute(&buffer, d)
	if err != nil {
		return "", err
	}
	return buffer.String(), nil
}
<commit_msg>exoscale: destroy the SSH keypair when removing the machine<commit_after>package exoscale

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"path/filepath"
	"strings"
	"text/template"
	"time"

	"github.com/codegangsta/cli"
	"github.com/docker/machine/drivers"
	"github.com/docker/machine/log"
	"github.com/docker/machine/provider"
	"github.com/docker/machine/state"
	"github.com/pyr/egoscale/src/egoscale"
)

type Driver struct {
	URL              string
	ApiKey           string
	ApiSecretKey     string
	InstanceProfile  string
	DiskSize         int
	Image            string
	SecurityGroup    string
	AvailabilityZone string
	MachineName      string
	KeyPair          string
	IPAddress        string
	PublicKey        string
	Id               string
	CaCertPath       string
	PrivateKeyPath   string
	DriverKeyPath    string
	SwarmMaster      bool
	SwarmHost        string
	SwarmDiscovery   string
	storePath        string
}

func init() {
	drivers.Register("exoscale", &drivers.RegisteredDriver{
		New:            NewDriver,
		GetCreateFlags: GetCreateFlags,
	})
}

// RegisterCreateFlags registers the flags this driver adds to
// "docker hosts create"
func GetCreateFlags() []cli.Flag {
	return []cli.Flag{
		cli.StringFlag{
			EnvVar: "EXOSCALE_ENDPOINT",
			Name:   "exoscale-url",
			Usage:  "exoscale API endpoint",
		},
		cli.StringFlag{
			EnvVar: "EXOSCALE_API_KEY",
			Name:   "exoscale-api-key",
			Usage:  "exoscale API key",
		},
		cli.StringFlag{
			EnvVar: "EXOSCALE_API_SECRET",
			Name:   "exoscale-api-secret-key",
			Usage:  "exoscale API secret key",
		},
		cli.StringFlag{
			EnvVar: "EXOSCALE_INSTANCE_PROFILE",
			Name:   "exoscale-instance-profile",
			Value:  "small",
			Usage:  "exoscale instance profile (small, medium, large, ...)",
		},
		cli.IntFlag{
			EnvVar: "EXOSCALE_DISK_SIZE",
			Name:   "exoscale-disk-size",
			Value:  50,
			Usage:  "exoscale disk size (10, 50, 100, 200, 400)",
		},
		cli.StringFlag{
			EnvVar: "EXSOCALE_IMAGE",
			Name:   "exoscale-image",
			Value:  "ubuntu-14.04",
			Usage:  "exoscale image template",
		},
		cli.StringFlag{
			EnvVar: "EXOSCALE_SECURITY_GROUP",
			Name:   "exoscale-security-group",
			Value:  "docker-machine",
			Usage:  "exoscale security group",
		},
		cli.StringFlag{
			EnvVar: "EXOSCALE_AVAILABILITY_ZONE",
			Name:   "exoscale-availability-zone",
			Value:  "ch-gva-2",
			Usage:  "exoscale availibility zone",
		},
	}
}

func NewDriver(machineName string, storePath string, caCert string, privateKey string) (drivers.Driver, error) {
	return &Driver{MachineName: machineName, storePath: storePath, CaCertPath: caCert, PrivateKeyPath: privateKey}, nil
}

func (d *Driver) AuthorizePort(ports []*drivers.Port) error {
	return nil
}

func (d *Driver) DeauthorizePort(ports []*drivers.Port) error {
	return nil
}

func (d *Driver) GetMachineName() string {
	return d.MachineName
}

func (d *Driver) GetSSHHostname() (string, error) {
	return d.GetIP()
}

func (d *Driver) GetSSHKeyPath() string {
	return filepath.Join(d.storePath, "id_rsa")
}

func (d *Driver) GetSSHPort() (int, error) {
	return 22, nil
}

func (d *Driver) GetSSHUsername() string {
	return "ubuntu"
}

func (d *Driver) GetProviderType() provider.ProviderType {
	return provider.Remote
}

func (d *Driver) DriverName() string {
	return "exoscale"
}

func (d *Driver) SetConfigFromFlags(flags drivers.DriverOptions) error {
	d.URL = flags.String("exoscale-endpoint")
	d.ApiKey = flags.String("exoscale-api-key")
	d.ApiSecretKey = flags.String("exoscale-api-secret-key")
	d.InstanceProfile = flags.String("exoscale-instance-profile")
	d.DiskSize = flags.Int("exoscale-disk-size")
	d.Image = flags.String("exoscale-image")
	d.SecurityGroup = flags.String("exoscale-security-group")
	d.AvailabilityZone = flags.String("exoscale-availability-zone")
	d.SwarmMaster = flags.Bool("swarm-master")
	d.SwarmHost = flags.String("swarm-host")
	d.SwarmDiscovery = flags.String("swarm-discovery")

	if d.URL == "" {
		d.URL = "https://api.exoscale.ch/compute"
	}
	if d.ApiKey == "" || d.ApiSecretKey == "" {
		return fmt.Errorf("Please specify an API key (--exoscale-api-key) and an API secret key (--exoscale-api-secret-key).")
	}

	return nil
}

func (d *Driver) GetURL() (string, error) {
	ip, err := d.GetIP()
	if err != nil {
		return "", err
	}
	return fmt.Sprintf("tcp://%s:2376", ip), nil
}

func (d *Driver) GetIP() (string, error) {
	if d.IPAddress == "" {
		return "", fmt.Errorf("IP address is not set")
	}
	return d.IPAddress, nil
}

func (d *Driver) GetState() (state.State, error) {
	client := egoscale.NewClient(d.URL, d.ApiKey, d.ApiSecretKey)
	vm, err := client.GetVirtualMachine(d.Id)
	if err != nil {
		return state.Error, err
	}
	switch vm.State {
	case "Starting":
		return state.Starting, nil
	case "Running":
		return state.Running, nil
	case "Stopping":
		return state.Running, nil
	case "Stopped":
		return state.Stopped, nil
	case "Destroyed":
		return state.Stopped, nil
	case "Expunging":
		return state.Stopped, nil
	case "Migrating":
		return state.Paused, nil
	case "Error":
		return state.Error, nil
	case "Unknown":
		return state.Error, nil
	case "Shutdowned":
		return state.Stopped, nil
	}
	return state.None, nil
}

func (d *Driver) PreCreateCheck() error {
	return nil
}

func (d *Driver) Create() error {
	log.Infof("Querying exoscale for the requested parameters...")
	client := egoscale.NewClient(d.URL, d.ApiKey, d.ApiSecretKey)
	topology, err := client.GetTopology()
	if err != nil {
		return err
	}

	// Availability zone UUID
	zone, ok := topology.Zones[d.AvailabilityZone]
	if !ok {
		return fmt.Errorf("Availability zone %v doesn't exist",
			d.AvailabilityZone)
	}
	log.Debugf("Availability zone %v = %s", d.AvailabilityZone, zone)

	// Image UUID
	var tpl string
	images, ok := topology.Images[strings.ToLower(d.Image)]
	if ok {
		tpl, ok = images[d.DiskSize]
	}
	if !ok {
		return fmt.Errorf("Unable to find image %v with size %d",
			d.Image, d.DiskSize)
	}
	log.Debugf("Image %v(%d) = %s", d.Image, d.DiskSize, tpl)

	// Profile UUID
	profile, ok := topology.Profiles[strings.ToLower(d.InstanceProfile)]
	if !ok {
		return fmt.Errorf("Unable to find the %s profile",
			d.InstanceProfile)
	}
	log.Debugf("Profile %v = %s", d.InstanceProfile, profile)

	// Security group
	sg, ok := topology.SecurityGroups[d.SecurityGroup]
	if !ok {
		log.Infof("Security group %v does not exist, create it",
			d.SecurityGroup)
		rules := []egoscale.SecurityGroupRule{
			{
				SecurityGroupId: "",
				Cidr:            "0.0.0.0/0",
				Protocol:        "TCP",
				Port:            22,
			},
			{
				SecurityGroupId: "",
				Cidr:            "0.0.0.0/0",
				Protocol:        "TCP",
				Port:            2376,
			},
			{
				SecurityGroupId: "",
				Cidr:            "0.0.0.0/0",
				Protocol:        "ICMP",
				IcmpType:        8,
				IcmpCode:        0,
			},
		}
		sgresp, err := client.CreateSecurityGroupWithRules(d.SecurityGroup,
			rules,
			make([]egoscale.SecurityGroupRule, 0, 0))
		if err != nil {
			return err
		}
		sg = sgresp.Id
	}
	log.Debugf("Security group %v = %s", d.SecurityGroup, sg)

	log.Infof("Generate an SSH keypair...")
	keypairName := fmt.Sprintf("docker-machine-%s", d.MachineName)
	kpresp, err := client.CreateKeypair(keypairName)
	if err != nil {
		return err
	}
	err = ioutil.WriteFile(d.GetSSHKeyPath(), []byte(kpresp.Privatekey), 0600)
	if err != nil {
		return err
	}
	d.KeyPair = keypairName

	log.Infof("Spawn exoscale host...")

	userdata, err := d.getCloudInit()
	if err != nil {
		return err
	}
	log.Debugf("Using the following cloud-init file:")
	log.Debugf("%s", userdata)

	machineProfile := egoscale.MachineProfile{
		Template:        tpl,
		ServiceOffering: profile,
		SecurityGroups:  []string{sg},
		Userdata:        userdata,
		Zone:            zone,
		Keypair:         d.KeyPair,
		Name:            d.MachineName,
	}

	cvmresp, err := client.CreateVirtualMachine(machineProfile)
	if err != nil {
		return err
	}

	vm, err := d.waitForVM(client, cvmresp)
	if err != nil {
		return err
	}
	d.IPAddress = vm.Nic[0].Ipaddress
	d.Id = vm.Id

	return nil
}

func (d *Driver) Start() error {
	vmstate, err := d.GetState()
	if err != nil {
		return err
	}
	if vmstate == state.Running || vmstate == state.Starting {
		log.Infof("Host is already running or starting")
		return nil
	}

	client := egoscale.NewClient(d.URL, d.ApiKey, d.ApiSecretKey)
	svmresp, err := client.StartVirtualMachine(d.Id)
	if err != nil {
		return err
	}
	_, err = d.waitForVM(client, svmresp)
	if err != nil {
		return err
	}
	return nil
}

func (d *Driver) Stop() error {
	vmstate, err := d.GetState()
	if err != nil {
		return err
	}
	if vmstate == state.Stopped {
		log.Infof("Host is already stopped")
		return nil
	}

	client := egoscale.NewClient(d.URL, d.ApiKey, d.ApiSecretKey)
	svmresp, err := client.StopVirtualMachine(d.Id)
	if err != nil {
		return err
	}
	_, err = d.waitForVM(client, svmresp)
	if err != nil {
		return err
	}
	return nil
}

func (d *Driver) Remove() error {
	client := egoscale.NewClient(d.URL, d.ApiKey, d.ApiSecretKey)

	// Destroy the SSH key
	_, err := client.DeleteKeypair(d.KeyPair)
	if err != nil {
		return err
	}

	// Destroy the virtual machine
	dvmresp, err := client.DestroyVirtualMachine(d.Id)
	if err != nil {
		return err
	}
	_, err = d.waitForVM(client, dvmresp)
	if err != nil {
		return err
	}
	return nil
}

func (d *Driver) Restart() error {
	vmstate, err := d.GetState()
	if err != nil {
		return err
	}
	if vmstate == state.Stopped {
		return fmt.Errorf("Host is stopped, use start command to start it")
	}

	client := egoscale.NewClient(d.URL, d.ApiKey, d.ApiSecretKey)
	svmresp, err := client.RebootVirtualMachine(d.Id)
	if err != nil {
		return err
	}
	_, err = d.waitForVM(client, svmresp)
	if err != nil {
		return err
	}

	return nil
}

func (d *Driver) Kill() error {
	return d.Stop()
}

func (d *Driver) waitForVM(client *egoscale.Client, jobid string) (*egoscale.DeployVirtualMachineResponse, error) {
	log.Infof("Waiting for VM...")
	maxRepeats := 60
	i := 0
	var resp *egoscale.QueryAsyncJobResultResponse
	var err error
	for ; i < maxRepeats; i++ {
		resp, err = client.PollAsyncJob(jobid)
		if err != nil {
			return nil, err
		}

		if resp.Jobstatus == 1 {
			break
		}
		time.Sleep(2 * time.Second)
	}
	if i == maxRepeats {
		return nil, fmt.Errorf("Timeout while waiting for VM")
	}
	vm, err := client.AsyncToVirtualMachine(*resp)
	if err != nil {
		return nil, err
	}

	return vm, nil
}

// Build a cloud-init user data string that will install and run
// docker.
func (d *Driver) getCloudInit() (string, error) {
	const tpl = `#cloud-config
manage_etc_hosts: true
fqdn: {{ .MachineName }}
resize_rootfs: true
`
	var buffer bytes.Buffer

	tmpl, err := template.New("cloud-init").Parse(tpl)
	if err != nil {
		return "", err
	}
	err = tmpl.Execute(&buffer, d)
	if err != nil {
		return "", err
	}
	return buffer.String(), nil
}
<|endoftext|>


<commit_before>package req

import (
	"encoding/json"
	"encoding/xml"
	"fmt"
	"io"
	"net/http"
	"os"
	"regexp"
	"time"
)

// Resp represents a request with it's response
type Resp struct {
	r      *Req
	req    *http.Request
	resp   *http.Response
	client *http.Client
	*multipartHelper
	reqBody  []byte
	respBody []byte
	cost     time.Duration
}

// Cost returns time spent by the request
func (r *Resp) Cost() time.Duration {
	return r.cost
}

// Request returns *http.Request
func (r *Resp) Request() *http.Request {
	return r.req
}

// Response returns *http.Response
func (r *Resp) Response() *http.Response {
	return r.resp
}

// Bytes returns response body as []byte
func (r *Resp) Bytes() []byte {
	return r.respBody
}

// String returns response body as string
func (r *Resp) String() string {
	return string(r.respBody)
}

// ToJSON convert json response body to struct or map
func (r *Resp) ToJSON(v interface{}) error {
	return json.Unmarshal(r.respBody, v)
}

// ToXML convert xml response body to struct or map
func (r *Resp) ToXML(v interface{}) error {
	return xml.Unmarshal(r.respBody, v)
}

// ToFile download the response body to file
func (r *Resp) ToFile(name string) error {
	file, err := os.Create(name)
	defer file.Close()
	if err != nil {
		return err
	}
	if r.respBody != nil {
		_, err = file.Write(r.respBody)
		if err != nil {
			return err
		}
	} else {
		_, err = io.Copy(file, r.resp.Body)
		if err != nil {
			return err
		}
	}
	return nil
}

var regNewline = regexp.MustCompile(`\n|\r`)

func (r *Resp) Format(s fmt.State, verb rune) {
	if r == nil || r.req == nil {
		return
	}
	req := r.req
	if s.Flag('+') { // include header and format pretty.
		fmt.Fprint(s, r.dump())
	} else if s.Flag('-') { // keep all informations in one line.
		fmt.Fprint(s, req.Method, " ", req.URL.String())
		if r.r.flag&Lcost != 0 {
			fmt.Fprint(s, " ", r.cost.String())
		}
		if r.r.flag&LreqBody != 0 {
			if len(r.reqBody) > 0 {
				str := regNewline.ReplaceAllString(string(r.reqBody), " ")
				fmt.Fprint(s, " ", str)
			} else {
				fmt.Fprint(s, " ******")
			}
		}
		if r.r.flag&LrespBody != 0 {
			if len(r.respBody) > 0 {
				str := regNewline.ReplaceAllString(string(r.respBody), " ")
				fmt.Fprint(s, " ", str)
			} else {
				fmt.Fprint(s, " ******")
			}
		}
	} else { // auto
		fmt.Fprint(s, req.Method, " ", req.URL.String())
		if r.r.flag&Lcost != 0 {
			fmt.Fprint(s, " ", r.cost.String())
		}
		respBody := r.respBody
		if (len(r.reqBody) > 0 && regNewline.Match(r.reqBody)) || (len(respBody) > 0 && regNewline.Match(respBody)) { // pretty format
			if len(r.reqBody) > 0 {
				fmt.Fprint(s, "\n", string(r.reqBody))
			}
			if len(respBody) > 0 {
				fmt.Fprint(s, "\n", string(respBody))
			}
		} else {
			if len(r.reqBody) > 0 {
				fmt.Fprint(s, " ", string(r.reqBody))
			}
			if len(respBody) > 0 {
				fmt.Fprint(s, " ", string(respBody))
			}
		}
	}

}
<commit_msg>fix #9<commit_after>package req

import (
	"encoding/json"
	"encoding/xml"
	"fmt"
	"io"
	"net/http"
	"os"
	"regexp"
	"time"
)

// Resp represents a request with it's response
type Resp struct {
	r      *Req
	req    *http.Request
	resp   *http.Response
	client *http.Client
	*multipartHelper
	reqBody  []byte
	respBody []byte
	cost     time.Duration
}

// Cost returns time spent by the request
func (r *Resp) Cost() time.Duration {
	return r.cost
}

// Request returns *http.Request
func (r *Resp) Request() *http.Request {
	return r.req
}

// Response returns *http.Response
func (r *Resp) Response() *http.Response {
	return r.resp
}

// Bytes returns response body as []byte
func (r *Resp) Bytes() []byte {
	return r.respBody
}

// String returns response body as string
func (r *Resp) String() string {
	return string(r.respBody)
}

// ToJSON convert json response body to struct or map
func (r *Resp) ToJSON(v interface{}) error {
	return json.Unmarshal(r.respBody, v)
}

// ToXML convert xml response body to struct or map
func (r *Resp) ToXML(v interface{}) error {
	return xml.Unmarshal(r.respBody, v)
}

// ToFile download the response body to file
func (r *Resp) ToFile(name string) error {
	file, err := os.Create(name)
	if err != nil {
		return err
	}
	defer file.Close()
	if r.respBody != nil {
		_, err = file.Write(r.respBody)
		if err != nil {
			return err
		}
	} else {
		_, err = io.Copy(file, r.resp.Body)
		if err != nil {
			return err
		}
	}
	return nil
}

var regNewline = regexp.MustCompile(`\n|\r`)

func (r *Resp) Format(s fmt.State, verb rune) {
	if r == nil || r.req == nil {
		return
	}
	req := r.req
	if s.Flag('+') { // include header and format pretty.
		fmt.Fprint(s, r.dump())
	} else if s.Flag('-') { // keep all informations in one line.
		fmt.Fprint(s, req.Method, " ", req.URL.String())
		if r.r.flag&Lcost != 0 {
			fmt.Fprint(s, " ", r.cost.String())
		}
		if r.r.flag&LreqBody != 0 {
			if len(r.reqBody) > 0 {
				str := regNewline.ReplaceAllString(string(r.reqBody), " ")
				fmt.Fprint(s, " ", str)
			} else {
				fmt.Fprint(s, " ******")
			}
		}
		if r.r.flag&LrespBody != 0 {
			if len(r.respBody) > 0 {
				str := regNewline.ReplaceAllString(string(r.respBody), " ")
				fmt.Fprint(s, " ", str)
			} else {
				fmt.Fprint(s, " ******")
			}
		}
	} else { // auto
		fmt.Fprint(s, req.Method, " ", req.URL.String())
		if r.r.flag&Lcost != 0 {
			fmt.Fprint(s, " ", r.cost.String())
		}
		respBody := r.respBody
		if (len(r.reqBody) > 0 && regNewline.Match(r.reqBody)) || (len(respBody) > 0 && regNewline.Match(respBody)) { // pretty format
			if len(r.reqBody) > 0 {
				fmt.Fprint(s, "\n", string(r.reqBody))
			}
			if len(respBody) > 0 {
				fmt.Fprint(s, "\n", string(respBody))
			}
		} else {
			if len(r.reqBody) > 0 {
				fmt.Fprint(s, " ", string(r.reqBody))
			}
			if len(respBody) > 0 {
				fmt.Fprint(s, " ", string(respBody))
			}
		}
	}

}
<|endoftext|>


<commit_before>package exoscale

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"path/filepath"
	"strings"
	"text/template"
	"time"

	"github.com/codegangsta/cli"
	"github.com/docker/machine/drivers"
	"github.com/docker/machine/log"
	"github.com/docker/machine/state"
	"github.com/docker/machine/utils"
	"github.com/pyr/egoscale/src/egoscale"
)

type Driver struct {
	URL              string
	ApiKey           string
	ApiSecretKey     string
	InstanceProfile  string
	DiskSize         int
	Image            string
	SecurityGroup    string
	AvailabilityZone string
	MachineName      string
	KeyPair          string
	IPAddress        string
	PublicKey        string
	Id               string
	CaCertPath       string
	PrivateKeyPath   string
	DriverKeyPath    string
	SwarmMaster      bool
	SwarmHost        string
	SwarmDiscovery   string
	storePath        string
}

func init() {
	drivers.Register("exoscale", &drivers.RegisteredDriver{
		New:            NewDriver,
		GetCreateFlags: GetCreateFlags,
	})
}

// RegisterCreateFlags registers the flags this driver adds to
// "docker hosts create"
func GetCreateFlags() []cli.Flag {
	return []cli.Flag{
		cli.StringFlag{
			EnvVar: "EXOSCALE_ENDPOINT",
			Name:   "exoscale-url",
			Usage:  "exoscale API endpoint",
		},
		cli.StringFlag{
			EnvVar: "EXOSCALE_API_KEY",
			Name:   "exoscale-api-key",
			Usage:  "exoscale API key",
		},
		cli.StringFlag{
			EnvVar: "EXOSCALE_API_SECRET",
			Name:   "exoscale-api-secret-key",
			Usage:  "exoscale API secret key",
		},
		cli.StringFlag{
			EnvVar: "EXOSCALE_INSTANCE_PROFILE",
			Name:   "exoscale-instance-profile",
			Value:  "small",
			Usage:  "exoscale instance profile (small, medium, large, ...)",
		},
		cli.IntFlag{
			EnvVar: "EXOSCALE_DISK_SIZE",
			Name:   "exoscale-disk-size",
			Value:  50,
			Usage:  "exoscale disk size (10, 50, 100, 200, 400)",
		},
		cli.StringFlag{
			EnvVar: "EXSOCALE_IMAGE",
			Name:   "exoscale-image",
			Value:  "ubuntu-14.04",
			Usage:  "exoscale image template",
		},
		cli.StringFlag{
			EnvVar: "EXOSCALE_SECURITY_GROUP",
			Name:   "exoscale-security-group",
			Value:  "docker-machine",
			Usage:  "exoscale security group",
		},
		cli.StringFlag{
			EnvVar: "EXOSCALE_AVAILABILITY_ZONE",
			Name:   "exoscale-availability-zone",
			Value:  "ch-gva-2",
			Usage:  "exoscale availibility zone",
		},
	}
}

func NewDriver(machineName string, storePath string, caCert string, privateKey string) (drivers.Driver, error) {
	return &Driver{MachineName: machineName, storePath: storePath, CaCertPath: caCert, PrivateKeyPath: privateKey}, nil
}

func (d *Driver) AuthorizePort(ports []*drivers.Port) error {
	return nil
}

func (d *Driver) DeauthorizePort(ports []*drivers.Port) error {
	return nil
}

func (d *Driver) GetMachineName() string {
	return d.MachineName
}

func (d *Driver) GetSSHHostname() (string, error) {
	return d.GetIP()
}

func (d *Driver) GetSSHKeyPath() string {
	return filepath.Join(d.storePath, "id_rsa")
}

func (d *Driver) GetSSHPort() (int, error) {
	return 22, nil
}

func (d *Driver) GetSSHUsername() string {
	return "ubuntu"
}

func (d *Driver) DriverName() string {
	return "exoscale"
}

func (d *Driver) SetConfigFromFlags(flags drivers.DriverOptions) error {
	d.URL = flags.String("exoscale-endpoint")
	d.ApiKey = flags.String("exoscale-api-key")
	d.ApiSecretKey = flags.String("exoscale-api-secret-key")
	d.InstanceProfile = flags.String("exoscale-instance-profile")
	d.DiskSize = flags.Int("exoscale-disk-size")
	d.Image = flags.String("exoscale-image")
	d.SecurityGroup = flags.String("exoscale-security-group")
	d.AvailabilityZone = flags.String("exoscale-availability-zone")
	d.SwarmMaster = flags.Bool("swarm-master")
	d.SwarmHost = flags.String("swarm-host")
	d.SwarmDiscovery = flags.String("swarm-discovery")

	if d.URL == "" {
		d.URL = "https://api.exoscale.ch/compute"
	}
	if d.ApiKey == "" || d.ApiSecretKey == "" {
		return fmt.Errorf("Please specify an API key (--exoscale-api-key) and an API secret key (--exoscale-api-secret-key).")
	}

	return nil
}

func (d *Driver) GetURL() (string, error) {
	ip, err := d.GetIP()
	if err != nil {
		return "", err
	}
	return fmt.Sprintf("tcp://%s:2376", ip), nil
}

func (d *Driver) GetIP() (string, error) {
	if d.IPAddress == "" {
		return "", fmt.Errorf("IP address is not set")
	}
	return d.IPAddress, nil
}

func (d *Driver) GetState() (state.State, error) {
	client := egoscale.NewClient(d.URL, d.ApiKey, d.ApiSecretKey)
	vm, err := client.GetVirtualMachine(d.Id)
	if err != nil {
		return state.Error, err
	}
	switch vm.State {
	case "Starting":
		return state.Starting, nil
	case "Running":
		return state.Running, nil
	case "Stopping":
		return state.Running, nil
	case "Stopped":
		return state.Stopped, nil
	case "Destroyed":
		return state.Stopped, nil
	case "Expunging":
		return state.Stopped, nil
	case "Migrating":
		return state.Paused, nil
	case "Error":
		return state.Error, nil
	case "Unknown":
		return state.Error, nil
	case "Shutdowned":
		return state.Stopped, nil
	}
	return state.None, nil
}

func (d *Driver) PreCreateCheck() error {
	return nil
}

func (d *Driver) Create() error {
	log.Infof("Querying exoscale for the requested parameters...")
	client := egoscale.NewClient(d.URL, d.ApiKey, d.ApiSecretKey)
	topology, err := client.GetTopology()
	if err != nil {
		return err
	}

	// Availability zone UUID
	zone, ok := topology.Zones[d.AvailabilityZone]
	if !ok {
		return fmt.Errorf("Availability zone %v doesn't exist",
			d.AvailabilityZone)
	}
	log.Debugf("Availability zone %v = %s", d.AvailabilityZone, zone)

	// Image UUID
	var tpl string
	images, ok := topology.Images[strings.ToLower(d.Image)]
	if ok {
		tpl, ok = images[d.DiskSize]
	}
	if !ok {
		return fmt.Errorf("Unable to find image %v with size %d",
			d.Image, d.DiskSize)
	}
	log.Debugf("Image %v(%d) = %s", d.Image, d.DiskSize, tpl)

	// Profile UUID
	profile, ok := topology.Profiles[strings.ToLower(d.InstanceProfile)]
	if !ok {
		return fmt.Errorf("Unable to find the %s profile",
			d.InstanceProfile)
	}
	log.Debugf("Profile %v = %s", d.InstanceProfile, profile)

	// Security group
	sg, ok := topology.SecurityGroups[d.SecurityGroup]
	if !ok {
		log.Infof("Security group %v does not exist, create it",
			d.SecurityGroup)
		rules := []egoscale.SecurityGroupRule{
			{
				SecurityGroupId: "",
				Cidr:            "0.0.0.0/0",
				Protocol:        "TCP",
				Port:            22,
			},
			{
				SecurityGroupId: "",
				Cidr:            "0.0.0.0/0",
				Protocol:        "TCP",
				Port:            2376,
			},
			{
				SecurityGroupId: "",
				Cidr:            "0.0.0.0/0",
				Protocol:        "TCP",
				Port:            3376,
			},
			{
				SecurityGroupId: "",
				Cidr:            "0.0.0.0/0",
				Protocol:        "ICMP",
				IcmpType:        8,
				IcmpCode:        0,
			},
		}
		sgresp, err := client.CreateSecurityGroupWithRules(d.SecurityGroup,
			rules,
			make([]egoscale.SecurityGroupRule, 0, 0))
		if err != nil {
			return err
		}
		sg = sgresp.Id
	}
	log.Debugf("Security group %v = %s", d.SecurityGroup, sg)

	log.Infof("Generate an SSH keypair...")
	keypairName := fmt.Sprintf("docker-machine-%s", d.MachineName)
	kpresp, err := client.CreateKeypair(keypairName)
	if err != nil {
		return err
	}
	err = ioutil.WriteFile(d.GetSSHKeyPath(), []byte(kpresp.Privatekey), 0600)
	if err != nil {
		return err
	}
	d.KeyPair = keypairName

	log.Infof("Spawn exoscale host...")

	userdata, err := d.getCloudInit()
	if err != nil {
		return err
	}
	log.Debugf("Using the following cloud-init file:")
	log.Debugf("%s", userdata)

	machineProfile := egoscale.MachineProfile{
		Template:        tpl,
		ServiceOffering: profile,
		SecurityGroups:  []string{sg},
		Userdata:        userdata,
		Zone:            zone,
		Keypair:         d.KeyPair,
		Name:            d.MachineName,
	}

	cvmresp, err := client.CreateVirtualMachine(machineProfile)
	if err != nil {
		return err
	}

	vm, err := d.waitForVM(client, cvmresp)
	if err != nil {
		return err
	}
	d.IPAddress = vm.Nic[0].Ipaddress
	d.Id = vm.Id

	return nil
}

func (d *Driver) Start() error {
	vmstate, err := d.GetState()
	if err != nil {
		return err
	}
	if vmstate == state.Running || vmstate == state.Starting {
		log.Infof("Host is already running or starting")
		return nil
	}

	client := egoscale.NewClient(d.URL, d.ApiKey, d.ApiSecretKey)
	svmresp, err := client.StartVirtualMachine(d.Id)
	if err != nil {
		return err
	}
	_, err = d.waitForVM(client, svmresp)
	if err != nil {
		return err
	}
	return nil
}

func (d *Driver) Stop() error {
	vmstate, err := d.GetState()
	if err != nil {
		return err
	}
	if vmstate == state.Stopped {
		log.Infof("Host is already stopped")
		return nil
	}

	client := egoscale.NewClient(d.URL, d.ApiKey, d.ApiSecretKey)
	svmresp, err := client.StopVirtualMachine(d.Id)
	if err != nil {
		return err
	}
	if _, err = d.waitForVM(client, svmresp); err != nil {
		return err
	}
	return nil
}

func (d *Driver) Remove() error {
	client := egoscale.NewClient(d.URL, d.ApiKey, d.ApiSecretKey)

	// Destroy the SSH key
	if _, err := client.DeleteKeypair(d.KeyPair); err != nil {
		return err
	}

	// Destroy the virtual machine
	dvmresp, err := client.DestroyVirtualMachine(d.Id)
	if err != nil {
		return err
	}
	_, err = d.waitForVM(client, dvmresp)
	if err != nil {
		return err
	}
	return nil
}

func (d *Driver) Restart() error {
	vmstate, err := d.GetState()
	if err != nil {
		return err
	}
	if vmstate == state.Stopped {
		return fmt.Errorf("Host is stopped, use start command to start it")
	}

	client := egoscale.NewClient(d.URL, d.ApiKey, d.ApiSecretKey)
	svmresp, err := client.RebootVirtualMachine(d.Id)
	if err != nil {
		return err
	}
	if _, err = d.waitForVM(client, svmresp); err != nil {
		return err
	}

	return nil
}

func (d *Driver) Kill() error {
	return d.Stop()
}

func (d *Driver) waitForVM(client *egoscale.Client, jobid string) (*egoscale.DeployVirtualMachineResponse, error) {
	log.Infof("Waiting for VM...")
	var resp *egoscale.QueryAsyncJobResultResponse
	var err error
	if err = utils.WaitForSpecificOrError(func() (bool, error) {
		resp, err = client.PollAsyncJob(jobid)
		if err != nil {
			return true, err
		}
		switch resp.Jobstatus {
		case 0: // Job is still in progress
		case 1: // Job has successfully completed
			return true, nil
		case 2: // Job has failed to complete
			return true, fmt.Errorf("Operation failed to complete")
		default: // Some other code
		}
		return false, nil
	}, 60, 2*time.Second); err != nil {
		return nil, err
	}
	vm, err := client.AsyncToVirtualMachine(*resp)
	if err != nil {
		return nil, err
	}

	return vm, nil
}

// Build a cloud-init user data string that will install and run
// docker.
func (d *Driver) getCloudInit() (string, error) {
	const tpl = `#cloud-config
manage_etc_hosts: true
fqdn: {{ .MachineName }}
resize_rootfs: true
`
	var buffer bytes.Buffer

	tmpl, err := template.New("cloud-init").Parse(tpl)
	if err != nil {
		return "", err
	}
	err = tmpl.Execute(&buffer, d)
	if err != nil {
		return "", err
	}
	return buffer.String(), nil
}
<commit_msg>exoscale: move job waiting into a separate function<commit_after>package exoscale

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"path/filepath"
	"strings"
	"text/template"
	"time"

	"github.com/codegangsta/cli"
	"github.com/docker/machine/drivers"
	"github.com/docker/machine/log"
	"github.com/docker/machine/state"
	"github.com/docker/machine/utils"
	"github.com/pyr/egoscale/src/egoscale"
)

type Driver struct {
	URL              string
	ApiKey           string
	ApiSecretKey     string
	InstanceProfile  string
	DiskSize         int
	Image            string
	SecurityGroup    string
	AvailabilityZone string
	MachineName      string
	KeyPair          string
	IPAddress        string
	PublicKey        string
	Id               string
	CaCertPath       string
	PrivateKeyPath   string
	DriverKeyPath    string
	SwarmMaster      bool
	SwarmHost        string
	SwarmDiscovery   string
	storePath        string
}

func init() {
	drivers.Register("exoscale", &drivers.RegisteredDriver{
		New:            NewDriver,
		GetCreateFlags: GetCreateFlags,
	})
}

// RegisterCreateFlags registers the flags this driver adds to
// "docker hosts create"
func GetCreateFlags() []cli.Flag {
	return []cli.Flag{
		cli.StringFlag{
			EnvVar: "EXOSCALE_ENDPOINT",
			Name:   "exoscale-url",
			Usage:  "exoscale API endpoint",
		},
		cli.StringFlag{
			EnvVar: "EXOSCALE_API_KEY",
			Name:   "exoscale-api-key",
			Usage:  "exoscale API key",
		},
		cli.StringFlag{
			EnvVar: "EXOSCALE_API_SECRET",
			Name:   "exoscale-api-secret-key",
			Usage:  "exoscale API secret key",
		},
		cli.StringFlag{
			EnvVar: "EXOSCALE_INSTANCE_PROFILE",
			Name:   "exoscale-instance-profile",
			Value:  "small",
			Usage:  "exoscale instance profile (small, medium, large, ...)",
		},
		cli.IntFlag{
			EnvVar: "EXOSCALE_DISK_SIZE",
			Name:   "exoscale-disk-size",
			Value:  50,
			Usage:  "exoscale disk size (10, 50, 100, 200, 400)",
		},
		cli.StringFlag{
			EnvVar: "EXSOCALE_IMAGE",
			Name:   "exoscale-image",
			Value:  "ubuntu-14.04",
			Usage:  "exoscale image template",
		},
		cli.StringFlag{
			EnvVar: "EXOSCALE_SECURITY_GROUP",
			Name:   "exoscale-security-group",
			Value:  "docker-machine",
			Usage:  "exoscale security group",
		},
		cli.StringFlag{
			EnvVar: "EXOSCALE_AVAILABILITY_ZONE",
			Name:   "exoscale-availability-zone",
			Value:  "ch-gva-2",
			Usage:  "exoscale availibility zone",
		},
	}
}

func NewDriver(machineName string, storePath string, caCert string, privateKey string) (drivers.Driver, error) {
	return &Driver{MachineName: machineName, storePath: storePath, CaCertPath: caCert, PrivateKeyPath: privateKey}, nil
}

func (d *Driver) AuthorizePort(ports []*drivers.Port) error {
	return nil
}

func (d *Driver) DeauthorizePort(ports []*drivers.Port) error {
	return nil
}

func (d *Driver) GetMachineName() string {
	return d.MachineName
}

func (d *Driver) GetSSHHostname() (string, error) {
	return d.GetIP()
}

func (d *Driver) GetSSHKeyPath() string {
	return filepath.Join(d.storePath, "id_rsa")
}

func (d *Driver) GetSSHPort() (int, error) {
	return 22, nil
}

func (d *Driver) GetSSHUsername() string {
	return "ubuntu"
}

func (d *Driver) DriverName() string {
	return "exoscale"
}

func (d *Driver) SetConfigFromFlags(flags drivers.DriverOptions) error {
	d.URL = flags.String("exoscale-endpoint")
	d.ApiKey = flags.String("exoscale-api-key")
	d.ApiSecretKey = flags.String("exoscale-api-secret-key")
	d.InstanceProfile = flags.String("exoscale-instance-profile")
	d.DiskSize = flags.Int("exoscale-disk-size")
	d.Image = flags.String("exoscale-image")
	d.SecurityGroup = flags.String("exoscale-security-group")
	d.AvailabilityZone = flags.String("exoscale-availability-zone")
	d.SwarmMaster = flags.Bool("swarm-master")
	d.SwarmHost = flags.String("swarm-host")
	d.SwarmDiscovery = flags.String("swarm-discovery")

	if d.URL == "" {
		d.URL = "https://api.exoscale.ch/compute"
	}
	if d.ApiKey == "" || d.ApiSecretKey == "" {
		return fmt.Errorf("Please specify an API key (--exoscale-api-key) and an API secret key (--exoscale-api-secret-key).")
	}

	return nil
}

func (d *Driver) GetURL() (string, error) {
	ip, err := d.GetIP()
	if err != nil {
		return "", err
	}
	return fmt.Sprintf("tcp://%s:2376", ip), nil
}

func (d *Driver) GetIP() (string, error) {
	if d.IPAddress == "" {
		return "", fmt.Errorf("IP address is not set")
	}
	return d.IPAddress, nil
}

func (d *Driver) GetState() (state.State, error) {
	client := egoscale.NewClient(d.URL, d.ApiKey, d.ApiSecretKey)
	vm, err := client.GetVirtualMachine(d.Id)
	if err != nil {
		return state.Error, err
	}
	switch vm.State {
	case "Starting":
		return state.Starting, nil
	case "Running":
		return state.Running, nil
	case "Stopping":
		return state.Running, nil
	case "Stopped":
		return state.Stopped, nil
	case "Destroyed":
		return state.Stopped, nil
	case "Expunging":
		return state.Stopped, nil
	case "Migrating":
		return state.Paused, nil
	case "Error":
		return state.Error, nil
	case "Unknown":
		return state.Error, nil
	case "Shutdowned":
		return state.Stopped, nil
	}
	return state.None, nil
}

func (d *Driver) PreCreateCheck() error {
	return nil
}

func (d *Driver) Create() error {
	log.Infof("Querying exoscale for the requested parameters...")
	client := egoscale.NewClient(d.URL, d.ApiKey, d.ApiSecretKey)
	topology, err := client.GetTopology()
	if err != nil {
		return err
	}

	// Availability zone UUID
	zone, ok := topology.Zones[d.AvailabilityZone]
	if !ok {
		return fmt.Errorf("Availability zone %v doesn't exist",
			d.AvailabilityZone)
	}
	log.Debugf("Availability zone %v = %s", d.AvailabilityZone, zone)

	// Image UUID
	var tpl string
	images, ok := topology.Images[strings.ToLower(d.Image)]
	if ok {
		tpl, ok = images[d.DiskSize]
	}
	if !ok {
		return fmt.Errorf("Unable to find image %v with size %d",
			d.Image, d.DiskSize)
	}
	log.Debugf("Image %v(%d) = %s", d.Image, d.DiskSize, tpl)

	// Profile UUID
	profile, ok := topology.Profiles[strings.ToLower(d.InstanceProfile)]
	if !ok {
		return fmt.Errorf("Unable to find the %s profile",
			d.InstanceProfile)
	}
	log.Debugf("Profile %v = %s", d.InstanceProfile, profile)

	// Security group
	sg, ok := topology.SecurityGroups[d.SecurityGroup]
	if !ok {
		log.Infof("Security group %v does not exist, create it",
			d.SecurityGroup)
		rules := []egoscale.SecurityGroupRule{
			{
				SecurityGroupId: "",
				Cidr:            "0.0.0.0/0",
				Protocol:        "TCP",
				Port:            22,
			},
			{
				SecurityGroupId: "",
				Cidr:            "0.0.0.0/0",
				Protocol:        "TCP",
				Port:            2376,
			},
			{
				SecurityGroupId: "",
				Cidr:            "0.0.0.0/0",
				Protocol:        "TCP",
				Port:            3376,
			},
			{
				SecurityGroupId: "",
				Cidr:            "0.0.0.0/0",
				Protocol:        "ICMP",
				IcmpType:        8,
				IcmpCode:        0,
			},
		}
		sgresp, err := client.CreateSecurityGroupWithRules(d.SecurityGroup,
			rules,
			make([]egoscale.SecurityGroupRule, 0, 0))
		if err != nil {
			return err
		}
		sg = sgresp.Id
	}
	log.Debugf("Security group %v = %s", d.SecurityGroup, sg)

	log.Infof("Generate an SSH keypair...")
	keypairName := fmt.Sprintf("docker-machine-%s", d.MachineName)
	kpresp, err := client.CreateKeypair(keypairName)
	if err != nil {
		return err
	}
	err = ioutil.WriteFile(d.GetSSHKeyPath(), []byte(kpresp.Privatekey), 0600)
	if err != nil {
		return err
	}
	d.KeyPair = keypairName

	log.Infof("Spawn exoscale host...")

	userdata, err := d.getCloudInit()
	if err != nil {
		return err
	}
	log.Debugf("Using the following cloud-init file:")
	log.Debugf("%s", userdata)

	machineProfile := egoscale.MachineProfile{
		Template:        tpl,
		ServiceOffering: profile,
		SecurityGroups:  []string{sg},
		Userdata:        userdata,
		Zone:            zone,
		Keypair:         d.KeyPair,
		Name:            d.MachineName,
	}

	cvmresp, err := client.CreateVirtualMachine(machineProfile)
	if err != nil {
		return err
	}

	vm, err := d.waitForVM(client, cvmresp)
	if err != nil {
		return err
	}
	d.IPAddress = vm.Nic[0].Ipaddress
	d.Id = vm.Id

	return nil
}

func (d *Driver) Start() error {
	vmstate, err := d.GetState()
	if err != nil {
		return err
	}
	if vmstate == state.Running || vmstate == state.Starting {
		log.Infof("Host is already running or starting")
		return nil
	}

	client := egoscale.NewClient(d.URL, d.ApiKey, d.ApiSecretKey)
	svmresp, err := client.StartVirtualMachine(d.Id)
	if err != nil {
		return err
	}
	if err = d.waitForJob(client, svmresp); err != nil {
		return err
	}
	return nil
}

func (d *Driver) Stop() error {
	vmstate, err := d.GetState()
	if err != nil {
		return err
	}
	if vmstate == state.Stopped {
		log.Infof("Host is already stopped")
		return nil
	}

	client := egoscale.NewClient(d.URL, d.ApiKey, d.ApiSecretKey)
	svmresp, err := client.StopVirtualMachine(d.Id)
	if err != nil {
		return err
	}
	if err = d.waitForJob(client, svmresp); err != nil {
		return err
	}
	return nil
}

func (d *Driver) Remove() error {
	client := egoscale.NewClient(d.URL, d.ApiKey, d.ApiSecretKey)

	// Destroy the SSH key
	if _, err := client.DeleteKeypair(d.KeyPair); err != nil {
		return err
	}

	// Destroy the virtual machine
	dvmresp, err := client.DestroyVirtualMachine(d.Id)
	if err != nil {
		return err
	}
	if err = d.waitForJob(client, dvmresp); err != nil {
		return err
	}
	return nil
}

func (d *Driver) Restart() error {
	vmstate, err := d.GetState()
	if err != nil {
		return err
	}
	if vmstate == state.Stopped {
		return fmt.Errorf("Host is stopped, use start command to start it")
	}

	client := egoscale.NewClient(d.URL, d.ApiKey, d.ApiSecretKey)
	svmresp, err := client.RebootVirtualMachine(d.Id)
	if err != nil {
		return err
	}
	if err = d.waitForJob(client, svmresp); err != nil {
		return err
	}

	return nil
}

func (d *Driver) Kill() error {
	return d.Stop()
}

func (d *Driver) jobIsDone(client *egoscale.Client, jobid string) (bool, error) {
	resp, err := client.PollAsyncJob(jobid)
	if err != nil {
		return true, err
	}
	switch resp.Jobstatus {
	case 0: // Job is still in progress
	case 1: // Job has successfully completed
		return true, nil
	case 2: // Job has failed to complete
		return true, fmt.Errorf("Operation failed to complete")
	default: // Some other code
	}
	return false, nil
}

func (d *Driver) waitForJob(client *egoscale.Client, jobid string) error {
	log.Infof("Waiting for job to complete...")
	return utils.WaitForSpecificOrError(func() (bool, error) {
		return d.jobIsDone(client, jobid)
	}, 60, 2*time.Second)
}

func (d *Driver) waitForVM(client *egoscale.Client, jobid string) (*egoscale.DeployVirtualMachineResponse, error) {
	if err := d.waitForJob(client, jobid); err != nil {
		return nil, err
	}
	resp, err := client.PollAsyncJob(jobid)
	if err != nil {
		return nil, err
	}
	vm, err := client.AsyncToVirtualMachine(*resp)
	if err != nil {
		return nil, err
	}

	return vm, nil
}

// Build a cloud-init user data string that will install and run
// docker.
func (d *Driver) getCloudInit() (string, error) {
	const tpl = `#cloud-config
manage_etc_hosts: true
fqdn: {{ .MachineName }}
resize_rootfs: true
`
	var buffer bytes.Buffer

	tmpl, err := template.New("cloud-init").Parse(tpl)
	if err != nil {
		return "", err
	}
	err = tmpl.Execute(&buffer, d)
	if err != nil {
		return "", err
	}
	return buffer.String(), nil
}
<|endoftext|>


<commit_before>package mysql

import (
	"time"

	"github.com/funkygao/dbus/engine"
	"github.com/funkygao/dbus/pkg/cluster"
	"github.com/funkygao/dbus/pkg/model"
	"github.com/funkygao/dbus/pkg/myslave"
	conf "github.com/funkygao/jsconf"
	log "github.com/funkygao/log4go"
)

// MysqlbinlogInput is an input plugin that pretends to be a mysql instance's
// slave and consumes mysql binlog events.
type MysqlbinlogInput struct {
	maxEventLength int
	stopChan       chan struct{}

	slave *myslave.MySlave
	cf    *conf.Conf
}

func (this *MysqlbinlogInput) Init(config *conf.Conf) {
	this.maxEventLength = config.Int("max_event_length", (1<<20)-100)
	this.cf = config
	this.stopChan = make(chan struct{})
}

func (this *MysqlbinlogInput) Stop(r engine.InputRunner) {
	log.Debug("[%s] stopping...", r.Name())

	close(this.stopChan)
	if this.slave != nil {
		this.slave.StopReplication()
	}
}

func (this *MysqlbinlogInput) MySlave() *myslave.MySlave {
	return this.slave
}

// used only for dbc: ugly design
func (this *MysqlbinlogInput) ConnectMyslave(dsn string) {
	this.slave = myslave.New("", dsn, engine.Globals().ZrootCheckpoint).LoadConfig(this.cf)
}

func (this *MysqlbinlogInput) OnAck(pack *engine.Packet) error {
	return this.slave.MarkAsProcessed(pack.Payload.(*model.RowsEvent))
}

func (this *MysqlbinlogInput) Run(r engine.InputRunner, h engine.PluginHelper) error {
	name := r.Name()
	backoff := time.Second * 5

	globals := engine.Globals()
	var myResources []cluster.Resource
	resourcesCh := r.Resources()

	for {
	RESTART_REPLICATION:

		// wait till got some resource
		for {
			if len(myResources) != 0 {
				log.Trace("[%s] bingo! %d: %+v", name, len(myResources), myResources)
				break
			}

			log.Trace("[%s] awaiting resources", name)

			select {
			case <-this.stopChan:
				log.Debug("[%s] yes sir!", name)
				return nil
			case myResources = <-resourcesCh:
			}
		}

		dsn := myResources[0].DSN() // MysqlbinlogInput only consumes 1 resource
		log.Trace("[%s] starting replication from %s...", name, dsn)
		this.slave = myslave.New(name, dsn, globals.ZrootCheckpoint).LoadConfig(this.cf)
		if err := this.slave.AssertValidRowFormat(); err != nil {
			// err might be: read initial handshake error
			panic(err)
		}

		if img, err := this.slave.BinlogRowImage(); err != nil {
			log.Error("[%s] %v", name, err)
		} else {
			log.Trace("[%s] binlog row image=%s", name, img)
		}

		ready := make(chan struct{})
		go this.slave.StartReplication(ready)
		select {
		case <-ready:
		case <-this.stopChan:
			log.Debug("[%s] yes sir!", name)
			return nil
		}

		// TODO declare the real ownership of the resource

		rows := this.slave.Events()
		errors := this.slave.Errors()
		for {
			select {
			case <-this.stopChan:
				log.Debug("[%s] yes sir!", name)
				return nil

			case err := <-errors:
				// e,g.
				// ERROR 1236 (HY000): Could not find first log file name in binary log index file
				// ERROR 1236 (HY000): Could not open log file
				// read initial handshake error, caused by Too many connections
				log.Error("[%s] backoff %s: %v, stop from %s", name, backoff, err, dsn)
				this.slave.StopReplication()

				// myResources not changed, so next round still consume the same resources

				select {
				case <-time.After(backoff):
				case <-this.stopChan:
					log.Debug("[%s] yes sir!", name)
					return nil
				}
				goto RESTART_REPLICATION

			case pack, ok := <-r.InChan():
				if !ok {
					log.Debug("[%s] yes sir!", name)
					return nil
				}

				select {
				case err := <-errors:
					// TODO is this necessary?
					log.Error("[%s] backoff %s: %v, stop from %s", name, backoff, err, dsn)
					this.slave.StopReplication()

					select {
					case <-time.After(backoff):
					case <-this.stopChan:
						log.Debug("[%s] yes sir!", name)
						return nil
					}
					goto RESTART_REPLICATION

				case myResources = <-resourcesCh:
					log.Trace("[%s] cluster rebalanced, stop from %s", name, dsn)
					this.slave.StopReplication()
					goto RESTART_REPLICATION

				case row, ok := <-rows:
					if !ok {
						log.Info("[%s] event stream closed", name)
						return nil
					}

					if row.Length() < this.maxEventLength {
						pack.Payload = row
						r.Inject(pack)
					} else {
						// TODO this.slave.MarkAsProcessed(r), also consider batcher partial failure
						log.Warn("[%s] ignored len=%d %s", name, row.Length(), row.MetaInfo())
						pack.Recycle()
					}

				case <-this.stopChan:
					log.Debug("[%s] yes sir!", name)
					return nil
				}
			}
		}
	}

	return nil
}
<commit_msg>swtich to standalone mode when dsn is defined in conf<commit_after>package mysql

import (
	"time"

	"github.com/funkygao/dbus/engine"
	"github.com/funkygao/dbus/pkg/cluster"
	"github.com/funkygao/dbus/pkg/model"
	"github.com/funkygao/dbus/pkg/myslave"
	conf "github.com/funkygao/jsconf"
	log "github.com/funkygao/log4go"
)

// MysqlbinlogInput is an input plugin that pretends to be a mysql instance's
// slave and consumes mysql binlog events.
type MysqlbinlogInput struct {
	maxEventLength int
	stopChan       chan struct{}

	slave *myslave.MySlave
	cf    *conf.Conf
}

func (this *MysqlbinlogInput) Init(config *conf.Conf) {
	this.maxEventLength = config.Int("max_event_length", (1<<20)-100)
	this.cf = config
	this.stopChan = make(chan struct{})
}

func (this *MysqlbinlogInput) Stop(r engine.InputRunner) {
	log.Debug("[%s] stopping...", r.Name())

	close(this.stopChan)
	if this.slave != nil {
		this.slave.StopReplication()
	}
}

func (this *MysqlbinlogInput) MySlave() *myslave.MySlave {
	return this.slave
}

// used only for dbc: ugly design
func (this *MysqlbinlogInput) ConnectMyslave(dsn string) {
	this.slave = myslave.New("", dsn, engine.Globals().ZrootCheckpoint).LoadConfig(this.cf)
}

func (this *MysqlbinlogInput) OnAck(pack *engine.Packet) error {
	return this.slave.MarkAsProcessed(pack.Payload.(*model.RowsEvent))
}

func (this *MysqlbinlogInput) Run(r engine.InputRunner, h engine.PluginHelper) error {
	if dsn := r.Conf().String("dsn", ""); len(dsn) > 0 {
		return this.runStandalone(dsn, r, h)
	}

	return this.runCluster(r, h)
}

func (this *MysqlbinlogInput) runStandalone(dsn string, r engine.InputRunner, h engine.PluginHelper) error {
	return nil
}

func (this *MysqlbinlogInput) runCluster(r engine.InputRunner, h engine.PluginHelper) error {
	name := r.Name()
	backoff := time.Second * 5

	globals := engine.Globals()
	var myResources []cluster.Resource
	resourcesCh := r.Resources()

	for {
	RESTART_REPLICATION:

		// wait till got some resource
		for {
			if len(myResources) != 0 {
				log.Trace("[%s] bingo! %d: %+v", name, len(myResources), myResources)
				break
			}

			log.Trace("[%s] awaiting resources", name)

			select {
			case <-this.stopChan:
				log.Debug("[%s] yes sir!", name)
				return nil
			case myResources = <-resourcesCh:
			}
		}

		dsn := myResources[0].DSN() // MysqlbinlogInput only consumes 1 resource
		log.Trace("[%s] starting replication from %s...", name, dsn)
		this.slave = myslave.New(name, dsn, globals.ZrootCheckpoint).LoadConfig(this.cf)
		if err := this.slave.AssertValidRowFormat(); err != nil {
			// err might be: read initial handshake error
			panic(err)
		}

		if img, err := this.slave.BinlogRowImage(); err != nil {
			log.Error("[%s] %v", name, err)
		} else {
			log.Trace("[%s] binlog row image=%s", name, img)
		}

		ready := make(chan struct{})
		go this.slave.StartReplication(ready)
		select {
		case <-ready:
		case <-this.stopChan:
			log.Debug("[%s] yes sir!", name)
			return nil
		}

		// TODO declare the real ownership of the resource

		rows := this.slave.Events()
		errors := this.slave.Errors()
		for {
			select {
			case <-this.stopChan:
				log.Debug("[%s] yes sir!", name)
				return nil

			case err := <-errors:
				// e,g.
				// ERROR 1236 (HY000): Could not find first log file name in binary log index file
				// ERROR 1236 (HY000): Could not open log file
				// read initial handshake error, caused by Too many connections
				log.Error("[%s] backoff %s: %v, stop from %s", name, backoff, err, dsn)
				this.slave.StopReplication()

				// myResources not changed, so next round still consume the same resources

				select {
				case <-time.After(backoff):
				case <-this.stopChan:
					log.Debug("[%s] yes sir!", name)
					return nil
				}
				goto RESTART_REPLICATION

			case pack, ok := <-r.InChan():
				if !ok {
					log.Debug("[%s] yes sir!", name)
					return nil
				}

				select {
				case err := <-errors:
					// TODO is this necessary?
					log.Error("[%s] backoff %s: %v, stop from %s", name, backoff, err, dsn)
					this.slave.StopReplication()

					select {
					case <-time.After(backoff):
					case <-this.stopChan:
						log.Debug("[%s] yes sir!", name)
						return nil
					}
					goto RESTART_REPLICATION

				case myResources = <-resourcesCh:
					log.Trace("[%s] cluster rebalanced, stop from %s", name, dsn)
					this.slave.StopReplication()
					goto RESTART_REPLICATION

				case row, ok := <-rows:
					if !ok {
						log.Info("[%s] event stream closed", name)
						return nil
					}

					if row.Length() < this.maxEventLength {
						pack.Payload = row
						r.Inject(pack)
					} else {
						// TODO this.slave.MarkAsProcessed(r), also consider batcher partial failure
						log.Warn("[%s] ignored len=%d %s", name, row.Length(), row.MetaInfo())
						pack.Recycle()
					}

				case <-this.stopChan:
					log.Debug("[%s] yes sir!", name)
					return nil
				}
			}
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package cloudwatch

import (
	"fmt"
	"strings"
	"sync"
	"time"

	"github.com/aws/aws-sdk-go/aws"

	"github.com/aws/aws-sdk-go/service/cloudwatch"

	"github.com/influxdata/telegraf"
	"github.com/influxdata/telegraf/internal"
	internalaws "github.com/influxdata/telegraf/internal/config/aws"
	"github.com/influxdata/telegraf/internal/errchan"
	"github.com/influxdata/telegraf/internal/limiter"
	"github.com/influxdata/telegraf/plugins/inputs"
)

type (
	CloudWatch struct {
		Region    string `toml:"region"`
		AccessKey string `toml:"access_key"`
		SecretKey string `toml:"secret_key"`
		RoleARN   string `toml:"role_arn"`
		Profile   string `toml:"profile"`
		Filename  string `toml:"shared_credential_file"`
		Token     string `toml:"token"`

		Period      internal.Duration `toml:"period"`
		Delay       internal.Duration `toml:"delay"`
		Namespace   string            `toml:"namespace"`
		Metrics     []*Metric         `toml:"metrics"`
		CacheTTL    internal.Duration `toml:"cache_ttl"`
		RateLimit   string            `toml:"ratelimit"`
		client      cloudwatchClient
		metricCache *MetricCache
	}

	Metric struct {
		MetricNames []string     `toml:"names"`
		Dimensions  []*Dimension `toml:"dimensions"`
	}

	Dimension struct {
		Name  string `toml:"name"`
		Value string `toml:"value"`
	}

	MetricCache struct {
		TTL     time.Duration
		Fetched time.Time
		Metrics []*cloudwatch.Metric
	}

	cloudwatchClient interface {
		ListMetrics(*cloudwatch.ListMetricsInput) (*cloudwatch.ListMetricsOutput, error)
		GetMetricStatistics(*cloudwatch.GetMetricStatisticsInput) (*cloudwatch.GetMetricStatisticsOutput, error)
	}
)

func (c *CloudWatch) SampleConfig() string {
	return `
  ## Amazon Region
  region = 'us-east-1'

  ## Amazon Credentials
  ## Credentials are loaded in the following order
  ## 1) Assumed credentials via STS if role_arn is specified
  ## 2) explicit credentials from 'access_key' and 'secret_key'
  ## 3) shared profile from 'profile'
  ## 4) environment variables
  ## 5) shared credentials file
  ## 6) EC2 Instance Profile
  #access_key = ""
  #secret_key = ""
  #token = ""
  #role_arn = ""
  #profile = ""
  #shared_credential_file = ""

  ## Requested CloudWatch aggregation Period (required - must be a multiple of 60s)
  period = '1m'

  ## Collection Delay (required - must account for metrics availability via CloudWatch API)
  delay = '1m'

  ## Recomended: use metric 'interval' that is a multiple of 'period' to avoid
  ## gaps or overlap in pulled data
  interval = '1m'

  ## Configure the TTL for the internal cache of metrics.
  ## Defaults to 1 hr if not specified
  #cache_ttl = '10m'

  ## Metric Statistic Namespace (required)
  namespace = 'AWS/ELB'

  ## Metrics to Pull (optional)
  ## Defaults to all Metrics in Namespace if nothing is provided
  ## Refreshes Namespace available metrics every 1h
  #[[inputs.cloudwatch.metrics]]
  #  names = ['Latency', 'RequestCount']
  #
  #  ## Dimension filters for Metric (optional)
  #  [[inputs.cloudwatch.metrics.dimensions]]
  #    name = 'LoadBalancerName'
  #    value = 'p-example'
`
}

func (c *CloudWatch) Description() string {
	return "Pull Metric Statistics from Amazon CloudWatch"
}

func (c *CloudWatch) Gather(acc telegraf.Accumulator) error {
	if c.client == nil {
		c.initializeCloudWatch()
	}

	var metrics []*cloudwatch.Metric

	// check for provided metric filter
	if c.Metrics != nil {
		metrics = []*cloudwatch.Metric{}
		for _, m := range c.Metrics {
			if !hasWilcard(m.Dimensions) {
				dimensions := make([]*cloudwatch.Dimension, len(m.Dimensions))
				for k, d := range m.Dimensions {
					fmt.Printf("Dimension [%s]:[%s]\n", d.Name, d.Value)
					dimensions[k] = &cloudwatch.Dimension{
						Name:  aws.String(d.Name),
						Value: aws.String(d.Value),
					}
				}
				for _, name := range m.MetricNames {
					metrics = append(metrics, &cloudwatch.Metric{
						Namespace:  aws.String(c.Namespace),
						MetricName: aws.String(name),
						Dimensions: dimensions,
					})
				}
			} else {
				allMetrics, err := c.fetchNamespaceMetrics()
				if err != nil {
					return err
				}
				for _, name := range m.MetricNames {
					for _, metric := range allMetrics {
						if isSelected(metric, m.Dimensions) {
							metrics = append(metrics, &cloudwatch.Metric{
								Namespace:  aws.String(c.Namespace),
								MetricName: aws.String(name),
								Dimensions: metric.Dimensions,
							})
						}
					}
				}
			}

		}
	} else {
		var err error
		metrics, err = c.fetchNamespaceMetrics()
		if err != nil {
			return err
		}
	}

	metricCount := len(metrics)
	errChan := errchan.New(metricCount)

	now := time.Now()

	// limit concurrency or we can easily exhaust user connection limit
	// see cloudwatch API request limits:
	// http://docs.aws.amazon.com/AmazonCloudWatch/latest/DeveloperGuide/cloudwatch_limits.html
	lmtr := limiter.NewRateLimiter(c.RateLimit, time.Second)
	defer lmtr.Stop()
	var wg sync.WaitGroup
	wg.Add(len(metrics))
	for _, m := range metrics {
		<-lmtr.C
		go func(inm *cloudwatch.Metric) {
			defer wg.Done()
			c.gatherMetric(acc, inm, now, errChan.C)
		}(m)
	}
	wg.Wait()

	return errChan.Error()
}

func init() {
	inputs.Add("cloudwatch", func() telegraf.Input {
		ttl, _ := time.ParseDuration("1hr")
		return &CloudWatch{
			CacheTTL: internal.Duration{Duration: ttl},
		}
	})
}

/*
 * Initialize CloudWatch client
 */
func (c *CloudWatch) initializeCloudWatch() error {
	credentialConfig := &internalaws.CredentialConfig{
		Region:    c.Region,
		AccessKey: c.AccessKey,
		SecretKey: c.SecretKey,
		RoleARN:   c.RoleARN,
		Profile:   c.Profile,
		Filename:  c.Filename,
		Token:     c.Token,
	}
	configProvider := credentialConfig.Credentials()

	c.client = cloudwatch.New(configProvider)
	return nil
}

/*
 * Fetch available metrics for given CloudWatch Namespace
 */
func (c *CloudWatch) fetchNamespaceMetrics() (metrics []*cloudwatch.Metric, err error) {
	if c.metricCache != nil && c.metricCache.IsValid() {
		metrics = c.metricCache.Metrics
		return
	}

	metrics = []*cloudwatch.Metric{}

	var token *string
	for more := true; more; {
		params := &cloudwatch.ListMetricsInput{
			Namespace:  aws.String(c.Namespace),
			Dimensions: []*cloudwatch.DimensionFilter{},
			NextToken:  token,
			MetricName: nil,
		}

		resp, err := c.client.ListMetrics(params)
		if err != nil {
			return nil, err
		}

		metrics = append(metrics, resp.Metrics...)

		token = resp.NextToken
		more = token != nil
	}

	c.metricCache = &MetricCache{
		Metrics: metrics,
		Fetched: time.Now(),
		TTL:     c.CacheTTL.Duration,
	}

	return
}

/*
 * Gather given Metric and emit any error
 */
func (c *CloudWatch) gatherMetric(
	acc telegraf.Accumulator,
	metric *cloudwatch.Metric,
	now time.Time,
	errChan chan error,
) {
	params := c.getStatisticsInput(metric, now)
	resp, err := c.client.GetMetricStatistics(params)
	if err != nil {
		errChan <- err
		return
	}

	for _, point := range resp.Datapoints {
		tags := map[string]string{
			"region": c.Region,
			"unit":   snakeCase(*point.Unit),
		}

		for _, d := range metric.Dimensions {
			tags[snakeCase(*d.Name)] = *d.Value
		}

		// record field for each statistic
		fields := map[string]interface{}{}

		if point.Average != nil {
			fields[formatField(*metric.MetricName, cloudwatch.StatisticAverage)] = *point.Average
		}
		if point.Maximum != nil {
			fields[formatField(*metric.MetricName, cloudwatch.StatisticMaximum)] = *point.Maximum
		}
		if point.Minimum != nil {
			fields[formatField(*metric.MetricName, cloudwatch.StatisticMinimum)] = *point.Minimum
		}
		if point.SampleCount != nil {
			fields[formatField(*metric.MetricName, cloudwatch.StatisticSampleCount)] = *point.SampleCount
		}
		if point.Sum != nil {
			fields[formatField(*metric.MetricName, cloudwatch.StatisticSum)] = *point.Sum
		}

		acc.AddFields(formatMeasurement(c.Namespace), fields, tags, *point.Timestamp)
	}

	errChan <- nil
}

/*
 * Formatting helpers
 */
func formatField(metricName string, statistic string) string {
	return fmt.Sprintf("%s_%s", snakeCase(metricName), snakeCase(statistic))
}

func formatMeasurement(namespace string) string {
	namespace = strings.Replace(namespace, "/", "_", -1)
	namespace = snakeCase(namespace)
	return fmt.Sprintf("cloudwatch_%s", namespace)
}

func snakeCase(s string) string {
	s = internal.SnakeCase(s)
	s = strings.Replace(s, "__", "_", -1)
	return s
}

/*
 * Map Metric to *cloudwatch.GetMetricStatisticsInput for given timeframe
 */
func (c *CloudWatch) getStatisticsInput(metric *cloudwatch.Metric, now time.Time) *cloudwatch.GetMetricStatisticsInput {
	end := now.Add(-c.Delay.Duration)

	input := &cloudwatch.GetMetricStatisticsInput{
		StartTime:  aws.Time(end.Add(-c.Period.Duration)),
		EndTime:    aws.Time(end),
		MetricName: metric.MetricName,
		Namespace:  metric.Namespace,
		Period:     aws.Int64(int64(c.Period.Duration.Seconds())),
		Dimensions: metric.Dimensions,
		Statistics: []*string{
			aws.String(cloudwatch.StatisticAverage),
			aws.String(cloudwatch.StatisticMaximum),
			aws.String(cloudwatch.StatisticMinimum),
			aws.String(cloudwatch.StatisticSum),
			aws.String(cloudwatch.StatisticSampleCount)},
	}
	return input
}

/*
 * Check Metric Cache validity
 */
func (c *MetricCache) IsValid() bool {
	return c.Metrics != nil && time.Since(c.Fetched) < c.TTL
}

func hasWilcard(dimensions []*Dimension) bool {
	for _, d := range dimensions {
		if d.Value == "" || d.Value == "*" {
			return true
		}
	}
	return false
}

func isSelected(metric *cloudwatch.Metric, dimensions []*Dimension) bool {
	if len(metric.Dimensions) != len(dimensions) {
		return false
	}
	for _, d := range dimensions {
		selected := false
		for _, d2 := range metric.Dimensions {
			if d.Name == *d2.Name {
				if d.Value == "" || d.Value == "*" || d.Value == *d2.Value {
					selected = true
				}
			}
		}
		if !selected {
			return false
		}
	}
	return true
}
<commit_msg>ahem, apparently limiter wants an int<commit_after>package cloudwatch

import (
	"fmt"
	"strings"
	"sync"
	"time"

	"github.com/aws/aws-sdk-go/aws"

	"github.com/aws/aws-sdk-go/service/cloudwatch"

	"github.com/influxdata/telegraf"
	"github.com/influxdata/telegraf/internal"
	internalaws "github.com/influxdata/telegraf/internal/config/aws"
	"github.com/influxdata/telegraf/internal/errchan"
	"github.com/influxdata/telegraf/internal/limiter"
	"github.com/influxdata/telegraf/plugins/inputs"
)

type (
	CloudWatch struct {
		Region    string `toml:"region"`
		AccessKey string `toml:"access_key"`
		SecretKey string `toml:"secret_key"`
		RoleARN   string `toml:"role_arn"`
		Profile   string `toml:"profile"`
		Filename  string `toml:"shared_credential_file"`
		Token     string `toml:"token"`

		Period      internal.Duration `toml:"period"`
		Delay       internal.Duration `toml:"delay"`
		Namespace   string            `toml:"namespace"`
		Metrics     []*Metric         `toml:"metrics"`
		CacheTTL    internal.Duration `toml:"cache_ttl"`
		RateLimit   int               `toml:"ratelimit"`
		client      cloudwatchClient
		metricCache *MetricCache
	}

	Metric struct {
		MetricNames []string     `toml:"names"`
		Dimensions  []*Dimension `toml:"dimensions"`
	}

	Dimension struct {
		Name  string `toml:"name"`
		Value string `toml:"value"`
	}

	MetricCache struct {
		TTL     time.Duration
		Fetched time.Time
		Metrics []*cloudwatch.Metric
	}

	cloudwatchClient interface {
		ListMetrics(*cloudwatch.ListMetricsInput) (*cloudwatch.ListMetricsOutput, error)
		GetMetricStatistics(*cloudwatch.GetMetricStatisticsInput) (*cloudwatch.GetMetricStatisticsOutput, error)
	}
)

func (c *CloudWatch) SampleConfig() string {
	return `
  ## Amazon Region
  region = 'us-east-1'

  ## Amazon Credentials
  ## Credentials are loaded in the following order
  ## 1) Assumed credentials via STS if role_arn is specified
  ## 2) explicit credentials from 'access_key' and 'secret_key'
  ## 3) shared profile from 'profile'
  ## 4) environment variables
  ## 5) shared credentials file
  ## 6) EC2 Instance Profile
  #access_key = ""
  #secret_key = ""
  #token = ""
  #role_arn = ""
  #profile = ""
  #shared_credential_file = ""

  ## Requested CloudWatch aggregation Period (required - must be a multiple of 60s)
  period = '1m'

  ## Collection Delay (required - must account for metrics availability via CloudWatch API)
  delay = '1m'

  ## Recomended: use metric 'interval' that is a multiple of 'period' to avoid
  ## gaps or overlap in pulled data
  interval = '1m'

  ## Configure the TTL for the internal cache of metrics.
  ## Defaults to 1 hr if not specified
  #cache_ttl = '10m'

  ## Metric Statistic Namespace (required)
  namespace = 'AWS/ELB'

  ## Metrics to Pull (optional)
  ## Defaults to all Metrics in Namespace if nothing is provided
  ## Refreshes Namespace available metrics every 1h
  #[[inputs.cloudwatch.metrics]]
  #  names = ['Latency', 'RequestCount']
  #
  #  ## Dimension filters for Metric (optional)
  #  [[inputs.cloudwatch.metrics.dimensions]]
  #    name = 'LoadBalancerName'
  #    value = 'p-example'
`
}

func (c *CloudWatch) Description() string {
	return "Pull Metric Statistics from Amazon CloudWatch"
}

func (c *CloudWatch) Gather(acc telegraf.Accumulator) error {
	if c.client == nil {
		c.initializeCloudWatch()
	}

	var metrics []*cloudwatch.Metric

	// check for provided metric filter
	if c.Metrics != nil {
		metrics = []*cloudwatch.Metric{}
		for _, m := range c.Metrics {
			if !hasWilcard(m.Dimensions) {
				dimensions := make([]*cloudwatch.Dimension, len(m.Dimensions))
				for k, d := range m.Dimensions {
					fmt.Printf("Dimension [%s]:[%s]\n", d.Name, d.Value)
					dimensions[k] = &cloudwatch.Dimension{
						Name:  aws.String(d.Name),
						Value: aws.String(d.Value),
					}
				}
				for _, name := range m.MetricNames {
					metrics = append(metrics, &cloudwatch.Metric{
						Namespace:  aws.String(c.Namespace),
						MetricName: aws.String(name),
						Dimensions: dimensions,
					})
				}
			} else {
				allMetrics, err := c.fetchNamespaceMetrics()
				if err != nil {
					return err
				}
				for _, name := range m.MetricNames {
					for _, metric := range allMetrics {
						if isSelected(metric, m.Dimensions) {
							metrics = append(metrics, &cloudwatch.Metric{
								Namespace:  aws.String(c.Namespace),
								MetricName: aws.String(name),
								Dimensions: metric.Dimensions,
							})
						}
					}
				}
			}

		}
	} else {
		var err error
		metrics, err = c.fetchNamespaceMetrics()
		if err != nil {
			return err
		}
	}

	metricCount := len(metrics)
	errChan := errchan.New(metricCount)

	now := time.Now()

	// limit concurrency or we can easily exhaust user connection limit
	// see cloudwatch API request limits:
	// http://docs.aws.amazon.com/AmazonCloudWatch/latest/DeveloperGuide/cloudwatch_limits.html
	lmtr := limiter.NewRateLimiter(c.RateLimit, time.Second)
	defer lmtr.Stop()
	var wg sync.WaitGroup
	wg.Add(len(metrics))
	for _, m := range metrics {
		<-lmtr.C
		go func(inm *cloudwatch.Metric) {
			defer wg.Done()
			c.gatherMetric(acc, inm, now, errChan.C)
		}(m)
	}
	wg.Wait()

	return errChan.Error()
}

func init() {
	inputs.Add("cloudwatch", func() telegraf.Input {
		ttl, _ := time.ParseDuration("1hr")
		return &CloudWatch{
			CacheTTL: internal.Duration{Duration: ttl},
		}
	})
}

/*
 * Initialize CloudWatch client
 */
func (c *CloudWatch) initializeCloudWatch() error {
	credentialConfig := &internalaws.CredentialConfig{
		Region:    c.Region,
		AccessKey: c.AccessKey,
		SecretKey: c.SecretKey,
		RoleARN:   c.RoleARN,
		Profile:   c.Profile,
		Filename:  c.Filename,
		Token:     c.Token,
	}
	configProvider := credentialConfig.Credentials()

	c.client = cloudwatch.New(configProvider)
	return nil
}

/*
 * Fetch available metrics for given CloudWatch Namespace
 */
func (c *CloudWatch) fetchNamespaceMetrics() (metrics []*cloudwatch.Metric, err error) {
	if c.metricCache != nil && c.metricCache.IsValid() {
		metrics = c.metricCache.Metrics
		return
	}

	metrics = []*cloudwatch.Metric{}

	var token *string
	for more := true; more; {
		params := &cloudwatch.ListMetricsInput{
			Namespace:  aws.String(c.Namespace),
			Dimensions: []*cloudwatch.DimensionFilter{},
			NextToken:  token,
			MetricName: nil,
		}

		resp, err := c.client.ListMetrics(params)
		if err != nil {
			return nil, err
		}

		metrics = append(metrics, resp.Metrics...)

		token = resp.NextToken
		more = token != nil
	}

	c.metricCache = &MetricCache{
		Metrics: metrics,
		Fetched: time.Now(),
		TTL:     c.CacheTTL.Duration,
	}

	return
}

/*
 * Gather given Metric and emit any error
 */
func (c *CloudWatch) gatherMetric(
	acc telegraf.Accumulator,
	metric *cloudwatch.Metric,
	now time.Time,
	errChan chan error,
) {
	params := c.getStatisticsInput(metric, now)
	resp, err := c.client.GetMetricStatistics(params)
	if err != nil {
		errChan <- err
		return
	}

	for _, point := range resp.Datapoints {
		tags := map[string]string{
			"region": c.Region,
			"unit":   snakeCase(*point.Unit),
		}

		for _, d := range metric.Dimensions {
			tags[snakeCase(*d.Name)] = *d.Value
		}

		// record field for each statistic
		fields := map[string]interface{}{}

		if point.Average != nil {
			fields[formatField(*metric.MetricName, cloudwatch.StatisticAverage)] = *point.Average
		}
		if point.Maximum != nil {
			fields[formatField(*metric.MetricName, cloudwatch.StatisticMaximum)] = *point.Maximum
		}
		if point.Minimum != nil {
			fields[formatField(*metric.MetricName, cloudwatch.StatisticMinimum)] = *point.Minimum
		}
		if point.SampleCount != nil {
			fields[formatField(*metric.MetricName, cloudwatch.StatisticSampleCount)] = *point.SampleCount
		}
		if point.Sum != nil {
			fields[formatField(*metric.MetricName, cloudwatch.StatisticSum)] = *point.Sum
		}

		acc.AddFields(formatMeasurement(c.Namespace), fields, tags, *point.Timestamp)
	}

	errChan <- nil
}

/*
 * Formatting helpers
 */
func formatField(metricName string, statistic string) string {
	return fmt.Sprintf("%s_%s", snakeCase(metricName), snakeCase(statistic))
}

func formatMeasurement(namespace string) string {
	namespace = strings.Replace(namespace, "/", "_", -1)
	namespace = snakeCase(namespace)
	return fmt.Sprintf("cloudwatch_%s", namespace)
}

func snakeCase(s string) string {
	s = internal.SnakeCase(s)
	s = strings.Replace(s, "__", "_", -1)
	return s
}

/*
 * Map Metric to *cloudwatch.GetMetricStatisticsInput for given timeframe
 */
func (c *CloudWatch) getStatisticsInput(metric *cloudwatch.Metric, now time.Time) *cloudwatch.GetMetricStatisticsInput {
	end := now.Add(-c.Delay.Duration)

	input := &cloudwatch.GetMetricStatisticsInput{
		StartTime:  aws.Time(end.Add(-c.Period.Duration)),
		EndTime:    aws.Time(end),
		MetricName: metric.MetricName,
		Namespace:  metric.Namespace,
		Period:     aws.Int64(int64(c.Period.Duration.Seconds())),
		Dimensions: metric.Dimensions,
		Statistics: []*string{
			aws.String(cloudwatch.StatisticAverage),
			aws.String(cloudwatch.StatisticMaximum),
			aws.String(cloudwatch.StatisticMinimum),
			aws.String(cloudwatch.StatisticSum),
			aws.String(cloudwatch.StatisticSampleCount)},
	}
	return input
}

/*
 * Check Metric Cache validity
 */
func (c *MetricCache) IsValid() bool {
	return c.Metrics != nil && time.Since(c.Fetched) < c.TTL
}

func hasWilcard(dimensions []*Dimension) bool {
	for _, d := range dimensions {
		if d.Value == "" || d.Value == "*" {
			return true
		}
	}
	return false
}

func isSelected(metric *cloudwatch.Metric, dimensions []*Dimension) bool {
	if len(metric.Dimensions) != len(dimensions) {
		return false
	}
	for _, d := range dimensions {
		selected := false
		for _, d2 := range metric.Dimensions {
			if d.Name == *d2.Name {
				if d.Value == "" || d.Value == "*" || d.Value == *d2.Value {
					selected = true
				}
			}
		}
		if !selected {
			return false
		}
	}
	return true
}
<|endoftext|>


<commit_before>package core

import (
	"testing"

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

func TestIncludes(t *testing.T) {
	label1 := BuildLabel{PackageName: "src/core", Name: "..."}
	label2 := BuildLabel{PackageName: "src/parse", Name: "parse"}
	assert.False(t, label1.includes(label2))
	label2 = BuildLabel{PackageName: "src/core", Name: "core_test"}
	assert.True(t, label1.includes(label2))
}

func TestIncludesSubstring(t *testing.T) {
	label1 := BuildLabel{PackageName: "third_party/python", Name: "..."}
	label2 := BuildLabel{PackageName: "third_party/python3", Name: "six"}
	assert.False(t, label1.includes(label2))
}

func TestIncludesSubpackages(t *testing.T) {
	label1 := BuildLabel{PackageName: "", Name: "..."}
	label2 := BuildLabel{PackageName: "third_party/python3", Name: "six"}
	assert.True(t, label1.includes(label2))
}
<commit_msg>add test<commit_after>package core

import (
	"testing"

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

func TestIncludes(t *testing.T) {
	label1 := BuildLabel{PackageName: "src/core", Name: "..."}
	label2 := BuildLabel{PackageName: "src/parse", Name: "parse"}
	assert.False(t, label1.includes(label2))
	label2 = BuildLabel{PackageName: "src/core", Name: "core_test"}
	assert.True(t, label1.includes(label2))
}

func TestIncludesSubstring(t *testing.T) {
	label1 := BuildLabel{PackageName: "third_party/python", Name: "..."}
	label2 := BuildLabel{PackageName: "third_party/python3", Name: "six"}
	assert.False(t, label1.includes(label2))
}

func TestIncludesSubpackages(t *testing.T) {
	label1 := BuildLabel{PackageName: "", Name: "..."}
	label2 := BuildLabel{PackageName: "third_party/python3", Name: "six"}
	assert.True(t, label1.includes(label2))
}

func TestParent(t *testing.T) {
	label := BuildLabel{PackageName: "src/core", Name: "core"}
	assert.Equal(t, label, label.Parent())
	label2 := BuildLabel{PackageName: "src/core", Name: "_core#src"}
	assert.Equal(t, label, label2.Parent())
	label3 := BuildLabel{PackageName: "src/core", Name: "_core_something"}
	assert.Equal(t, label3, label3.Parent())
}
<|endoftext|>


<commit_before>package lzma2

import (
	"fmt"
	"testing"

	"github.com/uli-go/xz/lzma"
)

func TestChunkTypeString(t *testing.T) {
	tests := [...]struct {
		c chunkType
		s string
	}{
		{cEOS, "EOS"},
		{cU, "U"},
		{cUD, "UD"},
		{cL, "L"},
		{cLR, "LR"},
		{cLRN, "LRN"},
		{cLRND, "LRND"},
	}
	for _, c := range tests {
		s := fmt.Sprintf("%v", c.c)
		if s != c.s {
			t.Errorf("got %s; want %s", s, c.s)
		}
	}
}

func TestHeaderChunkType(t *testing.T) {
	tests := []struct {
		h byte
		c chunkType
	}{
		{h: 0, c: cEOS},
		{h: 1, c: cUD},
		{h: 2, c: cU},
		{h: 1<<7 | 0x1f, c: cL},
		{h: 1<<7 | 1<<5 | 0x1f, c: cLR},
		{h: 1<<7 | 1<<6 | 0x1f, c: cLRN},
		{h: 1<<7 | 1<<6 | 1<<5 | 0x1f, c: cLRND},
		{h: 1<<7 | 1<<6 | 1<<5, c: cLRND},
	}
	if _, err := headerChunkType(3); err == nil {
		t.Fatalf("headerChunkType(%d) got %v; want %v",
			3, err, errHeaderByte)
	}
	for _, tc := range tests {
		c, err := headerChunkType(tc.h)
		if err != nil {
			t.Fatalf("headerChunkType error %s", err)
		}
		if c != tc.c {
			t.Errorf("got %s; want %s", c, tc.c)
		}
	}
}

func TestHeaderLen(t *testing.T) {
	tests := []struct {
		c chunkType
		n int
	}{
		{cEOS, 1}, {cU, 3}, {cUD, 3}, {cL, 5}, {cLR, 5}, {cLRN, 6},
		{cLRND, 6},
	}
	for _, tc := range tests {
		n := headerLen(tc.c)
		if n != tc.n {
			t.Errorf("header length for %s %d; want %d",
				tc.c, n, tc.n)
		}
	}
}

func TestMarshalling(t *testing.T) {
	props, err := lzma.NewProperties(3, 0, 2)
	if err != nil {
		t.Fatalf("NewProperties(3, 0, 2) error %s", err)
	}

	var h, g chunkHeader
	for c := cEOS; c <= cLRND; c++ {
		h.ctype = c
		if c >= cLRN {
			h.props = props
		}
		data, err := h.MarshalBinary()
		if err != nil {
			t.Fatalf("MarshalBinary for %v error %s", &h, err)
		}
		if err = g.UnmarshalBinary(data); err != nil {
			t.Fatalf("UnmarshalBinary error %s", err)
		}
		if g != h {
			t.Fatalf("got %v; want %v", g, h)
		}
	}
}
<commit_msg>lzma2: resorted tests for cU<commit_after>package lzma2

import (
	"fmt"
	"testing"

	"github.com/uli-go/xz/lzma"
)

func TestChunkTypeString(t *testing.T) {
	tests := [...]struct {
		c chunkType
		s string
	}{
		{cEOS, "EOS"},
		{cUD, "UD"},
		{cU, "U"},
		{cL, "L"},
		{cLR, "LR"},
		{cLRN, "LRN"},
		{cLRND, "LRND"},
	}
	for _, c := range tests {
		s := fmt.Sprintf("%v", c.c)
		if s != c.s {
			t.Errorf("got %s; want %s", s, c.s)
		}
	}
}

func TestHeaderChunkType(t *testing.T) {
	tests := []struct {
		h byte
		c chunkType
	}{
		{h: 0, c: cEOS},
		{h: 1, c: cUD},
		{h: 2, c: cU},
		{h: 1<<7 | 0x1f, c: cL},
		{h: 1<<7 | 1<<5 | 0x1f, c: cLR},
		{h: 1<<7 | 1<<6 | 0x1f, c: cLRN},
		{h: 1<<7 | 1<<6 | 1<<5 | 0x1f, c: cLRND},
		{h: 1<<7 | 1<<6 | 1<<5, c: cLRND},
	}
	if _, err := headerChunkType(3); err == nil {
		t.Fatalf("headerChunkType(%d) got %v; want %v",
			3, err, errHeaderByte)
	}
	for _, tc := range tests {
		c, err := headerChunkType(tc.h)
		if err != nil {
			t.Fatalf("headerChunkType error %s", err)
		}
		if c != tc.c {
			t.Errorf("got %s; want %s", c, tc.c)
		}
	}
}

func TestHeaderLen(t *testing.T) {
	tests := []struct {
		c chunkType
		n int
	}{
		{cEOS, 1}, {cU, 3}, {cUD, 3}, {cL, 5}, {cLR, 5}, {cLRN, 6},
		{cLRND, 6},
	}
	for _, tc := range tests {
		n := headerLen(tc.c)
		if n != tc.n {
			t.Errorf("header length for %s %d; want %d",
				tc.c, n, tc.n)
		}
	}
}

func TestMarshalling(t *testing.T) {
	props, err := lzma.NewProperties(3, 0, 2)
	if err != nil {
		t.Fatalf("NewProperties(3, 0, 2) error %s", err)
	}

	var h, g chunkHeader
	for c := cEOS; c <= cLRND; c++ {
		h.ctype = c
		if c >= cLRN {
			h.props = props
		}
		data, err := h.MarshalBinary()
		if err != nil {
			t.Fatalf("MarshalBinary for %v error %s", &h, err)
		}
		if err = g.UnmarshalBinary(data); err != nil {
			t.Fatalf("UnmarshalBinary error %s", err)
		}
		if g != h {
			t.Fatalf("got %v; want %v", g, h)
		}
	}
}
<|endoftext|>


<commit_before>package lzma2

import (
	"fmt"
	"testing"

	"github.com/uli-go/xz/lzma"
)

func TestChunkTypeString(t *testing.T) {
	tests := [...]struct {
		c chunkType
		s string
	}{
		{cEOS, "EOS"},
		{cUD, "UD"},
		{cU, "U"},
		{cL, "L"},
		{cLR, "LR"},
		{cLRN, "LRN"},
		{cLRND, "LRND"},
	}
	for _, c := range tests {
		s := fmt.Sprintf("%v", c.c)
		if s != c.s {
			t.Errorf("got %s; want %s", s, c.s)
		}
	}
}

func TestHeaderChunkType(t *testing.T) {
	tests := []struct {
		h byte
		c chunkType
	}{
		{h: 0, c: cEOS},
		{h: 1, c: cUD},
		{h: 2, c: cU},
		{h: 1<<7 | 0x1f, c: cL},
		{h: 1<<7 | 1<<5 | 0x1f, c: cLR},
		{h: 1<<7 | 1<<6 | 0x1f, c: cLRN},
		{h: 1<<7 | 1<<6 | 1<<5 | 0x1f, c: cLRND},
		{h: 1<<7 | 1<<6 | 1<<5, c: cLRND},
	}
	if _, err := headerChunkType(3); err == nil {
		t.Fatalf("headerChunkType(%d) got %v; want %v",
			3, err, errHeaderByte)
	}
	for _, tc := range tests {
		c, err := headerChunkType(tc.h)
		if err != nil {
			t.Fatalf("headerChunkType error %s", err)
		}
		if c != tc.c {
			t.Errorf("got %s; want %s", c, tc.c)
		}
	}
}

func TestHeaderLen(t *testing.T) {
	tests := []struct {
		c chunkType
		n int
	}{
		{cEOS, 1}, {cU, 3}, {cUD, 3}, {cL, 5}, {cLR, 5}, {cLRN, 6},
		{cLRND, 6},
	}
	for _, tc := range tests {
		n := headerLen(tc.c)
		if n != tc.n {
			t.Errorf("header length for %s %d; want %d",
				tc.c, n, tc.n)
		}
	}
}

func chunkHeaderSamples(t *testing.T) []chunkHeader {
	props, err := lzma.NewProperties(3, 0, 2)
	if err != nil {
		t.Fatalf("NewProperties(3, 0, 2) error %s", err)
	}

	headers := make([]chunkHeader, 0, 12)
	for c := cEOS; c <= cLRND; c++ {
		var h chunkHeader
		h.ctype = c
		if c >= cUD {
			h.unpacked = 0x0304
		}
		if c >= cL {
			h.packed = 0x0201
		}
		if c >= cLRN {
			h.props = props
		}
		headers = append(headers, h)
	}
	return headers
}

func TestChunkHeaderMarshalling(t *testing.T) {
	for _, h := range chunkHeaderSamples(t) {
		data, err := h.MarshalBinary()
		if err != nil {
			t.Fatalf("MarshalBinary for %v error %s", &h, err)
		}
		var g chunkHeader
		if err = g.UnmarshalBinary(data); err != nil {
			t.Fatalf("UnmarshalBinary error %s", err)
		}
		if g != h {
			t.Fatalf("got %v; want %v", g, h)
		}
	}
}
<commit_msg>lzma2: added TestReadChunkHeader<commit_after>package lzma2

import (
	"bytes"
	"fmt"
	"testing"

	"github.com/uli-go/xz/lzma"
)

func TestChunkTypeString(t *testing.T) {
	tests := [...]struct {
		c chunkType
		s string
	}{
		{cEOS, "EOS"},
		{cUD, "UD"},
		{cU, "U"},
		{cL, "L"},
		{cLR, "LR"},
		{cLRN, "LRN"},
		{cLRND, "LRND"},
	}
	for _, c := range tests {
		s := fmt.Sprintf("%v", c.c)
		if s != c.s {
			t.Errorf("got %s; want %s", s, c.s)
		}
	}
}

func TestHeaderChunkType(t *testing.T) {
	tests := []struct {
		h byte
		c chunkType
	}{
		{h: 0, c: cEOS},
		{h: 1, c: cUD},
		{h: 2, c: cU},
		{h: 1<<7 | 0x1f, c: cL},
		{h: 1<<7 | 1<<5 | 0x1f, c: cLR},
		{h: 1<<7 | 1<<6 | 0x1f, c: cLRN},
		{h: 1<<7 | 1<<6 | 1<<5 | 0x1f, c: cLRND},
		{h: 1<<7 | 1<<6 | 1<<5, c: cLRND},
	}
	if _, err := headerChunkType(3); err == nil {
		t.Fatalf("headerChunkType(%d) got %v; want %v",
			3, err, errHeaderByte)
	}
	for _, tc := range tests {
		c, err := headerChunkType(tc.h)
		if err != nil {
			t.Fatalf("headerChunkType error %s", err)
		}
		if c != tc.c {
			t.Errorf("got %s; want %s", c, tc.c)
		}
	}
}

func TestHeaderLen(t *testing.T) {
	tests := []struct {
		c chunkType
		n int
	}{
		{cEOS, 1}, {cU, 3}, {cUD, 3}, {cL, 5}, {cLR, 5}, {cLRN, 6},
		{cLRND, 6},
	}
	for _, tc := range tests {
		n := headerLen(tc.c)
		if n != tc.n {
			t.Errorf("header length for %s %d; want %d",
				tc.c, n, tc.n)
		}
	}
}

func chunkHeaderSamples(t *testing.T) []chunkHeader {
	props, err := lzma.NewProperties(3, 0, 2)
	if err != nil {
		t.Fatalf("NewProperties(3, 0, 2) error %s", err)
	}

	headers := make([]chunkHeader, 0, 12)
	for c := cEOS; c <= cLRND; c++ {
		var h chunkHeader
		h.ctype = c
		if c >= cUD {
			h.unpacked = 0x0304
		}
		if c >= cL {
			h.packed = 0x0201
		}
		if c >= cLRN {
			h.props = props
		}
		headers = append(headers, h)
	}
	return headers
}

func TestChunkHeaderMarshalling(t *testing.T) {
	for _, h := range chunkHeaderSamples(t) {
		data, err := h.MarshalBinary()
		if err != nil {
			t.Fatalf("MarshalBinary for %v error %s", h, err)
		}
		var g chunkHeader
		if err = g.UnmarshalBinary(data); err != nil {
			t.Fatalf("UnmarshalBinary error %s", err)
		}
		if g != h {
			t.Fatalf("got %v; want %v", g, h)
		}
	}
}

func TestReadChunkHeader(t *testing.T) {
	for _, h := range chunkHeaderSamples(t) {
		data, err := h.MarshalBinary()
		if err != nil {
			t.Fatalf("MarshalBinary for %v error %s", h, err)
		}
		r := bytes.NewReader(data)
		g, err := readChunkHeader(r)
		if err != nil {
			t.Fatalf("readChunkHeader for %v error %s", h, err)
		}
		if *g != h {
			t.Fatalf("got %v; want %v", g, h)
		}
	}
}
<|endoftext|>