<commit_before>package marshal

import (
	//"errors"
	"reflect"
	"strings"
	"fmt"

	"github.com/xitongsys/parquet-go/common"
	"github.com/xitongsys/parquet-go/layout"
	"github.com/xitongsys/parquet-go/schema"
	"github.com/xitongsys/parquet-go/types"
)

//Record Map KeyValue pair
type KeyValue struct {
	Key   reflect.Value
	Value reflect.Value
}

type MapRecord struct {
	KeyValues []KeyValue
	Index     int
}

//Convert the table map to objects slice. desInterface is a slice of pointers of objects
func Unmarshal(tableMap *map[string]*layout.Table, bgn int, end int, dstInterface interface{}, schemaHandler *schema.SchemaHandler, prefixPath string) (err error) {
	/*
	defer func() {
		if r := recover(); r != nil {
			switch x := r.(type) {
			case string:
				err = errors.New(x)
			case error:
				err = x
			default:
				err = errors.New("unknown error")
			}
		}
	}()
	*/

	tableNeeds := make(map[string]*layout.Table)
	tableBgn, tableEnd := make(map[string]int), make(map[string]int)
	for name, table := range *tableMap {
		if !strings.HasPrefix(name, prefixPath) {
			continue
		}

		tableNeeds[name] = table

		ln := len(table.Values)
		num := -1
		tableBgn[name], tableEnd[name] = -1, -1
		for i := 0; i < ln; i++ {
			if table.RepetitionLevels[i] == 0 {
				num++
				if num == bgn {
					tableBgn[name] = i
				}
				if num == end {
					tableEnd[name] = i
					break
				}
			}
		}

		if tableEnd[name] < 0 {
			tableEnd[name] = ln
		}
		if tableBgn[name] < 0 {
			return
		}
	}

	mapRecords := make(map[reflect.Value][]KeyValue)

	for name, table := range tableNeeds {
		path := table.Path
		bgn := tableBgn[name]
		end := tableEnd[name]
		schemaIndex := schemaHandler.MapIndex[common.PathToStr(path)]
		pT, cT := schemaHandler.SchemaElements[schemaIndex].Type, schemaHandler.SchemaElements[schemaIndex].ConvertedType

		repetitionLevels, definitionLevels := make([]int32, len(path)), make([]int32, len(path))
		for i := 0; i<len(path); i++ {
			repetitionLevels[i], _ = schemaHandler.MaxRepetitionLevel(path[:i+1])
			definitionLevels[i], _ = schemaHandler.MaxDefinitionLevel(path[:i+1])
		}
		repetitionIndexs := make([]int32, len(path))
		for i := 0; i < len(path); i++ {
			repetitionIndexs[i] = -1
		}

		fmt.Println("========", name)

		for i := bgn; i < end; i++ {
			rl, dl, val := table.RepetitionLevels[i], table.DefinitionLevels[i], table.Values[i]
			po, index := reflect.ValueOf(dstInterface).Elem(), 0

			for index < len(path) {
				if po.Type().Kind() == reflect.Slice {
					if po.IsNil() {
						po.Set(reflect.MakeSlice(po.Type(), 0, 0))
					}

					if rl == repetitionLevels[index] || repetitionIndexs[index] < 0 {
						repetitionIndexs[index]++
					}

					if repetitionIndexs[index] >= int32(po.Len()) {
						potmp := reflect.Append(po, reflect.New(po.Type().Elem()).Elem())
						po.Set(potmp)
					}

					fmt.Println("=====", po.Len(), repetitionIndexs[index])

					po = po.Index(int(repetitionIndexs[index]))

				} else if po.Type().Kind() == reflect.Map {
					if po.IsNil() {
						po.Set(reflect.MakeMap(po.Type()))
					}

					if _, ok := mapRecords[po]; !ok {
						mapRecords[po] = make([]KeyValue, 0)
					}

					index++
					if definitionLevels[index] > dl {
						break
					}

					if rl == repetitionLevels[index] {
						repetitionIndexs[index]++
					}

					if repetitionIndexs[index] >= int32(len(mapRecords[po])) {
						mapRecords[po] = append(mapRecords[po], KeyValue{})
					}

					if path[index + 1] == "key" {
						po = mapRecords[po][repetitionIndexs[index]].Key

					}else {
						po = mapRecords[po][repetitionIndexs[index]].Value
					}

					index++
					if definitionLevels[index] > dl {
						break
					}

				} else if po.Type().Kind() == reflect.Ptr {
					if po.IsNil() {
						po.Set(reflect.New(po.Type().Elem()))
					}

					po = po.Elem()

				} else if po.Type().Kind() == reflect.Struct {
					index++
					if definitionLevels[index] > dl {
						break;
					}
					po = po.FieldByName(path[index])

				} else {
					po.Set(reflect.ValueOf(types.ParquetTypeToGoType(val, pT, cT)))
					break
				}
			}
		}
	}

	for po, kvs := range mapRecords {
		for _, kv := range kvs {
			po.SetMapIndex(kv.Key, kv.Value)
		}
	}

	return nil
}
<commit_msg>debug<commit_after>package marshal

import (
	//"errors"
	"reflect"
	"strings"

	"github.com/xitongsys/parquet-go/common"
	"github.com/xitongsys/parquet-go/layout"
	"github.com/xitongsys/parquet-go/schema"
	"github.com/xitongsys/parquet-go/types"
)

//Record Map KeyValue pair
type KeyValue struct {
	Key   reflect.Value
	Value reflect.Value
}

type MapRecord struct {
	KeyValues []KeyValue
	Index     int
}

type SliceRecord struct {
	Values	[]reflect.Value
	Index	int
}

//Convert the table map to objects slice. desInterface is a slice of pointers of objects
func Unmarshal(tableMap *map[string]*layout.Table, bgn int, end int, dstInterface interface{}, schemaHandler *schema.SchemaHandler, prefixPath string) (err error) {
	/*
	defer func() {
		if r := recover(); r != nil {
			switch x := r.(type) {
			case string:
				err = errors.New(x)
			case error:
				err = x
			default:
				err = errors.New("unknown error")
			}
		}
	}()
	*/

	tableNeeds := make(map[string]*layout.Table)
	tableBgn, tableEnd := make(map[string]int), make(map[string]int)
	for name, table := range *tableMap {
		if !strings.HasPrefix(name, prefixPath) {
			continue
		}

		tableNeeds[name] = table

		ln := len(table.Values)
		num := -1
		tableBgn[name], tableEnd[name] = -1, -1
		for i := 0; i < ln; i++ {
			if table.RepetitionLevels[i] == 0 {
				num++
				if num == bgn {
					tableBgn[name] = i
				}
				if num == end {
					tableEnd[name] = i
					break
				}
			}
		}

		if tableEnd[name] < 0 {
			tableEnd[name] = ln
		}
		if tableBgn[name] < 0 {
			return
		}
	}

	mapRecords := make(map[reflect.Value]*MapRecord)
	mapRecordsStack := make([]reflect.Value, 0)
	sliceRecords := make(map[reflect.Value]*SliceRecord)
	sliceRecordsStack := make([]reflect.Value, 0)
	root := reflect.ValueOf(dstInterface).Elem()

	for name, table := range tableNeeds {
		path := table.Path
		bgn := tableBgn[name]
		end := tableEnd[name]
		schemaIndex := schemaHandler.MapIndex[common.PathToStr(path)]
		pT, cT := schemaHandler.SchemaElements[schemaIndex].Type, schemaHandler.SchemaElements[schemaIndex].ConvertedType

		repetitionLevels, definitionLevels := make([]int32, len(path)), make([]int32, len(path))
		for i := 0; i<len(path); i++ {
			repetitionLevels[i], _ = schemaHandler.MaxRepetitionLevel(path[:i+1])
			definitionLevels[i], _ = schemaHandler.MaxDefinitionLevel(path[:i+1])
		}

		for _, rc := range sliceRecords {
			rc.Index = -1
		}
		for _, rc := range mapRecords {
			rc.Index = -1
		}

		for i := bgn; i < end; i++ {
			rl, dl, val := table.RepetitionLevels[i], table.DefinitionLevels[i], table.Values[i]
			po, index := root, 0
			for index < len(path) {
				if po.Type().Kind() == reflect.Slice {
					if po.IsNil() {
						po.Set(reflect.MakeSlice(po.Type(), 0, 0))
					}

					if _, ok := sliceRecords[po]; !ok {
						sliceRecords[po] = &SliceRecord{
							Values:	[]reflect.Value{},
							Index:	-1,
						}
						sliceRecordsStack = append(sliceRecordsStack, po)
					}

					if rl == repetitionLevels[index] || sliceRecords[po].Index < 0 {
						sliceRecords[po].Index++
					}

					if sliceRecords[po].Index >= len(sliceRecords[po].Values) {
						sliceRecords[po].Values = append(sliceRecords[po].Values, reflect.New(po.Type().Elem()).Elem())
					}

					po = sliceRecords[po].Values[sliceRecords[po].Index]

				} else if po.Type().Kind() == reflect.Map {
					if po.IsNil() {
						po.Set(reflect.MakeMap(po.Type()))
					}

					if _, ok := mapRecords[po]; !ok {
						mapRecords[po] = &MapRecord{
							KeyValues:	[]KeyValue{},
							Index:		-1,
						}
						mapRecordsStack = append(mapRecordsStack, po)
					}

					index++
					if definitionLevels[index] > dl {
						break
					}

					if rl == repetitionLevels[index] || mapRecords[po].Index < 0 {
						mapRecords[po].Index++
					}

					if mapRecords[po].Index >= len(mapRecords[po].KeyValues) {
						mapRecords[po].KeyValues = append(mapRecords[po].KeyValues,
							KeyValue{
								Key: reflect.New(po.Type().Key()).Elem(),
								Value: reflect.New(po.Type().Elem()).Elem(),
						})
					}

					if path[index + 1] == "key" {
						po = mapRecords[po].KeyValues[mapRecords[po].Index].Key

					}else {
						po = mapRecords[po].KeyValues[mapRecords[po].Index].Value
					}

					index++
					if definitionLevels[index] > dl {
						break
					}

				} else if po.Type().Kind() == reflect.Ptr {
					if po.IsNil() {
						po.Set(reflect.New(po.Type().Elem()))
					}

					po = po.Elem()

				} else if po.Type().Kind() == reflect.Struct {
					index++
					if definitionLevels[index] > dl {
						break;
					}
					po = po.FieldByName(path[index])

				} else {
					po.Set(reflect.ValueOf(types.ParquetTypeToGoType(val, pT, cT)))
					break
				}
			}
		}
	}

	for i := len(sliceRecordsStack) - 1; i >= 0; i-- {
		po := sliceRecordsStack[i]
		vs := sliceRecords[po]
		potmp := reflect.Append(po, vs.Values...)
		po.Set(potmp)
	}

	for i := len(mapRecordsStack) - 1; i >= 0; i-- {
		po := mapRecordsStack[i]
		for _, kv := range mapRecords[po].KeyValues {
			po.SetMapIndex(kv.Key, kv.Value)
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package radix

import (
	"crypto/sha1"
	"encoding/hex"
	"io"
	"strings"

	"github.com/mediocregopher/radix.v2/resp"
)

// Action is an entity which can perform one or more tasks using a Conn
type Action interface {
	// Key returns a key which will be acted on. If the Action will act on more
	// than one key than any one can be returned. If no keys will be acted on
	// then nil should be returned.
	Key() []byte

	// Run actually performs the action using the given Conn
	Run(c Conn) error
}

// RawCmd implements the Action interface and describes a single redis command
// to be performed.
type RawCmd struct {
	// The name of the redis command to be performed. Always required
	Cmd []byte

	// The keys being operated on, and may be left empty if the command doesn't
	// operate on any specific key(s) (e.g.  SCAN)
	// TODO singular?
	Keys [][]byte

	// Args are any extra arguments to the command and can be almost any thing
	// TODO more deets
	Args []interface{}

	// Pointer value into which results from the command will be unmarshalled.
	// The Into method can be used to set this as well. See the Decoder docs
	// for more on unmarshalling
	Rcv interface{}
}

// Cmd returns an initialized RawCmd, populating the fields with the given
// values. Use CmdNoKey for commands which don't have an actual key (e.g. MULTI
// or PING). You can chain the Into method to conveniently set a result
// receiver.
func Cmd(cmd, key string, args ...interface{}) RawCmd {
	return RawCmd{
		Cmd:  []byte(cmd),
		Keys: [][]byte{[]byte(key)},
		Args: args,
	}
}

// CmdNoKey is like Cmd, but the returned RawCmd will not have its Keys field
// set.
func CmdNoKey(cmd string, args ...interface{}) RawCmd {
	return RawCmd{
		Cmd:  []byte(cmd),
		Args: args,
	}
}

// Into returns a RawCmd with all the same fields as the original, except the
// Rcv field set to the given value.
func (rc RawCmd) Into(rcv interface{}) RawCmd {
	rc.Rcv = rcv
	return rc
}

// Key implements the Key method of the Action interface.
func (rc RawCmd) Key() []byte {
	if len(rc.Keys) == 0 {
		return nil
	}
	return rc.Keys[0]
}

// MarshalRESP implements the resp.Marshaler interface.
// TODO describe how commands are written
func (rc RawCmd) MarshalRESP(p *resp.Pool, w io.Writer) error {
	var err error
	marshal := func(m resp.Marshaler) {
		if err != nil {
			return
		}
		err = m.MarshalRESP(p, w)
	}

	a := resp.Any{
		I:                     rc.Args,
		MarshalBulkString:     true,
		MarshalNoArrayHeaders: true,
	}
	marshal(resp.ArrayHeader{N: 1 + len(rc.Keys) + a.NumElems()})
	marshal(resp.BulkString{B: rc.Cmd})
	for _, k := range rc.Keys {
		marshal(resp.BulkString{B: k})
	}
	marshal(a)
	return err
}

// Run implements the Run method of the Action interface. It writes the RawCmd
// to the Conn, and unmarshals the result into the Rcv field (if set).
func (rc RawCmd) Run(conn Conn) error {
	if err := conn.Encode(rc); err != nil {
		return err
	}

	// Any will discard the data if its I is nil
	rcva := resp.Any{I: rc.Rcv}
	return conn.Decode(&rcva)
}

////////////////////////////////////////////////////////////////////////////////

var (
	evalsha = []byte("EVALSHA")
	eval    = []byte("EVAL")
)

// RawLuaCmd is an Action similar to RawCmd, but it runs a lua script on the
// redis server instead of a single Cmd. See redis' EVAL docs for more on how
// that works.
type RawLuaCmd struct {
	// The actual lua script which will be run.
	Script string

	// The keys being operated on, and may be left empty if the command doesn't
	// operate on any specific key(s)
	Keys []string

	// Args are any extra arguments to the command and can be almost any thing
	// TODO more deets
	Args []interface{}

	// Pointer value into which results from the command will be unmarshalled.
	// The Into method can be used to set this as well. See the Decoder docs
	// for more on unmarshalling
	Rcv interface{}
}

// LuaCmd returns an initialized RawLuraCmd, populating the fields with the given
// values. You can chain the Into method to conveniently set a result receiver.
func LuaCmd(script string, keys []string, args ...interface{}) RawLuaCmd {
	return RawLuaCmd{
		Script: script,
		Keys:   keys,
		Args:   args,
	}
}

// Into returns a RawLuaCmd with all the same fields as the original, except the
// Rcv field set to the given value.
func (rlc RawLuaCmd) Into(rcv interface{}) RawLuaCmd {
	rlc.Rcv = rcv
	return rlc
}

// Key implements the Key method of the Action interface.
func (rlc RawLuaCmd) Key() []byte {
	if len(rlc.Keys) == 0 {
		return nil
	}
	return []byte(rlc.Keys[0])
}

// Run implements the Run method of the Action interface. It will first attempt
// to perform the command using an EVALSHA, but will fallback to a normal EVAL
// if that doesn't work.
func (rlc RawLuaCmd) Run(conn Conn) error {
	// TODO if the Marshaler interface handled low-level operations, like we need
	// here, better then we could probably get rid of all this weird slice
	// shifting nonsense.

	rc := RawCmd{
		Cmd:  evalsha,
		Args: rlc.Args,
		Rcv:  rlc.Rcv,
	}

	// TODO alloc here probably isn't necessary
	sumRaw := sha1.Sum([]byte(rlc.Script))
	sum := hex.EncodeToString(sumRaw[:])

	// shift the arguments in place to allow room for the script hash, key
	// count, and keys themselves
	// TODO this isn't great because if the Args are re-allocated we don't
	// actually give back to the original Cmd
	shiftBy := 2 + len(rlc.Keys)
	for i := 0; i < shiftBy; i++ {
		rc.Args = append(rc.Args, nil)
	}
	copy(rc.Args[shiftBy:], rc.Args)
	rc.Args[0] = sum
	rc.Args[1] = len(rlc.Keys)
	for i := range rlc.Keys {
		rc.Args[i+2] = rlc.Keys[i]
	}

	err := rc.Run(conn)
	if err != nil && strings.HasPrefix(err.Error(), "NOSCRIPT") {
		rc.Cmd = eval
		rc.Args[0] = rlc.Script
		err = rc.Run(conn)
	}
	return err
}

////////////////////////////////////////////////////////////////////////////////

// Pipeline is an Action which first writes multiple commands to a Conn in a
// single write, then reads their responses in a single step. This effectively
// reduces network delay into a single round-trip
//
//	var fooVal string
//	p := Pipeline{
//		Cmd("SET", "foo", "bar"),
//		Cmd("GET", "foo").Into(&fooVal),
//	}
//	if err := conn.Do(p); err != nil {
//		panic(err)
//	}
//	fmt.Printf("fooVal: %q\n", fooVal)
//
type Pipeline []RawCmd

// Run implements the method for the Action interface. It will write all RawCmds
// in it sequentially, then read all of their responses sequentially.
//
// If an error is encountered the error will be returned immediately.
func (p Pipeline) Run(c Conn) error {
	for _, cmd := range p {
		if err := c.Encode(cmd); err != nil {
			return err
		}
	}

	for _, cmd := range p {
		rcva := resp.Any{I: cmd.Rcv}
		if err := c.Decode(&rcva); err != nil {
			return err
		}
	}

	return nil
}
<commit_msg>streamline RawLuaCmd to not do crazy splice shifting stuff<commit_after>package radix

import (
	"crypto/sha1"
	"encoding/hex"
	"io"
	"strings"

	"github.com/mediocregopher/radix.v2/resp"
)

// Action is an entity which can perform one or more tasks using a Conn
type Action interface {
	// Key returns a key which will be acted on. If the Action will act on more
	// than one key than any one can be returned. If no keys will be acted on
	// then nil should be returned.
	Key() []byte

	// Run actually performs the action using the given Conn
	Run(c Conn) error
}

// RawCmd implements the Action interface and describes a single redis command
// to be performed.
type RawCmd struct {
	// The name of the redis command to be performed. Always required
	Cmd []byte

	// The keys being operated on, and may be left empty if the command doesn't
	// operate on any specific key(s) (e.g.  SCAN)
	// TODO singular?
	Keys [][]byte

	// Args are any extra arguments to the command and can be almost any thing
	// TODO more deets
	Args []interface{}

	// Pointer value into which results from the command will be unmarshalled.
	// The Into method can be used to set this as well. See the Decoder docs
	// for more on unmarshalling
	Rcv interface{}
}

// Cmd returns an initialized RawCmd, populating the fields with the given
// values. Use CmdNoKey for commands which don't have an actual key (e.g. MULTI
// or PING). You can chain the Into method to conveniently set a result
// receiver.
func Cmd(cmd, key string, args ...interface{}) RawCmd {
	return RawCmd{
		Cmd:  []byte(cmd),
		Keys: [][]byte{[]byte(key)},
		Args: args,
	}
}

// CmdNoKey is like Cmd, but the returned RawCmd will not have its Keys field
// set.
func CmdNoKey(cmd string, args ...interface{}) RawCmd {
	return RawCmd{
		Cmd:  []byte(cmd),
		Args: args,
	}
}

// Into returns a RawCmd with all the same fields as the original, except the
// Rcv field set to the given value.
func (rc RawCmd) Into(rcv interface{}) RawCmd {
	rc.Rcv = rcv
	return rc
}

// Key implements the Key method of the Action interface.
func (rc RawCmd) Key() []byte {
	if len(rc.Keys) == 0 {
		return nil
	}
	return rc.Keys[0]
}

// MarshalRESP implements the resp.Marshaler interface.
// TODO describe how commands are written
func (rc RawCmd) MarshalRESP(p *resp.Pool, w io.Writer) error {
	var err error
	marshal := func(m resp.Marshaler) {
		if err != nil {
			return
		}
		err = m.MarshalRESP(p, w)
	}

	a := resp.Any{
		I:                     rc.Args,
		MarshalBulkString:     true,
		MarshalNoArrayHeaders: true,
	}
	marshal(resp.ArrayHeader{N: 1 + len(rc.Keys) + a.NumElems()})
	marshal(resp.BulkString{B: rc.Cmd})
	for _, k := range rc.Keys {
		marshal(resp.BulkString{B: k})
	}
	marshal(a)
	return err
}

// Run implements the Run method of the Action interface. It writes the RawCmd
// to the Conn, and unmarshals the result into the Rcv field (if set).
func (rc RawCmd) Run(conn Conn) error {
	if err := conn.Encode(rc); err != nil {
		return err
	}

	// Any will discard the data if its I is nil
	rcva := resp.Any{I: rc.Rcv}
	return conn.Decode(&rcva)
}

////////////////////////////////////////////////////////////////////////////////

var (
	evalsha = []byte("EVALSHA")
	eval    = []byte("EVAL")
)

// RawLuaCmd is an Action similar to RawCmd, but it runs a lua script on the
// redis server instead of a single Cmd. See redis' EVAL docs for more on how
// that works.
type RawLuaCmd struct {
	// The actual lua script which will be run.
	Script string

	// The keys being operated on, and may be left empty if the command doesn't
	// operate on any specific key(s)
	Keys []string

	// Args are any extra arguments to the command and can be almost any thing
	// TODO more deets
	Args []interface{}

	// Pointer value into which results from the command will be unmarshalled.
	// The Into method can be used to set this as well. See the Decoder docs
	// for more on unmarshalling
	Rcv interface{}
}

// LuaCmd returns an initialized RawLuraCmd, populating the fields with the given
// values. You can chain the Into method to conveniently set a result receiver.
func LuaCmd(script string, keys []string, args ...interface{}) RawLuaCmd {
	return RawLuaCmd{
		Script: script,
		Keys:   keys,
		Args:   args,
	}
}

// Into returns a RawLuaCmd with all the same fields as the original, except the
// Rcv field set to the given value.
func (rlc RawLuaCmd) Into(rcv interface{}) RawLuaCmd {
	rlc.Rcv = rcv
	return rlc
}

// Key implements the Key method of the Action interface.
func (rlc RawLuaCmd) Key() []byte {
	if len(rlc.Keys) == 0 {
		return nil
	}
	return []byte(rlc.Keys[0])
}

type mRawLuaCmd struct {
	RawLuaCmd
	eval bool
}

func (mrlc mRawLuaCmd) MarshalRESP(p *resp.Pool, w io.Writer) error {
	var err error
	marshal := func(m resp.Marshaler) {
		if err != nil {
			return
		}
		err = m.MarshalRESP(p, w)
	}

	a := resp.Any{
		I:                     mrlc.Args,
		MarshalBulkString:     true,
		MarshalNoArrayHeaders: true,
	}
	numKeys := len(mrlc.Keys)

	// EVAL(SHA) script/sum numkeys keys... args...
	marshal(resp.ArrayHeader{N: 3 + numKeys + a.NumElems()})
	if mrlc.eval {
		marshal(resp.BulkString{B: eval})
		marshal(resp.BulkString{B: []byte(mrlc.Script)})
	} else {
		// TODO alloc here isn't great
		sumRaw := sha1.Sum([]byte(mrlc.Script))
		sum := hex.EncodeToString(sumRaw[:])
		marshal(resp.BulkString{B: evalsha})
		marshal(resp.BulkString{B: []byte(sum)})
	}
	marshal(resp.Any{I: numKeys, MarshalBulkString: true})
	for _, k := range mrlc.Keys {
		marshal(resp.BulkString{B: []byte(k)})
	}
	marshal(a)
	return err
}

func (mrlc mRawLuaCmd) Run(conn Conn) error {
	if err := conn.Encode(mrlc); err != nil {
		return err
	}
	rcva := resp.Any{I: mrlc.Rcv}
	return conn.Decode(&rcva)
}

// Run implements the Run method of the Action interface. It will first attempt
// to perform the command using an EVALSHA, but will fallback to a normal EVAL
// if that doesn't work.
func (rlc RawLuaCmd) Run(conn Conn) error {
	err := mRawLuaCmd{RawLuaCmd: rlc}.Run(conn)
	if err != nil && strings.HasPrefix(err.Error(), "NOSCRIPT") {
		err = mRawLuaCmd{RawLuaCmd: rlc, eval: true}.Run(conn)
	}
	return err
}

////////////////////////////////////////////////////////////////////////////////

// Pipeline is an Action which first writes multiple commands to a Conn in a
// single write, then reads their responses in a single step. This effectively
// reduces network delay into a single round-trip
//
//	var fooVal string
//	p := Pipeline{
//		Cmd("SET", "foo", "bar"),
//		Cmd("GET", "foo").Into(&fooVal),
//	}
//	if err := conn.Do(p); err != nil {
//		panic(err)
//	}
//	fmt.Printf("fooVal: %q\n", fooVal)
//
type Pipeline []RawCmd

// Run implements the method for the Action interface. It will write all RawCmds
// in it sequentially, then read all of their responses sequentially.
//
// If an error is encountered the error will be returned immediately.
func (p Pipeline) Run(c Conn) error {
	for _, cmd := range p {
		if err := c.Encode(cmd); err != nil {
			return err
		}
	}

	for _, cmd := range p {
		rcva := resp.Any{I: cmd.Rcv}
		if err := c.Decode(&rcva); err != nil {
			return err
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package seqdb

import (
	"sync"
	"log"
	"net"
	"time"
	"strings"
)

type Service struct {
	ch chan bool
	waitGroup *sync.WaitGroup
}

// Make a new Service.
func NewService() *Service {
	s := &Service{
		ch:        make(chan bool),
		waitGroup: &sync.WaitGroup{},
	}
	s.waitGroup.Add(1)
	return s
}

// Accept connections and spawn a goroutine to serve each one.  Stop listening
// if anything is received on the service's channel.
func (s *Service) Serve(listener *net.TCPListener) {
	defer s.waitGroup.Done()
	for {
		select {
		case <-s.ch:
			log.Println("stopping listening on", listener.Addr())
			listener.Close()
			return
		default:
		}

		listener.SetDeadline(time.Now().Add(1e9))
		conn, err := listener.AcceptTCP()

		if nil != err {
			if opErr, ok := err.(*net.OpError); ok && opErr.Timeout() {
				continue
			}
			log.Println(err)
		}

		log.Println(conn.RemoteAddr(), "connected")
		s.waitGroup.Add(1)

		go s.serve(conn)
	}
}

// Stop the service by closing the service's channel.  Block until the service
// is really stopped.
func (s *Service) Stop() {
	close(s.ch)
	s.waitGroup.Wait()
}

// Serve a connection by reading and writing what was read.  That's right, this
// is an echo service.  Stop reading and writing if anything is received on the
// service's channel but only after writing what was read.
func (s *Service) serve(conn *net.TCPConn) {
	defer conn.Close()
	defer s.waitGroup.Done()
	for {

		select {
		case <-s.ch:
			log.Println("disconnecting", conn.RemoteAddr())
			return
		default:
		}

		conn.SetDeadline(time.Now().Add(10*1e9))

		buf := make([]byte, 4096)

		n, err := conn.Read(buf)

		if err != nil {
			log.Printf("Cannot read from buffer: %v\r\n", err)

		}

		s := strings.Trim(string(buf[:n]), " \r\n")

		// Parse the command
		code, response := ParseCommand(s)

		// QUIT command was issued
		if code == -1 { break }

		// Write response
		_, err = conn.Write([]byte(response))

		if err != nil {
			log.Printf("Error writing buffer: %v\r\n", err)
		}
	}
}<commit_msg>link added<commit_after>package seqdb

/*
	Based on http://rcrowley.org/articles/golang-graceful-stop.html
 */

import (
	"sync"
	"log"
	"net"
	"time"
	"strings"
)

type Service struct {
	ch chan bool
	waitGroup *sync.WaitGroup
}

// Make a new Service.
func NewService() *Service {
	s := &Service{
		ch:        make(chan bool),
		waitGroup: &sync.WaitGroup{},
	}
	s.waitGroup.Add(1)
	return s
}

// Accept connections and spawn a goroutine to serve each one.  Stop listening
// if anything is received on the service's channel.
func (s *Service) Serve(listener *net.TCPListener) {
	defer s.waitGroup.Done()
	for {
		select {
		case <-s.ch:
			log.Println("stopping listening on", listener.Addr())
			listener.Close()
			return
		default:
		}

		listener.SetDeadline(time.Now().Add(1e9))
		conn, err := listener.AcceptTCP()

		if nil != err {
			if opErr, ok := err.(*net.OpError); ok && opErr.Timeout() {
				continue
			}
			log.Println(err)
		}

		log.Println(conn.RemoteAddr(), "connected")
		s.waitGroup.Add(1)

		go s.serve(conn)
	}
}

// Stop the service by closing the service's channel.  Block until the service
// is really stopped.
func (s *Service) Stop() {
	close(s.ch)
	s.waitGroup.Wait()
}

// Serve a connection by reading and writing what was read.  That's right, this
// is an echo service.  Stop reading and writing if anything is received on the
// service's channel but only after writing what was read.
func (s *Service) serve(conn *net.TCPConn) {
	defer conn.Close()
	defer s.waitGroup.Done()
	for {

		select {
		case <-s.ch:
			log.Println("disconnecting", conn.RemoteAddr())
			return
		default:
		}

		conn.SetDeadline(time.Now().Add(10*1e9))

		buf := make([]byte, 4096)

		n, err := conn.Read(buf)

		if err != nil {
			log.Printf("Cannot read from buffer: %v\r\n", err)

		}

		s := strings.Trim(string(buf[:n]), " \r\n")

		// Parse the command
		code, response := ParseCommand(s)

		// QUIT command was issued
		if code == -1 { break }

		// Write response
		_, err = conn.Write([]byte(response))

		if err != nil {
			log.Printf("Error writing buffer: %v\r\n", err)
		}
	}
}<|endoftext|>


<commit_before>// Copyright (c) Ilia Kravets, 2015. All rights reserved. PROVIDED "AS IS"
// WITHOUT ANY WARRANTY, EXPRESS OR IMPLIED. See LICENSE file for details.

package packet

import (
	"log"

	"github.com/google/gopacket"

	"my/errs"
)

type DecodingLayerFactory interface {
	Create(gopacket.LayerType) gopacket.DecodingLayer
	SupportedLayers() gopacket.LayerClass
}

type ReusingLayerParser struct {
	first      gopacket.LayerType
	factories  map[gopacket.LayerType]DecodingLayerFactory
	layerCache map[gopacket.LayerType]*[]gopacket.DecodingLayer
	df         gopacket.DecodeFeedback
	Truncated  bool
	tps        []TypedPayload
}

func NewReusingLayerParser(first gopacket.LayerType, factories ...DecodingLayerFactory) *ReusingLayerParser {
	p := &ReusingLayerParser{
		factories:  make(map[gopacket.LayerType]DecodingLayerFactory),
		layerCache: make(map[gopacket.LayerType]*[]gopacket.DecodingLayer),
		first:      first,
		tps:        make([]TypedPayload, 100),
	}
	p.df = p // cast this once to the interface
	for _, f := range factories {
		p.AddDecodingLayerFactory(f)
	}
	p.AddDecodingLayerFactory(UnknownDecodingLayerFactory)
	return p
}

func (p *ReusingLayerParser) AddDecodingLayerFactory(f DecodingLayerFactory) {
	for _, typ := range f.SupportedLayers().LayerTypes() {
		if p.factories[typ] == nil {
			p.factories[typ] = f
		}
	}
}

func (p *ReusingLayerParser) DecodeLayers(data []byte, decoded *[]gopacket.DecodingLayer) (err error) {
	defer errs.PassE(&err)
	p.Truncated = false

	dlEnd := 0
	p.tps[0] = TypedPayload{Type: p.first, Payload: data}
	for tpsStart, tpsEnd := 0, 1; tpsStart < tpsEnd; tpsStart++ {
		tp := &p.tps[tpsStart]
		if tp.Type == gopacket.LayerTypeZero {
			continue
		}
		if dlEnd == len(*decoded) {
			*decoded = append(*decoded, nil)
		}
		layer := (*decoded)[dlEnd]
		layer, err := p.tryDecode(tp, layer)
		if err != nil {
			log.Printf("error: %s\n", err)
			tp.Type = gopacket.LayerTypeDecodeFailure
			layer, err = p.tryDecode(tp, layer)
			errs.CheckE(err)
		}
		(*decoded)[dlEnd] = layer
		dlEnd++

		if dml, ok := layer.(DecodingMultiLayer); ok {
			nls := dml.NextLayers()
			p.checkTpsSize(tpsEnd, len(nls))
			copy(p.tps[tpsEnd:], nls)
			tpsEnd += len(nls)
		} else if len(layer.LayerPayload()) > 0 {
			p.checkTpsSize(tpsEnd, 1)
			p.tps[tpsEnd] = TypedPayload{Type: layer.NextLayerType(), Payload: layer.LayerPayload()}
			tpsEnd++
		}
	}
	p.recycle((*decoded)[dlEnd:])
	*decoded = (*decoded)[:dlEnd]
	return nil
}

func (p *ReusingLayerParser) checkTpsSize(tpsEnd int, add int) {
	if len(p.tps) < tpsEnd+add {
		oldTps := p.tps[:tpsEnd]
		p.tps = make([]TypedPayload, len(p.tps)*2+add)
		copy(p.tps, oldTps)
	}
}

func (p *ReusingLayerParser) tryDecode(tp *TypedPayload, oldLayer gopacket.DecodingLayer) (layer gopacket.DecodingLayer, err error) {
	if oldLayer != nil && oldLayer.CanDecode().Contains(tp.Type) {
		layer = oldLayer
	} else {
		if oldLayer != nil {
			p.recycleOne(oldLayer)
		}
		if layer, err = p.getDecodingLayer(tp.Type); err != nil {
			return
		}
	}
	if err = layer.DecodeFromBytes(tp.Payload, p.df); err != nil {
		p.recycleOne(layer)
		layer = nil
	}
	return
}

func (p *ReusingLayerParser) recycle(layers []gopacket.DecodingLayer) {
	for _, layer := range layers {
		p.recycleOne(layer)
	}
}
func (p *ReusingLayerParser) recycleOne(layer gopacket.DecodingLayer) {
	typ := layer.CanDecode().LayerTypes()[0]
	lc := p.layerCache[typ]
	if lc == nil {
		lc = &[]gopacket.DecodingLayer{}
		p.layerCache[typ] = lc
	}
	*lc = append(*lc, layer)
}
func (p *ReusingLayerParser) getDecodingLayer(typ gopacket.LayerType) (layer gopacket.DecodingLayer, err error) {
	if layers, ok := p.layerCache[typ]; ok && len(*layers) > 0 {
		layer = (*layers)[len(*layers)-1]
		*layers = (*layers)[:len(*layers)-1]
		return
	}
	if factory, ok := p.factories[typ]; ok {
		layer = factory.Create(typ)
	} else {
		err = gopacket.UnsupportedLayerType(typ)
	}
	return
}

func (p *ReusingLayerParser) SetTruncated() {
	p.Truncated = true
}

type SingleDecodingLayerFactory struct {
	layerType gopacket.LayerType
	create    func() gopacket.DecodingLayer
}

var _ DecodingLayerFactory = &SingleDecodingLayerFactory{}

func NewSingleDecodingLayerFactory(layerType gopacket.LayerType, create func() gopacket.DecodingLayer) *SingleDecodingLayerFactory {
	return &SingleDecodingLayerFactory{
		layerType: layerType,
		create:    create,
	}
}
func (f *SingleDecodingLayerFactory) Create(layerType gopacket.LayerType) gopacket.DecodingLayer {
	errs.Check(layerType == f.layerType)
	return f.create()
}
func (f *SingleDecodingLayerFactory) SupportedLayers() gopacket.LayerClass {
	return f.layerType
}

type UnknownDecodingLayer struct {
	Data []byte
}

var (
	_ gopacket.Layer         = &UnknownDecodingLayer{}
	_ gopacket.DecodingLayer = &UnknownDecodingLayer{}
)

func (d *UnknownDecodingLayer) DecodeFromBytes(data []byte, df gopacket.DecodeFeedback) error {
	d.Data = data
	return nil
}
func (d *UnknownDecodingLayer) CanDecode() gopacket.LayerClass    { return d.LayerType() }
func (d *UnknownDecodingLayer) NextLayerType() gopacket.LayerType { return gopacket.LayerTypeZero }
func (d *UnknownDecodingLayer) LayerType() gopacket.LayerType     { return gopacket.LayerTypeDecodeFailure }
func (d *UnknownDecodingLayer) LayerContents() []byte             { return d.Data }
func (d *UnknownDecodingLayer) LayerPayload() []byte              { return nil }

var UnknownDecodingLayerFactory = NewSingleDecodingLayerFactory(
	gopacket.LayerTypeDecodeFailure,
	func() gopacket.DecodingLayer { return &UnknownDecodingLayer{} },
)

var _ = log.Ldate
<commit_msg>packet:trivial: add commented out logging<commit_after>// Copyright (c) Ilia Kravets, 2015. All rights reserved. PROVIDED "AS IS"
// WITHOUT ANY WARRANTY, EXPRESS OR IMPLIED. See LICENSE file for details.

package packet

import (
	"log"

	"github.com/google/gopacket"

	"my/errs"
)

type DecodingLayerFactory interface {
	Create(gopacket.LayerType) gopacket.DecodingLayer
	SupportedLayers() gopacket.LayerClass
}

type ReusingLayerParser struct {
	first      gopacket.LayerType
	factories  map[gopacket.LayerType]DecodingLayerFactory
	layerCache map[gopacket.LayerType]*[]gopacket.DecodingLayer
	df         gopacket.DecodeFeedback
	Truncated  bool
	tps        []TypedPayload
}

func NewReusingLayerParser(first gopacket.LayerType, factories ...DecodingLayerFactory) *ReusingLayerParser {
	p := &ReusingLayerParser{
		factories:  make(map[gopacket.LayerType]DecodingLayerFactory),
		layerCache: make(map[gopacket.LayerType]*[]gopacket.DecodingLayer),
		first:      first,
		tps:        make([]TypedPayload, 100),
	}
	p.df = p // cast this once to the interface
	for _, f := range factories {
		p.AddDecodingLayerFactory(f)
	}
	p.AddDecodingLayerFactory(UnknownDecodingLayerFactory)
	return p
}

func (p *ReusingLayerParser) AddDecodingLayerFactory(f DecodingLayerFactory) {
	for _, typ := range f.SupportedLayers().LayerTypes() {
		if p.factories[typ] == nil {
			p.factories[typ] = f
		}
	}
}

func (p *ReusingLayerParser) DecodeLayers(data []byte, decoded *[]gopacket.DecodingLayer) (err error) {
	defer errs.PassE(&err)
	p.Truncated = false

	dlEnd := 0
	p.tps[0] = TypedPayload{Type: p.first, Payload: data}
	for tpsStart, tpsEnd := 0, 1; tpsStart < tpsEnd; tpsStart++ {
		tp := &p.tps[tpsStart]
		if tp.Type == gopacket.LayerTypeZero {
			continue
		}
		if dlEnd == len(*decoded) {
			*decoded = append(*decoded, nil)
		}
		layer := (*decoded)[dlEnd]
		layer, err := p.tryDecode(tp, layer)
		if err != nil {
			log.Printf("error: %s\n", err)
			tp.Type = gopacket.LayerTypeDecodeFailure
			layer, err = p.tryDecode(tp, layer)
			errs.CheckE(err)
		}
		(*decoded)[dlEnd] = layer
		//log.Printf("decoded %T %v\n", layer, layer)
		dlEnd++

		if dml, ok := layer.(DecodingMultiLayer); ok {
			nls := dml.NextLayers()
			p.checkTpsSize(tpsEnd, len(nls))
			copy(p.tps[tpsEnd:], nls)
			tpsEnd += len(nls)
		} else if len(layer.LayerPayload()) > 0 {
			p.checkTpsSize(tpsEnd, 1)
			p.tps[tpsEnd] = TypedPayload{Type: layer.NextLayerType(), Payload: layer.LayerPayload()}
			tpsEnd++
		}
	}
	p.recycle((*decoded)[dlEnd:])
	*decoded = (*decoded)[:dlEnd]
	return nil
}

func (p *ReusingLayerParser) checkTpsSize(tpsEnd int, add int) {
	if len(p.tps) < tpsEnd+add {
		oldTps := p.tps[:tpsEnd]
		p.tps = make([]TypedPayload, len(p.tps)*2+add)
		copy(p.tps, oldTps)
	}
}

func (p *ReusingLayerParser) tryDecode(tp *TypedPayload, oldLayer gopacket.DecodingLayer) (layer gopacket.DecodingLayer, err error) {
	if oldLayer != nil && oldLayer.CanDecode().Contains(tp.Type) {
		layer = oldLayer
	} else {
		if oldLayer != nil {
			p.recycleOne(oldLayer)
		}
		if layer, err = p.getDecodingLayer(tp.Type); err != nil {
			return
		}
	}
	if err = layer.DecodeFromBytes(tp.Payload, p.df); err != nil {
		p.recycleOne(layer)
		layer = nil
	}
	return
}

func (p *ReusingLayerParser) recycle(layers []gopacket.DecodingLayer) {
	for _, layer := range layers {
		p.recycleOne(layer)
	}
}
func (p *ReusingLayerParser) recycleOne(layer gopacket.DecodingLayer) {
	typ := layer.CanDecode().LayerTypes()[0]
	lc := p.layerCache[typ]
	if lc == nil {
		lc = &[]gopacket.DecodingLayer{}
		p.layerCache[typ] = lc
	}
	*lc = append(*lc, layer)
}
func (p *ReusingLayerParser) getDecodingLayer(typ gopacket.LayerType) (layer gopacket.DecodingLayer, err error) {
	if layers, ok := p.layerCache[typ]; ok && len(*layers) > 0 {
		layer = (*layers)[len(*layers)-1]
		*layers = (*layers)[:len(*layers)-1]
		return
	}
	if factory, ok := p.factories[typ]; ok {
		layer = factory.Create(typ)
	} else {
		err = gopacket.UnsupportedLayerType(typ)
	}
	return
}

func (p *ReusingLayerParser) SetTruncated() {
	p.Truncated = true
}

type SingleDecodingLayerFactory struct {
	layerType gopacket.LayerType
	create    func() gopacket.DecodingLayer
}

var _ DecodingLayerFactory = &SingleDecodingLayerFactory{}

func NewSingleDecodingLayerFactory(layerType gopacket.LayerType, create func() gopacket.DecodingLayer) *SingleDecodingLayerFactory {
	return &SingleDecodingLayerFactory{
		layerType: layerType,
		create:    create,
	}
}
func (f *SingleDecodingLayerFactory) Create(layerType gopacket.LayerType) gopacket.DecodingLayer {
	errs.Check(layerType == f.layerType)
	return f.create()
}
func (f *SingleDecodingLayerFactory) SupportedLayers() gopacket.LayerClass {
	return f.layerType
}

type UnknownDecodingLayer struct {
	Data []byte
}

var (
	_ gopacket.Layer         = &UnknownDecodingLayer{}
	_ gopacket.DecodingLayer = &UnknownDecodingLayer{}
)

func (d *UnknownDecodingLayer) DecodeFromBytes(data []byte, df gopacket.DecodeFeedback) error {
	d.Data = data
	return nil
}
func (d *UnknownDecodingLayer) CanDecode() gopacket.LayerClass    { return d.LayerType() }
func (d *UnknownDecodingLayer) NextLayerType() gopacket.LayerType { return gopacket.LayerTypeZero }
func (d *UnknownDecodingLayer) LayerType() gopacket.LayerType     { return gopacket.LayerTypeDecodeFailure }
func (d *UnknownDecodingLayer) LayerContents() []byte             { return d.Data }
func (d *UnknownDecodingLayer) LayerPayload() []byte              { return nil }

var UnknownDecodingLayerFactory = NewSingleDecodingLayerFactory(
	gopacket.LayerTypeDecodeFailure,
	func() gopacket.DecodingLayer { return &UnknownDecodingLayer{} },
)

var _ = log.Ldate
<|endoftext|>


<commit_before>package state

import (
	"errors"
	"math/big"
	"sort"
	"sync"

	"github.com/bytom/common"
	"github.com/bytom/consensus"
	"github.com/bytom/consensus/difficulty"
	"github.com/bytom/protocol/bc"
	"github.com/bytom/protocol/bc/types"
)

// approxNodesPerDay is an approximation of the number of new blocks there are
// in a week on average.
const approxNodesPerDay = 24 * 24

// BlockNode represents a block within the block chain and is primarily used to
// aid in selecting the best chain to be the main chain.
type BlockNode struct {
	Parent  *BlockNode // parent is the parent block for this node.
	Hash    bc.Hash    // hash of the block.
	Seed    *bc.Hash   // seed hash of the block
	WorkSum *big.Int   // total amount of work in the chain up to

	Version                uint64
	Height                 uint64
	Timestamp              uint64
	Nonce                  uint64
	Bits                   uint64
	TransactionsMerkleRoot bc.Hash
	TransactionStatusHash  bc.Hash
}

func NewBlockNode(bh *types.BlockHeader, parent *BlockNode) (*BlockNode, error) {
	if bh.Height != 0 && parent == nil {
		return nil, errors.New("parent node can not be nil")
	}

	node := &BlockNode{
		Parent:    parent,
		Hash:      bh.Hash(),
		WorkSum:   difficulty.CalcWork(bh.Bits),
		Version:   bh.Version,
		Height:    bh.Height,
		Timestamp: bh.Timestamp,
		Nonce:     bh.Nonce,
		Bits:      bh.Bits,
		TransactionsMerkleRoot: bh.TransactionsMerkleRoot,
		TransactionStatusHash:  bh.TransactionStatusHash,
	}

	if bh.Height == 0 {
		node.Seed = consensus.InitialSeed
	} else {
		node.Seed = parent.CalcNextSeed()
		node.WorkSum = node.WorkSum.Add(parent.WorkSum, node.WorkSum)
	}
	return node, nil
}

// blockHeader convert a node to the header struct
func (node *BlockNode) BlockHeader() *types.BlockHeader {
	previousBlockHash := bc.Hash{}
	if node.Parent != nil {
		previousBlockHash = node.Parent.Hash
	}
	return &types.BlockHeader{
		Version:           node.Version,
		Height:            node.Height,
		PreviousBlockHash: previousBlockHash,
		Timestamp:         node.Timestamp,
		Nonce:             node.Nonce,
		Bits:              node.Bits,
		BlockCommitment: types.BlockCommitment{
			TransactionsMerkleRoot: node.TransactionsMerkleRoot,
			TransactionStatusHash:  node.TransactionStatusHash,
		},
	}
}

func (node *BlockNode) CalcPastMedianTime() uint64 {
	timestamps := []uint64{}
	iterNode := node
	for i := 0; i < consensus.MedianTimeBlocks && iterNode != nil; i++ {
		timestamps = append(timestamps, iterNode.Timestamp)
		iterNode = iterNode.Parent
	}

	sort.Sort(common.TimeSorter(timestamps))
	return timestamps[len(timestamps)/2]
}

// CalcNextBits calculate the bits for next block
func (node *BlockNode) CalcNextBits() uint64 {
	if node.Height%consensus.BlocksPerRetarget != 0 || node.Height == 0 {
		return node.Bits
	}

	compareNode := node.Parent
	for compareNode.Height%consensus.BlocksPerRetarget != 0 {
		compareNode = compareNode.Parent
	}
	return difficulty.CalcNextRequiredDifficulty(node.BlockHeader(), compareNode.BlockHeader())
}

// CalcNextSeed calculate the seed for next block
func (node *BlockNode) CalcNextSeed() *bc.Hash {
	if node.Height == 0 {
		return consensus.InitialSeed
	}
	if node.Height%consensus.SeedPerRetarget == 0 {
		return &node.Hash
	}
	return node.Seed
}

// BlockIndex is the struct for help chain trace block chain as tree
type BlockIndex struct {
	sync.RWMutex

	index     map[bc.Hash]*BlockNode
	mainChain []*BlockNode
}

// NewBlockIndex will create a empty BlockIndex
func NewBlockIndex() *BlockIndex {
	return &BlockIndex{
		index:     make(map[bc.Hash]*BlockNode),
		mainChain: make([]*BlockNode, 0, approxNodesPerDay),
	}
}

// AddNode will add node to the index map
func (bi *BlockIndex) AddNode(node *BlockNode) {
	bi.Lock()
	bi.index[node.Hash] = node
	bi.Unlock()
}

// GetNode will search node from the index map
func (bi *BlockIndex) GetNode(hash *bc.Hash) *BlockNode {
	bi.RLock()
	defer bi.RUnlock()
	return bi.index[*hash]
}

func (bi *BlockIndex) BestNode() *BlockNode {
	bi.RLock()
	defer bi.RUnlock()
	return bi.mainChain[len(bi.mainChain)-1]
}

// BlockExist check does the block existed in blockIndex
func (bi *BlockIndex) BlockExist(hash *bc.Hash) bool {
	bi.RLock()
	_, ok := bi.index[*hash]
	bi.RUnlock()
	return ok
}

// TODO: THIS FUNCTION MIGHT BE DELETED
func (bi *BlockIndex) InMainchain(hash bc.Hash) bool {
	bi.RLock()
	defer bi.RUnlock()

	node, ok := bi.index[hash]
	if !ok {
		return false
	}
	return bi.nodeByHeight(node.Height) == node
}

func (bi *BlockIndex) nodeByHeight(height uint64) *BlockNode {
	if height >= uint64(len(bi.mainChain)) {
		return nil
	}
	return bi.mainChain[height]
}

// NodeByHeight returns the block node at the specified height.
func (bi *BlockIndex) NodeByHeight(height uint64) *BlockNode {
	bi.RLock()
	defer bi.RUnlock()
	return bi.nodeByHeight(height)
}

// SetMainChain will set the the mainChain array
func (bi *BlockIndex) SetMainChain(node *BlockNode) {
	bi.Lock()
	defer bi.Unlock()

	needed := node.Height + 1
	if uint64(cap(bi.mainChain)) < needed {
		nodes := make([]*BlockNode, needed, needed+approxNodesPerDay)
		copy(nodes, bi.mainChain)
		bi.mainChain = nodes
	} else {
		i := uint64(len(bi.mainChain))
		bi.mainChain = bi.mainChain[0:needed]
		for ; i < needed; i++ {
			bi.mainChain[i] = nil
		}
	}

	for node != nil && bi.mainChain[node.Height] != node {
		bi.mainChain[node.Height] = node
		node = node.Parent
	}
}
<commit_msg>fix comment (#1152)<commit_after>package state

import (
	"errors"
	"math/big"
	"sort"
	"sync"

	"github.com/bytom/common"
	"github.com/bytom/consensus"
	"github.com/bytom/consensus/difficulty"
	"github.com/bytom/protocol/bc"
	"github.com/bytom/protocol/bc/types"
)

// approxNodesPerDay is an approximation of the number of new blocks there are
// in a day on average.
const approxNodesPerDay = 24 * 24

// BlockNode represents a block within the block chain and is primarily used to
// aid in selecting the best chain to be the main chain.
type BlockNode struct {
	Parent  *BlockNode // parent is the parent block for this node.
	Hash    bc.Hash    // hash of the block.
	Seed    *bc.Hash   // seed hash of the block
	WorkSum *big.Int   // total amount of work in the chain up to

	Version                uint64
	Height                 uint64
	Timestamp              uint64
	Nonce                  uint64
	Bits                   uint64
	TransactionsMerkleRoot bc.Hash
	TransactionStatusHash  bc.Hash
}

func NewBlockNode(bh *types.BlockHeader, parent *BlockNode) (*BlockNode, error) {
	if bh.Height != 0 && parent == nil {
		return nil, errors.New("parent node can not be nil")
	}

	node := &BlockNode{
		Parent:    parent,
		Hash:      bh.Hash(),
		WorkSum:   difficulty.CalcWork(bh.Bits),
		Version:   bh.Version,
		Height:    bh.Height,
		Timestamp: bh.Timestamp,
		Nonce:     bh.Nonce,
		Bits:      bh.Bits,
		TransactionsMerkleRoot: bh.TransactionsMerkleRoot,
		TransactionStatusHash:  bh.TransactionStatusHash,
	}

	if bh.Height == 0 {
		node.Seed = consensus.InitialSeed
	} else {
		node.Seed = parent.CalcNextSeed()
		node.WorkSum = node.WorkSum.Add(parent.WorkSum, node.WorkSum)
	}
	return node, nil
}

// blockHeader convert a node to the header struct
func (node *BlockNode) BlockHeader() *types.BlockHeader {
	previousBlockHash := bc.Hash{}
	if node.Parent != nil {
		previousBlockHash = node.Parent.Hash
	}
	return &types.BlockHeader{
		Version:           node.Version,
		Height:            node.Height,
		PreviousBlockHash: previousBlockHash,
		Timestamp:         node.Timestamp,
		Nonce:             node.Nonce,
		Bits:              node.Bits,
		BlockCommitment: types.BlockCommitment{
			TransactionsMerkleRoot: node.TransactionsMerkleRoot,
			TransactionStatusHash:  node.TransactionStatusHash,
		},
	}
}

func (node *BlockNode) CalcPastMedianTime() uint64 {
	timestamps := []uint64{}
	iterNode := node
	for i := 0; i < consensus.MedianTimeBlocks && iterNode != nil; i++ {
		timestamps = append(timestamps, iterNode.Timestamp)
		iterNode = iterNode.Parent
	}

	sort.Sort(common.TimeSorter(timestamps))
	return timestamps[len(timestamps)/2]
}

// CalcNextBits calculate the bits for next block
func (node *BlockNode) CalcNextBits() uint64 {
	if node.Height%consensus.BlocksPerRetarget != 0 || node.Height == 0 {
		return node.Bits
	}

	compareNode := node.Parent
	for compareNode.Height%consensus.BlocksPerRetarget != 0 {
		compareNode = compareNode.Parent
	}
	return difficulty.CalcNextRequiredDifficulty(node.BlockHeader(), compareNode.BlockHeader())
}

// CalcNextSeed calculate the seed for next block
func (node *BlockNode) CalcNextSeed() *bc.Hash {
	if node.Height == 0 {
		return consensus.InitialSeed
	}
	if node.Height%consensus.SeedPerRetarget == 0 {
		return &node.Hash
	}
	return node.Seed
}

// BlockIndex is the struct for help chain trace block chain as tree
type BlockIndex struct {
	sync.RWMutex

	index     map[bc.Hash]*BlockNode
	mainChain []*BlockNode
}

// NewBlockIndex will create a empty BlockIndex
func NewBlockIndex() *BlockIndex {
	return &BlockIndex{
		index:     make(map[bc.Hash]*BlockNode),
		mainChain: make([]*BlockNode, 0, approxNodesPerDay),
	}
}

// AddNode will add node to the index map
func (bi *BlockIndex) AddNode(node *BlockNode) {
	bi.Lock()
	bi.index[node.Hash] = node
	bi.Unlock()
}

// GetNode will search node from the index map
func (bi *BlockIndex) GetNode(hash *bc.Hash) *BlockNode {
	bi.RLock()
	defer bi.RUnlock()
	return bi.index[*hash]
}

func (bi *BlockIndex) BestNode() *BlockNode {
	bi.RLock()
	defer bi.RUnlock()
	return bi.mainChain[len(bi.mainChain)-1]
}

// BlockExist check does the block existed in blockIndex
func (bi *BlockIndex) BlockExist(hash *bc.Hash) bool {
	bi.RLock()
	_, ok := bi.index[*hash]
	bi.RUnlock()
	return ok
}

// TODO: THIS FUNCTION MIGHT BE DELETED
func (bi *BlockIndex) InMainchain(hash bc.Hash) bool {
	bi.RLock()
	defer bi.RUnlock()

	node, ok := bi.index[hash]
	if !ok {
		return false
	}
	return bi.nodeByHeight(node.Height) == node
}

func (bi *BlockIndex) nodeByHeight(height uint64) *BlockNode {
	if height >= uint64(len(bi.mainChain)) {
		return nil
	}
	return bi.mainChain[height]
}

// NodeByHeight returns the block node at the specified height.
func (bi *BlockIndex) NodeByHeight(height uint64) *BlockNode {
	bi.RLock()
	defer bi.RUnlock()
	return bi.nodeByHeight(height)
}

// SetMainChain will set the the mainChain array
func (bi *BlockIndex) SetMainChain(node *BlockNode) {
	bi.Lock()
	defer bi.Unlock()

	needed := node.Height + 1
	if uint64(cap(bi.mainChain)) < needed {
		nodes := make([]*BlockNode, needed, needed+approxNodesPerDay)
		copy(nodes, bi.mainChain)
		bi.mainChain = nodes
	} else {
		i := uint64(len(bi.mainChain))
		bi.mainChain = bi.mainChain[0:needed]
		for ; i < needed; i++ {
			bi.mainChain[i] = nil
		}
	}

	for node != nil && bi.mainChain[node.Height] != node {
		bi.mainChain[node.Height] = node
		node = node.Parent
	}
}
<|endoftext|>


<commit_before>package ws

import (
	"net"
	"sync"
	"time"

	"github.com/v2ray/v2ray-core/common/log"
	"github.com/v2ray/v2ray-core/common/signal"
)

type AwaitingConnection struct {
	conn   *wsconn
	expire time.Time
}

func (this *AwaitingConnection) Expired() bool {
	return this.expire.Before(time.Now())
}

type ConnectionCache struct {
	sync.Mutex
	cache       map[string][]*AwaitingConnection
	cleanupOnce signal.Once
}

func NewConnectionCache() *ConnectionCache {
	return &ConnectionCache{
		cache: make(map[string][]*AwaitingConnection),
	}
}

func (this *ConnectionCache) Cleanup() {
	defer this.cleanupOnce.Reset()

	for len(this.cache) > 0 {
		time.Sleep(time.Second * 4)
		this.Lock()
		for key, value := range this.cache {
			size := len(value)
			changed := false
			for i := 0; i < size; {
				if value[i].Expired() {
					value[i].conn.Close()
					value[i] = value[size-1]
					size--
					changed = true
				} else {
					i++
				}
			}
			if changed {
				for i := size; i < len(value); i++ {
					value[i] = nil
				}
				value = value[:size]
				this.cache[key] = value
			}
		}
		this.Unlock()
	}
}

func (this *ConnectionCache) Recycle(dest string, conn *wsconn) {
	this.Lock()
	defer this.Unlock()

	aconn := &AwaitingConnection{
		conn:   conn,
		expire: time.Now().Add(time.Second * 180),
	}

	var list []*AwaitingConnection
	if v, found := this.cache[dest]; found {
		v = append(v, aconn)
		list = v
	} else {
		list = []*AwaitingConnection{aconn}
	}
	this.cache[dest] = list

	go this.cleanupOnce.Do(this.Cleanup)
}

func FindFirstValid(list []*AwaitingConnection) int {
	for idx, conn := range list {
		if !conn.Expired() && !conn.conn.connClosing {
			return idx
		}
		go conn.conn.Close()
	}
	return -1
}

func (this *ConnectionCache) Get(dest string) net.Conn {
	this.Lock()
	defer this.Unlock()

	list, found := this.cache[dest]
	if !found {
		return nil
	}

	firstValid := FindFirstValid(list)
	if firstValid == -1 {
		delete(this.cache, dest)
		return nil
	}
	res := list[firstValid].conn
	list = list[firstValid+1:]
	this.cache[dest] = list
	log.Debug("WS:Conn Cache used.")
	return res
}
<commit_msg>Set reuse limit<commit_after>package ws

import (
	"net"
	"sync"
	"time"

	"github.com/v2ray/v2ray-core/common/log"
	"github.com/v2ray/v2ray-core/common/signal"
)

type AwaitingConnection struct {
	conn   *wsconn
	expire time.Time
}

func (this *AwaitingConnection) Expired() bool {
	return this.expire.Before(time.Now())
}

type ConnectionCache struct {
	sync.Mutex
	cache       map[string][]*AwaitingConnection
	cleanupOnce signal.Once
}

func NewConnectionCache() *ConnectionCache {
	return &ConnectionCache{
		cache: make(map[string][]*AwaitingConnection),
	}
}

func (this *ConnectionCache) Cleanup() {
	defer this.cleanupOnce.Reset()

	for len(this.cache) > 0 {
		time.Sleep(time.Second * 4)
		this.Lock()
		for key, value := range this.cache {
			size := len(value)
			changed := false
			for i := 0; i < size; {
				if value[i].Expired() {
					value[i].conn.Close()
					value[i] = value[size-1]
					size--
					changed = true
				} else {
					i++
				}
			}
			if changed {
				for i := size; i < len(value); i++ {
					value[i] = nil
				}
				value = value[:size]
				this.cache[key] = value
			}
		}
		this.Unlock()
	}
}

func (this *ConnectionCache) Recycle(dest string, conn *wsconn) {
	this.Lock()
	defer this.Unlock()

	aconn := &AwaitingConnection{
		conn:   conn,
		expire: time.Now().Add(time.Second * 7),
	}

	var list []*AwaitingConnection
	if v, found := this.cache[dest]; found {
		v = append(v, aconn)
		list = v
	} else {
		list = []*AwaitingConnection{aconn}
	}
	this.cache[dest] = list

	go this.cleanupOnce.Do(this.Cleanup)
}

func FindFirstValid(list []*AwaitingConnection) int {
	for idx, conn := range list {
		if !conn.Expired() && !conn.conn.connClosing {
			return idx
		}
		go conn.conn.Close()
	}
	return -1
}

func (this *ConnectionCache) Get(dest string) net.Conn {
	this.Lock()
	defer this.Unlock()

	list, found := this.cache[dest]
	if !found {
		return nil
	}

	firstValid := FindFirstValid(list)
	if firstValid == -1 {
		delete(this.cache, dest)
		return nil
	}
	res := list[firstValid].conn
	list = list[firstValid+1:]
	this.cache[dest] = list
	log.Debug("WS:Conn Cache used.")
	return res
}
<|endoftext|>


<commit_before>package tracetcp

import (
	"fmt"
	"net"
	"syscall"
	"time"
)

func MakeTimeval(t time.Duration) syscall.Timeval {
	return syscall.NsecToTimeval(int64(t))
}

func FD_SET(p *syscall.FdSet, i int) {
	p.Bits[i/64] |= 1 << uint(i) % 64
}

func FD_ISSET(p *syscall.FdSet, i int) bool {
	return (p.Bits[i/64] & (1 << uint(i) % 64)) != 0
}

func FD_ZERO(p *syscall.FdSet) {
	for i := range p.Bits {
		p.Bits[i] = 0
	}
}

func LookupAddress(host string) ([4]byte, error) {
	var addr [4]byte

	addresses, err := net.LookupHost(host)
	if err != nil {
		return addr, err
	}

	ip, err := net.ResolveIPAddr("ip", addresses[0])
	if err != nil {
		return addr, err
	}

	copy(addr[:], ip.IP.To4())
	return addr, nil
}

func ToSockaddrInet4(ip net.IPAddr, port int) *syscall.SockaddrInet4 {
	var addr [4]byte
	copy(addr[:], ip.IP.To4())

	return &syscall.SockaddrInet4{Port: port, Addr: addr}
}

func ToIPAddrAndPort(saddr syscall.Sockaddr) (addr net.IPAddr, port int, err error) {

	if sa, ok := saddr.(*syscall.SockaddrInet4); ok {
		port = sa.Port
		addr.IP = append(addr.IP, sa.Addr[:]...)
	} else {
		err = fmt.Errorf("%s", "ToIPAddrAndPort: syscall.Sockaddr not a syscall.SockaddeInet4")
	}

	return
}
<commit_msg>building impl<commit_after>package tracetcp

import (
	"fmt"
	"net"
	"syscall"
	"time"
)

func MakeTimeval(t time.Duration) syscall.Timeval {
	return syscall.NsecToTimeval(int64(t))
}

func FD_SET(p *syscall.FdSet, i int) {
	p.Bits[i/64] |= 1 << uint(i) % 64
}

func FD_ISSET(p *syscall.FdSet, i int) bool {
	return (p.Bits[i/64] & (1 << uint(i) % 64)) != 0
}

func FD_ZERO(p *syscall.FdSet) {
	for i := range p.Bits {
		p.Bits[i] = 0
	}
}

func LookupAddress(host string) (*net.IPAddr, error) {
	addresses, err := net.LookupHost(host)
	if err != nil {
		return &net.IPAddr{}, err
	}

	ip, err := net.ResolveIPAddr("ip", addresses[0])
	if err != nil {
		return &net.IPAddr{}, err
	}
	return ip, nil
}

func ToSockaddrInet4(ip net.IPAddr, port int) *syscall.SockaddrInet4 {
	var addr [4]byte
	copy(addr[:], ip.IP.To4())

	return &syscall.SockaddrInet4{Port: port, Addr: addr}
}

func ToIPAddrAndPort(saddr syscall.Sockaddr) (addr net.IPAddr, port int, err error) {

	if sa, ok := saddr.(*syscall.SockaddrInet4); ok {
		port = sa.Port
		addr.IP = append(addr.IP, sa.Addr[:]...)
	} else {
		err = fmt.Errorf("%s", "ToIPAddrAndPort: syscall.Sockaddr not a syscall.SockaddeInet4")
	}

	return
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/base64"
	"encoding/json"
	"flag"
	"io"
	"io/ioutil"
	"log"
	"net"

	"golang.org/x/crypto/ssh"
)

type source int

const (
	client source = iota
	server
)

func (src source) String() string {
	switch src {
	case client:
		return "client"
	case server:
		return "server"
	default:
		return "unknown"
	}
}

func (src source) MarshalJSON() ([]byte, error) {
	return json.Marshal(src.String())
}

type channelLog struct {
	ChannelID int `json:"channel_id"`
}

type requestLog struct {
	Type      string `json:"type"`
	WantReply bool   `json:"want_reply"`
	Payload   string `json:"payload"`

	Accepted bool `json:"accepted"`
}

type logEntry interface {
	eventType() string
}

type globalRequestLog struct {
	requestLog

	Response string `json:"response"`
}

func (entry globalRequestLog) eventType() string {
	return "global_request"
}

type newChannelLog struct {
	Type      string `json:"type"`
	ExtraData string `json:"extra_data"`

	Accepted bool `json:"accepted"`
}

func (entry newChannelLog) eventType() string {
	return "new_channel"
}

type channelRequestLog struct {
	channelLog
	requestLog
}

func (entry channelRequestLog) eventType() string {
	return "channel_request"
}

type channelDataLog struct {
	channelLog
	Data string `json:"data"`
}

func (entry channelDataLog) eventType() string {
	return "channel_data"
}

type channelErrorLog struct {
	channelLog
	Data string `json:"data"`
}

func (entry channelErrorLog) eventType() string {
	return "channel_error"
}

type channelEOFLog struct {
	channelLog
}

func (entry channelEOFLog) eventType() string {
	return "channel_eof"
}

type channelCloseLog struct {
	channelLog
}

func (entry channelCloseLog) eventType() string {
	return "channel_close"
}

type connectionCloseLog struct{}

func (entry connectionCloseLog) eventType() string {
	return "connection_close"
}

func logEvent(entry logEntry, src source) {
	jsonBytes, err := json.Marshal(struct {
		Source    string   `json:"source"`
		EventType string   `json:"event_type"`
		Event     logEntry `json:"event"`
	}{
		Source:    src.String(),
		EventType: entry.eventType(),
		Event:     entry,
	})
	if err != nil {
		panic(err)
	}
	log.Printf("%s", jsonBytes)
}

func streamReader(reader io.Reader) <-chan string {
	input := make(chan string)
	go func() {
		defer close(input)
		buffer := make([]byte, 256)
		for {
			n, err := reader.Read(buffer)
			if n > 0 {
				input <- string(buffer[:n])
			}
			if err != nil {
				if err != io.EOF {
					panic(err)
				}
				return
			}
		}
	}()
	return input
}

func handleChannel(channelID int, clientChannel ssh.Channel, clientRequests <-chan *ssh.Request, serverChannel ssh.Channel, serverRequests <-chan *ssh.Request) {
	clientInputStream := streamReader(clientChannel)
	serverInputStream := streamReader(serverChannel)
	serverErrorStream := streamReader(serverChannel.Stderr())

	for clientInputStream != nil || clientRequests != nil || serverInputStream != nil || serverRequests != nil {
		select {
		case clientInput, ok := <-clientInputStream:
			if !ok {
				if serverInputStream != nil {
					logEvent(channelEOFLog{
						channelLog: channelLog{
							ChannelID: channelID,
						},
					}, client)
					if err := serverChannel.CloseWrite(); err != nil {
						panic(err)
					}
				}
				clientInputStream = nil
				continue
			}
			logEvent(channelDataLog{
				channelLog: channelLog{
					ChannelID: channelID,
				},
				Data: clientInput,
			}, client)
			if _, err := serverChannel.Write([]byte(clientInput)); err != nil {
				panic(err)
			}
		case clientRequest, ok := <-clientRequests:
			if !ok {
				if serverRequests != nil {
					logEvent(channelCloseLog{
						channelLog: channelLog{
							ChannelID: channelID,
						},
					}, client)
					if err := serverChannel.Close(); err != nil {
						panic(err)
					}
				}
				clientRequests = nil
				continue
			}
			accepted, err := serverChannel.SendRequest(clientRequest.Type, clientRequest.WantReply, clientRequest.Payload)
			if err != nil {
				panic(err)
			}
			logEvent(channelRequestLog{
				channelLog: channelLog{
					ChannelID: channelID,
				},
				requestLog: requestLog{
					Type:      clientRequest.Type,
					WantReply: clientRequest.WantReply,
					Payload:   base64.RawStdEncoding.EncodeToString(clientRequest.Payload),
					Accepted:  accepted,
				},
			}, client)
			if clientRequest.WantReply {
				if err := clientRequest.Reply(accepted, nil); err != nil {
					panic(err)
				}
			}
		case serverInput, ok := <-serverInputStream:
			if !ok {
				if clientInputStream != nil {
					logEvent(channelEOFLog{
						channelLog: channelLog{
							ChannelID: channelID,
						},
					}, server)
					if err := clientChannel.CloseWrite(); err != nil {
						panic(err)
					}
				}
				serverInputStream = nil
				continue
			}
			logEvent(channelDataLog{
				channelLog: channelLog{
					ChannelID: channelID,
				},
				Data: serverInput,
			}, server)
			if _, err := clientChannel.Write([]byte(serverInput)); err != nil {
				panic(err)
			}
		case serverError, ok := <-serverErrorStream:
			if !ok {
				serverErrorStream = nil
				continue
			}
			logEvent(channelErrorLog{
				channelLog: channelLog{
					ChannelID: channelID,
				},
				Data: serverError,
			}, server)
			if _, err := clientChannel.Stderr().Write([]byte(serverError)); err != nil {
				panic(err)
			}
		case serverRequest, ok := <-serverRequests:
			if !ok {
				if clientRequests != nil {
					logEvent(channelCloseLog{
						channelLog: channelLog{
							ChannelID: channelID,
						},
					}, server)
					if err := clientChannel.Close(); err != nil {
						panic(err)
					}
				}
				serverRequests = nil
				continue
			}
			accepted, err := clientChannel.SendRequest(serverRequest.Type, serverRequest.WantReply, serverRequest.Payload)
			if err != nil {
				panic(err)
			}
			logEvent(channelRequestLog{
				channelLog: channelLog{
					ChannelID: channelID,
				},
				requestLog: requestLog{
					Type:      serverRequest.Type,
					WantReply: serverRequest.WantReply,
					Payload:   base64.RawStdEncoding.EncodeToString(serverRequest.Payload),
					Accepted:  accepted,
				},
			}, server)
			if serverRequest.WantReply {
				if err := serverRequest.Reply(accepted, nil); err != nil {
					panic(err)
				}
			}
		}
	}
}

func handleConn(clientConn net.Conn, sshServerConfig *ssh.ServerConfig, serverAddress string, clientKey ssh.Signer) {
	clientSSHConn, clientNewChannels, clientRequests, err := ssh.NewServerConn(clientConn, sshServerConfig)
	if err != nil {
		panic(err)
	}

	serverConn, err := net.Dial("tcp", serverAddress)
	if err != nil {
		panic(err)
	}

	serverSSHConn, serverNewChannels, serverRequests, err := ssh.NewClientConn(serverConn, serverAddress, &ssh.ClientConfig{
		User:            clientSSHConn.User(),
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
		Auth: []ssh.AuthMethod{
			ssh.PublicKeys(clientKey),
		},
		ClientVersion: "SSH-2.0-OpenSSH_7.2",
	})
	if err != nil {
		panic(err)
	}

	channelID := 0

	for clientNewChannels != nil || clientRequests != nil || serverNewChannels != nil || serverRequests != nil {
		select {
		case clientNewChannel, ok := <-clientNewChannels:
			if !ok {
				clientNewChannels = nil
				if serverNewChannels != nil {
					logEvent(connectionCloseLog{}, client)
					if err := serverSSHConn.Close(); err != nil {
						panic(err)
					}
				}
				continue
			}
			serverChannel, serverChannelRequests, err := serverSSHConn.OpenChannel(clientNewChannel.ChannelType(), clientNewChannel.ExtraData())
			logEvent(newChannelLog{
				Type:      clientNewChannel.ChannelType(),
				ExtraData: base64.RawStdEncoding.EncodeToString(clientNewChannel.ExtraData()),
				Accepted:  err == nil,
			}, client)
			if err != nil {
				if err, ok := err.(*ssh.OpenChannelError); ok {
					if err := clientNewChannel.Reject(err.Reason, err.Message); err != nil {
						panic(err)
					}
					continue
				}
				panic(err)
			}
			clientChannel, clientChannelRequests, err := clientNewChannel.Accept()
			if err != nil {
				panic(err)
			}
			go handleChannel(channelID, clientChannel, clientChannelRequests, serverChannel, serverChannelRequests)
		case clientRequest, ok := <-clientRequests:
			if !ok {
				clientRequests = nil
				continue
			}
			if clientRequest.Type == "no-more-sessions@openssh.com" {
				logEvent(globalRequestLog{
					requestLog: requestLog{
						Type:      clientRequest.Type,
						WantReply: clientRequest.WantReply,
						Payload:   base64.RawStdEncoding.EncodeToString(clientRequest.Payload),
						Accepted:  clientRequest.WantReply,
					},
					Response: base64.RawStdEncoding.EncodeToString([]byte{}),
				}, client)
				continue
			}
			accepted, response, err := serverSSHConn.SendRequest(clientRequest.Type, clientRequest.WantReply, clientRequest.Payload)
			if err != nil {
				panic(err)
			}
			logEvent(globalRequestLog{
				requestLog: requestLog{
					Type:      clientRequest.Type,
					WantReply: clientRequest.WantReply,
					Payload:   base64.RawStdEncoding.EncodeToString(clientRequest.Payload),
					Accepted:  accepted,
				},
				Response: base64.RawStdEncoding.EncodeToString(response),
			}, client)
			if err := clientRequest.Reply(accepted, response); err != nil {
				panic(err)
			}
		case serverNewChannel, ok := <-serverNewChannels:
			if !ok {
				if clientNewChannels != nil {
					logEvent(connectionCloseLog{}, server)
					if err := clientSSHConn.Close(); err != nil {
						panic(err)
					}
				}
				serverNewChannels = nil
				continue
			}
			panic(serverNewChannel.ChannelType())
		case serverRequest, ok := <-serverRequests:
			if !ok {
				serverRequests = nil
				continue
			}
			accepted, response, err := clientSSHConn.SendRequest(serverRequest.Type, serverRequest.WantReply, serverRequest.Payload)
			logEvent(globalRequestLog{
				requestLog: requestLog{
					Type:      serverRequest.Type,
					WantReply: serverRequest.WantReply,
					Payload:   base64.RawStdEncoding.EncodeToString(serverRequest.Payload),
					Accepted:  accepted,
				},
				Response: base64.RawStdEncoding.EncodeToString(response),
			}, server)
			if err != nil {
				panic(err)
			}
			if err := serverRequest.Reply(accepted, response); err != nil {
				panic(err)
			}
		}
		channelID++
	}
}

func main() {
	listenAddress := flag.String("listen_address", "127.0.0.1:2022", "listen address")
	hostKeyFile := flag.String("host_key_file", "", "host key file")
	serverAddress := flag.String("server_address", "127.0.0.1:22", "server address")
	clientKeyFile := flag.String("client_key_file", "", "client key file")
	flag.Parse()
	if *listenAddress == "" {
		panic("listen address is required")
	}
	if *hostKeyFile == "" {
		panic("host key file is required")
	}
	if *serverAddress == "" {
		panic("server address is required")
	}
	if *clientKeyFile == "" {
		panic("client key file is required")
	}

	log.SetFlags(0)

	serverConfig := &ssh.ServerConfig{
		NoClientAuth:  true,
		ServerVersion: "SSH-2.0-OpenSSH_7.2",
	}
	hostKeyBytes, err := ioutil.ReadFile(*hostKeyFile)
	if err != nil {
		panic(err)
	}
	hostKey, err := ssh.ParsePrivateKey(hostKeyBytes)
	if err != nil {
		panic(err)
	}
	serverConfig.AddHostKey(hostKey)

	clientKeyBytes, err := ioutil.ReadFile(*clientKeyFile)
	if err != nil {
		panic(err)
	}
	clientKey, err := ssh.ParsePrivateKey(clientKeyBytes)
	if err != nil {
		panic(err)
	}

	listener, err := net.Listen("tcp", *listenAddress)
	if err != nil {
		panic(err)
	}
	defer listener.Close()
	for {
		conn, err := listener.Accept()
		if err != nil {
			panic(err)
		}
		go handleConn(conn, serverConfig, *serverAddress, clientKey)
	}
}
<commit_msg>testproxy: log to stdout<commit_after>package main

import (
	"encoding/base64"
	"encoding/json"
	"flag"
	"io"
	"io/ioutil"
	"log"
	"net"

	"golang.org/x/crypto/ssh"
)

type source int

const (
	client source = iota
	server
)

func (src source) String() string {
	switch src {
	case client:
		return "client"
	case server:
		return "server"
	default:
		return "unknown"
	}
}

func (src source) MarshalJSON() ([]byte, error) {
	return json.Marshal(src.String())
}

type channelLog struct {
	ChannelID int `json:"channel_id"`
}

type requestLog struct {
	Type      string `json:"type"`
	WantReply bool   `json:"want_reply"`
	Payload   string `json:"payload"`

	Accepted bool `json:"accepted"`
}

type logEntry interface {
	eventType() string
}

type globalRequestLog struct {
	requestLog

	Response string `json:"response"`
}

func (entry globalRequestLog) eventType() string {
	return "global_request"
}

type newChannelLog struct {
	Type      string `json:"type"`
	ExtraData string `json:"extra_data"`

	Accepted bool `json:"accepted"`
}

func (entry newChannelLog) eventType() string {
	return "new_channel"
}

type channelRequestLog struct {
	channelLog
	requestLog
}

func (entry channelRequestLog) eventType() string {
	return "channel_request"
}

type channelDataLog struct {
	channelLog
	Data string `json:"data"`
}

func (entry channelDataLog) eventType() string {
	return "channel_data"
}

type channelErrorLog struct {
	channelLog
	Data string `json:"data"`
}

func (entry channelErrorLog) eventType() string {
	return "channel_error"
}

type channelEOFLog struct {
	channelLog
}

func (entry channelEOFLog) eventType() string {
	return "channel_eof"
}

type channelCloseLog struct {
	channelLog
}

func (entry channelCloseLog) eventType() string {
	return "channel_close"
}

type connectionCloseLog struct{}

func (entry connectionCloseLog) eventType() string {
	return "connection_close"
}

func logEvent(entry logEntry, src source) {
	jsonBytes, err := json.Marshal(struct {
		Source    string   `json:"source"`
		EventType string   `json:"event_type"`
		Event     logEntry `json:"event"`
	}{
		Source:    src.String(),
		EventType: entry.eventType(),
		Event:     entry,
	})
	if err != nil {
		panic(err)
	}
	log.Printf("%s", jsonBytes)
}

func streamReader(reader io.Reader) <-chan string {
	input := make(chan string)
	go func() {
		defer close(input)
		buffer := make([]byte, 256)
		for {
			n, err := reader.Read(buffer)
			if n > 0 {
				input <- string(buffer[:n])
			}
			if err != nil {
				if err != io.EOF {
					panic(err)
				}
				return
			}
		}
	}()
	return input
}

func handleChannel(channelID int, clientChannel ssh.Channel, clientRequests <-chan *ssh.Request, serverChannel ssh.Channel, serverRequests <-chan *ssh.Request) {
	clientInputStream := streamReader(clientChannel)
	serverInputStream := streamReader(serverChannel)
	serverErrorStream := streamReader(serverChannel.Stderr())

	for clientInputStream != nil || clientRequests != nil || serverInputStream != nil || serverRequests != nil {
		select {
		case clientInput, ok := <-clientInputStream:
			if !ok {
				if serverInputStream != nil {
					logEvent(channelEOFLog{
						channelLog: channelLog{
							ChannelID: channelID,
						},
					}, client)
					if err := serverChannel.CloseWrite(); err != nil {
						panic(err)
					}
				}
				clientInputStream = nil
				continue
			}
			logEvent(channelDataLog{
				channelLog: channelLog{
					ChannelID: channelID,
				},
				Data: clientInput,
			}, client)
			if _, err := serverChannel.Write([]byte(clientInput)); err != nil {
				panic(err)
			}
		case clientRequest, ok := <-clientRequests:
			if !ok {
				if serverRequests != nil {
					logEvent(channelCloseLog{
						channelLog: channelLog{
							ChannelID: channelID,
						},
					}, client)
					if err := serverChannel.Close(); err != nil {
						panic(err)
					}
				}
				clientRequests = nil
				continue
			}
			accepted, err := serverChannel.SendRequest(clientRequest.Type, clientRequest.WantReply, clientRequest.Payload)
			if err != nil {
				panic(err)
			}
			logEvent(channelRequestLog{
				channelLog: channelLog{
					ChannelID: channelID,
				},
				requestLog: requestLog{
					Type:      clientRequest.Type,
					WantReply: clientRequest.WantReply,
					Payload:   base64.RawStdEncoding.EncodeToString(clientRequest.Payload),
					Accepted:  accepted,
				},
			}, client)
			if clientRequest.WantReply {
				if err := clientRequest.Reply(accepted, nil); err != nil {
					panic(err)
				}
			}
		case serverInput, ok := <-serverInputStream:
			if !ok {
				if clientInputStream != nil {
					logEvent(channelEOFLog{
						channelLog: channelLog{
							ChannelID: channelID,
						},
					}, server)
					if err := clientChannel.CloseWrite(); err != nil {
						panic(err)
					}
				}
				serverInputStream = nil
				continue
			}
			logEvent(channelDataLog{
				channelLog: channelLog{
					ChannelID: channelID,
				},
				Data: serverInput,
			}, server)
			if _, err := clientChannel.Write([]byte(serverInput)); err != nil {
				panic(err)
			}
		case serverError, ok := <-serverErrorStream:
			if !ok {
				serverErrorStream = nil
				continue
			}
			logEvent(channelErrorLog{
				channelLog: channelLog{
					ChannelID: channelID,
				},
				Data: serverError,
			}, server)
			if _, err := clientChannel.Stderr().Write([]byte(serverError)); err != nil {
				panic(err)
			}
		case serverRequest, ok := <-serverRequests:
			if !ok {
				if clientRequests != nil {
					logEvent(channelCloseLog{
						channelLog: channelLog{
							ChannelID: channelID,
						},
					}, server)
					if err := clientChannel.Close(); err != nil {
						panic(err)
					}
				}
				serverRequests = nil
				continue
			}
			accepted, err := clientChannel.SendRequest(serverRequest.Type, serverRequest.WantReply, serverRequest.Payload)
			if err != nil {
				panic(err)
			}
			logEvent(channelRequestLog{
				channelLog: channelLog{
					ChannelID: channelID,
				},
				requestLog: requestLog{
					Type:      serverRequest.Type,
					WantReply: serverRequest.WantReply,
					Payload:   base64.RawStdEncoding.EncodeToString(serverRequest.Payload),
					Accepted:  accepted,
				},
			}, server)
			if serverRequest.WantReply {
				if err := serverRequest.Reply(accepted, nil); err != nil {
					panic(err)
				}
			}
		}
	}
}

func handleConn(clientConn net.Conn, sshServerConfig *ssh.ServerConfig, serverAddress string, clientKey ssh.Signer) {
	clientSSHConn, clientNewChannels, clientRequests, err := ssh.NewServerConn(clientConn, sshServerConfig)
	if err != nil {
		panic(err)
	}

	serverConn, err := net.Dial("tcp", serverAddress)
	if err != nil {
		panic(err)
	}

	serverSSHConn, serverNewChannels, serverRequests, err := ssh.NewClientConn(serverConn, serverAddress, &ssh.ClientConfig{
		User:            clientSSHConn.User(),
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
		Auth: []ssh.AuthMethod{
			ssh.PublicKeys(clientKey),
		},
		ClientVersion: "SSH-2.0-OpenSSH_7.2",
	})
	if err != nil {
		panic(err)
	}

	channelID := 0

	for clientNewChannels != nil || clientRequests != nil || serverNewChannels != nil || serverRequests != nil {
		select {
		case clientNewChannel, ok := <-clientNewChannels:
			if !ok {
				clientNewChannels = nil
				if serverNewChannels != nil {
					logEvent(connectionCloseLog{}, client)
					if err := serverSSHConn.Close(); err != nil {
						panic(err)
					}
				}
				continue
			}
			serverChannel, serverChannelRequests, err := serverSSHConn.OpenChannel(clientNewChannel.ChannelType(), clientNewChannel.ExtraData())
			logEvent(newChannelLog{
				Type:      clientNewChannel.ChannelType(),
				ExtraData: base64.RawStdEncoding.EncodeToString(clientNewChannel.ExtraData()),
				Accepted:  err == nil,
			}, client)
			if err != nil {
				if err, ok := err.(*ssh.OpenChannelError); ok {
					if err := clientNewChannel.Reject(err.Reason, err.Message); err != nil {
						panic(err)
					}
					continue
				}
				panic(err)
			}
			clientChannel, clientChannelRequests, err := clientNewChannel.Accept()
			if err != nil {
				panic(err)
			}
			go handleChannel(channelID, clientChannel, clientChannelRequests, serverChannel, serverChannelRequests)
		case clientRequest, ok := <-clientRequests:
			if !ok {
				clientRequests = nil
				continue
			}
			if clientRequest.Type == "no-more-sessions@openssh.com" {
				logEvent(globalRequestLog{
					requestLog: requestLog{
						Type:      clientRequest.Type,
						WantReply: clientRequest.WantReply,
						Payload:   base64.RawStdEncoding.EncodeToString(clientRequest.Payload),
						Accepted:  clientRequest.WantReply,
					},
					Response: base64.RawStdEncoding.EncodeToString([]byte{}),
				}, client)
				continue
			}
			accepted, response, err := serverSSHConn.SendRequest(clientRequest.Type, clientRequest.WantReply, clientRequest.Payload)
			if err != nil {
				panic(err)
			}
			logEvent(globalRequestLog{
				requestLog: requestLog{
					Type:      clientRequest.Type,
					WantReply: clientRequest.WantReply,
					Payload:   base64.RawStdEncoding.EncodeToString(clientRequest.Payload),
					Accepted:  accepted,
				},
				Response: base64.RawStdEncoding.EncodeToString(response),
			}, client)
			if err := clientRequest.Reply(accepted, response); err != nil {
				panic(err)
			}
		case serverNewChannel, ok := <-serverNewChannels:
			if !ok {
				if clientNewChannels != nil {
					logEvent(connectionCloseLog{}, server)
					if err := clientSSHConn.Close(); err != nil {
						panic(err)
					}
				}
				serverNewChannels = nil
				continue
			}
			panic(serverNewChannel.ChannelType())
		case serverRequest, ok := <-serverRequests:
			if !ok {
				serverRequests = nil
				continue
			}
			accepted, response, err := clientSSHConn.SendRequest(serverRequest.Type, serverRequest.WantReply, serverRequest.Payload)
			logEvent(globalRequestLog{
				requestLog: requestLog{
					Type:      serverRequest.Type,
					WantReply: serverRequest.WantReply,
					Payload:   base64.RawStdEncoding.EncodeToString(serverRequest.Payload),
					Accepted:  accepted,
				},
				Response: base64.RawStdEncoding.EncodeToString(response),
			}, server)
			if err != nil {
				panic(err)
			}
			if err := serverRequest.Reply(accepted, response); err != nil {
				panic(err)
			}
		}
		channelID++
	}
}

func main() {
	listenAddress := flag.String("listen_address", "127.0.0.1:2022", "listen address")
	hostKeyFile := flag.String("host_key_file", "", "host key file")
	serverAddress := flag.String("server_address", "127.0.0.1:22", "server address")
	clientKeyFile := flag.String("client_key_file", "", "client key file")
	flag.Parse()
	if *listenAddress == "" {
		panic("listen address is required")
	}
	if *hostKeyFile == "" {
		panic("host key file is required")
	}
	if *serverAddress == "" {
		panic("server address is required")
	}
	if *clientKeyFile == "" {
		panic("client key file is required")
	}

	log.SetFlags(0)
	log.SetOutput(os.Stdout)

	serverConfig := &ssh.ServerConfig{
		NoClientAuth:  true,
		ServerVersion: "SSH-2.0-OpenSSH_7.2",
	}
	hostKeyBytes, err := ioutil.ReadFile(*hostKeyFile)
	if err != nil {
		panic(err)
	}
	hostKey, err := ssh.ParsePrivateKey(hostKeyBytes)
	if err != nil {
		panic(err)
	}
	serverConfig.AddHostKey(hostKey)

	clientKeyBytes, err := ioutil.ReadFile(*clientKeyFile)
	if err != nil {
		panic(err)
	}
	clientKey, err := ssh.ParsePrivateKey(clientKeyBytes)
	if err != nil {
		panic(err)
	}

	listener, err := net.Listen("tcp", *listenAddress)
	if err != nil {
		panic(err)
	}
	defer listener.Close()
	for {
		conn, err := listener.Accept()
		if err != nil {
			panic(err)
		}
		go handleConn(conn, serverConfig, *serverAddress, clientKey)
	}
}
<|endoftext|>


<commit_before>// Copyright (c) 2014 David R. Jenni. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

/*
agodoc is a wrapper around godoc for use with Acme.
It shows the documentation of the identifier under the cursor.
*/
package main

import (
	"bytes"
	"errors"
	"fmt"
	"go/ast"
	"go/build"
	"go/parser"
	"go/token"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"

	"9fans.net/go/acme"
	"golang.org/x/tools/go/loader"
	"golang.org/x/tools/go/types"
)

type bodyReader struct{ *acme.Win }

func (r bodyReader) Read(data []byte) (int, error) {
	return r.Win.Read("body", data)
}

func main() {
	win, err := openWin()
	if err != nil {
		fmt.Fprintf(os.Stderr, "cannot open window: %v\n", err)
		os.Exit(1)
	}
	buf, err := ioutil.ReadAll(&bodyReader{win})
	if err != nil {
		fmt.Fprintf(os.Stderr, "cannot read content: %v\n", err)
		os.Exit(1)
	}
	filename, off, err := selection(win, buf)
	if err != nil {
		fmt.Fprintf(os.Stderr, "cannot get selection: %v\n", err)
		os.Exit(1)
	}

	fileInfos, err := ioutil.ReadDir(".")
	if err != nil {
		fmt.Fprintf(os.Stderr, "cannot read directory: %v\n", err)
		os.Exit(1)
	}
	var filenames []string
	for _, f := range fileInfos {
		if strings.HasSuffix(f.Name(), ".go") && f.Name() != filename {
			filenames = append(filenames, f.Name())
		}
	}
	path, ident, err := searchAtOff(off, string(buf), filenames...)
	if err != nil {
		fmt.Fprintf(os.Stderr, "cannot find identifier: %v\n", err)
		os.Exit(1)
	}
	if ident != "" {
		godoc(path, ident)
	} else {
		godoc(path)
	}
}

func openWin() (*acme.Win, error) {
	id, err := strconv.Atoi(os.Getenv("winid"))
	if err != nil {
		return nil, err
	}
	return acme.Open(id, nil)
}

func selection(win *acme.Win, buf []byte) (filename string, off int, err error) {
	filename, err = readFilename(win)
	if err != nil {
		return "", 0, err
	}
	filename = filepath.Base(filename)
	q0, _, err := readAddr(win)
	if err != nil {
		return "", 0, err
	}
	off, err = byteOffset(bytes.NewReader(buf), q0)
	if err != nil {
		return "", 0, err
	}
	return
}

func readFilename(win *acme.Win) (string, error) {
	b, err := win.ReadAll("tag")
	if err != nil {
		return "", err
	}
	tag := string(b)
	i := strings.Index(tag, " ")
	if i == -1 {
		return "", fmt.Errorf("cannot get filename from tag")
	}
	return tag[0:i], nil
}

func readAddr(win *acme.Win) (q0, q1 int, err error) {
	if _, _, err := win.ReadAddr(); err != nil {
		return 0, 0, err
	}
	if err := win.Ctl("addr=dot"); err != nil {
		return 0, 0, err
	}
	return win.ReadAddr()
}

func byteOffset(r io.RuneReader, off int) (bo int, err error) {
	for i := 0; i != off; i++ {
		_, s, err := r.ReadRune()
		if err != nil {
			return 0, err
		}
		bo += s
	}
	return
}

// searchAtOff returns the path and the identifier name at the given offset.
func searchAtOff(off int, src string, filenames ...string) (string, string, error) {
	var conf loader.Config
	var files []*ast.File

	for _, name := range filenames {
		f, err := parseFile(&conf, name, nil)
		if err != nil {
			return "", "", err
		}
		files = append(files, f)
	}
	f, err := parseFile(&conf, "", src)
	if err != nil {
		return "", "", err
	}
	files = append(files, f)

	for _, imp := range f.Imports {
		if isAtOff(conf.Fset, off, imp) {
			return strings.Trim(imp.Path.Value, "\""), "", nil
		}
	}

	ident := identAtOffset(conf.Fset, f, off)
	if ident == nil {
		return "", "", errors.New("no identifier here")
	}

	conf.CreateFromFiles("", files...)
	allowErrors(&conf)
	prg, err := conf.Load()
	if err != nil {
		return "", "", err
	}
	info := prg.Created[0]
	if obj := info.Uses[ident]; obj != nil {
		return fromObj(obj, conf.Fset, f, off)
	}
	if obj := info.Defs[ident]; obj != nil {
		return fromObj(obj, conf.Fset, f, off)
	}
	return "", "", errors.New("could not find identifier")
}

func allowErrors(conf *loader.Config) {
	ctxt := build.Default
	ctxt.CgoEnabled = false
	conf.Build = &ctxt
	conf.AllowErrors = true
	conf.ParserMode = parser.AllErrors
	conf.TypeChecker.Error = func(err error) {}
}

func parseFile(conf *loader.Config, name string, src interface{}) (*ast.File, error) {
	f, err := conf.ParseFile(name, src)
	if err != nil {
		return nil, err
	}
	if strings.HasSuffix(f.Name.Name, "_test") {
		f.Name.Name = f.Name.Name[:len(f.Name.Name)-len("_test")]
	}
	return f, nil
}

// currImportDir returns the current import directory.
func currImportDir() (string, error) {
	path, err := os.Getwd()
	if err != nil {
		return "", err
	}
	pkg, err := build.ImportDir(path, build.ImportComment)
	if err != nil {
		return "", err
	}
	return pkg.ImportPath, nil
}

// fromObj returns the path and the identifier name of an object.
func fromObj(obj types.Object, fset *token.FileSet, f *ast.File, off int) (string, string, error) {
	switch o := obj.(type) {
	case *types.Builtin:
		return "builtin", o.Name(), nil
	case *types.PkgName:
		return o.Imported().Path(), "", nil
	case *types.Const, *types.Func, *types.Nil, *types.TypeName, *types.Var:
		if obj.Pkg() == nil {
			return "builtin", obj.Name(), nil
		}
		if !isImported(fset, f, off) {
			impDir, err := currImportDir()
			return impDir, obj.Name(), err
		}
		return obj.Pkg().Path(), obj.Name(), nil
	default:
		return "", "", fmt.Errorf("cannot print documentation of %v\n", obj)
	}
}

// isImported returns true whether the identifier at the given offset is imported or not.
func isImported(fset *token.FileSet, f *ast.File, off int) (isImp bool) {
	ast.Inspect(f, func(node ast.Node) bool {
		switch s := node.(type) {
		case *ast.SelectorExpr:
			if isAtOff(fset, off, s.Sel) {
				if ident, ok := s.X.(*ast.Ident); ok {
					for _, imp := range f.Imports {
						if imp.Name != nil && imp.Name.Name == ident.Name {
							isImp = true
							return false
						}
						parts := strings.Split(strings.Trim(imp.Path.Value, "\""), "/")
						if parts[len(parts)-1] == ident.Name {
							isImp = true
						}
					}
				}
				return false
			}
		}
		return true
	})
	return isImp
}

// identAtOffset returns the identifier at an offset in a file.
func identAtOffset(fset *token.FileSet, f *ast.File, off int) (ident *ast.Ident) {
	ast.Inspect(f, func(node ast.Node) bool {
		switch n := node.(type) {
		case *ast.Ident:
			if isAtOff(fset, off, node) {
				ident = n
				return false
			}
		}
		return true
	})
	return ident
}

// isAtOff reports whether a node is at a given offset in a file set.
func isAtOff(fset *token.FileSet, off int, n ast.Node) bool {
	return fset.Position(n.Pos()).Offset <= off && off <= fset.Position(n.End()).Offset
}

// godoc runs the godoc command with the given arguments.
func godoc(args ...string) {
	c := exec.Command("godoc", args...)
	c.Stderr, c.Stdout = os.Stderr, os.Stderr
	if err := c.Run(); err != nil {
		fmt.Fprintf(os.Stderr, "godoc failed: %v\n", err)
		os.Exit(1)
	}
}
<commit_msg>Use new go/types import path.<commit_after>// Copyright (c) 2014 David R. Jenni. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

/*
agodoc is a wrapper around godoc for use with Acme.
It shows the documentation of the identifier under the cursor.
*/
package main

import (
	"bytes"
	"errors"
	"fmt"
	"go/ast"
	"go/build"
	"go/parser"
	"go/token"
	"go/types"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"

	"9fans.net/go/acme"
	"golang.org/x/tools/go/loader"
)

type bodyReader struct{ *acme.Win }

func (r bodyReader) Read(data []byte) (int, error) {
	return r.Win.Read("body", data)
}

func main() {
	win, err := openWin()
	if err != nil {
		fmt.Fprintf(os.Stderr, "cannot open window: %v\n", err)
		os.Exit(1)
	}
	buf, err := ioutil.ReadAll(&bodyReader{win})
	if err != nil {
		fmt.Fprintf(os.Stderr, "cannot read content: %v\n", err)
		os.Exit(1)
	}
	filename, off, err := selection(win, buf)
	if err != nil {
		fmt.Fprintf(os.Stderr, "cannot get selection: %v\n", err)
		os.Exit(1)
	}

	fileInfos, err := ioutil.ReadDir(".")
	if err != nil {
		fmt.Fprintf(os.Stderr, "cannot read directory: %v\n", err)
		os.Exit(1)
	}
	var filenames []string
	for _, f := range fileInfos {
		if strings.HasSuffix(f.Name(), ".go") && f.Name() != filename {
			filenames = append(filenames, f.Name())
		}
	}
	path, ident, err := searchAtOff(off, string(buf), filenames...)
	if err != nil {
		fmt.Fprintf(os.Stderr, "cannot find identifier: %v\n", err)
		os.Exit(1)
	}
	if ident != "" {
		godoc(path, ident)
	} else {
		godoc(path)
	}
}

func openWin() (*acme.Win, error) {
	id, err := strconv.Atoi(os.Getenv("winid"))
	if err != nil {
		return nil, err
	}
	return acme.Open(id, nil)
}

func selection(win *acme.Win, buf []byte) (filename string, off int, err error) {
	filename, err = readFilename(win)
	if err != nil {
		return "", 0, err
	}
	filename = filepath.Base(filename)
	q0, _, err := readAddr(win)
	if err != nil {
		return "", 0, err
	}
	off, err = byteOffset(bytes.NewReader(buf), q0)
	if err != nil {
		return "", 0, err
	}
	return
}

func readFilename(win *acme.Win) (string, error) {
	b, err := win.ReadAll("tag")
	if err != nil {
		return "", err
	}
	tag := string(b)
	i := strings.Index(tag, " ")
	if i == -1 {
		return "", fmt.Errorf("cannot get filename from tag")
	}
	return tag[0:i], nil
}

func readAddr(win *acme.Win) (q0, q1 int, err error) {
	if _, _, err := win.ReadAddr(); err != nil {
		return 0, 0, err
	}
	if err := win.Ctl("addr=dot"); err != nil {
		return 0, 0, err
	}
	return win.ReadAddr()
}

func byteOffset(r io.RuneReader, off int) (bo int, err error) {
	for i := 0; i != off; i++ {
		_, s, err := r.ReadRune()
		if err != nil {
			return 0, err
		}
		bo += s
	}
	return
}

// searchAtOff returns the path and the identifier name at the given offset.
func searchAtOff(off int, src string, filenames ...string) (string, string, error) {
	var conf loader.Config
	var files []*ast.File

	for _, name := range filenames {
		f, err := parseFile(&conf, name, nil)
		if err != nil {
			return "", "", err
		}
		files = append(files, f)
	}
	f, err := parseFile(&conf, "", src)
	if err != nil {
		return "", "", err
	}
	files = append(files, f)

	for _, imp := range f.Imports {
		if isAtOff(conf.Fset, off, imp) {
			return strings.Trim(imp.Path.Value, "\""), "", nil
		}
	}

	ident := identAtOffset(conf.Fset, f, off)
	if ident == nil {
		return "", "", errors.New("no identifier here")
	}

	conf.CreateFromFiles("", files...)
	allowErrors(&conf)
	prg, err := conf.Load()
	if err != nil {
		return "", "", err
	}
	info := prg.Created[0]
	if obj := info.Uses[ident]; obj != nil {
		return fromObj(obj, conf.Fset, f, off)
	}
	if obj := info.Defs[ident]; obj != nil {
		return fromObj(obj, conf.Fset, f, off)
	}
	return "", "", errors.New("could not find identifier")
}

func allowErrors(conf *loader.Config) {
	ctxt := build.Default
	ctxt.CgoEnabled = false
	conf.Build = &ctxt
	conf.AllowErrors = true
	conf.ParserMode = parser.AllErrors
	conf.TypeChecker.Error = func(err error) {}
}

func parseFile(conf *loader.Config, name string, src interface{}) (*ast.File, error) {
	f, err := conf.ParseFile(name, src)
	if err != nil {
		return nil, err
	}
	if strings.HasSuffix(f.Name.Name, "_test") {
		f.Name.Name = f.Name.Name[:len(f.Name.Name)-len("_test")]
	}
	return f, nil
}

// currImportDir returns the current import directory.
func currImportDir() (string, error) {
	path, err := os.Getwd()
	if err != nil {
		return "", err
	}
	pkg, err := build.ImportDir(path, build.ImportComment)
	if err != nil {
		return "", err
	}
	return pkg.ImportPath, nil
}

// fromObj returns the path and the identifier name of an object.
func fromObj(obj types.Object, fset *token.FileSet, f *ast.File, off int) (string, string, error) {
	switch o := obj.(type) {
	case *types.Builtin:
		return "builtin", o.Name(), nil
	case *types.PkgName:
		return o.Imported().Path(), "", nil
	case *types.Const, *types.Func, *types.Nil, *types.TypeName, *types.Var:
		if obj.Pkg() == nil {
			return "builtin", obj.Name(), nil
		}
		if !isImported(fset, f, off) {
			impDir, err := currImportDir()
			return impDir, obj.Name(), err
		}
		return obj.Pkg().Path(), obj.Name(), nil
	default:
		return "", "", fmt.Errorf("cannot print documentation of %v\n", obj)
	}
}

// isImported returns true whether the identifier at the given offset is imported or not.
func isImported(fset *token.FileSet, f *ast.File, off int) (isImp bool) {
	ast.Inspect(f, func(node ast.Node) bool {
		switch s := node.(type) {
		case *ast.SelectorExpr:
			if isAtOff(fset, off, s.Sel) {
				if ident, ok := s.X.(*ast.Ident); ok {
					for _, imp := range f.Imports {
						if imp.Name != nil && imp.Name.Name == ident.Name {
							isImp = true
							return false
						}
						parts := strings.Split(strings.Trim(imp.Path.Value, "\""), "/")
						if parts[len(parts)-1] == ident.Name {
							isImp = true
						}
					}
				}
				return false
			}
		}
		return true
	})
	return isImp
}

// identAtOffset returns the identifier at an offset in a file.
func identAtOffset(fset *token.FileSet, f *ast.File, off int) (ident *ast.Ident) {
	ast.Inspect(f, func(node ast.Node) bool {
		switch n := node.(type) {
		case *ast.Ident:
			if isAtOff(fset, off, node) {
				ident = n
				return false
			}
		}
		return true
	})
	return ident
}

// isAtOff reports whether a node is at a given offset in a file set.
func isAtOff(fset *token.FileSet, off int, n ast.Node) bool {
	return fset.Position(n.Pos()).Offset <= off && off <= fset.Position(n.End()).Offset
}

// godoc runs the godoc command with the given arguments.
func godoc(args ...string) {
	c := exec.Command("godoc", args...)
	c.Stderr, c.Stdout = os.Stderr, os.Stderr
	if err := c.Run(); err != nil {
		fmt.Fprintf(os.Stderr, "godoc failed: %v\n", err)
		os.Exit(1)
	}
}
<|endoftext|>