<commit_before>package cloudwatch

import (
	"encoding/json"
	"errors"
	"io/ioutil"
	"strings"
	"sync"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awsutil"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/credentials/ec2rolecreds"
	"github.com/aws/aws-sdk-go/aws/ec2metadata"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/cloudwatch"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/aws/aws-sdk-go/service/sts"
	"github.com/grafana/grafana/pkg/log"
	"github.com/grafana/grafana/pkg/middleware"
	m "github.com/grafana/grafana/pkg/models"
)

type actionHandler func(*cwRequest, *middleware.Context)

var actionHandlers map[string]actionHandler

type cwRequest struct {
	Region     string `json:"region"`
	Action     string `json:"action"`
	Body       []byte `json:"-"`
	DataSource *m.DataSource
}

type datasourceInfo struct {
	Profile       string
	Region        string
	AssumeRoleArn string
	Namespace     string

	AccessKey string
	SecretKey string
}

func (req *cwRequest) GetDatasourceInfo() *datasourceInfo {
	assumeRoleArn := req.DataSource.JsonData.Get("assumeRoleArn").MustString()
	accessKey := ""
	secretKey := ""

	for key, value := range req.DataSource.SecureJsonData.Decrypt() {
		if key == "accessKey" {
			accessKey = value
		}
		if key == "secretKey" {
			secretKey = value
		}
	}

	return &datasourceInfo{
		AssumeRoleArn: assumeRoleArn,
		Region:        req.Region,
		Profile:       req.DataSource.Database,
		AccessKey:     accessKey,
		SecretKey:     secretKey,
	}
}

func init() {
	actionHandlers = map[string]actionHandler{
		"GetMetricStatistics":     handleGetMetricStatistics,
		"ListMetrics":             handleListMetrics,
		"DescribeAlarms":          handleDescribeAlarms,
		"DescribeAlarmsForMetric": handleDescribeAlarmsForMetric,
		"DescribeAlarmHistory":    handleDescribeAlarmHistory,
		"DescribeInstances":       handleDescribeInstances,
		"__GetRegions":            handleGetRegions,
		"__GetNamespaces":         handleGetNamespaces,
		"__GetMetrics":            handleGetMetrics,
		"__GetDimensions":         handleGetDimensions,
	}
}

type cache struct {
	credential *credentials.Credentials
	expiration *time.Time
}

var awsCredentialCache map[string]cache = make(map[string]cache)
var credentialCacheLock sync.RWMutex

func getCredentials(dsInfo *datasourceInfo) *credentials.Credentials {
	cacheKey := dsInfo.Profile + ":" + dsInfo.AssumeRoleArn
	credentialCacheLock.RLock()
	if _, ok := awsCredentialCache[cacheKey]; ok {
		if awsCredentialCache[cacheKey].expiration != nil &&
			(*awsCredentialCache[cacheKey].expiration).After(time.Now().UTC()) {
			result := awsCredentialCache[cacheKey].credential
			credentialCacheLock.RUnlock()
			return result
		}
	}
	credentialCacheLock.RUnlock()

	accessKeyId := ""
	secretAccessKey := ""
	sessionToken := ""
	var expiration *time.Time
	expiration = nil
	if strings.Index(dsInfo.AssumeRoleArn, "arn:aws:iam:") == 0 {
		params := &sts.AssumeRoleInput{
			RoleArn:         aws.String(dsInfo.AssumeRoleArn),
			RoleSessionName: aws.String("GrafanaSession"),
			DurationSeconds: aws.Int64(900),
		}

		stsSess := session.New()
		stsCreds := credentials.NewChainCredentials(
			[]credentials.Provider{
				&credentials.EnvProvider{},
				&credentials.SharedCredentialsProvider{Filename: "", Profile: dsInfo.Profile},
				&ec2rolecreds.EC2RoleProvider{Client: ec2metadata.New(stsSess), ExpiryWindow: 5 * time.Minute},
			})
		stsConfig := &aws.Config{
			Region:      aws.String(dsInfo.Region),
			Credentials: stsCreds,
		}

		svc := sts.New(session.New(stsConfig), stsConfig)
		resp, err := svc.AssumeRole(params)
		if err != nil {
			// ignore
			log.Error(3, "CloudWatch: Failed to assume role", err)
		}
		if resp.Credentials != nil {
			accessKeyId = *resp.Credentials.AccessKeyId
			secretAccessKey = *resp.Credentials.SecretAccessKey
			sessionToken = *resp.Credentials.SessionToken
			expiration = resp.Credentials.Expiration
		}
	}

	sess := session.New()
	creds := credentials.NewChainCredentials(
		[]credentials.Provider{
			&credentials.StaticProvider{Value: credentials.Value{
				AccessKeyID:     accessKeyId,
				SecretAccessKey: secretAccessKey,
				SessionToken:    sessionToken,
			}},
			&credentials.EnvProvider{},
			&credentials.StaticProvider{Value: credentials.Value{
				AccessKeyID:     dsInfo.AccessKey,
				SecretAccessKey: dsInfo.SecretKey,
			}},
			&credentials.SharedCredentialsProvider{Filename: "", Profile: dsInfo.Profile},
			&ec2rolecreds.EC2RoleProvider{Client: ec2metadata.New(sess), ExpiryWindow: 5 * time.Minute},
		})

	credentialCacheLock.Lock()
	awsCredentialCache[cacheKey] = cache{
		credential: creds,
		expiration: expiration,
	}
	credentialCacheLock.Unlock()

	return creds
}

func getAwsConfig(req *cwRequest) *aws.Config {
	cfg := &aws.Config{
		Region:      aws.String(req.Region),
		Credentials: getCredentials(req.GetDatasourceInfo()),
	}
	return cfg
}

func handleGetMetricStatistics(req *cwRequest, c *middleware.Context) {
	cfg := getAwsConfig(req)
	svc := cloudwatch.New(session.New(cfg), cfg)

	reqParam := &struct {
		Parameters struct {
			Namespace          string                  `json:"namespace"`
			MetricName         string                  `json:"metricName"`
			Dimensions         []*cloudwatch.Dimension `json:"dimensions"`
			Statistics         []*string               `json:"statistics"`
			ExtendedStatistics []*string               `json:"extendedStatistics"`
			StartTime          int64                   `json:"startTime"`
			EndTime            int64                   `json:"endTime"`
			Period             int64                   `json:"period"`
		} `json:"parameters"`
	}{}
	json.Unmarshal(req.Body, reqParam)

	params := &cloudwatch.GetMetricStatisticsInput{
		Namespace:          aws.String(reqParam.Parameters.Namespace),
		MetricName:         aws.String(reqParam.Parameters.MetricName),
		Dimensions:         reqParam.Parameters.Dimensions,
		Statistics:         reqParam.Parameters.Statistics,
		ExtendedStatistics: reqParam.Parameters.ExtendedStatistics,
		StartTime:          aws.Time(time.Unix(reqParam.Parameters.StartTime, 0)),
		EndTime:            aws.Time(time.Unix(reqParam.Parameters.EndTime, 0)),
		Period:             aws.Int64(reqParam.Parameters.Period),
	}
	if len(reqParam.Parameters.Statistics) != 0 {
		params.Statistics = reqParam.Parameters.Statistics
	}
	if len(reqParam.Parameters.ExtendedStatistics) != 0 {
		params.ExtendedStatistics = reqParam.Parameters.ExtendedStatistics
	}

	resp, err := svc.GetMetricStatistics(params)
	if err != nil {
		c.JsonApiErr(500, "Unable to call AWS API", err)
		return
	}

	c.JSON(200, resp)
}

func handleListMetrics(req *cwRequest, c *middleware.Context) {
	cfg := getAwsConfig(req)
	svc := cloudwatch.New(session.New(cfg), cfg)

	reqParam := &struct {
		Parameters struct {
			Namespace  string                        `json:"namespace"`
			MetricName string                        `json:"metricName"`
			Dimensions []*cloudwatch.DimensionFilter `json:"dimensions"`
		} `json:"parameters"`
	}{}
	json.Unmarshal(req.Body, reqParam)

	params := &cloudwatch.ListMetricsInput{
		Namespace:  aws.String(reqParam.Parameters.Namespace),
		MetricName: aws.String(reqParam.Parameters.MetricName),
		Dimensions: reqParam.Parameters.Dimensions,
	}

	var resp cloudwatch.ListMetricsOutput
	err := svc.ListMetricsPages(params,
		func(page *cloudwatch.ListMetricsOutput, lastPage bool) bool {
			metrics, _ := awsutil.ValuesAtPath(page, "Metrics")
			for _, metric := range metrics {
				resp.Metrics = append(resp.Metrics, metric.(*cloudwatch.Metric))
			}
			return !lastPage
		})
	if err != nil {
		c.JsonApiErr(500, "Unable to call AWS API", err)
		return
	}

	c.JSON(200, resp)
}

func handleDescribeAlarms(req *cwRequest, c *middleware.Context) {
	cfg := getAwsConfig(req)
	svc := cloudwatch.New(session.New(cfg), cfg)

	reqParam := &struct {
		Parameters struct {
			ActionPrefix    string    `json:"actionPrefix"`
			AlarmNamePrefix string    `json:"alarmNamePrefix"`
			AlarmNames      []*string `json:"alarmNames"`
			StateValue      string    `json:"stateValue"`
		} `json:"parameters"`
	}{}
	json.Unmarshal(req.Body, reqParam)

	params := &cloudwatch.DescribeAlarmsInput{
		MaxRecords: aws.Int64(100),
	}
	if reqParam.Parameters.ActionPrefix != "" {
		params.ActionPrefix = aws.String(reqParam.Parameters.ActionPrefix)
	}
	if reqParam.Parameters.AlarmNamePrefix != "" {
		params.AlarmNamePrefix = aws.String(reqParam.Parameters.AlarmNamePrefix)
	}
	if len(reqParam.Parameters.AlarmNames) != 0 {
		params.AlarmNames = reqParam.Parameters.AlarmNames
	}
	if reqParam.Parameters.StateValue != "" {
		params.StateValue = aws.String(reqParam.Parameters.StateValue)
	}

	resp, err := svc.DescribeAlarms(params)
	if err != nil {
		c.JsonApiErr(500, "Unable to call AWS API", err)
		return
	}

	c.JSON(200, resp)
}

func handleDescribeAlarmsForMetric(req *cwRequest, c *middleware.Context) {
	cfg := getAwsConfig(req)
	svc := cloudwatch.New(session.New(cfg), cfg)

	reqParam := &struct {
		Parameters struct {
			Namespace         string                  `json:"namespace"`
			MetricName        string                  `json:"metricName"`
			Dimensions        []*cloudwatch.Dimension `json:"dimensions"`
			Statistic         string                  `json:"statistic"`
			ExtendedStatistic string                  `json:"extendedStatistic"`
			Period            int64                   `json:"period"`
		} `json:"parameters"`
	}{}
	json.Unmarshal(req.Body, reqParam)

	params := &cloudwatch.DescribeAlarmsForMetricInput{
		Namespace:  aws.String(reqParam.Parameters.Namespace),
		MetricName: aws.String(reqParam.Parameters.MetricName),
		Period:     aws.Int64(reqParam.Parameters.Period),
	}
	if len(reqParam.Parameters.Dimensions) != 0 {
		params.Dimensions = reqParam.Parameters.Dimensions
	}
	if reqParam.Parameters.Statistic != "" {
		params.Statistic = aws.String(reqParam.Parameters.Statistic)
	}
	if reqParam.Parameters.ExtendedStatistic != "" {
		params.ExtendedStatistic = aws.String(reqParam.Parameters.ExtendedStatistic)
	}

	resp, err := svc.DescribeAlarmsForMetric(params)
	if err != nil {
		c.JsonApiErr(500, "Unable to call AWS API", err)
		return
	}

	c.JSON(200, resp)
}

func handleDescribeAlarmHistory(req *cwRequest, c *middleware.Context) {
	cfg := getAwsConfig(req)
	svc := cloudwatch.New(session.New(cfg), cfg)

	reqParam := &struct {
		Parameters struct {
			AlarmName       string `json:"alarmName"`
			HistoryItemType string `json:"historyItemType"`
			StartDate       int64  `json:"startDate"`
			EndDate         int64  `json:"endDate"`
		} `json:"parameters"`
	}{}
	json.Unmarshal(req.Body, reqParam)

	params := &cloudwatch.DescribeAlarmHistoryInput{
		AlarmName: aws.String(reqParam.Parameters.AlarmName),
		StartDate: aws.Time(time.Unix(reqParam.Parameters.StartDate, 0)),
		EndDate:   aws.Time(time.Unix(reqParam.Parameters.EndDate, 0)),
	}
	if reqParam.Parameters.HistoryItemType != "" {
		params.HistoryItemType = aws.String(reqParam.Parameters.HistoryItemType)
	}

	resp, err := svc.DescribeAlarmHistory(params)
	if err != nil {
		c.JsonApiErr(500, "Unable to call AWS API", err)
		return
	}

	c.JSON(200, resp)
}

func handleDescribeInstances(req *cwRequest, c *middleware.Context) {
	cfg := getAwsConfig(req)
	svc := ec2.New(session.New(cfg), cfg)

	reqParam := &struct {
		Parameters struct {
			Filters     []*ec2.Filter `json:"filters"`
			InstanceIds []*string     `json:"instanceIds"`
		} `json:"parameters"`
	}{}
	json.Unmarshal(req.Body, reqParam)

	params := &ec2.DescribeInstancesInput{}
	if len(reqParam.Parameters.Filters) > 0 {
		params.Filters = reqParam.Parameters.Filters
	}
	if len(reqParam.Parameters.InstanceIds) > 0 {
		params.InstanceIds = reqParam.Parameters.InstanceIds
	}

	var resp ec2.DescribeInstancesOutput
	err := svc.DescribeInstancesPages(params,
		func(page *ec2.DescribeInstancesOutput, lastPage bool) bool {
			reservations, _ := awsutil.ValuesAtPath(page, "Reservations")
			for _, reservation := range reservations {
				resp.Reservations = append(resp.Reservations, reservation.(*ec2.Reservation))
			}
			return !lastPage
		})
	if err != nil {
		c.JsonApiErr(500, "Unable to call AWS API", err)
		return
	}

	c.JSON(200, resp)
}

func HandleRequest(c *middleware.Context, ds *m.DataSource) {
	var req cwRequest
	req.Body, _ = ioutil.ReadAll(c.Req.Request.Body)
	req.DataSource = ds
	json.Unmarshal(req.Body, &req)

	if handler, found := actionHandlers[req.Action]; !found {
		c.JsonApiErr(500, "Unexpected AWS Action", errors.New(req.Action))
		return
	} else {
		handler(&req, c)
	}
}
<commit_msg>fix cloudwatch percentile (#6989)<commit_after>package cloudwatch

import (
	"encoding/json"
	"errors"
	"io/ioutil"
	"strings"
	"sync"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awsutil"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/credentials/ec2rolecreds"
	"github.com/aws/aws-sdk-go/aws/ec2metadata"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/cloudwatch"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/aws/aws-sdk-go/service/sts"
	"github.com/grafana/grafana/pkg/log"
	"github.com/grafana/grafana/pkg/middleware"
	m "github.com/grafana/grafana/pkg/models"
)

type actionHandler func(*cwRequest, *middleware.Context)

var actionHandlers map[string]actionHandler

type cwRequest struct {
	Region     string `json:"region"`
	Action     string `json:"action"`
	Body       []byte `json:"-"`
	DataSource *m.DataSource
}

type datasourceInfo struct {
	Profile       string
	Region        string
	AssumeRoleArn string
	Namespace     string

	AccessKey string
	SecretKey string
}

func (req *cwRequest) GetDatasourceInfo() *datasourceInfo {
	assumeRoleArn := req.DataSource.JsonData.Get("assumeRoleArn").MustString()
	accessKey := ""
	secretKey := ""

	for key, value := range req.DataSource.SecureJsonData.Decrypt() {
		if key == "accessKey" {
			accessKey = value
		}
		if key == "secretKey" {
			secretKey = value
		}
	}

	return &datasourceInfo{
		AssumeRoleArn: assumeRoleArn,
		Region:        req.Region,
		Profile:       req.DataSource.Database,
		AccessKey:     accessKey,
		SecretKey:     secretKey,
	}
}

func init() {
	actionHandlers = map[string]actionHandler{
		"GetMetricStatistics":     handleGetMetricStatistics,
		"ListMetrics":             handleListMetrics,
		"DescribeAlarms":          handleDescribeAlarms,
		"DescribeAlarmsForMetric": handleDescribeAlarmsForMetric,
		"DescribeAlarmHistory":    handleDescribeAlarmHistory,
		"DescribeInstances":       handleDescribeInstances,
		"__GetRegions":            handleGetRegions,
		"__GetNamespaces":         handleGetNamespaces,
		"__GetMetrics":            handleGetMetrics,
		"__GetDimensions":         handleGetDimensions,
	}
}

type cache struct {
	credential *credentials.Credentials
	expiration *time.Time
}

var awsCredentialCache map[string]cache = make(map[string]cache)
var credentialCacheLock sync.RWMutex

func getCredentials(dsInfo *datasourceInfo) *credentials.Credentials {
	cacheKey := dsInfo.Profile + ":" + dsInfo.AssumeRoleArn
	credentialCacheLock.RLock()
	if _, ok := awsCredentialCache[cacheKey]; ok {
		if awsCredentialCache[cacheKey].expiration != nil &&
			(*awsCredentialCache[cacheKey].expiration).After(time.Now().UTC()) {
			result := awsCredentialCache[cacheKey].credential
			credentialCacheLock.RUnlock()
			return result
		}
	}
	credentialCacheLock.RUnlock()

	accessKeyId := ""
	secretAccessKey := ""
	sessionToken := ""
	var expiration *time.Time
	expiration = nil
	if strings.Index(dsInfo.AssumeRoleArn, "arn:aws:iam:") == 0 {
		params := &sts.AssumeRoleInput{
			RoleArn:         aws.String(dsInfo.AssumeRoleArn),
			RoleSessionName: aws.String("GrafanaSession"),
			DurationSeconds: aws.Int64(900),
		}

		stsSess := session.New()
		stsCreds := credentials.NewChainCredentials(
			[]credentials.Provider{
				&credentials.EnvProvider{},
				&credentials.SharedCredentialsProvider{Filename: "", Profile: dsInfo.Profile},
				&ec2rolecreds.EC2RoleProvider{Client: ec2metadata.New(stsSess), ExpiryWindow: 5 * time.Minute},
			})
		stsConfig := &aws.Config{
			Region:      aws.String(dsInfo.Region),
			Credentials: stsCreds,
		}

		svc := sts.New(session.New(stsConfig), stsConfig)
		resp, err := svc.AssumeRole(params)
		if err != nil {
			// ignore
			log.Error(3, "CloudWatch: Failed to assume role", err)
		}
		if resp.Credentials != nil {
			accessKeyId = *resp.Credentials.AccessKeyId
			secretAccessKey = *resp.Credentials.SecretAccessKey
			sessionToken = *resp.Credentials.SessionToken
			expiration = resp.Credentials.Expiration
		}
	}

	sess := session.New()
	creds := credentials.NewChainCredentials(
		[]credentials.Provider{
			&credentials.StaticProvider{Value: credentials.Value{
				AccessKeyID:     accessKeyId,
				SecretAccessKey: secretAccessKey,
				SessionToken:    sessionToken,
			}},
			&credentials.EnvProvider{},
			&credentials.StaticProvider{Value: credentials.Value{
				AccessKeyID:     dsInfo.AccessKey,
				SecretAccessKey: dsInfo.SecretKey,
			}},
			&credentials.SharedCredentialsProvider{Filename: "", Profile: dsInfo.Profile},
			&ec2rolecreds.EC2RoleProvider{Client: ec2metadata.New(sess), ExpiryWindow: 5 * time.Minute},
		})

	credentialCacheLock.Lock()
	awsCredentialCache[cacheKey] = cache{
		credential: creds,
		expiration: expiration,
	}
	credentialCacheLock.Unlock()

	return creds
}

func getAwsConfig(req *cwRequest) *aws.Config {
	cfg := &aws.Config{
		Region:      aws.String(req.Region),
		Credentials: getCredentials(req.GetDatasourceInfo()),
	}
	return cfg
}

func handleGetMetricStatistics(req *cwRequest, c *middleware.Context) {
	cfg := getAwsConfig(req)
	svc := cloudwatch.New(session.New(cfg), cfg)

	reqParam := &struct {
		Parameters struct {
			Namespace          string                  `json:"namespace"`
			MetricName         string                  `json:"metricName"`
			Dimensions         []*cloudwatch.Dimension `json:"dimensions"`
			Statistics         []*string               `json:"statistics"`
			ExtendedStatistics []*string               `json:"extendedStatistics"`
			StartTime          int64                   `json:"startTime"`
			EndTime            int64                   `json:"endTime"`
			Period             int64                   `json:"period"`
		} `json:"parameters"`
	}{}
	json.Unmarshal(req.Body, reqParam)

	params := &cloudwatch.GetMetricStatisticsInput{
		Namespace:  aws.String(reqParam.Parameters.Namespace),
		MetricName: aws.String(reqParam.Parameters.MetricName),
		Dimensions: reqParam.Parameters.Dimensions,
		StartTime:  aws.Time(time.Unix(reqParam.Parameters.StartTime, 0)),
		EndTime:    aws.Time(time.Unix(reqParam.Parameters.EndTime, 0)),
		Period:     aws.Int64(reqParam.Parameters.Period),
	}
	if len(reqParam.Parameters.Statistics) != 0 {
		params.Statistics = reqParam.Parameters.Statistics
	}
	if len(reqParam.Parameters.ExtendedStatistics) != 0 {
		params.ExtendedStatistics = reqParam.Parameters.ExtendedStatistics
	}

	resp, err := svc.GetMetricStatistics(params)
	if err != nil {
		c.JsonApiErr(500, "Unable to call AWS API", err)
		return
	}

	c.JSON(200, resp)
}

func handleListMetrics(req *cwRequest, c *middleware.Context) {
	cfg := getAwsConfig(req)
	svc := cloudwatch.New(session.New(cfg), cfg)

	reqParam := &struct {
		Parameters struct {
			Namespace  string                        `json:"namespace"`
			MetricName string                        `json:"metricName"`
			Dimensions []*cloudwatch.DimensionFilter `json:"dimensions"`
		} `json:"parameters"`
	}{}
	json.Unmarshal(req.Body, reqParam)

	params := &cloudwatch.ListMetricsInput{
		Namespace:  aws.String(reqParam.Parameters.Namespace),
		MetricName: aws.String(reqParam.Parameters.MetricName),
		Dimensions: reqParam.Parameters.Dimensions,
	}

	var resp cloudwatch.ListMetricsOutput
	err := svc.ListMetricsPages(params,
		func(page *cloudwatch.ListMetricsOutput, lastPage bool) bool {
			metrics, _ := awsutil.ValuesAtPath(page, "Metrics")
			for _, metric := range metrics {
				resp.Metrics = append(resp.Metrics, metric.(*cloudwatch.Metric))
			}
			return !lastPage
		})
	if err != nil {
		c.JsonApiErr(500, "Unable to call AWS API", err)
		return
	}

	c.JSON(200, resp)
}

func handleDescribeAlarms(req *cwRequest, c *middleware.Context) {
	cfg := getAwsConfig(req)
	svc := cloudwatch.New(session.New(cfg), cfg)

	reqParam := &struct {
		Parameters struct {
			ActionPrefix    string    `json:"actionPrefix"`
			AlarmNamePrefix string    `json:"alarmNamePrefix"`
			AlarmNames      []*string `json:"alarmNames"`
			StateValue      string    `json:"stateValue"`
		} `json:"parameters"`
	}{}
	json.Unmarshal(req.Body, reqParam)

	params := &cloudwatch.DescribeAlarmsInput{
		MaxRecords: aws.Int64(100),
	}
	if reqParam.Parameters.ActionPrefix != "" {
		params.ActionPrefix = aws.String(reqParam.Parameters.ActionPrefix)
	}
	if reqParam.Parameters.AlarmNamePrefix != "" {
		params.AlarmNamePrefix = aws.String(reqParam.Parameters.AlarmNamePrefix)
	}
	if len(reqParam.Parameters.AlarmNames) != 0 {
		params.AlarmNames = reqParam.Parameters.AlarmNames
	}
	if reqParam.Parameters.StateValue != "" {
		params.StateValue = aws.String(reqParam.Parameters.StateValue)
	}

	resp, err := svc.DescribeAlarms(params)
	if err != nil {
		c.JsonApiErr(500, "Unable to call AWS API", err)
		return
	}

	c.JSON(200, resp)
}

func handleDescribeAlarmsForMetric(req *cwRequest, c *middleware.Context) {
	cfg := getAwsConfig(req)
	svc := cloudwatch.New(session.New(cfg), cfg)

	reqParam := &struct {
		Parameters struct {
			Namespace         string                  `json:"namespace"`
			MetricName        string                  `json:"metricName"`
			Dimensions        []*cloudwatch.Dimension `json:"dimensions"`
			Statistic         string                  `json:"statistic"`
			ExtendedStatistic string                  `json:"extendedStatistic"`
			Period            int64                   `json:"period"`
		} `json:"parameters"`
	}{}
	json.Unmarshal(req.Body, reqParam)

	params := &cloudwatch.DescribeAlarmsForMetricInput{
		Namespace:  aws.String(reqParam.Parameters.Namespace),
		MetricName: aws.String(reqParam.Parameters.MetricName),
		Period:     aws.Int64(reqParam.Parameters.Period),
	}
	if len(reqParam.Parameters.Dimensions) != 0 {
		params.Dimensions = reqParam.Parameters.Dimensions
	}
	if reqParam.Parameters.Statistic != "" {
		params.Statistic = aws.String(reqParam.Parameters.Statistic)
	}
	if reqParam.Parameters.ExtendedStatistic != "" {
		params.ExtendedStatistic = aws.String(reqParam.Parameters.ExtendedStatistic)
	}

	resp, err := svc.DescribeAlarmsForMetric(params)
	if err != nil {
		c.JsonApiErr(500, "Unable to call AWS API", err)
		return
	}

	c.JSON(200, resp)
}

func handleDescribeAlarmHistory(req *cwRequest, c *middleware.Context) {
	cfg := getAwsConfig(req)
	svc := cloudwatch.New(session.New(cfg), cfg)

	reqParam := &struct {
		Parameters struct {
			AlarmName       string `json:"alarmName"`
			HistoryItemType string `json:"historyItemType"`
			StartDate       int64  `json:"startDate"`
			EndDate         int64  `json:"endDate"`
		} `json:"parameters"`
	}{}
	json.Unmarshal(req.Body, reqParam)

	params := &cloudwatch.DescribeAlarmHistoryInput{
		AlarmName: aws.String(reqParam.Parameters.AlarmName),
		StartDate: aws.Time(time.Unix(reqParam.Parameters.StartDate, 0)),
		EndDate:   aws.Time(time.Unix(reqParam.Parameters.EndDate, 0)),
	}
	if reqParam.Parameters.HistoryItemType != "" {
		params.HistoryItemType = aws.String(reqParam.Parameters.HistoryItemType)
	}

	resp, err := svc.DescribeAlarmHistory(params)
	if err != nil {
		c.JsonApiErr(500, "Unable to call AWS API", err)
		return
	}

	c.JSON(200, resp)
}

func handleDescribeInstances(req *cwRequest, c *middleware.Context) {
	cfg := getAwsConfig(req)
	svc := ec2.New(session.New(cfg), cfg)

	reqParam := &struct {
		Parameters struct {
			Filters     []*ec2.Filter `json:"filters"`
			InstanceIds []*string     `json:"instanceIds"`
		} `json:"parameters"`
	}{}
	json.Unmarshal(req.Body, reqParam)

	params := &ec2.DescribeInstancesInput{}
	if len(reqParam.Parameters.Filters) > 0 {
		params.Filters = reqParam.Parameters.Filters
	}
	if len(reqParam.Parameters.InstanceIds) > 0 {
		params.InstanceIds = reqParam.Parameters.InstanceIds
	}

	var resp ec2.DescribeInstancesOutput
	err := svc.DescribeInstancesPages(params,
		func(page *ec2.DescribeInstancesOutput, lastPage bool) bool {
			reservations, _ := awsutil.ValuesAtPath(page, "Reservations")
			for _, reservation := range reservations {
				resp.Reservations = append(resp.Reservations, reservation.(*ec2.Reservation))
			}
			return !lastPage
		})
	if err != nil {
		c.JsonApiErr(500, "Unable to call AWS API", err)
		return
	}

	c.JSON(200, resp)
}

func HandleRequest(c *middleware.Context, ds *m.DataSource) {
	var req cwRequest
	req.Body, _ = ioutil.ReadAll(c.Req.Request.Body)
	req.DataSource = ds
	json.Unmarshal(req.Body, &req)

	if handler, found := actionHandlers[req.Action]; !found {
		c.JsonApiErr(500, "Unexpected AWS Action", errors.New(req.Action))
		return
	} else {
		handler(&req, c)
	}
}
<|endoftext|>


<commit_before>// Copyright 2020 Google LLC
//
// 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 issueapi

import (
	"context"
	"errors"
	"fmt"
	"net/http"
	"strings"
	"time"

	"github.com/google/exposure-notifications-server/pkg/timeutils"
	"github.com/google/exposure-notifications-verification-server/pkg/api"
	"github.com/google/exposure-notifications-verification-server/pkg/controller"
	"github.com/google/exposure-notifications-verification-server/pkg/database"
	"github.com/google/exposure-notifications-verification-server/pkg/observability"
	"github.com/google/exposure-notifications-verification-server/pkg/otp"
	"github.com/google/exposure-notifications-verification-server/pkg/sms"

	"go.opencensus.io/stats"
	"go.opencensus.io/tag"
)

type dateParseSettings struct {
	Name          string
	ParseError    string
	ValidateError string
}

var (
	// Cache the UTC time.Location, to speed runtime.
	utc *time.Location

	onsetSettings = dateParseSettings{
		Name:          "symptom onset",
		ParseError:    "FAILED_TO_PROCESS_SYMPTOM_ONSET_DATE",
		ValidateError: "SYMPTOM_ONSET_DATE_NOT_IN_VALID_RANGE",
	}
	testSettings = dateParseSettings{
		Name:          "test",
		ParseError:    "FAILED_TO_PROCESS_TEST_DATE",
		ValidateError: "TEST_DATE_NOT_IN_VALID_RANGE",
	}
)

func init() {
	var err error
	utc, err = time.LoadLocation("UTC")
	if err != nil {
		panic("should have found UTC")
	}
}

// validateDate validates the date given -- returning the time or an error.
func validateDate(date, minDate, maxDate time.Time, tzOffset int) (*time.Time, error) {
	// Check that all our dates are utc.
	if date.Location() != utc || minDate.Location() != utc || maxDate.Location() != utc {
		return nil, errors.New("dates weren't in UTC")
	}

	// If we're dealing with a timezone where the offset is earlier than this one,
	// we loosen up the lower bound. We might have the following circumstance:
	//
	//    Server time: UTC Aug 1, 12:01 AM
	//    Client time: UTC July 30, 11:01 PM (ie, tzOffset = -30)
	//
	// In this circumstance, we'll have the following:
	//
	//    minTime: UTC July 31, maxTime: Aug 1, clientTime: July 30.
	//
	// which would be an error. Loosening up the lower bound, by a day, keeps us
	// all ok.
	if tzOffset < 0 {
		if m := minDate.Add(-24 * time.Hour); m.After(date) {
			return nil, fmt.Errorf("date %v before min %v", date, m)
		}
	} else if minDate.After(date) {
		return nil, fmt.Errorf("date %v before min %v", date, minDate)
	}
	if date.After(maxDate) {
		return nil, fmt.Errorf("date %v after max %v", date, maxDate)
	}
	return &date, nil
}

func (c *Controller) HandleIssue() http.Handler {
	logger := c.logger.Named("issueapi.HandleIssue")

	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		ctx := observability.WithBuildInfo(r.Context())

		var blame = observability.BlameNone
		var result = observability.ResultOK()

		defer func(blame, result *tag.Mutator) {
			ctx, err := tag.New(ctx, *blame, *result)
			if err != nil {
				c.logger.Warnw("failed to create context with additional tags", "error", err)
				// NOTE: do not return here. We should log it as success.
			}
			stats.Record(ctx, mRequest.M(1))
		}(&blame, &result)

		var request api.IssueCodeRequest
		if err := controller.BindJSON(w, r, &request); err != nil {
			c.h.RenderJSON(w, http.StatusBadRequest, api.Error(err))
			blame = observability.BlameClient
			result = observability.ResultError("FAILED_TO_PARSE_JSON_REQUEST")
			return
		}

		authApp, user, err := c.getAuthorizationFromContext(r)
		if err != nil {
			c.h.RenderJSON(w, http.StatusUnauthorized, api.Error(err))
			blame = observability.BlameClient
			result = observability.ResultError("MISSING_AUTHORIZED_APP")
			return
		}

		var realm *database.Realm
		if authApp != nil {
			realm, err = authApp.Realm(c.db)
			if err != nil {
				c.h.RenderJSON(w, http.StatusUnauthorized, nil)
				blame = observability.BlameClient
				result = observability.ResultError("UNAUTHORIZED")
				return
			}
		} else {
			// if it's a user logged in, we can pull realm from the context.
			realm = controller.RealmFromContext(ctx)
		}
		if realm == nil {
			c.h.RenderJSON(w, http.StatusBadRequest, api.Errorf("missing realm"))
			blame = observability.BlameServer
			result = observability.ResultError("MISSING_REALM")
			return
		}

		// Add realm so that metrics are groupable on a per-realm basis.
		ctx = observability.WithRealmID(ctx, realm.ID)

		// If this realm requires a date but no date was specified, return an error.
		if realm.RequireDate && request.SymptomDate == "" && request.TestDate == "" {
			c.h.RenderJSON(w, http.StatusBadRequest, api.Errorf("missing either test or symptom date").WithCode(api.ErrMissingDate))
			blame = observability.BlameClient
			result = observability.ResultError("MISSING_REQUIRED_FIELDS")
			return
		}

		// Validate that the request with the provided test type is valid for this realm.
		if !realm.ValidTestType(request.TestType) {
			c.h.RenderJSON(w, http.StatusBadRequest,
				api.Errorf("unsupported test type: %v", request.TestType))
			blame = observability.BlameClient
			result = observability.ResultError("UNSUPPORTED_TEST_TYPE")
			return
		}

		// Verify SMS configuration if phone was provided
		var smsProvider sms.Provider
		if request.Phone != "" {
			smsProvider, err = realm.SMSProvider(c.db)
			if err != nil {
				logger.Errorw("failed to get sms provider", "error", err)
				c.h.RenderJSON(w, http.StatusInternalServerError, api.Errorf("failed to get sms provider"))
				blame = observability.BlameServer
				result = observability.ResultError("FAILED_TO_GET_SMS_PROVIDER")
				return
			}
			if smsProvider == nil {
				err := fmt.Errorf("phone provided, but no sms provider is configured")
				c.h.RenderJSON(w, http.StatusBadRequest, api.Error(err))
				blame = observability.BlameServer
				result = observability.ResultError("FAILED_TO_GET_SMS_PROVIDER")
				return
			}
		}

		// Verify the test type
		request.TestType = strings.ToLower(request.TestType)
		if _, ok := c.validTestType[request.TestType]; !ok {
			c.h.RenderJSON(w, http.StatusBadRequest, api.Errorf("invalid test type"))
			blame = observability.BlameClient
			result = observability.ResultError("INVALID_TEST_TYPE")
			return
		}

		// Set up parallel arrays to leverage the observability reporting and connect the parse / valdiation errors
		// to the correct date.
		parsedDates := make([]*time.Time, 2)
		input := []string{request.SymptomDate, request.TestDate}
		dateSettings := []*dateParseSettings{&onsetSettings, &testSettings}
		for i, d := range input {
			if d != "" {
				if parsed, err := time.Parse("2006-01-02", d); err != nil {
					c.h.RenderJSON(w, http.StatusBadRequest, api.Errorf("failed to process %s date: %v", dateSettings[i].Name, err))
					blame = observability.BlameClient
					result = observability.ResultError(dateSettings[i].ParseError)
					return
				} else {
					// Max date is today (UTC time) and min date is AllowedTestAge ago, truncated.
					maxDate := timeutils.UTCMidnight(time.Now())
					minDate := timeutils.Midnight(maxDate.Add(-1 * c.config.GetAllowedSymptomAge()))

					validatedDate, err := validateDate(parsed, minDate, maxDate, int(request.TZOffset))
					if err != nil {
						err := fmt.Errorf("%s date must be on/after %v and on/before %v %v",
							dateSettings[i].Name,
							minDate.Format("2006-01-02"),
							maxDate.Format("2006-01-02"),
							parsed.Format("2006-01-02"),
						)
						c.h.RenderJSON(w, http.StatusBadRequest, api.Error(err))
						blame = observability.BlameClient
						result = observability.ResultError(dateSettings[i].ValidateError)
						return
					}
					parsedDates[i] = validatedDate
				}
			}
		}

		// If we got this far, we're about to issue a code - take from the limiter
		// to ensure this is permitted.
		if realm.AbusePreventionEnabled {
			key, err := realm.QuotaKey(c.config.GetRateLimitConfig().HMACKey)
			if err != nil {
				controller.InternalError(w, r, c.h, err)
				blame = observability.BlameServer
				result = observability.ResultError("FAILED_TO_GENERATE_HMAC")
				return
			}
			limit, remaining, reset, ok, err := c.limiter.Take(ctx, key)
			if err != nil {
				logger.Errorw("failed to take from limiter", "error", err)
				c.h.RenderJSON(w, http.StatusInternalServerError, api.Errorf("failed to verify realm stats, please try again"))
				blame = observability.BlameServer
				result = observability.ResultError("FAILED_TO_TAKE_FROM_LIMITER")
				return
			}
			c.recordCapacity(ctx, limit, remaining)
			if !ok {
				logger.Warnw("realm has exceeded daily quota",
					"realm", realm.ID,
					"limit", limit,
					"reset", reset)

				if c.config.GetEnforceRealmQuotas() {
					c.h.RenderJSON(w, http.StatusTooManyRequests, api.Errorf("exceeded realm quota"))
					blame = observability.BlameClient
					result = observability.ResultError("QUOTA_EXCEEDED")
					return
				}
			}
		}

		now := time.Now().UTC()
		expiryTime := now.Add(realm.CodeDuration.Duration)
		longExpiryTime := now.Add(realm.LongCodeDuration.Duration)
		if request.Phone == "" || smsProvider == nil {
			// If this isn't going to be send via SMS, make the long code expiration time same as short.
			// This is because the long code will never be shown or sent.
			longExpiryTime = expiryTime
		}

		// Generate verification code
		codeRequest := otp.Request{
			DB:             c.db,
			ShortLength:    realm.CodeLength,
			ShortExpiresAt: expiryTime,
			LongLength:     realm.LongCodeLength,
			LongExpiresAt:  longExpiryTime,
			TestType:       request.TestType,
			SymptomDate:    parsedDates[0],
			TestDate:       parsedDates[1],
			MaxSymptomAge:  c.config.GetAllowedSymptomAge(),
			IssuingUser:    user,
			IssuingApp:     authApp,
			RealmID:        realm.ID,
		}

		code, longCode, uuid, err := codeRequest.Issue(ctx, c.config.GetCollisionRetryCount())
		if err != nil {
			logger.Errorw("failed to issue code", "error", err)
			c.h.RenderJSON(w, http.StatusInternalServerError, api.Errorf("failed to generate otp code, please try again"))
			blame = observability.BlameServer
			result = observability.ResultError("FAILED_TO_ISSUE_CODE")
			return
		}

		if request.Phone != "" && smsProvider != nil {
			message := realm.BuildSMSText(code, longCode, c.config.GetENXRedirectDomain())
			if err := smsProvider.SendSMS(ctx, request.Phone, message); err != nil {
				// Delete the token
				if err := c.db.DeleteVerificationCode(code); err != nil {
					logger.Errorw("failed to delete verification code", "error", err)
					// fallthrough to the error
				}

				logger.Errorw("failed to send sms", "error", err)
				stats.RecordWithTags(ctx, []tag.Mutator{observability.ResultNotOK()}, mSMSRequest.M(1))
				c.h.RenderJSON(w, http.StatusInternalServerError, api.Errorf("failed to send sms"))
				blame = observability.BlameServer
				result = observability.ResultError("FAILED_TO_SEND_SMS")
				return
			}
			stats.RecordWithTags(ctx, []tag.Mutator{observability.ResultOK()}, mSMSRequest.M(1))
		}

		c.h.RenderJSON(w, http.StatusOK,
			&api.IssueCodeResponse{
				UUID:                   uuid,
				VerificationCode:       code,
				ExpiresAt:              expiryTime.Format(time.RFC1123),
				ExpiresAtTimestamp:     expiryTime.UTC().Unix(),
				LongExpiresAt:          longExpiryTime.Format(time.RFC1123),
				LongExpiresAtTimestamp: longExpiryTime.UTC().Unix(),
			})
	})
}

func (c *Controller) getAuthorizationFromContext(r *http.Request) (*database.AuthorizedApp, *database.User, error) {
	ctx := r.Context()

	authorizedApp := controller.AuthorizedAppFromContext(ctx)
	currentUser := controller.UserFromContext(ctx)

	if authorizedApp == nil && currentUser == nil {
		return nil, nil, fmt.Errorf("unable to identify authorized requestor")
	}

	return authorizedApp, currentUser, nil
}

func (c *Controller) recordCapacity(ctx context.Context, limit, remaining uint64) {
	stats.Record(ctx, mRealmTokenRemaining.M(int64(remaining)))

	issued := uint64(limit) - remaining

	capacity := float64(issued) / float64(limit)
	stats.Record(ctx, mRealmTokenCapacity.M(capacity))

	stats.Record(ctx, mRealmTokenUsed.M(1))
	stats.RecordWithTags(ctx, []tag.Mutator{tokenAvailableTag()}, mRealmToken.M(int64(remaining)))
	stats.RecordWithTags(ctx, []tag.Mutator{tokenLimitTag()}, mRealmToken.M(int64(limit)))
}
<commit_msg>Try and fix use case of burst capacity (#913)<commit_after>// Copyright 2020 Google LLC
//
// 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 issueapi

import (
	"context"
	"errors"
	"fmt"
	"net/http"
	"strings"
	"time"

	"github.com/google/exposure-notifications-server/pkg/timeutils"
	"github.com/google/exposure-notifications-verification-server/pkg/api"
	"github.com/google/exposure-notifications-verification-server/pkg/controller"
	"github.com/google/exposure-notifications-verification-server/pkg/database"
	"github.com/google/exposure-notifications-verification-server/pkg/observability"
	"github.com/google/exposure-notifications-verification-server/pkg/otp"
	"github.com/google/exposure-notifications-verification-server/pkg/sms"

	"go.opencensus.io/stats"
	"go.opencensus.io/tag"
)

type dateParseSettings struct {
	Name          string
	ParseError    string
	ValidateError string
}

var (
	// Cache the UTC time.Location, to speed runtime.
	utc *time.Location

	onsetSettings = dateParseSettings{
		Name:          "symptom onset",
		ParseError:    "FAILED_TO_PROCESS_SYMPTOM_ONSET_DATE",
		ValidateError: "SYMPTOM_ONSET_DATE_NOT_IN_VALID_RANGE",
	}
	testSettings = dateParseSettings{
		Name:          "test",
		ParseError:    "FAILED_TO_PROCESS_TEST_DATE",
		ValidateError: "TEST_DATE_NOT_IN_VALID_RANGE",
	}
)

func init() {
	var err error
	utc, err = time.LoadLocation("UTC")
	if err != nil {
		panic("should have found UTC")
	}
}

// validateDate validates the date given -- returning the time or an error.
func validateDate(date, minDate, maxDate time.Time, tzOffset int) (*time.Time, error) {
	// Check that all our dates are utc.
	if date.Location() != utc || minDate.Location() != utc || maxDate.Location() != utc {
		return nil, errors.New("dates weren't in UTC")
	}

	// If we're dealing with a timezone where the offset is earlier than this one,
	// we loosen up the lower bound. We might have the following circumstance:
	//
	//    Server time: UTC Aug 1, 12:01 AM
	//    Client time: UTC July 30, 11:01 PM (ie, tzOffset = -30)
	//
	// In this circumstance, we'll have the following:
	//
	//    minTime: UTC July 31, maxTime: Aug 1, clientTime: July 30.
	//
	// which would be an error. Loosening up the lower bound, by a day, keeps us
	// all ok.
	if tzOffset < 0 {
		if m := minDate.Add(-24 * time.Hour); m.After(date) {
			return nil, fmt.Errorf("date %v before min %v", date, m)
		}
	} else if minDate.After(date) {
		return nil, fmt.Errorf("date %v before min %v", date, minDate)
	}
	if date.After(maxDate) {
		return nil, fmt.Errorf("date %v after max %v", date, maxDate)
	}
	return &date, nil
}

func (c *Controller) HandleIssue() http.Handler {
	logger := c.logger.Named("issueapi.HandleIssue")

	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		ctx := observability.WithBuildInfo(r.Context())

		var blame = observability.BlameNone
		var result = observability.ResultOK()

		defer func(blame, result *tag.Mutator) {
			ctx, err := tag.New(ctx, *blame, *result)
			if err != nil {
				c.logger.Warnw("failed to create context with additional tags", "error", err)
				// NOTE: do not return here. We should log it as success.
			}
			stats.Record(ctx, mRequest.M(1))
		}(&blame, &result)

		var request api.IssueCodeRequest
		if err := controller.BindJSON(w, r, &request); err != nil {
			c.h.RenderJSON(w, http.StatusBadRequest, api.Error(err))
			blame = observability.BlameClient
			result = observability.ResultError("FAILED_TO_PARSE_JSON_REQUEST")
			return
		}

		authApp, user, err := c.getAuthorizationFromContext(r)
		if err != nil {
			c.h.RenderJSON(w, http.StatusUnauthorized, api.Error(err))
			blame = observability.BlameClient
			result = observability.ResultError("MISSING_AUTHORIZED_APP")
			return
		}

		var realm *database.Realm
		if authApp != nil {
			realm, err = authApp.Realm(c.db)
			if err != nil {
				c.h.RenderJSON(w, http.StatusUnauthorized, nil)
				blame = observability.BlameClient
				result = observability.ResultError("UNAUTHORIZED")
				return
			}
		} else {
			// if it's a user logged in, we can pull realm from the context.
			realm = controller.RealmFromContext(ctx)
		}
		if realm == nil {
			c.h.RenderJSON(w, http.StatusBadRequest, api.Errorf("missing realm"))
			blame = observability.BlameServer
			result = observability.ResultError("MISSING_REALM")
			return
		}

		// Add realm so that metrics are groupable on a per-realm basis.
		ctx = observability.WithRealmID(ctx, realm.ID)

		// If this realm requires a date but no date was specified, return an error.
		if realm.RequireDate && request.SymptomDate == "" && request.TestDate == "" {
			c.h.RenderJSON(w, http.StatusBadRequest, api.Errorf("missing either test or symptom date").WithCode(api.ErrMissingDate))
			blame = observability.BlameClient
			result = observability.ResultError("MISSING_REQUIRED_FIELDS")
			return
		}

		// Validate that the request with the provided test type is valid for this realm.
		if !realm.ValidTestType(request.TestType) {
			c.h.RenderJSON(w, http.StatusBadRequest,
				api.Errorf("unsupported test type: %v", request.TestType))
			blame = observability.BlameClient
			result = observability.ResultError("UNSUPPORTED_TEST_TYPE")
			return
		}

		// Verify SMS configuration if phone was provided
		var smsProvider sms.Provider
		if request.Phone != "" {
			smsProvider, err = realm.SMSProvider(c.db)
			if err != nil {
				logger.Errorw("failed to get sms provider", "error", err)
				c.h.RenderJSON(w, http.StatusInternalServerError, api.Errorf("failed to get sms provider"))
				blame = observability.BlameServer
				result = observability.ResultError("FAILED_TO_GET_SMS_PROVIDER")
				return
			}
			if smsProvider == nil {
				err := fmt.Errorf("phone provided, but no sms provider is configured")
				c.h.RenderJSON(w, http.StatusBadRequest, api.Error(err))
				blame = observability.BlameServer
				result = observability.ResultError("FAILED_TO_GET_SMS_PROVIDER")
				return
			}
		}

		// Verify the test type
		request.TestType = strings.ToLower(request.TestType)
		if _, ok := c.validTestType[request.TestType]; !ok {
			c.h.RenderJSON(w, http.StatusBadRequest, api.Errorf("invalid test type"))
			blame = observability.BlameClient
			result = observability.ResultError("INVALID_TEST_TYPE")
			return
		}

		// Set up parallel arrays to leverage the observability reporting and connect the parse / valdiation errors
		// to the correct date.
		parsedDates := make([]*time.Time, 2)
		input := []string{request.SymptomDate, request.TestDate}
		dateSettings := []*dateParseSettings{&onsetSettings, &testSettings}
		for i, d := range input {
			if d != "" {
				if parsed, err := time.Parse("2006-01-02", d); err != nil {
					c.h.RenderJSON(w, http.StatusBadRequest, api.Errorf("failed to process %s date: %v", dateSettings[i].Name, err))
					blame = observability.BlameClient
					result = observability.ResultError(dateSettings[i].ParseError)
					return
				} else {
					// Max date is today (UTC time) and min date is AllowedTestAge ago, truncated.
					maxDate := timeutils.UTCMidnight(time.Now())
					minDate := timeutils.Midnight(maxDate.Add(-1 * c.config.GetAllowedSymptomAge()))

					validatedDate, err := validateDate(parsed, minDate, maxDate, int(request.TZOffset))
					if err != nil {
						err := fmt.Errorf("%s date must be on/after %v and on/before %v %v",
							dateSettings[i].Name,
							minDate.Format("2006-01-02"),
							maxDate.Format("2006-01-02"),
							parsed.Format("2006-01-02"),
						)
						c.h.RenderJSON(w, http.StatusBadRequest, api.Error(err))
						blame = observability.BlameClient
						result = observability.ResultError(dateSettings[i].ValidateError)
						return
					}
					parsedDates[i] = validatedDate
				}
			}
		}

		// If we got this far, we're about to issue a code - take from the limiter
		// to ensure this is permitted.
		if realm.AbusePreventionEnabled {
			key, err := realm.QuotaKey(c.config.GetRateLimitConfig().HMACKey)
			if err != nil {
				controller.InternalError(w, r, c.h, err)
				blame = observability.BlameServer
				result = observability.ResultError("FAILED_TO_GENERATE_HMAC")
				return
			}
			limit, remaining, reset, ok, err := c.limiter.Take(ctx, key)
			if err != nil {
				logger.Errorw("failed to take from limiter", "error", err)
				c.h.RenderJSON(w, http.StatusInternalServerError, api.Errorf("failed to verify realm stats, please try again"))
				blame = observability.BlameServer
				result = observability.ResultError("FAILED_TO_TAKE_FROM_LIMITER")
				return
			}

			// Override limit if there has been a burst for this realm.
			if remaining > limit {
				remaining = limit
			}

			c.recordCapacity(ctx, limit, remaining)
			if !ok {
				logger.Warnw("realm has exceeded daily quota",
					"realm", realm.ID,
					"limit", limit,
					"reset", reset)

				if c.config.GetEnforceRealmQuotas() {
					c.h.RenderJSON(w, http.StatusTooManyRequests, api.Errorf("exceeded realm quota"))
					blame = observability.BlameClient
					result = observability.ResultError("QUOTA_EXCEEDED")
					return
				}
			}
		}

		now := time.Now().UTC()
		expiryTime := now.Add(realm.CodeDuration.Duration)
		longExpiryTime := now.Add(realm.LongCodeDuration.Duration)
		if request.Phone == "" || smsProvider == nil {
			// If this isn't going to be send via SMS, make the long code expiration time same as short.
			// This is because the long code will never be shown or sent.
			longExpiryTime = expiryTime
		}

		// Generate verification code
		codeRequest := otp.Request{
			DB:             c.db,
			ShortLength:    realm.CodeLength,
			ShortExpiresAt: expiryTime,
			LongLength:     realm.LongCodeLength,
			LongExpiresAt:  longExpiryTime,
			TestType:       request.TestType,
			SymptomDate:    parsedDates[0],
			TestDate:       parsedDates[1],
			MaxSymptomAge:  c.config.GetAllowedSymptomAge(),
			IssuingUser:    user,
			IssuingApp:     authApp,
			RealmID:        realm.ID,
		}

		code, longCode, uuid, err := codeRequest.Issue(ctx, c.config.GetCollisionRetryCount())
		if err != nil {
			logger.Errorw("failed to issue code", "error", err)
			c.h.RenderJSON(w, http.StatusInternalServerError, api.Errorf("failed to generate otp code, please try again"))
			blame = observability.BlameServer
			result = observability.ResultError("FAILED_TO_ISSUE_CODE")
			return
		}

		if request.Phone != "" && smsProvider != nil {
			message := realm.BuildSMSText(code, longCode, c.config.GetENXRedirectDomain())
			if err := smsProvider.SendSMS(ctx, request.Phone, message); err != nil {
				// Delete the token
				if err := c.db.DeleteVerificationCode(code); err != nil {
					logger.Errorw("failed to delete verification code", "error", err)
					// fallthrough to the error
				}

				logger.Errorw("failed to send sms", "error", err)
				stats.RecordWithTags(ctx, []tag.Mutator{observability.ResultNotOK()}, mSMSRequest.M(1))
				c.h.RenderJSON(w, http.StatusInternalServerError, api.Errorf("failed to send sms"))
				blame = observability.BlameServer
				result = observability.ResultError("FAILED_TO_SEND_SMS")
				return
			}
			stats.RecordWithTags(ctx, []tag.Mutator{observability.ResultOK()}, mSMSRequest.M(1))
		}

		c.h.RenderJSON(w, http.StatusOK,
			&api.IssueCodeResponse{
				UUID:                   uuid,
				VerificationCode:       code,
				ExpiresAt:              expiryTime.Format(time.RFC1123),
				ExpiresAtTimestamp:     expiryTime.UTC().Unix(),
				LongExpiresAt:          longExpiryTime.Format(time.RFC1123),
				LongExpiresAtTimestamp: longExpiryTime.UTC().Unix(),
			})
	})
}

func (c *Controller) getAuthorizationFromContext(r *http.Request) (*database.AuthorizedApp, *database.User, error) {
	ctx := r.Context()

	authorizedApp := controller.AuthorizedAppFromContext(ctx)
	currentUser := controller.UserFromContext(ctx)

	if authorizedApp == nil && currentUser == nil {
		return nil, nil, fmt.Errorf("unable to identify authorized requestor")
	}

	return authorizedApp, currentUser, nil
}

func (c *Controller) recordCapacity(ctx context.Context, limit, remaining uint64) {
	stats.Record(ctx, mRealmTokenRemaining.M(int64(remaining)))

	issued := uint64(limit) - remaining

	capacity := float64(issued) / float64(limit)
	stats.Record(ctx, mRealmTokenCapacity.M(capacity))

	stats.Record(ctx, mRealmTokenUsed.M(1))
	stats.RecordWithTags(ctx, []tag.Mutator{tokenAvailableTag()}, mRealmToken.M(int64(remaining)))
	stats.RecordWithTags(ctx, []tag.Mutator{tokenLimitTag()}, mRealmToken.M(int64(limit)))
}
<|endoftext|>


<commit_before>/*
Copyright 2020 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 register

import "fmt"

// Log represents the different types of logs that can be output as JSON
// This includes: Step, Download, DownloadProgress, Warning, Info, Error
type Log interface {
	Type() string
}

// Step represents a normal step in minikube execution
type Step struct {
	data map[string]string
}

// Type returns the cloud events compatible type of this struct
func (s *Step) Type() string {
	return "io.k8s.sigs.minikube.step"
}

// NewStep returns a new step type
func NewStep(message string) *Step {
	return &Step{data: map[string]string{
		"totalsteps":  Reg.totalSteps(),
		"currentstep": Reg.currentStep(),
		"message":     message,
		"name":        string(Reg.current),
	}}
}

// Download will be used to notify the user that a download has begun
type Download struct {
	data map[string]string
}

// Type returns the cloud events compatible type of this struct
func (s *Download) Type() string {
	return "io.k8s.sigs.minikube.download"
}

// NewDownload returns a new download type
func NewDownload(artifact string) *Download {
	return &Download{data: map[string]string{
		"totalsteps":  Reg.totalSteps(),
		"currentstep": Reg.currentStep(),
		"artifact":    artifact,
	}}
}

// DownloadProgress will be used to notify the user around the progress of a download
type DownloadProgress struct {
	data map[string]string
}

// Type returns the cloud events compatible type of this struct
func (s *DownloadProgress) Type() string {
	return "io.k8s.sigs.minikube.download.progress"
}

// NewDownloadProgress returns a new download progress type
func NewDownloadProgress(artifact, progress string) *DownloadProgress {
	return &DownloadProgress{data: map[string]string{
		"totalsteps":  Reg.totalSteps(),
		"currentstep": Reg.currentStep(),
		"progress":    progress,
		"artifact":    artifact,
	}}
}

// Warning will be used to notify the user of warnings
type Warning struct {
	data map[string]string
}

// NewWarning returns a new warning type
func NewWarning(warning string) *Warning {
	return &Warning{
		map[string]string{
			"message": warning,
		},
	}
}

// Type returns the cloud events compatible type of this struct
func (s *Warning) Type() string {
	return "io.k8s.sigs.minikube.warning"
}

// Info will be used to notify users of any extra info (env variables, options)
type Info struct {
	data map[string]string
}

func (s *Info) Type() string {
	return "io.k8s.sigs.minikube.info"
}

// NewInfo returns a new Info type
func NewInfo(message string) *Info {
	return &Info{
		map[string]string{
			"message": message,
		},
	}
}

// Error will be used to notify the user of errors
type Error struct {
	data map[string]string
}

func NewError(err string) *Error {
	return &Error{
		map[string]string{
			"message": err,
		},
	}
}

func NewErrorExitCode(err string, exitcode int, additionalData ...map[string]string) *Error {
	e := NewError(err)
	e.data["exitcode"] = fmt.Sprintf("%v", exitcode)
	for _, a := range additionalData {
		for k, v := range a {
			e.data[k] = v
		}
	}
	return e
}

func (s *Error) Type() string {
	return "io.k8s.sigs.minikube.error"
}
<commit_msg>add log<commit_after>/*
Copyright 2020 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 register

import "fmt"

// Log represents the different types of logs that can be output as JSON
// This includes: Step, Download, DownloadProgress, Warning, Info, Error
type Log interface {
	Type() string
}

// Step represents a normal step in minikube execution
type Step struct {
	data map[string]string
}

// Type returns the cloud events compatible type of this struct
func (s *Step) Type() string {
	return "io.k8s.sigs.minikube.step"
}

// NewStep returns a new step type
func NewStep(message string) *Step {
	return &Step{data: map[string]string{
		"totalsteps":  Reg.totalSteps(),
		"currentstep": Reg.currentStep(),
		"message":     message,
		"name":        string(Reg.current),
	}}
}

// Download will be used to notify the user that a download has begun
type Download struct {
	data map[string]string
}

// Type returns the cloud events compatible type of this struct
func (s *Download) Type() string {
	return "io.k8s.sigs.minikube.download"
}

// NewDownload returns a new download type
func NewDownload(artifact string) *Download {
	return &Download{data: map[string]string{
		"totalsteps":  Reg.totalSteps(),
		"currentstep": Reg.currentStep(),
		"artifact":    artifact,
	}}
}

// DownloadProgress will be used to notify the user around the progress of a download
type DownloadProgress struct {
	data map[string]string
}

// Type returns the cloud events compatible type of this struct
func (s *DownloadProgress) Type() string {
	return "io.k8s.sigs.minikube.download.progress"
}

// NewDownloadProgress returns a new download progress type
func NewDownloadProgress(artifact, progress string) *DownloadProgress {
	return &DownloadProgress{data: map[string]string{
		"totalsteps":  Reg.totalSteps(),
		"currentstep": Reg.currentStep(),
		"progress":    progress,
		"artifact":    artifact,
	}}
}

// Warning will be used to notify the user of warnings
type Warning struct {
	data map[string]string
}

// NewWarning returns a new warning type
func NewWarning(warning string) *Warning {
	return &Warning{
		map[string]string{
			"message": warning,
		},
	}
}

// Type returns the cloud events compatible type of this struct
func (s *Warning) Type() string {
	return "io.k8s.sigs.minikube.warning"
}

// Info will be used to notify users of any extra info (env variables, options)
type Info struct {
	data map[string]string
}

func (s *Info) Type() string {
	return "io.k8s.sigs.minikube.info"
}

// NewInfo returns a new Info type
func NewInfo(message string) *Info {
	return &Info{
		map[string]string{
			"message": message,
		},
	}
}

// Error will be used to notify the user of errors
type Error struct {
	data map[string]string
}

func NewError(err string) *Error {
	return &Error{
		map[string]string{
			"message": err,
		},
	}
}

// NewErrorExitCode returns an error that has an associated exit code
func NewErrorExitCode(err string, exitcode int, additionalData ...map[string]string) *Error {
	e := NewError(err)
	e.data["exitcode"] = fmt.Sprintf("%v", exitcode)
	for _, a := range additionalData {
		for k, v := range a {
			e.data[k] = v
		}
	}
	return e
}

func (s *Error) Type() string {
	return "io.k8s.sigs.minikube.error"
}
<|endoftext|>


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

package main

import (
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"os"
	"os/signal"
	"strconv"
	"syscall"

	"github.com/gorilla/mux"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
	"github.com/uber/jaeger-lib/metrics"
	"github.com/uber/tchannel-go"
	"github.com/uber/tchannel-go/thrift"
	"go.uber.org/zap"

	basicB "github.com/jaegertracing/jaeger/cmd/builder"
	"github.com/jaegertracing/jaeger/cmd/collector/app"
	"github.com/jaegertracing/jaeger/cmd/collector/app/builder"
	"github.com/jaegertracing/jaeger/cmd/collector/app/sampling"
	"github.com/jaegertracing/jaeger/cmd/collector/app/zipkin"
	"github.com/jaegertracing/jaeger/cmd/env"
	"github.com/jaegertracing/jaeger/cmd/flags"
	"github.com/jaegertracing/jaeger/pkg/config"
	"github.com/jaegertracing/jaeger/pkg/healthcheck"
	pMetrics "github.com/jaegertracing/jaeger/pkg/metrics"
	"github.com/jaegertracing/jaeger/pkg/recoveryhandler"
	"github.com/jaegertracing/jaeger/pkg/version"
	ss "github.com/jaegertracing/jaeger/plugin/sampling/strategystore"
	"github.com/jaegertracing/jaeger/plugin/storage"
	jc "github.com/jaegertracing/jaeger/thrift-gen/jaeger"
	sc "github.com/jaegertracing/jaeger/thrift-gen/sampling"
	zc "github.com/jaegertracing/jaeger/thrift-gen/zipkincore"
)

const serviceName = "jaeger-collector"

func main() {
	var signalsChannel = make(chan os.Signal)
	signal.Notify(signalsChannel, os.Interrupt, syscall.SIGTERM)

	storageFactory, err := storage.NewFactory(storage.FactoryConfigFromEnvAndCLI(os.Args, os.Stderr))
	if err != nil {
		log.Fatalf("Cannot initialize storage factory: %v", err)
	}
	strategyStoreFactory, err := ss.NewFactory(ss.FactoryConfigFromEnv())
	if err != nil {
		log.Fatalf("Cannot initialize sampling strategy store factory: %v", err)
	}
	v := viper.New()
	command := &cobra.Command{
		Use:   "jaeger-collector",
		Short: "Jaeger collector receives and processes traces from Jaeger agents and clients",
		Long:  `Jaeger collector receives traces from Jaeger agents and runs them through a processing pipeline.`,
		RunE: func(cmd *cobra.Command, args []string) error {
			err := flags.TryLoadConfigFile(v)
			if err != nil {
				return err
			}

			sFlags := new(flags.SharedFlags).InitFromViper(v)
			logger, err := sFlags.NewLogger(zap.NewProductionConfig())
			if err != nil {
				return err
			}
			hc, err := sFlags.NewHealthCheck(logger)
			if err != nil {
				logger.Fatal("Could not start the health check server.", zap.Error(err))
			}

			builderOpts := new(builder.CollectorOptions).InitFromViper(v)

			mBldr := new(pMetrics.Builder).InitFromViper(v)
			baseFactory, err := mBldr.CreateMetricsFactory("jaeger")
			if err != nil {
				logger.Fatal("Cannot create metrics factory.", zap.Error(err))
			}

			storageFactory.InitFromViper(v)
			if err := storageFactory.Initialize(baseFactory, logger); err != nil {
				logger.Fatal("Failed to init storage factory", zap.Error(err))
			}
			spanWriter, err := storageFactory.CreateSpanWriter()
			if err != nil {
				logger.Fatal("Failed to create span writer", zap.Error(err))
			}

			metricsFactory := baseFactory.Namespace("collector", nil)
			handlerBuilder, err := builder.NewSpanHandlerBuilder(
				builderOpts,
				spanWriter,
				basicB.Options.LoggerOption(logger),
				basicB.Options.MetricsFactoryOption(metricsFactory),
			)
			if err != nil {
				logger.Fatal("Unable to set up builder", zap.Error(err))
			}

			ch, err := tchannel.NewChannel(serviceName, &tchannel.ChannelOptions{})
			if err != nil {
				logger.Fatal("Unable to create new TChannel", zap.Error(err))
			}
			server := thrift.NewServer(ch)
			zipkinSpansHandler, jaegerBatchesHandler := handlerBuilder.BuildHandlers()
			server.Register(jc.NewTChanCollectorServer(jaegerBatchesHandler))
			server.Register(zc.NewTChanZipkinCollectorServer(zipkinSpansHandler))

			samplingHandler := initializeSamplingHandler(strategyStoreFactory, v, metricsFactory, logger)
			server.Register(sc.NewTChanSamplingManagerServer(samplingHandler))

			portStr := ":" + strconv.Itoa(builderOpts.CollectorPort)
			listener, err := net.Listen("tcp", portStr)
			if err != nil {
				logger.Fatal("Unable to start listening on channel", zap.Error(err))
			}
			ch.Serve(listener)

			r := mux.NewRouter()
			apiHandler := app.NewAPIHandler(jaegerBatchesHandler)
			apiHandler.RegisterRoutes(r)
			if h := mBldr.Handler(); h != nil {
				logger.Info("Registering metrics handler with HTTP server", zap.String("route", mBldr.HTTPRoute))
				r.Handle(mBldr.HTTPRoute, h)
			}
			httpPortStr := ":" + strconv.Itoa(builderOpts.CollectorHTTPPort)
			recoveryHandler := recoveryhandler.NewRecoveryHandler(logger, true)

			go startZipkinHTTPAPI(logger, builderOpts.CollectorZipkinHTTPPort, zipkinSpansHandler, recoveryHandler)

			logger.Info("Starting HTTP server", zap.Int("http-port", builderOpts.CollectorHTTPPort))

			go func() {
				if err := http.ListenAndServe(httpPortStr, recoveryHandler(r)); err != nil {
					logger.Fatal("Could not launch service", zap.Error(err))
				}
				hc.Set(healthcheck.Unavailable)
			}()

			hc.Ready()
			<-signalsChannel
			logger.Info("Shutting down")
			if closer, ok := spanWriter.(io.Closer); ok {
				err := closer.Close()
				if err != nil {
					logger.Error("Failed to close span writer", zap.Error(err))
				}
			}

			logger.Info("Shutdown complete")
			return nil
		},
	}

	command.AddCommand(version.Command())
	command.AddCommand(env.Command())

	flags.SetDefaultHealthCheckPort(builder.CollectorDefaultHealthCheckHTTPPort)

	config.AddFlags(
		v,
		command,
		flags.AddConfigFileFlag,
		flags.AddFlags,
		builder.AddFlags,
		storageFactory.AddFlags,
		pMetrics.AddFlags,
		strategyStoreFactory.AddFlags,
	)

	if err := command.Execute(); err != nil {
		fmt.Println(err.Error())
		os.Exit(1)
	}
}

func startZipkinHTTPAPI(
	logger *zap.Logger,
	zipkinPort int,
	zipkinSpansHandler app.ZipkinSpansHandler,
	recoveryHandler func(http.Handler) http.Handler,
) {
	if zipkinPort != 0 {
		zHandler := zipkin.NewAPIHandler(zipkinSpansHandler)
		r := mux.NewRouter()
		zHandler.RegisterRoutes(r)

		httpPortStr := ":" + strconv.Itoa(zipkinPort)
		logger.Info("Listening for Zipkin HTTP traffic", zap.Int("zipkin.http-port", zipkinPort))

		if err := http.ListenAndServe(httpPortStr, recoveryHandler(r)); err != nil {
			logger.Fatal("Could not launch service", zap.Error(err))
		}
	}
}

func initializeSamplingHandler(
	samplingStrategyStoreFactory *ss.Factory,
	v *viper.Viper,
	metricsFactory metrics.Factory,
	logger *zap.Logger,
) sampling.Handler {
	samplingStrategyStoreFactory.InitFromViper(v)
	if err := samplingStrategyStoreFactory.Initialize(metricsFactory, logger); err != nil {
		logger.Fatal("Failed to init sampling strategy store factory", zap.Error(err))
	}
	strategyStore, err := samplingStrategyStoreFactory.CreateStrategyStore()
	if err != nil {
		logger.Fatal("Failed to create sampling strategy store", zap.Error(err))
	}
	return sampling.NewHandler(strategyStore)
}
<commit_msg>Log jaeger-collector tchannel port (#1136)<commit_after>// Copyright (c) 2017 Uber Technologies, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"os"
	"os/signal"
	"strconv"
	"syscall"

	"github.com/gorilla/mux"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
	"github.com/uber/jaeger-lib/metrics"
	"github.com/uber/tchannel-go"
	"github.com/uber/tchannel-go/thrift"
	"go.uber.org/zap"

	basicB "github.com/jaegertracing/jaeger/cmd/builder"
	"github.com/jaegertracing/jaeger/cmd/collector/app"
	"github.com/jaegertracing/jaeger/cmd/collector/app/builder"
	"github.com/jaegertracing/jaeger/cmd/collector/app/sampling"
	"github.com/jaegertracing/jaeger/cmd/collector/app/zipkin"
	"github.com/jaegertracing/jaeger/cmd/env"
	"github.com/jaegertracing/jaeger/cmd/flags"
	"github.com/jaegertracing/jaeger/pkg/config"
	"github.com/jaegertracing/jaeger/pkg/healthcheck"
	pMetrics "github.com/jaegertracing/jaeger/pkg/metrics"
	"github.com/jaegertracing/jaeger/pkg/recoveryhandler"
	"github.com/jaegertracing/jaeger/pkg/version"
	ss "github.com/jaegertracing/jaeger/plugin/sampling/strategystore"
	"github.com/jaegertracing/jaeger/plugin/storage"
	jc "github.com/jaegertracing/jaeger/thrift-gen/jaeger"
	sc "github.com/jaegertracing/jaeger/thrift-gen/sampling"
	zc "github.com/jaegertracing/jaeger/thrift-gen/zipkincore"
)

const serviceName = "jaeger-collector"

func main() {
	var signalsChannel = make(chan os.Signal)
	signal.Notify(signalsChannel, os.Interrupt, syscall.SIGTERM)

	storageFactory, err := storage.NewFactory(storage.FactoryConfigFromEnvAndCLI(os.Args, os.Stderr))
	if err != nil {
		log.Fatalf("Cannot initialize storage factory: %v", err)
	}
	strategyStoreFactory, err := ss.NewFactory(ss.FactoryConfigFromEnv())
	if err != nil {
		log.Fatalf("Cannot initialize sampling strategy store factory: %v", err)
	}
	v := viper.New()
	command := &cobra.Command{
		Use:   "jaeger-collector",
		Short: "Jaeger collector receives and processes traces from Jaeger agents and clients",
		Long:  `Jaeger collector receives traces from Jaeger agents and runs them through a processing pipeline.`,
		RunE: func(cmd *cobra.Command, args []string) error {
			err := flags.TryLoadConfigFile(v)
			if err != nil {
				return err
			}

			sFlags := new(flags.SharedFlags).InitFromViper(v)
			logger, err := sFlags.NewLogger(zap.NewProductionConfig())
			if err != nil {
				return err
			}
			hc, err := sFlags.NewHealthCheck(logger)
			if err != nil {
				logger.Fatal("Could not start the health check server.", zap.Error(err))
			}

			builderOpts := new(builder.CollectorOptions).InitFromViper(v)

			mBldr := new(pMetrics.Builder).InitFromViper(v)
			baseFactory, err := mBldr.CreateMetricsFactory("jaeger")
			if err != nil {
				logger.Fatal("Cannot create metrics factory.", zap.Error(err))
			}

			storageFactory.InitFromViper(v)
			if err := storageFactory.Initialize(baseFactory, logger); err != nil {
				logger.Fatal("Failed to init storage factory", zap.Error(err))
			}
			spanWriter, err := storageFactory.CreateSpanWriter()
			if err != nil {
				logger.Fatal("Failed to create span writer", zap.Error(err))
			}

			metricsFactory := baseFactory.Namespace("collector", nil)
			handlerBuilder, err := builder.NewSpanHandlerBuilder(
				builderOpts,
				spanWriter,
				basicB.Options.LoggerOption(logger),
				basicB.Options.MetricsFactoryOption(metricsFactory),
			)
			if err != nil {
				logger.Fatal("Unable to set up builder", zap.Error(err))
			}

			ch, err := tchannel.NewChannel(serviceName, &tchannel.ChannelOptions{})
			if err != nil {
				logger.Fatal("Unable to create new TChannel", zap.Error(err))
			}
			server := thrift.NewServer(ch)
			zipkinSpansHandler, jaegerBatchesHandler := handlerBuilder.BuildHandlers()
			server.Register(jc.NewTChanCollectorServer(jaegerBatchesHandler))
			server.Register(zc.NewTChanZipkinCollectorServer(zipkinSpansHandler))

			samplingHandler := initializeSamplingHandler(strategyStoreFactory, v, metricsFactory, logger)
			server.Register(sc.NewTChanSamplingManagerServer(samplingHandler))

			portStr := ":" + strconv.Itoa(builderOpts.CollectorPort)
			logger.Info("Starting jaeger-collector TChannel server", zap.Int("port", builderOpts.CollectorPort))

			listener, err := net.Listen("tcp", portStr)
			if err != nil {
				logger.Fatal("Unable to start listening on channel", zap.Error(err))
			}
			ch.Serve(listener)

			r := mux.NewRouter()
			apiHandler := app.NewAPIHandler(jaegerBatchesHandler)
			apiHandler.RegisterRoutes(r)
			if h := mBldr.Handler(); h != nil {
				logger.Info("Registering metrics handler with HTTP server", zap.String("route", mBldr.HTTPRoute))
				r.Handle(mBldr.HTTPRoute, h)
			}
			httpPortStr := ":" + strconv.Itoa(builderOpts.CollectorHTTPPort)
			recoveryHandler := recoveryhandler.NewRecoveryHandler(logger, true)

			go startZipkinHTTPAPI(logger, builderOpts.CollectorZipkinHTTPPort, zipkinSpansHandler, recoveryHandler)

			logger.Info("Starting jaeger-collector HTTP server", zap.Int("http-port", builderOpts.CollectorHTTPPort))

			go func() {
				if err := http.ListenAndServe(httpPortStr, recoveryHandler(r)); err != nil {
					logger.Fatal("Could not launch service", zap.Error(err))
				}
				hc.Set(healthcheck.Unavailable)
			}()

			hc.Ready()
			<-signalsChannel
			logger.Info("Shutting down")
			if closer, ok := spanWriter.(io.Closer); ok {
				err := closer.Close()
				if err != nil {
					logger.Error("Failed to close span writer", zap.Error(err))
				}
			}

			logger.Info("Shutdown complete")
			return nil
		},
	}

	command.AddCommand(version.Command())
	command.AddCommand(env.Command())

	flags.SetDefaultHealthCheckPort(builder.CollectorDefaultHealthCheckHTTPPort)

	config.AddFlags(
		v,
		command,
		flags.AddConfigFileFlag,
		flags.AddFlags,
		builder.AddFlags,
		storageFactory.AddFlags,
		pMetrics.AddFlags,
		strategyStoreFactory.AddFlags,
	)

	if err := command.Execute(); err != nil {
		fmt.Println(err.Error())
		os.Exit(1)
	}
}

func startZipkinHTTPAPI(
	logger *zap.Logger,
	zipkinPort int,
	zipkinSpansHandler app.ZipkinSpansHandler,
	recoveryHandler func(http.Handler) http.Handler,
) {
	if zipkinPort != 0 {
		zHandler := zipkin.NewAPIHandler(zipkinSpansHandler)
		r := mux.NewRouter()
		zHandler.RegisterRoutes(r)

		httpPortStr := ":" + strconv.Itoa(zipkinPort)
		logger.Info("Listening for Zipkin HTTP traffic", zap.Int("zipkin.http-port", zipkinPort))

		if err := http.ListenAndServe(httpPortStr, recoveryHandler(r)); err != nil {
			logger.Fatal("Could not launch service", zap.Error(err))
		}
	}
}

func initializeSamplingHandler(
	samplingStrategyStoreFactory *ss.Factory,
	v *viper.Viper,
	metricsFactory metrics.Factory,
	logger *zap.Logger,
) sampling.Handler {
	samplingStrategyStoreFactory.InitFromViper(v)
	if err := samplingStrategyStoreFactory.Initialize(metricsFactory, logger); err != nil {
		logger.Fatal("Failed to init sampling strategy store factory", zap.Error(err))
	}
	strategyStore, err := samplingStrategyStoreFactory.CreateStrategyStore()
	if err != nil {
		logger.Fatal("Failed to create sampling strategy store", zap.Error(err))
	}
	return sampling.NewHandler(strategyStore)
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package aws

import (
	"context"
	"fmt"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/ec2metadata"
	"github.com/aws/aws-sdk-go/aws/request"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/aws/aws-sdk-go/service/ec2/ec2iface"
	corev1 "k8s.io/api/core/v1"
	expirationcache "k8s.io/client-go/tools/cache"
	"k8s.io/klog/v2"

	"k8s.io/kops/pkg/nodeidentity"
)

const (
	// CloudTagInstanceGroupName is a cloud tag that defines the instance group name
	// This is used by the aws nodeidentifier to securely identify the node instancegroup
	CloudTagInstanceGroupName = "kops.k8s.io/instancegroup"
	// ClusterAutoscalerNodeTemplateLabel is the prefix used on node labels when copying to cloud tags.
	ClusterAutoscalerNodeTemplateLabel = "k8s.io/cluster-autoscaler/node-template/label/"
	// The expiration time of nodeidentity.Info cache.
	cacheTTL           = 60 * time.Minute
	KarpenterNodeLabel = "karpenter.sh/"
)

// nodeIdentifier identifies a node from EC2
type nodeIdentifier struct {
	// client is the ec2 interface
	ec2Client ec2iface.EC2API

	// cache is a cache of nodeidentity.Info
	cache expirationcache.Store
	// cacheEnabled indicates if caching should be used
	cacheEnabled bool
}

// New creates and returns a nodeidentity.Identifier for Nodes running on AWS
func New(CacheNodeidentityInfo bool) (nodeidentity.Identifier, error) {
	config := aws.NewConfig()
	config = config.WithCredentialsChainVerboseErrors(true)

	s, err := session.NewSession(config)
	if err != nil {
		return nil, fmt.Errorf("error starting new AWS session: %v", err)
	}
	s.Handlers.Send.PushFront(func(r *request.Request) {
		// Log requests
		klog.V(4).Infof("AWS API Request: %s/%s", r.ClientInfo.ServiceName, r.Operation.Name)
	})

	metadata := ec2metadata.New(s, config)

	region, err := metadata.Region()
	if err != nil {
		return nil, fmt.Errorf("error querying ec2 metadata service (for region): %v", err)
	}

	ec2Client := ec2.New(s, config.WithRegion(region))

	return &nodeIdentifier{
		ec2Client:    ec2Client,
		cache:        expirationcache.NewTTLStore(stringKeyFunc, cacheTTL),
		cacheEnabled: CacheNodeidentityInfo,
	}, nil
}

// stringKeyFunc is a string as cache key function
func stringKeyFunc(obj interface{}) (string, error) {
	key := obj.(*nodeidentity.Info).InstanceID
	return key, nil
}

// IdentifyNode queries AWS for the node identity information
func (i *nodeIdentifier) IdentifyNode(ctx context.Context, node *corev1.Node) (*nodeidentity.Info, error) {
	providerID := node.Spec.ProviderID
	if providerID == "" {
		return nil, fmt.Errorf("providerID was not set for node %s", node.Name)
	}
	if !strings.HasPrefix(providerID, "aws://") {
		return nil, fmt.Errorf("providerID %q not recognized for node %s", providerID, node.Name)
	}

	tokens := strings.Split(strings.TrimPrefix(providerID, "aws://"), "/")
	if len(tokens) != 3 {
		return nil, fmt.Errorf("providerID %q not recognized for node %s", providerID, node.Name)
	}

	// zone := tokens[1]
	instanceID := tokens[2]

	// If caching is enabled try pulling nodeidentity.Info from cache before
	// doing a EC2 API call.
	if i.cacheEnabled {
		obj, exists, err := i.cache.GetByKey(instanceID)
		if err != nil {
			klog.Warningf("Nodeidentity info cache lookup failure: %v", err)
		}
		if exists {
			return obj.(*nodeidentity.Info), nil
		}
	}

	// Based on node-authorizer code
	instance, err := i.getInstance(instanceID)
	if err != nil {
		return nil, err
	}

	instanceState := "?"
	if instance.State != nil {
		instanceState = aws.StringValue(instance.State.Name)
	}
	if instanceState != ec2.InstanceStateNameRunning && instanceState != ec2.InstanceStateNamePending {
		return nil, fmt.Errorf("found instance %q, but state is %q", instanceID, instanceState)
	}

	labels := map[string]string{}
	if instance.InstanceLifecycle != nil {
		labels[fmt.Sprintf("node-role.kubernetes.io/%s-worker", *instance.InstanceLifecycle)] = "true"
	}

	info := &nodeidentity.Info{
		InstanceID: instanceID,
		Labels:     labels,
	}

	isKarpenterManaged := false
	for _, tag := range instance.Tags {
		key := aws.StringValue(tag.Key)
		if strings.HasPrefix(key, ClusterAutoscalerNodeTemplateLabel) {
			info.Labels[strings.TrimPrefix(aws.StringValue(tag.Key), ClusterAutoscalerNodeTemplateLabel)] = aws.StringValue(tag.Value)
		}
		if strings.HasPrefix(key, KarpenterNodeLabel) {
			isKarpenterManaged = true
		}
	}
	if isKarpenterManaged {
		info.Labels["karpenter.sh/provisioner-name"] = info.Labels[CloudTagInstanceGroupName]
	}

	// If caching is enabled add the nodeidentity.Info to cache.
	if i.cacheEnabled {
		err = i.cache.Add(info)
		if err != nil {
			klog.Warningf("Failed to add node identity info to cache: %v", err)
		}
	}

	return info, nil
}

// getInstance queries EC2 for the instance with the specified ID, returning an error if not found
func (i *nodeIdentifier) getInstance(instanceID string) (*ec2.Instance, error) {
	// Based on node-authorizer code
	resp, err := i.ec2Client.DescribeInstances(&ec2.DescribeInstancesInput{
		InstanceIds: aws.StringSlice([]string{instanceID}),
	})
	if err != nil {
		return nil, fmt.Errorf("error from ec2 DescribeInstances request: %v", err)
	}

	// @check we found some instances
	if len(resp.Reservations) <= 0 || len(resp.Reservations[0].Instances) <= 0 {
		return nil, fmt.Errorf("missing instance id: %s", instanceID)
	}
	if len(resp.Reservations[0].Instances) > 1 {
		return nil, fmt.Errorf("found multiple instances with instance id: %s", instanceID)
	}

	instance := resp.Reservations[0].Instances[0]
	return instance, nil
}
<commit_msg>Don't explicitly set karpenter provisioner label<commit_after>/*
Copyright 2019 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package aws

import (
	"context"
	"fmt"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/ec2metadata"
	"github.com/aws/aws-sdk-go/aws/request"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/aws/aws-sdk-go/service/ec2/ec2iface"
	corev1 "k8s.io/api/core/v1"
	expirationcache "k8s.io/client-go/tools/cache"
	"k8s.io/klog/v2"

	"k8s.io/kops/pkg/nodeidentity"
)

const (
	// CloudTagInstanceGroupName is a cloud tag that defines the instance group name
	// This is used by the aws nodeidentifier to securely identify the node instancegroup
	CloudTagInstanceGroupName = "kops.k8s.io/instancegroup"
	// ClusterAutoscalerNodeTemplateLabel is the prefix used on node labels when copying to cloud tags.
	ClusterAutoscalerNodeTemplateLabel = "k8s.io/cluster-autoscaler/node-template/label/"
	// The expiration time of nodeidentity.Info cache.
	cacheTTL           = 60 * time.Minute
	KarpenterNodeLabel = "karpenter.sh/"
)

// nodeIdentifier identifies a node from EC2
type nodeIdentifier struct {
	// client is the ec2 interface
	ec2Client ec2iface.EC2API

	// cache is a cache of nodeidentity.Info
	cache expirationcache.Store
	// cacheEnabled indicates if caching should be used
	cacheEnabled bool
}

// New creates and returns a nodeidentity.Identifier for Nodes running on AWS
func New(CacheNodeidentityInfo bool) (nodeidentity.Identifier, error) {
	config := aws.NewConfig()
	config = config.WithCredentialsChainVerboseErrors(true)

	s, err := session.NewSession(config)
	if err != nil {
		return nil, fmt.Errorf("error starting new AWS session: %v", err)
	}
	s.Handlers.Send.PushFront(func(r *request.Request) {
		// Log requests
		klog.V(4).Infof("AWS API Request: %s/%s", r.ClientInfo.ServiceName, r.Operation.Name)
	})

	metadata := ec2metadata.New(s, config)

	region, err := metadata.Region()
	if err != nil {
		return nil, fmt.Errorf("error querying ec2 metadata service (for region): %v", err)
	}

	ec2Client := ec2.New(s, config.WithRegion(region))

	return &nodeIdentifier{
		ec2Client:    ec2Client,
		cache:        expirationcache.NewTTLStore(stringKeyFunc, cacheTTL),
		cacheEnabled: CacheNodeidentityInfo,
	}, nil
}

// stringKeyFunc is a string as cache key function
func stringKeyFunc(obj interface{}) (string, error) {
	key := obj.(*nodeidentity.Info).InstanceID
	return key, nil
}

// IdentifyNode queries AWS for the node identity information
func (i *nodeIdentifier) IdentifyNode(ctx context.Context, node *corev1.Node) (*nodeidentity.Info, error) {
	providerID := node.Spec.ProviderID
	if providerID == "" {
		return nil, fmt.Errorf("providerID was not set for node %s", node.Name)
	}
	if !strings.HasPrefix(providerID, "aws://") {
		return nil, fmt.Errorf("providerID %q not recognized for node %s", providerID, node.Name)
	}

	tokens := strings.Split(strings.TrimPrefix(providerID, "aws://"), "/")
	if len(tokens) != 3 {
		return nil, fmt.Errorf("providerID %q not recognized for node %s", providerID, node.Name)
	}

	// zone := tokens[1]
	instanceID := tokens[2]

	// If caching is enabled try pulling nodeidentity.Info from cache before
	// doing a EC2 API call.
	if i.cacheEnabled {
		obj, exists, err := i.cache.GetByKey(instanceID)
		if err != nil {
			klog.Warningf("Nodeidentity info cache lookup failure: %v", err)
		}
		if exists {
			return obj.(*nodeidentity.Info), nil
		}
	}

	// Based on node-authorizer code
	instance, err := i.getInstance(instanceID)
	if err != nil {
		return nil, err
	}

	instanceState := "?"
	if instance.State != nil {
		instanceState = aws.StringValue(instance.State.Name)
	}
	if instanceState != ec2.InstanceStateNameRunning && instanceState != ec2.InstanceStateNamePending {
		return nil, fmt.Errorf("found instance %q, but state is %q", instanceID, instanceState)
	}

	labels := map[string]string{}
	if instance.InstanceLifecycle != nil {
		labels[fmt.Sprintf("node-role.kubernetes.io/%s-worker", *instance.InstanceLifecycle)] = "true"
	}

	info := &nodeidentity.Info{
		InstanceID: instanceID,
		Labels:     labels,
	}

	for _, tag := range instance.Tags {
		key := aws.StringValue(tag.Key)
		if strings.HasPrefix(key, ClusterAutoscalerNodeTemplateLabel) {
			info.Labels[strings.TrimPrefix(aws.StringValue(tag.Key), ClusterAutoscalerNodeTemplateLabel)] = aws.StringValue(tag.Value)
		}
	}

	// If caching is enabled add the nodeidentity.Info to cache.
	if i.cacheEnabled {
		err = i.cache.Add(info)
		if err != nil {
			klog.Warningf("Failed to add node identity info to cache: %v", err)
		}
	}

	return info, nil
}

// getInstance queries EC2 for the instance with the specified ID, returning an error if not found
func (i *nodeIdentifier) getInstance(instanceID string) (*ec2.Instance, error) {
	// Based on node-authorizer code
	resp, err := i.ec2Client.DescribeInstances(&ec2.DescribeInstancesInput{
		InstanceIds: aws.StringSlice([]string{instanceID}),
	})
	if err != nil {
		return nil, fmt.Errorf("error from ec2 DescribeInstances request: %v", err)
	}

	// @check we found some instances
	if len(resp.Reservations) <= 0 || len(resp.Reservations[0].Instances) <= 0 {
		return nil, fmt.Errorf("missing instance id: %s", instanceID)
	}
	if len(resp.Reservations[0].Instances) > 1 {
		return nil, fmt.Errorf("found multiple instances with instance id: %s", instanceID)
	}

	instance := resp.Reservations[0].Instances[0]
	return instance, nil
}
<|endoftext|>


<commit_before>/*
   Copyright 2020 Docker Compose CLI authors

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
*/

package compose

import (
	"context"
	"fmt"
	"time"

	"github.com/compose-spec/compose-go/types"
	"github.com/spf13/cobra"

	"github.com/docker/compose/v2/pkg/api"
)

type createOptions struct {
	Build         bool
	noBuild       bool
	Pull          string
	removeOrphans bool
	ignoreOrphans bool
	forceRecreate bool
	noRecreate    bool
	recreateDeps  bool
	noInherit     bool
	timeChanged   bool
	timeout       int
	quietPull     bool
}

func createCommand(p *projectOptions, backend api.Service) *cobra.Command {
	opts := createOptions{}
	cmd := &cobra.Command{
		Use:   "create [OPTIONS] [SERVICE...]",
		Short: "Creates containers for a service.",
		PreRunE: Adapt(func(ctx context.Context, args []string) error {
			if opts.Build && opts.noBuild {
				return fmt.Errorf("--build and --no-build are incompatible")
			}
			if opts.forceRecreate && opts.noRecreate {
				return fmt.Errorf("--force-recreate and --no-recreate are incompatible")
			}
			return nil
		}),
		RunE: p.WithProject(func(ctx context.Context, project *types.Project) error {
			opts.Apply(project)
			return backend.Create(ctx, project, api.CreateOptions{
				RemoveOrphans:        opts.removeOrphans,
				IgnoreOrphans:        opts.ignoreOrphans,
				Recreate:             opts.recreateStrategy(),
				RecreateDependencies: opts.dependenciesRecreateStrategy(),
				Inherit:              !opts.noInherit,
				Timeout:              opts.GetTimeout(),
				QuietPull:            false,
			})
		}),
		ValidArgsFunction: serviceCompletion(p),
	}
	flags := cmd.Flags()
	flags.BoolVar(&opts.Build, "build", false, "Build images before starting containers.")
	flags.BoolVar(&opts.noBuild, "no-build", false, "Don't build an image, even if it's missing.")
	flags.StringVar(&opts.Pull, "pull", "missing", `Pull image before running ("always"|"missing"|"never")`)
	flags.BoolVar(&opts.forceRecreate, "force-recreate", false, "Recreate containers even if their configuration and image haven't changed.")
	flags.BoolVar(&opts.noRecreate, "no-recreate", false, "If containers already exist, don't recreate them. Incompatible with --force-recreate.")
	return cmd
}

func (opts createOptions) recreateStrategy() string {
	if opts.noRecreate {
		return api.RecreateNever
	}
	if opts.forceRecreate {
		return api.RecreateForce
	}
	return api.RecreateDiverged
}

func (opts createOptions) dependenciesRecreateStrategy() string {
	if opts.noRecreate {
		return api.RecreateNever
	}
	if opts.recreateDeps {
		return api.RecreateForce
	}
	return api.RecreateDiverged
}

func (opts createOptions) GetTimeout() *time.Duration {
	if opts.timeChanged {
		t := time.Duration(opts.timeout) * time.Second
		return &t
	}
	return nil
}

func (opts createOptions) Apply(project *types.Project) {
	if opts.Pull != "" {
		for i, service := range project.Services {
			service.PullPolicy = opts.Pull
			project.Services[i] = service
		}
	}
	if opts.Build {
		for i, service := range project.Services {
			if service.Build == nil {
				continue
			}
			service.PullPolicy = types.PullPolicyBuild
			project.Services[i] = service
		}
	}
	if opts.noBuild {
		for i, service := range project.Services {
			service.Build = nil
			if service.Image == "" {
				service.Image = api.GetImageNameOrDefault(service, project.Name)
			}
			project.Services[i] = service
		}
	}
}
<commit_msg>don't apply default pull policy from command line if one is define in service configuration<commit_after>/*
   Copyright 2020 Docker Compose CLI authors

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
*/

package compose

import (
	"context"
	"fmt"
	"time"

	"github.com/compose-spec/compose-go/types"
	"github.com/spf13/cobra"

	"github.com/docker/compose/v2/pkg/api"
)

type createOptions struct {
	Build         bool
	noBuild       bool
	Pull          string
	pullChanged   bool
	removeOrphans bool
	ignoreOrphans bool
	forceRecreate bool
	noRecreate    bool
	recreateDeps  bool
	noInherit     bool
	timeChanged   bool
	timeout       int
	quietPull     bool
}

func createCommand(p *projectOptions, backend api.Service) *cobra.Command {
	opts := createOptions{}
	cmd := &cobra.Command{
		Use:   "create [OPTIONS] [SERVICE...]",
		Short: "Creates containers for a service.",
		PreRunE: AdaptCmd(func(ctx context.Context, cmd *cobra.Command, args []string) error {
			opts.pullChanged = cmd.Flags().Changed("pull")
			if opts.Build && opts.noBuild {
				return fmt.Errorf("--build and --no-build are incompatible")
			}
			if opts.forceRecreate && opts.noRecreate {
				return fmt.Errorf("--force-recreate and --no-recreate are incompatible")
			}
			return nil
		}),
		RunE: p.WithProject(func(ctx context.Context, project *types.Project) error {
			opts.Apply(project)
			return backend.Create(ctx, project, api.CreateOptions{
				RemoveOrphans:        opts.removeOrphans,
				IgnoreOrphans:        opts.ignoreOrphans,
				Recreate:             opts.recreateStrategy(),
				RecreateDependencies: opts.dependenciesRecreateStrategy(),
				Inherit:              !opts.noInherit,
				Timeout:              opts.GetTimeout(),
				QuietPull:            false,
			})
		}),
		ValidArgsFunction: serviceCompletion(p),
	}
	flags := cmd.Flags()
	flags.BoolVar(&opts.Build, "build", false, "Build images before starting containers.")
	flags.BoolVar(&opts.noBuild, "no-build", false, "Don't build an image, even if it's missing.")
	flags.StringVar(&opts.Pull, "pull", "missing", `Pull image before running ("always"|"missing"|"never")`)
	flags.BoolVar(&opts.forceRecreate, "force-recreate", false, "Recreate containers even if their configuration and image haven't changed.")
	flags.BoolVar(&opts.noRecreate, "no-recreate", false, "If containers already exist, don't recreate them. Incompatible with --force-recreate.")
	return cmd
}

func (opts createOptions) recreateStrategy() string {
	if opts.noRecreate {
		return api.RecreateNever
	}
	if opts.forceRecreate {
		return api.RecreateForce
	}
	return api.RecreateDiverged
}

func (opts createOptions) dependenciesRecreateStrategy() string {
	if opts.noRecreate {
		return api.RecreateNever
	}
	if opts.recreateDeps {
		return api.RecreateForce
	}
	return api.RecreateDiverged
}

func (opts createOptions) GetTimeout() *time.Duration {
	if opts.timeChanged {
		t := time.Duration(opts.timeout) * time.Second
		return &t
	}
	return nil
}

func (opts createOptions) Apply(project *types.Project) {
	if opts.pullChanged {
		for i, service := range project.Services {
			service.PullPolicy = opts.Pull
			project.Services[i] = service
		}
	}
	if opts.Build {
		for i, service := range project.Services {
			if service.Build == nil {
				continue
			}
			service.PullPolicy = types.PullPolicyBuild
			project.Services[i] = service
		}
	}
	if opts.noBuild {
		for i, service := range project.Services {
			service.Build = nil
			if service.Image == "" {
				service.Image = api.GetImageNameOrDefault(service, project.Name)
			}
			project.Services[i] = service
		}
	}
}
<|endoftext|>