<commit_before>package fasthttp

import (
	"errors"
	"net"
	"strconv"
	"sync"
	"sync/atomic"
	"time"
)

var (
	dial          = (&tcpDialer{}).NewDial()
	dialDualStack = (&tcpDialer{DualStack: true}).NewDial()
)

// Dial dials the given TCP addr using tcp4.
//
// This function has the following additional features comparing to net.Dial:
//
//   * It reduces load on DNS resolver by caching resolved TCP addressed
//     for one minute.
//   * It dials all the resolved TCP addresses in round-robin manner until
//     connection is established. This may be useful if certain addresses
//     are temporarily unreachable.
//   * It returns ErrDialTimeout if connection cannot be established during
//     DefaultDialTimeout seconds.
//
// This dialer is intended for custom code wrapping before passing
// to Client.Dial or HostClient.Dial.
//
// For instance, per-host counters and/or limits may be implemented
// by such wrappers.
//
// The addr passed to the function must contain port. Example addr values:
//
//     * foobar.baz:443
//     * foo.bar:80
//     * aaa.com:8080
func Dial(addr string) (net.Conn, error) {
	return dial(addr)
}

// DialDualStack dials the given TCP addr using both tcp4 and tcp6.
//
// This function has the following additional features comparing to net.Dial:
//
//   * It reduces load on DNS resolver by caching resolved TCP addressed
//     for one minute.
//   * It dials all the resolved TCP addresses in round-robin manner until
//     connection is established. This may be useful if certain addresses
//     are temporarily unreachable.
//   * It returns ErrDialTimeout if connection cannot be established during
//     DefaultDialTimeout seconds.
//
// This dialer is intended for custom code wrapping before passing
// to Client.Dial or HostClient.Dial.
//
// For instance, per-host counters and/or limits may be implemented
// by such wrappers.
//
// The addr passed to the function must contain port. Example addr values:
//
//     * foobar.baz:443
//     * foo.bar:80
//     * aaa.com:8080
func DialDualStack(addr string) (net.Conn, error) {
	return dialDualStack(addr)
}

type tcpDialer struct {
	DualStack bool

	tcpAddrsLock sync.Mutex
	tcpAddrsMap  map[string]*tcpAddrEntry
}

func (d *tcpDialer) NewDial() DialFunc {
	if d.tcpAddrsMap != nil {
		panic("BUG: NewDial() already called")
	}

	d.tcpAddrsMap = make(map[string]*tcpAddrEntry)
	go d.tcpAddrsClean()

	return func(addr string) (net.Conn, error) {
		addrs, idx, err := d.getTCPAddrs(addr)
		if err != nil {
			return nil, err
		}
		network := "tcp4"
		if d.DualStack {
			network = "tcp"
		}

		var conn net.Conn
		startTime := time.Now()
		n := uint32(len(addrs))
		for n > 0 {
			conn, err = tryDial(network, &addrs[idx%n])
			if err == nil {
				return conn, nil
			}
			if time.Since(startTime) > DefaultDialTimeout {
				return nil, ErrDialTimeout
			}
			idx++
			n--
		}
		return nil, err
	}
}

func tryDial(network string, addr *net.TCPAddr) (net.Conn, error) {
	ch := make(chan dialResult, 1)
	go func() {
		var dr dialResult
		dr.conn, dr.err = net.DialTCP(network, nil, addr)
		ch <- dr
	}()
	select {
	case dr := <-ch:
		return dr.conn, dr.err
	case <-time.After(DefaultDialTimeout):
		return nil, ErrDialTimeout
	}
}

type dialResult struct {
	conn net.Conn
	err  error
}

// ErrDialTimeout is returned when TCP dialing is timed out.
var ErrDialTimeout = errors.New("dialing to the given TCP address timed out")

// DefaultDialTimeout is timeout used by Dial and DialDualStack
// for establishing TCP connections.
const DefaultDialTimeout = 10 * time.Second

type tcpAddrEntry struct {
	addrs    []net.TCPAddr
	addrsIdx uint32

	resolveTime time.Time
	pending     bool
}

const tcpAddrsCacheDuration = time.Minute

func (d *tcpDialer) tcpAddrsClean() {
	expireDuration := 2 * tcpAddrsCacheDuration
	for {
		time.Sleep(time.Second)
		t := time.Now()

		d.tcpAddrsLock.Lock()
		for k, e := range d.tcpAddrsMap {
			if t.Sub(e.resolveTime) > expireDuration {
				delete(d.tcpAddrsMap, k)
			}
		}
		d.tcpAddrsLock.Unlock()
	}
}

func (d *tcpDialer) getTCPAddrs(addr string) ([]net.TCPAddr, uint32, error) {
	d.tcpAddrsLock.Lock()
	e := d.tcpAddrsMap[addr]
	if e != nil && !e.pending && time.Since(e.resolveTime) > tcpAddrsCacheDuration {
		e.pending = true
		e = nil
	}
	d.tcpAddrsLock.Unlock()

	if e == nil {
		addrs, err := resolveTCPAddrs(addr, d.DualStack)
		if err != nil {
			d.tcpAddrsLock.Lock()
			e = d.tcpAddrsMap[addr]
			if e != nil && e.pending {
				e.pending = false
			}
			d.tcpAddrsLock.Unlock()
			return nil, 0, err
		}

		e = &tcpAddrEntry{
			addrs:       addrs,
			resolveTime: time.Now(),
		}

		d.tcpAddrsLock.Lock()
		d.tcpAddrsMap[addr] = e
		d.tcpAddrsLock.Unlock()
	}

	idx := uint32(0)
	if len(e.addrs) > 0 {
		idx = atomic.AddUint32(&e.addrsIdx, 1)
	}
	return e.addrs, idx, nil
}

func resolveTCPAddrs(addr string, dualStack bool) ([]net.TCPAddr, error) {
	host, portS, err := net.SplitHostPort(addr)
	if err != nil {
		return nil, err
	}
	port, err := strconv.Atoi(portS)
	if err != nil {
		return nil, err
	}

	ips, err := net.LookupIP(host)
	if err != nil {
		return nil, err
	}

	n := len(ips)
	addrs := make([]net.TCPAddr, 0, n)
	for i := 0; i < n; i++ {
		ip := ips[i]
		if !dualStack && ip.To4() == nil {
			continue
		}
		addrs = append(addrs, net.TCPAddr{
			IP:   ip,
			Port: port,
		})
	}
	return addrs, nil
}
<commit_msg>Added DialTimeout and DialDualStackTimeout functions<commit_after>package fasthttp

import (
	"errors"
	"net"
	"strconv"
	"sync"
	"sync/atomic"
	"time"
)

// Dial dials the given TCP addr using tcp4.
//
// This function has the following additional features comparing to net.Dial:
//
//   * It reduces load on DNS resolver by caching resolved TCP addressed
//     for DefaultDNSCacheDuration.
//   * It dials all the resolved TCP addresses in round-robin manner until
//     connection is established. This may be useful if certain addresses
//     are temporarily unreachable.
//   * It returns ErrDialTimeout if connection cannot be established during
//     DefaultDialTimeout seconds. Use DialTimeout for customizing dial timeout.
//
// This dialer is intended for custom code wrapping before passing
// to Client.Dial or HostClient.Dial.
//
// For instance, per-host counters and/or limits may be implemented
// by such wrappers.
//
// The addr passed to the function must contain port. Example addr values:
//
//     * foobar.baz:443
//     * foo.bar:80
//     * aaa.com:8080
func Dial(addr string) (net.Conn, error) {
	return getDialer(DefaultDialTimeout, false)(addr)
}

// DialTimeout dials the given TCP addr using tcp4 using the given timeout.
//
// This function has the following additional features comparing to net.Dial:
//
//   * It reduces load on DNS resolver by caching resolved TCP addressed
//     for DefaultDNSCacheDuration.
//   * It dials all the resolved TCP addresses in round-robin manner until
//     connection is established. This may be useful if certain addresses
//     are temporarily unreachable.
//
// This dialer is intended for custom code wrapping before passing
// to Client.Dial or HostClient.Dial.
//
// For instance, per-host counters and/or limits may be implemented
// by such wrappers.
//
// The addr passed to the function must contain port. Example addr values:
//
//     * foobar.baz:443
//     * foo.bar:80
//     * aaa.com:8080
func DialTimeout(addr string, timeout time.Duration) (net.Conn, error) {
	return getDialer(timeout, false)(addr)
}

// DialDualStack dials the given TCP addr using both tcp4 and tcp6.
//
// This function has the following additional features comparing to net.Dial:
//
//   * It reduces load on DNS resolver by caching resolved TCP addressed
//     for DefaultDNSCacheDuration.
//   * It dials all the resolved TCP addresses in round-robin manner until
//     connection is established. This may be useful if certain addresses
//     are temporarily unreachable.
//   * It returns ErrDialTimeout if connection cannot be established during
//     DefaultDialTimeout seconds. Use DialDualStackTimeout for custom dial
//     timeout.
//
// This dialer is intended for custom code wrapping before passing
// to Client.Dial or HostClient.Dial.
//
// For instance, per-host counters and/or limits may be implemented
// by such wrappers.
//
// The addr passed to the function must contain port. Example addr values:
//
//     * foobar.baz:443
//     * foo.bar:80
//     * aaa.com:8080
func DialDualStack(addr string) (net.Conn, error) {
	return getDialer(DefaultDialTimeout, true)(addr)
}

// DialDualStackTimeout dials the given TCP addr using both tcp4 and tcp6
// using the given timeout.
//
// This function has the following additional features comparing to net.Dial:
//
//   * It reduces load on DNS resolver by caching resolved TCP addressed
//     for DefaultDNSCacheDuration.
//   * It dials all the resolved TCP addresses in round-robin manner until
//     connection is established. This may be useful if certain addresses
//     are temporarily unreachable.
//
// This dialer is intended for custom code wrapping before passing
// to Client.Dial or HostClient.Dial.
//
// For instance, per-host counters and/or limits may be implemented
// by such wrappers.
//
// The addr passed to the function must contain port. Example addr values:
//
//     * foobar.baz:443
//     * foo.bar:80
//     * aaa.com:8080
func DialDualStackTimeout(addr string, timeout time.Duration) (net.Conn, error) {
	return getDialer(timeout, true)(addr)
}

func getDialer(timeout time.Duration, dualStack bool) DialFunc {
	if timeout <= 0 {
		timeout = DefaultDialTimeout
	}
	timeoutRounded := int(timeout.Seconds()*10 + 9)

	m := dialMap
	if dualStack {
		m = dialDualStackMap
	}

	dialMapLock.Lock()
	d := m[timeoutRounded]
	if d == nil {
		dialer := dialerStd
		if dualStack {
			dialer = dialerDualStack
		}
		d = dialer.NewDial(timeout)
		m[timeoutRounded] = d
	}
	dialMapLock.Unlock()
	return d
}

var (
	dialerStd       = &tcpDialer{}
	dialerDualStack = &tcpDialer{DualStack: true}

	dialMap          = make(map[int]DialFunc)
	dialDualStackMap = make(map[int]DialFunc)
	dialMapLock      sync.Mutex
)

type tcpDialer struct {
	DualStack bool

	tcpAddrsLock sync.Mutex
	tcpAddrsMap  map[string]*tcpAddrEntry

	once sync.Once
}

func (d *tcpDialer) NewDial(timeout time.Duration) DialFunc {
	d.once.Do(func() {
		d.tcpAddrsMap = make(map[string]*tcpAddrEntry)
		go d.tcpAddrsClean()
	})

	return func(addr string) (net.Conn, error) {
		addrs, idx, err := d.getTCPAddrs(addr)
		if err != nil {
			return nil, err
		}
		network := "tcp4"
		if d.DualStack {
			network = "tcp"
		}

		var conn net.Conn
		startTime := time.Now()
		n := uint32(len(addrs))
		for n > 0 {
			conn, err = tryDial(network, &addrs[idx%n], timeout)
			if err == nil {
				return conn, nil
			}
			timeout -= time.Since(startTime)
			if timeout < 0 {
				return nil, ErrDialTimeout
			}
			idx++
			n--
		}
		return nil, err
	}
}

func tryDial(network string, addr *net.TCPAddr, timeout time.Duration) (net.Conn, error) {
	ch := make(chan dialResult, 1)
	go func() {
		var dr dialResult
		dr.conn, dr.err = net.DialTCP(network, nil, addr)
		ch <- dr
	}()
	select {
	case dr := <-ch:
		return dr.conn, dr.err
	case <-time.After(timeout):
		return nil, ErrDialTimeout
	}
}

type dialResult struct {
	conn net.Conn
	err  error
}

// ErrDialTimeout is returned when TCP dialing is timed out.
var ErrDialTimeout = errors.New("dialing to the given TCP address timed out")

// DefaultDialTimeout is timeout used by Dial and DialDualStack
// for establishing TCP connections.
const DefaultDialTimeout = 5 * time.Second

type tcpAddrEntry struct {
	addrs    []net.TCPAddr
	addrsIdx uint32

	resolveTime time.Time
	pending     bool
}

// DefaultDNSCacheDuration is the duration for caching resolved TCP addresses
// by Dial* functions.
const DefaultDNSCacheDuration = time.Minute

func (d *tcpDialer) tcpAddrsClean() {
	expireDuration := 2 * DefaultDNSCacheDuration
	for {
		time.Sleep(time.Second)
		t := time.Now()

		d.tcpAddrsLock.Lock()
		for k, e := range d.tcpAddrsMap {
			if t.Sub(e.resolveTime) > expireDuration {
				delete(d.tcpAddrsMap, k)
			}
		}
		d.tcpAddrsLock.Unlock()
	}
}

func (d *tcpDialer) getTCPAddrs(addr string) ([]net.TCPAddr, uint32, error) {
	d.tcpAddrsLock.Lock()
	e := d.tcpAddrsMap[addr]
	if e != nil && !e.pending && time.Since(e.resolveTime) > DefaultDNSCacheDuration {
		e.pending = true
		e = nil
	}
	d.tcpAddrsLock.Unlock()

	if e == nil {
		addrs, err := resolveTCPAddrs(addr, d.DualStack)
		if err != nil {
			d.tcpAddrsLock.Lock()
			e = d.tcpAddrsMap[addr]
			if e != nil && e.pending {
				e.pending = false
			}
			d.tcpAddrsLock.Unlock()
			return nil, 0, err
		}

		e = &tcpAddrEntry{
			addrs:       addrs,
			resolveTime: time.Now(),
		}

		d.tcpAddrsLock.Lock()
		d.tcpAddrsMap[addr] = e
		d.tcpAddrsLock.Unlock()
	}

	idx := uint32(0)
	if len(e.addrs) > 0 {
		idx = atomic.AddUint32(&e.addrsIdx, 1)
	}
	return e.addrs, idx, nil
}

func resolveTCPAddrs(addr string, dualStack bool) ([]net.TCPAddr, error) {
	host, portS, err := net.SplitHostPort(addr)
	if err != nil {
		return nil, err
	}
	port, err := strconv.Atoi(portS)
	if err != nil {
		return nil, err
	}

	ips, err := net.LookupIP(host)
	if err != nil {
		return nil, err
	}

	n := len(ips)
	addrs := make([]net.TCPAddr, 0, n)
	for i := 0; i < n; i++ {
		ip := ips[i]
		if !dualStack && ip.To4() == nil {
			continue
		}
		addrs = append(addrs, net.TCPAddr{
			IP:   ip,
			Port: port,
		})
	}
	return addrs, nil
}
<|endoftext|>


<commit_before>// +build linux darwin dragonfly freebsd netbsd openbsd rumprun

// Package tcplisten provides customizable TCP net.Listener with various
// performance-related options:
//
//   - SO_REUSEPORT. This option allows linear scaling server performance
//     on multi-CPU servers.
//     See https://www.nginx.com/blog/socket-sharding-nginx-release-1-9-1/ for details.
//
//   - TCP_DEFER_ACCEPT. This option expects the server reads from the accepted
//     connection before writing to them.
//
//   - TCP_FASTOPEN. See https://lwn.net/Articles/508865/ for details.
//
// The package is derived from https://github.com/kavu/go_reuseport .
package tcplisten

import (
	"errors"
	"fmt"
	"net"
	"os"
	"syscall"
)

// Config provides options to enable on the returned socket.
type Config struct {
	// ReusePort enables SO_REUSEPORT.
	ReusePort bool

	// DeferAccept enables TCP_DEFER_ACCEPT.
	DeferAccept bool

	// FastOpen enables TCP_FASTOPEN.
	FastOpen bool
}

// NewListener returns TCP listener with options set in the Config.
//
// Only tcp4 and tcp6 networks are supported.
func (cfg *Config) NewListener(network, addr string) (net.Listener, error) {
	sa, soType, err := getSockaddr(network, addr)
	if err != nil {
		return nil, err
	}

	syscall.ForkLock.RLock()
	fd, err := syscall.Socket(soType, syscall.SOCK_STREAM, syscall.IPPROTO_TCP)
	if err == nil {
		syscall.CloseOnExec(fd)
	}
	syscall.ForkLock.RUnlock()
	if err != nil {
		return nil, err
	}

	if err = cfg.fdSetup(fd, sa, addr); err != nil {
		syscall.Close(fd)
		return nil, err
	}

	name := fmt.Sprintf("reuseport.%d.%s.%s", os.Getpid(), network, addr)
	file := os.NewFile(uintptr(fd), name)
	ln, err := net.FileListener(file)
	if err != nil {
		file.Close()
		return nil, err
	}

	if err = file.Close(); err != nil {
		ln.Close()
		return nil, err
	}

	return ln, nil
}

func (cfg *Config) fdSetup(fd int, sa syscall.Sockaddr, addr string) error {
	var err error

	if err = syscall.SetsockoptInt(fd, syscall.SOL_SOCKET, syscall.SO_REUSEADDR, 1); err != nil {
		return fmt.Errorf("cannot enable SO_REUSEADDR: %s", err)
	}

	if cfg.ReusePort {
		if err = syscall.SetsockoptInt(fd, syscall.SOL_SOCKET, soReusePort, 1); err != nil {
			return fmt.Errorf("cannot enable SO_REUSEPORT: %s", err)
		}
	}

	if cfg.DeferAccept {
		if err = enableDeferAccept(fd); err != nil {
			return err
		}
	}

	if cfg.FastOpen {
		if err = enableFastOpen(fd); err != nil {
			return err
		}
	}

	if err = syscall.Bind(fd, sa); err != nil {
		return fmt.Errorf("cannot bind to %q: %s", addr, err)
	}

	if err = syscall.Listen(fd, syscall.SOMAXCONN); err != nil {
		return fmt.Errorf("cannot listen on %q: %s", addr, err)
	}

	return nil
}

func getSockaddr(network, addr string) (sa syscall.Sockaddr, soType int, err error) {
	if network != "tcp4" && network != "tcp6" {
		return nil, -1, errors.New("only tcp4 and tcp6 network is supported")
	}

	tcpAddr, err := net.ResolveTCPAddr(network, addr)
	if err != nil {
		return nil, -1, err
	}

	switch network {
	case "tcp4":
		var sa4 syscall.SockaddrInet4
		sa4.Port = tcpAddr.Port
		copy(sa4.Addr[:], tcpAddr.IP.To4())
		return &sa4, syscall.AF_INET, nil
	case "tcp6":
		var sa6 syscall.SockaddrInet6
		sa6.Port = tcpAddr.Port
		copy(sa6.Addr[:], tcpAddr.IP.To16())
		if tcpAddr.Zone != "" {
			ifi, err := net.InterfaceByName(tcpAddr.Zone)
			if err != nil {
				return nil, -1, err
			}
			sa6.ZoneId = uint32(ifi.Index)
		}
		return &sa6, syscall.AF_INET6, nil
	default:
		return nil, -1, errors.New("Unknown network type " + network)
	}
}
<commit_msg>documentation fix<commit_after>// +build linux darwin dragonfly freebsd netbsd openbsd rumprun

// Package tcplisten provides customizable TCP net.Listener with various
// performance-related options:
//
//   - SO_REUSEPORT. This option allows linear scaling server performance
//     on multi-CPU servers.
//     See https://www.nginx.com/blog/socket-sharding-nginx-release-1-9-1/ for details.
//
//   - TCP_DEFER_ACCEPT. This option expects the server reads from the accepted
//     connection before writing to them.
//
//   - TCP_FASTOPEN. See https://lwn.net/Articles/508865/ for details.
//
// The package is derived from https://github.com/kavu/go_reuseport .
package tcplisten

import (
	"errors"
	"fmt"
	"net"
	"os"
	"syscall"
)

// Config provides options to enable on the returned listener.
type Config struct {
	// ReusePort enables SO_REUSEPORT.
	ReusePort bool

	// DeferAccept enables TCP_DEFER_ACCEPT.
	DeferAccept bool

	// FastOpen enables TCP_FASTOPEN.
	FastOpen bool
}

// NewListener returns TCP listener with options set in the Config.
//
// Only tcp4 and tcp6 networks are supported.
func (cfg *Config) NewListener(network, addr string) (net.Listener, error) {
	sa, soType, err := getSockaddr(network, addr)
	if err != nil {
		return nil, err
	}

	syscall.ForkLock.RLock()
	fd, err := syscall.Socket(soType, syscall.SOCK_STREAM, syscall.IPPROTO_TCP)
	if err == nil {
		syscall.CloseOnExec(fd)
	}
	syscall.ForkLock.RUnlock()
	if err != nil {
		return nil, err
	}

	if err = cfg.fdSetup(fd, sa, addr); err != nil {
		syscall.Close(fd)
		return nil, err
	}

	name := fmt.Sprintf("reuseport.%d.%s.%s", os.Getpid(), network, addr)
	file := os.NewFile(uintptr(fd), name)
	ln, err := net.FileListener(file)
	if err != nil {
		file.Close()
		return nil, err
	}

	if err = file.Close(); err != nil {
		ln.Close()
		return nil, err
	}

	return ln, nil
}

func (cfg *Config) fdSetup(fd int, sa syscall.Sockaddr, addr string) error {
	var err error

	if err = syscall.SetsockoptInt(fd, syscall.SOL_SOCKET, syscall.SO_REUSEADDR, 1); err != nil {
		return fmt.Errorf("cannot enable SO_REUSEADDR: %s", err)
	}

	if cfg.ReusePort {
		if err = syscall.SetsockoptInt(fd, syscall.SOL_SOCKET, soReusePort, 1); err != nil {
			return fmt.Errorf("cannot enable SO_REUSEPORT: %s", err)
		}
	}

	if cfg.DeferAccept {
		if err = enableDeferAccept(fd); err != nil {
			return err
		}
	}

	if cfg.FastOpen {
		if err = enableFastOpen(fd); err != nil {
			return err
		}
	}

	if err = syscall.Bind(fd, sa); err != nil {
		return fmt.Errorf("cannot bind to %q: %s", addr, err)
	}

	if err = syscall.Listen(fd, syscall.SOMAXCONN); err != nil {
		return fmt.Errorf("cannot listen on %q: %s", addr, err)
	}

	return nil
}

func getSockaddr(network, addr string) (sa syscall.Sockaddr, soType int, err error) {
	if network != "tcp4" && network != "tcp6" {
		return nil, -1, errors.New("only tcp4 and tcp6 network is supported")
	}

	tcpAddr, err := net.ResolveTCPAddr(network, addr)
	if err != nil {
		return nil, -1, err
	}

	switch network {
	case "tcp4":
		var sa4 syscall.SockaddrInet4
		sa4.Port = tcpAddr.Port
		copy(sa4.Addr[:], tcpAddr.IP.To4())
		return &sa4, syscall.AF_INET, nil
	case "tcp6":
		var sa6 syscall.SockaddrInet6
		sa6.Port = tcpAddr.Port
		copy(sa6.Addr[:], tcpAddr.IP.To16())
		if tcpAddr.Zone != "" {
			ifi, err := net.InterfaceByName(tcpAddr.Zone)
			if err != nil {
				return nil, -1, err
			}
			sa6.ZoneId = uint32(ifi.Index)
		}
		return &sa6, syscall.AF_INET6, nil
	default:
		return nil, -1, errors.New("Unknown network type " + network)
	}
}
<|endoftext|>


<commit_before>package dispel

import (
	"bytes"
	"fmt"
	"io"
	"sort"
	"strings"
	"text/template"
	"unicode"
)

//go:generate -command asset go run ./asset.go
//go:generate asset --var=routesTmpl routes.go.tmpl
//go:generate asset --var=handlersTmpl handlers.go.tmpl
//go:generate asset --var=handlerfuncsTmpl handlerfuncs.go.tmpl
//go:generate asset --var=typesTmpl types.go.tmpl

var templatesMap = map[string]string{
	TemplateRoutes:       routesTmpl,
	TemplateHandlers:     handlersTmpl,
	TemplateHandlerfuncs: handlerfuncsTmpl,
	TemplateTypes:        typesTmpl,
}

// The templates available in a TemplateBundle.
const (
	TemplateRoutes       = "routes"
	TemplateHandlers     = "handlers"
	TemplateHandlerfuncs = "handlerfuncs"
	TemplateTypes        = "types"
)

// TemplateNames returns the same list than TemplateBundle.Names(), but doesn't require
// to create a Template instance.
func TemplateNames() []string {
	var a []string
	for name := range templatesMap {
		a = append(a, name)
	}
	sort.Strings(a)
	return a
}

func tmpl(a asset) string {
	return a.Content
}

// TemplateBundle represents a bundle of templates for generating the various dispel API source files.
type TemplateBundle struct {
	t  *template.Template
	sp *SchemaParser
}

// ExecuteTemplate executes the template named by name, using the ctx TemplateContext.
func (t *TemplateBundle) ExecuteTemplate(wr io.Writer, name string, ctx *TemplateContext) error {
	ctx.sp = t.sp
	return t.t.ExecuteTemplate(wr, name, ctx)
}

// ExecuteCustomTemplate parses then executes the template text, using the ctx TemplateContext.
func (t *TemplateBundle) ExecuteCustomTemplate(wr io.Writer, text string, ctx *TemplateContext) error {
	ctx.sp = t.sp
	tt, err := t.t.New("custom").Parse(text)
	if err != nil {
		return err
	}
	return tt.Execute(wr, ctx)
}

// Names returns the list of templates available in the Template bundle.
func (t *TemplateBundle) Names() []string {
	var a []string
	for _, tmpl := range t.t.Templates() {
		a = append(a, tmpl.Name())
	}
	return a
}

// TemplateContext represents the context passed to a Template.
type TemplateContext struct {
	Prgm                string   // name of the program generating the source
	PkgName             string   // package name for which source code is generated
	Routes              Routes   // routes parsed by the SchemaParser
	HandlerReceiverType string   // type which acts as the receiver of the handler funcs.
	ExistingHandlers    []string // list of existing handler funcs in the target package, with HandlerReceiverType as the receiver
	ExistingTypes       []string // list of existing types in the target package.

	sp *SchemaParser
}

func handlerFuncName(routeMethod string, routeName string) string {
	return strings.ToLower(routeMethod) + symbolName(routeName)
}

// NewTemplateBundle returns a new Template based on the SchemaParser.
func NewTemplateBundle(sp *SchemaParser) (*TemplateBundle, error) {
	t := template.New("").Funcs(template.FuncMap{
		"tolower":    strings.ToLower,
		"capitalize": capitalize,
		"symbolName": symbolName,
		"hasItem": func(a []string, s string) bool {
			for _, item := range a {
				if s == item {
					return true
				}
			}
			return false
		},
		"handlerFuncName": handlerFuncName,
		"allHandlerFuncsImplemented": func(routes Routes, existingHandlers []string) bool {
		LRoutesLoop:
			for _, route := range routes {
				fname := handlerFuncName(route.Method, route.Name)
				for _, h := range existingHandlers {
					if h == fname {
						continue LRoutesLoop
					}
				}
				return false
			}
			return true
		},
		"varname": func(s string) string {
			var buf bytes.Buffer
			var pickedFirstNonSymbolRune bool
			for _, r := range s {
				switch {
				case r == '*':
					continue
				case unicode.IsUpper(r):
					_, _ = buf.WriteRune(unicode.ToLower(r))
					pickedFirstNonSymbolRune = true
				default:
					if !pickedFirstNonSymbolRune {
						_, _ = buf.WriteRune(unicode.ToLower(r))
						pickedFirstNonSymbolRune = true
					}
				}
			}
			return buf.String()
		},
		"printTypeDef": func(j JSONType) string {
			return fmt.Sprintf("type %s %s", sp.JSONToGoType(j, false), sp.JSONToGoType(j, true))
		},
		"printTypeName": func(j JSONType) string {
			return sp.JSONToGoType(j, false)
		},
		"printSmartDerefType": func(j JSONType) string {
			// really smart
			switch j.(type) {
			case JSONObject:
				return "*" + sp.JSONToGoType(j, false)
			}
			return sp.JSONToGoType(j, false)
		},
	})
	for name, tmpl := range templatesMap {
		var err error
		t, err = t.New(name).Parse(tmpl)
		if err != nil {
			return nil, fmt.Errorf("template %s: %v", name, err)
		}
	}
	return &TemplateBundle{t: t, sp: sp}, nil
}
<commit_msg>Export the funcs used in TemplateBundle<commit_after>package dispel

import (
	"bytes"
	"fmt"
	"io"
	"sort"
	"strings"
	"text/template"
	"unicode"
)

//go:generate -command asset go run ./asset.go
//go:generate asset --var=routesTmpl routes.go.tmpl
//go:generate asset --var=handlersTmpl handlers.go.tmpl
//go:generate asset --var=handlerfuncsTmpl handlerfuncs.go.tmpl
//go:generate asset --var=typesTmpl types.go.tmpl

var templatesMap = map[string]string{
	TemplateRoutes:       routesTmpl,
	TemplateHandlers:     handlersTmpl,
	TemplateHandlerfuncs: handlerfuncsTmpl,
	TemplateTypes:        typesTmpl,
}

// The templates available in a TemplateBundle.
const (
	TemplateRoutes       = "routes"
	TemplateHandlers     = "handlers"
	TemplateHandlerfuncs = "handlerfuncs"
	TemplateTypes        = "types"
)

// TemplateNames returns the same list than TemplateBundle.Names(), but doesn't require
// to create a Template instance.
func TemplateNames() []string {
	var a []string
	for name := range templatesMap {
		a = append(a, name)
	}
	sort.Strings(a)
	return a
}

func tmpl(a asset) string {
	return a.Content
}

// TemplateBundle represents a bundle of templates for generating the various dispel API source files.
type TemplateBundle struct {
	t  *template.Template
	sp *SchemaParser
}

// ExecuteTemplate executes the template named by name, using the ctx TemplateContext.
func (t *TemplateBundle) ExecuteTemplate(wr io.Writer, name string, ctx *TemplateContext) error {
	ctx.sp = t.sp
	return t.t.ExecuteTemplate(wr, name, ctx)
}

// ExecuteCustomTemplate parses then executes the template text, using the ctx TemplateContext.
func (t *TemplateBundle) ExecuteCustomTemplate(wr io.Writer, text string, ctx *TemplateContext) error {
	ctx.sp = t.sp
	tt, err := t.t.New("custom").Parse(text)
	if err != nil {
		return err
	}
	return tt.Execute(wr, ctx)
}

// Names returns the list of templates available in the Template bundle.
func (t *TemplateBundle) Names() []string {
	var a []string
	for _, tmpl := range t.t.Templates() {
		a = append(a, tmpl.Name())
	}
	return a
}

// TemplateContext represents the context passed to a Template.
type TemplateContext struct {
	Prgm                string   // name of the program generating the source
	PkgName             string   // package name for which source code is generated
	Routes              Routes   // routes parsed by the SchemaParser
	HandlerReceiverType string   // type which acts as the receiver of the handler funcs.
	ExistingHandlers    []string // list of existing handler funcs in the target package, with HandlerReceiverType as the receiver
	ExistingTypes       []string // list of existing types in the target package.

	sp *SchemaParser
}

func handlerFuncName(routeMethod string, routeName string) string {
	return strings.ToLower(routeMethod) + symbolName(routeName)
}

// NewTemplateBundle returns a new Template based on the SchemaParser.
func NewTemplateBundle(sp *SchemaParser) (*TemplateBundle, error) {
	t := template.New("").Funcs(NewTemplateFuncMap(sp))
	for name, tmpl := range templatesMap {
		var err error
		t, err = t.New(name).Parse(tmpl)
		if err != nil {
			return nil, fmt.Errorf("template %s: %v", name, err)
		}
	}
	return &TemplateBundle{t: t, sp: sp}, nil
}

// NewTemplateFuncMap returns the template.FuncMap used in a TemplateBundle.
func NewTemplateFuncMap(sp *SchemaParser) template.FuncMap {
	return template.FuncMap{
		"tolower":    strings.ToLower,
		"capitalize": capitalize,
		"symbolName": symbolName,
		"hasItem": func(a []string, s string) bool {
			for _, item := range a {
				if s == item {
					return true
				}
			}
			return false
		},
		"handlerFuncName": handlerFuncName,
		"allHandlerFuncsImplemented": func(routes Routes, existingHandlers []string) bool {
		LRoutesLoop:
			for _, route := range routes {
				fname := handlerFuncName(route.Method, route.Name)
				for _, h := range existingHandlers {
					if h == fname {
						continue LRoutesLoop
					}
				}
				return false
			}
			return true
		},
		"varname": func(s string) string {
			var buf bytes.Buffer
			var pickedFirstNonSymbolRune bool
			for _, r := range s {
				switch {
				case r == '*':
					continue
				case unicode.IsUpper(r):
					_, _ = buf.WriteRune(unicode.ToLower(r))
					pickedFirstNonSymbolRune = true
				default:
					if !pickedFirstNonSymbolRune {
						_, _ = buf.WriteRune(unicode.ToLower(r))
						pickedFirstNonSymbolRune = true
					}
				}
			}
			return buf.String()
		},
		"printTypeDef": func(j JSONType) string {
			return fmt.Sprintf("type %s %s", sp.JSONToGoType(j, false), sp.JSONToGoType(j, true))
		},
		"printTypeName": func(j JSONType) string {
			return sp.JSONToGoType(j, false)
		},
		"printSmartDerefType": func(j JSONType) string {
			// really smart
			switch j.(type) {
			case JSONObject:
				return "*" + sp.JSONToGoType(j, false)
			}
			return sp.JSONToGoType(j, false)
		},
	}
}
<|endoftext|>


<commit_before>// Copyright 2019 Google Inc. All Rights Reserved.
// This file is available under the Apache license.

package mtail_test

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"net"
	"net/http"
	"os"
	"path/filepath"
	"runtime"
	"testing"
	"time"

	"github.com/golang/glog"
	"github.com/google/mtail/internal/metrics"
	"github.com/google/mtail/internal/mtail"
	"github.com/google/mtail/internal/testutil"
)

// makeServer makes a new Server for use in tests, but does not start
// the server. It returns the server, or any errors the new server creates.
func makeServer(tb testing.TB, pollInterval time.Duration, options ...mtail.Option) (*mtail.Server, error) {
	tb.Helper()
	ctx := context.Background()

	return mtail.New(ctx, metrics.NewStore(), options...)
}

// startUNIXSocketServer creates a new Server serving through a UNIX
// socket and starts it running. It returns the server, and a cleanup function.
// startUNIXSocketServer differs from TestStartServer in that it uses UNIX sockets
// (see the usage of mtail.BindUnixSocket) instead of TCP sockets.
func startUNIXSocketServer(tb testing.TB, pollInterval time.Duration, options ...mtail.Option) (*mtail.Server, func()) {
	tb.Helper()

	tmpDir := testutil.TestTempDir(tb)

	unixSocket := filepath.Join(tmpDir, "mtail_test.socket")
	options = append(options, mtail.BindUnixSocket(unixSocket))

	m, err := makeServer(tb, pollInterval, options...)
	testutil.FatalIfErr(tb, err)

	errc := make(chan error, 1)
	go func() {
		err := m.Run()
		errc <- err
	}()

	glog.Infof("check that server is listening")

	addr, err := net.ResolveUnixAddr("unix", unixSocket)
	testutil.FatalIfErr(tb, err)
	_, err = net.DialUnix("unix", nil, addr)
	testutil.FatalIfErr(tb, err)

	return m, func() {

		select {
		case err := <-errc:
			testutil.FatalIfErr(tb, err)
		case <-time.After(5 * time.Second):
			buf := make([]byte, 1<<16)
			n := runtime.Stack(buf, true)
			fmt.Fprintf(os.Stderr, "%s", buf[0:n])
			tb.Fatal("timeout waiting for shutdown")
		}
	}
}

// getMetricFromUNIXSocket fetches the name metrics from the Server serving through
// the UNIX socket sockPath, and returns the value of one named name.  Callers are
// responsible for type assertions on the returned value.
// TODO: move this method to testing.go in order to abstract away the http.Transport part from TestGetMetric
func getMetricFromUNIXSocket(tb testing.TB, sockPath, name string) interface{} {
	uri := "http://unix/debug/vars"
	client := &http.Client{
		Transport: &http.Transport{
			DialContext: func(_ context.Context, _, _ string) (net.Conn, error) {
				return net.Dial("unix", sockPath)
			},
		},
		Timeout: 5 * time.Second,
	}
	resp, err := client.Get(uri)
	if err != nil {
		tb.Fatal(err)
	}
	buf := new(bytes.Buffer)
	n, err := buf.ReadFrom(resp.Body)
	testutil.FatalIfErr(tb, err)
	glog.V(2).Infof("TestGetMetric: http client read %d bytes from debug/vars", n)
	var r map[string]interface{}
	if err := json.Unmarshal(buf.Bytes(), &r); err != nil {
		tb.Fatalf("%s: body was %s", err, buf.String())
	}
	glog.Infof("TestGetMetric: returned value for %s: %v", name, r[name])
	return r[name]
}

// expectMetricDelta checks to see if the difference between a and b is want;
// it assumes both values are float64s that came from a TestGetMetric.
func expectMetricDelta(tb testing.TB, a, b interface{}, want float64) {
	tb.Helper()
	if a == nil {
		a = 0.
	}
	if b == nil {
		b = 0.
	}
	delta := a.(float64) - b.(float64)
	if delta != want {
		tb.Errorf("Unexpected delta: got %v - %v = %g, want %g", a, b, delta, want)
	}
}

func TestBasicUNIXSockets(t *testing.T) {
	t.Skip("broken because unixSocket is in /var/run")
	unixSocket := "/var/run/mtail_test.socket"

	if testing.Verbose() {
		testutil.SetFlag(t, "vmodule", "tail=2,filestream=2")
	}
	logDir := testutil.TestTempDir(t)

	_, stopM := startUNIXSocketServer(t, 0, mtail.LogPathPatterns(logDir+"/*"), mtail.ProgramPath("../../examples/linecount.mtail"))
	defer stopM()

	startLineCount := getMetricFromUNIXSocket(t, unixSocket, "lines_total")
	startLogCount := getMetricFromUNIXSocket(t, unixSocket, "log_count")

	time.Sleep(1 * time.Second)

	logFile := filepath.Join(logDir, "log")

	f := testutil.TestOpenFile(t, logFile)

	for i := 1; i <= 3; i++ {
		testutil.WriteString(t, f, fmt.Sprintf("%d\n", i))
		time.Sleep(1 * time.Second)
	}

	endLineCount := getMetricFromUNIXSocket(t, unixSocket, "lines_total")
	endLogCount := getMetricFromUNIXSocket(t, unixSocket, "log_count")

	expectMetricDelta(t, endLineCount, startLineCount, 3)
	expectMetricDelta(t, endLogCount, startLogCount, 1)
}
<commit_msg>Reenable the unix socket listen test.<commit_after>// Copyright 2019 Google Inc. All Rights Reserved.
// This file is available under the Apache license.

package mtail_test

import (
	"net"
	"path/filepath"
	"testing"

	"github.com/golang/glog"
	"github.com/google/mtail/internal/mtail"
	"github.com/google/mtail/internal/testutil"
)

func TestBasicUNIXSockets(t *testing.T) {
	testutil.SkipIfShort(t)
	tmpDir := testutil.TestTempDir(t)
	sockListenAddr := filepath.Join(tmpDir, "mtail_test.sock")

	_, stopM := mtail.TestStartServer(t, 0, 1, mtail.LogPathPatterns(tmpDir+"/*"), mtail.ProgramPath("../../examples/linecount.mtail"), mtail.BindUnixSocket(sockListenAddr))
	defer stopM()

	glog.Infof("check that server is listening")

	addr, err := net.ResolveUnixAddr("unix", sockListenAddr)
	testutil.FatalIfErr(t, err)
	_, err = net.DialUnix("unix", nil, addr)
	testutil.FatalIfErr(t, err)
}
<|endoftext|>


<commit_before>package feedloggr

const tmplCSS string = `
body {
	margin: 30px 5%;
	line-height: 1.5;
	font-size: 14px;
	font-family: monospace;
	background-color: #FFF;
	color: #444;
}
a, a:visited {
	color: #444;
	text-decoration:none;
}
a:hover {
	color: #000;
	text-decoration: underline;
}

nav {
	text-align: center;
	margin-bottom: 20px;
}
nav > a:hover {
	background-color: #E6E6E6;
	border-color: #ADADAD;
}
section {
	background-color: #FFF;
	margin-bottom: 20px;
}
section > h1, nav > a {
	border: 1px solid #DDD;
	border-radius: 3px;
	padding: 10px;
	margin: 0;
	font-size: 16px;
	font-weight: bold;
	text-align: center;
}
section > h1 {
	background-color: #EEE;
}
section > ul, p {
	padding: 0;
	list-style: inside;
}
section > ul > li {
        margin-bottom: 2px;
}
section > ul > li > a:visited {
	color: #AAA;
}
footer {
	text-align: center;
	font-size: 12px;
}
.center {
	text-align: center;
}
`

const tmplPage string = `
<!doctype html>

	<meta charset="utf-8">
	<meta name="viewport" content="width=device-width">
	<title>{{.CurrentDate}} | Feedloggr</title>
	<link href="./style.css" rel="stylesheet" type="text/css">

		<
		Latest
		>

	{{range .Feeds}}

{{.Title}}

		{{if .Error}}

Error while updating feed:
{{.Error}}

		{{else}}

{{range .Items}}

{{.Title}}

		{{end}}

		{{end}}

	{{else}}

Sorry, no news for today!

	{{end}}

		Generated with Feedloggr

`
<commit_msg>Minor css fixes<commit_after>package feedloggr

const tmplCSS string = `
body {
	margin: 30px 5%;
	line-height: 1.5;
	font-size: 13px;
	font-family: monospace;
	background-color: #FFF;
	color: #444;
}
a, a:visited {
	color: #444;
	text-decoration: none;
}
a:hover {
	color: #000;
	text-decoration: underline;
}

nav {
	text-align: center;
	margin-bottom: 20px;
}
nav > a:hover {
	background-color: #E6E6E6;
	border-color: #ADADAD;
}
section {
	background-color: #FFF;
	margin-bottom: 20px;
}
section > h1, nav > a {
	border: 1px solid #DDD;
	border-radius: 3px;
	padding: 10px;
	margin: 0;
	font-size: 16px;
	font-weight: bold;
	text-align: center;
}
section > h1 {
	background-color: #EEE;
}
section > ul, p {
	padding: 0;
	list-style: "-" outside;
}
section > ul > li {
        margin-bottom: 5px;
}
section > ul > li > a:visited {
	color: #AAA;
}
footer {
	text-align: center;
	font-size: 12px;
}
`

const tmplPage string = `
<!doctype html>

	<meta charset="utf-8">
	<meta name="viewport" content="width=device-width">
	<title>{{.CurrentDate}} | Feedloggr</title>
	<link href="./style.css" rel="stylesheet" type="text/css">

		<
		Latest
		>

	{{range .Feeds}}

{{.Title}}

		{{if .Error}}

Error while updating feed:
{{.Error}}

		{{else}}

{{range .Items}}

{{.Title}}

		{{end}}

		{{end}}

	{{else}}

Sorry, no news for today!

	{{end}}

		Generated with Feedloggr

`
<|endoftext|>


<commit_before>package tera

import (
	"fmt"
	"testing"
	"time"
)

func logExecTime(start time.Time, prefix string) {
	elapsed_ms := time.Since(start) / time.Millisecond
	fmt.Printf("Performance: %s cost %d ms.\n", prefix, elapsed_ms)
}

func TestTera(*testing.T) {
	fmt.Println("Hello terago!")
	start := time.Now()
	client, c_err := NewClient("./tera.flag", "terago")
	defer client.Close()
	if c_err != nil {
		panic("tera.NewClient error: " + c_err.Error())
	}
	logExecTime(start, "NewClient")

	start = time.Now()
	table, t_err := client.OpenTable("terago")
	defer table.Close()
	if t_err != nil {
		panic("tera.OpenTable error: " + t_err.Error())
	}
	logExecTime(start, "OpenTable")

	{
		defer logExecTime(time.Now(), "PutKV")
		p_err := table.PutKV("hello", "terago", 10)
		if p_err != nil {
			panic("put key value error: " + p_err.Error())
		}
	}

	{
		defer logExecTime(time.Now(), "GetKV")
		// get an exist key value, return value
		value, g_err := table.GetKV("hello")
		if g_err != nil {
			panic("get key value error: " + g_err.Error())
		}
		fmt.Printf("get key[%s] value[%s].\n", "hello", value)
	}

	{
		defer logExecTime(time.Now(), "GetKV_NotExist")
		// get a not-exist key value, return "not found"
		_, g_err := table.GetKV("hell")
		if g_err == nil {
			panic("get key value should fail: " + g_err.Error())
		}
	}

	{
		defer logExecTime(time.Now(), "DeleteKV")
		d_err := table.DeleteKV("hello")
		if d_err != nil {
			panic("delete key value error: " + d_err.Error())
		}
	}

	_, g_err := table.GetKV("hello")
	if g_err == nil {
		panic("get key value should fail: " + g_err.Error())
	}
}
<commit_msg>modify unit test<commit_after>package tera

import (
	"fmt"
	"testing"
	"time"
)

func logExecTime(start time.Time, prefix string) {
	elapsed_ms := time.Since(start) / time.Millisecond
	fmt.Printf("Performance: %s cost %d ms.\n", prefix, elapsed_ms)
}

func TestTera(*testing.T) {
	fmt.Println("Hello terago!")
	start := time.Now()
	client, c_err := NewClient("tera.flag", "terago")
	defer client.Close()
	if c_err != nil {
		panic("tera.NewClient error: " + c_err.Error())
	}
	logExecTime(start, "NewClient")

	start = time.Now()
	table, t_err := client.OpenTable("terago")
	defer table.Close()
	if t_err != nil {
		panic("tera.OpenTable error: " + t_err.Error())
	}
	logExecTime(start, "OpenTable")

	start = time.Now()
	p_err := table.PutKV("hello", "terago", 10)
	if p_err != nil {
		panic("put key value error: " + p_err.Error())
	}
	logExecTime(start, "PutKV")

	start = time.Now()
	p_err = table.PutKVAsync("helloasync", "terago", 10)
	if p_err != nil {
		panic("put key value async error: " + p_err.Error())
	}
	logExecTime(start, "PutKVAsync")

	start = time.Now()
	// get an exist key value, return value
	value, g_err := table.GetKV("hello")
	if g_err != nil {
		panic("get key value error: " + g_err.Error())
	}
	fmt.Printf("get key[%s] value[%s].\n", "hello", value)
	logExecTime(start, "GetKV")

	start = time.Now()
	// get a not-exist key value, return "not found"
	_, g_err = table.GetKV("hell")
	if g_err == nil {
		panic("get key value should fail: " + g_err.Error())
	}
	logExecTime(start, "GetKV_NotExist")

	start = time.Now()
	d_err := table.DeleteKV("hello")
	if d_err != nil {
		panic("delete key value error: " + d_err.Error())
	}
	logExecTime(start, "DeleteKV")

	_, g_err = table.GetKV("hello")
	if g_err == nil {
		panic("get key value should fail: " + g_err.Error())
	}
}
<|endoftext|>


<commit_before>package sqs_test

import (
	"errors"
	"sync"
	"testing"

	"github.com/RichardKnop/machinery/v1"
	"github.com/RichardKnop/machinery/v1/brokers/iface"
	"github.com/RichardKnop/machinery/v1/brokers/sqs"
	"github.com/RichardKnop/machinery/v1/config"
	"github.com/RichardKnop/machinery/v1/retry"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/stretchr/testify/assert"

	awssqs "github.com/aws/aws-sdk-go/service/sqs"
)

var (
	testBroker           iface.Broker
	testAWSSQSBroker     *sqs.Broker
	errAWSSQSBroker      *sqs.Broker
	cnf                  *config.Config
	receiveMessageOutput *awssqs.ReceiveMessageOutput
)

func init() {
	testAWSSQSBroker = sqs.TestAWSSQSBroker
	errAWSSQSBroker = sqs.ErrAWSSQSBroker
	cnf = sqs.TestConf
	receiveMessageOutput = sqs.ReceiveMessageOutput
	testBroker = sqs.New(cnf)
}

func TestNewAWSSQSBroker(t *testing.T) {
	assert.IsType(t, testAWSSQSBroker, testBroker)
}

func TestPrivateFunc_continueReceivingMessages(t *testing.T) {
	qURL := testAWSSQSBroker.DefaultQueueURLForTest()
	deliveries := make(chan *awssqs.ReceiveMessageOutput)
	firstStep := make(chan int)
	nextStep := make(chan int)
	go func() {
		stopReceivingChan := testAWSSQSBroker.GetStopReceivingChanForTest()
		firstStep <- 1
		stopReceivingChan <- 1
	}()

	var (
		whetherContinue bool
		err             error
	)
	<-firstStep
	// Test the case that a signal was received from stopReceivingChan
	go func() {
		whetherContinue, err = testAWSSQSBroker.ContinueReceivingMessagesForTest(qURL, deliveries)
		nextStep <- 1
	}()
	<-nextStep
	assert.False(t, whetherContinue)
	assert.Nil(t, err)

	// Test the default condition
	whetherContinue, err = testAWSSQSBroker.ContinueReceivingMessagesForTest(qURL, deliveries)
	assert.True(t, whetherContinue)
	assert.Nil(t, err)

	// Test the error
	whetherContinue, err = errAWSSQSBroker.ContinueReceivingMessagesForTest(qURL, deliveries)
	assert.True(t, whetherContinue)
	assert.NotNil(t, err)

	// Test when there is no message
	outputCopy := *receiveMessageOutput
	receiveMessageOutput.Messages = []*awssqs.Message{}
	whetherContinue, err = testAWSSQSBroker.ContinueReceivingMessagesForTest(qURL, deliveries)
	assert.True(t, whetherContinue)
	assert.Nil(t, err)
	// recover original value
	*receiveMessageOutput = outputCopy

}

func TestPrivateFunc_consume(t *testing.T) {
	server1, err := machinery.NewServer(cnf)
	if err != nil {
		t.Fatal(err)
	}
	pool := make(chan struct{}, 0)
	wk := server1.NewWorker("sms_worker", 0)
	deliveries := make(chan *awssqs.ReceiveMessageOutput)
	outputCopy := *receiveMessageOutput
	outputCopy.Messages = []*awssqs.Message{}
	go func() { deliveries <- &outputCopy }()

	// an infinite loop will be executed only when there is no error
	err = testAWSSQSBroker.ConsumeForTest(deliveries, 0, wk, pool)
	assert.NotNil(t, err)

}

func TestPrivateFunc_consumeOne(t *testing.T) {
	server1, err := machinery.NewServer(cnf)
	if err != nil {
		t.Fatal(err)
	}
	wk := server1.NewWorker("sms_worker", 0)
	err = testAWSSQSBroker.ConsumeOneForTest(receiveMessageOutput, wk)
	assert.NotNil(t, err)

	outputCopy := *receiveMessageOutput
	outputCopy.Messages = []*awssqs.Message{}
	err = testAWSSQSBroker.ConsumeOneForTest(&outputCopy, wk)
	assert.NotNil(t, err)

	outputCopy.Messages = []*awssqs.Message{
		{
			Body: aws.String("foo message"),
		},
	}
	err = testAWSSQSBroker.ConsumeOneForTest(&outputCopy, wk)
	assert.NotNil(t, err)
}

func TestPrivateFunc_initializePool(t *testing.T) {
	concurrency := 9
	pool := make(chan struct{}, concurrency)
	testAWSSQSBroker.InitializePoolForTest(pool, concurrency)
	assert.Len(t, pool, concurrency)
}

func TestPrivateFunc_startConsuming(t *testing.T) {
	server1, err := machinery.NewServer(cnf)
	if err != nil {
		t.Fatal(err)
	}

	wk := server1.NewWorker("sms_worker", 0)
	retryFunc := testAWSSQSBroker.GetRetryFuncForTest()
	stopChan := testAWSSQSBroker.GetStopChanForTest()
	retryStopChan := testAWSSQSBroker.GetRetryStopChanForTest()
	assert.Nil(t, retryFunc)
	testAWSSQSBroker.StartConsumingForTest("fooTag", 1, wk)
	assert.IsType(t, retryFunc, retry.Closure())
	assert.Equal(t, len(stopChan), 0)
	assert.Equal(t, len(retryStopChan), 0)
}

func TestPrivateFuncDefaultQueueURL(t *testing.T) {
	qURL := testAWSSQSBroker.DefaultQueueURLForTest()

	assert.EqualValues(t, *qURL, "https://sqs.foo.amazonaws.com.cn/test_queue")
}

func TestPrivateFunc_stopReceiving(t *testing.T) {
	go testAWSSQSBroker.StopReceivingForTest()
	stopReceivingChan := testAWSSQSBroker.GetStopReceivingChanForTest()
	assert.NotNil(t, <-stopReceivingChan)
}

func TestPrivateFunc_receiveMessage(t *testing.T) {
	qURL := testAWSSQSBroker.DefaultQueueURLForTest()
	output, err := testAWSSQSBroker.ReceiveMessageForTest(qURL)
	assert.Nil(t, err)
	assert.Equal(t, receiveMessageOutput, output)
}

func TestPrivateFunc_consumeDeliveries(t *testing.T) {
	concurrency := 0
	pool := make(chan struct{}, concurrency)
	errorsChan := make(chan error)
	deliveries := make(chan *awssqs.ReceiveMessageOutput)
	server1, err := machinery.NewServer(cnf)
	if err != nil {
		t.Fatal(err)
	}
	wk := server1.NewWorker("sms_worker", 0)
	go func() { deliveries <- receiveMessageOutput }()
	whetherContinue, err := testAWSSQSBroker.ConsumeDeliveriesForTest(deliveries, concurrency, wk, pool, errorsChan)
	assert.True(t, whetherContinue)
	assert.Nil(t, err)

	go func() { errorsChan <- errors.New("foo error") }()
	whetherContinue, err = testAWSSQSBroker.ConsumeDeliveriesForTest(deliveries, concurrency, wk, pool, errorsChan)
	assert.False(t, whetherContinue)
	assert.NotNil(t, err)

	go func() { testAWSSQSBroker.GetStopChanForTest() <- 1 }()
	whetherContinue, err = testAWSSQSBroker.ConsumeDeliveriesForTest(deliveries, concurrency, wk, pool, errorsChan)
	assert.False(t, whetherContinue)
	assert.Nil(t, err)

	outputCopy := *receiveMessageOutput
	outputCopy.Messages = []*awssqs.Message{}
	go func() { deliveries <- &outputCopy }()
	whetherContinue, err = testAWSSQSBroker.ConsumeDeliveriesForTest(deliveries, concurrency, wk, pool, errorsChan)
	e := <-errorsChan
	assert.True(t, whetherContinue)
	assert.NotNil(t, e)
	assert.Nil(t, err)

	// using a wait group and a channel to fix the racing problem
	var wg sync.WaitGroup
	wg.Add(1)
	nextStep := make(chan bool, 1)
	go func() {
		defer wg.Done()
		// nextStep <- true runs after defer wg.Done(), to make sure the next go routine runs after this go routine
		nextStep <- true
		deliveries <- receiveMessageOutput
	}()
	if <-nextStep {
		// <-pool will block the routine in the following steps, so pool <- struct{}{} will be executed for sure
		go func() { wg.Wait(); pool <- struct{}{} }()
	}
	whetherContinue, err = testAWSSQSBroker.ConsumeDeliveriesForTest(deliveries, concurrency, wk, pool, errorsChan)
	// the pool shouldn't be consumed
	p := <-pool
	assert.True(t, whetherContinue)
	assert.NotNil(t, p)
	assert.Nil(t, err)
}

func TestPrivateFunc_deleteOne(t *testing.T) {
	err := testAWSSQSBroker.DeleteOneForTest(receiveMessageOutput)
	assert.Nil(t, err)

	err = errAWSSQSBroker.DeleteOneForTest(receiveMessageOutput)
	assert.NotNil(t, err)
}

func Test_CustomQueueName(t *testing.T) {
	server1, err := machinery.NewServer(cnf)
	if err != nil {
		t.Fatal(err)
	}

	wk := server1.NewWorker("test-worker", 0)
	qURL := testAWSSQSBroker.GetQueueURLForTest(wk)
	assert.Equal(t, qURL, testAWSSQSBroker.DefaultQueueURLForTest(), "")

	wk2 := server1.NewCustomQueueWorker("test-worker", 0, "my-custom-queue")
	qURL2 := testAWSSQSBroker.GetQueueURLForTest(wk2)
	assert.Equal(t, qURL2, testAWSSQSBroker.GetCustomQueueURL("my-custom-queue"), "")
}
<commit_msg>Add unit test to improve coverage<commit_after>package sqs_test

import (
	"context"
	"errors"
	"sync"
	"testing"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/stretchr/testify/assert"

	"github.com/RichardKnop/machinery/v1"
	"github.com/RichardKnop/machinery/v1/brokers/iface"
	"github.com/RichardKnop/machinery/v1/brokers/sqs"
	"github.com/RichardKnop/machinery/v1/config"
	"github.com/RichardKnop/machinery/v1/retry"

	awssqs "github.com/aws/aws-sdk-go/service/sqs"
)

var (
	testBroker           iface.Broker
	testAWSSQSBroker     *sqs.Broker
	errAWSSQSBroker      *sqs.Broker
	cnf                  *config.Config
	receiveMessageOutput *awssqs.ReceiveMessageOutput
)

func init() {
	testAWSSQSBroker = sqs.TestAWSSQSBroker
	errAWSSQSBroker = sqs.ErrAWSSQSBroker
	cnf = sqs.TestConf
	receiveMessageOutput = sqs.ReceiveMessageOutput
	testBroker = sqs.New(cnf)
}

func TestNewAWSSQSBroker(t *testing.T) {
	assert.IsType(t, testAWSSQSBroker, testBroker)
}

func TestPrivateFunc_continueReceivingMessages(t *testing.T) {
	qURL := testAWSSQSBroker.DefaultQueueURLForTest()
	deliveries := make(chan *awssqs.ReceiveMessageOutput)
	firstStep := make(chan int)
	nextStep := make(chan int)
	go func() {
		stopReceivingChan := testAWSSQSBroker.GetStopReceivingChanForTest()
		firstStep <- 1
		stopReceivingChan <- 1
	}()

	var (
		whetherContinue bool
		err             error
	)
	<-firstStep
	// Test the case that a signal was received from stopReceivingChan
	go func() {
		whetherContinue, err = testAWSSQSBroker.ContinueReceivingMessagesForTest(qURL, deliveries)
		nextStep <- 1
	}()
	<-nextStep
	assert.False(t, whetherContinue)
	assert.Nil(t, err)

	// Test the default condition
	whetherContinue, err = testAWSSQSBroker.ContinueReceivingMessagesForTest(qURL, deliveries)
	assert.True(t, whetherContinue)
	assert.Nil(t, err)

	// Test the error
	whetherContinue, err = errAWSSQSBroker.ContinueReceivingMessagesForTest(qURL, deliveries)
	assert.True(t, whetherContinue)
	assert.NotNil(t, err)

	// Test when there is no message
	outputCopy := *receiveMessageOutput
	receiveMessageOutput.Messages = []*awssqs.Message{}
	whetherContinue, err = testAWSSQSBroker.ContinueReceivingMessagesForTest(qURL, deliveries)
	assert.True(t, whetherContinue)
	assert.Nil(t, err)
	// recover original value
	*receiveMessageOutput = outputCopy

}

func TestPrivateFunc_consume(t *testing.T) {
	server1, err := machinery.NewServer(cnf)
	if err != nil {
		t.Fatal(err)
	}
	pool := make(chan struct{}, 0)
	wk := server1.NewWorker("sms_worker", 0)
	deliveries := make(chan *awssqs.ReceiveMessageOutput)
	outputCopy := *receiveMessageOutput
	outputCopy.Messages = []*awssqs.Message{}
	go func() { deliveries <- &outputCopy }()

	// an infinite loop will be executed only when there is no error
	err = testAWSSQSBroker.ConsumeForTest(deliveries, 0, wk, pool)
	assert.NotNil(t, err)

}

func TestPrivateFunc_consumeOne(t *testing.T) {
	server1, err := machinery.NewServer(cnf)
	if err != nil {
		t.Fatal(err)
	}
	wk := server1.NewWorker("sms_worker", 0)
	err = testAWSSQSBroker.ConsumeOneForTest(receiveMessageOutput, wk)
	assert.NotNil(t, err)

	outputCopy := *receiveMessageOutput
	outputCopy.Messages = []*awssqs.Message{}
	err = testAWSSQSBroker.ConsumeOneForTest(&outputCopy, wk)
	assert.NotNil(t, err)

	outputCopy.Messages = []*awssqs.Message{
		{
			Body: aws.String("foo message"),
		},
	}
	err = testAWSSQSBroker.ConsumeOneForTest(&outputCopy, wk)
	assert.NotNil(t, err)
}

func TestPrivateFunc_initializePool(t *testing.T) {
	concurrency := 9
	pool := make(chan struct{}, concurrency)
	testAWSSQSBroker.InitializePoolForTest(pool, concurrency)
	assert.Len(t, pool, concurrency)
}

func TestPrivateFunc_startConsuming(t *testing.T) {
	server1, err := machinery.NewServer(cnf)
	if err != nil {
		t.Fatal(err)
	}

	wk := server1.NewWorker("sms_worker", 0)
	retryFunc := testAWSSQSBroker.GetRetryFuncForTest()
	stopChan := testAWSSQSBroker.GetStopChanForTest()
	retryStopChan := testAWSSQSBroker.GetRetryStopChanForTest()
	assert.Nil(t, retryFunc)
	testAWSSQSBroker.StartConsumingForTest("fooTag", 1, wk)
	assert.IsType(t, retryFunc, retry.Closure())
	assert.Equal(t, len(stopChan), 0)
	assert.Equal(t, len(retryStopChan), 0)
}

func TestPrivateFuncDefaultQueueURL(t *testing.T) {
	qURL := testAWSSQSBroker.DefaultQueueURLForTest()

	assert.EqualValues(t, *qURL, "https://sqs.foo.amazonaws.com.cn/test_queue")
}

func TestPrivateFunc_stopReceiving(t *testing.T) {
	go testAWSSQSBroker.StopReceivingForTest()
	stopReceivingChan := testAWSSQSBroker.GetStopReceivingChanForTest()
	assert.NotNil(t, <-stopReceivingChan)
}

func TestPrivateFunc_receiveMessage(t *testing.T) {
	qURL := testAWSSQSBroker.DefaultQueueURLForTest()
	output, err := testAWSSQSBroker.ReceiveMessageForTest(qURL)
	assert.Nil(t, err)
	assert.Equal(t, receiveMessageOutput, output)
}

func TestPrivateFunc_consumeDeliveries(t *testing.T) {
	concurrency := 0
	pool := make(chan struct{}, concurrency)
	errorsChan := make(chan error)
	deliveries := make(chan *awssqs.ReceiveMessageOutput)
	server1, err := machinery.NewServer(cnf)
	if err != nil {
		t.Fatal(err)
	}
	wk := server1.NewWorker("sms_worker", 0)
	go func() { deliveries <- receiveMessageOutput }()
	whetherContinue, err := testAWSSQSBroker.ConsumeDeliveriesForTest(deliveries, concurrency, wk, pool, errorsChan)
	assert.True(t, whetherContinue)
	assert.Nil(t, err)

	go func() { errorsChan <- errors.New("foo error") }()
	whetherContinue, err = testAWSSQSBroker.ConsumeDeliveriesForTest(deliveries, concurrency, wk, pool, errorsChan)
	assert.False(t, whetherContinue)
	assert.NotNil(t, err)

	go func() { testAWSSQSBroker.GetStopChanForTest() <- 1 }()
	whetherContinue, err = testAWSSQSBroker.ConsumeDeliveriesForTest(deliveries, concurrency, wk, pool, errorsChan)
	assert.False(t, whetherContinue)
	assert.Nil(t, err)

	outputCopy := *receiveMessageOutput
	outputCopy.Messages = []*awssqs.Message{}
	go func() { deliveries <- &outputCopy }()
	whetherContinue, err = testAWSSQSBroker.ConsumeDeliveriesForTest(deliveries, concurrency, wk, pool, errorsChan)
	e := <-errorsChan
	assert.True(t, whetherContinue)
	assert.NotNil(t, e)
	assert.Nil(t, err)

	// using a wait group and a channel to fix the racing problem
	var wg sync.WaitGroup
	wg.Add(1)
	nextStep := make(chan bool, 1)
	go func() {
		defer wg.Done()
		// nextStep <- true runs after defer wg.Done(), to make sure the next go routine runs after this go routine
		nextStep <- true
		deliveries <- receiveMessageOutput
	}()
	if <-nextStep {
		// <-pool will block the routine in the following steps, so pool <- struct{}{} will be executed for sure
		go func() { wg.Wait(); pool <- struct{}{} }()
	}
	whetherContinue, err = testAWSSQSBroker.ConsumeDeliveriesForTest(deliveries, concurrency, wk, pool, errorsChan)
	// the pool shouldn't be consumed
	p := <-pool
	assert.True(t, whetherContinue)
	assert.NotNil(t, p)
	assert.Nil(t, err)
}

func TestPrivateFunc_deleteOne(t *testing.T) {
	err := testAWSSQSBroker.DeleteOneForTest(receiveMessageOutput)
	assert.Nil(t, err)

	err = errAWSSQSBroker.DeleteOneForTest(receiveMessageOutput)
	assert.NotNil(t, err)
}

func Test_CustomQueueName(t *testing.T) {
	server1, err := machinery.NewServer(cnf)
	if err != nil {
		t.Fatal(err)
	}

	wk := server1.NewWorker("test-worker", 0)
	qURL := testAWSSQSBroker.GetQueueURLForTest(wk)
	assert.Equal(t, qURL, testAWSSQSBroker.DefaultQueueURLForTest(), "")

	wk2 := server1.NewCustomQueueWorker("test-worker", 0, "my-custom-queue")
	qURL2 := testAWSSQSBroker.GetQueueURLForTest(wk2)
	assert.Equal(t, qURL2, testAWSSQSBroker.GetCustomQueueURL("my-custom-queue"), "")
}

func TestPrivateFunc_consumeWithConcurrency(t *testing.T) {

	msg := `{
        "UUID": "uuid-dummy-task",
        "Name": "test-task",
        "RoutingKey": "dummy-routing"
	}
	`

	testResp := "47f8b355-5115-4b45-b33a-439016400411"
	output := make(chan string) // The output channel

	cnf.ResultBackend = "eager"
	server1, err := machinery.NewServer(cnf)
	if err != nil {
		t.Fatal(err)
	}
	err = server1.RegisterTask("test-task", func(ctx context.Context) error {
		output <- testResp
		testAWSSQSBroker.StopConsuming()

		return nil
	})
	testAWSSQSBroker.SetRegisteredTaskNames([]string{"test-task"})
	assert.NoError(t, err)
	pool := make(chan struct{}, 1)
	pool <- struct{}{}
	wk := server1.NewWorker("sms_worker", 1)
	deliveries := make(chan *awssqs.ReceiveMessageOutput)
	outputCopy := *receiveMessageOutput
	outputCopy.Messages = []*awssqs.Message{}
	outputCopy.Messages = append(outputCopy.Messages, &awssqs.Message{MessageId: aws.String("test-sqs-msg1"), Body: aws.String(msg)})

	go func() {
		deliveries <- &outputCopy

	}()

	go func() {
		testAWSSQSBroker.StartConsumingForTest("test", 1, wk)
		err = testAWSSQSBroker.ConsumeForTest(deliveries, 1, wk, pool)
	}()

	select {
	case resp := <-output:
		assert.Equal(t, testResp, resp)
	case <-time.After(2 * time.Second):
		// call timed out
		t.Fatal("task not processed in 2 seconds")
	}
	//assert.Equal(t, testResp, resp)

}
<|endoftext|>


<commit_before>package autonat

import (
	"context"
	"sync"
	"time"

	pb "github.com/libp2p/go-libp2p-autonat/pb"

	ggio "github.com/gogo/protobuf/io"
	libp2p "github.com/libp2p/go-libp2p"
	host "github.com/libp2p/go-libp2p-host"
	inet "github.com/libp2p/go-libp2p-net"
	peer "github.com/libp2p/go-libp2p-peer"
	pstore "github.com/libp2p/go-libp2p-peerstore"
	ma "github.com/multiformats/go-multiaddr"
)

const P_CIRCUIT = 290

var (
	AutoNATServiceDialTimeout   = 42 * time.Second
	AutoNATServiceResetInterval = 1 * time.Minute

	AutoNATServiceThrottle = 3
)

// AutoNATService provides NAT autodetection services to other peers
type AutoNATService struct {
	ctx    context.Context
	dialer host.Host

	// rate limiter
	mx    sync.Mutex
	peers map[peer.ID]int
}

// NewAutoNATService creates a new AutoNATService instance attached to a host
func NewAutoNATService(ctx context.Context, h host.Host, opts ...libp2p.Option) (*AutoNATService, error) {
	opts = append(opts, libp2p.NoListenAddrs)
	dialer, err := libp2p.New(ctx, opts...)
	if err != nil {
		return nil, err
	}

	as := &AutoNATService{
		ctx:    ctx,
		dialer: dialer,
		peers:  make(map[peer.ID]int),
	}
	h.SetStreamHandler(AutoNATProto, as.handleStream)

	go as.resetPeers()

	return as, nil
}

func (as *AutoNATService) handleStream(s inet.Stream) {
	defer s.Close()

	pid := s.Conn().RemotePeer()
	log.Debugf("New stream from %s", pid.Pretty())

	r := ggio.NewDelimitedReader(s, inet.MessageSizeMax)
	w := ggio.NewDelimitedWriter(s)

	var req pb.Message
	var res pb.Message

	err := r.ReadMsg(&req)
	if err != nil {
		log.Debugf("Error reading message from %s: %s", pid.Pretty(), err.Error())
		s.Reset()
		return
	}

	t := req.GetType()
	if t != pb.Message_DIAL {
		log.Debugf("Unexpected message from %s: %s (%d)", pid.Pretty(), t.String(), t)
		s.Reset()
		return
	}

	dr := as.handleDial(pid, s.Conn().RemoteMultiaddr(), req.GetDial().GetPeer())
	res.Type = pb.Message_DIAL_RESPONSE.Enum()
	res.DialResponse = dr

	err = w.WriteMsg(&res)
	if err != nil {
		log.Debugf("Error writing response to %s: %s", pid.Pretty(), err.Error())
		s.Reset()
		return
	}
}

func (as *AutoNATService) handleDial(p peer.ID, obsaddr ma.Multiaddr, mpi *pb.Message_PeerInfo) *pb.Message_DialResponse {
	if mpi == nil {
		return newDialResponseError(pb.Message_E_BAD_REQUEST, "missing peer info")
	}

	mpid := mpi.GetId()
	if mpid != nil {
		mp, err := peer.IDFromBytes(mpid)
		if err != nil {
			return newDialResponseError(pb.Message_E_BAD_REQUEST, "bad peer id")
		}

		if mp != p {
			return newDialResponseError(pb.Message_E_BAD_REQUEST, "peer id mismatch")
		}
	}

	addrs := make([]ma.Multiaddr, 0)
	seen := make(map[string]struct{})

	// add observed addr to the list of addresses to dial
	if !as.skipDial(obsaddr) {
		addrs = append(addrs, obsaddr)
		seen[obsaddr.String()] = struct{}{}
	}

	for _, maddr := range mpi.GetAddrs() {
		addr, err := ma.NewMultiaddrBytes(maddr)
		if err != nil {
			log.Debugf("Error parsing multiaddr: %s", err.Error())
			continue
		}

		if as.skipDial(addr) {
			continue
		}

		str := addr.String()
		_, ok := seen[str]
		if ok {
			continue
		}

		addrs = append(addrs, addr)
		seen[str] = struct{}{}
	}

	if len(addrs) == 0 {
		return newDialResponseError(pb.Message_E_DIAL_ERROR, "no dialable addresses")
	}

	return as.doDial(pstore.PeerInfo{ID: p, Addrs: addrs})
}

func (as *AutoNATService) skipDial(addr ma.Multiaddr) bool {
	// skip relay addresses
	_, err := addr.ValueForProtocol(P_CIRCUIT)
	if err == nil {
		return true
	}

	// skip private network (unroutable) addresses
	if !isPublicAddr(addr) {
		return true
	}

	return false
}

func (as *AutoNATService) doDial(pi pstore.PeerInfo) *pb.Message_DialResponse {
	// rate limit check
	as.mx.Lock()
	count := as.peers[pi.ID]
	if count >= AutoNATServiceThrottle {
		as.mx.Unlock()
		return newDialResponseError(pb.Message_E_DIAL_REFUSED, "too many dials")
	}
	as.peers[pi.ID] = count + 1
	as.mx.Unlock()

	ctx, cancel := context.WithTimeout(as.ctx, AutoNATServiceDialTimeout)
	defer cancel()

	err := as.dialer.Connect(ctx, pi)
	if err != nil {
		log.Debugf("error dialing %s: %s", pi.ID.Pretty(), err.Error())
		// wait for the context to timeout to avoid leaking timing information
		// this renders the service ineffective as a port scanner
		<-ctx.Done()
		return newDialResponseError(pb.Message_E_DIAL_ERROR, "dial failed")
	}

	conns := as.dialer.Network().ConnsToPeer(pi.ID)
	if len(conns) == 0 {
		log.Errorf("supposedly connected to %s, but no connection to peer", pi.ID.Pretty())
		return newDialResponseError(pb.Message_E_INTERNAL_ERROR, "internal service error")
	}

	ra := conns[0].RemoteMultiaddr()
	as.dialer.Network().ClosePeer(pi.ID)
	return newDialResponseOK(ra)
}

func (as *AutoNATService) resetPeers() {
	ticker := time.NewTicker(AutoNATServiceResetInterval)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			as.mx.Lock()
			as.peers = make(map[peer.ID]int)
			as.mx.Unlock()

		case <-as.ctx.Done():
			return
		}
	}
}
<commit_msg>call AutoNATService.peers something else (reqs)<commit_after>package autonat

import (
	"context"
	"sync"
	"time"

	pb "github.com/libp2p/go-libp2p-autonat/pb"

	ggio "github.com/gogo/protobuf/io"
	libp2p "github.com/libp2p/go-libp2p"
	host "github.com/libp2p/go-libp2p-host"
	inet "github.com/libp2p/go-libp2p-net"
	peer "github.com/libp2p/go-libp2p-peer"
	pstore "github.com/libp2p/go-libp2p-peerstore"
	ma "github.com/multiformats/go-multiaddr"
)

const P_CIRCUIT = 290

var (
	AutoNATServiceDialTimeout   = 42 * time.Second
	AutoNATServiceResetInterval = 1 * time.Minute

	AutoNATServiceThrottle = 3
)

// AutoNATService provides NAT autodetection services to other peers
type AutoNATService struct {
	ctx    context.Context
	dialer host.Host

	// rate limiter
	mx   sync.Mutex
	reqs map[peer.ID]int
}

// NewAutoNATService creates a new AutoNATService instance attached to a host
func NewAutoNATService(ctx context.Context, h host.Host, opts ...libp2p.Option) (*AutoNATService, error) {
	opts = append(opts, libp2p.NoListenAddrs)
	dialer, err := libp2p.New(ctx, opts...)
	if err != nil {
		return nil, err
	}

	as := &AutoNATService{
		ctx:    ctx,
		dialer: dialer,
		reqs:   make(map[peer.ID]int),
	}
	h.SetStreamHandler(AutoNATProto, as.handleStream)

	go as.resetRateLimiter()

	return as, nil
}

func (as *AutoNATService) handleStream(s inet.Stream) {
	defer s.Close()

	pid := s.Conn().RemotePeer()
	log.Debugf("New stream from %s", pid.Pretty())

	r := ggio.NewDelimitedReader(s, inet.MessageSizeMax)
	w := ggio.NewDelimitedWriter(s)

	var req pb.Message
	var res pb.Message

	err := r.ReadMsg(&req)
	if err != nil {
		log.Debugf("Error reading message from %s: %s", pid.Pretty(), err.Error())
		s.Reset()
		return
	}

	t := req.GetType()
	if t != pb.Message_DIAL {
		log.Debugf("Unexpected message from %s: %s (%d)", pid.Pretty(), t.String(), t)
		s.Reset()
		return
	}

	dr := as.handleDial(pid, s.Conn().RemoteMultiaddr(), req.GetDial().GetPeer())
	res.Type = pb.Message_DIAL_RESPONSE.Enum()
	res.DialResponse = dr

	err = w.WriteMsg(&res)
	if err != nil {
		log.Debugf("Error writing response to %s: %s", pid.Pretty(), err.Error())
		s.Reset()
		return
	}
}

func (as *AutoNATService) handleDial(p peer.ID, obsaddr ma.Multiaddr, mpi *pb.Message_PeerInfo) *pb.Message_DialResponse {
	if mpi == nil {
		return newDialResponseError(pb.Message_E_BAD_REQUEST, "missing peer info")
	}

	mpid := mpi.GetId()
	if mpid != nil {
		mp, err := peer.IDFromBytes(mpid)
		if err != nil {
			return newDialResponseError(pb.Message_E_BAD_REQUEST, "bad peer id")
		}

		if mp != p {
			return newDialResponseError(pb.Message_E_BAD_REQUEST, "peer id mismatch")
		}
	}

	addrs := make([]ma.Multiaddr, 0)
	seen := make(map[string]struct{})

	// add observed addr to the list of addresses to dial
	if !as.skipDial(obsaddr) {
		addrs = append(addrs, obsaddr)
		seen[obsaddr.String()] = struct{}{}
	}

	for _, maddr := range mpi.GetAddrs() {
		addr, err := ma.NewMultiaddrBytes(maddr)
		if err != nil {
			log.Debugf("Error parsing multiaddr: %s", err.Error())
			continue
		}

		if as.skipDial(addr) {
			continue
		}

		str := addr.String()
		_, ok := seen[str]
		if ok {
			continue
		}

		addrs = append(addrs, addr)
		seen[str] = struct{}{}
	}

	if len(addrs) == 0 {
		return newDialResponseError(pb.Message_E_DIAL_ERROR, "no dialable addresses")
	}

	return as.doDial(pstore.PeerInfo{ID: p, Addrs: addrs})
}

func (as *AutoNATService) skipDial(addr ma.Multiaddr) bool {
	// skip relay addresses
	_, err := addr.ValueForProtocol(P_CIRCUIT)
	if err == nil {
		return true
	}

	// skip private network (unroutable) addresses
	if !isPublicAddr(addr) {
		return true
	}

	return false
}

func (as *AutoNATService) doDial(pi pstore.PeerInfo) *pb.Message_DialResponse {
	// rate limit check
	as.mx.Lock()
	count := as.reqs[pi.ID]
	if count >= AutoNATServiceThrottle {
		as.mx.Unlock()
		return newDialResponseError(pb.Message_E_DIAL_REFUSED, "too many dials")
	}
	as.reqs[pi.ID] = count + 1
	as.mx.Unlock()

	ctx, cancel := context.WithTimeout(as.ctx, AutoNATServiceDialTimeout)
	defer cancel()

	err := as.dialer.Connect(ctx, pi)
	if err != nil {
		log.Debugf("error dialing %s: %s", pi.ID.Pretty(), err.Error())
		// wait for the context to timeout to avoid leaking timing information
		// this renders the service ineffective as a port scanner
		<-ctx.Done()
		return newDialResponseError(pb.Message_E_DIAL_ERROR, "dial failed")
	}

	conns := as.dialer.Network().ConnsToPeer(pi.ID)
	if len(conns) == 0 {
		log.Errorf("supposedly connected to %s, but no connection to peer", pi.ID.Pretty())
		return newDialResponseError(pb.Message_E_INTERNAL_ERROR, "internal service error")
	}

	ra := conns[0].RemoteMultiaddr()
	as.dialer.Network().ClosePeer(pi.ID)
	return newDialResponseOK(ra)
}

func (as *AutoNATService) resetRateLimiter() {
	ticker := time.NewTicker(AutoNATServiceResetInterval)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			as.mx.Lock()
			as.reqs = make(map[peer.ID]int)
			as.mx.Unlock()

		case <-as.ctx.Done():
			return
		}
	}
}
<|endoftext|>


<commit_before>package text

import (
	"image/color"
	"math"
	"unicode"
	"unicode/utf8"

	"github.com/faiface/pixel"
	"golang.org/x/image/font"
)

var ASCII []rune

func init() {
	ASCII = make([]rune, unicode.MaxASCII-32)
	for i := range ASCII {
		ASCII[i] = rune(32 + i)
	}
}

func RangeTable(table *unicode.RangeTable) []rune {
	var runes []rune
	for _, rng := range table.R16 {
		for r := rng.Lo; r <= rng.Hi; r += rng.Stride {
			runes = append(runes, rune(r))
		}
	}
	for _, rng := range table.R32 {
		for r := rng.Lo; r <= rng.Hi; r += rng.Stride {
			runes = append(runes, rune(r))
		}
	}
	return runes
}

type Text struct {
	Orig pixel.Vec
	Dot  pixel.Vec

	atlas *Atlas

	lineHeight float64
	tabWidth   float64

	prevR rune
	glyph pixel.TrianglesData
	tris  pixel.TrianglesData
	d     pixel.Drawer
	trans *pixel.Batch
}

func New(face font.Face, runeSets ...[]rune) *Text {
	runes := []rune{unicode.ReplacementChar}
	for _, set := range runeSets {
		runes = append(runes, set...)
	}

	atlas := NewAtlas(face, runes)

	txt := &Text{
		atlas:      atlas,
		lineHeight: atlas.LineHeight(),
		tabWidth:   atlas.Glyph(' ').Advance * 4,
	}

	txt.glyph.SetLen(6)
	for i := range txt.glyph {
		txt.glyph[i].Color = pixel.Alpha(1)
		txt.glyph[i].Intensity = 1
	}

	txt.d.Picture = txt.atlas.pic
	txt.d.Triangles = &txt.tris
	txt.trans = pixel.NewBatch(&pixel.TrianglesData{}, atlas.pic)

	txt.Clear()

	return txt
}

func (txt *Text) Atlas() *Atlas {
	return txt.atlas
}

func (txt *Text) SetMatrix(m pixel.Matrix) {
	txt.trans.SetMatrix(m)
}

func (txt *Text) SetColorMask(c color.Color) {
	txt.trans.SetColorMask(c)
}

func (txt *Text) Color(c color.Color) {
	rgba := pixel.ToRGBA(c)
	for i := range txt.glyph {
		txt.glyph[i].Color = rgba
	}
}

func (txt *Text) LineHeight(scale float64) {
	txt.lineHeight = scale
}

func (txt *Text) TabWidth(width float64) {
	txt.tabWidth = width
}

func (txt *Text) Clear() {
	txt.prevR = -1
	txt.tris.SetLen(0)
	txt.d.Dirty()
}

func (txt *Text) Write(p []byte) (n int, err error) {
	n, err = len(p), nil // always returns this

	if len(p) == 0 {
		return
	}

	for len(p) > 0 {
		r, size := utf8.DecodeRune(p)
		p = p[size:]
		txt.WriteRune(r)
	}

	return
}

func (txt *Text) WriteString(s string) (n int, err error) {
	if len(s) == 0 {
		return
	}

	for _, r := range s {
		txt.WriteRune(r)
	}

	return len(s), nil
}

func (txt *Text) WriteByte(c byte) error {
	_, err := txt.WriteRune(rune(c))
	return err
}

func (txt *Text) WriteRune(r rune) (n int, err error) {
	n, err = utf8.RuneLen(r), nil // always returns this

	switch r {
	case '\n':
		txt.Dot -= pixel.Y(txt.lineHeight)
		txt.Dot = txt.Dot.WithX(txt.Orig.X())
		return
	case '\r':
		txt.Dot = txt.Dot.WithX(txt.Orig.X())
		return
	case '\t':
		rem := math.Mod(txt.Dot.X()-txt.Orig.X(), txt.tabWidth)
		rem = math.Mod(rem, rem+txt.tabWidth)
		if rem == 0 {
			rem = txt.tabWidth
		}
		txt.Dot += pixel.X(rem)
		return
	}

	if !txt.atlas.Contains(r) {
		r = unicode.ReplacementChar
	}
	if !txt.atlas.Contains(unicode.ReplacementChar) {
		return
	}

	glyph := txt.atlas.Glyph(r)

	if txt.prevR >= 0 {
		txt.Dot += pixel.X(txt.atlas.Kern(txt.prevR, r))
	}

	a := pixel.V(glyph.Frame.Min.X(), glyph.Frame.Min.Y())
	b := pixel.V(glyph.Frame.Max.X(), glyph.Frame.Min.Y())
	c := pixel.V(glyph.Frame.Max.X(), glyph.Frame.Max.Y())
	d := pixel.V(glyph.Frame.Min.X(), glyph.Frame.Max.Y())

	for i, v := range []pixel.Vec{a, b, c, a, c, d} {
		txt.glyph[i].Position = v - glyph.Orig + txt.Dot
		txt.glyph[i].Picture = v
	}

	txt.tris = append(txt.tris, txt.glyph...)

	txt.Dot += pixel.X(glyph.Advance)
	txt.prevR = r

	txt.d.Dirty()

	return
}

func (txt *Text) Draw(t pixel.Target) {
	txt.trans.Clear()
	txt.d.Draw(txt.trans)
	txt.trans.Draw(t)
}
<commit_msg>remove Batch from Text and optimize it<commit_after>package text

import (
	"image/color"
	"math"
	"unicode"
	"unicode/utf8"

	"github.com/faiface/pixel"
	"golang.org/x/image/font"
)

var ASCII []rune

func init() {
	ASCII = make([]rune, unicode.MaxASCII-32)
	for i := range ASCII {
		ASCII[i] = rune(32 + i)
	}
}

func RangeTable(table *unicode.RangeTable) []rune {
	var runes []rune
	for _, rng := range table.R16 {
		for r := rng.Lo; r <= rng.Hi; r += rng.Stride {
			runes = append(runes, rune(r))
		}
	}
	for _, rng := range table.R32 {
		for r := rng.Lo; r <= rng.Hi; r += rng.Stride {
			runes = append(runes, rune(r))
		}
	}
	return runes
}

type Text struct {
	Orig pixel.Vec
	Dot  pixel.Vec

	atlas *Atlas

	lineHeight float64
	tabWidth   float64

	prevR  rune
	bounds pixel.Rect
	glyph  pixel.TrianglesData
	tris   pixel.TrianglesData

	mat    pixel.Matrix
	col    pixel.RGBA
	trans  pixel.TrianglesData
	transD pixel.Drawer
	dirty  bool
}

func New(face font.Face, runeSets ...[]rune) *Text {
	runes := []rune{unicode.ReplacementChar}
	for _, set := range runeSets {
		runes = append(runes, set...)
	}

	atlas := NewAtlas(face, runes)

	txt := &Text{
		atlas:      atlas,
		lineHeight: atlas.LineHeight(),
		tabWidth:   atlas.Glyph(' ').Advance * 4,
		mat:        pixel.IM,
		col:        pixel.Alpha(1),
	}

	txt.glyph.SetLen(6)
	for i := range txt.glyph {
		txt.glyph[i].Color = pixel.Alpha(1)
		txt.glyph[i].Intensity = 1
	}

	txt.transD.Picture = txt.atlas.pic
	txt.transD.Triangles = &txt.trans

	txt.Clear()

	return txt
}

func (txt *Text) Atlas() *Atlas {
	return txt.atlas
}

func (txt *Text) SetMatrix(m pixel.Matrix) {
	if txt.mat != m {
		txt.mat = m
		txt.dirty = true
	}
}

func (txt *Text) SetColorMask(c color.Color) {
	rgba := pixel.ToRGBA(c)
	if txt.col != rgba {
		txt.col = rgba
		txt.dirty = true
	}
}

func (txt *Text) Bounds() pixel.Rect {
	return txt.bounds
}

func (txt *Text) Color(c color.Color) {
	rgba := pixel.ToRGBA(c)
	for i := range txt.glyph {
		txt.glyph[i].Color = rgba
	}
}

func (txt *Text) LineHeight(scale float64) {
	txt.lineHeight = scale
}

func (txt *Text) TabWidth(width float64) {
	txt.tabWidth = width
}

func (txt *Text) Clear() {
	txt.prevR = -1
	txt.bounds = pixel.Rect{}
	txt.tris.SetLen(0)
	txt.dirty = true
}

func (txt *Text) Write(p []byte) (n int, err error) {
	n, err = len(p), nil // always returns this

	if len(p) == 0 {
		return
	}

	for len(p) > 0 {
		r, size := utf8.DecodeRune(p)
		p = p[size:]
		txt.WriteRune(r)
	}

	return
}

func (txt *Text) WriteString(s string) (n int, err error) {
	if len(s) == 0 {
		return
	}

	for _, r := range s {
		txt.WriteRune(r)
	}

	return len(s), nil
}

func (txt *Text) WriteByte(c byte) error {
	//FIXME: this is not correct, what if I want to write a 4-byte rune byte by byte?
	_, err := txt.WriteRune(rune(c))
	return err
}

func (txt *Text) WriteRune(r rune) (n int, err error) {
	n, err = utf8.RuneLen(r), nil // always returns this

	switch r {
	case '\n':
		txt.Dot -= pixel.Y(txt.lineHeight)
		txt.Dot = txt.Dot.WithX(txt.Orig.X())
		return
	case '\r':
		txt.Dot = txt.Dot.WithX(txt.Orig.X())
		return
	case '\t':
		rem := math.Mod(txt.Dot.X()-txt.Orig.X(), txt.tabWidth)
		rem = math.Mod(rem, rem+txt.tabWidth)
		if rem == 0 {
			rem = txt.tabWidth
		}
		txt.Dot += pixel.X(rem)
		return
	}

	if !txt.atlas.Contains(r) {
		r = unicode.ReplacementChar
	}
	if !txt.atlas.Contains(unicode.ReplacementChar) {
		return
	}

	glyph := txt.atlas.Glyph(r)

	if txt.prevR >= 0 {
		txt.Dot += pixel.X(txt.atlas.Kern(txt.prevR, r))
	}

	glyphBounds := glyph.Frame.Moved(txt.Dot - glyph.Orig)
	if glyphBounds.H() > 0 {
		glyphBounds = glyphBounds.Resized(txt.Dot, pixel.V(
			glyphBounds.W(),
			txt.atlas.Ascent()+txt.atlas.Descent(),
		))
	}
	if txt.bounds.W()*txt.bounds.H() == 0 {
		txt.bounds = glyphBounds
	} else {
		txt.bounds = txt.bounds.Union(glyphBounds)
	}

	a := pixel.V(glyph.Frame.Min.X(), glyph.Frame.Min.Y())
	b := pixel.V(glyph.Frame.Max.X(), glyph.Frame.Min.Y())
	c := pixel.V(glyph.Frame.Max.X(), glyph.Frame.Max.Y())
	d := pixel.V(glyph.Frame.Min.X(), glyph.Frame.Max.Y())

	for i, v := range []pixel.Vec{a, b, c, a, c, d} {
		txt.glyph[i].Position = v - glyph.Orig + txt.Dot
		txt.glyph[i].Picture = v
	}

	txt.tris = append(txt.tris, txt.glyph...)

	txt.Dot += pixel.X(glyph.Advance)
	txt.prevR = r

	txt.dirty = true

	return
}

func (txt *Text) Draw(t pixel.Target) {
	if txt.dirty {
		txt.trans.SetLen(txt.tris.Len())
		txt.trans.Update(&txt.tris)
		for i := range txt.trans {
			txt.trans[i].Position = txt.mat.Project(txt.trans[i].Position)
			txt.trans[i].Color = txt.trans[i].Color.Mul(txt.col)
		}
		txt.transD.Dirty()
		txt.dirty = false
	}
	txt.transD.Draw(t)
}
<|endoftext|>


<commit_before>package text

import (
	"image/color"
	"math"
	"unicode"
	"unicode/utf8"

	"github.com/faiface/pixel"
	"golang.org/x/image/font"
)

// ASCII is a set of all ASCII runes. These runes are codepoints from 32 to 127 inclusive.
var ASCII []rune

func init() {
	ASCII = make([]rune, unicode.MaxASCII-32)
	for i := range ASCII {
		ASCII[i] = rune(32 + i)
	}
}

// RangeTable takes a *unicode.RangeTable and generates a set of runes contained within that
// RangeTable.
func RangeTable(table *unicode.RangeTable) []rune {
	var runes []rune
	for _, rng := range table.R16 {
		for r := rng.Lo; r <= rng.Hi; r += rng.Stride {
			runes = append(runes, rune(r))
		}
	}
	for _, rng := range table.R32 {
		for r := rng.Lo; r <= rng.Hi; r += rng.Stride {
			runes = append(runes, rune(r))
		}
	}
	return runes
}

// Text allows for effiecient and convenient text drawing.
//
// To create a Text object, use the New constructor:
//   txt := text.New(face, text.ASCII)
//
// As suggested by the constructor, a Text object is always associated with one font face and a
// fixed set of runes. For example, the Text we created above can draw text using the font face
// contained in the face variable and is capable of drawing ASCII characters.
//
// Here we create a Text object which can draw ASCII and Katakana characters:
//   txt := text.New(face, text.ASCII, text.RangeTable(unicode.Katakana))
//
// Similarly to IMDraw, Text functions as a buffer. It implements io.Writer interface, so writing
// text to it is really simple:
//   fmt.Print(txt, "Hello, world!")
//
// Newlines, tabs and carriage returns are supported.
//
// Finally, if we want the written text to show up on some other Target, we can draw it:
//   txt.Draw(target)
//
// Text exports two important fields: Orig and Dot. Dot is the position where the next character
// will be written. Dot is automatically moved when writing to a Text object, but you can also
// manipulate it manually. Orig specifies the text origin, usually the top-left dot position. Dot is
// always aligned to Orig when writing newlines.
//
// To reset the Dot to the Orig, just assign it:
//   txt.Dot = txt.Orig
type Text struct {
	// Orig specifies the text origin, usually the top-left dot position. Dot is always aligned
	// to Orig when writing newlines.
	Orig pixel.Vec

	// Dot is the position where the next character will be written. Dot is automatically moved
	// when writing to a Text object, but you can also manipulate it manually
	Dot pixel.Vec

	atlas *Atlas

	lineHeight float64
	tabWidth   float64

	buf    []byte
	prevR  rune
	bounds pixel.Rect
	glyph  pixel.TrianglesData
	tris   pixel.TrianglesData

	mat    pixel.Matrix
	col    pixel.RGBA
	trans  pixel.TrianglesData
	transD pixel.Drawer
	dirty  bool
}

// New creates a new Text capable of drawing runes contained in the provided rune sets, plus
// unicode.ReplacementChar using the provided font.Face. New automatically generates an Atlas for
// the Text.
//
// Do not destroy or close the font.Face after creating a Text. Although Text caches most of the
// stuff (pre-drawn glyphs, etc.), it still uses the face for a few things.
//
// Here we create a Text capable of drawing ASCII characters using the Go Regular font.
//   ttf, err := truetype.Parse(goregular.TTF)
//   if err != nil {
//       panic(err)
//   }
//   face := truetype.NewFace(ttf, &truetype.Options{
//       Size: 14,
//   })
//   txt := text.New(face, text.ASCII)
func New(face font.Face, runeSets ...[]rune) *Text {
	runes := []rune{unicode.ReplacementChar}
	for _, set := range runeSets {
		runes = append(runes, set...)
	}

	atlas := NewAtlas(face, runes)

	txt := &Text{
		atlas:      atlas,
		lineHeight: atlas.LineHeight(),
		tabWidth:   atlas.Glyph(' ').Advance * 4,
		mat:        pixel.IM,
		col:        pixel.Alpha(1),
	}

	txt.glyph.SetLen(6)
	for i := range txt.glyph {
		txt.glyph[i].Color = pixel.Alpha(1)
		txt.glyph[i].Intensity = 1
	}

	txt.transD.Picture = txt.atlas.pic
	txt.transD.Triangles = &txt.trans

	txt.Clear()

	return txt
}

// Atlas returns the underlying Text's Atlas containing all of the pre-drawn glyphs. The Atlas is
// also useful for getting values such as the recommended line height.
func (txt *Text) Atlas() *Atlas {
	return txt.atlas
}

// SetMatrix sets a Matrix by which the text will be transformed before drawing to another Target.
func (txt *Text) SetMatrix(m pixel.Matrix) {
	if txt.mat != m {
		txt.mat = m
		txt.dirty = true
	}
}

// SetColorMask sets a color by which the text will be masked before drawingto another Target.
func (txt *Text) SetColorMask(c color.Color) {
	rgba := pixel.ToRGBA(c)
	if txt.col != rgba {
		txt.col = rgba
		txt.dirty = true
	}
}

// Bounds returns the bounding box of the text currently written to the Text excluding whitespace.
//
// If the Text is empty, a zero rectangle is returned.
func (txt *Text) Bounds() pixel.Rect {
	return txt.bounds
}

// BoundsOf returns the bounding box of s if it was to be written to the Text right now.
func (txt *Text) BoundsOf(s string) pixel.Rect {
	dot := txt.Dot
	prevR := txt.prevR
	bounds := pixel.Rect{}

	for _, r := range s {
		var control bool
		dot, control = txt.controlRune(r, dot)
		if control {
			continue
		}

		var b pixel.Rect
		_, _, b, dot = txt.Atlas().DrawRune(prevR, r, dot)

		if bounds.W()*bounds.H() == 0 {
			bounds = b
		} else {
			bounds = bounds.Union(b)
		}

		prevR = r
	}

	return bounds
}

// Color sets the text color. This does not affect any previously written text.
func (txt *Text) Color(c color.Color) {
	rgba := pixel.ToRGBA(c)
	for i := range txt.glyph {
		txt.glyph[i].Color = rgba
	}
}

// LineHeight sets the vertical distance between two lines of text. This does not affect any
// previously written text.
//
// Example:
//   txt.LineHeight(1.5 * txt.Atlas().LineHeight())
func (txt *Text) LineHeight(height float64) {
	txt.lineHeight = height
}

// TabWidth sets the horizontal tab width. Tab characters will align to the multiples of this width.
//
// Example:
//   txt.TabWidth(8 * txt.Atlas().Glyph(' ').Advance)
func (txt *Text) TabWidth(width float64) {
	txt.tabWidth = width
}

// Clear removes all written text from the Text.
func (txt *Text) Clear() {
	txt.prevR = -1
	txt.bounds = pixel.Rect{}
	txt.tris.SetLen(0)
	txt.dirty = true
}

// Write writes a slice of bytes to the Text. This method never fails, always returns len(p), nil.
func (txt *Text) Write(p []byte) (n int, err error) {
	txt.buf = append(txt.buf, p...)
	txt.drawBuf()
	return len(p), nil
}

// WriteString writes a string to the Text. This method never fails, always returns len(s), nil.
func (txt *Text) WriteString(s string) (n int, err error) {
	txt.buf = append(txt.buf, s...)
	txt.drawBuf()
	return len(s), nil
}

// WriteByte writes a byte to the Text. This method never fails, always returns nil.
//
// Writing a multi-byte rune byte-by-byte is perfectly supported.
func (txt *Text) WriteByte(c byte) error {
	txt.buf = append(txt.buf, c)
	txt.drawBuf()
	return nil
}

// WriteRune writes a rune to the Text. This method never fails, always returns utf8.RuneLen(r), nil.
func (txt *Text) WriteRune(r rune) (n int, err error) {
	var b [4]byte
	n = utf8.EncodeRune(b[:], r)
	txt.buf = append(txt.buf, b[:n]...)
	txt.drawBuf()
	return n, nil
}

// Draw draws all text written to the Text to the provided Target. The text is transformed by the
// Text's matrix and color mask.
func (txt *Text) Draw(t pixel.Target) {
	if txt.dirty {
		txt.trans.SetLen(txt.tris.Len())
		txt.trans.Update(&txt.tris)

		for i := range txt.trans {
			txt.trans[i].Position = txt.mat.Project(txt.trans[i].Position)
			txt.trans[i].Color = txt.trans[i].Color.Mul(txt.col)
		}

		txt.transD.Dirty()
		txt.dirty = false
	}

	txt.transD.Draw(t)
}

// controlRune checks if r is a control rune (newline, tab, ...). If it is, a new dot position and
// true is returned. If r is not a control rune, the original dot and false is returned.
func (txt *Text) controlRune(r rune, dot pixel.Vec) (newDot pixel.Vec, control bool) {
	switch r {
	case '\n':
		dot -= pixel.Y(txt.lineHeight)
		dot = dot.WithX(txt.Orig.X())
	case '\r':
		dot = dot.WithX(txt.Orig.X())
	case '\t':
		rem := math.Mod(dot.X()-txt.Orig.X(), txt.tabWidth)
		rem = math.Mod(rem, rem+txt.tabWidth)
		if rem == 0 {
			rem = txt.tabWidth
		}
		dot += pixel.X(rem)
	default:
		return dot, false
	}
	return dot, true
}

func (txt *Text) drawBuf() {
	for utf8.FullRune(txt.buf) {
		r, size := utf8.DecodeRune(txt.buf)
		txt.buf = txt.buf[size:]

		var control bool
		txt.Dot, control = txt.controlRune(r, txt.Dot)
		if control {
			continue
		}

		var rect, frame, bounds pixel.Rect
		rect, frame, bounds, txt.Dot = txt.Atlas().DrawRune(txt.prevR, r, txt.Dot)

		txt.prevR = r

		rv := [...]pixel.Vec{pixel.V(rect.Min.X(), rect.Min.Y()),
			pixel.V(rect.Max.X(), rect.Min.Y()),
			pixel.V(rect.Max.X(), rect.Max.Y()),
			pixel.V(rect.Min.X(), rect.Max.Y()),
		}

		fv := [...]pixel.Vec{pixel.V(frame.Min.X(), frame.Min.Y()),
			pixel.V(frame.Max.X(), frame.Min.Y()),
			pixel.V(frame.Max.X(), frame.Max.Y()),
			pixel.V(frame.Min.X(), frame.Max.Y()),
		}

		for i, j := range [...]int{0, 1, 2, 0, 2, 3} {
			txt.glyph[i].Position = rv[j]
			txt.glyph[i].Picture = fv[j]
		}

		txt.tris = append(txt.tris, txt.glyph...)
		txt.dirty = true

		if txt.bounds.W()*txt.bounds.H() == 0 {
			txt.bounds = bounds
		} else {
			txt.bounds = txt.bounds.Union(bounds)
		}
	}
}
<commit_msg>add Text.Matrix and Text.ColorMask<commit_after>package text

import (
	"image/color"
	"math"
	"unicode"
	"unicode/utf8"

	"github.com/faiface/pixel"
	"golang.org/x/image/font"
)

// ASCII is a set of all ASCII runes. These runes are codepoints from 32 to 127 inclusive.
var ASCII []rune

func init() {
	ASCII = make([]rune, unicode.MaxASCII-32)
	for i := range ASCII {
		ASCII[i] = rune(32 + i)
	}
}

// RangeTable takes a *unicode.RangeTable and generates a set of runes contained within that
// RangeTable.
func RangeTable(table *unicode.RangeTable) []rune {
	var runes []rune
	for _, rng := range table.R16 {
		for r := rng.Lo; r <= rng.Hi; r += rng.Stride {
			runes = append(runes, rune(r))
		}
	}
	for _, rng := range table.R32 {
		for r := rng.Lo; r <= rng.Hi; r += rng.Stride {
			runes = append(runes, rune(r))
		}
	}
	return runes
}

// Text allows for effiecient and convenient text drawing.
//
// To create a Text object, use the New constructor:
//   txt := text.New(face, text.ASCII)
//
// As suggested by the constructor, a Text object is always associated with one font face and a
// fixed set of runes. For example, the Text we created above can draw text using the font face
// contained in the face variable and is capable of drawing ASCII characters.
//
// Here we create a Text object which can draw ASCII and Katakana characters:
//   txt := text.New(face, text.ASCII, text.RangeTable(unicode.Katakana))
//
// Similarly to IMDraw, Text functions as a buffer. It implements io.Writer interface, so writing
// text to it is really simple:
//   fmt.Print(txt, "Hello, world!")
//
// Newlines, tabs and carriage returns are supported.
//
// Finally, if we want the written text to show up on some other Target, we can draw it:
//   txt.Draw(target)
//
// Text exports two important fields: Orig and Dot. Dot is the position where the next character
// will be written. Dot is automatically moved when writing to a Text object, but you can also
// manipulate it manually. Orig specifies the text origin, usually the top-left dot position. Dot is
// always aligned to Orig when writing newlines.
//
// To reset the Dot to the Orig, just assign it:
//   txt.Dot = txt.Orig
type Text struct {
	// Orig specifies the text origin, usually the top-left dot position. Dot is always aligned
	// to Orig when writing newlines.
	Orig pixel.Vec

	// Dot is the position where the next character will be written. Dot is automatically moved
	// when writing to a Text object, but you can also manipulate it manually
	Dot pixel.Vec

	atlas *Atlas

	lineHeight float64
	tabWidth   float64

	buf    []byte
	prevR  rune
	bounds pixel.Rect
	glyph  pixel.TrianglesData
	tris   pixel.TrianglesData

	mat    pixel.Matrix
	col    pixel.RGBA
	trans  pixel.TrianglesData
	transD pixel.Drawer
	dirty  bool
}

// New creates a new Text capable of drawing runes contained in the provided rune sets, plus
// unicode.ReplacementChar using the provided font.Face. New automatically generates an Atlas for
// the Text.
//
// Do not destroy or close the font.Face after creating a Text. Although Text caches most of the
// stuff (pre-drawn glyphs, etc.), it still uses the face for a few things.
//
// Here we create a Text capable of drawing ASCII characters using the Go Regular font.
//   ttf, err := truetype.Parse(goregular.TTF)
//   if err != nil {
//       panic(err)
//   }
//   face := truetype.NewFace(ttf, &truetype.Options{
//       Size: 14,
//   })
//   txt := text.New(face, text.ASCII)
func New(face font.Face, runeSets ...[]rune) *Text {
	runes := []rune{unicode.ReplacementChar}
	for _, set := range runeSets {
		runes = append(runes, set...)
	}

	atlas := NewAtlas(face, runes)

	txt := &Text{
		atlas:      atlas,
		lineHeight: atlas.LineHeight(),
		tabWidth:   atlas.Glyph(' ').Advance * 4,
		mat:        pixel.IM,
		col:        pixel.Alpha(1),
	}

	txt.glyph.SetLen(6)
	for i := range txt.glyph {
		txt.glyph[i].Color = pixel.Alpha(1)
		txt.glyph[i].Intensity = 1
	}

	txt.transD.Picture = txt.atlas.pic
	txt.transD.Triangles = &txt.trans

	txt.Clear()

	return txt
}

// Atlas returns the underlying Text's Atlas containing all of the pre-drawn glyphs. The Atlas is
// also useful for getting values such as the recommended line height.
func (txt *Text) Atlas() *Atlas {
	return txt.atlas
}

// SetMatrix sets a Matrix by which the text will be transformed before drawing to another Target.
func (txt *Text) SetMatrix(m pixel.Matrix) {
	if txt.mat != m {
		txt.mat = m
		txt.dirty = true
	}
}

// Matrix returns the current Text's matrix.
func (txt *Text) Matrix() pixel.Matrix {
	return txt.mat
}

// SetColorMask sets a color by which the text will be masked before drawingto another Target.
func (txt *Text) SetColorMask(c color.Color) {
	rgba := pixel.ToRGBA(c)
	if txt.col != rgba {
		txt.col = rgba
		txt.dirty = true
	}
}

// ColorMask returns the current Text's color mask.
func (txt *Text) ColorMask() pixel.RGBA {
	return txt.col
}

// Bounds returns the bounding box of the text currently written to the Text excluding whitespace.
//
// If the Text is empty, a zero rectangle is returned.
func (txt *Text) Bounds() pixel.Rect {
	return txt.bounds
}

// BoundsOf returns the bounding box of s if it was to be written to the Text right now.
func (txt *Text) BoundsOf(s string) pixel.Rect {
	dot := txt.Dot
	prevR := txt.prevR
	bounds := pixel.Rect{}

	for _, r := range s {
		var control bool
		dot, control = txt.controlRune(r, dot)
		if control {
			continue
		}

		var b pixel.Rect
		_, _, b, dot = txt.Atlas().DrawRune(prevR, r, dot)

		if bounds.W()*bounds.H() == 0 {
			bounds = b
		} else {
			bounds = bounds.Union(b)
		}

		prevR = r
	}

	return bounds
}

// Color sets the text color. This does not affect any previously written text.
func (txt *Text) Color(c color.Color) {
	rgba := pixel.ToRGBA(c)
	for i := range txt.glyph {
		txt.glyph[i].Color = rgba
	}
}

// LineHeight sets the vertical distance between two lines of text. This does not affect any
// previously written text.
//
// Example:
//   txt.LineHeight(1.5 * txt.Atlas().LineHeight())
func (txt *Text) LineHeight(height float64) {
	txt.lineHeight = height
}

// TabWidth sets the horizontal tab width. Tab characters will align to the multiples of this width.
//
// Example:
//   txt.TabWidth(8 * txt.Atlas().Glyph(' ').Advance)
func (txt *Text) TabWidth(width float64) {
	txt.tabWidth = width
}

// Clear removes all written text from the Text.
func (txt *Text) Clear() {
	txt.prevR = -1
	txt.bounds = pixel.Rect{}
	txt.tris.SetLen(0)
	txt.dirty = true
}

// Write writes a slice of bytes to the Text. This method never fails, always returns len(p), nil.
func (txt *Text) Write(p []byte) (n int, err error) {
	txt.buf = append(txt.buf, p...)
	txt.drawBuf()
	return len(p), nil
}

// WriteString writes a string to the Text. This method never fails, always returns len(s), nil.
func (txt *Text) WriteString(s string) (n int, err error) {
	txt.buf = append(txt.buf, s...)
	txt.drawBuf()
	return len(s), nil
}

// WriteByte writes a byte to the Text. This method never fails, always returns nil.
//
// Writing a multi-byte rune byte-by-byte is perfectly supported.
func (txt *Text) WriteByte(c byte) error {
	txt.buf = append(txt.buf, c)
	txt.drawBuf()
	return nil
}

// WriteRune writes a rune to the Text. This method never fails, always returns utf8.RuneLen(r), nil.
func (txt *Text) WriteRune(r rune) (n int, err error) {
	var b [4]byte
	n = utf8.EncodeRune(b[:], r)
	txt.buf = append(txt.buf, b[:n]...)
	txt.drawBuf()
	return n, nil
}

// Draw draws all text written to the Text to the provided Target. The text is transformed by the
// Text's matrix and color mask.
func (txt *Text) Draw(t pixel.Target) {
	if txt.dirty {
		txt.trans.SetLen(txt.tris.Len())
		txt.trans.Update(&txt.tris)

		for i := range txt.trans {
			txt.trans[i].Position = txt.mat.Project(txt.trans[i].Position)
			txt.trans[i].Color = txt.trans[i].Color.Mul(txt.col)
		}

		txt.transD.Dirty()
		txt.dirty = false
	}

	txt.transD.Draw(t)
}

// controlRune checks if r is a control rune (newline, tab, ...). If it is, a new dot position and
// true is returned. If r is not a control rune, the original dot and false is returned.
func (txt *Text) controlRune(r rune, dot pixel.Vec) (newDot pixel.Vec, control bool) {
	switch r {
	case '\n':
		dot -= pixel.Y(txt.lineHeight)
		dot = dot.WithX(txt.Orig.X())
	case '\r':
		dot = dot.WithX(txt.Orig.X())
	case '\t':
		rem := math.Mod(dot.X()-txt.Orig.X(), txt.tabWidth)
		rem = math.Mod(rem, rem+txt.tabWidth)
		if rem == 0 {
			rem = txt.tabWidth
		}
		dot += pixel.X(rem)
	default:
		return dot, false
	}
	return dot, true
}

func (txt *Text) drawBuf() {
	for utf8.FullRune(txt.buf) {
		r, size := utf8.DecodeRune(txt.buf)
		txt.buf = txt.buf[size:]

		var control bool
		txt.Dot, control = txt.controlRune(r, txt.Dot)
		if control {
			continue
		}

		var rect, frame, bounds pixel.Rect
		rect, frame, bounds, txt.Dot = txt.Atlas().DrawRune(txt.prevR, r, txt.Dot)

		txt.prevR = r

		rv := [...]pixel.Vec{pixel.V(rect.Min.X(), rect.Min.Y()),
			pixel.V(rect.Max.X(), rect.Min.Y()),
			pixel.V(rect.Max.X(), rect.Max.Y()),
			pixel.V(rect.Min.X(), rect.Max.Y()),
		}

		fv := [...]pixel.Vec{pixel.V(frame.Min.X(), frame.Min.Y()),
			pixel.V(frame.Max.X(), frame.Min.Y()),
			pixel.V(frame.Max.X(), frame.Max.Y()),
			pixel.V(frame.Min.X(), frame.Max.Y()),
		}

		for i, j := range [...]int{0, 1, 2, 0, 2, 3} {
			txt.glyph[i].Position = rv[j]
			txt.glyph[i].Picture = fv[j]
		}

		txt.tris = append(txt.tris, txt.glyph...)
		txt.dirty = true

		if txt.bounds.W()*txt.bounds.H() == 0 {
			txt.bounds = bounds
		} else {
			txt.bounds = txt.bounds.Union(bounds)
		}
	}
}
<|endoftext|>