<commit_before>package sqlstmt

import (
	"bytes"
	"database/sql"
	"errors"
	"fmt"
	"reflect"
	"strings"

	"github.com/jjeffery/sqlstmt/private/column"
	"github.com/jjeffery/sqlstmt/private/scanner"
)

var errNotImplemented = errors.New("not implemented")

// InsertRowStmt inserts a single row. It is safe for concurrent
// access by multiple goroutines.
type InsertRowStmt struct {
	commonStmt
	autoIncrColumn *column.Info
}

// NewInsertRowStmt returns a new InsertRowStmt for the given
// row and SQL.
// It is safe for concurrent access by multiple goroutines.
// The dialect and naming conventions are inferred
// from DefaultSchema.
func NewInsertRowStmt(row interface{}, sql string) *InsertRowStmt {
	return newInsertRowStmt(DefaultSchema, row, sql)
}

func newInsertRowStmt(schema *Schema, row interface{}, sql string) *InsertRowStmt {
	stmt := &InsertRowStmt{}
	sql = checkSQL(sql, insertFormat)
	stmt.err = stmt.prepareExecRow(schema, row, sql)

	for _, col := range stmt.columns {
		if col.AutoIncrement {
			stmt.autoIncrColumn = col
			break
		}
	}

	if stmt.autoIncrColumn != nil {
		// Some DBs allow the auto-increment column to be specified.
		// Work out if this statement is doing this.
		for _, col := range stmt.inputs {
			if col == stmt.autoIncrColumn {
				// this statement is setting the auto-increment column explicitly
				stmt.autoIncrColumn = nil
				break
			}
		}
	}

	return stmt
}

// Exec executes the insert statement using the row as arguments.
func (stmt *InsertRowStmt) Exec(db Execer, row interface{}) error {
	// field for setting the auto-increment value
	var field reflect.Value
	if stmt.autoIncrColumn != nil {
		rowVal := reflect.ValueOf(row)
		field = stmt.autoIncrColumn.Index.ValueRW(rowVal)
		if !field.CanSet() {
			return fmt.Errorf("cannot set auto-increment value for type %s", rowVal.Type().Name())
		}
	}

	result, err := stmt.doExec(db, row)
	if err != nil {
		return err
	}

	if field.IsValid() {
		n, err := result.LastInsertId()
		if err != nil {
			return nil
		}
		// TODO: could catch a panic here if the type is not int8, 1nt16, int32, int64
		field.SetInt(n)
	}
	return nil
}

// ExecRowStmt updates or deletes a single row. It is safe for concurrent
// access by multiple goroutines.
type ExecRowStmt struct {
	commonStmt
}

// NewUpdateRowStmt returns a new ExecRowStmt for updating a single row.
// It is safe for concurrent access by multiple goroutines.
// The dialect and naming conventions are obtained from DefaultSchema.
func NewUpdateRowStmt(row interface{}, sql string) *ExecRowStmt {
	return newUpdateRowStmt(DefaultSchema, row, sql)
}

func newUpdateRowStmt(schema *Schema, row interface{}, sql string) *ExecRowStmt {
	stmt := &ExecRowStmt{}
	sql = checkSQL(sql, updateFormat)
	stmt.err = stmt.prepareExecRow(schema, row, sql)
	return stmt
}

// NewDeleteRowStmt returns a new ExecRowStmt for deleting a single
// row.
// It is safe for concurrent access by multiple goroutines.
// The dialect and naming conventions are obtained from DefaultSchema.
func NewDeleteRowStmt(row interface{}, sql string) *ExecRowStmt {
	return newDeleteRowStmt(DefaultSchema, row, sql)
}

func newDeleteRowStmt(schema *Schema, row interface{}, sql string) *ExecRowStmt {
	stmt := &ExecRowStmt{}
	sql = checkSQL(sql, deleteFormat)
	stmt.err = stmt.prepareExecRow(schema, row, sql)
	return stmt
}

// Exec executes the statement using the row as arguments. Returns the
// number of rows affected.
func (stmt *ExecRowStmt) Exec(db Execer, row interface{}) (int, error) {
	result, err := stmt.doExec(db, row)
	if err != nil {
		return 0, err
	}

	n, err := result.RowsAffected()
	if err != nil {
		return 0, err
	}

	return int(n), nil
}

// GetRowStmt executes a query that returns a single row.
// It is safe for concurrent access by multiple goroutines.
type GetRowStmt struct {
	commonStmt
}

// NewGetRowStmt returns a new GetRowStmt for the given
// row and SQL.
// It is safe for concurrent access by multiple goroutines.
// The dialect and naming conventions are obtained
// from DefaultSchema.
func NewGetRowStmt(row interface{}, sql string) *GetRowStmt {
	return newGetRowStmt(DefaultSchema, row, sql)
}

func newGetRowStmt(schema *Schema, row interface{}, sql string) *GetRowStmt {
	stmt := &GetRowStmt{}
	sql = checkSQL(sql, getFormat)
	stmt.err = stmt.prepareCommon(schema, row, sql)
	return stmt
}

// Get a single row into dest based on the fields populated in dest.
func (stmt *GetRowStmt) Get(db Queryer, dest interface{}) (int, error) {
	errorPtrType := func() error {
		expectedTypeName := stmt.expectedTypeName()
		return fmt.Errorf("expected dest to be *%s", expectedTypeName)
	}

	destValue := reflect.ValueOf(dest)

	if destValue.Kind() != reflect.Ptr {
		return 0, errorPtrType()
	}
	if destValue.IsNil() {
		return 0, errors.New("nil pointer passed")
	}

	rowValue := reflect.Indirect(destValue)
	rowType := rowValue.Type()
	if rowType != stmt.rowType {
		return 0, errorPtrType()
	}

	args, err := stmt.getArgs(dest)
	if err != nil {
		return 0, nil
	}

	if stmt.Logger != nil {
		msg := fmt.Sprintf("query=%q, args=%v\n", stmt.query, args)
		stmt.Logger.Print(msg)
	}
	rows, err := db.Query(stmt.query, args...)
	if err != nil {
		return 0, err
	}
	defer rows.Close()

	scanValues := make([]interface{}, len(stmt.outputs))

	if !rows.Next() {
		// no rows returned
		return 0, nil
	}

	for i, col := range stmt.outputs {
		cellValue := col.Index.ValueRW(rowValue)
		scanValues[i] = cellValue.Addr().Interface()
	}
	err = rows.Scan(scanValues...)
	if err != nil {
		return 0, err
	}

	return 1, nil
}

// SelectStmt executes a query that returns multiple rows.
// It is safe for concurrent access by multiple goroutines.
type SelectStmt struct {
	commonStmt
}

// NewSelectStmt returns a new SelectStmt for the given
// row and SQL. The dialect and naming conventions are obtained
// from DefaultSchema.
func NewSelectStmt(row interface{}, sql string) *SelectStmt {
	return newSelectStmt(DefaultSchema, row, sql)
}

func newSelectStmt(schema *Schema, row interface{}, sql string) *SelectStmt {
	stmt := &SelectStmt{}
	sql = checkSQL(sql, selectFormat)
	stmt.err = stmt.prepareCommon(schema, row, sql)
	if stmt.err == nil && len(stmt.inputs) > 0 {
		stmt.err = errors.New("unexpected inputs in query")
	}
	return stmt
}

// Select executes the statement's query and returns the resulting
// rows in the slice pointed to by dest. The args are for any
// placeholder parameters in the query.
func (stmt *SelectStmt) Select(db Queryer, dest interface{}, args ...interface{}) error {
	destValue := reflect.ValueOf(dest)

	if destValue.Kind() != reflect.Ptr {
		return stmt.errorSliceType()
	}
	if destValue.IsNil() {
		return errors.New("Select: nil pointer passed as dest")
	}

	sliceValue := reflect.Indirect(destValue)
	sliceType := sliceValue.Type()
	if sliceType.Kind() != reflect.Slice {
		return stmt.errorSliceType()
	}

	rowType := sliceType.Elem()
	isPtr := rowType.Kind() == reflect.Ptr
	if isPtr {
		rowType = rowType.Elem()
	}
	if rowType != stmt.rowType {
		return stmt.errorSliceType()
	}

	rows, err := db.Query(stmt.query, args...)
	if err != nil {
		return err
	}
	defer rows.Close()

	scanValues := make([]interface{}, len(stmt.columns))

	for rows.Next() {
		rowValuePtr := reflect.New(rowType)
		rowValue := reflect.Indirect(rowValuePtr)
		for i, col := range stmt.columns {
			cellValue := col.Index.ValueRW(rowValue)
			scanValues[i] = cellValue.Addr().Interface()
		}
		err = rows.Scan(scanValues...)
		if err != nil {
			return err
		}
		if isPtr {
			sliceValue.Set(reflect.Append(sliceValue, rowValuePtr))
		} else {
			sliceValue.Set(reflect.Append(sliceValue, rowValue))
		}
	}

	return rows.Err()
}

func (stmt *SelectStmt) errorSliceType() error {
	expectedType := stmt.expectedTypeName()
	return fmt.Errorf("Expected dest to be pointer to []%s or []*%s", expectedType, expectedType)
}

type commonStmt struct {
	// Logger is used for diagnostic logging.
	Logger Logger

	rowType    reflect.Type
	query      string
	dialect    Dialect
	convention Convention
	columns    []*column.Info
	inputs     []*column.Info
	outputs    []*column.Info
	err        error
}

// String prints the SQL query associated with the statement.
func (stmt *commonStmt) String() string {
	return stmt.query
}

func (stmt *commonStmt) prepareCommon(schema *Schema, row interface{}, sql string) error {
	stmt.rowType = reflect.TypeOf(row)
	if stmt.rowType.Kind() == reflect.Ptr {
		stmt.rowType = stmt.rowType.Elem()
	}
	stmt.columns = column.ListForType(stmt.rowType)
	stmt.convention = schema.convention()
	stmt.dialect = schema.dialect()
	stmt.Logger = schema.Logger
	if err := stmt.scanSQL(sql); err != nil {
		return err
	}
	if stmt.Logger != nil {
		msg := fmt.Sprintf("prepared=%q", stmt.query)
		stmt.Logger.Print(msg)
	}
	return nil
}

func (stmt *commonStmt) prepareExecRow(schema *Schema, row interface{}, sql string) error {
	if err := stmt.prepareCommon(schema, row, sql); err != nil {
		return err
	}
	if len(stmt.outputs) > 0 {
		return errors.New("unexpected query columns in exec statement")
	}
	return nil
}

func (stmt *commonStmt) addColumns(cols columnsT) {
	if cols.clause.isInput() {
		for _, col := range cols.filtered() {
			stmt.inputs = append(stmt.inputs, col)
		}
	} else if cols.clause.isOutput() {
		for _, col := range cols.filtered() {
			stmt.outputs = append(stmt.outputs, col)
		}
	}
}

func (stmt *commonStmt) scanSQL(query string) error {
	query = strings.TrimSpace(query)
	scan := scanner.New(strings.NewReader(query))
	columns := newColumns(stmt.columns, stmt.convention, stmt.dialect)
	var insertColumns *columnsT
	placeholderCount := 0
	var clause sqlClause
	var buf bytes.Buffer
	for {
		tok, lit := scan.Scan()
		if tok == scanner.EOF {
			break
		}
		switch tok {
		case scanner.WS:
			buf.WriteRune(' ')
		case scanner.COMMENT:
		case scanner.LITERAL, scanner.OP:
			buf.WriteString(lit)
		case scanner.PLACEHOLDER:
			placeholderCount++
			buf.WriteString(stmt.dialect.Placeholder(placeholderCount))
		case scanner.IDENT:
			if lit[0] == '{' {
				if !clause.acceptsColumns() {
					// invalid place to insert columns
					return fmt.Errorf("cannot expand %q in %q clause", lit, clause)
				}
				lit = strings.TrimSpace(scanner.Unquote(lit))
				if clause == clauseInsertValues {
					if lit != "" {
						return fmt.Errorf("columns for %q clause always match the %q clause",
							clause, clauseInsertColumns)
					}
					if insertColumns == nil {
						return fmt.Errorf("cannot expand %q clause because %q clause is missing",
							clause, clauseInsertColumns)
					}

					// change the clause but keep the filter and generate string
					cols := *insertColumns
					cols.clause = clause
					buf.WriteString(cols.String())
					stmt.addColumns(cols)
				} else {
					cols, err := columns.Parse(clause, lit)
					if err != nil {
						return fmt.Errorf("cannot expand %q in %q clause: %v", lit, clause, err)
					}
					buf.WriteString(cols.String())
					stmt.addColumns(cols)
					if clause == clauseInsertColumns {
						insertColumns = &cols
					}
				}
			} else if scanner.IsQuoted(lit) {
				lit = scanner.Unquote(lit)
				buf.WriteString(stmt.dialect.Quote(lit))
			} else {
				buf.WriteString(lit)

				// an unquoted identifer might be an SQL keyword
				clause = clause.nextClause(lit)
			}
		}
	}
	stmt.query = buf.String()
	return nil
}

func (stmt commonStmt) doExec(db Execer, row interface{}) (sql.Result, error) {
	if stmt.err != nil {
		return nil, stmt.err
	}
	args, err := stmt.getArgs(row)
	if err != nil {
		return nil, err
	}
	if stmt.Logger != nil {
		msg := fmt.Sprintf("query=%q, args=%v\n", stmt.query, args)
		stmt.Logger.Print(msg)
	}
	return db.Exec(stmt.query, args...)
}

func (stmt commonStmt) getArgs(row interface{}) ([]interface{}, error) {
	if stmt.err != nil {
		return nil, stmt.err
	}

	var args []interface{}

	rowVal := reflect.ValueOf(row)
	for rowVal.Type().Kind() == reflect.Ptr {
		rowVal = rowVal.Elem()
	}
	if rowVal.Type() != stmt.rowType {
		expectedType := stmt.expectedTypeName()
		return nil, fmt.Errorf("expected type %s or *(%s)", expectedType, expectedType)
	}

	for _, input := range stmt.inputs {
		args = append(args, input.Index.ValueRO(rowVal).Interface())
	}

	return args, nil
}

func (stmt commonStmt) expectedTypeName() string {
	return fmt.Sprintf("%s.%s", stmt.rowType.PkgPath(), stmt.rowType.Name())
}
<commit_msg>Fix bugs found in testing<commit_after>package sqlstmt

import (
	"bytes"
	"database/sql"
	"errors"
	"fmt"
	"reflect"
	"strings"

	"github.com/jjeffery/sqlstmt/private/column"
	"github.com/jjeffery/sqlstmt/private/scanner"
)

var errNotImplemented = errors.New("not implemented")

// InsertRowStmt inserts a single row. It is safe for concurrent
// access by multiple goroutines.
type InsertRowStmt struct {
	commonStmt
	autoIncrColumn *column.Info
}

// NewInsertRowStmt returns a new InsertRowStmt for the given
// row and SQL.
// It is safe for concurrent access by multiple goroutines.
// The dialect and naming conventions are inferred
// from DefaultSchema.
func NewInsertRowStmt(row interface{}, sql string) *InsertRowStmt {
	return newInsertRowStmt(DefaultSchema, row, sql)
}

func newInsertRowStmt(schema *Schema, row interface{}, sql string) *InsertRowStmt {
	stmt := &InsertRowStmt{}
	sql = checkSQL(sql, insertFormat)
	stmt.err = stmt.prepareExecRow(schema, row, sql)

	for _, col := range stmt.columns {
		if col.AutoIncrement {
			stmt.autoIncrColumn = col
			break
		}
	}

	if stmt.autoIncrColumn != nil {
		// Some DBs allow the auto-increment column to be specified.
		// Work out if this statement is doing this.
		for _, col := range stmt.inputs {
			if col == stmt.autoIncrColumn {
				// this statement is setting the auto-increment column explicitly
				stmt.autoIncrColumn = nil
				break
			}
		}
	}

	return stmt
}

// Exec executes the insert statement using the row as arguments.
func (stmt *InsertRowStmt) Exec(db Execer, row interface{}) error {
	// field for setting the auto-increment value
	var field reflect.Value
	if stmt.autoIncrColumn != nil {
		rowVal := reflect.ValueOf(row)
		field = stmt.autoIncrColumn.Index.ValueRW(rowVal)
		if !field.CanSet() {
			return fmt.Errorf("cannot set auto-increment value for type %s", rowVal.Type().Name())
		}
	}

	result, err := stmt.doExec(db, row)
	if err != nil {
		return err
	}

	if field.IsValid() {
		n, err := result.LastInsertId()
		if err != nil {
			return nil
		}
		// TODO: could catch a panic here if the type is not int8, 1nt16, int32, int64
		field.SetInt(n)
	}
	return nil
}

// ExecRowStmt updates or deletes a single row. It is safe for concurrent
// access by multiple goroutines.
type ExecRowStmt struct {
	commonStmt
}

// NewUpdateRowStmt returns a new ExecRowStmt for updating a single row.
// It is safe for concurrent access by multiple goroutines.
// The dialect and naming conventions are obtained from DefaultSchema.
func NewUpdateRowStmt(row interface{}, sql string) *ExecRowStmt {
	return newUpdateRowStmt(DefaultSchema, row, sql)
}

func newUpdateRowStmt(schema *Schema, row interface{}, sql string) *ExecRowStmt {
	stmt := &ExecRowStmt{}
	sql = checkSQL(sql, updateFormat)
	stmt.err = stmt.prepareExecRow(schema, row, sql)
	return stmt
}

// NewDeleteRowStmt returns a new ExecRowStmt for deleting a single
// row.
// It is safe for concurrent access by multiple goroutines.
// The dialect and naming conventions are obtained from DefaultSchema.
func NewDeleteRowStmt(row interface{}, sql string) *ExecRowStmt {
	return newDeleteRowStmt(DefaultSchema, row, sql)
}

func newDeleteRowStmt(schema *Schema, row interface{}, sql string) *ExecRowStmt {
	stmt := &ExecRowStmt{}
	sql = checkSQL(sql, deleteFormat)
	stmt.err = stmt.prepareExecRow(schema, row, sql)
	return stmt
}

// Exec executes the statement using the row as arguments. Returns the
// number of rows affected.
func (stmt *ExecRowStmt) Exec(db Execer, row interface{}) (int, error) {
	result, err := stmt.doExec(db, row)
	if err != nil {
		return 0, err
	}

	n, err := result.RowsAffected()
	if err != nil {
		return 0, err
	}

	return int(n), nil
}

// GetRowStmt executes a query that returns a single row.
// It is safe for concurrent access by multiple goroutines.
type GetRowStmt struct {
	commonStmt
}

// NewGetRowStmt returns a new GetRowStmt for the given
// row and SQL.
// It is safe for concurrent access by multiple goroutines.
// The dialect and naming conventions are obtained
// from DefaultSchema.
func NewGetRowStmt(row interface{}, sql string) *GetRowStmt {
	return newGetRowStmt(DefaultSchema, row, sql)
}

func newGetRowStmt(schema *Schema, row interface{}, sql string) *GetRowStmt {
	stmt := &GetRowStmt{}
	sql = checkSQL(sql, getFormat)
	stmt.err = stmt.prepareCommon(schema, row, sql)
	return stmt
}

// Get a single row into dest based on the fields populated in dest.
func (stmt *GetRowStmt) Get(db Queryer, dest interface{}) (int, error) {
	errorPtrType := func() error {
		expectedTypeName := stmt.expectedTypeName()
		return fmt.Errorf("expected dest to be *%s", expectedTypeName)
	}

	destValue := reflect.ValueOf(dest)

	if destValue.Kind() != reflect.Ptr {
		return 0, errorPtrType()
	}
	if destValue.IsNil() {
		return 0, errors.New("nil pointer passed")
	}

	rowValue := reflect.Indirect(destValue)
	rowType := rowValue.Type()
	if rowType != stmt.rowType {
		return 0, errorPtrType()
	}

	args, err := stmt.getArgs(dest)
	if err != nil {
		return 0, nil
	}

	if stmt.Logger != nil {
		msg := fmt.Sprintf("query=%q, args=%v\n", stmt.query, args)
		stmt.Logger.Print(msg)
	}
	rows, err := db.Query(stmt.query, args...)
	if err != nil {
		return 0, err
	}
	defer rows.Close()

	scanValues := make([]interface{}, len(stmt.outputs))

	if !rows.Next() {
		// no rows returned
		return 0, nil
	}

	for i, col := range stmt.outputs {
		cellValue := col.Index.ValueRW(rowValue)
		scanValues[i] = cellValue.Addr().Interface()
	}
	err = rows.Scan(scanValues...)
	if err != nil {
		return 0, err
	}

	return 1, nil
}

// SelectStmt executes a query that returns multiple rows.
// It is safe for concurrent access by multiple goroutines.
type SelectStmt struct {
	commonStmt
}

// NewSelectStmt returns a new SelectStmt for the given
// row and SQL. The dialect and naming conventions are obtained
// from DefaultSchema.
func NewSelectStmt(row interface{}, sql string) *SelectStmt {
	return newSelectStmt(DefaultSchema, row, sql)
}

func newSelectStmt(schema *Schema, row interface{}, sql string) *SelectStmt {
	stmt := &SelectStmt{}
	sql = checkSQL(sql, selectFormat)
	stmt.err = stmt.prepareCommon(schema, row, sql)
	if stmt.err == nil && len(stmt.inputs) > 0 {
		stmt.err = errors.New("unexpected inputs in query")
	}
	return stmt
}

// Select executes the statement's query and returns the resulting
// rows in the slice pointed to by dest. The args are for any
// placeholder parameters in the query.
func (stmt *SelectStmt) Select(db Queryer, dest interface{}, args ...interface{}) error {
	destValue := reflect.ValueOf(dest)

	if destValue.Kind() != reflect.Ptr {
		return stmt.errorSliceType()
	}
	if destValue.IsNil() {
		return errors.New("Select: nil pointer passed as dest")
	}

	sliceValue := reflect.Indirect(destValue)
	sliceType := sliceValue.Type()
	if sliceType.Kind() != reflect.Slice {
		return stmt.errorSliceType()
	}

	rowType := sliceType.Elem()
	isPtr := rowType.Kind() == reflect.Ptr
	if isPtr {
		rowType = rowType.Elem()
	}
	if rowType != stmt.rowType {
		return stmt.errorSliceType()
	}

	rows, err := db.Query(stmt.query, args...)
	if err != nil {
		return err
	}
	defer rows.Close()

	scanValues := make([]interface{}, len(stmt.columns))

	for rows.Next() {
		rowValuePtr := reflect.New(rowType)
		rowValue := reflect.Indirect(rowValuePtr)
		for i, col := range stmt.columns {
			cellValue := col.Index.ValueRW(rowValue)
			scanValues[i] = cellValue.Addr().Interface()
		}
		err = rows.Scan(scanValues...)
		if err != nil {
			return err
		}
		if isPtr {
			sliceValue.Set(reflect.Append(sliceValue, rowValuePtr))
		} else {
			sliceValue.Set(reflect.Append(sliceValue, rowValue))
		}
	}

	return rows.Err()
}

func (stmt *SelectStmt) errorSliceType() error {
	expectedType := stmt.expectedTypeName()
	return fmt.Errorf("Expected dest to be pointer to []%s or []*%s", expectedType, expectedType)
}

type commonStmt struct {
	// Logger is used for diagnostic logging.
	Logger Logger

	rowType    reflect.Type
	query      string
	dialect    Dialect
	convention Convention
	columns    []*column.Info
	inputs     []*column.Info
	outputs    []*column.Info
	err        error
}

// String prints the SQL query associated with the statement.
func (stmt *commonStmt) String() string {
	return stmt.query
}

func (stmt *commonStmt) prepareCommon(schema *Schema, row interface{}, sql string) error {
	stmt.rowType = reflect.TypeOf(row)
	if stmt.rowType.Kind() == reflect.Ptr {
		stmt.rowType = stmt.rowType.Elem()
	}
	stmt.columns = column.ListForType(stmt.rowType)
	stmt.convention = schema.convention()
	stmt.dialect = schema.dialect()
	stmt.Logger = schema.Logger
	if err := stmt.scanSQL(sql); err != nil {
		return err
	}
	if stmt.Logger != nil {
		msg := fmt.Sprintf("prepared=%q", stmt.query)
		stmt.Logger.Print(msg)
	}
	return nil
}

func (stmt *commonStmt) prepareExecRow(schema *Schema, row interface{}, sql string) error {
	if err := stmt.prepareCommon(schema, row, sql); err != nil {
		return err
	}
	if len(stmt.outputs) > 0 {
		return errors.New("unexpected query columns in exec statement")
	}
	return nil
}

func (stmt *commonStmt) addColumns(cols columnsT) {
	if cols.clause.isInput() {
		for _, col := range cols.filtered() {
			stmt.inputs = append(stmt.inputs, col)
		}
	} else if cols.clause.isOutput() {
		for _, col := range cols.filtered() {
			stmt.outputs = append(stmt.outputs, col)
		}
	}
}

// counter is used for keeping track of placeholders
type counter int

func (c *counter) Next() int {
	*c++
	return int(*c)
}

func (stmt *commonStmt) scanSQL(query string) error {
	query = strings.TrimSpace(query)
	scan := scanner.New(strings.NewReader(query))
	var counter counter
	columns := newColumns(stmt.columns, stmt.convention, stmt.dialect, counter.Next)
	var insertColumns *columnsT
	var clause sqlClause
	var buf bytes.Buffer
	for {
		tok, lit := scan.Scan()
		if tok == scanner.EOF {
			break
		}
		switch tok {
		case scanner.WS:
			buf.WriteRune(' ')
		case scanner.COMMENT:
			// strip comment
		case scanner.LITERAL, scanner.OP:
			buf.WriteString(lit)
		case scanner.PLACEHOLDER:
			buf.WriteString(stmt.dialect.Placeholder(counter.Next()))
		case scanner.IDENT:
			if lit[0] == '{' {
				if !clause.acceptsColumns() {
					// invalid place to insert columns
					return fmt.Errorf("cannot expand %q in %q clause", lit, clause)
				}
				lit = strings.TrimSpace(scanner.Unquote(lit))
				if clause == clauseInsertValues {
					if lit != "" {
						return fmt.Errorf("columns for %q clause always match the %q clause",
							clause, clauseInsertColumns)
					}
					if insertColumns == nil {
						return fmt.Errorf("cannot expand %q clause because %q clause is missing",
							clause, clauseInsertColumns)
					}

					// change the clause but keep the filter and generate string
					cols := *insertColumns
					cols.clause = clause
					buf.WriteString(cols.String())
					stmt.addColumns(cols)
				} else {
					cols, err := columns.Parse(clause, lit)
					if err != nil {
						return fmt.Errorf("cannot expand %q in %q clause: %v", lit, clause, err)
					}
					buf.WriteString(cols.String())
					stmt.addColumns(cols)
					if clause == clauseInsertColumns {
						insertColumns = &cols
					}
				}
			} else if scanner.IsQuoted(lit) {
				lit = scanner.Unquote(lit)
				buf.WriteString(stmt.dialect.Quote(lit))
			} else {
				buf.WriteString(lit)

				// an unquoted identifer might be an SQL keyword
				clause = clause.nextClause(lit)
			}
		}
	}
	stmt.query = buf.String()
	return nil
}

func (stmt commonStmt) doExec(db Execer, row interface{}) (sql.Result, error) {
	if stmt.err != nil {
		return nil, stmt.err
	}
	args, err := stmt.getArgs(row)
	if err != nil {
		return nil, err
	}
	if stmt.Logger != nil {
		msg := fmt.Sprintf("query=%q, args=%v\n", stmt.query, args)
		stmt.Logger.Print(msg)
	}
	return db.Exec(stmt.query, args...)
}

func (stmt commonStmt) getArgs(row interface{}) ([]interface{}, error) {
	if stmt.err != nil {
		return nil, stmt.err
	}

	var args []interface{}

	rowVal := reflect.ValueOf(row)
	for rowVal.Type().Kind() == reflect.Ptr {
		rowVal = rowVal.Elem()
	}
	if rowVal.Type() != stmt.rowType {
		expectedType := stmt.expectedTypeName()
		return nil, fmt.Errorf("expected type %s or *(%s)", expectedType, expectedType)
	}

	for _, input := range stmt.inputs {
		args = append(args, input.Index.ValueRO(rowVal).Interface())
	}

	return args, nil
}

func (stmt commonStmt) expectedTypeName() string {
	return fmt.Sprintf("%s.%s", stmt.rowType.PkgPath(), stmt.rowType.Name())
}
<|endoftext|>


<commit_before>package sxchange

import (
	"context"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/hex"
	"errors"
	"fmt"
	"hash/crc32"
	"io"
	"log"
	"net"
	"sync"
	"time"
	"unsafe"
)

const (
	verHI = 1
	verLO = 1
)

var (
	crc32q          = crc32.MakeTable(crc32.Koopman)
	errPartialWrite = errors.New("Partial write")
	errPartialRead  = errors.New("Partial read")
)

// CB type for callback function for each type of transfered data
type CB func([]byte, *Connection)

// DataTypeCB describes
type DataTypeCB struct {
	SizeBytes int8  // -1 for fixed size, 0 for no data (ping?), 1 for 0-255 bytes, 2 for 0-65535 bytes...
	FixedSize int32 // if some struct has fixed fize no other header needed
	Callback  CB    // which function to run after data received
}

// Connection informaion both for client or server
type Connection struct {
	conn         *net.TCPConn         // real connection
	AESKey       string               // AES-128/192/256 key in hex for connecton encryption
	eWriter      *cipher.StreamWriter // just StreamWriter object to simplify encryption
	eReader      *cipher.StreamReader // the same for reader
	writeMutex   sync.Mutex           // we need to lock on write operations
	Types        map[uint8]DataTypeCB // map of type->params&callback
	KeepAlive    time.Duration        // tcp keep-alive
	ReadTimeout  time.Duration        // we need to receive new message at least once per this duration
	WriteTimeout time.Duration        // maximal duration for every write operation
	MaxSize      uint32               // maximum number of bytes for one record/message
	Ctx          context.Context      // context for some values like server id
	CloseChan    chan interface{}     // channel for closing incoming messages circle
}

const (
	initialPacketSize = 1292 // we need to compare it manualy to be sure that compiler doesn't change aligment!
)

type initialPacket struct {
	Header    [8]uint8 // sxchngXY where X.Y is version of protocol
	MaxSize   uint32
	SizeBytes [256]int8
	FixedSize [256]int32
}

func (c *Connection) readAll(buf []byte, size int, ne bool) error {

	totalDeadline := time.Now().Add(c.ReadTimeout)
	c.conn.SetReadDeadline(totalDeadline)

	var (
		i   int
		err error
	)

	if ne || nil == c.eReader {
		i, err = io.ReadFull(c.conn, buf[:size])
	} else {
		i, err = io.ReadFull(c.eReader, buf[:size])
	}

	if nil != err {
		return err
	}
	if i != size {
		return errPartialRead
	}

	return nil
}

func (c *Connection) writeAll(buf []byte, size int, ne bool) error {
	totalDeadline := time.Now().Add(c.WriteTimeout)
	c.conn.SetWriteDeadline(totalDeadline)

	var (
		i   int
		err error
	)

	if ne || nil == c.eWriter {
		i, err = c.conn.Write(buf[:size])
	} else {
		i, err = c.eWriter.Write(buf[:size])
	}

	if nil != err {
		return err
	}
	if i != size {
		return errPartialWrite
	}

	return nil
}

func (c *Connection) encSetup() error {
	if "" != c.AESKey {
		aesKey, err := hex.DecodeString(c.AESKey)
		if err != nil {
			return err
		}

		block, err := aes.NewCipher(aesKey)
		if err != nil {
			return err
		}

		var iv [aes.BlockSize]byte
		if _, err := io.ReadFull(rand.Reader, iv[:]); nil != err {
			return err
		}

		c.conn.SetWriteDeadline(time.Now().Add(c.WriteTimeout))
		i, err := c.conn.Write(iv[:])
		if nil != err {
			return err
		}
		if i != len(iv) {
			return errPartialWrite
		}

		c.eWriter = &cipher.StreamWriter{S: cipher.NewCTR(block, iv[:]), W: c.conn}

		c.conn.SetReadDeadline(time.Now().Add(c.ReadTimeout))

		var ivread [aes.BlockSize]byte

		i, err = io.ReadFull(c.conn, ivread[:])
		if nil != err {
			return err
		}
		if i != len(ivread) {
			return errPartialRead
		}

		c.eReader = &cipher.StreamReader{S: cipher.NewCTR(block, ivread[:]), R: c.conn}
	}

	return nil
}

func (c *Connection) initConnection() error {
	if c.KeepAlive != 0 {
		c.conn.SetKeepAlivePeriod(c.KeepAlive)
		c.conn.SetKeepAlive(true)
	}

	// initial packets exchange
	initial := c.prepareInitialPacket()
	initialBuf := (*[initialPacketSize]byte)(unsafe.Pointer(&initial))[:]
	received := initialPacket{}
	receivedBuf := (*[initialPacketSize]byte)(unsafe.Pointer(&received))[:]

	// send out version
	err := c.writeAll(initialBuf, initialPacketSize, true)
	if nil != err {
		return err
	}

	err = c.readAll(receivedBuf, initialPacketSize, true)
	if nil != err {
		return err
	}

	if received.Header != initial.Header {
		return errors.New("HS error - different protocol version")
	}

	if received.MaxSize != initial.MaxSize {
		return errors.New("HS error - different message maximal size")
	}

	if (received.FixedSize != initial.FixedSize) || (received.MaxSize != received.MaxSize) {
		return errors.New("HS error - different set of datatypes")
	}

	err = c.encSetup()
	if nil != err {
		return err
	}

	return nil
}

// internal run function
func (c *Connection) run() error {

	defer c.conn.Close()

	var (
		sbuf [4]byte
		t    [1]byte
		msg  = make([]byte, 0, c.MaxSize)
		crcI uint32
		crcB = (*[4]byte)(unsafe.Pointer(&crcI))
	)

	for {

		select {
		case <-c.CloseChan:
			return nil
		default:
		}

		err := c.readAll(t[:], 1, false)
		if nil != err {
			return err
		}

		dt, ok := c.Types[t[0]]

		if !ok {
			return fmt.Errorf("Structure type %d is not defined", t[0])
		}

		var size2read int
		if dt.SizeBytes > 0 {
			err = c.readAll(sbuf[0:dt.SizeBytes], int(dt.SizeBytes), false)
			if nil != err {
				return err
			}
		}

		switch dt.SizeBytes {
		case -1:
			size2read = int(dt.FixedSize)
		case 0:
			size2read = 0
		case 1:
			size2read = int(sbuf[0])
		case 2:
			size2read = int(sbuf[0]) | (int(sbuf[1]) << 8)
		case 3:
			size2read = int(sbuf[0]) | (int(sbuf[1]) << 8) | (int(sbuf[2]) << 16)
		default:
			return fmt.Errorf("Structure type %d has invalid SizeBytes setting (%d)", t[0], dt.SizeBytes)
		}

		if size2read > int(c.MaxSize) {
			return fmt.Errorf("Size of message is greater than maximum (%d/%d)", size2read, c.MaxSize)
		}

		err = c.readAll(msg, size2read, false)
		if nil != err {
			return err
		}

		// one more check on done channel - we don't want to call callback after it has been closed
		select {
		case <-c.CloseChan:
			return nil
		default:
		}

		if 0 != size2read {
			err = c.readAll((*crcB)[:], 4, false)
			if nil != err {
				return err
			}
			if crc32.Checksum(msg[0:size2read], crc32q) != crcI {
				return errors.New("crc failed on read")
			}
		}

		// we need to allow types that only can be send by one of the side, but also defined
		if nil != dt.Callback {
			dt.Callback(msg[0:size2read], c)
		}
	}
}

// WriteMsg creates header and then writes msg buffer via TCP connection
func (c *Connection) WriteMsg(msgType uint8, msg []byte) error {

	dt, ok := c.Types[msgType]
	if !ok {
		return fmt.Errorf("Unknown msg type %d", msgType)
	}

	var (
		crcI uint32
		crcB = (*[4]byte)(unsafe.Pointer(&crcI))
	)

	header := [5]byte{}
	header[0] = msgType

	size2write := 0
	headerSize := 1

	switch dt.SizeBytes {
	case -1:
		size2write = int(dt.FixedSize)
	case 0:
		size2write = 0
	case 1:
		if len(msg) > 255 {
			return fmt.Errorf("Structure type %d has 1-byte size header but %d bytes givven", msgType, len(msg))
		}
		header[1] = byte(len(msg))
		size2write = len(msg)
		headerSize = 2
	case 2:
		if len(msg) > 255*255 {
			return fmt.Errorf("Structure type %d has 2-byte size header but %d bytes givven", msgType, len(msg))
		}
		header[1] = byte(len(msg) & 0xff)
		header[2] = byte((len(msg) & 0xff00) >> 8)
		size2write = len(msg)
		headerSize = 3

	case 3:
		if len(msg) > 255*255*255 {
			return fmt.Errorf("Structure type %d has 2-byte size header but %d bytes givven", msgType, len(msg))
		}
		header[1] = byte(len(msg) & 0xff)
		header[2] = byte((len(msg) & 0xff00) >> 8)
		header[3] = byte((len(msg) & 0xff0000) >> 16)
		size2write = len(msg)
		headerSize = 4

	default:
		return fmt.Errorf("Structure type %d has invalid SizeBytes setting (%d)", msgType, dt.SizeBytes)
	}

	if size2write > len(msg) {
		return fmt.Errorf("Unable to write %d bytes for structure %d, have only %d in buffer", msgType, size2write, len(msg))
	}

	c.writeMutex.Lock()
	defer c.writeMutex.Unlock()

	err := c.writeAll(header[:], headerSize, false)
	if nil != err {
		return err
	}

	err = c.writeAll(msg, size2write, false)
	if nil != err {
		return err
	}

	if size2write > 0 {
		crcI = crc32.Checksum(msg[0:size2write], crc32q)
		err = c.writeAll((*crcB)[:], 4, false)
		if nil != err {
			return err
		}
	}

	return nil
}

// Remote address retrive
func (c *Connection) Remote() net.Addr {
	if nil == c.conn {
		return nil
	}

	return c.conn.RemoteAddr()
}

func (c *Connection) prepareInitialPacket() initialPacket {
	result := initialPacket{MaxSize: c.MaxSize, Header: [8]byte{'s', 'x', 'c', 'h', 'n', 'g', verHI, verLO}}
	if "" != c.AESKey {
		result.Header[4] = 'E'
	}

	for i := 0; i < 256; i++ {
		result.FixedSize[i] = c.Types[uint8(i)].FixedSize
		result.SizeBytes[i] = c.Types[uint8(i)].SizeBytes
	}

	return result
}

func init() {
	if initialPacketSize != unsafe.Sizeof(initialPacket{}) {
		log.Fatalln("Golang uses different aligment withing structure, please modify sxchnge code!")
	}
}
<commit_msg>Add panic recover in connection.run<commit_after>package sxchange

import (
	"context"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/hex"
	"errors"
	"fmt"
	"hash/crc32"
	"io"
	"log"
	"net"
	"sync"
	"time"
	"unsafe"
)

const (
	verHI = 1
	verLO = 1
)

var (
	crc32q          = crc32.MakeTable(crc32.Koopman)
	errPartialWrite = errors.New("Partial write")
	errPartialRead  = errors.New("Partial read")
)

// CB type for callback function for each type of transfered data
type CB func([]byte, *Connection)

// DataTypeCB describes
type DataTypeCB struct {
	SizeBytes int8  // -1 for fixed size, 0 for no data (ping?), 1 for 0-255 bytes, 2 for 0-65535 bytes...
	FixedSize int32 // if some struct has fixed fize no other header needed
	Callback  CB    // which function to run after data received
}

// Connection informaion both for client or server
type Connection struct {
	conn         *net.TCPConn         // real connection
	AESKey       string               // AES-128/192/256 key in hex for connecton encryption
	eWriter      *cipher.StreamWriter // just StreamWriter object to simplify encryption
	eReader      *cipher.StreamReader // the same for reader
	writeMutex   sync.Mutex           // we need to lock on write operations
	Types        map[uint8]DataTypeCB // map of type->params&callback
	KeepAlive    time.Duration        // tcp keep-alive
	ReadTimeout  time.Duration        // we need to receive new message at least once per this duration
	WriteTimeout time.Duration        // maximal duration for every write operation
	MaxSize      uint32               // maximum number of bytes for one record/message
	Ctx          context.Context      // context for some values like server id
	CloseChan    chan interface{}     // channel for closing incoming messages circle
}

const (
	initialPacketSize = 1292 // we need to compare it manualy to be sure that compiler doesn't change aligment!
)

type initialPacket struct {
	Header    [8]uint8 // sxchngXY where X.Y is version of protocol
	MaxSize   uint32
	SizeBytes [256]int8
	FixedSize [256]int32
}

func (c *Connection) readAll(buf []byte, size int, ne bool) error {

	totalDeadline := time.Now().Add(c.ReadTimeout)
	c.conn.SetReadDeadline(totalDeadline)

	var (
		i   int
		err error
	)

	if ne || nil == c.eReader {
		i, err = io.ReadFull(c.conn, buf[:size])
	} else {
		i, err = io.ReadFull(c.eReader, buf[:size])
	}

	if nil != err {
		return err
	}
	if i != size {
		return errPartialRead
	}

	return nil
}

func (c *Connection) writeAll(buf []byte, size int, ne bool) error {
	totalDeadline := time.Now().Add(c.WriteTimeout)
	c.conn.SetWriteDeadline(totalDeadline)

	var (
		i   int
		err error
	)

	if ne || nil == c.eWriter {
		i, err = c.conn.Write(buf[:size])
	} else {
		i, err = c.eWriter.Write(buf[:size])
	}

	if nil != err {
		return err
	}
	if i != size {
		return errPartialWrite
	}

	return nil
}

func (c *Connection) encSetup() error {
	if "" != c.AESKey {
		aesKey, err := hex.DecodeString(c.AESKey)
		if err != nil {
			return err
		}

		block, err := aes.NewCipher(aesKey)
		if err != nil {
			return err
		}

		var iv [aes.BlockSize]byte
		if _, err := io.ReadFull(rand.Reader, iv[:]); nil != err {
			return err
		}

		c.conn.SetWriteDeadline(time.Now().Add(c.WriteTimeout))
		i, err := c.conn.Write(iv[:])
		if nil != err {
			return err
		}
		if i != len(iv) {
			return errPartialWrite
		}

		c.eWriter = &cipher.StreamWriter{S: cipher.NewCTR(block, iv[:]), W: c.conn}

		c.conn.SetReadDeadline(time.Now().Add(c.ReadTimeout))

		var ivread [aes.BlockSize]byte

		i, err = io.ReadFull(c.conn, ivread[:])
		if nil != err {
			return err
		}
		if i != len(ivread) {
			return errPartialRead
		}

		c.eReader = &cipher.StreamReader{S: cipher.NewCTR(block, ivread[:]), R: c.conn}
	}

	return nil
}

func (c *Connection) initConnection() error {
	if c.KeepAlive != 0 {
		c.conn.SetKeepAlivePeriod(c.KeepAlive)
		c.conn.SetKeepAlive(true)
	}

	// initial packets exchange
	initial := c.prepareInitialPacket()
	initialBuf := (*[initialPacketSize]byte)(unsafe.Pointer(&initial))[:]
	received := initialPacket{}
	receivedBuf := (*[initialPacketSize]byte)(unsafe.Pointer(&received))[:]

	// send out version
	err := c.writeAll(initialBuf, initialPacketSize, true)
	if nil != err {
		return err
	}

	err = c.readAll(receivedBuf, initialPacketSize, true)
	if nil != err {
		return err
	}

	if received.Header != initial.Header {
		return errors.New("HS error - different protocol version")
	}

	if received.MaxSize != initial.MaxSize {
		return errors.New("HS error - different message maximal size")
	}

	if (received.FixedSize != initial.FixedSize) || (received.MaxSize != received.MaxSize) {
		return errors.New("HS error - different set of datatypes")
	}

	err = c.encSetup()
	if nil != err {
		return err
	}

	return nil
}

// internal run function
func (c *Connection) run() error {

	defer c.conn.Close()

	defer func() {
		if r := recover(); r != nil {
			fmt.Println("Recovered in sxconnection.run", r)
		}
	}()

	var (
		sbuf [4]byte
		t    [1]byte
		msg  = make([]byte, 0, c.MaxSize)
		crcI uint32
		crcB = (*[4]byte)(unsafe.Pointer(&crcI))
	)

	for {

		select {
		case <-c.CloseChan:
			return nil
		default:
		}

		err := c.readAll(t[:], 1, false)
		if nil != err {
			return err
		}

		dt, ok := c.Types[t[0]]

		if !ok {
			return fmt.Errorf("Structure type %d is not defined", t[0])
		}

		var size2read int
		if dt.SizeBytes > 0 {
			err = c.readAll(sbuf[0:dt.SizeBytes], int(dt.SizeBytes), false)
			if nil != err {
				return err
			}
		}

		switch dt.SizeBytes {
		case -1:
			size2read = int(dt.FixedSize)
		case 0:
			size2read = 0
		case 1:
			size2read = int(sbuf[0])
		case 2:
			size2read = int(sbuf[0]) | (int(sbuf[1]) << 8)
		case 3:
			size2read = int(sbuf[0]) | (int(sbuf[1]) << 8) | (int(sbuf[2]) << 16)
		default:
			return fmt.Errorf("Structure type %d has invalid SizeBytes setting (%d)", t[0], dt.SizeBytes)
		}

		if size2read > int(c.MaxSize) {
			return fmt.Errorf("Size of message is greater than maximum (%d/%d)", size2read, c.MaxSize)
		}

		err = c.readAll(msg, size2read, false)
		if nil != err {
			return err
		}

		// one more check on done channel - we don't want to call callback after it has been closed
		select {
		case <-c.CloseChan:
			return nil
		default:
		}

		if 0 != size2read {
			err = c.readAll((*crcB)[:], 4, false)
			if nil != err {
				return err
			}
			if crc32.Checksum(msg[0:size2read], crc32q) != crcI {
				return errors.New("crc failed on read")
			}
		}

		// we need to allow types that only can be send by one of the side, but also defined
		if nil != dt.Callback {
			dt.Callback(msg[0:size2read], c)
		}
	}
}

// WriteMsg creates header and then writes msg buffer via TCP connection
func (c *Connection) WriteMsg(msgType uint8, msg []byte) error {

	dt, ok := c.Types[msgType]
	if !ok {
		return fmt.Errorf("Unknown msg type %d", msgType)
	}

	var (
		crcI uint32
		crcB = (*[4]byte)(unsafe.Pointer(&crcI))
	)

	header := [5]byte{}
	header[0] = msgType

	size2write := 0
	headerSize := 1

	switch dt.SizeBytes {
	case -1:
		size2write = int(dt.FixedSize)
	case 0:
		size2write = 0
	case 1:
		if len(msg) > 255 {
			return fmt.Errorf("Structure type %d has 1-byte size header but %d bytes givven", msgType, len(msg))
		}
		header[1] = byte(len(msg))
		size2write = len(msg)
		headerSize = 2
	case 2:
		if len(msg) > 255*255 {
			return fmt.Errorf("Structure type %d has 2-byte size header but %d bytes givven", msgType, len(msg))
		}
		header[1] = byte(len(msg) & 0xff)
		header[2] = byte((len(msg) & 0xff00) >> 8)
		size2write = len(msg)
		headerSize = 3

	case 3:
		if len(msg) > 255*255*255 {
			return fmt.Errorf("Structure type %d has 2-byte size header but %d bytes givven", msgType, len(msg))
		}
		header[1] = byte(len(msg) & 0xff)
		header[2] = byte((len(msg) & 0xff00) >> 8)
		header[3] = byte((len(msg) & 0xff0000) >> 16)
		size2write = len(msg)
		headerSize = 4

	default:
		return fmt.Errorf("Structure type %d has invalid SizeBytes setting (%d)", msgType, dt.SizeBytes)
	}

	if size2write > len(msg) {
		return fmt.Errorf("Unable to write %d bytes for structure %d, have only %d in buffer", msgType, size2write, len(msg))
	}

	c.writeMutex.Lock()
	defer c.writeMutex.Unlock()

	err := c.writeAll(header[:], headerSize, false)
	if nil != err {
		return err
	}

	err = c.writeAll(msg, size2write, false)
	if nil != err {
		return err
	}

	if size2write > 0 {
		crcI = crc32.Checksum(msg[0:size2write], crc32q)
		err = c.writeAll((*crcB)[:], 4, false)
		if nil != err {
			return err
		}
	}

	return nil
}

// Remote address retrive
func (c *Connection) Remote() net.Addr {
	if nil == c.conn {
		return nil
	}

	return c.conn.RemoteAddr()
}

func (c *Connection) prepareInitialPacket() initialPacket {
	result := initialPacket{MaxSize: c.MaxSize, Header: [8]byte{'s', 'x', 'c', 'h', 'n', 'g', verHI, verLO}}
	if "" != c.AESKey {
		result.Header[4] = 'E'
	}

	for i := 0; i < 256; i++ {
		result.FixedSize[i] = c.Types[uint8(i)].FixedSize
		result.SizeBytes[i] = c.Types[uint8(i)].SizeBytes
	}

	return result
}

func init() {
	if initialPacketSize != unsafe.Sizeof(initialPacket{}) {
		log.Fatalln("Golang uses different aligment withing structure, please modify sxchnge code!")
	}
}
<|endoftext|>


<commit_before>package goaws

import (
	"log"

	"github.com/Skarlso/go_aws_mine/config"
	"github.com/Skarlso/go_aws_mine/errorhandler"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/ec2"
)

// CreateEC2 testing AWS configuration.
func CreateEC2(ec2Config *config.EC2Config) {
	log.Println("Creating ec2 session.")
	sess := session.New(&aws.Config{Region: aws.String("eu-central-1")})
	ec2Client := ec2.New(sess, nil)
	runResult, err := ec2Client.RunInstances(&ec2.RunInstancesInput{
		ImageId:      &ec2Config.ImageID,
		DryRun:       &ec2Config.DryRun,
		MaxCount:     &ec2Config.MaxCount,
		MinCount:     &ec2Config.MinCount,
		InstanceType: &ec2Config.InstanceType,
		KeyName:      &ec2Config.KeyName,
		Monitoring:   &ec2.RunInstancesMonitoringEnabled{Enabled: &ec2Config.Monitoring.Enable},
	})
	log.Println("Instance created with id: ", *runResult.Instances[0].InstanceId)
	errorhandler.CheckError(err)
	ec2Id := aws.StringSlice([]string{*runResult.Instances[0].InstanceId})
	err = ec2Client.WaitUntilInstanceRunning(&ec2.DescribeInstancesInput{InstanceIds: ec2Id})
	if err != nil {
		errorhandler.CheckError(err)
	}
}
<commit_msg>Switched error handling<commit_after>package goaws

import (
	"log"

	"github.com/Skarlso/go_aws_mine/config"
	"github.com/Skarlso/go_aws_mine/errorhandler"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/ec2"
)

// CreateEC2 testing AWS configuration.
func CreateEC2(ec2Config *config.EC2Config) {
	log.Println("Creating ec2 session.")
	sess := session.New(&aws.Config{Region: aws.String("eu-central-1")})
	ec2Client := ec2.New(sess, nil)
	runResult, err := ec2Client.RunInstances(&ec2.RunInstancesInput{
		ImageId:      &ec2Config.ImageID,
		DryRun:       &ec2Config.DryRun,
		MaxCount:     &ec2Config.MaxCount,
		MinCount:     &ec2Config.MinCount,
		InstanceType: &ec2Config.InstanceType,
		KeyName:      &ec2Config.KeyName,
		Monitoring:   &ec2.RunInstancesMonitoringEnabled{Enabled: &ec2Config.Monitoring.Enable},
	})
	errorhandler.CheckError(err)
	log.Println("Instance created with id: ", *runResult.Instances[0].InstanceId)
	ec2Id := aws.StringSlice([]string{*runResult.Instances[0].InstanceId})
	err = ec2Client.WaitUntilInstanceRunning(&ec2.DescribeInstancesInput{InstanceIds: ec2Id})
	if err != nil {
		errorhandler.CheckError(err)
	}
}
<|endoftext|>


<commit_before>// (c) 2011-2013 Alexander Solovyov
// under terms of ISC license

package main

import (
	"bytes"
	"fmt"
	flags "github.com/jessevdk/go-flags"
	"github.com/wsxiaoys/terminal/color"
	"os"
	"path/filepath"
	"regexp"
)

var (
	Author  = "Alexander Solovyov"
	Version = "1.3"

	byteNewLine = []byte("\n")
)

var opts struct {
	IgnoreCase      bool     `short:"i" long:"ignore-case" description:"ignore pattern case"`
	OnlyName        bool     `short:"n" long:"filename" description:"print only filenames"`
	FindFiles       bool     `short:"f" long:"find-files" description:"search for files and not for text in them"`
	IgnoreFiles     []string `short:"x" long:"exclude" description:"exclude files that match the regexp from search" value-name:"RE"`
	AcceptFiles     []string `short:"o" long:"only" description:"search only in files that match the regexp" value-name:"RE"`
	SingleLine      bool     `short:"s" long:"singleline" description:"match on a single line (^/$ will be begginning/end of line)"`
	PlainText       bool     `short:"p" long:"plain" description:"search plain text"`
	Replace         *string  `short:"r" long:"replace" description:"replace found substrings with this string"`
	NoGlobalIgnores bool     `short:"I" long:"no-autoignore" description:"do not read .git/.hgignore files"`
	Force           bool     `long:"force" description:"force replacement in binary files"`
	Verbose         bool     `short:"v" long:"verbose" description:"be verbose (show non-fatal errors, like unreadable files)"`
	ShowVersion     bool     `short:"V" long:"version" description:"show version and exit"`
	ShowHelp        bool     `long:"help" description:"show this help message"`
	ShowColors		bool	 `short:"c" long:"colors" description:"show colors in output"`
}

func main() {
	argparser := flags.NewParser(&opts, flags.PrintErrors|flags.PassDoubleDash)

	args, err := argparser.Parse()
	if err != nil {
		return
	}

	if opts.ShowVersion {
		fmt.Printf("goreplace %s\n", Version)
		return
	}

	var noIgnores bool
	for _, item := range os.Args[1:] {
		if item == "-I" || item == "--no-autoignore" {
			noIgnores = true
		}
	}

	cwd, _ := os.Getwd()

	ignoreFileMatcher := NewMatcher(cwd, noIgnores)
	ignoreFileMatcher.Append(opts.IgnoreFiles)

	acceptedFileMatcher := NewGeneralMatcher([]string{}, []string{})
	if len(opts.AcceptFiles) > 0 {
		acceptedFileMatcher.Append(opts.AcceptFiles)
	} else {
		acceptedFileMatcher.Append([]string{".*"})
	}

	argparser.Usage = fmt.Sprintf("[OPTIONS] string-to-search\n\n%s",
		ignoreFileMatcher)

	if opts.ShowHelp || len(args) == 0 {
		argparser.WriteHelp(os.Stdout)
		return
	}

	arg := args[0]
	if opts.PlainText {
		arg = regexp.QuoteMeta(arg)
	}
	if opts.IgnoreCase {
		arg = "(?i:" + arg + ")"
	}

	pattern, err := regexp.Compile(arg)
	errhandle(err, true, "")

	if pattern.Match([]byte("")) {
		errhandle(fmt.Errorf("Your pattern matches empty string"), true, "")
	}

	searchFiles(pattern, ignoreFileMatcher, acceptedFileMatcher)
}

func errhandle(err error, exit bool, moreinfo string, a ...interface{}) bool {
	if err == nil {
		return false
	}
	fmt.Fprintf(os.Stderr, "%s\n%s\n", err, fmt.Sprintf(moreinfo, a...))
	if exit {
		os.Exit(1)
	}
	return true
}

func searchFiles(pattern *regexp.Regexp, ignoreFileMatcher Matcher,
	acceptedFileMatcher Matcher) {
	v := &GRVisitor{pattern, ignoreFileMatcher, acceptedFileMatcher, false}

	errors := make(chan error, 64)

	filepath.Walk(".", v.Walker(errors))

	select {
	case err := <-errors:
		if opts.Verbose {
			errhandle(err, false, "")
		}
	default:
	}
}

type GRVisitor struct {
	pattern             *regexp.Regexp
	ignoreFileMatcher   Matcher
	acceptedFileMatcher Matcher
	// Used to prevent sparse newline at the end of output
	prependNewLine bool
}

func (v *GRVisitor) Walker(errors chan<- error) filepath.WalkFunc {
	return func(fn string, fi os.FileInfo, err error) error {
		if err != nil {
			errors <- err
			return nil
		}

		// NOTE: if a directory is a symlink, filepath.Walk won't recurse inside
		if fi.Mode()&os.ModeSymlink != 0 {
			if fi, err = os.Stat(fn); err != nil {
				errors <- err
				return nil
			}
		}

		if fi.IsDir() {
			if !v.VisitDir(fn, fi) {
				return filepath.SkipDir
			}
			return nil
		}

		v.VisitFile(fn, fi)
		return nil
	}
}

func (v *GRVisitor) VisitDir(fn string, fi os.FileInfo) bool {
	return !v.ignoreFileMatcher.Match(fn, true)
}

func (v *GRVisitor) VisitFile(fn string, fi os.FileInfo) {
	if fi.Size() == 0 && !opts.FindFiles {
		return
	}

	if v.ignoreFileMatcher.Match(fn, false) {
		return
	}

	if !v.acceptedFileMatcher.Match(fn, false) {
		return
	}

	if fi.Size() >= 1024*1024*10 {
		fmt.Fprintf(os.Stderr, "Skipping %s, too big: %d\n", fn, fi.Size())
		return
	}

	if opts.FindFiles {
		v.SearchFileName(fn)
		return
	}

	f, content := v.GetFileAndContent(fn, fi)
	if f == nil {
		return
	}
	defer f.Close()

	if opts.Replace == nil {
		v.SearchFile(fn, content)
		return
	}

	changed, result := v.ReplaceInFile(fn, content)
	if changed {
		f.Seek(0, 0)
		n, err := f.Write(result)
		errhandle(err, true, "Error writing replacement in file %s", fn)
		if int64(n) < fi.Size() {
			err := f.Truncate(int64(n))
			errhandle(err, true, "Error truncating file to size %d", f)
		}
	}
}

func (v *GRVisitor) GetFileAndContent(fn string, fi os.FileInfo) (f *os.File, content []byte) {
	var err error
	var msg string

	if opts.Replace != nil {
		f, err = os.OpenFile(fn, os.O_RDWR, 0666)
		msg = "can't open file %s for reading and writing"
	} else {
		f, err = os.Open(fn)
		msg = "can't open file %s for reading"
	}

	if err != nil {
		if opts.Verbose {
			errhandle(err, false, msg, fn)
		}
		return
	}

	content = make([]byte, fi.Size())
	n, err := f.Read(content)
	errhandle(err, true, "can't read file %s", fn)
	if int64(n) != fi.Size() {
		errhandle(fmt.Errorf("Not whole file was read, only %d from %d",
			n, fi.Size()), true, "")
	}

	return
}

func (v *GRVisitor) SearchFile(fn string, content []byte) {
	lines := IntList([]int{})
	binary := false

	if bytes.IndexByte(content, 0) != -1 {
		binary = true
	}

	for _, info := range v.FindAllIndex(content) {
		if lines.Contains(info.num) {
			continue
		}

		if v.prependNewLine {
			fmt.Println("")
			v.prependNewLine = false
		}

		first := len(lines) == 0
		lines = append(lines, info.num)

		if first {
			if binary && !opts.OnlyName {
				fmt.Printf("Binary file %s matches\n", fn)
				break
			} else {
				if (opts.ShowColors) {
					color.Printf("@g%s\n", fn)
				} else {
					fmt.Printf("%s\n",fn)
				}
			}
		}

		if opts.OnlyName {
			return
		}

		if (opts.ShowColors) {
			color.Printf("@!@y%d:", info.num)
		} else {
			fmt.Printf("%d:", info.num)
		}
		colored := v.pattern.ReplaceAllStringFunc(string(info.line),
			func(wrap string) string {
				var res string
				if (opts.ShowColors) {
					res = color.Sprintf("@Y%s", wrap)
				} else {
					res = fmt.Sprintf("%s", wrap)
				}
				return res
			})
		fmt.Println(colored)
	}

	if len(lines) > 0 {
		v.prependNewLine = true
	}
}

func (v *GRVisitor) SearchFileName(fn string) {
	if !v.pattern.MatchString(fn) {
		return
	}
	colored := v.pattern.ReplaceAllStringFunc(fn,
		func(wrap string) string {
			var res string
			if (opts.ShowColors) {
				res = color.Sprintf("@Y%s", wrap)
			} else {
				res = fmt.Sprintf("%s", wrap)
			}
			return res
		})
	fmt.Println(colored)
}

func getSuffix(num int) string {
	if num > 1 {
		return "s"
	}
	return ""
}

func (v *GRVisitor) ReplaceInFile(fn string, content []byte) (changed bool, result []byte) {
	changed = false
	binary := false
	changenum := 0

	if opts.SingleLine {
		errhandle(
			fmt.Errorf("Can't handle singleline replacements yet"),
			true, "")
	}

	if opts.PlainText {
		errhandle(
			fmt.Errorf("Can't handle plain text replacements yet"),
			true, "")
	}

	if bytes.IndexByte(content, 0) != -1 {
		binary = true
	}

	result = v.pattern.ReplaceAllFunc(content, func(s []byte) []byte {
		if binary && !opts.Force {
			errhandle(
				fmt.Errorf("supply --force to force change of binary file"),
				false, "")
		}
		if !changed {
			changed = true
			if (opts.ShowColors) {
				color.Printf("@g%s", fn)
			} else {
				fmt.Printf("%s", fn)
			}
		}

		changenum += 1
		return []byte(*opts.Replace)
	})

	if changenum > 0 {
		if (opts.ShowColors) {
			color.Printf("@!@y - %d change%s made\n",
				changenum, getSuffix(changenum))
		} else {
			fmt.Printf(" - %d change%s made\n",
				changenum, getSuffix(changenum))
		}
	}

	return changed, result
}

type LineInfo struct {
	num  int
	line []byte
}

func (v *GRVisitor) FindAllIndex(content []byte) (res []*LineInfo) {
	if opts.SingleLine {
		return v.singlelineFindAllIndex(content)
	}

	linenum, last := 1, 0
	for _, bounds := range v.pattern.FindAllIndex(content, -1) {
		linenum += bytes.Count(content[last:bounds[0]], byteNewLine)
		last = bounds[0]
		begin, end := beginend(content, bounds[0], bounds[1])
		res = append(res, &LineInfo{linenum, content[begin:end]})
	}
	return res
}

func (v *GRVisitor) singlelineFindAllIndex(content []byte) (res []*LineInfo) {
	linenum, begin, end := 1, 0, 0
	for i := 0; i < len(content); i++ {
		if content[i] == '\n' {
			end = i
			line := content[begin:end]
			if v.pattern.Match(line) {
				res = append(res, &LineInfo{linenum, line})
			}
			linenum += 1
			begin = end + 1
		}
	}
	return res
}

// Given a []byte, start and finish of some inner slice, will find nearest
// newlines on both ends of this slice to contain this slice
func beginend(s []byte, start int, finish int) (begin int, end int) {
	begin = 0
	end = len(s)

	for i := start - 1; i >= 0; i-- {
		if s[i] == byteNewLine[0] {
			begin = i + 1
			break
		}
	}

	// -1 to check if current location is the end of a string
	for i := finish - 1; i < len(s); i++ {
		if s[i] == byteNewLine[0] {
			end = i
			break
		}
	}

	return begin, end
}

type IntList []int

func (il IntList) Contains(i int) bool {
	for _, x := range il {
		if x == i {
			return true
		}
	}
	return false
}
<commit_msg>Rename colors flag.<commit_after>// (c) 2011-2013 Alexander Solovyov
// under terms of ISC license

package main

import (
	"bytes"
	"fmt"
	flags "github.com/jessevdk/go-flags"
	"github.com/wsxiaoys/terminal/color"
	"os"
	"path/filepath"
	"regexp"
)

var (
	Author  = "Alexander Solovyov"
	Version = "1.3"

	byteNewLine = []byte("\n")
)

var opts struct {
	IgnoreCase      bool     `short:"i" long:"ignore-case" description:"ignore pattern case"`
	OnlyName        bool     `short:"n" long:"filename" description:"print only filenames"`
	FindFiles       bool     `short:"f" long:"find-files" description:"search for files and not for text in them"`
	IgnoreFiles     []string `short:"x" long:"exclude" description:"exclude files that match the regexp from search" value-name:"RE"`
	AcceptFiles     []string `short:"o" long:"only" description:"search only in files that match the regexp" value-name:"RE"`
	SingleLine      bool     `short:"s" long:"singleline" description:"match on a single line (^/$ will be begginning/end of line)"`
	PlainText       bool     `short:"p" long:"plain" description:"search plain text"`
	Replace         *string  `short:"r" long:"replace" description:"replace found substrings with this string"`
	NoGlobalIgnores bool     `short:"I" long:"no-autoignore" description:"do not read .git/.hgignore files"`
	Force           bool     `long:"force" description:"force replacement in binary files"`
	Verbose         bool     `short:"v" long:"verbose" description:"be verbose (show non-fatal errors, like unreadable files)"`
	ShowVersion     bool     `short:"V" long:"version" description:"show version and exit"`
	ShowHelp        bool     `long:"help" description:"show this help message"`
	NoColors		bool	 `short:"c" long:"no-colors" description:"do not show colors in output"`
}

func main() {
	argparser := flags.NewParser(&opts, flags.PrintErrors|flags.PassDoubleDash)

	args, err := argparser.Parse()
	if err != nil {
		return
	}

	if opts.ShowVersion {
		fmt.Printf("goreplace %s\n", Version)
		return
	}

	var noIgnores bool
	for _, item := range os.Args[1:] {
		if item == "-I" || item == "--no-autoignore" {
			noIgnores = true
		}
	}

	cwd, _ := os.Getwd()

	ignoreFileMatcher := NewMatcher(cwd, noIgnores)
	ignoreFileMatcher.Append(opts.IgnoreFiles)

	acceptedFileMatcher := NewGeneralMatcher([]string{}, []string{})
	if len(opts.AcceptFiles) > 0 {
		acceptedFileMatcher.Append(opts.AcceptFiles)
	} else {
		acceptedFileMatcher.Append([]string{".*"})
	}

	argparser.Usage = fmt.Sprintf("[OPTIONS] string-to-search\n\n%s",
		ignoreFileMatcher)

	if opts.ShowHelp || len(args) == 0 {
		argparser.WriteHelp(os.Stdout)
		return
	}

	arg := args[0]
	if opts.PlainText {
		arg = regexp.QuoteMeta(arg)
	}
	if opts.IgnoreCase {
		arg = "(?i:" + arg + ")"
	}

	pattern, err := regexp.Compile(arg)
	errhandle(err, true, "")

	if pattern.Match([]byte("")) {
		errhandle(fmt.Errorf("Your pattern matches empty string"), true, "")
	}

	searchFiles(pattern, ignoreFileMatcher, acceptedFileMatcher)
}

func errhandle(err error, exit bool, moreinfo string, a ...interface{}) bool {
	if err == nil {
		return false
	}
	fmt.Fprintf(os.Stderr, "%s\n%s\n", err, fmt.Sprintf(moreinfo, a...))
	if exit {
		os.Exit(1)
	}
	return true
}

func searchFiles(pattern *regexp.Regexp, ignoreFileMatcher Matcher,
	acceptedFileMatcher Matcher) {
	v := &GRVisitor{pattern, ignoreFileMatcher, acceptedFileMatcher, false}

	errors := make(chan error, 64)

	filepath.Walk(".", v.Walker(errors))

	select {
	case err := <-errors:
		if opts.Verbose {
			errhandle(err, false, "")
		}
	default:
	}
}

type GRVisitor struct {
	pattern             *regexp.Regexp
	ignoreFileMatcher   Matcher
	acceptedFileMatcher Matcher
	// Used to prevent sparse newline at the end of output
	prependNewLine bool
}

func (v *GRVisitor) Walker(errors chan<- error) filepath.WalkFunc {
	return func(fn string, fi os.FileInfo, err error) error {
		if err != nil {
			errors <- err
			return nil
		}

		// NOTE: if a directory is a symlink, filepath.Walk won't recurse inside
		if fi.Mode()&os.ModeSymlink != 0 {
			if fi, err = os.Stat(fn); err != nil {
				errors <- err
				return nil
			}
		}

		if fi.IsDir() {
			if !v.VisitDir(fn, fi) {
				return filepath.SkipDir
			}
			return nil
		}

		v.VisitFile(fn, fi)
		return nil
	}
}

func (v *GRVisitor) VisitDir(fn string, fi os.FileInfo) bool {
	return !v.ignoreFileMatcher.Match(fn, true)
}

func (v *GRVisitor) VisitFile(fn string, fi os.FileInfo) {
	if fi.Size() == 0 && !opts.FindFiles {
		return
	}

	if v.ignoreFileMatcher.Match(fn, false) {
		return
	}

	if !v.acceptedFileMatcher.Match(fn, false) {
		return
	}

	if fi.Size() >= 1024*1024*10 {
		fmt.Fprintf(os.Stderr, "Skipping %s, too big: %d\n", fn, fi.Size())
		return
	}

	if opts.FindFiles {
		v.SearchFileName(fn)
		return
	}

	f, content := v.GetFileAndContent(fn, fi)
	if f == nil {
		return
	}
	defer f.Close()

	if opts.Replace == nil {
		v.SearchFile(fn, content)
		return
	}

	changed, result := v.ReplaceInFile(fn, content)
	if changed {
		f.Seek(0, 0)
		n, err := f.Write(result)
		errhandle(err, true, "Error writing replacement in file %s", fn)
		if int64(n) < fi.Size() {
			err := f.Truncate(int64(n))
			errhandle(err, true, "Error truncating file to size %d", f)
		}
	}
}

func (v *GRVisitor) GetFileAndContent(fn string, fi os.FileInfo) (f *os.File, content []byte) {
	var err error
	var msg string

	if opts.Replace != nil {
		f, err = os.OpenFile(fn, os.O_RDWR, 0666)
		msg = "can't open file %s for reading and writing"
	} else {
		f, err = os.Open(fn)
		msg = "can't open file %s for reading"
	}

	if err != nil {
		if opts.Verbose {
			errhandle(err, false, msg, fn)
		}
		return
	}

	content = make([]byte, fi.Size())
	n, err := f.Read(content)
	errhandle(err, true, "can't read file %s", fn)
	if int64(n) != fi.Size() {
		errhandle(fmt.Errorf("Not whole file was read, only %d from %d",
			n, fi.Size()), true, "")
	}

	return
}

func (v *GRVisitor) SearchFile(fn string, content []byte) {
	lines := IntList([]int{})
	binary := false

	if bytes.IndexByte(content, 0) != -1 {
		binary = true
	}

	for _, info := range v.FindAllIndex(content) {
		if lines.Contains(info.num) {
			continue
		}

		if v.prependNewLine {
			fmt.Println("")
			v.prependNewLine = false
		}

		first := len(lines) == 0
		lines = append(lines, info.num)

		if first {
			if binary && !opts.OnlyName {
				fmt.Printf("Binary file %s matches\n", fn)
				break
			} else {
				if (!opts.NoColors) {
					color.Printf("@g%s\n", fn)
				} else {
					fmt.Printf("%s\n",fn)
				}
			}
		}

		if opts.OnlyName {
			return
		}

		if (!opts.NoColors) {
			color.Printf("@!@y%d:", info.num)
		} else {
			fmt.Printf("%d:", info.num)
		}
		colored := v.pattern.ReplaceAllStringFunc(string(info.line),
			func(wrap string) string {
				var res string
				if (!opts.NoColors) {
					res = color.Sprintf("@Y%s", wrap)
				} else {
					res = fmt.Sprintf("%s", wrap)
				}
				return res
			})
		fmt.Println(colored)
	}

	if len(lines) > 0 {
		v.prependNewLine = true
	}
}

func (v *GRVisitor) SearchFileName(fn string) {
	if !v.pattern.MatchString(fn) {
		return
	}
	colored := v.pattern.ReplaceAllStringFunc(fn,
		func(wrap string) string {
			var res string
			if (!opts.NoColors) {
				res = color.Sprintf("@Y%s", wrap)
			} else {
				res = fmt.Sprintf("%s", wrap)
			}
			return res
		})
	fmt.Println(colored)
}

func getSuffix(num int) string {
	if num > 1 {
		return "s"
	}
	return ""
}

func (v *GRVisitor) ReplaceInFile(fn string, content []byte) (changed bool, result []byte) {
	changed = false
	binary := false
	changenum := 0

	if opts.SingleLine {
		errhandle(
			fmt.Errorf("Can't handle singleline replacements yet"),
			true, "")
	}

	if opts.PlainText {
		errhandle(
			fmt.Errorf("Can't handle plain text replacements yet"),
			true, "")
	}

	if bytes.IndexByte(content, 0) != -1 {
		binary = true
	}

	result = v.pattern.ReplaceAllFunc(content, func(s []byte) []byte {
		if binary && !opts.Force {
			errhandle(
				fmt.Errorf("supply --force to force change of binary file"),
				false, "")
		}
		if !changed {
			changed = true
			if (!opts.NoColors) {
				color.Printf("@g%s", fn)
			} else {
				fmt.Printf("%s", fn)
			}
		}

		changenum += 1
		return []byte(*opts.Replace)
	})

	if changenum > 0 {
		if (!opts.NoColors) {
			color.Printf("@!@y - %d change%s made\n",
				changenum, getSuffix(changenum))
		} else {
			fmt.Printf(" - %d change%s made\n",
				changenum, getSuffix(changenum))
		}
	}

	return changed, result
}

type LineInfo struct {
	num  int
	line []byte
}

func (v *GRVisitor) FindAllIndex(content []byte) (res []*LineInfo) {
	if opts.SingleLine {
		return v.singlelineFindAllIndex(content)
	}

	linenum, last := 1, 0
	for _, bounds := range v.pattern.FindAllIndex(content, -1) {
		linenum += bytes.Count(content[last:bounds[0]], byteNewLine)
		last = bounds[0]
		begin, end := beginend(content, bounds[0], bounds[1])
		res = append(res, &LineInfo{linenum, content[begin:end]})
	}
	return res
}

func (v *GRVisitor) singlelineFindAllIndex(content []byte) (res []*LineInfo) {
	linenum, begin, end := 1, 0, 0
	for i := 0; i < len(content); i++ {
		if content[i] == '\n' {
			end = i
			line := content[begin:end]
			if v.pattern.Match(line) {
				res = append(res, &LineInfo{linenum, line})
			}
			linenum += 1
			begin = end + 1
		}
	}
	return res
}

// Given a []byte, start and finish of some inner slice, will find nearest
// newlines on both ends of this slice to contain this slice
func beginend(s []byte, start int, finish int) (begin int, end int) {
	begin = 0
	end = len(s)

	for i := start - 1; i >= 0; i-- {
		if s[i] == byteNewLine[0] {
			begin = i + 1
			break
		}
	}

	// -1 to check if current location is the end of a string
	for i := finish - 1; i < len(s); i++ {
		if s[i] == byteNewLine[0] {
			end = i
			break
		}
	}

	return begin, end
}

type IntList []int

func (il IntList) Contains(i int) bool {
	for _, x := range il {
		if x == i {
			return true
		}
	}
	return false
}
<|endoftext|>


<commit_before>package main

import (
	"net/http"
	"reflect"
	"sort"
	"testing"

	"github.com/ant0ine/go-json-rest/rest"
)

type MockResponseWriter struct {
	data interface{}
}

func (m *MockResponseWriter) Header() http.Header                      { return nil }
func (m *MockResponseWriter) WriteJson(v interface{}) error            { m.data = v; return nil }
func (m *MockResponseWriter) EncodeJson(v interface{}) ([]byte, error) { return nil, nil }
func (m *MockResponseWriter) WriteHeader(int)                          {}

func NewMockOptionsHandler(methods ...string) rest.HandlerFunc {
	return func(w rest.ResponseWriter, r *rest.Request) {
		sort.Strings(methods)
		w.WriteJson(methods)
	}
}

type RouteList []*rest.Route

func (slice RouteList) Less(i, j int) bool {
	return slice[i].HttpMethod+slice[i].PathExp < slice[j].HttpMethod+slice[j].PathExp
}

func (slice RouteList) Len() int {
	return len(slice)
}

func (slice RouteList) Swap(i, j int) {
	slice[i], slice[j] = slice[j], slice[i]
}

func TestAutogenOptionsRoutes(t *testing.T) {
	testList := []struct {
		out []*rest.Route
		in  []*rest.Route
	}{
		{nil, nil},
		{[]*rest.Route{}, []*rest.Route{}},
		{
			[]*rest.Route{
				rest.Get("/path", nil),
				rest.Options("/path", NewMockOptionsHandler("GET")),
			},
			[]*rest.Route{
				rest.Get("/path", nil),
			},
		},
		{
			[]*rest.Route{
				rest.Get("/path", nil),
				rest.Post("/path", nil),
				rest.Options("/path", NewMockOptionsHandler("GET", "POST")),
			},
			[]*rest.Route{
				rest.Get("/path", nil),
				rest.Post("/path", nil),
			},
		},
		{
			[]*rest.Route{
				rest.Get("/path", nil),
				rest.Options("/path", NewMockOptionsHandler("GET")),
				rest.Post("/path/path", nil),
				rest.Options("/path/path", NewMockOptionsHandler("POST")),
			},
			[]*rest.Route{
				rest.Get("/path", nil),
				rest.Post("/path/path", nil),
			},
		},
		{
			[]*rest.Route{
				rest.Get("/path", nil),
				rest.Put("/path", nil),
				rest.Options("/path", NewMockOptionsHandler("GET", "PUT")),
				rest.Post("/path/path", nil),
				rest.Options("/path/path", NewMockOptionsHandler("POST")),
			},
			[]*rest.Route{
				rest.Get("/path", nil),
				rest.Post("/path/path", nil),
				rest.Put("/path", nil),
			},
		},
	}

	for _, test := range testList {
		out := RouteList(AutogenOptionsRoutes(NewMockOptionsHandler, test.in...))

		if len(test.out) != len(out) {
			t.FailNow()
		}

		// Search requires sorted input
		sort.Sort(out)

		for _, expectedRoute := range test.out {

			idx := sort.Search(len(out), func(i int) bool {
				if out[i].HttpMethod+out[i].PathExp >= expectedRoute.HttpMethod+expectedRoute.PathExp {
					return true
				}
				return false
			})

			if idx == len(out) {
				t.FailNow()
			}

			if expectedRoute.HttpMethod == "OPTIONS" {
				exp := &MockResponseWriter{}
				expectedRoute.Func(exp, nil)

				recived := &MockResponseWriter{}
				out[idx].Func(recived, nil)

				if !reflect.DeepEqual(exp.data, recived.data) {
					t.FailNow()
				}
			}

		}
	}
}
<commit_msg>Improve TestAutogenOptionsRoutes imlementation.<commit_after>package main

import (
	"net/http"
	"reflect"
	"sort"
	"testing"

	"github.com/ant0ine/go-json-rest/rest"
)

type MockResponseWriter struct {
	methods []string
}

func (m *MockResponseWriter) Header() http.Header                      { return nil }
func (m *MockResponseWriter) WriteJson(v interface{}) error            { return nil }
func (m *MockResponseWriter) EncodeJson(v interface{}) ([]byte, error) { return nil, nil }
func (m *MockResponseWriter) WriteHeader(int)                          {}

func NewMockOptionsHandler(methods ...string) rest.HandlerFunc {
	return func(w rest.ResponseWriter, r *rest.Request) {
		sort.Strings(methods)
		mockWriter, _ := w.(*MockResponseWriter)
		mockWriter.methods = methods
	}
}

type RouteList []*rest.Route

func (slice RouteList) Less(i, j int) bool {
	return slice[i].HttpMethod+slice[i].PathExp < slice[j].HttpMethod+slice[j].PathExp
}

func (slice RouteList) Len() int {
	return len(slice)
}

func (slice RouteList) Swap(i, j int) {
	slice[i], slice[j] = slice[j], slice[i]
}

func TestAutogenOptionsRoutes(t *testing.T) {
	testList := []struct {
		out []*rest.Route
		in  []*rest.Route
	}{
		{nil, nil},
		{[]*rest.Route{}, []*rest.Route{}},
		{
			[]*rest.Route{
				rest.Get("/path", nil),
				rest.Options("/path", NewMockOptionsHandler("GET")),
			},
			[]*rest.Route{
				rest.Get("/path", nil),
			},
		},
		{
			[]*rest.Route{
				rest.Get("/path", nil),
				rest.Post("/path", nil),
				rest.Options("/path", NewMockOptionsHandler("GET", "POST")),
			},
			[]*rest.Route{
				rest.Get("/path", nil),
				rest.Post("/path", nil),
			},
		},
		{
			[]*rest.Route{
				rest.Get("/path", nil),
				rest.Options("/path", NewMockOptionsHandler("GET")),
				rest.Post("/path/path", nil),
				rest.Options("/path/path", NewMockOptionsHandler("POST")),
			},
			[]*rest.Route{
				rest.Get("/path", nil),
				rest.Post("/path/path", nil),
			},
		},
		{
			[]*rest.Route{
				rest.Get("/path", nil),
				rest.Put("/path", nil),
				rest.Options("/path", NewMockOptionsHandler("GET", "PUT")),
				rest.Post("/path/path", nil),
				rest.Options("/path/path", NewMockOptionsHandler("POST")),
			},
			[]*rest.Route{
				rest.Get("/path", nil),
				rest.Post("/path/path", nil),
				rest.Put("/path", nil),
			},
		},
	}

	for _, test := range testList {
		out := RouteList(AutogenOptionsRoutes(NewMockOptionsHandler, test.in...))

		if len(test.out) != len(out) {
			t.FailNow()
		}

		// Search requires sorted input
		sort.Sort(out)

		for _, expectedRoute := range test.out {

			idx := sort.Search(len(out), func(i int) bool {
				if out[i].HttpMethod+out[i].PathExp >= expectedRoute.HttpMethod+expectedRoute.PathExp {
					return true
				}
				return false
			})

			if idx == len(out) {
				t.FailNow()
			}

			if expectedRoute.HttpMethod == "OPTIONS" {
				exp := &MockResponseWriter{}
				expectedRoute.Func(exp, nil)

				recived := &MockResponseWriter{}
				out[idx].Func(recived, nil)

				if !reflect.DeepEqual(exp.methods, recived.methods) {
					t.FailNow()
				}
			}

		}
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2017 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 network

import (
	"fmt"
	"net/http"
	"time"

	computealpha "google.golang.org/api/compute/v0.alpha"

	"k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/util/wait"
	clientset "k8s.io/client-go/kubernetes"
	"k8s.io/kubernetes/pkg/cloudprovider"
	gcecloud "k8s.io/kubernetes/pkg/cloudprovider/providers/gce"
	"k8s.io/kubernetes/pkg/cloudprovider/providers/gce/cloud"
	"k8s.io/kubernetes/test/e2e/framework"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = SIGDescribe("Services [Feature:GCEAlphaFeature][Slow]", func() {
	f := framework.NewDefaultFramework("services")

	var cs clientset.Interface
	serviceLBNames := []string{}

	BeforeEach(func() {
		// This test suite requires the GCE environment.
		framework.SkipUnlessProviderIs("gce")
		cs = f.ClientSet
	})

	AfterEach(func() {
		if CurrentGinkgoTestDescription().Failed {
			framework.DescribeSvc(f.Namespace.Name)
		}
		for _, lb := range serviceLBNames {
			framework.Logf("cleaning gce resource for %s", lb)
			framework.CleanupServiceGCEResources(cs, lb, framework.TestContext.CloudConfig.Zone)
		}
		//reset serviceLBNames
		serviceLBNames = []string{}
	})
	It("should be able to create and tear down a standard-tier load balancer [Slow]", func() {
		lagTimeout := framework.LoadBalancerLagTimeoutDefault
		createTimeout := framework.GetServiceLoadBalancerCreationTimeout(cs)

		svcName := "net-tiers-svc"
		ns := f.Namespace.Name
		jig := framework.NewServiceTestJig(cs, svcName)

		By("creating a pod to be part of the service " + svcName)
		jig.RunOrFail(ns, nil)

		// Test 1: create a standard tiered LB for the Service.
		By("creating a Service of type LoadBalancer using the standard network tier")
		svc := jig.CreateTCPServiceOrFail(ns, func(svc *v1.Service) {
			svc.Spec.Type = v1.ServiceTypeLoadBalancer
			setNetworkTier(svc, gcecloud.NetworkTierAnnotationStandard)
		})
		// Verify that service has been updated properly.
		svcTier, err := gcecloud.GetServiceNetworkTier(svc)
		Expect(err).NotTo(HaveOccurred())
		Expect(svcTier).To(Equal(cloud.NetworkTierStandard))
		// Record the LB name for test cleanup.
		serviceLBNames = append(serviceLBNames, cloudprovider.GetLoadBalancerName(svc))

		// Wait and verify the LB.
		ingressIP := waitAndVerifyLBWithTier(jig, ns, svcName, "", createTimeout, lagTimeout)

		// Test 2: re-create a LB of a different tier for the updated Service.
		By("updating the Service to use the premium (default) tier")
		svc = jig.UpdateServiceOrFail(ns, svcName, func(svc *v1.Service) {
			clearNetworkTier(svc)
		})
		// Verify that service has been updated properly.
		svcTier, err = gcecloud.GetServiceNetworkTier(svc)
		Expect(err).NotTo(HaveOccurred())
		Expect(svcTier).To(Equal(cloud.NetworkTierDefault))

		// Wait until the ingress IP changes. Each tier has its own pool of
		// IPs, so changing tiers implies changing IPs.
		ingressIP = waitAndVerifyLBWithTier(jig, ns, svcName, ingressIP, createTimeout, lagTimeout)

		// Test 3: create a standard-tierd LB with a user-requested IP.
		By("reserving a static IP for the load balancer")
		requestedAddrName := fmt.Sprintf("e2e-ext-lb-net-tier-%s", framework.RunId)
		gceCloud, err := framework.GetGCECloud()
		Expect(err).NotTo(HaveOccurred())
		requestedIP, err := reserveAlphaRegionalAddress(gceCloud, requestedAddrName, cloud.NetworkTierStandard)
		Expect(err).NotTo(HaveOccurred(), "failed to reserve a STANDARD tiered address")
		defer func() {
			if requestedAddrName != "" {
				// Release GCE static address - this is not kube-managed and will not be automatically released.
				if err := gceCloud.DeleteRegionAddress(requestedAddrName, gceCloud.Region()); err != nil {
					framework.Logf("failed to release static IP address %q: %v", requestedAddrName, err)
				}
			}
		}()
		Expect(err).NotTo(HaveOccurred())
		framework.Logf("Allocated static IP to be used by the load balancer: %q", requestedIP)

		By("updating the Service to use the standard tier with a requested IP")
		svc = jig.UpdateServiceOrFail(ns, svc.Name, func(svc *v1.Service) {
			svc.Spec.LoadBalancerIP = requestedIP
			setNetworkTier(svc, gcecloud.NetworkTierAnnotationStandard)
		})
		// Verify that service has been updated properly.
		Expect(svc.Spec.LoadBalancerIP).To(Equal(requestedIP))
		svcTier, err = gcecloud.GetServiceNetworkTier(svc)
		Expect(err).NotTo(HaveOccurred())
		Expect(svcTier).To(Equal(cloud.NetworkTierStandard))

		// Wait until the ingress IP changes and verifies the LB.
		ingressIP = waitAndVerifyLBWithTier(jig, ns, svcName, ingressIP, createTimeout, lagTimeout)
	})
})

func waitAndVerifyLBWithTier(jig *framework.ServiceTestJig, ns, svcName, existingIP string, waitTimeout, checkTimeout time.Duration) string {
	var svc *v1.Service
	if existingIP == "" {
		// Creating the LB for the first time; wait for any ingress IP to show
		// up.
		svc = jig.WaitForNewIngressIPOrFail(ns, svcName, existingIP, waitTimeout)
	} else {
		// Re-creating the LB; wait for the ingress IP to change.
		svc = jig.WaitForNewIngressIPOrFail(ns, svcName, existingIP, waitTimeout)
	}

	svcPort := int(svc.Spec.Ports[0].Port)
	lbIngress := &svc.Status.LoadBalancer.Ingress[0]
	ingressIP := framework.GetIngressPoint(lbIngress)

	By("running sanity and reachability checks")
	if svc.Spec.LoadBalancerIP != "" {
		// Verify that the new ingress IP is the requested IP if it's set.
		Expect(ingressIP).To(Equal(svc.Spec.LoadBalancerIP))
	}
	jig.SanityCheckService(svc, v1.ServiceTypeLoadBalancer)
	// If the IP has been used by previous test, sometimes we get the lingering
	// 404 errors even after the LB is long gone. Tolerate and retry until the
	// the new LB is fully established since this feature is still Alpha in GCP.
	jig.TestReachableHTTPWithRetriableErrorCodes(ingressIP, svcPort, []int{http.StatusNotFound}, checkTimeout)

	// Verify the network tier matches the desired.
	svcNetTier, err := gcecloud.GetServiceNetworkTier(svc)
	Expect(err).NotTo(HaveOccurred())
	netTier, err := getLBNetworkTierByIP(ingressIP)
	Expect(err).NotTo(HaveOccurred(), "failed to get the network tier of the load balancer")
	Expect(netTier).To(Equal(svcNetTier))

	return ingressIP
}

func getLBNetworkTierByIP(ip string) (cloud.NetworkTier, error) {
	var rule *computealpha.ForwardingRule
	// Retry a few times to tolerate flakes.
	err := wait.PollImmediate(5*time.Second, 15*time.Second, func() (bool, error) {
		obj, err := getGCEForwardingRuleByIP(ip)
		if err != nil {
			return false, err
		}
		rule = obj
		return true, nil
	})
	if err != nil {
		return "", err
	}
	return cloud.NetworkTierGCEValueToType(rule.NetworkTier), nil
}

func getGCEForwardingRuleByIP(ip string) (*computealpha.ForwardingRule, error) {
	cloud, err := framework.GetGCECloud()
	if err != nil {
		return nil, err
	}
	ruleList, err := cloud.ListAlphaRegionForwardingRules(cloud.Region())
	if err != nil {
		return nil, err
	}
	for _, rule := range ruleList {
		if rule.IPAddress == ip {
			return rule, nil
		}
	}
	return nil, fmt.Errorf("forwarding rule with ip %q not found", ip)
}

func setNetworkTier(svc *v1.Service, tier string) {
	key := gcecloud.NetworkTierAnnotationKey
	if svc.ObjectMeta.Annotations == nil {
		svc.ObjectMeta.Annotations = map[string]string{}
	}
	svc.ObjectMeta.Annotations[key] = tier
}

func clearNetworkTier(svc *v1.Service) {
	key := gcecloud.NetworkTierAnnotationKey
	if svc.ObjectMeta.Annotations == nil {
		return
	}
	delete(svc.ObjectMeta.Annotations, key)
}

// TODO: add retries if this turns out to be flaky.
// TODO(#51665): remove this helper function once Network Tiers becomes beta.
func reserveAlphaRegionalAddress(cloud *gcecloud.GCECloud, name string, netTier cloud.NetworkTier) (string, error) {
	alphaAddr := &computealpha.Address{
		Name:        name,
		NetworkTier: netTier.ToGCEValue(),
	}

	if err := cloud.ReserveAlphaRegionAddress(alphaAddr, cloud.Region()); err != nil {
		return "", err
	}

	addr, err := cloud.GetRegionAddress(name, cloud.Region())
	if err != nil {
		return "", err
	}

	return addr.Address, nil
}
<commit_msg>requires string input<commit_after>/*
Copyright 2017 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 network

import (
	"fmt"
	"net/http"
	"time"

	computealpha "google.golang.org/api/compute/v0.alpha"

	"k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/util/wait"
	clientset "k8s.io/client-go/kubernetes"
	"k8s.io/kubernetes/pkg/cloudprovider"
	gcecloud "k8s.io/kubernetes/pkg/cloudprovider/providers/gce"
	"k8s.io/kubernetes/pkg/cloudprovider/providers/gce/cloud"
	"k8s.io/kubernetes/test/e2e/framework"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = SIGDescribe("Services [Feature:GCEAlphaFeature][Slow]", func() {
	f := framework.NewDefaultFramework("services")

	var cs clientset.Interface
	serviceLBNames := []string{}

	BeforeEach(func() {
		// This test suite requires the GCE environment.
		framework.SkipUnlessProviderIs("gce")
		cs = f.ClientSet
	})

	AfterEach(func() {
		if CurrentGinkgoTestDescription().Failed {
			framework.DescribeSvc(f.Namespace.Name)
		}
		for _, lb := range serviceLBNames {
			framework.Logf("cleaning gce resource for %s", lb)
			framework.CleanupServiceGCEResources(cs, lb, framework.TestContext.CloudConfig.Zone)
		}
		//reset serviceLBNames
		serviceLBNames = []string{}
	})
	It("should be able to create and tear down a standard-tier load balancer [Slow]", func() {
		lagTimeout := framework.LoadBalancerLagTimeoutDefault
		createTimeout := framework.GetServiceLoadBalancerCreationTimeout(cs)

		svcName := "net-tiers-svc"
		ns := f.Namespace.Name
		jig := framework.NewServiceTestJig(cs, svcName)

		By("creating a pod to be part of the service " + svcName)
		jig.RunOrFail(ns, nil)

		// Test 1: create a standard tiered LB for the Service.
		By("creating a Service of type LoadBalancer using the standard network tier")
		svc := jig.CreateTCPServiceOrFail(ns, func(svc *v1.Service) {
			svc.Spec.Type = v1.ServiceTypeLoadBalancer
			setNetworkTier(svc, string(gcecloud.NetworkTierAnnotationStandard))
		})
		// Verify that service has been updated properly.
		svcTier, err := gcecloud.GetServiceNetworkTier(svc)
		Expect(err).NotTo(HaveOccurred())
		Expect(svcTier).To(Equal(cloud.NetworkTierStandard))
		// Record the LB name for test cleanup.
		serviceLBNames = append(serviceLBNames, cloudprovider.GetLoadBalancerName(svc))

		// Wait and verify the LB.
		ingressIP := waitAndVerifyLBWithTier(jig, ns, svcName, "", createTimeout, lagTimeout)

		// Test 2: re-create a LB of a different tier for the updated Service.
		By("updating the Service to use the premium (default) tier")
		svc = jig.UpdateServiceOrFail(ns, svcName, func(svc *v1.Service) {
			clearNetworkTier(svc)
		})
		// Verify that service has been updated properly.
		svcTier, err = gcecloud.GetServiceNetworkTier(svc)
		Expect(err).NotTo(HaveOccurred())
		Expect(svcTier).To(Equal(cloud.NetworkTierDefault))

		// Wait until the ingress IP changes. Each tier has its own pool of
		// IPs, so changing tiers implies changing IPs.
		ingressIP = waitAndVerifyLBWithTier(jig, ns, svcName, ingressIP, createTimeout, lagTimeout)

		// Test 3: create a standard-tierd LB with a user-requested IP.
		By("reserving a static IP for the load balancer")
		requestedAddrName := fmt.Sprintf("e2e-ext-lb-net-tier-%s", framework.RunId)
		gceCloud, err := framework.GetGCECloud()
		Expect(err).NotTo(HaveOccurred())
		requestedIP, err := reserveAlphaRegionalAddress(gceCloud, requestedAddrName, cloud.NetworkTierStandard)
		Expect(err).NotTo(HaveOccurred(), "failed to reserve a STANDARD tiered address")
		defer func() {
			if requestedAddrName != "" {
				// Release GCE static address - this is not kube-managed and will not be automatically released.
				if err := gceCloud.DeleteRegionAddress(requestedAddrName, gceCloud.Region()); err != nil {
					framework.Logf("failed to release static IP address %q: %v", requestedAddrName, err)
				}
			}
		}()
		Expect(err).NotTo(HaveOccurred())
		framework.Logf("Allocated static IP to be used by the load balancer: %q", requestedIP)

		By("updating the Service to use the standard tier with a requested IP")
		svc = jig.UpdateServiceOrFail(ns, svc.Name, func(svc *v1.Service) {
			svc.Spec.LoadBalancerIP = requestedIP
			setNetworkTier(svc, string(gcecloud.NetworkTierAnnotationStandard))
		})
		// Verify that service has been updated properly.
		Expect(svc.Spec.LoadBalancerIP).To(Equal(requestedIP))
		svcTier, err = gcecloud.GetServiceNetworkTier(svc)
		Expect(err).NotTo(HaveOccurred())
		Expect(svcTier).To(Equal(cloud.NetworkTierStandard))

		// Wait until the ingress IP changes and verifies the LB.
		ingressIP = waitAndVerifyLBWithTier(jig, ns, svcName, ingressIP, createTimeout, lagTimeout)
	})
})

func waitAndVerifyLBWithTier(jig *framework.ServiceTestJig, ns, svcName, existingIP string, waitTimeout, checkTimeout time.Duration) string {
	var svc *v1.Service
	if existingIP == "" {
		// Creating the LB for the first time; wait for any ingress IP to show
		// up.
		svc = jig.WaitForNewIngressIPOrFail(ns, svcName, existingIP, waitTimeout)
	} else {
		// Re-creating the LB; wait for the ingress IP to change.
		svc = jig.WaitForNewIngressIPOrFail(ns, svcName, existingIP, waitTimeout)
	}

	svcPort := int(svc.Spec.Ports[0].Port)
	lbIngress := &svc.Status.LoadBalancer.Ingress[0]
	ingressIP := framework.GetIngressPoint(lbIngress)

	By("running sanity and reachability checks")
	if svc.Spec.LoadBalancerIP != "" {
		// Verify that the new ingress IP is the requested IP if it's set.
		Expect(ingressIP).To(Equal(svc.Spec.LoadBalancerIP))
	}
	jig.SanityCheckService(svc, v1.ServiceTypeLoadBalancer)
	// If the IP has been used by previous test, sometimes we get the lingering
	// 404 errors even after the LB is long gone. Tolerate and retry until the
	// the new LB is fully established since this feature is still Alpha in GCP.
	jig.TestReachableHTTPWithRetriableErrorCodes(ingressIP, svcPort, []int{http.StatusNotFound}, checkTimeout)

	// Verify the network tier matches the desired.
	svcNetTier, err := gcecloud.GetServiceNetworkTier(svc)
	Expect(err).NotTo(HaveOccurred())
	netTier, err := getLBNetworkTierByIP(ingressIP)
	Expect(err).NotTo(HaveOccurred(), "failed to get the network tier of the load balancer")
	Expect(netTier).To(Equal(svcNetTier))

	return ingressIP
}

func getLBNetworkTierByIP(ip string) (cloud.NetworkTier, error) {
	var rule *computealpha.ForwardingRule
	// Retry a few times to tolerate flakes.
	err := wait.PollImmediate(5*time.Second, 15*time.Second, func() (bool, error) {
		obj, err := getGCEForwardingRuleByIP(ip)
		if err != nil {
			return false, err
		}
		rule = obj
		return true, nil
	})
	if err != nil {
		return "", err
	}
	return cloud.NetworkTierGCEValueToType(rule.NetworkTier), nil
}

func getGCEForwardingRuleByIP(ip string) (*computealpha.ForwardingRule, error) {
	cloud, err := framework.GetGCECloud()
	if err != nil {
		return nil, err
	}
	ruleList, err := cloud.ListAlphaRegionForwardingRules(cloud.Region())
	if err != nil {
		return nil, err
	}
	for _, rule := range ruleList {
		if rule.IPAddress == ip {
			return rule, nil
		}
	}
	return nil, fmt.Errorf("forwarding rule with ip %q not found", ip)
}

func setNetworkTier(svc *v1.Service, tier string) {
	key := gcecloud.NetworkTierAnnotationKey
	if svc.ObjectMeta.Annotations == nil {
		svc.ObjectMeta.Annotations = map[string]string{}
	}
	svc.ObjectMeta.Annotations[key] = tier
}

func clearNetworkTier(svc *v1.Service) {
	key := gcecloud.NetworkTierAnnotationKey
	if svc.ObjectMeta.Annotations == nil {
		return
	}
	delete(svc.ObjectMeta.Annotations, key)
}

// TODO: add retries if this turns out to be flaky.
// TODO(#51665): remove this helper function once Network Tiers becomes beta.
func reserveAlphaRegionalAddress(cloud *gcecloud.GCECloud, name string, netTier cloud.NetworkTier) (string, error) {
	alphaAddr := &computealpha.Address{
		Name:        name,
		NetworkTier: netTier.ToGCEValue(),
	}

	if err := cloud.ReserveAlphaRegionAddress(alphaAddr, cloud.Region()); err != nil {
		return "", err
	}

	addr, err := cloud.GetRegionAddress(name, cloud.Region())
	if err != nil {
		return "", err
	}

	return addr.Address, nil
}
<|endoftext|>