<commit_before>package systests

import (
	"bytes"
	"net/http"
	"testing"
	"time"

	"golang.org/x/net/context"

	"github.com/keybase/client/go/client"
	"github.com/keybase/client/go/libkb"
	keybase1 "github.com/keybase/client/go/protocol/keybase1"
	"github.com/keybase/client/go/protocol/stellar1"
	"github.com/keybase/client/go/stellar"
	"github.com/keybase/client/go/teams"
	"github.com/keybase/stellarnet"
	"github.com/stellar/go/build"
	"github.com/stellar/go/clients/horizon"
	"github.com/stretchr/testify/require"
)

func TestStellarNoteRoundtripAndResets(t *testing.T) {
	ctx := newSMUContext(t)
	defer ctx.cleanup()

	// Sign up two users, bob and alice.
	alice := ctx.installKeybaseForUser("alice", 10)
	alice.signup()
	divDebug(ctx, "Signed up alice (%s)", alice.username)
	bob := ctx.installKeybaseForUser("bob", 10)
	bob.signup()
	divDebug(ctx, "Signed up bob (%s)", bob.username)

	t.Logf("note to self")
	encB64, err := stellar.NoteEncryptB64(context.Background(), alice.getPrimaryGlobalContext(), sampleNote(), nil)
	require.NoError(t, err)
	note, err := stellar.NoteDecryptB64(context.Background(), alice.getPrimaryGlobalContext(), encB64)
	require.NoError(t, err)
	require.Equal(t, sampleNote(), note)

	t.Logf("note to both users")
	other := bob.userVersion()
	encB64, err = stellar.NoteEncryptB64(context.Background(), alice.getPrimaryGlobalContext(), sampleNote(), &other)
	require.NoError(t, err)

	t.Logf("decrypt as self")
	note, err = stellar.NoteDecryptB64(context.Background(), alice.getPrimaryGlobalContext(), encB64)
	require.NoError(t, err)
	require.Equal(t, sampleNote(), note)

	t.Logf("decrypt as other")
	note, err = stellar.NoteDecryptB64(context.Background(), bob.getPrimaryGlobalContext(), encB64)
	require.NoError(t, err)
	require.Equal(t, sampleNote(), note)

	t.Logf("reset sender")
	alice.reset()
	divDebug(ctx, "Reset bob (%s)", bob.username)
	alice.loginAfterReset(10)
	divDebug(ctx, "Bob logged in after reset")

	t.Logf("fail to decrypt as post-reset self")
	note, err = stellar.NoteDecryptB64(context.Background(), alice.getPrimaryGlobalContext(), encB64)
	require.Error(t, err)
	require.Equal(t, "note not encrypted for logged-in user", err.Error())

	t.Logf("decrypt as other")
	note, err = stellar.NoteDecryptB64(context.Background(), bob.getPrimaryGlobalContext(), encB64)
	require.NoError(t, err)
	require.Equal(t, sampleNote(), note)
}

// Test took 38s on a dev server 2018-06-07
func TestStellarRelayAutoClaims(t *testing.T) {
	testStellarRelayAutoClaims(t, false, false)
}

// Test took 29s on a dev server 2018-06-07
func TestStellarRelayAutoClaimsWithPUK(t *testing.T) {
	testStellarRelayAutoClaims(t, true, true)
}

// Part 1:
// XLM is sent to a user before they have a [PUK / wallet].
// In the form of multiple relay payments.
// They then [get a PUK,] add a wallet, and enter the impteam,
// which all kick the autoclaim into gear.
//
// Part 2:
// A relay payment is sent to the user who already has a wallet.
// The funds should be claimed asap.
//
// To debug this test use log filter "stellar_test|poll-|AutoClaim|stellar.claim|pollfor"
func testStellarRelayAutoClaims(t *testing.T, startWithPUK, skipPart2 bool) {
	tt := newTeamTester(t)
	defer tt.cleanup()
	useStellarTestNet(t)

	alice := tt.addUser("alice")
	var bob *userPlusDevice
	if startWithPUK {
		bob = tt.addWalletlessUser("bob")
	} else {
		bob = tt.addPuklessUser("bob")
	}
	alice.kickTeamRekeyd()

	t.Logf("alice gets funded")
	res, err := alice.stellarClient.GetWalletAccountsLocal(context.Background(), 0)
	require.NoError(t, err)
	gift(t, res[0].AccountID)

	t.Logf("alice sends a first relay payment to bob P1")
	attachIdentifyUI(t, alice.tc.G, newSimpleIdentifyUI())
	cmd := client.CmdWalletSend{
		Contextified: libkb.NewContextified(alice.tc.G),
		Recipient:    bob.username,
		Amount:       "50",
	}
	for i := 0; i < retryCount; i++ {
		err = cmd.Run()
		if err == nil {
			break
		}
	}
	require.NoError(t, err)

	t.Logf("alice sends a second relay payment to bob P2")
	cmd = client.CmdWalletSend{
		Contextified: libkb.NewContextified(alice.tc.G),
		Recipient:    bob.username,
		Amount:       "30",
	}
	for i := 0; i < retryCount; i++ {
		err = cmd.Run()
		if err == nil {
			break
		}
	}
	require.NoError(t, err)

	t.Logf("get the impteam seqno to wait on later")
	team, _, _, err := teams.LookupImplicitTeam(context.Background(), alice.tc.G, alice.username+","+bob.username, false)
	require.NoError(t, err)
	nextSeqno := team.NextSeqno()

	if startWithPUK {
		t.Logf("bob gets a wallet")
		bob.tc.Tp.DisableAutoWallet = false
		bob.tc.G.GetStellar().CreateWalletSoft(context.Background())
	} else {
		t.Logf("bob gets a PUK and wallet")
		bob.perUserKeyUpgrade()
		bob.tc.G.GetStellar().CreateWalletSoft(context.Background())

		t.Logf("wait for alice to add bob to their impteam")
		alice.pollForTeamSeqnoLinkWithLoadArgs(keybase1.LoadTeamArg{ID: team.ID}, nextSeqno)
	}

	pollTime := 20 * time.Second
	if libkb.UseCITime(bob.tc.G) {
		// This test is especially slow.
		pollTime = 30 * time.Second
	}

	pollFor(t, "claims to complete", pollTime, bob.tc.G, func(i int) bool {
		res, err = bob.stellarClient.GetWalletAccountsLocal(context.Background(), 0)
		require.NoError(t, err)
		t.Logf("poll-1-%v: %v", i, res[0].BalanceDescription)
		if res[0].BalanceDescription == "0 XLM" {
			return false
		}
		if res[0].BalanceDescription == "49.9999800 XLM" {
			t.Logf("poll-1-%v: received T1 but not T2", i)
			return false
		}
		if res[0].BalanceDescription == "29.9999800 XLM" {
			t.Logf("poll-1-%v: received T2 but not T1", i)
			return false
		}
		t.Logf("poll-1-%v: received both payments", i)
		require.Equal(t, "79.9999700 XLM", res[0].BalanceDescription)
		return true
	})

	if skipPart2 {
		t.Logf("Skipping part 2")
		return
	}

	t.Logf("--------------------")
	t.Logf("Part 2: Alice sends a relay payment to bob who now already has a wallet")
	cmd = client.CmdWalletSend{
		Contextified: libkb.NewContextified(alice.tc.G),
		Recipient:    bob.username,
		Amount:       "10",
		ForceRelay:   true,
	}
	for i := 0; i < retryCount; i++ {
		err = cmd.Run()
		if err == nil {
			break
		}
	}
	require.NoError(t, err)

	pollFor(t, "final claim to complete", pollTime, bob.tc.G, func(i int) bool {
		res, err = bob.stellarClient.GetWalletAccountsLocal(context.Background(), 0)
		require.NoError(t, err)
		t.Logf("poll-2-%v: %v", i, res[0].BalanceDescription)
		if res[0].BalanceDescription == "79.9999700 XLM" {
			return false
		}
		t.Logf("poll-1-%v: received final payment", i)
		require.Equal(t, "89.9999600 XLM", res[0].BalanceDescription)
		return true
	})

}

func sampleNote() stellar1.NoteContents {
	return stellar1.NoteContents{
		Note:      "wizbang",
		StellarID: stellar1.TransactionID("6653fc2fdbc42ad51ccbe77ee0a3c29e258a5513c62fdc532cbfff91ab101abf"),
	}
}

// Friendbot sends someone XLM
func gift(t testing.TB, accountID stellar1.AccountID) {
	t.Logf("gift -> %v", accountID)
	url := "https://friendbot.stellar.org/?addr=" + accountID.String()
	t.Logf("gift url: %v", url)
	res, err := http.Get(url)
	require.NoError(t, err, "friendbot request error")
	bodyBuf := new(bytes.Buffer)
	bodyBuf.ReadFrom(res.Body)
	t.Logf("gift res: %v", bodyBuf.String())
	require.Equal(t, 200, res.StatusCode, "friendbot response status code")
}

func useStellarTestNet(t testing.TB) {
	stellarnet.SetClientAndNetwork(horizon.DefaultTestNetClient, build.TestNetwork)
}
<commit_msg>Retry friendbot gift (#13826)<commit_after>package systests

import (
	"bytes"
	"net/http"
	"testing"
	"time"

	"golang.org/x/net/context"

	"github.com/keybase/client/go/client"
	"github.com/keybase/client/go/libkb"
	keybase1 "github.com/keybase/client/go/protocol/keybase1"
	"github.com/keybase/client/go/protocol/stellar1"
	"github.com/keybase/client/go/stellar"
	"github.com/keybase/client/go/teams"
	"github.com/keybase/stellarnet"
	"github.com/stellar/go/build"
	"github.com/stellar/go/clients/horizon"
	"github.com/stretchr/testify/require"
)

func TestStellarNoteRoundtripAndResets(t *testing.T) {
	ctx := newSMUContext(t)
	defer ctx.cleanup()

	// Sign up two users, bob and alice.
	alice := ctx.installKeybaseForUser("alice", 10)
	alice.signup()
	divDebug(ctx, "Signed up alice (%s)", alice.username)
	bob := ctx.installKeybaseForUser("bob", 10)
	bob.signup()
	divDebug(ctx, "Signed up bob (%s)", bob.username)

	t.Logf("note to self")
	encB64, err := stellar.NoteEncryptB64(context.Background(), alice.getPrimaryGlobalContext(), sampleNote(), nil)
	require.NoError(t, err)
	note, err := stellar.NoteDecryptB64(context.Background(), alice.getPrimaryGlobalContext(), encB64)
	require.NoError(t, err)
	require.Equal(t, sampleNote(), note)

	t.Logf("note to both users")
	other := bob.userVersion()
	encB64, err = stellar.NoteEncryptB64(context.Background(), alice.getPrimaryGlobalContext(), sampleNote(), &other)
	require.NoError(t, err)

	t.Logf("decrypt as self")
	note, err = stellar.NoteDecryptB64(context.Background(), alice.getPrimaryGlobalContext(), encB64)
	require.NoError(t, err)
	require.Equal(t, sampleNote(), note)

	t.Logf("decrypt as other")
	note, err = stellar.NoteDecryptB64(context.Background(), bob.getPrimaryGlobalContext(), encB64)
	require.NoError(t, err)
	require.Equal(t, sampleNote(), note)

	t.Logf("reset sender")
	alice.reset()
	divDebug(ctx, "Reset bob (%s)", bob.username)
	alice.loginAfterReset(10)
	divDebug(ctx, "Bob logged in after reset")

	t.Logf("fail to decrypt as post-reset self")
	note, err = stellar.NoteDecryptB64(context.Background(), alice.getPrimaryGlobalContext(), encB64)
	require.Error(t, err)
	require.Equal(t, "note not encrypted for logged-in user", err.Error())

	t.Logf("decrypt as other")
	note, err = stellar.NoteDecryptB64(context.Background(), bob.getPrimaryGlobalContext(), encB64)
	require.NoError(t, err)
	require.Equal(t, sampleNote(), note)
}

// Test took 38s on a dev server 2018-06-07
func TestStellarRelayAutoClaims(t *testing.T) {
	testStellarRelayAutoClaims(t, false, false)
}

// Test took 29s on a dev server 2018-06-07
func TestStellarRelayAutoClaimsWithPUK(t *testing.T) {
	testStellarRelayAutoClaims(t, true, true)
}

// Part 1:
// XLM is sent to a user before they have a [PUK / wallet].
// In the form of multiple relay payments.
// They then [get a PUK,] add a wallet, and enter the impteam,
// which all kick the autoclaim into gear.
//
// Part 2:
// A relay payment is sent to the user who already has a wallet.
// The funds should be claimed asap.
//
// To debug this test use log filter "stellar_test|poll-|AutoClaim|stellar.claim|pollfor"
func testStellarRelayAutoClaims(t *testing.T, startWithPUK, skipPart2 bool) {
	tt := newTeamTester(t)
	defer tt.cleanup()
	useStellarTestNet(t)

	alice := tt.addUser("alice")
	var bob *userPlusDevice
	if startWithPUK {
		bob = tt.addWalletlessUser("bob")
	} else {
		bob = tt.addPuklessUser("bob")
	}
	alice.kickTeamRekeyd()

	t.Logf("alice gets funded")
	res, err := alice.stellarClient.GetWalletAccountsLocal(context.Background(), 0)
	require.NoError(t, err)
	gift(t, res[0].AccountID)

	t.Logf("alice sends a first relay payment to bob P1")
	attachIdentifyUI(t, alice.tc.G, newSimpleIdentifyUI())
	cmd := client.CmdWalletSend{
		Contextified: libkb.NewContextified(alice.tc.G),
		Recipient:    bob.username,
		Amount:       "50",
	}
	for i := 0; i < retryCount; i++ {
		err = cmd.Run()
		if err == nil {
			break
		}
	}
	require.NoError(t, err)

	t.Logf("alice sends a second relay payment to bob P2")
	cmd = client.CmdWalletSend{
		Contextified: libkb.NewContextified(alice.tc.G),
		Recipient:    bob.username,
		Amount:       "30",
	}
	for i := 0; i < retryCount; i++ {
		err = cmd.Run()
		if err == nil {
			break
		}
	}
	require.NoError(t, err)

	t.Logf("get the impteam seqno to wait on later")
	team, _, _, err := teams.LookupImplicitTeam(context.Background(), alice.tc.G, alice.username+","+bob.username, false)
	require.NoError(t, err)
	nextSeqno := team.NextSeqno()

	if startWithPUK {
		t.Logf("bob gets a wallet")
		bob.tc.Tp.DisableAutoWallet = false
		bob.tc.G.GetStellar().CreateWalletSoft(context.Background())
	} else {
		t.Logf("bob gets a PUK and wallet")
		bob.perUserKeyUpgrade()
		bob.tc.G.GetStellar().CreateWalletSoft(context.Background())

		t.Logf("wait for alice to add bob to their impteam")
		alice.pollForTeamSeqnoLinkWithLoadArgs(keybase1.LoadTeamArg{ID: team.ID}, nextSeqno)
	}

	pollTime := 20 * time.Second
	if libkb.UseCITime(bob.tc.G) {
		// This test is especially slow.
		pollTime = 30 * time.Second
	}

	pollFor(t, "claims to complete", pollTime, bob.tc.G, func(i int) bool {
		res, err = bob.stellarClient.GetWalletAccountsLocal(context.Background(), 0)
		require.NoError(t, err)
		t.Logf("poll-1-%v: %v", i, res[0].BalanceDescription)
		if res[0].BalanceDescription == "0 XLM" {
			return false
		}
		if res[0].BalanceDescription == "49.9999800 XLM" {
			t.Logf("poll-1-%v: received T1 but not T2", i)
			return false
		}
		if res[0].BalanceDescription == "29.9999800 XLM" {
			t.Logf("poll-1-%v: received T2 but not T1", i)
			return false
		}
		t.Logf("poll-1-%v: received both payments", i)
		require.Equal(t, "79.9999700 XLM", res[0].BalanceDescription)
		return true
	})

	if skipPart2 {
		t.Logf("Skipping part 2")
		return
	}

	t.Logf("--------------------")
	t.Logf("Part 2: Alice sends a relay payment to bob who now already has a wallet")
	cmd = client.CmdWalletSend{
		Contextified: libkb.NewContextified(alice.tc.G),
		Recipient:    bob.username,
		Amount:       "10",
		ForceRelay:   true,
	}
	for i := 0; i < retryCount; i++ {
		err = cmd.Run()
		if err == nil {
			break
		}
	}
	require.NoError(t, err)

	pollFor(t, "final claim to complete", pollTime, bob.tc.G, func(i int) bool {
		res, err = bob.stellarClient.GetWalletAccountsLocal(context.Background(), 0)
		require.NoError(t, err)
		t.Logf("poll-2-%v: %v", i, res[0].BalanceDescription)
		if res[0].BalanceDescription == "79.9999700 XLM" {
			return false
		}
		t.Logf("poll-1-%v: received final payment", i)
		require.Equal(t, "89.9999600 XLM", res[0].BalanceDescription)
		return true
	})

}

func sampleNote() stellar1.NoteContents {
	return stellar1.NoteContents{
		Note:      "wizbang",
		StellarID: stellar1.TransactionID("6653fc2fdbc42ad51ccbe77ee0a3c29e258a5513c62fdc532cbfff91ab101abf"),
	}
}

// Friendbot sends someone XLM
func gift(t testing.TB, accountID stellar1.AccountID) {
	t.Logf("gift -> %v", accountID)
	url := "https://friendbot.stellar.org/?addr=" + accountID.String()
	for i := 0; i < retryCount; i++ {
		t.Logf("gift url: %v", url)
		res, err := http.Get(url)
		if err != nil {
			t.Logf("http get %s error: %s", url, err)
			continue
		}
		bodyBuf := new(bytes.Buffer)
		bodyBuf.ReadFrom(res.Body)
		res.Body.Close()
		t.Logf("gift res: %v", bodyBuf.String())
		if res.StatusCode == 200 {
			return
		}
		t.Logf("gift status not ok: %d", res.StatusCode)
	}
	t.Fatalf("gift to %s failed after multiple attempts", accountID)
}

func useStellarTestNet(t testing.TB) {
	stellarnet.SetClientAndNetwork(horizon.DefaultTestNetClient, build.TestNetwork)
}
<|endoftext|>


<commit_before>package vtgate

import (
	"fmt"
	"sort"
	"strings"
	"sync"
	"time"

	log "github.com/golang/glog"
	pbq "github.com/youtube/vitess/go/vt/proto/query"
	pbt "github.com/youtube/vitess/go/vt/proto/topodata"
	"github.com/youtube/vitess/go/vt/tabletserver/tabletconn"
	"github.com/youtube/vitess/go/vt/topo"
	"golang.org/x/net/context"
)

// NewHealthCheck creates a new HealthCheck object.
func NewHealthCheck(connTimeout time.Duration, retryDelay time.Duration) *HealthCheck {
	hc := &HealthCheck{
		addrToConns: make(map[string]*healthCheckConn),
		targetToEPs: make(map[string]map[string][]*pbt.EndPoint),
		connTimeout: connTimeout,
		retryDelay:  retryDelay,
	}
	return hc
}

// HealthCheck performs health checking and notifies downstream components about any changes.
type HealthCheck struct {
	mu          sync.RWMutex                          // mu protects the two maps, not the parameters.
	addrToConns map[string]*healthCheckConn           // addrToConns maps from address to the healthCheckConn object.
	targetToEPs map[string]map[string][]*pbt.EndPoint // targetToEPs maps from keyspace/shard to a list of endpoints.
	connTimeout time.Duration
	retryDelay  time.Duration
}

// healthCheckConn contains details about an endpoint.
type healthCheckConn struct {
	mu         sync.RWMutex
	target     *pbq.Target
	cancelFunc context.CancelFunc
	conn       tabletconn.TabletConn
	stats      *pbq.RealtimeStats
}

// checkConn performs health checking on the given endpoint.
func (hc *HealthCheck) checkConn(endPoint *pbt.EndPoint) {
	ctx, cancelFunc := context.WithCancel(context.Background())
	hcc := &healthCheckConn{
		cancelFunc: cancelFunc,
	}

	// retry health check if it fails
	for {
		stream, errfunc, err := hcc.connect(ctx, hc, endPoint)
		if err != nil {
			select {
			case <-ctx.Done():
				return
			default:
			}
			log.Errorf("cannot connect to %+v: %v", endPoint, err)
			time.Sleep(hc.retryDelay)
			continue
		}
		for {
			err = hcc.processResponse(ctx, hc, endPoint, stream, errfunc)
			if err != nil {
				hcc.mu.Lock()
				hcc.conn.Close()
				hcc.conn = nil
				hcc.mu.Unlock()
				select {
				case <-ctx.Done():
					return
				default:
				}
				log.Errorf("error when streaming tablet health from %+v: %v", endPoint, err)
				time.Sleep(hc.retryDelay)
				break
			}
		}
	}
}

// func to connect to endpoint and start streaming.
// it makes sure we have a chance to signal the first connection attempt regardless whether it succeeds.
func (hcc *healthCheckConn) connect(ctx context.Context, hc *HealthCheck, endPoint *pbt.EndPoint) (<-chan *pbq.StreamHealthResponse, tabletconn.ErrFunc, error) {
	conn, err := tabletconn.GetDialer()(ctx, endPoint, "" /*keyspace*/, "" /*shard*/, pbt.TabletType_RDONLY, hc.connTimeout)
	if err != nil {
		return nil, nil, err
	}
	stream, errfunc, err := conn.StreamHealth(ctx)
	if err != nil {
		conn.Close()
		return nil, nil, err
	}
	err = hcc.processResponse(ctx, hc, endPoint, stream, errfunc)
	if err != nil {
		conn.Close()
		return nil, nil, err
	}
	hcc.mu.Lock()
	hcc.conn = conn
	hcc.mu.Unlock()
	return stream, errfunc, nil
}

// func to read one health check response, and notify downstream component
func (hcc *healthCheckConn) processResponse(ctx context.Context, hc *HealthCheck, endPoint *pbt.EndPoint, stream <-chan *pbq.StreamHealthResponse, errfunc tabletconn.ErrFunc) error {
	select {
	case <-ctx.Done():
		return ctx.Err()
	case shr, ok := <-stream:
		if !ok {
			return errfunc()
		}
		if shr.Target == nil || shr.RealtimeStats == nil || shr.RealtimeStats.HealthError != "" {
			return fmt.Errorf("health stats is not valid: %v", shr)
		}

		if hcc.target == nil {
			// The first time we see response for the endpoint.
			hcc.mu.Lock()
			hcc.target = shr.Target
			hcc.stats = shr.RealtimeStats
			hcc.mu.Unlock()
			hc.mu.Lock()
			key := hc.endPointToMapKey(endPoint)
			hc.addrToConns[key] = hcc
			hc.addEndPointToTargetProtected(hcc.target.Keyspace, hcc.target.Shard, endPoint)
			// TODO: notify downstream component for endpoint going up
			hc.mu.Unlock()
		} else {
			hcc.mu.Lock()
			hcc.target = shr.Target
			hcc.stats = shr.RealtimeStats
			hcc.mu.Unlock()
		}
		// TODO: notify downstream for tablettype and realtimestats
		return nil
	}
}

// AddEndPoint adds the endpoint, and starts health check.
// It notifies downstram components for endpoint going up after the first health check succeeds.
func (hc *HealthCheck) AddEndPoint(endPoint *pbt.EndPoint) {
	go hc.checkConn(endPoint)
}

// RemoveEndPoint removes the endpoint, and cancels the health check.
// It also notifies downstream components for endpoint going down.
func (hc *HealthCheck) RemoveEndPoint(endPoint *pbt.EndPoint) {
	hc.mu.Lock()
	defer hc.mu.Unlock()

	key := hc.endPointToMapKey(endPoint)
	hcc, ok := hc.addrToConns[key]
	if !ok {
		return
	}
	hcc.cancelFunc()
	delete(hc.addrToConns, key)
	if hcc.target != nil {
		hc.deleteEndPointFromTargetProtected(hcc.target.Keyspace, hcc.target.Shard, endPoint)
	}
	// TODO: notify downstream component for endpoint going down
}

// GetEndPointFromTarget returns all endpoints for the given keyspace/shard.
func (hc *HealthCheck) GetEndPointFromTarget(keyspace, shard string) []*pbt.EndPoint {
	hc.mu.RLock()
	defer hc.mu.RUnlock()
	shardMap, ok := hc.targetToEPs[keyspace]
	if !ok {
		return nil
	}
	epList, ok := shardMap[shard]
	if !ok {
		return nil
	}
	res := make([]*pbt.EndPoint, 0, 1)
	return append(res, epList...)
}

// addEndPointToTargetProtected adds the endpoint to the given target.
// LOCK_REQUIRED hc.mu
func (hc *HealthCheck) addEndPointToTargetProtected(keyspace, shard string, endPoint *pbt.EndPoint) {
	shardMap, ok := hc.targetToEPs[keyspace]
	if !ok {
		shardMap = make(map[string][]*pbt.EndPoint)
		hc.targetToEPs[keyspace] = shardMap
	}
	epList, ok := shardMap[shard]
	if !ok {
		epList = make([]*pbt.EndPoint, 0, 1)
	}
	for _, ep := range epList {
		if topo.EndPointEquality(ep, endPoint) {
			log.Warningf("endpoint is already up: %+v", endPoint)
			return
		}
	}
	shardMap[shard] = append(epList, endPoint)
}

// deleteEndPointFromTargetProtected deletes the endpoint for the given target.
// LOCK_REQUIRED hc.mu
func (hc *HealthCheck) deleteEndPointFromTargetProtected(keyspace, shard string, endPoint *pbt.EndPoint) {
	shardMap, ok := hc.targetToEPs[keyspace]
	if !ok {
		return
	}
	epList, ok := shardMap[shard]
	if !ok {
		return
	}
	for i, ep := range epList {
		if topo.EndPointEquality(ep, endPoint) {
			epList = append(epList[:i], epList[i+1:]...)
			shardMap[shard] = epList
			return
		}
	}
}

// endPointToMapKey creates a key to the map from endpoint's host and ports.
func (hc *HealthCheck) endPointToMapKey(endPoint *pbt.EndPoint) string {
	parts := make([]string, 0, 1)
	for name, port := range endPoint.PortMap {
		parts = append(parts, name+":"+string(port))
	}
	sort.Strings(parts)
	parts = append([]string{endPoint.Host}, parts...)
	return strings.Join(parts, ":")
}
<commit_msg>Complete the implementation.<commit_after>package vtgate

import (
	"fmt"
	"sort"
	"strings"
	"sync"
	"time"

	log "github.com/golang/glog"
	pbq "github.com/youtube/vitess/go/vt/proto/query"
	pbt "github.com/youtube/vitess/go/vt/proto/topodata"
	"github.com/youtube/vitess/go/vt/tabletserver/tabletconn"
	"github.com/youtube/vitess/go/vt/topo"
	"golang.org/x/net/context"
)

// HealthCheckStatsListener is the listener to receive health check stats update.
type HealthCheckStatsListener interface {
	StatsUpdate(endPoint *pbt.EndPoint, cell string, target *pbq.Target, stats *pbq.RealtimeStats)
}

// NewHealthCheck creates a new HealthCheck object.
func NewHealthCheck(listener HealthCheckStatsListener, connTimeout time.Duration, retryDelay time.Duration) *HealthCheck {
	return &HealthCheck{
		addrToConns: make(map[string]*healthCheckConn),
		targetToEPs: make(map[string]map[string]map[pbt.TabletType][]*pbt.EndPoint),
		listener:    listener,
		connTimeout: connTimeout,
		retryDelay:  retryDelay,
	}
}

// HealthCheck performs health checking and notifies downstream components about any changes.
type HealthCheck struct {
	// set at construction time
	listener    HealthCheckStatsListener
	connTimeout time.Duration
	retryDelay  time.Duration

	// mu protects all the following fields
	mu          sync.RWMutex
	addrToConns map[string]*healthCheckConn                              // addrToConns maps from address to the healthCheckConn object.
	targetToEPs map[string]map[string]map[pbt.TabletType][]*pbt.EndPoint // targetToEPs maps from keyspace/shard/tablettype to a list of endpoints.
}

// healthCheckConn contains details about an endpoint.
type healthCheckConn struct {
	// set at construction time
	cell       string
	cancelFunc context.CancelFunc

	// mu protects all the following fields
	mu     sync.RWMutex
	conn   tabletconn.TabletConn
	target *pbq.Target
	stats  *pbq.RealtimeStats
}

// checkConn performs health checking on the given endpoint.
func (hc *HealthCheck) checkConn(cell string, endPoint *pbt.EndPoint) {
	ctx, cancelFunc := context.WithCancel(context.Background())
	hcc := &healthCheckConn{
		cell:       cell,
		cancelFunc: cancelFunc,
	}

	// retry health check if it fails
	for {
		stream, errfunc, err := hcc.connect(ctx, hc, endPoint)
		if err != nil {
			select {
			case <-ctx.Done():
				return
			default:
			}
			log.Errorf("cannot connect to %+v: %v", endPoint, err)
			time.Sleep(hc.retryDelay)
			continue
		}
		for {
			err = hcc.processResponse(ctx, hc, endPoint, stream, errfunc)
			if err != nil {
				hcc.mu.Lock()
				hcc.conn.Close()
				hcc.conn = nil
				hcc.mu.Unlock()
				select {
				case <-ctx.Done():
					return
				default:
				}
				log.Errorf("error when streaming tablet health from %+v: %v", endPoint, err)
				time.Sleep(hc.retryDelay)
				break
			}
		}
	}
}

// connect creates connection to the endpoint and starts streaming.
func (hcc *healthCheckConn) connect(ctx context.Context, hc *HealthCheck, endPoint *pbt.EndPoint) (<-chan *pbq.StreamHealthResponse, tabletconn.ErrFunc, error) {
	conn, err := tabletconn.GetDialer()(ctx, endPoint, "" /*keyspace*/, "" /*shard*/, pbt.TabletType_RDONLY, hc.connTimeout)
	if err != nil {
		return nil, nil, err
	}
	stream, errfunc, err := conn.StreamHealth(ctx)
	if err != nil {
		conn.Close()
		return nil, nil, err
	}
	hcc.mu.Lock()
	hcc.conn = conn
	hcc.mu.Unlock()
	return stream, errfunc, nil
}

// processResponse reads one health check response, and notifies HealthCheckStatsListener.
func (hcc *healthCheckConn) processResponse(ctx context.Context, hc *HealthCheck, endPoint *pbt.EndPoint, stream <-chan *pbq.StreamHealthResponse, errfunc tabletconn.ErrFunc) error {
	select {
	case <-ctx.Done():
		return ctx.Err()
	case shr, ok := <-stream:
		if !ok {
			return errfunc()
		}
		if shr.Target == nil || shr.RealtimeStats == nil || shr.RealtimeStats.HealthError != "" {
			return fmt.Errorf("health stats is not valid: %v", shr)
		}

		if hcc.target == nil {
			// The first time we see response for the endpoint.
			hcc.mu.Lock()
			hcc.target = shr.Target
			hcc.stats = shr.RealtimeStats
			hcc.mu.Unlock()
			hc.mu.Lock()
			key := hc.endPointToMapKey(endPoint)
			hc.addrToConns[key] = hcc
			hc.addEndPointToTargetProtected(hcc.target, endPoint)
			hc.mu.Unlock()
		} else if hcc.target.TabletType != shr.Target.TabletType {
			hc.mu.Lock()
			hc.deleteEndPointFromTargetProtected(hcc.target, endPoint)
			hcc.mu.Lock()
			hcc.target = shr.Target
			hcc.stats = shr.RealtimeStats
			hcc.mu.Unlock()
			hc.addEndPointToTargetProtected(shr.Target, endPoint)
			hc.mu.Unlock()
		} else {
			hcc.mu.Lock()
			hcc.target = shr.Target
			hcc.stats = shr.RealtimeStats
			hcc.mu.Unlock()
		}
		// notify downstream for tablettype and realtimestats change
		hc.listener.StatsUpdate(endPoint, hcc.cell, hcc.target, hcc.stats)
		return nil
	}
}

// AddEndPoint adds the endpoint, and starts health check.
func (hc *HealthCheck) AddEndPoint(cell string, endPoint *pbt.EndPoint) {
	go hc.checkConn(cell, endPoint)
}

// RemoveEndPoint removes the endpoint, and stops the health check.
func (hc *HealthCheck) RemoveEndPoint(endPoint *pbt.EndPoint) {
	hc.mu.Lock()
	defer hc.mu.Unlock()

	key := hc.endPointToMapKey(endPoint)
	hcc, ok := hc.addrToConns[key]
	if !ok {
		return
	}
	hcc.cancelFunc()
	delete(hc.addrToConns, key)
	if hcc.target != nil {
		hc.deleteEndPointFromTargetProtected(hcc.target, endPoint)
	}
}

// GetEndPointsFromKeyspaceShard returns all endpoints for the given keyspace/shard.
func (hc *HealthCheck) GetEndPointsFromKeyspaceShard(keyspace, shard string) []*pbt.EndPoint {
	hc.mu.RLock()
	defer hc.mu.RUnlock()
	shardMap, ok := hc.targetToEPs[keyspace]
	if !ok {
		return nil
	}
	ttMap, ok := shardMap[shard]
	if !ok {
		return nil
	}
	res := make([]*pbt.EndPoint, 0, 1)
	for _, epList := range ttMap {
		res = append(res, epList...)
	}
	return res
}

// GetEndPointsFromTarget returns all endpoints for the given target.
func (hc *HealthCheck) GetEndPointsFromTarget(target *pbq.Target) []*pbt.EndPoint {
	hc.mu.RLock()
	defer hc.mu.RUnlock()
	shardMap, ok := hc.targetToEPs[target.Keyspace]
	if !ok {
		return nil
	}
	ttMap, ok := shardMap[target.Shard]
	if !ok {
		return nil
	}
	epList, ok := ttMap[target.TabletType]
	if !ok {
		return nil
	}
	res := make([]*pbt.EndPoint, 0, 1)
	return append(res, epList...)
}

// addEndPointToTargetProtected adds the endpoint to the given target.
// LOCK_REQUIRED hc.mu
func (hc *HealthCheck) addEndPointToTargetProtected(target *pbq.Target, endPoint *pbt.EndPoint) {
	shardMap, ok := hc.targetToEPs[target.Keyspace]
	if !ok {
		shardMap = make(map[string]map[pbt.TabletType][]*pbt.EndPoint)
		hc.targetToEPs[target.Keyspace] = shardMap
	}
	ttMap, ok := shardMap[target.Shard]
	if !ok {
		ttMap = make(map[pbt.TabletType][]*pbt.EndPoint)
		shardMap[target.Shard] = ttMap
	}
	epList, ok := ttMap[target.TabletType]
	if !ok {
		epList = make([]*pbt.EndPoint, 0, 1)
	}
	for _, ep := range epList {
		if topo.EndPointEquality(ep, endPoint) {
			log.Warningf("endpoint is already added: %+v", endPoint)
			return
		}
	}
	ttMap[target.TabletType] = append(epList, endPoint)
}

// deleteEndPointFromTargetProtected deletes the endpoint for the given target.
// LOCK_REQUIRED hc.mu
func (hc *HealthCheck) deleteEndPointFromTargetProtected(target *pbq.Target, endPoint *pbt.EndPoint) {
	shardMap, ok := hc.targetToEPs[target.Keyspace]
	if !ok {
		return
	}
	ttMap, ok := shardMap[target.Shard]
	if !ok {
		return
	}
	epList, ok := ttMap[target.TabletType]
	if !ok {
		return
	}
	for i, ep := range epList {
		if topo.EndPointEquality(ep, endPoint) {
			epList = append(epList[:i], epList[i+1:]...)
			ttMap[target.TabletType] = epList
			return
		}
	}
}

// endPointToMapKey creates a key to the map from endpoint's host and ports.
func (hc *HealthCheck) endPointToMapKey(endPoint *pbt.EndPoint) string {
	parts := make([]string, 0, 1)
	for name, port := range endPoint.PortMap {
		parts = append(parts, name+":"+string(port))
	}
	sort.Strings(parts)
	parts = append([]string{endPoint.Host}, parts...)
	return strings.Join(parts, ":")
}
<|endoftext|>


<commit_before>package ethutil

import (
	"math/big"
)

var BigInt0 *big.Int = big.NewInt(0)

// True
var BigTrue *big.Int = big.NewInt(1)

// False
var BigFalse *big.Int = big.NewInt(0)

// Big pow
//
// Returns the power of two big integers
func BigPow(a, b int) *big.Int {
	c := new(big.Int)
	c.Exp(big.NewInt(int64(a)), big.NewInt(int64(b)), big.NewInt(0))

	return c
}

// Big
//
// Shortcut for new(big.Int).SetString(..., 0)
func Big(num string) *big.Int {
	n := new(big.Int)
	n.SetString(num, 0)

	return n
}

// BigD
//
// Shortcut for new(big.Int).SetBytes(...)
func BigD(data []byte) *big.Int {
	n := new(big.Int)
	n.SetBytes(data)

	return n
}

// Big to bytes
//
// Returns the bytes of a big integer with the size specified by **base**
// Attempts to pad the byte array with zeros.
func BigToBytes(num *big.Int, base int) []byte {
	ret := make([]byte, base/8)

	if len(num.Bytes()) > base/8 {
		return num.Bytes()
	}

	return append(ret[:len(ret)-len(num.Bytes())], num.Bytes()...)
}

// Big copy
//
// Creates a copy of the given big integer
func BigCopy(src *big.Int) *big.Int {
	return new(big.Int).Set(src)
}

// Big max
//
// Returns the maximum size big integer
func BigMax(x, y *big.Int) *big.Int {
	if x.Cmp(y) <= 0 {
		return x
	}

	return y
}
<commit_msg>Fix BigMax to return the biggest number, not the smallest<commit_after>package ethutil

import (
	"math/big"
)

var BigInt0 *big.Int = big.NewInt(0)

// True
var BigTrue *big.Int = big.NewInt(1)

// False
var BigFalse *big.Int = big.NewInt(0)

// Big pow
//
// Returns the power of two big integers
func BigPow(a, b int) *big.Int {
	c := new(big.Int)
	c.Exp(big.NewInt(int64(a)), big.NewInt(int64(b)), big.NewInt(0))

	return c
}

// Big
//
// Shortcut for new(big.Int).SetString(..., 0)
func Big(num string) *big.Int {
	n := new(big.Int)
	n.SetString(num, 0)

	return n
}

// BigD
//
// Shortcut for new(big.Int).SetBytes(...)
func BigD(data []byte) *big.Int {
	n := new(big.Int)
	n.SetBytes(data)

	return n
}

// Big to bytes
//
// Returns the bytes of a big integer with the size specified by **base**
// Attempts to pad the byte array with zeros.
func BigToBytes(num *big.Int, base int) []byte {
	ret := make([]byte, base/8)

	if len(num.Bytes()) > base/8 {
		return num.Bytes()
	}

	return append(ret[:len(ret)-len(num.Bytes())], num.Bytes()...)
}

// Big copy
//
// Creates a copy of the given big integer
func BigCopy(src *big.Int) *big.Int {
	return new(big.Int).Set(src)
}

// Big max
//
// Returns the maximum size big integer
func BigMax(x, y *big.Int) *big.Int {
	if x.Cmp(y) <= 0 {
		return y
	}

	return x
}
<|endoftext|>


<commit_before>package gocqrs

import (
	"log"
	"strings"
)

const (
	Created = "Created"
	Updated = "Updated"
	Deleted = "Deleted"
)

var eventsNames = []string{"Created",
	"Updated",
	"Deleted",
}

type CRUDHandler struct {
	EntityName string `json:"entityName"`
}

func NewCRUDHandler(name string) CRUDHandler {
	var ch CRUDHandler
	if name == "" {
		log.Fatal("Invalid entity name to create CRUD handler")
		return ch
	}

	ch.EntityName = name
	return ch
}

// Handler CRUD events
func (ch CRUDHandler) Handle(ev Eventer, en *Entity) (StoreOptions, error) {
	var opt StoreOptions
	var err error

	switch ev.GetType() {
	case ch.CreateEvent():
		en.Version = 0
		en.ID = ev.GetId()
		en.Data = ev.GetData()
		opt.Create = true
	case ch.UpdateEvent():
		data := ev.GetData()
		log.Println(data)
		for k, d := range data {
			en.Data[k] = d
		}
	case ch.DeletedEvent():
	}
	return opt, err
}

func (ch CRUDHandler) EventName() []string {
	events := make([]string, 0)
	for _, p := range eventsNames {
		e := strings.Title(ch.EntityName) + p
		events = append(events, e)
	}

	return events
}

func (ch CRUDHandler) CreateEvent() string {
	return strings.Title(ch.EntityName) + "Created"
}

func (ch CRUDHandler) UpdateEvent() string {
	return strings.Title(ch.EntityName) + "Updated"
}

func (ch CRUDHandler) DeletedEvent() string {
	return strings.Title(ch.EntityName) + "Deleted"
}
<commit_msg>undelete event<commit_after>package gocqrs

import (
	"errors"
	"log"
	"strings"
)

const (
	Created = "Created"
	Updated = "Updated"
	Deleted = "Deleted"
)

var eventsNames = []string{"Created",
	"Updated",
	"Deleted",
	"Undeleted",
}

var (
	EntityDeleted = errors.New("Entity deleted, new to undeletet to update")
)

type CRUDHandler struct {
	EntityName string `json:"entityName"`
}

func NewCRUDHandler(name string) CRUDHandler {
	var ch CRUDHandler
	if name == "" {
		log.Fatal("Invalid entity name to create CRUD handler")
		return ch
	}

	ch.EntityName = name
	return ch
}

// Handler CRUD events
func (ch CRUDHandler) Handle(ev Eventer, en *Entity) (StoreOptions, error) {
	var opt StoreOptions
	var err error

	switch ev.GetType() {
	case ch.CreateEvent():
		if en.Deleted {
			return opt, EntityDeleted
		}
		en.Version = 0
		en.ID = ev.GetId()
		en.Data = ev.GetData()
		opt.Create = true
	case ch.UpdateEvent():
		if en.Deleted {
			return opt, EntityDeleted
		}
		data := ev.GetData()
		for k, d := range data {
			en.Data[k] = d
		}
	case ch.DeletedEvent():
		en.Deleted = true
	case ch.UnDeletedEvent():
		en.Deleted = false
	}
	return opt, err
}

func (ch CRUDHandler) EventName() []string {
	events := make([]string, 0)
	for _, p := range eventsNames {
		e := strings.Title(ch.EntityName) + p
		events = append(events, e)
	}

	return events
}

func (ch CRUDHandler) CreateEvent() string {
	return strings.Title(ch.EntityName) + "Created"
}

func (ch CRUDHandler) UpdateEvent() string {
	return strings.Title(ch.EntityName) + "Updated"
}

func (ch CRUDHandler) DeletedEvent() string {
	return strings.Title(ch.EntityName) + "Deleted"
}

func (ch CRUDHandler) UnDeletedEvent() string {
	return strings.Title(ch.EntityName) + "UnDeleted"
}
<|endoftext|>


<commit_before>// Package csrf generates and validates csrf tokens for martini.
// There are multiple methods of delivery including via a cookie or HTTP
// header.
// Validation occurs via a traditional hidden form key of "_csrf", or via
// a custom HTTP haeder "X-CSRFToken".
//
// package main
//
// import (
//     "github.com/codegangsta/martini"
//     "github.com/martini-contib/csrf"
//     "github.com/martini-contrib/render"
//     "github.com/martini-contib/sessions"
//     "net/http"dd
// )
//
// func main() {
//     m := martini.Classic()
//     store := sessions.NewCookieStore([]byte("secret123"))
//     m.Use(sessions.Sessions("my_session", store)
//     // Setup generation middleware.
//     m.Use(csrf.Generate(&csrf.Options{
//         Secret:     "token123",
//         SessionKey: "userId",
//     }))
//     m.Use(render.Renderer())
//
//     // Simulate the authentication of a session. If userId exists redirect
//     // to a form that requires csrf protection.
//     m.Get("/", func(s sessions.Session, r render.Render, x csrf.Csrf) {
//         if s.Get("userId") == nil {
//             r.Redirect("/login", 302)
//             return
//         }
//         r.Redirect("/protected", 302)
//     })
//
//     // Set userId for the session.
//     m.Get("/login", func(s sessions.Session, r render.Render) {
//         s.Set("userId", "123456")
//         r.Redirect("/", 302)
//     })
//
//     // Render a protected form. Passing a csrf token by calling x.GetToken()
//     m.Get("/protected", func(s sessions.Session, r render.Render, x csrf.Csrf) {
//         r.HTML(200, "protected", x.GetToken())
//     })
//
//     // Apply csrf validation to route.
//     m.Post("/protected", csrf.Validate, func(s sessions.Session, r render.Render) {
//         if u := s.Get("userId"); u != nil {
//             r.HTML(200, "result", "You submitted a valid token")
//             return
//         }
//         r.Redirect("/login", 401)
//     })
//
//     m.Run()
// }
package csrf

import (
	"code.google.com/p/xsrftoken"
	"fmt"
	"github.com/codegangsta/martini"
	"github.com/martini-contrib/sessions"
	"net/http"
	"regexp"
	"strings"
	"time"
)

// Csrf is used to get the current token and validate a suspect token.
type Csrf interface {
	// Return the token.
	GetToken() string
	// Validate by token.
	ValidToken(t string) bool
}

type csrf struct {
	// Token generated to pass via header, cookie, or hidden form value.
	Token string
	// This value must be unique per user.
	Id string
	// Secret used along with the unique id above to generate the Token.
	Secret string
}

// Returns the current token. This is typically used
// to populate a hidden form in an HTML template.
func (c *csrf) GetToken() string {
	return c.Token
}

// Validates the passed token against the existing Secret and Id.
func (c *csrf) ValidToken(t string) bool {
	return xsrftoken.Valid(t, c.Secret, c.Id, "POST")
}

// Maintains options to manage behavior of Generate.
type Options struct {
	// The global secret value used to generate Tokens.
	Secret string
	// Key used for getting the unique Id per user.
	SessionKey string
	// If true, send token via X-CSRFToken header.
	SetHeader bool
	// If true, send token via _csrf cookie.
	SetCookie bool
	// Set the Secure flag to true on the cookie.
	Secure bool
}

const domainReg = `/^\.?[a-z\d]+(?:(?:[a-z\d]*)|(?:[a-z\d\-]*[a-z\d]))(?:\.[a-z\d]+(?:(?:[a-z\d]*)|(?:[a-z\d\-]*[a-z\d])))*$/`

// Maps Csrf to each request. If this request is a Get request, it will generate a new token.
// Additionally, depending on options set, generated tokens will be sent via Header and/or Cookie.
func Generate(opts *Options) martini.Handler {
	return func(s sessions.Session, c martini.Context, r *http.Request, w http.ResponseWriter) {
		x := &csrf{Secret: opts.Secret}
		c.MapTo(x, (*Csrf)(nil))
		uid := s.Get(opts.SessionKey)
		if uid == nil {
			return
		}
		switch uid.(type) {
		case string:
			x.Id = uid.(string)
		default:
			return
		}
		// Don't set cookie or send header if this is not a get request
		// or was sen't via an api request.
		if r.Method == "GET" && r.Header.Get("X-API-Key") == "" {
			// If cookie present, map existing token, else generate a new one.
			if ex, err := r.Cookie("_csrf"); err == nil && ex.Value != "" {
				x.Token = ex.Value
			} else {
				x.Token = xsrftoken.Generate(x.Secret, x.Id, "POST")
				if opts.SetCookie {
					expire := time.Now().AddDate(0, 0, 1)
					// Verify the domain is valid. If it is not, set as empty.
					domain := strings.Split(r.Host, ":")[0]
					if ok, err := regexp.Match(domainReg, []byte(domain)); !ok || err != nil {
						domain = ""
					}
					cookie := &http.Cookie{
						Name:       "_csrf",
						Value:      x.Token,
						Path:       "/",
						Domain:     domain,
						Expires:    expire,
						RawExpires: expire.Format(time.UnixDate),
						MaxAge:     0,
						Secure:     opts.Secure,
						HttpOnly:   false,
						Raw:        fmt.Sprintf("_csrf=%s", x.Token),
						Unparsed:   []string{fmt.Sprintf("token=%s", x.Token)},
					}
					http.SetCookie(w, cookie)
				}
			}
			if opts.SetHeader {
				w.Header().Add("X-CSRFToken", x.Token)
			}
		}
	}
}

// Validate should be used as a per route middleware. It attempts to get a token from a "X-CSRFToken"
// HTTP header and then a "_csrf" form value. If one of these is found, the token will be validated
// using ValidToken. If this validation fails, http.StatusBadRequest is sent in the reply.
// If neither a header or form value is faound, http.StatusBadRequest is sent.
func Validate(r *http.Request, w http.ResponseWriter, x Csrf) {
	if token := r.Header.Get("X-CSRFToken"); token != "" {
		if !x.ValidToken(token) {
			http.Error(w, "Invalid X-CSRFToken", http.StatusBadRequest)
		}
		return
	}
	if token := r.FormValue("_csrf"); token != "" {
		if !x.ValidToken(token) {
			http.Error(w, "Invalid _csrf token", http.StatusBadRequest)
		}
		return
	}
	http.Error(w, "Bad Request", http.StatusBadRequest)
	return
}
<commit_msg>Fix package comments.<commit_after>// Package csrf generates and validates csrf tokens for martini.
// There are multiple methods of delivery including via a cookie or HTTP
// header.
// Validation occurs via a traditional hidden form key of "_csrf", or via
// a custom HTTP haeder "X-CSRFToken".
//
// package main
//
// import (
//     "github.com/codegangsta/martini"
//     "github.com/martini-contib/csrf"
//     "github.com/martini-contrib/render"
//     "github.com/martini-contib/sessions"
//     "net/http"dd
// )
//
// func main() {
//     m := martini.Classic()
//     store := sessions.NewCookieStore([]byte("secret123"))
//     m.Use(sessions.Sessions("my_session", store))
//     // Setup generation middleware.
//     m.Use(csrf.Generate(&csrf.Options{
//         Secret:     "token123",
//         SessionKey: "userId",
//     }))
//     m.Use(render.Renderer())
//
//     // Simulate the authentication of a session. If userId exists redirect
//     // to a form that requires csrf protection.
//     m.Get("/", func(s sessions.Session, r render.Render, x csrf.Csrf) {
//         if s.Get("userId") == nil {
//             r.Redirect("/login", 302)
//             return
//         }
//         r.Redirect("/protected", 302)
//     })
//
//     // Set userId for the session.
//     m.Get("/login", func(s sessions.Session, r render.Render) {
//         s.Set("userId", "123456")
//         r.Redirect("/", 302)
//     })
//
//     // Render a protected form. Passing a csrf token by calling x.GetToken()
//     m.Get("/protected", func(s sessions.Session, r render.Render, x csrf.Csrf) {
//         r.HTML(200, "protected", x.GetToken())
//     })
//
//     // Apply csrf validation to route.
//     m.Post("/protected", csrf.Validate, func(s sessions.Session, r render.Render) {
//         if u := s.Get("userId"); u != nil {
//             r.HTML(200, "result", "You submitted a valid token")
//             return
//         }
//         r.Redirect("/login", 401)
//     })
//
//     m.Run()
// }
package csrf

import (
	"code.google.com/p/xsrftoken"
	"fmt"
	"github.com/codegangsta/martini"
	"github.com/martini-contrib/sessions"
	"net/http"
	"regexp"
	"strings"
	"time"
)

// Csrf is used to get the current token and validate a suspect token.
type Csrf interface {
	// Return the token.
	GetToken() string
	// Validate by token.
	ValidToken(t string) bool
}

type csrf struct {
	// Token generated to pass via header, cookie, or hidden form value.
	Token string
	// This value must be unique per user.
	Id string
	// Secret used along with the unique id above to generate the Token.
	Secret string
}

// Returns the current token. This is typically used
// to populate a hidden form in an HTML template.
func (c *csrf) GetToken() string {
	return c.Token
}

// Validates the passed token against the existing Secret and Id.
func (c *csrf) ValidToken(t string) bool {
	return xsrftoken.Valid(t, c.Secret, c.Id, "POST")
}

// Maintains options to manage behavior of Generate.
type Options struct {
	// The global secret value used to generate Tokens.
	Secret string
	// Key used for getting the unique Id per user.
	SessionKey string
	// If true, send token via X-CSRFToken header.
	SetHeader bool
	// If true, send token via _csrf cookie.
	SetCookie bool
	// Set the Secure flag to true on the cookie.
	Secure bool
}

const domainReg = `/^\.?[a-z\d]+(?:(?:[a-z\d]*)|(?:[a-z\d\-]*[a-z\d]))(?:\.[a-z\d]+(?:(?:[a-z\d]*)|(?:[a-z\d\-]*[a-z\d])))*$/`

// Maps Csrf to each request. If this request is a Get request, it will generate a new token.
// Additionally, depending on options set, generated tokens will be sent via Header and/or Cookie.
func Generate(opts *Options) martini.Handler {
	return func(s sessions.Session, c martini.Context, r *http.Request, w http.ResponseWriter) {
		x := &csrf{Secret: opts.Secret}
		c.MapTo(x, (*Csrf)(nil))
		uid := s.Get(opts.SessionKey)
		if uid == nil {
			return
		}
		switch uid.(type) {
		case string:
			x.Id = uid.(string)
		default:
			return
		}
		// Don't set cookie or send header if this is not a get request
		// or was sen't via an api request.
		if r.Method == "GET" && r.Header.Get("X-API-Key") == "" {
			// If cookie present, map existing token, else generate a new one.
			if ex, err := r.Cookie("_csrf"); err == nil && ex.Value != "" {
				x.Token = ex.Value
			} else {
				x.Token = xsrftoken.Generate(x.Secret, x.Id, "POST")
				if opts.SetCookie {
					expire := time.Now().AddDate(0, 0, 1)
					// Verify the domain is valid. If it is not, set as empty.
					domain := strings.Split(r.Host, ":")[0]
					if ok, err := regexp.Match(domainReg, []byte(domain)); !ok || err != nil {
						domain = ""
					}
					cookie := &http.Cookie{
						Name:       "_csrf",
						Value:      x.Token,
						Path:       "/",
						Domain:     domain,
						Expires:    expire,
						RawExpires: expire.Format(time.UnixDate),
						MaxAge:     0,
						Secure:     opts.Secure,
						HttpOnly:   false,
						Raw:        fmt.Sprintf("_csrf=%s", x.Token),
						Unparsed:   []string{fmt.Sprintf("token=%s", x.Token)},
					}
					http.SetCookie(w, cookie)
				}
			}
			if opts.SetHeader {
				w.Header().Add("X-CSRFToken", x.Token)
			}
		}
	}
}

// Validate should be used as a per route middleware. It attempts to get a token from a "X-CSRFToken"
// HTTP header and then a "_csrf" form value. If one of these is found, the token will be validated
// using ValidToken. If this validation fails, http.StatusBadRequest is sent in the reply.
// If neither a header or form value is faound, http.StatusBadRequest is sent.
func Validate(r *http.Request, w http.ResponseWriter, x Csrf) {
	if token := r.Header.Get("X-CSRFToken"); token != "" {
		if !x.ValidToken(token) {
			http.Error(w, "Invalid X-CSRFToken", http.StatusBadRequest)
		}
		return
	}
	if token := r.FormValue("_csrf"); token != "" {
		if !x.ValidToken(token) {
			http.Error(w, "Invalid _csrf token", http.StatusBadRequest)
		}
		return
	}
	http.Error(w, "Bad Request", http.StatusBadRequest)
	return
}
<|endoftext|>


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

package sdb

import (
	"fmt"
	"github.com/jacobsa/aws/exp/sdb/conn"
)

type batchPutPair struct {
	Item ItemName
	Updates []PutUpdate
}

func getSortedPairs(updateMap map[ItemName][]PutUpdate) []batchPutPair

func validateUpdate(u PutUpdate) (err error) {
	// Make sure the attribute name is legal.
	if u.Name == "" {
		return fmt.Errorf("Invalid attribute name; names must be non-empty.")
	}

	if err = validateValue(string(u.Name)); err != nil {
		return fmt.Errorf("Invalid attribute name: %v", err)
	}

	// Make sure the attribute value is legal.
	if err = validateValue(string(u.Value)); err != nil {
		return fmt.Errorf("Invalid attribute value: %v", err)
	}

	return nil
}

func validateUpdates(updates []PutUpdate) (err error) {
	numUpdates := len(updates)
	if numUpdates == 0 || numUpdates > 256 {
		return fmt.Errorf("Illegal number of updates: %d", numUpdates)
	}

	for _, u := range updates {
		if err = validateUpdate(u); err != nil {
			return fmt.Errorf("Invalid update (%v): %v", err, u)
		}
	}

	return nil
}

func validatePrecondition(p Precondition) (err error) {
	// Make sure the attribute name is legal.
	if p.Name == "" {
		return fmt.Errorf("Invalid attribute name; names must be non-empty.")
	}

	if err = validateValue(string(p.Name)); err != nil {
		return fmt.Errorf("Invalid attribute name: %v", err)
	}

	// We require exactly one operand.
	if (p.Value == nil) == (p.Exists == nil) {
		return fmt.Errorf("Preconditions must contain exactly one of Value and Exists.")
	}

	// Make sure the attribute value is legal, if present.
	if p.Value != nil {
		if err = validateValue(string(*p.Value)); err != nil {
			return fmt.Errorf("Invalid attribute value: %v", err)
		}
	}

	return nil
}

func (d *domain) PutAttributes(
	item ItemName,
	updates []PutUpdate,
	preconditions []Precondition) (err error) {
	// Make sure the item name is legal.
	if item == "" {
		return fmt.Errorf("Invalid item name; names must be non-empty.")
	}

	if err = validateValue(string(item)); err != nil {
		return fmt.Errorf("Invalid item name: %v", err)
	}

	// Validate updates.
	if err = validateUpdates(updates); err != nil {
		return err
	}

	// Validate preconditions.
	for _, p := range preconditions {
		if err = validatePrecondition(p); err != nil {
			return fmt.Errorf("Invalid precondition (%v): %v", err, p)
		}
	}

	// Assemble an appropriate request.
	req := conn.Request{}
	req["DomainName"] = d.name
	req["ItemName"] = string(item)

	for i, u := range updates {
		keyPrefix := fmt.Sprintf("Attribute.%d.", i+1)
		req[keyPrefix + "Name"] = u.Name
		req[keyPrefix + "Value"] = u.Value

		if u.Replace {
			req[keyPrefix + "Replace"] = "true"
		}
	}

	for i, p := range preconditions {
		keyPrefix := fmt.Sprintf("Expected.%d.", i+1)
		req[keyPrefix + "Name"] = p.Name

		if p.Value != nil {
			req[keyPrefix + "Value"] = *p.Value
		} else if *p.Exists {
			req[keyPrefix + "Exists"] = "true"
		} else {
			req[keyPrefix + "Exists"] = "false"
		}
	}

	// Call the connection.
	if _, err = d.c.SendRequest(req); err != nil {
		return fmt.Errorf("SendRequest: %v", err)
	}

	return nil
}

func (d *domain) BatchPutAttributes(updateMap map[ItemName][]PutUpdate) (err error) {
	// Make sure the size of the request is legal.
	numItems := len(updateMap)
	if numItems == 0 || numItems > 25 {
		return fmt.Errorf("Illegal number of items: %d", numItems)
	}

	// Make sure each item name and set of updates is legal.
	for item, updates := range updateMap {
		if item == "" {
			return fmt.Errorf("Invalid item name; names must be non-empty.")
		}

		if err = validateValue(string(item)); err != nil {
			return fmt.Errorf("Invalid item name: %v", err)
		}

		if err = validateUpdates(updates); err != nil {
			return fmt.Errorf("Updates for item %s: %v", item, err)
		}
	}

	// Build a request.
	req := conn.Request{}
	req["DomainName"] = d.name

	pairs := getSortedPairs(updateMap)
	for i, pair := range pairs {
		itemPrefix := fmt.Sprintf("Item.%d.", i+1)
		req[itemPrefix + "ItemName"] = string(pair.Item)

		for j, u := range pair.Updates {
			updatePrefix := fmt.Sprintf("%sAttribute.%d.", itemPrefix, j+1)
			req[updatePrefix + "Name"] = u.Name
			req[updatePrefix + "Value"] = u.Value

			if u.Replace {
				req[updatePrefix + "Replace"] = "true"
			}
		}
	}

	// Call the connection.
	if _, err = d.c.SendRequest(req); err != nil {
		return fmt.Errorf("SendRequest: %v", err)
	}

	return nil
}
<commit_msg>Fixed a bug.<commit_after>// Copyright 2012 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package sdb

import (
	"fmt"
	"github.com/jacobsa/aws/exp/sdb/conn"
	"sort"
)

type batchPutPair struct {
	Item ItemName
	Updates []PutUpdate
}

type batchPutPairList []batchPutPair

func (l batchPutPairList) Len() int           { return len(l) }
func (l batchPutPairList) Less(i, j int) bool { return l[i].Item < l[j].Item }
func (l batchPutPairList) Swap(i, j int)      { l[j], l[i] = l[i], l[j] }

// Return the elements of the map sorted by item name.
func getSortedPairs(updateMap map[ItemName][]PutUpdate) batchPutPairList {
	res := batchPutPairList{}
	for item, updates := range updateMap {
		res = append(res, batchPutPair{item, updates})
	}

	sort.Sort(res)
	return res
}

func validateUpdate(u PutUpdate) (err error) {
	// Make sure the attribute name is legal.
	if u.Name == "" {
		return fmt.Errorf("Invalid attribute name; names must be non-empty.")
	}

	if err = validateValue(string(u.Name)); err != nil {
		return fmt.Errorf("Invalid attribute name: %v", err)
	}

	// Make sure the attribute value is legal.
	if err = validateValue(string(u.Value)); err != nil {
		return fmt.Errorf("Invalid attribute value: %v", err)
	}

	return nil
}

func validateUpdates(updates []PutUpdate) (err error) {
	numUpdates := len(updates)
	if numUpdates == 0 || numUpdates > 256 {
		return fmt.Errorf("Illegal number of updates: %d", numUpdates)
	}

	for _, u := range updates {
		if err = validateUpdate(u); err != nil {
			return fmt.Errorf("Invalid update (%v): %v", err, u)
		}
	}

	return nil
}

func validatePrecondition(p Precondition) (err error) {
	// Make sure the attribute name is legal.
	if p.Name == "" {
		return fmt.Errorf("Invalid attribute name; names must be non-empty.")
	}

	if err = validateValue(string(p.Name)); err != nil {
		return fmt.Errorf("Invalid attribute name: %v", err)
	}

	// We require exactly one operand.
	if (p.Value == nil) == (p.Exists == nil) {
		return fmt.Errorf("Preconditions must contain exactly one of Value and Exists.")
	}

	// Make sure the attribute value is legal, if present.
	if p.Value != nil {
		if err = validateValue(string(*p.Value)); err != nil {
			return fmt.Errorf("Invalid attribute value: %v", err)
		}
	}

	return nil
}

func (d *domain) PutAttributes(
	item ItemName,
	updates []PutUpdate,
	preconditions []Precondition) (err error) {
	// Make sure the item name is legal.
	if item == "" {
		return fmt.Errorf("Invalid item name; names must be non-empty.")
	}

	if err = validateValue(string(item)); err != nil {
		return fmt.Errorf("Invalid item name: %v", err)
	}

	// Validate updates.
	if err = validateUpdates(updates); err != nil {
		return err
	}

	// Validate preconditions.
	for _, p := range preconditions {
		if err = validatePrecondition(p); err != nil {
			return fmt.Errorf("Invalid precondition (%v): %v", err, p)
		}
	}

	// Assemble an appropriate request.
	req := conn.Request{}
	req["DomainName"] = d.name
	req["ItemName"] = string(item)

	for i, u := range updates {
		keyPrefix := fmt.Sprintf("Attribute.%d.", i+1)
		req[keyPrefix + "Name"] = u.Name
		req[keyPrefix + "Value"] = u.Value

		if u.Replace {
			req[keyPrefix + "Replace"] = "true"
		}
	}

	for i, p := range preconditions {
		keyPrefix := fmt.Sprintf("Expected.%d.", i+1)
		req[keyPrefix + "Name"] = p.Name

		if p.Value != nil {
			req[keyPrefix + "Value"] = *p.Value
		} else if *p.Exists {
			req[keyPrefix + "Exists"] = "true"
		} else {
			req[keyPrefix + "Exists"] = "false"
		}
	}

	// Call the connection.
	if _, err = d.c.SendRequest(req); err != nil {
		return fmt.Errorf("SendRequest: %v", err)
	}

	return nil
}

func (d *domain) BatchPutAttributes(updateMap map[ItemName][]PutUpdate) (err error) {
	// Make sure the size of the request is legal.
	numItems := len(updateMap)
	if numItems == 0 || numItems > 25 {
		return fmt.Errorf("Illegal number of items: %d", numItems)
	}

	// Make sure each item name and set of updates is legal.
	for item, updates := range updateMap {
		if item == "" {
			return fmt.Errorf("Invalid item name; names must be non-empty.")
		}

		if err = validateValue(string(item)); err != nil {
			return fmt.Errorf("Invalid item name: %v", err)
		}

		if err = validateUpdates(updates); err != nil {
			return fmt.Errorf("Updates for item %s: %v", item, err)
		}
	}

	// Build a request.
	req := conn.Request{}
	req["DomainName"] = d.name

	pairs := getSortedPairs(updateMap)
	for i, pair := range pairs {
		itemPrefix := fmt.Sprintf("Item.%d.", i+1)
		req[itemPrefix + "ItemName"] = string(pair.Item)

		for j, u := range pair.Updates {
			updatePrefix := fmt.Sprintf("%sAttribute.%d.", itemPrefix, j+1)
			req[updatePrefix + "Name"] = u.Name
			req[updatePrefix + "Value"] = u.Value

			if u.Replace {
				req[updatePrefix + "Replace"] = "true"
			}
		}
	}

	// Call the connection.
	if _, err = d.c.SendRequest(req); err != nil {
		return fmt.Errorf("SendRequest: %v", err)
	}

	return nil
}
<|endoftext|>


<commit_before>package expr

import (
	"strings"
	"github.com/garyburd/redigo/redis"
)

func redisGetHash(name string, key string, db int, c redis.Conn) (string, error) {
	c.Do("SELECT", db)
	v, err := c.Do("HGET", name, key)
	return redis.String(v, err)
}

func prepareMetric(metric string) string {
	parts := strings.Split(metric, ".")
	lastPart := parts[len(parts)-1]
	prefix := strings.Split(lastPart, ",")[0]
	return strings.Trim(prefix, ")")
}

func prepareKubeMetric(metric string) (string, string, string) {
	parts := strings.Split(metric, ".")
	suffix, name, item, key := "", "", "", ""
	if len(parts) > 8 {
		suffix = "." + strings.Join(parts[8:], ".")
	}
	if len(parts) == 8 {
		suffix = "." + parts[7]
	}
	if len(parts) > 7 {
		item = parts[7]
	}
	if len(parts) > 6 {
		name = parts[6]
	}
	if len(parts) > 4 {
		//parts[4] = node
		key = parts[4] + "_" + item
	}
	return name, key, suffix

}

func aliasByHash(metric string, redisHashName string, conn redis.Conn) string {
	if redisHashName == "kube" {
		name, key, suffix := prepareKubeMetric(metric)
		redisName, err := redisGetHash(name, key, 2, conn)
		if err != nil {
			return metric
		}
		return redisName + suffix
	} else {
		key := prepareMetric(metric)
		redisName, err := redisGetHash(redisHashName, key, 0, conn)
		if err != nil {
			return metric
		}
		return redisName
	}
}
<commit_msg>custom key for sql_query<commit_after>package expr

import (
	"strings"

	"github.com/garyburd/redigo/redis"
)

func redisGetHash(name string, key string, db int, c redis.Conn) (string, error) {
	c.Do("SELECT", db)
	v, err := c.Do("HGET", name, key)
	return redis.String(v, err)
}

func redisGetKey(key string, db int, c redis.Conn) (string, error) {
	c.Do("SELECT", db)
	v, err := c.Do("GET", key)
	return redis.String(v, err)
}

func prepareMetric(metric string) string {
	parts := strings.Split(metric, ".")
	lastPart := parts[len(parts)-1]
	prefix := strings.Split(lastPart, ",")[0]
	return strings.Trim(prefix, ")")
}

func prepareKubeMetric(metric string) (string, string, string) {
	parts := strings.Split(metric, ".")
	suffix, name, item, key := "", "", "", ""
	if len(parts) > 8 {
		suffix = "." + strings.Join(parts[8:], ".")
	}
	if len(parts) == 8 {
		suffix = "." + parts[7]
	}
	if len(parts) > 7 {
		item = parts[7]
	}
	if len(parts) > 6 {
		name = parts[6]
	}
	if len(parts) > 4 {
		//parts[4] = node
		key = parts[4] + "_" + item
	}
	return name, key, suffix

}

func aliasByHash(metric string, redisHashName string, conn redis.Conn) string {
	if redisHashName == "kube" {
		name, key, suffix := prepareKubeMetric(metric)
		redisName, err := redisGetHash(name, key, 2, conn)
		if err != nil {
			return metric
		}
		return redisName + suffix
	} else {
		key := prepareMetric(metric)
		var err error
		var redisName string
		if redisHashName == "sql_query" {
			redisName, err = redisGetKey(key, 5, conn)
		} else {
			redisName, err = redisGetHash(redisHashName, key, 0, conn)
		}
		if err != nil {
			return metric
		}
		return redisName
	}
}
<|endoftext|>


<commit_before>package consensus

import (
	"bytes"
	"errors"
	"fmt"

	"github.com/NebulousLabs/Sia/build"
	"github.com/NebulousLabs/Sia/crypto"
	"github.com/NebulousLabs/Sia/encoding"
	"github.com/NebulousLabs/Sia/types"
	"github.com/NebulousLabs/fastrand"

	"github.com/NebulousLabs/bolt"
)

// manageErr handles an error detected by the consistency checks.
func manageErr(tx *bolt.Tx, err error) {
	markInconsistency(tx)
	if build.DEBUG {
		panic(err)
	} else {
		fmt.Println(err)
	}
}

// consensusChecksum grabs a checksum of the consensus set by pushing all of
// the elements in sorted order into a merkle tree and taking the root. All
// consensus sets with the same current block should have identical consensus
// checksums.
func consensusChecksum(tx *bolt.Tx) crypto.Hash {
	// Create a checksum tree.
	tree := crypto.NewTree()

	// For all of the constant buckets, push every key and every value. Buckets
	// are sorted in byte-order, therefore this operation is deterministic.
	consensusSetBuckets := []*bolt.Bucket{
		tx.Bucket(BlockPath),
		tx.Bucket(SiacoinOutputs),
		tx.Bucket(FileContracts),
		tx.Bucket(SiafundOutputs),
		tx.Bucket(SiafundPool),
	}
	for i := range consensusSetBuckets {
		err := consensusSetBuckets[i].ForEach(func(k, v []byte) error {
			tree.Push(k)
			tree.Push(v)
			return nil
		})
		if err != nil {
			manageErr(tx, err)
		}
	}

	// Iterate through all the buckets looking for buckets prefixed with
	// prefixDSCO or prefixFCEX. Buckets are presented in byte-sorted order by
	// name.
	err := tx.ForEach(func(name []byte, b *bolt.Bucket) error {
		// If the bucket is not a delayed siacoin output bucket or a file
		// contract expiration bucket, skip.
		if !bytes.HasPrefix(name, prefixDSCO) && !bytes.HasPrefix(name, prefixFCEX) {
			return nil
		}

		// The bucket is a prefixed bucket - add all elements to the tree.
		return b.ForEach(func(k, v []byte) error {
			tree.Push(k)
			tree.Push(v)
			return nil
		})
	})
	if err != nil {
		manageErr(tx, err)
	}

	return tree.Root()
}

// checkSiacoinCount checks that the number of siacoins countable within the
// consensus set equal the expected number of siacoins for the block height.
func checkSiacoinCount(tx *bolt.Tx) {
	// Iterate through all the buckets looking for the delayed siacoin output
	// buckets, and check that they are for the correct heights.
	var dscoSiacoins types.Currency
	err := tx.ForEach(func(name []byte, b *bolt.Bucket) error {
		// Check if the bucket is a delayed siacoin output bucket.
		if !bytes.HasPrefix(name, prefixDSCO) {
			return nil
		}

		// Sum up the delayed outputs in this bucket.
		err := b.ForEach(func(_, delayedOutput []byte) error {
			var sco types.SiacoinOutput
			err := encoding.Unmarshal(delayedOutput, &sco)
			if err != nil {
				manageErr(tx, err)
			}
			dscoSiacoins = dscoSiacoins.Add(sco.Value)
			return nil
		})
		if err != nil {
			return err
		}
		return nil
	})
	if err != nil {
		manageErr(tx, err)
	}

	// Add all of the siacoin outputs.
	var scoSiacoins types.Currency
	err = tx.Bucket(SiacoinOutputs).ForEach(func(_, scoBytes []byte) error {
		var sco types.SiacoinOutput
		err := encoding.Unmarshal(scoBytes, &sco)
		if err != nil {
			manageErr(tx, err)
		}
		scoSiacoins = scoSiacoins.Add(sco.Value)
		return nil
	})
	if err != nil {
		manageErr(tx, err)
	}

	// Add all of the payouts from file contracts.
	var fcSiacoins types.Currency
	err = tx.Bucket(FileContracts).ForEach(func(_, fcBytes []byte) error {
		var fc types.FileContract
		err := encoding.Unmarshal(fcBytes, &fc)
		if err != nil {
			manageErr(tx, err)
		}
		var fcCoins types.Currency
		for _, output := range fc.ValidProofOutputs {
			fcCoins = fcCoins.Add(output.Value)
		}
		fcSiacoins = fcSiacoins.Add(fcCoins)
		return nil
	})
	if err != nil {
		manageErr(tx, err)
	}

	// Add all of the siafund claims.
	var claimSiacoins types.Currency
	err = tx.Bucket(SiafundOutputs).ForEach(func(_, sfoBytes []byte) error {
		var sfo types.SiafundOutput
		err := encoding.Unmarshal(sfoBytes, &sfo)
		if err != nil {
			manageErr(tx, err)
		}

		coinsPerFund := getSiafundPool(tx).Sub(sfo.ClaimStart)
		claimCoins := coinsPerFund.Mul(sfo.Value).Div(types.SiafundCount)
		claimSiacoins = claimSiacoins.Add(claimCoins)
		return nil
	})
	if err != nil {
		manageErr(tx, err)
	}

	expectedSiacoins := types.CalculateNumSiacoins(blockHeight(tx))
	totalSiacoins := dscoSiacoins.Add(scoSiacoins).Add(fcSiacoins).Add(claimSiacoins)
	if !totalSiacoins.Equals(expectedSiacoins) {
		diagnostics := fmt.Sprintf("Wrong number of siacoins\nDsco: %v\nSco: %v\nFc: %v\nClaim: %v\n", dscoSiacoins, scoSiacoins, fcSiacoins, claimSiacoins)
		if totalSiacoins.Cmp(expectedSiacoins) < 0 {
			diagnostics += fmt.Sprintf("total: %v\nexpected: %v\n expected is bigger: %v", totalSiacoins, expectedSiacoins, expectedSiacoins.Sub(totalSiacoins))
		} else {
			diagnostics += fmt.Sprintf("total: %v\nexpected: %v\n expected is bigger: %v", totalSiacoins, expectedSiacoins, totalSiacoins.Sub(expectedSiacoins))
		}
		manageErr(tx, errors.New(diagnostics))
	}
}

// checkSiafundCount checks that the number of siafunds countable within the
// consensus set equal the expected number of siafunds for the block height.
func checkSiafundCount(tx *bolt.Tx) {
	var total types.Currency
	err := tx.Bucket(SiafundOutputs).ForEach(func(_, siafundOutputBytes []byte) error {
		var sfo types.SiafundOutput
		err := encoding.Unmarshal(siafundOutputBytes, &sfo)
		if err != nil {
			manageErr(tx, err)
		}
		total = total.Add(sfo.Value)
		return nil
	})
	if err != nil {
		manageErr(tx, err)
	}
	if !total.Equals(types.SiafundCount) {
		manageErr(tx, errors.New("wrong number if siafunds in the consensus set"))
	}
}

// checkDSCOs scans the sets of delayed siacoin outputs and checks for
// consistency.
func checkDSCOs(tx *bolt.Tx) {
	// Create a map to track which delayed siacoin output maps exist, and
	// another map to track which ids have appeared in the dsco set.
	dscoTracker := make(map[types.BlockHeight]struct{})
	idMap := make(map[types.SiacoinOutputID]struct{})

	// Iterate through all the buckets looking for the delayed siacoin output
	// buckets, and check that they are for the correct heights.
	err := tx.ForEach(func(name []byte, b *bolt.Bucket) error {
		// If the bucket is not a delayed siacoin output bucket or a file
		// contract expiration bucket, skip.
		if !bytes.HasPrefix(name, prefixDSCO) {
			return nil
		}

		// Add the bucket to the dscoTracker.
		var height types.BlockHeight
		err := encoding.Unmarshal(name[len(prefixDSCO):], &height)
		if err != nil {
			manageErr(tx, err)
		}
		_, exists := dscoTracker[height]
		if exists {
			return errors.New("repeat dsco map")
		}
		dscoTracker[height] = struct{}{}

		var total types.Currency
		err = b.ForEach(func(idBytes, delayedOutput []byte) error {
			// Check that the output id has not appeared in another dsco.
			var id types.SiacoinOutputID
			copy(id[:], idBytes)
			_, exists := idMap[id]
			if exists {
				return errors.New("repeat delayed siacoin output")
			}
			idMap[id] = struct{}{}

			// Sum the funds in the bucket.
			var sco types.SiacoinOutput
			err := encoding.Unmarshal(delayedOutput, &sco)
			if err != nil {
				manageErr(tx, err)
			}
			total = total.Add(sco.Value)
			return nil
		})
		if err != nil {
			return err
		}

		// Check that the minimum value has been achieved - the coinbase from
		// an earlier block is guaranteed to be in the bucket.
		minimumValue := types.CalculateCoinbase(height - types.MaturityDelay)
		if total.Cmp(minimumValue) < 0 {
			return errors.New("total number of coins in the delayed output bucket is incorrect")
		}
		return nil
	})
	if err != nil {
		manageErr(tx, err)
	}

	// Check that all of the correct heights are represented.
	currentHeight := blockHeight(tx)
	expectedBuckets := 0
	for i := currentHeight + 1; i <= currentHeight+types.MaturityDelay; i++ {
		if i < types.MaturityDelay {
			continue
		}
		_, exists := dscoTracker[i]
		if !exists {
			manageErr(tx, errors.New("missing a dsco bucket"))
		}
		expectedBuckets++
	}
	if len(dscoTracker) != expectedBuckets {
		manageErr(tx, errors.New("too many dsco buckets"))
	}
}

// checkRevertApply reverts the most recent block, checking to see that the
// consensus set hash matches the hash obtained for the previous block. Then it
// applies the block again and checks that the consensus set hash matches the
// original consensus set hash.
func (cs *ConsensusSet) checkRevertApply(tx *bolt.Tx) {
	current := currentProcessedBlock(tx)
	// Don't perform the check if this block is the genesis block.
	if current.Block.ID() == cs.blockRoot.Block.ID() {
		return
	}

	parent, err := getBlockMap(tx, current.Block.ParentID)
	if err != nil {
		manageErr(tx, err)
	}
	if current.Height != parent.Height+1 {
		manageErr(tx, errors.New("parent structure of a block is incorrect"))
	}
	_, _, err = cs.forkBlockchain(tx, parent)
	if err != nil {
		manageErr(tx, err)
	}
	if consensusChecksum(tx) != parent.ConsensusChecksum {
		manageErr(tx, errors.New("consensus checksum mismatch after reverting"))
	}
	_, _, err = cs.forkBlockchain(tx, current)
	if err != nil {
		manageErr(tx, err)
	}
	if consensusChecksum(tx) != current.ConsensusChecksum {
		manageErr(tx, errors.New("consensus checksum mismatch after re-applying"))
	}
}

// checkConsistency runs a series of checks to make sure that the consensus set
// is consistent with some rules that should always be true.
func (cs *ConsensusSet) checkConsistency(tx *bolt.Tx) {
	if fastrand.Intn(25) == 0 {
		return
	}
	if cs.checkingConsistency {
		return
	}

	cs.checkingConsistency = true
	checkDSCOs(tx)
	checkSiacoinCount(tx)
	checkSiafundCount(tx)
	if build.DEBUG {
		cs.checkRevertApply(tx)
	}
	cs.checkingConsistency = false
}

// maybeCheckConsistency runs a consistency check with a small probability.
// Useful for detecting database corruption in production without needing to go
// through the extremely slow process of running a consistency check every
// block.
func (cs *ConsensusSet) maybeCheckConsistency(tx *bolt.Tx) {
	if fastrand.Intn(1000) == 0 {
		cs.checkConsistency(tx)
	}
}

// TODO: Check that every file contract has an expiration too, and that the
// number of file contracts + the number of expirations is equal.
<commit_msg>remove probabilistic consistency non-checking<commit_after>package consensus

import (
	"bytes"
	"errors"
	"fmt"

	"github.com/NebulousLabs/Sia/build"
	"github.com/NebulousLabs/Sia/crypto"
	"github.com/NebulousLabs/Sia/encoding"
	"github.com/NebulousLabs/Sia/types"
	"github.com/NebulousLabs/fastrand"

	"github.com/NebulousLabs/bolt"
)

// manageErr handles an error detected by the consistency checks.
func manageErr(tx *bolt.Tx, err error) {
	markInconsistency(tx)
	if build.DEBUG {
		panic(err)
	} else {
		fmt.Println(err)
	}
}

// consensusChecksum grabs a checksum of the consensus set by pushing all of
// the elements in sorted order into a merkle tree and taking the root. All
// consensus sets with the same current block should have identical consensus
// checksums.
func consensusChecksum(tx *bolt.Tx) crypto.Hash {
	// Create a checksum tree.
	tree := crypto.NewTree()

	// For all of the constant buckets, push every key and every value. Buckets
	// are sorted in byte-order, therefore this operation is deterministic.
	consensusSetBuckets := []*bolt.Bucket{
		tx.Bucket(BlockPath),
		tx.Bucket(SiacoinOutputs),
		tx.Bucket(FileContracts),
		tx.Bucket(SiafundOutputs),
		tx.Bucket(SiafundPool),
	}
	for i := range consensusSetBuckets {
		err := consensusSetBuckets[i].ForEach(func(k, v []byte) error {
			tree.Push(k)
			tree.Push(v)
			return nil
		})
		if err != nil {
			manageErr(tx, err)
		}
	}

	// Iterate through all the buckets looking for buckets prefixed with
	// prefixDSCO or prefixFCEX. Buckets are presented in byte-sorted order by
	// name.
	err := tx.ForEach(func(name []byte, b *bolt.Bucket) error {
		// If the bucket is not a delayed siacoin output bucket or a file
		// contract expiration bucket, skip.
		if !bytes.HasPrefix(name, prefixDSCO) && !bytes.HasPrefix(name, prefixFCEX) {
			return nil
		}

		// The bucket is a prefixed bucket - add all elements to the tree.
		return b.ForEach(func(k, v []byte) error {
			tree.Push(k)
			tree.Push(v)
			return nil
		})
	})
	if err != nil {
		manageErr(tx, err)
	}

	return tree.Root()
}

// checkSiacoinCount checks that the number of siacoins countable within the
// consensus set equal the expected number of siacoins for the block height.
func checkSiacoinCount(tx *bolt.Tx) {
	// Iterate through all the buckets looking for the delayed siacoin output
	// buckets, and check that they are for the correct heights.
	var dscoSiacoins types.Currency
	err := tx.ForEach(func(name []byte, b *bolt.Bucket) error {
		// Check if the bucket is a delayed siacoin output bucket.
		if !bytes.HasPrefix(name, prefixDSCO) {
			return nil
		}

		// Sum up the delayed outputs in this bucket.
		err := b.ForEach(func(_, delayedOutput []byte) error {
			var sco types.SiacoinOutput
			err := encoding.Unmarshal(delayedOutput, &sco)
			if err != nil {
				manageErr(tx, err)
			}
			dscoSiacoins = dscoSiacoins.Add(sco.Value)
			return nil
		})
		if err != nil {
			return err
		}
		return nil
	})
	if err != nil {
		manageErr(tx, err)
	}

	// Add all of the siacoin outputs.
	var scoSiacoins types.Currency
	err = tx.Bucket(SiacoinOutputs).ForEach(func(_, scoBytes []byte) error {
		var sco types.SiacoinOutput
		err := encoding.Unmarshal(scoBytes, &sco)
		if err != nil {
			manageErr(tx, err)
		}
		scoSiacoins = scoSiacoins.Add(sco.Value)
		return nil
	})
	if err != nil {
		manageErr(tx, err)
	}

	// Add all of the payouts from file contracts.
	var fcSiacoins types.Currency
	err = tx.Bucket(FileContracts).ForEach(func(_, fcBytes []byte) error {
		var fc types.FileContract
		err := encoding.Unmarshal(fcBytes, &fc)
		if err != nil {
			manageErr(tx, err)
		}
		var fcCoins types.Currency
		for _, output := range fc.ValidProofOutputs {
			fcCoins = fcCoins.Add(output.Value)
		}
		fcSiacoins = fcSiacoins.Add(fcCoins)
		return nil
	})
	if err != nil {
		manageErr(tx, err)
	}

	// Add all of the siafund claims.
	var claimSiacoins types.Currency
	err = tx.Bucket(SiafundOutputs).ForEach(func(_, sfoBytes []byte) error {
		var sfo types.SiafundOutput
		err := encoding.Unmarshal(sfoBytes, &sfo)
		if err != nil {
			manageErr(tx, err)
		}

		coinsPerFund := getSiafundPool(tx).Sub(sfo.ClaimStart)
		claimCoins := coinsPerFund.Mul(sfo.Value).Div(types.SiafundCount)
		claimSiacoins = claimSiacoins.Add(claimCoins)
		return nil
	})
	if err != nil {
		manageErr(tx, err)
	}

	expectedSiacoins := types.CalculateNumSiacoins(blockHeight(tx))
	totalSiacoins := dscoSiacoins.Add(scoSiacoins).Add(fcSiacoins).Add(claimSiacoins)
	if !totalSiacoins.Equals(expectedSiacoins) {
		diagnostics := fmt.Sprintf("Wrong number of siacoins\nDsco: %v\nSco: %v\nFc: %v\nClaim: %v\n", dscoSiacoins, scoSiacoins, fcSiacoins, claimSiacoins)
		if totalSiacoins.Cmp(expectedSiacoins) < 0 {
			diagnostics += fmt.Sprintf("total: %v\nexpected: %v\n expected is bigger: %v", totalSiacoins, expectedSiacoins, expectedSiacoins.Sub(totalSiacoins))
		} else {
			diagnostics += fmt.Sprintf("total: %v\nexpected: %v\n expected is bigger: %v", totalSiacoins, expectedSiacoins, totalSiacoins.Sub(expectedSiacoins))
		}
		manageErr(tx, errors.New(diagnostics))
	}
}

// checkSiafundCount checks that the number of siafunds countable within the
// consensus set equal the expected number of siafunds for the block height.
func checkSiafundCount(tx *bolt.Tx) {
	var total types.Currency
	err := tx.Bucket(SiafundOutputs).ForEach(func(_, siafundOutputBytes []byte) error {
		var sfo types.SiafundOutput
		err := encoding.Unmarshal(siafundOutputBytes, &sfo)
		if err != nil {
			manageErr(tx, err)
		}
		total = total.Add(sfo.Value)
		return nil
	})
	if err != nil {
		manageErr(tx, err)
	}
	if !total.Equals(types.SiafundCount) {
		manageErr(tx, errors.New("wrong number if siafunds in the consensus set"))
	}
}

// checkDSCOs scans the sets of delayed siacoin outputs and checks for
// consistency.
func checkDSCOs(tx *bolt.Tx) {
	// Create a map to track which delayed siacoin output maps exist, and
	// another map to track which ids have appeared in the dsco set.
	dscoTracker := make(map[types.BlockHeight]struct{})
	idMap := make(map[types.SiacoinOutputID]struct{})

	// Iterate through all the buckets looking for the delayed siacoin output
	// buckets, and check that they are for the correct heights.
	err := tx.ForEach(func(name []byte, b *bolt.Bucket) error {
		// If the bucket is not a delayed siacoin output bucket or a file
		// contract expiration bucket, skip.
		if !bytes.HasPrefix(name, prefixDSCO) {
			return nil
		}

		// Add the bucket to the dscoTracker.
		var height types.BlockHeight
		err := encoding.Unmarshal(name[len(prefixDSCO):], &height)
		if err != nil {
			manageErr(tx, err)
		}
		_, exists := dscoTracker[height]
		if exists {
			return errors.New("repeat dsco map")
		}
		dscoTracker[height] = struct{}{}

		var total types.Currency
		err = b.ForEach(func(idBytes, delayedOutput []byte) error {
			// Check that the output id has not appeared in another dsco.
			var id types.SiacoinOutputID
			copy(id[:], idBytes)
			_, exists := idMap[id]
			if exists {
				return errors.New("repeat delayed siacoin output")
			}
			idMap[id] = struct{}{}

			// Sum the funds in the bucket.
			var sco types.SiacoinOutput
			err := encoding.Unmarshal(delayedOutput, &sco)
			if err != nil {
				manageErr(tx, err)
			}
			total = total.Add(sco.Value)
			return nil
		})
		if err != nil {
			return err
		}

		// Check that the minimum value has been achieved - the coinbase from
		// an earlier block is guaranteed to be in the bucket.
		minimumValue := types.CalculateCoinbase(height - types.MaturityDelay)
		if total.Cmp(minimumValue) < 0 {
			return errors.New("total number of coins in the delayed output bucket is incorrect")
		}
		return nil
	})
	if err != nil {
		manageErr(tx, err)
	}

	// Check that all of the correct heights are represented.
	currentHeight := blockHeight(tx)
	expectedBuckets := 0
	for i := currentHeight + 1; i <= currentHeight+types.MaturityDelay; i++ {
		if i < types.MaturityDelay {
			continue
		}
		_, exists := dscoTracker[i]
		if !exists {
			manageErr(tx, errors.New("missing a dsco bucket"))
		}
		expectedBuckets++
	}
	if len(dscoTracker) != expectedBuckets {
		manageErr(tx, errors.New("too many dsco buckets"))
	}
}

// checkRevertApply reverts the most recent block, checking to see that the
// consensus set hash matches the hash obtained for the previous block. Then it
// applies the block again and checks that the consensus set hash matches the
// original consensus set hash.
func (cs *ConsensusSet) checkRevertApply(tx *bolt.Tx) {
	current := currentProcessedBlock(tx)
	// Don't perform the check if this block is the genesis block.
	if current.Block.ID() == cs.blockRoot.Block.ID() {
		return
	}

	parent, err := getBlockMap(tx, current.Block.ParentID)
	if err != nil {
		manageErr(tx, err)
	}
	if current.Height != parent.Height+1 {
		manageErr(tx, errors.New("parent structure of a block is incorrect"))
	}
	_, _, err = cs.forkBlockchain(tx, parent)
	if err != nil {
		manageErr(tx, err)
	}
	if consensusChecksum(tx) != parent.ConsensusChecksum {
		manageErr(tx, errors.New("consensus checksum mismatch after reverting"))
	}
	_, _, err = cs.forkBlockchain(tx, current)
	if err != nil {
		manageErr(tx, err)
	}
	if consensusChecksum(tx) != current.ConsensusChecksum {
		manageErr(tx, errors.New("consensus checksum mismatch after re-applying"))
	}
}

// checkConsistency runs a series of checks to make sure that the consensus set
// is consistent with some rules that should always be true.
func (cs *ConsensusSet) checkConsistency(tx *bolt.Tx) {
	if cs.checkingConsistency {
		return
	}

	cs.checkingConsistency = true
	checkDSCOs(tx)
	checkSiacoinCount(tx)
	checkSiafundCount(tx)
	if build.DEBUG {
		cs.checkRevertApply(tx)
	}
	cs.checkingConsistency = false
}

// maybeCheckConsistency runs a consistency check with a small probability.
// Useful for detecting database corruption in production without needing to go
// through the extremely slow process of running a consistency check every
// block.
func (cs *ConsensusSet) maybeCheckConsistency(tx *bolt.Tx) {
	if fastrand.Intn(1000) == 0 {
		cs.checkConsistency(tx)
	}
}

// TODO: Check that every file contract has an expiration too, and that the
// number of file contracts + the number of expirations is equal.
<|endoftext|>


<commit_before>package api

import (
	"errors"

	"github.com/jkomoros/boardgame"
	"github.com/jkomoros/boardgame/server/api/extendedgame"
	"github.com/jkomoros/boardgame/server/api/listing"
	"github.com/jkomoros/boardgame/server/api/users"
)

//Note: this is also duplicated in moves/seat_player.go
const playerToSeatRendevousDataType = "github.com/jkomoros/boardgame/server/api.PlayerToSeat"

//StorageManager extends the base boardgame.StorageManager with a few more
//methods necessary to make server work. When creating a new Server, you need
//to pass in a ServerStorageManager, which wraps one of these objects and thus
//implements these methods, too.
type StorageManager interface {

	//StorageManager extends the boardgame.StorageManager interface. Those
	//methods have two additional semantic expectations, however:
	//SaveGameAndCurrentState should create an ExtendedGameStorageRecord on
	//the first save of a game.
	boardgame.StorageManager

	//Name returns the name of the storage manager type, for example "memory", "bolt", or "mysql"
	Name() string

	//WithManagers is called during set up with references to all of the
	//managers. Will be called before Connect() is called.
	WithManagers(managers []*boardgame.GameManager)

	//Connect will be called before issuing any other substantive calls. The
	//config string is specific to the type of storage layer, which can be
	//interrogated with Nmae().
	Connect(config string) error

	//ExtendedGame is like Game(), but it returns an extended storage record
	//with additional fields necessary for Server.
	ExtendedGame(id string) (*extendedgame.StorageRecord, error)

	CombinedGame(id string) (*extendedgame.CombinedStorageRecord, error)

	//UpdateExtendedGame updates the extended game with the given Id.
	UpdateExtendedGame(id string, eGame *extendedgame.StorageRecord) error

	//Close should be called before the server is shut down.
	Close()

	//ListGames should list up to max games, in descending order based on the
	//LastActivity. If gameType is not "", only returns games that are that
	//gameType. If gameType is "", all gametypes are fine.
	ListGames(max int, list listing.Type, userID string, gameType string) []*extendedgame.CombinedStorageRecord

	//UserIDsForGame returns an array whose length equals game.NumPlayers.
	//Each one is either empty if there is no user in that slot yet, or the
	//uid representing the user.
	UserIDsForGame(gameID string) []string

	SetPlayerForGame(gameID string, playerIndex boardgame.PlayerIndex, userID string) error

	//Store or update all fields
	UpdateUser(user *users.StorageRecord) error

	GetUserByID(uid string) *users.StorageRecord

	GetUserByCookie(cookie string) *users.StorageRecord

	//If user is nil, the cookie should be deleted if it exists. If the user
	//does not yet exist, it should be added to the database.
	ConnectCookieToUser(cookie string, user *users.StorageRecord) error

	//Note: whenever you add methods here, also add them to boardgame/storage/test/StorageManager
}

//ServerStorageManager implements the ServerStorage interface by wrapping an
//object that supports StorageManager.
type ServerStorageManager struct {
	StorageManager
	server *Server
}

//NewServerStorageManager takes an object that implements StorageManager and
//wraps it.
func NewServerStorageManager(manager StorageManager) *ServerStorageManager {
	return &ServerStorageManager{
		manager,
		nil,
	}
}

//PlayerMoveApplied notifies all clients connected vie an active WebSocket for
//that game that the game has been modified.
func (s *ServerStorageManager) PlayerMoveApplied(game *boardgame.GameStorageRecord) error {

	//Do the wrapped manager's PlayerMoveApplied in case it has one.
	if err := s.StorageManager.PlayerMoveApplied(game); err != nil {
		return err
	}

	server := s.server

	if server == nil {
		return errors.New("no server configured. The storage manager should be added to a Server before it's used")
	}

	//Notify the web sockets that the game was changed
	server.notifier.gameChanged(game)

	return nil

}

//FetchInjectedDataForGame is where the server signals to SeatPlayer that
//there's a player to be seated.
func (s *ServerStorageManager) FetchInjectedDataForGame(gameID string, dataType string) interface{} {
	if dataType == playerToSeatRendevousDataType {
		//TODO: check to see if we have pending players for that gameID and return a moves/interfaces.SeatPlayerSignaler with a callback
		return nil
	}
	return s.StorageManager.FetchInjectedDataForGame(gameID, dataType)
}
<commit_msg>FetchInjectedDataForGame now returns the playersToSeat, so SeatPlayer should be signaled when there's a player to seat.<commit_after>package api

import (
	"errors"

	"github.com/jkomoros/boardgame"
	"github.com/jkomoros/boardgame/server/api/extendedgame"
	"github.com/jkomoros/boardgame/server/api/listing"
	"github.com/jkomoros/boardgame/server/api/users"
)

//Note: this is also duplicated in moves/seat_player.go
const playerToSeatRendevousDataType = "github.com/jkomoros/boardgame/server/api.PlayerToSeat"

//StorageManager extends the base boardgame.StorageManager with a few more
//methods necessary to make server work. When creating a new Server, you need
//to pass in a ServerStorageManager, which wraps one of these objects and thus
//implements these methods, too.
type StorageManager interface {

	//StorageManager extends the boardgame.StorageManager interface. Those
	//methods have two additional semantic expectations, however:
	//SaveGameAndCurrentState should create an ExtendedGameStorageRecord on
	//the first save of a game.
	boardgame.StorageManager

	//Name returns the name of the storage manager type, for example "memory", "bolt", or "mysql"
	Name() string

	//WithManagers is called during set up with references to all of the
	//managers. Will be called before Connect() is called.
	WithManagers(managers []*boardgame.GameManager)

	//Connect will be called before issuing any other substantive calls. The
	//config string is specific to the type of storage layer, which can be
	//interrogated with Nmae().
	Connect(config string) error

	//ExtendedGame is like Game(), but it returns an extended storage record
	//with additional fields necessary for Server.
	ExtendedGame(id string) (*extendedgame.StorageRecord, error)

	CombinedGame(id string) (*extendedgame.CombinedStorageRecord, error)

	//UpdateExtendedGame updates the extended game with the given Id.
	UpdateExtendedGame(id string, eGame *extendedgame.StorageRecord) error

	//Close should be called before the server is shut down.
	Close()

	//ListGames should list up to max games, in descending order based on the
	//LastActivity. If gameType is not "", only returns games that are that
	//gameType. If gameType is "", all gametypes are fine.
	ListGames(max int, list listing.Type, userID string, gameType string) []*extendedgame.CombinedStorageRecord

	//UserIDsForGame returns an array whose length equals game.NumPlayers.
	//Each one is either empty if there is no user in that slot yet, or the
	//uid representing the user.
	UserIDsForGame(gameID string) []string

	SetPlayerForGame(gameID string, playerIndex boardgame.PlayerIndex, userID string) error

	//Store or update all fields
	UpdateUser(user *users.StorageRecord) error

	GetUserByID(uid string) *users.StorageRecord

	GetUserByCookie(cookie string) *users.StorageRecord

	//If user is nil, the cookie should be deleted if it exists. If the user
	//does not yet exist, it should be added to the database.
	ConnectCookieToUser(cookie string, user *users.StorageRecord) error

	//Note: whenever you add methods here, also add them to boardgame/storage/test/StorageManager
}

//ServerStorageManager implements the ServerStorage interface by wrapping an
//object that supports StorageManager.
type ServerStorageManager struct {
	StorageManager
	server *Server
}

//NewServerStorageManager takes an object that implements StorageManager and
//wraps it.
func NewServerStorageManager(manager StorageManager) *ServerStorageManager {
	return &ServerStorageManager{
		manager,
		nil,
	}
}

//PlayerMoveApplied notifies all clients connected vie an active WebSocket for
//that game that the game has been modified.
func (s *ServerStorageManager) PlayerMoveApplied(game *boardgame.GameStorageRecord) error {

	//Do the wrapped manager's PlayerMoveApplied in case it has one.
	if err := s.StorageManager.PlayerMoveApplied(game); err != nil {
		return err
	}

	server := s.server

	if server == nil {
		return errors.New("no server configured. The storage manager should be added to a Server before it's used")
	}

	//Notify the web sockets that the game was changed
	server.notifier.gameChanged(game)

	return nil

}

//FetchInjectedDataForGame is where the server signals to SeatPlayer that
//there's a player to be seated.
func (s *ServerStorageManager) FetchInjectedDataForGame(gameID string, dataType string) interface{} {
	if dataType == playerToSeatRendevousDataType {
		slice := s.server.playersToSeat[gameID]
		if len(slice) > 0 {
			//The item's Committed() will remove itself from the list.
			return slice[0]
		}
	}
	return s.StorageManager.FetchInjectedDataForGame(gameID, dataType)
}
<|endoftext|>