<commit_before>package main

import (
	"errors"
	log "github.com/Sirupsen/logrus"
	"github.com/loadimpact/speedboat/js"
	"github.com/loadimpact/speedboat/lib"
	"github.com/loadimpact/speedboat/simple"
	"github.com/loadimpact/speedboat/stats"
	"github.com/loadimpact/speedboat/stats/accumulate"
	"github.com/loadimpact/speedboat/stats/influxdb"
	"github.com/loadimpact/speedboat/stats/writer"
	"github.com/urfave/cli"
	"golang.org/x/net/context"
	"gopkg.in/yaml.v2"
	"io/ioutil"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"time"
)

const (
	typeURL = "url"
	typeJS  = "js"
)

// Help text template
const helpTemplate = `NAME:
   {{.Name}} - {{.Usage}}

USAGE:
   {{if .UsageText}}{{.UsageText}}{{else}}{{.HelpName}} {{if .VisibleFlags}}[options] {{end}}filename|url{{end}}
   {{if .Version}}{{if not .HideVersion}}
VERSION:
   {{.Version}}
   {{end}}{{end}}{{if .VisibleFlags}}
OPTIONS:
   {{range .VisibleFlags}}{{.}}
   {{end}}{{end}}
`

var mVUs = stats.Stat{Name: "vus", Type: stats.GaugeType}

func pollVURamping(ctx context.Context, t lib.Test) <-chan int {
	ch := make(chan int)
	startTime := time.Now()

	go func() {
		defer close(ch)

		ticker := time.NewTicker(1 * time.Second)
		for {
			select {
			case <-ticker.C:
				ch <- t.VUsAt(time.Since(startTime))
			case <-ctx.Done():
				return
			}
		}
	}()

	return ch
}

func parseBackend(out string) (stats.Backend, error) {
	switch {
	case strings.HasPrefix(out, "influxdb+"):
		url := strings.TrimPrefix(out, "influxdb+")
		return influxdb.NewFromURL(url)
	default:
		return nil, errors.New("Unknown output destination")
	}
}

func parseStages(vus []string, total time.Duration) (stages []lib.TestStage, err error) {
	accountedTime := time.Duration(0)
	fluidStages := []int{}
	for i, spec := range vus {
		parts := strings.SplitN(spec, ":", 2)
		countParts := strings.SplitN(parts[0], "-", 2)

		stage := lib.TestStage{}

		// An absent first part means keep going from the last stage's end
		// If it's the first stage, just start with 0
		if countParts[0] == "" {
			if i > 0 {
				stage.StartVUs = stages[i-1].EndVUs
			}
		} else {
			start, err := strconv.ParseInt(countParts[0], 10, 64)
			if err != nil {
				return stages, err
			}
			stage.StartVUs = int(start)
		}

		// If an end is specified, use that, otherwise keep the VU level constant
		if len(countParts) > 1 && countParts[1] != "" {
			end, err := strconv.ParseInt(countParts[1], 10, 64)
			if err != nil {
				return stages, err
			}
			stage.EndVUs = int(end)
		} else {
			stage.EndVUs = stage.StartVUs
		}

		// If a time is specified, use that, otherwise mark the stage as "fluid", allotting it an
		// even slice of what time remains after all fixed stages are accounted for (may be 0)
		if len(parts) > 1 {
			duration, err := time.ParseDuration(parts[1])
			if err != nil {
				return stages, err
			}
			stage.Duration = duration
			accountedTime += duration
		} else {
			fluidStages = append(fluidStages, i)
		}

		stages = append(stages, stage)
	}

	// We're ignoring fluid stages if the fixed stages already take up all the allotted time
	// Otherwise, evenly divide the test's remaining time between all fluid stages
	if len(fluidStages) > 0 && accountedTime < total {
		fluidDuration := (total - accountedTime) / time.Duration(len(fluidStages))
		for _, i := range fluidStages {
			stage := stages[i]
			stage.Duration = fluidDuration
			stages[i] = stage
		}
	}

	return stages, nil
}

func guessType(arg string) string {
	switch {
	case strings.Contains(arg, "://"):
		return typeURL
	case strings.HasSuffix(arg, ".js"):
		return typeJS
	}
	return ""
}

func readAll(filename string) ([]byte, error) {
	if filename == "-" {
		return ioutil.ReadAll(os.Stdin)
	}

	return ioutil.ReadFile(filename)
}

func makeRunner(t lib.Test, filename, typ string) (lib.Runner, error) {
	if typ == typeURL {
		return simple.New(filename), nil
	}

	bytes, err := readAll(filename)
	if err != nil {
		return nil, err
	}

	switch typ {
	case typeJS:
		return js.New(filename, string(bytes)), nil
	default:
		return nil, errors.New("Type ambiguous, please specify -t/--type")
	}
}

func action(cc *cli.Context) error {
	if cc.IsSet("verbose") {
		log.SetLevel(log.DebugLevel)
	}

	for _, out := range cc.StringSlice("out") {
		backend, err := parseBackend(out)
		if err != nil {
			return cli.NewExitError(err.Error(), 1)
		}
		stats.DefaultRegistry.Backends = append(stats.DefaultRegistry.Backends, backend)
	}

	var formatter writer.Formatter
	switch cc.String("format") {
	case "":
	case "json":
		formatter = writer.JSONFormatter{}
	case "prettyjson":
		formatter = writer.PrettyJSONFormatter{}
	case "yaml":
		formatter = writer.YAMLFormatter{}
	default:
		return cli.NewExitError("Unknown output format", 1)
	}

	var summarizer *Summarizer
	if formatter != nil {
		filter := stats.MakeFilter(cc.StringSlice("exclude"), cc.StringSlice("select"))
		if cc.Bool("raw") {
			backend := &writer.Backend{
				Writer:    os.Stdout,
				Formatter: formatter,
			}
			backend.Filter = filter
			stats.DefaultRegistry.Backends = append(stats.DefaultRegistry.Backends, backend)
		} else {
			accumulator := accumulate.New()
			accumulator.Filter = filter
			accumulator.GroupBy = cc.StringSlice("group-by")
			stats.DefaultRegistry.Backends = append(stats.DefaultRegistry.Backends, accumulator)

			summarizer = &Summarizer{
				Accumulator: accumulator,
				Formatter:   formatter,
			}
		}
	}

	stages, err := parseStages(cc.StringSlice("vus"), cc.Duration("duration"))
	if err != nil {
		return cli.NewExitError(err.Error(), 1)
	}
	t := lib.Test{Stages: stages}

	var r lib.Runner
	switch len(cc.Args()) {
	case 0:
		cli.ShowAppHelp(cc)
		return nil
	case 1:
		filename := cc.Args()[0]
		typ := cc.String("type")
		if typ == "" {
			typ = guessType(filename)
		}

		if filename == "-" && typ == "" {
			typ = typeJS
		}

		runner, err := makeRunner(t, filename, typ)
		if err != nil {
			return cli.NewExitError(err.Error(), 1)
		}
		r = runner
	default:
		return cli.NewExitError("Too many arguments!", 1)
	}

	if cc.Bool("plan") {
		data, err := yaml.Marshal(map[string]interface{}{
			"stages": stages,
		})
		if err != nil {
			return cli.NewExitError(err.Error(), 1)
		}
		os.Stdout.Write(data)
		return nil
	}

	vus := lib.VUGroup{
		Pool: lib.VUPool{
			New: r.NewVU,
		},
		RunOnce: func(ctx context.Context, vu lib.VU) {
			if err := vu.RunOnce(ctx); err != nil {
				log.WithError(err).Error("Uncaught Error")
			}
		},
	}

	for i := 0; i < t.MaxVUs(); i++ {
		vu, err := vus.Pool.New()
		if err != nil {
			return cli.NewExitError(err.Error(), 1)
		}
		vus.Pool.Put(vu)
	}

	ctx, cancel := context.WithTimeout(context.Background(), t.TotalDuration())

	go func() {
		ticker := time.NewTicker(1 * time.Second)
		for {
			select {
			case <-ticker.C:
				if err := stats.Submit(); err != nil {
					log.WithError(err).Error("[Couldn't submit stats]")
				}
			case <-ctx.Done():
				return
			}
		}
	}()

	interval := cc.Duration("interval")
	if interval > 0 && summarizer != nil {
		go func() {
			ticker := time.NewTicker(interval)
			for {
				select {
				case <-ticker.C:
					if err := summarizer.Print(os.Stdout); err != nil {
						log.WithError(err).Error("Couldn't print statistics!")
					}
				case <-ctx.Done():
					return
				}
			}
		}()
	}

	go func() {
		quit := make(chan os.Signal)
		signal.Notify(quit)

		select {
		case <-quit:
			cancel()
		case <-ctx.Done():
		}
	}()

	if !cc.Bool("quiet") {
		log.WithFields(log.Fields{
			"at":     time.Now(),
			"length": t.TotalDuration(),
		}).Info("Starting test...")
	}

	vus.Start(ctx)
	scaleTo := pollVURamping(ctx, t)
mainLoop:
	for {
		select {
		case num := <-scaleTo:
			vus.Scale(num)
			stats.Add(stats.Sample{
				Stat:   &mVUs,
				Values: stats.Value(float64(num)),
			})
		case <-ctx.Done():
			break mainLoop
		}
	}

	vus.Stop()

	stats.Add(stats.Sample{Stat: &mVUs, Values: stats.Value(0)})
	stats.Submit()

	if summarizer != nil {
		if err := summarizer.Print(os.Stdout); err != nil {
			log.WithError(err).Error("Couldn't print statistics!")
		}
	}

	return nil
}

func main() {
	// Free up -v and -h for our own flags
	cli.VersionFlag.Name = "version"
	cli.HelpFlag.Name = "help, ?"
	cli.AppHelpTemplate = helpTemplate

	// Bootstrap the app from commandline flags
	app := cli.NewApp()
	app.Name = "speedboat"
	app.Usage = "A next-generation load generator"
	app.Version = "1.0.0-mvp1"
	app.Flags = []cli.Flag{
		cli.BoolFlag{
			Name:  "plan",
			Usage: "Don't run anything, just show the test plan",
		},
		cli.StringFlag{
			Name:  "type, t",
			Usage: "Input file type, if not evident (url or js)",
		},
		cli.StringSliceFlag{
			Name:  "vus, u",
			Usage: "Number of VUs to simulate",
			Value: &cli.StringSlice{"10"},
		},
		cli.DurationFlag{
			Name:  "duration, d",
			Usage: "Test duration",
			Value: time.Duration(10) * time.Second,
		},
		cli.BoolFlag{
			Name:  "verbose, v",
			Usage: "More verbose output",
		},
		cli.BoolFlag{
			Name:  "quiet, q",
			Usage: "Suppress the summary at the end of a test",
		},
		cli.StringFlag{
			Name:  "format, f",
			Usage: "Format for printed metrics (yaml, json, prettyjson)",
			Value: "yaml",
		},
		cli.DurationFlag{
			Name:  "interval, i",
			Usage: "Write periodic summaries",
		},
		cli.StringSliceFlag{
			Name:  "out, o",
			Usage: "Write metrics to a database",
		},
		cli.BoolFlag{
			Name:  "raw",
			Usage: "Instead of summaries, dump raw samples to stdout",
		},
		cli.StringSliceFlag{
			Name:  "select, s",
			Usage: "Include only named metrics",
		},
		cli.StringSliceFlag{
			Name:  "exclude, e",
			Usage: "Exclude named metrics",
		},
		cli.StringSliceFlag{
			Name:  "group-by, g",
			Usage: "Group metrics by tags",
		},
	}
	app.Action = action
	app.Run(os.Args)
}
<commit_msg>[fix] That was inserting an implicit 10-VU stage<commit_after>package main

import (
	"errors"
	log "github.com/Sirupsen/logrus"
	"github.com/loadimpact/speedboat/js"
	"github.com/loadimpact/speedboat/lib"
	"github.com/loadimpact/speedboat/simple"
	"github.com/loadimpact/speedboat/stats"
	"github.com/loadimpact/speedboat/stats/accumulate"
	"github.com/loadimpact/speedboat/stats/influxdb"
	"github.com/loadimpact/speedboat/stats/writer"
	"github.com/urfave/cli"
	"golang.org/x/net/context"
	"gopkg.in/yaml.v2"
	"io/ioutil"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"time"
)

const (
	typeURL = "url"
	typeJS  = "js"
)

// Help text template
const helpTemplate = `NAME:
   {{.Name}} - {{.Usage}}

USAGE:
   {{if .UsageText}}{{.UsageText}}{{else}}{{.HelpName}} {{if .VisibleFlags}}[options] {{end}}filename|url{{end}}
   {{if .Version}}{{if not .HideVersion}}
VERSION:
   {{.Version}}
   {{end}}{{end}}{{if .VisibleFlags}}
OPTIONS:
   {{range .VisibleFlags}}{{.}}
   {{end}}{{end}}
`

var mVUs = stats.Stat{Name: "vus", Type: stats.GaugeType}

func pollVURamping(ctx context.Context, t lib.Test) <-chan int {
	ch := make(chan int)
	startTime := time.Now()

	go func() {
		defer close(ch)

		ticker := time.NewTicker(1 * time.Second)
		for {
			select {
			case <-ticker.C:
				ch <- t.VUsAt(time.Since(startTime))
			case <-ctx.Done():
				return
			}
		}
	}()

	return ch
}

func parseBackend(out string) (stats.Backend, error) {
	switch {
	case strings.HasPrefix(out, "influxdb+"):
		url := strings.TrimPrefix(out, "influxdb+")
		return influxdb.NewFromURL(url)
	default:
		return nil, errors.New("Unknown output destination")
	}
}

func parseStages(vus []string, total time.Duration) (stages []lib.TestStage, err error) {
	if len(vus) == 0 {
		return []lib.TestStage{
			lib.TestStage{Duration: total, StartVUs: 10, EndVUs: 10},
		}, nil
	}

	accountedTime := time.Duration(0)
	fluidStages := []int{}
	for i, spec := range vus {
		parts := strings.SplitN(spec, ":", 2)
		countParts := strings.SplitN(parts[0], "-", 2)

		stage := lib.TestStage{}

		// An absent first part means keep going from the last stage's end
		// If it's the first stage, just start with 0
		if countParts[0] == "" {
			if i > 0 {
				stage.StartVUs = stages[i-1].EndVUs
			}
		} else {
			start, err := strconv.ParseInt(countParts[0], 10, 64)
			if err != nil {
				return stages, err
			}
			stage.StartVUs = int(start)
		}

		// If an end is specified, use that, otherwise keep the VU level constant
		if len(countParts) > 1 && countParts[1] != "" {
			end, err := strconv.ParseInt(countParts[1], 10, 64)
			if err != nil {
				return stages, err
			}
			stage.EndVUs = int(end)
		} else {
			stage.EndVUs = stage.StartVUs
		}

		// If a time is specified, use that, otherwise mark the stage as "fluid", allotting it an
		// even slice of what time remains after all fixed stages are accounted for (may be 0)
		if len(parts) > 1 {
			duration, err := time.ParseDuration(parts[1])
			if err != nil {
				return stages, err
			}
			stage.Duration = duration
			accountedTime += duration
		} else {
			fluidStages = append(fluidStages, i)
		}

		stages = append(stages, stage)
	}

	// We're ignoring fluid stages if the fixed stages already take up all the allotted time
	// Otherwise, evenly divide the test's remaining time between all fluid stages
	if len(fluidStages) > 0 && accountedTime < total {
		fluidDuration := (total - accountedTime) / time.Duration(len(fluidStages))
		for _, i := range fluidStages {
			stage := stages[i]
			stage.Duration = fluidDuration
			stages[i] = stage
		}
	}

	return stages, nil
}

func guessType(arg string) string {
	switch {
	case strings.Contains(arg, "://"):
		return typeURL
	case strings.HasSuffix(arg, ".js"):
		return typeJS
	}
	return ""
}

func readAll(filename string) ([]byte, error) {
	if filename == "-" {
		return ioutil.ReadAll(os.Stdin)
	}

	return ioutil.ReadFile(filename)
}

func makeRunner(t lib.Test, filename, typ string) (lib.Runner, error) {
	if typ == typeURL {
		return simple.New(filename), nil
	}

	bytes, err := readAll(filename)
	if err != nil {
		return nil, err
	}

	switch typ {
	case typeJS:
		return js.New(filename, string(bytes)), nil
	default:
		return nil, errors.New("Type ambiguous, please specify -t/--type")
	}
}

func action(cc *cli.Context) error {
	if cc.IsSet("verbose") {
		log.SetLevel(log.DebugLevel)
	}

	for _, out := range cc.StringSlice("out") {
		backend, err := parseBackend(out)
		if err != nil {
			return cli.NewExitError(err.Error(), 1)
		}
		stats.DefaultRegistry.Backends = append(stats.DefaultRegistry.Backends, backend)
	}

	var formatter writer.Formatter
	switch cc.String("format") {
	case "":
	case "json":
		formatter = writer.JSONFormatter{}
	case "prettyjson":
		formatter = writer.PrettyJSONFormatter{}
	case "yaml":
		formatter = writer.YAMLFormatter{}
	default:
		return cli.NewExitError("Unknown output format", 1)
	}

	var summarizer *Summarizer
	if formatter != nil {
		filter := stats.MakeFilter(cc.StringSlice("exclude"), cc.StringSlice("select"))
		if cc.Bool("raw") {
			backend := &writer.Backend{
				Writer:    os.Stdout,
				Formatter: formatter,
			}
			backend.Filter = filter
			stats.DefaultRegistry.Backends = append(stats.DefaultRegistry.Backends, backend)
		} else {
			accumulator := accumulate.New()
			accumulator.Filter = filter
			accumulator.GroupBy = cc.StringSlice("group-by")
			stats.DefaultRegistry.Backends = append(stats.DefaultRegistry.Backends, accumulator)

			summarizer = &Summarizer{
				Accumulator: accumulator,
				Formatter:   formatter,
			}
		}
	}

	stages, err := parseStages(cc.StringSlice("vus"), cc.Duration("duration"))
	if err != nil {
		return cli.NewExitError(err.Error(), 1)
	}
	t := lib.Test{Stages: stages}

	var r lib.Runner
	switch len(cc.Args()) {
	case 0:
		cli.ShowAppHelp(cc)
		return nil
	case 1:
		filename := cc.Args()[0]
		typ := cc.String("type")
		if typ == "" {
			typ = guessType(filename)
		}

		if filename == "-" && typ == "" {
			typ = typeJS
		}

		runner, err := makeRunner(t, filename, typ)
		if err != nil {
			return cli.NewExitError(err.Error(), 1)
		}
		r = runner
	default:
		return cli.NewExitError("Too many arguments!", 1)
	}

	if cc.Bool("plan") {
		data, err := yaml.Marshal(map[string]interface{}{
			"stages": stages,
		})
		if err != nil {
			return cli.NewExitError(err.Error(), 1)
		}
		os.Stdout.Write(data)
		return nil
	}

	vus := lib.VUGroup{
		Pool: lib.VUPool{
			New: r.NewVU,
		},
		RunOnce: func(ctx context.Context, vu lib.VU) {
			if err := vu.RunOnce(ctx); err != nil {
				log.WithError(err).Error("Uncaught Error")
			}
		},
	}

	for i := 0; i < t.MaxVUs(); i++ {
		vu, err := vus.Pool.New()
		if err != nil {
			return cli.NewExitError(err.Error(), 1)
		}
		vus.Pool.Put(vu)
	}

	ctx, cancel := context.WithTimeout(context.Background(), t.TotalDuration())

	go func() {
		ticker := time.NewTicker(1 * time.Second)
		for {
			select {
			case <-ticker.C:
				if err := stats.Submit(); err != nil {
					log.WithError(err).Error("[Couldn't submit stats]")
				}
			case <-ctx.Done():
				return
			}
		}
	}()

	interval := cc.Duration("interval")
	if interval > 0 && summarizer != nil {
		go func() {
			ticker := time.NewTicker(interval)
			for {
				select {
				case <-ticker.C:
					if err := summarizer.Print(os.Stdout); err != nil {
						log.WithError(err).Error("Couldn't print statistics!")
					}
				case <-ctx.Done():
					return
				}
			}
		}()
	}

	go func() {
		quit := make(chan os.Signal)
		signal.Notify(quit)

		select {
		case <-quit:
			cancel()
		case <-ctx.Done():
		}
	}()

	if !cc.Bool("quiet") {
		log.WithFields(log.Fields{
			"at":     time.Now(),
			"length": t.TotalDuration(),
		}).Info("Starting test...")
	}

	vus.Start(ctx)
	scaleTo := pollVURamping(ctx, t)
mainLoop:
	for {
		select {
		case num := <-scaleTo:
			vus.Scale(num)
			stats.Add(stats.Sample{
				Stat:   &mVUs,
				Values: stats.Value(float64(num)),
			})
		case <-ctx.Done():
			break mainLoop
		}
	}

	vus.Stop()

	stats.Add(stats.Sample{Stat: &mVUs, Values: stats.Value(0)})
	stats.Submit()

	if summarizer != nil {
		if err := summarizer.Print(os.Stdout); err != nil {
			log.WithError(err).Error("Couldn't print statistics!")
		}
	}

	return nil
}

func main() {
	// Free up -v and -h for our own flags
	cli.VersionFlag.Name = "version"
	cli.HelpFlag.Name = "help, ?"
	cli.AppHelpTemplate = helpTemplate

	// Bootstrap the app from commandline flags
	app := cli.NewApp()
	app.Name = "speedboat"
	app.Usage = "A next-generation load generator"
	app.Version = "1.0.0-mvp1"
	app.Flags = []cli.Flag{
		cli.BoolFlag{
			Name:  "plan",
			Usage: "Don't run anything, just show the test plan",
		},
		cli.StringFlag{
			Name:  "type, t",
			Usage: "Input file type, if not evident (url or js)",
		},
		cli.StringSliceFlag{
			Name:  "vus, u",
			Usage: "Number of VUs to simulate",
		},
		cli.DurationFlag{
			Name:  "duration, d",
			Usage: "Test duration",
			Value: time.Duration(10) * time.Second,
		},
		cli.BoolFlag{
			Name:  "verbose, v",
			Usage: "More verbose output",
		},
		cli.BoolFlag{
			Name:  "quiet, q",
			Usage: "Suppress the summary at the end of a test",
		},
		cli.StringFlag{
			Name:  "format, f",
			Usage: "Format for printed metrics (yaml, json, prettyjson)",
			Value: "yaml",
		},
		cli.DurationFlag{
			Name:  "interval, i",
			Usage: "Write periodic summaries",
		},
		cli.StringSliceFlag{
			Name:  "out, o",
			Usage: "Write metrics to a database",
		},
		cli.BoolFlag{
			Name:  "raw",
			Usage: "Instead of summaries, dump raw samples to stdout",
		},
		cli.StringSliceFlag{
			Name:  "select, s",
			Usage: "Include only named metrics",
		},
		cli.StringSliceFlag{
			Name:  "exclude, e",
			Usage: "Exclude named metrics",
		},
		cli.StringSliceFlag{
			Name:  "group-by, g",
			Usage: "Group metrics by tags",
		},
	}
	app.Action = action
	app.Run(os.Args)
}
<|endoftext|>


<commit_before>package datadog

import (
	"fmt"
	"github.com/stretchr/testify/assert"
	"net/http"
	"testing"
)

func Test_isRateLimited(t *testing.T) {
	tests := []struct {
		desc             string
		endpoint         string
		method           string
		isRateLimited    bool
		endpointFormated string
	}{
		{
			desc:             "is rate limited",
			endpoint:         "/v1/query?&query=avg:system.cpu.user{*}by{host}",
			method:           "GET",
			isRateLimited:    true,
			endpointFormated: "/v1/query",
		},
		{
			desc:             "is not rate limited",
			endpoint:         "/v1/series?api_key=12",
			method:           "POST",
			isRateLimited:    false,
			endpointFormated: "",
		},
		{
			desc:             "is rate limited but wrong method",
			endpoint:         "/v1/query?&query=avg:system.cpu.user{*}by{host}",
			method:           "POST",
			isRateLimited:    false,
			endpointFormated: "",
		},
	}
	for i, tt := range tests {
		t.Run(fmt.Sprintf("#%d %s", i, tt.desc), func(t *testing.T) {
			limited, edpt := isRateLimited(tt.method, tt.endpoint)
			assert.Equal(t, limited, tt.isRateLimited)
			assert.Equal(t, edpt, tt.endpointFormated)
		})
	}
}

func Test_updateRateLimits(t *testing.T) {
	// fake client to ensure that we are race free.
	client := Client{
		rateLimitingStats: make(map[string]RateLimit),
	}
	tests := []struct {
		desc   string
		api    string
		resp   *http.Response
		header RateLimit
		error  error
	}{
		{
			"nominal case query",
			"/v1/query",
			makeHeader("1", "2", "3", "4"),
			RateLimit{"1", "2", "3", "4"},
			nil,
		},
		{
			"nominal case logs",
			"/v1/logs-queries/list",
			makeHeader("2", "2", "1", "5"),
			RateLimit{"2", "2", "1", "5"},
			nil,
		},
		{
			"no response",
			"",
			nil,
			RateLimit{},
			fmt.Errorf("could not parse headers from the HTTP response."),
		},
		{
			"no header",
			"/v2/error",
			makeEmptyHeader(),
			RateLimit{},
			fmt.Errorf("could not parse headers from the HTTP response."),
		},
		{
			"update case query",
			"/v1/query",
			makeHeader("2", "4", "6", "4"),
			RateLimit{"2", "4", "6", "4"},
			nil,
		},
	}

	for i, tt := range tests {
		t.Run(fmt.Sprintf("#%d %s", i, tt.desc), func(t *testing.T) {
			err := client.updateRateLimits(tt.resp, tt.api)
			assert.Equal(t, tt.error, err)
			assert.Equal(t, tt.header, client.rateLimitingStats[tt.api])
		})
	}
}

func makeHeader(limit, period, reset, remaining string) *http.Response {
	h := http.Response{
		Header: make(map[string][]string),
	}
	h.Header.Set("X-RateLimit-Limit", limit)
	h.Header.Set("X-RateLimit-Reset", reset)
	h.Header.Set("X-RateLimit-Period", period)
	h.Header.Set("X-RateLimit-Remaining", remaining)
	return &h
}

func makeEmptyHeader() *http.Response {
	return &http.Response{
		Header: nil,
	}
}
<commit_msg>per @xornivore's review<commit_after>package datadog

import (
	"fmt"
	"github.com/stretchr/testify/assert"
	"net/http"
	"testing"
)

func Test_isRateLimited(t *testing.T) {
	tests := []struct {
		desc             string
		endpoint         string
		method           string
		isRateLimited    bool
		endpointFormated string
	}{
		{
			desc:             "is rate limited",
			endpoint:         "/v1/query?&query=avg:system.cpu.user{*}by{host}",
			method:           "GET",
			isRateLimited:    true,
			endpointFormated: "/v1/query",
		},
		{
			desc:             "is not rate limited",
			endpoint:         "/v1/series?api_key=12",
			method:           "POST",
			isRateLimited:    false,
			endpointFormated: "",
		},
		{
			desc:             "is rate limited but wrong method",
			endpoint:         "/v1/query?&query=avg:system.cpu.user{*}by{host}",
			method:           "POST",
			isRateLimited:    false,
			endpointFormated: "",
		},
	}
	for i, tt := range tests {
		t.Run(fmt.Sprintf("#%d %s", i, tt.desc), func(t *testing.T) {
			limited, edpt := isRateLimited(tt.method, tt.endpoint)
			assert.Equal(t, limited, tt.isRateLimited)
			assert.Equal(t, edpt, tt.endpointFormated)
		})
	}
}

func Test_updateRateLimits(t *testing.T) {
	// fake client to ensure that we are race free.
	client := Client{
		rateLimitingStats: make(map[string]RateLimit),
	}
	tests := []struct {
		desc   string
		api    string
		resp   *http.Response
		header RateLimit
		error  error
	}{
		{
			"nominal case query",
			"/v1/query",
			makeHeader(RateLimit{"1", "2", "3", "4"}),
			RateLimit{"1", "2", "3", "4"},
			nil,
		},
		{
			"nominal case logs",
			"/v1/logs-queries/list",
			makeHeader(RateLimit{"2", "2", "1", "5"}),
			RateLimit{"2", "2", "1", "5"},
			nil,
		},
		{
			"no response",
			"",
			nil,
			RateLimit{},
			fmt.Errorf("could not parse headers from the HTTP response."),
		},
		{
			"no header",
			"/v2/error",
			makeEmptyHeader(),
			RateLimit{},
			fmt.Errorf("could not parse headers from the HTTP response."),
		},
		{
			"update case query",
			"/v1/query",
			makeHeader(RateLimit{"2", "4", "6", "4"}),
			RateLimit{"2", "4", "6", "4"},
			nil,
		},
	}

	for i, tt := range tests {
		t.Run(fmt.Sprintf("#%d %s", i, tt.desc), func(t *testing.T) {
			err := client.updateRateLimits(tt.resp, tt.api)
			assert.Equal(t, tt.error, err)
			assert.Equal(t, tt.header, client.rateLimitingStats[tt.api])
		})
	}
}

func makeHeader(r RateLimit) *http.Response {
	h := http.Response{
		Header: make(map[string][]string),
	}
	h.Header.Set("X-RateLimit-Limit", r.Limit)
	h.Header.Set("X-RateLimit-Reset", r.Reset)
	h.Header.Set("X-RateLimit-Period", r.Period)
	h.Header.Set("X-RateLimit-Remaining", r.Remaining)
	return &h
}

func makeEmptyHeader() *http.Response {
	return &http.Response{
		Header: nil,
	}
}
<|endoftext|>


<commit_before>/*
Package bootstrap implements the capability to connect to an existing and online
Tinzenite peer network.

TODO: add encryption bootstrap capabilities
*/
package bootstrap

import (
	"io/ioutil"

	"github.com/tinzenite/channel"
	"github.com/tinzenite/shared"
)

/*
Create returns a struct that will allow to bootstrap to an existing
Tinzenite network. To actually start bootstrapping call Bootstrap.Start(address).
*/
func Create(path, localPeerName string, f Success) (*Bootstrap, error) {
	if shared.IsTinzenite(path) {
		return nil, shared.ErrIsTinzenite
	}
	// build structure
	err := shared.MakeDotTinzenite(path)
	if err != nil {
		return nil, err
	}
	// create object
	boot := &Bootstrap{
		path:   path,
		onDone: f}
	boot.cInterface = createChanInterface(boot)
	channel, err := channel.Create(localPeerName, nil, boot.cInterface)
	if err != nil {
		return nil, err
	}
	boot.channel = channel
	// get address for peer
	address, err := boot.channel.Address()
	if err != nil {
		return nil, err
	}
	// make peer (at correct location!)
	peer, err := shared.CreatePeer(localPeerName, address)
	if err != nil {
		return nil, err
	}
	boot.peer = peer
	// bg stuff
	boot.wg.Add(1)
	boot.stop = make(chan bool, 1)
	go boot.run()
	return boot, nil
}

/*
Load tries to load the given directory as a bootstrap object, allowing it to
connect to an existing network. To actually start bootstrapping call
Bootstrap.Start(address). NOTE: will fail if already connected to other peers!

TODO: strictly speaking we only need the selfpeer... look into this?
*/
func Load(path string, f Success) (*Bootstrap, error) {
	if !shared.IsTinzenite(path) {
		return nil, shared.ErrNotTinzenite
	}
	if !checkValidBootstrap(path) {
		return nil, errNotBootstrapCapable
	}
	// create object
	boot := &Bootstrap{
		path:   path,
		onDone: f}
	boot.cInterface = createChanInterface(boot)
	// load
	toxPeerDump, err := shared.LoadToxDump(path)
	if err != nil {
		return nil, err
	}
	boot.peer = toxPeerDump.SelfPeer
	channel, err := channel.Create(boot.peer.Name, toxPeerDump.ToxData, boot.cInterface)
	if err != nil {
		return nil, err
	}
	boot.channel = channel
	// bg stuff
	boot.wg.Add(1)
	boot.stop = make(chan bool, 1)
	go boot.run()
	return boot, nil
}

/*
checkValidBootstrap from given path. To be valid: must be .TINZENITEDIR, must
have only one peer (hopefully itself), and local/self.json must exit.
*/
func checkValidBootstrap(path string) bool {
	if !shared.IsTinzenite(path) {
		return false
	}
	tinzenPath := path + "/" + shared.TINZENITEDIR
	stats, err := ioutil.ReadDir(tinzenPath + "/" + shared.ORGDIR + "/" + shared.PEERSDIR)
	if err != nil {
		return false
	}
	exists, _ := shared.FileExists(tinzenPath + "/" + shared.LOCALDIR + "/" + shared.SELFPEERJSON)
	return len(stats) == 1 && exists
}
<commit_msg>applied trusted peer<commit_after>/*
Package bootstrap implements the capability to connect to an existing and online
Tinzenite peer network.

TODO: add encryption bootstrap capabilities
*/
package bootstrap

import (
	"io/ioutil"

	"github.com/tinzenite/channel"
	"github.com/tinzenite/shared"
)

/*
Create returns a struct that will allow to bootstrap to an existing
Tinzenite network. To actually start bootstrapping call Bootstrap.Start(address).
*/
func Create(path, localPeerName string, f Success) (*Bootstrap, error) {
	if shared.IsTinzenite(path) {
		return nil, shared.ErrIsTinzenite
	}
	// build structure
	err := shared.MakeDotTinzenite(path)
	if err != nil {
		return nil, err
	}
	// create object
	boot := &Bootstrap{
		path:   path,
		onDone: f}
	boot.cInterface = createChanInterface(boot)
	channel, err := channel.Create(localPeerName, nil, boot.cInterface)
	if err != nil {
		return nil, err
	}
	boot.channel = channel
	// get address for peer
	address, err := boot.channel.Address()
	if err != nil {
		return nil, err
	}
	// make peer (at correct location!)
	peer, err := shared.CreatePeer(localPeerName, address, false)
	if err != nil {
		return nil, err
	}
	boot.peer = peer
	// bg stuff
	boot.wg.Add(1)
	boot.stop = make(chan bool, 1)
	go boot.run()
	return boot, nil
}

/*
Load tries to load the given directory as a bootstrap object, allowing it to
connect to an existing network. To actually start bootstrapping call
Bootstrap.Start(address). NOTE: will fail if already connected to other peers!

TODO: strictly speaking we only need the selfpeer... look into this?
*/
func Load(path string, f Success) (*Bootstrap, error) {
	if !shared.IsTinzenite(path) {
		return nil, shared.ErrNotTinzenite
	}
	if !checkValidBootstrap(path) {
		return nil, errNotBootstrapCapable
	}
	// create object
	boot := &Bootstrap{
		path:   path,
		onDone: f}
	boot.cInterface = createChanInterface(boot)
	// load
	toxPeerDump, err := shared.LoadToxDump(path)
	if err != nil {
		return nil, err
	}
	boot.peer = toxPeerDump.SelfPeer
	channel, err := channel.Create(boot.peer.Name, toxPeerDump.ToxData, boot.cInterface)
	if err != nil {
		return nil, err
	}
	boot.channel = channel
	// bg stuff
	boot.wg.Add(1)
	boot.stop = make(chan bool, 1)
	go boot.run()
	return boot, nil
}

/*
checkValidBootstrap from given path. To be valid: must be .TINZENITEDIR, must
have only one peer (hopefully itself), and local/self.json must exit.
*/
func checkValidBootstrap(path string) bool {
	if !shared.IsTinzenite(path) {
		return false
	}
	tinzenPath := path + "/" + shared.TINZENITEDIR
	stats, err := ioutil.ReadDir(tinzenPath + "/" + shared.ORGDIR + "/" + shared.PEERSDIR)
	if err != nil {
		return false
	}
	exists, _ := shared.FileExists(tinzenPath + "/" + shared.LOCALDIR + "/" + shared.SELFPEERJSON)
	return len(stats) == 1 && exists
}
<|endoftext|>


<commit_before>package scheduler

import (
	"errors"
	"fmt"
	"sync"
	"time"

	"github.com/andres-erbsen/clock"
	"github.com/uber-common/bark"
	"golang.org/x/sync/syncmap"

	"code.uber.internal/go-common.git/x/log"
	"code.uber.internal/infra/kraken/.gen/go/p2p"
	"code.uber.internal/infra/kraken/lib/torrent/networkevent"
	"code.uber.internal/infra/kraken/lib/torrent/storage"
	"code.uber.internal/infra/kraken/torlib"
)

var (
	errPeerAlreadyDispatched   = errors.New("peer is already dispatched for the torrent")
	errPieceOutOfBounds        = errors.New("piece index out of bounds")
	errChunkNotSupported       = errors.New("reading / writing chunk of piece not supported")
	errRepeatedBitfieldMessage = errors.New("received repeated bitfield message")
)

type dispatcherFactory struct {
	Config               Config
	LocalPeerID          torlib.PeerID
	EventSender          eventSender
	Clock                clock.Clock
	NetworkEventProducer networkevent.Producer
}

// New creates a new dispatcher for the given torrent.
func (f *dispatcherFactory) New(t storage.Torrent) *dispatcher {
	d := &dispatcher{
		Torrent:              t,
		CreatedAt:            f.Clock.Now(),
		localPeerID:          f.LocalPeerID,
		eventSender:          f.EventSender,
		clock:                f.Clock,
		networkEventProducer: f.NetworkEventProducer,
	}
	if t.Complete() {
		d.complete.Do(func() { go d.eventSender.Send(completedDispatcherEvent{d}) })
	}
	return d
}

// dispatcher coordinates torrent state with sending / receiving messages between multiple
// peers. As such, dispatcher and torrent have a one-to-one relationship, while dispatcher
// and conn have a many-to-many relationship.
type dispatcher struct {
	Torrent     storage.Torrent
	CreatedAt   time.Time
	localPeerID torlib.PeerID
	clock       clock.Clock

	// Maps torlib.PeerID to *conn.
	conns syncmap.Map

	eventSender eventSender

	networkEventProducer networkevent.Producer

	mu              sync.Mutex // Protects the following fields:
	lastConnRemoved time.Time

	// complete ensures dispatcher only sends complete event once to scheduler
	complete sync.Once
}

func (d *dispatcher) LastConnRemoved() time.Time {
	d.mu.Lock()
	defer d.mu.Unlock()

	return d.lastConnRemoved
}

func (d *dispatcher) touchLastConnRemoved() {
	d.mu.Lock()
	defer d.mu.Unlock()

	d.lastConnRemoved = d.clock.Now()
}

// Empty returns true if the dispatcher has no conns.
func (d *dispatcher) Empty() bool {
	// syncmap.Map does not provide a length function, hence this poor man's
	// implementation of `len(d.conns) == 0`.
	empty := true
	d.conns.Range(func(k, v interface{}) bool {
		empty = false
		return false
	})
	return empty
}

// AddConn registers a new conn with the dispatcher.
func (d *dispatcher) AddConn(c *conn) error {
	if c.InfoHash != d.Torrent.InfoHash() {
		return errors.New("conn initialized for wrong torrent")
	}
	if _, ok := d.conns.LoadOrStore(c.PeerID, c); ok {
		return errors.New("conn already exists")
	}
	go d.sendInitialPieceRequests(c)
	go d.feed(c)
	return nil
}

// TearDown closes all dispatcher connections.
func (d *dispatcher) TearDown() {
	d.conns.Range(func(k, v interface{}) bool {
		conn := v.(*conn)
		conn.Close()
		return true
	})
}

func (d *dispatcher) String() string {
	return fmt.Sprintf("dispatcher(%s)", d.Torrent)
}

func (d *dispatcher) sendPieceRequest(c *conn, i int) error {
	d.logf(log.Fields{"conn": c, "piece": i}).Info("Sending piece request")
	return c.Send(newPieceRequestMessage(i, d.Torrent.PieceLength(i)))
}

func (d *dispatcher) sendInitialPieceRequests(c *conn) {
	for _, i := range d.Torrent.MissingPieces() {
		if !c.Bitfield.Has(i) {
			continue
		}
		if err := d.sendPieceRequest(c, i); err != nil {
			// Connection closed.
			break
		}
	}
}

// feed reads off of c and handles incoming messages. When c closes, the feed
// goroutine removes c from the dispatcher and exits.
func (d *dispatcher) feed(c *conn) {
	for msg := range c.Receiver() {
		if err := d.dispatch(c, msg); err != nil {
			d.log().Errorf("Error dispatching message: %s", err)
			// TODO Maybe close conn?
		}
	}
	d.logf(log.Fields{"peer": c.PeerID}).Debug("Dispatcher feeder exiting")
	d.conns.Delete(c.PeerID)
	d.touchLastConnRemoved()
}

func (d *dispatcher) dispatch(c *conn, msg *message) error {
	switch msg.Message.Type {
	case p2p.Message_ERROR:
		d.handleError(c, msg.Message.Error)
	case p2p.Message_ANNOUCE_PIECE:
		d.handleAnnouncePiece(c, msg.Message.AnnouncePiece)
	case p2p.Message_PIECE_REQUEST:
		d.handlePieceRequest(c, msg.Message.PieceRequest)
	case p2p.Message_PIECE_PAYLOAD:
		d.handlePiecePayload(c, msg.Message.PiecePayload, msg.Payload)
	case p2p.Message_CANCEL_PIECE:
		d.handleCancelPiece(c, msg.Message.CancelPiece)
	case p2p.Message_BITFIELD:
		d.handleBitfield(c, msg.Message.Bitfield)
	default:
		return fmt.Errorf("unknown message type: %d", msg.Message.Type)
	}
	return nil
}

func (d *dispatcher) handleError(c *conn, msg *p2p.ErrorMessage) {
	switch msg.Code {
	case p2p.ErrorMessage_PIECE_REQUEST_FAILED:
		d.log().Errorf("Piece request failed: %s", msg.Error)
		d.sendPieceRequest(c, int(msg.Index))
	}
}

func (d *dispatcher) handleAnnouncePiece(c *conn, msg *p2p.AnnouncePieceMessage) {
	if int(msg.Index) >= d.Torrent.NumPieces() {
		d.log().Errorf("Announce piece out of bounds: %d >= %d", msg.Index, d.Torrent.NumPieces())
		return
	}
	i := int(msg.Index)
	c.Bitfield.Set(i, true)
	if !d.Torrent.HasPiece(i) {
		d.sendPieceRequest(c, i)
	}
}

func (d *dispatcher) isFullPiece(i, offset, length int) bool {
	return offset == 0 && length == int(d.Torrent.PieceLength(i))
}

func (d *dispatcher) handlePieceRequest(c *conn, msg *p2p.PieceRequestMessage) {
	d.logf(log.Fields{"conn": c, "piece": msg.Index}).Debug("Received piece request")

	i := int(msg.Index)
	if !d.isFullPiece(i, int(msg.Offset), int(msg.Length)) {
		d.logf(log.Fields{"conn": c, "piece": i}).Error("Rejecting piece request: chunk not supported")
		c.Send(newErrorMessage(i, p2p.ErrorMessage_PIECE_REQUEST_FAILED, errChunkNotSupported))
		return
	}

	payload, err := d.Torrent.ReadPiece(i)
	if err != nil {
		d.logf(log.Fields{"conn": c, "piece": i}).Errorf("Error reading requested piece: %s", err)
		c.Send(newErrorMessage(i, p2p.ErrorMessage_PIECE_REQUEST_FAILED, err))
		return
	}
	if err := c.Send(newPiecePayloadMessage(i, payload)); err != nil {
		d.logf(log.Fields{"conn": c, "piece": i}).Errorf("Failed to send piece: %s", err)
		return
	}
	c.TouchLastPieceSent()
	d.logf(log.Fields{"conn": c, "piece": msg.Index}).Info("Sent piece")
}

func (d *dispatcher) handlePiecePayload(
	c *conn, msg *p2p.PiecePayloadMessage, payload []byte) {

	// TODO(codyg): re-request piece if write failed.

	i := int(msg.Index)
	if !d.isFullPiece(i, int(msg.Offset), int(msg.Length)) {
		d.logf(log.Fields{"conn": c, "piece": i}).Error("Rejecting piece payload: chunk not supported")
		return
	}
	if err := d.Torrent.WritePiece(payload, i); err != nil {
		if err != storage.ErrPieceComplete {
			d.logf(log.Fields{"conn": c, "piece": i}).Errorf("Error writing piece payload: %s", err)
		}
		return
	}
	d.logf(log.Fields{"conn": c, "piece": i}).Info("Received piece")
	d.networkEventProducer.Produce(
		networkevent.ReceivePieceEvent(d.Torrent.InfoHash(), d.localPeerID, c.PeerID, i))
	c.TouchLastGoodPieceReceived()
	if d.Torrent.Complete() {
		d.complete.Do(func() { go d.eventSender.Send(completedDispatcherEvent{d}) })
	}

	d.conns.Range(func(k, v interface{}) bool {
		if k.(torlib.PeerID) == c.PeerID {
			return true
		}
		cc := v.(*conn)

		d.logf(log.Fields{"conn": cc, "piece": i}).Debug("Announcing piece")

		// Ignore error -- this just means the connection was closed. The feed goroutine
		// for cc will clean up.
		// TODO(codyg): We need to slim down the number of peers we announce a new
		// piece to.  We could just rely on announcing to the tracker instead of flooding
		// the network with tons of announce piece requests.
		cc.Send(newAnnouncePieceMessage(i))

		return true
	})
}

func (d *dispatcher) handleCancelPiece(c *conn, msg *p2p.CancelPieceMessage) {
	// No-op: cancelling not supported because all received messages are synchronized,
	// therefore if we receive a cancel it is already too late -- we've already read
	// the piece.
}

func (d *dispatcher) handleBitfield(c *conn, msg *p2p.BitfieldMessage) {
	log.WithFields(log.Fields{
		"peer": c.PeerID,
	}).Error("Unexpected bitfield message from established conn")
}

func (d *dispatcher) logf(f log.Fields) bark.Logger {
	f["torrent"] = d.Torrent
	return log.WithFields(f)
}

func (d *dispatcher) log() bark.Logger {
	return d.logf(log.Fields{})
}
<commit_msg>Reduce log level for piece events<commit_after>package scheduler

import (
	"errors"
	"fmt"
	"sync"
	"time"

	"github.com/andres-erbsen/clock"
	"github.com/uber-common/bark"
	"golang.org/x/sync/syncmap"

	"code.uber.internal/go-common.git/x/log"
	"code.uber.internal/infra/kraken/.gen/go/p2p"
	"code.uber.internal/infra/kraken/lib/torrent/networkevent"
	"code.uber.internal/infra/kraken/lib/torrent/storage"
	"code.uber.internal/infra/kraken/torlib"
)

var (
	errPeerAlreadyDispatched   = errors.New("peer is already dispatched for the torrent")
	errPieceOutOfBounds        = errors.New("piece index out of bounds")
	errChunkNotSupported       = errors.New("reading / writing chunk of piece not supported")
	errRepeatedBitfieldMessage = errors.New("received repeated bitfield message")
)

type dispatcherFactory struct {
	Config               Config
	LocalPeerID          torlib.PeerID
	EventSender          eventSender
	Clock                clock.Clock
	NetworkEventProducer networkevent.Producer
}

// New creates a new dispatcher for the given torrent.
func (f *dispatcherFactory) New(t storage.Torrent) *dispatcher {
	d := &dispatcher{
		Torrent:              t,
		CreatedAt:            f.Clock.Now(),
		localPeerID:          f.LocalPeerID,
		eventSender:          f.EventSender,
		clock:                f.Clock,
		networkEventProducer: f.NetworkEventProducer,
	}
	if t.Complete() {
		d.complete.Do(func() { go d.eventSender.Send(completedDispatcherEvent{d}) })
	}
	return d
}

// dispatcher coordinates torrent state with sending / receiving messages between multiple
// peers. As such, dispatcher and torrent have a one-to-one relationship, while dispatcher
// and conn have a many-to-many relationship.
type dispatcher struct {
	Torrent     storage.Torrent
	CreatedAt   time.Time
	localPeerID torlib.PeerID
	clock       clock.Clock

	// Maps torlib.PeerID to *conn.
	conns syncmap.Map

	eventSender eventSender

	networkEventProducer networkevent.Producer

	mu              sync.Mutex // Protects the following fields:
	lastConnRemoved time.Time

	// complete ensures dispatcher only sends complete event once to scheduler
	complete sync.Once
}

func (d *dispatcher) LastConnRemoved() time.Time {
	d.mu.Lock()
	defer d.mu.Unlock()

	return d.lastConnRemoved
}

func (d *dispatcher) touchLastConnRemoved() {
	d.mu.Lock()
	defer d.mu.Unlock()

	d.lastConnRemoved = d.clock.Now()
}

// Empty returns true if the dispatcher has no conns.
func (d *dispatcher) Empty() bool {
	// syncmap.Map does not provide a length function, hence this poor man's
	// implementation of `len(d.conns) == 0`.
	empty := true
	d.conns.Range(func(k, v interface{}) bool {
		empty = false
		return false
	})
	return empty
}

// AddConn registers a new conn with the dispatcher.
func (d *dispatcher) AddConn(c *conn) error {
	if c.InfoHash != d.Torrent.InfoHash() {
		return errors.New("conn initialized for wrong torrent")
	}
	if _, ok := d.conns.LoadOrStore(c.PeerID, c); ok {
		return errors.New("conn already exists")
	}
	go d.sendInitialPieceRequests(c)
	go d.feed(c)
	return nil
}

// TearDown closes all dispatcher connections.
func (d *dispatcher) TearDown() {
	d.conns.Range(func(k, v interface{}) bool {
		conn := v.(*conn)
		conn.Close()
		return true
	})
}

func (d *dispatcher) String() string {
	return fmt.Sprintf("dispatcher(%s)", d.Torrent)
}

func (d *dispatcher) sendPieceRequest(c *conn, i int) error {
	d.logf(log.Fields{"conn": c, "piece": i}).Debug("Sending piece request")
	return c.Send(newPieceRequestMessage(i, d.Torrent.PieceLength(i)))
}

func (d *dispatcher) sendInitialPieceRequests(c *conn) {
	for _, i := range d.Torrent.MissingPieces() {
		if !c.Bitfield.Has(i) {
			continue
		}
		if err := d.sendPieceRequest(c, i); err != nil {
			// Connection closed.
			break
		}
	}
}

// feed reads off of c and handles incoming messages. When c closes, the feed
// goroutine removes c from the dispatcher and exits.
func (d *dispatcher) feed(c *conn) {
	for msg := range c.Receiver() {
		if err := d.dispatch(c, msg); err != nil {
			d.log().Errorf("Error dispatching message: %s", err)
			// TODO Maybe close conn?
		}
	}
	d.logf(log.Fields{"peer": c.PeerID}).Debug("Dispatcher feeder exiting")
	d.conns.Delete(c.PeerID)
	d.touchLastConnRemoved()
}

func (d *dispatcher) dispatch(c *conn, msg *message) error {
	switch msg.Message.Type {
	case p2p.Message_ERROR:
		d.handleError(c, msg.Message.Error)
	case p2p.Message_ANNOUCE_PIECE:
		d.handleAnnouncePiece(c, msg.Message.AnnouncePiece)
	case p2p.Message_PIECE_REQUEST:
		d.handlePieceRequest(c, msg.Message.PieceRequest)
	case p2p.Message_PIECE_PAYLOAD:
		d.handlePiecePayload(c, msg.Message.PiecePayload, msg.Payload)
	case p2p.Message_CANCEL_PIECE:
		d.handleCancelPiece(c, msg.Message.CancelPiece)
	case p2p.Message_BITFIELD:
		d.handleBitfield(c, msg.Message.Bitfield)
	default:
		return fmt.Errorf("unknown message type: %d", msg.Message.Type)
	}
	return nil
}

func (d *dispatcher) handleError(c *conn, msg *p2p.ErrorMessage) {
	switch msg.Code {
	case p2p.ErrorMessage_PIECE_REQUEST_FAILED:
		d.log().Errorf("Piece request failed: %s", msg.Error)
		d.sendPieceRequest(c, int(msg.Index))
	}
}

func (d *dispatcher) handleAnnouncePiece(c *conn, msg *p2p.AnnouncePieceMessage) {
	if int(msg.Index) >= d.Torrent.NumPieces() {
		d.log().Errorf("Announce piece out of bounds: %d >= %d", msg.Index, d.Torrent.NumPieces())
		return
	}
	i := int(msg.Index)
	c.Bitfield.Set(i, true)
	if !d.Torrent.HasPiece(i) {
		d.sendPieceRequest(c, i)
	}
}

func (d *dispatcher) isFullPiece(i, offset, length int) bool {
	return offset == 0 && length == int(d.Torrent.PieceLength(i))
}

func (d *dispatcher) handlePieceRequest(c *conn, msg *p2p.PieceRequestMessage) {
	d.logf(log.Fields{"conn": c, "piece": msg.Index}).Debug("Received piece request")

	i := int(msg.Index)
	if !d.isFullPiece(i, int(msg.Offset), int(msg.Length)) {
		d.logf(log.Fields{"conn": c, "piece": i}).Error("Rejecting piece request: chunk not supported")
		c.Send(newErrorMessage(i, p2p.ErrorMessage_PIECE_REQUEST_FAILED, errChunkNotSupported))
		return
	}

	payload, err := d.Torrent.ReadPiece(i)
	if err != nil {
		d.logf(log.Fields{"conn": c, "piece": i}).Errorf("Error reading requested piece: %s", err)
		c.Send(newErrorMessage(i, p2p.ErrorMessage_PIECE_REQUEST_FAILED, err))
		return
	}
	if err := c.Send(newPiecePayloadMessage(i, payload)); err != nil {
		d.logf(log.Fields{"conn": c, "piece": i}).Errorf("Failed to send piece: %s", err)
		return
	}
	c.TouchLastPieceSent()
	d.logf(log.Fields{"conn": c, "piece": msg.Index}).Debug("Sent piece")
}

func (d *dispatcher) handlePiecePayload(
	c *conn, msg *p2p.PiecePayloadMessage, payload []byte) {

	// TODO(codyg): re-request piece if write failed.

	i := int(msg.Index)
	if !d.isFullPiece(i, int(msg.Offset), int(msg.Length)) {
		d.logf(log.Fields{"conn": c, "piece": i}).Error("Rejecting piece payload: chunk not supported")
		return
	}
	if err := d.Torrent.WritePiece(payload, i); err != nil {
		if err != storage.ErrPieceComplete {
			d.logf(log.Fields{"conn": c, "piece": i}).Errorf("Error writing piece payload: %s", err)
		}
		return
	}
	d.logf(log.Fields{"conn": c, "piece": i}).Debug("Received piece")
	d.networkEventProducer.Produce(
		networkevent.ReceivePieceEvent(d.Torrent.InfoHash(), d.localPeerID, c.PeerID, i))
	c.TouchLastGoodPieceReceived()
	if d.Torrent.Complete() {
		d.complete.Do(func() { go d.eventSender.Send(completedDispatcherEvent{d}) })
	}

	d.conns.Range(func(k, v interface{}) bool {
		if k.(torlib.PeerID) == c.PeerID {
			return true
		}
		cc := v.(*conn)

		d.logf(log.Fields{"conn": cc, "piece": i}).Debug("Announcing piece")

		// Ignore error -- this just means the connection was closed. The feed goroutine
		// for cc will clean up.
		// TODO(codyg): We need to slim down the number of peers we announce a new
		// piece to.  We could just rely on announcing to the tracker instead of flooding
		// the network with tons of announce piece requests.
		cc.Send(newAnnouncePieceMessage(i))

		return true
	})
}

func (d *dispatcher) handleCancelPiece(c *conn, msg *p2p.CancelPieceMessage) {
	// No-op: cancelling not supported because all received messages are synchronized,
	// therefore if we receive a cancel it is already too late -- we've already read
	// the piece.
}

func (d *dispatcher) handleBitfield(c *conn, msg *p2p.BitfieldMessage) {
	log.WithFields(log.Fields{
		"peer": c.PeerID,
	}).Error("Unexpected bitfield message from established conn")
}

func (d *dispatcher) logf(f log.Fields) bark.Logger {
	f["torrent"] = d.Torrent
	return log.WithFields(f)
}

func (d *dispatcher) log() bark.Logger {
	return d.logf(log.Fields{})
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
)

const (
	Version         = "0.1.0"
	ExitCodeSuccess = 0
	ExitCodeError   = 1
)

type Mock struct {
	Name    string
	Meta    Meta
	Screens []Screen
	Launch  Launch
	Colors  []Color
	Strings []String
}

type Meta struct {
	Android Android
}

type Android struct {
	Package             string
	GradlePluginVersion string `json:"gradle_plugin_version"`
	BuildToolsVersion   string `json:"build_tools_version"`
	MinSdkVersion       int    `json:"min_sdk_version"`
	TargetSdkVersion    int    `json:"target_sdk_version"`
	CompileSdkVersion   string `json:"compile_sdk_version"`
	VersionCode         int    `json:"version_code"`
	VersionName         string `json:"version_name"`
}

type Screen struct {
	Id        string
	Name      string
	Layout    Layout
	Behaviors []Behavior
}

type Layout struct {
	Views []View
}

type View struct {
	Id     string
	Type   string
	Label  string
	SizeW  string `json:"size_w"`
	SizeH  string `json:"size_h"`
	AlignH string `json:"align_h"`
	AlignV string `json:"align_v"`
}

type Behavior struct {
	Trigger Trigger
	Action  Action
}

type Trigger struct {
	Type   string
	Widget string
}

type Action struct {
	Type    string
	Transit string
}

type Launch struct {
	Screen string
}

type Color struct {
	Id    string
	Value string
}

type String struct {
	Lang string
	Defs []Def
}

type Def struct {
	Id    string
	Value string
}

func main() {
	// Parse command
	if len(os.Args) == 1 {
		printUsage()
		os.Exit(ExitCodeError)
	}
	switch os.Args[1] {
	case "gen":
	case "version":
		printVersion()
		os.Exit(ExitCodeSuccess)
	case "help":
		fallthrough
	default:
		printUsage()
		os.Exit(ExitCodeError)
	}

	if len(os.Args) < 3 {
		printUsage()
		os.Exit(ExitCodeError)
	}
	genId := os.Args[2]

	mock := parseConfigs()
	gen(&mock, genId)
}

func printUsage() {
	fmt.Println(`mocker is a mock up framework for mobile apps.

Usage:
	mocker command [options]

Command:
	gen      generate source code (see 'Generator')
	help     show this help
	version  show version of mocker

Generator:
	mocker gen ID

	ID:
		ios  Objective-C code for iOS app
`)
}

func printVersion() {
	fmt.Println("mocker version \"" + Version + "\"")
}

func parseConfigs() (mock Mock) {
	filename := "Mockerfile"
	xmlFile, err := os.Open(filename)
	if err != nil {
		fmt.Println("Error opening file", err)
		return
	}
	defer xmlFile.Close()

	b, _ := ioutil.ReadAll(xmlFile)
	err = json.Unmarshal(b, &mock)
	if err != nil {
		fmt.Println("Error unmarshaling Mockerfile", err)
		return
	}

	return
}
<commit_msg>Added Android to usage.<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
)

const (
	Version         = "0.1.0"
	ExitCodeSuccess = 0
	ExitCodeError   = 1
)

type Mock struct {
	Name    string
	Meta    Meta
	Screens []Screen
	Launch  Launch
	Colors  []Color
	Strings []String
}

type Meta struct {
	Android Android
}

type Android struct {
	Package             string
	GradlePluginVersion string `json:"gradle_plugin_version"`
	BuildToolsVersion   string `json:"build_tools_version"`
	MinSdkVersion       int    `json:"min_sdk_version"`
	TargetSdkVersion    int    `json:"target_sdk_version"`
	CompileSdkVersion   string `json:"compile_sdk_version"`
	VersionCode         int    `json:"version_code"`
	VersionName         string `json:"version_name"`
}

type Screen struct {
	Id        string
	Name      string
	Layout    Layout
	Behaviors []Behavior
}

type Layout struct {
	Views []View
}

type View struct {
	Id     string
	Type   string
	Label  string
	SizeW  string `json:"size_w"`
	SizeH  string `json:"size_h"`
	AlignH string `json:"align_h"`
	AlignV string `json:"align_v"`
}

type Behavior struct {
	Trigger Trigger
	Action  Action
}

type Trigger struct {
	Type   string
	Widget string
}

type Action struct {
	Type    string
	Transit string
}

type Launch struct {
	Screen string
}

type Color struct {
	Id    string
	Value string
}

type String struct {
	Lang string
	Defs []Def
}

type Def struct {
	Id    string
	Value string
}

func main() {
	// Parse command
	if len(os.Args) == 1 {
		printUsage()
		os.Exit(ExitCodeError)
	}
	switch os.Args[1] {
	case "gen":
	case "version":
		printVersion()
		os.Exit(ExitCodeSuccess)
	case "help":
		fallthrough
	default:
		printUsage()
		os.Exit(ExitCodeError)
	}

	if len(os.Args) < 3 {
		printUsage()
		os.Exit(ExitCodeError)
	}
	genId := os.Args[2]

	mock := parseConfigs()
	gen(&mock, genId)
}

func printUsage() {
	fmt.Println(`mocker is a mock up framework for mobile apps.

Usage:
	mocker command [options]

Command:
	gen      generate source code (see 'Generator')
	help     show this help
	version  show version of mocker

Generator:
	mocker gen ID

	ID:
		android  Java and XML code for Android app
		ios      Objective-C code for iOS app
`)
}

func printVersion() {
	fmt.Println("mocker version \"" + Version + "\"")
}

func parseConfigs() (mock Mock) {
	filename := "Mockerfile"
	xmlFile, err := os.Open(filename)
	if err != nil {
		fmt.Println("Error opening file", err)
		return
	}
	defer xmlFile.Close()

	b, _ := ioutil.ReadAll(xmlFile)
	err = json.Unmarshal(b, &mock)
	if err != nil {
		fmt.Println("Error unmarshaling Mockerfile", err)
		return
	}

	return
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"os"
	"regexp"
	"sort"
	"time"

	"github.com/Sirupsen/logrus"
	humanize "github.com/dustin/go-humanize"
	"github.com/nlopes/slack"
)

const extractMsgGroupName = "msg"

// extract suffixes of Slack messages starting with @bot-name
var extractMsgPattern string = `(?m)^\s*<@%s>:?\s*(?P<` + extractMsgGroupName + `>.*)$`

// A standupMsg is a Slack message directed to the arriba bot (i.e. with a @botname prefix)
type standupMsg struct {
	ts   time.Time
	text string
}

// channelStandup contains the latest standup message of each user in a Slack channel.
type channelStandup map[string]standupMsg

// sortableChannelStandup is a channelStandup, sortable by the timestamp of its messages
// sortableChannelStandup implements sort.Interface to sort the keys of the channelStandup
type sortableChannelStandup struct {
	keys []string
	cs   channelStandup
}

func (s sortableChannelStandup) Swap(i, j int) { s.keys[i], s.keys[j] = s.keys[j], s.keys[i] }
func (s sortableChannelStandup) Len() int      { return len(s.keys) }
func (s sortableChannelStandup) Less(i, j int) bool {
	return s.cs[s.keys[i]].ts.After(s.cs[s.keys[j]].ts)
}

// getKeysByTimestamp returns the userIDs of the standup ordered by their message timestamp (newer first).
func (cs channelStandup) getKeysByTimestamp() []string {
	keys := make([]string, 0, len(cs))
	for k := range cs {
		keys = append(keys, k)
	}
	scs := sortableChannelStandup{
		cs:   cs,
		keys: keys,
	}
	sort.Sort(scs)
	return scs.keys
}

// standups contains the channelStandup of all Slack channels known to the bot.
type standups map[string]channelStandup

type arriba struct {
	rtm              *slack.RTM
	botID            string
	botName          string
	extractMsgRE     *regexp.Regexp
	historyDaysLimit int
	standups         standups
}

func newArriba(rtm *slack.RTM, historyDaysLimit int) arriba {
	return arriba{
		rtm:              rtm,
		historyDaysLimit: historyDaysLimit,
		standups:         make(standups),
	}
}

func parseSlackTimeStamp(ts string) (time.Time, error) {
	var seconds, milliseconds int64
	_, err := fmt.Sscanf(ts, "%d.%d", &seconds, &milliseconds)
	if err != nil {
		logrus.Warn("Can't parse timestamp ", ts)
		return time.Now(), err
	}
	return time.Unix(seconds, milliseconds*1000), nil
}

// extractStandupMsg parses Slack messages starting with @bot-name
func (a arriba) extractChannelStandupMsg(msg slack.Msg) (standupMsg, bool) {
	if msg.Type != "message" || msg.SubType != "" {
		return standupMsg{}, false
	}
	standupText := a.extractMsgRE.ReplaceAllString(msg.Text, "$"+extractMsgGroupName)
	if len(standupText) == len(msg.Text) {
		// Nothing was extracted
		return standupMsg{}, false
	}
	ts, err := parseSlackTimeStamp(msg.Timestamp)
	if err != nil {
		return standupMsg{}, false
	}
	return standupMsg{ts, standupText}, true
}

func (a arriba) retrieveChannelStandup(channelID string) (channelStandup, error) {
	params := slack.NewHistoryParameters()
	params.Count = 1000
	params.Oldest = fmt.Sprintf(
		"%d",
		time.Now().UTC().AddDate(-a.historyDaysLimit, 0, 0).Unix(),
	)
	// It would be way more efficient to use slack.SearchMsgs instead
	// of traversing the whole history, but that's not allowed for bots :(
	cstandup := make(channelStandup)
	for {
		history, error := a.rtm.GetChannelHistory(
			channelID,
			params,
		)
		if error != nil {
			return cstandup, error
		}

		for _, msg := range history.Messages {
			if _, ok := cstandup[msg.User]; ok {
				// we already have the latest standup message for this user
				continue
			}
			standupMsg, ok := a.extractChannelStandupMsg(msg.Msg)
			if ok && standupMsg.text != "" {
				cstandup[msg.User] = standupMsg
			}
		}

		if !history.HasMore {
			break
		}
		oldestMsg := history.Messages[len(history.Messages)-1]
		params.Oldest = oldestMsg.Timestamp
		params.Inclusive = false
	}
	return cstandup, nil
}

func (a arriba) retrieveStandups(channels []slack.Channel) {
	for _, channel := range channels {
		if channel.IsMember {
			logrus.Infof("Retrieveing standup for channel #%s (%s)", channel.Name, channel.ID)
			cstandup, err := a.retrieveChannelStandup(channel.ID)
			if err != nil {
				logrus.Errorf("Can't retrieve channel standup for channel #%s: %s", channel.Name, err)
			}
			a.standups[channel.ID] = cstandup
			logrus.Infof("Standup for channel #%s (%s) updated to %#v", channel.Name, channel.ID, cstandup)
		}
	}
}

func (a arriba) getUserName(userID string) string {
	info, err := a.rtm.GetUserInfo(userID)
	userName := "id" + userID
	if err != nil {
		logrus.Errorf("Couldn't get user information for user %s: %s", userID, err)
	} else {
		userName = info.Name
	}
	return userName
}

func (a arriba) prettyPrintChannelStandup(cstandup channelStandup) string {
	text := "¡Ándale! ¡Ándale! here's the standup status :tada:\n"
	for _, userID := range cstandup.getKeysByTimestamp() {
		standupMsg := cstandup[userID]
		humanTime := humanize.Time(standupMsg.ts)
		userName := a.getUserName(userID)
		text += fmt.Sprintf("*%s*: %s _(%s)_\n", userName, standupMsg.text, humanTime)
	}
	return text
}

func (a arriba) sendStatus(channelID string) {
	var statusText string
	if cstandup, ok := a.standups[channelID]; ok && len(cstandup) > 0 {
		statusText = a.prettyPrintChannelStandup(cstandup)
	} else {
		statusText = fmt.Sprintf("No standup messages found\nType a message starting with *@%s* to record your standup message", a.botName)
	}
	a.rtm.SendMessage(a.rtm.NewOutgoingMessage(statusText, channelID))

}

func (a arriba) updateLastStandup(channelID, userID string, msg standupMsg) {
	if _, ok := a.standups[channelID]; !ok {
		a.standups[channelID] = make(channelStandup)
	}
	a.standups[channelID][userID] = msg
	confirmationText := fmt.Sprintf("<@%s>: ¡Yeppa! standup status recorded :taco:", userID)
	a.rtm.SendMessage(a.rtm.NewOutgoingMessage(confirmationText, channelID))
}

func (a *arriba) handleConnectedEvent(ev *slack.ConnectedEvent) {
	if a.botID != "" {
		logrus.Warn("Received unexpected Connected event")
		return
	}
	logrus.Infof(
		"Connected as user %s (%s) to team %s (%s)",
		ev.Info.User.Name,
		ev.Info.User.ID,
		ev.Info.Team.Name,
		ev.Info.Team.ID,
	)
	a.botID = ev.Info.User.ID
	a.botName = ev.Info.User.Name
	a.extractMsgRE = regexp.MustCompile(fmt.Sprintf(extractMsgPattern, a.botID))
	a.retrieveStandups(ev.Info.Channels)
}

func (a arriba) handleMessageEvent(ev *slack.MessageEvent) {
	logrus.Debugf("Message received %+v", ev)
	if a.botID == "" {
		logrus.Warn("Received message event before finishing initialization")
		return
	}
	if ev.Channel == "" {
		logrus.Warn("Received message with empty channel")
		return
	}
	switch ev.Channel[0] {
	case 'C':
		smsg, ok := a.extractChannelStandupMsg(ev.Msg)
		if !ok {
			return
		}
		logrus.Infof("Received standup message in channel %s: %+v", ev.Channel, smsg)
		if smsg.text == "" {
			a.sendStatus(ev.Msg.Channel)
		} else {
			a.updateLastStandup(ev.Msg.Channel, ev.Msg.User, smsg)
		}

	case 'D':
		// Direct messages are not supported yet
	case 'G':
		// Group channels are not supported yet
	}
}

func (a arriba) run() {
	go a.rtm.ManageConnection()

	for {
		select {
		case msg := <-a.rtm.IncomingEvents:
			switch ev := msg.Data.(type) {
			case *slack.ConnectedEvent:
				a.handleConnectedEvent(ev)
			case *slack.MessageEvent:
				a.handleMessageEvent(ev)
			case *slack.RTMError:
				logrus.Error("Invalid credentials", ev.Error())
			case *slack.InvalidAuthEvent:
				logrus.Error("Invalid credentials")
				os.Exit(1)
			}
		}
	}
}

func usage() {
	fmt.Fprintf(os.Stderr, "usage: %s [ flags ] <SlackAPItoken>\n", os.Args[0])
	flag.PrintDefaults()
	fmt.Fprintf(os.Stderr, "You can obtain <SlackAPItoken> from https://<yourteam>.slack.com/services/new/bot\n")
}

func main() {
	var (
		debug            bool
		historyDaysLimit int
	)

	flag.Usage = usage
	flag.BoolVar(&debug, "debug", false, "Print debug information")
	flag.IntVar(&historyDaysLimit, "history-limit", 7, "History limit (in days)")
	flag.Parse()
	if len(flag.Args()) < 1 || historyDaysLimit < 1 {
		usage()
		os.Exit(1)
	}

	logrus.SetOutput(os.Stderr)
	if debug {
		logrus.SetLevel(logrus.DebugLevel)
	}

	api := slack.New(flag.Arg(0))
	api.SetDebug(debug)

	newArriba(api.NewRTM(), historyDaysLimit).run()
}
<commit_msg>Add support for private groups<commit_after>package main

import (
	"flag"
	"fmt"
	"os"
	"regexp"
	"sort"
	"time"

	"github.com/Sirupsen/logrus"
	humanize "github.com/dustin/go-humanize"
	"github.com/nlopes/slack"
)

const extractMsgGroupName = "msg"

// extract suffixes of Slack messages starting with @bot-name
var extractMsgPattern string = `(?m)^\s*<@%s>:?\s*(?P<` + extractMsgGroupName + `>.*)$`

// A standupMsg is a Slack message directed to the arriba bot (i.e. with a @botname prefix)
type standupMsg struct {
	ts   time.Time
	text string
}

// channelStandup contains the latest standup message of each user in a Slack channel.
type channelStandup map[string]standupMsg

// sortableChannelStandup is a channelStandup, sortable by the timestamp of its messages
// sortableChannelStandup implements sort.Interface to sort the keys of the channelStandup
type sortableChannelStandup struct {
	keys []string
	cs   channelStandup
}

func (s sortableChannelStandup) Swap(i, j int) { s.keys[i], s.keys[j] = s.keys[j], s.keys[i] }
func (s sortableChannelStandup) Len() int      { return len(s.keys) }
func (s sortableChannelStandup) Less(i, j int) bool {
	return s.cs[s.keys[i]].ts.After(s.cs[s.keys[j]].ts)
}

// getKeysByTimestamp returns the userIDs of the standup ordered by their message timestamp (newer first).
func (cs channelStandup) getKeysByTimestamp() []string {
	keys := make([]string, 0, len(cs))
	for k := range cs {
		keys = append(keys, k)
	}
	scs := sortableChannelStandup{
		cs:   cs,
		keys: keys,
	}
	sort.Sort(scs)
	return scs.keys
}

// standups contains the channelStandup of all Slack channels known to the bot.
type standups map[string]channelStandup

// channel provides a generic way to access the IDs, Names and history of both
// slack.Channel and slack.Group. Unfortunately nlopes/slack doesn't expose the
// underlying common type (groupConversation) and we cannot define methods for
// non-local types, which would allow to make things much cleaner ...
type channel struct {
	id, name   string
	getHistory func(*slack.RTM, slack.HistoryParameters) (*slack.History, error)
}

type arriba struct {
	rtm              *slack.RTM
	botID            string
	botName          string
	extractMsgRE     *regexp.Regexp
	historyDaysLimit int
	standups         standups
}

func newArriba(rtm *slack.RTM, historyDaysLimit int) arriba {
	return arriba{
		rtm:              rtm,
		historyDaysLimit: historyDaysLimit,
		standups:         make(standups),
	}
}

func parseSlackTimeStamp(ts string) (time.Time, error) {
	var seconds, milliseconds int64
	_, err := fmt.Sscanf(ts, "%d.%d", &seconds, &milliseconds)
	if err != nil {
		logrus.Warn("Can't parse timestamp ", ts)
		return time.Now(), err
	}
	return time.Unix(seconds, milliseconds*1000), nil
}

// extractStandupMsg parses Slack messages starting with @bot-name
func (a arriba) extractChannelStandupMsg(msg slack.Msg) (standupMsg, bool) {
	if msg.Type != "message" || msg.SubType != "" {
		return standupMsg{}, false
	}
	standupText := a.extractMsgRE.ReplaceAllString(msg.Text, "$"+extractMsgGroupName)
	if len(standupText) == len(msg.Text) {
		// Nothing was extracted
		return standupMsg{}, false
	}
	ts, err := parseSlackTimeStamp(msg.Timestamp)
	if err != nil {
		return standupMsg{}, false
	}
	return standupMsg{ts, standupText}, true
}

func (a arriba) retrieveChannelStandup(c channel) (channelStandup, error) {
	params := slack.NewHistoryParameters()
	params.Count = 1000
	params.Oldest = fmt.Sprintf(
		"%d",
		time.Now().UTC().AddDate(-a.historyDaysLimit, 0, 0).Unix(),
	)
	// It would be way more efficient to use slack.SearchMsgs instead
	// of traversing the whole history, but that's not allowed for bots :(
	cstandup := make(channelStandup)
	for {
		history, error := c.getHistory(a.rtm, params)
		if error != nil {
			return cstandup, error
		}

		for _, msg := range history.Messages {
			if _, ok := cstandup[msg.User]; ok {
				// we already have the latest standup message for this user
				continue
			}
			standupMsg, ok := a.extractChannelStandupMsg(msg.Msg)
			if ok && standupMsg.text != "" {
				cstandup[msg.User] = standupMsg
			}
		}

		if !history.HasMore {
			break
		}
		oldestMsg := history.Messages[len(history.Messages)-1]
		params.Oldest = oldestMsg.Timestamp
		params.Inclusive = false
	}
	return cstandup, nil
}

func (a arriba) retrieveStandups(channels []channel) {
	for _, channel := range channels {
		logrus.Infof("Retrieveing standup for channel #%s (%s)", channel.name, channel.id)
		cstandup, err := a.retrieveChannelStandup(channel)
		if err != nil {
			logrus.Errorf("Can't retrieve channel standup for channel #%s: %s", channel.name, err)
		}
		a.standups[channel.id] = cstandup
		logrus.Infof("Standup for channel #%s (%s) updated to %#v", channel.name, channel.id, cstandup)
	}
}

func (a arriba) getUserName(userID string) string {
	info, err := a.rtm.GetUserInfo(userID)
	userName := "id" + userID
	if err != nil {
		logrus.Errorf("Couldn't get user information for user %s: %s", userID, err)
	} else {
		userName = info.Name
	}
	return userName
}

func (a arriba) prettyPrintChannelStandup(cstandup channelStandup) string {
	text := "¡Ándale! ¡Ándale! here's the standup status :tada:\n"
	for _, userID := range cstandup.getKeysByTimestamp() {
		standupMsg := cstandup[userID]
		humanTime := humanize.Time(standupMsg.ts)
		userName := a.getUserName(userID)
		text += fmt.Sprintf("*%s*: %s _(%s)_\n", userName, standupMsg.text, humanTime)
	}
	return text
}

func (a arriba) sendStatus(channelID string) {
	var statusText string
	if cstandup, ok := a.standups[channelID]; ok && len(cstandup) > 0 {
		statusText = a.prettyPrintChannelStandup(cstandup)
	} else {
		statusText = fmt.Sprintf("No standup messages found\nType a message starting with *@%s* to record your standup message", a.botName)
	}
	a.rtm.SendMessage(a.rtm.NewOutgoingMessage(statusText, channelID))

}

func (a arriba) updateLastStandup(channelID, userID string, msg standupMsg) {
	if _, ok := a.standups[channelID]; !ok {
		a.standups[channelID] = make(channelStandup)
	}
	a.standups[channelID][userID] = msg
	confirmationText := fmt.Sprintf("<@%s>: ¡Yeppa! standup status recorded :taco:", userID)
	a.rtm.SendMessage(a.rtm.NewOutgoingMessage(confirmationText, channelID))
}

func (a *arriba) handleConnectedEvent(ev *slack.ConnectedEvent) {
	if a.botID != "" {
		logrus.Warn("Received unexpected Connected event")
		return
	}
	logrus.Infof(
		"Connected as user %s (%s) to team %s (%s)",
		ev.Info.User.Name,
		ev.Info.User.ID,
		ev.Info.Team.Name,
		ev.Info.Team.ID,
	)
	a.botID = ev.Info.User.ID
	a.botName = ev.Info.User.Name
	a.extractMsgRE = regexp.MustCompile(fmt.Sprintf(extractMsgPattern, a.botID))

	// Retrieve standups for public channels and private groups
	var channels []channel
	for _, c := range ev.Info.Channels {
		if c.IsMember {
			channel := channel{
				id:   c.ID,
				name: c.Name,
				getHistory: func(r *slack.RTM, params slack.HistoryParameters) (*slack.History, error) {
					return r.GetChannelHistory(c.ID, params)
				},
			}
			channels = append(channels, channel)
		}
	}
	for _, g := range ev.Info.Groups {
		channel := channel{
			id:   g.ID,
			name: g.Name,
			getHistory: func(r *slack.RTM, params slack.HistoryParameters) (*slack.History, error) {
				return r.GetGroupHistory(g.ID, params)
			},
		}
		channels = append(channels, channel)
	}
	a.retrieveStandups(channels)
}

func (a arriba) handleMessageEvent(ev *slack.MessageEvent) {
	logrus.Debugf("Message received %+v", ev)
	if a.botID == "" {
		logrus.Warn("Received message event before finishing initialization")
		return
	}
	if ev.Channel == "" {
		logrus.Warn("Received message with empty channel")
		return
	}
	switch ev.Channel[0] {
	case 'C', 'G':
		// Public and private (group) channels
		smsg, ok := a.extractChannelStandupMsg(ev.Msg)
		if !ok {
			return
		}
		logrus.Infof("Received standup message in channel %s: %+v", ev.Channel, smsg)
		if smsg.text == "" {
			a.sendStatus(ev.Msg.Channel)
		} else {
			a.updateLastStandup(ev.Msg.Channel, ev.Msg.User, smsg)
		}

	case 'D':
		// Direct messages are not supported yet
	}
}

func (a arriba) run() {
	go a.rtm.ManageConnection()

	for {
		select {
		case msg := <-a.rtm.IncomingEvents:
			switch ev := msg.Data.(type) {
			case *slack.ConnectedEvent:
				a.handleConnectedEvent(ev)
			case *slack.MessageEvent:
				a.handleMessageEvent(ev)
			case *slack.RTMError:
				logrus.Error("Invalid credentials", ev.Error())
			case *slack.InvalidAuthEvent:
				logrus.Error("Invalid credentials")
				os.Exit(1)
			}
		}
	}
}

func usage() {
	fmt.Fprintf(os.Stderr, "usage: %s [ flags ] <SlackAPItoken>\n", os.Args[0])
	flag.PrintDefaults()
	fmt.Fprintf(os.Stderr, "You can obtain <SlackAPItoken> from https://<yourteam>.slack.com/services/new/bot\n")
}

func main() {
	var (
		debug            bool
		historyDaysLimit int
	)

	flag.Usage = usage
	flag.BoolVar(&debug, "debug", false, "Print debug information")
	flag.IntVar(&historyDaysLimit, "history-limit", 7, "History limit (in days)")
	flag.Parse()
	if len(flag.Args()) < 1 || historyDaysLimit < 1 {
		usage()
		os.Exit(1)
	}

	logrus.SetOutput(os.Stderr)
	if debug {
		logrus.SetLevel(logrus.DebugLevel)
	}

	api := slack.New(flag.Arg(0))
	api.SetDebug(debug)

	newArriba(api.NewRTM(), historyDaysLimit).run()
}
<|endoftext|>


<commit_before>// Copyright 2016 carsonsx. All rights reserved.

package main

import (
	"bytes"
	"flag"
	"io"
	"log"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"time"
	"strings"
	"encoding/json"
	"fmt"
)

const (
	DEFAULT_PORT = "8011"
	FORM_FILE = "file"
	STORAGE_PATH = "files"
	DATE_FORMAT = "20060102"

	STATUS_SUCCESS = 0
	STATUS_ERROR = 1
)

var (
	port = DEFAULT_PORT
	upload_dir = STORAGE_PATH
)

type ResponseData struct {
	Status  int `json:"status"`
	Message string `json:"message"`
}

func logRequest(r *http.Request) {
	log.Printf("Request URI: %s", r.RequestURI)
}

func writeSuccess(w http.ResponseWriter, message string) {
	write(w, STATUS_SUCCESS, message)
}

func writeError(w http.ResponseWriter, err string) {
	write(w, STATUS_ERROR, err)
}

func write(w http.ResponseWriter, status int, message string) {
	data := ResponseData{status, message}
	b, err := json.Marshal(data)
	if err != nil {
		io.WriteString(w, err.Error())
		return
	}
	w.Write(b)
}

func welcome(w http.ResponseWriter, r *http.Request) {
	io.WriteString(w, "Welcome to HTTP File Server")
}

func upload(w http.ResponseWriter, r *http.Request) {
	logRequest(r)
	dir := strings.TrimPrefix(r.RequestURI, "/upload")
	dir = strings.TrimPrefix(dir, "/")
	byDate := r.FormValue("bydate")
	if byDate == "true" {
		dir = filepath.Join(time.Now().Format(DATE_FORMAT), dir)
	}
	storeDir := filepath.Join(upload_dir, dir)
	fp, err := extractFile(r, storeDir)
	if err != nil {
		writeError(w, err.Error())
	} else {
		writeSuccess(w, "uploaded " + filepath.Base(fp) + "")
	}
}

func remove(w http.ResponseWriter, r *http.Request) {
	logRequest(r)
	path := strings.TrimPrefix(r.RequestURI, "/remove")
	path = strings.TrimPrefix(path, "/")
	if path == "" {
		log.Println("path is required")
		writeError(w, "path is required")
		return
	}
	err := os.RemoveAll(filepath.Join(upload_dir, path))
	if err != nil {
		log.Println(err)
		writeError(w, err.Error())
	} else {
		output := "removed " + path
		log.Println(output)
		writeSuccess(w, output)
	}
}

func cmd(w http.ResponseWriter, r *http.Request) {
	logRequest(r)
	cmd := r.FormValue("cmd")
	args := r.Form["args"]
	pCmd := exec.Command(cmd, args...)
	var out bytes.Buffer
	pCmd.Stdout = &out
	err := pCmd.Run()
	if err != nil {
		writeError(w, err.Error())
	} else {
		writeSuccess(w, out.String())
	}
}

func main() {

	var showVersion bool

	p := flag.String("p", port, "port to serve on")
	d := flag.String("d", upload_dir, "the directory of static file to host")
	flag.BoolVar(&showVersion, "version", false, "Print version information.")
	flag.BoolVar(&showVersion, "v", false, "Print version information.")
	flag.Parse()

	if showVersion {
		fmt.Println(GetHumanVersion())
		return
	}

	port = *p
	upload_dir = *d
	http.HandleFunc("/upload", upload)
	http.HandleFunc("/upload/", upload)
	http.HandleFunc("/remove/", remove)
	http.Handle("/download/", http.StripPrefix("/download/", http.FileServer(http.Dir(upload_dir))))
	http.HandleFunc("/execute", cmd)
	http.HandleFunc("/", welcome)
	log.Println(GitCommit)
	log.Println(GitDescribe)
	log.Printf("Started HTTP File Server on port: %s\n", port)
	log.Fatalln(http.ListenAndServe(":" + port, nil))
}
<commit_msg>update version output<commit_after>// Copyright 2016 carsonsx. All rights reserved.

package main

import (
	"bytes"
	"flag"
	"io"
	"log"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"time"
	"strings"
	"encoding/json"
	"fmt"
)

const (
	DEFAULT_PORT = "8011"
	FORM_FILE = "file"
	STORAGE_PATH = "files"
	DATE_FORMAT = "20060102"

	STATUS_SUCCESS = 0
	STATUS_ERROR = 1
)

var (
	port = DEFAULT_PORT
	upload_dir = STORAGE_PATH
)

type ResponseData struct {
	Status  int `json:"status"`
	Message string `json:"message"`
}

func logRequest(r *http.Request) {
	log.Printf("Request URI: %s", r.RequestURI)
}

func writeSuccess(w http.ResponseWriter, message string) {
	write(w, STATUS_SUCCESS, message)
}

func writeError(w http.ResponseWriter, err string) {
	write(w, STATUS_ERROR, err)
}

func write(w http.ResponseWriter, status int, message string) {
	data := ResponseData{status, message}
	b, err := json.Marshal(data)
	if err != nil {
		io.WriteString(w, err.Error())
		return
	}
	w.Write(b)
}

func welcome(w http.ResponseWriter, r *http.Request) {
	io.WriteString(w, "Welcome to HTTP File Server")
}

func upload(w http.ResponseWriter, r *http.Request) {
	logRequest(r)
	dir := strings.TrimPrefix(r.RequestURI, "/upload")
	dir = strings.TrimPrefix(dir, "/")
	byDate := r.FormValue("bydate")
	if byDate == "true" {
		dir = filepath.Join(time.Now().Format(DATE_FORMAT), dir)
	}
	storeDir := filepath.Join(upload_dir, dir)
	fp, err := extractFile(r, storeDir)
	if err != nil {
		writeError(w, err.Error())
	} else {
		writeSuccess(w, "uploaded " + filepath.Base(fp) + "")
	}
}

func remove(w http.ResponseWriter, r *http.Request) {
	logRequest(r)
	path := strings.TrimPrefix(r.RequestURI, "/remove")
	path = strings.TrimPrefix(path, "/")
	if path == "" {
		log.Println("path is required")
		writeError(w, "path is required")
		return
	}
	err := os.RemoveAll(filepath.Join(upload_dir, path))
	if err != nil {
		log.Println(err)
		writeError(w, err.Error())
	} else {
		output := "removed " + path
		log.Println(output)
		writeSuccess(w, output)
	}
}

func cmd(w http.ResponseWriter, r *http.Request) {
	logRequest(r)
	cmd := r.FormValue("cmd")
	args := r.Form["args"]
	pCmd := exec.Command(cmd, args...)
	var out bytes.Buffer
	pCmd.Stdout = &out
	err := pCmd.Run()
	if err != nil {
		writeError(w, err.Error())
	} else {
		writeSuccess(w, out.String())
	}
}

func main() {

	var showVersion bool

	p := flag.String("p", port, "port to serve on")
	d := flag.String("d", upload_dir, "the directory of static file to host")
	flag.BoolVar(&showVersion, "version", false, "Print version information.")
	flag.BoolVar(&showVersion, "v", false, "Print version information.")
	flag.Parse()

	if showVersion {
		fmt.Printf("HFS %s", GetHumanVersion())
		return
	}

	port = *p
	upload_dir = *d
	http.HandleFunc("/upload", upload)
	http.HandleFunc("/upload/", upload)
	http.HandleFunc("/remove/", remove)
	http.Handle("/download/", http.StripPrefix("/download/", http.FileServer(http.Dir(upload_dir))))
	http.HandleFunc("/execute", cmd)
	http.HandleFunc("/", welcome)
	log.Println(GitCommit)
	log.Println(GitDescribe)
	log.Printf("Started HTTP File Server on port: %s\n", port)
	log.Fatalln(http.ListenAndServe(":" + port, nil))
}
<|endoftext|>


<commit_before>package main

import (
	"code.google.com/p/go-uuid/uuid"
	"encoding/json"
	"fmt"
	"github.com/gorilla/mux"
	gorillasession "github.com/gorilla/sessions"
	"github.com/ory-am/common/env"
	"github.com/ory-am/common/mgopath"
	"github.com/ory-am/event"
	gde "github.com/ory-am/gitdeploy/event"
	"github.com/ory-am/gitdeploy/job"
	gdLog "github.com/ory-am/gitdeploy/log"
	"github.com/ory-am/gitdeploy/sse"
	"github.com/ory-am/gitdeploy/storage"
	"github.com/ory-am/gitdeploy/storage/mongo"
	"github.com/ory-am/google-json-style-response/responder"
	"gopkg.in/mgo.v2"
	"log"
	"net/http"
	"os/exec"
	"regexp"
	"time"
	"runtime"
	"os"
    "strings"
)

const (
	sessionCurrentDeployment = "cdid"
	sessionName = "gdp"
)

var (
// API Version
	ApiVersion = "1.0"

// Generic configuration
	host = env.Getenv("HOST", "")
	port = env.Getenv("PORT", "7654")

	envAppTtl = env.Getenv("APP_TTL", "30m")
    envClusterConf = env.Getenv("FLYNN_CLUSTER_CONFIG", "")

// Configuration for CORS
	corsAllowOrigin = env.Getenv("CORS_ALLOW_ORIGIN", "http://localhost:9000")

	sessionStore = gorillasession.NewCookieStore([]byte(env.Getenv("SESSION_SECRET", "changme")))

// MongoDB
	envMongoPath = env.Getenv("MONGODB", "mongodb://localhost:27017/gitdeploy")
)

type deployRequest struct {
	Repository string `json:"repository"`
}

type appResponse struct {
	App string `json:"app"`
}

func main() {
	checkDependencies()
	eventManager := event.New()

	// mgo
	mongoDatabase, err := mgopath.Connect(envMongoPath)
	if err != nil {
		log.Fatal(err)
	}

	storage := mongo.New(mongoDatabase)
	eventManager.AttachListenerAggregate(storage)

	// SSE broker
	sseBroker := sse.New(storage)

	// Log listener
	eventManager.AttachListenerAggregate(new(gdLog.Listener))

	// Mux router
	r := mux.NewRouter()
	r.HandleFunc("/config", configHandler).Methods("GET")
	r.HandleFunc("/deployments", deployWrapperAction(sseBroker, eventManager, storage)).Methods("POST")
	r.HandleFunc("/deployments", setCORSHeaders).Methods("OPTIONS")
	r.HandleFunc("/deployments/{app:.+}/events", eventWrapperAction(sseBroker)).Methods("GET")
	r.HandleFunc("/apps/{app:.+}", getAppHandler(storage)).Methods("GET")
	r.PathPrefix("/").HandlerFunc(publicHandler("./app/dist"))
	http.Handle("/", r)

	go job.KillAppsOnHitList(storage)

	listen := fmt.Sprintf("%s:%s", host, port)
	log.Printf("Listening on %s", listen)
	log.Fatal(http.ListenAndServe(listen, nil))
}

func publicHandler(dir string) func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		path := dir + r.URL.Path
		if r.URL.Path == "/" {
			path = dir + "/index.html"
		}
		pattern := `!\.html|\.js|\.svg|\.css|\.png|\.jpg$`

		if f, err := os.Stat(path); err == nil {
			if !f.IsDir() {
				http.ServeFile(w, r, path)
				return
			} else {
                http.NotFound(w, r)
				return
			}
		}

		if matched, err := regexp.MatchString(pattern, path); err != nil {
			log.Printf("Could not exec regex: %s", err.Error())
		} else if !matched {
            http.ServeFile(w, r, dir + "/index.html")
            return
		} else {
			http.NotFound(w, r)
		}
	}
}

func configHandler(w http.ResponseWriter, r *http.Request) {
	setCORSHeaders(w, r)
	responseSuccess(w, struct {
		Now time.Time `json:"now"`
	}{
		time.Now(),
	})
}

func getAppHandler(store *mongo.MongoStorage) func(w http.ResponseWriter, r *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		if cleanUpSession(w, r) != nil {
			return
		}
		setCORSHeaders(w, r)
		vars := mux.Vars(r)
		id := vars["app"]
		app, err := store.GetApp(id)
		if err == mgo.ErrNotFound {
			responseError(w, http.StatusNotFound, fmt.Sprintf("App %s does not exist.", id))
			return
		} else if app.Killed {
			responseError(w, http.StatusNotFound, "App timed out and is no longer available.")
			return
		} else if err != nil {
			responseError(w, http.StatusInternalServerError, err.Error())
			return
		}
		logs, err := job.GetLogs(app.ID)
		if err != nil {
			responseError(w, http.StatusInternalServerError, err.Error())
			return
		}

		ps, err := job.GetPS(app.ID)
		if err != nil {
			responseError(w, http.StatusInternalServerError, err.Error())
			return
		}

		deployLogs, err := store.FindDeployLogsForApp(id)
		if err != nil {
			responseError(w, http.StatusInternalServerError, err.Error())
			return
		}

		responseSuccess(w, struct {
			*storage.App
			Logs       string                 `json:"logs"`
			DeployLogs []*storage.DeployEvent `json:"deployLogs"`
			PS         string                 `json:"ps"`
		}{app, logs, deployLogs, ps})
	}
}

func cleanUpSession(w http.ResponseWriter, r *http.Request) error {
	session, err := sessionStore.Get(r, sessionName)
	if err != nil {
		responseError(w, http.StatusBadRequest, "Please delete your cookies.")
		return err
	}
	delete(session.Values, sessionCurrentDeployment)
	if err := session.Save(r, w); err != nil {
		responseError(w, http.StatusInternalServerError, err.Error())
		return err
	}
	return nil
}

func eventWrapperAction(sseBroker *sse.Broker) func(w http.ResponseWriter, r *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		setCORSHeaders(w, r)
		sseBroker.EventHandler(w, r)
	}
}

func deployWrapperAction(sseBroker *sse.Broker, em *event.EventManager, store *mongo.MongoStorage) func(w http.ResponseWriter, r *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		setCORSHeaders(w, r)
		deployAction(w, r, sseBroker, em, store)
	}
}

func deployAction(w http.ResponseWriter, r *http.Request, sseBroker *sse.Broker, em *event.EventManager, store *mongo.MongoStorage) {
	session, err := sessionStore.Get(r, sessionName)
	if err != nil {
		responseError(w, http.StatusBadRequest, "Please delete your cookies.")
		return
	}

	// Check if the user is currently deploying an application and switch to that one.
	if v, ok := session.Values[sessionCurrentDeployment].(string); ok && len(v) > 0 {
		app, err := store.GetApp(v)
		if err != nil {
			cleanUpSession(w, r)
			log.Printf("Could not fetch app from cookie: %s", err.Error())
		} else {
			responseSuccess(w, app)
		}
	}

	// Parse body
	dr := new(deployRequest)
	decoder := json.NewDecoder(r.Body)
	decoder.Decode(dr)

	// Validate URL
	regExpression := "(https|http):\\/\\/github\\.com\\/[a-zA-Z0-9\\-\\_]+/[a-zA-Z0-9\\-\\_]+\\.git"
	if match, _ := regexp.MatchString(regExpression, dr.Repository); !match {
		responseError(w, http.StatusBadRequest, "I only support GitHub.")
		return
	}

	app := uuid.NewRandom().String()
	if ttl, err := time.ParseDuration(envAppTtl); err != nil {
		log.Println(err.Error())
		responseError(w, http.StatusInternalServerError, err.Error())
		return
	} else {
		appEntity, err := store.AddApp(app, time.Now().Add(ttl), dr.Repository, getIP(r))
		if err != nil {
			responseError(w, http.StatusInternalServerError, err.Error())
			return
		}
		session.Values[sessionCurrentDeployment] = appEntity.ID
		if err := session.Save(r, w); err != nil {
			responseError(w, http.StatusInternalServerError, err.Error())
			return
		}
		go runJobs(w, r, em, dr, appEntity, sseBroker, session, store)
		responseSuccess(w, appEntity)
	}
}

func runJobs(w http.ResponseWriter, r *http.Request, em *event.EventManager, dr *deployRequest, app *storage.App, sseBroker *sse.Broker, session *gorillasession.Session, store *mongo.MongoStorage) {
	sseBroker.OpenChannel(app.ID)
	sseBroker.Start(app.ID)
	defer func() {
		// Give the client the chance to read the output...
		time.Sleep(2 * time.Minute)
		sseBroker.CloseChannel(app.ID)
	}()

	em.Trigger("app.created", gde.New(app.ID, app.ID))

	destination, err := job.Clone(em, app.ID, dr.Repository)
	if err != nil {
		log.Printf("Error in job.clone %s: %s", app.ID, err.Error())
		return
	}

	if err = job.Parse(em, app.ID, destination); err != nil {
		log.Printf("Error in job.parse %s: %s", app.ID, err.Error())
		return
	}

	if err = job.Deploy(em, app.ID, destination); err != nil {
		log.Printf("Error in job.deploy %s: %s", app.ID, err.Error())
		return
	}

	cluster, err := job.GetCluster(em, app.ID)
	if err != nil {
		log.Printf("Error in job.deploy %s: %s", app.ID, err.Error())
		return
	}

	app.URL = fmt.Sprintf("%s.%s", app.ID, cluster.Host)
	if err := store.UpdateApp(app); err != nil {
		responseError(w, http.StatusInternalServerError, err.Error())
		log.Printf("Error %s: %s", app.ID, err.Error())
		return
	}

	log.Println("Deployment successful.")
	em.Trigger("app.deployed", gde.New(app.ID, app.URL))
}

// Set the different CORS headers required for CORS request
func setCORSHeaders(w http.ResponseWriter, r *http.Request) {
	w.Header().Add("Access-Control-Allow-Methods", "POST, GET, OPTIONS")
	w.Header().Set("Access-Control-Allow-Origin", corsAllowOrigin)
	w.Header().Set("Access-Control-Allow-Credentials", "true")
	w.Header().Set("Access-Control-Allow-Headers", "Origin, X-Requested-With, Content-Type, Accept")
}

func responseError(w http.ResponseWriter, code int, message string) {
	log.Printf("Error %d: %s", code, message)
	response := responder.New(ApiVersion)
	response.Write(w, response.Error(code, message))
}

func responseSuccess(w http.ResponseWriter, data interface{}) {
	response := responder.New(ApiVersion)
	response.Write(w, response.Success(data))
}

func checkDependencies() {
	go func() {
		_, err := exec.LookPath("git")
		if err != nil {
			log.Fatal("Git CLI is required but not installed or not in path.")
		}
	}()
	go checkIfFlynnExists()
}

func checkIfFlynnExists() {
	func() {
		_, err := exec.LookPath("flynn")
		if err != nil {
			if runtime.GOOS == "windows" {
				log.Fatal("Flynn CLI is required but not installed or not in path.")
			}
			log.Println("Could not find Flynn CLI, trying to install...")
			if o, err := exec.Command("sh", "bin/flynn-install.sh").CombinedOutput(); err != nil {
				log.Printf("Could not install Flynn CLI (%s): %s", err.Error(), o)
			} else if _, err := exec.LookPath("flynn"); err != nil {
				log.Fatal("Could not install Flynn CLI.")
			}
            log.Println("Flynn installed successfully!")
            log.Println("Adding flynn cluster...")
            log.Println(envClusterConf)
            args := append([]string{"cluster", "add"}, strings.Split(envClusterConf, " ")...)
            log.Printf("%s", args)
            if o, err := exec.Command("flynn", args...).CombinedOutput(); err != nil {
                log.Fatalf("Could not add cluster (status: %s) (output: %s) (args: %s)", err.Error(), o, args)
            }
		}
	}()
}

func getIP(r *http.Request) string {
    ip := removePort(r.RemoteAddr)
    if len(r.Header.Get("X-FORWARDED-FOR")) > 0 {
        ip = r.Header.Get("X-FORWARDED-FOR")
    }
    return ip
}

func removePort(ip string) string {
    split := strings.Split(ip, ":")
    if len(split) < 2 {
        return ip
    }
    return strings.Join(split[:len(split)-1], ":")
}
<commit_msg>Cluster adding improved<commit_after>package main

import (
	"code.google.com/p/go-uuid/uuid"
	"encoding/json"
	"fmt"
	"github.com/gorilla/mux"
	gorillasession "github.com/gorilla/sessions"
	"github.com/ory-am/common/env"
	"github.com/ory-am/common/mgopath"
	"github.com/ory-am/event"
	gde "github.com/ory-am/gitdeploy/event"
	"github.com/ory-am/gitdeploy/job"
	gdLog "github.com/ory-am/gitdeploy/log"
	"github.com/ory-am/gitdeploy/sse"
	"github.com/ory-am/gitdeploy/storage"
	"github.com/ory-am/gitdeploy/storage/mongo"
	"github.com/ory-am/google-json-style-response/responder"
	"gopkg.in/mgo.v2"
	"log"
	"net/http"
	"os/exec"
	"regexp"
	"time"
	"runtime"
	"os"
    "strings"
)

const (
	sessionCurrentDeployment = "cdid"
	sessionName = "gdp"
)

var (
// API Version
	ApiVersion = "1.0"

// Generic configuration
	host = env.Getenv("HOST", "")
	port = env.Getenv("PORT", "7654")

	envAppTtl = env.Getenv("APP_TTL", "30m")
    envClusterConf = env.Getenv("FLYNN_CLUSTER_CONFIG", "")

// Configuration for CORS
	corsAllowOrigin = env.Getenv("CORS_ALLOW_ORIGIN", "http://localhost:9000")

	sessionStore = gorillasession.NewCookieStore([]byte(env.Getenv("SESSION_SECRET", "changme")))

// MongoDB
	envMongoPath = env.Getenv("MONGODB", "mongodb://localhost:27017/gitdeploy")
)

type deployRequest struct {
	Repository string `json:"repository"`
}

type appResponse struct {
	App string `json:"app"`
}

func main() {
	checkDependencies()
	eventManager := event.New()

	// mgo
	mongoDatabase, err := mgopath.Connect(envMongoPath)
	if err != nil {
		log.Fatal(err)
	}

	storage := mongo.New(mongoDatabase)
	eventManager.AttachListenerAggregate(storage)

	// SSE broker
	sseBroker := sse.New(storage)

	// Log listener
	eventManager.AttachListenerAggregate(new(gdLog.Listener))

	// Mux router
	r := mux.NewRouter()
	r.HandleFunc("/config", configHandler).Methods("GET")
	r.HandleFunc("/deployments", deployWrapperAction(sseBroker, eventManager, storage)).Methods("POST")
	r.HandleFunc("/deployments", setCORSHeaders).Methods("OPTIONS")
	r.HandleFunc("/deployments/{app:.+}/events", eventWrapperAction(sseBroker)).Methods("GET")
	r.HandleFunc("/apps/{app:.+}", getAppHandler(storage)).Methods("GET")
	r.PathPrefix("/").HandlerFunc(publicHandler("./app/dist"))
	http.Handle("/", r)

	go job.KillAppsOnHitList(storage)

	listen := fmt.Sprintf("%s:%s", host, port)
	log.Printf("Listening on %s", listen)
	log.Fatal(http.ListenAndServe(listen, nil))
}

func publicHandler(dir string) func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		path := dir + r.URL.Path
		if r.URL.Path == "/" {
			path = dir + "/index.html"
		}
		pattern := `!\.html|\.js|\.svg|\.css|\.png|\.jpg$`

		if f, err := os.Stat(path); err == nil {
			if !f.IsDir() {
				http.ServeFile(w, r, path)
				return
			} else {
                http.NotFound(w, r)
				return
			}
		}

		if matched, err := regexp.MatchString(pattern, path); err != nil {
			log.Printf("Could not exec regex: %s", err.Error())
		} else if !matched {
            http.ServeFile(w, r, dir + "/index.html")
            return
		} else {
			http.NotFound(w, r)
		}
	}
}

func configHandler(w http.ResponseWriter, r *http.Request) {
	setCORSHeaders(w, r)
	responseSuccess(w, struct {
		Now time.Time `json:"now"`
	}{
		time.Now(),
	})
}

func getAppHandler(store *mongo.MongoStorage) func(w http.ResponseWriter, r *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		if cleanUpSession(w, r) != nil {
			return
		}
		setCORSHeaders(w, r)
		vars := mux.Vars(r)
		id := vars["app"]
		app, err := store.GetApp(id)
		if err == mgo.ErrNotFound {
			responseError(w, http.StatusNotFound, fmt.Sprintf("App %s does not exist.", id))
			return
		} else if app.Killed {
			responseError(w, http.StatusNotFound, "App timed out and is no longer available.")
			return
		} else if err != nil {
			responseError(w, http.StatusInternalServerError, err.Error())
			return
		}
		logs, err := job.GetLogs(app.ID)
		if err != nil {
			responseError(w, http.StatusInternalServerError, err.Error())
			return
		}

		ps, err := job.GetPS(app.ID)
		if err != nil {
			responseError(w, http.StatusInternalServerError, err.Error())
			return
		}

		deployLogs, err := store.FindDeployLogsForApp(id)
		if err != nil {
			responseError(w, http.StatusInternalServerError, err.Error())
			return
		}

		responseSuccess(w, struct {
			*storage.App
			Logs       string                 `json:"logs"`
			DeployLogs []*storage.DeployEvent `json:"deployLogs"`
			PS         string                 `json:"ps"`
		}{app, logs, deployLogs, ps})
	}
}

func cleanUpSession(w http.ResponseWriter, r *http.Request) error {
	session, err := sessionStore.Get(r, sessionName)
	if err != nil {
		responseError(w, http.StatusBadRequest, "Please delete your cookies.")
		return err
	}
	delete(session.Values, sessionCurrentDeployment)
	if err := session.Save(r, w); err != nil {
		responseError(w, http.StatusInternalServerError, err.Error())
		return err
	}
	return nil
}

func eventWrapperAction(sseBroker *sse.Broker) func(w http.ResponseWriter, r *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		setCORSHeaders(w, r)
		sseBroker.EventHandler(w, r)
	}
}

func deployWrapperAction(sseBroker *sse.Broker, em *event.EventManager, store *mongo.MongoStorage) func(w http.ResponseWriter, r *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		setCORSHeaders(w, r)
		deployAction(w, r, sseBroker, em, store)
	}
}

func deployAction(w http.ResponseWriter, r *http.Request, sseBroker *sse.Broker, em *event.EventManager, store *mongo.MongoStorage) {
	session, err := sessionStore.Get(r, sessionName)
	if err != nil {
		responseError(w, http.StatusBadRequest, "Please delete your cookies.")
		return
	}

	// Check if the user is currently deploying an application and switch to that one.
	if v, ok := session.Values[sessionCurrentDeployment].(string); ok && len(v) > 0 {
		app, err := store.GetApp(v)
		if err != nil {
			cleanUpSession(w, r)
			log.Printf("Could not fetch app from cookie: %s", err.Error())
		} else {
			responseSuccess(w, app)
		}
	}

	// Parse body
	dr := new(deployRequest)
	decoder := json.NewDecoder(r.Body)
	decoder.Decode(dr)

	// Validate URL
	regExpression := "(https|http):\\/\\/github\\.com\\/[a-zA-Z0-9\\-\\_]+/[a-zA-Z0-9\\-\\_]+\\.git"
	if match, _ := regexp.MatchString(regExpression, dr.Repository); !match {
		responseError(w, http.StatusBadRequest, "I only support GitHub.")
		return
	}

	app := uuid.NewRandom().String()
	if ttl, err := time.ParseDuration(envAppTtl); err != nil {
		log.Println(err.Error())
		responseError(w, http.StatusInternalServerError, err.Error())
		return
	} else {
		appEntity, err := store.AddApp(app, time.Now().Add(ttl), dr.Repository, getIP(r))
		if err != nil {
			responseError(w, http.StatusInternalServerError, err.Error())
			return
		}
		session.Values[sessionCurrentDeployment] = appEntity.ID
		if err := session.Save(r, w); err != nil {
			responseError(w, http.StatusInternalServerError, err.Error())
			return
		}
		go runJobs(w, r, em, dr, appEntity, sseBroker, session, store)
		responseSuccess(w, appEntity)
	}
}

func runJobs(w http.ResponseWriter, r *http.Request, em *event.EventManager, dr *deployRequest, app *storage.App, sseBroker *sse.Broker, session *gorillasession.Session, store *mongo.MongoStorage) {
	sseBroker.OpenChannel(app.ID)
	sseBroker.Start(app.ID)
	defer func() {
		// Give the client the chance to read the output...
		time.Sleep(2 * time.Minute)
		sseBroker.CloseChannel(app.ID)
	}()

	em.Trigger("app.created", gde.New(app.ID, app.ID))

	destination, err := job.Clone(em, app.ID, dr.Repository)
	if err != nil {
		log.Printf("Error in job.clone %s: %s", app.ID, err.Error())
		return
	}

	if err = job.Parse(em, app.ID, destination); err != nil {
		log.Printf("Error in job.parse %s: %s", app.ID, err.Error())
		return
	}

	if err = job.Deploy(em, app.ID, destination); err != nil {
		log.Printf("Error in job.deploy %s: %s", app.ID, err.Error())
		return
	}

	cluster, err := job.GetCluster(em, app.ID)
	if err != nil {
		log.Printf("Error in job.deploy %s: %s", app.ID, err.Error())
		return
	}

	app.URL = fmt.Sprintf("%s.%s", app.ID, cluster.Host)
	if err := store.UpdateApp(app); err != nil {
		responseError(w, http.StatusInternalServerError, err.Error())
		log.Printf("Error %s: %s", app.ID, err.Error())
		return
	}

	log.Println("Deployment successful.")
	em.Trigger("app.deployed", gde.New(app.ID, app.URL))
}

// Set the different CORS headers required for CORS request
func setCORSHeaders(w http.ResponseWriter, r *http.Request) {
	w.Header().Add("Access-Control-Allow-Methods", "POST, GET, OPTIONS")
	w.Header().Set("Access-Control-Allow-Origin", corsAllowOrigin)
	w.Header().Set("Access-Control-Allow-Credentials", "true")
	w.Header().Set("Access-Control-Allow-Headers", "Origin, X-Requested-With, Content-Type, Accept")
}

func responseError(w http.ResponseWriter, code int, message string) {
	log.Printf("Error %d: %s", code, message)
	response := responder.New(ApiVersion)
	response.Write(w, response.Error(code, message))
}

func responseSuccess(w http.ResponseWriter, data interface{}) {
	response := responder.New(ApiVersion)
	response.Write(w, response.Success(data))
}

func checkDependencies() {
	go func() {
		_, err := exec.LookPath("git")
		if err != nil {
			log.Fatal("Git CLI is required but not installed or not in path.")
		}
	}()
	go checkIfFlynnExists()
}

func checkIfFlynnExists() {
	func() {
		_, err := exec.LookPath("flynn")
		if err != nil {
			if runtime.GOOS == "windows" {
				log.Fatal("Flynn CLI is required but not installed or not in path.")
			}
			log.Println("Could not find Flynn CLI, trying to install...")
			if o, err := exec.Command("sh", "bin/flynn-install.sh").CombinedOutput(); err != nil {
				log.Printf("Could not install Flynn CLI (%s): %s", err.Error(), o)
			} else if _, err := exec.LookPath("flynn"); err != nil {
				log.Fatal("Could not install Flynn CLI.")
			}
            log.Println("Flynn installed successfully!")
            log.Println("Adding flynn cluster...")
            log.Println(envClusterConf)
            args := append([]string{"cluster", "add"}, strings.Split(envClusterConf, " ")...)
            log.Printf("%s", args)
            if o, err := exec.Command("flynn", args...).CombinedOutput(); err != nil {
                log.Fatalf("Could not add cluster (status: %s) (output: %s) (args: %s)", err.Error(), o, args)
            } else {
                log.Printf("Adding cluster successful: %s", o)
            }
		}
	}()
}

func getIP(r *http.Request) string {
    ip := removePort(r.RemoteAddr)
    if len(r.Header.Get("X-FORWARDED-FOR")) > 0 {
        ip = r.Header.Get("X-FORWARDED-FOR")
    }
    return ip
}

func removePort(ip string) string {
    split := strings.Split(ip, ":")
    if len(split) < 2 {
        return ip
    }
    return strings.Join(split[:len(split)-1], ":")
}
<|endoftext|>