<commit_before>// +build darwin freebsd linux netbsd openbsd
// +build cgo

package group

import (
	"fmt"
	"runtime"
	"strconv"
	"syscall"
	"unsafe"
)

/*
#include <unistd.h>
#include <sys/types.h>
#include <grp.h>
#include <stdlib.h>

static int mygetgrgid_r(int gid, struct group *grp,
	char *buf, size_t buflen, struct group **result) {
 return getgrgid_r(gid, grp, buf, buflen, result);
}

static char **next_mem(char **mem) {
	return ++mem;
}
*/
import "C"

func lookup(groupname string) (*Group, error) {
	return lookupUnix(-1, groupname, true)
}

func lookupId(gid string) (*Group, error) {
	i, e := strconv.Atoi(gid)
	if e != nil {
		return nil, e
	}
	return lookupUnix(i, "", false)
}

func lookupUnix(gid int, groupname string, lookupByName bool) (*Group, error) {
	var grp C.struct_group
	var result *C.struct_group

	var bufSize C.long
	if runtime.GOOS == "freebsd" {
		// FreeBSD doesn't have _SC_GETGR_R_SIZE_MAX
		// and just returns -1.  So just use the same
		// size that Linux returns
		// TODO: Confirm this!
		bufSize = 1024
	} else {
		bufSize = C.sysconf(C._SC_GETGR_R_SIZE_MAX)
		if bufSize <= 0 || bufSize > 1<<20 {
			return nil, fmt.Errorf("group: unreasonable _SC_GETGR_R_SIZE_MAX of %d", bufSize)
		}
	}
	buf := C.malloc(C.size_t(bufSize))
	defer C.free(buf)
	var rv C.int
	if lookupByName {
		nameC := C.CString(groupname)
		defer C.free(unsafe.Pointer(nameC))
		rv = C.getgrnam_r(nameC,
			&grp,
			(*C.char)(buf),
			C.size_t(bufSize),
			&result)
		if rv != 0 {
			return nil, fmt.Errorf("group: lookup groupname %s: %s", groupname, syscall.Errno(rv))
		}
		if result == nil {
			return nil, UnknownGroupError(groupname)
		}
	} else {
		// mygetgrgid_r is a wrapper around getgrgid_r to
		// avoid using gid_t because C.gid_t(gid) for
		// unknown reasons doesn't work on linux.
		rv = C.mygetgrgid_r(C.int(gid),
			&grp,
			(*C.char)(buf),
			C.size_t(bufSize),
			&result)
		if rv != 0 {
			return nil, fmt.Errorf("group: lookup groupid %d: %s", gid, syscall.Errno(rv))
		}
		if result == nil {
			return nil, UnknownGroupIdError(gid)
		}
	}

	members := make([]string, 0)
	if grp.gr_mem != nil {
		for mem := grp.gr_mem; *mem != nil; mem = C.next_mem(mem) {
			members = append(members, C.GoString(*mem))
		}
	}
	g := &Group{
		Gid:     strconv.Itoa(int(grp.gr_gid)),
		Name:    C.GoString(grp.gr_name),
		Members: members,
	}
	return g, nil
}
<commit_msg>Calculate member count and alloc members once without repeating append().<commit_after>// +build darwin freebsd linux netbsd openbsd
// +build cgo

package group

import (
	"fmt"
	"runtime"
	"strconv"
	"syscall"
	"unsafe"
)

/*
#include <unistd.h>
#include <sys/types.h>
#include <grp.h>
#include <stdlib.h>

static int mygetgrgid_r(int gid, struct group *grp,
	char *buf, size_t buflen, struct group **result) {
 return getgrgid_r(gid, grp, buf, buflen, result);
}

static char **next_mem(char **mem) {
	return ++mem;
}

static int member_count(char **mem) {
	int c = 0;
	for (; *mem; ++mem) {
		++c;
	}
	return c;
}
*/
import "C"

func lookup(groupname string) (*Group, error) {
	return lookupUnix(-1, groupname, true)
}

func lookupId(gid string) (*Group, error) {
	i, e := strconv.Atoi(gid)
	if e != nil {
		return nil, e
	}
	return lookupUnix(i, "", false)
}

func lookupUnix(gid int, groupname string, lookupByName bool) (*Group, error) {
	var grp C.struct_group
	var result *C.struct_group

	var bufSize C.long
	if runtime.GOOS == "freebsd" {
		// FreeBSD doesn't have _SC_GETGR_R_SIZE_MAX
		// and just returns -1.  So just use the same
		// size that Linux returns
		// TODO: Confirm this!
		bufSize = 1024
	} else {
		bufSize = C.sysconf(C._SC_GETGR_R_SIZE_MAX)
		if bufSize <= 0 || bufSize > 1<<20 {
			return nil, fmt.Errorf("group: unreasonable _SC_GETGR_R_SIZE_MAX of %d", bufSize)
		}
	}
	buf := C.malloc(C.size_t(bufSize))
	defer C.free(buf)
	var rv C.int
	if lookupByName {
		nameC := C.CString(groupname)
		defer C.free(unsafe.Pointer(nameC))
		rv = C.getgrnam_r(nameC,
			&grp,
			(*C.char)(buf),
			C.size_t(bufSize),
			&result)
		if rv != 0 {
			return nil, fmt.Errorf("group: lookup groupname %s: %s", groupname, syscall.Errno(rv))
		}
		if result == nil {
			return nil, UnknownGroupError(groupname)
		}
	} else {
		// mygetgrgid_r is a wrapper around getgrgid_r to
		// avoid using gid_t because C.gid_t(gid) for
		// unknown reasons doesn't work on linux.
		rv = C.mygetgrgid_r(C.int(gid),
			&grp,
			(*C.char)(buf),
			C.size_t(bufSize),
			&result)
		if rv != 0 {
			return nil, fmt.Errorf("group: lookup groupid %d: %s", gid, syscall.Errno(rv))
		}
		if result == nil {
			return nil, UnknownGroupIdError(gid)
		}
	}

	var members []string
	if grp.gr_mem != nil {
		members = make([]string, C.member_count(grp.gr_mem))
		i := 0
		for mem := grp.gr_mem; *mem != nil; mem = C.next_mem(mem) {
			members[i] = C.GoString(*mem)
			i += 1
		}
	} else {
		members = make([]string, 0)
	}
	g := &Group{
		Gid:     strconv.Itoa(int(grp.gr_gid)),
		Name:    C.GoString(grp.gr_name),
		Members: members,
	}
	return g, nil
}
<|endoftext|>


<commit_before><commit_msg>log configured hosts<commit_after><|endoftext|>


<commit_before><commit_msg>Add Rackspace ORD<commit_after><|endoftext|>


<commit_before>// Package db handles new connections and transactions with timeouts.
//
// Timeouts that can be defined:
//    Connect timeout         - sets the maximum wait for creating a new connection
//    Statement timeout       - sets the maximum wait for queries completion
//    New transaction timeout - sets the maximum wait for opening a new transaction
//
// This feature is useful when the servers must be kept up, even though your database servers are
// not stable. So the database clients can respond their users with kindly error messages or look
// for other database servers - fail fast is usually better than keep server resources busy.
package db

import (
	"database/sql"
	"errors"
	"fmt"
	"sync"
	"time"
)

var (
	// ErrNewTxTimedOut returned when takes too long to retrieve a transaction
	// from the database server.
	ErrNewTxTimedOut = errors.New("new transaction timed out")

	// ErrUnreachable returned when the database is current unreachable.
	ErrUnreachable = errors.New("unreachable database")
)

// PostgresDriver holds the postgres driver name. You should import the driver somewhere in your
// code. The default value for PostgresDriver was get from https://github.com/lib/pq. This driver
// isn't directly imported here to don't stick this library with only one type of Postgres driver,
// and to make it more flexible for testing.
//
//     import (
//       _ "github.com/lib/pq"
//     )
var PostgresDriver = "postgres"

// ConnParams is all the data needed to create a new connection with database.
// If you are looking for default values see NewConnParams() function that set
// defaults values to Data structure.
type ConnParams struct {
	Username     string
	Password     string
	DatabaseName string
	Host         string

	// ConnectTimeout is the timeout utilized to creating a new connection with database. It is not
	// recommended to use a timeout of less than 2 seconds.
	ConnectTimeout time.Duration
	// StatementTimeout is the timeout utilized to invalidate queries that last more than the
	// configured timeout.
	StatementTimeout time.Duration

	MaxIdleConnections int
	MaxOpenConnections int
}

// NewConnParams returns the database connection parameters with some default
// values.
func NewConnParams() ConnParams {
	return ConnParams{
		Host:               "127.0.0.1",
		ConnectTimeout:     3 * time.Second,
		StatementTimeout:   10 * time.Second,
		MaxIdleConnections: 16,
		MaxOpenConnections: 32,
	}
}

// ConnectPostgres connects to a postgres database using the values from d. In case of a successfully
// connection it returns a sql.DB and a nil error. In case of problem it returns a nil sql.DB and an
// error from sql.Open (standard library, see https://golang.org/pkg/database/sql/#Open)
func ConnectPostgres(d ConnParams) (db *sql.DB, err error) {
	// connect_timeout
	//
	// https://www.postgresql.org/docs/9.6/static/libpq-connect.html#LIBPQ-CONNECT-CONNECT-TIMEOUT
	// Maximum wait for connection, in seconds (write as a decimal integer string). Zero or not
	// specified means wait indefinitely. It is not recommended to use a timeout of less than 2
	// seconds.
	//
	// statement_timeout
	//
	// https://www.postgresql.org/docs/9.6/static/runtime-config-client.html#GUC-STATEMENT-TIMEOUT
	// Abort any statement that takes more than the specified number of milliseconds, starting from
	// the time the command arrives at the server from the client.
	connParams := fmt.Sprintf(
		"user=%s password=%s dbname=%s sslmode=disable host=%s connect_timeout=%d statement_timeout=%d",
		d.Username,
		d.Password,
		d.DatabaseName,
		d.Host,
		int(d.ConnectTimeout.Seconds()),
		int(d.StatementTimeout.Seconds()*1000),
	)

	if db, err = sql.Open(PostgresDriver, connParams); err != nil {
		return nil, err
	}

	db.SetMaxIdleConns(d.MaxIdleConnections)
	db.SetMaxOpenConns(d.MaxOpenConnections)
	return
}

// DB is a layer over the sql.DB to allow fail-fast approach when a database is
// unreachable and to allow carrying the transaction timeout.
type DB struct {
	*sql.DB
	checker   dbChecker
	txTimeout time.Duration
}

// NewDB build a DB layer over the sql.DB to add transaction timeout feature
// with fail-fast capability.
func NewDB(db *sql.DB, txTimeout time.Duration) *DB {
	return &DB{
		DB:        db,
		txTimeout: txTimeout,
	}
}

// Begin starts a new database transaction with a timeout. Starting a new
// transaction is not supposed to last too long, so a timeout of more than 3
// seconds it's usually not necessary.
func (db *DB) Begin() (*sql.Tx, error) {
	if db.checker.unreachable() {
		return nil, ErrUnreachable
	}

	tx, err := newTx(db.DB, db.txTimeout)
	if err == ErrNewTxTimedOut {
		go func() {
			db.checker.check(db.DB, db.txTimeout)
		}()
	}

	return tx, err
}

func newTx(db *sql.DB, timeout time.Duration) (*sql.Tx, error) {
	// The channels has size of 1 (buffered) to avoid keeping an unnecessary goroutine blocked in
	// memory. For example: a goroutine is spawn, and it returns via channel a new transaction or
	// an error. After spawning a goroutine the program blocks in the select statement waiting
	// until the first channel message. In case of a timeout message, the spawned goroutine will
	// put a message in one of this two channels (ch and chErr) and simply returns (die), the
	// program don't care about the messages, because it has already timed out. If the channels
	// were not buffered the goroutine would be blocked trying to put a message into the channel
	// until the program dies.
	ch := make(chan *sql.Tx, 1)
	chErr := make(chan error, 1)

	go func() {
		tx, err := db.Begin()
		if err != nil {
			chErr <- err
			return
		}

		ch <- tx
	}()

	select {
	case tx := <-ch:
		return tx, nil
	case err := <-chErr:
		return nil, err
	case <-time.After(timeout):
		return nil, ErrNewTxTimedOut
	}
}

// Unreachable fast way to check if a database connection is unreachable without
// creating a new transaction.
func Unreachable(db *DB) bool {
	return db.checker.unreachable()
}

type dbChecker struct {
	sync.RWMutex
	checking bool
}

func (d *dbChecker) unreachable() bool {
	d.RLock()
	defer d.RUnlock()
	return d.checking
}

func (d *dbChecker) start() {
	d.Lock()
	defer d.Unlock()
	d.checking = true
}

func (d *dbChecker) stop() {
	d.Lock()
	defer d.Unlock()
	d.checking = false
}

func (d *dbChecker) check(db *sql.DB, duration time.Duration) {
	if d.unreachable() {
		// already checking
		return
	}

	d.start()
	for range time.Tick(2 * time.Second) {
		tx, err := newTx(db, duration)
		if err != nil {
			continue
		}

		// db is back!
		d.stop()
		tx.Commit()
		return
	}
}
<commit_msg>Avoid more than 1 checker running at the same time<commit_after>// Package db handles new connections and transactions with timeouts.
//
// Timeouts that can be defined:
//    Connect timeout         - sets the maximum wait for creating a new connection
//    Statement timeout       - sets the maximum wait for queries completion
//    New transaction timeout - sets the maximum wait for opening a new transaction
//
// This feature is useful when the servers must be kept up, even though your database servers are
// not stable. So the database clients can respond their users with kindly error messages or look
// for other database servers - fail fast is usually better than keep server resources busy.
package db

import (
	"database/sql"
	"errors"
	"fmt"
	"sync"
	"time"
)

var (
	// ErrNewTxTimedOut returned when takes too long to retrieve a transaction
	// from the database server.
	ErrNewTxTimedOut = errors.New("new transaction timed out")

	// ErrUnreachable returned when the database is current unreachable.
	ErrUnreachable = errors.New("unreachable database")
)

// PostgresDriver holds the postgres driver name. You should import the driver somewhere in your
// code. The default value for PostgresDriver was get from https://github.com/lib/pq. This driver
// isn't directly imported here to don't stick this library with only one type of Postgres driver,
// and to make it more flexible for testing.
//
//     import (
//       _ "github.com/lib/pq"
//     )
var PostgresDriver = "postgres"

// ConnParams is all the data needed to create a new connection with database.
// If you are looking for default values see NewConnParams() function that set
// defaults values to Data structure.
type ConnParams struct {
	Username     string
	Password     string
	DatabaseName string
	Host         string

	// ConnectTimeout is the timeout utilized to creating a new connection with database. It is not
	// recommended to use a timeout of less than 2 seconds.
	ConnectTimeout time.Duration
	// StatementTimeout is the timeout utilized to invalidate queries that last more than the
	// configured timeout.
	StatementTimeout time.Duration

	MaxIdleConnections int
	MaxOpenConnections int
}

// NewConnParams returns the database connection parameters with some default
// values.
func NewConnParams() ConnParams {
	return ConnParams{
		Host:               "127.0.0.1",
		ConnectTimeout:     3 * time.Second,
		StatementTimeout:   10 * time.Second,
		MaxIdleConnections: 16,
		MaxOpenConnections: 32,
	}
}

// ConnectPostgres connects to a postgres database using the values from d. In case of a successfully
// connection it returns a sql.DB and a nil error. In case of problem it returns a nil sql.DB and an
// error from sql.Open (standard library, see https://golang.org/pkg/database/sql/#Open)
func ConnectPostgres(d ConnParams) (db *sql.DB, err error) {
	// connect_timeout
	//
	// https://www.postgresql.org/docs/9.6/static/libpq-connect.html#LIBPQ-CONNECT-CONNECT-TIMEOUT
	// Maximum wait for connection, in seconds (write as a decimal integer string). Zero or not
	// specified means wait indefinitely. It is not recommended to use a timeout of less than 2
	// seconds.
	//
	// statement_timeout
	//
	// https://www.postgresql.org/docs/9.6/static/runtime-config-client.html#GUC-STATEMENT-TIMEOUT
	// Abort any statement that takes more than the specified number of milliseconds, starting from
	// the time the command arrives at the server from the client.
	connParams := fmt.Sprintf(
		"user=%s password=%s dbname=%s sslmode=disable host=%s connect_timeout=%d statement_timeout=%d",
		d.Username,
		d.Password,
		d.DatabaseName,
		d.Host,
		int(d.ConnectTimeout.Seconds()),
		int(d.StatementTimeout.Seconds()*1000),
	)

	if db, err = sql.Open(PostgresDriver, connParams); err != nil {
		return nil, err
	}

	db.SetMaxIdleConns(d.MaxIdleConnections)
	db.SetMaxOpenConns(d.MaxOpenConnections)
	return
}

// DB is a layer over the sql.DB to allow fail-fast approach when a database is
// unreachable and to allow carrying the transaction timeout.
type DB struct {
	*sql.DB
	checker   dbChecker
	txTimeout time.Duration
}

// NewDB build a DB layer over the sql.DB to add transaction timeout feature
// with fail-fast capability.
func NewDB(db *sql.DB, txTimeout time.Duration) *DB {
	return &DB{
		DB:        db,
		txTimeout: txTimeout,
	}
}

// Begin starts a new database transaction with a timeout. Starting a new
// transaction is not supposed to last too long, so a timeout of more than 3
// seconds it's usually not necessary.
func (db *DB) Begin() (*sql.Tx, error) {
	if db.checker.unreachable() {
		return nil, ErrUnreachable
	}

	tx, err := newTx(db.DB, db.txTimeout)
	if err == ErrNewTxTimedOut {
		go func() {
			db.checker.check(db.DB, db.txTimeout)
		}()
	}

	return tx, err
}

func newTx(db *sql.DB, timeout time.Duration) (*sql.Tx, error) {
	// The channels has size of 1 (buffered) to avoid keeping an unnecessary goroutine blocked in
	// memory. For example: a goroutine is spawn, and it returns via channel a new transaction or
	// an error. After spawning a goroutine the program blocks in the select statement waiting
	// until the first channel message. In case of a timeout message, the spawned goroutine will
	// put a message in one of this two channels (ch and chErr) and simply returns (die), the
	// program don't care about the messages, because it has already timed out. If the channels
	// were not buffered the goroutine would be blocked trying to put a message into the channel
	// until the program dies.
	ch := make(chan *sql.Tx, 1)
	chErr := make(chan error, 1)

	go func() {
		tx, err := db.Begin()
		if err != nil {
			chErr <- err
			return
		}

		ch <- tx
	}()

	select {
	case tx := <-ch:
		return tx, nil
	case err := <-chErr:
		return nil, err
	case <-time.After(timeout):
		return nil, ErrNewTxTimedOut
	}
}

// Unreachable fast way to check if a database connection is unreachable without
// creating a new transaction.
func Unreachable(db *DB) bool {
	return db.checker.unreachable()
}

type dbChecker struct {
	sync.RWMutex
	checking bool
}

func (d *dbChecker) unreachable() bool {
	d.RLock()
	defer d.RUnlock()
	return d.checking
}

func (d *dbChecker) start() bool {
	d.Lock()
	defer d.Unlock()

	if d.checking {
		return false
	}

	d.checking = true
	return true
}

func (d *dbChecker) stop() {
	d.Lock()
	defer d.Unlock()
	d.checking = false
}

func (d *dbChecker) check(db *sql.DB, duration time.Duration) {
	if !d.start() {
		// already checking
		return
	}

	for range time.Tick(2 * time.Second) {
		tx, err := newTx(db, duration)
		if err != nil {
			continue
		}

		// db is back!
		d.stop()
		tx.Commit()
		return
	}
}
<|endoftext|>


<commit_before>package service

import (
	"testing"
	"time"

	"github.com/jonboulle/clockwork"
	"golang.org/x/net/context"

	"github.com/keybase/client/go/libkb"
	keybase1 "github.com/keybase/client/go/protocol"
	"github.com/keybase/gregor/protocol/gregor1"
)

func rekeySetup(tc libkb.TestContext) (gregor1.UID, *gregorHandler, *RekeyUIHandler) {
	tc.G.SetService()

	kbUID := keybase1.MakeTestUID(1)
	gUID := gregor1.UID(kbUID.ToBytes())
	did, err := libkb.NewDeviceID()
	if err != nil {
		tc.T.Fatal(err)
	}
	tc.G.Env.GetConfigWriter().SetUserConfig(libkb.NewUserConfig(kbUID, "", nil, did), true)

	h, err := newGregorHandler(tc.G)
	if err != nil {
		tc.T.Fatal(err)
	}

	rekeyHandler := NewRekeyUIHandler(tc.G, 0)
	rekeyHandler.alwaysAlive = true
	rekeyHandler.notifyComplete = make(chan int, 10)
	rekeyHandler.scorer = fakeScoreProblemFolders
	h.PushHandler(rekeyHandler)

	return gUID, h, rekeyHandler
}

func rekeyBroadcast(tc libkb.TestContext, gUID gregor1.UID, h *gregorHandler, body string) {
	msgID := gregor1.MsgID("my_random_id")
	m := gregor1.Message{
		Ibm_: &gregor1.InBandMessage{
			StateUpdate_: &gregor1.StateUpdateMessage{
				Md_: gregor1.Metadata{
					MsgID_: msgID,
					Ctime_: gregor1.ToTime(tc.G.Clock.Now()),
					Uid_:   gUID,
				},
				Creation_: &gregor1.Item{
					Category_: gregor1.Category("kbfs_tlf_rekey_needed"),
					Body_:     gregor1.Body(body),
				},
			},
		},
	}

	if err := h.BroadcastMessage(context.Background(), m); err != nil {
		tc.T.Fatal(err)
	}

}

func TestRekeyNeededMessageNoScores(t *testing.T) {
	tc := libkb.SetupTest(t, "gregor", 1)
	defer tc.Cleanup()

	gUID, h, _ := rekeySetup(tc)

	// Test that a broadcast with an empty body works.
	rekeyBroadcast(tc, gUID, h, `{}`)
}

const problemSet = `{
	"user": {
		"uid": "295a7eea607af32040647123732bc819",
		"username": "t_alice"
	},
	"kid": "011212121212121212121212121212121212121212121212121212121212121212120a",
	"tlfs": [
		{
			"tlf": {
				"tlfid": "folder",
				"name": "folder name",
				"writers": ["t_alice"],
				"readers": ["t_alice"],
				"isPrivate": true
			},
			"score": 300,
			"solutions": ["011313131313131313131313131313131313131313131313131313131313131313130a"]
		}
	]
}`

func TestRekeyNeededMessageWithScores(t *testing.T) {
	tc := libkb.SetupTest(t, "gregor", 1)
	defer tc.Cleanup()

	rkeyui := &fakeRekeyUI{}
	router := fakeUIRouter{
		rekeyUI: rkeyui,
	}
	tc.G.SetUIRouter(&router)

	clock := clockwork.NewFakeClock()
	tc.G.Clock = clock

	gUID, h, rekeyHandler := rekeySetup(tc)

	go func() {
		clock.BlockUntil(1)
		clock.Advance(3 * time.Second)
	}()

	rekeyBroadcast(tc, gUID, h, problemSet)

	<-rekeyHandler.notifyComplete

	// Test that refresh was called twice.
	if len(rkeyui.refreshArgs) != 2 {
		t.Fatalf("rkeyui refresh calls: %d, expected 2", len(rkeyui.refreshArgs))
	}

	// the first call should contain a TLF
	if len(rkeyui.refreshArgs[0].Tlfs) != 1 {
		t.Errorf("first refresh call, tlf count = %d, expected 1", len(rkeyui.refreshArgs[0].Tlfs))
	}
	// the second call should have updated the scores, and have no more TLFs in it.
	if len(rkeyui.refreshArgs[1].Tlfs) != 0 {
		t.Errorf("second/final refresh call, tlf count = %d, expected 0", len(rkeyui.refreshArgs[1].Tlfs))
	}
}

type fakeRekeyUI struct {
	sessionID   int
	refreshArgs []keybase1.RefreshArg
}

// A rekey is needed, but the user closes the rekey status window.
func TestRekeyNeededUserClose(t *testing.T) {
	tc := libkb.SetupTest(t, "gregor", 1)
	defer tc.Cleanup()

	rkeyui := &fakeRekeyUI{}
	router := fakeUIRouter{
		rekeyUI: rkeyui,
	}
	tc.G.SetUIRouter(&router)

	clock := clockwork.NewFakeClock()
	tc.G.Clock = clock

	gUID, h, rekeyHandler := rekeySetup(tc)

	rekeyBroadcast(tc, gUID, h, problemSet)

	clock.BlockUntil(1)
	h.RekeyStatusFinish(context.Background(), rkeyui.sessionID)
	clock.Advance(3 * time.Second)

	select {
	case <-rekeyHandler.notifyComplete:
	case <-time.After(2 * time.Second):
		t.Fatal("timeout waiting for rekeyHandler.notifyComplete")
	}

	// there should be one call to refresh to bring the window up, but then the RekeyStatusFinish call above
	// should close the window and stop the loop.
	if len(rkeyui.refreshArgs) != 1 {
		t.Fatalf("rkeyui refresh calls: %d, expected 1", len(rkeyui.refreshArgs))
	}

	if len(rkeyui.refreshArgs[0].Tlfs) != 1 {
		t.Errorf("first refresh call, tlf count = %d, expected 1", len(rkeyui.refreshArgs[0].Tlfs))
	}
}

func (f *fakeRekeyUI) DelegateRekeyUI(ctx context.Context) (int, error) {
	f.sessionID++
	return f.sessionID, nil
}

func (f *fakeRekeyUI) Refresh(ctx context.Context, arg keybase1.RefreshArg) error {
	f.refreshArgs = append(f.refreshArgs, arg)
	return nil
}

func fakeScoreProblemFolders(g *libkb.GlobalContext, existing keybase1.ProblemSet) (keybase1.ProblemSet, error) {
	// always return empty ProblemSet
	return keybase1.ProblemSet{}, nil
}
<commit_msg>Increase test timeout to 20s<commit_after>package service

import (
	"testing"
	"time"

	"github.com/jonboulle/clockwork"
	"golang.org/x/net/context"

	"github.com/keybase/client/go/libkb"
	keybase1 "github.com/keybase/client/go/protocol"
	"github.com/keybase/gregor/protocol/gregor1"
)

func rekeySetup(tc libkb.TestContext) (gregor1.UID, *gregorHandler, *RekeyUIHandler) {
	tc.G.SetService()

	kbUID := keybase1.MakeTestUID(1)
	gUID := gregor1.UID(kbUID.ToBytes())
	did, err := libkb.NewDeviceID()
	if err != nil {
		tc.T.Fatal(err)
	}
	tc.G.Env.GetConfigWriter().SetUserConfig(libkb.NewUserConfig(kbUID, "", nil, did), true)

	h, err := newGregorHandler(tc.G)
	if err != nil {
		tc.T.Fatal(err)
	}

	rekeyHandler := NewRekeyUIHandler(tc.G, 0)
	rekeyHandler.alwaysAlive = true
	rekeyHandler.notifyComplete = make(chan int, 10)
	rekeyHandler.scorer = fakeScoreProblemFolders
	h.PushHandler(rekeyHandler)

	return gUID, h, rekeyHandler
}

func rekeyBroadcast(tc libkb.TestContext, gUID gregor1.UID, h *gregorHandler, body string) {
	msgID := gregor1.MsgID("my_random_id")
	m := gregor1.Message{
		Ibm_: &gregor1.InBandMessage{
			StateUpdate_: &gregor1.StateUpdateMessage{
				Md_: gregor1.Metadata{
					MsgID_: msgID,
					Ctime_: gregor1.ToTime(tc.G.Clock.Now()),
					Uid_:   gUID,
				},
				Creation_: &gregor1.Item{
					Category_: gregor1.Category("kbfs_tlf_rekey_needed"),
					Body_:     gregor1.Body(body),
				},
			},
		},
	}

	if err := h.BroadcastMessage(context.Background(), m); err != nil {
		tc.T.Fatal(err)
	}

}

func TestRekeyNeededMessageNoScores(t *testing.T) {
	tc := libkb.SetupTest(t, "gregor", 1)
	defer tc.Cleanup()

	gUID, h, _ := rekeySetup(tc)

	// Test that a broadcast with an empty body works.
	rekeyBroadcast(tc, gUID, h, `{}`)
}

const problemSet = `{
	"user": {
		"uid": "295a7eea607af32040647123732bc819",
		"username": "t_alice"
	},
	"kid": "011212121212121212121212121212121212121212121212121212121212121212120a",
	"tlfs": [
		{
			"tlf": {
				"tlfid": "folder",
				"name": "folder name",
				"writers": ["t_alice"],
				"readers": ["t_alice"],
				"isPrivate": true
			},
			"score": 300,
			"solutions": ["011313131313131313131313131313131313131313131313131313131313131313130a"]
		}
	]
}`

func TestRekeyNeededMessageWithScores(t *testing.T) {
	tc := libkb.SetupTest(t, "gregor", 1)
	defer tc.Cleanup()

	rkeyui := &fakeRekeyUI{}
	router := fakeUIRouter{
		rekeyUI: rkeyui,
	}
	tc.G.SetUIRouter(&router)

	clock := clockwork.NewFakeClock()
	tc.G.Clock = clock

	gUID, h, rekeyHandler := rekeySetup(tc)

	go func() {
		clock.BlockUntil(1)
		clock.Advance(3 * time.Second)
	}()

	rekeyBroadcast(tc, gUID, h, problemSet)

	<-rekeyHandler.notifyComplete

	// Test that refresh was called twice.
	if len(rkeyui.refreshArgs) != 2 {
		t.Fatalf("rkeyui refresh calls: %d, expected 2", len(rkeyui.refreshArgs))
	}

	// the first call should contain a TLF
	if len(rkeyui.refreshArgs[0].Tlfs) != 1 {
		t.Errorf("first refresh call, tlf count = %d, expected 1", len(rkeyui.refreshArgs[0].Tlfs))
	}
	// the second call should have updated the scores, and have no more TLFs in it.
	if len(rkeyui.refreshArgs[1].Tlfs) != 0 {
		t.Errorf("second/final refresh call, tlf count = %d, expected 0", len(rkeyui.refreshArgs[1].Tlfs))
	}
}

type fakeRekeyUI struct {
	sessionID   int
	refreshArgs []keybase1.RefreshArg
}

// A rekey is needed, but the user closes the rekey status window.
func TestRekeyNeededUserClose(t *testing.T) {
	tc := libkb.SetupTest(t, "gregor", 1)
	defer tc.Cleanup()

	rkeyui := &fakeRekeyUI{}
	router := fakeUIRouter{
		rekeyUI: rkeyui,
	}
	tc.G.SetUIRouter(&router)

	clock := clockwork.NewFakeClock()
	tc.G.Clock = clock

	gUID, h, rekeyHandler := rekeySetup(tc)

	rekeyBroadcast(tc, gUID, h, problemSet)

	clock.BlockUntil(1)
	h.RekeyStatusFinish(context.Background(), rkeyui.sessionID)
	clock.Advance(3 * time.Second)

	select {
	case <-rekeyHandler.notifyComplete:
	case <-time.After(20 * time.Second):
		t.Fatal("timeout waiting for rekeyHandler.notifyComplete")
	}

	// there should be one call to refresh to bring the window up, but then the RekeyStatusFinish call above
	// should close the window and stop the loop.
	if len(rkeyui.refreshArgs) != 1 {
		t.Fatalf("rkeyui refresh calls: %d, expected 1", len(rkeyui.refreshArgs))
	}

	if len(rkeyui.refreshArgs[0].Tlfs) != 1 {
		t.Errorf("first refresh call, tlf count = %d, expected 1", len(rkeyui.refreshArgs[0].Tlfs))
	}
}

func (f *fakeRekeyUI) DelegateRekeyUI(ctx context.Context) (int, error) {
	f.sessionID++
	return f.sessionID, nil
}

func (f *fakeRekeyUI) Refresh(ctx context.Context, arg keybase1.RefreshArg) error {
	f.refreshArgs = append(f.refreshArgs, arg)
	return nil
}

func fakeScoreProblemFolders(g *libkb.GlobalContext, existing keybase1.ProblemSet) (keybase1.ProblemSet, error) {
	// always return empty ProblemSet
	return keybase1.ProblemSet{}, nil
}
<|endoftext|>


<commit_before>// Copyright 2015, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package worker

import "github.com/youtube/vitess/go/vt/throttler"

const (
	defaultOnline            = true
	defaultOffline           = true
	defaultSourceReaderCount = 10
	// defaultWriteQueryMaxRows is set to 640 rows based on the assumption that
	// one row will take 64 bytes on average and our max size is 4 kB * 10.
	defaultWriteQueryMaxRows = 640
	// defaultWriteQueryMaxSize caps the write statements which aggregate multiple
	// rows. This limit prevents e.g. that MySQL will OOM.
	defaultWriteQueryMaxSize = 4096 * 10
	// defaultDestinationPackCount is deprecated in favor of the writeQueryMax*
	// values and currently only used by VerticalSplitClone.
	// defaultDestinationPackCount is the number of StreamExecute responses which
	// will be aggreated into one transaction. See the vttablet flag
	// "-queryserver-config-stream-buffer-size" for the size (in bytes) of a
	// StreamExecute response. As of 06/2015, the default for it was 32 kB.
	// Note that higher values for this flag --destination_pack_count will
	// increase memory consumption in vtworker, vttablet and mysql.
	defaultDestinationPackCount    = 10
	defaultMinTableSizeForSplit    = 1024 * 1024
	defaultDestinationWriterCount  = 20
	defaultMinHealthyRdonlyTablets = 2
	defaultMaxTPS                  = throttler.MaxRateModuleDisabled
)
<commit_msg>worker: Set more reasonable defaults for --write-query-max-size and --write-query-max-rows.<commit_after>// Copyright 2015, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package worker

import "github.com/youtube/vitess/go/vt/throttler"

const (
	defaultOnline            = true
	defaultOffline           = true
	defaultSourceReaderCount = 10
	// defaultWriteQueryMaxRows aggregates up to 100 rows per INSERT or DELETE
	// query. Higher values are not recommended to avoid overloading MySQL.
	// The actual row count will be less if defaultWriteQueryMaxSize is reached
	// first, but always at least 1 row.
	defaultWriteQueryMaxRows = 100
	// defaultWriteQueryMaxSize caps the write queries which aggregate multiple
	// rows. This limit prevents e.g. that MySQL will OOM.
	defaultWriteQueryMaxSize = 1024 * 1024
	// defaultDestinationPackCount is deprecated in favor of the writeQueryMax*
	// values and currently only used by VerticalSplitClone.
	// defaultDestinationPackCount is the number of StreamExecute responses which
	// will be aggreated into one transaction. See the vttablet flag
	// "-queryserver-config-stream-buffer-size" for the size (in bytes) of a
	// StreamExecute response. As of 06/2015, the default for it was 32 kB.
	// Note that higher values for this flag --destination_pack_count will
	// increase memory consumption in vtworker, vttablet and mysql.
	defaultDestinationPackCount    = 10
	defaultMinTableSizeForSplit    = 1024 * 1024
	defaultDestinationWriterCount  = 20
	defaultMinHealthyRdonlyTablets = 2
	defaultMaxTPS                  = throttler.MaxRateModuleDisabled
)
<|endoftext|>


<commit_before>/*
 *
 * Copyright 2020 Google LLC
 *
 * 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
 *
 *     https://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 main

import (
	"context"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"time"

	"go.opencensus.io/plugin/ocgrpc"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials/insecure"
	xdscreds "google.golang.org/grpc/credentials/xds"
	walletpb "google.golang.org/grpc/grpc-wallet/grpc/examples/wallet"
	statspb "google.golang.org/grpc/grpc-wallet/grpc/examples/wallet/stats"
	"google.golang.org/grpc/grpc-wallet/observability"
	"google.golang.org/grpc/grpc-wallet/utility"
	"google.golang.org/grpc/metadata"

	_ "google.golang.org/grpc/xds" // To enable xds support.
)

var users = map[string]map[string]string{
	"Alice": {
		"authorization": "2bd806c9",
		"membership":    "premium",
	},
	"Bob": {
		"authorization": "81b637d8",
		"membership":    "normal",
	},
}

// TODO(easwars): Get rid of this struct and define flags as package globals.
type arguments struct {
	subcommand       string
	walletServer     string
	statsServer      string
	user             string
	watch            bool
	unaryWatch       bool
	gcpClientProject string
	route            string
	creds            string
	affinity         bool
}

var args arguments

// parseArguments parses the command line arguments using the flag package.
func parseArguments() {
	flags := flag.NewFlagSet("wallet_client", flag.ExitOnError)
	flags.StringVar(&args.walletServer, "wallet_server", "localhost:18881", "address of the wallet service, default 'localhost:18881'")
	flags.StringVar(&args.statsServer, "stats_server", "localhost:18882", "address of the stats service, default 'localhost:18882'")
	flags.StringVar(&args.user, "user", "Alice", "the name of the user account, default 'Alice'")
	flags.BoolVar(&args.watch, "watch", false, "if the balance/price should be watched rather than queried once, default false")
	flags.BoolVar(&args.unaryWatch, "unary_watch", false, "if the balance/price should be watched but with repeated unary RPCs rather than a streaming rpc, default false")
	flags.StringVar(&args.gcpClientProject, "gcp_client_project", "", "if set, metrics and traces will be sent to Cloud Monitoring and Cloud Trace")
	flags.StringVar(&args.route, "route", "", "a string value to set for the 'route' header, unset by default")
	flags.StringVar(&args.creds, "creds", "insecure", "type of transport credentials to use. Supported values include 'xds' and 'insecure', defaults to 'insecure'")
	flags.BoolVar(&args.affinity, "affinity", false, "if the requests will be sent with session affinity. defaults false")
	flags.Usage = func() {
		fmt.Fprintf(flag.CommandLine.Output(), "Usage of %s:\n", os.Args[0])
		fmt.Fprintf(flag.CommandLine.Output(), `
subcommands:
  balance
	print balance of the wallet
  price
	print price of grpc-coin

flags
`)
		flags.PrintDefaults()
	}

	if len(os.Args) < 2 {
		flags.Usage()
		log.Fatalln("no subcommand.")
	}
	args.subcommand = os.Args[1]
	if args.subcommand != "balance" && args.subcommand != "price" {
		flags.Usage()
		log.Fatalf("unrecognized subcommand '%s'.", args.subcommand)
	}
	flags.Parse(os.Args[2:])
	if args.user != "Alice" && args.user != "Bob" {
		flags.Usage()
		log.Fatalf("unrecognized user '%s'.", args.user)
	}
	if args.watch && args.unaryWatch {
		flags.Usage()
		log.Fatalln("unary_watch incompatible with watch.")
	}
	if args.subcommand == "price" && args.unaryWatch {
		flags.Usage()
		log.Fatalln("unary_watch incompatible with price subcommand.")
	}
}

const (
	affinityMDKey   = "session_id"
	affinityMDValue = "1234"
)

// createMetaData creates the metadata for an outgoing request based on the
// flags.
func createMetaData() (metadata.MD, error) {
	md, ok := users[args.user]
	if !ok {
		return nil, fmt.Errorf("unrecognized user: %v", args.user)
	}
	if args.route != "" {
		md["route"] = args.route
	}
	if args.affinity {
		md[affinityMDKey] = affinityMDValue
	}
	return metadata.New(md), nil
}

// handleBalanceResponse prints the data in a BalanceResponse.
func handleBalanceResponse(r *walletpb.BalanceResponse) {
	log.Printf("user: %s, total grpc-coin balance: %d.", args.user, r.GetBalance())
	for _, addrBalance := range r.GetAddresses() {
		log.Printf(" - address: %s, balance: %d.", addrBalance.Address, addrBalance.Balance)
	}
}

// balanceSubcommand handles the 'balance' subcommand.
func balanceSubcommand(opts []grpc.DialOption) {
	dialCtx, dialCancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer dialCancel()
	conn, err := grpc.DialContext(dialCtx, args.walletServer, opts...)
	if err != nil {
		log.Fatalf("Failed to connect to the Wallet server: %v", err)
	}
	defer conn.Close()

	c := walletpb.NewWalletClient(conn)
	md, err := createMetaData()
	if err != nil {
		log.Fatalf("error creating metadata: %v.", err)
	}

	ctx := metadata.NewOutgoingContext(context.Background(), md)
	if !args.watch {
		for {
			var header metadata.MD
			r, err := c.FetchBalance(ctx, &walletpb.BalanceRequest{IncludeBalancePerAddress: true}, grpc.Header(&header))
			if err != nil {
				log.Printf("failed to fetch balance: %v", err)
			} else {
				utility.PrintHostname(header)
				handleBalanceResponse(r)
			}
			if !args.unaryWatch {
				break
			}
			time.Sleep(time.Second)
		}
		return
	}
	s, err := c.WatchBalance(ctx, &walletpb.BalanceRequest{IncludeBalancePerAddress: true}, grpc.WaitForReady(true))
	if err != nil {
		log.Fatalf("failed to create stream: %v.", err)
	}
	header, err := s.Header()
	if err != nil {
		log.Fatalf("could not extract header: %v", err)
	}
	utility.PrintHostname(header)
	for {
		r, err := s.Recv()
		if err != nil {
			if err == io.EOF {
				break
			}
			log.Fatalf("failed to fetch balance: %v", err)
		}
		handleBalanceResponse(r)
	}
}

// priceSubcommand handles the 'price' subcommand.
func priceSubcommand(opts []grpc.DialOption) {
	dialCtx, dialCancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer dialCancel()
	conn, err := grpc.DialContext(dialCtx, args.statsServer, opts...)
	if err != nil {
		log.Fatalf("Failed to connect to the Wallet server: %v", err)
	}
	defer conn.Close()

	c := statspb.NewStatsClient(conn)
	md, err := createMetaData()
	if err != nil {
		log.Fatalf("error creating metadata: %v.", err)
	}

	ctx := metadata.NewOutgoingContext(context.Background(), md)
	if !args.watch {
		var header metadata.MD
		r, err := c.FetchPrice(ctx, &statspb.PriceRequest{}, grpc.Header(&header))
		if err != nil {
			log.Printf("failed to fetch price: %v", err)
		}
		utility.PrintHostname(header)
		log.Printf("grpc-coin price: %d.", r.GetPrice())
		return
	}
	s, err := c.WatchPrice(ctx, &statspb.PriceRequest{})
	if err != nil {
		log.Fatalf("failed to create stream: %v.", err)
	}
	header, err := s.Header()
	if err != nil {
		log.Fatalf("could not extract header: %v", err)
	}
	utility.PrintHostname(header)
	for {
		r, err := s.Recv()
		if err != nil {
			if err == io.EOF {
				break
			}
			log.Fatalf("failed to fetch price: %v", err)
		}
		log.Printf("grpc-coin price: %d.", r.GetPrice())
	}
}

func main() {
	parseArguments()

	// Parse credentials type from the command line to determine if xDS
	// credentials are to be used.
	xdsCreds, err := utility.ParseCredentialsType(args.creds)
	if err != nil {
		log.Fatal(err)
	}
	log.Printf("xDS credentials in use: %v", xdsCreds)

	// Use insecure credentials by default. But if xDS credentials are specified
	// on the command line, use xDS credentials with an insecure fallback.
	creds := insecure.NewCredentials()
	if xdsCreds {
		var err error
		creds, err = xdscreds.NewClientCredentials(xdscreds.ClientOptions{FallbackCreds: insecure.NewCredentials()})
		if err != nil {
			log.Fatalf("Failed to create client xDS credentials: %v", err)
		}
		log.Print("xDS credentials created")
	}
	opts := []grpc.DialOption{grpc.WithTransportCredentials(creds), grpc.WithBlock()}

	if args.gcpClientProject != "" {
		opts = append(opts, grpc.WithStatsHandler(new(ocgrpc.ClientHandler)))
		sd := observability.ConfigureStackdriver(args.gcpClientProject)
		defer sd.Flush()
		defer sd.StopMetricsExporter()
	}

	if args.subcommand == "balance" {
		balanceSubcommand(opts)
		return
	}
	priceSubcommand(opts)
}
<commit_msg>go/wallet_client: set affinity md value based on user (#85)<commit_after>/*
 *
 * Copyright 2020 Google LLC
 *
 * 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
 *
 *     https://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 main

import (
	"context"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"time"

	"go.opencensus.io/plugin/ocgrpc"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials/insecure"
	xdscreds "google.golang.org/grpc/credentials/xds"
	walletpb "google.golang.org/grpc/grpc-wallet/grpc/examples/wallet"
	statspb "google.golang.org/grpc/grpc-wallet/grpc/examples/wallet/stats"
	"google.golang.org/grpc/grpc-wallet/observability"
	"google.golang.org/grpc/grpc-wallet/utility"
	"google.golang.org/grpc/metadata"

	_ "google.golang.org/grpc/xds" // To enable xds support.
)

var users = map[string]map[string]string{
	"Alice": {
		"authorization": "2bd806c9",
		"membership":    "premium",
	},
	"Bob": {
		"authorization": "81b637d8",
		"membership":    "normal",
	},
}

// TODO(easwars): Get rid of this struct and define flags as package globals.
type arguments struct {
	subcommand       string
	walletServer     string
	statsServer      string
	user             string
	watch            bool
	unaryWatch       bool
	gcpClientProject string
	route            string
	creds            string
	affinity         bool
}

var args arguments

// parseArguments parses the command line arguments using the flag package.
func parseArguments() {
	flags := flag.NewFlagSet("wallet_client", flag.ExitOnError)
	flags.StringVar(&args.walletServer, "wallet_server", "localhost:18881", "address of the wallet service, default 'localhost:18881'")
	flags.StringVar(&args.statsServer, "stats_server", "localhost:18882", "address of the stats service, default 'localhost:18882'")
	flags.StringVar(&args.user, "user", "Alice", "the name of the user account, default 'Alice'")
	flags.BoolVar(&args.watch, "watch", false, "if the balance/price should be watched rather than queried once, default false")
	flags.BoolVar(&args.unaryWatch, "unary_watch", false, "if the balance/price should be watched but with repeated unary RPCs rather than a streaming rpc, default false")
	flags.StringVar(&args.gcpClientProject, "gcp_client_project", "", "if set, metrics and traces will be sent to Cloud Monitoring and Cloud Trace")
	flags.StringVar(&args.route, "route", "", "a string value to set for the 'route' header, unset by default")
	flags.StringVar(&args.creds, "creds", "insecure", "type of transport credentials to use. Supported values include 'xds' and 'insecure', defaults to 'insecure'")
	flags.BoolVar(&args.affinity, "affinity", false, "if the requests will be sent with session affinity. defaults false")
	flags.Usage = func() {
		fmt.Fprintf(flag.CommandLine.Output(), "Usage of %s:\n", os.Args[0])
		fmt.Fprintf(flag.CommandLine.Output(), `
subcommands:
  balance
	print balance of the wallet
  price
	print price of grpc-coin

flags
`)
		flags.PrintDefaults()
	}

	if len(os.Args) < 2 {
		flags.Usage()
		log.Fatalln("no subcommand.")
	}
	args.subcommand = os.Args[1]
	if args.subcommand != "balance" && args.subcommand != "price" {
		flags.Usage()
		log.Fatalf("unrecognized subcommand '%s'.", args.subcommand)
	}
	flags.Parse(os.Args[2:])
	if args.user != "Alice" && args.user != "Bob" {
		flags.Usage()
		log.Fatalf("unrecognized user '%s'.", args.user)
	}
	if args.watch && args.unaryWatch {
		flags.Usage()
		log.Fatalln("unary_watch incompatible with watch.")
	}
	if args.subcommand == "price" && args.unaryWatch {
		flags.Usage()
		log.Fatalln("unary_watch incompatible with price subcommand.")
	}
}

const affinityMDKey = "session_id"

var affinityMDValues = map[string]string{
	"Alice": "11111111",
	"Bob":   "22222222",
}

// createMetaData creates the metadata for an outgoing request based on the
// flags.
func createMetaData() (metadata.MD, error) {
	md, ok := users[args.user]
	if !ok {
		return nil, fmt.Errorf("unrecognized user: %v", args.user)
	}
	if args.route != "" {
		md["route"] = args.route
	}
	if args.affinity {
		md[affinityMDKey] = affinityMDValues[args.user]
	}
	return metadata.New(md), nil
}

// handleBalanceResponse prints the data in a BalanceResponse.
func handleBalanceResponse(r *walletpb.BalanceResponse) {
	log.Printf("user: %s, total grpc-coin balance: %d.", args.user, r.GetBalance())
	for _, addrBalance := range r.GetAddresses() {
		log.Printf(" - address: %s, balance: %d.", addrBalance.Address, addrBalance.Balance)
	}
}

// balanceSubcommand handles the 'balance' subcommand.
func balanceSubcommand(opts []grpc.DialOption) {
	dialCtx, dialCancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer dialCancel()
	conn, err := grpc.DialContext(dialCtx, args.walletServer, opts...)
	if err != nil {
		log.Fatalf("Failed to connect to the Wallet server: %v", err)
	}
	defer conn.Close()

	c := walletpb.NewWalletClient(conn)
	md, err := createMetaData()
	if err != nil {
		log.Fatalf("error creating metadata: %v.", err)
	}

	ctx := metadata.NewOutgoingContext(context.Background(), md)
	if !args.watch {
		for {
			var header metadata.MD
			r, err := c.FetchBalance(ctx, &walletpb.BalanceRequest{IncludeBalancePerAddress: true}, grpc.Header(&header))
			if err != nil {
				log.Printf("failed to fetch balance: %v", err)
			} else {
				utility.PrintHostname(header)
				handleBalanceResponse(r)
			}
			if !args.unaryWatch {
				break
			}
			time.Sleep(time.Second)
		}
		return
	}
	s, err := c.WatchBalance(ctx, &walletpb.BalanceRequest{IncludeBalancePerAddress: true}, grpc.WaitForReady(true))
	if err != nil {
		log.Fatalf("failed to create stream: %v.", err)
	}
	header, err := s.Header()
	if err != nil {
		log.Fatalf("could not extract header: %v", err)
	}
	utility.PrintHostname(header)
	for {
		r, err := s.Recv()
		if err != nil {
			if err == io.EOF {
				break
			}
			log.Fatalf("failed to fetch balance: %v", err)
		}
		handleBalanceResponse(r)
	}
}

// priceSubcommand handles the 'price' subcommand.
func priceSubcommand(opts []grpc.DialOption) {
	dialCtx, dialCancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer dialCancel()
	conn, err := grpc.DialContext(dialCtx, args.statsServer, opts...)
	if err != nil {
		log.Fatalf("Failed to connect to the Wallet server: %v", err)
	}
	defer conn.Close()

	c := statspb.NewStatsClient(conn)
	md, err := createMetaData()
	if err != nil {
		log.Fatalf("error creating metadata: %v.", err)
	}

	ctx := metadata.NewOutgoingContext(context.Background(), md)
	if !args.watch {
		var header metadata.MD
		r, err := c.FetchPrice(ctx, &statspb.PriceRequest{}, grpc.Header(&header))
		if err != nil {
			log.Printf("failed to fetch price: %v", err)
		}
		utility.PrintHostname(header)
		log.Printf("grpc-coin price: %d.", r.GetPrice())
		return
	}
	s, err := c.WatchPrice(ctx, &statspb.PriceRequest{})
	if err != nil {
		log.Fatalf("failed to create stream: %v.", err)
	}
	header, err := s.Header()
	if err != nil {
		log.Fatalf("could not extract header: %v", err)
	}
	utility.PrintHostname(header)
	for {
		r, err := s.Recv()
		if err != nil {
			if err == io.EOF {
				break
			}
			log.Fatalf("failed to fetch price: %v", err)
		}
		log.Printf("grpc-coin price: %d.", r.GetPrice())
	}
}

func main() {
	parseArguments()

	// Parse credentials type from the command line to determine if xDS
	// credentials are to be used.
	xdsCreds, err := utility.ParseCredentialsType(args.creds)
	if err != nil {
		log.Fatal(err)
	}
	log.Printf("xDS credentials in use: %v", xdsCreds)

	// Use insecure credentials by default. But if xDS credentials are specified
	// on the command line, use xDS credentials with an insecure fallback.
	creds := insecure.NewCredentials()
	if xdsCreds {
		var err error
		creds, err = xdscreds.NewClientCredentials(xdscreds.ClientOptions{FallbackCreds: insecure.NewCredentials()})
		if err != nil {
			log.Fatalf("Failed to create client xDS credentials: %v", err)
		}
		log.Print("xDS credentials created")
	}
	opts := []grpc.DialOption{grpc.WithTransportCredentials(creds), grpc.WithBlock()}

	if args.gcpClientProject != "" {
		opts = append(opts, grpc.WithStatsHandler(new(ocgrpc.ClientHandler)))
		sd := observability.ConfigureStackdriver(args.gcpClientProject)
		defer sd.Flush()
		defer sd.StopMetricsExporter()
	}

	if args.subcommand == "balance" {
		balanceSubcommand(opts)
		return
	}
	priceSubcommand(opts)
}
<|endoftext|>


<commit_before>package gobusters3

import (
	"bufio"
	"bytes"
	"context"
	"encoding/xml"
	"fmt"
	"net/http"
	"regexp"
	"strings"
	"text/tabwriter"

	"github.com/OJ/gobuster/v3/libgobuster"
)

// GobusterS3 is the main type to implement the interface
type GobusterS3 struct {
	options     *OptionsS3
	globalopts  *libgobuster.Options
	http        *libgobuster.HTTPClient
	bucketRegex *regexp.Regexp
}

// NewGobusterS3 creates a new initialized GobusterS3
func NewGobusterS3(cont context.Context, globalopts *libgobuster.Options, opts *OptionsS3) (*GobusterS3, error) {
	if globalopts == nil {
		return nil, fmt.Errorf("please provide valid global options")
	}

	if opts == nil {
		return nil, fmt.Errorf("please provide valid plugin options")
	}

	g := GobusterS3{
		options:    opts,
		globalopts: globalopts,
	}

	basicOptions := libgobuster.BasicHTTPOptions{
		Proxy:     opts.Proxy,
		Timeout:   opts.Timeout,
		UserAgent: opts.UserAgent,
	}

	httpOpts := libgobuster.HTTPOptions{
		BasicHTTPOptions: basicOptions,
		// needed so we can list bucket contents
		FollowRedirect: true,
	}

	h, err := libgobuster.NewHTTPClient(cont, &httpOpts)
	if err != nil {
		return nil, err
	}
	g.http = h
	g.bucketRegex = regexp.MustCompile(`^[a-z0-9\-.]{3,63}$`)

	return &g, nil
}

// Name should return the name of the plugin
func (s *GobusterS3) Name() string {
	return "S3 bucket enumeration"
}

// RequestsPerRun returns the number of requests this plugin makes per single wordlist item
func (s *GobusterS3) RequestsPerRun() int {
	return 1
}

// PreRun is the pre run implementation of GobusterS3
func (s *GobusterS3) PreRun() error {
	return nil
}

// Run is the process implementation of GobusterS3
func (s *GobusterS3) Run(word string, resChannel chan<- libgobuster.Result) error {
	// only check for valid bucket names
	if !s.isValidBucketName(word) {
		return nil
	}

	bucketURL := fmt.Sprintf("http://%s.s3.amazonaws.com/?max-keys=%d", word, s.options.MaxFilesToList)
	status, _, _, body, err := s.http.Request(bucketURL, libgobuster.RequestOptions{ReturnBody: true})
	if err != nil {
		return err
	}

	if status == nil || body == nil {
		return nil
	}

	// looks like 404 and 400 are the only negative status codes
	found := false
	switch *status {
	case http.StatusBadRequest:
	case http.StatusNotFound:
		found = false
	case http.StatusOK:
		// listing enabled
		found = true
		// parse xml
	default:
		// default to found as we use negative status codes
		found = true
	}

	// nothing found, bail out
	// may add the result later if we want to enable verbose output
	if !found {
		return nil
	}

	extraStr := ""
	if s.globalopts.Verbose {
		// get status
		if bytes.Contains(body, []byte("<Error>")) {
			awsError := AWSError{}
			err := xml.Unmarshal(body, &awsError)
			if err != nil {
				return fmt.Errorf("could not parse error xml: %w", err)
			}
			// https://docs.aws.amazon.com/AmazonS3/latest/API/ErrorResponses.html#ErrorCodeList
			extraStr = fmt.Sprintf("Error: %s (%s)", awsError.Message, awsError.Code)
		} else if bytes.Contains(body, []byte("<ListBucketResult ")) {
			// bucket listing enabled
			awsListing := AWSListing{}
			err := xml.Unmarshal(body, &awsListing)
			if err != nil {
				return fmt.Errorf("could not parse result xml: %w", err)
			}
			extraStr = "Bucket Listing enabled: "
			for _, x := range awsListing.Contents {
				extraStr += fmt.Sprintf("%s (%db), ", x.Key, x.Size)
			}
			extraStr = strings.TrimRight(extraStr, ", ")
		}
	}

	resChannel <- Result{
		Found:      found,
		BucketName: word,
		Status:     extraStr,
	}

	return nil
}

// GetConfigString returns the string representation of the current config
func (s *GobusterS3) GetConfigString() (string, error) {
	var buffer bytes.Buffer
	bw := bufio.NewWriter(&buffer)
	tw := tabwriter.NewWriter(bw, 0, 5, 3, ' ', 0)
	o := s.options

	if _, err := fmt.Fprintf(tw, "[+] Threads:\t%d\n", s.globalopts.Threads); err != nil {
		return "", err
	}

	if s.globalopts.Delay > 0 {
		if _, err := fmt.Fprintf(tw, "[+] Delay:\t%s\n", s.globalopts.Delay); err != nil {
			return "", err
		}
	}

	wordlist := "stdin (pipe)"
	if s.globalopts.Wordlist != "-" {
		wordlist = s.globalopts.Wordlist
	}
	if _, err := fmt.Fprintf(tw, "[+] Wordlist:\t%s\n", wordlist); err != nil {
		return "", err
	}

	if s.globalopts.PatternFile != "" {
		if _, err := fmt.Fprintf(tw, "[+] Patterns:\t%s (%d entries)\n", s.globalopts.PatternFile, len(s.globalopts.Patterns)); err != nil {
			return "", err
		}
	}

	if o.Proxy != "" {
		if _, err := fmt.Fprintf(tw, "[+] Proxy:\t%s\n", o.Proxy); err != nil {
			return "", err
		}
	}

	if o.UserAgent != "" {
		if _, err := fmt.Fprintf(tw, "[+] User Agent:\t%s\n", o.UserAgent); err != nil {
			return "", err
		}
	}

	if _, err := fmt.Fprintf(tw, "[+] Timeout:\t%s\n", o.Timeout.String()); err != nil {
		return "", err
	}

	if s.globalopts.Verbose {
		if _, err := fmt.Fprintf(tw, "[+] Verbose:\ttrue\n"); err != nil {
			return "", err
		}
	}

	if _, err := fmt.Fprintf(tw, "[+] Maximum files to list:\t%d\n", o.MaxFilesToList); err != nil {
		return "", err
	}

	if err := tw.Flush(); err != nil {
		return "", fmt.Errorf("error on tostring: %w", err)
	}

	if err := bw.Flush(); err != nil {
		return "", fmt.Errorf("error on tostring: %w", err)
	}

	return strings.TrimSpace(buffer.String()), nil
}

// https://docs.aws.amazon.com/AmazonS3/latest/dev/BucketRestrictions.html
func (s *GobusterS3) isValidBucketName(bucketName string) bool {
	if !s.bucketRegex.MatchString(bucketName) {
		return false
	}
	if strings.HasSuffix(bucketName, "-") ||
		strings.HasPrefix(bucketName, ".") ||
		strings.HasPrefix(bucketName, "-") ||
		strings.Contains(bucketName, "..") ||
		strings.Contains(bucketName, ".-") ||
		strings.Contains(bucketName, "-.") {
		return false
	}
	return true
}
<commit_msg>Change the S3 bucket protocol<commit_after>package gobusters3

import (
	"bufio"
	"bytes"
	"context"
	"encoding/xml"
	"fmt"
	"net/http"
	"regexp"
	"strings"
	"text/tabwriter"

	"github.com/OJ/gobuster/v3/libgobuster"
)

// GobusterS3 is the main type to implement the interface
type GobusterS3 struct {
	options     *OptionsS3
	globalopts  *libgobuster.Options
	http        *libgobuster.HTTPClient
	bucketRegex *regexp.Regexp
}

// NewGobusterS3 creates a new initialized GobusterS3
func NewGobusterS3(cont context.Context, globalopts *libgobuster.Options, opts *OptionsS3) (*GobusterS3, error) {
	if globalopts == nil {
		return nil, fmt.Errorf("please provide valid global options")
	}

	if opts == nil {
		return nil, fmt.Errorf("please provide valid plugin options")
	}

	g := GobusterS3{
		options:    opts,
		globalopts: globalopts,
	}

	basicOptions := libgobuster.BasicHTTPOptions{
		Proxy:     opts.Proxy,
		Timeout:   opts.Timeout,
		UserAgent: opts.UserAgent,
	}

	httpOpts := libgobuster.HTTPOptions{
		BasicHTTPOptions: basicOptions,
		// needed so we can list bucket contents
		FollowRedirect: true,
	}

	h, err := libgobuster.NewHTTPClient(cont, &httpOpts)
	if err != nil {
		return nil, err
	}
	g.http = h
	g.bucketRegex = regexp.MustCompile(`^[a-z0-9\-.]{3,63}$`)

	return &g, nil
}

// Name should return the name of the plugin
func (s *GobusterS3) Name() string {
	return "S3 bucket enumeration"
}

// RequestsPerRun returns the number of requests this plugin makes per single wordlist item
func (s *GobusterS3) RequestsPerRun() int {
	return 1
}

// PreRun is the pre run implementation of GobusterS3
func (s *GobusterS3) PreRun() error {
	return nil
}

// Run is the process implementation of GobusterS3
func (s *GobusterS3) Run(word string, resChannel chan<- libgobuster.Result) error {
	// only check for valid bucket names
	if !s.isValidBucketName(word) {
		return nil
	}

	bucketURL := fmt.Sprintf("https://%s.s3.amazonaws.com/?max-keys=%d", word, s.options.MaxFilesToList)
	status, _, _, body, err := s.http.Request(bucketURL, libgobuster.RequestOptions{ReturnBody: true})
	if err != nil {
		return err
	}

	if status == nil || body == nil {
		return nil
	}

	// looks like 404 and 400 are the only negative status codes
	found := false
	switch *status {
	case http.StatusBadRequest:
	case http.StatusNotFound:
		found = false
	case http.StatusOK:
		// listing enabled
		found = true
		// parse xml
	default:
		// default to found as we use negative status codes
		found = true
	}

	// nothing found, bail out
	// may add the result later if we want to enable verbose output
	if !found {
		return nil
	}

	extraStr := ""
	if s.globalopts.Verbose {
		// get status
		if bytes.Contains(body, []byte("<Error>")) {
			awsError := AWSError{}
			err := xml.Unmarshal(body, &awsError)
			if err != nil {
				return fmt.Errorf("could not parse error xml: %w", err)
			}
			// https://docs.aws.amazon.com/AmazonS3/latest/API/ErrorResponses.html#ErrorCodeList
			extraStr = fmt.Sprintf("Error: %s (%s)", awsError.Message, awsError.Code)
		} else if bytes.Contains(body, []byte("<ListBucketResult ")) {
			// bucket listing enabled
			awsListing := AWSListing{}
			err := xml.Unmarshal(body, &awsListing)
			if err != nil {
				return fmt.Errorf("could not parse result xml: %w", err)
			}
			extraStr = "Bucket Listing enabled: "
			for _, x := range awsListing.Contents {
				extraStr += fmt.Sprintf("%s (%db), ", x.Key, x.Size)
			}
			extraStr = strings.TrimRight(extraStr, ", ")
		}
	}

	resChannel <- Result{
		Found:      found,
		BucketName: word,
		Status:     extraStr,
	}

	return nil
}

// GetConfigString returns the string representation of the current config
func (s *GobusterS3) GetConfigString() (string, error) {
	var buffer bytes.Buffer
	bw := bufio.NewWriter(&buffer)
	tw := tabwriter.NewWriter(bw, 0, 5, 3, ' ', 0)
	o := s.options

	if _, err := fmt.Fprintf(tw, "[+] Threads:\t%d\n", s.globalopts.Threads); err != nil {
		return "", err
	}

	if s.globalopts.Delay > 0 {
		if _, err := fmt.Fprintf(tw, "[+] Delay:\t%s\n", s.globalopts.Delay); err != nil {
			return "", err
		}
	}

	wordlist := "stdin (pipe)"
	if s.globalopts.Wordlist != "-" {
		wordlist = s.globalopts.Wordlist
	}
	if _, err := fmt.Fprintf(tw, "[+] Wordlist:\t%s\n", wordlist); err != nil {
		return "", err
	}

	if s.globalopts.PatternFile != "" {
		if _, err := fmt.Fprintf(tw, "[+] Patterns:\t%s (%d entries)\n", s.globalopts.PatternFile, len(s.globalopts.Patterns)); err != nil {
			return "", err
		}
	}

	if o.Proxy != "" {
		if _, err := fmt.Fprintf(tw, "[+] Proxy:\t%s\n", o.Proxy); err != nil {
			return "", err
		}
	}

	if o.UserAgent != "" {
		if _, err := fmt.Fprintf(tw, "[+] User Agent:\t%s\n", o.UserAgent); err != nil {
			return "", err
		}
	}

	if _, err := fmt.Fprintf(tw, "[+] Timeout:\t%s\n", o.Timeout.String()); err != nil {
		return "", err
	}

	if s.globalopts.Verbose {
		if _, err := fmt.Fprintf(tw, "[+] Verbose:\ttrue\n"); err != nil {
			return "", err
		}
	}

	if _, err := fmt.Fprintf(tw, "[+] Maximum files to list:\t%d\n", o.MaxFilesToList); err != nil {
		return "", err
	}

	if err := tw.Flush(); err != nil {
		return "", fmt.Errorf("error on tostring: %w", err)
	}

	if err := bw.Flush(); err != nil {
		return "", fmt.Errorf("error on tostring: %w", err)
	}

	return strings.TrimSpace(buffer.String()), nil
}

// https://docs.aws.amazon.com/AmazonS3/latest/dev/BucketRestrictions.html
func (s *GobusterS3) isValidBucketName(bucketName string) bool {
	if !s.bucketRegex.MatchString(bucketName) {
		return false
	}
	if strings.HasSuffix(bucketName, "-") ||
		strings.HasPrefix(bucketName, ".") ||
		strings.HasPrefix(bucketName, "-") ||
		strings.Contains(bucketName, "..") ||
		strings.Contains(bucketName, ".-") ||
		strings.Contains(bucketName, "-.") {
		return false
	}
	return true
}
<|endoftext|>


<commit_before>package lib

import (
	. "github.com/r7kamura/gospel"
	"path"
	"runtime"
	"testing"
)

func TestInterpretor(t *testing.T) {
	Describe(t, "インタープリターのテスト", func() {
		data := []byte{'A', ' ', 'B', '\t', '\r', '\n'}
		Context("インスタンス生成", func() {
			It("インスタンスが生成できること", func() {
				actual := NewInterpreter(data)
				Expect(actual).To(Exist)
				Expect(actual.origin).To(Equal, data)
			})
		})
		Context("データを読み込んで命令を実行する", func() {
			It("不要なデータを排除してスペースやタブをそれぞれSとTに置き換えること", func() {
				data := []byte{' ', ' ', '	', ' ', ' ', ' ', ' ', ' ', '	', '\n'}
				expected := []byte{'S', 'S', 'T', 'S', 'S', 'S', 'S', 'S', 'T', '\n'}
				sut := NewInterpreter(data)
				dat, err := sut.ToChar()
				Expect(err).To(NotExist)
				Expect(dat).To(Equal, expected)
			})
			It("不要なデータを排除してスペースやタブをそれぞれSとTに置き換えること パート２", func() {
				data := []byte{' ', ' ', '	', ' ', ' ', ' ', ' ', '	', '	', '\n'}
				expected := []byte{'S', 'S', 'T', 'S', 'S', 'S', 'S', 'T', 'T', '\n'}
				sut := NewInterpreter(data)
				dat, err := sut.ToChar()
				Expect(err).To(NotExist)
				Expect(dat).To(Equal, expected)
			})
			It("不要なデータを排除して指定したコマンドが作成されること", func() {
				data := []byte{' ', ' ', '	', ' ', ' ', ' ', ' ', ' ', '	', '\n'}
				sut := NewInterpreter(data)
				err := sut.ToCode()
				Expect(err).To(NotExist)
				Expect(sut.commands.Len()).To(Equal, 1)
				Expect(sut.commands.Get(1)).To(Equal, NewSubCommandWithParam("stack", "push", 0x41))
			})
		})
		Context("不要な文字を排除する関数", func() {
			It("不要なデータ以外排除されていること", func() {
				sut := NewInterpreter(data)
				sut.filter()
				Expect(sut.source).To(Equal, []byte{' ', '\t', '\n'})
			})
			It("不要なデータ以外排除されていること（改行を増やす）", func() {
				data = append(data, '\n')
				sut := NewInterpreter(data)
				sut.filter()
				Expect(sut.source).To(Equal, []byte{' ', '\t', '\n', '\n'})
			})
		})
	})
}

func current_dir() string {
	_, fpath, _, _ := runtime.Caller(0)
	return path.Dir(fpath)
}
<commit_msg>指定されたデータをもとにコマンドが作成されたことを確認（テストパターンを追加）<commit_after>package lib

import (
	. "github.com/r7kamura/gospel"
	"path"
	"runtime"
	"testing"
)

func TestInterpretor(t *testing.T) {
	Describe(t, "インタープリターのテスト", func() {
		data := []byte{'A', ' ', 'B', '\t', '\r', '\n'}
		Context("インスタンス生成", func() {
			It("インスタンスが生成できること", func() {
				actual := NewInterpreter(data)
				Expect(actual).To(Exist)
				Expect(actual.origin).To(Equal, data)
			})
		})
		Context("データを読み込んで命令を実行する", func() {
			It("不要なデータを排除してスペースやタブをそれぞれSとTに置き換えること", func() {
				data := []byte{' ', ' ', '	', ' ', ' ', ' ', ' ', ' ', '	', '\n'}
				expected := []byte{'S', 'S', 'T', 'S', 'S', 'S', 'S', 'S', 'T', '\n'}
				sut := NewInterpreter(data)
				dat, err := sut.ToChar()
				Expect(err).To(NotExist)
				Expect(dat).To(Equal, expected)
			})
			It("不要なデータを排除してスペースやタブをそれぞれSとTに置き換えること パート２", func() {
				data := []byte{' ', ' ', '	', ' ', ' ', ' ', ' ', '	', '	', '\n'}
				expected := []byte{'S', 'S', 'T', 'S', 'S', 'S', 'S', 'T', 'T', '\n'}
				sut := NewInterpreter(data)
				dat, err := sut.ToChar()
				Expect(err).To(NotExist)
				Expect(dat).To(Equal, expected)
			})
			It("不要なデータを排除して指定したコマンドが作成されること", func() {
				data := []byte{' ', ' ', '	', ' ', ' ', ' ', ' ', ' ', '	', '\n'}
				sut := NewInterpreter(data)
				err := sut.ToCode()
				Expect(err).To(NotExist)
				Expect(sut.commands.Len()).To(Equal, 1)
				Expect(sut.commands.Get(1)).To(Equal, NewSubCommandWithParam("stack", "push", 0x41))
			})
			It("不要なデータを排除して指定したコマンドが作成されること", func() {
				data := []byte{' ', ' ', '	', ' ', ' ', ' ', ' ', '	', '	', '\n'}
				sut := NewInterpreter(data)
				err := sut.ToCode()
				Expect(err).To(NotExist)
				Expect(sut.commands.Len()).To(Equal, 1)
				Expect(sut.commands.Get(1)).To(Equal, NewSubCommandWithParam("stack", "push", 0x43))
			})
		})
		Context("不要な文字を排除する関数", func() {
			It("不要なデータ以外排除されていること", func() {
				sut := NewInterpreter(data)
				sut.filter()
				Expect(sut.source).To(Equal, []byte{' ', '\t', '\n'})
			})
			It("不要なデータ以外排除されていること（改行を増やす）", func() {
				data = append(data, '\n')
				sut := NewInterpreter(data)
				sut.filter()
				Expect(sut.source).To(Equal, []byte{' ', '\t', '\n', '\n'})
			})
		})
	})
}

func current_dir() string {
	_, fpath, _, _ := runtime.Caller(0)
	return path.Dir(fpath)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"
	"net/http"
	"strconv"
	"sync"
)

type dollars float32

func (d dollars) String() string {
	return fmt.Sprintf("$%.2f", d)
}

type database struct {
	sync.Mutex
	items map[string]dollars
}

func (db *database) list(w http.ResponseWriter, req *http.Request) {
	for item, price := range db.items {
		fmt.Fprintf(w, "%s: %s\n", item, price)
	}
}

func (db *database) price(w http.ResponseWriter, req *http.Request) {
	item := req.URL.Query().Get("item")
	if price, ok := db.items[item]; ok {
		fmt.Fprintf(w, "%s\n", price)
	} else {
		w.WriteHeader(http.StatusNotFound)
		fmt.Fprintf(w, "no such item: %q\n", item)
	}
}

func (db *database) update(w http.ResponseWriter, req *http.Request) {
	item, price, err := getParams(req)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		fmt.Fprint(w, err)
		return
	}

	if _, ok := db.items[item]; !ok {
		w.WriteHeader(http.StatusNotFound)
		fmt.Fprintf(w, "no such item: %q\n", item)
		return
	}
	db.Lock()
	db.items[item] = price
	db.Unlock() // Don't use defer to release lock earlier.

	w.WriteHeader(http.StatusOK)
}

func (db *database) create(w http.ResponseWriter, req *http.Request) {
	item, price, err := getParams(req)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		fmt.Fprint(w, err)
		return
	}

	if _, ok := db.items[item]; ok {
		w.WriteHeader(http.StatusBadRequest)
		fmt.Fprintf(w, "item already exists %q\n", item)
		return
	}
	db.Lock()
	db.items[item] = price
	db.Unlock() // Don't use defer to release lock earlier.

	w.WriteHeader(http.StatusOK)
}

func (db *database) remove(w http.ResponseWriter, req *http.Request) {
	item := req.URL.Query().Get("item")

	if _, ok := db.items[item]; !ok {
		w.WriteHeader(http.StatusNotFound)
		fmt.Fprintf(w, "no such item: %q\n", item)
		return
	}
	db.Lock()
	delete(db.items, item)
	db.Unlock() // Don't use defer to release lock earlier.

	w.WriteHeader(http.StatusOK)
}

func getParams(req *http.Request) (string, dollars, error) {
	item := req.URL.Query().Get("item")
	price := req.URL.Query().Get("price")
	priceFloat, err := strconv.ParseFloat(price, 32)
	if err != nil {
		return item, dollars(0), fmt.Errorf("invalid value for price: '%s'", price)
	}
	return item, dollars(priceFloat), nil
}

func main() {
	db := database{items: map[string]dollars{"shoes": 50, "socks": 5}}
	// Don't use here PUT, POST, DELETE due to task requirements.
	http.HandleFunc("/list", db.list)
	http.HandleFunc("/price", db.price)
	http.HandleFunc("/create", db.create)
	http.HandleFunc("/update", db.update)
	http.HandleFunc("/remove", db.remove)
	log.Fatal(http.ListenAndServe("localhost:8000", nil))
}
<commit_msg>[7.11] replaced 200 to 204<commit_after>package main

import (
	"fmt"
	"log"
	"net/http"
	"strconv"
	"sync"
)

type dollars float32

func (d dollars) String() string {
	return fmt.Sprintf("$%.2f", d)
}

type database struct {
	sync.Mutex
	items map[string]dollars
}

func (db *database) list(w http.ResponseWriter, req *http.Request) {
	for item, price := range db.items {
		fmt.Fprintf(w, "%s: %s\n", item, price)
	}
}

func (db *database) price(w http.ResponseWriter, req *http.Request) {
	item := req.URL.Query().Get("item")
	if price, ok := db.items[item]; ok {
		fmt.Fprintf(w, "%s\n", price)
	} else {
		w.WriteHeader(http.StatusNotFound)
		fmt.Fprintf(w, "no such item: %q\n", item)
	}
}

func (db *database) update(w http.ResponseWriter, req *http.Request) {
	item, price, err := getParams(req)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		fmt.Fprint(w, err)
		return
	}

	if _, ok := db.items[item]; !ok {
		w.WriteHeader(http.StatusNotFound)
		fmt.Fprintf(w, "no such item: %q\n", item)
		return
	}
	db.Lock()
	db.items[item] = price
	db.Unlock() // Don't use defer to release lock earlier.

	w.WriteHeader(http.StatusNoContent)
}

func (db *database) create(w http.ResponseWriter, req *http.Request) {
	item, price, err := getParams(req)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		fmt.Fprint(w, err)
		return
	}

	if _, ok := db.items[item]; ok {
		w.WriteHeader(http.StatusBadRequest)
		fmt.Fprintf(w, "item already exists %q\n", item)
		return
	}
	db.Lock()
	db.items[item] = price
	db.Unlock() // Don't use defer to release lock earlier.

	w.WriteHeader(http.StatusNoContent)
}

func (db *database) remove(w http.ResponseWriter, req *http.Request) {
	item := req.URL.Query().Get("item")

	if _, ok := db.items[item]; !ok {
		w.WriteHeader(http.StatusNotFound)
		fmt.Fprintf(w, "no such item: %q\n", item)
		return
	}
	db.Lock()
	delete(db.items, item)
	db.Unlock() // Don't use defer to release lock earlier.

	w.WriteHeader(http.StatusNoContent)
}

func getParams(req *http.Request) (string, dollars, error) {
	item := req.URL.Query().Get("item")
	price := req.URL.Query().Get("price")
	priceFloat, err := strconv.ParseFloat(price, 32)
	if err != nil {
		return item, dollars(0), fmt.Errorf("invalid value for price: '%s'", price)
	}
	return item, dollars(priceFloat), nil
}

func main() {
	db := database{items: map[string]dollars{"shoes": 50, "socks": 5}}
	// Don't use here PUT, POST, DELETE due to task requirements.
	http.HandleFunc("/list", db.list)
	http.HandleFunc("/price", db.price)
	http.HandleFunc("/create", db.create)
	http.HandleFunc("/update", db.update)
	http.HandleFunc("/remove", db.remove)
	log.Fatal(http.ListenAndServe("localhost:8000", nil))
}
<|endoftext|>


<commit_before>package main

import (
	"bosun.org/graphite"
	"encoding/json"
	"fmt"
	"github.com/Dieterbe/go-metrics"
	"github.com/raintank/raintank-metric/schema"
	"io/ioutil"
	"log"
	"math"
	"math/rand"
	"net"
	"net/http"
	"os"
	"strconv"
	"sync"
	"time"
	//"strings"
)

type Hit struct {
	Index  string                  `json:"_index"`
	Type   string                  `json:"_type"`
	Id     string                  `json:"_id"`
	Score  float64                 `json:"_score"`
	Source schema.MetricDefinition `json:"_source"`
}

type EsResult struct {
	Took     int
	TimedOut bool
	_shards  struct {
		total      int
		successful int
		failed     int
	}
	Hits struct {
		Total    int
		MaxScore int
		Hits     []Hit
	}
}

func perror(err error) {
	if err != nil {
		log.Fatal(err)
	}
}

type stat struct {
	def       schema.MetricDefinition
	firstSeen int64
}

var targets map[string]stat
var targetKeys []string
var targetsLock sync.Mutex

func init() {
	targets = make(map[string]stat)
	targetKeys = make([]string, 0)
}

func main() {
	if len(os.Args) != 4 {
		log.Fatal("usage: graphite-watcher <elasticsearch-addr> <metrics-addr> <graphite-addr>")
	}
	addr, _ := net.ResolveTCPAddr("tcp", os.Args[2])
	go metrics.Graphite(metrics.DefaultRegistry, 10e9, "graphite-watcher.", addr)
	lag := metrics.NewHistogram(metrics.NewExpDecaySample(1028, 0.015))
	metrics.Register("lag", lag)
	numMetrics := metrics.NewGauge()
	metrics.Register("num_metrics", numMetrics)
	nullPoints := metrics.NewCounter()
	metrics.Register("null_points", nullPoints)

	// for a metric to exist in ES at t=Y, there must at least have been 1 point for that metric
	// at a time X where X < Y.  Hence, we confidently say that if we see a metric at Y, we can
	// demand data to show up for that metric at >=Y
	// for our data check to be useful we need metrics to show up in ES soon after being in the pipeline,
	// which seems to be true (see nsqadmin)
	updateTargets := func() {
		getEsTick := time.NewTicker(time.Second * time.Duration(10))
		for ts := range getEsTick.C {
			tsUnix := ts.Unix()
			res, err := http.Get("http://" + os.Args[1] + "/metric/_search?q=*:*&size=10000000")
			perror(err)
			defer res.Body.Close()
			body, err := ioutil.ReadAll(res.Body)
			perror(err)
			var data EsResult
			err = json.Unmarshal(body, &data)
			perror(err)
			amount := len(data.Hits.Hits)
			numMetrics.Update(int64(amount))
			for _, h := range data.Hits.Hits {
				//fmt.Println(h.Source.Name)
				targetsLock.Lock()
				if _, ok := targets[h.Source.Name]; !ok {
					targets[h.Source.Name] = stat{h.Source, tsUnix}
					targetKeys = append(targetKeys, h.Source.Name)
				}
				targetsLock.Unlock()
			}
		}
	}

	test := func(wg *sync.WaitGroup, curTs int64, met stat) {
		defer wg.Done()
		g := graphite.HostHeader{Host: "http://" + os.Args[3] + "/render", Header: http.Header{}}
		g.Header.Add("X-Org-Id", strconv.FormatInt(int64(met.def.OrgId), 10))
		g.Header.Set("User-Agent", "graphite-watcher")
		q := graphite.Request{Targets: []string{met.def.Name}}
		series, err := g.Query(&q)
		if err != nil {
			log.Println("ERROR querying graphite:", err)
			return
		}
		for _, serie := range series {
			if met.def.Name != serie.Target {
				fmt.Println("ERROR: name != target name:", met.def.Name, serie.Target)
			}

			lastTs := int64(0)
			oldestNull := int64(math.MaxInt64)
			if len(serie.Datapoints) == 0 {
				fmt.Println("ERROR: series for", met.def.Name, "contains no points!")
			}
			for _, p := range serie.Datapoints {
				ts, err := p[1].Int64()
				if err != nil {
					fmt.Println("ERROR: could not parse timestamp", p)
				}
				if ts <= lastTs {
					fmt.Println("ERROR: timestamp must be bigger than last", lastTs, ts)
				}
				if lastTs == 0 && (ts < curTs-24*3600-60 || ts > curTs-24*3600+60) {
					fmt.Println("ERROR: first point", p, "should have been about 24h ago")
				}
				if lastTs != 0 && ts != lastTs+int64(met.def.Interval) {
					fmt.Println("ERROR: point", p, " is not interval ", met.def.Interval, "apart from previous point")
				}
				_, err = p[0].Float64()
				if err != nil && ts > met.firstSeen {
					if ts < oldestNull {
						oldestNull = ts
					}
					if ts < curTs-30 {
						nullPoints.Inc(1)
						// will show up in lag metric too
						fmt.Println("ERROR: ", met.def.Name, " at", curTs, "seeing a null for ts", p[1])
						// should this be a metric we track?
					}
				} else {
					// we saw a valid point, so reset oldestNull.
					oldestNull = int64(math.MaxInt64)
				}
				lastTs = ts
			}
			if lastTs < curTs-int64(met.def.Interval) || lastTs > curTs+int64(met.def.Interval) {
				fmt.Println("ERROR: last point at ", lastTs, "is out of range")
			}
			// if there was no null, we treat the point after the last one we had as null
			if oldestNull == math.MaxInt64 {
				oldestNull = lastTs + int64(met.def.Interval)
			}
			// lag cannot be < 0
			if oldestNull > curTs {
				oldestNull = curTs
			}
			// lag is from first null, even if there were non-nulls after it
			//fmt.Printf("%60s - lag %d\n", name, curTs-oldestNull)
			lag.Update(curTs - oldestNull)
		}
	}

	go updateTargets()

	tick := time.NewTicker(time.Millisecond * time.Duration(100))
	wg := &sync.WaitGroup{}
	for ts := range tick.C {
		var t stat
		targetsLock.Lock()
		if len(targetKeys) > 0 {
			key := targetKeys[rand.Intn(len(targets))]
			t = targets[key]
		}
		targetsLock.Unlock()
		if len(targetKeys) > 0 {
			wg.Add(1)
			go test(wg, ts.Unix(), t)
		}
	}
	wg.Wait()
}
<commit_msg>fix comments<commit_after>package main

import (
	"bosun.org/graphite"
	"encoding/json"
	"fmt"
	"github.com/Dieterbe/go-metrics"
	"github.com/raintank/raintank-metric/schema"
	"io/ioutil"
	"log"
	"math"
	"math/rand"
	"net"
	"net/http"
	"os"
	"strconv"
	"sync"
	"time"
	//"strings"
)

type Hit struct {
	Index  string                  `json:"_index"`
	Type   string                  `json:"_type"`
	Id     string                  `json:"_id"`
	Score  float64                 `json:"_score"`
	Source schema.MetricDefinition `json:"_source"`
}

type EsResult struct {
	Took     int
	TimedOut bool
	_shards  struct {
		total      int
		successful int
		failed     int
	}
	Hits struct {
		Total    int
		MaxScore int
		Hits     []Hit
	}
}

func perror(err error) {
	if err != nil {
		log.Fatal(err)
	}
}

type stat struct {
	def       schema.MetricDefinition
	firstSeen int64
}

var targets map[string]stat
var targetKeys []string
var targetsLock sync.Mutex

func init() {
	targets = make(map[string]stat)
	targetKeys = make([]string, 0)
}

func main() {
	if len(os.Args) != 4 {
		log.Fatal("usage: graphite-watcher <elasticsearch-addr> <metrics-addr> <graphite-addr>")
	}
	addr, _ := net.ResolveTCPAddr("tcp", os.Args[2])
	go metrics.Graphite(metrics.DefaultRegistry, 10e9, "graphite-watcher.", addr)
	lag := metrics.NewHistogram(metrics.NewExpDecaySample(1028, 0.015))
	metrics.Register("lag", lag)
	numMetrics := metrics.NewGauge()
	metrics.Register("num_metrics", numMetrics)
	nullPoints := metrics.NewCounter()
	metrics.Register("null_points", nullPoints)

	// for a metric to exist in ES at t=Y, there must at least have been 1 point for that metric
	// at a time X where X < Y.  Hence, we can confidently say that if we see a metric at Y, we can
	// demand data to show up for that metric at >=Y
	// for our data check to be useful we need metrics to show up in ES soon after being in the pipeline,
	// which seems to be true (see nsqadmin)
	updateTargets := func() {
		getEsTick := time.NewTicker(time.Second * time.Duration(10))
		for ts := range getEsTick.C {
			tsUnix := ts.Unix()
			res, err := http.Get("http://" + os.Args[1] + "/metric/_search?q=*:*&size=10000000")
			perror(err)
			defer res.Body.Close()
			body, err := ioutil.ReadAll(res.Body)
			perror(err)
			var data EsResult
			err = json.Unmarshal(body, &data)
			perror(err)
			amount := len(data.Hits.Hits)
			numMetrics.Update(int64(amount))
			for _, h := range data.Hits.Hits {
				//fmt.Println(h.Source.Name)
				targetsLock.Lock()
				if _, ok := targets[h.Source.Name]; !ok {
					targets[h.Source.Name] = stat{h.Source, tsUnix}
					targetKeys = append(targetKeys, h.Source.Name)
				}
				targetsLock.Unlock()
			}
		}
	}

	test := func(wg *sync.WaitGroup, curTs int64, met stat) {
		defer wg.Done()
		g := graphite.HostHeader{Host: "http://" + os.Args[3] + "/render", Header: http.Header{}}
		g.Header.Add("X-Org-Id", strconv.FormatInt(int64(met.def.OrgId), 10))
		g.Header.Set("User-Agent", "graphite-watcher")
		q := graphite.Request{Targets: []string{met.def.Name}}
		series, err := g.Query(&q)
		if err != nil {
			log.Println("ERROR querying graphite:", err)
			return
		}
		for _, serie := range series {
			if met.def.Name != serie.Target {
				fmt.Println("ERROR: name != target name:", met.def.Name, serie.Target)
			}

			lastTs := int64(0)
			oldestNull := int64(math.MaxInt64)
			if len(serie.Datapoints) == 0 {
				fmt.Println("ERROR: series for", met.def.Name, "contains no points!")
			}
			for _, p := range serie.Datapoints {
				ts, err := p[1].Int64()
				if err != nil {
					fmt.Println("ERROR: could not parse timestamp", p)
				}
				if ts <= lastTs {
					fmt.Println("ERROR: timestamp must be bigger than last", lastTs, ts)
				}
				if lastTs == 0 && (ts < curTs-24*3600-60 || ts > curTs-24*3600+60) {
					fmt.Println("ERROR: first point", p, "should have been about 24h ago")
				}
				if lastTs != 0 && ts != lastTs+int64(met.def.Interval) {
					fmt.Println("ERROR: point", p, " is not interval ", met.def.Interval, "apart from previous point")
				}
				_, err = p[0].Float64()
				if err != nil && ts > met.firstSeen {
					if ts < oldestNull {
						oldestNull = ts
					}
					if ts < curTs-30 {
						nullPoints.Inc(1)
						fmt.Println("ERROR: ", met.def.Name, " at", curTs, "seeing a null for ts", p[1])
					}
				} else {
					// we saw a valid point, so reset oldestNull.
					oldestNull = int64(math.MaxInt64)
				}
				lastTs = ts
			}
			if lastTs < curTs-int64(met.def.Interval) || lastTs > curTs+int64(met.def.Interval) {
				fmt.Println("ERROR: last point at ", lastTs, "is out of range")
			}
			// if there was no null, we treat the point after the last one we had as null
			if oldestNull == math.MaxInt64 {
				oldestNull = lastTs + int64(met.def.Interval)
			}
			// lag cannot be < 0
			if oldestNull > curTs {
				oldestNull = curTs
			}
			// lag is from first null after a range of data until now
			//fmt.Printf("%60s - lag %d\n", name, curTs-oldestNull)
			lag.Update(curTs - oldestNull)
		}
	}

	go updateTargets()

	tick := time.NewTicker(time.Millisecond * time.Duration(100))
	wg := &sync.WaitGroup{}
	for ts := range tick.C {
		var t stat
		targetsLock.Lock()
		if len(targetKeys) > 0 {
			key := targetKeys[rand.Intn(len(targets))]
			t = targets[key]
		}
		targetsLock.Unlock()
		if len(targetKeys) > 0 {
			wg.Add(1)
			go test(wg, ts.Unix(), t)
		}
	}
	wg.Wait()
}
<|endoftext|>


<commit_before>/*
	Package rrdialer implements a dialer which provides improved behaviour for
	hostnames with multiple IP addresses (aka. round-robin DNS).
	Unlike the default net.Dialer which divides the timeout between the multiple
	endpoints (IP addresses), the round-robin dialer tracks the historic
	performance of each endpoint to dial the fastest endpoint first and will
	concurrently dial other endpoints if unusually long connection times are
	observed.
*/
package rrdialer

import (
	"context"
	"net"
	"sync"
	"time"

	"github.com/Symantec/Dominator/lib/log"
)

type Dialer struct {
	dirname   string
	logger    log.DebugLogger
	rawDialer *net.Dialer
	waitGroup sync.WaitGroup
}

// New creates a new Dialer. The underlying raw dialer used to dial each
// endpoint is given by dialer. The directory in which endpoint statitics are
// written is given by cacheDir. If this is the empty string then the ".cache"
// subdirectory of the home directory is used. Log messages are written to
// logger. If the debug level is 3 or greater then all endpoints are dialed.
func New(dialer *net.Dialer, cacheDir string,
	logger log.DebugLogger) (*Dialer, error) {
	return newDialer(dialer, cacheDir, logger)
}

// Dial connects to the address on the named network.
func (d *Dialer) Dial(network, address string) (net.Conn, error) {
	return d.dialContext(context.Background(), network, address)
}

// DialContext connects to the address on the named network using the provided
// context.
func (d *Dialer) DialContext(ctx context.Context, network,
	address string) (net.Conn, error) {
	return d.dialContext(ctx, network, address)
}

// WaitForBackgroundResults will wait up to timeout for other endpoint
// connection attempts to complete, so that their performance statistics can be
// saved. It is recommended to call this just before the main function returns.
func (d *Dialer) WaitForBackgroundResults(timeout time.Duration) {
	d.waitForBackgroundResults(timeout)
}
<commit_msg>Add a blank line to lib/net/rrdialer package description.<commit_after>/*
	Package rrdialer implements a dialer which provides improved behaviour for
	hostnames with multiple IP addresses (aka. round-robin DNS).

	Unlike the default net.Dialer which divides the timeout between the multiple
	endpoints (IP addresses), the round-robin dialer tracks the historic
	performance of each endpoint to dial the fastest endpoint first and will
	concurrently dial other endpoints if unusually long connection times are
	observed.
*/
package rrdialer

import (
	"context"
	"net"
	"sync"
	"time"

	"github.com/Symantec/Dominator/lib/log"
)

type Dialer struct {
	dirname   string
	logger    log.DebugLogger
	rawDialer *net.Dialer
	waitGroup sync.WaitGroup
}

// New creates a new Dialer. The underlying raw dialer used to dial each
// endpoint is given by dialer. The directory in which endpoint statitics are
// written is given by cacheDir. If this is the empty string then the ".cache"
// subdirectory of the home directory is used. Log messages are written to
// logger. If the debug level is 3 or greater then all endpoints are dialed.
func New(dialer *net.Dialer, cacheDir string,
	logger log.DebugLogger) (*Dialer, error) {
	return newDialer(dialer, cacheDir, logger)
}

// Dial connects to the address on the named network.
func (d *Dialer) Dial(network, address string) (net.Conn, error) {
	return d.dialContext(context.Background(), network, address)
}

// DialContext connects to the address on the named network using the provided
// context.
func (d *Dialer) DialContext(ctx context.Context, network,
	address string) (net.Conn, error) {
	return d.dialContext(ctx, network, address)
}

// WaitForBackgroundResults will wait up to timeout for other endpoint
// connection attempts to complete, so that their performance statistics can be
// saved. It is recommended to call this just before the main function returns.
func (d *Dialer) WaitForBackgroundResults(timeout time.Duration) {
	d.waitForBackgroundResults(timeout)
}
<|endoftext|>


<commit_before>package main

import (
	".."
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"github.com/BurntSushi/toml"
	"github.com/bitly/go-nsq"
	log "github.com/Sirupsen/logrus"
	"net/http"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"
)

var (
	ip = flag.String("ip", "4.3.2.1", "Node ip address")
	configFile = flag.String("config", "haaas.conf", "Configuration file")
	version = flag.Bool("version", false, "Print current version")
	verbose = flag.Bool("verbose", false, "Log in verbose mode")
	config = nsq.NewConfig()
	properties *haaasd.Config
	daemon      *haaasd.Daemon
	producer    *nsq.Producer
	syslog        *haaasd.Syslog
)

func main() {
	log.SetFormatter(&log.TextFormatter{})
	flag.Parse()

	if *version {
		println(haaasd.AppVersion)
		os.Exit(0)
	}

	if *verbose {
		log.SetLevel(log.DebugLevel)
		log.WithField("loglevel", "debug").Info("Change loglevel")
	}

	loadProperties()

	daemon = haaasd.NewDaemon(properties)
	syslog = haaasd.NewSyslog(properties)
	syslog.Init()
	log.WithFields(log.Fields{
		"status": properties.Status,
		"id":    properties.NodeId(),
	}).Info("Starting haaasd")

	producer, _ = nsq.NewProducer(properties.ProducerAddr, config)

	initProducer()
	time.Sleep(1 * time.Second)

	var wg sync.WaitGroup
	// Start http API
	restApi := haaasd.NewRestApi(properties)
	go func() {
		defer wg.Done()
		wg.Add(1)
		err := restApi.Start()
		if err != nil {
			log.Fatal("Cannot start api")
		}
	}()

	// Start slave consumer
	go func() {
		defer wg.Done()
		wg.Add(1)
		consumer, _ := nsq.NewConsumer(fmt.Sprintf("commit_requested_%s", properties.ClusterId), properties.NodeId(), config)
		consumer.AddHandler(nsq.HandlerFunc(onCommitRequested))
		err := consumer.ConnectToNSQLookupd(properties.LookupdAddr)
		if err != nil {
			log.Panic("Could not connect")
		}
	}()

	// Start master consumer
	go func() {
		defer wg.Done()
		wg.Add(1)
		consumer, _ := nsq.NewConsumer(fmt.Sprintf("commit_slave_completed_%s", properties.ClusterId), properties.NodeId(), config)
		consumer.AddHandler(nsq.HandlerFunc(onCommitSlaveRequested))
		err := consumer.ConnectToNSQLookupd(properties.LookupdAddr)
		if err != nil {
			log.Panic("Could not connect")
		}
	}()

	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM, syscall.SIGQUIT)
	select {
	case signal := <-sigChan:
		log.Printf("Got signal: %v\n", signal)
	}
	restApi.Stop()

	log.Printf("Waiting on server to stop\n")
	wg.Wait()
}

func initProducer() {
	// Create required topics
	topics := []string{"commit_slave_completed", "commit_completed", "commit_failed"}
	channels := []string{"slave", "master"}
	topicChan := make(chan string, len(topics))
	for i := range topics {
		topicChan <- topics[i]
	}
	left := len(topics)
	for left > 0 {
		topic := <-topicChan
		log.WithField("topic", topic).Info("Creating topic")
		url := fmt.Sprintf("%s/topic/create?topic=%s_%s", properties.ProducerRestAddr, topic, properties.ClusterId)
		resp, err := http.PostForm(url, nil)
		if err != nil || resp.StatusCode != 200 {
			topicChan <- topic
			continue
		}
		for channel := range channels {
			log.WithField("channel", channels[channel]).Info("Creating channel")
			url := fmt.Sprintf("%s/channel/create?topic=%s_%s&channel=%s-%s", properties.ProducerRestAddr, topic, properties.ClusterId, properties.ClusterId, channels[channel])
			resp, err := http.PostForm(url, nil)
			if err != nil || resp.StatusCode != 200 {
				topicChan <- topic
				continue
			}
		}

		log.WithField("topic", topic).Info("Topic created")
		left--

	}
}

// loadProperties load properties file
func loadProperties() {
	properties = haaasd.DefaultConfig()
	if _, err := toml.DecodeFile(*configFile, properties); err != nil {
		log.Fatal(err)
		os.Exit(1)
	}
	properties.IpAddr = *ip
	len := len(properties.HapHome)
	if properties.HapHome[len - 1] == '/' {
		properties.HapHome = properties.HapHome[:len - 1]
	}
}

func filteredHandler(event string, message *nsq.Message, target string, f haaasd.HandlerFunc) error {
	defer message.Finish()
	match, err := daemon.Is(target)
	if err != nil {
		return err
	}

	if match {
		log.WithField("event", event).WithField("payload", message.Body).Debug("Handle event")
		data, err := bodyToData(message.Body)
		if err != nil {
			log.WithError(err).Error("Unable to read data")
			return err
		}
		f(data)
	} else {
		log.WithField("event", event).Debug("Ignore event")
	}

	return nil
}

func onCommitRequested(message *nsq.Message) error {
	return filteredHandler("commit_requested", message, "slave", reloadSlave)
}
func onCommitSlaveRequested(message *nsq.Message) error {
	return filteredHandler("commit_slave_completed_", message, "master", reloadMaster)
}

func reloadSlave(data *haaasd.EventMessage) error {
	hap := haaasd.NewHaproxy("slave",properties, data.Application, data.Platform, data.HapVersion)
	status, err := hap.ApplyConfiguration(data)
	if err == nil {
		if status != haaasd.UNCHANGED {
			syslog.Restart()
		}
		publishMessage("commit_slave_completed_", data)
	} else {
		log.WithFields(log.Fields{
			"correlationId": data.Correlationid,
			"application" : data.Application,
			"platform": data.Platform,
		}).WithError(err).Error("Commit failed")
		publishMessage("commit_failed_", map[string]string{"application": data.Application, "platform": data.Platform, "correlationid": data.Correlationid})
	}
	return nil
}

func reloadMaster(data *haaasd.EventMessage) error {
	hap := haaasd.NewHaproxy("master",properties, data.Application, data.Platform, data.HapVersion)
	status, err := hap.ApplyConfiguration(data)
	if err == nil {
		if status != haaasd.UNCHANGED {
			syslog.Restart()
		}
		publishMessage("commit_completed_", map[string]string{"application": data.Application, "platform": data.Platform, "correlationid": data.Correlationid})
	} else {
		log.WithFields(log.Fields{
			"correlationId": data.Correlationid,
			"application" : data.Application,
			"platform": data.Platform,
		}).WithError(err).Error("Commit failed")
		publishMessage("commit_failed_", map[string]string{"application": data.Application, "platform": data.Platform, "correlationid": data.Correlationid})
	}
	return nil
}

// Unmarshal json to EventMessage
func bodyToData(jsonStream []byte) (*haaasd.EventMessage, error) {
	dec := json.NewDecoder(bytes.NewReader(jsonStream))
	var message haaasd.EventMessage
	err := dec.Decode(&message)
	return &message, err
}

func publishMessage(topic_prefix string, data interface{}) error {
	jsonMsg, _ := json.Marshal(data)
	topic := topic_prefix + properties.ClusterId
	log.WithField("topic", topic).WithField("payload", string(jsonMsg)).Debug("Publish")
	return producer.Publish(topic, []byte(jsonMsg))
}
<commit_msg>Remove payload from event trace<commit_after>package main

import (
	".."
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"github.com/BurntSushi/toml"
	"github.com/bitly/go-nsq"
	log "github.com/Sirupsen/logrus"
	"net/http"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"
)

var (
	ip = flag.String("ip", "4.3.2.1", "Node ip address")
	configFile = flag.String("config", "haaas.conf", "Configuration file")
	version = flag.Bool("version", false, "Print current version")
	verbose = flag.Bool("verbose", false, "Log in verbose mode")
	config = nsq.NewConfig()
	properties *haaasd.Config
	daemon      *haaasd.Daemon
	producer    *nsq.Producer
	syslog        *haaasd.Syslog
)

func main() {
	log.SetFormatter(&log.TextFormatter{})
	flag.Parse()

	if *version {
		println(haaasd.AppVersion)
		os.Exit(0)
	}

	if *verbose {
		log.SetLevel(log.DebugLevel)
		log.WithField("loglevel", "debug").Info("Change loglevel")
	}

	loadProperties()

	daemon = haaasd.NewDaemon(properties)
	syslog = haaasd.NewSyslog(properties)
	syslog.Init()
	log.WithFields(log.Fields{
		"status": properties.Status,
		"id":    properties.NodeId(),
	}).Info("Starting haaasd")

	producer, _ = nsq.NewProducer(properties.ProducerAddr, config)

	initProducer()
	time.Sleep(1 * time.Second)

	var wg sync.WaitGroup
	// Start http API
	restApi := haaasd.NewRestApi(properties)
	go func() {
		defer wg.Done()
		wg.Add(1)
		err := restApi.Start()
		if err != nil {
			log.Fatal("Cannot start api")
		}
	}()

	// Start slave consumer
	go func() {
		defer wg.Done()
		wg.Add(1)
		consumer, _ := nsq.NewConsumer(fmt.Sprintf("commit_requested_%s", properties.ClusterId), properties.NodeId(), config)
		consumer.AddHandler(nsq.HandlerFunc(onCommitRequested))
		err := consumer.ConnectToNSQLookupd(properties.LookupdAddr)
		if err != nil {
			log.Panic("Could not connect")
		}
	}()

	// Start master consumer
	go func() {
		defer wg.Done()
		wg.Add(1)
		consumer, _ := nsq.NewConsumer(fmt.Sprintf("commit_slave_completed_%s", properties.ClusterId), properties.NodeId(), config)
		consumer.AddHandler(nsq.HandlerFunc(onCommitSlaveRequested))
		err := consumer.ConnectToNSQLookupd(properties.LookupdAddr)
		if err != nil {
			log.Panic("Could not connect")
		}
	}()

	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM, syscall.SIGQUIT)
	select {
	case signal := <-sigChan:
		log.Printf("Got signal: %v\n", signal)
	}
	restApi.Stop()

	log.Printf("Waiting on server to stop\n")
	wg.Wait()
}

func initProducer() {
	// Create required topics
	topics := []string{"commit_slave_completed", "commit_completed", "commit_failed"}
	channels := []string{"slave", "master"}
	topicChan := make(chan string, len(topics))
	for i := range topics {
		topicChan <- topics[i]
	}
	left := len(topics)
	for left > 0 {
		topic := <-topicChan
		log.WithField("topic", topic).Info("Creating topic")
		url := fmt.Sprintf("%s/topic/create?topic=%s_%s", properties.ProducerRestAddr, topic, properties.ClusterId)
		resp, err := http.PostForm(url, nil)
		if err != nil || resp.StatusCode != 200 {
			topicChan <- topic
			continue
		}
		for channel := range channels {
			log.WithField("channel", channels[channel]).Info("Creating channel")
			url := fmt.Sprintf("%s/channel/create?topic=%s_%s&channel=%s-%s", properties.ProducerRestAddr, topic, properties.ClusterId, properties.ClusterId, channels[channel])
			resp, err := http.PostForm(url, nil)
			if err != nil || resp.StatusCode != 200 {
				topicChan <- topic
				continue
			}
		}

		log.WithField("topic", topic).Info("Topic created")
		left--

	}
}

// loadProperties load properties file
func loadProperties() {
	properties = haaasd.DefaultConfig()
	if _, err := toml.DecodeFile(*configFile, properties); err != nil {
		log.Fatal(err)
		os.Exit(1)
	}
	properties.IpAddr = *ip
	len := len(properties.HapHome)
	if properties.HapHome[len - 1] == '/' {
		properties.HapHome = properties.HapHome[:len - 1]
	}
}

func filteredHandler(event string, message *nsq.Message, target string, f haaasd.HandlerFunc) error {
	defer message.Finish()
	match, err := daemon.Is(target)
	if err != nil {
		return err
	}

	if match {
		log.WithField("event", event).Debug("Handle event")
		data, err := bodyToData(message.Body)
		if err != nil {
			log.WithError(err).Error("Unable to read data")
			return err
		}
		f(data)
	} else {
		log.WithField("event", event).Debug("Ignore event")
	}

	return nil
}

func onCommitRequested(message *nsq.Message) error {
	return filteredHandler("commit_requested", message, "slave", reloadSlave)
}
func onCommitSlaveRequested(message *nsq.Message) error {
	return filteredHandler("commit_slave_completed_", message, "master", reloadMaster)
}

func reloadSlave(data *haaasd.EventMessage) error {
	hap := haaasd.NewHaproxy("slave",properties, data.Application, data.Platform, data.HapVersion)
	status, err := hap.ApplyConfiguration(data)
	if err == nil {
		if status != haaasd.UNCHANGED {
			syslog.Restart()
		}
		publishMessage("commit_slave_completed_", data)
	} else {
		log.WithFields(log.Fields{
			"correlationId": data.Correlationid,
			"application" : data.Application,
			"platform": data.Platform,
		}).WithError(err).Error("Commit failed")
		publishMessage("commit_failed_", map[string]string{"application": data.Application, "platform": data.Platform, "correlationid": data.Correlationid})
	}
	return nil
}

func reloadMaster(data *haaasd.EventMessage) error {
	hap := haaasd.NewHaproxy("master",properties, data.Application, data.Platform, data.HapVersion)
	status, err := hap.ApplyConfiguration(data)
	if err == nil {
		if status != haaasd.UNCHANGED {
			syslog.Restart()
		}
		publishMessage("commit_completed_", map[string]string{"application": data.Application, "platform": data.Platform, "correlationid": data.Correlationid})
	} else {
		log.WithFields(log.Fields{
			"correlationId": data.Correlationid,
			"application" : data.Application,
			"platform": data.Platform,
		}).WithError(err).Error("Commit failed")
		publishMessage("commit_failed_", map[string]string{"application": data.Application, "platform": data.Platform, "correlationid": data.Correlationid})
	}
	return nil
}

// Unmarshal json to EventMessage
func bodyToData(jsonStream []byte) (*haaasd.EventMessage, error) {
	dec := json.NewDecoder(bytes.NewReader(jsonStream))
	var message haaasd.EventMessage
	err := dec.Decode(&message)
	return &message, err
}

func publishMessage(topic_prefix string, data interface{}) error {
	jsonMsg, _ := json.Marshal(data)
	topic := topic_prefix + properties.ClusterId
	log.WithField("topic", topic).WithField("payload", string(jsonMsg)).Debug("Publish")
	return producer.Publish(topic, []byte(jsonMsg))
}
<|endoftext|>