<commit_before>package gcloudcleanup

import (
	"context"
	"errors"
	"math/rand"
	"os"
	"strings"
	"time"

	"cloud.google.com/go/storage"

	"contrib.go.opencensus.io/exporter/stackdriver"
	"go.opencensus.io/trace"

	"google.golang.org/api/compute/v1"
	"google.golang.org/api/option"
	"gopkg.in/urfave/cli.v2"

	"github.com/mihasya/go-metrics-librato"
	"github.com/rcrowley/go-metrics"
	"github.com/sirupsen/logrus"
	travismetrics "github.com/travis-ci/gcloud-cleanup/metrics"
	"github.com/travis-ci/gcloud-cleanup/ratelimit"
)

var (
	errInvalidInstancesMaxAge   = errors.New("invalid max age")
	errInvalidArchiveSampleRate = errors.New("invalid archive sample rate")
	errInvalidTraceSampleRate   = errors.New("invalid trace sample rate")
)

type CLI struct {
	c           *cli.Context
	ctx         context.Context
	cs          *compute.Service
	sc          *storage.Client
	log         *logrus.Logger
	rateLimiter ratelimit.RateLimiter

	instanceCleaner *instanceCleaner
	imageCleaner    *imageCleaner
}

func NewCLI(c *cli.Context) *CLI {
	log := logrus.New()
	log.Level = logrus.InfoLevel
	log.Formatter = &logrus.TextFormatter{DisableColors: true}

	return &CLI{
		c:   c,
		ctx: context.Background(),
		log: log,
	}
}

func (c *CLI) Run() error {
	c.setupLogger()
	c.setupRateLimiter()

	fields := logrus.Fields{}

	for _, name := range c.c.FlagNames() {
		fields[name] = c.c.Generic(name)
	}

	c.log.WithFields(fields).Debug("configuration")

	c.setupMetrics()

	err := c.setupOpenCensus(c.c.String("account-json"))
	if err != nil {
		c.log.WithField("err", err).Fatal("failed to set up opencensus")
	}

	err = c.setupComputeService(c.c.String("account-json"))
	if err != nil {
		c.log.WithField("err", err).Fatal("failed to set up compute service")
	}

	err = c.setupStorageClient(c.c.String("account-json"))
	if err != nil {
		c.log.WithField("err", err).Fatal("failed to set up storage client")
	}

	sleepDur := c.c.Duration("loop-sleep")
	if sleepDur == (0 * time.Second) {
		sleepDur = 5 * time.Minute
		c.log.WithField("loop_sleep", sleepDur).Info("default loop sleep set")
	}

	once := c.c.Bool("once")

	entities := c.c.StringSlice("entities")
	if len(entities) == 0 {
		entities = []string{"instances"}
		c.log.WithField("entities", entities).Info("default entities set")
	}

	entityMap := map[string]func() error{
		"instances": c.cleanupInstances,
		"images":    c.cleanupImages,
	}

	for {
		for _, entity := range entities {
			if f, ok := entityMap[entity]; ok {
				c.log.WithField("type", entity).Debug("entering entity loop")

				err := f()

				if err != nil {
					c.log.WithFields(logrus.Fields{
						"type": entity,
						"err":  err,
					}).Fatal("failure during entity cleanup")
				}
			} else {
				c.log.WithField("type", entity).Fatal("unknown entity type")
			}

			c.log.WithField("type", entity).Debug("done with entity loop")
		}

		if once {
			break
		}

		c.log.WithField("duration", sleepDur).Info("sleeping")
		time.Sleep(sleepDur)
	}
	return nil
}

func (c *CLI) setupComputeService(accountJSON string) error {
	cs, err := buildGoogleComputeService(accountJSON)
	c.cs = cs
	return err
}

func (c *CLI) setupStorageClient(accountJSON string) error {
	sc, err := buildGoogleStorageClient(c.ctx, accountJSON)
	c.sc = sc
	return err
}

func (c *CLI) setupLogger() {
	if c.c.Bool("debug") {
		c.log.Level = logrus.DebugLevel
	}
}

func (c *CLI) setupRateLimiter() {
	if c.c.String("rate-limit-redis-url") == "" {
		c.rateLimiter = ratelimit.NewNullRateLimiter()
		return
	}
	c.rateLimiter = ratelimit.NewRateLimiter(
		c.c.String("rate-limit-redis-url"),
		c.c.String("rate-limit-prefix"))
}

func (c *CLI) setupOpenCensus(accountJSON string) error {
	traceEnabled := c.c.Bool("opencensus-enable")

	if traceEnabled {
		creds, err := buildGoogleCloudCredentials(context.TODO(), accountJSON)
		if err != nil {
			return err
		}

		sd, err := stackdriver.NewExporter(stackdriver.Options{
			ProjectID: os.Getenv("GCLOUD_PROJECT"),
			TraceClientOptions: []option.ClientOption{
				option.WithCredentials(creds),
			},
			MonitoringClientOptions: []option.ClientOption{
				option.WithCredentials(creds),
			},
		})

		if err != nil {
			return err
		}

		defer sd.Flush()

		// Register/enable the trace exporter
		trace.RegisterExporter(sd)

		traceSampleRate := c.c.Int64("trace-sample-rate")
		if traceSampleRate <= 0 {
			c.log.WithFields(logrus.Fields{
				"trace_sample_rate": traceSampleRate,
			}).Error("trace sample rate must be positive")
			return errInvalidTraceSampleRate
		}

		trace.ApplyConfig(trace.Config{DefaultSampler: trace.ProbabilitySampler(1.0 / float64(traceSampleRate))})
		return nil
	}
	return nil
}

func (c *CLI) cleanupInstances() error {
	if c.instanceCleaner == nil {
		filters := c.c.StringSlice("instance-filters")
		if len(filters) == 0 {
			filters = []string{"name eq ^testing-gce.*"}
			c.log.WithField("filters", strings.Join(filters, ",")).Info("default filters set")
		}

		cutoffTime := time.Now().UTC().Add(-1 * c.c.Duration("instance-max-age"))

		if time.Now().UTC().Before(cutoffTime) {
			c.log.WithFields(logrus.Fields{
				"cutoff":  cutoffTime,
				"max_age": c.c.Duration("instance-max-age"),
			}).Error("invalid instance max age given")
			return errInvalidInstancesMaxAge
		}

		archiveSampleRate := c.c.Int64("archive-sample-rate")
		if archiveSampleRate <= 0 {
			c.log.WithFields(logrus.Fields{
				"sample_rate": archiveSampleRate,
			}).Error("archive sample rate must be positive")
			return errInvalidArchiveSampleRate
		}

		c.log.WithFields(logrus.Fields{
			"max_age":    c.c.Duration("instance-max-age"),
			"tick":       c.c.Duration("rate-tick-limit"),
			"project_id": c.c.String("project-id"),
			"filters":    strings.Join(filters, ","),
			"cutoff":     cutoffTime.Format(time.RFC3339),
		}).Debug("creating instance cleaner with")

		c.instanceCleaner = &instanceCleaner{
			ctx: c.ctx,
			cs:  c.cs,
			sc:  c.sc,
			log: c.log.WithField("component", "instance_cleaner"),

			rand: rand.New(rand.NewSource(time.Now().UnixNano())),

			projectID: c.c.String("project-id"),
			filters:   filters,

			archiveSerial:     c.c.Bool("archive-serial"),
			archiveBucket:     c.c.String("archive-bucket"),
			archiveSampleRate: archiveSampleRate,
			noop:              c.c.Bool("noop"),

			CutoffTime: cutoffTime,

			rateLimiter:       c.rateLimiter,
			rateLimitMaxCalls: uint64(c.c.Int("rate-limit-max-calls")),
			rateLimitDuration: c.c.Duration("rate-limit-duration"),
		}
	}

	c.instanceCleaner.CutoffTime = time.Now().UTC().Add(-1 * c.c.Duration("instance-max-age"))

	return c.instanceCleaner.Run()
}

func (i *CLI) setupMetrics() {
	go travismetrics.ReportMemstatsMetrics()

	if os.Getenv("LIBRATO_EMAIL") != "" && os.Getenv("LIBRATO_TOKEN") != "" && os.Getenv("LIBRATO_SOURCE") != "" {
		i.log.Info("starting librato metrics reporter")

		go librato.Librato(metrics.DefaultRegistry, time.Minute,
			os.Getenv("LIBRATO_EMAIL"), os.Getenv("LIBRATO_TOKEN"), os.Getenv("LIBRATO_SOURCE"),
			[]float64{0.50, 0.75, 0.90, 0.95, 0.99, 0.999, 1.0}, time.Millisecond)
	}
}

func (c *CLI) cleanupImages() error {
	if c.imageCleaner == nil {
		filters := c.c.StringSlice("image-filters")
		if len(filters) == 0 {
			filters = []string{"name eq ^travis-ci.*"}
			c.log.WithField("filters", strings.Join(filters, ",")).Info("default filters set")
		}

		c.imageCleaner = newImageCleaner(c.cs,
			c.log, c.rateLimiter, uint64(c.c.Int("rate-limit-max-calls")), c.c.Duration("rate-limit-duration"), c.c.String("project-id"),
			c.c.String("job-board-url"), filters, c.c.Bool("noop"))
	}

	return c.imageCleaner.Run()
}
<commit_msg>Inverse the feature flag checking to skip tracing<commit_after>package gcloudcleanup

import (
	"context"
	"errors"
	"math/rand"
	"os"
	"strings"
	"time"

	"cloud.google.com/go/storage"

	"contrib.go.opencensus.io/exporter/stackdriver"
	"go.opencensus.io/trace"

	"google.golang.org/api/compute/v1"
	"google.golang.org/api/option"
	"gopkg.in/urfave/cli.v2"

	"github.com/mihasya/go-metrics-librato"
	"github.com/rcrowley/go-metrics"
	"github.com/sirupsen/logrus"
	travismetrics "github.com/travis-ci/gcloud-cleanup/metrics"
	"github.com/travis-ci/gcloud-cleanup/ratelimit"
)

var (
	errInvalidInstancesMaxAge   = errors.New("invalid max age")
	errInvalidArchiveSampleRate = errors.New("invalid archive sample rate")
	errInvalidTraceSampleRate   = errors.New("invalid trace sample rate")
)

type CLI struct {
	c           *cli.Context
	ctx         context.Context
	cs          *compute.Service
	sc          *storage.Client
	log         *logrus.Logger
	rateLimiter ratelimit.RateLimiter

	instanceCleaner *instanceCleaner
	imageCleaner    *imageCleaner
}

func NewCLI(c *cli.Context) *CLI {
	log := logrus.New()
	log.Level = logrus.InfoLevel
	log.Formatter = &logrus.TextFormatter{DisableColors: true}

	return &CLI{
		c:   c,
		ctx: context.Background(),
		log: log,
	}
}

func (c *CLI) Run() error {
	c.setupLogger()
	c.setupRateLimiter()

	fields := logrus.Fields{}

	for _, name := range c.c.FlagNames() {
		fields[name] = c.c.Generic(name)
	}

	c.log.WithFields(fields).Debug("configuration")

	c.setupMetrics()

	err := c.setupOpenCensus(c.c.String("account-json"))
	if err != nil {
		c.log.WithField("err", err).Fatal("failed to set up opencensus")
	}

	err = c.setupComputeService(c.c.String("account-json"))
	if err != nil {
		c.log.WithField("err", err).Fatal("failed to set up compute service")
	}

	err = c.setupStorageClient(c.c.String("account-json"))
	if err != nil {
		c.log.WithField("err", err).Fatal("failed to set up storage client")
	}

	sleepDur := c.c.Duration("loop-sleep")
	if sleepDur == (0 * time.Second) {
		sleepDur = 5 * time.Minute
		c.log.WithField("loop_sleep", sleepDur).Info("default loop sleep set")
	}

	once := c.c.Bool("once")

	entities := c.c.StringSlice("entities")
	if len(entities) == 0 {
		entities = []string{"instances"}
		c.log.WithField("entities", entities).Info("default entities set")
	}

	entityMap := map[string]func() error{
		"instances": c.cleanupInstances,
		"images":    c.cleanupImages,
	}

	for {
		for _, entity := range entities {
			if f, ok := entityMap[entity]; ok {
				c.log.WithField("type", entity).Debug("entering entity loop")

				err := f()

				if err != nil {
					c.log.WithFields(logrus.Fields{
						"type": entity,
						"err":  err,
					}).Fatal("failure during entity cleanup")
				}
			} else {
				c.log.WithField("type", entity).Fatal("unknown entity type")
			}

			c.log.WithField("type", entity).Debug("done with entity loop")
		}

		if once {
			break
		}

		c.log.WithField("duration", sleepDur).Info("sleeping")
		time.Sleep(sleepDur)
	}
	return nil
}

func (c *CLI) setupComputeService(accountJSON string) error {
	cs, err := buildGoogleComputeService(accountJSON)
	c.cs = cs
	return err
}

func (c *CLI) setupStorageClient(accountJSON string) error {
	sc, err := buildGoogleStorageClient(c.ctx, accountJSON)
	c.sc = sc
	return err
}

func (c *CLI) setupLogger() {
	if c.c.Bool("debug") {
		c.log.Level = logrus.DebugLevel
	}
}

func (c *CLI) setupRateLimiter() {
	if c.c.String("rate-limit-redis-url") == "" {
		c.rateLimiter = ratelimit.NewNullRateLimiter()
		return
	}
	c.rateLimiter = ratelimit.NewRateLimiter(
		c.c.String("rate-limit-redis-url"),
		c.c.String("rate-limit-prefix"))
}

func (c *CLI) setupOpenCensus(accountJSON string) error {
	traceEnabled := c.c.Bool("opencensus-enable")

	if !traceEnabled {
		return nil
	}

	creds, err := buildGoogleCloudCredentials(context.TODO(), accountJSON)
	if err != nil {
		return err
	}

	sd, err := stackdriver.NewExporter(stackdriver.Options{
		ProjectID: os.Getenv("GCLOUD_PROJECT"),
		TraceClientOptions: []option.ClientOption{
			option.WithCredentials(creds),
		},
		MonitoringClientOptions: []option.ClientOption{
			option.WithCredentials(creds),
		},
	})

	if err != nil {
		return err
	}

	defer sd.Flush()

	// Register/enable the trace exporter
	trace.RegisterExporter(sd)

	traceSampleRate := c.c.Int64("trace-sample-rate")
	if traceSampleRate <= 0 {
		c.log.WithFields(logrus.Fields{
			"trace_sample_rate": traceSampleRate,
		}).Error("trace sample rate must be positive")
		return errInvalidTraceSampleRate
	}

	trace.ApplyConfig(trace.Config{DefaultSampler: trace.ProbabilitySampler(1.0 / float64(traceSampleRate))})

	return nil
}

func (c *CLI) cleanupInstances() error {
	if c.instanceCleaner == nil {
		filters := c.c.StringSlice("instance-filters")
		if len(filters) == 0 {
			filters = []string{"name eq ^testing-gce.*"}
			c.log.WithField("filters", strings.Join(filters, ",")).Info("default filters set")
		}

		cutoffTime := time.Now().UTC().Add(-1 * c.c.Duration("instance-max-age"))

		if time.Now().UTC().Before(cutoffTime) {
			c.log.WithFields(logrus.Fields{
				"cutoff":  cutoffTime,
				"max_age": c.c.Duration("instance-max-age"),
			}).Error("invalid instance max age given")
			return errInvalidInstancesMaxAge
		}

		archiveSampleRate := c.c.Int64("archive-sample-rate")
		if archiveSampleRate <= 0 {
			c.log.WithFields(logrus.Fields{
				"sample_rate": archiveSampleRate,
			}).Error("archive sample rate must be positive")
			return errInvalidArchiveSampleRate
		}

		c.log.WithFields(logrus.Fields{
			"max_age":    c.c.Duration("instance-max-age"),
			"tick":       c.c.Duration("rate-tick-limit"),
			"project_id": c.c.String("project-id"),
			"filters":    strings.Join(filters, ","),
			"cutoff":     cutoffTime.Format(time.RFC3339),
		}).Debug("creating instance cleaner with")

		c.instanceCleaner = &instanceCleaner{
			ctx: c.ctx,
			cs:  c.cs,
			sc:  c.sc,
			log: c.log.WithField("component", "instance_cleaner"),

			rand: rand.New(rand.NewSource(time.Now().UnixNano())),

			projectID: c.c.String("project-id"),
			filters:   filters,

			archiveSerial:     c.c.Bool("archive-serial"),
			archiveBucket:     c.c.String("archive-bucket"),
			archiveSampleRate: archiveSampleRate,
			noop:              c.c.Bool("noop"),

			CutoffTime: cutoffTime,

			rateLimiter:       c.rateLimiter,
			rateLimitMaxCalls: uint64(c.c.Int("rate-limit-max-calls")),
			rateLimitDuration: c.c.Duration("rate-limit-duration"),
		}
	}

	c.instanceCleaner.CutoffTime = time.Now().UTC().Add(-1 * c.c.Duration("instance-max-age"))

	return c.instanceCleaner.Run()
}

func (i *CLI) setupMetrics() {
	go travismetrics.ReportMemstatsMetrics()

	if os.Getenv("LIBRATO_EMAIL") != "" && os.Getenv("LIBRATO_TOKEN") != "" && os.Getenv("LIBRATO_SOURCE") != "" {
		i.log.Info("starting librato metrics reporter")

		go librato.Librato(metrics.DefaultRegistry, time.Minute,
			os.Getenv("LIBRATO_EMAIL"), os.Getenv("LIBRATO_TOKEN"), os.Getenv("LIBRATO_SOURCE"),
			[]float64{0.50, 0.75, 0.90, 0.95, 0.99, 0.999, 1.0}, time.Millisecond)
	}
}

func (c *CLI) cleanupImages() error {
	if c.imageCleaner == nil {
		filters := c.c.StringSlice("image-filters")
		if len(filters) == 0 {
			filters = []string{"name eq ^travis-ci.*"}
			c.log.WithField("filters", strings.Join(filters, ",")).Info("default filters set")
		}

		c.imageCleaner = newImageCleaner(c.cs,
			c.log, c.rateLimiter, uint64(c.c.Int("rate-limit-max-calls")), c.c.Duration("rate-limit-duration"), c.c.String("project-id"),
			c.c.String("job-board-url"), filters, c.c.Bool("noop"))
	}

	return c.imageCleaner.Run()
}
<|endoftext|>


<commit_before>package example

import (
	"context"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"sync"
	"time"

	"github.com/hashicorp/nomad/helper"

	log "github.com/hashicorp/go-hclog"
	"github.com/hashicorp/nomad/plugins/base"
	"github.com/hashicorp/nomad/plugins/device"
	"github.com/hashicorp/nomad/plugins/shared/hclspec"
	"github.com/hashicorp/nomad/plugins/shared/structs"
	"github.com/kr/pretty"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
)

const (
	// pluginName is the name of the plugin
	pluginName = "example-fs-device"

	// vendor is the vendor providing the devices
	vendor = "nomad"

	// deviceType is the type of device being returned
	deviceType = "file"

	// deviceName is the name of the devices being exposed
	deviceName = "mock"
)

var (
	// pluginInfo describes the plugin
	pluginInfo = &base.PluginInfoResponse{
		Type:             base.PluginTypeDevice,
		PluginApiVersion: "0.0.1", // XXX This should be an array and should be consts
		PluginVersion:    "0.1.0",
		Name:             pluginName,
	}

	// configSpec is the specification of the plugin's configuration
	configSpec = hclspec.NewObject(map[string]*hclspec.Spec{
		"dir": hclspec.NewDefault(
			hclspec.NewAttr("dir", "string", false),
			hclspec.NewLiteral("\".\""),
		),
		"list_period": hclspec.NewDefault(
			hclspec.NewAttr("list_period", "string", false),
			hclspec.NewLiteral("\"5s\""),
		),
		"unhealthy_perm": hclspec.NewDefault(
			hclspec.NewAttr("unhealthy_perm", "string", false),
			hclspec.NewLiteral("\"-rwxrwxrwx\""),
		),
	})
)

// Config contains configuration information for the plugin.
type Config struct {
	Dir           string `codec:"dir"`
	ListPeriod    string `codec:"list_period"`
	UnhealthyPerm string `codec:"unhealthy_perm"`
}

// FsDevice is an example device plugin. The device plugin exposes files as
// devices and periodically polls the directory for new files. If a file has a
// given file permission, it is considered unhealthy. This device plugin is
// purely for use as an example.
type FsDevice struct {
	logger log.Logger

	// deviceDir is the directory we expose as devices
	deviceDir string

	// unhealthyPerm is the permissions on a file we consider unhealthy
	unhealthyPerm string

	// listPeriod is how often we should list the device directory to detect new
	// devices
	listPeriod time.Duration

	// devices is the set of detected devices and maps whether they are healthy
	devices    map[string]bool
	deviceLock sync.RWMutex
}

// NewExampleDevice returns a new example device plugin.
func NewExampleDevice(log log.Logger) *FsDevice {
	return &FsDevice{
		logger:  log.Named(pluginName),
		devices: make(map[string]bool),
	}
}

// PluginInfo returns information describing the plugin.
func (d *FsDevice) PluginInfo() (*base.PluginInfoResponse, error) {
	return pluginInfo, nil
}

// ConfigSchema returns the plugins configuration schema.
func (d *FsDevice) ConfigSchema() (*hclspec.Spec, error) {
	return configSpec, nil
}

// SetConfig is used to set the configuration of the plugin.
func (d *FsDevice) SetConfig(data []byte, cfg *base.ClientAgentConfig) error {
	var config Config
	if err := base.MsgPackDecode(data, &config); err != nil {
		return err
	}

	// Save the device directory and the unhealthy permissions
	d.deviceDir = config.Dir
	d.unhealthyPerm = config.UnhealthyPerm

	// Convert the poll period
	period, err := time.ParseDuration(config.ListPeriod)
	if err != nil {
		return fmt.Errorf("failed to parse list period %q: %v", config.ListPeriod, err)
	}
	d.listPeriod = period

	d.logger.Debug("test debug")
	d.logger.Info("config set", "config", log.Fmt("% #v", pretty.Formatter(config)))
	return nil
}

// Fingerprint streams detected devices. If device changes are detected or the
// devices health changes, messages will be emitted.
func (d *FsDevice) Fingerprint(ctx context.Context) (<-chan *device.FingerprintResponse, error) {
	if d.deviceDir == "" {
		return nil, status.New(codes.Internal, "device directory not set in config").Err()
	}

	outCh := make(chan *device.FingerprintResponse)
	go d.fingerprint(ctx, outCh)
	return outCh, nil
}

// fingerprint is the long running goroutine that detects hardware
func (d *FsDevice) fingerprint(ctx context.Context, devices chan *device.FingerprintResponse) {
	defer close(devices)

	// Create a timer that will fire immediately for the first detection
	ticker := time.NewTimer(0)

	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			ticker.Reset(d.listPeriod)
		}

		d.logger.Trace("scanning for changes")

		files, err := ioutil.ReadDir(d.deviceDir)
		if err != nil {
			d.logger.Error("failed to list device directory", "error", err)
			devices <- device.NewFingerprintError(err)
			return
		}

		detected := d.diffFiles(files)
		if len(detected) == 0 {
			continue
		}

		devices <- device.NewFingerprint(getDeviceGroup(detected))

	}
}

func (d *FsDevice) diffFiles(files []os.FileInfo) []*device.Device {
	d.deviceLock.Lock()
	defer d.deviceLock.Unlock()

	// Build an unhealthy message
	unhealthyDesc := fmt.Sprintf("Device has bad permissions %q", d.unhealthyPerm)

	var changes bool
	fnames := make(map[string]struct{})
	for _, f := range files {
		name := f.Name()
		fnames[name] = struct{}{}
		if f.IsDir() {
			d.logger.Trace("skipping directory", "directory", name)
			continue
		}

		// Determine the health
		perms := f.Mode().Perm().String()
		healthy := perms != d.unhealthyPerm
		d.logger.Trace("checking health", "file perm", perms, "unhealthy perms", d.unhealthyPerm, "healthy", healthy)

		// See if we alreay have the device
		oldHealth, ok := d.devices[name]
		if ok && oldHealth == healthy {
			continue
		}

		// Health has changed or we have a new object
		changes = true
		d.devices[name] = healthy
	}

	for id := range d.devices {
		if _, ok := fnames[id]; !ok {
			delete(d.devices, id)
			changes = true
		}
	}

	// Nothing to do
	if !changes {
		return nil
	}

	// Build the devices
	detected := make([]*device.Device, 0, len(d.devices))
	for name, healthy := range d.devices {
		var desc string
		if !healthy {
			desc = unhealthyDesc
		}

		detected = append(detected, &device.Device{
			ID:         name,
			Healthy:    healthy,
			HealthDesc: desc,
		})
	}

	return detected
}

// getDeviceGroup is a helper to build the DeviceGroup given a set of devices.
func getDeviceGroup(devices []*device.Device) *device.DeviceGroup {
	return &device.DeviceGroup{
		Vendor:  vendor,
		Type:    deviceType,
		Name:    deviceName,
		Devices: devices,
	}
}

// Reserve returns information on how to mount the given devices.
func (d *FsDevice) Reserve(deviceIDs []string) (*device.ContainerReservation, error) {
	if len(deviceIDs) == 0 {
		return nil, status.New(codes.InvalidArgument, "no device ids given").Err()
	}

	resp := &device.ContainerReservation{}

	for _, id := range deviceIDs {
		// Check if the device is known
		if _, ok := d.devices[id]; !ok {
			return nil, status.Newf(codes.InvalidArgument, "unknown device %q", id).Err()
		}

		// Add a mount
		resp.Devices = append(resp.Devices, &device.DeviceSpec{
			TaskPath:    fmt.Sprintf("/dev/%s", id),
			HostPath:    filepath.Join(d.deviceDir, id),
			CgroupPerms: "rw",
		})
	}

	return resp, nil
}

// Stats streams statistics for the detected devices.
func (d *FsDevice) Stats(ctx context.Context, interval time.Duration) (<-chan *device.StatsResponse, error) {
	outCh := make(chan *device.StatsResponse)
	go d.stats(ctx, outCh, interval)
	return outCh, nil
}

// stats is the long running goroutine that streams device statistics
func (d *FsDevice) stats(ctx context.Context, stats chan *device.StatsResponse, interval time.Duration) {
	defer close(stats)

	// Create a timer that will fire immediately for the first detection
	ticker := time.NewTimer(0)

	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			ticker.Reset(interval)
		}

		deviceStats, err := d.collectStats()
		if err != nil {
			stats <- &device.StatsResponse{
				Error: err,
			}
			return
		}
		if deviceStats == nil {
			continue
		}

		stats <- &device.StatsResponse{
			Groups: []*device.DeviceGroupStats{deviceStats},
		}
	}
}

func (d *FsDevice) collectStats() (*device.DeviceGroupStats, error) {
	d.deviceLock.RLock()
	defer d.deviceLock.RUnlock()
	l := len(d.devices)
	if l == 0 {
		return nil, nil
	}

	now := time.Now()
	group := &device.DeviceGroupStats{
		Vendor:        vendor,
		Type:          deviceType,
		Name:          deviceName,
		InstanceStats: make(map[string]*device.DeviceStats, l),
	}

	for k := range d.devices {
		p := filepath.Join(d.deviceDir, k)
		f, err := os.Stat(p)
		if err != nil {
			return nil, fmt.Errorf("failed to stat %q: %v", p, err)
		}

		s := &device.DeviceStats{
			Summary: &structs.StatValue{
				IntNumeratorVal: helper.Int64ToPtr(f.Size()),
				Unit:            "bytes",
				Desc:            "Filesize in bytes",
			},
			Stats: &structs.StatObject{
				Attributes: map[string]*structs.StatValue{
					"size": {
						IntNumeratorVal: helper.Int64ToPtr(f.Size()),
						Unit:            "bytes",
						Desc:            "Filesize in bytes",
					},
					"modify_time": {
						StringVal: helper.StringToPtr(f.ModTime().String()),
						Desc:      "Last modified",
					},
					"mode": {
						StringVal: helper.StringToPtr(f.Mode().String()),
						Desc:      "File mode",
					},
				},
			},
			Timestamp: now,
		}

		group.InstanceStats[k] = s
	}

	return group, nil
}
<commit_msg>Use absolute path in example device plugin<commit_after>package example

import (
	"context"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"sync"
	"time"

	"github.com/hashicorp/nomad/helper"

	log "github.com/hashicorp/go-hclog"
	"github.com/hashicorp/nomad/plugins/base"
	"github.com/hashicorp/nomad/plugins/device"
	"github.com/hashicorp/nomad/plugins/shared/hclspec"
	"github.com/hashicorp/nomad/plugins/shared/structs"
	"github.com/kr/pretty"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
)

const (
	// pluginName is the name of the plugin
	pluginName = "example-fs-device"

	// vendor is the vendor providing the devices
	vendor = "nomad"

	// deviceType is the type of device being returned
	deviceType = "file"

	// deviceName is the name of the devices being exposed
	deviceName = "mock"
)

var (
	// pluginInfo describes the plugin
	pluginInfo = &base.PluginInfoResponse{
		Type:             base.PluginTypeDevice,
		PluginApiVersion: "0.0.1", // XXX This should be an array and should be consts
		PluginVersion:    "0.1.0",
		Name:             pluginName,
	}

	// configSpec is the specification of the plugin's configuration
	configSpec = hclspec.NewObject(map[string]*hclspec.Spec{
		"dir": hclspec.NewDefault(
			hclspec.NewAttr("dir", "string", false),
			hclspec.NewLiteral("\".\""),
		),
		"list_period": hclspec.NewDefault(
			hclspec.NewAttr("list_period", "string", false),
			hclspec.NewLiteral("\"5s\""),
		),
		"unhealthy_perm": hclspec.NewDefault(
			hclspec.NewAttr("unhealthy_perm", "string", false),
			hclspec.NewLiteral("\"-rwxrwxrwx\""),
		),
	})
)

// Config contains configuration information for the plugin.
type Config struct {
	Dir           string `codec:"dir"`
	ListPeriod    string `codec:"list_period"`
	UnhealthyPerm string `codec:"unhealthy_perm"`
}

// FsDevice is an example device plugin. The device plugin exposes files as
// devices and periodically polls the directory for new files. If a file has a
// given file permission, it is considered unhealthy. This device plugin is
// purely for use as an example.
type FsDevice struct {
	logger log.Logger

	// deviceDir is the directory we expose as devices
	deviceDir string

	// unhealthyPerm is the permissions on a file we consider unhealthy
	unhealthyPerm string

	// listPeriod is how often we should list the device directory to detect new
	// devices
	listPeriod time.Duration

	// devices is the set of detected devices and maps whether they are healthy
	devices    map[string]bool
	deviceLock sync.RWMutex
}

// NewExampleDevice returns a new example device plugin.
func NewExampleDevice(log log.Logger) *FsDevice {
	return &FsDevice{
		logger:  log.Named(pluginName),
		devices: make(map[string]bool),
	}
}

// PluginInfo returns information describing the plugin.
func (d *FsDevice) PluginInfo() (*base.PluginInfoResponse, error) {
	return pluginInfo, nil
}

// ConfigSchema returns the plugins configuration schema.
func (d *FsDevice) ConfigSchema() (*hclspec.Spec, error) {
	return configSpec, nil
}

// SetConfig is used to set the configuration of the plugin.
func (d *FsDevice) SetConfig(data []byte, cfg *base.ClientAgentConfig) error {
	var config Config
	if err := base.MsgPackDecode(data, &config); err != nil {
		return err
	}

	// Save the device directory and the unhealthy permissions
	d.deviceDir = config.Dir
	d.unhealthyPerm = config.UnhealthyPerm

	// Convert the poll period
	period, err := time.ParseDuration(config.ListPeriod)
	if err != nil {
		return fmt.Errorf("failed to parse list period %q: %v", config.ListPeriod, err)
	}
	d.listPeriod = period

	d.logger.Debug("test debug")
	d.logger.Info("config set", "config", log.Fmt("% #v", pretty.Formatter(config)))
	return nil
}

// Fingerprint streams detected devices. If device changes are detected or the
// devices health changes, messages will be emitted.
func (d *FsDevice) Fingerprint(ctx context.Context) (<-chan *device.FingerprintResponse, error) {
	if d.deviceDir == "" {
		return nil, status.New(codes.Internal, "device directory not set in config").Err()
	}

	outCh := make(chan *device.FingerprintResponse)
	go d.fingerprint(ctx, outCh)
	return outCh, nil
}

// fingerprint is the long running goroutine that detects hardware
func (d *FsDevice) fingerprint(ctx context.Context, devices chan *device.FingerprintResponse) {
	defer close(devices)

	// Create a timer that will fire immediately for the first detection
	ticker := time.NewTimer(0)

	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			ticker.Reset(d.listPeriod)
		}

		d.logger.Trace("scanning for changes")

		files, err := ioutil.ReadDir(d.deviceDir)
		if err != nil {
			d.logger.Error("failed to list device directory", "error", err)
			devices <- device.NewFingerprintError(err)
			return
		}

		detected := d.diffFiles(files)
		if len(detected) == 0 {
			continue
		}

		devices <- device.NewFingerprint(getDeviceGroup(detected))

	}
}

func (d *FsDevice) diffFiles(files []os.FileInfo) []*device.Device {
	d.deviceLock.Lock()
	defer d.deviceLock.Unlock()

	// Build an unhealthy message
	unhealthyDesc := fmt.Sprintf("Device has bad permissions %q", d.unhealthyPerm)

	var changes bool
	fnames := make(map[string]struct{})
	for _, f := range files {
		name := f.Name()
		fnames[name] = struct{}{}
		if f.IsDir() {
			d.logger.Trace("skipping directory", "directory", name)
			continue
		}

		// Determine the health
		perms := f.Mode().Perm().String()
		healthy := perms != d.unhealthyPerm
		d.logger.Trace("checking health", "file perm", perms, "unhealthy perms", d.unhealthyPerm, "healthy", healthy)

		// See if we alreay have the device
		oldHealth, ok := d.devices[name]
		if ok && oldHealth == healthy {
			continue
		}

		// Health has changed or we have a new object
		changes = true
		d.devices[name] = healthy
	}

	for id := range d.devices {
		if _, ok := fnames[id]; !ok {
			delete(d.devices, id)
			changes = true
		}
	}

	// Nothing to do
	if !changes {
		return nil
	}

	// Build the devices
	detected := make([]*device.Device, 0, len(d.devices))
	for name, healthy := range d.devices {
		var desc string
		if !healthy {
			desc = unhealthyDesc
		}

		detected = append(detected, &device.Device{
			ID:         name,
			Healthy:    healthy,
			HealthDesc: desc,
		})
	}

	return detected
}

// getDeviceGroup is a helper to build the DeviceGroup given a set of devices.
func getDeviceGroup(devices []*device.Device) *device.DeviceGroup {
	return &device.DeviceGroup{
		Vendor:  vendor,
		Type:    deviceType,
		Name:    deviceName,
		Devices: devices,
	}
}

// Reserve returns information on how to mount the given devices.
func (d *FsDevice) Reserve(deviceIDs []string) (*device.ContainerReservation, error) {
	if len(deviceIDs) == 0 {
		return nil, status.New(codes.InvalidArgument, "no device ids given").Err()
	}

	deviceDir, err := filepath.Abs(d.deviceDir)
	if err != nil {
		return nil, status.Newf(codes.Internal, "failed to load device dir abs path").Err()
	}

	resp := &device.ContainerReservation{}

	for _, id := range deviceIDs {
		// Check if the device is known
		if _, ok := d.devices[id]; !ok {
			return nil, status.Newf(codes.InvalidArgument, "unknown device %q", id).Err()
		}

		// Add a mount
		resp.Mounts = append(resp.Mounts, &device.Mount{
			TaskPath: fmt.Sprintf("/tmp/task-mounts/%s", id),
			HostPath: filepath.Join(deviceDir, id),
			ReadOnly: false,
		})
	}

	return resp, nil
}

// Stats streams statistics for the detected devices.
func (d *FsDevice) Stats(ctx context.Context, interval time.Duration) (<-chan *device.StatsResponse, error) {
	outCh := make(chan *device.StatsResponse)
	go d.stats(ctx, outCh, interval)
	return outCh, nil
}

// stats is the long running goroutine that streams device statistics
func (d *FsDevice) stats(ctx context.Context, stats chan *device.StatsResponse, interval time.Duration) {
	defer close(stats)

	// Create a timer that will fire immediately for the first detection
	ticker := time.NewTimer(0)

	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			ticker.Reset(interval)
		}

		deviceStats, err := d.collectStats()
		if err != nil {
			stats <- &device.StatsResponse{
				Error: err,
			}
			return
		}
		if deviceStats == nil {
			continue
		}

		stats <- &device.StatsResponse{
			Groups: []*device.DeviceGroupStats{deviceStats},
		}
	}
}

func (d *FsDevice) collectStats() (*device.DeviceGroupStats, error) {
	d.deviceLock.RLock()
	defer d.deviceLock.RUnlock()
	l := len(d.devices)
	if l == 0 {
		return nil, nil
	}

	now := time.Now()
	group := &device.DeviceGroupStats{
		Vendor:        vendor,
		Type:          deviceType,
		Name:          deviceName,
		InstanceStats: make(map[string]*device.DeviceStats, l),
	}

	for k := range d.devices {
		p := filepath.Join(d.deviceDir, k)
		f, err := os.Stat(p)
		if err != nil {
			return nil, fmt.Errorf("failed to stat %q: %v", p, err)
		}

		s := &device.DeviceStats{
			Summary: &structs.StatValue{
				IntNumeratorVal: helper.Int64ToPtr(f.Size()),
				Unit:            "bytes",
				Desc:            "Filesize in bytes",
			},
			Stats: &structs.StatObject{
				Attributes: map[string]*structs.StatValue{
					"size": {
						IntNumeratorVal: helper.Int64ToPtr(f.Size()),
						Unit:            "bytes",
						Desc:            "Filesize in bytes",
					},
					"modify_time": {
						StringVal: helper.StringToPtr(f.ModTime().String()),
						Desc:      "Last modified",
					},
					"mode": {
						StringVal: helper.StringToPtr(f.Mode().String()),
						Desc:      "File mode",
					},
				},
			},
			Timestamp: now,
		}

		group.InstanceStats[k] = s
	}

	return group, nil
}
<|endoftext|>


<commit_before>package lmdbcli

import (
	"flag"
	"fmt"
	"log"
	"os"
	"path"

	"github.com/szferi/gomdb"
)

var (
	pathFlag = flag.String("db", "", "Relative path to lmdb file")
	sizeFlag = flag.Float64("size", 2, "factor to allocate for growth or shrinkage")
	roFlag   = flag.Bool("ro", false, "open the database in read-only mode")
	minArgs  = map[string]int{"scan": 0, "stat": 0, "expand": 0, "exists": 1, "get": 1, "del": 1, "put": 2}
)

// Run golmdb using the directory containing the data as dbPath

func Run() {
	flag.Parse()

	if len(*pathFlag) == 0 && len(flag.Args()) == 1 {
		pathFlag = &flag.Args()[0]
	}
	if len(*pathFlag) == 0 {
		log.Fatal("-db must be specified")
	}

	size := uint64(1024 * 1024 * 32)
	if stat, err := os.Stat(path.Join(*pathFlag, "data.mdb")); err != nil {
		if os.IsNotExist(err) == false {
			log.Fatal("failed to stat data.mdb file: ", err)
		}
	} else {
		size = uint64(float64(stat.Size()) * *sizeFlag)
	}

	context := NewContext(*pathFlag, size, os.Stdout)
	defer context.Close()
	if err := context.SwitchDB(nil); err != nil {
		log.Fatal("could not select default database: ", err)
	}
	runShell(context)
}

func runShell(context *Context) {
	var err error
	for {
		fmt.Print(context.prompt)
		var fn, key, val string
		fmt.Scanln(&fn, &key, &val)

		if _, ok := minArgs[fn]; !ok {
			context.Write([]byte("error: invalid command"))
		} else if !checkNumArgs(fn, key, val) {
			context.Write([]byte("error: not enough arguments"))
		} else if fn == "get" {
			err = get(context, key)
		} else if fn == "exists" {
			err = exists(context, key)
		} else if fn == "del" {
			err = del(context, key)
		} else if fn == "put" {
			err = put(context, key, val)
		} else if fn == "scan" {
			err = scan(context)
		} else {
			return
		}
		if err != nil {
			context.Write([]byte(err.Error()))
		}
	}
}

func get(context *Context, key string) error {
	return context.WithinRead(func(txn *mdb.Txn) error {
		data, err := txn.Get(context.dbi, []byte(key))
		if err != nil {
			return err
		}
		context.Write(data)
		return nil
	})
}

func exists(context *Context, key string) error {
	return context.WithinRead(func(txn *mdb.Txn) error {
		_, err := txn.Get(context.dbi, []byte(key))
		if err != nil {
			context.Write([]byte("false"))
		} else {
			context.Write([]byte("true"))
		}
		return nil
	})
}

func del(context *Context, key string) error {
	return context.WithinWrite(func(txn *mdb.Txn) error {
		return txn.Del(context.dbi, []byte(key), nil)
	})
}

func put(context *Context, key, val string) error {
	return context.WithinWrite(func(txn *mdb.Txn) error {
		return txn.Put(context.dbi, []byte(key), []byte(val), 0)
	})
}

func scan(context *Context) error {
	return context.WithinRead(func(txn *mdb.Txn) error {
		cursor, err := txn.CursorOpen(context.dbi)
		if err != nil {
			return err
		}
		defer cursor.Close()
		for {
			key, val, err := cursor.Get(nil, nil, mdb.NEXT)
			if err == mdb.NotFound {
				return nil
			}
			if err != nil {
				return err
			}
			context.Write(key)
			context.Write(val)
		}
	})
}

func checkNumArgs(fn, key, val string) bool {
	if fn == "" {
		return false
	}
	n := 0
	if key != "" {
		n++
	}
	if val != "" {
		n++
	}
	if expected, ok := minArgs[fn]; ok {
		return n >= expected
	}
	return false
}
<commit_msg>better input and command parsing<commit_after>package lmdbcli

import (
	"bufio"
	"errors"
	"flag"
	"fmt"
	"log"
	"os"
	"path"
	"unicode"

	"github.com/szferi/gomdb"
)

var (
	pathFlag = flag.String("db", "", "Relative path to lmdb file")
	sizeFlag = flag.Float64("size", 2, "factor to allocate for growth or shrinkage")
	roFlag   = flag.Bool("ro", false, "open the database in read-only mode")
	minArgs  = map[string]int{"scan": 0, "stat": 0, "expand": 0, "exists": 1, "get": 1, "del": 1, "put": 2}
)

const (
	STATE_NONE int = iota
	STATE_WORD
	STATE_QUOTE
	STATE_ESCAPED
)

type Command struct {
	fn   string
	key  []byte
	val  []byte
	args [][]byte
}

// Run golmdb using the directory containing the data as dbPath

func Run() {
	flag.Parse()

	if len(*pathFlag) == 0 && len(flag.Args()) == 1 {
		pathFlag = &flag.Args()[0]
	}
	if len(*pathFlag) == 0 {
		log.Fatal("-db must be specified")
	}

	size := uint64(1024 * 1024 * 32)
	if stat, err := os.Stat(path.Join(*pathFlag, "data.mdb")); err != nil {
		if os.IsNotExist(err) == false {
			log.Fatal("failed to stat data.mdb file: ", err)
		}
	} else {
		size = uint64(float64(stat.Size()) * *sizeFlag)
	}

	context := NewContext(*pathFlag, size, os.Stdout)
	defer context.Close()
	if err := context.SwitchDB(nil); err != nil {
		log.Fatal("could not select default database: ", err)
	}
	runShell(context)
}

func runShell(context *Context) {
	var err error
	reader := bufio.NewReader(os.Stdin)
	for {
		fmt.Print(context.prompt)
		input, _ := reader.ReadSlice('\n')

		args := parseInput(input)
		if cmd, err1 := getCommand(args); err1 != nil {
			context.Write([]byte(err1.Error()))
		} else if cmd.fn == "get" {
			err = get(context, cmd.key)
		} else if cmd.fn == "exists" {
			err = exists(context, cmd.key)
		} else if cmd.fn == "del" {
			err = del(context, cmd.key)
		} else if cmd.fn == "put" {
			err = put(context, cmd.key, cmd.val)
		} else if cmd.fn == "scan" {
			err = scan(context)
		}
		if err != nil {
			context.Write([]byte(err.Error()))
		}
	}
}

func get(context *Context, key []byte) error {
	return context.WithinRead(func(txn *mdb.Txn) error {
		data, err := txn.Get(context.dbi, key)
		if err != nil {
			return err
		}
		context.Write(data)
		return nil
	})
}

func exists(context *Context, key []byte) error {
	return context.WithinRead(func(txn *mdb.Txn) error {
		_, err := txn.Get(context.dbi, key)
		if err != nil {
			context.Write([]byte("false"))
		} else {
			context.Write([]byte("true"))
		}
		return nil
	})
}

func del(context *Context, key []byte) error {
	return context.WithinWrite(func(txn *mdb.Txn) error {
		return txn.Del(context.dbi, key, nil)
	})
}

func put(context *Context, key, val []byte) error {
	return context.WithinWrite(func(txn *mdb.Txn) error {
		return txn.Put(context.dbi, key, val, 0)
	})
}

func scan(context *Context) error {
	return context.WithinRead(func(txn *mdb.Txn) error {
		cursor, err := txn.CursorOpen(context.dbi)
		if err != nil {
			return err
		}
		defer cursor.Close()
		for {
			key, val, err := cursor.Get(nil, nil, mdb.NEXT)
			if err == mdb.NotFound {
				return nil
			}
			if err != nil {
				return err
			}
			context.Write(key)
			context.Write(val)
		}
	})
}

// handle both space delimiters and arguments in quotations
// arguments are defined as contained by spaces ' arg ' or quotations '"arg"'
// forward slash escapes for nested quotations
func parseInput(in []byte) [][]byte {
	var results [][]byte
	var arg []byte
	state := STATE_NONE
	for _, b := range in {
		switch state {
		case STATE_NONE:
			if isQuote(b) {
				state = STATE_QUOTE
			} else if !isWhiteSpace(b) {
				arg = append(arg, b)
				state = STATE_WORD
			}
		case STATE_ESCAPED:
			arg = append(arg, b)
			state = STATE_QUOTE
		case STATE_WORD:
			if isWhiteSpace(b) {
				results = append(results, arg)
				arg = make([]byte, 0)
				state = STATE_NONE
			} else {
				arg = append(arg, b)
			}
		case STATE_QUOTE:
			if b == '\\' {
				state = STATE_ESCAPED
			} else if isQuote(b) {
				results = append(results, arg)
				arg = make([]byte, 0)
				state = STATE_NONE
			} else {
				arg = append(arg, b)
			}
		}
	}
	return results
}

func isWhiteSpace(b byte) bool {
	return unicode.IsSpace(rune(b))
}

func isQuote(b byte) bool {
	return b == '"' || b == '\''
}

func getCommand(args [][]byte) (Command, error) {
	numArgs := 0
	var cmd Command
	if len(args) == 0 {
		return cmd, errors.New("empty command")
	}
	fn := string(args[0])
	if _, ok := minArgs[fn]; !ok {
		return cmd, errors.New("invalid command")
	}
	var key, value []byte
	if len(args) >= 2 && len(args[1]) > 0 {
		key = args[1]
		numArgs++
	}
	if len(args) >= 3 && len(args[2]) > 0 {
		value = args[2]
		numArgs++
	}
	if numArgs < minArgs[fn] {
		return cmd, errors.New("not enough arguments")
	}
	return Command{
		fn:  fn,
		key: key,
		val: value,
	}, nil
}
<|endoftext|>


<commit_before>// Copyright 2015 John Kelley. All rights reserved.
// This source code is licensed under the Simplified BSD License

// Woozle is a simple DNS recursor with the ability to filter out
// certain queries. The author finds this useful for forcing Youtube
// over IPv4 instead of his Hurricane Electric IPv6 tunnel

package main

import (
	"fmt"
	"github.com/miekg/dns"
	"os"
	"os/signal"
	"syscall"
	"strings"
	"time"
)

// ================================================
// Configurables - TODO: move to CLI arg processing
// ================================================
const upstreamDNS = "10.10.10.1:53"
var filterDomainAAAA = []string{ "youtube.com.", "googlevideo.com." }

// ================================================
// Globals for tracking stats and caching
// ================================================
type DNSQuery struct {
	domain string
	queryType string
	filtered bool
}

type DomainStats struct {
	domain string
	frequency int
	filtered int
	queries map[string]int
}

var statPipe chan DNSQuery
var stats = map[string]*DomainStats{}
var totalQueries = 0
var timeStarted time.Time

// ================================================
// Helper functions
// ================================================

func serve(net string) {
	server := &dns.Server{Addr: ":53", Net: net, TsigSecret: nil}
	err := server.ListenAndServe()
	if err != nil {
		fmt.Printf("Failed to setup the "+net+" server: %s\n", err.Error())
		os.Exit(1)
	}
}

func getRootFromDomain(domain string) string {
	var root string
	components := strings.Split(domain, ".")
	idx := len(components)
	if idx > 2 {
		root = components[idx-3] + "." + components[idx-2]
	}
	return root
}

func dispStats() {
	fmt.Printf("\nQuery Statistics:\n")

	for d, s := range stats {
		fmt.Printf("%25s: %3d queries", d, s.frequency)
		if s.filtered > 0 {
			fmt.Printf(", %3d dropped", s.filtered)
		}
		fmt.Println()
	}
}

// ================================================
// Meat and potatoes
// ================================================

func handleRecurse(w dns.ResponseWriter, m *dns.Msg) {
	// collect stats
	statPipe <- DNSQuery{m.Question[0].Name, dns.TypeToString[m.Question[0].Qtype], false}

	// pass the query on to the upstream DNS server
	c := new(dns.Client)
	r, _, e := c.Exchange(m, upstreamDNS)
	if e != nil {
		fmt.Printf("%s Client query failed: %s\n", time.Now().Format("01/02 15:04"), e.Error())
	} else {
		w.WriteMsg(r)
	}
}

func filterAAAA(w dns.ResponseWriter, r *dns.Msg) {
	if r.Question[0].Qtype == dns.TypeAAAA {
		// collect stats
		statPipe <- DNSQuery{r.Question[0].Name, "AAAA", true}

		// send a blank reply
		m := new(dns.Msg)
		m.SetReply(r)
		w.WriteMsg(m)
	} else {
		handleRecurse(w, r)
	}
}

func handleStats(queryChan <-chan DNSQuery) {
	var domain string
	for query := range queryChan {
		totalQueries++
		domain = getRootFromDomain(query.domain)
		if nil == stats[domain] {
			stats[domain] = &DomainStats{domain:domain}
			stats[domain].queries = make(map[string]int)
		}
		stats[domain].frequency++
		if query.filtered {
			stats[domain].filtered++
		}
		stats[domain].queries[query.queryType] += 1
		//fmt.Printf("%s type %s (%d, %d)\n", query.domain, query.queryType, stats[query.domain].frequency, stats[query.domain].queries[query.queryType])
	}
	fmt.Printf("stats engine stopping\n");
}

func main() {
	// init update
	timeStarted = time.Now()
	lastSigInt := timeStarted.Add(time.Second * -31)

	// setup and kickoff stats collection
	statPipe = make(chan DNSQuery, 10)
	go handleStats(statPipe)

	// handler for filtering ipv6 AAAA records
	for _, domain := range filterDomainAAAA {
		dns.HandleFunc(domain, filterAAAA)
	}

	// default handler
	dns.HandleFunc(".", handleRecurse)

	// start server(s)
	//go serve("tcp")
	go serve("udp")

	// handle signals
	sig := make(chan os.Signal)
	signal.Notify(sig, syscall.SIGINT, syscall.SIGTERM)
forever:
	for s := range sig {
		fmt.Printf("\nUptime %s, %d queries performed", time.Since(timeStarted).String(), totalQueries)
		if time.Since(lastSigInt).Seconds() < 30 {
			fmt.Printf("\nSignal (%d) received, stopping\n", s)
			break forever
		} else {
			fmt.Printf(", send SIGINT again within 30s to quit\n")
			lastSigInt = time.Now();
			dispStats()
		}
	}
}
<commit_msg>Retry failed queries once after 10ms. Print queries that fail<commit_after>// Copyright 2015 John Kelley. All rights reserved.
// This source code is licensed under the Simplified BSD License

// Woozle is a simple DNS recursor with the ability to filter out
// certain queries. The author finds this useful for forcing Youtube
// over IPv4 instead of his Hurricane Electric IPv6 tunnel

package main

import (
	"fmt"
	"github.com/miekg/dns"
	"os"
	"os/signal"
	"syscall"
	"strings"
	"time"
)

// ================================================
// Configurables - TODO: move to CLI arg processing
// ================================================
const upstreamDNS = "10.10.10.1:53"
var filterDomainAAAA = []string{ "youtube.com.", "googlevideo.com." }

// ================================================
// Globals for tracking stats and caching
// ================================================
type DNSQuery struct {
	domain string
	queryType string
	filtered bool
}

type DomainStats struct {
	domain string
	frequency int
	filtered int
	queries map[string]int
}

var statPipe chan DNSQuery
var stats = map[string]*DomainStats{}
var totalQueries = 0
var timeStarted time.Time

// ================================================
// Helper functions
// ================================================

func serve(net string) {
	server := &dns.Server{Addr: ":53", Net: net, TsigSecret: nil}
	err := server.ListenAndServe()
	if err != nil {
		fmt.Printf("Failed to setup the "+net+" server: %s\n", err.Error())
		os.Exit(1)
	}
}

func getRootFromDomain(domain string) string {
	var root string
	components := strings.Split(domain, ".")
	idx := len(components)
	if idx > 2 {
		root = components[idx-3] + "." + components[idx-2]
	}
	return root
}

func dispStats() {
	fmt.Printf("\nQuery Statistics:\n")

	for d, s := range stats {
		fmt.Printf("%25s: %3d queries", d, s.frequency)
		if s.filtered > 0 {
			fmt.Printf(", %3d dropped", s.filtered)
		}
		fmt.Println()
	}
}

// ================================================
// Meat and potatoes
// ================================================

func handleRecurse(w dns.ResponseWriter, m *dns.Msg) {
	// collect stats
	statPipe <- DNSQuery{m.Question[0].Name, dns.TypeToString[m.Question[0].Qtype], false}

	// pass the query on to the upstream DNS server
	c := new(dns.Client)
	r, _, e := c.Exchange(m, upstreamDNS)
	if e != nil {
		fmt.Printf("%s retrying query for '%s'\n", time.Now().Format("01/02 15:04"), m.Question[0].Name);
		time.Sleep(time.Millisecond * 10);
		r, _, e = c.Exchange(m, upstreamDNS)
		if e != nil {
			fmt.Printf("%s Client query failed for '%s': %s\n", time.Now().Format("01/02 15:04"), m.Question[0].Name, e.Error())
			return
		}
	}
	w.WriteMsg(r)
}

func filterAAAA(w dns.ResponseWriter, r *dns.Msg) {
	if r.Question[0].Qtype == dns.TypeAAAA {
		// collect stats
		statPipe <- DNSQuery{r.Question[0].Name, "AAAA", true}

		// send a blank reply
		m := new(dns.Msg)
		m.SetReply(r)
		w.WriteMsg(m)
	} else {
		handleRecurse(w, r)
	}
}

func handleStats(queryChan <-chan DNSQuery) {
	var domain string
	for query := range queryChan {
		totalQueries++
		domain = getRootFromDomain(query.domain)
		if nil == stats[domain] {
			stats[domain] = &DomainStats{domain:domain}
			stats[domain].queries = make(map[string]int)
		}
		stats[domain].frequency++
		if query.filtered {
			stats[domain].filtered++
		}
		stats[domain].queries[query.queryType] += 1
		//fmt.Printf("%s type %s (%d, %d)\n", query.domain, query.queryType, stats[query.domain].frequency, stats[query.domain].queries[query.queryType])
	}
	fmt.Printf("stats engine stopping\n");
}

func main() {
	// init update
	timeStarted = time.Now()
	lastSigInt := timeStarted.Add(time.Second * -31)

	// setup and kickoff stats collection
	statPipe = make(chan DNSQuery, 10)
	go handleStats(statPipe)

	// handler for filtering ipv6 AAAA records
	for _, domain := range filterDomainAAAA {
		dns.HandleFunc(domain, filterAAAA)
	}

	// default handler
	dns.HandleFunc(".", handleRecurse)

	// start server(s)
	//go serve("tcp")
	go serve("udp")

	// handle signals
	sig := make(chan os.Signal)
	signal.Notify(sig, syscall.SIGINT, syscall.SIGTERM)
forever:
	for s := range sig {
		fmt.Printf("\nUptime %s, %d queries performed", time.Since(timeStarted).String(), totalQueries)
		if time.Since(lastSigInt).Seconds() < 30 {
			fmt.Printf("\nSignal (%d) received, stopping\n", s)
			break forever
		} else {
			fmt.Printf(", send SIGINT again within 30s to quit\n")
			lastSigInt = time.Now();
			dispStats()
		}
	}
}
<|endoftext|>


<commit_before>package cmd

import (
	"github.com/gobs/readline"

	"fmt"
	"strings"
)

type Command struct {
	Name string
	Help string
	Call func(string) bool
}

type Cmd struct {
	Prompt string

	PreLoop   func()
	PostLoop  func()
	PreCmd    func(string)
	PostCmd   func(string, bool) bool
	EmptyLine func()
	Default   func(string)

	Commands map[string]Command
}

func (cmd *Cmd) Init() {
	cmd.Commands = make(map[string]Command)

	if cmd.PreLoop == nil {
		cmd.PreLoop = func() {}
	}
	if cmd.PostLoop == nil {
		cmd.PostLoop = func() {}
	}
	if cmd.PreCmd == nil {
		cmd.PreCmd = func(string) {}
	}
	if cmd.PostCmd == nil {
		cmd.PostCmd = func(line string, stop bool) bool { return stop }
	}
	if cmd.EmptyLine == nil {
		cmd.EmptyLine = func() {}
	}
	if cmd.Default == nil {
		cmd.Default = func(line string) { fmt.Printf("invalid command: %v\n", line) }
	}

	cmd.Add(Command{"help", `list available commands`, cmd.Help})
}

func (cmd *Cmd) Add(command Command) {
	cmd.Commands[command.Name] = command
}

func (cmd *Cmd) Help(line string) (stop bool) {
	if len(line) == 0 {
		fmt.Println("Available commands:")

		for k, _ := range cmd.Commands {
			fmt.Println("    ", k)
		}
	} else {
		c, ok := cmd.Commands[line]
		if ok {
			if len(c.Help) > 0 {
				fmt.Println(c.Help)
			} else {
				fmt.Println("No help for ", line)
			}
		} else {
			fmt.Println("unknown command")
		}
	}
	return
}

func (cmd *Cmd) OneCmd(line string) (stop bool) {

	parts := strings.SplitN(line, " ", 2)
	cname := parts[0]

	command, ok := cmd.Commands[cname]

	if ok {
		var params string

		if len(parts) > 1 {
			params = strings.TrimSpace(parts[1])
		}

		stop = command.Call(params)
	} else {
		cmd.Default(line)
	}

	return
}

func (cmd *Cmd) CmdLoop() {
	if len(cmd.Prompt) == 0 {
		cmd.Prompt = "> "
	}

	cmd.PreLoop()

	// loop until ReadLine returns nil (signalling EOF)
	for {
		result := readline.ReadLine(&cmd.Prompt)
		if result == nil {
			break
		}

		line := strings.TrimSpace(*result)
		if line == "" {
			cmd.EmptyLine()
			continue
		}

		readline.AddHistory(*result) // allow user to recall this line

		cmd.PreCmd(line)

		stop := cmd.OneCmd(line)
		stop = cmd.PostCmd(line, stop)

		if stop {
			break
		}
	}

	cmd.PostLoop()
}
<commit_msg>Adding command completion<commit_after>package cmd

import (
	"github.com/gobs/readline"

	"fmt"
	"strings"
)

type Command struct {
	Name string
	Help string
	Call func(string) bool
}

type Completer struct {
	Words   []string
	Matches []string
}

func (c *Completer) Complete(prefix string, index int) string {
	if index == 0 {
		c.Matches = c.Matches[:0]
		no_prefix := len(prefix) == 0

		for _, w := range c.Words {
			if no_prefix || strings.HasPrefix(w, prefix) {
				c.Matches = append(c.Matches, w)
			}
		}
	}

	if index < len(c.Matches) {
		return c.Matches[index]
	} else {
		return ""
	}
}

func NewCompleter(words []string) (c *Completer) {
	c = new(Completer)
	c.Words = words
	c.Matches = make([]string, 0, len(c.Words))
	return
}

type Cmd struct {
	Prompt string

	PreLoop   func()
	PostLoop  func()
	PreCmd    func(string)
	PostCmd   func(string, bool) bool
	EmptyLine func()
	Default   func(string)

	Commands map[string]Command
}

func (cmd *Cmd) Init() {
	cmd.Commands = make(map[string]Command)

	if cmd.PreLoop == nil {
		cmd.PreLoop = func() {}
	}
	if cmd.PostLoop == nil {
		cmd.PostLoop = func() {}
	}
	if cmd.PreCmd == nil {
		cmd.PreCmd = func(string) {}
	}
	if cmd.PostCmd == nil {
		cmd.PostCmd = func(line string, stop bool) bool { return stop }
	}
	if cmd.EmptyLine == nil {
		cmd.EmptyLine = func() {}
	}
	if cmd.Default == nil {
		cmd.Default = func(line string) { fmt.Printf("invalid command: %v\n", line) }
	}

	cmd.Add(Command{"help", `list available commands`, cmd.Help})
}

func (cmd *Cmd) addCommandCompleter() {
	names := make([]string, 0, len(cmd.Commands))

	for n, _ := range cmd.Commands {
		names = append(names, n)
	}

	completer := NewCompleter(names)
	readline.SetCompletionEntryFunction(completer.Complete)
}

func (cmd *Cmd) Add(command Command) {
	cmd.Commands[command.Name] = command
}

func (cmd *Cmd) Help(line string) (stop bool) {
	if len(line) == 0 {
		fmt.Println("Available commands:")

		for k, _ := range cmd.Commands {
			fmt.Println("    ", k)
		}
	} else {
		c, ok := cmd.Commands[line]
		if ok {
			if len(c.Help) > 0 {
				fmt.Println(c.Help)
			} else {
				fmt.Println("No help for ", line)
			}
		} else {
			fmt.Println("unknown command")
		}
	}
	return
}

func (cmd *Cmd) OneCmd(line string) (stop bool) {

	parts := strings.SplitN(line, " ", 2)
	cname := parts[0]

	command, ok := cmd.Commands[cname]

	if ok {
		var params string

		if len(parts) > 1 {
			params = strings.TrimSpace(parts[1])
		}

		stop = command.Call(params)
	} else {
		cmd.Default(line)
	}

	return
}

func (cmd *Cmd) CmdLoop() {
	if len(cmd.Prompt) == 0 {
		cmd.Prompt = "> "
	}

	cmd.addCommandCompleter()

	cmd.PreLoop()

	// loop until ReadLine returns nil (signalling EOF)
	for {
		result := readline.ReadLine(&cmd.Prompt)
		if result == nil {
			break
		}

		line := strings.TrimSpace(*result)
		if line == "" {
			cmd.EmptyLine()
			continue
		}

		readline.AddHistory(*result) // allow user to recall this line

		cmd.PreCmd(line)

		stop := cmd.OneCmd(line)
		stop = cmd.PostCmd(line, stop)

		if stop {
			break
		}
	}

	cmd.PostLoop()
}
<|endoftext|>


<commit_before>/*
 * Copyright (c) 2013 Conformal Systems LLC <info@conformal.com>
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

package main

import (
	"encoding/json"
	"errors"
	"fmt"
	"github.com/conformal/btcjson"
	"github.com/conformal/btcutil"
	"github.com/conformal/btcwallet/tx"
	"github.com/conformal/btcwallet/wallet"
	"github.com/conformal/btcwire"
	"github.com/conformal/seelog"
	"os"
	"path/filepath"
	"sync"
	"time"
)

var (
	ErrNoWallet = errors.New("Wallet file does not exist.")
)

var (
	log     seelog.LoggerInterface = seelog.Default
	cfg     *config
	wallets = struct {
		sync.RWMutex
		m map[string]*BtcWallet
	}{
		m: make(map[string]*BtcWallet),
	}
)

func main() {
	tcfg, _, err := loadConfig()
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
	cfg = tcfg

	// Open wallet
	w, err := OpenWallet(cfg, "")
	if err != nil {
		log.Info(err.Error())
	} else {
		w.Track()
	}

	// Start HTTP server to listen and send messages to frontend and btcd
	// backend.  Try reconnection if connection failed.
	for {
		if err := ListenAndServe(); err == ConnRefused {
			// wait and try again.
			log.Info("Unable to connect to btcd. Retrying in 5 seconds.")
			time.Sleep(5 * time.Second)
		} else if err != nil {
			log.Info(err.Error())
			break
		}
	}
}

type BtcWallet struct {
	*wallet.Wallet
	mtx       sync.RWMutex
	dirty     bool
	UtxoStore struct {
		sync.RWMutex
		dirty bool
		s     tx.UtxoStore
	}
	TxStore struct {
		sync.RWMutex
		dirty bool
		s     tx.TxStore
	}
}

// walletdir returns the directory path which holds the wallet, utxo,
// and tx files.
func walletdir(cfg *config, account string) string {
	var wname string
	if account == "" {
		wname = "btcwallet"
	} else {
		wname = fmt.Sprintf("btcwallet-%s", account)
	}

	return filepath.Join(cfg.DataDir, wname)
}

func OpenWallet(cfg *config, account string) (*BtcWallet, error) {
	wdir := walletdir(cfg, account)
	fi, err := os.Stat(wdir)
	if err != nil {
		if os.IsNotExist(err) {
			// Attempt data directory creation
			if err = os.MkdirAll(wdir, 0700); err != nil {
				return nil, err
			}
		} else {
			return nil, err
		}
	} else {
		if !fi.IsDir() {
			return nil, fmt.Errorf("Data directory '%s' is not a directory.", cfg.DataDir)
		}
	}

	wfilepath := filepath.Join(wdir, "wallet.bin")
	txfilepath := filepath.Join(wdir, "tx.bin")
	utxofilepath := filepath.Join(wdir, "utxo.bin")
	var wfile, txfile, utxofile *os.File
	if wfile, err = os.Open(wfilepath); err != nil {
		if os.IsNotExist(err) {
			// Must create and save wallet first.
			return nil, ErrNoWallet
		} else {
			return nil, err
		}
	}
	defer wfile.Close()
	if txfile, err = os.Open(txfilepath); err != nil {
		if os.IsNotExist(err) {
			if txfile, err = os.Create(txfilepath); err != nil {
				return nil, err
			}
		} else {
			return nil, err
		}
	}
	defer txfile.Close()
	if utxofile, err = os.Open(utxofilepath); err != nil {
		if os.IsNotExist(err) {
			if utxofile, err = os.Create(utxofilepath); err != nil {
				return nil, err
			}
		} else {
			return nil, err
		}
	}
	defer utxofile.Close()

	wlt := new(wallet.Wallet)
	if _, err = wlt.ReadFrom(wfile); err != nil {
		return nil, err
	}

	var txs tx.TxStore
	if _, err = txs.ReadFrom(txfile); err != nil {
		return nil, err
	}

	var utxos tx.UtxoStore
	if _, err = utxos.ReadFrom(utxofile); err != nil {

	}

	w := &BtcWallet{
		Wallet: wlt,
	}
	w.UtxoStore.s = utxos
	w.TxStore.s = txs

	return w, nil
}

func (w *BtcWallet) Track() {
	wallets.Lock()
	name := w.Name()
	if wallets.m[name] == nil {
		wallets.m[name] = w
	}
	wallets.Unlock()

	for _, addr := range w.GetActiveAddresses() {
		go w.ReqNewTxsForAddress(addr)
	}
}

func (w *BtcWallet) RescanForAddress(addr string, blocks ...int) {
	seq.Lock()
	n := seq.n
	seq.n++
	seq.Unlock()

	params := []interface{}{addr}
	if len(blocks) > 0 {
		params = append(params, blocks[0])
	}
	if len(blocks) > 1 {
		params = append(params, blocks[1])
	}
	m := &btcjson.Message{
		Jsonrpc: "1.0",
		Id:      fmt.Sprintf("btcwallet(%v)", n),
		Method:  "rescan",
		Params:  params,
	}
	msg, _ := json.Marshal(m)

	replyHandlers.Lock()
	replyHandlers.m[n] = func(result interface{}) bool {
		// TODO(jrick)

		// btcd returns a nil result when the rescan is complete.
		// Returning true signals that this handler is finished
		// and can be removed.
		return result == nil
	}
	replyHandlers.Unlock()

	btcdMsgs <- msg
}

func (w *BtcWallet) ReqNewTxsForAddress(addr string) {
	seq.Lock()
	n := seq.n
	seq.n++
	seq.Unlock()

	m := &btcjson.Message{
		Jsonrpc: "1.0",
		Id:      fmt.Sprintf("btcwallet(%d)", n),
		Method:  "notifynewtxs",
		Params:  []interface{}{addr},
	}
	msg, _ := json.Marshal(m)

	replyHandlers.Lock()
	replyHandlers.m[n] = func(result interface{}) bool {
		// TODO(jrick): btcd also sends the block hash in the reply.
		// Do we want it saved as well?
		v, ok := result.(map[string]interface{})
		if !ok {
			log.Error("Tx Handler: Unexpected result type.")
			return false
		}
		sender58, ok := v["sender"].(string)
		if !ok {
			log.Error("Tx Handler: Unspecified sender.")
			return false
		}
		receiver58, ok := v["receiver"].(string)
		if !ok {
			log.Error("Tx Handler: Unspecified receiver.")
			return false
		}
		height, ok := v["height"].(float64)
		if !ok {
			log.Error("Tx Handler: Unspecified height.")
			return false
		}
		txhashBE, ok := v["txhash"].(string)
		if !ok {
			log.Error("Tx Handler: Unspecified transaction hash.")
			return false
		}
		amt, ok := v["amount"].(float64)
		if !ok {
			log.Error("Tx Handler: Unspecified amount.")
			return false
		}
		spent, ok := v["spent"].(bool)
		if !ok {
			log.Error("Tx Handler: Unspecified spent field.")
			return false
		}

		// btcd sends the tx  hashe as a BE string.  Convert to a
		// LE ShaHash.
		txhash, err := btcwire.NewShaHashFromStr(txhashBE)
		if err != nil {
			log.Error("Tx Handler: Tx hash string cannot be parsed: " + err.Error())
			return false
		}

		sender := btcutil.Base58Decode(sender58)
		receiver := btcutil.Base58Decode(receiver58)

		go func() {
			t := &tx.RecvTx{
				Amt: int64(amt),
			}
			copy(t.TxHash[:], txhash[:])
			copy(t.SenderAddr[:], sender)
			copy(t.ReceiverAddr[:], receiver)

			w.TxStore.Lock()
			txs := w.TxStore.s
			w.TxStore.s = append(txs, t)
			w.TxStore.dirty = true
			w.TxStore.Unlock()
		}()

		go func() {
			// Do not add output to utxo store if spent.
			if spent {
				return
			}

			u := &tx.Utxo{
				Amt:    int64(amt),
				Height: int64(height),
			}
			copy(u.TxHash[:], txhash[:])
			copy(u.Addr[:], receiver)

			w.UtxoStore.Lock()
			// All newly saved utxos are first classified as unconfirmed.
			utxos := w.UtxoStore.s.Unconfirmed
			w.UtxoStore.s.Unconfirmed = append(utxos, u)
			w.UtxoStore.dirty = true
			w.UtxoStore.Unlock()
		}()

		// Never remove this handler.
		return false
	}
	replyHandlers.Unlock()

	btcdMsgs <- msg
}

// Marshalling and unmarshalling byte arrays or slices results in ugly
// []interface{} slices where each interface{} is a float64.  This
// function unmangles this to return a byte slice.
func UnmangleJsonByteSlice(mangled []interface{}) (unmangled []byte) {
	unmangled = make([]byte, len(mangled))
	for i, _ := range mangled[:] {
		unmangled[i] = byte(mangled[i].(float64))
	}
	return unmangled
}
<commit_msg>Unbreak build caused by changes in tx pkg.<commit_after>/*
 * Copyright (c) 2013 Conformal Systems LLC <info@conformal.com>
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

package main

import (
	"encoding/json"
	"errors"
	"fmt"
	"github.com/conformal/btcjson"
	"github.com/conformal/btcutil"
	"github.com/conformal/btcwallet/tx"
	"github.com/conformal/btcwallet/wallet"
	"github.com/conformal/btcwire"
	"github.com/conformal/seelog"
	"os"
	"path/filepath"
	"sync"
	"time"
)

var (
	ErrNoWallet = errors.New("Wallet file does not exist.")
)

var (
	log     seelog.LoggerInterface = seelog.Default
	cfg     *config
	wallets = struct {
		sync.RWMutex
		m map[string]*BtcWallet
	}{
		m: make(map[string]*BtcWallet),
	}
)

func main() {
	tcfg, _, err := loadConfig()
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
	cfg = tcfg

	// Open wallet
	w, err := OpenWallet(cfg, "")
	if err != nil {
		log.Info(err.Error())
	} else {
		w.Track()
	}

	// Start HTTP server to listen and send messages to frontend and btcd
	// backend.  Try reconnection if connection failed.
	for {
		if err := ListenAndServe(); err == ConnRefused {
			// wait and try again.
			log.Info("Unable to connect to btcd. Retrying in 5 seconds.")
			time.Sleep(5 * time.Second)
		} else if err != nil {
			log.Info(err.Error())
			break
		}
	}
}

type BtcWallet struct {
	*wallet.Wallet
	mtx       sync.RWMutex
	dirty     bool
	UtxoStore struct {
		sync.RWMutex
		dirty bool
		s     tx.UtxoStore
	}
	TxStore struct {
		sync.RWMutex
		dirty bool
		s     tx.TxStore
	}
}

// walletdir returns the directory path which holds the wallet, utxo,
// and tx files.
func walletdir(cfg *config, account string) string {
	var wname string
	if account == "" {
		wname = "btcwallet"
	} else {
		wname = fmt.Sprintf("btcwallet-%s", account)
	}

	return filepath.Join(cfg.DataDir, wname)
}

func OpenWallet(cfg *config, account string) (*BtcWallet, error) {
	wdir := walletdir(cfg, account)
	fi, err := os.Stat(wdir)
	if err != nil {
		if os.IsNotExist(err) {
			// Attempt data directory creation
			if err = os.MkdirAll(wdir, 0700); err != nil {
				return nil, err
			}
		} else {
			return nil, err
		}
	} else {
		if !fi.IsDir() {
			return nil, fmt.Errorf("Data directory '%s' is not a directory.", cfg.DataDir)
		}
	}

	wfilepath := filepath.Join(wdir, "wallet.bin")
	txfilepath := filepath.Join(wdir, "tx.bin")
	utxofilepath := filepath.Join(wdir, "utxo.bin")
	var wfile, txfile, utxofile *os.File
	if wfile, err = os.Open(wfilepath); err != nil {
		if os.IsNotExist(err) {
			// Must create and save wallet first.
			return nil, ErrNoWallet
		} else {
			return nil, err
		}
	}
	defer wfile.Close()
	if txfile, err = os.Open(txfilepath); err != nil {
		if os.IsNotExist(err) {
			if txfile, err = os.Create(txfilepath); err != nil {
				return nil, err
			}
		} else {
			return nil, err
		}
	}
	defer txfile.Close()
	if utxofile, err = os.Open(utxofilepath); err != nil {
		if os.IsNotExist(err) {
			if utxofile, err = os.Create(utxofilepath); err != nil {
				return nil, err
			}
		} else {
			return nil, err
		}
	}
	defer utxofile.Close()

	wlt := new(wallet.Wallet)
	if _, err = wlt.ReadFrom(wfile); err != nil {
		return nil, err
	}

	var txs tx.TxStore
	if _, err = txs.ReadFrom(txfile); err != nil {
		return nil, err
	}

	var utxos tx.UtxoStore
	if _, err = utxos.ReadFrom(utxofile); err != nil {

	}

	w := &BtcWallet{
		Wallet: wlt,
	}
	w.UtxoStore.s = utxos
	w.TxStore.s = txs

	return w, nil
}

func (w *BtcWallet) Track() {
	wallets.Lock()
	name := w.Name()
	if wallets.m[name] == nil {
		wallets.m[name] = w
	}
	wallets.Unlock()

	for _, addr := range w.GetActiveAddresses() {
		go w.ReqNewTxsForAddress(addr)
	}
}

func (w *BtcWallet) RescanForAddress(addr string, blocks ...int) {
	seq.Lock()
	n := seq.n
	seq.n++
	seq.Unlock()

	params := []interface{}{addr}
	if len(blocks) > 0 {
		params = append(params, blocks[0])
	}
	if len(blocks) > 1 {
		params = append(params, blocks[1])
	}
	m := &btcjson.Message{
		Jsonrpc: "1.0",
		Id:      fmt.Sprintf("btcwallet(%v)", n),
		Method:  "rescan",
		Params:  params,
	}
	msg, _ := json.Marshal(m)

	replyHandlers.Lock()
	replyHandlers.m[n] = func(result interface{}) bool {
		// TODO(jrick)

		// btcd returns a nil result when the rescan is complete.
		// Returning true signals that this handler is finished
		// and can be removed.
		return result == nil
	}
	replyHandlers.Unlock()

	btcdMsgs <- msg
}

func (w *BtcWallet) ReqNewTxsForAddress(addr string) {
	seq.Lock()
	n := seq.n
	seq.n++
	seq.Unlock()

	m := &btcjson.Message{
		Jsonrpc: "1.0",
		Id:      fmt.Sprintf("btcwallet(%d)", n),
		Method:  "notifynewtxs",
		Params:  []interface{}{addr},
	}
	msg, _ := json.Marshal(m)

	replyHandlers.Lock()
	replyHandlers.m[n] = func(result interface{}) bool {
		// TODO(jrick): btcd also sends the block hash in the reply.
		// Do we want it saved as well?
		v, ok := result.(map[string]interface{})
		if !ok {
			log.Error("Tx Handler: Unexpected result type.")
			return false
		}
		sender58, ok := v["sender"].(string)
		if !ok {
			log.Error("Tx Handler: Unspecified sender.")
			return false
		}
		receiver58, ok := v["receiver"].(string)
		if !ok {
			log.Error("Tx Handler: Unspecified receiver.")
			return false
		}
		height, ok := v["height"].(float64)
		if !ok {
			log.Error("Tx Handler: Unspecified height.")
			return false
		}
		txhashBE, ok := v["txhash"].(string)
		if !ok {
			log.Error("Tx Handler: Unspecified transaction hash.")
			return false
		}
		amt, ok := v["amount"].(float64)
		if !ok {
			log.Error("Tx Handler: Unspecified amount.")
			return false
		}
		spent, ok := v["spent"].(bool)
		if !ok {
			log.Error("Tx Handler: Unspecified spent field.")
			return false
		}

		// btcd sends the tx  hashe as a BE string.  Convert to a
		// LE ShaHash.
		txhash, err := btcwire.NewShaHashFromStr(txhashBE)
		if err != nil {
			log.Error("Tx Handler: Tx hash string cannot be parsed: " + err.Error())
			return false
		}

		sender := btcutil.Base58Decode(sender58)
		receiver := btcutil.Base58Decode(receiver58)

		go func() {
			t := &tx.RecvTx{
				Amt: int64(amt),
			}
			copy(t.TxHash[:], txhash[:])
			copy(t.SenderAddr[:], sender)
			copy(t.ReceiverAddr[:], receiver)

			w.TxStore.Lock()
			txs := w.TxStore.s
			w.TxStore.s = append(txs, t)
			w.TxStore.dirty = true
			w.TxStore.Unlock()
		}()

		go func() {
			// Do not add output to utxo store if spent.
			if spent {
				return
			}

			u := &tx.Utxo{
				Amt:    int64(amt),
				Height: int64(height),
			}
			copy(u.Out.Hash[:], txhash[:])
			copy(u.Addr[:], receiver)

			w.UtxoStore.Lock()
			// All newly saved utxos are first classified as unconfirmed.
			utxos := w.UtxoStore.s.Unconfirmed
			w.UtxoStore.s.Unconfirmed = append(utxos, u)
			w.UtxoStore.dirty = true
			w.UtxoStore.Unlock()
		}()

		// Never remove this handler.
		return false
	}
	replyHandlers.Unlock()

	btcdMsgs <- msg
}

// Marshalling and unmarshalling byte arrays or slices results in ugly
// []interface{} slices where each interface{} is a float64.  This
// function unmangles this to return a byte slice.
func UnmangleJsonByteSlice(mangled []interface{}) (unmangled []byte) {
	unmangled = make([]byte, len(mangled))
	for i, _ := range mangled[:] {
		unmangled[i] = byte(mangled[i].(float64))
	}
	return unmangled
}
<|endoftext|>


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

/*
This file contains the command addons used used by the client via text commands.
*/

//
package main

import (
	"log"
)

type command struct {
	Desc    string
	Handler func(*client, []string) error
}

var cmdMap = make(map[string]command)

func init() {
	cmdMap["help"] = command{
		Desc: "help returns help information about available commands.",
		Handler: func(c *client, args []string) (e error) {
			if len(args) > 0 {
				if len(args) == 1 {
					cmds := ""
					for k := range cmdMap {
						cmds += " " + k
					}
					e = c.appendMsg("#msg-list", "Available commands:"+cmds)
				} else {
					if cmd, ok := cmdMap[args[1]]; ok {
						e = c.appendMsg("#msg-list", cmd.Desc)
					} else {
						e = c.appendMsg("#msg-list", "Command not available: "+args[1])
					}
				}
			}
			return
		},
	}
	//	cmdMap["motd"] = command{
	//		Desc: "motd prints the current message-of-the-day.",
	//		Handler: func(c *client, args []string) (e error) {
	//			if len(args) > 0 {
	//				// do something.
	//			}
	//			return
	//		},
	//	}
	cmdMap["clear"] = command{
		Desc: "clear the current terminal's content",
		Handler: func(c *client, args []string) (e error) {
			if len(args) > 0 {
				c.innerHTML("#msg-list", " ")
			}
			return
		},
	}
	cmdMap["login"] = command{
		Desc: "login lets you log into a registered user account.",
		Handler: func(c *client, args []string) (e error) {
			if len(args) > 0 {
				if len(args) == 1 {
					e = c.appendMsg("#msg-list", "Usage: login <name>")
				} else {
					name := args[1]
					if isName(name) {
						if userExists(name) {
							pass, e := c.promptSecure("#msg-txt", "Please enter your password")
							if e == nil && len(pass) > 0 {
								e = c.user.login(name, pass)
								if e != nil {
									log.Println("login error:", e)
									e = c.appendMsg("#msg-list", "Login failed")
								} else {
									e = c.innerHTML("#status-box", ""+c.user.Name+"")
									if e == nil {
										e = c.appendMsg("#msg-list", "Welcome back, "+c.user.Name)
									}
								}
							}
						} else {
							e = c.appendMsg("#msg-list", "User does not exist")
						}
					} else {
						e = c.appendMsg("#msg-list", "Invalid characters in name")
					}
				}
			}
			return
		},
	}
	cmdMap["register"] = command{
		Desc: "register a user account",
		Handler: func(c *client, args []string) (e error) {
			if len(args) > 1 {
				name := args[1]
				if isName(name) {
					if !userExists(name) {
						email, e := c.prompt("Enter your email address")
						if e == nil && isEmail(email) {
							pass1, e := c.promptSecure("#msg-txt", "Enter a good password")
							if e == nil {
								pass2, e := c.promptSecure("#msg-txt", "Re-enter your password")
								if e == nil && pass1 == pass2 {
									c.user.Email = email
									c.user.Name = name
									e = c.user.save(name, pass1, email)
									if e == nil {
										e = c.appendMsg("#msg-list", "User account created (don't forget your password!)")
									} else {
										e = c.appendMsg("#msg-list", e.Error())
									}
								} else {
									e = c.appendMsg("#msg-list", "Failed! Passwords did not match")
								}
							}
						} else {
							e = c.appendMsg("#msg-list", "Bad email address")
						}
					} else {
						e = c.appendMsg("#msg-list", "User already exists")
					}
				} else {
					e = c.appendMsg("#msg-list", "Invalid characters in name")
				}
			} else {
				e = c.appendMsg("#msg-list", "Usage: register <name>")
			}
			return
		},
	}
}
<commit_msg>Added logout command.<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

/*
This file contains the command addons used used by the client via text commands.
*/

//
package main

import (
	"log"
)

type command struct {
	Desc    string
	Handler func(*client, []string) error
}

var cmdMap = make(map[string]command)

func init() {
	cmdMap["help"] = command{
		Desc: "help returns help information about available commands.",
		Handler: func(c *client, args []string) (e error) {
			if len(args) > 0 {
				if len(args) == 1 {
					cmds := ""
					for k := range cmdMap {
						cmds += " " + k
					}
					e = c.appendMsg("#msg-list", "Available commands:"+cmds)
				} else {
					if cmd, ok := cmdMap[args[1]]; ok {
						e = c.appendMsg("#msg-list", cmd.Desc)
					} else {
						e = c.appendMsg("#msg-list", "Command not available: "+args[1])
					}
				}
			}
			return
		},
	}
	//	cmdMap["motd"] = command{
	//		Desc: "motd prints the current message-of-the-day.",
	//		Handler: func(c *client, args []string) (e error) {
	//			if len(args) > 0 {
	//				// do something.
	//			}
	//			return
	//		},
	//	}
	cmdMap["clear"] = command{
		Desc: "clear the current terminal's content",
		Handler: func(c *client, args []string) (e error) {
			if len(args) > 0 {
				c.innerHTML("#msg-list", " ")
			}
			return
		},
	}
	cmdMap["login"] = command{
		Desc: "login lets you log into a registered user account.",
		Handler: func(c *client, args []string) (e error) {
			if len(args) > 0 {
				if len(args) == 1 {
					e = c.appendMsg("#msg-list", "Usage: login <name>")
				} else {
					name := args[1]
					if isName(name) {
						if userExists(name) {
							pass, e := c.promptSecure("#msg-txt", "Please enter your password")
							if e == nil && len(pass) > 0 {
								e = c.user.login(name, pass)
								if e != nil {
									log.Println("login error:", e)
									e = c.appendMsg("#msg-list", "Login failed")
								} else {
									e = c.innerHTML("#status-box", ""+c.user.Name+"")
									if e == nil {
										e = c.appendMsg("#msg-list", "Welcome back, "+c.user.Name)
									}
								}
							}
						} else {
							e = c.appendMsg("#msg-list", "User does not exist")
						}
					} else {
						e = c.appendMsg("#msg-list", "Invalid characters in name")
					}
				}
			}
			return
		},
	}
	cmdMap["logout"] = command{
		Desc: "logout lets you log out of the connected user account.",
		Handler: func(c *client, args []string) (e error) {
			if c.user.auth == true {
				e = c.innerHTML("#status-box", "Guest")
				c.user = user{Name: "Guest"}
			} else {
				e = c.appendMsg("#msg-list", "Not logged in.")
			}
			return
		},
	}
	cmdMap["register"] = command{
		Desc: "register a user account",
		Handler: func(c *client, args []string) (e error) {
			if len(args) > 1 {
				name := args[1]
				if isName(name) {
					if !userExists(name) {
						email, e := c.prompt("Enter your email address")
						if e == nil && isEmail(email) {
							pass1, e := c.promptSecure("#msg-txt", "Enter a good password")
							if e == nil {
								pass2, e := c.promptSecure("#msg-txt", "Re-enter your password")
								if e == nil && pass1 == pass2 {
									c.user.Email = email
									c.user.Name = name
									e = c.user.save(name, pass1, email)
									if e == nil {
										e = c.appendMsg("#msg-list", "User account created (don't forget your password!)")
									} else {
										e = c.appendMsg("#msg-list", e.Error())
									}
								} else {
									e = c.appendMsg("#msg-list", "Failed! Passwords did not match")
								}
							}
						} else {
							e = c.appendMsg("#msg-list", "Bad email address")
						}
					} else {
						e = c.appendMsg("#msg-list", "User already exists")
					}
				} else {
					e = c.appendMsg("#msg-list", "Invalid characters in name")
				}
			} else {
				e = c.appendMsg("#msg-list", "Usage: register <name>")
			}
			return
		},
	}
}
<|endoftext|>


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

package execute

import (
	"context"
	"fmt"

	"go.chromium.org/luci/common/logging"
	"google.golang.org/protobuf/proto"

	cfgpb "go.chromium.org/luci/cv/api/config/v2"
	"go.chromium.org/luci/cv/api/recipe/v1"
	"go.chromium.org/luci/cv/internal/tryjob"
)

// retriability indicates whether a tryjob is allowed or denied to retry.
type retriability int

const (
	// The tryjob cannot be retried.
	retryDenied retriability = iota + 1
	// The tryjob needs to be retried if quota allows.
	retryAllowed
	// The tryjob failed, but was not triggered by this run, so trigger with
	// no regard for retry quota.
	retryAllowedIgnoreQuota
)

// canRetry determines whether an individual attempt can be retried.
func canRetry(definition *tryjob.Definition, attempt *tryjob.ExecutionState_Execution_Attempt) retriability {
	switch {
	// If this run did not trigger the tryjob, do not count its failure
	// towards retry quota. Also ignores the output_retry_denied property.
	case attempt.Reused:
		return retryAllowedIgnoreQuota
	// Check if the tryjob forbids retry.
	case attempt.Result.GetOutput().GetRetry() == recipe.Output_OUTPUT_RETRY_DENIED:
		return retryDenied
	default:
		return retryAllowed
	}
}

func ensureTryjobCriticalAndFailed(def *tryjob.Definition, attempt *tryjob.ExecutionState_Execution_Attempt) {
	switch {
	case !def.Critical:
		panic(fmt.Errorf("calling canRetryAll for non critical tryjob: %s", def))
	case attempt.Status != tryjob.Status_ENDED:
		panic(fmt.Errorf("calling canRetryAll for non ended tryjob: %s", def))
	case attempt.Result.Status == tryjob.Result_SUCCEEDED:
		panic(fmt.Errorf("calling canRetryAll for succeeded tryjob: %s", def))
	}
}

// canRetryAll checks whether all failed critical tryjobs can be retried.
//
// Panics when any provided index points to a tryjob that is not critical or
// its latest attempt doesn't fail.
//
// Always updates the consumed quota for failed executions.
func (e *Executor) canRetryAll(
	ctx context.Context,
	execState *tryjob.ExecutionState,
	failedIndices []int,
) bool {
	if !isRetryEnabled(execState.GetRequirement().GetRetryConfig()) {
		e.logRetryDenied(ctx, execState, failedIndices, "retry is not enabled in the config")
		return false
	}

	var totalUsedQuota int32
	canRetryAll := true
	for _, execution := range execState.GetExecutions() {
		totalUsedQuota += execution.GetUsedQuota()
	}
	retryConfig := execState.GetRequirement().GetRetryConfig()
	for _, idx := range failedIndices {
		definition := execState.GetRequirement().GetDefinitions()[idx]
		exec := execState.GetExecutions()[idx]
		attempt := tryjob.LatestAttempt(exec)
		ensureTryjobCriticalAndFailed(definition, attempt)
		switch canRetry(definition, attempt) {
		case retryDenied:
			canRetryAll = false
			e.logRetryDenied(ctx, execState, []int{idx}, "tryjob explicitly denies retry in its output")
		case retryAllowedIgnoreQuota:
		case retryAllowed:
			var quotaToRetry int32
			switch attempt.Result.Status {
			case tryjob.Result_TIMEOUT:
				quotaToRetry = retryConfig.GetTimeoutWeight()
			case tryjob.Result_FAILED_PERMANENTLY:
				quotaToRetry = retryConfig.GetFailureWeight()
			case tryjob.Result_FAILED_TRANSIENTLY:
				quotaToRetry = retryConfig.GetTransientFailureWeight()
			default:
				logging.Errorf(ctx, "fallback to use failure status because of unexpected attempt result status %s", attempt.Result.Status)
				quotaToRetry = retryConfig.GetFailureWeight()
			}
			exec.UsedQuota += quotaToRetry
			totalUsedQuota += quotaToRetry
			if exec.UsedQuota > retryConfig.GetSingleQuota() {
				canRetryAll = false
				e.logRetryDenied(ctx, execState, []int{idx}, "insufficient quota")
			}
		default:
			panic(fmt.Errorf("unknown retriability"))
		}
	}

	if canRetryAll && totalUsedQuota > retryConfig.GetGlobalQuota() {
		canRetryAll = false
		e.logRetryDenied(ctx, execState, failedIndices, "insufficient global quota")
	}
	return canRetryAll
}

var emptyRetryConfig = &cfgpb.Verifiers_Tryjob_RetryConfig{}

func isRetryEnabled(rc *cfgpb.Verifiers_Tryjob_RetryConfig) bool {
	switch {
	case rc == nil:
		return false
	case proto.Equal(rc, emptyRetryConfig):
		return false
	default:
		return true
	}
}
<commit_msg>cv: Minor refactoring, reorder funcs in execute/retry.go<commit_after>// Copyright 2022 The LUCI Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package execute

import (
	"context"
	"fmt"

	"go.chromium.org/luci/common/logging"
	"google.golang.org/protobuf/proto"

	cfgpb "go.chromium.org/luci/cv/api/config/v2"
	"go.chromium.org/luci/cv/api/recipe/v1"
	"go.chromium.org/luci/cv/internal/tryjob"
)

// retriability indicates whether a tryjob is allowed or denied to retry.
type retriability int

const (
	// The tryjob cannot be retried.
	retryDenied retriability = iota + 1
	// The tryjob needs to be retried if quota allows.
	retryAllowed
	// The tryjob failed, but was not triggered by this run, so trigger with
	// no regard for retry quota.
	retryAllowedIgnoreQuota
)

// canRetryAll checks whether all failed critical tryjobs can be retried.
//
// Panics when any provided index points to a tryjob that is not critical or
// its latest attempt doesn't fail.
//
// Always updates the consumed quota for failed executions.
func (e *Executor) canRetryAll(
	ctx context.Context,
	execState *tryjob.ExecutionState,
	failedIndices []int,
) bool {
	if !isRetryEnabled(execState.GetRequirement().GetRetryConfig()) {
		e.logRetryDenied(ctx, execState, failedIndices, "retry is not enabled in the config")
		return false
	}

	var totalUsedQuota int32
	canRetryAll := true
	for _, execution := range execState.GetExecutions() {
		totalUsedQuota += execution.GetUsedQuota()
	}
	retryConfig := execState.GetRequirement().GetRetryConfig()
	for _, idx := range failedIndices {
		definition := execState.GetRequirement().GetDefinitions()[idx]
		exec := execState.GetExecutions()[idx]
		attempt := tryjob.LatestAttempt(exec)
		ensureTryjobCriticalAndFailed(definition, attempt)
		switch canRetry(definition, attempt) {
		case retryDenied:
			canRetryAll = false
			e.logRetryDenied(ctx, execState, []int{idx}, "tryjob explicitly denies retry in its output")
		case retryAllowedIgnoreQuota:
		case retryAllowed:
			var quotaToRetry int32
			switch attempt.Result.Status {
			case tryjob.Result_TIMEOUT:
				quotaToRetry = retryConfig.GetTimeoutWeight()
			case tryjob.Result_FAILED_PERMANENTLY:
				quotaToRetry = retryConfig.GetFailureWeight()
			case tryjob.Result_FAILED_TRANSIENTLY:
				quotaToRetry = retryConfig.GetTransientFailureWeight()
			default:
				logging.Errorf(ctx, "fallback to use failure status because of unexpected attempt result status %s", attempt.Result.Status)
				quotaToRetry = retryConfig.GetFailureWeight()
			}
			exec.UsedQuota += quotaToRetry
			totalUsedQuota += quotaToRetry
			if exec.UsedQuota > retryConfig.GetSingleQuota() {
				canRetryAll = false
				e.logRetryDenied(ctx, execState, []int{idx}, "insufficient quota")
			}
		default:
			panic(fmt.Errorf("unknown retriability"))
		}
	}

	if canRetryAll && totalUsedQuota > retryConfig.GetGlobalQuota() {
		canRetryAll = false
		e.logRetryDenied(ctx, execState, failedIndices, "insufficient global quota")
	}
	return canRetryAll
}

func ensureTryjobCriticalAndFailed(def *tryjob.Definition, attempt *tryjob.ExecutionState_Execution_Attempt) {
	switch {
	case !def.Critical:
		panic(fmt.Errorf("calling canRetryAll for non critical tryjob: %s", def))
	case attempt.Status != tryjob.Status_ENDED:
		panic(fmt.Errorf("calling canRetryAll for non ended tryjob: %s", def))
	case attempt.Result.Status == tryjob.Result_SUCCEEDED:
		panic(fmt.Errorf("calling canRetryAll for succeeded tryjob: %s", def))
	}
}

// canRetry determines whether an individual attempt can be retried.
func canRetry(definition *tryjob.Definition, attempt *tryjob.ExecutionState_Execution_Attempt) retriability {
	switch {
	// If this run did not trigger the tryjob, do not count its failure
	// towards retry quota. Also ignores the output_retry_denied property.
	case attempt.Reused:
		return retryAllowedIgnoreQuota
	// Check if the tryjob forbids retry.
	case attempt.Result.GetOutput().GetRetry() == recipe.Output_OUTPUT_RETRY_DENIED:
		return retryDenied
	default:
		return retryAllowed
	}
}

var emptyRetryConfig = &cfgpb.Verifiers_Tryjob_RetryConfig{}

func isRetryEnabled(rc *cfgpb.Verifiers_Tryjob_RetryConfig) bool {
	switch {
	case rc == nil:
		return false
	case proto.Equal(rc, emptyRetryConfig):
		return false
	default:
		return true
	}
}
<|endoftext|>