<commit_before>package convey

import (
	"fmt"
	"reflect"
	"strings"
)

import (
	"github.com/jacobsa/oglematchers"
)

// assertion is an alias for a function with a signature that the convey.So()
// method can handle. Any future or custom assertions should conform to this
// method signature. The return value should be an empty string if the assertion
// passes and a well-formed failure message if not.
type assertion func(actual interface{}, expected ...interface{}) string

// ShouldEqual receives exactly two parameters and does an equality check.
func ShouldEqual(actual interface{}, expected ...interface{}) string {
	if message := onlyOne(expected); message != success {
		return message
	}
	return shouldEqual(actual, expected[0])
}
func shouldEqual(actual, expected interface{}) (message string) {
	defer func() {
		if r := recover(); r != nil {
			message = fmt.Sprintf(shouldHaveBeenEqual, actual, expected)
		}
	}()

	if matchError := oglematchers.Equals(expected).Matches(actual); matchError != nil {
		return fmt.Sprintf(shouldHaveBeenEqual, actual, expected)
	}

	return success
}

// ShouldNotEqual receives exactly two parameters and does an inequality check.
func ShouldNotEqual(actual interface{}, expected ...interface{}) string {
	if fail := onlyOne(expected); fail != success {
		return fail
	} else if ShouldEqual(actual, expected[0]) == success {
		return fmt.Sprintf(shouldNotHaveBeenEqual, actual, expected[0])
	}
	return success
}

// ShouldResemble receives exactly two parameters and does a deep equal check (see reflect.DeepEqual)
func ShouldResemble(actual interface{}, expected ...interface{}) string {
	if message := onlyOne(expected); message != success {
		return message
	}

	matcher := oglematchers.DeepEquals(expected[0])
	matchError := matcher.Matches(actual)
	if matchError != nil {
		return fmt.Sprintf(shouldHaveResembled, actual, expected[0])
	}

	return success
}

// ShouldNotResemble receives exactly two parameters and does an inverse deep equal check (see reflect.DeepEqual)
func ShouldNotResemble(actual interface{}, expected ...interface{}) string {
	if message := onlyOne(expected); message != success {
		return message
	} else if ShouldResemble(actual, expected[0]) == success {
		return fmt.Sprintf(shouldNotHaveResembled, actual, expected[0])
	}
	return success
}

// ShouldPointTo receives exactly two parameters and checks to see that they point to the same address.
func ShouldPointTo(actual interface{}, expected ...interface{}) string {
	if message := onlyOne(expected); message != success {
		return message
	}
	return shouldPointTo(actual, expected[0])

}
func shouldPointTo(actual, expected interface{}) string {
	actualValue := reflect.ValueOf(actual)
	expectedValue := reflect.ValueOf(expected)

	if ShouldNotBeNil(actual) != success {
		return fmt.Sprintf(shouldHaveBeenNonNilPointer, "first", "nil")
	} else if ShouldNotBeNil(expected) != success {
		return fmt.Sprintf(shouldHaveBeenNonNilPointer, "second", "nil")
	} else if actualValue.Kind() != reflect.Ptr {
		return fmt.Sprintf(shouldHaveBeenNonNilPointer, "first", "not")
	} else if expectedValue.Kind() != reflect.Ptr {
		return fmt.Sprintf(shouldHaveBeenNonNilPointer, "second", "not")
	} else if ShouldEqual(actualValue.Pointer(), expectedValue.Pointer()) != success {
		return fmt.Sprintf(shouldHavePointedTo,
			actual, reflect.ValueOf(actual).Pointer(),
			expected, reflect.ValueOf(expected).Pointer())
	}
	return success
}

// ShouldNotPointTo receives exactly two parameters and checks to see that they point to different addresess.
func ShouldNotPointTo(actual interface{}, expected ...interface{}) string {
	if message := onlyOne(expected); message != success {
		return message
	}
	compare := ShouldPointTo(actual, expected[0])
	if strings.HasPrefix(compare, shouldBePointers) {
		return compare
	} else if compare == success {
		return fmt.Sprintf(shouldNotHavePointedTo, actual, expected[0], reflect.ValueOf(actual).Pointer())
	}
	return success
}

// ShouldBeNil receives a single parameter and ensures that it is nil.
func ShouldBeNil(actual interface{}, expected ...interface{}) string {
	if fail := none(expected); fail != success {
		return fail
	} else if actual == nil {
		return success
	} else if interfaceIsNilPointer(actual) {
		return success
	}
	return fmt.Sprintf(shouldHaveBeenNil, actual)
}
func interfaceIsNilPointer(actual interface{}) bool {
	value := reflect.ValueOf(actual)
	return value.Kind() == reflect.Ptr && value.Pointer() == 0
}

// ShouldNotBeNil receives a single parameter and ensures that it is not nil.
func ShouldNotBeNil(actual interface{}, expected ...interface{}) string {
	if fail := none(expected); fail != success {
		return fail
	} else if ShouldBeNil(actual) == success {
		return fmt.Sprintf(shouldNotHaveBeenNil, actual)
	}
	return success
}

// ShouldBeTrue receives a single parameter and ensures that it is true.
func ShouldBeTrue(actual interface{}, expected ...interface{}) string {
	if fail := none(expected); fail != success {
		return fail
	} else if actual != true {
		return fmt.Sprintf(shouldHaveBeenTrue, actual)
	}
	return success
}

// ShouldBeFalse receives a single parameter and ensures that it is false.
func ShouldBeFalse(actual interface{}, expected ...interface{}) string {
	if fail := none(expected); fail != success {
		return fail
	} else if actual != false {
		return fmt.Sprintf(shouldHaveBeenFalse, actual)
	}
	return success
}

// ShouldBeGreaterThan receives exactly two parameters and ensures that the first is greater than the second.
func ShouldBeGreaterThan(actual interface{}, expected ...interface{}) string {
	if fail := onlyOne(expected); fail != success {
		return fail
	}

	if matchError := oglematchers.GreaterThan(expected[0]).Matches(actual); matchError != nil {
		return fmt.Sprintf(shouldHaveBeenGreater, actual, expected[0])
	}
	return success
}

// ShouldBeGreaterThanOrEqualTo receives exactly two parameters and ensures that the first is greater than or equal to the second.
func ShouldBeGreaterThanOrEqualTo(actual interface{}, expected ...interface{}) string {
	if fail := onlyOne(expected); fail != success {
		return fail
	} else if matchError := oglematchers.GreaterOrEqual(expected[0]).Matches(actual); matchError != nil {
		return fmt.Sprintf(shouldHaveBeenGreaterOrEqual, actual, expected[0])
	}
	return success
}

// ShouldBeLessThan receives exactly two parameters and ensures that the first is less than the second.
func ShouldBeLessThan(actual interface{}, expected ...interface{}) string {
	if fail := onlyOne(expected); fail != success {
		return fail
	} else if matchError := oglematchers.LessThan(expected[0]).Matches(actual); matchError != nil {
		return fmt.Sprintf(shouldHaveBeenLess, actual, expected[0])
	}
	return success
}

// ShouldBeLessThan receives exactly two parameters and ensures that the first is less than or equal to the second.
func ShouldBeLessThanOrEqualTo(actual interface{}, expected ...interface{}) string {
	if fail := onlyOne(expected); fail != success {
		return fail
	} else if matchError := oglematchers.LessOrEqual(expected[0]).Matches(actual); matchError != nil {
		return fmt.Sprintf(shouldHaveBeenLess, actual, expected[0])
	}
	return success
}

// ShouldBeBetween receives exactly three parameters: an actual value, a lower bound, and an upper bound.
// It ensures that the actual value is between both bounds or is equal to one of the bounds.
func ShouldBeBetween(actual interface{}, expected ...interface{}) string {
	if fail := onlyTwo(expected); fail != success {
		return fail
	}
	return "asdf"
}
<commit_msg>New referencing the smartystreets fork of oglematchers.<commit_after>package convey

import (
	"fmt"
	"reflect"
	"strings"
)

import (
	"github.com/smartystreets/oglematchers"
)

// assertion is an alias for a function with a signature that the convey.So()
// method can handle. Any future or custom assertions should conform to this
// method signature. The return value should be an empty string if the assertion
// passes and a well-formed failure message if not.
type assertion func(actual interface{}, expected ...interface{}) string

// ShouldEqual receives exactly two parameters and does an equality check.
func ShouldEqual(actual interface{}, expected ...interface{}) string {
	if message := onlyOne(expected); message != success {
		return message
	}
	return shouldEqual(actual, expected[0])
}
func shouldEqual(actual, expected interface{}) (message string) {
	defer func() {
		if r := recover(); r != nil {
			message = fmt.Sprintf(shouldHaveBeenEqual, actual, expected)
		}
	}()

	if matchError := oglematchers.Equals(expected).Matches(actual); matchError != nil {
		return fmt.Sprintf(shouldHaveBeenEqual, actual, expected)
	}

	return success
}

// ShouldNotEqual receives exactly two parameters and does an inequality check.
func ShouldNotEqual(actual interface{}, expected ...interface{}) string {
	if fail := onlyOne(expected); fail != success {
		return fail
	} else if ShouldEqual(actual, expected[0]) == success {
		return fmt.Sprintf(shouldNotHaveBeenEqual, actual, expected[0])
	}
	return success
}

// ShouldResemble receives exactly two parameters and does a deep equal check (see reflect.DeepEqual)
func ShouldResemble(actual interface{}, expected ...interface{}) string {
	if message := onlyOne(expected); message != success {
		return message
	}

	matcher := oglematchers.DeepEquals(expected[0])
	matchError := matcher.Matches(actual)
	if matchError != nil {
		return fmt.Sprintf(shouldHaveResembled, actual, expected[0])
	}

	return success
}

// ShouldNotResemble receives exactly two parameters and does an inverse deep equal check (see reflect.DeepEqual)
func ShouldNotResemble(actual interface{}, expected ...interface{}) string {
	if message := onlyOne(expected); message != success {
		return message
	} else if ShouldResemble(actual, expected[0]) == success {
		return fmt.Sprintf(shouldNotHaveResembled, actual, expected[0])
	}
	return success
}

// ShouldPointTo receives exactly two parameters and checks to see that they point to the same address.
func ShouldPointTo(actual interface{}, expected ...interface{}) string {
	if message := onlyOne(expected); message != success {
		return message
	}
	return shouldPointTo(actual, expected[0])

}
func shouldPointTo(actual, expected interface{}) string {
	actualValue := reflect.ValueOf(actual)
	expectedValue := reflect.ValueOf(expected)

	if ShouldNotBeNil(actual) != success {
		return fmt.Sprintf(shouldHaveBeenNonNilPointer, "first", "nil")
	} else if ShouldNotBeNil(expected) != success {
		return fmt.Sprintf(shouldHaveBeenNonNilPointer, "second", "nil")
	} else if actualValue.Kind() != reflect.Ptr {
		return fmt.Sprintf(shouldHaveBeenNonNilPointer, "first", "not")
	} else if expectedValue.Kind() != reflect.Ptr {
		return fmt.Sprintf(shouldHaveBeenNonNilPointer, "second", "not")
	} else if ShouldEqual(actualValue.Pointer(), expectedValue.Pointer()) != success {
		return fmt.Sprintf(shouldHavePointedTo,
			actual, reflect.ValueOf(actual).Pointer(),
			expected, reflect.ValueOf(expected).Pointer())
	}
	return success
}

// ShouldNotPointTo receives exactly two parameters and checks to see that they point to different addresess.
func ShouldNotPointTo(actual interface{}, expected ...interface{}) string {
	if message := onlyOne(expected); message != success {
		return message
	}
	compare := ShouldPointTo(actual, expected[0])
	if strings.HasPrefix(compare, shouldBePointers) {
		return compare
	} else if compare == success {
		return fmt.Sprintf(shouldNotHavePointedTo, actual, expected[0], reflect.ValueOf(actual).Pointer())
	}
	return success
}

// ShouldBeNil receives a single parameter and ensures that it is nil.
func ShouldBeNil(actual interface{}, expected ...interface{}) string {
	if fail := none(expected); fail != success {
		return fail
	} else if actual == nil {
		return success
	} else if interfaceIsNilPointer(actual) {
		return success
	}
	return fmt.Sprintf(shouldHaveBeenNil, actual)
}
func interfaceIsNilPointer(actual interface{}) bool {
	value := reflect.ValueOf(actual)
	return value.Kind() == reflect.Ptr && value.Pointer() == 0
}

// ShouldNotBeNil receives a single parameter and ensures that it is not nil.
func ShouldNotBeNil(actual interface{}, expected ...interface{}) string {
	if fail := none(expected); fail != success {
		return fail
	} else if ShouldBeNil(actual) == success {
		return fmt.Sprintf(shouldNotHaveBeenNil, actual)
	}
	return success
}

// ShouldBeTrue receives a single parameter and ensures that it is true.
func ShouldBeTrue(actual interface{}, expected ...interface{}) string {
	if fail := none(expected); fail != success {
		return fail
	} else if actual != true {
		return fmt.Sprintf(shouldHaveBeenTrue, actual)
	}
	return success
}

// ShouldBeFalse receives a single parameter and ensures that it is false.
func ShouldBeFalse(actual interface{}, expected ...interface{}) string {
	if fail := none(expected); fail != success {
		return fail
	} else if actual != false {
		return fmt.Sprintf(shouldHaveBeenFalse, actual)
	}
	return success
}

// ShouldBeGreaterThan receives exactly two parameters and ensures that the first is greater than the second.
func ShouldBeGreaterThan(actual interface{}, expected ...interface{}) string {
	if fail := onlyOne(expected); fail != success {
		return fail
	}

	if matchError := oglematchers.GreaterThan(expected[0]).Matches(actual); matchError != nil {
		return fmt.Sprintf(shouldHaveBeenGreater, actual, expected[0])
	}
	return success
}

// ShouldBeGreaterThanOrEqualTo receives exactly two parameters and ensures that the first is greater than or equal to the second.
func ShouldBeGreaterThanOrEqualTo(actual interface{}, expected ...interface{}) string {
	if fail := onlyOne(expected); fail != success {
		return fail
	} else if matchError := oglematchers.GreaterOrEqual(expected[0]).Matches(actual); matchError != nil {
		return fmt.Sprintf(shouldHaveBeenGreaterOrEqual, actual, expected[0])
	}
	return success
}

// ShouldBeLessThan receives exactly two parameters and ensures that the first is less than the second.
func ShouldBeLessThan(actual interface{}, expected ...interface{}) string {
	if fail := onlyOne(expected); fail != success {
		return fail
	} else if matchError := oglematchers.LessThan(expected[0]).Matches(actual); matchError != nil {
		return fmt.Sprintf(shouldHaveBeenLess, actual, expected[0])
	}
	return success
}

// ShouldBeLessThan receives exactly two parameters and ensures that the first is less than or equal to the second.
func ShouldBeLessThanOrEqualTo(actual interface{}, expected ...interface{}) string {
	if fail := onlyOne(expected); fail != success {
		return fail
	} else if matchError := oglematchers.LessOrEqual(expected[0]).Matches(actual); matchError != nil {
		return fmt.Sprintf(shouldHaveBeenLess, actual, expected[0])
	}
	return success
}

// ShouldBeBetween receives exactly three parameters: an actual value, a lower bound, and an upper bound.
// It ensures that the actual value is between both bounds or is equal to one of the bounds.
func ShouldBeBetween(actual interface{}, expected ...interface{}) string {
	if fail := onlyTwo(expected); fail != success {
		return fail
	}
	return "asdf"
}
<|endoftext|>


<commit_before>package whale

import (
	"strings"
	"time"
	"strconv"
)

import (
	"butler/null"
	"gokogiri/xpath"
	"golog"
	"rubex"
	tp "tritium/proto"
	"steno"
	"go-cache"
	"go-cache/arc"
)

type Whale struct {
	Log *golog.Logger
	Debugger steno.Debugger
	RegexpCache       cache.Cache
	XPathCache        cache.Cache
}

type EngineContext struct {
	Functions                []*Function
	Types                    []string
	Exports                  [][]string
	Logs                     []string
	ImportedFiles            []string
	Env                      map[string]string
	MatchStack               []string
	MatchShouldContinueStack []bool
	Yields                   []*YieldBlock
	*Whale
	*tp.Transform
	Rrules []*tp.RewriteRule

	InnerReplacer     *rubex.Regexp
	HeaderContentType *rubex.Regexp

	// Debug info
	Filename          string
	HadError          bool
	Deadline time.Time
	Mobjects []MemoryObject
	MessagePath string
}

const OutputBufferSize = 500 * 1024 //500KB
const defaultMobjects = 4
const TimeoutError = "EngineTimeout"

func NewEngine(logger *golog.Logger, debugger steno.Debugger) *Whale {
	e := &Whale{
		Log: logger,
		Debugger: debugger,
		RegexpCache:              arc.NewARCache(1000),
		XPathCache:               arc.NewARCache(1000),
	}
	e.RegexpCache.SetCleanFunc(CleanRegexpObject)
	e.XPathCache.SetCleanFunc(CleanXpathExpObject)
	return e
}

func NewEngineCtx(eng *Whale, vars map[string]string, transform *tp.Transform, rrules []*tp.RewriteRule, deadline time.Time, messagePath string) (ctx *EngineContext) {
	ctx = &EngineContext{
		Whale:                    eng,
		Exports:                  make([][]string, 0),
		Logs:                     make([]string, 0),
		Env:                      vars,
		Transform:                transform,
		Rrules:                   rrules,
		MatchStack:               make([]string, 0),
		MatchShouldContinueStack: make([]bool, 0),
		Yields:                   make([]*YieldBlock, 0),
		HadError:                 false,

		Deadline:                 deadline,
		Mobjects:                 make([]MemoryObject, 0, defaultMobjects),
		MessagePath:              messagePath,
	}
	return
}

func (eng *Whale) Free() {
	eng.RegexpCache.Reset()
	eng.XPathCache.Reset()
}

func (eng *Whale) Run(transform *tp.Transform, rrules []*tp.RewriteRule, input interface{}, vars map[string]string, deadline time.Time, customer, project, messagePath string) (output string, exports [][]string, logs []string) {
	ctx := NewEngineCtx(eng, vars, transform, rrules, deadline, messagePath)
	defer ctx.Free()
	ctx.Yields = append(ctx.Yields, &YieldBlock{Vars: make(map[string]interface{})})
	ctx.UsePackage(transform.Pkg)
	scope := &Scope{Value: input.(string)}
	obj := transform.Objects[0]
	ctx.Filename = null.GetString(obj.Name)
	ctx.RunInstruction(scope, obj.Root)
	output = scope.Value.(string)
	exports = ctx.Exports
	logs = ctx.Logs
	return
}

func (eng *Whale) GetCacheStats() (int, int, int, int) {
	return eng.RegexpCache.GetHitRate(), eng.RegexpCache.GetUsageRate(), eng.XPathCache.GetHitRate(), eng.XPathCache.GetUsageRate()
}

func (ctx *EngineContext) Free() {
    for _, o := range ctx.Mobjects {
		if o != nil {
			o.Free()
		}
	}
}

func (ctx *EngineContext) RunInstruction(scope *Scope, ins *tp.Instruction) (returnValue interface{}) {
	defer func() {
		//TODO Stack traces seem to get truncated on syslog...
		if x := recover(); x != nil {
			err, ok := x.(error)
			errString := ""
			if ok {
				errString = err.Error()
			} else {
				errString = x.(string)
			}
			if errString != TimeoutError {
				if ctx.HadError == false {
					ctx.HadError = true
					errString = errString + "\n" + ins.Type.String() + " " + null.GetString(ins.Value) + "\n\n\nTritium Stack\n=========\n\n"
				}
				errString = errString + ctx.FileAndLine(ins) + "\n"
			}
			panic(errString)
		}
	}()

	ctx.Whale.Debugger.TrapInstruction(ctx.MessagePath, ctx.Filename, ctx.Env, ins, scope.Value, scope.Index)

	if time.Now().After(ctx.Deadline) {
		panic(TimeoutError)
	}

	// If our object is invalid, then skip it
	if null.GetBool(ins.IsValid) == false {
		panic("Invalid instruction. Should have stopped before linking!")
	}
	indent := ""
	for i := 0; i < len(ctx.Yields); i++ {
		indent += "\t"
	}

	returnValue = ""
	switch *ins.Type {
	case tp.Instruction_BLOCK:
		for _, child := range ins.Children {
			returnValue = ctx.RunInstruction(scope, child)
		}
	case tp.Instruction_TEXT:
		returnValue = null.GetString(ins.Value)
	case tp.Instruction_LOCAL_VAR:
		name := null.GetString(ins.Value)
		vars := ctx.Vars()
		if len(ins.Arguments) > 0 {
			vars[name] = ctx.RunInstruction(scope, ins.Arguments[0])
		}
		if len(ins.Children) > 0 {
			ts := &Scope{Value: ctx.Vars()[name]}
			for _, child := range ins.Children {
				ctx.RunInstruction(ts, child)
			}
			vars[name] = ts.Value
		}
		returnValue = vars[name]
	case tp.Instruction_IMPORT:
		obj := ctx.Objects[int(null.GetInt32(ins.ObjectId))]
		curFile := ctx.Filename
		ctx.Filename = null.GetString(obj.Name)
		ctx.Whale.Debugger.LogImport(ctx.MessagePath, ctx.Filename, curFile, int(null.GetInt32(ins.LineNumber)))
		ctx.AddLog("__IMPORT__FILE__:" + ctx.Filename)
		for _, child := range obj.Root.Children {
			ctx.RunInstruction(scope, child)
		}
		ctx.Whale.Debugger.LogImportDone(ctx.MessagePath, ctx.Filename, curFile, int(null.GetInt32(ins.LineNumber)))
		ctx.Filename = curFile
	case tp.Instruction_FUNCTION_CALL:
		fun := ctx.Functions[int(null.GetInt32(ins.FunctionId))]
		args := make([]interface{}, len(ins.Arguments))
		for i, argIns := range ins.Arguments {
			args[i] = ctx.RunInstruction(scope, argIns)
		}

		if null.GetBool(fun.BuiltIn) {
			if f := builtInFunctions[fun.Name]; f != nil {
				returnValue = f(ctx, scope, ins, args)
				if returnValue == nil {
					returnValue = ""
				}
			} else {
				panic("missing function: " + fun.Name)
			}
		} else {
			// We are using a user-defined function
			//println("Resetting localvar")
			// Setup the new local var
			curFile := ctx.Filename
			ctx.Filename = fun.GetFilename()
			vars := make(map[string]interface{}, len(args))
			for i, arg := range fun.Args {
				vars[null.GetString(arg.Name)] = args[i]
			}
			yieldBlock := &YieldBlock{
				Ins:  ins,
				Vars: vars,
			}
			// PUSH!
			ctx.PushYieldBlock(yieldBlock)
			for _, child := range fun.Instruction.Children {
				returnValue = ctx.RunInstruction(scope, child)
			}
			// POP!
			ctx.PopYieldBlock()
			ctx.Filename = curFile
		}
	}

	return
}

func (ctx *EngineContext) ShouldContinue() (result bool) {
	if len(ctx.MatchShouldContinueStack) > 0 {
		result = ctx.MatchShouldContinueStack[len(ctx.MatchShouldContinueStack)-1]
	} else {
		result = false
	}
	return
}

func (ctx *EngineContext) MatchTarget() string {
	return ctx.MatchStack[len(ctx.MatchStack)-1]
}

func (ctx *EngineContext) PushYieldBlock(b *YieldBlock) {
	ctx.Yields = append(ctx.Yields, b)
}

func (ctx *EngineContext) PopYieldBlock() (b *YieldBlock) {
	num := len(ctx.Yields)
	if num > 0 {
		b = ctx.Yields[num-1]
		ctx.Yields = ctx.Yields[:num-1]
	}
	return
}

func (ctx *EngineContext) HasYieldBlock() bool {
	return len(ctx.Yields) > 0
}

func (ctx *EngineContext) TopYieldBlock() (b *YieldBlock) {
	num := len(ctx.Yields)
	if num > 0 {
		b = ctx.Yields[num-1]
	}
	return
}

func (ctx *EngineContext) Vars() map[string]interface{} {
	b := ctx.TopYieldBlock()
	if b != nil {
		return b.Vars
	}
	return nil
}

func (ctx *EngineContext) FileAndLine(ins *tp.Instruction) string {
	lineNum := strconv.Itoa(int(null.GetInt32(ins.LineNumber)))
	return (ctx.Filename + ":" + lineNum)
}

func (ctx *EngineContext) UsePackage(pkg *tp.Package) {
	ctx.Types = make([]string, len(pkg.Types))
	for i, t := range pkg.Types {
		ctx.Types[i] = null.GetString(t.Name)
	}

	ctx.Functions = make([]*Function, len(pkg.Functions))
	for i, f := range pkg.Functions {
		name := null.GetString(f.Name)
		for _, a := range f.Args {
			typeString := ctx.Types[int(null.GetInt32(a.TypeId))]
			name = name + "." + typeString
		}
		fun := &Function{
			Name:     name,
			Function: f,
		}
		ctx.Functions[i] = fun
	}
}

func (ctx *EngineContext) GetRegexp(pattern, options string) (r *rubex.Regexp) {
	sig := pattern + "/" + options
	object, err := ctx.RegexpCache.Get(sig)
	if err != nil {
		mode := rubex.ONIG_OPTION_DEFAULT
		if strings.Index(options, "i") >= 0 {
			mode = rubex.ONIG_OPTION_IGNORECASE
		}
		if strings.Index(options, "m") >= 0 {
			mode = rubex.ONIG_OPTION_MULTILINE
		}
		var err error
		r, err = rubex.NewRegexp(pattern, mode)
		if err == nil {
			//ctx.AddMemoryObject(r)
			ctx.RegexpCache.Set(sig, &RegexpObject{Regexp:r})
		}
		return r
	}
	return object.(*RegexpObject).Regexp
}

func (ctx *EngineContext) GetXpathExpr(p string) (e *xpath.Expression) {
	object, err := ctx.XPathCache.Get(p)
	if err != nil {
		e = xpath.Compile(p)
		if e != nil {
			//ctx.AddMemoryObject(e)
			ctx.XPathCache.Set(p, &XpathExpObject{Expression: e})
		} else {
			ctx.AddLog("Invalid XPath used: " + p)
		}
		return e
	}
	return object.(*XpathExpObject).Expression
}

func (ctx *EngineContext) AddExport(exports []string) {
	ctx.Exports = append(ctx.Exports, exports)
}

func (ctx *EngineContext) AddLog(log string) int {
	//ctx.Log.Info("TRITIUM: " + log)
	index := len(ctx.Logs)
	ctx.Logs = append(ctx.Logs, log)
	return index
}

func (ctx *EngineContext) UpdateLog(index int, log string) {
	//ctx.Log.Info("TRITIUM: " + log)
	if index >= 0 && index < len(ctx.Logs) {
		ctx.Logs[index] = log
	}
}

func (ctx *EngineContext) SetEnv(key, val string) {
	ctx.Env[key] = val
}

func (ctx *EngineContext) GetEnv(key string) (val string) {
	val = ctx.Env[key]
	return
}

func (ctx *EngineContext) SetVar(key string, val interface{}) {
	b := ctx.TopYieldBlock()
	if b != nil {
		b.Vars[key] = val
	}
}

func (ctx *EngineContext) GetVar(key string) (val interface{}) {
	b := ctx.TopYieldBlock()
	if b != nil {
		val = b.Vars[key]
	}
	return
}

func (ctx *EngineContext) GetInnerReplacer() (r *rubex.Regexp) {
	return ctx.GetRegexp(`[\\$](\d)`, "i")
}

func (ctx *EngineContext) GetHeaderContentTypeRegex() (r *rubex.Regexp) {
	return ctx.GetRegexp(`<meta\s+http-equiv="content-type"\s+content="(.*?)"`, "i")
}

func (ctx *EngineContext) GetOutputBuffer() (b []byte) {
	//b = ctx.OutputBuffer
	return
}

func (ctx *EngineContext) Logger() (logger *golog.Logger) {
	logger = ctx.Log
	return
}

func (ctx *EngineContext) PushMatchStack(match string) {
	ctx.MatchStack = append(ctx.MatchStack, match)
}

func (ctx *EngineContext) PopMatchStack() (match string) {
	if num := len(ctx.MatchStack); num > 0 {
		match = ctx.MatchStack[num-1]
		ctx.MatchStack = ctx.MatchStack[:num-1]
	}
	return
}

func (ctx *EngineContext) PushShouldContinueStack(cont bool) {
	ctx.MatchShouldContinueStack = append(ctx.MatchShouldContinueStack, cont)
}

func (ctx *EngineContext) PopShouldContinueStack() (cont bool) {
	if num := len(ctx.MatchShouldContinueStack); num > 0 {
		cont = ctx.MatchShouldContinueStack[num-1]
		ctx.MatchShouldContinueStack = ctx.MatchShouldContinueStack[:num-1]
	}
	return
}

func (ctx *EngineContext) SetShouldContinue(cont bool) {
	if num := len(ctx.MatchShouldContinueStack); num > 0 {
		ctx.MatchShouldContinueStack[num-1] = cont
	}
}

func (ctx *EngineContext) AddMemoryObject(o MemoryObject) {
	ctx.Mobjects = append(ctx.Mobjects, o)
}
<commit_msg>we need to stay in the caller's file until the call actually happens<commit_after>package whale

import (
	"strings"
	"time"
	"strconv"
)

import (
	"butler/null"
	"gokogiri/xpath"
	"golog"
	"rubex"
	tp "tritium/proto"
	"steno"
	"go-cache"
	"go-cache/arc"
)

type Whale struct {
	Log *golog.Logger
	Debugger steno.Debugger
	RegexpCache       cache.Cache
	XPathCache        cache.Cache
}

type EngineContext struct {
	Functions                []*Function
	Types                    []string
	Exports                  [][]string
	Logs                     []string
	ImportedFiles            []string
	Env                      map[string]string
	MatchStack               []string
	MatchShouldContinueStack []bool
	Yields                   []*YieldBlock
	*Whale
	*tp.Transform
	Rrules []*tp.RewriteRule

	InnerReplacer     *rubex.Regexp
	HeaderContentType *rubex.Regexp

	// Debug info
	Filename          string
	HadError          bool
	Deadline time.Time
	Mobjects []MemoryObject
	MessagePath string
}

const OutputBufferSize = 500 * 1024 //500KB
const defaultMobjects = 4
const TimeoutError = "EngineTimeout"

func NewEngine(logger *golog.Logger, debugger steno.Debugger) *Whale {
	e := &Whale{
		Log: logger,
		Debugger: debugger,
		RegexpCache:              arc.NewARCache(1000),
		XPathCache:               arc.NewARCache(1000),
	}
	e.RegexpCache.SetCleanFunc(CleanRegexpObject)
	e.XPathCache.SetCleanFunc(CleanXpathExpObject)
	return e
}

func NewEngineCtx(eng *Whale, vars map[string]string, transform *tp.Transform, rrules []*tp.RewriteRule, deadline time.Time, messagePath string) (ctx *EngineContext) {
	ctx = &EngineContext{
		Whale:                    eng,
		Exports:                  make([][]string, 0),
		Logs:                     make([]string, 0),
		Env:                      vars,
		Transform:                transform,
		Rrules:                   rrules,
		MatchStack:               make([]string, 0),
		MatchShouldContinueStack: make([]bool, 0),
		Yields:                   make([]*YieldBlock, 0),
		HadError:                 false,

		Deadline:                 deadline,
		Mobjects:                 make([]MemoryObject, 0, defaultMobjects),
		MessagePath:              messagePath,
	}
	return
}

func (eng *Whale) Free() {
	eng.RegexpCache.Reset()
	eng.XPathCache.Reset()
}

func (eng *Whale) Run(transform *tp.Transform, rrules []*tp.RewriteRule, input interface{}, vars map[string]string, deadline time.Time, customer, project, messagePath string) (output string, exports [][]string, logs []string) {
	ctx := NewEngineCtx(eng, vars, transform, rrules, deadline, messagePath)
	defer ctx.Free()
	ctx.Yields = append(ctx.Yields, &YieldBlock{Vars: make(map[string]interface{})})
	ctx.UsePackage(transform.Pkg)
	scope := &Scope{Value: input.(string)}
	obj := transform.Objects[0]
	ctx.Filename = null.GetString(obj.Name)
	ctx.RunInstruction(scope, obj.Root)
	output = scope.Value.(string)
	exports = ctx.Exports
	logs = ctx.Logs
	return
}

func (eng *Whale) GetCacheStats() (int, int, int, int) {
	return eng.RegexpCache.GetHitRate(), eng.RegexpCache.GetUsageRate(), eng.XPathCache.GetHitRate(), eng.XPathCache.GetUsageRate()
}

func (ctx *EngineContext) Free() {
    for _, o := range ctx.Mobjects {
		if o != nil {
			o.Free()
		}
	}
}

func (ctx *EngineContext) RunInstruction(scope *Scope, ins *tp.Instruction) (returnValue interface{}) {
	defer func() {
		//TODO Stack traces seem to get truncated on syslog...
		if x := recover(); x != nil {
			err, ok := x.(error)
			errString := ""
			if ok {
				errString = err.Error()
			} else {
				errString = x.(string)
			}
			if errString != TimeoutError {
				if ctx.HadError == false {
					ctx.HadError = true
					errString = errString + "\n" + ins.Type.String() + " " + null.GetString(ins.Value) + "\n\n\nTritium Stack\n=========\n\n"
				}
				errString = errString + ctx.FileAndLine(ins) + "\n"
			}
			panic(errString)
		}
	}()

	ctx.Whale.Debugger.TrapInstruction(ctx.MessagePath, ctx.Filename, ctx.Env, ins, scope.Value, scope.Index)

	if time.Now().After(ctx.Deadline) {
		panic(TimeoutError)
	}

	// If our object is invalid, then skip it
	if null.GetBool(ins.IsValid) == false {
		panic("Invalid instruction. Should have stopped before linking!")
	}
	indent := ""
	for i := 0; i < len(ctx.Yields); i++ {
		indent += "\t"
	}

	returnValue = ""
	switch *ins.Type {
	case tp.Instruction_BLOCK:
		for _, child := range ins.Children {
			returnValue = ctx.RunInstruction(scope, child)
		}
	case tp.Instruction_TEXT:
		returnValue = null.GetString(ins.Value)
	case tp.Instruction_LOCAL_VAR:
		name := null.GetString(ins.Value)
		vars := ctx.Vars()
		if len(ins.Arguments) > 0 {
			vars[name] = ctx.RunInstruction(scope, ins.Arguments[0])
		}
		if len(ins.Children) > 0 {
			ts := &Scope{Value: ctx.Vars()[name]}
			for _, child := range ins.Children {
				ctx.RunInstruction(ts, child)
			}
			vars[name] = ts.Value
		}
		returnValue = vars[name]
	case tp.Instruction_IMPORT:
		obj := ctx.Objects[int(null.GetInt32(ins.ObjectId))]
		curFile := ctx.Filename
		ctx.Filename = null.GetString(obj.Name)
		ctx.Whale.Debugger.LogImport(ctx.MessagePath, ctx.Filename, curFile, int(null.GetInt32(ins.LineNumber)))
		ctx.AddLog("__IMPORT__FILE__:" + ctx.Filename)
		for _, child := range obj.Root.Children {
			ctx.RunInstruction(scope, child)
		}
		ctx.Whale.Debugger.LogImportDone(ctx.MessagePath, ctx.Filename, curFile, int(null.GetInt32(ins.LineNumber)))
		ctx.Filename = curFile
	case tp.Instruction_FUNCTION_CALL:
		fun := ctx.Functions[int(null.GetInt32(ins.FunctionId))]
		args := make([]interface{}, len(ins.Arguments))
		for i, argIns := range ins.Arguments {
			args[i] = ctx.RunInstruction(scope, argIns)
		}

		if null.GetBool(fun.BuiltIn) {
			if f := builtInFunctions[fun.Name]; f != nil {
				returnValue = f(ctx, scope, ins, args)
				if returnValue == nil {
					returnValue = ""
				}
			} else {
				panic("missing function: " + fun.Name)
			}
		} else {
			// We are using a user-defined function
			//println("Resetting localvar")
			// Setup the new local var
			vars := make(map[string]interface{}, len(args))
			for i, arg := range fun.Args {
				vars[null.GetString(arg.Name)] = args[i]
			}
			yieldBlock := &YieldBlock{
				Ins:  ins,
				Vars: vars,
			}
			// PUSH!
			ctx.PushYieldBlock(yieldBlock)
			curFile := ctx.Filename
			ctx.Filename = fun.GetFilename()
			for _, child := range fun.Instruction.Children {
				returnValue = ctx.RunInstruction(scope, child)
			}
			ctx.Filename = curFile
			// POP!
			ctx.PopYieldBlock()

		}
	}

	return
}

func (ctx *EngineContext) ShouldContinue() (result bool) {
	if len(ctx.MatchShouldContinueStack) > 0 {
		result = ctx.MatchShouldContinueStack[len(ctx.MatchShouldContinueStack)-1]
	} else {
		result = false
	}
	return
}

func (ctx *EngineContext) MatchTarget() string {
	return ctx.MatchStack[len(ctx.MatchStack)-1]
}

func (ctx *EngineContext) PushYieldBlock(b *YieldBlock) {
	ctx.Yields = append(ctx.Yields, b)
}

func (ctx *EngineContext) PopYieldBlock() (b *YieldBlock) {
	num := len(ctx.Yields)
	if num > 0 {
		b = ctx.Yields[num-1]
		ctx.Yields = ctx.Yields[:num-1]
	}
	return
}

func (ctx *EngineContext) HasYieldBlock() bool {
	return len(ctx.Yields) > 0
}

func (ctx *EngineContext) TopYieldBlock() (b *YieldBlock) {
	num := len(ctx.Yields)
	if num > 0 {
		b = ctx.Yields[num-1]
	}
	return
}

func (ctx *EngineContext) Vars() map[string]interface{} {
	b := ctx.TopYieldBlock()
	if b != nil {
		return b.Vars
	}
	return nil
}

func (ctx *EngineContext) FileAndLine(ins *tp.Instruction) string {
	lineNum := strconv.Itoa(int(null.GetInt32(ins.LineNumber)))
	return (ctx.Filename + ":" + lineNum)
}

func (ctx *EngineContext) UsePackage(pkg *tp.Package) {
	ctx.Types = make([]string, len(pkg.Types))
	for i, t := range pkg.Types {
		ctx.Types[i] = null.GetString(t.Name)
	}

	ctx.Functions = make([]*Function, len(pkg.Functions))
	for i, f := range pkg.Functions {
		name := null.GetString(f.Name)
		for _, a := range f.Args {
			typeString := ctx.Types[int(null.GetInt32(a.TypeId))]
			name = name + "." + typeString
		}
		fun := &Function{
			Name:     name,
			Function: f,
		}
		ctx.Functions[i] = fun
	}
}

func (ctx *EngineContext) GetRegexp(pattern, options string) (r *rubex.Regexp) {
	sig := pattern + "/" + options
	object, err := ctx.RegexpCache.Get(sig)
	if err != nil {
		mode := rubex.ONIG_OPTION_DEFAULT
		if strings.Index(options, "i") >= 0 {
			mode = rubex.ONIG_OPTION_IGNORECASE
		}
		if strings.Index(options, "m") >= 0 {
			mode = rubex.ONIG_OPTION_MULTILINE
		}
		var err error
		r, err = rubex.NewRegexp(pattern, mode)
		if err == nil {
			//ctx.AddMemoryObject(r)
			ctx.RegexpCache.Set(sig, &RegexpObject{Regexp:r})
		}
		return r
	}
	return object.(*RegexpObject).Regexp
}

func (ctx *EngineContext) GetXpathExpr(p string) (e *xpath.Expression) {
	object, err := ctx.XPathCache.Get(p)
	if err != nil {
		e = xpath.Compile(p)
		if e != nil {
			//ctx.AddMemoryObject(e)
			ctx.XPathCache.Set(p, &XpathExpObject{Expression: e})
		} else {
			ctx.AddLog("Invalid XPath used: " + p)
		}
		return e
	}
	return object.(*XpathExpObject).Expression
}

func (ctx *EngineContext) AddExport(exports []string) {
	ctx.Exports = append(ctx.Exports, exports)
}

func (ctx *EngineContext) AddLog(log string) int {
	//ctx.Log.Info("TRITIUM: " + log)
	index := len(ctx.Logs)
	ctx.Logs = append(ctx.Logs, log)
	return index
}

func (ctx *EngineContext) UpdateLog(index int, log string) {
	//ctx.Log.Info("TRITIUM: " + log)
	if index >= 0 && index < len(ctx.Logs) {
		ctx.Logs[index] = log
	}
}

func (ctx *EngineContext) SetEnv(key, val string) {
	ctx.Env[key] = val
}

func (ctx *EngineContext) GetEnv(key string) (val string) {
	val = ctx.Env[key]
	return
}

func (ctx *EngineContext) SetVar(key string, val interface{}) {
	b := ctx.TopYieldBlock()
	if b != nil {
		b.Vars[key] = val
	}
}

func (ctx *EngineContext) GetVar(key string) (val interface{}) {
	b := ctx.TopYieldBlock()
	if b != nil {
		val = b.Vars[key]
	}
	return
}

func (ctx *EngineContext) GetInnerReplacer() (r *rubex.Regexp) {
	return ctx.GetRegexp(`[\\$](\d)`, "i")
}

func (ctx *EngineContext) GetHeaderContentTypeRegex() (r *rubex.Regexp) {
	return ctx.GetRegexp(`<meta\s+http-equiv="content-type"\s+content="(.*?)"`, "i")
}

func (ctx *EngineContext) GetOutputBuffer() (b []byte) {
	//b = ctx.OutputBuffer
	return
}

func (ctx *EngineContext) Logger() (logger *golog.Logger) {
	logger = ctx.Log
	return
}

func (ctx *EngineContext) PushMatchStack(match string) {
	ctx.MatchStack = append(ctx.MatchStack, match)
}

func (ctx *EngineContext) PopMatchStack() (match string) {
	if num := len(ctx.MatchStack); num > 0 {
		match = ctx.MatchStack[num-1]
		ctx.MatchStack = ctx.MatchStack[:num-1]
	}
	return
}

func (ctx *EngineContext) PushShouldContinueStack(cont bool) {
	ctx.MatchShouldContinueStack = append(ctx.MatchShouldContinueStack, cont)
}

func (ctx *EngineContext) PopShouldContinueStack() (cont bool) {
	if num := len(ctx.MatchShouldContinueStack); num > 0 {
		cont = ctx.MatchShouldContinueStack[num-1]
		ctx.MatchShouldContinueStack = ctx.MatchShouldContinueStack[:num-1]
	}
	return
}

func (ctx *EngineContext) SetShouldContinue(cont bool) {
	if num := len(ctx.MatchShouldContinueStack); num > 0 {
		ctx.MatchShouldContinueStack[num-1] = cont
	}
}

func (ctx *EngineContext) AddMemoryObject(o MemoryObject) {
	ctx.Mobjects = append(ctx.Mobjects, o)
}
<|endoftext|>


<commit_before>// Copyright 2014 Vic Demuzere
//
// Use of this source code is governed by the MIT license.

package irc

import (
	"reflect"
	"testing"
)

var messageTests = [...]*struct {
	parsed     *Message
	rawMessage string
	rawPrefix  string
	hostmask   bool // Is it very clear that the prefix is a hostname?
	server     bool // Is the prefix a servername?
}{
	{
		parsed: &Message{
			Prefix: &Prefix{
				Name: "syrk",
				User: "kalt",
				Host: "millennium.stealth.net",
			},
			Command:  "QUIT",
			Trailing: "Gone to have lunch",
		},
		rawMessage: ":syrk!kalt@millennium.stealth.net QUIT :Gone to have lunch",
		rawPrefix:  "syrk!kalt@millennium.stealth.net",
		hostmask:   true,
	},
	{
		parsed: &Message{
			Prefix: &Prefix{
				Name: "Trillian",
			},
			Command:  "SQUIT",
			Params:   []string{"cm22.eng.umd.edu"},
			Trailing: "Server out of control",
		},
		rawMessage: ":Trillian SQUIT cm22.eng.umd.edu :Server out of control",
		rawPrefix:  "Trillian",
		server:     true,
	},
	{
		parsed: &Message{
			Prefix: &Prefix{
				Name: "WiZ",
				User: "jto",
				Host: "tolsun.oulu.fi",
			},
			Command: "JOIN",
			Params:  []string{"#Twilight_zone"},
		},
		rawMessage: ":WiZ!jto@tolsun.oulu.fi JOIN #Twilight_zone",
		rawPrefix:  "WiZ!jto@tolsun.oulu.fi",
		hostmask:   true,
	},
	{
		parsed: &Message{
			Prefix: &Prefix{
				Name: "WiZ",
				User: "jto",
				Host: "tolsun.oulu.fi",
			},
			Command:  "PART",
			Params:   []string{"#playzone"},
			Trailing: "I lost",
		},
		rawMessage: ":WiZ!jto@tolsun.oulu.fi PART #playzone :I lost",
		rawPrefix:  "WiZ!jto@tolsun.oulu.fi",
		hostmask:   true,
	},
	{
		parsed: &Message{
			Prefix: &Prefix{
				Name: "WiZ",
				User: "jto",
				Host: "tolsun.oulu.fi",
			},
			Command: "MODE",
			Params:  []string{"#eu-opers", "-l"},
		},
		rawMessage: ":WiZ!jto@tolsun.oulu.fi MODE #eu-opers -l",
		rawPrefix:  "WiZ!jto@tolsun.oulu.fi",
		hostmask:   true,
	},
	{
		parsed: &Message{
			Command: "MODE",
			Params:  []string{"&oulu", "+b", "*!*@*.edu", "+e", "*!*@*.bu.edu"},
		},
		rawMessage: "MODE &oulu +b *!*@*.edu +e *!*@*.bu.edu",
	},
	{
		parsed: &Message{
			Command:  "PRIVMSG",
			Params:   []string{"#channel"},
			Trailing: "Message with :colons!",
		},
		rawMessage: "PRIVMSG #channel :Message with :colons!",
	},
	{
		parsed: &Message{
			Prefix: &Prefix{
				Name: "irc.vives.lan",
			},
			Command:  "251",
			Params:   []string{"test"},
			Trailing: "There are 2 users and 0 services on 1 servers",
		},
		rawMessage: ":irc.vives.lan 251 test :There are 2 users and 0 services on 1 servers",
		rawPrefix:  "irc.vives.lan",
		server:     true,
	},
	{
		parsed: &Message{
			Prefix: &Prefix{
				Name: "irc.vives.lan",
			},
			Command:  "376",
			Params:   []string{"test"},
			Trailing: "End of MOTD command",
		},
		rawMessage: ":irc.vives.lan 376 test :End of MOTD command",
		rawPrefix:  "irc.vives.lan",
		server:     true,
	},
	{
		parsed: &Message{
			Prefix: &Prefix{
				Name: "irc.vives.lan",
			},
			Command:  "250",
			Params:   []string{"test"},
			Trailing: "Highest connection count: 1 (1 connections received)",
		},
		rawMessage: ":irc.vives.lan 250 test :Highest connection count: 1 (1 connections received)",
		rawPrefix:  "irc.vives.lan",
		server:     true,
	},
	{
		parsed: &Message{
			Prefix: &Prefix{
				Name: "sorcix",
				User: "~sorcix",
				Host: "sorcix.users.quakenet.org",
			},
			Command:  "PRIVMSG",
			Params:   []string{"#viveslan"},
			Trailing: "\001ACTION is testing CTCP messages!\001",
		},
		rawMessage: ":sorcix!~sorcix@sorcix.users.quakenet.org PRIVMSG #viveslan :\001ACTION is testing CTCP messages!\001",
		rawPrefix:  "sorcix!~sorcix@sorcix.users.quakenet.org",
		hostmask:   true,
	},
	{
		parsed: &Message{
			Prefix: &Prefix{
				Name: "sorcix",
				User: "~sorcix",
				Host: "sorcix.users.quakenet.org",
			},
			Command:  "NOTICE",
			Params:   []string{"midnightfox"},
			Trailing: "\001PONG 1234567890\001",
		},
		rawMessage: ":sorcix!~sorcix@sorcix.users.quakenet.org NOTICE midnightfox :\001PONG 1234567890\001",
		rawPrefix:  "sorcix!~sorcix@sorcix.users.quakenet.org",
		hostmask:   true,
	},
	{
		parsed: &Message{
			Prefix: &Prefix{
				Name: "a",
				User: "b",
				Host: "c",
			},
			Command: "QUIT",
		},
		rawMessage: ":a!b@c QUIT",
		rawPrefix:  "a!b@c",
		hostmask:   true,
	},
	{
		parsed: &Message{
			Prefix: &Prefix{
				Name: "a",
				User: "b",
			},
			Command:  "PRIVMSG",
			Trailing: "message",
		},
		rawMessage: ":a!b PRIVMSG :message",
		rawPrefix:  "a!b",
	},
	{
		parsed: &Message{
			Prefix: &Prefix{
				Name: "a",
				Host: "c",
			},
			Command:  "NOTICE",
			Trailing: ":::Hey!",
		},
		rawMessage: ":a@c NOTICE ::::Hey!",
		rawPrefix:  "a@c",
	},
	{
		parsed: &Message{
			Prefix: &Prefix{
				Name: "nick",
			},
			Command:  "PRIVMSG",
			Params:   []string{"$@"},
			Trailing: "This message contains a\ttab!",
		},
		rawMessage: ":nick PRIVMSG $@ :This message contains a\ttab!",
		rawPrefix:  "nick",
	},
	{
		parsed: &Message{
			Command:  "TEST",
			Params:   []string{"$@", "", "param"},
			Trailing: "Trailing",
		},
		rawMessage: "TEST $@  param :Trailing",
	},
	{
		rawMessage: ": PRIVMSG test :Invalid message with empty prefix.",
		rawPrefix:  "",
	},
	{
		rawMessage: ":  PRIVMSG test :Invalid message with space prefix",
		rawPrefix:  " ",
	},
	{
		parsed: &Message{
			Command:       "TOPIC",
			Params:        []string{"#foo"},
			Trailing:      "",
		},
		rawMessage: "TOPIC #foo",
		rawPrefix:  "",
	},
	{
		parsed: &Message{
			Command:       "TOPIC",
			Params:        []string{"#foo"},
			Trailing:      "",
			EmptyTrailing: true,
		},
		rawMessage: "TOPIC #foo :",
		rawPrefix:  "",
	},
}

// -----
// PREFIX
// -----

func TestPrefix_IsHostmask(t *testing.T) {

	for i, test := range messageTests {

		// Skip tests that have no prefix
		if test.parsed == nil || test.parsed.Prefix == nil {
			continue
		}

		if test.hostmask && !test.parsed.Prefix.IsHostmask() {
			t.Errorf("Prefix %d should be recognized as a hostmask!", i)
		}

	}
}

func TestPrefix_IsServer(t *testing.T) {

	for i, test := range messageTests {

		// Skip tests that have no prefix
		if test.parsed == nil || test.parsed.Prefix == nil {
			continue
		}

		if test.server && !test.parsed.Prefix.IsServer() {
			t.Errorf("Prefix %d should be recognized as a server!", i)
		}

	}
}

func TestPrefix_String(t *testing.T) {
	var s string

	for i, test := range messageTests {

		// Skip tests that have no prefix
		if test.parsed == nil || test.parsed.Prefix == nil {
			continue
		}

		// Convert the prefix
		s = test.parsed.Prefix.String()

		// Result should be the same as the value in rawMessage.
		if s != test.rawPrefix {
			t.Errorf("Failed to stringify prefix %d:", i)
			t.Logf("Output: %s", s)
			t.Logf("Expected: %s", test.rawPrefix)
		}
	}
}

func TestPrefix_Len(t *testing.T) {
	var l int

	for i, test := range messageTests {

		// Skip tests that have no prefix
		if test.parsed == nil || test.parsed.Prefix == nil {
			continue
		}

		l = test.parsed.Prefix.Len()

		// Result should be the same as the value in rawMessage.
		if l != len(test.rawPrefix) {
			t.Errorf("Failed to calculate prefix length %d:", i)
			t.Logf("Output: %d", l)
			t.Logf("Expected: %d", len(test.rawPrefix))
		}
	}
}

func TestParsePrefix(t *testing.T) {
	var p *Prefix

	for i, test := range messageTests {

		// Skip tests that have no prefix
		if test.parsed == nil || test.parsed.Prefix == nil {
			continue
		}

		// Parse the prefix
		p = ParsePrefix(test.rawPrefix)

		// Result struct should be the same as the value in parsed.
		if *p != *test.parsed.Prefix {
			t.Errorf("Failed to parse prefix %d:", i)
			t.Logf("Output: %#v", p)
			t.Logf("Expected: %#v", test.parsed.Prefix)
		}
	}
}

// -----
// MESSAGE
// -----

func TestMessage_String(t *testing.T) {
	var s string

	for i, test := range messageTests {

		// Skip tests that have no valid struct
		if test.parsed == nil {
			continue
		}

		// Convert the prefix
		s = test.parsed.String()

		// Result should be the same as the value in rawMessage.
		if s != test.rawMessage {
			t.Errorf("Failed to stringify message %d:", i)
			t.Logf("Output: %s", s)
			t.Logf("Expected: %s", test.rawMessage)
		}
	}
}

func TestMessage_Len(t *testing.T) {
	var l int

	for i, test := range messageTests {

		// Skip tests that have no valid struct
		if test.parsed == nil {
			continue
		}

		l = test.parsed.Len()

		// Result should be the same as the value in rawMessage.
		if l != len(test.rawMessage) {
			t.Errorf("Failed to calculate message length %d:", i)
			t.Logf("Output: %d", l)
			t.Logf("Expected: %d", len(test.rawMessage))
		}
	}
}

func TestParseMessage(t *testing.T) {
	var p *Message

	for i, test := range messageTests {

		// Parse the prefix
		p = ParseMessage(test.rawMessage)

		// Result struct should be the same as the value in parsed.
		if !reflect.DeepEqual(p, test.parsed) {
			t.Errorf("Failed to parse message %d:", i)
			t.Logf("Output: %#v", p)
			t.Logf("Expected: %#v", test.parsed)
		}
	}
}

// -----
// BENCHMARK
// -----

func BenchmarkPrefix_String_short(b *testing.B) {
	b.ReportAllocs()

	prefix := new(Prefix)
	prefix.Name = "Namename"

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		prefix.String()
	}
}
func BenchmarkPrefix_String_long(b *testing.B) {
	b.ReportAllocs()

	prefix := new(Prefix)
	prefix.Name = "Namename"
	prefix.User = "Username"
	prefix.Host = "Hostname"

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		prefix.String()
	}
}
func BenchmarkParsePrefix_short(b *testing.B) {
	b.ReportAllocs()

	for i := 0; i < b.N; i++ {
		ParsePrefix("Namename")
	}
}
func BenchmarkParsePrefix_long(b *testing.B) {
	b.ReportAllocs()

	for i := 0; i < b.N; i++ {
		ParsePrefix("Namename!Username@Hostname")
	}
}
func BenchmarkMessage_String(b *testing.B) {
	b.ReportAllocs()

	for i := 0; i < b.N; i++ {
		messageTests[0].parsed.String()
	}
}
func BenchmarkParseMessage_short(b *testing.B) {
	b.ReportAllocs()

	for i := 0; i < b.N; i++ {
		ParseMessage("COMMAND arg1 :Message\r\n")
	}
}
func BenchmarkParseMessage_medium(b *testing.B) {
	b.ReportAllocs()

	for i := 0; i < b.N; i++ {
		ParseMessage(":Namename COMMAND arg6 arg7 :Message message message\r\n")
	}
}
func BenchmarkParseMessage_long(b *testing.B) {
	b.ReportAllocs()

	for i := 0; i < b.N; i++ {
		ParseMessage(":Namename!username@hostname COMMAND arg1 arg2 arg3 arg4 arg5 arg6 arg7 :Message message message message message\r\n")
	}
}
<commit_msg>Additional test that decodes and encodes again.<commit_after>// Copyright 2014 Vic Demuzere
//
// Use of this source code is governed by the MIT license.

package irc

import (
	"reflect"
	"testing"
)

var messageTests = [...]*struct {
	parsed     *Message
	rawMessage string
	rawPrefix  string
	hostmask   bool // Is it very clear that the prefix is a hostname?
	server     bool // Is the prefix a servername?
}{
	{
		parsed: &Message{
			Prefix: &Prefix{
				Name: "syrk",
				User: "kalt",
				Host: "millennium.stealth.net",
			},
			Command:  "QUIT",
			Trailing: "Gone to have lunch",
		},
		rawMessage: ":syrk!kalt@millennium.stealth.net QUIT :Gone to have lunch",
		rawPrefix:  "syrk!kalt@millennium.stealth.net",
		hostmask:   true,
	},
	{
		parsed: &Message{
			Prefix: &Prefix{
				Name: "Trillian",
			},
			Command:  "SQUIT",
			Params:   []string{"cm22.eng.umd.edu"},
			Trailing: "Server out of control",
		},
		rawMessage: ":Trillian SQUIT cm22.eng.umd.edu :Server out of control",
		rawPrefix:  "Trillian",
		server:     true,
	},
	{
		parsed: &Message{
			Prefix: &Prefix{
				Name: "WiZ",
				User: "jto",
				Host: "tolsun.oulu.fi",
			},
			Command: "JOIN",
			Params:  []string{"#Twilight_zone"},
		},
		rawMessage: ":WiZ!jto@tolsun.oulu.fi JOIN #Twilight_zone",
		rawPrefix:  "WiZ!jto@tolsun.oulu.fi",
		hostmask:   true,
	},
	{
		parsed: &Message{
			Prefix: &Prefix{
				Name: "WiZ",
				User: "jto",
				Host: "tolsun.oulu.fi",
			},
			Command:  "PART",
			Params:   []string{"#playzone"},
			Trailing: "I lost",
		},
		rawMessage: ":WiZ!jto@tolsun.oulu.fi PART #playzone :I lost",
		rawPrefix:  "WiZ!jto@tolsun.oulu.fi",
		hostmask:   true,
	},
	{
		parsed: &Message{
			Prefix: &Prefix{
				Name: "WiZ",
				User: "jto",
				Host: "tolsun.oulu.fi",
			},
			Command: "MODE",
			Params:  []string{"#eu-opers", "-l"},
		},
		rawMessage: ":WiZ!jto@tolsun.oulu.fi MODE #eu-opers -l",
		rawPrefix:  "WiZ!jto@tolsun.oulu.fi",
		hostmask:   true,
	},
	{
		parsed: &Message{
			Command: "MODE",
			Params:  []string{"&oulu", "+b", "*!*@*.edu", "+e", "*!*@*.bu.edu"},
		},
		rawMessage: "MODE &oulu +b *!*@*.edu +e *!*@*.bu.edu",
	},
	{
		parsed: &Message{
			Command:  "PRIVMSG",
			Params:   []string{"#channel"},
			Trailing: "Message with :colons!",
		},
		rawMessage: "PRIVMSG #channel :Message with :colons!",
	},
	{
		parsed: &Message{
			Prefix: &Prefix{
				Name: "irc.vives.lan",
			},
			Command:  "251",
			Params:   []string{"test"},
			Trailing: "There are 2 users and 0 services on 1 servers",
		},
		rawMessage: ":irc.vives.lan 251 test :There are 2 users and 0 services on 1 servers",
		rawPrefix:  "irc.vives.lan",
		server:     true,
	},
	{
		parsed: &Message{
			Prefix: &Prefix{
				Name: "irc.vives.lan",
			},
			Command:  "376",
			Params:   []string{"test"},
			Trailing: "End of MOTD command",
		},
		rawMessage: ":irc.vives.lan 376 test :End of MOTD command",
		rawPrefix:  "irc.vives.lan",
		server:     true,
	},
	{
		parsed: &Message{
			Prefix: &Prefix{
				Name: "irc.vives.lan",
			},
			Command:  "250",
			Params:   []string{"test"},
			Trailing: "Highest connection count: 1 (1 connections received)",
		},
		rawMessage: ":irc.vives.lan 250 test :Highest connection count: 1 (1 connections received)",
		rawPrefix:  "irc.vives.lan",
		server:     true,
	},
	{
		parsed: &Message{
			Prefix: &Prefix{
				Name: "sorcix",
				User: "~sorcix",
				Host: "sorcix.users.quakenet.org",
			},
			Command:  "PRIVMSG",
			Params:   []string{"#viveslan"},
			Trailing: "\001ACTION is testing CTCP messages!\001",
		},
		rawMessage: ":sorcix!~sorcix@sorcix.users.quakenet.org PRIVMSG #viveslan :\001ACTION is testing CTCP messages!\001",
		rawPrefix:  "sorcix!~sorcix@sorcix.users.quakenet.org",
		hostmask:   true,
	},
	{
		parsed: &Message{
			Prefix: &Prefix{
				Name: "sorcix",
				User: "~sorcix",
				Host: "sorcix.users.quakenet.org",
			},
			Command:  "NOTICE",
			Params:   []string{"midnightfox"},
			Trailing: "\001PONG 1234567890\001",
		},
		rawMessage: ":sorcix!~sorcix@sorcix.users.quakenet.org NOTICE midnightfox :\001PONG 1234567890\001",
		rawPrefix:  "sorcix!~sorcix@sorcix.users.quakenet.org",
		hostmask:   true,
	},
	{
		parsed: &Message{
			Prefix: &Prefix{
				Name: "a",
				User: "b",
				Host: "c",
			},
			Command: "QUIT",
		},
		rawMessage: ":a!b@c QUIT",
		rawPrefix:  "a!b@c",
		hostmask:   true,
	},
	{
		parsed: &Message{
			Prefix: &Prefix{
				Name: "a",
				User: "b",
			},
			Command:  "PRIVMSG",
			Trailing: "message",
		},
		rawMessage: ":a!b PRIVMSG :message",
		rawPrefix:  "a!b",
	},
	{
		parsed: &Message{
			Prefix: &Prefix{
				Name: "a",
				Host: "c",
			},
			Command:  "NOTICE",
			Trailing: ":::Hey!",
		},
		rawMessage: ":a@c NOTICE ::::Hey!",
		rawPrefix:  "a@c",
	},
	{
		parsed: &Message{
			Prefix: &Prefix{
				Name: "nick",
			},
			Command:  "PRIVMSG",
			Params:   []string{"$@"},
			Trailing: "This message contains a\ttab!",
		},
		rawMessage: ":nick PRIVMSG $@ :This message contains a\ttab!",
		rawPrefix:  "nick",
	},
	{
		parsed: &Message{
			Command:  "TEST",
			Params:   []string{"$@", "", "param"},
			Trailing: "Trailing",
		},
		rawMessage: "TEST $@  param :Trailing",
	},
	{
		rawMessage: ": PRIVMSG test :Invalid message with empty prefix.",
		rawPrefix:  "",
	},
	{
		rawMessage: ":  PRIVMSG test :Invalid message with space prefix",
		rawPrefix:  " ",
	},
	{
		parsed: &Message{
			Command:       "TOPIC",
			Params:        []string{"#foo"},
			Trailing:      "",
		},
		rawMessage: "TOPIC #foo",
		rawPrefix:  "",
	},
	{
		parsed: &Message{
			Command:       "TOPIC",
			Params:        []string{"#foo"},
			Trailing:      "",
			EmptyTrailing: true,
		},
		rawMessage: "TOPIC #foo :",
		rawPrefix:  "",
	},
}

// -----
// PREFIX
// -----

func TestPrefix_IsHostmask(t *testing.T) {

	for i, test := range messageTests {

		// Skip tests that have no prefix
		if test.parsed == nil || test.parsed.Prefix == nil {
			continue
		}

		if test.hostmask && !test.parsed.Prefix.IsHostmask() {
			t.Errorf("Prefix %d should be recognized as a hostmask!", i)
		}

	}
}

func TestPrefix_IsServer(t *testing.T) {

	for i, test := range messageTests {

		// Skip tests that have no prefix
		if test.parsed == nil || test.parsed.Prefix == nil {
			continue
		}

		if test.server && !test.parsed.Prefix.IsServer() {
			t.Errorf("Prefix %d should be recognized as a server!", i)
		}

	}
}

func TestPrefix_String(t *testing.T) {
	var s string

	for i, test := range messageTests {

		// Skip tests that have no prefix
		if test.parsed == nil || test.parsed.Prefix == nil {
			continue
		}

		// Convert the prefix
		s = test.parsed.Prefix.String()

		// Result should be the same as the value in rawMessage.
		if s != test.rawPrefix {
			t.Errorf("Failed to stringify prefix %d:", i)
			t.Logf("Output: %s", s)
			t.Logf("Expected: %s", test.rawPrefix)
		}
	}
}

func TestPrefix_Len(t *testing.T) {
	var l int

	for i, test := range messageTests {

		// Skip tests that have no prefix
		if test.parsed == nil || test.parsed.Prefix == nil {
			continue
		}

		l = test.parsed.Prefix.Len()

		// Result should be the same as the value in rawMessage.
		if l != len(test.rawPrefix) {
			t.Errorf("Failed to calculate prefix length %d:", i)
			t.Logf("Output: %d", l)
			t.Logf("Expected: %d", len(test.rawPrefix))
		}
	}
}

func TestParsePrefix(t *testing.T) {
	var p *Prefix

	for i, test := range messageTests {

		// Skip tests that have no prefix
		if test.parsed == nil || test.parsed.Prefix == nil {
			continue
		}

		// Parse the prefix
		p = ParsePrefix(test.rawPrefix)

		// Result struct should be the same as the value in parsed.
		if *p != *test.parsed.Prefix {
			t.Errorf("Failed to parse prefix %d:", i)
			t.Logf("Output: %#v", p)
			t.Logf("Expected: %#v", test.parsed.Prefix)
		}
	}
}

// -----
// MESSAGE
// -----

func TestMessage_String(t *testing.T) {
	var s string

	for i, test := range messageTests {

		// Skip tests that have no valid struct
		if test.parsed == nil {
			continue
		}

		// Convert the prefix
		s = test.parsed.String()

		// Result should be the same as the value in rawMessage.
		if s != test.rawMessage {
			t.Errorf("Failed to stringify message %d:", i)
			t.Logf("Output: %s", s)
			t.Logf("Expected: %s", test.rawMessage)
		}
	}
}

func TestMessage_Len(t *testing.T) {
	var l int

	for i, test := range messageTests {

		// Skip tests that have no valid struct
		if test.parsed == nil {
			continue
		}

		l = test.parsed.Len()

		// Result should be the same as the value in rawMessage.
		if l != len(test.rawMessage) {
			t.Errorf("Failed to calculate message length %d:", i)
			t.Logf("Output: %d", l)
			t.Logf("Expected: %d", len(test.rawMessage))
		}
	}
}

func TestParseMessage(t *testing.T) {
	var p *Message

	for i, test := range messageTests {

		// Parse the prefix
		p = ParseMessage(test.rawMessage)

		// Result struct should be the same as the value in parsed.
		if !reflect.DeepEqual(p, test.parsed) {
			t.Errorf("Failed to parse message %d:", i)
			t.Logf("Output: %#v", p)
			t.Logf("Expected: %#v", test.parsed)
		}
	}
}

// -----
// MESSAGE DECODE -> ENCODE
// -----

func TestMessageDecodeEncode(t *testing.T) {
	var (
		p *Message
		s string
	)

	for i, test := range messageTests {

		// Skip invalid messages
		if test.parsed == nil {
			continue
		}

		// Decode the message, then encode it again.
		p = ParseMessage(test.rawMessage)
		s = p.String()

		// Result struct should be the same as the original.
		if s != test.rawMessage {
			t.Errorf("Message %d failed decode-encode sequence!", i)
		}
	}
}

// -----
// BENCHMARK
// -----

func BenchmarkPrefix_String_short(b *testing.B) {
	b.ReportAllocs()

	prefix := new(Prefix)
	prefix.Name = "Namename"

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		prefix.String()
	}
}
func BenchmarkPrefix_String_long(b *testing.B) {
	b.ReportAllocs()

	prefix := new(Prefix)
	prefix.Name = "Namename"
	prefix.User = "Username"
	prefix.Host = "Hostname"

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		prefix.String()
	}
}
func BenchmarkParsePrefix_short(b *testing.B) {
	b.ReportAllocs()

	for i := 0; i < b.N; i++ {
		ParsePrefix("Namename")
	}
}
func BenchmarkParsePrefix_long(b *testing.B) {
	b.ReportAllocs()

	for i := 0; i < b.N; i++ {
		ParsePrefix("Namename!Username@Hostname")
	}
}
func BenchmarkMessage_String(b *testing.B) {
	b.ReportAllocs()

	for i := 0; i < b.N; i++ {
		messageTests[0].parsed.String()
	}
}
func BenchmarkParseMessage_short(b *testing.B) {
	b.ReportAllocs()

	for i := 0; i < b.N; i++ {
		ParseMessage("COMMAND arg1 :Message\r\n")
	}
}
func BenchmarkParseMessage_medium(b *testing.B) {
	b.ReportAllocs()

	for i := 0; i < b.N; i++ {
		ParseMessage(":Namename COMMAND arg6 arg7 :Message message message\r\n")
	}
}
func BenchmarkParseMessage_long(b *testing.B) {
	b.ReportAllocs()

	for i := 0; i < b.N; i++ {
		ParseMessage(":Namename!username@hostname COMMAND arg1 arg2 arg3 arg4 arg5 arg6 arg7 :Message message message message message\r\n")
	}
}
<|endoftext|>


<commit_before>// Package counter ...
package counter

import (
	"github.com/nanobox-io/nanobox/models"
	"github.com/nanobox-io/nanobox/util/locker"
)

// Increment will increment a counter by ID
func Increment(id string) (int, error) {
	locker.GlobalLock()
	defer locker.GlobalUnlock()

	counter, _ := models.FindCounterByID(id)

	//
	if counter.Count < 0 {
		counter.Count = 0
	}

	counter.Count = counter.Count + 1

	//
	err := counter.Save()

	return counter.Count, err
}

// Decrement will decrement a counter by ID
func Decrement(id string) (int, error) {
	locker.GlobalLock()
	defer locker.GlobalUnlock()

	counter, _ := models.FindCounterByID(id)

	//
	if counter.Count <= 0 {
		counter.Count = 0
	} else {
		counter.Count = counter.Count - 1
	}

	//
	err := counter.Save()

	return counter.Count, err
}

// Get returns the current counter value by ID
func Get(id string) (int, error) {

	counter, err := models.FindCounterByID(id)

	return counter.Count, err
}

// Reset resets a counter by ID; we can simply delete the counter model
func Reset(id string) error {
	locker.GlobalLock()
	defer locker.GlobalUnlock()

	counter := models.Counter{ID: id}
	return counter.Delete()
}
<commit_msg>Fix counter increment and decrement actions<commit_after>// Package counter ...
package counter

import (
	"github.com/nanobox-io/nanobox/models"
	"github.com/nanobox-io/nanobox/util/locker"
)

// Increment will increment a counter by ID
func Increment(id string) (int, error) {
	locker.GlobalLock()
	defer locker.GlobalUnlock()

	counter, _ := models.FindCounterByID(id)

	// set the ID in case the counter was empty
	counter.ID = id

	//
	if counter.Count < 0 {
		counter.Count = 0
	}

	counter.Count = counter.Count + 1

	//
	err := counter.Save()

	return counter.Count, err
}

// Decrement will decrement a counter by ID
func Decrement(id string) (int, error) {
	locker.GlobalLock()
	defer locker.GlobalUnlock()

	counter, _ := models.FindCounterByID(id)

	// set the ID in case the counter was empty
	counter.ID = id

	//
	if counter.Count <= 0 {
		counter.Count = 0
	} else {
		counter.Count = counter.Count - 1
	}

	//
	err := counter.Save()

	return counter.Count, err
}

// Get returns the current counter value by ID
func Get(id string) (int, error) {

	counter, err := models.FindCounterByID(id)

	return counter.Count, err
}

// Reset resets a counter by ID; we can simply delete the counter model
func Reset(id string) error {
	locker.GlobalLock()
	defer locker.GlobalUnlock()

	counter := models.Counter{ID: id}
	return counter.Delete()
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"log"
	"strings"

	"github.com/fweikert/continuous-integration/metrics/metrics"

	"github.com/fweikert/continuous-integration/metrics/clients"
	"github.com/fweikert/continuous-integration/metrics/publishers"
	"github.com/fweikert/continuous-integration/metrics/service"
)

var (
	bkOrg          = flag.String("buildkite_org", "bazel", "Buildkite organization slug")
	bkApiToken     = flag.String("buildkite_token", "", "Buildkite API access token that grants read access. See https://buildkite.com/docs/apis/rest-api#authentication")
	bkDebug        = flag.Bool("debug", false, "Enable debugging")
	pipelineString = flag.String("pipelines", "", "Comma separated list of slugs of pipelines whose performance statistics should be exported.")
	ghOrg          = flag.String("github_org", "bazelbuild", "Name of the GitHub organization.")
	ghRepo         = flag.String("github_repo", "bazelbuild", "Name of the GitHub repository.")
	ghApiToken     = flag.String("github_token", "", "Access token for the GitHub API.")
	sqlUser        = flag.String("sql_user", "", "User name for the CloudSQL publisher.")
	sqlPassword    = flag.String("sql_password", "", "Password for the CloudSQL publisher.")
	sqlInstance    = flag.String("sql_instance", "", "Instance name for the CloudSQL publisher.")
	sqlDatabase    = flag.String("sql_database", "metrics", "Name of the SQL database.")
	sqlLocalPort   = flag.Int("sql_local_port", 3306, "Port of the SQL database when testing locally. Requires the Cloud SQL proxy to be installed and running.")
)

const megaByte = 1024 * 1024

func handleError(metricName string, err error) {
	fmt.Printf("[%s] %v", metricName, err)
}

func main() {
	flag.Parse()

	if strings.TrimSpace(*pipelineString) == "" {
		log.Fatalf("No pipelines were specified.")
	}
	pipelines := strings.Split(*pipelineString, ",")

	bk, err := clients.CreateBuildkiteClient(*bkOrg, *bkApiToken, *bkDebug)
	if err != nil {
		log.Fatalf("Cannot create Buildkite client: %v", err)
	}

	cloudSql, err := publishers.CreateCloudSqlPublisher(*sqlUser, *sqlPassword, *sqlInstance, *sqlDatabase, *sqlLocalPort)
	if err != nil {
		log.Fatalf("Failed to set up Cloud SQL publisher: %v", err)
	}

	pipelinePerformance := metrics.CreatePipelinePerformance(bk, pipelines...)
	workerAvailability := metrics.CreateWorkerAvailability(bk)
	releaseDownloads := metrics.CreateReleaseDownloads(*ghOrg, *ghRepo, *ghApiToken, megaByte)

	srv := service.CreateService(handleError)
	srv.AddMetric(pipelinePerformance, 120, cloudSql)
	srv.AddMetric(workerAvailability, 60, cloudSql)
	srv.AddMetric(releaseDownloads, 3600, cloudSql)

	err = cloudSql.RegisterMetric(releaseDownloads)
	if err != nil {
		log.Fatalf("Could not register metric: %v", err)
	}

	ds, err := releaseDownloads.Collect()
	if err != nil {
		log.Fatalf("Could not collect download statistics: %v", err)
	}
	err = cloudSql.Publish(releaseDownloads.Name(), ds)
	if err != nil {
		log.Fatalf("Failed to publish download statistics: %v", err)
	}

	//srv.Start()
	//time.Sleep(30 * time.Second)
	//srv.Stop()
}
<commit_msg>Test all metrics locally.<commit_after>package main

import (
	"flag"
	"fmt"
	"log"
	"strings"

	"github.com/fweikert/continuous-integration/metrics/metrics"

	"github.com/fweikert/continuous-integration/metrics/clients"
	"github.com/fweikert/continuous-integration/metrics/publishers"
	"github.com/fweikert/continuous-integration/metrics/service"
)

var (
	bkOrg          = flag.String("buildkite_org", "bazel", "Buildkite organization slug")
	bkApiToken     = flag.String("buildkite_token", "", "Buildkite API access token that grants read access. See https://buildkite.com/docs/apis/rest-api#authentication")
	bkDebug        = flag.Bool("debug", false, "Enable debugging")
	pipelineString = flag.String("pipelines", "", "Comma separated list of slugs of pipelines whose performance statistics should be exported.")
	ghOrg          = flag.String("github_org", "bazelbuild", "Name of the GitHub organization.")
	ghRepo         = flag.String("github_repo", "bazelbuild", "Name of the GitHub repository.")
	ghApiToken     = flag.String("github_token", "", "Access token for the GitHub API.")
	sqlUser        = flag.String("sql_user", "", "User name for the CloudSQL publisher.")
	sqlPassword    = flag.String("sql_password", "", "Password for the CloudSQL publisher.")
	sqlInstance    = flag.String("sql_instance", "", "Instance name for the CloudSQL publisher.")
	sqlDatabase    = flag.String("sql_database", "metrics", "Name of the SQL database.")
	sqlLocalPort   = flag.Int("sql_local_port", 3306, "Port of the SQL database when testing locally. Requires the Cloud SQL proxy to be installed and running.")
)

const megaByte = 1024 * 1024

func handleError(metricName string, err error) {
	fmt.Printf("[%s] %v", metricName, err)
}

func main() {
	flag.Parse()

	if strings.TrimSpace(*pipelineString) == "" {
		log.Fatalf("No pipelines were specified.")
	}
	pipelines := strings.Split(*pipelineString, ",")

	bk, err := clients.CreateBuildkiteClient(*bkOrg, *bkApiToken, *bkDebug)
	if err != nil {
		log.Fatalf("Cannot create Buildkite client: %v", err)
	}

	cloudSql, err := publishers.CreateCloudSqlPublisher(*sqlUser, *sqlPassword, *sqlInstance, *sqlDatabase, *sqlLocalPort)
	if err != nil {
		log.Fatalf("Failed to set up Cloud SQL publisher: %v", err)
	}

	pipelinePerformance := metrics.CreatePipelinePerformance(bk, pipelines...)
	workerAvailability := metrics.CreateWorkerAvailability(bk)
	releaseDownloads := metrics.CreateReleaseDownloads(*ghOrg, *ghRepo, *ghApiToken, megaByte)

	srv := service.CreateService(handleError)
	srv.AddMetric(pipelinePerformance, 120, cloudSql)
	srv.AddMetric(workerAvailability, 60, cloudSql)
	srv.AddMetric(releaseDownloads, 3600, cloudSql)

	// Test code - please delete later.
	for _, m := range []metrics.Metric{pipelinePerformance, workerAvailability, releaseDownloads} {
		name := m.Name()
		err = cloudSql.RegisterMetric(m)
		if err != nil {
			log.Fatalf("Could not register metric %s: %v", name, err)
		}

		ds, err := m.Collect()
		if err != nil {
			log.Fatalf("Could not collect data for metric %s: %v", name, err)
		}
		err = cloudSql.Publish(name, ds)
		if err != nil {
			log.Fatalf("Failed to publish metric %s: %v", name, err)
		}
	}

	//srv.Start()
	//time.Sleep(30 * time.Second)
	//srv.Stop()
}
<|endoftext|>


<commit_before>// Scanner for generating workable ip:port pairs for IP cameras from DNS
// name patterns
// (c) 2013 Victor Agababov (vagababov@gmail.com), Sergey Shekyan (shekyan@gmail.com).
package main

import (
  "flag"
  "fmt"
  "io/ioutil"
  "net"
  "strings"
  "time"
)

const (
  server = "Netwave IP Camera"
)

var (
  localhost = net.IPv4(127, 0, 0, 1)

  ports       = flag.String("ports", "80", "comma separated ports to scan")
  parsedPorts = strings.Split(*ports, ",")

  suffix      = flag.String("domain", ".localhost", "domain to scan, e.g. .example.com")
  connectionTimeout = flag.Duration("conn_timeout", 5*time.Second, "timeout duration")
  processTimeout    = flag.Duration("process_timeout", 5*time.Second, "read/write analyze duration")

  parallelizm = flag.Int("parallelizm", 25, "number of concurrent resolutions")
  limitChan   = make(chan int, *parallelizm)
)

func process(response string) bool {
  t := strings.Split(response, "\r\n")
  for _, i := range t {
    if strings.HasPrefix(i, "Server:") && strings.Contains(i, "Camera") {
      return true
    }
  }
  return false
}

func resolve(prefix string) {
  full := prefix + *suffix
  if addr, err := net.ResolveIPAddr("ip", full); err == nil {
    if addr.IP.Equal(localhost) {
      fmt.Printf("%v %v localhost\n", full, addr)
    } else {
      for _, p := range parsedPorts {
        conn, err := net.DialTimeout("tcp", addr.String()+":"+p, *connectionTimeout)
        if err == nil {
                                        // Make sure the connection will be closed.
          conn.SetDeadline(time.Now().Add(*processTimeout))
          _, err = conn.Write([]byte("GET / HTTP/1.0\r\n\r\n"))
          if err == nil {
            result, err := ioutil.ReadAll(conn)
            if err == nil {
              response := string(result)
              if process(response) {
                fmt.Printf("%v %v success %s\n", full, addr, p)
              } else {
                fmt.Printf("%v %v other %s\n", full, addr, p)
              }
            } else {
              fmt.Printf("%v %v ping %s\n", full, addr, p)
            }
          } else {
                      fmt.Printf("%v %v ping %s\n", full, addr, p)
          }
                                        conn.Close()
        } else {
          fmt.Printf("%v %v failure %s\n", full, addr, p)
        }
      }
    }
  }
  <-limitChan
}

func main() {
  flag.Parse()
  parsedPorts = strings.Split(*ports, ",")
  s1, s2 := 'a', 'a'
  for {
    for {
      for i := 0; i < 10000; i++ {
        t := fmt.Sprintf("%c%c%04d", s1, s2, i)
        limitChan <- i
        go resolve(t)
      }
      s2++
      if s2 == 'z' {
        s2 = 'a'
        break
      }
    }
    s1++
    if s1 > 'z' {
      break
    }
  }
}

<commit_msg>added camera version logging<commit_after>// Scanner for generating workable ip:port pairs for IP cameras from DNS
// name patterns
// (c) 2013 Victor Agababov (vagababov@gmail.com), Sergey Shekyan (shekyan@gmail.com).
package main

import (
  "flag"
  "fmt"
  "io/ioutil"
  "net"
  "strings"
  "time"
)

const (
  server = "Netwave IP Camera"
)

var (
  localhost = net.IPv4(127, 0, 0, 1)

  ports       = flag.String("ports", "80", "comma separated ports to scan")
  parsedPorts = strings.Split(*ports, ",")

  suffix      = flag.String("domain", ".localhost", "domain to scan, e.g. .example.com")
  connectionTimeout = flag.Duration("conn_timeout", 5*time.Second, "timeout duration")
  processTimeout    = flag.Duration("process_timeout", 5*time.Second, "read/write analyze duration")

  parallelizm = flag.Int("parallelizm", 25, "number of concurrent resolutions")
  limitChan   = make(chan int, *parallelizm)
)

func process(response string) bool {
  t := strings.Split(response, "\r\n")
  for _, i := range t {
    if strings.HasPrefix(i, "Server:") && strings.Contains(i, "Camera") {
      return true
    }
  }
  return false
}

func get_version(response string) string {

  body := strings.Split(response, "\r\n\r\n")
  t := strings.Split(body[1], "\n")
  for _, i := range t {
    if strings.HasPrefix(i, "var sys_ver=") {
      return strings.Split(i, "=")[1]
    }
  }
  return "undefined"
}

func resolve(prefix string) {
  full := prefix + *suffix
  if addr, err := net.ResolveIPAddr("ip", full); err == nil {
    if addr.IP.Equal(localhost) {
      fmt.Printf("%v %v localhost\n", full, addr)
    } else {
      for _, p := range parsedPorts {
        conn, err := net.DialTimeout("tcp", addr.String()+":"+p, *connectionTimeout)
        if err == nil {
                                        // Make sure the connection will be closed.
          conn.SetDeadline(time.Now().Add(*processTimeout))
          _, err = conn.Write([]byte("GET /get_status.cgi HTTP/1.0\r\n\r\n"))
          if err == nil {
            result, err := ioutil.ReadAll(conn)
            if err == nil {
              response := string(result)
              if process(response) {
                fmt.Printf("%v %v success %s %s\n", full, addr, p, get_version(response))
              } else {
                fmt.Printf("%v %v other %s\n", full, addr, p)
              }
            } else {
              fmt.Printf("%v %v ping %s\n", full, addr, p)
            }
          } else {
                      fmt.Printf("%v %v ping %s\n", full, addr, p)
          }
                                        conn.Close()
        } else {
          fmt.Printf("%v %v failure %s\n", full, addr, p)
        }
      }
    }
  }
  <-limitChan
}

func main() {
  flag.Parse()
  parsedPorts = strings.Split(*ports, ",")
  s1, s2 := 'a', 'a'
  for {
    for {
      for i := 0; i < 10000; i++ {
        t := fmt.Sprintf("%c%c%04d", s1, s2, i)
        limitChan <- i
        go resolve(t)
      }
      s2++
      if s2 == 'z' {
        s2 = 'a'
        break
      }
    }
    s1++
    if s1 > 'z' {
      break
    }
  }
}

<|endoftext|>


<commit_before>//go:generate struct-markdown
//go:generate mapstructure-to-hcl2 -type IAPConfig

package googlecompute

import (
	"bufio"
	"bytes"
	"context"
	"fmt"
	"io"
	"log"
	"os"
	"os/exec"
	"runtime"
	"strconv"
	"strings"
	"text/template"
	"time"

	"github.com/hashicorp/packer/common/net"
	"github.com/hashicorp/packer/common/retry"
	"github.com/hashicorp/packer/helper/communicator"
	"github.com/hashicorp/packer/helper/multistep"
	"github.com/hashicorp/packer/packer"
	"github.com/hashicorp/packer/packer/tmp"
)

// StepStartTunnel represents a Packer build step that launches an IAP tunnel
type IAPConfig struct {
	// Whether to use an IAP proxy.
	// Prerequisites and limitations for using IAP:
	// - You must manually enable the IAP API in the Google Cloud console.
	// - You must have the gcloud sdk installed on the computer running Packer.
	// - You must be using a Service Account with a credentials file (using the
	//	 account_file option in the Packer template)
	// - This is currently only implemented for the SSH communicator, not the
	//   WinRM Communicator.
	// - You must add the given service account to project level IAP permissions
	//   in https://console.cloud.google.com/security/iap. To do so, click
	//   "project" > "SSH and TCP resoures" > "All Tunnel Resources" >
	//   "Add Member". Then add your service account and choose the role
	//   "IAP-secured Tunnel User" and add any conditions you may care about.
	IAP bool `mapstructure:"use_iap" required:"false"`
	// Which port to connect the local end of the IAM localhost proxy to. If
	// left blank, Packer will choose a port for you from available ports.
	IAPLocalhostPort int `mapstructure:"iap_localhost_port"`
	// What "hashbang" to use to invoke script that sets up gcloud.
	// Default: "/bin/sh"
	IAPHashBang string `mapstructure:"iap_hashbang" required:"false"`
	// What file extension to use for script that sets up gcloud.
	// Default: ".sh"
	IAPExt string `mapstructure:"iap_ext" required:"false"`
}

type TunnelDriver interface {
	StartTunnel(context.Context, string) error
	StopTunnel()
}

func RunTunnelCommand(cmd *exec.Cmd) error {
	// set stdout and stderr so we can read what's going on.
	var stdout, stderr bytes.Buffer
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr
	err := cmd.Start()
	if err != nil {
		err := fmt.Errorf("Error calling gcloud sdk to launch IAP tunnel: %s",
			err)
		return err
	}

	// Give tunnel 30 seconds to either launch, or return an error.
	// Unfortunately, the SDK doesn't provide any official acknowledgment that
	// the tunnel is launched when it's not being run through a TTY so we
	// are just trusting here that 30s is enough to know whether the tunnel
	// launch was going to fail. Yep, feels icky to me too. But I spent an
	// afternoon trying to figure out how to get the SDK to actually send
	// the "Listening on port [n]" line I see when I run it manually, and I
	// can't justify spending more time than that on aesthetics.
	for i := 0; i < 30; i++ {
		time.Sleep(1 * time.Second)

		lineStderr, err := stderr.ReadString('\n')
		if err != nil && err != io.EOF {
			log.Printf("Err from scanning stderr is %s", err)
			return fmt.Errorf("Error reading stderr from tunnel launch: %s", err)
		}
		if lineStderr != "" {
			log.Printf("stderr: %s", lineStderr)
		}

		lineStdout, err := stdout.ReadString('\n')
		if err != nil && err != io.EOF {
			log.Printf("Err from scanning stdout is %s", err)
			return fmt.Errorf("Error reading stdout from tunnel launch: %s", err)
		}
		if lineStdout != "" {
			log.Printf("stdout: %s", lineStdout)
		}

		if strings.Contains(lineStderr, "ERROR") {
			// 4033: Either you don't have permission to access the instance,
			// the instance doesn't exist, or the instance is stopped.
			// The two sub-errors we may see while the permissions settle are
			// "not authorized" and "failed to connect to backend," but after
			// about a minute of retries this goes away and we're able to
			// connect.
			// 4003: "failed to connect to backend". Network blip.
			if strings.Contains(lineStderr, "4033") || strings.Contains(lineStderr, "4003") {
				return RetryableTunnelError{lineStderr}
			} else {
				log.Printf("NOT RETRYABLE: %s", lineStderr)
				return fmt.Errorf("Non-retryable tunnel error: %s", lineStderr)
			}
		}
	}

	log.Printf("No error detected after tunnel launch; continuing...")
	return nil
}

type RetryableTunnelError struct {
	s string
}

func (e RetryableTunnelError) Error() string {
	return "Tunnel start: " + e.s
}

type StepStartTunnel struct {
	IAPConf     *IAPConfig
	CommConf    *communicator.Config
	AccountFile string
	ProjectId   string

	tunnelDriver TunnelDriver
}

func (s *StepStartTunnel) ConfigureLocalHostPort(ctx context.Context) error {
	if s.IAPConf.IAPLocalhostPort == 0 {
		log.Printf("Finding an available TCP port for IAP proxy")
		l, err := net.ListenRangeConfig{
			Min:     8000,
			Max:     9000,
			Addr:    "0.0.0.0",
			Network: "tcp",
		}.Listen(ctx)

		if err != nil {
			err := fmt.Errorf("error finding an available port to initiate a session tunnel: %s", err)
			return err
		}

		s.IAPConf.IAPLocalhostPort = l.Port
		l.Close()
		log.Printf("Setting up proxy to listen on localhost at %d",
			s.IAPConf.IAPLocalhostPort)
	}
	return nil
}

func (s *StepStartTunnel) createTempGcloudScript(args []string) (string, error) {
	// Generate temp script that contains both gcloud auth and gcloud compute
	// iap launch call.

	// Create temp file.
	tf, err := tmp.File("gcloud-setup")
	if err != nil {
		return "", fmt.Errorf("Error preparing gcloud setup script: %s", err)
	}
	defer tf.Close()
	// Write our contents to it
	writer := bufio.NewWriter(tf)

	if s.IAPConf.IAPHashBang != "" {
		s.IAPConf.IAPHashBang = fmt.Sprintf("#!%s\n", s.IAPConf.IAPHashBang)
		log.Printf("[INFO] (google): Prepending inline gcloud setup script with %s",
			s.IAPConf.IAPHashBang)
		_, err = writer.WriteString(s.IAPConf.IAPHashBang)
		if err != nil {
			return "", fmt.Errorf("Error preparing inline hashbang: %s", err)
		}

	}

	launchTemplate := `
gcloud auth activate-service-account --key-file='{{.AccountFile}}'
gcloud config set project {{.ProjectID}}
{{.Args}}
`
	if runtime.GOOS == "windows" {
		launchTemplate = `
call gcloud auth activate-service-account --key-file {{.AccountFile}}
call gcloud config set project {{.ProjectID}}
call {{.Args}}
`
	}
	// call command
	args = append([]string{"gcloud"}, args...)
	argString := strings.Join(args, " ")

	var tpl = template.Must(template.New("createTunnel").Parse(launchTemplate))
	var b bytes.Buffer

	opts := map[string]string{
		"AccountFile": s.AccountFile,
		"ProjectID":   s.ProjectId,
		"Args":        argString,
	}

	err = tpl.Execute(&b, opts)
	if err != nil {
		fmt.Println(err)
	}

	if _, err := writer.WriteString(b.String()); err != nil {
		return "", fmt.Errorf("Error preparing gcloud shell script: %s", err)
	}

	if err := writer.Flush(); err != nil {
		return "", fmt.Errorf("Error preparing shell script: %s", err)
	}
	// Have to close temp file before renaming it or Windows will complain.
	tf.Close()
	err = os.Chmod(tf.Name(), 0700)
	if err != nil {
		log.Printf("[ERROR] (google): error modifying permissions of temp script file: %s", err.Error())
	}

	// figure out what extension the file should have, and rename it.
	tempScriptFileName := tf.Name()
	if s.IAPConf.IAPExt != "" {
		err := os.Rename(tempScriptFileName, fmt.Sprintf("%s%s", tempScriptFileName, s.IAPConf.IAPExt))
		if err != nil {
			return "", fmt.Errorf("Error setting the correct temp file extension: %s", err)
		}
		tempScriptFileName = fmt.Sprintf("%s%s", tempScriptFileName, s.IAPConf.IAPExt)
	}

	return tempScriptFileName, nil
}

// Run executes the Packer build step that creates an IAP tunnel.
func (s *StepStartTunnel) Run(ctx context.Context, state multistep.StateBag) multistep.StepAction {
	if !s.IAPConf.IAP {
		log.Printf("Skipping step launch IAP tunnel; \"iap\" is false.")
		return multistep.ActionContinue
	}

	// shell out to create the tunnel.
	ui := state.Get("ui").(packer.Ui)
	instanceName := state.Get("instance_name").(string)
	c := state.Get("config").(*Config)

	ui.Say("Step Launch IAP Tunnel...")

	err := s.ConfigureLocalHostPort(ctx)
	if err != nil {
		state.Put("error", err)
		ui.Error(err.Error())
		return multistep.ActionHalt
	}

	// Generate list of args to use to call gcloud cli.
	args := []string{"compute", "start-iap-tunnel", instanceName,
		strconv.Itoa(s.CommConf.Port()),
		fmt.Sprintf("--local-host-port=localhost:%d", s.IAPConf.IAPLocalhostPort),
		"--zone", c.Zone,
	}

	// This is the port the IAP tunnel listens on, on localhost.
	// TODO make setting LocalHostPort optional
	s.CommConf.SSHPort = s.IAPConf.IAPLocalhostPort

	log.Printf("Creating tunnel launch script with args %#v", args)
	// Create temp file that contains both gcloud authentication, and gcloud
	// proxy setup call.
	tempScriptFileName, err := s.createTempGcloudScript(args)
	if err != nil {
		state.Put("error", err)
		ui.Error(err.Error())
		return multistep.ActionHalt
	}
	defer os.Remove(tempScriptFileName)

	s.tunnelDriver = NewTunnelDriver()

	err = retry.Config{
		Tries: 11,
		ShouldRetry: func(err error) bool {
			switch err.(type) {
			case RetryableTunnelError:
				return true
			default:
				return false
			}
		},
		RetryDelay: (&retry.Backoff{InitialBackoff: 200 * time.Millisecond, MaxBackoff: 30 * time.Second, Multiplier: 2}).Linear,
	}.Run(ctx, func(ctx context.Context) error {
		// tunnel launcher/destroyer has to be different on windows vs. unix.
		err := s.tunnelDriver.StartTunnel(ctx, tempScriptFileName)
		return err
	})
	if err != nil {
		state.Put("error", err)
		ui.Error(err.Error())
		return multistep.ActionHalt
	}

	return multistep.ActionContinue
}

// Cleanup stops the IAP tunnel and cleans up processes.
func (s *StepStartTunnel) Cleanup(state multistep.StateBag) {
	if !s.IAPConf.IAP {
		log.Printf("Skipping cleanup of IAP tunnel; \"iap\" is false.")
		return
	}
	if s.tunnelDriver != nil {
		s.tunnelDriver.StopTunnel()
	}
}
<commit_msg>make windows work with account files containing spaces<commit_after>//go:generate struct-markdown
//go:generate mapstructure-to-hcl2 -type IAPConfig

package googlecompute

import (
	"bufio"
	"bytes"
	"context"
	"fmt"
	"io"
	"log"
	"os"
	"os/exec"
	"runtime"
	"strconv"
	"strings"
	"text/template"
	"time"

	"github.com/hashicorp/packer/common/net"
	"github.com/hashicorp/packer/common/retry"
	"github.com/hashicorp/packer/helper/communicator"
	"github.com/hashicorp/packer/helper/multistep"
	"github.com/hashicorp/packer/packer"
	"github.com/hashicorp/packer/packer/tmp"
)

// StepStartTunnel represents a Packer build step that launches an IAP tunnel
type IAPConfig struct {
	// Whether to use an IAP proxy.
	// Prerequisites and limitations for using IAP:
	// - You must manually enable the IAP API in the Google Cloud console.
	// - You must have the gcloud sdk installed on the computer running Packer.
	// - You must be using a Service Account with a credentials file (using the
	//	 account_file option in the Packer template)
	// - This is currently only implemented for the SSH communicator, not the
	//   WinRM Communicator.
	// - You must add the given service account to project level IAP permissions
	//   in https://console.cloud.google.com/security/iap. To do so, click
	//   "project" > "SSH and TCP resoures" > "All Tunnel Resources" >
	//   "Add Member". Then add your service account and choose the role
	//   "IAP-secured Tunnel User" and add any conditions you may care about.
	IAP bool `mapstructure:"use_iap" required:"false"`
	// Which port to connect the local end of the IAM localhost proxy to. If
	// left blank, Packer will choose a port for you from available ports.
	IAPLocalhostPort int `mapstructure:"iap_localhost_port"`
	// What "hashbang" to use to invoke script that sets up gcloud.
	// Default: "/bin/sh"
	IAPHashBang string `mapstructure:"iap_hashbang" required:"false"`
	// What file extension to use for script that sets up gcloud.
	// Default: ".sh"
	IAPExt string `mapstructure:"iap_ext" required:"false"`
}

type TunnelDriver interface {
	StartTunnel(context.Context, string) error
	StopTunnel()
}

func RunTunnelCommand(cmd *exec.Cmd) error {
	// set stdout and stderr so we can read what's going on.
	var stdout, stderr bytes.Buffer
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr
	err := cmd.Start()
	if err != nil {
		err := fmt.Errorf("Error calling gcloud sdk to launch IAP tunnel: %s",
			err)
		return err
	}

	// Give tunnel 30 seconds to either launch, or return an error.
	// Unfortunately, the SDK doesn't provide any official acknowledgment that
	// the tunnel is launched when it's not being run through a TTY so we
	// are just trusting here that 30s is enough to know whether the tunnel
	// launch was going to fail. Yep, feels icky to me too. But I spent an
	// afternoon trying to figure out how to get the SDK to actually send
	// the "Listening on port [n]" line I see when I run it manually, and I
	// can't justify spending more time than that on aesthetics.
	for i := 0; i < 30; i++ {
		time.Sleep(1 * time.Second)

		lineStderr, err := stderr.ReadString('\n')
		if err != nil && err != io.EOF {
			log.Printf("Err from scanning stderr is %s", err)
			return fmt.Errorf("Error reading stderr from tunnel launch: %s", err)
		}
		if lineStderr != "" {
			log.Printf("stderr: %s", lineStderr)
		}

		lineStdout, err := stdout.ReadString('\n')
		if err != nil && err != io.EOF {
			log.Printf("Err from scanning stdout is %s", err)
			return fmt.Errorf("Error reading stdout from tunnel launch: %s", err)
		}
		if lineStdout != "" {
			log.Printf("stdout: %s", lineStdout)
		}

		if strings.Contains(lineStderr, "ERROR") {
			// 4033: Either you don't have permission to access the instance,
			// the instance doesn't exist, or the instance is stopped.
			// The two sub-errors we may see while the permissions settle are
			// "not authorized" and "failed to connect to backend," but after
			// about a minute of retries this goes away and we're able to
			// connect.
			// 4003: "failed to connect to backend". Network blip.
			if strings.Contains(lineStderr, "4033") || strings.Contains(lineStderr, "4003") {
				return RetryableTunnelError{lineStderr}
			} else {
				log.Printf("NOT RETRYABLE: %s", lineStderr)
				return fmt.Errorf("Non-retryable tunnel error: %s", lineStderr)
			}
		}
	}

	log.Printf("No error detected after tunnel launch; continuing...")
	return nil
}

type RetryableTunnelError struct {
	s string
}

func (e RetryableTunnelError) Error() string {
	return "Tunnel start: " + e.s
}

type StepStartTunnel struct {
	IAPConf     *IAPConfig
	CommConf    *communicator.Config
	AccountFile string
	ProjectId   string

	tunnelDriver TunnelDriver
}

func (s *StepStartTunnel) ConfigureLocalHostPort(ctx context.Context) error {
	if s.IAPConf.IAPLocalhostPort == 0 {
		log.Printf("Finding an available TCP port for IAP proxy")
		l, err := net.ListenRangeConfig{
			Min:     8000,
			Max:     9000,
			Addr:    "0.0.0.0",
			Network: "tcp",
		}.Listen(ctx)

		if err != nil {
			err := fmt.Errorf("error finding an available port to initiate a session tunnel: %s", err)
			return err
		}

		s.IAPConf.IAPLocalhostPort = l.Port
		l.Close()
		log.Printf("Setting up proxy to listen on localhost at %d",
			s.IAPConf.IAPLocalhostPort)
	}
	return nil
}

func (s *StepStartTunnel) createTempGcloudScript(args []string) (string, error) {
	// Generate temp script that contains both gcloud auth and gcloud compute
	// iap launch call.

	// Create temp file.
	tf, err := tmp.File("gcloud-setup")
	if err != nil {
		return "", fmt.Errorf("Error preparing gcloud setup script: %s", err)
	}
	defer tf.Close()
	// Write our contents to it
	writer := bufio.NewWriter(tf)

	if s.IAPConf.IAPHashBang != "" {
		s.IAPConf.IAPHashBang = fmt.Sprintf("#!%s\n", s.IAPConf.IAPHashBang)
		log.Printf("[INFO] (google): Prepending inline gcloud setup script with %s",
			s.IAPConf.IAPHashBang)
		_, err = writer.WriteString(s.IAPConf.IAPHashBang)
		if err != nil {
			return "", fmt.Errorf("Error preparing inline hashbang: %s", err)
		}

	}

	launchTemplate := `
gcloud auth activate-service-account --key-file='{{.AccountFile}}'
gcloud config set project {{.ProjectID}}
{{.Args}}
`
	if runtime.GOOS == "windows" {
		launchTemplate = `
call gcloud auth activate-service-account --key-file "{{.AccountFile}}"
call gcloud config set project {{.ProjectID}}
call {{.Args}}
`
	}
	// call command
	args = append([]string{"gcloud"}, args...)
	argString := strings.Join(args, " ")

	var tpl = template.Must(template.New("createTunnel").Parse(launchTemplate))
	var b bytes.Buffer

	opts := map[string]string{
		"AccountFile": s.AccountFile,
		"ProjectID":   s.ProjectId,
		"Args":        argString,
	}

	err = tpl.Execute(&b, opts)
	if err != nil {
		fmt.Println(err)
	}

	if _, err := writer.WriteString(b.String()); err != nil {
		return "", fmt.Errorf("Error preparing gcloud shell script: %s", err)
	}

	if err := writer.Flush(); err != nil {
		return "", fmt.Errorf("Error preparing shell script: %s", err)
	}
	// Have to close temp file before renaming it or Windows will complain.
	tf.Close()
	err = os.Chmod(tf.Name(), 0700)
	if err != nil {
		log.Printf("[ERROR] (google): error modifying permissions of temp script file: %s", err.Error())
	}

	// figure out what extension the file should have, and rename it.
	tempScriptFileName := tf.Name()
	if s.IAPConf.IAPExt != "" {
		err := os.Rename(tempScriptFileName, fmt.Sprintf("%s%s", tempScriptFileName, s.IAPConf.IAPExt))
		if err != nil {
			return "", fmt.Errorf("Error setting the correct temp file extension: %s", err)
		}
		tempScriptFileName = fmt.Sprintf("%s%s", tempScriptFileName, s.IAPConf.IAPExt)
	}

	return tempScriptFileName, nil
}

// Run executes the Packer build step that creates an IAP tunnel.
func (s *StepStartTunnel) Run(ctx context.Context, state multistep.StateBag) multistep.StepAction {
	if !s.IAPConf.IAP {
		log.Printf("Skipping step launch IAP tunnel; \"iap\" is false.")
		return multistep.ActionContinue
	}

	// shell out to create the tunnel.
	ui := state.Get("ui").(packer.Ui)
	instanceName := state.Get("instance_name").(string)
	c := state.Get("config").(*Config)

	ui.Say("Step Launch IAP Tunnel...")

	err := s.ConfigureLocalHostPort(ctx)
	if err != nil {
		state.Put("error", err)
		ui.Error(err.Error())
		return multistep.ActionHalt
	}

	// Generate list of args to use to call gcloud cli.
	args := []string{"compute", "start-iap-tunnel", instanceName,
		strconv.Itoa(s.CommConf.Port()),
		fmt.Sprintf("--local-host-port=localhost:%d", s.IAPConf.IAPLocalhostPort),
		"--zone", c.Zone,
	}

	// This is the port the IAP tunnel listens on, on localhost.
	// TODO make setting LocalHostPort optional
	s.CommConf.SSHPort = s.IAPConf.IAPLocalhostPort

	log.Printf("Creating tunnel launch script with args %#v", args)
	// Create temp file that contains both gcloud authentication, and gcloud
	// proxy setup call.
	tempScriptFileName, err := s.createTempGcloudScript(args)
	if err != nil {
		state.Put("error", err)
		ui.Error(err.Error())
		return multistep.ActionHalt
	}
	defer os.Remove(tempScriptFileName)

	s.tunnelDriver = NewTunnelDriver()

	err = retry.Config{
		Tries: 11,
		ShouldRetry: func(err error) bool {
			switch err.(type) {
			case RetryableTunnelError:
				return true
			default:
				return false
			}
		},
		RetryDelay: (&retry.Backoff{InitialBackoff: 200 * time.Millisecond, MaxBackoff: 30 * time.Second, Multiplier: 2}).Linear,
	}.Run(ctx, func(ctx context.Context) error {
		// tunnel launcher/destroyer has to be different on windows vs. unix.
		err := s.tunnelDriver.StartTunnel(ctx, tempScriptFileName)
		return err
	})
	if err != nil {
		state.Put("error", err)
		ui.Error(err.Error())
		return multistep.ActionHalt
	}

	return multistep.ActionContinue
}

// Cleanup stops the IAP tunnel and cleans up processes.
func (s *StepStartTunnel) Cleanup(state multistep.StateBag) {
	if !s.IAPConf.IAP {
		log.Printf("Skipping cleanup of IAP tunnel; \"iap\" is false.")
		return
	}
	if s.tunnelDriver != nil {
		s.tunnelDriver.StopTunnel()
	}
}
<|endoftext|>