<commit_before>package globaldns

import (
	"context"
	"fmt"
	"strings"

	"github.com/rancher/rancher/pkg/namespace"
	v1Rancher "github.com/rancher/types/apis/core/v1"
	v1beta1Rancher "github.com/rancher/types/apis/extensions/v1beta1"
	"github.com/rancher/types/apis/management.cattle.io/v3"
	projectv3 "github.com/rancher/types/apis/project.cattle.io/v3"
	"github.com/rancher/types/config"
	"github.com/sirupsen/logrus"
	k8serrors "k8s.io/apimachinery/pkg/api/errors"

	"k8s.io/api/extensions/v1beta1"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/runtime"
)

const (
	annotationGlobalDNS         = "rancher.io/globalDNS.hostname"
	appSelectorLabel            = "io.cattle.field/appId"
	projectSelectorLabel        = "field.cattle.io/projectId"
	UserIngressControllerName   = "globaldns-useringress-controller"
	UserGlobalDNSControllerName = "user-controller-watching-globaldns"
)

type UserIngressController struct {
	ingressLister         v1beta1Rancher.IngressLister
	globalDNSs            v3.GlobalDNSInterface
	globalDNSLister       v3.GlobalDNSLister
	appLister             projectv3.AppLister
	namespaceLister       v1Rancher.NamespaceLister
	multiclusterappLister v3.MultiClusterAppLister
	clusterName           string
}

func newUserIngressController(ctx context.Context, clusterContext *config.UserContext) *UserIngressController {
	n := &UserIngressController{
		ingressLister:         clusterContext.Extensions.Ingresses("").Controller().Lister(),
		globalDNSs:            clusterContext.Management.Management.GlobalDNSs(""),
		globalDNSLister:       clusterContext.Management.Management.GlobalDNSs("").Controller().Lister(),
		appLister:             clusterContext.Management.Project.Apps("").Controller().Lister(),
		namespaceLister:       clusterContext.Core.Namespaces("").Controller().Lister(),
		multiclusterappLister: clusterContext.Management.Management.MultiClusterApps("").Controller().Lister(),
		clusterName:           clusterContext.ClusterName,
	}
	return n
}

func Register(ctx context.Context, clusterContext *config.UserContext) {
	n := newUserIngressController(ctx, clusterContext)
	clusterContext.Extensions.Ingresses("").AddHandler(ctx, UserIngressControllerName, n.sync)
	g := newUserGlobalDNSController(clusterContext)
	clusterContext.Management.Management.GlobalDNSs("").AddHandler(ctx, UserGlobalDNSControllerName, g.sync)
}

func (ic *UserIngressController) sync(key string, obj *v1beta1.Ingress) (runtime.Object, error) {
	if obj == nil {
		return ic.reconcileAllGlobalDNSs()
	}
	//if there are no globaldns cr, skip this run. If error, return error

	noGlobalDNS, err := ic.noGlobalDNS()
	if err != nil {
		return nil, err
	}

	if noGlobalDNS {
		return nil, nil
	}

	annotations := obj.Annotations
	//look for globalDNS annotation, if found load the GlobalDNS if there are Ingress endpoints
	if annotations[annotationGlobalDNS] != "" && len(obj.Status.LoadBalancer.Ingress) > 0 {
		fqdnRequested := annotations[annotationGlobalDNS]
		//check if the corresponding GlobalDNS CR is present
		globalDNS, err := ic.findGlobalDNS(fqdnRequested)

		if globalDNS == nil || (err != nil && k8serrors.IsNotFound(err)) {
			return nil, nil
		}

		if err != nil {
			return nil, fmt.Errorf("UserIngressController: Error %v while finding  GlobalDNS resource for FQDN requested %v", err, fqdnRequested)
		}

		//check if 'multiclusterappID' or targetProjects on GlobalDNS CR matches the ingress's app/project
		if obj.DeletionTimestamp == nil {
			var checkPassed bool
			if globalDNS.Spec.MultiClusterAppName != "" {
				checkPassed, err = ic.checkForMultiClusterApp(obj, globalDNS)
			} else if len(globalDNS.Spec.ProjectNames) > 0 {
				//check if 'projectNames' on GlobalDNS CR matches to the user's project for multiclusterapp
				checkPassed, err = ic.checkForProjects(obj, globalDNS)
			}
			if err != nil {
				return nil, err
			}
			if !checkPassed {
				logrus.Debugf("UserIngressController: Not enqueing update on globaldns %v, for ingress key %v", globalDNS, key)
				return nil, nil
			}
		}
		//update endpoints on GlobalDNS status field by enqueueing update on GlobalDNS
		ic.globalDNSs.Controller().Enqueue(namespace.GlobalNamespace, globalDNS.Name)

		return nil, err
	}
	return nil, nil
}

func (ic *UserIngressController) noGlobalDNS() (bool, error) {
	gd, err := ic.globalDNSLister.List(namespace.GlobalNamespace, labels.NewSelector())
	if err != nil {
		return true, err
	}

	return len(gd) == 0, nil
}

func (ic *UserIngressController) reconcileAllGlobalDNSs() (runtime.Object, error) {
	globalDNSList, err := ic.globalDNSLister.List(namespace.GlobalNamespace, labels.NewSelector())
	if err != nil {
		return nil, err
	}

	for _, globalDNSObj := range globalDNSList {
		//call update on each GlobalDNS obj that refers to this current cluster
		targetsCluster, err := ic.doesGlobalDNSTargetCurrentCluster(globalDNSObj)
		if err != nil {
			return nil, err
		}
		if targetsCluster {
			//enqueue it to the globalDNS controller
			ic.globalDNSs.Controller().Enqueue(namespace.GlobalNamespace, globalDNSObj.Name)
		}
	}
	return nil, nil
}

func (ic *UserIngressController) doesGlobalDNSTargetCurrentCluster(globalDNS *v3.GlobalDNS) (bool, error) {
	var targetProjectNames []string

	if globalDNS.Spec.MultiClusterAppName != "" {
		mcappName, err := getMultiClusterAppName(globalDNS.Spec.MultiClusterAppName)
		if err != nil {
			return false, err
		}
		mcapp, err := ic.multiclusterappLister.Get(namespace.GlobalNamespace, mcappName)
		if err != nil && k8serrors.IsNotFound(err) {
			logrus.Debugf("UserIngressController: reconcileallgdns, multiclusterapp not found by name %v, might be marked for deletion", mcappName)
			//pass the check and let the controller continue in this case, it will enqueue update to globaldns and the other controller will figure out the endpoint updates
			return true, nil
		}
		if err != nil {
			return false, err
		}
		for _, t := range mcapp.Spec.Targets {
			targetProjectNames = append(targetProjectNames, t.ProjectName)
		}
	} else if len(globalDNS.Spec.ProjectNames) > 0 {
		targetProjectNames = append(targetProjectNames, globalDNS.Spec.ProjectNames...)
	}

	// go through target projects and check if its part of the current cluster
	for _, projectNameSet := range targetProjectNames {
		split := strings.SplitN(projectNameSet, ":", 2)
		if len(split) != 2 {
			return false, fmt.Errorf("UserIngressController: Error in splitting project Name %v", projectNameSet)
		}
		// check if the project in this iteration belongs to the same cluster in current context
		if split[0] == ic.clusterName {
			return true, nil
		}
	}

	return false, nil
}

func (ic *UserIngressController) findGlobalDNS(fqdnRequested string) (*v3.GlobalDNS, error) {
	allGlobalDNSs, err := ic.globalDNSLister.List(namespace.GlobalNamespace, labels.NewSelector())
	if err != nil {
		return nil, fmt.Errorf("UserIngressController: Error listing GlobalDNSs %v", err)
	}

	for _, gd := range allGlobalDNSs {
		if strings.EqualFold(gd.Spec.FQDN, fqdnRequested) {
			return gd, nil
		}
	}

	return nil, nil
}

func (ic *UserIngressController) isProjectApproved(projectsApproved []string, project string) bool {
	for _, listedProject := range projectsApproved {
		split := strings.SplitN(listedProject, ":", 2)
		if len(split) != 2 {
			logrus.Errorf("UserIngressController: Error in splitting project ID %v", listedProject)
			return false
		}
		listedProjectName := split[1]
		if strings.EqualFold(listedProjectName, project) {
			return true
		}
	}
	return false
}

func (ic *UserIngressController) checkForMultiClusterApp(obj *v1beta1.Ingress, globalDNS *v3.GlobalDNS) (bool, error) {
	ingressLabels := obj.Labels
	appID := ingressLabels[appSelectorLabel]

	if appID != "" {
		//find the app CR
		// go through all projects from multiclusterapp's targets
		mcappName, err := getMultiClusterAppName(globalDNS.Spec.MultiClusterAppName)
		if err != nil {
			return false, err
		}
		mcapp, err := ic.multiclusterappLister.Get(namespace.GlobalNamespace, mcappName)
		if err != nil && k8serrors.IsNotFound(err) {
			logrus.Debugf("UserIngressController: multiclusterapp not found by name %v", mcappName)
			//pass the check and let the controller continue in this case, it will enqueue update to globaldns and the other controller will figure out the endpoint updates
			return true, nil
		}
		if err != nil {
			return false, fmt.Errorf("UserIngressController: Error %v listing multiclusterapp by name %v", err, mcappName)
		}
		for _, t := range mcapp.Spec.Targets {
			split := strings.SplitN(t.ProjectName, ":", 2)
			if len(split) != 2 {
				return false, fmt.Errorf("error in splitting project ID %v", t.ProjectName)
			}
			if split[0] != ic.clusterName {
				continue
			}
			projectNS := split[1]
			userApp, err := ic.appLister.Get(projectNS, appID)
			if err != nil && k8serrors.IsNotFound(err) {
				logrus.Debugf("UserIngressController: The app %v for this ingress is not under the target namespace %v of multiclusterapp", projectNS, appID)
				return false, nil
			} else if err != nil {
				return false, fmt.Errorf("UserIngressController: Error finding the App with the Id %v under namespace %v", appID, projectNS)
			}

			if strings.EqualFold(userApp.Spec.MultiClusterAppName, mcappName) {
				return true, nil
			}
		}
	}
	return false, nil
}

func (ic *UserIngressController) checkForProjects(obj *v1beta1.Ingress, globalDNS *v3.GlobalDNS) (bool, error) {
	ns, err := ic.namespaceLister.Get("", obj.Namespace)
	if err != nil {
		return false, fmt.Errorf("UserIngressController: Cannot find the App's namespace with the Id %v, error: %v", obj.Namespace, err)
	}
	nameSpaceProject := ns.ObjectMeta.Labels[projectSelectorLabel]

	logrus.Debugf("UserIngressController: check if Ingress's project %v belongs to GlobalDNS target projects %v", nameSpaceProject, globalDNS.Spec.ProjectNames)

	if ic.isProjectApproved(globalDNS.Spec.ProjectNames, nameSpaceProject) {
		return true, nil
	}

	return false, nil
}
<commit_msg>On ingress updates always reconcile globalDNSs of the cluster<commit_after>package globaldns

import (
	"context"
	"fmt"
	"strings"

	"github.com/rancher/rancher/pkg/namespace"
	v1Rancher "github.com/rancher/types/apis/core/v1"
	v1beta1Rancher "github.com/rancher/types/apis/extensions/v1beta1"
	"github.com/rancher/types/apis/management.cattle.io/v3"
	projectv3 "github.com/rancher/types/apis/project.cattle.io/v3"
	"github.com/rancher/types/config"
	"github.com/sirupsen/logrus"
	k8serrors "k8s.io/apimachinery/pkg/api/errors"

	"k8s.io/api/extensions/v1beta1"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/runtime"
)

const (
	annotationGlobalDNS         = "rancher.io/globalDNS.hostname"
	appSelectorLabel            = "io.cattle.field/appId"
	projectSelectorLabel        = "field.cattle.io/projectId"
	UserIngressControllerName   = "globaldns-useringress-controller"
	UserGlobalDNSControllerName = "user-controller-watching-globaldns"
)

type UserIngressController struct {
	ingressLister         v1beta1Rancher.IngressLister
	globalDNSs            v3.GlobalDNSInterface
	globalDNSLister       v3.GlobalDNSLister
	appLister             projectv3.AppLister
	namespaceLister       v1Rancher.NamespaceLister
	multiclusterappLister v3.MultiClusterAppLister
	clusterName           string
}

func newUserIngressController(ctx context.Context, clusterContext *config.UserContext) *UserIngressController {
	n := &UserIngressController{
		ingressLister:         clusterContext.Extensions.Ingresses("").Controller().Lister(),
		globalDNSs:            clusterContext.Management.Management.GlobalDNSs(""),
		globalDNSLister:       clusterContext.Management.Management.GlobalDNSs("").Controller().Lister(),
		appLister:             clusterContext.Management.Project.Apps("").Controller().Lister(),
		namespaceLister:       clusterContext.Core.Namespaces("").Controller().Lister(),
		multiclusterappLister: clusterContext.Management.Management.MultiClusterApps("").Controller().Lister(),
		clusterName:           clusterContext.ClusterName,
	}
	return n
}

func Register(ctx context.Context, clusterContext *config.UserContext) {
	n := newUserIngressController(ctx, clusterContext)
	clusterContext.Extensions.Ingresses("").AddHandler(ctx, UserIngressControllerName, n.sync)
	g := newUserGlobalDNSController(clusterContext)
	clusterContext.Management.Management.GlobalDNSs("").AddHandler(ctx, UserGlobalDNSControllerName, g.sync)
}

func (ic *UserIngressController) sync(key string, obj *v1beta1.Ingress) (runtime.Object, error) {
	return ic.reconcileAllGlobalDNSs()
}

func (ic *UserIngressController) noGlobalDNS() (bool, error) {
	gd, err := ic.globalDNSLister.List(namespace.GlobalNamespace, labels.NewSelector())
	if err != nil {
		return true, err
	}

	return len(gd) == 0, nil
}

func (ic *UserIngressController) reconcileAllGlobalDNSs() (runtime.Object, error) {
	globalDNSList, err := ic.globalDNSLister.List(namespace.GlobalNamespace, labels.NewSelector())
	if err != nil {
		return nil, err
	}

	for _, globalDNSObj := range globalDNSList {
		//call update on each GlobalDNS obj that refers to this current cluster
		targetsCluster, err := ic.doesGlobalDNSTargetCurrentCluster(globalDNSObj)
		if err != nil {
			return nil, err
		}
		if targetsCluster {
			//enqueue it to the globalDNS controller
			ic.globalDNSs.Controller().Enqueue(namespace.GlobalNamespace, globalDNSObj.Name)
		}
	}
	return nil, nil
}

func (ic *UserIngressController) doesGlobalDNSTargetCurrentCluster(globalDNS *v3.GlobalDNS) (bool, error) {
	var targetProjectNames []string

	if globalDNS.Spec.MultiClusterAppName != "" {
		mcappName, err := getMultiClusterAppName(globalDNS.Spec.MultiClusterAppName)
		if err != nil {
			return false, err
		}
		mcapp, err := ic.multiclusterappLister.Get(namespace.GlobalNamespace, mcappName)
		if err != nil && k8serrors.IsNotFound(err) {
			logrus.Debugf("UserIngressController: reconcileallgdns, multiclusterapp not found by name %v, might be marked for deletion", mcappName)
			//pass the check and let the controller continue in this case, it will enqueue update to globaldns and the other controller will figure out the endpoint updates
			return true, nil
		}
		if err != nil {
			return false, err
		}
		for _, t := range mcapp.Spec.Targets {
			targetProjectNames = append(targetProjectNames, t.ProjectName)
		}
	} else if len(globalDNS.Spec.ProjectNames) > 0 {
		targetProjectNames = append(targetProjectNames, globalDNS.Spec.ProjectNames...)
	}

	// go through target projects and check if its part of the current cluster
	for _, projectNameSet := range targetProjectNames {
		split := strings.SplitN(projectNameSet, ":", 2)
		if len(split) != 2 {
			return false, fmt.Errorf("UserIngressController: Error in splitting project Name %v", projectNameSet)
		}
		// check if the project in this iteration belongs to the same cluster in current context
		if split[0] == ic.clusterName {
			return true, nil
		}
	}

	return false, nil
}
<|endoftext|>


<commit_before>package main

import (
	"io/ioutil"
	"log"
	"net/url"
	"os"
	"path/filepath"
	"time"

	"github.com/alphagov/govuk_crawler_worker/http_crawler"
	"github.com/alphagov/govuk_crawler_worker/queue"
	"github.com/alphagov/govuk_crawler_worker/ttl_hash_set"
	"github.com/alphagov/govuk_crawler_worker/util"
	"github.com/streadway/amqp"
)

func AcknowledgeItem(inbound <-chan *CrawlerMessageItem, ttlHashSet *ttl_hash_set.TTLHashSet) {
	for item := range inbound {
		start := time.Now()
		url := item.URL()

		_, err := ttlHashSet.Add(url)
		if err != nil {
			item.Reject(false)
			log.Println("Acknowledge failed (rejecting):", url, err)
			continue
		}

		item.Ack(false)
		log.Println("Acknowledged:", url)

		util.StatsDTiming("acknowledge_item", start, time.Now())
	}
}

func CrawlURL(crawlChannel <-chan *CrawlerMessageItem, crawler *http_crawler.Crawler, crawlerThreads int) <-chan *CrawlerMessageItem {
	if crawlerThreads <= 0 {
		panic("cannot start a negative or zero number of crawler threads")
	}

	extractChannel := make(chan *CrawlerMessageItem, 2)

	crawlLoop := func(
		crawl <-chan *CrawlerMessageItem,
		extract chan<- *CrawlerMessageItem,
		crawler *http_crawler.Crawler,
	) {
		for item := range crawl {
			start := time.Now()
			u, err := url.Parse(item.URL())
			if err != nil {
				item.Reject(false)
				log.Println("Couldn't crawl, invalid URL (rejecting):", item.URL(), err)
				continue
			}
			log.Println("Crawling URL:", u)

			body, err := crawler.Crawl(u)
			if err != nil {
				if err == http_crawler.RetryRequest5XXError || err == http_crawler.RetryRequest429Error {
					item.Reject(true)
					log.Println("Couldn't crawl (requeueing):", u.String(), err)

					if err == http_crawler.RetryRequest429Error {
						sleepTime := 5 * time.Second

						// Back off from crawling for a few seconds.
						log.Println("Sleeping for: ", sleepTime, " seconds. Received 429 HTTP status")
						time.Sleep(sleepTime)
					}
				} else {
					item.Reject(false)
					log.Println("Couldn't crawl (rejecting):", u.String(), err)
				}

				continue
			}

			item.HTMLBody = body

			if item.IsHTML() {
				extract <- item
			} else {
				item.Ack(false)
			}

			util.StatsDTiming("crawl_url", start, time.Now())
		}
	}

	for i := 1; i <= crawlerThreads; i++ {
		go crawlLoop(crawlChannel, extractChannel, crawler)
	}

	return extractChannel
}

func WriteItemToDisk(basePath string, crawlChannel <-chan *CrawlerMessageItem) <-chan *CrawlerMessageItem {
	extractChannel := make(chan *CrawlerMessageItem, 2)

	writeLoop := func(
		crawl <-chan *CrawlerMessageItem,
		extract chan<- *CrawlerMessageItem,
	) {
		for item := range crawl {
			start := time.Now()
			relativeFilePath, err := item.RelativeFilePath()

			if err != nil {
				item.Reject(false)
				log.Println("Couldn't write to disk (rejecting):", err)
				continue
			}

			filePath := filepath.Join(basePath, relativeFilePath)
			basePath := filepath.Dir(filePath)
			err = os.MkdirAll(basePath, 0755)

			if err != nil {
				item.Reject(false)
				log.Println("Couldn't write to disk (rejecting):", filePath, err)
				continue
			}

			err = ioutil.WriteFile(filePath, item.HTMLBody, 0644)

			if err != nil {
				item.Reject(false)
				log.Println("Couldn't write to disk (rejecting):", filePath, err)
				continue
			}

			log.Println("Wrote URL body to disk for:", item.URL())
			extract <- item

			util.StatsDTiming("write_to_disk", start, time.Now())
		}
	}

	go writeLoop(crawlChannel, extractChannel)

	return extractChannel
}

func ExtractURLs(extractChannel <-chan *CrawlerMessageItem) (<-chan string, <-chan *CrawlerMessageItem) {
	publishChannel := make(chan string, 100)
	acknowledgeChannel := make(chan *CrawlerMessageItem, 1)

	extractLoop := func(
		extract <-chan *CrawlerMessageItem,
		publish chan<- string,
		acknowledge chan<- *CrawlerMessageItem,
	) {
		for item := range extract {
			start := time.Now()
			urls, err := item.ExtractURLs()
			if err != nil {
				item.Reject(false)
				log.Println("ExtractURLs (rejecting):", string(item.Body), err)
			}

			log.Println("Extracted URLs:", len(urls))

			for _, u := range urls {
				publish <- u.String()
			}

			acknowledge <- item

			util.StatsDTiming("extract_urls", start, time.Now())
		}
	}

	go extractLoop(extractChannel, publishChannel, acknowledgeChannel)

	return publishChannel, acknowledgeChannel
}

func PublishURLs(ttlHashSet *ttl_hash_set.TTLHashSet, queueManager *queue.QueueManager, publish <-chan string) {
	for url := range publish {
		start := time.Now()
		exists, err := ttlHashSet.Exists(url)

		if err != nil {
			log.Println("Couldn't check existence of URL:", url, err)
		}

		if !exists {
			err = queueManager.Publish("#", "text/plain", url)
			if err != nil {
				log.Fatalln("Delivery failed:", url, err)
			}
		}

		util.StatsDGauge("publish_urls", int64(len(publish)))
		util.StatsDTiming("publish_urls", start, time.Now())
	}
}

func ReadFromQueue(inboundChannel <-chan amqp.Delivery, rootURL *url.URL, ttlHashSet *ttl_hash_set.TTLHashSet, blacklistPaths []string) chan *CrawlerMessageItem {
	outboundChannel := make(chan *CrawlerMessageItem, 2)

	readLoop := func(
		inbound <-chan amqp.Delivery,
		outbound chan<- *CrawlerMessageItem,
		ttlHashSet *ttl_hash_set.TTLHashSet,
		blacklistPaths []string,
	) {
		for item := range inbound {
			start := time.Now()
			message := NewCrawlerMessageItem(item, rootURL, blacklistPaths)

			exists, err := ttlHashSet.Exists(message.URL())
			if err != nil {
				item.Reject(true)
				log.Println("Couldn't check existence of (rejecting):", message.URL(), err)
				continue
			}

			if !exists {
				outbound <- message
			} else {
				log.Println("URL already crawled:", message.URL())
				item.Ack(false)
			}

			util.StatsDTiming("read_from_queue", start, time.Now())
		}
	}

	go readLoop(inboundChannel, outboundChannel, ttlHashSet, blacklistPaths)

	return outboundChannel
}
<commit_msg>Continue when we reject a message for not extracting URLs<commit_after>package main

import (
	"io/ioutil"
	"log"
	"net/url"
	"os"
	"path/filepath"
	"time"

	"github.com/alphagov/govuk_crawler_worker/http_crawler"
	"github.com/alphagov/govuk_crawler_worker/queue"
	"github.com/alphagov/govuk_crawler_worker/ttl_hash_set"
	"github.com/alphagov/govuk_crawler_worker/util"
	"github.com/streadway/amqp"
)

func AcknowledgeItem(inbound <-chan *CrawlerMessageItem, ttlHashSet *ttl_hash_set.TTLHashSet) {
	for item := range inbound {
		start := time.Now()
		url := item.URL()

		_, err := ttlHashSet.Add(url)
		if err != nil {
			item.Reject(false)
			log.Println("Acknowledge failed (rejecting):", url, err)
			continue
		}

		item.Ack(false)
		log.Println("Acknowledged:", url)

		util.StatsDTiming("acknowledge_item", start, time.Now())
	}
}

func CrawlURL(crawlChannel <-chan *CrawlerMessageItem, crawler *http_crawler.Crawler, crawlerThreads int) <-chan *CrawlerMessageItem {
	if crawlerThreads <= 0 {
		panic("cannot start a negative or zero number of crawler threads")
	}

	extractChannel := make(chan *CrawlerMessageItem, 2)

	crawlLoop := func(
		crawl <-chan *CrawlerMessageItem,
		extract chan<- *CrawlerMessageItem,
		crawler *http_crawler.Crawler,
	) {
		for item := range crawl {
			start := time.Now()
			u, err := url.Parse(item.URL())
			if err != nil {
				item.Reject(false)
				log.Println("Couldn't crawl, invalid URL (rejecting):", item.URL(), err)
				continue
			}
			log.Println("Crawling URL:", u)

			body, err := crawler.Crawl(u)
			if err != nil {
				if err == http_crawler.RetryRequest5XXError || err == http_crawler.RetryRequest429Error {
					item.Reject(true)
					log.Println("Couldn't crawl (requeueing):", u.String(), err)

					if err == http_crawler.RetryRequest429Error {
						sleepTime := 5 * time.Second

						// Back off from crawling for a few seconds.
						log.Println("Sleeping for: ", sleepTime, " seconds. Received 429 HTTP status")
						time.Sleep(sleepTime)
					}
				} else {
					item.Reject(false)
					log.Println("Couldn't crawl (rejecting):", u.String(), err)
				}

				continue
			}

			item.HTMLBody = body

			if item.IsHTML() {
				extract <- item
			} else {
				item.Ack(false)
			}

			util.StatsDTiming("crawl_url", start, time.Now())
		}
	}

	for i := 1; i <= crawlerThreads; i++ {
		go crawlLoop(crawlChannel, extractChannel, crawler)
	}

	return extractChannel
}

func WriteItemToDisk(basePath string, crawlChannel <-chan *CrawlerMessageItem) <-chan *CrawlerMessageItem {
	extractChannel := make(chan *CrawlerMessageItem, 2)

	writeLoop := func(
		crawl <-chan *CrawlerMessageItem,
		extract chan<- *CrawlerMessageItem,
	) {
		for item := range crawl {
			start := time.Now()
			relativeFilePath, err := item.RelativeFilePath()

			if err != nil {
				item.Reject(false)
				log.Println("Couldn't write to disk (rejecting):", err)
				continue
			}

			filePath := filepath.Join(basePath, relativeFilePath)
			basePath := filepath.Dir(filePath)
			err = os.MkdirAll(basePath, 0755)

			if err != nil {
				item.Reject(false)
				log.Println("Couldn't write to disk (rejecting):", filePath, err)
				continue
			}

			err = ioutil.WriteFile(filePath, item.HTMLBody, 0644)

			if err != nil {
				item.Reject(false)
				log.Println("Couldn't write to disk (rejecting):", filePath, err)
				continue
			}

			log.Println("Wrote URL body to disk for:", item.URL())
			extract <- item

			util.StatsDTiming("write_to_disk", start, time.Now())
		}
	}

	go writeLoop(crawlChannel, extractChannel)

	return extractChannel
}

func ExtractURLs(extractChannel <-chan *CrawlerMessageItem) (<-chan string, <-chan *CrawlerMessageItem) {
	publishChannel := make(chan string, 100)
	acknowledgeChannel := make(chan *CrawlerMessageItem, 1)

	extractLoop := func(
		extract <-chan *CrawlerMessageItem,
		publish chan<- string,
		acknowledge chan<- *CrawlerMessageItem,
	) {
		for item := range extract {
			start := time.Now()
			urls, err := item.ExtractURLs()
			if err != nil {
				item.Reject(false)
				log.Println("ExtractURLs (rejecting):", string(item.Body), err)

				continue
			}

			log.Println("Extracted URLs:", len(urls))

			for _, u := range urls {
				publish <- u.String()
			}

			acknowledge <- item

			util.StatsDTiming("extract_urls", start, time.Now())
		}
	}

	go extractLoop(extractChannel, publishChannel, acknowledgeChannel)

	return publishChannel, acknowledgeChannel
}

func PublishURLs(ttlHashSet *ttl_hash_set.TTLHashSet, queueManager *queue.QueueManager, publish <-chan string) {
	for url := range publish {
		start := time.Now()
		exists, err := ttlHashSet.Exists(url)

		if err != nil {
			log.Println("Couldn't check existence of URL:", url, err)
		}

		if !exists {
			err = queueManager.Publish("#", "text/plain", url)
			if err != nil {
				log.Fatalln("Delivery failed:", url, err)
			}
		}

		util.StatsDGauge("publish_urls", int64(len(publish)))
		util.StatsDTiming("publish_urls", start, time.Now())
	}
}

func ReadFromQueue(inboundChannel <-chan amqp.Delivery, rootURL *url.URL, ttlHashSet *ttl_hash_set.TTLHashSet, blacklistPaths []string) chan *CrawlerMessageItem {
	outboundChannel := make(chan *CrawlerMessageItem, 2)

	readLoop := func(
		inbound <-chan amqp.Delivery,
		outbound chan<- *CrawlerMessageItem,
		ttlHashSet *ttl_hash_set.TTLHashSet,
		blacklistPaths []string,
	) {
		for item := range inbound {
			start := time.Now()
			message := NewCrawlerMessageItem(item, rootURL, blacklistPaths)

			exists, err := ttlHashSet.Exists(message.URL())
			if err != nil {
				item.Reject(true)
				log.Println("Couldn't check existence of (rejecting):", message.URL(), err)
				continue
			}

			if !exists {
				outbound <- message
			} else {
				log.Println("URL already crawled:", message.URL())
				item.Ack(false)
			}

			util.StatsDTiming("read_from_queue", start, time.Now())
		}
	}

	go readLoop(inboundChannel, outboundChannel, ttlHashSet, blacklistPaths)

	return outboundChannel
}
<|endoftext|>


<commit_before>package create_container_step

import (
	"github.com/cloudfoundry-incubator/runtime-schema/models"
	steno "github.com/cloudfoundry/gosteno"
	"github.com/cloudfoundry-incubator/gordon"
)

type ContainerStep struct {
	runOnce         *models.RunOnce
	logger          *steno.Logger
	wardenClient    gordon.Client
	containerHandle *string
}

func New(
	runOnce *models.RunOnce,
	logger *steno.Logger,
	wardenClient gordon.Client,
	containerHandle *string,
) *ContainerStep {
	return &ContainerStep{
		runOnce:         runOnce,
		logger:          logger,
		wardenClient:    wardenClient,
		containerHandle: containerHandle,
	}
}

func (step ContainerStep) Perform() error {
	createResponse, err := step.wardenClient.Create()
	if err != nil {
		step.logger.Errord(
			map[string]interface{}{
				"runonce-guid": step.runOnce.Guid,
				"error":        err.Error(),
			},
			"runonce.container-create.failed",
		)

		return err
	}

	*step.containerHandle = createResponse.GetHandle()

	return nil
}

func (step ContainerStep) Cancel() {}

func (step ContainerStep) Cleanup() {
	_, err := step.wardenClient.Destroy(*step.containerHandle)
	if err != nil {
		step.logger.Errord(
			map[string]interface{}{
				"runonce-guid": step.runOnce.Guid,
				"handle":       step.runOnce.ContainerHandle,
				"error":        err.Error(),
			},
			"runonce.container-destroy.failed",
		)
	}
}
<commit_msg>pass nil to Create<commit_after>package create_container_step

import (
	"github.com/cloudfoundry-incubator/gordon"
	"github.com/cloudfoundry-incubator/runtime-schema/models"
	steno "github.com/cloudfoundry/gosteno"
)

type ContainerStep struct {
	runOnce         *models.RunOnce
	logger          *steno.Logger
	wardenClient    gordon.Client
	containerHandle *string
}

func New(
	runOnce *models.RunOnce,
	logger *steno.Logger,
	wardenClient gordon.Client,
	containerHandle *string,
) *ContainerStep {
	return &ContainerStep{
		runOnce:         runOnce,
		logger:          logger,
		wardenClient:    wardenClient,
		containerHandle: containerHandle,
	}
}

func (step ContainerStep) Perform() error {
	createResponse, err := step.wardenClient.Create(nil)
	if err != nil {
		step.logger.Errord(
			map[string]interface{}{
				"runonce-guid": step.runOnce.Guid,
				"error":        err.Error(),
			},
			"runonce.container-create.failed",
		)

		return err
	}

	*step.containerHandle = createResponse.GetHandle()

	return nil
}

func (step ContainerStep) Cancel() {}

func (step ContainerStep) Cleanup() {
	_, err := step.wardenClient.Destroy(*step.containerHandle)
	if err != nil {
		step.logger.Errord(
			map[string]interface{}{
				"runonce-guid": step.runOnce.Guid,
				"handle":       step.runOnce.ContainerHandle,
				"error":        err.Error(),
			},
			"runonce.container-destroy.failed",
		)
	}
}
<|endoftext|>


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

package ogletest_test

import (
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"path"
	"os"
	"os/exec"
	"regexp"
	"strings"
	"testing"
)

var dumpNew = flag.Bool("dump_new", false, "Dump new golden files.")
var objDir string

////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////

// buildPackage build a copy of the (possibly locally modified) package, and
// return a path to an include path where it can be found.
func buildPackage() (string, error) {
	// HACK(jacobsa): The gotest command reads the GCFLAGS environment variable,
	// which allows us to give a -I flag to 6g to let it find the locally built
	// package. However 6l cannot find it, and there is no linker flags
	// environment variable. So we actually install the package here, instead of
	// just building it locally.
	//
	// TODO(jacobsa): When the new 'go' tool comes out, check whether this still
	// replies to its 'test' command. If so, file an issue to get a linker flags
	// environment variable added.
	cmd := exec.Command("gomake", "install")
	output, err := cmd.CombinedOutput()

	if err != nil {
		return "", errors.New(fmt.Sprintf("%v:\n%s", err, output))
	}

	wd, err := os.Getwd()
	if err != nil {
		return "", errors.New("os.Getwd: " + err.Error())
	}

	return path.Join(wd, "_obj"), nil
}

// getCaseNames looks for integration test cases as files in the test_cases
// directory.
func getCaseNames() ([]string, error) {
	// Open the test cases directory.
	dir, err := os.Open("test_cases")
	if err != nil {
		return nil, errors.New(fmt.Sprintf("Opening dir: %v", err))
	}

	// Get a list of the names in the directory.
	names, err := dir.Readdirnames(0)
	if err != nil {
		return nil, errors.New(fmt.Sprintf("Readdirnames: %v", err))
	}

	// Filter the names.
	result := make([]string, len(names))
	resultLen := 0
	for _, name := range names {
		// Skip golden files and hidden files.
		if strings.HasPrefix(name, "golden.") || strings.HasPrefix(name, ".") {
			continue
		}

		// Check for the right format.
		if !strings.HasSuffix(name, "_test.go") {
			return nil, errors.New(fmt.Sprintf("Unexpected file: %s", name))
		}

		// Store the name minus the extension.
		result[resultLen] = name[:len(name) - 3]
		resultLen++
	}

	return result[:resultLen], nil
}

func writeContentsToFileOrDie(contents []byte, path string) {
	if err := ioutil.WriteFile(path, contents, 0600); err != nil {
		panic("iotuil.WriteFile: " + err.Error())
	}
}

func readFileOrDie(path string) []byte {
	contents, err := ioutil.ReadFile(path)
	if err != nil {
		panic("iotuil.ReadFile: " + err.Error())
	}

	return contents
}

// cleanOutput transforms the supplied output so that it no longer contains
// information that changes from run to run, making the golden tests less
// flaky.
func cleanOutput(o []byte, tempDir string) []byte {
	// Replace references to the temporary directory.
	o = []byte(strings.Replace(string(o), tempDir, "/tmp/dir", -1))

	// Replace things that look like line numbers and process counters in stack
	// traces.
	stackFrameRe := regexp.MustCompile(`\S+\.go:\d+ \(0x[0-9a-f]+\)`)
	o = stackFrameRe.ReplaceAll(o, []byte("some_file.go:0 (0x00000)"))

	// Replace unstable timings in gotest fail messages.
	timingRe := regexp.MustCompile(`--- FAIL: .* \(\d\.\d{2} seconds\)`)
	o = timingRe.ReplaceAll(o, []byte("--- FAIL: sometest (0.00 seconds)"))

	return o
}

// runTestCase runs the case with the supplied name (e.g. "passing_test"), and
// returns its output and exit code.
func runTestCase(name string) ([]byte, int, error) {
	// Create a temporary directory for the test files.
	tempDir, err := ioutil.TempDir("", "ogletest_integration_test")
	if err != nil {
		return nil, 0, errors.New("ioutil.TempDir: " + err.Error())
	}

	defer os.RemoveAll(tempDir)

	// Create a makefile within the directory.
	makefileContents := "include $(GOROOT)/src/Make.inc\n" +
		"TARG = github.com/jacobsa/ogletest/foobar\n" +
		"GOFILES = " + name + ".go\n" +
		"GCFLAGS += -I" + objDir + "\n" +
		"include $(GOROOT)/src/Make.pkg\n"

	writeContentsToFileOrDie([]byte(makefileContents), path.Join(tempDir, "Makefile"))

	// Create the test source file.
	sourceFile := name + ".go"
	testContents := readFileOrDie(path.Join("test_cases", sourceFile))
	writeContentsToFileOrDie(testContents, path.Join(tempDir, sourceFile))

	// Invoke gotest. Special case: pass a test filter to the filtered_test case.
	cmd := exec.Command("gotest")
	if name == "filtered_test" {
		cmd.Args = append(cmd.Args, "--ogletest.run=Test(Bar|Baz)")
	}

	cmd.Dir = tempDir
	output, err := cmd.CombinedOutput()

	// Did the process exist with zero code?
	if err == nil {
		return output, 0, nil
	}

	// Make sure the process actually exited.
	exitError, ok := err.(*exec.ExitError)
	if !ok || !exitError.Exited() {
		return nil, 0, errors.New("exec.Command.Output: " + err.Error())
	}

	output = cleanOutput(output, tempDir)
	return output, exitError.ExitStatus(), nil
}

// checkGolden file checks the supplied actual output for the named test case
// against the golden file for that case. If requested by the user, it rewrites
// the golden file on failure.
func checkAgainstGoldenFile(caseName string, output []byte) bool {
	goldenFile := path.Join("test_cases", "golden." + caseName)
	goldenContents := readFileOrDie(goldenFile)

	result := string(output) == string(goldenContents)
	if !result && *dumpNew {
		writeContentsToFileOrDie(output, goldenFile)
	}

	return result
}

////////////////////////////////////////////////////////////
// Tests
////////////////////////////////////////////////////////////

func TestGoldenFiles(t *testing.T) {
	// Ensure the local package is up to date, and get a path to it. This will
	// prevent the test cases from using the installed version, which may be out
	// of date.
	var err error
	objDir, err = buildPackage()

	// We expect there to be at least one case.
	caseNames, err := getCaseNames()
	if err != nil || len(caseNames) == 0 {
		t.Fatalf("Error getting cases: %v", err)
	}

	// Run each test case.
	for _, caseName := range caseNames {
		// Run the test case.
		output, exitCode, err := runTestCase(caseName)
		if err != nil {
			t.Fatalf("Running test case %s: %v", caseName, err)
		}

		// Check the status code. We assume all test cases fail except for passing_test.
		shouldPass := caseName == "passing_test"
		didPass := exitCode == 0
		if shouldPass != didPass {
			t.Errorf("Bad exit code for test case %s: %d", caseName, exitCode)
		}

		// Check the output against the golden file.
		if !checkAgainstGoldenFile(caseName, output) {
			t.Errorf("Output for test case %s doesn't match golden file.", caseName)
		}
	}
}
<commit_msg>Added createTempPackageDir for issue #19.<commit_after>// Copyright 2011 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package ogletest_test

import (
	"errors"
	"flag"
	"fmt"
	"go/build"
	"io/ioutil"
	"path"
	"os"
	"os/exec"
	"regexp"
	"strings"
	"testing"
)

var dumpNew = flag.Bool("dump_new", false, "Dump new golden files.")
var objDir string

////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////

// getCaseNames looks for integration test cases as files in the test_cases
// directory.
func getCaseNames() ([]string, error) {
	// Open the test cases directory.
	dir, err := os.Open("test_cases")
	if err != nil {
		return nil, errors.New(fmt.Sprintf("Opening dir: %v", err))
	}

	// Get a list of the names in the directory.
	names, err := dir.Readdirnames(0)
	if err != nil {
		return nil, errors.New(fmt.Sprintf("Readdirnames: %v", err))
	}

	// Filter the names.
	result := make([]string, len(names))
	resultLen := 0
	for _, name := range names {
		// Skip golden files and hidden files.
		if strings.HasPrefix(name, "golden.") || strings.HasPrefix(name, ".") {
			continue
		}

		// Check for the right format.
		if !strings.HasSuffix(name, "_test.go") {
			return nil, errors.New(fmt.Sprintf("Unexpected file: %s", name))
		}

		// Store the name minus the extension.
		result[resultLen] = name[:len(name) - 3]
		resultLen++
	}

	return result[:resultLen], nil
}

func writeContentsToFileOrDie(contents []byte, path string) {
	if err := ioutil.WriteFile(path, contents, 0600); err != nil {
		panic("ioutil.WriteFile: " + err.Error())
	}
}

func readFileOrDie(path string) []byte {
	contents, err := ioutil.ReadFile(path)
	if err != nil {
		panic("ioutil.ReadFile: " + err.Error())
	}

	return contents
}

// cleanOutput transforms the supplied output so that it no longer contains
// information that changes from run to run, making the golden tests less
// flaky.
func cleanOutput(o []byte, tempDir string) []byte {
	// Replace references to the temporary directory.
	o = []byte(strings.Replace(string(o), tempDir, "/tmp/dir", -1))

	// Replace things that look like line numbers and process counters in stack
	// traces.
	stackFrameRe := regexp.MustCompile(`\S+\.go:\d+ \(0x[0-9a-f]+\)`)
	o = stackFrameRe.ReplaceAll(o, []byte("some_file.go:0 (0x00000)"))

	// Replace unstable timings in gotest fail messages.
	timingRe := regexp.MustCompile(`--- FAIL: .* \(\d\.\d{2} seconds\)`)
	o = timingRe.ReplaceAll(o, []byte("--- FAIL: sometest (0.00 seconds)"))

	return o
}

// Create a temporary package directory somewhere that 'go test' can find, and
// return the directory and package name.
func createTempPackageDir(caseName string) (dir, pkg string) {
	const ogletestPkg = "github.com/jacobsa/ogletest"

	// Figure out where the local source code for ogletest is.
	ogletestTree, _, err := build.FindTree(ogletestPkg)
	if err != nil { panic("Finding ogletest package: " + err.Error()) }

	// Create a temporary directory underneath this.
	prefix := fmt.Sprintf("tmp-%s", caseName)
	dir, err = ioutil.TempDir(ogletestTree.Path, prefix)
	if err != nil { panic("ioutil.TempDir: " + err.Error()) }

	pkg = path.Join("github.com/jacobsa/ogletest", dir[len(ogletestTree.Path):])
	return
}

// runTestCase runs the case with the supplied name (e.g. "passing_test"), and
// returns its output and exit code.
func runTestCase(name string) ([]byte, int, error) {
	// Create a temporary directory for the test files.
	tempDir, err := ioutil.TempDir("", "ogletest_integration_test")
	if err != nil {
		return nil, 0, errors.New("ioutil.TempDir: " + err.Error())
	}

	defer os.RemoveAll(tempDir)

	// Create a makefile within the directory.
	makefileContents := "include $(GOROOT)/src/Make.inc\n" +
		"TARG = github.com/jacobsa/ogletest/foobar\n" +
		"GOFILES = " + name + ".go\n" +
		"GCFLAGS += -I" + objDir + "\n" +
		"include $(GOROOT)/src/Make.pkg\n"

	writeContentsToFileOrDie([]byte(makefileContents), path.Join(tempDir, "Makefile"))

	// Create the test source file.
	sourceFile := name + ".go"
	testContents := readFileOrDie(path.Join("test_cases", sourceFile))
	writeContentsToFileOrDie(testContents, path.Join(tempDir, sourceFile))

	// Invoke gotest. Special case: pass a test filter to the filtered_test case.
	cmd := exec.Command("gotest")
	if name == "filtered_test" {
		cmd.Args = append(cmd.Args, "--ogletest.run=Test(Bar|Baz)")
	}

	cmd.Dir = tempDir
	output, err := cmd.CombinedOutput()

	// Did the process exist with zero code?
	if err == nil {
		return output, 0, nil
	}

	// Make sure the process actually exited.
	exitError, ok := err.(*exec.ExitError)
	if !ok || !exitError.Exited() {
		return nil, 0, errors.New("exec.Command.Output: " + err.Error())
	}

	output = cleanOutput(output, tempDir)
	return output, exitError.ExitStatus(), nil
}

// checkGolden file checks the supplied actual output for the named test case
// against the golden file for that case. If requested by the user, it rewrites
// the golden file on failure.
func checkAgainstGoldenFile(caseName string, output []byte) bool {
	goldenFile := path.Join("test_cases", "golden." + caseName)
	goldenContents := readFileOrDie(goldenFile)

	result := string(output) == string(goldenContents)
	if !result && *dumpNew {
		writeContentsToFileOrDie(output, goldenFile)
	}

	return result
}

////////////////////////////////////////////////////////////
// Tests
////////////////////////////////////////////////////////////

func TestGoldenFiles(t *testing.T) {
	// We expect there to be at least one case.
	caseNames, err := getCaseNames()
	if err != nil || len(caseNames) == 0 {
		t.Fatalf("Error getting cases: %v", err)
	}

	// Run each test case.
	for _, caseName := range caseNames {
		// Run the test case.
		output, exitCode, err := runTestCase(caseName)
		if err != nil {
			t.Fatalf("Running test case %s: %v", caseName, err)
		}

		// Check the status code. We assume all test cases fail except for
		// passing_test.
		shouldPass := caseName == "passing_test"
		didPass := exitCode == 0
		if shouldPass != didPass {
			t.Errorf("Bad exit code for test case %s: %d", caseName, exitCode)
		}

		// Check the output against the golden file.
		if !checkAgainstGoldenFile(caseName, output) {
			t.Errorf("Output for test case %s doesn't match golden file.", caseName)
		}
	}
}
<|endoftext|>


<commit_before>package prometheus_client

import (
	"context"
	"fmt"
	"log"
	"net/http"
	"regexp"
	"sort"
	"strings"
	"sync"
	"time"

	"github.com/influxdata/telegraf"
	"github.com/influxdata/telegraf/internal"
	"github.com/influxdata/telegraf/plugins/outputs"
	"github.com/prometheus/client_golang/prometheus"
)

var invalidNameCharRE = regexp.MustCompile(`[^a-zA-Z0-9_]`)

// SampleID uniquely identifies a Sample
type SampleID string

// Sample represents the current value of a series.
type Sample struct {
	// Labels are the Prometheus labels.
	Labels map[string]string
	// Value is the value in the Prometheus output.
	Value float64
	// Expiration is the deadline that this Sample is valid until.
	Expiration time.Time
}

// MetricFamily contains the data required to build valid prometheus Metrics.
type MetricFamily struct {
	// Samples are the Sample belonging to this MetricFamily.
	Samples map[SampleID]*Sample
	// Type of the Value.
	ValueType prometheus.ValueType
	// LabelSet is the label counts for all Samples.
	LabelSet map[string]int
}

type PrometheusClient struct {
	Listen             string
	ExpirationInterval internal.Duration `toml:"expiration_interval"`

	server *http.Server

	sync.Mutex
	// fam is the non-expired MetricFamily by Prometheus metric name.
	fam map[string]*MetricFamily
	// now returns the current time.
	now func() time.Time
}

var sampleConfig = `
  ## Address to listen on
  # listen = ":9273"

  ## Interval to expire metrics and not deliver to prometheus, 0 == no expiration
  # expiration_interval = "60s"
`

func (p *PrometheusClient) Start() error {
	prometheus.Register(p)

	if p.Listen == "" {
		p.Listen = "localhost:9273"
	}

	mux := http.NewServeMux()
	mux.Handle("/metrics", prometheus.Handler())

	p.server = &http.Server{
		Addr:    p.Listen,
		Handler: mux,
	}

	go func() {
		if err := p.server.ListenAndServe(); err != nil {
			log.Printf("E! Error creating prometheus metric endpoint, err: %s\n",
				err.Error())
		}
	}()
	return nil
}

func (p *PrometheusClient) Stop() {
	// plugin gets cleaned up in Close() already.
}

func (p *PrometheusClient) Connect() error {
	// This service output does not need to make any further connections
	return nil
}

func (p *PrometheusClient) Close() error {
	ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
	defer cancel()
	return p.server.Shutdown(ctx)
}

func (p *PrometheusClient) SampleConfig() string {
	return sampleConfig
}

func (p *PrometheusClient) Description() string {
	return "Configuration for the Prometheus client to spawn"
}

// Implements prometheus.Collector
func (p *PrometheusClient) Describe(ch chan<- *prometheus.Desc) {
	prometheus.NewGauge(prometheus.GaugeOpts{Name: "Dummy", Help: "Dummy"}).Describe(ch)
}

// Expire removes Samples that have expired.
func (p *PrometheusClient) Expire() {
	now := p.now()
	for name, family := range p.fam {
		for key, sample := range family.Samples {
			if p.ExpirationInterval.Duration != 0 && now.After(sample.Expiration) {
				for k, _ := range sample.Labels {
					family.LabelSet[k]--
				}
				delete(family.Samples, key)

				if len(family.Samples) == 0 {
					delete(p.fam, name)
				}
			}
		}
	}
}

// Collect implements prometheus.Collector
func (p *PrometheusClient) Collect(ch chan<- prometheus.Metric) {
	p.Lock()
	defer p.Unlock()

	p.Expire()

	for name, family := range p.fam {
		// Get list of all labels on MetricFamily
		var labelNames []string
		for k, v := range family.LabelSet {
			if v > 0 {
				labelNames = append(labelNames, k)
			}
		}
		desc := prometheus.NewDesc(name, "Telegraf collected metric", labelNames, nil)

		for _, sample := range family.Samples {
			// Get labels for this sample; unset labels will be set to the
			// empty string
			var labels []string
			for _, label := range labelNames {
				v := sample.Labels[label]
				labels = append(labels, v)
			}

			metric, err := prometheus.NewConstMetric(desc, family.ValueType, sample.Value, labels...)
			if err != nil {
				log.Printf("E! Error creating prometheus metric, "+
					"key: %s, labels: %v,\nerr: %s\n",
					name, labels, err.Error())
			}

			ch <- metric
		}
	}
}

func sanitize(value string) string {
	return invalidNameCharRE.ReplaceAllString(value, "_")
}

func valueType(tt telegraf.ValueType) prometheus.ValueType {
	switch tt {
	case telegraf.Counter:
		return prometheus.CounterValue
	case telegraf.Gauge:
		return prometheus.GaugeValue
	default:
		return prometheus.UntypedValue
	}
}

// CreateSampleID creates a SampleID based on the tags of a telegraf.Metric.
func CreateSampleID(tags map[string]string) SampleID {
	pairs := make([]string, 0, len(tags))
	for k, v := range tags {
		pairs = append(pairs, fmt.Sprintf("%s=%s", k, v))
	}
	sort.Strings(pairs)
	return SampleID(strings.Join(pairs, ","))
}

func (p *PrometheusClient) Write(metrics []telegraf.Metric) error {
	p.Lock()
	defer p.Unlock()

	now := p.now()

	for _, point := range metrics {
		tags := point.Tags()
		vt := valueType(point.Type())
		sampleID := CreateSampleID(tags)

		labels := make(map[string]string)
		for k, v := range tags {
			labels[sanitize(k)] = v
		}

		for fn, fv := range point.Fields() {
			// Ignore string and bool fields.
			var value float64
			switch fv := fv.(type) {
			case int64:
				value = float64(fv)
			case float64:
				value = fv
			default:
				continue
			}

			sample := &Sample{
				Labels:     labels,
				Value:      value,
				Expiration: now.Add(p.ExpirationInterval.Duration),
			}

			// Special handling of value field; supports passthrough from
			// the prometheus input.
			var mname string
			if fn == "value" {
				mname = sanitize(point.Name())
			} else {
				mname = sanitize(fmt.Sprintf("%s_%s", point.Name(), fn))
			}

			var fam *MetricFamily
			var ok bool
			if fam, ok = p.fam[mname]; !ok {
				fam = &MetricFamily{
					Samples:   make(map[SampleID]*Sample),
					ValueType: vt,
					LabelSet:  make(map[string]int),
				}
				p.fam[mname] = fam
			} else {
				if fam.ValueType != vt {
					// Don't return an error since this would be a permanent error
					log.Printf("Mixed ValueType for measurement %q; dropping point", point.Name())
					break
				}
			}

			for k, _ := range sample.Labels {
				fam.LabelSet[k]++
			}

			fam.Samples[sampleID] = sample
		}
	}
	return nil
}

func init() {
	outputs.Add("prometheus_client", func() telegraf.Output {
		return &PrometheusClient{
			ExpirationInterval: internal.Duration{Duration: time.Second * 60},
			fam:                make(map[string]*MetricFamily),
			now:                time.Now,
		}
	})
}
<commit_msg>Fix prometheus output cannot be reloaded (#3053)<commit_after>package prometheus_client

import (
	"context"
	"fmt"
	"log"
	"net/http"
	"regexp"
	"sort"
	"strings"
	"sync"
	"time"

	"github.com/influxdata/telegraf"
	"github.com/influxdata/telegraf/internal"
	"github.com/influxdata/telegraf/plugins/outputs"
	"github.com/prometheus/client_golang/prometheus"
)

var invalidNameCharRE = regexp.MustCompile(`[^a-zA-Z0-9_]`)

// SampleID uniquely identifies a Sample
type SampleID string

// Sample represents the current value of a series.
type Sample struct {
	// Labels are the Prometheus labels.
	Labels map[string]string
	// Value is the value in the Prometheus output.
	Value float64
	// Expiration is the deadline that this Sample is valid until.
	Expiration time.Time
}

// MetricFamily contains the data required to build valid prometheus Metrics.
type MetricFamily struct {
	// Samples are the Sample belonging to this MetricFamily.
	Samples map[SampleID]*Sample
	// Type of the Value.
	ValueType prometheus.ValueType
	// LabelSet is the label counts for all Samples.
	LabelSet map[string]int
}

type PrometheusClient struct {
	Listen             string
	ExpirationInterval internal.Duration `toml:"expiration_interval"`

	server *http.Server

	sync.Mutex
	// fam is the non-expired MetricFamily by Prometheus metric name.
	fam map[string]*MetricFamily
	// now returns the current time.
	now func() time.Time
}

var sampleConfig = `
  ## Address to listen on
  # listen = ":9273"

  ## Interval to expire metrics and not deliver to prometheus, 0 == no expiration
  # expiration_interval = "60s"
`

func (p *PrometheusClient) Start() error {
	prometheus.Register(p)

	if p.Listen == "" {
		p.Listen = "localhost:9273"
	}

	mux := http.NewServeMux()
	mux.Handle("/metrics", prometheus.Handler())

	p.server = &http.Server{
		Addr:    p.Listen,
		Handler: mux,
	}

	go func() {
		if err := p.server.ListenAndServe(); err != nil {
			if err != http.ErrServerClosed {
				log.Printf("E! Error creating prometheus metric endpoint, err: %s\n",
					err.Error())
			}
		}
	}()
	return nil
}

func (p *PrometheusClient) Stop() {
	// plugin gets cleaned up in Close() already.
}

func (p *PrometheusClient) Connect() error {
	// This service output does not need to make any further connections
	return nil
}

func (p *PrometheusClient) Close() error {
	ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
	defer cancel()
	err := p.server.Shutdown(ctx)
	prometheus.Unregister(p)
	return err
}

func (p *PrometheusClient) SampleConfig() string {
	return sampleConfig
}

func (p *PrometheusClient) Description() string {
	return "Configuration for the Prometheus client to spawn"
}

// Implements prometheus.Collector
func (p *PrometheusClient) Describe(ch chan<- *prometheus.Desc) {
	prometheus.NewGauge(prometheus.GaugeOpts{Name: "Dummy", Help: "Dummy"}).Describe(ch)
}

// Expire removes Samples that have expired.
func (p *PrometheusClient) Expire() {
	now := p.now()
	for name, family := range p.fam {
		for key, sample := range family.Samples {
			if p.ExpirationInterval.Duration != 0 && now.After(sample.Expiration) {
				for k, _ := range sample.Labels {
					family.LabelSet[k]--
				}
				delete(family.Samples, key)

				if len(family.Samples) == 0 {
					delete(p.fam, name)
				}
			}
		}
	}
}

// Collect implements prometheus.Collector
func (p *PrometheusClient) Collect(ch chan<- prometheus.Metric) {
	p.Lock()
	defer p.Unlock()

	p.Expire()

	for name, family := range p.fam {
		// Get list of all labels on MetricFamily
		var labelNames []string
		for k, v := range family.LabelSet {
			if v > 0 {
				labelNames = append(labelNames, k)
			}
		}
		desc := prometheus.NewDesc(name, "Telegraf collected metric", labelNames, nil)

		for _, sample := range family.Samples {
			// Get labels for this sample; unset labels will be set to the
			// empty string
			var labels []string
			for _, label := range labelNames {
				v := sample.Labels[label]
				labels = append(labels, v)
			}

			metric, err := prometheus.NewConstMetric(desc, family.ValueType, sample.Value, labels...)
			if err != nil {
				log.Printf("E! Error creating prometheus metric, "+
					"key: %s, labels: %v,\nerr: %s\n",
					name, labels, err.Error())
			}

			ch <- metric
		}
	}
}

func sanitize(value string) string {
	return invalidNameCharRE.ReplaceAllString(value, "_")
}

func valueType(tt telegraf.ValueType) prometheus.ValueType {
	switch tt {
	case telegraf.Counter:
		return prometheus.CounterValue
	case telegraf.Gauge:
		return prometheus.GaugeValue
	default:
		return prometheus.UntypedValue
	}
}

// CreateSampleID creates a SampleID based on the tags of a telegraf.Metric.
func CreateSampleID(tags map[string]string) SampleID {
	pairs := make([]string, 0, len(tags))
	for k, v := range tags {
		pairs = append(pairs, fmt.Sprintf("%s=%s", k, v))
	}
	sort.Strings(pairs)
	return SampleID(strings.Join(pairs, ","))
}

func (p *PrometheusClient) Write(metrics []telegraf.Metric) error {
	p.Lock()
	defer p.Unlock()

	now := p.now()

	for _, point := range metrics {
		tags := point.Tags()
		vt := valueType(point.Type())
		sampleID := CreateSampleID(tags)

		labels := make(map[string]string)
		for k, v := range tags {
			labels[sanitize(k)] = v
		}

		for fn, fv := range point.Fields() {
			// Ignore string and bool fields.
			var value float64
			switch fv := fv.(type) {
			case int64:
				value = float64(fv)
			case float64:
				value = fv
			default:
				continue
			}

			sample := &Sample{
				Labels:     labels,
				Value:      value,
				Expiration: now.Add(p.ExpirationInterval.Duration),
			}

			// Special handling of value field; supports passthrough from
			// the prometheus input.
			var mname string
			if fn == "value" {
				mname = sanitize(point.Name())
			} else {
				mname = sanitize(fmt.Sprintf("%s_%s", point.Name(), fn))
			}

			var fam *MetricFamily
			var ok bool
			if fam, ok = p.fam[mname]; !ok {
				fam = &MetricFamily{
					Samples:   make(map[SampleID]*Sample),
					ValueType: vt,
					LabelSet:  make(map[string]int),
				}
				p.fam[mname] = fam
			} else {
				if fam.ValueType != vt {
					// Don't return an error since this would be a permanent error
					log.Printf("Mixed ValueType for measurement %q; dropping point", point.Name())
					break
				}
			}

			for k, _ := range sample.Labels {
				fam.LabelSet[k]++
			}

			fam.Samples[sampleID] = sample
		}
	}
	return nil
}

func init() {
	outputs.Add("prometheus_client", func() telegraf.Output {
		return &PrometheusClient{
			ExpirationInterval: internal.Duration{Duration: time.Second * 60},
			fam:                make(map[string]*MetricFamily),
			now:                time.Now,
		}
	})
}
<|endoftext|>


<commit_before>package lib

import (
	"crypto/hmac"
	"crypto/sha1"
	"encoding/base64"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"strings"
	"time"
)

const (
	//ENDPOINT   string = "oss-cn-beijing.aliyuncs.com"
	ENDPOINT   string = "oss-cn-beijing-internal.aliyuncs.com"
	GMT_LAYOUT string = "Mon, 02 Jan 2006 15:04:05 GMT"

	METHOD_GET    string = "GET"
	METHOD_PUT    string = "PUT"
	METHOD_DELETE string = "DELETE"
	METHOD_POST   string = "POST"
)

var (
	logger func(interface{})
)

type oss struct {
	access_id   string
	access_key  string
	urlString   string
	method      string
	bucket      string
	contentType string
	resource    string
	content     io.Reader
	req         *http.Request
	ossHeaders  map[string]string
}

func New(access_id, access_key string) *oss {
	return &oss{access_id: access_id, access_key: access_key, ossHeaders: make(map[string]string)}
}

func (o *oss) SetBucket(bucket string) {
	o.bucket = bucket
}

func (o *oss) SetMethod(method string) {
	o.method = method
}

func (o *oss) SetContentType(contentType string) {
	o.contentType = contentType
}

func (o *oss) SetResource(resource string) {
	o.resource = resource
}

func (o *oss) SetContent(content string) {
	o.content = strings.NewReader(content)
}

func (o *oss) SetOSSHeader(key, value string) {
	o.ossHeaders[key] = value
}

func (o *oss) PrepReq() {
	o.urlString = fmt.Sprintf("http://%s%s%s", ENDPOINT, o.bucket, o.resource)
	req, err := http.NewRequest(o.method, o.urlString, o.content)
	if err != nil {
		logger(err)
	}
	location, err := time.LoadLocation("GMT")
	now := time.Now().In(location).Format(GMT_LAYOUT)
	signature := generationSign(o.method, now, o.bucket, o.resource, o.access_key, o.contentType, o.ossHeaders)
	req.Header.Set("Date", now)
	req.Header.Set("Host", ENDPOINT)
	req.Header.Set("Authorization", fmt.Sprintf("OSS %s:%s", o.access_id, signature))
	req.Header.Set("Content-Type", o.contentType)

	for k, v := range o.ossHeaders {
		req.Header.Set(strings.ToUpper(k), v)
	}
	o.req = req
}

func (o *oss) Do() (int, http.Header, []byte) {
	client := &http.Client{}
	res, err := client.Do(o.req)
	if err != nil {
		logger(err)
	}
	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		logger(err)
	}

	if res.StatusCode != 200 {
		logger(body)
	}

	return res.StatusCode, res.Header, body
}

func generationSign(method, date, bucket, resource, access_key, content_type string, headers map[string]string) string {
	content_md5 := ""
	CanonicalizedOSSHeaders := ""
	for k, v := range headers {
		CanonicalizedOSSHeaders = fmt.Sprintf("%s%s:%s\n", CanonicalizedOSSHeaders, k, v)
	}
	CanonicalizedResource := fmt.Sprintf("%s%s", bucket, resource)

	sign := fmt.Sprintf("%s\n%s\n%s\n%s\n%s%s", method, content_md5, content_type, date, CanonicalizedOSSHeaders, CanonicalizedResource)

	hmacencoder := hmac.New(sha1.New, []byte(access_key))
	hmacencoder.Write([]byte(sign))

	signature := base64.StdEncoding.EncodeToString(hmacencoder.Sum(nil))
	return signature
}

var stdoss *oss

func Init(access_id, access_key string, mlogger func(interface{})) {
	stdoss = New(access_id, access_key)
	logger = mlogger
}

func MkDir(bucket, resource string) int {
	stdoss.SetOSSHeader("x-oss-acl", "public-read")
	stdoss.SetMethod(METHOD_PUT)
	stdoss.SetBucket(bucket)
	stdoss.SetContentType(MIMETYPE["default"])
	stdoss.SetContent("")
	stdoss.SetResource(fmt.Sprintf("%s/", resource))
	stdoss.PrepReq()
	status, _, _ := stdoss.Do()
	return status
}

func RmDir(bucket, resource string) int {
	stdoss.SetMethod(METHOD_DELETE)
	stdoss.SetBucket(bucket)
	stdoss.SetContentType(MIMETYPE["default"])
	stdoss.SetContent("")
	stdoss.SetResource(fmt.Sprintf("%s/", resource))
	stdoss.PrepReq()
	status, _, _ := stdoss.Do()
	return status
}

func Create(bucket, resource, content string) int {
	stdoss.SetOSSHeader("x-oss-acl", "public-read")
	stdoss.SetMethod(METHOD_PUT)
	stdoss.SetBucket(bucket)
	stdoss.SetContentType(MIMETYPE["jpg"])
	stdoss.SetContent(content)
	stdoss.SetResource(resource)
	stdoss.PrepReq()
	status, _, _ := stdoss.Do()
	return status
}

func Remove(bucket, resource string) int {
	stdoss.SetMethod(METHOD_DELETE)
	stdoss.SetBucket(bucket)
	stdoss.SetContentType(MIMETYPE["default"])
	stdoss.SetContent("")
	stdoss.SetResource(resource)
	stdoss.PrepReq()
	status, _, _ := stdoss.Do()
	return status
}

func CreateBucket(bucket string) int {
	stdoss.SetOSSHeader("x-oss-acl", "public-read")
	stdoss.SetMethod(METHOD_PUT)
	stdoss.SetBucket(bucket)
	stdoss.SetContentType(MIMETYPE["default"])
	stdoss.SetContent("")
	stdoss.SetResource("")
	stdoss.PrepReq()
	status, _, _ := stdoss.Do()
	return status
}

func RemoveBucket(bucket string) int {
	stdoss.SetMethod(METHOD_DELETE)
	stdoss.SetBucket(bucket)
	stdoss.SetContentType(MIMETYPE["default"])
	stdoss.SetContent("")
	stdoss.SetResource("")
	stdoss.PrepReq()
	status, _, _ := stdoss.Do()
	return status
}
<commit_msg>debug<commit_after>package lib

import (
	"crypto/hmac"
	"crypto/sha1"
	"encoding/base64"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"strings"
	"time"
)

const (
	//ENDPOINT string = "oss-cn-beijing.aliyuncs.com"
	ENDPOINT   string = "oss-cn-beijing-internal.aliyuncs.com"
	GMT_LAYOUT string = "Mon, 02 Jan 2006 15:04:05 GMT"

	METHOD_GET    string = "GET"
	METHOD_PUT    string = "PUT"
	METHOD_DELETE string = "DELETE"
	METHOD_POST   string = "POST"
)

var (
	logger func(interface{})
)

type oss struct {
	access_id   string
	access_key  string
	urlString   string
	method      string
	bucket      string
	contentType string
	resource    string
	content     io.Reader
	req         *http.Request
	ossHeaders  map[string]string
}

func New(access_id, access_key string) *oss {
	return &oss{access_id: access_id, access_key: access_key, ossHeaders: make(map[string]string)}
}

func (o *oss) SetBucket(bucket string) {
	o.bucket = bucket
}

func (o *oss) SetMethod(method string) {
	o.method = method
}

func (o *oss) SetContentType(contentType string) {
	o.contentType = contentType
}

func (o *oss) SetResource(resource string) {
	o.resource = resource
}

func (o *oss) SetContent(content string) {
	o.content = strings.NewReader(content)
}

func (o *oss) SetOSSHeader(key, value string) {
	o.ossHeaders[key] = value
}

func (o *oss) PrepReq() {
	o.urlString = fmt.Sprintf("http://%s%s%s", ENDPOINT, o.bucket, o.resource)
	req, err := http.NewRequest(o.method, o.urlString, o.content)
	if err != nil {
		logger(err)
	}
	location, err := time.LoadLocation("GMT")
	now := time.Now().In(location).Format(GMT_LAYOUT)
	signature := generationSign(o.method, now, o.bucket, o.resource, o.access_key, o.contentType, o.ossHeaders)
	req.Header.Set("Date", now)
	req.Header.Set("Host", ENDPOINT)
	req.Header.Set("Authorization", fmt.Sprintf("OSS %s:%s", o.access_id, signature))
	req.Header.Set("Content-Type", o.contentType)

	for k, v := range o.ossHeaders {
		req.Header.Set(strings.ToUpper(k), v)
	}
	o.req = req
}

func (o *oss) Do() (int, http.Header, []byte) {
	client := &http.Client{}
	res, err := client.Do(o.req)
	if err != nil {
		logger(err)
		return http.StatusInternalServerError, make(http.Header), nil
	}
	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		logger(err)
	}

	if res.StatusCode != 200 {
		logger(body)
	}

	return res.StatusCode, res.Header, body
}

func generationSign(method, date, bucket, resource, access_key, content_type string, headers map[string]string) string {
	content_md5 := ""
	CanonicalizedOSSHeaders := ""
	for k, v := range headers {
		CanonicalizedOSSHeaders = fmt.Sprintf("%s%s:%s\n", CanonicalizedOSSHeaders, k, v)
	}
	CanonicalizedResource := fmt.Sprintf("%s%s", bucket, resource)

	sign := fmt.Sprintf("%s\n%s\n%s\n%s\n%s%s", method, content_md5, content_type, date, CanonicalizedOSSHeaders, CanonicalizedResource)

	hmacencoder := hmac.New(sha1.New, []byte(access_key))
	hmacencoder.Write([]byte(sign))

	signature := base64.StdEncoding.EncodeToString(hmacencoder.Sum(nil))
	return signature
}

var stdoss *oss

func Init(access_id, access_key string, mlogger func(interface{})) {
	stdoss = New(access_id, access_key)
	logger = mlogger
}

func MkDir(bucket, resource string) int {
	stdoss.SetOSSHeader("x-oss-acl", "public-read")
	stdoss.SetMethod(METHOD_PUT)
	stdoss.SetBucket(bucket)
	stdoss.SetContentType(MIMETYPE["default"])
	stdoss.SetContent("")
	stdoss.SetResource(fmt.Sprintf("%s/", resource))
	stdoss.PrepReq()
	status, _, _ := stdoss.Do()
	return status
}

func RmDir(bucket, resource string) int {
	stdoss.SetMethod(METHOD_DELETE)
	stdoss.SetBucket(bucket)
	stdoss.SetContentType(MIMETYPE["default"])
	stdoss.SetContent("")
	stdoss.SetResource(fmt.Sprintf("%s/", resource))
	stdoss.PrepReq()
	status, _, _ := stdoss.Do()
	return status
}

func Create(bucket, resource, content string) int {
	stdoss.SetOSSHeader("x-oss-acl", "public-read")
	stdoss.SetMethod(METHOD_PUT)
	stdoss.SetBucket(bucket)
	stdoss.SetContentType(MIMETYPE["jpg"])
	stdoss.SetContent(content)
	stdoss.SetResource(resource)
	stdoss.PrepReq()
	status, _, _ := stdoss.Do()
	return status
}

func Remove(bucket, resource string) int {
	stdoss.SetMethod(METHOD_DELETE)
	stdoss.SetBucket(bucket)
	stdoss.SetContentType(MIMETYPE["default"])
	stdoss.SetContent("")
	stdoss.SetResource(resource)
	stdoss.PrepReq()
	status, _, _ := stdoss.Do()
	return status
}

func CreateBucket(bucket string) int {
	stdoss.SetOSSHeader("x-oss-acl", "public-read")
	stdoss.SetMethod(METHOD_PUT)
	stdoss.SetBucket(bucket)
	stdoss.SetContentType(MIMETYPE["default"])
	stdoss.SetContent("")
	stdoss.SetResource("")
	stdoss.PrepReq()
	status, _, _ := stdoss.Do()
	return status
}

func RemoveBucket(bucket string) int {
	stdoss.SetMethod(METHOD_DELETE)
	stdoss.SetBucket(bucket)
	stdoss.SetContentType(MIMETYPE["default"])
	stdoss.SetContent("")
	stdoss.SetResource("")
	stdoss.PrepReq()
	status, _, _ := stdoss.Do()
	return status
}
<|endoftext|>


<commit_before>// Copyright 2017 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.

package gps

import (
	"bytes"
	"context"
	"fmt"
	"os/exec"
	"sync"
	"sync/atomic"
	"time"

	"github.com/Masterminds/vcs"
)

// monitoredCmd wraps a cmd and will keep monitoring the process until it
// finishes, the provided context is canceled, or a certain amount of time has
// passed and the command showed no signs of activity.
type monitoredCmd struct {
	cmd     *exec.Cmd
	timeout time.Duration
	stdout  *activityBuffer
	stderr  *activityBuffer
}

// noProgressError indicates that the monitored process was terminated due to
// exceeding exceeding the progress timeout.
type noProgressError struct {
	timeout time.Duration
}

// killCmdError indicates that an error occurred while sending a kill signal to
// the monitored process.
type killCmdError struct {
	err error
}

func newMonitoredCmd(cmd *exec.Cmd, timeout time.Duration) *monitoredCmd {
	stdout, stderr := newActivityBuffer(), newActivityBuffer()
	cmd.Stdout, cmd.Stderr = stdout, stderr
	return &monitoredCmd{
		cmd:     cmd,
		timeout: timeout,
		stdout:  stdout,
		stderr:  stderr,
	}
}

// run will wait for the command to finish and return the error, if any. If the
// command does not show any progress, as indicated by writing to stdout or
// stderr, for more than the specified timeout, the process will be killed.
func (c *monitoredCmd) run(ctx context.Context) error {
	// Check for cancellation before even starting
	if ctx.Err() != nil {
		return ctx.Err()
	}

	err := c.cmd.Start()
	if err != nil {
		return err
	}

	ticker := time.NewTicker(c.timeout)
	defer ticker.Stop()

	// Atomic marker to track proc exit state. Guards against bad channel
	// select receive order, where a tick or context cancellation could come
	// in at the same time as process completion, but one of the former are
	// picked first; in such a case, cmd.Process could(?) be nil by the time we
	// call signal methods on it.
	var isDone *int32 = new(int32)
	done := make(chan error, 1)

	go func() {
		// Wait() can only be called once, so this must act as the completion
		// indicator for both normal *and* signal-induced termination.
		done <- c.cmd.Wait()
		atomic.CompareAndSwapInt32(isDone, 0, 1)
	}()

	var killerr error
selloop:
	for {
		select {
		case err := <-done:
			return err
		case <-ticker.C:
			if !atomic.CompareAndSwapInt32(isDone, 1, 1) && c.hasTimedOut() {
				if err := killProcess(c.cmd, isDone); err != nil {
					killerr = &killCmdError{err}
				} else {
					killerr = &noProgressError{c.timeout}
				}
				break selloop
			}
		case <-ctx.Done():
			if !atomic.CompareAndSwapInt32(isDone, 1, 1) {
				if err := killProcess(c.cmd, isDone); err != nil {
					killerr = &killCmdError{err}
				} else {
					killerr = ctx.Err()
				}
				break selloop
			}
		}
	}

	// This is only reachable on the signal-induced termination path, so block
	// until a message comes through the channel indicating that the command has
	// exited.
	//
	// TODO(sdboyer) if the signaling process errored (resulting in a
	// killCmdError stored in killerr), is it possible that this receive could
	// block forever on some kind of hung process?
	<-done
	return killerr
}

func (c *monitoredCmd) hasTimedOut() bool {
	t := time.Now().Add(-c.timeout)
	return c.stderr.lastActivity().Before(t) &&
		c.stdout.lastActivity().Before(t)
}

func (c *monitoredCmd) combinedOutput(ctx context.Context) ([]byte, error) {
	c.cmd.Stderr = c.stdout
	if err := c.run(ctx); err != nil {
		return c.stdout.Bytes(), err
	}

	return c.stdout.Bytes(), nil
}

// activityBuffer is a buffer that keeps track of the last time a Write
// operation was performed on it.
type activityBuffer struct {
	sync.Mutex
	buf               *bytes.Buffer
	lastActivityStamp time.Time
}

func newActivityBuffer() *activityBuffer {
	return &activityBuffer{
		buf: bytes.NewBuffer(nil),
	}
}

func (b *activityBuffer) Write(p []byte) (int, error) {
	b.Lock()
	defer b.Unlock()

	b.lastActivityStamp = time.Now()

	return b.buf.Write(p)
}

func (b *activityBuffer) String() string {
	b.Lock()
	defer b.Unlock()

	return b.buf.String()
}

func (b *activityBuffer) Bytes() []byte {
	b.Lock()
	defer b.Unlock()

	return b.buf.Bytes()
}

func (b *activityBuffer) lastActivity() time.Time {
	b.Lock()
	defer b.Unlock()

	return b.lastActivityStamp
}

func (e noProgressError) Error() string {
	return fmt.Sprintf("command killed after %s of no activity", e.timeout)
}

func (e killCmdError) Error() string {
	return fmt.Sprintf("error killing command: %s", e.err)
}

func runFromCwd(ctx context.Context, timeout time.Duration, cmd string, args ...string) ([]byte, error) {
	c := newMonitoredCmd(exec.Command(cmd, args...), timeout)
	return c.combinedOutput(ctx)
}

func runFromRepoDir(ctx context.Context, repo vcs.Repo, timeout time.Duration, cmd string, args ...string) ([]byte, error) {
	c := newMonitoredCmd(repo.CmdFromDir(cmd, args...), timeout)
	return c.combinedOutput(ctx)
}

const (
	// expensiveCmdTimeout is meant to be used in a command that is expensive
	// in terms of computation and we know it will take long or one that uses
	// the network, such as clones, updates, ....
	expensiveCmdTimeout = 2 * time.Minute
	// defaultCmdTimeout is just an umbrella value for all other commands that
	// should not take much.
	defaultCmdTimeout = 10 * time.Second
)
<commit_msg>Increase default subcommand timeout to 30s<commit_after>// Copyright 2017 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.

package gps

import (
	"bytes"
	"context"
	"fmt"
	"os/exec"
	"sync"
	"sync/atomic"
	"time"

	"github.com/Masterminds/vcs"
)

// monitoredCmd wraps a cmd and will keep monitoring the process until it
// finishes, the provided context is canceled, or a certain amount of time has
// passed and the command showed no signs of activity.
type monitoredCmd struct {
	cmd     *exec.Cmd
	timeout time.Duration
	stdout  *activityBuffer
	stderr  *activityBuffer
}

// noProgressError indicates that the monitored process was terminated due to
// exceeding exceeding the progress timeout.
type noProgressError struct {
	timeout time.Duration
}

// killCmdError indicates that an error occurred while sending a kill signal to
// the monitored process.
type killCmdError struct {
	err error
}

func newMonitoredCmd(cmd *exec.Cmd, timeout time.Duration) *monitoredCmd {
	stdout, stderr := newActivityBuffer(), newActivityBuffer()
	cmd.Stdout, cmd.Stderr = stdout, stderr
	return &monitoredCmd{
		cmd:     cmd,
		timeout: timeout,
		stdout:  stdout,
		stderr:  stderr,
	}
}

// run will wait for the command to finish and return the error, if any. If the
// command does not show any progress, as indicated by writing to stdout or
// stderr, for more than the specified timeout, the process will be killed.
func (c *monitoredCmd) run(ctx context.Context) error {
	// Check for cancellation before even starting
	if ctx.Err() != nil {
		return ctx.Err()
	}

	err := c.cmd.Start()
	if err != nil {
		return err
	}

	ticker := time.NewTicker(c.timeout)
	defer ticker.Stop()

	// Atomic marker to track proc exit state. Guards against bad channel
	// select receive order, where a tick or context cancellation could come
	// in at the same time as process completion, but one of the former are
	// picked first; in such a case, cmd.Process could(?) be nil by the time we
	// call signal methods on it.
	var isDone *int32 = new(int32)
	done := make(chan error, 1)

	go func() {
		// Wait() can only be called once, so this must act as the completion
		// indicator for both normal *and* signal-induced termination.
		done <- c.cmd.Wait()
		atomic.CompareAndSwapInt32(isDone, 0, 1)
	}()

	var killerr error
selloop:
	for {
		select {
		case err := <-done:
			return err
		case <-ticker.C:
			if !atomic.CompareAndSwapInt32(isDone, 1, 1) && c.hasTimedOut() {
				if err := killProcess(c.cmd, isDone); err != nil {
					killerr = &killCmdError{err}
				} else {
					killerr = &noProgressError{c.timeout}
				}
				break selloop
			}
		case <-ctx.Done():
			if !atomic.CompareAndSwapInt32(isDone, 1, 1) {
				if err := killProcess(c.cmd, isDone); err != nil {
					killerr = &killCmdError{err}
				} else {
					killerr = ctx.Err()
				}
				break selloop
			}
		}
	}

	// This is only reachable on the signal-induced termination path, so block
	// until a message comes through the channel indicating that the command has
	// exited.
	//
	// TODO(sdboyer) if the signaling process errored (resulting in a
	// killCmdError stored in killerr), is it possible that this receive could
	// block forever on some kind of hung process?
	<-done
	return killerr
}

func (c *monitoredCmd) hasTimedOut() bool {
	t := time.Now().Add(-c.timeout)
	return c.stderr.lastActivity().Before(t) &&
		c.stdout.lastActivity().Before(t)
}

func (c *monitoredCmd) combinedOutput(ctx context.Context) ([]byte, error) {
	c.cmd.Stderr = c.stdout
	if err := c.run(ctx); err != nil {
		return c.stdout.Bytes(), err
	}

	return c.stdout.Bytes(), nil
}

// activityBuffer is a buffer that keeps track of the last time a Write
// operation was performed on it.
type activityBuffer struct {
	sync.Mutex
	buf               *bytes.Buffer
	lastActivityStamp time.Time
}

func newActivityBuffer() *activityBuffer {
	return &activityBuffer{
		buf: bytes.NewBuffer(nil),
	}
}

func (b *activityBuffer) Write(p []byte) (int, error) {
	b.Lock()
	defer b.Unlock()

	b.lastActivityStamp = time.Now()

	return b.buf.Write(p)
}

func (b *activityBuffer) String() string {
	b.Lock()
	defer b.Unlock()

	return b.buf.String()
}

func (b *activityBuffer) Bytes() []byte {
	b.Lock()
	defer b.Unlock()

	return b.buf.Bytes()
}

func (b *activityBuffer) lastActivity() time.Time {
	b.Lock()
	defer b.Unlock()

	return b.lastActivityStamp
}

func (e noProgressError) Error() string {
	return fmt.Sprintf("command killed after %s of no activity", e.timeout)
}

func (e killCmdError) Error() string {
	return fmt.Sprintf("error killing command: %s", e.err)
}

func runFromCwd(ctx context.Context, timeout time.Duration, cmd string, args ...string) ([]byte, error) {
	c := newMonitoredCmd(exec.Command(cmd, args...), timeout)
	return c.combinedOutput(ctx)
}

func runFromRepoDir(ctx context.Context, repo vcs.Repo, timeout time.Duration, cmd string, args ...string) ([]byte, error) {
	c := newMonitoredCmd(repo.CmdFromDir(cmd, args...), timeout)
	return c.combinedOutput(ctx)
}

const (
	// expensiveCmdTimeout is meant to be used in a command that is expensive
	// in terms of computation and we know it will take long or one that uses
	// the network, such as clones, updates, ....
	expensiveCmdTimeout = 2 * time.Minute
	// defaultCmdTimeout is just an umbrella value for all other commands that
	// should not take much.
	defaultCmdTimeout = 30 * time.Second
)
<|endoftext|>


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

package oem

import (
	"fmt"
	"net/url"

	"github.com/coreos/ignition/config/types"
	"github.com/coreos/ignition/internal/log"
	"github.com/coreos/ignition/internal/providers"
	"github.com/coreos/ignition/internal/providers/azure"
	"github.com/coreos/ignition/internal/providers/digitalocean"
	"github.com/coreos/ignition/internal/providers/ec2"
	"github.com/coreos/ignition/internal/providers/file"
	"github.com/coreos/ignition/internal/providers/gce"
	"github.com/coreos/ignition/internal/providers/noop"
	"github.com/coreos/ignition/internal/providers/openstack"
	"github.com/coreos/ignition/internal/providers/packet"
	"github.com/coreos/ignition/internal/providers/qemu"
	"github.com/coreos/ignition/internal/providers/virtualbox"
	"github.com/coreos/ignition/internal/providers/vmware"
	"github.com/coreos/ignition/internal/registry"
	"github.com/coreos/ignition/internal/resource"

	"github.com/vincent-petithory/dataurl"
)

// Config represents a set of options that map to a particular OEM.
type Config struct {
	name              string
	fetch             providers.FuncFetchConfig
	newFetcher        providers.FuncNewFetcher
	baseConfig        types.Config
	defaultUserConfig types.Config
}

func (c Config) Name() string {
	return c.name
}

func (c Config) FetchFunc() providers.FuncFetchConfig {
	return c.fetch
}

func (c Config) NewFetcherFunc() providers.FuncNewFetcher {
	if c.newFetcher != nil {
		return c.newFetcher
	}
	return func(l *log.Logger, c *resource.HttpClient) (resource.Fetcher, error) {
		return resource.Fetcher{
			Logger: l,
			Client: c,
		}, nil
	}
}

func (c Config) BaseConfig() types.Config {
	return c.baseConfig
}

func (c Config) DefaultUserConfig() types.Config {
	return c.defaultUserConfig
}

var configs = registry.Create("oem configs")

func init() {
	configs.Register(Config{
		name:  "azure",
		fetch: azure.FetchConfig,
		baseConfig: types.Config{
			Systemd: types.Systemd{
				Units: []types.Unit{
					{Enable: true, Name: "waagent.service"},
					{Name: "etcd2.service", Dropins: []types.Dropin{
						{Name: "10-oem.conf", Contents: "[Service]\nEnvironment=ETCD_ELECTION_TIMEOUT=1200\n"},
					}},
				},
			},
			Storage: types.Storage{Files: []types.File{serviceFromOem("waagent.service")}},
		},
		defaultUserConfig: types.Config{Systemd: types.Systemd{Units: []types.Unit{userCloudInit("Azure", "azure")}}},
	})
	configs.Register(Config{
		name:  "cloudsigma",
		fetch: noop.FetchConfig,
	})
	configs.Register(Config{
		name:  "cloudstack",
		fetch: noop.FetchConfig,
	})
	configs.Register(Config{
		name:  "digitalocean",
		fetch: digitalocean.FetchConfig,
		baseConfig: types.Config{
			Systemd: types.Systemd{
				Units: []types.Unit{{Enable: true, Name: "coreos-metadata-sshkeys@.service"}},
			},
		},
		defaultUserConfig: types.Config{Systemd: types.Systemd{Units: []types.Unit{userCloudInit("DigitalOcean", "digitalocean")}}},
	})
	configs.Register(Config{
		name:  "brightbox",
		fetch: openstack.FetchConfig,
		defaultUserConfig: types.Config{
			Systemd: types.Systemd{
				Units: []types.Unit{
					{Mask: true, Name: "user-configdrive.service"},
					{Mask: true, Name: "user-configvirtfs.service"},
					userCloudInit("BrightBox", "ec2-compat"),
				},
			},
		},
	})
	configs.Register(Config{
		name:  "openstack",
		fetch: openstack.FetchConfig,
		defaultUserConfig: types.Config{
			Systemd: types.Systemd{
				Units: []types.Unit{
					{Mask: true, Name: "user-configdrive.service"},
					{Mask: true, Name: "user-configvirtfs.service"},
					userCloudInit("OpenStack", "ec2-compat"),
				},
			},
		},
	})
	configs.Register(Config{
		name:       "ec2",
		fetch:      ec2.FetchConfig,
		newFetcher: ec2.NewFetcher,
		baseConfig: types.Config{
			Systemd: types.Systemd{
				Units: []types.Unit{
					{Enable: true, Name: "coreos-metadata-sshkeys@.service"},
					{Name: "etcd2.service", Dropins: []types.Dropin{
						{Name: "10-oem.conf", Contents: "[Service]\nEnvironment=ETCD_ELECTION_TIMEOUT=1200\n"},
					}},
				},
			},
		},
		defaultUserConfig: types.Config{
			Systemd: types.Systemd{
				Units: []types.Unit{
					{Mask: true, Name: "user-configdrive.service"},
					{Mask: true, Name: "user-configvirtfs.service"},
					userCloudInit("EC2", "ec2-compat"),
				},
			},
		},
	})
	configs.Register(Config{
		name:  "exoscale",
		fetch: noop.FetchConfig,
	})
	configs.Register(Config{
		name:  "gce",
		fetch: gce.FetchConfig,
		baseConfig: types.Config{
			Systemd: types.Systemd{
				Units: []types.Unit{
					{Enable: true, Name: "coreos-metadata-sshkeys@.service"},
					{Enable: true, Name: "oem-gce.service"},
				},
			},
			Storage: types.Storage{
				Files: []types.File{
					serviceFromOem("oem-gce.service"),
					{
						Node: types.Node{
							Filesystem: "root",
							Path:       "/etc/hosts",
						},
						FileEmbedded1: types.FileEmbedded1{
							Mode:     0444,
							Contents: contentsFromString("169.254.169.254 metadata\n127.0.0.1 localhost\n"),
						},
					},
					{
						Node: types.Node{
							Filesystem: "root",
							Path:       "/etc/profile.d/google-cloud-sdk.sh",
						},
						FileEmbedded1: types.FileEmbedded1{
							Mode: 0444,
							Contents: contentsFromString(`#!/bin/sh
alias gcloud="(docker images google/cloud-sdk || docker pull google/cloud-sdk) > /dev/null;docker run -t -i --net="host" -v $HOME/.config:/.config -v /var/run/docker.sock:/var/run/doker.sock -v /usr/bin/docker:/usr/bin/docker google/cloud-sdk gcloud"
alias gcutil="(docker images google/cloud-sdk || docker pull google/cloud-sdk) > /dev/null;docker run -t -i --net="host" -v $HOME/.config:/.config google/cloud-sdk gcutil"
alias gsutil="(docker images google/cloud-sdk || docker pull google/cloud-sdk) > /dev/null;docker run -t -i --net="host" -v $HOME/.config:/.config google/cloud-sdk gsutil"
`),
						},
					},
				},
			},
		},
		defaultUserConfig: types.Config{Systemd: types.Systemd{Units: []types.Unit{userCloudInit("GCE", "gce")}}},
	})
	configs.Register(Config{
		name:  "hyperv",
		fetch: noop.FetchConfig,
	})
	configs.Register(Config{
		name:  "niftycloud",
		fetch: noop.FetchConfig,
	})
	configs.Register(Config{
		name:  "packet",
		fetch: packet.FetchConfig,
		baseConfig: types.Config{
			Systemd: types.Systemd{
				Units: []types.Unit{
					{Enable: true, Name: "coreos-metadata-sshkeys@.service"},
					{Enable: true, Name: "packet-phone-home.service"},
				},
			},
			Storage: types.Storage{Files: []types.File{serviceFromOem("packet-phone-home.service")}},
		},
		defaultUserConfig: types.Config{Systemd: types.Systemd{Units: []types.Unit{userCloudInit("Packet", "packet")}}},
	})
	configs.Register(Config{
		name:  "pxe",
		fetch: noop.FetchConfig,
	})
	configs.Register(Config{
		name:  "rackspace",
		fetch: noop.FetchConfig,
	})
	configs.Register(Config{
		name:  "rackspace-onmetal",
		fetch: noop.FetchConfig,
	})
	configs.Register(Config{
		name:  "vagrant",
		fetch: noop.FetchConfig,
	})
	configs.Register(Config{
		name:  "virtualbox",
		fetch: virtualbox.FetchConfig,
	})
	configs.Register(Config{
		name:  "vmware",
		fetch: vmware.FetchConfig,
		baseConfig: types.Config{
			Systemd: types.Systemd{Units: []types.Unit{{Enable: true, Name: "vmtoolsd.service"}}},
			Storage: types.Storage{Files: []types.File{serviceFromOem("vmtoolsd.service")}},
		},
		defaultUserConfig: types.Config{Systemd: types.Systemd{Units: []types.Unit{userCloudInit("VMware", "vmware")}}},
	})
	configs.Register(Config{
		name:  "interoute",
		fetch: noop.FetchConfig,
	})
	configs.Register(Config{
		name:  "qemu",
		fetch: qemu.FetchConfig,
	})
	configs.Register(Config{
		name:  "file",
		fetch: file.FetchConfig,
	})
}

func Get(name string) (config Config, ok bool) {
	config, ok = configs.Get(name).(Config)
	return
}

func MustGet(name string) Config {
	if config, ok := Get(name); ok {
		return config
	} else {
		panic(fmt.Sprintf("invalid OEM name %q provided", name))
	}
}

func Names() (names []string) {
	return configs.Names()
}

func contentsFromString(data string) types.FileContents {
	return types.FileContents{
		Source: (&url.URL{
			Scheme: "data",
			Opaque: "," + dataurl.EscapeString(data),
		}).String(),
	}
}

func contentsFromOem(path string) types.FileContents {
	return types.FileContents{
		Source: (&url.URL{
			Scheme: "oem",
			Path:   path,
		}).String(),
	}
}

func userCloudInit(name string, oem string) types.Unit {
	contents := `[Unit]
Description=Cloudinit from %s metadata

[Service]
Type=oneshot
ExecStart=/usr/bin/coreos-cloudinit --oem=%s

[Install]
WantedBy=multi-user.target
`

	return types.Unit{
		Name:     "oem-cloudinit.service",
		Enable:   true,
		Contents: fmt.Sprintf(contents, name, oem),
	}
}

func serviceFromOem(unit string) types.File {
	return types.File{
		Node: types.Node{
			Filesystem: "root",
			Path:       "/etc/systemd/system/" + unit,
		},
		FileEmbedded1: types.FileEmbedded1{
			Mode:     0444,
			Contents: contentsFromOem("/units/" + unit),
		},
	}
}
<commit_msg>oem: add vagrant-virtualbox oem-id<commit_after>// Copyright 2015 CoreOS, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package oem

import (
	"fmt"
	"net/url"

	"github.com/coreos/ignition/config/types"
	"github.com/coreos/ignition/internal/log"
	"github.com/coreos/ignition/internal/providers"
	"github.com/coreos/ignition/internal/providers/azure"
	"github.com/coreos/ignition/internal/providers/digitalocean"
	"github.com/coreos/ignition/internal/providers/ec2"
	"github.com/coreos/ignition/internal/providers/file"
	"github.com/coreos/ignition/internal/providers/gce"
	"github.com/coreos/ignition/internal/providers/noop"
	"github.com/coreos/ignition/internal/providers/openstack"
	"github.com/coreos/ignition/internal/providers/packet"
	"github.com/coreos/ignition/internal/providers/qemu"
	"github.com/coreos/ignition/internal/providers/virtualbox"
	"github.com/coreos/ignition/internal/providers/vmware"
	"github.com/coreos/ignition/internal/registry"
	"github.com/coreos/ignition/internal/resource"

	"github.com/vincent-petithory/dataurl"
)

// Config represents a set of options that map to a particular OEM.
type Config struct {
	name              string
	fetch             providers.FuncFetchConfig
	newFetcher        providers.FuncNewFetcher
	baseConfig        types.Config
	defaultUserConfig types.Config
}

func (c Config) Name() string {
	return c.name
}

func (c Config) FetchFunc() providers.FuncFetchConfig {
	return c.fetch
}

func (c Config) NewFetcherFunc() providers.FuncNewFetcher {
	if c.newFetcher != nil {
		return c.newFetcher
	}
	return func(l *log.Logger, c *resource.HttpClient) (resource.Fetcher, error) {
		return resource.Fetcher{
			Logger: l,
			Client: c,
		}, nil
	}
}

func (c Config) BaseConfig() types.Config {
	return c.baseConfig
}

func (c Config) DefaultUserConfig() types.Config {
	return c.defaultUserConfig
}

var configs = registry.Create("oem configs")

func init() {
	configs.Register(Config{
		name:  "azure",
		fetch: azure.FetchConfig,
		baseConfig: types.Config{
			Systemd: types.Systemd{
				Units: []types.Unit{
					{Enable: true, Name: "waagent.service"},
					{Name: "etcd2.service", Dropins: []types.Dropin{
						{Name: "10-oem.conf", Contents: "[Service]\nEnvironment=ETCD_ELECTION_TIMEOUT=1200\n"},
					}},
				},
			},
			Storage: types.Storage{Files: []types.File{serviceFromOem("waagent.service")}},
		},
		defaultUserConfig: types.Config{Systemd: types.Systemd{Units: []types.Unit{userCloudInit("Azure", "azure")}}},
	})
	configs.Register(Config{
		name:  "cloudsigma",
		fetch: noop.FetchConfig,
	})
	configs.Register(Config{
		name:  "cloudstack",
		fetch: noop.FetchConfig,
	})
	configs.Register(Config{
		name:  "digitalocean",
		fetch: digitalocean.FetchConfig,
		baseConfig: types.Config{
			Systemd: types.Systemd{
				Units: []types.Unit{{Enable: true, Name: "coreos-metadata-sshkeys@.service"}},
			},
		},
		defaultUserConfig: types.Config{Systemd: types.Systemd{Units: []types.Unit{userCloudInit("DigitalOcean", "digitalocean")}}},
	})
	configs.Register(Config{
		name:  "brightbox",
		fetch: openstack.FetchConfig,
		defaultUserConfig: types.Config{
			Systemd: types.Systemd{
				Units: []types.Unit{
					{Mask: true, Name: "user-configdrive.service"},
					{Mask: true, Name: "user-configvirtfs.service"},
					userCloudInit("BrightBox", "ec2-compat"),
				},
			},
		},
	})
	configs.Register(Config{
		name:  "openstack",
		fetch: openstack.FetchConfig,
		defaultUserConfig: types.Config{
			Systemd: types.Systemd{
				Units: []types.Unit{
					{Mask: true, Name: "user-configdrive.service"},
					{Mask: true, Name: "user-configvirtfs.service"},
					userCloudInit("OpenStack", "ec2-compat"),
				},
			},
		},
	})
	configs.Register(Config{
		name:       "ec2",
		fetch:      ec2.FetchConfig,
		newFetcher: ec2.NewFetcher,
		baseConfig: types.Config{
			Systemd: types.Systemd{
				Units: []types.Unit{
					{Enable: true, Name: "coreos-metadata-sshkeys@.service"},
					{Name: "etcd2.service", Dropins: []types.Dropin{
						{Name: "10-oem.conf", Contents: "[Service]\nEnvironment=ETCD_ELECTION_TIMEOUT=1200\n"},
					}},
				},
			},
		},
		defaultUserConfig: types.Config{
			Systemd: types.Systemd{
				Units: []types.Unit{
					{Mask: true, Name: "user-configdrive.service"},
					{Mask: true, Name: "user-configvirtfs.service"},
					userCloudInit("EC2", "ec2-compat"),
				},
			},
		},
	})
	configs.Register(Config{
		name:  "exoscale",
		fetch: noop.FetchConfig,
	})
	configs.Register(Config{
		name:  "gce",
		fetch: gce.FetchConfig,
		baseConfig: types.Config{
			Systemd: types.Systemd{
				Units: []types.Unit{
					{Enable: true, Name: "coreos-metadata-sshkeys@.service"},
					{Enable: true, Name: "oem-gce.service"},
				},
			},
			Storage: types.Storage{
				Files: []types.File{
					serviceFromOem("oem-gce.service"),
					{
						Node: types.Node{
							Filesystem: "root",
							Path:       "/etc/hosts",
						},
						FileEmbedded1: types.FileEmbedded1{
							Mode:     0444,
							Contents: contentsFromString("169.254.169.254 metadata\n127.0.0.1 localhost\n"),
						},
					},
					{
						Node: types.Node{
							Filesystem: "root",
							Path:       "/etc/profile.d/google-cloud-sdk.sh",
						},
						FileEmbedded1: types.FileEmbedded1{
							Mode: 0444,
							Contents: contentsFromString(`#!/bin/sh
alias gcloud="(docker images google/cloud-sdk || docker pull google/cloud-sdk) > /dev/null;docker run -t -i --net="host" -v $HOME/.config:/.config -v /var/run/docker.sock:/var/run/doker.sock -v /usr/bin/docker:/usr/bin/docker google/cloud-sdk gcloud"
alias gcutil="(docker images google/cloud-sdk || docker pull google/cloud-sdk) > /dev/null;docker run -t -i --net="host" -v $HOME/.config:/.config google/cloud-sdk gcutil"
alias gsutil="(docker images google/cloud-sdk || docker pull google/cloud-sdk) > /dev/null;docker run -t -i --net="host" -v $HOME/.config:/.config google/cloud-sdk gsutil"
`),
						},
					},
				},
			},
		},
		defaultUserConfig: types.Config{Systemd: types.Systemd{Units: []types.Unit{userCloudInit("GCE", "gce")}}},
	})
	configs.Register(Config{
		name:  "hyperv",
		fetch: noop.FetchConfig,
	})
	configs.Register(Config{
		name:  "niftycloud",
		fetch: noop.FetchConfig,
	})
	configs.Register(Config{
		name:  "packet",
		fetch: packet.FetchConfig,
		baseConfig: types.Config{
			Systemd: types.Systemd{
				Units: []types.Unit{
					{Enable: true, Name: "coreos-metadata-sshkeys@.service"},
					{Enable: true, Name: "packet-phone-home.service"},
				},
			},
			Storage: types.Storage{Files: []types.File{serviceFromOem("packet-phone-home.service")}},
		},
		defaultUserConfig: types.Config{Systemd: types.Systemd{Units: []types.Unit{userCloudInit("Packet", "packet")}}},
	})
	configs.Register(Config{
		name:  "pxe",
		fetch: noop.FetchConfig,
	})
	configs.Register(Config{
		name:  "rackspace",
		fetch: noop.FetchConfig,
	})
	configs.Register(Config{
		name:  "rackspace-onmetal",
		fetch: noop.FetchConfig,
	})
	configs.Register(Config{
		name:  "vagrant",
		fetch: noop.FetchConfig,
	})
	configs.Register(Config{
		name:  "vagrant-virtualbox",
		fetch: virtualbox.FetchConfig,
	})
	configs.Register(Config{
		name:  "virtualbox",
		fetch: virtualbox.FetchConfig,
	})
	configs.Register(Config{
		name:  "vmware",
		fetch: vmware.FetchConfig,
		baseConfig: types.Config{
			Systemd: types.Systemd{Units: []types.Unit{{Enable: true, Name: "vmtoolsd.service"}}},
			Storage: types.Storage{Files: []types.File{serviceFromOem("vmtoolsd.service")}},
		},
		defaultUserConfig: types.Config{Systemd: types.Systemd{Units: []types.Unit{userCloudInit("VMware", "vmware")}}},
	})
	configs.Register(Config{
		name:  "interoute",
		fetch: noop.FetchConfig,
	})
	configs.Register(Config{
		name:  "qemu",
		fetch: qemu.FetchConfig,
	})
	configs.Register(Config{
		name:  "file",
		fetch: file.FetchConfig,
	})
}

func Get(name string) (config Config, ok bool) {
	config, ok = configs.Get(name).(Config)
	return
}

func MustGet(name string) Config {
	if config, ok := Get(name); ok {
		return config
	} else {
		panic(fmt.Sprintf("invalid OEM name %q provided", name))
	}
}

func Names() (names []string) {
	return configs.Names()
}

func contentsFromString(data string) types.FileContents {
	return types.FileContents{
		Source: (&url.URL{
			Scheme: "data",
			Opaque: "," + dataurl.EscapeString(data),
		}).String(),
	}
}

func contentsFromOem(path string) types.FileContents {
	return types.FileContents{
		Source: (&url.URL{
			Scheme: "oem",
			Path:   path,
		}).String(),
	}
}

func userCloudInit(name string, oem string) types.Unit {
	contents := `[Unit]
Description=Cloudinit from %s metadata

[Service]
Type=oneshot
ExecStart=/usr/bin/coreos-cloudinit --oem=%s

[Install]
WantedBy=multi-user.target
`

	return types.Unit{
		Name:     "oem-cloudinit.service",
		Enable:   true,
		Contents: fmt.Sprintf(contents, name, oem),
	}
}

func serviceFromOem(unit string) types.File {
	return types.File{
		Node: types.Node{
			Filesystem: "root",
			Path:       "/etc/systemd/system/" + unit,
		},
		FileEmbedded1: types.FileEmbedded1{
			Mode:     0444,
			Contents: contentsFromOem("/units/" + unit),
		},
	}
}
<|endoftext|>