<commit_before>package cloudwatch

import (
	"context"
	"errors"
	"regexp"
	"sort"
	"strconv"
	"strings"
	"time"

	"github.com/grafana/grafana/pkg/log"
	"github.com/grafana/grafana/pkg/models"
	"github.com/grafana/grafana/pkg/setting"
	"github.com/grafana/grafana/pkg/tsdb"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/request"
	"github.com/aws/aws-sdk-go/service/cloudwatch"
	"github.com/grafana/grafana/pkg/components/null"
	"github.com/grafana/grafana/pkg/components/simplejson"
	"github.com/grafana/grafana/pkg/metrics"
)

type CloudWatchExecutor struct {
	*models.DataSource
}

type DatasourceInfo struct {
	Profile       string
	Region        string
	AuthType      string
	AssumeRoleArn string
	Namespace     string

	AccessKey string
	SecretKey string
}

func NewCloudWatchExecutor(dsInfo *models.DataSource) (tsdb.TsdbQueryEndpoint, error) {
	return &CloudWatchExecutor{}, nil
}

var (
	plog               log.Logger
	standardStatistics map[string]bool
	aliasFormat        *regexp.Regexp
)

func init() {
	plog = log.New("tsdb.cloudwatch")
	tsdb.RegisterTsdbQueryEndpoint("cloudwatch", NewCloudWatchExecutor)
	standardStatistics = map[string]bool{
		"Average":     true,
		"Maximum":     true,
		"Minimum":     true,
		"Sum":         true,
		"SampleCount": true,
	}
	aliasFormat = regexp.MustCompile(`\{\{\s*(.+?)\s*\}\}`)
}

func (e *CloudWatchExecutor) Query(ctx context.Context, dsInfo *models.DataSource, queryContext *tsdb.TsdbQuery) (*tsdb.Response, error) {
	var result *tsdb.Response
	e.DataSource = dsInfo
	queryType := queryContext.Queries[0].Model.Get("type").MustString("")
	var err error

	switch queryType {
	case "metricFindQuery":
		result, err = e.executeMetricFindQuery(ctx, queryContext)
		break
	case "annotationQuery":
		result, err = e.executeAnnotationQuery(ctx, queryContext)
		break
	case "timeSeriesQuery":
		fallthrough
	default:
		result, err = e.executeTimeSeriesQuery(ctx, queryContext)
		break
	}

	return result, err
}

func (e *CloudWatchExecutor) executeTimeSeriesQuery(ctx context.Context, queryContext *tsdb.TsdbQuery) (*tsdb.Response, error) {
	result := &tsdb.Response{
		Results: make(map[string]*tsdb.QueryResult),
	}

	errCh := make(chan error, 1)
	resCh := make(chan *tsdb.QueryResult, 1)

	currentlyExecuting := 0
	for i, model := range queryContext.Queries {
		queryType := model.Model.Get("type").MustString()
		if queryType != "timeSeriesQuery" && queryType != "" {
			continue
		}
		currentlyExecuting++
		go func(refId string, index int) {
			queryRes, err := e.executeQuery(ctx, queryContext.Queries[index].Model, queryContext)
			currentlyExecuting--
			if err != nil {
				errCh <- err
			} else {
				queryRes.RefId = refId
				resCh <- queryRes
			}
		}(model.RefId, i)
	}

	for currentlyExecuting != 0 {
		select {
		case res := <-resCh:
			result.Results[res.RefId] = res
		case err := <-errCh:
			return result, err
		case <-ctx.Done():
			return result, ctx.Err()
		}
	}

	return result, nil
}

func (e *CloudWatchExecutor) executeQuery(ctx context.Context, parameters *simplejson.Json, queryContext *tsdb.TsdbQuery) (*tsdb.QueryResult, error) {
	query, err := parseQuery(parameters)
	if err != nil {
		return nil, err
	}

	client, err := e.getClient(query.Region)
	if err != nil {
		return nil, err
	}

	startTime, err := queryContext.TimeRange.ParseFrom()
	if err != nil {
		return nil, err
	}

	endTime, err := queryContext.TimeRange.ParseTo()
	if err != nil {
		return nil, err
	}

	params := &cloudwatch.GetMetricStatisticsInput{
		Namespace:  aws.String(query.Namespace),
		MetricName: aws.String(query.MetricName),
		Dimensions: query.Dimensions,
		Period:     aws.Int64(int64(query.Period)),
		StartTime:  aws.Time(startTime),
		EndTime:    aws.Time(endTime),
	}
	if len(query.Statistics) > 0 {
		params.Statistics = query.Statistics
	}
	if len(query.ExtendedStatistics) > 0 {
		params.ExtendedStatistics = query.ExtendedStatistics
	}

	if setting.Env == setting.DEV {
		plog.Debug("CloudWatch query", "raw query", params)
	}

	resp, err := client.GetMetricStatisticsWithContext(ctx, params, request.WithResponseReadTimeout(10*time.Second))
	if err != nil {
		return nil, err
	}
	metrics.M_Aws_CloudWatch_GetMetricStatistics.Inc()

	queryRes, err := parseResponse(resp, query)
	if err != nil {
		return nil, err
	}

	return queryRes, nil
}

func parseDimensions(model *simplejson.Json) ([]*cloudwatch.Dimension, error) {
	var result []*cloudwatch.Dimension

	for k, v := range model.Get("dimensions").MustMap() {
		kk := k
		if vv, ok := v.(string); ok {
			result = append(result, &cloudwatch.Dimension{
				Name:  &kk,
				Value: &vv,
			})
		} else {
			return nil, errors.New("failed to parse")
		}
	}

	sort.Slice(result, func(i, j int) bool {
		return *result[i].Name < *result[j].Name
	})
	return result, nil
}

func parseStatistics(model *simplejson.Json) ([]string, []string, error) {
	var statistics []string
	var extendedStatistics []string

	for _, s := range model.Get("statistics").MustArray() {
		if ss, ok := s.(string); ok {
			if _, isStandard := standardStatistics[ss]; isStandard {
				statistics = append(statistics, ss)
			} else {
				extendedStatistics = append(extendedStatistics, ss)
			}
		} else {
			return nil, nil, errors.New("failed to parse")
		}
	}

	return statistics, extendedStatistics, nil
}

func parseQuery(model *simplejson.Json) (*CloudWatchQuery, error) {
	region, err := model.Get("region").String()
	if err != nil {
		return nil, err
	}

	namespace, err := model.Get("namespace").String()
	if err != nil {
		return nil, err
	}

	metricName, err := model.Get("metricName").String()
	if err != nil {
		return nil, err
	}

	dimensions, err := parseDimensions(model)
	if err != nil {
		return nil, err
	}

	statistics, extendedStatistics, err := parseStatistics(model)
	if err != nil {
		return nil, err
	}

	p := model.Get("period").MustString("")
	if p == "" {
		if namespace == "AWS/EC2" {
			p = "300"
		} else {
			p = "60"
		}
	}

	period := 300
	if regexp.MustCompile(`^\d+$`).Match([]byte(p)) {
		period, err = strconv.Atoi(p)
		if err != nil {
			return nil, err
		}
	} else {
		d, err := time.ParseDuration(p)
		if err != nil {
			return nil, err
		}
		period = int(d.Seconds())
	}

	alias := model.Get("alias").MustString("{{metric}}_{{stat}}")

	return &CloudWatchQuery{
		Region:             region,
		Namespace:          namespace,
		MetricName:         metricName,
		Dimensions:         dimensions,
		Statistics:         aws.StringSlice(statistics),
		ExtendedStatistics: aws.StringSlice(extendedStatistics),
		Period:             period,
		Alias:              alias,
	}, nil
}

func formatAlias(query *CloudWatchQuery, stat string, dimensions map[string]string) string {
	data := map[string]string{}
	data["region"] = query.Region
	data["namespace"] = query.Namespace
	data["metric"] = query.MetricName
	data["stat"] = stat
	for k, v := range dimensions {
		data[k] = v
	}

	result := aliasFormat.ReplaceAllFunc([]byte(query.Alias), func(in []byte) []byte {
		labelName := strings.Replace(string(in), "{{", "", 1)
		labelName = strings.Replace(labelName, "}}", "", 1)
		labelName = strings.TrimSpace(labelName)
		if val, exists := data[labelName]; exists {
			return []byte(val)
		}

		return in
	})

	return string(result)
}

func parseResponse(resp *cloudwatch.GetMetricStatisticsOutput, query *CloudWatchQuery) (*tsdb.QueryResult, error) {
	queryRes := tsdb.NewQueryResult()

	var value float64
	for _, s := range append(query.Statistics, query.ExtendedStatistics...) {
		series := tsdb.TimeSeries{
			Tags: map[string]string{},
		}
		for _, d := range query.Dimensions {
			series.Tags[*d.Name] = *d.Value
		}
		series.Name = formatAlias(query, *s, series.Tags)

		lastTimestamp := make(map[string]time.Time)
		sort.Slice(resp.Datapoints, func(i, j int) bool {
			return (*resp.Datapoints[i].Timestamp).Before(*resp.Datapoints[j].Timestamp)
		})
		for _, v := range resp.Datapoints {
			switch *s {
			case "Average":
				value = *v.Average
			case "Maximum":
				value = *v.Maximum
			case "Minimum":
				value = *v.Minimum
			case "Sum":
				value = *v.Sum
			case "SampleCount":
				value = *v.SampleCount
			default:
				if strings.Index(*s, "p") == 0 && v.ExtendedStatistics[*s] != nil {
					value = *v.ExtendedStatistics[*s]
				}
			}

			// terminate gap of data points
			timestamp := *v.Timestamp
			if _, ok := lastTimestamp[*s]; ok {
				nextTimestampFromLast := lastTimestamp[*s].Add(time.Duration(query.Period) * time.Second)
				for timestamp.After(nextTimestampFromLast) {
					series.Points = append(series.Points, tsdb.NewTimePoint(null.FloatFromPtr(nil), float64(nextTimestampFromLast.Unix()*1000)))
					nextTimestampFromLast = nextTimestampFromLast.Add(time.Duration(query.Period) * time.Second)
				}
			}
			lastTimestamp[*s] = timestamp

			series.Points = append(series.Points, tsdb.NewTimePoint(null.FloatFrom(value), float64(timestamp.Unix()*1000)))
		}

		queryRes.Series = append(queryRes.Series, &series)
	}

	return queryRes, nil
}
<commit_msg>fix default alias<commit_after>package cloudwatch

import (
	"context"
	"errors"
	"regexp"
	"sort"
	"strconv"
	"strings"
	"time"

	"github.com/grafana/grafana/pkg/log"
	"github.com/grafana/grafana/pkg/models"
	"github.com/grafana/grafana/pkg/setting"
	"github.com/grafana/grafana/pkg/tsdb"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/request"
	"github.com/aws/aws-sdk-go/service/cloudwatch"
	"github.com/grafana/grafana/pkg/components/null"
	"github.com/grafana/grafana/pkg/components/simplejson"
	"github.com/grafana/grafana/pkg/metrics"
)

type CloudWatchExecutor struct {
	*models.DataSource
}

type DatasourceInfo struct {
	Profile       string
	Region        string
	AuthType      string
	AssumeRoleArn string
	Namespace     string

	AccessKey string
	SecretKey string
}

func NewCloudWatchExecutor(dsInfo *models.DataSource) (tsdb.TsdbQueryEndpoint, error) {
	return &CloudWatchExecutor{}, nil
}

var (
	plog               log.Logger
	standardStatistics map[string]bool
	aliasFormat        *regexp.Regexp
)

func init() {
	plog = log.New("tsdb.cloudwatch")
	tsdb.RegisterTsdbQueryEndpoint("cloudwatch", NewCloudWatchExecutor)
	standardStatistics = map[string]bool{
		"Average":     true,
		"Maximum":     true,
		"Minimum":     true,
		"Sum":         true,
		"SampleCount": true,
	}
	aliasFormat = regexp.MustCompile(`\{\{\s*(.+?)\s*\}\}`)
}

func (e *CloudWatchExecutor) Query(ctx context.Context, dsInfo *models.DataSource, queryContext *tsdb.TsdbQuery) (*tsdb.Response, error) {
	var result *tsdb.Response
	e.DataSource = dsInfo
	queryType := queryContext.Queries[0].Model.Get("type").MustString("")
	var err error

	switch queryType {
	case "metricFindQuery":
		result, err = e.executeMetricFindQuery(ctx, queryContext)
		break
	case "annotationQuery":
		result, err = e.executeAnnotationQuery(ctx, queryContext)
		break
	case "timeSeriesQuery":
		fallthrough
	default:
		result, err = e.executeTimeSeriesQuery(ctx, queryContext)
		break
	}

	return result, err
}

func (e *CloudWatchExecutor) executeTimeSeriesQuery(ctx context.Context, queryContext *tsdb.TsdbQuery) (*tsdb.Response, error) {
	result := &tsdb.Response{
		Results: make(map[string]*tsdb.QueryResult),
	}

	errCh := make(chan error, 1)
	resCh := make(chan *tsdb.QueryResult, 1)

	currentlyExecuting := 0
	for i, model := range queryContext.Queries {
		queryType := model.Model.Get("type").MustString()
		if queryType != "timeSeriesQuery" && queryType != "" {
			continue
		}
		currentlyExecuting++
		go func(refId string, index int) {
			queryRes, err := e.executeQuery(ctx, queryContext.Queries[index].Model, queryContext)
			currentlyExecuting--
			if err != nil {
				errCh <- err
			} else {
				queryRes.RefId = refId
				resCh <- queryRes
			}
		}(model.RefId, i)
	}

	for currentlyExecuting != 0 {
		select {
		case res := <-resCh:
			result.Results[res.RefId] = res
		case err := <-errCh:
			return result, err
		case <-ctx.Done():
			return result, ctx.Err()
		}
	}

	return result, nil
}

func (e *CloudWatchExecutor) executeQuery(ctx context.Context, parameters *simplejson.Json, queryContext *tsdb.TsdbQuery) (*tsdb.QueryResult, error) {
	query, err := parseQuery(parameters)
	if err != nil {
		return nil, err
	}

	client, err := e.getClient(query.Region)
	if err != nil {
		return nil, err
	}

	startTime, err := queryContext.TimeRange.ParseFrom()
	if err != nil {
		return nil, err
	}

	endTime, err := queryContext.TimeRange.ParseTo()
	if err != nil {
		return nil, err
	}

	params := &cloudwatch.GetMetricStatisticsInput{
		Namespace:  aws.String(query.Namespace),
		MetricName: aws.String(query.MetricName),
		Dimensions: query.Dimensions,
		Period:     aws.Int64(int64(query.Period)),
		StartTime:  aws.Time(startTime),
		EndTime:    aws.Time(endTime),
	}
	if len(query.Statistics) > 0 {
		params.Statistics = query.Statistics
	}
	if len(query.ExtendedStatistics) > 0 {
		params.ExtendedStatistics = query.ExtendedStatistics
	}

	if setting.Env == setting.DEV {
		plog.Debug("CloudWatch query", "raw query", params)
	}

	resp, err := client.GetMetricStatisticsWithContext(ctx, params, request.WithResponseReadTimeout(10*time.Second))
	if err != nil {
		return nil, err
	}
	metrics.M_Aws_CloudWatch_GetMetricStatistics.Inc()

	queryRes, err := parseResponse(resp, query)
	if err != nil {
		return nil, err
	}

	return queryRes, nil
}

func parseDimensions(model *simplejson.Json) ([]*cloudwatch.Dimension, error) {
	var result []*cloudwatch.Dimension

	for k, v := range model.Get("dimensions").MustMap() {
		kk := k
		if vv, ok := v.(string); ok {
			result = append(result, &cloudwatch.Dimension{
				Name:  &kk,
				Value: &vv,
			})
		} else {
			return nil, errors.New("failed to parse")
		}
	}

	sort.Slice(result, func(i, j int) bool {
		return *result[i].Name < *result[j].Name
	})
	return result, nil
}

func parseStatistics(model *simplejson.Json) ([]string, []string, error) {
	var statistics []string
	var extendedStatistics []string

	for _, s := range model.Get("statistics").MustArray() {
		if ss, ok := s.(string); ok {
			if _, isStandard := standardStatistics[ss]; isStandard {
				statistics = append(statistics, ss)
			} else {
				extendedStatistics = append(extendedStatistics, ss)
			}
		} else {
			return nil, nil, errors.New("failed to parse")
		}
	}

	return statistics, extendedStatistics, nil
}

func parseQuery(model *simplejson.Json) (*CloudWatchQuery, error) {
	region, err := model.Get("region").String()
	if err != nil {
		return nil, err
	}

	namespace, err := model.Get("namespace").String()
	if err != nil {
		return nil, err
	}

	metricName, err := model.Get("metricName").String()
	if err != nil {
		return nil, err
	}

	dimensions, err := parseDimensions(model)
	if err != nil {
		return nil, err
	}

	statistics, extendedStatistics, err := parseStatistics(model)
	if err != nil {
		return nil, err
	}

	p := model.Get("period").MustString("")
	if p == "" {
		if namespace == "AWS/EC2" {
			p = "300"
		} else {
			p = "60"
		}
	}

	period := 300
	if regexp.MustCompile(`^\d+$`).Match([]byte(p)) {
		period, err = strconv.Atoi(p)
		if err != nil {
			return nil, err
		}
	} else {
		d, err := time.ParseDuration(p)
		if err != nil {
			return nil, err
		}
		period = int(d.Seconds())
	}

	alias := model.Get("alias").MustString()
	if alias == "" {
		alias = "{{metric}}_{{stat}}"
	}

	return &CloudWatchQuery{
		Region:             region,
		Namespace:          namespace,
		MetricName:         metricName,
		Dimensions:         dimensions,
		Statistics:         aws.StringSlice(statistics),
		ExtendedStatistics: aws.StringSlice(extendedStatistics),
		Period:             period,
		Alias:              alias,
	}, nil
}

func formatAlias(query *CloudWatchQuery, stat string, dimensions map[string]string) string {
	data := map[string]string{}
	data["region"] = query.Region
	data["namespace"] = query.Namespace
	data["metric"] = query.MetricName
	data["stat"] = stat
	for k, v := range dimensions {
		data[k] = v
	}

	result := aliasFormat.ReplaceAllFunc([]byte(query.Alias), func(in []byte) []byte {
		labelName := strings.Replace(string(in), "{{", "", 1)
		labelName = strings.Replace(labelName, "}}", "", 1)
		labelName = strings.TrimSpace(labelName)
		if val, exists := data[labelName]; exists {
			return []byte(val)
		}

		return in
	})

	return string(result)
}

func parseResponse(resp *cloudwatch.GetMetricStatisticsOutput, query *CloudWatchQuery) (*tsdb.QueryResult, error) {
	queryRes := tsdb.NewQueryResult()

	var value float64
	for _, s := range append(query.Statistics, query.ExtendedStatistics...) {
		series := tsdb.TimeSeries{
			Tags: map[string]string{},
		}
		for _, d := range query.Dimensions {
			series.Tags[*d.Name] = *d.Value
		}
		series.Name = formatAlias(query, *s, series.Tags)

		lastTimestamp := make(map[string]time.Time)
		sort.Slice(resp.Datapoints, func(i, j int) bool {
			return (*resp.Datapoints[i].Timestamp).Before(*resp.Datapoints[j].Timestamp)
		})
		for _, v := range resp.Datapoints {
			switch *s {
			case "Average":
				value = *v.Average
			case "Maximum":
				value = *v.Maximum
			case "Minimum":
				value = *v.Minimum
			case "Sum":
				value = *v.Sum
			case "SampleCount":
				value = *v.SampleCount
			default:
				if strings.Index(*s, "p") == 0 && v.ExtendedStatistics[*s] != nil {
					value = *v.ExtendedStatistics[*s]
				}
			}

			// terminate gap of data points
			timestamp := *v.Timestamp
			if _, ok := lastTimestamp[*s]; ok {
				nextTimestampFromLast := lastTimestamp[*s].Add(time.Duration(query.Period) * time.Second)
				for timestamp.After(nextTimestampFromLast) {
					series.Points = append(series.Points, tsdb.NewTimePoint(null.FloatFromPtr(nil), float64(nextTimestampFromLast.Unix()*1000)))
					nextTimestampFromLast = nextTimestampFromLast.Add(time.Duration(query.Period) * time.Second)
				}
			}
			lastTimestamp[*s] = timestamp

			series.Points = append(series.Points, tsdb.NewTimePoint(null.FloatFrom(value), float64(timestamp.Unix()*1000)))
		}

		queryRes.Series = append(queryRes.Series, &series)
	}

	return queryRes, nil
}
<|endoftext|>


<commit_before>package service

import (
	"database/sql"
	"log"
	"net"
	"os"
	"time"

	pb "github.com/otoolep/go-grpc-pg/proto"

	"golang.org/x/net/context"
	"google.golang.org/grpc"
)

// Service represents a gRPC service that communicates with a database backend.
type Service struct {
	grpc *grpc.Server
	db   *sql.DB

	ln   net.Listener
	addr string

	logger *log.Logger
}

// New returns an instantiated service.
func New(addr string, db *sql.DB) *Service {
	s := Service{
		grpc:   grpc.NewServer(),
		db:     db,
		addr:   addr,
		logger: log.New(os.Stderr, "[service] ", log.LstdFlags),
	}

	pb.RegisterDBProviderServer(s.grpc, (*gprcService)(&s))
	return &s
}

// Addr returns the address on which the service is listening.
func (s *Service) Addr() string {
	return s.ln.Addr().String()
}

// Open opens the service, starting it listening on the configured address.
func (s *Service) Open() error {
	ln, err := net.Listen("tcp", s.addr)
	if err != nil {
		return err
	}
	s.ln = ln
	s.logger.Println("listening on", s.ln.Addr().String())

	go func() {
		err := s.grpc.Serve(s.ln)
		if err != nil {
			s.logger.Println("gRPC Serve() returned:", err.Error())
		}
	}()

	return nil
}

// Close closes the service.
func (s *Service) Close() error {
	s.grpc.GracefulStop()
	s.logger.Println("gRPC server stopped")
	return nil
}

// gprcService is an unexported type, that is the same type as Service.
//
// Having the methods that the gRPC service requires on this type means that even though
// the methods are exported, since the type is not, these methods are not visible outside
// this package.
type gprcService Service

// Query implements the Query interface of the gRPC service.
func (g *gprcService) Query(c context.Context, q *pb.QueryRequest) (*pb.QueryResponse, error) {
	start := time.Now()
	rows, err := g.db.Query(q.Stmt)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	// Get the column names.
	cols, err := rows.Columns()
	if err != nil {
		return nil, err
	}

	response := pb.QueryResponse{
		Columns: cols,
	}

	// Iterate through each row returned by the query.
	for rows.Next() {
		row := make([]string, len(cols))
		// Get a set of pointers to the strings allocated above.
		rowI := make([]interface{}, len(cols))
		for i := range row {
			rowI[i] = &row[i]
		}

		if err := rows.Scan(rowI...); err != nil {
			return nil, err
		}

		// Add the latest rows to existing rows.
		response.Rows = append(response.Rows, &pb.Row{Values: row})
	}

	g.logger.Printf(`query '%s' completed in %s, %d %s returned`,
		q.Stmt, time.Since(start), len(response.Rows), prettyRows(int64(len(response.Rows))))
	return &response, nil
}

// Exec implements the Exec interface of the gRPC service.
func (g *gprcService) Exec(c context.Context, e *pb.ExecRequest) (*pb.ExecResponse, error) {
	start := time.Now()
	r, err := g.db.Exec(e.Stmt)
	if err != nil {
		return nil, err
	}

	lid, err := r.LastInsertId()
	if err != nil {
		// Not all databases support LastInsertId()
		lid = -1
	}
	ra, err := r.RowsAffected()
	if err != nil {
		return nil, err
	}

	g.logger.Printf(`exec '%s' completed in %s, %d %s affected`,
		e.Stmt, time.Since(start), ra, prettyRows(ra))
	return &pb.ExecResponse{
		LastInsertId: lid,
		RowsAffected: ra,
	}, nil
}

// prettyRows returns a singular or plural form of "row", depending on n.
func prettyRows(n int64) string {
	if n == 1 {
		return "row"
	}
	return "rows"
}
<commit_msg>Change type name from gprcService to grpcService for readability<commit_after>package service

import (
	"database/sql"
	"log"
	"net"
	"os"
	"time"

	pb "github.com/otoolep/go-grpc-pg/proto"

	"golang.org/x/net/context"
	"google.golang.org/grpc"
)

// Service represents a gRPC service that communicates with a database backend.
type Service struct {
	grpc *grpc.Server
	db   *sql.DB

	ln   net.Listener
	addr string

	logger *log.Logger
}

// New returns an instantiated service.
func New(addr string, db *sql.DB) *Service {
	s := Service{
		grpc:   grpc.NewServer(),
		db:     db,
		addr:   addr,
		logger: log.New(os.Stderr, "[service] ", log.LstdFlags),
	}

	pb.RegisterDBProviderServer(s.grpc, (*grpcService)(&s))
	return &s
}

// Addr returns the address on which the service is listening.
func (s *Service) Addr() string {
	return s.ln.Addr().String()
}

// Open opens the service, starting it listening on the configured address.
func (s *Service) Open() error {
	ln, err := net.Listen("tcp", s.addr)
	if err != nil {
		return err
	}
	s.ln = ln
	s.logger.Println("listening on", s.ln.Addr().String())

	go func() {
		err := s.grpc.Serve(s.ln)
		if err != nil {
			s.logger.Println("gRPC Serve() returned:", err.Error())
		}
	}()

	return nil
}

// Close closes the service.
func (s *Service) Close() error {
	s.grpc.GracefulStop()
	s.logger.Println("gRPC server stopped")
	return nil
}

// grpcService is an unexported type, that is the same type as Service.
//
// Having the methods that the gRPC service requires on this type means that even though
// the methods are exported, since the type is not, these methods are not visible outside
// this package.
type grpcService Service

// Query implements the Query interface of the gRPC service.
func (g *grpcService) Query(c context.Context, q *pb.QueryRequest) (*pb.QueryResponse, error) {
	start := time.Now()
	rows, err := g.db.Query(q.Stmt)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	// Get the column names.
	cols, err := rows.Columns()
	if err != nil {
		return nil, err
	}

	response := pb.QueryResponse{
		Columns: cols,
	}

	// Iterate through each row returned by the query.
	for rows.Next() {
		row := make([]string, len(cols))
		// Get a set of pointers to the strings allocated above.
		rowI := make([]interface{}, len(cols))
		for i := range row {
			rowI[i] = &row[i]
		}

		if err := rows.Scan(rowI...); err != nil {
			return nil, err
		}

		// Add the latest rows to existing rows.
		response.Rows = append(response.Rows, &pb.Row{Values: row})
	}

	g.logger.Printf(`query '%s' completed in %s, %d %s returned`,
		q.Stmt, time.Since(start), len(response.Rows), prettyRows(int64(len(response.Rows))))
	return &response, nil
}

// Exec implements the Exec interface of the gRPC service.
func (g *grpcService) Exec(c context.Context, e *pb.ExecRequest) (*pb.ExecResponse, error) {
	start := time.Now()
	r, err := g.db.Exec(e.Stmt)
	if err != nil {
		return nil, err
	}

	lid, err := r.LastInsertId()
	if err != nil {
		// Not all databases support LastInsertId()
		lid = -1
	}
	ra, err := r.RowsAffected()
	if err != nil {
		return nil, err
	}

	g.logger.Printf(`exec '%s' completed in %s, %d %s affected`,
		e.Stmt, time.Since(start), ra, prettyRows(ra))
	return &pb.ExecResponse{
		LastInsertId: lid,
		RowsAffected: ra,
	}, nil
}

// prettyRows returns a singular or plural form of "row", depending on n.
func prettyRows(n int64) string {
	if n == 1 {
		return "row"
	}
	return "rows"
}
<|endoftext|>


<commit_before>package service

import (
	"fmt"
	"github.com/ninjasphere/app-presets/model"
	"github.com/ninjasphere/go-ninja/api"
	"github.com/ninjasphere/go-ninja/logger"
)

type PresetsService struct {
	Model       *model.Presets
	Save        func(*model.Presets)
	Conn        *ninja.Connection
	Log         *logger.Logger
	initialized bool
}

func (ps *PresetsService) Init() error {
	if ps.Log == nil {
		return fmt.Errorf("illegal state: no logger")
	}
	if ps.Model == nil {
		return fmt.Errorf("illegal state: Model is nil")
	}
	if ps.Save == nil {
		return fmt.Errorf("illegal state: Save is nil")
	}
	if ps.Conn == nil {
		return fmt.Errorf("illegal state: Conn is nil")
	}
	ps.initialized = true
	return nil
}

func (ps *PresetsService) Destroy() error {
	ps.initialized = false
	return nil
}

func (ps *PresetsService) checkInit() {
	if ps.Log == nil {
		ps.Log = logger.GetLogger("com.ninja.app-presets")
	}
	if !ps.initialized {
		ps.Log.Fatalf("illegal state: the service is not initialized")
	}
}

// see: http://schema.ninjablocks.com/service/presets#listPresetable
func (ps *PresetsService) ListPresetable(scope string) ([]*model.ThingState, error) {
	ps.checkInit()
	return make([]*model.ThingState, 0, 0), fmt.Errorf("unimplemented function: ListPresetable")
}

// see: http://schema.ninjablocks.com/service/presets#fetchScenes
func (ps *PresetsService) FetchScenes(scope string) ([]*model.Scene, error) {
	ps.checkInit()
	collect := make([]*model.Scene, 0, 0)
	for _, m := range ps.Model.Scenes {
		if m.Scope == scope {
			collect = append(collect, m)
		}
	}
	return collect, nil
}

// see: http://schema.ninjablocks.com/service/presets#fetchScene
func (ps *PresetsService) FetchScene(id string) (*model.Scene, error) {
	return nil, fmt.Errorf("unimplemented function: FetchScene")
}

// see: http://schema.ninjablocks.com/service/presets#storeScene
func (ps *PresetsService) StoreScene(model *model.Scene) error {
	return fmt.Errorf("unimplemented function: StoreScene")
}

// see: http://schema.ninjablocks.com/service/presets#applyScene
func (ps *PresetsService) ApplyScene(id string) error {
	return fmt.Errorf("unimplemented function: ApplyScene")
}
<commit_msg>Implement FetchScene.<commit_after>package service

import (
	"fmt"
	"github.com/ninjasphere/app-presets/model"
	"github.com/ninjasphere/go-ninja/api"
	"github.com/ninjasphere/go-ninja/logger"
)

type PresetsService struct {
	Model       *model.Presets
	Save        func(*model.Presets)
	Conn        *ninja.Connection
	Log         *logger.Logger
	initialized bool
}

func (ps *PresetsService) Init() error {
	if ps.Log == nil {
		return fmt.Errorf("illegal state: no logger")
	}
	if ps.Model == nil {
		return fmt.Errorf("illegal state: Model is nil")
	}
	if ps.Save == nil {
		return fmt.Errorf("illegal state: Save is nil")
	}
	if ps.Conn == nil {
		return fmt.Errorf("illegal state: Conn is nil")
	}
	ps.initialized = true
	return nil
}

func (ps *PresetsService) Destroy() error {
	ps.initialized = false
	return nil
}

func (ps *PresetsService) checkInit() {
	if ps.Log == nil {
		ps.Log = logger.GetLogger("com.ninja.app-presets")
	}
	if !ps.initialized {
		ps.Log.Fatalf("illegal state: the service is not initialized")
	}
}

// see: http://schema.ninjablocks.com/service/presets#listPresetable
func (ps *PresetsService) ListPresetable(scope string) ([]*model.ThingState, error) {
	ps.checkInit()
	return make([]*model.ThingState, 0, 0), fmt.Errorf("unimplemented function: ListPresetable")
}

// see: http://schema.ninjablocks.com/service/presets#fetchScenes
func (ps *PresetsService) FetchScenes(scope string) ([]*model.Scene, error) {
	ps.checkInit()
	collect := make([]*model.Scene, 0, 0)
	for _, m := range ps.Model.Scenes {
		if m.Scope == scope {
			collect = append(collect, m)
		}
	}
	return collect, nil
}

// see: http://schema.ninjablocks.com/service/presets#fetchScene
func (ps *PresetsService) FetchScene(id string) (*model.Scene, error) {
	ps.checkInit()
	for _, m := range ps.Model.Scenes {
		if m.ID == id {
			return m, nil
		}
	}
	return nil, fmt.Errorf("No such scene: %s", id)
}

// see: http://schema.ninjablocks.com/service/presets#storeScene
func (ps *PresetsService) StoreScene(model *model.Scene) error {
	return fmt.Errorf("unimplemented function: StoreScene")
}

// see: http://schema.ninjablocks.com/service/presets#applyScene
func (ps *PresetsService) ApplyScene(id string) error {
	return fmt.Errorf("unimplemented function: ApplyScene")
}
<|endoftext|>


<commit_before>package session

import (
	"encoding/gob"
	"errors"
	"net/http"

	"github.com/RichardKnop/go-oauth2-server/config"
	"github.com/gorilla/sessions"
)

// Service wraps session functionality
type Service struct {
	sessionStore   sessions.Store
	sessionOptions *sessions.Options
	session        *sessions.Session
	r              *http.Request
	w              http.ResponseWriter
}

// UserSession has user data stored in a session after logging in
type UserSession struct {
	ClientID     string
	Username     string
	AccessToken  string
	RefreshToken string
}

const (
	storageSessionName  = "go_oauth2_server_session"
	userSessionKey      = "go_oauth2_server_user"
)

var (
	errSessonNotStarted = errors.New("Session not started")
)

func init() {
	// Register a new datatype for storage in sessions
	gob.Register(new(UserSession))
}

// NewService starts a new Service instance
func NewService(cnf *config.Config, r *http.Request, w http.ResponseWriter) *Service {
	return &Service{
		// Session cookie storage
		sessionStore: sessions.NewCookieStore([]byte(cnf.Session.Secret)),
		// Session options
		sessionOptions: &sessions.Options{
			Path:     cnf.Session.Path,
			MaxAge:   cnf.Session.MaxAge,
			HttpOnly: cnf.Session.HTTPOnly,
		},
		r: r,
		w: w,
	}
}

// StartSession starts a new session. This method must be called before other
// public methods of this struct as it sets the internal session object
func (s *Service) StartSession() error {
	session, err := s.sessionStore.Get(s.r, storageSessionName)
	if err != nil {
		return err
	}
	s.session = session
	return nil
}

// GetUserSession returns the user session
func (s *Service) GetUserSession() (*UserSession, error) {
	// Make sure StartSession has been called
	if s.session == nil {
		return nil, errSessonNotStarted
	}

	// Retrieve our user session struct and type-assert it
	userSession, ok := s.session.Values[userSessionKey].(*UserSession)
	if !ok {
		return nil, errors.New("User session type assertion error")
	}

	return userSession, nil
}

// SetUserSession saves the user session
func (s *Service) SetUserSession(userSession *UserSession) error {
	// Make sure StartSession has been called
	if s.session == nil {
		return errSessonNotStarted
	}

	// Set a new user session
	s.session.Values[userSessionKey] = userSession
	return s.session.Save(s.r, s.w)
}

// ClearUserSession deletes the user session
func (s *Service) ClearUserSession() error {
	// Make sure StartSession has been called
	if s.session == nil {
		return errSessonNotStarted
	}

	// Delete the user session
	delete(s.session.Values, userSessionKey)
	return s.session.Save(s.r, s.w)
}

// SetFlashMessage sets a flash message,
// useful for displaying an error after 302 redirection
func (s *Service) SetFlashMessage(msg string) error {
	// Make sure StartSession has been called
	if s.session == nil {
		return errSessonNotStarted
	}

	// Add the flash message
	s.session.AddFlash(msg)
	return s.session.Save(s.r, s.w)
}

// GetFlashMessage returns the first flash message
func (s *Service) GetFlashMessage() (interface{}, error) {
	// Make sure StartSession has been called
	if s.session == nil {
		return nil, errSessonNotStarted
	}

	// Get the last flash message from the stack
	if flashes := s.session.Flashes(); len(flashes) > 0 {
		// We need to save the session, otherwise the flash message won't be removed
		s.session.Save(s.r, s.w)
		return flashes[0], nil
	}

	// No flash messages in the stack
	return nil, nil
}
<commit_msg>Fixed a session bug introduced in the previous commit.<commit_after>package session

import (
	"encoding/gob"
	"errors"
	"net/http"

	"github.com/RichardKnop/go-oauth2-server/config"
	"github.com/gorilla/sessions"
)

// Service wraps session functionality
type Service struct {
	sessionStore   sessions.Store
	sessionOptions *sessions.Options
	session        *sessions.Session
	r              *http.Request
	w              http.ResponseWriter
}

// UserSession has user data stored in a session after logging in
type UserSession struct {
	ClientID     string
	Username     string
	AccessToken  string
	RefreshToken string
}

var (
	storageSessionName  = "go_oauth2_server_session"
	userSessionKey      = "go_oauth2_server_user"
	errSessonNotStarted = errors.New("Session not started")
)

func init() {
	// Register a new datatype for storage in sessions
	gob.Register(new(UserSession))
}

// NewService starts a new Service instance
func NewService(cnf *config.Config, r *http.Request, w http.ResponseWriter) *Service {
	return &Service{
		// Session cookie storage
		sessionStore: sessions.NewCookieStore([]byte(cnf.Session.Secret)),
		// Session options
		sessionOptions: &sessions.Options{
			Path:     cnf.Session.Path,
			MaxAge:   cnf.Session.MaxAge,
			HttpOnly: cnf.Session.HTTPOnly,
		},
		r: r,
		w: w,
	}
}

// StartSession starts a new session. This method must be called before other
// public methods of this struct as it sets the internal session object
func (s *Service) StartSession() error {
	session, err := s.sessionStore.Get(s.r, storageSessionName)
	if err != nil {
		return err
	}
	s.session = session
	return nil
}

// GetUserSession returns the user session
func (s *Service) GetUserSession() (*UserSession, error) {
	// Make sure StartSession has been called
	if s.session == nil {
		return nil, errSessonNotStarted
	}

	// Retrieve our user session struct and type-assert it
	userSession, ok := s.session.Values[userSessionKey].(*UserSession)
	if !ok {
		return nil, errors.New("User session type assertion error")
	}

	return userSession, nil
}

// SetUserSession saves the user session
func (s *Service) SetUserSession(userSession *UserSession) error {
	// Make sure StartSession has been called
	if s.session == nil {
		return errSessonNotStarted
	}

	// Set a new user session
	s.session.Values[userSessionKey] = userSession
	return s.session.Save(s.r, s.w)
}

// ClearUserSession deletes the user session
func (s *Service) ClearUserSession() error {
	// Make sure StartSession has been called
	if s.session == nil {
		return errSessonNotStarted
	}

	// Delete the user session
	delete(s.session.Values, userSessionKey)
	return s.session.Save(s.r, s.w)
}

// SetFlashMessage sets a flash message,
// useful for displaying an error after 302 redirection
func (s *Service) SetFlashMessage(msg string) error {
	// Make sure StartSession has been called
	if s.session == nil {
		return errSessonNotStarted
	}

	// Add the flash message
	s.session.AddFlash(msg)
	return s.session.Save(s.r, s.w)
}

// GetFlashMessage returns the first flash message
func (s *Service) GetFlashMessage() (interface{}, error) {
	// Make sure StartSession has been called
	if s.session == nil {
		return nil, errSessonNotStarted
	}

	// Get the last flash message from the stack
	if flashes := s.session.Flashes(); len(flashes) > 0 {
		// We need to save the session, otherwise the flash message won't be removed
		s.session.Save(s.r, s.w)
		return flashes[0], nil
	}

	// No flash messages in the stack
	return nil, nil
}
<|endoftext|>


<commit_before>// Copyright (c) nano Author. All Rights Reserved.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.

package session

import (
	"errors"
	"net"
	"sync"
	"sync/atomic"
	"time"

	"github.com/lonnng/nano/service"
)

type NetworkEntity interface {
	Push(route string, v interface{}) error
	Response(v interface{}) error
	Close() error
	RemoteAddr() net.Addr
}

var (
	ErrIllegalUID = errors.New("illegal uid")
)

// This session type as argument pass to Handler method, is a proxy session
// for frontend session in frontend server or backend session in backend
// server, correspond frontend session or backend session id as a field
// will be store in type instance
//
// This is user sessions, does not contain raw sockets information
type Session struct {
	sync.RWMutex                        // protect data
	id           int64                  // session global unique id
	uid          int64                  // binding user id
	LastRID      uint                   // last request id
	lastTime     int64                  // last heartbeat time
	Entity       NetworkEntity          // raw session id, agent in frontend server, or acceptor in backend server
	data         map[string]interface{} // session data store
}

// Create new session instance
func New(entity NetworkEntity) *Session {
	return &Session{
		id:       service.Connections.SessionID(),
		Entity:   entity,
		data:     make(map[string]interface{}),
		lastTime: time.Now().Unix(),
	}
}

// Push message to session
func (s *Session) Push(route string, v interface{}) error {
	return s.Entity.Push(route, v)
}

// Response message to session
func (s *Session) Response(v interface{}) error {
	return s.Entity.Response(v)
}

func (s *Session) ID() int64 {
	return s.id
}

func (s *Session) Uid() int64 {
	return atomic.LoadInt64(&s.uid)
}

func (s *Session) Bind(uid int64) error {
	if uid < 1 {
		return ErrIllegalUID
	}

	atomic.StoreInt64(&s.uid, uid)
	return nil
}

func (s *Session) Close() {
	s.Entity.Close()
}

func (s *Session) Remove(key string) {
	s.Lock()
	defer s.Unlock()

	delete(s.data, key)
}

func (s *Session) Set(key string, value interface{}) {
	s.Lock()
	defer s.Unlock()

	s.data[key] = value
}

func (s *Session) HasKey(key string) bool {
	s.RLock()
	defer s.RUnlock()

	_, has := s.data[key]
	return has
}

func (s *Session) Int(key string) int {
	s.RLock()
	defer s.RUnlock()

	v, ok := s.data[key]
	if !ok {
		return 0
	}

	value, ok := v.(int)
	if !ok {
		return 0
	}
	return value
}

func (s *Session) Int8(key string) int8 {
	s.RLock()
	defer s.RUnlock()

	v, ok := s.data[key]
	if !ok {
		return 0
	}

	value, ok := v.(int8)
	if !ok {
		return 0
	}
	return value
}

func (s *Session) Int16(key string) int16 {
	s.RLock()
	defer s.RUnlock()

	v, ok := s.data[key]
	if !ok {
		return 0
	}

	value, ok := v.(int16)
	if !ok {
		return 0
	}
	return value
}

func (s *Session) Int32(key string) int32 {
	s.RLock()
	defer s.RUnlock()

	v, ok := s.data[key]
	if !ok {
		return 0
	}

	value, ok := v.(int32)
	if !ok {
		return 0
	}
	return value
}

func (s *Session) Int64(key string) int64 {
	s.RLock()
	defer s.RUnlock()

	v, ok := s.data[key]
	if !ok {
		return 0
	}

	value, ok := v.(int64)
	if !ok {
		return 0
	}
	return value
}

func (s *Session) Uint(key string) uint {
	s.RLock()
	defer s.RUnlock()

	v, ok := s.data[key]
	if !ok {
		return 0
	}

	value, ok := v.(uint)
	if !ok {
		return 0
	}
	return value
}

func (s *Session) Uint8(key string) uint8 {
	s.RLock()
	defer s.RUnlock()

	v, ok := s.data[key]
	if !ok {
		return 0
	}

	value, ok := v.(uint8)
	if !ok {
		return 0
	}
	return value
}

func (s *Session) Uint16(key string) uint16 {
	s.RLock()
	defer s.RUnlock()

	v, ok := s.data[key]
	if !ok {
		return 0
	}

	value, ok := v.(uint16)
	if !ok {
		return 0
	}
	return value
}

func (s *Session) Uint32(key string) uint32 {
	s.RLock()
	defer s.RUnlock()

	v, ok := s.data[key]
	if !ok {
		return 0
	}

	value, ok := v.(uint32)
	if !ok {
		return 0
	}
	return value
}

func (s *Session) Uint64(key string) uint64 {
	s.RLock()
	defer s.RUnlock()

	v, ok := s.data[key]
	if !ok {
		return 0
	}

	value, ok := v.(uint64)
	if !ok {
		return 0
	}
	return value
}

func (s *Session) Float32(key string) float32 {
	s.RLock()
	defer s.RUnlock()

	v, ok := s.data[key]
	if !ok {
		return 0
	}

	value, ok := v.(float32)
	if !ok {
		return 0
	}
	return value
}

func (s *Session) Float64(key string) float64 {
	s.RLock()
	defer s.RUnlock()

	v, ok := s.data[key]
	if !ok {
		return 0
	}

	value, ok := v.(float64)
	if !ok {
		return 0
	}
	return value
}

func (s *Session) String(key string) string {
	s.RLock()
	defer s.RUnlock()

	v, ok := s.data[key]
	if !ok {
		return ""
	}

	value, ok := v.(string)
	if !ok {
		return ""
	}
	return value
}

func (s *Session) Value(key string) interface{} {
	s.RLock()
	defer s.RUnlock()

	return s.data[key]
}

// Retrieve all session state
func (s *Session) State() map[string]interface{} {
	s.RLock()
	defer s.RUnlock()

	return s.data
}

// Restore session state after reconnect
func (s *Session) Restore(data map[string]interface{}) {
	s.data = data
}

func (s *Session) Clear() {
	s.Lock()
	defer s.Unlock()

	s.data = map[string]interface{}{}
}
<commit_msg>session related docs<commit_after>// Copyright (c) nano Author. All Rights Reserved.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.

package session

import (
	"errors"
	"net"
	"sync"
	"sync/atomic"
	"time"

	"github.com/lonnng/nano/service"
)

type NetworkEntity interface {
	Push(route string, v interface{}) error
	Response(v interface{}) error
	Close() error
	RemoteAddr() net.Addr
}

var (
	ErrIllegalUID = errors.New("illegal uid")
)

// Session represents a client session which could storage temp data during low-level
// keep connected, all data will be released when the low-level connection was broken.
// Session instance related to the client will be passed to Handler method as the first
// parameter.
type Session struct {
	sync.RWMutex                        // protect data
	id           int64                  // session global unique id
	uid          int64                  // binding user id
	LastRID      uint                   // last request id
	lastTime     int64                  // last heartbeat time
	Entity       NetworkEntity          // low-level network entity
	data         map[string]interface{} // session data store
}

// New returns a new session instance
// a NetworkEntity represent low-level network instace
func New(entity NetworkEntity) *Session {
	return &Session{
		id:       service.Connections.SessionID(),
		Entity:   entity,
		data:     make(map[string]interface{}),
		lastTime: time.Now().Unix(),
	}
}

// Push message to client
func (s *Session) Push(route string, v interface{}) error {
	return s.Entity.Push(route, v)
}

// Response message to client
func (s *Session) Response(v interface{}) error {
	return s.Entity.Response(v)
}

// ID returns the session id
func (s *Session) ID() int64 {
	return s.id
}

// Uid returns UID that bind to current session
func (s *Session) Uid() int64 {
	return atomic.LoadInt64(&s.uid)
}

// Bind bind UID to current session
func (s *Session) Bind(uid int64) error {
	if uid < 1 {
		return ErrIllegalUID
	}

	atomic.StoreInt64(&s.uid, uid)
	return nil
}

// Close terminate current session, session related data will not be released,
// all related data should be Clear explicitly in Session closed callback
func (s *Session) Close() {
	s.Entity.Close()
}

// Remove delete data associated with the key from session storage
func (s *Session) Remove(key string) {
	s.Lock()
	defer s.Unlock()

	delete(s.data, key)
}

// Set associates value with the key in session storage
func (s *Session) Set(key string, value interface{}) {
	s.Lock()
	defer s.Unlock()

	s.data[key] = value
}

// HasKey decides whether a key has associated value
func (s *Session) HasKey(key string) bool {
	s.RLock()
	defer s.RUnlock()

	_, has := s.data[key]
	return has
}

// Int returns the value associated with the key as a int.
func (s *Session) Int(key string) int {
	s.RLock()
	defer s.RUnlock()

	v, ok := s.data[key]
	if !ok {
		return 0
	}

	value, ok := v.(int)
	if !ok {
		return 0
	}
	return value
}

// Int8 returns the value associated with the key as a int8.
func (s *Session) Int8(key string) int8 {
	s.RLock()
	defer s.RUnlock()

	v, ok := s.data[key]
	if !ok {
		return 0
	}

	value, ok := v.(int8)
	if !ok {
		return 0
	}
	return value
}

// Int16 returns the value associated with the key as a int16.
func (s *Session) Int16(key string) int16 {
	s.RLock()
	defer s.RUnlock()

	v, ok := s.data[key]
	if !ok {
		return 0
	}

	value, ok := v.(int16)
	if !ok {
		return 0
	}
	return value
}

// Int32 returns the value associated with the key as a int32.
func (s *Session) Int32(key string) int32 {
	s.RLock()
	defer s.RUnlock()

	v, ok := s.data[key]
	if !ok {
		return 0
	}

	value, ok := v.(int32)
	if !ok {
		return 0
	}
	return value
}

// Int64 returns the value associated with the key as a int64.
func (s *Session) Int64(key string) int64 {
	s.RLock()
	defer s.RUnlock()

	v, ok := s.data[key]
	if !ok {
		return 0
	}

	value, ok := v.(int64)
	if !ok {
		return 0
	}
	return value
}

// Uint returns the value associated with the key as a uint.
func (s *Session) Uint(key string) uint {
	s.RLock()
	defer s.RUnlock()

	v, ok := s.data[key]
	if !ok {
		return 0
	}

	value, ok := v.(uint)
	if !ok {
		return 0
	}
	return value
}

// Uint8 returns the value associated with the key as a uint8.
func (s *Session) Uint8(key string) uint8 {
	s.RLock()
	defer s.RUnlock()

	v, ok := s.data[key]
	if !ok {
		return 0
	}

	value, ok := v.(uint8)
	if !ok {
		return 0
	}
	return value
}

// Uint16 returns the value associated with the key as a uint16.
func (s *Session) Uint16(key string) uint16 {
	s.RLock()
	defer s.RUnlock()

	v, ok := s.data[key]
	if !ok {
		return 0
	}

	value, ok := v.(uint16)
	if !ok {
		return 0
	}
	return value
}

// Uint32 returns the value associated with the key as a uint32.
func (s *Session) Uint32(key string) uint32 {
	s.RLock()
	defer s.RUnlock()

	v, ok := s.data[key]
	if !ok {
		return 0
	}

	value, ok := v.(uint32)
	if !ok {
		return 0
	}
	return value
}

// Uint64 returns the value associated with the key as a uint64.
func (s *Session) Uint64(key string) uint64 {
	s.RLock()
	defer s.RUnlock()

	v, ok := s.data[key]
	if !ok {
		return 0
	}

	value, ok := v.(uint64)
	if !ok {
		return 0
	}
	return value
}

// Float32 returns the value associated with the key as a float32.
func (s *Session) Float32(key string) float32 {
	s.RLock()
	defer s.RUnlock()

	v, ok := s.data[key]
	if !ok {
		return 0
	}

	value, ok := v.(float32)
	if !ok {
		return 0
	}
	return value
}

// Float64 returns the value associated with the key as a float64.
func (s *Session) Float64(key string) float64 {
	s.RLock()
	defer s.RUnlock()

	v, ok := s.data[key]
	if !ok {
		return 0
	}

	value, ok := v.(float64)
	if !ok {
		return 0
	}
	return value
}

// String returns the value associated with the key as a string.
func (s *Session) String(key string) string {
	s.RLock()
	defer s.RUnlock()

	v, ok := s.data[key]
	if !ok {
		return ""
	}

	value, ok := v.(string)
	if !ok {
		return ""
	}
	return value
}

// String returns the value associated with the key as a interface{}.
func (s *Session) Value(key string) interface{} {
	s.RLock()
	defer s.RUnlock()

	return s.data[key]
}

// State returns all session state
func (s *Session) State() map[string]interface{} {
	s.RLock()
	defer s.RUnlock()

	return s.data
}

// Restore session state after reconnect
func (s *Session) Restore(data map[string]interface{}) {
	s.data = data
}

// Clear releases all data related to current session
func (s *Session) Clear() {
	s.Lock()
	defer s.Unlock()

	s.data = map[string]interface{}{}
}
<|endoftext|>


<commit_before>// Copyright (c) 2013 Kelsey Hightower. All rights reserved.
// Use of this source code is governed by the Apache License, Version 2.0
// that can be found in the LICENSE file.
package etcd

import (
	"errors"
	"strings"

	goetcd "github.com/coreos/go-etcd/etcd"
)

// Client is a wrapper around the etcd client
type Client struct {
	client *goetcd.Client
}

// NewEtcdClient returns an *etcd.Client with a connection to named machines.
// It returns an error if a connection to the cluster cannot be made.
func NewEtcdClient(machines []string, cert, key string, caCert string) (*Client, error) {
	var c *goetcd.Client
	var err error
	if cert != "" && key != "" {
		c, err = goetcd.NewTLSClient(machines, cert, key, caCert)
		if err != nil {
			return &Client{c}, err
		}
	} else {
		c = goetcd.NewClient(machines)
	}
	success := c.SetCluster(machines)
	if !success {
		return &Client{c}, errors.New("cannot connect to etcd cluster: " + strings.Join(machines, ","))
	}
	return &Client{c}, nil
}

// GetValues queries etcd for keys prefixed by prefix.
func (c *Client) GetValues(keys []string) (map[string]string, error) {
	vars := make(map[string]string)
	for _, key := range keys {
		resp, err := c.client.Get(key, true, true)
		if err != nil {
			return vars, err
		}
		err = nodeWalk(resp.Node, vars)
		if err != nil {
			return vars, err
		}
	}
	return vars, nil
}

// nodeWalk recursively descends nodes, updating vars.
func nodeWalk(node *goetcd.Node, vars map[string]string) error {
	if node != nil {
		key := node.Key
		if !node.Dir {
			vars[key] = node.Value
		} else {
			for _, node := range node.Nodes {
				nodeWalk(node, vars)
			}
		}
	}
	return nil
}
<commit_msg>fix(etcd/client.go): add DialTimeout for 0.6.x<commit_after>// Copyright (c) 2013 Kelsey Hightower. All rights reserved.
// Use of this source code is governed by the Apache License, Version 2.0
// that can be found in the LICENSE file.
package etcd

import (
	"errors"
	"strings"
	"time"

	goetcd "github.com/coreos/go-etcd/etcd"
)

// Client is a wrapper around the etcd client
type Client struct {
	client *goetcd.Client
}

// NewEtcdClient returns an *etcd.Client with a connection to named machines.
// It returns an error if a connection to the cluster cannot be made.
func NewEtcdClient(machines []string, cert, key string, caCert string) (*Client, error) {
	var c *goetcd.Client
	var err error
	if cert != "" && key != "" {
		c, err = goetcd.NewTLSClient(machines, cert, key, caCert)
		if err != nil {
			return &Client{c}, err
		}
	} else {
		c = goetcd.NewClient(machines)
	}
	// Configure the DialTimeout, since 1 second is often too short
	c.SetDialTimeout(time.Duration(3) * time.Second)
	success := c.SetCluster(machines)
	if !success {
		return &Client{c}, errors.New("cannot connect to etcd cluster: " + strings.Join(machines, ","))
	}
	return &Client{c}, nil
}

// GetValues queries etcd for keys prefixed by prefix.
func (c *Client) GetValues(keys []string) (map[string]string, error) {
	vars := make(map[string]string)
	for _, key := range keys {
		resp, err := c.client.Get(key, true, true)
		if err != nil {
			return vars, err
		}
		err = nodeWalk(resp.Node, vars)
		if err != nil {
			return vars, err
		}
	}
	return vars, nil
}

// nodeWalk recursively descends nodes, updating vars.
func nodeWalk(node *goetcd.Node, vars map[string]string) error {
	if node != nil {
		key := node.Key
		if !node.Dir {
			vars[key] = node.Value
		} else {
			for _, node := range node.Nodes {
				nodeWalk(node, vars)
			}
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package shared

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

	"github.com/pkg/errors"
	"gopkg.in/robfig/cron.v2"

	"github.com/lxc/lxd/shared/units"
	"github.com/lxc/lxd/shared/validate"
)

// InstanceAction indicates the type of action being performed.
type InstanceAction string

// InstanceAction types.
const (
	Stop     InstanceAction = "stop"
	Start    InstanceAction = "start"
	Restart  InstanceAction = "restart"
	Freeze   InstanceAction = "freeze"
	Unfreeze InstanceAction = "unfreeze"
)

// ConfigVolatilePrefix indicates the prefix used for volatile config keys.
const ConfigVolatilePrefix = "volatile."

// IsRootDiskDevice returns true if the given device representation is configured as root disk for
// an instance. It typically get passed a specific entry of api.Instance.Devices.
func IsRootDiskDevice(device map[string]string) bool {
	// Root disk devices also need a non-empty "pool" property, but we can't check that here
	// because this function is used with clients talking to older servers where there was no
	// concept of a storage pool, and also it is used for migrating from old to new servers.
	// The validation of the non-empty "pool" property is done inside the disk device itself.
	if device["type"] == "disk" && device["path"] == "/" && device["source"] == "" {
		return true
	}

	return false
}

// GetRootDiskDevice returns the instance device that is configured as root disk.
// Returns the device name and device config map.
func GetRootDiskDevice(devices map[string]map[string]string) (string, map[string]string, error) {
	var devName string
	var dev map[string]string

	for n, d := range devices {
		if IsRootDiskDevice(d) {
			if devName != "" {
				return "", nil, fmt.Errorf("More than one root device found")
			}

			devName = n
			dev = d
		}
	}

	if devName != "" {
		return devName, dev, nil
	}

	return "", nil, fmt.Errorf("No root device could be found")
}

// HugePageSizeKeys is a list of known hugepage size configuration keys.
var HugePageSizeKeys = [...]string{"limits.hugepages.64KB", "limits.hugepages.1MB", "limits.hugepages.2MB", "limits.hugepages.1GB"}

// HugePageSizeSuffix contains the list of known hugepage size suffixes.
var HugePageSizeSuffix = [...]string{"64KB", "1MB", "2MB", "1GB"}

// KnownInstanceConfigKeys maps all fully defined, well-known config keys
// to an appropriate checker function, which validates whether or not a
// given value is syntactically legal.
var KnownInstanceConfigKeys = map[string]func(value string) error{
	"boot.autostart":             validate.Optional(validate.IsBool),
	"boot.autostart.delay":       validate.Optional(validate.IsInt64),
	"boot.autostart.priority":    validate.Optional(validate.IsInt64),
	"boot.stop.priority":         validate.Optional(validate.IsInt64),
	"boot.host_shutdown_timeout": validate.Optional(validate.IsInt64),

	"limits.cpu": func(value string) error {
		if value == "" {
			return nil
		}

		// Validate the character set
		match, _ := regexp.MatchString("^[-,0-9]*$", value)
		if !match {
			return fmt.Errorf("Invalid CPU limit syntax")
		}

		// Validate first character
		if strings.HasPrefix(value, "-") || strings.HasPrefix(value, ",") {
			return fmt.Errorf("CPU limit can't start with a separator")
		}

		// Validate last character
		if strings.HasSuffix(value, "-") || strings.HasSuffix(value, ",") {
			return fmt.Errorf("CPU limit can't end with a separator")
		}

		return nil
	},
	"limits.cpu.allowance": func(value string) error {
		if value == "" {
			return nil
		}

		if strings.HasSuffix(value, "%") {
			// Percentage based allocation
			_, err := strconv.Atoi(strings.TrimSuffix(value, "%"))
			if err != nil {
				return err
			}

			return nil
		}

		// Time based allocation
		fields := strings.SplitN(value, "/", 2)
		if len(fields) != 2 {
			return fmt.Errorf("Invalid allowance: %s", value)
		}

		_, err := strconv.Atoi(strings.TrimSuffix(fields[0], "ms"))
		if err != nil {
			return err
		}

		_, err = strconv.Atoi(strings.TrimSuffix(fields[1], "ms"))
		if err != nil {
			return err
		}

		return nil
	},
	"limits.cpu.priority": validate.Optional(validate.IsPriority),

	"limits.disk.priority": validate.Optional(validate.IsPriority),

	"limits.hugepages.64KB": validate.Optional(validate.IsSize),
	"limits.hugepages.1MB":  validate.Optional(validate.IsSize),
	"limits.hugepages.2MB":  validate.Optional(validate.IsSize),
	"limits.hugepages.1GB":  validate.Optional(validate.IsSize),

	"limits.memory": func(value string) error {
		if value == "" {
			return nil
		}

		if strings.HasSuffix(value, "%") {
			_, err := strconv.ParseInt(strings.TrimSuffix(value, "%"), 10, 64)
			if err != nil {
				return err
			}

			return nil
		}

		_, err := units.ParseByteSizeString(value)
		if err != nil {
			return err
		}

		return nil
	},
	"limits.memory.enforce": func(value string) error {
		return validate.IsOneOf(value, []string{"soft", "hard"})
	},
	"limits.memory.swap":          validate.Optional(validate.IsBool),
	"limits.memory.swap.priority": validate.Optional(validate.IsPriority),
	"limits.memory.hugepages":     validate.Optional(validate.IsBool),

	"limits.network.priority": validate.Optional(validate.IsPriority),

	"limits.processes": validate.Optional(validate.IsInt64),

	"linux.kernel_modules": validate.IsAny,

	"migration.incremental.memory":            validate.Optional(validate.IsBool),
	"migration.incremental.memory.iterations": validate.Optional(validate.IsUint32),
	"migration.incremental.memory.goal":       validate.Optional(validate.IsUint32),

	"nvidia.runtime":             validate.Optional(validate.IsBool),
	"nvidia.driver.capabilities": validate.IsAny,
	"nvidia.require.cuda":        validate.IsAny,
	"nvidia.require.driver":      validate.IsAny,

	"security.nesting":       validate.Optional(validate.IsBool),
	"security.privileged":    validate.Optional(validate.IsBool),
	"security.devlxd":        validate.Optional(validate.IsBool),
	"security.devlxd.images": validate.Optional(validate.IsBool),

	"security.protection.delete": validate.Optional(validate.IsBool),
	"security.protection.shift":  validate.Optional(validate.IsBool),

	"security.idmap.base":     validate.Optional(validate.IsUint32),
	"security.idmap.isolated": validate.Optional(validate.IsBool),
	"security.idmap.size":     validate.Optional(validate.IsUint32),

	"security.secureboot": validate.Optional(validate.IsBool),

	"security.syscalls.allow":                   validate.IsAny,
	"security.syscalls.blacklist_default":       validate.Optional(validate.IsBool),
	"security.syscalls.blacklist_compat":        validate.Optional(validate.IsBool),
	"security.syscalls.blacklist":               validate.IsAny,
	"security.syscalls.deny_default":            validate.Optional(validate.IsBool),
	"security.syscalls.deny_compat":             validate.Optional(validate.IsBool),
	"security.syscalls.deny":                    validate.IsAny,
	"security.syscalls.intercept.bpf":           validate.Optional(validate.IsBool),
	"security.syscalls.intercept.bpf.devices":   validate.Optional(validate.IsBool),
	"security.syscalls.intercept.mknod":         validate.Optional(validate.IsBool),
	"security.syscalls.intercept.mount":         validate.Optional(validate.IsBool),
	"security.syscalls.intercept.mount.allowed": validate.IsAny,
	"security.syscalls.intercept.mount.fuse":    validate.IsAny,
	"security.syscalls.intercept.mount.shift":   validate.Optional(validate.IsBool),
	"security.syscalls.intercept.setxattr":      validate.Optional(validate.IsBool),
	"security.syscalls.whitelist":               validate.IsAny,

	"snapshots.schedule": func(value string) error {
		if value == "" {
			return nil
		}

		if len(strings.Split(value, " ")) != 5 {
			return fmt.Errorf("Schedule must be of the form: <minute> <hour> <day-of-month> <month> <day-of-week>")
		}

		_, err := cron.Parse(fmt.Sprintf("* %s", value))
		if err != nil {
			return errors.Wrap(err, "Error parsing schedule")
		}

		return nil
	},
	"snapshots.schedule.stopped": validate.Optional(validate.IsBool),
	"snapshots.pattern":          validate.IsAny,
	"snapshots.expiry": func(value string) error {
		// Validate expression
		_, err := GetSnapshotExpiry(time.Time{}, value)
		return err
	},

	// Caller is responsible for full validation of any raw.* value
	"raw.apparmor": validate.IsAny,
	"raw.idmap":    validate.IsAny,
	"raw.lxc":      validate.IsAny,
	"raw.qemu":     validate.IsAny,
	"raw.seccomp":  validate.IsAny,

	"volatile.apply_template":   validate.IsAny,
	"volatile.base_image":       validate.IsAny,
	"volatile.last_state.idmap": validate.IsAny,
	"volatile.last_state.power": validate.IsAny,
	"volatile.idmap.base":       validate.IsAny,
	"volatile.idmap.current":    validate.IsAny,
	"volatile.idmap.next":       validate.IsAny,
	"volatile.apply_quota":      validate.IsAny,
	"volatile.uuid":             validate.Optional(validate.IsUUID),
}

// ConfigKeyChecker returns a function that will check whether or not
// a provide value is valid for the associate config key.  Returns an
// error if the key is not known.  The checker function only performs
// syntactic checking of the value, semantic and usage checking must
// be done by the caller.  User defined keys are always considered to
// be valid, e.g. user.* and environment.* keys.
func ConfigKeyChecker(key string) (func(value string) error, error) {
	if f, ok := KnownInstanceConfigKeys[key]; ok {
		return f, nil
	}

	if strings.HasPrefix(key, ConfigVolatilePrefix) {
		if strings.HasSuffix(key, ".hwaddr") {
			return validate.IsAny, nil
		}

		if strings.HasSuffix(key, ".name") {
			return validate.IsAny, nil
		}

		if strings.HasSuffix(key, ".host_name") {
			return validate.IsAny, nil
		}

		if strings.HasSuffix(key, ".mtu") {
			return validate.IsAny, nil
		}

		if strings.HasSuffix(key, ".created") {
			return validate.IsAny, nil
		}

		if strings.HasSuffix(key, ".id") {
			return validate.IsAny, nil
		}

		if strings.HasSuffix(key, ".vlan") {
			return validate.IsAny, nil
		}

		if strings.HasSuffix(key, ".spoofcheck") {
			return validate.IsAny, nil
		}

		if strings.HasSuffix(key, ".apply_quota") {
			return validate.IsAny, nil
		}

		if strings.HasSuffix(key, ".ceph_rbd") {
			return validate.IsAny, nil
		}

		if strings.HasSuffix(key, ".driver") {
			return validate.IsAny, nil
		}

		if strings.HasSuffix(key, ".uuid") {
			return validate.IsAny, nil
		}
	}

	if strings.HasPrefix(key, "environment.") {
		return validate.IsAny, nil
	}

	if strings.HasPrefix(key, "user.") {
		return validate.IsAny, nil
	}

	if strings.HasPrefix(key, "image.") {
		return validate.IsAny, nil
	}

	if strings.HasPrefix(key, "limits.kernel.") &&
		(len(key) > len("limits.kernel.")) {
		return validate.IsAny, nil
	}

	return nil, fmt.Errorf("Unknown configuration key: %s", key)
}

// InstanceGetParentAndSnapshotName returns the parent instance name, snapshot name,
// and whether it actually was a snapshot name.
func InstanceGetParentAndSnapshotName(name string) (string, string, bool) {
	fields := strings.SplitN(name, SnapshotDelimiter, 2)
	if len(fields) == 1 {
		return name, "", false
	}

	return fields[0], fields[1], true
}

// InstanceIncludeWhenCopying is used to decide whether to include a config item or not when copying an instance.
// The remoteCopy argument indicates if the copy is remote (i.e between LXD nodes) as this affects the keys kept.
func InstanceIncludeWhenCopying(configKey string, remoteCopy bool) bool {
	if configKey == "volatile.base_image" {
		return true // Include volatile.base_image always as it can help optimize copies.
	}

	if configKey == "volatile.last_state.idmap" && !remoteCopy {
		return true // Include volatile.last_state.idmap when doing local copy to avoid needless remapping.
	}

	if strings.HasPrefix(configKey, ConfigVolatilePrefix) {
		return false // Exclude all other volatile keys.
	}

	return true // Keep all other keys.
}
<commit_msg>shared/instance: Adds ErrNoRootDisk error var and returns it from GetRootDiskDevice<commit_after>package shared

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

	"github.com/pkg/errors"
	"gopkg.in/robfig/cron.v2"

	"github.com/lxc/lxd/shared/units"
	"github.com/lxc/lxd/shared/validate"
)

// InstanceAction indicates the type of action being performed.
type InstanceAction string

// InstanceAction types.
const (
	Stop     InstanceAction = "stop"
	Start    InstanceAction = "start"
	Restart  InstanceAction = "restart"
	Freeze   InstanceAction = "freeze"
	Unfreeze InstanceAction = "unfreeze"
)

// ConfigVolatilePrefix indicates the prefix used for volatile config keys.
const ConfigVolatilePrefix = "volatile."

// IsRootDiskDevice returns true if the given device representation is configured as root disk for
// an instance. It typically get passed a specific entry of api.Instance.Devices.
func IsRootDiskDevice(device map[string]string) bool {
	// Root disk devices also need a non-empty "pool" property, but we can't check that here
	// because this function is used with clients talking to older servers where there was no
	// concept of a storage pool, and also it is used for migrating from old to new servers.
	// The validation of the non-empty "pool" property is done inside the disk device itself.
	if device["type"] == "disk" && device["path"] == "/" && device["source"] == "" {
		return true
	}

	return false
}

// ErrNoRootDisk means there is no root disk device found.
var ErrNoRootDisk = fmt.Errorf("No root device could be found")

// GetRootDiskDevice returns the instance device that is configured as root disk.
// Returns the device name and device config map.
func GetRootDiskDevice(devices map[string]map[string]string) (string, map[string]string, error) {
	var devName string
	var dev map[string]string

	for n, d := range devices {
		if IsRootDiskDevice(d) {
			if devName != "" {
				return "", nil, fmt.Errorf("More than one root device found")
			}

			devName = n
			dev = d
		}
	}

	if devName != "" {
		return devName, dev, nil
	}

	return "", nil, ErrNoRootDisk
}

// HugePageSizeKeys is a list of known hugepage size configuration keys.
var HugePageSizeKeys = [...]string{"limits.hugepages.64KB", "limits.hugepages.1MB", "limits.hugepages.2MB", "limits.hugepages.1GB"}

// HugePageSizeSuffix contains the list of known hugepage size suffixes.
var HugePageSizeSuffix = [...]string{"64KB", "1MB", "2MB", "1GB"}

// KnownInstanceConfigKeys maps all fully defined, well-known config keys
// to an appropriate checker function, which validates whether or not a
// given value is syntactically legal.
var KnownInstanceConfigKeys = map[string]func(value string) error{
	"boot.autostart":             validate.Optional(validate.IsBool),
	"boot.autostart.delay":       validate.Optional(validate.IsInt64),
	"boot.autostart.priority":    validate.Optional(validate.IsInt64),
	"boot.stop.priority":         validate.Optional(validate.IsInt64),
	"boot.host_shutdown_timeout": validate.Optional(validate.IsInt64),

	"limits.cpu": func(value string) error {
		if value == "" {
			return nil
		}

		// Validate the character set
		match, _ := regexp.MatchString("^[-,0-9]*$", value)
		if !match {
			return fmt.Errorf("Invalid CPU limit syntax")
		}

		// Validate first character
		if strings.HasPrefix(value, "-") || strings.HasPrefix(value, ",") {
			return fmt.Errorf("CPU limit can't start with a separator")
		}

		// Validate last character
		if strings.HasSuffix(value, "-") || strings.HasSuffix(value, ",") {
			return fmt.Errorf("CPU limit can't end with a separator")
		}

		return nil
	},
	"limits.cpu.allowance": func(value string) error {
		if value == "" {
			return nil
		}

		if strings.HasSuffix(value, "%") {
			// Percentage based allocation
			_, err := strconv.Atoi(strings.TrimSuffix(value, "%"))
			if err != nil {
				return err
			}

			return nil
		}

		// Time based allocation
		fields := strings.SplitN(value, "/", 2)
		if len(fields) != 2 {
			return fmt.Errorf("Invalid allowance: %s", value)
		}

		_, err := strconv.Atoi(strings.TrimSuffix(fields[0], "ms"))
		if err != nil {
			return err
		}

		_, err = strconv.Atoi(strings.TrimSuffix(fields[1], "ms"))
		if err != nil {
			return err
		}

		return nil
	},
	"limits.cpu.priority": validate.Optional(validate.IsPriority),

	"limits.disk.priority": validate.Optional(validate.IsPriority),

	"limits.hugepages.64KB": validate.Optional(validate.IsSize),
	"limits.hugepages.1MB":  validate.Optional(validate.IsSize),
	"limits.hugepages.2MB":  validate.Optional(validate.IsSize),
	"limits.hugepages.1GB":  validate.Optional(validate.IsSize),

	"limits.memory": func(value string) error {
		if value == "" {
			return nil
		}

		if strings.HasSuffix(value, "%") {
			_, err := strconv.ParseInt(strings.TrimSuffix(value, "%"), 10, 64)
			if err != nil {
				return err
			}

			return nil
		}

		_, err := units.ParseByteSizeString(value)
		if err != nil {
			return err
		}

		return nil
	},
	"limits.memory.enforce": func(value string) error {
		return validate.IsOneOf(value, []string{"soft", "hard"})
	},
	"limits.memory.swap":          validate.Optional(validate.IsBool),
	"limits.memory.swap.priority": validate.Optional(validate.IsPriority),
	"limits.memory.hugepages":     validate.Optional(validate.IsBool),

	"limits.network.priority": validate.Optional(validate.IsPriority),

	"limits.processes": validate.Optional(validate.IsInt64),

	"linux.kernel_modules": validate.IsAny,

	"migration.incremental.memory":            validate.Optional(validate.IsBool),
	"migration.incremental.memory.iterations": validate.Optional(validate.IsUint32),
	"migration.incremental.memory.goal":       validate.Optional(validate.IsUint32),

	"nvidia.runtime":             validate.Optional(validate.IsBool),
	"nvidia.driver.capabilities": validate.IsAny,
	"nvidia.require.cuda":        validate.IsAny,
	"nvidia.require.driver":      validate.IsAny,

	"security.nesting":       validate.Optional(validate.IsBool),
	"security.privileged":    validate.Optional(validate.IsBool),
	"security.devlxd":        validate.Optional(validate.IsBool),
	"security.devlxd.images": validate.Optional(validate.IsBool),

	"security.protection.delete": validate.Optional(validate.IsBool),
	"security.protection.shift":  validate.Optional(validate.IsBool),

	"security.idmap.base":     validate.Optional(validate.IsUint32),
	"security.idmap.isolated": validate.Optional(validate.IsBool),
	"security.idmap.size":     validate.Optional(validate.IsUint32),

	"security.secureboot": validate.Optional(validate.IsBool),

	"security.syscalls.allow":                   validate.IsAny,
	"security.syscalls.blacklist_default":       validate.Optional(validate.IsBool),
	"security.syscalls.blacklist_compat":        validate.Optional(validate.IsBool),
	"security.syscalls.blacklist":               validate.IsAny,
	"security.syscalls.deny_default":            validate.Optional(validate.IsBool),
	"security.syscalls.deny_compat":             validate.Optional(validate.IsBool),
	"security.syscalls.deny":                    validate.IsAny,
	"security.syscalls.intercept.bpf":           validate.Optional(validate.IsBool),
	"security.syscalls.intercept.bpf.devices":   validate.Optional(validate.IsBool),
	"security.syscalls.intercept.mknod":         validate.Optional(validate.IsBool),
	"security.syscalls.intercept.mount":         validate.Optional(validate.IsBool),
	"security.syscalls.intercept.mount.allowed": validate.IsAny,
	"security.syscalls.intercept.mount.fuse":    validate.IsAny,
	"security.syscalls.intercept.mount.shift":   validate.Optional(validate.IsBool),
	"security.syscalls.intercept.setxattr":      validate.Optional(validate.IsBool),
	"security.syscalls.whitelist":               validate.IsAny,

	"snapshots.schedule": func(value string) error {
		if value == "" {
			return nil
		}

		if len(strings.Split(value, " ")) != 5 {
			return fmt.Errorf("Schedule must be of the form: <minute> <hour> <day-of-month> <month> <day-of-week>")
		}

		_, err := cron.Parse(fmt.Sprintf("* %s", value))
		if err != nil {
			return errors.Wrap(err, "Error parsing schedule")
		}

		return nil
	},
	"snapshots.schedule.stopped": validate.Optional(validate.IsBool),
	"snapshots.pattern":          validate.IsAny,
	"snapshots.expiry": func(value string) error {
		// Validate expression
		_, err := GetSnapshotExpiry(time.Time{}, value)
		return err
	},

	// Caller is responsible for full validation of any raw.* value
	"raw.apparmor": validate.IsAny,
	"raw.idmap":    validate.IsAny,
	"raw.lxc":      validate.IsAny,
	"raw.qemu":     validate.IsAny,
	"raw.seccomp":  validate.IsAny,

	"volatile.apply_template":   validate.IsAny,
	"volatile.base_image":       validate.IsAny,
	"volatile.last_state.idmap": validate.IsAny,
	"volatile.last_state.power": validate.IsAny,
	"volatile.idmap.base":       validate.IsAny,
	"volatile.idmap.current":    validate.IsAny,
	"volatile.idmap.next":       validate.IsAny,
	"volatile.apply_quota":      validate.IsAny,
	"volatile.uuid":             validate.Optional(validate.IsUUID),
}

// ConfigKeyChecker returns a function that will check whether or not
// a provide value is valid for the associate config key.  Returns an
// error if the key is not known.  The checker function only performs
// syntactic checking of the value, semantic and usage checking must
// be done by the caller.  User defined keys are always considered to
// be valid, e.g. user.* and environment.* keys.
func ConfigKeyChecker(key string) (func(value string) error, error) {
	if f, ok := KnownInstanceConfigKeys[key]; ok {
		return f, nil
	}

	if strings.HasPrefix(key, ConfigVolatilePrefix) {
		if strings.HasSuffix(key, ".hwaddr") {
			return validate.IsAny, nil
		}

		if strings.HasSuffix(key, ".name") {
			return validate.IsAny, nil
		}

		if strings.HasSuffix(key, ".host_name") {
			return validate.IsAny, nil
		}

		if strings.HasSuffix(key, ".mtu") {
			return validate.IsAny, nil
		}

		if strings.HasSuffix(key, ".created") {
			return validate.IsAny, nil
		}

		if strings.HasSuffix(key, ".id") {
			return validate.IsAny, nil
		}

		if strings.HasSuffix(key, ".vlan") {
			return validate.IsAny, nil
		}

		if strings.HasSuffix(key, ".spoofcheck") {
			return validate.IsAny, nil
		}

		if strings.HasSuffix(key, ".apply_quota") {
			return validate.IsAny, nil
		}

		if strings.HasSuffix(key, ".ceph_rbd") {
			return validate.IsAny, nil
		}

		if strings.HasSuffix(key, ".driver") {
			return validate.IsAny, nil
		}

		if strings.HasSuffix(key, ".uuid") {
			return validate.IsAny, nil
		}
	}

	if strings.HasPrefix(key, "environment.") {
		return validate.IsAny, nil
	}

	if strings.HasPrefix(key, "user.") {
		return validate.IsAny, nil
	}

	if strings.HasPrefix(key, "image.") {
		return validate.IsAny, nil
	}

	if strings.HasPrefix(key, "limits.kernel.") &&
		(len(key) > len("limits.kernel.")) {
		return validate.IsAny, nil
	}

	return nil, fmt.Errorf("Unknown configuration key: %s", key)
}

// InstanceGetParentAndSnapshotName returns the parent instance name, snapshot name,
// and whether it actually was a snapshot name.
func InstanceGetParentAndSnapshotName(name string) (string, string, bool) {
	fields := strings.SplitN(name, SnapshotDelimiter, 2)
	if len(fields) == 1 {
		return name, "", false
	}

	return fields[0], fields[1], true
}

// InstanceIncludeWhenCopying is used to decide whether to include a config item or not when copying an instance.
// The remoteCopy argument indicates if the copy is remote (i.e between LXD nodes) as this affects the keys kept.
func InstanceIncludeWhenCopying(configKey string, remoteCopy bool) bool {
	if configKey == "volatile.base_image" {
		return true // Include volatile.base_image always as it can help optimize copies.
	}

	if configKey == "volatile.last_state.idmap" && !remoteCopy {
		return true // Include volatile.last_state.idmap when doing local copy to avoid needless remapping.
	}

	if strings.HasPrefix(configKey, ConfigVolatilePrefix) {
		return false // Exclude all other volatile keys.
	}

	return true // Keep all other keys.
}
<|endoftext|>