<commit_before>package main

import (
	"reflect"
	"testing"
)

var identifierTests = []struct {
	values map[string]map[string]string
	id     Identifier
	found  bool
	value  string
}{
	{
		values: nil,
		id:     Identifier{key: "foo"},
		found:  false,
	},
	{
		values: map[string]map[string]string{"": {"foo": "example"}},
		id:     Identifier{key: "foo"},
		found:  true,
		value:  "example",
	},
	{
		values: map[string]map[string]string{"": {"foo": "example"}},
		id:     Identifier{key: "bar"},
		found:  false,
	},
	{
		values: map[string]map[string]string{"example": {"foo": "example"}},
		id:     Identifier{scope: "example", key: "foo"},
		found:  true,
		value:  "example",
	},
	{
		values: map[string]map[string]string{"": {"foo": "example"}},
		id:     Identifier{scope: "example", key: "foo"},
		found:  false,
	},
	{
		values: map[string]map[string]string{"example": {"foo": "example"}},
		id:     Identifier{scope: "example", key: "bar"},
		found:  false,
	},
}

func Test_found(t *testing.T) {
	for _, test := range identifierTests {
		got := found(test.values, &test.id)
		if got != test.found {
			t.Errorf("found not correct for %+v (found: %+v, got: %+v)", test.id, test.found, got)
		}
	}
}

func Test_collect(t *testing.T) {
	ids := []*Identifier{
		&Identifier{scope: "foo", key: "foo"},
		&Identifier{scope: "foo", key: "bar"},
		&Identifier{scope: "zoo", key: "foo"},
		&Identifier{scope: "foo", key: "foo"},
		&Identifier{scope: "foo", key: "baz"},
		&Identifier{scope: "qux", key: "bar"},
	}
	expectedFoo := &IdentifierGroup{
		scope: "foo",
		keys:  []string{"foo", "bar", "baz"},
	}
	expectedBar := &IdentifierGroup{
		scope: "bar",
		keys:  nil,
	}
	idgFoo := collect(ids, "foo")
	if !reflect.DeepEqual(idgFoo, expectedFoo) {
		t.Errorf("collect not correct (expected: %+v, got: %+v)", expectedFoo, idgFoo)
	}
	idgBar := collect(ids, "bar")
	if !reflect.DeepEqual(idgBar, expectedBar) {
		t.Errorf("collect not correct (expected: %+v, got: %+v)", expectedBar, idgBar)
	}
}

func Test_insert(t *testing.T) {
	values := make(map[string]map[string]string)
	id := &Identifier{key: "foo"}
	value := "bar"
	insert(values, id, value)
	v, ok := values[""]["foo"]
	if !ok {
		t.Errorf("insert failed for %+v", id)
	}
	if v != value {
		t.Errorf("insert not correctly for %+v (found: %+v, got: %+v)", id, v, value)
	}
	id = &Identifier{scope: "foo", key: "bar"}
	value = "example"
	insert(values, id, value)
	v, ok = values["foo"]["bar"]
	if !ok {
		t.Errorf("insert failed for %+v", id)
	}
	if v != value {
		t.Errorf("insert not correctly for %+v (found: %+v, got: %+v)", id, v, value)
	}
}

func Test_lookup(t *testing.T) {
	for _, test := range identifierTests {
		got := lookup(test.values, &test.id)
		if got != test.value {
			t.Errorf("lookup not correct for %+v (found: %+v, got: %+v)", test.id, test.value, got)
		}
	}
}
<commit_msg>add tests for empty<commit_after>package main

import (
	"reflect"
	"testing"
)

var identifierTests = []struct {
	values map[string]map[string]string
	id     Identifier
	found  bool
	value  string
}{
	{
		values: nil,
		id:     Identifier{key: "foo"},
		found:  false,
	},
	{
		values: map[string]map[string]string{"": {"foo": "example"}},
		id:     Identifier{key: "foo"},
		found:  true,
		value:  "example",
	},
	{
		values: map[string]map[string]string{"": {"foo": "example"}},
		id:     Identifier{key: "bar"},
		found:  false,
	},
	{
		values: map[string]map[string]string{"example": {"foo": "example"}},
		id:     Identifier{scope: "example", key: "foo"},
		found:  true,
		value:  "example",
	},
	{
		values: map[string]map[string]string{"": {"foo": "example"}},
		id:     Identifier{scope: "example", key: "foo"},
		found:  false,
	},
	{
		values: map[string]map[string]string{"example": {"foo": "example"}},
		id:     Identifier{scope: "example", key: "bar"},
		found:  false,
	},
}

func Test_found(t *testing.T) {
	for _, test := range identifierTests {
		got := found(test.values, &test.id)
		if got != test.found {
			t.Errorf("found not correct for %+v (found: %+v, got: %+v)", test.id, test.found, got)
		}
	}
}

func Test_collect(t *testing.T) {
	ids := []*Identifier{
		&Identifier{scope: "foo", key: "foo"},
		&Identifier{scope: "foo", key: "bar"},
		&Identifier{scope: "zoo", key: "foo"},
		&Identifier{scope: "foo", key: "foo"},
		&Identifier{scope: "foo", key: "baz"},
		&Identifier{scope: "qux", key: "bar"},
	}
	expectedFoo := &IdentifierGroup{
		scope: "foo",
		keys:  []string{"foo", "bar", "baz"},
	}
	expectedBar := &IdentifierGroup{
		scope: "bar",
		keys:  nil,
	}
	idgFoo := collect(ids, "foo")
	if !reflect.DeepEqual(idgFoo, expectedFoo) {
		t.Errorf("collect not correct (expected: %+v, got: %+v)", expectedFoo, idgFoo)
	}
	idgBar := collect(ids, "bar")
	if !reflect.DeepEqual(idgBar, expectedBar) {
		t.Errorf("collect not correct (expected: %+v, got: %+v)", expectedBar, idgBar)
	}
}

func Test_insert(t *testing.T) {
	values := make(map[string]map[string]string)
	id := &Identifier{key: "foo"}
	value := "bar"
	insert(values, id, value)
	v, ok := values[""]["foo"]
	if !ok {
		t.Errorf("insert failed for %+v", id)
	}
	if v != value {
		t.Errorf("insert not correctly for %+v (found: %+v, got: %+v)", id, v, value)
	}
	id = &Identifier{scope: "foo", key: "bar"}
	value = "example"
	insert(values, id, value)
	v, ok = values["foo"]["bar"]
	if !ok {
		t.Errorf("insert failed for %+v", id)
	}
	if v != value {
		t.Errorf("insert not correctly for %+v (found: %+v, got: %+v)", id, v, value)
	}
}

func Test_empty(t *testing.T) {
	tests := []struct {
		values   map[string]map[string]string
		expected bool
	}{
		{
			values:   nil,
			expected: true,
		},
		{
			values: map[string]map[string]string{
				"foo": map[string]string{},
			},
			expected: true,
		},
		{
			values: map[string]map[string]string{
				"foo": map[string]string{
					"bar": "",
					"baz": "",
				},
			},
			expected: true,
		},
		{
			values: map[string]map[string]string{
				"foo": map[string]string{
					"bar": "",
					"baz": "",
				},
				"bar": map[string]string{
					"qux": "quux",
				},
			},
			expected: false,
		},
	}
	for _, test := range tests {
		got := empty(test.values)
		if got != test.expected {
			t.Errorf("empty not correctly for %+v (expected: %+v, got: %+v)", test.values, test.expected, got)
		}
	}
}

func Test_lookup(t *testing.T) {
	for _, test := range identifierTests {
		got := lookup(test.values, &test.id)
		if got != test.value {
			t.Errorf("lookup not correct for %+v (expected: %+v, got: %+v)", test.id, test.value, got)
		}
	}
}
<|endoftext|>


<commit_before>package ofsclient

import (
	"encoding/binary"
	"math/rand"
	"net"
	"sync"
)

type client struct {
	mutex     sync.Mutex
	conn      net.Conn
	servers   []string
	authToken string
	idSize    int
}

// NewClient creates an instance of client which implements Client
func NewClient(authToken string, idSize int, servers ...string) (Client, error) {
	if len(servers) == 0 {
		return nil, ErrNoServers
	}

	if len(authToken) > 255 {
		return nil, ErrAuthTokenTooLarge
	}

	return &client{authToken: authToken, servers: servers, idSize: idSize}, nil
}

func (c *client) Close() {
	if c.conn != nil {
		c.conn.Close()
		c.conn = nil
	}
}

// connect selects an available server endpoint and establishes a connection
// and initiates authentication
func (c *client) connect() (err error) {
	// select a server
	n := rand.Intn(len(c.servers))

	// connect to server
	c.conn, err = net.Dial("tcp", c.servers[n])
	if err != nil {
		return
	}

	// authenticate (as necessary)
	return c.authenticate()
}

// authenticate writes an authentication header and auth token to the server IFF
// len(client.authToken) > 0
func (c *client) authenticate() (err error) {
	// do we need to auth?
	if len(c.authToken) > 0 {
		// create the header 0xFFFFFFnn where nn is the length of the auth token
		authBytes := make([]byte, 4)
		binary.BigEndian.PutUint32(authBytes[0:4], uint32(0xFFFFFF00|len(c.authToken)))

		// write the auth header
		_, err = c.conn.Write(authBytes)
		if err != nil {
			return
		}

		// write the auth token
		_, err = c.conn.Write([]byte(c.authToken))
		// @NOTE: if err != nil We should probably assume the token was bad and return ErrAuthFailed

		return
	}
	return
}

func (c *client) StreamIDBytes(count int, buffer []byte, callback func(int, []byte) error) (totalAllocated int, err error) {
	bufferCount := len(buffer) / c.idSize
	if bufferCount == 0 {
		return 0, CreateBadArgumentError("buffer", "The buffer is too small (%d bytes) to hold a single ID (%d bytes)", len(buffer), c.idSize)
	}

	c.mutex.Lock()
	defer c.mutex.Unlock()

	// on exit, if there was an error, clear the connection
	defer func() {
		if err != nil {
			c.conn = nil
		}
	}()

	// if we don't have a connection establish one
	if c.conn == nil {
		err = c.connect()
		if err != nil {
			return
		}
	}

	// create the command header, which is the count of ids
	countBytes := make([]byte, 4)

	for count > 0 {
		var readCount int

		// max # if ids we should read
		if count-bufferCount > 0 {
			readCount = bufferCount
		} else {
			readCount = count
		}

		binary.BigEndian.PutUint32(countBytes, uint32(readCount))

		// write the count
		_, err = c.conn.Write(countBytes)
		if err != nil {
			return
		}

		var bytesRead int
		// the number of bytes we expect to read
		expectedBytes := readCount * c.idSize

		// read the # of bytes for readCount ids
		bytesRead, err = c.conn.Read(buffer[0:expectedBytes])
		if err != nil {
			return totalAllocated, err
		}

		// did we read all the byes we expected?
		if bytesRead != expectedBytes {
			return totalAllocated, ErrShortRead
		}

		// update our counts prior to callback
		totalAllocated += readCount
		count -= readCount

		// stream readCount ids back to caller
		err = callback(readCount, buffer)
		if err != nil {
			return
		}
	}

	return
}

// GenerateIDBytes generates count number of overt-flake identifiers in raw/byte form
func (c *client) GenerateIDBytes(count int) (ids []byte, err error) {
	c.mutex.Lock()
	defer c.mutex.Unlock()

	// on exit, if there was an error, clear the connection
	defer func() {
		if err != nil {
			c.conn = nil
		}
	}()

	// if we don't have a connection establish one
	if c.conn == nil {
		err = c.connect()
		if err != nil {
			return
		}
	}

	// create the command header, which is the count of ids
	countBytes := make([]byte, 4)
	binary.BigEndian.PutUint32(countBytes, uint32(count))

	// write the count
	_, err = c.conn.Write(countBytes)
	if err != nil {
		return
	}

	// the number of bytes we expect to read
	expectedBytes := count * c.idSize

	// create a bytes buffer to accumulate the ids
	ids = make([]byte, expectedBytes)
	// read the ids
	bytesRead, err := c.conn.Read(ids)
	if err != nil {
		return
	}

	// did we read all the byes we expected?
	if bytesRead != expectedBytes {
		return nil, ErrShortRead
	}

	return
}

// GenerateID generates a single ID in []byte form
func (c *client) GenerateID() (id []byte, err error) {
	return c.GenerateIDBytes(1)
}
<commit_msg>restructure to use ServerEntry re-write connect to not be random and to try each server in order (ordered by priority) fix issue where I was using conn.Read and it was not reading all the bytes and my changes to address that didn't seem to work, so going with io.ReadFull(conn,buffer) which is the correct answer anyway<commit_after>package ofsclient

import (
	"encoding/binary"
	"fmt"
	"io"
	"net"
	"sync"
)

type client struct {
	mutex       sync.Mutex
	conn        net.Conn
	servers     []ServerEntry
	idSize      int
	serverIndex int
}

// NewClient creates an instance of client which implements Client
func NewClient(idSize int, servers []ServerEntry) (Client, error) {
	if len(servers) == 0 {
		return nil, ErrNoServers
	}

	for _, server := range servers {
		if len(server.Auth) > 255 {
			return nil, ErrAuthTokenTooLarge
		}
	}

	return &client{servers: servers, idSize: idSize}, nil
}

func (c *client) Close() {
	if c.conn != nil {
		c.conn.Close()
		c.conn = nil
	}
}

// connect selects an available server endpoint and establishes a connection
// and initiates authentication
func (c *client) connect() (err error) {
	// server entries are assumed to be priority ordered
	for index, server := range c.servers {
		// connect to server
		c.conn, err = net.Dial("tcp", server.Server)
		if err != nil {
			continue
		}

		// authenticate (as necessary)
		err = c.authenticate(server.Auth)
		// connected and auth'ed? excellent
		if err == nil {
			c.serverIndex = index
			return nil
		}
	}

	return ErrNoServerConnection
}

// authenticate writes an authentication header and auth token to the server IFF
// len(client.authToken) > 0
func (c *client) authenticate(authToken string) (err error) {
	// do we need to auth?
	if len(authToken) > 0 {
		// create the header 0xFFFFFFnn where nn is the length of the auth token
		authBytes := make([]byte, 4)
		binary.BigEndian.PutUint32(authBytes[0:4], uint32(0xFFFFFF00|len(authToken)))

		// write the auth header
		_, err = c.conn.Write(authBytes)
		if err != nil {
			return
		}

		// write the auth token
		_, err = c.conn.Write([]byte(authToken))
		// @NOTE: if err != nil We should probably assume the token was bad and return ErrAuthFailed

		return
	}
	return
}

func (c *client) StreamIDBytes(count int, buffer []byte, callback func(int, []byte) error) (totalAllocated int, err error) {
	bufferCount := len(buffer) / c.idSize
	if bufferCount == 0 {
		return 0, CreateBadArgumentError("buffer", "The buffer is too small (%d bytes) to hold a single ID (%d bytes)", len(buffer), c.idSize)
	}

	c.mutex.Lock()
	defer c.mutex.Unlock()

	// on exit, if there was an error, clear the connection
	defer func() {
		if err != nil {
			c.conn = nil
		}
	}()

	// if we don't have a connection establish one
	if c.conn == nil {
		err = c.connect()
		if err != nil {
			return
		}
	}

	// create the command header, which is the count of ids
	countBytes := make([]byte, 4)

	for count > 0 {
		var readCount int

		// max # if ids we should read
		if count-bufferCount > 0 {
			readCount = bufferCount
		} else {
			readCount = count
		}

		binary.BigEndian.PutUint32(countBytes, uint32(readCount))

		// write the count
		_, err = c.conn.Write(countBytes)
		if err != nil {
			return
		}

		var bytesRead int
		// the number of bytes we expect to read
		expectedBytes := readCount * c.idSize

		// read the # of bytes for readCount ids
		bytesRead, err = io.ReadFull(c.conn, buffer[0:expectedBytes])
		if err != nil {
			return totalAllocated, err
		}

		// did we read all the byes we expected?
		if bytesRead != expectedBytes {
			return totalAllocated, ErrShortRead
		}

		// update our counts prior to callback
		totalAllocated += readCount
		count -= readCount

		// stream readCount ids back to caller
		err = callback(readCount, buffer)
		if err != nil {
			return
		}
	}

	return
}

// GenerateIDBytes generates count number of overt-flake identifiers in raw/byte form
func (c *client) GenerateIDBytes(count int) (ids []byte, err error) {
	c.mutex.Lock()
	defer c.mutex.Unlock()

	// on exit, if there was an error, clear the connection
	defer func() {
		if err != nil {
			c.conn = nil
		}
	}()

	// if we don't have a connection establish one
	if c.conn == nil {
		err = c.connect()
		if err != nil {
			return
		}
	}

	// create the command header, which is the count of ids
	countBytes := make([]byte, 4)
	binary.BigEndian.PutUint32(countBytes, uint32(count))

	// write the count
	_, err = c.conn.Write(countBytes)
	if err != nil {
		return
	}

	// the number of bytes we expect to read
	expectedBytes := count * c.idSize

	// create a bytes buffer to accumulate the ids
	ids = make([]byte, expectedBytes)

	var bytesRead int

	// read the ids
	bytesRead, err = io.ReadFull(c.conn, ids)
	if err != nil {
		return
	}

	// did we read all the byes we expected?
	if bytesRead != expectedBytes {
		fmt.Printf("Expected %d bytes but only read %d for id.Size %d\n", expectedBytes, bytesRead, c.idSize)
		return nil, ErrShortRead
	}

	return
}

// GenerateID generates a single ID in []byte form
func (c *client) GenerateID() (id []byte, err error) {
	return c.GenerateIDBytes(1)
}
<|endoftext|>


<commit_before>package main

import (
    "golang.org/x/crypto/pbkdf2"
	"golang.org/x/crypto/scrypt"
	"bytes"
	"crypto/sha256"
	"fmt"
	"time"
	"os"
	"math/rand"
	"github.com/vsergeev/btckeygenie/btckey"
)

// source: http://stackoverflow.com/a/31832326/1025599
const letterBytes = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"
const (
    letterIdxBits = 6                    // 6 bits to represent a letter index
    letterIdxMask = 1<<letterIdxBits - 1 // All 1-bits, as many as letterIdxBits
    letterIdxMax  = 63 / letterIdxBits   // # of letter indices fitting in 63 bits
)

func RandStringBytesMaskImpr(n int) string {
    b := make([]byte, n)
    // A rand.Int63() generates 63 random bits, enough for letterIdxMax letters!
    for i, cache, remain := n-1, rand.Int63(), letterIdxMax; i >= 0; {
        if remain == 0 {
            cache, remain = rand.Int63(), letterIdxMax
        }
        if idx := int(cache & letterIdxMask); idx < len(letterBytes) {
            b[i] = letterBytes[idx]
            i--
        }
        cache >>= letterIdxBits
        remain--
    }

    return string(b)
}
// end source

func main () {
	rand.Seed(time.Now().UTC().UnixNano()) // seed rand

	var address string
	saltValue := ""

	if len(os.Args) >= 2 {
		address = os.Args[1]
		if len(os.Args) == 3 {
			saltValue = os.Args[2]
		} else {
			saltValue = "";
		}
	} else {
		fmt.Printf("Usage: %s [Address] [Salt - optional]\n\n", os.Args[0])
		os.Exit(0)
	}

	fmt.Printf("Using address \"%s\" and salt \"%s\"\n", address, saltValue)

	tries := 0
	start := time.Now()
	for {
		passphraseValue := RandStringBytesMaskImpr(8)
		result := bruteforce(passphraseValue, saltValue, address);
		if result != "" {
			fmt.Printf("Found! Passphrase %s\n", passphraseValue)
			os.Exit(0)
		} else {
			tries += 1
			fmt.Printf("\rTried %d passphrases in %s [last passphrase: %s]", tries, time.Since(start), passphraseValue)
		}
	}
}

func bruteforce(passphraseValue string, saltValue string, address string) string {
	var priv btckey.PrivateKey
	var err error

    pass := fmt.Sprint(passphraseValue, "\x01")
    salt := fmt.Sprint(saltValue, "\x01")
    key, _ := scrypt.Key([]byte(pass), []byte(salt), 262144, 8, 1, 32)
    pass = fmt.Sprint(passphraseValue, "\x02")
    salt = fmt.Sprint(saltValue, "\x02")
    key2 := pbkdf2.Key([]byte(pass), []byte(salt), 65536, 32, sha256.New)

    var result bytes.Buffer
    for i := 0; i < len(key); i++ {
        result.WriteByte(key[i] ^ key2[i])
    }

	err = priv.FromBytes(result.Bytes())
	if err != nil {
		fmt.Printf("Error importing private key: %s [%s]\n", err, passphraseValue)
		return ""
	}

	address_uncompressed := priv.ToAddressUncompressed()

	if (address_uncompressed == address) {
		return passphraseValue
	}

	return ""
}
<commit_msg>Fix the RNG<commit_after>package main

import (
    "golang.org/x/crypto/pbkdf2"
	"golang.org/x/crypto/scrypt"
	"bytes"
	"crypto/sha256"
	"fmt"
	"time"
	"os"
	"math/rand"
	"github.com/vsergeev/btckeygenie/btckey"
)

const letterBytes = "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"
func random(r *rand.Rand, n int) string {
    b := make([]byte, n)
	for i := range b {
        b[i] = letterBytes[r.Intn(62)]
    }

    return string(b)
}

func main () {
	r := rand.New(rand.NewSource(time.Now().Unix()))

	var address string
	saltValue := ""

	if len(os.Args) >= 2 {
		address = os.Args[1]
		if len(os.Args) == 3 {
			saltValue = os.Args[2]
		} else {
			saltValue = "";
		}
	} else {
		fmt.Printf("Usage: %s [Address] [Salt - optional]\n\n", os.Args[0])
		os.Exit(0)
	}

	fmt.Printf("Using address \"%s\" and salt \"%s\"\n", address, saltValue)

	tries := 0
	start := time.Now()
	for {
		passphraseValue := random(r, 8)
		result := bruteforce(passphraseValue, saltValue, address);
		if result != "" {
			fmt.Printf("Found! Passphrase %s\n", passphraseValue)
			os.Exit(0)
		} else {
			tries += 1
			fmt.Printf("\rTried %d passphrases in %s [last passphrase: %s]", tries, time.Since(start), passphraseValue)
		}
	}
}

func bruteforce(passphraseValue string, saltValue string, address string) string {
	var priv btckey.PrivateKey
	var err error

    pass := fmt.Sprint(passphraseValue, "\x01")
    salt := fmt.Sprint(saltValue, "\x01")
    key, _ := scrypt.Key([]byte(pass), []byte(salt), 262144, 8, 1, 32)
    pass = fmt.Sprint(passphraseValue, "\x02")
    salt = fmt.Sprint(saltValue, "\x02")
    key2 := pbkdf2.Key([]byte(pass), []byte(salt), 65536, 32, sha256.New)

    var result bytes.Buffer
    for i := 0; i < len(key); i++ {
        result.WriteByte(key[i] ^ key2[i])
    }

	err = priv.FromBytes(result.Bytes())
	if err != nil {
		fmt.Printf("Error importing private key: %s [%s]\n", err, passphraseValue)
		return ""
	}

	address_uncompressed := priv.ToAddressUncompressed()

	if (address_uncompressed == address) {
		return passphraseValue
	}

	return ""
}
<|endoftext|>


<commit_before>package tree

import (
	"fmt"
	"math/rand"
	"reflect"
	"testing"
)

// Define a function Build(records []Record) (*Node, error)
// where Record is a struct containing int fields ID and Parent
// and Node is a struct containing int field ID and []*Node field Children.

var successTestCases = []struct {
	name     string
	input    []Record
	expected *Node
}{
	{
		name:     "empty input",
		input:    []Record{},
		expected: nil,
	},
	{
		name: "one node",
		input: []Record{
			{ID: 0},
		},
		expected: &Node{
			ID: 0,
		},
	},
	{
		name: "three nodes in order",
		input: []Record{
			{ID: 0},
			{ID: 1, Parent: 0},
			{ID: 2, Parent: 0},
		},
		expected: &Node{
			ID: 0,
			Children: []*Node{
				{ID: 1},
				{ID: 2},
			},
		},
	},
	{
		name: "three nodes in reverse order",
		input: []Record{
			{ID: 2, Parent: 0},
			{ID: 1, Parent: 0},
			{ID: 0},
		},
		expected: &Node{
			ID: 0,
			Children: []*Node{
				{ID: 1},
				{ID: 2},
			},
		},
	},
	{
		name: "more than two children",
		input: []Record{
			{ID: 3, Parent: 0},
			{ID: 2, Parent: 0},
			{ID: 1, Parent: 0},
			{ID: 0},
		},
		expected: &Node{
			ID: 0,
			Children: []*Node{
				{ID: 1},
				{ID: 2},
				{ID: 3},
			},
		},
	},
	{
		name: "binary tree",
		input: []Record{
			{ID: 5, Parent: 1},
			{ID: 3, Parent: 2},
			{ID: 2, Parent: 0},
			{ID: 4, Parent: 1},
			{ID: 1, Parent: 0},
			{ID: 0},
			{ID: 6, Parent: 2},
		},
		expected: &Node{
			ID: 0,
			Children: []*Node{
				{
					ID: 1,
					Children: []*Node{
						{ID: 4},
						{ID: 5},
					},
				},
				{
					ID: 2,
					Children: []*Node{
						{ID: 3},
						{ID: 6},
					},
				},
			},
		},
	},
	{
		name: "unbalanced tree",
		input: []Record{
			{ID: 5, Parent: 2},
			{ID: 3, Parent: 2},
			{ID: 2, Parent: 0},
			{ID: 4, Parent: 1},
			{ID: 1, Parent: 0},
			{ID: 0},
			{ID: 6, Parent: 2},
		},
		expected: &Node{
			ID: 0,
			Children: []*Node{
				{
					ID: 1,
					Children: []*Node{
						{ID: 4},
					},
				},
				{
					ID: 2,
					Children: []*Node{
						{ID: 3},
						{ID: 5},
						{ID: 6},
					},
				},
			},
		},
	},
}

var failureTestCases = []struct {
	name  string
	input []Record
}{
	{
		name: "root node has parent",
		input: []Record{
			{ID: 0, Parent: 1},
			{ID: 1, Parent: 0},
		},
	},
	{
		name: "no root node",
		input: []Record{
			{ID: 1, Parent: 0},
		},
	},
	{
		name: "non-continuous",
		input: []Record{
			{ID: 2, Parent: 0},
			{ID: 4, Parent: 2},
			{ID: 1, Parent: 0},
			{ID: 0},
		},
	},
	{
		name: "cycle directly",
		input: []Record{
			{ID: 5, Parent: 2},
			{ID: 3, Parent: 2},
			{ID: 2, Parent: 2},
			{ID: 4, Parent: 1},
			{ID: 1, Parent: 0},
			{ID: 0},
			{ID: 6, Parent: 3},
		},
	},
	{
		name: "cycle indirectly",
		input: []Record{
			{ID: 5, Parent: 2},
			{ID: 3, Parent: 2},
			{ID: 2, Parent: 6},
			{ID: 4, Parent: 1},
			{ID: 1, Parent: 0},
			{ID: 0},
			{ID: 6, Parent: 3},
		},
	},
	{
		name: "higher id parent of lower id",
		input: []Record{
			{ID: 0},
			{ID: 2, Parent: 0},
			{ID: 1, Parent: 2},
		},
	},
}

func (n Node) String() string {
	return fmt.Sprintf("%d:%s", n.ID, n.Children)
}

func TestMakeTreeSuccess(t *testing.T) {
	for _, tt := range successTestCases {
		actual, err := Build(tt.input)
		if err != nil {
			var _ error = err
			t.Fatalf("Build for test case %q returned error %q. Error not expected.",
				tt.name, err)
		}
		if !reflect.DeepEqual(actual, tt.expected) {
			t.Fatalf("Build for test case %q returned %s but was expected to return %s.",
				tt.name, actual, tt.expected)
		}
	}
}

func TestMakeTreeFailure(t *testing.T) {
	for _, tt := range failureTestCases {
		actual, err := Build(tt.input)
		if err == nil {
			t.Fatalf("Build for test case %q returned %s but was expected to fail.",
				tt.name, actual)
		}
	}
}

func shuffleRecords(records []Record) []Record {
	rand := rand.New(rand.NewSource(42))
	newRecords := make([]Record, len(records))
	for i, idx := range rand.Perm(len(records)) {
		newRecords[i] = records[idx]
	}
	return newRecords
}

// Binary tree
func makeTwoTreeRecords() []Record {
	records := make([]Record, 1<<16)
	for i := range records {
		if i == 0 {
			records[i] = Record{ID: 0}
		} else {
			records[i] = Record{ID: i, Parent: i >> 1}
		}
	}
	return shuffleRecords(records)
}

var twoTreeRecords = makeTwoTreeRecords()

func BenchmarkTwoTree(b *testing.B) {
	for i := 0; i < b.N; i++ {
		Build(twoTreeRecords)
	}
}

// Each node but the root node and leaf nodes has ten children.
func makeTenTreeRecords() []Record {
	records := make([]Record, 10000)
	for i := range records {
		if i == 0 {
			records[i] = Record{ID: 0}
		} else {
			records[i] = Record{ID: i, Parent: i / 10}
		}
	}
	return shuffleRecords(records)
}

var tenTreeRecords = makeTenTreeRecords()

func BenchmarkTenTree(b *testing.B) {
	for i := 0; i < b.N; i++ {
		Build(tenTreeRecords)
	}
}

func makeShallowRecords() []Record {
	records := make([]Record, 10000)
	for i := range records {
		records[i] = Record{ID: i, Parent: 0}
	}
	return shuffleRecords(records)
}

var shallowRecords = makeShallowRecords()

func BenchmarkShallowTree(b *testing.B) {
	for i := 0; i < b.N; i++ {
		Build(shallowRecords)
	}
}
<commit_msg>tree-building: add test cases for duplicate nodes (#1002)<commit_after>package tree

import (
	"fmt"
	"math/rand"
	"reflect"
	"testing"
)

// Define a function Build(records []Record) (*Node, error)
// where Record is a struct containing int fields ID and Parent
// and Node is a struct containing int field ID and []*Node field Children.

var successTestCases = []struct {
	name     string
	input    []Record
	expected *Node
}{
	{
		name:     "empty input",
		input:    []Record{},
		expected: nil,
	},
	{
		name: "one node",
		input: []Record{
			{ID: 0},
		},
		expected: &Node{
			ID: 0,
		},
	},
	{
		name: "three nodes in order",
		input: []Record{
			{ID: 0},
			{ID: 1, Parent: 0},
			{ID: 2, Parent: 0},
		},
		expected: &Node{
			ID: 0,
			Children: []*Node{
				{ID: 1},
				{ID: 2},
			},
		},
	},
	{
		name: "three nodes in reverse order",
		input: []Record{
			{ID: 2, Parent: 0},
			{ID: 1, Parent: 0},
			{ID: 0},
		},
		expected: &Node{
			ID: 0,
			Children: []*Node{
				{ID: 1},
				{ID: 2},
			},
		},
	},
	{
		name: "more than two children",
		input: []Record{
			{ID: 3, Parent: 0},
			{ID: 2, Parent: 0},
			{ID: 1, Parent: 0},
			{ID: 0},
		},
		expected: &Node{
			ID: 0,
			Children: []*Node{
				{ID: 1},
				{ID: 2},
				{ID: 3},
			},
		},
	},
	{
		name: "binary tree",
		input: []Record{
			{ID: 5, Parent: 1},
			{ID: 3, Parent: 2},
			{ID: 2, Parent: 0},
			{ID: 4, Parent: 1},
			{ID: 1, Parent: 0},
			{ID: 0},
			{ID: 6, Parent: 2},
		},
		expected: &Node{
			ID: 0,
			Children: []*Node{
				{
					ID: 1,
					Children: []*Node{
						{ID: 4},
						{ID: 5},
					},
				},
				{
					ID: 2,
					Children: []*Node{
						{ID: 3},
						{ID: 6},
					},
				},
			},
		},
	},
	{
		name: "unbalanced tree",
		input: []Record{
			{ID: 5, Parent: 2},
			{ID: 3, Parent: 2},
			{ID: 2, Parent: 0},
			{ID: 4, Parent: 1},
			{ID: 1, Parent: 0},
			{ID: 0},
			{ID: 6, Parent: 2},
		},
		expected: &Node{
			ID: 0,
			Children: []*Node{
				{
					ID: 1,
					Children: []*Node{
						{ID: 4},
					},
				},
				{
					ID: 2,
					Children: []*Node{
						{ID: 3},
						{ID: 5},
						{ID: 6},
					},
				},
			},
		},
	},
}

var failureTestCases = []struct {
	name  string
	input []Record
}{
	{
		name: "root node has parent",
		input: []Record{
			{ID: 0, Parent: 1},
			{ID: 1, Parent: 0},
		},
	},
	{
		name: "no root node",
		input: []Record{
			{ID: 1, Parent: 0},
		},
	},
	{
		name: "duplicate node",
		input: []Record{
			{ID: 0, Parent: 0},
			{ID: 1, Parent: 0},
			{ID: 1, Parent: 0},
		},
	},
	{
		name: "duplicate root",
		input: []Record{
			{ID: 0, Parent: 0},
			{ID: 0, Parent: 0},
		},
	},
	{
		name: "non-continuous",
		input: []Record{
			{ID: 2, Parent: 0},
			{ID: 4, Parent: 2},
			{ID: 1, Parent: 0},
			{ID: 0},
		},
	},
	{
		name: "cycle directly",
		input: []Record{
			{ID: 5, Parent: 2},
			{ID: 3, Parent: 2},
			{ID: 2, Parent: 2},
			{ID: 4, Parent: 1},
			{ID: 1, Parent: 0},
			{ID: 0},
			{ID: 6, Parent: 3},
		},
	},
	{
		name: "cycle indirectly",
		input: []Record{
			{ID: 5, Parent: 2},
			{ID: 3, Parent: 2},
			{ID: 2, Parent: 6},
			{ID: 4, Parent: 1},
			{ID: 1, Parent: 0},
			{ID: 0},
			{ID: 6, Parent: 3},
		},
	},
	{
		name: "higher id parent of lower id",
		input: []Record{
			{ID: 0},
			{ID: 2, Parent: 0},
			{ID: 1, Parent: 2},
		},
	},
}

func (n Node) String() string {
	return fmt.Sprintf("%d:%s", n.ID, n.Children)
}

func TestMakeTreeSuccess(t *testing.T) {
	for _, tt := range successTestCases {
		actual, err := Build(tt.input)
		if err != nil {
			var _ error = err
			t.Fatalf("Build for test case %q returned error %q. Error not expected.",
				tt.name, err)
		}
		if !reflect.DeepEqual(actual, tt.expected) {
			t.Fatalf("Build for test case %q returned %s but was expected to return %s.",
				tt.name, actual, tt.expected)
		}
	}
}

func TestMakeTreeFailure(t *testing.T) {
	for _, tt := range failureTestCases {
		actual, err := Build(tt.input)
		if err == nil {
			t.Fatalf("Build for test case %q returned %s but was expected to fail.",
				tt.name, actual)
		}
	}
}

func shuffleRecords(records []Record) []Record {
	rand := rand.New(rand.NewSource(42))
	newRecords := make([]Record, len(records))
	for i, idx := range rand.Perm(len(records)) {
		newRecords[i] = records[idx]
	}
	return newRecords
}

// Binary tree
func makeTwoTreeRecords() []Record {
	records := make([]Record, 1<<16)
	for i := range records {
		if i == 0 {
			records[i] = Record{ID: 0}
		} else {
			records[i] = Record{ID: i, Parent: i >> 1}
		}
	}
	return shuffleRecords(records)
}

var twoTreeRecords = makeTwoTreeRecords()

func BenchmarkTwoTree(b *testing.B) {
	for i := 0; i < b.N; i++ {
		Build(twoTreeRecords)
	}
}

// Each node but the root node and leaf nodes has ten children.
func makeTenTreeRecords() []Record {
	records := make([]Record, 10000)
	for i := range records {
		if i == 0 {
			records[i] = Record{ID: 0}
		} else {
			records[i] = Record{ID: i, Parent: i / 10}
		}
	}
	return shuffleRecords(records)
}

var tenTreeRecords = makeTenTreeRecords()

func BenchmarkTenTree(b *testing.B) {
	for i := 0; i < b.N; i++ {
		Build(tenTreeRecords)
	}
}

func makeShallowRecords() []Record {
	records := make([]Record, 10000)
	for i := range records {
		records[i] = Record{ID: i, Parent: 0}
	}
	return shuffleRecords(records)
}

var shallowRecords = makeShallowRecords()

func BenchmarkShallowTree(b *testing.B) {
	for i := 0; i < b.N; i++ {
		Build(shallowRecords)
	}
}
<|endoftext|>


<commit_before>package tree

import (
	"fmt"
	"math/rand"
	"reflect"
	"testing"
)

// Define a function Build(records []Record) (*Node, error)
// where Record is a struct containing int fields ID and Parent
// and Node is a struct containing int field ID and []*Node field Children.
//
// Also define a testVersion with a value that matches
// the targetTestVersion here.

const targetTestVersion = 4

var successTestCases = []struct {
	name     string
	input    []Record
	expected *Node
}{
	{
		name:     "empty input",
		input:    []Record{},
		expected: nil,
	},
	{
		name: "one node",
		input: []Record{
			{ID: 0},
		},
		expected: &Node{
			ID: 0,
		},
	},
	{
		name: "three nodes in order",
		input: []Record{
			{ID: 0},
			{ID: 1, Parent: 0},
			{ID: 2, Parent: 0},
		},
		expected: &Node{
			ID: 0,
			Children: []*Node{
				{ID: 1},
				{ID: 2},
			},
		},
	},
	{
		name: "three nodes in reverse order",
		input: []Record{
			{ID: 2, Parent: 0},
			{ID: 1, Parent: 0},
			{ID: 0},
		},
		expected: &Node{
			ID: 0,
			Children: []*Node{
				{ID: 1},
				{ID: 2},
			},
		},
	},
	{
		name: "more than two children",
		input: []Record{
			{ID: 3, Parent: 0},
			{ID: 2, Parent: 0},
			{ID: 1, Parent: 0},
			{ID: 0},
		},
		expected: &Node{
			ID: 0,
			Children: []*Node{
				{ID: 1},
				{ID: 2},
				{ID: 3},
			},
		},
	},
	{
		name: "binary tree",
		input: []Record{
			{ID: 5, Parent: 1},
			{ID: 3, Parent: 2},
			{ID: 2, Parent: 0},
			{ID: 4, Parent: 1},
			{ID: 1, Parent: 0},
			{ID: 0},
			{ID: 6, Parent: 2},
		},
		expected: &Node{
			ID: 0,
			Children: []*Node{
				{
					ID: 1,
					Children: []*Node{
						{ID: 4},
						{ID: 5},
					},
				},
				{
					ID: 2,
					Children: []*Node{
						{ID: 3},
						{ID: 6},
					},
				},
			},
		},
	},
	{
		name: "unbalanced tree",
		input: []Record{
			{ID: 5, Parent: 2},
			{ID: 3, Parent: 2},
			{ID: 2, Parent: 0},
			{ID: 4, Parent: 1},
			{ID: 1, Parent: 0},
			{ID: 0},
			{ID: 6, Parent: 2},
		},
		expected: &Node{
			ID: 0,
			Children: []*Node{
				{
					ID: 1,
					Children: []*Node{
						{ID: 4},
					},
				},
				{
					ID: 2,
					Children: []*Node{
						{ID: 3},
						{ID: 5},
						{ID: 6},
					},
				},
			},
		},
	},
}

var failureTestCases = []struct {
	name  string
	input []Record
}{
	{
		name: "root node has parent",
		input: []Record{
			{ID: 0, Parent: 1},
			{ID: 1, Parent: 0},
		},
	},
	{
		name: "no root node",
		input: []Record{
			{ID: 1, Parent: 0},
		},
	},
	{
		name: "non-continuous",
		input: []Record{
			{ID: 2, Parent: 0},
			{ID: 4, Parent: 2},
			{ID: 1, Parent: 0},
			{ID: 0},
		},
	},
	{
		name: "cycle directly",
		input: []Record{
			{ID: 5, Parent: 2},
			{ID: 3, Parent: 2},
			{ID: 2, Parent: 2},
			{ID: 4, Parent: 1},
			{ID: 1, Parent: 0},
			{ID: 0},
			{ID: 6, Parent: 3},
		},
	},
	{
		name: "cycle indirectly",
		input: []Record{
			{ID: 5, Parent: 2},
			{ID: 3, Parent: 2},
			{ID: 2, Parent: 6},
			{ID: 4, Parent: 1},
			{ID: 1, Parent: 0},
			{ID: 0},
			{ID: 6, Parent: 3},
		},
	},
	{
		name: "higher id parent of lower id",
		input: []Record{
			{ID: 0},
			{ID: 2, Parent: 0},
			{ID: 1, Parent: 2},
		},
	},
}

func (n Node) String() string {
	return fmt.Sprintf("%d:%s", n.ID, n.Children)
}

func TestTestVersion(t *testing.T) {
	if testVersion != targetTestVersion {
		t.Fatalf("Found testVersion = %v, want %v", testVersion, targetTestVersion)
	}
}

func TestMakeTreeSuccess(t *testing.T) {
	for _, tt := range successTestCases {
		actual, err := Build(tt.input)
		if err != nil {
			var _ error = err
			t.Fatalf("Build for test case %q returned error %q. Error not expected.",
				tt.name, err)
		}
		if !reflect.DeepEqual(actual, tt.expected) {
			t.Fatalf("Build for test case %q returned %s but was expected to return %s.",
				tt.name, tt.expected, actual)
		}
	}
}

func TestMakeTreeFailure(t *testing.T) {
	for _, tt := range failureTestCases {
		actual, err := Build(tt.input)
		if err == nil {
			t.Fatalf("Build for test case %q returned %s but was expected to fail.",
				tt.name, actual)
		}
	}
}

func shuffleRecords(records []Record) []Record {
	rand := rand.New(rand.NewSource(42))
	newRecords := make([]Record, len(records))
	for i, idx := range rand.Perm(len(records)) {
		newRecords[i] = records[idx]
	}
	return newRecords
}

// Binary tree
func makeTwoTreeRecords() []Record {
	records := make([]Record, 1<<16)
	for i := range records {
		if i == 0 {
			records[i] = Record{ID: 0}
		} else {
			records[i] = Record{ID: i, Parent: i >> 1}
		}
	}
	return shuffleRecords(records)
}

var twoTreeRecords = makeTwoTreeRecords()

func BenchmarkTwoTree(b *testing.B) {
	for i := 0; i < b.N; i++ {
		Build(twoTreeRecords)
	}
}

// Each node but the root node and leaf nodes has ten children.
func makeTenTreeRecords() []Record {
	records := make([]Record, 10000)
	for i := range records {
		if i == 0 {
			records[i] = Record{ID: 0}
		} else {
			records[i] = Record{ID: i, Parent: i / 10}
		}
	}
	return shuffleRecords(records)
}

var tenTreeRecords = makeTenTreeRecords()

func BenchmarkTenTree(b *testing.B) {
	for i := 0; i < b.N; i++ {
		Build(tenTreeRecords)
	}
}

func makeShallowRecords() []Record {
	records := make([]Record, 10000)
	for i := range records {
		records[i] = Record{ID: i, Parent: 0}
	}
	return shuffleRecords(records)
}

var shallowRecords = makeShallowRecords()

func BenchmarkShallowTree(b *testing.B) {
	for i := 0; i < b.N; i++ {
		Build(shallowRecords)
	}
}
<commit_msg>Fix test error output order<commit_after>package tree

import (
	"fmt"
	"math/rand"
	"reflect"
	"testing"
)

// Define a function Build(records []Record) (*Node, error)
// where Record is a struct containing int fields ID and Parent
// and Node is a struct containing int field ID and []*Node field Children.
//
// Also define a testVersion with a value that matches
// the targetTestVersion here.

const targetTestVersion = 4

var successTestCases = []struct {
	name     string
	input    []Record
	expected *Node
}{
	{
		name:     "empty input",
		input:    []Record{},
		expected: nil,
	},
	{
		name: "one node",
		input: []Record{
			{ID: 0},
		},
		expected: &Node{
			ID: 0,
		},
	},
	{
		name: "three nodes in order",
		input: []Record{
			{ID: 0},
			{ID: 1, Parent: 0},
			{ID: 2, Parent: 0},
		},
		expected: &Node{
			ID: 0,
			Children: []*Node{
				{ID: 1},
				{ID: 2},
			},
		},
	},
	{
		name: "three nodes in reverse order",
		input: []Record{
			{ID: 2, Parent: 0},
			{ID: 1, Parent: 0},
			{ID: 0},
		},
		expected: &Node{
			ID: 0,
			Children: []*Node{
				{ID: 1},
				{ID: 2},
			},
		},
	},
	{
		name: "more than two children",
		input: []Record{
			{ID: 3, Parent: 0},
			{ID: 2, Parent: 0},
			{ID: 1, Parent: 0},
			{ID: 0},
		},
		expected: &Node{
			ID: 0,
			Children: []*Node{
				{ID: 1},
				{ID: 2},
				{ID: 3},
			},
		},
	},
	{
		name: "binary tree",
		input: []Record{
			{ID: 5, Parent: 1},
			{ID: 3, Parent: 2},
			{ID: 2, Parent: 0},
			{ID: 4, Parent: 1},
			{ID: 1, Parent: 0},
			{ID: 0},
			{ID: 6, Parent: 2},
		},
		expected: &Node{
			ID: 0,
			Children: []*Node{
				{
					ID: 1,
					Children: []*Node{
						{ID: 4},
						{ID: 5},
					},
				},
				{
					ID: 2,
					Children: []*Node{
						{ID: 3},
						{ID: 6},
					},
				},
			},
		},
	},
	{
		name: "unbalanced tree",
		input: []Record{
			{ID: 5, Parent: 2},
			{ID: 3, Parent: 2},
			{ID: 2, Parent: 0},
			{ID: 4, Parent: 1},
			{ID: 1, Parent: 0},
			{ID: 0},
			{ID: 6, Parent: 2},
		},
		expected: &Node{
			ID: 0,
			Children: []*Node{
				{
					ID: 1,
					Children: []*Node{
						{ID: 4},
					},
				},
				{
					ID: 2,
					Children: []*Node{
						{ID: 3},
						{ID: 5},
						{ID: 6},
					},
				},
			},
		},
	},
}

var failureTestCases = []struct {
	name  string
	input []Record
}{
	{
		name: "root node has parent",
		input: []Record{
			{ID: 0, Parent: 1},
			{ID: 1, Parent: 0},
		},
	},
	{
		name: "no root node",
		input: []Record{
			{ID: 1, Parent: 0},
		},
	},
	{
		name: "non-continuous",
		input: []Record{
			{ID: 2, Parent: 0},
			{ID: 4, Parent: 2},
			{ID: 1, Parent: 0},
			{ID: 0},
		},
	},
	{
		name: "cycle directly",
		input: []Record{
			{ID: 5, Parent: 2},
			{ID: 3, Parent: 2},
			{ID: 2, Parent: 2},
			{ID: 4, Parent: 1},
			{ID: 1, Parent: 0},
			{ID: 0},
			{ID: 6, Parent: 3},
		},
	},
	{
		name: "cycle indirectly",
		input: []Record{
			{ID: 5, Parent: 2},
			{ID: 3, Parent: 2},
			{ID: 2, Parent: 6},
			{ID: 4, Parent: 1},
			{ID: 1, Parent: 0},
			{ID: 0},
			{ID: 6, Parent: 3},
		},
	},
	{
		name: "higher id parent of lower id",
		input: []Record{
			{ID: 0},
			{ID: 2, Parent: 0},
			{ID: 1, Parent: 2},
		},
	},
}

func (n Node) String() string {
	return fmt.Sprintf("%d:%s", n.ID, n.Children)
}

func TestTestVersion(t *testing.T) {
	if testVersion != targetTestVersion {
		t.Fatalf("Found testVersion = %v, want %v", testVersion, targetTestVersion)
	}
}

func TestMakeTreeSuccess(t *testing.T) {
	for _, tt := range successTestCases {
		actual, err := Build(tt.input)
		if err != nil {
			var _ error = err
			t.Fatalf("Build for test case %q returned error %q. Error not expected.",
				tt.name, err)
		}
		if !reflect.DeepEqual(actual, tt.expected) {
			t.Fatalf("Build for test case %q returned %s but was expected to return %s.",
				tt.name, actual, tt.expected)
		}
	}
}

func TestMakeTreeFailure(t *testing.T) {
	for _, tt := range failureTestCases {
		actual, err := Build(tt.input)
		if err == nil {
			t.Fatalf("Build for test case %q returned %s but was expected to fail.",
				tt.name, actual)
		}
	}
}

func shuffleRecords(records []Record) []Record {
	rand := rand.New(rand.NewSource(42))
	newRecords := make([]Record, len(records))
	for i, idx := range rand.Perm(len(records)) {
		newRecords[i] = records[idx]
	}
	return newRecords
}

// Binary tree
func makeTwoTreeRecords() []Record {
	records := make([]Record, 1<<16)
	for i := range records {
		if i == 0 {
			records[i] = Record{ID: 0}
		} else {
			records[i] = Record{ID: i, Parent: i >> 1}
		}
	}
	return shuffleRecords(records)
}

var twoTreeRecords = makeTwoTreeRecords()

func BenchmarkTwoTree(b *testing.B) {
	for i := 0; i < b.N; i++ {
		Build(twoTreeRecords)
	}
}

// Each node but the root node and leaf nodes has ten children.
func makeTenTreeRecords() []Record {
	records := make([]Record, 10000)
	for i := range records {
		if i == 0 {
			records[i] = Record{ID: 0}
		} else {
			records[i] = Record{ID: i, Parent: i / 10}
		}
	}
	return shuffleRecords(records)
}

var tenTreeRecords = makeTenTreeRecords()

func BenchmarkTenTree(b *testing.B) {
	for i := 0; i < b.N; i++ {
		Build(tenTreeRecords)
	}
}

func makeShallowRecords() []Record {
	records := make([]Record, 10000)
	for i := range records {
		records[i] = Record{ID: i, Parent: 0}
	}
	return shuffleRecords(records)
}

var shallowRecords = makeShallowRecords()

func BenchmarkShallowTree(b *testing.B) {
	for i := 0; i < b.N; i++ {
		Build(shallowRecords)
	}
}
<|endoftext|>


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

import (
	"context"
	"log"
	"net"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"go.opentelemetry.io/otel"
	"go.opentelemetry.io/otel/sdk/trace/tracetest"
)

func TestNewAgentClientUDPWithParamsBadHostport(t *testing.T) {
	agentClient, err := newAgentClientUDP(agentClientUDPParams{
		Host: "blahblah",
		Port: "",
	})
	assert.Error(t, err)
	assert.Nil(t, agentClient)
}

func TestNewAgentClientUDPWithParams(t *testing.T) {
	mockServer, err := newUDPListener()
	require.NoError(t, err)
	defer mockServer.Close()
	host, port, err := net.SplitHostPort(mockServer.LocalAddr().String())
	assert.NoError(t, err)

	agentClient, err := newAgentClientUDP(agentClientUDPParams{
		Host:                host,
		Port:                port,
		MaxPacketSize:       25000,
		AttemptReconnecting: true,
	})
	assert.NoError(t, err)
	assert.NotNil(t, agentClient)
	assert.Equal(t, 25000, agentClient.maxPacketSize)

	if assert.IsType(t, &reconnectingUDPConn{}, agentClient.connUDP) {
		assert.Equal(t, (*log.Logger)(nil), agentClient.connUDP.(*reconnectingUDPConn).logger)
	}

	assert.NoError(t, agentClient.Close())
}

func TestNewAgentClientUDPWithParamsDefaults(t *testing.T) {
	mockServer, err := newUDPListener()
	require.NoError(t, err)
	defer mockServer.Close()
	host, port, err := net.SplitHostPort(mockServer.LocalAddr().String())
	assert.NoError(t, err)

	agentClient, err := newAgentClientUDP(agentClientUDPParams{
		Host:                host,
		Port:                port,
		AttemptReconnecting: true,
	})
	assert.NoError(t, err)
	assert.NotNil(t, agentClient)
	assert.Equal(t, udpPacketMaxLength, agentClient.maxPacketSize)

	if assert.IsType(t, &reconnectingUDPConn{}, agentClient.connUDP) {
		assert.Equal(t, (*log.Logger)(nil), agentClient.connUDP.(*reconnectingUDPConn).logger)
	}

	assert.NoError(t, agentClient.Close())
}

func TestNewAgentClientUDPWithParamsReconnectingDisabled(t *testing.T) {
	mockServer, err := newUDPListener()
	require.NoError(t, err)
	defer mockServer.Close()
	host, port, err := net.SplitHostPort(mockServer.LocalAddr().String())
	assert.NoError(t, err)

	agentClient, err := newAgentClientUDP(agentClientUDPParams{
		Host:                host,
		Port:                port,
		Logger:              nil,
		AttemptReconnecting: false,
	})
	assert.NoError(t, err)
	assert.NotNil(t, agentClient)
	assert.Equal(t, udpPacketMaxLength, agentClient.maxPacketSize)

	assert.IsType(t, &net.UDPConn{}, agentClient.connUDP)

	assert.NoError(t, agentClient.Close())
}

type errorHandler struct{ t *testing.T }

func (eh errorHandler) Handle(err error) { assert.NoError(eh.t, err) }

func TestJaegerAgentUDPLimitBatching(t *testing.T) {
	otel.SetErrorHandler(errorHandler{t})

	// 1500 spans, size 79559, does not fit within one UDP packet with the default size of 65000.
	n := 1500
	s := make(tracetest.SpanStubs, n).Snapshots()

	exp, err := NewRawExporter(
		WithAgentEndpoint(WithAgentHost("localhost"), WithAgentPort("6831")),
	)
	require.NoError(t, err)

	ctx := context.Background()
	assert.NoError(t, exp.ExportSpans(ctx, s))
	assert.NoError(t, exp.Shutdown(ctx))
}

// generateALargeSpan generates a span with a long name.
func generateALargeSpan() tracetest.SpanStub {
	return tracetest.SpanStub{
		Name: "a-longer-name-that-makes-it-exceeds-limit",
	}
}

func TestSpanExceedsMaxPacketLimit(t *testing.T) {
	otel.SetErrorHandler(errorHandler{t})

	// 106 is the serialized size of a span with default values.
	maxSize := 106

	largeSpans := tracetest.SpanStubs{generateALargeSpan(), {}}.Snapshots()
	normalSpans := tracetest.SpanStubs{{}, {}}.Snapshots()

	exp, err := NewRawExporter(
		WithAgentEndpoint(WithAgentHost("localhost"), WithAgentPort("6831"), WithMaxPacketSize(maxSize+1)),
	)
	require.NoError(t, err)

	ctx := context.Background()
	assert.Error(t, exp.ExportSpans(ctx, largeSpans))
	assert.NoError(t, exp.ExportSpans(ctx, normalSpans))
	assert.NoError(t, exp.Shutdown(ctx))
}

func TestEmitBatchWithMultipleErrors(t *testing.T) {
	otel.SetErrorHandler(errorHandler{t})

	span := generateALargeSpan()
	largeSpans := tracetest.SpanStubs{span, span}.Snapshots()
	// make max packet size smaller than span
	maxSize := len(span.Name)
	exp, err := NewRawExporter(
		WithAgentEndpoint(WithAgentHost("localhost"), WithAgentPort("6831"), WithMaxPacketSize(maxSize)),
	)
	require.NoError(t, err)

	ctx := context.Background()
	err = exp.ExportSpans(ctx, largeSpans)
	assert.Error(t, err)
	require.Contains(t, err.Error(), "multiple errors")
}
<commit_msg>Use mock server as jaeger agent listener. (#1930)<commit_after>// Copyright The OpenTelemetry Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package jaeger

import (
	"context"
	"log"
	"net"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"go.opentelemetry.io/otel"
	"go.opentelemetry.io/otel/sdk/trace/tracetest"
)

func TestNewAgentClientUDPWithParamsBadHostport(t *testing.T) {
	agentClient, err := newAgentClientUDP(agentClientUDPParams{
		Host: "blahblah",
		Port: "",
	})
	assert.Error(t, err)
	assert.Nil(t, agentClient)
}

func TestNewAgentClientUDPWithParams(t *testing.T) {
	mockServer, err := newUDPListener()
	require.NoError(t, err)
	defer mockServer.Close()
	host, port, err := net.SplitHostPort(mockServer.LocalAddr().String())
	assert.NoError(t, err)

	agentClient, err := newAgentClientUDP(agentClientUDPParams{
		Host:                host,
		Port:                port,
		MaxPacketSize:       25000,
		AttemptReconnecting: true,
	})
	assert.NoError(t, err)
	assert.NotNil(t, agentClient)
	assert.Equal(t, 25000, agentClient.maxPacketSize)

	if assert.IsType(t, &reconnectingUDPConn{}, agentClient.connUDP) {
		assert.Equal(t, (*log.Logger)(nil), agentClient.connUDP.(*reconnectingUDPConn).logger)
	}

	assert.NoError(t, agentClient.Close())
}

func TestNewAgentClientUDPWithParamsDefaults(t *testing.T) {
	mockServer, err := newUDPListener()
	require.NoError(t, err)
	defer mockServer.Close()
	host, port, err := net.SplitHostPort(mockServer.LocalAddr().String())
	assert.NoError(t, err)

	agentClient, err := newAgentClientUDP(agentClientUDPParams{
		Host:                host,
		Port:                port,
		AttemptReconnecting: true,
	})
	assert.NoError(t, err)
	assert.NotNil(t, agentClient)
	assert.Equal(t, udpPacketMaxLength, agentClient.maxPacketSize)

	if assert.IsType(t, &reconnectingUDPConn{}, agentClient.connUDP) {
		assert.Equal(t, (*log.Logger)(nil), agentClient.connUDP.(*reconnectingUDPConn).logger)
	}

	assert.NoError(t, agentClient.Close())
}

func TestNewAgentClientUDPWithParamsReconnectingDisabled(t *testing.T) {
	mockServer, err := newUDPListener()
	require.NoError(t, err)
	defer mockServer.Close()
	host, port, err := net.SplitHostPort(mockServer.LocalAddr().String())
	assert.NoError(t, err)

	agentClient, err := newAgentClientUDP(agentClientUDPParams{
		Host:                host,
		Port:                port,
		Logger:              nil,
		AttemptReconnecting: false,
	})
	assert.NoError(t, err)
	assert.NotNil(t, agentClient)
	assert.Equal(t, udpPacketMaxLength, agentClient.maxPacketSize)

	assert.IsType(t, &net.UDPConn{}, agentClient.connUDP)

	assert.NoError(t, agentClient.Close())
}

type errorHandler struct{ t *testing.T }

func (eh errorHandler) Handle(err error) { assert.NoError(eh.t, err) }

func TestJaegerAgentUDPLimitBatching(t *testing.T) {
	otel.SetErrorHandler(errorHandler{t})

	mockServer, err := newUDPListener()
	require.NoError(t, err)
	defer mockServer.Close()
	host, port, err := net.SplitHostPort(mockServer.LocalAddr().String())
	assert.NoError(t, err)

	// 1500 spans, size 79559, does not fit within one UDP packet with the default size of 65000.
	n := 1500
	s := make(tracetest.SpanStubs, n).Snapshots()

	exp, err := NewRawExporter(
		WithAgentEndpoint(WithAgentHost(host), WithAgentPort(port)),
	)
	require.NoError(t, err)

	ctx := context.Background()
	assert.NoError(t, exp.ExportSpans(ctx, s))
	assert.NoError(t, exp.Shutdown(ctx))
}

// generateALargeSpan generates a span with a long name.
func generateALargeSpan() tracetest.SpanStub {
	return tracetest.SpanStub{
		Name: "a-longer-name-that-makes-it-exceeds-limit",
	}
}

func TestSpanExceedsMaxPacketLimit(t *testing.T) {
	otel.SetErrorHandler(errorHandler{t})

	mockServer, err := newUDPListener()
	require.NoError(t, err)
	defer mockServer.Close()
	host, port, err := net.SplitHostPort(mockServer.LocalAddr().String())
	assert.NoError(t, err)

	// 106 is the serialized size of a span with default values.
	maxSize := 106

	largeSpans := tracetest.SpanStubs{generateALargeSpan(), {}}.Snapshots()
	normalSpans := tracetest.SpanStubs{{}, {}}.Snapshots()

	exp, err := NewRawExporter(
		WithAgentEndpoint(WithAgentHost(host), WithAgentPort(port), WithMaxPacketSize(maxSize+1)),
	)
	require.NoError(t, err)

	ctx := context.Background()
	assert.Error(t, exp.ExportSpans(ctx, largeSpans))
	assert.NoError(t, exp.ExportSpans(ctx, normalSpans))
	assert.NoError(t, exp.Shutdown(ctx))
}

func TestEmitBatchWithMultipleErrors(t *testing.T) {
	otel.SetErrorHandler(errorHandler{t})

	mockServer, err := newUDPListener()
	require.NoError(t, err)
	defer mockServer.Close()
	host, port, err := net.SplitHostPort(mockServer.LocalAddr().String())
	assert.NoError(t, err)

	span := generateALargeSpan()
	largeSpans := tracetest.SpanStubs{span, span}.Snapshots()
	// make max packet size smaller than span
	maxSize := len(span.Name)
	exp, err := NewRawExporter(
		WithAgentEndpoint(WithAgentHost(host), WithAgentPort(port), WithMaxPacketSize(maxSize)),
	)
	require.NoError(t, err)

	ctx := context.Background()
	err = exp.ExportSpans(ctx, largeSpans)
	assert.Error(t, err)
	require.Contains(t, err.Error(), "multiple errors")
}
<|endoftext|>


<commit_before>package shards

import (
	// "time"
	"fmt"
	"github.com/trendrr/goshire/dynmap"
	"time"
)

// A router table.
// The table is considered generally immutable.  If any changes occur a new table should
// be generated and propagated.
type RouterTable struct {

	//The name of the service. should be unique.  (Example: "trendrrdb")
	Service string

	//The revision # of the router table
	//this should be always increasing so greater revision means more upto date router table
	//this should typically be a timestamp
	Revision int64

	//total # of partitions
	TotalPartitions int

	//Replication Factory
	ReplicationFactor int

	//This is the params that we partition by
	//multi key partitions will be accessed in order, separated by "|" then hashed.
	PartitionKeys []string

	//entries organized by partition
	//index in the array is the partition
	EntriesPartition [][]*RouterEntry

	//The unique entries
	Entries []*RouterEntry

	//serialized dynmap
	DynMap *dynmap.DynMap
}

func NewRouterTable(service string) *RouterTable {
	return &RouterTable{
		Service:           service,
		DynMap:            dynmap.NewDynMap(),
		Revision:          time.Now().Unix(),
		TotalPartitions:   0,
		ReplicationFactor: 2,
	}
}

// updates the revision to now
func (this *RouterTable) UpdateRevision() (previous, current int64) {
	prev := this.Revision
	this.Revision = time.Now().Unix()
	//reset the dynmap
	this.DynMap = dynmap.NewDynMap()
	return prev, this.Revision
}

//Rebuilds this router table, if you changed anything, you should call this and use the newly built
//table
func (this *RouterTable) Rebuild() (*RouterTable, error) {
	total := 0
	//set TotalPartitions
	for _, e := range this.Entries {
		total += len(e.Partitions)
	}
	this.TotalPartitions = total
	mp := this.toDynMap()

	table, err := ToRouterTable(mp)
	if err != nil {
		return nil, err
	}
	return table, nil
}

// Adds one or more new entries
// These entries will replace any entries in the current routertable.
// a new router table is returned.
func (this *RouterTable) AddEntries(entry ...*RouterEntry) (*RouterTable, error) {
	//copy the router table
	routerTable, err := this.Rebuild()
	if err != nil {
		return routerTable, err
	}

	entries := make([]*RouterEntry, 0)

	//build a new entries array with every *except the ones we are adding
	for _, e := range routerTable.Entries {
		//check if this entry matches any of the entries we are adding.
		found := false
		for _, en := range entry {
			if e.Id() == en.Id() {
				found = true
			}
		}
		if !found {
			entries = append(entries, e)
		}
	}
	//now add the new ones
	for _, en := range entry {
		entries = append(entries, en)
	}
	routerTable.Entries = entries
	routerTable.UpdateRevision()
	routerTable, err = routerTable.Rebuild()
	return routerTable, err
}

func (this *RouterTable) FindEntry(id string) (*RouterEntry, bool) {
	for _, e := range this.Entries {
		if e.Id() == id {
			return e, true
		}
	}
	return nil, false
}

// Creates a new router table from the dynmap passed in
func ToRouterTable(mp *dynmap.DynMap) (*RouterTable, error) {
	t := &RouterTable{}

	var ok bool
	t.Service, ok = mp.GetString("service")
	if !ok {
		return nil, fmt.Errorf("No Service in the table %s", mp)
	}

	t.Revision, ok = mp.GetInt64("revision")
	if !ok {
		return nil, fmt.Errorf("No Revision in the table %s", mp)
	}

	t.TotalPartitions, ok = mp.GetInt("total_partitions")
	if !ok {
		return nil, fmt.Errorf("No total_partitions in the table %s", mp)
	}

	t.ReplicationFactor, ok = mp.GetInt("replication_factor")
	if !ok {
		return nil, fmt.Errorf("No replication_factor in the table %s", mp)
	}

	t.PartitionKeys, ok = mp.GetStringSlice("partition_keys")
	if !ok {
		//do nothing, right?
	}

	//fill the entries
	t.Entries = make([]*RouterEntry, 0)
	entryMaps, ok := mp.GetDynMapSlice("entries")
	if !ok {
		return nil, fmt.Errorf("Bad entries in the table %s", mp)
	}
	for _, em := range entryMaps {
		routerEntry, err := ToRouterEntry(em)
		if err != nil {
			return nil, err
		}
		t.Entries = append(t.Entries, routerEntry)
	}

	// set up the partition to entry mapping
	partitionCount := 0
	entriesPartition := make([]*RouterEntry, t.TotalPartitions)
	for _, e := range t.Entries {
		for _, p := range e.Partitions {
			if p >= t.TotalPartitions {
				return nil, fmt.Errorf("Bad Partition entry (greater then total partitions): %s", e)
			}
			entriesPartition[p] = e
			partitionCount++
		}
	}

	if partitionCount != t.TotalPartitions {
		return nil, fmt.Errorf("Bad table, some partitions un accounted for")
	}

	t.DynMap = t.toDynMap()

	t.EntriesPartition = make([][]*RouterEntry, t.TotalPartitions)
	//Now setup the replication partitions.
	for _, e := range t.Entries {

		for _, p := range e.Partitions {
			pRep, err := t.repPartitions(p, e)
			if err != nil {
				return nil, fmt.Errorf("Bad table (%s)", err)
			}
			entries := make([]*RouterEntry, len(pRep)+1)
			entries[0] = e
			for i := 1; i < len(entries); i++ {
				entries[i] = entriesPartition[pRep[i-1]]
				e.PartitionsMap[pRep[i-1]] = false
			}
			t.EntriesPartition[p] = entries
		}
	}

	return t, nil
}

//internal use.
//skips the cached version
func (this *RouterTable) toDynMap() *dynmap.DynMap {
	mp := dynmap.NewDynMap()
	mp.Put("service", this.Service)
	mp.Put("revision", this.Revision)
	mp.Put("total_partitions", this.TotalPartitions)
	mp.Put("replication_factor", this.ReplicationFactor)
	mp.Put("partition_keys", this.PartitionKeys)

	entries := make([]*dynmap.DynMap, 0)
	for _, e := range this.Entries {
		entries = append(entries, e.ToDynMap())
	}
	mp.Put("entries", entries)
	this.DynMap = mp
	return mp
}

// Translate to a DynMap of the form:
// {
//     "service" : "trendrrdb",
//     "revision" : 898775762309309,
//     "total_partitions" : 256,
//     "entries" : [
//         {/*router entry 1*/},
//         {/*router entry 2*/}
//     ]
// }
func (this *RouterTable) ToDynMap() *dynmap.DynMap {
	if this.DynMap != nil && len(this.DynMap.Map) > 0 {
		return this.DynMap
	}
	return this.toDynMap()
}

//gets the partitions that should replicate this master.
func (this *RouterTable) repPartitions(partition int, entry *RouterEntry) ([]int, error) {
	//This method could be much better optimized, but
	//it is fairly rare, so we wont worry about it..
	entries := make([]int, 0)
	if partition >= this.TotalPartitions {
		return entries, fmt.Errorf("Requested partition %d is out of bounds (%d) ", partition, this.TotalPartitions)
	}

	for i := 1; i < this.TotalPartitions; i++ {
		par := (i + partition) % this.TotalPartitions
		v, ok := entry.PartitionsMap[par]
		if ok && v {
			//this is master.  skip to next one
			continue
		}
		entries = append(entries, par)
		//we check if len < repfactor -1 (minus one because we still need the current partition)
		if len(entries) == this.ReplicationFactor-1 {
			return entries, nil
		}
	}
	return entries, nil

}

// Gets the entries associated with the given partition
// [0] should be the master entry, and there should be
// table.ReplicationFactor number of entries
func (this *RouterTable) PartitionEntries(partition int) ([]*RouterEntry, error) {
	if partition >= this.TotalPartitions {
		return make([]*RouterEntry, 0), fmt.Errorf("Requested partition %d is out of bounds (%d) ", partition, this.TotalPartitions)
	}
	return this.EntriesPartition[partition], nil
}

type RouterEntry struct {
	//The address of this entry
	Address  string
	JsonPort int
	HttpPort int
	BinPort  int

	//list of partitions this entry is responsible for (master only)
	Partitions []int

	//Map of all the partitions this entry is responsible for.  true indicates master, false otherwise
	PartitionsMap map[int]bool
}

// Creates a new router entry from the dynmap passed in
func ToRouterEntry(mp *dynmap.DynMap) (*RouterEntry, error) {
	e := &RouterEntry{
		PartitionsMap: make(map[int]bool),
	}
	var ok bool
	e.Address, ok = mp.GetString("address")
	if !ok {
		return nil, fmt.Errorf("No Address in Entry: %s", mp)
	}

	e.JsonPort = mp.MustInt("ports.json", 0)
	e.HttpPort = mp.MustInt("ports.http", 0)
	e.BinPort = mp.MustInt("ports.bin", 0)

	e.Partitions, ok = mp.GetIntSlice("partitions")
	if !ok {
		e.Partitions = make([]int, 0)
	}
	for _, p := range e.Partitions {
		e.PartitionsMap[p] = true
	}
	return e, nil
}

//Id for this entry.  current is address:jsonport
func (this *RouterEntry) Id() string {
	return fmt.Sprintf("%s:%d", this.Address, this.JsonPort)
}

// Translate to a DynMap of the form:
// {
//     "id" : "localhost:8009"
//     "address" : "localhost",
//     "ports" : {
//         "json" : 8009,
//         "http" : 8010,
//	       "bin" : 8011
//     }
//     "partitions" : [1,2,3,4,5,6,7,8,9]
// }
func (this *RouterEntry) ToDynMap() *dynmap.DynMap {
	mp := dynmap.NewDynMap()
	mp.Put("address", this.Address)
	if this.JsonPort > 0 {
		mp.PutWithDot("ports.json", this.JsonPort)
	}

	if this.HttpPort > 0 {
		mp.PutWithDot("ports.http", this.HttpPort)
	}

	if this.BinPort > 0 {
		mp.PutWithDot("ports.bin", this.BinPort)
	}

	mp.Put("id", this.Id())
	mp.Put("partitions", this.Partitions)
	return mp
}
<commit_msg>better defaults<commit_after>package shards

import (
	// "time"
	"fmt"
	"github.com/trendrr/goshire/dynmap"
	"time"
)

// A router table.
// The table is considered generally immutable.  If any changes occur a new table should
// be generated and propagated.
type RouterTable struct {

	//The name of the service. should be unique.  (Example: "trendrrdb")
	Service string

	//The revision # of the router table
	//this should be always increasing so greater revision means more upto date router table
	//this should typically be a timestamp
	Revision int64

	//total # of partitions
	TotalPartitions int

	//Replication Factory
	ReplicationFactor int

	//This is the params that we partition by
	//multi key partitions will be accessed in order, separated by "|" then hashed.
	PartitionKeys []string

	//entries organized by partition
	//index in the array is the partition
	EntriesPartition [][]*RouterEntry

	//The unique entries
	Entries []*RouterEntry

	//serialized dynmap
	DynMap *dynmap.DynMap
}

func NewRouterTable(service string) *RouterTable {
	return &RouterTable{
		Service:           service,
		DynMap:            dynmap.NewDynMap(),
		Revision:          int64(0),
		TotalPartitions:   512,
		ReplicationFactor: 1,
	}
}

// updates the revision to now
func (this *RouterTable) UpdateRevision() (previous, current int64) {
	prev := this.Revision
	this.Revision = time.Now().Unix()
	//reset the dynmap
	this.DynMap = dynmap.NewDynMap()
	return prev, this.Revision
}

//Rebuilds this router table, if you changed anything, you should call this and use the newly built
//table
func (this *RouterTable) Rebuild() (*RouterTable, error) {
	total := 0
	//set TotalPartitions
	for _, e := range this.Entries {
		total += len(e.Partitions)
	}
	this.TotalPartitions = total
	mp := this.toDynMap()

	table, err := ToRouterTable(mp)
	if err != nil {
		return nil, err
	}
	return table, nil
}

// Adds one or more new entries
// These entries will replace any entries in the current routertable.
// a new router table is returned.
func (this *RouterTable) AddEntries(entry ...*RouterEntry) (*RouterTable, error) {
	//copy the router table
	routerTable, err := this.Rebuild()
	if err != nil {
		return routerTable, err
	}

	entries := make([]*RouterEntry, 0)

	//build a new entries array with every *except the ones we are adding
	for _, e := range routerTable.Entries {
		//check if this entry matches any of the entries we are adding.
		found := false
		for _, en := range entry {
			if e.Id() == en.Id() {
				found = true
			}
		}
		if !found {
			entries = append(entries, e)
		}
	}
	//now add the new ones
	for _, en := range entry {
		entries = append(entries, en)
	}
	routerTable.Entries = entries
	routerTable.UpdateRevision()
	routerTable, err = routerTable.Rebuild()
	return routerTable, err
}

func (this *RouterTable) FindEntry(id string) (*RouterEntry, bool) {
	for _, e := range this.Entries {
		if e.Id() == id {
			return e, true
		}
	}
	return nil, false
}

// Creates a new router table from the dynmap passed in
func ToRouterTable(mp *dynmap.DynMap) (*RouterTable, error) {
	t := &RouterTable{}

	var ok bool
	t.Service, ok = mp.GetString("service")
	if !ok {
		return nil, fmt.Errorf("No Service in the table %s", mp)
	}

	t.Revision, ok = mp.GetInt64("revision")
	if !ok {
		return nil, fmt.Errorf("No Revision in the table %s", mp)
	}

	t.TotalPartitions, ok = mp.GetInt("total_partitions")
	if !ok {
		return nil, fmt.Errorf("No total_partitions in the table %s", mp)
	}

	t.ReplicationFactor, ok = mp.GetInt("replication_factor")
	if !ok {
		return nil, fmt.Errorf("No replication_factor in the table %s", mp)
	}

	t.PartitionKeys, ok = mp.GetStringSlice("partition_keys")
	if !ok {
		//do nothing, right?
	}

	//fill the entries
	t.Entries = make([]*RouterEntry, 0)
	entryMaps, ok := mp.GetDynMapSlice("entries")
	if !ok {
		return nil, fmt.Errorf("Bad entries in the table %s", mp)
	}
	for _, em := range entryMaps {
		routerEntry, err := ToRouterEntry(em)
		if err != nil {
			return nil, err
		}
		t.Entries = append(t.Entries, routerEntry)
	}

	// set up the partition to entry mapping
	partitionCount := 0
	entriesPartition := make([]*RouterEntry, t.TotalPartitions)
	for _, e := range t.Entries {
		for _, p := range e.Partitions {
			if p >= t.TotalPartitions {
				return nil, fmt.Errorf("Bad Partition entry (greater then total partitions): %s", e)
			}
			entriesPartition[p] = e
			partitionCount++
		}
	}

	if partitionCount != t.TotalPartitions {
		return nil, fmt.Errorf("Bad table, some partitions un accounted for")
	}

	t.DynMap = t.toDynMap()

	t.EntriesPartition = make([][]*RouterEntry, t.TotalPartitions)
	//Now setup the replication partitions.
	for _, e := range t.Entries {

		for _, p := range e.Partitions {
			pRep, err := t.repPartitions(p, e)
			if err != nil {
				return nil, fmt.Errorf("Bad table (%s)", err)
			}
			entries := make([]*RouterEntry, len(pRep)+1)
			entries[0] = e
			for i := 1; i < len(entries); i++ {
				entries[i] = entriesPartition[pRep[i-1]]
				e.PartitionsMap[pRep[i-1]] = false
			}
			t.EntriesPartition[p] = entries
		}
	}

	return t, nil
}

//internal use.
//skips the cached version
func (this *RouterTable) toDynMap() *dynmap.DynMap {
	mp := dynmap.NewDynMap()
	mp.Put("service", this.Service)
	mp.Put("revision", this.Revision)
	mp.Put("total_partitions", this.TotalPartitions)
	mp.Put("replication_factor", this.ReplicationFactor)
	mp.Put("partition_keys", this.PartitionKeys)

	entries := make([]*dynmap.DynMap, 0)
	for _, e := range this.Entries {
		entries = append(entries, e.ToDynMap())
	}
	mp.Put("entries", entries)
	this.DynMap = mp
	return mp
}

// Translate to a DynMap of the form:
// {
//     "service" : "trendrrdb",
//     "revision" : 898775762309309,
//     "total_partitions" : 256,
//     "entries" : [
//         {/*router entry 1*/},
//         {/*router entry 2*/}
//     ]
// }
func (this *RouterTable) ToDynMap() *dynmap.DynMap {
	if this.DynMap != nil && len(this.DynMap.Map) > 0 {
		return this.DynMap
	}
	return this.toDynMap()
}

//gets the partitions that should replicate this master.
func (this *RouterTable) repPartitions(partition int, entry *RouterEntry) ([]int, error) {
	//This method could be much better optimized, but
	//it is fairly rare, so we wont worry about it..
	entries := make([]int, 0)
	if partition >= this.TotalPartitions {
		return entries, fmt.Errorf("Requested partition %d is out of bounds (%d) ", partition, this.TotalPartitions)
	}

	for i := 1; i < this.TotalPartitions; i++ {
		par := (i + partition) % this.TotalPartitions
		v, ok := entry.PartitionsMap[par]
		if ok && v {
			//this is master.  skip to next one
			continue
		}
		entries = append(entries, par)
		//we check if len < repfactor -1 (minus one because we still need the current partition)
		if len(entries) == this.ReplicationFactor-1 {
			return entries, nil
		}
	}
	return entries, nil

}

// Gets the entries associated with the given partition
// [0] should be the master entry, and there should be
// table.ReplicationFactor number of entries
func (this *RouterTable) PartitionEntries(partition int) ([]*RouterEntry, error) {
	if partition >= this.TotalPartitions {
		return make([]*RouterEntry, 0), fmt.Errorf("Requested partition %d is out of bounds (%d) ", partition, this.TotalPartitions)
	}
	return this.EntriesPartition[partition], nil
}

type RouterEntry struct {
	//The address of this entry
	Address  string
	JsonPort int
	HttpPort int
	BinPort  int

	//list of partitions this entry is responsible for (master only)
	Partitions []int

	//Map of all the partitions this entry is responsible for.  true indicates master, false otherwise
	PartitionsMap map[int]bool
}

// Creates a new router entry from the dynmap passed in
func ToRouterEntry(mp *dynmap.DynMap) (*RouterEntry, error) {
	e := &RouterEntry{
		PartitionsMap: make(map[int]bool),
	}
	var ok bool
	e.Address, ok = mp.GetString("address")
	if !ok {
		return nil, fmt.Errorf("No Address in Entry: %s", mp)
	}

	e.JsonPort = mp.MustInt("ports.json", 0)
	e.HttpPort = mp.MustInt("ports.http", 0)
	e.BinPort = mp.MustInt("ports.bin", 0)

	e.Partitions, ok = mp.GetIntSlice("partitions")
	if !ok {
		e.Partitions = make([]int, 0)
	}
	for _, p := range e.Partitions {
		e.PartitionsMap[p] = true
	}
	return e, nil
}

//Id for this entry.  current is address:jsonport
func (this *RouterEntry) Id() string {
	return fmt.Sprintf("%s:%d", this.Address, this.JsonPort)
}

// Translate to a DynMap of the form:
// {
//     "id" : "localhost:8009"
//     "address" : "localhost",
//     "ports" : {
//         "json" : 8009,
//         "http" : 8010,
//	       "bin" : 8011
//     }
//     "partitions" : [1,2,3,4,5,6,7,8,9]
// }
func (this *RouterEntry) ToDynMap() *dynmap.DynMap {
	mp := dynmap.NewDynMap()
	mp.Put("address", this.Address)
	if this.JsonPort > 0 {
		mp.PutWithDot("ports.json", this.JsonPort)
	}

	if this.HttpPort > 0 {
		mp.PutWithDot("ports.http", this.HttpPort)
	}

	if this.BinPort > 0 {
		mp.PutWithDot("ports.bin", this.BinPort)
	}

	mp.Put("id", this.Id())
	mp.Put("partitions", this.Partitions)
	return mp
}
<|endoftext|>


<commit_before>package user

import (
	jwt "github.com/dgrijalva/jwt-go"
	"golang.org/x/crypto/bcrypt"
	"regexp"
	"strings"
	"time"

	"github.com/eirka/eirka-libs/db"
	e "github.com/eirka/eirka-libs/errors"
)

const (
	// jwt claim keys
	jwt_claim_issuer    = "iss"
	jwt_claim_issued    = "iat"
	jwt_claim_notbefore = "nbf"
	jwt_claim_expire    = "exp"
	jwt_claim_user_id   = "user_id"
	// jwt issuer
	jwt_issuer = "pram"
	// jwt expire days
	jwt_expire_days = 90

	// the username validation regex
	username = `^([a-zA-Z0-9]+[\s_-]?)+$`
)

var regexUsername = regexp.MustCompile(username)

// reserved name list
var reservedNameList = map[string]bool{
	"admin":          true,
	"administrator":  true,
	"administration": true,
	"support":        true,
	"mod":            true,
	"moderator":      true,
	"anon":           true,
	"anonymous":      true,
	"root":           true,
	"webmaster":      true,
	"username":       true,
	"user":           true,
}

// user struct
type User struct {
	Id              uint
	Name            string
	IsAuthenticated bool
	hash            []byte
	isPasswordValid bool
}

// create a user struct
func DefaultUser() User {
	return User{
		Id:              1,
		IsAuthenticated: false,
	}
}

// sets the user id
func (u *User) SetId(uid uint) {
	u.Id = uid
	return
}

// sets authenticated
func (u *User) SetAuthenticated() {
	// do not set auth for the wrong users
	if u.Id == 0 || u.Id == 1 {
		return
	}

	u.IsAuthenticated = true
	return
}

// check user struct validity
func (u *User) IsValid() bool {

	// this isnt a real user id
	if u.Id == 0 {
		return false
	}

	// the anon account can never be authenticated
	if u.Id == 1 && u.IsAuthenticated {
		return false
	}

	// a user can never be unauthenticated
	if u.Id != 1 && !u.IsAuthenticated {
		return false
	}

	return true
}

// checks if the name is valid
func IsValidName(name string) bool {

	if reservedNameList[strings.ToLower(strings.TrimSpace(name))] {
		return false
	}

	return regexUsername.MatchString(name)
}

// will get the password and name from the database for an instantiated user
func (u *User) Password() (err error) {

	// check user struct validity
	if !u.IsValid() {
		return e.ErrUserNotValid
	}

	// Get Database handle
	dbase, err := db.GetDb()
	if err != nil {
		return
	}

	// get hashed password from database
	err = dbase.QueryRow("select user_name, user_password from users where user_id = ?", u.Id).Scan(&u.Name, &u.hash)
	if err != nil {
		return
	}

	return
}

// will get the password and user id from the database for a user name
func (u *User) FromName(name string) (err error) {

	// password length cant be 0
	if len(name) == 0 {
		return e.ErrUserNotValid
	}

	// Get Database handle
	dbase, err := db.GetDb()
	if err != nil {
		return
	}

	// get hashed password from database
	err = dbase.QueryRow("select user_id, user_password from users where user_name = ?", name).Scan(&u.Id, &u.hash)
	if err != nil {
		return
	}

	u.SetAuthenticated()

	if !u.IsValid() {
		return e.ErrUserNotValid
	}

	return

}

// check for duplicate name before registering
func CheckDuplicate(name string) (check bool) {

	// password length cant be 0
	if len(name) == 0 {
		return true
	}

	// Get Database handle
	dbase, err := db.GetDb()
	if err != nil {
		return true
	}

	// this will return true if there is a user
	err = dbase.QueryRow("select count(*) from users where user_name = ?", name).Scan(&check)
	if err != nil {
		return true
	}

	return

}

// Creates a JWT token with our claims
func (u *User) CreateToken() (newtoken string, err error) {

	// error if theres no secret set
	if Secret == "" {
		err = e.ErrNoSecret
		return
	}

	// check user struct validity
	if !u.IsValid() {
		err = e.ErrUserNotValid
		return
	}

	// a token should never be created
	if u.Id == 0 || u.Id == 1 {
		err = e.ErrUserNotValid
		return
	}

	// a token should never be created
	if !u.IsAuthenticated {
		err = e.ErrUserNotValid
		return
	}

	// check if password was valid
	if !u.isPasswordValid {
		err = e.ErrInvalidPassword
		return
	}

	// Create the token
	token := jwt.New(jwt.SigningMethodHS256)
	// the current time
	now := time.Now()

	// Set our claims
	token.Claims[jwt_claim_issuer] = jwt_issuer
	token.Claims[jwt_claim_issued] = now.Unix()
	token.Claims[jwt_claim_notbefore] = now.Unix()
	token.Claims[jwt_claim_expire] = now.Add(time.Hour * 24 * jwt_expire_days).Unix()
	token.Claims[jwt_claim_user_id] = u.Id

	// Sign and get the complete encoded token as a string
	newtoken, err = token.SignedString([]byte(Secret))
	if err != nil {
		return
	}

	return

}

// get the user info from id
func (u *User) IsAuthorized(ib uint) bool {

	var err error

	if !u.IsValid() {
		return false
	}

	// check for invalid stuff
	if ib == 0 {
		return false
	}

	// Get Database handle
	dbase, err := db.GetDb()
	if err != nil {
		return false
	}

	// holds our role
	var role uint

	// get data from users table
	err = dbase.QueryRow(`SELECT COALESCE((SELECT MAX(role_id) FROM user_ib_role_map WHERE user_ib_role_map.user_id = users.user_id AND ib_id = ?),user_role_map.role_id) as role
    FROM users
    INNER JOIN user_role_map ON (user_role_map.user_id = users.user_id)
    WHERE users.user_id = ?`, ib, u.Id).Scan(&role)
	if err != nil {
		return false
	}

	switch role {
	case 3:
		return true
	case 4:
		return true
	default:
		return false
	}

	return false

}

// compare password to
func (u *User) ComparePassword(password string) bool {

	// password length cant be 0
	if len(password) == 0 {
		return false
	}

	// if the hash wasnt populated
	if len(u.hash) == 0 {
		return false
	}

	// compare the stored hash with the provided password
	err := bcrypt.CompareHashAndPassword(u.hash, []byte(password))

	// we only want jwt tokens to be created after a valid password has been given
	u.isPasswordValid = err == nil

	return u.isPasswordValid

}

// hash a given password
func HashPassword(password string) (hash []byte, err error) {

	// password length cant be 0
	if len(password) == 0 {
		err = e.ErrInvalidPassword
		return
	}

	return bcrypt.GenerateFromPassword([]byte(password), bcrypt.DefaultCost)
}
<commit_msg>add user interface<commit_after>package user

import (
	jwt "github.com/dgrijalva/jwt-go"
	"golang.org/x/crypto/bcrypt"
	"regexp"
	"strings"
	"time"

	"github.com/eirka/eirka-libs/db"
	e "github.com/eirka/eirka-libs/errors"
)

const (
	// jwt claim keys
	jwt_claim_issuer    = "iss"
	jwt_claim_issued    = "iat"
	jwt_claim_notbefore = "nbf"
	jwt_claim_expire    = "exp"
	jwt_claim_user_id   = "user_id"
	// jwt issuer
	jwt_issuer = "pram"
	// jwt expire days
	jwt_expire_days = 90

	// the username validation regex
	username = `^([a-zA-Z0-9]+[\s_-]?)+$`
)

var regexUsername = regexp.MustCompile(username)

// reserved name list
var reservedNameList = map[string]bool{
	"admin":          true,
	"administrator":  true,
	"administration": true,
	"support":        true,
	"mod":            true,
	"moderator":      true,
	"anon":           true,
	"anonymous":      true,
	"root":           true,
	"webmaster":      true,
	"username":       true,
	"user":           true,
}

type Authenticator interface {
	IsValid() bool
	IsAuthorized(ib uint) bool
	SetId(uid uint)
	SetAuthenticated()
	Password() (err error)
	ComparePassword(password string) bool
	FromName(name string) (err error)
	CreateToken() (newtoken string, err error)
}

// user struct
type User struct {
	Id              uint
	Name            string
	IsAuthenticated bool
	hash            []byte
	isPasswordValid bool
}

var _ = Authenticator(&User{})

// create a user struct
func DefaultUser() *User {
	return &User{
		Id:              1,
		IsAuthenticated: false,
	}
}

// sets the user id
func (u *User) SetId(uid uint) {
	u.Id = uid
	return
}

// sets authenticated
func (u *User) SetAuthenticated() {
	// do not set auth for the wrong users
	if u.Id == 0 || u.Id == 1 {
		return
	}

	u.IsAuthenticated = true
	return
}

// check user struct validity
func (u *User) IsValid() bool {

	// this isnt a real user id
	if u.Id == 0 {
		return false
	}

	// the anon account can never be authenticated
	if u.Id == 1 && u.IsAuthenticated {
		return false
	}

	// a user can never be unauthenticated
	if u.Id != 1 && !u.IsAuthenticated {
		return false
	}

	return true
}

// checks if the name is valid
func IsValidName(name string) bool {

	if reservedNameList[strings.ToLower(strings.TrimSpace(name))] {
		return false
	}

	return regexUsername.MatchString(name)
}

// will get the password and name from the database for an instantiated user
func (u *User) Password() (err error) {

	// check user struct validity
	if !u.IsValid() {
		return e.ErrUserNotValid
	}

	// Get Database handle
	dbase, err := db.GetDb()
	if err != nil {
		return
	}

	// get hashed password from database
	err = dbase.QueryRow("select user_name, user_password from users where user_id = ?", u.Id).Scan(&u.Name, &u.hash)
	if err != nil {
		return
	}

	return
}

// will get the password and user id from the database for a user name
func (u *User) FromName(name string) (err error) {

	// password length cant be 0
	if len(name) == 0 {
		return e.ErrUserNotValid
	}

	// Get Database handle
	dbase, err := db.GetDb()
	if err != nil {
		return
	}

	// get hashed password from database
	err = dbase.QueryRow("select user_id, user_password from users where user_name = ?", name).Scan(&u.Id, &u.hash)
	if err != nil {
		return
	}

	u.SetAuthenticated()

	if !u.IsValid() {
		return e.ErrUserNotValid
	}

	return

}

// check for duplicate name before registering
func CheckDuplicate(name string) (check bool) {

	// password length cant be 0
	if len(name) == 0 {
		return true
	}

	// Get Database handle
	dbase, err := db.GetDb()
	if err != nil {
		return true
	}

	// this will return true if there is a user
	err = dbase.QueryRow("select count(*) from users where user_name = ?", name).Scan(&check)
	if err != nil {
		return true
	}

	return

}

// Creates a JWT token with our claims
func (u *User) CreateToken() (newtoken string, err error) {

	// error if theres no secret set
	if Secret == "" {
		err = e.ErrNoSecret
		return
	}

	// check user struct validity
	if !u.IsValid() {
		err = e.ErrUserNotValid
		return
	}

	// a token should never be created
	if u.Id == 0 || u.Id == 1 {
		err = e.ErrUserNotValid
		return
	}

	// a token should never be created
	if !u.IsAuthenticated {
		err = e.ErrUserNotValid
		return
	}

	// check if password was valid
	if !u.isPasswordValid {
		err = e.ErrInvalidPassword
		return
	}

	// Create the token
	token := jwt.New(jwt.SigningMethodHS256)
	// the current time
	now := time.Now()

	// Set our claims
	token.Claims[jwt_claim_issuer] = jwt_issuer
	token.Claims[jwt_claim_issued] = now.Unix()
	token.Claims[jwt_claim_notbefore] = now.Unix()
	token.Claims[jwt_claim_expire] = now.Add(time.Hour * 24 * jwt_expire_days).Unix()
	token.Claims[jwt_claim_user_id] = u.Id

	// Sign and get the complete encoded token as a string
	newtoken, err = token.SignedString([]byte(Secret))
	if err != nil {
		return
	}

	return

}

// get the user info from id
func (u *User) IsAuthorized(ib uint) bool {

	var err error

	if !u.IsValid() {
		return false
	}

	// check for invalid stuff
	if ib == 0 {
		return false
	}

	// Get Database handle
	dbase, err := db.GetDb()
	if err != nil {
		return false
	}

	// holds our role
	var role uint

	// get data from users table
	err = dbase.QueryRow(`SELECT COALESCE((SELECT MAX(role_id) FROM user_ib_role_map WHERE user_ib_role_map.user_id = users.user_id AND ib_id = ?),user_role_map.role_id) as role
    FROM users
    INNER JOIN user_role_map ON (user_role_map.user_id = users.user_id)
    WHERE users.user_id = ?`, ib, u.Id).Scan(&role)
	if err != nil {
		return false
	}

	switch role {
	case 3:
		return true
	case 4:
		return true
	default:
		return false
	}

	return false

}

// compare password to
func (u *User) ComparePassword(password string) bool {

	// password length cant be 0
	if len(password) == 0 {
		return false
	}

	// if the hash wasnt populated
	if len(u.hash) == 0 {
		return false
	}

	// compare the stored hash with the provided password
	err := bcrypt.CompareHashAndPassword(u.hash, []byte(password))

	// we only want jwt tokens to be created after a valid password has been given
	u.isPasswordValid = err == nil

	return u.isPasswordValid

}

// hash a given password
func HashPassword(password string) (hash []byte, err error) {

	// password length cant be 0
	if len(password) == 0 {
		err = e.ErrInvalidPassword
		return
	}

	return bcrypt.GenerateFromPassword([]byte(password), bcrypt.DefaultCost)
}
<|endoftext|>


<commit_before>package dondeestas

import (
	"bytes"
	"encoding/json"
	"io"
	"io/ioutil"
	"testing"
)

func stringToReadCloser(str string) io.ReadCloser {
	return ioutil.NopCloser(bytes.NewReader(bytes.NewBufferString(str).Bytes()))
}

func TestReadCloserToJson(t *testing.T) {
    // TODO: make this just a JSON object, not a Person
	expectedPerson, _ := createRandomPerson()
	expectedPersonJson, _ := json.Marshal(expectedPerson)
	expectedPersonStr := string(expectedPersonJson)

	var person Person
	if err := ReadCloserJsonToStruct(stringToReadCloser(expectedPersonStr), person); err != nil {
		t.Fatalf("Encountered error when retrieving json from string: %s", err)
	}

	t.Skip("TODO")

	// Test forcing the function to read a closed stream
	/*
    req = httptest.NewRequest("GET", "http://blah.com/foo", bytes.NewBufferString(expectedPersonStr))
	if err := req.Body.Close(); err != nil {
		t.Fatalf("Could not close test request body!")
	}
	if err := ReadCloserJsonToStruct(req.Body, person); err == nil {
		t.Error("Did not receive expected error when reading closed stream")
	}

	// Incorrect JSON object
	if err := ReadCloserJsonToStruct(ioutil.NopCloser(bytes.NewReader(bytes.NewBufferString(`{"id":"foo"}`).Bytes())), person); err == nil  {
		t.Error("Did not receive expected error on bad JSON unmarshalling")
	}
    */
}
<commit_msg>Renaming the function because...<commit_after>package dondeestas

import (
	"bytes"
	"encoding/json"
	"io"
	"io/ioutil"
	"testing"
)

func stringToReadCloser(str string) io.ReadCloser {
	return ioutil.NopCloser(bytes.NewReader(bytes.NewBufferString(str).Bytes()))
}

func TestReadCloserJsonToStruct(t *testing.T) {
    // TODO: make this just a JSON object, not a Person
	expectedPerson, _ := createRandomPerson()
	expectedPersonJson, _ := json.Marshal(expectedPerson)
	expectedPersonStr := string(expectedPersonJson)

	var person Person
	if err := readCloserJsonToStruct(stringToReadCloser(expectedPersonStr), person); err != nil {
		t.Fatalf("Encountered error when retrieving json from string: %s", err)
	}

	t.Skip("TODO")

	// Test forcing the function to read a closed stream
	/*
    req = httptest.NewRequest("GET", "http://blah.com/foo", bytes.NewBufferString(expectedPersonStr))
	if err := req.Body.Close(); err != nil {
		t.Fatalf("Could not close test request body!")
	}
	if err := readCloserJsonToStruct(req.Body, person); err == nil {
		t.Error("Did not receive expected error when reading closed stream")
	}

	// Incorrect JSON object
	if err := readCloserJsonToStruct(ioutil.NopCloser(bytes.NewReader(bytes.NewBufferString(`{"id":"foo"}`).Bytes())), person); err == nil  {
		t.Error("Did not receive expected error on bad JSON unmarshalling")
	}
    */
}
<|endoftext|>


<commit_before>package main

import (
	"log"
	"regexp"

	"github.com/hashicorp/consul/api"
)

// get a client handle for a specified address (or the local agent if "")
func getClient(address string) (*api.Client, error) {
	config := api.DefaultConfig()
	if address != "" {
		config.Address = address
	}
	return api.NewClient(config)
}

// get a list of all services, limit to those matching the search criteria
func getList(serviceString string, tag string) []*api.ServiceEntry {
	client, err := getClient("")
	if err != nil {
		log.Fatalf("Unable to get a consul client connection: %s\n", err)
	}

	serviceList, _, err := client.Catalog().Services(nil)
	if err != nil {
		log.Fatalf("Unable to get list of services from catalog: %s", err)
	}

	nodeServiceMap := make(map[string]map[string]*api.ServiceEntry)

	for svc := range serviceList {
		entries, _, err := client.Health().Service(svc, tag, false, nil)
		if err != nil {
			log.Fatalf("Unable to query for service \"%s\" from catalog: %s", svc, err)
		}
		for _, entry := range entries {
			if _, ok := nodeServiceMap[entry.Node.Node]; !ok {
				nodeServiceMap[entry.Node.Node] = make(map[string]*api.ServiceEntry)
			}
			nodeServiceMap[entry.Node.Node][entry.Service.ID] = entry
		}
	}

	// get a handle to the health endpoint and pre-calculate the regexp
	var re *regexp.Regexp
	if serviceString != "" {
		re = regexp.MustCompile(serviceString)
	}

	// prepare a slice to hold the result list
	seOut := make([]*api.ServiceEntry, 0)

	for _, services := range nodeServiceMap {
		for serviceID, service := range services {
			match := true
			if re != nil {
				idStr := re.FindString(serviceID)
				nameStr := re.FindString(service.Service.ID)
				match = idStr != "" || nameStr != ""
			}
			if match {
				seOut = append(seOut, service)
			}
		}
	}

	return seOut
}

// A service entry is considered healty if all the eligible checks are passing.
// serfChecks are not eligible
func isHealthy(se *api.ServiceEntry) bool {
	if se == nil || se.Checks == nil || len(se.Checks) == 0 {
		// No checks = failing
		return false
	}

	var healthy bool = true
	var eligible int = 0
	for _, c := range se.Checks {
		if c.Name == "serfHealth" {
			continue
		}
		// All found checks have to be passing
		healthy = healthy && (c.Status == "passing")
		eligible++
	}

	// No eligible checks were found
	if eligible == 0 {
		return false
	}

	return healthy
}
<commit_msg>simplifying a bit.<commit_after>package main

import (
	"log"
	"regexp"

	"github.com/hashicorp/consul/api"
)

// get a client handle for a specified address (or the local agent if "")
func getClient(address string) (*api.Client, error) {
	config := api.DefaultConfig()
	if address != "" {
		config.Address = address
	}
	return api.NewClient(config)
}

// get a list of all services, limit to those matching the search criteria
func getList(serviceString string, tag string) []*api.ServiceEntry {
	client, err := getClient("")
	if err != nil {
		log.Fatalf("Unable to get a consul client connection: %s\n", err)
	}

	serviceList, _, err := client.Catalog().Services(nil)
	if err != nil {
		log.Fatalf("Unable to get list of services from catalog: %s", err)
	}

	nodeServiceMap := make(map[string]map[string]*api.ServiceEntry)

	for svc := range serviceList {
		entries, _, err := client.Health().Service(svc, tag, false, nil)
		if err != nil {
			log.Fatalf("Unable to query for service \"%s\" from catalog: %s", svc, err)
		}
		for _, entry := range entries {
			if _, ok := nodeServiceMap[entry.Node.Node]; !ok {
				nodeServiceMap[entry.Node.Node] = make(map[string]*api.ServiceEntry)
			}
			nodeServiceMap[entry.Node.Node][entry.Service.ID] = entry
		}
	}

	seOut := make([]*api.ServiceEntry, 0)
	if serviceString == "" {
		for _, services := range nodeServiceMap {
			for _, service := range services {
				seOut = append(seOut, service)
			}
		}
	} else {
		var serviceID string
		var services map[string]*api.ServiceEntry
		var service *api.ServiceEntry
		re := regexp.MustCompile(serviceString)
		for _, services = range nodeServiceMap {
			for serviceID, service = range services {
				if re.FindString(serviceID) != "" || re.FindString(service.Service.ID) != "" {
					seOut = append(seOut, service)
				}
			}
		}
	}

	return seOut
}

// A service entry is considered healty if all the eligible checks are passing.
// serfChecks are not eligible
func isHealthy(se *api.ServiceEntry) bool {
	if se == nil || se.Checks == nil || len(se.Checks) == 0 {
		// No checks = failing
		return false
	}

	healthy := true
	eligible := 0
	for _, c := range se.Checks {
		if c.Name == "serfHealth" {
			continue
		}
		// All found checks have to be passing
		healthy = healthy && (c.Status == "passing")
		eligible++
	}

	// No eligible checks were found
	if eligible == 0 {
		return false
	}

	return healthy
}
<|endoftext|>