<commit_before>package api

import (
	"fmt"
	"github.com/Aptomi/aptomi/pkg/engine/apply/action"
	"github.com/Aptomi/aptomi/pkg/engine/apply/action/component"
	"github.com/Aptomi/aptomi/pkg/engine/diff"
	"github.com/Aptomi/aptomi/pkg/engine/resolve"
	"github.com/Aptomi/aptomi/pkg/event"
	"github.com/Aptomi/aptomi/pkg/lang"
	"github.com/Aptomi/aptomi/pkg/plugin"
	"github.com/Aptomi/aptomi/pkg/runtime"
	"github.com/Sirupsen/logrus"
	"github.com/julienschmidt/httprouter"
	"net/http"
	"strings"
)

// DependencyQueryFlag determines whether to query just dependency deployment status, or both deployment + readiness/health checks status
type DependencyQueryFlag string

const (
	// DependencyQueryDeploymentStatusOnly prescribes only to query dependency deployment status (i.e. actual state = desired state)
	DependencyQueryDeploymentStatusOnly DependencyQueryFlag = "deployed"

	// DependencyQueryDeploymentStatusAndReadiness prescribes to query both dependency deployment status (i.e. actual state = desired state), as well as readiness status (i.e. health checks = passing)
	DependencyQueryDeploymentStatusAndReadiness DependencyQueryFlag = "ready"
)

// DependenciesStatusObject is an informational data structure with Kind and Constructor for DependenciesStatus
var DependenciesStatusObject = &runtime.Info{
	Kind:        "dependencies-status",
	Constructor: func() runtime.Object { return &DependenciesStatus{} },
}

// DependenciesStatus is a struct which holds status information for a set of given dependencies
type DependenciesStatus struct {
	runtime.TypeKind `yaml:",inline"`

	// map containing status by dependency
	Status map[string]*DependencyStatus
}

// DependencyStatus is a struct which holds status information for an individual dependency
type DependencyStatus struct {
	Found     bool
	Deployed  bool
	Ready     bool
	Endpoints map[string]map[string]string
}

func (api *coreAPI) handleDependencyStatusGet(writer http.ResponseWriter, request *http.Request, params httprouter.Params) {
	// parse query mode flag (deployment status vs. readiness status) as well as the list of dependency IDs
	flag := DependencyQueryFlag(params.ByName("queryFlag"))
	dependencyIds := strings.Split(params.ByName("idList"), ",")

	// load the latest policy
	policy, _, errPolicy := api.store.GetPolicy(runtime.LastGen)
	if errPolicy != nil {
		panic(fmt.Sprintf("error while loading latest policy from the store: %s", errPolicy))
	}

	// initialize result
	result := &DependenciesStatus{
		TypeKind: DependenciesStatusObject.GetTypeKind(),
		Status:   make(map[string]*DependencyStatus),
	}
	for _, depID := range dependencyIds {
		parts := strings.Split(depID, "^")
		dObj, err := policy.GetObject(lang.DependencyObject.Kind, parts[1], parts[0])
		if dObj == nil || err != nil {
			dKey := runtime.KeyFromParts(parts[0], lang.DependencyObject.Kind, parts[1])
			result.Status[dKey] = &DependencyStatus{
				Found:     false,
				Deployed:  false,
				Ready:     false,
				Endpoints: make(map[string]map[string]string),
			}
			continue
		}

		d := dObj.(*lang.Dependency)
		result.Status[runtime.KeyForStorable(d)] = &DependencyStatus{
			Found:     true,
			Deployed:  true,
			Ready:     false,
			Endpoints: make(map[string]map[string]string),
		}
	}

	// load actual and desired states
	desiredState := resolve.NewPolicyResolver(policy, api.externalData, event.NewLog(logrus.WarnLevel, "api-dependencies-status")).ResolveAllDependencies()
	actualState, err := api.store.GetActualState()
	if err != nil {
		panic(fmt.Sprintf("can't load actual state from the store: %s", err))
	}

	// fetch deployment status for dependencies
	fetchDeploymentStatusForDependencies(result, actualState, desiredState)

	// fetch readiness status for dependencies, if we were asked to do so
	if flag == DependencyQueryDeploymentStatusAndReadiness {
		plugins := api.pluginRegistryFactory()
		fetchReadinessStatusForDependencies(result, plugins, policy, actualState, desiredState)
	}

	// fetch endpoints for dependencies
	fetchEndpointsForDependencies(result, actualState)

	// return the result back
	api.contentType.WriteOne(writer, request, result)
}

func fetchDeploymentStatusForDependencies(result *DependenciesStatus, actualState *resolve.PolicyResolution, desiredState *resolve.PolicyResolution) {
	// compare desired vs. actual state and see what's the dependency status for every provided dependency ID
	diff.NewPolicyResolutionDiff(desiredState, actualState).ActionPlan.Apply(
		action.WrapSequential(func(act action.Base) error {
			// if it's attach action is pending on component, let's see which particular dependency it affects
			if dAction, ok := act.(*component.AttachDependencyAction); ok {
				// reset status of this particular dependency to false
				if _, affected := result.Status[dAction.DependencyID]; affected {
					result.Status[dAction.DependencyID].Deployed = false
					return nil
				}
			}

			// if it's detach action is pending on component, let's see which particular dependency it affects
			if dAction, ok := act.(*component.DetachDependencyAction); ok {
				// reset status of this particular dependency to false
				if _, affected := result.Status[dAction.DependencyID]; affected {
					result.Status[dAction.DependencyID].Deployed = false
					return nil
				}
			}

			key, ok := act.DescribeChanges()["key"].(string)
			if ok && len(key) > 0 {
				// we found a component in diff, which is affected by the action. let's see if any of the dependencies are affected
				affectedDepKeys := make(map[string]bool)
				{
					prevInstance := actualState.ComponentInstanceMap[key]
					if prevInstance != nil {
						for dKey := range prevInstance.DependencyKeys {
							affectedDepKeys[dKey] = true
						}
					}
				}
				{
					nextInstance := desiredState.ComponentInstanceMap[key]
					if nextInstance != nil {
						for dKey := range nextInstance.DependencyKeys {
							affectedDepKeys[dKey] = true
						}
					}
				}

				// if our dependency is affected, reset its deployed status to false (because actions are pending)
				for dKey := range affectedDepKeys {
					if _, ok := result.Status[dKey]; ok {
						result.Status[dKey].Deployed = false
					}
				}
			}
			return nil
		}),
		action.NewApplyResultUpdaterImpl(),
	)

}

func fetchReadinessStatusForDependencies(result *DependenciesStatus, plugins plugin.Registry, policy *lang.Policy, actualState *resolve.PolicyResolution, desiredState *resolve.PolicyResolution) {
	for _, instance := range actualState.ComponentInstanceMap {
		for dKey := range instance.DependencyKeys {
			if _, ok := result.Status[dKey]; ok {
				codePlugin, err := pluginForComponentInstance(instance, policy, plugins)
				if err != nil {
					panic(fmt.Sprintf("Can't get plugin for component instance %s: %s", instance.GetKey(), err))
				}
				if codePlugin == nil {
					continue
				}

				instanceStatus, err := codePlugin.Status(
					&plugin.CodePluginInvocationParams{
						DeployName:   instance.GetDeployName(),
						Params:       instance.CalculatedCodeParams,
						PluginParams: map[string]string{plugin.ParamTargetSuffix: instance.Metadata.Key.TargetSuffix},
						EventLog:     event.NewLog(logrus.WarnLevel, "resources-status"),
					},
				)
				if err != nil {
					panic(fmt.Sprintf("Error while getting deployment resources status for component instance %s: %s", instance.GetKey(), err))
				}

				result.Status[dKey].Ready = result.Status[dKey].Deployed && instanceStatus
			}
		}
	}
}

func fetchEndpointsForDependencies(result *DependenciesStatus, actualState *resolve.PolicyResolution) {
	for _, instance := range actualState.ComponentInstanceMap {
		for dKey := range instance.DependencyKeys {
			if _, ok := result.Status[dKey]; ok {
				if len(instance.Endpoints) > 0 {
					result.Status[dKey].Endpoints[instance.GetName()] = instance.Endpoints
				}
			}
		}
	}
}
<commit_msg>fixed dependency status retrieval (in parallel)<commit_after>package api

import (
	"fmt"
	"github.com/Aptomi/aptomi/pkg/engine/apply/action"
	"github.com/Aptomi/aptomi/pkg/engine/apply/action/component"
	"github.com/Aptomi/aptomi/pkg/engine/diff"
	"github.com/Aptomi/aptomi/pkg/engine/resolve"
	"github.com/Aptomi/aptomi/pkg/event"
	"github.com/Aptomi/aptomi/pkg/lang"
	"github.com/Aptomi/aptomi/pkg/plugin"
	"github.com/Aptomi/aptomi/pkg/runtime"
	"github.com/Sirupsen/logrus"
	"github.com/julienschmidt/httprouter"
	"net/http"
	"runtime/debug"
	"strings"
	"sync"
)

// DependencyQueryFlag determines whether to query just dependency deployment status, or both deployment + readiness/health checks status
type DependencyQueryFlag string

const (
	// DependencyQueryDeploymentStatusOnly prescribes only to query dependency deployment status (i.e. actual state = desired state)
	DependencyQueryDeploymentStatusOnly DependencyQueryFlag = "deployed"

	// DependencyQueryDeploymentStatusAndReadiness prescribes to query both dependency deployment status (i.e. actual state = desired state), as well as readiness status (i.e. health checks = passing)
	DependencyQueryDeploymentStatusAndReadiness DependencyQueryFlag = "ready"
)

// DependenciesStatusObject is an informational data structure with Kind and Constructor for DependenciesStatus
var DependenciesStatusObject = &runtime.Info{
	Kind:        "dependencies-status",
	Constructor: func() runtime.Object { return &DependenciesStatus{} },
}

// DependenciesStatus is a struct which holds status information for a set of given dependencies
type DependenciesStatus struct {
	runtime.TypeKind `yaml:",inline"`

	// map containing status by dependency
	Status map[string]*DependencyStatus
}

// DependencyStatus is a struct which holds status information for an individual dependency
type DependencyStatus struct {
	Found     bool
	Deployed  bool
	Ready     bool
	Endpoints map[string]map[string]string
}

func (api *coreAPI) handleDependencyStatusGet(writer http.ResponseWriter, request *http.Request, params httprouter.Params) {
	// parse query mode flag (deployment status vs. readiness status) as well as the list of dependency IDs
	flag := DependencyQueryFlag(params.ByName("queryFlag"))
	dependencyIds := strings.Split(params.ByName("idList"), ",")

	// load the latest policy
	policy, _, errPolicy := api.store.GetPolicy(runtime.LastGen)
	if errPolicy != nil {
		panic(fmt.Sprintf("error while loading latest policy from the store: %s", errPolicy))
	}

	// initialize result
	result := &DependenciesStatus{
		TypeKind: DependenciesStatusObject.GetTypeKind(),
		Status:   make(map[string]*DependencyStatus),
	}
	for _, depID := range dependencyIds {
		parts := strings.Split(depID, "^")
		dObj, err := policy.GetObject(lang.DependencyObject.Kind, parts[1], parts[0])
		if dObj == nil || err != nil {
			dKey := runtime.KeyFromParts(parts[0], lang.DependencyObject.Kind, parts[1])
			result.Status[dKey] = &DependencyStatus{
				Found:     false,
				Deployed:  false,
				Ready:     false,
				Endpoints: make(map[string]map[string]string),
			}
			continue
		}

		d := dObj.(*lang.Dependency)
		result.Status[runtime.KeyForStorable(d)] = &DependencyStatus{
			Found:     true,
			Deployed:  true,
			Ready:     true,
			Endpoints: make(map[string]map[string]string),
		}
	}

	// load actual and desired states
	desiredState := resolve.NewPolicyResolver(policy, api.externalData, event.NewLog(logrus.WarnLevel, "api-dependencies-status")).ResolveAllDependencies()
	actualState, err := api.store.GetActualState()
	if err != nil {
		panic(fmt.Sprintf("can't load actual state from the store: %s", err))
	}

	// fetch deployment status for dependencies
	fetchDeploymentStatusForDependencies(result, actualState, desiredState)

	// fetch readiness status for dependencies, if we were asked to do so
	if flag == DependencyQueryDeploymentStatusAndReadiness {
		plugins := api.pluginRegistryFactory()
		fetchReadinessStatusForDependencies(result, plugins, policy, actualState)
	}

	// fetch endpoints for dependencies
	fetchEndpointsForDependencies(result, actualState)

	// return the result back
	api.contentType.WriteOne(writer, request, result)
}

func fetchDeploymentStatusForDependencies(result *DependenciesStatus, actualState *resolve.PolicyResolution, desiredState *resolve.PolicyResolution) {
	// compare desired vs. actual state and see what's the dependency status for every provided dependency ID
	diff.NewPolicyResolutionDiff(desiredState, actualState).ActionPlan.Apply(
		action.WrapSequential(func(act action.Base) error {
			// if it's attach action is pending on component, let's see which particular dependency it affects
			if dAction, ok := act.(*component.AttachDependencyAction); ok {
				// reset status of this particular dependency to false
				if _, affected := result.Status[dAction.DependencyID]; affected {
					result.Status[dAction.DependencyID].Deployed = false
					return nil
				}
			}

			// if it's detach action is pending on component, let's see which particular dependency it affects
			if dAction, ok := act.(*component.DetachDependencyAction); ok {
				// reset status of this particular dependency to false
				if _, affected := result.Status[dAction.DependencyID]; affected {
					result.Status[dAction.DependencyID].Deployed = false
					return nil
				}
			}

			key, ok := act.DescribeChanges()["key"].(string)
			if ok && len(key) > 0 {
				// we found a component in diff, which is affected by the action. let's see if any of the dependencies are affected
				affectedDepKeys := make(map[string]bool)
				{
					prevInstance := actualState.ComponentInstanceMap[key]
					if prevInstance != nil {
						for dKey := range prevInstance.DependencyKeys {
							affectedDepKeys[dKey] = true
						}
					}
				}
				{
					nextInstance := desiredState.ComponentInstanceMap[key]
					if nextInstance != nil {
						for dKey := range nextInstance.DependencyKeys {
							affectedDepKeys[dKey] = true
						}
					}
				}

				// if our dependency is affected, reset its deployed status to false (because actions are pending)
				for dKey := range affectedDepKeys {
					if _, ok := result.Status[dKey]; ok {
						result.Status[dKey].Deployed = false
					}
				}
			}
			return nil
		}),
		action.NewApplyResultUpdaterImpl(),
	)

	// TODO: once endpoints are retrieved in a separate go routing, we need to uncomment this

	/*
		for _, instance := range actualState.ComponentInstanceMap {
			if instance.IsCode && !instance.EndpointsUpToDate {
				for dKey := range instance.DependencyKeys {
					if _, ok := result.Status[dKey]; ok {
						result.Status[dKey].Deployed = false
					}
				}
			}
		}
	*/
}

func fetchReadinessStatusForDependencies(result *DependenciesStatus, plugins plugin.Registry, policy *lang.Policy, actualState *resolve.PolicyResolution) {
	// if dependency is not deployed, it means it's not ready
	for dKey := range result.Status {
		result.Status[dKey].Ready = result.Status[dKey].Ready && result.Status[dKey].Deployed
	}

	// update readiness
	dUpdateMutex := sync.Mutex{}
	var wg sync.WaitGroup
	errors := make(chan error, 1)
	for _, instance := range actualState.ComponentInstanceMap {
		// if component instance is not code, skip it
		if !instance.IsCode {
			continue
		}

		// we only need to query status of this component, if at least one dependency is still Ready
		foundDependenciesToCheck := false
		for dKey := range instance.DependencyKeys {
			if _, ok := result.Status[dKey]; ok && result.Status[dKey].Ready {
				foundDependenciesToCheck = true
				break
			}
		}

		// if we don't need to query status of this component instance, let's just return
		if !foundDependenciesToCheck {
			return
		}

		wg.Add(1)
		go func(instance *resolve.ComponentInstance) {
			// make sure we are converting panics into errors
			defer wg.Done()
			defer func() {
				if err := recover(); err != nil {
					select {
					case errors <- fmt.Errorf("panic: %s\n%s", err, string(debug.Stack())):
						// message sent
					default:
						// error was already there before, do nothing (but we have to keep an empty default block)
					}
				}
			}()

			// query status of this component instance
			codePlugin, err := pluginForComponentInstance(instance, policy, plugins)
			if err != nil {
				panic(fmt.Sprintf("Can't get plugin for component instance %s: %s", instance.GetKey(), err))
			}
			if codePlugin == nil {
				return
			}

			instanceStatus, err := codePlugin.Status(
				&plugin.CodePluginInvocationParams{
					DeployName:   instance.GetDeployName(),
					Params:       instance.CalculatedCodeParams,
					PluginParams: map[string]string{plugin.ParamTargetSuffix: instance.Metadata.Key.TargetSuffix},
					EventLog:     event.NewLog(logrus.WarnLevel, "resources-status"),
				},
			)
			if err != nil {
				panic(fmt.Sprintf("Error while getting deployment resources status for component instance %s: %s", instance.GetKey(), err))
			}

			// update status of dependencies
			dUpdateMutex.Lock()
			defer dUpdateMutex.Unlock()
			for dKey := range instance.DependencyKeys {
				if _, ok := result.Status[dKey]; ok {
					result.Status[dKey].Ready = result.Status[dKey].Ready && instanceStatus
				}
			}
		}(instance)
	}

	// wait until all go routines are over
	wg.Wait()

	// see if there were any errors
	select {
	case err := <-errors:
		panic(err)
	default:
		// no error, do nothing (but we have to keep an empty default block)
	}
}

func fetchEndpointsForDependencies(result *DependenciesStatus, actualState *resolve.PolicyResolution) {
	for _, instance := range actualState.ComponentInstanceMap {
		for dKey := range instance.DependencyKeys {
			if _, ok := result.Status[dKey]; ok {
				if len(instance.Endpoints) > 0 {
					result.Status[dKey].Endpoints[instance.GetName()] = instance.Endpoints
				}
			}
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package discoverspaces

import (
	"github.com/juju/errors"
	"github.com/juju/juju/api/discoverspaces"
	"github.com/juju/juju/apiserver/params"
	"github.com/juju/juju/environs"
	"github.com/juju/juju/worker"
	"github.com/juju/loggo"
	"github.com/juju/names"
	"github.com/juju/utils/set"
	"launchpad.net/tomb"
)

var logger = loggo.GetLogger("juju.discoverspaces")

type discoverspacesWorker struct {
	api      *discoverspaces.API
	tomb     tomb.Tomb
	observer *worker.EnvironObserver
}

// NewWorker returns a worker
func NewWorker(api *discoverspaces.API) worker.Worker {
	dw := &discoverspacesWorker{
		api: api,
	}
	go func() {
		defer dw.tomb.Done()
		dw.tomb.Kill(dw.loop())
	}()
	return dw
}

func (dw *discoverspacesWorker) Kill() {
	dw.tomb.Kill(nil)
}

func (dw *discoverspacesWorker) Wait() error {
	return dw.tomb.Wait()
}

func (dw *discoverspacesWorker) loop() (err error) {
	dw.observer, err = worker.NewEnvironObserver(dw.api)
	if err != nil {
		return err
	}
	defer func() {
		obsErr := worker.Stop(dw.observer)
		if err == nil {
			err = obsErr
		}
	}()
	environ := dw.observer.Environ()
	networkingEnviron, ok := environs.SupportsNetworking(environ)

	if ok {
		err = dw.handleSubnets(networkingEnviron)
		if err != nil {
			return errors.Trace(err)
		}
	}

	// TODO(mfoord): we'll have a watcher here checking if we need to
	// update the spaces/subnets definition.
	dying := dw.tomb.Dying()
	for {
		select {
		case <-dying:
			return nil
		}
	}
	return nil
}

func (dw *discoverspacesWorker) handleSubnets(env environs.NetworkingEnviron) error {
	ok, err := env.SupportsSpaceDiscovery()
	if err != nil {
		return errors.Trace(err)
	}
	if !ok {
		// Nothing to do.
		return nil
	}
	providerSpaces, err := env.Spaces()
	if err != nil {
		return errors.Trace(err)
	}
	listSpacesResult, err := dw.api.ListSpaces()
	if err != nil {
		return errors.Trace(err)
	}
	stateSpaceMap := make(map[string]params.ProviderSpace)
	spaceNames := make(set.Strings)
	for _, space := range listSpacesResult.Results {
		stateSpaceMap[space.ProviderId] = space
		spaceNames.Add(space.Name)
	}

	// TODO(mfoord): we need to delete spaces and subnets that no longer
	// exist, so long as they're not in use.
	for _, space := range providerSpaces {
		_, ok := stateSpaceMap[string(space.ProviderId)]
		if !ok {
			spaceTag, err := names.ParseSpaceTag(string(space.ProviderId))
			if err != nil {
				// XXX generate a valid name here
				return errors.Trace(err)
			}
			// We need to create the space.
			args := params.CreateSpacesParams{
				Spaces: []params.CreateSpaceParams{{
					Public:   false,
					SpaceTag: spaceTag.String(),
				}}}
			// XXX check the error result too.
			_, err = dw.api.CreateSpaces(args)
			if err != nil {
				return errors.Trace(err)
			}
		}
		// TODO(mfoord): currently no way of removing subnets, or
		// changing the space they're in, so we can only add ones we
		// don't already know about.
		for _, subnet := range space.Subnets {
			spaceTag, err := names.ParseSpaceTag(space.Name)
			if err != nil {
				return errors.Trace(err)
			}
			subnetTag, err := names.ParseSubnetTag(subnet.CIDR)
			if err != nil {
				return errors.Trace(err)
			}

			args := params.AddSubnetsParams{
				Subnets: []params.AddSubnetParams{{
					SubnetTag:        subnetTag.String(),
					SubnetProviderId: string(subnet.ProviderId),
					SpaceTag:         spaceTag.String(),
					Zones:            subnet.AvailabilityZones,
				}}}
			// XXX check the error result too.
			_, err = dw.api.AddSubnets(args)
			if err != nil {
				return errors.Trace(err)
			}
		}
	}
	return nil
}
<commit_msg>Comment updates<commit_after>// Copyright 2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package discoverspaces

import (
	"github.com/juju/errors"
	"github.com/juju/juju/api/discoverspaces"
	"github.com/juju/juju/apiserver/params"
	"github.com/juju/juju/environs"
	"github.com/juju/juju/worker"
	"github.com/juju/loggo"
	"github.com/juju/names"
	"github.com/juju/utils/set"
	"launchpad.net/tomb"
)

var logger = loggo.GetLogger("juju.discoverspaces")

type discoverspacesWorker struct {
	api      *discoverspaces.API
	tomb     tomb.Tomb
	observer *worker.EnvironObserver
}

// NewWorker returns a worker
func NewWorker(api *discoverspaces.API) worker.Worker {
	dw := &discoverspacesWorker{
		api: api,
	}
	go func() {
		defer dw.tomb.Done()
		dw.tomb.Kill(dw.loop())
	}()
	return dw
}

func (dw *discoverspacesWorker) Kill() {
	dw.tomb.Kill(nil)
}

func (dw *discoverspacesWorker) Wait() error {
	return dw.tomb.Wait()
}

func (dw *discoverspacesWorker) loop() (err error) {
	dw.observer, err = worker.NewEnvironObserver(dw.api)
	if err != nil {
		return err
	}
	defer func() {
		obsErr := worker.Stop(dw.observer)
		if err == nil {
			err = obsErr
		}
	}()
	environ := dw.observer.Environ()
	networkingEnviron, ok := environs.SupportsNetworking(environ)

	if ok {
		err = dw.handleSubnets(networkingEnviron)
		if err != nil {
			return errors.Trace(err)
		}
	}

	// TODO(mfoord): we'll have a watcher here checking if we need to
	// update the spaces/subnets definition.
	dying := dw.tomb.Dying()
	for {
		select {
		case <-dying:
			return nil
		}
	}
	return nil
}

func (dw *discoverspacesWorker) handleSubnets(env environs.NetworkingEnviron) error {
	ok, err := env.SupportsSpaceDiscovery()
	if err != nil {
		return errors.Trace(err)
	}
	if !ok {
		// Nothing to do.
		return nil
	}
	providerSpaces, err := env.Spaces()
	if err != nil {
		return errors.Trace(err)
	}
	listSpacesResult, err := dw.api.ListSpaces()
	if err != nil {
		return errors.Trace(err)
	}
	stateSpaceMap := make(map[string]params.ProviderSpace)
	spaceNames := make(set.Strings)
	for _, space := range listSpacesResult.Results {
		stateSpaceMap[space.ProviderId] = space
		spaceNames.Add(space.Name)
	}

	// TODO(mfoord): we need to delete spaces and subnets that no longer
	// exist, so long as they're not in use.
	for _, space := range providerSpaces {
		_, ok := stateSpaceMap[string(space.ProviderId)]
		if !ok {
			// XXX skip spaces with no subnets(?)
			spaceTag, err := names.ParseSpaceTag(string(space.ProviderId))
			if err != nil {
				// XXX generate a valid name here
				return errors.Trace(err)
			}
			// We need to create the space.
			args := params.CreateSpacesParams{
				Spaces: []params.CreateSpaceParams{{
					Public:   false,
					SpaceTag: spaceTag.String(),
				}}}
			// XXX check the error result too.
			_, err = dw.api.CreateSpaces(args)
			if err != nil {
				return errors.Trace(err)
			}
		}
		// TODO(mfoord): currently no way of removing subnets, or
		// changing the space they're in, so we can only add ones we
		// don't already know about.
		for _, subnet := range space.Subnets {
			spaceTag, err := names.ParseSpaceTag(space.Name)
			if err != nil {
				return errors.Trace(err)
			}
			subnetTag, err := names.ParseSubnetTag(subnet.CIDR)
			if err != nil {
				return errors.Trace(err)
			}

			args := params.AddSubnetsParams{
				Subnets: []params.AddSubnetParams{{
					SubnetTag:        subnetTag.String(),
					SubnetProviderId: string(subnet.ProviderId),
					SpaceTag:         spaceTag.String(),
					Zones:            subnet.AvailabilityZones,
				}}}
			// XXX check the error result too.
			_, err = dw.api.AddSubnets(args)
			if err != nil {
				return errors.Trace(err)
			}
		}
	}
	return nil
}
<|endoftext|>


<commit_before>/*
 * Datadog API for Go
 *
 * Please see the included LICENSE file for licensing information.
 *
 * Copyright 2013 by authors and contributors.
 */

package datadog

import (
	"encoding/json"
	"fmt"
	"net/url"
	"strconv"
	"strings"
)

type Period struct {
	Seconds *json.Number `json:"seconds,omitempty"`
	Text    *string      `json:"text,omitempty"`
	Value   *string      `json:"value,omitempty"`
	Name    *string      `json:"name,omitempty"`
	Unit    *string      `json:"unit,omitempty"`
}

type LogSet struct {
	ID   *json.Number `json:"id,omitempty"`
	Name *string      `json:"name,omitempty"`
}

type TimeRange struct {
	To   *json.Number `json:"to,omitempty"`
	From *json.Number `json:"from,omitempty"`
	Live *bool        `json:"live,omitempty"`
}

type QueryConfig struct {
	LogSet        *LogSet    `json:"logset,omitempty"`
	TimeRange     *TimeRange `json:"timeRange,omitempty"`
	QueryString   *string    `json:"queryString,omitempty"`
	QueryIsFailed *bool      `json:"queryIsFailed,omitempty"`
}

type ThresholdCount struct {
	Ok               *json.Number `json:"ok,omitempty"`
	Critical         *json.Number `json:"critical,omitempty"`
	Warning          *json.Number `json:"warning,omitempty"`
	Unknown          *json.Number `json:"unknown,omitempty"`
	CriticalRecovery *json.Number `json:"critical_recovery,omitempty"`
	WarningRecovery  *json.Number `json:"warning_recovery,omitempty"`
	Period           *Period      `json:"period,omitempty"`
	TimeAggregator   *string      `json:"timeAggregator,omitempty"`
}

type ThresholdWindows struct {
	RecoveryWindow *string `json:"recovery_window,omitempty"`
	TriggerWindow  *string `json:"trigger_window,omitempty"`
}

type NoDataTimeframe int

func (tf *NoDataTimeframe) UnmarshalJSON(data []byte) error {
	s := string(data)
	if s == "false" || s == "null" {
		*tf = 0
	} else {
		i, err := strconv.ParseInt(s, 10, 32)
		if err != nil {
			return err
		}
		*tf = NoDataTimeframe(i)
	}
	return nil
}

type Options struct {
	NoDataTimeframe   NoDataTimeframe   `json:"no_data_timeframe,omitempty"`
	NotifyAudit       *bool             `json:"notify_audit,omitempty"`
	NotifyNoData      *bool             `json:"notify_no_data,omitempty"`
	RenotifyInterval  *int              `json:"renotify_interval,omitempty"`
	NewHostDelay      *int              `json:"new_host_delay,omitempty"`
	EvaluationDelay   *int              `json:"evaluation_delay,omitempty"`
	Silenced          map[string]int    `json:"silenced,omitempty"`
	TimeoutH          *int              `json:"timeout_h,omitempty"`
	EscalationMessage *string           `json:"escalation_message,omitempty"`
	Thresholds        *ThresholdCount   `json:"thresholds,omitempty"`
	ThresholdWindows  *ThresholdWindows `json:"threshold_windows,omitempty"`
	IncludeTags       *bool             `json:"include_tags,omitempty"`
	RequireFullWindow *bool             `json:"require_full_window,omitempty"`
	Locked            *bool             `json:"locked,omitempty"`
	EnableLogsSample  *bool             `json:"enable_logs_sample,omitempty"`
	QueryConfig       *QueryConfig      `json:"queryConfig,omitempty"`
}

type TriggeringValue struct {
	FromTs *int `json:"from_ts,omitempty"`
	ToTs   *int `json:"to_ts,omitempty"`
	Value  *int `json:"value,omitempty"`
}

type GroupData struct {
	LastNoDataTs    *int             `json:"last_nodata_ts,omitempty"`
	LastNotifiedTs  *int             `json:"last_notified_ts,omitempty"`
	LastResolvedTs  *int             `json:"last_resolved_ts,omitempty"`
	LastTriggeredTs *int             `json:"last_triggered_ts,omitempty"`
	Name            *string          `json:"name,omitempty"`
	Status          *string          `json:"status,omitempty"`
	TriggeringValue *TriggeringValue `json:"triggering_value,omitempty"`
}

type State struct {
	Groups map[string]GroupData `json:"groups,omitempty"`
}

// Monitor allows watching a metric or check that you care about,
// notifying your team when some defined threshold is exceeded
type Monitor struct {
	Creator              *Creator `json:"creator,omitempty"`
	Id                   *int     `json:"id,omitempty"`
	Type                 *string  `json:"type,omitempty"`
	Query                *string  `json:"query,omitempty"`
	Name                 *string  `json:"name,omitempty"`
	Message              *string  `json:"message,omitempty"`
	OverallState         *string  `json:"overall_state,omitempty"`
	OverallStateModified *string  `json:"overall_state_modified,omitempty"`
	Tags                 []string `json:"tags"`
	Options              *Options `json:"options,omitempty"`
	State                State    `json:"state,omitempty"`
}

// Creator contains the creator of the monitor
type Creator struct {
	Email  *string `json:"email,omitempty"`
	Handle *string `json:"handle,omitempty"`
	Id     *int    `json:"id,omitempty"`
	Name   *string `json:"name,omitempty"`
}

// MuteMonitorScope specifies which scope to mute and when to end the mute
type MuteMonitorScope struct {
	Scope *string `json:"scope,omitempty"`
	End   *int    `json:"end,omitempty"`
}

// UnmuteMonitorScopes specifies which scope(s) to unmute
type UnmuteMonitorScopes struct {
	Scope     *string `json:"scope,omitempty"`
	AllScopes *bool   `json:"all_scopes,omitempty"`
}

// reqMonitors receives a slice of all monitors
type reqMonitors struct {
	Monitors []Monitor `json:"monitors,omitempty"`
}

// CreateMonitor adds a new monitor to the system. This returns a pointer to a
// monitor so you can pass that to UpdateMonitor later if needed
func (client *Client) CreateMonitor(monitor *Monitor) (*Monitor, error) {
	var out Monitor
	// TODO: is this more pretty of frowned upon?
	if err := client.doJsonRequest("POST", "/v1/monitor", monitor, &out); err != nil {
		return nil, err
	}
	return &out, nil
}

// UpdateMonitor takes a monitor that was previously retrieved through some method
// and sends it back to the server
func (client *Client) UpdateMonitor(monitor *Monitor) error {
	return client.doJsonRequest("PUT", fmt.Sprintf("/v1/monitor/%d", *monitor.Id),
		monitor, nil)
}

// GetMonitor retrieves a monitor by identifier
func (client *Client) GetMonitor(id int) (*Monitor, error) {
	var out Monitor
	if err := client.doJsonRequest("GET", fmt.Sprintf("/v1/monitor/%d", id), nil, &out); err != nil {
		return nil, err
	}
	return &out, nil
}

// GetMonitorsByName retrieves monitors by name
func (client *Client) GetMonitorsByName(name string) ([]Monitor, error) {
	return client.GetMonitorsWithOptions(MonitorQueryOpts{Name: &name})
}

// GetMonitorsByTags retrieves monitors by a slice of tags
func (client *Client) GetMonitorsByTags(tags []string) ([]Monitor, error) {
	return client.GetMonitorsWithOptions(MonitorQueryOpts{Tags: tags})
}

// GetMonitorsByMonitorTags retrieves monitors by a slice of monitor tags
func (client *Client) GetMonitorsByMonitorTags(tags []string) ([]Monitor, error) {
	return client.GetMonitorsWithOptions(MonitorQueryOpts{MonitorTags: tags})
}

// DeleteMonitor removes a monitor from the system
func (client *Client) DeleteMonitor(id int) error {
	return client.doJsonRequest("DELETE", fmt.Sprintf("/v1/monitor/%d", id),
		nil, nil)
}

// GetMonitors returns a slice of all monitors
func (client *Client) GetMonitors() ([]Monitor, error) {
	return client.GetMonitorsWithOptions(MonitorQueryOpts{})
}

// MonitorQueryOpts contains the options supported by
// https://docs.datadoghq.com/api/?lang=bash#get-all-monitor-details
type MonitorQueryOpts struct {
	GroupStates   []string
	Name          *string
	Tags          []string
	MonitorTags   []string
	WithDowntimes *bool
}

// GetMonitorsWithOptions returns a slice of all monitors
// It supports all the options for querying
func (client *Client) GetMonitorsWithOptions(opts MonitorQueryOpts) ([]Monitor, error) {
	var out reqMonitors
	var query []string
	if len(opts.Tags) > 0 {
		value := fmt.Sprintf("tags=%v", strings.Join(opts.Tags, ","))
		query = append(query, value)
	}

	if len(opts.GroupStates) > 0 {
		value := fmt.Sprintf("group_states=%v", strings.Join(opts.GroupStates, ","))
		query = append(query, value)
	}

	if len(opts.MonitorTags) > 0 {
		value := fmt.Sprintf("monitor_tags=%v", strings.Join(opts.MonitorTags, ","))
		query = append(query, value)
	}

	if v, ok := opts.GetWithDowntimesOk(); ok {
		query = append(query, fmt.Sprintf("with_downtimes=%t", v))
	}

	if v, ok := opts.GetNameOk(); ok {
		query = append(query, fmt.Sprintf("name=%s", v))
	}

	queryString, err := url.ParseQuery(strings.Join(query, "&"))
	if err != nil {
		return nil, err
	}
	err = client.doJsonRequest("GET", fmt.Sprintf("/v1/monitor?%v", queryString.Encode()), nil, &out.Monitors)
	if err != nil {
		return nil, err
	}
	return out.Monitors, nil
}

// MuteMonitors turns off monitoring notifications
func (client *Client) MuteMonitors() error {
	return client.doJsonRequest("POST", "/v1/monitor/mute_all", nil, nil)
}

// UnmuteMonitors turns on monitoring notifications
func (client *Client) UnmuteMonitors() error {
	return client.doJsonRequest("POST", "/v1/monitor/unmute_all", nil, nil)
}

// MuteMonitor turns off monitoring notifications for a monitor
func (client *Client) MuteMonitor(id int) error {
	return client.doJsonRequest("POST", fmt.Sprintf("/v1/monitor/%d/mute", id), nil, nil)
}

// MuteMonitorScope turns off monitoring notifications for a monitor for a given scope
func (client *Client) MuteMonitorScope(id int, muteMonitorScope *MuteMonitorScope) error {
	return client.doJsonRequest("POST", fmt.Sprintf("/v1/monitor/%d/mute", id), muteMonitorScope, nil)
}

// UnmuteMonitor turns on monitoring notifications for a monitor
func (client *Client) UnmuteMonitor(id int) error {
	return client.doJsonRequest("POST", fmt.Sprintf("/v1/monitor/%d/unmute", id), nil, nil)
}

// UnmuteMonitorScopes is similar to UnmuteMonitor, but provides finer-grained control to unmuting
func (client *Client) UnmuteMonitorScopes(id int, unmuteMonitorScopes *UnmuteMonitorScopes) error {
	return client.doJsonRequest("POST", fmt.Sprintf("/v1/monitor/%d/unmute", id), unmuteMonitorScopes, nil)
}
<commit_msg>Adds force delete monitor (#311)<commit_after>/*
 * Datadog API for Go
 *
 * Please see the included LICENSE file for licensing information.
 *
 * Copyright 2013 by authors and contributors.
 */

package datadog

import (
	"encoding/json"
	"fmt"
	"net/url"
	"strconv"
	"strings"
)

type Period struct {
	Seconds *json.Number `json:"seconds,omitempty"`
	Text    *string      `json:"text,omitempty"`
	Value   *string      `json:"value,omitempty"`
	Name    *string      `json:"name,omitempty"`
	Unit    *string      `json:"unit,omitempty"`
}

type LogSet struct {
	ID   *json.Number `json:"id,omitempty"`
	Name *string      `json:"name,omitempty"`
}

type TimeRange struct {
	To   *json.Number `json:"to,omitempty"`
	From *json.Number `json:"from,omitempty"`
	Live *bool        `json:"live,omitempty"`
}

type QueryConfig struct {
	LogSet        *LogSet    `json:"logset,omitempty"`
	TimeRange     *TimeRange `json:"timeRange,omitempty"`
	QueryString   *string    `json:"queryString,omitempty"`
	QueryIsFailed *bool      `json:"queryIsFailed,omitempty"`
}

type ThresholdCount struct {
	Ok               *json.Number `json:"ok,omitempty"`
	Critical         *json.Number `json:"critical,omitempty"`
	Warning          *json.Number `json:"warning,omitempty"`
	Unknown          *json.Number `json:"unknown,omitempty"`
	CriticalRecovery *json.Number `json:"critical_recovery,omitempty"`
	WarningRecovery  *json.Number `json:"warning_recovery,omitempty"`
	Period           *Period      `json:"period,omitempty"`
	TimeAggregator   *string      `json:"timeAggregator,omitempty"`
}

type ThresholdWindows struct {
	RecoveryWindow *string `json:"recovery_window,omitempty"`
	TriggerWindow  *string `json:"trigger_window,omitempty"`
}

type NoDataTimeframe int

func (tf *NoDataTimeframe) UnmarshalJSON(data []byte) error {
	s := string(data)
	if s == "false" || s == "null" {
		*tf = 0
	} else {
		i, err := strconv.ParseInt(s, 10, 32)
		if err != nil {
			return err
		}
		*tf = NoDataTimeframe(i)
	}
	return nil
}

type Options struct {
	NoDataTimeframe   NoDataTimeframe   `json:"no_data_timeframe,omitempty"`
	NotifyAudit       *bool             `json:"notify_audit,omitempty"`
	NotifyNoData      *bool             `json:"notify_no_data,omitempty"`
	RenotifyInterval  *int              `json:"renotify_interval,omitempty"`
	NewHostDelay      *int              `json:"new_host_delay,omitempty"`
	EvaluationDelay   *int              `json:"evaluation_delay,omitempty"`
	Silenced          map[string]int    `json:"silenced,omitempty"`
	TimeoutH          *int              `json:"timeout_h,omitempty"`
	EscalationMessage *string           `json:"escalation_message,omitempty"`
	Thresholds        *ThresholdCount   `json:"thresholds,omitempty"`
	ThresholdWindows  *ThresholdWindows `json:"threshold_windows,omitempty"`
	IncludeTags       *bool             `json:"include_tags,omitempty"`
	RequireFullWindow *bool             `json:"require_full_window,omitempty"`
	Locked            *bool             `json:"locked,omitempty"`
	EnableLogsSample  *bool             `json:"enable_logs_sample,omitempty"`
	QueryConfig       *QueryConfig      `json:"queryConfig,omitempty"`
}

type TriggeringValue struct {
	FromTs *int `json:"from_ts,omitempty"`
	ToTs   *int `json:"to_ts,omitempty"`
	Value  *int `json:"value,omitempty"`
}

type GroupData struct {
	LastNoDataTs    *int             `json:"last_nodata_ts,omitempty"`
	LastNotifiedTs  *int             `json:"last_notified_ts,omitempty"`
	LastResolvedTs  *int             `json:"last_resolved_ts,omitempty"`
	LastTriggeredTs *int             `json:"last_triggered_ts,omitempty"`
	Name            *string          `json:"name,omitempty"`
	Status          *string          `json:"status,omitempty"`
	TriggeringValue *TriggeringValue `json:"triggering_value,omitempty"`
}

type State struct {
	Groups map[string]GroupData `json:"groups,omitempty"`
}

// Monitor allows watching a metric or check that you care about,
// notifying your team when some defined threshold is exceeded
type Monitor struct {
	Creator              *Creator `json:"creator,omitempty"`
	Id                   *int     `json:"id,omitempty"`
	Type                 *string  `json:"type,omitempty"`
	Query                *string  `json:"query,omitempty"`
	Name                 *string  `json:"name,omitempty"`
	Message              *string  `json:"message,omitempty"`
	OverallState         *string  `json:"overall_state,omitempty"`
	OverallStateModified *string  `json:"overall_state_modified,omitempty"`
	Tags                 []string `json:"tags"`
	Options              *Options `json:"options,omitempty"`
	State                State    `json:"state,omitempty"`
}

// Creator contains the creator of the monitor
type Creator struct {
	Email  *string `json:"email,omitempty"`
	Handle *string `json:"handle,omitempty"`
	Id     *int    `json:"id,omitempty"`
	Name   *string `json:"name,omitempty"`
}

// MuteMonitorScope specifies which scope to mute and when to end the mute
type MuteMonitorScope struct {
	Scope *string `json:"scope,omitempty"`
	End   *int    `json:"end,omitempty"`
}

// UnmuteMonitorScopes specifies which scope(s) to unmute
type UnmuteMonitorScopes struct {
	Scope     *string `json:"scope,omitempty"`
	AllScopes *bool   `json:"all_scopes,omitempty"`
}

// reqMonitors receives a slice of all monitors
type reqMonitors struct {
	Monitors []Monitor `json:"monitors,omitempty"`
}

// CreateMonitor adds a new monitor to the system. This returns a pointer to a
// monitor so you can pass that to UpdateMonitor later if needed
func (client *Client) CreateMonitor(monitor *Monitor) (*Monitor, error) {
	var out Monitor
	// TODO: is this more pretty of frowned upon?
	if err := client.doJsonRequest("POST", "/v1/monitor", monitor, &out); err != nil {
		return nil, err
	}
	return &out, nil
}

// UpdateMonitor takes a monitor that was previously retrieved through some method
// and sends it back to the server
func (client *Client) UpdateMonitor(monitor *Monitor) error {
	return client.doJsonRequest("PUT", fmt.Sprintf("/v1/monitor/%d", *monitor.Id),
		monitor, nil)
}

// GetMonitor retrieves a monitor by identifier
func (client *Client) GetMonitor(id int) (*Monitor, error) {
	var out Monitor
	if err := client.doJsonRequest("GET", fmt.Sprintf("/v1/monitor/%d", id), nil, &out); err != nil {
		return nil, err
	}
	return &out, nil
}

// GetMonitorsByName retrieves monitors by name
func (client *Client) GetMonitorsByName(name string) ([]Monitor, error) {
	return client.GetMonitorsWithOptions(MonitorQueryOpts{Name: &name})
}

// GetMonitorsByTags retrieves monitors by a slice of tags
func (client *Client) GetMonitorsByTags(tags []string) ([]Monitor, error) {
	return client.GetMonitorsWithOptions(MonitorQueryOpts{Tags: tags})
}

// GetMonitorsByMonitorTags retrieves monitors by a slice of monitor tags
func (client *Client) GetMonitorsByMonitorTags(tags []string) ([]Monitor, error) {
	return client.GetMonitorsWithOptions(MonitorQueryOpts{MonitorTags: tags})
}

// DeleteMonitor removes a monitor from the system
func (client *Client) DeleteMonitor(id int) error {
	return client.doJsonRequest("DELETE", fmt.Sprintf("/v1/monitor/%d", id),
		nil, nil)
}

// ForceDeleteMonitor removes a monitor from the system, even if it's linked to SLOs or group monitors
func (client *Client) ForceDeleteMonitor(id int) error {
	return client.doJsonRequest("DELETE", fmt.Sprintf("/v1/monitor/%d?force=true", id),
		nil, nil)
}

// GetMonitors returns a slice of all monitors
func (client *Client) GetMonitors() ([]Monitor, error) {
	return client.GetMonitorsWithOptions(MonitorQueryOpts{})
}

// MonitorQueryOpts contains the options supported by
// https://docs.datadoghq.com/api/?lang=bash#get-all-monitor-details
type MonitorQueryOpts struct {
	GroupStates   []string
	Name          *string
	Tags          []string
	MonitorTags   []string
	WithDowntimes *bool
}

// GetMonitorsWithOptions returns a slice of all monitors
// It supports all the options for querying
func (client *Client) GetMonitorsWithOptions(opts MonitorQueryOpts) ([]Monitor, error) {
	var out reqMonitors
	var query []string
	if len(opts.Tags) > 0 {
		value := fmt.Sprintf("tags=%v", strings.Join(opts.Tags, ","))
		query = append(query, value)
	}

	if len(opts.GroupStates) > 0 {
		value := fmt.Sprintf("group_states=%v", strings.Join(opts.GroupStates, ","))
		query = append(query, value)
	}

	if len(opts.MonitorTags) > 0 {
		value := fmt.Sprintf("monitor_tags=%v", strings.Join(opts.MonitorTags, ","))
		query = append(query, value)
	}

	if v, ok := opts.GetWithDowntimesOk(); ok {
		query = append(query, fmt.Sprintf("with_downtimes=%t", v))
	}

	if v, ok := opts.GetNameOk(); ok {
		query = append(query, fmt.Sprintf("name=%s", v))
	}

	queryString, err := url.ParseQuery(strings.Join(query, "&"))
	if err != nil {
		return nil, err
	}
	err = client.doJsonRequest("GET", fmt.Sprintf("/v1/monitor?%v", queryString.Encode()), nil, &out.Monitors)
	if err != nil {
		return nil, err
	}
	return out.Monitors, nil
}

// MuteMonitors turns off monitoring notifications
func (client *Client) MuteMonitors() error {
	return client.doJsonRequest("POST", "/v1/monitor/mute_all", nil, nil)
}

// UnmuteMonitors turns on monitoring notifications
func (client *Client) UnmuteMonitors() error {
	return client.doJsonRequest("POST", "/v1/monitor/unmute_all", nil, nil)
}

// MuteMonitor turns off monitoring notifications for a monitor
func (client *Client) MuteMonitor(id int) error {
	return client.doJsonRequest("POST", fmt.Sprintf("/v1/monitor/%d/mute", id), nil, nil)
}

// MuteMonitorScope turns off monitoring notifications for a monitor for a given scope
func (client *Client) MuteMonitorScope(id int, muteMonitorScope *MuteMonitorScope) error {
	return client.doJsonRequest("POST", fmt.Sprintf("/v1/monitor/%d/mute", id), muteMonitorScope, nil)
}

// UnmuteMonitor turns on monitoring notifications for a monitor
func (client *Client) UnmuteMonitor(id int) error {
	return client.doJsonRequest("POST", fmt.Sprintf("/v1/monitor/%d/unmute", id), nil, nil)
}

// UnmuteMonitorScopes is similar to UnmuteMonitor, but provides finer-grained control to unmuting
func (client *Client) UnmuteMonitorScopes(id int, unmuteMonitorScopes *UnmuteMonitorScopes) error {
	return client.doJsonRequest("POST", fmt.Sprintf("/v1/monitor/%d/unmute", id), unmuteMonitorScopes, nil)
}
<|endoftext|>


<commit_before>// Copyright 2010 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// TODO(gri) consider making this a separate package outside the go directory.

package token

import (
	"fmt"
	"sort"
	"sync"
)

// Position describes an arbitrary source position
// including the file, line, and column location.
// A Position is valid if the line number is > 0.
//
type Position struct {
	Filename string // filename, if any
	Offset   int    // offset, starting at 0
	Line     int    // line number, starting at 1
	Column   int    // column number, starting at 1 (character count)
}

// IsValid returns true if the position is valid.
func (pos *Position) IsValid() bool { return pos.Line > 0 }

// String returns a string in one of several forms:
//
//	file:line:column    valid position with file name
//	line:column         valid position without file name
//	file                invalid position with file name
//	-                   invalid position without file name
//
func (pos Position) String() string {
	s := pos.Filename
	if pos.IsValid() {
		if s != "" {
			s += ":"
		}
		s += fmt.Sprintf("%d:%d", pos.Line, pos.Column)
	}
	if s == "" {
		s = "-"
	}
	return s
}

// Pos is a compact encoding of a source position within a file set.
// It can be converted into a Position for a more convenient, but much
// larger, representation.
//
// The Pos value for a given file is a number in the range [base, base+size],
// where base and size are specified when adding the file to the file set via
// AddFile.
//
// To create the Pos value for a specific source offset, first add
// the respective file to the current file set (via FileSet.AddFile)
// and then call File.Pos(offset) for that file. Given a Pos value p
// for a specific file set fset, the corresponding Position value is
// obtained by calling fset.Position(p).
//
// Pos values can be compared directly with the usual comparison operators:
// If two Pos values p and q are in the same file, comparing p and q is
// equivalent to comparing the respective source file offsets. If p and q
// are in different files, p < q is true if the file implied by p was added
// to the respective file set before the file implied by q.
//
type Pos int

// The zero value for Pos is NoPos; there is no file and line information
// associated with it, and NoPos().IsValid() is false. NoPos is always
// smaller than any other Pos value. The corresponding Position value
// for NoPos is the zero value for Position.
//
const NoPos Pos = 0

// IsValid returns true if the position is valid.
func (p Pos) IsValid() bool {
	return p != NoPos
}

func searchFiles(a []*File, x int) int {
	return sort.Search(len(a), func(i int) bool { return a[i].base > x }) - 1
}

func (s *FileSet) file(p Pos) *File {
	if i := searchFiles(s.files, int(p)); i >= 0 {
		f := s.files[i]
		// f.base <= int(p) by definition of searchFiles
		if int(p) <= f.base+f.size {
			return f
		}
	}
	return nil
}

// File returns the file which contains the position p.
// If no such file is found (for instance for p == NoPos),
// the result is nil.
//
func (s *FileSet) File(p Pos) (f *File) {
	if p != NoPos {
		s.mutex.RLock()
		f = s.file(p)
		s.mutex.RUnlock()
	}
	return
}

func (f *File) position(p Pos) (pos Position) {
	offset := int(p) - f.base
	pos.Offset = offset
	pos.Filename, pos.Line, pos.Column = f.info(offset)
	return
}

// Position converts a Pos in the fileset into a general Position.
func (s *FileSet) Position(p Pos) (pos Position) {
	if p != NoPos {
		// TODO(gri) consider optimizing the case where p
		//           is in the last file addded, or perhaps
		//           looked at - will eliminate one level
		//           of search
		s.mutex.RLock()
		if f := s.file(p); f != nil {
			pos = f.position(p)
		}
		s.mutex.RUnlock()
	}
	return
}

type lineInfo struct {
	offset   int
	filename string
	line     int
}

// AddLineInfo adds alternative file and line number information for
// a given file offset. The offset must be larger than the offset for
// the previously added alternative line info and smaller than the
// file size; otherwise the information is ignored.
//
// AddLineInfo is typically used to register alternative position
// information for //line filename:line comments in source files.
//
func (f *File) AddLineInfo(offset int, filename string, line int) {
	f.set.mutex.Lock()
	if i := len(f.infos); i == 0 || f.infos[i-1].offset < offset && offset < f.size {
		f.infos = append(f.infos, lineInfo{offset, filename, line})
	}
	f.set.mutex.Unlock()
}

// A File is a handle for a file belonging to a FileSet.
// A File has a name, size, and line offset table.
//
type File struct {
	set  *FileSet
	name string // file name as provided to AddFile
	base int    // Pos value range for this file is [base...base+size]
	size int    // file size as provided to AddFile

	// lines and infos are protected by set.mutex
	lines []int
	infos []lineInfo
}

// Name returns the file name of file f as registered with AddFile.
func (f *File) Name() string {
	return f.name
}

// Base returns the base offset of file f as registered with AddFile.
func (f *File) Base() int {
	return f.base
}

// Size returns the size of file f as registered with AddFile.
func (f *File) Size() int {
	return f.size
}

// LineCount returns the number of lines in file f.
func (f *File) LineCount() int {
	f.set.mutex.RLock()
	n := len(f.lines)
	f.set.mutex.RUnlock()
	return n
}

// AddLine adds the line offset for a new line.
// The line offset must be larger than the offset for the previous line
// and smaller than the file size; otherwise the line offset is ignored.
//
func (f *File) AddLine(offset int) {
	f.set.mutex.Lock()
	if i := len(f.lines); (i == 0 || f.lines[i-1] < offset) && offset < f.size {
		f.lines = append(f.lines, offset)
	}
	f.set.mutex.Unlock()
}

// SetLines sets the line offsets for a file and returns true if successful.
// The line offsets are the offsets of the first character of each line;
// for instance for the content "ab\nc\n" the line offsets are {0, 3}.
// An empty file has an empty line offset table.
// Each line offset must be larger than the offset for the previous line
// and smaller than the file size; otherwise SetLines fails and returns
// false.
//
func (f *File) SetLines(lines []int) bool {
	// verify validity of lines table
	size := f.size
	for i, offset := range lines {
		if i > 0 && offset <= lines[i-1] || size <= offset {
			return false
		}
	}

	// set lines table
	f.set.mutex.Lock()
	f.lines = lines
	f.set.mutex.Unlock()
	return true
}

// SetLinesForContent sets the line offsets for the given file content.
func (f *File) SetLinesForContent(content []byte) {
	var lines []int
	line := 0
	for offset, b := range content {
		if line >= 0 {
			lines = append(lines, line)
		}
		line = -1
		if b == '\n' {
			line = offset + 1
		}
	}

	// set lines table
	f.set.mutex.Lock()
	f.lines = lines
	f.set.mutex.Unlock()
}

// Pos returns the Pos value for the given file offset;
// the offset must be <= f.Size().
// f.Pos(f.Offset(p)) == p.
//
func (f *File) Pos(offset int) Pos {
	if offset > f.size {
		panic("illegal file offset")
	}
	return Pos(f.base + offset)
}

// Offset returns the offset for the given file position p;
// p must be a valid Pos value in that file.
// f.Offset(f.Pos(offset)) == offset.
//
func (f *File) Offset(p Pos) int {
	if int(p) < f.base || int(p) > f.base+f.size {
		panic("illegal Pos value")
	}
	return int(p) - f.base
}

// Line returns the line number for the given file position p;
// p must be a Pos value in that file or NoPos.
//
func (f *File) Line(p Pos) int {
	// TODO(gri) this can be implemented much more efficiently
	return f.Position(p).Line
}

// Position returns the Position value for the given file position p;
// p must be a Pos value in that file or NoPos.
//
func (f *File) Position(p Pos) (pos Position) {
	if p != NoPos {
		if int(p) < f.base || int(p) > f.base+f.size {
			panic("illegal Pos value")
		}
		pos = f.position(p)
	}
	return
}

func searchUints(a []int, x int) int {
	return sort.Search(len(a), func(i int) bool { return a[i] > x }) - 1
}

func searchLineInfos(a []lineInfo, x int) int {
	return sort.Search(len(a), func(i int) bool { return a[i].offset > x }) - 1
}

// info returns the file name, line, and column number for a file offset.
func (f *File) info(offset int) (filename string, line, column int) {
	filename = f.name
	if i := searchUints(f.lines, offset); i >= 0 {
		line, column = i+1, offset-f.lines[i]+1
	}
	if i := searchLineInfos(f.infos, offset); i >= 0 {
		alt := &f.infos[i]
		filename = alt.filename
		if i := searchUints(f.lines, alt.offset); i >= 0 {
			line += alt.line - i - 1
		}
	}
	return
}

// A FileSet represents a set of source files.
// Methods of file sets are synchronized; multiple goroutines
// may invoke them concurrently.
//
type FileSet struct {
	mutex sync.RWMutex  // protects the file set
	base  int           // base offset for the next file
	files []*File       // list of files in the order added to the set
	index map[*File]int // file -> files index for quick lookup
}

// NewFileSet creates a new file set.
func NewFileSet() *FileSet {
	s := new(FileSet)
	s.base = 1 // 0 == NoPos
	s.index = make(map[*File]int)
	return s
}

// Base returns the minimum base offset that must be provided to
// AddFile when adding the next file.
//
func (s *FileSet) Base() int {
	s.mutex.RLock()
	b := s.base
	s.mutex.RUnlock()
	return b

}

// AddFile adds a new file with a given filename, base offset, and file size
// to the file set s and returns the file. Multiple files may have the same
// name. The base offset must not be smaller than the FileSet's Base(), and
// size must not be negative.
//
// Adding the file will set the file set's Base() value to base + size + 1
// as the minimum base value for the next file. The following relationship
// exists between a Pos value p for a given file offset offs:
//
//	int(p) = base + offs
//
// with offs in the range [0, size] and thus p in the range [base, base+size].
// For convenience, File.Pos may be used to create file-specific position
// values from a file offset.
//
func (s *FileSet) AddFile(filename string, base, size int) *File {
	s.mutex.Lock()
	defer s.mutex.Unlock()
	if base < s.base || size < 0 {
		panic("illegal base or size")
	}
	// base >= s.base && size >= 0
	f := &File{s, filename, base, size, []int{0}, nil}
	base += size + 1 // +1 because EOF also has a position
	if base < 0 {
		panic("token.Pos offset overflow (> 2G of source code in file set)")
	}
	// add the file to the file set
	s.base = base
	s.index[f] = len(s.files)
	s.files = append(s.files, f)
	return f
}

// Files returns the files added to the file set.
func (s *FileSet) Files() <-chan *File {
	ch := make(chan *File)
	go func() {
		for i := 0; ; i++ {
			var f *File
			s.mutex.RLock()
			if i < len(s.files) {
				f = s.files[i]
			}
			s.mutex.RUnlock()
			if f == nil {
				break
			}
			ch <- f
		}
		close(ch)
	}()
	return ch
}
<commit_msg>go/token: faster FileSet.Position implementation<commit_after>// Copyright 2010 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// TODO(gri) consider making this a separate package outside the go directory.

package token

import (
	"fmt"
	"sort"
	"sync"
)

// Position describes an arbitrary source position
// including the file, line, and column location.
// A Position is valid if the line number is > 0.
//
type Position struct {
	Filename string // filename, if any
	Offset   int    // offset, starting at 0
	Line     int    // line number, starting at 1
	Column   int    // column number, starting at 1 (character count)
}

// IsValid returns true if the position is valid.
func (pos *Position) IsValid() bool { return pos.Line > 0 }

// String returns a string in one of several forms:
//
//	file:line:column    valid position with file name
//	line:column         valid position without file name
//	file                invalid position with file name
//	-                   invalid position without file name
//
func (pos Position) String() string {
	s := pos.Filename
	if pos.IsValid() {
		if s != "" {
			s += ":"
		}
		s += fmt.Sprintf("%d:%d", pos.Line, pos.Column)
	}
	if s == "" {
		s = "-"
	}
	return s
}

// Pos is a compact encoding of a source position within a file set.
// It can be converted into a Position for a more convenient, but much
// larger, representation.
//
// The Pos value for a given file is a number in the range [base, base+size],
// where base and size are specified when adding the file to the file set via
// AddFile.
//
// To create the Pos value for a specific source offset, first add
// the respective file to the current file set (via FileSet.AddFile)
// and then call File.Pos(offset) for that file. Given a Pos value p
// for a specific file set fset, the corresponding Position value is
// obtained by calling fset.Position(p).
//
// Pos values can be compared directly with the usual comparison operators:
// If two Pos values p and q are in the same file, comparing p and q is
// equivalent to comparing the respective source file offsets. If p and q
// are in different files, p < q is true if the file implied by p was added
// to the respective file set before the file implied by q.
//
type Pos int

// The zero value for Pos is NoPos; there is no file and line information
// associated with it, and NoPos().IsValid() is false. NoPos is always
// smaller than any other Pos value. The corresponding Position value
// for NoPos is the zero value for Position.
//
const NoPos Pos = 0

// IsValid returns true if the position is valid.
func (p Pos) IsValid() bool {
	return p != NoPos
}

func searchFiles(a []*File, x int) int {
	return sort.Search(len(a), func(i int) bool { return a[i].base > x }) - 1
}

func (s *FileSet) file(p Pos) *File {
	if f := s.last; f != nil && f.base <= int(p) && int(p) <= f.base+f.size {
		return f
	}
	if i := searchFiles(s.files, int(p)); i >= 0 {
		f := s.files[i]
		// f.base <= int(p) by definition of searchFiles
		if int(p) <= f.base+f.size {
			s.last = f
			return f
		}
	}
	return nil
}

// File returns the file which contains the position p.
// If no such file is found (for instance for p == NoPos),
// the result is nil.
//
func (s *FileSet) File(p Pos) (f *File) {
	if p != NoPos {
		s.mutex.RLock()
		f = s.file(p)
		s.mutex.RUnlock()
	}
	return
}

func (f *File) position(p Pos) (pos Position) {
	offset := int(p) - f.base
	pos.Offset = offset
	pos.Filename, pos.Line, pos.Column = f.info(offset)
	return
}

// Position converts a Pos in the fileset into a general Position.
func (s *FileSet) Position(p Pos) (pos Position) {
	if p != NoPos {
		// TODO(gri) consider optimizing the case where p
		//           is in the last file addded, or perhaps
		//           looked at - will eliminate one level
		//           of search
		s.mutex.RLock()
		if f := s.file(p); f != nil {
			pos = f.position(p)
		}
		s.mutex.RUnlock()
	}
	return
}

type lineInfo struct {
	offset   int
	filename string
	line     int
}

// AddLineInfo adds alternative file and line number information for
// a given file offset. The offset must be larger than the offset for
// the previously added alternative line info and smaller than the
// file size; otherwise the information is ignored.
//
// AddLineInfo is typically used to register alternative position
// information for //line filename:line comments in source files.
//
func (f *File) AddLineInfo(offset int, filename string, line int) {
	f.set.mutex.Lock()
	if i := len(f.infos); i == 0 || f.infos[i-1].offset < offset && offset < f.size {
		f.infos = append(f.infos, lineInfo{offset, filename, line})
	}
	f.set.mutex.Unlock()
}

// A File is a handle for a file belonging to a FileSet.
// A File has a name, size, and line offset table.
//
type File struct {
	set  *FileSet
	name string // file name as provided to AddFile
	base int    // Pos value range for this file is [base...base+size]
	size int    // file size as provided to AddFile

	// lines and infos are protected by set.mutex
	lines []int
	infos []lineInfo
}

// Name returns the file name of file f as registered with AddFile.
func (f *File) Name() string {
	return f.name
}

// Base returns the base offset of file f as registered with AddFile.
func (f *File) Base() int {
	return f.base
}

// Size returns the size of file f as registered with AddFile.
func (f *File) Size() int {
	return f.size
}

// LineCount returns the number of lines in file f.
func (f *File) LineCount() int {
	f.set.mutex.RLock()
	n := len(f.lines)
	f.set.mutex.RUnlock()
	return n
}

// AddLine adds the line offset for a new line.
// The line offset must be larger than the offset for the previous line
// and smaller than the file size; otherwise the line offset is ignored.
//
func (f *File) AddLine(offset int) {
	f.set.mutex.Lock()
	if i := len(f.lines); (i == 0 || f.lines[i-1] < offset) && offset < f.size {
		f.lines = append(f.lines, offset)
	}
	f.set.mutex.Unlock()
}

// SetLines sets the line offsets for a file and returns true if successful.
// The line offsets are the offsets of the first character of each line;
// for instance for the content "ab\nc\n" the line offsets are {0, 3}.
// An empty file has an empty line offset table.
// Each line offset must be larger than the offset for the previous line
// and smaller than the file size; otherwise SetLines fails and returns
// false.
//
func (f *File) SetLines(lines []int) bool {
	// verify validity of lines table
	size := f.size
	for i, offset := range lines {
		if i > 0 && offset <= lines[i-1] || size <= offset {
			return false
		}
	}

	// set lines table
	f.set.mutex.Lock()
	f.lines = lines
	f.set.mutex.Unlock()
	return true
}

// SetLinesForContent sets the line offsets for the given file content.
func (f *File) SetLinesForContent(content []byte) {
	var lines []int
	line := 0
	for offset, b := range content {
		if line >= 0 {
			lines = append(lines, line)
		}
		line = -1
		if b == '\n' {
			line = offset + 1
		}
	}

	// set lines table
	f.set.mutex.Lock()
	f.lines = lines
	f.set.mutex.Unlock()
}

// Pos returns the Pos value for the given file offset;
// the offset must be <= f.Size().
// f.Pos(f.Offset(p)) == p.
//
func (f *File) Pos(offset int) Pos {
	if offset > f.size {
		panic("illegal file offset")
	}
	return Pos(f.base + offset)
}

// Offset returns the offset for the given file position p;
// p must be a valid Pos value in that file.
// f.Offset(f.Pos(offset)) == offset.
//
func (f *File) Offset(p Pos) int {
	if int(p) < f.base || int(p) > f.base+f.size {
		panic("illegal Pos value")
	}
	return int(p) - f.base
}

// Line returns the line number for the given file position p;
// p must be a Pos value in that file or NoPos.
//
func (f *File) Line(p Pos) int {
	// TODO(gri) this can be implemented much more efficiently
	return f.Position(p).Line
}

// Position returns the Position value for the given file position p;
// p must be a Pos value in that file or NoPos.
//
func (f *File) Position(p Pos) (pos Position) {
	if p != NoPos {
		if int(p) < f.base || int(p) > f.base+f.size {
			panic("illegal Pos value")
		}
		pos = f.position(p)
	}
	return
}

func searchInts(a []int, x int) int {
	// This function body is a manually inlined version of:
	//
	//   return sort.Search(len(a), func(i int) bool { return a[i] > x }) - 1
	//
	// With better compiler optimizations, this may not be needed in the
	// future, but at the moment this change improves the go/printer
	// benchmark performance by ~30%. This has a direct impact on the
	// speed of gofmt and thus seems worthwhile (2011-04-29).
	i, j := 0, len(a)
	for i < j {
		h := i + (j-i)/2 // avoid overflow when computing h
		// i ≤ h < j
		if a[h] <= x {
			i = h + 1
		} else {
			j = h
		}
	}
	return i - 1
}

func searchLineInfos(a []lineInfo, x int) int {
	return sort.Search(len(a), func(i int) bool { return a[i].offset > x }) - 1
}

// info returns the file name, line, and column number for a file offset.
func (f *File) info(offset int) (filename string, line, column int) {
	filename = f.name
	if i := searchInts(f.lines, offset); i >= 0 {
		line, column = i+1, offset-f.lines[i]+1
	}
	if len(f.infos) > 0 {
		// almost no files have extra line infos
		if i := searchLineInfos(f.infos, offset); i >= 0 {
			alt := &f.infos[i]
			filename = alt.filename
			if i := searchInts(f.lines, alt.offset); i >= 0 {
				line += alt.line - i - 1
			}
		}
	}
	return
}

// A FileSet represents a set of source files.
// Methods of file sets are synchronized; multiple goroutines
// may invoke them concurrently.
//
type FileSet struct {
	mutex sync.RWMutex // protects the file set
	base  int          // base offset for the next file
	files []*File      // list of files in the order added to the set
	last  *File        // cache of last file looked up
}

// NewFileSet creates a new file set.
func NewFileSet() *FileSet {
	s := new(FileSet)
	s.base = 1 // 0 == NoPos
	return s
}

// Base returns the minimum base offset that must be provided to
// AddFile when adding the next file.
//
func (s *FileSet) Base() int {
	s.mutex.RLock()
	b := s.base
	s.mutex.RUnlock()
	return b

}

// AddFile adds a new file with a given filename, base offset, and file size
// to the file set s and returns the file. Multiple files may have the same
// name. The base offset must not be smaller than the FileSet's Base(), and
// size must not be negative.
//
// Adding the file will set the file set's Base() value to base + size + 1
// as the minimum base value for the next file. The following relationship
// exists between a Pos value p for a given file offset offs:
//
//	int(p) = base + offs
//
// with offs in the range [0, size] and thus p in the range [base, base+size].
// For convenience, File.Pos may be used to create file-specific position
// values from a file offset.
//
func (s *FileSet) AddFile(filename string, base, size int) *File {
	s.mutex.Lock()
	defer s.mutex.Unlock()
	if base < s.base || size < 0 {
		panic("illegal base or size")
	}
	// base >= s.base && size >= 0
	f := &File{s, filename, base, size, []int{0}, nil}
	base += size + 1 // +1 because EOF also has a position
	if base < 0 {
		panic("token.Pos offset overflow (> 2G of source code in file set)")
	}
	// add the file to the file set
	s.base = base
	s.files = append(s.files, f)
	s.last = f
	return f
}

// Files returns the files added to the file set.
func (s *FileSet) Files() <-chan *File {
	ch := make(chan *File)
	go func() {
		for i := 0; ; i++ {
			var f *File
			s.mutex.RLock()
			if i < len(s.files) {
				f = s.files[i]
			}
			s.mutex.RUnlock()
			if f == nil {
				break
			}
			ch <- f
		}
		close(ch)
	}()
	return ch
}
<|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/.
//
// Contributor: Zack Mullaly zmullaly@mozilla.com [:zack]

package agents

import (
	"fmt"

	migdb "github.com/mozilla/mig/database"
)

type PersistHeartbeatPostgres struct {
	db *migdb.DB
}

func NewPersistHeartbeatPostgres(db *migdb.DB) PersistHeartbeatPostgres {
	return PersistHeartbeatPostgres{
		db: db,
	}
}

func (persist PersistHeartbeatPostgres) PersistHeartbeat(heartbeat Heartbeat) error {
	fmt.Printf("POST /heartbeat got heartbeat %v\n", heartbeat)

	agent := heartbeat.ToMigAgent()
	err := persist.db.InsertAgent(agent, nil)

	return err
}

// _dontrun invokes a goroutine that updates the agent table when a heartbeat
// message would have been handled by the scheduler. For now we're holding onto
// the code as a reference
func _dontrun() {
	go func() {
		// if an agent already exists in database, we update it, otherwise we insert it
		agent, err := ctx.DB.AgentByQueueAndPID(agt.QueueLoc, agt.PID)
		if err != nil {
			agt.DestructionTime = time.Date(9998, time.January, 11, 11, 11, 11, 11, time.UTC)
			agt.Status = mig.AgtStatusOnline
			// create a new agent, set starttime to now
			agt.StartTime = time.Now()
			err = ctx.DB.InsertAgent(agt, nil)
			if err != nil {
				ctx.Channels.Log <- mig.Log{Desc: fmt.Sprintf("Heartbeat DB insertion failed with error '%v' for agent '%s'", err, agt.Name)}.Err()
			}
		} else {
			// the agent exists in database. reuse the existing ID, and keep the status if it was
			// previously set to destroyed, otherwise set status to online
			agt.ID = agent.ID
			if agt.Status == mig.AgtStatusDestroyed {
				agt.Status = agent.Status
			} else {
				agt.Status = mig.AgtStatusOnline
			}
			// If the refresh time is newer than what we know for the agent, replace
			// the agent in the database with the newer information. We want to keep
			// history here, so don't want to just update the information in the
			// existing row.
			//
			// Note: with older agents which might not send a refresh time, the refresh
			// time will be interpreted as the zero value, and the agents should just
			// update using UpdateAgentHeartbeat()
			if agt.RefreshTS.IsZero() {
				ctx.Channels.Log <- mig.Log{Desc: fmt.Sprintf("agent '%v' not sending refresh time, perhaps an older version?", agt.Name)}.Warning()
			}
			cutoff := agent.RefreshTS.Add(15 * time.Second)
			if !agt.RefreshTS.IsZero() && agt.RefreshTS.After(cutoff) {
				ctx.Channels.Log <- mig.Log{Desc: fmt.Sprintf("replacing refreshed agent for agent '%v'", agt.Name)}.Info()
				err = ctx.DB.ReplaceRefreshedAgent(agt)
				if err != nil {
					ctx.Channels.Log <- mig.Log{Desc: fmt.Sprintf("Heartbeat DB update failed (refresh) with error '%v' for agent '%s'", err, agt.Name)}.Err()
				}
			} else {
				err = ctx.DB.UpdateAgentHeartbeat(agt)
				if err != nil {
					ctx.Channels.Log <- mig.Log{Desc: fmt.Sprintf("Heartbeat DB update failed with error '%v' for agent '%s'", err, agt.Name)}.Err()
				}
			}
			// if the agent that exists in the database has a status of 'destroyed'
			// we should not be received a heartbeat from it. so, if detectmultiagents
			// is set in the scheduler configuration, we pass the agent queue over to the
			// routine than handles the destruction of agents
			if agent.Status == mig.AgtStatusDestroyed && ctx.Agent.DetectMultiAgents {
				ctx.Channels.DetectDupAgents <- agent.QueueLoc
			}
		}
	}()
}
<commit_msg>Update agent information in place when a heartbeat is received<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/.
//
// Contributor: Zack Mullaly zmullaly@mozilla.com [:zack]

package agents

import (
	"fmt"

	migdb "github.com/mozilla/mig/database"
)

type PersistHeartbeatPostgres struct {
	db *migdb.DB
}

func NewPersistHeartbeatPostgres(db *migdb.DB) PersistHeartbeatPostgres {
	return PersistHeartbeatPostgres{
		db: db,
	}
}

func (persist PersistHeartbeatPostgres) PersistHeartbeat(heartbeat Heartbeat) error {
	fmt.Printf("POST /heartbeat got heartbeat %v\n", heartbeat)

	agent := heartbeat.ToMigAgent()
	//err := persist.db.InsertAgent(agent, nil)
	agent, err := persist.db.AgentByQueueAndPID(agent.QueueLoc, agent.PID)
	if err != nil {
		return err
	}

	return persist.db.UpdateAgentHeartbeat(agent)
}

// _dontrun invokes a goroutine that updates the agent table when a heartbeat
// message would have been handled by the scheduler. For now we're holding onto
// the code as a reference
func _dontrun() {
	go func() {
		// if an agent already exists in database, we update it, otherwise we insert it
		agent, err := ctx.DB.AgentByQueueAndPID(agt.QueueLoc, agt.PID)
		if err != nil {
			agt.DestructionTime = time.Date(9998, time.January, 11, 11, 11, 11, 11, time.UTC)
			agt.Status = mig.AgtStatusOnline
			// create a new agent, set starttime to now
			agt.StartTime = time.Now()
			err = ctx.DB.InsertAgent(agt, nil)
			if err != nil {
				ctx.Channels.Log <- mig.Log{Desc: fmt.Sprintf("Heartbeat DB insertion failed with error '%v' for agent '%s'", err, agt.Name)}.Err()
			}
		} else {
			// the agent exists in database. reuse the existing ID, and keep the status if it was
			// previously set to destroyed, otherwise set status to online
			agt.ID = agent.ID
			if agt.Status == mig.AgtStatusDestroyed {
				agt.Status = agent.Status
			} else {
				agt.Status = mig.AgtStatusOnline
			}
			// If the refresh time is newer than what we know for the agent, replace
			// the agent in the database with the newer information. We want to keep
			// history here, so don't want to just update the information in the
			// existing row.
			//
			// Note: with older agents which might not send a refresh time, the refresh
			// time will be interpreted as the zero value, and the agents should just
			// update using UpdateAgentHeartbeat()
			if agt.RefreshTS.IsZero() {
				ctx.Channels.Log <- mig.Log{Desc: fmt.Sprintf("agent '%v' not sending refresh time, perhaps an older version?", agt.Name)}.Warning()
			}
			cutoff := agent.RefreshTS.Add(15 * time.Second)
			if !agt.RefreshTS.IsZero() && agt.RefreshTS.After(cutoff) {
				ctx.Channels.Log <- mig.Log{Desc: fmt.Sprintf("replacing refreshed agent for agent '%v'", agt.Name)}.Info()
				err = ctx.DB.ReplaceRefreshedAgent(agt)
				if err != nil {
					ctx.Channels.Log <- mig.Log{Desc: fmt.Sprintf("Heartbeat DB update failed (refresh) with error '%v' for agent '%s'", err, agt.Name)}.Err()
				}
			} else {
				err = ctx.DB.UpdateAgentHeartbeat(agt)
				if err != nil {
					ctx.Channels.Log <- mig.Log{Desc: fmt.Sprintf("Heartbeat DB update failed with error '%v' for agent '%s'", err, agt.Name)}.Err()
				}
			}
			// if the agent that exists in the database has a status of 'destroyed'
			// we should not be received a heartbeat from it. so, if detectmultiagents
			// is set in the scheduler configuration, we pass the agent queue over to the
			// routine than handles the destruction of agents
			if agent.Status == mig.AgtStatusDestroyed && ctx.Agent.DetectMultiAgents {
				ctx.Channels.DetectDupAgents <- agent.QueueLoc
			}
		}
	}()
}
<|endoftext|>


<commit_before>// Copyright 2018 ThoughtWorks, Inc.

// This file is part of Gauge.

// Gauge is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.

// Gauge is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with Gauge.  If not, see <http://www.gnu.org/licenses/>.

package runner

import (
	"context"
	"fmt"
	"io"
	"net"
	"os/exec"

	gm "github.com/getgauge/gauge/gauge_messages"
	"github.com/getgauge/gauge/manifest"
	"google.golang.org/grpc"
)

// GrpcRunner handles grpc messages.
type GrpcRunner struct {
	cmd    *exec.Cmd
	conn   *grpc.ClientConn
	Client gm.LspServiceClient
}

// ExecuteMessageWithTimeout process reuqest and give back the response
func (r *GrpcRunner) ExecuteMessageWithTimeout(message *gm.Message) (*gm.Message, error) {
	switch message.MessageType {
	case gm.Message_CacheFileRequest:
		r.Client.CacheFile(context.Background(), message.CacheFileRequest)
		return &gm.Message{}, nil
	case gm.Message_StepNamesRequest:
		response, err := r.Client.GetStepNames(context.Background(), message.StepNamesRequest)
		return &gm.Message{StepNamesResponse: response}, err
	case gm.Message_StepPositionsRequest:
		response, err := r.Client.GetStepPositions(context.Background(), message.StepPositionsRequest)
		return &gm.Message{StepPositionsResponse: response}, err
	case gm.Message_ImplementationFileListRequest:
		response, err := r.Client.GetImplementationFiles(context.Background(), &gm.Empty{})
		return &gm.Message{ImplementationFileListResponse: response}, err
	case gm.Message_StubImplementationCodeRequest:
		response, err := r.Client.ImplementStub(context.Background(), message.StubImplementationCodeRequest)
		return &gm.Message{FileDiff: response}, err
	case gm.Message_StepValidateRequest:
		response, err := r.Client.ValidateStep(context.Background(), message.StepValidateRequest)
		return &gm.Message{MessageType: gm.Message_StepValidateResponse, StepValidateResponse: response}, err
	case gm.Message_RefactorRequest:
		response, err := r.Client.Refactor(context.Background(), message.RefactorRequest)
		return &gm.Message{MessageType: gm.Message_RefactorResponse, RefactorResponse: response}, err
	case gm.Message_StepNameRequest:
		response, err := r.Client.GetStepName(context.Background(), message.StepNameRequest)
		return &gm.Message{MessageType: gm.Message_StepNameResponse, StepNameResponse: response}, err
	case gm.Message_ImplementationFileGlobPatternRequest:
		response, err := r.Client.GetGlobPatterns(context.Background(), &gm.Empty{})
		return &gm.Message{MessageType: gm.Message_ImplementationFileGlobPatternRequest, ImplementationFileGlobPatternResponse: response}, err
	default:
		return nil, nil
	}
}

func (r *GrpcRunner) ExecuteAndGetStatus(m *gm.Message) *gm.ProtoExecutionResult {
	return nil
}
func (r *GrpcRunner) IsProcessRunning() bool {
	return false
}

// Kill closes the grpc connection and kills the process
func (r *GrpcRunner) Kill() error {
	r.Client.KillProcess(context.Background(), &gm.KillProcessRequest{})
	if err := r.conn.Close(); err != nil {
		return err
	}
	// TODO: wait for process to exit or kill forcefully after runner kill timeout
	if err := r.cmd.Process.Kill(); err != nil {
		return err
	}
	return nil
}

func (r *GrpcRunner) Connection() net.Conn {
	return nil
}

func (r *GrpcRunner) IsMultithreaded() bool {
	return false
}

func (r *GrpcRunner) Pid() int {
	return 0
}

// ConnectToGrpcRunner makes a connection with grpc server
func ConnectToGrpcRunner(manifest *manifest.Manifest, outFile io.Writer) (*GrpcRunner, error) {
	// TODO: Remove hardcoded port
	cmd, _, err := runRunnerCommand(manifest, "54545", false, outFile)
	if err != nil {
		return nil, err
	}
	conn, err := grpc.Dial(fmt.Sprintf("127.0.0.1:54545"), grpc.WithInsecure(), grpc.WithBlock())
	if err != nil {
		return nil, err
	}
	return &GrpcRunner{Client: gm.NewLspServiceClient(conn), cmd: cmd, conn: conn}, nil
}
<commit_msg>Removed hard coded port, extracting port from runner's stdout.<commit_after>// Copyright 2018 ThoughtWorks, Inc.

// This file is part of Gauge.

// Gauge is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.

// Gauge is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with Gauge.  If not, see <http://www.gnu.org/licenses/>.

package runner

import (
	"context"
	"fmt"
	"io"
	"net"
	"os/exec"
	"strings"

	gm "github.com/getgauge/gauge/gauge_messages"
	"github.com/getgauge/gauge/manifest"
	"google.golang.org/grpc"
)

const (
	portPrefix = "Listening on port:"
	host       = "127.0.0.1"
)

// GrpcRunner handles grpc messages.
type GrpcRunner struct {
	cmd    *exec.Cmd
	conn   *grpc.ClientConn
	Client gm.LspServiceClient
}

// ExecuteMessageWithTimeout process reuqest and give back the response
func (r *GrpcRunner) ExecuteMessageWithTimeout(message *gm.Message) (*gm.Message, error) {
	switch message.MessageType {
	case gm.Message_CacheFileRequest:
		r.Client.CacheFile(context.Background(), message.CacheFileRequest)
		return &gm.Message{}, nil
	case gm.Message_StepNamesRequest:
		response, err := r.Client.GetStepNames(context.Background(), message.StepNamesRequest)
		return &gm.Message{StepNamesResponse: response}, err
	case gm.Message_StepPositionsRequest:
		response, err := r.Client.GetStepPositions(context.Background(), message.StepPositionsRequest)
		return &gm.Message{StepPositionsResponse: response}, err
	case gm.Message_ImplementationFileListRequest:
		response, err := r.Client.GetImplementationFiles(context.Background(), &gm.Empty{})
		return &gm.Message{ImplementationFileListResponse: response}, err
	case gm.Message_StubImplementationCodeRequest:
		response, err := r.Client.ImplementStub(context.Background(), message.StubImplementationCodeRequest)
		return &gm.Message{FileDiff: response}, err
	case gm.Message_StepValidateRequest:
		response, err := r.Client.ValidateStep(context.Background(), message.StepValidateRequest)
		return &gm.Message{MessageType: gm.Message_StepValidateResponse, StepValidateResponse: response}, err
	case gm.Message_RefactorRequest:
		response, err := r.Client.Refactor(context.Background(), message.RefactorRequest)
		return &gm.Message{MessageType: gm.Message_RefactorResponse, RefactorResponse: response}, err
	case gm.Message_StepNameRequest:
		response, err := r.Client.GetStepName(context.Background(), message.StepNameRequest)
		return &gm.Message{MessageType: gm.Message_StepNameResponse, StepNameResponse: response}, err
	case gm.Message_ImplementationFileGlobPatternRequest:
		response, err := r.Client.GetGlobPatterns(context.Background(), &gm.Empty{})
		return &gm.Message{MessageType: gm.Message_ImplementationFileGlobPatternRequest, ImplementationFileGlobPatternResponse: response}, err
	default:
		return nil, nil
	}
}

func (r *GrpcRunner) ExecuteAndGetStatus(m *gm.Message) *gm.ProtoExecutionResult {
	return nil
}
func (r *GrpcRunner) IsProcessRunning() bool {
	return false
}

// Kill closes the grpc connection and kills the process
func (r *GrpcRunner) Kill() error {
	r.Client.KillProcess(context.Background(), &gm.KillProcessRequest{})
	if err := r.conn.Close(); err != nil {
		return err
	}
	// TODO: wait for process to exit or kill forcefully after runner kill timeout
	if err := r.cmd.Process.Kill(); err != nil {
		return err
	}
	return nil
}

func (r *GrpcRunner) Connection() net.Conn {
	return nil
}

func (r *GrpcRunner) IsMultithreaded() bool {
	return false
}

func (r *GrpcRunner) Pid() int {
	return 0
}

type customWriter struct {
	file io.Writer
	port chan string
}

func (w customWriter) Write(p []byte) (n int, err error) {
	if strings.Contains(string(p), portPrefix) {
		w.port <- strings.TrimSuffix(strings.Split(string(p), portPrefix)[1], "\n")
	}
	return w.file.Write(p)
}

// ConnectToGrpcRunner makes a connection with grpc server
func ConnectToGrpcRunner(manifest *manifest.Manifest, outFile io.Writer) (*GrpcRunner, error) {
	portChan := make(chan string)
	cmd, _, err := runRunnerCommand(manifest, "0", false, customWriter{file: outFile, port: portChan})
	if err != nil {
		return nil, err
	}
	port := <-portChan
	close(portChan)
	conn, err := grpc.Dial(fmt.Sprintf("%s:%s", host, port), grpc.WithInsecure(), grpc.WithBlock())
	if err != nil {
		return nil, err
	}
	return &GrpcRunner{Client: gm.NewLspServiceClient(conn), cmd: cmd, conn: conn}, nil
}
<|endoftext|>


<commit_before>package jsonutil

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
)

var (
	MarshalPrefix = ""
	MarshalIndent = "    "
)

type mustMarhshalError struct {
	MustMarhshalError string `json:"must_marshal_error"`
}

func MustMarshal(i interface{}, embedError bool) []byte {
	bytes, err := json.Marshal(i)
	if err != nil {
		if embedError {
			e := mustMarhshalError{
				MustMarhshalError: err.Error(),
			}
			bytes, err := json.Marshal(e)
			if err != nil {
				panic(err)
			}
			return bytes
		}
		panic(err)
	}
	return bytes
}

func MustMarshalString(i interface{}, embedError bool) string {
	return string(MustMarshal(i, embedError))
}

func MustMarshalIndent(i interface{}, prefix, indent string, embedError bool) []byte {
	bytes, err := json.MarshalIndent(i, prefix, indent)
	if err != nil {
		panic(err)
	}
	return bytes
}

func PrettyPrint(b []byte) ([]byte, error) {
	var out bytes.Buffer
	err := json.Indent(&out, b, MarshalPrefix, MarshalIndent)
	return out.Bytes(), err
}

func MarshalBase64(i interface{}) (string, error) {
	data, err := json.Marshal(i)
	if err != nil {
		return "", err
	}
	return base64.StdEncoding.EncodeToString(data), nil
}
<commit_msg>add jsonutil.UnmarshalIoReader<commit_after>package jsonutil

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"io"
	"io/ioutil"
)

var (
	MarshalPrefix = ""
	MarshalIndent = "    "
)

type mustMarhshalError struct {
	MustMarhshalError string `json:"must_marshal_error"`
}

func MustMarshal(i interface{}, embedError bool) []byte {
	bytes, err := json.Marshal(i)
	if err != nil {
		if embedError {
			e := mustMarhshalError{
				MustMarhshalError: err.Error(),
			}
			bytes, err := json.Marshal(e)
			if err != nil {
				panic(err)
			}
			return bytes
		}
		panic(err)
	}
	return bytes
}

func MustMarshalString(i interface{}, embedError bool) string {
	return string(MustMarshal(i, embedError))
}

func MustMarshalIndent(i interface{}, prefix, indent string, embedError bool) []byte {
	bytes, err := json.MarshalIndent(i, prefix, indent)
	if err != nil {
		panic(err)
	}
	return bytes
}

func UnmarshalIoReader(r io.Reader, data interface{}) error {
	b, err := ioutil.ReadAll(r)
	if err != nil {
		return err
	}
	return json.Unmarshal(b, data)
}

func PrettyPrint(b []byte) ([]byte, error) {
	var out bytes.Buffer
	err := json.Indent(&out, b, MarshalPrefix, MarshalIndent)
	return out.Bytes(), err
}

func MarshalBase64(i interface{}) (string, error) {
	data, err := json.Marshal(i)
	if err != nil {
		return "", err
	}
	return base64.StdEncoding.EncodeToString(data), nil
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"strconv"
	"strings"

	log "github.com/sirupsen/logrus"
)

// CSPReport is the structure of the HTTP payload the system receives.
type CSPReport struct {
	Body CSPReportBody `json:"csp-report"`
}

// CSPReportBody contains the fields that are nested within the
// violation report.
type CSPReportBody struct {
	DocumentURI        string      `json:"document-uri"`
	Referrer           string      `json:"referrer"`
	BlockedURI         string      `json:"blocked-uri"`
	ViolatedDirective  string      `json:"violated-directive"`
	EffectiveDirective string      `json:"effective-directive"`
	OriginalPolicy     string      `json:"original-policy"`
	Disposition        string      `json:"disposition"`
	ScriptSample       string      `json:"script-sample"`
	StatusCode         interface{} `json:"status-code"`
}

var (
	// Rev is set at build time and holds the revision that the package
	// was created at.
	Rev = "dev"

	// Flag for toggling verbose output.
	debugFlag bool

	// Flag for toggling output format.
	outputFormat string

	// Flag for health check url.
	healthCheckPath = "/_healthcheck"

	// Shared defaults for the logger output. This ensures that we are
	// using the same keys for the `FieldKey` values across both formatters.
	logFieldMapDefaults = log.FieldMap{
		log.FieldKeyTime:  "timestamp",
		log.FieldKeyLevel: "level",
		log.FieldKeyMsg:   "message",
	}

	// Path to file which has blocked URI's per line.
	blockedURIfile string

	// Default URI Filter list.
	ignoredBlockedURIs = []string{
		"resource://",
		"chromenull://",
		"chrome-extension://",
		"safari-extension://",
		"mxjscall://",
		"webviewprogressproxy://",
		"res://",
		"mx://",
		"safari-resource://",
		"chromeinvoke://",
		"chromeinvokeimmediate://",
		"mbinit://",
		"opera://",
		"ms-appx://",
		"ms-appx-web://",
		"localhost",
		"127.0.0.1",
		"none://",
		"about:blank",
		"android-webview",
		"ms-browser-extension",
		"wvjbscheme://__wvjb_queue_message__",
		"nativebaiduhd://adblock",
		"bdvideo://error",
	}

	// TCP Port to listen on.
	listenPort int
)

func init() {
	log.SetOutput(os.Stdout)
	log.SetLevel(log.InfoLevel)
}

func trimEmptyAndComments(s []string) []string {
	var r []string
	for _, str := range s {
		if str == "" {
			continue
		}

		// ignore comments
		if strings.HasPrefix(str, "#") {
			continue
		}

		r = append(r, str)
	}
	return r
}

func main() {
	version := flag.Bool("version", false, "Display the version")
	flag.BoolVar(&debugFlag, "debug", false, "Output additional logging for debugging")
	flag.StringVar(&outputFormat, "output-format", "text", "Define how the violation reports are formatted for output.\nDefaults to 'text'. Valid options are 'text' or 'json'")
	flag.StringVar(&blockedURIfile, "filter-file", "", "Blocked URI Filter file")
	flag.IntVar(&listenPort, "port", 8080, "Port to listen on")
	flag.StringVar(&healthCheckPath, "health-check-path", healthCheckPath, "Health checker path")

	flag.Parse()

	if *version {
		fmt.Printf("csp-collector (%s)\n", Rev)
		os.Exit(0)
	}

	if blockedURIfile != "" {
		content, err := ioutil.ReadFile(blockedURIfile)
		if err != nil {
			fmt.Printf("Error reading Blocked File list: %s", blockedURIfile)
		}
		ignoredBlockedURIs = trimEmptyAndComments(strings.Split(string(content), "\n"))
	}

	if debugFlag {
		log.SetLevel(log.DebugLevel)
	}

	if outputFormat == "json" {
		log.SetFormatter(&log.JSONFormatter{
			FieldMap: logFieldMapDefaults,
		})
	} else {
		log.SetFormatter(&log.TextFormatter{
			FullTimestamp:          true,
			DisableLevelTruncation: true,
			QuoteEmptyFields:       true,
			DisableColors:          true,
			FieldMap:               logFieldMapDefaults,
		})
	}

	log.Debug("Starting up...")
	if blockedURIfile != "" {
		log.Debugf("Using Filter list from file at: %s\n", blockedURIfile)
	} else {
		log.Debug("Using Filter list from internal list")
	}
	log.Debugf("Blocked URI List: %s", ignoredBlockedURIs)
	log.Debugf("Listening on TCP Port: %s", strconv.Itoa(listenPort))

	http.HandleFunc("/", handleViolationReport)
	log.Fatal(http.ListenAndServe(fmt.Sprintf(":%s", strconv.Itoa(listenPort)), nil))
}

func handleViolationReport(w http.ResponseWriter, r *http.Request) {
	if r.Method == "GET" && r.URL.Path == healthCheckPath {
		w.WriteHeader(http.StatusOK)
		return
	}

	if r.Method != "POST" {
		w.WriteHeader(http.StatusMethodNotAllowed)
		log.WithFields(log.Fields{
			"http_method": r.Method,
		}).Debug("Received invalid HTTP method")
		return
	}

	decoder := json.NewDecoder(r.Body)
	var report CSPReport

	err := decoder.Decode(&report)
	if err != nil {
		w.WriteHeader(http.StatusUnprocessableEntity)
		log.Debug(fmt.Sprintf("Unable to decode invalid JSON payload: %s", err))
		return
	}
	defer r.Body.Close()

	reportValidation := validateViolation(report)
	if reportValidation != nil {
		http.Error(w, reportValidation.Error(), http.StatusBadRequest)
		log.Debug(fmt.Sprintf("Received invalid payload: %s", reportValidation.Error()))
		return
	}

	metadatas, gotMetadata := r.URL.Query()["metadata"]
	var metadata string
	if gotMetadata {
		metadata = metadatas[0]
	}

	log.WithFields(log.Fields{
		"document_uri":        report.Body.DocumentURI,
		"referrer":            report.Body.Referrer,
		"blocked_uri":         report.Body.BlockedURI,
		"violated_directive":  report.Body.ViolatedDirective,
		"effective_directive": report.Body.EffectiveDirective,
		"original_policy":     report.Body.OriginalPolicy,
		"disposition":         report.Body.Disposition,
		"script_sample":       report.Body.ScriptSample,
		"status_code":         report.Body.StatusCode,
		"metadata":            metadata,
	}).Info()
}

func validateViolation(r CSPReport) error {
	for _, value := range ignoredBlockedURIs {
		if strings.HasPrefix(r.Body.BlockedURI, value) {
			err := fmt.Errorf("blocked URI ('%s') is an invalid resource", value)
			return err
		}
	}

	if !strings.HasPrefix(r.Body.DocumentURI, "http") {
		return fmt.Errorf("document URI ('%s') is invalid", r.Body.DocumentURI)
	}

	return nil
}
<commit_msg>Use log library more consistently<commit_after>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"strconv"
	"strings"

	log "github.com/sirupsen/logrus"
)

// CSPReport is the structure of the HTTP payload the system receives.
type CSPReport struct {
	Body CSPReportBody `json:"csp-report"`
}

// CSPReportBody contains the fields that are nested within the
// violation report.
type CSPReportBody struct {
	DocumentURI        string      `json:"document-uri"`
	Referrer           string      `json:"referrer"`
	BlockedURI         string      `json:"blocked-uri"`
	ViolatedDirective  string      `json:"violated-directive"`
	EffectiveDirective string      `json:"effective-directive"`
	OriginalPolicy     string      `json:"original-policy"`
	Disposition        string      `json:"disposition"`
	ScriptSample       string      `json:"script-sample"`
	StatusCode         interface{} `json:"status-code"`
}

var (
	// Rev is set at build time and holds the revision that the package
	// was created at.
	Rev = "dev"

	// Flag for toggling verbose output.
	debugFlag bool

	// Flag for toggling output format.
	outputFormat string

	// Flag for health check url.
	healthCheckPath = "/_healthcheck"

	// Shared defaults for the logger output. This ensures that we are
	// using the same keys for the `FieldKey` values across both formatters.
	logFieldMapDefaults = log.FieldMap{
		log.FieldKeyTime:  "timestamp",
		log.FieldKeyLevel: "level",
		log.FieldKeyMsg:   "message",
	}

	// Path to file which has blocked URI's per line.
	blockedURIfile string

	// Default URI Filter list.
	ignoredBlockedURIs = []string{
		"resource://",
		"chromenull://",
		"chrome-extension://",
		"safari-extension://",
		"mxjscall://",
		"webviewprogressproxy://",
		"res://",
		"mx://",
		"safari-resource://",
		"chromeinvoke://",
		"chromeinvokeimmediate://",
		"mbinit://",
		"opera://",
		"ms-appx://",
		"ms-appx-web://",
		"localhost",
		"127.0.0.1",
		"none://",
		"about:blank",
		"android-webview",
		"ms-browser-extension",
		"wvjbscheme://__wvjb_queue_message__",
		"nativebaiduhd://adblock",
		"bdvideo://error",
	}

	// TCP Port to listen on.
	listenPort int
)

func init() {
	log.SetOutput(os.Stdout)
	log.SetLevel(log.InfoLevel)
}

func trimEmptyAndComments(s []string) []string {
	var r []string
	for _, str := range s {
		if str == "" {
			continue
		}

		// ignore comments
		if strings.HasPrefix(str, "#") {
			continue
		}

		r = append(r, str)
	}
	return r
}

func main() {
	version := flag.Bool("version", false, "Display the version")
	flag.BoolVar(&debugFlag, "debug", false, "Output additional logging for debugging")
	flag.StringVar(&outputFormat, "output-format", "text", "Define how the violation reports are formatted for output.\nDefaults to 'text'. Valid options are 'text' or 'json'")
	flag.StringVar(&blockedURIfile, "filter-file", "", "Blocked URI Filter file")
	flag.IntVar(&listenPort, "port", 8080, "Port to listen on")
	flag.StringVar(&healthCheckPath, "health-check-path", healthCheckPath, "Health checker path")

	flag.Parse()

	if *version {
		fmt.Printf("csp-collector (%s)\n", Rev)
		os.Exit(0)
	}

	if debugFlag {
		log.SetLevel(log.DebugLevel)
	}

	if outputFormat == "json" {
		log.SetFormatter(&log.JSONFormatter{
			FieldMap: logFieldMapDefaults,
		})
	} else {
		log.SetFormatter(&log.TextFormatter{
			FullTimestamp:          true,
			DisableLevelTruncation: true,
			QuoteEmptyFields:       true,
			DisableColors:          true,
			FieldMap:               logFieldMapDefaults,
		})
	}

	log.Debug("Starting up...")
	if blockedURIfile != "" {
		log.Debugf("Using Filter list from file at: %s\n", blockedURIfile)

		content, err := ioutil.ReadFile(blockedURIfile)
		if err != nil {
			log.Fatalf("Error reading Blocked File list: %s", blockedURIfile)
		}
		ignoredBlockedURIs = trimEmptyAndComments(strings.Split(string(content), "\n"))
	} else {
		log.Debug("Using Filter list from internal list")
	}

	log.Debugf("Blocked URI List: %s", ignoredBlockedURIs)
	log.Debugf("Listening on TCP Port: %s", strconv.Itoa(listenPort))

	http.HandleFunc("/", handleViolationReport)
	log.Fatal(http.ListenAndServe(fmt.Sprintf(":%s", strconv.Itoa(listenPort)), nil))
}

func handleViolationReport(w http.ResponseWriter, r *http.Request) {
	if r.Method == "GET" && r.URL.Path == healthCheckPath {
		w.WriteHeader(http.StatusOK)
		return
	}

	if r.Method != "POST" {
		w.WriteHeader(http.StatusMethodNotAllowed)
		log.WithFields(log.Fields{
			"http_method": r.Method,
		}).Debug("Received invalid HTTP method")
		return
	}

	decoder := json.NewDecoder(r.Body)
	var report CSPReport

	err := decoder.Decode(&report)
	if err != nil {
		w.WriteHeader(http.StatusUnprocessableEntity)
		log.Debugf("Unable to decode invalid JSON payload: %s", err)
		return
	}
	defer r.Body.Close()

	reportValidation := validateViolation(report)
	if reportValidation != nil {
		http.Error(w, reportValidation.Error(), http.StatusBadRequest)
		log.Debugf("Received invalid payload: %s", reportValidation.Error())
		return
	}

	metadatas, gotMetadata := r.URL.Query()["metadata"]
	var metadata string
	if gotMetadata {
		metadata = metadatas[0]
	}

	log.WithFields(log.Fields{
		"document_uri":        report.Body.DocumentURI,
		"referrer":            report.Body.Referrer,
		"blocked_uri":         report.Body.BlockedURI,
		"violated_directive":  report.Body.ViolatedDirective,
		"effective_directive": report.Body.EffectiveDirective,
		"original_policy":     report.Body.OriginalPolicy,
		"disposition":         report.Body.Disposition,
		"script_sample":       report.Body.ScriptSample,
		"status_code":         report.Body.StatusCode,
		"metadata":            metadata,
	}).Info()
}

func validateViolation(r CSPReport) error {
	for _, value := range ignoredBlockedURIs {
		if strings.HasPrefix(r.Body.BlockedURI, value) {
			err := fmt.Errorf("blocked URI ('%s') is an invalid resource", value)
			return err
		}
	}

	if !strings.HasPrefix(r.Body.DocumentURI, "http") {
		return fmt.Errorf("document URI ('%s') is invalid", r.Body.DocumentURI)
	}

	return nil
}
<|endoftext|>