<commit_before>package admin

import (
	"fmt"
	"reflect"
	"strconv"
	"strings"
	"time"

	"github.com/jinzhu/gorm"
	"github.com/jinzhu/now"
	"github.com/qor/qor"
	"github.com/qor/qor/resource"
	"github.com/qor/qor/roles"
	"github.com/qor/qor/utils"
)

type Resource struct {
	resource.Resource
	admin         *Admin
	Config        *Config
	Metas         []*Meta
	actions       []*Action
	scopes        []*Scope
	filters       map[string]*Filter
	searchAttrs   []string
	indexAttrs    []string
	newAttrs      []string
	editAttrs     []string
	showAttrs     []string
	cachedMetas   *map[string][]*Meta
	SearchHandler func(keyword string, context *qor.Context) *gorm.DB
}

func (res *Resource) Meta(meta *Meta) {
	if res.GetMeta(meta.Name) != nil {
		utils.ExitWithMsg("Duplicated meta %v defined for resource %v", meta.Name, res.Name)
	}

	meta.base = res
	meta.updateMeta()
	res.Metas = append(res.Metas, meta)
}

func (res Resource) GetAdmin() *Admin {
	return res.admin
}

func (res Resource) ToParam() string {
	return utils.ToParamString(res.Name)
}

func (res Resource) UseTheme(theme string) []string {
	if res.Config != nil {
		res.Config.Themes = append(res.Config.Themes, theme)
		return res.Config.Themes
	}
	return []string{}
}

func (res *Resource) convertObjectToMap(context *Context, value interface{}, kind string) interface{} {
	reflectValue := reflect.Indirect(reflect.ValueOf(value))
	switch reflectValue.Kind() {
	case reflect.Slice:
		values := []interface{}{}
		for i := 0; i < reflectValue.Len(); i++ {
			values = append(values, res.convertObjectToMap(context, reflectValue.Index(i).Interface(), kind))
		}
		return values
	case reflect.Struct:
		var metas []*Meta
		if kind == "index" {
			metas = res.indexMetas()
		} else if kind == "show" {
			metas = res.showMetas()
		}

		values := map[string]interface{}{}
		for _, meta := range metas {
			if meta.HasPermission(roles.Read, context.Context) {
				value := meta.GetValuer()(value, context.Context)
				if meta.Resource != nil {
					value = meta.Resource.(*Resource).convertObjectToMap(context, value, kind)
				}
				values[meta.GetName()] = value
			}
		}
		return values
	default:
		panic(fmt.Sprintf("Can't convert %v (%v) to map", reflectValue, reflectValue.Kind()))
	}
}

func (res *Resource) Decode(context *qor.Context, value interface{}) (errs []error) {
	return resource.Decode(context, value, res)
}

func (res *Resource) IndexAttrs(columns ...string) {
	res.indexAttrs = columns
}

func (res *Resource) NewAttrs(columns ...string) {
	res.newAttrs = columns
}

func (res *Resource) EditAttrs(columns ...string) {
	res.editAttrs = columns
}

func (res *Resource) ShowAttrs(columns ...string) {
	res.showAttrs = columns
}

func (res *Resource) SearchAttrs(columns ...string) []string {
	if len(columns) > 0 {
		res.searchAttrs = columns
		res.SearchHandler = func(keyword string, context *qor.Context) *gorm.DB {
			db := context.GetDB()
			var conditions []string
			var keywords []interface{}
			scope := db.NewScope(res.Value)

			for _, column := range columns {
				if field, ok := scope.FieldByName(column); ok {
					switch field.Field.Kind() {
					case reflect.String:
						conditions = append(conditions, fmt.Sprintf("upper(%v) like upper(?)", scope.Quote(field.DBName)))
						keywords = append(keywords, "%"+keyword+"%")
					case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64, reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
						if _, err := strconv.Atoi(keyword); err == nil {
							conditions = append(conditions, fmt.Sprintf("%v = ?", scope.Quote(field.DBName)))
							keywords = append(keywords, keyword)
						}
					case reflect.Float32, reflect.Float64:
						if _, err := strconv.ParseFloat(keyword, 64); err == nil {
							conditions = append(conditions, fmt.Sprintf("%v = ?", scope.Quote(field.DBName)))
							keywords = append(keywords, keyword)
						}
					case reflect.Struct:
						// time ?
						if _, ok := field.Field.Interface().(time.Time); ok {
							if parsedTime, err := now.Parse(keyword); err == nil {
								conditions = append(conditions, fmt.Sprintf("%v = ?", scope.Quote(field.DBName)))
								keywords = append(keywords, parsedTime)
							}
						}
					case reflect.Ptr:
						// time ?
						if _, ok := field.Field.Interface().(*time.Time); ok {
							if parsedTime, err := now.Parse(keyword); err == nil {
								conditions = append(conditions, fmt.Sprintf("%v = ?", scope.Quote(field.DBName)))
								keywords = append(keywords, parsedTime)
							}
						}
					default:
						conditions = append(conditions, fmt.Sprintf("%v = ?", scope.Quote(field.DBName)))
						keywords = append(keywords, keyword)
					}
				}
			}

			if len(conditions) > 0 {
				return context.GetDB().Where(strings.Join(conditions, " OR "), keywords...)
			} else {
				return context.GetDB()
			}
		}
	}
	return res.searchAttrs
}

func (res *Resource) getCachedMetas(cacheKey string, fc func() []resource.Metaor) []*Meta {
	if res.cachedMetas == nil {
		res.cachedMetas = &map[string][]*Meta{}
	}

	if values, ok := (*res.cachedMetas)[cacheKey]; ok {
		return values
	} else {
		values := fc()
		var metas []*Meta
		for _, value := range values {
			metas = append(metas, value.(*Meta))
		}
		(*res.cachedMetas)[cacheKey] = metas
		return metas
	}
}

func (res *Resource) GetMetas(_attrs ...[]string) []resource.Metaor {
	var attrs []string
	for _, value := range _attrs {
		if value != nil {
			attrs = value
			break
		}
	}

	if attrs == nil {
		scope := &gorm.Scope{Value: res.Value}
		structFields := scope.GetModelStruct().StructFields
		attrs = []string{}

	Fields:
		for _, field := range structFields {
			for _, meta := range res.Metas {
				if field.Name == meta.Alias {
					attrs = append(attrs, meta.Name)
					continue Fields
				}
			}

			if field.IsForeignKey {
				continue
			}

			for _, value := range []string{"CreatedAt", "UpdatedAt", "DeletedAt"} {
				if value == field.Name {
					continue Fields
				}
			}

			attrs = append(attrs, field.Name)
		}

	MetaIncluded:
		for _, meta := range res.Metas {
			for _, attr := range attrs {
				if attr == meta.Alias || attr == meta.Name {
					continue MetaIncluded
				}
			}
			attrs = append(attrs, meta.Name)
		}
	}

	primaryKey := res.PrimaryFieldName()

	metas := []resource.Metaor{}
	for _, attr := range attrs {
		var meta *Meta
		for _, m := range res.Metas {
			if m.GetName() == attr {
				meta = m
				break
			}
		}

		if meta == nil {
			meta = &Meta{}
			meta.Name = attr
			meta.base = res
			if attr == primaryKey {
				meta.Type = "hidden"
			}
			meta.updateMeta()
		}

		metas = append(metas, meta)
	}

	return metas
}

func (res *Resource) GetMeta(name string) *Meta {
	for _, meta := range res.Metas {
		if meta.Name == name || meta.GetFieldName() == name {
			return meta
		}
	}
	return nil
}

func (res *Resource) indexMetas() []*Meta {
	return res.getCachedMetas("index_metas", func() []resource.Metaor {
		return res.GetMetas(res.indexAttrs, res.showAttrs)
	})
}

func (res *Resource) newMetas() []*Meta {
	return res.getCachedMetas("new_metas", func() []resource.Metaor {
		return res.GetMetas(res.newAttrs, res.editAttrs)
	})
}

func (res *Resource) editMetas() []*Meta {
	return res.getCachedMetas("edit_metas", func() []resource.Metaor {
		return res.GetMetas(res.editAttrs)
	})
}

func (res *Resource) showMetas() []*Meta {
	return res.getCachedMetas("show_metas", func() []resource.Metaor {
		return res.GetMetas(res.showAttrs, res.editAttrs)
	})
}

func (res *Resource) allMetas() []*Meta {
	return res.getCachedMetas("all_metas", func() []resource.Metaor {
		return res.GetMetas()
	})
}

func (res *Resource) allowedMetas(attrs []*Meta, context *Context, roles ...roles.PermissionMode) []*Meta {
	var metas = []*Meta{}
	for _, meta := range attrs {
		for _, role := range roles {
			if meta.HasPermission(role, context.Context) {
				metas = append(metas, meta)
				break
			}
		}
	}
	return metas
}

func (res *Resource) HasPermission(mode roles.PermissionMode, context *qor.Context) bool {
	if res.Config == nil || res.Config.Permission == nil {
		return true
	}
	return res.Config.Permission.HasPermission(mode, context.Roles...)
}
<commit_msg>Possible to ignore some columns<commit_after>package admin

import (
	"fmt"
	"reflect"
	"strconv"
	"strings"
	"time"

	"github.com/jinzhu/gorm"
	"github.com/jinzhu/now"
	"github.com/qor/qor"
	"github.com/qor/qor/resource"
	"github.com/qor/qor/roles"
	"github.com/qor/qor/utils"
)

type Resource struct {
	resource.Resource
	admin         *Admin
	Config        *Config
	Metas         []*Meta
	actions       []*Action
	scopes        []*Scope
	filters       map[string]*Filter
	searchAttrs   []string
	indexAttrs    []string
	newAttrs      []string
	editAttrs     []string
	showAttrs     []string
	cachedMetas   *map[string][]*Meta
	SearchHandler func(keyword string, context *qor.Context) *gorm.DB
}

func (res *Resource) Meta(meta *Meta) {
	if res.GetMeta(meta.Name) != nil {
		utils.ExitWithMsg("Duplicated meta %v defined for resource %v", meta.Name, res.Name)
	}

	meta.base = res
	meta.updateMeta()
	res.Metas = append(res.Metas, meta)
}

func (res Resource) GetAdmin() *Admin {
	return res.admin
}

func (res Resource) ToParam() string {
	return utils.ToParamString(res.Name)
}

func (res Resource) UseTheme(theme string) []string {
	if res.Config != nil {
		res.Config.Themes = append(res.Config.Themes, theme)
		return res.Config.Themes
	}
	return []string{}
}

func (res *Resource) convertObjectToMap(context *Context, value interface{}, kind string) interface{} {
	reflectValue := reflect.Indirect(reflect.ValueOf(value))
	switch reflectValue.Kind() {
	case reflect.Slice:
		values := []interface{}{}
		for i := 0; i < reflectValue.Len(); i++ {
			values = append(values, res.convertObjectToMap(context, reflectValue.Index(i).Interface(), kind))
		}
		return values
	case reflect.Struct:
		var metas []*Meta
		if kind == "index" {
			metas = res.indexMetas()
		} else if kind == "show" {
			metas = res.showMetas()
		}

		values := map[string]interface{}{}
		for _, meta := range metas {
			if meta.HasPermission(roles.Read, context.Context) {
				value := meta.GetValuer()(value, context.Context)
				if meta.Resource != nil {
					value = meta.Resource.(*Resource).convertObjectToMap(context, value, kind)
				}
				values[meta.GetName()] = value
			}
		}
		return values
	default:
		panic(fmt.Sprintf("Can't convert %v (%v) to map", reflectValue, reflectValue.Kind()))
	}
}

func (res *Resource) Decode(context *qor.Context, value interface{}) (errs []error) {
	return resource.Decode(context, value, res)
}

func (res *Resource) IndexAttrs(columns ...string) []string {
	if len(columns) > 0 {
		res.indexAttrs = columns
	}
	return res.indexAttrs
}

func (res *Resource) NewAttrs(columns ...string) []string {
	if len(columns) > 0 {
		res.newAttrs = columns
	}
	return res.newAttrs
}

func (res *Resource) EditAttrs(columns ...string) []string {
	if len(columns) > 0 {
		res.editAttrs = columns
	}
	return res.editAttrs
}

func (res *Resource) ShowAttrs(columns ...string) []string {
	if len(columns) > 0 {
		res.showAttrs = columns
	}
	return res.showAttrs
}

func (res *Resource) SearchAttrs(columns ...string) []string {
	if len(columns) > 0 {
		res.searchAttrs = columns
		res.SearchHandler = func(keyword string, context *qor.Context) *gorm.DB {
			db := context.GetDB()
			var conditions []string
			var keywords []interface{}
			scope := db.NewScope(res.Value)

			for _, column := range columns {
				if field, ok := scope.FieldByName(column); ok {
					switch field.Field.Kind() {
					case reflect.String:
						conditions = append(conditions, fmt.Sprintf("upper(%v) like upper(?)", scope.Quote(field.DBName)))
						keywords = append(keywords, "%"+keyword+"%")
					case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64, reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
						if _, err := strconv.Atoi(keyword); err == nil {
							conditions = append(conditions, fmt.Sprintf("%v = ?", scope.Quote(field.DBName)))
							keywords = append(keywords, keyword)
						}
					case reflect.Float32, reflect.Float64:
						if _, err := strconv.ParseFloat(keyword, 64); err == nil {
							conditions = append(conditions, fmt.Sprintf("%v = ?", scope.Quote(field.DBName)))
							keywords = append(keywords, keyword)
						}
					case reflect.Struct:
						// time ?
						if _, ok := field.Field.Interface().(time.Time); ok {
							if parsedTime, err := now.Parse(keyword); err == nil {
								conditions = append(conditions, fmt.Sprintf("%v = ?", scope.Quote(field.DBName)))
								keywords = append(keywords, parsedTime)
							}
						}
					case reflect.Ptr:
						// time ?
						if _, ok := field.Field.Interface().(*time.Time); ok {
							if parsedTime, err := now.Parse(keyword); err == nil {
								conditions = append(conditions, fmt.Sprintf("%v = ?", scope.Quote(field.DBName)))
								keywords = append(keywords, parsedTime)
							}
						}
					default:
						conditions = append(conditions, fmt.Sprintf("%v = ?", scope.Quote(field.DBName)))
						keywords = append(keywords, keyword)
					}
				}
			}

			if len(conditions) > 0 {
				return context.GetDB().Where(strings.Join(conditions, " OR "), keywords...)
			} else {
				return context.GetDB()
			}
		}
	}
	return res.searchAttrs
}

func (res *Resource) getCachedMetas(cacheKey string, fc func() []resource.Metaor) []*Meta {
	if res.cachedMetas == nil {
		res.cachedMetas = &map[string][]*Meta{}
	}

	if values, ok := (*res.cachedMetas)[cacheKey]; ok {
		return values
	} else {
		values := fc()
		var metas []*Meta
		for _, value := range values {
			metas = append(metas, value.(*Meta))
		}
		(*res.cachedMetas)[cacheKey] = metas
		return metas
	}
}

func (res *Resource) GetMetas(_attrs ...[]string) []resource.Metaor {
	var attrs, ignoredAttrs []string
	for _, value := range _attrs {
		if value != nil {
			for _, v := range value {
				if strings.HasPrefix(v, "-") {
					ignoredAttrs = append(ignoredAttrs, strings.TrimLeft(v, "-"))
				} else {
					attrs = append(attrs, v)
				}
			}
			break
		}
	}

	if attrs == nil {
		scope := &gorm.Scope{Value: res.Value}
		structFields := scope.GetModelStruct().StructFields
		attrs = []string{}

	Fields:
		for _, field := range structFields {
			for _, attr := range ignoredAttrs {
				if attr == field.Name {
					continue Fields
				}
			}

			for _, meta := range res.Metas {
				if field.Name == meta.Alias {
					attrs = append(attrs, meta.Name)
					continue Fields
				}
			}

			if field.IsForeignKey {
				continue
			}

			for _, value := range []string{"CreatedAt", "UpdatedAt", "DeletedAt"} {
				if value == field.Name {
					continue Fields
				}
			}

			attrs = append(attrs, field.Name)
		}

	MetaIncluded:
		for _, meta := range res.Metas {
			for _, attr := range ignoredAttrs {
				if attr == meta.Name || attr == meta.Alias {
					continue MetaIncluded
				}
			}

			for _, attr := range attrs {
				if attr == meta.Alias || attr == meta.Name {
					continue MetaIncluded
				}
			}
			attrs = append(attrs, meta.Name)
		}
	}

	primaryKey := res.PrimaryFieldName()

	metas := []resource.Metaor{}
	for _, attr := range attrs {
		var meta *Meta
		for _, m := range res.Metas {
			if m.GetName() == attr {
				meta = m
				break
			}
		}

		if meta == nil {
			meta = &Meta{}
			meta.Name = attr
			meta.base = res
			if attr == primaryKey {
				meta.Type = "hidden"
			}
			meta.updateMeta()
		}

		metas = append(metas, meta)
	}

	return metas
}

func (res *Resource) GetMeta(name string) *Meta {
	for _, meta := range res.Metas {
		if meta.Name == name || meta.GetFieldName() == name {
			return meta
		}
	}
	return nil
}

func (res *Resource) indexMetas() []*Meta {
	return res.getCachedMetas("index_metas", func() []resource.Metaor {
		return res.GetMetas(res.indexAttrs, res.showAttrs)
	})
}

func (res *Resource) newMetas() []*Meta {
	return res.getCachedMetas("new_metas", func() []resource.Metaor {
		return res.GetMetas(res.newAttrs, res.editAttrs)
	})
}

func (res *Resource) editMetas() []*Meta {
	return res.getCachedMetas("edit_metas", func() []resource.Metaor {
		return res.GetMetas(res.editAttrs)
	})
}

func (res *Resource) showMetas() []*Meta {
	return res.getCachedMetas("show_metas", func() []resource.Metaor {
		return res.GetMetas(res.showAttrs, res.editAttrs)
	})
}

func (res *Resource) allMetas() []*Meta {
	return res.getCachedMetas("all_metas", func() []resource.Metaor {
		return res.GetMetas()
	})
}

func (res *Resource) allowedMetas(attrs []*Meta, context *Context, roles ...roles.PermissionMode) []*Meta {
	var metas = []*Meta{}
	for _, meta := range attrs {
		for _, role := range roles {
			if meta.HasPermission(role, context.Context) {
				metas = append(metas, meta)
				break
			}
		}
	}
	return metas
}

func (res *Resource) HasPermission(mode roles.PermissionMode, context *qor.Context) bool {
	if res.Config == nil || res.Config.Permission == nil {
		return true
	}
	return res.Config.Permission.HasPermission(mode, context.Roles...)
}
<|endoftext|>


<commit_before>package etcd

import (
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"errors"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"os"
	"path"
	"time"
)

// See SetConsistency for how to use these constants.
const (
	// Using strings rather than iota because the consistency level
	// could be persisted to disk, so it'd be better to use
	// human-readable values.
	STRONG_CONSISTENCY = "STRONG"
	WEAK_CONSISTENCY   = "WEAK"
)

const (
	defaultBufferSize = 10
)

type Config struct {
	CertFile    string        `json:"certFile"`
	KeyFile     string        `json:"keyFile"`
	CaCertFile  []string      `json:"caCertFiles"`
	DialTimeout time.Duration `json:"timeout"`
	Consistency string        `json:"consistency"`
}

type Client struct {
	config      Config   `json:"config"`
	cluster     *Cluster `json:"cluster"`
	httpClient  *http.Client
	persistence io.Writer
	cURLch      chan string
	// CheckRetry can be used to control the policy for failed requests
	// and modify the cluster if needed.
	// The client calls it before sending requests again, and
	// stops retrying if CheckRetry returns some error. The cases that
	// this function needs to handle include no response and unexpected
	// http status code of response.
	// If CheckRetry is nil, client will call the default one
	// `DefaultCheckRetry`.
	// Argument cluster is the etcd.Cluster object that these requests have been made on.
	// Argument numReqs is the number of http.Requests that have been made so far.
	// Argument lastResp is the http.Responses from the last request.
	// Argument err is the reason of the failure.
	CheckRetry func(cluster *Cluster, numReqs int,
		lastResp http.Response, err error) error
}

// NewClient create a basic client that is configured to be used
// with the given machine list.
func NewClient(machines []string) *Client {
	config := Config{
		// default timeout is one second
		DialTimeout: time.Second,
		// default consistency level is STRONG
		Consistency: STRONG_CONSISTENCY,
	}

	client := &Client{
		cluster: NewCluster(machines),
		config:  config,
	}

	client.initHTTPClient()
	client.saveConfig()

	return client
}

// NewTLSClient create a basic client with TLS configuration
func NewTLSClient(machines []string, cert, key, caCert string) (*Client, error) {
	// overwrite the default machine to use https
	if len(machines) == 0 {
		machines = []string{"https://127.0.0.1:4001"}
	}

	config := Config{
		// default timeout is one second
		DialTimeout: time.Second,
		// default consistency level is STRONG
		Consistency: STRONG_CONSISTENCY,
		CertFile:    cert,
		KeyFile:     key,
		CaCertFile:  make([]string, 0),
	}

	client := &Client{
		cluster: NewCluster(machines),
		config:  config,
	}

	err := client.initHTTPSClient(cert, key)
	if err != nil {
		return nil, err
	}

	err = client.AddRootCA(caCert)

	client.saveConfig()

	return client, nil
}

// NewClientFromFile creates a client from a given file path.
// The given file is expected to use the JSON format.
func NewClientFromFile(fpath string) (*Client, error) {
	fi, err := os.Open(fpath)
	if err != nil {
		return nil, err
	}

	defer func() {
		if err := fi.Close(); err != nil {
			panic(err)
		}
	}()

	return NewClientFromReader(fi)
}

// NewClientFromReader creates a Client configured from a given reader.
// The configuration is expected to use the JSON format.
func NewClientFromReader(reader io.Reader) (*Client, error) {
	c := new(Client)

	b, err := ioutil.ReadAll(reader)
	if err != nil {
		return nil, err
	}

	err = json.Unmarshal(b, c)
	if err != nil {
		return nil, err
	}
	if c.config.CertFile == "" {
		c.initHTTPClient()
	} else {
		err = c.initHTTPSClient(c.config.CertFile, c.config.KeyFile)
	}

	if err != nil {
		return nil, err
	}

	for _, caCert := range c.config.CaCertFile {
		if err := c.AddRootCA(caCert); err != nil {
			return nil, err
		}
	}

	return c, nil
}

// Override the Client's HTTP Transport object
func (c *Client) SetTransport(tr *http.Transport) {
	c.httpClient.Transport = tr
}

// initHTTPClient initializes a HTTP client for etcd client
func (c *Client) initHTTPClient() {
	tr := &http.Transport{
		Dial: c.dial,
		TLSClientConfig: &tls.Config{
			InsecureSkipVerify: true,
		},
	}
	c.httpClient = &http.Client{Transport: tr}
}

// initHTTPClient initializes a HTTPS client for etcd client
func (c *Client) initHTTPSClient(cert, key string) error {
	if cert == "" || key == "" {
		return errors.New("Require both cert and key path")
	}

	tlsCert, err := tls.LoadX509KeyPair(cert, key)
	if err != nil {
		return err
	}

	tlsConfig := &tls.Config{
		Certificates:       []tls.Certificate{tlsCert},
		InsecureSkipVerify: true,
	}

	tr := &http.Transport{
		TLSClientConfig: tlsConfig,
		Dial:            c.dial,
	}

	c.httpClient = &http.Client{Transport: tr}
	return nil
}

// SetPersistence sets a writer to which the config will be
// written every time it's changed.
func (c *Client) SetPersistence(writer io.Writer) {
	c.persistence = writer
}

// SetConsistency changes the consistency level of the client.
//
// When consistency is set to STRONG_CONSISTENCY, all requests,
// including GET, are sent to the leader.  This means that, assuming
// the absence of leader failures, GET requests are guaranteed to see
// the changes made by previous requests.
//
// When consistency is set to WEAK_CONSISTENCY, other requests
// are still sent to the leader, but GET requests are sent to a
// random server from the server pool.  This reduces the read
// load on the leader, but it's not guaranteed that the GET requests
// will see changes made by previous requests (they might have not
// yet been committed on non-leader servers).
func (c *Client) SetConsistency(consistency string) error {
	if !(consistency == STRONG_CONSISTENCY || consistency == WEAK_CONSISTENCY) {
		return errors.New("The argument must be either STRONG_CONSISTENCY or WEAK_CONSISTENCY.")
	}
	c.config.Consistency = consistency
	return nil
}

// AddRootCA adds a root CA cert for the etcd client
func (c *Client) AddRootCA(caCert string) error {
	if c.httpClient == nil {
		return errors.New("Client has not been initialized yet!")
	}

	certBytes, err := ioutil.ReadFile(caCert)
	if err != nil {
		return err
	}

	tr, ok := c.httpClient.Transport.(*http.Transport)

	if !ok {
		panic("AddRootCA(): Transport type assert should not fail")
	}

	if tr.TLSClientConfig.RootCAs == nil {
		caCertPool := x509.NewCertPool()
		ok = caCertPool.AppendCertsFromPEM(certBytes)
		if ok {
			tr.TLSClientConfig.RootCAs = caCertPool
		}
		tr.TLSClientConfig.InsecureSkipVerify = false
	} else {
		ok = tr.TLSClientConfig.RootCAs.AppendCertsFromPEM(certBytes)
	}

	if !ok {
		err = errors.New("Unable to load caCert")
	}

	c.config.CaCertFile = append(c.config.CaCertFile, caCert)
	c.saveConfig()

	return err
}

// SetCluster updates cluster information using the given machine list.
func (c *Client) SetCluster(machines []string) bool {
	success := c.internalSyncCluster(machines)
	return success
}

func (c *Client) GetCluster() []string {
	return c.cluster.Machines
}

// SyncCluster updates the cluster information using the internal machine list.
func (c *Client) SyncCluster() bool {
	return c.internalSyncCluster(c.cluster.Machines)
}

// internalSyncCluster syncs cluster information using the given machine list.
func (c *Client) internalSyncCluster(machines []string) bool {
	for _, machine := range machines {
		httpPath := c.createHttpPath(machine, path.Join(version, "machines"))
		resp, err := c.httpClient.Get(httpPath)
		if err != nil {
			// try another machine in the cluster
			continue
		} else {
			b, err := ioutil.ReadAll(resp.Body)
			resp.Body.Close()
			if err != nil {
				// try another machine in the cluster
				continue
			}

			// update Machines List
			c.cluster.updateFromStr(string(b))

			// update leader
			// the first one in the machine list is the leader
			c.cluster.switchLeader(0)

			logger.Debug("sync.machines ", c.cluster.Machines)
			c.saveConfig()
			return true
		}
	}
	return false
}

// createHttpPath creates a complete HTTP URL.
// serverName should contain both the host name and a port number, if any.
func (c *Client) createHttpPath(serverName string, _path string) string {
	u, err := url.Parse(serverName)
	if err != nil {
		panic(err)
	}

	u.Path = path.Join(u.Path, _path)

	if u.Scheme == "" {
		u.Scheme = "http"
	}
	return u.String()
}

// dial attempts to open a TCP connection to the provided address, explicitly
// enabling keep-alives with a one-second interval.
func (c *Client) dial(network, addr string) (net.Conn, error) {
	conn, err := net.DialTimeout(network, addr, c.config.DialTimeout)
	if err != nil {
		return nil, err
	}

	tcpConn, ok := conn.(*net.TCPConn)
	if !ok {
		return nil, errors.New("Failed type-assertion of net.Conn as *net.TCPConn")
	}

	// Keep TCP alive to check whether or not the remote machine is down
	if err = tcpConn.SetKeepAlive(true); err != nil {
		return nil, err
	}

	if err = tcpConn.SetKeepAlivePeriod(time.Second); err != nil {
		return nil, err
	}

	return tcpConn, nil
}

func (c *Client) OpenCURL() {
	c.cURLch = make(chan string, defaultBufferSize)
}

func (c *Client) CloseCURL() {
	c.cURLch = nil
}

func (c *Client) sendCURL(command string) {
	go func() {
		select {
		case c.cURLch <- command:
		default:
		}
	}()
}

func (c *Client) RecvCURL() string {
	return <-c.cURLch
}

// saveConfig saves the current config using c.persistence.
func (c *Client) saveConfig() error {
	if c.persistence != nil {
		b, err := json.Marshal(c)
		if err != nil {
			return err
		}

		_, err = c.persistence.Write(b)
		if err != nil {
			return err
		}
	}

	return nil
}

// MarshalJSON implements the Marshaller interface
// as defined by the standard JSON package.
func (c *Client) MarshalJSON() ([]byte, error) {
	b, err := json.Marshal(struct {
		Config  Config   `json:"config"`
		Cluster *Cluster `json:"cluster"`
	}{
		Config:  c.config,
		Cluster: c.cluster,
	})

	if err != nil {
		return nil, err
	}

	return b, nil
}

// UnmarshalJSON implements the Unmarshaller interface
// as defined by the standard JSON package.
func (c *Client) UnmarshalJSON(b []byte) error {
	temp := struct {
		Config  Config   `json:"config"`
		Cluster *Cluster `json:"cluster"`
	}{}
	err := json.Unmarshal(b, &temp)
	if err != nil {
		return err
	}

	c.cluster = temp.Cluster
	c.config = temp.Config
	return nil
}
<commit_msg>feat(client): add SetDialTimeout method<commit_after>package etcd

import (
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"errors"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"os"
	"path"
	"time"
)

// See SetConsistency for how to use these constants.
const (
	// Using strings rather than iota because the consistency level
	// could be persisted to disk, so it'd be better to use
	// human-readable values.
	STRONG_CONSISTENCY = "STRONG"
	WEAK_CONSISTENCY   = "WEAK"
)

const (
	defaultBufferSize = 10
)

type Config struct {
	CertFile    string        `json:"certFile"`
	KeyFile     string        `json:"keyFile"`
	CaCertFile  []string      `json:"caCertFiles"`
	DialTimeout time.Duration `json:"timeout"`
	Consistency string        `json:"consistency"`
}

type Client struct {
	config      Config   `json:"config"`
	cluster     *Cluster `json:"cluster"`
	httpClient  *http.Client
	persistence io.Writer
	cURLch      chan string
	// CheckRetry can be used to control the policy for failed requests
	// and modify the cluster if needed.
	// The client calls it before sending requests again, and
	// stops retrying if CheckRetry returns some error. The cases that
	// this function needs to handle include no response and unexpected
	// http status code of response.
	// If CheckRetry is nil, client will call the default one
	// `DefaultCheckRetry`.
	// Argument cluster is the etcd.Cluster object that these requests have been made on.
	// Argument numReqs is the number of http.Requests that have been made so far.
	// Argument lastResp is the http.Responses from the last request.
	// Argument err is the reason of the failure.
	CheckRetry func(cluster *Cluster, numReqs int,
		lastResp http.Response, err error) error
}

// NewClient create a basic client that is configured to be used
// with the given machine list.
func NewClient(machines []string) *Client {
	config := Config{
		// default timeout is one second
		DialTimeout: time.Second,
		// default consistency level is STRONG
		Consistency: STRONG_CONSISTENCY,
	}

	client := &Client{
		cluster: NewCluster(machines),
		config:  config,
	}

	client.initHTTPClient()
	client.saveConfig()

	return client
}

// NewTLSClient create a basic client with TLS configuration
func NewTLSClient(machines []string, cert, key, caCert string) (*Client, error) {
	// overwrite the default machine to use https
	if len(machines) == 0 {
		machines = []string{"https://127.0.0.1:4001"}
	}

	config := Config{
		// default timeout is one second
		DialTimeout: time.Second,
		// default consistency level is STRONG
		Consistency: STRONG_CONSISTENCY,
		CertFile:    cert,
		KeyFile:     key,
		CaCertFile:  make([]string, 0),
	}

	client := &Client{
		cluster: NewCluster(machines),
		config:  config,
	}

	err := client.initHTTPSClient(cert, key)
	if err != nil {
		return nil, err
	}

	err = client.AddRootCA(caCert)

	client.saveConfig()

	return client, nil
}

// NewClientFromFile creates a client from a given file path.
// The given file is expected to use the JSON format.
func NewClientFromFile(fpath string) (*Client, error) {
	fi, err := os.Open(fpath)
	if err != nil {
		return nil, err
	}

	defer func() {
		if err := fi.Close(); err != nil {
			panic(err)
		}
	}()

	return NewClientFromReader(fi)
}

// NewClientFromReader creates a Client configured from a given reader.
// The configuration is expected to use the JSON format.
func NewClientFromReader(reader io.Reader) (*Client, error) {
	c := new(Client)

	b, err := ioutil.ReadAll(reader)
	if err != nil {
		return nil, err
	}

	err = json.Unmarshal(b, c)
	if err != nil {
		return nil, err
	}
	if c.config.CertFile == "" {
		c.initHTTPClient()
	} else {
		err = c.initHTTPSClient(c.config.CertFile, c.config.KeyFile)
	}

	if err != nil {
		return nil, err
	}

	for _, caCert := range c.config.CaCertFile {
		if err := c.AddRootCA(caCert); err != nil {
			return nil, err
		}
	}

	return c, nil
}

// Override the Client's HTTP Transport object
func (c *Client) SetTransport(tr *http.Transport) {
	c.httpClient.Transport = tr
}

// initHTTPClient initializes a HTTP client for etcd client
func (c *Client) initHTTPClient() {
	tr := &http.Transport{
		Dial: c.dial,
		TLSClientConfig: &tls.Config{
			InsecureSkipVerify: true,
		},
	}
	c.httpClient = &http.Client{Transport: tr}
}

// initHTTPClient initializes a HTTPS client for etcd client
func (c *Client) initHTTPSClient(cert, key string) error {
	if cert == "" || key == "" {
		return errors.New("Require both cert and key path")
	}

	tlsCert, err := tls.LoadX509KeyPair(cert, key)
	if err != nil {
		return err
	}

	tlsConfig := &tls.Config{
		Certificates:       []tls.Certificate{tlsCert},
		InsecureSkipVerify: true,
	}

	tr := &http.Transport{
		TLSClientConfig: tlsConfig,
		Dial:            c.dial,
	}

	c.httpClient = &http.Client{Transport: tr}
	return nil
}

// SetPersistence sets a writer to which the config will be
// written every time it's changed.
func (c *Client) SetPersistence(writer io.Writer) {
	c.persistence = writer
}

// SetConsistency changes the consistency level of the client.
//
// When consistency is set to STRONG_CONSISTENCY, all requests,
// including GET, are sent to the leader.  This means that, assuming
// the absence of leader failures, GET requests are guaranteed to see
// the changes made by previous requests.
//
// When consistency is set to WEAK_CONSISTENCY, other requests
// are still sent to the leader, but GET requests are sent to a
// random server from the server pool.  This reduces the read
// load on the leader, but it's not guaranteed that the GET requests
// will see changes made by previous requests (they might have not
// yet been committed on non-leader servers).
func (c *Client) SetConsistency(consistency string) error {
	if !(consistency == STRONG_CONSISTENCY || consistency == WEAK_CONSISTENCY) {
		return errors.New("The argument must be either STRONG_CONSISTENCY or WEAK_CONSISTENCY.")
	}
	c.config.Consistency = consistency
	return nil
}

// Sets the DialTimeout value
func (c *Client) SetDialTimeout(d time.Duration) {
	c.config.DialTimeout = d
}

// AddRootCA adds a root CA cert for the etcd client
func (c *Client) AddRootCA(caCert string) error {
	if c.httpClient == nil {
		return errors.New("Client has not been initialized yet!")
	}

	certBytes, err := ioutil.ReadFile(caCert)
	if err != nil {
		return err
	}

	tr, ok := c.httpClient.Transport.(*http.Transport)

	if !ok {
		panic("AddRootCA(): Transport type assert should not fail")
	}

	if tr.TLSClientConfig.RootCAs == nil {
		caCertPool := x509.NewCertPool()
		ok = caCertPool.AppendCertsFromPEM(certBytes)
		if ok {
			tr.TLSClientConfig.RootCAs = caCertPool
		}
		tr.TLSClientConfig.InsecureSkipVerify = false
	} else {
		ok = tr.TLSClientConfig.RootCAs.AppendCertsFromPEM(certBytes)
	}

	if !ok {
		err = errors.New("Unable to load caCert")
	}

	c.config.CaCertFile = append(c.config.CaCertFile, caCert)
	c.saveConfig()

	return err
}

// SetCluster updates cluster information using the given machine list.
func (c *Client) SetCluster(machines []string) bool {
	success := c.internalSyncCluster(machines)
	return success
}

func (c *Client) GetCluster() []string {
	return c.cluster.Machines
}

// SyncCluster updates the cluster information using the internal machine list.
func (c *Client) SyncCluster() bool {
	return c.internalSyncCluster(c.cluster.Machines)
}

// internalSyncCluster syncs cluster information using the given machine list.
func (c *Client) internalSyncCluster(machines []string) bool {
	for _, machine := range machines {
		httpPath := c.createHttpPath(machine, path.Join(version, "machines"))
		resp, err := c.httpClient.Get(httpPath)
		if err != nil {
			// try another machine in the cluster
			continue
		} else {
			b, err := ioutil.ReadAll(resp.Body)
			resp.Body.Close()
			if err != nil {
				// try another machine in the cluster
				continue
			}

			// update Machines List
			c.cluster.updateFromStr(string(b))

			// update leader
			// the first one in the machine list is the leader
			c.cluster.switchLeader(0)

			logger.Debug("sync.machines ", c.cluster.Machines)
			c.saveConfig()
			return true
		}
	}
	return false
}

// createHttpPath creates a complete HTTP URL.
// serverName should contain both the host name and a port number, if any.
func (c *Client) createHttpPath(serverName string, _path string) string {
	u, err := url.Parse(serverName)
	if err != nil {
		panic(err)
	}

	u.Path = path.Join(u.Path, _path)

	if u.Scheme == "" {
		u.Scheme = "http"
	}
	return u.String()
}

// dial attempts to open a TCP connection to the provided address, explicitly
// enabling keep-alives with a one-second interval.
func (c *Client) dial(network, addr string) (net.Conn, error) {
	conn, err := net.DialTimeout(network, addr, c.config.DialTimeout)
	if err != nil {
		return nil, err
	}

	tcpConn, ok := conn.(*net.TCPConn)
	if !ok {
		return nil, errors.New("Failed type-assertion of net.Conn as *net.TCPConn")
	}

	// Keep TCP alive to check whether or not the remote machine is down
	if err = tcpConn.SetKeepAlive(true); err != nil {
		return nil, err
	}

	if err = tcpConn.SetKeepAlivePeriod(time.Second); err != nil {
		return nil, err
	}

	return tcpConn, nil
}

func (c *Client) OpenCURL() {
	c.cURLch = make(chan string, defaultBufferSize)
}

func (c *Client) CloseCURL() {
	c.cURLch = nil
}

func (c *Client) sendCURL(command string) {
	go func() {
		select {
		case c.cURLch <- command:
		default:
		}
	}()
}

func (c *Client) RecvCURL() string {
	return <-c.cURLch
}

// saveConfig saves the current config using c.persistence.
func (c *Client) saveConfig() error {
	if c.persistence != nil {
		b, err := json.Marshal(c)
		if err != nil {
			return err
		}

		_, err = c.persistence.Write(b)
		if err != nil {
			return err
		}
	}

	return nil
}

// MarshalJSON implements the Marshaller interface
// as defined by the standard JSON package.
func (c *Client) MarshalJSON() ([]byte, error) {
	b, err := json.Marshal(struct {
		Config  Config   `json:"config"`
		Cluster *Cluster `json:"cluster"`
	}{
		Config:  c.config,
		Cluster: c.cluster,
	})

	if err != nil {
		return nil, err
	}

	return b, nil
}

// UnmarshalJSON implements the Unmarshaller interface
// as defined by the standard JSON package.
func (c *Client) UnmarshalJSON(b []byte) error {
	temp := struct {
		Config  Config   `json:"config"`
		Cluster *Cluster `json:"cluster"`
	}{}
	err := json.Unmarshal(b, &temp)
	if err != nil {
		return err
	}

	c.cluster = temp.Cluster
	c.config = temp.Config
	return nil
}
<|endoftext|>


<commit_before>package integration_test

import (
	"encoding/json"
	"fmt"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gexec"
	"golang.org/x/crypto/ssh"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"os/exec"
	"path/filepath"
	"srcd.works/go-git.v4"
	gitssh "srcd.works/go-git.v4/plumbing/transport/ssh"
)

var _ = Describe("Claimer", func() {
	var claimer string
	var gitDir string

	BeforeSuite(func() {
		var err error

		gitDir, err = ioutil.TempDir("", "claimer-integration-tests")
		Expect(err).NotTo(HaveOccurred())

		claimer, err = gexec.Build(filepath.Join("github.com", "mdelillo", "claimer"))
		Expect(err).NotTo(HaveOccurred())
	})

	AfterSuite(func() {
		gexec.CleanupBuildArtifacts()
		os.RemoveAll(gitDir)
	})

	It("claims and releases locks", func() {
		apiToken := getEnv("CLAIMER_TEST_API_TOKEN")
		repoUrl := getEnv("CLAIMER_TEST_REPO_URL")
		deployKey := getEnv("CLAIMER_TEST_DEPLOY_KEY")
		channelId := getEnv("CLAIMER_TEST_CHANNEL_ID")
		botId := getEnv("CLAIMER_TEST_BOT_ID")

		signer, err := ssh.ParsePrivateKey([]byte(deployKey))
		Expect(err).NotTo(HaveOccurred())

		_, err = git.PlainClone(gitDir, false, &git.CloneOptions{
			URL: repoUrl,
			Auth: &gitssh.PublicKeys{
				User:   "git",
				Signer: signer,
			},
		})
		Expect(err).NotTo(HaveOccurred())

		resetClaimerTestPool(gitDir)

		claimerCommand := exec.Command(
			claimer,
			"-apiToken", apiToken,
			"-repoUrl", repoUrl,
			"-deployKey", deployKey,
		)
		session, err := gexec.Start(claimerCommand, GinkgoWriter, GinkgoWriter)
		Expect(err).NotTo(HaveOccurred())

		By("Checking the status")
		postSlackMessage(fmt.Sprintf("<@%s> status", botId), channelId, apiToken)
		Eventually(func() string { return latestSlackMessage(channelId, apiToken) }, "10s").
			Should(Equal("*Claimed:* \n*Unclaimed:* pool-1, pool-2"))

		By("Claiming pool-1")
		postSlackMessage(fmt.Sprintf("<@%s> claim pool-1", botId), channelId, apiToken)
		Eventually(func() string { return latestSlackMessage(channelId, apiToken) }, "10s").
			Should(Equal("Claimed pool-1"))
		updateGitRepo(gitDir)
		Expect(filepath.Join(gitDir, "pool-1", "claimed", "lock-a")).To(BeAnExistingFile())
		Expect(filepath.Join(gitDir, "pool-1", "unclaimed", "lock-a")).NotTo(BeAnExistingFile())

		By("Checking the status")
		postSlackMessage(fmt.Sprintf("<@%s> status", botId), channelId, apiToken)
		Eventually(func() string { return latestSlackMessage(channelId, apiToken) }, "10s").
			Should(Equal("*Claimed:* pool-1\n*Unclaimed:* pool-2"))

		By("Trying to claim pool-1 again")
		postSlackMessage(fmt.Sprintf("<@%s> claim pool-1", botId), channelId, apiToken)
		Eventually(func() string { return latestSlackMessage(channelId, apiToken) }, "10s").
			Should(Equal("pool-1 is not available"))

		By("Releasing pool-1")
		postSlackMessage(fmt.Sprintf("<@%s> release pool-1", botId), channelId, apiToken)
		Eventually(func() string { return latestSlackMessage(channelId, apiToken) }, "10s").
			Should(Equal("Released pool-1"))
		updateGitRepo(gitDir)
		Expect(filepath.Join(gitDir, "pool-1", "unclaimed", "lock-a")).To(BeAnExistingFile())
		Expect(filepath.Join(gitDir, "pool-1", "claimed", "lock-a")).NotTo(BeAnExistingFile())

		By("Checking the status")
		postSlackMessage(fmt.Sprintf("<@%s> status", botId), channelId, apiToken)
		Eventually(func() string { return latestSlackMessage(channelId, apiToken) }, "10s").
			Should(Equal("*Claimed:* \n*Unclaimed:* pool-1, pool-2"))

		By("Trying to release pool-1 again")
		postSlackMessage(fmt.Sprintf("<@%s> release pool-1", botId), channelId, apiToken)
		Eventually(func() string { return latestSlackMessage(channelId, apiToken) }, "10s").
			Should(Equal("pool-1 is not claimed"))

		By("Trying to claim non-existent pool")
		postSlackMessage(fmt.Sprintf("<@%s> claim non-existent-pool", botId), channelId, apiToken)
		Eventually(func() string { return latestSlackMessage(channelId, apiToken) }, "10s").
			Should(Equal("non-existent-pool is not available"))

		By("Trying to release non-existent-pool")
		postSlackMessage(fmt.Sprintf("<@%s> release non-existent-pool", botId), channelId, apiToken)
		Eventually(func() string { return latestSlackMessage(channelId, apiToken) }, "10s").
			Should(Equal("non-existent-pool is not claimed"))

		// @claimer claim pool-2
		// assert about error
		// assert about repo
		// @claimer status
		// assert about status

		session.Terminate().Wait()

		resetClaimerTestPool(gitDir)
	})
})

func getEnv(name string) string {
	value, ok := os.LookupEnv(name)
	if !ok {
		Fail(fmt.Sprintf("%s must be set", name))
	}
	return value
}

func postSlackMessage(text, channelId, apiToken string) {
	resp, err := http.PostForm(
		"https://slack.com/api/chat.postMessage",
		url.Values{
			"token":    {apiToken},
			"channel":  {channelId},
			"text":     {text},
			"as_user":  {"false"},
			"username": {"claimer-integration-test"},
		},
	)
	Expect(err).NotTo(HaveOccurred())
	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)
	Expect(err).NotTo(HaveOccurred())

	var slackResponse struct {
		Ok    bool
		Error string
	}
	Expect(json.Unmarshal(body, &slackResponse)).To(Succeed())

	Expect(slackResponse.Ok).To(BeTrue(), fmt.Sprintf("Posting to slack failed: %s", slackResponse.Error))
}

func latestSlackMessage(channelId, apiToken string) string {
	resp, err := http.PostForm(
		"https://slack.com/api/channels.history",
		url.Values{
			"token":   {apiToken},
			"channel": {channelId},
			"count":   {"1"},
		},
	)
	Expect(err).NotTo(HaveOccurred())
	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)
	Expect(err).NotTo(HaveOccurred())

	var slackResponse struct {
		Ok       bool
		Messages []struct {
			Text string
		}
		Error string
	}
	Expect(json.Unmarshal(body, &slackResponse)).To(Succeed())

	Expect(slackResponse.Ok).To(BeTrue(), fmt.Sprintf("Getting message from slack failed: %s", slackResponse.Error))

	return slackResponse.Messages[0].Text
}

func updateGitRepo(gitDir string) {
	runGitCommand(gitDir, "fetch")
	runGitCommand(gitDir, "reset", "--hard", "origin/master")
}

func resetClaimerTestPool(gitDir string) {
	runGitCommand(gitDir, "checkout", "master")
	runGitCommand(gitDir, "reset", "--hard", "initial-state")
	runGitCommand(gitDir, "push", "--force", "origin", "master")

}

func runGitCommand(dir string, args ...string) {
	cmd := exec.Command("git", args...)
	cmd.Dir = dir
	ExpectWithOffset(1, cmd.Run()).To(Succeed())
}
<commit_msg>Ensure process is killed after integration test<commit_after>package integration_test

import (
	"encoding/json"
	"fmt"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gexec"
	"golang.org/x/crypto/ssh"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"os/exec"
	"path/filepath"
	"srcd.works/go-git.v4"
	gitssh "srcd.works/go-git.v4/plumbing/transport/ssh"
)

var _ = Describe("Claimer", func() {
	var (
		claimer string
		gitDir string
	)

	BeforeSuite(func() {
		var err error

		gitDir, err = ioutil.TempDir("", "claimer-integration-tests")
		Expect(err).NotTo(HaveOccurred())

		claimer, err = gexec.Build(filepath.Join("github.com", "mdelillo", "claimer"))
		Expect(err).NotTo(HaveOccurred())
	})

	AfterSuite(func() {
		gexec.KillAndWait()
		gexec.CleanupBuildArtifacts()
		os.RemoveAll(gitDir)
	})

	It("claims and releases locks", func() {
		apiToken := getEnv("CLAIMER_TEST_API_TOKEN")
		repoUrl := getEnv("CLAIMER_TEST_REPO_URL")
		deployKey := getEnv("CLAIMER_TEST_DEPLOY_KEY")
		channelId := getEnv("CLAIMER_TEST_CHANNEL_ID")
		botId := getEnv("CLAIMER_TEST_BOT_ID")

		signer, err := ssh.ParsePrivateKey([]byte(deployKey))
		Expect(err).NotTo(HaveOccurred())

		_, err = git.PlainClone(gitDir, false, &git.CloneOptions{
			URL: repoUrl,
			Auth: &gitssh.PublicKeys{
				User:   "git",
				Signer: signer,
			},
		})
		Expect(err).NotTo(HaveOccurred())
		defer resetClaimerTestPool(gitDir)

		resetClaimerTestPool(gitDir)

		claimerCommand := exec.Command(
			claimer,
			"-apiToken", apiToken,
			"-repoUrl", repoUrl,
			"-deployKey", deployKey,
		)
		_, err = gexec.Start(claimerCommand, GinkgoWriter, GinkgoWriter)
		Expect(err).NotTo(HaveOccurred())

		By("Checking the status")
		postSlackMessage(fmt.Sprintf("<@%s> status", botId), channelId, apiToken)
		Eventually(func() string { return latestSlackMessage(channelId, apiToken) }, "10s").
			Should(Equal("*Claimed:* \n*Unclaimed:* pool-1, pool-2"))

		By("Claiming pool-1")
		postSlackMessage(fmt.Sprintf("<@%s> claim pool-1", botId), channelId, apiToken)
		Eventually(func() string { return latestSlackMessage(channelId, apiToken) }, "10s").
			Should(Equal("Claimed pool-1"))
		updateGitRepo(gitDir)
		Expect(filepath.Join(gitDir, "pool-1", "claimed", "lock-a")).To(BeAnExistingFile())
		Expect(filepath.Join(gitDir, "pool-1", "unclaimed", "lock-a")).NotTo(BeAnExistingFile())

		By("Checking the status")
		postSlackMessage(fmt.Sprintf("<@%s> status", botId), channelId, apiToken)
		Eventually(func() string { return latestSlackMessage(channelId, apiToken) }, "10s").
			Should(Equal("*Claimed:* pool-1\n*Unclaimed:* pool-2"))

		By("Trying to claim pool-1 again")
		postSlackMessage(fmt.Sprintf("<@%s> claim pool-1", botId), channelId, apiToken)
		Eventually(func() string { return latestSlackMessage(channelId, apiToken) }, "10s").
			Should(Equal("pool-1 is not available"))

		By("Releasing pool-1")
		postSlackMessage(fmt.Sprintf("<@%s> release pool-1", botId), channelId, apiToken)
		Eventually(func() string { return latestSlackMessage(channelId, apiToken) }, "10s").
			Should(Equal("Released pool-1"))
		updateGitRepo(gitDir)
		Expect(filepath.Join(gitDir, "pool-1", "unclaimed", "lock-a")).To(BeAnExistingFile())
		Expect(filepath.Join(gitDir, "pool-1", "claimed", "lock-a")).NotTo(BeAnExistingFile())

		By("Checking the status")
		postSlackMessage(fmt.Sprintf("<@%s> status", botId), channelId, apiToken)
		Eventually(func() string { return latestSlackMessage(channelId, apiToken) }, "10s").
			Should(Equal("*Claimed:* \n*Unclaimed:* pool-1, pool-2"))

		By("Trying to release pool-1 again")
		postSlackMessage(fmt.Sprintf("<@%s> release pool-1", botId), channelId, apiToken)
		Eventually(func() string { return latestSlackMessage(channelId, apiToken) }, "10s").
			Should(Equal("pool-1 is not claimed"))

		By("Trying to claim non-existent pool")
		postSlackMessage(fmt.Sprintf("<@%s> claim non-existent-pool", botId), channelId, apiToken)
		Eventually(func() string { return latestSlackMessage(channelId, apiToken) }, "10s").
			Should(Equal("non-existent-pool is not available"))

		By("Trying to release non-existent-pool")
		postSlackMessage(fmt.Sprintf("<@%s> release non-existent-pool", botId), channelId, apiToken)
		Eventually(func() string { return latestSlackMessage(channelId, apiToken) }, "10s").
			Should(Equal("non-existent-pool is not claimed"))

		// @claimer claim pool-2
		// assert about error
		// assert about repo
		// @claimer status
		// assert about status
	})
})

func getEnv(name string) string {
	value, ok := os.LookupEnv(name)
	if !ok {
		Fail(fmt.Sprintf("%s must be set", name))
	}
	return value
}

func postSlackMessage(text, channelId, apiToken string) {
	resp, err := http.PostForm(
		"https://slack.com/api/chat.postMessage",
		url.Values{
			"token":    {apiToken},
			"channel":  {channelId},
			"text":     {text},
			"as_user":  {"false"},
			"username": {"claimer-integration-test"},
		},
	)
	Expect(err).NotTo(HaveOccurred())
	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)
	Expect(err).NotTo(HaveOccurred())

	var slackResponse struct {
		Ok    bool
		Error string
	}
	Expect(json.Unmarshal(body, &slackResponse)).To(Succeed())

	Expect(slackResponse.Ok).To(BeTrue(), fmt.Sprintf("Posting to slack failed: %s", slackResponse.Error))
}

func latestSlackMessage(channelId, apiToken string) string {
	resp, err := http.PostForm(
		"https://slack.com/api/channels.history",
		url.Values{
			"token":   {apiToken},
			"channel": {channelId},
			"count":   {"1"},
		},
	)
	Expect(err).NotTo(HaveOccurred())
	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)
	Expect(err).NotTo(HaveOccurred())

	var slackResponse struct {
		Ok       bool
		Messages []struct {
			Text string
		}
		Error string
	}
	Expect(json.Unmarshal(body, &slackResponse)).To(Succeed())

	Expect(slackResponse.Ok).To(BeTrue(), fmt.Sprintf("Getting message from slack failed: %s", slackResponse.Error))

	return slackResponse.Messages[0].Text
}

func updateGitRepo(gitDir string) {
	runGitCommand(gitDir, "fetch")
	runGitCommand(gitDir, "reset", "--hard", "origin/master")
}

func resetClaimerTestPool(gitDir string) {
	runGitCommand(gitDir, "checkout", "master")
	runGitCommand(gitDir, "reset", "--hard", "initial-state")
	runGitCommand(gitDir, "push", "--force", "origin", "master")

}

func runGitCommand(dir string, args ...string) {
	cmd := exec.Command("git", args...)
	cmd.Dir = dir
	ExpectWithOffset(1, cmd.Run()).To(Succeed())
}
<|endoftext|>


<commit_before>package goldpushk

// The contents of this file are goldpushk's source of truth, specifically the DeployableUnitSet
// returned by ProductionDeployableUnits().

const (
	// Gold instances.
	Angle             Instance = "angle"
	Chrome            Instance = "chrome"
	ChromePublic      Instance = "chrome-public"
	ChromiumOSTastDev Instance = "cros-tast-dev"
	Flutter           Instance = "flutter"
	FlutterEngine     Instance = "flutter-engine"
	Fuchsia           Instance = "fuchsia"
	FuchsiaPublic     Instance = "fuchsia-public"
	Lottie            Instance = "lottie"
	Pdfium            Instance = "pdfium"
	Skia              Instance = "skia"
	SkiaInfra         Instance = "skia-infra"
	SkiaPublic        Instance = "skia-public"

	// Gold services.
	BaselineServer  Service = "baselineserver"
	DiffCalculator  Service = "diffcalculator"
	DiffServer      Service = "diffserver"
	IngestionBT     Service = "ingestion-bt"
	Frontend        Service = "frontend"
	GitilesFollower Service = "gitilesfollower"

	// Testing Gold instances.
	TestInstance1     Instance = "goldpushk-test1"
	TestInstance2     Instance = "goldpushk-test2"
	TestCorpInstance1 Instance = "goldpushk-corp-test1"
	TestCorpInstance2 Instance = "goldpushk-corp-test2"

	// Testing Gold services.
	HealthyTestServer  Service = "healthy-server"
	CrashingTestServer Service = "crashing-server"
)

var (
	// knownPublicInstances is the set of Gold instances that are public.
	//
	// Note: consider rearchitecting this file in a manner that does not require any global state,
	// especially if we add more public instances in the future. For some potential ideas, see Kevin's
	// comments here: https://skia-review.googlesource.com/c/buildbot/+/243778.
	knownPublicInstances = []Instance{
		ChromePublic, SkiaPublic,
	}
)

// ProductionDeployableUnits returns the DeployableUnitSet that will be used as the source of truth
// across all of goldpushk.
func ProductionDeployableUnits() DeployableUnitSet {
	s := DeployableUnitSet{
		knownInstances: []Instance{
			Angle,
			Chrome,
			ChromePublic,
			ChromiumOSTastDev,
			Flutter,
			FlutterEngine,
			Fuchsia,
			FuchsiaPublic,
			Lottie,
			Pdfium,
			Skia,
			SkiaInfra,
			SkiaPublic,
		},
		knownServices: []Service{
			BaselineServer,
			DiffCalculator,
			DiffServer,
			Frontend,
			GitilesFollower,
			IngestionBT,
		},
	}

	// Add common services to all known instances.
	for _, instance := range s.knownInstances {
		if isPublicInstance(instance) {
			// There is only one service for public view instances: - frontend.
			s.add(instance, Frontend)
		} else {
			// Add common services for regular instances.
			s.add(instance, DiffCalculator)
			s.add(instance, DiffServer)
			s.add(instance, Frontend)
			// See skbug.com/11367
			if instance != ChromiumOSTastDev {
				s.add(instance, GitilesFollower)
			}
			s.add(instance, IngestionBT)
		}
	}

	// Add BaselineServer to the instances that require it.
	publicInstancesNeedingBaselineServer := []Instance{
		Angle, Chrome, ChromiumOSTastDev, Flutter, FlutterEngine, FuchsiaPublic, SkiaInfra,
	}
	for _, instance := range publicInstancesNeedingBaselineServer {
		s.add(instance, BaselineServer)
	}

	// Internal baseline options.
	s.addWithOptions(Fuchsia, BaselineServer, DeploymentOptions{
		internal: true,
	})

	// Overwrite common services for "fuchsia" instance, which need to run on skia-corp.
	s.addWithOptions(Fuchsia, DiffCalculator, DeploymentOptions{internal: true})
	s.addWithOptions(Fuchsia, DiffServer, DeploymentOptions{internal: true})
	s.addWithOptions(Fuchsia, Frontend, DeploymentOptions{internal: true})
	s.addWithOptions(Fuchsia, GitilesFollower, DeploymentOptions{internal: true})
	s.addWithOptions(Fuchsia, IngestionBT, DeploymentOptions{internal: true})
	return s
}

// isPublicInstance returns true if the given instance is in knownPublicInstances.
func isPublicInstance(instance Instance) bool {
	for _, i := range knownPublicInstances {
		if i == instance {
			return true
		}
	}
	return false
}

// TestingDeployableUnits returns a DeployableUnitSet comprised of test services that can be
// deployed without disrupting any production services for the purpose of testing goldpushk.
func TestingDeployableUnits() DeployableUnitSet {
	s := DeployableUnitSet{
		knownInstances: []Instance{
			TestInstance1,
			TestInstance2,
			TestCorpInstance1,
			TestCorpInstance2,
		},
		knownServices: []Service{
			HealthyTestServer,
			CrashingTestServer,
		},
	}

	addHealthyServerInstance := func(instance Instance, service Service, internal bool) {
		s.addWithOptions(instance, service, DeploymentOptions{
			internal: internal,
		})
	}

	addHealthyServerInstance(TestInstance1, HealthyTestServer, false)
	s.add(TestInstance1, CrashingTestServer)
	addHealthyServerInstance(TestInstance2, HealthyTestServer, false)
	s.add(TestInstance2, CrashingTestServer)
	addHealthyServerInstance(TestCorpInstance1, HealthyTestServer, true)
	s.addWithOptions(TestCorpInstance1, CrashingTestServer, DeploymentOptions{internal: true})
	addHealthyServerInstance(TestCorpInstance2, HealthyTestServer, true)
	s.addWithOptions(TestCorpInstance2, CrashingTestServer, DeploymentOptions{internal: true})

	return s
}
<commit_msg>[gold] Make fuchsia gitilesfollower a manual update<commit_after>package goldpushk

// The contents of this file are goldpushk's source of truth, specifically the DeployableUnitSet
// returned by ProductionDeployableUnits().

const (
	// Gold instances.
	Angle             Instance = "angle"
	Chrome            Instance = "chrome"
	ChromePublic      Instance = "chrome-public"
	ChromiumOSTastDev Instance = "cros-tast-dev"
	Flutter           Instance = "flutter"
	FlutterEngine     Instance = "flutter-engine"
	Fuchsia           Instance = "fuchsia"
	FuchsiaPublic     Instance = "fuchsia-public"
	Lottie            Instance = "lottie"
	Pdfium            Instance = "pdfium"
	Skia              Instance = "skia"
	SkiaInfra         Instance = "skia-infra"
	SkiaPublic        Instance = "skia-public"

	// Gold services.
	BaselineServer  Service = "baselineserver"
	DiffCalculator  Service = "diffcalculator"
	DiffServer      Service = "diffserver"
	IngestionBT     Service = "ingestion-bt"
	Frontend        Service = "frontend"
	GitilesFollower Service = "gitilesfollower"

	// Testing Gold instances.
	TestInstance1     Instance = "goldpushk-test1"
	TestInstance2     Instance = "goldpushk-test2"
	TestCorpInstance1 Instance = "goldpushk-corp-test1"
	TestCorpInstance2 Instance = "goldpushk-corp-test2"

	// Testing Gold services.
	HealthyTestServer  Service = "healthy-server"
	CrashingTestServer Service = "crashing-server"
)

var (
	// knownPublicInstances is the set of Gold instances that are public.
	//
	// Note: consider rearchitecting this file in a manner that does not require any global state,
	// especially if we add more public instances in the future. For some potential ideas, see Kevin's
	// comments here: https://skia-review.googlesource.com/c/buildbot/+/243778.
	knownPublicInstances = []Instance{
		ChromePublic, SkiaPublic,
	}
)

// ProductionDeployableUnits returns the DeployableUnitSet that will be used as the source of truth
// across all of goldpushk.
func ProductionDeployableUnits() DeployableUnitSet {
	s := DeployableUnitSet{
		knownInstances: []Instance{
			Angle,
			Chrome,
			ChromePublic,
			ChromiumOSTastDev,
			Flutter,
			FlutterEngine,
			Fuchsia,
			FuchsiaPublic,
			Lottie,
			Pdfium,
			Skia,
			SkiaInfra,
			SkiaPublic,
		},
		knownServices: []Service{
			BaselineServer,
			DiffCalculator,
			DiffServer,
			Frontend,
			GitilesFollower,
			IngestionBT,
		},
	}

	// Add common services to all known instances.
	for _, instance := range s.knownInstances {
		if isPublicInstance(instance) {
			// There is only one service for public view instances: - frontend.
			s.add(instance, Frontend)
		} else {
			// Add common services for regular instances.
			s.add(instance, DiffCalculator)
			s.add(instance, DiffServer)
			s.add(instance, Frontend)
			// See skbug.com/11367 for ChromiumOSTastDev.
			// See http://review.skia.org/376843 for Fuchsia
			if instance != ChromiumOSTastDev && instance != Fuchsia {
				s.add(instance, GitilesFollower)
			}
			s.add(instance, IngestionBT)
		}
	}

	// Add BaselineServer to the instances that require it.
	publicInstancesNeedingBaselineServer := []Instance{
		Angle, Chrome, ChromiumOSTastDev, Flutter, FlutterEngine, FuchsiaPublic, SkiaInfra,
	}
	for _, instance := range publicInstancesNeedingBaselineServer {
		s.add(instance, BaselineServer)
	}

	// Internal baseline options.
	s.addWithOptions(Fuchsia, BaselineServer, DeploymentOptions{
		internal: true,
	})

	// Overwrite common services for "fuchsia" instance, which need to run on skia-corp.
	s.addWithOptions(Fuchsia, DiffCalculator, DeploymentOptions{internal: true})
	s.addWithOptions(Fuchsia, DiffServer, DeploymentOptions{internal: true})
	s.addWithOptions(Fuchsia, Frontend, DeploymentOptions{internal: true})
	s.addWithOptions(Fuchsia, GitilesFollower, DeploymentOptions{internal: true})
	s.addWithOptions(Fuchsia, IngestionBT, DeploymentOptions{internal: true})
	return s
}

// isPublicInstance returns true if the given instance is in knownPublicInstances.
func isPublicInstance(instance Instance) bool {
	for _, i := range knownPublicInstances {
		if i == instance {
			return true
		}
	}
	return false
}

// TestingDeployableUnits returns a DeployableUnitSet comprised of test services that can be
// deployed without disrupting any production services for the purpose of testing goldpushk.
func TestingDeployableUnits() DeployableUnitSet {
	s := DeployableUnitSet{
		knownInstances: []Instance{
			TestInstance1,
			TestInstance2,
			TestCorpInstance1,
			TestCorpInstance2,
		},
		knownServices: []Service{
			HealthyTestServer,
			CrashingTestServer,
		},
	}

	addHealthyServerInstance := func(instance Instance, service Service, internal bool) {
		s.addWithOptions(instance, service, DeploymentOptions{
			internal: internal,
		})
	}

	addHealthyServerInstance(TestInstance1, HealthyTestServer, false)
	s.add(TestInstance1, CrashingTestServer)
	addHealthyServerInstance(TestInstance2, HealthyTestServer, false)
	s.add(TestInstance2, CrashingTestServer)
	addHealthyServerInstance(TestCorpInstance1, HealthyTestServer, true)
	s.addWithOptions(TestCorpInstance1, CrashingTestServer, DeploymentOptions{internal: true})
	addHealthyServerInstance(TestCorpInstance2, HealthyTestServer, true)
	s.addWithOptions(TestCorpInstance2, CrashingTestServer, DeploymentOptions{internal: true})

	return s
}
<|endoftext|>


<commit_before>// Package gongflow provides server support for ng-flow (https://github.com/flowjs/ng-flow).  If you want a way to instantly
// test it out, grab the gongflow-demo package: http://godoc.org/github.com/patdek/gongflow-demo
package gongflow

import (
	"errors"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"path"
	"strconv"
	"time"
)

var (
	DefaultDirPermissions     os.FileMode = 0777
	DefaultFilePermissions    os.FileMode = 0777
	ErrNoTempDir                          = errors.New("gongflow: the temporary directory doesn't exist")
	ErrCantCreateDir                      = errors.New("gongflow: can't create a directory under the temp directory")
	ErrCantWriteFile                      = errors.New("gongflow: can't write to a file under the temp directory")
	ErrCantReadFile                       = errors.New("gongflow: can't read a file under the temp directory (or got back bad data)")
	ErrCantDelete                         = errors.New("gongflow: can't delete a file/directory under the temp directory")
	alreadyCheckedDirectory               = false
	lastCheckedDirectoryError error       = nil
)

// NgFlowData is all the data listed in the "How do I set it up with my server?" section of the ng-flow
// README.md https://github.com/flowjs/flow.js/blob/master/README.md
type NgFlowData struct {
	flowChunkNumber  int    // The index of the chunk in the current upload. First chunk is 1 (no base-0 counting here).
	flowTotalChunks  int    // The total number of chunks.
	flowChunkSize    int    // The general chunk size. Using this value and flowTotalSize you can calculate the total number of chunks. The "final chunk" can be anything less than 2x chunk size.
	flowTotalSize    int    // The total file size.
	flowIdentifier   string // A unique identifier for the file contained in the request.
	flowFilename     string // The original file name (since a bug in Firefox results in the file name not being transmitted in chunk multipart posts).
	flowRelativePath string // The file's relative path when selecting a directory (defaults to file name in all browsers except Chrome)
}

// PartFlowData does exactly what it says on the tin, it extracts all the flow data from a request object and puts
// it into a nice little struct for you
func PartFlowData(r *http.Request) (NgFlowData, error) {
	var err error
	ngfd := NgFlowData{}
	ngfd.flowChunkNumber, err = strconv.Atoi(r.FormValue("flowChunkNumber"))
	if err != nil {
		return ngfd, errors.New("Bad flowChunkNumber")
	}
	ngfd.flowTotalChunks, err = strconv.Atoi(r.FormValue("flowTotalChunks"))
	if err != nil {
		return ngfd, errors.New("Bad flowTotalChunks")
	}
	ngfd.flowChunkSize, err = strconv.Atoi(r.FormValue("flowChunkSize"))
	if err != nil {
		return ngfd, errors.New("Bad flowChunkSize")
	}
	ngfd.flowTotalSize, err = strconv.Atoi(r.FormValue("flowTotalSize"))
	if err != nil {
		return ngfd, errors.New("Bad flowTotalSize")
	}
	ngfd.flowIdentifier = r.FormValue("flowIdentifier")
	if ngfd.flowIdentifier == "" {
		return ngfd, errors.New("Bad flowIdentifier")
	}
	ngfd.flowFilename = r.FormValue("flowFilename")
	if ngfd.flowFilename == "" {
		return ngfd, errors.New("Bad flowFilename")
	}
	ngfd.flowRelativePath = r.FormValue("flowRelativePath")
	if ngfd.flowRelativePath == "" {
		return ngfd, errors.New("Bad flowRelativePath")
	}
	return ngfd, nil
}

// PartUpload is used to handle a POST from ng-flow, it will return an empty string for part upload (incomplete) and when
// all the parts have been uploaded, it will return the path to the reconstituted file.  So, you can just keep calling it
// until you get back the path to a file.
func PartUpload(tempDir string, ngfd NgFlowData, r *http.Request) (string, error) {
	err := checkDirectory(tempDir)
	if err != nil {
		return "", err
	}
	fileDir, chunkFile := buildPathParts(tempDir, ngfd)
	err = storePart(fileDir, chunkFile, ngfd, r)
	if err != nil {
		return "", errors.New("Unable to store part" + err.Error())
	}
	if allPartsUploaded(tempDir, ngfd) {
		file, err := combineParts(fileDir, ngfd)
		if err != nil {
			return "", err
		}
		return file, nil
	}
	return "", nil
}

// PartStatus is used to handle a GET from ng-flow, it will return a (message, 200) for when it already has a part, and it
// will return a (message, 404 | 500) when a part is incomplete or not started.
func PartStatus(tempDir string, ngfd NgFlowData) (string, int) {
	err := checkDirectory(tempDir)
	if err != nil {
		return "Directory is broken: " + err.Error(), 500
	}
	_, chunkFile := buildPathParts(tempDir, ngfd)
	flowChunkNumberString := strconv.Itoa(ngfd.flowChunkNumber)
	dat, err := ioutil.ReadFile(chunkFile)
	if err != nil {
		return "The part " + ngfd.flowIdentifier + ":" + flowChunkNumberString + " isn't started yet!", 404
	}
	// An exception for large last chunks, according to ng-flow the last chunk can be anywhere less
	// than 2x the chunk size unless you haave forceChunkSize on... seems like idiocy to me, but alright.
	if ngfd.flowChunkNumber != ngfd.flowTotalChunks && ngfd.flowChunkSize != len(dat) {
		return "The part " + ngfd.flowIdentifier + ":" + flowChunkNumberString + " is the wrong size!", 500
	}

	return "The part " + ngfd.flowIdentifier + ":" + flowChunkNumberString + " looks great!", 200
}

// PartsCleanup is used to go through the tempDir and remove any parts and directories older than
// than the timeoutDur, best to set this VERY conservatively.
func PartsCleanup(tempDir string, timeoutDur time.Duration) error {
	files, err := ioutil.ReadDir(tempDir)
	if err != nil {
		return err
	}
	for _, f := range files {
		fl := path.Join(tempDir, f.Name())
		finfo, err := os.Stat(fl)
		if err != nil {
			return err
		}

		log.Println(f.Name())
		log.Println(time.Now().Sub(finfo.ModTime()))
		if time.Now().Sub(finfo.ModTime()) > timeoutDur {
			err = os.RemoveAll(fl)
			if err != nil {
				return err
			}
		}
	}
	return nil
}

// buildPathParts simply builds the paths to the ID of the upload, and to the specific Chunk
func buildPathParts(tempDir string, ngfd NgFlowData) (string, string) {
	filePath := path.Join(tempDir, ngfd.flowIdentifier)
	chunkFile := path.Join(filePath, strconv.Itoa(ngfd.flowChunkNumber))
	return filePath, chunkFile
}

// combineParts will take the chunks uploaded, and combined them into a single file with the
// name as uploaded from the NgFlowData, and it will clean up the parts as it goes.
func combineParts(fileDir string, ngfd NgFlowData) (string, error) {
	combinedName := path.Join(fileDir, ngfd.flowFilename)
	cn, err := os.Create(combinedName)
	if err != nil {
		return "", err
	}
	defer cn.Close()

	files, err := ioutil.ReadDir(fileDir)
	if err != nil {
		return "", err
	}
	for _, f := range files {
		fl := path.Join(fileDir, f.Name())
		dat, err := ioutil.ReadFile(fl)
		if err != nil {
			return "", err
		}
		_, err = cn.Write(dat)
		if err != nil {
			return "", err
		}
		if fl != combinedName { // we don't want to delete the file we just created
			err = os.Remove(fl)
			if err != nil {
				return "", err
			}
		}
	}
	return combinedName, nil
}

// allPartsUploaded checks if the file is completely uploaded (based on total size)
func allPartsUploaded(tempDir string, ngfd NgFlowData) bool {
	partsPath := path.Join(tempDir, ngfd.flowIdentifier)
	files, err := ioutil.ReadDir(partsPath)
	if err != nil {
		log.Println(err)
	}
	totalSize := int64(0)
	for _, f := range files {
		fi, err := os.Stat(path.Join(partsPath, f.Name()))
		if err != nil {
			log.Println(err)
		}
		totalSize += fi.Size()
	}
	if totalSize == int64(ngfd.flowTotalSize) {
		return true
	}
	return false
}

// storePart puts the part in the request into the right place on disk
func storePart(tempDir string, tempFile string, ngfd NgFlowData, r *http.Request) error {
	err := os.MkdirAll(tempDir, DefaultDirPermissions)
	if err != nil {
		return errors.New("Bad directory")
	}
	file, _, err := r.FormFile("file")
	if err != nil {
		return errors.New("Can't access file field")
	}
	data, err := ioutil.ReadAll(file)
	if err != nil {
		return errors.New("Can't read file")
	}
	err = ioutil.WriteFile(tempFile, data, DefaultDirPermissions)
	if err != nil {
		return errors.New("Can't write file")
	}
	return nil
}

// checkDirectory makes sure that we have all the needed permissions to the temp directory to
// read/write/delete.  Expensive operation, so it only does it once.
func checkDirectory(d string) error {
	if alreadyCheckedDirectory {
		return lastCheckedDirectoryError
	}

	alreadyCheckedDirectory = true

	if !directoryExists(d) {
		lastCheckedDirectoryError = ErrNoTempDir
		return lastCheckedDirectoryError
	}

	testName := "5d58061677944334bb616ba19cec5cc4"
	testPart := "42"
	contentName := "foobie"
	testContent := `For instance, on the planet Earth, man had always assumed that he was more intelligent than
	dolphins because he had achieved so much—the wheel, New York, wars and so on—whilst all the dolphins had
	ever done was muck about in the water having a good time. But conversely, the dolphins had always believed
	that they were far more intelligent than man—for precisely the same reasons.`

	p := path.Join(d, testName, testPart)
	err := os.MkdirAll(p, DefaultDirPermissions)
	if err != nil {
		lastCheckedDirectoryError = ErrCantCreateDir
		return lastCheckedDirectoryError
	}

	f := path.Join(p, contentName)
	err = ioutil.WriteFile(f, []byte(testContent), DefaultFilePermissions)
	if err != nil {
		lastCheckedDirectoryError = ErrCantWriteFile
		return lastCheckedDirectoryError
	}

	b, err := ioutil.ReadFile(f)
	if err != nil {
		lastCheckedDirectoryError = ErrCantReadFile
		return lastCheckedDirectoryError
	}
	if string(b) != testContent {
		lastCheckedDirectoryError = ErrCantReadFile // TODO: This should probably be a different error
		return lastCheckedDirectoryError
	}

	err = os.RemoveAll(path.Join(d, testName))
	if err != nil {
		lastCheckedDirectoryError = ErrCantDelete
		return lastCheckedDirectoryError
	}

	if os.TempDir() == d {
		log.Println("You should really have a directory just for upload temp (different from system temp).  It is OK, but consider making a subdirectory for it.")
	}

	return nil
}

// directoryExists checks if the directory exists of course!
func directoryExists(d string) bool {
	finfo, err := os.Stat(d)
	if err == nil && finfo.IsDir() {
		return true
	}
	return false
}
<commit_msg>Named parts chunks to match ng-flow<commit_after>// Package gongflow provides server support for ng-flow (https://github.com/flowjs/ng-flow).  If you want a way to instantly
// test it out, grab the gongflow-demo package: http://godoc.org/github.com/patdek/gongflow-demo
package gongflow

import (
	"errors"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"path"
	"strconv"
	"time"
)

var (
	DefaultDirPermissions     os.FileMode = 0777
	DefaultFilePermissions    os.FileMode = 0777
	ErrNoTempDir                          = errors.New("gongflow: the temporary directory doesn't exist")
	ErrCantCreateDir                      = errors.New("gongflow: can't create a directory under the temp directory")
	ErrCantWriteFile                      = errors.New("gongflow: can't write to a file under the temp directory")
	ErrCantReadFile                       = errors.New("gongflow: can't read a file under the temp directory (or got back bad data)")
	ErrCantDelete                         = errors.New("gongflow: can't delete a file/directory under the temp directory")
	alreadyCheckedDirectory               = false
	lastCheckedDirectoryError error       = nil
)

// NgFlowData is all the data listed in the "How do I set it up with my server?" section of the ng-flow
// README.md https://github.com/flowjs/flow.js/blob/master/README.md
type NgFlowData struct {
	flowChunkNumber  int    // The index of the chunk in the current upload. First chunk is 1 (no base-0 counting here).
	flowTotalChunks  int    // The total number of chunks.
	flowChunkSize    int    // The general chunk size. Using this value and flowTotalSize you can calculate the total number of chunks. The "final chunk" can be anything less than 2x chunk size.
	flowTotalSize    int    // The total file size.
	flowIdentifier   string // A unique identifier for the file contained in the request.
	flowFilename     string // The original file name (since a bug in Firefox results in the file name not being transmitted in chunk multichunk posts).
	flowRelativePath string // The file's relative path when selecting a directory (defaults to file name in all browsers except Chrome)
}

// ChunkFlowData does exactly what it says on the tin, it extracts all the flow data from a request object and puts
// it into a nice little struct for you
func ChunkFlowData(r *http.Request) (NgFlowData, error) {
	var err error
	ngfd := NgFlowData{}
	ngfd.flowChunkNumber, err = strconv.Atoi(r.FormValue("flowChunkNumber"))
	if err != nil {
		return ngfd, errors.New("Bad flowChunkNumber")
	}
	ngfd.flowTotalChunks, err = strconv.Atoi(r.FormValue("flowTotalChunks"))
	if err != nil {
		return ngfd, errors.New("Bad flowTotalChunks")
	}
	ngfd.flowChunkSize, err = strconv.Atoi(r.FormValue("flowChunkSize"))
	if err != nil {
		return ngfd, errors.New("Bad flowChunkSize")
	}
	ngfd.flowTotalSize, err = strconv.Atoi(r.FormValue("flowTotalSize"))
	if err != nil {
		return ngfd, errors.New("Bad flowTotalSize")
	}
	ngfd.flowIdentifier = r.FormValue("flowIdentifier")
	if ngfd.flowIdentifier == "" {
		return ngfd, errors.New("Bad flowIdentifier")
	}
	ngfd.flowFilename = r.FormValue("flowFilename")
	if ngfd.flowFilename == "" {
		return ngfd, errors.New("Bad flowFilename")
	}
	ngfd.flowRelativePath = r.FormValue("flowRelativePath")
	if ngfd.flowRelativePath == "" {
		return ngfd, errors.New("Bad flowRelativePath")
	}
	return ngfd, nil
}

// ChunkUpload is used to handle a POST from ng-flow, it will return an empty string for chunk upload (incomplete) and when
// all the chunks have been uploaded, it will return the path to the reconstituted file.  So, you can just keep calling it
// until you get back the path to a file.
func ChunkUpload(tempDir string, ngfd NgFlowData, r *http.Request) (string, error) {
	err := checkDirectory(tempDir)
	if err != nil {
		return "", err
	}
	fileDir, chunkFile := buildPathChunks(tempDir, ngfd)
	err = storeChunk(fileDir, chunkFile, ngfd, r)
	if err != nil {
		return "", errors.New("Unable to store chunk" + err.Error())
	}
	if allChunksUploaded(tempDir, ngfd) {
		file, err := combineChunks(fileDir, ngfd)
		if err != nil {
			return "", err
		}
		return file, nil
	}
	return "", nil
}

// ChunkStatus is used to handle a GET from ng-flow, it will return a (message, 200) for when it already has a chunk, and it
// will return a (message, 404 | 500) when a chunk is incomplete or not started.
func ChunkStatus(tempDir string, ngfd NgFlowData) (string, int) {
	err := checkDirectory(tempDir)
	if err != nil {
		return "Directory is broken: " + err.Error(), 500
	}
	_, chunkFile := buildPathChunks(tempDir, ngfd)
	flowChunkNumberString := strconv.Itoa(ngfd.flowChunkNumber)
	dat, err := ioutil.ReadFile(chunkFile)
	if err != nil {
		return "The chunk " + ngfd.flowIdentifier + ":" + flowChunkNumberString + " isn't started yet!", 404
	}
	// An exception for large last chunks, according to ng-flow the last chunk can be anywhere less
	// than 2x the chunk size unless you haave forceChunkSize on... seems like idiocy to me, but alright.
	if ngfd.flowChunkNumber != ngfd.flowTotalChunks && ngfd.flowChunkSize != len(dat) {
		return "The chunk " + ngfd.flowIdentifier + ":" + flowChunkNumberString + " is the wrong size!", 500
	}

	return "The chunk " + ngfd.flowIdentifier + ":" + flowChunkNumberString + " looks great!", 200
}

// ChunksCleanup is used to go through the tempDir and remove any chunks and directories older than
// than the timeoutDur, best to set this VERY conservatively.
func ChunksCleanup(tempDir string, timeoutDur time.Duration) error {
	files, err := ioutil.ReadDir(tempDir)
	if err != nil {
		return err
	}
	for _, f := range files {
		fl := path.Join(tempDir, f.Name())
		finfo, err := os.Stat(fl)
		if err != nil {
			return err
		}

		log.Println(f.Name())
		log.Println(time.Now().Sub(finfo.ModTime()))
		if time.Now().Sub(finfo.ModTime()) > timeoutDur {
			err = os.RemoveAll(fl)
			if err != nil {
				return err
			}
		}
	}
	return nil
}

// buildPathChunks simply builds the paths to the ID of the upload, and to the specific Chunk
func buildPathChunks(tempDir string, ngfd NgFlowData) (string, string) {
	filePath := path.Join(tempDir, ngfd.flowIdentifier)
	chunkFile := path.Join(filePath, strconv.Itoa(ngfd.flowChunkNumber))
	return filePath, chunkFile
}

// combineChunks will take the chunks uploaded, and combined them into a single file with the
// name as uploaded from the NgFlowData, and it will clean up the chunks as it goes.
func combineChunks(fileDir string, ngfd NgFlowData) (string, error) {
	combinedName := path.Join(fileDir, ngfd.flowFilename)
	cn, err := os.Create(combinedName)
	if err != nil {
		return "", err
	}
	defer cn.Close()

	files, err := ioutil.ReadDir(fileDir)
	if err != nil {
		return "", err
	}
	for _, f := range files {
		fl := path.Join(fileDir, f.Name())
		dat, err := ioutil.ReadFile(fl)
		if err != nil {
			return "", err
		}
		_, err = cn.Write(dat)
		if err != nil {
			return "", err
		}
		if fl != combinedName { // we don't want to delete the file we just created
			err = os.Remove(fl)
			if err != nil {
				return "", err
			}
		}
	}
	return combinedName, nil
}

// allChunksUploaded checks if the file is completely uploaded (based on total size)
func allChunksUploaded(tempDir string, ngfd NgFlowData) bool {
	chunksPath := path.Join(tempDir, ngfd.flowIdentifier)
	files, err := ioutil.ReadDir(chunksPath)
	if err != nil {
		log.Println(err)
	}
	totalSize := int64(0)
	for _, f := range files {
		fi, err := os.Stat(path.Join(chunksPath, f.Name()))
		if err != nil {
			log.Println(err)
		}
		totalSize += fi.Size()
	}
	if totalSize == int64(ngfd.flowTotalSize) {
		return true
	}
	return false
}

// storeChunk puts the chunk in the request into the right place on disk
func storeChunk(tempDir string, tempFile string, ngfd NgFlowData, r *http.Request) error {
	err := os.MkdirAll(tempDir, DefaultDirPermissions)
	if err != nil {
		return errors.New("Bad directory")
	}
	file, _, err := r.FormFile("file")
	if err != nil {
		return errors.New("Can't access file field")
	}
	data, err := ioutil.ReadAll(file)
	if err != nil {
		return errors.New("Can't read file")
	}
	err = ioutil.WriteFile(tempFile, data, DefaultDirPermissions)
	if err != nil {
		return errors.New("Can't write file")
	}
	return nil
}

// checkDirectory makes sure that we have all the needed permissions to the temp directory to
// read/write/delete.  Expensive operation, so it only does it once.
func checkDirectory(d string) error {
	if alreadyCheckedDirectory {
		return lastCheckedDirectoryError
	}

	alreadyCheckedDirectory = true

	if !directoryExists(d) {
		lastCheckedDirectoryError = ErrNoTempDir
		return lastCheckedDirectoryError
	}

	testName := "5d58061677944334bb616ba19cec5cc4"
	testChunk := "42"
	contentName := "foobie"
	testContent := `For instance, on the planet Earth, man had always assumed that he was more intelligent than
	dolphins because he had achieved so much—the wheel, New York, wars and so on—whilst all the dolphins had
	ever done was muck about in the water having a good time. But conversely, the dolphins had always believed
	that they were far more intelligent than man—for precisely the same reasons.`

	p := path.Join(d, testName, testChunk)
	err := os.MkdirAll(p, DefaultDirPermissions)
	if err != nil {
		lastCheckedDirectoryError = ErrCantCreateDir
		return lastCheckedDirectoryError
	}

	f := path.Join(p, contentName)
	err = ioutil.WriteFile(f, []byte(testContent), DefaultFilePermissions)
	if err != nil {
		lastCheckedDirectoryError = ErrCantWriteFile
		return lastCheckedDirectoryError
	}

	b, err := ioutil.ReadFile(f)
	if err != nil {
		lastCheckedDirectoryError = ErrCantReadFile
		return lastCheckedDirectoryError
	}
	if string(b) != testContent {
		lastCheckedDirectoryError = ErrCantReadFile // TODO: This should probably be a different error
		return lastCheckedDirectoryError
	}

	err = os.RemoveAll(path.Join(d, testName))
	if err != nil {
		lastCheckedDirectoryError = ErrCantDelete
		return lastCheckedDirectoryError
	}

	if os.TempDir() == d {
		log.Println("You should really have a directory just for upload temp (different from system temp).  It is OK, but consider making a subdirectory for it.")
	}

	return nil
}

// directoryExists checks if the directory exists of course!
func directoryExists(d string) bool {
	finfo, err := os.Stat(d)
	if err == nil && finfo.IsDir() {
		return true
	}
	return false
}
<|endoftext|>


<commit_before>package thingsif

import (
	"encoding/json"
	"fmt"
	"time"
)

type DbMessage struct {
	Time  string  `json:"time"`
	DevId string  `json:"device_id"`
	Raw   string  `json:"raw"`
	Bat   float64 `json:"bat"`
	Humd  float64 `json:"humd"`
	Temp  float64 `json:"temp"`
	Rain  float64 `json:"rain,omitempty"`
	Pres  float64 `json:"pres,omitempty"`
	Valid bool    `json:"valid"`
}

type DbList struct {
	List []*DbMessage
}

func (mq *mqttCli) GetLast7days() ([]DbMessage, error) {
	fmt.Printf("Start get %v\n", time.Now())
	url := "https://" + mq.conf.Username + ".data.thethingsnetwork.org/api/v2/query?last=7d"
	messages := make([]DbMessage, 0)
	body, err := getHttpBody(url, mq.conf.Password)
	if err != nil {
		return messages, fmt.Errorf("Failed to get history: %v", err)
	}
	err = json.Unmarshal(body, &messages)
	if err != nil {
		return messages, fmt.Errorf("Failed to get history: %v", err)
	}
	fmt.Printf("done get\n", time.Now())
	return messages, nil
}
<commit_msg>Fixed broken print<commit_after>package thingsif

import (
	"encoding/json"
	"fmt"
	"time"
)

type DbMessage struct {
	Time  string  `json:"time"`
	DevId string  `json:"device_id"`
	Raw   string  `json:"raw"`
	Bat   float64 `json:"bat"`
	Humd  float64 `json:"humd"`
	Temp  float64 `json:"temp"`
	Rain  float64 `json:"rain,omitempty"`
	Pres  float64 `json:"pres,omitempty"`
	Valid bool    `json:"valid"`
}

type DbList struct {
	List []*DbMessage
}

func (mq *mqttCli) GetLast7days() ([]DbMessage, error) {
	fmt.Printf("Start get %v\n", time.Now())
	url := "https://" + mq.conf.Username + ".data.thethingsnetwork.org/api/v2/query?last=7d"
	messages := make([]DbMessage, 0)
	body, err := getHttpBody(url, mq.conf.Password)
	if err != nil {
		return messages, fmt.Errorf("Failed to get history: %v", err)
	}
	err = json.Unmarshal(body, &messages)
	if err != nil {
		return messages, fmt.Errorf("Failed to get history: %v", err)
	}
	fmt.Printf("done get %v\n", time.Now())
	return messages, nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"strings"
	"testing"

	"github.com/hashicorp/consul-template/config"
	"github.com/hashicorp/consul-template/dependency"
)

func TestRunner_appendSecrets(t *testing.T) {
	t.Parallel()

	secretValue1 := "somevalue"
	secretValue2 := "somevalue2"

	cases := []struct {
		name     string
		path     string
		noPrefix *bool
		data     *dependency.Secret
		notFound bool
	}{
		{
			name:     "kv1_secret",
			path:     "kv/bar",
			noPrefix: config.Bool(false),
			data: &dependency.Secret{
				Data: map[string]interface{}{
					"key_field":  secretValue1,
					"key_field2": secretValue2,
				},
			},
			notFound: false,
		},
		{
			name:     "kv2_secret",
			path:     "secret/data/foo",
			noPrefix: config.Bool(false),
			data: &dependency.Secret{
				Data: map[string]interface{}{
					"metadata": map[string]interface{}{
						"destroyed": bool(false),
						"version":   "1",
					},
					"data": map[string]interface{}{
						"key_field":  secretValue1,
						"key_field2": secretValue2,
					},
				},
			},
			notFound: false,
		},
		{
			name:     "kv2_secret_destroyed",
			path:     "secret/data/foo",
			noPrefix: config.Bool(false),
			data: &dependency.Secret{
				Data: map[string]interface{}{
					"metadata": map[string]interface{}{
						"destroyed": bool(true),
						"version":   "2",
					},
					"data": nil,
				},
			},
			notFound: true,
		},
		{
			name:     "kv2 secret noprefix excludes path",
			path:     "secret/data/foo",
			noPrefix: config.Bool(true),
			data: &dependency.Secret{
				Data: map[string]interface{}{
					"metadata": map[string]interface{}{
						"destroyed": bool(false),
						"version":   "1",
					},
					"data": map[string]interface{}{
						"key_field":  secretValue1,
						"key_field2": secretValue2,
					},
				},
			},
			notFound: false,
		},
		{
			name:     "kv2 secret false noprefix includes path",
			path:     "secret/data/foo",
			noPrefix: config.Bool(false),
			data: &dependency.Secret{
				Data: map[string]interface{}{
					"metadata": map[string]interface{}{
						"destroyed": bool(false),
						"version":   "1",
					},
					"data": map[string]interface{}{
						"key_field":  secretValue1,
						"key_field2": secretValue2,
					},
				},
			},
			notFound: false,
		},
		{
			name:     "kv2 secret default noprefix includes path",
			path:     "secret/data/foo",
			noPrefix: nil,
			data: &dependency.Secret{
				Data: map[string]interface{}{
					"metadata": map[string]interface{}{
						"destroyed": bool(false),
						"version":   "1",
					},
					"data": map[string]interface{}{
						"key_field":  secretValue1,
						"key_field2": secretValue2,
					},
				},
			},
			notFound: false,
		},
		{
			name:     "int secret skipped",
			path:     "kv/foo",
			noPrefix: config.Bool(false),
			data: &dependency.Secret{
				Data: map[string]interface{}{
					"key_field":  1,
					"key_field2": 1,
				},
			},
			notFound: true,
		},
	}

	for _, tc := range cases {
		t.Run(fmt.Sprintf("%s", tc.name), func(t *testing.T) {
			cfg := Config{
				Secrets: &PrefixConfigs{
					&PrefixConfig{
						Path:     config.String(tc.path),
						NoPrefix: tc.noPrefix,
					},
				},
			}
			c := DefaultConfig().Merge(&cfg)
			r, err := NewRunner(c, true)
			if err != nil {
				t.Fatal(err)
			}
			vrq, err := dependency.NewVaultReadQuery(tc.path)
			if err != nil {
				t.Fatal(err)
			}
			env := make(map[string]string)
			appendError := r.appendSecrets(env, vrq, tc.data)
			if appendError != nil {
				t.Fatalf("got err: %s", appendError)
			}

			if len(env) > 2 {
				t.Fatalf("Expected only 2 values in this test")
			}

			keyName1 := "key_field"
			if tc.noPrefix == nil || !*tc.noPrefix {
				keyName1 = tc.path + "_" + keyName1
			}
			keyName1 = strings.Replace(keyName1, "/", "_", -1)

			var value string
			value, ok := env[keyName1]
			if !ok && !tc.notFound {
				t.Fatalf("expected (%s) key, but was not found", keyName1)
			}
			if ok && tc.notFound {
				t.Fatalf("expected to not find key, but (%s) was found", keyName1)
			}
			if ok && value != secretValue1 {
				t.Fatalf("values didn't match, expected (%s), got (%s)", secretValue1, value)
			}

			keyName2 := "key_field2"
			if tc.noPrefix == nil || !*tc.noPrefix {
				keyName2 = tc.path + "_" + keyName2
			}
			keyName2 = strings.Replace(keyName2, "/", "_", -1)

			value, ok = env[keyName2]
			if !ok && !tc.notFound {
				t.Fatalf("expected (%s) key, but was not found", keyName2)
			}
			if ok && tc.notFound {
				t.Fatalf("expected to not find key, but (%s) was found", keyName2)
			}
			if ok && value != secretValue2 {
				t.Fatalf("values didn't match, expected (%s), got (%s)", secretValue2, value)
			}
		})
	}
}
<commit_msg>Add former tests as well<commit_after>package main

import (
	"fmt"
	"reflect"
	"strings"
	"testing"

	"github.com/hashicorp/consul-template/config"
	"github.com/hashicorp/consul-template/dependency"
)

func TestRunner_appendSecrets(t *testing.T) {
	t.Parallel()

	secretValue1 := "somevalue"
	secretValue2 := "somevalue2"

	cases := []struct {
		name     string
		path     string
		noPrefix *bool
		data     *dependency.Secret
		notFound bool
	}{
		{
			name:     "kv1_secret",
			path:     "kv/bar",
			noPrefix: config.Bool(false),
			data: &dependency.Secret{
				Data: map[string]interface{}{
					"key_field":  secretValue1,
					"key_field2": secretValue2,
				},
			},
			notFound: false,
		},
		{
			name:     "kv2_secret",
			path:     "secret/data/foo",
			noPrefix: config.Bool(false),
			data: &dependency.Secret{
				Data: map[string]interface{}{
					"metadata": map[string]interface{}{
						"destroyed": bool(false),
						"version":   "1",
					},
					"data": map[string]interface{}{
						"key_field":  secretValue1,
						"key_field2": secretValue2,
					},
				},
			},
			notFound: false,
		},
		{
			name:     "kv2_secret_destroyed",
			path:     "secret/data/foo",
			noPrefix: config.Bool(false),
			data: &dependency.Secret{
				Data: map[string]interface{}{
					"metadata": map[string]interface{}{
						"destroyed": bool(true),
						"version":   "2",
					},
					"data": nil,
				},
			},
			notFound: true,
		},
		{
			name:     "kv2 secret noprefix excludes path",
			path:     "secret/data/foo",
			noPrefix: config.Bool(true),
			data: &dependency.Secret{
				Data: map[string]interface{}{
					"metadata": map[string]interface{}{
						"destroyed": bool(false),
						"version":   "1",
					},
					"data": map[string]interface{}{
						"key_field":  secretValue1,
						"key_field2": secretValue2,
					},
				},
			},
			notFound: false,
		},
		{
			name:     "kv2 secret false noprefix includes path",
			path:     "secret/data/foo",
			noPrefix: config.Bool(false),
			data: &dependency.Secret{
				Data: map[string]interface{}{
					"metadata": map[string]interface{}{
						"destroyed": bool(false),
						"version":   "1",
					},
					"data": map[string]interface{}{
						"key_field":  secretValue1,
						"key_field2": secretValue2,
					},
				},
			},
			notFound: false,
		},
		{
			name:     "kv2 secret default noprefix includes path",
			path:     "secret/data/foo",
			noPrefix: nil,
			data: &dependency.Secret{
				Data: map[string]interface{}{
					"metadata": map[string]interface{}{
						"destroyed": bool(false),
						"version":   "1",
					},
					"data": map[string]interface{}{
						"key_field":  secretValue1,
						"key_field2": secretValue2,
					},
				},
			},
			notFound: false,
		},
		{
			name:     "int secret skipped",
			path:     "kv/foo",
			noPrefix: config.Bool(false),
			data: &dependency.Secret{
				Data: map[string]interface{}{
					"key_field":  1,
					"key_field2": 1,
				},
			},
			notFound: true,
		},
	}

	for _, tc := range cases {
		t.Run(fmt.Sprintf("%s", tc.name), func(t *testing.T) {
			cfg := Config{
				Secrets: &PrefixConfigs{
					&PrefixConfig{
						Path:     config.String(tc.path),
						NoPrefix: tc.noPrefix,
					},
				},
			}
			c := DefaultConfig().Merge(&cfg)
			r, err := NewRunner(c, true)
			if err != nil {
				t.Fatal(err)
			}
			vrq, err := dependency.NewVaultReadQuery(tc.path)
			if err != nil {
				t.Fatal(err)
			}
			env := make(map[string]string)
			appendError := r.appendSecrets(env, vrq, tc.data)
			if appendError != nil {
				t.Fatalf("got err: %s", appendError)
			}

			if len(env) > 2 {
				t.Fatalf("Expected only 2 values in this test")
			}

			keyName1 := "key_field"
			if tc.noPrefix == nil || !*tc.noPrefix {
				keyName1 = tc.path + "_" + keyName1
			}
			keyName1 = strings.Replace(keyName1, "/", "_", -1)

			var value string
			value, ok := env[keyName1]
			if !ok && !tc.notFound {
				t.Fatalf("expected (%s) key, but was not found", keyName1)
			}
			if ok && tc.notFound {
				t.Fatalf("expected to not find key, but (%s) was found", keyName1)
			}
			if ok && value != secretValue1 {
				t.Fatalf("values didn't match, expected (%s), got (%s)", secretValue1, value)
			}

			keyName2 := "key_field2"
			if tc.noPrefix == nil || !*tc.noPrefix {
				keyName2 = tc.path + "_" + keyName2
			}
			keyName2 = strings.Replace(keyName2, "/", "_", -1)

			value, ok = env[keyName2]
			if !ok && !tc.notFound {
				t.Fatalf("expected (%s) key, but was not found", keyName2)
			}
			if ok && tc.notFound {
				t.Fatalf("expected to not find key, but (%s) was found", keyName2)
			}
			if ok && value != secretValue2 {
				t.Fatalf("values didn't match, expected (%s), got (%s)", secretValue2, value)
			}
		})
	}
}

func TestRunner_appendPrefixes(t *testing.T) {
	t.Parallel()

	cases := []struct {
		name     string
		path     string
		noPrefix *bool
		data     []*dependency.KeyPair
		keyName  string
	}{
		{
			name:     "false noprefix appends path",
			path:     "app/my_service",
			noPrefix: config.Bool(false),
			data: []*dependency.KeyPair{
				&dependency.KeyPair{
					Key:   "mykey",
					Value: "myValue",
				},
			},
			keyName: "app_my_service_mykey",
		},
		{
			name:     "true noprefix excludes path",
			path:     "app/my_service",
			noPrefix: config.Bool(true),
			data: []*dependency.KeyPair{
				&dependency.KeyPair{
					Key:   "mykey",
					Value: "myValue",
				},
			},
			keyName: "mykey",
		},
		{
			name:     "null noprefix excludes path",
			path:     "app/my_service",
			noPrefix: nil,
			data: []*dependency.KeyPair{
				&dependency.KeyPair{
					Key:   "mykey",
					Value: "myValue",
				},
			},
			keyName: "mykey",
		},
	}

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			cfg := Config{
				Prefixes: &PrefixConfigs{
					&PrefixConfig{
						Path:     config.String(tc.path),
						NoPrefix: tc.noPrefix,
					},
				},
			}
			c := DefaultConfig().Merge(&cfg)
			r, err := NewRunner(c, true)
			if err != nil {
				t.Fatal(err)
			}
			kvq, err := dependency.NewKVListQuery(tc.path)
			if err != nil {
				t.Fatal(err)
			}
			env := make(map[string]string)
			appendError := r.appendPrefixes(env, kvq, tc.data)
			if appendError != nil {
				t.Fatalf("got err: %s", appendError)
			}

			if len(env) > 1 {
				t.Fatalf("Expected only 1 value in this test")
			}

			var value string
			value, ok := env[tc.keyName]
			if !ok {
				t.Fatalf("expected (%s) key, but was not found", tc.keyName)
			}
			if ok && value != tc.data[0].Value {
				t.Fatalf("values didn't match, expected (%s), got (%s)", tc.data[0].Value, value)
			}
		})
	}
}

func TestRunner_configEnv(t *testing.T) {
	t.Parallel()

	tt := []struct {
		name      string
		env       map[string]string
		pristine  bool
		custom    []string
		whitelist []string
		blacklist []string
		output    map[string]string
	}{
		{
			name:     "pristine env with no custom vars yields empty env",
			env:      map[string]string{"PATH": "/bin"},
			pristine: true,
			output:   map[string]string{},
		},
		{
			name:     "pristine env with custom vars only keeps custom vars",
			env:      map[string]string{"PATH": "/bin"},
			pristine: true,
			custom:   []string{"GOPATH=/usr/go"},
			output:   map[string]string{"GOPATH": "/usr/go"},
		},
		{
			name:   "custom vars overwrite input vars",
			env:    map[string]string{"PATH": "/bin"},
			custom: []string{"PATH=/usr/bin"},
			output: map[string]string{"PATH": "/usr/bin"},
		},
		{
			name:      "whitelist filters input by key",
			env:       map[string]string{"GOPATH": "/usr/go", "GO111MODULES": "true", "PATH": "/bin"},
			whitelist: []string{"GO*"},
			output:    map[string]string{"GOPATH": "/usr/go", "GO111MODULES": "true"},
		},
		{
			name:      "blacklist takes precedence over whitelist",
			env:       map[string]string{"GOPATH": "/usr/go", "PATH": "/bin", "EDITOR": "vi"},
			whitelist: []string{"GO*", "EDITOR"},
			blacklist: []string{"GO*"},
			output:    map[string]string{"EDITOR": "vi"},
		},
		{
			name:      "custom takes precedence over blacklist",
			env:       map[string]string{"PATH": "/bin", "EDITOR": "vi"},
			blacklist: []string{"EDITOR*"},
			custom:    []string{"EDITOR=nvim"},
			output:    map[string]string{"EDITOR": "nvim", "PATH": "/bin"},
		},
	}

	for _, tc := range tt {
		t.Run(tc.name, func(t *testing.T) {
			cfg := Config{
				Exec: &config.ExecConfig{
					Env: &config.EnvConfig{
						Pristine:  &tc.pristine,
						Blacklist: tc.blacklist,
						Whitelist: tc.whitelist,
						Custom:    tc.custom,
					},
				},
			}
			c := DefaultConfig().Merge(&cfg)
			r, err := NewRunner(c, true)
			if err != nil {
				t.Fatal(err)
			}
			result := r.applyConfigEnv(tc.env)

			if !reflect.DeepEqual(result, tc.output) {
				t.Fatalf("expected: %v\n got: %v", tc.output, result)
			}
		})
	}
}
<|endoftext|>