<commit_before>package cluster

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"math/rand"
	"net/http"
	"sync"
	"time"

	"github.com/lxc/lxd/lxd/db"
	"github.com/lxc/lxd/lxd/db/query"
	"github.com/lxc/lxd/lxd/task"
	"github.com/lxc/lxd/shared"
	log "github.com/lxc/lxd/shared/log15"
	"github.com/lxc/lxd/shared/logger"
	"github.com/pkg/errors"
)

// APIHeartbeatMember contains specific cluster node info.
type APIHeartbeatMember struct {
	ID            int64     // ID field value in nodes table.
	Address       string    // Host and Port of node.
	RaftID        uint64    // ID field value in raft_nodes table, zero if non-raft node.
	RaftRole      int       // Node role in the raft cluster, from the raft_nodes table
	Raft          bool      // Deprecated, use non-zero RaftID instead to indicate raft node.
	LastHeartbeat time.Time // Last time we received a successful response from node.
	Online        bool      // Calculated from offline threshold and LastHeatbeat time.
	updated       bool      // Has node been updated during this heartbeat run. Not sent to nodes.
}

// APIHeartbeatVersion contains max versions for all nodes in cluster.
type APIHeartbeatVersion struct {
	Schema        int
	APIExtensions int
}

// APIHeartbeat contains data sent to nodes in heartbeat.
type APIHeartbeat struct {
	sync.Mutex // Used to control access to Members maps.
	cluster    *db.Cluster
	Members    map[int64]APIHeartbeatMember
	Version    APIHeartbeatVersion
	Time       time.Time

	// Indicates if heartbeat contains a fresh set of node states.
	// This can be used to indicate to the receiving node that the state is fresh enough to
	// trigger node refresh activies (such as forkdns).
	FullStateList bool
}

// Update updates an existing APIHeartbeat struct with the raft and all node states supplied.
// If allNodes provided is an empty set then this is considered a non-full state list.
func (hbState *APIHeartbeat) Update(fullStateList bool, raftNodes []db.RaftNode, allNodes []db.NodeInfo, offlineThreshold time.Duration) {
	var maxSchemaVersion, maxAPIExtensionsVersion int

	if hbState.Members == nil {
		hbState.Members = make(map[int64]APIHeartbeatMember)
	}

	// If we've been supplied a fresh set of node states, this is a full state list.
	hbState.FullStateList = fullStateList

	raftNodeMap := make(map[string]db.RaftNode)

	// Convert raftNodes to a map keyed on address for lookups later.
	for _, raftNode := range raftNodes {
		raftNodeMap[raftNode.Address] = raftNode
	}

	// Add nodes (overwrites any nodes with same ID in map with fresh data).
	for _, node := range allNodes {
		member := APIHeartbeatMember{
			ID:            node.ID,
			Address:       node.Address,
			LastHeartbeat: node.Heartbeat,
			Online:        !node.Heartbeat.Before(time.Now().Add(-offlineThreshold)),
		}

		if raftNode, exists := raftNodeMap[member.Address]; exists {
			member.Raft = true // Deprecated
			member.RaftID = raftNode.ID
			member.RaftRole = int(raftNode.Role)
			delete(raftNodeMap, member.Address) // Used to check any remaining later.
		}

		// Add to the members map using the node ID (not the Raft Node ID).
		hbState.Members[node.ID] = member

		// Keep a record of highest APIExtensions and Schema version seen in all nodes.
		if node.APIExtensions > maxAPIExtensionsVersion {
			maxAPIExtensionsVersion = node.APIExtensions
		}

		if node.Schema > maxSchemaVersion {
			maxSchemaVersion = node.Schema
		}
	}

	hbState.Version = APIHeartbeatVersion{
		Schema:        maxSchemaVersion,
		APIExtensions: maxAPIExtensionsVersion,
	}

	if len(raftNodeMap) > 0 && hbState.cluster != nil {
		hbState.cluster.Transaction(func(tx *db.ClusterTx) error {
			for addr, raftNode := range raftNodeMap {
				_, err := tx.GetPendingNodeByAddress(addr)
				if err != nil {
					logger.Errorf("Unaccounted raft node(s) not found in 'nodes' table for heartbeat: %+v", raftNode)
				}
			}

			return nil
		})
	}

	return
}

// Send sends heartbeat requests to the nodes supplied and updates heartbeat state.
func (hbState *APIHeartbeat) Send(ctx context.Context, networkCert *shared.CertInfo, serverCert *shared.CertInfo, localAddress string, nodes []db.NodeInfo, spreadDuration time.Duration) {
	heartbeatsWg := sync.WaitGroup{}
	sendHeartbeat := func(nodeID int64, address string, spreadDuration time.Duration, heartbeatData *APIHeartbeat) {
		defer heartbeatsWg.Done()

		if spreadDuration > 0 {
			// Spread in time by waiting up to 3s less than the interval.
			spreadDurationMs := int(spreadDuration.Milliseconds())
			spreadRange := spreadDurationMs - 3000

			if spreadRange > 0 {
				time.Sleep(time.Duration(rand.Intn(spreadRange)) * time.Millisecond)
			}
		}

		// Update timestamp to current, used for time skew detection
		heartbeatData.Time = time.Now().UTC()

		err := HeartbeatNode(ctx, address, networkCert, serverCert, heartbeatData)
		if err == nil {
			heartbeatData.Lock()
			// Ensure only update nodes that exist in Members already.
			hbNode, existing := hbState.Members[nodeID]
			if !existing {
				return
			}

			hbNode.LastHeartbeat = time.Now()
			hbNode.Online = true
			hbNode.updated = true
			heartbeatData.Members[nodeID] = hbNode
			heartbeatData.Unlock()
			logger.Debug("Successful heartbeat", log.Ctx{"address": address})
		} else {
			logger.Warn("Failed heartbeat", log.Ctx{"address": address, "err": err})
		}
	}

	for _, node := range nodes {
		// Special case for the local node - just record the time now.
		if node.Address == localAddress {
			hbState.Lock()
			hbNode := hbState.Members[node.ID]
			hbNode.LastHeartbeat = time.Now()
			hbNode.Online = true
			hbNode.updated = true
			hbState.Members[node.ID] = hbNode
			hbState.Unlock()
			continue
		}

		// Parallelize the rest.
		heartbeatsWg.Add(1)
		go sendHeartbeat(node.ID, node.Address, spreadDuration, hbState)
	}

	heartbeatsWg.Wait()
}

// HeartbeatTask returns a task function that performs leader-initiated heartbeat
// checks against all LXD nodes in the cluster.
//
// It will update the heartbeat timestamp column of the nodes table
// accordingly, and also notify them of the current list of database nodes.
func HeartbeatTask(gateway *Gateway) (task.Func, task.Schedule) {
	// Since the database APIs are blocking we need to wrap the core logic
	// and run it in a goroutine, so we can abort as soon as the context expires.
	heartbeatWrapper := func(ctx context.Context) {
		ch := make(chan struct{})
		go func() {
			gateway.heartbeat(ctx, false)
			ch <- struct{}{}
		}()
		select {
		case <-ch:
		case <-ctx.Done():
		}
	}

	schedule := func() (time.Duration, error) {
		return task.Every(gateway.heartbeatInterval())()
	}

	return heartbeatWrapper, schedule
}

// heartbeatInterval returns heartbeat interval to use.
func (g *Gateway) heartbeatInterval() time.Duration {
	threshold := g.HeartbeatOfflineThreshold
	if threshold <= 0 {
		threshold = db.DefaultOfflineThreshold
	}

	return threshold / 2
}

func (g *Gateway) heartbeat(ctx context.Context, initialHeartbeat bool) {
	// Avoid concurent heartbeat loops.
	// This is possible when both the task and the out of band heartbeat
	// round from a dqlite connection both kick in at the same time.
	g.heartbeatLock.Lock()
	defer g.heartbeatLock.Unlock()

	if g.Cluster == nil || g.server == nil || g.memoryDial != nil {
		// We're not a raft node or we're not clustered
		return
	}

	raftNodes, err := g.currentRaftNodes()
	if err != nil {
		if errors.Cause(err) == ErrNotLeader {
			return
		}

		logger.Error("Failed to get current raft members", log.Ctx{"err": err})
		return
	}

	if initialHeartbeat {
		logger.Debug("Starting heartbeat round (initial)")
	} else {
		logger.Debug("Starting heartbeat round")
	}

	// Replace the local raft_nodes table immediately because it
	// might miss a row containing ourselves, since we might have
	// been elected leader before the former leader had chance to
	// send us a fresh update through the heartbeat pool.
	logger.Debug("Heartbeat updating local raft members", log.Ctx{"members": raftNodes})
	err = g.db.Transaction(func(tx *db.NodeTx) error {
		return tx.ReplaceRaftNodes(raftNodes)
	})
	if err != nil {
		logger.Warn("Failed to replace local raft members", log.Ctx{"err": err})
		return
	}

	startTime := time.Now()

	var allNodes []db.NodeInfo
	var localAddress string // Address of this node
	err = g.Cluster.Transaction(func(tx *db.ClusterTx) error {
		var err error
		allNodes, err = tx.GetNodes()
		if err != nil {
			return err
		}

		localAddress, err = tx.GetLocalNodeAddress()
		if err != nil {
			return err
		}

		return nil
	})
	if err != nil {
		logger.Warn("Failed to get current cluster members", log.Ctx{"err": err})
		return
	}

	heartbeatInterval := g.heartbeatInterval()

	// Cumulative set of node states (will be written back to database once done).
	hbState := &APIHeartbeat{cluster: g.Cluster}

	// If this leader node hasn't sent a heartbeat recently, then its node state records
	// are likely out of date, this can happen when a node becomes a leader.
	// Send stale set to all nodes in database to get a fresh set of active nodes.
	if initialHeartbeat {
		hbState.Update(false, raftNodes, allNodes, g.HeartbeatOfflineThreshold)
		hbState.Send(ctx, g.networkCert, g.serverCert(), localAddress, allNodes, 0)

		// We have the latest set of node states now, lets send that state set to all nodes.
		hbState.FullStateList = true
		hbState.Send(ctx, g.networkCert, g.serverCert(), localAddress, allNodes, 0)
	} else {
		hbState.Update(true, raftNodes, allNodes, g.HeartbeatOfflineThreshold)
		hbState.Send(ctx, g.networkCert, g.serverCert(), localAddress, allNodes, heartbeatInterval)
	}

	// Look for any new node which appeared since sending last heartbeat.
	var currentNodes []db.NodeInfo
	err = g.Cluster.Transaction(func(tx *db.ClusterTx) error {
		var err error
		currentNodes, err = tx.GetNodes()
		if err != nil {
			return err
		}

		return nil
	})
	if err != nil {
		logger.Warn("Failed to get current cluster members", log.Ctx{"err": err})
		return
	}

	newNodes := []db.NodeInfo{}
	for _, currentNode := range currentNodes {
		existing := false
		for _, node := range allNodes {
			if node.Address == currentNode.Address && node.ID == currentNode.ID {
				existing = true
				break
			}
		}

		if !existing {
			// We found a new node
			allNodes = append(allNodes, currentNode)
			newNodes = append(newNodes, currentNode)
		}
	}

	// If any new nodes found, send heartbeat to just them (with full node state).
	if len(newNodes) > 0 {
		hbState.Update(true, raftNodes, allNodes, g.HeartbeatOfflineThreshold)
		hbState.Send(ctx, g.networkCert, g.serverCert(), localAddress, newNodes, 0)
	}

	// If the context has been cancelled, return immediately.
	err = ctx.Err()
	if err != nil {
		logger.Warn("Aborting heartbeat round", log.Ctx{"err": err})
		return
	}

	err = query.Retry(func() error {
		return g.Cluster.Transaction(func(tx *db.ClusterTx) error {
			for _, node := range hbState.Members {
				if !node.updated {
					continue
				}

				err := tx.SetNodeHeartbeat(node.Address, node.LastHeartbeat)
				if err != nil && errors.Cause(err) != db.ErrNoSuchObject {
					return errors.Wrapf(err, "Failed updating heartbeat time for member %q", node.Address)
				}
			}

			return nil
		})
	})
	if err != nil {
		logger.Error("Failed updating cluster heartbeats", log.Ctx{"err": err})
		return
	}

	// If full node state was sent and node refresh task is specified, run it async.
	if g.HeartbeatNodeHook != nil {
		go g.HeartbeatNodeHook(hbState)
	}

	duration := time.Now().Sub(startTime)
	if duration > heartbeatInterval {
		logger.Warn("Heartbeat round duration greater than heartbeat interval", log.Ctx{"duration": duration, "interval": heartbeatInterval})
	}

	// Update last leader heartbeat time so next time a full node state list can be sent (if not this time).
	logger.Debug("Completed heartbeat round", log.Ctx{"duration": duration})
}

// HeartbeatNode performs a single heartbeat request against the node with the given address.
func HeartbeatNode(taskCtx context.Context, address string, networkCert *shared.CertInfo, serverCert *shared.CertInfo, heartbeatData *APIHeartbeat) error {
	logger.Debug("Sending heartbeat request", log.Ctx{"address": address})

	config, err := tlsClientConfig(networkCert, serverCert)
	if err != nil {
		return err
	}

	timeout := 2 * time.Second
	url := fmt.Sprintf("https://%s%s", address, databaseEndpoint)
	transport, cleanup := tlsTransport(config)
	defer cleanup()
	client := &http.Client{
		Transport: transport,
		Timeout:   timeout,
	}

	buffer := bytes.Buffer{}
	heartbeatData.Lock()
	err = json.NewEncoder(&buffer).Encode(heartbeatData)
	heartbeatData.Unlock()
	if err != nil {
		return err
	}

	request, err := http.NewRequest("PUT", url, bytes.NewReader(buffer.Bytes()))
	if err != nil {
		return err
	}
	setDqliteVersionHeader(request)

	// Use 1s later timeout to give HTTP client chance timeout with more useful info.
	ctx, cancel := context.WithTimeout(taskCtx, timeout+time.Second)
	defer cancel()
	request = request.WithContext(ctx)
	request.Close = true // Immediately close the connection after the request is done

	response, err := client.Do(request)
	if err != nil {
		return errors.Wrap(err, "Failed to send heartbeat request")
	}
	defer response.Body.Close()

	if response.StatusCode != http.StatusOK {
		return fmt.Errorf("Heartbeat request failed with status code: %s", response.Status)
	}

	return nil
}
<commit_msg>lxd/cluster/heartbeat: Introduces heartbeatMode type and constants for heartbeat modes<commit_after>package cluster

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"math/rand"
	"net/http"
	"sync"
	"time"

	"github.com/lxc/lxd/lxd/db"
	"github.com/lxc/lxd/lxd/db/query"
	"github.com/lxc/lxd/lxd/task"
	"github.com/lxc/lxd/shared"
	log "github.com/lxc/lxd/shared/log15"
	"github.com/lxc/lxd/shared/logger"
	"github.com/pkg/errors"
)

type heartbeatMode int

const (
	hearbeatNormal heartbeatMode = iota
	hearbeatImmediate
	hearbeatInitial
)

// APIHeartbeatMember contains specific cluster node info.
type APIHeartbeatMember struct {
	ID            int64     // ID field value in nodes table.
	Address       string    // Host and Port of node.
	RaftID        uint64    // ID field value in raft_nodes table, zero if non-raft node.
	RaftRole      int       // Node role in the raft cluster, from the raft_nodes table
	Raft          bool      // Deprecated, use non-zero RaftID instead to indicate raft node.
	LastHeartbeat time.Time // Last time we received a successful response from node.
	Online        bool      // Calculated from offline threshold and LastHeatbeat time.
	updated       bool      // Has node been updated during this heartbeat run. Not sent to nodes.
}

// APIHeartbeatVersion contains max versions for all nodes in cluster.
type APIHeartbeatVersion struct {
	Schema        int
	APIExtensions int
}

// APIHeartbeat contains data sent to nodes in heartbeat.
type APIHeartbeat struct {
	sync.Mutex // Used to control access to Members maps.
	cluster    *db.Cluster
	Members    map[int64]APIHeartbeatMember
	Version    APIHeartbeatVersion
	Time       time.Time

	// Indicates if heartbeat contains a fresh set of node states.
	// This can be used to indicate to the receiving node that the state is fresh enough to
	// trigger node refresh activies (such as forkdns).
	FullStateList bool
}

// Update updates an existing APIHeartbeat struct with the raft and all node states supplied.
// If allNodes provided is an empty set then this is considered a non-full state list.
func (hbState *APIHeartbeat) Update(fullStateList bool, raftNodes []db.RaftNode, allNodes []db.NodeInfo, offlineThreshold time.Duration) {
	var maxSchemaVersion, maxAPIExtensionsVersion int

	if hbState.Members == nil {
		hbState.Members = make(map[int64]APIHeartbeatMember)
	}

	// If we've been supplied a fresh set of node states, this is a full state list.
	hbState.FullStateList = fullStateList

	raftNodeMap := make(map[string]db.RaftNode)

	// Convert raftNodes to a map keyed on address for lookups later.
	for _, raftNode := range raftNodes {
		raftNodeMap[raftNode.Address] = raftNode
	}

	// Add nodes (overwrites any nodes with same ID in map with fresh data).
	for _, node := range allNodes {
		member := APIHeartbeatMember{
			ID:            node.ID,
			Address:       node.Address,
			LastHeartbeat: node.Heartbeat,
			Online:        !node.Heartbeat.Before(time.Now().Add(-offlineThreshold)),
		}

		if raftNode, exists := raftNodeMap[member.Address]; exists {
			member.Raft = true // Deprecated
			member.RaftID = raftNode.ID
			member.RaftRole = int(raftNode.Role)
			delete(raftNodeMap, member.Address) // Used to check any remaining later.
		}

		// Add to the members map using the node ID (not the Raft Node ID).
		hbState.Members[node.ID] = member

		// Keep a record of highest APIExtensions and Schema version seen in all nodes.
		if node.APIExtensions > maxAPIExtensionsVersion {
			maxAPIExtensionsVersion = node.APIExtensions
		}

		if node.Schema > maxSchemaVersion {
			maxSchemaVersion = node.Schema
		}
	}

	hbState.Version = APIHeartbeatVersion{
		Schema:        maxSchemaVersion,
		APIExtensions: maxAPIExtensionsVersion,
	}

	if len(raftNodeMap) > 0 && hbState.cluster != nil {
		hbState.cluster.Transaction(func(tx *db.ClusterTx) error {
			for addr, raftNode := range raftNodeMap {
				_, err := tx.GetPendingNodeByAddress(addr)
				if err != nil {
					logger.Errorf("Unaccounted raft node(s) not found in 'nodes' table for heartbeat: %+v", raftNode)
				}
			}

			return nil
		})
	}

	return
}

// Send sends heartbeat requests to the nodes supplied and updates heartbeat state.
func (hbState *APIHeartbeat) Send(ctx context.Context, networkCert *shared.CertInfo, serverCert *shared.CertInfo, localAddress string, nodes []db.NodeInfo, spreadDuration time.Duration) {
	heartbeatsWg := sync.WaitGroup{}
	sendHeartbeat := func(nodeID int64, address string, spreadDuration time.Duration, heartbeatData *APIHeartbeat) {
		defer heartbeatsWg.Done()

		if spreadDuration > 0 {
			// Spread in time by waiting up to 3s less than the interval.
			spreadDurationMs := int(spreadDuration.Milliseconds())
			spreadRange := spreadDurationMs - 3000

			if spreadRange > 0 {
				time.Sleep(time.Duration(rand.Intn(spreadRange)) * time.Millisecond)
			}
		}

		// Update timestamp to current, used for time skew detection
		heartbeatData.Time = time.Now().UTC()

		err := HeartbeatNode(ctx, address, networkCert, serverCert, heartbeatData)
		if err == nil {
			heartbeatData.Lock()
			// Ensure only update nodes that exist in Members already.
			hbNode, existing := hbState.Members[nodeID]
			if !existing {
				return
			}

			hbNode.LastHeartbeat = time.Now()
			hbNode.Online = true
			hbNode.updated = true
			heartbeatData.Members[nodeID] = hbNode
			heartbeatData.Unlock()
			logger.Debug("Successful heartbeat", log.Ctx{"address": address})
		} else {
			logger.Warn("Failed heartbeat", log.Ctx{"address": address, "err": err})
		}
	}

	for _, node := range nodes {
		// Special case for the local node - just record the time now.
		if node.Address == localAddress {
			hbState.Lock()
			hbNode := hbState.Members[node.ID]
			hbNode.LastHeartbeat = time.Now()
			hbNode.Online = true
			hbNode.updated = true
			hbState.Members[node.ID] = hbNode
			hbState.Unlock()
			continue
		}

		// Parallelize the rest.
		heartbeatsWg.Add(1)
		go sendHeartbeat(node.ID, node.Address, spreadDuration, hbState)
	}

	heartbeatsWg.Wait()
}

// HeartbeatTask returns a task function that performs leader-initiated heartbeat
// checks against all LXD nodes in the cluster.
//
// It will update the heartbeat timestamp column of the nodes table
// accordingly, and also notify them of the current list of database nodes.
func HeartbeatTask(gateway *Gateway) (task.Func, task.Schedule) {
	// Since the database APIs are blocking we need to wrap the core logic
	// and run it in a goroutine, so we can abort as soon as the context expires.
	heartbeatWrapper := func(ctx context.Context) {
		ch := make(chan struct{})
		go func() {
			gateway.heartbeat(ctx, false)
			ch <- struct{}{}
		}()
		select {
		case <-ch:
		case <-ctx.Done():
		}
	}

	schedule := func() (time.Duration, error) {
		return task.Every(gateway.heartbeatInterval())()
	}

	return heartbeatWrapper, schedule
}

// heartbeatInterval returns heartbeat interval to use.
func (g *Gateway) heartbeatInterval() time.Duration {
	threshold := g.HeartbeatOfflineThreshold
	if threshold <= 0 {
		threshold = db.DefaultOfflineThreshold
	}

	return threshold / 2
}

func (g *Gateway) heartbeat(ctx context.Context, initialHeartbeat bool) {
	// Avoid concurent heartbeat loops.
	// This is possible when both the task and the out of band heartbeat
	// round from a dqlite connection both kick in at the same time.
	g.heartbeatLock.Lock()
	defer g.heartbeatLock.Unlock()

	if g.Cluster == nil || g.server == nil || g.memoryDial != nil {
		// We're not a raft node or we're not clustered
		return
	}

	raftNodes, err := g.currentRaftNodes()
	if err != nil {
		if errors.Cause(err) == ErrNotLeader {
			return
		}

		logger.Error("Failed to get current raft members", log.Ctx{"err": err})
		return
	}

	if initialHeartbeat {
		logger.Debug("Starting heartbeat round (initial)")
	} else {
		logger.Debug("Starting heartbeat round")
	}

	// Replace the local raft_nodes table immediately because it
	// might miss a row containing ourselves, since we might have
	// been elected leader before the former leader had chance to
	// send us a fresh update through the heartbeat pool.
	logger.Debug("Heartbeat updating local raft members", log.Ctx{"members": raftNodes})
	err = g.db.Transaction(func(tx *db.NodeTx) error {
		return tx.ReplaceRaftNodes(raftNodes)
	})
	if err != nil {
		logger.Warn("Failed to replace local raft members", log.Ctx{"err": err})
		return
	}

	startTime := time.Now()

	var allNodes []db.NodeInfo
	var localAddress string // Address of this node
	err = g.Cluster.Transaction(func(tx *db.ClusterTx) error {
		var err error
		allNodes, err = tx.GetNodes()
		if err != nil {
			return err
		}

		localAddress, err = tx.GetLocalNodeAddress()
		if err != nil {
			return err
		}

		return nil
	})
	if err != nil {
		logger.Warn("Failed to get current cluster members", log.Ctx{"err": err})
		return
	}

	heartbeatInterval := g.heartbeatInterval()

	// Cumulative set of node states (will be written back to database once done).
	hbState := &APIHeartbeat{cluster: g.Cluster}

	// If this leader node hasn't sent a heartbeat recently, then its node state records
	// are likely out of date, this can happen when a node becomes a leader.
	// Send stale set to all nodes in database to get a fresh set of active nodes.
	if initialHeartbeat {
		hbState.Update(false, raftNodes, allNodes, g.HeartbeatOfflineThreshold)
		hbState.Send(ctx, g.networkCert, g.serverCert(), localAddress, allNodes, 0)

		// We have the latest set of node states now, lets send that state set to all nodes.
		hbState.FullStateList = true
		hbState.Send(ctx, g.networkCert, g.serverCert(), localAddress, allNodes, 0)
	} else {
		hbState.Update(true, raftNodes, allNodes, g.HeartbeatOfflineThreshold)
		hbState.Send(ctx, g.networkCert, g.serverCert(), localAddress, allNodes, heartbeatInterval)
	}

	// Look for any new node which appeared since sending last heartbeat.
	var currentNodes []db.NodeInfo
	err = g.Cluster.Transaction(func(tx *db.ClusterTx) error {
		var err error
		currentNodes, err = tx.GetNodes()
		if err != nil {
			return err
		}

		return nil
	})
	if err != nil {
		logger.Warn("Failed to get current cluster members", log.Ctx{"err": err})
		return
	}

	newNodes := []db.NodeInfo{}
	for _, currentNode := range currentNodes {
		existing := false
		for _, node := range allNodes {
			if node.Address == currentNode.Address && node.ID == currentNode.ID {
				existing = true
				break
			}
		}

		if !existing {
			// We found a new node
			allNodes = append(allNodes, currentNode)
			newNodes = append(newNodes, currentNode)
		}
	}

	// If any new nodes found, send heartbeat to just them (with full node state).
	if len(newNodes) > 0 {
		hbState.Update(true, raftNodes, allNodes, g.HeartbeatOfflineThreshold)
		hbState.Send(ctx, g.networkCert, g.serverCert(), localAddress, newNodes, 0)
	}

	// If the context has been cancelled, return immediately.
	err = ctx.Err()
	if err != nil {
		logger.Warn("Aborting heartbeat round", log.Ctx{"err": err})
		return
	}

	err = query.Retry(func() error {
		return g.Cluster.Transaction(func(tx *db.ClusterTx) error {
			for _, node := range hbState.Members {
				if !node.updated {
					continue
				}

				err := tx.SetNodeHeartbeat(node.Address, node.LastHeartbeat)
				if err != nil && errors.Cause(err) != db.ErrNoSuchObject {
					return errors.Wrapf(err, "Failed updating heartbeat time for member %q", node.Address)
				}
			}

			return nil
		})
	})
	if err != nil {
		logger.Error("Failed updating cluster heartbeats", log.Ctx{"err": err})
		return
	}

	// If full node state was sent and node refresh task is specified, run it async.
	if g.HeartbeatNodeHook != nil {
		go g.HeartbeatNodeHook(hbState)
	}

	duration := time.Now().Sub(startTime)
	if duration > heartbeatInterval {
		logger.Warn("Heartbeat round duration greater than heartbeat interval", log.Ctx{"duration": duration, "interval": heartbeatInterval})
	}

	// Update last leader heartbeat time so next time a full node state list can be sent (if not this time).
	logger.Debug("Completed heartbeat round", log.Ctx{"duration": duration})
}

// HeartbeatNode performs a single heartbeat request against the node with the given address.
func HeartbeatNode(taskCtx context.Context, address string, networkCert *shared.CertInfo, serverCert *shared.CertInfo, heartbeatData *APIHeartbeat) error {
	logger.Debug("Sending heartbeat request", log.Ctx{"address": address})

	config, err := tlsClientConfig(networkCert, serverCert)
	if err != nil {
		return err
	}

	timeout := 2 * time.Second
	url := fmt.Sprintf("https://%s%s", address, databaseEndpoint)
	transport, cleanup := tlsTransport(config)
	defer cleanup()
	client := &http.Client{
		Transport: transport,
		Timeout:   timeout,
	}

	buffer := bytes.Buffer{}
	heartbeatData.Lock()
	err = json.NewEncoder(&buffer).Encode(heartbeatData)
	heartbeatData.Unlock()
	if err != nil {
		return err
	}

	request, err := http.NewRequest("PUT", url, bytes.NewReader(buffer.Bytes()))
	if err != nil {
		return err
	}
	setDqliteVersionHeader(request)

	// Use 1s later timeout to give HTTP client chance timeout with more useful info.
	ctx, cancel := context.WithTimeout(taskCtx, timeout+time.Second)
	defer cancel()
	request = request.WithContext(ctx)
	request.Close = true // Immediately close the connection after the request is done

	response, err := client.Do(request)
	if err != nil {
		return errors.Wrap(err, "Failed to send heartbeat request")
	}
	defer response.Body.Close()

	if response.StatusCode != http.StatusOK {
		return fmt.Errorf("Heartbeat request failed with status code: %s", response.Status)
	}

	return nil
}
<|endoftext|>


<commit_before>package approvals

import (
	"bytes"
	"encoding/json"
	"encoding/xml"
	"fmt"
	"io"
	"os"
	"reflect"
	"regexp"
	"strings"

	"github.com/approvals/go-approval-tests/reporters"
	"github.com/approvals/go-approval-tests/utils"
)

var (
	defaultReporter            = reporters.NewDiffReporter()
	defaultFrontLoadedReporter = reporters.NewFrontLoadedReporter()
	defaultFolder              = ""
)

// Failable is an interface wrapper around testing.T
type Failable interface {
	Fail()
	Fatal(args ...interface{})
	Fatalf(format string, args ...interface{})
	Name() string
	Log(args ...interface{})
	Logf(format string, args ...interface{})
	Helper()
}

// VerifyWithExtension Example:
//   VerifyWithExtension(t, strings.NewReader("Hello"), ".json")
func VerifyWithExtension(t Failable, reader io.Reader, extWithDot string, opts ...verifyOptions) {
	t.Helper()
	namer := getApprovalName(t)

	if len(opts) > 0 {
		b, err := io.ReadAll(reader)
		if err != nil {
			// TODO: do something
		}

		var result string
		for _, o := range opts {
			for _, sb := range o.scrubbers {
				result = sb(string(b))
			}
		}

		reader = strings.NewReader(result)
	}

	reporter := getReporter()
	err := namer.compare(namer.getApprovalFile(extWithDot), namer.getReceivedFile(extWithDot), reader)
	if err != nil {
		reporter.Report(namer.getApprovalFile(extWithDot), namer.getReceivedFile(extWithDot))
		t.Log("Failed Approval: received does not match approved.")
		t.Fail()
	} else {
		_ = os.Remove(namer.getReceivedFile(extWithDot))
	}
}

// Verify Example:
//   Verify(t, strings.NewReader("Hello"))
func Verify(t Failable, reader io.Reader, opts ...verifyOptions) {
	t.Helper()

	if len(opts) > 1 {
		panic("Please use fluent syntax for options, see documentation for more information")
	}

	var extWithDot string
	if len(opts) == 0 || opts[0].extWithDot != "" {
		extWithDot = ".txt"
	} else {
		extWithDot = opts[0].extWithDot
	}

	namer := getApprovalName(t)

	if len(opts) > 0 {
		b, err := io.ReadAll(reader)
		if err != nil {
			// TODO: do something
		}

		var result string
		for _, o := range opts {
			for _, sb := range o.scrubbers {
				result = sb(string(b))
			}
		}

		reader = strings.NewReader(result)
	}

	reporter := getReporter()
	err := namer.compare(namer.getApprovalFile(extWithDot), namer.getReceivedFile(extWithDot), reader)
	if err != nil {
		reporter.Report(namer.getApprovalFile(extWithDot), namer.getReceivedFile(extWithDot))
		t.Log("Failed Approval: received does not match approved.")
		t.Fail()
	} else {
		_ = os.Remove(namer.getReceivedFile(extWithDot))
	}
}

// VerifyString stores the passed string into the received file and confirms
// that it matches the approved local file. On failure, it will launch a reporter.
func VerifyString(t Failable, s string, opts ...verifyOptions) {
	t.Helper()
	for _, o := range opts {
		for _, sb := range o.scrubbers {
			s = sb(s)
		}
	}
	reader := strings.NewReader(s)
	Verify(t, reader)
}

// VerifyXMLStruct Example:
//   VerifyXMLStruct(t, xml)
func VerifyXMLStruct(t Failable, obj interface{}, opts ...verifyOptions) {
	t.Helper()
	xmlContent, err := xml.MarshalIndent(obj, "", "  ")
	if err != nil {
		tip := ""
		if reflect.TypeOf(obj).Name() == "" {
			tip = "when using anonymous types be sure to include\n  XMLName xml.Name `xml:\"Your_Name_Here\"`\n"
		}
		message := fmt.Sprintf("error while pretty printing XML\n%verror:\n  %v\nXML:\n  %v\n", tip, err, obj)
		VerifyWithExtension(t, strings.NewReader(message), ".xml", opts...)
	} else {
		VerifyWithExtension(t, bytes.NewReader(xmlContent), ".xml", opts...)
	}
}

// VerifyXMLBytes Example:
//   VerifyXMLBytes(t, []byte("<Test/>"))
func VerifyXMLBytes(t Failable, bs []byte, opts ...verifyOptions) {
	t.Helper()
	type node struct {
		Attr     []xml.Attr
		XMLName  xml.Name
		Children []node `xml:",any"`
		Text     string `xml:",chardata"`
	}
	x := node{}

	err := xml.Unmarshal(bs, &x)
	if err != nil {
		message := fmt.Sprintf("error while parsing XML\nerror:\n  %s\nXML:\n  %s\n", err, string(bs))
		VerifyWithExtension(t, strings.NewReader(message), ".xml", opts...)
	} else {
		VerifyXMLStruct(t, x, opts...)
	}
}

// VerifyJSONStruct Example:
//   VerifyJSONStruct(t, json)
func VerifyJSONStruct(t Failable, obj interface{}, opts ...verifyOptions) {
	t.Helper()

	jsonb, err := json.MarshalIndent(obj, "", "  ")
	if err != nil {
		message := fmt.Sprintf("error while pretty printing JSON\nerror:\n  %s\nJSON:\n  %s\n", err, obj)
		VerifyWithExtension(t, strings.NewReader(message), ".json", opts...)
	} else {
		VerifyWithExtension(t, bytes.NewReader(jsonb), ".json", opts...)
	}
}

// VerifyJSONBytes Example:
//   VerifyJSONBytes(t, []byte("{ \"Greeting\": \"Hello\" }"))
func VerifyJSONBytes(t Failable, bs []byte, opts ...verifyOptions) {
	t.Helper()
	var obj map[string]interface{}
	err := json.Unmarshal(bs, &obj)
	if err != nil {
		message := fmt.Sprintf("error while parsing JSON\nerror:\n  %s\nJSON:\n  %s\n", err, string(bs))
		VerifyWithExtension(t, strings.NewReader(message), ".json", opts...)
	} else {
		VerifyJSONStruct(t, obj, opts...)
	}
}

// VerifyMap Example:
//   VerifyMap(t, map[string][string] { "dog": "bark" })
func VerifyMap(t Failable, m interface{}, opts ...verifyOptions) {
	t.Helper()
	outputText := utils.PrintMap(m)
	VerifyString(t, outputText, opts...)
}

// VerifyArray Example:
//   VerifyArray(t, []string{"dog", "cat"})
func VerifyArray(t Failable, array interface{}, opts ...verifyOptions) {
	t.Helper()
	outputText := utils.PrintArray(array)
	VerifyString(t, outputText, opts...)
}

// VerifyAll Example:
//   VerifyAll(t, "uppercase", []string("dog", "cat"}, func(x interface{}) string { return strings.ToUpper(x.(string)) })
func VerifyAll(t Failable, header string, collection interface{}, transform func(interface{}) string, opts ...verifyOptions) {
	t.Helper()
	if len(header) != 0 {
		header = fmt.Sprintf("%s\n\n\n", header)
	}

	outputText := header + strings.Join(utils.MapToString(collection, transform), "\n")
	VerifyString(t, outputText, opts...)
}

type reporterCloser struct {
	reporter reporters.Reporter
}

func (s *reporterCloser) Close() error {
	defaultReporter = s.reporter
	return nil
}

type frontLoadedReporterCloser struct {
	reporter reporters.Reporter
}

func (s *frontLoadedReporterCloser) Close() error {
	defaultFrontLoadedReporter = s.reporter
	return nil
}

// UseReporter configures which reporter to use on failure.
// Add at the test or method level to configure your reporter.
//
// The following examples shows how to use a reporter for all of your test cases
// in a package directory through go's setup feature.
//
//
//	func TestMain(m *testing.M) {
//		r := approvals.UseReporter(reporters.NewBeyondCompareReporter())
//		defer r.Close()
//
//		os.Exit(m.Run())
//	}
//
func UseReporter(reporter reporters.Reporter) io.Closer {
	closer := &reporterCloser{
		reporter: defaultReporter,
	}

	defaultReporter = reporter
	return closer
}

// UseFrontLoadedReporter configures reporters ahead of all other reporters to
// handle situations like CI. These reporters usually prevent reporting in
// scenarios that are headless.
func UseFrontLoadedReporter(reporter reporters.Reporter) io.Closer {
	closer := &frontLoadedReporterCloser{
		reporter: defaultFrontLoadedReporter,
	}

	defaultFrontLoadedReporter = reporter
	return closer
}

func getReporter() reporters.Reporter {
	return reporters.NewFirstWorkingReporter(
		defaultFrontLoadedReporter,
		defaultReporter,
	)
}

// UseFolder configures which folder to use to store approval files.
// By default, the approval files will be stored at the same level as the code.
//
// The following examples shows how to use the idiomatic 'testdata' folder
// for all of your test cases in a package directory.
//
//
//	func TestMain(m *testing.M) {
//		approvals.UseFolder("testdata")
//
//		os.Exit(m.Run())
//	}
//
func UseFolder(f string) {
	defaultFolder = f
}

type scrubber func(s string) string

// verifyOptions can be accessed via the approvals.Options() API enabling configuration of scrubbers
type verifyOptions struct {
	scrubbers  []scrubber
	extWithDot string
}

// Options enables providing individual Verify functions with customisations such as scrubbers
func Options() verifyOptions {
	return verifyOptions{}
}

// WithRegexScrubber allows you to 'scrub' dynamic data such as timestamps within your test input
// and replace it with a static placeholder
func (v verifyOptions) WithRegexScrubber(scrubber *regexp.Regexp, replacer string) verifyOptions {
	v.scrubbers = append(v.scrubbers, func(s string) string {
		return scrubber.ReplaceAllString(s, replacer)
	})
	return v
}

// WithExtension overrides the default file extension (.txt) for approval files.
func (v verifyOptions) WithExtension(extension string) verifyOptions {
	v.extWithDot = extension
	return v
}
<commit_msg>Deprecate VerifyWithExtension and use Verify instead<commit_after>package approvals

import (
	"bytes"
	"encoding/json"
	"encoding/xml"
	"fmt"
	"io"
	"os"
	"reflect"
	"regexp"
	"strings"

	"github.com/approvals/go-approval-tests/reporters"
	"github.com/approvals/go-approval-tests/utils"
)

var (
	defaultReporter            = reporters.NewDiffReporter()
	defaultFrontLoadedReporter = reporters.NewFrontLoadedReporter()
	defaultFolder              = ""
)

// Failable is an interface wrapper around testing.T
type Failable interface {
	Fail()
	Fatal(args ...interface{})
	Fatalf(format string, args ...interface{})
	Name() string
	Log(args ...interface{})
	Logf(format string, args ...interface{})
	Helper()
}

// VerifyWithExtension Example:
//   VerifyWithExtension(t, strings.NewReader("Hello"), ".json")
// Deprecated: Please use Verify with the Options() fluent syntax.
func VerifyWithExtension(t Failable, reader io.Reader, extWithDot string, opts ...verifyOptions) {
	t.Helper()

	var v verifyOptions
	if len(opts) == 0 {
		v = Options()
	} else {
		v = opts[0]
	}

	v = v.WithExtension(extWithDot)
	Verify(t, reader, v)
}

// Verify Example:
//   Verify(t, strings.NewReader("Hello"))
func Verify(t Failable, reader io.Reader, opts ...verifyOptions) {
	t.Helper()

	if len(opts) > 1 {
		panic("Please use fluent syntax for options, see documentation for more information")
	}

	var extWithDot string
	if len(opts) == 0 || opts[0].extWithDot == "" {
		extWithDot = ".txt"
	} else {
		extWithDot = opts[0].extWithDot
	}

	namer := getApprovalName(t)

	if len(opts) > 0 {
		b, err := io.ReadAll(reader)
		if err != nil {
			// TODO: do something
		}

		result := string(b)
		for _, o := range opts {
			for _, sb := range o.scrubbers {
				result = sb(result)
			}
		}

		reader = strings.NewReader(result)
	}

	reporter := getReporter()
	err := namer.compare(namer.getApprovalFile(extWithDot), namer.getReceivedFile(extWithDot), reader)
	if err != nil {
		reporter.Report(namer.getApprovalFile(extWithDot), namer.getReceivedFile(extWithDot))
		t.Log("Failed Approval: received does not match approved.")
		t.Fail()
	} else {
		_ = os.Remove(namer.getReceivedFile(extWithDot))
	}
}

// VerifyString stores the passed string into the received file and confirms
// that it matches the approved local file. On failure, it will launch a reporter.
func VerifyString(t Failable, s string, opts ...verifyOptions) {
	t.Helper()
	for _, o := range opts {
		for _, sb := range o.scrubbers {
			s = sb(s)
		}
	}
	reader := strings.NewReader(s)
	Verify(t, reader)
}

// VerifyXMLStruct Example:
//   VerifyXMLStruct(t, xml)
func VerifyXMLStruct(t Failable, obj interface{}, opts ...verifyOptions) {
	t.Helper()
	xmlContent, err := xml.MarshalIndent(obj, "", "  ")
	if err != nil {
		tip := ""
		if reflect.TypeOf(obj).Name() == "" {
			tip = "when using anonymous types be sure to include\n  XMLName xml.Name `xml:\"Your_Name_Here\"`\n"
		}
		message := fmt.Sprintf("error while pretty printing XML\n%verror:\n  %v\nXML:\n  %v\n", tip, err, obj)
		VerifyWithExtension(t, strings.NewReader(message), ".xml", opts...)
	} else {
		VerifyWithExtension(t, bytes.NewReader(xmlContent), ".xml", opts...)
	}
}

// VerifyXMLBytes Example:
//   VerifyXMLBytes(t, []byte("<Test/>"))
func VerifyXMLBytes(t Failable, bs []byte, opts ...verifyOptions) {
	t.Helper()
	type node struct {
		Attr     []xml.Attr
		XMLName  xml.Name
		Children []node `xml:",any"`
		Text     string `xml:",chardata"`
	}
	x := node{}

	err := xml.Unmarshal(bs, &x)
	if err != nil {
		message := fmt.Sprintf("error while parsing XML\nerror:\n  %s\nXML:\n  %s\n", err, string(bs))
		VerifyWithExtension(t, strings.NewReader(message), ".xml", opts...)
	} else {
		VerifyXMLStruct(t, x, opts...)
	}
}

// VerifyJSONStruct Example:
//   VerifyJSONStruct(t, json)
func VerifyJSONStruct(t Failable, obj interface{}, opts ...verifyOptions) {
	t.Helper()

	jsonb, err := json.MarshalIndent(obj, "", "  ")
	if err != nil {
		message := fmt.Sprintf("error while pretty printing JSON\nerror:\n  %s\nJSON:\n  %s\n", err, obj)
		VerifyWithExtension(t, strings.NewReader(message), ".json", opts...)
	} else {
		VerifyWithExtension(t, bytes.NewReader(jsonb), ".json", opts...)
	}
}

// VerifyJSONBytes Example:
//   VerifyJSONBytes(t, []byte("{ \"Greeting\": \"Hello\" }"))
func VerifyJSONBytes(t Failable, bs []byte, opts ...verifyOptions) {
	t.Helper()
	var obj map[string]interface{}
	err := json.Unmarshal(bs, &obj)
	if err != nil {
		message := fmt.Sprintf("error while parsing JSON\nerror:\n  %s\nJSON:\n  %s\n", err, string(bs))
		VerifyWithExtension(t, strings.NewReader(message), ".json", opts...)
	} else {
		VerifyJSONStruct(t, obj, opts...)
	}
}

// VerifyMap Example:
//   VerifyMap(t, map[string][string] { "dog": "bark" })
func VerifyMap(t Failable, m interface{}, opts ...verifyOptions) {
	t.Helper()
	outputText := utils.PrintMap(m)
	VerifyString(t, outputText, opts...)
}

// VerifyArray Example:
//   VerifyArray(t, []string{"dog", "cat"})
func VerifyArray(t Failable, array interface{}, opts ...verifyOptions) {
	t.Helper()
	outputText := utils.PrintArray(array)
	VerifyString(t, outputText, opts...)
}

// VerifyAll Example:
//   VerifyAll(t, "uppercase", []string("dog", "cat"}, func(x interface{}) string { return strings.ToUpper(x.(string)) })
func VerifyAll(t Failable, header string, collection interface{}, transform func(interface{}) string, opts ...verifyOptions) {
	t.Helper()
	if len(header) != 0 {
		header = fmt.Sprintf("%s\n\n\n", header)
	}

	outputText := header + strings.Join(utils.MapToString(collection, transform), "\n")
	VerifyString(t, outputText, opts...)
}

type reporterCloser struct {
	reporter reporters.Reporter
}

func (s *reporterCloser) Close() error {
	defaultReporter = s.reporter
	return nil
}

type frontLoadedReporterCloser struct {
	reporter reporters.Reporter
}

func (s *frontLoadedReporterCloser) Close() error {
	defaultFrontLoadedReporter = s.reporter
	return nil
}

// UseReporter configures which reporter to use on failure.
// Add at the test or method level to configure your reporter.
//
// The following examples shows how to use a reporter for all of your test cases
// in a package directory through go's setup feature.
//
//
//	func TestMain(m *testing.M) {
//		r := approvals.UseReporter(reporters.NewBeyondCompareReporter())
//		defer r.Close()
//
//		os.Exit(m.Run())
//	}
//
func UseReporter(reporter reporters.Reporter) io.Closer {
	closer := &reporterCloser{
		reporter: defaultReporter,
	}

	defaultReporter = reporter
	return closer
}

// UseFrontLoadedReporter configures reporters ahead of all other reporters to
// handle situations like CI. These reporters usually prevent reporting in
// scenarios that are headless.
func UseFrontLoadedReporter(reporter reporters.Reporter) io.Closer {
	closer := &frontLoadedReporterCloser{
		reporter: defaultFrontLoadedReporter,
	}

	defaultFrontLoadedReporter = reporter
	return closer
}

func getReporter() reporters.Reporter {
	return reporters.NewFirstWorkingReporter(
		defaultFrontLoadedReporter,
		defaultReporter,
	)
}

// UseFolder configures which folder to use to store approval files.
// By default, the approval files will be stored at the same level as the code.
//
// The following examples shows how to use the idiomatic 'testdata' folder
// for all of your test cases in a package directory.
//
//
//	func TestMain(m *testing.M) {
//		approvals.UseFolder("testdata")
//
//		os.Exit(m.Run())
//	}
//
func UseFolder(f string) {
	defaultFolder = f
}

type scrubber func(s string) string

// verifyOptions can be accessed via the approvals.Options() API enabling configuration of scrubbers
type verifyOptions struct {
	scrubbers  []scrubber
	extWithDot string
}

// Options enables providing individual Verify functions with customisations such as scrubbers
func Options() verifyOptions {
	return verifyOptions{}
}

// WithRegexScrubber allows you to 'scrub' dynamic data such as timestamps within your test input
// and replace it with a static placeholder
func (v verifyOptions) WithRegexScrubber(scrubber *regexp.Regexp, replacer string) verifyOptions {
	v.scrubbers = append(v.scrubbers, func(s string) string {
		return scrubber.ReplaceAllString(s, replacer)
	})
	return v
}

// WithExtension overrides the default file extension (.txt) for approval files.
func (v verifyOptions) WithExtension(extension string) verifyOptions {
	v.extWithDot = extension
	return v
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"os"

	"koding/kites/kloud/keycreator"

	"github.com/koding/kite/kitekey"
	"github.com/satori/go.uuid"
)

var (
	username   = flag.String("username", "", "kite.key username")
	kontrolURL = flag.String("kontrolurl", "", "kite.key kontrol URL")
	pem        = flag.String("pem", "", "private key to encrypt/decrypt kite.key")
	pub        = flag.String("pub", "", "kite.key kontrol key")
	id         = flag.String("id", "", "use fixed kite ID instead of generated one")
	file       = flag.String("file", "", "kite.key path")
)

func die(v ...interface{}) {
	fmt.Fprintln(os.Stderr, v...)
	os.Exit(1)
}

func main() {
	flag.Parse()

	var err error

	if *file != "" && *pem != "" && *pub != "" {
		err = create()
	} else {
		err = show()
	}

	if err != nil {
		die(err)
	}
}

func show() error {
	tok, err := kitekey.ParseFile(*file)
	if err != nil {
		return fmt.Errorf("reading %q failed: %s", *file, err)
	}

	return json.NewEncoder(os.Stdout).Encode(tok.Claims)
}

func create() error {
	pemBytes, err := ioutil.ReadFile(*pem)
	if err != nil {
		return fmt.Errorf("reading -pem key failed: %s", err)
	}

	pubBytes, err := ioutil.ReadFile(*pub)
	if err != nil {
		return fmt.Errorf("reading -pub key failed: %s", err)
	}

	if *id == "" {
		*id = uuid.NewV4().String()
	}

	k := &keycreator.Key{
		KontrolURL:        *kontrolURL,
		KontrolPrivateKey: string(bytes.TrimSpace(pemBytes)),
		KontrolPublicKey:  string(bytes.TrimSpace(pubBytes)),
	}

	kiteKey, err := k.Create(*username, *id)
	if err != nil {
		return fmt.Errorf("signing failed: %s", err)
	}

	if *file == "" {
		fmt.Println(kiteKey)
		return nil
	}

	return ioutil.WriteFile(*file, bytes.TrimSpace([]byte(kiteKey)), 0644)
}
<commit_msg>kitekey: add support for reading tokens with -file/-pub<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"os"

	"koding/kites/kloud/keycreator"

	"github.com/dgrijalva/jwt-go"
	"github.com/koding/kite/kitekey"
	"github.com/satori/go.uuid"
)

var (
	username   = flag.String("username", "", "kite.key username")
	kontrolURL = flag.String("kontrolurl", "", "kite.key kontrol URL")
	pem        = flag.String("pem", "", "private key to encrypt/decrypt kite.key")
	pub        = flag.String("pub", "", "kite.key kontrol key")
	id         = flag.String("id", "", "use fixed kite ID instead of generated one")
	file       = flag.String("file", "", "kite.key path")
)

func die(v ...interface{}) {
	fmt.Fprintln(os.Stderr, v...)
	os.Exit(1)
}

func main() {
	flag.Parse()

	var err error

	if *file != "" && *pem != "" && *pub != "" {
		err = create()
	} else {
		err = show()
	}

	if err != nil {
		die(err)
	}
}

func show() (err error) {
	var tok *jwt.Token

	if *pub != "" {
		key, err := ioutil.ReadFile(*file)
		if err != nil {
			return err
		}

		p, err := ioutil.ReadFile(*pub)
		if err != nil {
			return err
		}

		// try if it's a kite.key, if yes read kontrolKey
		if tok, err := jwt.Parse(string(p), kitekey.GetKontrolKey); err == nil {
			p = []byte(tok.Claims["kontrolKey"].(string))
		}

		p = bytes.TrimSpace(p)

		pubFn := func(*jwt.Token) (interface{}, error) {
			return p, nil
		}

		tok, err = jwt.Parse(string(key), pubFn)

	} else {
		tok, err = kitekey.ParseFile(*file)
	}
	if err != nil {
		return fmt.Errorf("reading %q failed: %s", *file, err)
	}

	return json.NewEncoder(os.Stdout).Encode(tok.Claims)
}

func create() error {
	pemBytes, err := ioutil.ReadFile(*pem)
	if err != nil {
		return fmt.Errorf("reading -pem key failed: %s", err)
	}

	pubBytes, err := ioutil.ReadFile(*pub)
	if err != nil {
		return fmt.Errorf("reading -pub key failed: %s", err)
	}

	if *id == "" {
		*id = uuid.NewV4().String()
	}

	k := &keycreator.Key{
		KontrolURL:        *kontrolURL,
		KontrolPrivateKey: string(bytes.TrimSpace(pemBytes)),
		KontrolPublicKey:  string(bytes.TrimSpace(pubBytes)),
	}

	kiteKey, err := k.Create(*username, *id)
	if err != nil {
		return fmt.Errorf("signing failed: %s", err)
	}

	if *file == "" {
		fmt.Println(kiteKey)
		return nil
	}

	return ioutil.WriteFile(*file, bytes.TrimSpace([]byte(kiteKey)), 0644)
}
<|endoftext|>


<commit_before>package cobbles

import (
	"io/ioutil"
	"path"
	// "fmt"
	"sort"
	"strings"
    // "reflect"
)

type Bundle struct {
	dimensions     []map[string]interface{}
	settings       map[string][]byte
	dimensionPaths map[string]string
}

type Context struct {
	Settings map[string]string
}

// Loads the given directory and returns a new Bundle.
func New(dirpath string) *Bundle {

	this := &Bundle{
		settings: map[string][]byte{},
	}

	files, err := ioutil.ReadDir(dirpath)
	if err != nil {
		panic(err)
	}

	for f := range files {
		file := files[f]
		if file.IsDir() == false {
			filepath := path.Join(dirpath, file.Name())
			if path.Base(filepath) == "dimensions.yaml" {
				this.loadDimensionsFile(filepath)
			} else if path.Ext(filepath) == ".yaml" {
				this.loadSettingsFile(filepath)
			}
		}
	}

	return this
}

// Loads a dimensions YAML file into the Bundle.
func (this *Bundle) loadDimensionsFile(filepath string) {
    if this.dimensions != nil {
        panic("A dimensions file has already been loaded.")
    }
	b, err := ioutil.ReadFile(filepath)
	if err != nil {
		panic(err)
	}
	this.dimensions = make([]map[string]interface{}, 4)
	fromYaml(b, &this.dimensions)
}

// Loads a settings YAML file into the Bundle.
func (this *Bundle) loadSettingsFile(filepath string) {
	settings, err := ioutil.ReadFile(filepath)
	if err != nil {
		panic(err)
	}
	name := path.Base(filepath)
	key := this.makeLookupPath(name[:strings.Index(name, ".")])
	if _, ok := this.settings[key]; ok {
		panic("A Settings group with the key [" + key + "] has already loaded.")
	}
	this.settings[key] = settings
}

// Returns a slice of ordered lookup strings.
func (this *Bundle) getLookupPaths(context string) []string {
    return []string{}
}

// Takes the unmarshalled YAML "dimensions" slice and returns it as a flattened map.
func (this *Bundle) flattenDimensions(dimensions []map[string]interface{}) map[string]string {

	build := make(map[string]string)

	for _, set := range dimensions {
		for key, dimension := range set {
			this.flattenDimension(key, dimension.(map[interface{}]interface{}), build)
		}
	}

	return build
}

// Takes an unmarshalled YAML "dimensions" map and adds it as a flattened string to the given "build" map.
func (this *Bundle) flattenDimension(prefix string, dimension map[interface{}]interface{}, build map[string]string) {

    for k, nextDimension := range dimension {
        key := k.(string)
        newPrefix := key
        if prefix != "" {
            newPrefix = prefix + "/" + key
        }
        build[newPrefix] = key
        // If nextDimension is a value and is a map.
        if nextDimension != nil {
            this.flattenDimension(newPrefix, nextDimension.(map[interface{}]interface{}), build);
        }
    }
}

// Takes the given context string and returns it as an ordered lookup path.
func (this *Bundle) makeLookupPath(context string) string {
	if context == "master" {
		return "master"
	}
	parts := strings.Split(context, ",")
	sort.Strings(parts)
	return strings.Join(parts, ",")
}

// Reads the configuration for the given context into the given configuration interface.
func (this *Bundle) Read(context string, config interface{}) {
    lookupPaths := this.getLookupPaths(context)
    for _, path := range lookupPaths {
        if yaml, ok := this.settings[path]; ok {
            fromYaml(yaml, config)
        }
    }
}

// Returns the bundle as a YAML byte slice.
func (this *Bundle) Debug() []byte {
	type dump struct {
		Dimensions []map[string]interface{}
		Settings   map[string]interface{}
	}
	d := dump{this.dimensions, map[string]interface{}{}}
	for k, v := range this.settings {
		var i interface{}
		fromYaml(v, &i)
		d.Settings[k] = i
	}
	return toYaml(d)
}
<commit_msg>Adding placeholder methods<commit_after>package cobbles

import (
	"io/ioutil"
	"path"
	// "fmt"
	"sort"
	"strings"
    // "reflect"
)

type Bundle struct {
	dimensions     []map[string]interface{}
	settings       map[string][]byte
	dimensionPaths map[string]string
}

type Context struct {
	Settings map[string]string
}

// Loads the given directory and returns a new Bundle.
func New(dirpath string) *Bundle {

	this := &Bundle{
		settings: map[string][]byte{},
	}

	files, err := ioutil.ReadDir(dirpath)
	if err != nil {
		panic(err)
	}

	for f := range files {
		file := files[f]
		if file.IsDir() == false {
			filepath := path.Join(dirpath, file.Name())
			if path.Base(filepath) == "dimensions.yaml" {
				this.loadDimensionsFile(filepath)
			} else if path.Ext(filepath) == ".yaml" {
				this.loadSettingsFile(filepath)
			}
		}
	}

	return this
}

// Loads a dimensions YAML file into the Bundle.
func (this *Bundle) loadDimensionsFile(filepath string) {
    if this.dimensions != nil {
        panic("A dimensions file has already been loaded.")
    }
	b, err := ioutil.ReadFile(filepath)
	if err != nil {
		panic(err)
	}
	this.dimensions = make([]map[string]interface{}, 4)
	fromYaml(b, &this.dimensions)
}

// Loads a settings YAML file into the Bundle.
func (this *Bundle) loadSettingsFile(filepath string) {
	settings, err := ioutil.ReadFile(filepath)
	if err != nil {
		panic(err)
	}
	name := path.Base(filepath)
	key := this.makeLookupPath(name[:strings.Index(name, ".")])
	if _, ok := this.settings[key]; ok {
		panic("A Settings group with the key [" + key + "] has already loaded.")
	}
	this.settings[key] = settings
}

// Returns a slice of ordered lookup strings.
func (this *Bundle) getLookupPaths(context string) []string {
    return []string{}
}

// Replaces any substitutions found in the final configuration.
func (this *Bundle) applySubstitutions(config interface{}) {
    // ...
}

// Takes the unmarshalled YAML "dimensions" slice and returns it as a flattened map.
func (this *Bundle) flattenDimensions(dimensions []map[string]interface{}) map[string]string {

	build := make(map[string]string)

	for _, set := range dimensions {
		for key, dimension := range set {
			this.flattenDimension(key, dimension.(map[interface{}]interface{}), build)
		}
	}

	return build
}

// Takes an unmarshalled YAML "dimensions" map and adds it as a flattened string to the given "build" map.
func (this *Bundle) flattenDimension(prefix string, dimension map[interface{}]interface{}, build map[string]string) {

    for k, nextDimension := range dimension {
        key := k.(string)
        newPrefix := key
        if prefix != "" {
            newPrefix = prefix + "/" + key
        }
        build[newPrefix] = key
        // If nextDimension is a value and is a map.
        if nextDimension != nil {
            this.flattenDimension(newPrefix, nextDimension.(map[interface{}]interface{}), build);
        }
    }
}

// Takes the given context string and returns it as an ordered lookup path.
func (this *Bundle) makeLookupPath(context string) string {
	if context == "master" {
		return "master"
	}
	parts := strings.Split(context, ",")
	sort.Strings(parts)
	return strings.Join(parts, ",")
}

// Reads the configuration for the given context into the given configuration interface.
func (this *Bundle) Read(context string, config interface{}) {
    lookupPaths := this.getLookupPaths(context)
    for _, path := range lookupPaths {
        if yaml, ok := this.settings[path]; ok {
            fromYaml(yaml, config)
        }
    }
    this.applySubstitutions(config)
}

// Returns the bundle as a YAML byte slice.
func (this *Bundle) Debug() []byte {
	type dump struct {
		Dimensions []map[string]interface{}
		Settings   map[string]interface{}
	}
	d := dump{this.dimensions, map[string]interface{}{}}
	for k, v := range this.settings {
		var i interface{}
		fromYaml(v, &i)
		d.Settings[k] = i
	}
	return toYaml(d)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io"
	"os"
	"strings"
	"time"

	"github.com/reconquest/prefixwriter-go"
)

var (
	syncProtocolPrefix      = "ORGALORG"
	syncProtocolHello       = "HELLO"
	syncProtocolNode        = "NODE"
	syncProtocolNodeCurrent = "CURRENT"
	syncProtocolStart       = "START"
	syncProtocolSync        = "SYNC"
)

// syncProtocol handles SYNC protocol described in the main.go.
//
// It will handle protocol over all connected nodes.
type syncProtocol struct {
	// output represents writer, that should be connected to stdins of
	// all connected nodes.
	output io.WriteCloser

	// prefix is a unique string which prefixes every protocol message.
	prefix string
}

// newSyncProtocol returns syncProtocol instantiated with unique prefix.
func newSyncProtocol() *syncProtocol {
	return &syncProtocol{
		prefix: fmt.Sprintf(
			"%s:%d",
			syncProtocolPrefix,
			time.Now().UnixNano(),
		),
	}
}

// Init starts protocol and sends HELLO message to the writer. Specified writer
// will be used in all further communications.
func (protocol *syncProtocol) Init(output io.WriteCloser) error {
	protocol.output = prefixwriter.New(output, protocol.prefix+" ")

	_, err := io.WriteString(
		protocol.output,
		syncProtocolHello+"\n",
	)
	if err != nil {
		return protocolSuspendEOF(err)
	}

	return nil
}

// SendNode sends to the writer serialized representation of specified node as
// NODE message.
func (protocol *syncProtocol) SendNode(
	current *remoteExecutionNode,
	neighbor *remoteExecutionNode,
) error {
	var line = syncProtocolNode + " " + neighbor.String()

	if current == neighbor {
		line += " " + syncProtocolNodeCurrent
	}

	_, err := io.WriteString(current.stdin, line+"\n")
	if err != nil {
		fmt.Fprintf(os.Stderr, "XXXXXX sync_protocol.go:73 err: %s\n", err)
		return protocolSuspendEOF(err)
	}

	return nil
}

// SendStart sends START message to the writer.
func (protocol *syncProtocol) SendStart() error {
	_, err := io.WriteString(
		protocol.output,
		syncProtocolStart+"\n",
	)
	if err != nil {
		return protocolSuspendEOF(err)
	}

	return nil
}

// IsSyncCommand will return true,  if specified line looks like incoming
// SYNC message from the remote node.
func (protocol *syncProtocol) IsSyncCommand(line string) bool {
	return strings.HasPrefix(line, protocol.prefix+" "+syncProtocolSync)
}

// SendSync sends SYNC message to the writer, tagging it as sent from node,
// described by given source and adding optional description for the given
// SYNC phase taken by extraction it from the original SYNC message, sent
// by node.
func (protocol *syncProtocol) SendSync(
	source fmt.Stringer,
	sync string,
) error {
	data := strings.TrimSpace(
		strings.TrimPrefix(sync, protocol.prefix+" "+syncProtocolSync),
	)

	_, err := io.WriteString(
		protocol.output,
		syncProtocolSync+" "+source.String()+" "+data+"\n",
	)

	if err != nil {
		return protocolSuspendEOF(err)
	}

	return nil
}

// Suspend EOF for be compatible with simple commands, that are not support
// protocol, and therefore can close exit earlier, than protocol is initiated.
func protocolSuspendEOF(err error) error {
	if err == io.EOF {
		return nil
	}

	return err
}
<commit_msg>xxx: remove debug message<commit_after>package main

import (
	"fmt"
	"io"
	"strings"
	"time"

	"github.com/reconquest/prefixwriter-go"
)

var (
	syncProtocolPrefix      = "ORGALORG"
	syncProtocolHello       = "HELLO"
	syncProtocolNode        = "NODE"
	syncProtocolNodeCurrent = "CURRENT"
	syncProtocolStart       = "START"
	syncProtocolSync        = "SYNC"
)

// syncProtocol handles SYNC protocol described in the main.go.
//
// It will handle protocol over all connected nodes.
type syncProtocol struct {
	// output represents writer, that should be connected to stdins of
	// all connected nodes.
	output io.WriteCloser

	// prefix is a unique string which prefixes every protocol message.
	prefix string
}

// newSyncProtocol returns syncProtocol instantiated with unique prefix.
func newSyncProtocol() *syncProtocol {
	return &syncProtocol{
		prefix: fmt.Sprintf(
			"%s:%d",
			syncProtocolPrefix,
			time.Now().UnixNano(),
		),
	}
}

// Init starts protocol and sends HELLO message to the writer. Specified writer
// will be used in all further communications.
func (protocol *syncProtocol) Init(output io.WriteCloser) error {
	protocol.output = prefixwriter.New(output, protocol.prefix+" ")

	_, err := io.WriteString(
		protocol.output,
		syncProtocolHello+"\n",
	)
	if err != nil {
		return protocolSuspendEOF(err)
	}

	return nil
}

// SendNode sends to the writer serialized representation of specified node as
// NODE message.
func (protocol *syncProtocol) SendNode(
	current *remoteExecutionNode,
	neighbor *remoteExecutionNode,
) error {
	var line = syncProtocolNode + " " + neighbor.String()

	if current == neighbor {
		line += " " + syncProtocolNodeCurrent
	}

	_, err := io.WriteString(current.stdin, line+"\n")
	if err != nil {
		return protocolSuspendEOF(err)
	}

	return nil
}

// SendStart sends START message to the writer.
func (protocol *syncProtocol) SendStart() error {
	_, err := io.WriteString(
		protocol.output,
		syncProtocolStart+"\n",
	)
	if err != nil {
		return protocolSuspendEOF(err)
	}

	return nil
}

// IsSyncCommand will return true,  if specified line looks like incoming
// SYNC message from the remote node.
func (protocol *syncProtocol) IsSyncCommand(line string) bool {
	return strings.HasPrefix(line, protocol.prefix+" "+syncProtocolSync)
}

// SendSync sends SYNC message to the writer, tagging it as sent from node,
// described by given source and adding optional description for the given
// SYNC phase taken by extraction it from the original SYNC message, sent
// by node.
func (protocol *syncProtocol) SendSync(
	source fmt.Stringer,
	sync string,
) error {
	data := strings.TrimSpace(
		strings.TrimPrefix(sync, protocol.prefix+" "+syncProtocolSync),
	)

	_, err := io.WriteString(
		protocol.output,
		syncProtocolSync+" "+source.String()+" "+data+"\n",
	)

	if err != nil {
		return protocolSuspendEOF(err)
	}

	return nil
}

// Suspend EOF for be compatible with simple commands, that are not support
// protocol, and therefore can close exit earlier, than protocol is initiated.
func protocolSuspendEOF(err error) error {
	if err == io.EOF {
		return nil
	}

	return err
}
<|endoftext|>


<commit_before>package taggolib

import (
	"bytes"
	"io/ioutil"
	"reflect"
	"testing"
)

// TestNew verifies that New creates the proper parser for an example input stream
func TestNew(t *testing.T) {
	// Read in test files
	flacFile, err := ioutil.ReadFile("./test/tone16bit.flac")
	if err != nil {
		t.Fatalf("Could not open test FLAC: %v", err)
	}

	// Table of tests
	var tests = []struct {
		stream     []byte
		parser     Parser
		err        error
		tags       []string
		properties []int
	}{
		// Check for FLAC file, with hardcoded expected tags and properties
		{flacFile, &FLACParser{}, nil, []string{"Artist", "Album", "Title"}, []int{5, 202, 16, 44100}},

		// Check for an unknown format
		{[]byte("nonsense"), nil, ErrUnknownFormat, nil, nil},
	}

	// Iterate all tests
	for _, test := range tests {
		// Generate a io.ReadSeeker
		reader := bytes.NewReader(test.stream)

		// Attempt to create a parser for the reader
		parser, err := New(reader)
		if err != nil {
			// If an error occurred, check if it was expected
			if err != test.err {
				t.Fatalf("unexpected error: %v", err)
			}
		}

		// Verify that the proper parser type was created
		if reflect.TypeOf(parser) != reflect.TypeOf(test.parser) {
			t.Fatalf("unexpected parser type: %v", reflect.TypeOf(parser))
		}

		// Check for valid tags
		if test.tags != nil {
			// Artist
			if parser.Artist() != test.tags[0] {
				t.Fatalf("mismatched tag Artist: %v != %v", parser.Artist(), test.tags[0])
			}

			// Album
			if parser.Album() != test.tags[1] {
				t.Fatalf("mismatched tag Album: %v != %v", parser.Album(), test.tags[1])
			}

			// Title
			if parser.Title() != test.tags[2] {
				t.Fatalf("mismatched tag Title: %v != %v", parser.Title(), test.tags[2])
			}
		}

		// Check for valid properties
		if test.properties != nil {
			// Duration
			if int(parser.Duration().Seconds()) != test.properties[0] {
				t.Fatalf("mismatched property Duration: %v != %v", parser.Duration().Seconds(), test.properties[0])
			}

			// Bitrate
			if parser.Bitrate() != test.properties[1] {
				t.Fatalf("mismatched property Bitrate: %v != %v", parser.Bitrate(), test.properties[1])
			}

			// BitDepth
			if parser.BitDepth() != test.properties[2] {
				t.Fatalf("mismatched property BitDepth: %v != %v", parser.BitDepth(), test.properties[2])
			}

			// SampleRate
			if parser.SampleRate() != test.properties[3] {
				t.Fatalf("mismatched property SampleRate: %v != %v", parser.SampleRate(), test.properties[3])
			}
		}
	}
}
<commit_msg>Add Encoder checking to tests<commit_after>package taggolib

import (
	"bytes"
	"io/ioutil"
	"reflect"
	"testing"
)

// TestNew verifies that New creates the proper parser for an example input stream
func TestNew(t *testing.T) {
	// Read in test files
	flacFile, err := ioutil.ReadFile("./test/tone16bit.flac")
	if err != nil {
		t.Fatalf("Could not open test FLAC: %v", err)
	}

	// Table of tests
	var tests = []struct {
		stream     []byte
		parser     Parser
		err        error
		encoder    string
		tags       []string
		properties []int
	}{
		// Check for FLAC file, with hardcoded expected tags and properties
		{flacFile, &FLACParser{}, nil, "reference libFLAC 1.1.4 20070213", []string{"Artist", "Album", "Title"}, []int{5, 202, 16, 44100}},

		// Check for an unknown format
		{[]byte("nonsense"), nil, ErrUnknownFormat, "", nil, nil},
	}

	// Iterate all tests
	for _, test := range tests {
		// Generate a io.ReadSeeker
		reader := bytes.NewReader(test.stream)

		// Attempt to create a parser for the reader
		parser, err := New(reader)
		if err != nil {
			// If an error occurred, check if it was expected
			if err != test.err {
				t.Fatalf("unexpected error: %v", err)
			}
		}

		// Verify that the proper parser type was created
		if reflect.TypeOf(parser) != reflect.TypeOf(test.parser) {
			t.Fatalf("unexpected parser type: %v", reflect.TypeOf(parser))
		}

		// Discard nil parser
		if parser == nil {
			continue
		}

		// Check for valid encoder
		if parser.Encoder() != test.encoder {
			t.Fatalf("mismatched Encoder: %v != %v", parser.Encoder(), test.encoder)
		}

		// Check for valid tags
		if test.tags != nil {
			// Artist
			if parser.Artist() != test.tags[0] {
				t.Fatalf("mismatched tag Artist: %v != %v", parser.Artist(), test.tags[0])
			}

			// Album
			if parser.Album() != test.tags[1] {
				t.Fatalf("mismatched tag Album: %v != %v", parser.Album(), test.tags[1])
			}

			// Title
			if parser.Title() != test.tags[2] {
				t.Fatalf("mismatched tag Title: %v != %v", parser.Title(), test.tags[2])
			}
		}

		// Check for valid properties
		if test.properties != nil {
			// Duration
			if int(parser.Duration().Seconds()) != test.properties[0] {
				t.Fatalf("mismatched property Duration: %v != %v", parser.Duration().Seconds(), test.properties[0])
			}

			// Bitrate
			if parser.Bitrate() != test.properties[1] {
				t.Fatalf("mismatched property Bitrate: %v != %v", parser.Bitrate(), test.properties[1])
			}

			// BitDepth
			if parser.BitDepth() != test.properties[2] {
				t.Fatalf("mismatched property BitDepth: %v != %v", parser.BitDepth(), test.properties[2])
			}

			// SampleRate
			if parser.SampleRate() != test.properties[3] {
				t.Fatalf("mismatched property SampleRate: %v != %v", parser.SampleRate(), test.properties[3])
			}
		}
	}
}
<|endoftext|>


<commit_before>package auth

import (
	"github.com/gorilla/securecookie"
	"net/http"
	"encoding/gob"
	"sparkle"
	"strings"
)

const (
	authDataKey string = "Sparkle.Auth"
	authCookieName string = "Sparkle.Auth.Cookie"
)

type authData struct {
	UserIdentifier string
	Addr string
}

var sc *securecookie.SecureCookie

func init() {
	gob.Register(&authData{})
	sc = securecookie.New(securecookie.GenerateRandomKey(32), securecookie.GenerateRandomKey(32))
}

func AuthInit() {
	sparkle.AddRequestInitHook(authInitRequestHook)
}

func authInitRequestHook(w http.ResponseWriter, r *http.Request, c *sparkle.Context) error {
	cookie, err := r.Cookie(authCookieName)
	if err != nil {
		return nil
	}

	value := &authData{}
	if err := sc.Decode(authCookieName, cookie.Value, &value); err != nil {
		// Despite decoding failing, we'll assume this just means we're not authenticated
		// So no need to return an error
		return nil
	}

	c.Set(authDataKey, value)
}

func getUserIP(r *http.Request) string {
	var portSeperatorIndex = strings.Index(r.RemoteAddr, ":")
	if (portSeperatorIndex == -1) {
		return r.RemoteAddr
	} else {
		return r.RemoteAddr[:portSeperatorIndex]
	}
}

func (a *authData) isValid(r *http.Request) bool {
	return getUserIP(r) == a.Addr &&
		   a.UserIdentifier != ""
}

func (c *sparkle.Context) getAuthData() *authData {
	data = c.Get(authDataKey)

	if result, ok := data.(authData); !ok {
		return nil
	}

	return result;
}

// Returns a boolean indictating whether the current context can be considered
// authenticated
func IsAuthenticated(c *sparkle.Context) bool {
	if auth := c.Get(authDataKey); auth == nil {
		return false
	}

	return auth.isValid()
}

// Get the user identifier of the authenticated user, or nil if there is no
// authenticated user
func AuthenticatedAs(c *sparkle.Context) string {
	if auth := c.Get(authDataKey); auth == nil {
		return nil
	}

	if !auth.isValid() {
		return nil
	}

	return auth.UserIdentifier
}

// Sets an Authentication Cookie for a given path, and userIdentifier
func Authenticate(c *sparkle.Context, path string, userIdentifier string) error {
	r := c.Request()
	w := c.ResponseWriter()

	value := &authData{ userIdentifier, r.getUserIP() }
	if encoded, err := s.Encode(authCookieName, value); err != nil {
		return error
	}

	http.SetCookie(w, &http.Cookie{
		Name: authCookieName,
		Value: encoded,
		Path: path,
	})

	c.Set(authDataKey, auth)
}
<commit_msg>Auth now compiles.<commit_after>package auth

import (
	"github.com/gorilla/securecookie"
	"net/http"
	"encoding/gob"
	"sparkle"
	"strings"
)

const (
	authDataKey string = "Sparkle.Auth"
	authCookieName string = "Sparkle.Auth.Cookie"
)

type authData struct {
	UserIdentifier string
	Addr string
}

var sc *securecookie.SecureCookie

func init() {
	gob.Register(&authData{})
	sc = securecookie.New(securecookie.GenerateRandomKey(32), securecookie.GenerateRandomKey(32))
}

func AuthInit() {
	sparkle.AddRequestInitHook(authInitRequestHook)
}

func authInitRequestHook(w http.ResponseWriter, r *http.Request, c *sparkle.Context) error {
	cookie, err := r.Cookie(authCookieName)
	if err != nil {
		return nil
	}

	value := &authData{}
	if err := sc.Decode(authCookieName, cookie.Value, &value); err != nil {
		// Despite decoding failing, we'll assume this just means we're not authenticated
		// So no need to return an error
		return nil
	}

	c.Set(authDataKey, value)
	return nil
}

func getUserIP(r *http.Request) string {
	var portSeperatorIndex = strings.Index(r.RemoteAddr, ":")
	if (portSeperatorIndex == -1) {
		return r.RemoteAddr
	}

	return r.RemoteAddr[:portSeperatorIndex]
}

func (a *authData) isValid(r *http.Request) bool {
	return getUserIP(r) == a.Addr &&
		   a.UserIdentifier != ""
}

func getAuthData(c *sparkle.Context) *authData {
	data := c.Get(authDataKey)

	result, ok := data.(*authData);
	if !ok {
		return nil
	}

	return result
}

// Returns a boolean indictating whether the current context can be considered
// authenticated
func IsAuthenticated(c *sparkle.Context) bool {
	auth := getAuthData(c)
	if auth == nil {
		return false
	}

	return auth.isValid(c.Request())
}

// Get the user identifier of the authenticated user, or nil if there is no
// authenticated user
func AuthenticatedAs(c *sparkle.Context) string {
	auth := getAuthData(c)

	if auth == nil || !auth.isValid(c.Request()) {
		return ""
	}

	return auth.UserIdentifier
}

// Sets an Authentication Cookie for a given path, and userIdentifier
func Authenticate(c *sparkle.Context, path string, userIdentifier string) error {
	r := c.Request()
	w := c.ResponseWriter()

	value := &authData{ userIdentifier, getUserIP(r) }
	encoded, err := sc.Encode(authCookieName, value)
	if err != nil {
		return err
	}

	http.SetCookie(w, &http.Cookie{
		Name: authCookieName,
		Value: encoded,
		Path: path,
	})

	c.Set(authDataKey, value)
	return nil
}
<|endoftext|>


<commit_before>/*-
 * Copyright 2015 Square Inc.
 *
 * 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 main

import (
	"crypto/x509"
	"crypto/x509/pkix"
	"net"
	"testing"

	"github.com/spiffe/go-spiffe/uri"
	"github.com/stretchr/testify/assert"
)

var fakeChains = [][]*x509.Certificate{
	{
		{
			Subject: pkix.Name{
				CommonName:         "gopher",
				OrganizationalUnit: []string{"triangle", "circle"},
			},
			DNSNames:    []string{"circle"},
			IPAddresses: []net.IP{net.IPv4(192, 168, 99, 100)},
			Extensions: []pkix.Extension{
				{
					Id:       uri.OidExtensionSubjectAltName,
					Value:    getURISANFromString("spiffe://example.org/path/service"),
					Critical: false,
				}},
		},
	},
}

func getURISANFromString(s string) (uriSAN []byte) {
	uriSAN, err := uri.MarshalUriSANs([]string{s})
	panicOnError(err)

	return uriSAN
}

func TestAuthorizeNotVerified(t *testing.T) {
	*serverAllowAll = true
	*serverAllowedCNs = []string{}
	*serverAllowedOUs = []string{}
	*serverAllowedDNSs = []string{}
	*serverAllowedIPs = []net.IP{}
	*serverAllowedURIs = []string{}

	assert.NotNil(t, verifyPeerCertificateServer(nil, nil), "conn w/o cert should be rejected")
}

func TestAuthorizeReject(t *testing.T) {
	*serverAllowAll = false
	*serverAllowedCNs = []string{"test"}
	*serverAllowedOUs = []string{"test"}
	*serverAllowedDNSs = []string{"test"}
	*serverAllowedIPs = []net.IP{}
	*serverAllowedURIs = []string{"test"}

	assert.NotNil(t, verifyPeerCertificateServer(nil, fakeChains), "should reject cert w/o matching CN/OU")
}

func TestAuthorizeAllowAll(t *testing.T) {
	*serverAllowAll = true
	*serverAllowedCNs = []string{}
	*serverAllowedOUs = []string{}
	*serverAllowedDNSs = []string{}
	*serverAllowedIPs = []net.IP{}

	assert.Nil(t, verifyPeerCertificateServer(nil, fakeChains), "allow-all should always allow authed clients")
}

func TestAuthorizeAllowCN(t *testing.T) {
	*serverAllowAll = false
	*serverAllowedCNs = []string{"gopher"}
	*serverAllowedOUs = []string{}
	*serverAllowedDNSs = []string{}
	*serverAllowedIPs = []net.IP{}
	*serverAllowedURIs = []string{}

	assert.Nil(t, verifyPeerCertificateServer(nil, fakeChains), "allow-cn should allow clients with matching CN")
}

func TestAuthorizeAllowOU(t *testing.T) {
	*serverAllowAll = false
	*serverAllowedCNs = []string{}
	*serverAllowedOUs = []string{"circle"}
	*serverAllowedDNSs = []string{}
	*serverAllowedIPs = []net.IP{}
	*serverAllowedURIs = []string{}

	assert.Nil(t, verifyPeerCertificateServer(nil, fakeChains), "allow-ou should allow clients with matching OU")
}

func TestAuthorizeAllowDNS(t *testing.T) {
	*serverAllowAll = false
	*serverAllowedCNs = []string{}
	*serverAllowedOUs = []string{}
	*serverAllowedDNSs = []string{"circle"}
	*serverAllowedIPs = []net.IP{}
	*serverAllowedURIs = []string{}

	assert.Nil(t, verifyPeerCertificateServer(nil, fakeChains), "allow-dns-san should allow clients with matching DNS SAN")
}

func TestAuthorizeAllowIP(t *testing.T) {
	*serverAllowAll = false
	*serverAllowedCNs = []string{}
	*serverAllowedOUs = []string{}
	*serverAllowedDNSs = []string{}
	*serverAllowedIPs = []net.IP{net.IPv4(192, 168, 99, 100)}
	*serverAllowedURIs = []string{}

	assert.Nil(t, verifyPeerCertificateServer(nil, fakeChains), "allow-ip-san should allow clients with matching IP SAN")
}

func TestVerifyAllowCN(t *testing.T) {
	*clientAllowedCNs = []string{"gopher"}
	*clientAllowedOUs = []string{}
	*clientAllowedDNSs = []string{}
	*clientAllowedIPs = []net.IP{}

	assert.Nil(t, verifyPeerCertificateClient(nil, fakeChains), "verify-cn should allow servers with matching CN")
}

func TestVerifyAllowOU(t *testing.T) {
	*clientAllowedCNs = []string{}
	*clientAllowedOUs = []string{"circle"}
	*clientAllowedDNSs = []string{}
	*clientAllowedIPs = []net.IP{}

	assert.Nil(t, verifyPeerCertificateClient(nil, fakeChains), "verify-ou should allow servers with matching OU")
}

func TestVerifyAllowDNS(t *testing.T) {
	*clientAllowedCNs = []string{}
	*clientAllowedOUs = []string{}
	*clientAllowedDNSs = []string{"circle"}
	*clientAllowedIPs = []net.IP{}

	assert.Nil(t, verifyPeerCertificateClient(nil, fakeChains), "verify-dns-san should allow servers with matching DNS SAN")
}

func TestVerifyAllowIP(t *testing.T) {
	*clientAllowedCNs = []string{}
	*clientAllowedOUs = []string{}
	*clientAllowedDNSs = []string{}
	*clientAllowedIPs = []net.IP{net.IPv4(192, 168, 99, 100)}

	assert.Nil(t, verifyPeerCertificateClient(nil, fakeChains), "verify-ip-san should allow servers with matching IP SAN")
}

func TestVerifyRejectCN(t *testing.T) {
	*clientAllowedCNs = []string{"test"}
	*clientAllowedOUs = []string{}
	*clientAllowedDNSs = []string{}
	*clientAllowedIPs = []net.IP{}

	assert.NotNil(t, verifyPeerCertificateClient(nil, fakeChains), "should reject cert w/o matching CN")
}

func TestVerifyRejectOU(t *testing.T) {
	*clientAllowedCNs = []string{}
	*clientAllowedOUs = []string{"test"}
	*clientAllowedDNSs = []string{}
	*clientAllowedIPs = []net.IP{}

	assert.NotNil(t, verifyPeerCertificateClient(nil, fakeChains), "should reject cert w/o matching OU")
}

func TestVerifyRejectDNS(t *testing.T) {
	*clientAllowedCNs = []string{}
	*clientAllowedOUs = []string{}
	*clientAllowedDNSs = []string{"test"}
	*clientAllowedIPs = []net.IP{}

	assert.NotNil(t, verifyPeerCertificateClient(nil, fakeChains), "should reject cert w/o matching DNS SAN")
}

func TestVerifyRejectIP(t *testing.T) {
	*clientAllowedCNs = []string{}
	*clientAllowedOUs = []string{}
	*clientAllowedDNSs = []string{}
	*clientAllowedIPs = []net.IP{net.IPv4(1, 1, 1, 1)}

	assert.NotNil(t, verifyPeerCertificateClient(nil, fakeChains), "should reject cert w/o matching IP SAN")
}

func TestAuthorizeAllowURI(t *testing.T) {
	*serverAllowAll = false
	*serverAllowedCNs = []string{}
	*serverAllowedOUs = []string{}
	*serverAllowedDNSs = []string{}
	*serverAllowedIPs = []net.IP{}
	*serverAllowedURIs = []string{"spiffe://example.org/path/service"}

	assert.Nil(t, verifyPeerCertificateServer(nil, fakeChains), "allow-uri-san should allow clients with matching URI SAN")
}

func TestAuthorizeRejectURI(t *testing.T) {
	*serverAllowAll = false
	*serverAllowedCNs = []string{}
	*serverAllowedOUs = []string{}
	*serverAllowedDNSs = []string{}
	*serverAllowedIPs = []net.IP{}
	*serverAllowedURIs = []string{"spiffe://invalid"}

	assert.NotNil(t, verifyPeerCertificateServer(nil, fakeChains), "should reject cert w/o matching URI")
}

func TestVerifyAllowURI(t *testing.T) {
	*clientAllowedCNs = []string{}
	*clientAllowedOUs = []string{}
	*clientAllowedDNSs = []string{}
	*clientAllowedIPs = []net.IP{}
	*clientAllowedURIs = []string{"spiffe://example.org/path/service"}

	assert.Nil(t, verifyPeerCertificateClient(nil, fakeChains), "verify-uri-san should allow clients with matching URI SAN")
}

func TestVerifyRejectURI(t *testing.T) {
	*clientAllowedCNs = []string{}
	*clientAllowedOUs = []string{}
	*clientAllowedDNSs = []string{}
	*clientAllowedIPs = []net.IP{}
	*clientAllowedURIs = []string{"spiffe://invalid"}

	assert.NotNil(t, verifyPeerCertificateClient(nil, fakeChains), "should reject cert w/o matching URI")
}
<commit_msg>Cosmetic change in auth_test.go<commit_after>/*-
 * Copyright 2015 Square Inc.
 *
 * 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 main

import (
	"crypto/x509"
	"crypto/x509/pkix"
	"net"
	"testing"

	"github.com/spiffe/go-spiffe/uri"
	"github.com/stretchr/testify/assert"
)

var fakeChains = [][]*x509.Certificate{
	{
		{
			Subject: pkix.Name{
				CommonName:         "gopher",
				OrganizationalUnit: []string{"triangle", "circle"},
			},
			DNSNames:    []string{"circle"},
			IPAddresses: []net.IP{net.IPv4(192, 168, 99, 100)},
			Extensions: []pkix.Extension{
				{
					Id:       uri.OidExtensionSubjectAltName,
					Value:    getURISANFromString("scheme://valid/path"),
					Critical: false,
				}},
		},
	},
}

func getURISANFromString(s string) (uriSAN []byte) {
	uriSAN, err := uri.MarshalUriSANs([]string{s})
	panicOnError(err)

	return uriSAN
}

func TestAuthorizeNotVerified(t *testing.T) {
	*serverAllowAll = true
	*serverAllowedCNs = []string{}
	*serverAllowedOUs = []string{}
	*serverAllowedDNSs = []string{}
	*serverAllowedIPs = []net.IP{}
	*serverAllowedURIs = []string{}

	assert.NotNil(t, verifyPeerCertificateServer(nil, nil), "conn w/o cert should be rejected")
}

func TestAuthorizeReject(t *testing.T) {
	*serverAllowAll = false
	*serverAllowedCNs = []string{"test"}
	*serverAllowedOUs = []string{"test"}
	*serverAllowedDNSs = []string{"test"}
	*serverAllowedIPs = []net.IP{}
	*serverAllowedURIs = []string{"test"}

	assert.NotNil(t, verifyPeerCertificateServer(nil, fakeChains), "should reject cert w/o matching CN/OU")
}

func TestAuthorizeAllowAll(t *testing.T) {
	*serverAllowAll = true
	*serverAllowedCNs = []string{}
	*serverAllowedOUs = []string{}
	*serverAllowedDNSs = []string{}
	*serverAllowedIPs = []net.IP{}

	assert.Nil(t, verifyPeerCertificateServer(nil, fakeChains), "allow-all should always allow authed clients")
}

func TestAuthorizeAllowCN(t *testing.T) {
	*serverAllowAll = false
	*serverAllowedCNs = []string{"gopher"}
	*serverAllowedOUs = []string{}
	*serverAllowedDNSs = []string{}
	*serverAllowedIPs = []net.IP{}
	*serverAllowedURIs = []string{}

	assert.Nil(t, verifyPeerCertificateServer(nil, fakeChains), "allow-cn should allow clients with matching CN")
}

func TestAuthorizeAllowOU(t *testing.T) {
	*serverAllowAll = false
	*serverAllowedCNs = []string{}
	*serverAllowedOUs = []string{"circle"}
	*serverAllowedDNSs = []string{}
	*serverAllowedIPs = []net.IP{}
	*serverAllowedURIs = []string{}

	assert.Nil(t, verifyPeerCertificateServer(nil, fakeChains), "allow-ou should allow clients with matching OU")
}

func TestAuthorizeAllowDNS(t *testing.T) {
	*serverAllowAll = false
	*serverAllowedCNs = []string{}
	*serverAllowedOUs = []string{}
	*serverAllowedDNSs = []string{"circle"}
	*serverAllowedIPs = []net.IP{}
	*serverAllowedURIs = []string{}

	assert.Nil(t, verifyPeerCertificateServer(nil, fakeChains), "allow-dns-san should allow clients with matching DNS SAN")
}

func TestAuthorizeAllowIP(t *testing.T) {
	*serverAllowAll = false
	*serverAllowedCNs = []string{}
	*serverAllowedOUs = []string{}
	*serverAllowedDNSs = []string{}
	*serverAllowedIPs = []net.IP{net.IPv4(192, 168, 99, 100)}
	*serverAllowedURIs = []string{}

	assert.Nil(t, verifyPeerCertificateServer(nil, fakeChains), "allow-ip-san should allow clients with matching IP SAN")
}

func TestAuthorizeAllowURI(t *testing.T) {
	*serverAllowAll = false
	*serverAllowedCNs = []string{}
	*serverAllowedOUs = []string{}
	*serverAllowedDNSs = []string{}
	*serverAllowedIPs = []net.IP{}
	*serverAllowedURIs = []string{"scheme://valid/path"}

	assert.Nil(t, verifyPeerCertificateServer(nil, fakeChains), "allow-uri-san should allow clients with matching URI SAN")
}

func TestAuthorizeRejectURI(t *testing.T) {
	*serverAllowAll = false
	*serverAllowedCNs = []string{}
	*serverAllowedOUs = []string{}
	*serverAllowedDNSs = []string{}
	*serverAllowedIPs = []net.IP{}
	*serverAllowedURIs = []string{"schema://invalid/path"}

	assert.NotNil(t, verifyPeerCertificateServer(nil, fakeChains), "should reject cert w/o matching URI")
}

func TestVerifyAllowCN(t *testing.T) {
	*clientAllowedCNs = []string{"gopher"}
	*clientAllowedOUs = []string{}
	*clientAllowedDNSs = []string{}
	*clientAllowedIPs = []net.IP{}

	assert.Nil(t, verifyPeerCertificateClient(nil, fakeChains), "verify-cn should allow servers with matching CN")
}

func TestVerifyAllowOU(t *testing.T) {
	*clientAllowedCNs = []string{}
	*clientAllowedOUs = []string{"circle"}
	*clientAllowedDNSs = []string{}
	*clientAllowedIPs = []net.IP{}

	assert.Nil(t, verifyPeerCertificateClient(nil, fakeChains), "verify-ou should allow servers with matching OU")
}

func TestVerifyAllowDNS(t *testing.T) {
	*clientAllowedCNs = []string{}
	*clientAllowedOUs = []string{}
	*clientAllowedDNSs = []string{"circle"}
	*clientAllowedIPs = []net.IP{}

	assert.Nil(t, verifyPeerCertificateClient(nil, fakeChains), "verify-dns-san should allow servers with matching DNS SAN")
}

func TestVerifyAllowIP(t *testing.T) {
	*clientAllowedCNs = []string{}
	*clientAllowedOUs = []string{}
	*clientAllowedDNSs = []string{}
	*clientAllowedIPs = []net.IP{net.IPv4(192, 168, 99, 100)}

	assert.Nil(t, verifyPeerCertificateClient(nil, fakeChains), "verify-ip-san should allow servers with matching IP SAN")
}

func TestVerifyRejectCN(t *testing.T) {
	*clientAllowedCNs = []string{"test"}
	*clientAllowedOUs = []string{}
	*clientAllowedDNSs = []string{}
	*clientAllowedIPs = []net.IP{}

	assert.NotNil(t, verifyPeerCertificateClient(nil, fakeChains), "should reject cert w/o matching CN")
}

func TestVerifyRejectOU(t *testing.T) {
	*clientAllowedCNs = []string{}
	*clientAllowedOUs = []string{"test"}
	*clientAllowedDNSs = []string{}
	*clientAllowedIPs = []net.IP{}

	assert.NotNil(t, verifyPeerCertificateClient(nil, fakeChains), "should reject cert w/o matching OU")
}

func TestVerifyRejectDNS(t *testing.T) {
	*clientAllowedCNs = []string{}
	*clientAllowedOUs = []string{}
	*clientAllowedDNSs = []string{"test"}
	*clientAllowedIPs = []net.IP{}

	assert.NotNil(t, verifyPeerCertificateClient(nil, fakeChains), "should reject cert w/o matching DNS SAN")
}

func TestVerifyRejectIP(t *testing.T) {
	*clientAllowedCNs = []string{}
	*clientAllowedOUs = []string{}
	*clientAllowedDNSs = []string{}
	*clientAllowedIPs = []net.IP{net.IPv4(1, 1, 1, 1)}

	assert.NotNil(t, verifyPeerCertificateClient(nil, fakeChains), "should reject cert w/o matching IP SAN")
}

func TestVerifyAllowURI(t *testing.T) {
	*clientAllowedCNs = []string{}
	*clientAllowedOUs = []string{}
	*clientAllowedDNSs = []string{}
	*clientAllowedIPs = []net.IP{}
	*clientAllowedURIs = []string{"scheme://valid/path"}

	assert.Nil(t, verifyPeerCertificateClient(nil, fakeChains), "verify-uri-san should allow clients with matching URI SAN")
}

func TestVerifyRejectURI(t *testing.T) {
	*clientAllowedCNs = []string{}
	*clientAllowedOUs = []string{}
	*clientAllowedDNSs = []string{}
	*clientAllowedIPs = []net.IP{}
	*clientAllowedURIs = []string{"scheme://invalid/path"}

	assert.NotNil(t, verifyPeerCertificateClient(nil, fakeChains), "should reject cert w/o matching URI")
}
<|endoftext|>


<commit_before>package venom

import (
	"errors"
	"fmt"
	"reflect"
	"strings"

	"github.com/spf13/pflag"
)

const (
	SquashFlagsTag = "++"
)

// Parse name for mapstructure tags i.e. fetch banana from:
//
// type Foo struct {
//     foo int `mapstructure:"banana"`
// }
//
func parseMapstructureTag(tag string) (string, bool) {
	parts := strings.SplitN(tag, ",", 2)
	if len(parts) == 0 {
		return "", false
	}
	return parts[0], true
}

func parseTag(tag string) (string, string, string) {
	parts := strings.SplitN(tag, ",", 3)

	// flag: bar, b, Some barness -> flag: bar,b,Some barness
	for i, p := range parts {
		parts[i] = strings.TrimSpace(p)
	}

	switch len(parts) {
	case 1:
		// flag: b
		if len(parts[0]) == 1 {
			return "", parts[0], ""
		}
		// flag: bar
		return parts[0], "", ""
	case 2:
		// flag: b,Some barness
		if len(parts[0]) == 1 {
			return "", parts[0], parts[1]
		}
		// flag: bar,b
		if len(parts[1]) == 1 {
			return parts[0], parts[1], ""
		}
		// flag: bar,Some barness
		return parts[0], "", parts[1]
	case 3:
		// flag: bar,b,Some barness
		return parts[0], parts[1], parts[2]
	default:
		return "", "", ""
	}
}

func DefineFlags(config interface{}) *pflag.FlagSet {
	flags, err := NewFlags(config)
	if err != nil {
		panic(err)
	}
	return flags
}

func NewFlags(config interface{}) (*pflag.FlagSet, error) {
	a := flagsFactory{
		tags: []string{"flag", "pflag"},
	}
	return a.createFlags(config)
}

type flagsFactory struct {
	tags []string
}

func (a flagsFactory) lookupTag(tag reflect.StructTag) (string, bool) {
	for _, name := range a.tags {
		v, ok := tag.Lookup(name)
		if ok {
			return v, true
		}
	}
	return "", false
}

func (a flagsFactory) createFlags(defaults interface{}) (*pflag.FlagSet, error) {
	var flags pflag.FlagSet

	//
	// Remove one level of indirection.
	//
	v := reflect.ValueOf(defaults)
	if v.Kind() == reflect.Ptr {
		v = reflect.Indirect(v)
	}

	//
	// Make sure we end up with a struct.
	//
	if v.Kind() != reflect.Struct {
		return nil, errors.New("Struct or pointer to struct expected")
	}

	//
	// For every tagged struct field create a flag.
	//
	for i := 0; i < v.Type().NumField(); i++ {
		structField := v.Type().Field(i)
		fieldType := structField.Type
		fieldValue := v.Field(i)

		tag, ok := a.lookupTag(structField.Tag)
		if !ok {
			continue
		}

		//
		// This means we want to squash all flags from the inner structure so they appear as is they are defined
		// in the outer structure.
		//
		if tag == SquashFlagsTag {
			if fieldType.Kind() != reflect.Struct {
				return nil, fmt.Errorf(`flag:"%s" is supported only for inner structs but is set on: %s`, SquashFlagsTag, fieldType)
			}
			innerFlags, err := a.createFlags(fieldValue.Interface())
			if err != nil {
				return nil, err
			}
			flags.AddFlagSet(innerFlags)
			continue
		}

		//
		// In case we have mapstructure defined it must have exactly the same name as flag has.
		//
		mapTag, ok := structField.Tag.Lookup("mapstructure")
		if ok {
			mapName, ok := parseMapstructureTag(mapTag)
			if ok && !strings.HasPrefix(tag, mapName+",") {
				return nil, fmt.Errorf(`Both "mapstructure" and "flag" tags must have equal names but are different for field: %s`, structField.Name)
				continue
			}
		}

		err := addFlagForTag(&flags, tag, fieldValue, fieldType)
		if err != nil {
			return nil, err
		}
	}

	return &flags, nil
}

func addFlagForTag(flags *pflag.FlagSet, tag string, fieldValue reflect.Value, fieldType reflect.Type) error {
	name, shorthand, usage := parseTag(tag)
	switch fieldType.Kind() {
	case reflect.Bool:
		value := bool(fieldValue.Bool())
		flags.BoolP(name, shorthand, value, usage)
	case reflect.Int:
		value := int(fieldValue.Int())
		flags.IntP(name, shorthand, value, usage)
	case reflect.Int8:
		value := int8(fieldValue.Int())
		flags.Int8P(name, shorthand, value, usage)
	case reflect.Int16:
		value := int32(fieldValue.Int())
		flags.Int32P(name, shorthand, value, usage) // Not a typo, pflags doesn't have Int16
	case reflect.Int32:
		value := int32(fieldValue.Int())
		flags.Int32P(name, shorthand, value, usage)
	case reflect.Int64:
		value := int64(fieldValue.Int())
		flags.Int64P(name, shorthand, value, usage)
	case reflect.Uint:
		value := uint(fieldValue.Uint())
		flags.UintP(name, shorthand, value, usage)
	case reflect.Uint8:
		value := uint8(fieldValue.Uint())
		flags.Uint8P(name, shorthand, value, usage)
	case reflect.Uint16:
		value := uint16(fieldValue.Uint())
		flags.Uint16P(name, shorthand, value, usage)
	case reflect.Uint32:
		value := uint32(fieldValue.Uint())
		flags.Uint32P(name, shorthand, value, usage)
	case reflect.Uint64:
		value := uint64(fieldValue.Uint())
		flags.Uint64P(name, shorthand, value, usage)
	case reflect.Float32:
		value := float32(fieldValue.Float())
		flags.Float32P(name, shorthand, value, usage)
	case reflect.Float64:
		value := float64(fieldValue.Float())
		flags.Float64P(name, shorthand, value, usage)
	case reflect.String:
		value := string(fieldValue.String())
		flags.StringP(name, shorthand, value, usage)
	default:
		return fmt.Errorf("Unsupported type for field with flag tag %q: %s", name, fieldType)
	}
	return nil
}
<commit_msg>Make sure we catch equality as well<commit_after>package venom

import (
	"errors"
	"fmt"
	"reflect"
	"strings"

	"github.com/spf13/pflag"
)

const (
	SquashFlagsTag = "++"
)

// Parse name for mapstructure tags i.e. fetch banana from:
//
// type Foo struct {
//     foo int `mapstructure:"banana"`
// }
//
func parseMapstructureTag(tag string) (string, bool) {
	parts := strings.SplitN(tag, ",", 2)
	if len(parts) == 0 {
		return "", false
	}
	return parts[0], true
}

func parseTag(tag string) (string, string, string) {
	parts := strings.SplitN(tag, ",", 3)

	// flag: bar, b, Some barness -> flag: bar,b,Some barness
	for i, p := range parts {
		parts[i] = strings.TrimSpace(p)
	}

	switch len(parts) {
	case 1:
		// flag: b
		if len(parts[0]) == 1 {
			return "", parts[0], ""
		}
		// flag: bar
		return parts[0], "", ""
	case 2:
		// flag: b,Some barness
		if len(parts[0]) == 1 {
			return "", parts[0], parts[1]
		}
		// flag: bar,b
		if len(parts[1]) == 1 {
			return parts[0], parts[1], ""
		}
		// flag: bar,Some barness
		return parts[0], "", parts[1]
	case 3:
		// flag: bar,b,Some barness
		return parts[0], parts[1], parts[2]
	default:
		return "", "", ""
	}
}

func DefineFlags(config interface{}) *pflag.FlagSet {
	flags, err := NewFlags(config)
	if err != nil {
		panic(err)
	}
	return flags
}

func NewFlags(config interface{}) (*pflag.FlagSet, error) {
	a := flagsFactory{
		tags: []string{"flag", "pflag"},
	}
	return a.createFlags(config)
}

type flagsFactory struct {
	tags []string
}

func (a flagsFactory) lookupTag(tag reflect.StructTag) (string, bool) {
	for _, name := range a.tags {
		v, ok := tag.Lookup(name)
		if ok {
			return v, true
		}
	}
	return "", false
}

func (a flagsFactory) createFlags(defaults interface{}) (*pflag.FlagSet, error) {
	var flags pflag.FlagSet

	//
	// Remove one level of indirection.
	//
	v := reflect.ValueOf(defaults)
	if v.Kind() == reflect.Ptr {
		v = reflect.Indirect(v)
	}

	//
	// Make sure we end up with a struct.
	//
	if v.Kind() != reflect.Struct {
		return nil, errors.New("Struct or pointer to struct expected")
	}

	//
	// For every tagged struct field create a flag.
	//
	for i := 0; i < v.Type().NumField(); i++ {
		structField := v.Type().Field(i)
		fieldType := structField.Type
		fieldValue := v.Field(i)

		tag, ok := a.lookupTag(structField.Tag)
		if !ok {
			continue
		}

		//
		// This means we want to squash all flags from the inner structure so they appear as is they are defined
		// in the outer structure.
		//
		if tag == SquashFlagsTag {
			if fieldType.Kind() != reflect.Struct {
				return nil, fmt.Errorf(`flag:"%s" is supported only for inner structs but is set on: %s`, SquashFlagsTag, fieldType)
			}
			innerFlags, err := a.createFlags(fieldValue.Interface())
			if err != nil {
				return nil, err
			}
			flags.AddFlagSet(innerFlags)
			continue
		}

		//
		// In case we have mapstructure defined it must have exactly the same name as flag has.
		//
		mapTag, ok := structField.Tag.Lookup("mapstructure")
		if ok {
			mapName, ok := parseMapstructureTag(mapTag)
			if ok && !(tag == mapName || strings.HasPrefix(tag, mapName+",")) {
				return nil, fmt.Errorf(`Both "mapstructure" and "flag" tags must have equal names but are different for field: %s`, structField.Name)
				continue
			}
		}

		err := addFlagForTag(&flags, tag, fieldValue, fieldType)
		if err != nil {
			return nil, err
		}
	}

	return &flags, nil
}

func addFlagForTag(flags *pflag.FlagSet, tag string, fieldValue reflect.Value, fieldType reflect.Type) error {
	name, shorthand, usage := parseTag(tag)
	switch fieldType.Kind() {
	case reflect.Bool:
		value := bool(fieldValue.Bool())
		flags.BoolP(name, shorthand, value, usage)
	case reflect.Int:
		value := int(fieldValue.Int())
		flags.IntP(name, shorthand, value, usage)
	case reflect.Int8:
		value := int8(fieldValue.Int())
		flags.Int8P(name, shorthand, value, usage)
	case reflect.Int16:
		value := int32(fieldValue.Int())
		flags.Int32P(name, shorthand, value, usage) // Not a typo, pflags doesn't have Int16
	case reflect.Int32:
		value := int32(fieldValue.Int())
		flags.Int32P(name, shorthand, value, usage)
	case reflect.Int64:
		value := int64(fieldValue.Int())
		flags.Int64P(name, shorthand, value, usage)
	case reflect.Uint:
		value := uint(fieldValue.Uint())
		flags.UintP(name, shorthand, value, usage)
	case reflect.Uint8:
		value := uint8(fieldValue.Uint())
		flags.Uint8P(name, shorthand, value, usage)
	case reflect.Uint16:
		value := uint16(fieldValue.Uint())
		flags.Uint16P(name, shorthand, value, usage)
	case reflect.Uint32:
		value := uint32(fieldValue.Uint())
		flags.Uint32P(name, shorthand, value, usage)
	case reflect.Uint64:
		value := uint64(fieldValue.Uint())
		flags.Uint64P(name, shorthand, value, usage)
	case reflect.Float32:
		value := float32(fieldValue.Float())
		flags.Float32P(name, shorthand, value, usage)
	case reflect.Float64:
		value := float64(fieldValue.Float())
		flags.Float64P(name, shorthand, value, usage)
	case reflect.String:
		value := string(fieldValue.String())
		flags.StringP(name, shorthand, value, usage)
	default:
		return fmt.Errorf("Unsupported type for field with flag tag %q: %s", name, fieldType)
	}
	return nil
}
<|endoftext|>