<commit_before>package aws

import (
	"encoding/base64"
	"encoding/json"
	"fmt"
	"strconv"

	"github.com/aws/aws-sdk-go/aws"
	gocf "github.com/crewjam/go-cloudformation"
	"github.com/google/uuid"
	"github.com/pchico83/i2kit/cli/providers/aws/ec2"
	"github.com/pchico83/i2kit/cli/schemas/compose"
	"github.com/pchico83/i2kit/cli/schemas/environment"
	"github.com/pchico83/i2kit/cli/schemas/service"
)

var amisPerRegion = map[string]string{
	"us-west-2": "ami-7293320a",
}

// Translate an i2kit service to a AWS CloudFormation template
func Translate(s *service.Service, e *environment.Environment, config *aws.Config) (string, error) {
	ami, ok := amisPerRegion[e.Provider.Region]
	if !ok {
		return "", fmt.Errorf("Region'%s' is not supported", e.Provider.Region)
	}
	t := gocf.NewTemplate()
	err := loadASG(t, s, e, ami, config)
	if err != nil {
		return "", err
	}
	err = loadELB(t, s, e)
	if err != nil {
		return "", err
	}
	loadIAM(t, s, e)
	loadLogGroup(t, s, e)
	loadRoute53(t, s, e)
	marshalledTemplate := []byte("")
	if marshalledTemplate, err = json.Marshal(t); err != nil {
		return "", err
	}
	return string(marshalledTemplate), nil
}

func loadASG(t *gocf.Template, s *service.Service, e *environment.Environment, ami string, config *aws.Config) error {
	vpc, err := ec2.GetVPC(e, config)
	if err != nil {
		return err
	}
	domain := e.Provider.HostedZone[:len(e.Provider.HostedZone)-1]
	encodedCompose, err := compose.Create(s, domain)
	if err != nil {
		return err
	}
	subnets := &gocf.StringListExpr{Literal: []*gocf.StringExpr{}}
	for _, item := range e.Provider.Subnets {
		subnets.Literal = append(subnets.Literal, gocf.String(*item))
	}
	replicas := strconv.Itoa(s.Replicas)
	asg := &gocf.AutoScalingAutoScalingGroup{
		HealthCheckGracePeriod:  gocf.Integer(15),
		LaunchConfigurationName: gocf.Ref("LaunchConfig").String(),
		LoadBalancerNames:       gocf.StringList(gocf.Ref("ELB")),
		MaxSize:                 gocf.String(replicas),
		MinSize:                 gocf.String(replicas),
		VPCZoneIdentifier:       subnets,
	}
	creationPolicy := &gocf.CreationPolicy{
		ResourceSignal: &gocf.CreationPolicyResourceSignal{
			Count:   gocf.Integer(int64(s.Replicas)),
			Timeout: gocf.String("PT15M"),
		},
	}
	updatePolicy := &gocf.UpdatePolicy{
		AutoScalingRollingUpdate: &gocf.UpdatePolicyAutoScalingRollingUpdate{
			MaxBatchSize:          gocf.Integer(1),
			PauseTime:             gocf.String("PT15M"),
			WaitOnResourceSignals: gocf.Bool(true),
			SuspendProcesses: gocf.StringList(
				gocf.String("HealthCheck"),
				gocf.String("ReplaceUnhealthy"),
				gocf.String("AZRebalance"),
				gocf.String("AlarmNotification"),
				gocf.String("ScheduledActions"),
			),
		},
	}
	asgResource := &gocf.Resource{
		Properties:     asg,
		CreationPolicy: creationPolicy,
		UpdatePolicy:   updatePolicy,
	}
	t.Resources["ASG"] = asgResource

	instanceIngressRules := gocf.EC2SecurityGroupRuleList{}
	loadbalancerIngressRules := gocf.EC2SecurityGroupRuleList{}
	for _, container := range s.Containers {
		for _, port := range container.Ports {
			instancePortNumber, _ := strconv.ParseInt(port.InstancePort, 10, 64)
			instanceIngressRules = append(
				instanceIngressRules,
				gocf.EC2SecurityGroupRule{
					SourceSecurityGroupIdXXSecurityGroupIngressXOnlyX: gocf.Ref("ELBSecurityGroup").String(),
					IpProtocol: gocf.String("tcp"),
					FromPort:   gocf.Integer(instancePortNumber),
					ToPort:     gocf.Integer(instancePortNumber),
				})
			portNumber, _ := strconv.ParseInt(port.Port, 10, 64)
			loadbalancerIngressRules = append(
				loadbalancerIngressRules,
				gocf.EC2SecurityGroupRule{
					CidrIp:     gocf.String("0.0.0.0/0"),
					IpProtocol: gocf.String("tcp"),
					FromPort:   gocf.Integer(portNumber),
					ToPort:     gocf.Integer(portNumber),
				})
		}
	}
	securityGroups := []gocf.Stringable{gocf.String(e.Provider.SecurityGroup)}
	if len(instanceIngressRules) > 0 {
		securityGroup := &gocf.EC2SecurityGroup{
			GroupDescription:     gocf.String(fmt.Sprintf("Instance Security Group for %s.%s", s.Name, e.Provider.Name)),
			SecurityGroupIngress: &instanceIngressRules,
			VpcId:                gocf.String(vpc),
		}
		t.AddResource("InstanceSecurityGroup", securityGroup)
		securityGroups = append(securityGroups, gocf.Ref("InstanceSecurityGroup").String())
		securityGroup = &gocf.EC2SecurityGroup{
			GroupDescription:     gocf.String(fmt.Sprintf("ELB Security Group for %s.%s", s.Name, e.Provider.Name)),
			SecurityGroupIngress: &loadbalancerIngressRules,
			VpcId:                gocf.String(vpc),
		}
		t.AddResource("ELBSecurityGroup", securityGroup)
	}

	launchConfig := &gocf.AutoScalingLaunchConfiguration{
		ImageId:            gocf.String(ami),
		InstanceType:       gocf.String(s.GetSize(e)),
		KeyName:            gocf.String(e.Provider.Keypair),
		SecurityGroups:     securityGroups,
		IamInstanceProfile: gocf.Ref("InstanceProfile").String(),
		UserData:           gocf.String(userData(s.Name, encodedCompose, e)),
	}
	t.AddResource("LaunchConfig", launchConfig)
	return nil
}

func userData(containerName, encodedCompose string, e *environment.Environment) string {
	uniqueOperationID := uuid.New().String()
	value := fmt.Sprintf(
		`#!/bin/bash

set -e
INSTANCE_ID=$(curl http://169.254.169.254/latest/meta-data/instance-id)
sudo docker run \
	--name %s \
	-e COMPOSE=%s \
	-e CONFIG=%s \
	-e UNIQUE_OPERATION_ID=%s \
	-e STACK=%s \
	-e REGION=%s \
	-v /var/run/docker.sock:/var/run/docker.sock \
	--log-driver=awslogs \
	--log-opt awslogs-region=%s \
	--log-opt awslogs-group=i2kit-%s \
	--log-opt tag=$INSTANCE_ID \
	riberaproject/agent`,
		containerName,
		encodedCompose,
		e.B64DockerConfig(),
		uniqueOperationID,
		containerName,
		e.Provider.Region,
		e.Provider.Region,
		containerName,
	)
	return base64.StdEncoding.EncodeToString([]byte(value))
}

func loadELB(t *gocf.Template, s *service.Service, e *environment.Environment) error {
	healthCheckPort := ""
	listeners := gocf.ElasticLoadBalancingListenerList{}
	for _, container := range s.Containers {
		for _, port := range container.Ports {
			if healthCheckPort == "" {
				healthCheckPort = port.InstancePort
			}
			certificate := port.Certificate
			if certificate == "" && (port.Protocol == "HTTPS" || port.Protocol == "SSL") {
				if e.Provider.Certificate == "" {
					return fmt.Errorf("Port '%s:%s' requires a certificate", port.Protocol, port.Port)
				}
				certificate = e.Provider.Certificate
			}
			listeners = append(listeners, gocf.ElasticLoadBalancingListener{
				InstancePort:     gocf.String(port.InstancePort),
				InstanceProtocol: gocf.String(port.InstanceProtocol),
				LoadBalancerPort: gocf.String(port.Port),
				Protocol:         gocf.String(port.Protocol),
				SSLCertificateId: gocf.String(certificate),
			})
		}
	}
	if len(listeners) == 0 {
		return nil
	}
	subnets := &gocf.StringListExpr{Literal: []*gocf.StringExpr{}}
	for _, item := range e.Provider.Subnets {
		subnets.Literal = append(subnets.Literal, gocf.String(*item))
	}
	securityGroups := []gocf.Stringable{gocf.String(e.Provider.SecurityGroup), gocf.Ref("ELBSecurityGroup").String()}
	elb := &gocf.ElasticLoadBalancingLoadBalancer{
		LoadBalancerName: gocf.String(s.Name),
		Subnets:          subnets,
		SecurityGroups:   securityGroups,
		HealthCheck: &gocf.ElasticLoadBalancingHealthCheck{
			HealthyThreshold:   gocf.String("2"),
			Interval:           gocf.String("15"),
			Target:             gocf.String(fmt.Sprintf("TCP:%s", healthCheckPort)),
			Timeout:            gocf.String("10"),
			UnhealthyThreshold: gocf.String("2"),
		},
	}
	elb.Listeners = &listeners
	t.Outputs["elbURL"] = &gocf.Output{
		Description: "The URL of the stack",
		Value:       gocf.GetAtt("ELB", "DNSName"),
	}
	t.Outputs["elbName"] = &gocf.Output{
		Description: "Load balancer name",
		Value:       gocf.Ref("ELB"),
	}
	t.AddResource("ELB", elb)
	return nil
}

func loadIAM(t *gocf.Template, s *service.Service, e *environment.Environment) {
	policy := gocf.IAMPolicies{
		PolicyName: gocf.String(s.Name),
		PolicyDocument: &map[string]interface{}{
			"Version": "2012-10-17",
			"Statement": map[string]interface{}{
				"Effect":   "Allow",
				"Action":   []string{"logs:CreateLogStream", "logs:PutLogEvents"},
				"Resource": fmt.Sprintf("arn:aws:logs:%s:*:log-group:i2kit-%s:log-stream:i-*", e.Provider.Region, s.Name),
			},
		},
	}
	role := &gocf.IAMRole{
		AssumeRolePolicyDocument: &map[string]interface{}{
			"Statement": map[string]interface{}{
				"Effect":    "Allow",
				"Principal": map[string]interface{}{"Service": []string{"ec2.amazonaws.com"}},
				"Action":    []string{"sts:AssumeRole"},
			},
		},
		Path:     gocf.String("/"),
		Policies: &gocf.IAMPoliciesList{policy},
	}
	t.AddResource("Role", role)
	instanceProfile := &gocf.IAMInstanceProfile{
		Path:  gocf.String("/"),
		Roles: gocf.StringList(gocf.Ref("Role")),
	}
	t.AddResource("InstanceProfile", instanceProfile)
}

func loadLogGroup(t *gocf.Template, s *service.Service, e *environment.Environment) {
	logGroup := &gocf.LogsLogGroup{
		LogGroupName:    gocf.String(fmt.Sprintf("i2kit-%s", s.Name)),
		RetentionInDays: gocf.Integer(30),
	}
	t.AddResource("LogGroup", logGroup)
}

func loadRoute53(t *gocf.Template, s *service.Service, e *environment.Environment) {
	recordName := fmt.Sprintf("%s.%s", s.Name, e.Provider.HostedZone)
	recordSet := &gocf.Route53RecordSet{
		HostedZoneName:  gocf.String(e.Provider.HostedZone),
		Name:            gocf.String(recordName),
		Type:            gocf.String("CNAME"),
		TTL:             gocf.String("900"),
		ResourceRecords: gocf.StringList(gocf.GetAtt("ELB", "DNSName")),
	}
	t.AddResource("DNSRecord", recordSet)
}
<commit_msg>Switch to ecs base image (#59)<commit_after>package aws

import (
	"encoding/base64"
	"encoding/json"
	"fmt"
	"strconv"

	"github.com/aws/aws-sdk-go/aws"
	gocf "github.com/crewjam/go-cloudformation"
	"github.com/google/uuid"
	"github.com/pchico83/i2kit/cli/providers/aws/ec2"
	"github.com/pchico83/i2kit/cli/schemas/compose"
	"github.com/pchico83/i2kit/cli/schemas/environment"
	"github.com/pchico83/i2kit/cli/schemas/service"
)

var amisPerRegion = map[string]string{
	"us-east-2":      "ami-1b90a67e",
	"us-east-1":      "ami-cb17d8b6",
	"us-west-2":      "ami-05b5277d",
	"us-west-1":      "ami-9cbbaffc",
	"eu-west-3":      "ami-914afcec",
	"eu-west-2":      "ami-a48d6bc3",
	"eu-west-1":      "ami-bfb5fec6",
	"eu-central-1":   "ami-ac055447",
	"ap-northeast-2": "ami-ba74d8d4",
	"ap-northeast-1": "ami-5add893c",
	"ap-southeast-2": "ami-4cc5072e",
	"ap-southeast-1": "ami-acbcefd0",
	"ca-central-1":   "ami-a535b2c1",
	"ap-south-1":     "ami-2149114e",
	"sa-east-1":      "ami-d3bce9bf",
}

// Translate an i2kit service to a AWS CloudFormation template
func Translate(s *service.Service, e *environment.Environment, config *aws.Config) (string, error) {
	ami, ok := amisPerRegion[e.Provider.Region]
	if !ok {
		return "", fmt.Errorf("Region'%s' is not supported", e.Provider.Region)
	}
	t := gocf.NewTemplate()
	err := loadASG(t, s, e, ami, config)
	if err != nil {
		return "", err
	}
	err = loadELB(t, s, e)
	if err != nil {
		return "", err
	}
	loadIAM(t, s, e)
	loadLogGroup(t, s, e)
	loadRoute53(t, s, e)
	marshalledTemplate := []byte("")
	if marshalledTemplate, err = json.Marshal(t); err != nil {
		return "", err
	}
	return string(marshalledTemplate), nil
}

func loadASG(t *gocf.Template, s *service.Service, e *environment.Environment, ami string, config *aws.Config) error {
	vpc, err := ec2.GetVPC(e, config)
	if err != nil {
		return err
	}
	domain := e.Provider.HostedZone[:len(e.Provider.HostedZone)-1]
	encodedCompose, err := compose.Create(s, domain)
	if err != nil {
		return err
	}
	subnets := &gocf.StringListExpr{Literal: []*gocf.StringExpr{}}
	for _, item := range e.Provider.Subnets {
		subnets.Literal = append(subnets.Literal, gocf.String(*item))
	}
	replicas := strconv.Itoa(s.Replicas)
	asg := &gocf.AutoScalingAutoScalingGroup{
		HealthCheckGracePeriod:  gocf.Integer(15),
		LaunchConfigurationName: gocf.Ref("LaunchConfig").String(),
		LoadBalancerNames:       gocf.StringList(gocf.Ref("ELB")),
		MaxSize:                 gocf.String(replicas),
		MinSize:                 gocf.String(replicas),
		VPCZoneIdentifier:       subnets,
	}
	creationPolicy := &gocf.CreationPolicy{
		ResourceSignal: &gocf.CreationPolicyResourceSignal{
			Count:   gocf.Integer(int64(s.Replicas)),
			Timeout: gocf.String("PT15M"),
		},
	}
	updatePolicy := &gocf.UpdatePolicy{
		AutoScalingRollingUpdate: &gocf.UpdatePolicyAutoScalingRollingUpdate{
			MaxBatchSize:          gocf.Integer(1),
			PauseTime:             gocf.String("PT15M"),
			WaitOnResourceSignals: gocf.Bool(true),
			SuspendProcesses: gocf.StringList(
				gocf.String("HealthCheck"),
				gocf.String("ReplaceUnhealthy"),
				gocf.String("AZRebalance"),
				gocf.String("AlarmNotification"),
				gocf.String("ScheduledActions"),
			),
		},
	}
	asgResource := &gocf.Resource{
		Properties:     asg,
		CreationPolicy: creationPolicy,
		UpdatePolicy:   updatePolicy,
	}
	t.Resources["ASG"] = asgResource

	instanceIngressRules := gocf.EC2SecurityGroupRuleList{}
	loadbalancerIngressRules := gocf.EC2SecurityGroupRuleList{}
	for _, container := range s.Containers {
		for _, port := range container.Ports {
			instancePortNumber, _ := strconv.ParseInt(port.InstancePort, 10, 64)
			instanceIngressRules = append(
				instanceIngressRules,
				gocf.EC2SecurityGroupRule{
					SourceSecurityGroupIdXXSecurityGroupIngressXOnlyX: gocf.Ref("ELBSecurityGroup").String(),
					IpProtocol: gocf.String("tcp"),
					FromPort:   gocf.Integer(instancePortNumber),
					ToPort:     gocf.Integer(instancePortNumber),
				})
			portNumber, _ := strconv.ParseInt(port.Port, 10, 64)
			loadbalancerIngressRules = append(
				loadbalancerIngressRules,
				gocf.EC2SecurityGroupRule{
					CidrIp:     gocf.String("0.0.0.0/0"),
					IpProtocol: gocf.String("tcp"),
					FromPort:   gocf.Integer(portNumber),
					ToPort:     gocf.Integer(portNumber),
				})
		}
	}
	securityGroups := []gocf.Stringable{gocf.String(e.Provider.SecurityGroup)}
	if len(instanceIngressRules) > 0 {
		securityGroup := &gocf.EC2SecurityGroup{
			GroupDescription:     gocf.String(fmt.Sprintf("Instance Security Group for %s.%s", s.Name, e.Provider.Name)),
			SecurityGroupIngress: &instanceIngressRules,
			VpcId:                gocf.String(vpc),
		}
		t.AddResource("InstanceSecurityGroup", securityGroup)
		securityGroups = append(securityGroups, gocf.Ref("InstanceSecurityGroup").String())
		securityGroup = &gocf.EC2SecurityGroup{
			GroupDescription:     gocf.String(fmt.Sprintf("ELB Security Group for %s.%s", s.Name, e.Provider.Name)),
			SecurityGroupIngress: &loadbalancerIngressRules,
			VpcId:                gocf.String(vpc),
		}
		t.AddResource("ELBSecurityGroup", securityGroup)
	}

	launchConfig := &gocf.AutoScalingLaunchConfiguration{
		ImageId:            gocf.String(ami),
		InstanceType:       gocf.String(s.GetSize(e)),
		KeyName:            gocf.String(e.Provider.Keypair),
		SecurityGroups:     securityGroups,
		IamInstanceProfile: gocf.Ref("InstanceProfile").String(),
		UserData:           gocf.String(userData(s.Name, encodedCompose, e)),
	}
	t.AddResource("LaunchConfig", launchConfig)
	return nil
}

func userData(containerName, encodedCompose string, e *environment.Environment) string {
	uniqueOperationID := uuid.New().String()
	value := fmt.Sprintf(
		`#!/bin/bash

set -e
INSTANCE_ID=$(curl http://169.254.169.254/latest/meta-data/instance-id)
sudo docker run \
	--name %s \
	-e COMPOSE=%s \
	-e CONFIG=%s \
	-e UNIQUE_OPERATION_ID=%s \
	-e STACK=%s \
	-e REGION=%s \
	-v /var/run/docker.sock:/var/run/docker.sock \
	--log-driver=awslogs \
	--log-opt awslogs-region=%s \
	--log-opt awslogs-group=i2kit-%s \
	--log-opt tag=$INSTANCE_ID \
	riberaproject/agent`,
		containerName,
		encodedCompose,
		e.B64DockerConfig(),
		uniqueOperationID,
		containerName,
		e.Provider.Region,
		e.Provider.Region,
		containerName,
	)
	return base64.StdEncoding.EncodeToString([]byte(value))
}

func loadELB(t *gocf.Template, s *service.Service, e *environment.Environment) error {
	healthCheckPort := ""
	listeners := gocf.ElasticLoadBalancingListenerList{}
	for _, container := range s.Containers {
		for _, port := range container.Ports {
			if healthCheckPort == "" {
				healthCheckPort = port.InstancePort
			}
			certificate := port.Certificate
			if certificate == "" && (port.Protocol == "HTTPS" || port.Protocol == "SSL") {
				if e.Provider.Certificate == "" {
					return fmt.Errorf("Port '%s:%s' requires a certificate", port.Protocol, port.Port)
				}
				certificate = e.Provider.Certificate
			}
			listeners = append(listeners, gocf.ElasticLoadBalancingListener{
				InstancePort:     gocf.String(port.InstancePort),
				InstanceProtocol: gocf.String(port.InstanceProtocol),
				LoadBalancerPort: gocf.String(port.Port),
				Protocol:         gocf.String(port.Protocol),
				SSLCertificateId: gocf.String(certificate),
			})
		}
	}
	if len(listeners) == 0 {
		return nil
	}
	subnets := &gocf.StringListExpr{Literal: []*gocf.StringExpr{}}
	for _, item := range e.Provider.Subnets {
		subnets.Literal = append(subnets.Literal, gocf.String(*item))
	}
	securityGroups := []gocf.Stringable{gocf.String(e.Provider.SecurityGroup), gocf.Ref("ELBSecurityGroup").String()}
	elb := &gocf.ElasticLoadBalancingLoadBalancer{
		LoadBalancerName: gocf.String(s.Name),
		Subnets:          subnets,
		SecurityGroups:   securityGroups,
		HealthCheck: &gocf.ElasticLoadBalancingHealthCheck{
			HealthyThreshold:   gocf.String("2"),
			Interval:           gocf.String("15"),
			Target:             gocf.String(fmt.Sprintf("TCP:%s", healthCheckPort)),
			Timeout:            gocf.String("10"),
			UnhealthyThreshold: gocf.String("2"),
		},
	}
	elb.Listeners = &listeners
	t.Outputs["elbURL"] = &gocf.Output{
		Description: "The URL of the stack",
		Value:       gocf.GetAtt("ELB", "DNSName"),
	}
	t.Outputs["elbName"] = &gocf.Output{
		Description: "Load balancer name",
		Value:       gocf.Ref("ELB"),
	}
	t.AddResource("ELB", elb)
	return nil
}

func loadIAM(t *gocf.Template, s *service.Service, e *environment.Environment) {
	policy := gocf.IAMPolicies{
		PolicyName: gocf.String(s.Name),
		PolicyDocument: &map[string]interface{}{
			"Version": "2012-10-17",
			"Statement": map[string]interface{}{
				"Effect":   "Allow",
				"Action":   []string{"logs:CreateLogStream", "logs:PutLogEvents"},
				"Resource": fmt.Sprintf("arn:aws:logs:%s:*:log-group:i2kit-%s:log-stream:i-*", e.Provider.Region, s.Name),
			},
		},
	}
	role := &gocf.IAMRole{
		AssumeRolePolicyDocument: &map[string]interface{}{
			"Statement": map[string]interface{}{
				"Effect":    "Allow",
				"Principal": map[string]interface{}{"Service": []string{"ec2.amazonaws.com"}},
				"Action":    []string{"sts:AssumeRole"},
			},
		},
		Path:     gocf.String("/"),
		Policies: &gocf.IAMPoliciesList{policy},
	}
	t.AddResource("Role", role)
	instanceProfile := &gocf.IAMInstanceProfile{
		Path:  gocf.String("/"),
		Roles: gocf.StringList(gocf.Ref("Role")),
	}
	t.AddResource("InstanceProfile", instanceProfile)
}

func loadLogGroup(t *gocf.Template, s *service.Service, e *environment.Environment) {
	logGroup := &gocf.LogsLogGroup{
		LogGroupName:    gocf.String(fmt.Sprintf("i2kit-%s", s.Name)),
		RetentionInDays: gocf.Integer(30),
	}
	t.AddResource("LogGroup", logGroup)
}

func loadRoute53(t *gocf.Template, s *service.Service, e *environment.Environment) {
	recordName := fmt.Sprintf("%s.%s", s.Name, e.Provider.HostedZone)
	recordSet := &gocf.Route53RecordSet{
		HostedZoneName:  gocf.String(e.Provider.HostedZone),
		Name:            gocf.String(recordName),
		Type:            gocf.String("CNAME"),
		TTL:             gocf.String("900"),
		ResourceRecords: gocf.StringList(gocf.GetAtt("ELB", "DNSName")),
	}
	t.AddResource("DNSRecord", recordSet)
}
<|endoftext|>


<commit_before>package utils

import (
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"os"
	"strings"
)

var configFile string
var hc = http.Client{}

func Get(endpoint string) []byte {
	response, err := http.Get(endpoint)
	if err != nil {
		fmt.Print(err.Error())
		os.Exit(1)
	}

	// Response
	responseData, err := ioutil.ReadAll(response.Body)
	if err != nil {
		log.Fatal(err)
	}
	return responseData
}

func Post(endpoint string, form url.Values) ([]byte, error) {

	req, err := http.NewRequest("POST", endpoint, strings.NewReader(form.Encode()))
	req.PostForm = form
	req.Header.Add("Content-Type", "application/x-www-form-urlencoded")

	// Fetch
	response, err := hc.Do(req)

	responseData, err := ioutil.ReadAll(response.Body)
	if err != nil {
		log.Fatal(err)
	}

	return responseData, err
}
<commit_msg>#21 Added Error struct for future error handling efforts<commit_after>package utils

import (
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"os"
	"strings"
)

var configFile string
var hc = http.Client{}

// Error .. an error returned from the API
type Error struct {
	Message  string `json:"message"`
	Code     int    `json:"code"`
	MoreInfo string `json:"moreInfo"`
}

// ErrorsList .. a list of errors returned from the API
type ErrorsList struct {
	errorList []Error
}

// Get .. issues an HTTP GET request
func Get(endpoint string) []byte {
	response, err := http.Get(endpoint)
	if err != nil {
		fmt.Print(err.Error())
		os.Exit(1)
	}

	// Response
	responseData, err := ioutil.ReadAll(response.Body)
	if err != nil {
		log.Fatal(err)
	}
	return responseData
}

// Post .. issues an HTTP POST request
func Post(endpoint string, form url.Values) ([]byte, error) {

	req, err := http.NewRequest("POST", endpoint, strings.NewReader(form.Encode()))
	req.PostForm = form
	req.Header.Add("Content-Type", "application/x-www-form-urlencoded")

	// Fetch
	response, err := hc.Do(req)
	if err != nil {
		log.Fatal(err)
	}

	responseData, err := ioutil.ReadAll(response.Body)
	if err != nil {
		log.Fatal(err)
	}

	return responseData, err
}
<|endoftext|>


<commit_before>package flotilla

import (
	"net/http"
	"os"
	"path/filepath"
)

func engineStaticFile(dirs []string, requested string) (file http.File, exists bool) {
	file, exists = file, false
	for _, dir := range dirs {
		filepath.Walk(dir, func(path string, _ os.FileInfo, _ error) (err error) {
			if filepath.Base(path) == requested {
				f, _ := os.Open(path)
				file, exists = f, true
				//fmt.Printf("\n{path: %+v, base: %+v, file: %+v, requested: %+v, exists: %t}\n", path, filepath.Base(path), file, requested, exists)
			}
			return err
		})

	}
	return file, exists
}

func engineAssetFile(requested string, engine *Engine) (file http.File, exists bool) {
	file, exists = file, false
	f, err := engine.Assets.Get(requested)
	if err == nil {
		file, exists = f, true
		//fmt.Printf("\n\n{requested: %+v, file: %+v, exists: %t}\n", requested, file, exists)
	}
	return file, exists
}

func handleStatic(c *Ctx) {
	var f http.File
	var exists bool = false
	requested := filepath.Base(c.Request.URL.Path)
	f, exists = engineStaticFile(c.Engine.StaticDirs(), requested)
	if !exists {
		f, exists = engineAssetFile(requested, c.Engine)
	}
	if exists {
		c.ServeFile(f)
	} else {
		c.Abort(404)
	}
}
<commit_msg>lists in conf; StaticDirs() method to bridge conf and otherwise set; rewriting static handler to be less a pain in the ass<commit_after>package flotilla

import (
	"os"
	"path/filepath"
)

func engineStaticFile(requested string, c *Ctx) (exists bool) {
	exists = false
	for _, dir := range c.Engine.StaticDirs() {
		filepath.Walk(dir, func(path string, _ os.FileInfo, _ error) (err error) {
			if filepath.Base(path) == requested {
				f, _ := os.Open(path)
				c.ServeFile(f)
				exists = true
				//fmt.Printf("\n{path: %+v, base: %+v, file: %+v, requested: %+v, exists: %t}\n", path, filepath.Base(path), file, requested, exists)
			}
			return err
		})

	}
	return exists
}

func engineAssetFile(requested string, c *Ctx) (exists bool) {
	exists = false
	f, err := c.Engine.Assets.Get(requested)
	if err == nil {
		c.ServeFile(f)
		exists = true
		//fmt.Printf("\n\n{requested: %+v, file: %+v, exists: %t}\n", requested, file, exists)
	}
	return exists
}

func staticfile() {}

func handleStatic(c *Ctx) {
	var exists bool = false
	requested := filepath.Base(c.Request.URL.Path)
	exists = engineStaticFile(requested, c)
	if !exists {
		exists = engineAssetFile(requested, c)
	}
	if !exists {
		c.Abort(404)
	}
}
<|endoftext|>


<commit_before>// statspipe is a metrics pipeline
package main

import (
	"bufio"
	"bytes"
	"flag"
	"fmt"
	"io"
	"log"
	"math"
	"net"
	//"regexp"
	"sort"
	"strconv"
	"sync"
	"sync/atomic"
	"time"

	//"github.com/davecgh/go-spew/spew"
	"github.com/davecheney/profile"
)

const FlushInterval = time.Duration(10 * time.Second)
const BufSize = 8192

// Metric Types
const Counter = "c"
const Gauge = "g"
const Timer = "ms"

//-----------------------------------------------------------------------------

// Command line flags
var (
	listen   = flag.String("listen", ":8125", "Listener address")
	graphite = flag.String("graphite", "localhost:2003", "Graphite server address")

	// Profiling
	cpuprofile   = flag.Bool("cpuprofile", false, "Enable CPU profiling")
	memprofile   = flag.Bool("memprofile", false, "Enable memory profiling")
	blockprofile = flag.Bool("blockprofile", false, "Enable block profiling")

	debug = flag.Bool("debug", false, "Enable debug mode")
)

//-----------------------------------------------------------------------------
// Data structures

// Metric is a numeric data point
type Metric struct {
	Bucket string
	Value  interface{}
	Type   string
}

// Metrics should be in statsd format. Metric names may not have spaces.
//
//     <metric_name>:<metric_value>|<metric_type>|@<sample_rate>
//
// Note: The sample rate is optional
// var statsPattern = regexp.MustCompile(`[\w\.]+:-?\d+\|(?:c|ms|g)(?:\|\@[\d\.]+)?`)

// In is a channel for processing metrics
var In = make(chan *Metric)

// counters holds all of the counter metrics
var counters = struct {
	sync.RWMutex
	m map[string]int64
}{m: make(map[string]int64)}

// gauges holds all of the gauge metrics
var gauges = struct {
	sync.RWMutex
	m map[string]uint64
}{m: make(map[string]uint64)}

// Timers is a list of uints
type Timers []uint64

// timers holds all of the timer metrics
var timers = struct {
	sync.RWMutex
	m map[string]Timers
}{m: make(map[string]Timers)}

// Internal metrics
type Stats struct {
	IngressRate     int64
	IngressMetrics  int64
	IngressCounters int64
	IngressGauges   int64
	IngressTimers   int64
}

var stats = &Stats{}

// TODO: move this to command line option
var Percentiles = []int{5, 95}

//-----------------------------------------------------------------------------

// Implement the sort interface for Timers
func (t Timers) Len() int           { return len(t) }
func (t Timers) Swap(i, j int)      { t[i], t[j] = t[j], t[i] }
func (t Timers) Less(i, j int) bool { return t[i] < t[j] }

//-----------------------------------------------------------------------------

// ListenUDP creates a UDP listener
func ListenUDP(addr string) error {
	var buf = make([]byte, 1024)
	ln, err := net.ResolveUDPAddr("udp", addr)

	if err != nil {
		return err
	}

	sock, err := net.ListenUDP("udp", ln)

	if err != nil {
		return err
	}

	log.Printf("Listening on UDP %s\n", ln)

	for {
		n, raddr, err := sock.ReadFromUDP(buf[:])

		if err != nil {
			// TODO: handle error
			continue
		}

		if *debug {
			log.Printf("DEBUG: Received UDP message: bytes=%d client=%s",
				n, raddr)
		}

		go handleUdpMessage(buf)
	}
}

func handleUdpMessage(buf []byte) {
	tokens := bytes.Split(buf, []byte("\n"))

	for _, token := range tokens {
		handleMessage(token)
	}
}

// ListenTCP creates a TCP listener
func ListenTCP(addr string) error {
	l, err := net.Listen("tcp", addr)

	if err != nil {
		return err
	}

	defer l.Close()
	log.Printf("Listening on TCP %s\n", l.Addr())

	for {
		conn, err := l.Accept()

		if err != nil {
			// TODO: handle error
			continue
		}

		go handleConnection(conn)
	}
}

// handleConnection handles a single client connection
func handleConnection(conn net.Conn) {
	defer conn.Close()
	r := bufio.NewReader(conn)

	// Incoming metrics should be separated by a newline
	for {
		line, err := r.ReadBytes('\n')

		if err != nil {
			if err == io.EOF {
				break
			} else {
				// TODO: handle error
			}
		}

		if *debug {
			log.Printf("DEBUG: Received TCP message: bytes=%d client=%s",
				len(line), conn.RemoteAddr())
		}

		handleMessage(line)
	}
}

// Handle an event message
func handleMessage(buf []byte) {
	// According to the statsd protocol, metrics should be separated by a
	// newline. This parser isn't quite as strict since it may be receiving
	// metrics from clients that aren't proper statsd clients (e.g. syslog).
	// Instead, find metrics by looking for tokens between spaces and then
	// use parseMetric to create an actual metric that can be aggregated.
	tokens := bytes.Split(buf, []byte(" "))

	for _, token := range tokens {
		// metrics must have a : and | at a minimum
		if bytes.Index(token, []byte(":")) == -1 ||
			bytes.Index(token, []byte("|")) == -1 {
			continue
		}

		metric, err := parseMetric(token)

		if err != nil {
			log.Printf("ERROR: Unable to parse metric %s: %s",
				token, err)
			continue
		}

		// Send metric off for processing
		In <- metric
	}
}

// parseMetric parses a raw metric into a Metric struct
func parseMetric(b []byte) (*Metric, error) {
	i := bytes.Index(b, []byte(":"))
	j := bytes.Index(b, []byte("|"))
	k := bytes.Index(b, []byte("@"))
	v := b[i+1 : j]

	// End position of the metric type is the end of the byte slice
	// if no sample was sent.
	tEnd := len(b)
	var sampleRate float64 = 1

	// Indicates that a sample rate was sent as part of the metric
	if k > -1 {
		tEnd = k - 1 // Use -1 because of the | before the @
		sr := b[(k + 1):len(b)]
		var err error
		sampleRate, err = strconv.ParseFloat(string(sr), 64)

		if err != nil {
			return nil, err
		}
	}

	m := &Metric{
		Bucket: string(b[0:i]),
		Type:   string(b[j+1 : tEnd]),
	}

	switch m.Type {
	case Counter:
		val, err := strconv.ParseInt(string(v), 10, 64)

		if err != nil {
			return nil, err
		}

		m.Value = int64(float64(val) / sampleRate)

	case Gauge, Timer:
		val, err := strconv.ParseUint(string(v), 10, 64)

		if err != nil {
			return nil, err
		}

		m.Value = val

	default:
		err := fmt.Errorf("unable to create metric for type %q", m.Type)

		return nil, err
	}

	return m, nil
}

// processMetrics updates new metrics and flushes aggregates to Graphite
func processMetrics() {
	ticker := time.NewTicker(FlushInterval)

	for {
		select {
		case <-ticker.C:
			flushMetrics()
		case m := <-In:
			if *debug {
				log.Printf("DEBUG: Received metric for processing: %+v", m)
			}

			atomic.AddInt64(&stats.IngressMetrics, 1)

			switch m.Type {
			case Counter:
				counters.Lock()
				counters.m[m.Bucket] += m.Value.(int64)
				counters.Unlock()
				atomic.AddInt64(&stats.IngressCounters, 1)

			case Gauge:
				gauges.Lock()
				gauges.m[m.Bucket] = m.Value.(uint64)
				gauges.Unlock()
				atomic.AddInt64(&stats.IngressGauges, 1)

			case Timer:
				timers.Lock()
				_, ok := timers.m[m.Bucket]

				if !ok {
					var t Timers
					timers.m[m.Bucket] = t
				}

				timers.m[m.Bucket] = append(timers.m[m.Bucket], m.Value.(uint64))
				timers.Unlock()
				atomic.AddInt64(&stats.IngressTimers, 1)

			default:
				if *debug {
					log.Printf("DEBUG: Unable to process unknown metric type %q", m.Type)
				}

			}

			if *debug {
				log.Printf("DEBUG: Finished processing metric: %+v", m)
			}
		}
	}
}

// flushMetrics sends metrics to Graphite
func flushMetrics() {
	var buf bytes.Buffer
	now := time.Now().Unix()

	log.Printf("%+v", stats)

	// Build buffer of stats
	flushCounters(&buf, now)
	flushGauges(&buf, now)
	flushTimers(&buf, now)
	flushInternalStats(&buf, now)

	// Send metrics to Graphite
	sendGraphite(&buf)
}

// flushInternalStats writes the internal stats to the buffer
func flushInternalStats(buf *bytes.Buffer, now int64) {
	//fmt.Fprintf(buf, "statsd.metrics.per_second %d %d\n", v, now)
	fmt.Fprintf(buf, "statsd.metrics.count %d %d\n",
		atomic.LoadInt64(&stats.IngressMetrics), now)
	fmt.Fprintf(buf, "statsd.counters.count %d %d\n",
		atomic.LoadInt64(&stats.IngressCounters), now)
	fmt.Fprintf(buf, "statsd.gauges.count %d %d\n",
		atomic.LoadInt64(&stats.IngressGauges), now)
	fmt.Fprintf(buf, "statsd.timers.count %d %d\n",
		atomic.LoadInt64(&stats.IngressTimers), now)

	// Clear internal metrics
	atomic.StoreInt64(&stats.IngressMetrics, 0)
	atomic.StoreInt64(&stats.IngressCounters, 0)
	atomic.StoreInt64(&stats.IngressGauges, 0)
	atomic.StoreInt64(&stats.IngressTimers, 0)
}

// flushCounters writes the counters to the buffer
func flushCounters(buf *bytes.Buffer, now int64) {
	counters.Lock()
	defer counters.Unlock()

	for k, v := range counters.m {
		fmt.Fprintf(buf, "%s %d %d\n", k, v, now)
		delete(counters.m, k)
	}
}

// flushGauges writes the gauges to the buffer
func flushGauges(buf *bytes.Buffer, now int64) {
	gauges.Lock()
	defer gauges.Unlock()

	for k, v := range gauges.m {
		fmt.Fprintf(buf, "%s %d %d\n", k, v, now)
		delete(gauges.m, k)
	}
}

// flushTimers writes the timers and aggregate statistics to the buffer
func flushTimers(buf *bytes.Buffer, now int64) {
	timers.RLock()
	defer timers.RUnlock()
	var n int64

	for k, t := range timers.m {
		count := len(t)

		// Skip processing if there are no timer values
		if count < 1 {
			break
		}

		var sum uint64

		for _, v := range t {
			sum += v
			n++
		}

		// Linear average (mean)
		mean := float64(sum) / float64(count)

		// Min and Max
		sort.Sort(t)
		min := t[0]
		max := t[len(t)-1]

		// Write out all derived stats
		fmt.Fprintf(buf, "%s.count %d %d\n", k, count, now)
		fmt.Fprintf(buf, "%s.mean %f %d\n", k, mean, now)
		fmt.Fprintf(buf, "%s.lower %d %d\n", k, min, now)
		fmt.Fprintf(buf, "%s.upper %d %d\n", k, max, now)

		// Calculate and write out percentiles
		for _, pct := range Percentiles {
			p := perc(t, pct)
			fmt.Fprintf(buf, "%s.perc%d %f %d\n", k, pct, p, now)
		}

		delete(timers.m, k)
	}
}

// percentile calculates Nth percentile of a list of values
func perc(values []uint64, pct int) float64 {
	p := float64(pct) / float64(100)
	n := float64(len(values))
	i := math.Ceil(p*n) - 1

	return float64(values[int(i)])
}

// sendGraphite sends metrics to graphite
func sendGraphite(buf *bytes.Buffer) {
	log.Printf("Sending metrics to Graphite: bytes=%d host=%s",
		buf.Len(), *graphite)
	t0 := time.Now()

	conn, err := net.Dial("tcp", *graphite)

	if err != nil {
		log.Printf("ERROR: Unable to connect to graphite: %s", err)
		return
	}

	w := bufio.NewWriter(conn)
	n, err := buf.WriteTo(w)

	if err != nil {
		log.Printf("ERROR: Unable to write to graphite: %s", err)
	}

	w.Flush()
	conn.Close()

	log.Printf("Finished sending metrics to Graphite: bytes=%d host=%s duration=%s",
		n, conn.RemoteAddr(), time.Now().Sub(t0))
}

//-----------------------------------------------------------------------------

func main() {
	flag.Parse()

	// Profiling
	cfg := profile.Config{
		CPUProfile:   *cpuprofile,
		MemProfile:   *memprofile,
		BlockProfile: *blockprofile,
		ProfilePath:  ".",
	}

	p := profile.Start(&cfg)
	defer p.Stop()

	// Process metrics as they arrive
	go processMetrics()

	// Setup listeners
	var wg sync.WaitGroup
	wg.Add(2)

	go func() {
		defer wg.Done()
		log.Fatal(ListenUDP(*listen))
	}()

	go func() {
		defer wg.Done()
		log.Fatal(ListenTCP(*listen))
	}()

	wg.Wait()
}
<commit_msg>Additional debug logging<commit_after>// statspipe is a metrics pipeline
package main

import (
	"bufio"
	"bytes"
	"flag"
	"fmt"
	"io"
	"log"
	"math"
	"net"
	//"regexp"
	"sort"
	"strconv"
	"sync"
	"sync/atomic"
	"time"

	//"github.com/davecgh/go-spew/spew"
	"github.com/davecheney/profile"
)

//-----------------------------------------------------------------------------

const FlushInterval = time.Duration(10 * time.Second)
const BufSize = 8192

// Metric Types
const Counter = "c"
const Gauge = "g"
const Timer = "ms"

//-----------------------------------------------------------------------------

// Command line flags
var (
	listen   = flag.String("listen", ":8125", "Listener address")
	graphite = flag.String("graphite", "localhost:2003", "Graphite server address")

	// Profiling
	cpuprofile   = flag.Bool("cpuprofile", false, "Enable CPU profiling")
	memprofile   = flag.Bool("memprofile", false, "Enable memory profiling")
	blockprofile = flag.Bool("blockprofile", false, "Enable block profiling")

	debug = flag.Bool("debug", false, "Enable debug mode")
)

//-----------------------------------------------------------------------------
// Data structures

// Metric is a numeric data point
type Metric struct {
	Bucket string
	Value  interface{}
	Type   string
}

// Metrics should be in statsd format. Metric names may not have spaces.
//
//     <metric_name>:<metric_value>|<metric_type>|@<sample_rate>
//
// Note: The sample rate is optional
// var statsPattern = regexp.MustCompile(`[\w\.]+:-?\d+\|(?:c|ms|g)(?:\|\@[\d\.]+)?`)

// In is a channel for processing metrics
var In = make(chan *Metric)

// counters holds all of the counter metrics
var counters = struct {
	sync.RWMutex
	m map[string]int64
}{m: make(map[string]int64)}

// gauges holds all of the gauge metrics
var gauges = struct {
	sync.RWMutex
	m map[string]uint64
}{m: make(map[string]uint64)}

// Timers is a list of uints
type Timers []uint64

// timers holds all of the timer metrics
var timers = struct {
	sync.RWMutex
	m map[string]Timers
}{m: make(map[string]Timers)}

// Internal metrics
type Stats struct {
	IngressRate     int64
	IngressMetrics  int64
	IngressCounters int64
	IngressGauges   int64
	IngressTimers   int64
}

var stats = &Stats{}

// TODO: move this to command line option
var Percentiles = []int{5, 95}

//-----------------------------------------------------------------------------

// Implement the sort interface for Timers
func (t Timers) Len() int           { return len(t) }
func (t Timers) Swap(i, j int)      { t[i], t[j] = t[j], t[i] }
func (t Timers) Less(i, j int) bool { return t[i] < t[j] }

//-----------------------------------------------------------------------------

// ListenUDP creates a UDP listener
func ListenUDP(addr string) error {
	var buf = make([]byte, 1024)
	ln, err := net.ResolveUDPAddr("udp", addr)

	if err != nil {
		return err
	}

	sock, err := net.ListenUDP("udp", ln)

	if err != nil {
		return err
	}

	log.Printf("Listening on UDP %s\n", ln)

	for {
		n, raddr, err := sock.ReadFromUDP(buf[:])

		if err != nil {
			// TODO: handle error
			continue
		}

		if *debug {
			log.Printf("DEBUG: Received UDP message: bytes=%d client=%s",
				n, raddr)
		}

		go handleUdpMessage(buf)
	}
}

func handleUdpMessage(buf []byte) {
	tokens := bytes.Split(buf, []byte("\n"))

	for _, token := range tokens {
		handleMessage(token)
	}
}

// ListenTCP creates a TCP listener
func ListenTCP(addr string) error {
	l, err := net.Listen("tcp", addr)

	if err != nil {
		return err
	}

	defer l.Close()
	log.Printf("Listening on TCP %s\n", l.Addr())

	for {
		conn, err := l.Accept()

		if err != nil {
			// TODO: handle error
			continue
		}

		go handleConnection(conn)
	}
}

// handleConnection handles a single client connection
func handleConnection(conn net.Conn) {
	defer conn.Close()
	r := bufio.NewReader(conn)

	// Incoming metrics should be separated by a newline
	for {
		line, err := r.ReadBytes('\n')

		if err != nil {
			if err == io.EOF {
				break
			} else {
				// TODO: handle error
			}
		}

		if *debug {
			log.Printf("DEBUG: Received TCP message: bytes=%d client=%s",
				len(line), conn.RemoteAddr())
		}

		handleMessage(line)
	}
}

// Handle an event message
func handleMessage(buf []byte) {
	// According to the statsd protocol, metrics should be separated by a
	// newline. This parser isn't quite as strict since it may be receiving
	// metrics from clients that aren't proper statsd clients (e.g. syslog).
	// Instead, find metrics by looking for tokens between spaces and then
	// use parseMetric to create an actual metric that can be aggregated.
	tokens := bytes.Split(buf, []byte(" "))

	for _, token := range tokens {
		// metrics must have a : and | at a minimum
		if bytes.Index(token, []byte(":")) == -1 ||
			bytes.Index(token, []byte("|")) == -1 {
			continue
		}

		if *debug {
			log.Printf("DEBUG: Parsing metric from token: %q", string(token))
		}

		metric, err := parseMetric(token)

		if err != nil {
			log.Printf("ERROR: Unable to parse metric %q: %s",
				token, err)
			continue
		}

		// Send metric off for processing
		In <- metric

		if *debug {
			log.Printf("DEBUG: Queued metric for processing: %+v", metric)
		}
	}
}

// parseMetric parses a raw metric into a Metric struct
func parseMetric(b []byte) (*Metric, error) {
	i := bytes.Index(b, []byte(":"))
	j := bytes.Index(b, []byte("|"))
	k := bytes.Index(b, []byte("@"))
	v := b[i+1 : j]

	// End position of the metric type is the end of the byte slice
	// if no sample was sent.
	tEnd := len(b)
	var sampleRate float64 = 1

	// Indicates that a sample rate was sent as part of the metric
	if k > -1 {
		tEnd = k - 1 // Use -1 because of the | before the @
		sr := b[(k + 1):len(b)]
		var err error
		sampleRate, err = strconv.ParseFloat(string(sr), 64)

		if err != nil {
			return nil, err
		}
	}

	m := &Metric{
		Bucket: string(b[0:i]),
		Type:   string(b[j+1 : tEnd]),
	}

	switch m.Type {
	case Counter:
		val, err := strconv.ParseInt(string(v), 10, 64)

		if err != nil {
			return nil, err
		}

		m.Value = int64(float64(val) / sampleRate)

	case Gauge, Timer:
		val, err := strconv.ParseUint(string(v), 10, 64)

		if err != nil {
			return nil, err
		}

		m.Value = val

	default:
		err := fmt.Errorf("unable to create metric for type %q", m.Type)

		return nil, err
	}

	return m, nil
}

// processMetrics updates new metrics and flushes aggregates to Graphite
func processMetrics() {
	ticker := time.NewTicker(FlushInterval)

	for {
		select {
		case <-ticker.C:
			flushMetrics()
		case m := <-In:
			if *debug {
				log.Printf("DEBUG: Received metric for processing: %+v", m)
			}

			atomic.AddInt64(&stats.IngressMetrics, 1)

			switch m.Type {
			case Counter:
				counters.Lock()
				counters.m[m.Bucket] += m.Value.(int64)
				counters.Unlock()
				atomic.AddInt64(&stats.IngressCounters, 1)

			case Gauge:
				gauges.Lock()
				gauges.m[m.Bucket] = m.Value.(uint64)
				gauges.Unlock()
				atomic.AddInt64(&stats.IngressGauges, 1)

			case Timer:
				timers.Lock()
				_, ok := timers.m[m.Bucket]

				if !ok {
					var t Timers
					timers.m[m.Bucket] = t
				}

				timers.m[m.Bucket] = append(timers.m[m.Bucket], m.Value.(uint64))
				timers.Unlock()
				atomic.AddInt64(&stats.IngressTimers, 1)

			default:
				if *debug {
					log.Printf("DEBUG: Unable to process unknown metric type %q", m.Type)
				}

			}

			if *debug {
				log.Printf("DEBUG: Finished processing metric: %+v", m)
			}
		}
	}
}

// flushMetrics sends metrics to Graphite
func flushMetrics() {
	var buf bytes.Buffer
	now := time.Now().Unix()

	log.Printf("%+v", stats)

	// Build buffer of stats
	flushCounters(&buf, now)
	flushGauges(&buf, now)
	flushTimers(&buf, now)
	flushInternalStats(&buf, now)

	// Send metrics to Graphite
	sendGraphite(&buf)
}

// flushInternalStats writes the internal stats to the buffer
func flushInternalStats(buf *bytes.Buffer, now int64) {
	//fmt.Fprintf(buf, "statsd.metrics.per_second %d %d\n", v, now)
	fmt.Fprintf(buf, "statsd.metrics.count %d %d\n",
		atomic.LoadInt64(&stats.IngressMetrics), now)
	fmt.Fprintf(buf, "statsd.counters.count %d %d\n",
		atomic.LoadInt64(&stats.IngressCounters), now)
	fmt.Fprintf(buf, "statsd.gauges.count %d %d\n",
		atomic.LoadInt64(&stats.IngressGauges), now)
	fmt.Fprintf(buf, "statsd.timers.count %d %d\n",
		atomic.LoadInt64(&stats.IngressTimers), now)

	// Clear internal metrics
	atomic.StoreInt64(&stats.IngressMetrics, 0)
	atomic.StoreInt64(&stats.IngressCounters, 0)
	atomic.StoreInt64(&stats.IngressGauges, 0)
	atomic.StoreInt64(&stats.IngressTimers, 0)
}

// flushCounters writes the counters to the buffer
func flushCounters(buf *bytes.Buffer, now int64) {
	counters.Lock()
	defer counters.Unlock()

	for k, v := range counters.m {
		fmt.Fprintf(buf, "%s %d %d\n", k, v, now)
		delete(counters.m, k)
	}
}

// flushGauges writes the gauges to the buffer
func flushGauges(buf *bytes.Buffer, now int64) {
	gauges.Lock()
	defer gauges.Unlock()

	for k, v := range gauges.m {
		fmt.Fprintf(buf, "%s %d %d\n", k, v, now)
		delete(gauges.m, k)
	}
}

// flushTimers writes the timers and aggregate statistics to the buffer
func flushTimers(buf *bytes.Buffer, now int64) {
	timers.RLock()
	defer timers.RUnlock()
	var n int64

	for k, t := range timers.m {
		count := len(t)

		// Skip processing if there are no timer values
		if count < 1 {
			break
		}

		var sum uint64

		for _, v := range t {
			sum += v
			n++
		}

		// Linear average (mean)
		mean := float64(sum) / float64(count)

		// Min and Max
		sort.Sort(t)
		min := t[0]
		max := t[len(t)-1]

		// Write out all derived stats
		fmt.Fprintf(buf, "%s.count %d %d\n", k, count, now)
		fmt.Fprintf(buf, "%s.mean %f %d\n", k, mean, now)
		fmt.Fprintf(buf, "%s.lower %d %d\n", k, min, now)
		fmt.Fprintf(buf, "%s.upper %d %d\n", k, max, now)

		// Calculate and write out percentiles
		for _, pct := range Percentiles {
			p := perc(t, pct)
			fmt.Fprintf(buf, "%s.perc%d %f %d\n", k, pct, p, now)
		}

		delete(timers.m, k)
	}
}

// percentile calculates Nth percentile of a list of values
func perc(values []uint64, pct int) float64 {
	p := float64(pct) / float64(100)
	n := float64(len(values))
	i := math.Ceil(p*n) - 1

	return float64(values[int(i)])
}

// sendGraphite sends metrics to graphite
func sendGraphite(buf *bytes.Buffer) {
	log.Printf("Sending metrics to Graphite: bytes=%d host=%s",
		buf.Len(), *graphite)
	t0 := time.Now()

	conn, err := net.Dial("tcp", *graphite)

	if err != nil {
		log.Printf("ERROR: Unable to connect to graphite: %s", err)
		return
	}

	w := bufio.NewWriter(conn)
	n, err := buf.WriteTo(w)

	if err != nil {
		log.Printf("ERROR: Unable to write to graphite: %s", err)
	}

	w.Flush()
	conn.Close()

	log.Printf("Finished sending metrics to Graphite: bytes=%d host=%s duration=%s",
		n, conn.RemoteAddr(), time.Now().Sub(t0))
}

//-----------------------------------------------------------------------------

func main() {
	flag.Parse()

	// Profiling
	cfg := profile.Config{
		CPUProfile:   *cpuprofile,
		MemProfile:   *memprofile,
		BlockProfile: *blockprofile,
		ProfilePath:  ".",
	}

	p := profile.Start(&cfg)
	defer p.Stop()

	// Process metrics as they arrive
	go processMetrics()

	// Setup listeners
	var wg sync.WaitGroup
	wg.Add(2)

	go func() {
		defer wg.Done()
		log.Fatal(ListenUDP(*listen))
	}()

	go func() {
		defer wg.Done()
		log.Fatal(ListenTCP(*listen))
	}()

	wg.Wait()
}
<|endoftext|>


<commit_before>// Copyright 2019 The Kubernetes Authors.
// SPDX-License-Identifier: Apache-2.0

package configcobra

import (
	"flag"

	"github.com/spf13/cobra"
	"k8s.io/apimachinery/pkg/api/meta"
	"k8s.io/cli-runtime/pkg/genericclioptions"
	"k8s.io/client-go/discovery"
	"k8s.io/client-go/rest"
	"k8s.io/client-go/tools/clientcmd"
	"k8s.io/kubectl/pkg/cmd/util"
	"sigs.k8s.io/cli-utils/cmd/apply"
	"sigs.k8s.io/cli-utils/cmd/destroy"
	"sigs.k8s.io/cli-utils/cmd/diff"
	"sigs.k8s.io/cli-utils/cmd/initcmd"
	"sigs.k8s.io/cli-utils/cmd/preview"
	"sigs.k8s.io/cli-utils/pkg/util/factory"
)

func GetLive(name string) *cobra.Command {
	cmd := &cobra.Command{
		Use:   "live",
		Short: "Commands for reading and writing resources to a cluster.",
	}
	ioStreams := genericclioptions.IOStreams{
		In:     cmd.InOrStdin(),
		Out:    cmd.OutOrStdout(),
		ErrOut: cmd.ErrOrStderr(),
	}

	flags := cmd.PersistentFlags()
	kubeConfigFlags := genericclioptions.NewConfigFlags(true).WithDeprecatedPasswordFlag()
	kubeConfigFlags.AddFlags(flags)
	userAgentKubeConfigFlags := &UserAgentKubeConfigFlags{
		Delegate:  kubeConfigFlags,
		UserAgent: "kustomize",
	}
	matchVersionKubeConfigFlags := util.NewMatchVersionFlags(
		&factory.CachingRESTClientGetter{
			Delegate: userAgentKubeConfigFlags,
		},
	)
	matchVersionKubeConfigFlags.AddFlags(cmd.PersistentFlags())
	cmd.PersistentFlags().AddGoFlagSet(flag.CommandLine)
	f := util.NewFactory(matchVersionKubeConfigFlags)

	applyCmd := apply.ApplyCommand(f, ioStreams)
	_ = applyCmd.Flags().MarkHidden("no-prune")

	cmd.AddCommand(
		applyCmd,
		initcmd.NewCmdInit(f, ioStreams),
		preview.GetPreviewRunner(f, ioStreams).Command,
		diff.NewCmdDiff(f, ioStreams),
		destroy.GetDestroyRunner(f, ioStreams).Command)
	return cmd
}

type UserAgentKubeConfigFlags struct {
	Delegate  genericclioptions.RESTClientGetter
	UserAgent string
}

func (u *UserAgentKubeConfigFlags) ToRESTConfig() (*rest.Config, error) {
	clientConfig, err := u.Delegate.ToRESTConfig()
	if err != nil {
		return nil, err
	}
	if u.UserAgent != "" {
		clientConfig.UserAgent = u.UserAgent
	}
	return clientConfig, nil
}

func (u *UserAgentKubeConfigFlags) ToDiscoveryClient() (discovery.CachedDiscoveryInterface, error) {
	return u.Delegate.ToDiscoveryClient()
}

func (u *UserAgentKubeConfigFlags) ToRESTMapper() (meta.RESTMapper, error) {
	return u.Delegate.ToRESTMapper()
}

func (u *UserAgentKubeConfigFlags) ToRawKubeConfigLoader() clientcmd.ClientConfig {
	return u.Delegate.ToRawKubeConfigLoader()
}
<commit_msg>Small change to live commands retrieval<commit_after>// Copyright 2019 The Kubernetes Authors.
// SPDX-License-Identifier: Apache-2.0

package configcobra

import (
	"flag"

	"github.com/spf13/cobra"
	"k8s.io/apimachinery/pkg/api/meta"
	"k8s.io/cli-runtime/pkg/genericclioptions"
	"k8s.io/client-go/discovery"
	"k8s.io/client-go/rest"
	"k8s.io/client-go/tools/clientcmd"
	"k8s.io/kubectl/pkg/cmd/util"
	"sigs.k8s.io/cli-utils/cmd/apply"
	"sigs.k8s.io/cli-utils/cmd/destroy"
	"sigs.k8s.io/cli-utils/cmd/diff"
	"sigs.k8s.io/cli-utils/cmd/initcmd"
	"sigs.k8s.io/cli-utils/cmd/preview"
	"sigs.k8s.io/cli-utils/pkg/util/factory"
)

func GetLive(name string) *cobra.Command {
	cmd := &cobra.Command{
		Use:   "live",
		Short: "Commands for reading and writing resources to a cluster.",
	}
	ioStreams := genericclioptions.IOStreams{
		In:     cmd.InOrStdin(),
		Out:    cmd.OutOrStdout(),
		ErrOut: cmd.ErrOrStderr(),
	}

	flags := cmd.PersistentFlags()
	kubeConfigFlags := genericclioptions.NewConfigFlags(true).WithDeprecatedPasswordFlag()
	kubeConfigFlags.AddFlags(flags)
	userAgentKubeConfigFlags := &UserAgentKubeConfigFlags{
		Delegate:  kubeConfigFlags,
		UserAgent: "kustomize",
	}
	matchVersionKubeConfigFlags := util.NewMatchVersionFlags(
		&factory.CachingRESTClientGetter{
			Delegate: userAgentKubeConfigFlags,
		},
	)
	matchVersionKubeConfigFlags.AddFlags(cmd.PersistentFlags())
	cmd.PersistentFlags().AddGoFlagSet(flag.CommandLine)
	f := util.NewFactory(matchVersionKubeConfigFlags)

	applyCmd := apply.ApplyCommand(f, ioStreams)
	_ = applyCmd.Flags().MarkHidden("no-prune")

	cmd.AddCommand(
		applyCmd,
		initcmd.NewCmdInit(f, ioStreams),
		preview.PreviewCommand(f, ioStreams),
		diff.NewCmdDiff(f, ioStreams),
		destroy.DestroyCommand(f, ioStreams))
	return cmd
}

type UserAgentKubeConfigFlags struct {
	Delegate  genericclioptions.RESTClientGetter
	UserAgent string
}

func (u *UserAgentKubeConfigFlags) ToRESTConfig() (*rest.Config, error) {
	clientConfig, err := u.Delegate.ToRESTConfig()
	if err != nil {
		return nil, err
	}
	if u.UserAgent != "" {
		clientConfig.UserAgent = u.UserAgent
	}
	return clientConfig, nil
}

func (u *UserAgentKubeConfigFlags) ToDiscoveryClient() (discovery.CachedDiscoveryInterface, error) {
	return u.Delegate.ToDiscoveryClient()
}

func (u *UserAgentKubeConfigFlags) ToRESTMapper() (meta.RESTMapper, error) {
	return u.Delegate.ToRESTMapper()
}

func (u *UserAgentKubeConfigFlags) ToRawKubeConfigLoader() clientcmd.ClientConfig {
	return u.Delegate.ToRawKubeConfigLoader()
}
<|endoftext|>


<commit_before>// Package migrate provides a tool to help migrate data from InfluxDB 1.x to 2.x
package migrate

import (
	"errors"
	"fmt"
	"os"
	"os/user"
	"path/filepath"
	"time"

	"github.com/influxdata/influxdb"
	"github.com/influxdata/influxdb/internal/fs"
	"github.com/influxdata/influxdb/kit/cli"
	"github.com/influxdata/influxdb/models"
	"github.com/influxdata/influxdb/tsdb/migrate"
	"github.com/spf13/cobra"
)

var Command = &cobra.Command{
	Use:   "migrate",
	Short: "Migrate from InfluxDB >= 1.7.x to InfluxDB 2.x",
	Long: `
TODO(edd)
`,
	Args: cobra.ExactArgs(0),
	RunE: migrateE,
}

var flags struct {
	basePath1x string // base path of 1.x installation
	db         string
	rp         string
	from       string // migrate only data at least as old as this (RFC3339Nano format)
	to         string // migrate only data at least as young as this (RFC3339Nano format)
	migrateHot bool   // migrate hot shards (can leave data behind)
	basePath2x string // base path of 2.x installation (defaults to ~/.influxdbv2)
	destOrg    string // destination 2.x organisation (base-16 format)

	dryRun  bool // enable dry-run mode (don't do any migration)
	verbose bool // enable verbose logging
}

// influx1Dir retrieves the InfluxDB 1.x directory.
func influx1Dir() (string, error) {
	var dir string
	// By default, store meta and data files in current users home directory
	u, err := user.Current()
	if err == nil {
		dir = u.HomeDir
	} else if home := os.Getenv("HOME"); home != "" {
		dir = home
	} else {
		wd, err := os.Getwd()
		if err != nil {
			return "", err
		}
		dir = wd
	}
	dir = filepath.Join(dir, ".influxdb")

	return dir, nil
}

func init() {
	v1Dir, err := influx1Dir()
	if err != nil {
		panic(fmt.Errorf("failed to determine default InfluxDB 1.x directory: %s", err))
	}

	v2Dir, err := fs.InfluxDir()
	if err != nil {
		panic(fmt.Errorf("failed to determine default InfluxDB 2.x directory: %s", err))
	}

	opts := []cli.Opt{
		{
			DestP:   &flags.basePath1x,
			Flag:    "influxdb-1x-path",
			Default: v1Dir,
			Desc:    "path to 1.x InfluxDB",
		},
		{
			DestP:   &flags.basePath2x,
			Flag:    "influxdb-2x-path",
			Default: v2Dir,
			Desc:    "path to 2.x InfluxDB",
		},
		{
			DestP:   &flags.db,
			Flag:    "db",
			Default: "",
			Desc:    "only import the provided 1.x database",
		},
		{
			DestP:   &flags.rp,
			Flag:    "rp",
			Default: "",
			Desc:    "only import the provided 1.x retention policy. --db must be set",
		},
		{
			DestP:   &flags.from,
			Flag:    "from",
			Default: "",
			Desc:    "earliest point to import",
		},
		{
			DestP:   &flags.to,
			Flag:    "to",
			Default: "",
			Desc:    "latest point to import",
		},
		{
			DestP:   &flags.destOrg,
			Flag:    "org-id",
			Default: "",
			Desc:    "destination 2.x organization id (required)",
		},
		{
			DestP:   &flags.dryRun,
			Flag:    "dry-run",
			Default: false,
			Desc:    "simulate migration without running it",
		},
		{
			DestP:   &flags.migrateHot,
			Flag:    "migrate-hot-shards",
			Default: false,
			Desc:    "migrate all shards including hot ones. Can leave unsnapshotted data behind",
		},
		{
			DestP:   &flags.verbose,
			Flag:    "verbose",
			Default: false,
			Desc:    "enable verbose logging",
		},
	}
	cli.BindOptions(Command, opts)
}

func migrateE(cmd *cobra.Command, args []string) error {
	if flags.destOrg == "" {
		return errors.New("destination organization must be set")
	} else if flags.rp != "" && flags.db == "" {
		return errors.New("source database empty. Cannot filter by retention policy")
	}

	destOrg, err := influxdb.IDFromString(flags.destOrg)
	if err != nil {
		return err
	}

	fromNano := models.MinNanoTime
	if flags.from != "" {
		from, err := time.Parse(time.RFC3339Nano, flags.from)
		if err != nil {
			return err
		}
		fromNano = from.UnixNano()

		// TODO: enable this feature...
		_ = fromNano
		return errors.New("--from flag is currently unsupported; only all-time can be migrated")
	}

	toNano := models.MaxNanoTime
	if flags.to != "" {
		to, err := time.Parse(time.RFC3339Nano, flags.to)
		if err != nil {
			return err
		}
		toNano = to.UnixNano()

		// TODO: enable this feature...
		_ = toNano
		return errors.New("--to flag is currently unsupported; only all-time can be migrated")
	}

	// TODO: enable this feature...
	if flags.migrateHot {
		return errors.New("--migrate-hot-shards flag is currently unsupported; fully compact all shards before migrating")
	}

	migrator := migrate.NewMigrator(migrate.Config{
		SourcePath:      flags.basePath1x,
		DestPath:        flags.basePath2x,
		From:            fromNano,
		To:              toNano,
		MigrateHotShard: flags.migrateHot,
		Stdout:          os.Stdout,
		VerboseLogging:  flags.verbose,
		DestOrg:         *destOrg,
		DryRun:          flags.dryRun,
	})
	return migrator.Process1xShards(flags.db, flags.rp)
}
<commit_msg>refactor: add some links to current issues<commit_after>// Package migrate provides a tool to help migrate data from InfluxDB 1.x to 2.x
package migrate

import (
	"errors"
	"fmt"
	"os"
	"os/user"
	"path/filepath"
	"time"

	"github.com/influxdata/influxdb"
	"github.com/influxdata/influxdb/internal/fs"
	"github.com/influxdata/influxdb/kit/cli"
	"github.com/influxdata/influxdb/models"
	"github.com/influxdata/influxdb/tsdb/migrate"
	"github.com/spf13/cobra"
)

var Command = &cobra.Command{
	Use:   "migrate",
	Short: "Migrate from InfluxDB >= 1.7.x to InfluxDB 2.x",
	Long: `NOTE!

⚠⚠ This tool is in a very early prototype stage. It can corrupt your 2.x TSM data.⚠⚠

Please see the following issues for progress towards making this tool production
ready:

- Ensure retention policy duration carried over to 2.x bucket: https://github.com/influxdata/influxdb/issues/17257
- Support migrating arbitrary time-ranges: https://github.com/influxdata/influxdb/issues/17249
- Ensure hot shards not migrated by default: https://github.com/influxdata/influxdb/issues/17250
- Handle cases where multiple 1.x shards have different field types for same series: https://github.com/influxdata/influxdb/issues/17251

This tool allows an operator to migrate their TSM data from an OSS 1.x server
into an OSS 2.x server in an offline manner.

It is very important when running this tool that the 2.x server is not running,
and ideally the 1.x server is not running, or at least, it is not writing into
any of the shards that will be migrated.
`,
	Args: cobra.ExactArgs(0),
	RunE: migrateE,
}

var flags struct {
	basePath1x string // base path of 1.x installation
	db         string
	rp         string
	from       string // migrate only data at least as old as this (RFC3339Nano format)
	to         string // migrate only data at least as young as this (RFC3339Nano format)
	migrateHot bool   // migrate hot shards (can leave data behind)
	basePath2x string // base path of 2.x installation (defaults to ~/.influxdbv2)
	destOrg    string // destination 2.x organisation (base-16 format)

	dryRun  bool // enable dry-run mode (don't do any migration)
	verbose bool // enable verbose logging
}

// influx1Dir retrieves the InfluxDB 1.x directory.
func influx1Dir() (string, error) {
	var dir string
	// By default, store meta and data files in current users home directory
	u, err := user.Current()
	if err == nil {
		dir = u.HomeDir
	} else if home := os.Getenv("HOME"); home != "" {
		dir = home
	} else {
		wd, err := os.Getwd()
		if err != nil {
			return "", err
		}
		dir = wd
	}
	dir = filepath.Join(dir, ".influxdb")

	return dir, nil
}

func init() {
	v1Dir, err := influx1Dir()
	if err != nil {
		panic(fmt.Errorf("failed to determine default InfluxDB 1.x directory: %s", err))
	}

	v2Dir, err := fs.InfluxDir()
	if err != nil {
		panic(fmt.Errorf("failed to determine default InfluxDB 2.x directory: %s", err))
	}

	opts := []cli.Opt{
		{
			DestP:   &flags.basePath1x,
			Flag:    "influxdb-1x-path",
			Default: v1Dir,
			Desc:    "path to 1.x InfluxDB",
		},
		{
			DestP:   &flags.basePath2x,
			Flag:    "influxdb-2x-path",
			Default: v2Dir,
			Desc:    "path to 2.x InfluxDB",
		},
		{
			DestP:   &flags.db,
			Flag:    "db",
			Default: "",
			Desc:    "only import the provided 1.x database",
		},
		{
			DestP:   &flags.rp,
			Flag:    "rp",
			Default: "",
			Desc:    "only import the provided 1.x retention policy. --db must be set",
		},
		{
			DestP:   &flags.from,
			Flag:    "from",
			Default: "",
			Desc:    "earliest point to import",
		},
		{
			DestP:   &flags.to,
			Flag:    "to",
			Default: "",
			Desc:    "latest point to import",
		},
		{
			DestP:   &flags.destOrg,
			Flag:    "org-id",
			Default: "",
			Desc:    "destination 2.x organization id (required)",
		},
		{
			DestP:   &flags.dryRun,
			Flag:    "dry-run",
			Default: false,
			Desc:    "simulate migration without running it",
		},
		{
			DestP:   &flags.migrateHot,
			Flag:    "migrate-hot-shards",
			Default: false,
			Desc:    "migrate all shards including hot ones. Can leave unsnapshotted data behind",
		},
		{
			DestP:   &flags.verbose,
			Flag:    "verbose",
			Default: false,
			Desc:    "enable verbose logging",
		},
	}
	cli.BindOptions(Command, opts)
}

func migrateE(cmd *cobra.Command, args []string) error {
	if flags.destOrg == "" {
		return errors.New("destination organization must be set")
	} else if flags.rp != "" && flags.db == "" {
		return errors.New("source database empty. Cannot filter by retention policy")
	}

	destOrg, err := influxdb.IDFromString(flags.destOrg)
	if err != nil {
		return err
	}

	fromNano := models.MinNanoTime
	if flags.from != "" {
		from, err := time.Parse(time.RFC3339Nano, flags.from)
		if err != nil {
			return err
		}
		fromNano = from.UnixNano()

		// TODO: enable this feature...
		_ = fromNano
		return errors.New("--from flag is currently unsupported; only all-time can be migrated")
	}

	toNano := models.MaxNanoTime
	if flags.to != "" {
		to, err := time.Parse(time.RFC3339Nano, flags.to)
		if err != nil {
			return err
		}
		toNano = to.UnixNano()

		// TODO: enable this feature...
		_ = toNano
		return errors.New("--to flag is currently unsupported; only all-time can be migrated")
	}

	// TODO: enable this feature...
	if flags.migrateHot {
		return errors.New("--migrate-hot-shards flag is currently unsupported; fully compact all shards before migrating")
	}

	migrator := migrate.NewMigrator(migrate.Config{
		SourcePath:      flags.basePath1x,
		DestPath:        flags.basePath2x,
		From:            fromNano,
		To:              toNano,
		MigrateHotShard: flags.migrateHot,
		Stdout:          os.Stdout,
		VerboseLogging:  flags.verbose,
		DestOrg:         *destOrg,
		DryRun:          flags.dryRun,
	})
	return migrator.Process1xShards(flags.db, flags.rp)
}
<|endoftext|>


<commit_before>/*
Copyright 2016 The Kubernetes Authors 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.
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 cmd

import (
	"bytes"
	"io"
	"os"

	"github.com/pkg/errors"
	"github.com/spf13/cobra"
	"k8s.io/minikube/pkg/minikube/exit"
	"k8s.io/minikube/pkg/minikube/out"
	"k8s.io/minikube/pkg/minikube/reason"
)

const longDescription = `Outputs minikube shell completion for the given shell (bash, zsh or fish)

	This depends on the bash-completion binary.  Example installation instructions:
	OS X:
		$ brew install bash-completion
		$ source $(brew --prefix)/etc/bash_completion
		$ minikube completion bash > ~/.minikube-completion  # for bash users
		$ minikube completion zsh > ~/.minikube-completion  # for zsh users
		$ source ~/.minikube-completion
		$ minikube completion fish > ~/.config/fish/completions/minikube.fish # for fish users
	Ubuntu:
		$ apt-get install bash-completion
		$ source /etc/bash_completion
		$ source <(minikube completion bash) # for bash users
		$ source <(minikube completion zsh) # for zsh users
		$ minikube completion fish > ~/.config/fish/completions/minikube.fish # for fish users

	Additionally, you may want to output the completion to a file and source in your .bashrc

	Note for zsh users: [1] zsh completions are only supported in versions of zsh >= 5.2
	Note for fish users: [2] please refer to this docs for more details https://fishshell.com/docs/current/#tab-completion
`

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

var completionCmd = &cobra.Command{
	Use:   "completion SHELL",
	Short: "Generate command completion for a shell",
	Long:  longDescription,
	Run: func(cmd *cobra.Command, args []string) {
		if len(args) != 1 {
			exit.Message(reason.Usage, "Usage: minikube completion SHELL")
		}
		if args[0] != "bash" && args[0] != "zsh" && args[0] != "fish" {
			exit.Message(reason.Usage, "Sorry, completion support is not yet implemented for {{.name}}", out.V{"name": args[0]})
		}
	},
}

var bashCmd = &cobra.Command{
	Use:   "bash",
	Short: "bash completion.",
	Long:  "Generate command completion for bash.",
	Run: func(cmd *cobra.Command, args []string) {
		err := GenerateBashCompletion(os.Stdout, cmd.Root())
		if err != nil {
			exit.Error(reason.InternalCompletion, "bash completion failed", err)
		}
	},
}

var zshCmd = &cobra.Command{
	Use:   "zsh",
	Short: "zsh completion.",
	Long:  "Generate command completion for zsh.",
	Run: func(cmd *cobra.Command, args []string) {
		err := GenerateZshCompletion(os.Stdout, cmd.Root())
		if err != nil {
			exit.Error(reason.InternalCompletion, "zsh completion failed", err)
		}
	},
}

var fishCmd = &cobra.Command{
	Use:   "fish",
	Short: "fish completion.",
	Long:  "Generate command completion for fish .",
	Run: func(cmd *cobra.Command, args []string) {
		err := GenerateFishCompletion(os.Stdout, cmd.Root())
		if err != nil {
			exit.Error(reason.InternalCompletion, "fish completion failed", err)
		}
	},
}

func init() {
	completionCmd.AddCommand(bashCmd)
	completionCmd.AddCommand(zshCmd)
	completionCmd.AddCommand(fishCmd)
}

// GenerateBashCompletion generates the completion for the bash shell
func GenerateBashCompletion(w io.Writer, cmd *cobra.Command) error {
	_, err := w.Write([]byte(boilerPlate))
	if err != nil {
		return err
	}

	err = cmd.GenBashCompletion(w)
	if err != nil {
		return errors.Wrap(err, "Error generating bash completion")
	}

	return nil
}

// GenerateZshCompletion generates the completion for the zsh shell
func GenerateZshCompletion(out io.Writer, cmd *cobra.Command) error {
	zshAutoloadTag := `#compdef minikube
`

	zshInitialization := `
__minikube_bash_source() {
	alias shopt=':'
	alias _expand=_bash_expand
	alias _complete=_bash_comp
	emulate -L sh
	setopt kshglob noshglob braceexpand
	source "$@"
}
__minikube_type() {
	# -t is not supported by zsh
	if [ "$1" == "-t" ]; then
		shift
		# fake Bash 4 to disable "complete -o nospace". Instead
		# "compopt +-o nospace" is used in the code to toggle trailing
		# spaces. We don't support that, but leave trailing spaces on
		# all the time
		if [ "$1" = "__minikube_compopt" ]; then
			echo builtin
			return 0
		fi
	fi
	type "$@"
}
__minikube_compgen() {
	local completions w
	completions=( $(compgen "$@") ) || return $?
	# filter by given word as prefix
	while [[ "$1" = -* && "$1" != -- ]]; do
		shift
		shift
	done
	if [[ "$1" == -- ]]; then
		shift
	fi
	for w in "${completions[@]}"; do
		if [[ "${w}" = "$1"* ]]; then
			echo "${w}"
		fi
	done
}
__minikube_compopt() {
	true # don't do anything. Not supported by bashcompinit in zsh
}
__minikube_declare() {
	if [ "$1" == "-F" ]; then
		whence -w "$@"
	else
		builtin declare "$@"
	fi
}
__minikube_ltrim_colon_completions()
{
	if [[ "$1" == *:* && "$COMP_WORDBREAKS" == *:* ]]; then
		# Remove colon-word prefix from COMPREPLY items
		local colon_word=${1%${1##*:}}
		local i=${#COMPREPLY[*]}
		while [[ $((--i)) -ge 0 ]]; do
			COMPREPLY[$i]=${COMPREPLY[$i]#"$colon_word"}
		done
	fi
}
__minikube_get_comp_words_by_ref() {
	cur="${COMP_WORDS[COMP_CWORD]}"
	prev="${COMP_WORDS[${COMP_CWORD}-1]}"
	words=("${COMP_WORDS[@]}")
	cword=("${COMP_CWORD[@]}")
}
__minikube_filedir() {
	local RET OLD_IFS w qw
	__debug "_filedir $@ cur=$cur"
	if [[ "$1" = \~* ]]; then
		# somehow does not work. Maybe, zsh does not call this at all
		eval echo "$1"
		return 0
	fi
	OLD_IFS="$IFS"
	IFS=$'\n'
	if [ "$1" = "-d" ]; then
		shift
		RET=( $(compgen -d) )
	else
		RET=( $(compgen -f) )
	fi
	IFS="$OLD_IFS"
	IFS="," __debug "RET=${RET[@]} len=${#RET[@]}"
	for w in ${RET[@]}; do
		if [[ ! "${w}" = "${cur}"* ]]; then
			continue
		fi
		if eval "[[ \"\${w}\" = *.$1 || -d \"\${w}\" ]]"; then
			qw="$(__minikube_quote "${w}")"
			if [ -d "${w}" ]; then
				COMPREPLY+=("${qw}/")
			else
				COMPREPLY+=("${qw}")
			fi
		fi
	done
}
__minikube_quote() {
	if [[ $1 == \'* || $1 == \"* ]]; then
		# Leave out first character
		printf %q "${1:1}"
	else
		printf %q "$1"
	fi
}
autoload -U +X bashcompinit && bashcompinit
# use word boundary patterns for BSD or GNU sed
LWORD='[[:<:]]'
RWORD='[[:>:]]'
if sed --help 2>&1 | grep -q GNU; then
	LWORD='\<'
	RWORD='\>'
fi
__minikube_convert_bash_to_zsh() {
	sed \
	-e 's/declare -F/whence -w/' \
	-e 's/_get_comp_words_by_ref "\$@"/_get_comp_words_by_ref "\$*"/' \
	-e 's/local \([a-zA-Z0-9_]*\)=/local \1; \1=/' \
	-e 's/flags+=("\(--.*\)=")/flags+=("\1"); two_word_flags+=("\1")/' \
	-e 's/must_have_one_flag+=("\(--.*\)=")/must_have_one_flag+=("\1")/' \
	-e "s/${LWORD}_filedir${RWORD}/__minikube_filedir/g" \
	-e "s/${LWORD}_get_comp_words_by_ref${RWORD}/__minikube_get_comp_words_by_ref/g" \
	-e "s/${LWORD}__ltrim_colon_completions${RWORD}/__minikube_ltrim_colon_completions/g" \
	-e "s/${LWORD}compgen${RWORD}/__minikube_compgen/g" \
	-e "s/${LWORD}compopt${RWORD}/__minikube_compopt/g" \
	-e "s/${LWORD}declare${RWORD}/__minikube_declare/g" \
	-e "s/\\\$(type${RWORD}/\$(__minikube_type/g" \
	<<'BASH_COMPLETION_EOF'
`

	_, err := out.Write([]byte(zshAutoloadTag))
	if err != nil {
		return err
	}

	_, err = out.Write([]byte(boilerPlate))
	if err != nil {
		return err
	}

	_, err = out.Write([]byte(zshInitialization))
	if err != nil {
		return err
	}

	buf := new(bytes.Buffer)
	err = cmd.GenZshCompletion(buf)
	if err != nil {
		return errors.Wrap(err, "Error generating zsh completion")
	}
	_, err = out.Write(buf.Bytes())
	if err != nil {
		return err
	}

	zshTail := `
BASH_COMPLETION_EOF
}
__minikube_bash_source <(__minikube_convert_bash_to_zsh)
`
	_, err = out.Write([]byte(zshTail))
	if err != nil {
		return err
	}

	return nil
}

// GenerateFishCompletion generates the completion for the bash shell
func GenerateFishCompletion(w io.Writer, cmd *cobra.Command) error {
	_, err := w.Write([]byte(boilerPlate))
	if err != nil {
		return err
	}

	err = cmd.GenFishCompletion(w, true)
	if err != nil {
		return errors.Wrap(err, "Error generating fish completion")
	}

	return nil
}
<commit_msg>fix: revert wrong fix and re-fix zsh completion<commit_after>/*
Copyright 2016 The Kubernetes Authors 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.
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 cmd

import (
	"bytes"
	"io"
	"os"

	"github.com/pkg/errors"
	"github.com/spf13/cobra"
	"k8s.io/minikube/pkg/minikube/exit"
	"k8s.io/minikube/pkg/minikube/out"
	"k8s.io/minikube/pkg/minikube/reason"
)

const longDescription = `Outputs minikube shell completion for the given shell (bash, zsh or fish)

	This depends on the bash-completion binary.  Example installation instructions:
	OS X:
		$ brew install bash-completion
		$ source $(brew --prefix)/etc/bash_completion
		$ minikube completion bash > ~/.minikube-completion  # for bash users
		$ minikube completion zsh > ~/.minikube-completion  # for zsh users
		$ source ~/.minikube-completion
		$ minikube completion fish > ~/.config/fish/completions/minikube.fish # for fish users
	Ubuntu:
		$ apt-get install bash-completion
		$ source /etc/bash_completion
		$ source <(minikube completion bash) # for bash users
		$ source <(minikube completion zsh) # for zsh users
		$ minikube completion fish > ~/.config/fish/completions/minikube.fish # for fish users

	Additionally, you may want to output the completion to a file and source in your .bashrc

	Note for zsh users: [1] zsh completions are only supported in versions of zsh >= 5.2
	Note for fish users: [2] please refer to this docs for more details https://fishshell.com/docs/current/#tab-completion
`

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

var completionCmd = &cobra.Command{
	Use:   "completion SHELL",
	Short: "Generate command completion for a shell",
	Long:  longDescription,
	Run: func(cmd *cobra.Command, args []string) {
		if len(args) != 1 {
			exit.Message(reason.Usage, "Usage: minikube completion SHELL")
		}
		if args[0] != "bash" && args[0] != "zsh" && args[0] != "fish" {
			exit.Message(reason.Usage, "Sorry, completion support is not yet implemented for {{.name}}", out.V{"name": args[0]})
		}
	},
}

var bashCmd = &cobra.Command{
	Use:   "bash",
	Short: "bash completion.",
	Long:  "Generate command completion for bash.",
	Run: func(cmd *cobra.Command, args []string) {
		err := GenerateBashCompletion(os.Stdout, cmd.Root())
		if err != nil {
			exit.Error(reason.InternalCompletion, "bash completion failed", err)
		}
	},
}

var zshCmd = &cobra.Command{
	Use:   "zsh",
	Short: "zsh completion.",
	Long:  "Generate command completion for zsh.",
	Run: func(cmd *cobra.Command, args []string) {
		err := GenerateZshCompletion(os.Stdout, cmd.Root())
		if err != nil {
			exit.Error(reason.InternalCompletion, "zsh completion failed", err)
		}
	},
}

var fishCmd = &cobra.Command{
	Use:   "fish",
	Short: "fish completion.",
	Long:  "Generate command completion for fish .",
	Run: func(cmd *cobra.Command, args []string) {
		err := GenerateFishCompletion(os.Stdout, cmd.Root())
		if err != nil {
			exit.Error(reason.InternalCompletion, "fish completion failed", err)
		}
	},
}

func init() {
	completionCmd.AddCommand(bashCmd)
	completionCmd.AddCommand(zshCmd)
	completionCmd.AddCommand(fishCmd)
}

// GenerateBashCompletion generates the completion for the bash shell
func GenerateBashCompletion(w io.Writer, cmd *cobra.Command) error {
	_, err := w.Write([]byte(boilerPlate))
	if err != nil {
		return err
	}

	err = cmd.GenBashCompletion(w)
	if err != nil {
		return errors.Wrap(err, "Error generating bash completion")
	}

	return nil
}

// GenerateZshCompletion generates the completion for the zsh shell
func GenerateZshCompletion(out io.Writer, cmd *cobra.Command) error {
	zshAutoloadTag := `#compdef minikube
`

	zshInitialization := `
__minikube_bash_source() {
	alias shopt=':'
	alias _expand=_bash_expand
	alias _complete=_bash_comp
	emulate -L sh
	setopt kshglob noshglob braceexpand
	source "$@"
}
__minikube_type() {
	# -t is not supported by zsh
	if [ "$1" == "-t" ]; then
		shift
		# fake Bash 4 to disable "complete -o nospace". Instead
		# "compopt +-o nospace" is used in the code to toggle trailing
		# spaces. We don't support that, but leave trailing spaces on
		# all the time
		if [ "$1" = "__minikube_compopt" ]; then
			echo builtin
			return 0
		fi
	fi
	type "$@"
}
__minikube_compgen() {
	local completions w
	completions=( $(compgen "$@") ) || return $?
	# filter by given word as prefix
	while [[ "$1" = -* && "$1" != -- ]]; do
		shift
		shift
	done
	if [[ "$1" == -- ]]; then
		shift
	fi
	for w in "${completions[@]}"; do
		if [[ "${w}" = "$1"* ]]; then
			echo "${w}"
		fi
	done
}
__minikube_compopt() {
	true # don't do anything. Not supported by bashcompinit in zsh
}
__minikube_declare() {
	if [ "$1" == "-F" ]; then
		whence -w "$@"
	else
		builtin declare "$@"
	fi
}
__minikube_ltrim_colon_completions()
{
	if [[ "$1" == *:* && "$COMP_WORDBREAKS" == *:* ]]; then
		# Remove colon-word prefix from COMPREPLY items
		local colon_word=${1%${1##*:}}
		local i=${#COMPREPLY[*]}
		while [[ $((--i)) -ge 0 ]]; do
			COMPREPLY[$i]=${COMPREPLY[$i]#"$colon_word"}
		done
	fi
}
__minikube_get_comp_words_by_ref() {
	cur="${COMP_WORDS[COMP_CWORD]}"
	prev="${COMP_WORDS[${COMP_CWORD}-1]}"
	words=("${COMP_WORDS[@]}")
	cword=("${COMP_CWORD[@]}")
}
__minikube_filedir() {
	local RET OLD_IFS w qw
	__debug "_filedir $@ cur=$cur"
	if [[ "$1" = \~* ]]; then
		# somehow does not work. Maybe, zsh does not call this at all
		eval echo "$1"
		return 0
	fi
	OLD_IFS="$IFS"
	IFS=$'\n'
	if [ "$1" = "-d" ]; then
		shift
		RET=( $(compgen -d) )
	else
		RET=( $(compgen -f) )
	fi
	IFS="$OLD_IFS"
	IFS="," __debug "RET=${RET[@]} len=${#RET[@]}"
	for w in ${RET[@]}; do
		if [[ ! "${w}" = "${cur}"* ]]; then
			continue
		fi
		if eval "[[ \"\${w}\" = *.$1 || -d \"\${w}\" ]]"; then
			qw="$(__minikube_quote "${w}")"
			if [ -d "${w}" ]; then
				COMPREPLY+=("${qw}/")
			else
				COMPREPLY+=("${qw}")
			fi
		fi
	done
}
__minikube_quote() {
	if [[ $1 == \'* || $1 == \"* ]]; then
		# Leave out first character
		printf %q "${1:1}"
	else
		printf %q "$1"
	fi
}
autoload -U +X bashcompinit && bashcompinit
# use word boundary patterns for BSD or GNU sed
LWORD='[[:<:]]'
RWORD='[[:>:]]'
if sed --help 2>&1 | grep -q GNU; then
	LWORD='\<'
	RWORD='\>'
fi
__minikube_convert_bash_to_zsh() {
	sed \
	-e 's/declare -F/whence -w/' \
	-e 's/_get_comp_words_by_ref "\$@"/_get_comp_words_by_ref "\$*"/' \
	-e 's/local \([a-zA-Z0-9_]*\)=/local \1; \1=/' \
	-e 's/flags+=("\(--.*\)=")/flags+=("\1"); two_word_flags+=("\1")/' \
	-e 's/must_have_one_flag+=("\(--.*\)=")/must_have_one_flag+=("\1")/' \
	-e "s/${LWORD}_filedir${RWORD}/__minikube_filedir/g" \
	-e "s/${LWORD}_get_comp_words_by_ref${RWORD}/__minikube_get_comp_words_by_ref/g" \
	-e "s/${LWORD}__ltrim_colon_completions${RWORD}/__minikube_ltrim_colon_completions/g" \
	-e "s/${LWORD}compgen${RWORD}/__minikube_compgen/g" \
	-e "s/${LWORD}compopt${RWORD}/__minikube_compopt/g" \
	-e "s/${LWORD}declare${RWORD}/__minikube_declare/g" \
	-e "s/\\\$(type${RWORD}/\$(__minikube_type/g" \
	-e "s/aliashash\[\"\([a-z]*\)\"\]/aliashash[\1]/g" \
	<<'BASH_COMPLETION_EOF'
`

	_, err := out.Write([]byte(zshAutoloadTag))
	if err != nil {
		return err
	}

	_, err = out.Write([]byte(boilerPlate))
	if err != nil {
		return err
	}

	_, err = out.Write([]byte(zshInitialization))
	if err != nil {
		return err
	}

	buf := new(bytes.Buffer)
	err = cmd.GenBashCompletion(buf)
	if err != nil {
		return errors.Wrap(err, "Error generating zsh completion")
	}
	_, err = out.Write(buf.Bytes())
	if err != nil {
		return err
	}

	zshTail := `
BASH_COMPLETION_EOF
}
__minikube_bash_source <(__minikube_convert_bash_to_zsh)
`
	_, err = out.Write([]byte(zshTail))
	if err != nil {
		return err
	}

	return nil
}

// GenerateFishCompletion generates the completion for the bash shell
func GenerateFishCompletion(w io.Writer, cmd *cobra.Command) error {
	_, err := w.Write([]byte(boilerPlate))
	if err != nil {
		return err
	}

	err = cmd.GenFishCompletion(w, true)
	if err != nil {
		return errors.Wrap(err, "Error generating fish completion")
	}

	return nil
}
<|endoftext|>


<commit_before>/*
 *
 * k6 - a next-generation load testing tool
 * Copyright (C) 2016 Load Impact
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

package http

import (
	"context"
	"errors"
	"fmt"
	"net/url"
	"strings"

	"github.com/dop251/goja"

	"github.com/loadimpact/k6/js/common"
	"github.com/loadimpact/k6/js/modules/k6/html"
	"github.com/loadimpact/k6/lib/netext/httpext"
)

// Response is a representation of an HTTP response to be returned to the goja VM
type Response struct {
	*httpext.Response `js:"-"`
	h                 *HTTP
}

// processResponse stores the body as an ArrayBuffer if indicated by
// respType. This is done here instead of in httpext.readResponseBody to avoid
// a reverse dependency on js/common or goja.
func processResponse(ctx context.Context, resp *httpext.Response, respType httpext.ResponseType) {
	if respType == httpext.ResponseTypeBinary {
		rt := common.GetRuntime(ctx)
		resp.Body = rt.NewArrayBuffer(resp.Body.([]byte))
	}
}

func (h *HTTP) responseFromHttpext(resp *httpext.Response) *Response {
	return &Response{Response: resp, h: h}
}

// JSON parses the body of a response as json and returns it to the goja VM
func (res *Response) JSON(selector ...string) goja.Value {
	v, err := res.Response.JSON(selector...)
	if err != nil {
		common.Throw(common.GetRuntime(res.GetCtx()), err)
	}
	if v == nil {
		return goja.Undefined()
	}
	return common.GetRuntime(res.GetCtx()).ToValue(v)
}

// HTML returns the body as an html.Selection
func (res *Response) HTML(selector ...string) html.Selection {
	var body string
	switch b := res.Body.(type) {
	case []byte:
		body = string(b)
	case string:
		body = b
	default:
		common.Throw(common.GetRuntime(res.GetCtx()), errors.New("invalid response type"))
	}

	sel, err := html.HTML{}.ParseHTML(res.GetCtx(), body)
	if err != nil {
		common.Throw(common.GetRuntime(res.GetCtx()), err)
	}
	sel.URL = res.URL
	if len(selector) > 0 {
		sel = sel.Find(selector[0])
	}
	return sel
}

// SubmitForm parses the body as an html looking for a from and then submitting it
// TODO: document the actual arguments that can be provided
func (res *Response) SubmitForm(args ...goja.Value) (*Response, error) {
	rt := common.GetRuntime(res.GetCtx())

	formSelector := "form"
	submitSelector := "[type=\"submit\"]"
	var fields map[string]goja.Value
	requestParams := goja.Null()
	if len(args) > 0 {
		params := args[0].ToObject(rt)
		for _, k := range params.Keys() {
			switch k {
			case "formSelector":
				formSelector = params.Get(k).String()
			case "submitSelector":
				submitSelector = params.Get(k).String()
			case "fields":
				if rt.ExportTo(params.Get(k), &fields) != nil {
					fields = nil
				}
			case "params":
				requestParams = params.Get(k)
			}
		}
	}

	form := res.HTML(formSelector)
	if form.Size() == 0 {
		common.Throw(rt, fmt.Errorf("no form found for selector '%s' in response '%s'", formSelector, res.URL))
	}

	methodAttr := form.Attr("method")
	var requestMethod string
	if methodAttr == goja.Undefined() {
		// Use GET by default
		requestMethod = HTTP_METHOD_GET
	} else {
		requestMethod = strings.ToUpper(methodAttr.String())
	}

	responseURL, err := url.Parse(res.URL)
	if err != nil {
		common.Throw(rt, err)
	}

	actionAttr := form.Attr("action")
	var requestURL *url.URL
	if actionAttr == goja.Undefined() {
		// Use the url of the response if no action is set
		requestURL = responseURL
	} else {
		actionURL, err := url.Parse(actionAttr.String())
		if err != nil {
			common.Throw(rt, err)
		}
		requestURL = responseURL.ResolveReference(actionURL)
	}

	// Set the body based on the form values
	values := form.SerializeObject()

	// Set the name + value of the submit button
	submit := form.Find(submitSelector)
	submitName := submit.Attr("name")
	submitValue := submit.Val()
	if submitName != goja.Undefined() && submitValue != goja.Undefined() {
		values[submitName.String()] = submitValue
	}

	// Set the values supplied in the arguments, overriding automatically set values
	for k, v := range fields {
		values[k] = v
	}

	if requestMethod == HTTP_METHOD_GET {
		q := url.Values{}
		for k, v := range values {
			q.Add(k, v.String())
		}
		requestURL.RawQuery = q.Encode()
		return res.h.Request(res.GetCtx(), requestMethod, rt.ToValue(requestURL.String()), goja.Null(), requestParams)
	}
	return res.h.Request(res.GetCtx(), requestMethod, rt.ToValue(requestURL.String()), rt.ToValue(values), requestParams)
}

// ClickLink parses the body as an html, looks for a link and than makes a request as if the link was
// clicked
func (res *Response) ClickLink(args ...goja.Value) (*Response, error) {
	rt := common.GetRuntime(res.GetCtx())

	selector := "a[href]"
	requestParams := goja.Null()
	if len(args) > 0 {
		params := args[0].ToObject(rt)
		for _, k := range params.Keys() {
			switch k {
			case "selector":
				selector = params.Get(k).String()
			case "params":
				requestParams = params.Get(k)
			}
		}
	}

	responseURL, err := url.Parse(res.URL)
	if err != nil {
		common.Throw(rt, err)
	}

	link := res.HTML(selector)
	if link.Size() == 0 {
		common.Throw(rt, fmt.Errorf("no element found for selector '%s' in response '%s'", selector, res.URL))
	}
	hrefAttr := link.Attr("href")
	if hrefAttr == goja.Undefined() {
		common.Throw(rt, fmt.Errorf("no valid href attribute value found on element '%s' in response '%s'", selector, res.URL))
	}
	hrefURL, err := url.Parse(hrefAttr.String())
	if err != nil {
		common.Throw(rt, err)
	}
	requestURL := responseURL.ResolveReference(hrefURL)

	return res.h.Get(res.GetCtx(), rt.ToValue(requestURL.String()), requestParams)
}
<commit_msg>Use ToBytes to convert response bodies in Response.HTML()<commit_after>/*
 *
 * k6 - a next-generation load testing tool
 * Copyright (C) 2016 Load Impact
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

package http

import (
	"context"
	"fmt"
	"net/url"
	"strings"

	"github.com/dop251/goja"

	"github.com/loadimpact/k6/js/common"
	"github.com/loadimpact/k6/js/modules/k6/html"
	"github.com/loadimpact/k6/lib/netext/httpext"
)

// Response is a representation of an HTTP response to be returned to the goja VM
type Response struct {
	*httpext.Response `js:"-"`
	h                 *HTTP
}

// processResponse stores the body as an ArrayBuffer if indicated by
// respType. This is done here instead of in httpext.readResponseBody to avoid
// a reverse dependency on js/common or goja.
func processResponse(ctx context.Context, resp *httpext.Response, respType httpext.ResponseType) {
	if respType == httpext.ResponseTypeBinary {
		rt := common.GetRuntime(ctx)
		resp.Body = rt.NewArrayBuffer(resp.Body.([]byte))
	}
}

func (h *HTTP) responseFromHttpext(resp *httpext.Response) *Response {
	return &Response{Response: resp, h: h}
}

// JSON parses the body of a response as json and returns it to the goja VM
func (res *Response) JSON(selector ...string) goja.Value {
	v, err := res.Response.JSON(selector...)
	if err != nil {
		common.Throw(common.GetRuntime(res.GetCtx()), err)
	}
	if v == nil {
		return goja.Undefined()
	}
	return common.GetRuntime(res.GetCtx()).ToValue(v)
}

// HTML returns the body as an html.Selection
func (res *Response) HTML(selector ...string) html.Selection {
	body, err := common.ToBytes(res.Body)
	if err != nil {
		common.Throw(common.GetRuntime(res.GetCtx()), err)
	}

	sel, err := html.HTML{}.ParseHTML(res.GetCtx(), string(body))
	if err != nil {
		common.Throw(common.GetRuntime(res.GetCtx()), err)
	}
	sel.URL = res.URL
	if len(selector) > 0 {
		sel = sel.Find(selector[0])
	}
	return sel
}

// SubmitForm parses the body as an html looking for a from and then submitting it
// TODO: document the actual arguments that can be provided
func (res *Response) SubmitForm(args ...goja.Value) (*Response, error) {
	rt := common.GetRuntime(res.GetCtx())

	formSelector := "form"
	submitSelector := "[type=\"submit\"]"
	var fields map[string]goja.Value
	requestParams := goja.Null()
	if len(args) > 0 {
		params := args[0].ToObject(rt)
		for _, k := range params.Keys() {
			switch k {
			case "formSelector":
				formSelector = params.Get(k).String()
			case "submitSelector":
				submitSelector = params.Get(k).String()
			case "fields":
				if rt.ExportTo(params.Get(k), &fields) != nil {
					fields = nil
				}
			case "params":
				requestParams = params.Get(k)
			}
		}
	}

	form := res.HTML(formSelector)
	if form.Size() == 0 {
		common.Throw(rt, fmt.Errorf("no form found for selector '%s' in response '%s'", formSelector, res.URL))
	}

	methodAttr := form.Attr("method")
	var requestMethod string
	if methodAttr == goja.Undefined() {
		// Use GET by default
		requestMethod = HTTP_METHOD_GET
	} else {
		requestMethod = strings.ToUpper(methodAttr.String())
	}

	responseURL, err := url.Parse(res.URL)
	if err != nil {
		common.Throw(rt, err)
	}

	actionAttr := form.Attr("action")
	var requestURL *url.URL
	if actionAttr == goja.Undefined() {
		// Use the url of the response if no action is set
		requestURL = responseURL
	} else {
		actionURL, err := url.Parse(actionAttr.String())
		if err != nil {
			common.Throw(rt, err)
		}
		requestURL = responseURL.ResolveReference(actionURL)
	}

	// Set the body based on the form values
	values := form.SerializeObject()

	// Set the name + value of the submit button
	submit := form.Find(submitSelector)
	submitName := submit.Attr("name")
	submitValue := submit.Val()
	if submitName != goja.Undefined() && submitValue != goja.Undefined() {
		values[submitName.String()] = submitValue
	}

	// Set the values supplied in the arguments, overriding automatically set values
	for k, v := range fields {
		values[k] = v
	}

	if requestMethod == HTTP_METHOD_GET {
		q := url.Values{}
		for k, v := range values {
			q.Add(k, v.String())
		}
		requestURL.RawQuery = q.Encode()
		return res.h.Request(res.GetCtx(), requestMethod, rt.ToValue(requestURL.String()), goja.Null(), requestParams)
	}
	return res.h.Request(res.GetCtx(), requestMethod, rt.ToValue(requestURL.String()), rt.ToValue(values), requestParams)
}

// ClickLink parses the body as an html, looks for a link and than makes a request as if the link was
// clicked
func (res *Response) ClickLink(args ...goja.Value) (*Response, error) {
	rt := common.GetRuntime(res.GetCtx())

	selector := "a[href]"
	requestParams := goja.Null()
	if len(args) > 0 {
		params := args[0].ToObject(rt)
		for _, k := range params.Keys() {
			switch k {
			case "selector":
				selector = params.Get(k).String()
			case "params":
				requestParams = params.Get(k)
			}
		}
	}

	responseURL, err := url.Parse(res.URL)
	if err != nil {
		common.Throw(rt, err)
	}

	link := res.HTML(selector)
	if link.Size() == 0 {
		common.Throw(rt, fmt.Errorf("no element found for selector '%s' in response '%s'", selector, res.URL))
	}
	hrefAttr := link.Attr("href")
	if hrefAttr == goja.Undefined() {
		common.Throw(rt, fmt.Errorf("no valid href attribute value found on element '%s' in response '%s'", selector, res.URL))
	}
	hrefURL, err := url.Parse(hrefAttr.String())
	if err != nil {
		common.Throw(rt, err)
	}
	requestURL := responseURL.ResolveReference(hrefURL)

	return res.h.Get(res.GetCtx(), rt.ToValue(requestURL.String()), requestParams)
}
<|endoftext|>