<commit_before>package dht

import (
	"context"
	"errors"
	"fmt"
	"sort"
	"sync"
	"time"

	"github.com/sirupsen/logrus"

	"github.com/nimona/go-nimona/mesh"
	"github.com/nimona/go-nimona/wire"
)

var (
	ErrNotFound = errors.New("not found")
)

const (
	wireExtention        = "dht"
	closestPeersToReturn = 8
	maxQueryTime         = time.Second * 5
)

// DHT is the struct that implements the dht protocol
type DHT struct {
	peerID         string
	store          *Store
	wire           wire.Wire
	registry       mesh.Registry
	queries        sync.Map
	refreshBuckets bool
}

func NewDHT(wr wire.Wire, pr mesh.Registry, peerID string, refreshBuckets bool, bootstrapPeerIDs ...string) (*DHT, error) {
	// create new kv store
	store, _ := newStore()

	// Create DHT node
	nd := &DHT{
		peerID:         peerID,
		store:          store,
		wire:           wr,
		registry:       pr,
		queries:        sync.Map{},
		refreshBuckets: refreshBuckets,
	}

	for _, peerID := range bootstrapPeerIDs {
		nd.registry.PutPeerInfo(&mesh.PeerInfo{
			ID: "bootstrap",
			Protocols: map[string][]string{
				"wire": []string{
					fmt.Sprintf("tcp:%s:26801/yamux/router/wire", peerID),
				},
			},
		})
	}

	wr.HandleExtensionEvents("dht", nd.handleMessage)

	go nd.refresh()
	return nd, nil
}

func (nd *DHT) refresh() {
	// TODO our init process is a bit messed up and registry doesn't know
	// about the peer's protocols instantly
	for len(nd.registry.GetLocalPeerInfo().Protocols) == 0 {
		time.Sleep(time.Millisecond * 250)
	}
	for {
		peerInfo := nd.registry.GetLocalPeerInfo()
		closestPeers, err := nd.FindPeersClosestTo(peerInfo.ID, closestPeersToReturn)
		if err != nil {
			logrus.WithError(err).Warnf("refresh could not get peers ids")
			return
		}

		resp := messageGetPeerInfo{
			SenderPeerInfo: *peerInfo,
			PeerID:         peerInfo.ID,
		}
		ctx := context.Background()
		peerIDs := getPeerIDsFromPeerInfos(closestPeers)
		nd.wire.Send(ctx, wireExtention, PayloadTypeGetPeerInfo, resp, peerIDs)
		time.Sleep(time.Second * 999)
	}
}

func (nd *DHT) handleMessage(message *wire.Message) error {
	// logrus.Debug("Got message", message.String())

	senderPeerInfo := &messageSenderPeerInfo{}
	if err := message.DecodePayload(senderPeerInfo); err == nil {
		nd.registry.PutPeerInfo(&senderPeerInfo.SenderPeerInfo)
	}

	switch message.PayloadType {
	case PayloadTypeGetPeerInfo:
		nd.handleGetPeerInfo(message)
	case PayloadTypePutPeerInfo:
		nd.handlePutPeerInfo(message)
	case PayloadTypeGetProviders:
		nd.handleGetProviders(message)
	case PayloadTypePutProviders:
		nd.handlePutProviders(message)
	case PayloadTypeGetValue:
		nd.handleGetValue(message)
	case PayloadTypePutValue:
		nd.handlePutValue(message)
	default:
		logrus.WithField("message.PayloadType", message.PayloadType).Warn("Payload type not known")
		return nil
	}
	return nil
}

func (nd *DHT) handleGetPeerInfo(message *wire.Message) {
	payload := &messageGetPeerInfo{}
	if err := message.DecodePayload(payload); err != nil {
		return
	}

	peerInfo, err := nd.registry.GetPeerInfo(payload.PeerID)
	if err != nil {
		return
	}

	closestPeers, _ := nd.FindPeersClosestTo(payload.PeerID, closestPeersToReturn)
	resp := messagePutPeerInfo{
		SenderPeerInfo: *nd.registry.GetLocalPeerInfo(),
		RequestID:      payload.RequestID,
		PeerID:         payload.PeerID,
		PeerInfo:       *peerInfo,
		ClosestPeers:   closestPeers,
	}

	ctx := context.Background()
	to := []string{message.From}
	nd.wire.Send(ctx, wireExtention, PayloadTypePutPeerInfo, resp, to)
}

func (nd *DHT) handlePutPeerInfo(message *wire.Message) {
	payload := &messagePutPeerInfo{}
	if err := message.DecodePayload(payload); err != nil {
		return
	}

	nd.registry.PutPeerInfo(&payload.PeerInfo)
	for _, peerInfo := range payload.ClosestPeers {
		nd.registry.PutPeerInfo(peerInfo)
	}

	if payload.RequestID == "" {
		return
	}

	q, exists := nd.queries.Load(payload.RequestID)
	if !exists {
		return
	}

	q.(*query).incomingMessages <- payload
}

func (nd *DHT) handleGetProviders(message *wire.Message) {
	payload := &messageGetProviders{}
	if err := message.DecodePayload(payload); err != nil {
		return
	}

	providers, err := nd.store.GetProviders(payload.Key)
	if err != nil {
		return
	}

	closestPeers, _ := nd.FindPeersClosestTo(payload.Key, closestPeersToReturn)
	resp := messagePutProviders{
		SenderPeerInfo: *nd.registry.GetLocalPeerInfo(),
		RequestID:      payload.RequestID,
		Key:            payload.Key,
		PeerIDs:        providers,
		ClosestPeers:   closestPeers,
	}

	ctx := context.Background()
	to := []string{message.From}
	nd.wire.Send(ctx, wireExtention, PayloadTypePutProviders, resp, to)
}

func (nd *DHT) handlePutProviders(message *wire.Message) {
	payload := &messagePutProviders{}
	if err := message.DecodePayload(payload); err != nil {
		return
	}

	for _, peerInfo := range payload.ClosestPeers {
		nd.registry.PutPeerInfo(peerInfo)
	}

	if err := nd.store.PutProvider(payload.Key, payload.PeerIDs...); err != nil {
		return
	}

	if payload.RequestID == "" {
		return
	}

	q, exists := nd.queries.Load(payload.RequestID)
	if !exists {
		return
	}

	q.(*query).incomingMessages <- payload
}

func (nd *DHT) handleGetValue(message *wire.Message) {
	payload := &messageGetValue{}
	if err := message.DecodePayload(payload); err != nil {
		return
	}

	value, _ := nd.store.GetValue(payload.Key)

	closestPeers, _ := nd.FindPeersClosestTo(payload.Key, closestPeersToReturn)
	resp := messagePutValue{
		SenderPeerInfo: *nd.registry.GetLocalPeerInfo(),
		RequestID:      payload.RequestID,
		Key:            payload.Key,
		Value:          value,
		ClosestPeers:   closestPeers,
	}

	ctx := context.Background()
	to := []string{message.From}
	nd.wire.Send(ctx, wireExtention, PayloadTypePutValue, resp, to)
}

func (nd *DHT) handlePutValue(message *wire.Message) {
	payload := &messagePutValue{}
	if err := message.DecodePayload(payload); err != nil {
		return
	}

	for _, peerInfo := range payload.ClosestPeers {
		nd.registry.PutPeerInfo(peerInfo)
	}

	if err := nd.store.PutValue(payload.Key, payload.Value); err != nil {
		return
	}

	if payload.RequestID == "" {
		return
	}

	q, exists := nd.queries.Load(payload.RequestID)
	if !exists {
		return
	}

	q.(*query).incomingMessages <- payload
}

// FindPeersClosestTo returns an array of n peers closest to the given key by xor distance
func (nd *DHT) FindPeersClosestTo(tk string, n int) ([]*mesh.PeerInfo, error) {
	// place to hold the results
	rks := []*mesh.PeerInfo{}

	htk := hash(tk)

	peerInfos, _ := nd.registry.GetAllPeerInfo()

	// slice to hold the distances
	dists := []distEntry{}
	for _, peerInfo := range peerInfos {
		// calculate distance
		de := distEntry{
			key:      peerInfo.ID,
			dist:     xor([]byte(htk), []byte(hash(peerInfo.ID))),
			peerInfo: peerInfo,
		}
		exists := false
		for _, ee := range dists {
			if ee.key == peerInfo.ID {
				exists = true
				break
			}
		}
		if !exists {
			dists = append(dists, de)
		}
	}

	// sort the distances
	sort.Slice(dists, func(i, j int) bool {
		return lessIntArr(dists[i].dist, dists[j].dist)
	})

	if n > len(dists) {
		n = len(dists)
	}

	// append n the first n number of keys
	for _, de := range dists {
		rks = append(rks, de.peerInfo)
		n--
		if n == 0 {
			break
		}
	}

	return rks, nil
}

func (nd *DHT) GetPeerInfo(ctx context.Context, key string) (*mesh.PeerInfo, error) {
	q := &query{
		dht:              nd,
		id:               mesh.RandStringBytesMaskImprSrc(8),
		key:              key,
		queryType:        PeerInfoQuery,
		incomingMessages: make(chan interface{}),
		outgoingMessages: make(chan interface{}),
	}

	nd.queries.Store(q.id, q)

	go q.Run(ctx)

	for {
		select {
		case value := <-q.outgoingMessages:
			return value.(*mesh.PeerInfo), nil
		case <-time.After(maxQueryTime):
			return nil, ErrNotFound
		case <-ctx.Done():
			return nil, ErrNotFound
		}
	}
}

func (nd *DHT) PutValue(ctx context.Context, key, value string) error {
	if err := nd.store.PutValue(key, value); err != nil {
		return err
	}

	closestPeers, _ := nd.FindPeersClosestTo(key, closestPeersToReturn)
	resp := messagePutValue{
		SenderPeerInfo: *nd.registry.GetLocalPeerInfo(),
		Key:            key,
		Value:          value,
	}

	closestPeerIDs := getPeerIDsFromPeerInfos(closestPeers)
	return nd.wire.Send(ctx, wireExtention, PayloadTypePutValue, resp, closestPeerIDs)
}

func (nd *DHT) GetValue(ctx context.Context, key string) (string, error) {
	q := &query{
		dht:              nd,
		id:               mesh.RandStringBytesMaskImprSrc(8),
		key:              key,
		queryType:        ValueQuery,
		incomingMessages: make(chan interface{}),
		outgoingMessages: make(chan interface{}),
	}

	nd.queries.Store(q.id, q)

	go q.Run(ctx)

	for {
		select {
		case value := <-q.outgoingMessages:
			valueStr, ok := value.(string)
			if !ok {
				continue
			}
			return valueStr, nil
		case <-time.After(maxQueryTime):
			return "", ErrNotFound
		case <-ctx.Done():
			return "", ErrNotFound
		}
	}
}

// TODO Find a better name for this
func (nd *DHT) PutProviders(ctx context.Context, key string) error {
	localPeerID := nd.registry.GetLocalPeerInfo().ID
	if err := nd.store.PutProvider(key, localPeerID); err != nil {
		return err
	}

	closestPeers, _ := nd.FindPeersClosestTo(key, closestPeersToReturn)
	resp := messagePutProviders{
		SenderPeerInfo: *nd.registry.GetLocalPeerInfo(),
		Key:            key,
		PeerIDs:        []string{localPeerID},
	}

	closestPeerIDs := getPeerIDsFromPeerInfos(closestPeers)
	return nd.wire.Send(ctx, wireExtention, PayloadTypePutProviders, resp, closestPeerIDs)
}

func (nd *DHT) GetProviders(ctx context.Context, key string) ([]string, error) {
	q := &query{
		dht:              nd,
		id:               mesh.RandStringBytesMaskImprSrc(8),
		key:              key,
		queryType:        ProviderQuery,
		incomingMessages: make(chan interface{}),
		outgoingMessages: make(chan interface{}),
	}

	nd.queries.Store(q.id, q)

	go q.Run(ctx)

	providers := []string{}
	for {
		select {
		case values := <-q.outgoingMessages:
			valuesStr, ok := values.([]string)
			if !ok {
				continue
			}
			providers = append(providers, valuesStr...)
		case <-time.After(maxQueryTime):
			return providers, nil
		case <-ctx.Done():
			return providers, nil
		}
	}
}

func (nd *DHT) GetAllProviders() (map[string][]string, error) {
	return nd.store.GetAllProviders()
}

func (nd *DHT) GetAllValues() (map[string]string, error) {
	return nd.store.GetAllValues()
}

func getPeerIDsFromPeerInfos(peerInfos []*mesh.PeerInfo) []string {
	peerIDs := []string{}
	for _, peerInfo := range peerInfos {
		peerIDs = append(peerIDs, peerInfo.ID)
	}
	return peerIDs
}
<commit_msg>Reduce refresh time<commit_after>package dht

import (
	"context"
	"errors"
	"fmt"
	"sort"
	"sync"
	"time"

	"github.com/sirupsen/logrus"

	"github.com/nimona/go-nimona/mesh"
	"github.com/nimona/go-nimona/wire"
)

var (
	ErrNotFound = errors.New("not found")
)

const (
	wireExtention        = "dht"
	closestPeersToReturn = 8
	maxQueryTime         = time.Second * 5
)

// DHT is the struct that implements the dht protocol
type DHT struct {
	peerID         string
	store          *Store
	wire           wire.Wire
	registry       mesh.Registry
	queries        sync.Map
	refreshBuckets bool
}

func NewDHT(wr wire.Wire, pr mesh.Registry, peerID string, refreshBuckets bool, bootstrapPeerIDs ...string) (*DHT, error) {
	// create new kv store
	store, _ := newStore()

	// Create DHT node
	nd := &DHT{
		peerID:         peerID,
		store:          store,
		wire:           wr,
		registry:       pr,
		queries:        sync.Map{},
		refreshBuckets: refreshBuckets,
	}

	for _, peerID := range bootstrapPeerIDs {
		nd.registry.PutPeerInfo(&mesh.PeerInfo{
			ID: "bootstrap",
			Protocols: map[string][]string{
				"wire": []string{
					fmt.Sprintf("tcp:%s:26801/yamux/router/wire", peerID),
				},
			},
		})
	}

	wr.HandleExtensionEvents("dht", nd.handleMessage)

	go nd.refresh()
	return nd, nil
}

func (nd *DHT) refresh() {
	// TODO our init process is a bit messed up and registry doesn't know
	// about the peer's protocols instantly
	for len(nd.registry.GetLocalPeerInfo().Protocols) == 0 {
		time.Sleep(time.Millisecond * 250)
	}
	for {
		peerInfo := nd.registry.GetLocalPeerInfo()
		closestPeers, err := nd.FindPeersClosestTo(peerInfo.ID, closestPeersToReturn)
		if err != nil {
			logrus.WithError(err).Warnf("refresh could not get peers ids")
			return
		}

		resp := messageGetPeerInfo{
			SenderPeerInfo: *peerInfo,
			PeerID:         peerInfo.ID,
		}
		ctx := context.Background()
		peerIDs := getPeerIDsFromPeerInfos(closestPeers)
		nd.wire.Send(ctx, wireExtention, PayloadTypeGetPeerInfo, resp, peerIDs)
		time.Sleep(time.Second * 30)
	}
}

func (nd *DHT) handleMessage(message *wire.Message) error {
	// logrus.Debug("Got message", message.String())

	senderPeerInfo := &messageSenderPeerInfo{}
	if err := message.DecodePayload(senderPeerInfo); err == nil {
		nd.registry.PutPeerInfo(&senderPeerInfo.SenderPeerInfo)
	}

	switch message.PayloadType {
	case PayloadTypeGetPeerInfo:
		nd.handleGetPeerInfo(message)
	case PayloadTypePutPeerInfo:
		nd.handlePutPeerInfo(message)
	case PayloadTypeGetProviders:
		nd.handleGetProviders(message)
	case PayloadTypePutProviders:
		nd.handlePutProviders(message)
	case PayloadTypeGetValue:
		nd.handleGetValue(message)
	case PayloadTypePutValue:
		nd.handlePutValue(message)
	default:
		logrus.WithField("message.PayloadType", message.PayloadType).Warn("Payload type not known")
		return nil
	}
	return nil
}

func (nd *DHT) handleGetPeerInfo(message *wire.Message) {
	payload := &messageGetPeerInfo{}
	if err := message.DecodePayload(payload); err != nil {
		return
	}

	peerInfo, err := nd.registry.GetPeerInfo(payload.PeerID)
	if err != nil {
		return
	}

	closestPeers, _ := nd.FindPeersClosestTo(payload.PeerID, closestPeersToReturn)
	resp := messagePutPeerInfo{
		SenderPeerInfo: *nd.registry.GetLocalPeerInfo(),
		RequestID:      payload.RequestID,
		PeerID:         payload.PeerID,
		PeerInfo:       *peerInfo,
		ClosestPeers:   closestPeers,
	}

	ctx := context.Background()
	to := []string{message.From}
	nd.wire.Send(ctx, wireExtention, PayloadTypePutPeerInfo, resp, to)
}

func (nd *DHT) handlePutPeerInfo(message *wire.Message) {
	payload := &messagePutPeerInfo{}
	if err := message.DecodePayload(payload); err != nil {
		return
	}

	nd.registry.PutPeerInfo(&payload.PeerInfo)
	for _, peerInfo := range payload.ClosestPeers {
		nd.registry.PutPeerInfo(peerInfo)
	}

	if payload.RequestID == "" {
		return
	}

	q, exists := nd.queries.Load(payload.RequestID)
	if !exists {
		return
	}

	q.(*query).incomingMessages <- payload
}

func (nd *DHT) handleGetProviders(message *wire.Message) {
	payload := &messageGetProviders{}
	if err := message.DecodePayload(payload); err != nil {
		return
	}

	providers, err := nd.store.GetProviders(payload.Key)
	if err != nil {
		return
	}

	closestPeers, _ := nd.FindPeersClosestTo(payload.Key, closestPeersToReturn)
	resp := messagePutProviders{
		SenderPeerInfo: *nd.registry.GetLocalPeerInfo(),
		RequestID:      payload.RequestID,
		Key:            payload.Key,
		PeerIDs:        providers,
		ClosestPeers:   closestPeers,
	}

	ctx := context.Background()
	to := []string{message.From}
	nd.wire.Send(ctx, wireExtention, PayloadTypePutProviders, resp, to)
}

func (nd *DHT) handlePutProviders(message *wire.Message) {
	payload := &messagePutProviders{}
	if err := message.DecodePayload(payload); err != nil {
		return
	}

	for _, peerInfo := range payload.ClosestPeers {
		nd.registry.PutPeerInfo(peerInfo)
	}

	if err := nd.store.PutProvider(payload.Key, payload.PeerIDs...); err != nil {
		return
	}

	if payload.RequestID == "" {
		return
	}

	q, exists := nd.queries.Load(payload.RequestID)
	if !exists {
		return
	}

	q.(*query).incomingMessages <- payload
}

func (nd *DHT) handleGetValue(message *wire.Message) {
	payload := &messageGetValue{}
	if err := message.DecodePayload(payload); err != nil {
		return
	}

	value, _ := nd.store.GetValue(payload.Key)

	closestPeers, _ := nd.FindPeersClosestTo(payload.Key, closestPeersToReturn)
	resp := messagePutValue{
		SenderPeerInfo: *nd.registry.GetLocalPeerInfo(),
		RequestID:      payload.RequestID,
		Key:            payload.Key,
		Value:          value,
		ClosestPeers:   closestPeers,
	}

	ctx := context.Background()
	to := []string{message.From}
	nd.wire.Send(ctx, wireExtention, PayloadTypePutValue, resp, to)
}

func (nd *DHT) handlePutValue(message *wire.Message) {
	payload := &messagePutValue{}
	if err := message.DecodePayload(payload); err != nil {
		return
	}

	for _, peerInfo := range payload.ClosestPeers {
		nd.registry.PutPeerInfo(peerInfo)
	}

	if err := nd.store.PutValue(payload.Key, payload.Value); err != nil {
		return
	}

	if payload.RequestID == "" {
		return
	}

	q, exists := nd.queries.Load(payload.RequestID)
	if !exists {
		return
	}

	q.(*query).incomingMessages <- payload
}

// FindPeersClosestTo returns an array of n peers closest to the given key by xor distance
func (nd *DHT) FindPeersClosestTo(tk string, n int) ([]*mesh.PeerInfo, error) {
	// place to hold the results
	rks := []*mesh.PeerInfo{}

	htk := hash(tk)

	peerInfos, _ := nd.registry.GetAllPeerInfo()

	// slice to hold the distances
	dists := []distEntry{}
	for _, peerInfo := range peerInfos {
		// calculate distance
		de := distEntry{
			key:      peerInfo.ID,
			dist:     xor([]byte(htk), []byte(hash(peerInfo.ID))),
			peerInfo: peerInfo,
		}
		exists := false
		for _, ee := range dists {
			if ee.key == peerInfo.ID {
				exists = true
				break
			}
		}
		if !exists {
			dists = append(dists, de)
		}
	}

	// sort the distances
	sort.Slice(dists, func(i, j int) bool {
		return lessIntArr(dists[i].dist, dists[j].dist)
	})

	if n > len(dists) {
		n = len(dists)
	}

	// append n the first n number of keys
	for _, de := range dists {
		rks = append(rks, de.peerInfo)
		n--
		if n == 0 {
			break
		}
	}

	return rks, nil
}

func (nd *DHT) GetPeerInfo(ctx context.Context, key string) (*mesh.PeerInfo, error) {
	q := &query{
		dht:              nd,
		id:               mesh.RandStringBytesMaskImprSrc(8),
		key:              key,
		queryType:        PeerInfoQuery,
		incomingMessages: make(chan interface{}),
		outgoingMessages: make(chan interface{}),
	}

	nd.queries.Store(q.id, q)

	go q.Run(ctx)

	for {
		select {
		case value := <-q.outgoingMessages:
			return value.(*mesh.PeerInfo), nil
		case <-time.After(maxQueryTime):
			return nil, ErrNotFound
		case <-ctx.Done():
			return nil, ErrNotFound
		}
	}
}

func (nd *DHT) PutValue(ctx context.Context, key, value string) error {
	if err := nd.store.PutValue(key, value); err != nil {
		return err
	}

	closestPeers, _ := nd.FindPeersClosestTo(key, closestPeersToReturn)
	resp := messagePutValue{
		SenderPeerInfo: *nd.registry.GetLocalPeerInfo(),
		Key:            key,
		Value:          value,
	}

	closestPeerIDs := getPeerIDsFromPeerInfos(closestPeers)
	return nd.wire.Send(ctx, wireExtention, PayloadTypePutValue, resp, closestPeerIDs)
}

func (nd *DHT) GetValue(ctx context.Context, key string) (string, error) {
	q := &query{
		dht:              nd,
		id:               mesh.RandStringBytesMaskImprSrc(8),
		key:              key,
		queryType:        ValueQuery,
		incomingMessages: make(chan interface{}),
		outgoingMessages: make(chan interface{}),
	}

	nd.queries.Store(q.id, q)

	go q.Run(ctx)

	for {
		select {
		case value := <-q.outgoingMessages:
			valueStr, ok := value.(string)
			if !ok {
				continue
			}
			return valueStr, nil
		case <-time.After(maxQueryTime):
			return "", ErrNotFound
		case <-ctx.Done():
			return "", ErrNotFound
		}
	}
}

// TODO Find a better name for this
func (nd *DHT) PutProviders(ctx context.Context, key string) error {
	localPeerID := nd.registry.GetLocalPeerInfo().ID
	if err := nd.store.PutProvider(key, localPeerID); err != nil {
		return err
	}

	closestPeers, _ := nd.FindPeersClosestTo(key, closestPeersToReturn)
	resp := messagePutProviders{
		SenderPeerInfo: *nd.registry.GetLocalPeerInfo(),
		Key:            key,
		PeerIDs:        []string{localPeerID},
	}

	closestPeerIDs := getPeerIDsFromPeerInfos(closestPeers)
	return nd.wire.Send(ctx, wireExtention, PayloadTypePutProviders, resp, closestPeerIDs)
}

func (nd *DHT) GetProviders(ctx context.Context, key string) ([]string, error) {
	q := &query{
		dht:              nd,
		id:               mesh.RandStringBytesMaskImprSrc(8),
		key:              key,
		queryType:        ProviderQuery,
		incomingMessages: make(chan interface{}),
		outgoingMessages: make(chan interface{}),
	}

	nd.queries.Store(q.id, q)

	go q.Run(ctx)

	providers := []string{}
	for {
		select {
		case values := <-q.outgoingMessages:
			valuesStr, ok := values.([]string)
			if !ok {
				continue
			}
			providers = append(providers, valuesStr...)
		case <-time.After(maxQueryTime):
			return providers, nil
		case <-ctx.Done():
			return providers, nil
		}
	}
}

func (nd *DHT) GetAllProviders() (map[string][]string, error) {
	return nd.store.GetAllProviders()
}

func (nd *DHT) GetAllValues() (map[string]string, error) {
	return nd.store.GetAllValues()
}

func getPeerIDsFromPeerInfos(peerInfos []*mesh.PeerInfo) []string {
	peerIDs := []string{}
	for _, peerInfo := range peerInfos {
		peerIDs = append(peerIDs, peerInfo.ID)
	}
	return peerIDs
}
<|endoftext|>


<commit_before>package quorum

import (
	"bytes"
	"common"
	"common/crypto"
	"common/log"
	"encoding/gob"
	"fmt"
	"time"
)

// All information that needs to be passed between participants each block
type heartbeat struct {
	entropyStage1 crypto.TruncatedHash
	entropyStage2 common.Entropy
}

// Contains a heartbeat that has been signed iteratively, is a key part of the
// signed solution to the Byzantine Generals Problem
type signedHeartbeat struct {
	heartbeat     *heartbeat
	heartbeatHash crypto.TruncatedHash
	signatories   []byte             // a list of everyone who's seen the heartbeat
	signatures    []crypto.Signature // their corresponding signatures
}

// Using the current State, newHeartbeat() creates a heartbeat that fulfills all
// of the requirements of the quorum.
func (s *State) newHeartbeat() (hb *heartbeat, err error) {
	hb = new(heartbeat)

	// Fetch value used to produce EntropyStage1 in prev. heartbeat
	hb.entropyStage2 = s.storedEntropyStage2

	// Generate EntropyStage2 for next heartbeat
	entropy, err := crypto.RandomByteSlice(common.EntropyVolume)
	if err != nil {
		return
	}
	copy(s.storedEntropyStage2[:], entropy) // convert entropy from slice to byte array

	// Use EntropyStage2 to generate EntropyStage1 for this heartbeat
	hb.entropyStage1, err = crypto.CalculateTruncatedHash(s.storedEntropyStage2[:])
	if err != nil {
		return
	}

	// more code will be added here

	return
}

// Convert heartbeat to []byte
func (hb *heartbeat) GobEncode() (gobHeartbeat []byte, err error) {
	// test for bad input
	if hb == nil {
		err = fmt.Errorf("Cannot Encode a nil object")
		return
	}

	w := new(bytes.Buffer)
	encoder := gob.NewEncoder(w)
	err = encoder.Encode(hb.entropyStage1)
	if err != nil {
		return
	}
	err = encoder.Encode(hb.entropyStage2)
	if err != nil {
		return
	}

	gobHeartbeat = w.Bytes()
	return
}

// Convert []byte to heartbeat
func (hb *heartbeat) GobDecode(gobHeartbeat []byte) (err error) {
	r := bytes.NewBuffer(gobHeartbeat)
	decoder := gob.NewDecoder(r)
	err = decoder.Decode(&hb.entropyStage1)
	if err != nil {
		return
	}
	err = decoder.Decode(&hb.entropyStage2)
	if err != nil {
		return
	}
	return
}

// take new heartbeat (our own), sign it, and package it into a signedHearteat
// I'm pretty sure this only follows a newHeartbeat() call; they can be merged
func (s *State) signHeartbeat(hb *heartbeat) (sh *signedHeartbeat, err error) {
	sh = new(signedHeartbeat)

	// confirm heartbeat and hash
	sh.heartbeat = hb
	gobHb, err := hb.GobEncode()
	if err != nil {
		return
	}
	sh.heartbeatHash, err = crypto.CalculateTruncatedHash([]byte(gobHb))
	if err != nil {
		return
	}

	// fill out sigantures
	sh.signatures = make([]crypto.Signature, 1)
	signedHb, err := crypto.Sign(s.secretKey, string(sh.heartbeatHash[:]))
	if err != nil {
		return
	}
	sh.signatures[0] = signedHb.Signature
	sh.signatories = make([]byte, 1)
	sh.signatories[0] = s.participantIndex
	return
}

// Takes a signed heartbeat and broadcasts it to the quorum
func (s *State) announceSignedHeartbeat(sh *signedHeartbeat) (err error) {
	msh, err := sh.GobEncode()
	if err != nil {
		return
	}

	payload := append([]byte(string(incomingSignedHeartbeat)), msh...)
	s.broadcast(payload)
	return
}

// handleSignedHeartbeat takes the payload of an incoming message of type
// 'incomingSignedHeartbeat' and verifies it according to rules established by
// the specification.
//
// The return code is currently purely for the testing suite, the numbers
// have been chosen arbitrarily
func (s *State) handleSignedHeartbeat(payload []byte) (returnCode int) {
	// covert payload to SignedHeartbeat
	sh := new(signedHeartbeat)
	err := sh.GobDecode(payload)
	if err != nil {
		log.Infoln("Received bad message SignedHeartbeat: ", err)
		returnCode = 11
		return
	}

	// Check that the slices of signatures and signatories are of the same length
	if len(sh.signatures) != len(sh.signatories) {
		log.Infoln("SignedHeartbeat has mismatched signatures")
		returnCode = 1
		return
	}

	// check that there are not too many signatures and signatories
	if len(sh.signatories) > common.QuorumSize {
		log.Infoln("Received an over-signed signedHeartbeat")
		returnCode = 12
		return
	}

	s.stepLock.Lock() // prevents a benign race condition; is here to follow best practices
	currentStep := s.currentStep
	s.stepLock.Unlock()
	// s.CurrentStep must be less than or equal to len(sh.Signatories), unless
	// there is a new block and s.CurrentStep is common.QuorumSize
	if currentStep > len(sh.signatories) {
		if currentStep == common.QuorumSize && len(sh.signatories) == 1 {
			// by waiting common.StepDuration, the new block will be compiled
			time.Sleep(common.StepDuration)
			// now continue to rest of function
		} else {
			log.Infoln("Received an out-of-sync SignedHeartbeat")
			returnCode = 2
			return
		}
	}

	// Check bounds on first signatory
	if int(sh.signatories[0]) >= common.QuorumSize {
		log.Infoln("Received an out of bounds index")
		returnCode = 9
		return
	}

	// we are starting to read from memory, initiate locks
	s.participantsLock.RLock()
	s.heartbeatsLock.Lock()
	defer s.participantsLock.RUnlock()
	defer s.heartbeatsLock.Unlock()

	// check that first sigatory is a participant
	if s.participants[sh.signatories[0]] == nil {
		log.Infoln("Received heartbeat from non-participant")
		returnCode = 10
		return
	}

	// Check if we have already received this heartbeat
	_, exists := s.heartbeats[sh.signatories[0]][sh.heartbeatHash]
	if exists {
		returnCode = 8
		return
	}

	// Check if we already have two heartbeats from this host
	if len(s.heartbeats[sh.signatories[0]]) >= 2 {
		log.Infoln("Received many invalid heartbeats from one host")
		returnCode = 13
		return
	}

	// iterate through the signatures and make sure each is legal
	var signedMessage crypto.SignedMessage // grows each iteration
	signedMessage.Message = string(sh.heartbeatHash[:])
	previousSignatories := make(map[byte]bool) // which signatories have already signed
	for i, signatory := range sh.signatories {
		// Check bounds on the signatory
		if int(signatory) >= common.QuorumSize {
			log.Infoln("Received an out of bounds index")
			returnCode = 9
			return
		}

		// Verify that the signatory is a participant in the quorum
		if s.participants[signatory] == nil {
			log.Infoln("Received a heartbeat signed by an invalid signatory")
			returnCode = 4
			return
		}

		// Verify that the signatory has only been seen once in the current SignedHeartbeat
		if previousSignatories[signatory] {
			log.Infoln("Received a double-signed heartbeat")
			returnCode = 5
			return
		}

		// record that we've seen this signatory in the current SignedHeartbeat
		previousSignatories[signatory] = true

		// verify the signature
		signedMessage.Signature = sh.signatures[i]
		time.Sleep(15 * time.Millisecond) // dirty partial fix for a bug of terrible consequences
		verification, err := crypto.Verify(s.participants[signatory].publicKey, signedMessage)
		if err != nil {
			log.Fatalln(err)
			return
		}

		// check status of verification
		if !verification {
			log.Infoln("Received invalid signature in SignedHeartbeat")
			returnCode = 6
			return
		}

		// throwing the signature into the message here makes code cleaner in the loop
		// and after we sign it to send it to everyone else
		signedMessage.Message = signedMessage.CombinedMessage()
	}

	// Add heartbeat to list of seen heartbeats
	s.heartbeats[sh.signatories[0]][sh.heartbeatHash] = sh.heartbeat

	// Sign the stack of signatures and send it to all hosts
	signedMessage, err = crypto.Sign(s.secretKey, signedMessage.Message)
	if err != nil {
		log.Fatalln(err)
	}

	// add our signature to the signedHeartbeat
	sh.signatures = append(sh.signatures, signedMessage.Signature)
	sh.signatories = append(sh.signatories, s.participantIndex)

	// broadcast the message to the quorum
	err = s.announceSignedHeartbeat(sh)
	if err != nil {
		log.Fatalln(err)
	}

	returnCode = 0
	return
}

func (sh *signedHeartbeat) GobEncode() (gobSignedHeartbeat []byte, err error) {
	// error check the input
	if sh == nil {
		err = fmt.Errorf("Cannot encode a nil object")
		return
	}

	w := new(bytes.Buffer)
	encoder := gob.NewEncoder(w)
	err = encoder.Encode(sh.heartbeat)
	if err != nil {
		return
	}
	err = encoder.Encode(sh.heartbeatHash)
	if err != nil {
		return
	}
	err = encoder.Encode(sh.signatories)
	if err != nil {
		return
	}
	err = encoder.Encode(sh.signatures)
	if err != nil {
		return
	}

	gobSignedHeartbeat = w.Bytes()
	return
}

func (shb *signedHeartbeat) GobDecode(gobSignedHeartbeat []byte) (err error) {
	if gobSignedHeartbeat == nil {
		err = fmt.Errorf("cannot decode a nil byte slice")
		return
	}

	r := bytes.NewBuffer(gobSignedHeartbeat)
	decoder := gob.NewDecoder(r)
	err = decoder.Decode(&shb.heartbeat)
	if err != nil {
		return
	}
	err = decoder.Decode(&shb.heartbeatHash)
	if err != nil {
		return
	}
	err = decoder.Decode(&shb.signatories)
	if err != nil {
		return
	}
	err = decoder.Decode(&shb.signatures)
	if err != nil {
		return
	}

	return
}

// participants are processed in a random order each block, determied by the
// entropy for the block. participantOrdering() deterministically picks that
// order, using entropy from the state.
func (s *State) participantOrdering() (participantOrdering [common.QuorumSize]byte) {
	// create an in-order list of participants
	for i := range participantOrdering {
		participantOrdering[i] = byte(i)
	}

	// shuffle the list of participants
	for i := range participantOrdering {
		newIndex, err := s.randInt(i, common.QuorumSize)
		if err != nil {
			log.Fatalln(err)
		}
		tmp := participantOrdering[newIndex]
		participantOrdering[newIndex] = participantOrdering[i]
		participantOrdering[i] = tmp
	}

	return
}

// Removes all traces of a participant from the State
func (s *State) tossParticipant(pi byte) {
	// remove from s.Participants
	s.participants[pi] = nil

	// remove from s.PreviousEntropyStage1
	var emptyEntropy common.Entropy
	zeroHash, err := crypto.CalculateTruncatedHash(emptyEntropy[:])
	if err != nil {
		log.Fatal(err)
	}
	s.previousEntropyStage1[pi] = zeroHash

	// nil map in s.Heartbeats
	s.heartbeats[pi] = nil
}

// Update the state according to the information presented in the heartbeat
// processHeartbeat uses return codes for testing purposes
func (s *State) processHeartbeat(hb *heartbeat, i byte) int {
	// compare EntropyStage2 to the hash from the previous heartbeat
	expectedHash, err := crypto.CalculateTruncatedHash(hb.entropyStage2[:])
	if err != nil {
		log.Fatalln(err)
	}
	if expectedHash != s.previousEntropyStage1[i] {
		s.tossParticipant(i)
		return 1
	}

	print("Confirming Participant ")
	println(i)

	// Add the EntropyStage2 to UpcomingEntropy
	th, err := crypto.CalculateTruncatedHash(append(s.upcomingEntropy[:], hb.entropyStage2[:]...))
	s.upcomingEntropy = common.Entropy(th)

	// store entropyStage1 to compare with next heartbeat from this participant
	s.previousEntropyStage1[i] = hb.entropyStage1

	return 0
}

// compile() takes the list of heartbeats and uses them to advance the state.
func (s *State) compile() {
	// fetch a participant ordering
	participantOrdering := s.participantOrdering()

	// Lock down s.participants and s.heartbeats for editing
	s.participantsLock.Lock()
	s.heartbeatsLock.Lock()

	// Read heartbeats, process them, then archive them.
	for _, participant := range participantOrdering {
		if s.participants[participant] == nil {
			continue
		}

		// each participant must submit exactly 1 heartbeat
		if len(s.heartbeats[participant]) != 1 {
			s.tossParticipant(participant)
			continue
		}

		// this is the only way I know to access the only element of a map;
		// the key is unknown
		for _, hb := range s.heartbeats[participant] {
			s.processHeartbeat(hb, participant)
		}

		// archive heartbeats (unimplemented)

		// clear heartbeat list for next block
		s.heartbeats[participant] = make(map[crypto.TruncatedHash]*heartbeat)
	}

	s.participantsLock.Unlock()
	s.heartbeatsLock.Unlock()

	// move UpcomingEntropy to CurrentEntropy
	s.currentEntropy = s.upcomingEntropy

	// generate, sign, and announce new heartbeat
	hb, err := s.newHeartbeat()
	if err != nil {
		log.Fatalln(err)
	}
	shb, err := s.signHeartbeat(hb)
	if err != nil {
		log.Fatalln(err)
	}
	err = s.announceSignedHeartbeat(shb)
	if err != nil {
		log.Fatalln(err)
	}
}

// Tick() updates s.CurrentStep, and calls compile() when all steps are complete
func (s *State) tick() {
	// Every common.StepDuration, advance the state stage
	ticker := time.Tick(common.StepDuration)
	for _ = range ticker {
		s.stepLock.Lock()
		if s.currentStep == common.QuorumSize {
			println("compiling")
			s.compile()
			s.currentStep = 1
		} else {
			println("stepping")
			s.currentStep += 1
		}
		s.stepLock.Unlock()
	}
}
<commit_msg>update consensus for RPC<commit_after>package quorum

import (
	"bytes"
	"common"
	"common/crypto"
	"common/log"
	"encoding/gob"
	"fmt"
	"time"
)

// All information that needs to be passed between participants each block
type heartbeat struct {
	entropyStage1 crypto.TruncatedHash
	entropyStage2 common.Entropy
}

// Contains a heartbeat that has been signed iteratively, is a key part of the
// signed solution to the Byzantine Generals Problem
type signedHeartbeat struct {
	heartbeat     *heartbeat
	heartbeatHash crypto.TruncatedHash
	signatories   []byte             // a list of everyone who's seen the heartbeat
	signatures    []crypto.Signature // their corresponding signatures
}

// Using the current State, newHeartbeat() creates a heartbeat that fulfills all
// of the requirements of the quorum.
func (s *State) newHeartbeat() (hb *heartbeat, err error) {
	hb = new(heartbeat)

	// Fetch value used to produce EntropyStage1 in prev. heartbeat
	hb.entropyStage2 = s.storedEntropyStage2

	// Generate EntropyStage2 for next heartbeat
	entropy, err := crypto.RandomByteSlice(common.EntropyVolume)
	if err != nil {
		return
	}
	copy(s.storedEntropyStage2[:], entropy) // convert entropy from slice to byte array

	// Use EntropyStage2 to generate EntropyStage1 for this heartbeat
	hb.entropyStage1, err = crypto.CalculateTruncatedHash(s.storedEntropyStage2[:])
	if err != nil {
		return
	}

	// more code will be added here

	return
}

// Convert heartbeat to []byte
func (hb *heartbeat) GobEncode() (gobHeartbeat []byte, err error) {
	// test for bad input
	if hb == nil {
		err = fmt.Errorf("Cannot Encode a nil object")
		return
	}

	w := new(bytes.Buffer)
	encoder := gob.NewEncoder(w)
	err = encoder.Encode(hb.entropyStage1)
	if err != nil {
		return
	}
	err = encoder.Encode(hb.entropyStage2)
	if err != nil {
		return
	}

	gobHeartbeat = w.Bytes()
	return
}

// Convert []byte to heartbeat
func (hb *heartbeat) GobDecode(gobHeartbeat []byte) (err error) {
	r := bytes.NewBuffer(gobHeartbeat)
	decoder := gob.NewDecoder(r)
	err = decoder.Decode(&hb.entropyStage1)
	if err != nil {
		return
	}
	err = decoder.Decode(&hb.entropyStage2)
	if err != nil {
		return
	}
	return
}

// take new heartbeat (our own), sign it, and package it into a signedHearteat
// I'm pretty sure this only follows a newHeartbeat() call; they can be merged
func (s *State) signHeartbeat(hb *heartbeat) (sh *signedHeartbeat, err error) {
	sh = new(signedHeartbeat)

	// confirm heartbeat and hash
	sh.heartbeat = hb
	gobHb, err := hb.GobEncode()
	if err != nil {
		return
	}
	sh.heartbeatHash, err = crypto.CalculateTruncatedHash([]byte(gobHb))
	if err != nil {
		return
	}

	// fill out sigantures
	sh.signatures = make([]crypto.Signature, 1)
	signedHb, err := crypto.Sign(s.secretKey, string(sh.heartbeatHash[:]))
	if err != nil {
		return
	}
	sh.signatures[0] = signedHb.Signature
	sh.signatories = make([]byte, 1)
	sh.signatories[0] = s.self.index
	return
}

// Takes a signed heartbeat and broadcasts it to the quorum
func (s *State) announceSignedHeartbeat(sh *signedHeartbeat) (err error) {
	s.broadcast(&common.Message{
		Proc: "State.HandleSignedHeartbeat",
		Args: *sh,
		Resp: nil,
	})
	return
}

// HandleSignedHeartbeat takes the payload of an incoming message of type
// 'incomingSignedHeartbeat' and verifies it according to rules established by
// the specification.
//
// The return code is currently purely for the testing suite, the numbers
// have been chosen arbitrarily
func (s *State) HandleSignedHeartbeat(sh signedHeartbeat, returnCode *int) error {
	// Check that the slices of signatures and signatories are of the same length
	if len(sh.signatures) != len(sh.signatories) {
		log.Infoln("SignedHeartbeat has mismatched signatures")
		*returnCode = 1
		return fmt.Errorf("SignedHeartbeat has mismatched signatures")
	}

	// check that there are not too many signatures and signatories
	if len(sh.signatories) > common.QuorumSize {
		log.Infoln("Received an over-signed signedHeartbeat")
		*returnCode = 12
		return fmt.Errorf("Received an over-signed signedHeartbeat")
	}

	s.stepLock.Lock() // prevents a benign race condition; is here to follow best practices
	currentStep := s.currentStep
	s.stepLock.Unlock()
	// s.CurrentStep must be less than or equal to len(sh.Signatories), unless
	// there is a new block and s.CurrentStep is common.QuorumSize
	if currentStep > len(sh.signatories) {
		if currentStep == common.QuorumSize && len(sh.signatories) == 1 {
			// by waiting common.StepDuration, the new block will be compiled
			time.Sleep(common.StepDuration)
			// now continue to rest of function
		} else {
			log.Infoln("Received an out-of-sync SignedHeartbeat")
			*returnCode = 2
			return fmt.Errorf("Received an out-of-sync SignedHeartbeat")
		}
	}

	// Check bounds on first signatory
	if int(sh.signatories[0]) >= common.QuorumSize {
		log.Infoln("Received an out of bounds index")
		*returnCode = 9
		return fmt.Errorf("Received an out of bounds index")
	}

	// we are starting to read from memory, initiate locks
	s.participantsLock.RLock()
	s.heartbeatsLock.Lock()
	defer s.participantsLock.RUnlock()
	defer s.heartbeatsLock.Unlock()

	// check that first sigatory is a participant
	if s.participants[sh.signatories[0]] == nil {
		log.Infoln("Received heartbeat from non-participant")
		*returnCode = 10
		return fmt.Errorf("Received heartbeat from non-participant")
	}

	// Check if we have already received this heartbeat
	_, exists := s.heartbeats[sh.signatories[0]][sh.heartbeatHash]
	if exists {
		*returnCode = 8
		return nil
	}

	// Check if we already have two heartbeats from this host
	if len(s.heartbeats[sh.signatories[0]]) >= 2 {
		log.Infoln("Received many invalid heartbeats from one host")
		*returnCode = 13
		return fmt.Errorf("Received many invalid heartbeats from one host")
	}

	// iterate through the signatures and make sure each is legal
	var signedMessage crypto.SignedMessage // grows each iteration
	signedMessage.Message = string(sh.heartbeatHash[:])
	previousSignatories := make(map[byte]bool) // which signatories have already signed
	for i, signatory := range sh.signatories {
		// Check bounds on the signatory
		if int(signatory) >= common.QuorumSize {
			log.Infoln("Received an out of bounds index")
			*returnCode = 9
			return fmt.Errorf("Received an out of bounds index")
		}

		// Verify that the signatory is a participant in the quorum
		if s.participants[signatory] == nil {
			log.Infoln("Received a heartbeat signed by an invalid signatory")
			*returnCode = 4
			return fmt.Errorf("Received a heartbeat signed by an invalid signatory")
		}

		// Verify that the signatory has only been seen once in the current SignedHeartbeat
		if previousSignatories[signatory] {
			log.Infoln("Received a double-signed heartbeat")
			*returnCode = 5
			return fmt.Errorf("Received a double-signed heartbeat")
		}

		// record that we've seen this signatory in the current SignedHeartbeat
		previousSignatories[signatory] = true

		// verify the signature
		signedMessage.Signature = sh.signatures[i]
		time.Sleep(15 * time.Millisecond) // dirty partial fix for a bug of terrible consequences
		verification, err := crypto.Verify(s.participants[signatory].publicKey, signedMessage)
		if err != nil {
			log.Fatalln(err)
			return err
		}

		// check status of verification
		if !verification {
			log.Infoln("Received invalid signature in SignedHeartbeat")
			*returnCode = 6
			return fmt.Errorf("Received invalid signature in SignedHeartbeat")
		}

		// throwing the signature into the message here makes code cleaner in the loop
		// and after we sign it to send it to everyone else
		signedMessage.Message = signedMessage.CombinedMessage()
	}

	// Add heartbeat to list of seen heartbeats
	s.heartbeats[sh.signatories[0]][sh.heartbeatHash] = sh.heartbeat

	// Sign the stack of signatures and send it to all hosts
	signedMessage, err := crypto.Sign(s.secretKey, signedMessage.Message)
	if err != nil {
		log.Fatalln(err)
	}

	// add our signature to the signedHeartbeat
	sh.signatures = append(sh.signatures, signedMessage.Signature)
	sh.signatories = append(sh.signatories, s.self.index)

	// broadcast the message to the quorum
	err = s.announceSignedHeartbeat(&sh)
	if err != nil {
		log.Fatalln(err)
		return err
	}

	*returnCode = 0
	return nil
}

func (sh *signedHeartbeat) GobEncode() (gobSignedHeartbeat []byte, err error) {
	// error check the input
	if sh == nil {
		err = fmt.Errorf("Cannot encode a nil object")
		return
	}

	w := new(bytes.Buffer)
	encoder := gob.NewEncoder(w)
	err = encoder.Encode(sh.heartbeat)
	if err != nil {
		return
	}
	err = encoder.Encode(sh.heartbeatHash)
	if err != nil {
		return
	}
	err = encoder.Encode(sh.signatories)
	if err != nil {
		return
	}
	err = encoder.Encode(sh.signatures)
	if err != nil {
		return
	}

	gobSignedHeartbeat = w.Bytes()
	return
}

func (shb *signedHeartbeat) GobDecode(gobSignedHeartbeat []byte) (err error) {
	if gobSignedHeartbeat == nil {
		err = fmt.Errorf("cannot decode a nil byte slice")
		return
	}

	r := bytes.NewBuffer(gobSignedHeartbeat)
	decoder := gob.NewDecoder(r)
	err = decoder.Decode(&shb.heartbeat)
	if err != nil {
		return
	}
	err = decoder.Decode(&shb.heartbeatHash)
	if err != nil {
		return
	}
	err = decoder.Decode(&shb.signatories)
	if err != nil {
		return
	}
	err = decoder.Decode(&shb.signatures)
	if err != nil {
		return
	}

	return
}

// participants are processed in a random order each block, determied by the
// entropy for the block. participantOrdering() deterministically picks that
// order, using entropy from the state.
func (s *State) participantOrdering() (participantOrdering [common.QuorumSize]byte) {
	// create an in-order list of participants
	for i := range participantOrdering {
		participantOrdering[i] = byte(i)
	}

	// shuffle the list of participants
	for i := range participantOrdering {
		newIndex, err := s.randInt(i, common.QuorumSize)
		if err != nil {
			log.Fatalln(err)
		}
		tmp := participantOrdering[newIndex]
		participantOrdering[newIndex] = participantOrdering[i]
		participantOrdering[i] = tmp
	}

	return
}

// Removes all traces of a participant from the State
func (s *State) tossParticipant(pi byte) {
	// remove from s.Participants
	s.participants[pi] = nil

	// remove from s.PreviousEntropyStage1
	var emptyEntropy common.Entropy
	zeroHash, err := crypto.CalculateTruncatedHash(emptyEntropy[:])
	if err != nil {
		log.Fatal(err)
	}
	s.previousEntropyStage1[pi] = zeroHash

	// nil map in s.Heartbeats
	s.heartbeats[pi] = nil
}

// Update the state according to the information presented in the heartbeat
// processHeartbeat uses return codes for testing purposes
func (s *State) processHeartbeat(hb *heartbeat, i byte) int {
	// compare EntropyStage2 to the hash from the previous heartbeat
	expectedHash, err := crypto.CalculateTruncatedHash(hb.entropyStage2[:])
	if err != nil {
		log.Fatalln(err)
	}
	if expectedHash != s.previousEntropyStage1[i] {
		s.tossParticipant(i)
		return 1
	}

	print("Confirming Participant ")
	println(i)

	// Add the EntropyStage2 to UpcomingEntropy
	th, err := crypto.CalculateTruncatedHash(append(s.upcomingEntropy[:], hb.entropyStage2[:]...))
	s.upcomingEntropy = common.Entropy(th)

	// store entropyStage1 to compare with next heartbeat from this participant
	s.previousEntropyStage1[i] = hb.entropyStage1

	return 0
}

// compile() takes the list of heartbeats and uses them to advance the state.
func (s *State) compile() {
	// fetch a participant ordering
	participantOrdering := s.participantOrdering()

	// Lock down s.participants and s.heartbeats for editing
	s.participantsLock.Lock()
	s.heartbeatsLock.Lock()

	// Read heartbeats, process them, then archive them.
	for _, participant := range participantOrdering {
		if s.participants[participant] == nil {
			continue
		}

		// each participant must submit exactly 1 heartbeat
		if len(s.heartbeats[participant]) != 1 {
			s.tossParticipant(participant)
			continue
		}

		// this is the only way I know to access the only element of a map;
		// the key is unknown
		for _, hb := range s.heartbeats[participant] {
			s.processHeartbeat(hb, participant)
		}

		// archive heartbeats (unimplemented)

		// clear heartbeat list for next block
		s.heartbeats[participant] = make(map[crypto.TruncatedHash]*heartbeat)
	}

	s.participantsLock.Unlock()
	s.heartbeatsLock.Unlock()

	// move UpcomingEntropy to CurrentEntropy
	s.currentEntropy = s.upcomingEntropy

	// generate, sign, and announce new heartbeat
	hb, err := s.newHeartbeat()
	if err != nil {
		log.Fatalln(err)
	}
	shb, err := s.signHeartbeat(hb)
	if err != nil {
		log.Fatalln(err)
	}
	err = s.announceSignedHeartbeat(shb)
	if err != nil {
		log.Fatalln(err)
	}
}

// Tick() updates s.CurrentStep, and calls compile() when all steps are complete
func (s *State) tick() {
	// Every common.StepDuration, advance the state stage
	ticker := time.Tick(common.StepDuration)
	for _ = range ticker {
		s.stepLock.Lock()
		if s.currentStep == common.QuorumSize {
			println("compiling")
			s.compile()
			s.currentStep = 1
		} else {
			println("stepping")
			s.currentStep += 1
		}
		s.stepLock.Unlock()
	}
}
<|endoftext|>


<commit_before>/*
A Go reader and writer for files in the random access `dictzip` format.
*/
package dictzip

//. Imports

import (
	"bytes"
	"compress/flate"
	"fmt"
	"hash/crc32"
	"io"
	"os"
	"sync"
	"time"
)

//. Writer

/*
Levels range from 1 (BestSpeed) to 9 (BestCompression), Level 0 (NoCompression), -1 (DefaultCompression)
*/
func Write(r io.Reader, filename string, level int) error {

	const blocksize = 58315

	crc := crc32.NewIEEE()
	isize := 0

	var buf bytes.Buffer
	fw, err := flate.NewWriter(&buf, level)
	if err != nil {
		return err
	}
	sizes := make([]int, 0)
	b := make([]byte, blocksize)
	total := 0
	eof := false
	for !eof {
		n, err := readfull(r, b)
		if err != nil {
			if err != io.EOF {
				return err
			} else {
				eof = true
			}
		}
		if n > 0 {
			crc.Write(b[:n])
			isize += n

			fw.Write(b[:n])
			fw.Flush()
			fw.Reset(&buf)

			l := buf.Len()
			sizes = append(sizes, l-total)
			total = l
		}
	}
	fw.Close()

	fp, err := os.Create(filename)
	if err != nil {
		return err
	}
	defer fp.Close()

	xfl := byte(0)
	if level == flate.BestCompression {
		xfl = 2
	} else if level == flate.BestSpeed {
		xfl = 4
	}
	now := time.Now().Unix()
	_, err = fp.Write([]byte{
		31, 139, 8, 4,
		byte(now & 255), byte((now >> 8) & 255), byte((now >> 16) & 255), byte((now >> 24) & 255),
		xfl, 255})
	if err != nil {
		return err
	}

	xlen := 10 + 2*len(sizes)
	ln := 6 + 2*len(sizes)
	_, err = fp.Write([]byte{
		byte(xlen & 255), byte((xlen >> 8) & 255),
		'R', 'A', byte(ln & 255), byte((ln >> 8) & 255),
		1, 0,
		byte(blocksize & 255), byte((blocksize >> 8) & 255),
		byte(len(sizes) & 255), byte((len(sizes) >> 8) & 255)})
	if err != nil {
		return err
	}
	for _, o := range sizes {
		_, err = fp.Write([]byte{byte(o & 255), byte((o >> 8) & 255)})
		if err != nil {
			return err
		}
	}

	_, err = fp.Write(buf.Bytes())
	if err != nil {
		return err
	}

	c := crc.Sum32()
	_, err = fp.Write([]byte{
		byte(c & 255), byte((c >> 8) & 255), byte((c >> 16) & 255), byte((c >> 24) & 255),
		byte(isize & 255), byte((isize >> 8) & 255), byte((isize >> 16) & 255), byte((isize >> 24) & 255),
	})
	if err != nil {
		return err
	}

	return nil

}

//. Reader

type Reader struct {
	fp        io.ReadSeeker
	offsets   []int64
	blocksize int64
	lock      sync.Mutex
}

func NewReader(rs io.ReadSeeker) (*Reader, error) {

	dz := &Reader{fp: rs}

	_, err := dz.fp.Seek(0, 0)
	if err != nil {
		return nil, err
	}

	metadata := []byte{}

	p := 0

	h := make([]byte, 10)
	n, err := readfull(dz.fp, h)
	if err != nil {
		return nil, err
	}
	p += n

	if h[0] != 31 || h[1] != 139 {
		return nil, fmt.Errorf("Invalid header: %02X %02X\n", h[0], h[1])
	}

	if h[2] != 8 {
		return nil, fmt.Errorf("Unknown compression method:", h[2])
	}

	flg := h[3]

	if flg&4 != 0 {
		h := make([]byte, 2)
		n, err := readfull(dz.fp, h)
		if err != nil {
			return nil, err
		}
		p += n

		xlen := int(h[0]) + 256*int(h[1])
		h = make([]byte, xlen)
		n, err = readfull(dz.fp, h)
		if err != nil {
			return nil, err
		}
		p += n

		for q := 0; q < len(h); {
			si1 := h[q]
			si2 := h[q+1]
			ln := int(h[q+2]) + 256*int(h[q+3])

			if si1 == 'R' && si2 == 'A' {
				metadata = h[q+4 : q+4+ln]
			}

			q += 4 + ln
		}

	}

	// skip file name (8), file comment (16)
	for _, f := range []byte{8, 16} {
		if flg&f != 0 {
			h := make([]byte, 1)
			for {
				n, err := readfull(dz.fp, h)
				if err != nil {
					return nil, err
				}
				p += n
				if h[0] == 0 {
					break
				}
			}
		}
	}

	if flg&2 != 0 {
		h := make([]byte, 2)
		n, err := readfull(dz.fp, h)
		if err != nil {
			return nil, err
		}
		p += n
	}

	if len(metadata) < 6 {
		return nil, fmt.Errorf("Missing dictzip metadata")
	}

	version := int(metadata[0]) + 256*int(metadata[1])

	if version != 1 {
		return nil, fmt.Errorf("Unknown dictzip version:", version)
	}

	dz.blocksize = int64(metadata[2]) + 256*int64(metadata[3])
	blockcnt := int(metadata[4]) + 256*int(metadata[5])

	dz.offsets = make([]int64, blockcnt+1)
	dz.offsets[0] = int64(p)
	for i := 0; i < blockcnt; i++ {
		dz.offsets[i+1] = dz.offsets[i] + int64(metadata[6+2*i]) + 256*int64(metadata[7+2*i])
	}

	return dz, nil

}

func (dz *Reader) Get(start, size int64) ([]byte, error) {

	dz.lock.Lock()
	defer dz.lock.Unlock()

	start1 := dz.blocksize * (start / dz.blocksize)
	size1 := size + (start - start1)

	_, err := dz.fp.Seek(dz.offsets[start/dz.blocksize], 0)
	if err != nil {
		return nil, err
	}
	rd := flate.NewReader(dz.fp)

	data := make([]byte, size1)
	_, err = readfull(rd, data)
	if err != nil {
		return nil, err
	}

	return data[start-start1:], nil
}

// Using start and size in base64 notation, such as used by the dictunzip program.
func (dz *Reader) GetB64(start, size string) ([]byte, error) {
	start2, err := decode(start)
	if err != nil {
		return nil, err
	}
	size2, err := decode(size)
	if err != nil {
		return nil, err
	}
	return dz.Get(start2, size2)
}

//. Helper function

func readfull(fp io.Reader, buf []byte) (int, error) {
	ln := len(buf)
	for p := 0; p < ln; {
		n, err := fp.Read(buf[p:])
		p += n
		if err != nil {
			if err != io.EOF || p < ln {
				return p, err
			} else {
				return p, nil
			}
		}
	}
	return ln, nil
}

//. Base64 decoder

var (
	list = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/"

	index = []uint64{
		99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99,
		99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99,
		99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 62, 99, 99, 99, 63,
		52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 99, 99, 99, 99, 99, 99,
		99, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14,
		15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 99, 99, 99, 99, 99,
		99, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40,
		41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 99, 99, 99, 99, 99,
		99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99,
		99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99,
		99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99,
		99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99,
		99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99,
		99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99,
		99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99,
		99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99,
	}
)

func decode(val string) (int64, error) {
	var result uint64
	var offset uint64

	for i := len(val) - 1; i >= 0; i-- {
		tmp := index[val[i]]
		if tmp == 99 {
			return 0, fmt.Errorf("Illegal character in base64 value: %v", val[i:i+1])
		}

		if (tmp<<offset)>>offset != tmp {
			return 0, fmt.Errorf("Type uint64 cannot store decoded base64 value: %v", val)
		}

		result |= tmp << offset
		offset += 6
	}
	return int64(result), nil
}
<commit_msg>header tekst<commit_after>/*
Package dictzip provides a reader and writer for files in the random access `dictzip` format.
*/
package dictzip

//. Imports

import (
	"bytes"
	"compress/flate"
	"fmt"
	"hash/crc32"
	"io"
	"os"
	"sync"
	"time"
)

//. Writer

/*
Levels range from 1 (BestSpeed) to 9 (BestCompression), Level 0 (NoCompression), -1 (DefaultCompression)
*/
func Write(r io.Reader, filename string, level int) error {

	const blocksize = 58315

	crc := crc32.NewIEEE()
	isize := 0

	var buf bytes.Buffer
	fw, err := flate.NewWriter(&buf, level)
	if err != nil {
		return err
	}
	sizes := make([]int, 0)
	b := make([]byte, blocksize)
	total := 0
	eof := false
	for !eof {
		n, err := readfull(r, b)
		if err != nil {
			if err != io.EOF {
				return err
			} else {
				eof = true
			}
		}
		if n > 0 {
			crc.Write(b[:n])
			isize += n

			fw.Write(b[:n])
			fw.Flush()
			fw.Reset(&buf)

			l := buf.Len()
			sizes = append(sizes, l-total)
			total = l
		}
	}
	fw.Close()

	fp, err := os.Create(filename)
	if err != nil {
		return err
	}
	defer fp.Close()

	xfl := byte(0)
	if level == flate.BestCompression {
		xfl = 2
	} else if level == flate.BestSpeed {
		xfl = 4
	}
	now := time.Now().Unix()
	_, err = fp.Write([]byte{
		31, 139, 8, 4,
		byte(now & 255), byte((now >> 8) & 255), byte((now >> 16) & 255), byte((now >> 24) & 255),
		xfl, 255})
	if err != nil {
		return err
	}

	xlen := 10 + 2*len(sizes)
	ln := 6 + 2*len(sizes)
	_, err = fp.Write([]byte{
		byte(xlen & 255), byte((xlen >> 8) & 255),
		'R', 'A', byte(ln & 255), byte((ln >> 8) & 255),
		1, 0,
		byte(blocksize & 255), byte((blocksize >> 8) & 255),
		byte(len(sizes) & 255), byte((len(sizes) >> 8) & 255)})
	if err != nil {
		return err
	}
	for _, o := range sizes {
		_, err = fp.Write([]byte{byte(o & 255), byte((o >> 8) & 255)})
		if err != nil {
			return err
		}
	}

	_, err = fp.Write(buf.Bytes())
	if err != nil {
		return err
	}

	c := crc.Sum32()
	_, err = fp.Write([]byte{
		byte(c & 255), byte((c >> 8) & 255), byte((c >> 16) & 255), byte((c >> 24) & 255),
		byte(isize & 255), byte((isize >> 8) & 255), byte((isize >> 16) & 255), byte((isize >> 24) & 255),
	})
	if err != nil {
		return err
	}

	return nil

}

//. Reader

type Reader struct {
	fp        io.ReadSeeker
	offsets   []int64
	blocksize int64
	lock      sync.Mutex
}

func NewReader(rs io.ReadSeeker) (*Reader, error) {

	dz := &Reader{fp: rs}

	_, err := dz.fp.Seek(0, 0)
	if err != nil {
		return nil, err
	}

	metadata := []byte{}

	p := 0

	h := make([]byte, 10)
	n, err := readfull(dz.fp, h)
	if err != nil {
		return nil, err
	}
	p += n

	if h[0] != 31 || h[1] != 139 {
		return nil, fmt.Errorf("Invalid header: %02X %02X\n", h[0], h[1])
	}

	if h[2] != 8 {
		return nil, fmt.Errorf("Unknown compression method:", h[2])
	}

	flg := h[3]

	if flg&4 != 0 {
		h := make([]byte, 2)
		n, err := readfull(dz.fp, h)
		if err != nil {
			return nil, err
		}
		p += n

		xlen := int(h[0]) + 256*int(h[1])
		h = make([]byte, xlen)
		n, err = readfull(dz.fp, h)
		if err != nil {
			return nil, err
		}
		p += n

		for q := 0; q < len(h); {
			si1 := h[q]
			si2 := h[q+1]
			ln := int(h[q+2]) + 256*int(h[q+3])

			if si1 == 'R' && si2 == 'A' {
				metadata = h[q+4 : q+4+ln]
			}

			q += 4 + ln
		}

	}

	// skip file name (8), file comment (16)
	for _, f := range []byte{8, 16} {
		if flg&f != 0 {
			h := make([]byte, 1)
			for {
				n, err := readfull(dz.fp, h)
				if err != nil {
					return nil, err
				}
				p += n
				if h[0] == 0 {
					break
				}
			}
		}
	}

	if flg&2 != 0 {
		h := make([]byte, 2)
		n, err := readfull(dz.fp, h)
		if err != nil {
			return nil, err
		}
		p += n
	}

	if len(metadata) < 6 {
		return nil, fmt.Errorf("Missing dictzip metadata")
	}

	version := int(metadata[0]) + 256*int(metadata[1])

	if version != 1 {
		return nil, fmt.Errorf("Unknown dictzip version:", version)
	}

	dz.blocksize = int64(metadata[2]) + 256*int64(metadata[3])
	blockcnt := int(metadata[4]) + 256*int(metadata[5])

	dz.offsets = make([]int64, blockcnt+1)
	dz.offsets[0] = int64(p)
	for i := 0; i < blockcnt; i++ {
		dz.offsets[i+1] = dz.offsets[i] + int64(metadata[6+2*i]) + 256*int64(metadata[7+2*i])
	}

	return dz, nil

}

func (dz *Reader) Get(start, size int64) ([]byte, error) {

	dz.lock.Lock()
	defer dz.lock.Unlock()

	start1 := dz.blocksize * (start / dz.blocksize)
	size1 := size + (start - start1)

	_, err := dz.fp.Seek(dz.offsets[start/dz.blocksize], 0)
	if err != nil {
		return nil, err
	}
	rd := flate.NewReader(dz.fp)

	data := make([]byte, size1)
	_, err = readfull(rd, data)
	if err != nil {
		return nil, err
	}

	return data[start-start1:], nil
}

// Using start and size in base64 notation, such as used by the dictunzip program.
func (dz *Reader) GetB64(start, size string) ([]byte, error) {
	start2, err := decode(start)
	if err != nil {
		return nil, err
	}
	size2, err := decode(size)
	if err != nil {
		return nil, err
	}
	return dz.Get(start2, size2)
}

//. Helper function

func readfull(fp io.Reader, buf []byte) (int, error) {
	ln := len(buf)
	for p := 0; p < ln; {
		n, err := fp.Read(buf[p:])
		p += n
		if err != nil {
			if err != io.EOF || p < ln {
				return p, err
			} else {
				return p, nil
			}
		}
	}
	return ln, nil
}

//. Base64 decoder

var (
	list = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/"

	index = []uint64{
		99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99,
		99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99,
		99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 62, 99, 99, 99, 63,
		52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 99, 99, 99, 99, 99, 99,
		99, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14,
		15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 99, 99, 99, 99, 99,
		99, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40,
		41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 99, 99, 99, 99, 99,
		99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99,
		99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99,
		99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99,
		99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99,
		99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99,
		99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99,
		99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99,
		99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99, 99,
	}
)

func decode(val string) (int64, error) {
	var result uint64
	var offset uint64

	for i := len(val) - 1; i >= 0; i-- {
		tmp := index[val[i]]
		if tmp == 99 {
			return 0, fmt.Errorf("Illegal character in base64 value: %v", val[i:i+1])
		}

		if (tmp<<offset)>>offset != tmp {
			return 0, fmt.Errorf("Type uint64 cannot store decoded base64 value: %v", val)
		}

		result |= tmp << offset
		offset += 6
	}
	return int64(result), nil
}
<|endoftext|>


<commit_before>package env

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

type Environmenter interface {
	Set(string, string) error
	Get(string) string
	List() []string

	// Set or change the namespace. Doesn't do anything in Dev.
	SetNamespace(string) Environmenter
}

// DefaultConfig() abstracts out the api.DefaultConfig call
// entirely. This is designed so in default usecases you will
// only ever need to import this one package.
// Obviously if you require more power you have the ability to import
// github.com/hashicorp/consul/api and configure your client as needed.
func DefaultConfig() *consul.Config {
	return consul.DefaultConfig()
}
<commit_msg>Comments<commit_after>// Package Env allows developers to use the
// [consul](https://github.com/hashicorp/consul) api
// as a storage for environment variables. This package utilizes
// build flags to allow developers to degrade their application
// to a dev environment that relies on OS environment variables,
// without the need to change their source code.
//
// At the production level this package abstracts out the
// consul.KV() api and at the dev level it abstracts out
// the core `os` package.
package env // import "github.com/kyani-inc/consul/env"

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

type Environmenter interface {
	Set(string, string) error
	Get(string) string
	List() []string

	// Set or change the namespace. Doesn't do anything in Dev.
	SetNamespace(string) Environmenter
}

// DefaultConfig() abstracts out the api.DefaultConfig call
// entirely. This is designed so in default usecases you will
// only ever need to import this one package.
// Obviously if you require more power you have the ability to import
// github.com/hashicorp/consul/api and configure your client as needed.
func DefaultConfig() *consul.Config {
	return consul.DefaultConfig()
}
<|endoftext|>


<commit_before>package rwvfs

import (
	"bytes"
	"io"
	"os"
	pathpkg "path"
	"path/filepath"
	"strings"

	"golang.org/x/tools/godoc/vfs"
	"golang.org/x/tools/godoc/vfs/mapfs"
)

// Map returns a new FileSystem from the provided map. Map keys should be
// forward slash-separated pathnames and not contain a leading slash.
func Map(m map[string]string) FileSystem {
	fs := mapFS{
		m:          m,
		dirs:       map[string]struct{}{"": struct{}{}},
		FileSystem: mapfs.New(m),
	}

	// Create initial dirs.
	for path := range m {
		if err := MkdirAll(fs, filepath.Dir(path)); err != nil {
			panic(err.Error())
		}
	}

	return fs
}

type mapFS struct {
	m    map[string]string
	dirs map[string]struct{}
	vfs.FileSystem
}

func (mfs mapFS) Create(path string) (io.WriteCloser, error) {
	// Mimic behavior of OS filesystem: truncate to empty string upon creation;
	// immediately update string values with writes.
	path = filename(path)
	mfs.m[path] = ""
	return &mapFile{m: mfs.m, path: path}, nil
}

func filename(p string) string {
	if p == "." {
		return "/"
	}
	return strings.TrimPrefix(p, "/")
}

// slashdir returns path.Dir(p), but special-cases paths not beginning
// with a slash to be in the root.
func slashdir(p string) string {
	d := pathpkg.Dir(p)
	if d == "." {
		return "/"
	}
	if strings.HasPrefix(p, "/") {
		return d
	}
	return "/" + d
}

func slash(p string) string {
	if p == "." {
		return "/"
	}
	return "/" + strings.TrimPrefix(p, "/")
}

type mapFile struct {
	buf  bytes.Buffer
	m    map[string]string
	path string
}

func (f *mapFile) Write(p []byte) (int, error) {
	return f.buf.Write(p)
}

func (f *mapFile) Close() error {
	if f.m == nil {
		// duplicate closes are noop
		return nil
	}
	f.m[f.path] = f.buf.String()
	f.buf.Reset()
	f.m = nil
	return nil
}

func (mfs mapFS) lstat(p string) (os.FileInfo, error) {
	// proxy mapfs.mapFS.Lstat to not return errors for empty directories
	// created with Mkdir
	p = filename(p)
	fi, err := mfs.FileSystem.Lstat(p)
	if os.IsNotExist(err) {
		_, ok := mfs.dirs[p]
		if ok {
			return fileInfo{name: pathpkg.Base(p), dir: true}, nil
		}
	}
	return fi, err
}

func (mfs mapFS) Lstat(p string) (os.FileInfo, error) {
	fi, err := mfs.lstat(p)
	if err != nil {
		err = &os.PathError{Op: "lstat", Path: p, Err: err}
	}
	return fi, err
}

func (mfs mapFS) Stat(p string) (os.FileInfo, error) {
	fi, err := mfs.lstat(p)
	if err != nil {
		err = &os.PathError{Op: "stat", Path: p, Err: err}
	}
	return fi, err
}

func (mfs mapFS) ReadDir(p string) ([]os.FileInfo, error) {
	// proxy mapfs.mapFS.ReadDir to not return errors for empty directories
	// created with Mkdir
	p = filename(p)
	fis, err := mfs.FileSystem.ReadDir(p)
	if os.IsNotExist(err) {
		_, ok := mfs.dirs[p]
		if ok {
			// return a list of subdirs and files (the underlying ReadDir impl
			// fails here because it thinks the directories don't exist).
			fis = nil
			for dir, _ := range mfs.dirs {
				if filepath.Dir(dir) == p || (p == "" && filepath.Dir(dir) == "." && dir != "." && dir != "") {
					fis = append(fis, newDirInfo(dir))
				}
			}
			for fn, b := range mfs.m {
				if slashdir(fn) == "/"+p {
					fis = append(fis, newFileInfo(fn, b))
				}
			}
			return fis, nil
		}
	}
	return fis, err
}

func fileInfoNames(fis []os.FileInfo) []string {
	names := make([]string, len(fis))
	for i, fi := range fis {
		names[i] = fi.Name()
	}
	return names
}

func (mfs mapFS) Mkdir(name string) error {
	name = filename(name)
	if _, err := mfs.Stat(slashdir(name)); err != nil {
		return err
	}
	fi, _ := mfs.Stat(name)
	if fi != nil {
		return os.ErrExist
	}
	mfs.dirs[name] = struct{}{}
	return nil
}

func (mfs mapFS) Remove(name string) error {
	name = filename(name)
	delete(mfs.dirs, name)
	delete(mfs.m, name)
	return nil
}
<commit_msg>wrap Open in mapfs to provide nicer error msg<commit_after>package rwvfs

import (
	"bytes"
	"io"
	"os"
	pathpkg "path"
	"path/filepath"
	"strings"

	"golang.org/x/tools/godoc/vfs"
	"golang.org/x/tools/godoc/vfs/mapfs"
)

// Map returns a new FileSystem from the provided map. Map keys should be
// forward slash-separated pathnames and not contain a leading slash.
func Map(m map[string]string) FileSystem {
	fs := mapFS{
		m:          m,
		dirs:       map[string]struct{}{"": struct{}{}},
		FileSystem: mapfs.New(m),
	}

	// Create initial dirs.
	for path := range m {
		if err := MkdirAll(fs, filepath.Dir(path)); err != nil {
			panic(err.Error())
		}
	}

	return fs
}

type mapFS struct {
	m    map[string]string
	dirs map[string]struct{}
	vfs.FileSystem
}

func (mfs mapFS) Open(path string) (vfs.ReadSeekCloser, error) {
	f, err := mfs.FileSystem.Open(path)
	if err != nil {
		return nil, &os.PathError{Op: "open", Path: path, Err: err}
	}
	return f, nil
}

func (mfs mapFS) Create(path string) (io.WriteCloser, error) {
	// Mimic behavior of OS filesystem: truncate to empty string upon creation;
	// immediately update string values with writes.
	path = filename(path)
	mfs.m[path] = ""
	return &mapFile{m: mfs.m, path: path}, nil
}

func filename(p string) string {
	if p == "." {
		return "/"
	}
	return strings.TrimPrefix(p, "/")
}

// slashdir returns path.Dir(p), but special-cases paths not beginning
// with a slash to be in the root.
func slashdir(p string) string {
	d := pathpkg.Dir(p)
	if d == "." {
		return "/"
	}
	if strings.HasPrefix(p, "/") {
		return d
	}
	return "/" + d
}

func slash(p string) string {
	if p == "." {
		return "/"
	}
	return "/" + strings.TrimPrefix(p, "/")
}

type mapFile struct {
	buf  bytes.Buffer
	m    map[string]string
	path string
}

func (f *mapFile) Write(p []byte) (int, error) {
	return f.buf.Write(p)
}

func (f *mapFile) Close() error {
	if f.m == nil {
		// duplicate closes are noop
		return nil
	}
	f.m[f.path] = f.buf.String()
	f.buf.Reset()
	f.m = nil
	return nil
}

func (mfs mapFS) lstat(p string) (os.FileInfo, error) {
	// proxy mapfs.mapFS.Lstat to not return errors for empty directories
	// created with Mkdir
	p = filename(p)
	fi, err := mfs.FileSystem.Lstat(p)
	if os.IsNotExist(err) {
		_, ok := mfs.dirs[p]
		if ok {
			return fileInfo{name: pathpkg.Base(p), dir: true}, nil
		}
	}
	return fi, err
}

func (mfs mapFS) Lstat(p string) (os.FileInfo, error) {
	fi, err := mfs.lstat(p)
	if err != nil {
		err = &os.PathError{Op: "lstat", Path: p, Err: err}
	}
	return fi, err
}

func (mfs mapFS) Stat(p string) (os.FileInfo, error) {
	fi, err := mfs.lstat(p)
	if err != nil {
		err = &os.PathError{Op: "stat", Path: p, Err: err}
	}
	return fi, err
}

func (mfs mapFS) ReadDir(p string) ([]os.FileInfo, error) {
	// proxy mapfs.mapFS.ReadDir to not return errors for empty directories
	// created with Mkdir
	p = filename(p)
	fis, err := mfs.FileSystem.ReadDir(p)
	if os.IsNotExist(err) {
		_, ok := mfs.dirs[p]
		if ok {
			// return a list of subdirs and files (the underlying ReadDir impl
			// fails here because it thinks the directories don't exist).
			fis = nil
			for dir, _ := range mfs.dirs {
				if filepath.Dir(dir) == p || (p == "" && filepath.Dir(dir) == "." && dir != "." && dir != "") {
					fis = append(fis, newDirInfo(dir))
				}
			}
			for fn, b := range mfs.m {
				if slashdir(fn) == "/"+p {
					fis = append(fis, newFileInfo(fn, b))
				}
			}
			return fis, nil
		}
	}
	return fis, err
}

func fileInfoNames(fis []os.FileInfo) []string {
	names := make([]string, len(fis))
	for i, fi := range fis {
		names[i] = fi.Name()
	}
	return names
}

func (mfs mapFS) Mkdir(name string) error {
	name = filename(name)
	if _, err := mfs.Stat(slashdir(name)); err != nil {
		return err
	}
	fi, _ := mfs.Stat(name)
	if fi != nil {
		return os.ErrExist
	}
	mfs.dirs[name] = struct{}{}
	return nil
}

func (mfs mapFS) Remove(name string) error {
	name = filename(name)
	delete(mfs.dirs, name)
	delete(mfs.m, name)
	return nil
}
<|endoftext|>


<commit_before>package lz4block

import (
	"bytes"
	"strings"
	"testing"
)

func TestBlockDecode(t *testing.T) {
	appendLen := func(p []byte, size int) []byte {
		for size > 0xFF {
			p = append(p, 0xFF)
			size -= 0xFF
		}

		p = append(p, byte(size))
		return p
	}

	emitSeq := func(lit string, offset uint16, matchLen int) []byte {
		var b byte
		litLen := len(lit)
		if litLen < 15 {
			b = byte(litLen << 4)
			litLen = -1
		} else {
			b = 0xF0
			litLen -= 15
		}

		if matchLen < 4 || offset == 0 {
			out := []byte{b}
			if litLen >= 0 {
				out = appendLen(out, litLen)
			}
			return append(out, lit...)
		}

		matchLen -= 4
		if matchLen < 15 {
			b |= byte(matchLen)
			matchLen = -1
		} else {
			b |= 0x0F
			matchLen -= 15
		}

		out := []byte{b}
		if litLen >= 0 {
			out = appendLen(out, litLen)
		}

		if len(lit) > 0 {
			out = append(out, lit...)
		}

		out = append(out, byte(offset), byte(offset>>8))

		if matchLen >= 0 {
			out = appendLen(out, matchLen)
		}

		return out
	}
	concat := func(in ...[]byte) []byte {
		var p []byte
		for _, b := range in {
			p = append(p, b...)
		}
		return p
	}

	tests := []struct {
		name string
		src  []byte
		exp  []byte
	}{
		{
			"empty_input",
			[]byte{0},
			[]byte{},
		},
		{
			"literal_only_short",
			emitSeq("hello", 0, 0),
			[]byte("hello"),
		},
		{
			"literal_only_long",
			emitSeq(strings.Repeat("A", 15+255+255+1), 0, 0),
			bytes.Repeat([]byte("A"), 15+255+255+1),
		},
		{
			"literal_only_long_1",
			emitSeq(strings.Repeat("A", 15), 0, 0),
			bytes.Repeat([]byte("A"), 15),
		},
		{
			"repeat_match_len",
			emitSeq("a", 1, 4),
			[]byte("aaaaa"),
		},
		{
			"repeat_match_len_2_seq",
			concat(emitSeq("a", 1, 4), emitSeq("B", 1, 4)),
			[]byte("aaaaaBBBBB"),
		},
		{
			"long_match",
			emitSeq("A", 1, 16),
			bytes.Repeat([]byte("A"), 17),
		},
		{
			"repeat_match_log_len_2_seq",
			concat(emitSeq("a", 1, 15), emitSeq("B", 1, 15), emitSeq("end", 0, 0)),
			[]byte(strings.Repeat("a", 16) + strings.Repeat("B", 16) + "end"),
		},
		{
			"fuzz-3e4ce8cc0da392ca5a353b6ffef6d08f400ac5f9",
			[]byte("0000\x01\x00"),
			[]byte("0000000"),
		},
	}

	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			buf := make([]byte, len(test.exp))
			n := decodeBlock(buf, test.src, nil)
			if n < 0 {
				t.Log(n)
			}
			if !bytes.Equal(buf, test.exp) {
				t.Fatalf("expected %q got %q", test.exp, buf)
			}
		})
	}
}

func TestDecodeBlockInvalid(t *testing.T) {
	t.Parallel()

	for _, test := range []struct {
		name string
		src  string
		size int // Output size to try.
	}{
		{
			"empty_input",
			"",
			100,
		},
		{
			"final_lit_too_short",
			"\x20a", // litlen = 2 but only a single-byte literal
			100,
		},
		{
			"write_beyond_len_dst",
			"\x1b0\x01\x00000000000000",
			len("\x1b0\x01\x00000000000000"),
		},
	} {
		t.Run(test.name, func(t *testing.T) {
			dst := make([]byte, test.size+8)
			for i := range dst {
				dst[i] = byte(i)
			}
			dst = dst[:test.size]

			r := decodeBlock(dst, []byte(test.src), nil)
			if r >= 0 {
				t.Errorf("no error for %s", test.name)
			}

			dst = dst[:cap(dst)]
			for i := test.size; i < len(dst); i++ {
				if dst[i] != byte(i) {
					t.Error("decodeBlock wrote out of bounds")
					break
				}
			}
		})
	}
}
<commit_msg>Unit test decodeBlock with preset dictionaries<commit_after>package lz4block

import (
	"bytes"
	"strings"
	"testing"
)

func TestBlockDecode(t *testing.T) {
	appendLen := func(p []byte, size int) []byte {
		for size > 0xFF {
			p = append(p, 0xFF)
			size -= 0xFF
		}

		p = append(p, byte(size))
		return p
	}

	emitSeq := func(lit string, offset uint16, matchLen int) []byte {
		var b byte
		litLen := len(lit)
		if litLen < 15 {
			b = byte(litLen << 4)
			litLen = -1
		} else {
			b = 0xF0
			litLen -= 15
		}

		if matchLen < 4 || offset == 0 {
			out := []byte{b}
			if litLen >= 0 {
				out = appendLen(out, litLen)
			}
			return append(out, lit...)
		}

		matchLen -= 4
		if matchLen < 15 {
			b |= byte(matchLen)
			matchLen = -1
		} else {
			b |= 0x0F
			matchLen -= 15
		}

		out := []byte{b}
		if litLen >= 0 {
			out = appendLen(out, litLen)
		}

		if len(lit) > 0 {
			out = append(out, lit...)
		}

		out = append(out, byte(offset), byte(offset>>8))

		if matchLen >= 0 {
			out = appendLen(out, matchLen)
		}

		return out
	}
	concat := func(in ...[]byte) []byte {
		var p []byte
		for _, b := range in {
			p = append(p, b...)
		}
		return p
	}

	tests := []struct {
		name string
		src  []byte
		exp  []byte
	}{
		{
			"empty_input",
			[]byte{0},
			[]byte{},
		},
		{
			"literal_only_short",
			emitSeq("hello", 0, 0),
			[]byte("hello"),
		},
		{
			"literal_only_long",
			emitSeq(strings.Repeat("A", 15+255+255+1), 0, 0),
			bytes.Repeat([]byte("A"), 15+255+255+1),
		},
		{
			"literal_only_long_1",
			emitSeq(strings.Repeat("A", 15), 0, 0),
			bytes.Repeat([]byte("A"), 15),
		},
		{
			"repeat_match_len",
			emitSeq("a", 1, 4),
			[]byte("aaaaa"),
		},
		{
			"repeat_match_len_2_seq",
			concat(emitSeq("a", 1, 4), emitSeq("B", 1, 4)),
			[]byte("aaaaaBBBBB"),
		},
		{
			"long_match",
			emitSeq("A", 1, 16),
			bytes.Repeat([]byte("A"), 17),
		},
		{
			"repeat_match_log_len_2_seq",
			concat(emitSeq("a", 1, 15), emitSeq("B", 1, 15), emitSeq("end", 0, 0)),
			[]byte(strings.Repeat("a", 16) + strings.Repeat("B", 16) + "end"),
		},
		{
			"fuzz-3e4ce8cc0da392ca5a353b6ffef6d08f400ac5f9",
			[]byte("0000\x01\x00"),
			[]byte("0000000"),
		},
	}

	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			buf := make([]byte, len(test.exp))
			n := decodeBlock(buf, test.src, nil)
			if n < 0 {
				t.Log(n)
			}
			if !bytes.Equal(buf, test.exp) {
				t.Fatalf("expected %q got %q", test.exp, buf)
			}
		})
	}
}

func TestDecodeBlockInvalid(t *testing.T) {
	t.Parallel()

	for _, test := range []struct {
		name string
		src  string
		size int // Output size to try.
	}{
		{
			"empty_input",
			"",
			100,
		},
		{
			"final_lit_too_short",
			"\x20a", // litlen = 2 but only a single-byte literal
			100,
		},
		{
			"write_beyond_len_dst",
			"\x1b0\x01\x00000000000000",
			len("\x1b0\x01\x00000000000000"),
		},
	} {
		t.Run(test.name, func(t *testing.T) {
			dst := make([]byte, test.size+8)
			for i := range dst {
				dst[i] = byte(i)
			}
			dst = dst[:test.size]

			r := decodeBlock(dst, []byte(test.src), nil)
			if r >= 0 {
				t.Errorf("no error for %s", test.name)
			}

			dst = dst[:cap(dst)]
			for i := test.size; i < len(dst); i++ {
				if dst[i] != byte(i) {
					t.Error("decodeBlock wrote out of bounds")
					break
				}
			}
		})
	}
}

func TestDecodeWithDict(t *testing.T) {
	for _, c := range []struct {
		src, dict, want string // If want == "", expect an error.
	}{
		// Entire match in dictionary.
		{"\x11b\x0a\x00\x401234", "barbazquux", "barbaz1234"},

		// First part in dictionary, rest in dst.
		{"\x35foo\x09\x00\x401234", "0barbaz", "foobarbazfoo1234"},

		// Offset points before start of dictionary.
		{"\x35foo\xff\xff\x401234", "XYZ", ""},
	} {
		// Add sentinel for debugging.
		dict := []byte(c.dict + "ABCD")[:len(c.dict)]
		dst := make([]byte, 100)

		r := decodeBlock(dst, []byte(c.src), dict)

		if c.want == "" {
			if r >= 0 {
				t.Error("expected an error, got", r)
			}
		} else {
			switch {
			case r < 0:
				t.Errorf("error %d for %q", r, c.want)
			case string(dst[:r]) != c.want:
				t.Errorf("expected %q, got %q", c.want, dst[:r])
			}
		}
	}
}
<|endoftext|>


<commit_before>package provider

import (
	"fmt"

	"github.com/hashicorp/terraform/helper/schema"
	"github.com/juju/gomaasapi"
	"github.com/roblox/terraform-provider-maas/pkg/gmaw"
	"github.com/roblox/terraform-provider-maas/pkg/maas/entity"
)

// DataSubnet provides a lookup for a MaaS Subnet
func DataSubnet() *schema.Resource {
	return &schema.Resource{
		Read: dataSubnetRead,

		Schema: map[string]*schema.Schema{
			"name": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
			},
			"vlan": &schema.Schema{
				Type:     schema.TypeInt,
				Optional: true,
			},
			"cidr": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
			},
			"rdns_mode": &schema.Schema{
				Type:     schema.TypeInt,
				Computed: true,
			},
			"gateway_ip": &schema.Schema{
				Type:     schema.TypeString,
				Computed: true,
			},
			"dns_servers": &schema.Schema{
				Type:     schema.TypeList,
				Computed: true,
				Elem: &schema.Schema{
					Type: schema.TypeString,
				},
			},
		},
	}
}

func dataSubnetRead(d *schema.ResourceData, m interface{}) error {
	mo := m.(*gomaasapi.MAASObject)
	res, err := gmaw.NewSubnets(mo).Get()
	if err != nil {
		return err
	}

	for idx := range res {
		if !dataSubnetIsMatch(d, &res[idx]) {
			continue
		}
		if err := d.Set("name", res[idx].Name); err != nil {
			return err
		}
		if err := d.Set("vlan", res[idx].VLAN.ID); err != nil {
			return err
		}
		if err := d.Set("cidr", res[idx].CIDR); err != nil {
			return err
		}
		if err := d.Set("rdns_mode", res[idx].RDNSMode); err != nil {
			return err
		}
		if err := d.Set("gateway_ip", res[idx].GatewayIP); err != nil {
			return err
		}
		if err := d.Set("dns_servers", res[idx].DNSServers); err != nil {
			return err
		}
		d.SetId(fmt.Sprintf("%d", res[idx].ID))
		return nil
	}
	return fmt.Errorf("could not find matching subnet")
}

func dataSubnetIsMatch(d *schema.ResourceData, res *entity.Subnet) bool {
	name, vlan, cidr := d.Get("name").(string), d.Get("vlan").(int), d.Get("cidr").(string)
	if !(name == "" || name == res.Name) {
		return false
	}
	if !(vlan == 0 || vlan == res.VLAN.ID) {
		return false
	}
	return (cidr == "" || cidr == res.CIDR)
}
<commit_msg>Jordan actually knows go, thank goodness<commit_after>package provider

import (
	"fmt"
	"strconv"

	"github.com/hashicorp/terraform/helper/schema"
	"github.com/juju/gomaasapi"
	"github.com/roblox/terraform-provider-maas/pkg/gmaw"
	"github.com/roblox/terraform-provider-maas/pkg/maas/entity"
)

// DataSubnet provides a lookup for a MaaS Subnet
func DataSubnet() *schema.Resource {
	return &schema.Resource{
		Read: dataSubnetRead,

		Schema: map[string]*schema.Schema{
			"name": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
			},
			"vlan": &schema.Schema{
				Type:     schema.TypeInt,
				Optional: true,
			},
			"cidr": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
			},
			"rdns_mode": &schema.Schema{
				Type:     schema.TypeInt,
				Computed: true,
			},
			"gateway_ip": &schema.Schema{
				Type:     schema.TypeString,
				Computed: true,
			},
			"dns_servers": &schema.Schema{
				Type:     schema.TypeList,
				Computed: true,
				Elem: &schema.Schema{
					Type: schema.TypeString,
				},
			},
		},
	}
}

func dataSubnetRead(d *schema.ResourceData, m interface{}) error {
	mo := m.(*gomaasapi.MAASObject)
	res, err := gmaw.NewSubnets(mo).Get()
	if err != nil {
		return err
	}

	for idx := range res {
		if !dataSubnetIsMatch(d, &res[idx]) {
			continue
		}
		if err := d.Set("name", res[idx].Name); err != nil {
			return err
		}
		if err := d.Set("vlan", res[idx].VLAN.ID); err != nil {
			return err
		}
		if err := d.Set("cidr", res[idx].CIDR); err != nil {
			return err
		}
		if err := d.Set("rdns_mode", res[idx].RDNSMode); err != nil {
			return err
		}
		if err := d.Set("gateway_ip", res[idx].GatewayIP); err != nil {
			return err
		}
		if err := d.Set("dns_servers", res[idx].DNSServers); err != nil {
			return err
		}
		d.SetId(strconv.Itoa(res[idx].ID))
		return nil
	}
	return fmt.Errorf("could not find matching subnet")
}

func dataSubnetIsMatch(d *schema.ResourceData, res *entity.Subnet) bool {
	name, vlan, cidr := d.Get("name").(string), d.Get("vlan").(int), d.Get("cidr").(string)
	if !(name == "" || name == res.Name) {
		return false
	}
	if !(vlan == 0 || vlan == res.VLAN.ID) {
		return false
	}
	return (cidr == "" || cidr == res.CIDR)
}
<|endoftext|>