<commit_before>package eval

// Builtin functions.

import (
	"bufio"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	"reflect"
	"runtime"
	"strconv"
)

var builtinFns []*BuiltinFn

func init() {
	// Needed to work around init loop.
	builtinFns = []*BuiltinFn{
		&BuiltinFn{":", nop},
		&BuiltinFn{"true", nop},

		&BuiltinFn{"print", wrapFn(print)},
		&BuiltinFn{"println", wrapFn(println)},

		&BuiltinFn{"into-lines", wrapFn(intoLines)},
		&BuiltinFn{"from-lines", wrapFn(fromLines)},

		&BuiltinFn{"rat", wrapFn(ratFn)},

		&BuiltinFn{"put", put},
		&BuiltinFn{"put-all", wrapFn(putAll)},
		&BuiltinFn{"unpack", wrapFn(unpack)},

		&BuiltinFn{"from-json", wrapFn(fromJSON)},

		&BuiltinFn{"typeof", typeof},

		&BuiltinFn{"fail", wrapFn(fail)},
		&BuiltinFn{"multi-error", wrapFn(multiErrorFn)},
		&BuiltinFn{"return", wrapFn(returnFn)},
		&BuiltinFn{"break", wrapFn(breakFn)},
		&BuiltinFn{"continue", wrapFn(continueFn)},

		&BuiltinFn{"each", wrapFn(each)},

		&BuiltinFn{"cd", cd},
		&BuiltinFn{"dirs", wrapFn(dirs)},
		&BuiltinFn{"jump", wrapFn(jump)},

		&BuiltinFn{"source", wrapFn(source)},

		&BuiltinFn{"+", wrapFn(plus)},
		&BuiltinFn{"-", wrapFn(minus)},
		&BuiltinFn{"mul", wrapFn(times)},
		&BuiltinFn{"div", wrapFn(divide)},

		&BuiltinFn{"=", eq},
		&BuiltinFn{"deepeq", deepeq},

		&BuiltinFn{"take", wrapFn(take)},
		&BuiltinFn{"drop", wrapFn(drop)},

		&BuiltinFn{"bind", wrapFn(bind)},
		&BuiltinFn{"le", wrapFn(le)},

		&BuiltinFn{"-stack", wrapFn(_stack)},
	}
}

var (
	ErrArgs              = errors.New("args error")
	ErrInput             = errors.New("input error")
	ErrStoreNotConnected = errors.New("store not connected")
	ErrNoMatchingDir     = errors.New("no matching directory")
	ErrNoEditor          = errors.New("no line editor")
)

var (
	evalCtxType = reflect.TypeOf((*evalCtx)(nil))
	valueType   = reflect.TypeOf((*Value)(nil)).Elem()
)

// wrapFn wraps an inner function into one suitable as a builtin function. It
// generates argument checking and conversion code according to the signature
// of the inner function. The inner function must accept evalCtx* as the first
// argument and return an exitus.
func wrapFn(inner interface{}) func(*evalCtx, []Value) {
	type_ := reflect.TypeOf(inner)
	if type_.In(0) != evalCtxType {
		panic("bad func")
	}

	requiredArgs := type_.NumIn() - 1
	isVariadic := type_.IsVariadic()
	var variadicType reflect.Type
	if isVariadic {
		requiredArgs--
		variadicType = type_.In(type_.NumIn() - 1).Elem()
		if !supportedIn(variadicType) {
			panic("bad func argument")
		}
	}

	for i := 0; i < requiredArgs; i++ {
		if !supportedIn(type_.In(i + 1)) {
			panic("bad func argument")
		}
	}

	return func(ec *evalCtx, args []Value) {
		if len(args) < requiredArgs || (!isVariadic && len(args) > requiredArgs) {
			throw(ErrArgs)
		}
		callArgs := make([]reflect.Value, len(args)+1)
		callArgs[0] = reflect.ValueOf(ec)

		ok := convertArgs(args[:requiredArgs], callArgs[1:],
			func(i int) reflect.Type { return type_.In(i + 1) })
		if !ok {
			throw(ErrArgs)
		}
		if isVariadic {
			ok := convertArgs(args[requiredArgs:], callArgs[1+requiredArgs:],
				func(i int) reflect.Type { return variadicType })
			if !ok {
				throw(ErrArgs)
			}
		}
		reflect.ValueOf(inner).Call(callArgs)
	}
}

func supportedIn(t reflect.Type) bool {
	return t.Kind() == reflect.String ||
		t.Kind() == reflect.Int || t.Kind() == reflect.Float64 ||
		t.Implements(valueType)
}

func convertArgs(args []Value, callArgs []reflect.Value, callType func(int) reflect.Type) bool {
	for i, arg := range args {
		var callArg interface{}
		switch callType(i).Kind() {
		case reflect.String:
			callArg = ToString(arg)
		case reflect.Int:
			var err error
			callArg, err = toInt(arg)
			if err != nil {
				return false
			}
		case reflect.Float64:
			var err error
			callArg, err = toFloat(arg)
			if err != nil {
				return false
				// return err
			}
		default:
			if reflect.TypeOf(arg).ConvertibleTo(callType(i)) {
				callArg = arg
			} else {
				return false
				// return argsError
			}
		}
		callArgs[i] = reflect.ValueOf(callArg)
	}
	return true
}

func nop(ec *evalCtx, args []Value) {
}

func put(ec *evalCtx, args []Value) {
	out := ec.ports[1].Chan
	for _, a := range args {
		out <- a
	}
}

func putAll(ec *evalCtx, lists ...List) {
	out := ec.ports[1].Chan
	for _, list := range lists {
		for _, x := range *list.inner {
			out <- x
		}
	}
}

func typeof(ec *evalCtx, args []Value) {
	out := ec.ports[1].Chan
	for _, a := range args {
		out <- String(a.Type().String())
	}
}

func fail(ec *evalCtx, arg Value) {
	throw(errors.New(ToString(arg)))
}

func multiErrorFn(ec *evalCtx, args ...Error) {
	throw(multiError{args})
}

func returnFn(ec *evalCtx) {
	throw(Return)
}

func breakFn(ec *evalCtx) {
	throw(Break)
}

func continueFn(ec *evalCtx) {
	throw(Continue)
}

func print(ec *evalCtx, args ...string) {
	out := ec.ports[1].File
	for i, arg := range args {
		if i > 0 {
			out.WriteString(" ")
		}
		out.WriteString(arg)
	}
}

func println(ec *evalCtx, args ...string) {
	print(ec, args...)
	ec.ports[1].File.WriteString("\n")
}

func intoLines(ec *evalCtx) {
	in := ec.ports[0].Chan
	out := ec.ports[1].File

	for v := range in {
		fmt.Fprintln(out, ToString(v))
	}
}

func fromLines(ec *evalCtx) {
	in := ec.ports[0].File
	out := ec.ports[1].Chan

	bufferedIn := bufio.NewReader(in)
	for {
		line, err := bufferedIn.ReadString('\n')
		if err == io.EOF {
			return
		} else if err != nil {
			throw(err)
		}
		out <- String(line[:len(line)-1])
	}
}

func ratFn(ec *evalCtx, arg Value) {
	out := ec.ports[1].Chan
	r, err := ToRat(arg)
	if err != nil {
		throw(err)
	}
	out <- r
}

// unpack takes any number of tables and output their list elements.
func unpack(ec *evalCtx) {
	in := ec.ports[0].Chan
	out := ec.ports[1].Chan

	for v := range in {
		if list, ok := v.(List); !ok {
			throw(ErrInput)
		} else {
			for _, e := range *list.inner {
				out <- e
			}
		}
	}
}

// fromJSON parses a stream of JSON data into Value's.
func fromJSON(ec *evalCtx) {
	in := ec.ports[0].File
	out := ec.ports[1].Chan

	dec := json.NewDecoder(in)
	var v interface{}
	for {
		err := dec.Decode(&v)
		if err != nil {
			if err == io.EOF {
				return
			}
			throw(err)
		}
		out <- FromJSONInterface(v)
	}
}

// each takes a single closure and applies it to all input values.
func each(ec *evalCtx, f *Closure) {
	in := ec.ports[0].Chan
in:
	for v := range in {
		newec := ec.fork("closure of each")
		ex := newec.PCall(f, []Value{v})
		ClosePorts(newec.ports)

		switch ex {
		case nil, Continue:
			// nop
		case Break:
			break in
		default:
			throw(ex)
		}
	}
}

func cd(ec *evalCtx, args []Value) {
	var dir string
	if len(args) == 0 {
		dir = mustGetHome("")
	} else if len(args) == 1 {
		dir = ToString(args[0])
	} else {
		throw(ErrArgs)
	}

	cdInner(dir, ec)
}

func cdInner(dir string, ec *evalCtx) {
	err := os.Chdir(dir)
	if err != nil {
		throw(err)
	}
	if ec.store != nil {
		pwd, err := os.Getwd()
		// BUG(xiaq): Possible error of os.Getwd after cd-ing is ignored.
		if err == nil {
			ec.store.AddDir(pwd)
		}
	}
}

func dirs(ec *evalCtx) {
	if ec.store == nil {
		throw(ErrStoreNotConnected)
	}
	dirs, err := ec.store.ListDirs()
	if err != nil {
		throw(errors.New("store error: " + err.Error()))
	}
	out := ec.ports[1].Chan
	for _, dir := range dirs {
		out <- Map{&map[Value]Value{
			String("path"):  String(dir.Path),
			String("score"): String(fmt.Sprint(dir.Score)),
		}}
	}
}

func jump(ec *evalCtx, arg string) {
	if ec.store == nil {
		throw(ErrStoreNotConnected)
	}
	dirs, err := ec.store.FindDirs(arg)
	if err != nil {
		throw(errors.New("store error: " + err.Error()))
	}
	if len(dirs) == 0 {
		throw(ErrNoMatchingDir)
	}
	dir := dirs[0].Path
	err = os.Chdir(dir)
	// TODO(xiaq): Remove directories that no longer exist
	if err != nil {
		throw(err)
	}
	ec.store.AddDir(dir)
}

func source(ec *evalCtx, fname string) {
	ec.Source(fname)
}

func toFloat(arg Value) (float64, error) {
	arg, ok := arg.(String)
	if !ok {
		return 0, fmt.Errorf("must be string")
	}
	num, err := strconv.ParseFloat(string(arg.(String)), 64)
	if err != nil {
		return 0, err
	}
	return num, nil
}

func toInt(arg Value) (int, error) {
	arg, ok := arg.(String)
	if !ok {
		return 0, fmt.Errorf("must be string")
	}
	num, err := strconv.Atoi(string(arg.(String)))
	if err != nil {
		return 0, err
	}
	return num, nil
}

func plus(ec *evalCtx, nums ...float64) {
	out := ec.ports[1].Chan
	sum := 0.0
	for _, f := range nums {
		sum += f
	}
	out <- String(fmt.Sprintf("%g", sum))
}

func minus(ec *evalCtx, sum float64, nums ...float64) {
	out := ec.ports[1].Chan
	for _, f := range nums {
		sum -= f
	}
	out <- String(fmt.Sprintf("%g", sum))
}

func times(ec *evalCtx, nums ...float64) {
	out := ec.ports[1].Chan
	prod := 1.0
	for _, f := range nums {
		prod *= f
	}
	out <- String(fmt.Sprintf("%g", prod))
}

func divide(ec *evalCtx, prod float64, nums ...float64) {
	out := ec.ports[1].Chan
	for _, f := range nums {
		prod /= f
	}
	out <- String(fmt.Sprintf("%g", prod))
}

func eq(ec *evalCtx, args []Value) {
	out := ec.ports[1].Chan
	if len(args) == 0 {
		throw(ErrArgs)
	}
	for i := 0; i+1 < len(args); i++ {
		if args[i] != args[i+1] {
			out <- Bool(false)
			return
		}
	}
	out <- Bool(true)
}

func deepeq(ec *evalCtx, args []Value) {
	out := ec.ports[1].Chan
	if len(args) == 0 {
		throw(ErrArgs)
	}
	for i := 0; i+1 < len(args); i++ {
		if !DeepEq(args[i], args[i+1]) {
			out <- Bool(false)
			return
		}
	}
	out <- Bool(true)
}

func take(ec *evalCtx, n int) {
	in := ec.ports[0].Chan
	out := ec.ports[1].Chan

	i := 0
	for v := range in {
		if i >= n {
			break
		}
		i++
		out <- v
	}
}

func drop(ec *evalCtx, n int) {
	in := ec.ports[0].Chan
	out := ec.ports[1].Chan

	for i := 0; i < n; i++ {
		<-in
	}
	for v := range in {
		out <- v
	}
}

func bind(ec *evalCtx, key string, function Value) {
	if ec.Editor == nil {
		throw(ErrNoEditor)
	}
	maybeThrow(ec.Editor.Bind(key, function))
}

func le(ec *evalCtx, name string, args ...Value) {
	if ec.Editor == nil {
		throw(ErrNoEditor)
	}
	maybeThrow(ec.Editor.Call(name, args))
}

func _stack(ec *evalCtx) {
	out := ec.ports[1].File

	// XXX dup with main.go
	buf := make([]byte, 1024)
	for runtime.Stack(buf, true) == cap(buf) {
		buf = make([]byte, cap(buf)*2)
	}
	out.Write(buf)
}
<commit_msg>Add a builtin -sleep.<commit_after>package eval

// Builtin functions.

import (
	"bufio"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	"reflect"
	"runtime"
	"strconv"
	"time"
)

var builtinFns []*BuiltinFn

func init() {
	// Needed to work around init loop.
	builtinFns = []*BuiltinFn{
		&BuiltinFn{":", nop},
		&BuiltinFn{"true", nop},

		&BuiltinFn{"print", wrapFn(print)},
		&BuiltinFn{"println", wrapFn(println)},

		&BuiltinFn{"into-lines", wrapFn(intoLines)},
		&BuiltinFn{"from-lines", wrapFn(fromLines)},

		&BuiltinFn{"rat", wrapFn(ratFn)},

		&BuiltinFn{"put", put},
		&BuiltinFn{"put-all", wrapFn(putAll)},
		&BuiltinFn{"unpack", wrapFn(unpack)},

		&BuiltinFn{"from-json", wrapFn(fromJSON)},

		&BuiltinFn{"typeof", typeof},

		&BuiltinFn{"fail", wrapFn(fail)},
		&BuiltinFn{"multi-error", wrapFn(multiErrorFn)},
		&BuiltinFn{"return", wrapFn(returnFn)},
		&BuiltinFn{"break", wrapFn(breakFn)},
		&BuiltinFn{"continue", wrapFn(continueFn)},

		&BuiltinFn{"each", wrapFn(each)},

		&BuiltinFn{"cd", cd},
		&BuiltinFn{"dirs", wrapFn(dirs)},
		&BuiltinFn{"jump", wrapFn(jump)},

		&BuiltinFn{"source", wrapFn(source)},

		&BuiltinFn{"+", wrapFn(plus)},
		&BuiltinFn{"-", wrapFn(minus)},
		&BuiltinFn{"mul", wrapFn(times)},
		&BuiltinFn{"div", wrapFn(divide)},

		&BuiltinFn{"=", eq},
		&BuiltinFn{"deepeq", deepeq},

		&BuiltinFn{"take", wrapFn(take)},
		&BuiltinFn{"drop", wrapFn(drop)},

		&BuiltinFn{"bind", wrapFn(bind)},
		&BuiltinFn{"le", wrapFn(le)},

		&BuiltinFn{"-sleep", wrapFn(_sleep)},
		&BuiltinFn{"-stack", wrapFn(_stack)},
	}
}

var (
	ErrArgs              = errors.New("args error")
	ErrInput             = errors.New("input error")
	ErrStoreNotConnected = errors.New("store not connected")
	ErrNoMatchingDir     = errors.New("no matching directory")
	ErrNoEditor          = errors.New("no line editor")
)

var (
	evalCtxType = reflect.TypeOf((*evalCtx)(nil))
	valueType   = reflect.TypeOf((*Value)(nil)).Elem()
)

// wrapFn wraps an inner function into one suitable as a builtin function. It
// generates argument checking and conversion code according to the signature
// of the inner function. The inner function must accept evalCtx* as the first
// argument and return an exitus.
func wrapFn(inner interface{}) func(*evalCtx, []Value) {
	type_ := reflect.TypeOf(inner)
	if type_.In(0) != evalCtxType {
		panic("bad func")
	}

	requiredArgs := type_.NumIn() - 1
	isVariadic := type_.IsVariadic()
	var variadicType reflect.Type
	if isVariadic {
		requiredArgs--
		variadicType = type_.In(type_.NumIn() - 1).Elem()
		if !supportedIn(variadicType) {
			panic("bad func argument")
		}
	}

	for i := 0; i < requiredArgs; i++ {
		if !supportedIn(type_.In(i + 1)) {
			panic("bad func argument")
		}
	}

	return func(ec *evalCtx, args []Value) {
		if len(args) < requiredArgs || (!isVariadic && len(args) > requiredArgs) {
			throw(ErrArgs)
		}
		callArgs := make([]reflect.Value, len(args)+1)
		callArgs[0] = reflect.ValueOf(ec)

		ok := convertArgs(args[:requiredArgs], callArgs[1:],
			func(i int) reflect.Type { return type_.In(i + 1) })
		if !ok {
			throw(ErrArgs)
		}
		if isVariadic {
			ok := convertArgs(args[requiredArgs:], callArgs[1+requiredArgs:],
				func(i int) reflect.Type { return variadicType })
			if !ok {
				throw(ErrArgs)
			}
		}
		reflect.ValueOf(inner).Call(callArgs)
	}
}

func supportedIn(t reflect.Type) bool {
	return t.Kind() == reflect.String ||
		t.Kind() == reflect.Int || t.Kind() == reflect.Float64 ||
		t.Implements(valueType)
}

func convertArgs(args []Value, callArgs []reflect.Value, callType func(int) reflect.Type) bool {
	for i, arg := range args {
		var callArg interface{}
		switch callType(i).Kind() {
		case reflect.String:
			callArg = ToString(arg)
		case reflect.Int:
			var err error
			callArg, err = toInt(arg)
			if err != nil {
				return false
			}
		case reflect.Float64:
			var err error
			callArg, err = toFloat(arg)
			if err != nil {
				return false
				// return err
			}
		default:
			if reflect.TypeOf(arg).ConvertibleTo(callType(i)) {
				callArg = arg
			} else {
				return false
				// return argsError
			}
		}
		callArgs[i] = reflect.ValueOf(callArg)
	}
	return true
}

func nop(ec *evalCtx, args []Value) {
}

func put(ec *evalCtx, args []Value) {
	out := ec.ports[1].Chan
	for _, a := range args {
		out <- a
	}
}

func putAll(ec *evalCtx, lists ...List) {
	out := ec.ports[1].Chan
	for _, list := range lists {
		for _, x := range *list.inner {
			out <- x
		}
	}
}

func typeof(ec *evalCtx, args []Value) {
	out := ec.ports[1].Chan
	for _, a := range args {
		out <- String(a.Type().String())
	}
}

func fail(ec *evalCtx, arg Value) {
	throw(errors.New(ToString(arg)))
}

func multiErrorFn(ec *evalCtx, args ...Error) {
	throw(multiError{args})
}

func returnFn(ec *evalCtx) {
	throw(Return)
}

func breakFn(ec *evalCtx) {
	throw(Break)
}

func continueFn(ec *evalCtx) {
	throw(Continue)
}

func print(ec *evalCtx, args ...string) {
	out := ec.ports[1].File
	for i, arg := range args {
		if i > 0 {
			out.WriteString(" ")
		}
		out.WriteString(arg)
	}
}

func println(ec *evalCtx, args ...string) {
	print(ec, args...)
	ec.ports[1].File.WriteString("\n")
}

func intoLines(ec *evalCtx) {
	in := ec.ports[0].Chan
	out := ec.ports[1].File

	for v := range in {
		fmt.Fprintln(out, ToString(v))
	}
}

func fromLines(ec *evalCtx) {
	in := ec.ports[0].File
	out := ec.ports[1].Chan

	bufferedIn := bufio.NewReader(in)
	for {
		line, err := bufferedIn.ReadString('\n')
		if err == io.EOF {
			return
		} else if err != nil {
			throw(err)
		}
		out <- String(line[:len(line)-1])
	}
}

func ratFn(ec *evalCtx, arg Value) {
	out := ec.ports[1].Chan
	r, err := ToRat(arg)
	if err != nil {
		throw(err)
	}
	out <- r
}

// unpack takes any number of tables and output their list elements.
func unpack(ec *evalCtx) {
	in := ec.ports[0].Chan
	out := ec.ports[1].Chan

	for v := range in {
		if list, ok := v.(List); !ok {
			throw(ErrInput)
		} else {
			for _, e := range *list.inner {
				out <- e
			}
		}
	}
}

// fromJSON parses a stream of JSON data into Value's.
func fromJSON(ec *evalCtx) {
	in := ec.ports[0].File
	out := ec.ports[1].Chan

	dec := json.NewDecoder(in)
	var v interface{}
	for {
		err := dec.Decode(&v)
		if err != nil {
			if err == io.EOF {
				return
			}
			throw(err)
		}
		out <- FromJSONInterface(v)
	}
}

// each takes a single closure and applies it to all input values.
func each(ec *evalCtx, f *Closure) {
	in := ec.ports[0].Chan
in:
	for v := range in {
		newec := ec.fork("closure of each")
		ex := newec.PCall(f, []Value{v})
		ClosePorts(newec.ports)

		switch ex {
		case nil, Continue:
			// nop
		case Break:
			break in
		default:
			throw(ex)
		}
	}
}

func cd(ec *evalCtx, args []Value) {
	var dir string
	if len(args) == 0 {
		dir = mustGetHome("")
	} else if len(args) == 1 {
		dir = ToString(args[0])
	} else {
		throw(ErrArgs)
	}

	cdInner(dir, ec)
}

func cdInner(dir string, ec *evalCtx) {
	err := os.Chdir(dir)
	if err != nil {
		throw(err)
	}
	if ec.store != nil {
		pwd, err := os.Getwd()
		// BUG(xiaq): Possible error of os.Getwd after cd-ing is ignored.
		if err == nil {
			ec.store.AddDir(pwd)
		}
	}
}

func dirs(ec *evalCtx) {
	if ec.store == nil {
		throw(ErrStoreNotConnected)
	}
	dirs, err := ec.store.ListDirs()
	if err != nil {
		throw(errors.New("store error: " + err.Error()))
	}
	out := ec.ports[1].Chan
	for _, dir := range dirs {
		out <- Map{&map[Value]Value{
			String("path"):  String(dir.Path),
			String("score"): String(fmt.Sprint(dir.Score)),
		}}
	}
}

func jump(ec *evalCtx, arg string) {
	if ec.store == nil {
		throw(ErrStoreNotConnected)
	}
	dirs, err := ec.store.FindDirs(arg)
	if err != nil {
		throw(errors.New("store error: " + err.Error()))
	}
	if len(dirs) == 0 {
		throw(ErrNoMatchingDir)
	}
	dir := dirs[0].Path
	err = os.Chdir(dir)
	// TODO(xiaq): Remove directories that no longer exist
	if err != nil {
		throw(err)
	}
	ec.store.AddDir(dir)
}

func source(ec *evalCtx, fname string) {
	ec.Source(fname)
}

func toFloat(arg Value) (float64, error) {
	arg, ok := arg.(String)
	if !ok {
		return 0, fmt.Errorf("must be string")
	}
	num, err := strconv.ParseFloat(string(arg.(String)), 64)
	if err != nil {
		return 0, err
	}
	return num, nil
}

func toInt(arg Value) (int, error) {
	arg, ok := arg.(String)
	if !ok {
		return 0, fmt.Errorf("must be string")
	}
	num, err := strconv.Atoi(string(arg.(String)))
	if err != nil {
		return 0, err
	}
	return num, nil
}

func plus(ec *evalCtx, nums ...float64) {
	out := ec.ports[1].Chan
	sum := 0.0
	for _, f := range nums {
		sum += f
	}
	out <- String(fmt.Sprintf("%g", sum))
}

func minus(ec *evalCtx, sum float64, nums ...float64) {
	out := ec.ports[1].Chan
	for _, f := range nums {
		sum -= f
	}
	out <- String(fmt.Sprintf("%g", sum))
}

func times(ec *evalCtx, nums ...float64) {
	out := ec.ports[1].Chan
	prod := 1.0
	for _, f := range nums {
		prod *= f
	}
	out <- String(fmt.Sprintf("%g", prod))
}

func divide(ec *evalCtx, prod float64, nums ...float64) {
	out := ec.ports[1].Chan
	for _, f := range nums {
		prod /= f
	}
	out <- String(fmt.Sprintf("%g", prod))
}

func eq(ec *evalCtx, args []Value) {
	out := ec.ports[1].Chan
	if len(args) == 0 {
		throw(ErrArgs)
	}
	for i := 0; i+1 < len(args); i++ {
		if args[i] != args[i+1] {
			out <- Bool(false)
			return
		}
	}
	out <- Bool(true)
}

func deepeq(ec *evalCtx, args []Value) {
	out := ec.ports[1].Chan
	if len(args) == 0 {
		throw(ErrArgs)
	}
	for i := 0; i+1 < len(args); i++ {
		if !DeepEq(args[i], args[i+1]) {
			out <- Bool(false)
			return
		}
	}
	out <- Bool(true)
}

func take(ec *evalCtx, n int) {
	in := ec.ports[0].Chan
	out := ec.ports[1].Chan

	i := 0
	for v := range in {
		if i >= n {
			break
		}
		i++
		out <- v
	}
}

func drop(ec *evalCtx, n int) {
	in := ec.ports[0].Chan
	out := ec.ports[1].Chan

	for i := 0; i < n; i++ {
		<-in
	}
	for v := range in {
		out <- v
	}
}

func bind(ec *evalCtx, key string, function Value) {
	if ec.Editor == nil {
		throw(ErrNoEditor)
	}
	maybeThrow(ec.Editor.Bind(key, function))
}

func le(ec *evalCtx, name string, args ...Value) {
	if ec.Editor == nil {
		throw(ErrNoEditor)
	}
	maybeThrow(ec.Editor.Call(name, args))
}

func _sleep(ec *evalCtx, t float64) {
	time.Sleep(time.Duration(t) * time.Second)
}

func _stack(ec *evalCtx) {
	out := ec.ports[1].File

	// XXX dup with main.go
	buf := make([]byte, 1024)
	for runtime.Stack(buf, true) == cap(buf) {
		buf = make([]byte, cap(buf)*2)
	}
	out.Write(buf)
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"github.com/rkoesters/xkcd"
	"log"
	"os"
	"strconv"
)

func main() {
	var comic *xkcd.Comic
	var err error

	switch len(os.Args) {
	case 1:
		comic, err = xkcd.GetCurrent()
		if err != nil {
			log.Fatal(err)
		}
	case 2:
		num, err := strconv.Atoi(os.Args[1])
		if err != nil {
			log.Fatal(err)
		}

		comic, err = xkcd.Get(num)
		if err != nil {
			log.Fatal(err)
		}
	default:
		log.Fatal("args")
	}

	e := json.NewEncoder(os.Stdout)
	err = e.Encode(comic)
	if err != nil {
		log.Fatal(err)
	}
}
<commit_msg>Refactored xkcd-json.<commit_after>package main

import (
	"encoding/json"
	"flag"
	"github.com/rkoesters/xkcd"
	"log"
	"os"
)

var number = flag.Int("n", 0, "Comic number.")

func main() {
	flag.Parse()

	var comic *xkcd.Comic
	var err error

	if *number == 0 {
		comic, err = xkcd.GetCurrent()
	} else {
		comic, err = xkcd.Get(*number)
	}

	if err != nil {
		log.Fatal(err)
	}

	e := json.NewEncoder(os.Stdout)
	err = e.Encode(comic)
	if err != nil {
		log.Fatal(err)
	}
}
<|endoftext|>


<commit_before>package main

import (
    "fmt"
    "github.com/elvinyung/gash"
)

func main() {
    p := gash.KvPair{"samplekey", "samplevalue"}
    fmt.Println(p.Key)

    t := gash.CreateSimpleHash(1000, gash.Djb2)
    t.Insert("foo", 4)
    fmt.Println(t.Find("foo"))
    tg.Insert("foo", "asdfofsdja")
    fmt.Println(t.Find("foo"))
}
<commit_msg>Fix typo<commit_after>package main

import (
    "fmt"
    "github.com/elvinyung/gash"
)

func main() {
    p := gash.KvPair{"samplekey", "samplevalue"}
    fmt.Println(p.Key)

    t := gash.CreateSimpleHash(1000, gash.Djb2)
    t.Insert("foo", 4)
    fmt.Println(t.Find("foo"))
    t.Insert("foo", "asdfofsdja")
    fmt.Println(t.Find("foo"))
}
<|endoftext|>


<commit_before>package connections

import (
	"github.com/gorilla/websocket"
	"github.com/sirupsen/logrus"

	"github.com/ivan1993spb/snake-server/game"
	"github.com/ivan1993spb/snake-server/objects/snake"
)

type ConnectionWorker struct {
	conn      *websocket.Conn
	logger    *logrus.Logger
	chStop    chan struct{}
	chStopErr chan error
	chOutput  chan OutputMessage
	chInput   chan InputMessageType

	flagStarted bool
	flagStopped bool
}

func NewConnectionWorker(conn *websocket.Conn, logger *logrus.Logger) *ConnectionWorker {
	return &ConnectionWorker{
		conn:      conn,
		logger:    logger,
		chStopErr: make(chan error, 0),
	}
}

func (cw *ConnectionWorker) Start(game *game.Game) error {
	//return cw.conn.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, "close"), time.Now().Add(time.Second))

	if cw.flagStarted {
		// Return error
		return nil
	}

	cw.startRead()
	cw.startWrite()

	// Start connection read
	// Start connection write
	// Listen game events
	// Parse input messages (?) and send game commands

	s, _ := snake.CreateSnake(game.World())
	go game.RunObserver(s, 16)
	//game.RunObserver(cw, 16)

	cw.flagStarted = true

	return <-cw.chStopErr
}

func (cw *ConnectionWorker) startRead() {
	go func() {
		for {
			//messageType, r, err := cw.conn.NextReader()
			_, _, err := cw.conn.NextReader()
			if err != nil {
				cw.chStopErr <- err
				return
			}
		}
	}()
}

func (cw *ConnectionWorker) startWrite() {
	go func() {
		for {
			select {
			case message := <-cw.chOutput:
				cw.logger.Info(message)
				//err := cw.writeMessage(message)
				// TODO: Handler error.
			case <-cw.chStop:
				return
			}

		}
	}()
}

func (cw *ConnectionWorker) writeMessage(message *OutputMessage) error {
	w, err := cw.conn.NextWriter(websocket.TextMessage)
	if err != nil {
		return nil
	}

	// TODO: Write message.

	// TODO: Return error.

	w.Close()
	return nil
}

func (cw *ConnectionWorker) Run(ch <-chan game.Event) {
	for {
		select {
		case event := <-ch:
			cw.logger.Info(event)
		case <-cw.chStop:
			return
		}
	}
}
<commit_msg>Create connection worker channels and goroutines structure for input and output<commit_after>package connections

import (
	"sync"
	"time"

	"github.com/gorilla/websocket"
	"github.com/pquerna/ffjson/ffjson"
	"github.com/sirupsen/logrus"

	"github.com/ivan1993spb/snake-server/game"
)

type ConnectionWorker struct {
	conn        *websocket.Conn
	logger      *logrus.Logger
	chStop      chan struct{}
	chStopErr   chan error
	chOutput    chan OutputMessage
	chsInput    []chan InputMessage
	chsInputMux *sync.RWMutex

	flagStarted bool
	flagStopped bool
}

func NewConnectionWorker(conn *websocket.Conn, logger *logrus.Logger) *ConnectionWorker {
	return &ConnectionWorker{
		conn:      conn,
		logger:    logger,
		chStopErr: make(chan error, 0),
	}
}

func (cw *ConnectionWorker) Start(game *game.Game) error {
	//return cw.conn.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, "close"), time.Now().Add(time.Second))

	if cw.flagStarted {
		// Return error
		return nil
	}

	// Start connection read
	// Start connection write
	// Listen game events
	// Parse input messages (?) and send game commands

	//game.RunObserver(cw, 16)

	cw.flagStarted = true

	return <-cw.chStopErr
}

func (cw *ConnectionWorker) read() <-chan []byte {
	// TODO: Create buffer.
	chout := make(chan []byte, 0)

	go func() {
		defer close(chout)

		for {
			messageType, data, err := cw.conn.ReadMessage()
			if err != nil {
				// TODO: Handle error?
				return
			}

			if websocket.TextMessage != messageType {
				// TODO: Handle case - unexpected message type?
				continue
			}

			chout <- data
		}
	}()

	return chout
}

func (cw *ConnectionWorker) decode(chin <-chan []byte, stop <-chan struct{}) <-chan InputMessage {
	// TODO: Create buffer.
	chout := make(chan InputMessage, 0)

	go func() {
		defer close(chout)

		var decoder = ffjson.NewDecoder()

		for {
			select {
			case data := <-chin:
				var inputMessage *InputMessage
				if err := decoder.DecodeFast(data, &inputMessage); err != nil {
					// TODO: Handler error.
				} else {
					chout <- *inputMessage
				}
			case <-stop:
				return
			}
		}
	}()

	return chout
}

func (cw *ConnectionWorker) broadcastInputMessage(chin <-chan InputMessage, stop <-chan struct{}) {
	go func() {
		for {
			select {
			case inputMessage := <-chin:
				cw.chsInputMux.RLock()
				for _, ch := range cw.chsInput {
					select {
					case ch <- inputMessage:
					case <-stop:
						return
					}
				}
				cw.chsInputMux.RUnlock()
			case <-stop:
				return
			}
		}
	}()
}

func (cw *ConnectionWorker) Input(stop <-chan struct{}) <-chan InputMessage {
	// TODO: Create buffer.
	chProxy := make(chan InputMessage, 0)

	cw.chsInputMux.Lock()
	cw.chsInput = append(cw.chsInput, chProxy)
	cw.chsInputMux.Unlock()

	// TODO: Create buffer.
	chout := make(chan InputMessage, 0)

	go func() {
		defer close(chout)
		defer func() {
			cw.chsInputMux.Lock()
			for i := range cw.chsInput {
				if cw.chsInput[i] == chProxy {
					cw.chsInput = append(cw.chsInput[:i], cw.chsInput[i+1:]...)
					close(chProxy)
					break
				}
			}
			cw.chsInputMux.Unlock()
		}()

		for {
			select {
			case <-stop:
				return
			case <-cw.chStop:
				return
			case inputMessage := <-chProxy:
				// TODO: Create timeout.
				cw.sendInputMessage(chout, inputMessage, stop, time.Second)
			}
		}
	}()

	return chout
}

func (cw *ConnectionWorker) sendInputMessage(ch chan InputMessage, inputMessage InputMessage, stop <-chan struct{}, timeout time.Duration) {
	var timer = time.NewTimer(timeout)
	defer timer.Stop()
	if cap(ch) == 0 {
		select {
		case ch <- inputMessage:
		case <-cw.chStop:
		case <-stop:
		case <-timer.C:
		}
	} else {
		for {
			select {
			case ch <- inputMessage:
				return
			case <-cw.chStop:
				return
			case <-stop:
				return
			case <-timer.C:
				return
			default:
				if len(ch) == cap(ch) {
					<-ch
				}
			}
		}
	}
}

func (cw *ConnectionWorker) write(chin <-chan []byte, stop <-chan struct{}) {
	go func() {
		for {
			select {
			case data := <-chin:
				if err := cw.conn.WriteMessage(websocket.TextMessage, data); err != nil {
					// TODO: Handler error.
				}
			case <-stop:
				return
			}
		}
	}()
}

func (cw *ConnectionWorker) outputMessage(outputMessage OutputMessage) {
	select {
	case cw.chOutput <- outputMessage:
	case <-cw.chStop:
	}
}

func (cw *ConnectionWorker) encode(chin <-chan OutputMessage, stop <-chan struct{}) <-chan []byte {
	// TODO: Create buffer.
	chout := make(chan []byte, 0)

	go func() {
		defer close(chout)

		for {
			select {
			case message := <-chin:
				if data, err := ffjson.MarshalFast(message); err != nil {
					// TODO: Handler error.
				} else {
					chout <- data
				}
			case <-stop:
				return
			}
		}
	}()

	return chout
}

func (cw *ConnectionWorker) listen(chin <-chan game.Event, stop <-chan struct{}) {
	for {
		select {
		case event := <-chin:
			// TODO: Do stuff.
			cw.logger.Info(event)
		case <-stop:
			return
		}
	}
}
<|endoftext|>


<commit_before>package duktape

/*
# include "duktape.h"
static void go_duk_eval_string(duk_context *ctx, const char *str) {
  return duk_eval_string(ctx, str);
}
extern duk_ret_t goFuncCall(duk_context *ctx);
*/
import "C"
import "errors"
import "fmt"
import "log"
import "regexp"
import "time"
import "unsafe"

const (
	DUK_TYPE_NONE Type = iota
	DUK_TYPE_UNDEFINED
	DUK_TYPE_NULL
	DUK_TYPE_BOOLEAN
	DUK_TYPE_NUMBER
	DUK_TYPE_STRING
	DUK_TYPE_OBJECT
	DUK_TYPE_BUFFER
	DUK_TYPE_POINTER
)

const goFuncCallName = "__goFuncCall__"

type Type int

func (t Type) IsNone() bool      { return t == DUK_TYPE_NONE }
func (t Type) IsUndefined() bool { return t == DUK_TYPE_UNDEFINED }
func (t Type) IsNull() bool      { return t == DUK_TYPE_NULL }
func (t Type) IsBool() bool      { return t == DUK_TYPE_BOOLEAN }
func (t Type) IsNumber() bool    { return t == DUK_TYPE_NUMBER }
func (t Type) IsString() bool    { return t == DUK_TYPE_STRING }
func (t Type) IsObject() bool    { return t == DUK_TYPE_OBJECT }
func (t Type) IsBuffer() bool    { return t == DUK_TYPE_BUFFER }
func (t Type) IsPointer() bool   { return t == DUK_TYPE_POINTER }

type Context struct {
	duk_context unsafe.Pointer
}

// Returns initialized duktape context object
func NewContext() *Context {
	ctx := &Context{
		duk_context: C.duk_create_heap(nil, nil, nil, nil, nil),
	}
	ctx.defineGoFuncCall()
	return ctx
}

//export goFuncCall
func goFuncCall(ctx unsafe.Pointer) C.duk_ret_t {
	c := &Context{ctx}
	fmt.Printf("goFuncCall with context: %#v\n", c)
	if c.GetTop() == 0 {
		// unexpected call, without function name's hash
		panic("Go function call without arguments is not supported")
		return C.DUK_RET_UNSUPPORTED_ERROR
	}
	if !c.GetType(0).IsString() {
		// unexpected type of function name's hash
		panic("Wrong type of function's key argument")
		return C.DUK_RET_EVAL_ERROR
	}
	hash := c.GetString(0)
	if fn, ok := goFuncMap[hash]; ok {
		return C.duk_ret_t(fn(c))
	}
	log.Printf("hash is: %s", hash)
	panic("Unimplemented")
	return C.DUK_RET_UNIMPLEMENTED_ERROR
}

func getKeyFor(funcName string) string {
	c := 0
	key := fmt.Sprintf("__%s_%d%d__", funcName, time.Now().Nanosecond(), c)
	for {
		if _, ok := goFuncMap[key]; ok {
			c++
			key = fmt.Sprintf("__%s_%d%d__", funcName, time.Now().Nanosecond(), c)
			continue
		}
		break
	}
	return key
}

var goFuncMap = map[string]func(*Context) int{}
var reFuncName = regexp.MustCompile("^[a-z][a-z0-9]*([A-Z][a-z0-9]*)*$")

func (d *Context) PushGoFunc(name string, fn func(*Context) int) error {
	if !reFuncName.MatchString(name) {
		return errors.New("Malformed function name '" + name + "'")
	}
	key := getKeyFor(name)
	goFuncMap[key] = fn

	d.EvalString(fmt.Sprintf(`
      function %s (){
        %s.apply(this, ['%s'].concat(Array.prototype.slice.apply(arguments)));
      };
  `, name, goFuncCallName, key))
	d.Pop()
	return nil
}

func (d *Context) defineGoFuncCall() {
	d.PushGlobalObject()
	d.PushCFunction((*[0]byte)(C.goFuncCall), int(C.DUK_VARARGS))
	d.PutPropString(-2, goFuncCallName)
	d.Pop()
}

func (d *Context) GetTop() int {
	return int(C.duk_get_top(d.duk_context))
}

func (d *Context) PushCFunction(fn *[0]byte, nargs int) {
	C.duk_push_c_function(
		d.duk_context,
		fn,
		C.duk_idx_t(nargs),
	)
}

func (d *Context) EvalString(script string) {
	str := C.CString(script)
	defer C.free(unsafe.Pointer(str))
	C.go_duk_eval_string(d.duk_context, str)
}

func (d *Context) Pop() {
	C.duk_pop(d.duk_context)
}

func (d *Context) GetType(i int) Type {
	return Type(C.duk_get_type(d.duk_context, C.duk_idx_t(i)))
}

func (d *Context) GetString(i int) string {
	if d.GetType(i).IsString() {
		if s := C.duk_get_string(d.duk_context, C.duk_idx_t(i)); s != nil {
			return C.GoString(s)
		}
	}
	return ""
}

func (d *Context) PushGlobalObject() {
	C.duk_push_global_object(d.duk_context)
}

func (d *Context) GetNumber(i int) float64 {
	return float64(C.duk_get_number(d.duk_context, C.duk_idx_t(i)))
}

func (d *Context) PushNumber(i float64) {
	C.duk_push_number(d.duk_context, C.duk_double_t(i))
}

func (d *Context) PutPropString(i int, prop string) {
	str := C.CString(prop)
	defer C.free(unsafe.Pointer(str))
	C.duk_put_prop_string(d.duk_context, C.duk_idx_t(i), str)
}

func (d *Context) DestroyHeap() {
	C.duk_destroy_heap(d.duk_context)
}
<commit_msg>PushString method added<commit_after>package duktape

/*
# include "duktape.h"
static void go_duk_eval_string(duk_context *ctx, const char *str) {
  return duk_eval_string(ctx, str);
}
extern duk_ret_t goFuncCall(duk_context *ctx);
*/
import "C"
import "errors"
import "fmt"
import "log"
import "regexp"
import "time"
import "unsafe"

const (
	DUK_TYPE_NONE Type = iota
	DUK_TYPE_UNDEFINED
	DUK_TYPE_NULL
	DUK_TYPE_BOOLEAN
	DUK_TYPE_NUMBER
	DUK_TYPE_STRING
	DUK_TYPE_OBJECT
	DUK_TYPE_BUFFER
	DUK_TYPE_POINTER
)

const goFuncCallName = "__goFuncCall__"

type Type int

func (t Type) IsNone() bool      { return t == DUK_TYPE_NONE }
func (t Type) IsUndefined() bool { return t == DUK_TYPE_UNDEFINED }
func (t Type) IsNull() bool      { return t == DUK_TYPE_NULL }
func (t Type) IsBool() bool      { return t == DUK_TYPE_BOOLEAN }
func (t Type) IsNumber() bool    { return t == DUK_TYPE_NUMBER }
func (t Type) IsString() bool    { return t == DUK_TYPE_STRING }
func (t Type) IsObject() bool    { return t == DUK_TYPE_OBJECT }
func (t Type) IsBuffer() bool    { return t == DUK_TYPE_BUFFER }
func (t Type) IsPointer() bool   { return t == DUK_TYPE_POINTER }

type Context struct {
	duk_context unsafe.Pointer
}

// Returns initialized duktape context object
func NewContext() *Context {
	ctx := &Context{
		duk_context: C.duk_create_heap(nil, nil, nil, nil, nil),
	}
	ctx.defineGoFuncCall()
	return ctx
}

//export goFuncCall
func goFuncCall(ctx unsafe.Pointer) C.duk_ret_t {
	c := &Context{ctx}
	fmt.Printf("goFuncCall with context: %#v\n", c)
	if c.GetTop() == 0 {
		// unexpected call, without function name's hash
		panic("Go function call without arguments is not supported")
		return C.DUK_RET_UNSUPPORTED_ERROR
	}
	if !c.GetType(0).IsString() {
		// unexpected type of function name's hash
		panic("Wrong type of function's key argument")
		return C.DUK_RET_EVAL_ERROR
	}
	hash := c.GetString(0)
	if fn, ok := goFuncMap[hash]; ok {
		return C.duk_ret_t(fn(c))
	}
	log.Printf("hash is: %s", hash)
	panic("Unimplemented")
	return C.DUK_RET_UNIMPLEMENTED_ERROR
}

func getKeyFor(funcName string) string {
	c := 0
	key := fmt.Sprintf("__%s_%d%d__", funcName, time.Now().Nanosecond(), c)
	for {
		if _, ok := goFuncMap[key]; ok {
			c++
			key = fmt.Sprintf("__%s_%d%d__", funcName, time.Now().Nanosecond(), c)
			continue
		}
		break
	}
	return key
}

var goFuncMap = map[string]func(*Context) int{}
var reFuncName = regexp.MustCompile("^[a-z][a-z0-9]*([A-Z][a-z0-9]*)*$")

func (d *Context) PushGoFunc(name string, fn func(*Context) int) error {
	if !reFuncName.MatchString(name) {
		return errors.New("Malformed function name '" + name + "'")
	}
	key := getKeyFor(name)
	goFuncMap[key] = fn

	d.EvalString(fmt.Sprintf(`
      function %s (){
        %s.apply(this, ['%s'].concat(Array.prototype.slice.apply(arguments)));
      };
  `, name, goFuncCallName, key))
	d.Pop()
	return nil
}

func (d *Context) defineGoFuncCall() {
	d.PushGlobalObject()
	d.PushCFunction((*[0]byte)(C.goFuncCall), int(C.DUK_VARARGS))
	d.PutPropString(-2, goFuncCallName)
	d.Pop()
}

func (d *Context) GetTop() int {
	return int(C.duk_get_top(d.duk_context))
}

func (d *Context) PushCFunction(fn *[0]byte, nargs int) {
	C.duk_push_c_function(
		d.duk_context,
		fn,
		C.duk_idx_t(nargs),
	)
}

func (d *Context) EvalString(script string) {
	str := C.CString(script)
	defer C.free(unsafe.Pointer(str))
	C.go_duk_eval_string(d.duk_context, str)
}

func (d *Context) Pop() {
	C.duk_pop(d.duk_context)
}

func (d *Context) GetType(i int) Type {
	return Type(C.duk_get_type(d.duk_context, C.duk_idx_t(i)))
}

func (d *Context) GetString(i int) string {
	if d.GetType(i).IsString() {
		if s := C.duk_get_string(d.duk_context, C.duk_idx_t(i)); s != nil {
			return C.GoString(s)
		}
	}
	return ""
}

func (d *Context) PushGlobalObject() {
	C.duk_push_global_object(d.duk_context)
}

func (d *Context) GetNumber(i int) float64 {
	return float64(C.duk_get_number(d.duk_context, C.duk_idx_t(i)))
}

func (d *Context) PushNumber(i float64) {
	C.duk_push_number(d.duk_context, C.duk_double_t(i))
}

func (d *Context) PushString(str string) {
	s := C.CString(str)
	defer C.free(unsafe.Pointer(s))
	C.duk_push_string(d.duk_context, s)
}

func (d *Context) PutPropString(i int, prop string) {
	str := C.CString(prop)
	defer C.free(unsafe.Pointer(str))
	C.duk_put_prop_string(d.duk_context, C.duk_idx_t(i), str)
}

func (d *Context) DestroyHeap() {
	C.duk_destroy_heap(d.duk_context)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/rackspace/gophercloud/openstack/identity"
	"github.com/rackspace/gophercloud/openstack/utils"
	"github.com/rackspace/gophercloud/rackspace/monitoring"
	"github.com/rackspace/gophercloud/rackspace/monitoring/notificationPlans"
)

func main() {
	ao, err := utils.AuthOptions()
	if err != nil {
		panic(err)
	}

	ao.AllowReauth = true
	r, err := identity.Authenticate(ao)
	if err != nil {
		panic(err)
	}

	// Find the cloud monitoring API

	sc, err := identity.GetServiceCatalog(r)
	if err != nil {
		panic(err)
	}

	ces, err := sc.CatalogEntries()
	if err != nil {
		panic(err)
	}

	monUrl, err := findMonitoringEndpoint(ces)
	if err != nil {
		panic(err)
	}

	// Build ourselves an interface to cloud monitoring!

	np := notificationPlans.NewClient(monitoring.Options{
		Endpoint:       monUrl,
		AuthOptions:    ao,
		Authentication: r,
	})

	// Try to delete a bogus notification plan

	dr, err := np.Delete("ajkhdlkajhdflkajshdf")
	if err != nil {
		fmt.Printf("%#v\n", err)
	}
	fmt.Printf("%#v\n", dr)
}

func findMonitoringEndpoint(ces []identity.CatalogEntry) (string, error) {
	for _, ce := range ces {
		if ce.Type == "rax:monitor" {
			return ce.Endpoints[0].PublicURL, nil
		}
	}
	return "", fmt.Errorf("No monitoring API in the service catalog")
}
<commit_msg>Forgot to remove unused file.<commit_after><|endoftext|>


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

// This file is a Firewall manager,
// as of now responsibilities of the Firewall manager is to
// 1. Create per tenant per segment chains
// 2. Divert traffic from base chains into ones created above
// 3. Install Firewall rules into managed chains (restrictive by default)
// 3.1 There is a set of permissive rules per chain
// 3.2 There is also u32 rule per chain that allows traffic between tenant endpoints
//
// Firewall created per user request as follows
// - Firewall := NewFirewall(NetIf)
// where NetIf is a struct defined in netif.go
//
// There are 4 public function available
// - Firewall.CreateChains() - creates managed Firewall chains
// - Firewall.DiverTrafficToPaniIptablesChain() - diverts traffic to/from/via interface
// - Firewall.CreateRules() - installs basic permissive rules
// - Firewall.CreateU32Rules() - installs u32 rules

import (
	"fmt"
	"log"
	"net"
	"strconv"
	"strings"
)

const (
	numberOfDefaultChains = 3 // INPUT,OUTPUT,FORWARD
	inputChainIndex       = 0
	outputChainIndex      = 1
	forwardChainIndex     = 2
)

// Firewall type describes state of firewall rules for the given endpoint.
// Currently methods of this type have no critical sections because
// Firewall won't do anything if per-tenant chains already created.
// But any new features must be safe for concurrent execution.
type Firewall struct {
	chains        [numberOfDefaultChains]FirewallChain
	u32Filter     string
	chainPrefix   string
	interfaceName string
	Agent         *Agent // access field to Agent
}

// FirewallChain describes state of the particular firewall chain.
type FirewallChain struct {
	baseChain string
	direction string
	rules     []string
	chainName string
}

// NewFirewallChain initializes new a firewall chain.
func NewFirewallChain(baseChain string, direction string, rules []string, chainName string) *FirewallChain {
	return &FirewallChain{baseChain, direction, rules, chainName}
}

// NewFirewall returns fully initialized firewall struct.
func NewFirewall(netif NetIf, agent *Agent) (*Firewall, error) {
	fw := new(Firewall)
	fw.Agent = agent
	if err := fw.Init(netif); err != nil {
		return nil, err
	}
	return fw, nil
}

// Collection of firewall rules.
type firewallRules []string

// Add appends a new firewall rule to the collection of firewall rules.
func (r *firewallRules) Add(content string) {
	*r = append(*r, content)
}

// prepareChainName returns a chain name with tenant-segment specific prefix.
func (fw *Firewall) prepareChainName(chainName string) string {
	return fmt.Sprintf("%s%s", fw.chainPrefix, chainName)
}

// Init initializes current firewall with a data from the given endpoint.
func (fw *Firewall) Init(netif NetIf) error {
	var err error
	fw.u32Filter, fw.chainPrefix, err = fw.prepareU32Rules(netif.Ip)
	if err != nil {
		// TODO need personalized error here, or even panic
		return err
	}
	fw.interfaceName = netif.Name

	// Default permissive rules to allow ssh/ICMP
	// from the Host to the Endpoint. Normally wildcard U32 filter would
	// block all the Host to Endpoint traffic.
	inputRules := new(firewallRules)
	inputRules.Add(fmt.Sprintf("-d %s", fw.Agent.networkConfig.currentHostIP))
	inputRules.Add(fmt.Sprintf("-d %s/%d", fw.Agent.networkConfig.currentHostGW, fw.Agent.networkConfig.currentHostGWNetSize))
	inputRules.Add("-p udp --sport 68 --dport 67 -d 255.255.255.255")
	inputRules.Add(fmt.Sprintf("-p tcp -m tcp --sport 22 -d %s", fw.Agent.networkConfig.currentHostIP))

	outputRules := new(firewallRules)
	outputRules.Add(fmt.Sprintf("-d %s", fw.Agent.networkConfig.currentHostIP))
	outputRules.Add(fmt.Sprintf("-d %s/%d", fw.Agent.networkConfig.currentHostGW, fw.Agent.networkConfig.currentHostGWNetSize))
	outputRules.Add(fmt.Sprintf("-p tcp -m tcp --sport 22 -d %s", fw.Agent.networkConfig.currentHostIP))

	forwardRules := new(firewallRules)

	c1 := NewFirewallChain("INPUT", "i", *inputRules, fw.prepareChainName("INPUT"))
	c2 := NewFirewallChain("OUTPUT", "o", *outputRules, fw.prepareChainName("OUTPUT"))
	c3 := NewFirewallChain("FORWARD", "i", *forwardRules, fw.prepareChainName("FORWARD"))
	chains := [numberOfDefaultChains]FirewallChain{*c1, *c2, *c3}
	fw.chains = chains
	return nil
}

// isChainExist verifies if given iptables chain exists.
// Returns true chain exists.
func (fw *Firewall) isChainExist(chain int) bool {
	cmd := "/sbin/iptables"
	args := []string{"-L", fw.chains[chain].chainName}
	_, err := fw.Agent.Helper.Executor.Exec(cmd, args)
	if err != nil {
		return false
	}
	return true
}

// detectMissingChains checks which Firewall chains haven't been created yet.
func (fw *Firewall) detectMissingChains() []int {
	var ret []int
	for chain := range fw.chains {
		log.Print("Testing chain", chain)
		if !fw.isChainExist(chain) {
			log.Print(">> Testing chain success", chain)
			ret = append(ret, chain)
		}
	}
	return ret
}

// CreateChains creates Firewall chains such as
// pani-T0S0-OUTPUT, pani-T0S0-FORWARD, pani-T0S0-INPUT.
func (fw *Firewall) CreateChains(newChains []int) error {
	for chain := range newChains {
		cmd := "/sbin/iptables"
		args := []string{"-N", fw.chains[chain].chainName}
		_, err := fw.Agent.Helper.Executor.Exec(cmd, args)
		if err != nil {
			return err
		}
	}
	return nil
}

// DiverTrafficToPaniIptablesChain injects iptables rules to send traffic into the PANI chain.
func (fw *Firewall) DiverTrafficToPaniIptablesChain(chain int) error {
	// Should be like that
	// iptables -A INPUT -i tap1234 -j PANI-T0S1-INPUT
	log.Print("Diverting traffic in", chain)
	baseChain := fw.chains[chain].baseChain
	direction := fmt.Sprintf("-%s", fw.chains[chain].direction)
	chainName := fw.chains[chain].chainName
	cmd := "/sbin/iptables"
	args := []string{"-A", baseChain, direction, fw.interfaceName, "-j", chainName}
	_, err := fw.Agent.Helper.Executor.Exec(cmd, args)
	if err != nil {
		log.Print("Diverting traffic failed", chain)
		return err
	}
	log.Print("Diverting traffic success", chain)
	return nil
}

// CreateRules creates permissive iptables rules for the given PANI chain
// to allow a traffic to flow between the Host and Endpoint.
func (fw *Firewall) CreateRules(chain int) error {
	log.Print("Creating firewall rules for chain", chain)
	for rule := range fw.chains[chain].rules {
		chainName := fw.chains[chain].chainName
		cmd := "/sbin/iptables"
		args := []string{"-A", chainName}
		args = append(args, strings.Split(fw.chains[chain].rules[rule], " ")...)
		args = append(args, []string{"-j", "ACCEPT"}...)
		_, err := fw.Agent.Helper.Executor.Exec(cmd, args)
		if err != nil {
			return err
			log.Print("Creating firewall rules failed")
		}
	}
	log.Print("Creating firewall rules success")
	return nil
}

// CreateU32Rules creates wildcard iptables rules for the given PANI chain.
// These rules serve to restrict traffic between segments and tenants.
func (fw *Firewall) CreateU32Rules(chain int) error {
	log.Print("Creating U32 firewall rules for chain", chain)
	chainName := fw.chains[chain].chainName
	cmd := "/sbin/iptables"
	args := []string{"-A", chainName, "-m", "u32", "--u32", fw.u32Filter, "-j", "ACCEPT"}
	_, err := fw.Agent.Helper.Executor.Exec(cmd, args)
	if err != nil {
		return err
		log.Print("Creating U32 firewall rules failed")
	}
	log.Print("Creating U32 firewall rules failed for chain", chain)
	return nil
}

// prepareTenantSegmentMask returns integer representation of a netmask
// for tenant+segment bits in pseudo network.
func (fw *Firewall) prepareTenantSegmentMask() uint64 {
	var res uint64
	tenantBits := fw.Agent.networkConfig.TenantBits()
	segmentBits := fw.Agent.networkConfig.SegmentBits()
	combinedTSBits := tenantBits + segmentBits
	endpointBits := fw.Agent.networkConfig.EndpointBits()
	res = ((1 << combinedTSBits) - 1) << endpointBits
	return res
}

// ipToInt transforms IP address from net.IP form to integer form.
// Stolen from IPAM/config, should really be in some shared library.
func ipToInt(ip net.IP) uint64 {
	return uint64(ip[12])<<24 | uint64(ip[13])<<16 | uint64(ip[14])<<8 | uint64(ip[15])
}

// PseudoNetNetmaskInt returns integer representation of pseudo net netmask.
func (fw *Firewall) PseudoNetNetmaskInt() (uint64, error) {
	cidr, err := fw.Agent.networkConfig.PNetCIDR()
	if err != nil {
		return 0, err
	}
	pNetMaskInt, err := MaskToInt(cidr.Mask)
	if err != nil {
		return 0, err
	}
	return pNetMaskInt, nil
}

// MaskToInt converts net.IPMask to integer.
// TODO Not strictly firewall method, maybe put in different place.
func MaskToInt(mask net.IPMask) (uint64, error) {
	var imask uint64
	m, err := strconv.ParseInt(mask.String(), 16, 64)
	imask = uint64(m)
	if err != nil {
		return 0, err
	}
	return imask, nil
}

// prepareU32Rules generates Firewall rules for U32 iptables module.
// This rules implemet PANI tenant/segment filtering
//   Return the filter rules for the iptables u32 module.
//   Goal: Filter out any traffic that does not have the same tenant and segment
//   bits in the destination address as the interface itself.
//   These bits can be extracted from the IP address: This is the address that
//   we are assigning to the interface. The function is to be called when the
//   interface is set up. The passed-in address therefore can be trusted: It is
//   not taken from a packet.
//      Example:
//      ipAddr = "10.0.1.4"
//
//      Return:
//      filter = '12&0xFF00FF00=0xA000100 && 16&0xFF00FF00=0xA000100'
//      chainPrefix = 'pani-T0S1-'
//
//   Also return the chain-prefix we'll use for this interface. This is
//   typically a string such as:
//       pani-T<tenant-id>S<segment-id>-
//   For example, with tenant 1 and segment 2, this would be:
//       pani-T1S2-
func (fw *Firewall) prepareU32Rules(ipAddr net.IP) (string, string, error) {
	fullMask, err := fw.prepareNetmaskBits()
	if err != nil {
		return "", "", err
	}
	addr := ipToInt(ipAddr)
	if err != nil {
		return "", "", err
	}
	filter1 := fmt.Sprintf("0x%X=0x%X", fullMask, addr&fullMask)
	filter := fmt.Sprintf("12&%s && 16&%s", filter1, filter1)
	tenantID := fw.extractTenantID(addr)
	segmentID := fw.extractSegmentID(addr)
	chainPrefix := fmt.Sprintf("pani-T%dS%d-", tenantID, segmentID)
	return filter, chainPrefix, nil
}

// prepareNetmaskBits returns integer representation of pseudo network netmask.
func (fw *Firewall) prepareNetmaskBits() (uint64, error) {
	iCidrMask, err := fw.PseudoNetNetmaskInt()
	if err != nil {
		return 0, err
	}
	combinedTSMask := fw.prepareTenantSegmentMask()
	res := iCidrMask | combinedTSMask
	return res, nil
}

// extractSegmentID extracts segment id from the given ip address.
// This is possible because segment id encoded in the ip address.
func (fw *Firewall) extractSegmentID(addr uint64) uint64 {
	endpointBits := fw.Agent.networkConfig.EndpointBits()
	segmentBits := fw.Agent.networkConfig.SegmentBits()
	sid := (addr >> endpointBits) & ((1 << segmentBits) - 1)
	return sid
}

// extractTenantID extracts tenant id from given the ip address.
// This is possible because tenant id encoded in the ip address.
func (fw *Firewall) extractTenantID(addr uint64) uint64 {
	endpointBits := fw.Agent.networkConfig.EndpointBits()
	segmentBits := fw.Agent.networkConfig.SegmentBits()
	tenantBits := fw.Agent.networkConfig.TenantBits()
	tid := (addr >> (endpointBits + segmentBits)) & ((1 << tenantBits) - 1)
	return tid
}

// provisionFirewallRules orchestrates Firewall to satisfy request
// to provision new endpoint.
// Creates per-tenant, per-segment iptables chains, diverts
// all traffic to/from/through netif.name interface to a proper chains.
// Currently tested with pani ML2 driver.
func provisionFirewallRules(netif NetIf, agent *Agent) error {
	log.Print("Firewall: Initializing")
	fw, err := NewFirewall(netif, agent)
	if err != nil {
		log.Fatal("Failed to initialize firewall ", err)
	}

	missingChains := fw.detectMissingChains()
	log.Print("Firewall: creating chains")
	err = fw.CreateChains(missingChains)
	if err != nil {
		return err
	}
	for chain := range missingChains {
		if err := fw.DiverTrafficToPaniIptablesChain(chain); err != nil {
			return err
		}
		if err := fw.CreateRules(chain); err != nil {
			return err
		}
		if err := fw.CreateU32Rules(chain); err != nil {
			return err
		}
	}
	return nil
}
<commit_msg>Refactoring and comments in firewall.go<commit_after>// Copyright (c) 2015 Pani Networks
// All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License"); you may
// not use this file except in compliance with the License. You may obtain
// a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
//  distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
// WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
// License for the specific language governing permissions and limitations
// under the License.

// This file is a Firewall manager,
// as of now responsibilities of the Firewall manager is to
// 1. Create per tenant per segment chains
// 2. Divert traffic from base chains into ones created above
// 3. Install Firewall rules into managed chains (restrictive by default)
// 3.1 There is a set of permissive rules per chain
// 3.2 There is also u32 rule per chain that allows traffic between tenant endpoints
//
// Firewall created per user request as follows
// - Firewall := NewFirewall(NetIf)
// where NetIf is a struct defined in netif.go
//
// There are 4 public function available
// - Firewall.CreateChains() - creates managed Firewall chains
// - Firewall.DivertTrafficToPaniIptablesChain() - diverts traffic to/from/via interface
// - Firewall.CreateRules() - installs basic permissive rules
// - Firewall.CreateU32Rules() - installs u32 rules
package agent

import (
	"fmt"
	"log"
	"net"
	"strconv"
	"strings"
)

const (
	numberOfDefaultChains = 3 // INPUT,OUTPUT,FORWARD
	inputChainIndex       = 0
	outputChainIndex      = 1
	forwardChainIndex     = 2
)

// Firewall describes state of firewall rules for the given endpoint.
// Currently methods of this type have no critical sections because
// Firewall won't do anything if per-tenant chains already created.
// But any new features must be safe for concurrent execution.
type Firewall struct {
	chains        [numberOfDefaultChains]FirewallChain
	u32Filter     string
	chainPrefix   string
	interfaceName string
	Agent         *Agent // access field to Agent
}

// FirewallChain describes state of the particular firewall chain.
type FirewallChain struct {
	baseChain string
	direction string
	rules     []string
	chainName string
}

// NewFirewallChain initializes a new firewall chain.
func NewFirewallChain(baseChain string, direction string, rules []string, chainName string) *FirewallChain {
	return &FirewallChain{baseChain, direction, rules, chainName}
}

// NewFirewall returns fully initialized firewall struct, with rules and chains
// configured for diven endpoint.
func NewFirewall(netif NetIf, agent *Agent) (*Firewall, error) {
	fw := new(Firewall)
	fw.Agent = agent
	if err := fw.Init(netif); err != nil {
		return nil, err
	}
	return fw, nil
}

// Collection of firewall rules.
type firewallRules []string

// Add appends a new firewall rule to the collection of firewall rules.
func (r *firewallRules) Add(content string) {
	*r = append(*r, content)
}

// prepareChainName returns a chain name with tenant-segment specific prefix.
func (fw *Firewall) prepareChainName(chainName string) string {
	return fmt.Sprintf("%s%s", fw.chainPrefix, chainName)
}

// Init initializes current firewall with a data from the given endpoint.
func (fw *Firewall) Init(netif NetIf) error {
	var err error
	fw.u32Filter, fw.chainPrefix, err = fw.prepareU32Rules(netif.Ip)
	if err != nil {
		// TODO need personalized error here, or even panic
		return err
	}
	fw.interfaceName = netif.Name

	// Default permissive rules to allow ssh/ICMP
	// from the Host to the Endpoint. Normally wildcard U32 filter would
	// block all the Host to Endpoint traffic.
	inputRules := new(firewallRules)
	inputRules.Add(fmt.Sprintf("-d %s", fw.Agent.networkConfig.currentHostIP))
	inputRules.Add(fmt.Sprintf("-d %s/%d", fw.Agent.networkConfig.currentHostGW, fw.Agent.networkConfig.currentHostGWNetSize))
	inputRules.Add("-p udp --sport 68 --dport 67 -d 255.255.255.255")
	inputRules.Add(fmt.Sprintf("-p tcp -m tcp --sport 22 -d %s", fw.Agent.networkConfig.currentHostIP))

	outputRules := new(firewallRules)
	outputRules.Add(fmt.Sprintf("-d %s", fw.Agent.networkConfig.currentHostIP))
	outputRules.Add(fmt.Sprintf("-d %s/%d", fw.Agent.networkConfig.currentHostGW, fw.Agent.networkConfig.currentHostGWNetSize))
	outputRules.Add(fmt.Sprintf("-p tcp -m tcp --sport 22 -d %s", fw.Agent.networkConfig.currentHostIP))

	forwardRules := new(firewallRules)

	c1 := NewFirewallChain("INPUT", "i", *inputRules, fw.prepareChainName("INPUT"))
	c2 := NewFirewallChain("OUTPUT", "o", *outputRules, fw.prepareChainName("OUTPUT"))
	c3 := NewFirewallChain("FORWARD", "i", *forwardRules, fw.prepareChainName("FORWARD"))
	chains := [numberOfDefaultChains]FirewallChain{*c1, *c2, *c3}
	fw.chains = chains
	return nil
}

// isChainExist verifies if given iptables chain exists.
// Returns true chain exists.
func (fw *Firewall) isChainExist(chain int) bool {
	cmd := "/sbin/iptables"
	args := []string{"-L", fw.chains[chain].chainName}
	_, err := fw.Agent.Helper.Executor.Exec(cmd, args)
	if err != nil {
		return false
	}
	return true
}

// detectMissingChains checks which Firewall chains haven't been created yet.
func (fw *Firewall) detectMissingChains() []int {
	var ret []int
	for chain := range fw.chains {
		log.Print("Testing chain", chain)
		if !fw.isChainExist(chain) {
			log.Print(">> Testing chain success", chain)
			ret = append(ret, chain)
		}
	}
	return ret
}

// CreateChains creates Firewall chains such as
// pani-T0S0-OUTPUT, pani-T0S0-FORWARD, pani-T0S0-INPUT.
func (fw *Firewall) CreateChains(newChains []int) error {
	for chain := range newChains {
		cmd := "/sbin/iptables"
		args := []string{"-N", fw.chains[chain].chainName}
		_, err := fw.Agent.Helper.Executor.Exec(cmd, args)
		if err != nil {
			return err
		}
	}
	return nil
}

// DivertTrafficToPaniIptablesChain injects iptables rules to send traffic into the PANI chain.
func (fw *Firewall) DivertTrafficToPaniIptablesChain(chain int) error {
	// Should be like that
	// iptables -A INPUT -i tap1234 -j PANI-T0S1-INPUT
	log.Print("Diverting traffic in", chain)
	baseChain := fw.chains[chain].baseChain
	direction := fmt.Sprintf("-%s", fw.chains[chain].direction)
	chainName := fw.chains[chain].chainName
	cmd := "/sbin/iptables"
	args := []string{"-A", baseChain, direction, fw.interfaceName, "-j", chainName}
	_, err := fw.Agent.Helper.Executor.Exec(cmd, args)
	if err != nil {
		log.Print("Diverting traffic failed", chain)
		return err
	}
	log.Print("Diverting traffic success", chain)
	return nil
}

// CreateRules creates permissive iptables rules for the given PANI chain
// to allow a traffic to flow between the Host and Endpoint.
func (fw *Firewall) CreateRules(chain int) error {
	log.Print("Creating firewall rules for chain", chain)
	for rule := range fw.chains[chain].rules {
		chainName := fw.chains[chain].chainName
		cmd := "/sbin/iptables"
		args := []string{"-A", chainName}
		args = append(args, strings.Split(fw.chains[chain].rules[rule], " ")...)
		args = append(args, []string{"-j", "ACCEPT"}...)
		_, err := fw.Agent.Helper.Executor.Exec(cmd, args)
		if err != nil {
			return err
			log.Print("Creating firewall rules failed")
		}
	}
	log.Print("Creating firewall rules success")
	return nil
}

// CreateU32Rules creates wildcard iptables rules for the given PANI chain.
// These rules serve to restrict traffic between segments and tenants.
func (fw *Firewall) CreateU32Rules(chain int) error {
	log.Print("Creating U32 firewall rules for chain", chain)
	chainName := fw.chains[chain].chainName
	cmd := "/sbin/iptables"
	args := []string{"-A", chainName, "-m", "u32", "--u32", fw.u32Filter, "-j", "ACCEPT"}
	_, err := fw.Agent.Helper.Executor.Exec(cmd, args)
	if err != nil {
		return err
		log.Print("Creating U32 firewall rules failed")
	}
	log.Print("Creating U32 firewall rules failed for chain", chain)
	return nil
}

// prepareTenantSegmentMask returns integer representation of a netmask
// for tenant+segment bits in pseudo network.
func (fw *Firewall) prepareTenantSegmentMask() uint64 {
	var res uint64
	tenantBits := fw.Agent.networkConfig.TenantBits()
	segmentBits := fw.Agent.networkConfig.SegmentBits()
	combinedTSBits := tenantBits + segmentBits
	endpointBits := fw.Agent.networkConfig.EndpointBits()
	res = ((1 << combinedTSBits) - 1) << endpointBits
	return res
}

// ipToInt transforms IP address from net.IP form to integer form.
// Stolen from IPAM/config, should really be in some shared library.
func ipToInt(ip net.IP) uint64 {
	return uint64(ip[12])<<24 | uint64(ip[13])<<16 | uint64(ip[14])<<8 | uint64(ip[15])
}

// PseudoNetNetmaskInt returns integer representation of pseudo net netmask.
func (fw *Firewall) PseudoNetNetmaskInt() (uint64, error) {
	cidr, err := fw.Agent.networkConfig.PNetCIDR()
	if err != nil {
		return 0, err
	}
	pNetMaskInt, err := MaskToInt(cidr.Mask)
	if err != nil {
		return 0, err
	}
	return pNetMaskInt, nil
}

// MaskToInt converts net.IPMask to integer.
// TODO Not strictly firewall method, maybe put in different place.
func MaskToInt(mask net.IPMask) (uint64, error) {
	var imask uint64
	m, err := strconv.ParseInt(mask.String(), 16, 64)
	imask = uint64(m)
	if err != nil {
		return 0, err
	}
	return imask, nil
}

// prepareU32Rules generates Firewall rules for U32 iptables module.
// This rules implemet PANI tenant/segment filtering
//   Return the filter rules for the iptables u32 module.
//   Goal: Filter out any traffic that does not have the same tenant and segment
//   bits in the destination address as the interface itself.
//   These bits can be extracted from the IP address: This is the address that
//   we are assigning to the interface. The function is to be called when the
//   interface is set up. The passed-in address therefore can be trusted: It is
//   not taken from a packet.
//      Example:
//      ipAddr = "10.0.1.4"
//
//      Return:
//      filter = '12&0xFF00FF00=0xA000100 && 16&0xFF00FF00=0xA000100'
//      chainPrefix = 'pani-T0S1-'
//
//   Also return the chain-prefix we'll use for this interface. This is
//   typically a string such as:
//       pani-T<tenant-id>S<segment-id>-
//   For example, with tenant 1 and segment 2, this would be:
//       pani-T1S2-
func (fw *Firewall) prepareU32Rules(ipAddr net.IP) (string, string, error) {
	fullMask, err := fw.prepareNetmaskBits()
	if err != nil {
		return "", "", err
	}
	addr := ipToInt(ipAddr)
	if err != nil {
		return "", "", err
	}
	filter1 := fmt.Sprintf("0x%X=0x%X", fullMask, addr&fullMask)
	filter := fmt.Sprintf("12&%s && 16&%s", filter1, filter1)
	tenantID := fw.extractTenantID(addr)
	segmentID := fw.extractSegmentID(addr)
	chainPrefix := fmt.Sprintf("pani-T%dS%d-", tenantID, segmentID)
	return filter, chainPrefix, nil
}

// prepareNetmaskBits returns integer representation of pseudo network netmask.
func (fw *Firewall) prepareNetmaskBits() (uint64, error) {
	iCidrMask, err := fw.PseudoNetNetmaskInt()
	if err != nil {
		return 0, err
	}
	combinedTSMask := fw.prepareTenantSegmentMask()
	res := iCidrMask | combinedTSMask
	return res, nil
}

// extractSegmentID extracts segment id from the given ip address.
// This is possible because segment id encoded in the ip address.
func (fw *Firewall) extractSegmentID(addr uint64) uint64 {
	endpointBits := fw.Agent.networkConfig.EndpointBits()
	segmentBits := fw.Agent.networkConfig.SegmentBits()
	sid := (addr >> endpointBits) & ((1 << segmentBits) - 1)
	return sid
}

// extractTenantID extracts tenant id from given the ip address.
// This is possible because tenant id encoded in the ip address.
func (fw *Firewall) extractTenantID(addr uint64) uint64 {
	endpointBits := fw.Agent.networkConfig.EndpointBits()
	segmentBits := fw.Agent.networkConfig.SegmentBits()
	tenantBits := fw.Agent.networkConfig.TenantBits()
	tid := (addr >> (endpointBits + segmentBits)) & ((1 << tenantBits) - 1)
	return tid
}

// provisionFirewallRules orchestrates Firewall to satisfy request
// to provision new endpoint.
// Creates per-tenant, per-segment iptables chains, diverts
// all traffic to/from/through netif.name interface to a proper chains.
// Currently tested with pani ML2 driver.
func provisionFirewallRules(netif NetIf, agent *Agent) error {
	log.Print("Firewall: Initializing")
	fw, err := NewFirewall(netif, agent)
	if err != nil {
		log.Fatal("Failed to initialize firewall ", err)
	}

	missingChains := fw.detectMissingChains()
	log.Print("Firewall: creating chains")
	err = fw.CreateChains(missingChains)
	if err != nil {
		return err
	}
	for chain := range missingChains {
		if err := fw.DivertTrafficToPaniIptablesChain(chain); err != nil {
			return err
		}
		if err := fw.CreateRules(chain); err != nil {
			return err
		}
		if err := fw.CreateU32Rules(chain); err != nil {
			return err
		}
	}
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2017 The go-interpreter Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package exec provides functions for executing WebAssembly bytecode.
package exec

import (
	"encoding/binary"
	"errors"
	"fmt"
	"math"

	"github.com/go-interpreter/wagon/disasm"
	"github.com/go-interpreter/wagon/exec/internal/compile"
	"github.com/go-interpreter/wagon/wasm"
	ops "github.com/go-interpreter/wagon/wasm/operators"
)

var (
	// ErrMultipleLinearMemories is returned by (*VM).NewVM when the module
	// has more then one entries in the linear memory space.
	ErrMultipleLinearMemories = errors.New("exec: more than one linear memories in module")
	// ErrInvalidArgumentCount is returned by (*VM).ExecCode when an invalid
	// number of arguments to the WebAssembly function are passed to it.
	ErrInvalidArgumentCount = errors.New("exec: invalid number of arguments to function")
)

// InvalidReturnTypeError is returned by (*VM).ExecCode when the module
// specifies an invalid return type value for the executed function.
type InvalidReturnTypeError int8

func (e InvalidReturnTypeError) Error() string {
	return fmt.Sprintf("Function has invalid return value_type: %d", int8(e))
}

// InvalidFunctionIndexError is returned by (*VM).ExecCode when the function
// index provided is invalid.
type InvalidFunctionIndexError int64

func (e InvalidFunctionIndexError) Error() string {
	return fmt.Sprintf("Invalid index to function index space: %d", int64(e))
}

type context struct {
	stack   []uint64
	locals  []uint64
	code    []byte
	pc      int64
	curFunc int64
}

// VM is the execution context for executing WebAssembly bytecode.
type VM struct {
	ctx context

	module        *wasm.Module
	globals       []uint64
	memory        []byte
	compiledFuncs []compiledFunction

	funcTable [256]func()
}

// As per the WebAssembly spec: https://github.com/WebAssembly/design/blob/27ac254c854994103c24834a994be16f74f54186/Semantics.md#linear-memory
const wasmPageSize = 65536 // (64 KB)

var endianess = binary.LittleEndian

// NewVM creates a new VM from a given module. If the module defines a
// start function, it will be executed.
func NewVM(module *wasm.Module) (*VM, error) {
	var vm VM

	if module.Memory != nil && len(module.Memory.Entries) != 0 {
		if len(module.Memory.Entries) > 1 {
			return nil, ErrMultipleLinearMemories
		}
		vm.memory = make([]byte, uint(module.Memory.Entries[0].Limits.Initial)*wasmPageSize)
		copy(vm.memory, module.LinearMemoryIndexSpace[0])
	}

	vm.compiledFuncs = make([]compiledFunction, len(module.FunctionIndexSpace))
	vm.globals = make([]uint64, len(module.GlobalIndexSpace))
	vm.newFuncTable()
	vm.module = module

	for i, fn := range module.FunctionIndexSpace {
		disassembly, err := disasm.Disassemble(fn, module)
		if err != nil {
			return nil, err
		}

		totalLocalVars := 0
		totalLocalVars += len(fn.Sig.ParamTypes)
		for _, entry := range fn.Body.Locals {
			totalLocalVars += int(entry.Count)
		}
		code, table := compile.Compile(disassembly.Code)
		vm.compiledFuncs[i] = compiledFunction{
			code:           code,
			branchTables:   table,
			maxDepth:       disassembly.MaxDepth,
			totalLocalVars: totalLocalVars,
			args:           len(fn.Sig.ParamTypes),
			returns:        len(fn.Sig.ReturnTypes) != 0,
		}
	}

	for i, global := range module.GlobalIndexSpace {
		val, err := module.ExecInitExpr(global.Init)
		if err != nil {
			return nil, err
		}
		switch v := val.(type) {
		case int32:
			vm.globals[i] = uint64(v)
		case int64:
			vm.globals[i] = uint64(v)
		case float32:
			vm.globals[i] = uint64(math.Float32bits(v))
		case float64:
			vm.globals[i] = uint64(math.Float64bits(v))
		}
	}

	if module.Start != nil {
		_, err := vm.ExecCode(int64(module.Start.Index))
		if err != nil {
			return nil, err
		}
	}

	return &vm, nil
}

func (vm *VM) pushBool(v bool) {
	if v {
		vm.pushUint64(1)
	} else {
		vm.pushUint64(0)
	}
}

func (vm *VM) fetchBool() bool {
	return vm.fetchInt8() != 0
}

func (vm *VM) fetchInt8() int8 {
	i := int8(vm.ctx.code[vm.ctx.pc])
	vm.ctx.pc++
	return i
}

func (vm *VM) fetchUint32() uint32 {
	v := endianess.Uint32(vm.ctx.code[vm.ctx.pc:])
	vm.ctx.pc += 4
	return v
}

func (vm *VM) fetchInt32() int32 {
	return int32(vm.fetchUint32())
}

func (vm *VM) fetchFloat32() float32 {
	return math.Float32frombits(vm.fetchUint32())
}

func (vm *VM) fetchUint64() uint64 {
	v := endianess.Uint64(vm.ctx.code[vm.ctx.pc:])
	vm.ctx.pc += 8
	return v
}

func (vm *VM) fetchInt64() int64 {
	return int64(vm.fetchUint64())
}

func (vm *VM) fetchFloat64() float64 {
	return math.Float64frombits(vm.fetchUint64())
}

func (vm *VM) popUint64() uint64 {
	i := vm.ctx.stack[len(vm.ctx.stack)-1]
	vm.ctx.stack = vm.ctx.stack[:len(vm.ctx.stack)-1]
	return i
}

func (vm *VM) popInt64() int64 {
	return int64(vm.popUint64())
}

func (vm *VM) popFloat64() float64 {
	return math.Float64frombits(vm.popUint64())
}

func (vm *VM) popUint32() uint32 {
	return uint32(vm.popUint64())
}

func (vm *VM) popInt32() int32 {
	return int32(vm.popUint32())
}

func (vm *VM) popFloat32() float32 {
	return math.Float32frombits(vm.popUint32())
}

func (vm *VM) pushUint64(i uint64) {
	vm.ctx.stack = append(vm.ctx.stack, i)
}

func (vm *VM) pushInt64(i int64) {
	vm.pushUint64(uint64(i))
}

func (vm *VM) pushFloat64(f float64) {
	vm.pushUint64(math.Float64bits(f))
}

func (vm *VM) pushUint32(i uint32) {
	vm.pushUint64(uint64(i))
}

func (vm *VM) pushInt32(i int32) {
	vm.pushUint64(uint64(i))
}

func (vm *VM) pushFloat32(f float32) {
	vm.pushUint32(math.Float32bits(f))
}

// ExecCode calls the function with the given index and arguments.
// fnIndex should be a valid index into the function index space of
// the VM's module.
func (vm *VM) ExecCode(fnIndex int64, args ...uint64) (interface{}, error) {
	if int(fnIndex) > len(vm.compiledFuncs) {
		return nil, InvalidFunctionIndexError(fnIndex)
	}
	if len(vm.module.GetFunction(int(fnIndex)).Sig.ParamTypes) != len(args) {
		return nil, ErrInvalidArgumentCount
	}
	compiled := vm.compiledFuncs[fnIndex]
	if len(vm.ctx.stack) < compiled.maxDepth {
		vm.ctx.stack = make([]uint64, compiled.maxDepth)
	}
	vm.ctx.locals = make([]uint64, compiled.totalLocalVars)
	vm.ctx.pc = 0
	vm.ctx.code = compiled.code

	for i, arg := range args {
		vm.ctx.locals[i] = arg
	}

	var rtrn interface{}
	res := vm.execCode(compiled)
	if compiled.returns {
		rtrnType := vm.module.GetFunction(int(fnIndex)).Sig.ReturnTypes[0]
		switch rtrnType {
		case wasm.ValueTypeI32:
			rtrn = uint32(res)
		case wasm.ValueTypeI64:
			rtrn = uint64(res)
		case wasm.ValueTypeF32:
			rtrn = math.Float32frombits(uint32(res))
		case wasm.ValueTypeF64:
			rtrn = math.Float64frombits(res)
		default:
			return nil, InvalidReturnTypeError(rtrnType)
		}
	}

	return rtrn, nil
}

func (vm *VM) execCode(compiled compiledFunction) uint64 {
outer:
	for int(vm.ctx.pc) < len(vm.ctx.code) {
		op := vm.ctx.code[vm.ctx.pc]
		vm.ctx.pc++
		switch op {
		case ops.Return:
			break outer
		case compile.OpJmp:
			vm.ctx.pc = vm.fetchInt64()
			continue
		case compile.OpJmpZ:
			target := vm.fetchInt64()
			if vm.popUint32() == 0 {
				vm.ctx.pc = target
				continue
			}
		case compile.OpJmpNz:
			target := vm.fetchInt64()
			preserveTop := vm.fetchBool()
			discard := vm.fetchInt64()
			if vm.popUint32() != 0 {
				vm.ctx.pc = target
				var top uint64
				if preserveTop {
					top = vm.ctx.stack[len(vm.ctx.stack)-1]
				}
				vm.ctx.stack = vm.ctx.stack[:len(vm.ctx.stack)-int(discard)]
				if preserveTop {
					vm.pushUint64(top)
				}
				continue
			}
		case ops.BrTable:
			// println("brtable")
			index := vm.fetchInt64()
			label := vm.popInt32()
			table := vm.compiledFuncs[vm.ctx.curFunc].branchTables[index]
			var target compile.Target
			if label < int32(len(table.Targets)) {
				target = table.Targets[int32(label)]
			} else {
				target = table.DefaultTarget
			}
			vm.ctx.pc = target.Addr
			var top uint64
			if target.PreserveTop {
				top = vm.ctx.stack[len(vm.ctx.stack)-1]
			}
			vm.ctx.stack = vm.ctx.stack[:len(vm.ctx.stack)-int(target.Discard)]
			if target.PreserveTop {
				vm.pushUint64(top)
			}
			continue
		case compile.OpDiscard:
			place := vm.fetchInt64()
			vm.ctx.stack = vm.ctx.stack[:len(vm.ctx.stack)-int(place)]
		case compile.OpDiscardPreserveTop:
			top := vm.ctx.stack[len(vm.ctx.stack)-1]
			place := vm.fetchInt64()
			vm.ctx.stack = vm.ctx.stack[:len(vm.ctx.stack)-int(place)]
			vm.pushUint64(top)
		default:
			vm.funcTable[op]()
		}
	}

	if compiled.returns {
		return vm.ctx.stack[len(vm.ctx.stack)-1]
	}
	return 0
}
<commit_msg>exec: Set length of stack to 0, set curFunc, let br_table to return<commit_after>// Copyright 2017 The go-interpreter Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package exec provides functions for executing WebAssembly bytecode.
package exec

import (
	"encoding/binary"
	"errors"
	"fmt"
	"math"

	"github.com/go-interpreter/wagon/disasm"
	"github.com/go-interpreter/wagon/exec/internal/compile"
	"github.com/go-interpreter/wagon/wasm"
	ops "github.com/go-interpreter/wagon/wasm/operators"
)

var (
	// ErrMultipleLinearMemories is returned by (*VM).NewVM when the module
	// has more then one entries in the linear memory space.
	ErrMultipleLinearMemories = errors.New("exec: more than one linear memories in module")
	// ErrInvalidArgumentCount is returned by (*VM).ExecCode when an invalid
	// number of arguments to the WebAssembly function are passed to it.
	ErrInvalidArgumentCount = errors.New("exec: invalid number of arguments to function")
)

// InvalidReturnTypeError is returned by (*VM).ExecCode when the module
// specifies an invalid return type value for the executed function.
type InvalidReturnTypeError int8

func (e InvalidReturnTypeError) Error() string {
	return fmt.Sprintf("Function has invalid return value_type: %d", int8(e))
}

// InvalidFunctionIndexError is returned by (*VM).ExecCode when the function
// index provided is invalid.
type InvalidFunctionIndexError int64

func (e InvalidFunctionIndexError) Error() string {
	return fmt.Sprintf("Invalid index to function index space: %d", int64(e))
}

type context struct {
	stack   []uint64
	locals  []uint64
	code    []byte
	pc      int64
	curFunc int64
}

// VM is the execution context for executing WebAssembly bytecode.
type VM struct {
	ctx context

	module        *wasm.Module
	globals       []uint64
	memory        []byte
	compiledFuncs []compiledFunction

	funcTable [256]func()
}

// As per the WebAssembly spec: https://github.com/WebAssembly/design/blob/27ac254c854994103c24834a994be16f74f54186/Semantics.md#linear-memory
const wasmPageSize = 65536 // (64 KB)

var endianess = binary.LittleEndian

// NewVM creates a new VM from a given module. If the module defines a
// start function, it will be executed.
func NewVM(module *wasm.Module) (*VM, error) {
	var vm VM

	if module.Memory != nil && len(module.Memory.Entries) != 0 {
		if len(module.Memory.Entries) > 1 {
			return nil, ErrMultipleLinearMemories
		}
		vm.memory = make([]byte, uint(module.Memory.Entries[0].Limits.Initial)*wasmPageSize)
		copy(vm.memory, module.LinearMemoryIndexSpace[0])
	}

	vm.compiledFuncs = make([]compiledFunction, len(module.FunctionIndexSpace))
	vm.globals = make([]uint64, len(module.GlobalIndexSpace))
	vm.newFuncTable()
	vm.module = module

	for i, fn := range module.FunctionIndexSpace {
		disassembly, err := disasm.Disassemble(fn, module)
		if err != nil {
			return nil, err
		}

		totalLocalVars := 0
		totalLocalVars += len(fn.Sig.ParamTypes)
		for _, entry := range fn.Body.Locals {
			totalLocalVars += int(entry.Count)
		}
		code, table := compile.Compile(disassembly.Code)
		vm.compiledFuncs[i] = compiledFunction{
			code:           code,
			branchTables:   table,
			maxDepth:       disassembly.MaxDepth,
			totalLocalVars: totalLocalVars,
			args:           len(fn.Sig.ParamTypes),
			returns:        len(fn.Sig.ReturnTypes) != 0,
		}
	}

	for i, global := range module.GlobalIndexSpace {
		val, err := module.ExecInitExpr(global.Init)
		if err != nil {
			return nil, err
		}
		switch v := val.(type) {
		case int32:
			vm.globals[i] = uint64(v)
		case int64:
			vm.globals[i] = uint64(v)
		case float32:
			vm.globals[i] = uint64(math.Float32bits(v))
		case float64:
			vm.globals[i] = uint64(math.Float64bits(v))
		}
	}

	if module.Start != nil {
		_, err := vm.ExecCode(int64(module.Start.Index))
		if err != nil {
			return nil, err
		}
	}

	return &vm, nil
}

func (vm *VM) pushBool(v bool) {
	if v {
		vm.pushUint64(1)
	} else {
		vm.pushUint64(0)
	}
}

func (vm *VM) fetchBool() bool {
	return vm.fetchInt8() != 0
}

func (vm *VM) fetchInt8() int8 {
	i := int8(vm.ctx.code[vm.ctx.pc])
	vm.ctx.pc++
	return i
}

func (vm *VM) fetchUint32() uint32 {
	v := endianess.Uint32(vm.ctx.code[vm.ctx.pc:])
	vm.ctx.pc += 4
	return v
}

func (vm *VM) fetchInt32() int32 {
	return int32(vm.fetchUint32())
}

func (vm *VM) fetchFloat32() float32 {
	return math.Float32frombits(vm.fetchUint32())
}

func (vm *VM) fetchUint64() uint64 {
	v := endianess.Uint64(vm.ctx.code[vm.ctx.pc:])
	vm.ctx.pc += 8
	return v
}

func (vm *VM) fetchInt64() int64 {
	return int64(vm.fetchUint64())
}

func (vm *VM) fetchFloat64() float64 {
	return math.Float64frombits(vm.fetchUint64())
}

func (vm *VM) popUint64() uint64 {
	i := vm.ctx.stack[len(vm.ctx.stack)-1]
	vm.ctx.stack = vm.ctx.stack[:len(vm.ctx.stack)-1]
	return i
}

func (vm *VM) popInt64() int64 {
	return int64(vm.popUint64())
}

func (vm *VM) popFloat64() float64 {
	return math.Float64frombits(vm.popUint64())
}

func (vm *VM) popUint32() uint32 {
	return uint32(vm.popUint64())
}

func (vm *VM) popInt32() int32 {
	return int32(vm.popUint32())
}

func (vm *VM) popFloat32() float32 {
	return math.Float32frombits(vm.popUint32())
}

func (vm *VM) pushUint64(i uint64) {
	vm.ctx.stack = append(vm.ctx.stack, i)
}

func (vm *VM) pushInt64(i int64) {
	vm.pushUint64(uint64(i))
}

func (vm *VM) pushFloat64(f float64) {
	vm.pushUint64(math.Float64bits(f))
}

func (vm *VM) pushUint32(i uint32) {
	vm.pushUint64(uint64(i))
}

func (vm *VM) pushInt32(i int32) {
	vm.pushUint64(uint64(i))
}

func (vm *VM) pushFloat32(f float32) {
	vm.pushUint32(math.Float32bits(f))
}

// ExecCode calls the function with the given index and arguments.
// fnIndex should be a valid index into the function index space of
// the VM's module.
func (vm *VM) ExecCode(fnIndex int64, args ...uint64) (interface{}, error) {
	if int(fnIndex) > len(vm.compiledFuncs) {
		return nil, InvalidFunctionIndexError(fnIndex)
	}
	if len(vm.module.GetFunction(int(fnIndex)).Sig.ParamTypes) != len(args) {
		return nil, ErrInvalidArgumentCount
	}
	compiled := vm.compiledFuncs[fnIndex]
	if len(vm.ctx.stack) < compiled.maxDepth {
		vm.ctx.stack = make([]uint64, 0, compiled.maxDepth)
	}
	vm.ctx.locals = make([]uint64, compiled.totalLocalVars)
	vm.ctx.pc = 0
	vm.ctx.code = compiled.code
	vm.ctx.curFunc = fnIndex

	for i, arg := range args {
		vm.ctx.locals[i] = arg
	}

	var rtrn interface{}
	res := vm.execCode(compiled)
	if compiled.returns {
		rtrnType := vm.module.GetFunction(int(fnIndex)).Sig.ReturnTypes[0]
		switch rtrnType {
		case wasm.ValueTypeI32:
			rtrn = uint32(res)
		case wasm.ValueTypeI64:
			rtrn = uint64(res)
		case wasm.ValueTypeF32:
			rtrn = math.Float32frombits(uint32(res))
		case wasm.ValueTypeF64:
			rtrn = math.Float64frombits(res)
		default:
			return nil, InvalidReturnTypeError(rtrnType)
		}
	}

	return rtrn, nil
}

func (vm *VM) execCode(compiled compiledFunction) uint64 {
outer:
	for int(vm.ctx.pc) < len(vm.ctx.code) {
		op := vm.ctx.code[vm.ctx.pc]
		vm.ctx.pc++
		switch op {
		case ops.Return:
			break outer
		case compile.OpJmp:
			vm.ctx.pc = vm.fetchInt64()
			continue
		case compile.OpJmpZ:
			target := vm.fetchInt64()
			if vm.popUint32() == 0 {
				vm.ctx.pc = target
				continue
			}
		case compile.OpJmpNz:
			target := vm.fetchInt64()
			preserveTop := vm.fetchBool()
			discard := vm.fetchInt64()
			if vm.popUint32() != 0 {
				vm.ctx.pc = target
				var top uint64
				if preserveTop {
					top = vm.ctx.stack[len(vm.ctx.stack)-1]
				}
				vm.ctx.stack = vm.ctx.stack[:len(vm.ctx.stack)-int(discard)]
				if preserveTop {
					vm.pushUint64(top)
				}
				continue
			}
		case ops.BrTable:
			index := vm.fetchInt64()
			label := vm.popInt32()
			table := vm.compiledFuncs[vm.ctx.curFunc].branchTables[index]
			var target compile.Target
			if label >= 0 && label < int32(len(table.Targets)) {
				target = table.Targets[int32(label)]
			} else {
				target = table.DefaultTarget
			}

			if target.Return {
				break outer
			}
			vm.ctx.pc = target.Addr
			var top uint64
			if target.PreserveTop {
				top = vm.ctx.stack[len(vm.ctx.stack)-1]
			}
			vm.ctx.stack = vm.ctx.stack[:len(vm.ctx.stack)-int(target.Discard)]
			if target.PreserveTop {
				vm.pushUint64(top)
			}
			continue
		case compile.OpDiscard:
			place := vm.fetchInt64()
			vm.ctx.stack = vm.ctx.stack[:len(vm.ctx.stack)-int(place)]
		case compile.OpDiscardPreserveTop:
			top := vm.ctx.stack[len(vm.ctx.stack)-1]
			place := vm.fetchInt64()
			vm.ctx.stack = vm.ctx.stack[:len(vm.ctx.stack)-int(place)]
			vm.pushUint64(top)
		default:
			vm.funcTable[op]()
		}
	}

	if compiled.returns {
		return vm.ctx.stack[len(vm.ctx.stack)-1]
	}
	return 0
}
<|endoftext|>


<commit_before>// +build !confonly

package http

import (
	"bufio"
	"context"
	"encoding/base64"
	"io"
	"net/http"
	"net/url"
	"sync"

	"golang.org/x/net/http2"

	"v2ray.com/core"
	"v2ray.com/core/common"
	"v2ray.com/core/common/buf"
	"v2ray.com/core/common/net"
	"v2ray.com/core/common/protocol"
	"v2ray.com/core/common/retry"
	"v2ray.com/core/common/session"
	"v2ray.com/core/common/signal"
	"v2ray.com/core/common/task"
	"v2ray.com/core/features/policy"
	"v2ray.com/core/transport"
	"v2ray.com/core/transport/internet"
	"v2ray.com/core/transport/internet/tls"
)

type Client struct {
	serverPicker  protocol.ServerPicker
	policyManager policy.Manager
}

type h2Conn struct {
	rawConn net.Conn
	h2Conn  *http2.ClientConn
}

var (
	cachedH2Mutex sync.Mutex
	cachedH2Conns map[net.Destination]h2Conn
)

// NewClient create a new http client based on the given config.
func NewClient(ctx context.Context, config *ClientConfig) (*Client, error) {
	serverList := protocol.NewServerList()
	for _, rec := range config.Server {
		s, err := protocol.NewServerSpecFromPB(*rec)
		if err != nil {
			return nil, newError("failed to get server spec").Base(err)
		}
		serverList.AddServer(s)
	}
	if serverList.Size() == 0 {
		return nil, newError("0 target server")
	}

	v := core.MustFromContext(ctx)
	return &Client{
		serverPicker:  protocol.NewRoundRobinServerPicker(serverList),
		policyManager: v.GetFeature(policy.ManagerType()).(policy.Manager),
	}, nil
}

// Process implements proxy.Outbound.Process. We first create a socket tunnel via HTTP CONNECT method, then redirect all inbound traffic to that tunnel.
func (c *Client) Process(ctx context.Context, link *transport.Link, dialer internet.Dialer) error {
	outbound := session.OutboundFromContext(ctx)
	if outbound == nil || !outbound.Target.IsValid() {
		return newError("target not specified.")
	}
	target := outbound.Target
	targetAddr := target.NetAddr()

	if target.Network == net.Network_UDP {
		return newError("UDP is not supported by HTTP outbound")
	}

	var user *protocol.MemoryUser
	var conn internet.Connection

	if err := retry.ExponentialBackoff(5, 100).On(func() error {
		server := c.serverPicker.PickServer()
		dest := server.Destination()
		user = server.PickUser()

		netConn, err := setUpHTTPTunnel(ctx, dest, targetAddr, user, dialer)
		if netConn != nil {
			conn = internet.Connection(netConn)
		}
		return err
	}); err != nil {
		return newError("failed to find an available destination").Base(err)
	}

	defer func() {
		if err := conn.Close(); err != nil {
			newError("failed to closed connection").Base(err).WriteToLog(session.ExportIDToError(ctx))
		}
	}()

	p := c.policyManager.ForLevel(0)
	if user != nil {
		p = c.policyManager.ForLevel(user.Level)
	}

	ctx, cancel := context.WithCancel(ctx)
	timer := signal.CancelAfterInactivity(ctx, cancel, p.Timeouts.ConnectionIdle)

	requestFunc := func() error {
		defer timer.SetTimeout(p.Timeouts.DownlinkOnly)
		return buf.Copy(link.Reader, buf.NewWriter(conn), buf.UpdateActivity(timer))
	}
	responseFunc := func() error {
		defer timer.SetTimeout(p.Timeouts.UplinkOnly)
		return buf.Copy(buf.NewReader(conn), link.Writer, buf.UpdateActivity(timer))
	}

	var responseDonePost = task.OnSuccess(responseFunc, task.Close(link.Writer))
	if err := task.Run(ctx, requestFunc, responseDonePost); err != nil {
		return newError("connection ends").Base(err)
	}

	return nil
}

// setUpHTTPTunnel will create a socket tunnel via HTTP CONNECT method
func setUpHTTPTunnel(ctx context.Context, dest net.Destination, target string, user *protocol.MemoryUser, dialer internet.Dialer) (net.Conn, error) {
	req := &http.Request{
		Method: http.MethodConnect,
		URL:    &url.URL{Host: target},
		Header: http.Header{"Proxy-Connection": []string{"Keep-Alive"}},
		Host:   target,
	}

	if user != nil && user.Account != nil {
		account := user.Account.(*Account)
		auth := account.GetUsername() + ":" + account.GetPassword()
		req.Header.Set("Proxy-Authorization", "Basic "+base64.StdEncoding.EncodeToString([]byte(auth)))
	}

	connectHTTP1 := func(rawConn net.Conn) (net.Conn, error) {
		err := req.Write(rawConn)
		if err != nil {
			rawConn.Close()
			return nil, err
		}

		resp, err := http.ReadResponse(bufio.NewReader(rawConn), req)
		if err != nil {
			rawConn.Close()
			return nil, err
		}

		if resp.StatusCode != http.StatusOK {
			rawConn.Close()
			return nil, newError("Proxy responded with non 200 code: " + resp.Status)
		}
		return rawConn, nil
	}

	connectHTTP2 := func(rawConn net.Conn, h2clientConn *http2.ClientConn) (net.Conn, error) {
		pr, pw := io.Pipe()
		req.Body = pr

		resp, err := h2clientConn.RoundTrip(req)
		if err != nil {
			rawConn.Close()
			return nil, err
		}

		if resp.StatusCode != http.StatusOK {
			rawConn.Close()
			return nil, newError("Proxy responded with non 200 code: " + resp.Status)
		}
		return newHTTP2Conn(rawConn, pw, resp.Body), nil
	}

	cachedH2Mutex.Lock()
	defer cachedH2Mutex.Unlock()

	if cachedConn, found := cachedH2Conns[dest]; found {
		rc, cc := cachedConn.rawConn, cachedConn.h2Conn
		if cc.CanTakeNewRequest() {
			proxyConn, err := connectHTTP2(rc, cc)
			if err != nil {
				return nil, err
			}

			return proxyConn, nil
		}
	}

	rawConn, err := dialer.Dial(ctx, dest)
	if err != nil {
		return nil, err
	}

	nextProto := ""
	if tlsConn, ok := rawConn.(*tls.Conn); ok {
		if err := tlsConn.Handshake(); err != nil {
			rawConn.Close()
			return nil, err
		}
		nextProto = tlsConn.ConnectionState().NegotiatedProtocol
	}

	switch nextProto {
	case "", "http/1.1":
		return connectHTTP1(rawConn)
	case "h2":
		t := http2.Transport{}
		h2clientConn, err := t.NewClientConn(rawConn)
		if err != nil {
			rawConn.Close()
			return nil, err
		}

		proxyConn, err := connectHTTP2(rawConn, h2clientConn)
		if err != nil {
			rawConn.Close()
			return nil, err
		}

		if cachedH2Conns == nil {
			cachedH2Conns = make(map[net.Destination]h2Conn)
		}

		cachedH2Conns[dest] = h2Conn{
			rawConn: rawConn,
			h2Conn:  h2clientConn,
		}

		return proxyConn, err
	default:
		return nil, newError("negotiated unsupported application layer protocol: " + nextProto)
	}
}

func newHTTP2Conn(c net.Conn, pipedReqBody *io.PipeWriter, respBody io.ReadCloser) net.Conn {
	return &http2Conn{Conn: c, in: pipedReqBody, out: respBody}
}

type http2Conn struct {
	net.Conn
	in  *io.PipeWriter
	out io.ReadCloser
}

func (h *http2Conn) Read(p []byte) (n int, err error) {
	return h.out.Read(p)
}

func (h *http2Conn) Write(p []byte) (n int, err error) {
	return h.in.Write(p)
}

func (h *http2Conn) Close() error {
	h.in.Close()
	return h.out.Close()
}

func init() {
	common.Must(common.RegisterConfig((*ClientConfig)(nil), func(ctx context.Context, config interface{}) (interface{}, error) {
		return NewClient(ctx, config.(*ClientConfig))
	}))
}
<commit_msg>Fix http outbound<commit_after>// +build !confonly

package http

import (
	"bufio"
	"context"
	"encoding/base64"
	"io"
	"net/http"
	"net/url"
	"sync"

	"golang.org/x/net/http2"

	"v2ray.com/core"
	"v2ray.com/core/common"
	"v2ray.com/core/common/buf"
	"v2ray.com/core/common/net"
	"v2ray.com/core/common/protocol"
	"v2ray.com/core/common/retry"
	"v2ray.com/core/common/session"
	"v2ray.com/core/common/signal"
	"v2ray.com/core/common/task"
	"v2ray.com/core/features/policy"
	"v2ray.com/core/transport"
	"v2ray.com/core/transport/internet"
	"v2ray.com/core/transport/internet/tls"
)

type Client struct {
	serverPicker  protocol.ServerPicker
	policyManager policy.Manager
}

type h2Conn struct {
	rawConn net.Conn
	h2Conn  *http2.ClientConn
}

var (
	cachedH2Mutex sync.Mutex
	cachedH2Conns map[net.Destination]h2Conn
)

// NewClient create a new http client based on the given config.
func NewClient(ctx context.Context, config *ClientConfig) (*Client, error) {
	serverList := protocol.NewServerList()
	for _, rec := range config.Server {
		s, err := protocol.NewServerSpecFromPB(*rec)
		if err != nil {
			return nil, newError("failed to get server spec").Base(err)
		}
		serverList.AddServer(s)
	}
	if serverList.Size() == 0 {
		return nil, newError("0 target server")
	}

	v := core.MustFromContext(ctx)
	return &Client{
		serverPicker:  protocol.NewRoundRobinServerPicker(serverList),
		policyManager: v.GetFeature(policy.ManagerType()).(policy.Manager),
	}, nil
}

// Process implements proxy.Outbound.Process. We first create a socket tunnel via HTTP CONNECT method, then redirect all inbound traffic to that tunnel.
func (c *Client) Process(ctx context.Context, link *transport.Link, dialer internet.Dialer) error {
	outbound := session.OutboundFromContext(ctx)
	if outbound == nil || !outbound.Target.IsValid() {
		return newError("target not specified.")
	}
	target := outbound.Target
	targetAddr := target.NetAddr()

	if target.Network == net.Network_UDP {
		return newError("UDP is not supported by HTTP outbound")
	}

	var user *protocol.MemoryUser
	var conn internet.Connection

	if err := retry.ExponentialBackoff(5, 100).On(func() error {
		server := c.serverPicker.PickServer()
		dest := server.Destination()
		user = server.PickUser()

		netConn, err := setUpHTTPTunnel(ctx, dest, targetAddr, user, dialer)
		if netConn != nil {
			conn = internet.Connection(netConn)
		}
		return err
	}); err != nil {
		return newError("failed to find an available destination").Base(err)
	}

	defer func() {
		if err := conn.Close(); err != nil {
			newError("failed to closed connection").Base(err).WriteToLog(session.ExportIDToError(ctx))
		}
	}()

	p := c.policyManager.ForLevel(0)
	if user != nil {
		p = c.policyManager.ForLevel(user.Level)
	}

	ctx, cancel := context.WithCancel(ctx)
	timer := signal.CancelAfterInactivity(ctx, cancel, p.Timeouts.ConnectionIdle)

	requestFunc := func() error {
		defer timer.SetTimeout(p.Timeouts.DownlinkOnly)
		return buf.Copy(link.Reader, buf.NewWriter(conn), buf.UpdateActivity(timer))
	}
	responseFunc := func() error {
		defer timer.SetTimeout(p.Timeouts.UplinkOnly)
		return buf.Copy(buf.NewReader(conn), link.Writer, buf.UpdateActivity(timer))
	}

	var responseDonePost = task.OnSuccess(responseFunc, task.Close(link.Writer))
	if err := task.Run(ctx, requestFunc, responseDonePost); err != nil {
		return newError("connection ends").Base(err)
	}

	return nil
}

// setUpHTTPTunnel will create a socket tunnel via HTTP CONNECT method
func setUpHTTPTunnel(ctx context.Context, dest net.Destination, target string, user *protocol.MemoryUser, dialer internet.Dialer) (net.Conn, error) {
	req := &http.Request{
		Method: http.MethodConnect,
		URL:    &url.URL{Host: target},
		Header: http.Header{"Proxy-Connection": []string{"Keep-Alive"}},
		Host:   target,
	}

	if user != nil && user.Account != nil {
		account := user.Account.(*Account)
		auth := account.GetUsername() + ":" + account.GetPassword()
		req.Header.Set("Proxy-Authorization", "Basic "+base64.StdEncoding.EncodeToString([]byte(auth)))
	}

	connectHTTP1 := func(rawConn net.Conn) (net.Conn, error) {
		err := req.Write(rawConn)
		if err != nil {
			rawConn.Close()
			return nil, err
		}

		resp, err := http.ReadResponse(bufio.NewReader(rawConn), req)
		if err != nil {
			rawConn.Close()
			return nil, err
		}

		if resp.StatusCode != http.StatusOK {
			rawConn.Close()
			return nil, newError("Proxy responded with non 200 code: " + resp.Status)
		}
		return rawConn, nil
	}

	connectHTTP2 := func(rawConn net.Conn, h2clientConn *http2.ClientConn) (net.Conn, error) {
		pr, pw := io.Pipe()
		req.Body = pr

		resp, err := h2clientConn.RoundTrip(req)
		if err != nil {
			rawConn.Close()
			return nil, err
		}

		if resp.StatusCode != http.StatusOK {
			rawConn.Close()
			return nil, newError("Proxy responded with non 200 code: " + resp.Status)
		}
		return newHTTP2Conn(rawConn, pw, resp.Body), nil
	}

	cachedH2Mutex.Lock()
	defer cachedH2Mutex.Unlock()

	if cachedConn, found := cachedH2Conns[dest]; found {
		rc, cc := cachedConn.rawConn, cachedConn.h2Conn
		if cc.CanTakeNewRequest() {
			proxyConn, err := connectHTTP2(rc, cc)
			if err != nil {
				return nil, err
			}

			return proxyConn, nil
		}
	}

	rawConn, err := dialer.Dial(ctx, dest)
	if err != nil {
		return nil, err
	}

	iConn := rawConn
	if statConn, ok := iConn.(*internet.StatCouterConnection); ok {
		iConn = statConn.Connection
	}

	nextProto := ""
	if tlsConn, ok := iConn.(*tls.Conn); ok {
		if err := tlsConn.Handshake(); err != nil {
			rawConn.Close()
			return nil, err
		}
		nextProto = tlsConn.ConnectionState().NegotiatedProtocol
	}

	switch nextProto {
	case "", "http/1.1":
		return connectHTTP1(rawConn)
	case "h2":
		t := http2.Transport{}
		h2clientConn, err := t.NewClientConn(rawConn)
		if err != nil {
			rawConn.Close()
			return nil, err
		}

		proxyConn, err := connectHTTP2(rawConn, h2clientConn)
		if err != nil {
			rawConn.Close()
			return nil, err
		}

		if cachedH2Conns == nil {
			cachedH2Conns = make(map[net.Destination]h2Conn)
		}

		cachedH2Conns[dest] = h2Conn{
			rawConn: rawConn,
			h2Conn:  h2clientConn,
		}

		return proxyConn, err
	default:
		return nil, newError("negotiated unsupported application layer protocol: " + nextProto)
	}
}

func newHTTP2Conn(c net.Conn, pipedReqBody *io.PipeWriter, respBody io.ReadCloser) net.Conn {
	return &http2Conn{Conn: c, in: pipedReqBody, out: respBody}
}

type http2Conn struct {
	net.Conn
	in  *io.PipeWriter
	out io.ReadCloser
}

func (h *http2Conn) Read(p []byte) (n int, err error) {
	return h.out.Read(p)
}

func (h *http2Conn) Write(p []byte) (n int, err error) {
	return h.in.Write(p)
}

func (h *http2Conn) Close() error {
	h.in.Close()
	return h.out.Close()
}

func init() {
	common.Must(common.RegisterConfig((*ClientConfig)(nil), func(ctx context.Context, config interface{}) (interface{}, error) {
		return NewClient(ctx, config.(*ClientConfig))
	}))
}
<|endoftext|>