<commit_before>package proxy

import (
	"fmt"
	"reflect"
	"strconv"
	"strings"
	"sync"

	"github.com/juju/errors"
	log "github.com/ngaut/logging"
	"github.com/wandoulabs/cm/hack"
	. "github.com/wandoulabs/cm/mysql"
	"github.com/wandoulabs/cm/sqlparser"
	"github.com/wandoulabs/cm/sqltypes"
	"github.com/wandoulabs/cm/vt/schema"
	"github.com/wandoulabs/cm/vt/tabletserver"
	"github.com/wandoulabs/cm/vt/tabletserver/planbuilder"
)

func applyFilter(columnNumbers []int, input RowValue) (output RowValue) {
	output = make(RowValue, len(columnNumbers))
	for colIndex, colPointer := range columnNumbers {
		if colPointer >= 0 {
			output[colIndex] = input[colPointer]
		}
	}

	return output
}

func (c *Conn) handleQuery(sql string) (err error) {
	sql = strings.TrimRight(sql, ";")
	stmt, err := sqlparser.Parse(sql)
	if err != nil {
		return errors.Errorf(`parse sql "%s" error`, sql)
	}

	switch v := stmt.(type) {
	case *sqlparser.Select:
		return c.handleSelect(v, sql, nil)
	case *sqlparser.Insert:
		return c.handleExec(stmt, sql, nil)
	case *sqlparser.Update:
		return c.handleExec(stmt, sql, nil)
	case *sqlparser.Delete:
		return c.handleExec(stmt, sql, nil)
	case *sqlparser.Set:
		return c.handleSet(v)
	case *sqlparser.SimpleSelect:
		return c.handleSimpleSelect(sql, v)
	default:
		return errors.Errorf("statement %T not support now, %+v, %s", stmt, stmt, sql)
	}

	return nil
}

func (c *Conn) getShardList(stmt sqlparser.Statement, bindVars map[string]interface{}) ([]*Node, error) {
	if c.schema == nil {
		return nil, NewDefaultError(ER_NO_DB_ERROR)
	}

	var n []*Node
	names := c.server.getNodeNames()
	if len(names) > 0 {
		n = append(n, c.server.getNode(names[0]))
	}
	return n, nil
}

func (c *Conn) getConn(n *Node, isSelect bool) (co *SqlConn, err error) {
	if !c.needBeginTx() {
		if isSelect {
			co, err = n.getSelectConn()
		} else {
			co, err = n.getMasterConn()
		}
		if err != nil {
			return
		}
	} else {
		var ok bool
		c.Lock()
		co, ok = c.txConns[n]
		c.Unlock()

		if !ok {
			if co, err = n.getMasterConn(); err != nil {
				return
			}

			if err = co.Begin(); err != nil {
				return
			}

			c.Lock()
			c.txConns[n] = co
			c.Unlock()
		}
	}

	if err = co.UseDB(c.schema.db); err != nil {
		return
	}

	if err = co.SetCharset(c.charset); err != nil {
		return
	}

	return
}

func (c *Conn) getShardConns(isSelect bool, stmt sqlparser.Statement, bindVars map[string]interface{}) ([]*SqlConn, error) {
	nodes, err := c.getShardList(stmt, bindVars)
	if err != nil {
		return nil, err
	} else if nodes == nil {
		return nil, nil
	}

	conns := make([]*SqlConn, 0, len(nodes))

	var co *SqlConn
	for _, n := range nodes {
		co, err = c.getConn(n, isSelect)
		if err != nil {
			log.Info(err)
			break
		}

		conns = append(conns, co)
	}

	return conns, errors.Trace(err)
}

func (c *Conn) executeInShard(conns []*SqlConn, sql string, args []interface{}) ([]*Result, error) {
	var wg sync.WaitGroup
	wg.Add(len(conns))

	rs := make([]interface{}, len(conns))
	f := func(rs []interface{}, i int, co *SqlConn) {
		r, err := co.Execute(sql, args...)
		if err != nil {
			rs[i] = err
		} else {
			rs[i] = r
		}

		wg.Done()
	}

	for i, co := range conns {
		go f(rs, i, co)
	}

	wg.Wait()

	var err error
	r := make([]*Result, len(conns))
	for i, v := range rs {
		if e, ok := v.(error); ok {
			err = e
			break
		}
		r[i] = rs[i].(*Result)
	}

	return r, errors.Trace(err)
}

func (c *Conn) closeShardConns(conns []*SqlConn, rollback bool) {
	if c.isInTransaction() {
		return
	}

	for _, co := range conns {
		if rollback {
			co.Rollback()
		}

		co.Close()
	}
}

func (c *Conn) newEmptyResultset(stmt *sqlparser.Select) *Resultset {
	r := new(Resultset)
	r.Fields = make([]*Field, len(stmt.SelectExprs))

	for i, expr := range stmt.SelectExprs {
		r.Fields[i] = &Field{}
		switch e := expr.(type) {
		case *sqlparser.StarExpr:
			r.Fields[i].Name = []byte("*")
		case *sqlparser.NonStarExpr:
			if e.As != nil {
				r.Fields[i].Name = e.As
				r.Fields[i].OrgName = hack.Slice(nstring(e.Expr))
			} else {
				r.Fields[i].Name = hack.Slice(nstring(e.Expr))
			}
		default:
			r.Fields[i].Name = hack.Slice(nstring(e))
		}
	}

	r.Values = make([]RowValue, 0)
	r.RowDatas = make([]RowData, 0)

	return r
}

func makeBindVars(args []interface{}) map[string]interface{} {
	bindVars := make(map[string]interface{}, len(args))

	for i, v := range args {
		bindVars[fmt.Sprintf("v%d", i+1)] = v
	}

	return bindVars
}

func (c *Conn) getTableSchema(tableName string) (table *schema.Table, ok bool) {
	ti := c.server.autoSchamas[c.db].GetTable(tableName)
	if ti == nil {
		return nil, false
	}

	log.Infof("%+v", ti.Table)

	return ti.Table, true
}

func (c *Conn) getTableInfo(tableName string) *tabletserver.TableInfo {
	return c.server.autoSchamas[c.db].GetTable(tableName)
}

func (c *Conn) getPlanAndTableInfo(sql string) (*planbuilder.ExecPlan, *tabletserver.TableInfo, error) {
	plan, err := planbuilder.GetExecPlan(sql, c.getTableSchema)
	if err != nil {
		return nil, nil, errors.Trace(err)
	}

	log.Infof("%+v", plan)

	ti := c.getTableInfo(plan.TableName)
	if ti == nil {
		return plan, nil, errors.Errorf("unsupport sql %s", sql)
	}

	return plan, ti, nil
}

func pkValuesToStrings(pkValues []interface{}) []string {
	s := make([]string, 0)
	for _, values := range pkValues {
		switch v := values.(type) {
		case sqltypes.Value:
			s = append(s, v.String())
		case []interface{}:
			for _, value := range v {
				s = append(s, value.(sqltypes.Value).String())
			}
		default:
			log.Fatal(v, reflect.TypeOf(v))
		}
	}

	return s
}

func (c *Conn) handleSelect(stmt *sqlparser.Select, sql string, args []interface{}) error {
	// handle cache
	plan, ti, err := c.getPlanAndTableInfo(sql)
	if err != nil {
		return errors.Trace(err)
	}

	//todo: fix hard code
	keys := pkValuesToStrings(plan.PKValues)
	items := ti.Cache.Get(keys)
	if len(items) == len(keys) { //all cache hint
		var fields []string //construct field name
		for _, i := range plan.ColumnNumbers {
			c := ti.Columns[i]
			fields = append(fields, c.Name)
		}

		var values []RowValue
		for _, key := range keys {
			row, ok := items[key]
			if !ok {
				log.Fatal("should never happend")
			}
			retValue := applyFilter(plan.ColumnNumbers, row.Row)
			values = append(values, retValue)

			log.Info("hit cache!", sql)
		}

		r, err := c.buildResultset(fields, values)
		if err != nil {
			log.Error(err)
			return errors.Trace(err)
		}

		return c.writeResultset(c.status, r)
	}

	if plan.PlanId == planbuilder.PLAN_PK_IN {
		if len(keys) > 1 {
			break
		}

		pk := ti.Columns[ti.PKColumns[0]]
		rowsql := fmt.Sprintf("select * from %s where %s = %s", plan.TableName, pk.Name, plan.PKValues[0])
		log.Info(rowsql)

		result, err := c.server.autoSchamas[c.db].Exec(rowsql)
		if err != nil {
			return errors.Trace(err)
		}

		log.Infof("%s, %+v, %+v", sql, plan, stmt)
		retValues := applyFilter(plan.ColumnNumbers, result.Values[0])
		log.Infof("%+v", retValues)

		//just simple cache just now
		if len(result.Values) == 1 && len(keys) == 1 {
			ti.Cache.Set(keys[0], result.Values[0], 0)
		}
	}

	bindVars := makeBindVars(args)

	conns, err := c.getShardConns(true, stmt, bindVars)
	if err != nil {
		return errors.Trace(err)
	} else if conns == nil {
		r := c.newEmptyResultset(stmt)
		return c.writeResultset(c.status, r)
	}

	var rs []*Result
	rs, err = c.executeInShard(conns, sql, args)
	c.closeShardConns(conns, false)
	if err == nil {
		err = c.mergeSelectResult(rs, stmt)
	}

	return errors.Trace(err)
}

func (c *Conn) handleExec(stmt sqlparser.Statement, sql string, args []interface{}) error {
	bindVars := makeBindVars(args)

	conns, err := c.getShardConns(false, stmt, bindVars)
	if err != nil {
		return errors.Trace(err)
	} else if conns == nil {
		return c.writeOK(nil)
	}

	var rs []*Result
	if len(conns) == 1 {
		rs, err = c.executeInShard(conns, sql, args)
	} else {
		log.Warning("not implement yet")

	}

	c.closeShardConns(conns, err != nil)

	if err == nil {
		err = c.mergeExecResult(rs)
	}

	return errors.Trace(err)
}

func (c *Conn) mergeExecResult(rs []*Result) error {
	r := new(Result)

	for _, v := range rs {
		r.Status |= v.Status
		r.AffectedRows += v.AffectedRows
		if r.InsertId == 0 {
			r.InsertId = v.InsertId
		} else if r.InsertId > v.InsertId {

			r.InsertId = v.InsertId
		}
	}

	if r.InsertId > 0 {
		c.lastInsertId = int64(r.InsertId)
	}

	c.affectedRows = int64(r.AffectedRows)

	return c.writeOK(r)
}

func (c *Conn) mergeSelectResult(rs []*Result, stmt *sqlparser.Select) error {
	r := rs[0].Resultset

	status := c.status | rs[0].Status

	for i := 1; i < len(rs); i++ {
		status |= rs[i].Status

		for j := range rs[i].Values {
			r.Values = append(r.Values, rs[i].Values[j])
			r.RowDatas = append(r.RowDatas, rs[i].RowDatas[j])
		}
	}

	c.sortSelectResult(r, stmt)

	if err := c.limitSelectResult(r, stmt); err != nil {
		return errors.Trace(err)
	}

	return c.writeResultset(status, r)
}

func (c *Conn) sortSelectResult(r *Resultset, stmt *sqlparser.Select) error {
	if stmt.OrderBy == nil {
		return nil
	}

	sk := make([]SortKey, len(stmt.OrderBy))

	for i, o := range stmt.OrderBy {
		sk[i].Name = nstring(o.Expr)
		sk[i].Direction = o.Direction
	}

	return r.Sort(sk)
}

func (c *Conn) limitSelectResult(r *Resultset, stmt *sqlparser.Select) error {
	if stmt.Limit == nil {
		return nil
	}

	var offset, count int64
	var err error
	if stmt.Limit.Offset == nil {
		offset = 0
	} else {
		if o, ok := stmt.Limit.Offset.(sqlparser.NumVal); !ok {
			return errors.Errorf("invalid select limit %s", nstring(stmt.Limit))
		} else {
			if offset, err = strconv.ParseInt(hack.String([]byte(o)), 10, 64); err != nil {
				return errors.Trace(err)
			}
		}
	}

	if o, ok := stmt.Limit.Rowcount.(sqlparser.NumVal); !ok {
		return errors.Errorf("invalid limit %s", nstring(stmt.Limit))
	} else {
		if count, err = strconv.ParseInt(hack.String([]byte(o)), 10, 64); err != nil {
			return errors.Trace(err)
		} else if count < 0 {
			return errors.Errorf("invalid limit %s", nstring(stmt.Limit))
		}
	}

	if offset+count > int64(len(r.Values)) {
		count = int64(len(r.Values)) - offset
	}

	r.Values = r.Values[offset : offset+count]
	r.RowDatas = r.RowDatas[offset : offset+count]

	return nil
}
<commit_msg>more cache & refactor<commit_after>package proxy

import (
	"fmt"
	"reflect"
	"strconv"
	"strings"
	"sync"

	"github.com/juju/errors"
	log "github.com/ngaut/logging"
	"github.com/wandoulabs/cm/hack"
	. "github.com/wandoulabs/cm/mysql"
	"github.com/wandoulabs/cm/sqlparser"
	"github.com/wandoulabs/cm/sqltypes"
	"github.com/wandoulabs/cm/vt/schema"
	"github.com/wandoulabs/cm/vt/tabletserver"
	"github.com/wandoulabs/cm/vt/tabletserver/planbuilder"
)

func applyFilter(columnNumbers []int, input RowValue) (output RowValue) {
	output = make(RowValue, len(columnNumbers))
	for colIndex, colPointer := range columnNumbers {
		if colPointer >= 0 {
			output[colIndex] = input[colPointer]
		}
	}

	return output
}

func (c *Conn) handleQuery(sql string) (err error) {
	sql = strings.TrimRight(sql, ";")
	stmt, err := sqlparser.Parse(sql)
	if err != nil {
		return errors.Errorf(`parse sql "%s" error`, sql)
	}

	switch v := stmt.(type) {
	case *sqlparser.Select:
		return c.handleSelect(v, sql, nil)
	case *sqlparser.Insert:
		return c.handleExec(stmt, sql, nil)
	case *sqlparser.Update:
		return c.handleExec(stmt, sql, nil)
	case *sqlparser.Delete:
		return c.handleExec(stmt, sql, nil)
	case *sqlparser.Set:
		return c.handleSet(v)
	case *sqlparser.SimpleSelect:
		return c.handleSimpleSelect(sql, v)
	default:
		return errors.Errorf("statement %T not support now, %+v, %s", stmt, stmt, sql)
	}

	return nil
}

func (c *Conn) getShardList(stmt sqlparser.Statement, bindVars map[string]interface{}) ([]*Node, error) {
	if c.schema == nil {
		return nil, NewDefaultError(ER_NO_DB_ERROR)
	}

	var n []*Node
	names := c.server.getNodeNames()
	if len(names) > 0 {
		n = append(n, c.server.getNode(names[0]))
	}
	return n, nil
}

func (c *Conn) getConn(n *Node, isSelect bool) (co *SqlConn, err error) {
	if !c.needBeginTx() {
		if isSelect {
			co, err = n.getSelectConn()
		} else {
			co, err = n.getMasterConn()
		}
		if err != nil {
			return
		}
	} else {
		var ok bool
		c.Lock()
		co, ok = c.txConns[n]
		c.Unlock()

		if !ok {
			if co, err = n.getMasterConn(); err != nil {
				return
			}

			if err = co.Begin(); err != nil {
				return
			}

			c.Lock()
			c.txConns[n] = co
			c.Unlock()
		}
	}

	if err = co.UseDB(c.schema.db); err != nil {
		return
	}

	if err = co.SetCharset(c.charset); err != nil {
		return
	}

	return
}

func (c *Conn) getShardConns(isSelect bool, stmt sqlparser.Statement, bindVars map[string]interface{}) ([]*SqlConn, error) {
	nodes, err := c.getShardList(stmt, bindVars)
	if err != nil {
		return nil, err
	} else if nodes == nil {
		return nil, nil
	}

	conns := make([]*SqlConn, 0, len(nodes))

	var co *SqlConn
	for _, n := range nodes {
		co, err = c.getConn(n, isSelect)
		if err != nil {
			log.Info(err)
			break
		}

		conns = append(conns, co)
	}

	return conns, errors.Trace(err)
}

func (c *Conn) executeInShard(conns []*SqlConn, sql string, args []interface{}) ([]*Result, error) {
	var wg sync.WaitGroup
	wg.Add(len(conns))

	rs := make([]interface{}, len(conns))
	f := func(rs []interface{}, i int, co *SqlConn) {
		r, err := co.Execute(sql, args...)
		if err != nil {
			rs[i] = err
		} else {
			rs[i] = r
		}

		wg.Done()
	}

	for i, co := range conns {
		go f(rs, i, co)
	}

	wg.Wait()

	var err error
	r := make([]*Result, len(conns))
	for i, v := range rs {
		if e, ok := v.(error); ok {
			err = e
			break
		}
		r[i] = rs[i].(*Result)
	}

	return r, errors.Trace(err)
}

func (c *Conn) closeShardConns(conns []*SqlConn, rollback bool) {
	if c.isInTransaction() {
		return
	}

	for _, co := range conns {
		if rollback {
			co.Rollback()
		}

		co.Close()
	}
}

func (c *Conn) newEmptyResultset(stmt *sqlparser.Select) *Resultset {
	r := new(Resultset)
	r.Fields = make([]*Field, len(stmt.SelectExprs))

	for i, expr := range stmt.SelectExprs {
		r.Fields[i] = &Field{}
		switch e := expr.(type) {
		case *sqlparser.StarExpr:
			r.Fields[i].Name = []byte("*")
		case *sqlparser.NonStarExpr:
			if e.As != nil {
				r.Fields[i].Name = e.As
				r.Fields[i].OrgName = hack.Slice(nstring(e.Expr))
			} else {
				r.Fields[i].Name = hack.Slice(nstring(e.Expr))
			}
		default:
			r.Fields[i].Name = hack.Slice(nstring(e))
		}
	}

	r.Values = make([]RowValue, 0)
	r.RowDatas = make([]RowData, 0)

	return r
}

func makeBindVars(args []interface{}) map[string]interface{} {
	bindVars := make(map[string]interface{}, len(args))

	for i, v := range args {
		bindVars[fmt.Sprintf("v%d", i+1)] = v
	}

	return bindVars
}

func (c *Conn) getTableSchema(tableName string) (table *schema.Table, ok bool) {
	ti := c.server.autoSchamas[c.db].GetTable(tableName)
	if ti == nil {
		return nil, false
	}

	log.Infof("%+v", ti.Table)

	return ti.Table, true
}

func (c *Conn) getTableInfo(tableName string) *tabletserver.TableInfo {
	return c.server.autoSchamas[c.db].GetTable(tableName)
}

func (c *Conn) getPlanAndTableInfo(sql string) (*planbuilder.ExecPlan, *tabletserver.TableInfo, error) {
	plan, err := planbuilder.GetExecPlan(sql, c.getTableSchema)
	if err != nil {
		return nil, nil, errors.Trace(err)
	}

	log.Infof("%+v", plan)

	ti := c.getTableInfo(plan.TableName)
	if ti == nil {
		return plan, nil, errors.Errorf("unsupport sql %s", sql)
	}

	return plan, ti, nil
}

func pkValuesToStrings(pkValues []interface{}) []string {
	s := make([]string, 0)
	for _, values := range pkValues {
		switch v := values.(type) {
		case sqltypes.Value:
			s = append(s, v.String())
		case []interface{}:
			for _, value := range v {
				s = append(s, value.(sqltypes.Value).String())
			}
		default:
			log.Fatal(v, reflect.TypeOf(v))
		}
	}

	return s
}

func getFieldNames(plan *planbuilder.ExecPlan, ti *tabletserver.TableInfo) []string {
	fields := make([]string, 0, len(plan.ColumnNumbers)) //construct field name
	for _, i := range plan.ColumnNumbers {
		c := ti.Columns[i]
		fields = append(fields, c.Name)
	}

	return fields
}

func (c *Conn) handleSelect(stmt *sqlparser.Select, sql string, args []interface{}) error {
	// handle cache
	plan, ti, err := c.getPlanAndTableInfo(sql)
	if err != nil {
		return errors.Trace(err)
	}

	//todo: fix hard code
	keys := pkValuesToStrings(plan.PKValues)
	items := ti.Cache.Get(keys)
	if len(items) == len(keys) { //all cache hint
		var values []RowValue
		for _, key := range keys {
			row, ok := items[key]
			if !ok {
				log.Fatal("should never happend")
			}
			retValue := applyFilter(plan.ColumnNumbers, row.Row)
			values = append(values, retValue)

			log.Info("hit cache!", sql)
		}

		r, err := c.buildResultset(getFieldNames(plan, ti), values)
		if err != nil {
			log.Error(err)
			return errors.Trace(err)
		}

		return c.writeResultset(c.status, r)
	}

	if plan.PlanId == planbuilder.PLAN_PK_IN && len(keys) == 1 {
		pk := ti.Columns[ti.PKColumns[0]]
		rowsql := fmt.Sprintf("select * from %s where %s = %s", plan.TableName, pk.Name, plan.PKValues[0])
		log.Info(rowsql)

		result, err := c.server.autoSchamas[c.db].Exec(rowsql)
		if err != nil {
			return errors.Trace(err)
		}

		log.Infof("%s, %+v, %+v", sql, plan, stmt)
		retValues := applyFilter(plan.ColumnNumbers, result.Values[0])
		log.Infof("%+v", retValues)

		//just simple cache just now
		if len(result.Values) == 1 && len(keys) == 1 {
			ti.Cache.Set(keys[0], result.Values[0], 0)
		}

		var values []RowValue
		values = append(values, retValues)
		r, err := c.buildResultset(getFieldNames(plan, ti), values)
		if err != nil {
			log.Error(err)
			return errors.Trace(err)
		}

		return c.writeResultset(c.status, r)

	}

	bindVars := makeBindVars(args)

	conns, err := c.getShardConns(true, stmt, bindVars)
	if err != nil {
		return errors.Trace(err)
	} else if conns == nil {
		r := c.newEmptyResultset(stmt)
		return c.writeResultset(c.status, r)
	}

	var rs []*Result
	rs, err = c.executeInShard(conns, sql, args)
	c.closeShardConns(conns, false)
	if err == nil {
		err = c.mergeSelectResult(rs, stmt)
	}

	return errors.Trace(err)
}

func (c *Conn) handleExec(stmt sqlparser.Statement, sql string, args []interface{}) error {
	bindVars := makeBindVars(args)

	conns, err := c.getShardConns(false, stmt, bindVars)
	if err != nil {
		return errors.Trace(err)
	} else if conns == nil {
		return c.writeOK(nil)
	}

	var rs []*Result
	if len(conns) == 1 {
		rs, err = c.executeInShard(conns, sql, args)
	} else {
		log.Warning("not implement yet")

	}

	c.closeShardConns(conns, err != nil)

	if err == nil {
		err = c.mergeExecResult(rs)
	}

	return errors.Trace(err)
}

func (c *Conn) mergeExecResult(rs []*Result) error {
	r := new(Result)

	for _, v := range rs {
		r.Status |= v.Status
		r.AffectedRows += v.AffectedRows
		if r.InsertId == 0 {
			r.InsertId = v.InsertId
		} else if r.InsertId > v.InsertId {

			r.InsertId = v.InsertId
		}
	}

	if r.InsertId > 0 {
		c.lastInsertId = int64(r.InsertId)
	}

	c.affectedRows = int64(r.AffectedRows)

	return c.writeOK(r)
}

func (c *Conn) mergeSelectResult(rs []*Result, stmt *sqlparser.Select) error {
	r := rs[0].Resultset

	status := c.status | rs[0].Status

	for i := 1; i < len(rs); i++ {
		status |= rs[i].Status

		for j := range rs[i].Values {
			r.Values = append(r.Values, rs[i].Values[j])
			r.RowDatas = append(r.RowDatas, rs[i].RowDatas[j])
		}
	}

	c.sortSelectResult(r, stmt)

	if err := c.limitSelectResult(r, stmt); err != nil {
		return errors.Trace(err)
	}

	return c.writeResultset(status, r)
}

func (c *Conn) sortSelectResult(r *Resultset, stmt *sqlparser.Select) error {
	if stmt.OrderBy == nil {
		return nil
	}

	sk := make([]SortKey, len(stmt.OrderBy))

	for i, o := range stmt.OrderBy {
		sk[i].Name = nstring(o.Expr)
		sk[i].Direction = o.Direction
	}

	return r.Sort(sk)
}

func (c *Conn) limitSelectResult(r *Resultset, stmt *sqlparser.Select) error {
	if stmt.Limit == nil {
		return nil
	}

	var offset, count int64
	var err error
	if stmt.Limit.Offset == nil {
		offset = 0
	} else {
		if o, ok := stmt.Limit.Offset.(sqlparser.NumVal); !ok {
			return errors.Errorf("invalid select limit %s", nstring(stmt.Limit))
		} else {
			if offset, err = strconv.ParseInt(hack.String([]byte(o)), 10, 64); err != nil {
				return errors.Trace(err)
			}
		}
	}

	if o, ok := stmt.Limit.Rowcount.(sqlparser.NumVal); !ok {
		return errors.Errorf("invalid limit %s", nstring(stmt.Limit))
	} else {
		if count, err = strconv.ParseInt(hack.String([]byte(o)), 10, 64); err != nil {
			return errors.Trace(err)
		} else if count < 0 {
			return errors.Errorf("invalid limit %s", nstring(stmt.Limit))
		}
	}

	if offset+count > int64(len(r.Values)) {
		count = int64(len(r.Values)) - offset
	}

	r.Values = r.Values[offset : offset+count]
	r.RowDatas = r.RowDatas[offset : offset+count]

	return nil
}
<|endoftext|>


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

package proxyserver

import (
	"flag"
	"fmt"
	"io"
	"net/http"
	_ "net/http/pprof"
	"strings"
	"time"

	"github.com/troubling/hummingbird/client"
	"github.com/troubling/hummingbird/common"
	"github.com/troubling/hummingbird/common/conf"
	"github.com/troubling/hummingbird/common/ring"
	"github.com/troubling/hummingbird/common/srv"
	"github.com/troubling/hummingbird/proxyserver/middleware"

	"github.com/justinas/alice"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/uber-go/tally"
	promreporter "github.com/uber-go/tally/prometheus"
	"go.uber.org/zap"
)

type ProxyServer struct {
	logger            srv.LowLevelLogger
	logLevel          zap.AtomicLevel
	mc                ring.MemcacheRing
	accountAutoCreate bool
	proxyDirectClient *client.ProxyDirectClient
	metricsCloser     io.Closer
}

func (server *ProxyServer) Type() string {
	return "proxy"
}

func (server *ProxyServer) Background(flags *flag.FlagSet) chan struct{} {
	return nil
}

func (server *ProxyServer) Finalize() {
	if server.metricsCloser != nil {
		server.metricsCloser.Close()
	}
}

func (server *ProxyServer) GetHandler(config conf.Config, metricsPrefix string) http.Handler {
	var metricsScope tally.Scope
	metricsScope, server.metricsCloser = tally.NewRootScope(tally.ScopeOptions{
		Prefix:         metricsPrefix,
		Tags:           map[string]string{},
		CachedReporter: promreporter.NewReporter(promreporter.Options{}),
		Separator:      promreporter.DefaultSeparator,
	}, time.Second)
	router := srv.NewRouter()
	router.Get("/metrics", prometheus.Handler())
	router.Get("/loglevel", server.logLevel)
	router.Put("/loglevel", server.logLevel)
	router.Get("/v1/:account/:container/*obj", http.HandlerFunc(server.ObjectGetHandler))
	router.Head("/v1/:account/:container/*obj", http.HandlerFunc(server.ObjectHeadHandler))
	router.Put("/v1/:account/:container/*obj", http.HandlerFunc(server.ObjectPutHandler))
	router.Delete("/v1/:account/:container/*obj", http.HandlerFunc(server.ObjectDeleteHandler))
	router.Post("/v1/:account/:container/*obj", http.HandlerFunc(server.ObjectPostHandler))
	router.Options("/v1/:account/:container/*obj", http.HandlerFunc(server.OptionsHandler))

	router.Get("/v1/:account/:container", http.HandlerFunc(server.ContainerGetHandler))
	router.Get("/v1/:account/:container/", http.HandlerFunc(server.ContainerGetHandler))
	router.Head("/v1/:account/:container", http.HandlerFunc(server.ContainerHeadHandler))
	router.Head("/v1/:account/:container/", http.HandlerFunc(server.ContainerHeadHandler))
	router.Put("/v1/:account/:container", http.HandlerFunc(server.ContainerPutHandler))
	router.Put("/v1/:account/:container/", http.HandlerFunc(server.ContainerPutHandler))
	router.Delete("/v1/:account/:container", http.HandlerFunc(server.ContainerDeleteHandler))
	router.Delete("/v1/:account/:container/", http.HandlerFunc(server.ContainerDeleteHandler))
	router.Post("/v1/:account/:container", http.HandlerFunc(server.ContainerPostHandler))
	router.Post("/v1/:account/:container/", http.HandlerFunc(server.ContainerPostHandler))
	router.Options("/v1/:account/:container", http.HandlerFunc(server.OptionsHandler))
	router.Options("/v1/:account/:container/", http.HandlerFunc(server.OptionsHandler))

	router.Get("/v1/:account", http.HandlerFunc(server.AccountGetHandler))
	router.Get("/v1/:account/", http.HandlerFunc(server.AccountGetHandler))
	router.Head("/v1/:account", http.HandlerFunc(server.AccountHeadHandler))
	router.Head("/v1/:account/", http.HandlerFunc(server.AccountHeadHandler))
	router.Put("/v1/:account", http.HandlerFunc(server.AccountPutHandler))
	router.Put("/v1/:account/", http.HandlerFunc(server.AccountPutHandler))
	router.Delete("/v1/:account", http.HandlerFunc(server.AccountDeleteHandler))
	router.Delete("/v1/:account/", http.HandlerFunc(server.AccountDeleteHandler))
	router.Post("/v1/:account", http.HandlerFunc(server.AccountPostHandler))
	router.Post("/v1/:account/", http.HandlerFunc(server.AccountPostHandler))
	router.Options("/v1/:account", http.HandlerFunc(server.OptionsHandler))
	router.Options("/v1/:account/", http.HandlerFunc(server.OptionsHandler))

	if config.GetBool("proxy-server", "pprof_enabled", false) {
		router.Get("/debug/pprof/:parm", http.DefaultServeMux)
		router.Post("/debug/pprof/:parm", http.DefaultServeMux)
	}

	tempAuth := config.GetBool("proxy-server", "tempauth_enabled", true)
	var middlewares []struct {
		construct func(conf.Section, tally.Scope) (func(http.Handler) http.Handler, error)
		section   string
	}
	// TODO: make this all dynamical and stuff
	if tempAuth {
		middlewares = []struct {
			construct func(conf.Section, tally.Scope) (func(http.Handler) http.Handler, error)
			section   string
		}{
			{middleware.NewCatchError, "filter:catch_errors"},
			{middleware.NewHealthcheck, "filter:healthcheck"},
			{middleware.NewRequestLogger, "filter:proxy-logging"},
			{middleware.NewCors, "filter:cors"}, // TODO: i dont want to have to have a seciton for this
			{middleware.NewFormPost, "filter:formpost"},
			{middleware.NewTempURL, "filter:tempurl"},
			{middleware.NewTempAuth, "filter:tempauth"},
			{middleware.NewBulk, "filter:bulk"},
			{middleware.NewMultirange, "filter:multirange"},
			{middleware.NewRatelimiter, "filter:ratelimit"},
			{middleware.NewStaticWeb, "filter:staticweb"},
			{middleware.NewCopyMiddleware, "filter:copy"},
			{middleware.NewVersionedWrites, "filter:versioned_writes"},
			{middleware.NewXlo, "filter:slo"},
		}
	} else {
		middlewares = []struct {
			construct func(conf.Section, tally.Scope) (func(http.Handler) http.Handler, error)
			section   string
		}{
			{middleware.NewCatchError, "filter:catch_errors"},
			{middleware.NewHealthcheck, "filter:healthcheck"},
			{middleware.NewRequestLogger, "filter:proxy-logging"},
			{middleware.NewCors, "filter:cors"},
			{middleware.NewFormPost, "filter:formpost"},
			{middleware.NewTempURL, "filter:tempurl"},
			{middleware.NewAuthToken, "filter:authtoken"},
			{middleware.NewKeystoneAuth, "filter:keystoneauth"},
			{middleware.NewBulk, "filter:bulk"},
			{middleware.NewMultirange, "filter:multirange"},
			{middleware.NewRatelimiter, "filter:ratelimit"},
			{middleware.NewStaticWeb, "filter:staticweb"},
			{middleware.NewCopyMiddleware, "filter:copy"},
			{middleware.NewVersionedWrites, "filter:versioned_writes"},
			{middleware.NewXlo, "filter:slo"},
		}
	}
	pipeline := alice.New(middleware.NewContext(config.GetBool("debug", "debug_x_source_code", false),
		server.mc, server.logger, server.proxyDirectClient))
	for _, m := range middlewares {
		mid, err := m.construct(config.GetSection(m.section), metricsScope)
		if err != nil {
			// TODO: propagate error upwards instead of panicking
			panic("Unable to construct middleware")
		}
		pipeline = pipeline.Append(mid)
	}
	return pipeline.Then(router)
}

func NewServer(serverconf conf.Config, flags *flag.FlagSet, cnf srv.ConfigLoader) (*srv.IpPort, srv.Server, srv.LowLevelLogger, error) {
	var err error
	var ipPort *srv.IpPort
	server := &ProxyServer{}
	server.mc, err = ring.NewMemcacheRingFromConfig(serverconf)
	if err != nil {
		return ipPort, nil, nil, err
	}

	bindIP := serverconf.GetDefault("DEFAULT", "bind_ip", "0.0.0.0")
	bindPort := int(serverconf.GetInt("DEFAULT", "bind_port", common.DefaultProxyServerPort))
	certFile := serverconf.GetDefault("DEFAULT", "cert_file", "")
	keyFile := serverconf.GetDefault("DEFAULT", "key_file", "")

	logLevelString := serverconf.GetDefault("app:proxy-server", "log_level", "INFO")
	server.logLevel = zap.NewAtomicLevel()
	server.logLevel.UnmarshalText([]byte(strings.ToLower(logLevelString)))
	server.accountAutoCreate = serverconf.GetBool("app:proxy-server", "account_autocreate", false)
	if server.logger, err = srv.SetupLogger("proxy-server", &server.logLevel, flags); err != nil {
		return ipPort, nil, nil, fmt.Errorf("Error setting up logger: %v", err)
	}
	policies, err := cnf.GetPolicies()
	if err != nil {
		return ipPort, nil, nil, err
	}
	server.proxyDirectClient, err = client.NewProxyDirectClient(policies, cnf, server.logger, certFile, keyFile)
	if err != nil {
		return ipPort, nil, nil, fmt.Errorf("Error setting up proxyDirectClient: %v", err)
	}
	info := map[string]interface{}{
		"version":                  common.Version,
		"strict_cors_mode":         true,
		"policies":                 policies.GetPolicyInfo(),
		"account_autocreate":       server.accountAutoCreate,
		"allow_account_management": true,
	}
	for k, v := range common.DEFAULT_CONSTRAINTS {
		info[k] = v
	}
	middleware.RegisterInfo("swift", info)
	ipPort = &srv.IpPort{Ip: bindIP, Port: int(bindPort), CertFile: certFile, KeyFile: keyFile}
	return ipPort, server, server.logger, nil
}
<commit_msg>HUM-406 Protect proxy server's /loglevel etc.<commit_after>//  Copyright (c) 2015 Rackspace
//
//  Licensed under the Apache License, Version 2.0 (the "License");
//  you may not use this file except in compliance with the License.
//  You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
//  Unless required by applicable law or agreed to in writing, software
//  distributed under the License is distributed on an "AS IS" BASIS,
//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
//  implied.
//  See the License for the specific language governing permissions and
//  limitations under the License.

package proxyserver

import (
	"flag"
	"fmt"
	"io"
	"net/http"
	_ "net/http/pprof"
	"path"
	"strings"
	"time"

	"github.com/troubling/hummingbird/client"
	"github.com/troubling/hummingbird/common"
	"github.com/troubling/hummingbird/common/conf"
	"github.com/troubling/hummingbird/common/ring"
	"github.com/troubling/hummingbird/common/srv"
	"github.com/troubling/hummingbird/proxyserver/middleware"

	"github.com/justinas/alice"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/uber-go/tally"
	promreporter "github.com/uber-go/tally/prometheus"
	"go.uber.org/zap"
)

type ProxyServer struct {
	logger            srv.LowLevelLogger
	logLevel          zap.AtomicLevel
	mc                ring.MemcacheRing
	accountAutoCreate bool
	proxyDirectClient *client.ProxyDirectClient
	metricsCloser     io.Closer
}

func (server *ProxyServer) Type() string {
	return "proxy"
}

func (server *ProxyServer) Background(flags *flag.FlagSet) chan struct{} {
	return nil
}

func (server *ProxyServer) Finalize() {
	if server.metricsCloser != nil {
		server.metricsCloser.Close()
	}
}

func (server *ProxyServer) GetHandler(config conf.Config, metricsPrefix string) http.Handler {
	obfuscatedPrefix, _ := config.Get("proxy-server", "obfuscated_prefix")
	var metricsScope tally.Scope
	metricsScope, server.metricsCloser = tally.NewRootScope(tally.ScopeOptions{
		Prefix:         metricsPrefix,
		Tags:           map[string]string{},
		CachedReporter: promreporter.NewReporter(promreporter.Options{}),
		Separator:      promreporter.DefaultSeparator,
	}, time.Second)
	router := srv.NewRouter()
	if obfuscatedPrefix != "" {
		router.Get(path.Join("/", obfuscatedPrefix, "metrics"), prometheus.Handler())
		router.Get(path.Join("/", obfuscatedPrefix, "loglevel"), server.logLevel)
		router.Put(path.Join("/", obfuscatedPrefix, "loglevel"), server.logLevel)
		router.Get(path.Join("/", obfuscatedPrefix, "debug/pprof/:parm"), http.DefaultServeMux)
		router.Post(path.Join("/", obfuscatedPrefix, "debug/pprof/:parm"), http.DefaultServeMux)
	}
	router.Get("/v1/:account/:container/*obj", http.HandlerFunc(server.ObjectGetHandler))
	router.Head("/v1/:account/:container/*obj", http.HandlerFunc(server.ObjectHeadHandler))
	router.Put("/v1/:account/:container/*obj", http.HandlerFunc(server.ObjectPutHandler))
	router.Delete("/v1/:account/:container/*obj", http.HandlerFunc(server.ObjectDeleteHandler))
	router.Post("/v1/:account/:container/*obj", http.HandlerFunc(server.ObjectPostHandler))
	router.Options("/v1/:account/:container/*obj", http.HandlerFunc(server.OptionsHandler))

	router.Get("/v1/:account/:container", http.HandlerFunc(server.ContainerGetHandler))
	router.Get("/v1/:account/:container/", http.HandlerFunc(server.ContainerGetHandler))
	router.Head("/v1/:account/:container", http.HandlerFunc(server.ContainerHeadHandler))
	router.Head("/v1/:account/:container/", http.HandlerFunc(server.ContainerHeadHandler))
	router.Put("/v1/:account/:container", http.HandlerFunc(server.ContainerPutHandler))
	router.Put("/v1/:account/:container/", http.HandlerFunc(server.ContainerPutHandler))
	router.Delete("/v1/:account/:container", http.HandlerFunc(server.ContainerDeleteHandler))
	router.Delete("/v1/:account/:container/", http.HandlerFunc(server.ContainerDeleteHandler))
	router.Post("/v1/:account/:container", http.HandlerFunc(server.ContainerPostHandler))
	router.Post("/v1/:account/:container/", http.HandlerFunc(server.ContainerPostHandler))
	router.Options("/v1/:account/:container", http.HandlerFunc(server.OptionsHandler))
	router.Options("/v1/:account/:container/", http.HandlerFunc(server.OptionsHandler))

	router.Get("/v1/:account", http.HandlerFunc(server.AccountGetHandler))
	router.Get("/v1/:account/", http.HandlerFunc(server.AccountGetHandler))
	router.Head("/v1/:account", http.HandlerFunc(server.AccountHeadHandler))
	router.Head("/v1/:account/", http.HandlerFunc(server.AccountHeadHandler))
	router.Put("/v1/:account", http.HandlerFunc(server.AccountPutHandler))
	router.Put("/v1/:account/", http.HandlerFunc(server.AccountPutHandler))
	router.Delete("/v1/:account", http.HandlerFunc(server.AccountDeleteHandler))
	router.Delete("/v1/:account/", http.HandlerFunc(server.AccountDeleteHandler))
	router.Post("/v1/:account", http.HandlerFunc(server.AccountPostHandler))
	router.Post("/v1/:account/", http.HandlerFunc(server.AccountPostHandler))
	router.Options("/v1/:account", http.HandlerFunc(server.OptionsHandler))
	router.Options("/v1/:account/", http.HandlerFunc(server.OptionsHandler))

	tempAuth := config.GetBool("proxy-server", "tempauth_enabled", true)
	var middlewares []struct {
		construct func(conf.Section, tally.Scope) (func(http.Handler) http.Handler, error)
		section   string
	}
	// TODO: make this all dynamical and stuff
	if tempAuth {
		middlewares = []struct {
			construct func(conf.Section, tally.Scope) (func(http.Handler) http.Handler, error)
			section   string
		}{
			{middleware.NewCatchError, "filter:catch_errors"},
			{middleware.NewHealthcheck, "filter:healthcheck"},
			{middleware.NewRequestLogger, "filter:proxy-logging"},
			{middleware.NewCors, "filter:cors"}, // TODO: i dont want to have to have a seciton for this
			{middleware.NewFormPost, "filter:formpost"},
			{middleware.NewTempURL, "filter:tempurl"},
			{middleware.NewTempAuth, "filter:tempauth"},
			{middleware.NewBulk, "filter:bulk"},
			{middleware.NewMultirange, "filter:multirange"},
			{middleware.NewRatelimiter, "filter:ratelimit"},
			{middleware.NewStaticWeb, "filter:staticweb"},
			{middleware.NewCopyMiddleware, "filter:copy"},
			{middleware.NewVersionedWrites, "filter:versioned_writes"},
			{middleware.NewXlo, "filter:slo"},
		}
	} else {
		middlewares = []struct {
			construct func(conf.Section, tally.Scope) (func(http.Handler) http.Handler, error)
			section   string
		}{
			{middleware.NewCatchError, "filter:catch_errors"},
			{middleware.NewHealthcheck, "filter:healthcheck"},
			{middleware.NewRequestLogger, "filter:proxy-logging"},
			{middleware.NewCors, "filter:cors"},
			{middleware.NewFormPost, "filter:formpost"},
			{middleware.NewTempURL, "filter:tempurl"},
			{middleware.NewAuthToken, "filter:authtoken"},
			{middleware.NewKeystoneAuth, "filter:keystoneauth"},
			{middleware.NewBulk, "filter:bulk"},
			{middleware.NewMultirange, "filter:multirange"},
			{middleware.NewRatelimiter, "filter:ratelimit"},
			{middleware.NewStaticWeb, "filter:staticweb"},
			{middleware.NewCopyMiddleware, "filter:copy"},
			{middleware.NewVersionedWrites, "filter:versioned_writes"},
			{middleware.NewXlo, "filter:slo"},
		}
	}
	pipeline := alice.New(middleware.NewContext(config.GetBool("debug", "debug_x_source_code", false),
		server.mc, server.logger, server.proxyDirectClient))
	for _, m := range middlewares {
		mid, err := m.construct(config.GetSection(m.section), metricsScope)
		if err != nil {
			// TODO: propagate error upwards instead of panicking
			panic("Unable to construct middleware")
		}
		pipeline = pipeline.Append(mid)
	}
	return pipeline.Then(router)
}

func NewServer(serverconf conf.Config, flags *flag.FlagSet, cnf srv.ConfigLoader) (*srv.IpPort, srv.Server, srv.LowLevelLogger, error) {
	var err error
	var ipPort *srv.IpPort
	server := &ProxyServer{}
	server.mc, err = ring.NewMemcacheRingFromConfig(serverconf)
	if err != nil {
		return ipPort, nil, nil, err
	}

	bindIP := serverconf.GetDefault("DEFAULT", "bind_ip", "0.0.0.0")
	bindPort := int(serverconf.GetInt("DEFAULT", "bind_port", common.DefaultProxyServerPort))
	certFile := serverconf.GetDefault("DEFAULT", "cert_file", "")
	keyFile := serverconf.GetDefault("DEFAULT", "key_file", "")

	logLevelString := serverconf.GetDefault("app:proxy-server", "log_level", "INFO")
	server.logLevel = zap.NewAtomicLevel()
	server.logLevel.UnmarshalText([]byte(strings.ToLower(logLevelString)))
	server.accountAutoCreate = serverconf.GetBool("app:proxy-server", "account_autocreate", false)
	if server.logger, err = srv.SetupLogger("proxy-server", &server.logLevel, flags); err != nil {
		return ipPort, nil, nil, fmt.Errorf("Error setting up logger: %v", err)
	}
	policies, err := cnf.GetPolicies()
	if err != nil {
		return ipPort, nil, nil, err
	}
	server.proxyDirectClient, err = client.NewProxyDirectClient(policies, cnf, server.logger, certFile, keyFile)
	if err != nil {
		return ipPort, nil, nil, fmt.Errorf("Error setting up proxyDirectClient: %v", err)
	}
	info := map[string]interface{}{
		"version":                  common.Version,
		"strict_cors_mode":         true,
		"policies":                 policies.GetPolicyInfo(),
		"account_autocreate":       server.accountAutoCreate,
		"allow_account_management": true,
	}
	for k, v := range common.DEFAULT_CONSTRAINTS {
		info[k] = v
	}
	middleware.RegisterInfo("swift", info)
	ipPort = &srv.IpPort{Ip: bindIP, Port: int(bindPort), CertFile: certFile, KeyFile: keyFile}
	return ipPort, server, server.logger, nil
}
<|endoftext|>


<commit_before>package main

import (
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"os"
	"reflect"
	"strings"
	"testing"
	"time"
)

func Test_parseCGIHeaders(t *testing.T) {
	data := []struct {
		in      string
		out     string
		headers map[string]string
	}{
		{
			in:      "Some text",
			out:     "Some text",
			headers: map[string]string{},
		},
		{
			in:      "Location: url\n\nSome text",
			out:     "Some text",
			headers: map[string]string{"Location": "url"},
		},
		{
			in:      "Location: url\n\n",
			out:     "",
			headers: map[string]string{"Location": "url"},
		},
		{
			in:      "Location: url\nX-Name:  x-value\n\nSome text",
			out:     "Some text",
			headers: map[string]string{"Location": "url", "X-Name": "x-value"},
		},
	}

	for i, item := range data {
		out, headers := parseCGIHeaders(item.in)
		if !reflect.DeepEqual(item.headers, headers) || item.out != out {
			t.Errorf("%d:\nexpected: %s / %#v\nreal    : %s / %#v", i, item.out, item.headers, out, headers)
		}
	}
}

func Test_getShellAndParams(t *testing.T) {
	shell, params, err := getShellAndParams("ls", "sh", false)
	if shell != "sh" || !reflect.DeepEqual(params, []string{"-c", "ls"}) || err != nil {
		t.Errorf("1. getShellAndParams() failed")
	}

	shell, params, err = getShellAndParams("ls", "bash", false)
	if shell != "bash" || !reflect.DeepEqual(params, []string{"-c", "ls"}) || err != nil {
		t.Errorf("3. getShellAndParams() failed")
	}

	shell, params, err = getShellAndParams("ls -l -a", "", false)
	if shell != "ls" || !reflect.DeepEqual(params, []string{"-l", "-a"}) || err != nil {
		t.Errorf("4. getShellAndParams() failed")
	}

	shell, params, err = getShellAndParams("ls -l 'a b'", "", false)
	if shell != "ls" || !reflect.DeepEqual(params, []string{"-l", "a b"}) || err != nil {
		t.Errorf("5. getShellAndParams() failed")
	}

	_, _, err = getShellAndParams("ls '-l", "", false)
	if err == nil {
		t.Errorf("6. getShellAndParams() failed")
	}
}

func Test_getShellAndParams_windows(t *testing.T) {
	shell, params, err := getShellAndParams("ls", "sh", true)
	if shell != "cmd" || !reflect.DeepEqual(params, []string{"/C", "ls"}) || err != nil {
		t.Errorf("2. getShellAndParams() failed")
	}
}

func httpRequest(method string, url string, postData string) ([]byte, error) {
	var postDataReader io.Reader
	if method == "POST" && len(postData) > 0 {
		postDataReader = strings.NewReader(postData)
	}

	request, err := http.NewRequest(method, url, postDataReader)
	if err != nil {
		return nil, err
	}

	request.Header.Set("X-Real-Ip", "127.0.0.1")
	client := &http.Client{}
	res, err := client.Do(request)
	if err != nil {
		return nil, err
	}

	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return nil, err
	}

	err = res.Body.Close()
	if err != nil {
		return nil, err
	}

	return body, nil
}

func getFreePort(t *testing.T) string {
	listen, _ := net.Listen("tcp", ":0")
	parts := strings.Split(listen.Addr().String(), ":")
	err := listen.Close()
	if err != nil {
		t.Errorf("getFreePort() failed")
	}

	return parts[len(parts)-1]
}

func testHTTP(t *testing.T, method, url, postData string, fn func(body string) bool, message string) {
	res, err := httpRequest(method, url, postData)
	if err != nil {
		t.Errorf("%s, get %s failed: %s", message, url, err)
	}
	if !fn(string(res)) {
		t.Errorf("%s failed", message)
	}
}

func Test_main(t *testing.T) {
	port := getFreePort(t)
	os.Args = []string{"shell2http",
		"-add-exit",
		"-cache=1",
		"-cgi",
		// "-export-all-vars",
		"-export-vars=HOME",
		"-one-thread",
		"-shell=",
		"-log=/dev/null",
		"-port=" + port,
		"/echo", "echo 123",
		"/form", "echo var=$v_var",
		"/error", "/ not exists cmd",
		"/post", "cat",
		"/redirect", `echo "Location: /` + "\n" + `"`,
	}
	go main()
	time.Sleep(100 * time.Millisecond) // wait for up http server

	// hide stderr
	oldStderr := os.Stderr // keep backup of the real stderr
	newStderr, err := os.Open("/dev/null")
	if err != nil {
		t.Errorf("open /dev/null: %s", err)
	}
	os.Stderr = newStderr
	defer func() {
		os.Stderr = oldStderr
		err := newStderr.Close()
		if err != nil {
			t.Errorf("Stderr Close failed: %s", err)
		}
	}()

	testHTTP(t, "GET", "http://localhost:"+port+"/", "",
		func(res string) bool { return len(res) > 0 && strings.HasPrefix(res, "<!DOCTYPE html>") },
		"1. get /",
	)

	testHTTP(t, "GET", "http://localhost:"+port+"/echo", "",
		func(res string) bool { return res == "123\n" },
		"2. echo",
	)

	testHTTP(t, "GET", "http://localhost:"+port+"/echo", "",
		func(res string) bool { return res == "123\n" },
		"3. echo from cache",
	)

	testHTTP(t, "GET", "http://localhost:"+port+"/404", "",
		func(res string) bool { return strings.HasPrefix(res, "404 page not found") },
		"4. 404",
	)

	testHTTP(t, "GET", "http://localhost:"+port+"/error", "",
		func(res string) bool { return strings.HasPrefix(res, "exec error:") },
		"5. error",
	)

	testHTTP(t, "GET", "http://localhost:"+port+"/redirect", "",
		func(res string) bool { return strings.HasPrefix(res, "<!DOCTYPE html>") },
		"6. redirect",
	)

	testHTTP(t, "POST", "http://localhost:"+port+"/post", "X-header: value\n\ntext",
		func(res string) bool { return strings.HasPrefix(res, "text") },
		"7. POST",
	)
}
<commit_msg>Added test for errChain()<commit_after>package main

import (
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"os"
	"reflect"
	"strings"
	"testing"
	"time"
)

func Test_parseCGIHeaders(t *testing.T) {
	data := []struct {
		in      string
		out     string
		headers map[string]string
	}{
		{
			in:      "Some text",
			out:     "Some text",
			headers: map[string]string{},
		},
		{
			in:      "Location: url\n\nSome text",
			out:     "Some text",
			headers: map[string]string{"Location": "url"},
		},
		{
			in:      "Location: url\n\n",
			out:     "",
			headers: map[string]string{"Location": "url"},
		},
		{
			in:      "Location: url\nX-Name:  x-value\n\nSome text",
			out:     "Some text",
			headers: map[string]string{"Location": "url", "X-Name": "x-value"},
		},
	}

	for i, item := range data {
		out, headers := parseCGIHeaders(item.in)
		if !reflect.DeepEqual(item.headers, headers) || item.out != out {
			t.Errorf("%d:\nexpected: %s / %#v\nreal    : %s / %#v", i, item.out, item.headers, out, headers)
		}
	}
}

func Test_getShellAndParams(t *testing.T) {
	shell, params, err := getShellAndParams("ls", "sh", false)
	if shell != "sh" || !reflect.DeepEqual(params, []string{"-c", "ls"}) || err != nil {
		t.Errorf("1. getShellAndParams() failed")
	}

	shell, params, err = getShellAndParams("ls", "bash", false)
	if shell != "bash" || !reflect.DeepEqual(params, []string{"-c", "ls"}) || err != nil {
		t.Errorf("3. getShellAndParams() failed")
	}

	shell, params, err = getShellAndParams("ls -l -a", "", false)
	if shell != "ls" || !reflect.DeepEqual(params, []string{"-l", "-a"}) || err != nil {
		t.Errorf("4. getShellAndParams() failed")
	}

	shell, params, err = getShellAndParams("ls -l 'a b'", "", false)
	if shell != "ls" || !reflect.DeepEqual(params, []string{"-l", "a b"}) || err != nil {
		t.Errorf("5. getShellAndParams() failed")
	}

	_, _, err = getShellAndParams("ls '-l", "", false)
	if err == nil {
		t.Errorf("6. getShellAndParams() failed")
	}
}

func Test_getShellAndParams_windows(t *testing.T) {
	shell, params, err := getShellAndParams("ls", "sh", true)
	if shell != "cmd" || !reflect.DeepEqual(params, []string{"/C", "ls"}) || err != nil {
		t.Errorf("2. getShellAndParams() failed")
	}
}

func Test_errChain(t *testing.T) {
	err := errChain()
	if err != nil {
		t.Errorf("1. errChain() empty failed")
	}

	err = errChain(func() error { return nil })
	if err != nil {
		t.Errorf("2. errChain() failed")
	}

	err = errChain(func() error { return nil }, func() error { return nil })
	if err != nil {
		t.Errorf("3. errChain() failed")
	}

	err = errChain(func() error { return fmt.Errorf("error") })
	if err == nil {
		t.Errorf("4. errChain() failed")
	}

	err = errChain(func() error { return nil }, func() error { return fmt.Errorf("error") })
	if err == nil {
		t.Errorf("5. errChain() failed")
	}

	var1 := false
	err = errChain(func() error { return fmt.Errorf("error") }, func() error { var1 = true; return nil })
	if err == nil || var1 {
		t.Errorf("6. errChain() failed")
	}
}

func httpRequest(method string, url string, postData string) ([]byte, error) {
	var postDataReader io.Reader
	if method == "POST" && len(postData) > 0 {
		postDataReader = strings.NewReader(postData)
	}

	request, err := http.NewRequest(method, url, postDataReader)
	if err != nil {
		return nil, err
	}

	request.Header.Set("X-Real-Ip", "127.0.0.1")
	client := &http.Client{}
	res, err := client.Do(request)
	if err != nil {
		return nil, err
	}

	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return nil, err
	}

	err = res.Body.Close()
	if err != nil {
		return nil, err
	}

	return body, nil
}

func getFreePort(t *testing.T) string {
	listen, _ := net.Listen("tcp", ":0")
	parts := strings.Split(listen.Addr().String(), ":")
	err := listen.Close()
	if err != nil {
		t.Errorf("getFreePort() failed")
	}

	return parts[len(parts)-1]
}

func testHTTP(t *testing.T, method, url, postData string, fn func(body string) bool, message string) {
	res, err := httpRequest(method, url, postData)
	if err != nil {
		t.Errorf("%s, get %s failed: %s", message, url, err)
	}
	if !fn(string(res)) {
		t.Errorf("%s failed", message)
	}
}

func Test_main(t *testing.T) {
	port := getFreePort(t)
	os.Args = []string{"shell2http",
		"-add-exit",
		"-cache=1",
		"-cgi",
		// "-export-all-vars",
		"-export-vars=HOME",
		"-one-thread",
		"-shell=",
		"-log=/dev/null",
		"-port=" + port,
		"/echo", "echo 123",
		"/form", "echo var=$v_var",
		"/error", "/ not exists cmd",
		"/post", "cat",
		"/redirect", `echo "Location: /` + "\n" + `"`,
	}
	go main()
	time.Sleep(100 * time.Millisecond) // wait for up http server

	// hide stderr
	oldStderr := os.Stderr // keep backup of the real stderr
	newStderr, err := os.Open("/dev/null")
	if err != nil {
		t.Errorf("open /dev/null: %s", err)
	}
	os.Stderr = newStderr
	defer func() {
		os.Stderr = oldStderr
		err := newStderr.Close()
		if err != nil {
			t.Errorf("Stderr Close failed: %s", err)
		}
	}()

	testHTTP(t, "GET", "http://localhost:"+port+"/", "",
		func(res string) bool { return len(res) > 0 && strings.HasPrefix(res, "<!DOCTYPE html>") },
		"1. get /",
	)

	testHTTP(t, "GET", "http://localhost:"+port+"/echo", "",
		func(res string) bool { return res == "123\n" },
		"2. echo",
	)

	testHTTP(t, "GET", "http://localhost:"+port+"/echo", "",
		func(res string) bool { return res == "123\n" },
		"3. echo from cache",
	)

	testHTTP(t, "GET", "http://localhost:"+port+"/404", "",
		func(res string) bool { return strings.HasPrefix(res, "404 page not found") },
		"4. 404",
	)

	testHTTP(t, "GET", "http://localhost:"+port+"/error", "",
		func(res string) bool { return strings.HasPrefix(res, "exec error:") },
		"5. error",
	)

	testHTTP(t, "GET", "http://localhost:"+port+"/redirect", "",
		func(res string) bool { return strings.HasPrefix(res, "<!DOCTYPE html>") },
		"6. redirect",
	)

	testHTTP(t, "POST", "http://localhost:"+port+"/post", "X-header: value\n\ntext",
		func(res string) bool { return strings.HasPrefix(res, "text") },
		"7. POST",
	)
}
<|endoftext|>


<commit_before>// Ogo

package ogo

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

	"github.com/Odinman/ogo/utils"
	//"gopkg.in/redis.v3"
)

const (
	TD_TOTAL_FIELD = "_total_"
)

type TemporaryCounter struct {
	Name          string // 计数器名称
	Total         string // 汇总字段
	TotalAmount   float64
	Journal       string  //流水号
	JournalAmount float64 //金额
}

/* {{{ func CacheSet(key, value string, expire int) error
 *
 */
func CacheSet(key, value string, expire int) error {
	if cc := ClusterClient(); cc != nil {
		return cc.Set(key, value, time.Duration(expire)*time.Second).Err()
	} else {
		return fmt.Errorf("not found cache client")
	}
}

/* }}} */

/* {{{ func CacheIncrByFloat(key string, value float64) error
 *
 */
func CacheIncrByFloat(key string, value float64) error {
	if cc := ClusterClient(); cc != nil {
		return cc.IncrByFloat(key, value).Err()
	} else {
		return fmt.Errorf("incr %s failed", key)
	}
}

/* }}} */

/* {{{ func CacheGet(key) (string,error)
 *
 */
func CacheGet(key string) (string, error) {
	if cc := ClusterClient(); cc != nil {
		return cc.Get(key).Result()
	} else {
		return "", fmt.Errorf("not found cache client")
	}
}

/* }}} */

/* {{{ func GetLock(key string) (string, error)
 * 获取锁
 */
func GetLock(key string) (string, error) {
	if cc := ClusterClient(); cc != nil {
		cs := utils.NewUUID()
		ts := time.Now().Unix()
		val := fmt.Sprint(ts, ",", cs)
		var tried int
		for tried <= 3 {
			tried++
			if err := cc.SetNX(key, val, 600*time.Second).Err(); err != nil { //10分钟自动消失
				Debug("[GetLock][key: %s][val: %s]", key, val)
				// lock exist
				if old, err := cc.Get(key).Result(); err == nil {
					Debug("[GetLock][key: %s][old_val: %s]", key, old)
					vs := strings.SplitN(old, ",", 2)
					if ots, _ := strconv.Atoi(vs[0]); int(ts)-ots > 60 { //1分钟过期
						//过期了,抢
						if getSet := cc.GetSet(key, val); getSet.Err() == nil {
							gl := getSet.Val()
							if gl == old {
								//抢到了,返回checksum
								return cs, nil
							} else {
								//没抢到,并且还覆盖了人家抢到的锁,(可能会产生问题)
								//cc.Set(key, gl, 600*time.Second)
							}
						} else { //奇怪的情况
							Debug("[GetLock][key: %s][old_val: %s][ots: %d][strange_failed]", key, old, ots)
							return "", err
						}
					} else {
						Debug("[GetLock][key: %s][old_val: %s][ots: %d][not_expired_failed]", key, old, ots)
					}
				} else { //奇怪的情况
					return "", err
				}
			} else {
				Debug("[GetLock][key: %s][val: %s][not_exists_ok]", key, val)
				return cs, nil
			}
		}
		time.Sleep(50 * time.Millisecond) //50ms
	}
	return "", fmt.Errorf("can't get lock")
}

/* }}} */

/* {{{ func ReleaseLock(key string) error
 * 释放锁
 */
func ReleaseLock(key string) error {
	if cc := ClusterClient(); cc != nil {
		if cur, err := cc.Get(key).Result(); err == nil {
			//vs := strings.SplitN(cur, ",", 2)
			Debug("[lock_val: %s]", cur)
			return cc.Del(key).Err()
		}
	}
	return fmt.Errorf("release wrong")
}

/* }}} */

/* {{{ func (tc *TemporaryCounter) Incr(amount float64) error
 * 临时扣除(就是在TemporaryCounter增加一笔)
 */
func (tc *TemporaryCounter) Incr(amount float64) error {
	if cc := ClusterClient(); cc != nil {
		if tc.Name == "" || tc.Journal == "" || amount == 0 { //必须有名称以及流水号
			return fmt.Errorf("can't incr")
		}
		if tc.Total == "" {
			tc.Total = TD_TOTAL_FIELD
		}
		if cc.HExists(tc.Name, tc.Journal).Val() { // 流水号存在,说明已经处理
			return nil
		} else {
			if cc.HSet(tc.Name, tc.Journal, strconv.FormatFloat(amount, 'f', 6, 64)).Val() { // set 成功
				tc.JournalAmount = amount
				return cc.HIncrByFloat(tc.Name, tc.Total, amount).Err()
			}
		}
	} else {
		return fmt.Errorf("not found cache client")
	}
	return fmt.Errorf("incr failed")
}

/* }}} */

/* {{{ func (tc *TemporaryCounter) GetTotalAmount() (float64, error)
 * 获取临时扣除总额
 */
func (tc *TemporaryCounter) GetTotalAmount() (float64, error) {
	if cc := ClusterClient(); cc != nil {
		if tc.Name == "" { //必须有名称
			return 0, fmt.Errorf("can't get")
		}
		if tc.Total == "" {
			tc.Total = TD_TOTAL_FIELD
		}
		if tas, err := cc.HGet(tc.Name, tc.Total).Result(); err != nil {
			return 0, err
		} else {
			return strconv.ParseFloat(tas, 64)
		}
	} else {
		return 0, fmt.Errorf("not found cache client")
	}
}

/* }}} */

/* {{{ func (tc *TemporaryCounter) ClearJournal(journal string) error
 * 清除临时扣除
 */
func (tc *TemporaryCounter) ClearJournal(journal string) error {
	if cc := ClusterClient(); cc != nil {
		if tc.Name == "" || journal == "" { //必须有名称以及流水号
			return fmt.Errorf("can't clear")
		}
		if tc.Total == "" {
			tc.Total = TD_TOTAL_FIELD
		}
		if as, err := cc.HGet(tc.Name, journal).Result(); err != nil {
			return err
		} else {
			amount, _ := strconv.ParseFloat(as, 64)
			if err := cc.HDel(tc.Name, journal).Err(); err != nil {
				return err
			} else {
				return cc.HIncrByFloat(tc.Name, tc.Total, 0-amount).Err()
			}
		}
	} else {
		return fmt.Errorf("not found cache client")
	}
	return fmt.Errorf("deduct failed")
}

/* }}} */
<commit_msg>update<commit_after>// Ogo

package ogo

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

	"github.com/Odinman/ogo/utils"
	//"gopkg.in/redis.v3"
)

const (
	TD_TOTAL_FIELD = "_total_"
)

type TemporaryCounter struct {
	Name          string // 计数器名称
	Total         string // 汇总字段
	TotalAmount   float64
	Journal       string  //流水号
	JournalAmount float64 //金额
}

/* {{{ func CacheSet(key, value string, expire int) error
 *
 */
func CacheSet(key, value string, expire int) error {
	if cc := ClusterClient(); cc != nil {
		return cc.Set(key, value, time.Duration(expire)*time.Second).Err()
	} else {
		return fmt.Errorf("not found cache client")
	}
}

/* }}} */

/* {{{ func CacheIncrByFloat(key string, value float64) error
 *
 */
func CacheIncrByFloat(key string, value float64) error {
	if cc := ClusterClient(); cc != nil {
		return cc.IncrByFloat(key, value).Err()
	} else {
		return fmt.Errorf("incr %s failed", key)
	}
}

/* }}} */

/* {{{ func CacheGet(key) (string,error)
 *
 */
func CacheGet(key string) (string, error) {
	if cc := ClusterClient(); cc != nil {
		return cc.Get(key).Result()
	} else {
		return "", fmt.Errorf("not found cache client")
	}
}

/* }}} */

/* {{{ func GetLock(key string) (string, error)
 * 获取锁
 */
func GetLock(key string) (string, error) {
	if cc := ClusterClient(); cc != nil {
		cs := utils.NewUUID()
		ts := time.Now().Unix()
		val := fmt.Sprint(ts, ",", cs)
		var tried int
		for tried <= 3 {
			tried++
			//if err := cc.SetNX(key, val, 600*time.Second).Err(); err != nil { //10分钟自动消失
			if err := cc.SetNX(key, val).Err(); err != nil { //10分钟自动消失
				Debug("[GetLock][key: %s][val: %s]", key, val)
				// lock exist
				if old, err := cc.Get(key).Result(); err == nil {
					Debug("[GetLock][key: %s][old_val: %s]", key, old)
					vs := strings.SplitN(old, ",", 2)
					if ots, _ := strconv.Atoi(vs[0]); int(ts)-ots > 60 { //1分钟过期
						//过期了,抢
						if getSet := cc.GetSet(key, val); getSet.Err() == nil {
							gl := getSet.Val()
							if gl == old {
								//抢到了,返回checksum
								return cs, nil
							} else {
								//没抢到,并且还覆盖了人家抢到的锁,(可能会产生问题)
								//cc.Set(key, gl, 600*time.Second)
							}
						} else { //奇怪的情况
							Debug("[GetLock][key: %s][old_val: %s][ots: %d][strange_failed]", key, old, ots)
							return "", err
						}
					} else {
						Debug("[GetLock][key: %s][old_val: %s][ots: %d][not_expired_failed]", key, old, ots)
					}
				} else { //奇怪的情况
					return "", err
				}
			} else {
				Debug("[GetLock][key: %s][val: %s][not_exists_ok]", key, val)
				return cs, nil
			}
		}
		time.Sleep(50 * time.Millisecond) //50ms
	}
	return "", fmt.Errorf("can't get lock")
}

/* }}} */

/* {{{ func ReleaseLock(key string) error
 * 释放锁
 */
func ReleaseLock(key string) error {
	if cc := ClusterClient(); cc != nil {
		if cur, err := cc.Get(key).Result(); err == nil {
			//vs := strings.SplitN(cur, ",", 2)
			Debug("[lock_val: %s]", cur)
			return cc.Del(key).Err()
		}
	}
	return fmt.Errorf("release wrong")
}

/* }}} */

/* {{{ func (tc *TemporaryCounter) Incr(amount float64) error
 * 临时扣除(就是在TemporaryCounter增加一笔)
 */
func (tc *TemporaryCounter) Incr(amount float64) error {
	if cc := ClusterClient(); cc != nil {
		if tc.Name == "" || tc.Journal == "" || amount == 0 { //必须有名称以及流水号
			return fmt.Errorf("can't incr")
		}
		if tc.Total == "" {
			tc.Total = TD_TOTAL_FIELD
		}
		if cc.HExists(tc.Name, tc.Journal).Val() { // 流水号存在,说明已经处理
			return nil
		} else {
			if cc.HSet(tc.Name, tc.Journal, strconv.FormatFloat(amount, 'f', 6, 64)).Val() { // set 成功
				tc.JournalAmount = amount
				return cc.HIncrByFloat(tc.Name, tc.Total, amount).Err()
			}
		}
	} else {
		return fmt.Errorf("not found cache client")
	}
	return fmt.Errorf("incr failed")
}

/* }}} */

/* {{{ func (tc *TemporaryCounter) GetTotalAmount() (float64, error)
 * 获取临时扣除总额
 */
func (tc *TemporaryCounter) GetTotalAmount() (float64, error) {
	if cc := ClusterClient(); cc != nil {
		if tc.Name == "" { //必须有名称
			return 0, fmt.Errorf("can't get")
		}
		if tc.Total == "" {
			tc.Total = TD_TOTAL_FIELD
		}
		if tas, err := cc.HGet(tc.Name, tc.Total).Result(); err != nil {
			return 0, err
		} else {
			return strconv.ParseFloat(tas, 64)
		}
	} else {
		return 0, fmt.Errorf("not found cache client")
	}
}

/* }}} */

/* {{{ func (tc *TemporaryCounter) ClearJournal(journal string) error
 * 清除临时扣除
 */
func (tc *TemporaryCounter) ClearJournal(journal string) error {
	if cc := ClusterClient(); cc != nil {
		if tc.Name == "" || journal == "" { //必须有名称以及流水号
			return fmt.Errorf("can't clear")
		}
		if tc.Total == "" {
			tc.Total = TD_TOTAL_FIELD
		}
		if as, err := cc.HGet(tc.Name, journal).Result(); err != nil {
			return err
		} else {
			amount, _ := strconv.ParseFloat(as, 64)
			if err := cc.HDel(tc.Name, journal).Err(); err != nil {
				return err
			} else {
				return cc.HIncrByFloat(tc.Name, tc.Total, 0-amount).Err()
			}
		}
	} else {
		return fmt.Errorf("not found cache client")
	}
	return fmt.Errorf("deduct failed")
}

/* }}} */
<|endoftext|>


<commit_before>package cmd

import (
	"fmt"
	"github.com/imwithye/gitignore/template"
	_ "github.com/imwithye/gitignore/template/all"
	"github.com/spf13/cobra"
	"io/ioutil"
	"strings"
)

var addCmdQuite, addCmdVerbose bool
var addCmd = &cobra.Command{
	Use:   "add [templates]",
	Short: "Add multi git gitignore templates to .gitignore",
	Run: func(cmd *cobra.Command, args []string) {
		if len(args) < 1 {
			cmd.Help()
			return
		}
		ignores := []string{}
		for _, arg := range args {
			t, err := template.GetTemplate(arg)
			if err != nil {
				if addCmdQuite {
					continue
				}
				fmt.Println("*", err)
				continue
			}
			if !addCmdQuite && addCmdVerbose {
				fmt.Println("-", t.Path)
			}
			ignores = append(ignores, t.Ignore)
		}
		if len(ignores) > 0 {
			ioutil.WriteFile(".gitignore", []byte(strings.Join(ignores, "\n")), 0644)
		} else {
			fmt.Println("*", "nothing will be added")
		}
	},
}

func init() {
	RootCmd.AddCommand(addCmd)
	addCmd.Flags().BoolVarP(&addCmdQuite, "quite", "q", false, "quite mode")
	addCmd.Flags().BoolVarP(&addCmdVerbose, "verbose", "v", false, "verbose mode")
}
<commit_msg>add line sep between different ignores<commit_after>package cmd

import (
	"fmt"
	"github.com/imwithye/gitignore/template"
	_ "github.com/imwithye/gitignore/template/all"
	"github.com/spf13/cobra"
	"io/ioutil"
	"strings"
)

var addCmdQuite, addCmdVerbose bool
var addCmd = &cobra.Command{
	Use:   "add [templates]",
	Short: "Add multi git gitignore templates to .gitignore",
	Run: func(cmd *cobra.Command, args []string) {
		if len(args) < 1 {
			cmd.Help()
			return
		}
		ignores := []string{}
		for _, arg := range args {
			t, err := template.GetTemplate(arg)
			if err != nil {
				if addCmdQuite {
					continue
				}
				fmt.Println("*", err)
				continue
			}
			if !addCmdQuite && addCmdVerbose {
				fmt.Println("-", t.Path)
			}
			ignores = append(ignores, t.Ignore)
		}
		if len(ignores) > 0 {
			ioutil.WriteFile(".gitignore", []byte(strings.Join(ignores, "\n\n")), 0644)
		} else {
			fmt.Println("*", "nothing will be added")
		}
	},
}

func init() {
	RootCmd.AddCommand(addCmd)
	addCmd.Flags().BoolVarP(&addCmdQuite, "quite", "q", false, "quite mode")
	addCmd.Flags().BoolVarP(&addCmdVerbose, "verbose", "v", false, "verbose mode")
}
<|endoftext|>


<commit_before>package cmd

import (
	"fmt"
	"os"
	"os/signal"
	"syscall"

	ccli "github.com/micro/cli"
	"github.com/micro/go-micro"
	"github.com/micro/go-micro/config/cmd"
	"github.com/micro/go-micro/runtime"
	"github.com/micro/go-micro/util/log"
	"github.com/micro/micro/api"
	"github.com/micro/micro/bot"
	"github.com/micro/micro/broker"
	"github.com/micro/micro/cli"
	"github.com/micro/micro/health"
	"github.com/micro/micro/monitor"
	"github.com/micro/micro/network"
	"github.com/micro/micro/new"
	"github.com/micro/micro/plugin"
	"github.com/micro/micro/plugin/build"
	"github.com/micro/micro/proxy"
	"github.com/micro/micro/registry"
	"github.com/micro/micro/router"
	"github.com/micro/micro/server"
	"github.com/micro/micro/service"
	"github.com/micro/micro/tunnel"
	"github.com/micro/micro/web"

	// include usage
	_ "github.com/micro/micro/internal/usage"
)

var (
	name        = "micro"
	description = "A microservice runtime"
	version     = "1.11.0"
)

func init() {
	plugin.Register(build.Flags())
}

func setup(app *ccli.App) {
	app.Flags = append(app.Flags,
		ccli.BoolFlag{
			Name:   "enable_acme",
			Usage:  "Enables ACME support via Let's Encrypt. ACME hosts should also be specified.",
			EnvVar: "MICRO_ENABLE_ACME",
		},
		ccli.StringFlag{
			Name:   "acme_hosts",
			Usage:  "Comma separated list of hostnames to manage ACME certs for",
			EnvVar: "MICRO_ACME_HOSTS",
		},
		ccli.StringFlag{
			Name:   "acme_provider",
			Usage:  "The provider that will be used to get communicate with Let's Encrypt. Valid options: autocert",
			EnvVar: "MICRO_ACME_PROVIDER",
		},
		ccli.BoolFlag{
			Name:   "enable_tls",
			Usage:  "Enable TLS support. Expects cert and key file to be specified",
			EnvVar: "MICRO_ENABLE_TLS",
		},
		ccli.StringFlag{
			Name:   "tls_cert_file",
			Usage:  "Path to the TLS Certificate file",
			EnvVar: "MICRO_TLS_CERT_FILE",
		},
		ccli.StringFlag{
			Name:   "tls_key_file",
			Usage:  "Path to the TLS Key file",
			EnvVar: "MICRO_TLS_KEY_FILE",
		},
		ccli.StringFlag{
			Name:   "tls_client_ca_file",
			Usage:  "Path to the TLS CA file to verify clients against",
			EnvVar: "MICRO_TLS_CLIENT_CA_FILE",
		},
		ccli.StringFlag{
			Name:   "api_address",
			Usage:  "Set the api address e.g 0.0.0.0:8080",
			EnvVar: "MICRO_API_ADDRESS",
		},
		ccli.StringFlag{
			Name:   "proxy_address",
			Usage:  "Proxy requests via the HTTP address specified",
			EnvVar: "MICRO_PROXY_ADDRESS",
		},
		ccli.StringFlag{
			Name:   "web_address",
			Usage:  "Set the web UI address e.g 0.0.0.0:8082",
			EnvVar: "MICRO_WEB_ADDRESS",
		},
		ccli.StringFlag{
			Name:   "network",
			Usage:  "Set the micro network name: local, go.micro",
			EnvVar: "MICRO_NETWORK",
		},
		ccli.StringFlag{
			Name:   "network_address",
			Usage:  "Set the micro network address e.g. :9093",
			EnvVar: "MICRO_NETWORK_ADDRESS",
		},
		ccli.StringFlag{
			Name:   "router_address",
			Usage:  "Set the micro router address e.g. :8084",
			EnvVar: "MICRO_ROUTER_ADDRESS",
		},
		ccli.StringFlag{
			Name:   "gateway_address",
			Usage:  "Set the micro default gateway address e.g. :9094",
			EnvVar: "MICRO_GATEWAY_ADDRESS",
		},
		ccli.StringFlag{
			Name:   "tunnel_address",
			Usage:  "Set the micro tunnel address e.g. :8083",
			EnvVar: "MICRO_TUNNEL_ADDRESS",
		},
		ccli.StringFlag{
			Name:   "api_handler",
			Usage:  "Specify the request handler to be used for mapping HTTP requests to services; {api, proxy, rpc}",
			EnvVar: "MICRO_API_HANDLER",
		},
		ccli.StringFlag{
			Name:   "api_namespace",
			Usage:  "Set the namespace used by the API e.g. com.example.api",
			EnvVar: "MICRO_API_NAMESPACE",
		},
		ccli.StringFlag{
			Name:   "web_namespace",
			Usage:  "Set the namespace used by the Web proxy e.g. com.example.web",
			EnvVar: "MICRO_WEB_NAMESPACE",
		},
		ccli.BoolFlag{
			Name:   "enable_stats",
			Usage:  "Enable stats",
			EnvVar: "MICRO_ENABLE_STATS",
		},
		ccli.BoolTFlag{
			Name:   "report_usage",
			Usage:  "Report usage statistics",
			EnvVar: "MICRO_REPORT_USAGE",
		},
	)

	plugins := plugin.Plugins()

	for _, p := range plugins {
		if flags := p.Flags(); len(flags) > 0 {
			app.Flags = append(app.Flags, flags...)
		}

		if cmds := p.Commands(); len(cmds) > 0 {
			app.Commands = append(app.Commands, cmds...)
		}
	}

	before := app.Before

	app.Before = func(ctx *ccli.Context) error {
		if len(ctx.String("api_handler")) > 0 {
			api.Handler = ctx.String("api_handler")
		}
		if len(ctx.String("api_address")) > 0 {
			api.Address = ctx.String("api_address")
		}
		if len(ctx.String("proxy_address")) > 0 {
			proxy.Address = ctx.String("proxy_address")
		}
		if len(ctx.String("web_address")) > 0 {
			web.Address = ctx.String("web_address")
		}
		if len(ctx.String("network_address")) > 0 {
			server.Network = ctx.String("network_address")
		}
		if len(ctx.String("router_address")) > 0 {
			router.Router = ctx.String("router_address")
		}
		if len(ctx.String("tunnel_address")) > 0 {
			tunnel.Address = ctx.String("tunnel_address")
		}
		if len(ctx.String("api_namespace")) > 0 {
			api.Namespace = ctx.String("api_namespace")
		}
		if len(ctx.String("web_namespace")) > 0 {
			web.Namespace = ctx.String("web_namespace")
		}

		for _, p := range plugins {
			if err := p.Init(ctx); err != nil {
				return err
			}
		}

		// now do previous before
		return before(ctx)
	}
}

// Init initialised the command line
func Init(options ...micro.Option) {
	Setup(cmd.App(), options...)

	cmd.Init(
		cmd.Name(name),
		cmd.Description(description),
		cmd.Version(version),
	)
}

// Setup sets up a cli.App
func Setup(app *ccli.App, options ...micro.Option) {
	// Add the various commands
	app.Commands = append(app.Commands, api.Commands(options...)...)
	app.Commands = append(app.Commands, bot.Commands()...)
	app.Commands = append(app.Commands, cli.Commands()...)
	app.Commands = append(app.Commands, broker.Commands(options...)...)
	app.Commands = append(app.Commands, health.Commands(options...)...)
	app.Commands = append(app.Commands, proxy.Commands(options...)...)
	app.Commands = append(app.Commands, monitor.Commands(options...)...)
	app.Commands = append(app.Commands, router.Commands(options...)...)
	app.Commands = append(app.Commands, tunnel.Commands(options...)...)
	app.Commands = append(app.Commands, network.Commands(options...)...)
	app.Commands = append(app.Commands, registry.Commands(options...)...)
	app.Commands = append(app.Commands, server.Commands(options...)...)
	app.Commands = append(app.Commands, service.Commands(options...)...)
	app.Commands = append(app.Commands, new.Commands()...)
	app.Commands = append(app.Commands, build.Commands()...)
	app.Commands = append(app.Commands, web.Commands(options...)...)

	// boot micro
	app.Action = func(context *ccli.Context) {
		log.Name("micro")

		// get the network flag
		network := context.GlobalString("network")

		// pass through the environment
		// TODO: perhaps don't do this
		env := os.Environ()

		switch network {
		case "local":
			// no op for now
			log.Info("Setting local network")
		default:
			log.Info("Setting global network")
			// set the resolver to use https://micro.mu/network
			env = append(env, "MICRO_NETWORK_RESOLVER=http")
		}

		log.Info("Loading core services")

		services := []string{
			"registry", // :8000
			"broker",   // :8001
			"tunnel",   // :8083
			"router",   // :8084
			"network",  // :8085
			"proxy",    // :8081
			"monitor",  // :????
			"api",      // :8080
			"web",      // :8082
			"bot",      // :????
		}

		for _, service := range services {
			name := fmt.Sprintf("go.micro.%s", service)
			log.Infof("Registering %s\n", name)

			args := []runtime.CreateOption{
				runtime.WithCommand(os.Args[0], service),
				runtime.WithEnv(env),
				runtime.WithOutput(os.Stdout),
			}

			// register the service
			runtime.Create(&runtime.Service{
				Name: name,
			}, args...)
		}

		shutdown := make(chan os.Signal, 1)
		signal.Notify(shutdown, syscall.SIGTERM, syscall.SIGINT, syscall.SIGQUIT)

		log.Info("Starting service runtime")

		// start the runtime
		if err := runtime.Start(); err != nil {
			log.Fatal(err)
		}

		log.Info("Service runtime started")

		// TODO: should we launch the console?
		// start the console
		// cli.Init(context)

		select {
		case <-shutdown:
			log.Info("Shutdown signal received")
			log.Info("Stopping service runtime")
		}

		// stop all the things
		if err := runtime.Stop(); err != nil {
			log.Fatal(err)
		}

		log.Info("Service runtime shutdown")

		// exit success
		os.Exit(0)

	}

	setup(app)
}
<commit_msg>Fix typo in acme_provider flag<commit_after>package cmd

import (
	"fmt"
	"os"
	"os/signal"
	"syscall"

	ccli "github.com/micro/cli"
	"github.com/micro/go-micro"
	"github.com/micro/go-micro/config/cmd"
	"github.com/micro/go-micro/runtime"
	"github.com/micro/go-micro/util/log"
	"github.com/micro/micro/api"
	"github.com/micro/micro/bot"
	"github.com/micro/micro/broker"
	"github.com/micro/micro/cli"
	"github.com/micro/micro/health"
	"github.com/micro/micro/monitor"
	"github.com/micro/micro/network"
	"github.com/micro/micro/new"
	"github.com/micro/micro/plugin"
	"github.com/micro/micro/plugin/build"
	"github.com/micro/micro/proxy"
	"github.com/micro/micro/registry"
	"github.com/micro/micro/router"
	"github.com/micro/micro/server"
	"github.com/micro/micro/service"
	"github.com/micro/micro/tunnel"
	"github.com/micro/micro/web"

	// include usage
	_ "github.com/micro/micro/internal/usage"
)

var (
	name        = "micro"
	description = "A microservice runtime"
	version     = "1.11.0"
)

func init() {
	plugin.Register(build.Flags())
}

func setup(app *ccli.App) {
	app.Flags = append(app.Flags,
		ccli.BoolFlag{
			Name:   "enable_acme",
			Usage:  "Enables ACME support via Let's Encrypt. ACME hosts should also be specified.",
			EnvVar: "MICRO_ENABLE_ACME",
		},
		ccli.StringFlag{
			Name:   "acme_hosts",
			Usage:  "Comma separated list of hostnames to manage ACME certs for",
			EnvVar: "MICRO_ACME_HOSTS",
		},
		ccli.StringFlag{
			Name:   "acme_provider",
			Usage:  "The provider that will be used to communicate with Let's Encrypt. Valid options: autocert",
			EnvVar: "MICRO_ACME_PROVIDER",
		},
		ccli.BoolFlag{
			Name:   "enable_tls",
			Usage:  "Enable TLS support. Expects cert and key file to be specified",
			EnvVar: "MICRO_ENABLE_TLS",
		},
		ccli.StringFlag{
			Name:   "tls_cert_file",
			Usage:  "Path to the TLS Certificate file",
			EnvVar: "MICRO_TLS_CERT_FILE",
		},
		ccli.StringFlag{
			Name:   "tls_key_file",
			Usage:  "Path to the TLS Key file",
			EnvVar: "MICRO_TLS_KEY_FILE",
		},
		ccli.StringFlag{
			Name:   "tls_client_ca_file",
			Usage:  "Path to the TLS CA file to verify clients against",
			EnvVar: "MICRO_TLS_CLIENT_CA_FILE",
		},
		ccli.StringFlag{
			Name:   "api_address",
			Usage:  "Set the api address e.g 0.0.0.0:8080",
			EnvVar: "MICRO_API_ADDRESS",
		},
		ccli.StringFlag{
			Name:   "proxy_address",
			Usage:  "Proxy requests via the HTTP address specified",
			EnvVar: "MICRO_PROXY_ADDRESS",
		},
		ccli.StringFlag{
			Name:   "web_address",
			Usage:  "Set the web UI address e.g 0.0.0.0:8082",
			EnvVar: "MICRO_WEB_ADDRESS",
		},
		ccli.StringFlag{
			Name:   "network",
			Usage:  "Set the micro network name: local, go.micro",
			EnvVar: "MICRO_NETWORK",
		},
		ccli.StringFlag{
			Name:   "network_address",
			Usage:  "Set the micro network address e.g. :9093",
			EnvVar: "MICRO_NETWORK_ADDRESS",
		},
		ccli.StringFlag{
			Name:   "router_address",
			Usage:  "Set the micro router address e.g. :8084",
			EnvVar: "MICRO_ROUTER_ADDRESS",
		},
		ccli.StringFlag{
			Name:   "gateway_address",
			Usage:  "Set the micro default gateway address e.g. :9094",
			EnvVar: "MICRO_GATEWAY_ADDRESS",
		},
		ccli.StringFlag{
			Name:   "tunnel_address",
			Usage:  "Set the micro tunnel address e.g. :8083",
			EnvVar: "MICRO_TUNNEL_ADDRESS",
		},
		ccli.StringFlag{
			Name:   "api_handler",
			Usage:  "Specify the request handler to be used for mapping HTTP requests to services; {api, proxy, rpc}",
			EnvVar: "MICRO_API_HANDLER",
		},
		ccli.StringFlag{
			Name:   "api_namespace",
			Usage:  "Set the namespace used by the API e.g. com.example.api",
			EnvVar: "MICRO_API_NAMESPACE",
		},
		ccli.StringFlag{
			Name:   "web_namespace",
			Usage:  "Set the namespace used by the Web proxy e.g. com.example.web",
			EnvVar: "MICRO_WEB_NAMESPACE",
		},
		ccli.BoolFlag{
			Name:   "enable_stats",
			Usage:  "Enable stats",
			EnvVar: "MICRO_ENABLE_STATS",
		},
		ccli.BoolTFlag{
			Name:   "report_usage",
			Usage:  "Report usage statistics",
			EnvVar: "MICRO_REPORT_USAGE",
		},
	)

	plugins := plugin.Plugins()

	for _, p := range plugins {
		if flags := p.Flags(); len(flags) > 0 {
			app.Flags = append(app.Flags, flags...)
		}

		if cmds := p.Commands(); len(cmds) > 0 {
			app.Commands = append(app.Commands, cmds...)
		}
	}

	before := app.Before

	app.Before = func(ctx *ccli.Context) error {
		if len(ctx.String("api_handler")) > 0 {
			api.Handler = ctx.String("api_handler")
		}
		if len(ctx.String("api_address")) > 0 {
			api.Address = ctx.String("api_address")
		}
		if len(ctx.String("proxy_address")) > 0 {
			proxy.Address = ctx.String("proxy_address")
		}
		if len(ctx.String("web_address")) > 0 {
			web.Address = ctx.String("web_address")
		}
		if len(ctx.String("network_address")) > 0 {
			server.Network = ctx.String("network_address")
		}
		if len(ctx.String("router_address")) > 0 {
			router.Router = ctx.String("router_address")
		}
		if len(ctx.String("tunnel_address")) > 0 {
			tunnel.Address = ctx.String("tunnel_address")
		}
		if len(ctx.String("api_namespace")) > 0 {
			api.Namespace = ctx.String("api_namespace")
		}
		if len(ctx.String("web_namespace")) > 0 {
			web.Namespace = ctx.String("web_namespace")
		}

		for _, p := range plugins {
			if err := p.Init(ctx); err != nil {
				return err
			}
		}

		// now do previous before
		return before(ctx)
	}
}

// Init initialised the command line
func Init(options ...micro.Option) {
	Setup(cmd.App(), options...)

	cmd.Init(
		cmd.Name(name),
		cmd.Description(description),
		cmd.Version(version),
	)
}

// Setup sets up a cli.App
func Setup(app *ccli.App, options ...micro.Option) {
	// Add the various commands
	app.Commands = append(app.Commands, api.Commands(options...)...)
	app.Commands = append(app.Commands, bot.Commands()...)
	app.Commands = append(app.Commands, cli.Commands()...)
	app.Commands = append(app.Commands, broker.Commands(options...)...)
	app.Commands = append(app.Commands, health.Commands(options...)...)
	app.Commands = append(app.Commands, proxy.Commands(options...)...)
	app.Commands = append(app.Commands, monitor.Commands(options...)...)
	app.Commands = append(app.Commands, router.Commands(options...)...)
	app.Commands = append(app.Commands, tunnel.Commands(options...)...)
	app.Commands = append(app.Commands, network.Commands(options...)...)
	app.Commands = append(app.Commands, registry.Commands(options...)...)
	app.Commands = append(app.Commands, server.Commands(options...)...)
	app.Commands = append(app.Commands, service.Commands(options...)...)
	app.Commands = append(app.Commands, new.Commands()...)
	app.Commands = append(app.Commands, build.Commands()...)
	app.Commands = append(app.Commands, web.Commands(options...)...)

	// boot micro
	app.Action = func(context *ccli.Context) {
		log.Name("micro")

		// get the network flag
		network := context.GlobalString("network")

		// pass through the environment
		// TODO: perhaps don't do this
		env := os.Environ()

		switch network {
		case "local":
			// no op for now
			log.Info("Setting local network")
		default:
			log.Info("Setting global network")
			// set the resolver to use https://micro.mu/network
			env = append(env, "MICRO_NETWORK_RESOLVER=http")
		}

		log.Info("Loading core services")

		services := []string{
			"registry", // :8000
			"broker",   // :8001
			"tunnel",   // :8083
			"router",   // :8084
			"network",  // :8085
			"proxy",    // :8081
			"monitor",  // :????
			"api",      // :8080
			"web",      // :8082
			"bot",      // :????
		}

		for _, service := range services {
			name := fmt.Sprintf("go.micro.%s", service)
			log.Infof("Registering %s\n", name)

			args := []runtime.CreateOption{
				runtime.WithCommand(os.Args[0], service),
				runtime.WithEnv(env),
				runtime.WithOutput(os.Stdout),
			}

			// register the service
			runtime.Create(&runtime.Service{
				Name: name,
			}, args...)
		}

		shutdown := make(chan os.Signal, 1)
		signal.Notify(shutdown, syscall.SIGTERM, syscall.SIGINT, syscall.SIGQUIT)

		log.Info("Starting service runtime")

		// start the runtime
		if err := runtime.Start(); err != nil {
			log.Fatal(err)
		}

		log.Info("Service runtime started")

		// TODO: should we launch the console?
		// start the console
		// cli.Init(context)

		select {
		case <-shutdown:
			log.Info("Shutdown signal received")
			log.Info("Stopping service runtime")
		}

		// stop all the things
		if err := runtime.Stop(); err != nil {
			log.Fatal(err)
		}

		log.Info("Service runtime shutdown")

		// exit success
		os.Exit(0)

	}

	setup(app)
}
<|endoftext|>


<commit_before>/*
Copyright 2016 The Kubernetes Authors.

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

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

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

package cloudup

import (
	"fmt"

	"github.com/blang/semver"
	"github.com/golang/glog"
	"k8s.io/kops/pkg/apis/kops"
	"k8s.io/kops/pkg/apis/kops/util"
	"k8s.io/kops/pkg/apis/kops/validation"
	"k8s.io/kops/upup/pkg/fi"
	"k8s.io/kops/upup/pkg/fi/cloudup/awsup"
	"k8s.io/kops/util/pkg/reflectutils"
)

// Default Machine types for various types of instance group machine
const (
	defaultNodeMachineTypeGCE     = "n1-standard-2"
	defaultNodeMachineTypeVSphere = "vsphere_node"
	defaultNodeMachineTypeDO      = "s-2vcpu-4gb"

	defaultBastionMachineTypeGCE     = "f1-micro"
	defaultBastionMachineTypeVSphere = "vsphere_bastion"

	defaultMasterMachineTypeGCE     = "n1-standard-1"
	defaultMasterMachineTypeVSphere = "vsphere_master"
	defaultMasterMachineTypeDO      = "s-2vcpu-2gb"

	defaultVSphereNodeImage = "kops_ubuntu_16_04.ova"
	defaultDONodeImage      = "coreos-stable"
)

var awsDedicatedInstanceExceptions = map[string]bool{
	"t2.nano":   true,
	"t2.micro":  true,
	"t2.small":  true,
	"t2.medium": true,
	"t2.large":  true,
	"t2.xlarge": true,
}

// PopulateInstanceGroupSpec sets default values in the InstanceGroup
// The InstanceGroup is simpler than the cluster spec, so we just populate in place (like the rest of k8s)
func PopulateInstanceGroupSpec(cluster *kops.Cluster, input *kops.InstanceGroup, channel *kops.Channel) (*kops.InstanceGroup, error) {
	err := validation.ValidateInstanceGroup(input)
	if err != nil {
		return nil, err
	}

	ig := &kops.InstanceGroup{}
	reflectutils.JsonMergeStruct(ig, input)

	// TODO: Clean up
	if ig.IsMaster() {
		if ig.Spec.MachineType == "" {
			ig.Spec.MachineType, err = defaultMachineType(cluster, ig)
			if err != nil {
				return nil, fmt.Errorf("error assigning default machine type for masters: %v", err)
			}

		}
		if ig.Spec.MinSize == nil {
			ig.Spec.MinSize = fi.Int32(1)
		}
		if ig.Spec.MaxSize == nil {
			ig.Spec.MaxSize = fi.Int32(1)
		}
	} else if ig.Spec.Role == kops.InstanceGroupRoleBastion {
		if ig.Spec.MachineType == "" {
			ig.Spec.MachineType, err = defaultMachineType(cluster, ig)
			if err != nil {
				return nil, fmt.Errorf("error assigning default machine type for bastions: %v", err)
			}
		}
		if ig.Spec.MinSize == nil {
			ig.Spec.MinSize = fi.Int32(1)
		}
		if ig.Spec.MaxSize == nil {
			ig.Spec.MaxSize = fi.Int32(1)
		}
	} else {
		if ig.Spec.MachineType == "" {
			ig.Spec.MachineType, err = defaultMachineType(cluster, ig)
			if err != nil {
				return nil, fmt.Errorf("error assigning default machine type for nodes: %v", err)
			}
		}
		if ig.Spec.MinSize == nil {
			ig.Spec.MinSize = fi.Int32(2)
		}
		if ig.Spec.MaxSize == nil {
			ig.Spec.MaxSize = fi.Int32(2)
		}
	}

	if ig.Spec.Image == "" {
		ig.Spec.Image = defaultImage(cluster, channel)
	}

	if ig.Spec.Tenancy != "" && ig.Spec.Tenancy != "default" {
		switch kops.CloudProviderID(cluster.Spec.CloudProvider) {
		case kops.CloudProviderAWS:
			if _, ok := awsDedicatedInstanceExceptions[ig.Spec.MachineType]; ok {
				return nil, fmt.Errorf("Invalid dedicated instance type: %s", ig.Spec.MachineType)
			}
		default:
			glog.Warning("Trying to set tenancy on non-AWS environment")
		}
	}

	if ig.IsMaster() {
		if len(ig.Spec.Subnets) == 0 {
			return nil, fmt.Errorf("Master InstanceGroup %s did not specify any Subnets", ig.ObjectMeta.Name)
		}
	} else if ig.Spec.Role == kops.InstanceGroupRoleBastion {
		if len(ig.Spec.Subnets) == 0 {
			for _, subnet := range cluster.Spec.Subnets {
				if subnet.Type == kops.SubnetTypeUtility {
					ig.Spec.Subnets = append(ig.Spec.Subnets, subnet.Name)
				}
			}
		}
	} else {
		if len(ig.Spec.Subnets) == 0 {
			for _, subnet := range cluster.Spec.Subnets {
				if subnet.Type != kops.SubnetTypeUtility {
					ig.Spec.Subnets = append(ig.Spec.Subnets, subnet.Name)
				}
			}
		}
	}

	if len(ig.Spec.Subnets) == 0 {
		return nil, fmt.Errorf("unable to infer any Subnets for InstanceGroup %s ", ig.ObjectMeta.Name)
	}

	return ig, nil
}

// defaultMachineType returns the default MachineType for the instance group, based on the cloudprovider
func defaultMachineType(cluster *kops.Cluster, ig *kops.InstanceGroup) (string, error) {
	switch kops.CloudProviderID(cluster.Spec.CloudProvider) {
	case kops.CloudProviderAWS:
		cloud, err := BuildCloud(cluster)
		if err != nil {
			return "", fmt.Errorf("error building cloud for AWS cluster: %v", err)
		}

		instanceType, err := cloud.(awsup.AWSCloud).DefaultInstanceType(cluster, ig)
		if err != nil {
			return "", fmt.Errorf("error finding default machine type: %v", err)
		}
		return instanceType, nil

	case kops.CloudProviderGCE:
		switch ig.Spec.Role {
		case kops.InstanceGroupRoleMaster:
			return defaultMasterMachineTypeGCE, nil

		case kops.InstanceGroupRoleNode:
			return defaultNodeMachineTypeGCE, nil

		case kops.InstanceGroupRoleBastion:
			return defaultBastionMachineTypeGCE, nil
		}

	case kops.CloudProviderDO:
		switch ig.Spec.Role {
		case kops.InstanceGroupRoleMaster:
			return defaultMasterMachineTypeDO, nil

		case kops.InstanceGroupRoleNode:
			return defaultNodeMachineTypeDO, nil

		}

	case kops.CloudProviderVSphere:
		switch ig.Spec.Role {
		case kops.InstanceGroupRoleMaster:
			return defaultMasterMachineTypeVSphere, nil

		case kops.InstanceGroupRoleNode:
			return defaultNodeMachineTypeVSphere, nil

		case kops.InstanceGroupRoleBastion:
			return defaultBastionMachineTypeVSphere, nil
		}
	}

	glog.V(2).Infof("Cannot set default MachineType for CloudProvider=%q, Role=%q", cluster.Spec.CloudProvider, ig.Spec.Role)
	return "", nil
}

// defaultImage returns the default Image, based on the cloudprovider
func defaultImage(cluster *kops.Cluster, channel *kops.Channel) string {
	if channel != nil {
		var kubernetesVersion *semver.Version
		if cluster.Spec.KubernetesVersion != "" {
			var err error
			kubernetesVersion, err = util.ParseKubernetesVersion(cluster.Spec.KubernetesVersion)
			if err != nil {
				glog.Warningf("cannot parse KubernetesVersion %q in cluster", cluster.Spec.KubernetesVersion)
			}
		}
		if kubernetesVersion != nil {
			image := channel.FindImage(kops.CloudProviderID(cluster.Spec.CloudProvider), *kubernetesVersion)
			if image != nil {
				return image.Name
			}
		}
	}

	switch kops.CloudProviderID(cluster.Spec.CloudProvider) {
	case kops.CloudProviderDO:
		return defaultDONodeImage
	case kops.CloudProviderVSphere:
		return defaultVSphereNodeImage
	}

	glog.Infof("Cannot set default Image for CloudProvider=%q", cluster.Spec.CloudProvider)
	return ""
}
<commit_msg>Determine the default flavours<commit_after>/*
Copyright 2016 The Kubernetes Authors.

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

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

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

package cloudup

import (
	"fmt"

	"github.com/blang/semver"
	"github.com/golang/glog"
	"k8s.io/kops/pkg/apis/kops"
	"k8s.io/kops/pkg/apis/kops/util"
	"k8s.io/kops/pkg/apis/kops/validation"
	"k8s.io/kops/upup/pkg/fi"
	"k8s.io/kops/upup/pkg/fi/cloudup/awsup"
	"k8s.io/kops/util/pkg/reflectutils"
	"k8s.io/kops/upup/pkg/fi/cloudup/openstack"
	"k8s.io/kops/upup/pkg/fi/utils"
)

// Default Machine types for various types of instance group machine
const (
	defaultNodeMachineTypeGCE     = "n1-standard-2"
	defaultNodeMachineTypeVSphere = "vsphere_node"
	defaultNodeMachineTypeDO      = "s-2vcpu-4gb"

	defaultBastionMachineTypeGCE     = "f1-micro"
	defaultBastionMachineTypeVSphere = "vsphere_bastion"

	defaultMasterMachineTypeGCE     = "n1-standard-1"
	defaultMasterMachineTypeVSphere = "vsphere_master"
	defaultMasterMachineTypeDO      = "s-2vcpu-2gb"

	defaultVSphereNodeImage = "kops_ubuntu_16_04.ova"
	defaultDONodeImage      = "coreos-stable"
)

var awsDedicatedInstanceExceptions = map[string]bool{
	"t2.nano":   true,
	"t2.micro":  true,
	"t2.small":  true,
	"t2.medium": true,
	"t2.large":  true,
	"t2.xlarge": true,
}

// PopulateInstanceGroupSpec sets default values in the InstanceGroup
// The InstanceGroup is simpler than the cluster spec, so we just populate in place (like the rest of k8s)
func PopulateInstanceGroupSpec(cluster *kops.Cluster, input *kops.InstanceGroup, channel *kops.Channel) (*kops.InstanceGroup, error) {
	err := validation.ValidateInstanceGroup(input)
	if err != nil {
		return nil, err
	}

	ig := &kops.InstanceGroup{}
	reflectutils.JsonMergeStruct(ig, input)

	// TODO: Clean up
	if ig.IsMaster() {
		if ig.Spec.MachineType == "" {
			ig.Spec.MachineType, err = defaultMachineType(cluster, ig)
			if err != nil {
				return nil, fmt.Errorf("error assigning default machine type for masters: %v", err)
			}

		}
		if ig.Spec.MinSize == nil {
			ig.Spec.MinSize = fi.Int32(1)
		}
		if ig.Spec.MaxSize == nil {
			ig.Spec.MaxSize = fi.Int32(1)
		}
	} else if ig.Spec.Role == kops.InstanceGroupRoleBastion {
		if ig.Spec.MachineType == "" {
			ig.Spec.MachineType, err = defaultMachineType(cluster, ig)
			if err != nil {
				return nil, fmt.Errorf("error assigning default machine type for bastions: %v", err)
			}
		}
		if ig.Spec.MinSize == nil {
			ig.Spec.MinSize = fi.Int32(1)
		}
		if ig.Spec.MaxSize == nil {
			ig.Spec.MaxSize = fi.Int32(1)
		}
	} else {
		if ig.Spec.MachineType == "" {
			ig.Spec.MachineType, err = defaultMachineType(cluster, ig)
			if err != nil {
				return nil, fmt.Errorf("error assigning default machine type for nodes: %v", err)
			}
		}
		if ig.Spec.MinSize == nil {
			ig.Spec.MinSize = fi.Int32(2)
		}
		if ig.Spec.MaxSize == nil {
			ig.Spec.MaxSize = fi.Int32(2)
		}
	}

	if ig.Spec.Image == "" {
		ig.Spec.Image = defaultImage(cluster, channel)
	}

	if ig.Spec.Tenancy != "" && ig.Spec.Tenancy != "default" {
		switch kops.CloudProviderID(cluster.Spec.CloudProvider) {
		case kops.CloudProviderAWS:
			if _, ok := awsDedicatedInstanceExceptions[ig.Spec.MachineType]; ok {
				return nil, fmt.Errorf("Invalid dedicated instance type: %s", ig.Spec.MachineType)
			}
		default:
			glog.Warning("Trying to set tenancy on non-AWS environment")
		}
	}

	if ig.IsMaster() {
		if len(ig.Spec.Subnets) == 0 {
			return nil, fmt.Errorf("Master InstanceGroup %s did not specify any Subnets", ig.ObjectMeta.Name)
		}
	} else if ig.Spec.Role == kops.InstanceGroupRoleBastion {
		if len(ig.Spec.Subnets) == 0 {
			for _, subnet := range cluster.Spec.Subnets {
				if subnet.Type == kops.SubnetTypeUtility {
					ig.Spec.Subnets = append(ig.Spec.Subnets, subnet.Name)
				}
			}
		}
	} else {
		if len(ig.Spec.Subnets) == 0 {
			for _, subnet := range cluster.Spec.Subnets {
				if subnet.Type != kops.SubnetTypeUtility {
					ig.Spec.Subnets = append(ig.Spec.Subnets, subnet.Name)
				}
			}
		}
	}

	if len(ig.Spec.Subnets) == 0 {
		return nil, fmt.Errorf("unable to infer any Subnets for InstanceGroup %s ", ig.ObjectMeta.Name)
	}

	return ig, nil
}

// defaultMachineType returns the default MachineType for the instance group, based on the cloudprovider
func defaultMachineType(cluster *kops.Cluster, ig *kops.InstanceGroup) (string, error) {
	switch kops.CloudProviderID(cluster.Spec.CloudProvider) {
	case kops.CloudProviderAWS:
		cloud, err := BuildCloud(cluster)
		if err != nil {
			return "", fmt.Errorf("error building cloud for AWS cluster: %v", err)
		}

		instanceType, err := cloud.(awsup.AWSCloud).DefaultInstanceType(cluster, ig)
		if err != nil {
			return "", fmt.Errorf("error finding default machine type: %v", err)
		}
		return instanceType, nil

	case kops.CloudProviderGCE:
		switch ig.Spec.Role {
		case kops.InstanceGroupRoleMaster:
			return defaultMasterMachineTypeGCE, nil

		case kops.InstanceGroupRoleNode:
			return defaultNodeMachineTypeGCE, nil

		case kops.InstanceGroupRoleBastion:
			return defaultBastionMachineTypeGCE, nil
		}

	case kops.CloudProviderDO:
		switch ig.Spec.Role {
		case kops.InstanceGroupRoleMaster:
			return defaultMasterMachineTypeDO, nil

		case kops.InstanceGroupRoleNode:
			return defaultNodeMachineTypeDO, nil

		}

	case kops.CloudProviderVSphere:
		switch ig.Spec.Role {
		case kops.InstanceGroupRoleMaster:
			return defaultMasterMachineTypeVSphere, nil

		case kops.InstanceGroupRoleNode:
			return defaultNodeMachineTypeVSphere, nil

		case kops.InstanceGroupRoleBastion:
			return defaultBastionMachineTypeVSphere, nil
		}

	case kops.CloudProviderOpenstack:
		cloud, err := BuildCloud(cluster)
		if err != nil {
			return "", fmt.Errorf("error building cloud for Openstack cluster: %v", err)
		}

		instanceType, err := cloud.(openstack.OpenstackCloud).DefaultInstanceType(cluster, ig)
		if err != nil {
			return "", fmt.Errorf("error finding default machine type: %v", err)
		}
		return instanceType, nil

	}

	glog.V(2).Infof("Cannot set default MachineType for CloudProvider=%q, Role=%q", cluster.Spec.CloudProvider, ig.Spec.Role)
	return "", nil
}

// defaultImage returns the default Image, based on the cloudprovider
func defaultImage(cluster *kops.Cluster, channel *kops.Channel) string {
	if channel != nil {
		var kubernetesVersion *semver.Version
		if cluster.Spec.KubernetesVersion != "" {
			var err error
			kubernetesVersion, err = util.ParseKubernetesVersion(cluster.Spec.KubernetesVersion)
			if err != nil {
				glog.Warningf("cannot parse KubernetesVersion %q in cluster", cluster.Spec.KubernetesVersion)
			}
		}
		if kubernetesVersion != nil {
			image := channel.FindImage(kops.CloudProviderID(cluster.Spec.CloudProvider), *kubernetesVersion)
			if image != nil {
				return image.Name
			}
		}
	}

	switch kops.CloudProviderID(cluster.Spec.CloudProvider) {
	case kops.CloudProviderDO:
		return defaultDONodeImage
	case kops.CloudProviderVSphere:
		return defaultVSphereNodeImage
	}

	glog.Infof("Cannot set default Image for CloudProvider=%q", cluster.Spec.CloudProvider)
	return ""
}
<|endoftext|>