<commit_before>/*
Package prometheus provides middleware to add Prometheus metrics.

Example:
```
package main
import (
    "github.com/labstack/echo/v4"
    "github.com/labstack/echo-contrib/prometheus"
)
func main() {
    e := echo.New()
    // Enable metrics middleware
    p := prometheus.NewPrometheus("echo", nil)
    p.Use(e)

    e.Logger.Fatal(e.Start(":1323"))
}
```
*/
package prometheus

import (
	"bytes"
	"errors"
	"net/http"
	"os"
	"strconv"
	"time"

	"github.com/labstack/echo/v4"
	"github.com/labstack/echo/v4/middleware"
	"github.com/labstack/gommon/log"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"github.com/prometheus/common/expfmt"
)

var defaultMetricPath = "/metrics"
var defaultSubsystem = "echo"

// Standard default metrics
//	counter, counter_vec, gauge, gauge_vec,
//	histogram, histogram_vec, summary, summary_vec
var reqCnt = &Metric{
	ID:          "reqCnt",
	Name:        "requests_total",
	Description: "How many HTTP requests processed, partitioned by status code and HTTP method.",
	Type:        "counter_vec",
	Args:        []string{"code", "method", "host", "url"}}

var reqDur = &Metric{
	ID:          "reqDur",
	Name:        "request_duration_seconds",
	Description: "The HTTP request latencies in seconds.",
	Args:        []string{"code", "method", "url"},
	Type:        "histogram_vec"}

var resSz = &Metric{
	ID:          "resSz",
	Name:        "response_size_bytes",
	Description: "The HTTP response sizes in bytes.",
	Args:        []string{"code", "method", "url"},
	Type:        "histogram_vec"}

var reqSz = &Metric{
	ID:          "reqSz",
	Name:        "request_size_bytes",
	Description: "The HTTP request sizes in bytes.",
	Args:        []string{"code", "method", "url"},
	Type:        "histogram_vec"}

var standardMetrics = []*Metric{
	reqCnt,
	reqDur,
	resSz,
	reqSz,
}

/*
RequestCounterLabelMappingFunc is a function which can be supplied to the middleware to control
the cardinality of the request counter's "url" label, which might be required in some contexts.
For instance, if for a "/customer/:name" route you don't want to generate a time series for every
possible customer name, you could use this function:

func(c echo.Context) string {
	url := c.Request.URL.Path
	for _, p := range c.Params {
		if p.Key == "name" {
			url = strings.Replace(url, p.Value, ":name", 1)
			break
		}
	}
	return url
}

which would map "/customer/alice" and "/customer/bob" to their template "/customer/:name".
It can also be applied for the "Host" label
*/
type RequestCounterLabelMappingFunc func(c echo.Context) string

// Metric is a definition for the name, description, type, ID, and
// prometheus.Collector type (i.e. CounterVec, Summary, etc) of each metric
type Metric struct {
	MetricCollector prometheus.Collector
	ID              string
	Name            string
	Description     string
	Type            string
	Args            []string
}

// Prometheus contains the metrics gathered by the instance and its path
type Prometheus struct {
	reqCnt               *prometheus.CounterVec
	reqDur, reqSz, resSz *prometheus.HistogramVec
	router               *echo.Echo
	listenAddress        string
	Ppg                  PushGateway

	MetricsList []*Metric
	MetricsPath string
	Subsystem   string
	Skipper     middleware.Skipper

	RequestCounterURLLabelMappingFunc  RequestCounterLabelMappingFunc
	RequestCounterHostLabelMappingFunc RequestCounterLabelMappingFunc

	// Context string to use as a prometheus URL label
	URLLabelFromContext string
}

// PushGateway contains the configuration for pushing to a Prometheus pushgateway (optional)
type PushGateway struct {
	// Push interval in seconds
	//lint:ignore ST1011 renaming would be breaking change
	PushIntervalSeconds time.Duration

	// Push Gateway URL in format http://domain:port
	// where JOBNAME can be any string of your choice
	PushGatewayURL string

	// pushgateway job name, defaults to "echo"
	Job string
}

// NewPrometheus generates a new set of metrics with a certain subsystem name
func NewPrometheus(subsystem string, skipper middleware.Skipper, customMetricsList ...[]*Metric) *Prometheus {
	var metricsList []*Metric
	if skipper == nil {
		skipper = middleware.DefaultSkipper
	}

	if len(customMetricsList) > 1 {
		panic("Too many args. NewPrometheus( string, <optional []*Metric> ).")
	} else if len(customMetricsList) == 1 {
		metricsList = customMetricsList[0]
	}

	metricsList = append(metricsList, standardMetrics...)

	p := &Prometheus{
		MetricsList: metricsList,
		MetricsPath: defaultMetricPath,
		Subsystem:   defaultSubsystem,
		Skipper:     skipper,
		RequestCounterURLLabelMappingFunc: func(c echo.Context) string {
			return c.Path() // i.e. by default do nothing, i.e. return URL as is
		},
		RequestCounterHostLabelMappingFunc: func(c echo.Context) string {
			return c.Request().Host
		},
	}

	p.registerMetrics(subsystem)

	return p
}

// SetPushGateway sends metrics to a remote pushgateway exposed on pushGatewayURL
// every pushInterval. Metrics are fetched from
func (p *Prometheus) SetPushGateway(pushGatewayURL string, pushInterval time.Duration) {
	p.Ppg.PushGatewayURL = pushGatewayURL
	p.Ppg.PushIntervalSeconds = pushInterval
	p.startPushTicker()
}

// SetPushGatewayJob job name, defaults to "echo"
func (p *Prometheus) SetPushGatewayJob(j string) {
	p.Ppg.Job = j
}

// SetListenAddress for exposing metrics on address. If not set, it will be exposed at the
// same address of the echo engine that is being used
// func (p *Prometheus) SetListenAddress(address string) {
// 	p.listenAddress = address
// 	if p.listenAddress != "" {
// 		p.router = echo.Echo().Router()
// 	}
// }

// SetListenAddressWithRouter for using a separate router to expose metrics. (this keeps things like GET /metrics out of
// your content's access log).
// func (p *Prometheus) SetListenAddressWithRouter(listenAddress string, r *echo.Echo) {
// 	p.listenAddress = listenAddress
// 	if len(p.listenAddress) > 0 {
// 		p.router = r
// 	}
// }

// SetMetricsPath set metrics paths
func (p *Prometheus) SetMetricsPath(e *echo.Echo) {
	if p.listenAddress != "" {
		p.router.GET(p.MetricsPath, prometheusHandler())
		p.runServer()
	} else {
		e.GET(p.MetricsPath, prometheusHandler())
	}
}

func (p *Prometheus) runServer() {
	if p.listenAddress != "" {
		go p.router.Start(p.listenAddress)
	}
}

func (p *Prometheus) getMetrics() []byte {
	out := &bytes.Buffer{}
	metricFamilies, _ := prometheus.DefaultGatherer.Gather()
	for i := range metricFamilies {
		expfmt.MetricFamilyToText(out, metricFamilies[i])

	}
	return out.Bytes()
}

func (p *Prometheus) getPushGatewayURL() string {
	h, _ := os.Hostname()
	if p.Ppg.Job == "" {
		p.Ppg.Job = "echo"
	}
	return p.Ppg.PushGatewayURL + "/metrics/job/" + p.Ppg.Job + "/instance/" + h
}

func (p *Prometheus) sendMetricsToPushGateway(metrics []byte) {
	req, err := http.NewRequest("POST", p.getPushGatewayURL(), bytes.NewBuffer(metrics))
	if err != nil {
		log.Errorf("failed to create push gateway request: %v", err)
		return
	}
	client := &http.Client{}
	if _, err = client.Do(req); err != nil {
		log.Errorf("Error sending to push gateway: %v", err)
	}
}

func (p *Prometheus) startPushTicker() {
	ticker := time.NewTicker(time.Second * p.Ppg.PushIntervalSeconds)
	go func() {
		for range ticker.C {
			p.sendMetricsToPushGateway(p.getMetrics())
		}
	}()
}

// NewMetric associates prometheus.Collector based on Metric.Type
func NewMetric(m *Metric, subsystem string) prometheus.Collector {
	var metric prometheus.Collector
	switch m.Type {
	case "counter_vec":
		metric = prometheus.NewCounterVec(
			prometheus.CounterOpts{
				Subsystem: subsystem,
				Name:      m.Name,
				Help:      m.Description,
			},
			m.Args,
		)
	case "counter":
		metric = prometheus.NewCounter(
			prometheus.CounterOpts{
				Subsystem: subsystem,
				Name:      m.Name,
				Help:      m.Description,
			},
		)
	case "gauge_vec":
		metric = prometheus.NewGaugeVec(
			prometheus.GaugeOpts{
				Subsystem: subsystem,
				Name:      m.Name,
				Help:      m.Description,
			},
			m.Args,
		)
	case "gauge":
		metric = prometheus.NewGauge(
			prometheus.GaugeOpts{
				Subsystem: subsystem,
				Name:      m.Name,
				Help:      m.Description,
			},
		)
	case "histogram_vec":
		metric = prometheus.NewHistogramVec(
			prometheus.HistogramOpts{
				Subsystem: subsystem,
				Name:      m.Name,
				Help:      m.Description,
			},
			m.Args,
		)
	case "histogram":
		metric = prometheus.NewHistogram(
			prometheus.HistogramOpts{
				Subsystem: subsystem,
				Name:      m.Name,
				Help:      m.Description,
			},
		)
	case "summary_vec":
		metric = prometheus.NewSummaryVec(
			prometheus.SummaryOpts{
				Subsystem: subsystem,
				Name:      m.Name,
				Help:      m.Description,
			},
			m.Args,
		)
	case "summary":
		metric = prometheus.NewSummary(
			prometheus.SummaryOpts{
				Subsystem: subsystem,
				Name:      m.Name,
				Help:      m.Description,
			},
		)
	}
	return metric
}

func (p *Prometheus) registerMetrics(subsystem string) {

	for _, metricDef := range p.MetricsList {
		metric := NewMetric(metricDef, subsystem)
		if err := prometheus.Register(metric); err != nil {
			log.Errorf("%s could not be registered in Prometheus: %v", metricDef.Name, err)
		}
		switch metricDef {
		case reqCnt:
			p.reqCnt = metric.(*prometheus.CounterVec)
		case reqDur:
			p.reqDur = metric.(*prometheus.HistogramVec)
		case resSz:
			p.resSz = metric.(*prometheus.HistogramVec)
		case reqSz:
			p.reqSz = metric.(*prometheus.HistogramVec)
		}
		metricDef.MetricCollector = metric
	}
}

// Use adds the middleware to the Echo engine.
func (p *Prometheus) Use(e *echo.Echo) {
	e.Use(p.HandlerFunc)
	p.SetMetricsPath(e)
}

// HandlerFunc defines handler function for middleware
func (p *Prometheus) HandlerFunc(next echo.HandlerFunc) echo.HandlerFunc {
	return func(c echo.Context) error {
		if c.Path() == p.MetricsPath {
			return next(c)
		}
		if p.Skipper(c) {
			return next(c)
		}

		start := time.Now()
		reqSz := computeApproximateRequestSize(c.Request())

		err := next(c)

		status := c.Response().Status
		if err != nil {
			var httpError *echo.HTTPError
			if errors.As(err, &httpError) {
				status = httpError.Code
			}
			if status == 0 || status == http.StatusOK {
				status = http.StatusInternalServerError
			}
		}

		elapsed := float64(time.Since(start)) / float64(time.Second)

		url := p.RequestCounterURLLabelMappingFunc(c)
		if len(p.URLLabelFromContext) > 0 {
			u := c.Get(p.URLLabelFromContext)
			if u == nil {
				u = "unknown"
			}
			url = u.(string)
		}

		statusStr := strconv.Itoa(status)
		p.reqDur.WithLabelValues(statusStr, c.Request().Method, url).Observe(elapsed)
		p.reqCnt.WithLabelValues(statusStr, c.Request().Method, p.RequestCounterHostLabelMappingFunc(c), url).Inc()
		p.reqSz.WithLabelValues(statusStr, c.Request().Method, url).Observe(float64(reqSz))

		resSz := float64(c.Response().Size)
		p.resSz.WithLabelValues(statusStr, c.Request().Method, url).Observe(resSz)

		return err
	}
}

func prometheusHandler() echo.HandlerFunc {
	h := promhttp.Handler()
	return func(c echo.Context) error {
		h.ServeHTTP(c.Response(), c.Request())
		return nil
	}
}

func computeApproximateRequestSize(r *http.Request) int {
	s := 0
	if r.URL != nil {
		s = len(r.URL.Path)
	}

	s += len(r.Method)
	s += len(r.Proto)
	for name, values := range r.Header {
		s += len(name)
		for _, value := range values {
			s += len(value)
		}
	}
	s += len(r.Host)

	// N.B. r.Form and r.MultipartForm are assumed to be included in r.URL.

	if r.ContentLength != -1 {
		s += int(r.ContentLength)
	}
	return s
}
<commit_msg>Add proper buckets for request/response sizes.<commit_after>/*
Package prometheus provides middleware to add Prometheus metrics.

Example:
```
package main
import (
    "github.com/labstack/echo/v4"
    "github.com/labstack/echo-contrib/prometheus"
)
func main() {
    e := echo.New()
    // Enable metrics middleware
    p := prometheus.NewPrometheus("echo", nil)
    p.Use(e)

    e.Logger.Fatal(e.Start(":1323"))
}
```
*/
package prometheus

import (
	"bytes"
	"errors"
	"net/http"
	"os"
	"strconv"
	"time"

	"github.com/labstack/echo/v4"
	"github.com/labstack/echo/v4/middleware"
	"github.com/labstack/gommon/log"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"github.com/prometheus/common/expfmt"
)

var defaultMetricPath = "/metrics"
var defaultSubsystem = "echo"

const (
	_          = iota // ignore first value by assigning to blank identifier
	KB float64 = 1 << (10 * iota)
	MB
	GB
	TB
)

// reqDurBuckets is the buckets for request duration. Here, we use the prometheus defaults
// which are for ~10s request length max: []float64{.005, .01, .025, .05, .1, .25, .5, 1, 2.5, 5, 10}
var reqDurBuckets = prometheus.DefBuckets

// reqSzBuckets is the buckets for request size. Here we define a spectrom from 1KB thru 1NB up to 10MB.
var reqSzBuckets = []float64{1.0 * KB, 2.0 * KB, 5.0 * KB, 10.0 * KB, 100 * KB, 500 * KB, 1.0 * MB, 2.5 * MB, 5.0 * MB, 10.0 * MB}

// resSzBuckets is the buckets for response size. Here we define a spectrom from 1KB thru 1NB up to 10MB.
var resSzBuckets = []float64{1.0 * KB, 2.0 * KB, 5.0 * KB, 10.0 * KB, 100 * KB, 500 * KB, 1.0 * MB, 2.5 * MB, 5.0 * MB, 10.0 * MB}

// Standard default metrics
//	counter, counter_vec, gauge, gauge_vec,
//	histogram, histogram_vec, summary, summary_vec
var reqCnt = &Metric{
	ID:          "reqCnt",
	Name:        "requests_total",
	Description: "How many HTTP requests processed, partitioned by status code and HTTP method.",
	Type:        "counter_vec",
	Args:        []string{"code", "method", "host", "url"}}

var reqDur = &Metric{
	ID:          "reqDur",
	Name:        "request_duration_seconds",
	Description: "The HTTP request latencies in seconds.",
	Args:        []string{"code", "method", "url"},
	Type:        "histogram_vec",
	Buckets:     reqDurBuckets}

var resSz = &Metric{
	ID:          "resSz",
	Name:        "response_size_bytes",
	Description: "The HTTP response sizes in bytes.",
	Args:        []string{"code", "method", "url"},
	Type:        "histogram_vec",
	Buckets:     resSzBuckets}

var reqSz = &Metric{
	ID:          "reqSz",
	Name:        "request_size_bytes",
	Description: "The HTTP request sizes in bytes.",
	Args:        []string{"code", "method", "url"},
	Type:        "histogram_vec",
	Buckets:     reqSzBuckets}

var standardMetrics = []*Metric{
	reqCnt,
	reqDur,
	resSz,
	reqSz,
}

/*
RequestCounterLabelMappingFunc is a function which can be supplied to the middleware to control
the cardinality of the request counter's "url" label, which might be required in some contexts.
For instance, if for a "/customer/:name" route you don't want to generate a time series for every
possible customer name, you could use this function:

func(c echo.Context) string {
	url := c.Request.URL.Path
	for _, p := range c.Params {
		if p.Key == "name" {
			url = strings.Replace(url, p.Value, ":name", 1)
			break
		}
	}
	return url
}

which would map "/customer/alice" and "/customer/bob" to their template "/customer/:name".
It can also be applied for the "Host" label
*/
type RequestCounterLabelMappingFunc func(c echo.Context) string

// Metric is a definition for the name, description, type, ID, and
// prometheus.Collector type (i.e. CounterVec, Summary, etc) of each metric
type Metric struct {
	MetricCollector prometheus.Collector
	ID              string
	Name            string
	Description     string
	Type            string
	Args            []string
	Buckets         []float64
}

// Prometheus contains the metrics gathered by the instance and its path
type Prometheus struct {
	reqCnt               *prometheus.CounterVec
	reqDur, reqSz, resSz *prometheus.HistogramVec
	router               *echo.Echo
	listenAddress        string
	Ppg                  PushGateway

	MetricsList []*Metric
	MetricsPath string
	Subsystem   string
	Skipper     middleware.Skipper

	RequestCounterURLLabelMappingFunc  RequestCounterLabelMappingFunc
	RequestCounterHostLabelMappingFunc RequestCounterLabelMappingFunc

	// Context string to use as a prometheus URL label
	URLLabelFromContext string
}

// PushGateway contains the configuration for pushing to a Prometheus pushgateway (optional)
type PushGateway struct {
	// Push interval in seconds
	//lint:ignore ST1011 renaming would be breaking change
	PushIntervalSeconds time.Duration

	// Push Gateway URL in format http://domain:port
	// where JOBNAME can be any string of your choice
	PushGatewayURL string

	// pushgateway job name, defaults to "echo"
	Job string
}

// NewPrometheus generates a new set of metrics with a certain subsystem name
func NewPrometheus(subsystem string, skipper middleware.Skipper, customMetricsList ...[]*Metric) *Prometheus {
	var metricsList []*Metric
	if skipper == nil {
		skipper = middleware.DefaultSkipper
	}

	if len(customMetricsList) > 1 {
		panic("Too many args. NewPrometheus( string, <optional []*Metric> ).")
	} else if len(customMetricsList) == 1 {
		metricsList = customMetricsList[0]
	}

	metricsList = append(metricsList, standardMetrics...)

	p := &Prometheus{
		MetricsList: metricsList,
		MetricsPath: defaultMetricPath,
		Subsystem:   defaultSubsystem,
		Skipper:     skipper,
		RequestCounterURLLabelMappingFunc: func(c echo.Context) string {
			return c.Path() // i.e. by default do nothing, i.e. return URL as is
		},
		RequestCounterHostLabelMappingFunc: func(c echo.Context) string {
			return c.Request().Host
		},
	}

	p.registerMetrics(subsystem)

	return p
}

// SetPushGateway sends metrics to a remote pushgateway exposed on pushGatewayURL
// every pushInterval. Metrics are fetched from
func (p *Prometheus) SetPushGateway(pushGatewayURL string, pushInterval time.Duration) {
	p.Ppg.PushGatewayURL = pushGatewayURL
	p.Ppg.PushIntervalSeconds = pushInterval
	p.startPushTicker()
}

// SetPushGatewayJob job name, defaults to "echo"
func (p *Prometheus) SetPushGatewayJob(j string) {
	p.Ppg.Job = j
}

// SetListenAddress for exposing metrics on address. If not set, it will be exposed at the
// same address of the echo engine that is being used
// func (p *Prometheus) SetListenAddress(address string) {
// 	p.listenAddress = address
// 	if p.listenAddress != "" {
// 		p.router = echo.Echo().Router()
// 	}
// }

// SetListenAddressWithRouter for using a separate router to expose metrics. (this keeps things like GET /metrics out of
// your content's access log).
// func (p *Prometheus) SetListenAddressWithRouter(listenAddress string, r *echo.Echo) {
// 	p.listenAddress = listenAddress
// 	if len(p.listenAddress) > 0 {
// 		p.router = r
// 	}
// }

// SetMetricsPath set metrics paths
func (p *Prometheus) SetMetricsPath(e *echo.Echo) {
	if p.listenAddress != "" {
		p.router.GET(p.MetricsPath, prometheusHandler())
		p.runServer()
	} else {
		e.GET(p.MetricsPath, prometheusHandler())
	}
}

func (p *Prometheus) runServer() {
	if p.listenAddress != "" {
		go p.router.Start(p.listenAddress)
	}
}

func (p *Prometheus) getMetrics() []byte {
	out := &bytes.Buffer{}
	metricFamilies, _ := prometheus.DefaultGatherer.Gather()
	for i := range metricFamilies {
		expfmt.MetricFamilyToText(out, metricFamilies[i])

	}
	return out.Bytes()
}

func (p *Prometheus) getPushGatewayURL() string {
	h, _ := os.Hostname()
	if p.Ppg.Job == "" {
		p.Ppg.Job = "echo"
	}
	return p.Ppg.PushGatewayURL + "/metrics/job/" + p.Ppg.Job + "/instance/" + h
}

func (p *Prometheus) sendMetricsToPushGateway(metrics []byte) {
	req, err := http.NewRequest("POST", p.getPushGatewayURL(), bytes.NewBuffer(metrics))
	if err != nil {
		log.Errorf("failed to create push gateway request: %v", err)
		return
	}
	client := &http.Client{}
	if _, err = client.Do(req); err != nil {
		log.Errorf("Error sending to push gateway: %v", err)
	}
}

func (p *Prometheus) startPushTicker() {
	ticker := time.NewTicker(time.Second * p.Ppg.PushIntervalSeconds)
	go func() {
		for range ticker.C {
			p.sendMetricsToPushGateway(p.getMetrics())
		}
	}()
}

// NewMetric associates prometheus.Collector based on Metric.Type
func NewMetric(m *Metric, subsystem string) prometheus.Collector {
	var metric prometheus.Collector
	switch m.Type {
	case "counter_vec":
		metric = prometheus.NewCounterVec(
			prometheus.CounterOpts{
				Subsystem: subsystem,
				Name:      m.Name,
				Help:      m.Description,
			},
			m.Args,
		)
	case "counter":
		metric = prometheus.NewCounter(
			prometheus.CounterOpts{
				Subsystem: subsystem,
				Name:      m.Name,
				Help:      m.Description,
			},
		)
	case "gauge_vec":
		metric = prometheus.NewGaugeVec(
			prometheus.GaugeOpts{
				Subsystem: subsystem,
				Name:      m.Name,
				Help:      m.Description,
			},
			m.Args,
		)
	case "gauge":
		metric = prometheus.NewGauge(
			prometheus.GaugeOpts{
				Subsystem: subsystem,
				Name:      m.Name,
				Help:      m.Description,
			},
		)
	case "histogram_vec":
		metric = prometheus.NewHistogramVec(
			prometheus.HistogramOpts{
				Subsystem: subsystem,
				Name:      m.Name,
				Help:      m.Description,
				Buckets:   m.Buckets,
			},
			m.Args,
		)
	case "histogram":
		metric = prometheus.NewHistogram(
			prometheus.HistogramOpts{
				Subsystem: subsystem,
				Name:      m.Name,
				Help:      m.Description,
				Buckets:   m.Buckets,
			},
		)
	case "summary_vec":
		metric = prometheus.NewSummaryVec(
			prometheus.SummaryOpts{
				Subsystem: subsystem,
				Name:      m.Name,
				Help:      m.Description,
			},
			m.Args,
		)
	case "summary":
		metric = prometheus.NewSummary(
			prometheus.SummaryOpts{
				Subsystem: subsystem,
				Name:      m.Name,
				Help:      m.Description,
			},
		)
	}
	return metric
}

func (p *Prometheus) registerMetrics(subsystem string) {

	for _, metricDef := range p.MetricsList {
		metric := NewMetric(metricDef, subsystem)
		if err := prometheus.Register(metric); err != nil {
			log.Errorf("%s could not be registered in Prometheus: %v", metricDef.Name, err)
		}
		switch metricDef {
		case reqCnt:
			p.reqCnt = metric.(*prometheus.CounterVec)
		case reqDur:
			p.reqDur = metric.(*prometheus.HistogramVec)
		case resSz:
			p.resSz = metric.(*prometheus.HistogramVec)
		case reqSz:
			p.reqSz = metric.(*prometheus.HistogramVec)
		}
		metricDef.MetricCollector = metric
	}
}

// Use adds the middleware to the Echo engine.
func (p *Prometheus) Use(e *echo.Echo) {
	e.Use(p.HandlerFunc)
	p.SetMetricsPath(e)
}

// HandlerFunc defines handler function for middleware
func (p *Prometheus) HandlerFunc(next echo.HandlerFunc) echo.HandlerFunc {
	return func(c echo.Context) error {
		if c.Path() == p.MetricsPath {
			return next(c)
		}
		if p.Skipper(c) {
			return next(c)
		}

		start := time.Now()
		reqSz := computeApproximateRequestSize(c.Request())

		err := next(c)

		status := c.Response().Status
		if err != nil {
			var httpError *echo.HTTPError
			if errors.As(err, &httpError) {
				status = httpError.Code
			}
			if status == 0 || status == http.StatusOK {
				status = http.StatusInternalServerError
			}
		}

		elapsed := float64(time.Since(start)) / float64(time.Second)

		url := p.RequestCounterURLLabelMappingFunc(c)
		if len(p.URLLabelFromContext) > 0 {
			u := c.Get(p.URLLabelFromContext)
			if u == nil {
				u = "unknown"
			}
			url = u.(string)
		}

		statusStr := strconv.Itoa(status)
		p.reqDur.WithLabelValues(statusStr, c.Request().Method, url).Observe(elapsed)
		p.reqCnt.WithLabelValues(statusStr, c.Request().Method, p.RequestCounterHostLabelMappingFunc(c), url).Inc()
		p.reqSz.WithLabelValues(statusStr, c.Request().Method, url).Observe(float64(reqSz))

		resSz := float64(c.Response().Size)
		p.resSz.WithLabelValues(statusStr, c.Request().Method, url).Observe(resSz)

		return err
	}
}

func prometheusHandler() echo.HandlerFunc {
	h := promhttp.Handler()
	return func(c echo.Context) error {
		h.ServeHTTP(c.Response(), c.Request())
		return nil
	}
}

func computeApproximateRequestSize(r *http.Request) int {
	s := 0
	if r.URL != nil {
		s = len(r.URL.Path)
	}

	s += len(r.Method)
	s += len(r.Proto)
	for name, values := range r.Header {
		s += len(name)
		for _, value := range values {
			s += len(value)
		}
	}
	s += len(r.Host)

	// N.B. r.Form and r.MultipartForm are assumed to be included in r.URL.

	if r.ContentLength != -1 {
		s += int(r.ContentLength)
	}
	return s
}
<|endoftext|>


<commit_before>package hammy

import (
	"math/rand"
	"time"
	"fmt"
)

//Main data processor implementation
type RequestHandlerImpl struct {
	//Interface for triggers retriving
	TGetter TriggersGetter
	//Interface for state storage
	SKeeper StateKeeper
	//Interface for executer
	Exec Executer
	//Interface for command commiter
	CBProcessor CmdBufferProcessor
}

//Internal struct for processing result
type objectProcessResult struct {
	Key string
	Err error
}

func (rh *RequestHandlerImpl) Handle(data IncomingData) (errs map[string]error) {
	//Allocate return value
	errs = make(map[string]error)

	if (len(data) == 0) {
		return
	}

	//Step 1: Loading triggers
	keys := make([]string, len(data))
	i := 0
	for k, _ := range data {
		keys[i] = k
		i++
	}

	triggers, err := rh.TGetter.MGet(keys)
	if err != nil {
		for _, k := range keys {
			errs[k] = fmt.Errorf("Failed to load trigger: %v", err)
		}
		return
	}

	if len(triggers) == 0 {
		return
	}

	//fix keys array
	keys = make([]string, len(triggers))
	i = 0
	for k := range triggers {
		keys[i] = k
		i++
	}

	//Step 2: Loading state
	states := rh.SKeeper.MGet(keys)
	if len(states) != len(triggers) {
		panic(fmt.Sprintf("Invalid length of StateKeeper answer (%v states for %v triggers)",
				len(states), len(triggers)))
	}

	//Step 3: Start trigger processing
	c := make(chan objectProcessResult) //result chanel
	runningTasks := 0
	for key, tr := range triggers {
		state := states[key]
		if state.Err != nil {
			errs[key] = fmt.Errorf("Failed to load state: %v", state.Err)
		} else {
			go rh.processTrigger(key, tr,
				state.State, state.Cas, data[key], c)
			runningTasks++
		}
	}

	//Step 4: Waiting until all triggers done
	if runningTasks > 0 {
		for res := range c {
			if res.Err != nil {
				errs[res.Key] = res.Err
			}
			runningTasks--
			if runningTasks == 0 {
				break
			}
		}
	}

	return
}

func (rh *RequestHandlerImpl) processTrigger(
		key string,
		trigger string,
		state State,
		cas *uint64,
		data IncomingObjectData,
		c chan objectProcessResult) {
//
	ret := objectProcessResult{Key: key, Err: nil}
	defer func() {
		c <- ret
	}()
//
	newState, cmdb, err := rh.Exec.ProcessTrigger(key, trigger, &state, data)
	if err != nil {
		ret.Err = fmt.Errorf("Trigger processor error: %v", err)
		return
	}

	//TODO: check state

	retry, err := rh.SKeeper.Set(key, *newState, cas)
	if err != nil {
		ret.Err = fmt.Errorf("StateKeeper set operation failed: %v", err)
		return
	}
	if retry {
		//Ooops! Collision!
		//Starting loop of retries
		rand.Seed( time.Now().UTC().UnixNano())
		for {
			//random sleep
			t := rand.Intn(100)
			if t > 50 {
				time.Sleep(time.Duration(t) * time.Millisecond)
			}

			//Retry...
			ans := rh.SKeeper.Get(key)
			if ans.Err != nil {
				ret.Err = fmt.Errorf("StateKeeper get operation failed: %v", ans.Err)
				return
			}
			state = ans.State
			cas = ans.Cas

			newState, cmdb, err = rh.Exec.ProcessTrigger(key, trigger, &state, data)
			if err != nil {
				ret.Err = fmt.Errorf("Trigger processor error: %v", err)
				return
			}

			//TODO: check state

			retry, err = rh.SKeeper.Set(key, *newState, cas)
			if err != nil {
				ret.Err = fmt.Errorf("StateKeeper set operation failed: %v", err)
				return
			}
			if !retry {
				//Work is done! No more need to retry!
				break
			}
		}
	}

	//Commit cmdbuffer
	err = rh.CBProcessor.Process(key, cmdb)
	if err != nil {
		ret.Err = fmt.Errorf("Failed to process cmdbuffer: %v", err)
		return
	}

	return
}
<commit_msg>opimization for stateless triggers<commit_after>package hammy

import (
	"math/rand"
	"time"
	"fmt"
	"reflect"
)

//Main data processor implementation
type RequestHandlerImpl struct {
	//Interface for triggers retriving
	TGetter TriggersGetter
	//Interface for state storage
	SKeeper StateKeeper
	//Interface for executer
	Exec Executer
	//Interface for command commiter
	CBProcessor CmdBufferProcessor
}

//Internal struct for processing result
type objectProcessResult struct {
	Key string
	Err error
}

func (rh *RequestHandlerImpl) Handle(data IncomingData) (errs map[string]error) {
	//Allocate return value
	errs = make(map[string]error)

	if (len(data) == 0) {
		return
	}

	//Step 1: Loading triggers
	keys := make([]string, len(data))
	i := 0
	for k, _ := range data {
		keys[i] = k
		i++
	}

	triggers, err := rh.TGetter.MGet(keys)
	if err != nil {
		for _, k := range keys {
			errs[k] = fmt.Errorf("Failed to load trigger: %v", err)
		}
		return
	}

	if len(triggers) == 0 {
		return
	}

	//fix keys array
	keys = make([]string, len(triggers))
	i = 0
	for k := range triggers {
		keys[i] = k
		i++
	}

	//Step 2: Loading state
	states := rh.SKeeper.MGet(keys)
	if len(states) != len(triggers) {
		panic(fmt.Sprintf("Invalid length of StateKeeper answer (%v states for %v triggers)",
				len(states), len(triggers)))
	}

	//Step 3: Start trigger processing
	c := make(chan objectProcessResult) //result chanel
	runningTasks := 0
	for key, tr := range triggers {
		state := states[key]
		if state.Err != nil {
			errs[key] = fmt.Errorf("Failed to load state: %v", state.Err)
		} else {
			go rh.processTrigger(key, tr,
				state.State, state.Cas, data[key], c)
			runningTasks++
		}
	}

	//Step 4: Waiting until all triggers done
	if runningTasks > 0 {
		for res := range c {
			if res.Err != nil {
				errs[res.Key] = res.Err
			}
			runningTasks--
			if runningTasks == 0 {
				break
			}
		}
	}

	return
}

func (rh *RequestHandlerImpl) processTrigger(
		key string,
		trigger string,
		state State,
		cas *uint64,
		data IncomingObjectData,
		c chan objectProcessResult) {
//
	ret := objectProcessResult{Key: key, Err: nil}
	defer func() {
		c <- ret
	}()
//
	newState, cmdb, err := rh.Exec.ProcessTrigger(key, trigger, &state, data)
	if err != nil {
		ret.Err = fmt.Errorf("Trigger processor error: %v", err)
		return
	}

	//Save state if it changes
	if !StatesEq(newState, &state) {
		retry, err := rh.SKeeper.Set(key, *newState, cas)
		if err != nil {
			ret.Err = fmt.Errorf("StateKeeper set operation failed: %v", err)
			return
		}
		if retry {
			//Ooops! Collision!
			//Starting loop of retries
			rand.Seed( time.Now().UTC().UnixNano())
			for {
				//random sleep
				t := rand.Intn(100)
				if t > 50 {
					time.Sleep(time.Duration(t) * time.Millisecond)
				}

				//Retry...
				ans := rh.SKeeper.Get(key)
				if ans.Err != nil {
					ret.Err = fmt.Errorf("StateKeeper get operation failed: %v", ans.Err)
					return
				}
				state = ans.State
				cas = ans.Cas

				newState, cmdb, err = rh.Exec.ProcessTrigger(key, trigger, &state, data)
				if err != nil {
					ret.Err = fmt.Errorf("Trigger processor error: %v", err)
					return
				}

				//TODO: check state
				//Save state if it changes
				if !StatesEq(newState, &state) {
					retry, err = rh.SKeeper.Set(key, *newState, cas)
					if err != nil {
						ret.Err = fmt.Errorf("StateKeeper set operation failed: %v", err)
						return
					}
					if !retry {
						//Work is done! No more need to retry!
						break
					}
				} else {
					// %)
					break
				}
			}
		}
	}

	//Commit cmdbuffer
	err = rh.CBProcessor.Process(key, cmdb)
	if err != nil {
		ret.Err = fmt.Errorf("Failed to process cmdbuffer: %v", err)
		return
	}

	return
}

//Compare two States
//true if equal, false otherwise
func StatesEq(a *State, b *State) bool {
	return reflect.DeepEqual(a, b)
}
<|endoftext|>


<commit_before>package properties

import (
	"log"

	"github.com/spf13/pflag"
	"github.com/spf13/viper"
	_ "github.com/spf13/viper/remote"
)

// Properties is a struct that allow to handle application
// properties seting and acessing. It wraps spf13/viper library
type Properties struct {
	*viper.Viper
	Config Config
}

// Properties constructor
// Settings default values if need
func New(config ...Config) *Properties {
	var c Config

	if len(config) == 0 {
		c = NewConfig()
	} else {
		c = config[0]
	}
	c.InitConfig()

	prop := Properties{Config: c, Viper: viper.New()}
	prop.init()

	return &prop
}

//initializes the properties instance - make calls to p.Viper library to initilize configuration.
func (p Properties) init() {

	//gives an instance of viper to Properties instance

	//Set config file name

	if p.Config.ConfigName != "" {
		p.Viper.SetConfigName(p.Config.ConfigName)
	}

	//set config file type for remote K/V stores
	if p.Config.ConfigType != "" {
		p.Viper.SetConfigType(p.Config.ConfigType)
	}

	//set the lookup pathes for config files
	if p.Config.ConfigPathes != nil && len(p.Config.ConfigPathes) > 0 {
		for _, path := range p.Config.ConfigPathes {
			p.Viper.AddConfigPath(path)
		}
		err := p.Viper.ReadInConfig()
		if err != nil {
			log.Panic(err)
		}
	}

	//Bind Env vars :
	if p.Config.EnvVars != nil && len(p.Config.EnvVars) > 0 {
		for _, envVar := range p.Config.EnvVars {
			p.Viper.BindEnv(envVar)
		}
	}

	//Set remote providers
	if p.Config.Providers != nil && len(p.Config.Providers) > 0 {
		for _, provider := range p.Config.Providers {
			if provider.KeyFile != "" {
				p.Viper.AddSecureRemoteProvider(provider.Name, provider.Url, provider.Path, provider.KeyFile)
			} else {
				p.Viper.AddRemoteProvider(provider.Name, provider.Url, provider.Path)
			}
		}
		err := p.Viper.ReadRemoteConfig()
		if err == nil {
			log.Panic(err)
		}
	}

	//Bind flags
	if p.Config.Flags != nil && len(p.Config.Flags) > 0 {
		for _, flag := range p.Config.Flags {
			pflag.String(flag.Name, flag.Default, flag.Usage)
			p.Viper.BindPFlag(flag.Name, pflag.Lookup(flag.Name))
		}
		pflag.Parse()
	}
}

// GetOrDie get key, if not found panic
func (p Properties) GetOrDie(key string) interface{} {
	if v := p.Get(key); v == nil {
		log.Panicf("Required property %s is not found!!", key)
	} else {
		return v
	}
	return nil
}

// GetOrDie get key, if not found panic
func (p Properties) GetSubOrDie(key string) *viper.Viper {
	if v := p.Sub(key); v == nil {
		log.Panicf("Required sub properties \"%s\" are not found!!", key)
	} else {
		return v
	}
	return nil
}

// GetStringOrDefault get string or a default value
func (p Properties) GetStringOrDefault(key string, dlft string) string {
	if v := p.GetString(key); v == "" {
		return dlft
	} else {
		return v
	}
}

// TryLoadRemoteProperties try load configuration from remote througth Viper
func (p Properties) TryLoadRemoteProperties() {
	var name = p.GetString("remote.name")
	var url = p.GetString("remote.url")
	var path = p.GetString("remote.path")
	var key = p.GetString("remote.key")
	if name != "" && url != "" && path != "" {
		if key != "" {
			p.AddSecureRemoteProvider(name, url, path, key)
		} else {
			p.AddRemoteProvider(name, url, path)
		}

		err := p.Viper.ReadRemoteConfig()
		if err != nil {
			// Handle errors reading the config file
			log.Panicf("Fatal error config Remote provider: %s \n", err)
		}
	}

}

// GetDefaultModeProperties get A Default Property set for classic app based on Flag with
// Config file : app.json in current directory from Flag
// With HOME and PWD from env
func (c Config) GetDefaultModeProperties() *Properties {
	props := New(c)
	return props
}

// // A Base Property set for quick init with
// // Config file : app.json in current directory from Properties (no flag listener)
// func (c Config) GetSimpleProperties() *Properties {
// 	props := New(c)
// 	props.Set(ConfigDirTag, DefaultConfigDir)
// 	props.Set(ConfigNameTag, DefaultConfigName)
// 	props.Set(ConfigTypeTag, DefaultConfigType)
// 	return props
// }

// DefaultLoadModeProperties is an Helper to Load Properties and merge it with mode related Properties
// defaultPath will be use by default if user not provide a ConfigDirTag in command line
// func DefaultLoadModeProperties(defaultPath string) *Properties {
// 	return LoadModeProperties(defaultPath, DefaultConfigMode, GetDefaultModeProperties(), true)
// }

// CheckRunInTestEnvironment return true if this application is running with 'go test'
func CheckRunInTestEnvironment() bool {
	if pflag.Lookup("test.v") == nil {
		return false
	} else {
		return true
	}
}

// Helper to Load Properties and merge it with mode related Properties
// path will be use by default if user not provide a ConfigDirTag in command line
// defaultMode will be use by default if user not provide a ModeTag in command line
// props is used as properties base.
// panicOnModeLoad if true, when loading mode properties failed call "panic" otherwise "warning"
func (props *Properties) LoadModeProperties(panicOnModeLoad bool) *Properties {

	var configName = props.GetStringOrDefault(ConfigNameTag, props.Config.ConfigName)
	//var configDir = props.GetStringOrDefault(ConfigDirTag, props.Config.DefaultConfigDir)
	var configType = props.GetStringOrDefault(ConfigTypeTag, props.Config.ConfigType)

	props.SetConfigType(configType)
	//props.AddConfigPath(configDir)
	var modeStr = props.GetString(ModeTag)
	if modeStr == "" {
		if isInTest := CheckRunInTestEnvironment(); isInTest == true {
			modeStr = props.Config.TestModeTag
		} else {
			modeStr = props.Config.DefaultConfigMode
		}
		if modeStr == "" {
			log.Panic("Mode is not set !")
		}
	}

	modeConfigName := modeStr + "." + configName
	props.SetConfigName(modeConfigName)

	err := props.MergeInConfig() // Find and read the config file
	if err != nil {
		if panicOnModeLoad {
			log.Panicf("Fatal error config mode %s : %s \n", modeConfigName, err)
		} else {
			log.Printf("Fatal error config mode %s : %s \n", modeConfigName, err)
			return props
		}
	}

	// var modeStr = props.GetString(ModeTag, defaultMode)
	// if modeStr == "" {
	// 	if isInTest := CheckRunInTestEnvironment(); isInTest == true {
	// 		modeStr = isInTest
	// 	} else {
	// 		log.Panic("Mode is not set !")
	// 	}
	// }

	// var modeConfigFilePath = configDir + "/" + modeStr + "." + configName + "." + configType
	// file, err := afero.ReadFile(afero.NewOsFs(), modeConfigFilePath)
	// if err != nil {
	// 	if panicOnModeLoad {
	// 		log.Panicf("Fatal error reading mode file %s : %s \n", modeConfigFilePath, err)
	// 	} else {
	// 		log.Printf("Error reading mode file %s : %s \n", modeConfigFilePath, err)
	// 		return props
	// 	}
	// }

	// err = props.MergeConfig(bytes.NewReader(file))
	// if err != nil {
	// 	log.Panicf("Fatal error merging mode file %s : %s \n", modeConfigFilePath, err)
	// }
	return props

}<commit_msg>missing cleanup<commit_after>package properties

import (
	"log"

	"github.com/spf13/pflag"
	"github.com/spf13/viper"
	_ "github.com/spf13/viper/remote"
)

// Properties is a struct that allow to handle application
// properties seting and acessing. It wraps spf13/viper library
type Properties struct {
	*viper.Viper
	Config Config
}

// Properties constructor
// Settings default values if need
func New(config ...Config) *Properties {
	var c Config

	if len(config) == 0 {
		c = NewConfig()
	} else {
		c = config[0]
	}
	c.InitConfig()

	prop := Properties{Config: c, Viper: viper.New()}
	prop.init()

	return &prop
}

//initializes the properties instance - make calls to p.Viper library to initilize configuration.
func (p Properties) init() {

	//gives an instance of viper to Properties instance

	//Set config file name

	if p.Config.ConfigName != "" {
		p.Viper.SetConfigName(p.Config.ConfigName)
	}

	//set config file type for remote K/V stores
	if p.Config.ConfigType != "" {
		p.Viper.SetConfigType(p.Config.ConfigType)
	}

	//set the lookup pathes for config files
	if p.Config.ConfigPathes != nil && len(p.Config.ConfigPathes) > 0 {
		for _, path := range p.Config.ConfigPathes {
			p.Viper.AddConfigPath(path)
		}
		err := p.Viper.ReadInConfig()
		if err != nil {
			log.Panic(err)
		}
	}

	//Bind Env vars :
	if p.Config.EnvVars != nil && len(p.Config.EnvVars) > 0 {
		for _, envVar := range p.Config.EnvVars {
			p.Viper.BindEnv(envVar)
		}
	}

	//Set remote providers
	if p.Config.Providers != nil && len(p.Config.Providers) > 0 {
		for _, provider := range p.Config.Providers {
			if provider.KeyFile != "" {
				p.Viper.AddSecureRemoteProvider(provider.Name, provider.Url, provider.Path, provider.KeyFile)
			} else {
				p.Viper.AddRemoteProvider(provider.Name, provider.Url, provider.Path)
			}
		}
		err := p.Viper.ReadRemoteConfig()
		if err == nil {
			log.Panic(err)
		}
	}

	//Bind flags
	if p.Config.Flags != nil && len(p.Config.Flags) > 0 {
		for _, flag := range p.Config.Flags {
			pflag.String(flag.Name, flag.Default, flag.Usage)
			p.Viper.BindPFlag(flag.Name, pflag.Lookup(flag.Name))
		}
		pflag.Parse()
	}
}

// GetOrDie get key, if not found panic
func (p Properties) GetOrDie(key string) interface{} {
	if v := p.Get(key); v == nil {
		log.Panicf("Required property %s is not found!!", key)
	} else {
		return v
	}
	return nil
}

// GetOrDie get key, if not found panic
func (p Properties) GetSubOrDie(key string) *viper.Viper {
	if v := p.Sub(key); v == nil {
		log.Panicf("Required sub properties \"%s\" are not found!!", key)
	} else {
		return v
	}
	return nil
}

// GetStringOrDefault get string or a default value
func (p Properties) GetStringOrDefault(key string, dlft string) string {
	if v := p.GetString(key); v == "" {
		return dlft
	} else {
		return v
	}
}

// TryLoadRemoteProperties try load configuration from remote througth Viper
func (p Properties) TryLoadRemoteProperties() {
	var name = p.GetString("remote.name")
	var url = p.GetString("remote.url")
	var path = p.GetString("remote.path")
	var key = p.GetString("remote.key")
	if name != "" && url != "" && path != "" {
		if key != "" {
			p.AddSecureRemoteProvider(name, url, path, key)
		} else {
			p.AddRemoteProvider(name, url, path)
		}

		err := p.Viper.ReadRemoteConfig()
		if err != nil {
			// Handle errors reading the config file
			log.Panicf("Fatal error config Remote provider: %s \n", err)
		}
	}

}

// GetDefaultModeProperties get A Default Property set for classic app based on Flag with
// Config file : app.json in current directory from Flag
// With HOME and PWD from env
func (c Config) GetDefaultModeProperties() *Properties {
	props := New(c)
	return props
}

// CheckRunInTestEnvironment return true if this application is running with 'go test'
func CheckRunInTestEnvironment() bool {
	if pflag.Lookup("test.v") == nil {
		return false
	} else {
		return true
	}
}

// Helper to Load Properties and merge it with mode related Properties
// path will be use by default if user not provide a ConfigDirTag in command line
// defaultMode will be use by default if user not provide a ModeTag in command line
// props is used as properties base.
// panicOnModeLoad if true, when loading mode properties failed call "panic" otherwise "warning"
func (props *Properties) LoadModeProperties(panicOnModeLoad bool) *Properties {

	var configName = props.GetStringOrDefault(ConfigNameTag, props.Config.ConfigName)
	var configType = props.GetStringOrDefault(ConfigTypeTag, props.Config.ConfigType)

	props.SetConfigType(configType)
	var modeStr = props.GetString(ModeTag)
	if modeStr == "" {
		if isInTest := CheckRunInTestEnvironment(); isInTest == true {
			modeStr = props.Config.TestModeTag
		} else {
			modeStr = props.Config.DefaultConfigMode
		}
		if modeStr == "" {
			log.Panic("Mode is not set !")
		}
	}

	modeConfigName := modeStr + "." + configName
	props.SetConfigName(modeConfigName)

	err := props.MergeInConfig() // Find and read the config file
	if err != nil {
		if panicOnModeLoad {
			log.Panicf("Fatal error config mode %s : %s \n", modeConfigName, err)
		} else {
			log.Printf("Fatal error config mode %s : %s \n", modeConfigName, err)
			return props
		}
	}

	return props

}<|endoftext|>


<commit_before>package owned

import "fmt"

func Example() {
	ch := New(nil)

	// Set an arbitrary value.
	ch <- func(_ interface{}) interface{} {
		return 42
	}

	// Print current value.
	ch <- func(v interface{}) interface{} {
		fmt.Println("Current value:", v)
		return v
	}

	// Atomic compare and swap
	swapped := make(chan bool)
	ch <- func(v interface{}) interface{} {
		if v.(int) == 42 {
			swapped <- true
			return 23
		}
		swapped <- false
		return v
	}
	if <-swapped {
		fmt.Println("Swapped")
	} else {
		fmt.Println("Didn't swap")
	}

	// Output:
	// Current value: 42
	// Swapped
}
<commit_msg>owned: Show methods in example<commit_after>package owned

import "fmt"

func Example() {
	ch := New(nil)

	// Set an arbitrary value.
	ch.Set(42)
	// Equivalent:
	ch <- func(_ interface{}) interface{} {
		return 42
	}

	// Print current value.
	fmt.Println("Current value:", ch.Get())
	// Equivalent:
	ch <- func(v interface{}) interface{} {
		fmt.Println("Current value:", v)
		return v
	}

	// Atomic compare and swap
	if ch.CAS(42, 23) {
		fmt.Println("Swapped")
	} else {
		fmt.Println("Didn't swap")
	}
	// Equivalent (but doesn't swap, as the comparison fails):
	swapped := make(chan bool)
	ch <- func(v interface{}) interface{} {
		if v.(int) == 42 {
			swapped <- true
			return 23
		}
		swapped <- false
		return v
	}
	if <-swapped {
		fmt.Println("Swapped")
	} else {
		fmt.Println("Didn't swap")
	}

	// Output:
	// Current value: 42
	// Current value: 42
	// Swapped
	// Didn't swap
}
<|endoftext|>


<commit_before>package daemon

import (
	"bufio"
	"fmt"
	"io"
	"os/exec"
	"os/user"
	"path"
	"sync"
	"syscall"

	"github.com/subgraph/oz"
	"github.com/subgraph/oz/fs"
	"github.com/subgraph/oz/network"
	"github.com/subgraph/oz/xpra"

	"github.com/op/go-logging"
	"github.com/subgraph/oz/oz-init"
)

const initPath = "/usr/local/bin/oz-init"

type Sandbox struct {
	daemon  *daemonState
	id      int
	display int
	profile *oz.Profile
	init    *exec.Cmd
	cred    *syscall.Credential
	fs      *fs.Filesystem
	stderr  io.ReadCloser
	addr    string
	xpra    *xpra.Xpra
	ready   sync.WaitGroup
	network *network.SandboxNetwork
}

/*
func findSandbox(id int) *Sandbox {
	for _, sb := range sandboxes {
		if sb.id == id {
			return sb
		}
	}
	return nil
}
*/

func createInitCommand(name, chroot string, env []string, uid uint32, display int, stn *network.SandboxNetwork, nettype string) *exec.Cmd {
	cmd := exec.Command(initPath)
	cmd.Dir = "/"

	cloneFlags := uintptr(syscall.CLONE_NEWNS)
	cloneFlags |= syscall.CLONE_NEWIPC
	cloneFlags |= syscall.CLONE_NEWPID
	cloneFlags |= syscall.CLONE_NEWUTS

	if nettype != "host" {
		cloneFlags |= syscall.CLONE_NEWNET
	}

	cmd.SysProcAttr = &syscall.SysProcAttr{
		Chroot:     chroot,
		Cloneflags: cloneFlags,
	}
	cmd.Env = []string{
		"INIT_PROFILE=" + name,
		fmt.Sprintf("INIT_UID=%d", uid),
	}

	if stn.Ip != "" {
		cmd.Env = append(cmd.Env, "INIT_ADDR="+stn.Ip)
		cmd.Env = append(cmd.Env, "INIT_VHOST="+stn.VethHost)
		cmd.Env = append(cmd.Env, "INIT_VGUEST="+stn.VethGuest)
		cmd.Env = append(cmd.Env, "INIT_GATEWAY="+stn.Gateway.String()+"/"+stn.Class)
	}

	cmd.Env = append(cmd.Env, fmt.Sprintf("INIT_DISPLAY=%d", display))

	for _, e := range env {
		cmd.Env = append(cmd.Env, ozinit.EnvPrefix+e)
	}

	return cmd
}

func (d *daemonState) launch(p *oz.Profile, env []string, uid, gid uint32, log *logging.Logger) (*Sandbox, error) {
	u, err := user.LookupId(fmt.Sprintf("%d", uid))
	if err != nil {
		return nil, fmt.Errorf("failed to lookup user for uid=%d: %v", uid, err)
	}
	fs := fs.NewFromProfile(p, u, d.config.SandboxPath, d.config.UseFullDev, d.log)
	if err := fs.Setup(d.config.ProfileDir); err != nil {
		return nil, err
	}
	display := 0
	if p.XServer.Enabled && p.Networking.Nettype == "host" {
		display = d.nextDisplay
		d.nextDisplay += 1
	}

	stn := new(network.SandboxNetwork)
	if p.Networking.Nettype == "bridge" {
		stn, err = network.PrepareSandboxNetwork(d.network, log)
		if err != nil {
			return nil, fmt.Errorf("Unable to prepare veth network: %+v", err)
		}
	}

	cmd := createInitCommand(p.Name, fs.Root(), env, uid, display, stn, p.Networking.Nettype)
	log.Debug("Command environment: %+v", cmd.Env)
	pp, err := cmd.StderrPipe()
	if err != nil {
		fs.Cleanup()
		return nil, fmt.Errorf("error creating stderr pipe for init process: %v", err)

	}

	if err := cmd.Start(); err != nil {
		fs.Cleanup()
		return nil, fmt.Errorf("Unable to start process: %+v", err)
	}

	sbox := &Sandbox{
		daemon:  d,
		id:      d.nextSboxId,
		display: display,
		profile: p,
		init:    cmd,
		cred:    &syscall.Credential{Uid: uid, Gid: gid},
		fs:      fs,
		addr:    path.Join(fs.Root(), "tmp", "oz-init-control"),
		stderr:  pp,
		network: stn,
	}

	if p.Networking.Nettype == "bridge" {
		if err := network.NetInit(stn, d.network, cmd.Process.Pid, log); err != nil {
			cmd.Process.Kill()
			fs.Cleanup()
			return nil, fmt.Errorf("Unable to create veth networking: %+v", err)
		}
	}

	sbox.ready.Add(1)
	go sbox.logMessages()
	if sbox.profile.XServer.Enabled {
		go func() {
			sbox.ready.Wait()
			go sbox.startXpraClient()
		}()
	}
	d.nextSboxId += 1
	d.sandboxes = append(d.sandboxes, sbox)
	return sbox, nil
}

func (sbox *Sandbox) remove(log *logging.Logger) {
	sboxes := []*Sandbox{}
	for _, sb := range sbox.daemon.sandboxes {
		if sb == sbox {
			sb.fs.Cleanup()
			if sb.profile.Networking.Nettype == "bridge" {
				sb.network.Cleanup(log)
			}
		} else {
			sboxes = append(sboxes, sb)
		}
	}
	sbox.daemon.sandboxes = sboxes
}

func (sbox *Sandbox) logMessages() {
	scanner := bufio.NewScanner(sbox.stderr)
	seenOk := false
	for scanner.Scan() {
		line := scanner.Text()
		if line == "OK" && !seenOk {
			sbox.daemon.log.Info("oz-init (%s) is ready", sbox.profile.Name)
			seenOk = true
			sbox.ready.Done()
		} else if len(line) > 1 {
			sbox.logLine(line)
		}
	}
	sbox.stderr.Close()
}

func (sbox *Sandbox) logLine(line string) {
	if len(line) < 2 {
		return
	}
	f := sbox.getLogFunc(line[0])
	msg := line[2:]
	if f != nil {
		f("[%s] %s", sbox.profile.Name, msg)
	} else {
		sbox.daemon.log.Info("[%s] %s", sbox.profile.Name, line)
	}
}

func (sbox *Sandbox) getLogFunc(c byte) func(string, ...interface{}) {
	log := sbox.daemon.log
	switch c {
	case 'D':
		return log.Debug
	case 'I':
		return log.Info
	case 'N':
		return log.Notice
	case 'W':
		return log.Warning
	case 'E':
		return log.Error
	case 'C':
		return log.Critical
	}
	return nil
}

func (sbox *Sandbox) startXpraClient() {
	sbox.xpra = xpra.NewClient(
		&sbox.profile.XServer,
		uint64(sbox.display),
		sbox.cred,
		sbox.fs.Xpra(),
		sbox.profile.Name,
		sbox.daemon.log)

	if sbox.daemon.config.LogXpra {
		sbox.setupXpraLogging()
	}
	if err := sbox.xpra.Process.Start(); err != nil {
		sbox.daemon.Warning("Failed to start xpra client: %v", err)
	}
}

func (sbox *Sandbox) setupXpraLogging() {
	stdout, err := sbox.xpra.Process.StdoutPipe()
	if err != nil {
		sbox.daemon.Warning("Failed to create xpra stdout pipe: %v", err)
		return
	}
	stderr, err := sbox.xpra.Process.StderrPipe()
	if err != nil {
		stdout.Close()
		sbox.daemon.Warning("Failed to create xpra stderr pipe: %v", err)
	}
	go sbox.logPipeOutput(stdout, "xpra-stdout")
	go sbox.logPipeOutput(stderr, "xpra-stderr")
}

func (sbox *Sandbox) logPipeOutput(p io.Reader, label string) {
	scanner := bufio.NewScanner(p)
	for scanner.Scan() {
		line := scanner.Text()
		sbox.daemon.log.Info("(%s) %s", label, line)
	}
}
<commit_msg>get rid of commented out code<commit_after>package daemon

import (
	"bufio"
	"fmt"
	"io"
	"os/exec"
	"os/user"
	"path"
	"sync"
	"syscall"

	"github.com/subgraph/oz"
	"github.com/subgraph/oz/fs"
	"github.com/subgraph/oz/network"
	"github.com/subgraph/oz/xpra"

	"github.com/op/go-logging"
	"github.com/subgraph/oz/oz-init"
)

const initPath = "/usr/local/bin/oz-init"

type Sandbox struct {
	daemon  *daemonState
	id      int
	display int
	profile *oz.Profile
	init    *exec.Cmd
	cred    *syscall.Credential
	fs      *fs.Filesystem
	stderr  io.ReadCloser
	addr    string
	xpra    *xpra.Xpra
	ready   sync.WaitGroup
	network *network.SandboxNetwork
}

func createInitCommand(name, chroot string, env []string, uid uint32, display int, stn *network.SandboxNetwork, nettype string) *exec.Cmd {
	cmd := exec.Command(initPath)
	cmd.Dir = "/"

	cloneFlags := uintptr(syscall.CLONE_NEWNS)
	cloneFlags |= syscall.CLONE_NEWIPC
	cloneFlags |= syscall.CLONE_NEWPID
	cloneFlags |= syscall.CLONE_NEWUTS

	if nettype != "host" {
		cloneFlags |= syscall.CLONE_NEWNET
	}

	cmd.SysProcAttr = &syscall.SysProcAttr{
		Chroot:     chroot,
		Cloneflags: cloneFlags,
	}
	cmd.Env = []string{
		"INIT_PROFILE=" + name,
		fmt.Sprintf("INIT_UID=%d", uid),
	}

	if stn.Ip != "" {
		cmd.Env = append(cmd.Env, "INIT_ADDR="+stn.Ip)
		cmd.Env = append(cmd.Env, "INIT_VHOST="+stn.VethHost)
		cmd.Env = append(cmd.Env, "INIT_VGUEST="+stn.VethGuest)
		cmd.Env = append(cmd.Env, "INIT_GATEWAY="+stn.Gateway.String()+"/"+stn.Class)
	}

	cmd.Env = append(cmd.Env, fmt.Sprintf("INIT_DISPLAY=%d", display))

	for _, e := range env {
		cmd.Env = append(cmd.Env, ozinit.EnvPrefix+e)
	}

	return cmd
}

func (d *daemonState) launch(p *oz.Profile, env []string, uid, gid uint32, log *logging.Logger) (*Sandbox, error) {
	u, err := user.LookupId(fmt.Sprintf("%d", uid))
	if err != nil {
		return nil, fmt.Errorf("failed to lookup user for uid=%d: %v", uid, err)
	}
	fs := fs.NewFromProfile(p, u, d.config.SandboxPath, d.config.UseFullDev, d.log)
	if err := fs.Setup(d.config.ProfileDir); err != nil {
		return nil, err
	}
	display := 0
	if p.XServer.Enabled && p.Networking.Nettype == "host" {
		display = d.nextDisplay
		d.nextDisplay += 1
	}

	stn := new(network.SandboxNetwork)
	if p.Networking.Nettype == "bridge" {
		stn, err = network.PrepareSandboxNetwork(d.network, log)
		if err != nil {
			return nil, fmt.Errorf("Unable to prepare veth network: %+v", err)
		}
	}

	cmd := createInitCommand(p.Name, fs.Root(), env, uid, display, stn, p.Networking.Nettype)
	log.Debug("Command environment: %+v", cmd.Env)
	pp, err := cmd.StderrPipe()
	if err != nil {
		fs.Cleanup()
		return nil, fmt.Errorf("error creating stderr pipe for init process: %v", err)

	}

	if err := cmd.Start(); err != nil {
		fs.Cleanup()
		return nil, fmt.Errorf("Unable to start process: %+v", err)
	}

	sbox := &Sandbox{
		daemon:  d,
		id:      d.nextSboxId,
		display: display,
		profile: p,
		init:    cmd,
		cred:    &syscall.Credential{Uid: uid, Gid: gid},
		fs:      fs,
		addr:    path.Join(fs.Root(), "tmp", "oz-init-control"),
		stderr:  pp,
		network: stn,
	}

	if p.Networking.Nettype == "bridge" {
		if err := network.NetInit(stn, d.network, cmd.Process.Pid, log); err != nil {
			cmd.Process.Kill()
			fs.Cleanup()
			return nil, fmt.Errorf("Unable to create veth networking: %+v", err)
		}
	}

	sbox.ready.Add(1)
	go sbox.logMessages()
	if sbox.profile.XServer.Enabled {
		go func() {
			sbox.ready.Wait()
			go sbox.startXpraClient()
		}()
	}
	d.nextSboxId += 1
	d.sandboxes = append(d.sandboxes, sbox)
	return sbox, nil
}

func (sbox *Sandbox) remove(log *logging.Logger) {
	sboxes := []*Sandbox{}
	for _, sb := range sbox.daemon.sandboxes {
		if sb == sbox {
			sb.fs.Cleanup()
			if sb.profile.Networking.Nettype == "bridge" {
				sb.network.Cleanup(log)
			}
		} else {
			sboxes = append(sboxes, sb)
		}
	}
	sbox.daemon.sandboxes = sboxes
}

func (sbox *Sandbox) logMessages() {
	scanner := bufio.NewScanner(sbox.stderr)
	seenOk := false
	for scanner.Scan() {
		line := scanner.Text()
		if line == "OK" && !seenOk {
			sbox.daemon.log.Info("oz-init (%s) is ready", sbox.profile.Name)
			seenOk = true
			sbox.ready.Done()
		} else if len(line) > 1 {
			sbox.logLine(line)
		}
	}
	sbox.stderr.Close()
}

func (sbox *Sandbox) logLine(line string) {
	if len(line) < 2 {
		return
	}
	f := sbox.getLogFunc(line[0])
	msg := line[2:]
	if f != nil {
		f("[%s] %s", sbox.profile.Name, msg)
	} else {
		sbox.daemon.log.Info("[%s] %s", sbox.profile.Name, line)
	}
}

func (sbox *Sandbox) getLogFunc(c byte) func(string, ...interface{}) {
	log := sbox.daemon.log
	switch c {
	case 'D':
		return log.Debug
	case 'I':
		return log.Info
	case 'N':
		return log.Notice
	case 'W':
		return log.Warning
	case 'E':
		return log.Error
	case 'C':
		return log.Critical
	}
	return nil
}

func (sbox *Sandbox) startXpraClient() {
	sbox.xpra = xpra.NewClient(
		&sbox.profile.XServer,
		uint64(sbox.display),
		sbox.cred,
		sbox.fs.Xpra(),
		sbox.profile.Name,
		sbox.daemon.log)

	if sbox.daemon.config.LogXpra {
		sbox.setupXpraLogging()
	}
	if err := sbox.xpra.Process.Start(); err != nil {
		sbox.daemon.Warning("Failed to start xpra client: %v", err)
	}
}

func (sbox *Sandbox) setupXpraLogging() {
	stdout, err := sbox.xpra.Process.StdoutPipe()
	if err != nil {
		sbox.daemon.Warning("Failed to create xpra stdout pipe: %v", err)
		return
	}
	stderr, err := sbox.xpra.Process.StderrPipe()
	if err != nil {
		stdout.Close()
		sbox.daemon.Warning("Failed to create xpra stderr pipe: %v", err)
	}
	go sbox.logPipeOutput(stdout, "xpra-stdout")
	go sbox.logPipeOutput(stderr, "xpra-stderr")
}

func (sbox *Sandbox) logPipeOutput(p io.Reader, label string) {
	scanner := bufio.NewScanner(p)
	for scanner.Scan() {
		line := scanner.Text()
		sbox.daemon.log.Info("(%s) %s", label, line)
	}
}
<|endoftext|>


<commit_before>package daemon

import (
	"bufio"
	"fmt"
	"io"
	"os/exec"
	"os/user"
	"path"
	"sync"
	"syscall"

	"github.com/subgraph/oz"
	"github.com/subgraph/oz/fs"
	"github.com/subgraph/oz/network"
	"github.com/subgraph/oz/xpra"

	"github.com/op/go-logging"
	"github.com/subgraph/oz/oz-init"
)

const initPath = "/usr/local/bin/oz-init"

type Sandbox struct {
	daemon  *daemonState
	id      int
	display int
	profile *oz.Profile
	init    *exec.Cmd
	cred    *syscall.Credential
	fs      *fs.Filesystem
	stderr  io.ReadCloser
	addr    string
	xpra    *xpra.Xpra
	ready   sync.WaitGroup
	network *network.SandboxNetwork
}

func createInitCommand(name, chroot string, env []string, uid uint32, display int, stn *network.SandboxNetwork, nettype string) *exec.Cmd {
	cmd := exec.Command(initPath)
	cmd.Dir = "/"

	cloneFlags := uintptr(syscall.CLONE_NEWNS)
	cloneFlags |= syscall.CLONE_NEWIPC
	cloneFlags |= syscall.CLONE_NEWPID
	cloneFlags |= syscall.CLONE_NEWUTS

	if nettype != "host" {
		cloneFlags |= syscall.CLONE_NEWNET
	}

	cmd.SysProcAttr = &syscall.SysProcAttr{
		Chroot:     chroot,
		Cloneflags: cloneFlags,
	}
	cmd.Env = []string{
		"INIT_PROFILE=" + name,
		fmt.Sprintf("INIT_UID=%d", uid),
	}

	if stn.Ip != "" {
		cmd.Env = append(cmd.Env, "INIT_ADDR="+stn.Ip)
		cmd.Env = append(cmd.Env, "INIT_VHOST="+stn.VethHost)
		cmd.Env = append(cmd.Env, "INIT_VGUEST="+stn.VethGuest)
		cmd.Env = append(cmd.Env, "INIT_GATEWAY="+stn.Gateway.String()+"/"+stn.Class)
	}

	cmd.Env = append(cmd.Env, fmt.Sprintf("INIT_DISPLAY=%d", display))

	for _, e := range env {
		cmd.Env = append(cmd.Env, ozinit.EnvPrefix+e)
	}

	return cmd
}

func (d *daemonState) launch(p *oz.Profile, env []string, uid, gid uint32, log *logging.Logger) (*Sandbox, error) {
	u, err := user.LookupId(fmt.Sprintf("%d", uid))
	if err != nil {
		return nil, fmt.Errorf("failed to lookup user for uid=%d: %v", uid, err)
	}
	fs := fs.NewFromProfile(p, u, d.config.SandboxPath, d.config.UseFullDev, d.log)
	if err := fs.Setup(d.config.ProfileDir); err != nil {
		return nil, err
	}
	display := 0
	if p.XServer.Enabled && p.Networking.Nettype == "host" {
		display = d.nextDisplay
		d.nextDisplay += 1
	}

	stn := new(network.SandboxNetwork)
	if p.Networking.Nettype == "bridge" {
		stn, err = network.PrepareSandboxNetwork(d.network, log)
		if err != nil {
			return nil, fmt.Errorf("Unable to prepare veth network: %+v", err)
		}
	}

	cmd := createInitCommand(p.Name, fs.Root(), env, uid, display, stn, p.Networking.Nettype)
	log.Debug("Command environment: %+v", cmd.Env)
	pp, err := cmd.StderrPipe()
	if err != nil {
		fs.Cleanup()
		return nil, fmt.Errorf("error creating stderr pipe for init process: %v", err)

	}

	if err := cmd.Start(); err != nil {
		fs.Cleanup()
		return nil, fmt.Errorf("Unable to start process: %+v", err)
	}

	sbox := &Sandbox{
		daemon:  d,
		id:      d.nextSboxId,
		display: display,
		profile: p,
		init:    cmd,
		cred:    &syscall.Credential{Uid: uid, Gid: gid},
		fs:      fs,
		addr:    path.Join(fs.Root(), "tmp", "oz-init-control"),
		stderr:  pp,
		network: stn,
	}

	if p.Networking.Nettype == "bridge" {
		if err := network.NetInit(stn, d.network, cmd.Process.Pid, log); err != nil {
			cmd.Process.Kill()
			fs.Cleanup()
			return nil, fmt.Errorf("Unable to create veth networking: %+v", err)
		}
	}

	sbox.ready.Add(1)
	go sbox.logMessages()
	if sbox.profile.XServer.Enabled {
		go func() {
			sbox.ready.Wait()
			go sbox.startXpraClient()
		}()
	}
	d.nextSboxId += 1
	d.sandboxes = append(d.sandboxes, sbox)
	return sbox, nil
}

func (sbox *Sandbox) remove(log *logging.Logger) {
	sboxes := []*Sandbox{}
	for _, sb := range sbox.daemon.sandboxes {
		if sb == sbox {
			sb.fs.Cleanup()
			if sb.profile.Networking.Nettype == "bridge" {
				sb.network.Cleanup(log)
			}
		} else {
			sboxes = append(sboxes, sb)
		}
	}
	sbox.daemon.sandboxes = sboxes
}

func (sbox *Sandbox) logMessages() {
	scanner := bufio.NewScanner(sbox.stderr)
	seenOk := false
	for scanner.Scan() {
		line := scanner.Text()
		if line == "OK" && !seenOk {
			sbox.daemon.log.Info("oz-init (%s) is ready", sbox.profile.Name)
			seenOk = true
			sbox.ready.Done()
		} else if len(line) > 1 {
			sbox.logLine(line)
		}
	}
	sbox.stderr.Close()
}

func (sbox *Sandbox) logLine(line string) {
	if len(line) < 2 {
		return
	}
	f := sbox.getLogFunc(line[0])
	msg := line[2:]
	if f != nil {
		f("[%s] %s", sbox.profile.Name, msg)
	} else {
		sbox.daemon.log.Info("[%s] %s", sbox.profile.Name, line)
	}
}

func (sbox *Sandbox) getLogFunc(c byte) func(string, ...interface{}) {
	log := sbox.daemon.log
	switch c {
	case 'D':
		return log.Debug
	case 'I':
		return log.Info
	case 'N':
		return log.Notice
	case 'W':
		return log.Warning
	case 'E':
		return log.Error
	case 'C':
		return log.Critical
	}
	return nil
}

func (sbox *Sandbox) startXpraClient() {
	sbox.xpra = xpra.NewClient(
		&sbox.profile.XServer,
		uint64(sbox.display),
		sbox.cred,
		sbox.fs.Xpra(),
		sbox.profile.Name,
		sbox.daemon.log)

	if sbox.daemon.config.LogXpra {
		sbox.setupXpraLogging()
	}
	if err := sbox.xpra.Process.Start(); err != nil {
		sbox.daemon.Warning("Failed to start xpra client: %v", err)
	}
}

func (sbox *Sandbox) setupXpraLogging() {
	stdout, err := sbox.xpra.Process.StdoutPipe()
	if err != nil {
		sbox.daemon.Warning("Failed to create xpra stdout pipe: %v", err)
		return
	}
	stderr, err := sbox.xpra.Process.StderrPipe()
	if err != nil {
		stdout.Close()
		sbox.daemon.Warning("Failed to create xpra stderr pipe: %v", err)
	}
	go sbox.logPipeOutput(stdout, "xpra-stdout")
	go sbox.logPipeOutput(stderr, "xpra-stderr")
}

func (sbox *Sandbox) logPipeOutput(p io.Reader, label string) {
	scanner := bufio.NewScanner(p)
	for scanner.Scan() {
		line := scanner.Text()
		sbox.daemon.log.Info("(%s) %s", label, line)
	}
}
<commit_msg>whitespace<commit_after>package daemon

import (
	"bufio"
	"fmt"
	"io"
	"os/exec"
	"os/user"
	"path"
	"sync"
	"syscall"

	"github.com/subgraph/oz"
	"github.com/subgraph/oz/fs"
	"github.com/subgraph/oz/network"
	"github.com/subgraph/oz/xpra"

	"github.com/op/go-logging"
	"github.com/subgraph/oz/oz-init"
)

const initPath = "/usr/local/bin/oz-init"

type Sandbox struct {
	daemon  *daemonState
	id      int
	display int
	profile *oz.Profile
	init    *exec.Cmd
	cred    *syscall.Credential
	fs      *fs.Filesystem
	stderr  io.ReadCloser
	addr    string
	xpra    *xpra.Xpra
	ready   sync.WaitGroup
	network *network.SandboxNetwork
}

func createInitCommand(name, chroot string, env []string, uid uint32, display int, stn *network.SandboxNetwork, nettype string) *exec.Cmd {
	cmd := exec.Command(initPath)
	cmd.Dir = "/"

	cloneFlags := uintptr(syscall.CLONE_NEWNS)
	cloneFlags |= syscall.CLONE_NEWIPC
	cloneFlags |= syscall.CLONE_NEWPID
	cloneFlags |= syscall.CLONE_NEWUTS

	if nettype != "host" {
		cloneFlags |= syscall.CLONE_NEWNET
	}

	cmd.SysProcAttr = &syscall.SysProcAttr{
		Chroot:     chroot,
		Cloneflags: cloneFlags,
	}
	cmd.Env = []string{
		"INIT_PROFILE=" + name,
		fmt.Sprintf("INIT_UID=%d", uid),
	}

	if stn.Ip != "" {
		cmd.Env = append(cmd.Env, "INIT_ADDR="+stn.Ip)
		cmd.Env = append(cmd.Env, "INIT_VHOST="+stn.VethHost)
		cmd.Env = append(cmd.Env, "INIT_VGUEST="+stn.VethGuest)
		cmd.Env = append(cmd.Env, "INIT_GATEWAY="+stn.Gateway.String()+"/"+stn.Class)
	}

	cmd.Env = append(cmd.Env, fmt.Sprintf("INIT_DISPLAY=%d", display))

	for _, e := range env {
		cmd.Env = append(cmd.Env, ozinit.EnvPrefix+e)
	}

	return cmd
}

func (d *daemonState) launch(p *oz.Profile, env []string, uid, gid uint32, log *logging.Logger) (*Sandbox, error) {
	u, err := user.LookupId(fmt.Sprintf("%d", uid))
	if err != nil {
		return nil, fmt.Errorf("failed to lookup user for uid=%d: %v", uid, err)
	}
	fs := fs.NewFromProfile(p, u, d.config.SandboxPath, d.config.UseFullDev, d.log)
	if err := fs.Setup(d.config.ProfileDir); err != nil {
		return nil, err
	}
	display := 0
	if p.XServer.Enabled && p.Networking.Nettype == "host" {
		display = d.nextDisplay
		d.nextDisplay += 1
	}

	stn := new(network.SandboxNetwork)
	if p.Networking.Nettype == "bridge" {
		stn, err = network.PrepareSandboxNetwork(d.network, log)
		if err != nil {
			return nil, fmt.Errorf("Unable to prepare veth network: %+v", err)
		}
	}

	cmd := createInitCommand(p.Name, fs.Root(), env, uid, display, stn, p.Networking.Nettype)
	log.Debug("Command environment: %+v", cmd.Env)
	pp, err := cmd.StderrPipe()
	if err != nil {
		fs.Cleanup()
		return nil, fmt.Errorf("error creating stderr pipe for init process: %v", err)

	}

	if err := cmd.Start(); err != nil {
		fs.Cleanup()
		return nil, fmt.Errorf("Unable to start process: %+v", err)
	}

	sbox := &Sandbox{
		daemon:  d,
		id:      d.nextSboxId,
		display: display,
		profile: p,
		init:    cmd,
		cred:    &syscall.Credential{Uid: uid, Gid: gid},
		fs:      fs,
		addr:    path.Join(fs.Root(), "tmp", "oz-init-control"),
		stderr:  pp,
		network: stn,
	}

	if p.Networking.Nettype == "bridge" {
		if err := network.NetInit(stn, d.network, cmd.Process.Pid, log); err != nil {
			cmd.Process.Kill()
			fs.Cleanup()
			return nil, fmt.Errorf("Unable to create veth networking: %+v", err)
		}
	}

	sbox.ready.Add(1)
	go sbox.logMessages()
	if sbox.profile.XServer.Enabled {
		go func() {
			sbox.ready.Wait()
			go sbox.startXpraClient()
		}()
	}
	d.nextSboxId += 1
	d.sandboxes = append(d.sandboxes, sbox)
	return sbox, nil
}

func (sbox *Sandbox) remove(log *logging.Logger) {
	sboxes := []*Sandbox{}
	for _, sb := range sbox.daemon.sandboxes {
		if sb == sbox {
			sb.fs.Cleanup()
			if sb.profile.Networking.Nettype == "bridge" {
				sb.network.Cleanup(log)
			}
		} else {
			sboxes = append(sboxes, sb)
		}
	}
	sbox.daemon.sandboxes = sboxes
}

func (sbox *Sandbox) logMessages() {
	scanner := bufio.NewScanner(sbox.stderr)
	seenOk := false
	for scanner.Scan() {
		line := scanner.Text()
		if line == "OK" && !seenOk {
			sbox.daemon.log.Info("oz-init (%s) is ready", sbox.profile.Name)
			seenOk = true
			sbox.ready.Done()
		} else if len(line) > 1 {
			sbox.logLine(line)
		}
	}
	sbox.stderr.Close()
}

func (sbox *Sandbox) logLine(line string) {
	if len(line) < 2 {
		return
	}
	f := sbox.getLogFunc(line[0])
	msg := line[2:]
	if f != nil {
		f("[%s] %s", sbox.profile.Name, msg)
	} else {
		sbox.daemon.log.Info("[%s] %s", sbox.profile.Name, line)
	}
}

func (sbox *Sandbox) getLogFunc(c byte) func(string, ...interface{}) {
	log := sbox.daemon.log
	switch c {
	case 'D':
		return log.Debug
	case 'I':
		return log.Info
	case 'N':
		return log.Notice
	case 'W':
		return log.Warning
	case 'E':
		return log.Error
	case 'C':
		return log.Critical
	}
	return nil
}

func (sbox *Sandbox) startXpraClient() {
	sbox.xpra = xpra.NewClient(
		&sbox.profile.XServer,
		uint64(sbox.display),
		sbox.cred,
		sbox.fs.Xpra(),
		sbox.profile.Name,
		sbox.daemon.log)

	if sbox.daemon.config.LogXpra {
		sbox.setupXpraLogging()
	}
	if err := sbox.xpra.Process.Start(); err != nil {
		sbox.daemon.Warning("Failed to start xpra client: %v", err)
	}
}

func (sbox *Sandbox) setupXpraLogging() {
	stdout, err := sbox.xpra.Process.StdoutPipe()
	if err != nil {
		sbox.daemon.Warning("Failed to create xpra stdout pipe: %v", err)
		return
	}
	stderr, err := sbox.xpra.Process.StderrPipe()
	if err != nil {
		stdout.Close()
		sbox.daemon.Warning("Failed to create xpra stderr pipe: %v", err)
	}
	go sbox.logPipeOutput(stdout, "xpra-stdout")
	go sbox.logPipeOutput(stderr, "xpra-stderr")
}

func (sbox *Sandbox) logPipeOutput(p io.Reader, label string) {
	scanner := bufio.NewScanner(p)
	for scanner.Scan() {
		line := scanner.Text()
		sbox.daemon.log.Info("(%s) %s", label, line)
	}
}
<|endoftext|>


<commit_before>package acceptance_test

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"lib/models"
	"lib/policy_client"
	"math/rand"
	"net"
	"net/http"
	"os"
	"strings"
	"sync"
	"time"

	"code.cloudfoundry.org/lager/lagertest"

	"github.com/cloudfoundry-incubator/cf-test-helpers/cf"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gexec"
	"github.com/pivotal-cf-experimental/rainmaker"
)

const Timeout_Push = 5 * time.Minute
const Timeout_Short = 10 * time.Second

var ports []int

func getSubnet(ip string) string {
	return strings.Split(ip, ".")[2]
}

func isSameCell(sourceIP, destIP string) bool {
	return getSubnet(sourceIP) == getSubnet(destIP)
}

var _ = Describe("connectivity between containers on the overlay network", func() {
	Describe("networking policy", func() {
		var (
			appProxy     string
			appRegistry  string
			appsTest     []string
			orgName      string
			spaceName    string
			appInstances int
			applications int
		)

		BeforeEach(func() {
			appInstances = testConfig.AppInstances
			applications = testConfig.Applications

			appProxy = fmt.Sprintf("proxy-%d", rand.Int31())
			appRegistry = fmt.Sprintf("registry-%d", rand.Int31())
			for i := 0; i < applications; i++ {
				appsTest = append(appsTest, fmt.Sprintf("tick-%d-%d", i, rand.Int31()))
			}

			ports = []int{8080}
			for i := 0; i < testConfig.Policies; i++ {
				ports = append(ports, 7000+i)
			}

			Auth(testConfig.TestUser, testConfig.TestUserPassword)

			orgName = "test-org"
			Expect(cf.Cf("create-org", orgName).Wait(Timeout_Push)).To(gexec.Exit(0))
			Expect(cf.Cf("target", "-o", orgName).Wait(Timeout_Push)).To(gexec.Exit(0))

			spaceName = "test-space"
			Expect(cf.Cf("create-space", spaceName).Wait(Timeout_Push)).To(gexec.Exit(0))
			Expect(cf.Cf("target", "-o", orgName, "-s", spaceName).Wait(Timeout_Push)).To(gexec.Exit(0))
		})

		AfterEach(func() {
			appReport(appProxy, Timeout_Short)
			appReport(appRegistry, Timeout_Short)
			appsReport(appsTest, Timeout_Short)

			// clean up everything
			Expect(cf.Cf("delete-org", orgName, "-f").Wait(Timeout_Push)).To(gexec.Exit(0))
		})

		It("allows the user to configure policies", func(done Done) {
			By("pushing the registry app and proxy app")
			var setupWG sync.WaitGroup
			setupWG.Add(2)
			go func() {
				defer GinkgoRecover()
				pushRegistryApp(appRegistry)
				setupWG.Done()
			}()

			go func() {
				defer GinkgoRecover()
				pushApp(appProxy)
				setupWG.Done()
			}()
			setupWG.Wait()

			By("pushing the tick apps")
			newManifest := modifyTickManifest(appRegistry)
			runWithTimeout("push tick apps", 5*time.Minute, func() {
				pushAppsOfType(appsTest, "tick", newManifest)
			})

			By("scaling the tick apps")
			scaleApps(appsTest, appInstances)

			By("checking that all test app instances have registered themselves")
			checkRegistry(appRegistry, 60*time.Second, 500*time.Millisecond, len(appsTest)*appInstances)

			appIPs := getAppIPs(appRegistry)

			By("checking that the connection fails")
			runWithTimeout("check connection failures", 5*time.Minute, func() {
				assertConnectionFails(appProxy, appIPs, ports)
			})

			By("creating policies")
			doAllPolicies("create", appProxy, appsTest, ports)

			By(fmt.Sprintf("checking that %s can reach %s", appProxy, appsTest))
			runWithTimeout("check connection success", 5*time.Minute, func() {
				assertConnectionSucceeds(appProxy, appIPs, ports)
			})

			dumpStats(appProxy, config.AppsDomain)

			By("deleting policies")
			doAllPolicies("delete", appProxy, appsTest, ports)

			By(fmt.Sprintf("checking that %s can NOT reach %s", appProxy, appsTest))
			runWithTimeout("check connection failures, again", 5*time.Minute, func() {
				assertConnectionFails(appProxy, appIPs, ports)
			})

			By("checking that reflex no longer reports deleted instances")
			scaleApps(appsTest, 1 /* instances */)
			checkRegistry(appRegistry, 60*time.Second, 500*time.Millisecond, len(appsTest))

			close(done)
		}, 30*60 /* <-- overall spec timeout in seconds */)
	})
})

func getToken() string {
	session := cf.Cf("oauth-token").Wait(2 * Timeout_Short)
	Expect(session.Wait(Timeout_Short)).To(gexec.Exit(0))
	return string(session.Out.Contents())
}

func getGuids(sourceAppName string, dstAppNames []string) (string, []string) {
	dstGuids := []string{}
	sourceGuid := ""
	token := getToken()
	token = strings.TrimPrefix(token, "bearer ")
	appsClient := rainmaker.NewApplicationsService(rainmaker.Config{Host: "http://" + config.ApiEndpoint})

	appsList, err := appsClient.List(token)
	Expect(err).NotTo(HaveOccurred())

	for {
		for _, app := range appsList.Applications {
			if app.Name == sourceAppName {
				sourceGuid = app.GUID
				continue
			}
			for _, tickAppName := range dstAppNames {
				if app.Name == tickAppName {
					dstGuids = append(dstGuids, app.GUID)
					break
				}
			}
		}
		if appsList.HasNextPage() {
			appsList, err = appsList.Next(token)
			Expect(err).NotTo(HaveOccurred())
		} else {
			break
		}
	}

	Expect(sourceGuid).NotTo(BeEmpty())
	Expect(dstGuids).To(HaveLen(len(dstAppNames)))

	return sourceGuid, dstGuids
}

func doAllPolicies(action string, source string, dstList []string, dstPorts []int) {
	policyClient := policy_client.New(lagertest.NewTestLogger("test"), &http.Client{}, "http://"+config.ApiEndpoint)
	sourceGuid, dstGuids := getGuids(source, dstList)
	policies := []models.Policy{}
	for _, dstGuid := range dstGuids {
		for _, port := range dstPorts {
			policies = append(policies, models.Policy{
				Source: models.Source{
					ID: sourceGuid,
				},
				Destination: models.Destination{
					ID:       dstGuid,
					Port:     port,
					Protocol: "tcp",
				},
			})
		}
	}
	token := getToken()
	if action == "create" {
		Expect(policyClient.AddPolicies(policies, token)).To(Succeed())
	} else if action == "delete" {
		Expect(policyClient.DeletePolicies(policies, token)).To(Succeed())
	}
}

func runWithTimeout(operation string, timeout time.Duration, work func()) {
	done := make(chan bool)
	go func() {
		fmt.Printf("starting %s\n", operation)
		work()
		done <- true
	}()

	select {
	case <-done:
		fmt.Printf("completed %s\n", operation)
		return
	case <-time.After(timeout):
		Fail("timeout on " + operation)
	}
}

func dumpStats(host, domain string) {
	respBytes, err := httpGetBytes(fmt.Sprintf("http://%s.%s/stats", host, domain))
	Expect(err).NotTo(HaveOccurred())

	fmt.Printf("STATS: %s\n", string(respBytes))
	netStatsFile := os.Getenv("NETWORK_STATS_FILE")
	if netStatsFile != "" {
		Expect(ioutil.WriteFile(netStatsFile, respBytes, 0600)).To(Succeed())
	}
}

type RegistryInstancesResponse struct {
	Instances []struct {
		ServiceName string `json:"service_name"`
		Endpoint    struct {
			Value string `json:"value"`
		} `json:"endpoint"`
	} `json:"instances"`
}

func getInstancesFromA8(registry string) (*RegistryInstancesResponse, error) {
	respBytes, err := httpGetBytes(fmt.Sprintf("http://%s.%s/api/v1/instances", registry, config.AppsDomain))
	if err != nil {
		return nil, err
	}

	var instancesResponse RegistryInstancesResponse
	err = json.Unmarshal(respBytes, &instancesResponse)
	if err != nil {
		return nil, err
	}
	return &instancesResponse, nil
}

func checkRegistry(registry string, timeout, pollingInterval time.Duration, totalInstances int) {
	registeredApps := func() (int, error) {
		instancesResponse, err := getInstancesFromA8(registry)
		if err != nil {
			return 0, err
		}
		return len(instancesResponse.Instances), nil
	}

	Eventually(registeredApps, timeout, pollingInterval).Should(Equal(totalInstances))
}

func getAppIPs(registry string) []string {
	instancesResponse, err := getInstancesFromA8(registry)
	Expect(err).NotTo(HaveOccurred())

	ips := []string{}
	for _, instance := range instancesResponse.Instances {
		ip, _, err := net.SplitHostPort(instance.Endpoint.Value)
		Expect(err).NotTo(HaveOccurred())
		ips = append(ips, ip)
	}
	return ips
}

func assertConnectionSucceeds(sourceApp string, destApps []string, ports []int) {
	workPoolRun(5, destApps, func(appIP string) {
		for _, port := range ports {
			assertSingleConnection(appIP, port, sourceApp, true)
		}
	})
}

func assertConnectionFails(sourceApp string, destApps []string, ports []int) {
	workPoolRun(5, destApps, func(appIP string) {
		for _, port := range ports {
			assertSingleConnection(appIP, port, sourceApp, false)
		}
	})
}

func assertSingleConnection(destIP string, port int, sourceAppName string, shouldSucceed bool) {
	proxyTest := func() (string, error) {
		respBytes, err := httpGetBytes(fmt.Sprintf("http://%s.%s/proxy/%s:%d", sourceAppName, config.AppsDomain, destIP, port))
		if err != nil {
			return "", err
		}
		return string(respBytes), nil
	}
	if shouldSucceed {
		By(fmt.Sprintf("eventually proxy should reach %s at port %d", destIP, port))
		Eventually(proxyTest, 10*time.Second, 500*time.Millisecond).Should(ContainSubstring(`application_name`))
	} else {
		By(fmt.Sprintf("eventually proxy should NOT reach %s at port %d", destIP, port))
		Eventually(proxyTest, 10*time.Second, 500*time.Millisecond).Should(ContainSubstring("request failed"))
	}
}

var httpClient = &http.Client{
	Transport: &http.Transport{
		DisableKeepAlives: true,
		Dial: (&net.Dialer{
			Timeout:   4 * time.Second,
			KeepAlive: 0,
		}).Dial,
	},
}

func httpGetBytes(url string) ([]byte, error) {
	resp, err := httpClient.Get(url)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	respBytes, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}
	return respBytes, nil
}
<commit_msg>Don't log tokens to stdout in tests<commit_after>package acceptance_test

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"lib/models"
	"lib/policy_client"
	"math/rand"
	"net"
	"net/http"
	"os"
	"os/exec"
	"strings"
	"sync"
	"time"

	"code.cloudfoundry.org/lager/lagertest"

	"github.com/cloudfoundry-incubator/cf-test-helpers/cf"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gexec"
	"github.com/pivotal-cf-experimental/rainmaker"
)

const Timeout_Push = 5 * time.Minute
const Timeout_Short = 10 * time.Second

var ports []int

func getSubnet(ip string) string {
	return strings.Split(ip, ".")[2]
}

func isSameCell(sourceIP, destIP string) bool {
	return getSubnet(sourceIP) == getSubnet(destIP)
}

var _ = Describe("connectivity between containers on the overlay network", func() {
	Describe("networking policy", func() {
		var (
			appProxy     string
			appRegistry  string
			appsTest     []string
			orgName      string
			spaceName    string
			appInstances int
			applications int
		)

		BeforeEach(func() {
			appInstances = testConfig.AppInstances
			applications = testConfig.Applications

			appProxy = fmt.Sprintf("proxy-%d", rand.Int31())
			appRegistry = fmt.Sprintf("registry-%d", rand.Int31())
			for i := 0; i < applications; i++ {
				appsTest = append(appsTest, fmt.Sprintf("tick-%d-%d", i, rand.Int31()))
			}

			ports = []int{8080}
			for i := 0; i < testConfig.Policies; i++ {
				ports = append(ports, 7000+i)
			}

			Auth(testConfig.TestUser, testConfig.TestUserPassword)

			orgName = "test-org"
			Expect(cf.Cf("create-org", orgName).Wait(Timeout_Push)).To(gexec.Exit(0))
			Expect(cf.Cf("target", "-o", orgName).Wait(Timeout_Push)).To(gexec.Exit(0))

			spaceName = "test-space"
			Expect(cf.Cf("create-space", spaceName).Wait(Timeout_Push)).To(gexec.Exit(0))
			Expect(cf.Cf("target", "-o", orgName, "-s", spaceName).Wait(Timeout_Push)).To(gexec.Exit(0))
		})

		AfterEach(func() {
			appReport(appProxy, Timeout_Short)
			appReport(appRegistry, Timeout_Short)
			appsReport(appsTest, Timeout_Short)

			// clean up everything
			Expect(cf.Cf("delete-org", orgName, "-f").Wait(Timeout_Push)).To(gexec.Exit(0))
		})

		It("allows the user to configure policies", func(done Done) {
			By("pushing the registry app and proxy app")
			var setupWG sync.WaitGroup
			setupWG.Add(2)
			go func() {
				defer GinkgoRecover()
				pushRegistryApp(appRegistry)
				setupWG.Done()
			}()

			go func() {
				defer GinkgoRecover()
				pushApp(appProxy)
				setupWG.Done()
			}()
			setupWG.Wait()

			By("pushing the tick apps")
			newManifest := modifyTickManifest(appRegistry)
			runWithTimeout("push tick apps", 5*time.Minute, func() {
				pushAppsOfType(appsTest, "tick", newManifest)
			})

			By("scaling the tick apps")
			scaleApps(appsTest, appInstances)

			By("checking that all test app instances have registered themselves")
			checkRegistry(appRegistry, 60*time.Second, 500*time.Millisecond, len(appsTest)*appInstances)

			appIPs := getAppIPs(appRegistry)

			By("checking that the connection fails")
			runWithTimeout("check connection failures", 5*time.Minute, func() {
				assertConnectionFails(appProxy, appIPs, ports)
			})

			By("creating policies")
			doAllPolicies("create", appProxy, appsTest, ports)

			By(fmt.Sprintf("checking that %s can reach %s", appProxy, appsTest))
			runWithTimeout("check connection success", 5*time.Minute, func() {
				assertConnectionSucceeds(appProxy, appIPs, ports)
			})

			dumpStats(appProxy, config.AppsDomain)

			By("deleting policies")
			doAllPolicies("delete", appProxy, appsTest, ports)

			By(fmt.Sprintf("checking that %s can NOT reach %s", appProxy, appsTest))
			runWithTimeout("check connection failures, again", 5*time.Minute, func() {
				assertConnectionFails(appProxy, appIPs, ports)
			})

			By("checking that reflex no longer reports deleted instances")
			scaleApps(appsTest, 1 /* instances */)
			checkRegistry(appRegistry, 60*time.Second, 500*time.Millisecond, len(appsTest))

			close(done)
		}, 30*60 /* <-- overall spec timeout in seconds */)
	})
})

func getToken() string {
	By("getting token")
	cmd := exec.Command("cf", "oauth-token")
	session, err := gexec.Start(cmd, nil, nil)
	Expect(err).NotTo(HaveOccurred())
	Eventually(session.Wait(2 * Timeout_Short)).Should(gexec.Exit(0))
	return string(session.Out.Contents())
}

func getGuids(sourceAppName string, dstAppNames []string) (string, []string) {
	dstGuids := []string{}
	sourceGuid := ""
	token := getToken()
	token = strings.TrimPrefix(token, "bearer ")
	appsClient := rainmaker.NewApplicationsService(rainmaker.Config{Host: "http://" + config.ApiEndpoint})

	appsList, err := appsClient.List(token)
	Expect(err).NotTo(HaveOccurred())

	for {
		for _, app := range appsList.Applications {
			if app.Name == sourceAppName {
				sourceGuid = app.GUID
				continue
			}
			for _, tickAppName := range dstAppNames {
				if app.Name == tickAppName {
					dstGuids = append(dstGuids, app.GUID)
					break
				}
			}
		}
		if appsList.HasNextPage() {
			appsList, err = appsList.Next(token)
			Expect(err).NotTo(HaveOccurred())
		} else {
			break
		}
	}

	Expect(sourceGuid).NotTo(BeEmpty())
	Expect(dstGuids).To(HaveLen(len(dstAppNames)))

	return sourceGuid, dstGuids
}

func doAllPolicies(action string, source string, dstList []string, dstPorts []int) {
	policyClient := policy_client.New(lagertest.NewTestLogger("test"), &http.Client{}, "http://"+config.ApiEndpoint)
	sourceGuid, dstGuids := getGuids(source, dstList)
	policies := []models.Policy{}
	for _, dstGuid := range dstGuids {
		for _, port := range dstPorts {
			policies = append(policies, models.Policy{
				Source: models.Source{
					ID: sourceGuid,
				},
				Destination: models.Destination{
					ID:       dstGuid,
					Port:     port,
					Protocol: "tcp",
				},
			})
		}
	}
	token := getToken()
	if action == "create" {
		Expect(policyClient.AddPolicies(policies, token)).To(Succeed())
	} else if action == "delete" {
		Expect(policyClient.DeletePolicies(policies, token)).To(Succeed())
	}
}

func runWithTimeout(operation string, timeout time.Duration, work func()) {
	done := make(chan bool)
	go func() {
		fmt.Printf("starting %s\n", operation)
		work()
		done <- true
	}()

	select {
	case <-done:
		fmt.Printf("completed %s\n", operation)
		return
	case <-time.After(timeout):
		Fail("timeout on " + operation)
	}
}

func dumpStats(host, domain string) {
	respBytes, err := httpGetBytes(fmt.Sprintf("http://%s.%s/stats", host, domain))
	Expect(err).NotTo(HaveOccurred())

	fmt.Printf("STATS: %s\n", string(respBytes))
	netStatsFile := os.Getenv("NETWORK_STATS_FILE")
	if netStatsFile != "" {
		Expect(ioutil.WriteFile(netStatsFile, respBytes, 0600)).To(Succeed())
	}
}

type RegistryInstancesResponse struct {
	Instances []struct {
		ServiceName string `json:"service_name"`
		Endpoint    struct {
			Value string `json:"value"`
		} `json:"endpoint"`
	} `json:"instances"`
}

func getInstancesFromA8(registry string) (*RegistryInstancesResponse, error) {
	respBytes, err := httpGetBytes(fmt.Sprintf("http://%s.%s/api/v1/instances", registry, config.AppsDomain))
	if err != nil {
		return nil, err
	}

	var instancesResponse RegistryInstancesResponse
	err = json.Unmarshal(respBytes, &instancesResponse)
	if err != nil {
		return nil, err
	}
	return &instancesResponse, nil
}

func checkRegistry(registry string, timeout, pollingInterval time.Duration, totalInstances int) {
	registeredApps := func() (int, error) {
		instancesResponse, err := getInstancesFromA8(registry)
		if err != nil {
			return 0, err
		}
		return len(instancesResponse.Instances), nil
	}

	Eventually(registeredApps, timeout, pollingInterval).Should(Equal(totalInstances))
}

func getAppIPs(registry string) []string {
	instancesResponse, err := getInstancesFromA8(registry)
	Expect(err).NotTo(HaveOccurred())

	ips := []string{}
	for _, instance := range instancesResponse.Instances {
		ip, _, err := net.SplitHostPort(instance.Endpoint.Value)
		Expect(err).NotTo(HaveOccurred())
		ips = append(ips, ip)
	}
	return ips
}

func assertConnectionSucceeds(sourceApp string, destApps []string, ports []int) {
	workPoolRun(5, destApps, func(appIP string) {
		for _, port := range ports {
			assertSingleConnection(appIP, port, sourceApp, true)
		}
	})
}

func assertConnectionFails(sourceApp string, destApps []string, ports []int) {
	workPoolRun(5, destApps, func(appIP string) {
		for _, port := range ports {
			assertSingleConnection(appIP, port, sourceApp, false)
		}
	})
}

func assertSingleConnection(destIP string, port int, sourceAppName string, shouldSucceed bool) {
	proxyTest := func() (string, error) {
		respBytes, err := httpGetBytes(fmt.Sprintf("http://%s.%s/proxy/%s:%d", sourceAppName, config.AppsDomain, destIP, port))
		if err != nil {
			return "", err
		}
		return string(respBytes), nil
	}
	if shouldSucceed {
		By(fmt.Sprintf("eventually proxy should reach %s at port %d", destIP, port))
		Eventually(proxyTest, 10*time.Second, 500*time.Millisecond).Should(ContainSubstring(`application_name`))
	} else {
		By(fmt.Sprintf("eventually proxy should NOT reach %s at port %d", destIP, port))
		Eventually(proxyTest, 10*time.Second, 500*time.Millisecond).Should(ContainSubstring("request failed"))
	}
}

var httpClient = &http.Client{
	Transport: &http.Transport{
		DisableKeepAlives: true,
		Dial: (&net.Dialer{
			Timeout:   4 * time.Second,
			KeepAlive: 0,
		}).Dial,
	},
}

func httpGetBytes(url string) ([]byte, error) {
	resp, err := httpClient.Get(url)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	respBytes, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}
	return respBytes, nil
}
<|endoftext|>


<commit_before>package atlas

import (
	"github.com/hashicorp/atlas-go/v1"
	"github.com/hashicorp/terraform/helper/schema"
	"github.com/hashicorp/terraform/terraform"
)

const (
	// defaultAtlasServer is the default endpoint for Atlas if
	// none is specified.
	defaultAtlasServer = "https://atlas.hashicorp.com"
)

// Provider returns a terraform.ResourceProvider.
func Provider() terraform.ResourceProvider {
	return &schema.Provider{
		Schema: map[string]*schema.Schema{
			"token": &schema.Schema{
				Type:        schema.TypeString,
				Required:    true,
				DefaultFunc: schema.EnvDefaultFunc("ATLAS_TOKEN", nil),
				Description: descriptions["token"],
			},

			"address": &schema.Schema{
				Type:        schema.TypeString,
				Optional:    true,
				DefaultFunc: schema.EnvDefaultFunc("ATLAS_ADDRESS", defaultAtlasServer),
				Description: descriptions["address"],
			},
		},

		DataSourcesMap: map[string]*schema.Resource{
			"atlas_artifact": dataSourceAtlasArtifact(),
		},

		ResourcesMap: map[string]*schema.Resource{
			"atlas_artifact": resourceArtifact(),
		},

		ConfigureFunc: providerConfigure,
	}
}

func providerConfigure(d *schema.ResourceData) (interface{}, error) {
	var err error
	client := atlas.DefaultClient()
	if v := d.Get("address").(string); v != "" {
		client, err = atlas.NewClient(v)
		if err != nil {
			return nil, err
		}
	}
	client.Token = d.Get("token").(string)

	return client, nil
}

var descriptions map[string]string

func init() {
	descriptions = map[string]string{
		"address": "The address of the Atlas server. If blank, the public\n" +
			"server at atlas.hashicorp.com will be used.",

		"token": "The access token for reading artifacts. This is required\n" +
			"if reading private artifacts.",
	}
}
<commit_msg>Add the VersionHeader to the atlas provider too<commit_after>package atlas

import (
	"github.com/hashicorp/atlas-go/v1"
	"github.com/hashicorp/terraform/helper/schema"
	"github.com/hashicorp/terraform/terraform"
)

const (
	// defaultAtlasServer is the default endpoint for Atlas if
	// none is specified.
	defaultAtlasServer = "https://atlas.hashicorp.com"
)

// Provider returns a terraform.ResourceProvider.
func Provider() terraform.ResourceProvider {
	return &schema.Provider{
		Schema: map[string]*schema.Schema{
			"token": &schema.Schema{
				Type:        schema.TypeString,
				Required:    true,
				DefaultFunc: schema.EnvDefaultFunc("ATLAS_TOKEN", nil),
				Description: descriptions["token"],
			},

			"address": &schema.Schema{
				Type:        schema.TypeString,
				Optional:    true,
				DefaultFunc: schema.EnvDefaultFunc("ATLAS_ADDRESS", defaultAtlasServer),
				Description: descriptions["address"],
			},
		},

		DataSourcesMap: map[string]*schema.Resource{
			"atlas_artifact": dataSourceAtlasArtifact(),
		},

		ResourcesMap: map[string]*schema.Resource{
			"atlas_artifact": resourceArtifact(),
		},

		ConfigureFunc: providerConfigure,
	}
}

func providerConfigure(d *schema.ResourceData) (interface{}, error) {
	var err error
	client := atlas.DefaultClient()
	if v := d.Get("address").(string); v != "" {
		client, err = atlas.NewClient(v)
		if err != nil {
			return nil, err
		}
	}
	client.DefaultHeader.Set(terraform.VersionHeader, terraform.Version)
	client.Token = d.Get("token").(string)

	return client, nil
}

var descriptions map[string]string

func init() {
	descriptions = map[string]string{
		"address": "The address of the Atlas server. If blank, the public\n" +
			"server at atlas.hashicorp.com will be used.",

		"token": "The access token for reading artifacts. This is required\n" +
			"if reading private artifacts.",
	}
}
<|endoftext|>