<commit_before>// Package consistent provides a consistent hashing function.
//
// Consistent hashing is often used to distribute requests to a changing set of servers.  For example,
// say you have some cache servers cacheA, cacheB, and cacheC.  You want to decide which cache server
// to use to look up information on a user.
//
// You could use a typical hash table and hash the user id
// to one of cacheA, cacheB, or cacheC.  But with a typical hash table, if you add or remove a server,
// almost all keys will get remapped to different results, which basically could bring your service
// to a grinding halt while the caches get rebuilt.
//
// With a consistent hash, adding or removing a server drastically reduces the number of keys that
// get remapped.
//
// Read more about consistent hashing on wikipedia:  http://en.wikipedia.org/wiki/Consistent_hashing
//
package consistent

import (
	"errors"
	"fmt"
	"hash/crc32"
	"sort"
)

type uints []uint32
// Len returns the length of the uints array.
func (x uints) Len() int           { return len(x) }
// Less returns true if element i is less than element j.
func (x uints) Less(i, j int) bool { return x[i] < x[j] }
// Swap exchanges elements i and j.
func (x uints) Swap(i, j int)      { x[i], x[j] = x[j], x[i] }

// ErrEmptyCircle is the error returned when trying to get an element when nothing has been added to hash.
var ErrEmptyCircle = errors.New("empty circle")

// Consistent holds the information about the members of the consistent hash circle.
type Consistent struct {
	circle           map[uint32]string
	sortedHashes     uints
	NumberOfReplicas int
	count            int64
}

// New creates a new Consistent object with a default setting of 20 replicas for each entry.
//
// To change the number of replicas, set NumberOfReplicas before adding entries.
func New() *Consistent {
	c := new(Consistent)
	c.NumberOfReplicas = 20
	c.circle = make(map[uint32]string)
	return c
}

// eltKey generates a string key for an element with an index.
func (c *Consistent) eltKey(elt string, idx int) string {
        return fmt.Sprintf("%s|%d", elt, idx)
}

// Add inserts a string element in the consistent hash.
func (c *Consistent) Add(elt string) {
	for i := 0; i < c.NumberOfReplicas; i++ {
		c.circle[c.hashKey(c.eltKey(elt, i))] = elt
	}
	c.updateSortedHashes()
	c.count++
}

// Remove removes an element from the hash.
func (c *Consistent) Remove(elt string) {
	for i := 0; i < c.NumberOfReplicas; i++ {
		delete(c.circle, c.hashKey(c.eltKey(elt, i)))
	}
	c.updateSortedHashes()
	c.count--
}

// Get returns an element close to where name hashes to in the circle.
func (c *Consistent) Get(name string) (string, error) {
	if len(c.circle) == 0 {
		return "", ErrEmptyCircle
	}
	key := c.hashKey(name)
	f := func(x int) bool {
		return c.sortedHashes[x] > key
	}
	i := sort.Search(len(c.sortedHashes), f)
	if i >= len(c.sortedHashes) {
		i = 0
	}
	return c.circle[c.sortedHashes[i]], nil
}

// Get two returns the two closest distinct elements to the name input in the circle.
func (c *Consistent) GetTwo(name string) (string, string, error) {
	if len(c.circle) == 0 {
		return "", "", ErrEmptyCircle
	}
	key := c.hashKey(name)
	f := func(x int) bool {
		return c.sortedHashes[x] > key
	}
	i := sort.Search(len(c.sortedHashes), f)
	if i >= len(c.sortedHashes) {
		i = 0
	}
	a := c.circle[c.sortedHashes[i]]

	if c.count == 1 {
		return a, "", nil
	}

	start := i
	i++
	if i >= len(c.sortedHashes) {
		i = 0
	}
	var b string
	for i != start {
		b = c.circle[c.sortedHashes[i]]
		if b != a {
			break
		}
		i++
		if i >= len(c.sortedHashes) {
			i = 0
		}
	}
	return a, b, nil
}

func (c *Consistent) hashKey(key string) uint32 {
	return crc32.ChecksumIEEE([]byte(key))
}

func (c *Consistent) updateSortedHashes() {
	hashes := uints(nil)
	for k := range c.circle {
		hashes = append(hashes, k)
	}
	sort.Sort(hashes)
	c.sortedHashes = hashes
}

<commit_msg>fixed comment<commit_after>// Package consistent provides a consistent hashing function.
//
// Consistent hashing is often used to distribute requests to a changing set of servers.  For example,
// say you have some cache servers cacheA, cacheB, and cacheC.  You want to decide which cache server
// to use to look up information on a user.
//
// You could use a typical hash table and hash the user id
// to one of cacheA, cacheB, or cacheC.  But with a typical hash table, if you add or remove a server,
// almost all keys will get remapped to different results, which basically could bring your service
// to a grinding halt while the caches get rebuilt.
//
// With a consistent hash, adding or removing a server drastically reduces the number of keys that
// get remapped.
//
// Read more about consistent hashing on wikipedia:  http://en.wikipedia.org/wiki/Consistent_hashing
//
package consistent

import (
	"errors"
	"fmt"
	"hash/crc32"
	"sort"
)

type uints []uint32
// Len returns the length of the uints array.
func (x uints) Len() int           { return len(x) }
// Less returns true if element i is less than element j.
func (x uints) Less(i, j int) bool { return x[i] < x[j] }
// Swap exchanges elements i and j.
func (x uints) Swap(i, j int)      { x[i], x[j] = x[j], x[i] }

// ErrEmptyCircle is the error returned when trying to get an element when nothing has been added to hash.
var ErrEmptyCircle = errors.New("empty circle")

// Consistent holds the information about the members of the consistent hash circle.
type Consistent struct {
	circle           map[uint32]string
	sortedHashes     uints
	NumberOfReplicas int
	count            int64
}

// New creates a new Consistent object with a default setting of 20 replicas for each entry.
//
// To change the number of replicas, set NumberOfReplicas before adding entries.
func New() *Consistent {
	c := new(Consistent)
	c.NumberOfReplicas = 20
	c.circle = make(map[uint32]string)
	return c
}

// eltKey generates a string key for an element with an index.
func (c *Consistent) eltKey(elt string, idx int) string {
        return fmt.Sprintf("%s|%d", elt, idx)
}

// Add inserts a string element in the consistent hash.
func (c *Consistent) Add(elt string) {
	for i := 0; i < c.NumberOfReplicas; i++ {
		c.circle[c.hashKey(c.eltKey(elt, i))] = elt
	}
	c.updateSortedHashes()
	c.count++
}

// Remove removes an element from the hash.
func (c *Consistent) Remove(elt string) {
	for i := 0; i < c.NumberOfReplicas; i++ {
		delete(c.circle, c.hashKey(c.eltKey(elt, i)))
	}
	c.updateSortedHashes()
	c.count--
}

// Get returns an element close to where name hashes to in the circle.
func (c *Consistent) Get(name string) (string, error) {
	if len(c.circle) == 0 {
		return "", ErrEmptyCircle
	}
	key := c.hashKey(name)
	f := func(x int) bool {
		return c.sortedHashes[x] > key
	}
	i := sort.Search(len(c.sortedHashes), f)
	if i >= len(c.sortedHashes) {
		i = 0
	}
	return c.circle[c.sortedHashes[i]], nil
}

// GetTwo returns the two closest distinct elements to the name input in the circle.
func (c *Consistent) GetTwo(name string) (string, string, error) {
	if len(c.circle) == 0 {
		return "", "", ErrEmptyCircle
	}
	key := c.hashKey(name)
	f := func(x int) bool {
		return c.sortedHashes[x] > key
	}
	i := sort.Search(len(c.sortedHashes), f)
	if i >= len(c.sortedHashes) {
		i = 0
	}
	a := c.circle[c.sortedHashes[i]]

	if c.count == 1 {
		return a, "", nil
	}

	start := i
	i++
	if i >= len(c.sortedHashes) {
		i = 0
	}
	var b string
	for i != start {
		b = c.circle[c.sortedHashes[i]]
		if b != a {
			break
		}
		i++
		if i >= len(c.sortedHashes) {
			i = 0
		}
	}
	return a, b, nil
}

func (c *Consistent) hashKey(key string) uint32 {
	return crc32.ChecksumIEEE([]byte(key))
}

func (c *Consistent) updateSortedHashes() {
	hashes := uints(nil)
	for k := range c.circle {
		hashes = append(hashes, k)
	}
	sort.Sort(hashes)
	c.sortedHashes = hashes
}

<|endoftext|>


<commit_before>package consul

import (
	"crypto/tls"
	"fmt"
	"io"
	"math/rand"
	"net"
	"strings"
	"time"

	"github.com/armon/go-metrics"
	"github.com/hashicorp/consul/consul/structs"
	"github.com/hashicorp/go-msgpack/codec"
	"github.com/hashicorp/yamux"
	"github.com/inconshreveable/muxado"
)

type RPCType byte

const (
	rpcConsul RPCType = iota
	rpcRaft
	rpcMultiplex
	rpcTLS
	rpcMultiplexV2
)

const (
	// maxQueryTime is used to bound the limit of a blocking query
	maxQueryTime = 600 * time.Second

	// Warn if the Raft command is larger than this.
	// If it's over 1MB something is probably being abusive.
	raftWarnSize = 1024 * 1024

	// enqueueLimit caps how long we will wait to enqueue
	// a new Raft command. Something is probably wrong if this
	// value is ever reached. However, it prevents us from blocking
	// the requesting goroutine forever.
	enqueueLimit = 30 * time.Second
)

// listen is used to listen for incoming RPC connections
func (s *Server) listen() {
	for {
		// Accept a connection
		conn, err := s.rpcListener.Accept()
		if err != nil {
			if s.shutdown {
				return
			}
			s.logger.Printf("[ERR] consul.rpc: failed to accept RPC conn: %v", err)
			continue
		}

		go s.handleConn(conn, false)
		metrics.IncrCounter([]string{"consul", "rpc", "accept_conn"}, 1)
	}
}

// handleConn is used to determine if this is a Raft or
// Consul type RPC connection and invoke the correct handler
func (s *Server) handleConn(conn net.Conn, isTLS bool) {
	// Read a single byte
	buf := make([]byte, 1)
	if _, err := conn.Read(buf); err != nil {
		if err != io.EOF {
			s.logger.Printf("[ERR] consul.rpc: failed to read byte: %v", err)
		}
		conn.Close()
		return
	}

	// Enforce TLS if VerifyIncoming is set
	if s.config.VerifyIncoming && !isTLS && RPCType(buf[0]) != rpcTLS {
		s.logger.Printf("[WARN] consul.rpc: Non-TLS connection attempted with VerifyIncoming set")
		conn.Close()
		return
	}

	// Switch on the byte
	switch RPCType(buf[0]) {
	case rpcConsul:
		s.handleConsulConn(conn)

	case rpcRaft:
		metrics.IncrCounter([]string{"consul", "rpc", "raft_handoff"}, 1)
		s.raftLayer.Handoff(conn)

	case rpcMultiplex:
		s.handleMultiplex(conn)

	case rpcTLS:
		if s.rpcTLS == nil {
			s.logger.Printf("[WARN] consul.rpc: TLS connection attempted, server not configured for TLS")
			conn.Close()
			return
		}
		conn = tls.Server(conn, s.rpcTLS)
		s.handleConn(conn, true)

	case rpcMultiplexV2:
		s.handleMultiplexV2(conn)

	default:
		s.logger.Printf("[ERR] consul.rpc: unrecognized RPC byte: %v", buf[0])
		conn.Close()
		return
	}
}

// handleMultiplex is used to multiplex a single incoming connection
// using the Muxado multiplexer
func (s *Server) handleMultiplex(conn net.Conn) {
	defer conn.Close()
	server := muxado.Server(conn)
	for {
		sub, err := server.Accept()
		if err != nil {
			if !strings.Contains(err.Error(), "closed") {
				s.logger.Printf("[ERR] consul.rpc: multiplex conn accept failed: %v", err)
			}
			return
		}
		go s.handleConsulConn(sub)
	}
}

// handleMultiplexV2 is used to multiplex a single incoming connection
// using the Yamux multiplexer
func (s *Server) handleMultiplexV2(conn net.Conn) {
	defer conn.Close()
	conf := yamux.DefaultConfig()
	conf.LogOutput = s.config.LogOutput
	server, _ := yamux.Server(conn, conf)
	for {
		sub, err := server.Accept()
		if err != nil {
			if err != io.EOF {
				s.logger.Printf("[ERR] consul.rpc: multiplex conn accept failed: %v", err)
			}
			return
		}
		go s.handleConsulConn(sub)
	}
}

// handleConsulConn is used to service a single Consul RPC connection
func (s *Server) handleConsulConn(conn net.Conn) {
	defer conn.Close()
	rpcCodec := codec.GoRpc.ServerCodec(conn, msgpackHandle)
	for {
		select {
		case <-s.shutdownCh:
			return
		default:
		}

		if err := s.rpcServer.ServeRequest(rpcCodec); err != nil {
			if err != io.EOF && !strings.Contains(err.Error(), "closed") {
				s.logger.Printf("[ERR] consul.rpc: RPC error: %v (%v)", err, conn)
				metrics.IncrCounter([]string{"consul", "rpc", "request_error"}, 1)
			}
			return
		}
		metrics.IncrCounter([]string{"consul", "rpc", "request"}, 1)
	}
}

// forward is used to forward to a remote DC or to forward to the local leader
// Returns a bool of if forwarding was performed, as well as any error
func (s *Server) forward(method string, info structs.RPCInfo, args interface{}, reply interface{}) (bool, error) {
	// Handle DC forwarding
	dc := info.RequestDatacenter()
	if dc != s.config.Datacenter {
		err := s.forwardDC(method, dc, args, reply)
		return true, err
	}

	// Check if we can allow a stale read
	if info.IsRead() && info.AllowStaleRead() {
		return false, nil
	}

	// Handle leader forwarding
	if !s.IsLeader() {
		err := s.forwardLeader(method, args, reply)
		return true, err
	}
	return false, nil
}

// forwardLeader is used to forward an RPC call to the leader, or fail if no leader
func (s *Server) forwardLeader(method string, args interface{}, reply interface{}) error {
	// Get the leader
	leader := s.raft.Leader()
	if leader == "" {
		return structs.ErrNoLeader
	}

	// Lookup the server
	s.localLock.RLock()
	server := s.localConsuls[leader]
	s.localLock.RUnlock()

	// Handle a missing server
	if server == nil {
		return structs.ErrNoLeader
	}
	return s.connPool.RPC(s.config.Datacenter, server.Addr, server.Version, method, args, reply)
}

// forwardDC is used to forward an RPC call to a remote DC, or fail if no servers
func (s *Server) forwardDC(method, dc string, args interface{}, reply interface{}) error {
	// Bail if we can't find any servers
	s.remoteLock.RLock()
	servers := s.remoteConsuls[dc]
	if len(servers) == 0 {
		s.remoteLock.RUnlock()
		s.logger.Printf("[WARN] consul.rpc: RPC request for DC '%s', no path found", dc)
		return structs.ErrNoDCPath
	}

	// Select a random addr
	offset := rand.Int31() % int32(len(servers))
	server := servers[offset]
	s.remoteLock.RUnlock()

	// Forward to remote Consul
	metrics.IncrCounter([]string{"consul", "rpc", "cross-dc", dc}, 1)
	return s.connPool.RPC(dc, server.Addr, server.Version, method, args, reply)
}

// globalRPC is used to forward an RPC request to one server in each datacenter.
// This will only error for RPC-related errors. Otherwise, application-level
// errors can be sent in the response objects.
func (s *Server) globalRPC(method string, args interface{},
	reply structs.CompoundResponse) error {

	errorCh := make(chan error)
	respCh := make(chan interface{})

	// Make a new request into each datacenter
	for dc, _ := range s.remoteConsuls {
		go func(dc string) {
			rr := reply.New()
			if err := s.forwardDC(method, dc, args, &rr); err != nil {
				errorCh <- err
				return
			}
			respCh <- rr
		}(dc)
	}

	replies, total := 0, len(s.remoteConsuls)
	for replies < total {
		select {
		case err := <-errorCh:
			return err
		case rr := <-respCh:
			reply.Add(rr)
			replies++
		}
	}
	return nil
}

// raftApply is used to encode a message, run it through raft, and return
// the FSM response along with any errors
func (s *Server) raftApply(t structs.MessageType, msg interface{}) (interface{}, error) {
	buf, err := structs.Encode(t, msg)
	if err != nil {
		return nil, fmt.Errorf("Failed to encode request: %v", err)
	}

	// Warn if the command is very large
	if n := len(buf); n > raftWarnSize {
		s.logger.Printf("[WARN] consul: Attempting to apply large raft entry (%d bytes)", n)
	}

	future := s.raft.Apply(buf, enqueueLimit)
	if err := future.Error(); err != nil {
		return nil, err
	}

	return future.Response(), nil
}

// blockingRPC is used for queries that need to wait for a
// minimum index. This is used to block and wait for changes.
func (s *Server) blockingRPC(b *structs.QueryOptions, m *structs.QueryMeta,
	tables MDBTables, run func() error) error {
	opts := blockingRPCOptions{
		queryOpts: b,
		queryMeta: m,
		tables:    tables,
		run:       run,
	}
	return s.blockingRPCOpt(&opts)
}

// blockingRPCOptions is used to parameterize blockingRPCOpt since
// it takes so many options. It should be prefered over blockingRPC.
type blockingRPCOptions struct {
	queryOpts *structs.QueryOptions
	queryMeta *structs.QueryMeta
	tables    MDBTables
	kvWatch   bool
	kvPrefix  string
	run       func() error
}

// blockingRPCOpt is the replacement for blockingRPC as it allows
// for more parameterization easily. It should be prefered over blockingRPC.
func (s *Server) blockingRPCOpt(opts *blockingRPCOptions) error {
	var timeout <-chan time.Time
	var notifyCh chan struct{}
	var state *StateStore

	// Fast path non-blocking
	if opts.queryOpts.MinQueryIndex == 0 {
		goto RUN_QUERY
	}

	// Sanity check that we have tables to block on
	if len(opts.tables) == 0 && !opts.kvWatch {
		panic("no tables to block on")
	}

	// Restrict the max query time, and ensure there is always one
	if opts.queryOpts.MaxQueryTime > maxQueryTime {
		opts.queryOpts.MaxQueryTime = maxQueryTime
	} else if opts.queryOpts.MaxQueryTime <= 0 {
		opts.queryOpts.MaxQueryTime = maxQueryTime
	}

	// Setup a query timeout
	timeout = time.After(opts.queryOpts.MaxQueryTime)

	// Setup the notify channel
	notifyCh = make(chan struct{}, 1)

	// Ensure we tear down any watchers on return
	state = s.fsm.State()
	defer func() {
		state.StopWatch(opts.tables, notifyCh)
		if opts.kvWatch {
			state.StopWatchKV(opts.kvPrefix, notifyCh)
		}
	}()

REGISTER_NOTIFY:
	// Register the notification channel. This may be done
	// multiple times if we have not reached the target wait index.
	state.Watch(opts.tables, notifyCh)
	if opts.kvWatch {
		state.WatchKV(opts.kvPrefix, notifyCh)
	}

RUN_QUERY:
	// Update the query meta data
	s.setQueryMeta(opts.queryMeta)

	// Check if query must be consistent
	if opts.queryOpts.RequireConsistent {
		if err := s.consistentRead(); err != nil {
			return err
		}
	}

	// Run the query function
	metrics.IncrCounter([]string{"consul", "rpc", "query"}, 1)
	err := opts.run()

	// Check for minimum query time
	if err == nil && opts.queryMeta.Index > 0 && opts.queryMeta.Index <= opts.queryOpts.MinQueryIndex {
		select {
		case <-notifyCh:
			goto REGISTER_NOTIFY
		case <-timeout:
		}
	}
	return err
}

// setQueryMeta is used to populate the QueryMeta data for an RPC call
func (s *Server) setQueryMeta(m *structs.QueryMeta) {
	if s.IsLeader() {
		m.LastContact = 0
		m.KnownLeader = true
	} else {
		m.LastContact = time.Now().Sub(s.raft.LastContact())
		m.KnownLeader = (s.raft.Leader() != "")
	}
}

// consistentRead is used to ensure we do not perform a stale
// read. This is done by verifying leadership before the read.
func (s *Server) consistentRead() error {
	defer metrics.MeasureSince([]string{"consul", "rpc", "consistentRead"}, time.Now())
	future := s.raft.VerifyLeader()
	return future.Error()
}
<commit_msg>consul: proactively clear timers<commit_after>package consul

import (
	"crypto/tls"
	"fmt"
	"io"
	"math/rand"
	"net"
	"strings"
	"time"

	"github.com/armon/go-metrics"
	"github.com/hashicorp/consul/consul/structs"
	"github.com/hashicorp/go-msgpack/codec"
	"github.com/hashicorp/yamux"
	"github.com/inconshreveable/muxado"
)

type RPCType byte

const (
	rpcConsul RPCType = iota
	rpcRaft
	rpcMultiplex
	rpcTLS
	rpcMultiplexV2
)

const (
	// maxQueryTime is used to bound the limit of a blocking query
	maxQueryTime = 600 * time.Second

	// Warn if the Raft command is larger than this.
	// If it's over 1MB something is probably being abusive.
	raftWarnSize = 1024 * 1024

	// enqueueLimit caps how long we will wait to enqueue
	// a new Raft command. Something is probably wrong if this
	// value is ever reached. However, it prevents us from blocking
	// the requesting goroutine forever.
	enqueueLimit = 30 * time.Second
)

// listen is used to listen for incoming RPC connections
func (s *Server) listen() {
	for {
		// Accept a connection
		conn, err := s.rpcListener.Accept()
		if err != nil {
			if s.shutdown {
				return
			}
			s.logger.Printf("[ERR] consul.rpc: failed to accept RPC conn: %v", err)
			continue
		}

		go s.handleConn(conn, false)
		metrics.IncrCounter([]string{"consul", "rpc", "accept_conn"}, 1)
	}
}

// handleConn is used to determine if this is a Raft or
// Consul type RPC connection and invoke the correct handler
func (s *Server) handleConn(conn net.Conn, isTLS bool) {
	// Read a single byte
	buf := make([]byte, 1)
	if _, err := conn.Read(buf); err != nil {
		if err != io.EOF {
			s.logger.Printf("[ERR] consul.rpc: failed to read byte: %v", err)
		}
		conn.Close()
		return
	}

	// Enforce TLS if VerifyIncoming is set
	if s.config.VerifyIncoming && !isTLS && RPCType(buf[0]) != rpcTLS {
		s.logger.Printf("[WARN] consul.rpc: Non-TLS connection attempted with VerifyIncoming set")
		conn.Close()
		return
	}

	// Switch on the byte
	switch RPCType(buf[0]) {
	case rpcConsul:
		s.handleConsulConn(conn)

	case rpcRaft:
		metrics.IncrCounter([]string{"consul", "rpc", "raft_handoff"}, 1)
		s.raftLayer.Handoff(conn)

	case rpcMultiplex:
		s.handleMultiplex(conn)

	case rpcTLS:
		if s.rpcTLS == nil {
			s.logger.Printf("[WARN] consul.rpc: TLS connection attempted, server not configured for TLS")
			conn.Close()
			return
		}
		conn = tls.Server(conn, s.rpcTLS)
		s.handleConn(conn, true)

	case rpcMultiplexV2:
		s.handleMultiplexV2(conn)

	default:
		s.logger.Printf("[ERR] consul.rpc: unrecognized RPC byte: %v", buf[0])
		conn.Close()
		return
	}
}

// handleMultiplex is used to multiplex a single incoming connection
// using the Muxado multiplexer
func (s *Server) handleMultiplex(conn net.Conn) {
	defer conn.Close()
	server := muxado.Server(conn)
	for {
		sub, err := server.Accept()
		if err != nil {
			if !strings.Contains(err.Error(), "closed") {
				s.logger.Printf("[ERR] consul.rpc: multiplex conn accept failed: %v", err)
			}
			return
		}
		go s.handleConsulConn(sub)
	}
}

// handleMultiplexV2 is used to multiplex a single incoming connection
// using the Yamux multiplexer
func (s *Server) handleMultiplexV2(conn net.Conn) {
	defer conn.Close()
	conf := yamux.DefaultConfig()
	conf.LogOutput = s.config.LogOutput
	server, _ := yamux.Server(conn, conf)
	for {
		sub, err := server.Accept()
		if err != nil {
			if err != io.EOF {
				s.logger.Printf("[ERR] consul.rpc: multiplex conn accept failed: %v", err)
			}
			return
		}
		go s.handleConsulConn(sub)
	}
}

// handleConsulConn is used to service a single Consul RPC connection
func (s *Server) handleConsulConn(conn net.Conn) {
	defer conn.Close()
	rpcCodec := codec.GoRpc.ServerCodec(conn, msgpackHandle)
	for {
		select {
		case <-s.shutdownCh:
			return
		default:
		}

		if err := s.rpcServer.ServeRequest(rpcCodec); err != nil {
			if err != io.EOF && !strings.Contains(err.Error(), "closed") {
				s.logger.Printf("[ERR] consul.rpc: RPC error: %v (%v)", err, conn)
				metrics.IncrCounter([]string{"consul", "rpc", "request_error"}, 1)
			}
			return
		}
		metrics.IncrCounter([]string{"consul", "rpc", "request"}, 1)
	}
}

// forward is used to forward to a remote DC or to forward to the local leader
// Returns a bool of if forwarding was performed, as well as any error
func (s *Server) forward(method string, info structs.RPCInfo, args interface{}, reply interface{}) (bool, error) {
	// Handle DC forwarding
	dc := info.RequestDatacenter()
	if dc != s.config.Datacenter {
		err := s.forwardDC(method, dc, args, reply)
		return true, err
	}

	// Check if we can allow a stale read
	if info.IsRead() && info.AllowStaleRead() {
		return false, nil
	}

	// Handle leader forwarding
	if !s.IsLeader() {
		err := s.forwardLeader(method, args, reply)
		return true, err
	}
	return false, nil
}

// forwardLeader is used to forward an RPC call to the leader, or fail if no leader
func (s *Server) forwardLeader(method string, args interface{}, reply interface{}) error {
	// Get the leader
	leader := s.raft.Leader()
	if leader == "" {
		return structs.ErrNoLeader
	}

	// Lookup the server
	s.localLock.RLock()
	server := s.localConsuls[leader]
	s.localLock.RUnlock()

	// Handle a missing server
	if server == nil {
		return structs.ErrNoLeader
	}
	return s.connPool.RPC(s.config.Datacenter, server.Addr, server.Version, method, args, reply)
}

// forwardDC is used to forward an RPC call to a remote DC, or fail if no servers
func (s *Server) forwardDC(method, dc string, args interface{}, reply interface{}) error {
	// Bail if we can't find any servers
	s.remoteLock.RLock()
	servers := s.remoteConsuls[dc]
	if len(servers) == 0 {
		s.remoteLock.RUnlock()
		s.logger.Printf("[WARN] consul.rpc: RPC request for DC '%s', no path found", dc)
		return structs.ErrNoDCPath
	}

	// Select a random addr
	offset := rand.Int31() % int32(len(servers))
	server := servers[offset]
	s.remoteLock.RUnlock()

	// Forward to remote Consul
	metrics.IncrCounter([]string{"consul", "rpc", "cross-dc", dc}, 1)
	return s.connPool.RPC(dc, server.Addr, server.Version, method, args, reply)
}

// globalRPC is used to forward an RPC request to one server in each datacenter.
// This will only error for RPC-related errors. Otherwise, application-level
// errors can be sent in the response objects.
func (s *Server) globalRPC(method string, args interface{},
	reply structs.CompoundResponse) error {

	errorCh := make(chan error)
	respCh := make(chan interface{})

	// Make a new request into each datacenter
	for dc, _ := range s.remoteConsuls {
		go func(dc string) {
			rr := reply.New()
			if err := s.forwardDC(method, dc, args, &rr); err != nil {
				errorCh <- err
				return
			}
			respCh <- rr
		}(dc)
	}

	replies, total := 0, len(s.remoteConsuls)
	for replies < total {
		select {
		case err := <-errorCh:
			return err
		case rr := <-respCh:
			reply.Add(rr)
			replies++
		}
	}
	return nil
}

// raftApply is used to encode a message, run it through raft, and return
// the FSM response along with any errors
func (s *Server) raftApply(t structs.MessageType, msg interface{}) (interface{}, error) {
	buf, err := structs.Encode(t, msg)
	if err != nil {
		return nil, fmt.Errorf("Failed to encode request: %v", err)
	}

	// Warn if the command is very large
	if n := len(buf); n > raftWarnSize {
		s.logger.Printf("[WARN] consul: Attempting to apply large raft entry (%d bytes)", n)
	}

	future := s.raft.Apply(buf, enqueueLimit)
	if err := future.Error(); err != nil {
		return nil, err
	}

	return future.Response(), nil
}

// blockingRPC is used for queries that need to wait for a
// minimum index. This is used to block and wait for changes.
func (s *Server) blockingRPC(b *structs.QueryOptions, m *structs.QueryMeta,
	tables MDBTables, run func() error) error {
	opts := blockingRPCOptions{
		queryOpts: b,
		queryMeta: m,
		tables:    tables,
		run:       run,
	}
	return s.blockingRPCOpt(&opts)
}

// blockingRPCOptions is used to parameterize blockingRPCOpt since
// it takes so many options. It should be prefered over blockingRPC.
type blockingRPCOptions struct {
	queryOpts *structs.QueryOptions
	queryMeta *structs.QueryMeta
	tables    MDBTables
	kvWatch   bool
	kvPrefix  string
	run       func() error
}

// blockingRPCOpt is the replacement for blockingRPC as it allows
// for more parameterization easily. It should be prefered over blockingRPC.
func (s *Server) blockingRPCOpt(opts *blockingRPCOptions) error {
	var timeout *time.Timer
	var notifyCh chan struct{}
	var state *StateStore

	// Fast path non-blocking
	if opts.queryOpts.MinQueryIndex == 0 {
		goto RUN_QUERY
	}

	// Sanity check that we have tables to block on
	if len(opts.tables) == 0 && !opts.kvWatch {
		panic("no tables to block on")
	}

	// Restrict the max query time, and ensure there is always one
	if opts.queryOpts.MaxQueryTime > maxQueryTime {
		opts.queryOpts.MaxQueryTime = maxQueryTime
	} else if opts.queryOpts.MaxQueryTime <= 0 {
		opts.queryOpts.MaxQueryTime = maxQueryTime
	}

	// Setup a query timeout
	timeout = time.NewTimer(opts.queryOpts.MaxQueryTime)

	// Setup the notify channel
	notifyCh = make(chan struct{}, 1)

	// Ensure we tear down any watchers on return
	state = s.fsm.State()
	defer func() {
		timeout.Stop()
		state.StopWatch(opts.tables, notifyCh)
		if opts.kvWatch {
			state.StopWatchKV(opts.kvPrefix, notifyCh)
		}
	}()

REGISTER_NOTIFY:
	// Register the notification channel. This may be done
	// multiple times if we have not reached the target wait index.
	state.Watch(opts.tables, notifyCh)
	if opts.kvWatch {
		state.WatchKV(opts.kvPrefix, notifyCh)
	}

RUN_QUERY:
	// Update the query meta data
	s.setQueryMeta(opts.queryMeta)

	// Check if query must be consistent
	if opts.queryOpts.RequireConsistent {
		if err := s.consistentRead(); err != nil {
			return err
		}
	}

	// Run the query function
	metrics.IncrCounter([]string{"consul", "rpc", "query"}, 1)
	err := opts.run()

	// Check for minimum query time
	if err == nil && opts.queryMeta.Index > 0 && opts.queryMeta.Index <= opts.queryOpts.MinQueryIndex {
		select {
		case <-notifyCh:
			goto REGISTER_NOTIFY
		case <-timeout.C:
		}
	}
	return err
}

// setQueryMeta is used to populate the QueryMeta data for an RPC call
func (s *Server) setQueryMeta(m *structs.QueryMeta) {
	if s.IsLeader() {
		m.LastContact = 0
		m.KnownLeader = true
	} else {
		m.LastContact = time.Now().Sub(s.raft.LastContact())
		m.KnownLeader = (s.raft.Leader() != "")
	}
}

// consistentRead is used to ensure we do not perform a stale
// read. This is done by verifying leadership before the read.
func (s *Server) consistentRead() error {
	defer metrics.MeasureSince([]string{"consul", "rpc", "consistentRead"}, time.Now())
	future := s.raft.VerifyLeader()
	return future.Error()
}
<|endoftext|>


<commit_before>package gospell

import (
	"bufio"
	"fmt"
	"io"
	"regexp"
	"strconv"
	"strings"
)

// AffixType is either an affix prefix or suffix
type AffixType int

// specific Affix types
const (
	Prefix AffixType = iota
	Suffix
)

// Affix is a rule for affix (adding prefixes or suffixes)
type Affix struct {
	Type         AffixType // either PFX or SFX
	CrossProduct bool
	Rules        []Rule
}

// Expand provides all variations of a given word based on this affix rule
func (a Affix) Expand(word string, out []string) []string {
	for _, r := range a.Rules {
		if r.Matcher != nil && !r.Matcher.MatchString(word) {
			continue
		}
		if a.Type == Prefix {
			out = append(out, r.AffixText+word)
			// TODO is does Strip apply to prefixes too?
		} else {
			stripWord := word
			if r.Strip != "" && strings.HasSuffix(word, r.Strip) {
				stripWord = word[:len(word)-len(r.Strip)]
			}
			out = append(out, stripWord+r.AffixText)
		}
	}
	return out
}

// Rule is a Affix rule
type Rule struct {
	Strip     string
	AffixText string         // suffix or prefix text to add
	Matcher   *regexp.Regexp // matcher to see if this rule applies or not
}

// AFFFile is a partial representation of a Hunspell AFF file.
type AFFFile struct {
	Flag              string
	TryChars          string
	WordChars         string
	IconvReplacements [][2]string
	Replacements      [][2]string
	AffixMap          map[rune]Affix
	CompoundMin       int
}

// Expand expands a word/affix
func (a AFFFile) Expand(wordAffix string, out []string) ([]string, error) {
	out = out[:0]
	idx := strings.Index(wordAffix, "/")

	// not found
	if idx == -1 {
		out = append(out, wordAffix)
		return out, nil
	}
	if idx == 0 || idx+1 == len(wordAffix) {
		return nil, fmt.Errorf("Slash char found in first or last position")
	}
	// safe
	word, keyString := wordAffix[:idx], wordAffix[idx+1:]
	out = append(out, word)
	prefixes := make([]Affix, 0, 5)
	suffixes := make([]Affix, 0, 5)
	for _, key := range keyString {
		// want keyString to []?something?
		// then iterate over that
		af, ok := a.AffixMap[key]
		if !ok {
			return nil, fmt.Errorf("unable to find affix key %v", key)
		}
		if !af.CrossProduct {
			out = af.Expand(word, out)
			continue
		}
		if af.Type == Prefix {
			prefixes = append(prefixes, af)
		} else {
			suffixes = append(suffixes, af)
		}
	}

	// expand all suffixes with out any prefixes
	for _, suf := range suffixes {
		out = suf.Expand(word, out)
	}
	for _, pre := range prefixes {
		prewords := pre.Expand(word, nil)
		out = append(out, prewords...)

		// now do cross product
		for _, suf := range suffixes {
			for _, w := range prewords {
				out = suf.Expand(w, out)
			}
		}
	}
	return out, nil
}

func isCrossProduct(val string) (bool, error) {
	switch val {
	case "Y":
		return true, nil
	case "N":
		return false, nil
	}
	return false, fmt.Errorf("CrossProduct is not Y or N: got %q", val)
}

// NewAFF reads an Hunspell AFF file
func NewAFF(file io.Reader) (*AFFFile, error) {
	aff := AFFFile{
		Flag:        "ASCII",
		AffixMap:    make(map[rune]Affix),
		CompoundMin: 3, // default in Hunspell
	}
	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		line := scanner.Text()
		parts := strings.Fields(line)
		if len(parts) == 0 {
			continue
		}
		switch parts[0] {
		case "#":
			continue
		case "TRY":
			if len(parts) != 2 {
				return nil, fmt.Errorf("TRY stanza had %d fields, expected 2", len(parts))
			}
			aff.TryChars = parts[1]
		case "ICONV":
			// if only 2 fields, then its the first stanza that just provides a count
			//  we don't care, as we dynamically allocate
			if len(parts) == 2 {
				continue
			}
			if len(parts) != 3 {
				return nil, fmt.Errorf("ICONV stanza had %d fields, expected 2", len(parts))
			}
			// we have 3
			aff.IconvReplacements = append(aff.IconvReplacements, [2]string{parts[1], parts[2]})
		case "REP":
			// if only 2 fields, then its the first stanza that just provides a count
			//  we don't care, as we dynamically allocate
			if len(parts) == 2 {
				continue
			}
			if len(parts) != 3 {
				return nil, fmt.Errorf("REP stanza had %d fields, expected 2", len(parts))
			}
			// we have 3
			aff.Replacements = append(aff.Replacements, [2]string{parts[1], parts[2]})
		case "COMPOUNDMIN":
			if len(parts) != 2 {
				return nil, fmt.Errorf("COMPOUNDMIN stanza had %d fields, expected 2", len(parts))
			}
			val, err := strconv.ParseInt(parts[1], 10, 64)
			if err != nil {
				return nil, fmt.Errorf("COMPOUNDMIN stanza had %q expected number", parts[1])
			}
			aff.CompoundMin = int(val)
		case "WORDCHARS":
			if len(parts) != 2 {
				return nil, fmt.Errorf("WORDCHAR stanza had %d fields, expected 2", len(parts))
			}
			aff.WordChars = parts[1]
		case "FLAG":
			if len(parts) != 2 {
				return nil, fmt.Errorf("FLAG stanza had %d, expected 1", len(parts))
			}
			aff.Flag = parts[1]
			return nil, fmt.Errorf("FLAG stanza not yet supported")
		case "PFX", "SFX":
			atype := Prefix
			if parts[0] == "SFX" {
				atype = Suffix
			}

			switch len(parts) {
			case 4:
				cross, err := isCrossProduct(parts[2])
				if err != nil {
					return nil, err
				}
				// this is a new Affix!
				a := Affix{
					Type:         atype,
					CrossProduct: cross,
				}
				flag := rune(parts[1][0])
				aff.AffixMap[flag] = a
			case 5:
				// does this need to be split out into suffix and prefix?
				flag := rune(parts[1][0])
				a, ok := aff.AffixMap[flag]
				if !ok {
					return nil, fmt.Errorf("Got rules for flag %q but no definition", flag)
				}

				strip := ""
				if parts[2] != "0" {
					strip = parts[2]
				}

				var matcher *regexp.Regexp
				var err error
				pat := parts[4]
				if pat != "." {
					if a.Type == Prefix {
						pat = "^" + pat
					} else {
						pat = pat + "$"
					}
					matcher, err = regexp.Compile(pat)
					if err != nil {
						return nil, fmt.Errorf("Unable to compile %s", pat)
					}
				}

				a.Rules = append(a.Rules, Rule{
					Strip:     strip,
					AffixText: parts[3],
					Matcher:   matcher,
				})
				aff.AffixMap[flag] = a
			default:
				return nil, fmt.Errorf("%s stanza had %d fields, expected 4 or 5", parts[0], len(parts))
			}
		default:
			// nothing
		}
	}

	if err := scanner.Err(); err != nil {
		return nil, err
	}

	return &aff, nil
}
<commit_msg>relax validation for numbers, for now<commit_after>package gospell

import (
	"bufio"
	"fmt"
	"io"
	"regexp"
	"strconv"
	"strings"
)

// AffixType is either an affix prefix or suffix
type AffixType int

// specific Affix types
const (
	Prefix AffixType = iota
	Suffix
)

// Affix is a rule for affix (adding prefixes or suffixes)
type Affix struct {
	Type         AffixType // either PFX or SFX
	CrossProduct bool
	Rules        []Rule
}

// Expand provides all variations of a given word based on this affix rule
func (a Affix) Expand(word string, out []string) []string {
	for _, r := range a.Rules {
		if r.Matcher != nil && !r.Matcher.MatchString(word) {
			continue
		}
		if a.Type == Prefix {
			out = append(out, r.AffixText+word)
			// TODO is does Strip apply to prefixes too?
		} else {
			stripWord := word
			if r.Strip != "" && strings.HasSuffix(word, r.Strip) {
				stripWord = word[:len(word)-len(r.Strip)]
			}
			out = append(out, stripWord+r.AffixText)
		}
	}
	return out
}

// Rule is a Affix rule
type Rule struct {
	Strip     string
	AffixText string         // suffix or prefix text to add
	Matcher   *regexp.Regexp // matcher to see if this rule applies or not
}

// AFFFile is a partial representation of a Hunspell AFF file.
type AFFFile struct {
	Flag              string
	TryChars          string
	WordChars         string
	IconvReplacements [][2]string
	Replacements      [][2]string
	AffixMap          map[rune]Affix
	CompoundMin       int
}

// Expand expands a word/affix
func (a AFFFile) Expand(wordAffix string, out []string) ([]string, error) {
	out = out[:0]
	idx := strings.Index(wordAffix, "/")

	// not found
	if idx == -1 {
		out = append(out, wordAffix)
		return out, nil
	}
	if idx == 0 || idx+1 == len(wordAffix) {
		return nil, fmt.Errorf("Slash char found in first or last position")
	}
	// safe
	word, keyString := wordAffix[:idx], wordAffix[idx+1:]
	out = append(out, word)
	prefixes := make([]Affix, 0, 5)
	suffixes := make([]Affix, 0, 5)
	for _, key := range keyString {
		// want keyString to []?something?
		// then iterate over that
		af, ok := a.AffixMap[key]
		if !ok {
			continue
			//return nil, fmt.Errorf("unable to find affix key %v", key)
		}
		if !af.CrossProduct {
			out = af.Expand(word, out)
			continue
		}
		if af.Type == Prefix {
			prefixes = append(prefixes, af)
		} else {
			suffixes = append(suffixes, af)
		}
	}

	// expand all suffixes with out any prefixes
	for _, suf := range suffixes {
		out = suf.Expand(word, out)
	}
	for _, pre := range prefixes {
		prewords := pre.Expand(word, nil)
		out = append(out, prewords...)

		// now do cross product
		for _, suf := range suffixes {
			for _, w := range prewords {
				out = suf.Expand(w, out)
			}
		}
	}
	return out, nil
}

func isCrossProduct(val string) (bool, error) {
	switch val {
	case "Y":
		return true, nil
	case "N":
		return false, nil
	}
	return false, fmt.Errorf("CrossProduct is not Y or N: got %q", val)
}

// NewAFF reads an Hunspell AFF file
func NewAFF(file io.Reader) (*AFFFile, error) {
	aff := AFFFile{
		Flag:        "ASCII",
		AffixMap:    make(map[rune]Affix),
		CompoundMin: 3, // default in Hunspell
	}
	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		line := scanner.Text()
		parts := strings.Fields(line)
		if len(parts) == 0 {
			continue
		}
		switch parts[0] {
		case "#":
			continue
		case "TRY":
			if len(parts) != 2 {
				return nil, fmt.Errorf("TRY stanza had %d fields, expected 2", len(parts))
			}
			aff.TryChars = parts[1]
		case "ICONV":
			// if only 2 fields, then its the first stanza that just provides a count
			//  we don't care, as we dynamically allocate
			if len(parts) == 2 {
				continue
			}
			if len(parts) != 3 {
				return nil, fmt.Errorf("ICONV stanza had %d fields, expected 2", len(parts))
			}
			// we have 3
			aff.IconvReplacements = append(aff.IconvReplacements, [2]string{parts[1], parts[2]})
		case "REP":
			// if only 2 fields, then its the first stanza that just provides a count
			//  we don't care, as we dynamically allocate
			if len(parts) == 2 {
				continue
			}
			if len(parts) != 3 {
				return nil, fmt.Errorf("REP stanza had %d fields, expected 2", len(parts))
			}
			// we have 3
			aff.Replacements = append(aff.Replacements, [2]string{parts[1], parts[2]})
		case "COMPOUNDMIN":
			if len(parts) != 2 {
				return nil, fmt.Errorf("COMPOUNDMIN stanza had %d fields, expected 2", len(parts))
			}
			val, err := strconv.ParseInt(parts[1], 10, 64)
			if err != nil {
				return nil, fmt.Errorf("COMPOUNDMIN stanza had %q expected number", parts[1])
			}
			aff.CompoundMin = int(val)
		case "WORDCHARS":
			if len(parts) != 2 {
				return nil, fmt.Errorf("WORDCHAR stanza had %d fields, expected 2", len(parts))
			}
			aff.WordChars = parts[1]
		case "FLAG":
			if len(parts) != 2 {
				return nil, fmt.Errorf("FLAG stanza had %d, expected 1", len(parts))
			}
			aff.Flag = parts[1]
			return nil, fmt.Errorf("FLAG stanza not yet supported")
		case "PFX", "SFX":
			atype := Prefix
			if parts[0] == "SFX" {
				atype = Suffix
			}

			switch len(parts) {
			case 4:
				cross, err := isCrossProduct(parts[2])
				if err != nil {
					return nil, err
				}
				// this is a new Affix!
				a := Affix{
					Type:         atype,
					CrossProduct: cross,
				}
				flag := rune(parts[1][0])
				aff.AffixMap[flag] = a
			case 5:
				// does this need to be split out into suffix and prefix?
				flag := rune(parts[1][0])
				a, ok := aff.AffixMap[flag]
				if !ok {
					return nil, fmt.Errorf("Got rules for flag %q but no definition", flag)
				}

				strip := ""
				if parts[2] != "0" {
					strip = parts[2]
				}

				var matcher *regexp.Regexp
				var err error
				pat := parts[4]
				if pat != "." {
					if a.Type == Prefix {
						pat = "^" + pat
					} else {
						pat = pat + "$"
					}
					matcher, err = regexp.Compile(pat)
					if err != nil {
						return nil, fmt.Errorf("Unable to compile %s", pat)
					}
				}

				a.Rules = append(a.Rules, Rule{
					Strip:     strip,
					AffixText: parts[3],
					Matcher:   matcher,
				})
				aff.AffixMap[flag] = a
			default:
				return nil, fmt.Errorf("%s stanza had %d fields, expected 4 or 5", parts[0], len(parts))
			}
		default:
			// nothing
		}
	}

	if err := scanner.Err(); err != nil {
		return nil, err
	}

	return &aff, nil
}
<|endoftext|>


<commit_before>package log

import (
	"bufio"
	"fmt"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/mgutz/ansi"
)

type sourceLine struct {
	lineno int
	line   string
}

type frameInfo struct {
	filename     string
	lineno       int
	method       string
	context      []*sourceLine
	contextLines int
}

func (ci *frameInfo) readSource(contextLines int) error {
	if ci.lineno == 0 || disableCallstack {
		return nil
	}
	start := maxInt(1, ci.lineno-contextLines)
	end := ci.lineno + contextLines

	f, err := os.Open(ci.filename)
	if err != nil {
		// if we can't read a file, it means user is running this in production
		disableCallstack = true
		return err
	}
	defer f.Close()

	lineno := 1
	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		if start <= lineno && lineno <= end {
			line := scanner.Text()
			line = expandTabs(line, 4)
			ci.context = append(ci.context, &sourceLine{lineno: lineno, line: line})
		}
		lineno++
	}

	if err := scanner.Err(); err != nil {
		InternalLog.Warn("scanner error", "file", ci.filename, "err", err)
	}
	return nil
}

func (ci *frameInfo) String(color string, sourceColor string) string {
	buf := pool.Get()
	defer pool.Put(buf)

	if disableCallstack {
		buf.WriteString(color)
		buf.WriteString(Separator)
		buf.WriteString(indent)
		buf.WriteString(ci.filename)
		buf.WriteRune(':')
		buf.WriteString(strconv.Itoa(ci.lineno))
		return buf.String()
	}

	// skip anything in the logxi package
	if isLogxiCode(ci.filename) {
		return ""
	}

	// make path relative to current working directory or home
	tildeFilename, err := filepath.Rel(wd, ci.filename)
	if err != nil {
		InternalLog.Warn("Could not make path relative", "path", ci.filename)
		return ""
	}
	// ../../../ is too complex.  Make path relative to home
	if strings.HasPrefix(tildeFilename, strings.Repeat(".."+string(os.PathSeparator), 3)) {
		tildeFilename = strings.Replace(tildeFilename, home, "~", 1)
	}

	buf.WriteString(color)
	buf.WriteString(Separator)
	buf.WriteString(indent)
	buf.WriteString("in ")
	buf.WriteString(ci.method)
	buf.WriteString("(")
	buf.WriteString(tildeFilename)
	buf.WriteRune(':')
	buf.WriteString(strconv.Itoa(ci.lineno))
	buf.WriteString(")")

	if ci.contextLines == -1 {
		return buf.String()
	}
	buf.WriteString("\n")

	// the width of the printed line number
	var linenoWidth int
	// trim spaces at start of source code based on common spaces
	var skipSpaces = 1000

	// calculate width of lineno and number of leading spaces that can be
	// removed
	for _, li := range ci.context {
		linenoWidth = maxInt(linenoWidth, len(fmt.Sprintf("%d", li.lineno)))
		index := indexOfNonSpace(li.line)
		if index > -1 && index < skipSpaces {
			skipSpaces = index
		}
	}

	for _, li := range ci.context {
		var format string
		format = fmt.Sprintf("%%s%%%dd:  %%s\n", linenoWidth)

		if li.lineno == ci.lineno {
			buf.WriteString(color)
			if ci.contextLines > 2 {
				format = fmt.Sprintf("%%s=> %%%dd:  %%s\n", linenoWidth)
			}
		} else {
			buf.WriteString(sourceColor)
			if ci.contextLines > 2 {
				// account for "=> "
				format = fmt.Sprintf("%%s%%%dd:  %%s\n", linenoWidth+3)
			}
		}
		// trim spaces at start
		idx := minInt(len(li.line), skipSpaces)
		buf.WriteString(fmt.Sprintf(format, Separator+indent+indent, li.lineno, li.line[idx:]))
	}
	// get rid of last \n
	buf.Truncate(buf.Len() - 1)
	buf.WriteString(ansi.Reset)
	return buf.String()
}

// parseDebugStack parases a stack created by debug.Stack()
//
// This is what the string looks like
// /Users/mgutz/go/src/github.com/mgutz/logxi/v1/jsonFormatter.go:45 (0x5fa70)
// 	(*JSONFormatter).writeError: jf.writeString(buf, err.Error()+"\n"+string(debug.Stack()))
// /Users/mgutz/go/src/github.com/mgutz/logxi/v1/jsonFormatter.go:82 (0x5fdc3)
// 	(*JSONFormatter).appendValue: jf.writeError(buf, err)
// /Users/mgutz/go/src/github.com/mgutz/logxi/v1/jsonFormatter.go:109 (0x605ca)
// 	(*JSONFormatter).set: jf.appendValue(buf, val)
// ...
// /Users/mgutz/goroot/src/runtime/asm_amd64.s:2232 (0x38bf1)
// 	goexit:
func parseDebugStack(stack string, skip int, ignoreRuntime bool) []*frameInfo {
	lines := strings.Split(stack, "\n")
	frames := []*frameInfo{}
	for i := skip * 2; i < len(lines); i += 2 {
		ci := &frameInfo{}
		sourceLine := lines[i]
		if sourceLine == "" {
			break
		}
		if ignoreRuntime && strings.Contains(sourceLine, filepath.Join("src", "runtime")) {
			break
		}

		colon := strings.Index(sourceLine, ":")
		slash := strings.Index(sourceLine, "/")
		if colon < slash {
			// must be on Windows where paths look like c:/foo/bar.go:lineno
			colon = strings.Index(sourceLine[slash:], ":") + slash
		}
		space := strings.Index(sourceLine, " ")
		ci.filename = sourceLine[0:colon]
		numstr := sourceLine[colon+1 : space]
		lineno, err := strconv.Atoi(numstr)
		if err != nil {
			InternalLog.Warn("Could not parse line number", "sourceLine", sourceLine, "numstr", numstr)
			continue
		}
		ci.lineno = lineno

		methodLine := lines[i+1]
		colon = strings.Index(methodLine, ":")
		ci.method = strings.Trim(methodLine[0:colon], "\t ")
		frames = append(frames, ci)
	}
	return frames
}

// parseDebugStack parases a stack created by debug.Stack()
//
// This is what the string looks like
// /Users/mgutz/go/src/github.com/mgutz/logxi/v1/jsonFormatter.go:45 (0x5fa70)
// 	(*JSONFormatter).writeError: jf.writeString(buf, err.Error()+"\n"+string(debug.Stack()))
// /Users/mgutz/go/src/github.com/mgutz/logxi/v1/jsonFormatter.go:82 (0x5fdc3)
// 	(*JSONFormatter).appendValue: jf.writeError(buf, err)
// /Users/mgutz/go/src/github.com/mgutz/logxi/v1/jsonFormatter.go:109 (0x605ca)
// 	(*JSONFormatter).set: jf.appendValue(buf, val)
// ...
// /Users/mgutz/goroot/src/runtime/asm_amd64.s:2232 (0x38bf1)
// 	goexit:
func trimDebugStack(stack string) string {
	buf := pool.Get()
	defer pool.Put(buf)
	lines := strings.Split(stack, "\n")
	for i := 0; i < len(lines); i += 2 {
		sourceLine := lines[i]
		if sourceLine == "" {
			break
		}

		colon := strings.Index(sourceLine, ":")
		slash := strings.Index(sourceLine, "/")
		if colon < slash {
			// must be on Windows where paths look like c:/foo/bar.go:lineno
			colon = strings.Index(sourceLine[slash:], ":") + slash
		}
		filename := sourceLine[0:colon]
		// skip anything in the logxi package
		if isLogxiCode(filename) {
			continue
		}
		buf.WriteString(sourceLine)
		buf.WriteRune('\n')
		buf.WriteString(lines[i+1])
		buf.WriteRune('\n')
	}
	return buf.String()
}

func parseLogxiStack(entry map[string]interface{}, skip int, ignoreRuntime bool) []*frameInfo {
	kv := entry[KeyMap.CallStack]
	if kv == nil {
		return nil
	}

	var frames []*frameInfo
	if stack, ok := kv.(string); ok {
		frames = parseDebugStack(stack, skip, ignoreRuntime)
	}
	return frames
}
<commit_msg>temporarily disable custom stack parsing; bug with embedded newlines<commit_after>package log

import (
	"bufio"
	"fmt"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/mgutz/ansi"
)

type sourceLine struct {
	lineno int
	line   string
}

type frameInfo struct {
	filename     string
	lineno       int
	method       string
	context      []*sourceLine
	contextLines int
}

func (ci *frameInfo) readSource(contextLines int) error {
	if ci.lineno == 0 || disableCallstack {
		return nil
	}
	start := maxInt(1, ci.lineno-contextLines)
	end := ci.lineno + contextLines

	f, err := os.Open(ci.filename)
	if err != nil {
		// if we can't read a file, it means user is running this in production
		disableCallstack = true
		return err
	}
	defer f.Close()

	lineno := 1
	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		if start <= lineno && lineno <= end {
			line := scanner.Text()
			line = expandTabs(line, 4)
			ci.context = append(ci.context, &sourceLine{lineno: lineno, line: line})
		}
		lineno++
	}

	if err := scanner.Err(); err != nil {
		InternalLog.Warn("scanner error", "file", ci.filename, "err", err)
	}
	return nil
}

func (ci *frameInfo) String(color string, sourceColor string) string {
	buf := pool.Get()
	defer pool.Put(buf)

	if disableCallstack {
		buf.WriteString(color)
		buf.WriteString(Separator)
		buf.WriteString(indent)
		buf.WriteString(ci.filename)
		buf.WriteRune(':')
		buf.WriteString(strconv.Itoa(ci.lineno))
		return buf.String()
	}

	// skip anything in the logxi package
	if isLogxiCode(ci.filename) {
		return ""
	}

	// make path relative to current working directory or home
	tildeFilename, err := filepath.Rel(wd, ci.filename)
	if err != nil {
		InternalLog.Warn("Could not make path relative", "path", ci.filename)
		return ""
	}
	// ../../../ is too complex.  Make path relative to home
	if strings.HasPrefix(tildeFilename, strings.Repeat(".."+string(os.PathSeparator), 3)) {
		tildeFilename = strings.Replace(tildeFilename, home, "~", 1)
	}

	buf.WriteString(color)
	buf.WriteString(Separator)
	buf.WriteString(indent)
	buf.WriteString("in ")
	buf.WriteString(ci.method)
	buf.WriteString("(")
	buf.WriteString(tildeFilename)
	buf.WriteRune(':')
	buf.WriteString(strconv.Itoa(ci.lineno))
	buf.WriteString(")")

	if ci.contextLines == -1 {
		return buf.String()
	}
	buf.WriteString("\n")

	// the width of the printed line number
	var linenoWidth int
	// trim spaces at start of source code based on common spaces
	var skipSpaces = 1000

	// calculate width of lineno and number of leading spaces that can be
	// removed
	for _, li := range ci.context {
		linenoWidth = maxInt(linenoWidth, len(fmt.Sprintf("%d", li.lineno)))
		index := indexOfNonSpace(li.line)
		if index > -1 && index < skipSpaces {
			skipSpaces = index
		}
	}

	for _, li := range ci.context {
		var format string
		format = fmt.Sprintf("%%s%%%dd:  %%s\n", linenoWidth)

		if li.lineno == ci.lineno {
			buf.WriteString(color)
			if ci.contextLines > 2 {
				format = fmt.Sprintf("%%s=> %%%dd:  %%s\n", linenoWidth)
			}
		} else {
			buf.WriteString(sourceColor)
			if ci.contextLines > 2 {
				// account for "=> "
				format = fmt.Sprintf("%%s%%%dd:  %%s\n", linenoWidth+3)
			}
		}
		// trim spaces at start
		idx := minInt(len(li.line), skipSpaces)
		buf.WriteString(fmt.Sprintf(format, Separator+indent+indent, li.lineno, li.line[idx:]))
	}
	// get rid of last \n
	buf.Truncate(buf.Len() - 1)
	buf.WriteString(ansi.Reset)
	return buf.String()
}

// parseDebugStack parases a stack created by debug.Stack()
//
// This is what the string looks like
// /Users/mgutz/go/src/github.com/mgutz/logxi/v1/jsonFormatter.go:45 (0x5fa70)
// 	(*JSONFormatter).writeError: jf.writeString(buf, err.Error()+"\n"+string(debug.Stack()))
// /Users/mgutz/go/src/github.com/mgutz/logxi/v1/jsonFormatter.go:82 (0x5fdc3)
// 	(*JSONFormatter).appendValue: jf.writeError(buf, err)
// /Users/mgutz/go/src/github.com/mgutz/logxi/v1/jsonFormatter.go:109 (0x605ca)
// 	(*JSONFormatter).set: jf.appendValue(buf, val)
// ...
// /Users/mgutz/goroot/src/runtime/asm_amd64.s:2232 (0x38bf1)
// 	goexit:
func parseDebugStack(stack string, skip int, ignoreRuntime bool) []*frameInfo {
	frames := []*frameInfo{}
	// BUG temporarily disable since there is a bug with embedded newlines
	if true {
		return frames
	}

	lines := strings.Split(stack, "\n")

	for i := skip * 2; i < len(lines); i += 2 {
		ci := &frameInfo{}
		sourceLine := lines[i]
		if sourceLine == "" {
			break
		}
		if ignoreRuntime && strings.Contains(sourceLine, filepath.Join("src", "runtime")) {
			break
		}

		colon := strings.Index(sourceLine, ":")
		slash := strings.Index(sourceLine, "/")
		if colon < slash {
			// must be on Windows where paths look like c:/foo/bar.go:lineno
			colon = strings.Index(sourceLine[slash:], ":") + slash
		}
		space := strings.Index(sourceLine, " ")
		ci.filename = sourceLine[0:colon]

		// BUG with callstack where the error message has embedded newlines
		// if colon > space {
		// 	fmt.Println("lines", lines)
		// }
		// fmt.Println("SOURCELINE", sourceLine, "len", len(sourceLine), "COLON", colon, "SPACE", space)
		numstr := sourceLine[colon+1 : space]
		lineno, err := strconv.Atoi(numstr)
		if err != nil {
			InternalLog.Warn("Could not parse line number", "sourceLine", sourceLine, "numstr", numstr)
			continue
		}
		ci.lineno = lineno

		methodLine := lines[i+1]
		colon = strings.Index(methodLine, ":")
		ci.method = strings.Trim(methodLine[0:colon], "\t ")
		frames = append(frames, ci)
	}
	return frames
}

// parseDebugStack parases a stack created by debug.Stack()
//
// This is what the string looks like
// /Users/mgutz/go/src/github.com/mgutz/logxi/v1/jsonFormatter.go:45 (0x5fa70)
// 	(*JSONFormatter).writeError: jf.writeString(buf, err.Error()+"\n"+string(debug.Stack()))
// /Users/mgutz/go/src/github.com/mgutz/logxi/v1/jsonFormatter.go:82 (0x5fdc3)
// 	(*JSONFormatter).appendValue: jf.writeError(buf, err)
// /Users/mgutz/go/src/github.com/mgutz/logxi/v1/jsonFormatter.go:109 (0x605ca)
// 	(*JSONFormatter).set: jf.appendValue(buf, val)
// ...
// /Users/mgutz/goroot/src/runtime/asm_amd64.s:2232 (0x38bf1)
// 	goexit:
func trimDebugStack(stack string) string {
	buf := pool.Get()
	defer pool.Put(buf)
	lines := strings.Split(stack, "\n")
	for i := 0; i < len(lines); i += 2 {
		sourceLine := lines[i]
		if sourceLine == "" {
			break
		}

		colon := strings.Index(sourceLine, ":")
		slash := strings.Index(sourceLine, "/")
		if colon < slash {
			// must be on Windows where paths look like c:/foo/bar.go:lineno
			colon = strings.Index(sourceLine[slash:], ":") + slash
		}
		filename := sourceLine[0:colon]
		// skip anything in the logxi package
		if isLogxiCode(filename) {
			continue
		}
		buf.WriteString(sourceLine)
		buf.WriteRune('\n')
		buf.WriteString(lines[i+1])
		buf.WriteRune('\n')
	}
	return buf.String()
}

func parseLogxiStack(entry map[string]interface{}, skip int, ignoreRuntime bool) []*frameInfo {
	kv := entry[KeyMap.CallStack]
	if kv == nil {
		return nil
	}

	var frames []*frameInfo
	if stack, ok := kv.(string); ok {
		frames = parseDebugStack(stack, skip, ignoreRuntime)
	}
	return frames
}
<|endoftext|>


<commit_before>// Package tree implements a basic balanced binary tree
package tree

import (
	"encoding/json"
	"fmt"
	"math/rand"
	"sync"
	"time"
)

// Node is a fundemental part of what makes a tree a tree. Many Nodes creates a tree
type Node struct {
	Left  *Node
	Right *Node
	Data  int
}

// Tree basic tree structure
type Tree struct {
	Root      *Node
	Total     int
	NodeCount int
}

var (
	// ErrPositiveIntegers reports that only positive intergers may be added to the tree
	ErrPositiveIntegers = fmt.Errorf("only postive integers may be added")
	// ErrNodeNotFound reports that a Node wasn't found
	ErrNodeNotFound = fmt.Errorf("Node not found")
)

// NewTree creates a pointer to the Tree struct
func NewTree() *Tree {
	return new(Tree)
}

// FindNode recursively looks for the node with the specified value
func (t *Tree) FindNode(data int) (err error) {
	newNode := Node{
		Data: data,
	}
	if t.Root != nil {
		if t.findNode(t.Root, newNode) != nil {
			return
		}
	}
	return ErrNodeNotFound
}

func (t *Tree) findNode(search *Node, target Node) *Node {
	var returnNode *Node
	if search == nil {
		return returnNode
	}
	if search.Data == target.Data {
		return search
	}
	returnNode = t.findNode(search.Left, target)
	if returnNode == nil {
		returnNode = t.findNode(search.Right, target)
	}
	return returnNode
}

// Add appends a new node to a branch in a balanced manner
func (t *Tree) Add(data int) (err error) {
	t.Total += data
	t.NodeCount++
	if data < 0 {
		return ErrPositiveIntegers
	}
	NodeToAdd := Node{
		Data: data,
	}
	if t.Root == nil {
		t.Root = new(Node)
	}
	if t.Root.Data == 0 {
		t.Root = &NodeToAdd
		return
	}
	t.add(t.Root, NodeToAdd)
	return
}

func (t *Tree) add(oldNode *Node, newNode Node) {
	if newNode.Data < oldNode.Data {
		if oldNode.Left == nil {
			oldNode.Left = &newNode
		} else {
			t.add(oldNode.Left, newNode)
		}
	} else if newNode.Data > oldNode.Data {
		if oldNode.Right == nil {
			oldNode.Right = &newNode
		} else {
			t.add(oldNode.Right, newNode)
		}
	}
	return
}

// InOrderTraversal prints out the values in order
func (t *Tree) InOrderTraversal() {
	if t.Root != nil {
		currentNode := t.Root
		if currentNode.Left == nil && currentNode.Right == nil {
			fmt.Println(currentNode.Data)
		} else {
			t.inOrderTraversal(currentNode)
		}
	}
	return
}

func (t *Tree) inOrderTraversal(n *Node) {
	if n.Left != nil {
		t.inOrderTraversal(n.Left)
	}
	fmt.Println(n.Data)
	if n.Right != nil {
		t.inOrderTraversal(n.Right)
	}
	return
}

// Traversal prints out the values by branch side, left, right, ect...
func (t *Tree) Traversal() {
	if t.Root != nil {
		currentNode := t.Root
		if currentNode.Left == nil && currentNode.Right == nil {
			fmt.Println(currentNode.Data)
		} else {
			t.traversal(currentNode)
		}
	}
	return
}

func (t *Tree) traversal(n *Node) {
	fmt.Println(n.Data)
	if n.Left != nil {
		t.traversal(n.Left)
	}
	if n.Right != nil {
		t.traversal(n.Right)
	}
	return
}

// Sum added up all the values stored in the Nodes.. It is a redundant function because total value is kept as a Tree
// value
func (t *Tree) Sum() (total int) {
	var wg sync.WaitGroup
	c := make(chan int, 100)
	if t.Root != nil {
		currentNode := t.Root
		if currentNode.Left == nil && currentNode.Right == nil {
			return 1
		}
		wg.Add(1)
		t.sum(currentNode, c, &wg)
	}
	go func() {
		wg.Wait()
		close(c)
	}()
	for n := range c {
		total += n
	}
	return total
}

func (t *Tree) sum(n *Node, counter chan int, wg *sync.WaitGroup) {
	defer wg.Done()
	if n.Left != nil {
		wg.Add(1)
		go t.sum(n.Left, counter, wg)
	}
	counter <- n.Data
	if n.Right != nil {
		wg.Add(1)
		go t.sum(n.Right, counter, wg)
	}
	return
}

// CountEdges returns the number of edges the tree contains
func (t *Tree) CountEdges() (edges int) {
	var wg sync.WaitGroup
	c := make(chan int, 100)
	if t.Root != nil {
		currentNode := t.Root
		if currentNode.Left == nil && currentNode.Right == nil {
			return 1
		}
		wg.Add(1)
		t.countEdges(currentNode, c, &wg)
	}

	go func() {
		wg.Wait()
		close(c)
	}()
	for n := range c {
		edges += n
	}
	return edges
}

func (t *Tree) countEdges(n *Node, counter chan int, wg *sync.WaitGroup) {
	defer wg.Done()
	if n.Left != nil {
		wg.Add(1)
		go t.countEdges(n.Left, counter, wg)
	}
	counter <- 1
	if n.Right != nil {
		wg.Add(1)
		go t.countEdges(n.Right, counter, wg)
	}
	return
}

// GenerateRandomTree uses time (time.Now().Unix()) to create enthorpy for a source of random numbers to append to the Tree
func (t *Tree) GenerateRandomTree(numberOfNodesToCreate int) (err error) {
	if numberOfNodesToCreate < 0 {
		return ErrPositiveIntegers
	}
	u := time.Now()
	source := rand.NewSource(u.Unix())
	r := rand.New(source)
	arr := r.Perm(numberOfNodesToCreate)
	for _, a := range arr {
		t.Add(a)
	}
	return
}

// GetRootData returns the data stored at the root, however this does not return the root Node
func (t *Tree) GetRootData() int {
	return t.Root.Data
}

// GetTreeTotal returns the sum of the collective nodes on the Tree
func (t *Tree) GetTreeTotal() int {
	return t.Total
}

// TreeToArray converts to the into an int slice
func (t *Tree) TreeToArray() []int {
	var wg sync.WaitGroup
	c := make(chan int, 100)
	arr := make([]int, 0, t.NodeCount)
	if t.Root != nil {
		currentNode := t.Root
		if currentNode.Left == nil && currentNode.Right == nil {
			return []int{currentNode.Data}
		}
		wg.Add(1)
		t.traversalGetVals(currentNode, c, &wg)
	}
	go func() {
		wg.Wait()
		close(c)
	}()
	for n := range c {
		arr = append(arr, n)
	}
	return arr
}

func (t *Tree) traversalGetVals(n *Node, c chan int, wg *sync.WaitGroup) {
	defer wg.Done()
	if n.Left != nil {
		c <- n.Left.Data
		wg.Add(1)
		go t.traversalGetVals(n.Left, c, wg)
	}
	if n.Right != nil {
		c <- n.Right.Data
		wg.Add(1)
		go t.traversalGetVals(n.Right, c, wg)
	}
	return
}

// ShiftRoot rebuilds the tree with a new root
func (t *Tree) ShiftRoot(newRoot int) {
	arr := t.TreeToArray()
	n := Tree{}
	n.Add(newRoot)
	for _, i := range arr {
		n.Add(i)
	}
	*t = n
}

// PrintTree uses json.MarshalIndent() to print the Tree in an organized fashion, which can then be analysized as a JSON
// object
func (t *Tree) PrintTree() {
	b, err := json.MarshalIndent(t, "", " ")
	if err != nil {
		panic(err)
	}
	fmt.Println(string(b))
}
<commit_msg>godocs<commit_after>// Package tree implements a basic balanced binary tree
package tree

import (
	"encoding/json"
	"fmt"
	"math/rand"
	"sync"
	"time"
)

// Node is a fundemental part of what makes a tree a tree. Many Nodes creates a tree
type Node struct {
	Left  *Node
	Right *Node
	Data  int
}

// Tree basic tree structure
type Tree struct {
	Root      *Node
	Total     int
	NodeCount int
}

var (
	// ErrPositiveIntegers reports that only positive intergers may be added to the tree
	ErrPositiveIntegers = fmt.Errorf("only postive integers may be added")
	// ErrNodeNotFound reports that a Node wasn't found
	ErrNodeNotFound = fmt.Errorf("Node not found")
)

// NewTree creates a pointer to the Tree struct
func NewTree() *Tree {
	return new(Tree)
}

// FindNode recursively looks for the Node with the specified value
func (t *Tree) FindNode(data int) (err error) {
	newNode := Node{
		Data: data,
	}
	if t.Root != nil {
		if t.findNode(t.Root, newNode) != nil {
			return
		}
	}
	return ErrNodeNotFound
}

func (t *Tree) findNode(search *Node, target Node) *Node {
	var returnNode *Node
	if search == nil {
		return returnNode
	}
	if search.Data == target.Data {
		return search
	}
	returnNode = t.findNode(search.Left, target)
	if returnNode == nil {
		returnNode = t.findNode(search.Right, target)
	}
	return returnNode
}

// Add appends a new Node to a branch in a balanced manner
func (t *Tree) Add(data int) (err error) {
	t.Total += data
	t.NodeCount++
	if data < 0 {
		return ErrPositiveIntegers
	}
	NodeToAdd := Node{
		Data: data,
	}
	if t.Root == nil {
		t.Root = new(Node)
	}
	if t.Root.Data == 0 {
		t.Root = &NodeToAdd
		return
	}
	t.add(t.Root, NodeToAdd)
	return
}

func (t *Tree) add(oldNode *Node, newNode Node) {
	if newNode.Data < oldNode.Data {
		if oldNode.Left == nil {
			oldNode.Left = &newNode
		} else {
			t.add(oldNode.Left, newNode)
		}
	} else if newNode.Data > oldNode.Data {
		if oldNode.Right == nil {
			oldNode.Right = &newNode
		} else {
			t.add(oldNode.Right, newNode)
		}
	}
	return
}

// InOrderTraversal prints out the values in order
func (t *Tree) InOrderTraversal() {
	if t.Root != nil {
		currentNode := t.Root
		if currentNode.Left == nil && currentNode.Right == nil {
			fmt.Println(currentNode.Data)
		} else {
			t.inOrderTraversal(currentNode)
		}
	}
	return
}

func (t *Tree) inOrderTraversal(n *Node) {
	if n.Left != nil {
		t.inOrderTraversal(n.Left)
	}
	fmt.Println(n.Data)
	if n.Right != nil {
		t.inOrderTraversal(n.Right)
	}
	return
}

// Traversal prints out the values by branch side, left, right, ect...
func (t *Tree) Traversal() {
	if t.Root != nil {
		currentNode := t.Root
		if currentNode.Left == nil && currentNode.Right == nil {
			fmt.Println(currentNode.Data)
		} else {
			t.traversal(currentNode)
		}
	}
	return
}

func (t *Tree) traversal(n *Node) {
	fmt.Println(n.Data)
	if n.Left != nil {
		t.traversal(n.Left)
	}
	if n.Right != nil {
		t.traversal(n.Right)
	}
	return
}

// Sum added up all the values stored in the Nodes.. It is a redundant function because total value is kept as a Tree
// value
func (t *Tree) Sum() (total int) {
	var wg sync.WaitGroup
	c := make(chan int, 100)
	if t.Root != nil {
		currentNode := t.Root
		if currentNode.Left == nil && currentNode.Right == nil {
			return 1
		}
		wg.Add(1)
		t.sum(currentNode, c, &wg)
	}
	go func() {
		wg.Wait()
		close(c)
	}()
	for n := range c {
		total += n
	}
	return total
}

func (t *Tree) sum(n *Node, counter chan int, wg *sync.WaitGroup) {
	defer wg.Done()
	if n.Left != nil {
		wg.Add(1)
		go t.sum(n.Left, counter, wg)
	}
	counter <- n.Data
	if n.Right != nil {
		wg.Add(1)
		go t.sum(n.Right, counter, wg)
	}
	return
}

// CountEdges returns the number of edges the tree contains
func (t *Tree) CountEdges() (edges int) {
	var wg sync.WaitGroup
	c := make(chan int, 100)
	if t.Root != nil {
		currentNode := t.Root
		if currentNode.Left == nil && currentNode.Right == nil {
			return 1
		}
		wg.Add(1)
		t.countEdges(currentNode, c, &wg)
	}

	go func() {
		wg.Wait()
		close(c)
	}()
	for n := range c {
		edges += n
	}
	return edges
}

func (t *Tree) countEdges(n *Node, counter chan int, wg *sync.WaitGroup) {
	defer wg.Done()
	if n.Left != nil {
		wg.Add(1)
		go t.countEdges(n.Left, counter, wg)
	}
	counter <- 1
	if n.Right != nil {
		wg.Add(1)
		go t.countEdges(n.Right, counter, wg)
	}
	return
}

// GenerateRandomTree uses time (time.Now().Unix()) to create enthorpy for a source of random numbers to append to the Tree
func (t *Tree) GenerateRandomTree(numberOfNodesToCreate int) (err error) {
	if numberOfNodesToCreate < 0 {
		return ErrPositiveIntegers
	}
	u := time.Now()
	source := rand.NewSource(u.Unix())
	r := rand.New(source)
	arr := r.Perm(numberOfNodesToCreate)
	for _, a := range arr {
		t.Add(a)
	}
	return
}

// GetRootData returns the data stored at the root, however this does not return the root Node
func (t *Tree) GetRootData() int {
	return t.Root.Data
}

// GetTreeTotal returns the sum of the collective nodes on the Tree
func (t *Tree) GetTreeTotal() int {
	return t.Total
}

// TreeToArray converts to Tree into an int slice
func (t *Tree) TreeToArray() []int {
	var wg sync.WaitGroup
	c := make(chan int, 100)
	arr := make([]int, 0, t.NodeCount)
	if t.Root != nil {
		currentNode := t.Root
		if currentNode.Left == nil && currentNode.Right == nil {
			return []int{currentNode.Data}
		}
		wg.Add(1)
		t.traversalGetVals(currentNode, c, &wg)
	}
	go func() {
		wg.Wait()
		close(c)
	}()
	for n := range c {
		arr = append(arr, n)
	}
	return arr
}

func (t *Tree) traversalGetVals(n *Node, c chan int, wg *sync.WaitGroup) {
	defer wg.Done()
	if n.Left != nil {
		c <- n.Left.Data
		wg.Add(1)
		go t.traversalGetVals(n.Left, c, wg)
	}
	if n.Right != nil {
		c <- n.Right.Data
		wg.Add(1)
		go t.traversalGetVals(n.Right, c, wg)
	}
	return
}

// ShiftRoot rebuilds the tree with a new root
func (t *Tree) ShiftRoot(newRoot int) {
	arr := t.TreeToArray()
	n := Tree{}
	n.Add(newRoot)
	for _, i := range arr {
		n.Add(i)
	}
	*t = n
}

// PrintTree uses json.MarshalIndent() to print the Tree in an organized fashion, which can then be analysized as a JSON
// object
func (t *Tree) PrintTree() {
	b, err := json.MarshalIndent(t, "", " ")
	if err != nil {
		panic(err)
	}
	fmt.Println(string(b))
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/rsa"
	"encoding/json"
	"net/http"

	"github.com/gorilla/mux"
	redisocket "github.com/syhlion/redisocket.v2"
)

func GetAllChannel(rsender *redisocket.Sender) func(w http.ResponseWriter, r *http.Request) {

	return func(w http.ResponseWriter, r *http.Request) {
		params := mux.Vars(r)
		channels, err := rsender.GetChannels(listenChannelPrefix, params["app_key"], "*")
		if err != nil {
			logger.GetRequestEntry(r).WithError(err).Warn("get redis error")
			w.WriteHeader(400)
			w.Write([]byte("get redis error"))
		}
		b, err := json.Marshal(channels)
		if err != nil {
			logger.GetRequestEntry(r).WithError(err).Warn("json marshal error")
			w.WriteHeader(400)
			w.Write([]byte("json marshal error"))
		}
		w.Header().Set("Content-Type", "application/json;charset=UTF-8")
		w.WriteHeader(http.StatusOK)
		w.Write(b)
		return
	}
}
func GetOnlineByChannel(rsender *redisocket.Sender) func(w http.ResponseWriter, r *http.Request) {

	return func(w http.ResponseWriter, r *http.Request) {
		params := mux.Vars(r)
		online, err := rsender.GetOnlineByChannel(listenChannelPrefix, params["app_key"], params["channel"])
		if err != nil {
			logger.GetRequestEntry(r).WithError(err).Warn("get redis error")
			w.WriteHeader(400)
			w.Write([]byte("get redis error"))
		}
		b, err := json.Marshal(online)
		if err != nil {
			logger.GetRequestEntry(r).WithError(err).Warn("json marshal error")
			w.WriteHeader(400)
			w.Write([]byte("json marshal error"))
		}
		w.Header().Set("Content-Type", "application/json;charset=UTF-8")
		w.WriteHeader(http.StatusOK)
		w.Write(b)
		return
	}
}
func GetOnline(rsender *redisocket.Sender) func(w http.ResponseWriter, r *http.Request) {

	return func(w http.ResponseWriter, r *http.Request) {
		params := mux.Vars(r)
		online, err := rsender.GetOnline(listenChannelPrefix, params["app_key"])
		if err != nil {
			logger.GetRequestEntry(r).WithError(err).Warn("get redis error")
			w.WriteHeader(400)
			w.Write([]byte("get redis error"))
		}
		b, err := json.Marshal(online)
		if err != nil {
			logger.GetRequestEntry(r).WithError(err).Warn("json marshal error")
			w.WriteHeader(400)
			w.Write([]byte("json marshal error"))
		}
		w.Header().Set("Content-Type", "application/json;charset=UTF-8")
		w.WriteHeader(http.StatusOK)
		w.Write(b)
		return
	}
}
func PushToSocket(rsender *redisocket.Sender) func(w http.ResponseWriter, r *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		params := mux.Vars(r)
		app_key := params["app_key"]
		if app_key == "" {
			logger.GetRequestEntry(r).Warn("empty param")
			w.WriteHeader(400)
			w.Write([]byte("empty param"))
			return
		}
		socket_id := params["socket_id"]
		data := r.FormValue("data")
		j := JsonCheck(data)
		rsender.PushToUid(listenChannelPrefix, app_key, socket_id, j)
		push := struct {
			SocketId string      `json:"socket_id"`
			Data     interface{} `json:"data"`
		}{
			SocketId: socket_id,
			Data:     data,
		}
		d, err := json.Marshal(push)
		if err != nil {
			logger.GetRequestEntry(r).Warn(err)
			w.WriteHeader(400)
			w.Write([]byte("data error"))
			return
		}
		w.Header().Set("Content-Type", "application/json;charset=UTF-8")
		w.WriteHeader(http.StatusOK)
		w.Write(d)
		return
	}
}
func PushToUser(rsender *redisocket.Sender) func(w http.ResponseWriter, r *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		params := mux.Vars(r)
		app_key := params["app_key"]
		if app_key == "" {
			logger.GetRequestEntry(r).Warn("empty param")
			w.WriteHeader(400)
			w.Write([]byte("empty param"))
			return
		}
		user_id := params["user_id"]
		data := r.FormValue("data")
		j := JsonCheck(data)
		rsender.PushToUid(listenChannelPrefix, app_key, user_id, j)
		push := struct {
			UserId string      `json:"user_id"`
			Data   interface{} `json:"data"`
		}{
			UserId: user_id,
			Data:   data,
		}
		d, err := json.Marshal(push)
		if err != nil {
			logger.GetRequestEntry(r).Warn(err)
			w.WriteHeader(400)
			w.Write([]byte("data error"))
			return
		}
		w.Header().Set("Content-Type", "application/json;charset=UTF-8")
		w.WriteHeader(http.StatusOK)
		w.Write(d)
		return
	}
}

func PushBatchMessage(rsender *redisocket.Sender) func(w http.ResponseWriter, r *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		params := mux.Vars(r)
		app_key := params["app_key"]
		if app_key == "" {
			logger.GetRequestEntry(r).Warn("empty param")
			w.WriteHeader(400)
			w.Write([]byte("empty param"))
			return
		}
		data := r.FormValue("batch_data")
		if data == "" {
			logger.GetRequestEntry(r).Warn("empty batch data")
			w.WriteHeader(400)
			w.Write([]byte("empty batch data"))
			return
		}
		batchData := make([]BatchData, 0)
		byteData := []byte(data)
		err := json.Unmarshal(byteData, &batchData)
		if err != nil {
			logger.GetRequestEntry(r).Warn(err)
			w.WriteHeader(400)
			w.Write([]byte("data error"))
			return
		}
		l := len(batchData)
		bd := make([]redisocket.BatchData, l)
		for _, data := range batchData {
			push := struct {
				Channel string      `json:"channel"`
				Event   string      `json:"event"`
				Data    interface{} `json:"data"`
			}{
				Channel: data.Channel,
				Event:   data.Event,
				Data:    data.Data,
			}
			d, err := json.Marshal(push)
			if err != nil {
				logger.GetRequestEntry(r).Warn(err)
				continue
			}
			b := redisocket.BatchData{
				Data:  d,
				Event: data.Channel,
			}
			bd = append(bd, b)

		}
		rsender.PushBatch(listenChannelPrefix, app_key, bd)
		response := struct {
			Total int `json:"total"`
			Cap   int `json:"cap"`
		}{
			Total: len(batchData),
			Cap:   len(byteData),
		}
		d, err := json.Marshal(response)
		if err != nil {
			w.WriteHeader(400)
			w.Write([]byte("response marshal error"))
			return
		}
		w.Header().Set("Content-Type", "application/json;charset=UTF-8")
		w.WriteHeader(http.StatusOK)
		w.Write(d)
		return
	}
}

func PushMessage(rsender *redisocket.Sender) func(w http.ResponseWriter, r *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		params := mux.Vars(r)
		app_key := params["app_key"]
		channel := params["channel"]
		event := params["event"]
		if app_key == "" || channel == "" || event == "" {
			logger.GetRequestEntry(r).Warn("empty param")
			w.WriteHeader(400)
			w.Write([]byte("empty param"))
			return
		}

		data := r.FormValue("data")
		if data == "" {
			logger.GetRequestEntry(r).Warn("empty data")
			w.WriteHeader(400)
			w.Write([]byte("empty data"))
			return
		}
		jsonData := JsonCheck(data)

		push := struct {
			Channel string      `json:"channel"`
			Event   string      `json:"event"`
			Data    interface{} `json:"data"`
		}{
			Channel: channel,
			Event:   event,
			Data:    jsonData,
		}
		d, err := json.Marshal(push)
		if err != nil {
			logger.GetRequestEntry(r).Warn(err)
			w.WriteHeader(400)
			w.Write([]byte("data error"))
			return
		}
		_, err = rsender.Push(listenChannelPrefix, app_key, channel, d)
		if err != nil {
			logger.GetRequestEntry(r).Warn(err)
			w.WriteHeader(400)
			w.Write([]byte("data error"))
			return
		}
		w.Header().Set("Content-Type", "application/json;charset=UTF-8")
		w.WriteHeader(http.StatusOK)
		w.Write(d)
		return
	}
}
func DecodeJWT(key *rsa.PublicKey) func(w http.ResponseWriter, r *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {

		data := r.FormValue("data")
		auth, err := Decode(key, data)
		if err != nil {
			logger.GetRequestEntry(r).Warnf("error:%s, post data:%s", err, data)
			w.WriteHeader(400)
			w.Write([]byte("data error"))
			return
		}
		if err = json.NewEncoder(w).Encode(auth); err != nil {
			logger.GetRequestEntry(r).Warnf("error:%s", err)
			w.WriteHeader(400)
			w.Write([]byte("parse error"))
		}
		return
	}
}
<commit_msg>fix master api<commit_after>package main

import (
	"crypto/rsa"
	"encoding/json"
	"net/http"

	"github.com/gorilla/mux"
	redisocket "github.com/syhlion/redisocket.v2"
)

func GetAllChannel(rsender *redisocket.Sender) func(w http.ResponseWriter, r *http.Request) {

	return func(w http.ResponseWriter, r *http.Request) {
		params := mux.Vars(r)
		channels, err := rsender.GetChannels(listenChannelPrefix, params["app_key"], "*")
		if err != nil {
			logger.GetRequestEntry(r).WithError(err).Warn("get redis error")
			w.WriteHeader(400)
			w.Write([]byte("get redis error"))
		}
		b, err := json.Marshal(channels)
		if err != nil {
			logger.GetRequestEntry(r).WithError(err).Warn("json marshal error")
			w.WriteHeader(400)
			w.Write([]byte("json marshal error"))
		}
		w.Header().Set("Content-Type", "application/json;charset=UTF-8")
		w.WriteHeader(http.StatusOK)
		w.Write(b)
		return
	}
}
func GetOnlineByChannel(rsender *redisocket.Sender) func(w http.ResponseWriter, r *http.Request) {

	return func(w http.ResponseWriter, r *http.Request) {
		params := mux.Vars(r)
		online, err := rsender.GetOnlineByChannel(listenChannelPrefix, params["app_key"], params["channel"])
		if err != nil {
			logger.GetRequestEntry(r).WithError(err).Warn("get redis error")
			w.WriteHeader(400)
			w.Write([]byte("get redis error"))
		}
		b, err := json.Marshal(online)
		if err != nil {
			logger.GetRequestEntry(r).WithError(err).Warn("json marshal error")
			w.WriteHeader(400)
			w.Write([]byte("json marshal error"))
		}
		w.Header().Set("Content-Type", "application/json;charset=UTF-8")
		w.WriteHeader(http.StatusOK)
		w.Write(b)
		return
	}
}
func GetOnline(rsender *redisocket.Sender) func(w http.ResponseWriter, r *http.Request) {

	return func(w http.ResponseWriter, r *http.Request) {
		params := mux.Vars(r)
		online, err := rsender.GetOnline(listenChannelPrefix, params["app_key"])
		if err != nil {
			logger.GetRequestEntry(r).WithError(err).Warn("get redis error")
			w.WriteHeader(400)
			w.Write([]byte("get redis error"))
		}
		b, err := json.Marshal(online)
		if err != nil {
			logger.GetRequestEntry(r).WithError(err).Warn("json marshal error")
			w.WriteHeader(400)
			w.Write([]byte("json marshal error"))
		}
		w.Header().Set("Content-Type", "application/json;charset=UTF-8")
		w.WriteHeader(http.StatusOK)
		w.Write(b)
		return
	}
}
func PushToSocket(rsender *redisocket.Sender) func(w http.ResponseWriter, r *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		params := mux.Vars(r)
		app_key := params["app_key"]
		if app_key == "" {
			logger.GetRequestEntry(r).Warn("empty param")
			w.WriteHeader(400)
			w.Write([]byte("empty param"))
			return
		}
		socket_id := params["socket_id"]
		data := r.FormValue("data")
		j := JsonCheck(data)
		rsender.PushToSid(listenChannelPrefix, app_key, socket_id, j)
		push := struct {
			SocketId string      `json:"socket_id"`
			Data     interface{} `json:"data"`
		}{
			SocketId: socket_id,
			Data:     data,
		}
		d, err := json.Marshal(push)
		if err != nil {
			logger.GetRequestEntry(r).Warn(err)
			w.WriteHeader(400)
			w.Write([]byte("data error"))
			return
		}
		w.Header().Set("Content-Type", "application/json;charset=UTF-8")
		w.WriteHeader(http.StatusOK)
		w.Write(d)
		return
	}
}
func PushToUser(rsender *redisocket.Sender) func(w http.ResponseWriter, r *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		params := mux.Vars(r)
		app_key := params["app_key"]
		if app_key == "" {
			logger.GetRequestEntry(r).Warn("empty param")
			w.WriteHeader(400)
			w.Write([]byte("empty param"))
			return
		}
		user_id := params["user_id"]
		data := r.FormValue("data")
		j := JsonCheck(data)
		rsender.PushToUid(listenChannelPrefix, app_key, user_id, j)
		push := struct {
			UserId string      `json:"user_id"`
			Data   interface{} `json:"data"`
		}{
			UserId: user_id,
			Data:   data,
		}
		d, err := json.Marshal(push)
		if err != nil {
			logger.GetRequestEntry(r).Warn(err)
			w.WriteHeader(400)
			w.Write([]byte("data error"))
			return
		}
		w.Header().Set("Content-Type", "application/json;charset=UTF-8")
		w.WriteHeader(http.StatusOK)
		w.Write(d)
		return
	}
}

func PushBatchMessage(rsender *redisocket.Sender) func(w http.ResponseWriter, r *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		params := mux.Vars(r)
		app_key := params["app_key"]
		if app_key == "" {
			logger.GetRequestEntry(r).Warn("empty param")
			w.WriteHeader(400)
			w.Write([]byte("empty param"))
			return
		}
		data := r.FormValue("batch_data")
		if data == "" {
			logger.GetRequestEntry(r).Warn("empty batch data")
			w.WriteHeader(400)
			w.Write([]byte("empty batch data"))
			return
		}
		batchData := make([]BatchData, 0)
		byteData := []byte(data)
		err := json.Unmarshal(byteData, &batchData)
		if err != nil {
			logger.GetRequestEntry(r).Warn(err)
			w.WriteHeader(400)
			w.Write([]byte("data error"))
			return
		}
		l := len(batchData)
		bd := make([]redisocket.BatchData, l)
		for _, data := range batchData {
			push := struct {
				Channel string      `json:"channel"`
				Event   string      `json:"event"`
				Data    interface{} `json:"data"`
			}{
				Channel: data.Channel,
				Event:   data.Event,
				Data:    data.Data,
			}
			d, err := json.Marshal(push)
			if err != nil {
				logger.GetRequestEntry(r).Warn(err)
				continue
			}
			b := redisocket.BatchData{
				Data:  d,
				Event: data.Channel,
			}
			bd = append(bd, b)

		}
		rsender.PushBatch(listenChannelPrefix, app_key, bd)
		response := struct {
			Total int `json:"total"`
			Cap   int `json:"cap"`
		}{
			Total: len(batchData),
			Cap:   len(byteData),
		}
		d, err := json.Marshal(response)
		if err != nil {
			w.WriteHeader(400)
			w.Write([]byte("response marshal error"))
			return
		}
		w.Header().Set("Content-Type", "application/json;charset=UTF-8")
		w.WriteHeader(http.StatusOK)
		w.Write(d)
		return
	}
}

func PushMessage(rsender *redisocket.Sender) func(w http.ResponseWriter, r *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		params := mux.Vars(r)
		app_key := params["app_key"]
		channel := params["channel"]
		event := params["event"]
		if app_key == "" || channel == "" || event == "" {
			logger.GetRequestEntry(r).Warn("empty param")
			w.WriteHeader(400)
			w.Write([]byte("empty param"))
			return
		}

		data := r.FormValue("data")
		if data == "" {
			logger.GetRequestEntry(r).Warn("empty data")
			w.WriteHeader(400)
			w.Write([]byte("empty data"))
			return
		}
		jsonData := JsonCheck(data)

		push := struct {
			Channel string      `json:"channel"`
			Event   string      `json:"event"`
			Data    interface{} `json:"data"`
		}{
			Channel: channel,
			Event:   event,
			Data:    jsonData,
		}
		d, err := json.Marshal(push)
		if err != nil {
			logger.GetRequestEntry(r).Warn(err)
			w.WriteHeader(400)
			w.Write([]byte("data error"))
			return
		}
		_, err = rsender.Push(listenChannelPrefix, app_key, channel, d)
		if err != nil {
			logger.GetRequestEntry(r).Warn(err)
			w.WriteHeader(400)
			w.Write([]byte("data error"))
			return
		}
		w.Header().Set("Content-Type", "application/json;charset=UTF-8")
		w.WriteHeader(http.StatusOK)
		w.Write(d)
		return
	}
}
func DecodeJWT(key *rsa.PublicKey) func(w http.ResponseWriter, r *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {

		data := r.FormValue("data")
		auth, err := Decode(key, data)
		if err != nil {
			logger.GetRequestEntry(r).Warnf("error:%s, post data:%s", err, data)
			w.WriteHeader(400)
			w.Write([]byte("data error"))
			return
		}
		if err = json.NewEncoder(w).Encode(auth); err != nil {
			logger.GetRequestEntry(r).Warnf("error:%s", err)
			w.WriteHeader(400)
			w.Write([]byte("parse error"))
		}
		return
	}
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/tls"
	"fmt"
	"log"
	"net/http"
	"os"
	"regexp"
	"strings"

	"github.com/franela/play-with-docker/config"
	"github.com/franela/play-with-docker/handlers"
	"github.com/franela/play-with-docker/services"
	"github.com/franela/play-with-docker/templates"
	gh "github.com/gorilla/handlers"
	"github.com/gorilla/mux"
	"github.com/miekg/dns"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"github.com/urfave/negroni"
	"github.com/yhat/wsutil"
)

func main() {

	config.ParseFlags()

	bypassCaptcha := len(os.Getenv("GOOGLE_RECAPTCHA_DISABLED")) > 0

	// Start the DNS server
	dns.HandleFunc(".", handleDnsRequest)
	udpDnsServer := &dns.Server{Addr: ":53", Net: "udp"}
	go func() {
		err := udpDnsServer.ListenAndServe()
		if err != nil {
			log.Fatal(err)
		}
	}()
	tcpDnsServer := &dns.Server{Addr: ":53", Net: "tcp"}
	go func() {
		err := tcpDnsServer.ListenAndServe()
		if err != nil {
			log.Fatal(err)
		}
	}()

	server := services.CreateWSServer()
	server.On("connection", handlers.WS)
	server.On("error", handlers.WSError)

	err := services.LoadSessionsFromDisk()
	if err != nil && !os.IsNotExist(err) {
		log.Fatal("Error decoding sessions from disk ", err)
	}

	r := mux.NewRouter()
	corsRouter := mux.NewRouter()

	// Reverse proxy (needs to be the first route, to make sure it is the first thing we check)
	proxyHandler := handlers.NewMultipleHostReverseProxy()
	websocketProxyHandler := handlers.NewMultipleHostWebsocketReverseProxy()

	proxyMultiplexer := func(rw http.ResponseWriter, r *http.Request) {
		if wsutil.IsWebSocketRequest(r) {
			websocketProxyHandler.ServeHTTP(rw, r)
		} else {
			proxyHandler.ServeHTTP(rw, r)
		}
	}

	corsHandler := gh.CORS(gh.AllowCredentials(), gh.AllowedHeaders([]string{"x-requested-with", "content-type"}), gh.AllowedOrigins([]string{"*"}))

	// Specific routes
	r.Host(`{subdomain:.*}{node:pwd[0-9]{1,3}_[0-9]{1,3}_[0-9]{1,3}_[0-9]{1,3}}-{port:[0-9]*}.{tld:.*}`).HandlerFunc(proxyMultiplexer)
	r.Host(`{subdomain:.*}{node:pwd[0-9]{1,3}_[0-9]{1,3}_[0-9]{1,3}_[0-9]{1,3}}.{tld:.*}`).HandlerFunc(proxyMultiplexer)
	r.HandleFunc("/ping", handlers.Ping).Methods("GET")
	corsRouter.HandleFunc("/sessions/{sessionId}", handlers.GetSession).Methods("GET")
	corsRouter.HandleFunc("/sessions/{sessionId}/instances", handlers.NewInstance).Methods("POST")
	corsRouter.HandleFunc("/sessions/{sessionId}/instances/{instanceName}", handlers.DeleteInstance).Methods("DELETE")
	r.HandleFunc("/sessions/{sessionId}/instances/{instanceName}/keys", handlers.SetKeys).Methods("POST")

	h := func(w http.ResponseWriter, r *http.Request) {
		http.ServeFile(w, r, "./www/index.html")
	}

	r.HandleFunc("/p/{sessionId}", h).Methods("GET")
	r.PathPrefix("/assets").Handler(http.FileServer(http.Dir("./www")))
	r.HandleFunc("/robots.txt", func(rw http.ResponseWriter, r *http.Request) {
		http.ServeFile(rw, r, "www/robots.txt")
	})
	r.HandleFunc("/sdk.js", func(rw http.ResponseWriter, r *http.Request) {
		http.ServeFile(rw, r, "www/sdk.js")
	})

	corsRouter.Handle("/sessions/{sessionId}/ws/", server)
	r.Handle("/metrics", promhttp.Handler())

	// Generic routes
	r.HandleFunc("/", func(rw http.ResponseWriter, r *http.Request) {
		if bypassCaptcha {
			http.ServeFile(rw, r, "./www/bypass.html")
		} else {
			welcome, tmplErr := templates.GetWelcomeTemplate()
			if tmplErr != nil {
				log.Fatal(tmplErr)
			}
			rw.Write(welcome)
		}
	}).Methods("GET")

	corsRouter.HandleFunc("/", handlers.NewSession).Methods("POST")

	n := negroni.Classic()
	r.PathPrefix("/").Handler(negroni.New(negroni.Wrap(corsHandler(corsRouter))))
	n.UseHandler(r)

	go func() {
		log.Println("Listening on port " + config.PortNumber)
		log.Fatal(http.ListenAndServe("0.0.0.0:"+config.PortNumber, n))
	}()

	ssl := mux.NewRouter()
	sslProxyHandler := handlers.NewSSLDaemonHandler()
	ssl.Host(`{subdomain:.*}{node:pwd[0-9]{1,3}_[0-9]{1,3}_[0-9]{1,3}_[0-9]{1,3}}-2375.{tld:.*}`).Handler(sslProxyHandler)
	log.Println("Listening TLS on port " + config.SSLPortNumber)

	s := &http.Server{Addr: "0.0.0.0:" + config.SSLPortNumber, Handler: ssl}
	s.TLSConfig = &tls.Config{}
	s.TLSConfig.GetCertificate = func(clientHello *tls.ClientHelloInfo) (*tls.Certificate, error) {

		chunks := strings.Split(clientHello.ServerName, ".")
		chunks = strings.Split(chunks[0], "-")
		ip := strings.Replace(strings.TrimPrefix(chunks[0], "pwd"), "_", ".", -1)
		i := services.FindInstanceByIP(ip)
		if i == nil {
			return nil, fmt.Errorf("Instance %s doesn't exist", clientHello.ServerName)
		}
		if i.GetCertificate() == nil {
			return nil, fmt.Errorf("Instance %s doesn't have a certificate", clientHello.ServerName)
		}
		return i.GetCertificate(), nil
	}
	log.Fatal(s.ListenAndServeTLS("", ""))
}

var dnsFilter = regexp.MustCompile(`pwd[0-9]{1,3}_[0-9]{1,3}_[0-9]{1,3}_[0-9]{1,3}`)

func handleDnsRequest(w dns.ResponseWriter, r *dns.Msg) {
	if len(r.Question) > 0 && dnsFilter.MatchString(r.Question[0].Name) {
		// this is something we know about and we should try to handle
		question := r.Question[0].Name
		domainChunks := strings.Split(question, ".")
		tldChunks := strings.Split(strings.TrimPrefix(domainChunks[0], "pwd"), "-")
		ip := strings.Replace(tldChunks[0], "_", ".", -1)

		m := new(dns.Msg)
		m.SetReply(r)
		m.Authoritative = true
		m.RecursionAvailable = true
		a, err := dns.NewRR(fmt.Sprintf("%s 60 IN A %s", question, ip))
		if err != nil {
			log.Fatal(err)
		}
		m.Answer = append(m.Answer, a)
		w.WriteMsg(m)
		return
	} else {
		// we have no information about this and we are not a recursive dns server, so we just fail so the client can fallback to the next dns server it has configured
		w.Close()
		// dns.HandleFailed(w, r)
		return
	}
}
<commit_msg>Make PWD dns recursive<commit_after>package main

import (
	"crypto/tls"
	"fmt"
	"log"
	"net"
	"net/http"
	"os"
	"regexp"
	"strings"

	"github.com/franela/play-with-docker/config"
	"github.com/franela/play-with-docker/handlers"
	"github.com/franela/play-with-docker/services"
	"github.com/franela/play-with-docker/templates"
	gh "github.com/gorilla/handlers"
	"github.com/gorilla/mux"
	"github.com/miekg/dns"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"github.com/urfave/negroni"
	"github.com/yhat/wsutil"
)

func main() {

	config.ParseFlags()

	bypassCaptcha := len(os.Getenv("GOOGLE_RECAPTCHA_DISABLED")) > 0

	// Start the DNS server
	dns.HandleFunc(".", handleDnsRequest)
	udpDnsServer := &dns.Server{Addr: ":53", Net: "udp"}
	go func() {
		err := udpDnsServer.ListenAndServe()
		if err != nil {
			log.Fatal(err)
		}
	}()
	tcpDnsServer := &dns.Server{Addr: ":53", Net: "tcp"}
	go func() {
		err := tcpDnsServer.ListenAndServe()
		if err != nil {
			log.Fatal(err)
		}
	}()

	server := services.CreateWSServer()
	server.On("connection", handlers.WS)
	server.On("error", handlers.WSError)

	err := services.LoadSessionsFromDisk()
	if err != nil && !os.IsNotExist(err) {
		log.Fatal("Error decoding sessions from disk ", err)
	}

	r := mux.NewRouter()
	corsRouter := mux.NewRouter()

	// Reverse proxy (needs to be the first route, to make sure it is the first thing we check)
	proxyHandler := handlers.NewMultipleHostReverseProxy()
	websocketProxyHandler := handlers.NewMultipleHostWebsocketReverseProxy()

	proxyMultiplexer := func(rw http.ResponseWriter, r *http.Request) {
		if wsutil.IsWebSocketRequest(r) {
			websocketProxyHandler.ServeHTTP(rw, r)
		} else {
			proxyHandler.ServeHTTP(rw, r)
		}
	}

	corsHandler := gh.CORS(gh.AllowCredentials(), gh.AllowedHeaders([]string{"x-requested-with", "content-type"}), gh.AllowedOrigins([]string{"*"}))

	// Specific routes
	r.Host(`{subdomain:.*}{node:pwd[0-9]{1,3}_[0-9]{1,3}_[0-9]{1,3}_[0-9]{1,3}}-{port:[0-9]*}.{tld:.*}`).HandlerFunc(proxyMultiplexer)
	r.Host(`{subdomain:.*}{node:pwd[0-9]{1,3}_[0-9]{1,3}_[0-9]{1,3}_[0-9]{1,3}}.{tld:.*}`).HandlerFunc(proxyMultiplexer)
	r.HandleFunc("/ping", handlers.Ping).Methods("GET")
	corsRouter.HandleFunc("/sessions/{sessionId}", handlers.GetSession).Methods("GET")
	corsRouter.HandleFunc("/sessions/{sessionId}/instances", handlers.NewInstance).Methods("POST")
	corsRouter.HandleFunc("/sessions/{sessionId}/instances/{instanceName}", handlers.DeleteInstance).Methods("DELETE")
	r.HandleFunc("/sessions/{sessionId}/instances/{instanceName}/keys", handlers.SetKeys).Methods("POST")

	h := func(w http.ResponseWriter, r *http.Request) {
		http.ServeFile(w, r, "./www/index.html")
	}

	r.HandleFunc("/p/{sessionId}", h).Methods("GET")
	r.PathPrefix("/assets").Handler(http.FileServer(http.Dir("./www")))
	r.HandleFunc("/robots.txt", func(rw http.ResponseWriter, r *http.Request) {
		http.ServeFile(rw, r, "www/robots.txt")
	})
	r.HandleFunc("/sdk.js", func(rw http.ResponseWriter, r *http.Request) {
		http.ServeFile(rw, r, "www/sdk.js")
	})

	corsRouter.Handle("/sessions/{sessionId}/ws/", server)
	r.Handle("/metrics", promhttp.Handler())

	// Generic routes
	r.HandleFunc("/", func(rw http.ResponseWriter, r *http.Request) {
		if bypassCaptcha {
			http.ServeFile(rw, r, "./www/bypass.html")
		} else {
			welcome, tmplErr := templates.GetWelcomeTemplate()
			if tmplErr != nil {
				log.Fatal(tmplErr)
			}
			rw.Write(welcome)
		}
	}).Methods("GET")

	corsRouter.HandleFunc("/", handlers.NewSession).Methods("POST")

	n := negroni.Classic()
	r.PathPrefix("/").Handler(negroni.New(negroni.Wrap(corsHandler(corsRouter))))
	n.UseHandler(r)

	go func() {
		log.Println("Listening on port " + config.PortNumber)
		log.Fatal(http.ListenAndServe("0.0.0.0:"+config.PortNumber, n))
	}()

	ssl := mux.NewRouter()
	sslProxyHandler := handlers.NewSSLDaemonHandler()
	ssl.Host(`{subdomain:.*}{node:pwd[0-9]{1,3}_[0-9]{1,3}_[0-9]{1,3}_[0-9]{1,3}}-2375.{tld:.*}`).Handler(sslProxyHandler)
	log.Println("Listening TLS on port " + config.SSLPortNumber)

	s := &http.Server{Addr: "0.0.0.0:" + config.SSLPortNumber, Handler: ssl}
	s.TLSConfig = &tls.Config{}
	s.TLSConfig.GetCertificate = func(clientHello *tls.ClientHelloInfo) (*tls.Certificate, error) {

		chunks := strings.Split(clientHello.ServerName, ".")
		chunks = strings.Split(chunks[0], "-")
		ip := strings.Replace(strings.TrimPrefix(chunks[0], "pwd"), "_", ".", -1)
		i := services.FindInstanceByIP(ip)
		if i == nil {
			return nil, fmt.Errorf("Instance %s doesn't exist", clientHello.ServerName)
		}
		if i.GetCertificate() == nil {
			return nil, fmt.Errorf("Instance %s doesn't have a certificate", clientHello.ServerName)
		}
		return i.GetCertificate(), nil
	}
	log.Fatal(s.ListenAndServeTLS("", ""))
}

var dnsFilter = regexp.MustCompile(`pwd[0-9]{1,3}_[0-9]{1,3}_[0-9]{1,3}_[0-9]{1,3}`)

func handleDnsRequest(w dns.ResponseWriter, r *dns.Msg) {
	if len(r.Question) > 0 && dnsFilter.MatchString(r.Question[0].Name) {
		// this is something we know about and we should try to handle
		question := r.Question[0].Name
		domainChunks := strings.Split(question, ".")
		tldChunks := strings.Split(strings.TrimPrefix(domainChunks[0], "pwd"), "-")
		ip := strings.Replace(tldChunks[0], "_", ".", -1)

		m := new(dns.Msg)
		m.SetReply(r)
		m.Authoritative = true
		m.RecursionAvailable = true
		a, err := dns.NewRR(fmt.Sprintf("%s 60 IN A %s", question, ip))
		if err != nil {
			log.Fatal(err)
		}
		m.Answer = append(m.Answer, a)
		w.WriteMsg(m)
		return
	} else {
		if len(r.Question) > 0 {
			question := r.Question[0].Name
			ips, err := net.LookupIP(question)
			if err != nil {
				// we have no information about this and we are not a recursive dns server, so we just fail so the client can fallback to the next dns server it has configured
				w.Close()
				// dns.HandleFailed(w, r)
				return
			}
			log.Printf("Not a PWD host. Looking up [%s] got [%s]\n", question, ips)
			m := new(dns.Msg)
			m.SetReply(r)
			m.Authoritative = true
			m.RecursionAvailable = true
			for _, ip := range ips {
				ipv4 := ip.To4()
				if ipv4 == nil {
					a, err := dns.NewRR(fmt.Sprintf("%s 60 IN AAAA %s", question, ip.String()))
					if err != nil {
						log.Fatal(err)
					}
					m.Answer = append(m.Answer, a)
				} else {
					a, err := dns.NewRR(fmt.Sprintf("%s 60 IN A %s", question, ipv4.String()))
					if err != nil {
						log.Fatal(err)
					}
					m.Answer = append(m.Answer, a)
				}
			}
			w.WriteMsg(m)
			return

		} else {
			// we have no information about this and we are not a recursive dns server, so we just fail so the client can fallback to the next dns server it has configured
			w.Close()
			// dns.HandleFailed(w, r)
			return
		}
	}
}
<|endoftext|>