<commit_before>package main

import (
	"os"
	"os/exec"
	"syscall"

	"github.com/fsnotify/fsnotify"
	"github.com/jessevdk/go-flags"
)

const (
	program = "Sentinel"
	version = "0.4.0"
	// ACTION is the environment variable for the type of notification triggered.
	ACTION = "SENTINEL_ACTION"
	// PATH is the environment variable for the type of notification triggered.
	PATH = "SENTINEL_PATH"
)

var opts struct {
	Verbose bool `short:"v" long:"verbose" description:"Print more details during operation, otherwise remain quiet until an error occurs."`
	Version bool `short:"V" long:"version" description:"Show program version and exit."`
	Flags   struct {
		Create bool `short:"c" long:"create" description:"Watch for file creation."`
		Write  bool `short:"w" long:"write" description:"Watch for file editing."`
		Delete bool `short:"d" long:"delete" description:"Watch for file deletion."`
		Rename bool `short:"r" long:"rename" description:"Watch for file renaming."`
		Chmod  bool `short:"m" long:"chmod" description:"Watch for file attribute changes (date or permissions)."`
		Loop   bool `short:"L" long:"loop" description:"Don't quit after each triggered event."`
	} `group:"Flags"`
	Commands struct {
		CreateAction string `short:"C" long:"createaction" description:"Script to run when a file is created." value-name:"CMD"`
		WriteAction  string `short:"W" long:"writeaction" description:"Script to run when a file is edited." value-name:"CMD"`
		DeleteAction string `short:"D" long:"deleteaction" description:"Script to run when a file is deleted." value-name:"CMD"`
		RenameAction string `short:"R" long:"renameaction" description:"Script to run when a file is renamed." value-name:"CMD"`
		ChmodAction  string `short:"M" long:"chmodaction" description:"Script to run when a file's date or permissions change." value-name:"CMD"`
	} `group:"Commands"`
	Args struct {
		Directory []string `positional-arg-name:"PATH"`
	} `positional-args:"yes"`
}

func main() {
	_, err := flags.Parse(&opts)
	if err != nil {
		return
	}

	if opts.Version {
		pr("%s %s\n", program, version)
		return
	}

	if len(opts.Args.Directory) == 0 {
		warn("No paths specified.")
	}
	var paths []string
	for _, d := range opts.Args.Directory {
		if !exists(d) {
			warn("Path %s does not exist.", d)
		}
		paths = append(paths, d)
	}

	// Default: Watch for any changes
	var flags fsnotify.Op

	if opts.Flags.Create {
		v("Watching for creation.\n")
		flags |= fsnotify.Create
	}

	if opts.Flags.Write {
		v("Watching for write.\n")
		flags |= fsnotify.Write
	}

	if opts.Flags.Delete {
		v("Watching for delete.\n")
		flags |= fsnotify.Remove
	}

	if opts.Flags.Rename {
		v("Watching for rename.\n")
		flags |= fsnotify.Rename
	}

	if opts.Flags.Chmod {
		v("Watching for permission changes.\n")
		flags |= fsnotify.Chmod
	}

	watcher, err := fsnotify.NewWatcher()
	if err != nil {
		return
	}
	defer watcher.Close()

	done := make(chan bool)
	go func() {
		for {
			select {
			case event := <-watcher.Events:
				if flags&event.Op&fsnotify.Create == fsnotify.Create && opts.Commands.CreateAction != "" {
					v("CREATE: Running '%s'\n", opts.Commands.CreateAction)
					os.Setenv(ACTION, "create")
					os.Setenv(PATH, event.Name)
					runCommand(opts.Commands.CreateAction)
					if !opts.Flags.Loop {
						done <- true
					}
				}
				if flags&event.Op&fsnotify.Write == fsnotify.Write && opts.Commands.WriteAction != "" {
					v("WRITE: Running '%s'\n", opts.Commands.WriteAction)
					os.Setenv(ACTION, "write")
					os.Setenv(PATH, event.Name)
					runCommand(opts.Commands.WriteAction)
					if !opts.Flags.Loop {
						done <- true
					}
				}
				if flags&event.Op&fsnotify.Remove == fsnotify.Remove && opts.Commands.DeleteAction != "" {
					v("REMOVE: Running '%s'\n", opts.Commands.DeleteAction)
					os.Setenv(ACTION, "delete")
					os.Setenv(PATH, event.Name)
					runCommand(opts.Commands.DeleteAction)
					if !opts.Flags.Loop {
						done <- true
					}
				}
				if flags&event.Op&fsnotify.Rename == fsnotify.Rename && opts.Commands.RenameAction != "" {
					v("RENAME: Running '%s'\n", opts.Commands.RenameAction)
					os.Setenv(ACTION, "rename")
					os.Setenv(PATH, event.Name)
					runCommand(opts.Commands.RenameAction)
					if !opts.Flags.Loop {
						done <- true
					}
				}
				if flags&event.Op&fsnotify.Chmod == fsnotify.Chmod && opts.Commands.ChmodAction != "" {
					v("CHMOD: Running '%s'\n", opts.Commands.ChmodAction)
					os.Setenv(ACTION, "chmod")
					os.Setenv(PATH, event.Name)
					runCommand(opts.Commands.ChmodAction)
					if !opts.Flags.Loop {
						done <- true
					}
				}
			case err := <-watcher.Errors:
				if err.Error() != "" {
					fatal("Error: ", err.Error())
				}
				done <- true
			}
		}
	}()

	for _, dir := range paths {
		v("* %s\n", dir)
		err = watcher.Add(dir)
		if err != nil {
			fatal(err.Error())
		}
	}

	// We'll never return from this without a break signal if in loop mode
	<-done
}

func runCommand(script string) {
	cmd := exec.Command("bash", script)
	err := cmd.Start()
	if err != nil {
		v("Error: %s\n", err)
	}

	err = cmd.Wait()
	if err != nil {
		exit, ok := err.(*exec.ExitError)
		if ok {
			status, ok := exit.Sys().(syscall.WaitStatus)
			if ok && status == 256 {
				os.Exit(0)
				v("Exit code: %d\n", status)
			}
		} else {
			v("Error: %s\n", err)
		}
	}

}
<commit_msg>Checking for return code 256 and 512 should work on Linux.<commit_after>package main

import (
	"os"
	"os/exec"
	"syscall"

	"github.com/fsnotify/fsnotify"
	"github.com/jessevdk/go-flags"
)

const (
	program = "Sentinel"
	version = "0.4.1"
	// ACTION is the environment variable for the type of notification triggered.
	ACTION = "SENTINEL_ACTION"
	// PATH is the environment variable for the type of notification triggered.
	PATH = "SENTINEL_PATH"
)

var opts struct {
	Verbose bool `short:"v" long:"verbose" description:"Print more details during operation, otherwise remain quiet until an error occurs."`
	Version bool `short:"V" long:"version" description:"Show program version and exit."`
	Flags   struct {
		Create bool `short:"c" long:"create" description:"Watch for file creation."`
		Write  bool `short:"w" long:"write" description:"Watch for file editing."`
		Delete bool `short:"d" long:"delete" description:"Watch for file deletion."`
		Rename bool `short:"r" long:"rename" description:"Watch for file renaming."`
		Chmod  bool `short:"m" long:"chmod" description:"Watch for file attribute changes (date or permissions)."`
		Loop   bool `short:"L" long:"loop" description:"Don't quit after each triggered event."`
	} `group:"Flags"`
	Commands struct {
		CreateAction string `short:"C" long:"createaction" description:"Script to run when a file is created." value-name:"CMD"`
		WriteAction  string `short:"W" long:"writeaction" description:"Script to run when a file is edited." value-name:"CMD"`
		DeleteAction string `short:"D" long:"deleteaction" description:"Script to run when a file is deleted." value-name:"CMD"`
		RenameAction string `short:"R" long:"renameaction" description:"Script to run when a file is renamed." value-name:"CMD"`
		ChmodAction  string `short:"M" long:"chmodaction" description:"Script to run when a file's date or permissions change." value-name:"CMD"`
	} `group:"Commands"`
	Args struct {
		Directory []string `positional-arg-name:"PATH"`
	} `positional-args:"yes"`
}

func main() {
	_, err := flags.Parse(&opts)
	if err != nil {
		return
	}

	if opts.Version {
		pr("%s %s\n", program, version)
		return
	}

	if len(opts.Args.Directory) == 0 {
		warn("No paths specified.")
	}
	var paths []string
	for _, d := range opts.Args.Directory {
		if !exists(d) {
			warn("Path %s does not exist.", d)
		}
		paths = append(paths, d)
	}

	// Default: Watch for any changes
	var flags fsnotify.Op

	if opts.Flags.Create {
		v("Watching for creation.\n")
		flags |= fsnotify.Create
	}

	if opts.Flags.Write {
		v("Watching for write.\n")
		flags |= fsnotify.Write
	}

	if opts.Flags.Delete {
		v("Watching for delete.\n")
		flags |= fsnotify.Remove
	}

	if opts.Flags.Rename {
		v("Watching for rename.\n")
		flags |= fsnotify.Rename
	}

	if opts.Flags.Chmod {
		v("Watching for permission changes.\n")
		flags |= fsnotify.Chmod
	}

	watcher, err := fsnotify.NewWatcher()
	if err != nil {
		return
	}
	defer watcher.Close()

	done := make(chan bool)
	go func() {
		for {
			select {
			case event := <-watcher.Events:
				if flags&event.Op&fsnotify.Create == fsnotify.Create && opts.Commands.CreateAction != "" {
					v("CREATE: Running '%s'\n", opts.Commands.CreateAction)
					os.Setenv(ACTION, "create")
					os.Setenv(PATH, event.Name)
					runCommand(opts.Commands.CreateAction)
					if !opts.Flags.Loop {
						done <- true
					}
				}
				if flags&event.Op&fsnotify.Write == fsnotify.Write && opts.Commands.WriteAction != "" {
					v("WRITE: Running '%s'\n", opts.Commands.WriteAction)
					os.Setenv(ACTION, "write")
					os.Setenv(PATH, event.Name)
					runCommand(opts.Commands.WriteAction)
					if !opts.Flags.Loop {
						done <- true
					}
				}
				if flags&event.Op&fsnotify.Remove == fsnotify.Remove && opts.Commands.DeleteAction != "" {
					v("REMOVE: Running '%s'\n", opts.Commands.DeleteAction)
					os.Setenv(ACTION, "delete")
					os.Setenv(PATH, event.Name)
					runCommand(opts.Commands.DeleteAction)
					if !opts.Flags.Loop {
						done <- true
					}
				}
				if flags&event.Op&fsnotify.Rename == fsnotify.Rename && opts.Commands.RenameAction != "" {
					v("RENAME: Running '%s'\n", opts.Commands.RenameAction)
					os.Setenv(ACTION, "rename")
					os.Setenv(PATH, event.Name)
					runCommand(opts.Commands.RenameAction)
					if !opts.Flags.Loop {
						done <- true
					}
				}
				if flags&event.Op&fsnotify.Chmod == fsnotify.Chmod && opts.Commands.ChmodAction != "" {
					v("CHMOD: Running '%s'\n", opts.Commands.ChmodAction)
					os.Setenv(ACTION, "chmod")
					os.Setenv(PATH, event.Name)
					runCommand(opts.Commands.ChmodAction)
					if !opts.Flags.Loop {
						done <- true
					}
				}
			case err := <-watcher.Errors:
				if err.Error() != "" {
					fatal("Error: ", err.Error())
				}
				done <- true
			}
		}
	}()

	for _, dir := range paths {
		v("* %s\n", dir)
		err = watcher.Add(dir)
		if err != nil {
			fatal(err.Error())
		}
	}

	// We'll never return from this without a break signal if in loop mode
	<-done
}

func runCommand(script string) {
	cmd := exec.Command("bash", script)
	err := cmd.Start()
	if err != nil {
		v("Error: %s\n", err)
	}

	err = cmd.Wait()
	if err != nil {
		exit, ok := err.(*exec.ExitError)
		if ok {
			status, ok := exit.Sys().(syscall.WaitStatus)
			if ok {
				if status == 256 || status == 512 {
					os.Exit(0)
					v("Exit code: %d\n", status)
				}
			}
		} else {
			v("Error: %s\n", err)
		}
	}

}
<|endoftext|>


<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/sampler"
	"github.com/loadimpact/speedboat/sampler/influxdb"
	"github.com/loadimpact/speedboat/sampler/stream"
	"github.com/loadimpact/speedboat/simple"
	"github.com/urfave/cli"
	"golang.org/x/net/context"
	"gopkg.in/yaml.v2"
	"io"
	"io/ioutil"
	"os"
	"strings"
	"time"
)

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

// Configure the global logger.
func configureLogging(c *cli.Context) {
	log.SetLevel(log.InfoLevel)
	if c.GlobalBool("verbose") {
		log.SetLevel(log.DebugLevel)
	}
}

// Configure the global sampler.
func configureSampler(c *cli.Context) {
	sampler.DefaultSampler.OnError = func(err error) {
		log.WithError(err).Error("[Sampler error]")
	}

	for _, output := range c.GlobalStringSlice("metrics") {
		parts := strings.SplitN(output, "+", 2)
		switch parts[0] {
		case "influxdb":
			out, err := influxdb.NewFromURL(parts[1])
			if err != nil {
				log.WithError(err).Fatal("Couldn't create InfluxDB client")
			}
			sampler.DefaultSampler.Outputs = append(sampler.DefaultSampler.Outputs, out)
		default:
			var writer io.WriteCloser
			switch output {
			case "stdout", "-":
				writer = os.Stdout
			default:
				file, err := os.Create(output)
				if err != nil {
					log.WithError(err).Fatal("Couldn't create output file")
				}
				writer = file
			}

			var out sampler.Output
			switch c.GlobalString("format") {
			case "json":
				out = &stream.JSONOutput{Output: writer}
			case "csv":
				out = &stream.CSVOutput{Output: writer}
			default:
				log.Fatal("Unknown output format")
			}
			sampler.DefaultSampler.Outputs = append(sampler.DefaultSampler.Outputs, out)
		}
	}
}

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

func parse(cc *cli.Context) (conf Config, err error) {
	if len(cc.Args()) == 0 {
		return conf, errors.New("Nothing to do!")
	}

	conf.VUs = cc.Int("vus")
	conf.Duration = cc.Duration("duration").String()

	arg := cc.Args()[0]
	argType := cc.String("type")
	if argType == "" {
		argType = guessType(arg)
	}

	switch argType {
	case typeYML:
		bytes, err := ioutil.ReadFile(cc.Args()[0])
		if err != nil {
			return conf, errors.New("Couldn't read config file")
		}
		if err := yaml.Unmarshal(bytes, &conf); err != nil {
			return conf, errors.New("Couldn't parse config file")
		}
	case typeURL:
		conf.URL = arg
	case typeJS:
		conf.Script = arg
	default:
		return conf, errors.New("Unsure of what to do, try specifying --type")
	}

	return conf, nil
}

func headlessController(c context.Context, t *lib.Test) <-chan int {
	ch := make(chan int)

	go func() {
		defer close(ch)

		select {
		case ch <- t.VUsAt(0):
		case <-c.Done():
			return
		}

		startTime := time.Now()
		ticker := time.NewTicker(100 * time.Millisecond)
		for {
			select {
			case <-ticker.C:
				ch <- t.VUsAt(time.Since(startTime))
			case <-c.Done():
				return
			}
		}
	}()

	return ch
}

func action(cc *cli.Context) error {
	if len(cc.Args()) == 0 {
		cli.ShowAppHelp(cc)
		return nil
	}

	conf, err := parse(cc)
	if err != nil {
		log.WithError(err).Fatal("Invalid arguments; see --help")
	}

	t, err := conf.MakeTest()
	if err != nil {
		log.WithError(err).Fatal("Configuration error")
	}

	// Inspect the test to find a suitable runner; additional ones can easily be added
	var runner lib.Runner
	switch {
	case t.Script == "":
		runner = simple.New(t)
	case strings.HasSuffix(t.Script, ".js"):
		src, err := ioutil.ReadFile(t.Script)
		if err != nil {
			log.WithError(err).Fatal("Couldn't read script")
		}
		runner = js.New(t, t.Script, string(src))
	default:
		log.Fatal("No suitable runner found!")
	}

	// Context that expires at the end of the test
	ctx, cancel := context.WithTimeout(context.Background(), t.TotalDuration())

	// Configure the VU logger
	logger := &log.Logger{
		Out:       os.Stderr,
		Level:     log.DebugLevel,
		Formatter: &log.TextFormatter{},
	}
	ctx = lib.WithLogger(ctx, logger)

	// Store metrics unless the --quiet flag is specified
	quiet := cc.Bool("quiet")
	sampler.DefaultSampler.Accumulate = !quiet

	// Commit metrics to any configured backends once per second
	go func() {
		ticker := time.NewTicker(1 * time.Second)
		for {
			select {
			case <-ticker.C:
				commitMetrics()
			case <-ctx.Done():
				return
			}
		}
	}()

	// Use a "headless controller" to scale VUs by polling the test ramp
	mVUs := sampler.Gauge("vus")
	vus := []context.CancelFunc{}
	for scale := range headlessController(ctx, &t) {
		for i := len(vus); i < scale; i++ {
			log.WithField("id", i).Debug("Spawning VU")
			vuCtx, vuCancel := context.WithCancel(ctx)
			vus = append(vus, vuCancel)

			go func(ctx context.Context) {
				defer func() {
					if v := recover(); v != nil {
						switch err := v.(type) {
						case lib.FlowControl:
							switch err {
							case lib.AbortTest:
								log.Error("Test aborted")
								cancel()
							}
						default:
							panic(err)
						}
					}
				}()

				vu, err := runner.NewVU()
				if err != nil {
					log.WithError(err).Error("Couldn't spawn VU")
					return
				}

				vu.Reconfigure(int64(i))
				for {
					select {
					case <-ctx.Done():
						return
					default:
						if err := vu.RunOnce(ctx); err != nil {
							log.WithError(err).Error("Script error")
						}
					}
				}
			}(vuCtx)
		}
		for i := len(vus); i > scale; i-- {
			log.WithField("id", i-1).Debug("Dropping VU")
			vus[i-1]()
			vus = vus[:i-1]
		}
		mVUs.Int(len(vus))
	}

	// Wait until the end of the test
	<-ctx.Done()

	// Print and commit final metrics
	if !quiet {
		printMetrics()
	}
	commitMetrics()
	closeMetrics()

	return nil
}

func main() {
	// Free up -v and -h for our own flags
	cli.VersionFlag.Name = "version"
	cli.HelpFlag.Name = "help, ?"
	cli.AppHelpTemplate = `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}}`

	// Bootstrap using action-registered commandline flags
	app := cli.NewApp()
	app.Commands = nil
	app.Name = "speedboat"
	app.Usage = "A next-generation load generator"
	app.Version = "1.0.0-mvp1"
	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:  "type, t",
			Usage: "Input file type, if not evident (url, yml or js)",
		},
		cli.IntFlag{
			Name:  "vus, u",
			Usage: "Number of VUs to simulate",
			Value: 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.StringSliceFlag{
			Name:  "metrics, m",
			Usage: "Write metrics to a file or database",
		},
		cli.StringFlag{
			Name:  "format, f",
			Usage: "Metric output format (json or csv)",
			Value: "json",
		},
	}
	app.Before = func(c *cli.Context) error {
		configureLogging(c)
		configureSampler(c)
		return nil
	}
	app.Action = action
	app.Run(os.Args)
}
<commit_msg>[fix] Made help template text a const at top of file<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/sampler"
	"github.com/loadimpact/speedboat/sampler/influxdb"
	"github.com/loadimpact/speedboat/sampler/stream"
	"github.com/loadimpact/speedboat/simple"
	"github.com/urfave/cli"
	"golang.org/x/net/context"
	"gopkg.in/yaml.v2"
	"io"
	"io/ioutil"
	"os"
	"strings"
	"time"
)

const (
	typeURL = "url"
	typeYML = "yml"
	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}}`
)

// Configure the global logger.
func configureLogging(c *cli.Context) {
	log.SetLevel(log.InfoLevel)
	if c.GlobalBool("verbose") {
		log.SetLevel(log.DebugLevel)
	}
}

// Configure the global sampler.
func configureSampler(c *cli.Context) {
	sampler.DefaultSampler.OnError = func(err error) {
		log.WithError(err).Error("[Sampler error]")
	}

	for _, output := range c.GlobalStringSlice("metrics") {
		parts := strings.SplitN(output, "+", 2)
		switch parts[0] {
		case "influxdb":
			out, err := influxdb.NewFromURL(parts[1])
			if err != nil {
				log.WithError(err).Fatal("Couldn't create InfluxDB client")
			}
			sampler.DefaultSampler.Outputs = append(sampler.DefaultSampler.Outputs, out)
		default:
			var writer io.WriteCloser
			switch output {
			case "stdout", "-":
				writer = os.Stdout
			default:
				file, err := os.Create(output)
				if err != nil {
					log.WithError(err).Fatal("Couldn't create output file")
				}
				writer = file
			}

			var out sampler.Output
			switch c.GlobalString("format") {
			case "json":
				out = &stream.JSONOutput{Output: writer}
			case "csv":
				out = &stream.CSVOutput{Output: writer}
			default:
				log.Fatal("Unknown output format")
			}
			sampler.DefaultSampler.Outputs = append(sampler.DefaultSampler.Outputs, out)
		}
	}
}

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

func parse(cc *cli.Context) (conf Config, err error) {
	if len(cc.Args()) == 0 {
		return conf, errors.New("Nothing to do!")
	}

	conf.VUs = cc.Int("vus")
	conf.Duration = cc.Duration("duration").String()

	arg := cc.Args()[0]
	argType := cc.String("type")
	if argType == "" {
		argType = guessType(arg)
	}

	switch argType {
	case typeYML:
		bytes, err := ioutil.ReadFile(cc.Args()[0])
		if err != nil {
			return conf, errors.New("Couldn't read config file")
		}
		if err := yaml.Unmarshal(bytes, &conf); err != nil {
			return conf, errors.New("Couldn't parse config file")
		}
	case typeURL:
		conf.URL = arg
	case typeJS:
		conf.Script = arg
	default:
		return conf, errors.New("Unsure of what to do, try specifying --type")
	}

	return conf, nil
}

func headlessController(c context.Context, t *lib.Test) <-chan int {
	ch := make(chan int)

	go func() {
		defer close(ch)

		select {
		case ch <- t.VUsAt(0):
		case <-c.Done():
			return
		}

		startTime := time.Now()
		ticker := time.NewTicker(100 * time.Millisecond)
		for {
			select {
			case <-ticker.C:
				ch <- t.VUsAt(time.Since(startTime))
			case <-c.Done():
				return
			}
		}
	}()

	return ch
}

func action(cc *cli.Context) error {
	if len(cc.Args()) == 0 {
		cli.ShowAppHelp(cc)
		return nil
	}

	conf, err := parse(cc)
	if err != nil {
		log.WithError(err).Fatal("Invalid arguments; see --help")
	}

	t, err := conf.MakeTest()
	if err != nil {
		log.WithError(err).Fatal("Configuration error")
	}

	// Inspect the test to find a suitable runner; additional ones can easily be added
	var runner lib.Runner
	switch {
	case t.Script == "":
		runner = simple.New(t)
	case strings.HasSuffix(t.Script, ".js"):
		src, err := ioutil.ReadFile(t.Script)
		if err != nil {
			log.WithError(err).Fatal("Couldn't read script")
		}
		runner = js.New(t, t.Script, string(src))
	default:
		log.Fatal("No suitable runner found!")
	}

	// Context that expires at the end of the test
	ctx, cancel := context.WithTimeout(context.Background(), t.TotalDuration())

	// Configure the VU logger
	logger := &log.Logger{
		Out:       os.Stderr,
		Level:     log.DebugLevel,
		Formatter: &log.TextFormatter{},
	}
	ctx = lib.WithLogger(ctx, logger)

	// Store metrics unless the --quiet flag is specified
	quiet := cc.Bool("quiet")
	sampler.DefaultSampler.Accumulate = !quiet

	// Commit metrics to any configured backends once per second
	go func() {
		ticker := time.NewTicker(1 * time.Second)
		for {
			select {
			case <-ticker.C:
				commitMetrics()
			case <-ctx.Done():
				return
			}
		}
	}()

	// Use a "headless controller" to scale VUs by polling the test ramp
	mVUs := sampler.Gauge("vus")
	vus := []context.CancelFunc{}
	for scale := range headlessController(ctx, &t) {
		for i := len(vus); i < scale; i++ {
			log.WithField("id", i).Debug("Spawning VU")
			vuCtx, vuCancel := context.WithCancel(ctx)
			vus = append(vus, vuCancel)

			go func(ctx context.Context) {
				defer func() {
					if v := recover(); v != nil {
						switch err := v.(type) {
						case lib.FlowControl:
							switch err {
							case lib.AbortTest:
								log.Error("Test aborted")
								cancel()
							}
						default:
							panic(err)
						}
					}
				}()

				vu, err := runner.NewVU()
				if err != nil {
					log.WithError(err).Error("Couldn't spawn VU")
					return
				}

				vu.Reconfigure(int64(i))
				for {
					select {
					case <-ctx.Done():
						return
					default:
						if err := vu.RunOnce(ctx); err != nil {
							log.WithError(err).Error("Script error")
						}
					}
				}
			}(vuCtx)
		}
		for i := len(vus); i > scale; i-- {
			log.WithField("id", i-1).Debug("Dropping VU")
			vus[i-1]()
			vus = vus[:i-1]
		}
		mVUs.Int(len(vus))
	}

	// Wait until the end of the test
	<-ctx.Done()

	// Print and commit final metrics
	if !quiet {
		printMetrics()
	}
	commitMetrics()
	closeMetrics()

	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 using action-registered commandline flags
	app := cli.NewApp()
	app.Commands = nil
	app.Name = "speedboat"
	app.Usage = "A next-generation load generator"
	app.Version = "1.0.0-mvp1"
	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:  "type, t",
			Usage: "Input file type, if not evident (url, yml or js)",
		},
		cli.IntFlag{
			Name:  "vus, u",
			Usage: "Number of VUs to simulate",
			Value: 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.StringSliceFlag{
			Name:  "metrics, m",
			Usage: "Write metrics to a file or database",
		},
		cli.StringFlag{
			Name:  "format, f",
			Usage: "Metric output format (json or csv)",
			Value: "json",
		},
	}
	app.Before = func(c *cli.Context) error {
		configureLogging(c)
		configureSampler(c)
		return nil
	}
	app.Action = action
	app.Run(os.Args)
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"image"
	"image/jpeg"
	"io"
	"net/http"
	"os"
	"path"
	"strings"
	"unicode"

	"golang.org/x/text/transform"
	"golang.org/x/text/unicode/norm"

	_ "image/gif"
	_ "image/png"
)

func main() {
	var (
		name    = flag.String("name", "", "name of the image")
		quality = flag.Int("quality", 100, "jpeg quality encoding")
	)
	flag.Parse()
	args := flag.Args()
	if len(args) != 1 {
		fmt.Fprintln(os.Stderr, "wimg: invalid arguments")
		flag.PrintDefaults()
		os.Exit(1)
	}
	err := run(args[0], *name, *quality)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}

func run(src, name string, quality int) error {
	if name == "" {
		name = baseWithoutExt(src)
	}
	name, err := normalize(name)
	if err != nil {
		return err
	}
	resp, err := http.Get(src)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	return save(name, resp.Body, quality)
}

func baseWithoutExt(src string) string {
	ext := path.Ext(src)
	name := path.Base(src)
	return strings.TrimSuffix(name, ext)
}

func normalize(name string) (string, error) {
	t := transform.Chain(norm.NFD, transform.RemoveFunc(remove), norm.NFC)
	name = strings.TrimSpace(name)
	name, _, err := transform.String(t, name)
	if err != nil {
		return "", err
	}
	name = strings.ToLower(name)
	name = strings.Replace(name, " ", "_", -1)
	return name, nil
}

func remove(r rune) bool {
	return unicode.Is(unicode.Mn, r)
}

func save(name string, r io.Reader, quality int) error {
	m, _, err := image.Decode(r)
	if err != nil {
		return err
	}
	rect := m.Bounds()
	filename := getSaveName(name, rect)
	f, err := os.Create(filename)
	if err != nil {
		return err
	}
	defer f.Close()
	return jpeg.Encode(f, m, &jpeg.Options{Quality: quality})
}

func getSaveName(name string, rect image.Rectangle) string {
	return fmt.Sprintf("%dx%d_%s.jpg", rect.Dx(), rect.Dy(), name)
}
<commit_msg>Save to specified directory.<commit_after>package main

import (
	"flag"
	"fmt"
	"image"
	"image/jpeg"
	"io"
	"net/http"
	"os"
	"path"
	"path/filepath"
	"strings"
	"unicode"

	"golang.org/x/text/transform"
	"golang.org/x/text/unicode/norm"

	_ "image/gif"
	_ "image/png"
)

func main() {
	var (
		dir     = flag.String("dir", "", "directory to save file (default: cwd)")
		name    = flag.String("name", "", "name of the image")
		quality = flag.Int("quality", 100, "jpeg quality encoding")
	)
	flag.Parse()
	args := flag.Args()
	if len(args) != 1 {
		fmt.Fprintln(os.Stderr, "wimg: invalid arguments")
		flag.PrintDefaults()
		os.Exit(1)
	}
	o := &options{
		dir:     *dir,
		name:    *name,
		quality: *quality,
	}
	err := run(args[0], o)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}

type options struct {
	dir     string
	name    string
	quality int
}

func (o *options) clean(src string) error {
	var err error
	if o.name == "" {
		o.name = baseWithoutExt(src)
	}
	o.name, err = normalize(o.name)
	if err != nil {
		return err
	}
	return nil
}

func run(src string, o *options) error {
	err := o.clean(src)
	if err != nil {
		return err
	}
	resp, err := http.Get(src)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	return save(resp.Body, o)
}

func baseWithoutExt(src string) string {
	ext := path.Ext(src)
	name := path.Base(src)
	return strings.TrimSuffix(name, ext)
}

func normalize(name string) (string, error) {
	t := transform.Chain(norm.NFD, transform.RemoveFunc(remove), norm.NFC)
	name = strings.TrimSpace(name)
	name, _, err := transform.String(t, name)
	if err != nil {
		return "", err
	}
	name = strings.ToLower(name)
	name = strings.Replace(name, " ", "_", -1)
	return name, nil
}

func remove(r rune) bool {
	return unicode.Is(unicode.Mn, r)
}

func save(r io.Reader, o *options) error {
	m, _, err := image.Decode(r)
	if err != nil {
		return err
	}
	rect := m.Bounds()
	filename := getSaveName(rect, o)
	f, err := os.Create(filename)
	if err != nil {
		return err
	}
	defer f.Close()
	return jpeg.Encode(f, m, &jpeg.Options{Quality: o.quality})
}

func getSaveName(rect image.Rectangle, o *options) string {
	name := fmt.Sprintf("%dx%d_%s.jpg", rect.Dx(), rect.Dy(), o.name)
	return filepath.Join(o.dir, name)
}
<|endoftext|>


<commit_before>package main

import (
	//"fmt"
	"html/template"
	"net/http"
	"path/filepath"
)

var templates = template.Must(template.ParseFiles(
	filepath.Join(myjournal.TemplateDIR, "templates/base.html"),
	filepath.Join(myjournal.TemplateDIR, "templates/sidebar.html"),
	filepath.Join(myjournal.TemplateDIR, "templates/content.html"),
	filepath.Join(myjournal.TemplateDIR, "templates/customjs.html"),
))

type Page struct {
	Title     string
	IsList    bool
	PrevId    string
	NextId    string
	HttpFQDN  string
	Search    string
	CssLookup map[string]string
	Navbar    interface{}
	Content   interface{}
}

func renderTemplate(w http.ResponseWriter, tmpl string, p *Page) {
	err := templates.Execute(w, p)
	//err := templates.ExecuteTemplate(w, tmpl+".html", p)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}
}
<commit_msg>removed redundent template dir<commit_after>package main

import (
	//"fmt"
	"html/template"
	"net/http"
	"path/filepath"
)

var templates = template.Must(template.ParseFiles(
	filepath.Join(myjournal.TemplateDIR, "base.html"),
	filepath.Join(myjournal.TemplateDIR, "sidebar.html"),
	filepath.Join(myjournal.TemplateDIR, "content.html"),
	filepath.Join(myjournal.TemplateDIR, "customjs.html"),
))

type Page struct {
	Title     string
	IsList    bool
	PrevId    string
	NextId    string
	HttpFQDN  string
	Search    string
	CssLookup map[string]string
	Navbar    interface{}
	Content   interface{}
}

func renderTemplate(w http.ResponseWriter, tmpl string, p *Page) {
	err := templates.Execute(w, p)
	//err := templates.ExecuteTemplate(w, tmpl+".html", p)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}
}
<|endoftext|>


<commit_before>// +build !linux, !arm

package wiringpi

import (
	"time"
	"log"
)

//noinspection ALL
const (
	WPI_MODE_PINS          = 0
	WPI_MODE_GPIO          = 1
	WPI_MODE_GPIO_SYS      = 2
	WPI_MODE_PIFACE        = 3
	WPI_MODE_UNINITIALISED = 4

	INPUT      = 0
	OUTPUT     = 1
	PWM_OUTPUT = 2
	GPIO_CLOCK = 3

	LOW  = 0
	HIGH = 1

	PUD_OFF  = 1
	PUD_DOWN = 2
	PUD_UP   = 3

	PWM_MODE_MS  = 0
	PWM_MODE_BAL = 1

	INT_EDGE_SETUP   = 0
	INT_EDGE_FALLING = 1
	INT_EDGE_RISING  = 2
	INT_EDGE_BOTH    = 3

	// Used for emulator:
	gpio_pin_count = 26
)

var(
	gpio_list [gpio_pin_count]int
	gpio_mode_list [gpio_pin_count]int
)

func PinToGpio(_ int) int {
	return 0
}

func Setup() error {
	log.Println("Warning: Running in emulation mode")

	for i:= 0; i < gpio_pin_count; i++ {
		gpio_list[i] = LOW
		gpio_mode_list[i] = MODE_IN
	}

	return nil
}

func PinMode(pin int, mode int) {
	if pin < gpio_pin_count {
		gpio_mode_list[pin] = mode
	}
}

func DigitalWrite(pin int, mode int) {
	if pin < gpio_pin_count {
		gpio_list[pin] = mode
	}
}

func DigitalRead(pin int) int {
	if pin < gpio_pin_count {
		return gpio_list[pin]
	}

	return LOW
}

func DigitalReadStr(pin int) string {
	if DigitalRead(pin) == LOW {
		return "LOW"
	}
	return "HIGH"
}

func GetMode(pin int) int {
	if pin < gpio_pin_count {
		return gpio_mode_list[pin]
	}
}

func GetModeStr(pin int) string {
	var mode = GetMode(pin)

	if mode > len(gpioModes) {
		return "INVALID"
	}

	return gpioModes[GetMode(pin)]
}

func Delay(ms int) {
	time.Sleep(time.Duration(ms) * time.Millisecond)
}

func DelayMicroseconds(microSec int) {
	time.Sleep(time.Duration(microSec) * time.Microsecond)
}

func WiringISR(pin int, mode int) chan int {
	if pin < gpio_pin_count {
		gpio_list[pin] = mode
	}
	return nil
}<commit_msg>Fixed return error in emulator<commit_after>// +build !linux, !arm

package wiringpi

import (
	"time"
	"log"
)

//noinspection ALL
const (
	WPI_MODE_PINS          = 0
	WPI_MODE_GPIO          = 1
	WPI_MODE_GPIO_SYS      = 2
	WPI_MODE_PIFACE        = 3
	WPI_MODE_UNINITIALISED = 4

	INPUT      = 0
	OUTPUT     = 1
	PWM_OUTPUT = 2
	GPIO_CLOCK = 3

	LOW  = 0
	HIGH = 1

	PUD_OFF  = 1
	PUD_DOWN = 2
	PUD_UP   = 3

	PWM_MODE_MS  = 0
	PWM_MODE_BAL = 1

	INT_EDGE_SETUP   = 0
	INT_EDGE_FALLING = 1
	INT_EDGE_RISING  = 2
	INT_EDGE_BOTH    = 3

	// Used for emulator:
	gpio_pin_count = 26
)

var(
	gpio_list [gpio_pin_count]int
	gpio_mode_list [gpio_pin_count]int
)

func PinToGpio(_ int) int {
	return 0
}

func Setup() error {
	log.Println("Warning: Running in emulation mode")

	for i:= 0; i < gpio_pin_count; i++ {
		gpio_list[i] = LOW
		gpio_mode_list[i] = MODE_IN
	}

	return nil
}

func PinMode(pin int, mode int) {
	if pin < gpio_pin_count {
		gpio_mode_list[pin] = mode
	}
}

func DigitalWrite(pin int, mode int) {
	if pin < gpio_pin_count {
		gpio_list[pin] = mode
	}
}

func DigitalRead(pin int) int {
	if pin < gpio_pin_count {
		return gpio_list[pin]
	}

	return LOW
}

func DigitalReadStr(pin int) string {
	if DigitalRead(pin) == LOW {
		return "LOW"
	}
	return "HIGH"
}

func GetMode(pin int) int {
	if pin < gpio_pin_count {
		return gpio_mode_list[pin]
	}

	return MODE_IN
}

func GetModeStr(pin int) string {
	var mode = GetMode(pin)

	if mode > len(gpioModes) {
		return "INVALID"
	}

	return gpioModes[GetMode(pin)]
}

func Delay(ms int) {
	time.Sleep(time.Duration(ms) * time.Millisecond)
}

func DelayMicroseconds(microSec int) {
	time.Sleep(time.Duration(microSec) * time.Microsecond)
}

func WiringISR(pin int, mode int) chan int {
	if pin < gpio_pin_count {
		gpio_list[pin] = mode
	}
	return nil
}<|endoftext|>


<commit_before>package pagination

const elipses = "..."

// Pagination contains the information necessary to print a proper pagination
type Pagination struct {
	numSections byte
	page        uint
	sections    [3]section
}

// Print converts the Pagination sections into a string.
//
// The pageFn func takes a page number and returns whatever text is needed for
// that page.
//
// The sep string is what is to appear between the sections.
func (p Pagination) Print(pageFn func(uint) string, sep string) string {
	str := ""
	for i := byte(0); i < p.numSections; i++ {
		if i != 0 {
			str += sep
		}
		for page := p.sections[i].First; page <= p.sections[i].Last; page++ {
			str += pageFn(page)
		}
	}
	return str
}

// HTML calls Print with a HTML based pageFn and a simple elipses.
//
// The urlBase will have the page number appended to it, so it needs to be
// formatted with this in mind.
func (p Pagination) HTML(urlBase string) string {
	return p.Print(func(page uint) string {
		numStr := itoa(page)
		if page == p.page {
			return numStr + " "
		}
		return "" + numStr + " "
	}, elipses)
}

// String stringifies the Sections with a simple pageFn
func (p Pagination) String() string {
	return p.Print(func(page uint) string {
		return itoa(page) + " "
	}, elipses)
}

// Section contains the information for a single section of a pagination.
type section struct {
	First, Last uint
}

func itoa(n uint) string {
	if n == 0 {
		return "0"
	}
	var digits [20]byte
	pos := 20
	for ; n > 0; n /= 10 {
		pos--
		digits[pos] = '0' + byte(n%10)
	}
	return string(digits[pos:])
}
<commit_msg>Off-by-one error<commit_after>package pagination

const elipses = "..."

// Pagination contains the information necessary to print a proper pagination
type Pagination struct {
	numSections byte
	page        uint
	sections    [3]section
}

// Print converts the Pagination sections into a string.
//
// The pageFn func takes a page number and returns whatever text is needed for
// that page.
//
// The sep string is what is to appear between the sections.
func (p Pagination) Print(pageFn func(uint) string, sep string) string {
	str := ""
	for i := byte(0); i < p.numSections; i++ {
		if i != 0 {
			str += sep
		}
		for page := p.sections[i].First; page <= p.sections[i].Last; page++ {
			str += pageFn(page)
		}
	}
	return str
}

// HTML calls Print with a HTML based pageFn and a simple elipses.
//
// The urlBase will have the page number appended to it, so it needs to be
// formatted with this in mind.
func (p Pagination) HTML(urlBase string) string {
	return p.Print(func(page uint) string {
		numStr := itoa(page + 1)
		if page == p.page {
			return numStr + " "
		}
		return "" + numStr + " "
	}, elipses)
}

// String stringifies the Sections with a simple pageFn
func (p Pagination) String() string {
	return p.Print(func(page uint) string {
		return itoa(page+1) + " "
	}, elipses)
}

// Section contains the information for a single section of a pagination.
type section struct {
	First, Last uint
}

func itoa(n uint) string {
	if n == 0 {
		return "0"
	}
	var digits [20]byte
	pos := 20
	for ; n > 0; n /= 10 {
		pos--
		digits[pos] = '0' + byte(n%10)
	}
	return string(digits[pos:])
}
<|endoftext|>


<commit_before>package render

import (
	"github.com/quarnster/util/text"
	"image/color"
	"sort"
)

type (
	// Colour represented by a underlying color.RGBA structure
	Colour color.RGBA

	// The Flavour struct contains the specific settings
	// used to style a particular Region.
	Flavour struct {
		Background Colour
		Foreground Colour
		Font       Font
		Flags      ViewRegionFlags
	}

	// The Recipe type groups text.RegionSets by their Flavour.
	// The idea is to allow large groups of text be rendered as
	// a single batch without any state changes inbetween the
	// batches.
	Recipe map[Flavour]text.RegionSet

	// A RenderUnit is just a Flavour and an associated Region.
	RenderUnit struct {
		Flavour Flavour
		Region  text.Region
	}

	// A TranscribedRecipe is a linear (in text.Regions) representation of a Recipe
	TranscribedRecipe []RenderUnit

	ColourScheme interface {
		// Takes a ViewRegions pointer as input and uses the data contained in it
		// to determine the Flavour it should be rendered with.
		Spice(*ViewRegions) Flavour
	}

	Renderer interface {
		// Renders the given Recipe
		Render(Recipe)
	}
)

// Transform takes a ColourScheme, a ViewRegionMap and a viewport as input.
//
// The viewport would be the text.Region of the current buffer that is visible to the user
// and any ViewRegions outside of this area are not forwarded for further processing.
//
// The remaining ViewRegions are then passed on to the ColourScheme for determining the exact Flavour
// for which that RegionSet should be styled, adding Regions of the same Flavour to the same RegionSet.
//
// Typically there are more ViewRegions available in a text buffer than there are unique Flavours in
// a ColourScheme, so this operation can be viewed as reducing the number of state changes required to
// display the text to the user.
//
// The final output, the Recipe, contains a mapping of all unique Flavours and that Flavour's
// associated RegionSet.
func Transform(scheme ColourScheme, data ViewRegionMap, viewport text.Region) Recipe {
	// TODO:
	// 	caret_blink := true
	// if b, ok := v.Settings().Get("caret_blink", true).(bool); ok {
	// 	caret_blink = b
	// }
	//
	// highlight_line := false
	// if b, ok := v.Settings().Get("highlight_line", highlight_line).(bool); ok {
	// 	highlight_line = b
	// }
	//	if b, ok := v.Settings().Get("inverse_caret_state", false).(bool); !b && ok {
	// 	if caret_style == termbox.AttrReverse {
	// 		caret_style = termbox.AttrUnderline
	// 	} else {
	// 		caret_style = termbox.AttrReverse
	// 	}
	// }
	// caret_style := termbox.AttrUnderline
	// if b, ok := v.Settings().Get("caret_style", "underline").(string); ok {
	// 	if b == "block" {
	// 		caret_style = termbox.AttrReverse
	// 	}
	// }

	data.Cull(viewport)
	recipe := make(Recipe)
	for _, v := range data {
		k := scheme.Spice(&v)
		rs := recipe[k]
		rs.AddAll(v.Regions.Regions())

		if rs.HasNonEmpty() {
			var rs2 text.RegionSet
			var last text.Region
			for i, r := range rs.Regions() {
				if i > 0 && r.Begin() == last.End() {
					rs2.Add(r.Cover(last))
				} else {
					rs2.Add(r)
				}
				last = r
			}
			recipe[k] = rs2
		}
	}
	return recipe
}

// Transcribing the Recipe creates a linear step-by-step
// representation of it, which might or might not
// make it easier for Renderers to work with.
func (r Recipe) Transcribe() (ret TranscribedRecipe) {
	for flav, set := range r {
		for _, r := range set.Regions() {
			ret = append(ret, RenderUnit{Flavour: flav, Region: r})
		}
	}
	sort.Sort(&ret)
	return
}

// Just used to satisfy the sort.Interface interface, typically not used otherwise.
func (r *TranscribedRecipe) Len() int {
	return len(*r)
}

// Just used to satisfy the sort.Interface interface, typically not used otherwise.
func (r *TranscribedRecipe) Less(i, j int) bool {
	a, b := (*r)[i].Region, (*r)[j].Region
	if a.Begin() == b.Begin() {
		return a.End() > b.End()
	}
	return a.Begin() < b.Begin()
}

// Just used to satisfy the sort.Interface interface, typically not used otherwise.
func (r *TranscribedRecipe) Swap(i, j int) {
	(*r)[i], (*r)[j] = (*r)[j], (*r)[i]
}
<commit_msg>frontend/html: primitive key events and reload via javascript.<commit_after>package render

import (
	"github.com/quarnster/util/text"
	"image/color"
	"sort"
)

type (
	// Colour represented by a underlying color.RGBA structure
	Colour color.RGBA

	// The Flavour struct contains the specific settings
	// used to style a particular Region.
	Flavour struct {
		Background Colour
		Foreground Colour
		Font       Font
		Flags      ViewRegionFlags
	}

	// The Recipe type groups text.RegionSets by their Flavour.
	// The idea is to allow large groups of text be rendered as
	// a single batch without any state changes inbetween the
	// batches.
	Recipe map[Flavour]text.RegionSet

	// A RenderUnit is just a Flavour and an associated Region.
	RenderUnit struct {
		Flavour Flavour
		Region  text.Region
	}

	// A TranscribedRecipe is a linear (in text.Regions) representation of a Recipe
	TranscribedRecipe []RenderUnit

	ColourScheme interface {
		// Takes a ViewRegions pointer as input and uses the data contained in it
		// to determine the Flavour it should be rendered with.
		Spice(*ViewRegions) Flavour
	}

	Renderer interface {
		// Renders the given Recipe
		Render(Recipe)
	}
)

// Transform takes a ColourScheme, a ViewRegionMap and a viewport as input.
//
// The viewport would be the text.Region of the current buffer that is visible to the user
// and any ViewRegions outside of this area are not forwarded for further processing.
//
// The remaining ViewRegions are then passed on to the ColourScheme for determining the exact Flavour
// for which that RegionSet should be styled, adding Regions of the same Flavour to the same RegionSet.
//
// Typically there are more ViewRegions available in a text buffer than there are unique Flavours in
// a ColourScheme, so this operation can be viewed as reducing the number of state changes required to
// display the text to the user.
//
// The final output, the Recipe, contains a mapping of all unique Flavours and that Flavour's
// associated RegionSet.
func Transform(scheme ColourScheme, data ViewRegionMap, viewport text.Region) Recipe {
	// TODO:
	// 	caret_blink := true
	// if b, ok := v.Settings().Get("caret_blink", true).(bool); ok {
	// 	caret_blink = b
	// }
	//
	// highlight_line := false
	// if b, ok := v.Settings().Get("highlight_line", highlight_line).(bool); ok {
	// 	highlight_line = b
	// }
	//	if b, ok := v.Settings().Get("inverse_caret_state", false).(bool); !b && ok {
	// 	if caret_style == termbox.AttrReverse {
	// 		caret_style = termbox.AttrUnderline
	// 	} else {
	// 		caret_style = termbox.AttrReverse
	// 	}
	// }
	// caret_style := termbox.AttrUnderline
	// if b, ok := v.Settings().Get("caret_style", "underline").(string); ok {
	// 	if b == "block" {
	// 		caret_style = termbox.AttrReverse
	// 	}
	// }

	data.Cull(viewport)
	recipe := make(Recipe)
	for _, v := range data {
		k := scheme.Spice(&v)
		rs := recipe[k]
		rs.AddAll(v.Regions.Regions())
		if rs.HasNonEmpty() {
			var last text.Region
			rgs := rs.Regions()
			for i, r := range rgs {
				if i > 0 && r.Begin() == last.End() {
					rs.Add(r.Cover(last))
				}
				last = r
			}
			recipe[k] = rs
		}
	}
	return recipe
}

// Transcribing the Recipe creates a linear step-by-step
// representation of it, which might or might not
// make it easier for Renderers to work with.
func (r Recipe) Transcribe() (ret TranscribedRecipe) {
	for flav, set := range r {
		for _, r := range set.Regions() {
			ret = append(ret, RenderUnit{Flavour: flav, Region: r})
		}
	}
	sort.Sort(&ret)
	return
}

// Just used to satisfy the sort.Interface interface, typically not used otherwise.
func (r *TranscribedRecipe) Len() int {
	return len(*r)
}

// Just used to satisfy the sort.Interface interface, typically not used otherwise.
func (r *TranscribedRecipe) Less(i, j int) bool {
	a, b := (*r)[i].Region, (*r)[j].Region
	if a.Begin() == b.Begin() {
		return a.End() > b.End()
	}
	return a.Begin() < b.Begin()
}

// Just used to satisfy the sort.Interface interface, typically not used otherwise.
func (r *TranscribedRecipe) Swap(i, j int) {
	(*r)[i], (*r)[j] = (*r)[j], (*r)[i]
}
<|endoftext|>


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

import (
	"github.com/WE-Development/mosel/api"
	"log"
)

type nodeRespHandler struct {
	cache *dataCache
}

func NewNodeRespHandler(cache *dataCache) (*nodeRespHandler,error) {
	return &nodeRespHandler{
		cache: cache,
	}, nil
}

func (handler nodeRespHandler) handleNodeResp(node string, resp api.NodeResponse) {
	//log.Println(resp)
	handler.cache.Add(node, resp.Time, resp.NodeInfo)
}
<commit_msg>remove old debug logging<commit_after>/*
 * Copyright 2016 Robin Engel
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package context

import (
	"github.com/WE-Development/mosel/api"
)

type nodeRespHandler struct {
	cache *dataCache
}

func NewNodeRespHandler(cache *dataCache) (*nodeRespHandler, error) {
	return &nodeRespHandler{
		cache: cache,
	}, nil
}

func (handler nodeRespHandler) handleNodeResp(node string, resp api.NodeResponse) {
	//log.Println(resp)
	handler.cache.Add(node, resp.Time, resp.NodeInfo)
}
<|endoftext|>


<commit_before>package tailtopic

import (
	"fmt"
	"os"
	"os/signal"
)

// TailTopic is the app/command
type TailTopic struct {
	consumer consumer
	decoder  decoder
	printer  printer
	messages chan *message
	closing  chan bool
}

// Start consuming, decoding and printing messages
func (tt *TailTopic) Start() {
	go tt.signalListening()
	go tt.messageListening()
	tt.consume()
}

func (tt *TailTopic) stop() {
	close(tt.closing)
}

func (tt *TailTopic) signalListening() {
	signals := make(chan os.Signal, 1)
	signal.Notify(signals, os.Kill, os.Interrupt)
	<-signals
	tt.stop()
}

func (tt *TailTopic) messageListening() {
	for msg := range tt.messages {
		msgVal, err := tt.decoder.decode(msg.Value)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Failed to decode message! %v\n", err)
			return
		}
		tt.printer.println(msgVal)
	}
}

func (tt *TailTopic) consume() {
	err := tt.consumer.consume(tt.messages, tt.closing)
	if err != nil {
		fmt.Println("Could not start consumer", err)
	}
}

// NewKafkaAvroTailTopic creates new TailTopic with Kafka consumer and Avro Decoder
func NewKafkaAvroTailTopic(topic, offset, broker, schemaregURI string) *TailTopic {
	return &TailTopic{
		&kafkaConsumer{topic, offset, broker},
		newAvroDecoder(schemaregURI),
		&console{},
		make(chan *message, 16),
		make(chan bool)}
}
<commit_msg>Direct error to stderr<commit_after>package tailtopic

import (
	"fmt"
	"os"
	"os/signal"
)

// TailTopic is the app/command
type TailTopic struct {
	consumer consumer
	decoder  decoder
	printer  printer
	messages chan *message
	closing  chan bool
}

// Start consuming, decoding and printing messages
func (tt *TailTopic) Start() {
	go tt.signalListening()
	go tt.messageListening()
	tt.consume()
}

func (tt *TailTopic) stop() {
	close(tt.closing)
}

func (tt *TailTopic) signalListening() {
	signals := make(chan os.Signal, 1)
	signal.Notify(signals, os.Kill, os.Interrupt)
	<-signals
	tt.stop()
}

func (tt *TailTopic) messageListening() {
	for msg := range tt.messages {
		msgVal, err := tt.decoder.decode(msg.Value)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Failed to decode message! %v\n", err)
			return
		}
		tt.printer.println(msgVal)
	}
}

func (tt *TailTopic) consume() {
	err := tt.consumer.consume(tt.messages, tt.closing)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to start consumer! %v\n", err)
	}
}

// NewKafkaAvroTailTopic creates new TailTopic with Kafka consumer and Avro Decoder
func NewKafkaAvroTailTopic(topic, offset, broker, schemaregURI string) *TailTopic {
	return &TailTopic{
		&kafkaConsumer{topic, offset, broker},
		newAvroDecoder(schemaregURI),
		&console{},
		make(chan *message, 16),
		make(chan bool)}
}
<|endoftext|>


<commit_before>package bdiscord

import (
	"bytes"
	"github.com/42wim/matterbridge/bridge/config"
	"github.com/42wim/matterbridge/bridge/helper"
	log "github.com/Sirupsen/logrus"
	"github.com/bwmarrin/discordgo"
	"regexp"
	"strings"
	"sync"
)

type bdiscord struct {
	c              *discordgo.Session
	Channels       []*discordgo.Channel
	Nick           string
	UseChannelID   bool
	userMemberMap  map[string]*discordgo.Member
	guildID        string
	webhookID      string
	webhookToken   string
	channelInfoMap map[string]*config.ChannelInfo
	sync.RWMutex
	*config.BridgeConfig
}

var flog *log.Entry
var protocol = "discord"

func init() {
	flog = log.WithFields(log.Fields{"module": protocol})
}

func New(cfg *config.BridgeConfig) *bdiscord {
	b := &bdiscord{BridgeConfig: cfg}
	b.userMemberMap = make(map[string]*discordgo.Member)
	b.channelInfoMap = make(map[string]*config.ChannelInfo)
	if b.Config.WebhookURL != "" {
		flog.Debug("Configuring Discord Incoming Webhook")
		b.webhookID, b.webhookToken = b.splitURL(b.Config.WebhookURL)
	}
	return b
}

func (b *bdiscord) Connect() error {
	var err error
	flog.Info("Connecting")
	if b.Config.WebhookURL == "" {
		flog.Info("Connecting using token")
	} else {
		flog.Info("Connecting using webhookurl (for posting) and token")
	}
	if !strings.HasPrefix(b.Config.Token, "Bot ") {
		b.Config.Token = "Bot " + b.Config.Token
	}
	b.c, err = discordgo.New(b.Config.Token)
	if err != nil {
		flog.Debugf("%#v", err)
		return err
	}
	flog.Info("Connection succeeded")
	b.c.AddHandler(b.messageCreate)
	b.c.AddHandler(b.memberUpdate)
	b.c.AddHandler(b.messageUpdate)
	b.c.AddHandler(b.messageDelete)
	err = b.c.Open()
	if err != nil {
		flog.Debugf("%#v", err)
		return err
	}
	guilds, err := b.c.UserGuilds(100, "", "")
	if err != nil {
		flog.Debugf("%#v", err)
		return err
	}
	userinfo, err := b.c.User("@me")
	if err != nil {
		flog.Debugf("%#v", err)
		return err
	}
	b.Nick = userinfo.Username
	for _, guild := range guilds {
		if guild.Name == b.Config.Server {
			b.Channels, err = b.c.GuildChannels(guild.ID)
			b.guildID = guild.ID
			if err != nil {
				flog.Debugf("%#v", err)
				return err
			}
		}
	}
	return nil
}

func (b *bdiscord) Disconnect() error {
	return nil
}

func (b *bdiscord) JoinChannel(channel config.ChannelInfo) error {
	b.channelInfoMap[channel.ID] = &channel
	idcheck := strings.Split(channel.Name, "ID:")
	if len(idcheck) > 1 {
		b.UseChannelID = true
	}
	return nil
}

func (b *bdiscord) Send(msg config.Message) (string, error) {
	flog.Debugf("Receiving %#v", msg)
	channelID := b.getChannelID(msg.Channel)
	if channelID == "" {
		flog.Errorf("Could not find channelID for %v", msg.Channel)
		return "", nil
	}
	if msg.Event == config.EVENT_USER_ACTION {
		msg.Text = "_" + msg.Text + "_"
	}

	wID := b.webhookID
	wToken := b.webhookToken
	if ci, ok := b.channelInfoMap[msg.Channel+b.Account]; ok {
		if ci.Options.WebhookURL != "" {
			wID, wToken = b.splitURL(ci.Options.WebhookURL)
		}
	}

	if wID == "" {
		flog.Debugf("Broadcasting using token (API)")
		if msg.Event == config.EVENT_MSG_DELETE {
			if msg.ID == "" {
				return "", nil
			}
			err := b.c.ChannelMessageDelete(channelID, msg.ID)
			return "", err
		}
		if msg.ID != "" {
			_, err := b.c.ChannelMessageEdit(channelID, msg.ID, msg.Username+msg.Text)
			return msg.ID, err
		}

		if msg.Extra != nil {
			for _, rmsg := range helper.HandleExtra(&msg, b.General) {
				b.c.ChannelMessageSend(channelID, rmsg.Username+rmsg.Text)
			}
			// check if we have files to upload (from slack, telegram or mattermost)
			if len(msg.Extra["file"]) > 0 {
				var err error
				for _, f := range msg.Extra["file"] {
					fi := f.(config.FileInfo)
					files := []*discordgo.File{}
					files = append(files, &discordgo.File{fi.Name, "", bytes.NewReader(*fi.Data)})
					_, err = b.c.ChannelMessageSendComplex(channelID, &discordgo.MessageSend{Content: msg.Username + fi.Comment, Files: files})
					if err != nil {
						flog.Errorf("file upload failed: %#v", err)
					}
				}
				return "", nil
			}
		}

		res, err := b.c.ChannelMessageSend(channelID, msg.Username+msg.Text)
		if err != nil {
			return "", err
		}
		return res.ID, err
	}
	flog.Debugf("Broadcasting using Webhook")
	err := b.c.WebhookExecute(
		wID,
		wToken,
		true,
		&discordgo.WebhookParams{
			Content:   msg.Text,
			Username:  msg.Username,
			AvatarURL: msg.Avatar,
		})
	return "", err
}

func (b *bdiscord) messageDelete(s *discordgo.Session, m *discordgo.MessageDelete) {
	rmsg := config.Message{Account: b.Account, ID: m.ID, Event: config.EVENT_MSG_DELETE, Text: config.EVENT_MSG_DELETE}
	rmsg.Channel = b.getChannelName(m.ChannelID)
	if b.UseChannelID {
		rmsg.Channel = "ID:" + m.ChannelID
	}
	flog.Debugf("Sending message from %s to gateway", b.Account)
	flog.Debugf("Message is %#v", rmsg)
	b.Remote <- rmsg
}

func (b *bdiscord) messageUpdate(s *discordgo.Session, m *discordgo.MessageUpdate) {
	if b.Config.EditDisable {
		return
	}
	// only when message is actually edited
	if m.Message.EditedTimestamp != "" {
		flog.Debugf("Sending edit message")
		m.Content = m.Content + b.Config.EditSuffix
		b.messageCreate(s, (*discordgo.MessageCreate)(m))
	}
}

func (b *bdiscord) messageCreate(s *discordgo.Session, m *discordgo.MessageCreate) {
	var err error
	// not relay our own messages
	if m.Author.Username == b.Nick {
		return
	}
	// if using webhooks, do not relay if it's ours
	if b.useWebhook() && m.Author.Bot && b.isWebhookID(m.Author.ID) {
		return
	}

	if len(m.Attachments) > 0 {
		for _, attach := range m.Attachments {
			m.Content = m.Content + "\n" + attach.URL
		}
	}

	var text string
	if m.Content != "" {
		flog.Debugf("Receiving message %#v", m.Message)
		/*
			if len(m.MentionRoles) > 0 {
				m.Message.Content = b.replaceRoleMentions(m.Message.Content)
			}
		*/
		m.Message.Content = b.stripCustomoji(m.Message.Content)
		m.Message.Content = b.replaceChannelMentions(m.Message.Content)
		text, err = m.ContentWithMoreMentionsReplaced(b.c)
		if err != nil {
			flog.Errorf("ContentWithMoreMentionsReplaced failed: %s", err)
			text = m.ContentWithMentionsReplaced()
		}
		//	text = m.ContentWithMentionsReplaced()
	}

	rmsg := config.Message{Account: b.Account, Avatar: "https://cdn.discordapp.com/avatars/" + m.Author.ID + "/" + m.Author.Avatar + ".jpg",
		UserID: m.Author.ID, ID: m.ID}

	rmsg.Channel = b.getChannelName(m.ChannelID)
	if b.UseChannelID {
		rmsg.Channel = "ID:" + m.ChannelID
	}

	if !b.Config.UseUserName {
		rmsg.Username = b.getNick(m.Author)
	} else {
		rmsg.Username = m.Author.Username
	}

	if b.Config.ShowEmbeds && m.Message.Embeds != nil {
		for _, embed := range m.Message.Embeds {
			text = text + "embed: " + embed.Title + " - " + embed.Description + " - " + embed.URL + "\n"
		}
	}

	// no empty messages
	if text == "" {
		return
	}

	text, ok := b.replaceAction(text)
	if ok {
		rmsg.Event = config.EVENT_USER_ACTION
	}

	rmsg.Text = text
	flog.Debugf("Sending message from %s on %s to gateway", m.Author.Username, b.Account)
	flog.Debugf("Message is %#v", rmsg)
	b.Remote <- rmsg
}

func (b *bdiscord) memberUpdate(s *discordgo.Session, m *discordgo.GuildMemberUpdate) {
	b.Lock()
	if _, ok := b.userMemberMap[m.Member.User.ID]; ok {
		flog.Debugf("%s: memberupdate: user %s (nick %s) changes nick to %s", b.Account, m.Member.User.Username, b.userMemberMap[m.Member.User.ID].Nick, m.Member.Nick)
	}
	b.userMemberMap[m.Member.User.ID] = m.Member
	b.Unlock()
}

func (b *bdiscord) getNick(user *discordgo.User) string {
	var err error
	b.Lock()
	defer b.Unlock()
	if _, ok := b.userMemberMap[user.ID]; ok {
		if b.userMemberMap[user.ID] != nil {
			if b.userMemberMap[user.ID].Nick != "" {
				// only return if nick is set
				return b.userMemberMap[user.ID].Nick
			}
			// otherwise return username
			return user.Username
		}
	}
	// if we didn't find nick, search for it
	member, err := b.c.GuildMember(b.guildID, user.ID)
	if err != nil {
		return user.Username
	}
	b.userMemberMap[user.ID] = member
	// only return if nick is set
	if b.userMemberMap[user.ID].Nick != "" {
		return b.userMemberMap[user.ID].Nick
	}
	return user.Username
}

func (b *bdiscord) getChannelID(name string) string {
	idcheck := strings.Split(name, "ID:")
	if len(idcheck) > 1 {
		return idcheck[1]
	}
	for _, channel := range b.Channels {
		if channel.Name == name {
			return channel.ID
		}
	}
	return ""
}

func (b *bdiscord) getChannelName(id string) string {
	for _, channel := range b.Channels {
		if channel.ID == id {
			return channel.Name
		}
	}
	return ""
}

func (b *bdiscord) replaceRoleMentions(text string) string {
	roles, err := b.c.GuildRoles(b.guildID)
	if err != nil {
		flog.Debugf("%#v", string(err.(*discordgo.RESTError).ResponseBody))
		return text
	}
	for _, role := range roles {
		text = strings.Replace(text, "<@&"+role.ID+">", "@"+role.Name, -1)
	}
	return text
}

func (b *bdiscord) replaceChannelMentions(text string) string {
	var err error
	re := regexp.MustCompile("<#[0-9]+>")
	text = re.ReplaceAllStringFunc(text, func(m string) string {
		channel := b.getChannelName(m[2 : len(m)-1])
		// if at first don't succeed, try again
		if channel == "" {
			b.Channels, err = b.c.GuildChannels(b.guildID)
			if err != nil {
				return "#unknownchannel"
			}
			channel = b.getChannelName(m[2 : len(m)-1])
			return "#" + channel
		}
		return "#" + channel
	})
	return text
}

func (b *bdiscord) replaceAction(text string) (string, bool) {
	if strings.HasPrefix(text, "_") && strings.HasSuffix(text, "_") {
		return strings.Replace(text, "_", "", -1), true
	}
	return text, false
}

func (b *bdiscord) stripCustomoji(text string) string {
	// <:doge:302803592035958784>
	re := regexp.MustCompile("<(:.*?:)[0-9]+>")
	return re.ReplaceAllString(text, `$1`)
}

// splitURL splits a webhookURL and returns the id and token
func (b *bdiscord) splitURL(url string) (string, string) {
	webhookURLSplit := strings.Split(url, "/")
	if len(webhookURLSplit) != 7 {
		log.Fatalf("%s is no correct discord WebhookURL", url)
	}
	return webhookURLSplit[len(webhookURLSplit)-2], webhookURLSplit[len(webhookURLSplit)-1]
}

// useWebhook returns true if we have a webhook defined somewhere
func (b *bdiscord) useWebhook() bool {
	if b.Config.WebhookURL != "" {
		return true
	}
	for _, channel := range b.channelInfoMap {
		if channel.Options.WebhookURL != "" {
			return true
		}
	}
	return false
}

// isWebhookID returns true if the specified id is used in a defined webhook
func (b *bdiscord) isWebhookID(id string) bool {
	if b.Config.WebhookURL != "" {
		wID, _ := b.splitURL(b.Config.WebhookURL)
		if wID == id {
			return true
		}
	}
	for _, channel := range b.channelInfoMap {
		if channel.Options.WebhookURL != "" {
			wID, _ := b.splitURL(channel.Options.WebhookURL)
			if wID == id {
				return true
			}
		}
	}
	return false
}
<commit_msg>Remove unused code<commit_after>package bdiscord

import (
	"bytes"
	"github.com/42wim/matterbridge/bridge/config"
	"github.com/42wim/matterbridge/bridge/helper"
	log "github.com/Sirupsen/logrus"
	"github.com/bwmarrin/discordgo"
	"regexp"
	"strings"
	"sync"
)

type bdiscord struct {
	c              *discordgo.Session
	Channels       []*discordgo.Channel
	Nick           string
	UseChannelID   bool
	userMemberMap  map[string]*discordgo.Member
	guildID        string
	webhookID      string
	webhookToken   string
	channelInfoMap map[string]*config.ChannelInfo
	sync.RWMutex
	*config.BridgeConfig
}

var flog *log.Entry
var protocol = "discord"

func init() {
	flog = log.WithFields(log.Fields{"module": protocol})
}

func New(cfg *config.BridgeConfig) *bdiscord {
	b := &bdiscord{BridgeConfig: cfg}
	b.userMemberMap = make(map[string]*discordgo.Member)
	b.channelInfoMap = make(map[string]*config.ChannelInfo)
	if b.Config.WebhookURL != "" {
		flog.Debug("Configuring Discord Incoming Webhook")
		b.webhookID, b.webhookToken = b.splitURL(b.Config.WebhookURL)
	}
	return b
}

func (b *bdiscord) Connect() error {
	var err error
	flog.Info("Connecting")
	if b.Config.WebhookURL == "" {
		flog.Info("Connecting using token")
	} else {
		flog.Info("Connecting using webhookurl (for posting) and token")
	}
	if !strings.HasPrefix(b.Config.Token, "Bot ") {
		b.Config.Token = "Bot " + b.Config.Token
	}
	b.c, err = discordgo.New(b.Config.Token)
	if err != nil {
		flog.Debugf("%#v", err)
		return err
	}
	flog.Info("Connection succeeded")
	b.c.AddHandler(b.messageCreate)
	b.c.AddHandler(b.memberUpdate)
	b.c.AddHandler(b.messageUpdate)
	b.c.AddHandler(b.messageDelete)
	err = b.c.Open()
	if err != nil {
		flog.Debugf("%#v", err)
		return err
	}
	guilds, err := b.c.UserGuilds(100, "", "")
	if err != nil {
		flog.Debugf("%#v", err)
		return err
	}
	userinfo, err := b.c.User("@me")
	if err != nil {
		flog.Debugf("%#v", err)
		return err
	}
	b.Nick = userinfo.Username
	for _, guild := range guilds {
		if guild.Name == b.Config.Server {
			b.Channels, err = b.c.GuildChannels(guild.ID)
			b.guildID = guild.ID
			if err != nil {
				flog.Debugf("%#v", err)
				return err
			}
		}
	}
	return nil
}

func (b *bdiscord) Disconnect() error {
	return nil
}

func (b *bdiscord) JoinChannel(channel config.ChannelInfo) error {
	b.channelInfoMap[channel.ID] = &channel
	idcheck := strings.Split(channel.Name, "ID:")
	if len(idcheck) > 1 {
		b.UseChannelID = true
	}
	return nil
}

func (b *bdiscord) Send(msg config.Message) (string, error) {
	flog.Debugf("Receiving %#v", msg)
	channelID := b.getChannelID(msg.Channel)
	if channelID == "" {
		flog.Errorf("Could not find channelID for %v", msg.Channel)
		return "", nil
	}
	if msg.Event == config.EVENT_USER_ACTION {
		msg.Text = "_" + msg.Text + "_"
	}

	wID := b.webhookID
	wToken := b.webhookToken
	if ci, ok := b.channelInfoMap[msg.Channel+b.Account]; ok {
		if ci.Options.WebhookURL != "" {
			wID, wToken = b.splitURL(ci.Options.WebhookURL)
		}
	}

	if wID == "" {
		flog.Debugf("Broadcasting using token (API)")
		if msg.Event == config.EVENT_MSG_DELETE {
			if msg.ID == "" {
				return "", nil
			}
			err := b.c.ChannelMessageDelete(channelID, msg.ID)
			return "", err
		}
		if msg.ID != "" {
			_, err := b.c.ChannelMessageEdit(channelID, msg.ID, msg.Username+msg.Text)
			return msg.ID, err
		}

		if msg.Extra != nil {
			for _, rmsg := range helper.HandleExtra(&msg, b.General) {
				b.c.ChannelMessageSend(channelID, rmsg.Username+rmsg.Text)
			}
			// check if we have files to upload (from slack, telegram or mattermost)
			if len(msg.Extra["file"]) > 0 {
				var err error
				for _, f := range msg.Extra["file"] {
					fi := f.(config.FileInfo)
					files := []*discordgo.File{}
					files = append(files, &discordgo.File{fi.Name, "", bytes.NewReader(*fi.Data)})
					_, err = b.c.ChannelMessageSendComplex(channelID, &discordgo.MessageSend{Content: msg.Username + fi.Comment, Files: files})
					if err != nil {
						flog.Errorf("file upload failed: %#v", err)
					}
				}
				return "", nil
			}
		}

		res, err := b.c.ChannelMessageSend(channelID, msg.Username+msg.Text)
		if err != nil {
			return "", err
		}
		return res.ID, err
	}
	flog.Debugf("Broadcasting using Webhook")
	err := b.c.WebhookExecute(
		wID,
		wToken,
		true,
		&discordgo.WebhookParams{
			Content:   msg.Text,
			Username:  msg.Username,
			AvatarURL: msg.Avatar,
		})
	return "", err
}

func (b *bdiscord) messageDelete(s *discordgo.Session, m *discordgo.MessageDelete) {
	rmsg := config.Message{Account: b.Account, ID: m.ID, Event: config.EVENT_MSG_DELETE, Text: config.EVENT_MSG_DELETE}
	rmsg.Channel = b.getChannelName(m.ChannelID)
	if b.UseChannelID {
		rmsg.Channel = "ID:" + m.ChannelID
	}
	flog.Debugf("Sending message from %s to gateway", b.Account)
	flog.Debugf("Message is %#v", rmsg)
	b.Remote <- rmsg
}

func (b *bdiscord) messageUpdate(s *discordgo.Session, m *discordgo.MessageUpdate) {
	if b.Config.EditDisable {
		return
	}
	// only when message is actually edited
	if m.Message.EditedTimestamp != "" {
		flog.Debugf("Sending edit message")
		m.Content = m.Content + b.Config.EditSuffix
		b.messageCreate(s, (*discordgo.MessageCreate)(m))
	}
}

func (b *bdiscord) messageCreate(s *discordgo.Session, m *discordgo.MessageCreate) {
	var err error
	// not relay our own messages
	if m.Author.Username == b.Nick {
		return
	}
	// if using webhooks, do not relay if it's ours
	if b.useWebhook() && m.Author.Bot && b.isWebhookID(m.Author.ID) {
		return
	}

	if len(m.Attachments) > 0 {
		for _, attach := range m.Attachments {
			m.Content = m.Content + "\n" + attach.URL
		}
	}

	var text string
	if m.Content != "" {
		flog.Debugf("Receiving message %#v", m.Message)
		m.Message.Content = b.stripCustomoji(m.Message.Content)
		m.Message.Content = b.replaceChannelMentions(m.Message.Content)
		text, err = m.ContentWithMoreMentionsReplaced(b.c)
		if err != nil {
			flog.Errorf("ContentWithMoreMentionsReplaced failed: %s", err)
			text = m.ContentWithMentionsReplaced()
		}
	}

	rmsg := config.Message{Account: b.Account, Avatar: "https://cdn.discordapp.com/avatars/" + m.Author.ID + "/" + m.Author.Avatar + ".jpg",
		UserID: m.Author.ID, ID: m.ID}

	rmsg.Channel = b.getChannelName(m.ChannelID)
	if b.UseChannelID {
		rmsg.Channel = "ID:" + m.ChannelID
	}

	if !b.Config.UseUserName {
		rmsg.Username = b.getNick(m.Author)
	} else {
		rmsg.Username = m.Author.Username
	}

	if b.Config.ShowEmbeds && m.Message.Embeds != nil {
		for _, embed := range m.Message.Embeds {
			text = text + "embed: " + embed.Title + " - " + embed.Description + " - " + embed.URL + "\n"
		}
	}

	// no empty messages
	if text == "" {
		return
	}

	text, ok := b.replaceAction(text)
	if ok {
		rmsg.Event = config.EVENT_USER_ACTION
	}

	rmsg.Text = text
	flog.Debugf("Sending message from %s on %s to gateway", m.Author.Username, b.Account)
	flog.Debugf("Message is %#v", rmsg)
	b.Remote <- rmsg
}

func (b *bdiscord) memberUpdate(s *discordgo.Session, m *discordgo.GuildMemberUpdate) {
	b.Lock()
	if _, ok := b.userMemberMap[m.Member.User.ID]; ok {
		flog.Debugf("%s: memberupdate: user %s (nick %s) changes nick to %s", b.Account, m.Member.User.Username, b.userMemberMap[m.Member.User.ID].Nick, m.Member.Nick)
	}
	b.userMemberMap[m.Member.User.ID] = m.Member
	b.Unlock()
}

func (b *bdiscord) getNick(user *discordgo.User) string {
	var err error
	b.Lock()
	defer b.Unlock()
	if _, ok := b.userMemberMap[user.ID]; ok {
		if b.userMemberMap[user.ID] != nil {
			if b.userMemberMap[user.ID].Nick != "" {
				// only return if nick is set
				return b.userMemberMap[user.ID].Nick
			}
			// otherwise return username
			return user.Username
		}
	}
	// if we didn't find nick, search for it
	member, err := b.c.GuildMember(b.guildID, user.ID)
	if err != nil {
		return user.Username
	}
	b.userMemberMap[user.ID] = member
	// only return if nick is set
	if b.userMemberMap[user.ID].Nick != "" {
		return b.userMemberMap[user.ID].Nick
	}
	return user.Username
}

func (b *bdiscord) getChannelID(name string) string {
	idcheck := strings.Split(name, "ID:")
	if len(idcheck) > 1 {
		return idcheck[1]
	}
	for _, channel := range b.Channels {
		if channel.Name == name {
			return channel.ID
		}
	}
	return ""
}

func (b *bdiscord) getChannelName(id string) string {
	for _, channel := range b.Channels {
		if channel.ID == id {
			return channel.Name
		}
	}
	return ""
}

func (b *bdiscord) replaceChannelMentions(text string) string {
	var err error
	re := regexp.MustCompile("<#[0-9]+>")
	text = re.ReplaceAllStringFunc(text, func(m string) string {
		channel := b.getChannelName(m[2 : len(m)-1])
		// if at first don't succeed, try again
		if channel == "" {
			b.Channels, err = b.c.GuildChannels(b.guildID)
			if err != nil {
				return "#unknownchannel"
			}
			channel = b.getChannelName(m[2 : len(m)-1])
			return "#" + channel
		}
		return "#" + channel
	})
	return text
}

func (b *bdiscord) replaceAction(text string) (string, bool) {
	if strings.HasPrefix(text, "_") && strings.HasSuffix(text, "_") {
		return strings.Replace(text, "_", "", -1), true
	}
	return text, false
}

func (b *bdiscord) stripCustomoji(text string) string {
	// <:doge:302803592035958784>
	re := regexp.MustCompile("<(:.*?:)[0-9]+>")
	return re.ReplaceAllString(text, `$1`)
}

// splitURL splits a webhookURL and returns the id and token
func (b *bdiscord) splitURL(url string) (string, string) {
	webhookURLSplit := strings.Split(url, "/")
	if len(webhookURLSplit) != 7 {
		log.Fatalf("%s is no correct discord WebhookURL", url)
	}
	return webhookURLSplit[len(webhookURLSplit)-2], webhookURLSplit[len(webhookURLSplit)-1]
}

// useWebhook returns true if we have a webhook defined somewhere
func (b *bdiscord) useWebhook() bool {
	if b.Config.WebhookURL != "" {
		return true
	}
	for _, channel := range b.channelInfoMap {
		if channel.Options.WebhookURL != "" {
			return true
		}
	}
	return false
}

// isWebhookID returns true if the specified id is used in a defined webhook
func (b *bdiscord) isWebhookID(id string) bool {
	if b.Config.WebhookURL != "" {
		wID, _ := b.splitURL(b.Config.WebhookURL)
		if wID == id {
			return true
		}
	}
	for _, channel := range b.channelInfoMap {
		if channel.Options.WebhookURL != "" {
			wID, _ := b.splitURL(channel.Options.WebhookURL)
			if wID == id {
				return true
			}
		}
	}
	return false
}
<|endoftext|>


<commit_before>package broker_test

import (
	"net/http"
	"net/http/httptest"

	"github.com/apoydence/talaria/broker"
	"github.com/apoydence/talaria/messages"
	"github.com/gogo/protobuf/proto"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("Connection", func() {
	var (
		recorder   *httptest.ResponseRecorder
		req        *http.Request
		connection *broker.Connection
		server     *httptest.Server
		mockServer *mockServer
	)

	BeforeEach(func() {
		recorder = httptest.NewRecorder()
		req, _ = http.NewRequest("GET", "http://some.url", nil)
		mockServer = newMockServer()
		server = httptest.NewServer(mockServer)
		var err error
		connection, err = broker.NewConnection("ws" + server.URL[4:])
		Expect(err).ToNot(HaveOccurred())
	})

	Describe("FetchFile", func() {

		AfterEach(func() {
			connection.Close()
			server.Close()
		})

		It("returns an error", func(done Done) {
			defer close(done)

			mockServer.serverCh <- buildError(99, "some-error")

			go func() {
				defer GinkgoRecover()
				err := connection.FetchFile(9, "some-file")
				Expect(err).ToNot(BeNil())
				Expect(err.Error()).To(Equal("some-error"))
			}()

			var clientMsg *messages.Client
			Eventually(mockServer.clientCh).Should(Receive(&clientMsg))
			Expect(clientMsg.GetMessageType()).To(Equal(messages.Client_FetchFile))
			Expect(clientMsg.FetchFile.GetName()).To(Equal("some-file"))
			Expect(clientMsg.FetchFile.GetFileId()).To(BeEquivalentTo(9))
		})

		It("returns the file ID", func(done Done) {
			defer close(done)

			mockServer.serverCh <- buildFileLocation(99)

			go func() {
				defer GinkgoRecover()
				err := connection.FetchFile(9, "some-file")
				Expect(err).To(BeNil())
			}()

			var clientMsg *messages.Client
			Eventually(mockServer.clientCh).Should(Receive(&clientMsg))
			Expect(clientMsg.GetMessageType()).To(Equal(messages.Client_FetchFile))
			Expect(clientMsg.FetchFile).ToNot(BeNil())
			Expect(clientMsg.FetchFile.GetName()).To(Equal("some-file"))
		})

		It("returns a redirect error", func(done Done) {
			defer close(done)

			mockServer.serverCh <- buildRemoteFileLocation(99, "http://some.uri")

			go func() {
				defer GinkgoRecover()
				err := connection.FetchFile(9, "some-file")
				Expect(err).ToNot(BeNil())
				Expect(err.Uri).To(Equal("http://some.uri"))
			}()

			var clientMsg *messages.Client
			Eventually(mockServer.clientCh).Should(Receive(&clientMsg))
			Expect(clientMsg.GetMessageType()).To(Equal(messages.Client_FetchFile))
			Expect(clientMsg.FetchFile).ToNot(BeNil())
			Expect(clientMsg.FetchFile.GetName()).To(Equal("some-file"))
		})
	})

	Describe("WriteToFile", func() {

		AfterEach(func() {
			connection.Close()
			server.Close()
		})

		It("returns an error", func(done Done) {
			defer close(done)

			expectedData := []byte("some-data")
			mockServer.serverCh <- buildError(99, "some-error")

			go func() {
				defer GinkgoRecover()
				_, err := connection.WriteToFile(8, expectedData)
				Expect(err).To(HaveOccurred())
				Expect(err.Error()).To(Equal("some-error"))
			}()

			var clientMsg *messages.Client
			Eventually(mockServer.clientCh).Should(Receive(&clientMsg))
			Expect(clientMsg.GetMessageType()).To(Equal(messages.Client_WriteToFile))
			Expect(clientMsg.WriteToFile).ToNot(BeNil())
			Expect(clientMsg.WriteToFile.GetFileId()).To(BeEquivalentTo(8))
			Expect(clientMsg.WriteToFile.GetData()).To(Equal(expectedData))
		})

		It("returns the new file offset", func(done Done) {
			defer close(done)

			expectedData := []byte("some-data")
			mockServer.serverCh <- buildFileOffset(99, 101)

			go func() {
				defer GinkgoRecover()
				fileOffset, err := connection.WriteToFile(8, expectedData)
				Expect(err).ToNot(HaveOccurred())
				Expect(fileOffset).To(BeEquivalentTo(101))
			}()

			var clientMsg *messages.Client
			Eventually(mockServer.clientCh).Should(Receive(&clientMsg))
			Expect(clientMsg.GetMessageType()).To(Equal(messages.Client_WriteToFile))
			Expect(clientMsg.WriteToFile).ToNot(BeNil())
			Expect(clientMsg.WriteToFile.GetFileId()).To(BeEquivalentTo(8))
			Expect(clientMsg.WriteToFile.GetData()).To(Equal(expectedData))
		})
	})

	Describe("ReadFromFile", func() {

		AfterEach(func() {
			connection.Close()
			server.Close()
		})

		It("returns an error", func(done Done) {
			defer close(done)

			mockServer.serverCh <- buildError(99, "some-error")

			go func() {
				defer GinkgoRecover()
				_, err := connection.ReadFromFile(8)
				Expect(err).To(HaveOccurred())
				Expect(err.Error()).To(Equal("some-error"))
			}()

			var clientMsg *messages.Client
			Eventually(mockServer.clientCh).Should(Receive(&clientMsg))
			Expect(clientMsg.GetMessageType()).To(Equal(messages.Client_ReadFromFile))
			Expect(clientMsg.ReadFromFile).ToNot(BeNil())
			Expect(clientMsg.ReadFromFile.GetFileId()).To(BeEquivalentTo(8))
		})

		It("returns the data and offset", func(done Done) {
			defer close(done)

			expectedData := []byte("some-data")
			mockServer.serverCh <- buildReadData(99, expectedData)

			go func() {
				defer GinkgoRecover()
				data, err := connection.ReadFromFile(8)
				Expect(err).ToNot(HaveOccurred())
				Expect(data).To(Equal(expectedData))
			}()

			var clientMsg *messages.Client
			Eventually(mockServer.clientCh).Should(Receive(&clientMsg))
			Expect(clientMsg.GetMessageType()).To(Equal(messages.Client_ReadFromFile))
			Expect(clientMsg.ReadFromFile).ToNot(BeNil())
			Expect(clientMsg.ReadFromFile.GetFileId()).To(BeEquivalentTo(8))
		})
	})

	Describe("Close", func() {
		It("closes the connection to the server", func(done Done) {
			defer close(done)
			connection.Close()
			server.Close()
		})
	})

})

func buildError(messageId uint64, errStr string) []byte {
	msgType := messages.Server_Error
	server := &messages.Server{
		MessageType: &msgType,
		MessageId:   proto.Uint64(messageId),
		Error: &messages.Error{
			Message: proto.String(errStr),
		},
	}

	data, err := server.Marshal()
	Expect(err).ToNot(HaveOccurred())
	return data
}

func buildFileLocation(messageId uint64) []byte {
	msgType := messages.Server_FileLocation
	server := &messages.Server{
		MessageType: &msgType,
		MessageId:   proto.Uint64(messageId),
		FileLocation: &messages.FileLocation{
			Local: proto.Bool(true),
		},
	}

	data, err := server.Marshal()
	Expect(err).ToNot(HaveOccurred())
	return data
}

func buildRemoteFileLocation(messageId uint64, uri string) []byte {
	msgType := messages.Server_FileLocation
	server := &messages.Server{
		MessageType: &msgType,
		MessageId:   proto.Uint64(messageId),
		FileLocation: &messages.FileLocation{
			Local: proto.Bool(false),
			Uri:   proto.String(uri),
		},
	}

	data, err := server.Marshal()
	Expect(err).ToNot(HaveOccurred())
	return data
}

func buildFileOffset(messageId uint64, offset int64) []byte {
	msgType := messages.Server_FileOffset
	server := &messages.Server{
		MessageType: &msgType,
		MessageId:   proto.Uint64(messageId),
		FileOffset: &messages.FileOffset{
			Offset: proto.Int64(offset),
		},
	}

	data, err := server.Marshal()
	Expect(err).ToNot(HaveOccurred())
	return data
}

func buildReadData(messageId uint64, data []byte) []byte {
	msgType := messages.Server_ReadData
	server := &messages.Server{
		MessageType: &msgType,
		MessageId:   proto.Uint64(messageId),
		ReadData: &messages.ReadData{
			Data: data,
		},
	}

	data, err := server.Marshal()
	Expect(err).ToNot(HaveOccurred())
	return data
}
<commit_msg>Clean up connection test<commit_after>package broker_test

import (
	"net/http/httptest"

	"github.com/apoydence/talaria/broker"
	"github.com/apoydence/talaria/messages"
	"github.com/gogo/protobuf/proto"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("Connection", func() {
	var (
		connection *broker.Connection
		server     *httptest.Server
		mockServer *mockServer
	)

	BeforeEach(func() {
		mockServer = newMockServer()
		server = httptest.NewServer(mockServer)
		var err error
		connection, err = broker.NewConnection("ws" + server.URL[4:])
		Expect(err).ToNot(HaveOccurred())
	})

	Describe("FetchFile", func() {

		AfterEach(func() {
			connection.Close()
			server.Close()
		})

		It("returns an error", func(done Done) {
			defer close(done)

			mockServer.serverCh <- buildError(99, "some-error")

			go func() {
				defer GinkgoRecover()
				err := connection.FetchFile(9, "some-file")
				Expect(err).ToNot(BeNil())
				Expect(err.Error()).To(Equal("some-error"))
			}()

			var clientMsg *messages.Client
			Eventually(mockServer.clientCh).Should(Receive(&clientMsg))
			Expect(clientMsg.GetMessageType()).To(Equal(messages.Client_FetchFile))
			Expect(clientMsg.FetchFile.GetName()).To(Equal("some-file"))
			Expect(clientMsg.FetchFile.GetFileId()).To(BeEquivalentTo(9))
		})

		It("returns the file ID", func(done Done) {
			defer close(done)

			mockServer.serverCh <- buildFileLocation(99)

			go func() {
				defer GinkgoRecover()
				err := connection.FetchFile(9, "some-file")
				Expect(err).To(BeNil())
			}()

			var clientMsg *messages.Client
			Eventually(mockServer.clientCh).Should(Receive(&clientMsg))
			Expect(clientMsg.GetMessageType()).To(Equal(messages.Client_FetchFile))
			Expect(clientMsg.FetchFile).ToNot(BeNil())
			Expect(clientMsg.FetchFile.GetName()).To(Equal("some-file"))
		})

		It("returns a redirect error", func(done Done) {
			defer close(done)

			mockServer.serverCh <- buildRemoteFileLocation(99, "http://some.uri")

			go func() {
				defer GinkgoRecover()
				err := connection.FetchFile(9, "some-file")
				Expect(err).ToNot(BeNil())
				Expect(err.Uri).To(Equal("http://some.uri"))
			}()

			var clientMsg *messages.Client
			Eventually(mockServer.clientCh).Should(Receive(&clientMsg))
			Expect(clientMsg.GetMessageType()).To(Equal(messages.Client_FetchFile))
			Expect(clientMsg.FetchFile).ToNot(BeNil())
			Expect(clientMsg.FetchFile.GetName()).To(Equal("some-file"))
		})
	})

	Describe("WriteToFile", func() {

		AfterEach(func() {
			connection.Close()
			server.Close()
		})

		It("returns an error", func(done Done) {
			defer close(done)

			expectedData := []byte("some-data")
			mockServer.serverCh <- buildError(99, "some-error")

			go func() {
				defer GinkgoRecover()
				_, err := connection.WriteToFile(8, expectedData)
				Expect(err).To(HaveOccurred())
				Expect(err.Error()).To(Equal("some-error"))
			}()

			var clientMsg *messages.Client
			Eventually(mockServer.clientCh).Should(Receive(&clientMsg))
			Expect(clientMsg.GetMessageType()).To(Equal(messages.Client_WriteToFile))
			Expect(clientMsg.WriteToFile).ToNot(BeNil())
			Expect(clientMsg.WriteToFile.GetFileId()).To(BeEquivalentTo(8))
			Expect(clientMsg.WriteToFile.GetData()).To(Equal(expectedData))
		})

		It("returns the new file offset", func(done Done) {
			defer close(done)

			expectedData := []byte("some-data")
			mockServer.serverCh <- buildFileOffset(99, 101)

			go func() {
				defer GinkgoRecover()
				fileOffset, err := connection.WriteToFile(8, expectedData)
				Expect(err).ToNot(HaveOccurred())
				Expect(fileOffset).To(BeEquivalentTo(101))
			}()

			var clientMsg *messages.Client
			Eventually(mockServer.clientCh).Should(Receive(&clientMsg))
			Expect(clientMsg.GetMessageType()).To(Equal(messages.Client_WriteToFile))
			Expect(clientMsg.WriteToFile).ToNot(BeNil())
			Expect(clientMsg.WriteToFile.GetFileId()).To(BeEquivalentTo(8))
			Expect(clientMsg.WriteToFile.GetData()).To(Equal(expectedData))
		})
	})

	Describe("ReadFromFile", func() {

		AfterEach(func() {
			connection.Close()
			server.Close()
		})

		It("returns an error", func(done Done) {
			defer close(done)

			mockServer.serverCh <- buildError(99, "some-error")

			go func() {
				defer GinkgoRecover()
				_, err := connection.ReadFromFile(8)
				Expect(err).To(HaveOccurred())
				Expect(err.Error()).To(Equal("some-error"))
			}()

			var clientMsg *messages.Client
			Eventually(mockServer.clientCh).Should(Receive(&clientMsg))
			Expect(clientMsg.GetMessageType()).To(Equal(messages.Client_ReadFromFile))
			Expect(clientMsg.ReadFromFile).ToNot(BeNil())
			Expect(clientMsg.ReadFromFile.GetFileId()).To(BeEquivalentTo(8))
		})

		It("returns the data and offset", func(done Done) {
			defer close(done)

			expectedData := []byte("some-data")
			mockServer.serverCh <- buildReadData(99, expectedData)

			go func() {
				defer GinkgoRecover()
				data, err := connection.ReadFromFile(8)
				Expect(err).ToNot(HaveOccurred())
				Expect(data).To(Equal(expectedData))
			}()

			var clientMsg *messages.Client
			Eventually(mockServer.clientCh).Should(Receive(&clientMsg))
			Expect(clientMsg.GetMessageType()).To(Equal(messages.Client_ReadFromFile))
			Expect(clientMsg.ReadFromFile).ToNot(BeNil())
			Expect(clientMsg.ReadFromFile.GetFileId()).To(BeEquivalentTo(8))
		})
	})

	Describe("Close", func() {
		It("closes the connection to the server", func(done Done) {
			defer close(done)
			connection.Close()
			server.Close()
		})
	})

})

func buildError(messageId uint64, errStr string) []byte {
	msgType := messages.Server_Error
	server := &messages.Server{
		MessageType: &msgType,
		MessageId:   proto.Uint64(messageId),
		Error: &messages.Error{
			Message: proto.String(errStr),
		},
	}

	data, err := server.Marshal()
	Expect(err).ToNot(HaveOccurred())
	return data
}

func buildFileLocation(messageId uint64) []byte {
	msgType := messages.Server_FileLocation
	server := &messages.Server{
		MessageType: &msgType,
		MessageId:   proto.Uint64(messageId),
		FileLocation: &messages.FileLocation{
			Local: proto.Bool(true),
		},
	}

	data, err := server.Marshal()
	Expect(err).ToNot(HaveOccurred())
	return data
}

func buildRemoteFileLocation(messageId uint64, uri string) []byte {
	msgType := messages.Server_FileLocation
	server := &messages.Server{
		MessageType: &msgType,
		MessageId:   proto.Uint64(messageId),
		FileLocation: &messages.FileLocation{
			Local: proto.Bool(false),
			Uri:   proto.String(uri),
		},
	}

	data, err := server.Marshal()
	Expect(err).ToNot(HaveOccurred())
	return data
}

func buildFileOffset(messageId uint64, offset int64) []byte {
	msgType := messages.Server_FileOffset
	server := &messages.Server{
		MessageType: &msgType,
		MessageId:   proto.Uint64(messageId),
		FileOffset: &messages.FileOffset{
			Offset: proto.Int64(offset),
		},
	}

	data, err := server.Marshal()
	Expect(err).ToNot(HaveOccurred())
	return data
}

func buildReadData(messageId uint64, data []byte) []byte {
	msgType := messages.Server_ReadData
	server := &messages.Server{
		MessageType: &msgType,
		MessageId:   proto.Uint64(messageId),
		ReadData: &messages.ReadData{
			Data: data,
		},
	}

	data, err := server.Marshal()
	Expect(err).ToNot(HaveOccurred())
	return data
}
<|endoftext|>