<commit_before>package render

import (
	"strings"

	"github.com/weaveworks/scope/probe/docker"
	"github.com/weaveworks/scope/probe/kubernetes"
	"github.com/weaveworks/scope/report"
)

// Constants are used in the tests.
const (
	UnmanagedID    = "unmanaged"
	UnmanagedMajor = "Unmanaged"
)

// UnmanagedIDPrefix is the prefix of unmanaged pseudo nodes
var UnmanagedIDPrefix = MakePseudoNodeID(UnmanagedID, "")

func renderKubernetesTopologies(rpt report.Report) bool {
	// Render if any k8s topology has any nodes
	topologies := []*report.Topology{
		&rpt.Pod,
		&rpt.Service,
		&rpt.Deployment,
		&rpt.DaemonSet,
		&rpt.StatefulSet,
		&rpt.CronJob,
		&rpt.PersistentVolume,
		&rpt.PersistentVolumeClaim,
		&rpt.StorageClass,
	}
	for _, t := range topologies {
		if len(t.Nodes) > 0 {
			return true
		}
	}
	return false
}

func isPauseContainer(n report.Node) bool {
	image, ok := n.Latest.Lookup(docker.ImageName)
	return ok && kubernetes.IsPauseImageName(image)
}

// PodRenderer is a Renderer which produces a renderable kubernetes
// graph by merging the container graph and the pods topology.
var PodRenderer = Memoise(ConditionalRenderer(renderKubernetesTopologies,
	MakeFilter(
		func(n report.Node) bool {
			state, ok := n.Latest.Lookup(kubernetes.State)
			return !ok || !(state == kubernetes.StateDeleted || state == kubernetes.StateFailed)
		},
		MakeReduce(
			PropagateSingleMetrics(report.Container,
				MakeMap(propagateHostID,
					Map2Parent{topologies: []string{report.Pod}, noParentsPseudoID: UnmanagedID,
						chainRenderer: MakeFilter(
							ComposeFilterFuncs(
								IsRunning,
								Complement(isPauseContainer),
							),
							ContainerWithImageNameRenderer,
						)},
				),
			),
			ConnectionJoin(MapPod2IP, report.Pod),
			KubernetesVolumesRenderer,
		),
	),
))

// Pods are not tagged with a Host ID, but their container children are.
// If n doesn't already have a host ID, copy it from one of the children
func propagateHostID(n report.Node) report.Node {
	if _, found := n.Latest.Lookup(report.HostNodeID); found {
		return n
	}
	var first *report.Node
	n.Children.ForEach(func(child report.Node) {
		if first == nil {
			first = &child
		}
	})
	if first != nil {
		n.Latest = n.Latest.Propagate(first.Latest, report.HostNodeID)
	}
	return n
}

// PodServiceRenderer is a Renderer which produces a renderable kubernetes services
// graph by merging the pods graph and the services topology.
//
// not memoised
var PodServiceRenderer = ConditionalRenderer(renderKubernetesTopologies,
	renderParents(
		report.Pod, []string{report.Service}, "",
		PodRenderer,
	),
)

// KubeControllerRenderer is a Renderer which combines all the 'controller' topologies.
// Pods with no controller are mapped to 'Unmanaged'
// We can't simply combine the rendered graphs of the high level objects as they would never
// have connections to each other.
//
// not memoised
var KubeControllerRenderer = ConditionalRenderer(renderKubernetesTopologies,
	renderParents(
		report.Pod, []string{report.Deployment, report.DaemonSet, report.StatefulSet, report.CronJob}, UnmanagedID,
		PodRenderer,
	),
)

// renderParents produces a 'standard' renderer for mapping from some child topology to some parent topologies,
// by taking a child renderer, mapping to parents, propagating single metrics, and joining with full parent topology.
// Other options are as per Map2Parent.
func renderParents(childTopology string, parentTopologies []string, noParentsPseudoID string, childRenderer Renderer) Renderer {
	selectors := make([]Renderer, len(parentTopologies))
	for i, topology := range parentTopologies {
		selectors[i] = TopologySelector(topology)
	}
	return MakeReduce(append(
		selectors,
		PropagateSingleMetrics(childTopology,
			Map2Parent{topologies: parentTopologies, noParentsPseudoID: noParentsPseudoID,
				chainRenderer: childRenderer},
		),
	)...)
}

// MapPod2IP maps pod nodes to their IP address.  This allows pods to
// be joined directly with the endpoint topology.
func MapPod2IP(m report.Node) []string {
	// if this pod belongs to the host's networking namespace
	// we cannot use its IP to attribute connections
	// (they could come from any other process on the host or DNAT-ed IPs)
	if _, ok := m.Latest.Lookup(kubernetes.IsInHostNetwork); ok {
		return nil
	}

	ip, ok := m.Latest.Lookup(kubernetes.IP)
	if !ok || ip == "" {
		return nil
	}
	return []string{report.MakeScopedEndpointNodeID("", ip, "")}
}

// Map2Parent is a Renderer which maps Nodes to some parent grouping.
type Map2Parent struct {
	// Renderer to chain from
	chainRenderer Renderer
	// The topology IDs to look for parents in
	topologies []string
	// Either the ID prefix of the pseudo node to use for nodes without
	// any parents in the group, eg. UnmanagedID, or "" to drop nodes without any parents.
	noParentsPseudoID string
}

// Render implements Renderer
func (m Map2Parent) Render(rpt report.Report) Nodes {
	input := m.chainRenderer.Render(rpt)
	ret := newJoinResults(nil)

	for _, n := range input.Nodes {
		// Uncontained becomes Unmanaged/whatever if noParentsPseudoID is set
		if m.noParentsPseudoID != "" && strings.HasPrefix(n.ID, UncontainedIDPrefix) {
			id := MakePseudoNodeID(m.noParentsPseudoID, n.ID[len(UncontainedIDPrefix):])
			ret.addChildAndChildren(n, id, Pseudo)
			continue
		}

		// Propagate all pseudo nodes
		if n.Topology == Pseudo {
			ret.passThrough(n)
			continue
		}

		added := false
		// For each topology, map to any parents we can find
		for _, topology := range m.topologies {
			if groupIDs, ok := n.Parents.Lookup(topology); ok {
				for _, id := range groupIDs {
					ret.addChildAndChildren(n, id, topology)
					added = true
				}
			}
		}

		if !added && m.noParentsPseudoID != "" {
			// Map to pseudo node
			id := MakePseudoNodeID(m.noParentsPseudoID, report.ExtractHostID(n))
			ret.addChildAndChildren(n, id, Pseudo)
		}
	}
	return ret.result(input)
}
<commit_msg>Fix pod host propagation<commit_after>package render

import (
	"strings"

	"github.com/weaveworks/scope/probe/docker"
	"github.com/weaveworks/scope/probe/kubernetes"
	"github.com/weaveworks/scope/report"
)

// Constants are used in the tests.
const (
	UnmanagedID    = "unmanaged"
	UnmanagedMajor = "Unmanaged"
)

// UnmanagedIDPrefix is the prefix of unmanaged pseudo nodes
var UnmanagedIDPrefix = MakePseudoNodeID(UnmanagedID, "")

func renderKubernetesTopologies(rpt report.Report) bool {
	// Render if any k8s topology has any nodes
	topologies := []*report.Topology{
		&rpt.Pod,
		&rpt.Service,
		&rpt.Deployment,
		&rpt.DaemonSet,
		&rpt.StatefulSet,
		&rpt.CronJob,
		&rpt.PersistentVolume,
		&rpt.PersistentVolumeClaim,
		&rpt.StorageClass,
	}
	for _, t := range topologies {
		if len(t.Nodes) > 0 {
			return true
		}
	}
	return false
}

func isPauseContainer(n report.Node) bool {
	image, ok := n.Latest.Lookup(docker.ImageName)
	return ok && kubernetes.IsPauseImageName(image)
}

// PodRenderer is a Renderer which produces a renderable kubernetes
// graph by merging the container graph and the pods topology.
var PodRenderer = Memoise(ConditionalRenderer(renderKubernetesTopologies,
	MakeFilter(
		func(n report.Node) bool {
			state, ok := n.Latest.Lookup(kubernetes.State)
			return !ok || !(state == kubernetes.StateDeleted || state == kubernetes.StateFailed)
		},
		MakeReduce(
			PropagateSingleMetrics(report.Container,
				MakeMap(propagatePodHost,
					Map2Parent{topologies: []string{report.Pod}, noParentsPseudoID: UnmanagedID,
						chainRenderer: MakeFilter(
							ComposeFilterFuncs(
								IsRunning,
								Complement(isPauseContainer),
							),
							ContainerWithImageNameRenderer,
						)},
				),
			),
			ConnectionJoin(MapPod2IP, report.Pod),
			KubernetesVolumesRenderer,
		),
	),
))

// Pods are not tagged with a Host parent, but their container children are.
// If n doesn't already have a host, copy it from one of the children
func propagatePodHost(n report.Node) report.Node {
	if n.Topology != report.Pod {
		return n
	} else if _, found := n.Parents.Lookup(report.Host); found {
		return n
	}
	done := false
	n.Children.ForEach(func(child report.Node) {
		if !done {
			if hosts, found := child.Parents.Lookup(report.Host); found {
				for _, h := range hosts {
					n = n.WithParent(report.Host, h)
				}
				done = true
			}
		}
	})
	return n
}

// PodServiceRenderer is a Renderer which produces a renderable kubernetes services
// graph by merging the pods graph and the services topology.
//
// not memoised
var PodServiceRenderer = ConditionalRenderer(renderKubernetesTopologies,
	renderParents(
		report.Pod, []string{report.Service}, "",
		PodRenderer,
	),
)

// KubeControllerRenderer is a Renderer which combines all the 'controller' topologies.
// Pods with no controller are mapped to 'Unmanaged'
// We can't simply combine the rendered graphs of the high level objects as they would never
// have connections to each other.
//
// not memoised
var KubeControllerRenderer = ConditionalRenderer(renderKubernetesTopologies,
	renderParents(
		report.Pod, []string{report.Deployment, report.DaemonSet, report.StatefulSet, report.CronJob}, UnmanagedID,
		PodRenderer,
	),
)

// renderParents produces a 'standard' renderer for mapping from some child topology to some parent topologies,
// by taking a child renderer, mapping to parents, propagating single metrics, and joining with full parent topology.
// Other options are as per Map2Parent.
func renderParents(childTopology string, parentTopologies []string, noParentsPseudoID string, childRenderer Renderer) Renderer {
	selectors := make([]Renderer, len(parentTopologies))
	for i, topology := range parentTopologies {
		selectors[i] = TopologySelector(topology)
	}
	return MakeReduce(append(
		selectors,
		PropagateSingleMetrics(childTopology,
			Map2Parent{topologies: parentTopologies, noParentsPseudoID: noParentsPseudoID,
				chainRenderer: childRenderer},
		),
	)...)
}

// MapPod2IP maps pod nodes to their IP address.  This allows pods to
// be joined directly with the endpoint topology.
func MapPod2IP(m report.Node) []string {
	// if this pod belongs to the host's networking namespace
	// we cannot use its IP to attribute connections
	// (they could come from any other process on the host or DNAT-ed IPs)
	if _, ok := m.Latest.Lookup(kubernetes.IsInHostNetwork); ok {
		return nil
	}

	ip, ok := m.Latest.Lookup(kubernetes.IP)
	if !ok || ip == "" {
		return nil
	}
	return []string{report.MakeScopedEndpointNodeID("", ip, "")}
}

// Map2Parent is a Renderer which maps Nodes to some parent grouping.
type Map2Parent struct {
	// Renderer to chain from
	chainRenderer Renderer
	// The topology IDs to look for parents in
	topologies []string
	// Either the ID prefix of the pseudo node to use for nodes without
	// any parents in the group, eg. UnmanagedID, or "" to drop nodes without any parents.
	noParentsPseudoID string
}

// Render implements Renderer
func (m Map2Parent) Render(rpt report.Report) Nodes {
	input := m.chainRenderer.Render(rpt)
	ret := newJoinResults(nil)

	for _, n := range input.Nodes {
		// Uncontained becomes Unmanaged/whatever if noParentsPseudoID is set
		if m.noParentsPseudoID != "" && strings.HasPrefix(n.ID, UncontainedIDPrefix) {
			id := MakePseudoNodeID(m.noParentsPseudoID, n.ID[len(UncontainedIDPrefix):])
			ret.addChildAndChildren(n, id, Pseudo)
			continue
		}

		// Propagate all pseudo nodes
		if n.Topology == Pseudo {
			ret.passThrough(n)
			continue
		}

		added := false
		// For each topology, map to any parents we can find
		for _, topology := range m.topologies {
			if groupIDs, ok := n.Parents.Lookup(topology); ok {
				for _, id := range groupIDs {
					ret.addChildAndChildren(n, id, topology)
					added = true
				}
			}
		}

		if !added && m.noParentsPseudoID != "" {
			// Map to pseudo node
			id := MakePseudoNodeID(m.noParentsPseudoID, report.ExtractHostID(n))
			ret.addChildAndChildren(n, id, Pseudo)
		}
	}
	return ret.result(input)
}
<|endoftext|>


<commit_before>package render

import (
	"strings"

	"github.com/weaveworks/scope/probe/docker"
	"github.com/weaveworks/scope/probe/kubernetes"
	"github.com/weaveworks/scope/report"
)

// Constants are used in the tests.
const (
	UnmanagedID    = "unmanaged"
	UnmanagedMajor = "Unmanaged"
)

func renderKubernetesTopologies(rpt report.Report) bool {
	return len(rpt.Pod.Nodes)+len(rpt.Service.Nodes)+len(rpt.Deployment.Nodes)+len(rpt.ReplicaSet.Nodes) >= 1
}

// PodRenderer is a Renderer which produces a renderable kubernetes
// graph by merging the container graph and the pods topology.
var PodRenderer = ConditionalRenderer(renderKubernetesTopologies,
	ApplyDecorators(MakeFilter(
		func(n report.Node) bool {
			state, ok := n.Latest.Lookup(kubernetes.State)
			return (!ok || state != kubernetes.StateDeleted)
		},
		MakeReduce(
			MakeMap(
				MapContainer2Pod,
				ContainerWithImageNameRenderer,
			),
			SelectPod,
		),
	)),
)

// PodServiceRenderer is a Renderer which produces a renderable kubernetes services
// graph by merging the pods graph and the services topology.
var PodServiceRenderer = ConditionalRenderer(renderKubernetesTopologies,
	ApplyDecorators(
		MakeReduce(
			MakeMap(
				Map2Service,
				PodRenderer,
			),
			SelectService,
		),
	),
)

// DeploymentRenderer is a Renderer which produces a renderable kubernetes deployments
// graph by merging the pods graph and the deployments topology.
var DeploymentRenderer = ConditionalRenderer(renderKubernetesTopologies,
	ApplyDecorators(
		MakeReduce(
			MakeMap(
				Map2Deployment,
				ReplicaSetRenderer,
			),
			SelectDeployment,
		),
	),
)

// ReplicaSetRenderer is a Renderer which produces a renderable kubernetes replica sets
// graph by merging the pods graph and the replica sets topology.
var ReplicaSetRenderer = ConditionalRenderer(renderKubernetesTopologies,
	ApplyDecorators(
		MakeReduce(
			MakeMap(
				Map2ReplicaSet,
				PodRenderer,
			),
			SelectReplicaSet,
		),
	),
)

// MapContainer2Pod maps container Nodes to pod
// Nodes.
//
// If this function is given a node without a kubernetes_pod_id
// (including other pseudo nodes), it will produce an "Unmanaged"
// pseudo node.
//
// Otherwise, this function will produce a node with the correct ID
// format for a container, but without any Major or Minor labels.
// It does not have enough info to do that, and the resulting graph
// must be merged with a container graph to get that info.
func MapContainer2Pod(n report.Node, _ report.Networks) report.Nodes {
	// Uncontained becomes unmanaged in the pods view
	if strings.HasPrefix(n.ID, MakePseudoNodeID(UncontainedID)) {
		id := MakePseudoNodeID(UnmanagedID, report.ExtractHostID(n))
		node := NewDerivedPseudoNode(id, n)
		return report.Nodes{id: node}
	}

	// Propagate all pseudo nodes
	if n.Topology == Pseudo {
		return report.Nodes{n.ID: n}
	}

	// Ignore non-running containers
	if state, ok := n.Latest.Lookup(docker.ContainerState); ok && state != docker.StateRunning {
		return report.Nodes{}
	}

	// Otherwise, if some some reason the container doesn't have a pod uid (maybe
	// slightly out of sync reports, or its not in a pod), make it part of unmanaged.
	uid, ok := n.Latest.Lookup(docker.LabelPrefix + "io.kubernetes.pod.uid")
	if !ok {
		id := MakePseudoNodeID(UnmanagedID, report.ExtractHostID(n))
		node := NewDerivedPseudoNode(id, n)
		return report.Nodes{id: node}
	}

	id := report.MakePodNodeID(uid)
	node := NewDerivedNode(id, n).
		WithTopology(report.Pod)
	node.Counters = node.Counters.Add(n.Topology, 1)
	return report.Nodes{id: node}
}

// The various ways of grouping pods
var (
	Map2Service    = Map2Parent(report.Service)
	Map2Deployment = Map2Parent(report.Deployment)
	Map2ReplicaSet = Map2Parent(report.ReplicaSet)
)

// Map2Parent maps Nodes to some parent grouping.
func Map2Parent(topology string) func(n report.Node, _ report.Networks) report.Nodes {
	return func(n report.Node, _ report.Networks) report.Nodes {
		// Propagate all pseudo nodes
		if n.Topology == Pseudo {
			return report.Nodes{n.ID: n}
		}

		// Otherwise, if some some reason the node doesn't have any of these ids
		// (maybe slightly out of sync reports, or its not in this group), just
		// drop it
		groupIDs, ok := n.Parents.Lookup(topology)
		if !ok {
			return report.Nodes{}
		}

		result := report.Nodes{}
		for _, id := range groupIDs {
			node := NewDerivedNode(id, n).WithTopology(topology)
			node.Counters = node.Counters.Add(n.Topology, 1)
			result[id] = node
		}
		return result
	}
}
<commit_msg>Propagate the pod counter from replicasets for deployments.<commit_after>package render

import (
	"strings"

	"github.com/weaveworks/scope/probe/docker"
	"github.com/weaveworks/scope/probe/kubernetes"
	"github.com/weaveworks/scope/report"
)

// Constants are used in the tests.
const (
	UnmanagedID    = "unmanaged"
	UnmanagedMajor = "Unmanaged"
)

func renderKubernetesTopologies(rpt report.Report) bool {
	return len(rpt.Pod.Nodes)+len(rpt.Service.Nodes)+len(rpt.Deployment.Nodes)+len(rpt.ReplicaSet.Nodes) >= 1
}

// PodRenderer is a Renderer which produces a renderable kubernetes
// graph by merging the container graph and the pods topology.
var PodRenderer = ConditionalRenderer(renderKubernetesTopologies,
	ApplyDecorators(MakeFilter(
		func(n report.Node) bool {
			state, ok := n.Latest.Lookup(kubernetes.State)
			return (!ok || state != kubernetes.StateDeleted)
		},
		MakeReduce(
			MakeMap(
				MapContainer2Pod,
				ContainerWithImageNameRenderer,
			),
			SelectPod,
		),
	)),
)

// PodServiceRenderer is a Renderer which produces a renderable kubernetes services
// graph by merging the pods graph and the services topology.
var PodServiceRenderer = ConditionalRenderer(renderKubernetesTopologies,
	ApplyDecorators(
		MakeReduce(
			MakeMap(
				Map2Service,
				PodRenderer,
			),
			SelectService,
		),
	),
)

// DeploymentRenderer is a Renderer which produces a renderable kubernetes deployments
// graph by merging the pods graph and the deployments topology.
var DeploymentRenderer = ConditionalRenderer(renderKubernetesTopologies,
	ApplyDecorators(
		MakeReduce(
			MakeMap(
				Map2Deployment,
				ReplicaSetRenderer,
			),
			SelectDeployment,
		),
	),
)

// ReplicaSetRenderer is a Renderer which produces a renderable kubernetes replica sets
// graph by merging the pods graph and the replica sets topology.
var ReplicaSetRenderer = ConditionalRenderer(renderKubernetesTopologies,
	ApplyDecorators(
		MakeReduce(
			MakeMap(
				Map2ReplicaSet,
				PodRenderer,
			),
			SelectReplicaSet,
		),
	),
)

// MapContainer2Pod maps container Nodes to pod
// Nodes.
//
// If this function is given a node without a kubernetes_pod_id
// (including other pseudo nodes), it will produce an "Unmanaged"
// pseudo node.
//
// Otherwise, this function will produce a node with the correct ID
// format for a container, but without any Major or Minor labels.
// It does not have enough info to do that, and the resulting graph
// must be merged with a container graph to get that info.
func MapContainer2Pod(n report.Node, _ report.Networks) report.Nodes {
	// Uncontained becomes unmanaged in the pods view
	if strings.HasPrefix(n.ID, MakePseudoNodeID(UncontainedID)) {
		id := MakePseudoNodeID(UnmanagedID, report.ExtractHostID(n))
		node := NewDerivedPseudoNode(id, n)
		return report.Nodes{id: node}
	}

	// Propagate all pseudo nodes
	if n.Topology == Pseudo {
		return report.Nodes{n.ID: n}
	}

	// Ignore non-running containers
	if state, ok := n.Latest.Lookup(docker.ContainerState); ok && state != docker.StateRunning {
		return report.Nodes{}
	}

	// Otherwise, if some some reason the container doesn't have a pod uid (maybe
	// slightly out of sync reports, or its not in a pod), make it part of unmanaged.
	uid, ok := n.Latest.Lookup(docker.LabelPrefix + "io.kubernetes.pod.uid")
	if !ok {
		id := MakePseudoNodeID(UnmanagedID, report.ExtractHostID(n))
		node := NewDerivedPseudoNode(id, n)
		return report.Nodes{id: node}
	}

	id := report.MakePodNodeID(uid)
	node := NewDerivedNode(id, n).
		WithTopology(report.Pod)
	node.Counters = node.Counters.Add(n.Topology, 1)
	return report.Nodes{id: node}
}

// The various ways of grouping pods
var (
	Map2Service    = Map2Parent(report.Service)
	Map2Deployment = Map2Parent(report.Deployment)
	Map2ReplicaSet = Map2Parent(report.ReplicaSet)
)

// Map2Parent maps Nodes to some parent grouping.
func Map2Parent(topology string) func(n report.Node, _ report.Networks) report.Nodes {
	return func(n report.Node, _ report.Networks) report.Nodes {
		// Propagate all pseudo nodes
		if n.Topology == Pseudo {
			return report.Nodes{n.ID: n}
		}

		// Otherwise, if some some reason the node doesn't have any of these ids
		// (maybe slightly out of sync reports, or its not in this group), just
		// drop it
		groupIDs, ok := n.Parents.Lookup(topology)
		if !ok {
			return report.Nodes{}
		}

		result := report.Nodes{}
		for _, id := range groupIDs {
			node := NewDerivedNode(id, n).WithTopology(topology)
			node.Counters = node.Counters.Add(n.Topology, 1)

			// When mapping replica(tionController)s(ets) to deployments
			// we must propagate the pod counter.
			if n.Topology != report.Pod {
				if count, ok := n.Counters.Lookup(report.Pod); ok {
					node.Counters = node.Counters.Add(report.Pod, count)
				}
			}

			result[id] = node
		}
		return result
	}
}
<|endoftext|>


<commit_before>package ethereum

import (
	"github.com/SmartPool/smartpool-client"
	"github.com/SmartPool/smartpool-client/protocol"
	"github.com/ethereum/go-ethereum/common"
	"math/big"
	"sync"
)

// TimestampClaimRepo only select shares that don't have most recent timestamp
// in order to make sure coming shares' counters are greater than selected
// shares
type TimestampClaimRepo struct {
	activeShares    []*Share
	recentTimestamp *big.Int
	noShares        uint64
	noRecentShares  uint64
	mu              sync.Mutex
}

// TODO: Load from persistent storage
func NewTimestampClaimRepo() *TimestampClaimRepo {
	return &TimestampClaimRepo{
		[]*Share{},
		big.NewInt(0),
		0,
		0,
		sync.Mutex{},
	}
}

func (cr *TimestampClaimRepo) AddShare(s smartpool.Share) {
	cr.mu.Lock()
	defer cr.mu.Unlock()
	share := s.(*Share)
	cr.activeShares = append(cr.activeShares, share)
	if share.Timestamp().Cmp(cr.recentTimestamp) == 0 {
		cr.noRecentShares++
	} else if share.Timestamp().Cmp(cr.recentTimestamp) < 0 {
		cr.noShares++
	} else if share.Timestamp().Cmp(cr.recentTimestamp) > 0 {
		cr.noShares += cr.noRecentShares
		cr.noRecentShares = 1
		cr.recentTimestamp = big.NewInt(0)
		cr.recentTimestamp.Add(share.Timestamp(), common.Big0)
	}
}

func (cr *TimestampClaimRepo) GetCurrentClaim(threshold int) smartpool.Claim {
	cr.mu.Lock()
	defer cr.mu.Unlock()
	smartpool.Output.Printf("Have %d eligible shares\n", cr.noShares)
	smartpool.Output.Printf("Current timestamp: 0x%s\n", cr.recentTimestamp.Text(16))
	smartpool.Output.Printf("Shares with current timestamp: %d\n", cr.noRecentShares)
	if cr.noShares < uint64(threshold) {
		return nil
	}
	c := protocol.NewClaim()
	newActiveShares := []*Share{}
	for _, s := range cr.activeShares {
		if s.Timestamp().Cmp(cr.recentTimestamp) < 0 {
			c.AddShare(s)
		} else {
			newActiveShares = append(newActiveShares, s)
		}
	}
	cr.activeShares = newActiveShares
	cr.noShares = 0
	return c
}
<commit_msg>fix wording<commit_after>package ethereum

import (
	"github.com/SmartPool/smartpool-client"
	"github.com/SmartPool/smartpool-client/protocol"
	"github.com/ethereum/go-ethereum/common"
	"math/big"
	"sync"
)

// TimestampClaimRepo only select shares that don't have most recent timestamp
// in order to make sure coming shares' counters are greater than selected
// shares
type TimestampClaimRepo struct {
	activeShares    []*Share
	recentTimestamp *big.Int
	noShares        uint64
	noRecentShares  uint64
	mu              sync.Mutex
}

// TODO: Load from persistent storage
func NewTimestampClaimRepo() *TimestampClaimRepo {
	return &TimestampClaimRepo{
		[]*Share{},
		big.NewInt(0),
		0,
		0,
		sync.Mutex{},
	}
}

func (cr *TimestampClaimRepo) AddShare(s smartpool.Share) {
	cr.mu.Lock()
	defer cr.mu.Unlock()
	share := s.(*Share)
	cr.activeShares = append(cr.activeShares, share)
	if share.Timestamp().Cmp(cr.recentTimestamp) == 0 {
		cr.noRecentShares++
	} else if share.Timestamp().Cmp(cr.recentTimestamp) < 0 {
		cr.noShares++
	} else if share.Timestamp().Cmp(cr.recentTimestamp) > 0 {
		cr.noShares += cr.noRecentShares
		cr.noRecentShares = 1
		cr.recentTimestamp = big.NewInt(0)
		cr.recentTimestamp.Add(share.Timestamp(), common.Big0)
	}
}

func (cr *TimestampClaimRepo) GetCurrentClaim(threshold int) smartpool.Claim {
	cr.mu.Lock()
	defer cr.mu.Unlock()
	smartpool.Output.Printf("Have %d valid shares\n", cr.noShares)
	smartpool.Output.Printf("Current timestamp: 0x%s\n", cr.recentTimestamp.Text(16))
	smartpool.Output.Printf("Shares with current timestamp: %d\n", cr.noRecentShares)
	if cr.noShares < uint64(threshold) {
		return nil
	}
	c := protocol.NewClaim()
	newActiveShares := []*Share{}
	for _, s := range cr.activeShares {
		if s.Timestamp().Cmp(cr.recentTimestamp) < 0 {
			c.AddShare(s)
		} else {
			newActiveShares = append(newActiveShares, s)
		}
	}
	cr.activeShares = newActiveShares
	cr.noShares = 0
	return c
}
<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

package simplepush

import (
	"sync"
	"sync/atomic"
)

var (
	Metrics map[string]int64
	metrex  sync.Mutex
)

func init() {
	Metrics = make(map[string]int64)
}

func MetricsSnapshot() map[string]int64 {
	defer metrex.Unlock()
	metrex.Lock()
	var oldMetrics map[string]int64
	// copy the old metrics
	for k, v := range Metrics {
		oldMetrics[k] = v
	}
	return oldMetrics
}

func get(metric string) (m int64) {
	defer metrex.Unlock()
	if m, exists := Metrics[metric]; !exists {
		metrex.Lock()
		m = int64(0)
		Metrics[metric] = m
	}
	return
}

func MetricIncrementBy(metric string, count int) {
	m := get(metric)
	atomic.AddInt64(&m, int64(count))
}

func MetricIncrement(metric string) {
	MetricIncrementBy(metric, 1)
}

func MetricDecrement(metric string) {
	MetricIncrementBy(metric, -1)
}

/*
type MetricsLogger struct {
	logger *util.HekaLogger
}

func NewMetricsLogger(logger *util.HekaLogger) (mlogger *MetricsLogger) {
	mlogger = &MetricsLogger{logger: logger}
	return
}

func (self *MetricsLogger) incr(value int64) int64 {
	atomic.AddInt64(&value, int64(1))
	return value
}
*/
<commit_msg>Fix to metric deref<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

package simplepush

import (
	"sync"
	"sync/atomic"
)

var (
	Metrics map[string]int64
	metrex  sync.Mutex
)

func init() {
	Metrics = make(map[string]int64)
}

func MetricsSnapshot() map[string]int64 {
	defer metrex.Unlock()
	metrex.Lock()
	var oldMetrics map[string]int64
	// copy the old metrics
	for k, v := range Metrics {
		oldMetrics[k] = v
	}
	return oldMetrics
}

func get(metric string) (m int64) {
	if m, exists := Metrics[metric]; !exists {
	    defer metrex.Unlock()
		metrex.Lock()
		m = int64(0)
		Metrics[metric] = m
	}
	return
}

func MetricIncrementBy(metric string, count int) {
	m := get(metric)
	atomic.AddInt64(&m, int64(count))
}

func MetricIncrement(metric string) {
	MetricIncrementBy(metric, 1)
}

func MetricDecrement(metric string) {
	MetricIncrementBy(metric, -1)
}

/*
type MetricsLogger struct {
	logger *util.HekaLogger
}

func NewMetricsLogger(logger *util.HekaLogger) (mlogger *MetricsLogger) {
	mlogger = &MetricsLogger{logger: logger}
	return
}

func (self *MetricsLogger) incr(value int64) int64 {
	atomic.AddInt64(&value, int64(1))
	return value
}
*/
<|endoftext|>


<commit_before>package test

import (
	"flag"
	"fmt"
	"log"
	"os"
	"regexp"
	"testing"
	"time"

	sous "github.com/opentable/sous/lib"
	"github.com/opentable/sous/util/docker_registry"
	"github.com/samsalisbury/semv"
	"github.com/stretchr/testify/assert"
)

var imageName string

func TestMain(m *testing.M) {
	log.Print("hello there")
	flag.Parse()
	os.Exit(wrapCompose(m))
}

func TestGetLabels(t *testing.T) {
	registerLabelledContainers()
	assert := assert.New(t)
	cl := docker_registry.NewClient()
	cl.BecomeFoolishlyTrusting()

	labels, err := cl.LabelsForImageName(imageName)

	assert.Nil(err)
	assert.Contains(labels, sous.DockerRepoLabel)
	resetSingularity()
}

func TestGetRunningDeploymentSet(t *testing.T) {
	//	sous.Log.Debug.SetFlags(sous.Log.Debug.Flags() | log.Ltime)
	//	sous.Log.Debug.SetOutput(os.Stderr)
	//	sous.Log.Debug.Print("Starting stderr output")
	assert := assert.New(t)

	registerLabelledContainers()
	drc := docker_registry.NewClient()
	drc.BecomeFoolishlyTrusting()
	nc := sous.NewNameCache(drc, "sqlite3", sous.InMemoryConnection("grds"))
	ra := sous.NewRectiAgent(nc)

	deps, which := deploymentWithRepo(assert, ra, "https://github.com/opentable/docker-grafana.git")
	if assert.Equal(3, len(deps)) {
		grafana := deps[which]
		assert.Equal(singularityURL, grafana.Cluster)
		assert.Regexp("^0\\.1", grafana.Resources["cpus"])    // XXX strings and floats...
		assert.Regexp("^100\\.", grafana.Resources["memory"]) // XXX strings and floats...
		assert.Equal("1", grafana.Resources["ports"])         // XXX strings and floats...
		assert.Equal(17, grafana.SourceVersion.Version.Patch)
		assert.Equal("91495f1b1630084e301241100ecf2e775f6b672c", grafana.SourceVersion.Version.Meta)
		assert.Equal(1, grafana.NumInstances)
		assert.Equal(sous.ManifestKindService, grafana.Kind)
	}

	resetSingularity()
}

func TestMissingImage(t *testing.T) {
	assert := assert.New(t)

	clusterDefs := sous.Defs{
		Clusters: sous.Clusters{
			singularityURL: sous.Cluster{
				BaseURL: singularityURL,
			},
		},
	}
	repoOne := "https://github.com/opentable/one.git"

	drc := docker_registry.NewClient()
	drc.BecomeFoolishlyTrusting()
	// easiest way to make sure that the manifest doesn't actually get registered
	dummyNc := sous.NewNameCache(drc, "sqlite3", sous.InMemoryConnection("bitbucket"))

	stateOne := sous.State{
		Defs: clusterDefs,
		Manifests: sous.Manifests{
			"one": manifest(dummyNc, "opentable/one", "test-one", repoOne, "1.1.1"),
		},
	}

	// ****
	nc := sous.NewNameCache(drc, "sqlite3", sous.InMemoryConnection("missingimage"))
	ra := sous.NewRectiAgent(nc)
	err := sous.Resolve(ra, stateOne)
	assert.Error(err)

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

	_, which := deploymentWithRepo(assert, ra, repoOne)
	assert.Equal(which, -1, "opentable/one was deployed")

	resetSingularity()
}

func TestResolve(t *testing.T) {
	assert := assert.New(t)

	clusterDefs := sous.Defs{
		Clusters: sous.Clusters{
			singularityURL: sous.Cluster{
				BaseURL: singularityURL,
			},
		},
	}
	repoOne := "https://github.com/opentable/one.git"
	repoTwo := "https://github.com/opentable/two.git"
	repoThree := "https://github.com/opentable/three.git"

	drc := docker_registry.NewClient()
	drc.BecomeFoolishlyTrusting()

	nc := sous.NewNameCache(drc, "sqlite3", sous.InMemoryConnection("testresolve"))
	ra := sous.NewRectiAgent(nc)

	stateOneTwo := sous.State{
		Defs: clusterDefs,
		Manifests: sous.Manifests{
			"one": manifest(nc, "opentable/one", "test-one", repoOne, "1.1.1"),
			"two": manifest(nc, "opentable/two", "test-two", repoTwo, "1.1.1"),
		},
	}
	stateTwoThree := sous.State{
		Defs: clusterDefs,
		Manifests: sous.Manifests{
			"two":   manifest(nc, "opentable/two", "test-two", repoTwo, "1.1.1"),
			"three": manifest(nc, "opentable/three", "test-three", repoThree, "1.1.1"),
		},
	}

	// ****
	log.Print("Resolving from nothing to one+two")
	err := sous.Resolve(ra, stateOneTwo)
	if err != nil {
		assert.Fail(err.Error())
	}
	// ****
	time.Sleep(3 * time.Second)

	deps, which := deploymentWithRepo(assert, ra, repoOne)
	if assert.NotEqual(which, -1, "opentable/one not successfully deployed") {
		one := deps[which]
		assert.Equal(1, one.NumInstances)
	}

	which = findRepo(deps, repoTwo)
	if assert.NotEqual(-1, which, "opentable/two not successfully deployed") {
		two := deps[which]
		assert.Equal(1, two.NumInstances)
	}

	// ****
	log.Println("Resolving from one+two to two+three")
	conflictRE := regexp.MustCompile(`Pending deploy already in progress`)

	// XXX Let's hope this is a temporary solution to a testing issue
	// The problem is laid out in DCOPS-7625
	for tries := 0; tries < 30; tries++ {
		err = sous.Resolve(ra, stateTwoThree)
		if err != nil {
			if !conflictRE.MatchString(err.Error()) {
				assert.Fail(err.Error())
			}
			log.Printf("Singularity conflict - waiting for previous deploy to complete - try #%d", tries+1)
			time.Sleep(1 * time.Second)
		}
	}

	if !assert.NoError(err) {
		assert.Fail(err)
	}
	// ****

	deps, which = deploymentWithRepo(assert, ra, repoTwo)
	if assert.NotEqual(-1, which, "opentable/two no longer deployed after resolve") {
		assert.Equal(1, deps[which].NumInstances)
	}

	which = findRepo(deps, repoThree)
	if assert.NotEqual(-1, which, "opentable/three not successfully deployed") {
		assert.Equal(1, deps[which].NumInstances)
		if assert.Len(deps[which].DeployConfig.Volumes, 1) {
			assert.Equal("RO", string(deps[which].DeployConfig.Volumes[0].Mode))
		}
	}

	which = findRepo(deps, repoOne)
	if which != -1 {
		assert.Equal(0, deps[which].NumInstances)
	}

	resetSingularity()
}

func deploymentWithRepo(assert *assert.Assertions, ra sous.RectificationClient, repo string) (sous.Deployments, int) {
	sc := sous.NewSetCollector(ra)
	deps, err := sc.GetRunningDeployment([]string{singularityURL})
	if assert.Nil(err) {
		return deps, findRepo(deps, repo)
	}
	return sous.Deployments{}, -1
}

func findRepo(deps sous.Deployments, repo string) int {
	for i := range deps {
		if deps[i] != nil {
			if deps[i].SourceVersion.RepoURL == sous.RepoURL(repo) {
				return i
			}
		}
	}
	return -1
}

func manifest(nc sous.ImageMapper, drepo, containerDir, sourceURL, version string) *sous.Manifest {
	sv := sous.SourceVersion{
		RepoURL:    sous.RepoURL(sourceURL),
		RepoOffset: sous.RepoOffset(""),
		Version:    semv.MustParse(version),
	}

	in := buildImageName(drepo, version)
	buildAndPushContainer(containerDir, in)

	nc.Insert(sv, in, "")

	return &sous.Manifest{
		Source: sous.SourceLocation{
			RepoURL:    sous.RepoURL(sourceURL),
			RepoOffset: sous.RepoOffset(""),
		},
		Owners: []string{`xyz`},
		Kind:   sous.ManifestKindService,
		Deployments: sous.DeploySpecs{
			singularityURL: sous.PartialDeploySpec{
				DeployConfig: sous.DeployConfig{
					Resources:    sous.Resources{}, //map[string]string
					Args:         []string{},
					Env:          sous.Env{}, //map[s]s
					NumInstances: 1,
					Volumes: sous.Volumes{
						&sous.Volume{"h", "c", sous.VolumeMode("RO")},
					},
				},
				Version: semv.MustParse(version),
				//clusterName: "it",
			},
		},
	}
}

func registerLabelledContainers() {
	registerAndDeploy(ip, "hello-labels", "hello-labels", []int32{})
	registerAndDeploy(ip, "hello-server-labels", "hello-server-labels", []int32{80})
	registerAndDeploy(ip, "grafana-repo", "grafana-labels", []int32{})
	imageName = fmt.Sprintf("%s/%s:%s", registryName, "grafana-repo", "latest")
}
<commit_msg>Fixes stupid error<commit_after>package test

import (
	"flag"
	"fmt"
	"log"
	"os"
	"regexp"
	"testing"
	"time"

	sous "github.com/opentable/sous/lib"
	"github.com/opentable/sous/util/docker_registry"
	"github.com/samsalisbury/semv"
	"github.com/stretchr/testify/assert"
)

var imageName string

func TestMain(m *testing.M) {
	log.Print("hello there")
	flag.Parse()
	os.Exit(wrapCompose(m))
}

func TestGetLabels(t *testing.T) {
	registerLabelledContainers()
	assert := assert.New(t)
	cl := docker_registry.NewClient()
	cl.BecomeFoolishlyTrusting()

	labels, err := cl.LabelsForImageName(imageName)

	assert.Nil(err)
	assert.Contains(labels, sous.DockerRepoLabel)
	resetSingularity()
}

func TestGetRunningDeploymentSet(t *testing.T) {
	//	sous.Log.Debug.SetFlags(sous.Log.Debug.Flags() | log.Ltime)
	//	sous.Log.Debug.SetOutput(os.Stderr)
	//	sous.Log.Debug.Print("Starting stderr output")
	assert := assert.New(t)

	registerLabelledContainers()
	drc := docker_registry.NewClient()
	drc.BecomeFoolishlyTrusting()
	nc := sous.NewNameCache(drc, "sqlite3", sous.InMemoryConnection("grds"))
	ra := sous.NewRectiAgent(nc)

	deps, which := deploymentWithRepo(assert, ra, "https://github.com/opentable/docker-grafana.git")
	if assert.Equal(3, len(deps)) {
		grafana := deps[which]
		assert.Equal(singularityURL, grafana.Cluster)
		assert.Regexp("^0\\.1", grafana.Resources["cpus"])    // XXX strings and floats...
		assert.Regexp("^100\\.", grafana.Resources["memory"]) // XXX strings and floats...
		assert.Equal("1", grafana.Resources["ports"])         // XXX strings and floats...
		assert.Equal(17, grafana.SourceVersion.Version.Patch)
		assert.Equal("91495f1b1630084e301241100ecf2e775f6b672c", grafana.SourceVersion.Version.Meta)
		assert.Equal(1, grafana.NumInstances)
		assert.Equal(sous.ManifestKindService, grafana.Kind)
	}

	resetSingularity()
}

func TestMissingImage(t *testing.T) {
	assert := assert.New(t)

	clusterDefs := sous.Defs{
		Clusters: sous.Clusters{
			singularityURL: sous.Cluster{
				BaseURL: singularityURL,
			},
		},
	}
	repoOne := "https://github.com/opentable/one.git"

	drc := docker_registry.NewClient()
	drc.BecomeFoolishlyTrusting()
	// easiest way to make sure that the manifest doesn't actually get registered
	dummyNc := sous.NewNameCache(drc, "sqlite3", sous.InMemoryConnection("bitbucket"))

	stateOne := sous.State{
		Defs: clusterDefs,
		Manifests: sous.Manifests{
			"one": manifest(dummyNc, "opentable/one", "test-one", repoOne, "1.1.1"),
		},
	}

	// ****
	nc := sous.NewNameCache(drc, "sqlite3", sous.InMemoryConnection("missingimage"))
	ra := sous.NewRectiAgent(nc)
	err := sous.Resolve(ra, stateOne)
	assert.Error(err)

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

	_, which := deploymentWithRepo(assert, ra, repoOne)
	assert.Equal(which, -1, "opentable/one was deployed")

	resetSingularity()
}

func TestResolve(t *testing.T) {
	assert := assert.New(t)

	clusterDefs := sous.Defs{
		Clusters: sous.Clusters{
			singularityURL: sous.Cluster{
				BaseURL: singularityURL,
			},
		},
	}
	repoOne := "https://github.com/opentable/one.git"
	repoTwo := "https://github.com/opentable/two.git"
	repoThree := "https://github.com/opentable/three.git"

	drc := docker_registry.NewClient()
	drc.BecomeFoolishlyTrusting()

	nc := sous.NewNameCache(drc, "sqlite3", sous.InMemoryConnection("testresolve"))
	ra := sous.NewRectiAgent(nc)

	stateOneTwo := sous.State{
		Defs: clusterDefs,
		Manifests: sous.Manifests{
			"one": manifest(nc, "opentable/one", "test-one", repoOne, "1.1.1"),
			"two": manifest(nc, "opentable/two", "test-two", repoTwo, "1.1.1"),
		},
	}
	stateTwoThree := sous.State{
		Defs: clusterDefs,
		Manifests: sous.Manifests{
			"two":   manifest(nc, "opentable/two", "test-two", repoTwo, "1.1.1"),
			"three": manifest(nc, "opentable/three", "test-three", repoThree, "1.1.1"),
		},
	}

	// ****
	log.Print("Resolving from nothing to one+two")
	err := sous.Resolve(ra, stateOneTwo)
	if err != nil {
		assert.Fail(err.Error())
	}
	// ****
	time.Sleep(3 * time.Second)

	deps, which := deploymentWithRepo(assert, ra, repoOne)
	if assert.NotEqual(which, -1, "opentable/one not successfully deployed") {
		one := deps[which]
		assert.Equal(1, one.NumInstances)
	}

	which = findRepo(deps, repoTwo)
	if assert.NotEqual(-1, which, "opentable/two not successfully deployed") {
		two := deps[which]
		assert.Equal(1, two.NumInstances)
	}

	// ****
	log.Println("Resolving from one+two to two+three")
	conflictRE := regexp.MustCompile(`Pending deploy already in progress`)

	// XXX Let's hope this is a temporary solution to a testing issue
	// The problem is laid out in DCOPS-7625
	for tries := 0; tries < 30; tries++ {
		err = sous.Resolve(ra, stateTwoThree)
		if err != nil {
			if !conflictRE.MatchString(err.Error()) {
				assert.Fail(err.Error())
			}
			log.Printf("Singularity conflict - waiting for previous deploy to complete - try #%d", tries+1)
			time.Sleep(1 * time.Second)
		}
	}

	if !assert.NoError(err) {
		assert.Fail(err.Error())
	}
	// ****

	deps, which = deploymentWithRepo(assert, ra, repoTwo)
	if assert.NotEqual(-1, which, "opentable/two no longer deployed after resolve") {
		assert.Equal(1, deps[which].NumInstances)
	}

	which = findRepo(deps, repoThree)
	if assert.NotEqual(-1, which, "opentable/three not successfully deployed") {
		assert.Equal(1, deps[which].NumInstances)
		if assert.Len(deps[which].DeployConfig.Volumes, 1) {
			assert.Equal("RO", string(deps[which].DeployConfig.Volumes[0].Mode))
		}
	}

	which = findRepo(deps, repoOne)
	if which != -1 {
		assert.Equal(0, deps[which].NumInstances)
	}

	resetSingularity()
}

func deploymentWithRepo(assert *assert.Assertions, ra sous.RectificationClient, repo string) (sous.Deployments, int) {
	sc := sous.NewSetCollector(ra)
	deps, err := sc.GetRunningDeployment([]string{singularityURL})
	if assert.Nil(err) {
		return deps, findRepo(deps, repo)
	}
	return sous.Deployments{}, -1
}

func findRepo(deps sous.Deployments, repo string) int {
	for i := range deps {
		if deps[i] != nil {
			if deps[i].SourceVersion.RepoURL == sous.RepoURL(repo) {
				return i
			}
		}
	}
	return -1
}

func manifest(nc sous.ImageMapper, drepo, containerDir, sourceURL, version string) *sous.Manifest {
	sv := sous.SourceVersion{
		RepoURL:    sous.RepoURL(sourceURL),
		RepoOffset: sous.RepoOffset(""),
		Version:    semv.MustParse(version),
	}

	in := buildImageName(drepo, version)
	buildAndPushContainer(containerDir, in)

	nc.Insert(sv, in, "")

	return &sous.Manifest{
		Source: sous.SourceLocation{
			RepoURL:    sous.RepoURL(sourceURL),
			RepoOffset: sous.RepoOffset(""),
		},
		Owners: []string{`xyz`},
		Kind:   sous.ManifestKindService,
		Deployments: sous.DeploySpecs{
			singularityURL: sous.PartialDeploySpec{
				DeployConfig: sous.DeployConfig{
					Resources:    sous.Resources{}, //map[string]string
					Args:         []string{},
					Env:          sous.Env{}, //map[s]s
					NumInstances: 1,
					Volumes: sous.Volumes{
						&sous.Volume{"h", "c", sous.VolumeMode("RO")},
					},
				},
				Version: semv.MustParse(version),
				//clusterName: "it",
			},
		},
	}
}

func registerLabelledContainers() {
	registerAndDeploy(ip, "hello-labels", "hello-labels", []int32{})
	registerAndDeploy(ip, "hello-server-labels", "hello-server-labels", []int32{80})
	registerAndDeploy(ip, "grafana-repo", "grafana-labels", []int32{})
	imageName = fmt.Sprintf("%s/%s:%s", registryName, "grafana-repo", "latest")
}
<|endoftext|>


<commit_before>// (C) 2014 Cybozu.  All rights reserved.
// Use of this source code is governed by a BSD-style license
// that can be found in the LICENSE file.

package kintone

import (
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"strconv"
	"strings"
	"time"
)

// Record represens a record in an application.
//
// Fields is a mapping between field IDs and fields.
// Although field types are shown as interface{}, they are guaranteed
// to be one of a *Field type in this package.
type Record struct {
	id       uint64
	revision int64
	Fields   map[string]interface{}
}

// NewRecord creates an instance of Record.
//
// The revision number is initialized to -1.
func NewRecord(fields map[string]interface{}) *Record {
	return &Record{0, -1, fields}
}

// MarshalJSON marshals field data of a record into JSON.
func (rec Record) MarshalJSON() ([]byte, error) {
	return json.Marshal(rec.Fields)
}

// Id returns the record number.
//
// A record number is unique within an application.
func (rec Record) Id() uint64 {
	return rec.id
}

// Revision returns the record revision number.
func (rec Record) Revision() int64 {
	return rec.revision
}

// Assert string list.
func stringList(l []interface{}) []string {
	sl := make([]string, len(l))
	for i, v := range l {
		sl[i] = v.(string)
	}
	return sl
}

// Convert user list.
func userList(l []interface{}) ([]User, error) {
	b, err := json.Marshal(l)
	if err != nil {
		return nil, err
	}
	var ul []User
	err = json.Unmarshal(b, &ul)
	if err != nil {
		return nil, err
	}
	return ul, nil
}

// Convert string "record number" into an integer.
func numericId(id string) (uint64, error) {
	n := strings.LastIndex(id, "-")
	if n != -1 {
		id = id[(n + 1):]
	}
	nid, err := strconv.ParseUint(id, 10, 64)
	if err != nil {
		return 0, err
	}
	return nid, nil
}

type recordData map[string]struct {
	Type  string      `json:"type"`
	Value interface{} `json:"value"`
}

func decodeRecordData(data recordData) (*Record, error) {
	fields := make(map[string]interface{})
	rec := &Record{0, -1, fields}
	for key, v := range data {
		switch v.Type {
		case FT_SINGLE_LINE_TEXT:
			fields[key] = SingleLineTextField(v.Value.(string))
		case FT_MULTI_LINE_TEXT:
			fields[key] = MultiLineTextField(v.Value.(string))
		case FT_RICH_TEXT:
			fields[key] = RichTextField(v.Value.(string))
		case FT_DECIMAL:
			fields[key] = DecimalField(v.Value.(string))
		case FT_CALC:
			fields[key] = CalcField(v.Value.(string))
		case FT_CHECK_BOX:
			fields[key] = CheckBoxField(stringList(v.Value.([]interface{})))
		case FT_RADIO:
			if v.Value == nil {
				fields[key] = RadioButtonField("")
			} else {
				fields[key] = RadioButtonField(v.Value.(string))
			}
		case FT_SINGLE_SELECT:
			if v.Value == nil {
				fields[key] = SingleSelectField{Valid: false}
			} else {
				fields[key] = SingleSelectField{v.Value.(string), true}
			}
		case FT_MULTI_SELECT:
			fields[key] = MultiSelectField(stringList(v.Value.([]interface{})))
		case FT_FILE:
			b1, err := json.Marshal(v.Value)
			if err != nil {
				return nil, err
			}
			var fl []File
			err = json.Unmarshal(b1, &fl)
			if err != nil {
				return nil, err
			}
			fields[key] = FileField(fl)
		case FT_LINK:
			fields[key] = LinkField(v.Value.(string))
		case FT_DATE:
			if v.Value == nil || v.Value == "" {
				fields[key] = DateField{Valid: false}
			} else {
				d, err := time.Parse("2006-01-02", v.Value.(string))
				if err != nil {
					return nil, fmt.Errorf("Invalid date: %v", v.Value)
				}
				fields[key] = DateField{d, true}
			}
		case FT_TIME:
			if v.Value == nil || v.Value == "" {
				fields[key] = TimeField{Valid: false}
			} else {
				t, err := time.Parse("15:04", v.Value.(string))
				if err != nil {
					t, err = time.Parse("15:04:05", v.Value.(string))
					if err != nil {
						return nil, fmt.Errorf("Invalid time: %v", v.Value)
					}
				}
				fields[key] = TimeField{t, true}
			}
		case FT_DATETIME:
			if v.Value == nil || v.Value == "" {
				fields[key] = DateTimeField{Valid: false}
			} else {
				if s, ok := v.Value.(string); ok {
					dt, err := time.Parse(time.RFC3339, s)
					if err != nil {
						return nil, fmt.Errorf("Invalid datetime: %v", v.Value)
					}
					fields[key] = DateTimeField{dt, true}
				}
			}
		case FT_USER:
			ul, err := userList(v.Value.([]interface{}))
			if err != nil {
				return nil, err
			}
			fields[key] = UserField(ul)
		case FT_CATEGORY:
			fields[key] = CategoryField(stringList(v.Value.([]interface{})))
		case FT_STATUS:
			fields[key] = StatusField(v.Value.(string))
		case FT_ASSIGNEE:
			al, err := userList(v.Value.([]interface{}))
			if err != nil {
				return nil, err
			}
			fields[key] = AssigneeField(al)
		case FT_RECNUM:
			if nid, err := numericId(v.Value.(string)); err != nil {
				return nil, err
			} else {
				rec.id = nid
			}
			fields[key] = RecordNumberField(v.Value.(string))
		case FT_CREATOR:
			creator := v.Value.(map[string]interface{})
			fields[key] = CreatorField{
				creator["code"].(string),
				creator["name"].(string),
			}
		case FT_CTIME:
			var ctime time.Time
			if ctime.UnmarshalText([]byte(v.Value.(string))) != nil {
				return nil, fmt.Errorf("Invalid datetime: %v", v.Value)
			}
			fields[key] = CreationTimeField(ctime)
		case FT_MODIFIER:
			modifier := v.Value.(map[string]interface{})
			fields[key] = ModifierField{
				modifier["code"].(string),
				modifier["name"].(string),
			}
		case FT_MTIME:
			var mtime time.Time
			if mtime.UnmarshalText([]byte(v.Value.(string))) != nil {
				return nil, fmt.Errorf("Invalid datetime: %v", v.Value)
			}
			fields[key] = CreationTimeField(mtime)
		case FT_SUBTABLE:
			b2, err := json.Marshal(v.Value)
			if err != nil {
				return nil, err
			}
			var stl []SubTableEntry
			err = json.Unmarshal(b2, &stl)
			if err != nil {
				return nil, err
			}
			fields[key] = SubTableField(stl)
		case FT_ID:
			id, err := strconv.ParseUint(v.Value.(string), 10, 64)
			if err != nil {
				return nil, fmt.Errorf("Invalid record ID: %v", v.Value)
			}
			rec.id = id
		case FT_REVISION:
			revision, err := strconv.ParseInt(v.Value.(string), 10, 64)
			if err != nil {
				return nil, fmt.Errorf("Invalid revision number: %v", v.Value)
			}
			rec.revision = revision
		default:
			log.Printf("Unknown type: %v", v.Type)
		}
	}
	return rec, nil
}

// DecodeRecords decodes JSON response for multi-get API.
func DecodeRecords(b []byte) ([]*Record, error) {
	var t struct {
		Records []recordData `json:"records"`
	}
	err := json.Unmarshal(b, &t)
	if err != nil {
		return nil, errors.New("Invalid JSON format")
	}
	rec_list := make([]*Record, len(t.Records))
	for i, rd := range t.Records {
		r, err := decodeRecordData(rd)
		if err != nil {
			return nil, err
		}
		rec_list[i] = r
	}
	return rec_list, nil
}

// DecodeRecord decodes JSON response for single-get API.
func DecodeRecord(b []byte) (*Record, error) {
	var t struct {
		RecordData recordData `json:"record"`
	}
	err := json.Unmarshal(b, &t)
	if err != nil {
		return nil, errors.New("Invalid JSON format")
	}
	rec, err := decodeRecordData(t.RecordData)
	if err != nil {
		return nil, err
	}
	return rec, nil
}
<commit_msg>add NewRecordWithId method<commit_after>// (C) 2014 Cybozu.  All rights reserved.
// Use of this source code is governed by a BSD-style license
// that can be found in the LICENSE file.

package kintone

import (
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"strconv"
	"strings"
	"time"
)

// Record represens a record in an application.
//
// Fields is a mapping between field IDs and fields.
// Although field types are shown as interface{}, they are guaranteed
// to be one of a *Field type in this package.
type Record struct {
	id       uint64
	revision int64
	Fields   map[string]interface{}
}

// NewRecord creates an instance of Record.
//
// The revision number is initialized to -1.
func NewRecord(fields map[string]interface{}) *Record {
	return &Record{0, -1, fields}
}

// NewRecord creates using an existing record id.
//
// The revision number is initialized to -1.
func NewRecordWithId(id uint64, fields map[string]interface{}) *Record {
	return &Record{id, -1, fields}
}

// MarshalJSON marshals field data of a record into JSON.
func (rec Record) MarshalJSON() ([]byte, error) {
	return json.Marshal(rec.Fields)
}

// Id returns the record number.
//
// A record number is unique within an application.
func (rec Record) Id() uint64 {
	return rec.id
}

// Revision returns the record revision number.
func (rec Record) Revision() int64 {
	return rec.revision
}

// Assert string list.
func stringList(l []interface{}) []string {
	sl := make([]string, len(l))
	for i, v := range l {
		sl[i] = v.(string)
	}
	return sl
}

// Convert user list.
func userList(l []interface{}) ([]User, error) {
	b, err := json.Marshal(l)
	if err != nil {
		return nil, err
	}
	var ul []User
	err = json.Unmarshal(b, &ul)
	if err != nil {
		return nil, err
	}
	return ul, nil
}

// Convert string "record number" into an integer.
func numericId(id string) (uint64, error) {
	n := strings.LastIndex(id, "-")
	if n != -1 {
		id = id[(n + 1):]
	}
	nid, err := strconv.ParseUint(id, 10, 64)
	if err != nil {
		return 0, err
	}
	return nid, nil
}

type recordData map[string]struct {
	Type  string      `json:"type"`
	Value interface{} `json:"value"`
}

func decodeRecordData(data recordData) (*Record, error) {
	fields := make(map[string]interface{})
	rec := &Record{0, -1, fields}
	for key, v := range data {
		switch v.Type {
		case FT_SINGLE_LINE_TEXT:
			fields[key] = SingleLineTextField(v.Value.(string))
		case FT_MULTI_LINE_TEXT:
			fields[key] = MultiLineTextField(v.Value.(string))
		case FT_RICH_TEXT:
			fields[key] = RichTextField(v.Value.(string))
		case FT_DECIMAL:
			fields[key] = DecimalField(v.Value.(string))
		case FT_CALC:
			fields[key] = CalcField(v.Value.(string))
		case FT_CHECK_BOX:
			fields[key] = CheckBoxField(stringList(v.Value.([]interface{})))
		case FT_RADIO:
			if v.Value == nil {
				fields[key] = RadioButtonField("")
			} else {
				fields[key] = RadioButtonField(v.Value.(string))
			}
		case FT_SINGLE_SELECT:
			if v.Value == nil {
				fields[key] = SingleSelectField{Valid: false}
			} else {
				fields[key] = SingleSelectField{v.Value.(string), true}
			}
		case FT_MULTI_SELECT:
			fields[key] = MultiSelectField(stringList(v.Value.([]interface{})))
		case FT_FILE:
			b1, err := json.Marshal(v.Value)
			if err != nil {
				return nil, err
			}
			var fl []File
			err = json.Unmarshal(b1, &fl)
			if err != nil {
				return nil, err
			}
			fields[key] = FileField(fl)
		case FT_LINK:
			fields[key] = LinkField(v.Value.(string))
		case FT_DATE:
			if v.Value == nil || v.Value == "" {
				fields[key] = DateField{Valid: false}
			} else {
				d, err := time.Parse("2006-01-02", v.Value.(string))
				if err != nil {
					return nil, fmt.Errorf("Invalid date: %v", v.Value)
				}
				fields[key] = DateField{d, true}
			}
		case FT_TIME:
			if v.Value == nil || v.Value == "" {
				fields[key] = TimeField{Valid: false}
			} else {
				t, err := time.Parse("15:04", v.Value.(string))
				if err != nil {
					t, err = time.Parse("15:04:05", v.Value.(string))
					if err != nil {
						return nil, fmt.Errorf("Invalid time: %v", v.Value)
					}
				}
				fields[key] = TimeField{t, true}
			}
		case FT_DATETIME:
			if v.Value == nil || v.Value == "" {
				fields[key] = DateTimeField{Valid: false}
			} else {
				if s, ok := v.Value.(string); ok {
					dt, err := time.Parse(time.RFC3339, s)
					if err != nil {
						return nil, fmt.Errorf("Invalid datetime: %v", v.Value)
					}
					fields[key] = DateTimeField{dt, true}
				}
			}
		case FT_USER:
			ul, err := userList(v.Value.([]interface{}))
			if err != nil {
				return nil, err
			}
			fields[key] = UserField(ul)
		case FT_CATEGORY:
			fields[key] = CategoryField(stringList(v.Value.([]interface{})))
		case FT_STATUS:
			fields[key] = StatusField(v.Value.(string))
		case FT_ASSIGNEE:
			al, err := userList(v.Value.([]interface{}))
			if err != nil {
				return nil, err
			}
			fields[key] = AssigneeField(al)
		case FT_RECNUM:
			if nid, err := numericId(v.Value.(string)); err != nil {
				return nil, err
			} else {
				rec.id = nid
			}
			fields[key] = RecordNumberField(v.Value.(string))
		case FT_CREATOR:
			creator := v.Value.(map[string]interface{})
			fields[key] = CreatorField{
				creator["code"].(string),
				creator["name"].(string),
			}
		case FT_CTIME:
			var ctime time.Time
			if ctime.UnmarshalText([]byte(v.Value.(string))) != nil {
				return nil, fmt.Errorf("Invalid datetime: %v", v.Value)
			}
			fields[key] = CreationTimeField(ctime)
		case FT_MODIFIER:
			modifier := v.Value.(map[string]interface{})
			fields[key] = ModifierField{
				modifier["code"].(string),
				modifier["name"].(string),
			}
		case FT_MTIME:
			var mtime time.Time
			if mtime.UnmarshalText([]byte(v.Value.(string))) != nil {
				return nil, fmt.Errorf("Invalid datetime: %v", v.Value)
			}
			fields[key] = CreationTimeField(mtime)
		case FT_SUBTABLE:
			b2, err := json.Marshal(v.Value)
			if err != nil {
				return nil, err
			}
			var stl []SubTableEntry
			err = json.Unmarshal(b2, &stl)
			if err != nil {
				return nil, err
			}
			fields[key] = SubTableField(stl)
		case FT_ID:
			id, err := strconv.ParseUint(v.Value.(string), 10, 64)
			if err != nil {
				return nil, fmt.Errorf("Invalid record ID: %v", v.Value)
			}
			rec.id = id
		case FT_REVISION:
			revision, err := strconv.ParseInt(v.Value.(string), 10, 64)
			if err != nil {
				return nil, fmt.Errorf("Invalid revision number: %v", v.Value)
			}
			rec.revision = revision
		default:
			log.Printf("Unknown type: %v", v.Type)
		}
	}
	return rec, nil
}

// DecodeRecords decodes JSON response for multi-get API.
func DecodeRecords(b []byte) ([]*Record, error) {
	var t struct {
		Records []recordData `json:"records"`
	}
	err := json.Unmarshal(b, &t)
	if err != nil {
		return nil, errors.New("Invalid JSON format")
	}
	rec_list := make([]*Record, len(t.Records))
	for i, rd := range t.Records {
		r, err := decodeRecordData(rd)
		if err != nil {
			return nil, err
		}
		rec_list[i] = r
	}
	return rec_list, nil
}

// DecodeRecord decodes JSON response for single-get API.
func DecodeRecord(b []byte) (*Record, error) {
	var t struct {
		RecordData recordData `json:"record"`
	}
	err := json.Unmarshal(b, &t)
	if err != nil {
		return nil, errors.New("Invalid JSON format")
	}
	rec, err := decodeRecordData(t.RecordData)
	if err != nil {
		return nil, err
	}
	return rec, nil
}
<|endoftext|>


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

package txtdirect

import (
	"context"
	"fmt"
	"log"
	"net/http"
	"net/url"
	"strconv"
	"strings"
)

type record struct {
	Version string
	To      string
	Code    int
	Type    string
	Use     []string
	Vcs     string
	Website string
	From    string
	Root    string
	Re      string
	Ref     bool
	Headers map[string]string
}

// getRecord uses the given host to find a TXT record
// and then parses the txt record and returns a TXTDirect record
// struct instance. It returns an error when it can't find any txt
// records or if the TXT record is not standard.
func getRecord(host string, c Config, w http.ResponseWriter, r *http.Request) (record, error) {
	txts, err := query(host, r.Context(), c)
	if err != nil {
		log.Printf("Initial DNS query failed: %s", err)
	}
	// if error present or record empty, jump into wildcards
	if err != nil || txts[0] == "" {
		hostSlice := strings.Split(host, ".")
		hostSlice[0] = "_"
		host = strings.Join(hostSlice, ".")
		txts, err = query(host, r.Context(), c)
		if err != nil {
			log.Printf("Wildcard DNS query failed: %s", err.Error())
			return record{}, err
		}
	}

	if len(txts) != 1 {
		return record{}, fmt.Errorf("could not parse TXT record with %d records", len(txts))
	}

	var rec record
	if rec, err = ParseRecord(txts[0], w, r, c); err != nil {
		return rec, fmt.Errorf("could not parse record: %s", err)
	}

	r = rec.addToContext(r)

	// Add the headers from record to the response
	if len(rec.Headers) != 0 {
		for header, val := range rec.Headers {
			w.Header().Set(header, val)
		}
	}

	return rec, nil
}

// ParseRecord takes a string containing the DNS TXT record and returns
// a TXTDirect record struct instance.
// It will return an error if the DNS TXT record is not standard or
// if the record type is not enabled in the TXTDirect's config.
func ParseRecord(str string, w http.ResponseWriter, req *http.Request, c Config) (record, error) {
	r := record{
		Headers: map[string]string{},
	}

	s := strings.Split(str, ";")

	// Trim whitespace both leading and trailing
	for i := range s {
		s[i] = strings.TrimSpace(s[i])
	}

	for _, l := range s {
		switch {
		case strings.HasPrefix(l, "code="):
			l = strings.TrimPrefix(l, "code=")
			i, err := strconv.Atoi(l)
			if err != nil {
				return record{}, fmt.Errorf("could not parse status code: %s", err)
			}
			r.Code = i

		case strings.HasPrefix(l, "from="):
			l = strings.TrimPrefix(l, "from=")
			l, err := parsePlaceholders(l, req, []string{})
			if err != nil {
				return record{}, err
			}
			r.From = l

		case strings.HasPrefix(l, "re="):
			l = strings.TrimPrefix(l, "re=")
			r.Re = l

		case strings.HasPrefix(l, "ref="):
			l, err := strconv.ParseBool(strings.TrimPrefix(l, "ref="))
			if err != nil {
				fallback(w, req, "global", http.StatusMovedPermanently, c)
				return record{}, err
			}
			r.Ref = l

		case strings.HasPrefix(l, "root="):
			l = strings.TrimPrefix(l, "root=")
			l = ParseURI(l, w, req, c)
			r.Root = l

		case strings.HasPrefix(l, "to="):
			l = strings.TrimPrefix(l, "to=")
			l, err := parsePlaceholders(l, req, []string{})
			if err != nil {
				return record{}, err
			}
			l = ParseURI(l, w, req, c)
			r.To = l

		case strings.HasPrefix(l, "type="):
			l = strings.TrimPrefix(l, "type=")
			r.Type = l

		case strings.HasPrefix(l, "use="):
			l = strings.TrimPrefix(l, "use=")
			if !strings.HasPrefix(l, "_redirect.") {
				return record{}, fmt.Errorf("The given zone address is invalid")
			}
			r.Use = append(r.Use, l)

		case strings.HasPrefix(l, "v="):
			l = strings.TrimPrefix(l, "v=")
			r.Version = l
			if r.Version != "txtv0" {
				return record{}, fmt.Errorf("unhandled version '%s'", r.Version)
			}
			log.Print("WARN: txtv0 is not suitable for production")

		case strings.HasPrefix(l, "vcs="):
			l = strings.TrimPrefix(l, "vcs=")
			r.Vcs = l

		case strings.HasPrefix(l, "website="):
			l = strings.TrimPrefix(l, "website=")
			l = ParseURI(l, w, req, c)
			r.Website = l
		case strings.HasPrefix(l, ">"):
			header := strings.Split(l, "=")
			h, err := url.PathUnescape(header[1])
			if err != nil {
				return record{}, err
			}
			r.Headers[header[0][1:]] = h
		default:
			tuple := strings.Split(l, "=")
			if len(tuple) != 2 {
				return record{}, fmt.Errorf("arbitrary data not allowed")
			}
			continue
		}
		if len(l) > 255 {
			return record{}, fmt.Errorf("TXT record cannot exceed the maximum of 255 characters")
		}
		if r.Type == "dockerv2" && r.To == "" {
			return record{}, fmt.Errorf("[txtdirect]: to= field is required in dockerv2 type")
		}
	}

	if r.Code == 0 {
		r.Code = http.StatusFound
	}

	if r.Type == "" {
		r.Type = "host"
	}

	if r.Type == "host" && r.To == "" {
		fallback(w, req, "global", http.StatusMovedPermanently, c)
		return record{}, nil
	}

	if !contains(c.Enable, r.Type) {
		return record{}, fmt.Errorf("%s type is not enabled in configuration", r.Type)
	}

	return r, nil
}

// Adds the given record to the request's context with "records" key.
func (rec record) addToContext(r *http.Request) *http.Request {
	// Fetch fallback config from context and add the record to it
	recordsContext := r.Context().Value("records")

	// Create a new records field in the context if it doesn't exist
	if recordsContext == nil {
		return r.WithContext(context.WithValue(r.Context(), "records", []record{rec}))
	}

	records := append(recordsContext.([]record), rec)

	// Replace the fallback config instance inside the request's context
	return r.WithContext(context.WithValue(r.Context(), "records", records))
}

// ParseURI parses the given URI and triggers fallback if the URI isn't valid
func ParseURI(uri string, w http.ResponseWriter, r *http.Request, c Config) string {
	url, err := url.Parse(uri)
	if err != nil {
		fallback(w, r, "global", http.StatusMovedPermanently, c)
		return ""
	}
	return url.String()
}

// UpstreamRecord will check all of the use= fields and sends a request to each
// upstream zone address and choses the first one that returns the final TXT
// record
func (rec *record) UpstreamRecord(c Config, w http.ResponseWriter, r *http.Request) (record, string, error) {
	var upstreamRec record
	var err error

	for _, zone := range rec.Use {
		upstreamRec, err = getRecord(zone, c, w, r)
		if err != nil {
			continue
		}
		return upstreamRec, zone, nil
	}

	return record{}, "", fmt.Errorf("Couldn't find any records from upstream")
}
<commit_msg>(record): Add defaults to records without use=<commit_after>/*
Copyright 2019 - The TXTDirect Authors
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
    http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package txtdirect

import (
	"context"
	"fmt"
	"log"
	"net/http"
	"net/url"
	"strconv"
	"strings"
)

type record struct {
	Version string
	To      string
	Code    int
	Type    string
	Use     []string
	Vcs     string
	Website string
	From    string
	Root    string
	Re      string
	Ref     bool
	Headers map[string]string
}

// getRecord uses the given host to find a TXT record
// and then parses the txt record and returns a TXTDirect record
// struct instance. It returns an error when it can't find any txt
// records or if the TXT record is not standard.
func getRecord(host string, c Config, w http.ResponseWriter, r *http.Request) (record, error) {
	txts, err := query(host, r.Context(), c)
	if err != nil {
		log.Printf("Initial DNS query failed: %s", err)
	}
	// if error present or record empty, jump into wildcards
	if err != nil || txts[0] == "" {
		hostSlice := strings.Split(host, ".")
		hostSlice[0] = "_"
		host = strings.Join(hostSlice, ".")
		txts, err = query(host, r.Context(), c)
		if err != nil {
			log.Printf("Wildcard DNS query failed: %s", err.Error())
			return record{}, err
		}
	}

	if len(txts) != 1 {
		return record{}, fmt.Errorf("could not parse TXT record with %d records", len(txts))
	}

	var rec record
	if rec, err = ParseRecord(txts[0], w, r, c); err != nil {
		return rec, fmt.Errorf("could not parse record: %s", err)
	}

	r = rec.addToContext(r)

	// Add the headers from record to the response
	if len(rec.Headers) != 0 {
		for header, val := range rec.Headers {
			w.Header().Set(header, val)
		}
	}

	return rec, nil
}

// ParseRecord takes a string containing the DNS TXT record and returns
// a TXTDirect record struct instance.
// It will return an error if the DNS TXT record is not standard or
// if the record type is not enabled in the TXTDirect's config.
func ParseRecord(str string, w http.ResponseWriter, req *http.Request, c Config) (record, error) {
	r := record{
		Headers: map[string]string{},
	}

	s := strings.Split(str, ";")

	// Trim whitespace both leading and trailing
	for i := range s {
		s[i] = strings.TrimSpace(s[i])
	}

	for _, l := range s {
		switch {
		case strings.HasPrefix(l, "code="):
			l = strings.TrimPrefix(l, "code=")
			i, err := strconv.Atoi(l)
			if err != nil {
				return record{}, fmt.Errorf("could not parse status code: %s", err)
			}
			r.Code = i

		case strings.HasPrefix(l, "from="):
			l = strings.TrimPrefix(l, "from=")
			l, err := parsePlaceholders(l, req, []string{})
			if err != nil {
				return record{}, err
			}
			r.From = l

		case strings.HasPrefix(l, "re="):
			l = strings.TrimPrefix(l, "re=")
			r.Re = l

		case strings.HasPrefix(l, "ref="):
			l, err := strconv.ParseBool(strings.TrimPrefix(l, "ref="))
			if err != nil {
				fallback(w, req, "global", http.StatusMovedPermanently, c)
				return record{}, err
			}
			r.Ref = l

		case strings.HasPrefix(l, "root="):
			l = strings.TrimPrefix(l, "root=")
			l = ParseURI(l, w, req, c)
			r.Root = l

		case strings.HasPrefix(l, "to="):
			l = strings.TrimPrefix(l, "to=")
			l, err := parsePlaceholders(l, req, []string{})
			if err != nil {
				return record{}, err
			}
			l = ParseURI(l, w, req, c)
			r.To = l

		case strings.HasPrefix(l, "type="):
			l = strings.TrimPrefix(l, "type=")
			r.Type = l

		case strings.HasPrefix(l, "use="):
			l = strings.TrimPrefix(l, "use=")
			if !strings.HasPrefix(l, "_redirect.") {
				return record{}, fmt.Errorf("The given zone address is invalid")
			}
			r.Use = append(r.Use, l)

		case strings.HasPrefix(l, "v="):
			l = strings.TrimPrefix(l, "v=")
			r.Version = l
			if r.Version != "txtv0" {
				return record{}, fmt.Errorf("unhandled version '%s'", r.Version)
			}
			log.Print("WARN: txtv0 is not suitable for production")

		case strings.HasPrefix(l, "vcs="):
			l = strings.TrimPrefix(l, "vcs=")
			r.Vcs = l

		case strings.HasPrefix(l, "website="):
			l = strings.TrimPrefix(l, "website=")
			l = ParseURI(l, w, req, c)
			r.Website = l
		case strings.HasPrefix(l, ">"):
			header := strings.Split(l, "=")
			h, err := url.PathUnescape(header[1])
			if err != nil {
				return record{}, err
			}
			r.Headers[header[0][1:]] = h
		default:
			tuple := strings.Split(l, "=")
			if len(tuple) != 2 {
				return record{}, fmt.Errorf("arbitrary data not allowed")
			}
			continue
		}
		if len(l) > 255 {
			return record{}, fmt.Errorf("TXT record cannot exceed the maximum of 255 characters")
		}
		if r.Type == "dockerv2" && r.To == "" {
			return record{}, fmt.Errorf("[txtdirect]: to= field is required in dockerv2 type")
		}
	}

	if r.Code == 0 {
		r.Code = http.StatusFound
	}

	// Only apply rules and default to records that doesn't point to a upstream record
	if len(r.Use) == 0 {
		if r.Type == "" {
			r.Type = "host"
		}

		if r.Type == "host" && r.To == "" {
			fallback(w, req, "global", http.StatusMovedPermanently, c)
			return record{}, nil
		}

		if !contains(c.Enable, r.Type) {
			return record{}, fmt.Errorf("%s type is not enabled in configuration", r.Type)
		}
	}

	return r, nil
}

// Adds the given record to the request's context with "records" key.
func (rec record) addToContext(r *http.Request) *http.Request {
	// Fetch fallback config from context and add the record to it
	recordsContext := r.Context().Value("records")

	// Create a new records field in the context if it doesn't exist
	if recordsContext == nil {
		return r.WithContext(context.WithValue(r.Context(), "records", []record{rec}))
	}

	records := append(recordsContext.([]record), rec)

	// Replace the fallback config instance inside the request's context
	return r.WithContext(context.WithValue(r.Context(), "records", records))
}

// ParseURI parses the given URI and triggers fallback if the URI isn't valid
func ParseURI(uri string, w http.ResponseWriter, r *http.Request, c Config) string {
	url, err := url.Parse(uri)
	if err != nil {
		fallback(w, r, "global", http.StatusMovedPermanently, c)
		return ""
	}
	return url.String()
}

// UpstreamRecord will check all of the use= fields and sends a request to each
// upstream zone address and choses the first one that returns the final TXT
// record
func (rec *record) UpstreamRecord(c Config, w http.ResponseWriter, r *http.Request) (record, string, error) {
	var upstreamRec record
	var err error

	for _, zone := range rec.Use {
		upstreamRec, err = getRecord(zone, c, w, r)
		if err != nil {
			continue
		}
		return upstreamRec, zone, nil
	}

	return record{}, "", fmt.Errorf("Couldn't find any records from upstream")
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"fmt"
	"os"
	"os/exec"
	"os/signal"
	"regexp"
	"strings"
	"sync"
	"syscall"
	"time"

	flag "github.com/cespare/pflag"
	"github.com/howeyc/fsnotify"
)

const defaultSubSymbol = "{}"

var (
	reflexes []*Reflex
	matchAll = regexp.MustCompile(".*")

	flagConf       string
	flagSequential bool
	flagDecoration string
	decoration     Decoration
	verbose        bool
	globalFlags    = flag.NewFlagSet("", flag.ContinueOnError)
	globalConfig   = &Config{}

	reflexID = 0
	stdout   = make(chan OutMsg, 1)

	cleanupMu = &sync.Mutex{}
)

func usage() {
	fmt.Fprintf(os.Stderr, `Usage: %s [OPTIONS] [COMMAND]

COMMAND is any command you'd like to run. Any instance of {} will be replaced
with the filename of the changed file. (The symbol may be changed with the
--substitute flag.)

OPTIONS are given below:
`, os.Args[0])

	globalFlags.PrintDefaults()

	fmt.Fprintln(os.Stderr, `
Examples:

    # Print each .txt file if it changes
    $ reflex -r '\.txt$' echo {}

    # Run 'make' if any of the .c files in this directory change:
    $ reflex -g '*.c' make

    # Build and run a server; rebuild and restart when .java files change:
    $ reflex -r '\.java$' -s -- sh -c 'make && java bin/Server'
`)
}

func init() {
	globalFlags.Usage = usage
	globalFlags.StringVarP(&flagConf, "config", "c", "", `
            A configuration file that describes how to run reflex
            (or '-' to read the configuration from stdin).`)
	globalFlags.BoolVarP(&verbose, "verbose", "v", false, `
            Verbose mode: print out more information about what reflex is doing.`)
	globalFlags.BoolVarP(&flagSequential, "sequential", "e", false, `
            Don't run multiple commands at the same time.`)
	globalFlags.StringVarP(&flagDecoration, "decoration", "d", "plain", `
            How to decorate command output. Choices: none, plain, fancy.`)
	globalConfig.registerFlags(globalFlags)
}

func anyNonGlobalsRegistered() bool {
	any := false
	walkFn := func(f *flag.Flag) {
		if !(f.Name == "config" || f.Name == "verbose" || f.Name == "sequential" || f.Name == "decoration") {
			any = any || true
		}
	}
	globalFlags.Visit(walkFn)
	return any
}

func parseMatchers(rs, gs string) (regex *regexp.Regexp, glob string, err error) {
	if rs == "" && gs == "" {
		return matchAll, "", nil
	}
	if rs == "" {
		return nil, gs, nil
	}
	if gs == "" {
		regex, err := regexp.Compile(rs)
		if err != nil {
			return nil, "", err
		}
		return regex, "", nil
	}
	return nil, "", errors.New("Both regex and glob specified.")
}

// This ties together a single reflex 'instance' so that multiple watches/commands can be handled together
// easily.
type Reflex struct {
	id           int
	source       string // Describes what config/line defines this Reflex
	startService bool
	backlog      Backlog
	regex        *regexp.Regexp
	glob         string
	useRegex     bool
	onlyFiles    bool
	onlyDirs     bool
	command      []string
	subSymbol    string

	done       chan struct{}
	rawChanges chan string
	filtered   chan string
	batched    chan string

	// Used for services (startService = true)
	cmd    *exec.Cmd
	tty    *os.File
	mu     *sync.Mutex // protects killed
	killed bool
}

// This function is not threadsafe.
func NewReflex(c *Config) (*Reflex, error) {
	regex, glob, err := parseMatchers(c.regex, c.glob)
	if err != nil {
		Fatalln("Error parsing glob/regex.\n" + err.Error())
	}
	if len(c.command) == 0 {
		return nil, errors.New("Must give command to execute.")
	}

	if c.subSymbol == "" {
		return nil, errors.New("Substitution symbol must be non-empty.")
	}

	substitution := false
	for _, part := range c.command {
		if strings.Contains(part, c.subSymbol) {
			substitution = true
			break
		}
	}

	var backlog Backlog
	if substitution {
		if c.startService {
			return nil, errors.New("Using --start-service does not work with a command that has a substitution symbol.")
		}
		backlog = NewUniqueFilesBacklog()
	} else {
		backlog = NewUnifiedBacklog()
	}

	if c.onlyFiles && c.onlyDirs {
		return nil, errors.New("Cannot specify both --only-files and --only-dirs.")
	}

	reflex := &Reflex{
		id:           reflexID,
		source:       c.source,
		startService: c.startService,
		backlog:      backlog,
		regex:        regex,
		glob:         glob,
		useRegex:     regex != nil,
		onlyFiles:    c.onlyFiles,
		onlyDirs:     c.onlyDirs,
		command:      c.command,
		subSymbol:    c.subSymbol,

		rawChanges: make(chan string),
		filtered:   make(chan string),
		batched:    make(chan string),

		mu: &sync.Mutex{},
	}
	reflexID++

	return reflex, nil
}

func (r *Reflex) PrintInfo() {
	fmt.Println("Reflex from", r.source)
	fmt.Println("| ID:", r.id)
	if r.regex == matchAll {
		fmt.Println("| No regex (-r) or glob (-g) given, so matching all file changes.")
	} else if r.useRegex {
		fmt.Println("| Regex:", r.regex)
	} else {
		fmt.Println("| Glob:", r.glob)
	}
	if r.onlyFiles {
		fmt.Println("| Only matching files.")
	} else if r.onlyDirs {
		fmt.Println("| Only matching directories.")
	}
	if !r.startService {
		fmt.Println("| Substitution symbol", r.subSymbol)
	}
	replacer := strings.NewReplacer(r.subSymbol, "<filename>")
	command := make([]string, len(r.command))
	for i, part := range r.command {
		command[i] = replacer.Replace(part)
	}
	fmt.Println("| Command:", command)
	fmt.Println("+---------")
}

func printGlobals() {
	fmt.Println("Globals set at commandline")
	walkFn := func(f *flag.Flag) {
		fmt.Printf("| --%s (-%s) '%s' (default: '%s')\n", f.Name, f.Shorthand, f.Value, f.DefValue)
	}
	globalFlags.Visit(walkFn)
	fmt.Println("+---------")
}

func cleanup(reason string) {
	cleanupMu.Lock()
	defer cleanupMu.Unlock()
	fmt.Println(reason)
	wg := &sync.WaitGroup{}
	for _, reflex := range reflexes {
		if reflex.done != nil {
			wg.Add(1)
			go func(reflex *Reflex) {
				terminate(reflex)
				wg.Done()
			}(reflex)
		}
	}
	wg.Wait()
	// Give just a little time to finish printing output.
	time.Sleep(10 * time.Millisecond)
	os.Exit(0)
}

func main() {
	if err := globalFlags.Parse(os.Args[1:]); err != nil {
		Fatalln(err)
	}
	globalConfig.command = globalFlags.Args()
	globalConfig.source = "[commandline]"
	if verbose {
		printGlobals()
	}
	switch strings.ToLower(flagDecoration) {
	case "none":
		decoration = DecorationNone
	case "plain":
		decoration = DecorationPlain
	case "fancy":
		decoration = DecorationFancy
	default:
		Fatalln(fmt.Sprintf("Invalid decoration %s. Choices: none, plain, fancy.", flagDecoration))
	}

	var configs []*Config
	if flagConf == "" {
		if flagSequential {
			Fatalln("Cannot set --sequential without --config (because you cannot specify multiple commands).")
		}
		configs = []*Config{globalConfig}
	} else {
		if anyNonGlobalsRegistered() {
			Fatalln("Cannot set other flags along with --config other than --sequential, --verbose, and --decoration.")
		}
		var err error
		configs, err = ReadConfigs(flagConf)
		if err != nil {
			Fatalln("Could not parse configs: ", err)
		}
	}

	for _, config := range configs {
		reflex, err := NewReflex(config)
		if err != nil {
			Fatalln("Could not make reflex for config:", err)
		}
		if verbose {
			reflex.PrintInfo()
		}
		reflexes = append(reflexes, reflex)
	}

	// Catch ctrl-c and make sure to kill off children.
	signals := make(chan os.Signal, 1)
	signal.Notify(signals, os.Interrupt)
	signal.Notify(signals, os.Signal(syscall.SIGTERM))
	go func() {
		s := <-signals
		reason := fmt.Sprintf("Interrupted (%s). Cleaning up children...", s)
		cleanup(reason)
	}()
	defer cleanup("Cleaning up.")

	watcher, err := fsnotify.NewWatcher()
	if err != nil {
		Fatalln(err)
	}
	defer watcher.Close()

	rawChanges := make(chan string)
	allRawChanges := make([]chan<- string, len(reflexes))
	done := make(chan error)
	for i, reflex := range reflexes {
		allRawChanges[i] = reflex.rawChanges
	}
	go watch(".", watcher, rawChanges, done)
	go broadcast(rawChanges, allRawChanges)

	go printOutput(stdout, os.Stdout)

	for _, reflex := range reflexes {
		go filterMatching(reflex.rawChanges, reflex.filtered, reflex)
		go batch(reflex.filtered, reflex.batched, reflex)
		go runEach(reflex.batched, reflex)
		if reflex.startService {
			// Easy hack to kick off the initial start.
			infoPrintln(reflex.id, "Starting service")
			runCommand(reflex, "", stdout)
		}
	}

	Fatalln(<-done)
}
<commit_msg>Don't bother unlocking mutex in exit func<commit_after>package main

import (
	"errors"
	"fmt"
	"os"
	"os/exec"
	"os/signal"
	"regexp"
	"strings"
	"sync"
	"syscall"
	"time"

	flag "github.com/cespare/pflag"
	"github.com/howeyc/fsnotify"
)

const defaultSubSymbol = "{}"

var (
	reflexes []*Reflex
	matchAll = regexp.MustCompile(".*")

	flagConf       string
	flagSequential bool
	flagDecoration string
	decoration     Decoration
	verbose        bool
	globalFlags    = flag.NewFlagSet("", flag.ContinueOnError)
	globalConfig   = &Config{}

	reflexID = 0
	stdout   = make(chan OutMsg, 1)

	cleanupMu = &sync.Mutex{}
)

func usage() {
	fmt.Fprintf(os.Stderr, `Usage: %s [OPTIONS] [COMMAND]

COMMAND is any command you'd like to run. Any instance of {} will be replaced
with the filename of the changed file. (The symbol may be changed with the
--substitute flag.)

OPTIONS are given below:
`, os.Args[0])

	globalFlags.PrintDefaults()

	fmt.Fprintln(os.Stderr, `
Examples:

    # Print each .txt file if it changes
    $ reflex -r '\.txt$' echo {}

    # Run 'make' if any of the .c files in this directory change:
    $ reflex -g '*.c' make

    # Build and run a server; rebuild and restart when .java files change:
    $ reflex -r '\.java$' -s -- sh -c 'make && java bin/Server'
`)
}

func init() {
	globalFlags.Usage = usage
	globalFlags.StringVarP(&flagConf, "config", "c", "", `
            A configuration file that describes how to run reflex
            (or '-' to read the configuration from stdin).`)
	globalFlags.BoolVarP(&verbose, "verbose", "v", false, `
            Verbose mode: print out more information about what reflex is doing.`)
	globalFlags.BoolVarP(&flagSequential, "sequential", "e", false, `
            Don't run multiple commands at the same time.`)
	globalFlags.StringVarP(&flagDecoration, "decoration", "d", "plain", `
            How to decorate command output. Choices: none, plain, fancy.`)
	globalConfig.registerFlags(globalFlags)
}

func anyNonGlobalsRegistered() bool {
	any := false
	walkFn := func(f *flag.Flag) {
		if !(f.Name == "config" || f.Name == "verbose" || f.Name == "sequential" || f.Name == "decoration") {
			any = any || true
		}
	}
	globalFlags.Visit(walkFn)
	return any
}

func parseMatchers(rs, gs string) (regex *regexp.Regexp, glob string, err error) {
	if rs == "" && gs == "" {
		return matchAll, "", nil
	}
	if rs == "" {
		return nil, gs, nil
	}
	if gs == "" {
		regex, err := regexp.Compile(rs)
		if err != nil {
			return nil, "", err
		}
		return regex, "", nil
	}
	return nil, "", errors.New("Both regex and glob specified.")
}

// This ties together a single reflex 'instance' so that multiple watches/commands can be handled together
// easily.
type Reflex struct {
	id           int
	source       string // Describes what config/line defines this Reflex
	startService bool
	backlog      Backlog
	regex        *regexp.Regexp
	glob         string
	useRegex     bool
	onlyFiles    bool
	onlyDirs     bool
	command      []string
	subSymbol    string

	done       chan struct{}
	rawChanges chan string
	filtered   chan string
	batched    chan string

	// Used for services (startService = true)
	cmd    *exec.Cmd
	tty    *os.File
	mu     *sync.Mutex // protects killed
	killed bool
}

// This function is not threadsafe.
func NewReflex(c *Config) (*Reflex, error) {
	regex, glob, err := parseMatchers(c.regex, c.glob)
	if err != nil {
		Fatalln("Error parsing glob/regex.\n" + err.Error())
	}
	if len(c.command) == 0 {
		return nil, errors.New("Must give command to execute.")
	}

	if c.subSymbol == "" {
		return nil, errors.New("Substitution symbol must be non-empty.")
	}

	substitution := false
	for _, part := range c.command {
		if strings.Contains(part, c.subSymbol) {
			substitution = true
			break
		}
	}

	var backlog Backlog
	if substitution {
		if c.startService {
			return nil, errors.New("Using --start-service does not work with a command that has a substitution symbol.")
		}
		backlog = NewUniqueFilesBacklog()
	} else {
		backlog = NewUnifiedBacklog()
	}

	if c.onlyFiles && c.onlyDirs {
		return nil, errors.New("Cannot specify both --only-files and --only-dirs.")
	}

	reflex := &Reflex{
		id:           reflexID,
		source:       c.source,
		startService: c.startService,
		backlog:      backlog,
		regex:        regex,
		glob:         glob,
		useRegex:     regex != nil,
		onlyFiles:    c.onlyFiles,
		onlyDirs:     c.onlyDirs,
		command:      c.command,
		subSymbol:    c.subSymbol,

		rawChanges: make(chan string),
		filtered:   make(chan string),
		batched:    make(chan string),

		mu: &sync.Mutex{},
	}
	reflexID++

	return reflex, nil
}

func (r *Reflex) PrintInfo() {
	fmt.Println("Reflex from", r.source)
	fmt.Println("| ID:", r.id)
	if r.regex == matchAll {
		fmt.Println("| No regex (-r) or glob (-g) given, so matching all file changes.")
	} else if r.useRegex {
		fmt.Println("| Regex:", r.regex)
	} else {
		fmt.Println("| Glob:", r.glob)
	}
	if r.onlyFiles {
		fmt.Println("| Only matching files.")
	} else if r.onlyDirs {
		fmt.Println("| Only matching directories.")
	}
	if !r.startService {
		fmt.Println("| Substitution symbol", r.subSymbol)
	}
	replacer := strings.NewReplacer(r.subSymbol, "<filename>")
	command := make([]string, len(r.command))
	for i, part := range r.command {
		command[i] = replacer.Replace(part)
	}
	fmt.Println("| Command:", command)
	fmt.Println("+---------")
}

func printGlobals() {
	fmt.Println("Globals set at commandline")
	walkFn := func(f *flag.Flag) {
		fmt.Printf("| --%s (-%s) '%s' (default: '%s')\n", f.Name, f.Shorthand, f.Value, f.DefValue)
	}
	globalFlags.Visit(walkFn)
	fmt.Println("+---------")
}

func cleanup(reason string) {
	cleanupMu.Lock()
	fmt.Println(reason)
	wg := &sync.WaitGroup{}
	for _, reflex := range reflexes {
		if reflex.done != nil {
			wg.Add(1)
			go func(reflex *Reflex) {
				terminate(reflex)
				wg.Done()
			}(reflex)
		}
	}
	wg.Wait()
	// Give just a little time to finish printing output.
	time.Sleep(10 * time.Millisecond)
	os.Exit(0)
}

func main() {
	if err := globalFlags.Parse(os.Args[1:]); err != nil {
		Fatalln(err)
	}
	globalConfig.command = globalFlags.Args()
	globalConfig.source = "[commandline]"
	if verbose {
		printGlobals()
	}
	switch strings.ToLower(flagDecoration) {
	case "none":
		decoration = DecorationNone
	case "plain":
		decoration = DecorationPlain
	case "fancy":
		decoration = DecorationFancy
	default:
		Fatalln(fmt.Sprintf("Invalid decoration %s. Choices: none, plain, fancy.", flagDecoration))
	}

	var configs []*Config
	if flagConf == "" {
		if flagSequential {
			Fatalln("Cannot set --sequential without --config (because you cannot specify multiple commands).")
		}
		configs = []*Config{globalConfig}
	} else {
		if anyNonGlobalsRegistered() {
			Fatalln("Cannot set other flags along with --config other than --sequential, --verbose, and --decoration.")
		}
		var err error
		configs, err = ReadConfigs(flagConf)
		if err != nil {
			Fatalln("Could not parse configs: ", err)
		}
	}

	for _, config := range configs {
		reflex, err := NewReflex(config)
		if err != nil {
			Fatalln("Could not make reflex for config:", err)
		}
		if verbose {
			reflex.PrintInfo()
		}
		reflexes = append(reflexes, reflex)
	}

	// Catch ctrl-c and make sure to kill off children.
	signals := make(chan os.Signal, 1)
	signal.Notify(signals, os.Interrupt)
	signal.Notify(signals, os.Signal(syscall.SIGTERM))
	go func() {
		s := <-signals
		reason := fmt.Sprintf("Interrupted (%s). Cleaning up children...", s)
		cleanup(reason)
	}()
	defer cleanup("Cleaning up.")

	watcher, err := fsnotify.NewWatcher()
	if err != nil {
		Fatalln(err)
	}
	defer watcher.Close()

	rawChanges := make(chan string)
	allRawChanges := make([]chan<- string, len(reflexes))
	done := make(chan error)
	for i, reflex := range reflexes {
		allRawChanges[i] = reflex.rawChanges
	}
	go watch(".", watcher, rawChanges, done)
	go broadcast(rawChanges, allRawChanges)

	go printOutput(stdout, os.Stdout)

	for _, reflex := range reflexes {
		go filterMatching(reflex.rawChanges, reflex.filtered, reflex)
		go batch(reflex.filtered, reflex.batched, reflex)
		go runEach(reflex.batched, reflex)
		if reflex.startService {
			// Easy hack to kick off the initial start.
			infoPrintln(reflex.id, "Starting service")
			runCommand(reflex, "", stdout)
		}
	}

	Fatalln(<-done)
}
<|endoftext|>


<commit_before>package agent

import (
	"fmt"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

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

// extra endpoints that should be tested, and their allowed methods
var extraTestEndpoints = map[string][]string{
	"/v1/query":             []string{"GET", "POST"},
	"/v1/query/":            []string{"GET", "PUT", "DELETE"},
	"/v1/query/xxx/execute": []string{"GET"},
	"/v1/query/xxx/explain": []string{"GET"},
}

// certain endpoints can't be unit tested.
func includePathInTest(path string) bool {
	var hanging = path == "/v1/status/peers" || path == "/v1/agent/monitor" || path == "/v1/agent/reload" // these hang
	var custom = path == "/v1/query" || path == "/v1/query/"                                              // these have custom logic
	return !(hanging || custom)
}

func TestHTTPAPI_MethodNotAllowed_OSS(t *testing.T) {

	a := NewTestAgent(t.Name(), `acl_datacenter = "dc1"`)
	a.Agent.LogWriter = logger.NewLogWriter(512)
	defer a.Shutdown()

	all := []string{"GET", "PUT", "POST", "DELETE", "HEAD", "OPTIONS"}
	client := http.Client{}

	testMethodNotAllowed := func(method string, path string, allowedMethods []string) {
		t.Run(method+" "+path, func(t *testing.T) {
			uri := fmt.Sprintf("http://%s%s", a.HTTPAddr(), path)
			req, _ := http.NewRequest(method, uri, nil)
			resp, err := client.Do(req)
			if err != nil {
				t.Fatal("client.Do failed: ", err)
			}

			allowed := method == "OPTIONS"
			for _, allowedMethod := range allowedMethods {
				if allowedMethod == method {
					allowed = true
					break
				}
			}

			if allowed && resp.StatusCode == http.StatusMethodNotAllowed {
				t.Fatalf("method allowed: got status code %d want any other code", resp.StatusCode)
			}
			if !allowed && resp.StatusCode != http.StatusMethodNotAllowed {
				t.Fatalf("method not allowed: got status code %d want %d", resp.StatusCode, http.StatusMethodNotAllowed)
			}
		})
	}

	for path, methods := range extraTestEndpoints {
		for _, method := range all {
			testMethodNotAllowed(method, path, methods)
		}
	}

	for path, methods := range allowedMethods {
		if includePathInTest(path) {
			for _, method := range all {
				testMethodNotAllowed(method, path, methods)
			}
		}
	}
}

func TestHTTPAPI_OptionMethod_OSS(t *testing.T) {
	a := NewTestAgent(t.Name(), `acl_datacenter = "dc1"`)
	a.Agent.LogWriter = logger.NewLogWriter(512)
	defer a.Shutdown()

	testOptionMethod := func(path string, methods []string) {
		t.Run("OPTIONS "+path, func(t *testing.T) {
			uri := fmt.Sprintf("http://%s%s", a.HTTPAddr(), path)
			req, _ := http.NewRequest("OPTIONS", uri, nil)
			resp := httptest.NewRecorder()
			a.srv.Handler.ServeHTTP(resp, req)
			allMethods := append([]string{"OPTIONS"}, methods...)

			if resp.Code != http.StatusOK {
				t.Fatalf("options request: got status code %d want %d", resp.Code, http.StatusOK)
			}

			optionsStr := resp.Header().Get("Allow")
			if optionsStr == "" {
				t.Fatalf("options request: got empty 'Allow' header")
			} else if optionsStr != strings.Join(allMethods, ",") {
				t.Fatalf("options request: got 'Allow' header value of %s want %s", optionsStr, allMethods)
			}
		})
	}

	for path, methods := range extraTestEndpoints {
		testOptionMethod(path, methods)
	}
	for path, methods := range allowedMethods {
		if includePathInTest(path) {
			testOptionMethod(path, methods)
		}
	}
}
<commit_msg>cleanup unit test code a bit<commit_after>package agent

import (
	"fmt"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

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

// extra endpoints that should be tested, and their allowed methods
var extraTestEndpoints = map[string][]string{
	"/v1/query":             []string{"GET", "POST"},
	"/v1/query/":            []string{"GET", "PUT", "DELETE"},
	"/v1/query/xxx/execute": []string{"GET"},
	"/v1/query/xxx/explain": []string{"GET"},
}

// These endpoints are ignored in unit testing for response codes
var ignoredEndpoints = []string{"/v1/status/peers","/v1/agent/monitor",  "/v1/agent/reload" }
// These have custom logic
var customEndpoints = []string{"/v1/query", "/v1/query/"}

// includePathInTest returns whether this path should be ignored for the purpose of testing its response code
func includePathInTest(path string) bool {
	ignored := false
	for _, p := range ignoredEndpoints {
		if p == path {
			ignored = true
			break
		}
	}
	for _, p := range customEndpoints {
		if p == path {
			ignored = true
			break
		}
	}

	return !ignored
}

func TestHTTPAPI_MethodNotAllowed_OSS(t *testing.T) {

	a := NewTestAgent(t.Name(), `acl_datacenter = "dc1"`)
	a.Agent.LogWriter = logger.NewLogWriter(512)
	defer a.Shutdown()

	all := []string{"GET", "PUT", "POST", "DELETE", "HEAD", "OPTIONS"}
	client := http.Client{}

	testMethodNotAllowed := func(method string, path string, allowedMethods []string) {
		t.Run(method+" "+path, func(t *testing.T) {
			uri := fmt.Sprintf("http://%s%s", a.HTTPAddr(), path)
			req, _ := http.NewRequest(method, uri, nil)
			resp, err := client.Do(req)
			if err != nil {
				t.Fatal("client.Do failed: ", err)
			}

			allowed := method == "OPTIONS"
			for _, allowedMethod := range allowedMethods {
				if allowedMethod == method {
					allowed = true
					break
				}
			}

			if allowed && resp.StatusCode == http.StatusMethodNotAllowed {
				t.Fatalf("method allowed: got status code %d want any other code", resp.StatusCode)
			}
			if !allowed && resp.StatusCode != http.StatusMethodNotAllowed {
				t.Fatalf("method not allowed: got status code %d want %d", resp.StatusCode, http.StatusMethodNotAllowed)
			}
		})
	}

	for path, methods := range extraTestEndpoints {
		for _, method := range all {
			testMethodNotAllowed(method, path, methods)
		}
	}

	for path, methods := range allowedMethods {
		if includePathInTest(path) {
			for _, method := range all {
				testMethodNotAllowed(method, path, methods)
			}
		}
	}
}

func TestHTTPAPI_OptionMethod_OSS(t *testing.T) {
	a := NewTestAgent(t.Name(), `acl_datacenter = "dc1"`)
	a.Agent.LogWriter = logger.NewLogWriter(512)
	defer a.Shutdown()

	testOptionMethod := func(path string, methods []string) {
		t.Run("OPTIONS "+path, func(t *testing.T) {
			uri := fmt.Sprintf("http://%s%s", a.HTTPAddr(), path)
			req, _ := http.NewRequest("OPTIONS", uri, nil)
			resp := httptest.NewRecorder()
			a.srv.Handler.ServeHTTP(resp, req)
			allMethods := append([]string{"OPTIONS"}, methods...)

			if resp.Code != http.StatusOK {
				t.Fatalf("options request: got status code %d want %d", resp.Code, http.StatusOK)
			}

			optionsStr := resp.Header().Get("Allow")
			if optionsStr == "" {
				t.Fatalf("options request: got empty 'Allow' header")
			} else if optionsStr != strings.Join(allMethods, ",") {
				t.Fatalf("options request: got 'Allow' header value of %s want %s", optionsStr, allMethods)
			}
		})
	}

	for path, methods := range extraTestEndpoints {
		testOptionMethod(path, methods)
	}
	for path, methods := range allowedMethods {
		if includePathInTest(path) {
			testOptionMethod(path, methods)
		}
	}
}
<|endoftext|>