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

// These examples demonstrate more intricate uses of the pflag package.
package pflag_test

import (
	"errors"
	"fmt"
	"strings"
	"time"

	"github.com/spf13/pflag"
)

// Example 1: A single string flag called "species" with default value "gopher".
var species = pflag.String("species", "gopher", "the species we are studying")

// Example 2: A flag with a shorthand letter.
var gopherType = pflag.StringP("gopher_type", "g", "pocket", "the variety of gopher")

// Example 3: A user-defined flag type, a slice of durations.
type interval []time.Duration

// String is the method to format the flag's value, part of the flag.Value interface.
// The String method's output will be used in diagnostics.
func (i *interval) String() string {
	return fmt.Sprint(*i)
}

func (i *interval) Type() string {
	return "interval"
}

// Set is the method to set the flag value, part of the flag.Value interface.
// Set's argument is a string to be parsed to set the flag.
// It's a comma-separated list, so we split it.
func (i *interval) Set(value string) error {
	// If we wanted to allow the flag to be set multiple times,
	// accumulating values, we would delete this if statement.
	// That would permit usages such as
	//	-deltaT 10s -deltaT 15s
	// and other combinations.
	if len(*i) > 0 {
		return errors.New("interval flag already set")
	}
	for _, dt := range strings.Split(value, ",") {
		duration, err := time.ParseDuration(dt)
		if err != nil {
			return err
		}
		*i = append(*i, duration)
	}
	return nil
}

// Define a flag to accumulate durations. Because it has a special type,
// we need to use the Var function and therefore create the flag during
// init.

var intervalFlag interval

func init() {
	// Tie the command-line flag to the intervalFlag variable and
	// set a usage message.
	pflag.Var(&intervalFlag, "deltaT", "comma-separated list of intervals to use between events")
}

func Example() {
	// All the interesting pieces are with the variables declared above, but
	// to enable the pflag package to see the flags defined there, one must
	// execute, typically at the start of main (not init!):
	//	pflag.Parse()
	// We don't run it here because this is not a main function and
	// the testing suite has already parsed the flags.
}

func ExampleShorthandLookup() {
	fs := pflag.NewFlagSet("Example", pflag.ContinueOnError)
	fs.BoolP("verbose", "v", false, "verbose output")

	name := "verbose"
	flag := fs.ShorthandLookup(name[:1])

	fmt.Println(flag.Name)
}
<commit_msg>Add example to FlagSet.ShorthandLookup<commit_after>// Copyright 2012 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// These examples demonstrate more intricate uses of the pflag package.
package pflag_test

import (
	"errors"
	"fmt"
	"strings"
	"time"

	"github.com/spf13/pflag"
)

// Example 1: A single string flag called "species" with default value "gopher".
var species = pflag.String("species", "gopher", "the species we are studying")

// Example 2: A flag with a shorthand letter.
var gopherType = pflag.StringP("gopher_type", "g", "pocket", "the variety of gopher")

// Example 3: A user-defined flag type, a slice of durations.
type interval []time.Duration

// String is the method to format the flag's value, part of the flag.Value interface.
// The String method's output will be used in diagnostics.
func (i *interval) String() string {
	return fmt.Sprint(*i)
}

func (i *interval) Type() string {
	return "interval"
}

// Set is the method to set the flag value, part of the flag.Value interface.
// Set's argument is a string to be parsed to set the flag.
// It's a comma-separated list, so we split it.
func (i *interval) Set(value string) error {
	// If we wanted to allow the flag to be set multiple times,
	// accumulating values, we would delete this if statement.
	// That would permit usages such as
	//	-deltaT 10s -deltaT 15s
	// and other combinations.
	if len(*i) > 0 {
		return errors.New("interval flag already set")
	}
	for _, dt := range strings.Split(value, ",") {
		duration, err := time.ParseDuration(dt)
		if err != nil {
			return err
		}
		*i = append(*i, duration)
	}
	return nil
}

// Define a flag to accumulate durations. Because it has a special type,
// we need to use the Var function and therefore create the flag during
// init.

var intervalFlag interval

func init() {
	// Tie the command-line flag to the intervalFlag variable and
	// set a usage message.
	pflag.Var(&intervalFlag, "deltaT", "comma-separated list of intervals to use between events")
}

func Example() {
	// All the interesting pieces are with the variables declared above, but
	// to enable the pflag package to see the flags defined there, one must
	// execute, typically at the start of main (not init!):
	//	pflag.Parse()
	// We don't run it here because this is not a main function and
	// the testing suite has already parsed the flags.
}

func ExampleShorthandLookup() {
	pflag.BoolP("verbose", "v", false, "verbose output")

	name := "verbose"
	flag := pflag.ShorthandLookup(name[:1])

	fmt.Println(flag.Name)
}

func ExampleFlagSet_ShorthandLookup() {
	fs := pflag.NewFlagSet("Example", pflag.ContinueOnError)
	fs.BoolP("verbose", "v", false, "verbose output")

	name := "verbose"
	flag := fs.ShorthandLookup(name[:1])

	fmt.Println(flag.Name)
}
<|endoftext|>


<commit_before>// Package crc16 implements the 16-bit cyclic redundancy check, or CRC-16, checksum.
//
// It provides parameters for the majority of well-known CRC-16 algorithms.
package crc16

import "github.com/sigurn/utils"

// Params represents parameters of CRC-16 algorithms.
// More information about algorithms parametrization and parameter descriptions
// can be found here - http://www.zlib.net/crc_v3.txt
type Params struct {
	Poly   uint16
	Init   uint16
	RefIn  bool
	RefOut bool
	XorOut uint16
	Check  uint16
	Name   string
}

// Predefined CRC-16 algorithms.
// List of algorithms with their parameters borrowed from here -  http://reveng.sourceforge.net/crc-catalogue/16.htm
//
// The variables can be used to create Table for the selected algorithm.
var (
	CRC16_ARC         = Params{0x8005, 0x0000, true, true, 0x0000, 0xBB3D, "CRC-16/ARC"}
	CRC16_AUG_CCITT   = Params{0x1021, 0x1D0F, false, false, 0x0000, 0xE5CC, "CRC-16/AUG-CCITT"}
	CRC16_BUYPASS     = Params{0x8005, 0x0000, false, false, 0x0000, 0xFEE8, "CRC-16/BUYPASS"}
	CRC16_CCITT_FALSE = Params{0x1021, 0xFFFF, false, false, 0x0000, 0x29B1, "CRC-16/CCITT-FALSE"}
	CRC16_CDMA2000    = Params{0xC867, 0xFFFF, false, false, 0x0000, 0x4C06, "CRC-16/CDMA2000"}
	CRC16_DDS_110     = Params{0x8005, 0x800D, false, false, 0x0000, 0x9ECF, "CRC-16/DDS-110"}
	CRC16_DECT_R      = Params{0x0589, 0x0000, false, false, 0x0001, 0x007E, "CRC-16/DECT-R"}
	CRC16_DECT_X      = Params{0x0589, 0x0000, false, false, 0x0000, 0x007F, "CRC-16/DECT-X"}
	CRC16_DNP         = Params{0x3D65, 0x0000, true, true, 0xFFFF, 0xEA82, "CRC-16/DNP"}
	CRC16_EN_13757    = Params{0x3D65, 0x0000, false, false, 0xFFFF, 0xC2B7, "CRC-16/EN-13757"}
	CRC16_GENIBUS     = Params{0x1021, 0xFFFF, false, false, 0xFFFF, 0xD64E, "CRC-16/GENIBUS"}
	CRC16_MAXIM       = Params{0x8005, 0x0000, true, true, 0xFFFF, 0x44C2, "CRC-16/MAXIM"}
	CRC16_MCRF4XX     = Params{0x1021, 0xFFFF, true, true, 0x0000, 0x6F91, "CRC-16/MCRF4XX"}
	CRC16_RIELLO      = Params{0x1021, 0xB2AA, true, true, 0x0000, 0x63D0, "CRC-16/RIELLO"}
	CRC16_T10_DIF     = Params{0x8BB7, 0x0000, false, false, 0x0000, 0xD0DB, "CRC-16/T10-DIF"}
	CRC16_TELEDISK    = Params{0xA097, 0x0000, false, false, 0x0000, 0x0FB3, "CRC-16/TELEDISK"}
	CRC16_TMS37157    = Params{0x1021, 0x89EC, true, true, 0x0000, 0x26B1, "CRC-16/TMS37157"}
	CRC16_USB         = Params{0x8005, 0xFFFF, true, true, 0xFFFF, 0xB4C8, "CRC-16/USB"}
	CRC16_CRC_A       = Params{0x1021, 0xC6C6, true, true, 0x0000, 0xBF05, "CRC-16/CRC-A"}
	CRC16_KERMIT      = Params{0x1021, 0x0000, true, true, 0x0000, 0x2189, "CRC-16/KERMIT"}
	CRC16_MODBUS      = Params{0x8005, 0xFFFF, true, true, 0x0000, 0x4B37, "CRC-16/MODBUS"}
	CRC16_X_25        = Params{0x1021, 0xFFFF, true, true, 0xFFFF, 0x906E, "CRC-16/X-25"}
	CRC16_XMODEM      = Params{0x1021, 0x0000, false, false, 0x0000, 0x31C3, "CRC-16/XMODEM"}
)

// Table is a 256-word table representing polinomial and algorithm settings for efficient processing.
type Table struct {
	params Params
	data   [256]uint16
}

// MakeTable returns the Table constructed from the specified algorithm.
func MakeTable(params Params) *Table {
	table := new(Table)
	table.params = params
	for n := 0; n < 256; n++ {
		crc := uint16(n) << 8
		for i := 0; i < 8; i++ {
			bit := (crc & 0x8000) != 0
			crc <<= 1
			if bit {
				crc ^= params.Poly
			}
		}
		table.data[n] = crc
	}
	return table
}

// Init returns the initial value for CRC register corresponding to the specified algorithm.
func Init(table *Table) uint16 {
	return table.params.Init
}

// Update returns the result of adding the bytes in data to the crc.
func Update(crc uint16, data []byte, table *Table) uint16 {
	for _, d := range data {
		if table.params.RefIn {
			d = utils.ReverseByte(d)
		}
		crc = crc<<8 ^ table.data[byte(crc>>8)^d]
	}
	return crc
}

// Complete returns the result of CRC calculation and post-calculation processing of the crc.
func Complete(crc uint16, table *Table) uint16 {
	if table.params.RefOut {
		return utils.ReverseUint16(crc) ^ table.params.XorOut
	}
	return crc ^ table.params.XorOut
}

// Checksum returns CRC checksum of data usign scpecified algorithm represented by the Table.
func Checksum(data []byte, table *Table) uint16 {
	crc := Init(table)
	crc = Update(crc, data, table)
	return Complete(crc, table)
}
<commit_msg>Use math/bits for bit operations and remove utils dependency.<commit_after>// Package crc16 implements the 16-bit cyclic redundancy check, or CRC-16, checksum.
//
// It provides parameters for the majority of well-known CRC-16 algorithms.
package crc16

import "math/bits"

// Params represents parameters of CRC-16 algorithms.
// More information about algorithms parametrization and parameter descriptions
// can be found here - http://www.zlib.net/crc_v3.txt
type Params struct {
	Poly   uint16
	Init   uint16
	RefIn  bool
	RefOut bool
	XorOut uint16
	Check  uint16
	Name   string
}

// Predefined CRC-16 algorithms.
// List of algorithms with their parameters borrowed from here -  http://reveng.sourceforge.net/crc-catalogue/16.htm
//
// The variables can be used to create Table for the selected algorithm.
var (
	CRC16_ARC         = Params{0x8005, 0x0000, true, true, 0x0000, 0xBB3D, "CRC-16/ARC"}
	CRC16_AUG_CCITT   = Params{0x1021, 0x1D0F, false, false, 0x0000, 0xE5CC, "CRC-16/AUG-CCITT"}
	CRC16_BUYPASS     = Params{0x8005, 0x0000, false, false, 0x0000, 0xFEE8, "CRC-16/BUYPASS"}
	CRC16_CCITT_FALSE = Params{0x1021, 0xFFFF, false, false, 0x0000, 0x29B1, "CRC-16/CCITT-FALSE"}
	CRC16_CDMA2000    = Params{0xC867, 0xFFFF, false, false, 0x0000, 0x4C06, "CRC-16/CDMA2000"}
	CRC16_DDS_110     = Params{0x8005, 0x800D, false, false, 0x0000, 0x9ECF, "CRC-16/DDS-110"}
	CRC16_DECT_R      = Params{0x0589, 0x0000, false, false, 0x0001, 0x007E, "CRC-16/DECT-R"}
	CRC16_DECT_X      = Params{0x0589, 0x0000, false, false, 0x0000, 0x007F, "CRC-16/DECT-X"}
	CRC16_DNP         = Params{0x3D65, 0x0000, true, true, 0xFFFF, 0xEA82, "CRC-16/DNP"}
	CRC16_EN_13757    = Params{0x3D65, 0x0000, false, false, 0xFFFF, 0xC2B7, "CRC-16/EN-13757"}
	CRC16_GENIBUS     = Params{0x1021, 0xFFFF, false, false, 0xFFFF, 0xD64E, "CRC-16/GENIBUS"}
	CRC16_MAXIM       = Params{0x8005, 0x0000, true, true, 0xFFFF, 0x44C2, "CRC-16/MAXIM"}
	CRC16_MCRF4XX     = Params{0x1021, 0xFFFF, true, true, 0x0000, 0x6F91, "CRC-16/MCRF4XX"}
	CRC16_RIELLO      = Params{0x1021, 0xB2AA, true, true, 0x0000, 0x63D0, "CRC-16/RIELLO"}
	CRC16_T10_DIF     = Params{0x8BB7, 0x0000, false, false, 0x0000, 0xD0DB, "CRC-16/T10-DIF"}
	CRC16_TELEDISK    = Params{0xA097, 0x0000, false, false, 0x0000, 0x0FB3, "CRC-16/TELEDISK"}
	CRC16_TMS37157    = Params{0x1021, 0x89EC, true, true, 0x0000, 0x26B1, "CRC-16/TMS37157"}
	CRC16_USB         = Params{0x8005, 0xFFFF, true, true, 0xFFFF, 0xB4C8, "CRC-16/USB"}
	CRC16_CRC_A       = Params{0x1021, 0xC6C6, true, true, 0x0000, 0xBF05, "CRC-16/CRC-A"}
	CRC16_KERMIT      = Params{0x1021, 0x0000, true, true, 0x0000, 0x2189, "CRC-16/KERMIT"}
	CRC16_MODBUS      = Params{0x8005, 0xFFFF, true, true, 0x0000, 0x4B37, "CRC-16/MODBUS"}
	CRC16_X_25        = Params{0x1021, 0xFFFF, true, true, 0xFFFF, 0x906E, "CRC-16/X-25"}
	CRC16_XMODEM      = Params{0x1021, 0x0000, false, false, 0x0000, 0x31C3, "CRC-16/XMODEM"}
)

// Table is a 256-word table representing polinomial and algorithm settings for efficient processing.
type Table struct {
	params Params
	data   [256]uint16
}

// MakeTable returns the Table constructed from the specified algorithm.
func MakeTable(params Params) *Table {
	table := new(Table)
	table.params = params
	for n := 0; n < 256; n++ {
		crc := uint16(n) << 8
		for i := 0; i < 8; i++ {
			bit := (crc & 0x8000) != 0
			crc <<= 1
			if bit {
				crc ^= params.Poly
			}
		}
		table.data[n] = crc
	}
	return table
}

// Init returns the initial value for CRC register corresponding to the specified algorithm.
func Init(table *Table) uint16 {
	return table.params.Init
}

// Update returns the result of adding the bytes in data to the crc.
func Update(crc uint16, data []byte, table *Table) uint16 {
	for _, d := range data {
		if table.params.RefIn {
			d = bits.Reverse8(d)
		}
		crc = crc<<8 ^ table.data[byte(crc>>8)^d]
	}
	return crc
}

// Complete returns the result of CRC calculation and post-calculation processing of the crc.
func Complete(crc uint16, table *Table) uint16 {
	if table.params.RefOut {
		return bits.Reverse16(crc) ^ table.params.XorOut
	}
	return crc ^ table.params.XorOut
}

// Checksum returns CRC checksum of data usign scpecified algorithm represented by the Table.
func Checksum(data []byte, table *Table) uint16 {
	crc := Init(table)
	crc = Update(crc, data, table)
	return Complete(crc, table)
}
<|endoftext|>


<commit_before>// Copyright 2015 Jimmy Zelinskie. All rights reserved.
// Use of this source code is governed by the BSD 2-Clause license,
// which can be found in the LICENSE file.

// Package dango defines a middleware abstraction for the composition of
// BitTorrent tracker functionality.
package dango

import (
	"net"
	"time"

	"golang.org/x/net/context"
)

// Peer represents one instance of a BitTorrent client currently participating
// in a Swarm.
type Peer struct {
	ID   string
	IP   net.IP
	Port uint32
	Left uint64
}

// PeerIterator represents a stream of Peers.
type PeerIterator interface {
	Next() bool
	Peer() Peer
}

// Event represents an event done by a BitTorrent client.
type Event string

const (
	// Started is the event sent by a BitTorrent client when it joins a Swarm.
	Started Event = "started"

	// Stopped is the event sent by a BitTorrent client when it leaves a Swarm.
	Stopped Event = "stopped"

	// Completed is the event sent by a BitTorrent client when it finishes
	// downloading all of the required chunks.
	Completed Event = "completed"
)

// Infohash is the hash of a set of files that are to be downloaded by a client
// participating in a Swarm.
type Infohash string

// Swarm represents the metadata of collection of Peers sharing a torrent.
type Swarm struct{ Complete, Incomplete, Downloaded uint64 }

// AnnounceIntervals represents the intervals at which a BitTorrent client is
// intended to announce to a Tracker.
type AnnounceIntervals struct{ AnnounceInterval, MinAnnounceInterval time.Duration }

// AnnounceRequest represents the transport-agnostic Announce request sent
// from a BitTorrent client to a Tracker.
type AnnounceRequest struct {
	Peer       Peer
	Infohash   Infohash
	Event      Event
	URL        string
	Downloaded uint64
	Uploaded   uint64
	Compact    bool
	NumWant    uint16
}

// AnnounceResponse represents the transport-agnostic Announce response sent
// from a Tracker to a BitTorrent client.
type AnnounceResponse struct {
	IPv4Peers, IPv6Peers PeerIterator
	AnnounceIntervals
	Swarm
	Compact bool
}

// AnnounceResponseWriter is used by an AnnounceHandler to construct a response.
type AnnounceResponseWriter interface {
	WriteAnnounceResponse(*AnnounceResponse)
	WriteError(error)
}

// AnnounceHandler is a function that operates on an Announce before a response
// has been written.
type AnnounceHandler func(context.Context, AnnounceResponseWriter, *AnnounceRequest) (context.Context, error)

// AnnounceMiddleware enables the extension of an AnnounceHandler via a closure.
type AnnounceMiddleware func(AnnounceHandler) AnnounceHandler

// ScrapeRequest represents a transport-agnostic Scrape request sent from a
// BitTorrent client to a tracker.
type ScrapeRequest []Infohash

// ScrapeResponse represents the transport-agnostic Scrape response sent from
// a tracker to a BitTorrent client.
type ScrapeResponse map[Infohash]Swarm

// ScrapeResponseWriter is used by a ScrapeHandler to construct a response.
type ScrapeResponseWriter interface {
	WriteScrapeResponse(*ScrapeResponse)
	WriteError(error)
}

// ScrapeHandler is a function that operates on a Scrape before a response
// has been written.
type ScrapeHandler func(context.Context, ScrapeResponseWriter, *ScrapeRequest) (context.Context, error)

// ScrapeMiddleware enables the extension of an ScrapeHandler via a closure.
type ScrapeMiddleware func(ScrapeHandler) ScrapeHandler

// AnnounceChain represents a composition of AnnounceMiddleware.
type AnnounceChain struct{ mw []AnnounceMiddleware }

// Append creates a new AnnounceChain by appending the provided middleware to
// the current chain.
func (c AnnounceChain) Append(mw ...AnnounceMiddleware) (nc AnnounceChain) {
	nc.mw = make([]AnnounceMiddleware, len(c.mw)+len(mw))
	copy(nc.mw[:len(c.mw)], c.mw)
	copy(nc.mw[len(c.mw):], mw)
	return
}

// Finalize creates an AnnounceHandler that is the composition of the chain.
func (c AnnounceChain) Finalize(final AnnounceHandler) AnnounceHandler {
	for i := len(c.mw) - 1; i >= 0; i-- {
		final = c.mw[i](final)
	}
	return final
}

// ScrapeChain represents a composition of ScrapeMiddleware.
type ScrapeChain struct{ mw []ScrapeMiddleware }

// Append creates a new ScrapeChain by appending the provided middleware to the
// current chain.
func (c ScrapeChain) Append(mw ...ScrapeMiddleware) (nc ScrapeChain) {
	nc.mw = make([]ScrapeMiddleware, len(c.mw)+len(mw))
	copy(nc.mw[:len(c.mw)], c.mw)
	copy(nc.mw[len(c.mw):], mw)
	return
}

// Finalize creates an ScrapeHandler that is the composition of the chain.
func (c ScrapeChain) Finalize(final ScrapeHandler) ScrapeHandler {
	for i := len(c.mw) - 1; i >= 0; i-- {
		final = c.mw[i](final)
	}
	return final
}

// ServeAnnounce creates an empty context and calls itself.
func (h AnnounceHandler) ServeAnnounce(w AnnounceResponseWriter, r *AnnounceRequest) {
	ctx := context.Background()
	h(ctx, w, r)
}

// ServeScrape creates an empty context and calls itself.
func (h ScrapeHandler) ServeScrape(w ScrapeResponseWriter, r *ScrapeRequest) {
	ctx := context.Background()
	h(ctx, w, r)
}

// NewTracker returns a Tracker given an AnnounceHandler and a ScrapeHandler.
func NewTracker(ah AnnounceHandler, sh ScrapeHandler) Tracker {
	return struct {
		AnnounceHandler
		ScrapeHandler
	}{
		AnnounceHandler: ah,
		ScrapeHandler:   sh,
	}
}

// Tracker represents a BitTorrent tracker.
type Tracker interface {
	ServeAnnounce(AnnounceResponseWriter, *AnnounceRequest)
	ServeScrape(ScrapeResponseWriter, *ScrapeRequest)
}
<commit_msg>move Left from Peer to AnnounceRequest<commit_after>// Copyright 2015 Jimmy Zelinskie. All rights reserved.
// Use of this source code is governed by the BSD 2-Clause license,
// which can be found in the LICENSE file.

// Package dango defines a middleware abstraction for the composition of
// BitTorrent tracker functionality.
package dango

import (
	"net"
	"time"

	"golang.org/x/net/context"
)

// Peer represents one instance of a BitTorrent client currently participating
// in a Swarm.
type Peer struct {
	ID   string
	IP   net.IP
	Port uint32
}

// PeerIterator represents a stream of Peers.
type PeerIterator interface {
	Next() bool
	Peer() Peer
}

// Event represents an event done by a BitTorrent client.
type Event string

const (
	// Started is the event sent by a BitTorrent client when it joins a Swarm.
	Started Event = "started"

	// Stopped is the event sent by a BitTorrent client when it leaves a Swarm.
	Stopped Event = "stopped"

	// Completed is the event sent by a BitTorrent client when it finishes
	// downloading all of the required chunks.
	Completed Event = "completed"
)

// Infohash is the hash of a set of files that are to be downloaded by a client
// participating in a Swarm.
type Infohash string

// Swarm represents the metadata of collection of Peers sharing a torrent.
type Swarm struct{ Complete, Incomplete, Downloaded uint64 }

// AnnounceIntervals represents the intervals at which a BitTorrent client is
// intended to announce to a Tracker.
type AnnounceIntervals struct{ AnnounceInterval, MinAnnounceInterval time.Duration }

// AnnounceRequest represents the transport-agnostic Announce request sent
// from a BitTorrent client to a Tracker.
type AnnounceRequest struct {
	Peer       Peer
	Infohash   Infohash
	Event      Event
	URL        string
	Downloaded uint64
	Uploaded   uint64
	Left       uint64
	Compact    bool
	NumWant    uint16
}

// AnnounceResponse represents the transport-agnostic Announce response sent
// from a Tracker to a BitTorrent client.
type AnnounceResponse struct {
	IPv4Peers, IPv6Peers PeerIterator
	AnnounceIntervals
	Swarm
	Compact bool
}

// AnnounceResponseWriter is used by an AnnounceHandler to construct a response.
type AnnounceResponseWriter interface {
	WriteAnnounceResponse(*AnnounceResponse)
	WriteError(error)
}

// AnnounceHandler is a function that operates on an Announce before a response
// has been written.
type AnnounceHandler func(context.Context, AnnounceResponseWriter, *AnnounceRequest) (context.Context, error)

// AnnounceMiddleware enables the extension of an AnnounceHandler via a closure.
type AnnounceMiddleware func(AnnounceHandler) AnnounceHandler

// ScrapeRequest represents a transport-agnostic Scrape request sent from a
// BitTorrent client to a tracker.
type ScrapeRequest []Infohash

// ScrapeResponse represents the transport-agnostic Scrape response sent from
// a tracker to a BitTorrent client.
type ScrapeResponse map[Infohash]Swarm

// ScrapeResponseWriter is used by a ScrapeHandler to construct a response.
type ScrapeResponseWriter interface {
	WriteScrapeResponse(*ScrapeResponse)
	WriteError(error)
}

// ScrapeHandler is a function that operates on a Scrape before a response
// has been written.
type ScrapeHandler func(context.Context, ScrapeResponseWriter, *ScrapeRequest) (context.Context, error)

// ScrapeMiddleware enables the extension of an ScrapeHandler via a closure.
type ScrapeMiddleware func(ScrapeHandler) ScrapeHandler

// AnnounceChain represents a composition of AnnounceMiddleware.
type AnnounceChain struct{ mw []AnnounceMiddleware }

// Append creates a new AnnounceChain by appending the provided middleware to
// the current chain.
func (c AnnounceChain) Append(mw ...AnnounceMiddleware) (nc AnnounceChain) {
	nc.mw = make([]AnnounceMiddleware, len(c.mw)+len(mw))
	copy(nc.mw[:len(c.mw)], c.mw)
	copy(nc.mw[len(c.mw):], mw)
	return
}

// Finalize creates an AnnounceHandler that is the composition of the chain.
func (c AnnounceChain) Finalize(final AnnounceHandler) AnnounceHandler {
	for i := len(c.mw) - 1; i >= 0; i-- {
		final = c.mw[i](final)
	}
	return final
}

// ScrapeChain represents a composition of ScrapeMiddleware.
type ScrapeChain struct{ mw []ScrapeMiddleware }

// Append creates a new ScrapeChain by appending the provided middleware to the
// current chain.
func (c ScrapeChain) Append(mw ...ScrapeMiddleware) (nc ScrapeChain) {
	nc.mw = make([]ScrapeMiddleware, len(c.mw)+len(mw))
	copy(nc.mw[:len(c.mw)], c.mw)
	copy(nc.mw[len(c.mw):], mw)
	return
}

// Finalize creates an ScrapeHandler that is the composition of the chain.
func (c ScrapeChain) Finalize(final ScrapeHandler) ScrapeHandler {
	for i := len(c.mw) - 1; i >= 0; i-- {
		final = c.mw[i](final)
	}
	return final
}

// ServeAnnounce creates an empty context and calls itself.
func (h AnnounceHandler) ServeAnnounce(w AnnounceResponseWriter, r *AnnounceRequest) {
	ctx := context.Background()
	h(ctx, w, r)
}

// ServeScrape creates an empty context and calls itself.
func (h ScrapeHandler) ServeScrape(w ScrapeResponseWriter, r *ScrapeRequest) {
	ctx := context.Background()
	h(ctx, w, r)
}

// NewTracker returns a Tracker given an AnnounceHandler and a ScrapeHandler.
func NewTracker(ah AnnounceHandler, sh ScrapeHandler) Tracker {
	return struct {
		AnnounceHandler
		ScrapeHandler
	}{
		AnnounceHandler: ah,
		ScrapeHandler:   sh,
	}
}

// Tracker represents a BitTorrent tracker.
type Tracker interface {
	ServeAnnounce(AnnounceResponseWriter, *AnnounceRequest)
	ServeScrape(ScrapeResponseWriter, *ScrapeRequest)
}
<|endoftext|>


<commit_before>package workerpool

import (
	"sync"
	"testing"
	"time"
)

const max = 20

func TestMaxWorkers(t *testing.T) {
	t.Parallel()

	wp := New(max)
	defer wp.Stop()

	started := make(chan struct{}, max)
	sync := make(chan struct{})

	// Start workers, and have them all wait on a channel before completing.
	for i := 0; i < max; i++ {
		wp.Submit(func() {
			started <- struct{}{}
			<-sync
		})
	}

	// Wait for all enqueued tasks to be dispatched to workers.
	timeout := time.After(5 * time.Second)
	for startCount := 0; startCount < max; {
		select {
		case <-started:
			startCount++
		case <-timeout:
			t.Fatal("timed out waiting for workers to start")
		}
	}

	// Release workers.
	close(sync)
}

func TestReuseWorkers(t *testing.T) {
	t.Parallel()

	wp := New(5)
	defer wp.Stop()

	sync := make(chan struct{})

	// Cause worker to be created, and available for reuse before next task.
	for i := 0; i < 10; i++ {
		wp.Submit(func() { <-sync })
		sync <- struct{}{}
		time.Sleep(100 * time.Millisecond)
	}

	// If the same worker was always reused, then only one worker would have
	// been created and there should only be one ready.
	if countReady(wp) > 1 {
		t.Fatal("Worker not reused")
	}
}

func TestWorkerTimeout(t *testing.T) {
	t.Parallel()

	wp := New(max)
	defer wp.Stop()

	sync := make(chan struct{})
	started := make(chan struct{}, max)
	// Cause workers to be created.  Workers wait on channel, keeping them busy
	// and causing the worker pool to create more.
	for i := 0; i < max; i++ {
		wp.Submit(func() {
			started <- struct{}{}
			<-sync
		})
	}

	// Wait for tasks to start.
	for i := 0; i < max; i++ {
		<-started
	}

	if anyReady(wp) {
		t.Fatal("number of ready workers should ber zero")
	}
	// Release workers.
	close(sync)

	if countReady(wp) != max {
		t.Fatal("Expected", max, "ready workers")
	}

	// Check that a worker timed out.
	time.Sleep((idleTimeoutSec + 1) * time.Second)
	if countReady(wp) != max-1 {
		t.Fatal("First worker did not timeout")
	}

	// Check that another worker timed out.
	time.Sleep((idleTimeoutSec + 1) * time.Second)
	if countReady(wp) != max-2 {
		t.Fatal("Second worker did not timeout")
	}
}

func TestStop(t *testing.T) {
	t.Parallel()

	wp := New(max)
	defer wp.Stop()

	started := make(chan struct{}, max)
	sync := make(chan struct{})

	// Start workers, and have them all wait on a channel before completing.
	for i := 0; i < max; i++ {
		wp.Submit(func() {
			started <- struct{}{}
			<-sync
		})
	}

	// Wait for all enqueued tasks to be dispatched to workers.
	timeout := time.After(5 * time.Second)
	for startCount := 0; startCount < max; {
		select {
		case <-started:
			startCount++
		case <-timeout:
			t.Fatal("timed out waiting for workers to start")
		}
	}

	// Release workers.
	close(sync)

	wp.Stop()
	if anyReady(wp) {
		t.Fatal("should have zero workers after stop")
	}
}

func TestSubmitWait(t *testing.T) {
	wp := New(1)
	defer wp.Stop()

	done1 := make(chan struct{})
	wp.Submit(func() {
		time.Sleep(100 * time.Millisecond)
		close(done1)
	})
	select {
	case <-done1:
		t.Fatal("Submit did not return immediately")
	default:
	}

	done2 := make(chan struct{})
	wp.SubmitWait(func() {
		time.Sleep(100 * time.Millisecond)
		close(done2)
	})
	select {
	case <-done2:
	default:
		t.Fatal("SubmitWait did not wait for function to execute")
	}
}

func anyReady(w *WorkerPool) bool {
	select {
	case wkCh := <-w.readyWorkers:
		w.readyWorkers <- wkCh
		return true
	default:
	}
	return false
}

func countReady(w *WorkerPool) int {
	// Try to pull max workers off of ready queue.
	timeout := time.After(5 * time.Second)
	readyTmp := make(chan chan func(), max)
	var readyCount int
	for i := 0; i < max; i++ {
		select {
		case wkCh := <-w.readyWorkers:
			readyTmp <- wkCh
			readyCount++
		case <-timeout:
			readyCount = i
			i = max
		}
	}

	// Restore ready workers.
	close(readyTmp)
	go func() {
		for r := range readyTmp {
			w.readyWorkers <- r
		}
	}()
	return readyCount
}

/*

Run benchmarking with: go test -bench '.'

*/

func BenchmarkEnqueue(b *testing.B) {
	wp := New(1)
	defer wp.Stop()
	releaseChan := make(chan struct{})

	b.ResetTimer()

	// Start workers, and have them all wait on a channel before completing.
	for i := 0; i < b.N; i++ {
		wp.Submit(func() { <-releaseChan })
	}
	close(releaseChan)
}

func BenchmarkExecute1Worker(b *testing.B) {
	wp := New(1)
	defer wp.Stop()
	var allDone sync.WaitGroup
	allDone.Add(b.N)

	b.ResetTimer()

	// Start workers, and have them all wait on a channel before completing.
	for i := 0; i < b.N; i++ {
		wp.Submit(func() {
			time.Sleep(time.Millisecond)
			allDone.Done()
		})
	}
	allDone.Wait()
}

func BenchmarkExecute2Worker(b *testing.B) {
	wp := New(2)
	defer wp.Stop()
	var allDone sync.WaitGroup
	allDone.Add(b.N)

	b.ResetTimer()

	// Start workers, and have them all wait on a channel before completing.
	for i := 0; i < b.N; i++ {
		wp.Submit(func() {
			time.Sleep(time.Millisecond)
			allDone.Done()
		})
	}
	allDone.Wait()
}

func BenchmarkExecute4Workers(b *testing.B) {
	wp := New(4)
	defer wp.Stop()
	var allDone sync.WaitGroup
	allDone.Add(b.N)

	b.ResetTimer()

	// Start workers, and have them all wait on a channel before completing.
	for i := 0; i < b.N; i++ {
		wp.Submit(func() {
			time.Sleep(time.Millisecond)
			allDone.Done()
		})
	}
	allDone.Wait()
}

func BenchmarkExecute16Workers(b *testing.B) {
	wp := New(16)
	defer wp.Stop()
	var allDone sync.WaitGroup
	allDone.Add(b.N)

	b.ResetTimer()

	// Start workers, and have them all wait on a channel before completing.
	for i := 0; i < b.N; i++ {
		wp.Submit(func() {
			time.Sleep(time.Millisecond)
			allDone.Done()
		})
	}
	allDone.Wait()
}

func BenchmarkExecute64Workers(b *testing.B) {
	wp := New(64)
	defer wp.Stop()
	var allDone sync.WaitGroup
	allDone.Add(b.N)

	b.ResetTimer()

	// Start workers, and have them all wait on a channel before completing.
	for i := 0; i < b.N; i++ {
		wp.Submit(func() {
			time.Sleep(time.Millisecond)
			allDone.Done()
		})
	}
	allDone.Wait()
}
<commit_msg>Complete test coverage<commit_after>package workerpool

import (
	"sync"
	"testing"
	"time"
)

const max = 20

func TestMaxWorkers(t *testing.T) {
	t.Parallel()

	wp := New(0)
	if wp.maxWorkers != 1 {
		t.Fatal("should have created one worker")
	}

	wp = New(max)
	defer wp.Stop()

	started := make(chan struct{}, max)
	sync := make(chan struct{})

	// Start workers, and have them all wait on a channel before completing.
	for i := 0; i < max; i++ {
		wp.Submit(func() {
			started <- struct{}{}
			<-sync
		})
	}

	// Wait for all enqueued tasks to be dispatched to workers.
	timeout := time.After(5 * time.Second)
	for startCount := 0; startCount < max; {
		select {
		case <-started:
			startCount++
		case <-timeout:
			t.Fatal("timed out waiting for workers to start")
		}
	}

	// Release workers.
	close(sync)
}

func TestReuseWorkers(t *testing.T) {
	t.Parallel()

	wp := New(5)
	defer wp.Stop()

	sync := make(chan struct{})

	// Cause worker to be created, and available for reuse before next task.
	for i := 0; i < 10; i++ {
		wp.Submit(func() { <-sync })
		sync <- struct{}{}
		time.Sleep(100 * time.Millisecond)
	}

	// If the same worker was always reused, then only one worker would have
	// been created and there should only be one ready.
	if countReady(wp) > 1 {
		t.Fatal("Worker not reused")
	}
}

func TestWorkerTimeout(t *testing.T) {
	t.Parallel()

	wp := New(max)
	defer wp.Stop()

	sync := make(chan struct{})
	started := make(chan struct{}, max)
	// Cause workers to be created.  Workers wait on channel, keeping them busy
	// and causing the worker pool to create more.
	for i := 0; i < max; i++ {
		wp.Submit(func() {
			started <- struct{}{}
			<-sync
		})
	}

	// Wait for tasks to start.
	for i := 0; i < max; i++ {
		<-started
	}

	if anyReady(wp) {
		t.Fatal("number of ready workers should ber zero")
	}
	// Release workers.
	close(sync)

	if countReady(wp) != max {
		t.Fatal("Expected", max, "ready workers")
	}

	// Check that a worker timed out.
	time.Sleep((idleTimeoutSec + 1) * time.Second)
	if countReady(wp) != max-1 {
		t.Fatal("First worker did not timeout")
	}

	// Check that another worker timed out.
	time.Sleep((idleTimeoutSec + 1) * time.Second)
	if countReady(wp) != max-2 {
		t.Fatal("Second worker did not timeout")
	}
}

func TestStop(t *testing.T) {
	t.Parallel()

	wp := New(max)
	defer wp.Stop()

	started := make(chan struct{}, max)
	sync := make(chan struct{})

	// Start workers, and have them all wait on a channel before completing.
	for i := 0; i < max; i++ {
		wp.Submit(func() {
			started <- struct{}{}
			<-sync
		})
	}

	// Wait for all enqueued tasks to be dispatched to workers.
	timeout := time.After(5 * time.Second)
	for startCount := 0; startCount < max; {
		select {
		case <-started:
			startCount++
		case <-timeout:
			t.Fatal("timed out waiting for workers to start")
		}
	}

	// Release workers.
	close(sync)

	wp.Stop()
	if anyReady(wp) {
		t.Fatal("should have zero workers after stop")
	}
}

func TestSubmitWait(t *testing.T) {
	wp := New(1)
	defer wp.Stop()

	// Check that these are noop.
	wp.Submit(nil)
	wp.SubmitWait(nil)

	done1 := make(chan struct{})
	wp.Submit(func() {
		time.Sleep(100 * time.Millisecond)
		close(done1)
	})
	select {
	case <-done1:
		t.Fatal("Submit did not return immediately")
	default:
	}

	done2 := make(chan struct{})
	wp.SubmitWait(func() {
		time.Sleep(100 * time.Millisecond)
		close(done2)
	})
	select {
	case <-done2:
	default:
		t.Fatal("SubmitWait did not wait for function to execute")
	}
}

func anyReady(w *WorkerPool) bool {
	select {
	case wkCh := <-w.readyWorkers:
		w.readyWorkers <- wkCh
		return true
	default:
	}
	return false
}

func countReady(w *WorkerPool) int {
	// Try to pull max workers off of ready queue.
	timeout := time.After(5 * time.Second)
	readyTmp := make(chan chan func(), max)
	var readyCount int
	for i := 0; i < max; i++ {
		select {
		case wkCh := <-w.readyWorkers:
			readyTmp <- wkCh
			readyCount++
		case <-timeout:
			readyCount = i
			i = max
		}
	}

	// Restore ready workers.
	close(readyTmp)
	go func() {
		for r := range readyTmp {
			w.readyWorkers <- r
		}
	}()
	return readyCount
}

/*

Run benchmarking with: go test -bench '.'

*/

func BenchmarkEnqueue(b *testing.B) {
	wp := New(1)
	defer wp.Stop()
	releaseChan := make(chan struct{})

	b.ResetTimer()

	// Start workers, and have them all wait on a channel before completing.
	for i := 0; i < b.N; i++ {
		wp.Submit(func() { <-releaseChan })
	}
	close(releaseChan)
}

func BenchmarkExecute1Worker(b *testing.B) {
	wp := New(1)
	defer wp.Stop()
	var allDone sync.WaitGroup
	allDone.Add(b.N)

	b.ResetTimer()

	// Start workers, and have them all wait on a channel before completing.
	for i := 0; i < b.N; i++ {
		wp.Submit(func() {
			time.Sleep(time.Millisecond)
			allDone.Done()
		})
	}
	allDone.Wait()
}

func BenchmarkExecute2Worker(b *testing.B) {
	wp := New(2)
	defer wp.Stop()
	var allDone sync.WaitGroup
	allDone.Add(b.N)

	b.ResetTimer()

	// Start workers, and have them all wait on a channel before completing.
	for i := 0; i < b.N; i++ {
		wp.Submit(func() {
			time.Sleep(time.Millisecond)
			allDone.Done()
		})
	}
	allDone.Wait()
}

func BenchmarkExecute4Workers(b *testing.B) {
	wp := New(4)
	defer wp.Stop()
	var allDone sync.WaitGroup
	allDone.Add(b.N)

	b.ResetTimer()

	// Start workers, and have them all wait on a channel before completing.
	for i := 0; i < b.N; i++ {
		wp.Submit(func() {
			time.Sleep(time.Millisecond)
			allDone.Done()
		})
	}
	allDone.Wait()
}

func BenchmarkExecute16Workers(b *testing.B) {
	wp := New(16)
	defer wp.Stop()
	var allDone sync.WaitGroup
	allDone.Add(b.N)

	b.ResetTimer()

	// Start workers, and have them all wait on a channel before completing.
	for i := 0; i < b.N; i++ {
		wp.Submit(func() {
			time.Sleep(time.Millisecond)
			allDone.Done()
		})
	}
	allDone.Wait()
}

func BenchmarkExecute64Workers(b *testing.B) {
	wp := New(64)
	defer wp.Stop()
	var allDone sync.WaitGroup
	allDone.Add(b.N)

	b.ResetTimer()

	// Start workers, and have them all wait on a channel before completing.
	for i := 0; i < b.N; i++ {
		wp.Submit(func() {
			time.Sleep(time.Millisecond)
			allDone.Done()
		})
	}
	allDone.Wait()
}
<|endoftext|>


<commit_before>// The MIT License (MIT)
//
// Copyright (c) 2014 wandoulabs
// Copyright (c) 2014 siddontang
//
// Permission is hereby granted, free of charge, to any person obtaining a copy of
// this software and associated documentation files (the "Software"), to deal in
// the Software without restriction, including without limitation the rights to
// use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
// the Software, and to permit persons to whom the Software is furnished to do so,
// subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.

// 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 server

import (
	"math/rand"
	"net"
	"sync"
	"sync/atomic"
	"time"

	"github.com/juju/errors"
	"github.com/ngaut/log"
	mysql "github.com/pingcap/tidb/mysqldef"
	"github.com/pingcap/tidb/util/arena"
)

var (
	baseConnID uint32 = 10000
)

// Server is the MySQL protocol server
type Server struct {
	cfg               *Config
	driver            IDriver
	listener          net.Listener
	rwlock            *sync.RWMutex
	concurrentLimiter *TokenLimiter
	clients           map[uint32]*clientConn
}

func (s *Server) getToken() *Token {
	return s.concurrentLimiter.Get()
}

func (s *Server) releaseToken(token *Token) {
	s.concurrentLimiter.Put(token)
}

// Generate a random string using ASCII characters but avoid seperator character.
// See: https://github.com/mysql/mysql-server/blob/5.7/mysys_ssl/crypt_genhash_impl.cc#L435
func randomBuf(size int) ([]byte, error) {
	buf := make([]byte, size)
	rand.Seed(time.Now().UTC().UnixNano())
	for i := 0; i < size; i++ {
		buf[i] &= 0x7f
		if buf[i] == 0 || buf[i] == byte('$') {
			buf[i]++
		}
	}
	return buf, nil
}

func (s *Server) newConn(conn net.Conn) (cc *clientConn, err error) {
	log.Info("newConn", conn.RemoteAddr().String())
	cc = &clientConn{
		conn:         conn,
		pkg:          newPacketIO(conn),
		server:       s,
		connectionID: atomic.AddUint32(&baseConnID, 1),
		collation:    mysql.DefaultCollationID,
		charset:      mysql.DefaultCharset,
		alloc:        arena.NewAllocator(32 * 1024),
	}
	cc.salt, err = randomBuf(20)
	return
}

func (s *Server) skipAuth() bool {
	return s.cfg.SkipAuth
}

// NewServer creates a new Server.
func NewServer(cfg *Config, driver IDriver) (*Server, error) {
	s := &Server{
		cfg:               cfg,
		driver:            driver,
		concurrentLimiter: NewTokenLimiter(100),
		rwlock:            &sync.RWMutex{},
		clients:           make(map[uint32]*clientConn),
	}

	var err error
	s.listener, err = net.Listen("tcp", s.cfg.Addr)
	if err != nil {
		return nil, errors.Trace(err)
	}

	log.Infof("Server run MySql Protocol Listen at [%s]", s.cfg.Addr)
	return s, nil
}

// Run runs the server.
func (s *Server) Run() error {
	for {
		conn, err := s.listener.Accept()
		if err != nil {
			log.Errorf("accept error %s", err.Error())
			return errors.Trace(err)
		}

		go s.onConn(conn)
	}
}

// Close closes the server.
func (s *Server) Close() {
	s.rwlock.Lock()
	defer s.rwlock.Unlock()

	if s.listener != nil {
		s.listener.Close()
		s.listener = nil
	}
}

func (s *Server) onConn(c net.Conn) {
	conn, err := s.newConn(c)
	if err != nil {
		log.Errorf("newConn error %s", errors.ErrorStack(err))
		return
	}
	if err := conn.handshake(); err != nil {
		log.Errorf("handshake error %s", errors.ErrorStack(err))
		c.Close()
		return
	}
	defer func() {
		log.Infof("close %s", conn)
	}()

	s.rwlock.Lock()
	s.clients[conn.connectionID] = conn
	s.rwlock.Unlock()

	conn.Run()
}
<commit_msg>tidb-server: Remove useless error<commit_after>// The MIT License (MIT)
//
// Copyright (c) 2014 wandoulabs
// Copyright (c) 2014 siddontang
//
// Permission is hereby granted, free of charge, to any person obtaining a copy of
// this software and associated documentation files (the "Software"), to deal in
// the Software without restriction, including without limitation the rights to
// use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
// the Software, and to permit persons to whom the Software is furnished to do so,
// subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.

// 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 server

import (
	"math/rand"
	"net"
	"sync"
	"sync/atomic"
	"time"

	"github.com/juju/errors"
	"github.com/ngaut/log"
	mysql "github.com/pingcap/tidb/mysqldef"
	"github.com/pingcap/tidb/util/arena"
)

var (
	baseConnID uint32 = 10000
)

// Server is the MySQL protocol server
type Server struct {
	cfg               *Config
	driver            IDriver
	listener          net.Listener
	rwlock            *sync.RWMutex
	concurrentLimiter *TokenLimiter
	clients           map[uint32]*clientConn
}

func (s *Server) getToken() *Token {
	return s.concurrentLimiter.Get()
}

func (s *Server) releaseToken(token *Token) {
	s.concurrentLimiter.Put(token)
}

// Generate a random string using ASCII characters but avoid seperator character.
// See: https://github.com/mysql/mysql-server/blob/5.7/mysys_ssl/crypt_genhash_impl.cc#L435
func randomBuf(size int) []byte {
	buf := make([]byte, size)
	rand.Seed(time.Now().UTC().UnixNano())
	for i := 0; i < size; i++ {
		buf[i] = byte(rand.Intn(127))
		if buf[i] == 0 || buf[i] == byte('$') {
			buf[i]++
		}
	}
	return buf
}

func (s *Server) newConn(conn net.Conn) (cc *clientConn, err error) {
	log.Info("newConn", conn.RemoteAddr().String())
	cc = &clientConn{
		conn:         conn,
		pkg:          newPacketIO(conn),
		server:       s,
		connectionID: atomic.AddUint32(&baseConnID, 1),
		collation:    mysql.DefaultCollationID,
		charset:      mysql.DefaultCharset,
		alloc:        arena.NewAllocator(32 * 1024),
	}
	cc.salt = randomBuf(20)
	return
}

func (s *Server) skipAuth() bool {
	return s.cfg.SkipAuth
}

// NewServer creates a new Server.
func NewServer(cfg *Config, driver IDriver) (*Server, error) {
	s := &Server{
		cfg:               cfg,
		driver:            driver,
		concurrentLimiter: NewTokenLimiter(100),
		rwlock:            &sync.RWMutex{},
		clients:           make(map[uint32]*clientConn),
	}

	var err error
	s.listener, err = net.Listen("tcp", s.cfg.Addr)
	if err != nil {
		return nil, errors.Trace(err)
	}

	log.Infof("Server run MySql Protocol Listen at [%s]", s.cfg.Addr)
	return s, nil
}

// Run runs the server.
func (s *Server) Run() error {
	for {
		conn, err := s.listener.Accept()
		if err != nil {
			log.Errorf("accept error %s", err.Error())
			return errors.Trace(err)
		}

		go s.onConn(conn)
	}
}

// Close closes the server.
func (s *Server) Close() {
	s.rwlock.Lock()
	defer s.rwlock.Unlock()

	if s.listener != nil {
		s.listener.Close()
		s.listener = nil
	}
}

func (s *Server) onConn(c net.Conn) {
	conn, err := s.newConn(c)
	if err != nil {
		log.Errorf("newConn error %s", errors.ErrorStack(err))
		return
	}
	if err := conn.handshake(); err != nil {
		log.Errorf("handshake error %s", errors.ErrorStack(err))
		c.Close()
		return
	}
	defer func() {
		log.Infof("close %s", conn)
	}()

	s.rwlock.Lock()
	s.clients[conn.connectionID] = conn
	s.rwlock.Unlock()

	conn.Run()
}
<|endoftext|>


<commit_before><commit_msg>Remove todo for version info as its not its own metric.<commit_after><|endoftext|>


<commit_before>package secure

import (
	"errors"
	"github.com/Comcast/webpa-common/secure/key"
	"github.com/SermoDigital/jose/jws"
	"github.com/SermoDigital/jose/jwt"
	"time"
)

var (
	ErrorNoProtectedHeader = errors.New("Missing protected header")
	ErrorNoSigningMethod   = errors.New("Signing method (alg) is missing or unrecognized")
)

// Validator describes the behavior of a type which can validate tokens
type Validator interface {
	// Validate asserts that the given token is valid, most often verifying
	// the credentials in the token.  A separate error is returned to indicate
	// any problems during validation, such as the inability to access a network resource.
	// In general, the contract of this method is that a Token passes validation
	// if and only if it returns BOTH true and a nil error.
	Validate(*Token) (bool, error)
}

// ValidatorFunc is a function type that implements Validator
type ValidatorFunc func(*Token) (bool, error)

func (v ValidatorFunc) Validate(token *Token) (bool, error) {
	return v(token)
}

// Validators is an aggregate Validator.  A Validators instance considers a token
// valid if any of its validators considers it valid.  An empty Validators rejects
// all tokens.
type Validators []Validator

func (v Validators) Validate(token *Token) (valid bool, err error) {
	for _, validator := range v {
		if valid, err = validator.Validate(token); valid && err == nil {
			return
		}
	}

	return
}

// ExactMatchValidator simply matches a token's value (exluding the prefix, such as "Basic"),
// to a string.
type ExactMatchValidator string

func (v ExactMatchValidator) Validate(token *Token) (bool, error) {
	return string(v) == token.value, nil
}

// JWSValidator provides validation for JWT tokens encoded as JWS.
type JWSValidator struct {
	DefaultKeyId  string
	Resolver      key.Resolver
	Parser        JWSParser
	JWTValidators []*jwt.Validator
}

func (v JWSValidator) Validate(token *Token) (valid bool, err error) {
	if token.Type() != Bearer {
		return
	}

	parser := v.Parser
	if parser == nil {
		parser = DefaultJWSParser
	}

	jwsToken, err := parser.ParseJWS(token)
	if err != nil {
		return
	}

	protected := jwsToken.Protected()
	if len(protected) == 0 {
		err = ErrorNoProtectedHeader
		return
	}

	alg, _ := protected.Get("alg").(string)
	signingMethod := jws.GetSigningMethod(alg)
	if signingMethod == nil {
		err = ErrorNoSigningMethod
		return
	}

	keyId, _ := protected.Get("kid").(string)
	if len(keyId) == 0 {
		keyId = v.DefaultKeyId
	}

	pair, err := v.Resolver.ResolveKey(keyId)
	if err != nil {
		return
	}

	if len(v.JWTValidators) > 0 {
		// all JWS implementations also implement jwt.JWT
		err = jwsToken.(jwt.JWT).Validate(pair.Public(), signingMethod, v.JWTValidators...)
	} else {
		err = jwsToken.Verify(pair.Public(), signingMethod)
	}

	valid = (err == nil)
	return
}

// JWTValidatorFactory is a configurable factory for *jwt.Validator instances
type JWTValidatorFactory struct {
	Expected  jwt.Claims `json:"expected"`
	ExpLeeway int        `json:"expLeeway"`
	NbfLeeway int        `json:"nbfLeeway"`
}

func (f *JWTValidatorFactory) expLeeway() time.Duration {
	if f.ExpLeeway > 0 {
		return time.Duration(f.ExpLeeway) * time.Second
	}

	return 0
}

func (f *JWTValidatorFactory) nbfLeeway() time.Duration {
	if f.NbfLeeway > 0 {
		return time.Duration(f.NbfLeeway) * time.Second
	}

	return 0
}

// New returns a jwt.Validator using the configuration expected claims (if any)
// and a validator function that checks the exp and nbf claims.
//
// The SermoDigital library doesn't appear to do anything with the EXP and NBF
// members of jwt.Validator, but this Factory Method populates them anyway.
func (f *JWTValidatorFactory) New(custom ...jwt.ValidateFunc) *jwt.Validator {
	expLeeway := f.expLeeway()
	nbfLeeway := f.nbfLeeway()

	var validateFunc jwt.ValidateFunc
	customCount := len(custom)
	if customCount > 0 {
		validateFunc = func(claims jwt.Claims) (err error) {
			err = claims.Validate(time.Now(), expLeeway, nbfLeeway)
			for index := 0; index < customCount && err == nil; index++ {
				err = custom[index](claims)
			}

			return
		}
	} else {
		// if no custom validate functions were passed, use a simpler function
		validateFunc = func(claims jwt.Claims) error {
			return claims.Validate(time.Now(), expLeeway, nbfLeeway)
		}
	}

	return &jwt.Validator{
		Expected: f.Expected,
		EXP:      expLeeway,
		NBF:      nbfLeeway,
		Fn:       validateFunc,
	}
}
<commit_msg>check for against multiple authorization strings (#31)<commit_after>package secure

import (
	"errors"
	"github.com/Comcast/webpa-common/secure/key"
	"github.com/SermoDigital/jose/jws"
	"github.com/SermoDigital/jose/jwt"
	"time"
	"strings"
)

var (
	ErrorNoProtectedHeader = errors.New("Missing protected header")
	ErrorNoSigningMethod   = errors.New("Signing method (alg) is missing or unrecognized")
)

// Validator describes the behavior of a type which can validate tokens
type Validator interface {
	// Validate asserts that the given token is valid, most often verifying
	// the credentials in the token.  A separate error is returned to indicate
	// any problems during validation, such as the inability to access a network resource.
	// In general, the contract of this method is that a Token passes validation
	// if and only if it returns BOTH true and a nil error.
	Validate(*Token) (bool, error)
}

// ValidatorFunc is a function type that implements Validator
type ValidatorFunc func(*Token) (bool, error)

func (v ValidatorFunc) Validate(token *Token) (bool, error) {
	return v(token)
}

// Validators is an aggregate Validator.  A Validators instance considers a token
// valid if any of its validators considers it valid.  An empty Validators rejects
// all tokens.
type Validators []Validator

func (v Validators) Validate(token *Token) (valid bool, err error) {
	for _, validator := range v {
		if valid, err = validator.Validate(token); valid && err == nil {
			return
		}
	}

	return
}

// ExactMatchValidator simply matches a token's value (exluding the prefix, such as "Basic"),
// to a string.
type ExactMatchValidator string

func (v ExactMatchValidator) Validate(token *Token) (bool, error) {
	for _, value := range strings.Split(string(v), ",") {
		if value == token.value {
			return true, nil
		}
	}

	return false, nil
}

// JWSValidator provides validation for JWT tokens encoded as JWS.
type JWSValidator struct {
	DefaultKeyId  string
	Resolver      key.Resolver
	Parser        JWSParser
	JWTValidators []*jwt.Validator
}

func (v JWSValidator) Validate(token *Token) (valid bool, err error) {
	if token.Type() != Bearer {
		return
	}

	parser := v.Parser
	if parser == nil {
		parser = DefaultJWSParser
	}

	jwsToken, err := parser.ParseJWS(token)
	if err != nil {
		return
	}

	protected := jwsToken.Protected()
	if len(protected) == 0 {
		err = ErrorNoProtectedHeader
		return
	}

	alg, _ := protected.Get("alg").(string)
	signingMethod := jws.GetSigningMethod(alg)
	if signingMethod == nil {
		err = ErrorNoSigningMethod
		return
	}

	keyId, _ := protected.Get("kid").(string)
	if len(keyId) == 0 {
		keyId = v.DefaultKeyId
	}

	pair, err := v.Resolver.ResolveKey(keyId)
	if err != nil {
		return
	}

	if len(v.JWTValidators) > 0 {
		// all JWS implementations also implement jwt.JWT
		err = jwsToken.(jwt.JWT).Validate(pair.Public(), signingMethod, v.JWTValidators...)
	} else {
		err = jwsToken.Verify(pair.Public(), signingMethod)
	}

	valid = (err == nil)
	return
}

// JWTValidatorFactory is a configurable factory for *jwt.Validator instances
type JWTValidatorFactory struct {
	Expected  jwt.Claims `json:"expected"`
	ExpLeeway int        `json:"expLeeway"`
	NbfLeeway int        `json:"nbfLeeway"`
}

func (f *JWTValidatorFactory) expLeeway() time.Duration {
	if f.ExpLeeway > 0 {
		return time.Duration(f.ExpLeeway) * time.Second
	}

	return 0
}

func (f *JWTValidatorFactory) nbfLeeway() time.Duration {
	if f.NbfLeeway > 0 {
		return time.Duration(f.NbfLeeway) * time.Second
	}

	return 0
}

// New returns a jwt.Validator using the configuration expected claims (if any)
// and a validator function that checks the exp and nbf claims.
//
// The SermoDigital library doesn't appear to do anything with the EXP and NBF
// members of jwt.Validator, but this Factory Method populates them anyway.
func (f *JWTValidatorFactory) New(custom ...jwt.ValidateFunc) *jwt.Validator {
	expLeeway := f.expLeeway()
	nbfLeeway := f.nbfLeeway()

	var validateFunc jwt.ValidateFunc
	customCount := len(custom)
	if customCount > 0 {
		validateFunc = func(claims jwt.Claims) (err error) {
			err = claims.Validate(time.Now(), expLeeway, nbfLeeway)
			for index := 0; index < customCount && err == nil; index++ {
				err = custom[index](claims)
			}

			return
		}
	} else {
		// if no custom validate functions were passed, use a simpler function
		validateFunc = func(claims jwt.Claims) error {
			return claims.Validate(time.Now(), expLeeway, nbfLeeway)
		}
	}

	return &jwt.Validator{
		Expected: f.Expected,
		EXP:      expLeeway,
		NBF:      nbfLeeway,
		Fn:       validateFunc,
	}
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"
	"log"
)

type mongoEngine struct {
	session *mgo.Session
	dbName  string
}

func (eng mongoEngine) Close() {
	// TODO
}

// NewMongoEngine returns an Engine based on a mongodb batabase backend
func NewMongoEngine(dbName string, hostPorts string) Engine {
	log.Printf("connecting to mongodb '%s'\n", hostPorts)
	s, err := mgo.Dial(hostPorts)
	if err != nil {
		panic(err)
	}
	s.SetMode(mgo.Monotonic, true)
	eng := &mongoEngine{
		session: s,
		dbName:  dbName,
	}

	eng.EnsureIndexes("organisations") //TODOL parameterise this

	return eng
}

func (eng *mongoEngine) EnsureIndexes(collection string) {
	c := eng.session.DB(eng.dbName).C(collection)

	eng.session.ResetIndexCache()

	// create collection if it's not there
	c.Create(&mgo.CollectionInfo{})

	err := c.EnsureIndex(mgo.Index{
		Key:        []string{"uuid"}, //TODO: parameterise this
		Unique:     true,
		DropDups:   true,
		Background: false,
		Sparse:     false,
	})
	if err != nil {
		panic(err)
	}

}

func (eng *mongoEngine) Drop(collection string) {
	err := eng.session.DB(eng.dbName).C(collection).DropCollection()
	if err != nil {
		log.Printf("failed to drop collection")
	}
	eng.EnsureIndexes(collection)
}

func (eng *mongoEngine) Write(collection string, id string, cont Document) error {
	coll := eng.session.DB(eng.dbName).C(collection)
	if id == "" {
		return errors.New("missing id")
	}
	_, err := coll.Upsert(bson.D{{"uuid", id}}, cont) //TODO: parameterise uuid
	if err != nil {
		log.Printf("insert failed: %v\n", err)
	}
	return nil
}

func (eng *mongoEngine) Count(collection string) int {
	coll := eng.session.DB(eng.dbName).C(collection)
	count, err := coll.Count()
	if err != nil {
		panic(err)
	}
	return count
}

func (eng *mongoEngine) Load(collection, id string) (bool, Document, error) {
	c := eng.session.DB(eng.dbName).C(collection)
	var content Document
	err := c.Find(bson.M{"uuid": id}).One(&content)
	if err == mgo.ErrNotFound {
		return false, Document{}, nil
	}
	if err != nil {
		return false, Document{}, err
	}
	return true, content, nil
}

func (eng mongoEngine) All(collection string, stopchan chan struct{}) (chan Document, error) {
	cont := make(chan Document)

	go func() {
		defer close(cont)
		coll := eng.session.DB(eng.dbName).C(collection)
		iter := coll.Find(nil).Iter()
		var result Document
		for iter.Next(&result) {
			select {
			case <-stopchan:
				break
			case cont <- result:
			}
		}
		if err := iter.Close(); err != nil {
			panic(err)
		}
	}()

	return cont, nil
}
<commit_msg>remove _id on read<commit_after>package main

import (
	"errors"
	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"
	"log"
)

type mongoEngine struct {
	session *mgo.Session
	dbName  string
}

func (eng mongoEngine) Close() {
	// TODO
}

// NewMongoEngine returns an Engine based on a mongodb batabase backend
func NewMongoEngine(dbName string, hostPorts string) Engine {
	log.Printf("connecting to mongodb '%s'\n", hostPorts)
	s, err := mgo.Dial(hostPorts)
	if err != nil {
		panic(err)
	}
	s.SetMode(mgo.Monotonic, true)
	eng := &mongoEngine{
		session: s,
		dbName:  dbName,
	}

	eng.EnsureIndexes("organisations") //TODOL parameterise this

	return eng
}

func (eng *mongoEngine) EnsureIndexes(collection string) {
	c := eng.session.DB(eng.dbName).C(collection)

	eng.session.ResetIndexCache()

	// create collection if it's not there
	c.Create(&mgo.CollectionInfo{})

	err := c.EnsureIndex(mgo.Index{
		Key:        []string{"uuid"}, //TODO: parameterise this
		Unique:     true,
		DropDups:   true,
		Background: false,
		Sparse:     false,
	})
	if err != nil {
		panic(err)
	}

}

func (eng *mongoEngine) Drop(collection string) {
	err := eng.session.DB(eng.dbName).C(collection).DropCollection()
	if err != nil {
		log.Printf("failed to drop collection")
	}
	eng.EnsureIndexes(collection)
}

func (eng *mongoEngine) Write(collection string, id string, cont Document) error {
	coll := eng.session.DB(eng.dbName).C(collection)
	if id == "" {
		return errors.New("missing id")
	}
	_, err := coll.Upsert(bson.D{{"uuid", id}}, cont) //TODO: parameterise uuid
	if err != nil {
		log.Printf("insert failed: %v\n", err)
	}
	return nil
}

func (eng *mongoEngine) Count(collection string) int {
	coll := eng.session.DB(eng.dbName).C(collection)
	count, err := coll.Count()
	if err != nil {
		panic(err)
	}
	return count
}

func (eng *mongoEngine) Load(collection, id string) (bool, Document, error) {
	c := eng.session.DB(eng.dbName).C(collection)
	var content Document
	err := c.Find(bson.M{"uuid": id}).One(&content)
	if err == mgo.ErrNotFound {
		return false, Document{}, nil
	}
	if err != nil {
		return false, Document{}, err
	}
	cleanup(content)
	return true, content, nil
}

func (eng mongoEngine) All(collection string, stopchan chan struct{}) (chan Document, error) {
	cont := make(chan Document)

	go func() {
		defer close(cont)
		coll := eng.session.DB(eng.dbName).C(collection)
		iter := coll.Find(nil).Iter()
		var result Document
		for iter.Next(&result) {
			cleanup(result)
			select {
			case <-stopchan:
				break
			case cont <- result:
			}
		}
		if err := iter.Close(); err != nil {
			panic(err)
		}
	}()

	return cont, nil
}

func cleanup(doc Document) {
	delete(doc, "_id")
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"fmt"
	"log"

	"github.com/kr/pretty"
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
	"github.com/pborman/uuid"
)

type mongoEngine struct {
	session *mgo.Session
	dbName  string
}

func (eng mongoEngine) Close() {
	// TODO
}

// NewMongoEngine returns an Engine based on a mongodb database backend
func NewMongoEngine(dbName string, collections Collections, hostPorts string) Engine {
	log.Printf("connecting to mongodb '%s'\n", hostPorts)
	s, err := mgo.Dial(hostPorts)
	if err != nil {
		panic(err)
	}
	s.SetMode(mgo.Monotonic, true)
	eng := &mongoEngine{
		session: s,
		dbName:  dbName,
	}

	for _, coll := range collections {
		eng.EnsureIndexes(coll)
	}

	return eng
}

func (eng *mongoEngine) EnsureIndexes(collection Collection) {
	c := eng.session.DB(eng.dbName).C(collection.name)

	eng.session.ResetIndexCache()

	// create collection if it's not there
	c.Create(&mgo.CollectionInfo{})

	err := c.EnsureIndex(mgo.Index{
		Key:        []string{collection.idPropertyName},
		Unique:     true,
		DropDups:   true,
		Background: false,
		Sparse:     false,
	})
	if err != nil {
		log.Printf("failed to create index by %s on %s : %v\n", collection.name, collection.idPropertyName, err.Error())
	}

}

func (eng *mongoEngine) Drop(collection Collection) {
	err := eng.session.DB(eng.dbName).C(collection.name).DropCollection()
	if err != nil {
		log.Printf("failed to drop collection")
	}
	eng.EnsureIndexes(collection)
}

func (eng *mongoEngine) Write(collection Collection, id string, cont Document) error {
	coll := eng.session.DB(eng.dbName).C(collection.name)
	if id == "" {
		return errors.New("missing id")
	}
	_, err := coll.Upsert(bson.D{{collection.idPropertyName, id}}, cont)
	if err != nil {
		log.Printf("insert failed: %v\n", err)
	}
	return nil
}

func (eng *mongoEngine) Count(collection Collection) int {
	coll := eng.session.DB(eng.dbName).C(collection.name)
	count, err := coll.Count()
	if err != nil {
		panic(err)
	}
	return count
}

func (eng *mongoEngine) Load(collection Collection, id string) (bool, Document, error) {
	c := eng.session.DB(eng.dbName).C(collection.name)
	var content Document
	err := c.Find(bson.M{collection.idPropertyName: id}).One(&content)
	if err == mgo.ErrNotFound {
		return false, Document{}, nil
	}
	if err != nil {
		return false, Document{}, err
	}
	cleanup(content)
	return true, content, nil
}

func (eng *mongoEngine) Delete(collection Collection, id string) error {
	c := eng.session.DB(eng.dbName).C(collection.name)
	err := c.Remove(bson.M{collection.idPropertyName: id})
	if err != mgo.ErrNotFound {
		return err
	}
	return nil
}

func (eng mongoEngine) All(collection Collection, stopchan chan struct{}) (chan Document, error) {
	cont := make(chan Document)

	go func() {
		defer close(cont)
		coll := eng.session.DB(eng.dbName).C(collection.name)
		iter := coll.Find(nil).Iter()
		result := &Document{}
		for iter.Next(result) {
			cleanup(*result)
			select {
			case <-stopchan:
				break
			case cont <- *result:
				result = &Document{}
			}
		}
		if err := iter.Close(); err != nil {
			panic(err)
		}
	}()

	return cont, nil
}

func (eng mongoEngine) Ids(collection Collection, stopchan chan struct{}) (chan string, error) {
	ids := make(chan string)
	fmt.Println("Start")
	fmt.Printf("config: %# v\n", pretty.Formatter(collection))

	go func() {
		defer close(ids)
		coll := eng.session.DB(eng.dbName).C(collection.name)
		iter := coll.Find(nil).Select(bson.M{collection.idPropertyName: true}).Iter()
		var result map[string]interface{}
		for iter.Next(&result) {
			select {
			case <-stopchan:
				break
			case ids <- getUUIDString(result[collection.idPropertyName]):
			}
		}
		if err := iter.Close(); err != nil {
			panic(err)
		}
	}()
	fmt.Println("End")
	return ids, nil
}

func getUUIDString(uuidValue interface{}) string {
	if uuidString, ok := uuidValue.(string); ok {
		return uuidString
	} else if binaryUUID, ok := uuidValue.(bson.Binary); ok {
		return uuid.UUID(binaryUUID.Data).String()
	} else{
		fmt.Printf("UUID field is in an unknown format!\n %# v", pretty.Formatter(uuidValue))
		return ""
	}
}

func cleanup(doc Document) {
	delete(doc, "_id")
}
<commit_msg>Removed debug output<commit_after>package main

import (
	"errors"
	"fmt"
	"log"

	"github.com/kr/pretty"
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
	"github.com/pborman/uuid"
)

type mongoEngine struct {
	session *mgo.Session
	dbName  string
}

func (eng mongoEngine) Close() {
	// TODO
}

// NewMongoEngine returns an Engine based on a mongodb database backend
func NewMongoEngine(dbName string, collections Collections, hostPorts string) Engine {
	log.Printf("connecting to mongodb '%s'\n", hostPorts)
	s, err := mgo.Dial(hostPorts)
	if err != nil {
		panic(err)
	}
	s.SetMode(mgo.Monotonic, true)
	eng := &mongoEngine{
		session: s,
		dbName:  dbName,
	}

	for _, coll := range collections {
		eng.EnsureIndexes(coll)
	}

	return eng
}

func (eng *mongoEngine) EnsureIndexes(collection Collection) {
	c := eng.session.DB(eng.dbName).C(collection.name)

	eng.session.ResetIndexCache()

	// create collection if it's not there
	c.Create(&mgo.CollectionInfo{})

	err := c.EnsureIndex(mgo.Index{
		Key:        []string{collection.idPropertyName},
		Unique:     true,
		DropDups:   true,
		Background: false,
		Sparse:     false,
	})
	if err != nil {
		log.Printf("failed to create index by %s on %s : %v\n", collection.name, collection.idPropertyName, err.Error())
	}

}

func (eng *mongoEngine) Drop(collection Collection) {
	err := eng.session.DB(eng.dbName).C(collection.name).DropCollection()
	if err != nil {
		log.Printf("failed to drop collection")
	}
	eng.EnsureIndexes(collection)
}

func (eng *mongoEngine) Write(collection Collection, id string, cont Document) error {
	coll := eng.session.DB(eng.dbName).C(collection.name)
	if id == "" {
		return errors.New("missing id")
	}
	_, err := coll.Upsert(bson.D{{collection.idPropertyName, id}}, cont)
	if err != nil {
		log.Printf("insert failed: %v\n", err)
	}
	return nil
}

func (eng *mongoEngine) Count(collection Collection) int {
	coll := eng.session.DB(eng.dbName).C(collection.name)
	count, err := coll.Count()
	if err != nil {
		panic(err)
	}
	return count
}

func (eng *mongoEngine) Load(collection Collection, id string) (bool, Document, error) {
	c := eng.session.DB(eng.dbName).C(collection.name)
	var content Document
	err := c.Find(bson.M{collection.idPropertyName: id}).One(&content)
	if err == mgo.ErrNotFound {
		return false, Document{}, nil
	}
	if err != nil {
		return false, Document{}, err
	}
	cleanup(content)
	return true, content, nil
}

func (eng *mongoEngine) Delete(collection Collection, id string) error {
	c := eng.session.DB(eng.dbName).C(collection.name)
	err := c.Remove(bson.M{collection.idPropertyName: id})
	if err != mgo.ErrNotFound {
		return err
	}
	return nil
}

func (eng mongoEngine) All(collection Collection, stopchan chan struct{}) (chan Document, error) {
	cont := make(chan Document)

	go func() {
		defer close(cont)
		coll := eng.session.DB(eng.dbName).C(collection.name)
		iter := coll.Find(nil).Iter()
		result := &Document{}
		for iter.Next(result) {
			cleanup(*result)
			select {
			case <-stopchan:
				break
			case cont <- *result:
				result = &Document{}
			}
		}
		if err := iter.Close(); err != nil {
			panic(err)
		}
	}()

	return cont, nil
}

func (eng mongoEngine) Ids(collection Collection, stopchan chan struct{}) (chan string, error) {
	ids := make(chan string)
	go func() {
		defer close(ids)
		coll := eng.session.DB(eng.dbName).C(collection.name)
		iter := coll.Find(nil).Select(bson.M{collection.idPropertyName: true}).Iter()
		var result map[string]interface{}
		for iter.Next(&result) {
			select {
			case <-stopchan:
				break
			case ids <- getUUIDString(result[collection.idPropertyName]):
			}
		}
		if err := iter.Close(); err != nil {
			panic(err)
		}
	}()
	return ids, nil
}

func getUUIDString(uuidValue interface{}) string {
	if uuidString, ok := uuidValue.(string); ok {
		return uuidString
	} else if binaryUUID, ok := uuidValue.(bson.Binary); ok {
		return uuid.UUID(binaryUUID.Data).String()
	} else{
		fmt.Printf("UUID field is in an unknown format!\n %# v", pretty.Formatter(uuidValue))
		return ""
	}
}

func cleanup(doc Document) {
	delete(doc, "_id")
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"fmt"
	"io/ioutil"
	"os/exec"
	"regexp"
	"runtime"
)

// This function checks if Git is installed. If not it tries to install it.
// It will return an error if unsuccessful.
func EnsureGit(logger Logger) error {
	logger.Log("Checking if git command is available...")

	path, err := exec.LookPath("git")
	if err == nil {
		logger.Log("Git was found: " + path)

		return nil
	} else {
		logger.Log("Could not find a \"git\" command. I will try to install it.")
	}

	switch runtime.GOOS {
	case "windows":
		fmt.Println("Hello from Windows")
		return nil
	case "linux":
		return LinuxInstallGit(logger)
	case "darwin":
		fmt.Println("Hello from Mac")
		return nil
	}

	return nil
}

// This function assumes we are on a linux system and tries to find the distro
// type (debian based, fedora based etc). If "git" command is not available,
// it will try to install git using the system's package manager. If that is
// not possible (e.g. permission denied), it will simply return and error.
// This list of commands can be useful: https://git-scm.com/download/linux
func LinuxInstallGit(logger Logger) error {
	distributorID, err := SystemCommand(
		[]string{"lsb_release", "-i"}, ioutil.Discard)
	if err != nil {
		logger.Log("Could not determine the Linux distribution.")
		logger.Log("Please install Git and run the agent again.")

		return err
	}

	distroFuncMap := map[string](func(Logger) error){
		"Debian": InstallGitOnDebian,
		"Ubuntu": InstallGitOnDebian,
		"Arch":   InstallGitOnArch,
	}
	for distro, function := range distroFuncMap {
		matched, err := regexp.MatchString(distro, distributorID.output)
		if err != nil {
			return err
		}

		if matched {
			logger.Log("We seem to be on " + distro + ".")
			return function(logger)
		}
	}

	// If not already returned:
	return errors.New("I don't know how to install git on your distribution.\n Please install Git and run the agent again.")
}

func InstallGitOnDebian(logger Logger) error {
	logger.Log("Trying with apt-get.")
	res, err := SystemCommand([]string{"apt-get", "install", "-y", "git"}, logger)
	if err == nil && !res.success {
		// Stderr is already written no need to return it.
		return errors.New("I wasn't able to install git. Please install it manually and run the agent again.")
	}

	return err
}

func InstallGitOnArch(logger Logger) error {
	logger.Log("Trying with pacman.")
	res, err := SystemCommand([]string{"pacman", "-S", "--noconfirm", "git"}, logger)
	if err == nil && !res.success {
		// Stderr is already written no need to return it.
		return errors.New("I wasn't able to install git. Please install it manually and run the agent again.")
	}

	return err
}
<commit_msg>Add placeholder methods for Windows and Mac (git installation)<commit_after>package main

import (
	"errors"
	"io/ioutil"
	"os/exec"
	"regexp"
	"runtime"
)

// This function checks if Git is installed. If not it tries to install it.
// It will return an error if unsuccessful.
func EnsureGit(logger Logger) error {
	logger.Log("Checking if git command is available...")

	path, err := exec.LookPath("git")
	if err == nil {
		logger.Log("Git was found: " + path)

		return nil
	} else {
		logger.Log("Could not find a \"git\" command. I will try to install it.")
	}

	switch runtime.GOOS {
	case "windows":
		return WindowsInstallGit(logger)
	case "linux":
		return LinuxInstallGit(logger)
	case "darwin":
		return MacInstallGit(logger)
	}

	return nil
}

func WindowsInstallGit(logger Logger) error {
	return errors.New("I don't know how to install Git on Windows. Please install it manually and run the agent again.")
}

func MacInstallGit(logger Logger) error {
	return errors.New("I don't know how to install Git on Mac. Please install it manually and run the agent again.")
}

// This function assumes we are on a linux system and tries to find the distro
// type (debian based, fedora based etc). If "git" command is not available,
// it will try to install git using the system's package manager. If that is
// not possible (e.g. permission denied), it will simply return and error.
// This list of commands can be useful: https://git-scm.com/download/linux
func LinuxInstallGit(logger Logger) error {
	distributorID, err := SystemCommand(
		[]string{"lsb_release", "-i"}, ioutil.Discard)
	if err != nil {
		logger.Log("Could not determine the Linux distribution.")
		logger.Log("Please install Git and run the agent again.")

		return err
	}

	distroFuncMap := map[string](func(Logger) error){
		"Debian": InstallGitOnDebian,
		"Ubuntu": InstallGitOnDebian,
		"Arch":   InstallGitOnArch,
	}
	for distro, function := range distroFuncMap {
		matched, err := regexp.MatchString(distro, distributorID.output)
		if err != nil {
			return err
		}

		if matched {
			logger.Log("We seem to be on " + distro + ".")
			return function(logger)
		}
	}

	// If not already returned:
	return errors.New("I don't know how to install git on your distribution.\n Please install Git and run the agent again.")
}

func InstallGitOnDebian(logger Logger) error {
	logger.Log("Trying with apt-get.")
	res, err := SystemCommand([]string{"apt-get", "install", "-y", "git"}, logger)
	if err == nil && !res.success {
		// Stderr is already written no need to return it.
		return errors.New("I wasn't able to install git. Please install it manually and run the agent again.")
	}

	return err
}

func InstallGitOnArch(logger Logger) error {
	logger.Log("Trying with pacman.")
	res, err := SystemCommand([]string{"pacman", "-S", "--noconfirm", "git"}, logger)
	if err == nil && !res.success {
		// Stderr is already written no need to return it.
		return errors.New("I wasn't able to install git. Please install it manually and run the agent again.")
	}

	return err
}
<|endoftext|>


<commit_before>package protocol

import (
	"errors"
	"sync"
	"sync/atomic"
	"time"

	"github.com/golang/groupcache/lru"
	log "github.com/sirupsen/logrus"

	"github.com/bytom/bytom/consensus"
	"github.com/bytom/bytom/consensus/bcrp"
	"github.com/bytom/bytom/event"
	"github.com/bytom/bytom/protocol/bc"
	"github.com/bytom/bytom/protocol/bc/types"
	"github.com/bytom/bytom/protocol/state"
)

// msg type
const (
	MsgNewTx = iota
	MsgRemoveTx
	logModule = "protocol"
)

var (
	maxCachedErrTxs = 1000
	maxMsgChSize    = 1000
	maxNewTxNum     = 10000
	maxOrphanNum    = 2000

	orphanTTL                = 10 * time.Minute
	orphanExpireScanInterval = 3 * time.Minute

	// ErrTransactionNotExist is the pre-defined error message
	ErrTransactionNotExist = errors.New("transaction are not existed in the mempool")
	// ErrPoolIsFull indicates the pool is full
	ErrPoolIsFull = errors.New("transaction pool reach the max number")
	// ErrDustTx indicates transaction is dust tx
	ErrDustTx = errors.New("transaction is dust tx")
)

type TxMsgEvent struct{ TxMsg *TxPoolMsg }

// TxDesc store tx and related info for mining strategy
type TxDesc struct {
	Tx     *types.Tx `json:"transaction"`
	Added  time.Time `json:"-"`
	Height uint64    `json:"-"`
	Weight uint64    `json:"-"`
	Fee    uint64    `json:"-"`
}

// TxPoolMsg is use for notify pool changes
type TxPoolMsg struct {
	*TxDesc
	MsgType int
}

type orphanTx struct {
	*TxDesc
	expiration time.Time
}

// TxPool is use for store the unconfirmed transaction
type TxPool struct {
	lastUpdated     int64
	mtx             sync.RWMutex
	store           state.Store
	pool            map[bc.Hash]*TxDesc
	utxo            map[bc.Hash]*types.Tx
	orphans         map[bc.Hash]*orphanTx
	orphansByPrev   map[bc.Hash]map[bc.Hash]*orphanTx
	errCache        *lru.Cache
	eventDispatcher *event.Dispatcher
}

// NewTxPool init a new TxPool
func NewTxPool(store state.Store, dispatcher *event.Dispatcher) *TxPool {
	tp := &TxPool{
		lastUpdated:     time.Now().Unix(),
		store:           store,
		pool:            make(map[bc.Hash]*TxDesc),
		utxo:            make(map[bc.Hash]*types.Tx),
		orphans:         make(map[bc.Hash]*orphanTx),
		orphansByPrev:   make(map[bc.Hash]map[bc.Hash]*orphanTx),
		errCache:        lru.New(maxCachedErrTxs),
		eventDispatcher: dispatcher,
	}
	go tp.orphanExpireWorker()
	return tp
}

// AddErrCache add a failed transaction record to lru cache
func (tp *TxPool) AddErrCache(txHash *bc.Hash, err error) {
	tp.mtx.Lock()
	defer tp.mtx.Unlock()

	tp.errCache.Add(txHash, err)
}

// ExpireOrphan expire all the orphans that before the input time range
func (tp *TxPool) ExpireOrphan(now time.Time) {
	tp.mtx.Lock()
	defer tp.mtx.Unlock()

	for hash, orphan := range tp.orphans {
		if orphan.expiration.Before(now) {
			tp.removeOrphan(&hash)
		}
	}
}

// GetErrCache return the error of the transaction
func (tp *TxPool) GetErrCache(txHash *bc.Hash) error {
	tp.mtx.Lock()
	defer tp.mtx.Unlock()

	v, ok := tp.errCache.Get(txHash)
	if !ok {
		return nil
	}
	return v.(error)
}

// RemoveTransaction remove a transaction from the pool
func (tp *TxPool) RemoveTransaction(txHash *bc.Hash) {
	tp.mtx.Lock()
	defer tp.mtx.Unlock()

	txD, ok := tp.pool[*txHash]
	if !ok {
		return
	}

	for _, output := range txD.Tx.ResultIds {
		delete(tp.utxo, *output)
	}
	delete(tp.pool, *txHash)

	atomic.StoreInt64(&tp.lastUpdated, time.Now().Unix())
	tp.eventDispatcher.Post(TxMsgEvent{TxMsg: &TxPoolMsg{TxDesc: txD, MsgType: MsgRemoveTx}})
	log.WithFields(log.Fields{"module": logModule, "tx_id": txHash}).Debug("remove tx from mempool")
}

// GetTransaction return the TxDesc by hash
func (tp *TxPool) GetTransaction(txHash *bc.Hash) (*TxDesc, error) {
	tp.mtx.RLock()
	defer tp.mtx.RUnlock()

	if txD, ok := tp.pool[*txHash]; ok {
		return txD, nil
	}
	return nil, ErrTransactionNotExist
}

// GetTransactions return all the transactions in the pool
func (tp *TxPool) GetTransactions() []*TxDesc {
	tp.mtx.RLock()
	defer tp.mtx.RUnlock()

	txDs := make([]*TxDesc, len(tp.pool))
	i := 0
	for _, desc := range tp.pool {
		txDs[i] = desc
		i++
	}
	return txDs
}

// IsTransactionInPool check wheather a transaction in pool or not
func (tp *TxPool) IsTransactionInPool(txHash *bc.Hash) bool {
	tp.mtx.RLock()
	defer tp.mtx.RUnlock()

	_, ok := tp.pool[*txHash]
	return ok
}

// IsTransactionInErrCache check wheather a transaction in errCache or not
func (tp *TxPool) IsTransactionInErrCache(txHash *bc.Hash) bool {
	tp.mtx.RLock()
	defer tp.mtx.RUnlock()

	_, ok := tp.errCache.Get(txHash)
	return ok
}

// HaveTransaction IsTransactionInErrCache check is  transaction in errCache or pool
func (tp *TxPool) HaveTransaction(txHash *bc.Hash) bool {
	return tp.IsTransactionInPool(txHash) || tp.IsTransactionInErrCache(txHash)
}

func isTransactionNoBtmInput(tx *types.Tx) bool {
	for _, input := range tx.TxData.Inputs {
		if input.AssetID() == *consensus.BTMAssetID {
			return false
		}
	}
	return true
}

func isTransactionZeroOutput(tx *types.Tx) bool {
	for _, output := range tx.TxData.Outputs {
		if output.Amount == uint64(0) {
			return true
		}
	}
	return false
}

func isInvalidBCRPTx(tx *types.Tx) bool {
	for _, output := range tx.TxData.Outputs {
		if bcrp.IsBCRPScript(output.ControlProgram) {
			if output.AssetId.String() != consensus.BTMAssetID.String() || output.Amount < consensus.BCRPRequiredBTMAmount {
				return true
			}
		}
	}
	return false
}

func (tp *TxPool) IsDust(tx *types.Tx) bool {
	return isTransactionNoBtmInput(tx) || isTransactionZeroOutput(tx) || isInvalidBCRPTx(tx)
}

func (tp *TxPool) processTransaction(tx *types.Tx, height, fee uint64) (bool, error) {
	tp.mtx.Lock()
	defer tp.mtx.Unlock()

	txD := &TxDesc{
		Tx:     tx,
		Weight: tx.SerializedSize,
		Height: height,
		Fee:    fee,
	}
	requireParents, err := tp.checkOrphanUtxos(tx)
	if err != nil {
		return false, err
	}

	if len(requireParents) > 0 {
		return true, tp.addOrphan(txD, requireParents)
	}

	if err := tp.addTransaction(txD); err != nil {
		return false, err
	}

	tp.processOrphans(txD)
	return false, nil
}

// ProcessTransaction is the main entry for txpool handle new tx, ignore dust tx.
func (tp *TxPool) ProcessTransaction(tx *types.Tx, height, fee uint64) (bool, error) {
	if tp.IsDust(tx) {
		log.WithFields(log.Fields{"module": logModule, "tx_id": tx.ID.String()}).Warn("dust tx")
		return false, nil
	}
	return tp.processTransaction(tx, height, fee)
}

func (tp *TxPool) addOrphan(txD *TxDesc, requireParents []*bc.Hash) error {
	if len(tp.orphans) >= maxOrphanNum {
		return ErrPoolIsFull
	}

	orphan := &orphanTx{txD, time.Now().Add(orphanTTL)}
	tp.orphans[txD.Tx.ID] = orphan
	for _, hash := range requireParents {
		if _, ok := tp.orphansByPrev[*hash]; !ok {
			tp.orphansByPrev[*hash] = make(map[bc.Hash]*orphanTx)
		}
		tp.orphansByPrev[*hash][txD.Tx.ID] = orphan
	}
	return nil
}

func (tp *TxPool) addTransaction(txD *TxDesc) error {
	if len(tp.pool) >= maxNewTxNum {
		return ErrPoolIsFull
	}

	tx := txD.Tx
	txD.Added = time.Now()
	tp.pool[tx.ID] = txD
	for _, id := range tx.ResultIds {
		_, err := tx.OriginalOutput(*id)
		if err != nil {
			// error due to it's a retirement, utxo doesn't care this output type so skip it
			continue
		}

		tp.utxo[*id] = tx
	}

	atomic.StoreInt64(&tp.lastUpdated, time.Now().Unix())
	tp.eventDispatcher.Post(TxMsgEvent{TxMsg: &TxPoolMsg{TxDesc: txD, MsgType: MsgNewTx}})
	log.WithFields(log.Fields{"module": logModule, "tx_id": tx.ID.String()}).Debug("Add tx to mempool")
	return nil
}

func (tp *TxPool) checkOrphanUtxos(tx *types.Tx) ([]*bc.Hash, error) {
	view := state.NewUtxoViewpoint()
	if err := tp.store.GetTransactionsUtxo(view, []*bc.Tx{tx.Tx}); err != nil {
		return nil, err
	}

	hashes := []*bc.Hash{}
	for _, hash := range tx.SpentOutputIDs {
		if !view.CanSpend(&hash) && tp.utxo[hash] == nil {
			hashes = append(hashes, &hash)
		}
	}
	return hashes, nil
}

func (tp *TxPool) orphanExpireWorker() {
	ticker := time.NewTicker(orphanExpireScanInterval)
	defer ticker.Stop()

	for now := range ticker.C {
		tp.ExpireOrphan(now)
	}
}

func (tp *TxPool) processOrphans(txD *TxDesc) {
	processOrphans := []*orphanTx{}
	addRely := func(tx *types.Tx) {
		for _, outHash := range tx.ResultIds {
			orphans, ok := tp.orphansByPrev[*outHash]
			if !ok {
				continue
			}

			for _, orphan := range orphans {
				processOrphans = append(processOrphans, orphan)
			}
			delete(tp.orphansByPrev, *outHash)
		}
	}

	addRely(txD.Tx)
	for ; len(processOrphans) > 0; processOrphans = processOrphans[1:] {
		processOrphan := processOrphans[0]
		requireParents, err := tp.checkOrphanUtxos(processOrphan.Tx)
		if err != nil {
			log.WithFields(log.Fields{"module": logModule, "err": err}).Error("processOrphans got unexpect error")
			continue
		}

		if len(requireParents) == 0 {
			addRely(processOrphan.Tx)
			tp.removeOrphan(&processOrphan.Tx.ID)
			tp.addTransaction(processOrphan.TxDesc)
		}
	}
}

func (tp *TxPool) removeOrphan(hash *bc.Hash) {
	orphan, ok := tp.orphans[*hash]
	if !ok {
		return
	}

	for _, spend := range orphan.Tx.SpentOutputIDs {
		orphans, ok := tp.orphansByPrev[spend]
		if !ok {
			continue
		}

		if delete(orphans, *hash); len(orphans) == 0 {
			delete(tp.orphansByPrev, spend)
		}
	}
	delete(tp.orphans, *hash)
}
<commit_msg>feat(bcrp): close bcrp temporary<commit_after>package protocol

import (
	"errors"
	"sync"
	"sync/atomic"
	"time"

	"github.com/golang/groupcache/lru"
	log "github.com/sirupsen/logrus"

	"github.com/bytom/bytom/consensus"
	"github.com/bytom/bytom/consensus/bcrp"
	"github.com/bytom/bytom/event"
	"github.com/bytom/bytom/protocol/bc"
	"github.com/bytom/bytom/protocol/bc/types"
	"github.com/bytom/bytom/protocol/state"
)

// msg type
const (
	MsgNewTx = iota
	MsgRemoveTx
	logModule = "protocol"
)

var (
	maxCachedErrTxs = 1000
	maxMsgChSize    = 1000
	maxNewTxNum     = 10000
	maxOrphanNum    = 2000

	orphanTTL                = 10 * time.Minute
	orphanExpireScanInterval = 3 * time.Minute

	// ErrTransactionNotExist is the pre-defined error message
	ErrTransactionNotExist = errors.New("transaction are not existed in the mempool")
	// ErrPoolIsFull indicates the pool is full
	ErrPoolIsFull = errors.New("transaction pool reach the max number")
	// ErrDustTx indicates transaction is dust tx
	ErrDustTx = errors.New("transaction is dust tx")
)

type TxMsgEvent struct{ TxMsg *TxPoolMsg }

// TxDesc store tx and related info for mining strategy
type TxDesc struct {
	Tx     *types.Tx `json:"transaction"`
	Added  time.Time `json:"-"`
	Height uint64    `json:"-"`
	Weight uint64    `json:"-"`
	Fee    uint64    `json:"-"`
}

// TxPoolMsg is use for notify pool changes
type TxPoolMsg struct {
	*TxDesc
	MsgType int
}

type orphanTx struct {
	*TxDesc
	expiration time.Time
}

// TxPool is use for store the unconfirmed transaction
type TxPool struct {
	lastUpdated     int64
	mtx             sync.RWMutex
	store           state.Store
	pool            map[bc.Hash]*TxDesc
	utxo            map[bc.Hash]*types.Tx
	orphans         map[bc.Hash]*orphanTx
	orphansByPrev   map[bc.Hash]map[bc.Hash]*orphanTx
	errCache        *lru.Cache
	eventDispatcher *event.Dispatcher
}

// NewTxPool init a new TxPool
func NewTxPool(store state.Store, dispatcher *event.Dispatcher) *TxPool {
	tp := &TxPool{
		lastUpdated:     time.Now().Unix(),
		store:           store,
		pool:            make(map[bc.Hash]*TxDesc),
		utxo:            make(map[bc.Hash]*types.Tx),
		orphans:         make(map[bc.Hash]*orphanTx),
		orphansByPrev:   make(map[bc.Hash]map[bc.Hash]*orphanTx),
		errCache:        lru.New(maxCachedErrTxs),
		eventDispatcher: dispatcher,
	}
	go tp.orphanExpireWorker()
	return tp
}

// AddErrCache add a failed transaction record to lru cache
func (tp *TxPool) AddErrCache(txHash *bc.Hash, err error) {
	tp.mtx.Lock()
	defer tp.mtx.Unlock()

	tp.errCache.Add(txHash, err)
}

// ExpireOrphan expire all the orphans that before the input time range
func (tp *TxPool) ExpireOrphan(now time.Time) {
	tp.mtx.Lock()
	defer tp.mtx.Unlock()

	for hash, orphan := range tp.orphans {
		if orphan.expiration.Before(now) {
			tp.removeOrphan(&hash)
		}
	}
}

// GetErrCache return the error of the transaction
func (tp *TxPool) GetErrCache(txHash *bc.Hash) error {
	tp.mtx.Lock()
	defer tp.mtx.Unlock()

	v, ok := tp.errCache.Get(txHash)
	if !ok {
		return nil
	}
	return v.(error)
}

// RemoveTransaction remove a transaction from the pool
func (tp *TxPool) RemoveTransaction(txHash *bc.Hash) {
	tp.mtx.Lock()
	defer tp.mtx.Unlock()

	txD, ok := tp.pool[*txHash]
	if !ok {
		return
	}

	for _, output := range txD.Tx.ResultIds {
		delete(tp.utxo, *output)
	}
	delete(tp.pool, *txHash)

	atomic.StoreInt64(&tp.lastUpdated, time.Now().Unix())
	tp.eventDispatcher.Post(TxMsgEvent{TxMsg: &TxPoolMsg{TxDesc: txD, MsgType: MsgRemoveTx}})
	log.WithFields(log.Fields{"module": logModule, "tx_id": txHash}).Debug("remove tx from mempool")
}

// GetTransaction return the TxDesc by hash
func (tp *TxPool) GetTransaction(txHash *bc.Hash) (*TxDesc, error) {
	tp.mtx.RLock()
	defer tp.mtx.RUnlock()

	if txD, ok := tp.pool[*txHash]; ok {
		return txD, nil
	}
	return nil, ErrTransactionNotExist
}

// GetTransactions return all the transactions in the pool
func (tp *TxPool) GetTransactions() []*TxDesc {
	tp.mtx.RLock()
	defer tp.mtx.RUnlock()

	txDs := make([]*TxDesc, len(tp.pool))
	i := 0
	for _, desc := range tp.pool {
		txDs[i] = desc
		i++
	}
	return txDs
}

// IsTransactionInPool check wheather a transaction in pool or not
func (tp *TxPool) IsTransactionInPool(txHash *bc.Hash) bool {
	tp.mtx.RLock()
	defer tp.mtx.RUnlock()

	_, ok := tp.pool[*txHash]
	return ok
}

// IsTransactionInErrCache check wheather a transaction in errCache or not
func (tp *TxPool) IsTransactionInErrCache(txHash *bc.Hash) bool {
	tp.mtx.RLock()
	defer tp.mtx.RUnlock()

	_, ok := tp.errCache.Get(txHash)
	return ok
}

// HaveTransaction IsTransactionInErrCache check is  transaction in errCache or pool
func (tp *TxPool) HaveTransaction(txHash *bc.Hash) bool {
	return tp.IsTransactionInPool(txHash) || tp.IsTransactionInErrCache(txHash)
}

func isTransactionNoBtmInput(tx *types.Tx) bool {
	for _, input := range tx.TxData.Inputs {
		if input.AssetID() == *consensus.BTMAssetID {
			return false
		}
	}
	return true
}

func isTransactionZeroOutput(tx *types.Tx) bool {
	for _, output := range tx.TxData.Outputs {
		if output.Amount == uint64(0) {
			return true
		}
	}
	return false
}

func isInvalidBCRPTx(tx *types.Tx) bool {
	for _, output := range tx.TxData.Outputs {
		if bcrp.IsBCRPScript(output.ControlProgram) {
			return true
		}
	}
	return false
}

func (tp *TxPool) IsDust(tx *types.Tx) bool {
	return isTransactionNoBtmInput(tx) || isTransactionZeroOutput(tx) || isInvalidBCRPTx(tx)
}

func (tp *TxPool) processTransaction(tx *types.Tx, height, fee uint64) (bool, error) {
	tp.mtx.Lock()
	defer tp.mtx.Unlock()

	txD := &TxDesc{
		Tx:     tx,
		Weight: tx.SerializedSize,
		Height: height,
		Fee:    fee,
	}
	requireParents, err := tp.checkOrphanUtxos(tx)
	if err != nil {
		return false, err
	}

	if len(requireParents) > 0 {
		return true, tp.addOrphan(txD, requireParents)
	}

	if err := tp.addTransaction(txD); err != nil {
		return false, err
	}

	tp.processOrphans(txD)
	return false, nil
}

// ProcessTransaction is the main entry for txpool handle new tx, ignore dust tx.
func (tp *TxPool) ProcessTransaction(tx *types.Tx, height, fee uint64) (bool, error) {
	if tp.IsDust(tx) {
		log.WithFields(log.Fields{"module": logModule, "tx_id": tx.ID.String()}).Warn("dust tx")
		return false, nil
	}
	return tp.processTransaction(tx, height, fee)
}

func (tp *TxPool) addOrphan(txD *TxDesc, requireParents []*bc.Hash) error {
	if len(tp.orphans) >= maxOrphanNum {
		return ErrPoolIsFull
	}

	orphan := &orphanTx{txD, time.Now().Add(orphanTTL)}
	tp.orphans[txD.Tx.ID] = orphan
	for _, hash := range requireParents {
		if _, ok := tp.orphansByPrev[*hash]; !ok {
			tp.orphansByPrev[*hash] = make(map[bc.Hash]*orphanTx)
		}
		tp.orphansByPrev[*hash][txD.Tx.ID] = orphan
	}
	return nil
}

func (tp *TxPool) addTransaction(txD *TxDesc) error {
	if len(tp.pool) >= maxNewTxNum {
		return ErrPoolIsFull
	}

	tx := txD.Tx
	txD.Added = time.Now()
	tp.pool[tx.ID] = txD
	for _, id := range tx.ResultIds {
		_, err := tx.OriginalOutput(*id)
		if err != nil {
			// error due to it's a retirement, utxo doesn't care this output type so skip it
			continue
		}

		tp.utxo[*id] = tx
	}

	atomic.StoreInt64(&tp.lastUpdated, time.Now().Unix())
	tp.eventDispatcher.Post(TxMsgEvent{TxMsg: &TxPoolMsg{TxDesc: txD, MsgType: MsgNewTx}})
	log.WithFields(log.Fields{"module": logModule, "tx_id": tx.ID.String()}).Debug("Add tx to mempool")
	return nil
}

func (tp *TxPool) checkOrphanUtxos(tx *types.Tx) ([]*bc.Hash, error) {
	view := state.NewUtxoViewpoint()
	if err := tp.store.GetTransactionsUtxo(view, []*bc.Tx{tx.Tx}); err != nil {
		return nil, err
	}

	hashes := []*bc.Hash{}
	for _, hash := range tx.SpentOutputIDs {
		if !view.CanSpend(&hash) && tp.utxo[hash] == nil {
			hashes = append(hashes, &hash)
		}
	}
	return hashes, nil
}

func (tp *TxPool) orphanExpireWorker() {
	ticker := time.NewTicker(orphanExpireScanInterval)
	defer ticker.Stop()

	for now := range ticker.C {
		tp.ExpireOrphan(now)
	}
}

func (tp *TxPool) processOrphans(txD *TxDesc) {
	processOrphans := []*orphanTx{}
	addRely := func(tx *types.Tx) {
		for _, outHash := range tx.ResultIds {
			orphans, ok := tp.orphansByPrev[*outHash]
			if !ok {
				continue
			}

			for _, orphan := range orphans {
				processOrphans = append(processOrphans, orphan)
			}
			delete(tp.orphansByPrev, *outHash)
		}
	}

	addRely(txD.Tx)
	for ; len(processOrphans) > 0; processOrphans = processOrphans[1:] {
		processOrphan := processOrphans[0]
		requireParents, err := tp.checkOrphanUtxos(processOrphan.Tx)
		if err != nil {
			log.WithFields(log.Fields{"module": logModule, "err": err}).Error("processOrphans got unexpect error")
			continue
		}

		if len(requireParents) == 0 {
			addRely(processOrphan.Tx)
			tp.removeOrphan(&processOrphan.Tx.ID)
			tp.addTransaction(processOrphan.TxDesc)
		}
	}
}

func (tp *TxPool) removeOrphan(hash *bc.Hash) {
	orphan, ok := tp.orphans[*hash]
	if !ok {
		return
	}

	for _, spend := range orphan.Tx.SpentOutputIDs {
		orphans, ok := tp.orphansByPrev[spend]
		if !ok {
			continue
		}

		if delete(orphans, *hash); len(orphans) == 0 {
			delete(tp.orphansByPrev, spend)
		}
	}
	delete(tp.orphans, *hash)
}
<|endoftext|>