<commit_before>// (c) 2019-2020, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package rpcchainvm

import (
	"golang.org/x/net/context"

	"google.golang.org/grpc"

	"github.com/hashicorp/go-plugin"

	"github.com/ava-labs/avalanchego/snow/engine/snowman/block"
	"github.com/ava-labs/avalanchego/vms/rpcchainvm/vmproto"
)

// Handshake is a common handshake that is shared by plugin and host.
var Handshake = plugin.HandshakeConfig{
	ProtocolVersion:  1,
	MagicCookieKey:   "VM_PLUGIN",
	MagicCookieValue: "dynamic",
}

// PluginMap is the map of plugins we can dispense.
var PluginMap = map[string]plugin.Plugin{
	"vm": &Plugin{},
}

// Plugin is the implementation of plugin.Plugin so we can serve/consume this.
// We also implement GRPCPlugin so that this plugin can be served over gRPC.
type Plugin struct {
	plugin.NetRPCUnsupportedPlugin
	// Concrete implementation, written in Go. This is only used for plugins
	// that are written in Go.
	vm block.ChainVM
}

// New creates a new plugin from the provided VM
func New(vm block.ChainVM) *Plugin { return &Plugin{vm: vm} }

// GRPCServer registers a new GRPC server.
func (p *Plugin) GRPCServer(broker *plugin.GRPCBroker, s *grpc.Server) error {
	vmproto.RegisterVMServer(s, NewServer(p.vm, broker))
	return nil
}

// GRPCClient returns a new GRPC client
func (p *Plugin) GRPCClient(ctx context.Context, broker *plugin.GRPCBroker, c *grpc.ClientConn) (interface{}, error) {
	return NewClient(vmproto.NewVMClient(c), broker), nil
}
<commit_msg>Bump protocol version of rpcchainvm<commit_after>// (c) 2019-2020, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package rpcchainvm

import (
	"golang.org/x/net/context"

	"google.golang.org/grpc"

	"github.com/hashicorp/go-plugin"

	"github.com/ava-labs/avalanchego/snow/engine/snowman/block"
	"github.com/ava-labs/avalanchego/vms/rpcchainvm/vmproto"
)

// Handshake is a common handshake that is shared by plugin and host.
var Handshake = plugin.HandshakeConfig{
	ProtocolVersion:  2,
	MagicCookieKey:   "VM_PLUGIN",
	MagicCookieValue: "dynamic",
}

// PluginMap is the map of plugins we can dispense.
var PluginMap = map[string]plugin.Plugin{
	"vm": &Plugin{},
}

// Plugin is the implementation of plugin.Plugin so we can serve/consume this.
// We also implement GRPCPlugin so that this plugin can be served over gRPC.
type Plugin struct {
	plugin.NetRPCUnsupportedPlugin
	// Concrete implementation, written in Go. This is only used for plugins
	// that are written in Go.
	vm block.ChainVM
}

// New creates a new plugin from the provided VM
func New(vm block.ChainVM) *Plugin { return &Plugin{vm: vm} }

// GRPCServer registers a new GRPC server.
func (p *Plugin) GRPCServer(broker *plugin.GRPCBroker, s *grpc.Server) error {
	vmproto.RegisterVMServer(s, NewServer(p.vm, broker))
	return nil
}

// GRPCClient returns a new GRPC client
func (p *Plugin) GRPCClient(ctx context.Context, broker *plugin.GRPCBroker, c *grpc.ClientConn) (interface{}, error) {
	return NewClient(vmproto.NewVMClient(c), broker), nil
}
<|endoftext|>


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

package ctmap

import (
	"bytes"
	"fmt"
	"math"
	"net"
	"unsafe"

	"github.com/cilium/cilium/pkg/bpf"
	"github.com/cilium/cilium/pkg/byteorder"
	"github.com/cilium/cilium/pkg/logging"
)

var log = logging.DefaultLogger

const (
	MapName6       = "cilium_ct6_"
	MapName4       = "cilium_ct4_"
	MapName6Global = MapName6 + "global"
	MapName4Global = MapName4 + "global"

	MapNumEntriesLocal  = 64000
	MapNumEntriesGlobal = 1000000

	TUPLE_F_OUT     = 0
	TUPLE_F_IN      = 1
	TUPLE_F_RELATED = 2

	// MaxTime specifies the last possible time for GCFilter.Time
	MaxTime = math.MaxUint32

	noAction = iota
	deleteEntry
)

type CtType int

// CtKey is the interface describing keys to the conntrack maps.
type CtKey interface {
	bpf.MapKey

	// ToNetwork converts fields to network byte order.
	ToNetwork() CtKey

	// ToHost converts fields to host byte order.
	ToHost() CtKey

	// Dumps contents of key to buffer. Returns true if successful.
	Dump(buffer *bytes.Buffer) bool
}

// CtEntry represents an entry in the connection tracking table.
type CtEntry struct {
	rx_packets uint64
	rx_bytes   uint64
	tx_packets uint64
	tx_bytes   uint64
	lifetime   uint32
	flags      uint16
	// revnat is in network byte order
	revnat     uint16
	unused     uint16
	src_sec_id uint32
}

// GetValuePtr returns the unsafe.Pointer for s.
func (c *CtEntry) GetValuePtr() unsafe.Pointer { return unsafe.Pointer(c) }

// String returns the readable format
func (c *CtEntry) String() string {
	return fmt.Sprintf("expires=%d rx_packets=%d rx_bytes=%d tx_packets=%d tx_bytes=%d flags=%x revnat=%d src_sec_id=%d\n",
		c.lifetime,
		c.rx_packets,
		c.rx_bytes,
		c.tx_packets,
		c.tx_bytes,
		c.flags,
		byteorder.NetworkToHost(c.revnat),
		c.src_sec_id)
}

// CtEntryDump represents the key and value contained in the conntrack map.
type CtEntryDump struct {
	Key   CtKey
	Value CtEntry
}

const (
	// GCFilterNone doesn't filter the CT entries
	GCFilterNone = iota
	// GCFilterByTime filters CT entries by time
	GCFilterByTime
)

// GCFilterType is the type of a filter.
type GCFilterType uint

// GCFilter contains the necessary fields to filter the CT maps.
// Filtering by endpoint requires both EndpointID to be > 0 and
// EndpointIP to be not nil.
type GCFilter struct {
	Type       GCFilterType
	Time       uint32
	EndpointID uint16
	EndpointIP net.IP
}

// NewGCFilterBy creates a new GCFilter of the given type.
func NewGCFilterBy(filterType GCFilterType) *GCFilter {
	return &GCFilter{
		Type: filterType,
	}
}

// TypeString returns the filter type in human readable way.
func (f *GCFilter) TypeString() string {
	switch f.Type {
	case GCFilterNone:
		return "none"
	case GCFilterByTime:
		return "timeout"
	default:
		return "(unknown)"
	}
}

// ToString iterates through Map m and writes the values of the ct entries in m
// to a string.
func ToString(m *bpf.Map, mapName string) (string, error) {
	var buffer bytes.Buffer
	entries, err := dumpToSlice(m, mapName)
	if err != nil {
		return "", err
	}
	for _, entry := range entries {
		if !entry.Key.ToHost().Dump(&buffer) {
			continue
		}

		value := entry.Value
		buffer.WriteString(
			fmt.Sprintf(" expires=%d rx_packets=%d rx_bytes=%d tx_packets=%d tx_bytes=%d flags=%x revnat=%d src_sec_id=%d\n",
				value.lifetime,
				value.rx_packets,
				value.rx_bytes,
				value.tx_packets,
				value.tx_bytes,
				value.flags,
				byteorder.NetworkToHost(value.revnat),
				value.src_sec_id,
			),
		)

	}
	return buffer.String(), nil
}

// DumpToSlice iterates through map m and returns a slice mapping each key to
// its value in m.
func dumpToSlice(m *bpf.Map, mapType string) ([]CtEntryDump, error) {
	entries := []CtEntryDump{}

	switch mapType {
	case MapName6, MapName6Global:
		var key, nextKey CtKey6Global
		for {
			err := m.GetNextKey(&key, &nextKey)
			if err != nil {
				break
			}

			entry, err := m.Lookup(&nextKey)
			if err != nil {
				return nil, err
			}
			ctEntry := entry.(*CtEntry)

			nK := nextKey
			eDump := CtEntryDump{Key: &nK, Value: *ctEntry}
			entries = append(entries, eDump)

			key = nextKey
		}

	case MapName4, MapName4Global:
		var key, nextKey CtKey4Global
		for {
			err := m.GetNextKey(&key, &nextKey)
			if err != nil {
				break
			}

			entry, err := m.Lookup(&nextKey)
			if err != nil {
				return nil, err
			}
			ctEntry := entry.(*CtEntry)

			nK := nextKey
			eDump := CtEntryDump{Key: &nK, Value: *ctEntry}
			entries = append(entries, eDump)

			key = nextKey
		}
	}
	return entries, nil
}

// doGC6 iterates through a CTv6 map and drops entries based on the given
// filter.
func doGC6(m *bpf.Map, filter *GCFilter) int {
	var (
		action, deleted int
		nextKey, tmpKey CtKey6Global
	)

	err := m.GetNextKey(&tmpKey, &nextKey)
	if err != nil {
		return 0
	}

	for {
		nextKeyValid := m.GetNextKey(&nextKey, &tmpKey)
		entryMap, err := m.Lookup(&nextKey)
		if err != nil {
			log.WithError(err).Error("error during map Lookup")
			break
		}

		entry := entryMap.(*CtEntry)

		// In CT entries, the source address of the conntrack entry (`saddr`) is
		// the destination of the packet received, therefore it's the packet's
		// destination IP
		action = filter.doFiltering(nextKey.daddr.IP(), nextKey.saddr.IP(), nextKey.sport, uint8(nextKey.nexthdr), nextKey.flags, entry)

		switch action {
		case deleteEntry:
			err := m.Delete(&nextKey)
			if err != nil {
				log.WithError(err).Errorf("Unable to delete CT entry %s", nextKey.String())
			} else {
				deleted++
			}
		}

		if nextKeyValid != nil {
			break
		}
		nextKey = tmpKey
	}
	return deleted
}

// doGC4 iterates through a CTv4 map and drops entries based on the given
// filter.
func doGC4(m *bpf.Map, filter *GCFilter) int {
	var (
		action, deleted int
		nextKey, tmpKey CtKey4Global
	)

	err := m.GetNextKey(&tmpKey, &nextKey)
	if err != nil {
		return 0
	}

	for true {
		nextKeyValid := m.GetNextKey(&nextKey, &tmpKey)
		entryMap, err := m.Lookup(&nextKey)
		if err != nil {
			log.WithError(err).Error("error during map Lookup")
			break
		}

		entry := entryMap.(*CtEntry)

		// In CT entries, the source address of the conntrack entry (`saddr`) is
		// the destination of the packet received, therefore it's the packet's
		// destination IP
		action = filter.doFiltering(nextKey.daddr.IP(), nextKey.saddr.IP(), nextKey.sport, uint8(nextKey.nexthdr), nextKey.flags, entry)

		switch action {
		case deleteEntry:
			err := m.Delete(&nextKey)
			if err != nil {
				log.WithError(err).Errorf("Unable to delete CT entry %s", nextKey.String())
			} else {
				deleted++
			}
		}

		if nextKeyValid != nil {
			break
		}
		nextKey = tmpKey
	}
	return deleted
}

func (f *GCFilter) doFiltering(srcIP net.IP, dstIP net.IP, dstPort uint16, nextHdr, flags uint8, entry *CtEntry) (action int) {
	// Delete all entries with a lifetime smaller than f timestamp.
	if f.Type == GCFilterByTime && entry.lifetime < f.Time {
		return deleteEntry
	}

	return noAction
}

// GC runs garbage collection for map m with name mapName with the given filter.
// It returns how many items were deleted from m.
func GC(m *bpf.Map, mapName string, filter *GCFilter) int {
	if filter.Type == GCFilterByTime {
		// If LRUHashtable, no need to garbage collect as LRUHashtable cleans itself up.
		// FIXME: GH-3239 LRU logic is not handling timeouts gracefully enough
		// if m.MapInfo.MapType == bpf.MapTypeLRUHash {
		// 	return 0
		// }
		t, _ := bpf.GetMtime()
		tsec := t / 1000000000
		filter.Time = uint32(tsec)
	}

	switch mapName {
	case MapName6, MapName6Global:
		return doGC6(m, filter)
	case MapName4, MapName4Global:
		return doGC4(m, filter)
	default:
		return 0
	}
}

// Flush runs garbage collection for map m with the name mapName, deleting all
// entries. The specified map must be already opened using bpf.OpenMap().
func Flush(m *bpf.Map, mapName string) int {
	filter := NewGCFilterBy(GCFilterByTime)
	filter.Time = MaxTime

	switch mapName {
	case MapName6, MapName6Global:
		return doGC6(m, filter)
	case MapName4, MapName4Global:
		return doGC4(m, filter)
	default:
		return 0
	}
}
<commit_msg>ctmap: Make GC bpf map dumps more robust.<commit_after>// Copyright 2016-2017 Authors of Cilium
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package ctmap

import (
	"bytes"
	"fmt"
	"math"
	"net"
	"unsafe"

	"github.com/cilium/cilium/pkg/bpf"
	"github.com/cilium/cilium/pkg/byteorder"
	"github.com/cilium/cilium/pkg/logging"
)

var log = logging.DefaultLogger

const (
	MapName6       = "cilium_ct6_"
	MapName4       = "cilium_ct4_"
	MapName6Global = MapName6 + "global"
	MapName4Global = MapName4 + "global"

	MapNumEntriesLocal  = 64000
	MapNumEntriesGlobal = 1000000

	TUPLE_F_OUT     = 0
	TUPLE_F_IN      = 1
	TUPLE_F_RELATED = 2

	// MaxTime specifies the last possible time for GCFilter.Time
	MaxTime = math.MaxUint32

	noAction = iota
	deleteEntry
)

type CtType int

// CtKey is the interface describing keys to the conntrack maps.
type CtKey interface {
	bpf.MapKey

	// ToNetwork converts fields to network byte order.
	ToNetwork() CtKey

	// ToHost converts fields to host byte order.
	ToHost() CtKey

	// Dumps contents of key to buffer. Returns true if successful.
	Dump(buffer *bytes.Buffer) bool
}

// CtEntry represents an entry in the connection tracking table.
type CtEntry struct {
	rx_packets uint64
	rx_bytes   uint64
	tx_packets uint64
	tx_bytes   uint64
	lifetime   uint32
	flags      uint16
	// revnat is in network byte order
	revnat     uint16
	unused     uint16
	src_sec_id uint32
}

// GetValuePtr returns the unsafe.Pointer for s.
func (c *CtEntry) GetValuePtr() unsafe.Pointer { return unsafe.Pointer(c) }

// String returns the readable format
func (c *CtEntry) String() string {
	return fmt.Sprintf("expires=%d rx_packets=%d rx_bytes=%d tx_packets=%d tx_bytes=%d flags=%x revnat=%d src_sec_id=%d\n",
		c.lifetime,
		c.rx_packets,
		c.rx_bytes,
		c.tx_packets,
		c.tx_bytes,
		c.flags,
		byteorder.NetworkToHost(c.revnat),
		c.src_sec_id)
}

// CtEntryDump represents the key and value contained in the conntrack map.
type CtEntryDump struct {
	Key   CtKey
	Value CtEntry
}

const (
	// GCFilterNone doesn't filter the CT entries
	GCFilterNone = iota
	// GCFilterByTime filters CT entries by time
	GCFilterByTime
)

// GCFilterType is the type of a filter.
type GCFilterType uint

// GCFilter contains the necessary fields to filter the CT maps.
// Filtering by endpoint requires both EndpointID to be > 0 and
// EndpointIP to be not nil.
type GCFilter struct {
	Type       GCFilterType
	Time       uint32
	EndpointID uint16
	EndpointIP net.IP
}

// NewGCFilterBy creates a new GCFilter of the given type.
func NewGCFilterBy(filterType GCFilterType) *GCFilter {
	return &GCFilter{
		Type: filterType,
	}
}

// TypeString returns the filter type in human readable way.
func (f *GCFilter) TypeString() string {
	switch f.Type {
	case GCFilterNone:
		return "none"
	case GCFilterByTime:
		return "timeout"
	default:
		return "(unknown)"
	}
}

// ToString iterates through Map m and writes the values of the ct entries in m
// to a string.
func ToString(m *bpf.Map, mapName string) (string, error) {
	var buffer bytes.Buffer
	entries, err := dumpToSlice(m, mapName)
	if err != nil {
		return "", err
	}
	for _, entry := range entries {
		if !entry.Key.ToHost().Dump(&buffer) {
			continue
		}

		value := entry.Value
		buffer.WriteString(
			fmt.Sprintf(" expires=%d rx_packets=%d rx_bytes=%d tx_packets=%d tx_bytes=%d flags=%x revnat=%d src_sec_id=%d\n",
				value.lifetime,
				value.rx_packets,
				value.rx_bytes,
				value.tx_packets,
				value.tx_bytes,
				value.flags,
				byteorder.NetworkToHost(value.revnat),
				value.src_sec_id,
			),
		)

	}
	return buffer.String(), nil
}

// DumpToSlice iterates through map m and returns a slice mapping each key to
// its value in m.
func dumpToSlice(m *bpf.Map, mapType string) ([]CtEntryDump, error) {
	entries := []CtEntryDump{}

	switch mapType {
	case MapName6, MapName6Global:
		var key, nextKey CtKey6Global
		for {
			err := m.GetNextKey(&key, &nextKey)
			if err != nil {
				break
			}

			entry, err := m.Lookup(&nextKey)
			if err != nil {
				return nil, err
			}
			ctEntry := entry.(*CtEntry)

			nK := nextKey
			eDump := CtEntryDump{Key: &nK, Value: *ctEntry}
			entries = append(entries, eDump)

			key = nextKey
		}

	case MapName4, MapName4Global:
		var key, nextKey CtKey4Global
		for {
			err := m.GetNextKey(&key, &nextKey)
			if err != nil {
				break
			}

			entry, err := m.Lookup(&nextKey)
			if err != nil {
				return nil, err
			}
			ctEntry := entry.(*CtEntry)

			nK := nextKey
			eDump := CtEntryDump{Key: &nK, Value: *ctEntry}
			entries = append(entries, eDump)

			key = nextKey
		}
	}
	return entries, nil
}

// doGC6 iterates through a CTv6 map and drops entries based on the given
// filter.
func doGC6(m *bpf.Map, filter *GCFilter) int {
	var (
		action, deleted              int
		prevKey, currentKey, nextKey CtKey6Global
	)

	// prevKey is initially invalid, causing GetNextKey to return the first key in the map as currentKey.
	prevKeyValid := false
	err := m.GetNextKey(&prevKey, &currentKey)
	if err != nil {
		// Map is empty, nothing to clean up.
		return 0
	}

	var count uint32
	for count = 1; count <= m.MapInfo.MaxEntries; count++ {
		// currentKey was returned by GetNextKey() so we know it existed in the map, but it may have been
		// deleted by a concurrent map operation. If currentKey is no longer in the map, nextKey will be
		// the first key in the map again. Use the nextKey only if we still find currentKey in the Lookup()
		// after the GetNextKey() call, this way we know nextKey is NOT the first key in the map.
		nextKeyValid := m.GetNextKey(&currentKey, &nextKey)
		entryMap, err := m.Lookup(&currentKey)
		if err != nil {
			// Restarting from a invalid key starts the iteration again from the beginning.
			// If we have a previously found key, try to restart from there instead
			if prevKeyValid {
				currentKey = prevKey
				// Restart from a given previous key only once, otherwise if the prevKey is
				// concurrently deleted we might loop forever trying to look it up.
				prevKeyValid = false
			} else {
				// Depending on exactly when currentKey was deleted from the map, nextKey may be the actual
				// keyelement after the deleted one, or the first element in the map.
				currentKey = nextKey
			}
			continue
		}

		entry := entryMap.(*CtEntry)

		// In CT entries, the source address of the conntrack entry (`saddr`) is
		// the destination of the packet received, therefore it's the packet's
		// destination IP
		action = filter.doFiltering(currentKey.daddr.IP(), currentKey.saddr.IP(), currentKey.sport,
			uint8(currentKey.nexthdr), currentKey.flags, entry)

		switch action {
		case deleteEntry:
			err := m.Delete(&currentKey)
			if err != nil {
				log.WithError(err).Errorf("Unable to delete CT entry %s", currentKey.String())
			} else {
				deleted++
			}
		}

		if nextKeyValid != nil {
			break
		}
		// remember the last found key
		prevKey = currentKey
		prevKeyValid = true
		// continue from the next key
		currentKey = nextKey
	}

	if count > m.MapInfo.MaxEntries {
		// TODO Add a metric we can bump and observe here.
		log.WithError(err).Warning("Garbage collection on IPv6 CT map failed to finish")
	}

	return deleted
}

// doGC4 iterates through a CTv4 map and drops entries based on the given
// filter.
func doGC4(m *bpf.Map, filter *GCFilter) int {
	var (
		action, deleted              int
		prevKey, currentKey, nextKey CtKey4Global
	)

	// prevKey is initially invalid, causing GetNextKey to return the first key in the map as currentKey.
	prevKeyValid := false
	err := m.GetNextKey(&prevKey, &currentKey)
	if err != nil {
		// Map is empty, nothing to clean up.
		return 0
	}

	var count uint32
	for count = 1; count <= m.MapInfo.MaxEntries; count++ {
		// currentKey was returned by GetNextKey() so we know it existed in the map, but it may have been
		// deleted by a concurrent map operation. If currentKey is no longer in the map, nextKey will be
		// the first key in the map again. Use the nextKey only if we still find currentKey in the Lookup()
		// after the GetNextKey() call, this way we know nextKey is NOT the first key in the map.
		nextKeyValid := m.GetNextKey(&currentKey, &nextKey)
		entryMap, err := m.Lookup(&currentKey)
		if err != nil {
			// Restarting from a invalid key starts the iteration again from the beginning.
			// If we have a previously found key, try to restart from there instead
			if prevKeyValid {
				currentKey = prevKey
				// Restart from a given previous key only once, otherwise if the prevKey is
				// concurrently deleted we might loop forever trying to look it up.
				prevKeyValid = false
			} else {
				// Depending on exactly when currentKey was deleted from the map, nextKey may be the actual
				// keyelement after the deleted one, or the first element in the map.
				currentKey = nextKey
			}
			continue
		}

		entry := entryMap.(*CtEntry)

		// In CT entries, the source address of the conntrack entry (`saddr`) is
		// the destination of the packet received, therefore it's the packet's
		// destination IP
		action = filter.doFiltering(currentKey.daddr.IP(), currentKey.saddr.IP(), currentKey.sport,
			uint8(currentKey.nexthdr), currentKey.flags, entry)

		switch action {
		case deleteEntry:
			err := m.Delete(&currentKey)
			if err != nil {
				log.WithError(err).Errorf("Unable to delete CT entry %s", currentKey.String())
			} else {
				deleted++
			}
		}

		if nextKeyValid != nil {
			break
		}
		// remember the last found key
		prevKey = currentKey
		prevKeyValid = true
		// continue from the next key
		currentKey = nextKey
	}

	if count > m.MapInfo.MaxEntries {
		// TODO Add a metric we can bump and observe here.
		log.WithError(err).Warning("Garbage collection on IPv4 CT map failed to finish")
	}

	return deleted
}

func (f *GCFilter) doFiltering(srcIP net.IP, dstIP net.IP, dstPort uint16, nextHdr, flags uint8, entry *CtEntry) (action int) {
	// Delete all entries with a lifetime smaller than f timestamp.
	if f.Type == GCFilterByTime && entry.lifetime < f.Time {
		return deleteEntry
	}

	return noAction
}

// GC runs garbage collection for map m with name mapName with the given filter.
// It returns how many items were deleted from m.
func GC(m *bpf.Map, mapName string, filter *GCFilter) int {
	if filter.Type == GCFilterByTime {
		// If LRUHashtable, no need to garbage collect as LRUHashtable cleans itself up.
		// FIXME: GH-3239 LRU logic is not handling timeouts gracefully enough
		// if m.MapInfo.MapType == bpf.MapTypeLRUHash {
		// 	return 0
		// }
		t, _ := bpf.GetMtime()
		tsec := t / 1000000000
		filter.Time = uint32(tsec)
	}

	switch mapName {
	case MapName6, MapName6Global:
		return doGC6(m, filter)
	case MapName4, MapName4Global:
		return doGC4(m, filter)
	default:
		return 0
	}
}

// Flush runs garbage collection for map m with the name mapName, deleting all
// entries. The specified map must be already opened using bpf.OpenMap().
func Flush(m *bpf.Map, mapName string) int {
	filter := NewGCFilterBy(GCFilterByTime)
	filter.Time = MaxTime

	switch mapName {
	case MapName6, MapName6Global:
		return doGC6(m, filter)
	case MapName4, MapName4Global:
		return doGC4(m, filter)
	default:
		return 0
	}
}
<|endoftext|>


<commit_before>package proxy

import (
	"encoding/json"

	"github.com/asaskevich/govalidator"
	"github.com/hellofresh/janus/pkg/router"
)

// Route is the container for a proxy and it's handlers
type Route struct {
	proxy    *Definition
	handlers []router.Constructor
}

type routeJSONProxy struct {
	Proxy *Definition `json:"proxy"`
}

// NewRoute creates an instance of Route
func NewRoute(proxy *Definition, handlers ...router.Constructor) *Route {
	return &Route{proxy, handlers}
}

// JSONMarshal encodes route struct to JSON
func (r *Route) JSONMarshal() ([]byte, error) {
	return json.Marshal(routeJSONProxy{r.proxy})
}

// JSONUnmarshalRoute decodes route struct from JSON
func JSONUnmarshalRoute(rawRoute []byte) (*Route, error) {
	var proxyRoute routeJSONProxy
	if err := json.Unmarshal(rawRoute, &proxyRoute); err != nil {
		return nil, err
	}
	return NewRoute(proxyRoute.Proxy), nil
}

// Definition defines proxy rules for a route
type Definition struct {
	PreserveHost        bool     `bson:"preserve_host" json:"preserve_host" mapstructure:"preserve_host"`
	ListenPath          string   `bson:"listen_path" json:"listen_path" mapstructure:"listen_path" valid:"required"`
	UpstreamURL         string   `bson:"upstream_url" json:"upstream_url" mapstructure:"upstream_url" valid:"url,required"`
	StripPath           bool     `bson:"strip_path" json:"strip_path" mapstructure:"strip_path"`
	AppendPath          bool     `bson:"append_path" json:"append_path" mapstructure:"append_path"`
	EnableLoadBalancing bool     `bson:"enable_load_balancing" json:"enable_load_balancing" mapstructure:"enable_load_balancing"`
	Methods             []string `bson:"methods" json:"methods"`
	Hosts               []string `bson:"hosts" json:"hosts"`
}

// Validate validates proxy data
func (d *Definition) Validate() (bool, error) {
	return govalidator.ValidateStruct(d)
}
<commit_msg>Added constructor fot proxy<commit_after>package proxy

import (
	"encoding/json"

	"github.com/asaskevich/govalidator"
	"github.com/hellofresh/janus/pkg/router"
)

// Route is the container for a proxy and it's handlers
type Route struct {
	proxy    *Definition
	handlers []router.Constructor
}

type routeJSONProxy struct {
	Proxy *Definition `json:"proxy"`
}

// NewRoute creates an instance of Route
func NewRoute(proxy *Definition, handlers ...router.Constructor) *Route {
	return &Route{proxy, handlers}
}

// JSONMarshal encodes route struct to JSON
func (r *Route) JSONMarshal() ([]byte, error) {
	return json.Marshal(routeJSONProxy{r.proxy})
}

// JSONUnmarshalRoute decodes route struct from JSON
func JSONUnmarshalRoute(rawRoute []byte) (*Route, error) {
	var proxyRoute routeJSONProxy
	if err := json.Unmarshal(rawRoute, &proxyRoute); err != nil {
		return nil, err
	}
	return NewRoute(proxyRoute.Proxy), nil
}

// Definition defines proxy rules for a route
type Definition struct {
	PreserveHost        bool     `bson:"preserve_host" json:"preserve_host" mapstructure:"preserve_host"`
	ListenPath          string   `bson:"listen_path" json:"listen_path" mapstructure:"listen_path" valid:"required"`
	UpstreamURL         string   `bson:"upstream_url" json:"upstream_url" mapstructure:"upstream_url" valid:"url,required"`
	StripPath           bool     `bson:"strip_path" json:"strip_path" mapstructure:"strip_path"`
	AppendPath          bool     `bson:"append_path" json:"append_path" mapstructure:"append_path"`
	EnableLoadBalancing bool     `bson:"enable_load_balancing" json:"enable_load_balancing" mapstructure:"enable_load_balancing"`
	Methods             []string `bson:"methods" json:"methods"`
	Hosts               []string `bson:"hosts" json:"hosts"`
}

// NewDefinition creates a new Proxy Definition with default values
func NewDefinition() *Definition {
	return &Definition{
		Methods: make([]string, 0),
		Hosts:   make([]string, 0),
	}
}

// Validate validates proxy data
func (d *Definition) Validate() (bool, error) {
	return govalidator.ValidateStruct(d)
}
<|endoftext|>


<commit_before>/*
Copyright 2017 The Kubernetes Authors.

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

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

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

package server

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"fmt"
	"io"
	"math"
	"math/big"
	"net"
	"os"
	"time"

	"github.com/pkg/errors"
	k8snet "k8s.io/apimachinery/pkg/util/net"
	"k8s.io/apimachinery/pkg/util/runtime"
	"k8s.io/client-go/tools/remotecommand"
	k8scert "k8s.io/client-go/util/cert"
	"k8s.io/kubernetes/pkg/kubelet/server/streaming"
	"k8s.io/utils/exec"

	ctrdutil "github.com/containerd/cri/pkg/containerd/util"
)

const (
	// OrganizationName is is the name of this organization, used for certificates etc.
	OrganizationName = "containerd"
	// CRIName is the common name of the CRI plugin
	CRIName = "cri"
)

func newStreamServer(c *criService, addr, port string) (streaming.Server, error) {
	if addr == "" {
		a, err := k8snet.ChooseBindAddress(nil)
		if err != nil {
			return nil, errors.Wrap(err, "failed to get stream server address")
		}
		addr = a.String()
	}
	config := streaming.DefaultConfig
	config.Addr = net.JoinHostPort(addr, port)
	runtime := newStreamRuntime(c)
	tlsCert, err := newTLSCert()
	if err != nil {
		return nil, errors.Wrap(err, "failed to generate tls certificate for stream server")
	}
	config.TLSConfig = &tls.Config{
		Certificates:       []tls.Certificate{tlsCert},
		InsecureSkipVerify: true,
	}
	return streaming.NewServer(config, runtime)
}

type streamRuntime struct {
	c *criService
}

func newStreamRuntime(c *criService) streaming.Runtime {
	return &streamRuntime{c: c}
}

// Exec executes a command inside the container. exec.ExitError is returned if the command
// returns non-zero exit code.
func (s *streamRuntime) Exec(containerID string, cmd []string, stdin io.Reader, stdout, stderr io.WriteCloser,
	tty bool, resize <-chan remotecommand.TerminalSize) error {
	exitCode, err := s.c.execInContainer(ctrdutil.NamespacedContext(), containerID, execOptions{
		cmd:    cmd,
		stdin:  stdin,
		stdout: stdout,
		stderr: stderr,
		tty:    tty,
		resize: resize,
	})
	if err != nil {
		return errors.Wrap(err, "failed to exec in container")
	}
	if *exitCode == 0 {
		return nil
	}
	return &exec.CodeExitError{
		Err:  errors.Errorf("error executing command %v, exit code %d", cmd, *exitCode),
		Code: int(*exitCode),
	}
}

func (s *streamRuntime) Attach(containerID string, in io.Reader, out, err io.WriteCloser, tty bool,
	resize <-chan remotecommand.TerminalSize) error {
	return s.c.attachContainer(ctrdutil.NamespacedContext(), containerID, in, out, err, tty, resize)
}

func (s *streamRuntime) PortForward(podSandboxID string, port int32, stream io.ReadWriteCloser) error {
	if port <= 0 || port > math.MaxUint16 {
		return errors.Errorf("invalid port %d", port)
	}
	return s.c.portForward(podSandboxID, port, stream)
}

// handleResizing spawns a goroutine that processes the resize channel, calling resizeFunc for each
// remotecommand.TerminalSize received from the channel. The resize channel must be closed elsewhere to stop the
// goroutine.
func handleResizing(resize <-chan remotecommand.TerminalSize, resizeFunc func(size remotecommand.TerminalSize)) {
	if resize == nil {
		return
	}

	go func() {
		defer runtime.HandleCrash()

		for {
			size, ok := <-resize
			if !ok {
				return
			}
			if size.Height < 1 || size.Width < 1 {
				continue
			}
			resizeFunc(size)
		}
	}()
}

// newTLSCert returns a tls.certificate loaded from a newly generated
// x509certificate from a newly generated rsa public/private key pair. The
// x509certificate is self signed.
// TODO (mikebrow): replace / rewrite this function to support using CA
// signing of the cetificate. Requires a security plan for kubernetes regarding
// CRI connections / streaming, etc. For example, kubernetes could configure or
// require a CA service and pass a configuration down through CRI.
func newTLSCert() (tls.Certificate, error) {
	fail := func(err error) (tls.Certificate, error) { return tls.Certificate{}, err }
	var years = 1 // duration of certificate

	// Generate new private key
	privKey, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return fail(errors.Wrap(err, "private key cannot be created"))
	}

	// Generate pem block using the private key
	keyPem := pem.EncodeToMemory(&pem.Block{
		Type:  k8scert.RSAPrivateKeyBlockType,
		Bytes: x509.MarshalPKCS1PrivateKey(privKey),
	})

	// Generate a new random serial number for certificate
	serialNumber, err := rand.Int(rand.Reader, new(big.Int).Lsh(big.NewInt(1), 128))
	if err != nil {
		return fail(errors.Wrap(err, "failed to generate serial number"))
	}
	hostName, err := os.Hostname()
	if err != nil {
		return fail(errors.Wrap(err, "failed to get hostname"))
	}
	addrs, err := net.InterfaceAddrs()
	if err != nil {
		return fail(errors.Wrap(err, "failed to get host IP addresses"))
	}

	// Configure and create new certificate
	tml := x509.Certificate{
		NotBefore:    time.Now(),
		NotAfter:     time.Now().AddDate(years, 0, 0),
		SerialNumber: serialNumber,
		Subject: pkix.Name{
			CommonName:   fmt.Sprintf("%s:%s:%s", OrganizationName, CRIName, hostName),
			Organization: []string{OrganizationName},
		},
		BasicConstraintsValid: true,
	}
	for _, addr := range addrs {
		var ip net.IP

		switch v := addr.(type) {
		case *net.IPNet:
			ip = v.IP
		case *net.IPAddr:
			ip = v.IP
		default:
			continue
		}

		tml.IPAddresses = append(tml.IPAddresses, ip)
		tml.DNSNames = append(tml.DNSNames, ip.String())
	}

	cert, err := x509.CreateCertificate(rand.Reader, &tml, &tml, &privKey.PublicKey, privKey)
	if err != nil {
		return fail(errors.Wrap(err, "certificate cannot be created"))
	}

	// Generate a pem block with the certificate
	certPem := pem.EncodeToMemory(&pem.Block{
		Type:  k8scert.CertificateBlockType,
		Bytes: cert,
	})

	// Load the tls certificate
	tlsCert, err := tls.X509KeyPair(certPem, keyPem)
	if err != nil {
		return fail(errors.Wrap(err, "certificate could not be loaded"))
	}

	return tlsCert, nil
}
<commit_msg>Make const private.<commit_after>/*
Copyright 2017 The Kubernetes Authors.

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

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

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

package server

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"fmt"
	"io"
	"math"
	"math/big"
	"net"
	"os"
	"time"

	"github.com/pkg/errors"
	k8snet "k8s.io/apimachinery/pkg/util/net"
	"k8s.io/apimachinery/pkg/util/runtime"
	"k8s.io/client-go/tools/remotecommand"
	k8scert "k8s.io/client-go/util/cert"
	"k8s.io/kubernetes/pkg/kubelet/server/streaming"
	"k8s.io/utils/exec"

	ctrdutil "github.com/containerd/cri/pkg/containerd/util"
)

const (
	// certOrganizationName is the name of this organization, used for certificates etc.
	certOrganizationName = "containerd"
	// certCommonName is the common name of the CRI plugin
	certCommonName = "cri"
)

func newStreamServer(c *criService, addr, port string) (streaming.Server, error) {
	if addr == "" {
		a, err := k8snet.ChooseBindAddress(nil)
		if err != nil {
			return nil, errors.Wrap(err, "failed to get stream server address")
		}
		addr = a.String()
	}
	config := streaming.DefaultConfig
	config.Addr = net.JoinHostPort(addr, port)
	runtime := newStreamRuntime(c)
	tlsCert, err := newTLSCert()
	if err != nil {
		return nil, errors.Wrap(err, "failed to generate tls certificate for stream server")
	}
	config.TLSConfig = &tls.Config{
		Certificates:       []tls.Certificate{tlsCert},
		InsecureSkipVerify: true,
	}
	return streaming.NewServer(config, runtime)
}

type streamRuntime struct {
	c *criService
}

func newStreamRuntime(c *criService) streaming.Runtime {
	return &streamRuntime{c: c}
}

// Exec executes a command inside the container. exec.ExitError is returned if the command
// returns non-zero exit code.
func (s *streamRuntime) Exec(containerID string, cmd []string, stdin io.Reader, stdout, stderr io.WriteCloser,
	tty bool, resize <-chan remotecommand.TerminalSize) error {
	exitCode, err := s.c.execInContainer(ctrdutil.NamespacedContext(), containerID, execOptions{
		cmd:    cmd,
		stdin:  stdin,
		stdout: stdout,
		stderr: stderr,
		tty:    tty,
		resize: resize,
	})
	if err != nil {
		return errors.Wrap(err, "failed to exec in container")
	}
	if *exitCode == 0 {
		return nil
	}
	return &exec.CodeExitError{
		Err:  errors.Errorf("error executing command %v, exit code %d", cmd, *exitCode),
		Code: int(*exitCode),
	}
}

func (s *streamRuntime) Attach(containerID string, in io.Reader, out, err io.WriteCloser, tty bool,
	resize <-chan remotecommand.TerminalSize) error {
	return s.c.attachContainer(ctrdutil.NamespacedContext(), containerID, in, out, err, tty, resize)
}

func (s *streamRuntime) PortForward(podSandboxID string, port int32, stream io.ReadWriteCloser) error {
	if port <= 0 || port > math.MaxUint16 {
		return errors.Errorf("invalid port %d", port)
	}
	return s.c.portForward(podSandboxID, port, stream)
}

// handleResizing spawns a goroutine that processes the resize channel, calling resizeFunc for each
// remotecommand.TerminalSize received from the channel. The resize channel must be closed elsewhere to stop the
// goroutine.
func handleResizing(resize <-chan remotecommand.TerminalSize, resizeFunc func(size remotecommand.TerminalSize)) {
	if resize == nil {
		return
	}

	go func() {
		defer runtime.HandleCrash()

		for {
			size, ok := <-resize
			if !ok {
				return
			}
			if size.Height < 1 || size.Width < 1 {
				continue
			}
			resizeFunc(size)
		}
	}()
}

// newTLSCert returns a tls.certificate loaded from a newly generated
// x509certificate from a newly generated rsa public/private key pair. The
// x509certificate is self signed.
// TODO (mikebrow): replace / rewrite this function to support using CA
// signing of the cetificate. Requires a security plan for kubernetes regarding
// CRI connections / streaming, etc. For example, kubernetes could configure or
// require a CA service and pass a configuration down through CRI.
func newTLSCert() (tls.Certificate, error) {
	fail := func(err error) (tls.Certificate, error) { return tls.Certificate{}, err }
	var years = 1 // duration of certificate

	// Generate new private key
	privKey, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return fail(errors.Wrap(err, "private key cannot be created"))
	}

	// Generate pem block using the private key
	keyPem := pem.EncodeToMemory(&pem.Block{
		Type:  k8scert.RSAPrivateKeyBlockType,
		Bytes: x509.MarshalPKCS1PrivateKey(privKey),
	})

	// Generate a new random serial number for certificate
	serialNumber, err := rand.Int(rand.Reader, new(big.Int).Lsh(big.NewInt(1), 128))
	if err != nil {
		return fail(errors.Wrap(err, "failed to generate serial number"))
	}
	hostName, err := os.Hostname()
	if err != nil {
		return fail(errors.Wrap(err, "failed to get hostname"))
	}
	addrs, err := net.InterfaceAddrs()
	if err != nil {
		return fail(errors.Wrap(err, "failed to get host IP addresses"))
	}

	// Configure and create new certificate
	tml := x509.Certificate{
		NotBefore:    time.Now(),
		NotAfter:     time.Now().AddDate(years, 0, 0),
		SerialNumber: serialNumber,
		Subject: pkix.Name{
			CommonName:   fmt.Sprintf("%s:%s:%s", certOrganizationName, certCommonName, hostName),
			Organization: []string{certOrganizationName},
		},
		BasicConstraintsValid: true,
	}
	for _, addr := range addrs {
		var ip net.IP

		switch v := addr.(type) {
		case *net.IPNet:
			ip = v.IP
		case *net.IPAddr:
			ip = v.IP
		default:
			continue
		}

		tml.IPAddresses = append(tml.IPAddresses, ip)
		tml.DNSNames = append(tml.DNSNames, ip.String())
	}

	cert, err := x509.CreateCertificate(rand.Reader, &tml, &tml, &privKey.PublicKey, privKey)
	if err != nil {
		return fail(errors.Wrap(err, "certificate cannot be created"))
	}

	// Generate a pem block with the certificate
	certPem := pem.EncodeToMemory(&pem.Block{
		Type:  k8scert.CertificateBlockType,
		Bytes: cert,
	})

	// Load the tls certificate
	tlsCert, err := tls.X509KeyPair(certPem, keyPem)
	if err != nil {
		return fail(errors.Wrap(err, "certificate could not be loaded"))
	}

	return tlsCert, nil
}
<|endoftext|>


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

package plans_test

import (
	"fmt"

	. "github.com/pingcap/check"
	"github.com/pingcap/tidb"
	"github.com/pingcap/tidb/expression/expressions"
	"github.com/pingcap/tidb/kv"
	"github.com/pingcap/tidb/model"
	mysql "github.com/pingcap/tidb/mysqldef"
	"github.com/pingcap/tidb/parser/opcode"
	"github.com/pingcap/tidb/plan/plans"
	"github.com/pingcap/tidb/rset/rsets"
	"github.com/pingcap/tidb/sessionctx/variable"
	"github.com/pingcap/tidb/stmt"
)

type testShowSuit struct {
	txn  kv.Transaction
	vars map[string]interface{}
}

// implement Context interface
func (p *testShowSuit) GetTxn(forceNew bool) (kv.Transaction, error) { return p.txn, nil }

func (p *testShowSuit) FinishTxn(rollback bool) error { return nil }

// SetValue saves a value associated with this context for key
func (p *testShowSuit) SetValue(key fmt.Stringer, value interface{}) {
	p.vars[key.String()] = value
}

// Value returns the value associated with this context for key
func (p *testShowSuit) Value(key fmt.Stringer) interface{} {
	return p.vars[key.String()]
}

// ClearValue clears the value associated with this context for key
func (p *testShowSuit) ClearValue(key fmt.Stringer) {}

var _ = Suite(&testShowSuit{})

func (p *testShowSuit) SetUpSuite(c *C) {
	var err error
	store, err := tidb.NewStore(tidb.EngineGoLevelDBMemory)
	c.Assert(err, IsNil)
	p.vars = map[string]interface{}{}
	p.txn, _ = store.Begin()
	variable.BindSessionVars(p)
}

func (p *testShowSuit) TestShowVariables(c *C) {
	pln := &plans.ShowPlan{
		Target:      stmt.ShowVariables,
		GlobalScope: true,
		Pattern: &expressions.PatternLike{
			Pattern: &expressions.Value{
				Val: "character_set_results",
			},
		},
	}
	fls := pln.GetFields()
	c.Assert(fls, HasLen, 2)
	c.Assert(fls[0].Name, Equals, "Variable_name")
	c.Assert(fls[1].Name, Equals, "Value")
	c.Assert(fls[0].Col.Tp, Equals, mysql.TypeVarchar)
	c.Assert(fls[0].Col.Tp, Equals, mysql.TypeVarchar)

	sessionVars := variable.GetSessionVars(p)
	ret := map[string]string{}
	rset := rsets.Recordset{
		Ctx:  p,
		Plan: pln,
	}
	rset.Do(func(data []interface{}) (bool, error) {
		ret[data[0].(string)] = data[1].(string)
		return true, nil
	})
	c.Assert(ret, HasLen, 1)
	v, ok := ret["character_set_results"]
	c.Assert(ok, IsTrue)
	c.Assert(v, Equals, "latin1")
	// Set session variable to utf8
	sessionVars.Systems["character_set_results"] = "utf8"
	pln.Close()
	rset.Do(func(data []interface{}) (bool, error) {
		ret[data[0].(string)] = data[1].(string)
		return true, nil
	})
	c.Assert(ret, HasLen, 1)
	v, ok = ret["character_set_results"]
	c.Assert(ok, IsTrue)
	// Show global varibale get latin1
	c.Assert(v, Equals, "latin1")

	pln.GlobalScope = false
	pln.Close()
	rset.Do(func(data []interface{}) (bool, error) {
		ret[data[0].(string)] = data[1].(string)
		return true, nil
	})
	c.Assert(ret, HasLen, 1)
	v, ok = ret["character_set_results"]
	c.Assert(ok, IsTrue)
	// Show session varibale get utf8
	c.Assert(v, Equals, "utf8")
	pln.Close()
	pln.Pattern = nil
	pln.Where = &expressions.BinaryOperation{
		L:  &expressions.Ident{CIStr: model.NewCIStr("Variable_name")},
		R:  expressions.Value{Val: "autocommit"},
		Op: opcode.EQ,
	}

	ret = map[string]string{}
	sessionVars.Systems["autocommit"] = "on"
	rset.Do(func(data []interface{}) (bool, error) {
		ret[data[0].(string)] = data[1].(string)
		return true, nil
	})

	c.Assert(ret, HasLen, 1)
	v, ok = ret["autocommit"]
	c.Assert(ok, IsTrue)
	c.Assert(v, Equals, "on")

	pln.Target = stmt.ShowWarnings
	fls = pln.GetFields()
	c.Assert(fls, HasLen, 3)
	c.Assert(fls[1].Col.Tp, Equals, mysql.TypeLong)

	pln.Target = stmt.ShowCharset
	fls = pln.GetFields()
	c.Assert(fls, HasLen, 4)
	c.Assert(fls[3].Col.Tp, Equals, mysql.TypeLonglong)
}

func (p *testShowSuit) TearDownSuite(c *C) {
	p.txn.Commit()
}
<commit_msg>plans: Address comment<commit_after>// Copyright 2015 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

package plans_test

import (
	"fmt"

	. "github.com/pingcap/check"
	"github.com/pingcap/tidb"
	"github.com/pingcap/tidb/expression/expressions"
	"github.com/pingcap/tidb/kv"
	"github.com/pingcap/tidb/model"
	mysql "github.com/pingcap/tidb/mysqldef"
	"github.com/pingcap/tidb/parser/opcode"
	"github.com/pingcap/tidb/plan/plans"
	"github.com/pingcap/tidb/rset/rsets"
	"github.com/pingcap/tidb/sessionctx/variable"
	"github.com/pingcap/tidb/stmt"
)

type testShowSuit struct {
	txn  kv.Transaction
	vars map[string]interface{}
}

// implement Context interface
func (p *testShowSuit) GetTxn(forceNew bool) (kv.Transaction, error) { return p.txn, nil }

func (p *testShowSuit) FinishTxn(rollback bool) error { return nil }

// SetValue saves a value associated with this context for key
func (p *testShowSuit) SetValue(key fmt.Stringer, value interface{}) {
	p.vars[key.String()] = value
}

// Value returns the value associated with this context for key
func (p *testShowSuit) Value(key fmt.Stringer) interface{} {
	return p.vars[key.String()]
}

// ClearValue clears the value associated with this context for key
func (p *testShowSuit) ClearValue(key fmt.Stringer) {}

var _ = Suite(&testShowSuit{})

func (p *testShowSuit) SetUpSuite(c *C) {
	var err error
	store, err := tidb.NewStore(tidb.EngineGoLevelDBMemory)
	c.Assert(err, IsNil)
	p.vars = map[string]interface{}{}
	p.txn, _ = store.Begin()
	variable.BindSessionVars(p)
}

func (p *testShowSuit) TestShowVariables(c *C) {
	pln := &plans.ShowPlan{
		Target:      stmt.ShowVariables,
		GlobalScope: true,
		Pattern: &expressions.PatternLike{
			Pattern: &expressions.Value{
				Val: "character_set_results",
			},
		},
	}
	fls := pln.GetFields()
	c.Assert(fls, HasLen, 2)
	c.Assert(fls[0].Name, Equals, "Variable_name")
	c.Assert(fls[1].Name, Equals, "Value")
	c.Assert(fls[0].Col.Tp, Equals, mysql.TypeVarchar)
	c.Assert(fls[1].Col.Tp, Equals, mysql.TypeVarchar)

	sessionVars := variable.GetSessionVars(p)
	ret := map[string]string{}
	rset := rsets.Recordset{
		Ctx:  p,
		Plan: pln,
	}
	rset.Do(func(data []interface{}) (bool, error) {
		ret[data[0].(string)] = data[1].(string)
		return true, nil
	})
	c.Assert(ret, HasLen, 1)
	v, ok := ret["character_set_results"]
	c.Assert(ok, IsTrue)
	c.Assert(v, Equals, "latin1")
	// Set session variable to utf8
	sessionVars.Systems["character_set_results"] = "utf8"
	pln.Close()
	rset.Do(func(data []interface{}) (bool, error) {
		ret[data[0].(string)] = data[1].(string)
		return true, nil
	})
	c.Assert(ret, HasLen, 1)
	v, ok = ret["character_set_results"]
	c.Assert(ok, IsTrue)
	// Show global varibale get latin1
	c.Assert(v, Equals, "latin1")

	pln.GlobalScope = false
	pln.Close()
	rset.Do(func(data []interface{}) (bool, error) {
		ret[data[0].(string)] = data[1].(string)
		return true, nil
	})
	c.Assert(ret, HasLen, 1)
	v, ok = ret["character_set_results"]
	c.Assert(ok, IsTrue)
	// Show session varibale get utf8
	c.Assert(v, Equals, "utf8")
	pln.Close()
	pln.Pattern = nil
	pln.Where = &expressions.BinaryOperation{
		L:  &expressions.Ident{CIStr: model.NewCIStr("Variable_name")},
		R:  expressions.Value{Val: "autocommit"},
		Op: opcode.EQ,
	}

	ret = map[string]string{}
	sessionVars.Systems["autocommit"] = "on"
	rset.Do(func(data []interface{}) (bool, error) {
		ret[data[0].(string)] = data[1].(string)
		return true, nil
	})

	c.Assert(ret, HasLen, 1)
	v, ok = ret["autocommit"]
	c.Assert(ok, IsTrue)
	c.Assert(v, Equals, "on")

	pln.Target = stmt.ShowWarnings
	fls = pln.GetFields()
	c.Assert(fls, HasLen, 3)
	c.Assert(fls[1].Col.Tp, Equals, mysql.TypeLong)

	pln.Target = stmt.ShowCharset
	fls = pln.GetFields()
	c.Assert(fls, HasLen, 4)
	c.Assert(fls[3].Col.Tp, Equals, mysql.TypeLonglong)
}

func (p *testShowSuit) TearDownSuite(c *C) {
	p.txn.Commit()
}
<|endoftext|>


<commit_before>// Copyright 2016 Factom Foundation
// Use of this source code is governed by the MIT
// license that can be found in the LICENSE file.

package factom_test

import (
	"bytes"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"net/http"
	"net/http/httptest"
	"testing"

	. "github.com/FactomProject/factom"
)

var ()

func TestNewChain(t *testing.T) {
	ent := new(Entry)
	ent.ChainID = ""
	ent.Content = []byte("This is a test Entry.")
	ent.ExtIDs = append(ent.ExtIDs, []byte("This is the first extid."))
	ent.ExtIDs = append(ent.ExtIDs, []byte("This is the second extid."))

	newChain := NewChain(ent)
	expectedID := "5a402200c5cf278e47905ce52d7d64529a0291829a7bd230072c5468be709069"

	if newChain.ChainID != expectedID {
		t.Errorf("expected:%s\nrecieved:%s", expectedID, newChain.ChainID)
	}
	t.Log(newChain.ChainID)

	cfb := NewChainFromBytes(ent.Content, ent.ExtIDs...)
	if cfb.ChainID != expectedID {
		t.Errorf("expected:%s\nrecieved:%s", expectedID, cfb.ChainID)
	}
	t.Log(cfb.ChainID)

	cfs := NewChainFromStrings(
		"This is a test Entry.",
		"This is the first extid.",
		"This is the second extid.",
	)
	if cfs.ChainID != expectedID {
		t.Errorf("expected:%s\nrecieved:%s", expectedID, cfs.ChainID)
	}
	t.Log(cfs.ChainID)
}

func TestIfExists(t *testing.T) {
	simlatedFactomdResponse := `{
  "jsonrpc": "2.0",
  "id": 0,
  "result": {
    "ChainHead": "f65f67774139fa78344dcdd302631a0d646db0c2be4d58e3e48b2a188c1b856c"
  }
}`

	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		fmt.Fprintln(w, simlatedFactomdResponse)
	}))
	defer ts.Close()

	url := ts.URL[7:]
	SetFactomdServer(url)

	expectedID := "f65f67774139fa78344dcdd302631a0d646db0c2be4d58e3e48b2a188c1b856c"
	//fmt.Println(ChainExists(expectedID))
	if ChainExists(expectedID) != true {
		t.Errorf("chain %s does not exist", expectedID)
	}
}

func TestIfNotExists(t *testing.T) {
	simlatedFactomdResponse := `{"jsonrpc":"2.0","id":0,"error":{"code":-32009,"message":"Missing Chain Head"}}`

	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		fmt.Fprintln(w, simlatedFactomdResponse)
	}))
	defer ts.Close()

	url := ts.URL[7:]
	SetFactomdServer(url)
	unexpectedID := "5a402200c5cf278e47905ce52d7d64529a0291829a7bd230072c5468be709069"

	if ChainExists(unexpectedID) != false {
		t.Errorf("chain %s shouldn't exist", unexpectedID)
	}
}

func TestComposeChainCommit(t *testing.T) {
	type response struct {
		Message string `json:"message"`
	}
	ecAddr, _ := GetECAddress("Es2Rf7iM6PdsqfYCo3D1tnAR65SkLENyWJG1deUzpRMQmbh9F3eG")
	ent := new(Entry)
	ent.ChainID = "954d5a49fd70d9b8bcdb35d252267829957f7ef7fa6c74f88419bdc5e82209f4"
	ent.Content = []byte("test!")
	ent.ExtIDs = append(ent.ExtIDs, []byte("test"))
	newChain := NewChain(ent)

	cCommit, _ := ComposeChainCommit(newChain, ecAddr)
	r := new(response)
	json.Unmarshal(cCommit.Params, r)
	binCommit, _ := hex.DecodeString(r.Message)
	t.Logf("%x", binCommit)

	//the commit has a timestamp which is updated new for each time it is called.  This means it is different after each call.
	//we will check the non-changing parts

	if len(binCommit) != 200 {
		t.Error("expected commit to be 200 bytes long, instead got", len(binCommit))
	}
	result := binCommit[0:1]
	expected := []byte{0x00}
	if !bytes.Equal(result, expected) {
		t.Errorf("expected:%s\nrecieved:%s", expected, result)
	}
	//skip the 6 bytes of the timestamp
	result = binCommit[7:136]
	expected, _ = hex.DecodeString("516870d4c0e1ee2d5f0d415e51fc10ae6b8d895561e9314afdc33048194d76f07cc61c8a81aea23d76ff6447689757dc1e36af66e300ce3e06b8d816c79acfd2285ed45081d5b8819a678d13c7c2d04f704b34c74e8aaecd9bd34609bee047200b3b6a27bcceb6a42d62a3a8d02a6f0d73653215771de243a63ac048a18b59da29")

	if !bytes.Equal(result, expected) {
		t.Errorf("expected:%s\nrecieved:%s", expected, result)
	}
}

func TestComposeChainReveal(t *testing.T) {

	ent := new(Entry)
	ent.ChainID = "954d5a49fd70d9b8bcdb35d252267829957f7ef7fa6c74f88419bdc5e82209f4"
	ent.Content = []byte("test!")
	ent.ExtIDs = append(ent.ExtIDs, []byte("test"))
	newChain := NewChain(ent)

	cReveal, _ := ComposeChainReveal(newChain)

	expectedResponse := `{"entry":"00954d5a49fd70d9b8bcdb35d252267829957f7ef7fa6c74f88419bdc5e82209f400060004746573747465737421"}`
	if expectedResponse != string(cReveal.Params) {
		t.Errorf("expected:%s\nrecieved:%s", expectedResponse, cReveal.Params)
	}
}

func TestCommitChain(t *testing.T) {
	simlatedFactomdResponse := `{
   "jsonrpc":"2.0",
   "id":0,
   "result":{
      "message":"Chain Commit Success",
      "txid":"76e123d133a841fe3e08c5e3f3d392f8431f2d7668890c03f003f541efa8fc61"
   }
}`

	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		fmt.Fprintln(w, simlatedFactomdResponse)
	}))
	defer ts.Close()

	url := ts.URL[7:]
	SetFactomdServer(url)

	ent := new(Entry)
	ent.ChainID = "954d5a49fd70d9b8bcdb35d252267829957f7ef7fa6c74f88419bdc5e82209f4"
	ent.Content = []byte("test!")
	ent.ExtIDs = append(ent.ExtIDs, []byte("test"))
	newChain := NewChain(ent)
	ecAddr, _ := GetECAddress("Es2Rf7iM6PdsqfYCo3D1tnAR65SkLENyWJG1deUzpRMQmbh9F3eG")

	expectedResponse := "76e123d133a841fe3e08c5e3f3d392f8431f2d7668890c03f003f541efa8fc61"
	response, _ := CommitChain(newChain, ecAddr)

	if expectedResponse != response {
		t.Errorf("expected:%s\nrecieved:%s", expectedResponse, response)
	}
	t.Log(response)
}

func TestRevealChain(t *testing.T) {
	simlatedFactomdResponse := `{
  "jsonrpc": "2.0",
  "id": 0,
  "result": {
    "message": "Entry Reveal Success",
    "entryhash": "f5c956749fc3eba4acc60fd485fb100e601070a44fcce54ff358d60669854734"
  }
}`

	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		fmt.Fprintln(w, simlatedFactomdResponse)
	}))
	defer ts.Close()

	url := ts.URL[7:]
	SetFactomdServer(url)

	ent := new(Entry)
	ent.ChainID = "954d5a49fd70d9b8bcdb35d252267829957f7ef7fa6c74f88419bdc5e82209f4"
	ent.Content = []byte("test!")
	ent.ExtIDs = append(ent.ExtIDs, []byte("test"))
	newChain := NewChain(ent)

	expectedResponse := "f5c956749fc3eba4acc60fd485fb100e601070a44fcce54ff358d60669854734"
	response, _ := RevealChain(newChain)

	if expectedResponse != response {
		t.Errorf("expected:%s\nrecieved:%s", expectedResponse, response)
	}
	t.Log(response)
}
<commit_msg>test updates<commit_after>// Copyright 2016 Factom Foundation
// Use of this source code is governed by the MIT
// license that can be found in the LICENSE file.

package factom_test

import (
	"bytes"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"net/http"
	"net/http/httptest"
	"testing"

	. "github.com/FactomProject/factom"
)

var ()

func TestNewChain(t *testing.T) {
	ent := new(Entry)
	ent.ChainID = ""
	ent.Content = []byte("This is a test Entry.")
	ent.ExtIDs = append(ent.ExtIDs, []byte("This is the first extid."))
	ent.ExtIDs = append(ent.ExtIDs, []byte("This is the second extid."))

	newChain := NewChain(ent)
	expectedID := "5a402200c5cf278e47905ce52d7d64529a0291829a7bd230072c5468be709069"

	if newChain.ChainID != expectedID {
		t.Errorf("expected:%s\nrecieved:%s", expectedID, newChain.ChainID)
	}
	t.Log(newChain.ChainID)

	cfb := NewChainFromBytes(ent.Content, ent.ExtIDs...)
	if cfb.ChainID != expectedID {
		t.Errorf("expected:%s\nrecieved:%s", expectedID, cfb.ChainID)
	}
	t.Log(cfb.ChainID)

	cfs := NewChainFromStrings(
		"This is a test Entry.",
		"This is the first extid.",
		"This is the second extid.",
	)
	if cfs.ChainID != expectedID {
		t.Errorf("expected:%s\nrecieved:%s", expectedID, cfs.ChainID)
	}
	t.Log(cfs.ChainID)
}

func TestIfExists(t *testing.T) {
	simlatedFactomdResponse := `{
  "jsonrpc": "2.0",
  "id": 0,
  "result": {
    "ChainHead": "f65f67774139fa78344dcdd302631a0d646db0c2be4d58e3e48b2a188c1b856c"
  }
}`

	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		fmt.Fprintln(w, simlatedFactomdResponse)
	}))
	defer ts.Close()

	url := ts.URL[7:]
	SetFactomdServer(url)

	expectedID := "f65f67774139fa78344dcdd302631a0d646db0c2be4d58e3e48b2a188c1b856c"
	//fmt.Println(ChainExists(expectedID))
	if ChainExists(expectedID) != true {
		t.Errorf("chain %s does not exist", expectedID)
	}
}

func TestIfNotExists(t *testing.T) {
	simlatedFactomdResponse := `{"jsonrpc":"2.0","id":0,"error":{"code":-32009,"message":"Missing Chain Head"}}`

	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		fmt.Fprintln(w, simlatedFactomdResponse)
	}))
	defer ts.Close()

	url := ts.URL[7:]
	SetFactomdServer(url)
	unexpectedID := "5a402200c5cf278e47905ce52d7d64529a0291829a7bd230072c5468be709069"

	if ChainExists(unexpectedID) != false {
		t.Errorf("chain %s shouldn't exist", unexpectedID)
	}
}

func TestComposeChainCommit(t *testing.T) {
	type response struct {
		Message string `json:"message"`
	}
	ecAddr, err := GetECAddress("Es2Rf7iM6PdsqfYCo3D1tnAR65SkLENyWJG1deUzpRMQmbh9F3eG")
	if err != nil {
		t.Error(err)
	}

	ent := new(Entry)
	ent.ChainID = "954d5a49fd70d9b8bcdb35d252267829957f7ef7fa6c74f88419bdc5e82209f4"
	ent.Content = []byte("test!")
	ent.ExtIDs = append(ent.ExtIDs, []byte("test"))
	newChain := NewChain(ent)

	cCommit, err := ComposeChainCommit(newChain, ecAddr)
	if err != nil {
		t.Error(err)
	}

	r := new(response)
	json.Unmarshal(cCommit.Params, r)
	binCommit, _ := hex.DecodeString(r.Message)
	t.Logf("%x", binCommit)

	//the commit has a timestamp which is updated new for each time it is called.  This means it is different after each call.
	//we will check the non-changing parts

	if len(binCommit) != 200 {
		t.Error("expected commit to be 200 bytes long, instead got", len(binCommit))
	}
	result := binCommit[0:1]
	expected := []byte{0x00}
	if !bytes.Equal(result, expected) {
		t.Errorf("expected:%s\nrecieved:%s", expected, result)
	}
	//skip the 6 bytes of the timestamp
	result = binCommit[7:136]
	expected, err = hex.DecodeString("516870d4c0e1ee2d5f0d415e51fc10ae6b8d895561e9314afdc33048194d76f07cc61c8a81aea23d76ff6447689757dc1e36af66e300ce3e06b8d816c79acfd2285ed45081d5b8819a678d13c7c2d04f704b34c74e8aaecd9bd34609bee047200b3b6a27bcceb6a42d62a3a8d02a6f0d73653215771de243a63ac048a18b59da29")
	if err != nil {
		t.Error(err)
	}

	if !bytes.Equal(result, expected) {
		t.Errorf("expected:%s\nrecieved:%s", expected, result)
	}
}

func TestComposeChainReveal(t *testing.T) {

	ent := new(Entry)
	ent.ChainID = "954d5a49fd70d9b8bcdb35d252267829957f7ef7fa6c74f88419bdc5e82209f4"
	ent.Content = []byte("test!")
	ent.ExtIDs = append(ent.ExtIDs, []byte("test"))
	newChain := NewChain(ent)

	cReveal, err := ComposeChainReveal(newChain)
	if err != nil {
		t.Error(err)
	}

	expectedResponse := `{"entry":"00954d5a49fd70d9b8bcdb35d252267829957f7ef7fa6c74f88419bdc5e82209f400060004746573747465737421"}`
	if expectedResponse != string(cReveal.Params) {
		t.Errorf("expected:%s\nrecieved:%s", expectedResponse, cReveal.Params)
	}
}

func TestCommitChain(t *testing.T) {
	simlatedFactomdResponse := `{
   "jsonrpc":"2.0",
   "id":0,
   "result":{
      "message":"Chain Commit Success",
      "txid":"76e123d133a841fe3e08c5e3f3d392f8431f2d7668890c03f003f541efa8fc61"
   }
}`

	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		fmt.Fprintln(w, simlatedFactomdResponse)
	}))
	defer ts.Close()

	url := ts.URL[7:]
	SetFactomdServer(url)

	ent := new(Entry)
	ent.ChainID = "954d5a49fd70d9b8bcdb35d252267829957f7ef7fa6c74f88419bdc5e82209f4"
	ent.Content = []byte("test!")
	ent.ExtIDs = append(ent.ExtIDs, []byte("test"))
	newChain := NewChain(ent)
	ecAddr, err := GetECAddress("Es2Rf7iM6PdsqfYCo3D1tnAR65SkLENyWJG1deUzpRMQmbh9F3eG")
	if err != nil {
		t.Error(err)
	}

	expectedResponse := "76e123d133a841fe3e08c5e3f3d392f8431f2d7668890c03f003f541efa8fc61"
	response, _ := CommitChain(newChain, ecAddr)

	if expectedResponse != response {
		t.Errorf("expected:%s\nrecieved:%s", expectedResponse, response)
	}
	t.Log(response)
}

func TestRevealChain(t *testing.T) {
	simlatedFactomdResponse := `{
  "jsonrpc": "2.0",
  "id": 0,
  "result": {
    "message": "Entry Reveal Success",
    "entryhash": "f5c956749fc3eba4acc60fd485fb100e601070a44fcce54ff358d60669854734"
  }
}`

	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		fmt.Fprintln(w, simlatedFactomdResponse)
	}))
	defer ts.Close()

	url := ts.URL[7:]
	SetFactomdServer(url)

	ent := new(Entry)
	ent.ChainID = "954d5a49fd70d9b8bcdb35d252267829957f7ef7fa6c74f88419bdc5e82209f4"
	ent.Content = []byte("test!")
	ent.ExtIDs = append(ent.ExtIDs, []byte("test"))
	newChain := NewChain(ent)

	expectedResponse := "f5c956749fc3eba4acc60fd485fb100e601070a44fcce54ff358d60669854734"
	response, err := RevealChain(newChain)
	if err != nil {
		t.Error(err)
	}

	if expectedResponse != response {
		t.Errorf("expected:%s\nrecieved:%s", expectedResponse, response)
	}
	t.Log(response)
}
<|endoftext|>


<commit_before>package calcium

import (
	"bufio"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"sync"

	"github.com/pkg/errors"
	"github.com/projecteru2/core/engine"
	enginetypes "github.com/projecteru2/core/engine/types"
	"github.com/projecteru2/core/log"
	"github.com/projecteru2/core/types"
	"github.com/projecteru2/core/utils"
	"golang.org/x/net/context"
)

var winchCommand = []byte{0x80}  // 128, non-ASCII
var escapeCommand = []byte{0x1d} // 29, ^]

type window struct {
	Height uint `json:"Row"`
	Width  uint `json:"Col"`
}

func execuateInside(ctx context.Context, client engine.API, ID, cmd, user string, env []string, privileged bool) ([]byte, error) {
	cmds := utils.MakeCommandLineArgs(cmd)
	execConfig := &enginetypes.ExecConfig{
		User:         user,
		Cmd:          cmds,
		Privileged:   privileged,
		Env:          env,
		AttachStderr: true,
		AttachStdout: true,
	}
	b := []byte{}
	execID, stdout, stderr, _, err := client.Execute(ctx, ID, execConfig)
	if err != nil {
		return nil, errors.WithStack(err)
	}

	for m := range processStdStream(ctx, stdout, stderr, bufio.ScanLines, byte('\n')) {
		b = append(b, m.Data...)
	}

	exitCode, err := client.ExecExitCode(ctx, execID)
	if err != nil {
		return b, errors.WithStack(err)
	}
	if exitCode != 0 {
		return b, errors.WithStack(fmt.Errorf("%s", b))
	}
	return b, nil
}

func distributionInspect(ctx context.Context, node *types.Node, image string, digests []string) bool {
	remoteDigest, err := node.Engine.ImageRemoteDigest(ctx, image)
	if err != nil {
		log.Errorf(ctx, "[distributionInspect] get manifest failed %v", err)
		return false
	}

	for _, digest := range digests {
		if digest == remoteDigest {
			log.Debugf(ctx, "[distributionInspect] Local digest %s", digest)
			log.Debugf(ctx, "[distributionInspect] Remote digest %s", remoteDigest)
			return true
		}
	}
	return false
}

// Pull an image
func pullImage(ctx context.Context, node *types.Node, image string) error {
	log.Infof(ctx, "[pullImage] Pulling image %s", image)
	if image == "" {
		return errors.WithStack(types.ErrNoImage)
	}

	// check local
	exists := false
	digests, err := node.Engine.ImageLocalDigests(ctx, image)
	if err != nil {
		log.Errorf(ctx, "[pullImage] Check image failed %v", err)
	} else {
		log.Debug(ctx, "[pullImage] Local Image exists")
		exists = true
	}

	if exists && distributionInspect(ctx, node, image, digests) {
		log.Debug(ctx, "[pullImage] Image cached, skip pulling")
		return nil
	}

	log.Info("[pullImage] Image not cached, pulling")
	rc, err := node.Engine.ImagePull(ctx, image, false)
	defer utils.EnsureReaderClosed(ctx, rc)
	if err != nil {
		log.Errorf(ctx, "[pullImage] Error during pulling image %s: %v", image, err)
		return errors.WithStack(err)
	}
	log.Infof(ctx, "[pullImage] Done pulling image %s", image)
	return nil
}

func makeCopyMessage(id, name, path string, err error, data io.ReadCloser) *types.CopyMessage {
	return &types.CopyMessage{
		ID:    id,
		Name:  name,
		Path:  path,
		Error: err,
		Data:  data,
	}
}

func processVirtualizationInStream(
	ctx context.Context,
	inStream io.WriteCloser,
	inCh <-chan []byte,
	resizeFunc func(height, width uint) error,
) <-chan struct{} { // nolint
	specialPrefixCallback := map[string]func([]byte){
		string(winchCommand): func(body []byte) {
			w := &window{}
			if err := json.Unmarshal(body, w); err != nil {
				log.Errorf(ctx, "[processVirtualizationInStream] invalid winch command: %q", body)
				return
			}
			if err := resizeFunc(w.Height, w.Width); err != nil {
				log.Errorf(ctx, "[processVirtualizationInStream] resize window error: %v", err)
				return
			}
		},

		string(escapeCommand): func(_ []byte) {
			inStream.Close()
		},
	}
	return rawProcessVirtualizationInStream(ctx, inStream, inCh, specialPrefixCallback)
}

func rawProcessVirtualizationInStream(
	ctx context.Context,
	inStream io.WriteCloser,
	inCh <-chan []byte,
	specialPrefixCallback map[string]func([]byte),
) <-chan struct{} {
	done := make(chan struct{})
	utils.SentryGo(func() {
		defer close(done)
		defer inStream.Close()

		for cmd := range inCh {
			if len(cmd) == 0 {
				continue
			}
			if f, ok := specialPrefixCallback[string(cmd[:1])]; ok {
				f(cmd[1:])
				continue
			}
			if _, err := inStream.Write(cmd); err != nil {
				log.Errorf(ctx, "[rawProcessVirtualizationInStream] failed to write virtual input stream: %v", err)
				return
			}
		}
	})

	return done
}

func processVirtualizationOutStream(
	ctx context.Context,
	outStream io.ReadCloser,
	splitFunc bufio.SplitFunc,
	split byte,

) <-chan []byte {
	outCh := make(chan []byte)
	utils.SentryGo(func() {
		defer close(outCh)
		if outStream == nil {
			return
		}
		defer outStream.Close()
		scanner := bufio.NewScanner(outStream)
		scanner.Split(splitFunc)
		for scanner.Scan() {
			bs := scanner.Bytes()
			if split != 0 {
				bs = append(bs, split)
			}
			outCh <- bs
		}
		if err := scanner.Err(); err != nil {
			log.Warnf(ctx, "[processVirtualizationOutStream] failed to read output from output stream: %v", err)
		}
	})
	return outCh
}

func processBuildImageStream(ctx context.Context, reader io.ReadCloser) chan *types.BuildImageMessage {
	ch := make(chan *types.BuildImageMessage)
	utils.SentryGo(func() {
		defer close(ch)
		defer utils.EnsureReaderClosed(ctx, reader)
		decoder := json.NewDecoder(reader)
		for {
			message := &types.BuildImageMessage{}
			err := decoder.Decode(message)
			if err != nil {
				if err != io.EOF {
					malformed, _ := ioutil.ReadAll(decoder.Buffered()) // TODO err check
					log.Errorf(ctx, "[processBuildImageStream] Decode image message failed %v, buffered: %s", err, string(malformed))
					message.Error = err.Error()
					ch <- message
				}
				break
			}
			ch <- message
		}
	})
	return ch
}

func processStdStream(ctx context.Context, stdout, stderr io.ReadCloser, splitFunc bufio.SplitFunc, split byte) chan types.StdStreamMessage {
	ch := make(chan types.StdStreamMessage)

	wg := sync.WaitGroup{}

	wg.Add(1)
	utils.SentryGo(func() {
		defer wg.Done()
		for data := range processVirtualizationOutStream(ctx, stdout, splitFunc, split) {
			ch <- types.StdStreamMessage{Data: data, StdStreamType: types.Stdout}
		}
	})

	wg.Add(1)
	utils.SentryGo(func() {
		defer wg.Done()
		for data := range processVirtualizationOutStream(ctx, stderr, splitFunc, split) {
			ch <- types.StdStreamMessage{Data: data, StdStreamType: types.Stderr}
		}
	})

	utils.SentryGo(func() {
		defer close(ch)
		wg.Wait()
	})

	return ch
}
<commit_msg>bugfix: consumer mustn't end until producer closes (#408)<commit_after>package calcium

import (
	"bufio"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"sync"

	"github.com/pkg/errors"
	"github.com/projecteru2/core/engine"
	enginetypes "github.com/projecteru2/core/engine/types"
	"github.com/projecteru2/core/log"
	"github.com/projecteru2/core/types"
	"github.com/projecteru2/core/utils"
	"golang.org/x/net/context"
)

var winchCommand = []byte{0x80}  // 128, non-ASCII
var escapeCommand = []byte{0x1d} // 29, ^]

type window struct {
	Height uint `json:"Row"`
	Width  uint `json:"Col"`
}

func execuateInside(ctx context.Context, client engine.API, ID, cmd, user string, env []string, privileged bool) ([]byte, error) {
	cmds := utils.MakeCommandLineArgs(cmd)
	execConfig := &enginetypes.ExecConfig{
		User:         user,
		Cmd:          cmds,
		Privileged:   privileged,
		Env:          env,
		AttachStderr: true,
		AttachStdout: true,
	}
	b := []byte{}
	execID, stdout, stderr, _, err := client.Execute(ctx, ID, execConfig)
	if err != nil {
		return nil, errors.WithStack(err)
	}

	for m := range processStdStream(ctx, stdout, stderr, bufio.ScanLines, byte('\n')) {
		b = append(b, m.Data...)
	}

	exitCode, err := client.ExecExitCode(ctx, execID)
	if err != nil {
		return b, errors.WithStack(err)
	}
	if exitCode != 0 {
		return b, errors.WithStack(fmt.Errorf("%s", b))
	}
	return b, nil
}

func distributionInspect(ctx context.Context, node *types.Node, image string, digests []string) bool {
	remoteDigest, err := node.Engine.ImageRemoteDigest(ctx, image)
	if err != nil {
		log.Errorf(ctx, "[distributionInspect] get manifest failed %v", err)
		return false
	}

	for _, digest := range digests {
		if digest == remoteDigest {
			log.Debugf(ctx, "[distributionInspect] Local digest %s", digest)
			log.Debugf(ctx, "[distributionInspect] Remote digest %s", remoteDigest)
			return true
		}
	}
	return false
}

// Pull an image
func pullImage(ctx context.Context, node *types.Node, image string) error {
	log.Infof(ctx, "[pullImage] Pulling image %s", image)
	if image == "" {
		return errors.WithStack(types.ErrNoImage)
	}

	// check local
	exists := false
	digests, err := node.Engine.ImageLocalDigests(ctx, image)
	if err != nil {
		log.Errorf(ctx, "[pullImage] Check image failed %v", err)
	} else {
		log.Debug(ctx, "[pullImage] Local Image exists")
		exists = true
	}

	if exists && distributionInspect(ctx, node, image, digests) {
		log.Debug(ctx, "[pullImage] Image cached, skip pulling")
		return nil
	}

	log.Info("[pullImage] Image not cached, pulling")
	rc, err := node.Engine.ImagePull(ctx, image, false)
	defer utils.EnsureReaderClosed(ctx, rc)
	if err != nil {
		log.Errorf(ctx, "[pullImage] Error during pulling image %s: %v", image, err)
		return errors.WithStack(err)
	}
	log.Infof(ctx, "[pullImage] Done pulling image %s", image)
	return nil
}

func makeCopyMessage(id, name, path string, err error, data io.ReadCloser) *types.CopyMessage {
	return &types.CopyMessage{
		ID:    id,
		Name:  name,
		Path:  path,
		Error: err,
		Data:  data,
	}
}

func processVirtualizationInStream(
	ctx context.Context,
	inStream io.WriteCloser,
	inCh <-chan []byte,
	resizeFunc func(height, width uint) error,
) <-chan struct{} { // nolint
	specialPrefixCallback := map[string]func([]byte){
		string(winchCommand): func(body []byte) {
			w := &window{}
			if err := json.Unmarshal(body, w); err != nil {
				log.Errorf(ctx, "[processVirtualizationInStream] invalid winch command: %q", body)
				return
			}
			if err := resizeFunc(w.Height, w.Width); err != nil {
				log.Errorf(ctx, "[processVirtualizationInStream] resize window error: %v", err)
				return
			}
		},

		string(escapeCommand): func(_ []byte) {
			inStream.Close()
		},
	}
	return rawProcessVirtualizationInStream(ctx, inStream, inCh, specialPrefixCallback)
}

func rawProcessVirtualizationInStream(
	ctx context.Context,
	inStream io.WriteCloser,
	inCh <-chan []byte,
	specialPrefixCallback map[string]func([]byte),
) <-chan struct{} {
	done := make(chan struct{})
	utils.SentryGo(func() {
		defer close(done)
		defer inStream.Close()

		for cmd := range inCh {
			if len(cmd) == 0 {
				continue
			}
			if f, ok := specialPrefixCallback[string(cmd[:1])]; ok {
				f(cmd[1:])
				continue
			}
			if _, err := inStream.Write(cmd); err != nil {
				log.Errorf(ctx, "[rawProcessVirtualizationInStream] failed to write virtual input stream: %v", err)
				continue
			}
		}
	})

	return done
}

func processVirtualizationOutStream(
	ctx context.Context,
	outStream io.ReadCloser,
	splitFunc bufio.SplitFunc,
	split byte,

) <-chan []byte {
	outCh := make(chan []byte)
	utils.SentryGo(func() {
		defer close(outCh)
		if outStream == nil {
			return
		}
		defer outStream.Close()
		scanner := bufio.NewScanner(outStream)
		scanner.Split(splitFunc)
		for scanner.Scan() {
			bs := scanner.Bytes()
			if split != 0 {
				bs = append(bs, split)
			}
			outCh <- bs
		}
		if err := scanner.Err(); err != nil {
			log.Warnf(ctx, "[processVirtualizationOutStream] failed to read output from output stream: %v", err)
		}
	})
	return outCh
}

func processBuildImageStream(ctx context.Context, reader io.ReadCloser) chan *types.BuildImageMessage {
	ch := make(chan *types.BuildImageMessage)
	utils.SentryGo(func() {
		defer close(ch)
		defer utils.EnsureReaderClosed(ctx, reader)
		decoder := json.NewDecoder(reader)
		for {
			message := &types.BuildImageMessage{}
			err := decoder.Decode(message)
			if err != nil {
				if err != io.EOF {
					malformed, _ := ioutil.ReadAll(decoder.Buffered()) // TODO err check
					log.Errorf(ctx, "[processBuildImageStream] Decode image message failed %v, buffered: %s", err, string(malformed))
					message.Error = err.Error()
					ch <- message
				}
				break
			}
			ch <- message
		}
	})
	return ch
}

func processStdStream(ctx context.Context, stdout, stderr io.ReadCloser, splitFunc bufio.SplitFunc, split byte) chan types.StdStreamMessage {
	ch := make(chan types.StdStreamMessage)

	wg := sync.WaitGroup{}

	wg.Add(1)
	utils.SentryGo(func() {
		defer wg.Done()
		for data := range processVirtualizationOutStream(ctx, stdout, splitFunc, split) {
			ch <- types.StdStreamMessage{Data: data, StdStreamType: types.Stdout}
		}
	})

	wg.Add(1)
	utils.SentryGo(func() {
		defer wg.Done()
		for data := range processVirtualizationOutStream(ctx, stderr, splitFunc, split) {
			ch <- types.StdStreamMessage{Data: data, StdStreamType: types.Stderr}
		}
	})

	utils.SentryGo(func() {
		defer close(ch)
		wg.Wait()
	})

	return ch
}
<|endoftext|>


<commit_before>package pkg

import (
	"encoding/json"
	"io/ioutil"
	"net/http"
	"regexp"
	"sync"

	"github.com/golang/protobuf/proto"

	"github.com/ketchuphq/ketchup/proto/ketchup/packages"
	"github.com/ketchuphq/ketchup/util/errors"
)

type Registry struct {
	URL      string
	Registry *packages.Registry

	mu sync.RWMutex
}

func (r *Registry) Proto() *packages.Registry {
	r.mu.RLock()
	defer r.mu.RUnlock()
	return proto.Clone(r.Registry).(*packages.Registry)
}

func (r *Registry) Sync() error {
	res, err := http.Get(r.URL)
	if err != nil {
		return errors.Wrap(err)
	}
	b, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return errors.Wrap(err)
	}
	repo := &packages.Registry{}
	err = json.Unmarshal(b, &repo)
	if err != nil {
		return errors.Wrap(err)
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	r.Registry = repo
	return nil
}

func (r *Registry) Search(name string) (*packages.Package, error) {
	err := r.Sync()
	if err != nil {
		return nil, err
	}
	r.mu.RLock()
	defer r.mu.RUnlock()
	for _, p := range r.Registry.Packages {
		if p.GetName() == name {
			return p, nil
		}
	}
	return nil, nil
}

func (r *Registry) Match(re *regexp.Regexp) ([]*packages.Package, error) {
	err := r.Sync()
	if err != nil {
		return nil, err
	}
	r.mu.RLock()
	defer r.mu.RUnlock()
	out := []*packages.Package{}
	for _, p := range r.Registry.Packages {
		if re.MatchString(p.GetName()) {
			out = append(out, p)
		}
	}
	return out, nil
}

// FetchDefaultRegistry fetches the default registry
func (m *Module) Registry(registryURL string) *Registry {
	return &Registry{URL: registryURL}
}

// press registry daemon should periodically scrape
func getGithubTags(p *packages.Package) {
	// paginate should cache
}

func getBitbucketTags(p *packages.Package) {
	// paginate should cache
}
<commit_msg>pkg: Comments and handle status code for registry syncing.<commit_after>package pkg

import (
	"encoding/json"
	"io/ioutil"
	"net/http"
	"regexp"
	"sync"

	"github.com/golang/protobuf/proto"

	"github.com/ketchuphq/ketchup/proto/ketchup/packages"
	"github.com/ketchuphq/ketchup/util/errors"
)

type Registry struct {
	URL      string
	Registry *packages.Registry

	mu sync.RWMutex
}

// Proto returns a clone of the underlying registry proto
func (r *Registry) Proto() *packages.Registry {
	r.mu.RLock()
	defer r.mu.RUnlock()
	return proto.Clone(r.Registry).(*packages.Registry)
}

// Sync the repo data from the source
func (r *Registry) Sync() error {
	res, err := http.Get(r.URL)
	if err != nil {
		return errors.Wrap(err)
	}
	if res.StatusCode > 299 {
		return errors.New("unexpected status code from %s: %d", r.URL, res.StatusCode)
	}
	b, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return errors.Wrap(err)
	}
	repo := &packages.Registry{}
	err = json.Unmarshal(b, &repo)
	if err != nil {
		return errors.Wrap(err)
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	r.Registry = repo
	return nil
}

// Search the repo for a package with the given name
func (r *Registry) Search(name string) (*packages.Package, error) {
	err := r.Sync()
	if err != nil {
		return nil, err
	}
	r.mu.RLock()
	defer r.mu.RUnlock()
	for _, p := range r.Registry.Packages {
		if p.GetName() == name {
			return p, nil
		}
	}
	return nil, nil
}

// Match searches the registry for all packages with name matching the
// given regex.
func (r *Registry) Match(re *regexp.Regexp) ([]*packages.Package, error) {
	err := r.Sync()
	if err != nil {
		return nil, err
	}
	r.mu.RLock()
	defer r.mu.RUnlock()
	out := []*packages.Package{}
	for _, p := range r.Registry.Packages {
		if re.MatchString(p.GetName()) {
			out = append(out, p)
		}
	}
	return out, nil
}

// Registry creates and returns a new registry for the given url
func (m *Module) Registry(registryURL string) *Registry {
	return &Registry{URL: registryURL}
}

// press registry daemon should periodically scrape
func getGithubTags(p *packages.Package) {
	// paginate should cache
}

func getBitbucketTags(p *packages.Package) {
	// paginate should cache
}
<|endoftext|>


<commit_before>package main

import (
	"bitbucket.org/anacrolix/go.torrent/dht"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"os/signal"
)

type pingResponse struct {
	addr string
	krpc dht.Msg
}

var (
	tableFileName = flag.String("tableFile", "", "name of file for storing node info")
	serveAddr     = flag.String("serveAddr", ":0", "local UDP address")
	infoHash      = flag.String("infoHash", "", "torrent infohash")

	s dht.Server
)

func loadTable() error {
	if *tableFileName == "" {
		return nil
	}
	f, err := os.Open(*tableFileName)
	if os.IsNotExist(err) {
		return nil
	}
	if err != nil {
		return fmt.Errorf("error opening table file: %s", err)
	}
	defer f.Close()
	added := 0
	for {
		b := make([]byte, dht.CompactNodeInfoLen)
		_, err := io.ReadFull(f, b)
		if err == io.EOF {
			break
		}
		if err != nil {
			return fmt.Errorf("error reading table file: %s", err)
		}
		var ni dht.NodeInfo
		err = ni.UnmarshalCompact(b)
		if err != nil {
			return fmt.Errorf("error unmarshaling compact node info: %s", err)
		}
		s.AddNode(ni)
		added++
	}
	log.Printf("loaded %d nodes from table file", added)
	return nil
}

func init() {
	log.SetFlags(log.LstdFlags | log.Lshortfile)
	flag.Parse()
	switch len(*infoHash) {
	case 20:
	case 40:
		if _, err := fmt.Sscanf(*infoHash, "%x", infoHash); err != nil {
			log.Fatal(err)
		}
	default:
		log.Fatal("require 20 byte infohash")
	}
	var err error
	s.Socket, err = net.ListenUDP("udp4", func() *net.UDPAddr {
		addr, err := net.ResolveUDPAddr("udp4", *serveAddr)
		if err != nil {
			log.Fatalf("error resolving serve addr: %s", err)
		}
		return addr
	}())
	if err != nil {
		log.Fatal(err)
	}
	s.Init()
	err = loadTable()
	if err != nil {
		log.Fatalf("error loading table: %s", err)
	}
	log.Printf("dht server on %s, ID is %q", s.Socket.LocalAddr(), s.IDString())
	setupSignals()
}

func saveTable() error {
	goodNodes := s.Nodes()
	if *tableFileName == "" {
		if len(goodNodes) != 0 {
			log.Printf("discarding %d good nodes!", len(goodNodes))
		}
		return nil
	}
	f, err := os.OpenFile(*tableFileName, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0666)
	if err != nil {
		return fmt.Errorf("error opening table file: %s", err)
	}
	defer f.Close()
	for _, nodeInfo := range goodNodes {
		var b [dht.CompactNodeInfoLen]byte
		err := nodeInfo.PutCompact(b[:])
		if err != nil {
			return fmt.Errorf("error compacting node info: %s", err)
		}
		_, err = f.Write(b[:])
		if err != nil {
			return fmt.Errorf("error writing compact node info: %s", err)
		}
	}
	log.Printf("saved %d nodes to table file", len(goodNodes))
	return nil
}

func setupSignals() {
	ch := make(chan os.Signal)
	signal.Notify(ch)
	go func() {
		<-ch
		s.StopServing()
	}()
}

func main() {
	// go s.Bootstrap()
	go func() {
		ps, err := s.GetPeers(*infoHash)
		if err != nil {
			log.Fatal(err)
		}
		for sl := range ps.Values {
			for _, p := range sl {
				fmt.Println(p)
			}
		}
		s.StopServing()
	}()
	err := s.Serve()
	if err := saveTable(); err != nil {
		log.Printf("error saving node table: %s", err)
	}
	if err != nil {
		log.Fatalf("error serving dht: %s", err)
	}
}
<commit_msg>dht-get-peers: Some improvements<commit_after>package main

import (
	"bitbucket.org/anacrolix/go.torrent/dht"
	"bitbucket.org/anacrolix/go.torrent/tracker"
	_ "bitbucket.org/anacrolix/go.torrent/util/profile"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"os/signal"
)

type pingResponse struct {
	addr string
	krpc dht.Msg
}

var (
	tableFileName = flag.String("tableFile", "", "name of file for storing node info")
	serveAddr     = flag.String("serveAddr", ":0", "local UDP address")
	infoHash      = flag.String("infoHash", "", "torrent infohash")

	s dht.Server
)

func loadTable() error {
	if *tableFileName == "" {
		return nil
	}
	f, err := os.Open(*tableFileName)
	if os.IsNotExist(err) {
		return nil
	}
	if err != nil {
		return fmt.Errorf("error opening table file: %s", err)
	}
	defer f.Close()
	added := 0
	for {
		b := make([]byte, dht.CompactNodeInfoLen)
		_, err := io.ReadFull(f, b)
		if err == io.EOF {
			break
		}
		if err != nil {
			return fmt.Errorf("error reading table file: %s", err)
		}
		var ni dht.NodeInfo
		err = ni.UnmarshalCompact(b)
		if err != nil {
			return fmt.Errorf("error unmarshaling compact node info: %s", err)
		}
		s.AddNode(ni)
		added++
	}
	log.Printf("loaded %d nodes from table file", added)
	return nil
}

func init() {
	log.SetFlags(log.LstdFlags | log.Lshortfile)
	flag.Parse()
	switch len(*infoHash) {
	case 20:
	case 40:
		_, err := fmt.Sscanf(*infoHash, "%x", infoHash)
		if err != nil {
			log.Fatal(err)
		}
	default:
		log.Fatal("require 20 byte infohash")
	}
	var err error
	s.Socket, err = net.ListenUDP("udp4", func() *net.UDPAddr {
		addr, err := net.ResolveUDPAddr("udp4", *serveAddr)
		if err != nil {
			log.Fatalf("error resolving serve addr: %s", err)
		}
		return addr
	}())
	if err != nil {
		log.Fatal(err)
	}
	s.Init()
	err = loadTable()
	if err != nil {
		log.Fatalf("error loading table: %s", err)
	}
	log.Printf("dht server on %s, ID is %q", s.Socket.LocalAddr(), s.IDString())
	setupSignals()
}

func saveTable() error {
	goodNodes := s.Nodes()
	if *tableFileName == "" {
		if len(goodNodes) != 0 {
			log.Printf("discarding %d good nodes!", len(goodNodes))
		}
		return nil
	}
	f, err := os.OpenFile(*tableFileName, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0666)
	if err != nil {
		return fmt.Errorf("error opening table file: %s", err)
	}
	defer f.Close()
	for _, nodeInfo := range goodNodes {
		var b [dht.CompactNodeInfoLen]byte
		err := nodeInfo.PutCompact(b[:])
		if err != nil {
			return fmt.Errorf("error compacting node info: %s", err)
		}
		_, err = f.Write(b[:])
		if err != nil {
			return fmt.Errorf("error writing compact node info: %s", err)
		}
	}
	log.Printf("saved %d nodes to table file", len(goodNodes))
	return nil
}

func setupSignals() {
	ch := make(chan os.Signal)
	signal.Notify(ch, os.Interrupt)
	go func() {
		<-ch
		s.StopServing()
	}()
}

func main() {
	go func() {
		defer s.StopServing()
		if err := s.Bootstrap(); err != nil {
			log.Printf("error bootstrapping: %s", err)
			return
		}
		saveTable()
		ps, err := s.GetPeers(*infoHash)
		if err != nil {
			log.Fatal(err)
		}
		seen := make(map[tracker.CompactPeer]struct{})
		for sl := range ps.Values {
			for _, p := range sl {
				if _, ok := seen[p]; ok {
					continue
				}
				seen[p] = struct{}{}
				fmt.Println((&net.UDPAddr{
					IP:   p.IP[:],
					Port: int(p.Port),
				}).String())
			}
		}
	}()
	err := s.Serve()
	if err := saveTable(); err != nil {
		log.Printf("error saving node table: %s", err)
	}
	if err != nil {
		log.Fatalf("error serving dht: %s", err)
	}
}
<|endoftext|>


<commit_before>package command

import (
	"encoding/json"
	"flag"
	"fmt"
	"net/http"
	"net/url"
	"sort"
	"strings"
	"time"

	"github.com/funkygao/gafka/ctx"
	"github.com/funkygao/gocli"
	"github.com/funkygao/golib/gofmt"
	"github.com/funkygao/golib/top"
	"github.com/ryanuber/columnize"
)

type Haproxy struct {
	Ui  cli.Ui
	Cmd string

	zone string
}

func (this *Haproxy) Run(args []string) (exitCode int) {
	var topMode bool
	cmdFlags := flag.NewFlagSet("haproxy", flag.ContinueOnError)
	cmdFlags.Usage = func() { this.Ui.Output(this.Help()) }
	cmdFlags.StringVar(&this.zone, "z", ctx.DefaultZone(), "")
	cmdFlags.BoolVar(&topMode, "top", true, "")
	if err := cmdFlags.Parse(args); err != nil {
		return 1
	}

	zone := ctx.Zone(this.zone)
	if topMode {
		header, _ := this.getStats(zone.HaProxyStatsUri[0])
		t := top.New(header, "%8s %4s %15s %15s %8s %6s %8s %10s %8s %8s %5s %7s %9s %6s")
		go func() {
			for {
				rows := make([]string, 0)
				for _, uri := range zone.HaProxyStatsUri {
					_, r := this.getStats(uri)
					rows = append(rows, r...)
				}
				t.Refresh(rows)

				time.Sleep(time.Second * 3)
			}
		}()
		if err := t.Start(); err != nil {
			panic(err)
		}
	} else {
		for _, uri := range zone.HaProxyStatsUri {
			this.fetchStats(uri)
		}
	}

	return
}

func (*Haproxy) Synopsis() string {
	return "Query ehaproxy cluster for load stats"
}

func (this *Haproxy) getStats(statsUri string) (header string, rows []string) {
	client := http.Client{Timeout: time.Second * 30}
	resp, err := client.Get(statsUri)
	swallow(err)
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		swallow(fmt.Errorf("fetch[%s] stats got status: %d", resp.StatusCode))
	}

	var records map[string]map[string]int64
	reader := json.NewDecoder(resp.Body)
	err = reader.Decode(&records)
	swallow(err)

	u, err := url.Parse(statsUri)
	swallow(err)
	var shortHostname string
	if strings.Contains(u.Host, ":") {
		u.Host = u.Host[:strings.Index(u.Host, ":")]
	}
	tuples := strings.SplitN(u.Host, ".", 4)
	if len(tuples) < 4 {
		shortHostname = u.Host
	} else {
		shortHostname = tuples[3]
	}
	if len(shortHostname) > 8 {
		shortHostname = shortHostname[:8]
	}

	sortedSvcs := make([]string, 0)
	for svc, _ := range records {
		sortedSvcs = append(sortedSvcs, svc)
	}
	sort.Strings(sortedSvcs)

	sortedCols := make([]string, 0)
	for k, _ := range records["pub"] {
		sortedCols = append(sortedCols, k)
	}
	sort.Strings(sortedCols)

	header = strings.Join(append([]string{"host", "svc"}, sortedCols...), "|")
	for _, svc := range sortedSvcs {
		stats := records[svc]

		var vals = []string{shortHostname, svc}
		for _, k := range sortedCols {
			v := stats[k]

			vals = append(vals, gofmt.Comma(v))
		}

		rows = append(rows, strings.Join(vals, "|"))
	}

	return
}

func (this *Haproxy) fetchStats(statsUri string) {
	client := http.Client{Timeout: time.Second * 30}
	resp, err := client.Get(statsUri)
	swallow(err)
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		swallow(fmt.Errorf("fetch[%s] stats got status: %d", resp.StatusCode))
	}

	var records map[string]map[string]int64
	reader := json.NewDecoder(resp.Body)
	err = reader.Decode(&records)
	swallow(err)

	u, err := url.Parse(statsUri)
	swallow(err)
	this.Ui.Info(u.Host)

	sortedSvcs := make([]string, 0)
	for svc, _ := range records {
		sortedSvcs = append(sortedSvcs, svc)
	}
	sort.Strings(sortedSvcs)

	sortedCols := make([]string, 0)
	for k, _ := range records["pub"] {
		sortedCols = append(sortedCols, k)
	}
	sort.Strings(sortedCols)

	lines := []string{strings.Join(append([]string{"svc"}, sortedCols...), "|")}
	for _, svc := range sortedSvcs {
		stats := records[svc]

		var vals = []string{svc}
		for _, k := range sortedCols {
			v := stats[k]

			vals = append(vals, gofmt.Comma(v))
		}

		lines = append(lines, strings.Join(vals, "|"))
	}

	this.Ui.Output(columnize.SimpleFormat(lines))
}

func (this *Haproxy) Help() string {
	help := fmt.Sprintf(`
Usage: %s haproxy [options]

    %s

Options:

    -z zone

    -top
      Top mode

`, this.Cmd, this.Synopsis())
	return strings.TrimSpace(help)
}
<commit_msg>refmt<commit_after>package command

import (
	"encoding/json"
	"flag"
	"fmt"
	"net/http"
	"net/url"
	"sort"
	"strings"
	"time"

	"github.com/funkygao/gafka/ctx"
	"github.com/funkygao/gocli"
	"github.com/funkygao/golib/gofmt"
	"github.com/funkygao/golib/top"
	"github.com/ryanuber/columnize"
)

type Haproxy struct {
	Ui  cli.Ui
	Cmd string

	zone string
}

func (this *Haproxy) Run(args []string) (exitCode int) {
	var topMode bool
	cmdFlags := flag.NewFlagSet("haproxy", flag.ContinueOnError)
	cmdFlags.Usage = func() { this.Ui.Output(this.Help()) }
	cmdFlags.StringVar(&this.zone, "z", ctx.DefaultZone(), "")
	cmdFlags.BoolVar(&topMode, "top", true, "")
	if err := cmdFlags.Parse(args); err != nil {
		return 1
	}

	zone := ctx.Zone(this.zone)
	if topMode {
		header, _ := this.getStats(zone.HaProxyStatsUri[0])
		t := top.New(header, "%8s %4s %21s %21s %9s %6s %8s %12s %8s %8s %7s %7s %14s %6s")
		go func() {
			for {
				rows := make([]string, 0)
				for _, uri := range zone.HaProxyStatsUri {
					_, r := this.getStats(uri)
					rows = append(rows, r...)
				}
				t.Refresh(rows)

				time.Sleep(time.Second * 3)
			}
		}()
		if err := t.Start(); err != nil {
			panic(err)
		}
	} else {
		for _, uri := range zone.HaProxyStatsUri {
			this.fetchStats(uri)
		}
	}

	return
}

func (*Haproxy) Synopsis() string {
	return "Query ehaproxy cluster for load stats"
}

func (this *Haproxy) getStats(statsUri string) (header string, rows []string) {
	client := http.Client{Timeout: time.Second * 30}
	resp, err := client.Get(statsUri)
	swallow(err)
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		swallow(fmt.Errorf("fetch[%s] stats got status: %d", resp.StatusCode))
	}

	var records map[string]map[string]int64
	reader := json.NewDecoder(resp.Body)
	err = reader.Decode(&records)
	swallow(err)

	u, err := url.Parse(statsUri)
	swallow(err)
	var shortHostname string
	if strings.Contains(u.Host, ":") {
		u.Host = u.Host[:strings.Index(u.Host, ":")]
	}
	tuples := strings.SplitN(u.Host, ".", 4)
	if len(tuples) < 4 {
		shortHostname = u.Host
	} else {
		shortHostname = tuples[3]
	}
	if len(shortHostname) > 8 {
		shortHostname = shortHostname[:8]
	}

	sortedSvcs := make([]string, 0)
	for svc, _ := range records {
		sortedSvcs = append(sortedSvcs, svc)
	}
	sort.Strings(sortedSvcs)

	sortedCols := make([]string, 0)
	for k, _ := range records["pub"] {
		sortedCols = append(sortedCols, k)
	}
	sort.Strings(sortedCols)

	header = strings.Join(append([]string{"host", "svc"}, sortedCols...), "|")
	for _, svc := range sortedSvcs {
		stats := records[svc]

		var vals = []string{shortHostname, svc}
		for _, k := range sortedCols {
			v := stats[k]

			vals = append(vals, gofmt.Comma(v))
		}

		rows = append(rows, strings.Join(vals, "|"))
	}

	return
}

func (this *Haproxy) fetchStats(statsUri string) {
	client := http.Client{Timeout: time.Second * 30}
	resp, err := client.Get(statsUri)
	swallow(err)
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		swallow(fmt.Errorf("fetch[%s] stats got status: %d", resp.StatusCode))
	}

	var records map[string]map[string]int64
	reader := json.NewDecoder(resp.Body)
	err = reader.Decode(&records)
	swallow(err)

	u, err := url.Parse(statsUri)
	swallow(err)
	this.Ui.Info(u.Host)

	sortedSvcs := make([]string, 0)
	for svc, _ := range records {
		sortedSvcs = append(sortedSvcs, svc)
	}
	sort.Strings(sortedSvcs)

	sortedCols := make([]string, 0)
	for k, _ := range records["pub"] {
		sortedCols = append(sortedCols, k)
	}
	sort.Strings(sortedCols)

	lines := []string{strings.Join(append([]string{"svc"}, sortedCols...), "|")}
	for _, svc := range sortedSvcs {
		stats := records[svc]

		var vals = []string{svc}
		for _, k := range sortedCols {
			v := stats[k]

			vals = append(vals, gofmt.Comma(v))
		}

		lines = append(lines, strings.Join(vals, "|"))
	}

	this.Ui.Output(columnize.SimpleFormat(lines))
}

func (this *Haproxy) Help() string {
	help := fmt.Sprintf(`
Usage: %s haproxy [options]

    %s

Options:

    -z zone

    -top
      Top mode

`, this.Cmd, this.Synopsis())
	return strings.TrimSpace(help)
}
<|endoftext|>